diff --git a/venv/bin/Activate.ps1 b/venv/bin/Activate.ps1 deleted file mode 100644 index b49d77ba44b24fe6d69f6bbe75139b3b5dc23075..0000000000000000000000000000000000000000 --- a/venv/bin/Activate.ps1 +++ /dev/null @@ -1,247 +0,0 @@ -<# -.Synopsis -Activate a Python virtual environment for the current PowerShell session. - -.Description -Pushes the python executable for a virtual environment to the front of the -$Env:PATH environment variable and sets the prompt to signify that you are -in a Python virtual environment. Makes use of the command line switches as -well as the `pyvenv.cfg` file values present in the virtual environment. - -.Parameter VenvDir -Path to the directory that contains the virtual environment to activate. The -default value for this is the parent of the directory that the Activate.ps1 -script is located within. - -.Parameter Prompt -The prompt prefix to display when this virtual environment is activated. By -default, this prompt is the name of the virtual environment folder (VenvDir) -surrounded by parentheses and followed by a single space (ie. '(.venv) '). - -.Example -Activate.ps1 -Activates the Python virtual environment that contains the Activate.ps1 script. - -.Example -Activate.ps1 -Verbose -Activates the Python virtual environment that contains the Activate.ps1 script, -and shows extra information about the activation as it executes. - -.Example -Activate.ps1 -VenvDir C:\Users\MyUser\Common\.venv -Activates the Python virtual environment located in the specified location. - -.Example -Activate.ps1 -Prompt "MyPython" -Activates the Python virtual environment that contains the Activate.ps1 script, -and prefixes the current prompt with the specified string (surrounded in -parentheses) while the virtual environment is active. - -.Notes -On Windows, it may be required to enable this Activate.ps1 script by setting the -execution policy for the user. You can do this by issuing the following PowerShell -command: - -PS C:\> Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser - -For more information on Execution Policies: -https://go.microsoft.com/fwlink/?LinkID=135170 - -#> -Param( - [Parameter(Mandatory = $false)] - [String] - $VenvDir, - [Parameter(Mandatory = $false)] - [String] - $Prompt -) - -<# Function declarations --------------------------------------------------- #> - -<# -.Synopsis -Remove all shell session elements added by the Activate script, including the -addition of the virtual environment's Python executable from the beginning of -the PATH variable. - -.Parameter NonDestructive -If present, do not remove this function from the global namespace for the -session. - -#> -function global:deactivate ([switch]$NonDestructive) { - # Revert to original values - - # The prior prompt: - if (Test-Path -Path Function:_OLD_VIRTUAL_PROMPT) { - Copy-Item -Path Function:_OLD_VIRTUAL_PROMPT -Destination Function:prompt - Remove-Item -Path Function:_OLD_VIRTUAL_PROMPT - } - - # The prior PYTHONHOME: - if (Test-Path -Path Env:_OLD_VIRTUAL_PYTHONHOME) { - Copy-Item -Path Env:_OLD_VIRTUAL_PYTHONHOME -Destination Env:PYTHONHOME - Remove-Item -Path Env:_OLD_VIRTUAL_PYTHONHOME - } - - # The prior PATH: - if (Test-Path -Path Env:_OLD_VIRTUAL_PATH) { - Copy-Item -Path Env:_OLD_VIRTUAL_PATH -Destination Env:PATH - Remove-Item -Path Env:_OLD_VIRTUAL_PATH - } - - # Just remove the VIRTUAL_ENV altogether: - if (Test-Path -Path Env:VIRTUAL_ENV) { - Remove-Item -Path env:VIRTUAL_ENV - } - - # Just remove VIRTUAL_ENV_PROMPT altogether. - if (Test-Path -Path Env:VIRTUAL_ENV_PROMPT) { - Remove-Item -Path env:VIRTUAL_ENV_PROMPT - } - - # Just remove the _PYTHON_VENV_PROMPT_PREFIX altogether: - if (Get-Variable -Name "_PYTHON_VENV_PROMPT_PREFIX" -ErrorAction SilentlyContinue) { - Remove-Variable -Name _PYTHON_VENV_PROMPT_PREFIX -Scope Global -Force - } - - # Leave deactivate function in the global namespace if requested: - if (-not $NonDestructive) { - Remove-Item -Path function:deactivate - } -} - -<# -.Description -Get-PyVenvConfig parses the values from the pyvenv.cfg file located in the -given folder, and returns them in a map. - -For each line in the pyvenv.cfg file, if that line can be parsed into exactly -two strings separated by `=` (with any amount of whitespace surrounding the =) -then it is considered a `key = value` line. The left hand string is the key, -the right hand is the value. - -If the value starts with a `'` or a `"` then the first and last character is -stripped from the value before being captured. - -.Parameter ConfigDir -Path to the directory that contains the `pyvenv.cfg` file. -#> -function Get-PyVenvConfig( - [String] - $ConfigDir -) { - Write-Verbose "Given ConfigDir=$ConfigDir, obtain values in pyvenv.cfg" - - # Ensure the file exists, and issue a warning if it doesn't (but still allow the function to continue). - $pyvenvConfigPath = Join-Path -Resolve -Path $ConfigDir -ChildPath 'pyvenv.cfg' -ErrorAction Continue - - # An empty map will be returned if no config file is found. - $pyvenvConfig = @{ } - - if ($pyvenvConfigPath) { - - Write-Verbose "File exists, parse `key = value` lines" - $pyvenvConfigContent = Get-Content -Path $pyvenvConfigPath - - $pyvenvConfigContent | ForEach-Object { - $keyval = $PSItem -split "\s*=\s*", 2 - if ($keyval[0] -and $keyval[1]) { - $val = $keyval[1] - - # Remove extraneous quotations around a string value. - if ("'""".Contains($val.Substring(0, 1))) { - $val = $val.Substring(1, $val.Length - 2) - } - - $pyvenvConfig[$keyval[0]] = $val - Write-Verbose "Adding Key: '$($keyval[0])'='$val'" - } - } - } - return $pyvenvConfig -} - - -<# Begin Activate script --------------------------------------------------- #> - -# Determine the containing directory of this script -$VenvExecPath = Split-Path -Parent $MyInvocation.MyCommand.Definition -$VenvExecDir = Get-Item -Path $VenvExecPath - -Write-Verbose "Activation script is located in path: '$VenvExecPath'" -Write-Verbose "VenvExecDir Fullname: '$($VenvExecDir.FullName)" -Write-Verbose "VenvExecDir Name: '$($VenvExecDir.Name)" - -# Set values required in priority: CmdLine, ConfigFile, Default -# First, get the location of the virtual environment, it might not be -# VenvExecDir if specified on the command line. -if ($VenvDir) { - Write-Verbose "VenvDir given as parameter, using '$VenvDir' to determine values" -} -else { - Write-Verbose "VenvDir not given as a parameter, using parent directory name as VenvDir." - $VenvDir = $VenvExecDir.Parent.FullName.TrimEnd("\\/") - Write-Verbose "VenvDir=$VenvDir" -} - -# Next, read the `pyvenv.cfg` file to determine any required value such -# as `prompt`. -$pyvenvCfg = Get-PyVenvConfig -ConfigDir $VenvDir - -# Next, set the prompt from the command line, or the config file, or -# just use the name of the virtual environment folder. -if ($Prompt) { - Write-Verbose "Prompt specified as argument, using '$Prompt'" -} -else { - Write-Verbose "Prompt not specified as argument to script, checking pyvenv.cfg value" - if ($pyvenvCfg -and $pyvenvCfg['prompt']) { - Write-Verbose " Setting based on value in pyvenv.cfg='$($pyvenvCfg['prompt'])'" - $Prompt = $pyvenvCfg['prompt']; - } - else { - Write-Verbose " Setting prompt based on parent's directory's name. (Is the directory name passed to venv module when creating the virtual environment)" - Write-Verbose " Got leaf-name of $VenvDir='$(Split-Path -Path $venvDir -Leaf)'" - $Prompt = Split-Path -Path $venvDir -Leaf - } -} - -Write-Verbose "Prompt = '$Prompt'" -Write-Verbose "VenvDir='$VenvDir'" - -# Deactivate any currently active virtual environment, but leave the -# deactivate function in place. -deactivate -nondestructive - -# Now set the environment variable VIRTUAL_ENV, used by many tools to determine -# that there is an activated venv. -$env:VIRTUAL_ENV = $VenvDir - -if (-not $Env:VIRTUAL_ENV_DISABLE_PROMPT) { - - Write-Verbose "Setting prompt to '$Prompt'" - - # Set the prompt to include the env name - # Make sure _OLD_VIRTUAL_PROMPT is global - function global:_OLD_VIRTUAL_PROMPT { "" } - Copy-Item -Path function:prompt -Destination function:_OLD_VIRTUAL_PROMPT - New-Variable -Name _PYTHON_VENV_PROMPT_PREFIX -Description "Python virtual environment prompt prefix" -Scope Global -Option ReadOnly -Visibility Public -Value $Prompt - - function global:prompt { - Write-Host -NoNewline -ForegroundColor Green "($_PYTHON_VENV_PROMPT_PREFIX) " - _OLD_VIRTUAL_PROMPT - } - $env:VIRTUAL_ENV_PROMPT = $Prompt -} - -# Clear PYTHONHOME -if (Test-Path -Path Env:PYTHONHOME) { - Copy-Item -Path Env:PYTHONHOME -Destination Env:_OLD_VIRTUAL_PYTHONHOME - Remove-Item -Path Env:PYTHONHOME -} - -# Add the venv to the PATH -Copy-Item -Path Env:PATH -Destination Env:_OLD_VIRTUAL_PATH -$Env:PATH = "$VenvExecDir$([System.IO.Path]::PathSeparator)$Env:PATH" diff --git a/venv/bin/accelerate b/venv/bin/accelerate deleted file mode 100644 index 399650c4d72f22bbeec63ece11ac898e2d23570a..0000000000000000000000000000000000000000 --- a/venv/bin/accelerate +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from accelerate.commands.accelerate_cli import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/accelerate-config b/venv/bin/accelerate-config deleted file mode 100644 index 4a61a2288d65ac8ba5f4bb04e294db019a13330b..0000000000000000000000000000000000000000 --- a/venv/bin/accelerate-config +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from accelerate.commands.config import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/accelerate-estimate-memory b/venv/bin/accelerate-estimate-memory deleted file mode 100644 index dd9069a99c68aa7cc6440f780909b55cf29e98c7..0000000000000000000000000000000000000000 --- a/venv/bin/accelerate-estimate-memory +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from accelerate.commands.estimate import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/accelerate-launch b/venv/bin/accelerate-launch deleted file mode 100644 index 53be5d2dbf6ea8acbb08dea05ac3bc52f4823dc5..0000000000000000000000000000000000000000 --- a/venv/bin/accelerate-launch +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from accelerate.commands.launch import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/accelerate-merge-weights b/venv/bin/accelerate-merge-weights deleted file mode 100644 index d70256f9a5e6f568e027a40ea3d90c93e8f06f0b..0000000000000000000000000000000000000000 --- a/venv/bin/accelerate-merge-weights +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from accelerate.commands.merge import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/activate b/venv/bin/activate deleted file mode 100644 index 6856b99d193a585ff1c0507892b7ec456d23d309..0000000000000000000000000000000000000000 --- a/venv/bin/activate +++ /dev/null @@ -1,69 +0,0 @@ -# This file must be used with "source bin/activate" *from bash* -# you cannot run it directly - -deactivate () { - # reset old environment variables - if [ -n "${_OLD_VIRTUAL_PATH:-}" ] ; then - PATH="${_OLD_VIRTUAL_PATH:-}" - export PATH - unset _OLD_VIRTUAL_PATH - fi - if [ -n "${_OLD_VIRTUAL_PYTHONHOME:-}" ] ; then - PYTHONHOME="${_OLD_VIRTUAL_PYTHONHOME:-}" - export PYTHONHOME - unset _OLD_VIRTUAL_PYTHONHOME - fi - - # This should detect bash and zsh, which have a hash command that must - # be called to get it to forget past commands. Without forgetting - # past commands the $PATH changes we made may not be respected - if [ -n "${BASH:-}" -o -n "${ZSH_VERSION:-}" ] ; then - hash -r 2> /dev/null - fi - - if [ -n "${_OLD_VIRTUAL_PS1:-}" ] ; then - PS1="${_OLD_VIRTUAL_PS1:-}" - export PS1 - unset _OLD_VIRTUAL_PS1 - fi - - unset VIRTUAL_ENV - unset VIRTUAL_ENV_PROMPT - if [ ! "${1:-}" = "nondestructive" ] ; then - # Self destruct! - unset -f deactivate - fi -} - -# unset irrelevant variables -deactivate nondestructive - -VIRTUAL_ENV="/workspace/ComfyUI/venv" -export VIRTUAL_ENV - -_OLD_VIRTUAL_PATH="$PATH" -PATH="$VIRTUAL_ENV/bin:$PATH" -export PATH - -# unset PYTHONHOME if set -# this will fail if PYTHONHOME is set to the empty string (which is bad anyway) -# could use `if (set -u; : $PYTHONHOME) ;` in bash -if [ -n "${PYTHONHOME:-}" ] ; then - _OLD_VIRTUAL_PYTHONHOME="${PYTHONHOME:-}" - unset PYTHONHOME -fi - -if [ -z "${VIRTUAL_ENV_DISABLE_PROMPT:-}" ] ; then - _OLD_VIRTUAL_PS1="${PS1:-}" - PS1="(venv) ${PS1:-}" - export PS1 - VIRTUAL_ENV_PROMPT="(venv) " - export VIRTUAL_ENV_PROMPT -fi - -# This should detect bash and zsh, which have a hash command that must -# be called to get it to forget past commands. Without forgetting -# past commands the $PATH changes we made may not be respected -if [ -n "${BASH:-}" -o -n "${ZSH_VERSION:-}" ] ; then - hash -r 2> /dev/null -fi diff --git a/venv/bin/activate.csh b/venv/bin/activate.csh deleted file mode 100644 index 2f00d786397ac022311d9029b4d083d28b5ae750..0000000000000000000000000000000000000000 --- a/venv/bin/activate.csh +++ /dev/null @@ -1,26 +0,0 @@ -# This file must be used with "source bin/activate.csh" *from csh*. -# You cannot run it directly. -# Created by Davide Di Blasi . -# Ported to Python 3.3 venv by Andrew Svetlov - -alias deactivate 'test $?_OLD_VIRTUAL_PATH != 0 && setenv PATH "$_OLD_VIRTUAL_PATH" && unset _OLD_VIRTUAL_PATH; rehash; test $?_OLD_VIRTUAL_PROMPT != 0 && set prompt="$_OLD_VIRTUAL_PROMPT" && unset _OLD_VIRTUAL_PROMPT; unsetenv VIRTUAL_ENV; unsetenv VIRTUAL_ENV_PROMPT; test "\!:*" != "nondestructive" && unalias deactivate' - -# Unset irrelevant variables. -deactivate nondestructive - -setenv VIRTUAL_ENV "/workspace/ComfyUI/venv" - -set _OLD_VIRTUAL_PATH="$PATH" -setenv PATH "$VIRTUAL_ENV/bin:$PATH" - - -set _OLD_VIRTUAL_PROMPT="$prompt" - -if (! "$?VIRTUAL_ENV_DISABLE_PROMPT") then - set prompt = "(venv) $prompt" - setenv VIRTUAL_ENV_PROMPT "(venv) " -endif - -alias pydoc python -m pydoc - -rehash diff --git a/venv/bin/activate.fish b/venv/bin/activate.fish deleted file mode 100644 index 41f0e4a4b084251bcc6c3f6282959932ad982de6..0000000000000000000000000000000000000000 --- a/venv/bin/activate.fish +++ /dev/null @@ -1,69 +0,0 @@ -# This file must be used with "source /bin/activate.fish" *from fish* -# (https://fishshell.com/); you cannot run it directly. - -function deactivate -d "Exit virtual environment and return to normal shell environment" - # reset old environment variables - if test -n "$_OLD_VIRTUAL_PATH" - set -gx PATH $_OLD_VIRTUAL_PATH - set -e _OLD_VIRTUAL_PATH - end - if test -n "$_OLD_VIRTUAL_PYTHONHOME" - set -gx PYTHONHOME $_OLD_VIRTUAL_PYTHONHOME - set -e _OLD_VIRTUAL_PYTHONHOME - end - - if test -n "$_OLD_FISH_PROMPT_OVERRIDE" - set -e _OLD_FISH_PROMPT_OVERRIDE - # prevents error when using nested fish instances (Issue #93858) - if functions -q _old_fish_prompt - functions -e fish_prompt - functions -c _old_fish_prompt fish_prompt - functions -e _old_fish_prompt - end - end - - set -e VIRTUAL_ENV - set -e VIRTUAL_ENV_PROMPT - if test "$argv[1]" != "nondestructive" - # Self-destruct! - functions -e deactivate - end -end - -# Unset irrelevant variables. -deactivate nondestructive - -set -gx VIRTUAL_ENV "/workspace/ComfyUI/venv" - -set -gx _OLD_VIRTUAL_PATH $PATH -set -gx PATH "$VIRTUAL_ENV/bin" $PATH - -# Unset PYTHONHOME if set. -if set -q PYTHONHOME - set -gx _OLD_VIRTUAL_PYTHONHOME $PYTHONHOME - set -e PYTHONHOME -end - -if test -z "$VIRTUAL_ENV_DISABLE_PROMPT" - # fish uses a function instead of an env var to generate the prompt. - - # Save the current fish_prompt function as the function _old_fish_prompt. - functions -c fish_prompt _old_fish_prompt - - # With the original prompt function renamed, we can override with our own. - function fish_prompt - # Save the return status of the last command. - set -l old_status $status - - # Output the venv prompt; color taken from the blue of the Python logo. - printf "%s%s%s" (set_color 4B8BBE) "(venv) " (set_color normal) - - # Restore the return status of the previous command. - echo "exit $old_status" | . - # Output the original/"old" prompt. - _old_fish_prompt - end - - set -gx _OLD_FISH_PROMPT_OVERRIDE "$VIRTUAL_ENV" - set -gx VIRTUAL_ENV_PROMPT "(venv) " -end diff --git a/venv/bin/alembic b/venv/bin/alembic deleted file mode 100644 index d39b3ce1e34df0705469f81d5141e9f4499856d2..0000000000000000000000000000000000000000 --- a/venv/bin/alembic +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from alembic.config import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/backend-test-tools b/venv/bin/backend-test-tools deleted file mode 100644 index 6fe579f6b1aec62ec9bc497e30c0e47362c86b06..0000000000000000000000000000000000000000 --- a/venv/bin/backend-test-tools +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from onnx.backend.test.cmd_tools import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/chardetect b/venv/bin/chardetect deleted file mode 100644 index 41d7624b9886b55e7c8f12261ec2b34190d3228b..0000000000000000000000000000000000000000 --- a/venv/bin/chardetect +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from chardet.cli.chardetect import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/check-model b/venv/bin/check-model deleted file mode 100644 index cfd70738ddbcb3b3bd621ddb60611ccb828dd7c5..0000000000000000000000000000000000000000 --- a/venv/bin/check-model +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from onnx.bin.checker import check_model -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(check_model()) diff --git a/venv/bin/check-node b/venv/bin/check-node deleted file mode 100644 index 569da1ff51b596cda1ecd53bfa65cfb4fa311900..0000000000000000000000000000000000000000 --- a/venv/bin/check-node +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from onnx.bin.checker import check_node -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(check_node()) diff --git a/venv/bin/coloredlogs b/venv/bin/coloredlogs deleted file mode 100644 index 2b7c4d7d00bd06bc2cec1b7b533a4a0ddfa6cbe9..0000000000000000000000000000000000000000 --- a/venv/bin/coloredlogs +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from coloredlogs.cli import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/cpuinfo b/venv/bin/cpuinfo deleted file mode 100644 index 5d8a1924ba40b93a89c16dc158cd3b1ac7b16203..0000000000000000000000000000000000000000 --- a/venv/bin/cpuinfo +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from cpuinfo import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/cygdb b/venv/bin/cygdb deleted file mode 100644 index 3232b8e913a6438c6623ba2277197b772dc6065b..0000000000000000000000000000000000000000 --- a/venv/bin/cygdb +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from Cython.Debugger.Cygdb import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/cython b/venv/bin/cython deleted file mode 100644 index 77cf9f3637fdc2348104cea3d65f714c016a969b..0000000000000000000000000000000000000000 --- a/venv/bin/cython +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from Cython.Compiler.Main import setuptools_main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(setuptools_main()) diff --git a/venv/bin/cythonize b/venv/bin/cythonize deleted file mode 100644 index 94d2ace221cc907da2755d3489153391ece9e9a3..0000000000000000000000000000000000000000 --- a/venv/bin/cythonize +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from Cython.Build.Cythonize import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/deepspeed b/venv/bin/deepspeed deleted file mode 100644 index 060e07dc1945b15a10c8a81088ec426231c8d302..0000000000000000000000000000000000000000 --- a/venv/bin/deepspeed +++ /dev/null @@ -1,6 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python - -from deepspeed.launcher.runner import main - -if __name__ == '__main__': - main() diff --git a/venv/bin/deepspeed.pt b/venv/bin/deepspeed.pt deleted file mode 100644 index c604309caa522e7b06d3cca72396a1044f9bae6f..0000000000000000000000000000000000000000 --- a/venv/bin/deepspeed.pt +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7d5bc66cee1ff65f80ad49f921cfb3746fe89f3edc000891eb3a7820c8703f4c -size 120 diff --git a/venv/bin/diffusers-cli b/venv/bin/diffusers-cli deleted file mode 100644 index a656223522f4a25fc5ad61ec27549f3f72c7ac43..0000000000000000000000000000000000000000 --- a/venv/bin/diffusers-cli +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from diffusers.commands.diffusers_cli import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/dotenv b/venv/bin/dotenv deleted file mode 100644 index 53ae8439714853773202b5e5f8751236c94c7504..0000000000000000000000000000000000000000 --- a/venv/bin/dotenv +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from dotenv.__main__ import cli -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(cli()) diff --git a/venv/bin/ds b/venv/bin/ds deleted file mode 100644 index 060e07dc1945b15a10c8a81088ec426231c8d302..0000000000000000000000000000000000000000 --- a/venv/bin/ds +++ /dev/null @@ -1,6 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python - -from deepspeed.launcher.runner import main - -if __name__ == '__main__': - main() diff --git a/venv/bin/ds_bench b/venv/bin/ds_bench deleted file mode 100644 index e48ba39a050defe747c9b34516dc5f4c52fa77a8..0000000000000000000000000000000000000000 --- a/venv/bin/ds_bench +++ /dev/null @@ -1,20 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python - -from benchmarks.communication.run_all import main -from benchmarks.communication.constants import * -from benchmarks.communication.utils import * -import os -import sys - -# Run the same file with deepspeed launcher. This is required since setuptools will auto-detect python files and insert a python shebang for both 'scripts' and 'entry_points', and this benchmarks require the DS launcher -required_env = ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"] -if not all(map(lambda v: v in os.environ, required_env)): - import subprocess - r = subprocess.check_output(["which", "ds_bench"]) - ds_bench_bin = r.decode('utf-8').strip() - safe_cmd = ["deepspeed", ds_bench_bin] + sys.argv[1:] - subprocess.run(safe_cmd) -else: - args = benchmark_parser().parse_args() - rank = args.local_rank - main(args, rank) diff --git a/venv/bin/ds_elastic b/venv/bin/ds_elastic deleted file mode 100644 index f054e3f6957459fee44b5d56f729fbf1dad4efc0..0000000000000000000000000000000000000000 --- a/venv/bin/ds_elastic +++ /dev/null @@ -1,40 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python - -import argparse -import json - -import deepspeed -from deepspeed.elasticity import compute_elastic_config - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('-c', '--config', type=str, help="DeepSpeed config json") - parser.add_argument('-w', '--world-size', type=int, default=0, help="Intended/current world size") - args = parser.parse_args() - ds_config = json.load(open(args.config, 'r')) - - ds_version = deepspeed.__version__ - - elastic_config = ds_config['elasticity'] - print('------------------------------------------') - print("Elasticity config:") - print('------------------------------------------') - print(json.dumps(elastic_config, indent=4, sort_keys=True)) - - if args.world_size > 0: - final_batch_size, valid_gpus, micro_batch_size = compute_elastic_config(ds_config=ds_config, - target_deepspeed_version=ds_version, - world_size=args.world_size) - print('------------------------------------------') - print(f"Calculated results for world size {args.world_size}:") - print('------------------------------------------') - print(f'final_batch_size .... {final_batch_size}') - print(f'valid_gpus .......... {valid_gpus}') - print(f'micro_batch_size .... {micro_batch_size}') - else: - final_batch_size, valid_gpus = compute_elastic_config(ds_config=ds_config, target_deepspeed_version=ds_version) - print('------------------------------------------') - print("Calculated results:") - print('------------------------------------------') - print(f'final_batch_size .... {final_batch_size}') - print(f'valid_gpus .......... {valid_gpus}') diff --git a/venv/bin/ds_io b/venv/bin/ds_io deleted file mode 100644 index fbb3ac37cef90b3d2d41c3bf803ed8a8b2c37def..0000000000000000000000000000000000000000 --- a/venv/bin/ds_io +++ /dev/null @@ -1,6 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python - -from deepspeed.nvme import ds_io_main - -if __name__ == '__main__': - ds_io_main() diff --git a/venv/bin/ds_nvme_tune b/venv/bin/ds_nvme_tune deleted file mode 100644 index 981b25819374ba078e07f0e00123d42cae7cd080..0000000000000000000000000000000000000000 --- a/venv/bin/ds_nvme_tune +++ /dev/null @@ -1,9 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python - -from deepspeed.nvme import sweep_main, generate_main, parse_sweep_arguments - -if __name__ == '__main__': - args = parse_sweep_arguments() - print(f"Running DeepNVMe performance tuning on {args.nvme_dir}") - sweep_main(args) - generate_main(args.log_dir) diff --git a/venv/bin/ds_report b/venv/bin/ds_report deleted file mode 100644 index 8d1ad8712249376ea5819fc36cf414bcdfced5eb..0000000000000000000000000000000000000000 --- a/venv/bin/ds_report +++ /dev/null @@ -1,6 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python - -from deepspeed.env_report import cli_main - -if __name__ == '__main__': - cli_main() diff --git a/venv/bin/ds_ssh b/venv/bin/ds_ssh deleted file mode 100644 index d89fc0b44e176c8acf422302308e6a865350ab49..0000000000000000000000000000000000000000 --- a/venv/bin/ds_ssh +++ /dev/null @@ -1,34 +0,0 @@ -#!/bin/bash - -# Copyright 2020 The Microsoft DeepSpeed Team - -command -v pdsh -if [ $? != 0 ]; then - echo "Cannot find pdsh, please install via 'apt-get install -y pdsh'" - exit 1 -fi - -hostfile=/job/hostfile - -while getopts "h?f:" opt; do - case "$opt" in - h|\?) - echo "-f : specify a hostfile, defaults to /job/hostfile" - exit 0 - ;; - f) - hostfile=$OPTARG - shift $((OPTIND-1)) - ;; - esac -done - -echo "hostfile=$hostfile" - -if [ -f $hostfile ]; then - hosts=`cat $hostfile | awk '{print $1}' | paste -sd "," -` - export PDSH_RCMD_TYPE=ssh - pdsh -w ${hosts} $@ -else - echo "Missing hostfile at ${hostfile}, unable to proceed" -fi diff --git a/venv/bin/dsr b/venv/bin/dsr deleted file mode 100644 index 8d1ad8712249376ea5819fc36cf414bcdfced5eb..0000000000000000000000000000000000000000 --- a/venv/bin/dsr +++ /dev/null @@ -1,6 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python - -from deepspeed.env_report import cli_main - -if __name__ == '__main__': - cli_main() diff --git a/venv/bin/f2py b/venv/bin/f2py deleted file mode 100644 index 31400035c5c7ead0690fd63465c2a443aa808111..0000000000000000000000000000000000000000 --- a/venv/bin/f2py +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from numpy.f2py.f2py2e import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/fonttools b/venv/bin/fonttools deleted file mode 100644 index cc5cb6fd89391f03ea1b4efa29ca824b41988e24..0000000000000000000000000000000000000000 --- a/venv/bin/fonttools +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from fontTools.__main__ import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/gguf-convert-endian b/venv/bin/gguf-convert-endian deleted file mode 100644 index faa3fceb421ef9c31bc573b6cc03a82bc3ff246b..0000000000000000000000000000000000000000 --- a/venv/bin/gguf-convert-endian +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from gguf.scripts.gguf_convert_endian import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/gguf-dump b/venv/bin/gguf-dump deleted file mode 100644 index 62c2139ecf38a02fcb85bae594170ca7869933fd..0000000000000000000000000000000000000000 --- a/venv/bin/gguf-dump +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from gguf.scripts.gguf_dump import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/gguf-editor-gui b/venv/bin/gguf-editor-gui deleted file mode 100644 index 4995aaa7ebec09bcf5289773666498f6938f6b7a..0000000000000000000000000000000000000000 --- a/venv/bin/gguf-editor-gui +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from gguf.scripts.gguf_editor_gui import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/gguf-new-metadata b/venv/bin/gguf-new-metadata deleted file mode 100644 index 86a516d87f72446fb410a49ef584ccd59437d0bf..0000000000000000000000000000000000000000 --- a/venv/bin/gguf-new-metadata +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from gguf.scripts.gguf_new_metadata import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/gguf-set-metadata b/venv/bin/gguf-set-metadata deleted file mode 100644 index 1cbe7b25f16f2160209d909626da7688ca32dab3..0000000000000000000000000000000000000000 --- a/venv/bin/gguf-set-metadata +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from gguf.scripts.gguf_set_metadata import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/hjson b/venv/bin/hjson deleted file mode 100644 index 6689d5f3f174cafcd902971b01e6486a353a5cbf..0000000000000000000000000000000000000000 --- a/venv/bin/hjson +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from hjson.tool import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/huggingface-cli b/venv/bin/huggingface-cli deleted file mode 100644 index 1581da59b63cb2e5d809c23b2a7836df0f6c8fd4..0000000000000000000000000000000000000000 --- a/venv/bin/huggingface-cli +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from huggingface_hub.commands.huggingface_cli import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/humanfriendly b/venv/bin/humanfriendly deleted file mode 100644 index 798c37a08524dfbb699c0b088a8699b3da421fe9..0000000000000000000000000000000000000000 --- a/venv/bin/humanfriendly +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from humanfriendly.cli import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/imageio_download_bin b/venv/bin/imageio_download_bin deleted file mode 100644 index ebd288e56d5d30333265bed632a68675ef044a95..0000000000000000000000000000000000000000 --- a/venv/bin/imageio_download_bin +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from imageio.__main__ import download_bin_main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(download_bin_main()) diff --git a/venv/bin/imageio_remove_bin b/venv/bin/imageio_remove_bin deleted file mode 100644 index 50854762e75c48ac0f317a863af3a41b4a9b3b69..0000000000000000000000000000000000000000 --- a/venv/bin/imageio_remove_bin +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from imageio.__main__ import remove_bin_main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(remove_bin_main()) diff --git a/venv/bin/insightface-cli b/venv/bin/insightface-cli deleted file mode 100644 index a2675a3610fb80bb423ed3e3148100a83a98f056..0000000000000000000000000000000000000000 --- a/venv/bin/insightface-cli +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from insightface.commands.insightface_cli import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/isympy b/venv/bin/isympy deleted file mode 100644 index 85b79269eb94e5488f5e14774ca299f0070b12f0..0000000000000000000000000000000000000000 --- a/venv/bin/isympy +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from isympy import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/lsm2bin b/venv/bin/lsm2bin deleted file mode 100644 index dffbc47f548ca9732747e2812d81af0c1df19dd9..0000000000000000000000000000000000000000 --- a/venv/bin/lsm2bin +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from tifffile.lsm2bin import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/mako-render b/venv/bin/mako-render deleted file mode 100644 index 8a6677167b9e7cb529ec6c742de06c63006dd23c..0000000000000000000000000000000000000000 --- a/venv/bin/mako-render +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from mako.cmd import cmdline -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(cmdline()) diff --git a/venv/bin/markdown-it b/venv/bin/markdown-it deleted file mode 100644 index 8a48cd5d1082066b32e48f1b30715e713a062fc2..0000000000000000000000000000000000000000 --- a/venv/bin/markdown-it +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from markdown_it.cli.parse import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/ninja b/venv/bin/ninja deleted file mode 100644 index 3dcf5604376f64e650c14891de82a719b59fcee5..0000000000000000000000000000000000000000 --- a/venv/bin/ninja +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:94f5318756260447ffb88d0bc31b67dbc07119f1dc0c8dc2cb58ea575a471025 -size 1085448 diff --git a/venv/bin/normalizer b/venv/bin/normalizer deleted file mode 100644 index d7b2663ae5eb0b88ff673c386cdf07a58750aa73..0000000000000000000000000000000000000000 --- a/venv/bin/normalizer +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from charset_normalizer import cli -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(cli.cli_detect()) diff --git a/venv/bin/numpy-config b/venv/bin/numpy-config deleted file mode 100644 index 41f4dd0ed5df5abb20b9c0e0b437a870a4fa9eff..0000000000000000000000000000000000000000 --- a/venv/bin/numpy-config +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from numpy._configtool import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/onnxruntime_test b/venv/bin/onnxruntime_test deleted file mode 100644 index bf851f57e30c5bd39848992fec3cbf34eb0b939b..0000000000000000000000000000000000000000 --- a/venv/bin/onnxruntime_test +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from onnxruntime.tools.onnxruntime_test import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/pip b/venv/bin/pip deleted file mode 100644 index 5691806511fcc6c5ee67abd2a4a7f13370e9eb2a..0000000000000000000000000000000000000000 --- a/venv/bin/pip +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from pip._internal.cli.main import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/pip3 b/venv/bin/pip3 deleted file mode 100644 index 5691806511fcc6c5ee67abd2a4a7f13370e9eb2a..0000000000000000000000000000000000000000 --- a/venv/bin/pip3 +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from pip._internal.cli.main import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/pip3.10 b/venv/bin/pip3.10 deleted file mode 100644 index 5691806511fcc6c5ee67abd2a4a7f13370e9eb2a..0000000000000000000000000000000000000000 --- a/venv/bin/pip3.10 +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from pip._internal.cli.main import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/proton b/venv/bin/proton deleted file mode 100644 index 399ebd57ec5148368ea536745058e2a7bccee81c..0000000000000000000000000000000000000000 --- a/venv/bin/proton +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from triton.profiler.proton import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/proton-viewer b/venv/bin/proton-viewer deleted file mode 100644 index 811aebe8f7e48f8e93668f0940d2546b896b1a86..0000000000000000000000000000000000000000 --- a/venv/bin/proton-viewer +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from triton.profiler.viewer import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/pyav b/venv/bin/pyav deleted file mode 100644 index befd971c6b7a97282cd9789187dadc52dbd43fe9..0000000000000000000000000000000000000000 --- a/venv/bin/pyav +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from av.__main__ import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/pyftmerge b/venv/bin/pyftmerge deleted file mode 100644 index 68d0d8b18ec7b9212cf33f87ba028f42653fc2f8..0000000000000000000000000000000000000000 --- a/venv/bin/pyftmerge +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from fontTools.merge import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/pyftsubset b/venv/bin/pyftsubset deleted file mode 100644 index 61cb0850c7c0f9aaeabb935a27b813ec38c7f1ac..0000000000000000000000000000000000000000 --- a/venv/bin/pyftsubset +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from fontTools.subset import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/pygmentize b/venv/bin/pygmentize deleted file mode 100644 index b819405e74fd476960353223dd1278ab778041f4..0000000000000000000000000000000000000000 --- a/venv/bin/pygmentize +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from pygments.cmdline import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/python b/venv/bin/python deleted file mode 100644 index 64e8728adfafdb3e95b983fb0960f9a09cd3bac9..0000000000000000000000000000000000000000 --- a/venv/bin/python +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:45692c3da2492563eabf0a8f5dc18d20dc9c34ffe3a18202563e00bae684be91 -size 5904904 diff --git a/venv/bin/python3 b/venv/bin/python3 deleted file mode 100644 index 64e8728adfafdb3e95b983fb0960f9a09cd3bac9..0000000000000000000000000000000000000000 --- a/venv/bin/python3 +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:45692c3da2492563eabf0a8f5dc18d20dc9c34ffe3a18202563e00bae684be91 -size 5904904 diff --git a/venv/bin/python3.10 b/venv/bin/python3.10 deleted file mode 100644 index 64e8728adfafdb3e95b983fb0960f9a09cd3bac9..0000000000000000000000000000000000000000 --- a/venv/bin/python3.10 +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:45692c3da2492563eabf0a8f5dc18d20dc9c34ffe3a18202563e00bae684be91 -size 5904904 diff --git a/venv/bin/sz_split b/venv/bin/sz_split deleted file mode 100644 index 2ee765ec44d947dc0d529113034b2d89a041833b..0000000000000000000000000000000000000000 --- a/venv/bin/sz_split +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from cli.split import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/sz_wc b/venv/bin/sz_wc deleted file mode 100644 index a515643d44096c581095ebdf7cda9403be666cca..0000000000000000000000000000000000000000 --- a/venv/bin/sz_wc +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from cli.wc import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/tiff2fsspec b/venv/bin/tiff2fsspec deleted file mode 100644 index 5d6da0bfd7b99746c0be7a7f7dc1dadc896fed2c..0000000000000000000000000000000000000000 --- a/venv/bin/tiff2fsspec +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from tifffile.tiff2fsspec import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/tiffcomment b/venv/bin/tiffcomment deleted file mode 100644 index 22dc754cc5727e4a049b7ae0e833717aa7c68d5b..0000000000000000000000000000000000000000 --- a/venv/bin/tiffcomment +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from tifffile.tiffcomment import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/tifffile b/venv/bin/tifffile deleted file mode 100644 index 1f519b980c5cb0cac9acddc70bf5c9317ce4639d..0000000000000000000000000000000000000000 --- a/venv/bin/tifffile +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from tifffile import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/tiny-agents b/venv/bin/tiny-agents deleted file mode 100644 index ca7cc3073e175146bccb02c80059b4ba03309dbf..0000000000000000000000000000000000000000 --- a/venv/bin/tiny-agents +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from huggingface_hub.inference._mcp.cli import app -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(app()) diff --git a/venv/bin/torchfrtrace b/venv/bin/torchfrtrace deleted file mode 100644 index 9d3c6afe47e1d9fdd80a6b1ad7943a15a8470af0..0000000000000000000000000000000000000000 --- a/venv/bin/torchfrtrace +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from tools.flight_recorder.fr_trace import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/torchrun b/venv/bin/torchrun deleted file mode 100644 index 328d24da5399bc29be842766f2e8fd7003f6e70b..0000000000000000000000000000000000000000 --- a/venv/bin/torchrun +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from torch.distributed.run import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/tqdm b/venv/bin/tqdm deleted file mode 100644 index 7b37f4a3be008ef4021dde868be47b535edc3f37..0000000000000000000000000000000000000000 --- a/venv/bin/tqdm +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from tqdm.cli import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/transformers b/venv/bin/transformers deleted file mode 100644 index eaf46c2b222f9781457bc7420706421fc59a813f..0000000000000000000000000000000000000000 --- a/venv/bin/transformers +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from transformers.commands.transformers_cli import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/transformers-cli b/venv/bin/transformers-cli deleted file mode 100644 index fdb2762a29618580f6af0661d285369644c37ce6..0000000000000000000000000000000000000000 --- a/venv/bin/transformers-cli +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from transformers.commands.transformers_cli import main_cli -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main_cli()) diff --git a/venv/bin/ttx b/venv/bin/ttx deleted file mode 100644 index db54df50bdb741c1152edd5c3e05c3650146dad2..0000000000000000000000000000000000000000 --- a/venv/bin/ttx +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from fontTools.ttx import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/typer b/venv/bin/typer deleted file mode 100644 index 4cc465c48fefc82dd20682a1d862d344ce45692d..0000000000000000000000000000000000000000 --- a/venv/bin/typer +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from typer.cli import main -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(main()) diff --git a/venv/bin/ultralytics b/venv/bin/ultralytics deleted file mode 100644 index 00dc536eff214a1eac72fd6cb685d722a53189f7..0000000000000000000000000000000000000000 --- a/venv/bin/ultralytics +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from ultralytics.cfg import entrypoint -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(entrypoint()) diff --git a/venv/bin/uv b/venv/bin/uv deleted file mode 100644 index 4ec0125893f3fb4642af33b0fa1d83e1404c4113..0000000000000000000000000000000000000000 --- a/venv/bin/uv +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:19af10f031be0587a0f8352375e15ea46740df9346843d834333a3e491899bd0 -size 45929808 diff --git a/venv/bin/uvx b/venv/bin/uvx deleted file mode 100644 index a5fc1c3d8f86feac17484e4738f39ac5f384599c..0000000000000000000000000000000000000000 --- a/venv/bin/uvx +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a119bc8313764e6cba793788e1c11b84e8c32e59c20dcd4d57aff71b07512694 -size 361008 diff --git a/venv/bin/yolo b/venv/bin/yolo deleted file mode 100644 index 00dc536eff214a1eac72fd6cb685d722a53189f7..0000000000000000000000000000000000000000 --- a/venv/bin/yolo +++ /dev/null @@ -1,8 +0,0 @@ -#!/workspace/ComfyUI/venv/bin/python -# -*- coding: utf-8 -*- -import re -import sys -from ultralytics.cfg import entrypoint -if __name__ == '__main__': - sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) - sys.exit(entrypoint()) diff --git a/venv/insightface/data/images/Tom_Hanks_54745.png b/venv/insightface/data/images/Tom_Hanks_54745.png deleted file mode 100644 index 906315d13fa29bb3a5ded3e162592f2c7f041b23..0000000000000000000000000000000000000000 Binary files a/venv/insightface/data/images/Tom_Hanks_54745.png and /dev/null differ diff --git a/venv/insightface/data/images/mask_black.jpg b/venv/insightface/data/images/mask_black.jpg deleted file mode 100644 index 0eab0df555c23f1e033537fe39f3c0c8303dd369..0000000000000000000000000000000000000000 Binary files a/venv/insightface/data/images/mask_black.jpg and /dev/null differ diff --git a/venv/insightface/data/images/mask_blue.jpg b/venv/insightface/data/images/mask_blue.jpg deleted file mode 100644 index f71336b9a0d3038ebd84e6995ebfbe54946fcbb4..0000000000000000000000000000000000000000 Binary files a/venv/insightface/data/images/mask_blue.jpg and /dev/null differ diff --git a/venv/insightface/data/images/mask_green.jpg b/venv/insightface/data/images/mask_green.jpg deleted file mode 100644 index ac2ad55f4fc580c915dfa4c157ca3bfc84e453f4..0000000000000000000000000000000000000000 Binary files a/venv/insightface/data/images/mask_green.jpg and /dev/null differ diff --git a/venv/insightface/data/images/mask_white.jpg b/venv/insightface/data/images/mask_white.jpg deleted file mode 100644 index 2148ab2d09fdee6e3f59315470e98ecfc54339e4..0000000000000000000000000000000000000000 Binary files a/venv/insightface/data/images/mask_white.jpg and /dev/null differ diff --git a/venv/insightface/data/images/t1.jpg b/venv/insightface/data/images/t1.jpg deleted file mode 100644 index 8fd6427a177bd01650c0150e9d02457c3a5dcddd..0000000000000000000000000000000000000000 --- a/venv/insightface/data/images/t1.jpg +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:47f682e945b659f93a9e490b9c9c4a2a864abe64dace9e1a2893845ddfd69489 -size 128824 diff --git a/venv/insightface/data/objects/meanshape_68.pkl b/venv/insightface/data/objects/meanshape_68.pkl deleted file mode 100644 index d5297e9e8ea5574298ddd287b058252e03aa18c1..0000000000000000000000000000000000000000 --- a/venv/insightface/data/objects/meanshape_68.pkl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:39ffecf84ba73f0d0d7e49380833ba88713c9fcdec51df4f7ac45a48b8f4cc51 -size 974 diff --git a/venv/insightface/thirdparty/face3d/mesh/cython/__pycache__/setup.cpython-310.pyc b/venv/insightface/thirdparty/face3d/mesh/cython/__pycache__/setup.cpython-310.pyc deleted file mode 100644 index 3dbcb79db2686564cd9e9233724c4619acbd2254..0000000000000000000000000000000000000000 Binary files a/venv/insightface/thirdparty/face3d/mesh/cython/__pycache__/setup.cpython-310.pyc and /dev/null differ diff --git a/venv/insightface/thirdparty/face3d/mesh/cython/mesh_core.h b/venv/insightface/thirdparty/face3d/mesh/cython/mesh_core.h deleted file mode 100644 index 3eb00290f36f8cbe27a17d88bca31fe718d77e2c..0000000000000000000000000000000000000000 --- a/venv/insightface/thirdparty/face3d/mesh/cython/mesh_core.h +++ /dev/null @@ -1,83 +0,0 @@ -#ifndef MESH_CORE_HPP_ -#define MESH_CORE_HPP_ - -#include -#include -#include -#include -#include -#include - -using namespace std; - -class point -{ - public: - float x; - float y; - - float dot(point p) - { - return this->x * p.x + this->y * p.y; - } - - point operator-(const point& p) - { - point np; - np.x = this->x - p.x; - np.y = this->y - p.y; - return np; - } - - point operator+(const point& p) - { - point np; - np.x = this->x + p.x; - np.y = this->y + p.y; - return np; - } - - point operator*(float s) - { - point np; - np.x = s * this->x; - np.y = s * this->y; - return np; - } -}; - - -bool isPointInTri(point p, point p0, point p1, point p2, int h, int w); -void get_point_weight(float* weight, point p, point p0, point p1, point p2); - -void _get_normal_core( - float* normal, float* tri_normal, int* triangles, - int ntri); - -void _rasterize_triangles_core( - float* vertices, int* triangles, - float* depth_buffer, int* triangle_buffer, float* barycentric_weight, - int nver, int ntri, - int h, int w); - -void _render_colors_core( - float* image, float* vertices, int* triangles, - float* colors, - float* depth_buffer, - int nver, int ntri, - int h, int w, int c); - -void _render_texture_core( - float* image, float* vertices, int* triangles, - float* texture, float* tex_coords, int* tex_triangles, - float* depth_buffer, - int nver, int tex_nver, int ntri, - int h, int w, int c, - int tex_h, int tex_w, int tex_c, - int mapping_type); - -void _write_obj_with_colors_texture(string filename, string mtl_name, - float* vertices, int* triangles, float* colors, float* uv_coords, - int nver, int ntri, int ntexver); - -#endif \ No newline at end of file diff --git a/venv/insightface/thirdparty/face3d/mesh/cython/mesh_core_cython.c b/venv/insightface/thirdparty/face3d/mesh/cython/mesh_core_cython.c deleted file mode 100644 index e113ac54ae03e4a79ea2a975ef33e67deb9d93d6..0000000000000000000000000000000000000000 --- a/venv/insightface/thirdparty/face3d/mesh/cython/mesh_core_cython.c +++ /dev/null @@ -1,9091 +0,0 @@ -/* Generated by Cython 0.28.2 */ - -/* BEGIN: Cython Metadata -{ - "distutils": { - "depends": [ - "insightface/thirdparty/face3d/mesh/cython/mesh_core.h" - ], - "include_dirs": [ - "insightface/thirdparty/face3d/mesh/cython" - ], - "name": "mesh_core_cython", - "sources": [ - "insightface/thirdparty/face3d/mesh/cython/mesh_core_cython.pyx" - ] - }, - "module_name": "mesh_core_cython" -} -END: Cython Metadata */ - -#define PY_SSIZE_T_CLEAN -#include "Python.h" -#ifndef Py_PYTHON_H - #error Python headers needed to compile C extensions, please install development version of Python. -#elif PY_VERSION_HEX < 0x02060000 || (0x03000000 <= PY_VERSION_HEX && PY_VERSION_HEX < 0x03030000) - #error Cython requires Python 2.6+ or Python 3.3+. -#else -#define CYTHON_ABI "0_28_2" -#define CYTHON_FUTURE_DIVISION 0 -#include -#ifndef offsetof - #define offsetof(type, member) ( (size_t) & ((type*)0) -> member ) -#endif -#if !defined(WIN32) && !defined(MS_WINDOWS) - #ifndef __stdcall - #define __stdcall - #endif - #ifndef __cdecl - #define __cdecl - #endif - #ifndef __fastcall - #define __fastcall - #endif -#endif -#ifndef DL_IMPORT - #define DL_IMPORT(t) t -#endif -#ifndef DL_EXPORT - #define DL_EXPORT(t) t -#endif -#define __PYX_COMMA , -#ifndef HAVE_LONG_LONG - #if PY_VERSION_HEX >= 0x02070000 - #define HAVE_LONG_LONG - #endif -#endif -#ifndef PY_LONG_LONG - #define PY_LONG_LONG LONG_LONG -#endif -#ifndef Py_HUGE_VAL - #define Py_HUGE_VAL HUGE_VAL -#endif -#ifdef PYPY_VERSION - #define CYTHON_COMPILING_IN_PYPY 1 - #define CYTHON_COMPILING_IN_PYSTON 0 - #define CYTHON_COMPILING_IN_CPYTHON 0 - #undef CYTHON_USE_TYPE_SLOTS - #define CYTHON_USE_TYPE_SLOTS 0 - #undef CYTHON_USE_PYTYPE_LOOKUP - #define CYTHON_USE_PYTYPE_LOOKUP 0 - #if PY_VERSION_HEX < 0x03050000 - #undef CYTHON_USE_ASYNC_SLOTS - #define CYTHON_USE_ASYNC_SLOTS 0 - #elif !defined(CYTHON_USE_ASYNC_SLOTS) - #define CYTHON_USE_ASYNC_SLOTS 1 - #endif - #undef CYTHON_USE_PYLIST_INTERNALS - #define CYTHON_USE_PYLIST_INTERNALS 0 - #undef CYTHON_USE_UNICODE_INTERNALS - #define CYTHON_USE_UNICODE_INTERNALS 0 - #undef CYTHON_USE_UNICODE_WRITER - #define CYTHON_USE_UNICODE_WRITER 0 - #undef CYTHON_USE_PYLONG_INTERNALS - #define CYTHON_USE_PYLONG_INTERNALS 0 - #undef CYTHON_AVOID_BORROWED_REFS - #define CYTHON_AVOID_BORROWED_REFS 1 - #undef CYTHON_ASSUME_SAFE_MACROS - #define CYTHON_ASSUME_SAFE_MACROS 0 - #undef CYTHON_UNPACK_METHODS - #define CYTHON_UNPACK_METHODS 0 - #undef CYTHON_FAST_THREAD_STATE - #define CYTHON_FAST_THREAD_STATE 0 - #undef CYTHON_FAST_PYCALL - #define CYTHON_FAST_PYCALL 0 - #undef CYTHON_PEP489_MULTI_PHASE_INIT - #define CYTHON_PEP489_MULTI_PHASE_INIT 0 - #undef CYTHON_USE_TP_FINALIZE - #define CYTHON_USE_TP_FINALIZE 0 -#elif defined(PYSTON_VERSION) - #define CYTHON_COMPILING_IN_PYPY 0 - #define CYTHON_COMPILING_IN_PYSTON 1 - #define CYTHON_COMPILING_IN_CPYTHON 0 - #ifndef CYTHON_USE_TYPE_SLOTS - #define CYTHON_USE_TYPE_SLOTS 1 - #endif - #undef CYTHON_USE_PYTYPE_LOOKUP - #define CYTHON_USE_PYTYPE_LOOKUP 0 - #undef CYTHON_USE_ASYNC_SLOTS - #define CYTHON_USE_ASYNC_SLOTS 0 - #undef CYTHON_USE_PYLIST_INTERNALS - #define CYTHON_USE_PYLIST_INTERNALS 0 - #ifndef CYTHON_USE_UNICODE_INTERNALS - #define CYTHON_USE_UNICODE_INTERNALS 1 - #endif - #undef CYTHON_USE_UNICODE_WRITER - #define CYTHON_USE_UNICODE_WRITER 0 - #undef CYTHON_USE_PYLONG_INTERNALS - #define CYTHON_USE_PYLONG_INTERNALS 0 - #ifndef CYTHON_AVOID_BORROWED_REFS - #define CYTHON_AVOID_BORROWED_REFS 0 - #endif - #ifndef CYTHON_ASSUME_SAFE_MACROS - #define CYTHON_ASSUME_SAFE_MACROS 1 - #endif - #ifndef CYTHON_UNPACK_METHODS - #define CYTHON_UNPACK_METHODS 1 - #endif - #undef CYTHON_FAST_THREAD_STATE - #define CYTHON_FAST_THREAD_STATE 0 - #undef CYTHON_FAST_PYCALL - #define CYTHON_FAST_PYCALL 0 - #undef CYTHON_PEP489_MULTI_PHASE_INIT - #define CYTHON_PEP489_MULTI_PHASE_INIT 0 - #undef CYTHON_USE_TP_FINALIZE - #define CYTHON_USE_TP_FINALIZE 0 -#else - #define CYTHON_COMPILING_IN_PYPY 0 - #define CYTHON_COMPILING_IN_PYSTON 0 - #define CYTHON_COMPILING_IN_CPYTHON 1 - #ifndef CYTHON_USE_TYPE_SLOTS - #define CYTHON_USE_TYPE_SLOTS 1 - #endif - #if PY_VERSION_HEX < 0x02070000 - #undef CYTHON_USE_PYTYPE_LOOKUP - #define CYTHON_USE_PYTYPE_LOOKUP 0 - #elif !defined(CYTHON_USE_PYTYPE_LOOKUP) - #define CYTHON_USE_PYTYPE_LOOKUP 1 - #endif - #if PY_MAJOR_VERSION < 3 - #undef CYTHON_USE_ASYNC_SLOTS - #define CYTHON_USE_ASYNC_SLOTS 0 - #elif !defined(CYTHON_USE_ASYNC_SLOTS) - #define CYTHON_USE_ASYNC_SLOTS 1 - #endif - #if PY_VERSION_HEX < 0x02070000 - #undef CYTHON_USE_PYLONG_INTERNALS - #define CYTHON_USE_PYLONG_INTERNALS 0 - #elif !defined(CYTHON_USE_PYLONG_INTERNALS) - #define CYTHON_USE_PYLONG_INTERNALS 1 - #endif - #ifndef CYTHON_USE_PYLIST_INTERNALS - #define CYTHON_USE_PYLIST_INTERNALS 1 - #endif - #ifndef CYTHON_USE_UNICODE_INTERNALS - #define CYTHON_USE_UNICODE_INTERNALS 1 - #endif - #if PY_VERSION_HEX < 0x030300F0 - #undef CYTHON_USE_UNICODE_WRITER - #define CYTHON_USE_UNICODE_WRITER 0 - #elif !defined(CYTHON_USE_UNICODE_WRITER) - #define CYTHON_USE_UNICODE_WRITER 1 - #endif - #ifndef CYTHON_AVOID_BORROWED_REFS - #define CYTHON_AVOID_BORROWED_REFS 0 - #endif - #ifndef CYTHON_ASSUME_SAFE_MACROS - #define CYTHON_ASSUME_SAFE_MACROS 1 - #endif - #ifndef CYTHON_UNPACK_METHODS - #define CYTHON_UNPACK_METHODS 1 - #endif - #ifndef CYTHON_FAST_THREAD_STATE - #define CYTHON_FAST_THREAD_STATE 1 - #endif - #ifndef CYTHON_FAST_PYCALL - #define CYTHON_FAST_PYCALL 1 - #endif - #ifndef CYTHON_PEP489_MULTI_PHASE_INIT - #define CYTHON_PEP489_MULTI_PHASE_INIT (0 && PY_VERSION_HEX >= 0x03050000) - #endif - #ifndef CYTHON_USE_TP_FINALIZE - #define CYTHON_USE_TP_FINALIZE (PY_VERSION_HEX >= 0x030400a1) - #endif -#endif -#if !defined(CYTHON_FAST_PYCCALL) -#define CYTHON_FAST_PYCCALL (CYTHON_FAST_PYCALL && PY_VERSION_HEX >= 0x030600B1) -#endif -#if CYTHON_USE_PYLONG_INTERNALS - #include "longintrepr.h" - #undef SHIFT - #undef BASE - #undef MASK -#endif -#ifndef __has_attribute - #define __has_attribute(x) 0 -#endif -#ifndef __has_cpp_attribute - #define __has_cpp_attribute(x) 0 -#endif -#ifndef CYTHON_RESTRICT - #if defined(__GNUC__) - #define CYTHON_RESTRICT __restrict__ - #elif defined(_MSC_VER) && _MSC_VER >= 1400 - #define CYTHON_RESTRICT __restrict - #elif defined (__STDC_VERSION__) && __STDC_VERSION__ >= 199901L - #define CYTHON_RESTRICT restrict - #else - #define CYTHON_RESTRICT - #endif -#endif -#ifndef CYTHON_UNUSED -# if defined(__GNUC__) -# if !(defined(__cplusplus)) || (__GNUC__ > 3 || (__GNUC__ == 3 && __GNUC_MINOR__ >= 4)) -# define CYTHON_UNUSED __attribute__ ((__unused__)) -# else -# define CYTHON_UNUSED -# endif -# elif defined(__ICC) || (defined(__INTEL_COMPILER) && !defined(_MSC_VER)) -# define CYTHON_UNUSED __attribute__ ((__unused__)) -# else -# define CYTHON_UNUSED -# endif -#endif -#ifndef CYTHON_MAYBE_UNUSED_VAR -# if defined(__cplusplus) - template void CYTHON_MAYBE_UNUSED_VAR( const T& ) { } -# else -# define CYTHON_MAYBE_UNUSED_VAR(x) (void)(x) -# endif -#endif -#ifndef CYTHON_NCP_UNUSED -# if CYTHON_COMPILING_IN_CPYTHON -# define CYTHON_NCP_UNUSED -# else -# define CYTHON_NCP_UNUSED CYTHON_UNUSED -# endif -#endif -#define __Pyx_void_to_None(void_result) ((void)(void_result), Py_INCREF(Py_None), Py_None) -#ifdef _MSC_VER - #ifndef _MSC_STDINT_H_ - #if _MSC_VER < 1300 - typedef unsigned char uint8_t; - typedef unsigned int uint32_t; - #else - typedef unsigned __int8 uint8_t; - typedef unsigned __int32 uint32_t; - #endif - #endif -#else - #include -#endif -#ifndef CYTHON_FALLTHROUGH - #if defined(__cplusplus) && __cplusplus >= 201103L - #if __has_cpp_attribute(fallthrough) - #define CYTHON_FALLTHROUGH [[fallthrough]] - #elif __has_cpp_attribute(clang::fallthrough) - #define CYTHON_FALLTHROUGH [[clang::fallthrough]] - #elif __has_cpp_attribute(gnu::fallthrough) - #define CYTHON_FALLTHROUGH [[gnu::fallthrough]] - #endif - #endif - #ifndef CYTHON_FALLTHROUGH - #if __has_attribute(fallthrough) - #define CYTHON_FALLTHROUGH __attribute__((fallthrough)) - #else - #define CYTHON_FALLTHROUGH - #endif - #endif - #if defined(__clang__ ) && defined(__apple_build_version__) - #if __apple_build_version__ < 7000000 - #undef CYTHON_FALLTHROUGH - #define CYTHON_FALLTHROUGH - #endif - #endif -#endif - -#ifndef CYTHON_INLINE - #if defined(__clang__) - #define CYTHON_INLINE __inline__ __attribute__ ((__unused__)) - #elif defined(__GNUC__) - #define CYTHON_INLINE __inline__ - #elif defined(_MSC_VER) - #define CYTHON_INLINE __inline - #elif defined (__STDC_VERSION__) && __STDC_VERSION__ >= 199901L - #define CYTHON_INLINE inline - #else - #define CYTHON_INLINE - #endif -#endif - -#if CYTHON_COMPILING_IN_PYPY && PY_VERSION_HEX < 0x02070600 && !defined(Py_OptimizeFlag) - #define Py_OptimizeFlag 0 -#endif -#define __PYX_BUILD_PY_SSIZE_T "n" -#define CYTHON_FORMAT_SSIZE_T "z" -#if PY_MAJOR_VERSION < 3 - #define __Pyx_BUILTIN_MODULE_NAME "__builtin__" - #define __Pyx_PyCode_New(a, k, l, s, f, code, c, n, v, fv, cell, fn, name, fline, lnos)\ - PyCode_New(a+k, l, s, f, code, c, n, v, fv, cell, fn, name, fline, lnos) - #define __Pyx_DefaultClassType PyClass_Type -#else - #define __Pyx_BUILTIN_MODULE_NAME "builtins" - #define __Pyx_PyCode_New(a, k, l, s, f, code, c, n, v, fv, cell, fn, name, fline, lnos)\ - PyCode_New(a, k, l, s, f, code, c, n, v, fv, cell, fn, name, fline, lnos) - #define __Pyx_DefaultClassType PyType_Type -#endif -#ifndef Py_TPFLAGS_CHECKTYPES - #define Py_TPFLAGS_CHECKTYPES 0 -#endif -#ifndef Py_TPFLAGS_HAVE_INDEX - #define Py_TPFLAGS_HAVE_INDEX 0 -#endif -#ifndef Py_TPFLAGS_HAVE_NEWBUFFER - #define Py_TPFLAGS_HAVE_NEWBUFFER 0 -#endif -#ifndef Py_TPFLAGS_HAVE_FINALIZE - #define Py_TPFLAGS_HAVE_FINALIZE 0 -#endif -#if PY_VERSION_HEX <= 0x030700A3 || !defined(METH_FASTCALL) - #ifndef METH_FASTCALL - #define METH_FASTCALL 0x80 - #endif - typedef PyObject *(*__Pyx_PyCFunctionFast) (PyObject *self, PyObject *const *args, Py_ssize_t nargs); - typedef PyObject *(*__Pyx_PyCFunctionFastWithKeywords) (PyObject *self, PyObject *const *args, - Py_ssize_t nargs, PyObject *kwnames); -#else - #define __Pyx_PyCFunctionFast _PyCFunctionFast - #define __Pyx_PyCFunctionFastWithKeywords _PyCFunctionFastWithKeywords -#endif -#if CYTHON_FAST_PYCCALL -#define __Pyx_PyFastCFunction_Check(func)\ - ((PyCFunction_Check(func) && (METH_FASTCALL == (PyCFunction_GET_FLAGS(func) & ~(METH_CLASS | METH_STATIC | METH_COEXIST | METH_KEYWORDS))))) -#else -#define __Pyx_PyFastCFunction_Check(func) 0 -#endif -#if CYTHON_COMPILING_IN_PYPY && !defined(PyObject_Malloc) - #define PyObject_Malloc(s) PyMem_Malloc(s) - #define PyObject_Free(p) PyMem_Free(p) - #define PyObject_Realloc(p) PyMem_Realloc(p) -#endif -#if CYTHON_COMPILING_IN_PYSTON - #define __Pyx_PyCode_HasFreeVars(co) PyCode_HasFreeVars(co) - #define __Pyx_PyFrame_SetLineNumber(frame, lineno) PyFrame_SetLineNumber(frame, lineno) -#else - #define __Pyx_PyCode_HasFreeVars(co) (PyCode_GetNumFree(co) > 0) - #define __Pyx_PyFrame_SetLineNumber(frame, lineno) (frame)->f_lineno = (lineno) -#endif -#if !CYTHON_FAST_THREAD_STATE || PY_VERSION_HEX < 0x02070000 - #define __Pyx_PyThreadState_Current PyThreadState_GET() -#elif PY_VERSION_HEX >= 0x03060000 - #define __Pyx_PyThreadState_Current _PyThreadState_UncheckedGet() -#elif PY_VERSION_HEX >= 0x03000000 - #define __Pyx_PyThreadState_Current PyThreadState_GET() -#else - #define __Pyx_PyThreadState_Current _PyThreadState_Current -#endif -#if PY_VERSION_HEX < 0x030700A2 && !defined(PyThread_tss_create) && !defined(Py_tss_NEEDS_INIT) -#include "pythread.h" -#define Py_tss_NEEDS_INIT 0 -typedef int Py_tss_t; -static CYTHON_INLINE int PyThread_tss_create(Py_tss_t *key) { - *key = PyThread_create_key(); - return 0; // PyThread_create_key reports success always -} -static CYTHON_INLINE Py_tss_t * PyThread_tss_alloc(void) { - Py_tss_t *key = (Py_tss_t *)PyObject_Malloc(sizeof(Py_tss_t)); - *key = Py_tss_NEEDS_INIT; - return key; -} -static CYTHON_INLINE void PyThread_tss_free(Py_tss_t *key) { - PyObject_Free(key); -} -static CYTHON_INLINE int PyThread_tss_is_created(Py_tss_t *key) { - return *key != Py_tss_NEEDS_INIT; -} -static CYTHON_INLINE void PyThread_tss_delete(Py_tss_t *key) { - PyThread_delete_key(*key); - *key = Py_tss_NEEDS_INIT; -} -static CYTHON_INLINE int PyThread_tss_set(Py_tss_t *key, void *value) { - return PyThread_set_key_value(*key, value); -} -static CYTHON_INLINE void * PyThread_tss_get(Py_tss_t *key) { - return PyThread_get_key_value(*key); -} -#endif // TSS (Thread Specific Storage) API -#if CYTHON_COMPILING_IN_CPYTHON || defined(_PyDict_NewPresized) -#define __Pyx_PyDict_NewPresized(n) ((n <= 8) ? PyDict_New() : _PyDict_NewPresized(n)) -#else -#define __Pyx_PyDict_NewPresized(n) PyDict_New() -#endif -#if PY_MAJOR_VERSION >= 3 || CYTHON_FUTURE_DIVISION - #define __Pyx_PyNumber_Divide(x,y) PyNumber_TrueDivide(x,y) - #define __Pyx_PyNumber_InPlaceDivide(x,y) PyNumber_InPlaceTrueDivide(x,y) -#else - #define __Pyx_PyNumber_Divide(x,y) PyNumber_Divide(x,y) - #define __Pyx_PyNumber_InPlaceDivide(x,y) PyNumber_InPlaceDivide(x,y) -#endif -#if CYTHON_COMPILING_IN_CPYTHON && PY_VERSION_HEX >= 0x030500A1 && CYTHON_USE_UNICODE_INTERNALS -#define __Pyx_PyDict_GetItemStr(dict, name) _PyDict_GetItem_KnownHash(dict, name, ((PyASCIIObject *) name)->hash) -#else -#define __Pyx_PyDict_GetItemStr(dict, name) PyDict_GetItem(dict, name) -#endif -#if PY_VERSION_HEX > 0x03030000 && defined(PyUnicode_KIND) - #define CYTHON_PEP393_ENABLED 1 - #define __Pyx_PyUnicode_READY(op) (likely(PyUnicode_IS_READY(op)) ?\ - 0 : _PyUnicode_Ready((PyObject *)(op))) - #define __Pyx_PyUnicode_GET_LENGTH(u) PyUnicode_GET_LENGTH(u) - #define __Pyx_PyUnicode_READ_CHAR(u, i) PyUnicode_READ_CHAR(u, i) - #define __Pyx_PyUnicode_MAX_CHAR_VALUE(u) PyUnicode_MAX_CHAR_VALUE(u) - #define __Pyx_PyUnicode_KIND(u) PyUnicode_KIND(u) - #define __Pyx_PyUnicode_DATA(u) PyUnicode_DATA(u) - #define __Pyx_PyUnicode_READ(k, d, i) PyUnicode_READ(k, d, i) - #define __Pyx_PyUnicode_WRITE(k, d, i, ch) PyUnicode_WRITE(k, d, i, ch) - #define __Pyx_PyUnicode_IS_TRUE(u) (0 != (likely(PyUnicode_IS_READY(u)) ? PyUnicode_GET_LENGTH(u) : PyUnicode_GET_SIZE(u))) -#else - #define CYTHON_PEP393_ENABLED 0 - #define PyUnicode_1BYTE_KIND 1 - #define PyUnicode_2BYTE_KIND 2 - #define PyUnicode_4BYTE_KIND 4 - #define __Pyx_PyUnicode_READY(op) (0) - #define __Pyx_PyUnicode_GET_LENGTH(u) PyUnicode_GET_SIZE(u) - #define __Pyx_PyUnicode_READ_CHAR(u, i) ((Py_UCS4)(PyUnicode_AS_UNICODE(u)[i])) - #define __Pyx_PyUnicode_MAX_CHAR_VALUE(u) ((sizeof(Py_UNICODE) == 2) ? 65535 : 1114111) - #define __Pyx_PyUnicode_KIND(u) (sizeof(Py_UNICODE)) - #define __Pyx_PyUnicode_DATA(u) ((void*)PyUnicode_AS_UNICODE(u)) - #define __Pyx_PyUnicode_READ(k, d, i) ((void)(k), (Py_UCS4)(((Py_UNICODE*)d)[i])) - #define __Pyx_PyUnicode_WRITE(k, d, i, ch) (((void)(k)), ((Py_UNICODE*)d)[i] = ch) - #define __Pyx_PyUnicode_IS_TRUE(u) (0 != PyUnicode_GET_SIZE(u)) -#endif -#if CYTHON_COMPILING_IN_PYPY - #define __Pyx_PyUnicode_Concat(a, b) PyNumber_Add(a, b) - #define __Pyx_PyUnicode_ConcatSafe(a, b) PyNumber_Add(a, b) -#else - #define __Pyx_PyUnicode_Concat(a, b) PyUnicode_Concat(a, b) - #define __Pyx_PyUnicode_ConcatSafe(a, b) ((unlikely((a) == Py_None) || unlikely((b) == Py_None)) ?\ - PyNumber_Add(a, b) : __Pyx_PyUnicode_Concat(a, b)) -#endif -#if CYTHON_COMPILING_IN_PYPY && !defined(PyUnicode_Contains) - #define PyUnicode_Contains(u, s) PySequence_Contains(u, s) -#endif -#if CYTHON_COMPILING_IN_PYPY && !defined(PyByteArray_Check) - #define PyByteArray_Check(obj) PyObject_TypeCheck(obj, &PyByteArray_Type) -#endif -#if CYTHON_COMPILING_IN_PYPY && !defined(PyObject_Format) - #define PyObject_Format(obj, fmt) PyObject_CallMethod(obj, "__format__", "O", fmt) -#endif -#define __Pyx_PyString_FormatSafe(a, b) ((unlikely((a) == Py_None)) ? PyNumber_Remainder(a, b) : __Pyx_PyString_Format(a, b)) -#define __Pyx_PyUnicode_FormatSafe(a, b) ((unlikely((a) == Py_None)) ? PyNumber_Remainder(a, b) : PyUnicode_Format(a, b)) -#if PY_MAJOR_VERSION >= 3 - #define __Pyx_PyString_Format(a, b) PyUnicode_Format(a, b) -#else - #define __Pyx_PyString_Format(a, b) PyString_Format(a, b) -#endif -#if PY_MAJOR_VERSION < 3 && !defined(PyObject_ASCII) - #define PyObject_ASCII(o) PyObject_Repr(o) -#endif -#if PY_MAJOR_VERSION >= 3 - #define PyBaseString_Type PyUnicode_Type - #define PyStringObject PyUnicodeObject - #define PyString_Type PyUnicode_Type - #define PyString_Check PyUnicode_Check - #define PyString_CheckExact PyUnicode_CheckExact -#endif -#if PY_MAJOR_VERSION >= 3 - #define __Pyx_PyBaseString_Check(obj) PyUnicode_Check(obj) - #define __Pyx_PyBaseString_CheckExact(obj) PyUnicode_CheckExact(obj) -#else - #define __Pyx_PyBaseString_Check(obj) (PyString_Check(obj) || PyUnicode_Check(obj)) - #define __Pyx_PyBaseString_CheckExact(obj) (PyString_CheckExact(obj) || PyUnicode_CheckExact(obj)) -#endif -#ifndef PySet_CheckExact - #define PySet_CheckExact(obj) (Py_TYPE(obj) == &PySet_Type) -#endif -#if CYTHON_ASSUME_SAFE_MACROS - #define __Pyx_PySequence_SIZE(seq) Py_SIZE(seq) -#else - #define __Pyx_PySequence_SIZE(seq) PySequence_Size(seq) -#endif -#if PY_MAJOR_VERSION >= 3 - #define PyIntObject PyLongObject - #define PyInt_Type PyLong_Type - #define PyInt_Check(op) PyLong_Check(op) - #define PyInt_CheckExact(op) PyLong_CheckExact(op) - #define PyInt_FromString PyLong_FromString - #define PyInt_FromUnicode PyLong_FromUnicode - #define PyInt_FromLong PyLong_FromLong - #define PyInt_FromSize_t PyLong_FromSize_t - #define PyInt_FromSsize_t PyLong_FromSsize_t - #define PyInt_AsLong PyLong_AsLong - #define PyInt_AS_LONG PyLong_AS_LONG - #define PyInt_AsSsize_t PyLong_AsSsize_t - #define PyInt_AsUnsignedLongMask PyLong_AsUnsignedLongMask - #define PyInt_AsUnsignedLongLongMask PyLong_AsUnsignedLongLongMask - #define PyNumber_Int PyNumber_Long -#endif -#if PY_MAJOR_VERSION >= 3 - #define PyBoolObject PyLongObject -#endif -#if PY_MAJOR_VERSION >= 3 && CYTHON_COMPILING_IN_PYPY - #ifndef PyUnicode_InternFromString - #define PyUnicode_InternFromString(s) PyUnicode_FromString(s) - #endif -#endif -#if PY_VERSION_HEX < 0x030200A4 - typedef long Py_hash_t; - #define __Pyx_PyInt_FromHash_t PyInt_FromLong - #define __Pyx_PyInt_AsHash_t PyInt_AsLong -#else - #define __Pyx_PyInt_FromHash_t PyInt_FromSsize_t - #define __Pyx_PyInt_AsHash_t PyInt_AsSsize_t -#endif -#if PY_MAJOR_VERSION >= 3 - #define __Pyx_PyMethod_New(func, self, klass) ((self) ? PyMethod_New(func, self) : (Py_INCREF(func), func)) -#else - #define __Pyx_PyMethod_New(func, self, klass) PyMethod_New(func, self, klass) -#endif -#if CYTHON_USE_ASYNC_SLOTS - #if PY_VERSION_HEX >= 0x030500B1 - #define __Pyx_PyAsyncMethodsStruct PyAsyncMethods - #define __Pyx_PyType_AsAsync(obj) (Py_TYPE(obj)->tp_as_async) - #else - #define __Pyx_PyType_AsAsync(obj) ((__Pyx_PyAsyncMethodsStruct*) (Py_TYPE(obj)->tp_reserved)) - #endif -#else - #define __Pyx_PyType_AsAsync(obj) NULL -#endif -#ifndef __Pyx_PyAsyncMethodsStruct - typedef struct { - unaryfunc am_await; - unaryfunc am_aiter; - unaryfunc am_anext; - } __Pyx_PyAsyncMethodsStruct; -#endif - -#if defined(WIN32) || defined(MS_WINDOWS) - #define _USE_MATH_DEFINES -#endif -#include -#ifdef NAN -#define __PYX_NAN() ((float) NAN) -#else -static CYTHON_INLINE float __PYX_NAN() { - float value; - memset(&value, 0xFF, sizeof(value)); - return value; -} -#endif -#if defined(__CYGWIN__) && defined(_LDBL_EQ_DBL) -#define __Pyx_truncl trunc -#else -#define __Pyx_truncl truncl -#endif - - -#define __PYX_ERR(f_index, lineno, Ln_error) \ -{ \ - __pyx_filename = __pyx_f[f_index]; __pyx_lineno = lineno; __pyx_clineno = __LINE__; goto Ln_error; \ -} - -#ifndef __PYX_EXTERN_C - #ifdef __cplusplus - #define __PYX_EXTERN_C extern "C" - #else - #define __PYX_EXTERN_C extern - #endif -#endif - -#define __PYX_HAVE__mesh_core_cython -#define __PYX_HAVE_API__mesh_core_cython -/* Early includes */ -#include -#include -#include "numpy/arrayobject.h" -#include "numpy/ufuncobject.h" -#include "ios" -#include "new" -#include "stdexcept" -#include "typeinfo" -#include -#include "mesh_core.h" -#ifdef _OPENMP -#include -#endif /* _OPENMP */ - -#if defined(PYREX_WITHOUT_ASSERTIONS) && !defined(CYTHON_WITHOUT_ASSERTIONS) -#define CYTHON_WITHOUT_ASSERTIONS -#endif - -typedef struct {PyObject **p; const char *s; const Py_ssize_t n; const char* encoding; - const char is_unicode; const char is_str; const char intern; } __Pyx_StringTabEntry; - -#define __PYX_DEFAULT_STRING_ENCODING_IS_ASCII 0 -#define __PYX_DEFAULT_STRING_ENCODING_IS_DEFAULT 0 -#define __PYX_DEFAULT_STRING_ENCODING "" -#define __Pyx_PyObject_FromString __Pyx_PyBytes_FromString -#define __Pyx_PyObject_FromStringAndSize __Pyx_PyBytes_FromStringAndSize -#define __Pyx_uchar_cast(c) ((unsigned char)c) -#define __Pyx_long_cast(x) ((long)x) -#define __Pyx_fits_Py_ssize_t(v, type, is_signed) (\ - (sizeof(type) < sizeof(Py_ssize_t)) ||\ - (sizeof(type) > sizeof(Py_ssize_t) &&\ - likely(v < (type)PY_SSIZE_T_MAX ||\ - v == (type)PY_SSIZE_T_MAX) &&\ - (!is_signed || likely(v > (type)PY_SSIZE_T_MIN ||\ - v == (type)PY_SSIZE_T_MIN))) ||\ - (sizeof(type) == sizeof(Py_ssize_t) &&\ - (is_signed || likely(v < (type)PY_SSIZE_T_MAX ||\ - v == (type)PY_SSIZE_T_MAX))) ) -#if defined (__cplusplus) && __cplusplus >= 201103L - #include - #define __Pyx_sst_abs(value) std::abs(value) -#elif SIZEOF_INT >= SIZEOF_SIZE_T - #define __Pyx_sst_abs(value) abs(value) -#elif SIZEOF_LONG >= SIZEOF_SIZE_T - #define __Pyx_sst_abs(value) labs(value) -#elif defined (_MSC_VER) - #define __Pyx_sst_abs(value) ((Py_ssize_t)_abs64(value)) -#elif defined (__STDC_VERSION__) && __STDC_VERSION__ >= 199901L - #define __Pyx_sst_abs(value) llabs(value) -#elif defined (__GNUC__) - #define __Pyx_sst_abs(value) __builtin_llabs(value) -#else - #define __Pyx_sst_abs(value) ((value<0) ? -value : value) -#endif -static CYTHON_INLINE const char* __Pyx_PyObject_AsString(PyObject*); -static CYTHON_INLINE const char* __Pyx_PyObject_AsStringAndSize(PyObject*, Py_ssize_t* length); -#define __Pyx_PyByteArray_FromString(s) PyByteArray_FromStringAndSize((const char*)s, strlen((const char*)s)) -#define __Pyx_PyByteArray_FromStringAndSize(s, l) PyByteArray_FromStringAndSize((const char*)s, l) -#define __Pyx_PyBytes_FromString PyBytes_FromString -#define __Pyx_PyBytes_FromStringAndSize PyBytes_FromStringAndSize -static CYTHON_INLINE PyObject* __Pyx_PyUnicode_FromString(const char*); -#if PY_MAJOR_VERSION < 3 - #define __Pyx_PyStr_FromString __Pyx_PyBytes_FromString - #define __Pyx_PyStr_FromStringAndSize __Pyx_PyBytes_FromStringAndSize -#else - #define __Pyx_PyStr_FromString __Pyx_PyUnicode_FromString - #define __Pyx_PyStr_FromStringAndSize __Pyx_PyUnicode_FromStringAndSize -#endif -#define __Pyx_PyBytes_AsWritableString(s) ((char*) PyBytes_AS_STRING(s)) -#define __Pyx_PyBytes_AsWritableSString(s) ((signed char*) PyBytes_AS_STRING(s)) -#define __Pyx_PyBytes_AsWritableUString(s) ((unsigned char*) PyBytes_AS_STRING(s)) -#define __Pyx_PyBytes_AsString(s) ((const char*) PyBytes_AS_STRING(s)) -#define __Pyx_PyBytes_AsSString(s) ((const signed char*) PyBytes_AS_STRING(s)) -#define __Pyx_PyBytes_AsUString(s) ((const unsigned char*) PyBytes_AS_STRING(s)) -#define __Pyx_PyObject_AsWritableString(s) ((char*) __Pyx_PyObject_AsString(s)) -#define __Pyx_PyObject_AsWritableSString(s) ((signed char*) __Pyx_PyObject_AsString(s)) -#define __Pyx_PyObject_AsWritableUString(s) ((unsigned char*) __Pyx_PyObject_AsString(s)) -#define __Pyx_PyObject_AsSString(s) ((const signed char*) __Pyx_PyObject_AsString(s)) -#define __Pyx_PyObject_AsUString(s) ((const unsigned char*) __Pyx_PyObject_AsString(s)) -#define __Pyx_PyObject_FromCString(s) __Pyx_PyObject_FromString((const char*)s) -#define __Pyx_PyBytes_FromCString(s) __Pyx_PyBytes_FromString((const char*)s) -#define __Pyx_PyByteArray_FromCString(s) __Pyx_PyByteArray_FromString((const char*)s) -#define __Pyx_PyStr_FromCString(s) __Pyx_PyStr_FromString((const char*)s) -#define __Pyx_PyUnicode_FromCString(s) __Pyx_PyUnicode_FromString((const char*)s) -static CYTHON_INLINE size_t __Pyx_Py_UNICODE_strlen(const Py_UNICODE *u) { - const Py_UNICODE *u_end = u; - while (*u_end++) ; - return (size_t)(u_end - u - 1); -} -#define __Pyx_PyUnicode_FromUnicode(u) PyUnicode_FromUnicode(u, __Pyx_Py_UNICODE_strlen(u)) -#define __Pyx_PyUnicode_FromUnicodeAndLength PyUnicode_FromUnicode -#define __Pyx_PyUnicode_AsUnicode PyUnicode_AsUnicode -#define __Pyx_NewRef(obj) (Py_INCREF(obj), obj) -#define __Pyx_Owned_Py_None(b) __Pyx_NewRef(Py_None) -#define __Pyx_PyBool_FromLong(b) ((b) ? __Pyx_NewRef(Py_True) : __Pyx_NewRef(Py_False)) -static CYTHON_INLINE int __Pyx_PyObject_IsTrue(PyObject*); -static CYTHON_INLINE PyObject* __Pyx_PyNumber_IntOrLong(PyObject* x); -#define __Pyx_PySequence_Tuple(obj)\ - (likely(PyTuple_CheckExact(obj)) ? __Pyx_NewRef(obj) : PySequence_Tuple(obj)) -static CYTHON_INLINE Py_ssize_t __Pyx_PyIndex_AsSsize_t(PyObject*); -static CYTHON_INLINE PyObject * __Pyx_PyInt_FromSize_t(size_t); -#if CYTHON_ASSUME_SAFE_MACROS -#define __pyx_PyFloat_AsDouble(x) (PyFloat_CheckExact(x) ? PyFloat_AS_DOUBLE(x) : PyFloat_AsDouble(x)) -#else -#define __pyx_PyFloat_AsDouble(x) PyFloat_AsDouble(x) -#endif -#define __pyx_PyFloat_AsFloat(x) ((float) __pyx_PyFloat_AsDouble(x)) -#if PY_MAJOR_VERSION >= 3 -#define __Pyx_PyNumber_Int(x) (PyLong_CheckExact(x) ? __Pyx_NewRef(x) : PyNumber_Long(x)) -#else -#define __Pyx_PyNumber_Int(x) (PyInt_CheckExact(x) ? __Pyx_NewRef(x) : PyNumber_Int(x)) -#endif -#define __Pyx_PyNumber_Float(x) (PyFloat_CheckExact(x) ? __Pyx_NewRef(x) : PyNumber_Float(x)) -#if PY_MAJOR_VERSION < 3 && __PYX_DEFAULT_STRING_ENCODING_IS_ASCII -static int __Pyx_sys_getdefaultencoding_not_ascii; -static int __Pyx_init_sys_getdefaultencoding_params(void) { - PyObject* sys; - PyObject* default_encoding = NULL; - PyObject* ascii_chars_u = NULL; - PyObject* ascii_chars_b = NULL; - const char* default_encoding_c; - sys = PyImport_ImportModule("sys"); - if (!sys) goto bad; - default_encoding = PyObject_CallMethod(sys, (char*) "getdefaultencoding", NULL); - Py_DECREF(sys); - if (!default_encoding) goto bad; - default_encoding_c = PyBytes_AsString(default_encoding); - if (!default_encoding_c) goto bad; - if (strcmp(default_encoding_c, "ascii") == 0) { - __Pyx_sys_getdefaultencoding_not_ascii = 0; - } else { - char ascii_chars[128]; - int c; - for (c = 0; c < 128; c++) { - ascii_chars[c] = c; - } - __Pyx_sys_getdefaultencoding_not_ascii = 1; - ascii_chars_u = PyUnicode_DecodeASCII(ascii_chars, 128, NULL); - if (!ascii_chars_u) goto bad; - ascii_chars_b = PyUnicode_AsEncodedString(ascii_chars_u, default_encoding_c, NULL); - if (!ascii_chars_b || !PyBytes_Check(ascii_chars_b) || memcmp(ascii_chars, PyBytes_AS_STRING(ascii_chars_b), 128) != 0) { - PyErr_Format( - PyExc_ValueError, - "This module compiled with c_string_encoding=ascii, but default encoding '%.200s' is not a superset of ascii.", - default_encoding_c); - goto bad; - } - Py_DECREF(ascii_chars_u); - Py_DECREF(ascii_chars_b); - } - Py_DECREF(default_encoding); - return 0; -bad: - Py_XDECREF(default_encoding); - Py_XDECREF(ascii_chars_u); - Py_XDECREF(ascii_chars_b); - return -1; -} -#endif -#if __PYX_DEFAULT_STRING_ENCODING_IS_DEFAULT && PY_MAJOR_VERSION >= 3 -#define __Pyx_PyUnicode_FromStringAndSize(c_str, size) PyUnicode_DecodeUTF8(c_str, size, NULL) -#else -#define __Pyx_PyUnicode_FromStringAndSize(c_str, size) PyUnicode_Decode(c_str, size, __PYX_DEFAULT_STRING_ENCODING, NULL) -#if __PYX_DEFAULT_STRING_ENCODING_IS_DEFAULT -static char* __PYX_DEFAULT_STRING_ENCODING; -static int __Pyx_init_sys_getdefaultencoding_params(void) { - PyObject* sys; - PyObject* default_encoding = NULL; - char* default_encoding_c; - sys = PyImport_ImportModule("sys"); - if (!sys) goto bad; - default_encoding = PyObject_CallMethod(sys, (char*) (const char*) "getdefaultencoding", NULL); - Py_DECREF(sys); - if (!default_encoding) goto bad; - default_encoding_c = PyBytes_AsString(default_encoding); - if (!default_encoding_c) goto bad; - __PYX_DEFAULT_STRING_ENCODING = (char*) malloc(strlen(default_encoding_c)); - if (!__PYX_DEFAULT_STRING_ENCODING) goto bad; - strcpy(__PYX_DEFAULT_STRING_ENCODING, default_encoding_c); - Py_DECREF(default_encoding); - return 0; -bad: - Py_XDECREF(default_encoding); - return -1; -} -#endif -#endif - - -/* Test for GCC > 2.95 */ -#if defined(__GNUC__) && (__GNUC__ > 2 || (__GNUC__ == 2 && (__GNUC_MINOR__ > 95))) - #define likely(x) __builtin_expect(!!(x), 1) - #define unlikely(x) __builtin_expect(!!(x), 0) -#else /* !__GNUC__ or GCC < 2.95 */ - #define likely(x) (x) - #define unlikely(x) (x) -#endif /* __GNUC__ */ -static CYTHON_INLINE void __Pyx_pretend_to_initialize(void* ptr) { (void)ptr; } - -static PyObject *__pyx_m = NULL; -static PyObject *__pyx_d; -static PyObject *__pyx_b; -static PyObject *__pyx_cython_runtime; -static PyObject *__pyx_empty_tuple; -static PyObject *__pyx_empty_bytes; -static PyObject *__pyx_empty_unicode; -static int __pyx_lineno; -static int __pyx_clineno = 0; -static const char * __pyx_cfilenm= __FILE__; -static const char *__pyx_filename; - -/* Header.proto */ -#if !defined(CYTHON_CCOMPLEX) - #if defined(__cplusplus) - #define CYTHON_CCOMPLEX 1 - #elif defined(_Complex_I) - #define CYTHON_CCOMPLEX 1 - #else - #define CYTHON_CCOMPLEX 0 - #endif -#endif -#if CYTHON_CCOMPLEX - #ifdef __cplusplus - #include - #else - #include - #endif -#endif -#if CYTHON_CCOMPLEX && !defined(__cplusplus) && defined(__sun__) && defined(__GNUC__) - #undef _Complex_I - #define _Complex_I 1.0fj -#endif - - -static const char *__pyx_f[] = { - "insightface/thirdparty/face3d/mesh/cython/mesh_core_cython.pyx", - "__init__.pxd", - "stringsource", - "type.pxd", -}; -/* BufferFormatStructs.proto */ -#define IS_UNSIGNED(type) (((type) -1) > 0) -struct __Pyx_StructField_; -#define __PYX_BUF_FLAGS_PACKED_STRUCT (1 << 0) -typedef struct { - const char* name; - struct __Pyx_StructField_* fields; - size_t size; - size_t arraysize[8]; - int ndim; - char typegroup; - char is_unsigned; - int flags; -} __Pyx_TypeInfo; -typedef struct __Pyx_StructField_ { - __Pyx_TypeInfo* type; - const char* name; - size_t offset; -} __Pyx_StructField; -typedef struct { - __Pyx_StructField* field; - size_t parent_offset; -} __Pyx_BufFmt_StackElem; -typedef struct { - __Pyx_StructField root; - __Pyx_BufFmt_StackElem* head; - size_t fmt_offset; - size_t new_count, enc_count; - size_t struct_alignment; - int is_complex; - char enc_type; - char new_packmode; - char enc_packmode; - char is_valid_array; -} __Pyx_BufFmt_Context; - - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":730 - * # in Cython to enable them only on the right systems. - * - * ctypedef npy_int8 int8_t # <<<<<<<<<<<<<< - * ctypedef npy_int16 int16_t - * ctypedef npy_int32 int32_t - */ -typedef npy_int8 __pyx_t_5numpy_int8_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":731 - * - * ctypedef npy_int8 int8_t - * ctypedef npy_int16 int16_t # <<<<<<<<<<<<<< - * ctypedef npy_int32 int32_t - * ctypedef npy_int64 int64_t - */ -typedef npy_int16 __pyx_t_5numpy_int16_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":732 - * ctypedef npy_int8 int8_t - * ctypedef npy_int16 int16_t - * ctypedef npy_int32 int32_t # <<<<<<<<<<<<<< - * ctypedef npy_int64 int64_t - * #ctypedef npy_int96 int96_t - */ -typedef npy_int32 __pyx_t_5numpy_int32_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":733 - * ctypedef npy_int16 int16_t - * ctypedef npy_int32 int32_t - * ctypedef npy_int64 int64_t # <<<<<<<<<<<<<< - * #ctypedef npy_int96 int96_t - * #ctypedef npy_int128 int128_t - */ -typedef npy_int64 __pyx_t_5numpy_int64_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":737 - * #ctypedef npy_int128 int128_t - * - * ctypedef npy_uint8 uint8_t # <<<<<<<<<<<<<< - * ctypedef npy_uint16 uint16_t - * ctypedef npy_uint32 uint32_t - */ -typedef npy_uint8 __pyx_t_5numpy_uint8_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":738 - * - * ctypedef npy_uint8 uint8_t - * ctypedef npy_uint16 uint16_t # <<<<<<<<<<<<<< - * ctypedef npy_uint32 uint32_t - * ctypedef npy_uint64 uint64_t - */ -typedef npy_uint16 __pyx_t_5numpy_uint16_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":739 - * ctypedef npy_uint8 uint8_t - * ctypedef npy_uint16 uint16_t - * ctypedef npy_uint32 uint32_t # <<<<<<<<<<<<<< - * ctypedef npy_uint64 uint64_t - * #ctypedef npy_uint96 uint96_t - */ -typedef npy_uint32 __pyx_t_5numpy_uint32_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":740 - * ctypedef npy_uint16 uint16_t - * ctypedef npy_uint32 uint32_t - * ctypedef npy_uint64 uint64_t # <<<<<<<<<<<<<< - * #ctypedef npy_uint96 uint96_t - * #ctypedef npy_uint128 uint128_t - */ -typedef npy_uint64 __pyx_t_5numpy_uint64_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":744 - * #ctypedef npy_uint128 uint128_t - * - * ctypedef npy_float32 float32_t # <<<<<<<<<<<<<< - * ctypedef npy_float64 float64_t - * #ctypedef npy_float80 float80_t - */ -typedef npy_float32 __pyx_t_5numpy_float32_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":745 - * - * ctypedef npy_float32 float32_t - * ctypedef npy_float64 float64_t # <<<<<<<<<<<<<< - * #ctypedef npy_float80 float80_t - * #ctypedef npy_float128 float128_t - */ -typedef npy_float64 __pyx_t_5numpy_float64_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":754 - * # The int types are mapped a bit surprising -- - * # numpy.int corresponds to 'l' and numpy.long to 'q' - * ctypedef npy_long int_t # <<<<<<<<<<<<<< - * ctypedef npy_longlong long_t - * ctypedef npy_longlong longlong_t - */ -typedef npy_long __pyx_t_5numpy_int_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":755 - * # numpy.int corresponds to 'l' and numpy.long to 'q' - * ctypedef npy_long int_t - * ctypedef npy_longlong long_t # <<<<<<<<<<<<<< - * ctypedef npy_longlong longlong_t - * - */ -typedef npy_longlong __pyx_t_5numpy_long_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":756 - * ctypedef npy_long int_t - * ctypedef npy_longlong long_t - * ctypedef npy_longlong longlong_t # <<<<<<<<<<<<<< - * - * ctypedef npy_ulong uint_t - */ -typedef npy_longlong __pyx_t_5numpy_longlong_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":758 - * ctypedef npy_longlong longlong_t - * - * ctypedef npy_ulong uint_t # <<<<<<<<<<<<<< - * ctypedef npy_ulonglong ulong_t - * ctypedef npy_ulonglong ulonglong_t - */ -typedef npy_ulong __pyx_t_5numpy_uint_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":759 - * - * ctypedef npy_ulong uint_t - * ctypedef npy_ulonglong ulong_t # <<<<<<<<<<<<<< - * ctypedef npy_ulonglong ulonglong_t - * - */ -typedef npy_ulonglong __pyx_t_5numpy_ulong_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":760 - * ctypedef npy_ulong uint_t - * ctypedef npy_ulonglong ulong_t - * ctypedef npy_ulonglong ulonglong_t # <<<<<<<<<<<<<< - * - * ctypedef npy_intp intp_t - */ -typedef npy_ulonglong __pyx_t_5numpy_ulonglong_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":762 - * ctypedef npy_ulonglong ulonglong_t - * - * ctypedef npy_intp intp_t # <<<<<<<<<<<<<< - * ctypedef npy_uintp uintp_t - * - */ -typedef npy_intp __pyx_t_5numpy_intp_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":763 - * - * ctypedef npy_intp intp_t - * ctypedef npy_uintp uintp_t # <<<<<<<<<<<<<< - * - * ctypedef npy_double float_t - */ -typedef npy_uintp __pyx_t_5numpy_uintp_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":765 - * ctypedef npy_uintp uintp_t - * - * ctypedef npy_double float_t # <<<<<<<<<<<<<< - * ctypedef npy_double double_t - * ctypedef npy_longdouble longdouble_t - */ -typedef npy_double __pyx_t_5numpy_float_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":766 - * - * ctypedef npy_double float_t - * ctypedef npy_double double_t # <<<<<<<<<<<<<< - * ctypedef npy_longdouble longdouble_t - * - */ -typedef npy_double __pyx_t_5numpy_double_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":767 - * ctypedef npy_double float_t - * ctypedef npy_double double_t - * ctypedef npy_longdouble longdouble_t # <<<<<<<<<<<<<< - * - * ctypedef npy_cfloat cfloat_t - */ -typedef npy_longdouble __pyx_t_5numpy_longdouble_t; -/* Declarations.proto */ -#if CYTHON_CCOMPLEX - #ifdef __cplusplus - typedef ::std::complex< float > __pyx_t_float_complex; - #else - typedef float _Complex __pyx_t_float_complex; - #endif -#else - typedef struct { float real, imag; } __pyx_t_float_complex; -#endif -static CYTHON_INLINE __pyx_t_float_complex __pyx_t_float_complex_from_parts(float, float); - -/* Declarations.proto */ -#if CYTHON_CCOMPLEX - #ifdef __cplusplus - typedef ::std::complex< double > __pyx_t_double_complex; - #else - typedef double _Complex __pyx_t_double_complex; - #endif -#else - typedef struct { double real, imag; } __pyx_t_double_complex; -#endif -static CYTHON_INLINE __pyx_t_double_complex __pyx_t_double_complex_from_parts(double, double); - - -/*--- Type declarations ---*/ - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":769 - * ctypedef npy_longdouble longdouble_t - * - * ctypedef npy_cfloat cfloat_t # <<<<<<<<<<<<<< - * ctypedef npy_cdouble cdouble_t - * ctypedef npy_clongdouble clongdouble_t - */ -typedef npy_cfloat __pyx_t_5numpy_cfloat_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":770 - * - * ctypedef npy_cfloat cfloat_t - * ctypedef npy_cdouble cdouble_t # <<<<<<<<<<<<<< - * ctypedef npy_clongdouble clongdouble_t - * - */ -typedef npy_cdouble __pyx_t_5numpy_cdouble_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":771 - * ctypedef npy_cfloat cfloat_t - * ctypedef npy_cdouble cdouble_t - * ctypedef npy_clongdouble clongdouble_t # <<<<<<<<<<<<<< - * - * ctypedef npy_cdouble complex_t - */ -typedef npy_clongdouble __pyx_t_5numpy_clongdouble_t; - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":773 - * ctypedef npy_clongdouble clongdouble_t - * - * ctypedef npy_cdouble complex_t # <<<<<<<<<<<<<< - * - * cdef inline object PyArray_MultiIterNew1(a): - */ -typedef npy_cdouble __pyx_t_5numpy_complex_t; - -/* --- Runtime support code (head) --- */ -/* Refnanny.proto */ -#ifndef CYTHON_REFNANNY - #define CYTHON_REFNANNY 0 -#endif -#if CYTHON_REFNANNY - typedef struct { - void (*INCREF)(void*, PyObject*, int); - void (*DECREF)(void*, PyObject*, int); - void (*GOTREF)(void*, PyObject*, int); - void (*GIVEREF)(void*, PyObject*, int); - void* (*SetupContext)(const char*, int, const char*); - void (*FinishContext)(void**); - } __Pyx_RefNannyAPIStruct; - static __Pyx_RefNannyAPIStruct *__Pyx_RefNanny = NULL; - static __Pyx_RefNannyAPIStruct *__Pyx_RefNannyImportAPI(const char *modname); - #define __Pyx_RefNannyDeclarations void *__pyx_refnanny = NULL; -#ifdef WITH_THREAD - #define __Pyx_RefNannySetupContext(name, acquire_gil)\ - if (acquire_gil) {\ - PyGILState_STATE __pyx_gilstate_save = PyGILState_Ensure();\ - __pyx_refnanny = __Pyx_RefNanny->SetupContext((name), __LINE__, __FILE__);\ - PyGILState_Release(__pyx_gilstate_save);\ - } else {\ - __pyx_refnanny = __Pyx_RefNanny->SetupContext((name), __LINE__, __FILE__);\ - } -#else - #define __Pyx_RefNannySetupContext(name, acquire_gil)\ - __pyx_refnanny = __Pyx_RefNanny->SetupContext((name), __LINE__, __FILE__) -#endif - #define __Pyx_RefNannyFinishContext()\ - __Pyx_RefNanny->FinishContext(&__pyx_refnanny) - #define __Pyx_INCREF(r) __Pyx_RefNanny->INCREF(__pyx_refnanny, (PyObject *)(r), __LINE__) - #define __Pyx_DECREF(r) __Pyx_RefNanny->DECREF(__pyx_refnanny, (PyObject *)(r), __LINE__) - #define __Pyx_GOTREF(r) __Pyx_RefNanny->GOTREF(__pyx_refnanny, (PyObject *)(r), __LINE__) - #define __Pyx_GIVEREF(r) __Pyx_RefNanny->GIVEREF(__pyx_refnanny, (PyObject *)(r), __LINE__) - #define __Pyx_XINCREF(r) do { if((r) != NULL) {__Pyx_INCREF(r); }} while(0) - #define __Pyx_XDECREF(r) do { if((r) != NULL) {__Pyx_DECREF(r); }} while(0) - #define __Pyx_XGOTREF(r) do { if((r) != NULL) {__Pyx_GOTREF(r); }} while(0) - #define __Pyx_XGIVEREF(r) do { if((r) != NULL) {__Pyx_GIVEREF(r);}} while(0) -#else - #define __Pyx_RefNannyDeclarations - #define __Pyx_RefNannySetupContext(name, acquire_gil) - #define __Pyx_RefNannyFinishContext() - #define __Pyx_INCREF(r) Py_INCREF(r) - #define __Pyx_DECREF(r) Py_DECREF(r) - #define __Pyx_GOTREF(r) - #define __Pyx_GIVEREF(r) - #define __Pyx_XINCREF(r) Py_XINCREF(r) - #define __Pyx_XDECREF(r) Py_XDECREF(r) - #define __Pyx_XGOTREF(r) - #define __Pyx_XGIVEREF(r) -#endif -#define __Pyx_XDECREF_SET(r, v) do {\ - PyObject *tmp = (PyObject *) r;\ - r = v; __Pyx_XDECREF(tmp);\ - } while (0) -#define __Pyx_DECREF_SET(r, v) do {\ - PyObject *tmp = (PyObject *) r;\ - r = v; __Pyx_DECREF(tmp);\ - } while (0) -#define __Pyx_CLEAR(r) do { PyObject* tmp = ((PyObject*)(r)); r = NULL; __Pyx_DECREF(tmp);} while(0) -#define __Pyx_XCLEAR(r) do { if((r) != NULL) {PyObject* tmp = ((PyObject*)(r)); r = NULL; __Pyx_DECREF(tmp);}} while(0) - -/* RaiseArgTupleInvalid.proto */ -static void __Pyx_RaiseArgtupleInvalid(const char* func_name, int exact, - Py_ssize_t num_min, Py_ssize_t num_max, Py_ssize_t num_found); - -/* RaiseDoubleKeywords.proto */ -static void __Pyx_RaiseDoubleKeywordsError(const char* func_name, PyObject* kw_name); - -/* ParseKeywords.proto */ -static int __Pyx_ParseOptionalKeywords(PyObject *kwds, PyObject **argnames[],\ - PyObject *kwds2, PyObject *values[], Py_ssize_t num_pos_args,\ - const char* function_name); - -/* ArgTypeTest.proto */ -#define __Pyx_ArgTypeTest(obj, type, none_allowed, name, exact)\ - ((likely((Py_TYPE(obj) == type) | (none_allowed && (obj == Py_None)))) ? 1 :\ - __Pyx__ArgTypeTest(obj, type, name, exact)) -static int __Pyx__ArgTypeTest(PyObject *obj, PyTypeObject *type, const char *name, int exact); - -/* IsLittleEndian.proto */ -static CYTHON_INLINE int __Pyx_Is_Little_Endian(void); - -/* BufferFormatCheck.proto */ -static const char* __Pyx_BufFmt_CheckString(__Pyx_BufFmt_Context* ctx, const char* ts); -static void __Pyx_BufFmt_Init(__Pyx_BufFmt_Context* ctx, - __Pyx_BufFmt_StackElem* stack, - __Pyx_TypeInfo* type); - -/* BufferGetAndValidate.proto */ -#define __Pyx_GetBufferAndValidate(buf, obj, dtype, flags, nd, cast, stack)\ - ((obj == Py_None || obj == NULL) ?\ - (__Pyx_ZeroBuffer(buf), 0) :\ - __Pyx__GetBufferAndValidate(buf, obj, dtype, flags, nd, cast, stack)) -static int __Pyx__GetBufferAndValidate(Py_buffer* buf, PyObject* obj, - __Pyx_TypeInfo* dtype, int flags, int nd, int cast, __Pyx_BufFmt_StackElem* stack); -static void __Pyx_ZeroBuffer(Py_buffer* buf); -static CYTHON_INLINE void __Pyx_SafeReleaseBuffer(Py_buffer* info); -static Py_ssize_t __Pyx_minusones[] = { -1, -1, -1, -1, -1, -1, -1, -1 }; -static Py_ssize_t __Pyx_zeros[] = { 0, 0, 0, 0, 0, 0, 0, 0 }; - -/* PyThreadStateGet.proto */ -#if CYTHON_FAST_THREAD_STATE -#define __Pyx_PyThreadState_declare PyThreadState *__pyx_tstate; -#define __Pyx_PyThreadState_assign __pyx_tstate = __Pyx_PyThreadState_Current; -#define __Pyx_PyErr_Occurred() __pyx_tstate->curexc_type -#else -#define __Pyx_PyThreadState_declare -#define __Pyx_PyThreadState_assign -#define __Pyx_PyErr_Occurred() PyErr_Occurred() -#endif - -/* PyErrFetchRestore.proto */ -#if CYTHON_FAST_THREAD_STATE -#define __Pyx_PyErr_Clear() __Pyx_ErrRestore(NULL, NULL, NULL) -#define __Pyx_ErrRestoreWithState(type, value, tb) __Pyx_ErrRestoreInState(PyThreadState_GET(), type, value, tb) -#define __Pyx_ErrFetchWithState(type, value, tb) __Pyx_ErrFetchInState(PyThreadState_GET(), type, value, tb) -#define __Pyx_ErrRestore(type, value, tb) __Pyx_ErrRestoreInState(__pyx_tstate, type, value, tb) -#define __Pyx_ErrFetch(type, value, tb) __Pyx_ErrFetchInState(__pyx_tstate, type, value, tb) -static CYTHON_INLINE void __Pyx_ErrRestoreInState(PyThreadState *tstate, PyObject *type, PyObject *value, PyObject *tb); -static CYTHON_INLINE void __Pyx_ErrFetchInState(PyThreadState *tstate, PyObject **type, PyObject **value, PyObject **tb); -#if CYTHON_COMPILING_IN_CPYTHON -#define __Pyx_PyErr_SetNone(exc) (Py_INCREF(exc), __Pyx_ErrRestore((exc), NULL, NULL)) -#else -#define __Pyx_PyErr_SetNone(exc) PyErr_SetNone(exc) -#endif -#else -#define __Pyx_PyErr_Clear() PyErr_Clear() -#define __Pyx_PyErr_SetNone(exc) PyErr_SetNone(exc) -#define __Pyx_ErrRestoreWithState(type, value, tb) PyErr_Restore(type, value, tb) -#define __Pyx_ErrFetchWithState(type, value, tb) PyErr_Fetch(type, value, tb) -#define __Pyx_ErrRestoreInState(tstate, type, value, tb) PyErr_Restore(type, value, tb) -#define __Pyx_ErrFetchInState(tstate, type, value, tb) PyErr_Fetch(type, value, tb) -#define __Pyx_ErrRestore(type, value, tb) PyErr_Restore(type, value, tb) -#define __Pyx_ErrFetch(type, value, tb) PyErr_Fetch(type, value, tb) -#endif - -/* PyObjectGetAttrStr.proto */ -#if CYTHON_USE_TYPE_SLOTS -static CYTHON_INLINE PyObject* __Pyx_PyObject_GetAttrStr(PyObject* obj, PyObject* attr_name); -#else -#define __Pyx_PyObject_GetAttrStr(o,n) PyObject_GetAttr(o,n) -#endif - -/* GetBuiltinName.proto */ -static PyObject *__Pyx_GetBuiltinName(PyObject *name); - -/* PyObjectCall.proto */ -#if CYTHON_COMPILING_IN_CPYTHON -static CYTHON_INLINE PyObject* __Pyx_PyObject_Call(PyObject *func, PyObject *arg, PyObject *kw); -#else -#define __Pyx_PyObject_Call(func, arg, kw) PyObject_Call(func, arg, kw) -#endif - -/* RaiseException.proto */ -static void __Pyx_Raise(PyObject *type, PyObject *value, PyObject *tb, PyObject *cause); - -/* PyCFunctionFastCall.proto */ -#if CYTHON_FAST_PYCCALL -static CYTHON_INLINE PyObject *__Pyx_PyCFunction_FastCall(PyObject *func, PyObject **args, Py_ssize_t nargs); -#else -#define __Pyx_PyCFunction_FastCall(func, args, nargs) (assert(0), NULL) -#endif - -/* PyFunctionFastCall.proto */ -#if CYTHON_FAST_PYCALL -#define __Pyx_PyFunction_FastCall(func, args, nargs)\ - __Pyx_PyFunction_FastCallDict((func), (args), (nargs), NULL) -#if 1 || PY_VERSION_HEX < 0x030600B1 -static PyObject *__Pyx_PyFunction_FastCallDict(PyObject *func, PyObject **args, int nargs, PyObject *kwargs); -#else -#define __Pyx_PyFunction_FastCallDict(func, args, nargs, kwargs) _PyFunction_FastCallDict(func, args, nargs, kwargs) -#endif -#endif - -/* PyObjectCallMethO.proto */ -#if CYTHON_COMPILING_IN_CPYTHON -static CYTHON_INLINE PyObject* __Pyx_PyObject_CallMethO(PyObject *func, PyObject *arg); -#endif - -/* PyObjectCallOneArg.proto */ -static CYTHON_INLINE PyObject* __Pyx_PyObject_CallOneArg(PyObject *func, PyObject *arg); - -/* DictGetItem.proto */ -#if PY_MAJOR_VERSION >= 3 && !CYTHON_COMPILING_IN_PYPY -static PyObject *__Pyx_PyDict_GetItem(PyObject *d, PyObject* key); -#define __Pyx_PyObject_Dict_GetItem(obj, name)\ - (likely(PyDict_CheckExact(obj)) ?\ - __Pyx_PyDict_GetItem(obj, name) : PyObject_GetItem(obj, name)) -#else -#define __Pyx_PyDict_GetItem(d, key) PyObject_GetItem(d, key) -#define __Pyx_PyObject_Dict_GetItem(obj, name) PyObject_GetItem(obj, name) -#endif - -/* RaiseTooManyValuesToUnpack.proto */ -static CYTHON_INLINE void __Pyx_RaiseTooManyValuesError(Py_ssize_t expected); - -/* RaiseNeedMoreValuesToUnpack.proto */ -static CYTHON_INLINE void __Pyx_RaiseNeedMoreValuesError(Py_ssize_t index); - -/* RaiseNoneIterError.proto */ -static CYTHON_INLINE void __Pyx_RaiseNoneNotIterableError(void); - -/* ExtTypeTest.proto */ -static CYTHON_INLINE int __Pyx_TypeTest(PyObject *obj, PyTypeObject *type); - -/* SaveResetException.proto */ -#if CYTHON_FAST_THREAD_STATE -#define __Pyx_ExceptionSave(type, value, tb) __Pyx__ExceptionSave(__pyx_tstate, type, value, tb) -static CYTHON_INLINE void __Pyx__ExceptionSave(PyThreadState *tstate, PyObject **type, PyObject **value, PyObject **tb); -#define __Pyx_ExceptionReset(type, value, tb) __Pyx__ExceptionReset(__pyx_tstate, type, value, tb) -static CYTHON_INLINE void __Pyx__ExceptionReset(PyThreadState *tstate, PyObject *type, PyObject *value, PyObject *tb); -#else -#define __Pyx_ExceptionSave(type, value, tb) PyErr_GetExcInfo(type, value, tb) -#define __Pyx_ExceptionReset(type, value, tb) PyErr_SetExcInfo(type, value, tb) -#endif - -/* PyErrExceptionMatches.proto */ -#if CYTHON_FAST_THREAD_STATE -#define __Pyx_PyErr_ExceptionMatches(err) __Pyx_PyErr_ExceptionMatchesInState(__pyx_tstate, err) -static CYTHON_INLINE int __Pyx_PyErr_ExceptionMatchesInState(PyThreadState* tstate, PyObject* err); -#else -#define __Pyx_PyErr_ExceptionMatches(err) PyErr_ExceptionMatches(err) -#endif - -/* GetException.proto */ -#if CYTHON_FAST_THREAD_STATE -#define __Pyx_GetException(type, value, tb) __Pyx__GetException(__pyx_tstate, type, value, tb) -static int __Pyx__GetException(PyThreadState *tstate, PyObject **type, PyObject **value, PyObject **tb); -#else -static int __Pyx_GetException(PyObject **type, PyObject **value, PyObject **tb); -#endif - -/* Import.proto */ -static PyObject *__Pyx_Import(PyObject *name, PyObject *from_list, int level); - -/* CLineInTraceback.proto */ -#ifdef CYTHON_CLINE_IN_TRACEBACK -#define __Pyx_CLineForTraceback(tstate, c_line) (((CYTHON_CLINE_IN_TRACEBACK)) ? c_line : 0) -#else -static int __Pyx_CLineForTraceback(PyThreadState *tstate, int c_line); -#endif - -/* CodeObjectCache.proto */ -typedef struct { - PyCodeObject* code_object; - int code_line; -} __Pyx_CodeObjectCacheEntry; -struct __Pyx_CodeObjectCache { - int count; - int max_count; - __Pyx_CodeObjectCacheEntry* entries; -}; -static struct __Pyx_CodeObjectCache __pyx_code_cache = {0,0,NULL}; -static int __pyx_bisect_code_objects(__Pyx_CodeObjectCacheEntry* entries, int count, int code_line); -static PyCodeObject *__pyx_find_code_object(int code_line); -static void __pyx_insert_code_object(int code_line, PyCodeObject* code_object); - -/* AddTraceback.proto */ -static void __Pyx_AddTraceback(const char *funcname, int c_line, - int py_line, const char *filename); - -/* BufferStructDeclare.proto */ -typedef struct { - Py_ssize_t shape, strides, suboffsets; -} __Pyx_Buf_DimInfo; -typedef struct { - size_t refcount; - Py_buffer pybuffer; -} __Pyx_Buffer; -typedef struct { - __Pyx_Buffer *rcbuffer; - char *data; - __Pyx_Buf_DimInfo diminfo[8]; -} __Pyx_LocalBuf_ND; - -#if PY_MAJOR_VERSION < 3 - static int __Pyx_GetBuffer(PyObject *obj, Py_buffer *view, int flags); - static void __Pyx_ReleaseBuffer(Py_buffer *view); -#else - #define __Pyx_GetBuffer PyObject_GetBuffer - #define __Pyx_ReleaseBuffer PyBuffer_Release -#endif - - -/* RealImag.proto */ -#if CYTHON_CCOMPLEX - #ifdef __cplusplus - #define __Pyx_CREAL(z) ((z).real()) - #define __Pyx_CIMAG(z) ((z).imag()) - #else - #define __Pyx_CREAL(z) (__real__(z)) - #define __Pyx_CIMAG(z) (__imag__(z)) - #endif -#else - #define __Pyx_CREAL(z) ((z).real) - #define __Pyx_CIMAG(z) ((z).imag) -#endif -#if defined(__cplusplus) && CYTHON_CCOMPLEX\ - && (defined(_WIN32) || defined(__clang__) || (defined(__GNUC__) && (__GNUC__ >= 5 || __GNUC__ == 4 && __GNUC_MINOR__ >= 4 )) || __cplusplus >= 201103) - #define __Pyx_SET_CREAL(z,x) ((z).real(x)) - #define __Pyx_SET_CIMAG(z,y) ((z).imag(y)) -#else - #define __Pyx_SET_CREAL(z,x) __Pyx_CREAL(z) = (x) - #define __Pyx_SET_CIMAG(z,y) __Pyx_CIMAG(z) = (y) -#endif - -/* Arithmetic.proto */ -#if CYTHON_CCOMPLEX - #define __Pyx_c_eq_float(a, b) ((a)==(b)) - #define __Pyx_c_sum_float(a, b) ((a)+(b)) - #define __Pyx_c_diff_float(a, b) ((a)-(b)) - #define __Pyx_c_prod_float(a, b) ((a)*(b)) - #define __Pyx_c_quot_float(a, b) ((a)/(b)) - #define __Pyx_c_neg_float(a) (-(a)) - #ifdef __cplusplus - #define __Pyx_c_is_zero_float(z) ((z)==(float)0) - #define __Pyx_c_conj_float(z) (::std::conj(z)) - #if 1 - #define __Pyx_c_abs_float(z) (::std::abs(z)) - #define __Pyx_c_pow_float(a, b) (::std::pow(a, b)) - #endif - #else - #define __Pyx_c_is_zero_float(z) ((z)==0) - #define __Pyx_c_conj_float(z) (conjf(z)) - #if 1 - #define __Pyx_c_abs_float(z) (cabsf(z)) - #define __Pyx_c_pow_float(a, b) (cpowf(a, b)) - #endif - #endif -#else - static CYTHON_INLINE int __Pyx_c_eq_float(__pyx_t_float_complex, __pyx_t_float_complex); - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_sum_float(__pyx_t_float_complex, __pyx_t_float_complex); - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_diff_float(__pyx_t_float_complex, __pyx_t_float_complex); - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_prod_float(__pyx_t_float_complex, __pyx_t_float_complex); - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_quot_float(__pyx_t_float_complex, __pyx_t_float_complex); - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_neg_float(__pyx_t_float_complex); - static CYTHON_INLINE int __Pyx_c_is_zero_float(__pyx_t_float_complex); - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_conj_float(__pyx_t_float_complex); - #if 1 - static CYTHON_INLINE float __Pyx_c_abs_float(__pyx_t_float_complex); - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_pow_float(__pyx_t_float_complex, __pyx_t_float_complex); - #endif -#endif - -/* Arithmetic.proto */ -#if CYTHON_CCOMPLEX - #define __Pyx_c_eq_double(a, b) ((a)==(b)) - #define __Pyx_c_sum_double(a, b) ((a)+(b)) - #define __Pyx_c_diff_double(a, b) ((a)-(b)) - #define __Pyx_c_prod_double(a, b) ((a)*(b)) - #define __Pyx_c_quot_double(a, b) ((a)/(b)) - #define __Pyx_c_neg_double(a) (-(a)) - #ifdef __cplusplus - #define __Pyx_c_is_zero_double(z) ((z)==(double)0) - #define __Pyx_c_conj_double(z) (::std::conj(z)) - #if 1 - #define __Pyx_c_abs_double(z) (::std::abs(z)) - #define __Pyx_c_pow_double(a, b) (::std::pow(a, b)) - #endif - #else - #define __Pyx_c_is_zero_double(z) ((z)==0) - #define __Pyx_c_conj_double(z) (conj(z)) - #if 1 - #define __Pyx_c_abs_double(z) (cabs(z)) - #define __Pyx_c_pow_double(a, b) (cpow(a, b)) - #endif - #endif -#else - static CYTHON_INLINE int __Pyx_c_eq_double(__pyx_t_double_complex, __pyx_t_double_complex); - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_sum_double(__pyx_t_double_complex, __pyx_t_double_complex); - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_diff_double(__pyx_t_double_complex, __pyx_t_double_complex); - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_prod_double(__pyx_t_double_complex, __pyx_t_double_complex); - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_quot_double(__pyx_t_double_complex, __pyx_t_double_complex); - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_neg_double(__pyx_t_double_complex); - static CYTHON_INLINE int __Pyx_c_is_zero_double(__pyx_t_double_complex); - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_conj_double(__pyx_t_double_complex); - #if 1 - static CYTHON_INLINE double __Pyx_c_abs_double(__pyx_t_double_complex); - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_pow_double(__pyx_t_double_complex, __pyx_t_double_complex); - #endif -#endif - -/* CIntToPy.proto */ -static CYTHON_INLINE PyObject* __Pyx_PyInt_From_int(int value); - -/* CIntToPy.proto */ -static CYTHON_INLINE PyObject* __Pyx_PyInt_From_enum__NPY_TYPES(enum NPY_TYPES value); - -/* CIntFromPy.proto */ -static CYTHON_INLINE int __Pyx_PyInt_As_int(PyObject *); - -/* CIntToPy.proto */ -static CYTHON_INLINE PyObject* __Pyx_PyInt_From_long(long value); - -/* CIntFromPy.proto */ -static CYTHON_INLINE long __Pyx_PyInt_As_long(PyObject *); - -/* FastTypeChecks.proto */ -#if CYTHON_COMPILING_IN_CPYTHON -#define __Pyx_TypeCheck(obj, type) __Pyx_IsSubtype(Py_TYPE(obj), (PyTypeObject *)type) -static CYTHON_INLINE int __Pyx_IsSubtype(PyTypeObject *a, PyTypeObject *b); -static CYTHON_INLINE int __Pyx_PyErr_GivenExceptionMatches(PyObject *err, PyObject *type); -static CYTHON_INLINE int __Pyx_PyErr_GivenExceptionMatches2(PyObject *err, PyObject *type1, PyObject *type2); -#else -#define __Pyx_TypeCheck(obj, type) PyObject_TypeCheck(obj, (PyTypeObject *)type) -#define __Pyx_PyErr_GivenExceptionMatches(err, type) PyErr_GivenExceptionMatches(err, type) -#define __Pyx_PyErr_GivenExceptionMatches2(err, type1, type2) (PyErr_GivenExceptionMatches(err, type1) || PyErr_GivenExceptionMatches(err, type2)) -#endif -#define __Pyx_PyException_Check(obj) __Pyx_TypeCheck(obj, PyExc_Exception) - -/* CheckBinaryVersion.proto */ -static int __Pyx_check_binary_version(void); - -/* PyIdentifierFromString.proto */ -#if !defined(__Pyx_PyIdentifier_FromString) -#if PY_MAJOR_VERSION < 3 - #define __Pyx_PyIdentifier_FromString(s) PyString_FromString(s) -#else - #define __Pyx_PyIdentifier_FromString(s) PyUnicode_FromString(s) -#endif -#endif - -/* ModuleImport.proto */ -static PyObject *__Pyx_ImportModule(const char *name); - -/* TypeImport.proto */ -static PyTypeObject *__Pyx_ImportType(const char *module_name, const char *class_name, size_t size, int strict); - -/* InitStrings.proto */ -static int __Pyx_InitStrings(__Pyx_StringTabEntry *t); - - -/* Module declarations from 'cpython.buffer' */ - -/* Module declarations from 'libc.string' */ - -/* Module declarations from 'libc.stdio' */ - -/* Module declarations from '__builtin__' */ - -/* Module declarations from 'cpython.type' */ -static PyTypeObject *__pyx_ptype_7cpython_4type_type = 0; - -/* Module declarations from 'cpython' */ - -/* Module declarations from 'cpython.object' */ - -/* Module declarations from 'cpython.ref' */ - -/* Module declarations from 'cpython.mem' */ - -/* Module declarations from 'numpy' */ - -/* Module declarations from 'numpy' */ -static PyTypeObject *__pyx_ptype_5numpy_dtype = 0; -static PyTypeObject *__pyx_ptype_5numpy_flatiter = 0; -static PyTypeObject *__pyx_ptype_5numpy_broadcast = 0; -static PyTypeObject *__pyx_ptype_5numpy_ndarray = 0; -static PyTypeObject *__pyx_ptype_5numpy_ufunc = 0; -static CYTHON_INLINE char *__pyx_f_5numpy__util_dtypestring(PyArray_Descr *, char *, char *, int *); /*proto*/ -static CYTHON_INLINE int __pyx_f_5numpy_import_array(void); /*proto*/ - -/* Module declarations from 'libcpp.string' */ - -/* Module declarations from 'mesh_core_cython' */ -static std::string __pyx_convert_string_from_py_std__in_string(PyObject *); /*proto*/ -static __Pyx_TypeInfo __Pyx_TypeInfo_float = { "float", NULL, sizeof(float), { 0 }, 0, 'R', 0, 0 }; -static __Pyx_TypeInfo __Pyx_TypeInfo_int = { "int", NULL, sizeof(int), { 0 }, 0, IS_UNSIGNED(int) ? 'U' : 'I', IS_UNSIGNED(int), 0 }; -#define __Pyx_MODULE_NAME "mesh_core_cython" -extern int __pyx_module_is_main_mesh_core_cython; -int __pyx_module_is_main_mesh_core_cython = 0; - -/* Implementation of 'mesh_core_cython' */ -static PyObject *__pyx_builtin_ValueError; -static PyObject *__pyx_builtin_range; -static PyObject *__pyx_builtin_RuntimeError; -static PyObject *__pyx_builtin_ImportError; -static const char __pyx_k_c[] = "c"; -static const char __pyx_k_h[] = "h"; -static const char __pyx_k_w[] = "w"; -static const char __pyx_k_np[] = "np"; -static const char __pyx_k_main[] = "__main__"; -static const char __pyx_k_ntri[] = "ntri"; -static const char __pyx_k_nver[] = "nver"; -static const char __pyx_k_test[] = "__test__"; -static const char __pyx_k_image[] = "image"; -static const char __pyx_k_numpy[] = "numpy"; -static const char __pyx_k_range[] = "range"; -static const char __pyx_k_tex_c[] = "tex_c"; -static const char __pyx_k_tex_h[] = "tex_h"; -static const char __pyx_k_tex_w[] = "tex_w"; -static const char __pyx_k_colors[] = "colors"; -static const char __pyx_k_import[] = "__import__"; -static const char __pyx_k_normal[] = "normal"; -static const char __pyx_k_ntexver[] = "ntexver"; -static const char __pyx_k_texture[] = "texture"; -static const char __pyx_k_filename[] = "filename"; -static const char __pyx_k_mtl_name[] = "mtl_name"; -static const char __pyx_k_tex_nver[] = "tex_nver"; -static const char __pyx_k_vertices[] = "vertices"; -static const char __pyx_k_triangles[] = "triangles"; -static const char __pyx_k_uv_coords[] = "uv_coords"; -static const char __pyx_k_ValueError[] = "ValueError"; -static const char __pyx_k_tex_coords[] = "tex_coords"; -static const char __pyx_k_tri_normal[] = "tri_normal"; -static const char __pyx_k_ImportError[] = "ImportError"; -static const char __pyx_k_RuntimeError[] = "RuntimeError"; -static const char __pyx_k_depth_buffer[] = "depth_buffer"; -static const char __pyx_k_mapping_type[] = "mapping_type"; -static const char __pyx_k_tex_triangles[] = "tex_triangles"; -static const char __pyx_k_get_normal_core[] = "get_normal_core"; -static const char __pyx_k_triangle_buffer[] = "triangle_buffer"; -static const char __pyx_k_mesh_core_cython[] = "mesh_core_cython"; -static const char __pyx_k_barycentric_weight[] = "barycentric_weight"; -static const char __pyx_k_cline_in_traceback[] = "cline_in_traceback"; -static const char __pyx_k_render_colors_core[] = "render_colors_core"; -static const char __pyx_k_render_texture_core[] = "render_texture_core"; -static const char __pyx_k_rasterize_triangles_core[] = "rasterize_triangles_core"; -static const char __pyx_k_ndarray_is_not_C_contiguous[] = "ndarray is not C contiguous"; -static const char __pyx_k_numpy_core_multiarray_failed_to[] = "numpy.core.multiarray failed to import"; -static const char __pyx_k_unknown_dtype_code_in_numpy_pxd[] = "unknown dtype code in numpy.pxd (%d)"; -static const char __pyx_k_Format_string_allocated_too_shor[] = "Format string allocated too short, see comment in numpy.pxd"; -static const char __pyx_k_Non_native_byte_order_not_suppor[] = "Non-native byte order not supported"; -static const char __pyx_k_insightface_thirdparty_face3d_me[] = "insightface/thirdparty/face3d/mesh/cython/mesh_core_cython.pyx"; -static const char __pyx_k_ndarray_is_not_Fortran_contiguou[] = "ndarray is not Fortran contiguous"; -static const char __pyx_k_numpy_core_umath_failed_to_impor[] = "numpy.core.umath failed to import"; -static const char __pyx_k_write_obj_with_colors_texture_co[] = "write_obj_with_colors_texture_core"; -static const char __pyx_k_Format_string_allocated_too_shor_2[] = "Format string allocated too short."; -static PyObject *__pyx_kp_u_Format_string_allocated_too_shor; -static PyObject *__pyx_kp_u_Format_string_allocated_too_shor_2; -static PyObject *__pyx_n_s_ImportError; -static PyObject *__pyx_kp_u_Non_native_byte_order_not_suppor; -static PyObject *__pyx_n_s_RuntimeError; -static PyObject *__pyx_n_s_ValueError; -static PyObject *__pyx_n_s_barycentric_weight; -static PyObject *__pyx_n_s_c; -static PyObject *__pyx_n_s_cline_in_traceback; -static PyObject *__pyx_n_s_colors; -static PyObject *__pyx_n_s_depth_buffer; -static PyObject *__pyx_n_s_filename; -static PyObject *__pyx_n_s_get_normal_core; -static PyObject *__pyx_n_s_h; -static PyObject *__pyx_n_s_image; -static PyObject *__pyx_n_s_import; -static PyObject *__pyx_kp_s_insightface_thirdparty_face3d_me; -static PyObject *__pyx_n_s_main; -static PyObject *__pyx_n_s_mapping_type; -static PyObject *__pyx_n_s_mesh_core_cython; -static PyObject *__pyx_n_s_mtl_name; -static PyObject *__pyx_kp_u_ndarray_is_not_C_contiguous; -static PyObject *__pyx_kp_u_ndarray_is_not_Fortran_contiguou; -static PyObject *__pyx_n_s_normal; -static PyObject *__pyx_n_s_np; -static PyObject *__pyx_n_s_ntexver; -static PyObject *__pyx_n_s_ntri; -static PyObject *__pyx_n_s_numpy; -static PyObject *__pyx_kp_s_numpy_core_multiarray_failed_to; -static PyObject *__pyx_kp_s_numpy_core_umath_failed_to_impor; -static PyObject *__pyx_n_s_nver; -static PyObject *__pyx_n_s_range; -static PyObject *__pyx_n_s_rasterize_triangles_core; -static PyObject *__pyx_n_s_render_colors_core; -static PyObject *__pyx_n_s_render_texture_core; -static PyObject *__pyx_n_s_test; -static PyObject *__pyx_n_s_tex_c; -static PyObject *__pyx_n_s_tex_coords; -static PyObject *__pyx_n_s_tex_h; -static PyObject *__pyx_n_s_tex_nver; -static PyObject *__pyx_n_s_tex_triangles; -static PyObject *__pyx_n_s_tex_w; -static PyObject *__pyx_n_s_texture; -static PyObject *__pyx_n_s_tri_normal; -static PyObject *__pyx_n_s_triangle_buffer; -static PyObject *__pyx_n_s_triangles; -static PyObject *__pyx_kp_u_unknown_dtype_code_in_numpy_pxd; -static PyObject *__pyx_n_s_uv_coords; -static PyObject *__pyx_n_s_vertices; -static PyObject *__pyx_n_s_w; -static PyObject *__pyx_n_s_write_obj_with_colors_texture_co; -static PyObject *__pyx_pf_16mesh_core_cython_get_normal_core(CYTHON_UNUSED PyObject *__pyx_self, PyArrayObject *__pyx_v_normal, PyArrayObject *__pyx_v_tri_normal, PyArrayObject *__pyx_v_triangles, int __pyx_v_ntri); /* proto */ -static PyObject *__pyx_pf_16mesh_core_cython_2rasterize_triangles_core(CYTHON_UNUSED PyObject *__pyx_self, PyArrayObject *__pyx_v_vertices, PyArrayObject *__pyx_v_triangles, PyArrayObject *__pyx_v_depth_buffer, PyArrayObject *__pyx_v_triangle_buffer, PyArrayObject *__pyx_v_barycentric_weight, int __pyx_v_nver, int __pyx_v_ntri, int __pyx_v_h, int __pyx_v_w); /* proto */ -static PyObject *__pyx_pf_16mesh_core_cython_4render_colors_core(CYTHON_UNUSED PyObject *__pyx_self, PyArrayObject *__pyx_v_image, PyArrayObject *__pyx_v_vertices, PyArrayObject *__pyx_v_triangles, PyArrayObject *__pyx_v_colors, PyArrayObject *__pyx_v_depth_buffer, int __pyx_v_nver, int __pyx_v_ntri, int __pyx_v_h, int __pyx_v_w, int __pyx_v_c); /* proto */ -static PyObject *__pyx_pf_16mesh_core_cython_6render_texture_core(CYTHON_UNUSED PyObject *__pyx_self, PyArrayObject *__pyx_v_image, PyArrayObject *__pyx_v_vertices, PyArrayObject *__pyx_v_triangles, PyArrayObject *__pyx_v_texture, PyArrayObject *__pyx_v_tex_coords, PyArrayObject *__pyx_v_tex_triangles, PyArrayObject *__pyx_v_depth_buffer, int __pyx_v_nver, int __pyx_v_tex_nver, int __pyx_v_ntri, int __pyx_v_h, int __pyx_v_w, int __pyx_v_c, int __pyx_v_tex_h, int __pyx_v_tex_w, int __pyx_v_tex_c, int __pyx_v_mapping_type); /* proto */ -static PyObject *__pyx_pf_16mesh_core_cython_8write_obj_with_colors_texture_core(CYTHON_UNUSED PyObject *__pyx_self, std::string __pyx_v_filename, std::string __pyx_v_mtl_name, PyArrayObject *__pyx_v_vertices, PyArrayObject *__pyx_v_triangles, PyArrayObject *__pyx_v_colors, PyArrayObject *__pyx_v_uv_coords, int __pyx_v_nver, int __pyx_v_ntri, int __pyx_v_ntexver); /* proto */ -static int __pyx_pf_5numpy_7ndarray___getbuffer__(PyArrayObject *__pyx_v_self, Py_buffer *__pyx_v_info, int __pyx_v_flags); /* proto */ -static void __pyx_pf_5numpy_7ndarray_2__releasebuffer__(PyArrayObject *__pyx_v_self, Py_buffer *__pyx_v_info); /* proto */ -static PyObject *__pyx_tuple_; -static PyObject *__pyx_tuple__2; -static PyObject *__pyx_tuple__3; -static PyObject *__pyx_tuple__4; -static PyObject *__pyx_tuple__5; -static PyObject *__pyx_tuple__6; -static PyObject *__pyx_tuple__7; -static PyObject *__pyx_tuple__8; -static PyObject *__pyx_tuple__9; -static PyObject *__pyx_tuple__10; -static PyObject *__pyx_tuple__12; -static PyObject *__pyx_tuple__14; -static PyObject *__pyx_tuple__16; -static PyObject *__pyx_tuple__18; -static PyObject *__pyx_codeobj__11; -static PyObject *__pyx_codeobj__13; -static PyObject *__pyx_codeobj__15; -static PyObject *__pyx_codeobj__17; -static PyObject *__pyx_codeobj__19; -/* Late includes */ - -/* "mesh_core_cython.pyx":40 - * int nver, int ntri, int ntexver) - * - * def get_normal_core(np.ndarray[float, ndim=2, mode = "c"] normal not None, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] tri_normal not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - -/* Python wrapper */ -static PyObject *__pyx_pw_16mesh_core_cython_1get_normal_core(PyObject *__pyx_self, PyObject *__pyx_args, PyObject *__pyx_kwds); /*proto*/ -static PyMethodDef __pyx_mdef_16mesh_core_cython_1get_normal_core = {"get_normal_core", (PyCFunction)__pyx_pw_16mesh_core_cython_1get_normal_core, METH_VARARGS|METH_KEYWORDS, 0}; -static PyObject *__pyx_pw_16mesh_core_cython_1get_normal_core(PyObject *__pyx_self, PyObject *__pyx_args, PyObject *__pyx_kwds) { - PyArrayObject *__pyx_v_normal = 0; - PyArrayObject *__pyx_v_tri_normal = 0; - PyArrayObject *__pyx_v_triangles = 0; - int __pyx_v_ntri; - PyObject *__pyx_r = 0; - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("get_normal_core (wrapper)", 0); - { - static PyObject **__pyx_pyargnames[] = {&__pyx_n_s_normal,&__pyx_n_s_tri_normal,&__pyx_n_s_triangles,&__pyx_n_s_ntri,0}; - PyObject* values[4] = {0,0,0,0}; - if (unlikely(__pyx_kwds)) { - Py_ssize_t kw_args; - const Py_ssize_t pos_args = PyTuple_GET_SIZE(__pyx_args); - switch (pos_args) { - case 4: values[3] = PyTuple_GET_ITEM(__pyx_args, 3); - CYTHON_FALLTHROUGH; - case 3: values[2] = PyTuple_GET_ITEM(__pyx_args, 2); - CYTHON_FALLTHROUGH; - case 2: values[1] = PyTuple_GET_ITEM(__pyx_args, 1); - CYTHON_FALLTHROUGH; - case 1: values[0] = PyTuple_GET_ITEM(__pyx_args, 0); - CYTHON_FALLTHROUGH; - case 0: break; - default: goto __pyx_L5_argtuple_error; - } - kw_args = PyDict_Size(__pyx_kwds); - switch (pos_args) { - case 0: - if (likely((values[0] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_normal)) != 0)) kw_args--; - else goto __pyx_L5_argtuple_error; - CYTHON_FALLTHROUGH; - case 1: - if (likely((values[1] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_tri_normal)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("get_normal_core", 1, 4, 4, 1); __PYX_ERR(0, 40, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 2: - if (likely((values[2] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_triangles)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("get_normal_core", 1, 4, 4, 2); __PYX_ERR(0, 40, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 3: - if (likely((values[3] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_ntri)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("get_normal_core", 1, 4, 4, 3); __PYX_ERR(0, 40, __pyx_L3_error) - } - } - if (unlikely(kw_args > 0)) { - if (unlikely(__Pyx_ParseOptionalKeywords(__pyx_kwds, __pyx_pyargnames, 0, values, pos_args, "get_normal_core") < 0)) __PYX_ERR(0, 40, __pyx_L3_error) - } - } else if (PyTuple_GET_SIZE(__pyx_args) != 4) { - goto __pyx_L5_argtuple_error; - } else { - values[0] = PyTuple_GET_ITEM(__pyx_args, 0); - values[1] = PyTuple_GET_ITEM(__pyx_args, 1); - values[2] = PyTuple_GET_ITEM(__pyx_args, 2); - values[3] = PyTuple_GET_ITEM(__pyx_args, 3); - } - __pyx_v_normal = ((PyArrayObject *)values[0]); - __pyx_v_tri_normal = ((PyArrayObject *)values[1]); - __pyx_v_triangles = ((PyArrayObject *)values[2]); - __pyx_v_ntri = __Pyx_PyInt_As_int(values[3]); if (unlikely((__pyx_v_ntri == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 43, __pyx_L3_error) - } - goto __pyx_L4_argument_unpacking_done; - __pyx_L5_argtuple_error:; - __Pyx_RaiseArgtupleInvalid("get_normal_core", 1, 4, 4, PyTuple_GET_SIZE(__pyx_args)); __PYX_ERR(0, 40, __pyx_L3_error) - __pyx_L3_error:; - __Pyx_AddTraceback("mesh_core_cython.get_normal_core", __pyx_clineno, __pyx_lineno, __pyx_filename); - __Pyx_RefNannyFinishContext(); - return NULL; - __pyx_L4_argument_unpacking_done:; - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_normal), __pyx_ptype_5numpy_ndarray, 0, "normal", 0))) __PYX_ERR(0, 40, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_tri_normal), __pyx_ptype_5numpy_ndarray, 0, "tri_normal", 0))) __PYX_ERR(0, 41, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_triangles), __pyx_ptype_5numpy_ndarray, 0, "triangles", 0))) __PYX_ERR(0, 42, __pyx_L1_error) - __pyx_r = __pyx_pf_16mesh_core_cython_get_normal_core(__pyx_self, __pyx_v_normal, __pyx_v_tri_normal, __pyx_v_triangles, __pyx_v_ntri); - - /* function exit code */ - goto __pyx_L0; - __pyx_L1_error:; - __pyx_r = NULL; - __pyx_L0:; - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -static PyObject *__pyx_pf_16mesh_core_cython_get_normal_core(CYTHON_UNUSED PyObject *__pyx_self, PyArrayObject *__pyx_v_normal, PyArrayObject *__pyx_v_tri_normal, PyArrayObject *__pyx_v_triangles, int __pyx_v_ntri) { - __Pyx_LocalBuf_ND __pyx_pybuffernd_normal; - __Pyx_Buffer __pyx_pybuffer_normal; - __Pyx_LocalBuf_ND __pyx_pybuffernd_tri_normal; - __Pyx_Buffer __pyx_pybuffer_tri_normal; - __Pyx_LocalBuf_ND __pyx_pybuffernd_triangles; - __Pyx_Buffer __pyx_pybuffer_triangles; - PyObject *__pyx_r = NULL; - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("get_normal_core", 0); - __pyx_pybuffer_normal.pybuffer.buf = NULL; - __pyx_pybuffer_normal.refcount = 0; - __pyx_pybuffernd_normal.data = NULL; - __pyx_pybuffernd_normal.rcbuffer = &__pyx_pybuffer_normal; - __pyx_pybuffer_tri_normal.pybuffer.buf = NULL; - __pyx_pybuffer_tri_normal.refcount = 0; - __pyx_pybuffernd_tri_normal.data = NULL; - __pyx_pybuffernd_tri_normal.rcbuffer = &__pyx_pybuffer_tri_normal; - __pyx_pybuffer_triangles.pybuffer.buf = NULL; - __pyx_pybuffer_triangles.refcount = 0; - __pyx_pybuffernd_triangles.data = NULL; - __pyx_pybuffernd_triangles.rcbuffer = &__pyx_pybuffer_triangles; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_normal.rcbuffer->pybuffer, (PyObject*)__pyx_v_normal, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 40, __pyx_L1_error) - } - __pyx_pybuffernd_normal.diminfo[0].strides = __pyx_pybuffernd_normal.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_normal.diminfo[0].shape = __pyx_pybuffernd_normal.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_normal.diminfo[1].strides = __pyx_pybuffernd_normal.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_normal.diminfo[1].shape = __pyx_pybuffernd_normal.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_tri_normal.rcbuffer->pybuffer, (PyObject*)__pyx_v_tri_normal, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 40, __pyx_L1_error) - } - __pyx_pybuffernd_tri_normal.diminfo[0].strides = __pyx_pybuffernd_tri_normal.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_tri_normal.diminfo[0].shape = __pyx_pybuffernd_tri_normal.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_tri_normal.diminfo[1].strides = __pyx_pybuffernd_tri_normal.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_tri_normal.diminfo[1].shape = __pyx_pybuffernd_tri_normal.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer, (PyObject*)__pyx_v_triangles, &__Pyx_TypeInfo_int, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 40, __pyx_L1_error) - } - __pyx_pybuffernd_triangles.diminfo[0].strides = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_triangles.diminfo[0].shape = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_triangles.diminfo[1].strides = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_triangles.diminfo[1].shape = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.shape[1]; - - /* "mesh_core_cython.pyx":45 - * int ntri - * ): - * _get_normal_core( # <<<<<<<<<<<<<< - * np.PyArray_DATA(normal), np.PyArray_DATA(tri_normal), np.PyArray_DATA(triangles), - * ntri) - */ - _get_normal_core(((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_normal))), ((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_tri_normal))), ((int *)PyArray_DATA(((PyArrayObject *)__pyx_v_triangles))), __pyx_v_ntri); - - /* "mesh_core_cython.pyx":40 - * int nver, int ntri, int ntexver) - * - * def get_normal_core(np.ndarray[float, ndim=2, mode = "c"] normal not None, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] tri_normal not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - - /* function exit code */ - __pyx_r = Py_None; __Pyx_INCREF(Py_None); - goto __pyx_L0; - __pyx_L1_error:; - { PyObject *__pyx_type, *__pyx_value, *__pyx_tb; - __Pyx_PyThreadState_declare - __Pyx_PyThreadState_assign - __Pyx_ErrFetch(&__pyx_type, &__pyx_value, &__pyx_tb); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_normal.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_tri_normal.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer); - __Pyx_ErrRestore(__pyx_type, __pyx_value, __pyx_tb);} - __Pyx_AddTraceback("mesh_core_cython.get_normal_core", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = NULL; - goto __pyx_L2; - __pyx_L0:; - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_normal.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_tri_normal.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer); - __pyx_L2:; - __Pyx_XGIVEREF(__pyx_r); - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "mesh_core_cython.pyx":49 - * ntri) - * - * def rasterize_triangles_core( # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - -/* Python wrapper */ -static PyObject *__pyx_pw_16mesh_core_cython_3rasterize_triangles_core(PyObject *__pyx_self, PyObject *__pyx_args, PyObject *__pyx_kwds); /*proto*/ -static PyMethodDef __pyx_mdef_16mesh_core_cython_3rasterize_triangles_core = {"rasterize_triangles_core", (PyCFunction)__pyx_pw_16mesh_core_cython_3rasterize_triangles_core, METH_VARARGS|METH_KEYWORDS, 0}; -static PyObject *__pyx_pw_16mesh_core_cython_3rasterize_triangles_core(PyObject *__pyx_self, PyObject *__pyx_args, PyObject *__pyx_kwds) { - PyArrayObject *__pyx_v_vertices = 0; - PyArrayObject *__pyx_v_triangles = 0; - PyArrayObject *__pyx_v_depth_buffer = 0; - PyArrayObject *__pyx_v_triangle_buffer = 0; - PyArrayObject *__pyx_v_barycentric_weight = 0; - int __pyx_v_nver; - int __pyx_v_ntri; - int __pyx_v_h; - int __pyx_v_w; - PyObject *__pyx_r = 0; - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("rasterize_triangles_core (wrapper)", 0); - { - static PyObject **__pyx_pyargnames[] = {&__pyx_n_s_vertices,&__pyx_n_s_triangles,&__pyx_n_s_depth_buffer,&__pyx_n_s_triangle_buffer,&__pyx_n_s_barycentric_weight,&__pyx_n_s_nver,&__pyx_n_s_ntri,&__pyx_n_s_h,&__pyx_n_s_w,0}; - PyObject* values[9] = {0,0,0,0,0,0,0,0,0}; - if (unlikely(__pyx_kwds)) { - Py_ssize_t kw_args; - const Py_ssize_t pos_args = PyTuple_GET_SIZE(__pyx_args); - switch (pos_args) { - case 9: values[8] = PyTuple_GET_ITEM(__pyx_args, 8); - CYTHON_FALLTHROUGH; - case 8: values[7] = PyTuple_GET_ITEM(__pyx_args, 7); - CYTHON_FALLTHROUGH; - case 7: values[6] = PyTuple_GET_ITEM(__pyx_args, 6); - CYTHON_FALLTHROUGH; - case 6: values[5] = PyTuple_GET_ITEM(__pyx_args, 5); - CYTHON_FALLTHROUGH; - case 5: values[4] = PyTuple_GET_ITEM(__pyx_args, 4); - CYTHON_FALLTHROUGH; - case 4: values[3] = PyTuple_GET_ITEM(__pyx_args, 3); - CYTHON_FALLTHROUGH; - case 3: values[2] = PyTuple_GET_ITEM(__pyx_args, 2); - CYTHON_FALLTHROUGH; - case 2: values[1] = PyTuple_GET_ITEM(__pyx_args, 1); - CYTHON_FALLTHROUGH; - case 1: values[0] = PyTuple_GET_ITEM(__pyx_args, 0); - CYTHON_FALLTHROUGH; - case 0: break; - default: goto __pyx_L5_argtuple_error; - } - kw_args = PyDict_Size(__pyx_kwds); - switch (pos_args) { - case 0: - if (likely((values[0] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_vertices)) != 0)) kw_args--; - else goto __pyx_L5_argtuple_error; - CYTHON_FALLTHROUGH; - case 1: - if (likely((values[1] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_triangles)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("rasterize_triangles_core", 1, 9, 9, 1); __PYX_ERR(0, 49, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 2: - if (likely((values[2] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_depth_buffer)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("rasterize_triangles_core", 1, 9, 9, 2); __PYX_ERR(0, 49, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 3: - if (likely((values[3] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_triangle_buffer)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("rasterize_triangles_core", 1, 9, 9, 3); __PYX_ERR(0, 49, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 4: - if (likely((values[4] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_barycentric_weight)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("rasterize_triangles_core", 1, 9, 9, 4); __PYX_ERR(0, 49, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 5: - if (likely((values[5] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_nver)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("rasterize_triangles_core", 1, 9, 9, 5); __PYX_ERR(0, 49, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 6: - if (likely((values[6] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_ntri)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("rasterize_triangles_core", 1, 9, 9, 6); __PYX_ERR(0, 49, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 7: - if (likely((values[7] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_h)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("rasterize_triangles_core", 1, 9, 9, 7); __PYX_ERR(0, 49, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 8: - if (likely((values[8] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_w)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("rasterize_triangles_core", 1, 9, 9, 8); __PYX_ERR(0, 49, __pyx_L3_error) - } - } - if (unlikely(kw_args > 0)) { - if (unlikely(__Pyx_ParseOptionalKeywords(__pyx_kwds, __pyx_pyargnames, 0, values, pos_args, "rasterize_triangles_core") < 0)) __PYX_ERR(0, 49, __pyx_L3_error) - } - } else if (PyTuple_GET_SIZE(__pyx_args) != 9) { - goto __pyx_L5_argtuple_error; - } else { - values[0] = PyTuple_GET_ITEM(__pyx_args, 0); - values[1] = PyTuple_GET_ITEM(__pyx_args, 1); - values[2] = PyTuple_GET_ITEM(__pyx_args, 2); - values[3] = PyTuple_GET_ITEM(__pyx_args, 3); - values[4] = PyTuple_GET_ITEM(__pyx_args, 4); - values[5] = PyTuple_GET_ITEM(__pyx_args, 5); - values[6] = PyTuple_GET_ITEM(__pyx_args, 6); - values[7] = PyTuple_GET_ITEM(__pyx_args, 7); - values[8] = PyTuple_GET_ITEM(__pyx_args, 8); - } - __pyx_v_vertices = ((PyArrayObject *)values[0]); - __pyx_v_triangles = ((PyArrayObject *)values[1]); - __pyx_v_depth_buffer = ((PyArrayObject *)values[2]); - __pyx_v_triangle_buffer = ((PyArrayObject *)values[3]); - __pyx_v_barycentric_weight = ((PyArrayObject *)values[4]); - __pyx_v_nver = __Pyx_PyInt_As_int(values[5]); if (unlikely((__pyx_v_nver == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 55, __pyx_L3_error) - __pyx_v_ntri = __Pyx_PyInt_As_int(values[6]); if (unlikely((__pyx_v_ntri == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 55, __pyx_L3_error) - __pyx_v_h = __Pyx_PyInt_As_int(values[7]); if (unlikely((__pyx_v_h == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 56, __pyx_L3_error) - __pyx_v_w = __Pyx_PyInt_As_int(values[8]); if (unlikely((__pyx_v_w == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 56, __pyx_L3_error) - } - goto __pyx_L4_argument_unpacking_done; - __pyx_L5_argtuple_error:; - __Pyx_RaiseArgtupleInvalid("rasterize_triangles_core", 1, 9, 9, PyTuple_GET_SIZE(__pyx_args)); __PYX_ERR(0, 49, __pyx_L3_error) - __pyx_L3_error:; - __Pyx_AddTraceback("mesh_core_cython.rasterize_triangles_core", __pyx_clineno, __pyx_lineno, __pyx_filename); - __Pyx_RefNannyFinishContext(); - return NULL; - __pyx_L4_argument_unpacking_done:; - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_vertices), __pyx_ptype_5numpy_ndarray, 0, "vertices", 0))) __PYX_ERR(0, 50, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_triangles), __pyx_ptype_5numpy_ndarray, 0, "triangles", 0))) __PYX_ERR(0, 51, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_depth_buffer), __pyx_ptype_5numpy_ndarray, 0, "depth_buffer", 0))) __PYX_ERR(0, 52, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_triangle_buffer), __pyx_ptype_5numpy_ndarray, 0, "triangle_buffer", 0))) __PYX_ERR(0, 53, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_barycentric_weight), __pyx_ptype_5numpy_ndarray, 0, "barycentric_weight", 0))) __PYX_ERR(0, 54, __pyx_L1_error) - __pyx_r = __pyx_pf_16mesh_core_cython_2rasterize_triangles_core(__pyx_self, __pyx_v_vertices, __pyx_v_triangles, __pyx_v_depth_buffer, __pyx_v_triangle_buffer, __pyx_v_barycentric_weight, __pyx_v_nver, __pyx_v_ntri, __pyx_v_h, __pyx_v_w); - - /* function exit code */ - goto __pyx_L0; - __pyx_L1_error:; - __pyx_r = NULL; - __pyx_L0:; - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -static PyObject *__pyx_pf_16mesh_core_cython_2rasterize_triangles_core(CYTHON_UNUSED PyObject *__pyx_self, PyArrayObject *__pyx_v_vertices, PyArrayObject *__pyx_v_triangles, PyArrayObject *__pyx_v_depth_buffer, PyArrayObject *__pyx_v_triangle_buffer, PyArrayObject *__pyx_v_barycentric_weight, int __pyx_v_nver, int __pyx_v_ntri, int __pyx_v_h, int __pyx_v_w) { - __Pyx_LocalBuf_ND __pyx_pybuffernd_barycentric_weight; - __Pyx_Buffer __pyx_pybuffer_barycentric_weight; - __Pyx_LocalBuf_ND __pyx_pybuffernd_depth_buffer; - __Pyx_Buffer __pyx_pybuffer_depth_buffer; - __Pyx_LocalBuf_ND __pyx_pybuffernd_triangle_buffer; - __Pyx_Buffer __pyx_pybuffer_triangle_buffer; - __Pyx_LocalBuf_ND __pyx_pybuffernd_triangles; - __Pyx_Buffer __pyx_pybuffer_triangles; - __Pyx_LocalBuf_ND __pyx_pybuffernd_vertices; - __Pyx_Buffer __pyx_pybuffer_vertices; - PyObject *__pyx_r = NULL; - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("rasterize_triangles_core", 0); - __pyx_pybuffer_vertices.pybuffer.buf = NULL; - __pyx_pybuffer_vertices.refcount = 0; - __pyx_pybuffernd_vertices.data = NULL; - __pyx_pybuffernd_vertices.rcbuffer = &__pyx_pybuffer_vertices; - __pyx_pybuffer_triangles.pybuffer.buf = NULL; - __pyx_pybuffer_triangles.refcount = 0; - __pyx_pybuffernd_triangles.data = NULL; - __pyx_pybuffernd_triangles.rcbuffer = &__pyx_pybuffer_triangles; - __pyx_pybuffer_depth_buffer.pybuffer.buf = NULL; - __pyx_pybuffer_depth_buffer.refcount = 0; - __pyx_pybuffernd_depth_buffer.data = NULL; - __pyx_pybuffernd_depth_buffer.rcbuffer = &__pyx_pybuffer_depth_buffer; - __pyx_pybuffer_triangle_buffer.pybuffer.buf = NULL; - __pyx_pybuffer_triangle_buffer.refcount = 0; - __pyx_pybuffernd_triangle_buffer.data = NULL; - __pyx_pybuffernd_triangle_buffer.rcbuffer = &__pyx_pybuffer_triangle_buffer; - __pyx_pybuffer_barycentric_weight.pybuffer.buf = NULL; - __pyx_pybuffer_barycentric_weight.refcount = 0; - __pyx_pybuffernd_barycentric_weight.data = NULL; - __pyx_pybuffernd_barycentric_weight.rcbuffer = &__pyx_pybuffer_barycentric_weight; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_vertices.rcbuffer->pybuffer, (PyObject*)__pyx_v_vertices, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 49, __pyx_L1_error) - } - __pyx_pybuffernd_vertices.diminfo[0].strides = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_vertices.diminfo[0].shape = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_vertices.diminfo[1].strides = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_vertices.diminfo[1].shape = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer, (PyObject*)__pyx_v_triangles, &__Pyx_TypeInfo_int, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 49, __pyx_L1_error) - } - __pyx_pybuffernd_triangles.diminfo[0].strides = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_triangles.diminfo[0].shape = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_triangles.diminfo[1].strides = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_triangles.diminfo[1].shape = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer, (PyObject*)__pyx_v_depth_buffer, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 49, __pyx_L1_error) - } - __pyx_pybuffernd_depth_buffer.diminfo[0].strides = __pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_depth_buffer.diminfo[0].shape = __pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_depth_buffer.diminfo[1].strides = __pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_depth_buffer.diminfo[1].shape = __pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_triangle_buffer.rcbuffer->pybuffer, (PyObject*)__pyx_v_triangle_buffer, &__Pyx_TypeInfo_int, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 49, __pyx_L1_error) - } - __pyx_pybuffernd_triangle_buffer.diminfo[0].strides = __pyx_pybuffernd_triangle_buffer.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_triangle_buffer.diminfo[0].shape = __pyx_pybuffernd_triangle_buffer.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_triangle_buffer.diminfo[1].strides = __pyx_pybuffernd_triangle_buffer.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_triangle_buffer.diminfo[1].shape = __pyx_pybuffernd_triangle_buffer.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_barycentric_weight.rcbuffer->pybuffer, (PyObject*)__pyx_v_barycentric_weight, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 49, __pyx_L1_error) - } - __pyx_pybuffernd_barycentric_weight.diminfo[0].strides = __pyx_pybuffernd_barycentric_weight.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_barycentric_weight.diminfo[0].shape = __pyx_pybuffernd_barycentric_weight.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_barycentric_weight.diminfo[1].strides = __pyx_pybuffernd_barycentric_weight.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_barycentric_weight.diminfo[1].shape = __pyx_pybuffernd_barycentric_weight.rcbuffer->pybuffer.shape[1]; - - /* "mesh_core_cython.pyx":58 - * int h, int w - * ): - * _rasterize_triangles_core( # <<<<<<<<<<<<<< - * np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), - * np.PyArray_DATA(depth_buffer), np.PyArray_DATA(triangle_buffer), np.PyArray_DATA(barycentric_weight), - */ - _rasterize_triangles_core(((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_vertices))), ((int *)PyArray_DATA(((PyArrayObject *)__pyx_v_triangles))), ((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_depth_buffer))), ((int *)PyArray_DATA(((PyArrayObject *)__pyx_v_triangle_buffer))), ((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_barycentric_weight))), __pyx_v_nver, __pyx_v_ntri, __pyx_v_h, __pyx_v_w); - - /* "mesh_core_cython.pyx":49 - * ntri) - * - * def rasterize_triangles_core( # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - - /* function exit code */ - __pyx_r = Py_None; __Pyx_INCREF(Py_None); - goto __pyx_L0; - __pyx_L1_error:; - { PyObject *__pyx_type, *__pyx_value, *__pyx_tb; - __Pyx_PyThreadState_declare - __Pyx_PyThreadState_assign - __Pyx_ErrFetch(&__pyx_type, &__pyx_value, &__pyx_tb); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_barycentric_weight.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_triangle_buffer.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_vertices.rcbuffer->pybuffer); - __Pyx_ErrRestore(__pyx_type, __pyx_value, __pyx_tb);} - __Pyx_AddTraceback("mesh_core_cython.rasterize_triangles_core", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = NULL; - goto __pyx_L2; - __pyx_L0:; - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_barycentric_weight.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_triangle_buffer.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_vertices.rcbuffer->pybuffer); - __pyx_L2:; - __Pyx_XGIVEREF(__pyx_r); - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "mesh_core_cython.pyx":64 - * h, w) - * - * def render_colors_core(np.ndarray[float, ndim=3, mode = "c"] image not None, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - -/* Python wrapper */ -static PyObject *__pyx_pw_16mesh_core_cython_5render_colors_core(PyObject *__pyx_self, PyObject *__pyx_args, PyObject *__pyx_kwds); /*proto*/ -static PyMethodDef __pyx_mdef_16mesh_core_cython_5render_colors_core = {"render_colors_core", (PyCFunction)__pyx_pw_16mesh_core_cython_5render_colors_core, METH_VARARGS|METH_KEYWORDS, 0}; -static PyObject *__pyx_pw_16mesh_core_cython_5render_colors_core(PyObject *__pyx_self, PyObject *__pyx_args, PyObject *__pyx_kwds) { - PyArrayObject *__pyx_v_image = 0; - PyArrayObject *__pyx_v_vertices = 0; - PyArrayObject *__pyx_v_triangles = 0; - PyArrayObject *__pyx_v_colors = 0; - PyArrayObject *__pyx_v_depth_buffer = 0; - int __pyx_v_nver; - int __pyx_v_ntri; - int __pyx_v_h; - int __pyx_v_w; - int __pyx_v_c; - PyObject *__pyx_r = 0; - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("render_colors_core (wrapper)", 0); - { - static PyObject **__pyx_pyargnames[] = {&__pyx_n_s_image,&__pyx_n_s_vertices,&__pyx_n_s_triangles,&__pyx_n_s_colors,&__pyx_n_s_depth_buffer,&__pyx_n_s_nver,&__pyx_n_s_ntri,&__pyx_n_s_h,&__pyx_n_s_w,&__pyx_n_s_c,0}; - PyObject* values[10] = {0,0,0,0,0,0,0,0,0,0}; - if (unlikely(__pyx_kwds)) { - Py_ssize_t kw_args; - const Py_ssize_t pos_args = PyTuple_GET_SIZE(__pyx_args); - switch (pos_args) { - case 10: values[9] = PyTuple_GET_ITEM(__pyx_args, 9); - CYTHON_FALLTHROUGH; - case 9: values[8] = PyTuple_GET_ITEM(__pyx_args, 8); - CYTHON_FALLTHROUGH; - case 8: values[7] = PyTuple_GET_ITEM(__pyx_args, 7); - CYTHON_FALLTHROUGH; - case 7: values[6] = PyTuple_GET_ITEM(__pyx_args, 6); - CYTHON_FALLTHROUGH; - case 6: values[5] = PyTuple_GET_ITEM(__pyx_args, 5); - CYTHON_FALLTHROUGH; - case 5: values[4] = PyTuple_GET_ITEM(__pyx_args, 4); - CYTHON_FALLTHROUGH; - case 4: values[3] = PyTuple_GET_ITEM(__pyx_args, 3); - CYTHON_FALLTHROUGH; - case 3: values[2] = PyTuple_GET_ITEM(__pyx_args, 2); - CYTHON_FALLTHROUGH; - case 2: values[1] = PyTuple_GET_ITEM(__pyx_args, 1); - CYTHON_FALLTHROUGH; - case 1: values[0] = PyTuple_GET_ITEM(__pyx_args, 0); - CYTHON_FALLTHROUGH; - case 0: break; - default: goto __pyx_L5_argtuple_error; - } - kw_args = PyDict_Size(__pyx_kwds); - switch (pos_args) { - case 0: - if (likely((values[0] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_image)) != 0)) kw_args--; - else goto __pyx_L5_argtuple_error; - CYTHON_FALLTHROUGH; - case 1: - if (likely((values[1] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_vertices)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_colors_core", 1, 10, 10, 1); __PYX_ERR(0, 64, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 2: - if (likely((values[2] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_triangles)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_colors_core", 1, 10, 10, 2); __PYX_ERR(0, 64, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 3: - if (likely((values[3] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_colors)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_colors_core", 1, 10, 10, 3); __PYX_ERR(0, 64, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 4: - if (likely((values[4] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_depth_buffer)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_colors_core", 1, 10, 10, 4); __PYX_ERR(0, 64, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 5: - if (likely((values[5] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_nver)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_colors_core", 1, 10, 10, 5); __PYX_ERR(0, 64, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 6: - if (likely((values[6] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_ntri)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_colors_core", 1, 10, 10, 6); __PYX_ERR(0, 64, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 7: - if (likely((values[7] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_h)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_colors_core", 1, 10, 10, 7); __PYX_ERR(0, 64, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 8: - if (likely((values[8] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_w)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_colors_core", 1, 10, 10, 8); __PYX_ERR(0, 64, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 9: - if (likely((values[9] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_c)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_colors_core", 1, 10, 10, 9); __PYX_ERR(0, 64, __pyx_L3_error) - } - } - if (unlikely(kw_args > 0)) { - if (unlikely(__Pyx_ParseOptionalKeywords(__pyx_kwds, __pyx_pyargnames, 0, values, pos_args, "render_colors_core") < 0)) __PYX_ERR(0, 64, __pyx_L3_error) - } - } else if (PyTuple_GET_SIZE(__pyx_args) != 10) { - goto __pyx_L5_argtuple_error; - } else { - values[0] = PyTuple_GET_ITEM(__pyx_args, 0); - values[1] = PyTuple_GET_ITEM(__pyx_args, 1); - values[2] = PyTuple_GET_ITEM(__pyx_args, 2); - values[3] = PyTuple_GET_ITEM(__pyx_args, 3); - values[4] = PyTuple_GET_ITEM(__pyx_args, 4); - values[5] = PyTuple_GET_ITEM(__pyx_args, 5); - values[6] = PyTuple_GET_ITEM(__pyx_args, 6); - values[7] = PyTuple_GET_ITEM(__pyx_args, 7); - values[8] = PyTuple_GET_ITEM(__pyx_args, 8); - values[9] = PyTuple_GET_ITEM(__pyx_args, 9); - } - __pyx_v_image = ((PyArrayObject *)values[0]); - __pyx_v_vertices = ((PyArrayObject *)values[1]); - __pyx_v_triangles = ((PyArrayObject *)values[2]); - __pyx_v_colors = ((PyArrayObject *)values[3]); - __pyx_v_depth_buffer = ((PyArrayObject *)values[4]); - __pyx_v_nver = __Pyx_PyInt_As_int(values[5]); if (unlikely((__pyx_v_nver == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 69, __pyx_L3_error) - __pyx_v_ntri = __Pyx_PyInt_As_int(values[6]); if (unlikely((__pyx_v_ntri == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 69, __pyx_L3_error) - __pyx_v_h = __Pyx_PyInt_As_int(values[7]); if (unlikely((__pyx_v_h == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 70, __pyx_L3_error) - __pyx_v_w = __Pyx_PyInt_As_int(values[8]); if (unlikely((__pyx_v_w == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 70, __pyx_L3_error) - __pyx_v_c = __Pyx_PyInt_As_int(values[9]); if (unlikely((__pyx_v_c == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 70, __pyx_L3_error) - } - goto __pyx_L4_argument_unpacking_done; - __pyx_L5_argtuple_error:; - __Pyx_RaiseArgtupleInvalid("render_colors_core", 1, 10, 10, PyTuple_GET_SIZE(__pyx_args)); __PYX_ERR(0, 64, __pyx_L3_error) - __pyx_L3_error:; - __Pyx_AddTraceback("mesh_core_cython.render_colors_core", __pyx_clineno, __pyx_lineno, __pyx_filename); - __Pyx_RefNannyFinishContext(); - return NULL; - __pyx_L4_argument_unpacking_done:; - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_image), __pyx_ptype_5numpy_ndarray, 0, "image", 0))) __PYX_ERR(0, 64, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_vertices), __pyx_ptype_5numpy_ndarray, 0, "vertices", 0))) __PYX_ERR(0, 65, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_triangles), __pyx_ptype_5numpy_ndarray, 0, "triangles", 0))) __PYX_ERR(0, 66, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_colors), __pyx_ptype_5numpy_ndarray, 0, "colors", 0))) __PYX_ERR(0, 67, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_depth_buffer), __pyx_ptype_5numpy_ndarray, 0, "depth_buffer", 0))) __PYX_ERR(0, 68, __pyx_L1_error) - __pyx_r = __pyx_pf_16mesh_core_cython_4render_colors_core(__pyx_self, __pyx_v_image, __pyx_v_vertices, __pyx_v_triangles, __pyx_v_colors, __pyx_v_depth_buffer, __pyx_v_nver, __pyx_v_ntri, __pyx_v_h, __pyx_v_w, __pyx_v_c); - - /* function exit code */ - goto __pyx_L0; - __pyx_L1_error:; - __pyx_r = NULL; - __pyx_L0:; - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -static PyObject *__pyx_pf_16mesh_core_cython_4render_colors_core(CYTHON_UNUSED PyObject *__pyx_self, PyArrayObject *__pyx_v_image, PyArrayObject *__pyx_v_vertices, PyArrayObject *__pyx_v_triangles, PyArrayObject *__pyx_v_colors, PyArrayObject *__pyx_v_depth_buffer, int __pyx_v_nver, int __pyx_v_ntri, int __pyx_v_h, int __pyx_v_w, int __pyx_v_c) { - __Pyx_LocalBuf_ND __pyx_pybuffernd_colors; - __Pyx_Buffer __pyx_pybuffer_colors; - __Pyx_LocalBuf_ND __pyx_pybuffernd_depth_buffer; - __Pyx_Buffer __pyx_pybuffer_depth_buffer; - __Pyx_LocalBuf_ND __pyx_pybuffernd_image; - __Pyx_Buffer __pyx_pybuffer_image; - __Pyx_LocalBuf_ND __pyx_pybuffernd_triangles; - __Pyx_Buffer __pyx_pybuffer_triangles; - __Pyx_LocalBuf_ND __pyx_pybuffernd_vertices; - __Pyx_Buffer __pyx_pybuffer_vertices; - PyObject *__pyx_r = NULL; - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("render_colors_core", 0); - __pyx_pybuffer_image.pybuffer.buf = NULL; - __pyx_pybuffer_image.refcount = 0; - __pyx_pybuffernd_image.data = NULL; - __pyx_pybuffernd_image.rcbuffer = &__pyx_pybuffer_image; - __pyx_pybuffer_vertices.pybuffer.buf = NULL; - __pyx_pybuffer_vertices.refcount = 0; - __pyx_pybuffernd_vertices.data = NULL; - __pyx_pybuffernd_vertices.rcbuffer = &__pyx_pybuffer_vertices; - __pyx_pybuffer_triangles.pybuffer.buf = NULL; - __pyx_pybuffer_triangles.refcount = 0; - __pyx_pybuffernd_triangles.data = NULL; - __pyx_pybuffernd_triangles.rcbuffer = &__pyx_pybuffer_triangles; - __pyx_pybuffer_colors.pybuffer.buf = NULL; - __pyx_pybuffer_colors.refcount = 0; - __pyx_pybuffernd_colors.data = NULL; - __pyx_pybuffernd_colors.rcbuffer = &__pyx_pybuffer_colors; - __pyx_pybuffer_depth_buffer.pybuffer.buf = NULL; - __pyx_pybuffer_depth_buffer.refcount = 0; - __pyx_pybuffernd_depth_buffer.data = NULL; - __pyx_pybuffernd_depth_buffer.rcbuffer = &__pyx_pybuffer_depth_buffer; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_image.rcbuffer->pybuffer, (PyObject*)__pyx_v_image, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 3, 0, __pyx_stack) == -1)) __PYX_ERR(0, 64, __pyx_L1_error) - } - __pyx_pybuffernd_image.diminfo[0].strides = __pyx_pybuffernd_image.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_image.diminfo[0].shape = __pyx_pybuffernd_image.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_image.diminfo[1].strides = __pyx_pybuffernd_image.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_image.diminfo[1].shape = __pyx_pybuffernd_image.rcbuffer->pybuffer.shape[1]; __pyx_pybuffernd_image.diminfo[2].strides = __pyx_pybuffernd_image.rcbuffer->pybuffer.strides[2]; __pyx_pybuffernd_image.diminfo[2].shape = __pyx_pybuffernd_image.rcbuffer->pybuffer.shape[2]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_vertices.rcbuffer->pybuffer, (PyObject*)__pyx_v_vertices, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 64, __pyx_L1_error) - } - __pyx_pybuffernd_vertices.diminfo[0].strides = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_vertices.diminfo[0].shape = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_vertices.diminfo[1].strides = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_vertices.diminfo[1].shape = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer, (PyObject*)__pyx_v_triangles, &__Pyx_TypeInfo_int, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 64, __pyx_L1_error) - } - __pyx_pybuffernd_triangles.diminfo[0].strides = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_triangles.diminfo[0].shape = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_triangles.diminfo[1].strides = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_triangles.diminfo[1].shape = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_colors.rcbuffer->pybuffer, (PyObject*)__pyx_v_colors, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 64, __pyx_L1_error) - } - __pyx_pybuffernd_colors.diminfo[0].strides = __pyx_pybuffernd_colors.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_colors.diminfo[0].shape = __pyx_pybuffernd_colors.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_colors.diminfo[1].strides = __pyx_pybuffernd_colors.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_colors.diminfo[1].shape = __pyx_pybuffernd_colors.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer, (PyObject*)__pyx_v_depth_buffer, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 64, __pyx_L1_error) - } - __pyx_pybuffernd_depth_buffer.diminfo[0].strides = __pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_depth_buffer.diminfo[0].shape = __pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_depth_buffer.diminfo[1].strides = __pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_depth_buffer.diminfo[1].shape = __pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer.shape[1]; - - /* "mesh_core_cython.pyx":72 - * int h, int w, int c - * ): - * _render_colors_core( # <<<<<<<<<<<<<< - * np.PyArray_DATA(image), np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), - * np.PyArray_DATA(colors), - */ - _render_colors_core(((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_image))), ((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_vertices))), ((int *)PyArray_DATA(((PyArrayObject *)__pyx_v_triangles))), ((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_colors))), ((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_depth_buffer))), __pyx_v_nver, __pyx_v_ntri, __pyx_v_h, __pyx_v_w, __pyx_v_c); - - /* "mesh_core_cython.pyx":64 - * h, w) - * - * def render_colors_core(np.ndarray[float, ndim=3, mode = "c"] image not None, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - - /* function exit code */ - __pyx_r = Py_None; __Pyx_INCREF(Py_None); - goto __pyx_L0; - __pyx_L1_error:; - { PyObject *__pyx_type, *__pyx_value, *__pyx_tb; - __Pyx_PyThreadState_declare - __Pyx_PyThreadState_assign - __Pyx_ErrFetch(&__pyx_type, &__pyx_value, &__pyx_tb); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_colors.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_image.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_vertices.rcbuffer->pybuffer); - __Pyx_ErrRestore(__pyx_type, __pyx_value, __pyx_tb);} - __Pyx_AddTraceback("mesh_core_cython.render_colors_core", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = NULL; - goto __pyx_L2; - __pyx_L0:; - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_colors.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_image.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_vertices.rcbuffer->pybuffer); - __pyx_L2:; - __Pyx_XGIVEREF(__pyx_r); - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "mesh_core_cython.pyx":79 - * h, w, c) - * - * def render_texture_core(np.ndarray[float, ndim=3, mode = "c"] image not None, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - -/* Python wrapper */ -static PyObject *__pyx_pw_16mesh_core_cython_7render_texture_core(PyObject *__pyx_self, PyObject *__pyx_args, PyObject *__pyx_kwds); /*proto*/ -static PyMethodDef __pyx_mdef_16mesh_core_cython_7render_texture_core = {"render_texture_core", (PyCFunction)__pyx_pw_16mesh_core_cython_7render_texture_core, METH_VARARGS|METH_KEYWORDS, 0}; -static PyObject *__pyx_pw_16mesh_core_cython_7render_texture_core(PyObject *__pyx_self, PyObject *__pyx_args, PyObject *__pyx_kwds) { - PyArrayObject *__pyx_v_image = 0; - PyArrayObject *__pyx_v_vertices = 0; - PyArrayObject *__pyx_v_triangles = 0; - PyArrayObject *__pyx_v_texture = 0; - PyArrayObject *__pyx_v_tex_coords = 0; - PyArrayObject *__pyx_v_tex_triangles = 0; - PyArrayObject *__pyx_v_depth_buffer = 0; - int __pyx_v_nver; - int __pyx_v_tex_nver; - int __pyx_v_ntri; - int __pyx_v_h; - int __pyx_v_w; - int __pyx_v_c; - int __pyx_v_tex_h; - int __pyx_v_tex_w; - int __pyx_v_tex_c; - int __pyx_v_mapping_type; - PyObject *__pyx_r = 0; - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("render_texture_core (wrapper)", 0); - { - static PyObject **__pyx_pyargnames[] = {&__pyx_n_s_image,&__pyx_n_s_vertices,&__pyx_n_s_triangles,&__pyx_n_s_texture,&__pyx_n_s_tex_coords,&__pyx_n_s_tex_triangles,&__pyx_n_s_depth_buffer,&__pyx_n_s_nver,&__pyx_n_s_tex_nver,&__pyx_n_s_ntri,&__pyx_n_s_h,&__pyx_n_s_w,&__pyx_n_s_c,&__pyx_n_s_tex_h,&__pyx_n_s_tex_w,&__pyx_n_s_tex_c,&__pyx_n_s_mapping_type,0}; - PyObject* values[17] = {0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0}; - if (unlikely(__pyx_kwds)) { - Py_ssize_t kw_args; - const Py_ssize_t pos_args = PyTuple_GET_SIZE(__pyx_args); - switch (pos_args) { - case 17: values[16] = PyTuple_GET_ITEM(__pyx_args, 16); - CYTHON_FALLTHROUGH; - case 16: values[15] = PyTuple_GET_ITEM(__pyx_args, 15); - CYTHON_FALLTHROUGH; - case 15: values[14] = PyTuple_GET_ITEM(__pyx_args, 14); - CYTHON_FALLTHROUGH; - case 14: values[13] = PyTuple_GET_ITEM(__pyx_args, 13); - CYTHON_FALLTHROUGH; - case 13: values[12] = PyTuple_GET_ITEM(__pyx_args, 12); - CYTHON_FALLTHROUGH; - case 12: values[11] = PyTuple_GET_ITEM(__pyx_args, 11); - CYTHON_FALLTHROUGH; - case 11: values[10] = PyTuple_GET_ITEM(__pyx_args, 10); - CYTHON_FALLTHROUGH; - case 10: values[9] = PyTuple_GET_ITEM(__pyx_args, 9); - CYTHON_FALLTHROUGH; - case 9: values[8] = PyTuple_GET_ITEM(__pyx_args, 8); - CYTHON_FALLTHROUGH; - case 8: values[7] = PyTuple_GET_ITEM(__pyx_args, 7); - CYTHON_FALLTHROUGH; - case 7: values[6] = PyTuple_GET_ITEM(__pyx_args, 6); - CYTHON_FALLTHROUGH; - case 6: values[5] = PyTuple_GET_ITEM(__pyx_args, 5); - CYTHON_FALLTHROUGH; - case 5: values[4] = PyTuple_GET_ITEM(__pyx_args, 4); - CYTHON_FALLTHROUGH; - case 4: values[3] = PyTuple_GET_ITEM(__pyx_args, 3); - CYTHON_FALLTHROUGH; - case 3: values[2] = PyTuple_GET_ITEM(__pyx_args, 2); - CYTHON_FALLTHROUGH; - case 2: values[1] = PyTuple_GET_ITEM(__pyx_args, 1); - CYTHON_FALLTHROUGH; - case 1: values[0] = PyTuple_GET_ITEM(__pyx_args, 0); - CYTHON_FALLTHROUGH; - case 0: break; - default: goto __pyx_L5_argtuple_error; - } - kw_args = PyDict_Size(__pyx_kwds); - switch (pos_args) { - case 0: - if (likely((values[0] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_image)) != 0)) kw_args--; - else goto __pyx_L5_argtuple_error; - CYTHON_FALLTHROUGH; - case 1: - if (likely((values[1] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_vertices)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 1); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 2: - if (likely((values[2] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_triangles)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 2); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 3: - if (likely((values[3] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_texture)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 3); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 4: - if (likely((values[4] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_tex_coords)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 4); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 5: - if (likely((values[5] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_tex_triangles)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 5); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 6: - if (likely((values[6] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_depth_buffer)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 6); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 7: - if (likely((values[7] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_nver)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 7); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 8: - if (likely((values[8] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_tex_nver)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 8); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 9: - if (likely((values[9] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_ntri)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 9); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 10: - if (likely((values[10] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_h)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 10); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 11: - if (likely((values[11] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_w)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 11); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 12: - if (likely((values[12] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_c)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 12); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 13: - if (likely((values[13] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_tex_h)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 13); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 14: - if (likely((values[14] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_tex_w)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 14); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 15: - if (likely((values[15] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_tex_c)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 15); __PYX_ERR(0, 79, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 16: - if (likely((values[16] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_mapping_type)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, 16); __PYX_ERR(0, 79, __pyx_L3_error) - } - } - if (unlikely(kw_args > 0)) { - if (unlikely(__Pyx_ParseOptionalKeywords(__pyx_kwds, __pyx_pyargnames, 0, values, pos_args, "render_texture_core") < 0)) __PYX_ERR(0, 79, __pyx_L3_error) - } - } else if (PyTuple_GET_SIZE(__pyx_args) != 17) { - goto __pyx_L5_argtuple_error; - } else { - values[0] = PyTuple_GET_ITEM(__pyx_args, 0); - values[1] = PyTuple_GET_ITEM(__pyx_args, 1); - values[2] = PyTuple_GET_ITEM(__pyx_args, 2); - values[3] = PyTuple_GET_ITEM(__pyx_args, 3); - values[4] = PyTuple_GET_ITEM(__pyx_args, 4); - values[5] = PyTuple_GET_ITEM(__pyx_args, 5); - values[6] = PyTuple_GET_ITEM(__pyx_args, 6); - values[7] = PyTuple_GET_ITEM(__pyx_args, 7); - values[8] = PyTuple_GET_ITEM(__pyx_args, 8); - values[9] = PyTuple_GET_ITEM(__pyx_args, 9); - values[10] = PyTuple_GET_ITEM(__pyx_args, 10); - values[11] = PyTuple_GET_ITEM(__pyx_args, 11); - values[12] = PyTuple_GET_ITEM(__pyx_args, 12); - values[13] = PyTuple_GET_ITEM(__pyx_args, 13); - values[14] = PyTuple_GET_ITEM(__pyx_args, 14); - values[15] = PyTuple_GET_ITEM(__pyx_args, 15); - values[16] = PyTuple_GET_ITEM(__pyx_args, 16); - } - __pyx_v_image = ((PyArrayObject *)values[0]); - __pyx_v_vertices = ((PyArrayObject *)values[1]); - __pyx_v_triangles = ((PyArrayObject *)values[2]); - __pyx_v_texture = ((PyArrayObject *)values[3]); - __pyx_v_tex_coords = ((PyArrayObject *)values[4]); - __pyx_v_tex_triangles = ((PyArrayObject *)values[5]); - __pyx_v_depth_buffer = ((PyArrayObject *)values[6]); - __pyx_v_nver = __Pyx_PyInt_As_int(values[7]); if (unlikely((__pyx_v_nver == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 86, __pyx_L3_error) - __pyx_v_tex_nver = __Pyx_PyInt_As_int(values[8]); if (unlikely((__pyx_v_tex_nver == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 86, __pyx_L3_error) - __pyx_v_ntri = __Pyx_PyInt_As_int(values[9]); if (unlikely((__pyx_v_ntri == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 86, __pyx_L3_error) - __pyx_v_h = __Pyx_PyInt_As_int(values[10]); if (unlikely((__pyx_v_h == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 87, __pyx_L3_error) - __pyx_v_w = __Pyx_PyInt_As_int(values[11]); if (unlikely((__pyx_v_w == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 87, __pyx_L3_error) - __pyx_v_c = __Pyx_PyInt_As_int(values[12]); if (unlikely((__pyx_v_c == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 87, __pyx_L3_error) - __pyx_v_tex_h = __Pyx_PyInt_As_int(values[13]); if (unlikely((__pyx_v_tex_h == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 88, __pyx_L3_error) - __pyx_v_tex_w = __Pyx_PyInt_As_int(values[14]); if (unlikely((__pyx_v_tex_w == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 88, __pyx_L3_error) - __pyx_v_tex_c = __Pyx_PyInt_As_int(values[15]); if (unlikely((__pyx_v_tex_c == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 88, __pyx_L3_error) - __pyx_v_mapping_type = __Pyx_PyInt_As_int(values[16]); if (unlikely((__pyx_v_mapping_type == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 89, __pyx_L3_error) - } - goto __pyx_L4_argument_unpacking_done; - __pyx_L5_argtuple_error:; - __Pyx_RaiseArgtupleInvalid("render_texture_core", 1, 17, 17, PyTuple_GET_SIZE(__pyx_args)); __PYX_ERR(0, 79, __pyx_L3_error) - __pyx_L3_error:; - __Pyx_AddTraceback("mesh_core_cython.render_texture_core", __pyx_clineno, __pyx_lineno, __pyx_filename); - __Pyx_RefNannyFinishContext(); - return NULL; - __pyx_L4_argument_unpacking_done:; - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_image), __pyx_ptype_5numpy_ndarray, 0, "image", 0))) __PYX_ERR(0, 79, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_vertices), __pyx_ptype_5numpy_ndarray, 0, "vertices", 0))) __PYX_ERR(0, 80, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_triangles), __pyx_ptype_5numpy_ndarray, 0, "triangles", 0))) __PYX_ERR(0, 81, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_texture), __pyx_ptype_5numpy_ndarray, 0, "texture", 0))) __PYX_ERR(0, 82, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_tex_coords), __pyx_ptype_5numpy_ndarray, 0, "tex_coords", 0))) __PYX_ERR(0, 83, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_tex_triangles), __pyx_ptype_5numpy_ndarray, 0, "tex_triangles", 0))) __PYX_ERR(0, 84, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_depth_buffer), __pyx_ptype_5numpy_ndarray, 0, "depth_buffer", 0))) __PYX_ERR(0, 85, __pyx_L1_error) - __pyx_r = __pyx_pf_16mesh_core_cython_6render_texture_core(__pyx_self, __pyx_v_image, __pyx_v_vertices, __pyx_v_triangles, __pyx_v_texture, __pyx_v_tex_coords, __pyx_v_tex_triangles, __pyx_v_depth_buffer, __pyx_v_nver, __pyx_v_tex_nver, __pyx_v_ntri, __pyx_v_h, __pyx_v_w, __pyx_v_c, __pyx_v_tex_h, __pyx_v_tex_w, __pyx_v_tex_c, __pyx_v_mapping_type); - - /* function exit code */ - goto __pyx_L0; - __pyx_L1_error:; - __pyx_r = NULL; - __pyx_L0:; - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -static PyObject *__pyx_pf_16mesh_core_cython_6render_texture_core(CYTHON_UNUSED PyObject *__pyx_self, PyArrayObject *__pyx_v_image, PyArrayObject *__pyx_v_vertices, PyArrayObject *__pyx_v_triangles, PyArrayObject *__pyx_v_texture, PyArrayObject *__pyx_v_tex_coords, PyArrayObject *__pyx_v_tex_triangles, PyArrayObject *__pyx_v_depth_buffer, int __pyx_v_nver, int __pyx_v_tex_nver, int __pyx_v_ntri, int __pyx_v_h, int __pyx_v_w, int __pyx_v_c, int __pyx_v_tex_h, int __pyx_v_tex_w, int __pyx_v_tex_c, int __pyx_v_mapping_type) { - __Pyx_LocalBuf_ND __pyx_pybuffernd_depth_buffer; - __Pyx_Buffer __pyx_pybuffer_depth_buffer; - __Pyx_LocalBuf_ND __pyx_pybuffernd_image; - __Pyx_Buffer __pyx_pybuffer_image; - __Pyx_LocalBuf_ND __pyx_pybuffernd_tex_coords; - __Pyx_Buffer __pyx_pybuffer_tex_coords; - __Pyx_LocalBuf_ND __pyx_pybuffernd_tex_triangles; - __Pyx_Buffer __pyx_pybuffer_tex_triangles; - __Pyx_LocalBuf_ND __pyx_pybuffernd_texture; - __Pyx_Buffer __pyx_pybuffer_texture; - __Pyx_LocalBuf_ND __pyx_pybuffernd_triangles; - __Pyx_Buffer __pyx_pybuffer_triangles; - __Pyx_LocalBuf_ND __pyx_pybuffernd_vertices; - __Pyx_Buffer __pyx_pybuffer_vertices; - PyObject *__pyx_r = NULL; - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("render_texture_core", 0); - __pyx_pybuffer_image.pybuffer.buf = NULL; - __pyx_pybuffer_image.refcount = 0; - __pyx_pybuffernd_image.data = NULL; - __pyx_pybuffernd_image.rcbuffer = &__pyx_pybuffer_image; - __pyx_pybuffer_vertices.pybuffer.buf = NULL; - __pyx_pybuffer_vertices.refcount = 0; - __pyx_pybuffernd_vertices.data = NULL; - __pyx_pybuffernd_vertices.rcbuffer = &__pyx_pybuffer_vertices; - __pyx_pybuffer_triangles.pybuffer.buf = NULL; - __pyx_pybuffer_triangles.refcount = 0; - __pyx_pybuffernd_triangles.data = NULL; - __pyx_pybuffernd_triangles.rcbuffer = &__pyx_pybuffer_triangles; - __pyx_pybuffer_texture.pybuffer.buf = NULL; - __pyx_pybuffer_texture.refcount = 0; - __pyx_pybuffernd_texture.data = NULL; - __pyx_pybuffernd_texture.rcbuffer = &__pyx_pybuffer_texture; - __pyx_pybuffer_tex_coords.pybuffer.buf = NULL; - __pyx_pybuffer_tex_coords.refcount = 0; - __pyx_pybuffernd_tex_coords.data = NULL; - __pyx_pybuffernd_tex_coords.rcbuffer = &__pyx_pybuffer_tex_coords; - __pyx_pybuffer_tex_triangles.pybuffer.buf = NULL; - __pyx_pybuffer_tex_triangles.refcount = 0; - __pyx_pybuffernd_tex_triangles.data = NULL; - __pyx_pybuffernd_tex_triangles.rcbuffer = &__pyx_pybuffer_tex_triangles; - __pyx_pybuffer_depth_buffer.pybuffer.buf = NULL; - __pyx_pybuffer_depth_buffer.refcount = 0; - __pyx_pybuffernd_depth_buffer.data = NULL; - __pyx_pybuffernd_depth_buffer.rcbuffer = &__pyx_pybuffer_depth_buffer; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_image.rcbuffer->pybuffer, (PyObject*)__pyx_v_image, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 3, 0, __pyx_stack) == -1)) __PYX_ERR(0, 79, __pyx_L1_error) - } - __pyx_pybuffernd_image.diminfo[0].strides = __pyx_pybuffernd_image.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_image.diminfo[0].shape = __pyx_pybuffernd_image.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_image.diminfo[1].strides = __pyx_pybuffernd_image.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_image.diminfo[1].shape = __pyx_pybuffernd_image.rcbuffer->pybuffer.shape[1]; __pyx_pybuffernd_image.diminfo[2].strides = __pyx_pybuffernd_image.rcbuffer->pybuffer.strides[2]; __pyx_pybuffernd_image.diminfo[2].shape = __pyx_pybuffernd_image.rcbuffer->pybuffer.shape[2]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_vertices.rcbuffer->pybuffer, (PyObject*)__pyx_v_vertices, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 79, __pyx_L1_error) - } - __pyx_pybuffernd_vertices.diminfo[0].strides = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_vertices.diminfo[0].shape = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_vertices.diminfo[1].strides = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_vertices.diminfo[1].shape = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer, (PyObject*)__pyx_v_triangles, &__Pyx_TypeInfo_int, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 79, __pyx_L1_error) - } - __pyx_pybuffernd_triangles.diminfo[0].strides = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_triangles.diminfo[0].shape = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_triangles.diminfo[1].strides = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_triangles.diminfo[1].shape = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_texture.rcbuffer->pybuffer, (PyObject*)__pyx_v_texture, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 3, 0, __pyx_stack) == -1)) __PYX_ERR(0, 79, __pyx_L1_error) - } - __pyx_pybuffernd_texture.diminfo[0].strides = __pyx_pybuffernd_texture.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_texture.diminfo[0].shape = __pyx_pybuffernd_texture.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_texture.diminfo[1].strides = __pyx_pybuffernd_texture.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_texture.diminfo[1].shape = __pyx_pybuffernd_texture.rcbuffer->pybuffer.shape[1]; __pyx_pybuffernd_texture.diminfo[2].strides = __pyx_pybuffernd_texture.rcbuffer->pybuffer.strides[2]; __pyx_pybuffernd_texture.diminfo[2].shape = __pyx_pybuffernd_texture.rcbuffer->pybuffer.shape[2]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_tex_coords.rcbuffer->pybuffer, (PyObject*)__pyx_v_tex_coords, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 79, __pyx_L1_error) - } - __pyx_pybuffernd_tex_coords.diminfo[0].strides = __pyx_pybuffernd_tex_coords.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_tex_coords.diminfo[0].shape = __pyx_pybuffernd_tex_coords.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_tex_coords.diminfo[1].strides = __pyx_pybuffernd_tex_coords.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_tex_coords.diminfo[1].shape = __pyx_pybuffernd_tex_coords.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_tex_triangles.rcbuffer->pybuffer, (PyObject*)__pyx_v_tex_triangles, &__Pyx_TypeInfo_int, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 79, __pyx_L1_error) - } - __pyx_pybuffernd_tex_triangles.diminfo[0].strides = __pyx_pybuffernd_tex_triangles.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_tex_triangles.diminfo[0].shape = __pyx_pybuffernd_tex_triangles.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_tex_triangles.diminfo[1].strides = __pyx_pybuffernd_tex_triangles.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_tex_triangles.diminfo[1].shape = __pyx_pybuffernd_tex_triangles.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer, (PyObject*)__pyx_v_depth_buffer, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 79, __pyx_L1_error) - } - __pyx_pybuffernd_depth_buffer.diminfo[0].strides = __pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_depth_buffer.diminfo[0].shape = __pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_depth_buffer.diminfo[1].strides = __pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_depth_buffer.diminfo[1].shape = __pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer.shape[1]; - - /* "mesh_core_cython.pyx":91 - * int mapping_type - * ): - * _render_texture_core( # <<<<<<<<<<<<<< - * np.PyArray_DATA(image), np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), - * np.PyArray_DATA(texture), np.PyArray_DATA(tex_coords), np.PyArray_DATA(tex_triangles), - */ - _render_texture_core(((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_image))), ((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_vertices))), ((int *)PyArray_DATA(((PyArrayObject *)__pyx_v_triangles))), ((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_texture))), ((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_tex_coords))), ((int *)PyArray_DATA(((PyArrayObject *)__pyx_v_tex_triangles))), ((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_depth_buffer))), __pyx_v_nver, __pyx_v_tex_nver, __pyx_v_ntri, __pyx_v_h, __pyx_v_w, __pyx_v_c, __pyx_v_tex_h, __pyx_v_tex_w, __pyx_v_tex_c, __pyx_v_mapping_type); - - /* "mesh_core_cython.pyx":79 - * h, w, c) - * - * def render_texture_core(np.ndarray[float, ndim=3, mode = "c"] image not None, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - - /* function exit code */ - __pyx_r = Py_None; __Pyx_INCREF(Py_None); - goto __pyx_L0; - __pyx_L1_error:; - { PyObject *__pyx_type, *__pyx_value, *__pyx_tb; - __Pyx_PyThreadState_declare - __Pyx_PyThreadState_assign - __Pyx_ErrFetch(&__pyx_type, &__pyx_value, &__pyx_tb); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_image.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_tex_coords.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_tex_triangles.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_texture.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_vertices.rcbuffer->pybuffer); - __Pyx_ErrRestore(__pyx_type, __pyx_value, __pyx_tb);} - __Pyx_AddTraceback("mesh_core_cython.render_texture_core", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = NULL; - goto __pyx_L2; - __pyx_L0:; - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_depth_buffer.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_image.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_tex_coords.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_tex_triangles.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_texture.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_vertices.rcbuffer->pybuffer); - __pyx_L2:; - __Pyx_XGIVEREF(__pyx_r); - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "mesh_core_cython.pyx":100 - * mapping_type) - * - * def write_obj_with_colors_texture_core(string filename, string mtl_name, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - -/* Python wrapper */ -static PyObject *__pyx_pw_16mesh_core_cython_9write_obj_with_colors_texture_core(PyObject *__pyx_self, PyObject *__pyx_args, PyObject *__pyx_kwds); /*proto*/ -static PyMethodDef __pyx_mdef_16mesh_core_cython_9write_obj_with_colors_texture_core = {"write_obj_with_colors_texture_core", (PyCFunction)__pyx_pw_16mesh_core_cython_9write_obj_with_colors_texture_core, METH_VARARGS|METH_KEYWORDS, 0}; -static PyObject *__pyx_pw_16mesh_core_cython_9write_obj_with_colors_texture_core(PyObject *__pyx_self, PyObject *__pyx_args, PyObject *__pyx_kwds) { - std::string __pyx_v_filename; - std::string __pyx_v_mtl_name; - PyArrayObject *__pyx_v_vertices = 0; - PyArrayObject *__pyx_v_triangles = 0; - PyArrayObject *__pyx_v_colors = 0; - PyArrayObject *__pyx_v_uv_coords = 0; - int __pyx_v_nver; - int __pyx_v_ntri; - int __pyx_v_ntexver; - PyObject *__pyx_r = 0; - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("write_obj_with_colors_texture_core (wrapper)", 0); - { - static PyObject **__pyx_pyargnames[] = {&__pyx_n_s_filename,&__pyx_n_s_mtl_name,&__pyx_n_s_vertices,&__pyx_n_s_triangles,&__pyx_n_s_colors,&__pyx_n_s_uv_coords,&__pyx_n_s_nver,&__pyx_n_s_ntri,&__pyx_n_s_ntexver,0}; - PyObject* values[9] = {0,0,0,0,0,0,0,0,0}; - if (unlikely(__pyx_kwds)) { - Py_ssize_t kw_args; - const Py_ssize_t pos_args = PyTuple_GET_SIZE(__pyx_args); - switch (pos_args) { - case 9: values[8] = PyTuple_GET_ITEM(__pyx_args, 8); - CYTHON_FALLTHROUGH; - case 8: values[7] = PyTuple_GET_ITEM(__pyx_args, 7); - CYTHON_FALLTHROUGH; - case 7: values[6] = PyTuple_GET_ITEM(__pyx_args, 6); - CYTHON_FALLTHROUGH; - case 6: values[5] = PyTuple_GET_ITEM(__pyx_args, 5); - CYTHON_FALLTHROUGH; - case 5: values[4] = PyTuple_GET_ITEM(__pyx_args, 4); - CYTHON_FALLTHROUGH; - case 4: values[3] = PyTuple_GET_ITEM(__pyx_args, 3); - CYTHON_FALLTHROUGH; - case 3: values[2] = PyTuple_GET_ITEM(__pyx_args, 2); - CYTHON_FALLTHROUGH; - case 2: values[1] = PyTuple_GET_ITEM(__pyx_args, 1); - CYTHON_FALLTHROUGH; - case 1: values[0] = PyTuple_GET_ITEM(__pyx_args, 0); - CYTHON_FALLTHROUGH; - case 0: break; - default: goto __pyx_L5_argtuple_error; - } - kw_args = PyDict_Size(__pyx_kwds); - switch (pos_args) { - case 0: - if (likely((values[0] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_filename)) != 0)) kw_args--; - else goto __pyx_L5_argtuple_error; - CYTHON_FALLTHROUGH; - case 1: - if (likely((values[1] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_mtl_name)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("write_obj_with_colors_texture_core", 1, 9, 9, 1); __PYX_ERR(0, 100, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 2: - if (likely((values[2] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_vertices)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("write_obj_with_colors_texture_core", 1, 9, 9, 2); __PYX_ERR(0, 100, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 3: - if (likely((values[3] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_triangles)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("write_obj_with_colors_texture_core", 1, 9, 9, 3); __PYX_ERR(0, 100, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 4: - if (likely((values[4] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_colors)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("write_obj_with_colors_texture_core", 1, 9, 9, 4); __PYX_ERR(0, 100, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 5: - if (likely((values[5] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_uv_coords)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("write_obj_with_colors_texture_core", 1, 9, 9, 5); __PYX_ERR(0, 100, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 6: - if (likely((values[6] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_nver)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("write_obj_with_colors_texture_core", 1, 9, 9, 6); __PYX_ERR(0, 100, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 7: - if (likely((values[7] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_ntri)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("write_obj_with_colors_texture_core", 1, 9, 9, 7); __PYX_ERR(0, 100, __pyx_L3_error) - } - CYTHON_FALLTHROUGH; - case 8: - if (likely((values[8] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_ntexver)) != 0)) kw_args--; - else { - __Pyx_RaiseArgtupleInvalid("write_obj_with_colors_texture_core", 1, 9, 9, 8); __PYX_ERR(0, 100, __pyx_L3_error) - } - } - if (unlikely(kw_args > 0)) { - if (unlikely(__Pyx_ParseOptionalKeywords(__pyx_kwds, __pyx_pyargnames, 0, values, pos_args, "write_obj_with_colors_texture_core") < 0)) __PYX_ERR(0, 100, __pyx_L3_error) - } - } else if (PyTuple_GET_SIZE(__pyx_args) != 9) { - goto __pyx_L5_argtuple_error; - } else { - values[0] = PyTuple_GET_ITEM(__pyx_args, 0); - values[1] = PyTuple_GET_ITEM(__pyx_args, 1); - values[2] = PyTuple_GET_ITEM(__pyx_args, 2); - values[3] = PyTuple_GET_ITEM(__pyx_args, 3); - values[4] = PyTuple_GET_ITEM(__pyx_args, 4); - values[5] = PyTuple_GET_ITEM(__pyx_args, 5); - values[6] = PyTuple_GET_ITEM(__pyx_args, 6); - values[7] = PyTuple_GET_ITEM(__pyx_args, 7); - values[8] = PyTuple_GET_ITEM(__pyx_args, 8); - } - __pyx_v_filename = __pyx_convert_string_from_py_std__in_string(values[0]); if (unlikely(PyErr_Occurred())) __PYX_ERR(0, 100, __pyx_L3_error) - __pyx_v_mtl_name = __pyx_convert_string_from_py_std__in_string(values[1]); if (unlikely(PyErr_Occurred())) __PYX_ERR(0, 100, __pyx_L3_error) - __pyx_v_vertices = ((PyArrayObject *)values[2]); - __pyx_v_triangles = ((PyArrayObject *)values[3]); - __pyx_v_colors = ((PyArrayObject *)values[4]); - __pyx_v_uv_coords = ((PyArrayObject *)values[5]); - __pyx_v_nver = __Pyx_PyInt_As_int(values[6]); if (unlikely((__pyx_v_nver == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 105, __pyx_L3_error) - __pyx_v_ntri = __Pyx_PyInt_As_int(values[7]); if (unlikely((__pyx_v_ntri == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 105, __pyx_L3_error) - __pyx_v_ntexver = __Pyx_PyInt_As_int(values[8]); if (unlikely((__pyx_v_ntexver == (int)-1) && PyErr_Occurred())) __PYX_ERR(0, 105, __pyx_L3_error) - } - goto __pyx_L4_argument_unpacking_done; - __pyx_L5_argtuple_error:; - __Pyx_RaiseArgtupleInvalid("write_obj_with_colors_texture_core", 1, 9, 9, PyTuple_GET_SIZE(__pyx_args)); __PYX_ERR(0, 100, __pyx_L3_error) - __pyx_L3_error:; - __Pyx_AddTraceback("mesh_core_cython.write_obj_with_colors_texture_core", __pyx_clineno, __pyx_lineno, __pyx_filename); - __Pyx_RefNannyFinishContext(); - return NULL; - __pyx_L4_argument_unpacking_done:; - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_vertices), __pyx_ptype_5numpy_ndarray, 0, "vertices", 0))) __PYX_ERR(0, 101, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_triangles), __pyx_ptype_5numpy_ndarray, 0, "triangles", 0))) __PYX_ERR(0, 102, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_colors), __pyx_ptype_5numpy_ndarray, 0, "colors", 0))) __PYX_ERR(0, 103, __pyx_L1_error) - if (unlikely(!__Pyx_ArgTypeTest(((PyObject *)__pyx_v_uv_coords), __pyx_ptype_5numpy_ndarray, 0, "uv_coords", 0))) __PYX_ERR(0, 104, __pyx_L1_error) - __pyx_r = __pyx_pf_16mesh_core_cython_8write_obj_with_colors_texture_core(__pyx_self, __pyx_v_filename, __pyx_v_mtl_name, __pyx_v_vertices, __pyx_v_triangles, __pyx_v_colors, __pyx_v_uv_coords, __pyx_v_nver, __pyx_v_ntri, __pyx_v_ntexver); - - /* function exit code */ - goto __pyx_L0; - __pyx_L1_error:; - __pyx_r = NULL; - __pyx_L0:; - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -static PyObject *__pyx_pf_16mesh_core_cython_8write_obj_with_colors_texture_core(CYTHON_UNUSED PyObject *__pyx_self, std::string __pyx_v_filename, std::string __pyx_v_mtl_name, PyArrayObject *__pyx_v_vertices, PyArrayObject *__pyx_v_triangles, PyArrayObject *__pyx_v_colors, PyArrayObject *__pyx_v_uv_coords, int __pyx_v_nver, int __pyx_v_ntri, int __pyx_v_ntexver) { - __Pyx_LocalBuf_ND __pyx_pybuffernd_colors; - __Pyx_Buffer __pyx_pybuffer_colors; - __Pyx_LocalBuf_ND __pyx_pybuffernd_triangles; - __Pyx_Buffer __pyx_pybuffer_triangles; - __Pyx_LocalBuf_ND __pyx_pybuffernd_uv_coords; - __Pyx_Buffer __pyx_pybuffer_uv_coords; - __Pyx_LocalBuf_ND __pyx_pybuffernd_vertices; - __Pyx_Buffer __pyx_pybuffer_vertices; - PyObject *__pyx_r = NULL; - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("write_obj_with_colors_texture_core", 0); - __pyx_pybuffer_vertices.pybuffer.buf = NULL; - __pyx_pybuffer_vertices.refcount = 0; - __pyx_pybuffernd_vertices.data = NULL; - __pyx_pybuffernd_vertices.rcbuffer = &__pyx_pybuffer_vertices; - __pyx_pybuffer_triangles.pybuffer.buf = NULL; - __pyx_pybuffer_triangles.refcount = 0; - __pyx_pybuffernd_triangles.data = NULL; - __pyx_pybuffernd_triangles.rcbuffer = &__pyx_pybuffer_triangles; - __pyx_pybuffer_colors.pybuffer.buf = NULL; - __pyx_pybuffer_colors.refcount = 0; - __pyx_pybuffernd_colors.data = NULL; - __pyx_pybuffernd_colors.rcbuffer = &__pyx_pybuffer_colors; - __pyx_pybuffer_uv_coords.pybuffer.buf = NULL; - __pyx_pybuffer_uv_coords.refcount = 0; - __pyx_pybuffernd_uv_coords.data = NULL; - __pyx_pybuffernd_uv_coords.rcbuffer = &__pyx_pybuffer_uv_coords; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_vertices.rcbuffer->pybuffer, (PyObject*)__pyx_v_vertices, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 100, __pyx_L1_error) - } - __pyx_pybuffernd_vertices.diminfo[0].strides = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_vertices.diminfo[0].shape = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_vertices.diminfo[1].strides = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_vertices.diminfo[1].shape = __pyx_pybuffernd_vertices.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer, (PyObject*)__pyx_v_triangles, &__Pyx_TypeInfo_int, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 100, __pyx_L1_error) - } - __pyx_pybuffernd_triangles.diminfo[0].strides = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_triangles.diminfo[0].shape = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_triangles.diminfo[1].strides = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_triangles.diminfo[1].shape = __pyx_pybuffernd_triangles.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_colors.rcbuffer->pybuffer, (PyObject*)__pyx_v_colors, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 100, __pyx_L1_error) - } - __pyx_pybuffernd_colors.diminfo[0].strides = __pyx_pybuffernd_colors.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_colors.diminfo[0].shape = __pyx_pybuffernd_colors.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_colors.diminfo[1].strides = __pyx_pybuffernd_colors.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_colors.diminfo[1].shape = __pyx_pybuffernd_colors.rcbuffer->pybuffer.shape[1]; - { - __Pyx_BufFmt_StackElem __pyx_stack[1]; - if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_uv_coords.rcbuffer->pybuffer, (PyObject*)__pyx_v_uv_coords, &__Pyx_TypeInfo_float, PyBUF_FORMAT| PyBUF_C_CONTIGUOUS, 2, 0, __pyx_stack) == -1)) __PYX_ERR(0, 100, __pyx_L1_error) - } - __pyx_pybuffernd_uv_coords.diminfo[0].strides = __pyx_pybuffernd_uv_coords.rcbuffer->pybuffer.strides[0]; __pyx_pybuffernd_uv_coords.diminfo[0].shape = __pyx_pybuffernd_uv_coords.rcbuffer->pybuffer.shape[0]; __pyx_pybuffernd_uv_coords.diminfo[1].strides = __pyx_pybuffernd_uv_coords.rcbuffer->pybuffer.strides[1]; __pyx_pybuffernd_uv_coords.diminfo[1].shape = __pyx_pybuffernd_uv_coords.rcbuffer->pybuffer.shape[1]; - - /* "mesh_core_cython.pyx":107 - * int nver, int ntri, int ntexver - * ): - * _write_obj_with_colors_texture(filename, mtl_name, # <<<<<<<<<<<<<< - * np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), np.PyArray_DATA(colors), np.PyArray_DATA(uv_coords), - * nver, ntri, ntexver) - */ - _write_obj_with_colors_texture(__pyx_v_filename, __pyx_v_mtl_name, ((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_vertices))), ((int *)PyArray_DATA(((PyArrayObject *)__pyx_v_triangles))), ((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_colors))), ((float *)PyArray_DATA(((PyArrayObject *)__pyx_v_uv_coords))), __pyx_v_nver, __pyx_v_ntri, __pyx_v_ntexver); - - /* "mesh_core_cython.pyx":100 - * mapping_type) - * - * def write_obj_with_colors_texture_core(string filename, string mtl_name, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - - /* function exit code */ - __pyx_r = Py_None; __Pyx_INCREF(Py_None); - goto __pyx_L0; - __pyx_L1_error:; - { PyObject *__pyx_type, *__pyx_value, *__pyx_tb; - __Pyx_PyThreadState_declare - __Pyx_PyThreadState_assign - __Pyx_ErrFetch(&__pyx_type, &__pyx_value, &__pyx_tb); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_colors.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_uv_coords.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_vertices.rcbuffer->pybuffer); - __Pyx_ErrRestore(__pyx_type, __pyx_value, __pyx_tb);} - __Pyx_AddTraceback("mesh_core_cython.write_obj_with_colors_texture_core", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = NULL; - goto __pyx_L2; - __pyx_L0:; - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_colors.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_triangles.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_uv_coords.rcbuffer->pybuffer); - __Pyx_SafeReleaseBuffer(&__pyx_pybuffernd_vertices.rcbuffer->pybuffer); - __pyx_L2:; - __Pyx_XGIVEREF(__pyx_r); - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":215 - * # experimental exception made for __getbuffer__ and __releasebuffer__ - * # -- the details of this may change. - * def __getbuffer__(ndarray self, Py_buffer* info, int flags): # <<<<<<<<<<<<<< - * # This implementation of getbuffer is geared towards Cython - * # requirements, and does not yet fulfill the PEP. - */ - -/* Python wrapper */ -static CYTHON_UNUSED int __pyx_pw_5numpy_7ndarray_1__getbuffer__(PyObject *__pyx_v_self, Py_buffer *__pyx_v_info, int __pyx_v_flags); /*proto*/ -static CYTHON_UNUSED int __pyx_pw_5numpy_7ndarray_1__getbuffer__(PyObject *__pyx_v_self, Py_buffer *__pyx_v_info, int __pyx_v_flags) { - int __pyx_r; - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("__getbuffer__ (wrapper)", 0); - __pyx_r = __pyx_pf_5numpy_7ndarray___getbuffer__(((PyArrayObject *)__pyx_v_self), ((Py_buffer *)__pyx_v_info), ((int)__pyx_v_flags)); - - /* function exit code */ - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -static int __pyx_pf_5numpy_7ndarray___getbuffer__(PyArrayObject *__pyx_v_self, Py_buffer *__pyx_v_info, int __pyx_v_flags) { - int __pyx_v_i; - int __pyx_v_ndim; - int __pyx_v_endian_detector; - int __pyx_v_little_endian; - int __pyx_v_t; - char *__pyx_v_f; - PyArray_Descr *__pyx_v_descr = 0; - int __pyx_v_offset; - int __pyx_r; - __Pyx_RefNannyDeclarations - int __pyx_t_1; - int __pyx_t_2; - PyObject *__pyx_t_3 = NULL; - int __pyx_t_4; - int __pyx_t_5; - int __pyx_t_6; - PyObject *__pyx_t_7 = NULL; - char *__pyx_t_8; - if (__pyx_v_info == NULL) { - PyErr_SetString(PyExc_BufferError, "PyObject_GetBuffer: view==NULL argument is obsolete"); - return -1; - } - __Pyx_RefNannySetupContext("__getbuffer__", 0); - __pyx_v_info->obj = Py_None; __Pyx_INCREF(Py_None); - __Pyx_GIVEREF(__pyx_v_info->obj); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":222 - * - * cdef int i, ndim - * cdef int endian_detector = 1 # <<<<<<<<<<<<<< - * cdef bint little_endian = ((&endian_detector)[0] != 0) - * - */ - __pyx_v_endian_detector = 1; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":223 - * cdef int i, ndim - * cdef int endian_detector = 1 - * cdef bint little_endian = ((&endian_detector)[0] != 0) # <<<<<<<<<<<<<< - * - * ndim = PyArray_NDIM(self) - */ - __pyx_v_little_endian = ((((char *)(&__pyx_v_endian_detector))[0]) != 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":225 - * cdef bint little_endian = ((&endian_detector)[0] != 0) - * - * ndim = PyArray_NDIM(self) # <<<<<<<<<<<<<< - * - * if ((flags & pybuf.PyBUF_C_CONTIGUOUS == pybuf.PyBUF_C_CONTIGUOUS) - */ - __pyx_v_ndim = PyArray_NDIM(__pyx_v_self); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":227 - * ndim = PyArray_NDIM(self) - * - * if ((flags & pybuf.PyBUF_C_CONTIGUOUS == pybuf.PyBUF_C_CONTIGUOUS) # <<<<<<<<<<<<<< - * and not PyArray_CHKFLAGS(self, NPY_C_CONTIGUOUS)): - * raise ValueError(u"ndarray is not C contiguous") - */ - __pyx_t_2 = (((__pyx_v_flags & PyBUF_C_CONTIGUOUS) == PyBUF_C_CONTIGUOUS) != 0); - if (__pyx_t_2) { - } else { - __pyx_t_1 = __pyx_t_2; - goto __pyx_L4_bool_binop_done; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":228 - * - * if ((flags & pybuf.PyBUF_C_CONTIGUOUS == pybuf.PyBUF_C_CONTIGUOUS) - * and not PyArray_CHKFLAGS(self, NPY_C_CONTIGUOUS)): # <<<<<<<<<<<<<< - * raise ValueError(u"ndarray is not C contiguous") - * - */ - __pyx_t_2 = ((!(PyArray_CHKFLAGS(__pyx_v_self, NPY_C_CONTIGUOUS) != 0)) != 0); - __pyx_t_1 = __pyx_t_2; - __pyx_L4_bool_binop_done:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":227 - * ndim = PyArray_NDIM(self) - * - * if ((flags & pybuf.PyBUF_C_CONTIGUOUS == pybuf.PyBUF_C_CONTIGUOUS) # <<<<<<<<<<<<<< - * and not PyArray_CHKFLAGS(self, NPY_C_CONTIGUOUS)): - * raise ValueError(u"ndarray is not C contiguous") - */ - if (unlikely(__pyx_t_1)) { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":229 - * if ((flags & pybuf.PyBUF_C_CONTIGUOUS == pybuf.PyBUF_C_CONTIGUOUS) - * and not PyArray_CHKFLAGS(self, NPY_C_CONTIGUOUS)): - * raise ValueError(u"ndarray is not C contiguous") # <<<<<<<<<<<<<< - * - * if ((flags & pybuf.PyBUF_F_CONTIGUOUS == pybuf.PyBUF_F_CONTIGUOUS) - */ - __pyx_t_3 = __Pyx_PyObject_Call(__pyx_builtin_ValueError, __pyx_tuple_, NULL); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 229, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __Pyx_Raise(__pyx_t_3, 0, 0, 0); - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __PYX_ERR(1, 229, __pyx_L1_error) - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":227 - * ndim = PyArray_NDIM(self) - * - * if ((flags & pybuf.PyBUF_C_CONTIGUOUS == pybuf.PyBUF_C_CONTIGUOUS) # <<<<<<<<<<<<<< - * and not PyArray_CHKFLAGS(self, NPY_C_CONTIGUOUS)): - * raise ValueError(u"ndarray is not C contiguous") - */ - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":231 - * raise ValueError(u"ndarray is not C contiguous") - * - * if ((flags & pybuf.PyBUF_F_CONTIGUOUS == pybuf.PyBUF_F_CONTIGUOUS) # <<<<<<<<<<<<<< - * and not PyArray_CHKFLAGS(self, NPY_F_CONTIGUOUS)): - * raise ValueError(u"ndarray is not Fortran contiguous") - */ - __pyx_t_2 = (((__pyx_v_flags & PyBUF_F_CONTIGUOUS) == PyBUF_F_CONTIGUOUS) != 0); - if (__pyx_t_2) { - } else { - __pyx_t_1 = __pyx_t_2; - goto __pyx_L7_bool_binop_done; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":232 - * - * if ((flags & pybuf.PyBUF_F_CONTIGUOUS == pybuf.PyBUF_F_CONTIGUOUS) - * and not PyArray_CHKFLAGS(self, NPY_F_CONTIGUOUS)): # <<<<<<<<<<<<<< - * raise ValueError(u"ndarray is not Fortran contiguous") - * - */ - __pyx_t_2 = ((!(PyArray_CHKFLAGS(__pyx_v_self, NPY_F_CONTIGUOUS) != 0)) != 0); - __pyx_t_1 = __pyx_t_2; - __pyx_L7_bool_binop_done:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":231 - * raise ValueError(u"ndarray is not C contiguous") - * - * if ((flags & pybuf.PyBUF_F_CONTIGUOUS == pybuf.PyBUF_F_CONTIGUOUS) # <<<<<<<<<<<<<< - * and not PyArray_CHKFLAGS(self, NPY_F_CONTIGUOUS)): - * raise ValueError(u"ndarray is not Fortran contiguous") - */ - if (unlikely(__pyx_t_1)) { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":233 - * if ((flags & pybuf.PyBUF_F_CONTIGUOUS == pybuf.PyBUF_F_CONTIGUOUS) - * and not PyArray_CHKFLAGS(self, NPY_F_CONTIGUOUS)): - * raise ValueError(u"ndarray is not Fortran contiguous") # <<<<<<<<<<<<<< - * - * info.buf = PyArray_DATA(self) - */ - __pyx_t_3 = __Pyx_PyObject_Call(__pyx_builtin_ValueError, __pyx_tuple__2, NULL); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 233, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __Pyx_Raise(__pyx_t_3, 0, 0, 0); - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __PYX_ERR(1, 233, __pyx_L1_error) - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":231 - * raise ValueError(u"ndarray is not C contiguous") - * - * if ((flags & pybuf.PyBUF_F_CONTIGUOUS == pybuf.PyBUF_F_CONTIGUOUS) # <<<<<<<<<<<<<< - * and not PyArray_CHKFLAGS(self, NPY_F_CONTIGUOUS)): - * raise ValueError(u"ndarray is not Fortran contiguous") - */ - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":235 - * raise ValueError(u"ndarray is not Fortran contiguous") - * - * info.buf = PyArray_DATA(self) # <<<<<<<<<<<<<< - * info.ndim = ndim - * if sizeof(npy_intp) != sizeof(Py_ssize_t): - */ - __pyx_v_info->buf = PyArray_DATA(__pyx_v_self); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":236 - * - * info.buf = PyArray_DATA(self) - * info.ndim = ndim # <<<<<<<<<<<<<< - * if sizeof(npy_intp) != sizeof(Py_ssize_t): - * # Allocate new buffer for strides and shape info. - */ - __pyx_v_info->ndim = __pyx_v_ndim; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":237 - * info.buf = PyArray_DATA(self) - * info.ndim = ndim - * if sizeof(npy_intp) != sizeof(Py_ssize_t): # <<<<<<<<<<<<<< - * # Allocate new buffer for strides and shape info. - * # This is allocated as one block, strides first. - */ - __pyx_t_1 = (((sizeof(npy_intp)) != (sizeof(Py_ssize_t))) != 0); - if (__pyx_t_1) { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":240 - * # Allocate new buffer for strides and shape info. - * # This is allocated as one block, strides first. - * info.strides = PyObject_Malloc(sizeof(Py_ssize_t) * 2 * ndim) # <<<<<<<<<<<<<< - * info.shape = info.strides + ndim - * for i in range(ndim): - */ - __pyx_v_info->strides = ((Py_ssize_t *)PyObject_Malloc((((sizeof(Py_ssize_t)) * 2) * ((size_t)__pyx_v_ndim)))); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":241 - * # This is allocated as one block, strides first. - * info.strides = PyObject_Malloc(sizeof(Py_ssize_t) * 2 * ndim) - * info.shape = info.strides + ndim # <<<<<<<<<<<<<< - * for i in range(ndim): - * info.strides[i] = PyArray_STRIDES(self)[i] - */ - __pyx_v_info->shape = (__pyx_v_info->strides + __pyx_v_ndim); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":242 - * info.strides = PyObject_Malloc(sizeof(Py_ssize_t) * 2 * ndim) - * info.shape = info.strides + ndim - * for i in range(ndim): # <<<<<<<<<<<<<< - * info.strides[i] = PyArray_STRIDES(self)[i] - * info.shape[i] = PyArray_DIMS(self)[i] - */ - __pyx_t_4 = __pyx_v_ndim; - __pyx_t_5 = __pyx_t_4; - for (__pyx_t_6 = 0; __pyx_t_6 < __pyx_t_5; __pyx_t_6+=1) { - __pyx_v_i = __pyx_t_6; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":243 - * info.shape = info.strides + ndim - * for i in range(ndim): - * info.strides[i] = PyArray_STRIDES(self)[i] # <<<<<<<<<<<<<< - * info.shape[i] = PyArray_DIMS(self)[i] - * else: - */ - (__pyx_v_info->strides[__pyx_v_i]) = (PyArray_STRIDES(__pyx_v_self)[__pyx_v_i]); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":244 - * for i in range(ndim): - * info.strides[i] = PyArray_STRIDES(self)[i] - * info.shape[i] = PyArray_DIMS(self)[i] # <<<<<<<<<<<<<< - * else: - * info.strides = PyArray_STRIDES(self) - */ - (__pyx_v_info->shape[__pyx_v_i]) = (PyArray_DIMS(__pyx_v_self)[__pyx_v_i]); - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":237 - * info.buf = PyArray_DATA(self) - * info.ndim = ndim - * if sizeof(npy_intp) != sizeof(Py_ssize_t): # <<<<<<<<<<<<<< - * # Allocate new buffer for strides and shape info. - * # This is allocated as one block, strides first. - */ - goto __pyx_L9; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":246 - * info.shape[i] = PyArray_DIMS(self)[i] - * else: - * info.strides = PyArray_STRIDES(self) # <<<<<<<<<<<<<< - * info.shape = PyArray_DIMS(self) - * info.suboffsets = NULL - */ - /*else*/ { - __pyx_v_info->strides = ((Py_ssize_t *)PyArray_STRIDES(__pyx_v_self)); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":247 - * else: - * info.strides = PyArray_STRIDES(self) - * info.shape = PyArray_DIMS(self) # <<<<<<<<<<<<<< - * info.suboffsets = NULL - * info.itemsize = PyArray_ITEMSIZE(self) - */ - __pyx_v_info->shape = ((Py_ssize_t *)PyArray_DIMS(__pyx_v_self)); - } - __pyx_L9:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":248 - * info.strides = PyArray_STRIDES(self) - * info.shape = PyArray_DIMS(self) - * info.suboffsets = NULL # <<<<<<<<<<<<<< - * info.itemsize = PyArray_ITEMSIZE(self) - * info.readonly = not PyArray_ISWRITEABLE(self) - */ - __pyx_v_info->suboffsets = NULL; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":249 - * info.shape = PyArray_DIMS(self) - * info.suboffsets = NULL - * info.itemsize = PyArray_ITEMSIZE(self) # <<<<<<<<<<<<<< - * info.readonly = not PyArray_ISWRITEABLE(self) - * - */ - __pyx_v_info->itemsize = PyArray_ITEMSIZE(__pyx_v_self); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":250 - * info.suboffsets = NULL - * info.itemsize = PyArray_ITEMSIZE(self) - * info.readonly = not PyArray_ISWRITEABLE(self) # <<<<<<<<<<<<<< - * - * cdef int t - */ - __pyx_v_info->readonly = (!(PyArray_ISWRITEABLE(__pyx_v_self) != 0)); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":253 - * - * cdef int t - * cdef char* f = NULL # <<<<<<<<<<<<<< - * cdef dtype descr = self.descr - * cdef int offset - */ - __pyx_v_f = NULL; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":254 - * cdef int t - * cdef char* f = NULL - * cdef dtype descr = self.descr # <<<<<<<<<<<<<< - * cdef int offset - * - */ - __pyx_t_3 = ((PyObject *)__pyx_v_self->descr); - __Pyx_INCREF(__pyx_t_3); - __pyx_v_descr = ((PyArray_Descr *)__pyx_t_3); - __pyx_t_3 = 0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":257 - * cdef int offset - * - * info.obj = self # <<<<<<<<<<<<<< - * - * if not PyDataType_HASFIELDS(descr): - */ - __Pyx_INCREF(((PyObject *)__pyx_v_self)); - __Pyx_GIVEREF(((PyObject *)__pyx_v_self)); - __Pyx_GOTREF(__pyx_v_info->obj); - __Pyx_DECREF(__pyx_v_info->obj); - __pyx_v_info->obj = ((PyObject *)__pyx_v_self); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":259 - * info.obj = self - * - * if not PyDataType_HASFIELDS(descr): # <<<<<<<<<<<<<< - * t = descr.type_num - * if ((descr.byteorder == c'>' and little_endian) or - */ - __pyx_t_1 = ((!(PyDataType_HASFIELDS(__pyx_v_descr) != 0)) != 0); - if (__pyx_t_1) { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":260 - * - * if not PyDataType_HASFIELDS(descr): - * t = descr.type_num # <<<<<<<<<<<<<< - * if ((descr.byteorder == c'>' and little_endian) or - * (descr.byteorder == c'<' and not little_endian)): - */ - __pyx_t_4 = __pyx_v_descr->type_num; - __pyx_v_t = __pyx_t_4; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":261 - * if not PyDataType_HASFIELDS(descr): - * t = descr.type_num - * if ((descr.byteorder == c'>' and little_endian) or # <<<<<<<<<<<<<< - * (descr.byteorder == c'<' and not little_endian)): - * raise ValueError(u"Non-native byte order not supported") - */ - __pyx_t_2 = ((__pyx_v_descr->byteorder == '>') != 0); - if (!__pyx_t_2) { - goto __pyx_L15_next_or; - } else { - } - __pyx_t_2 = (__pyx_v_little_endian != 0); - if (!__pyx_t_2) { - } else { - __pyx_t_1 = __pyx_t_2; - goto __pyx_L14_bool_binop_done; - } - __pyx_L15_next_or:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":262 - * t = descr.type_num - * if ((descr.byteorder == c'>' and little_endian) or - * (descr.byteorder == c'<' and not little_endian)): # <<<<<<<<<<<<<< - * raise ValueError(u"Non-native byte order not supported") - * if t == NPY_BYTE: f = "b" - */ - __pyx_t_2 = ((__pyx_v_descr->byteorder == '<') != 0); - if (__pyx_t_2) { - } else { - __pyx_t_1 = __pyx_t_2; - goto __pyx_L14_bool_binop_done; - } - __pyx_t_2 = ((!(__pyx_v_little_endian != 0)) != 0); - __pyx_t_1 = __pyx_t_2; - __pyx_L14_bool_binop_done:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":261 - * if not PyDataType_HASFIELDS(descr): - * t = descr.type_num - * if ((descr.byteorder == c'>' and little_endian) or # <<<<<<<<<<<<<< - * (descr.byteorder == c'<' and not little_endian)): - * raise ValueError(u"Non-native byte order not supported") - */ - if (unlikely(__pyx_t_1)) { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":263 - * if ((descr.byteorder == c'>' and little_endian) or - * (descr.byteorder == c'<' and not little_endian)): - * raise ValueError(u"Non-native byte order not supported") # <<<<<<<<<<<<<< - * if t == NPY_BYTE: f = "b" - * elif t == NPY_UBYTE: f = "B" - */ - __pyx_t_3 = __Pyx_PyObject_Call(__pyx_builtin_ValueError, __pyx_tuple__3, NULL); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 263, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __Pyx_Raise(__pyx_t_3, 0, 0, 0); - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __PYX_ERR(1, 263, __pyx_L1_error) - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":261 - * if not PyDataType_HASFIELDS(descr): - * t = descr.type_num - * if ((descr.byteorder == c'>' and little_endian) or # <<<<<<<<<<<<<< - * (descr.byteorder == c'<' and not little_endian)): - * raise ValueError(u"Non-native byte order not supported") - */ - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":264 - * (descr.byteorder == c'<' and not little_endian)): - * raise ValueError(u"Non-native byte order not supported") - * if t == NPY_BYTE: f = "b" # <<<<<<<<<<<<<< - * elif t == NPY_UBYTE: f = "B" - * elif t == NPY_SHORT: f = "h" - */ - switch (__pyx_v_t) { - case NPY_BYTE: - __pyx_v_f = ((char *)"b"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":265 - * raise ValueError(u"Non-native byte order not supported") - * if t == NPY_BYTE: f = "b" - * elif t == NPY_UBYTE: f = "B" # <<<<<<<<<<<<<< - * elif t == NPY_SHORT: f = "h" - * elif t == NPY_USHORT: f = "H" - */ - case NPY_UBYTE: - __pyx_v_f = ((char *)"B"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":266 - * if t == NPY_BYTE: f = "b" - * elif t == NPY_UBYTE: f = "B" - * elif t == NPY_SHORT: f = "h" # <<<<<<<<<<<<<< - * elif t == NPY_USHORT: f = "H" - * elif t == NPY_INT: f = "i" - */ - case NPY_SHORT: - __pyx_v_f = ((char *)"h"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":267 - * elif t == NPY_UBYTE: f = "B" - * elif t == NPY_SHORT: f = "h" - * elif t == NPY_USHORT: f = "H" # <<<<<<<<<<<<<< - * elif t == NPY_INT: f = "i" - * elif t == NPY_UINT: f = "I" - */ - case NPY_USHORT: - __pyx_v_f = ((char *)"H"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":268 - * elif t == NPY_SHORT: f = "h" - * elif t == NPY_USHORT: f = "H" - * elif t == NPY_INT: f = "i" # <<<<<<<<<<<<<< - * elif t == NPY_UINT: f = "I" - * elif t == NPY_LONG: f = "l" - */ - case NPY_INT: - __pyx_v_f = ((char *)"i"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":269 - * elif t == NPY_USHORT: f = "H" - * elif t == NPY_INT: f = "i" - * elif t == NPY_UINT: f = "I" # <<<<<<<<<<<<<< - * elif t == NPY_LONG: f = "l" - * elif t == NPY_ULONG: f = "L" - */ - case NPY_UINT: - __pyx_v_f = ((char *)"I"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":270 - * elif t == NPY_INT: f = "i" - * elif t == NPY_UINT: f = "I" - * elif t == NPY_LONG: f = "l" # <<<<<<<<<<<<<< - * elif t == NPY_ULONG: f = "L" - * elif t == NPY_LONGLONG: f = "q" - */ - case NPY_LONG: - __pyx_v_f = ((char *)"l"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":271 - * elif t == NPY_UINT: f = "I" - * elif t == NPY_LONG: f = "l" - * elif t == NPY_ULONG: f = "L" # <<<<<<<<<<<<<< - * elif t == NPY_LONGLONG: f = "q" - * elif t == NPY_ULONGLONG: f = "Q" - */ - case NPY_ULONG: - __pyx_v_f = ((char *)"L"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":272 - * elif t == NPY_LONG: f = "l" - * elif t == NPY_ULONG: f = "L" - * elif t == NPY_LONGLONG: f = "q" # <<<<<<<<<<<<<< - * elif t == NPY_ULONGLONG: f = "Q" - * elif t == NPY_FLOAT: f = "f" - */ - case NPY_LONGLONG: - __pyx_v_f = ((char *)"q"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":273 - * elif t == NPY_ULONG: f = "L" - * elif t == NPY_LONGLONG: f = "q" - * elif t == NPY_ULONGLONG: f = "Q" # <<<<<<<<<<<<<< - * elif t == NPY_FLOAT: f = "f" - * elif t == NPY_DOUBLE: f = "d" - */ - case NPY_ULONGLONG: - __pyx_v_f = ((char *)"Q"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":274 - * elif t == NPY_LONGLONG: f = "q" - * elif t == NPY_ULONGLONG: f = "Q" - * elif t == NPY_FLOAT: f = "f" # <<<<<<<<<<<<<< - * elif t == NPY_DOUBLE: f = "d" - * elif t == NPY_LONGDOUBLE: f = "g" - */ - case NPY_FLOAT: - __pyx_v_f = ((char *)"f"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":275 - * elif t == NPY_ULONGLONG: f = "Q" - * elif t == NPY_FLOAT: f = "f" - * elif t == NPY_DOUBLE: f = "d" # <<<<<<<<<<<<<< - * elif t == NPY_LONGDOUBLE: f = "g" - * elif t == NPY_CFLOAT: f = "Zf" - */ - case NPY_DOUBLE: - __pyx_v_f = ((char *)"d"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":276 - * elif t == NPY_FLOAT: f = "f" - * elif t == NPY_DOUBLE: f = "d" - * elif t == NPY_LONGDOUBLE: f = "g" # <<<<<<<<<<<<<< - * elif t == NPY_CFLOAT: f = "Zf" - * elif t == NPY_CDOUBLE: f = "Zd" - */ - case NPY_LONGDOUBLE: - __pyx_v_f = ((char *)"g"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":277 - * elif t == NPY_DOUBLE: f = "d" - * elif t == NPY_LONGDOUBLE: f = "g" - * elif t == NPY_CFLOAT: f = "Zf" # <<<<<<<<<<<<<< - * elif t == NPY_CDOUBLE: f = "Zd" - * elif t == NPY_CLONGDOUBLE: f = "Zg" - */ - case NPY_CFLOAT: - __pyx_v_f = ((char *)"Zf"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":278 - * elif t == NPY_LONGDOUBLE: f = "g" - * elif t == NPY_CFLOAT: f = "Zf" - * elif t == NPY_CDOUBLE: f = "Zd" # <<<<<<<<<<<<<< - * elif t == NPY_CLONGDOUBLE: f = "Zg" - * elif t == NPY_OBJECT: f = "O" - */ - case NPY_CDOUBLE: - __pyx_v_f = ((char *)"Zd"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":279 - * elif t == NPY_CFLOAT: f = "Zf" - * elif t == NPY_CDOUBLE: f = "Zd" - * elif t == NPY_CLONGDOUBLE: f = "Zg" # <<<<<<<<<<<<<< - * elif t == NPY_OBJECT: f = "O" - * else: - */ - case NPY_CLONGDOUBLE: - __pyx_v_f = ((char *)"Zg"); - break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":280 - * elif t == NPY_CDOUBLE: f = "Zd" - * elif t == NPY_CLONGDOUBLE: f = "Zg" - * elif t == NPY_OBJECT: f = "O" # <<<<<<<<<<<<<< - * else: - * raise ValueError(u"unknown dtype code in numpy.pxd (%d)" % t) - */ - case NPY_OBJECT: - __pyx_v_f = ((char *)"O"); - break; - default: - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":282 - * elif t == NPY_OBJECT: f = "O" - * else: - * raise ValueError(u"unknown dtype code in numpy.pxd (%d)" % t) # <<<<<<<<<<<<<< - * info.format = f - * return - */ - __pyx_t_3 = __Pyx_PyInt_From_int(__pyx_v_t); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 282, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __pyx_t_7 = PyUnicode_Format(__pyx_kp_u_unknown_dtype_code_in_numpy_pxd, __pyx_t_3); if (unlikely(!__pyx_t_7)) __PYX_ERR(1, 282, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_7); - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __pyx_t_3 = __Pyx_PyObject_CallOneArg(__pyx_builtin_ValueError, __pyx_t_7); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 282, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __Pyx_DECREF(__pyx_t_7); __pyx_t_7 = 0; - __Pyx_Raise(__pyx_t_3, 0, 0, 0); - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __PYX_ERR(1, 282, __pyx_L1_error) - break; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":283 - * else: - * raise ValueError(u"unknown dtype code in numpy.pxd (%d)" % t) - * info.format = f # <<<<<<<<<<<<<< - * return - * else: - */ - __pyx_v_info->format = __pyx_v_f; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":284 - * raise ValueError(u"unknown dtype code in numpy.pxd (%d)" % t) - * info.format = f - * return # <<<<<<<<<<<<<< - * else: - * info.format = PyObject_Malloc(_buffer_format_string_len) - */ - __pyx_r = 0; - goto __pyx_L0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":259 - * info.obj = self - * - * if not PyDataType_HASFIELDS(descr): # <<<<<<<<<<<<<< - * t = descr.type_num - * if ((descr.byteorder == c'>' and little_endian) or - */ - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":286 - * return - * else: - * info.format = PyObject_Malloc(_buffer_format_string_len) # <<<<<<<<<<<<<< - * info.format[0] = c'^' # Native data types, manual alignment - * offset = 0 - */ - /*else*/ { - __pyx_v_info->format = ((char *)PyObject_Malloc(0xFF)); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":287 - * else: - * info.format = PyObject_Malloc(_buffer_format_string_len) - * info.format[0] = c'^' # Native data types, manual alignment # <<<<<<<<<<<<<< - * offset = 0 - * f = _util_dtypestring(descr, info.format + 1, - */ - (__pyx_v_info->format[0]) = '^'; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":288 - * info.format = PyObject_Malloc(_buffer_format_string_len) - * info.format[0] = c'^' # Native data types, manual alignment - * offset = 0 # <<<<<<<<<<<<<< - * f = _util_dtypestring(descr, info.format + 1, - * info.format + _buffer_format_string_len, - */ - __pyx_v_offset = 0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":289 - * info.format[0] = c'^' # Native data types, manual alignment - * offset = 0 - * f = _util_dtypestring(descr, info.format + 1, # <<<<<<<<<<<<<< - * info.format + _buffer_format_string_len, - * &offset) - */ - __pyx_t_8 = __pyx_f_5numpy__util_dtypestring(__pyx_v_descr, (__pyx_v_info->format + 1), (__pyx_v_info->format + 0xFF), (&__pyx_v_offset)); if (unlikely(__pyx_t_8 == ((char *)NULL))) __PYX_ERR(1, 289, __pyx_L1_error) - __pyx_v_f = __pyx_t_8; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":292 - * info.format + _buffer_format_string_len, - * &offset) - * f[0] = c'\0' # Terminate format string # <<<<<<<<<<<<<< - * - * def __releasebuffer__(ndarray self, Py_buffer* info): - */ - (__pyx_v_f[0]) = '\x00'; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":215 - * # experimental exception made for __getbuffer__ and __releasebuffer__ - * # -- the details of this may change. - * def __getbuffer__(ndarray self, Py_buffer* info, int flags): # <<<<<<<<<<<<<< - * # This implementation of getbuffer is geared towards Cython - * # requirements, and does not yet fulfill the PEP. - */ - - /* function exit code */ - __pyx_r = 0; - goto __pyx_L0; - __pyx_L1_error:; - __Pyx_XDECREF(__pyx_t_3); - __Pyx_XDECREF(__pyx_t_7); - __Pyx_AddTraceback("numpy.ndarray.__getbuffer__", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = -1; - if (__pyx_v_info->obj != NULL) { - __Pyx_GOTREF(__pyx_v_info->obj); - __Pyx_DECREF(__pyx_v_info->obj); __pyx_v_info->obj = 0; - } - goto __pyx_L2; - __pyx_L0:; - if (__pyx_v_info->obj == Py_None) { - __Pyx_GOTREF(__pyx_v_info->obj); - __Pyx_DECREF(__pyx_v_info->obj); __pyx_v_info->obj = 0; - } - __pyx_L2:; - __Pyx_XDECREF((PyObject *)__pyx_v_descr); - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":294 - * f[0] = c'\0' # Terminate format string - * - * def __releasebuffer__(ndarray self, Py_buffer* info): # <<<<<<<<<<<<<< - * if PyArray_HASFIELDS(self): - * PyObject_Free(info.format) - */ - -/* Python wrapper */ -static CYTHON_UNUSED void __pyx_pw_5numpy_7ndarray_3__releasebuffer__(PyObject *__pyx_v_self, Py_buffer *__pyx_v_info); /*proto*/ -static CYTHON_UNUSED void __pyx_pw_5numpy_7ndarray_3__releasebuffer__(PyObject *__pyx_v_self, Py_buffer *__pyx_v_info) { - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("__releasebuffer__ (wrapper)", 0); - __pyx_pf_5numpy_7ndarray_2__releasebuffer__(((PyArrayObject *)__pyx_v_self), ((Py_buffer *)__pyx_v_info)); - - /* function exit code */ - __Pyx_RefNannyFinishContext(); -} - -static void __pyx_pf_5numpy_7ndarray_2__releasebuffer__(PyArrayObject *__pyx_v_self, Py_buffer *__pyx_v_info) { - __Pyx_RefNannyDeclarations - int __pyx_t_1; - __Pyx_RefNannySetupContext("__releasebuffer__", 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":295 - * - * def __releasebuffer__(ndarray self, Py_buffer* info): - * if PyArray_HASFIELDS(self): # <<<<<<<<<<<<<< - * PyObject_Free(info.format) - * if sizeof(npy_intp) != sizeof(Py_ssize_t): - */ - __pyx_t_1 = (PyArray_HASFIELDS(__pyx_v_self) != 0); - if (__pyx_t_1) { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":296 - * def __releasebuffer__(ndarray self, Py_buffer* info): - * if PyArray_HASFIELDS(self): - * PyObject_Free(info.format) # <<<<<<<<<<<<<< - * if sizeof(npy_intp) != sizeof(Py_ssize_t): - * PyObject_Free(info.strides) - */ - PyObject_Free(__pyx_v_info->format); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":295 - * - * def __releasebuffer__(ndarray self, Py_buffer* info): - * if PyArray_HASFIELDS(self): # <<<<<<<<<<<<<< - * PyObject_Free(info.format) - * if sizeof(npy_intp) != sizeof(Py_ssize_t): - */ - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":297 - * if PyArray_HASFIELDS(self): - * PyObject_Free(info.format) - * if sizeof(npy_intp) != sizeof(Py_ssize_t): # <<<<<<<<<<<<<< - * PyObject_Free(info.strides) - * # info.shape was stored after info.strides in the same block - */ - __pyx_t_1 = (((sizeof(npy_intp)) != (sizeof(Py_ssize_t))) != 0); - if (__pyx_t_1) { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":298 - * PyObject_Free(info.format) - * if sizeof(npy_intp) != sizeof(Py_ssize_t): - * PyObject_Free(info.strides) # <<<<<<<<<<<<<< - * # info.shape was stored after info.strides in the same block - * - */ - PyObject_Free(__pyx_v_info->strides); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":297 - * if PyArray_HASFIELDS(self): - * PyObject_Free(info.format) - * if sizeof(npy_intp) != sizeof(Py_ssize_t): # <<<<<<<<<<<<<< - * PyObject_Free(info.strides) - * # info.shape was stored after info.strides in the same block - */ - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":294 - * f[0] = c'\0' # Terminate format string - * - * def __releasebuffer__(ndarray self, Py_buffer* info): # <<<<<<<<<<<<<< - * if PyArray_HASFIELDS(self): - * PyObject_Free(info.format) - */ - - /* function exit code */ - __Pyx_RefNannyFinishContext(); -} - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":775 - * ctypedef npy_cdouble complex_t - * - * cdef inline object PyArray_MultiIterNew1(a): # <<<<<<<<<<<<<< - * return PyArray_MultiIterNew(1, a) - * - */ - -static CYTHON_INLINE PyObject *__pyx_f_5numpy_PyArray_MultiIterNew1(PyObject *__pyx_v_a) { - PyObject *__pyx_r = NULL; - __Pyx_RefNannyDeclarations - PyObject *__pyx_t_1 = NULL; - __Pyx_RefNannySetupContext("PyArray_MultiIterNew1", 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":776 - * - * cdef inline object PyArray_MultiIterNew1(a): - * return PyArray_MultiIterNew(1, a) # <<<<<<<<<<<<<< - * - * cdef inline object PyArray_MultiIterNew2(a, b): - */ - __Pyx_XDECREF(__pyx_r); - __pyx_t_1 = PyArray_MultiIterNew(1, ((void *)__pyx_v_a)); if (unlikely(!__pyx_t_1)) __PYX_ERR(1, 776, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_1); - __pyx_r = __pyx_t_1; - __pyx_t_1 = 0; - goto __pyx_L0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":775 - * ctypedef npy_cdouble complex_t - * - * cdef inline object PyArray_MultiIterNew1(a): # <<<<<<<<<<<<<< - * return PyArray_MultiIterNew(1, a) - * - */ - - /* function exit code */ - __pyx_L1_error:; - __Pyx_XDECREF(__pyx_t_1); - __Pyx_AddTraceback("numpy.PyArray_MultiIterNew1", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = 0; - __pyx_L0:; - __Pyx_XGIVEREF(__pyx_r); - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":778 - * return PyArray_MultiIterNew(1, a) - * - * cdef inline object PyArray_MultiIterNew2(a, b): # <<<<<<<<<<<<<< - * return PyArray_MultiIterNew(2, a, b) - * - */ - -static CYTHON_INLINE PyObject *__pyx_f_5numpy_PyArray_MultiIterNew2(PyObject *__pyx_v_a, PyObject *__pyx_v_b) { - PyObject *__pyx_r = NULL; - __Pyx_RefNannyDeclarations - PyObject *__pyx_t_1 = NULL; - __Pyx_RefNannySetupContext("PyArray_MultiIterNew2", 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":779 - * - * cdef inline object PyArray_MultiIterNew2(a, b): - * return PyArray_MultiIterNew(2, a, b) # <<<<<<<<<<<<<< - * - * cdef inline object PyArray_MultiIterNew3(a, b, c): - */ - __Pyx_XDECREF(__pyx_r); - __pyx_t_1 = PyArray_MultiIterNew(2, ((void *)__pyx_v_a), ((void *)__pyx_v_b)); if (unlikely(!__pyx_t_1)) __PYX_ERR(1, 779, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_1); - __pyx_r = __pyx_t_1; - __pyx_t_1 = 0; - goto __pyx_L0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":778 - * return PyArray_MultiIterNew(1, a) - * - * cdef inline object PyArray_MultiIterNew2(a, b): # <<<<<<<<<<<<<< - * return PyArray_MultiIterNew(2, a, b) - * - */ - - /* function exit code */ - __pyx_L1_error:; - __Pyx_XDECREF(__pyx_t_1); - __Pyx_AddTraceback("numpy.PyArray_MultiIterNew2", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = 0; - __pyx_L0:; - __Pyx_XGIVEREF(__pyx_r); - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":781 - * return PyArray_MultiIterNew(2, a, b) - * - * cdef inline object PyArray_MultiIterNew3(a, b, c): # <<<<<<<<<<<<<< - * return PyArray_MultiIterNew(3, a, b, c) - * - */ - -static CYTHON_INLINE PyObject *__pyx_f_5numpy_PyArray_MultiIterNew3(PyObject *__pyx_v_a, PyObject *__pyx_v_b, PyObject *__pyx_v_c) { - PyObject *__pyx_r = NULL; - __Pyx_RefNannyDeclarations - PyObject *__pyx_t_1 = NULL; - __Pyx_RefNannySetupContext("PyArray_MultiIterNew3", 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":782 - * - * cdef inline object PyArray_MultiIterNew3(a, b, c): - * return PyArray_MultiIterNew(3, a, b, c) # <<<<<<<<<<<<<< - * - * cdef inline object PyArray_MultiIterNew4(a, b, c, d): - */ - __Pyx_XDECREF(__pyx_r); - __pyx_t_1 = PyArray_MultiIterNew(3, ((void *)__pyx_v_a), ((void *)__pyx_v_b), ((void *)__pyx_v_c)); if (unlikely(!__pyx_t_1)) __PYX_ERR(1, 782, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_1); - __pyx_r = __pyx_t_1; - __pyx_t_1 = 0; - goto __pyx_L0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":781 - * return PyArray_MultiIterNew(2, a, b) - * - * cdef inline object PyArray_MultiIterNew3(a, b, c): # <<<<<<<<<<<<<< - * return PyArray_MultiIterNew(3, a, b, c) - * - */ - - /* function exit code */ - __pyx_L1_error:; - __Pyx_XDECREF(__pyx_t_1); - __Pyx_AddTraceback("numpy.PyArray_MultiIterNew3", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = 0; - __pyx_L0:; - __Pyx_XGIVEREF(__pyx_r); - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":784 - * return PyArray_MultiIterNew(3, a, b, c) - * - * cdef inline object PyArray_MultiIterNew4(a, b, c, d): # <<<<<<<<<<<<<< - * return PyArray_MultiIterNew(4, a, b, c, d) - * - */ - -static CYTHON_INLINE PyObject *__pyx_f_5numpy_PyArray_MultiIterNew4(PyObject *__pyx_v_a, PyObject *__pyx_v_b, PyObject *__pyx_v_c, PyObject *__pyx_v_d) { - PyObject *__pyx_r = NULL; - __Pyx_RefNannyDeclarations - PyObject *__pyx_t_1 = NULL; - __Pyx_RefNannySetupContext("PyArray_MultiIterNew4", 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":785 - * - * cdef inline object PyArray_MultiIterNew4(a, b, c, d): - * return PyArray_MultiIterNew(4, a, b, c, d) # <<<<<<<<<<<<<< - * - * cdef inline object PyArray_MultiIterNew5(a, b, c, d, e): - */ - __Pyx_XDECREF(__pyx_r); - __pyx_t_1 = PyArray_MultiIterNew(4, ((void *)__pyx_v_a), ((void *)__pyx_v_b), ((void *)__pyx_v_c), ((void *)__pyx_v_d)); if (unlikely(!__pyx_t_1)) __PYX_ERR(1, 785, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_1); - __pyx_r = __pyx_t_1; - __pyx_t_1 = 0; - goto __pyx_L0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":784 - * return PyArray_MultiIterNew(3, a, b, c) - * - * cdef inline object PyArray_MultiIterNew4(a, b, c, d): # <<<<<<<<<<<<<< - * return PyArray_MultiIterNew(4, a, b, c, d) - * - */ - - /* function exit code */ - __pyx_L1_error:; - __Pyx_XDECREF(__pyx_t_1); - __Pyx_AddTraceback("numpy.PyArray_MultiIterNew4", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = 0; - __pyx_L0:; - __Pyx_XGIVEREF(__pyx_r); - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":787 - * return PyArray_MultiIterNew(4, a, b, c, d) - * - * cdef inline object PyArray_MultiIterNew5(a, b, c, d, e): # <<<<<<<<<<<<<< - * return PyArray_MultiIterNew(5, a, b, c, d, e) - * - */ - -static CYTHON_INLINE PyObject *__pyx_f_5numpy_PyArray_MultiIterNew5(PyObject *__pyx_v_a, PyObject *__pyx_v_b, PyObject *__pyx_v_c, PyObject *__pyx_v_d, PyObject *__pyx_v_e) { - PyObject *__pyx_r = NULL; - __Pyx_RefNannyDeclarations - PyObject *__pyx_t_1 = NULL; - __Pyx_RefNannySetupContext("PyArray_MultiIterNew5", 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":788 - * - * cdef inline object PyArray_MultiIterNew5(a, b, c, d, e): - * return PyArray_MultiIterNew(5, a, b, c, d, e) # <<<<<<<<<<<<<< - * - * cdef inline tuple PyDataType_SHAPE(dtype d): - */ - __Pyx_XDECREF(__pyx_r); - __pyx_t_1 = PyArray_MultiIterNew(5, ((void *)__pyx_v_a), ((void *)__pyx_v_b), ((void *)__pyx_v_c), ((void *)__pyx_v_d), ((void *)__pyx_v_e)); if (unlikely(!__pyx_t_1)) __PYX_ERR(1, 788, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_1); - __pyx_r = __pyx_t_1; - __pyx_t_1 = 0; - goto __pyx_L0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":787 - * return PyArray_MultiIterNew(4, a, b, c, d) - * - * cdef inline object PyArray_MultiIterNew5(a, b, c, d, e): # <<<<<<<<<<<<<< - * return PyArray_MultiIterNew(5, a, b, c, d, e) - * - */ - - /* function exit code */ - __pyx_L1_error:; - __Pyx_XDECREF(__pyx_t_1); - __Pyx_AddTraceback("numpy.PyArray_MultiIterNew5", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = 0; - __pyx_L0:; - __Pyx_XGIVEREF(__pyx_r); - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":790 - * return PyArray_MultiIterNew(5, a, b, c, d, e) - * - * cdef inline tuple PyDataType_SHAPE(dtype d): # <<<<<<<<<<<<<< - * if PyDataType_HASSUBARRAY(d): - * return d.subarray.shape - */ - -static CYTHON_INLINE PyObject *__pyx_f_5numpy_PyDataType_SHAPE(PyArray_Descr *__pyx_v_d) { - PyObject *__pyx_r = NULL; - __Pyx_RefNannyDeclarations - int __pyx_t_1; - __Pyx_RefNannySetupContext("PyDataType_SHAPE", 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":791 - * - * cdef inline tuple PyDataType_SHAPE(dtype d): - * if PyDataType_HASSUBARRAY(d): # <<<<<<<<<<<<<< - * return d.subarray.shape - * else: - */ - __pyx_t_1 = (PyDataType_HASSUBARRAY(__pyx_v_d) != 0); - if (__pyx_t_1) { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":792 - * cdef inline tuple PyDataType_SHAPE(dtype d): - * if PyDataType_HASSUBARRAY(d): - * return d.subarray.shape # <<<<<<<<<<<<<< - * else: - * return () - */ - __Pyx_XDECREF(__pyx_r); - __Pyx_INCREF(((PyObject*)__pyx_v_d->subarray->shape)); - __pyx_r = ((PyObject*)__pyx_v_d->subarray->shape); - goto __pyx_L0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":791 - * - * cdef inline tuple PyDataType_SHAPE(dtype d): - * if PyDataType_HASSUBARRAY(d): # <<<<<<<<<<<<<< - * return d.subarray.shape - * else: - */ - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":794 - * return d.subarray.shape - * else: - * return () # <<<<<<<<<<<<<< - * - * cdef inline char* _util_dtypestring(dtype descr, char* f, char* end, int* offset) except NULL: - */ - /*else*/ { - __Pyx_XDECREF(__pyx_r); - __Pyx_INCREF(__pyx_empty_tuple); - __pyx_r = __pyx_empty_tuple; - goto __pyx_L0; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":790 - * return PyArray_MultiIterNew(5, a, b, c, d, e) - * - * cdef inline tuple PyDataType_SHAPE(dtype d): # <<<<<<<<<<<<<< - * if PyDataType_HASSUBARRAY(d): - * return d.subarray.shape - */ - - /* function exit code */ - __pyx_L0:; - __Pyx_XGIVEREF(__pyx_r); - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":796 - * return () - * - * cdef inline char* _util_dtypestring(dtype descr, char* f, char* end, int* offset) except NULL: # <<<<<<<<<<<<<< - * # Recursive utility function used in __getbuffer__ to get format - * # string. The new location in the format string is returned. - */ - -static CYTHON_INLINE char *__pyx_f_5numpy__util_dtypestring(PyArray_Descr *__pyx_v_descr, char *__pyx_v_f, char *__pyx_v_end, int *__pyx_v_offset) { - PyArray_Descr *__pyx_v_child = 0; - int __pyx_v_endian_detector; - int __pyx_v_little_endian; - PyObject *__pyx_v_fields = 0; - PyObject *__pyx_v_childname = NULL; - PyObject *__pyx_v_new_offset = NULL; - PyObject *__pyx_v_t = NULL; - char *__pyx_r; - __Pyx_RefNannyDeclarations - PyObject *__pyx_t_1 = NULL; - Py_ssize_t __pyx_t_2; - PyObject *__pyx_t_3 = NULL; - PyObject *__pyx_t_4 = NULL; - int __pyx_t_5; - int __pyx_t_6; - int __pyx_t_7; - long __pyx_t_8; - char *__pyx_t_9; - __Pyx_RefNannySetupContext("_util_dtypestring", 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":801 - * - * cdef dtype child - * cdef int endian_detector = 1 # <<<<<<<<<<<<<< - * cdef bint little_endian = ((&endian_detector)[0] != 0) - * cdef tuple fields - */ - __pyx_v_endian_detector = 1; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":802 - * cdef dtype child - * cdef int endian_detector = 1 - * cdef bint little_endian = ((&endian_detector)[0] != 0) # <<<<<<<<<<<<<< - * cdef tuple fields - * - */ - __pyx_v_little_endian = ((((char *)(&__pyx_v_endian_detector))[0]) != 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":805 - * cdef tuple fields - * - * for childname in descr.names: # <<<<<<<<<<<<<< - * fields = descr.fields[childname] - * child, new_offset = fields - */ - if (unlikely(__pyx_v_descr->names == Py_None)) { - PyErr_SetString(PyExc_TypeError, "'NoneType' object is not iterable"); - __PYX_ERR(1, 805, __pyx_L1_error) - } - __pyx_t_1 = __pyx_v_descr->names; __Pyx_INCREF(__pyx_t_1); __pyx_t_2 = 0; - for (;;) { - if (__pyx_t_2 >= PyTuple_GET_SIZE(__pyx_t_1)) break; - #if CYTHON_ASSUME_SAFE_MACROS && !CYTHON_AVOID_BORROWED_REFS - __pyx_t_3 = PyTuple_GET_ITEM(__pyx_t_1, __pyx_t_2); __Pyx_INCREF(__pyx_t_3); __pyx_t_2++; if (unlikely(0 < 0)) __PYX_ERR(1, 805, __pyx_L1_error) - #else - __pyx_t_3 = PySequence_ITEM(__pyx_t_1, __pyx_t_2); __pyx_t_2++; if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 805, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - #endif - __Pyx_XDECREF_SET(__pyx_v_childname, __pyx_t_3); - __pyx_t_3 = 0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":806 - * - * for childname in descr.names: - * fields = descr.fields[childname] # <<<<<<<<<<<<<< - * child, new_offset = fields - * - */ - if (unlikely(__pyx_v_descr->fields == Py_None)) { - PyErr_SetString(PyExc_TypeError, "'NoneType' object is not subscriptable"); - __PYX_ERR(1, 806, __pyx_L1_error) - } - __pyx_t_3 = __Pyx_PyDict_GetItem(__pyx_v_descr->fields, __pyx_v_childname); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 806, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - if (!(likely(PyTuple_CheckExact(__pyx_t_3))||((__pyx_t_3) == Py_None)||(PyErr_Format(PyExc_TypeError, "Expected %.16s, got %.200s", "tuple", Py_TYPE(__pyx_t_3)->tp_name), 0))) __PYX_ERR(1, 806, __pyx_L1_error) - __Pyx_XDECREF_SET(__pyx_v_fields, ((PyObject*)__pyx_t_3)); - __pyx_t_3 = 0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":807 - * for childname in descr.names: - * fields = descr.fields[childname] - * child, new_offset = fields # <<<<<<<<<<<<<< - * - * if (end - f) - (new_offset - offset[0]) < 15: - */ - if (likely(__pyx_v_fields != Py_None)) { - PyObject* sequence = __pyx_v_fields; - Py_ssize_t size = __Pyx_PySequence_SIZE(sequence); - if (unlikely(size != 2)) { - if (size > 2) __Pyx_RaiseTooManyValuesError(2); - else if (size >= 0) __Pyx_RaiseNeedMoreValuesError(size); - __PYX_ERR(1, 807, __pyx_L1_error) - } - #if CYTHON_ASSUME_SAFE_MACROS && !CYTHON_AVOID_BORROWED_REFS - __pyx_t_3 = PyTuple_GET_ITEM(sequence, 0); - __pyx_t_4 = PyTuple_GET_ITEM(sequence, 1); - __Pyx_INCREF(__pyx_t_3); - __Pyx_INCREF(__pyx_t_4); - #else - __pyx_t_3 = PySequence_ITEM(sequence, 0); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 807, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __pyx_t_4 = PySequence_ITEM(sequence, 1); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 807, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_4); - #endif - } else { - __Pyx_RaiseNoneNotIterableError(); __PYX_ERR(1, 807, __pyx_L1_error) - } - if (!(likely(((__pyx_t_3) == Py_None) || likely(__Pyx_TypeTest(__pyx_t_3, __pyx_ptype_5numpy_dtype))))) __PYX_ERR(1, 807, __pyx_L1_error) - __Pyx_XDECREF_SET(__pyx_v_child, ((PyArray_Descr *)__pyx_t_3)); - __pyx_t_3 = 0; - __Pyx_XDECREF_SET(__pyx_v_new_offset, __pyx_t_4); - __pyx_t_4 = 0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":809 - * child, new_offset = fields - * - * if (end - f) - (new_offset - offset[0]) < 15: # <<<<<<<<<<<<<< - * raise RuntimeError(u"Format string allocated too short, see comment in numpy.pxd") - * - */ - __pyx_t_4 = __Pyx_PyInt_From_int((__pyx_v_offset[0])); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 809, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_4); - __pyx_t_3 = PyNumber_Subtract(__pyx_v_new_offset, __pyx_t_4); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 809, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - __pyx_t_5 = __Pyx_PyInt_As_int(__pyx_t_3); if (unlikely((__pyx_t_5 == (int)-1) && PyErr_Occurred())) __PYX_ERR(1, 809, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __pyx_t_6 = ((((__pyx_v_end - __pyx_v_f) - ((int)__pyx_t_5)) < 15) != 0); - if (unlikely(__pyx_t_6)) { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":810 - * - * if (end - f) - (new_offset - offset[0]) < 15: - * raise RuntimeError(u"Format string allocated too short, see comment in numpy.pxd") # <<<<<<<<<<<<<< - * - * if ((child.byteorder == c'>' and little_endian) or - */ - __pyx_t_3 = __Pyx_PyObject_Call(__pyx_builtin_RuntimeError, __pyx_tuple__4, NULL); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 810, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __Pyx_Raise(__pyx_t_3, 0, 0, 0); - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __PYX_ERR(1, 810, __pyx_L1_error) - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":809 - * child, new_offset = fields - * - * if (end - f) - (new_offset - offset[0]) < 15: # <<<<<<<<<<<<<< - * raise RuntimeError(u"Format string allocated too short, see comment in numpy.pxd") - * - */ - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":812 - * raise RuntimeError(u"Format string allocated too short, see comment in numpy.pxd") - * - * if ((child.byteorder == c'>' and little_endian) or # <<<<<<<<<<<<<< - * (child.byteorder == c'<' and not little_endian)): - * raise ValueError(u"Non-native byte order not supported") - */ - __pyx_t_7 = ((__pyx_v_child->byteorder == '>') != 0); - if (!__pyx_t_7) { - goto __pyx_L8_next_or; - } else { - } - __pyx_t_7 = (__pyx_v_little_endian != 0); - if (!__pyx_t_7) { - } else { - __pyx_t_6 = __pyx_t_7; - goto __pyx_L7_bool_binop_done; - } - __pyx_L8_next_or:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":813 - * - * if ((child.byteorder == c'>' and little_endian) or - * (child.byteorder == c'<' and not little_endian)): # <<<<<<<<<<<<<< - * raise ValueError(u"Non-native byte order not supported") - * # One could encode it in the format string and have Cython - */ - __pyx_t_7 = ((__pyx_v_child->byteorder == '<') != 0); - if (__pyx_t_7) { - } else { - __pyx_t_6 = __pyx_t_7; - goto __pyx_L7_bool_binop_done; - } - __pyx_t_7 = ((!(__pyx_v_little_endian != 0)) != 0); - __pyx_t_6 = __pyx_t_7; - __pyx_L7_bool_binop_done:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":812 - * raise RuntimeError(u"Format string allocated too short, see comment in numpy.pxd") - * - * if ((child.byteorder == c'>' and little_endian) or # <<<<<<<<<<<<<< - * (child.byteorder == c'<' and not little_endian)): - * raise ValueError(u"Non-native byte order not supported") - */ - if (unlikely(__pyx_t_6)) { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":814 - * if ((child.byteorder == c'>' and little_endian) or - * (child.byteorder == c'<' and not little_endian)): - * raise ValueError(u"Non-native byte order not supported") # <<<<<<<<<<<<<< - * # One could encode it in the format string and have Cython - * # complain instead, BUT: < and > in format strings also imply - */ - __pyx_t_3 = __Pyx_PyObject_Call(__pyx_builtin_ValueError, __pyx_tuple__5, NULL); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 814, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __Pyx_Raise(__pyx_t_3, 0, 0, 0); - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __PYX_ERR(1, 814, __pyx_L1_error) - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":812 - * raise RuntimeError(u"Format string allocated too short, see comment in numpy.pxd") - * - * if ((child.byteorder == c'>' and little_endian) or # <<<<<<<<<<<<<< - * (child.byteorder == c'<' and not little_endian)): - * raise ValueError(u"Non-native byte order not supported") - */ - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":824 - * - * # Output padding bytes - * while offset[0] < new_offset: # <<<<<<<<<<<<<< - * f[0] = 120 # "x"; pad byte - * f += 1 - */ - while (1) { - __pyx_t_3 = __Pyx_PyInt_From_int((__pyx_v_offset[0])); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 824, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __pyx_t_4 = PyObject_RichCompare(__pyx_t_3, __pyx_v_new_offset, Py_LT); __Pyx_XGOTREF(__pyx_t_4); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 824, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_4); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 824, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - if (!__pyx_t_6) break; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":825 - * # Output padding bytes - * while offset[0] < new_offset: - * f[0] = 120 # "x"; pad byte # <<<<<<<<<<<<<< - * f += 1 - * offset[0] += 1 - */ - (__pyx_v_f[0]) = 0x78; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":826 - * while offset[0] < new_offset: - * f[0] = 120 # "x"; pad byte - * f += 1 # <<<<<<<<<<<<<< - * offset[0] += 1 - * - */ - __pyx_v_f = (__pyx_v_f + 1); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":827 - * f[0] = 120 # "x"; pad byte - * f += 1 - * offset[0] += 1 # <<<<<<<<<<<<<< - * - * offset[0] += child.itemsize - */ - __pyx_t_8 = 0; - (__pyx_v_offset[__pyx_t_8]) = ((__pyx_v_offset[__pyx_t_8]) + 1); - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":829 - * offset[0] += 1 - * - * offset[0] += child.itemsize # <<<<<<<<<<<<<< - * - * if not PyDataType_HASFIELDS(child): - */ - __pyx_t_8 = 0; - (__pyx_v_offset[__pyx_t_8]) = ((__pyx_v_offset[__pyx_t_8]) + __pyx_v_child->elsize); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":831 - * offset[0] += child.itemsize - * - * if not PyDataType_HASFIELDS(child): # <<<<<<<<<<<<<< - * t = child.type_num - * if end - f < 5: - */ - __pyx_t_6 = ((!(PyDataType_HASFIELDS(__pyx_v_child) != 0)) != 0); - if (__pyx_t_6) { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":832 - * - * if not PyDataType_HASFIELDS(child): - * t = child.type_num # <<<<<<<<<<<<<< - * if end - f < 5: - * raise RuntimeError(u"Format string allocated too short.") - */ - __pyx_t_4 = __Pyx_PyInt_From_int(__pyx_v_child->type_num); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 832, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_4); - __Pyx_XDECREF_SET(__pyx_v_t, __pyx_t_4); - __pyx_t_4 = 0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":833 - * if not PyDataType_HASFIELDS(child): - * t = child.type_num - * if end - f < 5: # <<<<<<<<<<<<<< - * raise RuntimeError(u"Format string allocated too short.") - * - */ - __pyx_t_6 = (((__pyx_v_end - __pyx_v_f) < 5) != 0); - if (unlikely(__pyx_t_6)) { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":834 - * t = child.type_num - * if end - f < 5: - * raise RuntimeError(u"Format string allocated too short.") # <<<<<<<<<<<<<< - * - * # Until ticket #99 is fixed, use integers to avoid warnings - */ - __pyx_t_4 = __Pyx_PyObject_Call(__pyx_builtin_RuntimeError, __pyx_tuple__6, NULL); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 834, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_4); - __Pyx_Raise(__pyx_t_4, 0, 0, 0); - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - __PYX_ERR(1, 834, __pyx_L1_error) - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":833 - * if not PyDataType_HASFIELDS(child): - * t = child.type_num - * if end - f < 5: # <<<<<<<<<<<<<< - * raise RuntimeError(u"Format string allocated too short.") - * - */ - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":837 - * - * # Until ticket #99 is fixed, use integers to avoid warnings - * if t == NPY_BYTE: f[0] = 98 #"b" # <<<<<<<<<<<<<< - * elif t == NPY_UBYTE: f[0] = 66 #"B" - * elif t == NPY_SHORT: f[0] = 104 #"h" - */ - __pyx_t_4 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_BYTE); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 837, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_4); - __pyx_t_3 = PyObject_RichCompare(__pyx_v_t, __pyx_t_4, Py_EQ); __Pyx_XGOTREF(__pyx_t_3); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 837, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_3); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 837, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 98; - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":838 - * # Until ticket #99 is fixed, use integers to avoid warnings - * if t == NPY_BYTE: f[0] = 98 #"b" - * elif t == NPY_UBYTE: f[0] = 66 #"B" # <<<<<<<<<<<<<< - * elif t == NPY_SHORT: f[0] = 104 #"h" - * elif t == NPY_USHORT: f[0] = 72 #"H" - */ - __pyx_t_3 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_UBYTE); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 838, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __pyx_t_4 = PyObject_RichCompare(__pyx_v_t, __pyx_t_3, Py_EQ); __Pyx_XGOTREF(__pyx_t_4); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 838, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_4); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 838, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 66; - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":839 - * if t == NPY_BYTE: f[0] = 98 #"b" - * elif t == NPY_UBYTE: f[0] = 66 #"B" - * elif t == NPY_SHORT: f[0] = 104 #"h" # <<<<<<<<<<<<<< - * elif t == NPY_USHORT: f[0] = 72 #"H" - * elif t == NPY_INT: f[0] = 105 #"i" - */ - __pyx_t_4 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_SHORT); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 839, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_4); - __pyx_t_3 = PyObject_RichCompare(__pyx_v_t, __pyx_t_4, Py_EQ); __Pyx_XGOTREF(__pyx_t_3); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 839, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_3); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 839, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 0x68; - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":840 - * elif t == NPY_UBYTE: f[0] = 66 #"B" - * elif t == NPY_SHORT: f[0] = 104 #"h" - * elif t == NPY_USHORT: f[0] = 72 #"H" # <<<<<<<<<<<<<< - * elif t == NPY_INT: f[0] = 105 #"i" - * elif t == NPY_UINT: f[0] = 73 #"I" - */ - __pyx_t_3 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_USHORT); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 840, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __pyx_t_4 = PyObject_RichCompare(__pyx_v_t, __pyx_t_3, Py_EQ); __Pyx_XGOTREF(__pyx_t_4); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 840, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_4); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 840, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 72; - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":841 - * elif t == NPY_SHORT: f[0] = 104 #"h" - * elif t == NPY_USHORT: f[0] = 72 #"H" - * elif t == NPY_INT: f[0] = 105 #"i" # <<<<<<<<<<<<<< - * elif t == NPY_UINT: f[0] = 73 #"I" - * elif t == NPY_LONG: f[0] = 108 #"l" - */ - __pyx_t_4 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_INT); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 841, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_4); - __pyx_t_3 = PyObject_RichCompare(__pyx_v_t, __pyx_t_4, Py_EQ); __Pyx_XGOTREF(__pyx_t_3); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 841, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_3); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 841, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 0x69; - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":842 - * elif t == NPY_USHORT: f[0] = 72 #"H" - * elif t == NPY_INT: f[0] = 105 #"i" - * elif t == NPY_UINT: f[0] = 73 #"I" # <<<<<<<<<<<<<< - * elif t == NPY_LONG: f[0] = 108 #"l" - * elif t == NPY_ULONG: f[0] = 76 #"L" - */ - __pyx_t_3 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_UINT); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 842, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __pyx_t_4 = PyObject_RichCompare(__pyx_v_t, __pyx_t_3, Py_EQ); __Pyx_XGOTREF(__pyx_t_4); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 842, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_4); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 842, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 73; - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":843 - * elif t == NPY_INT: f[0] = 105 #"i" - * elif t == NPY_UINT: f[0] = 73 #"I" - * elif t == NPY_LONG: f[0] = 108 #"l" # <<<<<<<<<<<<<< - * elif t == NPY_ULONG: f[0] = 76 #"L" - * elif t == NPY_LONGLONG: f[0] = 113 #"q" - */ - __pyx_t_4 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_LONG); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 843, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_4); - __pyx_t_3 = PyObject_RichCompare(__pyx_v_t, __pyx_t_4, Py_EQ); __Pyx_XGOTREF(__pyx_t_3); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 843, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_3); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 843, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 0x6C; - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":844 - * elif t == NPY_UINT: f[0] = 73 #"I" - * elif t == NPY_LONG: f[0] = 108 #"l" - * elif t == NPY_ULONG: f[0] = 76 #"L" # <<<<<<<<<<<<<< - * elif t == NPY_LONGLONG: f[0] = 113 #"q" - * elif t == NPY_ULONGLONG: f[0] = 81 #"Q" - */ - __pyx_t_3 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_ULONG); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 844, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __pyx_t_4 = PyObject_RichCompare(__pyx_v_t, __pyx_t_3, Py_EQ); __Pyx_XGOTREF(__pyx_t_4); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 844, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_4); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 844, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 76; - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":845 - * elif t == NPY_LONG: f[0] = 108 #"l" - * elif t == NPY_ULONG: f[0] = 76 #"L" - * elif t == NPY_LONGLONG: f[0] = 113 #"q" # <<<<<<<<<<<<<< - * elif t == NPY_ULONGLONG: f[0] = 81 #"Q" - * elif t == NPY_FLOAT: f[0] = 102 #"f" - */ - __pyx_t_4 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_LONGLONG); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 845, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_4); - __pyx_t_3 = PyObject_RichCompare(__pyx_v_t, __pyx_t_4, Py_EQ); __Pyx_XGOTREF(__pyx_t_3); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 845, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_3); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 845, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 0x71; - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":846 - * elif t == NPY_ULONG: f[0] = 76 #"L" - * elif t == NPY_LONGLONG: f[0] = 113 #"q" - * elif t == NPY_ULONGLONG: f[0] = 81 #"Q" # <<<<<<<<<<<<<< - * elif t == NPY_FLOAT: f[0] = 102 #"f" - * elif t == NPY_DOUBLE: f[0] = 100 #"d" - */ - __pyx_t_3 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_ULONGLONG); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 846, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __pyx_t_4 = PyObject_RichCompare(__pyx_v_t, __pyx_t_3, Py_EQ); __Pyx_XGOTREF(__pyx_t_4); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 846, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_4); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 846, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 81; - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":847 - * elif t == NPY_LONGLONG: f[0] = 113 #"q" - * elif t == NPY_ULONGLONG: f[0] = 81 #"Q" - * elif t == NPY_FLOAT: f[0] = 102 #"f" # <<<<<<<<<<<<<< - * elif t == NPY_DOUBLE: f[0] = 100 #"d" - * elif t == NPY_LONGDOUBLE: f[0] = 103 #"g" - */ - __pyx_t_4 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_FLOAT); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 847, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_4); - __pyx_t_3 = PyObject_RichCompare(__pyx_v_t, __pyx_t_4, Py_EQ); __Pyx_XGOTREF(__pyx_t_3); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 847, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_3); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 847, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 0x66; - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":848 - * elif t == NPY_ULONGLONG: f[0] = 81 #"Q" - * elif t == NPY_FLOAT: f[0] = 102 #"f" - * elif t == NPY_DOUBLE: f[0] = 100 #"d" # <<<<<<<<<<<<<< - * elif t == NPY_LONGDOUBLE: f[0] = 103 #"g" - * elif t == NPY_CFLOAT: f[0] = 90; f[1] = 102; f += 1 # Zf - */ - __pyx_t_3 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_DOUBLE); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 848, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __pyx_t_4 = PyObject_RichCompare(__pyx_v_t, __pyx_t_3, Py_EQ); __Pyx_XGOTREF(__pyx_t_4); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 848, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_4); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 848, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 0x64; - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":849 - * elif t == NPY_FLOAT: f[0] = 102 #"f" - * elif t == NPY_DOUBLE: f[0] = 100 #"d" - * elif t == NPY_LONGDOUBLE: f[0] = 103 #"g" # <<<<<<<<<<<<<< - * elif t == NPY_CFLOAT: f[0] = 90; f[1] = 102; f += 1 # Zf - * elif t == NPY_CDOUBLE: f[0] = 90; f[1] = 100; f += 1 # Zd - */ - __pyx_t_4 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_LONGDOUBLE); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 849, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_4); - __pyx_t_3 = PyObject_RichCompare(__pyx_v_t, __pyx_t_4, Py_EQ); __Pyx_XGOTREF(__pyx_t_3); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 849, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_3); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 849, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 0x67; - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":850 - * elif t == NPY_DOUBLE: f[0] = 100 #"d" - * elif t == NPY_LONGDOUBLE: f[0] = 103 #"g" - * elif t == NPY_CFLOAT: f[0] = 90; f[1] = 102; f += 1 # Zf # <<<<<<<<<<<<<< - * elif t == NPY_CDOUBLE: f[0] = 90; f[1] = 100; f += 1 # Zd - * elif t == NPY_CLONGDOUBLE: f[0] = 90; f[1] = 103; f += 1 # Zg - */ - __pyx_t_3 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_CFLOAT); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 850, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __pyx_t_4 = PyObject_RichCompare(__pyx_v_t, __pyx_t_3, Py_EQ); __Pyx_XGOTREF(__pyx_t_4); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 850, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_4); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 850, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 90; - (__pyx_v_f[1]) = 0x66; - __pyx_v_f = (__pyx_v_f + 1); - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":851 - * elif t == NPY_LONGDOUBLE: f[0] = 103 #"g" - * elif t == NPY_CFLOAT: f[0] = 90; f[1] = 102; f += 1 # Zf - * elif t == NPY_CDOUBLE: f[0] = 90; f[1] = 100; f += 1 # Zd # <<<<<<<<<<<<<< - * elif t == NPY_CLONGDOUBLE: f[0] = 90; f[1] = 103; f += 1 # Zg - * elif t == NPY_OBJECT: f[0] = 79 #"O" - */ - __pyx_t_4 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_CDOUBLE); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 851, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_4); - __pyx_t_3 = PyObject_RichCompare(__pyx_v_t, __pyx_t_4, Py_EQ); __Pyx_XGOTREF(__pyx_t_3); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 851, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_3); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 851, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 90; - (__pyx_v_f[1]) = 0x64; - __pyx_v_f = (__pyx_v_f + 1); - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":852 - * elif t == NPY_CFLOAT: f[0] = 90; f[1] = 102; f += 1 # Zf - * elif t == NPY_CDOUBLE: f[0] = 90; f[1] = 100; f += 1 # Zd - * elif t == NPY_CLONGDOUBLE: f[0] = 90; f[1] = 103; f += 1 # Zg # <<<<<<<<<<<<<< - * elif t == NPY_OBJECT: f[0] = 79 #"O" - * else: - */ - __pyx_t_3 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_CLONGDOUBLE); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 852, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __pyx_t_4 = PyObject_RichCompare(__pyx_v_t, __pyx_t_3, Py_EQ); __Pyx_XGOTREF(__pyx_t_4); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 852, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_4); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 852, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - if (__pyx_t_6) { - (__pyx_v_f[0]) = 90; - (__pyx_v_f[1]) = 0x67; - __pyx_v_f = (__pyx_v_f + 1); - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":853 - * elif t == NPY_CDOUBLE: f[0] = 90; f[1] = 100; f += 1 # Zd - * elif t == NPY_CLONGDOUBLE: f[0] = 90; f[1] = 103; f += 1 # Zg - * elif t == NPY_OBJECT: f[0] = 79 #"O" # <<<<<<<<<<<<<< - * else: - * raise ValueError(u"unknown dtype code in numpy.pxd (%d)" % t) - */ - __pyx_t_4 = __Pyx_PyInt_From_enum__NPY_TYPES(NPY_OBJECT); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 853, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_4); - __pyx_t_3 = PyObject_RichCompare(__pyx_v_t, __pyx_t_4, Py_EQ); __Pyx_XGOTREF(__pyx_t_3); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 853, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - __pyx_t_6 = __Pyx_PyObject_IsTrue(__pyx_t_3); if (unlikely(__pyx_t_6 < 0)) __PYX_ERR(1, 853, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - if (likely(__pyx_t_6)) { - (__pyx_v_f[0]) = 79; - goto __pyx_L15; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":855 - * elif t == NPY_OBJECT: f[0] = 79 #"O" - * else: - * raise ValueError(u"unknown dtype code in numpy.pxd (%d)" % t) # <<<<<<<<<<<<<< - * f += 1 - * else: - */ - /*else*/ { - __pyx_t_3 = PyUnicode_Format(__pyx_kp_u_unknown_dtype_code_in_numpy_pxd, __pyx_v_t); if (unlikely(!__pyx_t_3)) __PYX_ERR(1, 855, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_3); - __pyx_t_4 = __Pyx_PyObject_CallOneArg(__pyx_builtin_ValueError, __pyx_t_3); if (unlikely(!__pyx_t_4)) __PYX_ERR(1, 855, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_4); - __Pyx_DECREF(__pyx_t_3); __pyx_t_3 = 0; - __Pyx_Raise(__pyx_t_4, 0, 0, 0); - __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0; - __PYX_ERR(1, 855, __pyx_L1_error) - } - __pyx_L15:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":856 - * else: - * raise ValueError(u"unknown dtype code in numpy.pxd (%d)" % t) - * f += 1 # <<<<<<<<<<<<<< - * else: - * # Cython ignores struct boundary information ("T{...}"), - */ - __pyx_v_f = (__pyx_v_f + 1); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":831 - * offset[0] += child.itemsize - * - * if not PyDataType_HASFIELDS(child): # <<<<<<<<<<<<<< - * t = child.type_num - * if end - f < 5: - */ - goto __pyx_L13; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":860 - * # Cython ignores struct boundary information ("T{...}"), - * # so don't output it - * f = _util_dtypestring(child, f, end, offset) # <<<<<<<<<<<<<< - * return f - * - */ - /*else*/ { - __pyx_t_9 = __pyx_f_5numpy__util_dtypestring(__pyx_v_child, __pyx_v_f, __pyx_v_end, __pyx_v_offset); if (unlikely(__pyx_t_9 == ((char *)NULL))) __PYX_ERR(1, 860, __pyx_L1_error) - __pyx_v_f = __pyx_t_9; - } - __pyx_L13:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":805 - * cdef tuple fields - * - * for childname in descr.names: # <<<<<<<<<<<<<< - * fields = descr.fields[childname] - * child, new_offset = fields - */ - } - __Pyx_DECREF(__pyx_t_1); __pyx_t_1 = 0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":861 - * # so don't output it - * f = _util_dtypestring(child, f, end, offset) - * return f # <<<<<<<<<<<<<< - * - * - */ - __pyx_r = __pyx_v_f; - goto __pyx_L0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":796 - * return () - * - * cdef inline char* _util_dtypestring(dtype descr, char* f, char* end, int* offset) except NULL: # <<<<<<<<<<<<<< - * # Recursive utility function used in __getbuffer__ to get format - * # string. The new location in the format string is returned. - */ - - /* function exit code */ - __pyx_L1_error:; - __Pyx_XDECREF(__pyx_t_1); - __Pyx_XDECREF(__pyx_t_3); - __Pyx_XDECREF(__pyx_t_4); - __Pyx_AddTraceback("numpy._util_dtypestring", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = NULL; - __pyx_L0:; - __Pyx_XDECREF((PyObject *)__pyx_v_child); - __Pyx_XDECREF(__pyx_v_fields); - __Pyx_XDECREF(__pyx_v_childname); - __Pyx_XDECREF(__pyx_v_new_offset); - __Pyx_XDECREF(__pyx_v_t); - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":977 - * - * - * cdef inline void set_array_base(ndarray arr, object base): # <<<<<<<<<<<<<< - * cdef PyObject* baseptr - * if base is None: - */ - -static CYTHON_INLINE void __pyx_f_5numpy_set_array_base(PyArrayObject *__pyx_v_arr, PyObject *__pyx_v_base) { - PyObject *__pyx_v_baseptr; - __Pyx_RefNannyDeclarations - int __pyx_t_1; - int __pyx_t_2; - __Pyx_RefNannySetupContext("set_array_base", 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":979 - * cdef inline void set_array_base(ndarray arr, object base): - * cdef PyObject* baseptr - * if base is None: # <<<<<<<<<<<<<< - * baseptr = NULL - * else: - */ - __pyx_t_1 = (__pyx_v_base == Py_None); - __pyx_t_2 = (__pyx_t_1 != 0); - if (__pyx_t_2) { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":980 - * cdef PyObject* baseptr - * if base is None: - * baseptr = NULL # <<<<<<<<<<<<<< - * else: - * Py_INCREF(base) # important to do this before decref below! - */ - __pyx_v_baseptr = NULL; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":979 - * cdef inline void set_array_base(ndarray arr, object base): - * cdef PyObject* baseptr - * if base is None: # <<<<<<<<<<<<<< - * baseptr = NULL - * else: - */ - goto __pyx_L3; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":982 - * baseptr = NULL - * else: - * Py_INCREF(base) # important to do this before decref below! # <<<<<<<<<<<<<< - * baseptr = base - * Py_XDECREF(arr.base) - */ - /*else*/ { - Py_INCREF(__pyx_v_base); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":983 - * else: - * Py_INCREF(base) # important to do this before decref below! - * baseptr = base # <<<<<<<<<<<<<< - * Py_XDECREF(arr.base) - * arr.base = baseptr - */ - __pyx_v_baseptr = ((PyObject *)__pyx_v_base); - } - __pyx_L3:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":984 - * Py_INCREF(base) # important to do this before decref below! - * baseptr = base - * Py_XDECREF(arr.base) # <<<<<<<<<<<<<< - * arr.base = baseptr - * - */ - Py_XDECREF(__pyx_v_arr->base); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":985 - * baseptr = base - * Py_XDECREF(arr.base) - * arr.base = baseptr # <<<<<<<<<<<<<< - * - * cdef inline object get_array_base(ndarray arr): - */ - __pyx_v_arr->base = __pyx_v_baseptr; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":977 - * - * - * cdef inline void set_array_base(ndarray arr, object base): # <<<<<<<<<<<<<< - * cdef PyObject* baseptr - * if base is None: - */ - - /* function exit code */ - __Pyx_RefNannyFinishContext(); -} - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":987 - * arr.base = baseptr - * - * cdef inline object get_array_base(ndarray arr): # <<<<<<<<<<<<<< - * if arr.base is NULL: - * return None - */ - -static CYTHON_INLINE PyObject *__pyx_f_5numpy_get_array_base(PyArrayObject *__pyx_v_arr) { - PyObject *__pyx_r = NULL; - __Pyx_RefNannyDeclarations - int __pyx_t_1; - __Pyx_RefNannySetupContext("get_array_base", 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":988 - * - * cdef inline object get_array_base(ndarray arr): - * if arr.base is NULL: # <<<<<<<<<<<<<< - * return None - * else: - */ - __pyx_t_1 = ((__pyx_v_arr->base == NULL) != 0); - if (__pyx_t_1) { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":989 - * cdef inline object get_array_base(ndarray arr): - * if arr.base is NULL: - * return None # <<<<<<<<<<<<<< - * else: - * return arr.base - */ - __Pyx_XDECREF(__pyx_r); - __pyx_r = Py_None; __Pyx_INCREF(Py_None); - goto __pyx_L0; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":988 - * - * cdef inline object get_array_base(ndarray arr): - * if arr.base is NULL: # <<<<<<<<<<<<<< - * return None - * else: - */ - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":991 - * return None - * else: - * return arr.base # <<<<<<<<<<<<<< - * - * - */ - /*else*/ { - __Pyx_XDECREF(__pyx_r); - __Pyx_INCREF(((PyObject *)__pyx_v_arr->base)); - __pyx_r = ((PyObject *)__pyx_v_arr->base); - goto __pyx_L0; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":987 - * arr.base = baseptr - * - * cdef inline object get_array_base(ndarray arr): # <<<<<<<<<<<<<< - * if arr.base is NULL: - * return None - */ - - /* function exit code */ - __pyx_L0:; - __Pyx_XGIVEREF(__pyx_r); - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":996 - * # Versions of the import_* functions which are more suitable for - * # Cython code. - * cdef inline int import_array() except -1: # <<<<<<<<<<<<<< - * try: - * _import_array() - */ - -static CYTHON_INLINE int __pyx_f_5numpy_import_array(void) { - int __pyx_r; - __Pyx_RefNannyDeclarations - PyObject *__pyx_t_1 = NULL; - PyObject *__pyx_t_2 = NULL; - PyObject *__pyx_t_3 = NULL; - int __pyx_t_4; - PyObject *__pyx_t_5 = NULL; - PyObject *__pyx_t_6 = NULL; - PyObject *__pyx_t_7 = NULL; - PyObject *__pyx_t_8 = NULL; - __Pyx_RefNannySetupContext("import_array", 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":997 - * # Cython code. - * cdef inline int import_array() except -1: - * try: # <<<<<<<<<<<<<< - * _import_array() - * except Exception: - */ - { - __Pyx_PyThreadState_declare - __Pyx_PyThreadState_assign - __Pyx_ExceptionSave(&__pyx_t_1, &__pyx_t_2, &__pyx_t_3); - __Pyx_XGOTREF(__pyx_t_1); - __Pyx_XGOTREF(__pyx_t_2); - __Pyx_XGOTREF(__pyx_t_3); - /*try:*/ { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":998 - * cdef inline int import_array() except -1: - * try: - * _import_array() # <<<<<<<<<<<<<< - * except Exception: - * raise ImportError("numpy.core.multiarray failed to import") - */ - __pyx_t_4 = _import_array(); if (unlikely(__pyx_t_4 == ((int)-1))) __PYX_ERR(1, 998, __pyx_L3_error) - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":997 - * # Cython code. - * cdef inline int import_array() except -1: - * try: # <<<<<<<<<<<<<< - * _import_array() - * except Exception: - */ - } - __Pyx_XDECREF(__pyx_t_1); __pyx_t_1 = 0; - __Pyx_XDECREF(__pyx_t_2); __pyx_t_2 = 0; - __Pyx_XDECREF(__pyx_t_3); __pyx_t_3 = 0; - goto __pyx_L8_try_end; - __pyx_L3_error:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":999 - * try: - * _import_array() - * except Exception: # <<<<<<<<<<<<<< - * raise ImportError("numpy.core.multiarray failed to import") - * - */ - __pyx_t_4 = __Pyx_PyErr_ExceptionMatches(((PyObject *)(&((PyTypeObject*)PyExc_Exception)[0]))); - if (__pyx_t_4) { - __Pyx_AddTraceback("numpy.import_array", __pyx_clineno, __pyx_lineno, __pyx_filename); - if (__Pyx_GetException(&__pyx_t_5, &__pyx_t_6, &__pyx_t_7) < 0) __PYX_ERR(1, 999, __pyx_L5_except_error) - __Pyx_GOTREF(__pyx_t_5); - __Pyx_GOTREF(__pyx_t_6); - __Pyx_GOTREF(__pyx_t_7); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1000 - * _import_array() - * except Exception: - * raise ImportError("numpy.core.multiarray failed to import") # <<<<<<<<<<<<<< - * - * cdef inline int import_umath() except -1: - */ - __pyx_t_8 = __Pyx_PyObject_Call(__pyx_builtin_ImportError, __pyx_tuple__7, NULL); if (unlikely(!__pyx_t_8)) __PYX_ERR(1, 1000, __pyx_L5_except_error) - __Pyx_GOTREF(__pyx_t_8); - __Pyx_Raise(__pyx_t_8, 0, 0, 0); - __Pyx_DECREF(__pyx_t_8); __pyx_t_8 = 0; - __PYX_ERR(1, 1000, __pyx_L5_except_error) - } - goto __pyx_L5_except_error; - __pyx_L5_except_error:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":997 - * # Cython code. - * cdef inline int import_array() except -1: - * try: # <<<<<<<<<<<<<< - * _import_array() - * except Exception: - */ - __Pyx_XGIVEREF(__pyx_t_1); - __Pyx_XGIVEREF(__pyx_t_2); - __Pyx_XGIVEREF(__pyx_t_3); - __Pyx_ExceptionReset(__pyx_t_1, __pyx_t_2, __pyx_t_3); - goto __pyx_L1_error; - __pyx_L8_try_end:; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":996 - * # Versions of the import_* functions which are more suitable for - * # Cython code. - * cdef inline int import_array() except -1: # <<<<<<<<<<<<<< - * try: - * _import_array() - */ - - /* function exit code */ - __pyx_r = 0; - goto __pyx_L0; - __pyx_L1_error:; - __Pyx_XDECREF(__pyx_t_5); - __Pyx_XDECREF(__pyx_t_6); - __Pyx_XDECREF(__pyx_t_7); - __Pyx_XDECREF(__pyx_t_8); - __Pyx_AddTraceback("numpy.import_array", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = -1; - __pyx_L0:; - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1002 - * raise ImportError("numpy.core.multiarray failed to import") - * - * cdef inline int import_umath() except -1: # <<<<<<<<<<<<<< - * try: - * _import_umath() - */ - -static CYTHON_INLINE int __pyx_f_5numpy_import_umath(void) { - int __pyx_r; - __Pyx_RefNannyDeclarations - PyObject *__pyx_t_1 = NULL; - PyObject *__pyx_t_2 = NULL; - PyObject *__pyx_t_3 = NULL; - int __pyx_t_4; - PyObject *__pyx_t_5 = NULL; - PyObject *__pyx_t_6 = NULL; - PyObject *__pyx_t_7 = NULL; - PyObject *__pyx_t_8 = NULL; - __Pyx_RefNannySetupContext("import_umath", 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1003 - * - * cdef inline int import_umath() except -1: - * try: # <<<<<<<<<<<<<< - * _import_umath() - * except Exception: - */ - { - __Pyx_PyThreadState_declare - __Pyx_PyThreadState_assign - __Pyx_ExceptionSave(&__pyx_t_1, &__pyx_t_2, &__pyx_t_3); - __Pyx_XGOTREF(__pyx_t_1); - __Pyx_XGOTREF(__pyx_t_2); - __Pyx_XGOTREF(__pyx_t_3); - /*try:*/ { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1004 - * cdef inline int import_umath() except -1: - * try: - * _import_umath() # <<<<<<<<<<<<<< - * except Exception: - * raise ImportError("numpy.core.umath failed to import") - */ - __pyx_t_4 = _import_umath(); if (unlikely(__pyx_t_4 == ((int)-1))) __PYX_ERR(1, 1004, __pyx_L3_error) - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1003 - * - * cdef inline int import_umath() except -1: - * try: # <<<<<<<<<<<<<< - * _import_umath() - * except Exception: - */ - } - __Pyx_XDECREF(__pyx_t_1); __pyx_t_1 = 0; - __Pyx_XDECREF(__pyx_t_2); __pyx_t_2 = 0; - __Pyx_XDECREF(__pyx_t_3); __pyx_t_3 = 0; - goto __pyx_L8_try_end; - __pyx_L3_error:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1005 - * try: - * _import_umath() - * except Exception: # <<<<<<<<<<<<<< - * raise ImportError("numpy.core.umath failed to import") - * - */ - __pyx_t_4 = __Pyx_PyErr_ExceptionMatches(((PyObject *)(&((PyTypeObject*)PyExc_Exception)[0]))); - if (__pyx_t_4) { - __Pyx_AddTraceback("numpy.import_umath", __pyx_clineno, __pyx_lineno, __pyx_filename); - if (__Pyx_GetException(&__pyx_t_5, &__pyx_t_6, &__pyx_t_7) < 0) __PYX_ERR(1, 1005, __pyx_L5_except_error) - __Pyx_GOTREF(__pyx_t_5); - __Pyx_GOTREF(__pyx_t_6); - __Pyx_GOTREF(__pyx_t_7); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1006 - * _import_umath() - * except Exception: - * raise ImportError("numpy.core.umath failed to import") # <<<<<<<<<<<<<< - * - * cdef inline int import_ufunc() except -1: - */ - __pyx_t_8 = __Pyx_PyObject_Call(__pyx_builtin_ImportError, __pyx_tuple__8, NULL); if (unlikely(!__pyx_t_8)) __PYX_ERR(1, 1006, __pyx_L5_except_error) - __Pyx_GOTREF(__pyx_t_8); - __Pyx_Raise(__pyx_t_8, 0, 0, 0); - __Pyx_DECREF(__pyx_t_8); __pyx_t_8 = 0; - __PYX_ERR(1, 1006, __pyx_L5_except_error) - } - goto __pyx_L5_except_error; - __pyx_L5_except_error:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1003 - * - * cdef inline int import_umath() except -1: - * try: # <<<<<<<<<<<<<< - * _import_umath() - * except Exception: - */ - __Pyx_XGIVEREF(__pyx_t_1); - __Pyx_XGIVEREF(__pyx_t_2); - __Pyx_XGIVEREF(__pyx_t_3); - __Pyx_ExceptionReset(__pyx_t_1, __pyx_t_2, __pyx_t_3); - goto __pyx_L1_error; - __pyx_L8_try_end:; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1002 - * raise ImportError("numpy.core.multiarray failed to import") - * - * cdef inline int import_umath() except -1: # <<<<<<<<<<<<<< - * try: - * _import_umath() - */ - - /* function exit code */ - __pyx_r = 0; - goto __pyx_L0; - __pyx_L1_error:; - __Pyx_XDECREF(__pyx_t_5); - __Pyx_XDECREF(__pyx_t_6); - __Pyx_XDECREF(__pyx_t_7); - __Pyx_XDECREF(__pyx_t_8); - __Pyx_AddTraceback("numpy.import_umath", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = -1; - __pyx_L0:; - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1008 - * raise ImportError("numpy.core.umath failed to import") - * - * cdef inline int import_ufunc() except -1: # <<<<<<<<<<<<<< - * try: - * _import_umath() - */ - -static CYTHON_INLINE int __pyx_f_5numpy_import_ufunc(void) { - int __pyx_r; - __Pyx_RefNannyDeclarations - PyObject *__pyx_t_1 = NULL; - PyObject *__pyx_t_2 = NULL; - PyObject *__pyx_t_3 = NULL; - int __pyx_t_4; - PyObject *__pyx_t_5 = NULL; - PyObject *__pyx_t_6 = NULL; - PyObject *__pyx_t_7 = NULL; - PyObject *__pyx_t_8 = NULL; - __Pyx_RefNannySetupContext("import_ufunc", 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1009 - * - * cdef inline int import_ufunc() except -1: - * try: # <<<<<<<<<<<<<< - * _import_umath() - * except Exception: - */ - { - __Pyx_PyThreadState_declare - __Pyx_PyThreadState_assign - __Pyx_ExceptionSave(&__pyx_t_1, &__pyx_t_2, &__pyx_t_3); - __Pyx_XGOTREF(__pyx_t_1); - __Pyx_XGOTREF(__pyx_t_2); - __Pyx_XGOTREF(__pyx_t_3); - /*try:*/ { - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1010 - * cdef inline int import_ufunc() except -1: - * try: - * _import_umath() # <<<<<<<<<<<<<< - * except Exception: - * raise ImportError("numpy.core.umath failed to import") - */ - __pyx_t_4 = _import_umath(); if (unlikely(__pyx_t_4 == ((int)-1))) __PYX_ERR(1, 1010, __pyx_L3_error) - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1009 - * - * cdef inline int import_ufunc() except -1: - * try: # <<<<<<<<<<<<<< - * _import_umath() - * except Exception: - */ - } - __Pyx_XDECREF(__pyx_t_1); __pyx_t_1 = 0; - __Pyx_XDECREF(__pyx_t_2); __pyx_t_2 = 0; - __Pyx_XDECREF(__pyx_t_3); __pyx_t_3 = 0; - goto __pyx_L8_try_end; - __pyx_L3_error:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1011 - * try: - * _import_umath() - * except Exception: # <<<<<<<<<<<<<< - * raise ImportError("numpy.core.umath failed to import") - */ - __pyx_t_4 = __Pyx_PyErr_ExceptionMatches(((PyObject *)(&((PyTypeObject*)PyExc_Exception)[0]))); - if (__pyx_t_4) { - __Pyx_AddTraceback("numpy.import_ufunc", __pyx_clineno, __pyx_lineno, __pyx_filename); - if (__Pyx_GetException(&__pyx_t_5, &__pyx_t_6, &__pyx_t_7) < 0) __PYX_ERR(1, 1011, __pyx_L5_except_error) - __Pyx_GOTREF(__pyx_t_5); - __Pyx_GOTREF(__pyx_t_6); - __Pyx_GOTREF(__pyx_t_7); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1012 - * _import_umath() - * except Exception: - * raise ImportError("numpy.core.umath failed to import") # <<<<<<<<<<<<<< - */ - __pyx_t_8 = __Pyx_PyObject_Call(__pyx_builtin_ImportError, __pyx_tuple__9, NULL); if (unlikely(!__pyx_t_8)) __PYX_ERR(1, 1012, __pyx_L5_except_error) - __Pyx_GOTREF(__pyx_t_8); - __Pyx_Raise(__pyx_t_8, 0, 0, 0); - __Pyx_DECREF(__pyx_t_8); __pyx_t_8 = 0; - __PYX_ERR(1, 1012, __pyx_L5_except_error) - } - goto __pyx_L5_except_error; - __pyx_L5_except_error:; - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1009 - * - * cdef inline int import_ufunc() except -1: - * try: # <<<<<<<<<<<<<< - * _import_umath() - * except Exception: - */ - __Pyx_XGIVEREF(__pyx_t_1); - __Pyx_XGIVEREF(__pyx_t_2); - __Pyx_XGIVEREF(__pyx_t_3); - __Pyx_ExceptionReset(__pyx_t_1, __pyx_t_2, __pyx_t_3); - goto __pyx_L1_error; - __pyx_L8_try_end:; - } - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1008 - * raise ImportError("numpy.core.umath failed to import") - * - * cdef inline int import_ufunc() except -1: # <<<<<<<<<<<<<< - * try: - * _import_umath() - */ - - /* function exit code */ - __pyx_r = 0; - goto __pyx_L0; - __pyx_L1_error:; - __Pyx_XDECREF(__pyx_t_5); - __Pyx_XDECREF(__pyx_t_6); - __Pyx_XDECREF(__pyx_t_7); - __Pyx_XDECREF(__pyx_t_8); - __Pyx_AddTraceback("numpy.import_ufunc", __pyx_clineno, __pyx_lineno, __pyx_filename); - __pyx_r = -1; - __pyx_L0:; - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -/* "string.from_py":13 - * - * @cname("__pyx_convert_string_from_py_std__in_string") - * cdef string __pyx_convert_string_from_py_std__in_string(object o) except *: # <<<<<<<<<<<<<< - * cdef Py_ssize_t length - * cdef const char* data = __Pyx_PyObject_AsStringAndSize(o, &length) - */ - -static std::string __pyx_convert_string_from_py_std__in_string(PyObject *__pyx_v_o) { - Py_ssize_t __pyx_v_length; - char const *__pyx_v_data; - std::string __pyx_r; - __Pyx_RefNannyDeclarations - char const *__pyx_t_1; - __Pyx_RefNannySetupContext("__pyx_convert_string_from_py_std__in_string", 0); - - /* "string.from_py":15 - * cdef string __pyx_convert_string_from_py_std__in_string(object o) except *: - * cdef Py_ssize_t length - * cdef const char* data = __Pyx_PyObject_AsStringAndSize(o, &length) # <<<<<<<<<<<<<< - * return string(data, length) - * - */ - __pyx_t_1 = __Pyx_PyObject_AsStringAndSize(__pyx_v_o, (&__pyx_v_length)); if (unlikely(__pyx_t_1 == ((char const *)NULL))) __PYX_ERR(2, 15, __pyx_L1_error) - __pyx_v_data = __pyx_t_1; - - /* "string.from_py":16 - * cdef Py_ssize_t length - * cdef const char* data = __Pyx_PyObject_AsStringAndSize(o, &length) - * return string(data, length) # <<<<<<<<<<<<<< - * - * - */ - __pyx_r = std::string(__pyx_v_data, __pyx_v_length); - goto __pyx_L0; - - /* "string.from_py":13 - * - * @cname("__pyx_convert_string_from_py_std__in_string") - * cdef string __pyx_convert_string_from_py_std__in_string(object o) except *: # <<<<<<<<<<<<<< - * cdef Py_ssize_t length - * cdef const char* data = __Pyx_PyObject_AsStringAndSize(o, &length) - */ - - /* function exit code */ - __pyx_L1_error:; - __Pyx_AddTraceback("string.from_py.__pyx_convert_string_from_py_std__in_string", __pyx_clineno, __pyx_lineno, __pyx_filename); - __Pyx_pretend_to_initialize(&__pyx_r); - __pyx_L0:; - __Pyx_RefNannyFinishContext(); - return __pyx_r; -} - -static PyMethodDef __pyx_methods[] = { - {0, 0, 0, 0} -}; - -#if PY_MAJOR_VERSION >= 3 -#if CYTHON_PEP489_MULTI_PHASE_INIT -static PyObject* __pyx_pymod_create(PyObject *spec, PyModuleDef *def); /*proto*/ -static int __pyx_pymod_exec_mesh_core_cython(PyObject* module); /*proto*/ -static PyModuleDef_Slot __pyx_moduledef_slots[] = { - {Py_mod_create, (void*)__pyx_pymod_create}, - {Py_mod_exec, (void*)__pyx_pymod_exec_mesh_core_cython}, - {0, NULL} -}; -#endif - -static struct PyModuleDef __pyx_moduledef = { - PyModuleDef_HEAD_INIT, - "mesh_core_cython", - 0, /* m_doc */ - #if CYTHON_PEP489_MULTI_PHASE_INIT - 0, /* m_size */ - #else - -1, /* m_size */ - #endif - __pyx_methods /* m_methods */, - #if CYTHON_PEP489_MULTI_PHASE_INIT - __pyx_moduledef_slots, /* m_slots */ - #else - NULL, /* m_reload */ - #endif - NULL, /* m_traverse */ - NULL, /* m_clear */ - NULL /* m_free */ -}; -#endif - -static __Pyx_StringTabEntry __pyx_string_tab[] = { - {&__pyx_kp_u_Format_string_allocated_too_shor, __pyx_k_Format_string_allocated_too_shor, sizeof(__pyx_k_Format_string_allocated_too_shor), 0, 1, 0, 0}, - {&__pyx_kp_u_Format_string_allocated_too_shor_2, __pyx_k_Format_string_allocated_too_shor_2, sizeof(__pyx_k_Format_string_allocated_too_shor_2), 0, 1, 0, 0}, - {&__pyx_n_s_ImportError, __pyx_k_ImportError, sizeof(__pyx_k_ImportError), 0, 0, 1, 1}, - {&__pyx_kp_u_Non_native_byte_order_not_suppor, __pyx_k_Non_native_byte_order_not_suppor, sizeof(__pyx_k_Non_native_byte_order_not_suppor), 0, 1, 0, 0}, - {&__pyx_n_s_RuntimeError, __pyx_k_RuntimeError, sizeof(__pyx_k_RuntimeError), 0, 0, 1, 1}, - {&__pyx_n_s_ValueError, __pyx_k_ValueError, sizeof(__pyx_k_ValueError), 0, 0, 1, 1}, - {&__pyx_n_s_barycentric_weight, __pyx_k_barycentric_weight, sizeof(__pyx_k_barycentric_weight), 0, 0, 1, 1}, - {&__pyx_n_s_c, __pyx_k_c, sizeof(__pyx_k_c), 0, 0, 1, 1}, - {&__pyx_n_s_cline_in_traceback, __pyx_k_cline_in_traceback, sizeof(__pyx_k_cline_in_traceback), 0, 0, 1, 1}, - {&__pyx_n_s_colors, __pyx_k_colors, sizeof(__pyx_k_colors), 0, 0, 1, 1}, - {&__pyx_n_s_depth_buffer, __pyx_k_depth_buffer, sizeof(__pyx_k_depth_buffer), 0, 0, 1, 1}, - {&__pyx_n_s_filename, __pyx_k_filename, sizeof(__pyx_k_filename), 0, 0, 1, 1}, - {&__pyx_n_s_get_normal_core, __pyx_k_get_normal_core, sizeof(__pyx_k_get_normal_core), 0, 0, 1, 1}, - {&__pyx_n_s_h, __pyx_k_h, sizeof(__pyx_k_h), 0, 0, 1, 1}, - {&__pyx_n_s_image, __pyx_k_image, sizeof(__pyx_k_image), 0, 0, 1, 1}, - {&__pyx_n_s_import, __pyx_k_import, sizeof(__pyx_k_import), 0, 0, 1, 1}, - {&__pyx_kp_s_insightface_thirdparty_face3d_me, __pyx_k_insightface_thirdparty_face3d_me, sizeof(__pyx_k_insightface_thirdparty_face3d_me), 0, 0, 1, 0}, - {&__pyx_n_s_main, __pyx_k_main, sizeof(__pyx_k_main), 0, 0, 1, 1}, - {&__pyx_n_s_mapping_type, __pyx_k_mapping_type, sizeof(__pyx_k_mapping_type), 0, 0, 1, 1}, - {&__pyx_n_s_mesh_core_cython, __pyx_k_mesh_core_cython, sizeof(__pyx_k_mesh_core_cython), 0, 0, 1, 1}, - {&__pyx_n_s_mtl_name, __pyx_k_mtl_name, sizeof(__pyx_k_mtl_name), 0, 0, 1, 1}, - {&__pyx_kp_u_ndarray_is_not_C_contiguous, __pyx_k_ndarray_is_not_C_contiguous, sizeof(__pyx_k_ndarray_is_not_C_contiguous), 0, 1, 0, 0}, - {&__pyx_kp_u_ndarray_is_not_Fortran_contiguou, __pyx_k_ndarray_is_not_Fortran_contiguou, sizeof(__pyx_k_ndarray_is_not_Fortran_contiguou), 0, 1, 0, 0}, - {&__pyx_n_s_normal, __pyx_k_normal, sizeof(__pyx_k_normal), 0, 0, 1, 1}, - {&__pyx_n_s_np, __pyx_k_np, sizeof(__pyx_k_np), 0, 0, 1, 1}, - {&__pyx_n_s_ntexver, __pyx_k_ntexver, sizeof(__pyx_k_ntexver), 0, 0, 1, 1}, - {&__pyx_n_s_ntri, __pyx_k_ntri, sizeof(__pyx_k_ntri), 0, 0, 1, 1}, - {&__pyx_n_s_numpy, __pyx_k_numpy, sizeof(__pyx_k_numpy), 0, 0, 1, 1}, - {&__pyx_kp_s_numpy_core_multiarray_failed_to, __pyx_k_numpy_core_multiarray_failed_to, sizeof(__pyx_k_numpy_core_multiarray_failed_to), 0, 0, 1, 0}, - {&__pyx_kp_s_numpy_core_umath_failed_to_impor, __pyx_k_numpy_core_umath_failed_to_impor, sizeof(__pyx_k_numpy_core_umath_failed_to_impor), 0, 0, 1, 0}, - {&__pyx_n_s_nver, __pyx_k_nver, sizeof(__pyx_k_nver), 0, 0, 1, 1}, - {&__pyx_n_s_range, __pyx_k_range, sizeof(__pyx_k_range), 0, 0, 1, 1}, - {&__pyx_n_s_rasterize_triangles_core, __pyx_k_rasterize_triangles_core, sizeof(__pyx_k_rasterize_triangles_core), 0, 0, 1, 1}, - {&__pyx_n_s_render_colors_core, __pyx_k_render_colors_core, sizeof(__pyx_k_render_colors_core), 0, 0, 1, 1}, - {&__pyx_n_s_render_texture_core, __pyx_k_render_texture_core, sizeof(__pyx_k_render_texture_core), 0, 0, 1, 1}, - {&__pyx_n_s_test, __pyx_k_test, sizeof(__pyx_k_test), 0, 0, 1, 1}, - {&__pyx_n_s_tex_c, __pyx_k_tex_c, sizeof(__pyx_k_tex_c), 0, 0, 1, 1}, - {&__pyx_n_s_tex_coords, __pyx_k_tex_coords, sizeof(__pyx_k_tex_coords), 0, 0, 1, 1}, - {&__pyx_n_s_tex_h, __pyx_k_tex_h, sizeof(__pyx_k_tex_h), 0, 0, 1, 1}, - {&__pyx_n_s_tex_nver, __pyx_k_tex_nver, sizeof(__pyx_k_tex_nver), 0, 0, 1, 1}, - {&__pyx_n_s_tex_triangles, __pyx_k_tex_triangles, sizeof(__pyx_k_tex_triangles), 0, 0, 1, 1}, - {&__pyx_n_s_tex_w, __pyx_k_tex_w, sizeof(__pyx_k_tex_w), 0, 0, 1, 1}, - {&__pyx_n_s_texture, __pyx_k_texture, sizeof(__pyx_k_texture), 0, 0, 1, 1}, - {&__pyx_n_s_tri_normal, __pyx_k_tri_normal, sizeof(__pyx_k_tri_normal), 0, 0, 1, 1}, - {&__pyx_n_s_triangle_buffer, __pyx_k_triangle_buffer, sizeof(__pyx_k_triangle_buffer), 0, 0, 1, 1}, - {&__pyx_n_s_triangles, __pyx_k_triangles, sizeof(__pyx_k_triangles), 0, 0, 1, 1}, - {&__pyx_kp_u_unknown_dtype_code_in_numpy_pxd, __pyx_k_unknown_dtype_code_in_numpy_pxd, sizeof(__pyx_k_unknown_dtype_code_in_numpy_pxd), 0, 1, 0, 0}, - {&__pyx_n_s_uv_coords, __pyx_k_uv_coords, sizeof(__pyx_k_uv_coords), 0, 0, 1, 1}, - {&__pyx_n_s_vertices, __pyx_k_vertices, sizeof(__pyx_k_vertices), 0, 0, 1, 1}, - {&__pyx_n_s_w, __pyx_k_w, sizeof(__pyx_k_w), 0, 0, 1, 1}, - {&__pyx_n_s_write_obj_with_colors_texture_co, __pyx_k_write_obj_with_colors_texture_co, sizeof(__pyx_k_write_obj_with_colors_texture_co), 0, 0, 1, 1}, - {0, 0, 0, 0, 0, 0, 0} -}; -static int __Pyx_InitCachedBuiltins(void) { - __pyx_builtin_ValueError = __Pyx_GetBuiltinName(__pyx_n_s_ValueError); if (!__pyx_builtin_ValueError) __PYX_ERR(1, 229, __pyx_L1_error) - __pyx_builtin_range = __Pyx_GetBuiltinName(__pyx_n_s_range); if (!__pyx_builtin_range) __PYX_ERR(1, 242, __pyx_L1_error) - __pyx_builtin_RuntimeError = __Pyx_GetBuiltinName(__pyx_n_s_RuntimeError); if (!__pyx_builtin_RuntimeError) __PYX_ERR(1, 810, __pyx_L1_error) - __pyx_builtin_ImportError = __Pyx_GetBuiltinName(__pyx_n_s_ImportError); if (!__pyx_builtin_ImportError) __PYX_ERR(1, 1000, __pyx_L1_error) - return 0; - __pyx_L1_error:; - return -1; -} - -static int __Pyx_InitCachedConstants(void) { - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("__Pyx_InitCachedConstants", 0); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":229 - * if ((flags & pybuf.PyBUF_C_CONTIGUOUS == pybuf.PyBUF_C_CONTIGUOUS) - * and not PyArray_CHKFLAGS(self, NPY_C_CONTIGUOUS)): - * raise ValueError(u"ndarray is not C contiguous") # <<<<<<<<<<<<<< - * - * if ((flags & pybuf.PyBUF_F_CONTIGUOUS == pybuf.PyBUF_F_CONTIGUOUS) - */ - __pyx_tuple_ = PyTuple_Pack(1, __pyx_kp_u_ndarray_is_not_C_contiguous); if (unlikely(!__pyx_tuple_)) __PYX_ERR(1, 229, __pyx_L1_error) - __Pyx_GOTREF(__pyx_tuple_); - __Pyx_GIVEREF(__pyx_tuple_); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":233 - * if ((flags & pybuf.PyBUF_F_CONTIGUOUS == pybuf.PyBUF_F_CONTIGUOUS) - * and not PyArray_CHKFLAGS(self, NPY_F_CONTIGUOUS)): - * raise ValueError(u"ndarray is not Fortran contiguous") # <<<<<<<<<<<<<< - * - * info.buf = PyArray_DATA(self) - */ - __pyx_tuple__2 = PyTuple_Pack(1, __pyx_kp_u_ndarray_is_not_Fortran_contiguou); if (unlikely(!__pyx_tuple__2)) __PYX_ERR(1, 233, __pyx_L1_error) - __Pyx_GOTREF(__pyx_tuple__2); - __Pyx_GIVEREF(__pyx_tuple__2); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":263 - * if ((descr.byteorder == c'>' and little_endian) or - * (descr.byteorder == c'<' and not little_endian)): - * raise ValueError(u"Non-native byte order not supported") # <<<<<<<<<<<<<< - * if t == NPY_BYTE: f = "b" - * elif t == NPY_UBYTE: f = "B" - */ - __pyx_tuple__3 = PyTuple_Pack(1, __pyx_kp_u_Non_native_byte_order_not_suppor); if (unlikely(!__pyx_tuple__3)) __PYX_ERR(1, 263, __pyx_L1_error) - __Pyx_GOTREF(__pyx_tuple__3); - __Pyx_GIVEREF(__pyx_tuple__3); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":810 - * - * if (end - f) - (new_offset - offset[0]) < 15: - * raise RuntimeError(u"Format string allocated too short, see comment in numpy.pxd") # <<<<<<<<<<<<<< - * - * if ((child.byteorder == c'>' and little_endian) or - */ - __pyx_tuple__4 = PyTuple_Pack(1, __pyx_kp_u_Format_string_allocated_too_shor); if (unlikely(!__pyx_tuple__4)) __PYX_ERR(1, 810, __pyx_L1_error) - __Pyx_GOTREF(__pyx_tuple__4); - __Pyx_GIVEREF(__pyx_tuple__4); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":814 - * if ((child.byteorder == c'>' and little_endian) or - * (child.byteorder == c'<' and not little_endian)): - * raise ValueError(u"Non-native byte order not supported") # <<<<<<<<<<<<<< - * # One could encode it in the format string and have Cython - * # complain instead, BUT: < and > in format strings also imply - */ - __pyx_tuple__5 = PyTuple_Pack(1, __pyx_kp_u_Non_native_byte_order_not_suppor); if (unlikely(!__pyx_tuple__5)) __PYX_ERR(1, 814, __pyx_L1_error) - __Pyx_GOTREF(__pyx_tuple__5); - __Pyx_GIVEREF(__pyx_tuple__5); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":834 - * t = child.type_num - * if end - f < 5: - * raise RuntimeError(u"Format string allocated too short.") # <<<<<<<<<<<<<< - * - * # Until ticket #99 is fixed, use integers to avoid warnings - */ - __pyx_tuple__6 = PyTuple_Pack(1, __pyx_kp_u_Format_string_allocated_too_shor_2); if (unlikely(!__pyx_tuple__6)) __PYX_ERR(1, 834, __pyx_L1_error) - __Pyx_GOTREF(__pyx_tuple__6); - __Pyx_GIVEREF(__pyx_tuple__6); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1000 - * _import_array() - * except Exception: - * raise ImportError("numpy.core.multiarray failed to import") # <<<<<<<<<<<<<< - * - * cdef inline int import_umath() except -1: - */ - __pyx_tuple__7 = PyTuple_Pack(1, __pyx_kp_s_numpy_core_multiarray_failed_to); if (unlikely(!__pyx_tuple__7)) __PYX_ERR(1, 1000, __pyx_L1_error) - __Pyx_GOTREF(__pyx_tuple__7); - __Pyx_GIVEREF(__pyx_tuple__7); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1006 - * _import_umath() - * except Exception: - * raise ImportError("numpy.core.umath failed to import") # <<<<<<<<<<<<<< - * - * cdef inline int import_ufunc() except -1: - */ - __pyx_tuple__8 = PyTuple_Pack(1, __pyx_kp_s_numpy_core_umath_failed_to_impor); if (unlikely(!__pyx_tuple__8)) __PYX_ERR(1, 1006, __pyx_L1_error) - __Pyx_GOTREF(__pyx_tuple__8); - __Pyx_GIVEREF(__pyx_tuple__8); - - /* "../../../../jack_ssd/home/jack/anaconda3/lib/python3.6/site-packages/Cython/Includes/numpy/__init__.pxd":1012 - * _import_umath() - * except Exception: - * raise ImportError("numpy.core.umath failed to import") # <<<<<<<<<<<<<< - */ - __pyx_tuple__9 = PyTuple_Pack(1, __pyx_kp_s_numpy_core_umath_failed_to_impor); if (unlikely(!__pyx_tuple__9)) __PYX_ERR(1, 1012, __pyx_L1_error) - __Pyx_GOTREF(__pyx_tuple__9); - __Pyx_GIVEREF(__pyx_tuple__9); - - /* "mesh_core_cython.pyx":40 - * int nver, int ntri, int ntexver) - * - * def get_normal_core(np.ndarray[float, ndim=2, mode = "c"] normal not None, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] tri_normal not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - __pyx_tuple__10 = PyTuple_Pack(4, __pyx_n_s_normal, __pyx_n_s_tri_normal, __pyx_n_s_triangles, __pyx_n_s_ntri); if (unlikely(!__pyx_tuple__10)) __PYX_ERR(0, 40, __pyx_L1_error) - __Pyx_GOTREF(__pyx_tuple__10); - __Pyx_GIVEREF(__pyx_tuple__10); - __pyx_codeobj__11 = (PyObject*)__Pyx_PyCode_New(4, 0, 4, 0, CO_OPTIMIZED|CO_NEWLOCALS, __pyx_empty_bytes, __pyx_empty_tuple, __pyx_empty_tuple, __pyx_tuple__10, __pyx_empty_tuple, __pyx_empty_tuple, __pyx_kp_s_insightface_thirdparty_face3d_me, __pyx_n_s_get_normal_core, 40, __pyx_empty_bytes); if (unlikely(!__pyx_codeobj__11)) __PYX_ERR(0, 40, __pyx_L1_error) - - /* "mesh_core_cython.pyx":49 - * ntri) - * - * def rasterize_triangles_core( # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - __pyx_tuple__12 = PyTuple_Pack(9, __pyx_n_s_vertices, __pyx_n_s_triangles, __pyx_n_s_depth_buffer, __pyx_n_s_triangle_buffer, __pyx_n_s_barycentric_weight, __pyx_n_s_nver, __pyx_n_s_ntri, __pyx_n_s_h, __pyx_n_s_w); if (unlikely(!__pyx_tuple__12)) __PYX_ERR(0, 49, __pyx_L1_error) - __Pyx_GOTREF(__pyx_tuple__12); - __Pyx_GIVEREF(__pyx_tuple__12); - __pyx_codeobj__13 = (PyObject*)__Pyx_PyCode_New(9, 0, 9, 0, CO_OPTIMIZED|CO_NEWLOCALS, __pyx_empty_bytes, __pyx_empty_tuple, __pyx_empty_tuple, __pyx_tuple__12, __pyx_empty_tuple, __pyx_empty_tuple, __pyx_kp_s_insightface_thirdparty_face3d_me, __pyx_n_s_rasterize_triangles_core, 49, __pyx_empty_bytes); if (unlikely(!__pyx_codeobj__13)) __PYX_ERR(0, 49, __pyx_L1_error) - - /* "mesh_core_cython.pyx":64 - * h, w) - * - * def render_colors_core(np.ndarray[float, ndim=3, mode = "c"] image not None, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - __pyx_tuple__14 = PyTuple_Pack(10, __pyx_n_s_image, __pyx_n_s_vertices, __pyx_n_s_triangles, __pyx_n_s_colors, __pyx_n_s_depth_buffer, __pyx_n_s_nver, __pyx_n_s_ntri, __pyx_n_s_h, __pyx_n_s_w, __pyx_n_s_c); if (unlikely(!__pyx_tuple__14)) __PYX_ERR(0, 64, __pyx_L1_error) - __Pyx_GOTREF(__pyx_tuple__14); - __Pyx_GIVEREF(__pyx_tuple__14); - __pyx_codeobj__15 = (PyObject*)__Pyx_PyCode_New(10, 0, 10, 0, CO_OPTIMIZED|CO_NEWLOCALS, __pyx_empty_bytes, __pyx_empty_tuple, __pyx_empty_tuple, __pyx_tuple__14, __pyx_empty_tuple, __pyx_empty_tuple, __pyx_kp_s_insightface_thirdparty_face3d_me, __pyx_n_s_render_colors_core, 64, __pyx_empty_bytes); if (unlikely(!__pyx_codeobj__15)) __PYX_ERR(0, 64, __pyx_L1_error) - - /* "mesh_core_cython.pyx":79 - * h, w, c) - * - * def render_texture_core(np.ndarray[float, ndim=3, mode = "c"] image not None, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - __pyx_tuple__16 = PyTuple_Pack(17, __pyx_n_s_image, __pyx_n_s_vertices, __pyx_n_s_triangles, __pyx_n_s_texture, __pyx_n_s_tex_coords, __pyx_n_s_tex_triangles, __pyx_n_s_depth_buffer, __pyx_n_s_nver, __pyx_n_s_tex_nver, __pyx_n_s_ntri, __pyx_n_s_h, __pyx_n_s_w, __pyx_n_s_c, __pyx_n_s_tex_h, __pyx_n_s_tex_w, __pyx_n_s_tex_c, __pyx_n_s_mapping_type); if (unlikely(!__pyx_tuple__16)) __PYX_ERR(0, 79, __pyx_L1_error) - __Pyx_GOTREF(__pyx_tuple__16); - __Pyx_GIVEREF(__pyx_tuple__16); - __pyx_codeobj__17 = (PyObject*)__Pyx_PyCode_New(17, 0, 17, 0, CO_OPTIMIZED|CO_NEWLOCALS, __pyx_empty_bytes, __pyx_empty_tuple, __pyx_empty_tuple, __pyx_tuple__16, __pyx_empty_tuple, __pyx_empty_tuple, __pyx_kp_s_insightface_thirdparty_face3d_me, __pyx_n_s_render_texture_core, 79, __pyx_empty_bytes); if (unlikely(!__pyx_codeobj__17)) __PYX_ERR(0, 79, __pyx_L1_error) - - /* "mesh_core_cython.pyx":100 - * mapping_type) - * - * def write_obj_with_colors_texture_core(string filename, string mtl_name, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - __pyx_tuple__18 = PyTuple_Pack(9, __pyx_n_s_filename, __pyx_n_s_mtl_name, __pyx_n_s_vertices, __pyx_n_s_triangles, __pyx_n_s_colors, __pyx_n_s_uv_coords, __pyx_n_s_nver, __pyx_n_s_ntri, __pyx_n_s_ntexver); if (unlikely(!__pyx_tuple__18)) __PYX_ERR(0, 100, __pyx_L1_error) - __Pyx_GOTREF(__pyx_tuple__18); - __Pyx_GIVEREF(__pyx_tuple__18); - __pyx_codeobj__19 = (PyObject*)__Pyx_PyCode_New(9, 0, 9, 0, CO_OPTIMIZED|CO_NEWLOCALS, __pyx_empty_bytes, __pyx_empty_tuple, __pyx_empty_tuple, __pyx_tuple__18, __pyx_empty_tuple, __pyx_empty_tuple, __pyx_kp_s_insightface_thirdparty_face3d_me, __pyx_n_s_write_obj_with_colors_texture_co, 100, __pyx_empty_bytes); if (unlikely(!__pyx_codeobj__19)) __PYX_ERR(0, 100, __pyx_L1_error) - __Pyx_RefNannyFinishContext(); - return 0; - __pyx_L1_error:; - __Pyx_RefNannyFinishContext(); - return -1; -} - -static int __Pyx_InitGlobals(void) { - if (__Pyx_InitStrings(__pyx_string_tab) < 0) __PYX_ERR(0, 1, __pyx_L1_error); - return 0; - __pyx_L1_error:; - return -1; -} - -static int __Pyx_modinit_global_init_code(void); /*proto*/ -static int __Pyx_modinit_variable_export_code(void); /*proto*/ -static int __Pyx_modinit_function_export_code(void); /*proto*/ -static int __Pyx_modinit_type_init_code(void); /*proto*/ -static int __Pyx_modinit_type_import_code(void); /*proto*/ -static int __Pyx_modinit_variable_import_code(void); /*proto*/ -static int __Pyx_modinit_function_import_code(void); /*proto*/ - -static int __Pyx_modinit_global_init_code(void) { - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("__Pyx_modinit_global_init_code", 0); - /*--- Global init code ---*/ - __Pyx_RefNannyFinishContext(); - return 0; -} - -static int __Pyx_modinit_variable_export_code(void) { - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("__Pyx_modinit_variable_export_code", 0); - /*--- Variable export code ---*/ - __Pyx_RefNannyFinishContext(); - return 0; -} - -static int __Pyx_modinit_function_export_code(void) { - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("__Pyx_modinit_function_export_code", 0); - /*--- Function export code ---*/ - __Pyx_RefNannyFinishContext(); - return 0; -} - -static int __Pyx_modinit_type_init_code(void) { - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("__Pyx_modinit_type_init_code", 0); - /*--- Type init code ---*/ - __Pyx_RefNannyFinishContext(); - return 0; -} - -static int __Pyx_modinit_type_import_code(void) { - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("__Pyx_modinit_type_import_code", 0); - /*--- Type import code ---*/ - __pyx_ptype_7cpython_4type_type = __Pyx_ImportType(__Pyx_BUILTIN_MODULE_NAME, "type", - #if defined(PYPY_VERSION_NUM) && PYPY_VERSION_NUM < 0x050B0000 - sizeof(PyTypeObject), - #else - sizeof(PyHeapTypeObject), - #endif - 0); if (unlikely(!__pyx_ptype_7cpython_4type_type)) __PYX_ERR(3, 9, __pyx_L1_error) - __pyx_ptype_5numpy_dtype = __Pyx_ImportType("numpy", "dtype", sizeof(PyArray_Descr), 0); if (unlikely(!__pyx_ptype_5numpy_dtype)) __PYX_ERR(1, 164, __pyx_L1_error) - __pyx_ptype_5numpy_flatiter = __Pyx_ImportType("numpy", "flatiter", sizeof(PyArrayIterObject), 0); if (unlikely(!__pyx_ptype_5numpy_flatiter)) __PYX_ERR(1, 186, __pyx_L1_error) - __pyx_ptype_5numpy_broadcast = __Pyx_ImportType("numpy", "broadcast", sizeof(PyArrayMultiIterObject), 0); if (unlikely(!__pyx_ptype_5numpy_broadcast)) __PYX_ERR(1, 190, __pyx_L1_error) - __pyx_ptype_5numpy_ndarray = __Pyx_ImportType("numpy", "ndarray", sizeof(PyArrayObject), 0); if (unlikely(!__pyx_ptype_5numpy_ndarray)) __PYX_ERR(1, 199, __pyx_L1_error) - __pyx_ptype_5numpy_ufunc = __Pyx_ImportType("numpy", "ufunc", sizeof(PyUFuncObject), 0); if (unlikely(!__pyx_ptype_5numpy_ufunc)) __PYX_ERR(1, 872, __pyx_L1_error) - __Pyx_RefNannyFinishContext(); - return 0; - __pyx_L1_error:; - __Pyx_RefNannyFinishContext(); - return -1; -} - -static int __Pyx_modinit_variable_import_code(void) { - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("__Pyx_modinit_variable_import_code", 0); - /*--- Variable import code ---*/ - __Pyx_RefNannyFinishContext(); - return 0; -} - -static int __Pyx_modinit_function_import_code(void) { - __Pyx_RefNannyDeclarations - __Pyx_RefNannySetupContext("__Pyx_modinit_function_import_code", 0); - /*--- Function import code ---*/ - __Pyx_RefNannyFinishContext(); - return 0; -} - - -#if PY_MAJOR_VERSION < 3 -#ifdef CYTHON_NO_PYINIT_EXPORT -#define __Pyx_PyMODINIT_FUNC void -#else -#define __Pyx_PyMODINIT_FUNC PyMODINIT_FUNC -#endif -#else -#ifdef CYTHON_NO_PYINIT_EXPORT -#define __Pyx_PyMODINIT_FUNC PyObject * -#else -#define __Pyx_PyMODINIT_FUNC PyMODINIT_FUNC -#endif -#endif -#ifndef CYTHON_SMALL_CODE -#if defined(__clang__) - #define CYTHON_SMALL_CODE -#elif defined(__GNUC__) - #define CYTHON_SMALL_CODE __attribute__((optimize("Os"))) -#else - #define CYTHON_SMALL_CODE -#endif -#endif - - -#if PY_MAJOR_VERSION < 3 -__Pyx_PyMODINIT_FUNC initmesh_core_cython(void) CYTHON_SMALL_CODE; /*proto*/ -__Pyx_PyMODINIT_FUNC initmesh_core_cython(void) -#else -__Pyx_PyMODINIT_FUNC PyInit_mesh_core_cython(void) CYTHON_SMALL_CODE; /*proto*/ -__Pyx_PyMODINIT_FUNC PyInit_mesh_core_cython(void) -#if CYTHON_PEP489_MULTI_PHASE_INIT -{ - return PyModuleDef_Init(&__pyx_moduledef); -} -static int __Pyx_copy_spec_to_module(PyObject *spec, PyObject *moddict, const char* from_name, const char* to_name) { - PyObject *value = PyObject_GetAttrString(spec, from_name); - int result = 0; - if (likely(value)) { - result = PyDict_SetItemString(moddict, to_name, value); - Py_DECREF(value); - } else if (PyErr_ExceptionMatches(PyExc_AttributeError)) { - PyErr_Clear(); - } else { - result = -1; - } - return result; -} -static PyObject* __pyx_pymod_create(PyObject *spec, CYTHON_UNUSED PyModuleDef *def) { - PyObject *module = NULL, *moddict, *modname; - if (__pyx_m) - return __Pyx_NewRef(__pyx_m); - modname = PyObject_GetAttrString(spec, "name"); - if (unlikely(!modname)) goto bad; - module = PyModule_NewObject(modname); - Py_DECREF(modname); - if (unlikely(!module)) goto bad; - moddict = PyModule_GetDict(module); - if (unlikely(!moddict)) goto bad; - if (unlikely(__Pyx_copy_spec_to_module(spec, moddict, "loader", "__loader__") < 0)) goto bad; - if (unlikely(__Pyx_copy_spec_to_module(spec, moddict, "origin", "__file__") < 0)) goto bad; - if (unlikely(__Pyx_copy_spec_to_module(spec, moddict, "parent", "__package__") < 0)) goto bad; - if (unlikely(__Pyx_copy_spec_to_module(spec, moddict, "submodule_search_locations", "__path__") < 0)) goto bad; - return module; -bad: - Py_XDECREF(module); - return NULL; -} - - -static int __pyx_pymod_exec_mesh_core_cython(PyObject *__pyx_pyinit_module) -#endif -#endif -{ - PyObject *__pyx_t_1 = NULL; - int __pyx_t_2; - __Pyx_RefNannyDeclarations - #if CYTHON_PEP489_MULTI_PHASE_INIT - if (__pyx_m && __pyx_m == __pyx_pyinit_module) return 0; - #elif PY_MAJOR_VERSION >= 3 - if (__pyx_m) return __Pyx_NewRef(__pyx_m); - #endif - #if CYTHON_REFNANNY -__Pyx_RefNanny = __Pyx_RefNannyImportAPI("refnanny"); -if (!__Pyx_RefNanny) { - PyErr_Clear(); - __Pyx_RefNanny = __Pyx_RefNannyImportAPI("Cython.Runtime.refnanny"); - if (!__Pyx_RefNanny) - Py_FatalError("failed to import 'refnanny' module"); -} -#endif - __Pyx_RefNannySetupContext("__Pyx_PyMODINIT_FUNC PyInit_mesh_core_cython(void)", 0); - if (__Pyx_check_binary_version() < 0) __PYX_ERR(0, 1, __pyx_L1_error) - __pyx_empty_tuple = PyTuple_New(0); if (unlikely(!__pyx_empty_tuple)) __PYX_ERR(0, 1, __pyx_L1_error) - __pyx_empty_bytes = PyBytes_FromStringAndSize("", 0); if (unlikely(!__pyx_empty_bytes)) __PYX_ERR(0, 1, __pyx_L1_error) - __pyx_empty_unicode = PyUnicode_FromStringAndSize("", 0); if (unlikely(!__pyx_empty_unicode)) __PYX_ERR(0, 1, __pyx_L1_error) - #ifdef __Pyx_CyFunction_USED - if (__pyx_CyFunction_init() < 0) __PYX_ERR(0, 1, __pyx_L1_error) - #endif - #ifdef __Pyx_FusedFunction_USED - if (__pyx_FusedFunction_init() < 0) __PYX_ERR(0, 1, __pyx_L1_error) - #endif - #ifdef __Pyx_Coroutine_USED - if (__pyx_Coroutine_init() < 0) __PYX_ERR(0, 1, __pyx_L1_error) - #endif - #ifdef __Pyx_Generator_USED - if (__pyx_Generator_init() < 0) __PYX_ERR(0, 1, __pyx_L1_error) - #endif - #ifdef __Pyx_AsyncGen_USED - if (__pyx_AsyncGen_init() < 0) __PYX_ERR(0, 1, __pyx_L1_error) - #endif - #ifdef __Pyx_StopAsyncIteration_USED - if (__pyx_StopAsyncIteration_init() < 0) __PYX_ERR(0, 1, __pyx_L1_error) - #endif - /*--- Library function declarations ---*/ - /*--- Threads initialization code ---*/ - #if defined(__PYX_FORCE_INIT_THREADS) && __PYX_FORCE_INIT_THREADS - #ifdef WITH_THREAD /* Python build with threading support? */ - PyEval_InitThreads(); - #endif - #endif - /*--- Module creation code ---*/ - #if CYTHON_PEP489_MULTI_PHASE_INIT - __pyx_m = __pyx_pyinit_module; - Py_INCREF(__pyx_m); - #else - #if PY_MAJOR_VERSION < 3 - __pyx_m = Py_InitModule4("mesh_core_cython", __pyx_methods, 0, 0, PYTHON_API_VERSION); Py_XINCREF(__pyx_m); - #else - __pyx_m = PyModule_Create(&__pyx_moduledef); - #endif - if (unlikely(!__pyx_m)) __PYX_ERR(0, 1, __pyx_L1_error) - #endif - __pyx_d = PyModule_GetDict(__pyx_m); if (unlikely(!__pyx_d)) __PYX_ERR(0, 1, __pyx_L1_error) - Py_INCREF(__pyx_d); - __pyx_b = PyImport_AddModule(__Pyx_BUILTIN_MODULE_NAME); if (unlikely(!__pyx_b)) __PYX_ERR(0, 1, __pyx_L1_error) - __pyx_cython_runtime = PyImport_AddModule((char *) "cython_runtime"); if (unlikely(!__pyx_cython_runtime)) __PYX_ERR(0, 1, __pyx_L1_error) - #if CYTHON_COMPILING_IN_PYPY - Py_INCREF(__pyx_b); - #endif - if (PyObject_SetAttrString(__pyx_m, "__builtins__", __pyx_b) < 0) __PYX_ERR(0, 1, __pyx_L1_error); - /*--- Initialize various global constants etc. ---*/ - if (__Pyx_InitGlobals() < 0) __PYX_ERR(0, 1, __pyx_L1_error) - #if PY_MAJOR_VERSION < 3 && (__PYX_DEFAULT_STRING_ENCODING_IS_ASCII || __PYX_DEFAULT_STRING_ENCODING_IS_DEFAULT) - if (__Pyx_init_sys_getdefaultencoding_params() < 0) __PYX_ERR(0, 1, __pyx_L1_error) - #endif - if (__pyx_module_is_main_mesh_core_cython) { - if (PyObject_SetAttrString(__pyx_m, "__name__", __pyx_n_s_main) < 0) __PYX_ERR(0, 1, __pyx_L1_error) - } - #if PY_MAJOR_VERSION >= 3 - { - PyObject *modules = PyImport_GetModuleDict(); if (unlikely(!modules)) __PYX_ERR(0, 1, __pyx_L1_error) - if (!PyDict_GetItemString(modules, "mesh_core_cython")) { - if (unlikely(PyDict_SetItemString(modules, "mesh_core_cython", __pyx_m) < 0)) __PYX_ERR(0, 1, __pyx_L1_error) - } - } - #endif - /*--- Builtin init code ---*/ - if (__Pyx_InitCachedBuiltins() < 0) __PYX_ERR(0, 1, __pyx_L1_error) - /*--- Constants init code ---*/ - if (__Pyx_InitCachedConstants() < 0) __PYX_ERR(0, 1, __pyx_L1_error) - /*--- Global type/function init code ---*/ - (void)__Pyx_modinit_global_init_code(); - (void)__Pyx_modinit_variable_export_code(); - (void)__Pyx_modinit_function_export_code(); - (void)__Pyx_modinit_type_init_code(); - if (unlikely(__Pyx_modinit_type_import_code() != 0)) goto __pyx_L1_error; - (void)__Pyx_modinit_variable_import_code(); - (void)__Pyx_modinit_function_import_code(); - /*--- Execution code ---*/ - #if defined(__Pyx_Generator_USED) || defined(__Pyx_Coroutine_USED) - if (__Pyx_patch_abc() < 0) __PYX_ERR(0, 1, __pyx_L1_error) - #endif - - /* "mesh_core_cython.pyx":1 - * import numpy as np # <<<<<<<<<<<<<< - * cimport numpy as np - * from libcpp.string cimport string - */ - __pyx_t_1 = __Pyx_Import(__pyx_n_s_numpy, 0, -1); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 1, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_1); - if (PyDict_SetItem(__pyx_d, __pyx_n_s_np, __pyx_t_1) < 0) __PYX_ERR(0, 1, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_1); __pyx_t_1 = 0; - - /* "mesh_core_cython.pyx":6 - * - * # use the Numpy-C-API from Cython - * np.import_array() # <<<<<<<<<<<<<< - * - * # cdefine the signature of our c function - */ - __pyx_t_2 = __pyx_f_5numpy_import_array(); if (unlikely(__pyx_t_2 == ((int)-1))) __PYX_ERR(0, 6, __pyx_L1_error) - - /* "mesh_core_cython.pyx":40 - * int nver, int ntri, int ntexver) - * - * def get_normal_core(np.ndarray[float, ndim=2, mode = "c"] normal not None, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] tri_normal not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - __pyx_t_1 = PyCFunction_NewEx(&__pyx_mdef_16mesh_core_cython_1get_normal_core, NULL, __pyx_n_s_mesh_core_cython); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 40, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_1); - if (PyDict_SetItem(__pyx_d, __pyx_n_s_get_normal_core, __pyx_t_1) < 0) __PYX_ERR(0, 40, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_1); __pyx_t_1 = 0; - - /* "mesh_core_cython.pyx":49 - * ntri) - * - * def rasterize_triangles_core( # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - __pyx_t_1 = PyCFunction_NewEx(&__pyx_mdef_16mesh_core_cython_3rasterize_triangles_core, NULL, __pyx_n_s_mesh_core_cython); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 49, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_1); - if (PyDict_SetItem(__pyx_d, __pyx_n_s_rasterize_triangles_core, __pyx_t_1) < 0) __PYX_ERR(0, 49, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_1); __pyx_t_1 = 0; - - /* "mesh_core_cython.pyx":64 - * h, w) - * - * def render_colors_core(np.ndarray[float, ndim=3, mode = "c"] image not None, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - __pyx_t_1 = PyCFunction_NewEx(&__pyx_mdef_16mesh_core_cython_5render_colors_core, NULL, __pyx_n_s_mesh_core_cython); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 64, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_1); - if (PyDict_SetItem(__pyx_d, __pyx_n_s_render_colors_core, __pyx_t_1) < 0) __PYX_ERR(0, 64, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_1); __pyx_t_1 = 0; - - /* "mesh_core_cython.pyx":79 - * h, w, c) - * - * def render_texture_core(np.ndarray[float, ndim=3, mode = "c"] image not None, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - __pyx_t_1 = PyCFunction_NewEx(&__pyx_mdef_16mesh_core_cython_7render_texture_core, NULL, __pyx_n_s_mesh_core_cython); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 79, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_1); - if (PyDict_SetItem(__pyx_d, __pyx_n_s_render_texture_core, __pyx_t_1) < 0) __PYX_ERR(0, 79, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_1); __pyx_t_1 = 0; - - /* "mesh_core_cython.pyx":100 - * mapping_type) - * - * def write_obj_with_colors_texture_core(string filename, string mtl_name, # <<<<<<<<<<<<<< - * np.ndarray[float, ndim=2, mode = "c"] vertices not None, - * np.ndarray[int, ndim=2, mode="c"] triangles not None, - */ - __pyx_t_1 = PyCFunction_NewEx(&__pyx_mdef_16mesh_core_cython_9write_obj_with_colors_texture_core, NULL, __pyx_n_s_mesh_core_cython); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 100, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_1); - if (PyDict_SetItem(__pyx_d, __pyx_n_s_write_obj_with_colors_texture_co, __pyx_t_1) < 0) __PYX_ERR(0, 100, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_1); __pyx_t_1 = 0; - - /* "mesh_core_cython.pyx":1 - * import numpy as np # <<<<<<<<<<<<<< - * cimport numpy as np - * from libcpp.string cimport string - */ - __pyx_t_1 = __Pyx_PyDict_NewPresized(0); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 1, __pyx_L1_error) - __Pyx_GOTREF(__pyx_t_1); - if (PyDict_SetItem(__pyx_d, __pyx_n_s_test, __pyx_t_1) < 0) __PYX_ERR(0, 1, __pyx_L1_error) - __Pyx_DECREF(__pyx_t_1); __pyx_t_1 = 0; - - /* "string.from_py":13 - * - * @cname("__pyx_convert_string_from_py_std__in_string") - * cdef string __pyx_convert_string_from_py_std__in_string(object o) except *: # <<<<<<<<<<<<<< - * cdef Py_ssize_t length - * cdef const char* data = __Pyx_PyObject_AsStringAndSize(o, &length) - */ - - /*--- Wrapped vars code ---*/ - - goto __pyx_L0; - __pyx_L1_error:; - __Pyx_XDECREF(__pyx_t_1); - if (__pyx_m) { - if (__pyx_d) { - __Pyx_AddTraceback("init mesh_core_cython", 0, __pyx_lineno, __pyx_filename); - } - Py_DECREF(__pyx_m); __pyx_m = 0; - } else if (!PyErr_Occurred()) { - PyErr_SetString(PyExc_ImportError, "init mesh_core_cython"); - } - __pyx_L0:; - __Pyx_RefNannyFinishContext(); - #if CYTHON_PEP489_MULTI_PHASE_INIT - return (__pyx_m != NULL) ? 0 : -1; - #elif PY_MAJOR_VERSION >= 3 - return __pyx_m; - #else - return; - #endif -} - -/* --- Runtime support code --- */ -/* Refnanny */ -#if CYTHON_REFNANNY -static __Pyx_RefNannyAPIStruct *__Pyx_RefNannyImportAPI(const char *modname) { - PyObject *m = NULL, *p = NULL; - void *r = NULL; - m = PyImport_ImportModule((char *)modname); - if (!m) goto end; - p = PyObject_GetAttrString(m, (char *)"RefNannyAPI"); - if (!p) goto end; - r = PyLong_AsVoidPtr(p); -end: - Py_XDECREF(p); - Py_XDECREF(m); - return (__Pyx_RefNannyAPIStruct *)r; -} -#endif - -/* RaiseArgTupleInvalid */ -static void __Pyx_RaiseArgtupleInvalid( - const char* func_name, - int exact, - Py_ssize_t num_min, - Py_ssize_t num_max, - Py_ssize_t num_found) -{ - Py_ssize_t num_expected; - const char *more_or_less; - if (num_found < num_min) { - num_expected = num_min; - more_or_less = "at least"; - } else { - num_expected = num_max; - more_or_less = "at most"; - } - if (exact) { - more_or_less = "exactly"; - } - PyErr_Format(PyExc_TypeError, - "%.200s() takes %.8s %" CYTHON_FORMAT_SSIZE_T "d positional argument%.1s (%" CYTHON_FORMAT_SSIZE_T "d given)", - func_name, more_or_less, num_expected, - (num_expected == 1) ? "" : "s", num_found); -} - -/* RaiseDoubleKeywords */ -static void __Pyx_RaiseDoubleKeywordsError( - const char* func_name, - PyObject* kw_name) -{ - PyErr_Format(PyExc_TypeError, - #if PY_MAJOR_VERSION >= 3 - "%s() got multiple values for keyword argument '%U'", func_name, kw_name); - #else - "%s() got multiple values for keyword argument '%s'", func_name, - PyString_AsString(kw_name)); - #endif -} - -/* ParseKeywords */ -static int __Pyx_ParseOptionalKeywords( - PyObject *kwds, - PyObject **argnames[], - PyObject *kwds2, - PyObject *values[], - Py_ssize_t num_pos_args, - const char* function_name) -{ - PyObject *key = 0, *value = 0; - Py_ssize_t pos = 0; - PyObject*** name; - PyObject*** first_kw_arg = argnames + num_pos_args; - while (PyDict_Next(kwds, &pos, &key, &value)) { - name = first_kw_arg; - while (*name && (**name != key)) name++; - if (*name) { - values[name-argnames] = value; - continue; - } - name = first_kw_arg; - #if PY_MAJOR_VERSION < 3 - if (likely(PyString_CheckExact(key)) || likely(PyString_Check(key))) { - while (*name) { - if ((CYTHON_COMPILING_IN_PYPY || PyString_GET_SIZE(**name) == PyString_GET_SIZE(key)) - && _PyString_Eq(**name, key)) { - values[name-argnames] = value; - break; - } - name++; - } - if (*name) continue; - else { - PyObject*** argname = argnames; - while (argname != first_kw_arg) { - if ((**argname == key) || ( - (CYTHON_COMPILING_IN_PYPY || PyString_GET_SIZE(**argname) == PyString_GET_SIZE(key)) - && _PyString_Eq(**argname, key))) { - goto arg_passed_twice; - } - argname++; - } - } - } else - #endif - if (likely(PyUnicode_Check(key))) { - while (*name) { - int cmp = (**name == key) ? 0 : - #if !CYTHON_COMPILING_IN_PYPY && PY_MAJOR_VERSION >= 3 - (PyUnicode_GET_SIZE(**name) != PyUnicode_GET_SIZE(key)) ? 1 : - #endif - PyUnicode_Compare(**name, key); - if (cmp < 0 && unlikely(PyErr_Occurred())) goto bad; - if (cmp == 0) { - values[name-argnames] = value; - break; - } - name++; - } - if (*name) continue; - else { - PyObject*** argname = argnames; - while (argname != first_kw_arg) { - int cmp = (**argname == key) ? 0 : - #if !CYTHON_COMPILING_IN_PYPY && PY_MAJOR_VERSION >= 3 - (PyUnicode_GET_SIZE(**argname) != PyUnicode_GET_SIZE(key)) ? 1 : - #endif - PyUnicode_Compare(**argname, key); - if (cmp < 0 && unlikely(PyErr_Occurred())) goto bad; - if (cmp == 0) goto arg_passed_twice; - argname++; - } - } - } else - goto invalid_keyword_type; - if (kwds2) { - if (unlikely(PyDict_SetItem(kwds2, key, value))) goto bad; - } else { - goto invalid_keyword; - } - } - return 0; -arg_passed_twice: - __Pyx_RaiseDoubleKeywordsError(function_name, key); - goto bad; -invalid_keyword_type: - PyErr_Format(PyExc_TypeError, - "%.200s() keywords must be strings", function_name); - goto bad; -invalid_keyword: - PyErr_Format(PyExc_TypeError, - #if PY_MAJOR_VERSION < 3 - "%.200s() got an unexpected keyword argument '%.200s'", - function_name, PyString_AsString(key)); - #else - "%s() got an unexpected keyword argument '%U'", - function_name, key); - #endif -bad: - return -1; -} - -/* ArgTypeTest */ -static int __Pyx__ArgTypeTest(PyObject *obj, PyTypeObject *type, const char *name, int exact) -{ - if (unlikely(!type)) { - PyErr_SetString(PyExc_SystemError, "Missing type object"); - return 0; - } - else if (exact) { - #if PY_MAJOR_VERSION == 2 - if ((type == &PyBaseString_Type) && likely(__Pyx_PyBaseString_CheckExact(obj))) return 1; - #endif - } - else { - if (likely(__Pyx_TypeCheck(obj, type))) return 1; - } - PyErr_Format(PyExc_TypeError, - "Argument '%.200s' has incorrect type (expected %.200s, got %.200s)", - name, type->tp_name, Py_TYPE(obj)->tp_name); - return 0; -} - -/* IsLittleEndian */ -static CYTHON_INLINE int __Pyx_Is_Little_Endian(void) -{ - union { - uint32_t u32; - uint8_t u8[4]; - } S; - S.u32 = 0x01020304; - return S.u8[0] == 4; -} - -/* BufferFormatCheck */ -static void __Pyx_BufFmt_Init(__Pyx_BufFmt_Context* ctx, - __Pyx_BufFmt_StackElem* stack, - __Pyx_TypeInfo* type) { - stack[0].field = &ctx->root; - stack[0].parent_offset = 0; - ctx->root.type = type; - ctx->root.name = "buffer dtype"; - ctx->root.offset = 0; - ctx->head = stack; - ctx->head->field = &ctx->root; - ctx->fmt_offset = 0; - ctx->head->parent_offset = 0; - ctx->new_packmode = '@'; - ctx->enc_packmode = '@'; - ctx->new_count = 1; - ctx->enc_count = 0; - ctx->enc_type = 0; - ctx->is_complex = 0; - ctx->is_valid_array = 0; - ctx->struct_alignment = 0; - while (type->typegroup == 'S') { - ++ctx->head; - ctx->head->field = type->fields; - ctx->head->parent_offset = 0; - type = type->fields->type; - } -} -static int __Pyx_BufFmt_ParseNumber(const char** ts) { - int count; - const char* t = *ts; - if (*t < '0' || *t > '9') { - return -1; - } else { - count = *t++ - '0'; - while (*t >= '0' && *t < '9') { - count *= 10; - count += *t++ - '0'; - } - } - *ts = t; - return count; -} -static int __Pyx_BufFmt_ExpectNumber(const char **ts) { - int number = __Pyx_BufFmt_ParseNumber(ts); - if (number == -1) - PyErr_Format(PyExc_ValueError,\ - "Does not understand character buffer dtype format string ('%c')", **ts); - return number; -} -static void __Pyx_BufFmt_RaiseUnexpectedChar(char ch) { - PyErr_Format(PyExc_ValueError, - "Unexpected format string character: '%c'", ch); -} -static const char* __Pyx_BufFmt_DescribeTypeChar(char ch, int is_complex) { - switch (ch) { - case 'c': return "'char'"; - case 'b': return "'signed char'"; - case 'B': return "'unsigned char'"; - case 'h': return "'short'"; - case 'H': return "'unsigned short'"; - case 'i': return "'int'"; - case 'I': return "'unsigned int'"; - case 'l': return "'long'"; - case 'L': return "'unsigned long'"; - case 'q': return "'long long'"; - case 'Q': return "'unsigned long long'"; - case 'f': return (is_complex ? "'complex float'" : "'float'"); - case 'd': return (is_complex ? "'complex double'" : "'double'"); - case 'g': return (is_complex ? "'complex long double'" : "'long double'"); - case 'T': return "a struct"; - case 'O': return "Python object"; - case 'P': return "a pointer"; - case 's': case 'p': return "a string"; - case 0: return "end"; - default: return "unparseable format string"; - } -} -static size_t __Pyx_BufFmt_TypeCharToStandardSize(char ch, int is_complex) { - switch (ch) { - case '?': case 'c': case 'b': case 'B': case 's': case 'p': return 1; - case 'h': case 'H': return 2; - case 'i': case 'I': case 'l': case 'L': return 4; - case 'q': case 'Q': return 8; - case 'f': return (is_complex ? 8 : 4); - case 'd': return (is_complex ? 16 : 8); - case 'g': { - PyErr_SetString(PyExc_ValueError, "Python does not define a standard format string size for long double ('g').."); - return 0; - } - case 'O': case 'P': return sizeof(void*); - default: - __Pyx_BufFmt_RaiseUnexpectedChar(ch); - return 0; - } -} -static size_t __Pyx_BufFmt_TypeCharToNativeSize(char ch, int is_complex) { - switch (ch) { - case 'c': case 'b': case 'B': case 's': case 'p': return 1; - case 'h': case 'H': return sizeof(short); - case 'i': case 'I': return sizeof(int); - case 'l': case 'L': return sizeof(long); - #ifdef HAVE_LONG_LONG - case 'q': case 'Q': return sizeof(PY_LONG_LONG); - #endif - case 'f': return sizeof(float) * (is_complex ? 2 : 1); - case 'd': return sizeof(double) * (is_complex ? 2 : 1); - case 'g': return sizeof(long double) * (is_complex ? 2 : 1); - case 'O': case 'P': return sizeof(void*); - default: { - __Pyx_BufFmt_RaiseUnexpectedChar(ch); - return 0; - } - } -} -typedef struct { char c; short x; } __Pyx_st_short; -typedef struct { char c; int x; } __Pyx_st_int; -typedef struct { char c; long x; } __Pyx_st_long; -typedef struct { char c; float x; } __Pyx_st_float; -typedef struct { char c; double x; } __Pyx_st_double; -typedef struct { char c; long double x; } __Pyx_st_longdouble; -typedef struct { char c; void *x; } __Pyx_st_void_p; -#ifdef HAVE_LONG_LONG -typedef struct { char c; PY_LONG_LONG x; } __Pyx_st_longlong; -#endif -static size_t __Pyx_BufFmt_TypeCharToAlignment(char ch, CYTHON_UNUSED int is_complex) { - switch (ch) { - case '?': case 'c': case 'b': case 'B': case 's': case 'p': return 1; - case 'h': case 'H': return sizeof(__Pyx_st_short) - sizeof(short); - case 'i': case 'I': return sizeof(__Pyx_st_int) - sizeof(int); - case 'l': case 'L': return sizeof(__Pyx_st_long) - sizeof(long); -#ifdef HAVE_LONG_LONG - case 'q': case 'Q': return sizeof(__Pyx_st_longlong) - sizeof(PY_LONG_LONG); -#endif - case 'f': return sizeof(__Pyx_st_float) - sizeof(float); - case 'd': return sizeof(__Pyx_st_double) - sizeof(double); - case 'g': return sizeof(__Pyx_st_longdouble) - sizeof(long double); - case 'P': case 'O': return sizeof(__Pyx_st_void_p) - sizeof(void*); - default: - __Pyx_BufFmt_RaiseUnexpectedChar(ch); - return 0; - } -} -/* These are for computing the padding at the end of the struct to align - on the first member of the struct. This will probably the same as above, - but we don't have any guarantees. - */ -typedef struct { short x; char c; } __Pyx_pad_short; -typedef struct { int x; char c; } __Pyx_pad_int; -typedef struct { long x; char c; } __Pyx_pad_long; -typedef struct { float x; char c; } __Pyx_pad_float; -typedef struct { double x; char c; } __Pyx_pad_double; -typedef struct { long double x; char c; } __Pyx_pad_longdouble; -typedef struct { void *x; char c; } __Pyx_pad_void_p; -#ifdef HAVE_LONG_LONG -typedef struct { PY_LONG_LONG x; char c; } __Pyx_pad_longlong; -#endif -static size_t __Pyx_BufFmt_TypeCharToPadding(char ch, CYTHON_UNUSED int is_complex) { - switch (ch) { - case '?': case 'c': case 'b': case 'B': case 's': case 'p': return 1; - case 'h': case 'H': return sizeof(__Pyx_pad_short) - sizeof(short); - case 'i': case 'I': return sizeof(__Pyx_pad_int) - sizeof(int); - case 'l': case 'L': return sizeof(__Pyx_pad_long) - sizeof(long); -#ifdef HAVE_LONG_LONG - case 'q': case 'Q': return sizeof(__Pyx_pad_longlong) - sizeof(PY_LONG_LONG); -#endif - case 'f': return sizeof(__Pyx_pad_float) - sizeof(float); - case 'd': return sizeof(__Pyx_pad_double) - sizeof(double); - case 'g': return sizeof(__Pyx_pad_longdouble) - sizeof(long double); - case 'P': case 'O': return sizeof(__Pyx_pad_void_p) - sizeof(void*); - default: - __Pyx_BufFmt_RaiseUnexpectedChar(ch); - return 0; - } -} -static char __Pyx_BufFmt_TypeCharToGroup(char ch, int is_complex) { - switch (ch) { - case 'c': - return 'H'; - case 'b': case 'h': case 'i': - case 'l': case 'q': case 's': case 'p': - return 'I'; - case 'B': case 'H': case 'I': case 'L': case 'Q': - return 'U'; - case 'f': case 'd': case 'g': - return (is_complex ? 'C' : 'R'); - case 'O': - return 'O'; - case 'P': - return 'P'; - default: { - __Pyx_BufFmt_RaiseUnexpectedChar(ch); - return 0; - } - } -} -static void __Pyx_BufFmt_RaiseExpected(__Pyx_BufFmt_Context* ctx) { - if (ctx->head == NULL || ctx->head->field == &ctx->root) { - const char* expected; - const char* quote; - if (ctx->head == NULL) { - expected = "end"; - quote = ""; - } else { - expected = ctx->head->field->type->name; - quote = "'"; - } - PyErr_Format(PyExc_ValueError, - "Buffer dtype mismatch, expected %s%s%s but got %s", - quote, expected, quote, - __Pyx_BufFmt_DescribeTypeChar(ctx->enc_type, ctx->is_complex)); - } else { - __Pyx_StructField* field = ctx->head->field; - __Pyx_StructField* parent = (ctx->head - 1)->field; - PyErr_Format(PyExc_ValueError, - "Buffer dtype mismatch, expected '%s' but got %s in '%s.%s'", - field->type->name, __Pyx_BufFmt_DescribeTypeChar(ctx->enc_type, ctx->is_complex), - parent->type->name, field->name); - } -} -static int __Pyx_BufFmt_ProcessTypeChunk(__Pyx_BufFmt_Context* ctx) { - char group; - size_t size, offset, arraysize = 1; - if (ctx->enc_type == 0) return 0; - if (ctx->head->field->type->arraysize[0]) { - int i, ndim = 0; - if (ctx->enc_type == 's' || ctx->enc_type == 'p') { - ctx->is_valid_array = ctx->head->field->type->ndim == 1; - ndim = 1; - if (ctx->enc_count != ctx->head->field->type->arraysize[0]) { - PyErr_Format(PyExc_ValueError, - "Expected a dimension of size %zu, got %zu", - ctx->head->field->type->arraysize[0], ctx->enc_count); - return -1; - } - } - if (!ctx->is_valid_array) { - PyErr_Format(PyExc_ValueError, "Expected %d dimensions, got %d", - ctx->head->field->type->ndim, ndim); - return -1; - } - for (i = 0; i < ctx->head->field->type->ndim; i++) { - arraysize *= ctx->head->field->type->arraysize[i]; - } - ctx->is_valid_array = 0; - ctx->enc_count = 1; - } - group = __Pyx_BufFmt_TypeCharToGroup(ctx->enc_type, ctx->is_complex); - do { - __Pyx_StructField* field = ctx->head->field; - __Pyx_TypeInfo* type = field->type; - if (ctx->enc_packmode == '@' || ctx->enc_packmode == '^') { - size = __Pyx_BufFmt_TypeCharToNativeSize(ctx->enc_type, ctx->is_complex); - } else { - size = __Pyx_BufFmt_TypeCharToStandardSize(ctx->enc_type, ctx->is_complex); - } - if (ctx->enc_packmode == '@') { - size_t align_at = __Pyx_BufFmt_TypeCharToAlignment(ctx->enc_type, ctx->is_complex); - size_t align_mod_offset; - if (align_at == 0) return -1; - align_mod_offset = ctx->fmt_offset % align_at; - if (align_mod_offset > 0) ctx->fmt_offset += align_at - align_mod_offset; - if (ctx->struct_alignment == 0) - ctx->struct_alignment = __Pyx_BufFmt_TypeCharToPadding(ctx->enc_type, - ctx->is_complex); - } - if (type->size != size || type->typegroup != group) { - if (type->typegroup == 'C' && type->fields != NULL) { - size_t parent_offset = ctx->head->parent_offset + field->offset; - ++ctx->head; - ctx->head->field = type->fields; - ctx->head->parent_offset = parent_offset; - continue; - } - if ((type->typegroup == 'H' || group == 'H') && type->size == size) { - } else { - __Pyx_BufFmt_RaiseExpected(ctx); - return -1; - } - } - offset = ctx->head->parent_offset + field->offset; - if (ctx->fmt_offset != offset) { - PyErr_Format(PyExc_ValueError, - "Buffer dtype mismatch; next field is at offset %" CYTHON_FORMAT_SSIZE_T "d but %" CYTHON_FORMAT_SSIZE_T "d expected", - (Py_ssize_t)ctx->fmt_offset, (Py_ssize_t)offset); - return -1; - } - ctx->fmt_offset += size; - if (arraysize) - ctx->fmt_offset += (arraysize - 1) * size; - --ctx->enc_count; - while (1) { - if (field == &ctx->root) { - ctx->head = NULL; - if (ctx->enc_count != 0) { - __Pyx_BufFmt_RaiseExpected(ctx); - return -1; - } - break; - } - ctx->head->field = ++field; - if (field->type == NULL) { - --ctx->head; - field = ctx->head->field; - continue; - } else if (field->type->typegroup == 'S') { - size_t parent_offset = ctx->head->parent_offset + field->offset; - if (field->type->fields->type == NULL) continue; - field = field->type->fields; - ++ctx->head; - ctx->head->field = field; - ctx->head->parent_offset = parent_offset; - break; - } else { - break; - } - } - } while (ctx->enc_count); - ctx->enc_type = 0; - ctx->is_complex = 0; - return 0; -} -static PyObject * -__pyx_buffmt_parse_array(__Pyx_BufFmt_Context* ctx, const char** tsp) -{ - const char *ts = *tsp; - int i = 0, number; - int ndim = ctx->head->field->type->ndim; -; - ++ts; - if (ctx->new_count != 1) { - PyErr_SetString(PyExc_ValueError, - "Cannot handle repeated arrays in format string"); - return NULL; - } - if (__Pyx_BufFmt_ProcessTypeChunk(ctx) == -1) return NULL; - while (*ts && *ts != ')') { - switch (*ts) { - case ' ': case '\f': case '\r': case '\n': case '\t': case '\v': continue; - default: break; - } - number = __Pyx_BufFmt_ExpectNumber(&ts); - if (number == -1) return NULL; - if (i < ndim && (size_t) number != ctx->head->field->type->arraysize[i]) - return PyErr_Format(PyExc_ValueError, - "Expected a dimension of size %zu, got %d", - ctx->head->field->type->arraysize[i], number); - if (*ts != ',' && *ts != ')') - return PyErr_Format(PyExc_ValueError, - "Expected a comma in format string, got '%c'", *ts); - if (*ts == ',') ts++; - i++; - } - if (i != ndim) - return PyErr_Format(PyExc_ValueError, "Expected %d dimension(s), got %d", - ctx->head->field->type->ndim, i); - if (!*ts) { - PyErr_SetString(PyExc_ValueError, - "Unexpected end of format string, expected ')'"); - return NULL; - } - ctx->is_valid_array = 1; - ctx->new_count = 1; - *tsp = ++ts; - return Py_None; -} -static const char* __Pyx_BufFmt_CheckString(__Pyx_BufFmt_Context* ctx, const char* ts) { - int got_Z = 0; - while (1) { - switch(*ts) { - case 0: - if (ctx->enc_type != 0 && ctx->head == NULL) { - __Pyx_BufFmt_RaiseExpected(ctx); - return NULL; - } - if (__Pyx_BufFmt_ProcessTypeChunk(ctx) == -1) return NULL; - if (ctx->head != NULL) { - __Pyx_BufFmt_RaiseExpected(ctx); - return NULL; - } - return ts; - case ' ': - case '\r': - case '\n': - ++ts; - break; - case '<': - if (!__Pyx_Is_Little_Endian()) { - PyErr_SetString(PyExc_ValueError, "Little-endian buffer not supported on big-endian compiler"); - return NULL; - } - ctx->new_packmode = '='; - ++ts; - break; - case '>': - case '!': - if (__Pyx_Is_Little_Endian()) { - PyErr_SetString(PyExc_ValueError, "Big-endian buffer not supported on little-endian compiler"); - return NULL; - } - ctx->new_packmode = '='; - ++ts; - break; - case '=': - case '@': - case '^': - ctx->new_packmode = *ts++; - break; - case 'T': - { - const char* ts_after_sub; - size_t i, struct_count = ctx->new_count; - size_t struct_alignment = ctx->struct_alignment; - ctx->new_count = 1; - ++ts; - if (*ts != '{') { - PyErr_SetString(PyExc_ValueError, "Buffer acquisition: Expected '{' after 'T'"); - return NULL; - } - if (__Pyx_BufFmt_ProcessTypeChunk(ctx) == -1) return NULL; - ctx->enc_type = 0; - ctx->enc_count = 0; - ctx->struct_alignment = 0; - ++ts; - ts_after_sub = ts; - for (i = 0; i != struct_count; ++i) { - ts_after_sub = __Pyx_BufFmt_CheckString(ctx, ts); - if (!ts_after_sub) return NULL; - } - ts = ts_after_sub; - if (struct_alignment) ctx->struct_alignment = struct_alignment; - } - break; - case '}': - { - size_t alignment = ctx->struct_alignment; - ++ts; - if (__Pyx_BufFmt_ProcessTypeChunk(ctx) == -1) return NULL; - ctx->enc_type = 0; - if (alignment && ctx->fmt_offset % alignment) { - ctx->fmt_offset += alignment - (ctx->fmt_offset % alignment); - } - } - return ts; - case 'x': - if (__Pyx_BufFmt_ProcessTypeChunk(ctx) == -1) return NULL; - ctx->fmt_offset += ctx->new_count; - ctx->new_count = 1; - ctx->enc_count = 0; - ctx->enc_type = 0; - ctx->enc_packmode = ctx->new_packmode; - ++ts; - break; - case 'Z': - got_Z = 1; - ++ts; - if (*ts != 'f' && *ts != 'd' && *ts != 'g') { - __Pyx_BufFmt_RaiseUnexpectedChar('Z'); - return NULL; - } - CYTHON_FALLTHROUGH; - case 'c': case 'b': case 'B': case 'h': case 'H': case 'i': case 'I': - case 'l': case 'L': case 'q': case 'Q': - case 'f': case 'd': case 'g': - case 'O': case 'p': - if (ctx->enc_type == *ts && got_Z == ctx->is_complex && - ctx->enc_packmode == ctx->new_packmode) { - ctx->enc_count += ctx->new_count; - ctx->new_count = 1; - got_Z = 0; - ++ts; - break; - } - CYTHON_FALLTHROUGH; - case 's': - if (__Pyx_BufFmt_ProcessTypeChunk(ctx) == -1) return NULL; - ctx->enc_count = ctx->new_count; - ctx->enc_packmode = ctx->new_packmode; - ctx->enc_type = *ts; - ctx->is_complex = got_Z; - ++ts; - ctx->new_count = 1; - got_Z = 0; - break; - case ':': - ++ts; - while(*ts != ':') ++ts; - ++ts; - break; - case '(': - if (!__pyx_buffmt_parse_array(ctx, &ts)) return NULL; - break; - default: - { - int number = __Pyx_BufFmt_ExpectNumber(&ts); - if (number == -1) return NULL; - ctx->new_count = (size_t)number; - } - } - } -} - -/* BufferGetAndValidate */ - static CYTHON_INLINE void __Pyx_SafeReleaseBuffer(Py_buffer* info) { - if (unlikely(info->buf == NULL)) return; - if (info->suboffsets == __Pyx_minusones) info->suboffsets = NULL; - __Pyx_ReleaseBuffer(info); -} -static void __Pyx_ZeroBuffer(Py_buffer* buf) { - buf->buf = NULL; - buf->obj = NULL; - buf->strides = __Pyx_zeros; - buf->shape = __Pyx_zeros; - buf->suboffsets = __Pyx_minusones; -} -static int __Pyx__GetBufferAndValidate( - Py_buffer* buf, PyObject* obj, __Pyx_TypeInfo* dtype, int flags, - int nd, int cast, __Pyx_BufFmt_StackElem* stack) -{ - buf->buf = NULL; - if (unlikely(__Pyx_GetBuffer(obj, buf, flags) == -1)) { - __Pyx_ZeroBuffer(buf); - return -1; - } - if (unlikely(buf->ndim != nd)) { - PyErr_Format(PyExc_ValueError, - "Buffer has wrong number of dimensions (expected %d, got %d)", - nd, buf->ndim); - goto fail; - } - if (!cast) { - __Pyx_BufFmt_Context ctx; - __Pyx_BufFmt_Init(&ctx, stack, dtype); - if (!__Pyx_BufFmt_CheckString(&ctx, buf->format)) goto fail; - } - if (unlikely((unsigned)buf->itemsize != dtype->size)) { - PyErr_Format(PyExc_ValueError, - "Item size of buffer (%" CYTHON_FORMAT_SSIZE_T "d byte%s) does not match size of '%s' (%" CYTHON_FORMAT_SSIZE_T "d byte%s)", - buf->itemsize, (buf->itemsize > 1) ? "s" : "", - dtype->name, (Py_ssize_t)dtype->size, (dtype->size > 1) ? "s" : ""); - goto fail; - } - if (buf->suboffsets == NULL) buf->suboffsets = __Pyx_minusones; - return 0; -fail:; - __Pyx_SafeReleaseBuffer(buf); - return -1; -} - -/* PyErrFetchRestore */ - #if CYTHON_FAST_THREAD_STATE -static CYTHON_INLINE void __Pyx_ErrRestoreInState(PyThreadState *tstate, PyObject *type, PyObject *value, PyObject *tb) { - PyObject *tmp_type, *tmp_value, *tmp_tb; - tmp_type = tstate->curexc_type; - tmp_value = tstate->curexc_value; - tmp_tb = tstate->curexc_traceback; - tstate->curexc_type = type; - tstate->curexc_value = value; - tstate->curexc_traceback = tb; - Py_XDECREF(tmp_type); - Py_XDECREF(tmp_value); - Py_XDECREF(tmp_tb); -} -static CYTHON_INLINE void __Pyx_ErrFetchInState(PyThreadState *tstate, PyObject **type, PyObject **value, PyObject **tb) { - *type = tstate->curexc_type; - *value = tstate->curexc_value; - *tb = tstate->curexc_traceback; - tstate->curexc_type = 0; - tstate->curexc_value = 0; - tstate->curexc_traceback = 0; -} -#endif - -/* PyObjectGetAttrStr */ - #if CYTHON_USE_TYPE_SLOTS -static CYTHON_INLINE PyObject* __Pyx_PyObject_GetAttrStr(PyObject* obj, PyObject* attr_name) { - PyTypeObject* tp = Py_TYPE(obj); - if (likely(tp->tp_getattro)) - return tp->tp_getattro(obj, attr_name); -#if PY_MAJOR_VERSION < 3 - if (likely(tp->tp_getattr)) - return tp->tp_getattr(obj, PyString_AS_STRING(attr_name)); -#endif - return PyObject_GetAttr(obj, attr_name); -} -#endif - -/* GetBuiltinName */ - static PyObject *__Pyx_GetBuiltinName(PyObject *name) { - PyObject* result = __Pyx_PyObject_GetAttrStr(__pyx_b, name); - if (unlikely(!result)) { - PyErr_Format(PyExc_NameError, -#if PY_MAJOR_VERSION >= 3 - "name '%U' is not defined", name); -#else - "name '%.200s' is not defined", PyString_AS_STRING(name)); -#endif - } - return result; -} - -/* PyObjectCall */ - #if CYTHON_COMPILING_IN_CPYTHON -static CYTHON_INLINE PyObject* __Pyx_PyObject_Call(PyObject *func, PyObject *arg, PyObject *kw) { - PyObject *result; - ternaryfunc call = func->ob_type->tp_call; - if (unlikely(!call)) - return PyObject_Call(func, arg, kw); - if (unlikely(Py_EnterRecursiveCall((char*)" while calling a Python object"))) - return NULL; - result = (*call)(func, arg, kw); - Py_LeaveRecursiveCall(); - if (unlikely(!result) && unlikely(!PyErr_Occurred())) { - PyErr_SetString( - PyExc_SystemError, - "NULL result without error in PyObject_Call"); - } - return result; -} -#endif - -/* RaiseException */ - #if PY_MAJOR_VERSION < 3 -static void __Pyx_Raise(PyObject *type, PyObject *value, PyObject *tb, - CYTHON_UNUSED PyObject *cause) { - __Pyx_PyThreadState_declare - Py_XINCREF(type); - if (!value || value == Py_None) - value = NULL; - else - Py_INCREF(value); - if (!tb || tb == Py_None) - tb = NULL; - else { - Py_INCREF(tb); - if (!PyTraceBack_Check(tb)) { - PyErr_SetString(PyExc_TypeError, - "raise: arg 3 must be a traceback or None"); - goto raise_error; - } - } - if (PyType_Check(type)) { -#if CYTHON_COMPILING_IN_PYPY - if (!value) { - Py_INCREF(Py_None); - value = Py_None; - } -#endif - PyErr_NormalizeException(&type, &value, &tb); - } else { - if (value) { - PyErr_SetString(PyExc_TypeError, - "instance exception may not have a separate value"); - goto raise_error; - } - value = type; - type = (PyObject*) Py_TYPE(type); - Py_INCREF(type); - if (!PyType_IsSubtype((PyTypeObject *)type, (PyTypeObject *)PyExc_BaseException)) { - PyErr_SetString(PyExc_TypeError, - "raise: exception class must be a subclass of BaseException"); - goto raise_error; - } - } - __Pyx_PyThreadState_assign - __Pyx_ErrRestore(type, value, tb); - return; -raise_error: - Py_XDECREF(value); - Py_XDECREF(type); - Py_XDECREF(tb); - return; -} -#else -static void __Pyx_Raise(PyObject *type, PyObject *value, PyObject *tb, PyObject *cause) { - PyObject* owned_instance = NULL; - if (tb == Py_None) { - tb = 0; - } else if (tb && !PyTraceBack_Check(tb)) { - PyErr_SetString(PyExc_TypeError, - "raise: arg 3 must be a traceback or None"); - goto bad; - } - if (value == Py_None) - value = 0; - if (PyExceptionInstance_Check(type)) { - if (value) { - PyErr_SetString(PyExc_TypeError, - "instance exception may not have a separate value"); - goto bad; - } - value = type; - type = (PyObject*) Py_TYPE(value); - } else if (PyExceptionClass_Check(type)) { - PyObject *instance_class = NULL; - if (value && PyExceptionInstance_Check(value)) { - instance_class = (PyObject*) Py_TYPE(value); - if (instance_class != type) { - int is_subclass = PyObject_IsSubclass(instance_class, type); - if (!is_subclass) { - instance_class = NULL; - } else if (unlikely(is_subclass == -1)) { - goto bad; - } else { - type = instance_class; - } - } - } - if (!instance_class) { - PyObject *args; - if (!value) - args = PyTuple_New(0); - else if (PyTuple_Check(value)) { - Py_INCREF(value); - args = value; - } else - args = PyTuple_Pack(1, value); - if (!args) - goto bad; - owned_instance = PyObject_Call(type, args, NULL); - Py_DECREF(args); - if (!owned_instance) - goto bad; - value = owned_instance; - if (!PyExceptionInstance_Check(value)) { - PyErr_Format(PyExc_TypeError, - "calling %R should have returned an instance of " - "BaseException, not %R", - type, Py_TYPE(value)); - goto bad; - } - } - } else { - PyErr_SetString(PyExc_TypeError, - "raise: exception class must be a subclass of BaseException"); - goto bad; - } - if (cause) { - PyObject *fixed_cause; - if (cause == Py_None) { - fixed_cause = NULL; - } else if (PyExceptionClass_Check(cause)) { - fixed_cause = PyObject_CallObject(cause, NULL); - if (fixed_cause == NULL) - goto bad; - } else if (PyExceptionInstance_Check(cause)) { - fixed_cause = cause; - Py_INCREF(fixed_cause); - } else { - PyErr_SetString(PyExc_TypeError, - "exception causes must derive from " - "BaseException"); - goto bad; - } - PyException_SetCause(value, fixed_cause); - } - PyErr_SetObject(type, value); - if (tb) { -#if CYTHON_COMPILING_IN_PYPY - PyObject *tmp_type, *tmp_value, *tmp_tb; - PyErr_Fetch(&tmp_type, &tmp_value, &tmp_tb); - Py_INCREF(tb); - PyErr_Restore(tmp_type, tmp_value, tb); - Py_XDECREF(tmp_tb); -#else - PyThreadState *tstate = __Pyx_PyThreadState_Current; - PyObject* tmp_tb = tstate->curexc_traceback; - if (tb != tmp_tb) { - Py_INCREF(tb); - tstate->curexc_traceback = tb; - Py_XDECREF(tmp_tb); - } -#endif - } -bad: - Py_XDECREF(owned_instance); - return; -} -#endif - -/* PyCFunctionFastCall */ - #if CYTHON_FAST_PYCCALL -static CYTHON_INLINE PyObject * __Pyx_PyCFunction_FastCall(PyObject *func_obj, PyObject **args, Py_ssize_t nargs) { - PyCFunctionObject *func = (PyCFunctionObject*)func_obj; - PyCFunction meth = PyCFunction_GET_FUNCTION(func); - PyObject *self = PyCFunction_GET_SELF(func); - int flags = PyCFunction_GET_FLAGS(func); - assert(PyCFunction_Check(func)); - assert(METH_FASTCALL == (flags & ~(METH_CLASS | METH_STATIC | METH_COEXIST | METH_KEYWORDS))); - assert(nargs >= 0); - assert(nargs == 0 || args != NULL); - /* _PyCFunction_FastCallDict() must not be called with an exception set, - because it may clear it (directly or indirectly) and so the - caller loses its exception */ - assert(!PyErr_Occurred()); - if ((PY_VERSION_HEX < 0x030700A0) || unlikely(flags & METH_KEYWORDS)) { - return (*((__Pyx_PyCFunctionFastWithKeywords)meth)) (self, args, nargs, NULL); - } else { - return (*((__Pyx_PyCFunctionFast)meth)) (self, args, nargs); - } -} -#endif - -/* PyFunctionFastCall */ - #if CYTHON_FAST_PYCALL -#include "frameobject.h" -static PyObject* __Pyx_PyFunction_FastCallNoKw(PyCodeObject *co, PyObject **args, Py_ssize_t na, - PyObject *globals) { - PyFrameObject *f; - PyThreadState *tstate = __Pyx_PyThreadState_Current; - PyObject **fastlocals; - Py_ssize_t i; - PyObject *result; - assert(globals != NULL); - /* XXX Perhaps we should create a specialized - PyFrame_New() that doesn't take locals, but does - take builtins without sanity checking them. - */ - assert(tstate != NULL); - f = PyFrame_New(tstate, co, globals, NULL); - if (f == NULL) { - return NULL; - } - fastlocals = f->f_localsplus; - for (i = 0; i < na; i++) { - Py_INCREF(*args); - fastlocals[i] = *args++; - } - result = PyEval_EvalFrameEx(f,0); - ++tstate->recursion_depth; - Py_DECREF(f); - --tstate->recursion_depth; - return result; -} -#if 1 || PY_VERSION_HEX < 0x030600B1 -static PyObject *__Pyx_PyFunction_FastCallDict(PyObject *func, PyObject **args, int nargs, PyObject *kwargs) { - PyCodeObject *co = (PyCodeObject *)PyFunction_GET_CODE(func); - PyObject *globals = PyFunction_GET_GLOBALS(func); - PyObject *argdefs = PyFunction_GET_DEFAULTS(func); - PyObject *closure; -#if PY_MAJOR_VERSION >= 3 - PyObject *kwdefs; -#endif - PyObject *kwtuple, **k; - PyObject **d; - Py_ssize_t nd; - Py_ssize_t nk; - PyObject *result; - assert(kwargs == NULL || PyDict_Check(kwargs)); - nk = kwargs ? PyDict_Size(kwargs) : 0; - if (Py_EnterRecursiveCall((char*)" while calling a Python object")) { - return NULL; - } - if ( -#if PY_MAJOR_VERSION >= 3 - co->co_kwonlyargcount == 0 && -#endif - likely(kwargs == NULL || nk == 0) && - co->co_flags == (CO_OPTIMIZED | CO_NEWLOCALS | CO_NOFREE)) { - if (argdefs == NULL && co->co_argcount == nargs) { - result = __Pyx_PyFunction_FastCallNoKw(co, args, nargs, globals); - goto done; - } - else if (nargs == 0 && argdefs != NULL - && co->co_argcount == Py_SIZE(argdefs)) { - /* function called with no arguments, but all parameters have - a default value: use default values as arguments .*/ - args = &PyTuple_GET_ITEM(argdefs, 0); - result =__Pyx_PyFunction_FastCallNoKw(co, args, Py_SIZE(argdefs), globals); - goto done; - } - } - if (kwargs != NULL) { - Py_ssize_t pos, i; - kwtuple = PyTuple_New(2 * nk); - if (kwtuple == NULL) { - result = NULL; - goto done; - } - k = &PyTuple_GET_ITEM(kwtuple, 0); - pos = i = 0; - while (PyDict_Next(kwargs, &pos, &k[i], &k[i+1])) { - Py_INCREF(k[i]); - Py_INCREF(k[i+1]); - i += 2; - } - nk = i / 2; - } - else { - kwtuple = NULL; - k = NULL; - } - closure = PyFunction_GET_CLOSURE(func); -#if PY_MAJOR_VERSION >= 3 - kwdefs = PyFunction_GET_KW_DEFAULTS(func); -#endif - if (argdefs != NULL) { - d = &PyTuple_GET_ITEM(argdefs, 0); - nd = Py_SIZE(argdefs); - } - else { - d = NULL; - nd = 0; - } -#if PY_MAJOR_VERSION >= 3 - result = PyEval_EvalCodeEx((PyObject*)co, globals, (PyObject *)NULL, - args, nargs, - k, (int)nk, - d, (int)nd, kwdefs, closure); -#else - result = PyEval_EvalCodeEx(co, globals, (PyObject *)NULL, - args, nargs, - k, (int)nk, - d, (int)nd, closure); -#endif - Py_XDECREF(kwtuple); -done: - Py_LeaveRecursiveCall(); - return result; -} -#endif -#endif - -/* PyObjectCallMethO */ - #if CYTHON_COMPILING_IN_CPYTHON -static CYTHON_INLINE PyObject* __Pyx_PyObject_CallMethO(PyObject *func, PyObject *arg) { - PyObject *self, *result; - PyCFunction cfunc; - cfunc = PyCFunction_GET_FUNCTION(func); - self = PyCFunction_GET_SELF(func); - if (unlikely(Py_EnterRecursiveCall((char*)" while calling a Python object"))) - return NULL; - result = cfunc(self, arg); - Py_LeaveRecursiveCall(); - if (unlikely(!result) && unlikely(!PyErr_Occurred())) { - PyErr_SetString( - PyExc_SystemError, - "NULL result without error in PyObject_Call"); - } - return result; -} -#endif - -/* PyObjectCallOneArg */ - #if CYTHON_COMPILING_IN_CPYTHON -static PyObject* __Pyx__PyObject_CallOneArg(PyObject *func, PyObject *arg) { - PyObject *result; - PyObject *args = PyTuple_New(1); - if (unlikely(!args)) return NULL; - Py_INCREF(arg); - PyTuple_SET_ITEM(args, 0, arg); - result = __Pyx_PyObject_Call(func, args, NULL); - Py_DECREF(args); - return result; -} -static CYTHON_INLINE PyObject* __Pyx_PyObject_CallOneArg(PyObject *func, PyObject *arg) { -#if CYTHON_FAST_PYCALL - if (PyFunction_Check(func)) { - return __Pyx_PyFunction_FastCall(func, &arg, 1); - } -#endif - if (likely(PyCFunction_Check(func))) { - if (likely(PyCFunction_GET_FLAGS(func) & METH_O)) { - return __Pyx_PyObject_CallMethO(func, arg); -#if CYTHON_FAST_PYCCALL - } else if (PyCFunction_GET_FLAGS(func) & METH_FASTCALL) { - return __Pyx_PyCFunction_FastCall(func, &arg, 1); -#endif - } - } - return __Pyx__PyObject_CallOneArg(func, arg); -} -#else -static CYTHON_INLINE PyObject* __Pyx_PyObject_CallOneArg(PyObject *func, PyObject *arg) { - PyObject *result; - PyObject *args = PyTuple_Pack(1, arg); - if (unlikely(!args)) return NULL; - result = __Pyx_PyObject_Call(func, args, NULL); - Py_DECREF(args); - return result; -} -#endif - -/* DictGetItem */ - #if PY_MAJOR_VERSION >= 3 && !CYTHON_COMPILING_IN_PYPY -static PyObject *__Pyx_PyDict_GetItem(PyObject *d, PyObject* key) { - PyObject *value; - value = PyDict_GetItemWithError(d, key); - if (unlikely(!value)) { - if (!PyErr_Occurred()) { - PyObject* args = PyTuple_Pack(1, key); - if (likely(args)) - PyErr_SetObject(PyExc_KeyError, args); - Py_XDECREF(args); - } - return NULL; - } - Py_INCREF(value); - return value; -} -#endif - -/* RaiseTooManyValuesToUnpack */ - static CYTHON_INLINE void __Pyx_RaiseTooManyValuesError(Py_ssize_t expected) { - PyErr_Format(PyExc_ValueError, - "too many values to unpack (expected %" CYTHON_FORMAT_SSIZE_T "d)", expected); -} - -/* RaiseNeedMoreValuesToUnpack */ - static CYTHON_INLINE void __Pyx_RaiseNeedMoreValuesError(Py_ssize_t index) { - PyErr_Format(PyExc_ValueError, - "need more than %" CYTHON_FORMAT_SSIZE_T "d value%.1s to unpack", - index, (index == 1) ? "" : "s"); -} - -/* RaiseNoneIterError */ - static CYTHON_INLINE void __Pyx_RaiseNoneNotIterableError(void) { - PyErr_SetString(PyExc_TypeError, "'NoneType' object is not iterable"); -} - -/* ExtTypeTest */ - static CYTHON_INLINE int __Pyx_TypeTest(PyObject *obj, PyTypeObject *type) { - if (unlikely(!type)) { - PyErr_SetString(PyExc_SystemError, "Missing type object"); - return 0; - } - if (likely(__Pyx_TypeCheck(obj, type))) - return 1; - PyErr_Format(PyExc_TypeError, "Cannot convert %.200s to %.200s", - Py_TYPE(obj)->tp_name, type->tp_name); - return 0; -} - -/* SaveResetException */ - #if CYTHON_FAST_THREAD_STATE -static CYTHON_INLINE void __Pyx__ExceptionSave(PyThreadState *tstate, PyObject **type, PyObject **value, PyObject **tb) { - #if PY_VERSION_HEX >= 0x030700A2 - *type = tstate->exc_state.exc_type; - *value = tstate->exc_state.exc_value; - *tb = tstate->exc_state.exc_traceback; - #else - *type = tstate->exc_type; - *value = tstate->exc_value; - *tb = tstate->exc_traceback; - #endif - Py_XINCREF(*type); - Py_XINCREF(*value); - Py_XINCREF(*tb); -} -static CYTHON_INLINE void __Pyx__ExceptionReset(PyThreadState *tstate, PyObject *type, PyObject *value, PyObject *tb) { - PyObject *tmp_type, *tmp_value, *tmp_tb; - #if PY_VERSION_HEX >= 0x030700A2 - tmp_type = tstate->exc_state.exc_type; - tmp_value = tstate->exc_state.exc_value; - tmp_tb = tstate->exc_state.exc_traceback; - tstate->exc_state.exc_type = type; - tstate->exc_state.exc_value = value; - tstate->exc_state.exc_traceback = tb; - #else - tmp_type = tstate->exc_type; - tmp_value = tstate->exc_value; - tmp_tb = tstate->exc_traceback; - tstate->exc_type = type; - tstate->exc_value = value; - tstate->exc_traceback = tb; - #endif - Py_XDECREF(tmp_type); - Py_XDECREF(tmp_value); - Py_XDECREF(tmp_tb); -} -#endif - -/* PyErrExceptionMatches */ - #if CYTHON_FAST_THREAD_STATE -static int __Pyx_PyErr_ExceptionMatchesTuple(PyObject *exc_type, PyObject *tuple) { - Py_ssize_t i, n; - n = PyTuple_GET_SIZE(tuple); -#if PY_MAJOR_VERSION >= 3 - for (i=0; icurexc_type; - if (exc_type == err) return 1; - if (unlikely(!exc_type)) return 0; - if (unlikely(PyTuple_Check(err))) - return __Pyx_PyErr_ExceptionMatchesTuple(exc_type, err); - return __Pyx_PyErr_GivenExceptionMatches(exc_type, err); -} -#endif - -/* GetException */ - #if CYTHON_FAST_THREAD_STATE -static int __Pyx__GetException(PyThreadState *tstate, PyObject **type, PyObject **value, PyObject **tb) { -#else -static int __Pyx_GetException(PyObject **type, PyObject **value, PyObject **tb) { -#endif - PyObject *local_type, *local_value, *local_tb; -#if CYTHON_FAST_THREAD_STATE - PyObject *tmp_type, *tmp_value, *tmp_tb; - local_type = tstate->curexc_type; - local_value = tstate->curexc_value; - local_tb = tstate->curexc_traceback; - tstate->curexc_type = 0; - tstate->curexc_value = 0; - tstate->curexc_traceback = 0; -#else - PyErr_Fetch(&local_type, &local_value, &local_tb); -#endif - PyErr_NormalizeException(&local_type, &local_value, &local_tb); -#if CYTHON_FAST_THREAD_STATE - if (unlikely(tstate->curexc_type)) -#else - if (unlikely(PyErr_Occurred())) -#endif - goto bad; - #if PY_MAJOR_VERSION >= 3 - if (local_tb) { - if (unlikely(PyException_SetTraceback(local_value, local_tb) < 0)) - goto bad; - } - #endif - Py_XINCREF(local_tb); - Py_XINCREF(local_type); - Py_XINCREF(local_value); - *type = local_type; - *value = local_value; - *tb = local_tb; -#if CYTHON_FAST_THREAD_STATE - #if PY_VERSION_HEX >= 0x030700A2 - tmp_type = tstate->exc_state.exc_type; - tmp_value = tstate->exc_state.exc_value; - tmp_tb = tstate->exc_state.exc_traceback; - tstate->exc_state.exc_type = local_type; - tstate->exc_state.exc_value = local_value; - tstate->exc_state.exc_traceback = local_tb; - #else - tmp_type = tstate->exc_type; - tmp_value = tstate->exc_value; - tmp_tb = tstate->exc_traceback; - tstate->exc_type = local_type; - tstate->exc_value = local_value; - tstate->exc_traceback = local_tb; - #endif - Py_XDECREF(tmp_type); - Py_XDECREF(tmp_value); - Py_XDECREF(tmp_tb); -#else - PyErr_SetExcInfo(local_type, local_value, local_tb); -#endif - return 0; -bad: - *type = 0; - *value = 0; - *tb = 0; - Py_XDECREF(local_type); - Py_XDECREF(local_value); - Py_XDECREF(local_tb); - return -1; -} - -/* Import */ - static PyObject *__Pyx_Import(PyObject *name, PyObject *from_list, int level) { - PyObject *empty_list = 0; - PyObject *module = 0; - PyObject *global_dict = 0; - PyObject *empty_dict = 0; - PyObject *list; - #if PY_MAJOR_VERSION < 3 - PyObject *py_import; - py_import = __Pyx_PyObject_GetAttrStr(__pyx_b, __pyx_n_s_import); - if (!py_import) - goto bad; - #endif - if (from_list) - list = from_list; - else { - empty_list = PyList_New(0); - if (!empty_list) - goto bad; - list = empty_list; - } - global_dict = PyModule_GetDict(__pyx_m); - if (!global_dict) - goto bad; - empty_dict = PyDict_New(); - if (!empty_dict) - goto bad; - { - #if PY_MAJOR_VERSION >= 3 - if (level == -1) { - if (strchr(__Pyx_MODULE_NAME, '.')) { - module = PyImport_ImportModuleLevelObject( - name, global_dict, empty_dict, list, 1); - if (!module) { - if (!PyErr_ExceptionMatches(PyExc_ImportError)) - goto bad; - PyErr_Clear(); - } - } - level = 0; - } - #endif - if (!module) { - #if PY_MAJOR_VERSION < 3 - PyObject *py_level = PyInt_FromLong(level); - if (!py_level) - goto bad; - module = PyObject_CallFunctionObjArgs(py_import, - name, global_dict, empty_dict, list, py_level, NULL); - Py_DECREF(py_level); - #else - module = PyImport_ImportModuleLevelObject( - name, global_dict, empty_dict, list, level); - #endif - } - } -bad: - #if PY_MAJOR_VERSION < 3 - Py_XDECREF(py_import); - #endif - Py_XDECREF(empty_list); - Py_XDECREF(empty_dict); - return module; -} - -/* CLineInTraceback */ - #ifndef CYTHON_CLINE_IN_TRACEBACK -static int __Pyx_CLineForTraceback(CYTHON_UNUSED PyThreadState *tstate, int c_line) { - PyObject *use_cline; - PyObject *ptype, *pvalue, *ptraceback; -#if CYTHON_COMPILING_IN_CPYTHON - PyObject **cython_runtime_dict; -#endif - __Pyx_ErrFetchInState(tstate, &ptype, &pvalue, &ptraceback); -#if CYTHON_COMPILING_IN_CPYTHON - cython_runtime_dict = _PyObject_GetDictPtr(__pyx_cython_runtime); - if (likely(cython_runtime_dict)) { - use_cline = __Pyx_PyDict_GetItemStr(*cython_runtime_dict, __pyx_n_s_cline_in_traceback); - } else -#endif - { - PyObject *use_cline_obj = __Pyx_PyObject_GetAttrStr(__pyx_cython_runtime, __pyx_n_s_cline_in_traceback); - if (use_cline_obj) { - use_cline = PyObject_Not(use_cline_obj) ? Py_False : Py_True; - Py_DECREF(use_cline_obj); - } else { - PyErr_Clear(); - use_cline = NULL; - } - } - if (!use_cline) { - c_line = 0; - PyObject_SetAttr(__pyx_cython_runtime, __pyx_n_s_cline_in_traceback, Py_False); - } - else if (PyObject_Not(use_cline) != 0) { - c_line = 0; - } - __Pyx_ErrRestoreInState(tstate, ptype, pvalue, ptraceback); - return c_line; -} -#endif - -/* CodeObjectCache */ - static int __pyx_bisect_code_objects(__Pyx_CodeObjectCacheEntry* entries, int count, int code_line) { - int start = 0, mid = 0, end = count - 1; - if (end >= 0 && code_line > entries[end].code_line) { - return count; - } - while (start < end) { - mid = start + (end - start) / 2; - if (code_line < entries[mid].code_line) { - end = mid; - } else if (code_line > entries[mid].code_line) { - start = mid + 1; - } else { - return mid; - } - } - if (code_line <= entries[mid].code_line) { - return mid; - } else { - return mid + 1; - } -} -static PyCodeObject *__pyx_find_code_object(int code_line) { - PyCodeObject* code_object; - int pos; - if (unlikely(!code_line) || unlikely(!__pyx_code_cache.entries)) { - return NULL; - } - pos = __pyx_bisect_code_objects(__pyx_code_cache.entries, __pyx_code_cache.count, code_line); - if (unlikely(pos >= __pyx_code_cache.count) || unlikely(__pyx_code_cache.entries[pos].code_line != code_line)) { - return NULL; - } - code_object = __pyx_code_cache.entries[pos].code_object; - Py_INCREF(code_object); - return code_object; -} -static void __pyx_insert_code_object(int code_line, PyCodeObject* code_object) { - int pos, i; - __Pyx_CodeObjectCacheEntry* entries = __pyx_code_cache.entries; - if (unlikely(!code_line)) { - return; - } - if (unlikely(!entries)) { - entries = (__Pyx_CodeObjectCacheEntry*)PyMem_Malloc(64*sizeof(__Pyx_CodeObjectCacheEntry)); - if (likely(entries)) { - __pyx_code_cache.entries = entries; - __pyx_code_cache.max_count = 64; - __pyx_code_cache.count = 1; - entries[0].code_line = code_line; - entries[0].code_object = code_object; - Py_INCREF(code_object); - } - return; - } - pos = __pyx_bisect_code_objects(__pyx_code_cache.entries, __pyx_code_cache.count, code_line); - if ((pos < __pyx_code_cache.count) && unlikely(__pyx_code_cache.entries[pos].code_line == code_line)) { - PyCodeObject* tmp = entries[pos].code_object; - entries[pos].code_object = code_object; - Py_DECREF(tmp); - return; - } - if (__pyx_code_cache.count == __pyx_code_cache.max_count) { - int new_max = __pyx_code_cache.max_count + 64; - entries = (__Pyx_CodeObjectCacheEntry*)PyMem_Realloc( - __pyx_code_cache.entries, (size_t)new_max*sizeof(__Pyx_CodeObjectCacheEntry)); - if (unlikely(!entries)) { - return; - } - __pyx_code_cache.entries = entries; - __pyx_code_cache.max_count = new_max; - } - for (i=__pyx_code_cache.count; i>pos; i--) { - entries[i] = entries[i-1]; - } - entries[pos].code_line = code_line; - entries[pos].code_object = code_object; - __pyx_code_cache.count++; - Py_INCREF(code_object); -} - -/* AddTraceback */ - #include "compile.h" -#include "frameobject.h" -#include "traceback.h" -static PyCodeObject* __Pyx_CreateCodeObjectForTraceback( - const char *funcname, int c_line, - int py_line, const char *filename) { - PyCodeObject *py_code = 0; - PyObject *py_srcfile = 0; - PyObject *py_funcname = 0; - #if PY_MAJOR_VERSION < 3 - py_srcfile = PyString_FromString(filename); - #else - py_srcfile = PyUnicode_FromString(filename); - #endif - if (!py_srcfile) goto bad; - if (c_line) { - #if PY_MAJOR_VERSION < 3 - py_funcname = PyString_FromFormat( "%s (%s:%d)", funcname, __pyx_cfilenm, c_line); - #else - py_funcname = PyUnicode_FromFormat( "%s (%s:%d)", funcname, __pyx_cfilenm, c_line); - #endif - } - else { - #if PY_MAJOR_VERSION < 3 - py_funcname = PyString_FromString(funcname); - #else - py_funcname = PyUnicode_FromString(funcname); - #endif - } - if (!py_funcname) goto bad; - py_code = __Pyx_PyCode_New( - 0, - 0, - 0, - 0, - 0, - __pyx_empty_bytes, /*PyObject *code,*/ - __pyx_empty_tuple, /*PyObject *consts,*/ - __pyx_empty_tuple, /*PyObject *names,*/ - __pyx_empty_tuple, /*PyObject *varnames,*/ - __pyx_empty_tuple, /*PyObject *freevars,*/ - __pyx_empty_tuple, /*PyObject *cellvars,*/ - py_srcfile, /*PyObject *filename,*/ - py_funcname, /*PyObject *name,*/ - py_line, - __pyx_empty_bytes /*PyObject *lnotab*/ - ); - Py_DECREF(py_srcfile); - Py_DECREF(py_funcname); - return py_code; -bad: - Py_XDECREF(py_srcfile); - Py_XDECREF(py_funcname); - return NULL; -} -static void __Pyx_AddTraceback(const char *funcname, int c_line, - int py_line, const char *filename) { - PyCodeObject *py_code = 0; - PyFrameObject *py_frame = 0; - PyThreadState *tstate = __Pyx_PyThreadState_Current; - if (c_line) { - c_line = __Pyx_CLineForTraceback(tstate, c_line); - } - py_code = __pyx_find_code_object(c_line ? -c_line : py_line); - if (!py_code) { - py_code = __Pyx_CreateCodeObjectForTraceback( - funcname, c_line, py_line, filename); - if (!py_code) goto bad; - __pyx_insert_code_object(c_line ? -c_line : py_line, py_code); - } - py_frame = PyFrame_New( - tstate, /*PyThreadState *tstate,*/ - py_code, /*PyCodeObject *code,*/ - __pyx_d, /*PyObject *globals,*/ - 0 /*PyObject *locals*/ - ); - if (!py_frame) goto bad; - __Pyx_PyFrame_SetLineNumber(py_frame, py_line); - PyTraceBack_Here(py_frame); -bad: - Py_XDECREF(py_code); - Py_XDECREF(py_frame); -} - -#if PY_MAJOR_VERSION < 3 -static int __Pyx_GetBuffer(PyObject *obj, Py_buffer *view, int flags) { - if (PyObject_CheckBuffer(obj)) return PyObject_GetBuffer(obj, view, flags); - if (__Pyx_TypeCheck(obj, __pyx_ptype_5numpy_ndarray)) return __pyx_pw_5numpy_7ndarray_1__getbuffer__(obj, view, flags); - PyErr_Format(PyExc_TypeError, "'%.200s' does not have the buffer interface", Py_TYPE(obj)->tp_name); - return -1; -} -static void __Pyx_ReleaseBuffer(Py_buffer *view) { - PyObject *obj = view->obj; - if (!obj) return; - if (PyObject_CheckBuffer(obj)) { - PyBuffer_Release(view); - return; - } - if ((0)) {} - else if (__Pyx_TypeCheck(obj, __pyx_ptype_5numpy_ndarray)) __pyx_pw_5numpy_7ndarray_3__releasebuffer__(obj, view); - view->obj = NULL; - Py_DECREF(obj); -} -#endif - - - /* CIntFromPyVerify */ - #define __PYX_VERIFY_RETURN_INT(target_type, func_type, func_value)\ - __PYX__VERIFY_RETURN_INT(target_type, func_type, func_value, 0) -#define __PYX_VERIFY_RETURN_INT_EXC(target_type, func_type, func_value)\ - __PYX__VERIFY_RETURN_INT(target_type, func_type, func_value, 1) -#define __PYX__VERIFY_RETURN_INT(target_type, func_type, func_value, exc)\ - {\ - func_type value = func_value;\ - if (sizeof(target_type) < sizeof(func_type)) {\ - if (unlikely(value != (func_type) (target_type) value)) {\ - func_type zero = 0;\ - if (exc && unlikely(value == (func_type)-1 && PyErr_Occurred()))\ - return (target_type) -1;\ - if (is_unsigned && unlikely(value < zero))\ - goto raise_neg_overflow;\ - else\ - goto raise_overflow;\ - }\ - }\ - return (target_type) value;\ - } - -/* Declarations */ - #if CYTHON_CCOMPLEX - #ifdef __cplusplus - static CYTHON_INLINE __pyx_t_float_complex __pyx_t_float_complex_from_parts(float x, float y) { - return ::std::complex< float >(x, y); - } - #else - static CYTHON_INLINE __pyx_t_float_complex __pyx_t_float_complex_from_parts(float x, float y) { - return x + y*(__pyx_t_float_complex)_Complex_I; - } - #endif -#else - static CYTHON_INLINE __pyx_t_float_complex __pyx_t_float_complex_from_parts(float x, float y) { - __pyx_t_float_complex z; - z.real = x; - z.imag = y; - return z; - } -#endif - -/* Arithmetic */ - #if CYTHON_CCOMPLEX -#else - static CYTHON_INLINE int __Pyx_c_eq_float(__pyx_t_float_complex a, __pyx_t_float_complex b) { - return (a.real == b.real) && (a.imag == b.imag); - } - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_sum_float(__pyx_t_float_complex a, __pyx_t_float_complex b) { - __pyx_t_float_complex z; - z.real = a.real + b.real; - z.imag = a.imag + b.imag; - return z; - } - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_diff_float(__pyx_t_float_complex a, __pyx_t_float_complex b) { - __pyx_t_float_complex z; - z.real = a.real - b.real; - z.imag = a.imag - b.imag; - return z; - } - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_prod_float(__pyx_t_float_complex a, __pyx_t_float_complex b) { - __pyx_t_float_complex z; - z.real = a.real * b.real - a.imag * b.imag; - z.imag = a.real * b.imag + a.imag * b.real; - return z; - } - #if 1 - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_quot_float(__pyx_t_float_complex a, __pyx_t_float_complex b) { - if (b.imag == 0) { - return __pyx_t_float_complex_from_parts(a.real / b.real, a.imag / b.real); - } else if (fabsf(b.real) >= fabsf(b.imag)) { - if (b.real == 0 && b.imag == 0) { - return __pyx_t_float_complex_from_parts(a.real / b.real, a.imag / b.imag); - } else { - float r = b.imag / b.real; - float s = 1.0 / (b.real + b.imag * r); - return __pyx_t_float_complex_from_parts( - (a.real + a.imag * r) * s, (a.imag - a.real * r) * s); - } - } else { - float r = b.real / b.imag; - float s = 1.0 / (b.imag + b.real * r); - return __pyx_t_float_complex_from_parts( - (a.real * r + a.imag) * s, (a.imag * r - a.real) * s); - } - } - #else - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_quot_float(__pyx_t_float_complex a, __pyx_t_float_complex b) { - if (b.imag == 0) { - return __pyx_t_float_complex_from_parts(a.real / b.real, a.imag / b.real); - } else { - float denom = b.real * b.real + b.imag * b.imag; - return __pyx_t_float_complex_from_parts( - (a.real * b.real + a.imag * b.imag) / denom, - (a.imag * b.real - a.real * b.imag) / denom); - } - } - #endif - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_neg_float(__pyx_t_float_complex a) { - __pyx_t_float_complex z; - z.real = -a.real; - z.imag = -a.imag; - return z; - } - static CYTHON_INLINE int __Pyx_c_is_zero_float(__pyx_t_float_complex a) { - return (a.real == 0) && (a.imag == 0); - } - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_conj_float(__pyx_t_float_complex a) { - __pyx_t_float_complex z; - z.real = a.real; - z.imag = -a.imag; - return z; - } - #if 1 - static CYTHON_INLINE float __Pyx_c_abs_float(__pyx_t_float_complex z) { - #if !defined(HAVE_HYPOT) || defined(_MSC_VER) - return sqrtf(z.real*z.real + z.imag*z.imag); - #else - return hypotf(z.real, z.imag); - #endif - } - static CYTHON_INLINE __pyx_t_float_complex __Pyx_c_pow_float(__pyx_t_float_complex a, __pyx_t_float_complex b) { - __pyx_t_float_complex z; - float r, lnr, theta, z_r, z_theta; - if (b.imag == 0 && b.real == (int)b.real) { - if (b.real < 0) { - float denom = a.real * a.real + a.imag * a.imag; - a.real = a.real / denom; - a.imag = -a.imag / denom; - b.real = -b.real; - } - switch ((int)b.real) { - case 0: - z.real = 1; - z.imag = 0; - return z; - case 1: - return a; - case 2: - z = __Pyx_c_prod_float(a, a); - return __Pyx_c_prod_float(a, a); - case 3: - z = __Pyx_c_prod_float(a, a); - return __Pyx_c_prod_float(z, a); - case 4: - z = __Pyx_c_prod_float(a, a); - return __Pyx_c_prod_float(z, z); - } - } - if (a.imag == 0) { - if (a.real == 0) { - return a; - } else if (b.imag == 0) { - z.real = powf(a.real, b.real); - z.imag = 0; - return z; - } else if (a.real > 0) { - r = a.real; - theta = 0; - } else { - r = -a.real; - theta = atan2f(0, -1); - } - } else { - r = __Pyx_c_abs_float(a); - theta = atan2f(a.imag, a.real); - } - lnr = logf(r); - z_r = expf(lnr * b.real - theta * b.imag); - z_theta = theta * b.real + lnr * b.imag; - z.real = z_r * cosf(z_theta); - z.imag = z_r * sinf(z_theta); - return z; - } - #endif -#endif - -/* Declarations */ - #if CYTHON_CCOMPLEX - #ifdef __cplusplus - static CYTHON_INLINE __pyx_t_double_complex __pyx_t_double_complex_from_parts(double x, double y) { - return ::std::complex< double >(x, y); - } - #else - static CYTHON_INLINE __pyx_t_double_complex __pyx_t_double_complex_from_parts(double x, double y) { - return x + y*(__pyx_t_double_complex)_Complex_I; - } - #endif -#else - static CYTHON_INLINE __pyx_t_double_complex __pyx_t_double_complex_from_parts(double x, double y) { - __pyx_t_double_complex z; - z.real = x; - z.imag = y; - return z; - } -#endif - -/* Arithmetic */ - #if CYTHON_CCOMPLEX -#else - static CYTHON_INLINE int __Pyx_c_eq_double(__pyx_t_double_complex a, __pyx_t_double_complex b) { - return (a.real == b.real) && (a.imag == b.imag); - } - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_sum_double(__pyx_t_double_complex a, __pyx_t_double_complex b) { - __pyx_t_double_complex z; - z.real = a.real + b.real; - z.imag = a.imag + b.imag; - return z; - } - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_diff_double(__pyx_t_double_complex a, __pyx_t_double_complex b) { - __pyx_t_double_complex z; - z.real = a.real - b.real; - z.imag = a.imag - b.imag; - return z; - } - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_prod_double(__pyx_t_double_complex a, __pyx_t_double_complex b) { - __pyx_t_double_complex z; - z.real = a.real * b.real - a.imag * b.imag; - z.imag = a.real * b.imag + a.imag * b.real; - return z; - } - #if 1 - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_quot_double(__pyx_t_double_complex a, __pyx_t_double_complex b) { - if (b.imag == 0) { - return __pyx_t_double_complex_from_parts(a.real / b.real, a.imag / b.real); - } else if (fabs(b.real) >= fabs(b.imag)) { - if (b.real == 0 && b.imag == 0) { - return __pyx_t_double_complex_from_parts(a.real / b.real, a.imag / b.imag); - } else { - double r = b.imag / b.real; - double s = 1.0 / (b.real + b.imag * r); - return __pyx_t_double_complex_from_parts( - (a.real + a.imag * r) * s, (a.imag - a.real * r) * s); - } - } else { - double r = b.real / b.imag; - double s = 1.0 / (b.imag + b.real * r); - return __pyx_t_double_complex_from_parts( - (a.real * r + a.imag) * s, (a.imag * r - a.real) * s); - } - } - #else - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_quot_double(__pyx_t_double_complex a, __pyx_t_double_complex b) { - if (b.imag == 0) { - return __pyx_t_double_complex_from_parts(a.real / b.real, a.imag / b.real); - } else { - double denom = b.real * b.real + b.imag * b.imag; - return __pyx_t_double_complex_from_parts( - (a.real * b.real + a.imag * b.imag) / denom, - (a.imag * b.real - a.real * b.imag) / denom); - } - } - #endif - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_neg_double(__pyx_t_double_complex a) { - __pyx_t_double_complex z; - z.real = -a.real; - z.imag = -a.imag; - return z; - } - static CYTHON_INLINE int __Pyx_c_is_zero_double(__pyx_t_double_complex a) { - return (a.real == 0) && (a.imag == 0); - } - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_conj_double(__pyx_t_double_complex a) { - __pyx_t_double_complex z; - z.real = a.real; - z.imag = -a.imag; - return z; - } - #if 1 - static CYTHON_INLINE double __Pyx_c_abs_double(__pyx_t_double_complex z) { - #if !defined(HAVE_HYPOT) || defined(_MSC_VER) - return sqrt(z.real*z.real + z.imag*z.imag); - #else - return hypot(z.real, z.imag); - #endif - } - static CYTHON_INLINE __pyx_t_double_complex __Pyx_c_pow_double(__pyx_t_double_complex a, __pyx_t_double_complex b) { - __pyx_t_double_complex z; - double r, lnr, theta, z_r, z_theta; - if (b.imag == 0 && b.real == (int)b.real) { - if (b.real < 0) { - double denom = a.real * a.real + a.imag * a.imag; - a.real = a.real / denom; - a.imag = -a.imag / denom; - b.real = -b.real; - } - switch ((int)b.real) { - case 0: - z.real = 1; - z.imag = 0; - return z; - case 1: - return a; - case 2: - z = __Pyx_c_prod_double(a, a); - return __Pyx_c_prod_double(a, a); - case 3: - z = __Pyx_c_prod_double(a, a); - return __Pyx_c_prod_double(z, a); - case 4: - z = __Pyx_c_prod_double(a, a); - return __Pyx_c_prod_double(z, z); - } - } - if (a.imag == 0) { - if (a.real == 0) { - return a; - } else if (b.imag == 0) { - z.real = pow(a.real, b.real); - z.imag = 0; - return z; - } else if (a.real > 0) { - r = a.real; - theta = 0; - } else { - r = -a.real; - theta = atan2(0, -1); - } - } else { - r = __Pyx_c_abs_double(a); - theta = atan2(a.imag, a.real); - } - lnr = log(r); - z_r = exp(lnr * b.real - theta * b.imag); - z_theta = theta * b.real + lnr * b.imag; - z.real = z_r * cos(z_theta); - z.imag = z_r * sin(z_theta); - return z; - } - #endif -#endif - -/* CIntToPy */ - static CYTHON_INLINE PyObject* __Pyx_PyInt_From_int(int value) { - const int neg_one = (int) -1, const_zero = (int) 0; - const int is_unsigned = neg_one > const_zero; - if (is_unsigned) { - if (sizeof(int) < sizeof(long)) { - return PyInt_FromLong((long) value); - } else if (sizeof(int) <= sizeof(unsigned long)) { - return PyLong_FromUnsignedLong((unsigned long) value); -#ifdef HAVE_LONG_LONG - } else if (sizeof(int) <= sizeof(unsigned PY_LONG_LONG)) { - return PyLong_FromUnsignedLongLong((unsigned PY_LONG_LONG) value); -#endif - } - } else { - if (sizeof(int) <= sizeof(long)) { - return PyInt_FromLong((long) value); -#ifdef HAVE_LONG_LONG - } else if (sizeof(int) <= sizeof(PY_LONG_LONG)) { - return PyLong_FromLongLong((PY_LONG_LONG) value); -#endif - } - } - { - int one = 1; int little = (int)*(unsigned char *)&one; - unsigned char *bytes = (unsigned char *)&value; - return _PyLong_FromByteArray(bytes, sizeof(int), - little, !is_unsigned); - } -} - -/* CIntToPy */ - static CYTHON_INLINE PyObject* __Pyx_PyInt_From_enum__NPY_TYPES(enum NPY_TYPES value) { - const enum NPY_TYPES neg_one = (enum NPY_TYPES) -1, const_zero = (enum NPY_TYPES) 0; - const int is_unsigned = neg_one > const_zero; - if (is_unsigned) { - if (sizeof(enum NPY_TYPES) < sizeof(long)) { - return PyInt_FromLong((long) value); - } else if (sizeof(enum NPY_TYPES) <= sizeof(unsigned long)) { - return PyLong_FromUnsignedLong((unsigned long) value); -#ifdef HAVE_LONG_LONG - } else if (sizeof(enum NPY_TYPES) <= sizeof(unsigned PY_LONG_LONG)) { - return PyLong_FromUnsignedLongLong((unsigned PY_LONG_LONG) value); -#endif - } - } else { - if (sizeof(enum NPY_TYPES) <= sizeof(long)) { - return PyInt_FromLong((long) value); -#ifdef HAVE_LONG_LONG - } else if (sizeof(enum NPY_TYPES) <= sizeof(PY_LONG_LONG)) { - return PyLong_FromLongLong((PY_LONG_LONG) value); -#endif - } - } - { - int one = 1; int little = (int)*(unsigned char *)&one; - unsigned char *bytes = (unsigned char *)&value; - return _PyLong_FromByteArray(bytes, sizeof(enum NPY_TYPES), - little, !is_unsigned); - } -} - -/* CIntFromPy */ - static CYTHON_INLINE int __Pyx_PyInt_As_int(PyObject *x) { - const int neg_one = (int) -1, const_zero = (int) 0; - const int is_unsigned = neg_one > const_zero; -#if PY_MAJOR_VERSION < 3 - if (likely(PyInt_Check(x))) { - if (sizeof(int) < sizeof(long)) { - __PYX_VERIFY_RETURN_INT(int, long, PyInt_AS_LONG(x)) - } else { - long val = PyInt_AS_LONG(x); - if (is_unsigned && unlikely(val < 0)) { - goto raise_neg_overflow; - } - return (int) val; - } - } else -#endif - if (likely(PyLong_Check(x))) { - if (is_unsigned) { -#if CYTHON_USE_PYLONG_INTERNALS - const digit* digits = ((PyLongObject*)x)->ob_digit; - switch (Py_SIZE(x)) { - case 0: return (int) 0; - case 1: __PYX_VERIFY_RETURN_INT(int, digit, digits[0]) - case 2: - if (8 * sizeof(int) > 1 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 2 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(int, unsigned long, (((((unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(int) >= 2 * PyLong_SHIFT) { - return (int) (((((int)digits[1]) << PyLong_SHIFT) | (int)digits[0])); - } - } - break; - case 3: - if (8 * sizeof(int) > 2 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 3 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(int, unsigned long, (((((((unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(int) >= 3 * PyLong_SHIFT) { - return (int) (((((((int)digits[2]) << PyLong_SHIFT) | (int)digits[1]) << PyLong_SHIFT) | (int)digits[0])); - } - } - break; - case 4: - if (8 * sizeof(int) > 3 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 4 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(int, unsigned long, (((((((((unsigned long)digits[3]) << PyLong_SHIFT) | (unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(int) >= 4 * PyLong_SHIFT) { - return (int) (((((((((int)digits[3]) << PyLong_SHIFT) | (int)digits[2]) << PyLong_SHIFT) | (int)digits[1]) << PyLong_SHIFT) | (int)digits[0])); - } - } - break; - } -#endif -#if CYTHON_COMPILING_IN_CPYTHON - if (unlikely(Py_SIZE(x) < 0)) { - goto raise_neg_overflow; - } -#else - { - int result = PyObject_RichCompareBool(x, Py_False, Py_LT); - if (unlikely(result < 0)) - return (int) -1; - if (unlikely(result == 1)) - goto raise_neg_overflow; - } -#endif - if (sizeof(int) <= sizeof(unsigned long)) { - __PYX_VERIFY_RETURN_INT_EXC(int, unsigned long, PyLong_AsUnsignedLong(x)) -#ifdef HAVE_LONG_LONG - } else if (sizeof(int) <= sizeof(unsigned PY_LONG_LONG)) { - __PYX_VERIFY_RETURN_INT_EXC(int, unsigned PY_LONG_LONG, PyLong_AsUnsignedLongLong(x)) -#endif - } - } else { -#if CYTHON_USE_PYLONG_INTERNALS - const digit* digits = ((PyLongObject*)x)->ob_digit; - switch (Py_SIZE(x)) { - case 0: return (int) 0; - case -1: __PYX_VERIFY_RETURN_INT(int, sdigit, (sdigit) (-(sdigit)digits[0])) - case 1: __PYX_VERIFY_RETURN_INT(int, digit, +digits[0]) - case -2: - if (8 * sizeof(int) - 1 > 1 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 2 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(int, long, -(long) (((((unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(int) - 1 > 2 * PyLong_SHIFT) { - return (int) (((int)-1)*(((((int)digits[1]) << PyLong_SHIFT) | (int)digits[0]))); - } - } - break; - case 2: - if (8 * sizeof(int) > 1 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 2 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(int, unsigned long, (((((unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(int) - 1 > 2 * PyLong_SHIFT) { - return (int) ((((((int)digits[1]) << PyLong_SHIFT) | (int)digits[0]))); - } - } - break; - case -3: - if (8 * sizeof(int) - 1 > 2 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 3 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(int, long, -(long) (((((((unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(int) - 1 > 3 * PyLong_SHIFT) { - return (int) (((int)-1)*(((((((int)digits[2]) << PyLong_SHIFT) | (int)digits[1]) << PyLong_SHIFT) | (int)digits[0]))); - } - } - break; - case 3: - if (8 * sizeof(int) > 2 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 3 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(int, unsigned long, (((((((unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(int) - 1 > 3 * PyLong_SHIFT) { - return (int) ((((((((int)digits[2]) << PyLong_SHIFT) | (int)digits[1]) << PyLong_SHIFT) | (int)digits[0]))); - } - } - break; - case -4: - if (8 * sizeof(int) - 1 > 3 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 4 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(int, long, -(long) (((((((((unsigned long)digits[3]) << PyLong_SHIFT) | (unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(int) - 1 > 4 * PyLong_SHIFT) { - return (int) (((int)-1)*(((((((((int)digits[3]) << PyLong_SHIFT) | (int)digits[2]) << PyLong_SHIFT) | (int)digits[1]) << PyLong_SHIFT) | (int)digits[0]))); - } - } - break; - case 4: - if (8 * sizeof(int) > 3 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 4 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(int, unsigned long, (((((((((unsigned long)digits[3]) << PyLong_SHIFT) | (unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(int) - 1 > 4 * PyLong_SHIFT) { - return (int) ((((((((((int)digits[3]) << PyLong_SHIFT) | (int)digits[2]) << PyLong_SHIFT) | (int)digits[1]) << PyLong_SHIFT) | (int)digits[0]))); - } - } - break; - } -#endif - if (sizeof(int) <= sizeof(long)) { - __PYX_VERIFY_RETURN_INT_EXC(int, long, PyLong_AsLong(x)) -#ifdef HAVE_LONG_LONG - } else if (sizeof(int) <= sizeof(PY_LONG_LONG)) { - __PYX_VERIFY_RETURN_INT_EXC(int, PY_LONG_LONG, PyLong_AsLongLong(x)) -#endif - } - } - { -#if CYTHON_COMPILING_IN_PYPY && !defined(_PyLong_AsByteArray) - PyErr_SetString(PyExc_RuntimeError, - "_PyLong_AsByteArray() not available in PyPy, cannot convert large numbers"); -#else - int val; - PyObject *v = __Pyx_PyNumber_IntOrLong(x); - #if PY_MAJOR_VERSION < 3 - if (likely(v) && !PyLong_Check(v)) { - PyObject *tmp = v; - v = PyNumber_Long(tmp); - Py_DECREF(tmp); - } - #endif - if (likely(v)) { - int one = 1; int is_little = (int)*(unsigned char *)&one; - unsigned char *bytes = (unsigned char *)&val; - int ret = _PyLong_AsByteArray((PyLongObject *)v, - bytes, sizeof(val), - is_little, !is_unsigned); - Py_DECREF(v); - if (likely(!ret)) - return val; - } -#endif - return (int) -1; - } - } else { - int val; - PyObject *tmp = __Pyx_PyNumber_IntOrLong(x); - if (!tmp) return (int) -1; - val = __Pyx_PyInt_As_int(tmp); - Py_DECREF(tmp); - return val; - } -raise_overflow: - PyErr_SetString(PyExc_OverflowError, - "value too large to convert to int"); - return (int) -1; -raise_neg_overflow: - PyErr_SetString(PyExc_OverflowError, - "can't convert negative value to int"); - return (int) -1; -} - -/* CIntToPy */ - static CYTHON_INLINE PyObject* __Pyx_PyInt_From_long(long value) { - const long neg_one = (long) -1, const_zero = (long) 0; - const int is_unsigned = neg_one > const_zero; - if (is_unsigned) { - if (sizeof(long) < sizeof(long)) { - return PyInt_FromLong((long) value); - } else if (sizeof(long) <= sizeof(unsigned long)) { - return PyLong_FromUnsignedLong((unsigned long) value); -#ifdef HAVE_LONG_LONG - } else if (sizeof(long) <= sizeof(unsigned PY_LONG_LONG)) { - return PyLong_FromUnsignedLongLong((unsigned PY_LONG_LONG) value); -#endif - } - } else { - if (sizeof(long) <= sizeof(long)) { - return PyInt_FromLong((long) value); -#ifdef HAVE_LONG_LONG - } else if (sizeof(long) <= sizeof(PY_LONG_LONG)) { - return PyLong_FromLongLong((PY_LONG_LONG) value); -#endif - } - } - { - int one = 1; int little = (int)*(unsigned char *)&one; - unsigned char *bytes = (unsigned char *)&value; - return _PyLong_FromByteArray(bytes, sizeof(long), - little, !is_unsigned); - } -} - -/* CIntFromPy */ - static CYTHON_INLINE long __Pyx_PyInt_As_long(PyObject *x) { - const long neg_one = (long) -1, const_zero = (long) 0; - const int is_unsigned = neg_one > const_zero; -#if PY_MAJOR_VERSION < 3 - if (likely(PyInt_Check(x))) { - if (sizeof(long) < sizeof(long)) { - __PYX_VERIFY_RETURN_INT(long, long, PyInt_AS_LONG(x)) - } else { - long val = PyInt_AS_LONG(x); - if (is_unsigned && unlikely(val < 0)) { - goto raise_neg_overflow; - } - return (long) val; - } - } else -#endif - if (likely(PyLong_Check(x))) { - if (is_unsigned) { -#if CYTHON_USE_PYLONG_INTERNALS - const digit* digits = ((PyLongObject*)x)->ob_digit; - switch (Py_SIZE(x)) { - case 0: return (long) 0; - case 1: __PYX_VERIFY_RETURN_INT(long, digit, digits[0]) - case 2: - if (8 * sizeof(long) > 1 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 2 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(long, unsigned long, (((((unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(long) >= 2 * PyLong_SHIFT) { - return (long) (((((long)digits[1]) << PyLong_SHIFT) | (long)digits[0])); - } - } - break; - case 3: - if (8 * sizeof(long) > 2 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 3 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(long, unsigned long, (((((((unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(long) >= 3 * PyLong_SHIFT) { - return (long) (((((((long)digits[2]) << PyLong_SHIFT) | (long)digits[1]) << PyLong_SHIFT) | (long)digits[0])); - } - } - break; - case 4: - if (8 * sizeof(long) > 3 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 4 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(long, unsigned long, (((((((((unsigned long)digits[3]) << PyLong_SHIFT) | (unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(long) >= 4 * PyLong_SHIFT) { - return (long) (((((((((long)digits[3]) << PyLong_SHIFT) | (long)digits[2]) << PyLong_SHIFT) | (long)digits[1]) << PyLong_SHIFT) | (long)digits[0])); - } - } - break; - } -#endif -#if CYTHON_COMPILING_IN_CPYTHON - if (unlikely(Py_SIZE(x) < 0)) { - goto raise_neg_overflow; - } -#else - { - int result = PyObject_RichCompareBool(x, Py_False, Py_LT); - if (unlikely(result < 0)) - return (long) -1; - if (unlikely(result == 1)) - goto raise_neg_overflow; - } -#endif - if (sizeof(long) <= sizeof(unsigned long)) { - __PYX_VERIFY_RETURN_INT_EXC(long, unsigned long, PyLong_AsUnsignedLong(x)) -#ifdef HAVE_LONG_LONG - } else if (sizeof(long) <= sizeof(unsigned PY_LONG_LONG)) { - __PYX_VERIFY_RETURN_INT_EXC(long, unsigned PY_LONG_LONG, PyLong_AsUnsignedLongLong(x)) -#endif - } - } else { -#if CYTHON_USE_PYLONG_INTERNALS - const digit* digits = ((PyLongObject*)x)->ob_digit; - switch (Py_SIZE(x)) { - case 0: return (long) 0; - case -1: __PYX_VERIFY_RETURN_INT(long, sdigit, (sdigit) (-(sdigit)digits[0])) - case 1: __PYX_VERIFY_RETURN_INT(long, digit, +digits[0]) - case -2: - if (8 * sizeof(long) - 1 > 1 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 2 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(long, long, -(long) (((((unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(long) - 1 > 2 * PyLong_SHIFT) { - return (long) (((long)-1)*(((((long)digits[1]) << PyLong_SHIFT) | (long)digits[0]))); - } - } - break; - case 2: - if (8 * sizeof(long) > 1 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 2 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(long, unsigned long, (((((unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(long) - 1 > 2 * PyLong_SHIFT) { - return (long) ((((((long)digits[1]) << PyLong_SHIFT) | (long)digits[0]))); - } - } - break; - case -3: - if (8 * sizeof(long) - 1 > 2 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 3 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(long, long, -(long) (((((((unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(long) - 1 > 3 * PyLong_SHIFT) { - return (long) (((long)-1)*(((((((long)digits[2]) << PyLong_SHIFT) | (long)digits[1]) << PyLong_SHIFT) | (long)digits[0]))); - } - } - break; - case 3: - if (8 * sizeof(long) > 2 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 3 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(long, unsigned long, (((((((unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(long) - 1 > 3 * PyLong_SHIFT) { - return (long) ((((((((long)digits[2]) << PyLong_SHIFT) | (long)digits[1]) << PyLong_SHIFT) | (long)digits[0]))); - } - } - break; - case -4: - if (8 * sizeof(long) - 1 > 3 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 4 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(long, long, -(long) (((((((((unsigned long)digits[3]) << PyLong_SHIFT) | (unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(long) - 1 > 4 * PyLong_SHIFT) { - return (long) (((long)-1)*(((((((((long)digits[3]) << PyLong_SHIFT) | (long)digits[2]) << PyLong_SHIFT) | (long)digits[1]) << PyLong_SHIFT) | (long)digits[0]))); - } - } - break; - case 4: - if (8 * sizeof(long) > 3 * PyLong_SHIFT) { - if (8 * sizeof(unsigned long) > 4 * PyLong_SHIFT) { - __PYX_VERIFY_RETURN_INT(long, unsigned long, (((((((((unsigned long)digits[3]) << PyLong_SHIFT) | (unsigned long)digits[2]) << PyLong_SHIFT) | (unsigned long)digits[1]) << PyLong_SHIFT) | (unsigned long)digits[0]))) - } else if (8 * sizeof(long) - 1 > 4 * PyLong_SHIFT) { - return (long) ((((((((((long)digits[3]) << PyLong_SHIFT) | (long)digits[2]) << PyLong_SHIFT) | (long)digits[1]) << PyLong_SHIFT) | (long)digits[0]))); - } - } - break; - } -#endif - if (sizeof(long) <= sizeof(long)) { - __PYX_VERIFY_RETURN_INT_EXC(long, long, PyLong_AsLong(x)) -#ifdef HAVE_LONG_LONG - } else if (sizeof(long) <= sizeof(PY_LONG_LONG)) { - __PYX_VERIFY_RETURN_INT_EXC(long, PY_LONG_LONG, PyLong_AsLongLong(x)) -#endif - } - } - { -#if CYTHON_COMPILING_IN_PYPY && !defined(_PyLong_AsByteArray) - PyErr_SetString(PyExc_RuntimeError, - "_PyLong_AsByteArray() not available in PyPy, cannot convert large numbers"); -#else - long val; - PyObject *v = __Pyx_PyNumber_IntOrLong(x); - #if PY_MAJOR_VERSION < 3 - if (likely(v) && !PyLong_Check(v)) { - PyObject *tmp = v; - v = PyNumber_Long(tmp); - Py_DECREF(tmp); - } - #endif - if (likely(v)) { - int one = 1; int is_little = (int)*(unsigned char *)&one; - unsigned char *bytes = (unsigned char *)&val; - int ret = _PyLong_AsByteArray((PyLongObject *)v, - bytes, sizeof(val), - is_little, !is_unsigned); - Py_DECREF(v); - if (likely(!ret)) - return val; - } -#endif - return (long) -1; - } - } else { - long val; - PyObject *tmp = __Pyx_PyNumber_IntOrLong(x); - if (!tmp) return (long) -1; - val = __Pyx_PyInt_As_long(tmp); - Py_DECREF(tmp); - return val; - } -raise_overflow: - PyErr_SetString(PyExc_OverflowError, - "value too large to convert to long"); - return (long) -1; -raise_neg_overflow: - PyErr_SetString(PyExc_OverflowError, - "can't convert negative value to long"); - return (long) -1; -} - -/* FastTypeChecks */ - #if CYTHON_COMPILING_IN_CPYTHON -static int __Pyx_InBases(PyTypeObject *a, PyTypeObject *b) { - while (a) { - a = a->tp_base; - if (a == b) - return 1; - } - return b == &PyBaseObject_Type; -} -static CYTHON_INLINE int __Pyx_IsSubtype(PyTypeObject *a, PyTypeObject *b) { - PyObject *mro; - if (a == b) return 1; - mro = a->tp_mro; - if (likely(mro)) { - Py_ssize_t i, n; - n = PyTuple_GET_SIZE(mro); - for (i = 0; i < n; i++) { - if (PyTuple_GET_ITEM(mro, i) == (PyObject *)b) - return 1; - } - return 0; - } - return __Pyx_InBases(a, b); -} -#if PY_MAJOR_VERSION == 2 -static int __Pyx_inner_PyErr_GivenExceptionMatches2(PyObject *err, PyObject* exc_type1, PyObject* exc_type2) { - PyObject *exception, *value, *tb; - int res; - __Pyx_PyThreadState_declare - __Pyx_PyThreadState_assign - __Pyx_ErrFetch(&exception, &value, &tb); - res = exc_type1 ? PyObject_IsSubclass(err, exc_type1) : 0; - if (unlikely(res == -1)) { - PyErr_WriteUnraisable(err); - res = 0; - } - if (!res) { - res = PyObject_IsSubclass(err, exc_type2); - if (unlikely(res == -1)) { - PyErr_WriteUnraisable(err); - res = 0; - } - } - __Pyx_ErrRestore(exception, value, tb); - return res; -} -#else -static CYTHON_INLINE int __Pyx_inner_PyErr_GivenExceptionMatches2(PyObject *err, PyObject* exc_type1, PyObject *exc_type2) { - int res = exc_type1 ? __Pyx_IsSubtype((PyTypeObject*)err, (PyTypeObject*)exc_type1) : 0; - if (!res) { - res = __Pyx_IsSubtype((PyTypeObject*)err, (PyTypeObject*)exc_type2); - } - return res; -} -#endif -static CYTHON_INLINE int __Pyx_PyErr_GivenExceptionMatches(PyObject *err, PyObject* exc_type) { - if (likely(err == exc_type)) return 1; - if (likely(PyExceptionClass_Check(err))) { - return __Pyx_inner_PyErr_GivenExceptionMatches2(err, NULL, exc_type); - } - return PyErr_GivenExceptionMatches(err, exc_type); -} -static CYTHON_INLINE int __Pyx_PyErr_GivenExceptionMatches2(PyObject *err, PyObject *exc_type1, PyObject *exc_type2) { - if (likely(err == exc_type1 || err == exc_type2)) return 1; - if (likely(PyExceptionClass_Check(err))) { - return __Pyx_inner_PyErr_GivenExceptionMatches2(err, exc_type1, exc_type2); - } - return (PyErr_GivenExceptionMatches(err, exc_type1) || PyErr_GivenExceptionMatches(err, exc_type2)); -} -#endif - -/* CheckBinaryVersion */ - static int __Pyx_check_binary_version(void) { - char ctversion[4], rtversion[4]; - PyOS_snprintf(ctversion, 4, "%d.%d", PY_MAJOR_VERSION, PY_MINOR_VERSION); - PyOS_snprintf(rtversion, 4, "%s", Py_GetVersion()); - if (ctversion[0] != rtversion[0] || ctversion[2] != rtversion[2]) { - char message[200]; - PyOS_snprintf(message, sizeof(message), - "compiletime version %s of module '%.100s' " - "does not match runtime version %s", - ctversion, __Pyx_MODULE_NAME, rtversion); - return PyErr_WarnEx(NULL, message, 1); - } - return 0; -} - -/* ModuleImport */ - #ifndef __PYX_HAVE_RT_ImportModule -#define __PYX_HAVE_RT_ImportModule -static PyObject *__Pyx_ImportModule(const char *name) { - PyObject *py_name = 0; - PyObject *py_module = 0; - py_name = __Pyx_PyIdentifier_FromString(name); - if (!py_name) - goto bad; - py_module = PyImport_Import(py_name); - Py_DECREF(py_name); - return py_module; -bad: - Py_XDECREF(py_name); - return 0; -} -#endif - -/* TypeImport */ - #ifndef __PYX_HAVE_RT_ImportType -#define __PYX_HAVE_RT_ImportType -static PyTypeObject *__Pyx_ImportType(const char *module_name, const char *class_name, - size_t size, int strict) -{ - PyObject *py_module = 0; - PyObject *result = 0; - PyObject *py_name = 0; - char warning[200]; - Py_ssize_t basicsize; -#ifdef Py_LIMITED_API - PyObject *py_basicsize; -#endif - py_module = __Pyx_ImportModule(module_name); - if (!py_module) - goto bad; - py_name = __Pyx_PyIdentifier_FromString(class_name); - if (!py_name) - goto bad; - result = PyObject_GetAttr(py_module, py_name); - Py_DECREF(py_name); - py_name = 0; - Py_DECREF(py_module); - py_module = 0; - if (!result) - goto bad; - if (!PyType_Check(result)) { - PyErr_Format(PyExc_TypeError, - "%.200s.%.200s is not a type object", - module_name, class_name); - goto bad; - } -#ifndef Py_LIMITED_API - basicsize = ((PyTypeObject *)result)->tp_basicsize; -#else - py_basicsize = PyObject_GetAttrString(result, "__basicsize__"); - if (!py_basicsize) - goto bad; - basicsize = PyLong_AsSsize_t(py_basicsize); - Py_DECREF(py_basicsize); - py_basicsize = 0; - if (basicsize == (Py_ssize_t)-1 && PyErr_Occurred()) - goto bad; -#endif - if (!strict && (size_t)basicsize > size) { - PyOS_snprintf(warning, sizeof(warning), - "%s.%s size changed, may indicate binary incompatibility. Expected %zd, got %zd", - module_name, class_name, basicsize, size); - if (PyErr_WarnEx(NULL, warning, 0) < 0) goto bad; - } - else if ((size_t)basicsize != size) { - PyErr_Format(PyExc_ValueError, - "%.200s.%.200s has the wrong size, try recompiling. Expected %zd, got %zd", - module_name, class_name, basicsize, size); - goto bad; - } - return (PyTypeObject *)result; -bad: - Py_XDECREF(py_module); - Py_XDECREF(result); - return NULL; -} -#endif - -/* InitStrings */ - static int __Pyx_InitStrings(__Pyx_StringTabEntry *t) { - while (t->p) { - #if PY_MAJOR_VERSION < 3 - if (t->is_unicode) { - *t->p = PyUnicode_DecodeUTF8(t->s, t->n - 1, NULL); - } else if (t->intern) { - *t->p = PyString_InternFromString(t->s); - } else { - *t->p = PyString_FromStringAndSize(t->s, t->n - 1); - } - #else - if (t->is_unicode | t->is_str) { - if (t->intern) { - *t->p = PyUnicode_InternFromString(t->s); - } else if (t->encoding) { - *t->p = PyUnicode_Decode(t->s, t->n - 1, t->encoding, NULL); - } else { - *t->p = PyUnicode_FromStringAndSize(t->s, t->n - 1); - } - } else { - *t->p = PyBytes_FromStringAndSize(t->s, t->n - 1); - } - #endif - if (!*t->p) - return -1; - if (PyObject_Hash(*t->p) == -1) - return -1; - ++t; - } - return 0; -} - -static CYTHON_INLINE PyObject* __Pyx_PyUnicode_FromString(const char* c_str) { - return __Pyx_PyUnicode_FromStringAndSize(c_str, (Py_ssize_t)strlen(c_str)); -} -static CYTHON_INLINE const char* __Pyx_PyObject_AsString(PyObject* o) { - Py_ssize_t ignore; - return __Pyx_PyObject_AsStringAndSize(o, &ignore); -} -#if __PYX_DEFAULT_STRING_ENCODING_IS_ASCII || __PYX_DEFAULT_STRING_ENCODING_IS_DEFAULT -#if !CYTHON_PEP393_ENABLED -static const char* __Pyx_PyUnicode_AsStringAndSize(PyObject* o, Py_ssize_t *length) { - char* defenc_c; - PyObject* defenc = _PyUnicode_AsDefaultEncodedString(o, NULL); - if (!defenc) return NULL; - defenc_c = PyBytes_AS_STRING(defenc); -#if __PYX_DEFAULT_STRING_ENCODING_IS_ASCII - { - char* end = defenc_c + PyBytes_GET_SIZE(defenc); - char* c; - for (c = defenc_c; c < end; c++) { - if ((unsigned char) (*c) >= 128) { - PyUnicode_AsASCIIString(o); - return NULL; - } - } - } -#endif - *length = PyBytes_GET_SIZE(defenc); - return defenc_c; -} -#else -static CYTHON_INLINE const char* __Pyx_PyUnicode_AsStringAndSize(PyObject* o, Py_ssize_t *length) { - if (unlikely(__Pyx_PyUnicode_READY(o) == -1)) return NULL; -#if __PYX_DEFAULT_STRING_ENCODING_IS_ASCII - if (likely(PyUnicode_IS_ASCII(o))) { - *length = PyUnicode_GET_LENGTH(o); - return PyUnicode_AsUTF8(o); - } else { - PyUnicode_AsASCIIString(o); - return NULL; - } -#else - return PyUnicode_AsUTF8AndSize(o, length); -#endif -} -#endif -#endif -static CYTHON_INLINE const char* __Pyx_PyObject_AsStringAndSize(PyObject* o, Py_ssize_t *length) { -#if __PYX_DEFAULT_STRING_ENCODING_IS_ASCII || __PYX_DEFAULT_STRING_ENCODING_IS_DEFAULT - if ( -#if PY_MAJOR_VERSION < 3 && __PYX_DEFAULT_STRING_ENCODING_IS_ASCII - __Pyx_sys_getdefaultencoding_not_ascii && -#endif - PyUnicode_Check(o)) { - return __Pyx_PyUnicode_AsStringAndSize(o, length); - } else -#endif -#if (!CYTHON_COMPILING_IN_PYPY) || (defined(PyByteArray_AS_STRING) && defined(PyByteArray_GET_SIZE)) - if (PyByteArray_Check(o)) { - *length = PyByteArray_GET_SIZE(o); - return PyByteArray_AS_STRING(o); - } else -#endif - { - char* result; - int r = PyBytes_AsStringAndSize(o, &result, length); - if (unlikely(r < 0)) { - return NULL; - } else { - return result; - } - } -} -static CYTHON_INLINE int __Pyx_PyObject_IsTrue(PyObject* x) { - int is_true = x == Py_True; - if (is_true | (x == Py_False) | (x == Py_None)) return is_true; - else return PyObject_IsTrue(x); -} -static PyObject* __Pyx_PyNumber_IntOrLongWrongResultType(PyObject* result, const char* type_name) { -#if PY_MAJOR_VERSION >= 3 - if (PyLong_Check(result)) { - if (PyErr_WarnFormat(PyExc_DeprecationWarning, 1, - "__int__ returned non-int (type %.200s). " - "The ability to return an instance of a strict subclass of int " - "is deprecated, and may be removed in a future version of Python.", - Py_TYPE(result)->tp_name)) { - Py_DECREF(result); - return NULL; - } - return result; - } -#endif - PyErr_Format(PyExc_TypeError, - "__%.4s__ returned non-%.4s (type %.200s)", - type_name, type_name, Py_TYPE(result)->tp_name); - Py_DECREF(result); - return NULL; -} -static CYTHON_INLINE PyObject* __Pyx_PyNumber_IntOrLong(PyObject* x) { -#if CYTHON_USE_TYPE_SLOTS - PyNumberMethods *m; -#endif - const char *name = NULL; - PyObject *res = NULL; -#if PY_MAJOR_VERSION < 3 - if (likely(PyInt_Check(x) || PyLong_Check(x))) -#else - if (likely(PyLong_Check(x))) -#endif - return __Pyx_NewRef(x); -#if CYTHON_USE_TYPE_SLOTS - m = Py_TYPE(x)->tp_as_number; - #if PY_MAJOR_VERSION < 3 - if (m && m->nb_int) { - name = "int"; - res = m->nb_int(x); - } - else if (m && m->nb_long) { - name = "long"; - res = m->nb_long(x); - } - #else - if (likely(m && m->nb_int)) { - name = "int"; - res = m->nb_int(x); - } - #endif -#else - if (!PyBytes_CheckExact(x) && !PyUnicode_CheckExact(x)) { - res = PyNumber_Int(x); - } -#endif - if (likely(res)) { -#if PY_MAJOR_VERSION < 3 - if (unlikely(!PyInt_Check(res) && !PyLong_Check(res))) { -#else - if (unlikely(!PyLong_CheckExact(res))) { -#endif - return __Pyx_PyNumber_IntOrLongWrongResultType(res, name); - } - } - else if (!PyErr_Occurred()) { - PyErr_SetString(PyExc_TypeError, - "an integer is required"); - } - return res; -} -static CYTHON_INLINE Py_ssize_t __Pyx_PyIndex_AsSsize_t(PyObject* b) { - Py_ssize_t ival; - PyObject *x; -#if PY_MAJOR_VERSION < 3 - if (likely(PyInt_CheckExact(b))) { - if (sizeof(Py_ssize_t) >= sizeof(long)) - return PyInt_AS_LONG(b); - else - return PyInt_AsSsize_t(x); - } -#endif - if (likely(PyLong_CheckExact(b))) { - #if CYTHON_USE_PYLONG_INTERNALS - const digit* digits = ((PyLongObject*)b)->ob_digit; - const Py_ssize_t size = Py_SIZE(b); - if (likely(__Pyx_sst_abs(size) <= 1)) { - ival = likely(size) ? digits[0] : 0; - if (size == -1) ival = -ival; - return ival; - } else { - switch (size) { - case 2: - if (8 * sizeof(Py_ssize_t) > 2 * PyLong_SHIFT) { - return (Py_ssize_t) (((((size_t)digits[1]) << PyLong_SHIFT) | (size_t)digits[0])); - } - break; - case -2: - if (8 * sizeof(Py_ssize_t) > 2 * PyLong_SHIFT) { - return -(Py_ssize_t) (((((size_t)digits[1]) << PyLong_SHIFT) | (size_t)digits[0])); - } - break; - case 3: - if (8 * sizeof(Py_ssize_t) > 3 * PyLong_SHIFT) { - return (Py_ssize_t) (((((((size_t)digits[2]) << PyLong_SHIFT) | (size_t)digits[1]) << PyLong_SHIFT) | (size_t)digits[0])); - } - break; - case -3: - if (8 * sizeof(Py_ssize_t) > 3 * PyLong_SHIFT) { - return -(Py_ssize_t) (((((((size_t)digits[2]) << PyLong_SHIFT) | (size_t)digits[1]) << PyLong_SHIFT) | (size_t)digits[0])); - } - break; - case 4: - if (8 * sizeof(Py_ssize_t) > 4 * PyLong_SHIFT) { - return (Py_ssize_t) (((((((((size_t)digits[3]) << PyLong_SHIFT) | (size_t)digits[2]) << PyLong_SHIFT) | (size_t)digits[1]) << PyLong_SHIFT) | (size_t)digits[0])); - } - break; - case -4: - if (8 * sizeof(Py_ssize_t) > 4 * PyLong_SHIFT) { - return -(Py_ssize_t) (((((((((size_t)digits[3]) << PyLong_SHIFT) | (size_t)digits[2]) << PyLong_SHIFT) | (size_t)digits[1]) << PyLong_SHIFT) | (size_t)digits[0])); - } - break; - } - } - #endif - return PyLong_AsSsize_t(b); - } - x = PyNumber_Index(b); - if (!x) return -1; - ival = PyInt_AsSsize_t(x); - Py_DECREF(x); - return ival; -} -static CYTHON_INLINE PyObject * __Pyx_PyInt_FromSize_t(size_t ival) { - return PyInt_FromSize_t(ival); -} - - -#endif /* Py_PYTHON_H */ diff --git a/venv/insightface/thirdparty/face3d/mesh/cython/mesh_core_cython.pyx b/venv/insightface/thirdparty/face3d/mesh/cython/mesh_core_cython.pyx deleted file mode 100644 index 40f589a9b6ddfe5d42d62c242348380c0fe2242a..0000000000000000000000000000000000000000 --- a/venv/insightface/thirdparty/face3d/mesh/cython/mesh_core_cython.pyx +++ /dev/null @@ -1,109 +0,0 @@ -import numpy as np -cimport numpy as np -from libcpp.string cimport string - -# use the Numpy-C-API from Cython -np.import_array() - -# cdefine the signature of our c function -cdef extern from "mesh_core.h": - void _rasterize_triangles_core( - float* vertices, int* triangles, - float* depth_buffer, int* triangle_buffer, float* barycentric_weight, - int nver, int ntri, - int h, int w) - - void _render_colors_core( - float* image, float* vertices, int* triangles, - float* colors, - float* depth_buffer, - int nver, int ntri, - int h, int w, int c) - - void _render_texture_core( - float* image, float* vertices, int* triangles, - float* texture, float* tex_coords, int* tex_triangles, - float* depth_buffer, - int nver, int tex_nver, int ntri, - int h, int w, int c, - int tex_h, int tex_w, int tex_c, - int mapping_type) - - void _get_normal_core( - float* normal, float* tri_normal, int* triangles, - int ntri) - - void _write_obj_with_colors_texture(string filename, string mtl_name, - float* vertices, int* triangles, float* colors, float* uv_coords, - int nver, int ntri, int ntexver) - -def get_normal_core(np.ndarray[float, ndim=2, mode = "c"] normal not None, - np.ndarray[float, ndim=2, mode = "c"] tri_normal not None, - np.ndarray[int, ndim=2, mode="c"] triangles not None, - int ntri - ): - _get_normal_core( - np.PyArray_DATA(normal), np.PyArray_DATA(tri_normal), np.PyArray_DATA(triangles), - ntri) - -def rasterize_triangles_core( - np.ndarray[float, ndim=2, mode = "c"] vertices not None, - np.ndarray[int, ndim=2, mode="c"] triangles not None, - np.ndarray[float, ndim=2, mode = "c"] depth_buffer not None, - np.ndarray[int, ndim=2, mode = "c"] triangle_buffer not None, - np.ndarray[float, ndim=2, mode = "c"] barycentric_weight not None, - int nver, int ntri, - int h, int w - ): - _rasterize_triangles_core( - np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), - np.PyArray_DATA(depth_buffer), np.PyArray_DATA(triangle_buffer), np.PyArray_DATA(barycentric_weight), - nver, ntri, - h, w) - -def render_colors_core(np.ndarray[float, ndim=3, mode = "c"] image not None, - np.ndarray[float, ndim=2, mode = "c"] vertices not None, - np.ndarray[int, ndim=2, mode="c"] triangles not None, - np.ndarray[float, ndim=2, mode = "c"] colors not None, - np.ndarray[float, ndim=2, mode = "c"] depth_buffer not None, - int nver, int ntri, - int h, int w, int c - ): - _render_colors_core( - np.PyArray_DATA(image), np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), - np.PyArray_DATA(colors), - np.PyArray_DATA(depth_buffer), - nver, ntri, - h, w, c) - -def render_texture_core(np.ndarray[float, ndim=3, mode = "c"] image not None, - np.ndarray[float, ndim=2, mode = "c"] vertices not None, - np.ndarray[int, ndim=2, mode="c"] triangles not None, - np.ndarray[float, ndim=3, mode = "c"] texture not None, - np.ndarray[float, ndim=2, mode = "c"] tex_coords not None, - np.ndarray[int, ndim=2, mode="c"] tex_triangles not None, - np.ndarray[float, ndim=2, mode = "c"] depth_buffer not None, - int nver, int tex_nver, int ntri, - int h, int w, int c, - int tex_h, int tex_w, int tex_c, - int mapping_type - ): - _render_texture_core( - np.PyArray_DATA(image), np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), - np.PyArray_DATA(texture), np.PyArray_DATA(tex_coords), np.PyArray_DATA(tex_triangles), - np.PyArray_DATA(depth_buffer), - nver, tex_nver, ntri, - h, w, c, - tex_h, tex_w, tex_c, - mapping_type) - -def write_obj_with_colors_texture_core(string filename, string mtl_name, - np.ndarray[float, ndim=2, mode = "c"] vertices not None, - np.ndarray[int, ndim=2, mode="c"] triangles not None, - np.ndarray[float, ndim=2, mode = "c"] colors not None, - np.ndarray[float, ndim=2, mode = "c"] uv_coords not None, - int nver, int ntri, int ntexver - ): - _write_obj_with_colors_texture(filename, mtl_name, - np.PyArray_DATA(vertices), np.PyArray_DATA(triangles), np.PyArray_DATA(colors), np.PyArray_DATA(uv_coords), - nver, ntri, ntexver) diff --git a/venv/insightface/thirdparty/face3d/mesh/cython/setup.py b/venv/insightface/thirdparty/face3d/mesh/cython/setup.py deleted file mode 100644 index fc1ede1e777b369b1bbd9c94f6fa85ce037de3f9..0000000000000000000000000000000000000000 --- a/venv/insightface/thirdparty/face3d/mesh/cython/setup.py +++ /dev/null @@ -1,20 +0,0 @@ -''' -python setup.py build_ext -i -to compile -''' - -# setup.py -from distutils.core import setup, Extension -from Cython.Build import cythonize -from Cython.Distutils import build_ext -import numpy - -setup( - name = 'mesh_core_cython', - cmdclass={'build_ext': build_ext}, - ext_modules=[Extension("mesh_core_cython", - sources=["mesh_core_cython.pyx", "mesh_core.cpp"], - language='c++', - include_dirs=[numpy.get_include()])], -) - diff --git a/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/LICENSE.rst b/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/LICENSE.rst deleted file mode 100644 index 191ddaf31642db9d52a7f32826043c6550c299d8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/LICENSE.rst +++ /dev/null @@ -1,21 +0,0 @@ -The MIT License (MIT) - -Copyright (c) 2017 Laurent LAPORTE - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. \ No newline at end of file diff --git a/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/METADATA b/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/METADATA deleted file mode 100644 index 86dc1909809bbd0a14101b42d8fa3cedf6ba6703..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/METADATA +++ /dev/null @@ -1,195 +0,0 @@ -Metadata-Version: 2.2 -Name: Deprecated -Version: 1.2.18 -Summary: Python @deprecated decorator to deprecate old python classes, functions or methods. -Home-page: https://github.com/laurent-laporte-pro/deprecated -Author: Laurent LAPORTE -Author-email: laurent.laporte.pro@gmail.com -License: MIT -Project-URL: Documentation, https://deprecated.readthedocs.io/en/latest/ -Project-URL: Source, https://github.com/laurent-laporte-pro/deprecated -Project-URL: Bug Tracker, https://github.com/laurent-laporte-pro/deprecated/issues -Keywords: deprecate,deprecated,deprecation,warning,warn,decorator -Platform: any -Classifier: Development Status :: 5 - Production/Stable -Classifier: Environment :: Web Environment -Classifier: Intended Audience :: Developers -Classifier: License :: OSI Approved :: MIT License -Classifier: Operating System :: OS Independent -Classifier: Programming Language :: Python -Classifier: Programming Language :: Python :: 2 -Classifier: Programming Language :: Python :: 2.7 -Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3.4 -Classifier: Programming Language :: Python :: 3.5 -Classifier: Programming Language :: Python :: 3.6 -Classifier: Programming Language :: Python :: 3.7 -Classifier: Programming Language :: Python :: 3.8 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Programming Language :: Python :: 3.10 -Classifier: Programming Language :: Python :: 3.11 -Classifier: Programming Language :: Python :: 3.12 -Classifier: Topic :: Software Development :: Libraries :: Python Modules -Requires-Python: >=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.* -Description-Content-Type: text/x-rst -License-File: LICENSE.rst -Requires-Dist: wrapt<2,>=1.10 -Provides-Extra: dev -Requires-Dist: tox; extra == "dev" -Requires-Dist: PyTest; extra == "dev" -Requires-Dist: PyTest-Cov; extra == "dev" -Requires-Dist: bump2version<1; extra == "dev" -Requires-Dist: setuptools; python_version >= "3.12" and extra == "dev" -Dynamic: author -Dynamic: author-email -Dynamic: classifier -Dynamic: description -Dynamic: description-content-type -Dynamic: home-page -Dynamic: keywords -Dynamic: license -Dynamic: platform -Dynamic: project-url -Dynamic: provides-extra -Dynamic: requires-dist -Dynamic: requires-python -Dynamic: summary - - -Deprecated Library ------------------- - -Deprecated is Easy to Use -````````````````````````` - -If you need to mark a function or a method as deprecated, -you can use the ``@deprecated`` decorator: - -Save in a hello.py: - -.. code:: python - - from deprecated import deprecated - - - @deprecated(version='1.2.1', reason="You should use another function") - def some_old_function(x, y): - return x + y - - - class SomeClass(object): - @deprecated(version='1.3.0', reason="This method is deprecated") - def some_old_method(self, x, y): - return x + y - - - some_old_function(12, 34) - obj = SomeClass() - obj.some_old_method(5, 8) - - -And Easy to Setup -````````````````` - -And run it: - -.. code:: bash - - $ pip install Deprecated - $ python hello.py - hello.py:15: DeprecationWarning: Call to deprecated function (or staticmethod) some_old_function. - (You should use another function) -- Deprecated since version 1.2.0. - some_old_function(12, 34) - hello.py:17: DeprecationWarning: Call to deprecated method some_old_method. - (This method is deprecated) -- Deprecated since version 1.3.0. - obj.some_old_method(5, 8) - - -You can document your code -`````````````````````````` - -Have you ever wonder how to document that some functions, classes, methods, etc. are deprecated? -This is now possible with the integrated Sphinx directives: - -For instance, in hello_sphinx.py: - -.. code:: python - - from deprecated.sphinx import deprecated - from deprecated.sphinx import versionadded - from deprecated.sphinx import versionchanged - - - @versionadded(version='1.0', reason="This function is new") - def function_one(): - '''This is the function one''' - - - @versionchanged(version='1.0', reason="This function is modified") - def function_two(): - '''This is the function two''' - - - @deprecated(version='1.0', reason="This function will be removed soon") - def function_three(): - '''This is the function three''' - - - function_one() - function_two() - function_three() # warns - - help(function_one) - help(function_two) - help(function_three) - - -The result it immediate -``````````````````````` - -Run it: - -.. code:: bash - - $ python hello_sphinx.py - - hello_sphinx.py:23: DeprecationWarning: Call to deprecated function (or staticmethod) function_three. - (This function will be removed soon) -- Deprecated since version 1.0. - function_three() # warns - - Help on function function_one in module __main__: - - function_one() - This is the function one - - .. versionadded:: 1.0 - This function is new - - Help on function function_two in module __main__: - - function_two() - This is the function two - - .. versionchanged:: 1.0 - This function is modified - - Help on function function_three in module __main__: - - function_three() - This is the function three - - .. deprecated:: 1.0 - This function will be removed soon - - -Links -````` - -* `Python package index (PyPi) `_ -* `GitHub website `_ -* `Read The Docs `_ -* `EBook on Lulu.com `_ -* `StackOverFlow Q&A `_ -* `Development version - `_ - diff --git a/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/RECORD b/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/RECORD deleted file mode 100644 index a6578c4734b87ec6dbbc826ecafe96fb1280fb16..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/RECORD +++ /dev/null @@ -1,12 +0,0 @@ -Deprecated-1.2.18.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -Deprecated-1.2.18.dist-info/LICENSE.rst,sha256=HoPt0VvkGbXVveNy4yXlJ_9PmRX1SOfHUxS0H2aZ6Dw,1081 -Deprecated-1.2.18.dist-info/METADATA,sha256=4CrUw5Bl8_NsBuZYe0Nw-mIwQnVpT1CnmBYU9BqOuq8,5725 -Deprecated-1.2.18.dist-info/RECORD,, -Deprecated-1.2.18.dist-info/WHEEL,sha256=9Hm2OB-j1QcCUq9Jguht7ayGIIZBRTdOXD1qg9cCgPM,109 -Deprecated-1.2.18.dist-info/top_level.txt,sha256=nHbOYawKPQQE5lQl-toUB1JBRJjUyn_m_Mb8RVJ0RjA,11 -deprecated/__init__.py,sha256=yZNbmDKXF4PLtp_Ikdb_9ObJLkHuFSUHvqidFTKKGFM,351 -deprecated/__pycache__/__init__.cpython-310.pyc,, -deprecated/__pycache__/classic.cpython-310.pyc,, -deprecated/__pycache__/sphinx.cpython-310.pyc,, -deprecated/classic.py,sha256=7WXOt4Vf1NhrUznm8ypjS50CMyAdZwrGT58Lhb8fW14,10609 -deprecated/sphinx.py,sha256=cOKnXbDyFAwDr5O7HBEpgQrx-J-qfp57sfdK_LabDxs,11109 diff --git a/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/WHEEL b/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/WHEEL deleted file mode 100644 index eaea6f3b57e8165ab7d2673900a9939348ee5e42..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/WHEEL +++ /dev/null @@ -1,6 +0,0 @@ -Wheel-Version: 1.0 -Generator: setuptools (75.8.0) -Root-Is-Purelib: true -Tag: py2-none-any -Tag: py3-none-any - diff --git a/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/top_level.txt deleted file mode 100644 index 9f8d5502dae589d488c5107e99768ae5023bb4ea..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/Deprecated-1.2.18.dist-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -deprecated diff --git a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/AUTHORS b/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/AUTHORS deleted file mode 100644 index 45b14c961cdbbebb245ad77d0cd0c8b368bc9aca..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/AUTHORS +++ /dev/null @@ -1,59 +0,0 @@ -GitPython was originally written by Michael Trier. -GitPython 0.2 was partially (re)written by Sebastian Thiel, based on 0.1.6 and git-dulwich. - -Contributors are: - --Michael Trier --Alan Briolat --Florian Apolloner --David Aguilar --Jelmer Vernooij --Steve Frécinaux --Kai Lautaportti --Paul Sowden --Sebastian Thiel --Jonathan Chu --Vincent Driessen --Phil Elson --Bernard `Guyzmo` Pratz --Timothy B. Hartman --Konstantin Popov --Peter Jones --Anson Mansfield --Ken Odegard --Alexis Horgix Chotard --Piotr Babij --Mikuláš Poul --Charles Bouchard-Légaré --Yaroslav Halchenko --Tim Swast --William Luc Ritchie --David Host --A. Jesse Jiryu Davis --Steven Whitman --Stefan Stancu --César Izurieta --Arthur Milchior --Anil Khatri --JJ Graham --Ben Thayer --Dries Kennes --Pratik Anurag --Harmon --Liam Beguin --Ram Rachum --Alba Mendez --Robert Westman --Hugo van Kemenade --Hiroki Tokunaga --Julien Mauroy --Patrick Gerard --Luke Twist --Joseph Hale --Santos Gallegos --Wenhan Zhu --Eliah Kagan --Ethan Lin --Jonas Scharpf - -Portions derived from other open source works and are clearly marked. diff --git a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/LICENSE b/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/LICENSE deleted file mode 100644 index ba8a219fe1f27c10b50df8cd4f26c0ab833bbbc8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/LICENSE +++ /dev/null @@ -1,29 +0,0 @@ -Copyright (C) 2008, 2009 Michael Trier and contributors -All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions -are met: - -* Redistributions of source code must retain the above copyright -notice, this list of conditions and the following disclaimer. - -* Redistributions in binary form must reproduce the above copyright -notice, this list of conditions and the following disclaimer in the -documentation and/or other materials provided with the distribution. - -* Neither the name of the GitPython project nor the names of -its contributors may be used to endorse or promote products derived -from this software without specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS -"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT -LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR -A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT -OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, -SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED -TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR -PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF -LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING -NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS -SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/METADATA b/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/METADATA deleted file mode 100644 index 2589bfb3dee272a84f3b80aedaedf9a80c475499..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/METADATA +++ /dev/null @@ -1,295 +0,0 @@ -Metadata-Version: 2.1 -Name: GitPython -Version: 3.1.44 -Summary: GitPython is a Python library used to interact with Git repositories -Home-page: https://github.com/gitpython-developers/GitPython -Author: Sebastian Thiel, Michael Trier -Author-email: byronimo@gmail.com, mtrier@gmail.com -License: BSD-3-Clause -Classifier: Development Status :: 5 - Production/Stable -Classifier: Environment :: Console -Classifier: Intended Audience :: Developers -Classifier: License :: OSI Approved :: BSD License -Classifier: Operating System :: OS Independent -Classifier: Operating System :: POSIX -Classifier: Operating System :: Microsoft :: Windows -Classifier: Operating System :: MacOS :: MacOS X -Classifier: Typing :: Typed -Classifier: Programming Language :: Python -Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3.7 -Classifier: Programming Language :: Python :: 3.8 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Programming Language :: Python :: 3.10 -Classifier: Programming Language :: Python :: 3.11 -Classifier: Programming Language :: Python :: 3.12 -Requires-Python: >=3.7 -Description-Content-Type: text/markdown -License-File: LICENSE -License-File: AUTHORS -Requires-Dist: gitdb<5,>=4.0.1 -Requires-Dist: typing-extensions>=3.7.4.3; python_version < "3.8" -Provides-Extra: test -Requires-Dist: coverage[toml]; extra == "test" -Requires-Dist: ddt!=1.4.3,>=1.1.1; extra == "test" -Requires-Dist: mock; python_version < "3.8" and extra == "test" -Requires-Dist: mypy; extra == "test" -Requires-Dist: pre-commit; extra == "test" -Requires-Dist: pytest>=7.3.1; extra == "test" -Requires-Dist: pytest-cov; extra == "test" -Requires-Dist: pytest-instafail; extra == "test" -Requires-Dist: pytest-mock; extra == "test" -Requires-Dist: pytest-sugar; extra == "test" -Requires-Dist: typing-extensions; python_version < "3.11" and extra == "test" -Provides-Extra: doc -Requires-Dist: sphinx<7.2,>=7.1.2; extra == "doc" -Requires-Dist: sphinx_rtd_theme; extra == "doc" -Requires-Dist: sphinx-autodoc-typehints; extra == "doc" - -![Python package](https://github.com/gitpython-developers/GitPython/workflows/Python%20package/badge.svg) -[![Documentation Status](https://readthedocs.org/projects/gitpython/badge/?version=stable)](https://readthedocs.org/projects/gitpython/?badge=stable) -[![Packaging status](https://repology.org/badge/tiny-repos/python:gitpython.svg)](https://repology.org/metapackage/python:gitpython/versions) - -## [Gitoxide](https://github.com/Byron/gitoxide): A peek into the future… - -I started working on GitPython in 2009, back in the days when Python was 'my thing' and I had great plans with it. -Of course, back in the days, I didn't really know what I was doing and this shows in many places. Somewhat similar to -Python this happens to be 'good enough', but at the same time is deeply flawed and broken beyond repair. - -By now, GitPython is widely used and I am sure there is a good reason for that, it's something to be proud of and happy about. -The community is maintaining the software and is keeping it relevant for which I am absolutely grateful. For the time to come I am happy to continue maintaining GitPython, remaining hopeful that one day it won't be needed anymore. - -More than 15 years after my first meeting with 'git' I am still in excited about it, and am happy to finally have the tools and -probably the skills to scratch that itch of mine: implement `git` in a way that makes tool creation a piece of cake for most. - -If you like the idea and want to learn more, please head over to [gitoxide](https://github.com/Byron/gitoxide), an -implementation of 'git' in [Rust](https://www.rust-lang.org). - -*(Please note that `gitoxide` is not currently available for use in Python, and that Rust is required.)* - -## GitPython - -GitPython is a python library used to interact with git repositories, high-level like git-porcelain, -or low-level like git-plumbing. - -It provides abstractions of git objects for easy access of repository data often backed by calling the `git` -command-line program. - -### DEVELOPMENT STATUS - -This project is in **maintenance mode**, which means that - -- …there will be no feature development, unless these are contributed -- …there will be no bug fixes, unless they are relevant to the safety of users, or contributed -- …issues will be responded to with waiting times of up to a month - -The project is open to contributions of all kinds, as well as new maintainers. - -### REQUIREMENTS - -GitPython needs the `git` executable to be installed on the system and available in your -`PATH` for most operations. If it is not in your `PATH`, you can help GitPython find it -by setting the `GIT_PYTHON_GIT_EXECUTABLE=` environment variable. - -- Git (1.7.x or newer) -- Python >= 3.7 - -The list of dependencies are listed in `./requirements.txt` and `./test-requirements.txt`. -The installer takes care of installing them for you. - -### INSTALL - -GitPython and its required package dependencies can be installed in any of the following ways, all of which should typically be done in a [virtual environment](https://docs.python.org/3/tutorial/venv.html). - -#### From PyPI - -To obtain and install a copy [from PyPI](https://pypi.org/project/GitPython/), run: - -```sh -pip install GitPython -``` - -(A distribution package can also be downloaded for manual installation at [the PyPI page](https://pypi.org/project/GitPython/).) - -#### From downloaded source code - -If you have downloaded the source code, run this from inside the unpacked `GitPython` directory: - -```sh -pip install . -``` - -#### By cloning the source code repository - -To clone the [the GitHub repository](https://github.com/gitpython-developers/GitPython) from source to work on the code, you can do it like so: - -```sh -git clone https://github.com/gitpython-developers/GitPython -cd GitPython -./init-tests-after-clone.sh -``` - -On Windows, `./init-tests-after-clone.sh` can be run in a Git Bash shell. - -If you are cloning [your own fork](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/about-forks), then replace the above `git clone` command with one that gives the URL of your fork. Or use this [`gh`](https://cli.github.com/) command (assuming you have `gh` and your fork is called `GitPython`): - -```sh -gh repo clone GitPython -``` - -Having cloned the repo, create and activate your [virtual environment](https://docs.python.org/3/tutorial/venv.html). - -Then make an [editable install](https://pip.pypa.io/en/stable/topics/local-project-installs/#editable-installs): - -```sh -pip install -e ".[test]" -``` - -In the less common case that you do not want to install test dependencies, `pip install -e .` can be used instead. - -#### With editable *dependencies* (not preferred, and rarely needed) - -In rare cases, you may want to work on GitPython and one or both of its [gitdb](https://github.com/gitpython-developers/gitdb) and [smmap](https://github.com/gitpython-developers/smmap) dependencies at the same time, with changes in your local working copy of gitdb or smmap immediately reflected in the behavior of your local working copy of GitPython. This can be done by making editable installations of those dependencies in the same virtual environment where you install GitPython. - -If you want to do that *and* you want the versions in GitPython's git submodules to be used, then pass `-e git/ext/gitdb` and/or `-e git/ext/gitdb/gitdb/ext/smmap` to `pip install`. This can be done in any order, and in separate `pip install` commands or the same one, so long as `-e` appears before *each* path. For example, you can install GitPython, gitdb, and smmap editably in the currently active virtual environment this way: - -```sh -pip install -e ".[test]" -e git/ext/gitdb -e git/ext/gitdb/gitdb/ext/smmap -``` - -The submodules must have been cloned for that to work, but that will already be the case if you have run `./init-tests-after-clone.sh`. You can use `pip list` to check which packages are installed editably and which are installed normally. - -To reiterate, this approach should only rarely be used. For most development it is preferable to allow the gitdb and smmap dependencices to be retrieved automatically from PyPI in their latest stable packaged versions. - -### Limitations - -#### Leakage of System Resources - -GitPython is not suited for long-running processes (like daemons) as it tends to -leak system resources. It was written in a time where destructors (as implemented -in the `__del__` method) still ran deterministically. - -In case you still want to use it in such a context, you will want to search the -codebase for `__del__` implementations and call these yourself when you see fit. - -Another way assure proper cleanup of resources is to factor out GitPython into a -separate process which can be dropped periodically. - -#### Windows support - -See [Issue #525](https://github.com/gitpython-developers/GitPython/issues/525). - -### RUNNING TESTS - -_Important_: Right after cloning this repository, please be sure to have executed -the `./init-tests-after-clone.sh` script in the repository root. Otherwise -you will encounter test failures. - -#### Install test dependencies - -Ensure testing libraries are installed. This is taken care of already if you installed with: - -```sh -pip install -e ".[test]" -``` - -If you had installed with a command like `pip install -e .` instead, you can still run -the above command to add the testing dependencies. - -#### Test commands - -To test, run: - -```sh -pytest -``` - -To lint, and apply some linting fixes as well as automatic code formatting, run: - -```sh -pre-commit run --all-files -``` - -This includes the linting and autoformatting done by Ruff, as well as some other checks. - -To typecheck, run: - -```sh -mypy -``` - -#### CI (and tox) - -Style and formatting checks, and running tests on all the different supported Python versions, will be performed: - -- Upon submitting a pull request. -- On each push, *if* you have a fork with GitHub Actions enabled. -- Locally, if you run [`tox`](https://tox.wiki/) (this skips any Python versions you don't have installed). - -#### Configuration files - -Specific tools are all configured in the `./pyproject.toml` file: - -- `pytest` (test runner) -- `coverage.py` (code coverage) -- `ruff` (linter and formatter) -- `mypy` (type checker) - -Orchestration tools: - -- Configuration for `pre-commit` is in the `./.pre-commit-config.yaml` file. -- Configuration for `tox` is in `./tox.ini`. -- Configuration for GitHub Actions (CI) is in files inside `./.github/workflows/`. - -### Contributions - -Please have a look at the [contributions file][contributing]. - -### INFRASTRUCTURE - -- [User Documentation](http://gitpython.readthedocs.org) -- [Questions and Answers](http://stackexchange.com/filters/167317/gitpython) -- Please post on Stack Overflow and use the `gitpython` tag -- [Issue Tracker](https://github.com/gitpython-developers/GitPython/issues) - - Post reproducible bugs and feature requests as a new issue. - Please be sure to provide the following information if posting bugs: - - GitPython version (e.g. `import git; git.__version__`) - - Python version (e.g. `python --version`) - - The encountered stack-trace, if applicable - - Enough information to allow reproducing the issue - -### How to make a new release - -1. Update/verify the **version** in the `VERSION` file. -2. Update/verify that the `doc/source/changes.rst` changelog file was updated. It should include a link to the forthcoming release page: `https://github.com/gitpython-developers/GitPython/releases/tag/` -3. Commit everything. -4. Run `git tag -s ` to tag the version in Git. -5. _Optionally_ create and activate a [virtual environment](https://packaging.python.org/en/latest/guides/installing-using-pip-and-virtual-environments/#creating-a-virtual-environment). (Then the next step can install `build` and `twine`.) -6. Run `make release`. -7. Go to [GitHub Releases](https://github.com/gitpython-developers/GitPython/releases) and publish a new one with the recently pushed tag. Generate the changelog. - -### Projects using GitPython - -- [PyDriller](https://github.com/ishepard/pydriller) -- [Kivy Designer](https://github.com/kivy/kivy-designer) -- [Prowl](https://github.com/nettitude/Prowl) -- [Python Taint](https://github.com/python-security/pyt) -- [Buster](https://github.com/axitkhurana/buster) -- [git-ftp](https://github.com/ezyang/git-ftp) -- [Git-Pandas](https://github.com/wdm0006/git-pandas) -- [PyGitUp](https://github.com/msiemens/PyGitUp) -- [PyJFuzz](https://github.com/mseclab/PyJFuzz) -- [Loki](https://github.com/Neo23x0/Loki) -- [Omniwallet](https://github.com/OmniLayer/omniwallet) -- [GitViper](https://github.com/BeayemX/GitViper) -- [Git Gud](https://github.com/bthayer2365/git-gud) - -### LICENSE - -[3-Clause BSD License](https://opensource.org/license/bsd-3-clause/), also known as the New BSD License. See the [LICENSE file][license]. - -One file exclusively used for fuzz testing is subject to [a separate license, detailed here](./fuzzing/README.md#license). -This file is not included in the wheel or sdist packages published by the maintainers of GitPython. - -[contributing]: https://github.com/gitpython-developers/GitPython/blob/main/CONTRIBUTING.md -[license]: https://github.com/gitpython-developers/GitPython/blob/main/LICENSE diff --git a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/RECORD b/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/RECORD deleted file mode 100644 index 7785311309cb2f00f2f01a1b6aed9301a7b93e5e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/RECORD +++ /dev/null @@ -1,83 +0,0 @@ -GitPython-3.1.44.dist-info/AUTHORS,sha256=tZ9LuyBks2V2HKTPK7kCmtd9Guu_LyU1oZHvU0NiAok,2334 -GitPython-3.1.44.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -GitPython-3.1.44.dist-info/LICENSE,sha256=hvyUwyGpr7wRUUcTURuv3tIl8lEA3MD3NQ6CvCMbi-s,1503 -GitPython-3.1.44.dist-info/METADATA,sha256=0O_Fr2Y7A-DlPYhlbSxGjblBC2mWkw3USNUhyL80Ip8,13245 -GitPython-3.1.44.dist-info/RECORD,, -GitPython-3.1.44.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -GitPython-3.1.44.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91 -GitPython-3.1.44.dist-info/top_level.txt,sha256=0hzDuIp8obv624V3GmbqsagBWkk8ohtGU-Bc1PmTT0o,4 -git/__init__.py,sha256=nkQImgv-bWdiZOFDjzN-gbt93FoRHD0nY6_t9LQxy4Y,8899 -git/__pycache__/__init__.cpython-310.pyc,, -git/__pycache__/cmd.cpython-310.pyc,, -git/__pycache__/compat.cpython-310.pyc,, -git/__pycache__/config.cpython-310.pyc,, -git/__pycache__/db.cpython-310.pyc,, -git/__pycache__/diff.cpython-310.pyc,, -git/__pycache__/exc.cpython-310.pyc,, -git/__pycache__/remote.cpython-310.pyc,, -git/__pycache__/types.cpython-310.pyc,, -git/__pycache__/util.cpython-310.pyc,, -git/cmd.py,sha256=QwiaBy0mFbi9xjRKhRgUVK-_-K6xVdFqh9l0cxPqPSc,67724 -git/compat.py,sha256=y1E6y6O2q5r8clSlr8ZNmuIWG9nmHuehQEsVsmBffs8,4526 -git/config.py,sha256=vTUlK6d8ORqFqjOv4Vbq_Hm-5mp-jOAt1dkq0IdzJ3U,34933 -git/db.py,sha256=vIW9uWSbqu99zbuU2ZDmOhVOv1UPTmxrnqiCtRHCfjE,2368 -git/diff.py,sha256=wmpMCIdMiVOqreGVPOGYyO4gFboGOAicyrvvI7PPjEg,27095 -git/exc.py,sha256=Gc7g1pHpn8OmTse30NHmJVsBJ2CYH8LxaR8y8UA3lIM,7119 -git/index/__init__.py,sha256=i-Nqb8Lufp9aFbmxpQBORmmQnjEVVM1Pn58fsQkyGgQ,406 -git/index/__pycache__/__init__.cpython-310.pyc,, -git/index/__pycache__/base.cpython-310.pyc,, -git/index/__pycache__/fun.cpython-310.pyc,, -git/index/__pycache__/typ.cpython-310.pyc,, -git/index/__pycache__/util.cpython-310.pyc,, -git/index/base.py,sha256=nDD7XVLNbgBKpJMrrTVyHBy6NVLWgDkk7oUw6ZOegPc,60808 -git/index/fun.py,sha256=37cA3DBC9vpAnSVu5TGA072SnoF5XZOkOukExwlejHs,16736 -git/index/typ.py,sha256=uuKNwitUw83FhVaLSwo4pY7PHDQudtZTLJrLGym4jcI,6570 -git/index/util.py,sha256=fULi7GPG-MvprKrRCD5c15GNdzku_1E38We0d97WB3A,3659 -git/objects/__init__.py,sha256=O6ZL_olX7e5-8iIbKviRPkVSJxN37WA-EC0q9d48U5Y,637 -git/objects/__pycache__/__init__.cpython-310.pyc,, -git/objects/__pycache__/base.cpython-310.pyc,, -git/objects/__pycache__/blob.cpython-310.pyc,, -git/objects/__pycache__/commit.cpython-310.pyc,, -git/objects/__pycache__/fun.cpython-310.pyc,, -git/objects/__pycache__/tag.cpython-310.pyc,, -git/objects/__pycache__/tree.cpython-310.pyc,, -git/objects/__pycache__/util.cpython-310.pyc,, -git/objects/base.py,sha256=0dqNkSRVH0mk0-7ZKIkGBK7iNYrzLTVxwQFUd6CagsE,10277 -git/objects/blob.py,sha256=zwwq0KfOMYeP5J2tW5CQatoLyeqFRlfkxP1Vwx1h07s,1215 -git/objects/commit.py,sha256=GH1_83C9t7RGTukwozTHDgvxYQPRjTHhPDkXJyBbJyo,30553 -git/objects/fun.py,sha256=B4jCqhAjm6Hl79GK58FPzW1H9K6Wc7Tx0rssyWmAcEE,8935 -git/objects/submodule/__init__.py,sha256=6xySp767LVz3UylWgUalntS_nGXRuVzXxDuFAv_Wc2c,303 -git/objects/submodule/__pycache__/__init__.cpython-310.pyc,, -git/objects/submodule/__pycache__/base.cpython-310.pyc,, -git/objects/submodule/__pycache__/root.cpython-310.pyc,, -git/objects/submodule/__pycache__/util.cpython-310.pyc,, -git/objects/submodule/base.py,sha256=MQ-2xV8JznGwy2hLQv1aeQNgAkhBhgc5tdtClFL3DmE,63901 -git/objects/submodule/root.py,sha256=5eTtYNHasqdPq6q0oDCPr7IaO6uAHL3b4DxMoiO2LhE,20246 -git/objects/submodule/util.py,sha256=sQqAYaiSJdFkZa9NlAuK_wTsMNiS-kkQnQjvIoJtc_o,3509 -git/objects/tag.py,sha256=jAGESnpmTEv-dLakPzheT5ILZFFArcItnXYqfxfDrgc,4441 -git/objects/tree.py,sha256=jJH888SHiP4dGzE-ra1yenQOyya_0C_MkHr06c1gHpM,13849 -git/objects/util.py,sha256=Nlza4zLgdPmr_Yasyvvs6c1rKtW_wMxI6wDmQpQ3ufw,23846 -git/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -git/refs/__init__.py,sha256=DWlJNnsx-4jM_E-VycbP-FZUdn6iWhjnH_uZ_pZXBro,509 -git/refs/__pycache__/__init__.cpython-310.pyc,, -git/refs/__pycache__/head.cpython-310.pyc,, -git/refs/__pycache__/log.cpython-310.pyc,, -git/refs/__pycache__/reference.cpython-310.pyc,, -git/refs/__pycache__/remote.cpython-310.pyc,, -git/refs/__pycache__/symbolic.cpython-310.pyc,, -git/refs/__pycache__/tag.cpython-310.pyc,, -git/refs/head.py,sha256=SGa3N301HfAi79X6UR5Mcg7mO9TnCH3Bk549kHlJVaQ,10513 -git/refs/log.py,sha256=kXiuAgTo1DIuM_BfbDUk9gQ0YO-mutIMVdHv1_ES90o,12493 -git/refs/reference.py,sha256=l6mhF4YLSEwtjz6b9PpOQH-fkng7EYWMaJhkjn-2jXA,5630 -git/refs/remote.py,sha256=WwqV9T7BbYf3F_WZNUQivu9xktIIKGklCjDpwQrhD-A,2806 -git/refs/symbolic.py,sha256=c8zOwaqzcg-J-rGrpuWdvh8zwMvSUqAHghd4vJoYG_s,34552 -git/refs/tag.py,sha256=kgzV2vhpL4FD2TqHb0BJuMRAHgAvJF-TcoyWlaB-djQ,5010 -git/remote.py,sha256=pYn9dAlz-QwvNMWXD1M57pMPQitthOM86qTRK_cpTqU,46786 -git/repo/__init__.py,sha256=CILSVH36fX_WxVFSjD9o1WF5LgsNedPiJvSngKZqfVU,210 -git/repo/__pycache__/__init__.cpython-310.pyc,, -git/repo/__pycache__/base.cpython-310.pyc,, -git/repo/__pycache__/fun.cpython-310.pyc,, -git/repo/base.py,sha256=0GU6nKNdT8SYjDI5Y5DeZ1zCEX3tHeq1VW2MSpne05g,59891 -git/repo/fun.py,sha256=HSGC0-rqeKKx9fDg7JyQyMZgIwUWn-FnSZR_gRGpG-E,13573 -git/types.py,sha256=MQzIDEOnoueXGsAJF_0MgUc_osH7Eu0Sw3DQofYzCVE,10272 -git/util.py,sha256=2uAv34zZ_827-zJ3-D5ACrVH-4Q4EO_KLUTH23zi2AI,43770 diff --git a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/REQUESTED deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/WHEEL b/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/WHEEL deleted file mode 100644 index ae527e7d64811439e61b93aa375defb30e06edfe..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/WHEEL +++ /dev/null @@ -1,5 +0,0 @@ -Wheel-Version: 1.0 -Generator: setuptools (75.6.0) -Root-Is-Purelib: true -Tag: py3-none-any - diff --git a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/top_level.txt deleted file mode 100644 index 5664e303b5dc2e9ef8e14a0845d9486ec1920afd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/GitPython-3.1.44.dist-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -git diff --git a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/COPYING b/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/COPYING deleted file mode 100644 index 94a9ed024d3859793618152ea559a168bbcbb5e2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/COPYING +++ /dev/null @@ -1,674 +0,0 @@ - GNU GENERAL PUBLIC LICENSE - Version 3, 29 June 2007 - - Copyright (C) 2007 Free Software Foundation, Inc. - Everyone is permitted to copy and distribute verbatim copies - of this license document, but changing it is not allowed. - - Preamble - - The GNU General Public License is a free, copyleft license for -software and other kinds of works. - - The licenses for most software and other practical works are designed -to take away your freedom to share and change the works. By contrast, -the GNU General Public License is intended to guarantee your freedom to -share and change all versions of a program--to make sure it remains free -software for all its users. We, the Free Software Foundation, use the -GNU General Public License for most of our software; it applies also to -any other work released this way by its authors. You can apply it to -your programs, too. - - When we speak of free software, we are referring to freedom, not -price. Our General Public Licenses are designed to make sure that you -have the freedom to distribute copies of free software (and charge for -them if you wish), that you receive source code or can get it if you -want it, that you can change the software or use pieces of it in new -free programs, and that you know you can do these things. - - To protect your rights, we need to prevent others from denying you -these rights or asking you to surrender the rights. Therefore, you have -certain responsibilities if you distribute copies of the software, or if -you modify it: responsibilities to respect the freedom of others. - - For example, if you distribute copies of such a program, whether -gratis or for a fee, you must pass on to the recipients the same -freedoms that you received. You must make sure that they, too, receive -or can get the source code. And you must show them these terms so they -know their rights. - - Developers that use the GNU GPL protect your rights with two steps: -(1) assert copyright on the software, and (2) offer you this License -giving you legal permission to copy, distribute and/or modify it. - - For the developers' and authors' protection, the GPL clearly explains -that there is no warranty for this free software. For both users' and -authors' sake, the GPL requires that modified versions be marked as -changed, so that their problems will not be attributed erroneously to -authors of previous versions. - - Some devices are designed to deny users access to install or run -modified versions of the software inside them, although the manufacturer -can do so. This is fundamentally incompatible with the aim of -protecting users' freedom to change the software. The systematic -pattern of such abuse occurs in the area of products for individuals to -use, which is precisely where it is most unacceptable. Therefore, we -have designed this version of the GPL to prohibit the practice for those -products. If such problems arise substantially in other domains, we -stand ready to extend this provision to those domains in future versions -of the GPL, as needed to protect the freedom of users. - - Finally, every program is threatened constantly by software patents. -States should not allow patents to restrict development and use of -software on general-purpose computers, but in those that do, we wish to -avoid the special danger that patents applied to a free program could -make it effectively proprietary. To prevent this, the GPL assures that -patents cannot be used to render the program non-free. - - The precise terms and conditions for copying, distribution and -modification follow. - - TERMS AND CONDITIONS - - 0. Definitions. - - "This License" refers to version 3 of the GNU General Public License. - - "Copyright" also means copyright-like laws that apply to other kinds of -works, such as semiconductor masks. - - "The Program" refers to any copyrightable work licensed under this -License. Each licensee is addressed as "you". "Licensees" and -"recipients" may be individuals or organizations. - - To "modify" a work means to copy from or adapt all or part of the work -in a fashion requiring copyright permission, other than the making of an -exact copy. The resulting work is called a "modified version" of the -earlier work or a work "based on" the earlier work. - - A "covered work" means either the unmodified Program or a work based -on the Program. - - To "propagate" a work means to do anything with it that, without -permission, would make you directly or secondarily liable for -infringement under applicable copyright law, except executing it on a -computer or modifying a private copy. Propagation includes copying, -distribution (with or without modification), making available to the -public, and in some countries other activities as well. - - To "convey" a work means any kind of propagation that enables other -parties to make or receive copies. Mere interaction with a user through -a computer network, with no transfer of a copy, is not conveying. - - An interactive user interface displays "Appropriate Legal Notices" -to the extent that it includes a convenient and prominently visible -feature that (1) displays an appropriate copyright notice, and (2) -tells the user that there is no warranty for the work (except to the -extent that warranties are provided), that licensees may convey the -work under this License, and how to view a copy of this License. If -the interface presents a list of user commands or options, such as a -menu, a prominent item in the list meets this criterion. - - 1. Source Code. - - The "source code" for a work means the preferred form of the work -for making modifications to it. "Object code" means any non-source -form of a work. - - A "Standard Interface" means an interface that either is an official -standard defined by a recognized standards body, or, in the case of -interfaces specified for a particular programming language, one that -is widely used among developers working in that language. - - The "System Libraries" of an executable work include anything, other -than the work as a whole, that (a) is included in the normal form of -packaging a Major Component, but which is not part of that Major -Component, and (b) serves only to enable use of the work with that -Major Component, or to implement a Standard Interface for which an -implementation is available to the public in source code form. A -"Major Component", in this context, means a major essential component -(kernel, window system, and so on) of the specific operating system -(if any) on which the executable work runs, or a compiler used to -produce the work, or an object code interpreter used to run it. - - The "Corresponding Source" for a work in object code form means all -the source code needed to generate, install, and (for an executable -work) run the object code and to modify the work, including scripts to -control those activities. However, it does not include the work's -System Libraries, or general-purpose tools or generally available free -programs which are used unmodified in performing those activities but -which are not part of the work. For example, Corresponding Source -includes interface definition files associated with source files for -the work, and the source code for shared libraries and dynamically -linked subprograms that the work is specifically designed to require, -such as by intimate data communication or control flow between those -subprograms and other parts of the work. - - The Corresponding Source need not include anything that users -can regenerate automatically from other parts of the Corresponding -Source. - - The Corresponding Source for a work in source code form is that -same work. - - 2. Basic Permissions. - - All rights granted under this License are granted for the term of -copyright on the Program, and are irrevocable provided the stated -conditions are met. This License explicitly affirms your unlimited -permission to run the unmodified Program. The output from running a -covered work is covered by this License only if the output, given its -content, constitutes a covered work. This License acknowledges your -rights of fair use or other equivalent, as provided by copyright law. - - You may make, run and propagate covered works that you do not -convey, without conditions so long as your license otherwise remains -in force. You may convey covered works to others for the sole purpose -of having them make modifications exclusively for you, or provide you -with facilities for running those works, provided that you comply with -the terms of this License in conveying all material for which you do -not control copyright. Those thus making or running the covered works -for you must do so exclusively on your behalf, under your direction -and control, on terms that prohibit them from making any copies of -your copyrighted material outside their relationship with you. - - Conveying under any other circumstances is permitted solely under -the conditions stated below. Sublicensing is not allowed; section 10 -makes it unnecessary. - - 3. Protecting Users' Legal Rights From Anti-Circumvention Law. - - No covered work shall be deemed part of an effective technological -measure under any applicable law fulfilling obligations under article -11 of the WIPO copyright treaty adopted on 20 December 1996, or -similar laws prohibiting or restricting circumvention of such -measures. - - When you convey a covered work, you waive any legal power to forbid -circumvention of technological measures to the extent such circumvention -is effected by exercising rights under this License with respect to -the covered work, and you disclaim any intention to limit operation or -modification of the work as a means of enforcing, against the work's -users, your or third parties' legal rights to forbid circumvention of -technological measures. - - 4. Conveying Verbatim Copies. - - You may convey verbatim copies of the Program's source code as you -receive it, in any medium, provided that you conspicuously and -appropriately publish on each copy an appropriate copyright notice; -keep intact all notices stating that this License and any -non-permissive terms added in accord with section 7 apply to the code; -keep intact all notices of the absence of any warranty; and give all -recipients a copy of this License along with the Program. - - You may charge any price or no price for each copy that you convey, -and you may offer support or warranty protection for a fee. - - 5. Conveying Modified Source Versions. - - You may convey a work based on the Program, or the modifications to -produce it from the Program, in the form of source code under the -terms of section 4, provided that you also meet all of these conditions: - - a) The work must carry prominent notices stating that you modified - it, and giving a relevant date. - - b) The work must carry prominent notices stating that it is - released under this License and any conditions added under section - 7. This requirement modifies the requirement in section 4 to - "keep intact all notices". - - c) You must license the entire work, as a whole, under this - License to anyone who comes into possession of a copy. This - License will therefore apply, along with any applicable section 7 - additional terms, to the whole of the work, and all its parts, - regardless of how they are packaged. This License gives no - permission to license the work in any other way, but it does not - invalidate such permission if you have separately received it. - - d) If the work has interactive user interfaces, each must display - Appropriate Legal Notices; however, if the Program has interactive - interfaces that do not display Appropriate Legal Notices, your - work need not make them do so. - - A compilation of a covered work with other separate and independent -works, which are not by their nature extensions of the covered work, -and which are not combined with it such as to form a larger program, -in or on a volume of a storage or distribution medium, is called an -"aggregate" if the compilation and its resulting copyright are not -used to limit the access or legal rights of the compilation's users -beyond what the individual works permit. Inclusion of a covered work -in an aggregate does not cause this License to apply to the other -parts of the aggregate. - - 6. Conveying Non-Source Forms. - - You may convey a covered work in object code form under the terms -of sections 4 and 5, provided that you also convey the -machine-readable Corresponding Source under the terms of this License, -in one of these ways: - - a) Convey the object code in, or embodied in, a physical product - (including a physical distribution medium), accompanied by the - Corresponding Source fixed on a durable physical medium - customarily used for software interchange. - - b) Convey the object code in, or embodied in, a physical product - (including a physical distribution medium), accompanied by a - written offer, valid for at least three years and valid for as - long as you offer spare parts or customer support for that product - model, to give anyone who possesses the object code either (1) a - copy of the Corresponding Source for all the software in the - product that is covered by this License, on a durable physical - medium customarily used for software interchange, for a price no - more than your reasonable cost of physically performing this - conveying of source, or (2) access to copy the - Corresponding Source from a network server at no charge. - - c) Convey individual copies of the object code with a copy of the - written offer to provide the Corresponding Source. This - alternative is allowed only occasionally and noncommercially, and - only if you received the object code with such an offer, in accord - with subsection 6b. - - d) Convey the object code by offering access from a designated - place (gratis or for a charge), and offer equivalent access to the - Corresponding Source in the same way through the same place at no - further charge. You need not require recipients to copy the - Corresponding Source along with the object code. If the place to - copy the object code is a network server, the Corresponding Source - may be on a different server (operated by you or a third party) - that supports equivalent copying facilities, provided you maintain - clear directions next to the object code saying where to find the - Corresponding Source. Regardless of what server hosts the - Corresponding Source, you remain obligated to ensure that it is - available for as long as needed to satisfy these requirements. - - e) Convey the object code using peer-to-peer transmission, provided - you inform other peers where the object code and Corresponding - Source of the work are being offered to the general public at no - charge under subsection 6d. - - A separable portion of the object code, whose source code is excluded -from the Corresponding Source as a System Library, need not be -included in conveying the object code work. - - A "User Product" is either (1) a "consumer product", which means any -tangible personal property which is normally used for personal, family, -or household purposes, or (2) anything designed or sold for incorporation -into a dwelling. In determining whether a product is a consumer product, -doubtful cases shall be resolved in favor of coverage. For a particular -product received by a particular user, "normally used" refers to a -typical or common use of that class of product, regardless of the status -of the particular user or of the way in which the particular user -actually uses, or expects or is expected to use, the product. A product -is a consumer product regardless of whether the product has substantial -commercial, industrial or non-consumer uses, unless such uses represent -the only significant mode of use of the product. - - "Installation Information" for a User Product means any methods, -procedures, authorization keys, or other information required to install -and execute modified versions of a covered work in that User Product from -a modified version of its Corresponding Source. The information must -suffice to ensure that the continued functioning of the modified object -code is in no case prevented or interfered with solely because -modification has been made. - - If you convey an object code work under this section in, or with, or -specifically for use in, a User Product, and the conveying occurs as -part of a transaction in which the right of possession and use of the -User Product is transferred to the recipient in perpetuity or for a -fixed term (regardless of how the transaction is characterized), the -Corresponding Source conveyed under this section must be accompanied -by the Installation Information. But this requirement does not apply -if neither you nor any third party retains the ability to install -modified object code on the User Product (for example, the work has -been installed in ROM). - - The requirement to provide Installation Information does not include a -requirement to continue to provide support service, warranty, or updates -for a work that has been modified or installed by the recipient, or for -the User Product in which it has been modified or installed. Access to a -network may be denied when the modification itself materially and -adversely affects the operation of the network or violates the rules and -protocols for communication across the network. - - Corresponding Source conveyed, and Installation Information provided, -in accord with this section must be in a format that is publicly -documented (and with an implementation available to the public in -source code form), and must require no special password or key for -unpacking, reading or copying. - - 7. Additional Terms. - - "Additional permissions" are terms that supplement the terms of this -License by making exceptions from one or more of its conditions. -Additional permissions that are applicable to the entire Program shall -be treated as though they were included in this License, to the extent -that they are valid under applicable law. If additional permissions -apply only to part of the Program, that part may be used separately -under those permissions, but the entire Program remains governed by -this License without regard to the additional permissions. - - When you convey a copy of a covered work, you may at your option -remove any additional permissions from that copy, or from any part of -it. (Additional permissions may be written to require their own -removal in certain cases when you modify the work.) You may place -additional permissions on material, added by you to a covered work, -for which you have or can give appropriate copyright permission. - - Notwithstanding any other provision of this License, for material you -add to a covered work, you may (if authorized by the copyright holders of -that material) supplement the terms of this License with terms: - - a) Disclaiming warranty or limiting liability differently from the - terms of sections 15 and 16 of this License; or - - b) Requiring preservation of specified reasonable legal notices or - author attributions in that material or in the Appropriate Legal - Notices displayed by works containing it; or - - c) Prohibiting misrepresentation of the origin of that material, or - requiring that modified versions of such material be marked in - reasonable ways as different from the original version; or - - d) Limiting the use for publicity purposes of names of licensors or - authors of the material; or - - e) Declining to grant rights under trademark law for use of some - trade names, trademarks, or service marks; or - - f) Requiring indemnification of licensors and authors of that - material by anyone who conveys the material (or modified versions of - it) with contractual assumptions of liability to the recipient, for - any liability that these contractual assumptions directly impose on - those licensors and authors. - - All other non-permissive additional terms are considered "further -restrictions" within the meaning of section 10. If the Program as you -received it, or any part of it, contains a notice stating that it is -governed by this License along with a term that is a further -restriction, you may remove that term. If a license document contains -a further restriction but permits relicensing or conveying under this -License, you may add to a covered work material governed by the terms -of that license document, provided that the further restriction does -not survive such relicensing or conveying. - - If you add terms to a covered work in accord with this section, you -must place, in the relevant source files, a statement of the -additional terms that apply to those files, or a notice indicating -where to find the applicable terms. - - Additional terms, permissive or non-permissive, may be stated in the -form of a separately written license, or stated as exceptions; -the above requirements apply either way. - - 8. Termination. - - You may not propagate or modify a covered work except as expressly -provided under this License. Any attempt otherwise to propagate or -modify it is void, and will automatically terminate your rights under -this License (including any patent licenses granted under the third -paragraph of section 11). - - However, if you cease all violation of this License, then your -license from a particular copyright holder is reinstated (a) -provisionally, unless and until the copyright holder explicitly and -finally terminates your license, and (b) permanently, if the copyright -holder fails to notify you of the violation by some reasonable means -prior to 60 days after the cessation. - - Moreover, your license from a particular copyright holder is -reinstated permanently if the copyright holder notifies you of the -violation by some reasonable means, this is the first time you have -received notice of violation of this License (for any work) from that -copyright holder, and you cure the violation prior to 30 days after -your receipt of the notice. - - Termination of your rights under this section does not terminate the -licenses of parties who have received copies or rights from you under -this License. If your rights have been terminated and not permanently -reinstated, you do not qualify to receive new licenses for the same -material under section 10. - - 9. Acceptance Not Required for Having Copies. - - You are not required to accept this License in order to receive or -run a copy of the Program. Ancillary propagation of a covered work -occurring solely as a consequence of using peer-to-peer transmission -to receive a copy likewise does not require acceptance. However, -nothing other than this License grants you permission to propagate or -modify any covered work. These actions infringe copyright if you do -not accept this License. Therefore, by modifying or propagating a -covered work, you indicate your acceptance of this License to do so. - - 10. Automatic Licensing of Downstream Recipients. - - Each time you convey a covered work, the recipient automatically -receives a license from the original licensors, to run, modify and -propagate that work, subject to this License. You are not responsible -for enforcing compliance by third parties with this License. - - An "entity transaction" is a transaction transferring control of an -organization, or substantially all assets of one, or subdividing an -organization, or merging organizations. If propagation of a covered -work results from an entity transaction, each party to that -transaction who receives a copy of the work also receives whatever -licenses to the work the party's predecessor in interest had or could -give under the previous paragraph, plus a right to possession of the -Corresponding Source of the work from the predecessor in interest, if -the predecessor has it or can get it with reasonable efforts. - - You may not impose any further restrictions on the exercise of the -rights granted or affirmed under this License. For example, you may -not impose a license fee, royalty, or other charge for exercise of -rights granted under this License, and you may not initiate litigation -(including a cross-claim or counterclaim in a lawsuit) alleging that -any patent claim is infringed by making, using, selling, offering for -sale, or importing the Program or any portion of it. - - 11. Patents. - - A "contributor" is a copyright holder who authorizes use under this -License of the Program or a work on which the Program is based. The -work thus licensed is called the contributor's "contributor version". - - A contributor's "essential patent claims" are all patent claims -owned or controlled by the contributor, whether already acquired or -hereafter acquired, that would be infringed by some manner, permitted -by this License, of making, using, or selling its contributor version, -but do not include claims that would be infringed only as a -consequence of further modification of the contributor version. For -purposes of this definition, "control" includes the right to grant -patent sublicenses in a manner consistent with the requirements of -this License. - - Each contributor grants you a non-exclusive, worldwide, royalty-free -patent license under the contributor's essential patent claims, to -make, use, sell, offer for sale, import and otherwise run, modify and -propagate the contents of its contributor version. - - In the following three paragraphs, a "patent license" is any express -agreement or commitment, however denominated, not to enforce a patent -(such as an express permission to practice a patent or covenant not to -sue for patent infringement). To "grant" such a patent license to a -party means to make such an agreement or commitment not to enforce a -patent against the party. - - If you convey a covered work, knowingly relying on a patent license, -and the Corresponding Source of the work is not available for anyone -to copy, free of charge and under the terms of this License, through a -publicly available network server or other readily accessible means, -then you must either (1) cause the Corresponding Source to be so -available, or (2) arrange to deprive yourself of the benefit of the -patent license for this particular work, or (3) arrange, in a manner -consistent with the requirements of this License, to extend the patent -license to downstream recipients. "Knowingly relying" means you have -actual knowledge that, but for the patent license, your conveying the -covered work in a country, or your recipient's use of the covered work -in a country, would infringe one or more identifiable patents in that -country that you have reason to believe are valid. - - If, pursuant to or in connection with a single transaction or -arrangement, you convey, or propagate by procuring conveyance of, a -covered work, and grant a patent license to some of the parties -receiving the covered work authorizing them to use, propagate, modify -or convey a specific copy of the covered work, then the patent license -you grant is automatically extended to all recipients of the covered -work and works based on it. - - A patent license is "discriminatory" if it does not include within -the scope of its coverage, prohibits the exercise of, or is -conditioned on the non-exercise of one or more of the rights that are -specifically granted under this License. You may not convey a covered -work if you are a party to an arrangement with a third party that is -in the business of distributing software, under which you make payment -to the third party based on the extent of your activity of conveying -the work, and under which the third party grants, to any of the -parties who would receive the covered work from you, a discriminatory -patent license (a) in connection with copies of the covered work -conveyed by you (or copies made from those copies), or (b) primarily -for and in connection with specific products or compilations that -contain the covered work, unless you entered into that arrangement, -or that patent license was granted, prior to 28 March 2007. - - Nothing in this License shall be construed as excluding or limiting -any implied license or other defenses to infringement that may -otherwise be available to you under applicable patent law. - - 12. No Surrender of Others' Freedom. - - If conditions are imposed on you (whether by court order, agreement or -otherwise) that contradict the conditions of this License, they do not -excuse you from the conditions of this License. If you cannot convey a -covered work so as to satisfy simultaneously your obligations under this -License and any other pertinent obligations, then as a consequence you may -not convey it at all. For example, if you agree to terms that obligate you -to collect a royalty for further conveying from those to whom you convey -the Program, the only way you could satisfy both those terms and this -License would be to refrain entirely from conveying the Program. - - 13. Use with the GNU Affero General Public License. - - Notwithstanding any other provision of this License, you have -permission to link or combine any covered work with a work licensed -under version 3 of the GNU Affero General Public License into a single -combined work, and to convey the resulting work. The terms of this -License will continue to apply to the part which is the covered work, -but the special requirements of the GNU Affero General Public License, -section 13, concerning interaction through a network will apply to the -combination as such. - - 14. Revised Versions of this License. - - The Free Software Foundation may publish revised and/or new versions of -the GNU General Public License from time to time. Such new versions will -be similar in spirit to the present version, but may differ in detail to -address new problems or concerns. - - Each version is given a distinguishing version number. If the -Program specifies that a certain numbered version of the GNU General -Public License "or any later version" applies to it, you have the -option of following the terms and conditions either of that numbered -version or of any later version published by the Free Software -Foundation. If the Program does not specify a version number of the -GNU General Public License, you may choose any version ever published -by the Free Software Foundation. - - If the Program specifies that a proxy can decide which future -versions of the GNU General Public License can be used, that proxy's -public statement of acceptance of a version permanently authorizes you -to choose that version for the Program. - - Later license versions may give you additional or different -permissions. However, no additional obligations are imposed on any -author or copyright holder as a result of your choosing to follow a -later version. - - 15. Disclaimer of Warranty. - - THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY -APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT -HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY -OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, -THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR -PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM -IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF -ALL NECESSARY SERVICING, REPAIR OR CORRECTION. - - 16. Limitation of Liability. - - IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING -WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS -THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY -GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE -USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF -DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD -PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), -EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF -SUCH DAMAGES. - - 17. Interpretation of Sections 15 and 16. - - If the disclaimer of warranty and limitation of liability provided -above cannot be given local legal effect according to their terms, -reviewing courts shall apply local law that most closely approximates -an absolute waiver of all civil liability in connection with the -Program, unless a warranty or assumption of liability accompanies a -copy of the Program in return for a fee. - - END OF TERMS AND CONDITIONS - - How to Apply These Terms to Your New Programs - - If you develop a new program, and you want it to be of the greatest -possible use to the public, the best way to achieve this is to make it -free software which everyone can redistribute and change under these terms. - - To do so, attach the following notices to the program. It is safest -to attach them to the start of each source file to most effectively -state the exclusion of warranty; and each file should have at least -the "copyright" line and a pointer to where the full notice is found. - - - Copyright (C) - - This program is free software: you can redistribute it and/or modify - it under the terms of the GNU General Public License as published by - the Free Software Foundation, either version 3 of the License, or - (at your option) any later version. - - This program is distributed in the hope that it will be useful, - but WITHOUT ANY WARRANTY; without even the implied warranty of - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - GNU General Public License for more details. - - You should have received a copy of the GNU General Public License - along with this program. If not, see . - -Also add information on how to contact you by electronic and paper mail. - - If the program does terminal interaction, make it output a short -notice like this when it starts in an interactive mode: - - Copyright (C) - This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'. - This is free software, and you are welcome to redistribute it - under certain conditions; type `show c' for details. - -The hypothetical commands `show w' and `show c' should show the appropriate -parts of the General Public License. Of course, your program's commands -might be different; for a GUI interface, you would use an "about box". - - You should also get your employer (if you work as a programmer) or school, -if any, to sign a "copyright disclaimer" for the program, if necessary. -For more information on this, and how to apply and follow the GNU GPL, see -. - - The GNU General Public License does not permit incorporating your program -into proprietary programs. If your program is a subroutine library, you -may consider it more useful to permit linking proprietary applications with -the library. If this is what you want to do, use the GNU Lesser General -Public License instead of this License. But first, please read -. diff --git a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/COPYING.LESSER b/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/COPYING.LESSER deleted file mode 100644 index 65c5ca88a67c30becee01c5a8816d964b03862f9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/COPYING.LESSER +++ /dev/null @@ -1,165 +0,0 @@ - GNU LESSER GENERAL PUBLIC LICENSE - Version 3, 29 June 2007 - - Copyright (C) 2007 Free Software Foundation, Inc. - Everyone is permitted to copy and distribute verbatim copies - of this license document, but changing it is not allowed. - - - This version of the GNU Lesser General Public License incorporates -the terms and conditions of version 3 of the GNU General Public -License, supplemented by the additional permissions listed below. - - 0. Additional Definitions. - - As used herein, "this License" refers to version 3 of the GNU Lesser -General Public License, and the "GNU GPL" refers to version 3 of the GNU -General Public License. - - "The Library" refers to a covered work governed by this License, -other than an Application or a Combined Work as defined below. - - An "Application" is any work that makes use of an interface provided -by the Library, but which is not otherwise based on the Library. -Defining a subclass of a class defined by the Library is deemed a mode -of using an interface provided by the Library. - - A "Combined Work" is a work produced by combining or linking an -Application with the Library. The particular version of the Library -with which the Combined Work was made is also called the "Linked -Version". - - The "Minimal Corresponding Source" for a Combined Work means the -Corresponding Source for the Combined Work, excluding any source code -for portions of the Combined Work that, considered in isolation, are -based on the Application, and not on the Linked Version. - - The "Corresponding Application Code" for a Combined Work means the -object code and/or source code for the Application, including any data -and utility programs needed for reproducing the Combined Work from the -Application, but excluding the System Libraries of the Combined Work. - - 1. Exception to Section 3 of the GNU GPL. - - You may convey a covered work under sections 3 and 4 of this License -without being bound by section 3 of the GNU GPL. - - 2. Conveying Modified Versions. - - If you modify a copy of the Library, and, in your modifications, a -facility refers to a function or data to be supplied by an Application -that uses the facility (other than as an argument passed when the -facility is invoked), then you may convey a copy of the modified -version: - - a) under this License, provided that you make a good faith effort to - ensure that, in the event an Application does not supply the - function or data, the facility still operates, and performs - whatever part of its purpose remains meaningful, or - - b) under the GNU GPL, with none of the additional permissions of - this License applicable to that copy. - - 3. Object Code Incorporating Material from Library Header Files. - - The object code form of an Application may incorporate material from -a header file that is part of the Library. You may convey such object -code under terms of your choice, provided that, if the incorporated -material is not limited to numerical parameters, data structure -layouts and accessors, or small macros, inline functions and templates -(ten or fewer lines in length), you do both of the following: - - a) Give prominent notice with each copy of the object code that the - Library is used in it and that the Library and its use are - covered by this License. - - b) Accompany the object code with a copy of the GNU GPL and this license - document. - - 4. Combined Works. - - You may convey a Combined Work under terms of your choice that, -taken together, effectively do not restrict modification of the -portions of the Library contained in the Combined Work and reverse -engineering for debugging such modifications, if you also do each of -the following: - - a) Give prominent notice with each copy of the Combined Work that - the Library is used in it and that the Library and its use are - covered by this License. - - b) Accompany the Combined Work with a copy of the GNU GPL and this license - document. - - c) For a Combined Work that displays copyright notices during - execution, include the copyright notice for the Library among - these notices, as well as a reference directing the user to the - copies of the GNU GPL and this license document. - - d) Do one of the following: - - 0) Convey the Minimal Corresponding Source under the terms of this - License, and the Corresponding Application Code in a form - suitable for, and under terms that permit, the user to - recombine or relink the Application with a modified version of - the Linked Version to produce a modified Combined Work, in the - manner specified by section 6 of the GNU GPL for conveying - Corresponding Source. - - 1) Use a suitable shared library mechanism for linking with the - Library. A suitable mechanism is one that (a) uses at run time - a copy of the Library already present on the user's computer - system, and (b) will operate properly with a modified version - of the Library that is interface-compatible with the Linked - Version. - - e) Provide Installation Information, but only if you would otherwise - be required to provide such information under section 6 of the - GNU GPL, and only to the extent that such information is - necessary to install and execute a modified version of the - Combined Work produced by recombining or relinking the - Application with a modified version of the Linked Version. (If - you use option 4d0, the Installation Information must accompany - the Minimal Corresponding Source and Corresponding Application - Code. If you use option 4d1, you must provide the Installation - Information in the manner specified by section 6 of the GNU GPL - for conveying Corresponding Source.) - - 5. Combined Libraries. - - You may place library facilities that are a work based on the -Library side by side in a single library together with other library -facilities that are not Applications and are not covered by this -License, and convey such a combined library under terms of your -choice, if you do both of the following: - - a) Accompany the combined library with a copy of the same work based - on the Library, uncombined with any other library facilities, - conveyed under the terms of this License. - - b) Give prominent notice with the combined library that part of it - is a work based on the Library, and explaining where to find the - accompanying uncombined form of the same work. - - 6. Revised Versions of the GNU Lesser General Public License. - - The Free Software Foundation may publish revised and/or new versions -of the GNU Lesser General Public License from time to time. Such new -versions will be similar in spirit to the present version, but may -differ in detail to address new problems or concerns. - - Each version is given a distinguishing version number. If the -Library as you received it specifies that a certain numbered version -of the GNU Lesser General Public License "or any later version" -applies to it, you have the option of following the terms and -conditions either of that published version or of any later version -published by the Free Software Foundation. If the Library as you -received it does not specify a version number of the GNU Lesser -General Public License, you may choose any version of the GNU Lesser -General Public License ever published by the Free Software Foundation. - - If the Library as you received it specifies that a proxy can decide -whether future versions of the GNU Lesser General Public License shall -apply, that proxy's public statement of acceptance of any version is -permanent authorization for you to choose that version for the -Library. diff --git a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/METADATA b/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/METADATA deleted file mode 100644 index 5bca79f83a61a9bb9e2232a01746f7181bcc2311..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/METADATA +++ /dev/null @@ -1,103 +0,0 @@ -Metadata-Version: 2.2 -Name: PyGithub -Version: 2.6.1 -Summary: Use the full Github API v3 -Author-email: Vincent Jacques -Project-URL: Documentation, https://pygithub.readthedocs.io/en/stable/ -Project-URL: Source, https://github.com/pygithub/pygithub -Project-URL: Tracker, https://github.com/pygithub/pygithub/issues -Keywords: github -Classifier: Development Status :: 5 - Production/Stable -Classifier: Environment :: Web Environment -Classifier: Intended Audience :: Developers -Classifier: License :: OSI Approved :: GNU Library or Lesser General Public License (LGPL) -Classifier: Operating System :: OS Independent -Classifier: Programming Language :: Python -Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3.8 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Programming Language :: Python :: 3.10 -Classifier: Programming Language :: Python :: 3.11 -Classifier: Programming Language :: Python :: 3.12 -Classifier: Topic :: Software Development -Requires-Python: >=3.8 -Description-Content-Type: text/markdown -License-File: COPYING -License-File: COPYING.LESSER -Requires-Dist: pynacl>=1.4.0 -Requires-Dist: requests>=2.14.0 -Requires-Dist: pyjwt[crypto]>=2.4.0 -Requires-Dist: typing-extensions>=4.0.0 -Requires-Dist: urllib3>=1.26.0 -Requires-Dist: Deprecated -Provides-Extra: integrations - -# PyGitHub - -[![PyPI](https://img.shields.io/pypi/v/PyGithub.svg)](https://pypi.python.org/pypi/PyGithub) -![CI](https://github.com/PyGithub/PyGithub/workflows/CI/badge.svg) -[![readthedocs](https://img.shields.io/badge/docs-stable-brightgreen.svg?style=flat)](https://pygithub.readthedocs.io/en/stable/?badge=stable) -[![License](https://img.shields.io/badge/license-LGPL-blue.svg)](https://en.wikipedia.org/wiki/GNU_Lesser_General_Public_License) -[![Slack](https://img.shields.io/badge/Slack%20channel-%20%20-blue.svg)](https://join.slack.com/t/pygithub-project/shared_invite/zt-duj89xtx-uKFZtgAg209o6Vweqm8xeQ) -[![Open Source Helpers](https://www.codetriage.com/pygithub/pygithub/badges/users.svg)](https://www.codetriage.com/pygithub/pygithub) -[![codecov](https://codecov.io/gh/PyGithub/PyGithub/branch/master/graph/badge.svg)](https://codecov.io/gh/PyGithub/PyGithub) -[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) - -PyGitHub is a Python library to access the [GitHub REST API]. -This library enables you to manage [GitHub] resources such as repositories, user profiles, and organizations in your Python applications. - -[GitHub REST API]: https://docs.github.com/en/rest -[GitHub]: https://github.com - -## Install - -```bash -pip install PyGithub -``` - -## Simple Demo - -```python -from github import Github - -# Authentication is defined via github.Auth -from github import Auth - -# using an access token -auth = Auth.Token("access_token") - -# First create a Github instance: - -# Public Web Github -g = Github(auth=auth) - -# Github Enterprise with custom hostname -g = Github(base_url="https://{hostname}/api/v3", auth=auth) - -# Then play with your Github objects: -for repo in g.get_user().get_repos(): - print(repo.name) - -# To close connections after use -g.close() -``` - -## Documentation - -More information can be found on the [PyGitHub documentation site.](https://pygithub.readthedocs.io/en/stable/introduction.html) - -## Development - -### Contributing - -Long-term discussion and bug reports are maintained via GitHub Issues. -Code review is done via GitHub Pull Requests. - -For more information read [CONTRIBUTING.md]. - -[CONTRIBUTING.md]: https://github.com/PyGithub/PyGithub/blob/main/CONTRIBUTING.md - -### Maintainership - -We're actively seeking maintainers that will triage issues and pull requests and cut releases. -If you work on a project that leverages PyGitHub and have a vested interest in keeping the code alive and well, send an email to someone in the MAINTAINERS file. diff --git a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/RECORD b/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/RECORD deleted file mode 100644 index 611d67cdb717a2954c7f996c1daa3fab5652f849..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/RECORD +++ /dev/null @@ -1,337 +0,0 @@ -PyGithub-2.6.1.dist-info/COPYING,sha256=jOtLnuWt7d5Hsx6XXB2QxzrSe2sWWh3NgMfFRetluQM,35147 -PyGithub-2.6.1.dist-info/COPYING.LESSER,sha256=2n6rt7r999OuXp8iOqW9we7ORaxWncIbOwN1ILRGR2g,7651 -PyGithub-2.6.1.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -PyGithub-2.6.1.dist-info/METADATA,sha256=6GM2H6ECoFpjMHvmsQMTQvrkHNNqLQxqLrSaduB7pow,3889 -PyGithub-2.6.1.dist-info/RECORD,, -PyGithub-2.6.1.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -PyGithub-2.6.1.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91 -PyGithub-2.6.1.dist-info/top_level.txt,sha256=X_Dn8Q-zPGudQV37mRBM5uHhwhEsv4x1AeAuB1dfEls,7 -github/AccessToken.py,sha256=_Ryw-pme8fxvSY-2dqDlnH7x1akRrxry_EKWGex4mEU,6756 -github/AdvisoryBase.py,sha256=61onXFv5hVi0GpmUjKCPAJo8bRKJAwxA6cxix8V3M3I,7196 -github/AdvisoryCredit.py,sha256=Sth4hayVCtuAS8vkCBITY7iAQ9Z_2_QEh8vIYXtxuKs,5213 -github/AdvisoryCreditDetailed.py,sha256=0Jns6-Wgi2da_nAc_-cdKrdrcx6q1LTn73VXz-1kI24,5163 -github/AdvisoryVulnerability.py,sha256=kOYX8pNMO2YA2bA3j02w69GVUSBfqvOL4VXFcppT9ss,8607 -github/AdvisoryVulnerabilityPackage.py,sha256=Bh_BiNS34PBv5aZYFR9SUNUNdXRvhMPh1Nll11x3Tpg,4972 -github/AppAuthentication.py,sha256=8ItYtxH1Z3qUCFiI5pIrUgWkfxrNfDtv_PTLlhAUhCs,3362 -github/ApplicationOAuth.py,sha256=bXK9Vuw4d_VlgqfR3I4eDLflTWiJWqYDT2ee6W3TLPA,8200 -github/Artifact.py,sha256=5wWfNLP61VqVBBnmy0OPgX_XXezpyLw3UO3D7V53M8o,8364 -github/Auth.py,sha256=j4n343bx20ZYTFu7KjdBXIe8JvtGh-VxNaRwWGRBWDc,18777 -github/AuthenticatedUser.py,sha256=RPSbudt7Fwf93A6y0MagxCjqGS6cZ74IianKlvfVSfM,51604 -github/Authorization.py,sha256=X1HngQla9ail9EMVvd7hIYe3D7iD8bkP7VOEyEn5Oi8,8589 -github/AuthorizationApplication.py,sha256=AO_iJkla_Z-SlGBdjDgU4vdIhQvC-xyi3g9ZavwLPXY,3912 -github/Autolink.py,sha256=oca0RVXVIcjZJ-m9_2i01-8aaz7Ric8QUcWq-kc8AY8,4945 -github/Branch.py,sha256=xMHCSGF6atubtwMTWev9ImwRASMXQCLQVHxIcpjnRQU,32091 -github/BranchProtection.py,sha256=emtYLjbzekTi5VxvVO7fk1bwCdyY0fkxHBjEB3qxvHs,10599 -github/CVSS.py,sha256=O7nf9sgk1XVZ8n0CnHgxXZ_A5NH8oarvWfUIUTamL7I,4658 -github/CWE.py,sha256=liISfggat6zcylarfj0SB67SIBNuTHJjVcPmoOUuzxg,4072 -github/CheckRun.py,sha256=8oVyQ0iR2VyiLLRoHNR0ABXwwI6m8MI1LQWP1WJNNq4,12905 -github/CheckRunAnnotation.py,sha256=3ZmsqDEoxIX0ggVk87cczYeyQVv9WMjwZRP_3cAdR6o,5647 -github/CheckRunOutput.py,sha256=28uFE7IsNBxdm2nrpsRrXmkH-zKKR6J3JWgZa5b6wzU,5058 -github/CheckSuite.py,sha256=ZVbfdkDmmJxhwnD5LYBofRcOdHXjPHIp9JVX_IVHeQg,12650 -github/Clones.py,sha256=IDesxzVgVHIf4OyJeunH6--5V-fUn0z4a4YeSgFDcpA,4862 -github/CodeScanAlert.py,sha256=kzdKASKNHz4zqjLWuQ4j8bsk8ukDpWnRalRGq-QYpSU,7166 -github/CodeScanAlertInstance.py,sha256=A5OdwJBXmvYnTQgpFEMN5iadIMbwobWbL-JTd-Kanvo,5463 -github/CodeScanAlertInstanceLocation.py,sha256=5o4nPcL5VZgwys6EnvhGshXTOqplZgxDw7zB6wzNDTs,4539 -github/CodeScanRule.py,sha256=fqunSyh14WpluH95E5aKF7zCVLr8kv3El5rHNmXrPqA,5182 -github/CodeScanTool.py,sha256=z4gipqMPyBArAqX4bUvnINCc3P5cABzes1fY3D1_IWY,4545 -github/CodeSecurityConfig.py,sha256=Z36edlAPbKh3B5V4MmuJUyQaU0zvsaOvDnf2NENuE30,9907 -github/CodeSecurityConfigRepository.py,sha256=TvM3cdov7dpiJnBzsSpiWOatBMcSDxVP6cXl6ebiTJ0,3291 -github/Commit.py,sha256=jJO2eLxCfkGj3fKavpKEnda8Yn1hLqu_A7HEhKeualU,16463 -github/CommitCombinedStatus.py,sha256=rZIoDhiJtdSAhKlKIXcg2AaOedb5F8tymHPdgVqsRRE,5892 -github/CommitComment.py,sha256=KTavqBGP56fFVWjI_ycuR6aC1rvag5kMjXd3BAY2E1o,11391 -github/CommitStats.py,sha256=1YwhWBFBGzUE6wE8BVTYhMMIUDAVQaz6SDnc8Z8Ck6A,4012 -github/CommitStatus.py,sha256=u3ZpM_zwyQSPRlagL_nuNDWQ0X9WtU0KZhq2bRTFwL4,6973 -github/Comparison.py,sha256=II8s1H54jBhyEmgVBNHzM3qklbj1LnQHbhh9Ep7HscE,8010 -github/Consts.py,sha256=GJUQi2O4WHJTjNyB42fYj5x2eQFDrm8aQhfB7ahicRQ,10492 -github/ContentFile.py,sha256=8g-Qasve1GdVok9Ym_3sVUaQ2nPjT9FgI4xH75BPliE,12595 -github/Copilot.py,sha256=GoE5Goq3wI9Aeel7KsjodX_Dm4YYRGJla5GZM52J0gI,4532 -github/CopilotSeat.py,sha256=xvhOo-A6sAKkGMhiXWaZAcW4p9qIqE7FeU3BZPVf6BQ,4794 -github/DefaultCodeSecurityConfig.py,sha256=PGrYyWk7WdLK_NqP7mTW6VnPT0H4b8wXKPjzx0rBJik,4741 -github/DependabotAlert.py,sha256=0Hq4RgnanHzOsXOMSTG544g9NTbwG5SNFJBR1B255bo,7711 -github/DependabotAlertAdvisory.py,sha256=7TI5mRyg6PJmZUkLSdHL2Zypq6hKmc8FqeWShA61Cag,3591 -github/DependabotAlertDependency.py,sha256=2Cg3e6G2W6LajWZ-xBArSsRtz4zW5wab-8oKINU0h_Y,3665 -github/DependabotAlertVulnerability.py,sha256=eAt2Fy-REMD8lSzTSxyAtn6Nf9P0NEyTffRXc2sgKr4,3988 -github/Deployment.py,sha256=D1tYEGGmn62nqexHQzda5oVHDhBPihB0Y6o8dN57oBM,13654 -github/DeploymentStatus.py,sha256=qq-qMMBHAtw3hIIkMNiXbaD_7JdPpLvywHwF0lSb1Hs,9370 -github/DiscussionBase.py,sha256=54013rlbYF4_MItFXuH0AYzmuhMW6Z8kpvDVhc5Z1gQ,6784 -github/DiscussionCommentBase.py,sha256=Hnuox_kzfZ9xruwYWzRtqbDnl5Ehndr5TynkcNC2WlA,6835 -github/Download.py,sha256=MKVNGIqD4kZo5fhdF2HUnyweLnPcgs7nWuQSMAbo1Kk,11218 -github/Enterprise.py,sha256=x-qYVrFUFlE6ZR1LeZgg8eFkvrucLwGshJnn3EAJgQo,5327 -github/EnterpriseConsumedLicenses.py,sha256=eDO5MrobiG5bRXGGIYOnIl8hgftq3IO_IooOu4LjmMs,5762 -github/Environment.py,sha256=iynT0fU-E7m8Tt9jqa94QdOezGjchayxw2nfI3Lh1k8,13623 -github/EnvironmentDeploymentBranchPolicy.py,sha256=vZeU3Y4nQJSHJ3Lz_rd5PF_J7rJdvdWkIaXR330q2yY,4090 -github/EnvironmentProtectionRule.py,sha256=91IbsSTP2lW4n8xWd8FAjpJEe26frqqg-6yVputx-Gw,5574 -github/EnvironmentProtectionRuleReviewer.py,sha256=_HLqHz4nSJeaYnk38bbcb-283Hg7jQ1nmTnUDxehQd8,5217 -github/Event.py,sha256=hkNr9ps_wiJCCKsk1-3aBR0uhJBxronm7ebSGyoHTMs,6182 -github/File.py,sha256=aO9qr6VoVWVrcn7ej-mR0mq-nTbJLpFYK1jy94xKTCk,6684 -github/Gist.py,sha256=PPa2kPhY142xQ_hqXtBtKJjR_tu3UXkGKHCKaMIcri8,15230 -github/GistComment.py,sha256=Bh4j2lvIl7lmRtPYuneOAU9C2ENtGUTdL5SjDu6sYZQ,7187 -github/GistFile.py,sha256=RZWinuJBhSQ2V4K2jU9PLVEeMvp6nLWM3b4LxsOAItY,5168 -github/GistHistoryState.py,sha256=B2CY5KWBySuIqJfIWdUfSyHklRPiyuSi2lgdmVuvXIE,10660 -github/GitAuthor.py,sha256=KkvmNGbJLdBb2KaBI9b51YCHwqHVbhlNT-AkK5fW634,4925 -github/GitBlob.py,sha256=qIsFepFXDhtHeV7FePc2K9ArcZTEPqeTrH6DYydoVgA,5133 -github/GitCommit.py,sha256=P50SxhMzC41208FYNgt24q_5OCEjCer5pnv_eowNTIk,9603 -github/GitCommitVerification.py,sha256=mNB8LYL3HAK7H9j7tlnL7bsQNHFzIy-7ocFfVehkvF8,5178 -github/GitObject.py,sha256=gSKCYAv4aMnLYrIyfPbmwxuD1tFzdnmP6aLqul9YpJ4,4263 -github/GitRef.py,sha256=QdXAK6l-hEou-fz73agqmUhWGb2ubXHT8mB_UWeoFbk,5958 -github/GitRelease.py,sha256=mdi6VKI8fMPl3zqVG7fi3ixC8BXlm0BSCUJ3SNqZCuM,19855 -github/GitReleaseAsset.py,sha256=ej3ZvO5deyNABVOO6USbghjywYe5YHXbzRu9wYE2Qy8,9336 -github/GitTag.py,sha256=wOTNQy1WQcsm44SiA2Q4yyB-Y9HbjPO2K7pAR51IxR8,6682 -github/GitTree.py,sha256=iGdDGSup7MfQ3HuS4F5e7W_JWsU8dyz_XqdNqt39fR8,5460 -github/GitTreeElement.py,sha256=XiHf3uk9SBzdK7eqzVvuD8JzWzn2rg3LuS9hZzE-wAo,4975 -github/GithubApp.py,sha256=UaoMUb-6lwxz5Ho-rrixLl7qQLuwXUePICqGmr_ydgY,10064 -github/GithubException.py,sha256=OBCphJtOygdOID_-6pekI13tFoO1xPF_-B6TbvIP-jI,7175 -github/GithubIntegration.py,sha256=iO4OX27rJjr8BsVmwbue1IpJpMZlQJVUEeLu2OsT3S0,15186 -github/GithubObject.py,sha256=2R3MxaWqsBwSWmjN0X0ZYVEGXEyL5_qNsJabflNhj-k,22219 -github/GithubRetry.py,sha256=Rzz4pN-0x5tEZ9Z-q9VrhS04y6rficm9SnSmru6dLH0,12063 -github/GitignoreTemplate.py,sha256=el6be_upqLHbVc9N36K71LT1iR82bGBNftg7j1zpepM,4178 -github/GlobalAdvisory.py,sha256=K98M-E0zPg5MU60GFMJwFtH-fPsgwClQGaNLjNthdic,6500 -github/Hook.py,sha256=GWufDuLdPEtPYE__BumrZN-cITS41mD2c0dibd4SQ0U,10162 -github/HookDelivery.py,sha256=xassYw8R-YdLEzzsmP3l7_nmrd9ED57lI8ybe-Nd58w,9555 -github/HookDescription.py,sha256=uiAtb4k8CDPcNG63aal3zyQTTEer7b3cPqh9hm517Lk,4587 -github/HookResponse.py,sha256=y8lPbm9h9APbhrHQ0nKPBardAqTozQYin5tD0eHuIAo,4250 -github/InputFileContent.py,sha256=U8bCrl2z9KHQhWtxc-fwliIwJDZlcvUbzY8WbhKfIx0,3480 -github/InputGitAuthor.py,sha256=iW7vxND5eUq6Ull_aH7J8iMKfTZ2ejGLtoSoSGGDX2E,3820 -github/InputGitTreeElement.py,sha256=gAq3OTo0NOOmyjqqtNbwk3wagrmaOp1n-XK990tcYss,3933 -github/Installation.py,sha256=tuVqoD2xRT-2tjYfFlFf4MYr2k-NJMKGRmo6_ugfexk,12641 -github/InstallationAuthorization.py,sha256=4_nDj-7Ltcr1Wdc9yqs-sOuMH6BPA5Pcc_cI067bNZI,6901 -github/Invitation.py,sha256=3KGXH57B5VbU1dUk83AqE9GZ3rBeR79QtXN8WD7yY0o,7305 -github/Issue.py,sha256=7T-S9QZ1ifQEIRxRJ0HPRL9nVe5UpuIi3pJhBfOTTvs,30560 -github/IssueComment.py,sha256=bjwtq1QDFxq3q2jOeNt6fF3Ia85xvr0Xc240pmQ6Wps,13141 -github/IssueEvent.py,sha256=GXMUXdk0vT4EFT-l3HJoLt__7iUMB7WiPFb3jb4gJWM,12108 -github/IssuePullRequest.py,sha256=yYYgYohKY46v9Eq9eWNQAV7R8fsHrx8d6vNeUxOgx-Y,4869 -github/Label.py,sha256=Ohli5Ya-vaNtyKnUUoanRs8hcBvcD58pOcyh5nJkgwU,7617 -github/License.py,sha256=ffxavMjTTK1inKTmt96Avrv6d9sB0OUnV4tCaO5udmc,7992 -github/MainClass.py,sha256=mwPg14ViVjmlz9AAFLjlDttgAzC0AYvERVY76PsCpRc,49018 -github/Membership.py,sha256=PoZy0XAAtNPGvU2G1h90zBQzpYUJzvjDEhWJkVunCtA,6575 -github/MergedUpstream.py,sha256=Mxqfh8GV2aVs682hwnWsVala6O6OsuvkvL4PuD3vYLE,3372 -github/Migration.py,sha256=FD4X2ZyjH9zFIGE-wABt_pd7qY09Gs1BdPNSfoTAChc,12145 -github/Milestone.py,sha256=ScAUsM3UwHKmY5uUdI0XB-92fqAT-4xPrFm_uORsjiw,10985 -github/NamedEnterpriseUser.py,sha256=IGLTM3XL7gouxmbTxZDlfMzROiQu_bEJ7bNG-2DJTbs,10543 -github/NamedUser.py,sha256=8fbxtu98KHZ5TlcgaxQilmfidh-VuvGKSvVKBD-25cg,30543 -github/Notification.py,sha256=G3ZlYtbwEVX2K7SGMfaKNbCMXlEaEslIcN_Tohec2e4,8434 -github/NotificationSubject.py,sha256=9Vl6b_j98V4UhGIkHRXYj_v3bAUIH0qSbh-AJs-bJPU,4960 -github/Organization.py,sha256=jUoT6Pt0X6qiwuLKdv4_1O2d1ZlkhmRBAWtrgbh2iDY,95864 -github/OrganizationCustomProperty.py,sha256=7AgcQgTpOkYgv2mAmwvqmHEHFPaEHWLTDNqCWpMeJPc,7941 -github/OrganizationDependabotAlert.py,sha256=rI7b3NS4Dhs4Cz3SZBmZaF95qMGf5k0zsfUc9Rgqz2I,2983 -github/OrganizationSecret.py,sha256=c34t5QiHgYfPEYoj4I4j5HnvK0ALuR_Cj7Y-TGXPFDM,6336 -github/OrganizationVariable.py,sha256=qQcoHpGjdhb0IPD76GjL8Biqvgdeq_2nD3tO2IOu600,6065 -github/PaginatedList.py,sha256=GdBlOD9FxR2wVb8CQwtmG1t9UIy7jcL4i1Qr6gFMUH0,19548 -github/Path.py,sha256=ZmUgLnYu-WPHt54qhA2GjFZxFb1oR--lHKoiGrmgfRQ,4952 -github/Permissions.py,sha256=EqchnebEsf8zhxa1_ONIS2BlWPO-KmjUK6Q9J5W73Zk,5403 -github/Plan.py,sha256=FtYnYjHjsEGVM4gs6aTCSVPIHp6eKY2dtYW88yDC5BQ,5297 -github/Project.py,sha256=9_DwCeIUTU7QMYEuMnQk0g--qUy4puErSA9nQXguiWs,11991 -github/ProjectCard.py,sha256=ICxGVyuzNDaQ6oWe1O65g4eVJqSrZjh6xDGuh0isKsU,10881 -github/ProjectColumn.py,sha256=2dvDWGAT1fbE0Z75WOz9EevmHJVe3AfzOfrDSXgGgmU,10141 -github/PublicKey.py,sha256=gD5-H7eqCvYGW1MDJLIra7wz0HZI5br9GgUVqMtKuI0,6587 -github/PullRequest.py,sha256=yOWWfKFnvQcAJ0zGmecqFSNZU5UPdiPsvFnjIHx_6Kc,51191 -github/PullRequestComment.py,sha256=-y9Y8BspYRe8hgt5pSRFOCuv-HQ4P45XSAlpAIuT25s,17133 -github/PullRequestMergeStatus.py,sha256=o2x4A_YxXAF-1dMiOAZLP8rbfl-6nbsgkVgU5iYc8c0,4503 -github/PullRequestPart.py,sha256=8G3qy1aFvjdWsck7GErNS1heApGSOIHD0KQ_wVjf3vU,5385 -github/PullRequestReview.py,sha256=Zh1spNf3mST9VTfPVvDd052I-ksoGLX1EcBy7cqWOXE,7766 -github/Rate.py,sha256=zppWBE9wwGgg28TtGQZom6caG-kYz7Edw8O6bK4o-jk,5006 -github/RateLimit.py,sha256=aXmT5-B6W4AUQD1K79g012MCmmwITRaUG4VbHW6gIa8,7871 -github/Reaction.py,sha256=YauwrJTcoPhUfP85sgMg0FuLRvJceloqktJ6wfXesu8,5576 -github/Referrer.py,sha256=_ZLDv74O26rySNT8OWCYJRxlPw27V7bw_3bRZlVNMhw,4701 -github/RepoCodeSecurityConfig.py,sha256=TaMX2mYEW-6nsDbxYs-OAhnV2aveHGtiPa09mH-W4tk,4682 -github/Repository.py,sha256=wOlqESLJqZpohEclXLxpHS_-jDWohaVyzPplG58w_i8,204970 -github/RepositoryAdvisory.py,sha256=iu6LtELT85aC7PPxqDsOu1FFPZW_XUIls2Pf_233QFM,18310 -github/RepositoryDiscussion.py,sha256=imjaUm7riDnuNENvUVIB0GotoN2TXsXO3qthMcOOtU8,9163 -github/RepositoryDiscussionCategory.py,sha256=Wu8tplwDQdfzc8kfuKxMKDHPg3LOVLOg9pMFydyHNqQ,6803 -github/RepositoryDiscussionComment.py,sha256=xKCbqPHdN4Qa8MO02Q4hkjvid3qtpeUJgq99VmN_uoY,6581 -github/RepositoryKey.py,sha256=-g6f9PPaMwpTNCeGnfc787gSE7s3WH1jR8Vd-BjaRAA,7508 -github/RepositoryPreferences.py,sha256=WhXwIqydh4PewCjKt7WXGzmVZU3CZodkilLojSSESrM,4555 -github/Requester.py,sha256=0BuG_dB_aol27T-QEIieZrgqHxT8idx9GlTaxiSoRds,52603 -github/RequiredPullRequestReviews.py,sha256=13mLV-Fa2ZQj30v0Zl_5vOoaeZLoAIaMCz2PfN_XmdU,7704 -github/RequiredStatusChecks.py,sha256=QQuY8o1H6cv7HafvpBhu8Mm45hsYoroQtNdhwVHoFO8,4966 -github/Secret.py,sha256=OzFE9tp62nd02_rf1KMKxYpTzxV5DFvJKMRh4IJ5jls,5741 -github/SecurityAndAnalysis.py,sha256=31x-3HrAqWc0dfiZfnwcw-HmQIeXWvMOXFjWTFiXuJ0,8178 -github/SecurityAndAnalysisFeature.py,sha256=CJJ6rR4AzGigv4RfVbPCZKHq6USQAGS6b9BW8uE29XI,4258 -github/SelfHostedActionsRunner.py,sha256=60ucM62i80VDG9migGM0pFWeJ4seDczNGu_Tp4OanjU,5524 -github/SourceImport.py,sha256=Bcpc_MlsYB0VjwwuviLelh_3bSfZGxssnIQKm9PMqC0,11372 -github/Stargazer.py,sha256=7nNVcVeIqjK9SWr1QMG0M_fb6pHijPzm3AMLmGpL5ww,4825 -github/StatsCodeFrequency.py,sha256=R7XWZsjSExFjQvNWc9TZ2bJhBGLxywoU5LfVjKao7V0,4290 -github/StatsCommitActivity.py,sha256=YmuVoKMvfTcYpdqZgfYddZ1fIML_I8786-zXxS0vdU4,4472 -github/StatsContributor.py,sha256=7e46cpqLCfwCojxjzVVA4kBO-76IA50gZP8hbX5C2lw,5805 -github/StatsParticipation.py,sha256=TsW0s5CK2S042aZARHZtmhRntUIEY-BioaCsxa3mQxA,4244 -github/StatsPunchCard.py,sha256=JBAWWEUowsCiMAY98kT7iZRGQ2KUsq4n8b27zI2qQHA,4012 -github/Tag.py,sha256=1vPAqpjcJ9k-Gfl2L6zaiYEGeTE5JlmBvMrN5qw4O1g,5402 -github/Team.py,sha256=LKEFCKmx4oE7MnUV-6tkA0g5AX7lhPhuBb8KDV5I_Jg,23893 -github/TeamDiscussion.py,sha256=BLVmJYU3Er9VQO1005gzPvWqpy82GLULlDFFzriegRQ,6599 -github/TimelineEvent.py,sha256=QYM38PNDtiDeXZuQURHc86OGw6PflmXyUN6W0S4_e_I,7533 -github/TimelineEventSource.py,sha256=tf5x3asgFzz3z3SbOmINxBAA4BMuoSVS-X_ZkkHU4hA,4324 -github/Topic.py,sha256=codJNHBB5imQvI4bVJUXl7QH8NEgsJ_JA2YgtvPupg4,8373 -github/Traffic.py,sha256=RTkvDq-RyFRkveTpi_kwZ3qITFkz249pXh3g1SlENbg,4768 -github/UserKey.py,sha256=gRGKuveCFtZI_ZhId6hHKodS8El3snf9Mk7wMmCLEwM,6272 -github/Variable.py,sha256=3sOXF7nJlia_YEzTBF9nrZd5BPsWo5akXCiAeZE0xvc,5699 -github/View.py,sha256=9cP39jexYuQOhYDHheWCy1qOEVUlDG_0ACF1LJXf8fA,4845 -github/Workflow.py,sha256=YA3a37RAeIjBGmFn1kXicg_ZyV02WnOBGLe56KObIdI,12088 -github/WorkflowJob.py,sha256=FckU3PIOfHK-8RMlkwv9v_xNypGM82aG4-Gjy0TblJI,10652 -github/WorkflowRun.py,sha256=zm4lX7BCoPzpMuAiOB8CSgNzgQ4vFOyXKlSFlTLFL8c,19556 -github/WorkflowStep.py,sha256=vloX2qGP9tzjpeAIcU1B_EoY9TLn-Z15QZQnx8I4LeQ,5643 -github/__init__.py,sha256=62guPyYWH5dj0djnitPdPX-GpGa1xH4bc3WCYncXE60,5401 -github/__pycache__/AccessToken.cpython-310.pyc,, -github/__pycache__/AdvisoryBase.cpython-310.pyc,, -github/__pycache__/AdvisoryCredit.cpython-310.pyc,, -github/__pycache__/AdvisoryCreditDetailed.cpython-310.pyc,, -github/__pycache__/AdvisoryVulnerability.cpython-310.pyc,, -github/__pycache__/AdvisoryVulnerabilityPackage.cpython-310.pyc,, -github/__pycache__/AppAuthentication.cpython-310.pyc,, -github/__pycache__/ApplicationOAuth.cpython-310.pyc,, -github/__pycache__/Artifact.cpython-310.pyc,, -github/__pycache__/Auth.cpython-310.pyc,, -github/__pycache__/AuthenticatedUser.cpython-310.pyc,, -github/__pycache__/Authorization.cpython-310.pyc,, -github/__pycache__/AuthorizationApplication.cpython-310.pyc,, -github/__pycache__/Autolink.cpython-310.pyc,, -github/__pycache__/Branch.cpython-310.pyc,, -github/__pycache__/BranchProtection.cpython-310.pyc,, -github/__pycache__/CVSS.cpython-310.pyc,, -github/__pycache__/CWE.cpython-310.pyc,, -github/__pycache__/CheckRun.cpython-310.pyc,, -github/__pycache__/CheckRunAnnotation.cpython-310.pyc,, -github/__pycache__/CheckRunOutput.cpython-310.pyc,, -github/__pycache__/CheckSuite.cpython-310.pyc,, -github/__pycache__/Clones.cpython-310.pyc,, -github/__pycache__/CodeScanAlert.cpython-310.pyc,, -github/__pycache__/CodeScanAlertInstance.cpython-310.pyc,, -github/__pycache__/CodeScanAlertInstanceLocation.cpython-310.pyc,, -github/__pycache__/CodeScanRule.cpython-310.pyc,, -github/__pycache__/CodeScanTool.cpython-310.pyc,, -github/__pycache__/CodeSecurityConfig.cpython-310.pyc,, -github/__pycache__/CodeSecurityConfigRepository.cpython-310.pyc,, -github/__pycache__/Commit.cpython-310.pyc,, -github/__pycache__/CommitCombinedStatus.cpython-310.pyc,, -github/__pycache__/CommitComment.cpython-310.pyc,, -github/__pycache__/CommitStats.cpython-310.pyc,, -github/__pycache__/CommitStatus.cpython-310.pyc,, -github/__pycache__/Comparison.cpython-310.pyc,, -github/__pycache__/Consts.cpython-310.pyc,, -github/__pycache__/ContentFile.cpython-310.pyc,, -github/__pycache__/Copilot.cpython-310.pyc,, -github/__pycache__/CopilotSeat.cpython-310.pyc,, -github/__pycache__/DefaultCodeSecurityConfig.cpython-310.pyc,, -github/__pycache__/DependabotAlert.cpython-310.pyc,, -github/__pycache__/DependabotAlertAdvisory.cpython-310.pyc,, -github/__pycache__/DependabotAlertDependency.cpython-310.pyc,, -github/__pycache__/DependabotAlertVulnerability.cpython-310.pyc,, -github/__pycache__/Deployment.cpython-310.pyc,, -github/__pycache__/DeploymentStatus.cpython-310.pyc,, -github/__pycache__/DiscussionBase.cpython-310.pyc,, -github/__pycache__/DiscussionCommentBase.cpython-310.pyc,, -github/__pycache__/Download.cpython-310.pyc,, -github/__pycache__/Enterprise.cpython-310.pyc,, -github/__pycache__/EnterpriseConsumedLicenses.cpython-310.pyc,, -github/__pycache__/Environment.cpython-310.pyc,, -github/__pycache__/EnvironmentDeploymentBranchPolicy.cpython-310.pyc,, -github/__pycache__/EnvironmentProtectionRule.cpython-310.pyc,, -github/__pycache__/EnvironmentProtectionRuleReviewer.cpython-310.pyc,, -github/__pycache__/Event.cpython-310.pyc,, -github/__pycache__/File.cpython-310.pyc,, -github/__pycache__/Gist.cpython-310.pyc,, -github/__pycache__/GistComment.cpython-310.pyc,, -github/__pycache__/GistFile.cpython-310.pyc,, -github/__pycache__/GistHistoryState.cpython-310.pyc,, -github/__pycache__/GitAuthor.cpython-310.pyc,, -github/__pycache__/GitBlob.cpython-310.pyc,, -github/__pycache__/GitCommit.cpython-310.pyc,, -github/__pycache__/GitCommitVerification.cpython-310.pyc,, -github/__pycache__/GitObject.cpython-310.pyc,, -github/__pycache__/GitRef.cpython-310.pyc,, -github/__pycache__/GitRelease.cpython-310.pyc,, -github/__pycache__/GitReleaseAsset.cpython-310.pyc,, -github/__pycache__/GitTag.cpython-310.pyc,, -github/__pycache__/GitTree.cpython-310.pyc,, -github/__pycache__/GitTreeElement.cpython-310.pyc,, -github/__pycache__/GithubApp.cpython-310.pyc,, -github/__pycache__/GithubException.cpython-310.pyc,, -github/__pycache__/GithubIntegration.cpython-310.pyc,, -github/__pycache__/GithubObject.cpython-310.pyc,, -github/__pycache__/GithubRetry.cpython-310.pyc,, -github/__pycache__/GitignoreTemplate.cpython-310.pyc,, -github/__pycache__/GlobalAdvisory.cpython-310.pyc,, -github/__pycache__/Hook.cpython-310.pyc,, -github/__pycache__/HookDelivery.cpython-310.pyc,, -github/__pycache__/HookDescription.cpython-310.pyc,, -github/__pycache__/HookResponse.cpython-310.pyc,, -github/__pycache__/InputFileContent.cpython-310.pyc,, -github/__pycache__/InputGitAuthor.cpython-310.pyc,, -github/__pycache__/InputGitTreeElement.cpython-310.pyc,, -github/__pycache__/Installation.cpython-310.pyc,, -github/__pycache__/InstallationAuthorization.cpython-310.pyc,, -github/__pycache__/Invitation.cpython-310.pyc,, -github/__pycache__/Issue.cpython-310.pyc,, -github/__pycache__/IssueComment.cpython-310.pyc,, -github/__pycache__/IssueEvent.cpython-310.pyc,, -github/__pycache__/IssuePullRequest.cpython-310.pyc,, -github/__pycache__/Label.cpython-310.pyc,, -github/__pycache__/License.cpython-310.pyc,, -github/__pycache__/MainClass.cpython-310.pyc,, -github/__pycache__/Membership.cpython-310.pyc,, -github/__pycache__/MergedUpstream.cpython-310.pyc,, -github/__pycache__/Migration.cpython-310.pyc,, -github/__pycache__/Milestone.cpython-310.pyc,, -github/__pycache__/NamedEnterpriseUser.cpython-310.pyc,, -github/__pycache__/NamedUser.cpython-310.pyc,, -github/__pycache__/Notification.cpython-310.pyc,, -github/__pycache__/NotificationSubject.cpython-310.pyc,, -github/__pycache__/Organization.cpython-310.pyc,, -github/__pycache__/OrganizationCustomProperty.cpython-310.pyc,, -github/__pycache__/OrganizationDependabotAlert.cpython-310.pyc,, -github/__pycache__/OrganizationSecret.cpython-310.pyc,, -github/__pycache__/OrganizationVariable.cpython-310.pyc,, -github/__pycache__/PaginatedList.cpython-310.pyc,, -github/__pycache__/Path.cpython-310.pyc,, -github/__pycache__/Permissions.cpython-310.pyc,, -github/__pycache__/Plan.cpython-310.pyc,, -github/__pycache__/Project.cpython-310.pyc,, -github/__pycache__/ProjectCard.cpython-310.pyc,, -github/__pycache__/ProjectColumn.cpython-310.pyc,, -github/__pycache__/PublicKey.cpython-310.pyc,, -github/__pycache__/PullRequest.cpython-310.pyc,, -github/__pycache__/PullRequestComment.cpython-310.pyc,, -github/__pycache__/PullRequestMergeStatus.cpython-310.pyc,, -github/__pycache__/PullRequestPart.cpython-310.pyc,, -github/__pycache__/PullRequestReview.cpython-310.pyc,, -github/__pycache__/Rate.cpython-310.pyc,, -github/__pycache__/RateLimit.cpython-310.pyc,, -github/__pycache__/Reaction.cpython-310.pyc,, -github/__pycache__/Referrer.cpython-310.pyc,, -github/__pycache__/RepoCodeSecurityConfig.cpython-310.pyc,, -github/__pycache__/Repository.cpython-310.pyc,, -github/__pycache__/RepositoryAdvisory.cpython-310.pyc,, -github/__pycache__/RepositoryDiscussion.cpython-310.pyc,, -github/__pycache__/RepositoryDiscussionCategory.cpython-310.pyc,, -github/__pycache__/RepositoryDiscussionComment.cpython-310.pyc,, -github/__pycache__/RepositoryKey.cpython-310.pyc,, -github/__pycache__/RepositoryPreferences.cpython-310.pyc,, -github/__pycache__/Requester.cpython-310.pyc,, -github/__pycache__/RequiredPullRequestReviews.cpython-310.pyc,, -github/__pycache__/RequiredStatusChecks.cpython-310.pyc,, -github/__pycache__/Secret.cpython-310.pyc,, -github/__pycache__/SecurityAndAnalysis.cpython-310.pyc,, -github/__pycache__/SecurityAndAnalysisFeature.cpython-310.pyc,, -github/__pycache__/SelfHostedActionsRunner.cpython-310.pyc,, -github/__pycache__/SourceImport.cpython-310.pyc,, -github/__pycache__/Stargazer.cpython-310.pyc,, -github/__pycache__/StatsCodeFrequency.cpython-310.pyc,, -github/__pycache__/StatsCommitActivity.cpython-310.pyc,, -github/__pycache__/StatsContributor.cpython-310.pyc,, -github/__pycache__/StatsParticipation.cpython-310.pyc,, -github/__pycache__/StatsPunchCard.cpython-310.pyc,, -github/__pycache__/Tag.cpython-310.pyc,, -github/__pycache__/Team.cpython-310.pyc,, -github/__pycache__/TeamDiscussion.cpython-310.pyc,, -github/__pycache__/TimelineEvent.cpython-310.pyc,, -github/__pycache__/TimelineEventSource.cpython-310.pyc,, -github/__pycache__/Topic.cpython-310.pyc,, -github/__pycache__/Traffic.cpython-310.pyc,, -github/__pycache__/UserKey.cpython-310.pyc,, -github/__pycache__/Variable.cpython-310.pyc,, -github/__pycache__/View.cpython-310.pyc,, -github/__pycache__/Workflow.cpython-310.pyc,, -github/__pycache__/WorkflowJob.cpython-310.pyc,, -github/__pycache__/WorkflowRun.cpython-310.pyc,, -github/__pycache__/WorkflowStep.cpython-310.pyc,, -github/__pycache__/__init__.cpython-310.pyc,, -github/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 diff --git a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/REQUESTED deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/WHEEL b/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/WHEEL deleted file mode 100644 index 505164bc02d63fe6b0b3299f849a77c5f1beeb41..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/WHEEL +++ /dev/null @@ -1,5 +0,0 @@ -Wheel-Version: 1.0 -Generator: setuptools (75.8.0) -Root-Is-Purelib: true -Tag: py3-none-any - diff --git a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/top_level.txt deleted file mode 100644 index 345dcda8f5a6110b8025bdc8e03c3e903fd35e83..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyGithub-2.6.1.dist-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -github diff --git a/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/AUTHORS.rst b/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/AUTHORS.rst deleted file mode 100644 index 88e2b6ad752425c0c372874e5b3b7695408e7ee3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/AUTHORS.rst +++ /dev/null @@ -1,7 +0,0 @@ -Authors -======= - -``pyjwt`` is currently written and maintained by `Jose Padilla `_. -Originally written and maintained by `Jeff Lindsay `_. - -A full list of contributors can be found on GitHub’s `overview `_. diff --git a/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/LICENSE b/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/LICENSE deleted file mode 100644 index fd0ecbc889c067a562b3a745c9d2d7934f87d64b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -The MIT License (MIT) - -Copyright (c) 2015-2022 José Padilla - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/METADATA b/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/METADATA deleted file mode 100644 index f31b700e64c9dd3c2cf4ed16a262787fe61e6f8d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/METADATA +++ /dev/null @@ -1,106 +0,0 @@ -Metadata-Version: 2.1 -Name: PyJWT -Version: 2.10.1 -Summary: JSON Web Token implementation in Python -Author-email: Jose Padilla -License: MIT -Project-URL: Homepage, https://github.com/jpadilla/pyjwt -Keywords: json,jwt,security,signing,token,web -Classifier: Development Status :: 5 - Production/Stable -Classifier: Intended Audience :: Developers -Classifier: License :: OSI Approved :: MIT License -Classifier: Natural Language :: English -Classifier: Programming Language :: Python -Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3 :: Only -Classifier: Programming Language :: Python :: 3.9 -Classifier: Programming Language :: Python :: 3.10 -Classifier: Programming Language :: Python :: 3.11 -Classifier: Programming Language :: Python :: 3.12 -Classifier: Programming Language :: Python :: 3.13 -Classifier: Topic :: Utilities -Requires-Python: >=3.9 -Description-Content-Type: text/x-rst -License-File: LICENSE -License-File: AUTHORS.rst -Provides-Extra: crypto -Requires-Dist: cryptography>=3.4.0; extra == "crypto" -Provides-Extra: dev -Requires-Dist: coverage[toml]==5.0.4; extra == "dev" -Requires-Dist: cryptography>=3.4.0; extra == "dev" -Requires-Dist: pre-commit; extra == "dev" -Requires-Dist: pytest<7.0.0,>=6.0.0; extra == "dev" -Requires-Dist: sphinx; extra == "dev" -Requires-Dist: sphinx-rtd-theme; extra == "dev" -Requires-Dist: zope.interface; extra == "dev" -Provides-Extra: docs -Requires-Dist: sphinx; extra == "docs" -Requires-Dist: sphinx-rtd-theme; extra == "docs" -Requires-Dist: zope.interface; extra == "docs" -Provides-Extra: tests -Requires-Dist: coverage[toml]==5.0.4; extra == "tests" -Requires-Dist: pytest<7.0.0,>=6.0.0; extra == "tests" - -PyJWT -===== - -.. image:: https://github.com/jpadilla/pyjwt/workflows/CI/badge.svg - :target: https://github.com/jpadilla/pyjwt/actions?query=workflow%3ACI - -.. image:: https://img.shields.io/pypi/v/pyjwt.svg - :target: https://pypi.python.org/pypi/pyjwt - -.. image:: https://codecov.io/gh/jpadilla/pyjwt/branch/master/graph/badge.svg - :target: https://codecov.io/gh/jpadilla/pyjwt - -.. image:: https://readthedocs.org/projects/pyjwt/badge/?version=stable - :target: https://pyjwt.readthedocs.io/en/stable/ - -A Python implementation of `RFC 7519 `_. Original implementation was written by `@progrium `_. - -Sponsor -------- - -.. |auth0-logo| image:: https://github.com/user-attachments/assets/ee98379e-ee76-4bcb-943a-e25c4ea6d174 - :width: 160px - -+--------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| |auth0-logo| | If you want to quickly add secure token-based authentication to Python projects, feel free to check Auth0's Python SDK and free plan at `auth0.com/signup `_. | -+--------------+-----------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ - -Installing ----------- - -Install with **pip**: - -.. code-block:: console - - $ pip install PyJWT - - -Usage ------ - -.. code-block:: pycon - - >>> import jwt - >>> encoded = jwt.encode({"some": "payload"}, "secret", algorithm="HS256") - >>> print(encoded) - eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzb21lIjoicGF5bG9hZCJ9.4twFt5NiznN84AWoo1d7KO1T_yoc0Z6XOpOVswacPZg - >>> jwt.decode(encoded, "secret", algorithms=["HS256"]) - {'some': 'payload'} - -Documentation -------------- - -View the full docs online at https://pyjwt.readthedocs.io/en/stable/ - - -Tests ------ - -You can run tests from the project root after cloning with: - -.. code-block:: console - - $ tox diff --git a/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/RECORD b/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/RECORD deleted file mode 100644 index bcc25fd4c9f246b69129bb84143ccbb6743cec19..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/RECORD +++ /dev/null @@ -1,32 +0,0 @@ -PyJWT-2.10.1.dist-info/AUTHORS.rst,sha256=klzkNGECnu2_VY7At89_xLBF3vUSDruXk3xwgUBxzwc,322 -PyJWT-2.10.1.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -PyJWT-2.10.1.dist-info/LICENSE,sha256=eXp6ICMdTEM-nxkR2xcx0GtYKLmPSZgZoDT3wPVvXOU,1085 -PyJWT-2.10.1.dist-info/METADATA,sha256=EkewF6D6KU8SGaaQzVYfxUUU1P_gs_dp1pYTkoYvAx8,3990 -PyJWT-2.10.1.dist-info/RECORD,, -PyJWT-2.10.1.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91 -PyJWT-2.10.1.dist-info/top_level.txt,sha256=RP5DHNyJbMq2ka0FmfTgoSaQzh7e3r5XuCWCO8a00k8,4 -jwt/__init__.py,sha256=VB2vFKuboTjcDGeZ8r-UqK_dz3NsQSQEqySSICby8Xg,1711 -jwt/__pycache__/__init__.cpython-310.pyc,, -jwt/__pycache__/algorithms.cpython-310.pyc,, -jwt/__pycache__/api_jwk.cpython-310.pyc,, -jwt/__pycache__/api_jws.cpython-310.pyc,, -jwt/__pycache__/api_jwt.cpython-310.pyc,, -jwt/__pycache__/exceptions.cpython-310.pyc,, -jwt/__pycache__/help.cpython-310.pyc,, -jwt/__pycache__/jwk_set_cache.cpython-310.pyc,, -jwt/__pycache__/jwks_client.cpython-310.pyc,, -jwt/__pycache__/types.cpython-310.pyc,, -jwt/__pycache__/utils.cpython-310.pyc,, -jwt/__pycache__/warnings.cpython-310.pyc,, -jwt/algorithms.py,sha256=cKr-XEioe0mBtqJMCaHEswqVOA1Z8Purt5Sb3Bi-5BE,30409 -jwt/api_jwk.py,sha256=6F1r7rmm8V5qEnBKA_xMjS9R7VoANe1_BL1oD2FrAjE,4451 -jwt/api_jws.py,sha256=aM8vzqQf6mRrAw7bRy-Moj_pjWsKSVQyYK896AfMjJU,11762 -jwt/api_jwt.py,sha256=OGT4hok1l5A6FH_KdcrU5g6u6EQ8B7em0r9kGM9SYgA,14512 -jwt/exceptions.py,sha256=bUIOJ-v9tjopTLS-FYOTc3kFx5WP5IZt7ksN_HE1G9Q,1211 -jwt/help.py,sha256=vFdNzjQoAch04XCMYpCkyB2blaqHAGAqQrtf9nSPkdk,1808 -jwt/jwk_set_cache.py,sha256=hBKmN-giU7-G37L_XKgc_OZu2ah4wdbj1ZNG_GkoSE8,959 -jwt/jwks_client.py,sha256=p9b-IbQqo2tEge9Zit3oSPBFNePqwho96VLbnUrHUWs,4259 -jwt/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -jwt/types.py,sha256=VnhGv_VFu5a7_mrPoSCB7HaNLrJdhM8Sq1sSfEg0gLU,99 -jwt/utils.py,sha256=hxOjvDBheBYhz-RIPiEz7Q88dSUSTMzEdKE_Ww2VdJw,3640 -jwt/warnings.py,sha256=50XWOnyNsIaqzUJTk6XHNiIDykiL763GYA92MjTKmok,59 diff --git a/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/WHEEL b/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/WHEEL deleted file mode 100644 index ae527e7d64811439e61b93aa375defb30e06edfe..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/WHEEL +++ /dev/null @@ -1,5 +0,0 @@ -Wheel-Version: 1.0 -Generator: setuptools (75.6.0) -Root-Is-Purelib: true -Tag: py3-none-any - diff --git a/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/top_level.txt deleted file mode 100644 index 27ccc9bc3a9ca1cd3dab3b2530e67b3309fde50a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyJWT-2.10.1.dist-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -jwt diff --git a/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/LICENSE b/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/LICENSE deleted file mode 100644 index 91e18a62b67551a4d427e2eaee0d33dcab94e141..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/LICENSE +++ /dev/null @@ -1,174 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - -TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - -1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - -2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - -3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - -4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - -5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - -6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - -7. Disclaimer of Warranty. Unless required by applicable law or - 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 - 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 - risks associated with Your exercise of permissions under this License. - -8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - -9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. diff --git a/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/METADATA b/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/METADATA deleted file mode 100644 index fb216d4e8790acb9ccd057e2b3965accd9afa196..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/METADATA +++ /dev/null @@ -1,245 +0,0 @@ -Metadata-Version: 2.1 -Name: PyNaCl -Version: 1.5.0 -Summary: Python binding to the Networking and Cryptography (NaCl) library -Home-page: https://github.com/pyca/pynacl/ -Author: The PyNaCl developers -Author-email: cryptography-dev@python.org -License: Apache License 2.0 -Platform: UNKNOWN -Classifier: Programming Language :: Python :: Implementation :: CPython -Classifier: Programming Language :: Python :: Implementation :: PyPy -Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3.6 -Classifier: Programming Language :: Python :: 3.7 -Classifier: Programming Language :: Python :: 3.8 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Programming Language :: Python :: 3.10 -Requires-Python: >=3.6 -Requires-Dist: cffi (>=1.4.1) -Provides-Extra: docs -Requires-Dist: sphinx (>=1.6.5) ; extra == 'docs' -Requires-Dist: sphinx-rtd-theme ; extra == 'docs' -Provides-Extra: tests -Requires-Dist: pytest (!=3.3.0,>=3.2.1) ; extra == 'tests' -Requires-Dist: hypothesis (>=3.27.0) ; extra == 'tests' - -=============================================== -PyNaCl: Python binding to the libsodium library -=============================================== - -.. image:: https://img.shields.io/pypi/v/pynacl.svg - :target: https://pypi.org/project/PyNaCl/ - :alt: Latest Version - -.. image:: https://codecov.io/github/pyca/pynacl/coverage.svg?branch=main - :target: https://codecov.io/github/pyca/pynacl?branch=main - -.. image:: https://img.shields.io/pypi/pyversions/pynacl.svg - :target: https://pypi.org/project/PyNaCl/ - :alt: Compatible Python Versions - -PyNaCl is a Python binding to `libsodium`_, which is a fork of the -`Networking and Cryptography library`_. These libraries have a stated goal of -improving usability, security and speed. It supports Python 3.6+ as well as -PyPy 3. - -.. _libsodium: https://github.com/jedisct1/libsodium -.. _Networking and Cryptography library: https://nacl.cr.yp.to/ - -Features --------- - -* Digital signatures -* Secret-key encryption -* Public-key encryption -* Hashing and message authentication -* Password based key derivation and password hashing - -`Changelog`_ ------------- - -.. _Changelog: https://pynacl.readthedocs.io/en/stable/changelog/ - -Installation -============ - -Binary wheel install --------------------- - -PyNaCl ships as a binary wheel on macOS, Windows and Linux ``manylinux1`` [#many]_ , -so all dependencies are included. Make sure you have an up-to-date pip -and run: - -.. code-block:: console - - $ pip install pynacl - -Faster wheel build ------------------- - -You can define the environment variable ``LIBSODIUM_MAKE_ARGS`` to pass arguments to ``make`` -and enable `parallelization`_: - -.. code-block:: console - - $ LIBSODIUM_MAKE_ARGS=-j4 pip install pynacl - -Linux source build ------------------- - -PyNaCl relies on `libsodium`_, a portable C library. A copy is bundled -with PyNaCl so to install you can run: - -.. code-block:: console - - $ pip install pynacl - -If you'd prefer to use the version of ``libsodium`` provided by your -distribution, you can disable the bundled copy during install by running: - -.. code-block:: console - - $ SODIUM_INSTALL=system pip install pynacl - -.. warning:: Usage of the legacy ``easy_install`` command provided by setuptools - is generally discouraged, and is completely unsupported in PyNaCl's case. - -.. _parallelization: https://www.gnu.org/software/make/manual/html_node/Parallel.html - -.. _libsodium: https://github.com/jedisct1/libsodium - -.. [#many] `manylinux1 wheels `_ - are built on a baseline linux environment based on Centos 5.11 - and should work on most x86 and x86_64 glibc based linux environments. - -Changelog -========= - -1.5.0 (2022-01-07) ------------------- - -* **BACKWARDS INCOMPATIBLE:** Removed support for Python 2.7 and Python 3.5. -* **BACKWARDS INCOMPATIBLE:** We no longer distribute ``manylinux1`` - wheels. -* Added ``manylinux2014``, ``manylinux_2_24``, ``musllinux``, and macOS - ``universal2`` wheels (the latter supports macOS ``arm64``). -* Update ``libsodium`` to 1.0.18-stable (July 25, 2021 release). -* Add inline type hints. - -1.4.0 (2020-05-25) ------------------- - -* Update ``libsodium`` to 1.0.18. -* **BACKWARDS INCOMPATIBLE:** We no longer distribute 32-bit ``manylinux1`` - wheels. Continuing to produce them was a maintenance burden. -* Added support for Python 3.8, and removed support for Python 3.4. -* Add low level bindings for extracting the seed and the public key - from crypto_sign_ed25519 secret key -* Add low level bindings for deterministic random generation. -* Add ``wheel`` and ``setuptools`` setup_requirements in ``setup.py`` (#485) -* Fix checks on very slow builders (#481, #495) -* Add low-level bindings to ed25519 arithmetic functions -* Update low-level blake2b state implementation -* Fix wrong short-input behavior of SealedBox.decrypt() (#517) -* Raise CryptPrefixError exception instead of InvalidkeyError when trying - to check a password against a verifier stored in a unknown format (#519) -* Add support for minimal builds of libsodium. Trying to call functions - not available in a minimal build will raise an UnavailableError - exception. To compile a minimal build of the bundled libsodium, set - the SODIUM_INSTALL_MINIMAL environment variable to any non-empty - string (e.g. ``SODIUM_INSTALL_MINIMAL=1``) for setup. - -1.3.0 2018-09-26 ----------------- - -* Added support for Python 3.7. -* Update ``libsodium`` to 1.0.16. -* Run and test all code examples in PyNaCl docs through sphinx's - doctest builder. -* Add low-level bindings for chacha20-poly1305 AEAD constructions. -* Add low-level bindings for the chacha20-poly1305 secretstream constructions. -* Add low-level bindings for ed25519ph pre-hashed signing construction. -* Add low-level bindings for constant-time increment and addition - on fixed-precision big integers represented as little-endian - byte sequences. -* Add low-level bindings for the ISO/IEC 7816-4 compatible padding API. -* Add low-level bindings for libsodium's crypto_kx... key exchange - construction. -* Set hypothesis deadline to None in tests/test_pwhash.py to avoid - incorrect test failures on slower processor architectures. GitHub - issue #370 - -1.2.1 - 2017-12-04 ------------------- - -* Update hypothesis minimum allowed version. -* Infrastructure: add proper configuration for readthedocs builder - runtime environment. - -1.2.0 - 2017-11-01 ------------------- - -* Update ``libsodium`` to 1.0.15. -* Infrastructure: add jenkins support for automatic build of - ``manylinux1`` binary wheels -* Added support for ``SealedBox`` construction. -* Added support for ``argon2i`` and ``argon2id`` password hashing constructs - and restructured high-level password hashing implementation to expose - the same interface for all hashers. -* Added support for 128 bit ``siphashx24`` variant of ``siphash24``. -* Added support for ``from_seed`` APIs for X25519 keypair generation. -* Dropped support for Python 3.3. - -1.1.2 - 2017-03-31 ------------------- - -* reorder link time library search path when using bundled - libsodium - -1.1.1 - 2017-03-15 ------------------- - -* Fixed a circular import bug in ``nacl.utils``. - -1.1.0 - 2017-03-14 ------------------- - -* Dropped support for Python 2.6. -* Added ``shared_key()`` method on ``Box``. -* You can now pass ``None`` to ``nonce`` when encrypting with ``Box`` or - ``SecretBox`` and it will automatically generate a random nonce. -* Added support for ``siphash24``. -* Added support for ``blake2b``. -* Added support for ``scrypt``. -* Update ``libsodium`` to 1.0.11. -* Default to the bundled ``libsodium`` when compiling. -* All raised exceptions are defined mixing-in - ``nacl.exceptions.CryptoError`` - -1.0.1 - 2016-01-24 ------------------- - -* Fix an issue with absolute paths that prevented the creation of wheels. - -1.0 - 2016-01-23 ----------------- - -* PyNaCl has been ported to use the new APIs available in cffi 1.0+. - Due to this change we no longer support PyPy releases older than 2.6. -* Python 3.2 support has been dropped. -* Functions to convert between Ed25519 and Curve25519 keys have been added. - -0.3.0 - 2015-03-04 ------------------- - -* The low-level API (`nacl.c.*`) has been changed to match the - upstream NaCl C/C++ conventions (as well as those of other NaCl bindings). - The order of arguments and return values has changed significantly. To - avoid silent failures, `nacl.c` has been removed, and replaced with - `nacl.bindings` (with the new argument ordering). If you have code which - calls these functions (e.g. `nacl.c.crypto_box_keypair()`), you must review - the new docstrings and update your code/imports to match the new - conventions. - - diff --git a/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/RECORD b/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/RECORD deleted file mode 100644 index 6b18358443482290485bfc3be42d13f52aeb4723..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/RECORD +++ /dev/null @@ -1,68 +0,0 @@ -PyNaCl-1.5.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -PyNaCl-1.5.0.dist-info/LICENSE,sha256=0xdK1j5yHUydzLitQyCEiZLTFDabxGMZcgtYAskVP-k,9694 -PyNaCl-1.5.0.dist-info/METADATA,sha256=OJaXCiHgNRywLY9cj3X2euddUPZ4dnyyqAQMU01X4j0,8634 -PyNaCl-1.5.0.dist-info/RECORD,, -PyNaCl-1.5.0.dist-info/WHEEL,sha256=TIQeZFe3DwXBO5UGlCH1aKpf5Cx6FJLbIUqd-Sq2juI,185 -PyNaCl-1.5.0.dist-info/top_level.txt,sha256=wfdEOI_G2RIzmzsMyhpqP17HUh6Jcqi99to9aHLEslo,13 -nacl/__init__.py,sha256=0IUunzBT8_Jn0DUdHacBExOYeAEMggo8slkfjo7O0XM,1116 -nacl/__pycache__/__init__.cpython-310.pyc,, -nacl/__pycache__/encoding.cpython-310.pyc,, -nacl/__pycache__/exceptions.cpython-310.pyc,, -nacl/__pycache__/hash.cpython-310.pyc,, -nacl/__pycache__/hashlib.cpython-310.pyc,, -nacl/__pycache__/public.cpython-310.pyc,, -nacl/__pycache__/secret.cpython-310.pyc,, -nacl/__pycache__/signing.cpython-310.pyc,, -nacl/__pycache__/utils.cpython-310.pyc,, -nacl/_sodium.abi3.so,sha256=uJ6RwSnbb9wO4esR0bVUqrfFHtBOGm34IQIdmaE1fGY,2740136 -nacl/bindings/__init__.py,sha256=BDlStrds2EuUS4swOL4pnf92PWVS_CHRCptX3KhEX-s,16997 -nacl/bindings/__pycache__/__init__.cpython-310.pyc,, -nacl/bindings/__pycache__/crypto_aead.cpython-310.pyc,, -nacl/bindings/__pycache__/crypto_box.cpython-310.pyc,, -nacl/bindings/__pycache__/crypto_core.cpython-310.pyc,, -nacl/bindings/__pycache__/crypto_generichash.cpython-310.pyc,, -nacl/bindings/__pycache__/crypto_hash.cpython-310.pyc,, -nacl/bindings/__pycache__/crypto_kx.cpython-310.pyc,, -nacl/bindings/__pycache__/crypto_pwhash.cpython-310.pyc,, -nacl/bindings/__pycache__/crypto_scalarmult.cpython-310.pyc,, -nacl/bindings/__pycache__/crypto_secretbox.cpython-310.pyc,, -nacl/bindings/__pycache__/crypto_secretstream.cpython-310.pyc,, -nacl/bindings/__pycache__/crypto_shorthash.cpython-310.pyc,, -nacl/bindings/__pycache__/crypto_sign.cpython-310.pyc,, -nacl/bindings/__pycache__/randombytes.cpython-310.pyc,, -nacl/bindings/__pycache__/sodium_core.cpython-310.pyc,, -nacl/bindings/__pycache__/utils.cpython-310.pyc,, -nacl/bindings/crypto_aead.py,sha256=BIw1k_JCfr5ylZk0RF5rCFIM1fhfLkEa-aiWkrfffNE,15597 -nacl/bindings/crypto_box.py,sha256=Ox0NG2t4MsGhBAa7Kgah4o0gc99ULMsqkdX56ofOouY,10139 -nacl/bindings/crypto_core.py,sha256=6u9G3y7H-QrawO785UkFFFtwDoCkeHE63GOUl9p5-eA,13736 -nacl/bindings/crypto_generichash.py,sha256=9mX0DGIIzicr-uXrqFM1nU4tirasbixDwbcdfV7W1fc,8852 -nacl/bindings/crypto_hash.py,sha256=Rg1rsEwE3azhsQT-dNVPA4NB9VogJAKn1EfxYt0pPe0,2175 -nacl/bindings/crypto_kx.py,sha256=oZNVlNgROpHOa1XQ_uZe0tqIkdfuApeJlRnwR23_74k,6723 -nacl/bindings/crypto_pwhash.py,sha256=laVDo4xFUuGyEjtZAU510AklBF6ablBy7Z3HN1WDYjY,18848 -nacl/bindings/crypto_scalarmult.py,sha256=_DX-mst2uCnzjo6fP5HRTnhv1BC95B9gmJc3L_or16g,8244 -nacl/bindings/crypto_secretbox.py,sha256=KgZ1VvkCJDlQ85jtfe9c02VofPvuEgZEhWni-aX3MsM,2914 -nacl/bindings/crypto_secretstream.py,sha256=G0FgZS01qA5RzWzm5Bdms8Yy_lvgdZDoUYYBActPmvQ,11165 -nacl/bindings/crypto_shorthash.py,sha256=PQU7djHTLDGdVs-w_TsivjFHHp5EK5k2Yh6p-6z0T60,2603 -nacl/bindings/crypto_sign.py,sha256=53j2im9E4F79qT_2U8IfCAc3lzg0VMwEjvAPEUccVDg,10342 -nacl/bindings/randombytes.py,sha256=uBK3W4WcjgnjZdWanrX0fjYZpr9KHbBgNMl9rui-Ojc,1563 -nacl/bindings/sodium_core.py,sha256=9Y9CX--sq-TaPaQRPRpx8SWDSS9PJOja_Cqb-yqyJNQ,1039 -nacl/bindings/utils.py,sha256=KDwQnadXeNMbqEA1SmpNyCVo5k8MiUQa07QM66VzfXM,4298 -nacl/encoding.py,sha256=qTAPc2MXSkdh4cqDVY0ra6kHyViHMCmEo_re7cgGk5w,2915 -nacl/exceptions.py,sha256=GZH32aJtZgqCO4uz0LRsev8z0WyvAYuV3YVqT9AAQq4,2451 -nacl/hash.py,sha256=EYBOe6UVc9SUQINEmyuRSa1QGRSvdwdrBzTL1tdFLU8,6392 -nacl/hashlib.py,sha256=L5Fv75St8AMPvb-GhA4YqX5p1mC_Sb4HhC1NxNQMpJA,4400 -nacl/public.py,sha256=RVGCWQRjIJOmW-8sNrVLtsDjMMGx30i6UyfViGCnQNA,14792 -nacl/pwhash/__init__.py,sha256=XSDXd7wQHNLEHl0mkHfVb5lFQsp6ygHkhen718h0BSM,2675 -nacl/pwhash/__pycache__/__init__.cpython-310.pyc,, -nacl/pwhash/__pycache__/_argon2.cpython-310.pyc,, -nacl/pwhash/__pycache__/argon2i.cpython-310.pyc,, -nacl/pwhash/__pycache__/argon2id.cpython-310.pyc,, -nacl/pwhash/__pycache__/scrypt.cpython-310.pyc,, -nacl/pwhash/_argon2.py,sha256=jL1ChR9biwYh3RSuc-LJ2-W4DlVLHpir-XHGX8cpeJQ,1779 -nacl/pwhash/argon2i.py,sha256=IIvIuO9siKUu5-Wpz0SGiltLQv7Du_mi9BUE8INRK_4,4405 -nacl/pwhash/argon2id.py,sha256=H22i8O4j9Ws4L3JsXl9TRcJzDcyaVumhQRPzINAgJWM,4433 -nacl/pwhash/scrypt.py,sha256=fMr3Qht1a1EY8aebNNntfLRjinIPXtKYKKrrBhY5LDc,6986 -nacl/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -nacl/secret.py,sha256=kauBNuP-0rb3TjU2EMBMu5Vnmzjnscp1bRqMspy5LzU,12108 -nacl/signing.py,sha256=kbTEUyHLUMaNLv1nCjxzGxCs82Qs5w8gxE_CnEwPuIU,8337 -nacl/utils.py,sha256=gmlTD1x9ZNwzHd8LpALH1CHud-Htv8ejRb3y7TyS9f0,2341 diff --git a/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/WHEEL b/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/WHEEL deleted file mode 100644 index b93bb215c339b8e370dd4771640df50ecdf4b751..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/WHEEL +++ /dev/null @@ -1,7 +0,0 @@ -Wheel-Version: 1.0 -Generator: bdist_wheel (0.37.1) -Root-Is-Purelib: false -Tag: cp36-abi3-manylinux_2_17_x86_64 -Tag: cp36-abi3-manylinux2014_x86_64 -Tag: cp36-abi3-manylinux_2_24_x86_64 - diff --git a/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/top_level.txt deleted file mode 100644 index f52507f09b95c1d37c005528f28b345a8af187a4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/PyNaCl-1.5.0.dist-info/top_level.txt +++ /dev/null @@ -1,2 +0,0 @@ -_sodium -nacl diff --git a/venv/lib/python3.10/site-packages/_cffi_backend.cpython-310-x86_64-linux-gnu.so b/venv/lib/python3.10/site-packages/_cffi_backend.cpython-310-x86_64-linux-gnu.so deleted file mode 100644 index 9f40757cd1ce87ef45424e8fc926e5582cf87aba..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/_cffi_backend.cpython-310-x86_64-linux-gnu.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a5c8945700e8898906b6ea28b3b822e54d934937811d5a0392778408ccda39b2 -size 985520 diff --git a/venv/lib/python3.10/site-packages/_soundfile.py b/venv/lib/python3.10/site-packages/_soundfile.py deleted file mode 100644 index 660e26597a84d5f5f9bc8032635fedbe2e2eb26e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/_soundfile.py +++ /dev/null @@ -1,11 +0,0 @@ -# auto-generated file -import _cffi_backend - -ffi = _cffi_backend.FFI('_soundfile', - _version = 0x2601, - _types = b'\x00\x00\x12\x0D\x00\x00\x68\x03\x00\x00\x07\x01\x00\x00\x67\x03\x00\x00\x75\x03\x00\x00\x00\x0F\x00\x00\x12\x0D\x00\x00\x6A\x03\x00\x00\x07\x01\x00\x00\x03\x11\x00\x00\x00\x0F\x00\x00\x12\x0D\x00\x00\x07\x01\x00\x00\x07\x01\x00\x00\x03\x11\x00\x00\x07\x01\x00\x00\x00\x0F\x00\x00\x07\x0D\x00\x00\x69\x03\x00\x00\x00\x0F\x00\x00\x07\x0D\x00\x00\x12\x11\x00\x00\x07\x01\x00\x00\x00\x0F\x00\x00\x07\x0D\x00\x00\x07\x01\x00\x00\x00\x0F\x00\x00\x07\x0D\x00\x00\x00\x0F\x00\x00\x02\x0D\x00\x00\x67\x03\x00\x00\x00\x0F\x00\x00\x02\x0D\x00\x00\x12\x11\x00\x00\x00\x0F\x00\x00\x02\x0D\x00\x00\x12\x11\x00\x00\x6A\x03\x00\x00\x1C\x01\x00\x00\x00\x0F\x00\x00\x02\x0D\x00\x00\x12\x11\x00\x00\x07\x01\x00\x00\x07\x11\x00\x00\x00\x0F\x00\x00\x02\x0D\x00\x00\x12\x11\x00\x00\x07\x01\x00\x00\x04\x11\x00\x00\x07\x01\x00\x00\x00\x0F\x00\x00\x36\x0D\x00\x00\x12\x11\x00\x00\x6B\x03\x00\x00\x17\x01\x00\x00\x00\x0F\x00\x00\x36\x0D\x00\x00\x12\x11\x00\x00\x6F\x03\x00\x00\x17\x01\x00\x00\x00\x0F\x00\x00\x36\x0D\x00\x00\x12\x11\x00\x00\x02\x03\x00\x00\x17\x01\x00\x00\x00\x0F\x00\x00\x36\x0D\x00\x00\x12\x11\x00\x00\x17\x01\x00\x00\x07\x01\x00\x00\x00\x0F\x00\x00\x36\x0D\x00\x00\x12\x11\x00\x00\x74\x03\x00\x00\x17\x01\x00\x00\x00\x0F\x00\x00\x36\x0D\x00\x00\x12\x11\x00\x00\x04\x11\x00\x00\x17\x01\x00\x00\x00\x0F\x00\x00\x36\x0D\x00\x00\x17\x01\x00\x00\x07\x01\x00\x00\x04\x11\x00\x00\x00\x0F\x00\x00\x36\x0D\x00\x00\x04\x11\x00\x00\x00\x0F\x00\x00\x36\x0D\x00\x00\x04\x11\x00\x00\x17\x01\x00\x00\x04\x11\x00\x00\x00\x0F\x00\x00\x36\x0D\x00\x00\x75\x03\x00\x00\x17\x01\x00\x00\x04\x11\x00\x00\x00\x0F\x00\x00\x75\x0D\x00\x00\x12\x11\x00\x00\x00\x0F\x00\x00\x00\x09\x00\x00\x01\x09\x00\x00\x02\x09\x00\x00\x03\x09\x00\x00\x02\x01\x00\x00\x0E\x01\x00\x00\x00\x0B\x00\x00\x01\x0B\x00\x00\x02\x0B\x00\x00\x0D\x01\x00\x00\x51\x03\x00\x00\x56\x03\x00\x00\x59\x03\x00\x00\x5E\x03\x00\x00\x05\x01\x00\x00\x00\x01', - _globals = (b'\xFF\xFF\xFF\x0BSFC_FILE_TRUNCATE',4224,b'\xFF\xFF\xFF\x0BSFC_GET_FORMAT_INFO',4136,b'\xFF\xFF\xFF\x0BSFC_GET_FORMAT_MAJOR',4145,b'\xFF\xFF\xFF\x0BSFC_GET_FORMAT_MAJOR_COUNT',4144,b'\xFF\xFF\xFF\x0BSFC_GET_FORMAT_SUBTYPE',4147,b'\xFF\xFF\xFF\x0BSFC_GET_FORMAT_SUBTYPE_COUNT',4146,b'\xFF\xFF\xFF\x0BSFC_GET_LIB_VERSION',4096,b'\xFF\xFF\xFF\x0BSFC_GET_LOG_INFO',4097,b'\xFF\xFF\xFF\x0BSFC_SET_BITRATE_MODE',4869,b'\xFF\xFF\xFF\x0BSFC_SET_CLIPPING',4288,b'\xFF\xFF\xFF\x0BSFC_SET_COMPRESSION_LEVEL',4865,b'\xFF\xFF\xFF\x0BSFC_SET_SCALE_FLOAT_INT_READ',4116,b'\xFF\xFF\xFF\x0BSFC_SET_SCALE_INT_FLOAT_WRITE',4117,b'\xFF\xFF\xFF\x0BSFM_RDWR',48,b'\xFF\xFF\xFF\x0BSFM_READ',16,b'\xFF\xFF\xFF\x0BSFM_WRITE',32,b'\xFF\xFF\xFF\x0BSF_BITRATE_MODE_AVERAGE',1,b'\xFF\xFF\xFF\x0BSF_BITRATE_MODE_CONSTANT',0,b'\xFF\xFF\xFF\x0BSF_BITRATE_MODE_VARIABLE',2,b'\xFF\xFF\xFF\x0BSF_FALSE',0,b'\xFF\xFF\xFF\x0BSF_FORMAT_ENDMASK',805306368,b'\xFF\xFF\xFF\x0BSF_FORMAT_SUBMASK',65535,b'\xFF\xFF\xFF\x0BSF_FORMAT_TYPEMASK',268369920,b'\xFF\xFF\xFF\x0BSF_TRUE',1,b'\x00\x00\x20\x23sf_close',0,b'\x00\x00\x2D\x23sf_command',0,b'\x00\x00\x20\x23sf_error',0,b'\x00\x00\x18\x23sf_error_number',0,b'\x00\x00\x23\x23sf_error_str',0,b'\x00\x00\x1D\x23sf_format_check',0,b'\x00\x00\x14\x23sf_get_string',0,b'\x00\x00\x06\x23sf_open',0,b'\x00\x00\x0B\x23sf_open_fd',0,b'\x00\x00\x00\x23sf_open_virtual',0,b'\x00\x00\x20\x23sf_perror',0,b'\x00\x00\x33\x23sf_read_double',0,b'\x00\x00\x38\x23sf_read_float',0,b'\x00\x00\x3D\x23sf_read_int',0,b'\x00\x00\x4C\x23sf_read_raw',0,b'\x00\x00\x47\x23sf_read_short',0,b'\x00\x00\x4C\x23sf_readf_double',0,b'\x00\x00\x4C\x23sf_readf_float',0,b'\x00\x00\x4C\x23sf_readf_int',0,b'\x00\x00\x4C\x23sf_readf_short',0,b'\x00\x00\x42\x23sf_seek',0,b'\x00\x00\x28\x23sf_set_string',0,b'\x00\x00\x11\x23sf_strerror',0,b'\x00\x00\x1B\x23sf_version_string',0,b'\x00\x00\x33\x23sf_write_double',0,b'\x00\x00\x38\x23sf_write_float',0,b'\x00\x00\x3D\x23sf_write_int',0,b'\x00\x00\x4C\x23sf_write_raw',0,b'\x00\x00\x47\x23sf_write_short',0,b'\x00\x00\x63\x23sf_write_sync',0,b'\x00\x00\x4C\x23sf_writef_double',0,b'\x00\x00\x4C\x23sf_writef_float',0,b'\x00\x00\x4C\x23sf_writef_int',0,b'\x00\x00\x4C\x23sf_writef_short',0), - _struct_unions = ((b'\x00\x00\x00\x66\x00\x00\x00\x02SF_FORMAT_INFO',b'\x00\x00\x02\x11format',b'\x00\x00\x07\x11name',b'\x00\x00\x07\x11extension'),(b'\x00\x00\x00\x67\x00\x00\x00\x02SF_INFO',b'\x00\x00\x36\x11frames',b'\x00\x00\x02\x11samplerate',b'\x00\x00\x02\x11channels',b'\x00\x00\x02\x11format',b'\x00\x00\x02\x11sections',b'\x00\x00\x02\x11seekable'),(b'\x00\x00\x00\x68\x00\x00\x00\x02SF_VIRTUAL_IO',b'\x00\x00\x71\x11get_filelen',b'\x00\x00\x70\x11seek',b'\x00\x00\x72\x11read',b'\x00\x00\x73\x11write',b'\x00\x00\x71\x11tell'),(b'\x00\x00\x00\x69\x00\x00\x00\x10SNDFILE_tag',)), - _enums = (b'\x00\x00\x00\x6C\x00\x00\x00\x16$1\x00SF_FORMAT_SUBMASK,SF_FORMAT_TYPEMASK,SF_FORMAT_ENDMASK',b'\x00\x00\x00\x6D\x00\x00\x00\x16$2\x00SFC_GET_LIB_VERSION,SFC_GET_LOG_INFO,SFC_GET_FORMAT_INFO,SFC_GET_FORMAT_MAJOR_COUNT,SFC_GET_FORMAT_MAJOR,SFC_GET_FORMAT_SUBTYPE_COUNT,SFC_GET_FORMAT_SUBTYPE,SFC_FILE_TRUNCATE,SFC_SET_CLIPPING,SFC_SET_SCALE_FLOAT_INT_READ,SFC_SET_SCALE_INT_FLOAT_WRITE,SFC_SET_COMPRESSION_LEVEL,SFC_SET_BITRATE_MODE',b'\x00\x00\x00\x6E\x00\x00\x00\x16$3\x00SF_FALSE,SF_TRUE,SFM_READ,SFM_WRITE,SFM_RDWR,SF_BITRATE_MODE_CONSTANT,SF_BITRATE_MODE_AVERAGE,SF_BITRATE_MODE_VARIABLE'), - _typenames = (b'\x00\x00\x00\x66SF_FORMAT_INFO',b'\x00\x00\x00\x67SF_INFO',b'\x00\x00\x00\x68SF_VIRTUAL_IO',b'\x00\x00\x00\x69SNDFILE',b'\x00\x00\x00\x36sf_count_t',b'\x00\x00\x00\x71sf_vio_get_filelen',b'\x00\x00\x00\x72sf_vio_read',b'\x00\x00\x00\x70sf_vio_seek',b'\x00\x00\x00\x71sf_vio_tell',b'\x00\x00\x00\x73sf_vio_write'), -) diff --git a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/METADATA b/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/METADATA deleted file mode 100644 index 759c8515bad1a667f2d84ea773eb86a46c3e3581..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/METADATA +++ /dev/null @@ -1,397 +0,0 @@ -Metadata-Version: 2.4 -Name: accelerate -Version: 1.9.0 -Summary: Accelerate -Home-page: https://github.com/huggingface/accelerate -Author: The HuggingFace team -Author-email: zach.mueller@huggingface.co -License: Apache -Keywords: deep learning -Classifier: Development Status :: 5 - Production/Stable -Classifier: Intended Audience :: Developers -Classifier: Intended Audience :: Education -Classifier: Intended Audience :: Science/Research -Classifier: License :: OSI Approved :: Apache Software License -Classifier: Operating System :: OS Independent -Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence -Requires-Python: >=3.9.0 -Description-Content-Type: text/markdown -License-File: LICENSE -Requires-Dist: numpy<3.0.0,>=1.17 -Requires-Dist: packaging>=20.0 -Requires-Dist: psutil -Requires-Dist: pyyaml -Requires-Dist: torch>=2.0.0 -Requires-Dist: huggingface_hub>=0.21.0 -Requires-Dist: safetensors>=0.4.3 -Provides-Extra: quality -Requires-Dist: black~=23.1; extra == "quality" -Requires-Dist: hf-doc-builder>=0.3.0; extra == "quality" -Requires-Dist: ruff~=0.11.2; extra == "quality" -Provides-Extra: docs -Provides-Extra: test-prod -Requires-Dist: pytest<=8.0.0,>=7.2.0; extra == "test-prod" -Requires-Dist: pytest-xdist; extra == "test-prod" -Requires-Dist: pytest-subtests; extra == "test-prod" -Requires-Dist: parameterized; extra == "test-prod" -Requires-Dist: pytest-order; extra == "test-prod" -Provides-Extra: test-dev -Requires-Dist: datasets; extra == "test-dev" -Requires-Dist: diffusers; extra == "test-dev" -Requires-Dist: evaluate; extra == "test-dev" -Requires-Dist: torchdata>=0.8.0; extra == "test-dev" -Requires-Dist: torchpippy>=0.2.0; extra == "test-dev" -Requires-Dist: transformers; extra == "test-dev" -Requires-Dist: scipy; extra == "test-dev" -Requires-Dist: scikit-learn; extra == "test-dev" -Requires-Dist: tqdm; extra == "test-dev" -Requires-Dist: bitsandbytes; extra == "test-dev" -Requires-Dist: timm; extra == "test-dev" -Provides-Extra: testing -Requires-Dist: pytest<=8.0.0,>=7.2.0; extra == "testing" -Requires-Dist: pytest-xdist; extra == "testing" -Requires-Dist: pytest-subtests; extra == "testing" -Requires-Dist: parameterized; extra == "testing" -Requires-Dist: pytest-order; extra == "testing" -Requires-Dist: datasets; extra == "testing" -Requires-Dist: diffusers; extra == "testing" -Requires-Dist: evaluate; extra == "testing" -Requires-Dist: torchdata>=0.8.0; extra == "testing" -Requires-Dist: torchpippy>=0.2.0; extra == "testing" -Requires-Dist: transformers; extra == "testing" -Requires-Dist: scipy; extra == "testing" -Requires-Dist: scikit-learn; extra == "testing" -Requires-Dist: tqdm; extra == "testing" -Requires-Dist: bitsandbytes; extra == "testing" -Requires-Dist: timm; extra == "testing" -Provides-Extra: deepspeed -Requires-Dist: deepspeed; extra == "deepspeed" -Provides-Extra: rich -Requires-Dist: rich; extra == "rich" -Provides-Extra: test-fp8 -Requires-Dist: torchao; extra == "test-fp8" -Provides-Extra: test-trackers -Requires-Dist: wandb; extra == "test-trackers" -Requires-Dist: comet-ml; extra == "test-trackers" -Requires-Dist: tensorboard; extra == "test-trackers" -Requires-Dist: dvclive; extra == "test-trackers" -Requires-Dist: mlflow; extra == "test-trackers" -Requires-Dist: matplotlib; extra == "test-trackers" -Requires-Dist: swanlab; extra == "test-trackers" -Requires-Dist: trackio; extra == "test-trackers" -Provides-Extra: dev -Requires-Dist: black~=23.1; extra == "dev" -Requires-Dist: hf-doc-builder>=0.3.0; extra == "dev" -Requires-Dist: ruff~=0.11.2; extra == "dev" -Requires-Dist: pytest<=8.0.0,>=7.2.0; extra == "dev" -Requires-Dist: pytest-xdist; extra == "dev" -Requires-Dist: pytest-subtests; extra == "dev" -Requires-Dist: parameterized; extra == "dev" -Requires-Dist: pytest-order; extra == "dev" -Requires-Dist: datasets; extra == "dev" -Requires-Dist: diffusers; extra == "dev" -Requires-Dist: evaluate; extra == "dev" -Requires-Dist: torchdata>=0.8.0; extra == "dev" -Requires-Dist: torchpippy>=0.2.0; extra == "dev" -Requires-Dist: transformers; extra == "dev" -Requires-Dist: scipy; extra == "dev" -Requires-Dist: scikit-learn; extra == "dev" -Requires-Dist: tqdm; extra == "dev" -Requires-Dist: bitsandbytes; extra == "dev" -Requires-Dist: timm; extra == "dev" -Requires-Dist: rich; extra == "dev" -Provides-Extra: sagemaker -Requires-Dist: sagemaker; extra == "sagemaker" -Dynamic: author -Dynamic: author-email -Dynamic: classifier -Dynamic: description -Dynamic: description-content-type -Dynamic: home-page -Dynamic: keywords -Dynamic: license -Dynamic: license-file -Dynamic: provides-extra -Dynamic: requires-dist -Dynamic: requires-python -Dynamic: summary - - - -

-
- -
-

- -

- - License - Documentation - GitHub release - Contributor Covenant -

- -

-

Run your *raw* PyTorch training script on any kind of device -

- -

- -

- -## Easy to integrate - -🤗 Accelerate was created for PyTorch users who like to write the training loop of PyTorch models but are reluctant to write and maintain the boilerplate code needed to use multi-GPUs/TPU/fp16. - -🤗 Accelerate abstracts exactly and only the boilerplate code related to multi-GPUs/TPU/fp16 and leaves the rest of your code unchanged. - -Here is an example: - -```diff - import torch - import torch.nn.functional as F - from datasets import load_dataset -+ from accelerate import Accelerator - -+ accelerator = Accelerator() -- device = 'cpu' -+ device = accelerator.device - - model = torch.nn.Transformer().to(device) - optimizer = torch.optim.Adam(model.parameters()) - - dataset = load_dataset('my_dataset') - data = torch.utils.data.DataLoader(dataset, shuffle=True) - -+ model, optimizer, data = accelerator.prepare(model, optimizer, data) - - model.train() - for epoch in range(10): - for source, targets in data: - source = source.to(device) - targets = targets.to(device) - - optimizer.zero_grad() - - output = model(source) - loss = F.cross_entropy(output, targets) - -- loss.backward() -+ accelerator.backward(loss) - - optimizer.step() -``` - -As you can see in this example, by adding 5-lines to any standard PyTorch training script you can now run on any kind of single or distributed node setting (single CPU, single GPU, multi-GPUs and TPUs) as well as with or without mixed precision (fp8, fp16, bf16). - -In particular, the same code can then be run without modification on your local machine for debugging or your training environment. - -🤗 Accelerate even handles the device placement for you (which requires a few more changes to your code, but is safer in general), so you can even simplify your training loop further: - -```diff - import torch - import torch.nn.functional as F - from datasets import load_dataset -+ from accelerate import Accelerator - -- device = 'cpu' -+ accelerator = Accelerator() - -- model = torch.nn.Transformer().to(device) -+ model = torch.nn.Transformer() - optimizer = torch.optim.Adam(model.parameters()) - - dataset = load_dataset('my_dataset') - data = torch.utils.data.DataLoader(dataset, shuffle=True) - -+ model, optimizer, data = accelerator.prepare(model, optimizer, data) - - model.train() - for epoch in range(10): - for source, targets in data: -- source = source.to(device) -- targets = targets.to(device) - - optimizer.zero_grad() - - output = model(source) - loss = F.cross_entropy(output, targets) - -- loss.backward() -+ accelerator.backward(loss) - - optimizer.step() -``` - -Want to learn more? Check out the [documentation](https://huggingface.co/docs/accelerate) or have a look at our [examples](https://github.com/huggingface/accelerate/tree/main/examples). - -## Launching script - -🤗 Accelerate also provides an optional CLI tool that allows you to quickly configure and test your training environment before launching the scripts. No need to remember how to use `torch.distributed.run` or to write a specific launcher for TPU training! -On your machine(s) just run: - -```bash -accelerate config -``` - -and answer the questions asked. This will generate a config file that will be used automatically to properly set the default options when doing - -```bash -accelerate launch my_script.py --args_to_my_script -``` - -For instance, here is how you would run the GLUE example on the MRPC task (from the root of the repo): - -```bash -accelerate launch examples/nlp_example.py -``` - -This CLI tool is **optional**, and you can still use `python my_script.py` or `python -m torchrun my_script.py` at your convenience. - -You can also directly pass in the arguments you would to `torchrun` as arguments to `accelerate launch` if you wish to not run` accelerate config`. - -For example, here is how to launch on two GPUs: - -```bash -accelerate launch --multi_gpu --num_processes 2 examples/nlp_example.py -``` - -To learn more, check the CLI documentation available [here](https://huggingface.co/docs/accelerate/package_reference/cli). - -Or view the configuration zoo [here](https://github.com/huggingface/accelerate/blob/main/examples/config_yaml_templates/) - -## Launching multi-CPU run using MPI - -🤗 Here is another way to launch multi-CPU run using MPI. You can learn how to install Open MPI on [this page](https://www.open-mpi.org/faq/?category=building#easy-build). You can use Intel MPI or MVAPICH as well. -Once you have MPI setup on your cluster, just run: -```bash -accelerate config -``` -Answer the questions that are asked, selecting to run using multi-CPU, and answer "yes" when asked if you want accelerate to launch mpirun. -Then, use `accelerate launch` with your script like: -```bash -accelerate launch examples/nlp_example.py -``` -Alternatively, you can use mpirun directly, without using the CLI like: -```bash -mpirun -np 2 python examples/nlp_example.py -``` - -## Launching training using DeepSpeed - -🤗 Accelerate supports training on single/multiple GPUs using DeepSpeed. To use it, you don't need to change anything in your training code; you can set everything using just `accelerate config`. However, if you desire to tweak your DeepSpeed related args from your Python script, we provide you the `DeepSpeedPlugin`. - -```python -from accelerate import Accelerator, DeepSpeedPlugin - -# deepspeed needs to know your gradient accumulation steps beforehand, so don't forget to pass it -# Remember you still need to do gradient accumulation by yourself, just like you would have done without deepspeed -deepspeed_plugin = DeepSpeedPlugin(zero_stage=2, gradient_accumulation_steps=2) -accelerator = Accelerator(mixed_precision='fp16', deepspeed_plugin=deepspeed_plugin) - -# How to save your 🤗 Transformer? -accelerator.wait_for_everyone() -unwrapped_model = accelerator.unwrap_model(model) -unwrapped_model.save_pretrained(save_dir, save_function=accelerator.save, state_dict=accelerator.get_state_dict(model)) -``` - -Note: DeepSpeed support is experimental for now. In case you get into some problem, please open an issue. - -## Launching your training from a notebook - -🤗 Accelerate also provides a `notebook_launcher` function you can use in a notebook to launch a distributed training. This is especially useful for Colab or Kaggle notebooks with a TPU backend. Just define your training loop in a `training_function` then in your last cell, add: - -```python -from accelerate import notebook_launcher - -notebook_launcher(training_function) -``` - -An example can be found in [this notebook](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb). [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb) - -## Why should I use 🤗 Accelerate? - -You should use 🤗 Accelerate when you want to easily run your training scripts in a distributed environment without having to renounce full control over your training loop. This is not a high-level framework above PyTorch, just a thin wrapper so you don't have to learn a new library. In fact, the whole API of 🤗 Accelerate is in one class, the `Accelerator` object. - -## Why shouldn't I use 🤗 Accelerate? - -You shouldn't use 🤗 Accelerate if you don't want to write a training loop yourself. There are plenty of high-level libraries above PyTorch that will offer you that, 🤗 Accelerate is not one of them. - -## Frameworks using 🤗 Accelerate - -If you like the simplicity of 🤗 Accelerate but would prefer a higher-level abstraction around its capabilities, some frameworks and libraries that are built on top of 🤗 Accelerate are listed below: - -* [Amphion](https://github.com/open-mmlab/Amphion) is a toolkit for Audio, Music, and Speech Generation. Its purpose is to support reproducible research and help junior researchers and engineers get started in the field of audio, music, and speech generation research and development. -* [Animus](https://github.com/Scitator/animus) is a minimalistic framework to run machine learning experiments. Animus highlights common "breakpoints" in ML experiments and provides a unified interface for them within [IExperiment](https://github.com/Scitator/animus/blob/main/animus/core.py#L76). -* [Catalyst](https://github.com/catalyst-team/catalyst#getting-started) is a PyTorch framework for Deep Learning Research and Development. It focuses on reproducibility, rapid experimentation, and codebase reuse so you can create something new rather than write yet another train loop. Catalyst provides a [Runner](https://catalyst-team.github.io/catalyst/api/core.html#runner) to connect all parts of the experiment: hardware backend, data transformations, model training, and inference logic. -* [fastai](https://github.com/fastai/fastai#installing) is a PyTorch framework for Deep Learning that simplifies training fast and accurate neural nets using modern best practices. fastai provides a [Learner](https://docs.fast.ai/learner.html#Learner) to handle the training, fine-tuning, and inference of deep learning algorithms. -* [Finetuner](https://github.com/jina-ai/finetuner) is a service that enables models to create higher-quality embeddings for semantic search, visual similarity search, cross-modal text<->image search, recommendation systems, clustering, duplication detection, anomaly detection, or other uses. -* [InvokeAI](https://github.com/invoke-ai/InvokeAI) is a creative engine for Stable Diffusion models, offering industry-leading WebUI, terminal usage support, and serves as the foundation for many commercial products. -* [Kornia](https://kornia.readthedocs.io/en/latest/get-started/introduction.html) is a differentiable library that allows classical computer vision to be integrated into deep learning models. Kornia provides a [Trainer](https://kornia.readthedocs.io/en/latest/x.html#kornia.x.Trainer) with the specific purpose to train and fine-tune the supported deep learning algorithms within the library. -* [Open Assistant](https://projects.laion.ai/Open-Assistant/) is a chat-based assistant that understands tasks, can interact with their party systems, and retrieve information dynamically to do so. -* [pytorch-accelerated](https://github.com/Chris-hughes10/pytorch-accelerated) is a lightweight training library, with a streamlined feature set centered around a general-purpose [Trainer](https://pytorch-accelerated.readthedocs.io/en/latest/trainer.html), that places a huge emphasis on simplicity and transparency; enabling users to understand exactly what is going on under the hood, but without having to write and maintain the boilerplate themselves! -* [Stable Diffusion web UI](https://github.com/AUTOMATIC1111/stable-diffusion-webui) is an open-source browser-based easy-to-use interface based on the Gradio library for Stable Diffusion. -* [torchkeras](https://github.com/lyhue1991/torchkeras) is a simple tool for training pytorch model just in a keras style, a dynamic and beautiful plot is provided in notebook to monitor your loss or metric. -* [transformers](https://github.com/huggingface/transformers) as a tool for helping train state-of-the-art machine learning models in PyTorch, Tensorflow, and JAX. (Accelerate is the backend for the PyTorch side). - - -## Installation - -This repository is tested on Python 3.8+ and PyTorch 1.10.0+ - -You should install 🤗 Accelerate in a [virtual environment](https://docs.python.org/3/library/venv.html). If you're unfamiliar with Python virtual environments, check out the [user guide](https://packaging.python.org/guides/installing-using-pip-and-virtual-environments/). - -First, create a virtual environment with the version of Python you're going to use and activate it. - -Then, you will need to install PyTorch: refer to the [official installation page](https://pytorch.org/get-started/locally/#start-locally) regarding the specific install command for your platform. Then 🤗 Accelerate can be installed using pip as follows: - -```bash -pip install accelerate -``` - -## Supported integrations - -- CPU only -- multi-CPU on one node (machine) -- multi-CPU on several nodes (machines) -- single GPU -- multi-GPU on one node (machine) -- multi-GPU on several nodes (machines) -- TPU -- FP16/BFloat16 mixed precision -- FP8 mixed precision with [Transformer Engine](https://github.com/NVIDIA/TransformerEngine) or [MS-AMP](https://github.com/Azure/MS-AMP/) -- DeepSpeed support (Experimental) -- PyTorch Fully Sharded Data Parallel (FSDP) support (Experimental) -- Megatron-LM support (Experimental) - -## Citing 🤗 Accelerate - -If you use 🤗 Accelerate in your publication, please cite it by using the following BibTeX entry. - -```bibtex -@Misc{accelerate, - title = {Accelerate: Training and inference at scale made simple, efficient and adaptable.}, - author = {Sylvain Gugger and Lysandre Debut and Thomas Wolf and Philipp Schmid and Zachary Mueller and Sourab Mangrulkar and Marc Sun and Benjamin Bossan}, - howpublished = {\url{https://github.com/huggingface/accelerate}}, - year = {2022} -} -``` diff --git a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/RECORD b/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/RECORD deleted file mode 100644 index 1b303e57ceb3ebd16870a7970f7e6b1f66dc9ed8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/RECORD +++ /dev/null @@ -1,177 +0,0 @@ -../../../bin/accelerate,sha256=pJLu-G8APkkbztTcu7dVc9nJQqoWaX6tRJMFbGk4rYY,251 -../../../bin/accelerate-config,sha256=5kNl4YYg0LvQP_PSDMsBgOEznOasBtSkMlPRHBdKu14,243 -../../../bin/accelerate-estimate-memory,sha256=SjoEbmGByXY58ySjRkLX4r_lpsqDOevAaGewm4SDzQ8,245 -../../../bin/accelerate-launch,sha256=vIEs89n_ZJS1V7y2Ggvt_9KGNSNyFZatek118eKrvWI,243 -../../../bin/accelerate-merge-weights,sha256=U_3uPjQ1AVtg0rUvrT4mScxYcue2lKnEqYVBADZ2UcE,242 -accelerate-1.9.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -accelerate-1.9.0.dist-info/METADATA,sha256=nHhgqcSMeeHtgWVysVrwXnLx8yO-CqvYKFWNgQpTHsM,19866 -accelerate-1.9.0.dist-info/RECORD,, -accelerate-1.9.0.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -accelerate-1.9.0.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91 -accelerate-1.9.0.dist-info/entry_points.txt,sha256=Vpy8gUGfZ-1VnM2229fb8CpJNLBdMH_wtJ9PQ7b_2tQ,296 -accelerate-1.9.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357 -accelerate-1.9.0.dist-info/top_level.txt,sha256=esVfdxTidsjQ90zsN_rPpjLFJ4ijRlx4mnLrG09hlt4,11 -accelerate/__init__.py,sha256=DOfEIpDInu8FrrQgZuJXl_mj_H2XIPjcKacOidSd174,1504 -accelerate/__pycache__/__init__.cpython-310.pyc,, -accelerate/__pycache__/accelerator.cpython-310.pyc,, -accelerate/__pycache__/big_modeling.cpython-310.pyc,, -accelerate/__pycache__/checkpointing.cpython-310.pyc,, -accelerate/__pycache__/data_loader.cpython-310.pyc,, -accelerate/__pycache__/hooks.cpython-310.pyc,, -accelerate/__pycache__/inference.cpython-310.pyc,, -accelerate/__pycache__/launchers.cpython-310.pyc,, -accelerate/__pycache__/local_sgd.cpython-310.pyc,, -accelerate/__pycache__/logging.cpython-310.pyc,, -accelerate/__pycache__/memory_utils.cpython-310.pyc,, -accelerate/__pycache__/optimizer.cpython-310.pyc,, -accelerate/__pycache__/scheduler.cpython-310.pyc,, -accelerate/__pycache__/state.cpython-310.pyc,, -accelerate/__pycache__/tracking.cpython-310.pyc,, -accelerate/accelerator.py,sha256=Z0TCu7R50DQCvIk_VVmUrx0cPDs7PEWldDB585ByreE,181670 -accelerate/big_modeling.py,sha256=bqqO4Lw00YErTup7MUM6saOVHRFvMNOdNqAwg96fBG0,34765 -accelerate/checkpointing.py,sha256=ARi8ZNFlLQW2uz9hd2D5-5rRUNoY2zL7BSe8HQKBFjs,13957 -accelerate/commands/__init__.py,sha256=m1PPTDT4ziIAvM0-FDSgIMIZ69Konn126s6LwuzH6v8,606 -accelerate/commands/__pycache__/__init__.cpython-310.pyc,, -accelerate/commands/__pycache__/accelerate_cli.cpython-310.pyc,, -accelerate/commands/__pycache__/env.cpython-310.pyc,, -accelerate/commands/__pycache__/estimate.cpython-310.pyc,, -accelerate/commands/__pycache__/launch.cpython-310.pyc,, -accelerate/commands/__pycache__/merge.cpython-310.pyc,, -accelerate/commands/__pycache__/test.cpython-310.pyc,, -accelerate/commands/__pycache__/to_fsdp2.cpython-310.pyc,, -accelerate/commands/__pycache__/tpu.cpython-310.pyc,, -accelerate/commands/__pycache__/utils.cpython-310.pyc,, -accelerate/commands/accelerate_cli.py,sha256=SkwFad6Z1ZsGjtm7TiXFq8je-akshp_0WxX_6rGSBw8,1972 -accelerate/commands/config/__init__.py,sha256=iJK8dgj3pc5Vdr1E7UuGoFu-BlybyXLxYDoTg9gXngE,1645 -accelerate/commands/config/__pycache__/__init__.cpython-310.pyc,, -accelerate/commands/config/__pycache__/cluster.cpython-310.pyc,, -accelerate/commands/config/__pycache__/config.cpython-310.pyc,, -accelerate/commands/config/__pycache__/config_args.cpython-310.pyc,, -accelerate/commands/config/__pycache__/config_utils.cpython-310.pyc,, -accelerate/commands/config/__pycache__/default.cpython-310.pyc,, -accelerate/commands/config/__pycache__/sagemaker.cpython-310.pyc,, -accelerate/commands/config/__pycache__/update.cpython-310.pyc,, -accelerate/commands/config/cluster.py,sha256=Ln8AbiHX0fDLYV8Gxk3WjC2JL94CEAjuhhn6BKdQ1zE,37544 -accelerate/commands/config/config.py,sha256=FuRlQvOjgATEtyqOSsGD-KEtOCvACOHjs2C-krrtldk,3035 -accelerate/commands/config/config_args.py,sha256=hJo7E3J_T7MPNkgtyNjPilJuH-wJK09mGWF6CvGFYIQ,9970 -accelerate/commands/config/config_utils.py,sha256=lbV9iKAVSF5bUyQDGpzFduhH5JENYgAilPpRLLlCEq8,3299 -accelerate/commands/config/default.py,sha256=UY0dQ7w1peBkQe8QDCm3XJrTnsRcm_G-_hf44NzNAbU,6213 -accelerate/commands/config/sagemaker.py,sha256=SRLup9XrV_Z-23A3d8gZbkqbonDcscMDDmbpmytCzbY,10636 -accelerate/commands/config/update.py,sha256=NXW1J7GkUHpg71QlIXsmMB_0z8S8IZo2FWax5POwrhc,2395 -accelerate/commands/env.py,sha256=gRGwtPhMMWmZ5NTMaMrgHXaOzBV6SvAtP78TuuJvC6g,4336 -accelerate/commands/estimate.py,sha256=Qduq4xudVyIede37BMEe1rNhXf-rfW-MHV2KtwxdfEA,12585 -accelerate/commands/launch.py,sha256=h-d1W1y6NEwciGdzeLY-aTQG2N6dx3iBgcwSZZ_Yjn0,46889 -accelerate/commands/menu/__init__.py,sha256=uqSlBM0TFHBwzdv3p3SXfpAk1lZFp4h1a7mbBdscPHs,645 -accelerate/commands/menu/__pycache__/__init__.cpython-310.pyc,, -accelerate/commands/menu/__pycache__/cursor.cpython-310.pyc,, -accelerate/commands/menu/__pycache__/helpers.cpython-310.pyc,, -accelerate/commands/menu/__pycache__/input.cpython-310.pyc,, -accelerate/commands/menu/__pycache__/keymap.cpython-310.pyc,, -accelerate/commands/menu/__pycache__/selection_menu.cpython-310.pyc,, -accelerate/commands/menu/cursor.py,sha256=-lmpJVAzvNc0c3EOtSuLoKB59zqylVCbYyWLPnrOmvQ,2028 -accelerate/commands/menu/helpers.py,sha256=KrSB5fJjH4MUEUAQJ6bYaN16AYcnl9UalDrPD3DYeeg,1483 -accelerate/commands/menu/input.py,sha256=T8Mdd-Y_OURgqfDV9qZh4Wf6hmT22AneNtJzj4JA1Rk,2512 -accelerate/commands/menu/keymap.py,sha256=eXj-suyYs1m5dEHoUKN4mKAMLc8DWHnwhP6G6JSU0jQ,4086 -accelerate/commands/menu/selection_menu.py,sha256=bxy-DHaKKC6SCToOlMBv5_z0MdUzylEg6Sio9OuV3GM,4921 -accelerate/commands/merge.py,sha256=3Uo14LADvk8RJwZ_i6kIsGwWB3Ul5y8Uq47oHkDYqKU,2387 -accelerate/commands/test.py,sha256=YrPYEaAACOGZ6btn2MV6NbMSEdBUcMWADLbQWaZSHtk,2149 -accelerate/commands/to_fsdp2.py,sha256=gfbhoUT4qFB3LVDMNmckElgLG0yWm8aj_aofszeiJmM,5991 -accelerate/commands/tpu.py,sha256=KyxDP7IuveidZrbW4rx2s8Ku3o_ptI6tzwr_R7ck0os,5548 -accelerate/commands/utils.py,sha256=aT8xUCe2pCkFII7yZxcfaohEjgBAzMUM7WiD4UuWSOY,4150 -accelerate/data_loader.py,sha256=ecNfEwXrFqhBzUEaa6_IUvr-yMfPcj1AoATqSffyUtU,65318 -accelerate/hooks.py,sha256=y4eN7blCjOkXY9_mVBmkgiTKqJfasAKrHOddEHj91Vg,33293 -accelerate/inference.py,sha256=NLANdzXm5PwmDWbPYkFmoRoQSLLvuhfvIG33xfpapT0,7668 -accelerate/launchers.py,sha256=gKB8sH5YwV-hOM8ptbzHgm16729X6864vKdLOls4eUY,14176 -accelerate/local_sgd.py,sha256=aCj_yqXK_FhhZRWEpzXIkgXBERH6fC3HyrC3nsOj1uA,4160 -accelerate/logging.py,sha256=4XcgY_BV7Qn_enh2tZ-8fNtuaE_3n-LsYJbgwhRx_PI,5042 -accelerate/memory_utils.py,sha256=3R5LoeHl6GgTZ-IMPrDZMdaEehWarGdPqODushb-6pg,862 -accelerate/optimizer.py,sha256=p5qfCOMi8DAjvx6eECypo4ej3pvmQ_U_bgmrEAKc4hY,8158 -accelerate/scheduler.py,sha256=des_4M_Tt1W8gCYZZbLla0GHBEgJY3Wx2EGBQPTzeiY,4238 -accelerate/state.py,sha256=7GgdLMjGHLuv0t4SftscosFgbgIXTO0Itp1VwG9qdaM,56508 -accelerate/test_utils/__init__.py,sha256=ROl3bvBqmuI6ueRlPRNkNC-83yb6u6z_FFzrHJtq8rM,1809 -accelerate/test_utils/__pycache__/__init__.cpython-310.pyc,, -accelerate/test_utils/__pycache__/examples.cpython-310.pyc,, -accelerate/test_utils/__pycache__/testing.cpython-310.pyc,, -accelerate/test_utils/__pycache__/training.cpython-310.pyc,, -accelerate/test_utils/examples.py,sha256=CO3SlkkrORUjroreoPfBPwfcNCDlHSgNWIt81iz2oQ8,7257 -accelerate/test_utils/scripts/__init__.py,sha256=m1PPTDT4ziIAvM0-FDSgIMIZ69Konn126s6LwuzH6v8,606 -accelerate/test_utils/scripts/__pycache__/__init__.cpython-310.pyc,, -accelerate/test_utils/scripts/__pycache__/test_cli.cpython-310.pyc,, -accelerate/test_utils/scripts/__pycache__/test_ddp_comm_hook.cpython-310.pyc,, -accelerate/test_utils/scripts/__pycache__/test_distributed_data_loop.cpython-310.pyc,, -accelerate/test_utils/scripts/__pycache__/test_merge_weights.cpython-310.pyc,, -accelerate/test_utils/scripts/__pycache__/test_notebook.cpython-310.pyc,, -accelerate/test_utils/scripts/__pycache__/test_ops.cpython-310.pyc,, -accelerate/test_utils/scripts/__pycache__/test_script.cpython-310.pyc,, -accelerate/test_utils/scripts/__pycache__/test_sync.cpython-310.pyc,, -accelerate/test_utils/scripts/external_deps/__init__.py,sha256=m1PPTDT4ziIAvM0-FDSgIMIZ69Konn126s6LwuzH6v8,606 -accelerate/test_utils/scripts/external_deps/__pycache__/__init__.cpython-310.pyc,, -accelerate/test_utils/scripts/external_deps/__pycache__/test_checkpointing.cpython-310.pyc,, -accelerate/test_utils/scripts/external_deps/__pycache__/test_ds_multiple_model.cpython-310.pyc,, -accelerate/test_utils/scripts/external_deps/__pycache__/test_metrics.cpython-310.pyc,, -accelerate/test_utils/scripts/external_deps/__pycache__/test_peak_memory_usage.cpython-310.pyc,, -accelerate/test_utils/scripts/external_deps/__pycache__/test_performance.cpython-310.pyc,, -accelerate/test_utils/scripts/external_deps/__pycache__/test_pippy.cpython-310.pyc,, -accelerate/test_utils/scripts/external_deps/__pycache__/test_zero3_integration.cpython-310.pyc,, -accelerate/test_utils/scripts/external_deps/test_checkpointing.py,sha256=XHaNRmnrARd1izXFjWGi5UjYGas-4vqayW51jAHBPCA,10699 -accelerate/test_utils/scripts/external_deps/test_ds_multiple_model.py,sha256=Cg4-h0B4UcOQ5CxXjIdrsPVR5fFsWCv24DqZGjXEwW8,13790 -accelerate/test_utils/scripts/external_deps/test_metrics.py,sha256=Ev2XKaiwmznoxKujskAAuISGChW646MOiyf0CXEPb9Y,12168 -accelerate/test_utils/scripts/external_deps/test_peak_memory_usage.py,sha256=9Yn9Rc7d-yWr1fU0RagASPG5l8vrKeHVYbuYABbA-fU,12498 -accelerate/test_utils/scripts/external_deps/test_performance.py,sha256=4SW108BHEdpzDA_VY4B0GKKAdms4QxVlVywhZ-CZwRI,11721 -accelerate/test_utils/scripts/external_deps/test_pippy.py,sha256=rYF5bqSCe_pfylxDv2_2Q7S2mHIb9_6fGns-IcLgE_E,4789 -accelerate/test_utils/scripts/external_deps/test_zero3_integration.py,sha256=P9alBOHZ9Lfqs5LoRP7bCbXl-tnsNrBkvJZGseibBeA,1665 -accelerate/test_utils/scripts/test_cli.py,sha256=-p78NTJgKdeTNDqyrQuGEZZqMpnzdUHPyWZ05FMl5EI,1081 -accelerate/test_utils/scripts/test_ddp_comm_hook.py,sha256=k_-2MBjLKNdMGIcneTbuGd84K05Wp1GEQX6DUVF9UBw,3566 -accelerate/test_utils/scripts/test_distributed_data_loop.py,sha256=RUWTwd7DIpr2fl7JtKOsvTjMiJioTxO8FdSr2Lw_5uI,15137 -accelerate/test_utils/scripts/test_merge_weights.py,sha256=r1oqPOMOqLe5zWenWWTyl9cDbbkFJadpNNiPNBlgR18,5982 -accelerate/test_utils/scripts/test_notebook.py,sha256=qfIy3IvH74-kGn8nadBn_k7qrviqvsxy5ijsnUhuY6o,3894 -accelerate/test_utils/scripts/test_ops.py,sha256=Bcs-h8EMJwULTfbizlFN5qkv3JraWEpoSZWMn-HswiI,6265 -accelerate/test_utils/scripts/test_script.py,sha256=zns5-YNj5h3YHBk79XisiqDQfX2nl9UL8K_NS5IoqFg,36446 -accelerate/test_utils/scripts/test_sync.py,sha256=PDe8sYZLCL2LKjj_L9b-Bh2BjAjeii9EZ8sZNfuYx5s,18817 -accelerate/test_utils/testing.py,sha256=5cb7oMSn4y5vvcWaTF-1rWQ6YB6h_XaKKxOd81T8oM4,29698 -accelerate/test_utils/training.py,sha256=jO5YEIr34jAcnJ_9WNp_x3zuHzSam_I6IgMvmcGm7yI,6456 -accelerate/tracking.py,sha256=GNJRg-yFrDP5D4_Ru1PwJI3uBQdw_FuSH1_3fEFrlww,51450 -accelerate/utils/__init__.py,sha256=9oseM-5xSca7nKEQvlvdrMifxDVcPVp5tuK6KwAbfuo,8029 -accelerate/utils/__pycache__/__init__.cpython-310.pyc,, -accelerate/utils/__pycache__/ao.cpython-310.pyc,, -accelerate/utils/__pycache__/bnb.cpython-310.pyc,, -accelerate/utils/__pycache__/constants.cpython-310.pyc,, -accelerate/utils/__pycache__/dataclasses.cpython-310.pyc,, -accelerate/utils/__pycache__/deepspeed.cpython-310.pyc,, -accelerate/utils/__pycache__/environment.cpython-310.pyc,, -accelerate/utils/__pycache__/fsdp_utils.cpython-310.pyc,, -accelerate/utils/__pycache__/imports.cpython-310.pyc,, -accelerate/utils/__pycache__/launch.cpython-310.pyc,, -accelerate/utils/__pycache__/megatron_lm.cpython-310.pyc,, -accelerate/utils/__pycache__/memory.cpython-310.pyc,, -accelerate/utils/__pycache__/modeling.cpython-310.pyc,, -accelerate/utils/__pycache__/offload.cpython-310.pyc,, -accelerate/utils/__pycache__/operations.cpython-310.pyc,, -accelerate/utils/__pycache__/other.cpython-310.pyc,, -accelerate/utils/__pycache__/random.cpython-310.pyc,, -accelerate/utils/__pycache__/rich.cpython-310.pyc,, -accelerate/utils/__pycache__/torch_xla.cpython-310.pyc,, -accelerate/utils/__pycache__/tqdm.cpython-310.pyc,, -accelerate/utils/__pycache__/transformer_engine.cpython-310.pyc,, -accelerate/utils/__pycache__/versions.cpython-310.pyc,, -accelerate/utils/ao.py,sha256=yfi84ywNpXsKZ6di4jbmtd6bIvRi9bmygr3giN92_Pg,4777 -accelerate/utils/bnb.py,sha256=PdgdXwjhQ4Sg9orj4hlUiMQlPxSsdsiNfFVq4yBtl-g,20651 -accelerate/utils/constants.py,sha256=A5KmcIry76B9p4Fa8CuI5ugv7bVcyTc5g4IBU1cvky0,3516 -accelerate/utils/dataclasses.py,sha256=REC089NzoyU3TfL-Tx5xYgzY6Ks4MLClMUWQVgzcGBQ,134397 -accelerate/utils/deepspeed.py,sha256=gt-jwOpOtRs6N4sohuSO2I3x1qLoO6uYYjnPrKp6Zk4,14796 -accelerate/utils/environment.py,sha256=h0zacbBkAp9szltTf5-aTr5NcbVsQp7wl6DFWp8XNuI,15257 -accelerate/utils/fsdp_utils.py,sha256=D0RIBcvdWLdnpha0HS0L484GEYKSfDYCWvtBul737zg,36481 -accelerate/utils/imports.py,sha256=7_8VRwewcD2O9MlpQjR7F1rZvZs8P5b2rr9U0eD1-zU,16715 -accelerate/utils/launch.py,sha256=8kqF22yFrWVsuaEBsDGlpSSt_FR_8SBhVy0mtv4uUhQ,33421 -accelerate/utils/megatron_lm.py,sha256=ZMxcYAdGnMWla2PtB04ctb9mhMfKUP8qLYusGNjyScU,58054 -accelerate/utils/memory.py,sha256=N7QGA6wdx8Xw7F0V-0TFshJLiXWtRnoXgezNZesvinw,7213 -accelerate/utils/modeling.py,sha256=g4OsRSSnJQMJ7q-CnUg2hUWQ2M0-0reeevD_CX0rdKQ,95847 -accelerate/utils/offload.py,sha256=VFaL8oSJzqZ_47VuUQ69xZi9bF2heRSFoOSnnOxbGXc,7825 -accelerate/utils/operations.py,sha256=84SpXwKsbNHHePboL9gUqK6bA5vo0kyH_gLAdwnzE_s,31266 -accelerate/utils/other.py,sha256=fAUxeRbhs3y3827xRpqVWqgCPjE1Qt4Vxf1NMdNiWak,20283 -accelerate/utils/random.py,sha256=Xv_ZJm9eaC2Q7rgZy9OpOunKuTingMiDQCH00qhNVxE,6220 -accelerate/utils/rich.py,sha256=8JZX_uGMQX-BufdXxJpdne7BWd1KyLHSgbiGxrDMYr8,847 -accelerate/utils/torch_xla.py,sha256=Pq1tuqN0X_pWDVza6YgjfO45uoJdoRVRForLeLQzFus,1908 -accelerate/utils/tqdm.py,sha256=k8e9JnieTEQHCCNBaiBys7hPxWlEbyRASdIma-qy_X8,1657 -accelerate/utils/transformer_engine.py,sha256=PENq1zhDzENnFFfV4hLzhZiiLo28tQFvmuhpQPlc3S0,6594 -accelerate/utils/versions.py,sha256=UgmcbjBm--6CIx1ZamSAMjAK_B_2l48LbeaNygqej8M,2149 diff --git a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/REQUESTED deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/WHEEL b/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/WHEEL deleted file mode 100644 index 370ec5f1aa3524bfd31a19fd21e907335f48a7f6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/WHEEL +++ /dev/null @@ -1,5 +0,0 @@ -Wheel-Version: 1.0 -Generator: setuptools (78.1.1) -Root-Is-Purelib: true -Tag: py3-none-any - diff --git a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/entry_points.txt b/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/entry_points.txt deleted file mode 100644 index 8b9bf6b798b250a47a3febdf0e32c88507fbf86d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/entry_points.txt +++ /dev/null @@ -1,6 +0,0 @@ -[console_scripts] -accelerate = accelerate.commands.accelerate_cli:main -accelerate-config = accelerate.commands.config:main -accelerate-estimate-memory = accelerate.commands.estimate:main -accelerate-launch = accelerate.commands.launch:main -accelerate-merge-weights = accelerate.commands.merge:main diff --git a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/licenses/LICENSE b/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/licenses/LICENSE deleted file mode 100644 index 261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/licenses/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - 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 - 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 - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - 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. diff --git a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/top_level.txt deleted file mode 100644 index a9368375be0e0e13fdad0eea4b92541bd9e1f594..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate-1.9.0.dist-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -accelerate diff --git a/venv/lib/python3.10/site-packages/accelerate/__init__.py b/venv/lib/python3.10/site-packages/accelerate/__init__.py deleted file mode 100644 index 3f3e15e6578f524482415d54567bee4ec849d297..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/__init__.py +++ /dev/null @@ -1,50 +0,0 @@ -# Copyright 2020 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. -__version__ = "1.9.0" - -from .accelerator import Accelerator -from .big_modeling import ( - cpu_offload, - cpu_offload_with_hook, - disk_offload, - dispatch_model, - init_empty_weights, - init_on_device, - load_checkpoint_and_dispatch, -) -from .data_loader import skip_first_batches -from .inference import prepare_pippy -from .launchers import debug_launcher, notebook_launcher -from .state import PartialState -from .utils import ( - AutocastKwargs, - DataLoaderConfiguration, - DDPCommunicationHookType, - DeepSpeedPlugin, - DistributedDataParallelKwargs, - DistributedType, - FullyShardedDataParallelPlugin, - GradScalerKwargs, - InitProcessGroupKwargs, - ProfileKwargs, - find_executable_batch_size, - infer_auto_device_map, - is_rich_available, - load_checkpoint_in_model, - synchronize_rng_states, -) - - -if is_rich_available(): - from .utils import rich diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 0f4e8d36800c31275572a702f389595b29645384..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/accelerator.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/accelerator.cpython-310.pyc deleted file mode 100644 index ca9c80e66833e667c4bdf4ec693a4052fd3e55b1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/__pycache__/accelerator.cpython-310.pyc +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c005c697ef0add5080c18e222594cb3e894ade8e7a6e30d27b2365127f090b9b -size 126436 diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/big_modeling.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/big_modeling.cpython-310.pyc deleted file mode 100644 index 415cfb4fa7ad15c1087a1a66282dcff295558d6d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/__pycache__/big_modeling.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/checkpointing.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/checkpointing.cpython-310.pyc deleted file mode 100644 index 44d678348ddd4d3f52f6add5e221408b554b8317..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/__pycache__/checkpointing.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/data_loader.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/data_loader.cpython-310.pyc deleted file mode 100644 index af179aea28a58ed9503190b474d7f28cc31b6658..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/__pycache__/data_loader.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/hooks.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/hooks.cpython-310.pyc deleted file mode 100644 index f83422eec8e247230f1372b06d7f21a070dfba9e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/__pycache__/hooks.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/inference.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/inference.cpython-310.pyc deleted file mode 100644 index daf88a50cf84199c049f8a5ab11dfe71b616e822..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/__pycache__/inference.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/launchers.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/launchers.cpython-310.pyc deleted file mode 100644 index 3b7b527f9b90cf7a9ba853df3dcbda3ae36f1f04..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/__pycache__/launchers.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/local_sgd.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/local_sgd.cpython-310.pyc deleted file mode 100644 index a37d25efbc7117a99b857c57d20c4e680297f8c3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/__pycache__/local_sgd.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/logging.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/logging.cpython-310.pyc deleted file mode 100644 index 692de25352aa030fc0c2b8592ef1e9c176a07a96..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/__pycache__/logging.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/memory_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/memory_utils.cpython-310.pyc deleted file mode 100644 index 6244f6b7565c9965c5ffba90d8a44799cd7a3915..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/__pycache__/memory_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/optimizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/optimizer.cpython-310.pyc deleted file mode 100644 index 02063e049071461e8313e69378384a6ac7c00112..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/__pycache__/optimizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/scheduler.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/scheduler.cpython-310.pyc deleted file mode 100644 index 328965736ef97e7592f4ce15da327cbba997753a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/__pycache__/scheduler.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/state.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/state.cpython-310.pyc deleted file mode 100644 index 872206cc59d156e2db26ea8721d6c21f9999ada3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/__pycache__/state.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/__pycache__/tracking.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/__pycache__/tracking.cpython-310.pyc deleted file mode 100644 index 741056381c5553eacbc67213b052118075c1d190..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/__pycache__/tracking.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/accelerator.py b/venv/lib/python3.10/site-packages/accelerate/accelerator.py deleted file mode 100644 index 9297426e1fbacd80e670b16d2be6eab8813034ca..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/accelerator.py +++ /dev/null @@ -1,4012 +0,0 @@ -# Copyright 2021 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 contextlib -import functools -import json -import math -import os -import re -import shutil -import sys -import warnings -from collections import OrderedDict -from contextlib import contextmanager -from functools import partial -from types import MethodType -from typing import Any, Callable, Union - -import torch -import torch.utils.hooks as hooks -from huggingface_hub import split_torch_state_dict_into_shards - -from accelerate.utils.dataclasses import FP8BackendType - -from .checkpointing import load_accelerator_state, load_custom_state, save_accelerator_state, save_custom_state -from .data_loader import DataLoaderDispatcher, prepare_data_loader, skip_first_batches -from .logging import get_logger -from .optimizer import AcceleratedOptimizer -from .scheduler import AcceleratedScheduler -from .state import AcceleratorState, GradientState, PartialState -from .tracking import LOGGER_TYPE_TO_CLASS, GeneralTracker, filter_trackers -from .utils import ( - MODEL_NAME, - SAFE_WEIGHTS_INDEX_NAME, - SAFE_WEIGHTS_NAME, - SAFE_WEIGHTS_PATTERN_NAME, - WEIGHTS_INDEX_NAME, - WEIGHTS_NAME, - WEIGHTS_PATTERN_NAME, - AORecipeKwargs, - AutocastKwargs, - DataLoaderConfiguration, - DeepSpeedPlugin, - DistributedDataParallelKwargs, - DistributedType, - DynamoBackend, - FP8RecipeKwargs, - FullyShardedDataParallelPlugin, - GradientAccumulationPlugin, - GradScalerKwargs, - InitProcessGroupKwargs, - KwargsHandler, - LoggerType, - MegatronLMPlugin, - MSAMPRecipeKwargs, - PrecisionType, - ProfileKwargs, - ProjectConfiguration, - RNGType, - TERecipeKwargs, - TorchDynamoPlugin, - TorchTensorParallelPlugin, - apply_fp8_autowrap, - check_os_kernel, - clean_state_dict_for_safetensors, - compare_versions, - convert_model, - convert_model_to_fp8_ao, - convert_outputs_to_fp32, - ensure_weights_retied, - extract_model_from_parallel, - fsdp2_apply_ac, - fsdp2_canonicalize_names, - fsdp2_prepare_model, - fsdp2_switch_optimizer_parameters, - gather, - gather_object, - get_fsdp2_grad_scaler, - get_grad_scaler, - get_mixed_precision_context_manager, - get_pretty_name, - has_offloaded_params, - is_bf16_available, - is_bitsandbytes_multi_backend_available, - is_deepspeed_available, - is_ipex_available, - is_lomo_available, - is_megatron_lm_available, - is_mlu_available, - is_msamp_available, - is_musa_available, - is_npu_available, - is_torch_version, - is_torch_xla_available, - is_torchao_available, - is_transformer_engine_available, - is_xpu_available, - load_fsdp_model, - load_fsdp_optimizer, - model_has_dtensor, - pad_across_processes, - parse_choice_from_env, - recursively_apply, - reduce, - release_memory, - save, - save_fsdp_model, - save_fsdp_optimizer, - wait_for_everyone, -) -from .utils.constants import ( - BETA_TP_AVAILABLE_PYTORCH_VERSION, - BETA_TP_AVAILABLE_TRANSFORMERS_VERSION, - FSDP2_PYTORCH_VERSION, - FSDP_PYTORCH_VERSION, - PROFILE_PATTERN_NAME, - SCALER_NAME, -) -from .utils.modeling import get_state_dict_offloaded_model -from .utils.other import compile_regions, compile_regions_deepspeed, is_compiled_module - - -if is_deepspeed_available(): - from .utils import ( - DeepSpeedEngineWrapper, - DeepSpeedOptimizerWrapper, - DeepSpeedSchedulerWrapper, - DummyOptim, - DummyScheduler, - map_pytorch_optim_to_deepspeed, - ) - -if is_megatron_lm_available(): - from .utils import ( - MegatronEngine, - MegatronLMDummyDataLoader, - MegatronLMDummyScheduler, - MegatronLMOptimizerWrapper, - MegatronLMSchedulerWrapper, - megatron_lm_initialize, - megatron_lm_prepare_data_loader, - megatron_lm_prepare_model_optimizer_scheduler, - ) - -from torch.distributed.algorithms.join import Join - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - import torch_xla.distributed.xla_multiprocessing as xmp - - -if is_npu_available(check_device=False): - import torch_npu # noqa: F401 - - -try: - from torch.optim.lr_scheduler import LRScheduler -except ImportError: - from torch.optim.lr_scheduler import _LRScheduler as LRScheduler - -logger = get_logger(__name__) - -# Sentinel values for defaults -_split_batches = object() -_dispatch_batches = object() -_even_batches = object() -_use_seedable_sampler = object() - - -class Accelerator: - """ - Creates an instance of an accelerator for distributed training or mixed precision training. - - Args: - device_placement (`bool`, *optional*, defaults to `True`): - Whether or not the accelerator should put objects on device (tensors yielded by the dataloader, model, - etc...). - mixed_precision (`str`, *optional*): - Whether or not to use mixed precision training. Choose from 'no','fp16','bf16' or 'fp8'. Will default to - the value in the environment variable `ACCELERATE_MIXED_PRECISION`, which will use the default value in the - accelerate config of the current system or the flag passed with the `accelerate.launch` command. 'fp8' - requires the installation of transformers-engine. - gradient_accumulation_steps (`int`, *optional*, default to 1): - The number of steps that should pass before gradients are accumulated. A number > 1 should be combined with - `Accelerator.accumulate`. If not passed, will default to the value in the environment variable - `ACCELERATE_GRADIENT_ACCUMULATION_STEPS`. Can also be configured through a `GradientAccumulationPlugin`. - cpu (`bool`, *optional*): - Whether or not to force the script to execute on CPU. Will ignore GPU available if set to `True` and force - the execution on one process only. - dataloader_config (`DataLoaderConfiguration`, *optional*): - A configuration for how the dataloaders should be handled in distributed scenarios. - deepspeed_plugin ([`~utils.DeepSpeedPlugin`] or dict of `str`: [`~utils.DeepSpeedPlugin`], *optional*): - Tweak your DeepSpeed related args using this argument. This argument is optional and can be configured - directly using *accelerate config*. If using multiple plugins, use the configured `key` property of each - plugin to access them from `accelerator.state.get_deepspeed_plugin(key)`. Alias for `deepspeed_plugins`. - fsdp_plugin ([`~utils.FullyShardedDataParallelPlugin`], *optional*): - Tweak your FSDP related args using this argument. This argument is optional and can be configured directly - using *accelerate config* - torch_tp_plugin ([`~utils.TorchTensorParallelPlugin`], *optional*): - Tweak your torch tensor parallel. This argument is optional and can be configured directly using - *accelerate config* - megatron_lm_plugin ([`~utils.MegatronLMPlugin`], *optional*): - Tweak your MegatronLM related args using this argument. This argument is optional and can be configured - directly using *accelerate config* - rng_types (list of `str` or [`~utils.RNGType`]): - The list of random number generators to synchronize at the beginning of each iteration in your prepared - dataloaders. Should be one or several of: - - - `"torch"`: the base torch random number generator - - `"cuda"`: the CUDA random number generator (GPU only) - - `"xla"`: the XLA random number generator (TPU only) - - `"generator"`: the `torch.Generator` of the sampler (or batch sampler if there is no sampler in your - dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type. - - Will default to `["torch"]` for PyTorch versions <=1.5.1 and `["generator"]` for PyTorch versions >= 1.6. - log_with (list of `str`, [`~utils.LoggerType`] or [`~tracking.GeneralTracker`], *optional*): - A list of loggers to be setup for experiment tracking. Should be one or several of: - - - `"all"` - - `"tensorboard"` - - `"wandb"` - - `"trackio"` - - `"aim"` - - `"comet_ml"` - - `"mlflow"` - - `"dvclive"` - - `"swanlab"` - If `"all"` is selected, will pick up all available trackers in the environment and initialize them. Can - also accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`. - project_config ([`~utils.ProjectConfiguration`], *optional*): - A configuration for how saving the state can be handled. - project_dir (`str`, `os.PathLike`, *optional*): - A path to a directory for storing data such as logs of locally-compatible loggers and potentially saved - checkpoints. - step_scheduler_with_optimizer (`bool`, *optional*, defaults to `True`): - Set `True` if the learning rate scheduler is stepped at the same time as the optimizer, `False` if only - done under certain circumstances (at the end of each epoch, for instance). - kwargs_handlers (list of [`~utils.KwargsHandler`], *optional*) - A list of [`~utils.KwargsHandler`] to customize how the objects related to distributed training, profiling - or mixed precision are created. See [kwargs](kwargs) for more information. - dynamo_backend (`str` or [`~utils.DynamoBackend`], *optional*, defaults to `"no"`): - Set to one of the possible dynamo backends to optimize your training with torch dynamo. - dynamo_plugin ([`~utils.TorchDynamoPlugin`], *optional*): - A configuration for how torch dynamo should be handled, if more tweaking than just the `backend` or `mode` - is needed. - gradient_accumulation_plugin ([`~utils.GradientAccumulationPlugin`], *optional*): - A configuration for how gradient accumulation should be handled, if more tweaking than just the - `gradient_accumulation_steps` is needed. - - **Available attributes:** - - - **device** (`torch.device`) -- The device to use. - - **distributed_type** ([`~utils.DistributedType`]) -- The distributed training configuration. - - **local_process_index** (`int`) -- The process index on the current machine. - - **mixed_precision** (`str`) -- The configured mixed precision mode. - - **num_processes** (`int`) -- The total number of processes used for training. - - **optimizer_step_was_skipped** (`bool`) -- Whether or not the optimizer update was skipped (because of - gradient overflow in mixed precision), in which - case the learning rate should not be changed. - - **process_index** (`int`) -- The overall index of the current process among all processes. - - **state** ([`~state.AcceleratorState`]) -- The distributed setup state. - - **sync_gradients** (`bool`) -- Whether the gradients are currently being synced across all processes. - - **use_distributed** (`bool`) -- Whether the current configuration is for distributed training. - """ - - def __init__( - self, - device_placement: bool = True, - split_batches: bool = _split_batches, - mixed_precision: PrecisionType | str | None = None, - gradient_accumulation_steps: int = 1, - cpu: bool = False, - dataloader_config: DataLoaderConfiguration | None = None, - deepspeed_plugin: DeepSpeedPlugin | dict[str, DeepSpeedPlugin] | None = None, - fsdp_plugin: FullyShardedDataParallelPlugin | None = None, - torch_tp_plugin: TorchTensorParallelPlugin | None = None, - megatron_lm_plugin: MegatronLMPlugin | None = None, - rng_types: list[str | RNGType] | None = None, - log_with: str | LoggerType | GeneralTracker | list[str | LoggerType | GeneralTracker] | None = None, - project_dir: str | os.PathLike | None = None, - project_config: ProjectConfiguration | None = None, - gradient_accumulation_plugin: GradientAccumulationPlugin | None = None, - step_scheduler_with_optimizer: bool = True, - kwargs_handlers: list[KwargsHandler] | None = None, - dynamo_backend: DynamoBackend | str | None = None, - dynamo_plugin: TorchDynamoPlugin | None = None, - deepspeed_plugins: DeepSpeedPlugin | dict[str, DeepSpeedPlugin] | None = None, - ): - self.trackers = [] - if project_config is not None: - self.project_configuration = project_config - else: - self.project_configuration = ProjectConfiguration(project_dir=project_dir) - if project_dir is not None and self.project_dir is None: - self.project_configuration.set_directories(project_dir) - - if mixed_precision is not None: - mixed_precision = str(mixed_precision) - if mixed_precision not in PrecisionType: - raise ValueError( - f"Unknown mixed_precision mode: {mixed_precision}. Choose between {PrecisionType.list()}" - ) - - if dynamo_plugin is not None and dynamo_backend is not None: - raise ValueError("You cannot pass in both `dynamo_plugin` and `dynamo_backend`, please only pass in one.") - if dynamo_backend is not None: - dynamo_plugin = TorchDynamoPlugin(backend=dynamo_backend) - elif dynamo_plugin is None: - dynamo_plugin = TorchDynamoPlugin() - - if deepspeed_plugins is not None and deepspeed_plugin is not None: - raise ValueError("You cannot pass in both `deepspeed_plugins` and `deepspeed_plugin`.") - elif deepspeed_plugin is not None: - deepspeed_plugins = deepspeed_plugin - - if deepspeed_plugins is None: - # First check if we're creating another `Accelerator` w/o setting `deepspeed_plugin` - if AcceleratorState._shared_state != {} and AcceleratorState().distributed_type == DistributedType.DEEPSPEED: - deepspeed_plugins = AcceleratorState().deepspeed_plugins - else: - # init from env variables - deepspeed_plugins = ( - DeepSpeedPlugin() if os.environ.get("ACCELERATE_USE_DEEPSPEED", "false") == "true" else None - ) - else: - # If we're creating a second `Accelerator`, users shouldn't be passing in a `deepspeed_plugin` - if ( - AcceleratorState._shared_state != {} - and AcceleratorState().distributed_type == DistributedType.DEEPSPEED - and AcceleratorState().deepspeed_plugins is not None - ): - raise NotImplementedError( - "You cannot pass in a `deepspeed_plugin` when creating a second `Accelerator`. " - "Please make sure the first `Accelerator` is initialized with all the plugins you want to use." - ) - if isinstance(deepspeed_plugins, dict): - for plugin in deepspeed_plugins.values(): - if not isinstance(plugin, DeepSpeedPlugin): - raise TypeError("`deepspeed_plugin` must be a DeepSpeedPlugin object.") - - if deepspeed_plugins is not None: - os.environ["ACCELERATE_USE_DEEPSPEED"] = "true" # use DeepSpeed if plugin is provided - if not is_deepspeed_available(): - raise ImportError("DeepSpeed is not installed => run `pip install deepspeed` or build it from source.") - if is_mlu_available(): - if compare_versions("deepspeed", "<", "0.15.2"): - raise ImportError("DeepSpeed MLU version must be >= 0.15.2. Please update DeepSpeed.") - elif is_musa_available(): - if compare_versions("deepspeed", "<", "0.14.3"): - raise ImportError("DeepSpeed MUSA version must be >= 0.14.3. Please update DeepSpeed.") - elif compare_versions("deepspeed", "<", "0.9.3"): - raise ImportError("DeepSpeed version must be >= 0.9.3. Please update DeepSpeed.") - - self.deepspeed_engine_wrapped = None - - if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true" or isinstance( - fsdp_plugin, FullyShardedDataParallelPlugin - ): - if not is_torch_version(">=", FSDP_PYTORCH_VERSION): - raise ValueError(f"FSDP requires PyTorch >= {FSDP_PYTORCH_VERSION}") - - if isinstance(torch_tp_plugin, TorchTensorParallelPlugin): - if not is_torch_version(">=", BETA_TP_AVAILABLE_PYTORCH_VERSION): - raise ValueError(f"TP requires PyTorch >= {BETA_TP_AVAILABLE_PYTORCH_VERSION}") - - if not compare_versions("transformers", ">=", BETA_TP_AVAILABLE_TRANSFORMERS_VERSION): - raise ValueError(f"TP requires transformers >= {BETA_TP_AVAILABLE_TRANSFORMERS_VERSION}") - - if fsdp_plugin is None: # init from env variables - fsdp_plugin = ( - FullyShardedDataParallelPlugin() if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true" else None - ) - else: - if not isinstance(fsdp_plugin, FullyShardedDataParallelPlugin): - raise TypeError("`fsdp_plugin` must be a FullyShardedDataParallelPlugin object.") - os.environ["ACCELERATE_USE_FSDP"] = "true" # use FSDP if plugin is provided - - if fsdp_plugin is not None and fsdp_plugin.fsdp_version == 2: - if not is_torch_version(">=", FSDP2_PYTORCH_VERSION): - raise ImportError(f"FSDP2 requires PyTorch >= {FSDP2_PYTORCH_VERSION}") - - if torch_tp_plugin is not None and not isinstance(torch_tp_plugin, TorchTensorParallelPlugin): - raise TypeError("`torch_tp_plugin` must be a TorchTensorParallelPlugin object.") - - if megatron_lm_plugin is None: # init from env variables - megatron_lm_plugin = ( - MegatronLMPlugin() if os.environ.get("ACCELERATE_USE_MEGATRON_LM", "false") == "true" else None - ) - else: - if not isinstance(megatron_lm_plugin, MegatronLMPlugin): - raise TypeError("`megatron_lm_plugin` must be a MegatronLMPlugin object.") - os.environ["ACCELERATE_USE_MEGATRON_LM"] = "true" # use MegatronLM if plugin is provided - - if megatron_lm_plugin: - if not is_megatron_lm_available(): - raise ImportError("Megatron is not installed. please build it from source.") - - # Kwargs handlers - self.ddp_handler = None - self.scaler_handler = None - self.init_handler = None - self.fp8_recipe_handler = None - self.ao_recipe_handler = None - self.te_recipe_handler = None - self.msamp_recipe_handler = None - self.autocast_handler = None - self.profile_handler = None - self.has_lomo_optimizer = False - - found_handlers = set() - handler_class_to_attr = { - DistributedDataParallelKwargs: "ddp_handler", - GradScalerKwargs: "scaler_handler", - InitProcessGroupKwargs: "init_handler", - FP8RecipeKwargs: "fp8_recipe_handler", - AutocastKwargs: "autocast_handler", - ProfileKwargs: "profile_handler", - AORecipeKwargs: "ao_recipe_handler", - TERecipeKwargs: "te_recipe_handler", - MSAMPRecipeKwargs: "msamp_recipe_handler", - } - self.has_fp8_handler = False - if kwargs_handlers is not None: - for handler in kwargs_handlers: - assert isinstance(handler, KwargsHandler), ( - f"Unsupported kwargs handler passed: {handler}, must be one that inherits `accelerate.utils.KwargsHandler`." - ) - # Add the handler class to the set of found handlers - if handler.__class__ in found_handlers: - raise ValueError(f"You can only pass one {handler.__class__} in `kwargs_handlers`.") - found_handlers.add(handler.__class__) - handler_attr = handler_class_to_attr[handler.__class__] - setattr(self, handler_attr, handler) - if "recipe_handler" in handler_attr and not self.has_fp8_handler: - self.has_fp8_handler = True - - kwargs = self.init_handler.to_kwargs() if self.init_handler is not None else {} - self.state = AcceleratorState( - mixed_precision=mixed_precision, - cpu=cpu, - dynamo_plugin=dynamo_plugin, - deepspeed_plugin=deepspeed_plugins, - fsdp_plugin=fsdp_plugin, - torch_tp_plugin=torch_tp_plugin, - megatron_lm_plugin=megatron_lm_plugin, - _from_accelerator=True, - **kwargs, - ) - - self.fp8_enabled = self.state.mixed_precision == "fp8" or mixed_precision == "fp8" - - # Check for automatic FP8 recipe creation - if self.fp8_enabled and not self.has_fp8_handler: - if self.fp8_backend == FP8BackendType.AO: - self.ao_recipe_handler = AORecipeKwargs() - elif self.fp8_backend == FP8BackendType.TE: - self.te_recipe_handler = TERecipeKwargs() - elif self.fp8_backend == FP8BackendType.MSAMP: - self.msamp_recipe_handler = MSAMPRecipeKwargs() - elif self.fp8_backend == FP8BackendType.NO: - # Prioritize AO -> TE -> MSAMP - if is_torchao_available(): - logger.info("Found `torchao` installed, using it for FP8 training.") - self.ao_recipe_handler = AORecipeKwargs() - elif is_transformer_engine_available(): - logger.info("Found `transformer-engine` installed, using it for FP8 training.") - self.te_recipe_handler = TERecipeKwargs() - elif is_msamp_available(): - logger.info("Found `msamp` installed, using it for FP8 training.") - self.msamp_recipe_handler = MSAMPRecipeKwargs() - else: - raise ImportError( - "Tried to train with `fp8` and auto-detect backend, but no FP8-compatible backend was installed. " - "Valid backends are: `torchao`, `transformer-engine`, and `msamp`." - ) - self.has_fp8_handler = True - - self.delayed_fp8_autocast = False - if self.has_fp8_handler: - # We already check if FP8 is available during `self.state` - if not self.fp8_enabled and ( - self.distributed_type not in (DistributedType.FSDP, DistributedType.DEEPSPEED) - ): - raise ValueError("Passing in an FP8 configuration requires setting `mixed_precision='fp8'`.") - self.delayed_fp8_autocast = self.fp8_backend == "TE" and self.distributed_type in ( - DistributedType.MULTI_GPU, - DistributedType.FSDP, - ) - - # TODO: S1ro - this is probably gonna be a problem with other fp8 backends too - if ( - self.fp8_backend == FP8BackendType.AO - and self.state.distributed_type == DistributedType.FSDP - and self.state.fsdp_plugin.cpu_ram_efficient_loading - ): - raise ValueError( - "torchao with FSDP2 and cpu_ram_efficient_loading is not supported, setting `cpu_ram_efficient_loading` to False will fix the issue and work as intended." - ) - - trackers = filter_trackers(log_with, self.logging_dir) - if len(trackers) < 1 and log_with is not None: - warnings.warn(f"`log_with={log_with}` was passed but no supported trackers are currently installed.") - self.log_with = trackers - - if ( - (mixed_precision != "bf16") - and getattr(self.state, "downcast_bfloat", False) - and (self.state.distributedType != DistributedType.XLA) - ): - raise ValueError("Can only use `downcast_bf16` when using `mixed_precision='bf16'` and on a TPU") - - if gradient_accumulation_plugin is not None: - if gradient_accumulation_steps != 1: - raise ValueError( - "You can only pass one of `gradient_accumulation_steps` and `gradient_accumulation_plugin`. Please only pass in the created `GradientAccumulationPlugin` object." - ) - else: - gradient_accumulation_steps = int( - parse_choice_from_env("ACCELERATE_GRADIENT_ACCUMULATION_STEPS", gradient_accumulation_steps) - ) - gradient_accumulation_plugin = GradientAccumulationPlugin(num_steps=gradient_accumulation_steps) - - # If using DeepSpeed, update gradient accumulation steps from the DeepSpeed plugin - self.gradient_state = GradientState( - gradient_accumulation_plugin=gradient_accumulation_plugin, - ) - - self.device_placement = device_placement - if dataloader_config is None: - dataloader_config = DataLoaderConfiguration() - self.dataloader_config = dataloader_config - self.step_scheduler_with_optimizer = step_scheduler_with_optimizer - - # Mixed precision attributes - self.scaler = None - self.native_amp = False - if ( - self.state.mixed_precision == "fp16" - and self.device.type != "cpu" - and self.distributed_type not in (DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM) - ): - self.native_amp = True - if self.device.type not in ( - "xpu", - "cuda", - "npu", - "xla", - "mlu", - "musa", - "hpu", - "sdaa", - ) or is_torch_xla_available(check_is_tpu=True): - raise ValueError(f"fp16 mixed precision requires a GPU (not {self.device.type!r}).") - kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {} - - # FSDP2 doesn't use ShardedGradScaler, don't want to modify `get_grad_scaler`, rather create a simple utility - if self.is_fsdp2: - self.scaler = get_fsdp2_grad_scaler(**kwargs) - else: - self.scaler = get_grad_scaler(self.distributed_type, **kwargs) - - elif self.state.mixed_precision == "bf16" and self.distributed_type not in ( - DistributedType.DEEPSPEED, - DistributedType.MEGATRON_LM, - ): - if self.device.type in ["cpu", "xpu", "hpu"]: - self.native_amp = True - else: - self.native_amp = is_bf16_available(True) - if mixed_precision == "bf16" and not self.native_amp and not is_torch_xla_available(): - raise ValueError("bf16 mixed precision requires PyTorch >= 1.10 and a supported device.") - - # for DeepSpeed, self.state.mixed_precision is always "bf16", - # see https://github.com/huggingface/accelerate/blob/main/src/accelerate/state.py#L968 and - # https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1263. - elif self.fp8_enabled: - # We always enable `native_amp` for FP8 - self.native_amp = True - if self.fp8_backend == FP8BackendType.MSAMP: - if self.distributed_type == DistributedType.FSDP: - raise NotImplementedError( - "`accelerate` + `MS-AMP` + `FSDP` is not supported at this time. " - "Please consider using deepspeed, which is supported." - ) - elif self.distributed_type != DistributedType.DEEPSPEED: - # MS-AMP requires `GradScaler` even with bf16 autocast w/ single GPU or DDP: - self.scaler = get_grad_scaler(**kwargs) - - # Start of internal step tracking - self.step = 0 - - # Internal references to the training objects - self._optimizers = [] - self._models = [] - self._schedulers = [] - self._dataloaders = [] - self._custom_objects = [] - - # Hooks - self._load_model_state_pre_hook = OrderedDict() - self._save_model_state_pre_hook = OrderedDict() - - # RNG Types - self.rng_types = rng_types - if self.rng_types is None: - self.rng_types = ["generator"] - - # Set a flag tensor for early stopping and other breakpoints - self.flag_tensor = None - - check_os_kernel() - - @property - def deepspeed_plugin(self): - """ - Returns the currently active DeepSpeedPlugin. - - If using multiple plugins, the first one will be the active one by default. Manually call - `accelerator.state.select_deepspeed_plugin(key)` to activate a different plugin. - - If deepspeed is not enabled, this will return `None`. - """ - return self.state.deepspeed_plugin - - @property - def use_distributed(self): - """ - Whether the Accelerator is configured for distributed training - """ - return self.state.use_distributed - - @property - def distributed_type(self): - return self.state.distributed_type - - @property - def num_processes(self): - return self.state.num_processes - - @property - def process_index(self): - return self.state.process_index - - @property - def local_process_index(self): - return self.state.local_process_index - - @property - def device(self): - return self.state.device - - @property - def split_batches(self): - return self.dataloader_config.split_batches - - @property - def dispatch_batches(self): - return self.dataloader_config.dispatch_batches - - @property - def even_batches(self): - return self.dataloader_config.even_batches - - @even_batches.setter - def even_batches(self, value: bool): - self.dataloader_config.even_batches = value - - @property - def use_seedable_sampler(self): - return self.dataloader_config.use_seedable_sampler - - @property - def non_blocking(self): - return self.dataloader_config.non_blocking - - @property - def use_stateful_dataloader(self): - if hasattr(self.dataloader_config, "use_stateful_dataloader"): - return self.dataloader_config.use_stateful_dataloader - return False - - @property - def project_dir(self): - return self.project_configuration.project_dir - - @property - def logging_dir(self): - return self.project_configuration.logging_dir - - @property - def save_iteration(self): - return self.project_configuration.iteration - - @property - def is_main_process(self): - """True for one process only.""" - return self.state.is_main_process - - @property - def is_local_main_process(self): - """True for one process per server.""" - return self.state.is_local_main_process - - @property - def is_last_process(self): - return self.process_index == self.num_processes - 1 - - @property - def mixed_precision(self): - return self.state.mixed_precision - - @property - def is_fsdp2(self): - return self.state.is_fsdp2 - - @contextmanager - def split_between_processes(self, inputs: list | tuple | dict | torch.Tensor, apply_padding: bool = False): - """ - Splits `input` between `self.num_processes` quickly and can be then used on that process. Useful when doing - distributed inference, such as with different prompts. - - Note that when using a `dict`, all keys need to have the same number of elements. - - Args: - inputs (`list`, `tuple`, `torch.Tensor`, or `dict` of `list`/`tuple`/`torch.Tensor`): - The input to split between processes. - apply_padding (`bool`, `optional`, defaults to `False`): - Whether to apply padding by repeating the last element of the input so that all processes have the same - number of elements. Useful when trying to perform actions such as `Accelerator.gather()` on the outputs - or passing in less inputs than there are processes. If so, just remember to drop the padded elements - afterwards. - - Example: - - ```python - # Assume there are two processes - from accelerate import Accelerator - - accelerator = Accelerator() - with accelerator.split_between_processes(["A", "B", "C"]) as inputs: - print(inputs) - # Process 0 - ["A", "B"] - # Process 1 - ["C"] - - with accelerator.split_between_processes(["A", "B", "C"], apply_padding=True) as inputs: - print(inputs) - # Process 0 - ["A", "B"] - # Process 1 - ["C", "C"] - ``` - """ - with PartialState().split_between_processes(inputs, apply_padding=apply_padding) as inputs: - yield inputs - - def on_main_process(self, function: Callable[..., Any] = None): - """ - A decorator that will run the decorated function on the main process only. Can also be called using the - `PartialState` class. - - Args: - function (`Callable`): The function to decorate. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - - - >>> @accelerator.on_main_process - ... def print_something(): - ... print("This will be printed by process 0 only.") - - - >>> print_something() - "This will be printed by process 0 only" - ``` - """ - # For times when the `Accelerator` object itself utilizes this decorator. - if function is None: - if "Accelerator." in self.__qualname__: - function = self - else: - raise ValueError( - "The `on_main_process` decorator must be called with a function on an instantiated `Accelerator` object." - ) - - def _inner(*args, **kwargs): - return PartialState().on_main_process(function)(*args, **kwargs) - - return _inner - - def on_local_main_process(self, function: Callable[..., Any] = None): - """ - A decorator that will run the decorated function on the local main process only. Can also be called using the - `PartialState` class. - - Args: - function (`Callable`): The function to decorate. - - Example: - ```python - # Assume we have 2 servers with 4 processes each. - from accelerate import Accelerator - - accelerator = Accelerator() - - - @accelerator.on_local_main_process - def print_something(): - print("This will be printed by process 0 only on each server.") - - - print_something() - # On server 1: - "This will be printed by process 0 only" - # On server 2: - "This will be printed by process 0 only" - ``` - """ - # For times when the `Accelerator` object itself utilizes this decorator. - if function is None: - if "Accelerator." in self.__qualname__: - function = self - else: - raise ValueError( - "The `on_local_main_process` decorator must be called with a function on an instantiated `Accelerator` object." - ) - - def _inner(*args, **kwargs): - return PartialState().on_local_main_process(function)(*args, **kwargs) - - return _inner - - def on_last_process(self, function: Callable[..., Any]): - """ - A decorator that will run the decorated function on the last process only. Can also be called using the - `PartialState` class. - - Args: - function (`Callable`): The function to decorate. - - Example: - ```python - # Assume we have 4 processes. - from accelerate import Accelerator - - accelerator = Accelerator() - - - @accelerator.on_last_process - def print_something(): - print(f"Printed on process {accelerator.process_index}") - - - print_something() - "Printed on process 3" - ``` - """ - # For times when the `Accelerator` object itself utilizes this decorator. - if function is None: - if "Accelerator." in self.__qualname__: - function = self - else: - raise ValueError( - "The `on_last_process` decorator must be called with a function on an instantiated `Accelerator` object." - ) - - def _inner(*args, **kwargs): - return PartialState().on_last_process(function)(*args, **kwargs) - - return _inner - - def on_process(self, function: Callable[..., Any] = None, process_index: int = None): - """ - A decorator that will run the decorated function on a given process index only. Can also be called using the - `PartialState` class. - - Args: - function (`Callable`, `optional`): - The function to decorate. - process_index (`int`, `optional`): - The index of the process on which to run the function. - - Example: - ```python - # Assume we have 4 processes. - from accelerate import Accelerator - - accelerator = Accelerator() - - - @accelerator.on_process(process_index=2) - def print_something(): - print(f"Printed on process {accelerator.process_index}") - - - print_something() - "Printed on process 2" - ``` - """ - # Initial construction of the decorator. - if (self is not None) and (process_index is not None) and (function is None): - return partial(self.on_process, process_index=process_index) - # For times when the `Accelerator` object itself utilizes this decorator. - if function is None: - if "Accelerator." in self.__qualname__: - function = self - else: - raise ValueError( - "The `on_main_process` decorator must be called with a function on an instantiated `Accelerator` object." - ) - - def _inner(*args, **kwargs): - return PartialState().on_process(function, process_index)(*args, **kwargs) - - return _inner - - def on_local_process(self, function: Callable[..., Any] = None, local_process_index: int = None): - """ - A decorator that will run the decorated function on a given local process index only. Can also be called using - the `PartialState` class. - - Args: - function (`Callable`, *optional*): - The function to decorate. - local_process_index (`int`, *optional*): - The index of the local process on which to run the function. - - Example: - ```python - # Assume we have 2 servers with 4 processes each. - from accelerate import Accelerator - - accelerator = Accelerator() - - - @accelerator.on_local_process(local_process_index=2) - def print_something(): - print(f"Printed on process {accelerator.local_process_index}") - - - print_something() - # On server 1: - "Printed on process 2" - # On server 2: - "Printed on process 2" - ``` - """ - # Initial construction of the decorator. - if (self is not None) and (local_process_index is not None) and (function is None): - return partial(self.on_local_process, local_process_index=local_process_index) - # For times when the `Accelerator` object itself utilizes this decorator. - if function is None: - if "Accelerator." in self.__qualname__: - function = self - else: - raise ValueError( - "The `on_main_process` decorator must be called with a function on an instantiated `Accelerator` object." - ) - - def _inner(*args, **kwargs): - return PartialState().on_local_process(function, local_process_index)(*args, **kwargs) - - return _inner - - @contextmanager - def main_process_first(self): - """ - Lets the main process go first inside a with block. - - The other processes will enter the with block after the main process exits. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> with accelerator.main_process_first(): - ... # This will be printed first by process 0 then in a seemingly - ... # random order by the other processes. - ... print(f"This will be printed by process {accelerator.process_index}") - ``` - """ - with self.state.main_process_first(): - yield - - @contextmanager - def local_main_process_first(self): - """ - Lets the local main process go inside a with block. - - The other processes will enter the with block after the main process exits. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> with accelerator.local_main_process_first(): - ... # This will be printed first by local process 0 then in a seemingly - ... # random order by the other processes. - ... print(f"This will be printed by process {accelerator.local_process_index}") - ``` - """ - with self.state.local_main_process_first(): - yield - - @contextmanager - def no_sync(self, model): - """ - A context manager to disable gradient synchronizations across DDP processes by calling - `torch.nn.parallel.DistributedDataParallel.no_sync`. - - If `model` is not in DDP, this context manager does nothing - - Args: - model (`torch.nn.Module`): - PyTorch Module that was prepared with `Accelerator.prepare` - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> dataloader, model, optimizer = accelerator.prepare(dataloader, model, optimizer) - >>> input_a = next(iter(dataloader)) - >>> input_b = next(iter(dataloader)) - - >>> with accelerator.no_sync(): - ... outputs = model(input_a) - ... loss = loss_func(outputs) - ... accelerator.backward(loss) - ... # No synchronization across processes, only accumulate gradients - >>> outputs = model(input_b) - >>> accelerator.backward(loss) - >>> # Synchronization across all processes - >>> optimizer.step() - >>> optimizer.zero_grad() - ``` - """ - context = contextlib.nullcontext - if self.use_distributed: - if self.distributed_type != DistributedType.DEEPSPEED or self.state.deepspeed_plugin.zero_stage < 2: - context = getattr(model, "no_sync", context) - - with context(): - yield - - @staticmethod - @contextmanager - def trigger_sync_in_backward(model): - """Trigger the sync of the gradients in the next backward pass of the model after multiple forward passes under - `Accelerator.no_sync` (only applicable in multi-GPU scenarios). - - If the script is not launched in distributed mode, this context manager does nothing. - - Args: - model (`torch.nn.Module`): - The model for which to trigger the gradient synchronization. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> dataloader, model, optimizer = accelerator.prepare(dataloader, model, optimizer) - - >>> with accelerator.no_sync(): - ... loss_a = loss_func(model(input_a)) # first forward pass - ... loss_b = loss_func(model(input_b)) # second forward pass - >>> accelerator.backward(loss_a) # No synchronization across processes, only accumulate gradients - >>> with accelerator.trigger_sync_in_backward(model): - ... accelerator.backward(loss_b) # Synchronization across all processes - >>> optimizer.step() - >>> optimizer.zero_grad() - ``` - """ - if not isinstance(model, torch.nn.parallel.DistributedDataParallel): - yield - return - - old_require_backward_grad_sync = model.require_backward_grad_sync - old_require_forward_param_sync = model.require_forward_param_sync - - # EXPERIMENTAL: This will force grad sync during `backward()`, but it is unknown if it breaks other DDP features. - # https://github.com/pytorch/pytorch/blob/e1502c0cdbfd17548c612f25d5a65b1e4b86224d/torch/nn/parallel/distributed.py#L1453-L1466 - model.require_backward_grad_sync = True - model.require_forward_param_sync = True - # https://github.com/pytorch/pytorch/blob/e1502c0cdbfd17548c612f25d5a65b1e4b86224d/torch/csrc/distributed/c10d/reducer.cpp#L1371-L1402 - model.reducer.prepare_for_backward([]) - try: - yield - finally: - model.require_backward_grad_sync = old_require_backward_grad_sync - model.require_forward_param_sync = old_require_forward_param_sync - - def _do_sync(self): - "Sets the right `sync_gradients` context and either resets or increases `self.step`" - if self.gradient_state.sync_with_dataloader and self.gradient_state.end_of_dataloader: - self.step = 0 - self.gradient_state._set_sync_gradients(True) - else: - self.step += 1 - self.gradient_state._set_sync_gradients((self.step % self.gradient_state.num_steps) == 0) - - @property - def sync_gradients(self): - return self.gradient_state.sync_gradients - - @sync_gradients.setter - def sync_gradients(self, sync_gradients): - self.gradient_state.sync_gradients = sync_gradients - - @property - def gradient_accumulation_steps(self): - return self.gradient_state.num_steps - - @gradient_accumulation_steps.setter - def gradient_accumulation_steps(self, gradient_accumulation_steps): - self.gradient_state.plugin_kwargs.update({"num_steps": gradient_accumulation_steps}) - - @contextmanager - def accumulate(self, *models): - """ - A context manager that will lightly wrap around and perform gradient accumulation automatically - - Args: - *models (list of `torch.nn.Module`): - PyTorch Modules that were prepared with `Accelerator.prepare`. Models passed to `accumulate()` will - skip gradient syncing during backward pass in distributed training - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator(gradient_accumulation_steps=1) - >>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler) - - >>> for input, output in dataloader: - ... with accelerator.accumulate(model): - ... outputs = model(input) - ... loss = loss_func(outputs) - ... loss.backward() - ... optimizer.step() - ... scheduler.step() - ... optimizer.zero_grad() - ``` - """ - self._do_sync() - - allow_gradient_sync = ( - self.sync_gradients # must sync if sync gradients need to complete an optimizer step - or ( - # the no_sync context stops the gradients from reducing during distributed training - # bringing speedup (potentially at some costs). Here, no_sync can be prevented - # by setting sync_each_batch = True. - self.use_distributed # only relevant in distributed settings - and self.gradient_state.plugin_kwargs.get("sync_each_batch", False) - ) - ) - with contextlib.ExitStack() as cm_stack: - for m in models: - cm_stack.enter_context(contextlib.nullcontext() if allow_gradient_sync else self.no_sync(m)) - yield - - @contextmanager - def join_uneven_inputs(self, joinables, even_batches=None): - """ - A context manager that facilitates distributed training or evaluation on uneven inputs, which acts as a wrapper - around `torch.distributed.algorithms.join`. This is useful when the total batch size does not evenly divide the - length of the dataset. - - Args: - joinables (`list[torch.distributed.algorithms.Joinable]`): - A list of models or optimizers that subclass `torch.distributed.algorithms.Joinable`. Most commonly, a - PyTorch Module that was prepared with `Accelerator.prepare` for DistributedDataParallel training. - even_batches (`bool`, *optional*) - If set, this will override the value of `even_batches` set in the `Accelerator`. If it is not provided, - the default `Accelerator` value wil be used. - - - - `join_uneven_inputs` is only supported for Distributed Data Parallel training on multiple GPUs. For any other - configuration, this method will have no effect. - - - - - - Overidding `even_batches` will not affect iterable-style data loaders. - - - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator(even_batches=True) - >>> ddp_model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader) - - >>> with accelerator.join_uneven_inputs([ddp_model], even_batches=False): - ... for input, output in dataloader: - ... outputs = model(input) - ... loss = loss_func(outputs) - ... loss.backward() - ... optimizer.step() - ... optimizer.zero_grad() - ``` - """ - if self.distributed_type in ( - DistributedType.MULTI_GPU, - DistributedType.MULTI_NPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_XPU, - DistributedType.MULTI_HPU, - ): - dl_even_batches_values = [] - - if even_batches is not None: - iterable_dl_seen = False - # override value in batch sampler for map-style datasets - for dl_idx, dl in enumerate(self._dataloaders): - if isinstance(dl, DataLoaderDispatcher): - iterable_dl_seen = True - continue - dl_even_batches_values.append((dl_idx, dl.batch_sampler.even_batches)) - dl.batch_sampler.even_batches = even_batches - - if iterable_dl_seen: - warnings.warn( - "Overridding even_batches is only supported for map-style datasets, yet some dataloaders given were iterable" - ) - else: - even_batches = self.even_batches - - enable_join = False if even_batches else True - try: - with Join(joinables, enable=enable_join, throw_on_early_termination=False): - yield - finally: - # reset any batch samplers that have been modified - for dl_idx, even_batches_value in dl_even_batches_values: - self._dataloaders[dl_idx].batch_sampler.even_batches = even_batches_value - else: - # Even when disabled, Join expects models to subclass Joinable, so skip entirely for single process runs - if self.distributed_type != DistributedType.NO: - warnings.warn( - "Joining uneven inputs is only supported for multi-GPU training, as a result `join_uneven_inputs` will have no effect." - ) - - with contextlib.nullcontext(joinables): - yield - - def print(self, *args, **kwargs): - """ - Drop in replacement of `print()` to only print once per server. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> accelerator.print("Hello world!") - ``` - """ - self.state.print(*args, **kwargs) - - def _prepare_one(self, obj, first_pass=False, device_placement=None): - # First pass of preparation: DataLoader, model, optimizer - if first_pass: - if isinstance(obj, torch.utils.data.DataLoader): - return self.prepare_data_loader(obj, device_placement=device_placement) - elif isinstance(obj, torch.nn.Module): - return self.prepare_model(obj, device_placement=device_placement) - elif isinstance(obj, torch.optim.Optimizer): - optimizer = self.prepare_optimizer(obj, device_placement=device_placement) - return optimizer - # Second pass of preparation: LR scheduler (which need the full list of optimizers) - elif isinstance(obj, LRScheduler): - scheduler = self.prepare_scheduler(obj) - return scheduler - # Return the unprocessed object if previous criteria was not met - return obj - - def prepare(self, *args, device_placement=None): - """ - Prepare all objects passed in `args` for distributed training and mixed precision, then return them in the same - order. - - Args: - *args (list of objects): - Any of the following type of objects: - - - `torch.utils.data.DataLoader`: PyTorch Dataloader - - `torch.nn.Module`: PyTorch Module - - `torch.optim.Optimizer`: PyTorch Optimizer - - `torch.optim.lr_scheduler.LRScheduler`: PyTorch LR Scheduler - - device_placement (`list[bool]`, *optional*): - Used to customize whether automatic device placement should be performed for each object passed. Needs - to be a list of the same length as `args`. Not compatible with DeepSpeed or FSDP. - - - - You don't need to prepare a model if you only use it for inference without any kind of mixed precision - - - - Examples: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> # Assume a model, optimizer, data_loader and scheduler are defined - >>> model, optimizer, data_loader, scheduler = accelerator.prepare(model, optimizer, data_loader, scheduler) - ``` - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> # Assume a model, optimizer, data_loader and scheduler are defined - >>> device_placement = [True, True, False, False] - >>> # Will place the first two items passed in automatically to the right device but not the last two. - >>> model, optimizer, data_loader, scheduler = accelerator.prepare( - ... model, optimizer, data_loader, scheduler, device_placement=device_placement - ... ) - ``` - """ - if device_placement is None: - device_placement = [None for _ in args] - elif self.distributed_type in (DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM): - raise ValueError("You can't customize device placements with DeepSpeed or Megatron-LM.") - elif len(device_placement) != len(args): - raise ValueError( - f"`device_placement` should be a list with {len(args)} elements (the number of objects passed)." - ) - - for obj in args: - # TODO: Look at enabling native TP training directly with a proper config - if ( - isinstance(obj, torch.nn.Module) - and self.verify_device_map(obj) - and self.distributed_type != DistributedType.NO - and os.environ.get("ACCELERATE_BYPASS_DEVICE_MAP", "false") != "true" - ): - raise ValueError( - "You can't train a model that has been loaded with `device_map='auto'` in any distributed mode." - " Please rerun your script specifying `--num_processes=1` or by launching with `python {{myscript.py}}`." - ) - - if self.distributed_type == DistributedType.DEEPSPEED: - model_count = 0 - for obj in args: - if isinstance(obj, torch.nn.Module): - model_count += 1 - if model_count > 1: - raise AssertionError( - "You can't use same `Accelerator()` instance with multiple models when using DeepSpeed" - ) - - # On TPUs, putting the model on the XLA device will create new parameters, so the corresponding optimizer will - # have parameters disconnected from the model (so no training :-( ). - # If the model and optimizer have parameters on different devices we raise an error. - if self.distributed_type == DistributedType.XLA: - model_device, optimizer_device = self._get_devices() - if model_device is not None and optimizer_device is not None and model_device != optimizer_device: - raise ValueError( - "The model and the optimizer parameters are not on the same device, which probably means you " - "created an optimizer around your model **before** putting on the device. Make sure the line " - "model.to(device) is before the optimizer creation in your script or remove it entirely and use " - "the flag default value for `device_placement` in your `Accelerator` to let it handle that " - "part for you." - ) - - if self.is_fsdp2: - model_count = 0 - optimizer_count = 0 - for i, obj in enumerate(args): - if isinstance(obj, torch.nn.Module): - model_count += 1 - elif isinstance(obj, torch.optim.Optimizer): - optimizer_count += 1 - - # This needs to be written as such, so that passing other objects other than models/optimizers doesn't raise an error - if (model_count < 1 and optimizer_count > 0) or (model_count > 0 and optimizer_count < 1): - raise ValueError( - "When using FSDP2, a model and optimizer must be passed together to `Accelerator.prepare()`" - " as the optimizer needs to have its parameters modified after the model is converted." - ) - if model_count > 1: - raise ValueError("Only one model is supported when using FSDP2") - - # If we're dealing with device placement, this deals with that by... - tpu_should_fix_optimizer = self.device_placement and self.distributed_type == DistributedType.XLA - - if tpu_should_fix_optimizer: - # 1. grabbing old model parameters - old_named_params = self._get_named_parameters(*args, drop_refs=False) - - if self.distributed_type in [DistributedType.MULTI_CPU, DistributedType.MULTI_XPU, DistributedType.NO]: - if ( - is_torch_version("<", "2.7.0") - and (self.device.type == "cpu" or self.device.type == "xpu") - and self.state.use_ipex - ): - logger.warning( - "You are using lower version of PyTorch(< 2.7.0) with ipex acceleration on Intel CPU or XPU, Intel has upstreamed most of the optimizations into stock PyTorch from 2.7.0, we enourage you to install the latest stock PyTorch and enjoy the out-of-experience on Intel CPU/XPU." - ) - args = self._prepare_ipex(*args) - if self.fp8_backend == FP8BackendType.TE: - args = self._prepare_te(*args) - elif self.fp8_backend == FP8BackendType.AO: - args = self._prepare_ao(*args) - if self.distributed_type == DistributedType.DEEPSPEED: - result = self._prepare_deepspeed(*args) - elif self.distributed_type == DistributedType.MEGATRON_LM: - result = self._prepare_megatron_lm(*args) - elif self.is_fsdp2: - result = self._prepare_fsdp2(*args) - else: - if self.fp8_backend == FP8BackendType.MSAMP: - args, device_placement = self._prepare_msamp(*args, device_placement=device_placement) - result = tuple( - self._prepare_one(obj, first_pass=True, device_placement=d) for obj, d in zip(args, device_placement) - ) - result = tuple(self._prepare_one(obj, device_placement=d) for obj, d in zip(result, device_placement)) - if tpu_should_fix_optimizer: - # 2. grabbing new model parameters - new_named_params = self._get_named_parameters(*result) - # 3. building a map from the first to the second - mapping = {p: new_named_params[n] for n, p in old_named_params.items()} - # 4. using that map to update the parameters of the optimizer - for obj in result: - if isinstance(obj, torch.optim.Optimizer): - obj._switch_parameters(mapping) - - for item in result: - if any( - item in container - for container in (self._dataloaders, self._models, self._optimizers, self._schedulers) - ): - item._is_accelerate_prepared = True - - return result if len(result) > 1 else result[0] - - def _prepare_fsdp2(self, *args): - # First pass: prepare everything except schedulers (and model, which is prepared separately below) - result = [ - self._prepare_one(obj, first_pass=True) if not isinstance(obj, torch.nn.Module) else obj for obj in args - ] - - # Second pass: prepare schedulers - result = [self._prepare_one(obj) if not isinstance(obj, torch.nn.Module) else obj for obj in result] - - # Prepare the model - model_index, model = None, None - for i, obj in enumerate(result): - if isinstance(obj, torch.nn.Module): - model_index, model = i, obj - - # Invariant: if we have a model, we also have an optimizer (checked in `prepare`) - if model_index is None: - return tuple(result) - - # Needs to be done first, to make sure AC + fully_shard will work as expected - self.state.fsdp_plugin.set_auto_wrap_policy(model) - - # Apply AC if needed - if self.state.fsdp_plugin.activation_checkpointing: - model = fsdp2_apply_ac(self, model) - - # Apply compile if needed, has to be *after* applying AC - # Copied from: `accelerator.prepare_model` ~ L1804 - if self.state.dynamo_plugin.backend != DynamoBackend.NO and not is_compiled_module(model): - if self.state.dynamo_plugin.use_regional_compilation: - model = compile_regions(model, **self.state.dynamo_plugin.to_kwargs()) - else: - model = torch.compile(model, **self.state.dynamo_plugin.to_kwargs()) - - # Get old params and canonicalize - we cannonicalize to have the mapping easy - old_named_params = fsdp2_canonicalize_names(self._get_named_parameters(*tuple(result), drop_refs=True)) - - # Swap the optimizer parameters with empty, so `fully_shard` after will not allocate too much memory - for obj in result: - if isinstance(obj, torch.optim.Optimizer): - for param_group in obj.param_groups: - for i, p in enumerate(param_group["params"]): - # We drop a reference to the original param here, so that _move_states_to_device triggers a reallocation - # We reassign the data_ptr to the original param, so that we preserve the mapping to the new ones - param_group["params"][i] = torch.empty_like(p) - param_group["params"][i].data_ptr = p.data_ptr() - - self._models.append(model) - - # Prepare everything FSDP2 related for the model (except AC) - model = fsdp2_prepare_model(self, model) - - # Remove the old model from the list - if len(self._models) > 1 and (self._models[-2] is self._models[-1]): - del self._models[-2] - - # Replace the old model with the new one (shouldn't be needed as everything should be in place) - result[model_index] = model - - # Get new params and canonicalize - new_named_params = fsdp2_canonicalize_names(self._get_named_parameters(*result)) - # Build a map from old to new params - mapping = {p: new_named_params[n] for n, p in old_named_params.items()} - # Update the optimizer parameters - for obj in result: - if isinstance(obj, torch.optim.Optimizer): - fsdp2_switch_optimizer_parameters(obj, mapping) - - return result - - def prepare_model(self, model: torch.nn.Module, device_placement: bool = None, evaluation_mode: bool = False): - """ - Prepares a PyTorch model for training in any distributed setup. It is recommended to use - [`Accelerator.prepare`] instead. - - Args: - model (`torch.nn.Module`): - A PyTorch model to prepare. You don't need to prepare a model if it is used only for inference without - any kind of mixed precision - device_placement (`bool`, *optional*): - Whether or not to place the model on the proper device. Will default to `self.device_placement`. - evaluation_mode (`bool`, *optional*, defaults to `False`): - Whether or not to set the model for evaluation only, by just applying mixed precision and - `torch.compile` (if configured in the `Accelerator` object). - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> # Assume a model is defined - >>> model = accelerator.prepare_model(model) - ``` - """ - if device_placement is None: - device_placement = self.device_placement and self.distributed_type != DistributedType.FSDP - - self._models.append(model) - - # TODO: Look at enabling native TP training directly with a proper config - if ( - self.verify_device_map(model) - and self.distributed_type != DistributedType.NO - and os.environ.get("ACCELERATE_BYPASS_DEVICE_MAP", "false") != "true" - ): - raise ValueError( - "You can't train a model that has been loaded with `device_map='auto'` in any distributed mode." - " Please rerun your script specifying `--num_processes=1` or by launching with `python {{myscript.py}}`." - ) - - if self.native_amp: - model._original_forward = model.forward - autocast_context = get_mixed_precision_context_manager(self.native_amp, self.autocast_handler) - # NOTE: MS-AMP adds `__func__` already to `model.forward`, so we should always use `model.forward` - if self.fp8_backend == FP8BackendType.MSAMP or not hasattr(model.forward, "__func__"): - model_forward_func = model.forward - model.forward = convert_outputs_to_fp32(autocast_context(model_forward_func)) - else: - model_forward_func = model.forward.__func__ - new_forward = autocast_context(model_forward_func) - model.forward = MethodType(new_forward, model) - model.forward = MethodType(convert_outputs_to_fp32(model.forward.__func__), model) - - # We prepare TE after, allowing for bf16 autocast to happen first - if self.fp8_backend == FP8BackendType.TE and not self.delayed_fp8_autocast: - model = apply_fp8_autowrap(model, self.te_recipe_handler or self.fp8_recipe_handler) - - if (getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False)) and getattr( - model, "hf_device_map", False - ): - model_devices = set(model.hf_device_map.values()) - if len(model_devices) > 1 and self.distributed_type != DistributedType.NO: - raise ValueError( - "You can't train a model that has been loaded in 8-bit or 4-bit precision on multiple devices in any distributed mode." - " In order to use 8-bit or 4-bit models that have been loaded across multiple GPUs the solution is to use Naive Pipeline Parallelism." - " Therefore you should not specify that you are under any distributed regime in your accelerate config." - ) - elif len(model_devices) == 1: - current_device = list(model_devices)[0] - if isinstance(current_device, torch.device): - current_device_index = current_device.index - elif isinstance(current_device, str): - current_device_index = torch.device(current_device).index - else: - current_device_index = current_device - - if self.device.type == "cpu" and is_bitsandbytes_multi_backend_available(): - # bnb with multi-backend supports CPU which don't need to check index. - pass - elif torch.device(current_device_index) != self.device: - # if on the first device (GPU 0) we don't care - if (self.device.index is not None) or (current_device_index != 0): - raise ValueError( - "You can't train a model that has been loaded in 8-bit or 4-bit precision on a different device than the one " - "you're training on. Make sure you loaded the model on the correct device using for example `device_map={'':torch.cuda.current_device()}` or `device_map={'':torch.xpu.current_device()}`" - ) - if ( - ("cpu" in model_devices and not is_bitsandbytes_multi_backend_available()) - or ("cpu" in model_devices and is_xpu_available()) - or "disk" in model_devices - ): - raise ValueError( - "You can't train a model that has been loaded in 8-bit or 4-bit precision with CPU or disk offload. " - "If you want train the 8-bit or 4-bit model in CPU, please install bitsandbytes with multi-backend, see https://huggingface.co/docs/bitsandbytes/main/en/installation#multi-backend" - ) - elif device_placement and not self.verify_device_map(model): - model = model.to(self.device) - if not evaluation_mode: - if self.distributed_type in ( - DistributedType.MULTI_GPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_NPU, - DistributedType.MULTI_XPU, - DistributedType.MULTI_HPU, - ): - if model_has_dtensor(model): - raise ValueError( - "Your model contains `DTensor` parameters, which is incompatible with DDP. Maybe you loaded your model with `device_map='auto'`? Specify `device_map='cuda'` or 'cpu' instead." - ) - if any(p.requires_grad for p in model.parameters()): - kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {} - # TODO: Look at enabling native TP training directly with a proper config - if os.environ.get("ACCELERATE_BYPASS_DEVICE_MAP", "false") != "true": - if self.device.type == "hpu": - device_ids, output_device = [self.device.index], self.device.index - else: - device_ids, output_device = [self.local_process_index], self.local_process_index - else: - device_ids, output_device = None, None - - model = torch.nn.parallel.DistributedDataParallel( - model, device_ids=device_ids, output_device=output_device, **kwargs - ) - if self.ddp_handler is not None: - self.ddp_handler.register_comm_hook(model) - elif self.distributed_type == DistributedType.TP: - if not compare_versions("transformers", ">=", BETA_TP_AVAILABLE_TRANSFORMERS_VERSION): - raise ValueError(f"TP requires transformers >= {BETA_TP_AVAILABLE_TRANSFORMERS_VERSION}") - if not hasattr(model, "tp_size"): - raise NotImplementedError( - "Model should undergo tensor parallel before passing it to accelerate." - "You can use .from_pretrained(..., tp_plan='auto') if the model supports" - ) - if model.tp_size != self.state.torch_tp_plugin.tp_size: - raise ValueError( - f"tp_size in the plugin {self.state.torch_tp_plugin.tp_size} should be same as model's tp size {model.tp_size}" - ) - elif self.is_fsdp2: - raise ValueError( - "FSDP2 preparation should be done via `accelerate.prepare()`, as it requires a model and an optimizer." - ) - - elif self.distributed_type == DistributedType.FSDP: - # We need to fix the optimizer *before* sharding the model - from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP - - # Check if the model is already a FSDP model due to `Manual Wrapping` and if so, - # don't wrap it again - # In case the model is already compiled using PyTorch 2.0 and the wrapped model in it - # is a FSDP model, don't wrap it again - is_type_fsdp = isinstance(model, FSDP) or ( - is_compiled_module(model) and isinstance(model._orig_mod, FSDP) - ) - - if not is_type_fsdp: - self.state.fsdp_plugin.set_auto_wrap_policy(model) - fsdp_plugin = self.state.fsdp_plugin - - # need to ensure that params are re-tied after running - # param_init_fn - fsdp_plugin.param_init_fn = ensure_weights_retied( - fsdp_plugin.param_init_fn, - model, - self.device, - ) - - kwargs = { - # We fallback to reshard_after_forward if sharding_strategy is not set. - # We prerfer sharding_strategy to not break the behavior of the existing code. - # Deprecation warning has already been issued in `utils.dataclasses.py` - "sharding_strategy": fsdp_plugin.sharding_strategy or fsdp_plugin.reshard_after_forward, - "cpu_offload": fsdp_plugin.cpu_offload, - "auto_wrap_policy": fsdp_plugin.auto_wrap_policy, - "mixed_precision": fsdp_plugin.mixed_precision_policy, - "sync_module_states": fsdp_plugin.sync_module_states, - "backward_prefetch": fsdp_plugin.backward_prefetch, - "forward_prefetch": fsdp_plugin.forward_prefetch, - "use_orig_params": fsdp_plugin.use_orig_params, - "param_init_fn": fsdp_plugin.param_init_fn, - "ignored_modules": fsdp_plugin.ignored_modules, - "limit_all_gathers": fsdp_plugin.limit_all_gathers, - "device_id": self.device, - } - model = FSDP(model, **kwargs) - if fsdp_plugin.activation_checkpointing: - from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( - CheckpointImpl, - apply_activation_checkpointing, - checkpoint_wrapper, - ) - - apply_activation_checkpointing( - model, - checkpoint_wrapper_fn=functools.partial( - checkpoint_wrapper, - checkpoint_impl=CheckpointImpl.NO_REENTRANT, - ), - auto_wrap_policy=fsdp_plugin.auto_wrap_policy, - ) - - # In the event the model had been loaded in low precision, but - # mixed precision had also been activated, then we follow DeepSpeed's - # strategy to hold the parameters in full precision. - # - assume that trainer.args.bf16 and trainer.args.fp16 are already checked against - # fsdp_plugin.mixed_precision_policy. - # - NOTE: we do not check the mixed_precision attribute on the FSDP root wrapper. - # * this attribute will always set by init_utils.init_core_state so its always not None. - # * mixed_precision.param_dtype only regards _fwd_bwd_param_dtype - # * if model is loaded in 16bit, and even if mixed_precision.param_dtype is None, - # we still want to upcast the flat_param. - if self.mixed_precision != "no": # if mixed precision is set - upcasted_log = [] - for module in FSDP.fsdp_modules(model): - # Referencing DeepSpeed Zero3 - # - in Init, params are converted to 16bit while partitioning. - # - in accelerator.prepare, deepspeed.initialize is called to: - # * creates the DeepSpeedEngine. - # * since zero_optimization() is True , calls engine._configure_zero_optimizer. - # - # Inside the DeepSpeed Zero3 optimizer configuration, which initializes - # DeepSpeedZeroOptimizer_Stage3, during which: - # * trainable_param_groups are obtained from the attached optimizer - # (already partitioned in 16bit). - # * then _setup_for_real_optimizer -> _create_fp32_partitions - # which performs the fp32 upcasting. - - # To mimic DeepSeepds's casting in FSDP, we look at the (single) FlatParameter held - # within an FSDP wrapper. This FlatParameter will be seen by the optimizer. - # - even though there is a torch.device('meta') guard below, we - # expect _init_utils._init_param_handle_from_module to already - # sync the parameter. - - if not module._has_params: - continue # skip if FSDP module not managing parameters - param = module._flat_param - if ( - param.dtype != torch.float32 - and param.device != torch.device("meta") - and param.requires_grad - ): - # keep log of names_params that was upcasted - # NOTE: resorted to this because warnings.simplefilter("once") is somehow not working - name_param_log = (module.module.__class__.__name__, ", ".join(module._flat_param._fqns)) - if name_param_log not in upcasted_log: - upcasted_log.append(name_param_log) - - # this works because of FSDP's _runtime_utils.lazy_init. - # Have to be careful not to call anything before this that - # triggers lazy_init (e.g., _is_fsdp_root). - param.data = param.data.to(torch.float32) # upcasting - module._handle._orig_param_dtype = torch.float32 # update - - # report the warnings - # some messages can be quite repetitive, especially when reporting about layers that have identical architecture. - if self.is_main_process: - for name_log, param_log in upcasted_log: - warnings.warn( - f"Upcasted low precision parameters in {name_log} because mixed precision turned on in FSDP. " - f"Affects: {param_log}." - ) - - if len(upcasted_log) > 0: - warnings.warn( - "FSDP upcast of low precision parameters may affect the precision of model checkpoints." - ) - - # if the previous and current models are same, delete the previous one - if len(self._models) > 1 and (self._models[-2] is self._models[-1]): - del self._models[-2] - self._models[-1] = model - elif self.distributed_type == DistributedType.MULTI_CPU: - kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {} - model = torch.nn.parallel.DistributedDataParallel(model, **kwargs) - if self.ddp_handler is not None: - self.ddp_handler.register_comm_hook(model) - elif self.distributed_type == DistributedType.XLA and self.state.fork_launched: - model = xmp.MpModelWrapper(model).to(self.device) - # Now we can apply the FP8 autocast - if self.fp8_backend == FP8BackendType.TE and self.delayed_fp8_autocast: - model = apply_fp8_autowrap(model, self.te_recipe_handler or self.fp8_recipe_handler) - # torch.compile should be called last and only if the model isn't already compiled - if self.state.dynamo_plugin.backend != DynamoBackend.NO and not is_compiled_module(model): - if self.state.dynamo_plugin.use_regional_compilation: - model = compile_regions(model, **self.state.dynamo_plugin.to_kwargs()) - else: - model = torch.compile(model, **self.state.dynamo_plugin.to_kwargs()) - return model - - def _prepare_ao(self, *args): - if not is_torchao_available(): - raise ImportError( - "`torchao` was not found on your system or is too old of a version. Please ensure that `torchao >= 0.6.1` is installed" - ) - - if self.is_fsdp2: - models = [x for x in args if isinstance(x, torch.nn.Module)] - optimizers = [x for x in args if isinstance(x, torch.optim.Optimizer)] - for arg in args: - if isinstance(arg, torch.nn.Module): - convert_model_to_fp8_ao( - arg, - config=self.ao_recipe_handler.config, - module_filter_func=self.ao_recipe_handler.module_filter_func, - ) - - # Invariant: with FSDP2, optimizer is always passed to `prepare()` together with model - # We only precompute scales if float8 all gather is enabled, possibly can add a flag for this later - if self.is_fsdp2 and len(optimizers) > 0 and self.ao_recipe_handler.config.enable_fsdp_float8_all_gather: - from torchao.float8 import precompute_float8_dynamic_scale_for_fsdp - - optimizers[0].register_step_post_hook( - lambda *args, **kwargs: precompute_float8_dynamic_scale_for_fsdp(models[0]) - ) - - return args - - def _prepare_te(self, *args): - if not is_transformer_engine_available(): - raise ImportError( - "`transformer_engine` was not found on your system. Please ensure that `transformer_engine` is installed" - ) - model, optimizer = None, None - num_models, num_optimizers = 0, 0 - result = [obj for obj in args] - for obj in result: - if isinstance(obj, torch.nn.Module): - model = obj - num_models += 1 - elif isinstance(obj, (torch.optim.Optimizer)): - optimizer = obj - num_optimizers += 1 - if optimizer is None and model is None: - return result - elif optimizer is None or model is None: - raise ValueError( - "You must pass a model and an optimizer together to `accelerate.prepare()` when using TransformerEngine." - ) - elif num_models > 1 or num_optimizers > 1: - raise ValueError( - f"You can't use multiple models ({num_models}) or optimizers {num_optimizers} with TransformerEngine." - ) - old_named_params = self._get_named_parameters(model) - with torch.no_grad(): - convert_model(model) - new_named_params = self._get_named_parameters(model) - mapping = {p: new_named_params[n] for n, p in old_named_params.items()} - # We need to switch the optimizer params to the new params *after* the model is wrapped in FSDP - for param_group in optimizer.param_groups: - param_group["params"] = [mapping[p] for p in param_group["params"]] - - return result - - def _prepare_deepspeed(self, *args): - import deepspeed - - ds_initialize = deepspeed.initialize - if self.fp8_backend == FP8BackendType.MSAMP: - # MS-AMP requires DeepSpeed patches - from msamp import deepspeed as msamp_deepspeed - - ds_initialize = msamp_deepspeed.initialize - - deepspeed_plugin = self.deepspeed_plugin - - is_dataloader_present = any(isinstance(obj, torch.utils.data.DataLoader) for obj in args) - tp_size = deepspeed_plugin.deepspeed_config.get("tensor_parallel", {}).get("autotp_size", 0) - if tp_size > 1: - if not compare_versions("deepspeed", ">=", "0.16.4"): - raise ImportError( - "Deepspeed TP requires deepspeed >= 0.16.4, Please update DeepSpeed via `pip install deepspeed -U`." - ) - if not is_torch_version(">=", "2.2.0"): - raise ImportError( - "Tried to use TP, but `torch.distributed.device_mesh` requires PyTorch >= 2.2.0. Please upgrade your PyTorch version" - ) - from torch.distributed.device_mesh import init_device_mesh - - mesh_dim_name = "tp" - self.state.ds_device_mesh = init_device_mesh(self.device.type, (tp_size,), mesh_dim_names=(mesh_dim_name,)) - - result = [ - self._prepare_one(obj, first_pass=True) if isinstance(obj, torch.utils.data.DataLoader) else obj - for obj in args - ] - - if deepspeed_plugin.is_auto("train_micro_batch_size_per_gpu"): - if is_dataloader_present: - batch_sizes = [obj.batch_size for obj in args if hasattr(obj, "batch_size")] - if any(bs is None for bs in batch_sizes): - raise ValueError( - "At least one of the dataloaders passed to `accelerate.prepare()` has `None` as batch size. " - "Please set an integer value in `train_micro_batch_size_per_gpu` in the deepspeed config file " - "or assign integer value to `AcceleratorState().deepspeed_plugin.deepspeed_config['train_micro_batch_size_per_gpu']`." - ) - if self.split_batches: - batch_sizes = [batch_size // self.num_processes for batch_size in batch_sizes] - - batch_size_per_device = min(batch_sizes) if deepspeed_plugin.is_train_batch_min else max(batch_sizes) - if len(batch_sizes) > 1: - logger.info( - "Since you passed both train and evaluation dataloader, `is_train_batch_min` (here " - f"{deepspeed_plugin.is_train_batch_min} will decide the `train_batch_size` ({batch_size_per_device})." - ) - else: - raise ValueError( - "When using DeepSpeed, `accelerate.prepare()` requires you to pass at least one of training or evaluation dataloaders " - "with `batch_size` attribute returning an integer value " - "or alternatively set an integer value in `train_micro_batch_size_per_gpu` in the deepspeed config file " - "or assign integer value to `AcceleratorState().deepspeed_plugin.deepspeed_config['train_micro_batch_size_per_gpu']`." - ) - else: - batch_size_per_device = deepspeed_plugin.get_value("train_micro_batch_size_per_gpu") - - # handle `gradient_accumulation_steps` when the value is `auto` - deepspeed_plugin.fill_match( - "gradient_accumulation_steps", - must_match=False, - gradient_accumulation_steps=self.gradient_accumulation_steps, - ) - - deepspeed_gradient_accumulation_steps = deepspeed_plugin.get_value("gradient_accumulation_steps") - # update gradient_accumulation_steps if there is a mismatch - if deepspeed_gradient_accumulation_steps != self.gradient_accumulation_steps: - logger.warning( - f"Gradient accumulation steps mismatch: GradientAccumulationPlugin has {self.gradient_accumulation_steps}, " - f"DeepSpeed config has {deepspeed_gradient_accumulation_steps}. Using DeepSpeed's value." - ) - self.gradient_accumulation_steps = deepspeed_gradient_accumulation_steps - - config_kwargs = { - "gradient_clipping": 1.0, - "zero_optimization.stage3_gather_16bit_weights_on_model_save": False, - } - # This is skipped when preparing just a model - if batch_size_per_device is not None: - config_kwargs["train_micro_batch_size_per_gpu"] = batch_size_per_device - config_kwargs["train_batch_size"] = ( - batch_size_per_device * deepspeed_plugin.get_value("gradient_accumulation_steps") * self.num_processes - ) - - model = None - optimizer = None - scheduler = None - for obj in result: - if isinstance(obj, torch.nn.Module): - model = obj - elif isinstance(obj, (torch.optim.Optimizer, DummyOptim)): - optimizer = obj - elif (isinstance(obj, (LRScheduler, DummyScheduler))) or ( - type(obj).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES - ): - scheduler = obj - - if optimizer is not None: - if "optimizer" in deepspeed_plugin.deepspeed_config and not isinstance(optimizer, (DummyOptim)): - raise ValueError( - "You cannot specify an optimizer in the config file and in the code at the same time. " - "Please remove the optimizer from the config file or " - "create `accelerate.utils.DummyOptim` in the code." - ) - elif "optimizer" not in deepspeed_plugin.deepspeed_config and isinstance(optimizer, (DummyOptim)): - raise ValueError( - "You cannot create a `DummyOptim` without specifying an optimizer in the config file." - ) - - if isinstance(optimizer, (torch.optim.Optimizer)): - deepspeed_plugin.deepspeed_config["zero_allow_untested_optimizer"] = True - - if scheduler is not None: - if "scheduler" in deepspeed_plugin.deepspeed_config and not isinstance(scheduler, (DummyScheduler)): - raise ValueError( - "You cannot specify a scheduler in the config file and in the code at the same time. " - "Please remove the scheduler from the config file or " - "create `accelerate.utils.DummyScheduler` in the code." - ) - elif ( - "scheduler" not in deepspeed_plugin.deepspeed_config - and isinstance(scheduler, (DummyScheduler)) - and scheduler.lr_scheduler_callable is None - ): - raise ValueError( - "Either specify a scheduler in the config file or " - "pass in the `lr_scheduler_callable` parameter when using `accelerate.utils.DummyScheduler`." - ) - - if optimizer is not None and scheduler is not None: - if isinstance(optimizer, (DummyOptim)) and not isinstance(scheduler, (DummyScheduler)): - raise ValueError( - "You can only specify `accelerate.utils.DummyScheduler` in the code when using " - "`accelerate.utils.DummyOptim`." - ) - - if model is not None: - # If we are using FP8, we need to apply the autowrap now - if self.fp8_backend == FP8BackendType.TE: - model = apply_fp8_autowrap(model, self.fp8_recipe_handler) - # if the model is an MOE, set the appropriate MOE layers as leaf Z3 modules - deepspeed_plugin.set_moe_leaf_modules(model) - # deal with config keys that use `auto` value and rely on model's hidden_size - hidden_size_based_keys = [ - "zero_optimization.reduce_bucket_size", - "zero_optimization.stage3_prefetch_bucket_size", - "zero_optimization.stage3_param_persistence_threshold", - ] - hidden_size_auto_keys = [x for x in hidden_size_based_keys if deepspeed_plugin.is_auto(x)] - if len(hidden_size_auto_keys) > 0: - reasoning = ( - "therefore it's not possible to automatically fill out the following `auto` entries " - + f"in the DeepSpeed config file: {hidden_size_auto_keys}. You can fix that by replacing " - + "`auto` values for these keys with an integer value of your choice." - ) - if not hasattr(model, "config"): - raise ValueError("Can't find `model.config` entry, " + reasoning) - - if hasattr(model.config, "hidden_size"): - hidden_size = model.config.hidden_size - elif hasattr(model.config, "hidden_sizes"): - # if there are many hidden sizes pick the largest one - hidden_size = max(model.config.hidden_sizes) - else: - raise ValueError( - "Can find neither `model.config.hidden_size` nor `model.config.hidden_sizes`, " + reasoning - ) - - config_kwargs.update( - { - "zero_optimization.reduce_bucket_size": hidden_size * hidden_size, - "zero_optimization.stage3_prefetch_bucket_size": int(0.9 * hidden_size * hidden_size), - "zero_optimization.stage3_param_persistence_threshold": 10 * hidden_size, - } - ) - - if isinstance(optimizer, (DummyOptim)): - config_kwargs.update( - {"optimizer.params.lr": optimizer.lr, "optimizer.params.weight_decay": optimizer.weight_decay} - ) - if isinstance(scheduler, (DummyScheduler)) and scheduler.lr_scheduler_callable is None: - max_lr = ( - getattr(scheduler.optimizer, "lr", None) - if getattr(scheduler.optimizer, "defaults", None) is None - else scheduler.optimizer.defaults["lr"] - ) - config_kwargs.update( - { - "scheduler.params.warmup_min_lr": 0, - "scheduler.params.warmup_max_lr": max_lr, - "scheduler.params.warmup_num_steps": scheduler.warmup_num_steps, - } - ) - if scheduler.total_num_steps is not None: - config_kwargs["scheduler.params.total_num_steps"] = ( - math.ceil(scheduler.total_num_steps / self.num_processes) - if not self.split_batches - else scheduler.total_num_steps - ) - deepspeed_plugin.deepspeed_config_process(must_match=False, **config_kwargs) - self.deepspeed_config = deepspeed_plugin.deepspeed_config - kwargs = dict(model=model, config_params=self.deepspeed_config) - if optimizer is not None: - if isinstance(optimizer, (DummyOptim)): - kwargs["model_parameters"] = optimizer.params - if isinstance(scheduler, (DummyScheduler)) and scheduler.lr_scheduler_callable is not None: - kwargs["lr_scheduler"] = scheduler.lr_scheduler_callable - else: - if self.deepspeed_config["zero_optimization"].get("offload_optimizer", {}).get( - "device", "none" - ) != "none" and self.deepspeed_config.get("zero_force_ds_cpu_optimizer", True): - if self.device.type == "hpu" and os.environ.get("PT_HPU_LAZY_MODE", "1") == "1": - raise ValueError( - "You can't use an Offload Optimizer with HPU in Lazy Mode. " - "Please set the environment variable `PT_HPU_LAZY_MODE` to `0`." - ) - - optimizer = map_pytorch_optim_to_deepspeed(optimizer) - kwargs["optimizer"] = optimizer - if scheduler is not None: - if type(scheduler).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES: - kwargs["lr_scheduler"] = scheduler - - if self.device.type == "hpu": - # This env variable is initialized here to make sure it is set to "true" - # It should be done by the launcher but it does not work for multi-node runs - os.environ["DEEPSPEED_USE_HPU"] = "true" - - engine, optimizer, _, lr_scheduler = ds_initialize(**kwargs) - - if compare_versions("deepspeed", ">=", "0.14.4") and self.state.dynamo_plugin.backend != DynamoBackend.NO: - compile_kwargs = self.state.dynamo_plugin.to_kwargs() - if self.state.dynamo_plugin.use_regional_compilation: - compile_regions_deepspeed(engine.module, **compile_kwargs) - else: - engine.compile(backend=compile_kwargs.pop("backend"), compile_kwargs=compile_kwargs) - if optimizer is not None: - optimizer = DeepSpeedOptimizerWrapper(optimizer) - if scheduler is not None: - if lr_scheduler is None: - scheduler = AcceleratedScheduler( - scheduler, - optimizer, - step_with_optimizer=self.step_scheduler_with_optimizer, - split_batches=self.split_batches, - ) - else: - scheduler = DeepSpeedSchedulerWrapper(lr_scheduler, optimizer) - - for i in range(len(result)): - if isinstance(result[i], torch.nn.Module): - result[i] = engine - elif isinstance(result[i], (torch.optim.Optimizer, DummyOptim)): - result[i] = optimizer - elif (isinstance(result[i], (LRScheduler, DummyScheduler))) or ( - type(result[i]).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES - ): - result[i] = scheduler - # pointing for deepspeed_engine_wrapped.backward() - if self.deepspeed_engine_wrapped is None: - self.deepspeed_engine_wrapped = DeepSpeedEngineWrapper(engine) - else: - logger.warning( - "A wrapped DeepSpeed engine reference is currently tied for this `Accelerator()` instance. " - "If you want to call `accelerator.backward()` referencing a new model/engine, " - "please create a separate `Accelerator()` instance and call `accelerator.prepare()` on it." - ) - self._models.append(engine) - if optimizer is not None: - self._optimizers.append(optimizer) - if scheduler is not None: - self._schedulers.append(scheduler) - return tuple(result) - - def _prepare_megatron_lm(self, *args): - megatron_lm_plugin = self.state.megatron_lm_plugin - micro_batch_size = None - if not megatron_lm_plugin.megatron_dataset_flag: - batch_sizes = [obj.batch_size for obj in args if hasattr(obj, "batch_size")] - if len(batch_sizes) == 0: - raise ValueError( - "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using Megatron-LM." - ) - - micro_batch_size = min(batch_sizes) if megatron_lm_plugin.is_train_batch_min else max(batch_sizes) - if len(batch_sizes) > 1: - logger.info( - "Since you passed both train and evaluation dataloader, `is_train_batch_min` (here " - f"{megatron_lm_plugin.is_train_batch_min} will decide the `train_batch_size` ({micro_batch_size})." - ) - else: - for obj in args: - if isinstance(obj, MegatronLMDummyDataLoader): - micro_batch_size = obj.dataset_args["micro_batch_size"] - break - if micro_batch_size is not None: - dp_degree = self.num_processes // (megatron_lm_plugin.tp_degree * megatron_lm_plugin.pp_degree) - megatron_lm_plugin.set_training_args(micro_batch_size, dp_degree) - else: - raise ValueError( - "When you do not pass the dataloader parameter, the `data_parallel_size`, " - "`micro_batch_size`, and `global_batch_size` megatron parameters will not be updated." - ) - model = None - optimizer = None - scheduler = None - batch_data = None - for obj in args: - if isinstance(obj, torch.utils.data.DataLoader) and batch_data is None: - batch_data = next(iter(obj)) - elif isinstance(obj, torch.nn.Module): - model = obj - elif isinstance(obj, (torch.optim.Optimizer)): - optimizer = obj - elif isinstance(obj, (LRScheduler, MegatronLMDummyScheduler)): - scheduler = obj - - if model is not None: - megatron_lm_plugin.set_network_size_args(model, batch_data) - if optimizer is not None: - megatron_lm_plugin.set_optimizer_type(optimizer) - if scheduler is not None: - if not isinstance(scheduler, MegatronLMDummyScheduler): - raise ValueError( - "You can't use a custom scheduler with Megatron-LM. Please use the `accelerate.utils.MegatronLMDummyScheduler` instead." - ) - megatron_lm_plugin.set_scheduler_args(scheduler) - - # initialize megatron-lm - megatron_lm_initialize(self, args_defaults=megatron_lm_plugin.megatron_lm_default_args) - - (model, optimizer, scheduler) = megatron_lm_prepare_model_optimizer_scheduler(self) - self.wait_for_everyone() - - counter = 0 - result = [] - for obj in args: - if isinstance(obj, torch.utils.data.DataLoader): - result.append(megatron_lm_prepare_data_loader(self, obj)) - counter += 1 - elif isinstance(obj, MegatronLMDummyDataLoader): - if counter == 0: - obj.set_megatron_data_args() - dataloaders = megatron_lm_prepare_data_loader(self, obj) - result.append(dataloaders[counter]) - counter += 1 - else: - result.append(obj) - - if model is not None: - model = MegatronEngine(self, model, optimizer, scheduler) - if optimizer is not None: - optimizer = MegatronLMOptimizerWrapper(optimizer) - if scheduler is not None: - scheduler = MegatronLMSchedulerWrapper(scheduler, optimizer) - - for i in range(len(result)): - if isinstance(result[i], torch.nn.Module): - result[i] = model - elif isinstance(result[i], torch.optim.Optimizer): - result[i] = optimizer - elif isinstance(result[i], MegatronLMDummyScheduler): - result[i] = scheduler - - if model is not None: - self._models.append(model) - if len(self._models) > 1: - raise AssertionError( - "You can't use same `Accelerator()` instance with multiple models when using Megatron-LM" - ) - if optimizer is not None: - self._optimizers.append(optimizer) - if scheduler is not None: - self._schedulers.append(scheduler) - - return tuple(result) - - def _prepare_ipex(self, *args): - """ - Prepares model and optimizer for training with IPEX on CPU/XPU. This covers 3 cases, IPEX compiled with CPU - only support, IPEX compiled with XPU support and training with XPU pytorch backend available in stock pytorch - starting from version 2.4. - """ - - # ipex.optimize() is available only for IPEX, both IPEX-CPU and IPEX-XPU - if is_ipex_available(): - import intel_extension_for_pytorch as ipex - else: - raise ImportError( - "IPEX is not installed or IPEX's version does not match current PyTorch version. Please refer" - " to https://github.com/intel/intel-extension-for-pytorch." - ) - - models = [] - optimizers = [] - result = [obj for obj in args] - for i, obj in enumerate(result): - if isinstance(obj, torch.nn.Module): - model = obj - model.train() - models.append((i, model)) - elif isinstance(obj, (torch.optim.Optimizer)): - optimizers.append((i, obj)) - - # Impossible to determine what to do if multiple models and/or optimizers are provided - if len(optimizers) > 1 or (len(models) > 1 and len(optimizers) == 1): - raise ValueError( - "Prepare with IPEX expects either 1+ models and no optimizer OR a single model-optimizer pair." - ) - - # Nothing to do - if len(models) == 0 and len(optimizers) == 0: - return result - - dtype = torch.bfloat16 if self.state.mixed_precision == "bf16" else None - # Multiple models and no optimizer (inference) are provided - if len(models) > 0 and len(optimizers) == 0: - for i, model in models: - if self.device.type == "xpu" and next(model.parameters()).device.type == "cpu": - model = model.to(self.device) - model, _ = ipex.optimize(model, optimizer=None, dtype=dtype, inplace=True, level="O1") - # Replace in result - result[i] = model - - # A single model-optimizer pair (training) is provided - if len(models) == 1 and len(optimizers) == 1: - i_model, model = models[0] - i_optimizer, optimizer = optimizers[0] - if self.device.type == "xpu" and next(model.parameters()).device.type == "cpu": - model = model.to(self.device) - model, optimizer = ipex.optimize(model, optimizer=optimizer, dtype=dtype, inplace=True, level="O1") - # Replace in result - result[i_model] = model - result[i_optimizer] = optimizer - - return tuple(result) - - def _prepare_device_mesh(self): - """ - Prepare the device mesh for distributed training. The dataloader will determine how to load data based on the - device mesh. - """ - if self.state.torch_tp_plugin: - return self.state.torch_tp_plugin.torch_device_mesh - elif self.distributed_type == DistributedType.DEEPSPEED and hasattr(self.state, "ds_device_mesh"): - return self.state.ds_device_mesh - return None - - def _prepare_msamp(self, *args, device_placement): - if not is_msamp_available(): - raise ImportError( - "MS-AMP was not found on your system. Please ensure that MS-AMP is available " - " or choose `'te'` as the backend for FP8 mixed precision training." - ) - # We've already checked for FSDP + MS-AMP during `__init__` - import msamp - - model, optimizer = None, None - optimizer_index = None - num_models, num_optimizers = 0, 0 - result = [obj for obj in args] - for i, obj in enumerate(result): - if isinstance(obj, torch.nn.Module): - model = obj - num_models += 1 - elif isinstance(obj, (torch.optim.Optimizer)): - optimizer = obj - optimizer_index = i - num_optimizers += 1 - # DataLoader/Scheduler case - if optimizer is None and model is None: - return result, device_placement - elif optimizer is None or model is None: - raise ValueError( - "You must pass a model and an optimizer together to `accelerate.prepare()` when using MS-AMP." - ) - elif num_models > 1 or num_optimizers > 1: - raise ValueError( - f"You can't use multiple models ({num_models}) or optimizers {num_optimizers} with MS-AMP." - ) - else: - # DEPRECATE @ 2.0 - if self.fp8_recipe_handler is not None: - opt_level = self.fp8_recipe_handler.opt_level - else: - opt_level = self.msamp_recipe_handler.opt_level - model, optimizer = msamp.initialize(model, optimizer, opt_level=opt_level) - for i in range(len(result)): - if isinstance(result[i], torch.nn.Module): - result[i] = model - elif isinstance(result[i], (torch.optim.Optimizer)): - result[i] = optimizer - if optimizer_index is not None: - # NOTE: MS-AMP moves the optimizer, but *not* the model to the right device - device_placement[optimizer_index] = False - return tuple(result), device_placement - - def prepare_data_loader( - self, data_loader: torch.utils.data.DataLoader, device_placement=None, slice_fn_for_dispatch=None - ): - """ - Prepares a PyTorch DataLoader for training in any distributed setup. It is recommended to use - [`Accelerator.prepare`] instead. - - Args: - data_loader (`torch.utils.data.DataLoader`): - A vanilla PyTorch DataLoader to prepare - device_placement (`bool`, *optional*): - Whether or not to place the batches on the proper device in the prepared dataloader. Will default to - `self.device_placement`. - slice_fn_for_dispatch (`Callable`, *optional*`): - If passed, this function will be used to slice tensors across `num_processes`. Will default to - [`~utils.slice_tensors`]. This argument is used only when `dispatch_batches` is set to `True` and will - be ignored otherwise. - - Example: - - ```python - >>> import torch - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> data_loader = torch.utils.data.DataLoader(...) - >>> data_loader = accelerator.prepare_data_loader(data_loader, device_placement=True) - ``` - """ - # Ensure we can't double wrap a DataLoader due to `find_batch_size` - if getattr(data_loader, "_is_accelerate_prepared", False): - if data_loader not in self._dataloaders: - self._dataloaders.append(data_loader) - return data_loader - if device_placement is None: - device_placement = self.device_placement if self.distributed_type != DistributedType.XLA else False - - device_mesh = self._prepare_device_mesh() - - prepared_data_loader = prepare_data_loader( - data_loader, - self.device, - num_processes=self.num_processes, - process_index=self.process_index, - split_batches=self.split_batches, - put_on_device=device_placement, - rng_types=self.rng_types.copy(), - dispatch_batches=self.dispatch_batches, - even_batches=self.even_batches, - slice_fn_for_dispatch=slice_fn_for_dispatch, - use_seedable_sampler=self.use_seedable_sampler, - data_seed=self.dataloader_config.data_seed, - non_blocking=self.non_blocking, - use_stateful_dataloader=self.use_stateful_dataloader, - torch_device_mesh=device_mesh, - ) - self._dataloaders.append(prepared_data_loader) - return prepared_data_loader - - def prepare_optimizer(self, optimizer: torch.optim.Optimizer, device_placement=None): - """ - Prepares a PyTorch Optimizer for training in any distributed setup. It is recommended to use - [`Accelerator.prepare`] instead. - - Args: - optimizer (`torch.optim.Optimizer`): - A vanilla PyTorch optimizer to prepare - device_placement (`bool`, *optional*): - Whether or not to place the optimizer on the proper device. Will default to `self.device_placement`. - - Example: - - ```python - >>> import torch - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> optimizer = torch.optim.Adam(...) - >>> optimizer = accelerator.prepare_optimizer(optimizer, device_placement=True) - ``` - """ - if is_lomo_available(): - # We need to import locally to avoid circular imports since lomo imports stuff from - # transformers & accelerate - from lomo_optim import AdaLomo, Lomo - - # Support multiple optimizers: https://github.com/huggingface/accelerate/pull/2695#discussion_r1589164607 - self.has_lomo_optimizer |= isinstance(optimizer, (Lomo, AdaLomo)) - - # Ensure we can't double wrap an optimizer due to `find_batch_size` - if getattr(optimizer, "_is_accelerate_prepared", False): - if optimizer not in self._optimizers: - self._optimizers.append(optimizer) - return optimizer - if device_placement is None: - device_placement = self.device_placement - # NOTE: Special case with MS-AMP we do *not* pass in the scaler explicitly to the `AcceleratedOptimizer`, - # Their optimizer handles it for us. - scaler = None if self.fp8_backend == FP8BackendType.MSAMP else self.scaler - optimizer = AcceleratedOptimizer(optimizer, device_placement=device_placement, scaler=scaler) - self._optimizers.append(optimizer) - return optimizer - - def prepare_scheduler(self, scheduler: LRScheduler): - """ - Prepares a PyTorch Scheduler for training in any distributed setup. It is recommended to use - [`Accelerator.prepare`] instead. - - Args: - scheduler (`torch.optim.lr_scheduler.LRScheduler`): - A vanilla PyTorch scheduler to prepare - - Example: - - ```python - >>> import torch - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> optimizer = torch.optim.Adam(...) - >>> scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, ...) - >>> scheduler = accelerator.prepare_scheduler(scheduler) - ``` - """ - # Ensure we can't double wrap a scheduler due to `find_batch_size` - if getattr(scheduler, "_is_accelerate_prepared", False): - if scheduler not in self._schedulers: - self._schedulers.append(scheduler) - return scheduler - # We try to find the optimizer associated with `scheduler`, the default is the full list. - optimizer = self._optimizers - for opt in self._optimizers: - if getattr(scheduler, "optimizer", None) == opt.optimizer: - optimizer = opt - break - scheduler = AcceleratedScheduler( - scheduler, - optimizer, - step_with_optimizer=self.step_scheduler_with_optimizer, - split_batches=self.split_batches, - ) - self._schedulers.append(scheduler) - return scheduler - - def backward(self, loss, **kwargs): - """ - Scales the gradients in accordance to the `GradientAccumulationPlugin` and calls the correct `backward()` based - on the configuration. - - Should be used in lieu of `loss.backward()`. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator(gradient_accumulation_steps=2) - >>> outputs = model(inputs) - >>> loss = loss_fn(outputs, labels) - >>> accelerator.backward(loss) - ``` - """ - learning_rate = kwargs.get("learning_rate") - - if self.distributed_type != DistributedType.DEEPSPEED: - # deepspeed handles loss scaling by gradient_accumulation_steps in its `backward` - loss = loss / self.gradient_accumulation_steps - if self.distributed_type == DistributedType.DEEPSPEED: - self.deepspeed_engine_wrapped.backward(loss, sync_gradients=self.sync_gradients, **kwargs) - elif self.distributed_type == DistributedType.MEGATRON_LM: - return - elif self.scaler is not None: - self.scaler.scale(loss).backward(**kwargs) - elif learning_rate is not None and self.has_lomo_optimizer: - self.lomo_backward(loss, learning_rate) - else: - loss.backward(**kwargs) - - def set_trigger(self): - """ - Sets the internal trigger tensor to 1 on the current process. A latter check should follow using this which - will check across all processes. - - Note: - Does not require `wait_for_everyone()` - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> # Assume later in the training script - >>> # `should_do_breakpoint` is a custom function to monitor when to break, - >>> # e.g. when the loss is NaN - >>> if should_do_breakpoint(loss): - ... accelerator.set_trigger() - >>> # Assume later in the training script - >>> if accelerator.check_breakpoint(): - ... break - ``` - """ - self.flag_tensor = torch.tensor(1, device=self.device) - - def check_trigger(self): - """ - Checks if the internal trigger tensor has been set to 1 in any of the processes. If so, will return `True` and - reset the trigger tensor to 0. - - Note: - Does not require `wait_for_everyone()` - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> # Assume later in the training script - >>> # `should_do_breakpoint` is a custom function to monitor when to break, - >>> # e.g. when the loss is NaN - >>> if should_do_breakpoint(loss): - ... accelerator.set_trigger() - >>> # Assume later in the training script - >>> if accelerator.check_trigger(): - ... break - ``` - """ - # Now that we are outside `__init__`, we can initialize it if it is `None` on device - if self.flag_tensor is None: - self.flag_tensor = torch.tensor(0, device=self.device) - flag_tensor = self.reduce(self.flag_tensor) - if flag_tensor.item() >= 1: - self.flag_tensor = torch.tensor(0, device=self.device) - return True - return False - - def unscale_gradients(self, optimizer=None): - """ - Unscale the gradients in mixed precision training with AMP. This is a noop in all other settings. - - Likely should be called through [`Accelerator.clip_grad_norm_`] or [`Accelerator.clip_grad_value_`] - - Args: - optimizer (`torch.optim.Optimizer` or `list[torch.optim.Optimizer]`, *optional*): - The optimizer(s) for which to unscale gradients. If not set, will unscale gradients on all optimizers - that were passed to [`~Accelerator.prepare`]. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> model, optimizer = accelerator.prepare(model, optimizer) - >>> outputs = model(inputs) - >>> loss = loss_fn(outputs, labels) - >>> accelerator.backward(loss) - >>> accelerator.unscale_gradients(optimizer=optimizer) - ``` - """ - if self.native_amp and self.mixed_precision == "fp16": - if optimizer is None: - # TODO: this unscales all optimizers where we should only unscale the one where parameters are. - optimizer = self._optimizers - elif not isinstance(optimizer, (tuple, list)): - optimizer = [optimizer] - for opt in optimizer: - while isinstance(opt, AcceleratedOptimizer): - opt = opt.optimizer - self.scaler.unscale_(opt) - - def clip_grad_norm_(self, parameters, max_norm, norm_type=2): - """ - Should be used in place of `torch.nn.utils.clip_grad_norm_`. - - Returns: - `torch.Tensor`: Total norm of the parameter gradients (viewed as a single vector). - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator(gradient_accumulation_steps=2) - >>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler) - - >>> for input, target in dataloader: - ... optimizer.zero_grad() - ... output = model(input) - ... loss = loss_func(output, target) - ... accelerator.backward(loss) - ... if accelerator.sync_gradients: - ... accelerator.clip_grad_norm_(model.parameters(), max_grad_norm) - ... optimizer.step() - ``` - """ - if self.distributed_type == DistributedType.FSDP: - self.unscale_gradients() - parameters = [p for p in parameters] - for model in self._models: - if parameters == [p for p in model.parameters()]: - if not self.is_fsdp2: - return model.clip_grad_norm_(max_norm, norm_type) - else: - return torch.nn.utils.clip_grad_norm_( - parameters, max_norm, norm_type=norm_type - ) # viz: https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md - elif self.distributed_type == DistributedType.DEEPSPEED: - # DeepSpeed handles gradient clipping internally, but we can retrieve the gradient norm - if self.deepspeed_engine_wrapped is not None: - return self.deepspeed_engine_wrapped.get_global_grad_norm() - return None - elif self.distributed_type == DistributedType.XLA: - # Reduce gradients first for XLA - for acc_opt in self._optimizers: - if not acc_opt.gradient_state.is_xla_gradients_synced: - opt = acc_opt - while isinstance(opt, AcceleratedOptimizer): - opt = opt.optimizer - gradients = xm._fetch_gradients(opt) - # Use xm.all_reduce to perform an in-place all-reduce. Recusrsive all-reduce each tensor - # one by one in self.reduce is non-inplace. - xm.all_reduce("sum", gradients, scale=1.0 / self.num_processes) - # Set is_xla_gradients_synced to True to avoid all-reduce twice in the AcceleratedOptimizer step. - acc_opt.gradient_state.is_xla_gradients_synced = True - if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true": - self.unscale_gradients() - parameters = [p for p in parameters] - for model in self._models: - if parameters == [p for p in model.parameters()]: - return model.clip_grad_norm_(max_norm, norm_type) - self.unscale_gradients() - return torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type=norm_type) - - def clip_grad_value_(self, parameters, clip_value): - """ - Should be used in place of `torch.nn.utils.clip_grad_value_`. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator(gradient_accumulation_steps=2) - >>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler) - - >>> for input, target in dataloader: - ... optimizer.zero_grad() - ... output = model(input) - ... loss = loss_func(output, target) - ... accelerator.backward(loss) - ... if accelerator.sync_gradients: - ... accelerator.clip_grad_value_(model.parameters(), clip_value) - ... optimizer.step() - ``` - """ - if self.distributed_type in [DistributedType.DEEPSPEED, DistributedType.FSDP]: - raise Exception("DeepSpeed and FSDP do not support `clip_grad_value_`. Use `clip_grad_norm_` instead.") - self.unscale_gradients() - torch.nn.utils.clip_grad_value_(parameters, clip_value) - - def gather(self, tensor): - """ - Gather the values in *tensor* across all processes and concatenate them on the first dimension. Useful to - regroup the predictions from all processes when doing evaluation. - - Note: - This gather happens in all processes. - - Args: - tensor (`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`): - The tensors to gather across all processes. - - Returns: - `torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`: The gathered tensor(s). Note that the - first dimension of the result is *num_processes* multiplied by the first dimension of the input tensors. - - Example: - - ```python - >>> # Assuming four processes - >>> import torch - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> process_tensor = torch.tensor([accelerator.process_index]) - >>> gathered_tensor = accelerator.gather(process_tensor) - >>> gathered_tensor - tensor([0, 1, 2, 3]) - ``` - """ - return gather(tensor) - - def gather_for_metrics(self, input_data, use_gather_object=False): - """ - Gathers `input_data` and potentially drops duplicates in the last batch if on a distributed system. Should be - used for gathering the inputs and targets for metric calculation. - - Args: - input (`torch.Tensor`, `object`, a nested tuple/list/dictionary of `torch.Tensor`, or a nested tuple/list/dictionary of `object`): - The tensors or objects for calculating metrics across all processes - use_gather_object(`bool`): - Whether to forcibly use gather_object instead of gather (which is already done if all objects passed do - not contain tensors). This flag can be useful for gathering tensors with different sizes that we don't - want to pad and concatenate along the first dimension. Using it with GPU tensors is not well supported - and inefficient as it incurs GPU -> CPU transfer since tensors would be pickled. - - Example: - - ```python - >>> # Assuming two processes, with a batch size of 5 on a dataset with 9 samples - >>> import torch - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> dataloader = torch.utils.data.DataLoader(range(9), batch_size=5) - >>> dataloader = accelerator.prepare(dataloader) - >>> batch = next(iter(dataloader)) - >>> gathered_items = accelerator.gather_for_metrics(batch) - >>> len(gathered_items) - 9 - ``` - """ - - try: - recursively_apply(lambda x: x, input_data, error_on_other_type=True) - all_tensors = True - except TypeError: - all_tensors = False - - use_gather_object = use_gather_object or not all_tensors - - if use_gather_object: - data = gather_object(input_data) - else: - data = self.gather(input_data) - - try: - if self.gradient_state.end_of_dataloader: - # at the end of a dataloader, `gather_for_metrics` regresses to - # `gather` unless the dataset has a remainder so log. - if self.gradient_state.remainder == -1: - logger.info( - "The used dataset had no length, returning gathered tensors. You should drop the remainder yourself." - ) - return data - elif self.gradient_state.remainder > 0: - # Last batch needs to be truncated on distributed systems as it contains additional samples - def _adjust_samples(tensor): - return tensor[: self.gradient_state.remainder] - - if use_gather_object: - # gather_object put the objects in a list - return _adjust_samples(data) - else: - return recursively_apply(_adjust_samples, data) - else: # remainder is 0 - # no remainder even though at end of dataloader, so nothing to do. - return data - else: - # Not at the end of the dataloader, no need to adjust the tensors - return data - except Exception: - # Dataset had no length or raised an error - return data - - def reduce(self, tensor, reduction="sum", scale=1.0): - """ - Reduce the values in *tensor* across all processes based on *reduction*. - - Note: - All processes get the reduced value. - - Args: - tensor (`torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`): - The tensors to reduce across all processes. - reduction (`str`, *optional*, defaults to "sum"): - A reduction type, can be one of 'sum', 'mean', or 'none'. If 'none', will not perform any operation. - scale (`float`, *optional*, defaults to 1.0): - A default scaling value to be applied after the reduce, only valied on XLA. - - Returns: - `torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`: - The reduced tensor(s). - - Example: - - ```python - >>> # Assuming two processes - >>> import torch - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> process_tensor = torch.arange(accelerator.num_processes) + 1 + (2 * accelerator.process_index) - >>> process_tensor = process_tensor.to(accelerator.device) - >>> reduced_tensor = accelerator.reduce(process_tensor, reduction="sum") - >>> reduced_tensor - tensor([4, 6]) - ``` - """ - return reduce(tensor, reduction, scale) - - def pad_across_processes(self, tensor, dim=0, pad_index=0, pad_first=False): - """ - Recursively pad the tensors in a nested list/tuple/dictionary of tensors from all devices to the same size so - they can safely be gathered. - - Args: - tensor (nested list/tuple/dictionary of `torch.Tensor`): - The data to gather. - dim (`int`, *optional*, defaults to 0): - The dimension on which to pad. - pad_index (`int`, *optional*, defaults to 0): - The value with which to pad. - pad_first (`bool`, *optional*, defaults to `False`): - Whether to pad at the beginning or the end. - - Returns: - `torch.Tensor`, or a nested tuple/list/dictionary of `torch.Tensor`: - The padded tensor(s). - - Example: - - ```python - >>> # Assuming two processes, with the first processes having a tensor of size 1 and the second of size 2 - >>> import torch - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> process_tensor = torch.arange(accelerator.process_index + 1).to(accelerator.device) - >>> padded_tensor = accelerator.pad_across_processes(process_tensor) - >>> padded_tensor.shape - torch.Size([2]) - ``` - """ - return pad_across_processes(tensor, dim=dim, pad_index=pad_index, pad_first=pad_first) - - def unwrap_model(self, model, keep_fp32_wrapper: bool = True, keep_torch_compile: bool = True): - """ - Unwraps the `model` from the additional layer possible added by [`~Accelerator.prepare`]. Useful before saving - the model. - - Args: - model (`torch.nn.Module`): - The model to unwrap. - keep_fp32_wrapper (`bool`, *optional*, defaults to `True`): - Whether to not remove the mixed precision hook if it was added. - keep_torch_compile (`bool`, *optional*, defaults to `True`): - Whether to not unwrap compiled model if compiled. - Returns: - `torch.nn.Module`: The unwrapped model. - - Example: - - ```python - >>> # Assuming two GPU processes - >>> from torch.nn.parallel import DistributedDataParallel - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> model = accelerator.prepare(MyModel()) - >>> print(model.__class__.__name__) - DistributedDataParallel - - >>> model = accelerator.unwrap_model(model) - >>> print(model.__class__.__name__) - MyModel - ``` - """ - return extract_model_from_parallel(model, keep_fp32_wrapper, keep_torch_compile) - - def wait_for_everyone(self): - """ - Will stop the execution of the current process until every other process has reached that point (so this does - nothing when the script is only run in one process). Useful to do before saving a model. - - Example: - - ```python - >>> # Assuming two GPU processes - >>> import time - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> if accelerator.is_main_process: - ... time.sleep(2) - >>> else: - ... print("I'm waiting for the main process to finish its sleep...") - >>> accelerator.wait_for_everyone() - >>> # Should print on every process at the same time - >>> print("Everyone is here") - ``` - """ - wait_for_everyone() - - @on_main_process - def init_trackers(self, project_name: str, config: dict | None = None, init_kwargs: dict | None = {}): - """ - Initializes a run for all trackers stored in `self.log_with`, potentially with starting configurations - - Args: - project_name (`str`): - The name of the project. All trackers will save their data based on this - config (`dict`, *optional*): - Optional starting configuration to be logged. - init_kwargs (`dict`, *optional*): - A nested dictionary of kwargs to be passed to a specific tracker's `__init__` function. Should be - formatted like so: - ```python - {"wandb": {"tags": ["tag_a", "tag_b"]}} - ``` - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator(log_with="tensorboard") - >>> accelerator.init_trackers( - ... project_name="my_project", - ... config={"learning_rate": 0.001, "batch_size": 32}, - ... init_kwargs={"tensorboard": {"flush_secs": 60}}, - ... ) - ``` - """ - for tracker in self.log_with: - if issubclass(type(tracker), GeneralTracker): - # Custom trackers are already initialized - self.trackers.append(tracker) - else: - tracker_init = LOGGER_TYPE_TO_CLASS[str(tracker)] - if tracker_init.requires_logging_directory: - # We can skip this check since it was done in `__init__` - self.trackers.append( - tracker_init(project_name, self.logging_dir, **init_kwargs.get(str(tracker), {})) - ) - else: - self.trackers.append(tracker_init(project_name, **init_kwargs.get(str(tracker), {}))) - - for tracker in self.trackers: - tracker.start() - - if config is not None: - for tracker in self.trackers: - tracker.store_init_configuration(config) - - def get_tracker(self, name: str, unwrap: bool = False): - """ - Returns a `tracker` from `self.trackers` based on `name` on the main process only. - - Args: - name (`str`): - The name of a tracker, corresponding to the `.name` property. - unwrap (`bool`): - Whether to return the internal tracking mechanism or to return the wrapped tracker instead - (recommended). - - Returns: - `GeneralTracker`: The tracker corresponding to `name` if it exists. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator(log_with="tensorboard") - >>> accelerator.init_trackers("my_project") - >>> tensorboard_tracker = accelerator.get_tracker("tensorboard") - ``` - """ - if len(self.trackers) > 0: - for tracker in self.trackers: - if tracker.name == name: - return tracker.tracker if unwrap else tracker - raise ValueError(f"{name} is not an available tracker stored inside the `Accelerator`.") - # Handle tracker only made on main process - return GeneralTracker(_blank=True) - - @on_main_process - def log(self, values: dict, step: int | None = None, log_kwargs: dict | None = {}): - """ - Logs `values` to all stored trackers in `self.trackers` on the main process only. - - Args: - values (`dict`): - Values should be a dictionary-like object containing only types `int`, `float`, or `str`. - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - log_kwargs (`dict`, *optional*): - A nested dictionary of kwargs to be passed to a specific tracker's `log` function. Should be formatted - like so: - ```python - {"wandb": {"tags": ["tag_a", "tag_b"]}} - ``` - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator(log_with="tensorboard") - >>> accelerator.init_trackers("my_project") - >>> accelerator.log({"loss": 0.5, "accuracy": 0.9}) - ``` - """ - for tracker in self.trackers: - tracker.log(values, step=step, **log_kwargs.get(tracker.name, {})) - - def end_training(self): - """ - Runs any special end training behaviors, such as stopping trackers on the main process only or destoying - process group. Should always be called at the end of your script if using experiment tracking. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator(log_with="tensorboard") - >>> accelerator.init_trackers("my_project") - >>> # Do training - >>> accelerator.end_training() - ``` - """ - for tracker in self.trackers: - tracker.finish() - - self.state.destroy_process_group() - - def save(self, obj, f, safe_serialization=False): - """ - Save the object passed to disk once per machine. Use in place of `torch.save`. - - Args: - obj (`object`): The object to save. - f (`str` or `os.PathLike`): Where to save the content of `obj`. - safe_serialization (`bool`, *optional*, defaults to `False`): Whether to save `obj` using `safetensors` - - Note: - If `save_on_each_node` was passed in as a `ProjectConfiguration`, will save the object once per node, - rather than only once on the main node. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> arr = [0, 1, 2, 3] - >>> accelerator.save(arr, "array.pkl") - ``` - """ - save( - obj, - f, - save_on_each_node=self.project_configuration.save_on_each_node, - safe_serialization=safe_serialization, - ) - - def save_model( - self, - model: torch.nn.Module, - save_directory: Union[str, os.PathLike], - max_shard_size: Union[int, str] = "10GB", - safe_serialization: bool = True, - ): - """ - Save a model so that it can be re-loaded using load_checkpoint_in_model - - Arguments: - model: (`torch.nn.Module`): - Model to be saved. The model can be wrapped or unwraped. - save_directory (`str` or `os.PathLike`): - Directory to which to save. Will be created if it doesn't exist. - max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`): - 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 `"5MB"`). - - - - If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard - which will be bigger than `max_shard_size`. - - - - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> model = ... - >>> accelerator.save_model(model, save_directory) - ``` - """ - - if os.path.isfile(save_directory): - logger.error(f"Provided path ({save_directory}) should be a directory, not a file") - return - - # get the state_dict of the model - if any(has_offloaded_params(module) for module in model.modules()): - state_dict = get_state_dict_offloaded_model(model) - else: - if any(param.device == torch.device("meta") for param in model.parameters()): - raise RuntimeError("You can't save the model since some parameters are on the meta device.") - state_dict = self.get_state_dict(model) - - # Case: DeepSpeed zero3 gets gathered and `state_dict` is empty - if state_dict is None: - return - os.makedirs(save_directory, exist_ok=True) - - if safe_serialization: - state_dict = clean_state_dict_for_safetensors(state_dict) - weights_name = SAFE_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME - filename_pattern = SAFE_WEIGHTS_PATTERN_NAME if safe_serialization else WEIGHTS_PATTERN_NAME - - state_dict_split = split_torch_state_dict_into_shards( - state_dict, filename_pattern=filename_pattern, max_shard_size=max_shard_size - ) - - # Clean the folder from a previous save - for filename in os.listdir(save_directory): - full_filename = os.path.join(save_directory, filename) - # If we have a shard file that is not going to be replaced, we delete it, but only from the main process - # in distributed settings to avoid race conditions. - weights_no_suffix = weights_name.replace(".bin", "") - - # make sure that file to be deleted matches format of sharded file, e.g. pytorch_model-00001-of-00005 - filename_no_suffix = filename.replace(".bin", "") - reg = re.compile(r"(.*?)-\d{5}-of-\d{5}") - - if ( - filename.startswith(weights_no_suffix) - and os.path.isfile(full_filename) - and filename not in state_dict_split.filename_to_tensors.keys() - and reg.fullmatch(filename_no_suffix) is not None - and PartialState().is_main_process - ): - os.remove(full_filename) - - # Save the model - for filename, tensors in state_dict_split.filename_to_tensors.items(): - shard = {tensor: state_dict[tensor] for tensor in tensors} - self.save(shard, os.path.join(save_directory, filename), safe_serialization=safe_serialization) - - # Save index if sharded - if state_dict_split.is_sharded: - index = { - "metadata": state_dict_split.metadata, - "weight_map": state_dict_split.tensor_to_filename, - } - save_index_file = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME - save_index_file = os.path.join(save_directory, save_index_file) - with open(save_index_file, "w", encoding="utf-8") as f: - content = json.dumps(index, indent=2, sort_keys=True) + "\n" - f.write(content) - logger.info( - f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be " - f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the " - f"index located at {save_index_file}." - ) - else: - path_to_weights = os.path.join(save_directory, WEIGHTS_NAME) - logger.info(f"Model weights saved in {path_to_weights}") - - def register_save_state_pre_hook(self, hook: Callable[..., None]) -> hooks.RemovableHandle: - """ - Registers a pre hook to be run before `save_checkpoint` is called in [`Accelerator.save_state`]. - - Args: - hook (`Callable`): - A function to be called in [`Accelerator.save_state`] before `save_checkpoint`. - - The hook should have the following signature: - - `hook(models: list[torch.nn.Module], weights: list[dict[str, torch.Tensor]], input_dir: str) -> None` - - The `models` argument are the models as saved in the accelerator state under `accelerator._models`, `weigths` - argument are the state dicts of the `models`, and the `input_dir` argument is the `input_dir` argument passed - to [`Accelerator.load_state`]. - - - - Should only be used in conjunction with [`Accelerator.register_load_state_pre_hook`]. Can be useful to save - configurations in addition to model weights. Can also be used to overwrite model saving with a customized - method. In this case, make sure to remove already loaded weights from the weights list. - - - - Returns: - `torch.utils.hooks.RemovableHandle`: a handle that can be used to remove the added hook by calling - `handle.remove()` - """ - handle = hooks.RemovableHandle(self._save_model_state_pre_hook) - self._save_model_state_pre_hook[handle.id] = hook - return handle - - def save_state(self, output_dir: str = None, safe_serialization: bool = True, **save_model_func_kwargs): - """ - Saves the current states of the model, optimizer, scaler, RNG generators, and registered objects to a folder. - - If a `ProjectConfiguration` was passed to the `Accelerator` object with `automatic_checkpoint_naming` enabled - then checkpoints will be saved to `self.project_dir/checkpoints`. If the number of current saves is greater - than `total_limit` then the oldest save is deleted. Each checkpoint is saved in separate folders named - `checkpoint_`. - - Otherwise they are just saved to `output_dir`. - - - - Should only be used when wanting to save a checkpoint during training and restoring the state in the same - environment. - - - - Args: - output_dir (`str` or `os.PathLike`): - The name of the folder to save all relevant weights and states. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). - save_model_func_kwargs (`dict`, *optional*): - Additional keyword arguments for saving model which can be passed to the underlying save function, such - as optional arguments for DeepSpeed's `save_checkpoint` function. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> model, optimizer, lr_scheduler = ... - >>> model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler) - >>> accelerator.save_state(output_dir="my_checkpoint") - ``` - """ - if self.project_configuration.automatic_checkpoint_naming: - output_dir = os.path.join(self.project_dir, "checkpoints") - os.makedirs(output_dir, exist_ok=True) - if self.project_configuration.automatic_checkpoint_naming: - folders = [os.path.join(output_dir, folder) for folder in os.listdir(output_dir)] - if ( - self.project_configuration.total_limit is not None - and (len(folders) + 1 > self.project_configuration.total_limit) - and self.is_main_process - ): - - def _inner(folder): - return list(map(int, re.findall(r"[\/]?([0-9]+)(?=[^\/]*$)", folder)))[0] - - folders.sort(key=_inner) - logger.warning( - f"Deleting {len(folders) + 1 - self.project_configuration.total_limit} checkpoints to make room for new checkpoint." - ) - for folder in folders[: len(folders) + 1 - self.project_configuration.total_limit]: - shutil.rmtree(folder) - output_dir = os.path.join(output_dir, f"checkpoint_{self.save_iteration}") - if os.path.exists(output_dir): - raise ValueError( - f"Checkpoint directory {output_dir} ({self.save_iteration}) already exists. Please manually override `self.save_iteration` with what iteration to start with." - ) - self.wait_for_everyone() - os.makedirs(output_dir, exist_ok=True) - logger.info(f"Saving current state to {output_dir}") - - if self.distributed_type == DistributedType.XLA: - # Finish running the previous step before checkpointing - xm.mark_step() - - # Save the models taking care of FSDP and DeepSpeed nuances - weights = [] - for i, model in enumerate(self._models): - if self.distributed_type == DistributedType.FSDP: - logger.info("Saving FSDP model") - save_fsdp_model(self.state.fsdp_plugin, self, model, output_dir, i) - logger.info(f"FSDP Model saved to output dir {output_dir}") - elif self.distributed_type == DistributedType.DEEPSPEED: - logger.info("Saving DeepSpeed Model and Optimizer") - ckpt_id = f"{MODEL_NAME}" if i == 0 else f"{MODEL_NAME}_{i}" - model.save_checkpoint(output_dir, ckpt_id, **save_model_func_kwargs) - logger.info(f"DeepSpeed Model and Optimizer saved to output dir {os.path.join(output_dir, ckpt_id)}") - elif self.distributed_type == DistributedType.MEGATRON_LM: - logger.info("Saving Megatron-LM Model, Optimizer and Scheduler") - model.save_checkpoint(output_dir) - logger.info(f"Megatron-LM Model , Optimizer and Scheduler saved to output dir {output_dir}") - else: - weights.append(self.get_state_dict(model, unwrap=False)) - - # Save the optimizers taking care of FSDP and DeepSpeed nuances - optimizers = [] - if self.distributed_type == DistributedType.FSDP: - for i, opt in enumerate(self._optimizers): - logger.info("Saving FSDP Optimizer") - save_fsdp_optimizer(self.state.fsdp_plugin, self, opt, self._models[i], output_dir, i) - logger.info(f"FSDP Optimizer saved to output dir {output_dir}") - elif self.distributed_type not in [DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM]: - optimizers = self._optimizers - - # Save the lr schedulers taking care of DeepSpeed nuances - schedulers = [] - if self.distributed_type == DistributedType.DEEPSPEED: - for i, scheduler in enumerate(self._schedulers): - if isinstance(scheduler, DeepSpeedSchedulerWrapper): - continue - schedulers.append(scheduler) - elif self.distributed_type not in [DistributedType.MEGATRON_LM]: - schedulers = self._schedulers - - # Save the samplers of the dataloaders - dataloaders = self._dataloaders - - # Call model loading hooks that might have been registered with - # accelerator.register_model_state_hook - for hook in self._save_model_state_pre_hook.values(): - hook(self._models, weights, output_dir) - - save_location = save_accelerator_state( - output_dir, - weights, - optimizers, - schedulers, - dataloaders, - self.state.process_index, - self.step, - self.scaler, - save_on_each_node=self.project_configuration.save_on_each_node, - safe_serialization=safe_serialization, - ) - for i, obj in enumerate(self._custom_objects): - save_custom_state(obj, output_dir, i, save_on_each_node=self.project_configuration.save_on_each_node) - self.project_configuration.iteration += 1 - return save_location - - def register_load_state_pre_hook(self, hook: Callable[..., None]) -> hooks.RemovableHandle: - """ - Registers a pre hook to be run before [`load_checkpoint`] is called in [`Accelerator.load_state`]. - - Args: - hook (`Callable`): - A function to be called in [`Accelerator.load_state`] before `load_checkpoint`. - - The hook should have the following signature: - - `hook(models: list[torch.nn.Module], input_dir: str) -> None` - - The `models` argument are the models as saved in the accelerator state under `accelerator._models`, and the - `input_dir` argument is the `input_dir` argument passed to [`Accelerator.load_state`]. - - - - Should only be used in conjunction with [`Accelerator.register_save_state_pre_hook`]. Can be useful to load - configurations in addition to model weights. Can also be used to overwrite model loading with a customized - method. In this case, make sure to remove already loaded models from the models list. - - - - Returns: - `torch.utils.hooks.RemovableHandle`: a handle that can be used to remove the added hook by calling - `handle.remove()` - """ - handle = hooks.RemovableHandle(self._load_model_state_pre_hook) - self._load_model_state_pre_hook[handle.id] = hook - return handle - - def load_state(self, input_dir: str = None, load_kwargs: dict | None = None, **load_model_func_kwargs): - """ - Loads the current states of the model, optimizer, scaler, RNG generators, and registered objects. - - - - Should only be used in conjunction with [`Accelerator.save_state`]. If a file is not registered for - checkpointing, it will not be loaded if stored in the directory. - - - - Args: - input_dir (`str` or `os.PathLike`): - The name of the folder all relevant weights and states were saved in. Can be `None` if - `automatic_checkpoint_naming` is used, and will pick up from the latest checkpoint. - load_kwargs (`dict`, *optional*): - Additional keyword arguments for the underlying `load` function, such as optional arguments for - state_dict and optimizer on. - load_model_func_kwargs (`dict`, *optional*): - Additional keyword arguments for loading model which can be passed to the underlying load function, - such as optional arguments for DeepSpeed's `load_checkpoint` function or a `map_location` to load the - model and optimizer on. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> model, optimizer, lr_scheduler = ... - >>> model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler) - >>> accelerator.load_state("my_checkpoint") - ``` - """ - if input_dir is not None: - # Check if folder exists - input_dir = os.path.expanduser(input_dir) - if not os.path.isdir(input_dir): - raise ValueError(f"Tried to find {input_dir} but folder does not exist") - elif self.project_configuration.automatic_checkpoint_naming: - # Pick up from automatic checkpoint naming - input_dir = os.path.join(self.project_dir, "checkpoints") - folders = [os.path.join(input_dir, folder) for folder in os.listdir(input_dir)] - - def _inner(folder): - return list(map(int, re.findall(r"[\/]?([0-9]+)(?=[^\/]*$)", folder)))[0] - - folders.sort(key=_inner) - input_dir = folders[-1] - else: - raise ValueError("No input_dir provided and automatic checkpoint naming is disabled.") - logger.info(f"Loading states from {input_dir}") - - # Load the models taking care of FSDP and DeepSpeed nuances - models = [] - for i, model in enumerate(self._models): - if self.distributed_type == DistributedType.FSDP: - logger.info("Loading FSDP model") - load_fsdp_model(self.state.fsdp_plugin, self, model, input_dir, i) - logger.info(f"FSDP Model loaded from input dir {input_dir}") - elif self.distributed_type == DistributedType.DEEPSPEED: - logger.info("Loading DeepSpeed Model and Optimizer") - ckpt_id = f"{MODEL_NAME}" if i == 0 else f"{MODEL_NAME}_{i}" - model.load_checkpoint(input_dir, ckpt_id, **load_model_func_kwargs) - logger.info(f"DeepSpeed Model and Optimizer loaded from input dir {os.path.join(input_dir, ckpt_id)}") - elif self.distributed_type == DistributedType.MEGATRON_LM: - logger.info("Loading Megatron-LM Model, Optimizer and Scheduler") - model.load_checkpoint(input_dir) - logger.info(f"Megatron-LM Model , Optimizer and Scheduler loaded from input dir {input_dir}") - else: - models.append(model) - - # We need to load the scaler state before the optimizer for FSDP2 - # (`torch.distributed.checkpoint.set_optimizer_state_dict`) which we use to set the state of the optimizer calls `optimizer.step` on - # a dummy tensor, but since the scaler is not initialized, it will raise an error (the scaler exists but its `_scale` is None) - scaler = None - if self.scaler is not None and self.is_fsdp2: - input_scaler_file = os.path.join(input_dir, SCALER_NAME) - scaler_state = torch.load(input_scaler_file) - self.scaler.load_state_dict(scaler_state) - # We also need to call the `_lazy_init_scale_growth_tracker` to initialize the scaler, as it would else be called - # on the first call to scale - self.scaler._lazy_init_scale_growth_tracker(self.scaler._device) - logger.info("GradScaler state loaded successfully") - else: - scaler = self.scaler - - # Load the optimizers taking care of FSDP and DeepSpeed nuances - optimizers = [] - if self.distributed_type == DistributedType.FSDP: - for i, opt in enumerate(self._optimizers): - logger.info("Loading FSDP Optimizer") - load_fsdp_optimizer(self.state.fsdp_plugin, self, opt, self._models[i], input_dir, i) - logger.info(f"FSDP Optimizer loaded from input dir {input_dir}") - elif self.distributed_type not in [DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM]: - optimizers = self._optimizers - - # Load the lr schedulers taking care of DeepSpeed nuances - schedulers = [] - if self.distributed_type == DistributedType.DEEPSPEED: - for i, scheduler in enumerate(self._schedulers): - if isinstance(scheduler, DeepSpeedSchedulerWrapper): - continue - schedulers.append(scheduler) - elif self.distributed_type not in [DistributedType.MEGATRON_LM]: - schedulers = self._schedulers - - dataloaders = self._dataloaders - - # Call model loading hooks that might have been registered with - # accelerator.register_model_state_hook - for hook in self._load_model_state_pre_hook.values(): - hook(models, input_dir) - - map_location = load_model_func_kwargs.pop("map_location", None) - if map_location is None: - if self.num_processes > 1 and self.distributed_type in ( - DistributedType.MULTI_GPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_NPU, - DistributedType.MULTI_HPU, - ): - map_location = "on_device" - else: - map_location = "cpu" - - override_attributes = load_accelerator_state( - input_dir, - models, - optimizers, - schedulers, - dataloaders, - self.state.process_index, - scaler, - map_location, - load_kwargs, - **load_model_func_kwargs, - ) - if "step" in override_attributes: - self.step = override_attributes["step"] - custom_checkpoints = [ - f for f in os.listdir(input_dir) if re.search(r"^custom_checkpoint_\d+\.pkl$", f) is not None - ] - if len(custom_checkpoints) != len(self._custom_objects): - err = ( - f"Number of custom checkpoints in folder {input_dir} does not match the number of registered objects:" - ) - err += f"\n\tFound checkpoints: {len(custom_checkpoints)}" - err += f"\n\tRegistered objects: {len(self._custom_objects)}\n" - err += "Please make sure to only load checkpoints from folders that were created with the same set of registered objects," - err += "or avoid using `custom_checkpoint` in the filename for files in that same directory and load them in manually." - raise RuntimeError(err) - else: - logger.info(f"Loading in {len(custom_checkpoints)} custom states") - for index, obj in enumerate(self._custom_objects): - load_custom_state(obj, input_dir, index) - - def free_memory(self, *objects): - """ - Will release all references to the internal objects stored and call the garbage collector. You should call this - method between two trainings with different models/optimizers. Also will reset `Accelerator.step` to 0. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> model, optimizer, scheduler = ... - >>> model, optimizer, scheduler = accelerator.prepare(model, optimizer, scheduler) - >>> model, optimizer, scheduler = accelerator.free_memory(model, optimizer, scheduler) - ``` - """ - # Deepspeed needs a bit more prep that should be done first - if hasattr(self, "deepspeed_engine_wrapped"): - if self.deepspeed_engine_wrapped is not None: - self.deepspeed_engine_wrapped.engine.destroy() - self.deepspeed_engine_wrapped = None - objects = release_memory(*objects) - self._schedulers = [] - self._optimizers = [] - self._models = [] - self._dataloaders = [] - self.step = 0 - return objects - - def clear(self, *objects): - """ - Alias for [`Accelerate.free_memory`], releases all references to the internal objects stored and call the - garbage collector. You should call this method between two trainings with different models/optimizers. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> model, optimizer, scheduler = ... - >>> model, optimizer, scheduler = accelerator.prepare(model, optimizer, scheduler) - >>> model, optimizer, scheduler = accelerator.clear(model, optimizer, scheduler) - ``` - """ - return self.free_memory(*objects) - - def _get_named_parameters(self, *args, drop_refs=False): - named_parameters = {} - accessor_mapping = {} - for obj in args: - if isinstance(obj, torch.nn.Module): - obj = extract_model_from_parallel(obj) - if not drop_refs: - named_parameters.update({n: p for n, p in obj.named_parameters()}) - continue - - # we need this bit as `WeightWithDynamic...` returns 0 when `data_ptr()` is called, - # the underlying pointer is actually hidden in `_tensor` attribute - if self.fp8_backend == FP8BackendType.AO: - from torchao.float8.fsdp_utils import WeightWithDynamicFloat8CastTensor - - accessor_mapping[WeightWithDynamicFloat8CastTensor] = "_tensor" - - named_parameters.update( - { - n: getattr(p, accessor_mapping[type(p)]).data_ptr() - if type(p) in accessor_mapping - else p.data_ptr() - for n, p in obj.named_parameters() - } - ) - return named_parameters - - def _get_devices(self, *args): - model_device = None - optimizer_device = None - for obj in args: - # Loop through model parameters and stop at the first once we have its device. - if isinstance(obj, torch.nn.Module): - for param in obj.parameters(): - model_device = param.device - break - # Loop through optimizer parameters groups and stop at the first once we have its device. - if isinstance(obj, torch.optim.Optimizer): - for param_group in obj.param_groups: - if len(param_group["params"]) > 0: - optimizer_device = param_group["params"][0].device - break - return (model_device, optimizer_device) - - def get_state_dict(self, model, unwrap=True): - """ - Returns the state dictionary of a model sent through [`Accelerator.prepare`] potentially without full - precision. - - Args: - model (`torch.nn.Module`): - A PyTorch model sent through [`Accelerator.prepare`] - unwrap (`bool`, *optional*, defaults to `True`): - Whether to return the original underlying state_dict of `model` or to return the wrapped state_dict - - Returns: - `dict`: The state dictionary of the model potentially without full precision. - - Example: - - ```python - >>> import torch - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> net = torch.nn.Linear(2, 2) - >>> net = accelerator.prepare(net) - >>> state_dict = accelerator.get_state_dict(net) - ``` - """ - - if self.distributed_type == DistributedType.DEEPSPEED: - zero3_sharding = self.deepspeed_config["zero_optimization"]["stage"] == 3 - tp_sharding = self.deepspeed_config.get("tensor_parallel", {}).get("autotp_size", 0) > 1 - if zero3_sharding or tp_sharding: - if model.zero_gather_16bit_weights_on_model_save(): - if tp_sharding and not compare_versions("deepspeed", ">=", "0.16.4"): - raise ImportError( - "Deepspeed TP requires deepspeed >= 0.16.4, Please update DeepSpeed via `pip install deepspeed -U`." - ) - state_dict = ( - model._consolidated_16bit_state_dict() - if tp_sharding - else model._zero3_consolidated_16bit_state_dict() - ) - else: - raise ValueError( - "Cannot get 16bit model weights because `stage3_gather_16bit_weights_on_model_save` in DeepSpeed config is False. " - "To save the model weights in 16bit, set `stage3_gather_16bit_weights_on_model_save` to True in DeepSpeed config file or " - "set `zero3_save_16bit_model` to True when using `accelerate config`. " - "To save the full checkpoint, run `model.save_checkpoint(save_dir)` and use `zero_to_fp32.py` to recover weights." - ) - else: - from deepspeed.checkpoint.utils import clone_tensors_for_torch_save - - state_dict = clone_tensors_for_torch_save(self.unwrap_model(model).state_dict()) - elif self.is_fsdp2: - from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict - - options = StateDictOptions(full_state_dict=True, broadcast_from_rank0=True, cpu_offload=True) - state_dict = get_model_state_dict(model, options=options) - elif self.distributed_type == DistributedType.FSDP: - from torch.distributed.fsdp import FullStateDictConfig, StateDictType - from torch.distributed.fsdp import FullyShardedDataParallel as FSDP - - full_state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) - with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, full_state_dict_config): - state_dict = model.state_dict() - else: - if unwrap: - model = self.unwrap_model(model) - state_dict = model.state_dict() - - return state_dict - - def register_for_checkpointing(self, *objects): - """ - Makes note of `objects` and will save or load them in during `save_state` or `load_state`. - - These should be utilized when the state is being loaded or saved in the same script. It is not designed to be - used in different scripts. - - - - Every `object` must have a `load_state_dict` and `state_dict` function to be stored. - - - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> # Assume `CustomObject` has a `state_dict` and `load_state_dict` function. - >>> obj = CustomObject() - >>> accelerator.register_for_checkpointing(obj) - >>> accelerator.save_state("checkpoint.pt") - ``` - """ - invalid_objects = [] - for obj in objects: - if not hasattr(obj, "state_dict") or not hasattr(obj, "load_state_dict"): - invalid_objects.append(obj) - if len(invalid_objects) > 0: - err = "All `objects` must include a `state_dict` and `load_state_dict` function to be stored. The following inputs are invalid:" - for index, obj in enumerate(invalid_objects): - err += f"\n\t- Item at index {index}, `{get_pretty_name(obj)}`" - raise ValueError(err) - self._custom_objects.extend(objects) - - @contextmanager - def autocast(self, autocast_handler: AutocastKwargs = None): - """ - Will apply automatic mixed-precision inside the block inside this context manager, if it is enabled. Nothing - different will happen otherwise. - - A different `autocast_handler` can be passed in to override the one set in the `Accelerator` object. This is - useful in blocks under `autocast` where you want to revert to fp32. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator(mixed_precision="fp16") - >>> with accelerator.autocast(): - ... train() - ``` - """ - if autocast_handler is None: - autocast_handler = self.autocast_handler - autocast_context = get_mixed_precision_context_manager(self.native_amp, autocast_handler) - autocast_context.__enter__() - # TODO: should the `yield` be in a try/finally block? - yield - autocast_context.__exit__(*sys.exc_info()) - - @contextmanager - def profile(self, profile_handler: ProfileKwargs | None = None): - """ - Will profile the code inside the context manager. The profile will be saved to a Chrome Trace file if - `profile_handler.output_trace_dir` is set. - - A different `profile_handler` can be passed in to override the one set in the `Accelerator` object. - - Args: - profile_handler (`ProfileKwargs`, *optional*): - The profile handler to use for this context manager. If not passed, will use the one set in the - `Accelerator` object. - - Example: - - ```python - # Profile with default settings - from accelerate import Accelerator - from accelerate.utils import ProfileKwargs - - accelerator = Accelerator() - with accelerator.profile() as prof: - train() - accelerator.print(prof.key_averages().table()) - - - # Profile with the custom handler - def custom_handler(prof): - print(prof.key_averages().table(sort_by="self_cpu_time_total", row_limit=10)) - - - kwargs = ProfileKwargs(schedule_option=dict(wait=1, warmup=1, active=1), on_trace_ready=custom_handler) - accelerator = Accelerator(kwarg_handler=[kwargs]) - with accelerator.profile() as prof: - for _ in range(10): - train_iteration() - prof.step() - - - # Profile and export to Chrome Trace - kwargs = ProfileKwargs(output_trace_dir="output_trace") - accelerator = Accelerator(kwarg_handler=[kwargs]) - with accelerator.profile(): - train() - ``` - """ - profile_handler = profile_handler or self.profile_handler or ProfileKwargs() - - with profile_handler.build() as profiler: - yield profiler - - if profile_handler.output_trace_dir is None: - return - - os.makedirs(profile_handler.output_trace_dir, exist_ok=True) - profiler.export_chrome_trace( - os.path.join(profile_handler.output_trace_dir, PROFILE_PATTERN_NAME.format(suffix=self.process_index)) - ) - self.wait_for_everyone() - - @property - def optimizer_step_was_skipped(self): - """ - Whether or not the optimizer update was skipped (because of gradient overflow in mixed precision), in which - case the learning rate should not be changed. - """ - for optimizer in self._optimizers: - if optimizer.step_was_skipped: - return True - return False - - def skip_first_batches(self, dataloader, num_batches: int = 0): - """ - Creates a new `torch.utils.data.DataLoader` that will efficiently skip the first `num_batches`. - - Args: - dataloader (`torch.utils.data.DataLoader`): The data loader in which to skip batches. - num_batches (`int`, *optional*, defaults to 0): The number of batches to skip - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler) - >>> skipped_dataloader = accelerator.skip_first_batches(dataloader, num_batches=2) - >>> # for the first epoch only - >>> for input, target in skipped_dataloader: - ... optimizer.zero_grad() - ... output = model(input) - ... loss = loss_func(output, target) - ... accelerator.backward(loss) - ... optimizer.step() - - >>> # subsequent epochs - >>> for input, target in dataloader: - ... optimizer.zero_grad() - ... ... - ``` - """ - return skip_first_batches(dataloader, num_batches=num_batches) - - def __deepcopy__(self, memo): - logger.info("Deep copying the `Accelerator` object, note that this will point to the same original object.") - return self - - def verify_device_map(self, model: torch.nn.Module) -> bool: - """ - Verifies that `model` has not been prepared with big model inference with a device-map resembling `auto`. - """ - # Checks if any of the child modules has the attribute `hf_device_map` and this map has more than one entry. - for m in model.modules(): - if hasattr(m, "hf_device_map") and len(m.hf_device_map) > 1: - return True - - return False - - def lomo_backward(self, loss: torch.Tensor, learning_rate: float) -> None: - """ - Runs backward pass on LOMO optimizers. - """ - if is_lomo_available(): - # We need to import locally to avoid circular imports since lomo imports stuff from - # transformers & accelerate - from lomo_optim import AdaLomo, Lomo - - if learning_rate is None: - raise ValueError("A learning rate must be passed in order to call backward pass with LOMO optimizers.") - - _backward_called = False - - for optimizer in self._optimizers: - if isinstance(optimizer.optimizer, (Lomo, AdaLomo)): - optimizer.optimizer.fused_backward(loss, learning_rate) - _backward_called = True - - if not _backward_called: - raise ValueError( - "Backward pass not properly called on LOMO optimizers. Are you sure you passed a LOMO optimizer in accelerator.prepare()?" - ) - - @property - def fp8_backend(self) -> FP8BackendType: - "Returns the configured backend for training in FP8" - if self.has_fp8_handler: - if self.fp8_recipe_handler is not None: - return FP8BackendType(self.fp8_recipe_handler.backend) - elif self.ao_recipe_handler is not None: - return FP8BackendType.AO - elif self.te_recipe_handler is not None: - return FP8BackendType.TE - elif self.msamp_recipe_handler is not None: - return FP8BackendType.MSAMP - elif self.state.deepspeed_plugin is not None and self.state.deepspeed_plugin.enable_msamp: - return FP8BackendType.MSAMP - - return FP8BackendType(parse_choice_from_env("ACCELERATE_FP8_BACKEND", "NO")) diff --git a/venv/lib/python3.10/site-packages/accelerate/big_modeling.py b/venv/lib/python3.10/site-packages/accelerate/big_modeling.py deleted file mode 100644 index bddcaa8a0cca214e08aa13b9b8650e7ff7c475a9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/big_modeling.py +++ /dev/null @@ -1,749 +0,0 @@ -# Copyright 2022 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 logging -import os -import re -from contextlib import contextmanager -from functools import wraps -from typing import Optional, Union - -import torch -import torch.nn as nn - -from .hooks import ( - AlignDevicesHook, - CpuOffload, - LayerwiseCastingHook, - UserCpuOffloadHook, - add_hook_to_module, - attach_align_device_hook, - attach_align_device_hook_on_blocks, -) -from .utils import ( - OffloadedWeightsLoader, - check_cuda_p2p_ib_support, - check_device_map, - extract_submodules_state_dict, - find_tied_parameters, - get_balanced_memory, - infer_auto_device_map, - is_bnb_available, - is_mlu_available, - is_musa_available, - is_npu_available, - is_sdaa_available, - is_xpu_available, - load_checkpoint_in_model, - offload_state_dict, - parse_flag_from_env, - retie_parameters, -) -from .utils.constants import SUPPORTED_PYTORCH_LAYERS_FOR_UPCASTING -from .utils.other import recursive_getattr - - -logger = logging.getLogger(__name__) - - -@contextmanager -def init_empty_weights(include_buffers: bool = None): - """ - A context manager under which models are initialized with all parameters on the meta device, therefore creating an - empty model. Useful when just initializing the model would blow the available RAM. - - Args: - include_buffers (`bool`, *optional*): - Whether or not to also put all buffers on the meta device while initializing. - - Example: - - ```python - import torch.nn as nn - from accelerate import init_empty_weights - - # Initialize a model with 100 billions parameters in no time and without using any RAM. - with init_empty_weights(): - tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)]) - ``` - - - - Any model created under this context manager has no weights. As such you can't do something like - `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`]. - Make sure to overwrite the default device_map param for [`load_checkpoint_and_dispatch`], otherwise dispatch is not - called. - - - """ - if include_buffers is None: - include_buffers = parse_flag_from_env("ACCELERATE_INIT_INCLUDE_BUFFERS", False) - with init_on_device(torch.device("meta"), include_buffers=include_buffers) as f: - yield f - - -@contextmanager -def init_on_device(device: torch.device, include_buffers: bool = None): - """ - A context manager under which models are initialized with all parameters on the specified device. - - Args: - device (`torch.device`): - Device to initialize all parameters on. - include_buffers (`bool`, *optional*): - Whether or not to also put all buffers on the meta device while initializing. - - Example: - - ```python - import torch.nn as nn - from accelerate import init_on_device - - with init_on_device(device=torch.device("cuda")): - tst = nn.Linear(100, 100) # on `cuda` device - ``` - """ - if include_buffers is None: - include_buffers = parse_flag_from_env("ACCELERATE_INIT_INCLUDE_BUFFERS", False) - - if include_buffers: - with device: - yield - return - - old_register_parameter = nn.Module.register_parameter - if include_buffers: - old_register_buffer = nn.Module.register_buffer - - def register_empty_parameter(module, name, param): - old_register_parameter(module, name, param) - if param is not None: - param_cls = type(module._parameters[name]) - kwargs = module._parameters[name].__dict__ - kwargs["requires_grad"] = param.requires_grad - module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs) - - def register_empty_buffer(module, name, buffer, persistent=True): - old_register_buffer(module, name, buffer, persistent=persistent) - if buffer is not None: - module._buffers[name] = module._buffers[name].to(device) - - # Patch tensor creation - if include_buffers: - tensor_constructors_to_patch = { - torch_function_name: getattr(torch, torch_function_name) - for torch_function_name in ["empty", "zeros", "ones", "full"] - } - else: - tensor_constructors_to_patch = {} - - def patch_tensor_constructor(fn): - def wrapper(*args, **kwargs): - kwargs["device"] = device - return fn(*args, **kwargs) - - return wrapper - - try: - nn.Module.register_parameter = register_empty_parameter - if include_buffers: - nn.Module.register_buffer = register_empty_buffer - for torch_function_name in tensor_constructors_to_patch.keys(): - setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name))) - yield - finally: - nn.Module.register_parameter = old_register_parameter - if include_buffers: - nn.Module.register_buffer = old_register_buffer - for torch_function_name, old_torch_function in tensor_constructors_to_patch.items(): - setattr(torch, torch_function_name, old_torch_function) - - -def cpu_offload( - model: nn.Module, - execution_device: Optional[torch.device] = None, - offload_buffers: bool = False, - state_dict: Optional[dict[str, torch.Tensor]] = None, - preload_module_classes: Optional[list[str]] = None, -): - """ - Activates full CPU offload for a model. As a result, all parameters of the model will be offloaded and only one - copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that - state dict and put on the execution device passed as they are needed, then offloaded again. - - Args: - model (`torch.nn.Module`): - The model to offload. - execution_device (`torch.device`, *optional*): - The device on which the forward pass of the model will be executed (should be a GPU). Will default to the - model first parameter device. - offload_buffers (`bool`, *optional*, defaults to `False`): - Whether or not to offload the buffers with the model parameters. - state_dict (`Dict[str, torch.Tensor]`, *optional*): - The state dict of the model that will be kept on CPU. - preload_module_classes (`List[str]`, *optional*): - A list of classes whose instances should load all their weights (even in the submodules) at the beginning - of the forward. This should only be used for classes that have submodules which are registered but not - called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, - `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. - """ - if execution_device is None: - execution_device = next(iter(model.parameters())).device - if state_dict is None: - state_dict = {n: p.to("cpu") for n, p in model.state_dict().items()} - - add_hook_to_module(model, AlignDevicesHook(io_same_device=True), append=True) - attach_align_device_hook( - model, - execution_device=execution_device, - offload=True, - offload_buffers=offload_buffers, - weights_map=state_dict, - preload_module_classes=preload_module_classes, - ) - - return model - - -def cpu_offload_with_hook( - model: torch.nn.Module, - execution_device: Optional[Union[int, str, torch.device]] = None, - prev_module_hook: Optional[UserCpuOffloadHook] = None, -): - """ - Offloads a model on the CPU and puts it back to an execution device when executed. The difference with - [`cpu_offload`] is that the model stays on the execution device after the forward and is only offloaded again when - the `offload` method of the returned `hook` is called. Useful for pipelines running a model in a loop. - - Args: - model (`torch.nn.Module`): - The model to offload. - execution_device(`str`, `int` or `torch.device`, *optional*): - The device on which the model should be executed. Will default to the MPS device if it's available, then - GPU 0 if there is a GPU, and finally to the CPU. - prev_module_hook (`UserCpuOffloadHook`, *optional*): - The hook sent back by this function for a previous model in the pipeline you are running. If passed, its - offload method will be called just before the forward of the model to which this hook is attached. - - Example: - - ```py - model_1, hook_1 = cpu_offload_with_hook(model_1, cuda_device) - model_2, hook_2 = cpu_offload_with_hook(model_2, cuda_device, prev_module_hook=hook_1) - model_3, hook_3 = cpu_offload_with_hook(model_3, cuda_device, prev_module_hook=hook_2) - - hid_1 = model_1(input) - for i in range(50): - # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop. - hid_2 = model_2(hid_1) - # model2 is offloaded to the CPU just before this forward. - hid_3 = model_3(hid_3) - - # For model3, you need to manually call the hook offload method. - hook_3.offload() - ``` - """ - hook = CpuOffload(execution_device=execution_device, prev_module_hook=prev_module_hook) - add_hook_to_module(model, hook, append=True) - user_hook = UserCpuOffloadHook(model, hook) - return model, user_hook - - -def disk_offload( - model: nn.Module, - offload_dir: Union[str, os.PathLike], - execution_device: Optional[torch.device] = None, - offload_buffers: bool = False, - preload_module_classes: Optional[list[str]] = None, -): - """ - Activates full disk offload for a model. As a result, all parameters of the model will be offloaded as - memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and - put on the execution device passed as they are needed, then offloaded again. - - Args: - model (`torch.nn.Module`): The model to offload. - offload_dir (`str` or `os.PathLike`): - The folder in which to offload the model weights (or where the model weights are already offloaded). - execution_device (`torch.device`, *optional*): - The device on which the forward pass of the model will be executed (should be a GPU). Will default to the - model's first parameter device. - offload_buffers (`bool`, *optional*, defaults to `False`): - Whether or not to offload the buffers with the model parameters. - preload_module_classes (`List[str]`, *optional*): - A list of classes whose instances should load all their weights (even in the submodules) at the beginning - of the forward. This should only be used for classes that have submodules which are registered but not - called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, - `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. - """ - if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")): - offload_state_dict(offload_dir, model.state_dict()) - if execution_device is None: - execution_device = next(iter(model.parameters())).device - weights_map = OffloadedWeightsLoader(save_folder=offload_dir) - - add_hook_to_module(model, AlignDevicesHook(io_same_device=True), append=True) - attach_align_device_hook( - model, - execution_device=execution_device, - offload=True, - offload_buffers=offload_buffers, - weights_map=weights_map, - preload_module_classes=preload_module_classes, - ) - - return model - - -def dispatch_model( - model: nn.Module, - device_map: dict[str, Union[str, int, torch.device]], - main_device: Optional[torch.device] = None, - state_dict: Optional[dict[str, torch.Tensor]] = None, - offload_dir: Optional[Union[str, os.PathLike]] = None, - offload_index: Optional[dict[str, str]] = None, - offload_buffers: bool = False, - skip_keys: Optional[Union[str, list[str]]] = None, - preload_module_classes: Optional[list[str]] = None, - force_hooks: bool = False, -): - """ - Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on - the CPU or even the disk. - - Args: - model (`torch.nn.Module`): - The model to dispatch. - device_map (`Dict[str, Union[str, int, torch.device]]`): - A dictionary mapping module names in the models `state_dict` to the device they should go to. Note that - `"disk"` is accepted even if it's not a proper value for `torch.device`. - main_device (`str`, `int` or `torch.device`, *optional*): - The main execution device. Will default to the first device in the `device_map` different from `"cpu"` or - `"disk"`. - state_dict (`Dict[str, torch.Tensor]`, *optional*): - The state dict of the part of the model that will be kept on CPU. - offload_dir (`str` or `os.PathLike`): - The folder in which to offload the model weights (or where the model weights are already offloaded). - offload_index (`Dict`, *optional*): - A dictionary from weight name to their information (`dtype`/ `shape` or safetensors filename). Will default - to the index saved in `save_folder`. - offload_buffers (`bool`, *optional*, defaults to `False`): - Whether or not to offload the buffers with the model parameters. - skip_keys (`str` or `List[str]`, *optional*): - A list of keys to ignore when moving inputs or outputs between devices. - preload_module_classes (`List[str]`, *optional*): - A list of classes whose instances should load all their weights (even in the submodules) at the beginning - of the forward. This should only be used for classes that have submodules which are registered but not - called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, - `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. - force_hooks (`bool`, *optional*, defaults to `False`): - Whether or not to force device hooks to be attached to the model even if all layers are dispatched to a - single device. - """ - # Error early if the device map is incomplete. - check_device_map(model, device_map) - - # We need to force hook for quantized model that can't be moved with to() - if getattr(model, "quantization_method", "bitsandbytes") == "bitsandbytes": - # since bnb 0.43.2, we can move 4-bit model - if getattr(model, "is_loaded_in_8bit", False) or ( - getattr(model, "is_loaded_in_4bit", False) and not is_bnb_available(min_version="0.43.2") - ): - force_hooks = True - - # We attach hooks if the device_map has at least 2 different devices or if - # force_hooks is set to `True`. Otherwise, the model in already loaded - # in the unique device and the user can decide where to dispatch the model. - # If the model is quantized, we always force-dispatch the model - if (len(set(device_map.values())) > 1) or force_hooks: - if main_device is None: - if set(device_map.values()) == {"cpu"} or set(device_map.values()) == {"cpu", "disk"}: - main_device = "cpu" - else: - main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0] - - if main_device != "cpu": - cpu_modules = [name for name, device in device_map.items() if device == "cpu"] - if state_dict is None and len(cpu_modules) > 0: - state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules) - - disk_modules = [name for name, device in device_map.items() if device == "disk"] - if offload_dir is None and offload_index is None and len(disk_modules) > 0: - raise ValueError( - "We need an `offload_dir` to dispatch this model according to this `device_map`, the following submodules " - f"need to be offloaded: {', '.join(disk_modules)}." - ) - if ( - len(disk_modules) > 0 - and offload_index is None - and (not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json"))) - ): - disk_state_dict = extract_submodules_state_dict(model.state_dict(), disk_modules) - offload_state_dict(offload_dir, disk_state_dict) - - execution_device = { - name: main_device if device in ["cpu", "disk"] else device for name, device in device_map.items() - } - execution_device[""] = main_device - offloaded_devices = ["disk"] if main_device == "cpu" or main_device == "mps" else ["cpu", "disk"] - offload = {name: device in offloaded_devices for name, device in device_map.items()} - save_folder = offload_dir if len(disk_modules) > 0 else None - if state_dict is not None or save_folder is not None or offload_index is not None: - device = main_device if offload_index is not None else None - weights_map = OffloadedWeightsLoader( - state_dict=state_dict, save_folder=save_folder, index=offload_index, device=device - ) - else: - weights_map = None - - # When dispatching the model's parameters to the devices specified in device_map, we want to avoid allocating memory several times for the - # tied parameters. The dictionary tied_params_map keeps track of the already allocated data for a given tied parameter (represented by its - # original pointer) on each devices. - tied_params = find_tied_parameters(model) - - tied_params_map = {} - for group in tied_params: - for param_name in group: - # data_ptr() is enough here, as `find_tied_parameters` finds tied params simply by comparing `param1 is param2`, so we don't need - # to care about views of tensors through storage_offset. - data_ptr = recursive_getattr(model, param_name).data_ptr() - tied_params_map[data_ptr] = {} - - # Note: To handle the disk offloading case, we can not simply use weights_map[param_name].data_ptr() as the reference pointer, - # as we have no guarantee that safetensors' `file.get_tensor()` will always give the same pointer. - - attach_align_device_hook_on_blocks( - model, - execution_device=execution_device, - offload=offload, - offload_buffers=offload_buffers, - weights_map=weights_map, - skip_keys=skip_keys, - preload_module_classes=preload_module_classes, - tied_params_map=tied_params_map, - ) - - # warn if there is any params on the meta device - offloaded_devices_str = " and ".join( - [device for device in set(device_map.values()) if device in ("cpu", "disk")] - ) - if len(offloaded_devices_str) > 0: - logger.warning( - f"Some parameters are on the meta device because they were offloaded to the {offloaded_devices_str}." - ) - - # Attaching the hook may break tied weights, so we retie them - retie_parameters(model, tied_params) - - # add warning to cuda and to method - def add_warning(fn, model): - @wraps(fn) - def wrapper(*args, **kwargs): - warning_msg = "You shouldn't move a model that is dispatched using accelerate hooks." - if str(fn.__name__) == "to": - to_device = torch._C._nn._parse_to(*args, **kwargs)[0] - if to_device is not None: - logger.warning(warning_msg) - else: - logger.warning(warning_msg) - for param in model.parameters(): - if param.device == torch.device("meta"): - raise RuntimeError("You can't move a model that has some modules offloaded to cpu or disk.") - return fn(*args, **kwargs) - - return wrapper - - # Make sure to update _accelerate_added_attributes in hooks.py if you add any hook - model.to = add_warning(model.to, model) - if is_npu_available(): - model.npu = add_warning(model.npu, model) - elif is_mlu_available(): - model.mlu = add_warning(model.mlu, model) - elif is_sdaa_available(): - model.sdaa = add_warning(model.sdaa, model) - elif is_musa_available(): - model.musa = add_warning(model.musa, model) - elif is_xpu_available(): - model.xpu = add_warning(model.xpu, model) - else: - model.cuda = add_warning(model.cuda, model) - - # Check if we are using multi-gpus with RTX 4000 series - use_multi_gpu = len([device for device in set(device_map.values()) if device not in ("cpu", "disk")]) > 1 - if use_multi_gpu and not check_cuda_p2p_ib_support(): - logger.warning( - "We've detected an older driver with an RTX 4000 series GPU. These drivers have issues with P2P. " - "This can affect the multi-gpu inference when using accelerate device_map." - "Please make sure to update your driver to the latest version which resolves this." - ) - else: - device = list(device_map.values())[0] - # `torch.Tensor.to()` is not supported by `torch_npu` (see this [issue](https://github.com/Ascend/pytorch/issues/16)). - if is_npu_available() and isinstance(device, int): - device = f"npu:{device}" - elif is_mlu_available() and isinstance(device, int): - device = f"mlu:{device}" - elif is_sdaa_available() and isinstance(device, int): - device = f"sdaa:{device}" - elif is_musa_available() and isinstance(device, int): - device = f"musa:{device}" - if device != "disk": - model.to(device) - else: - raise ValueError( - "You are trying to offload the whole model to the disk. Please use the `disk_offload` function instead." - ) - # Convert OrderedDict back to dict for easier usage - model.hf_device_map = dict(device_map) - return model - - -def load_checkpoint_and_dispatch( - model: nn.Module, - checkpoint: Union[str, os.PathLike], - device_map: Optional[Union[str, dict[str, Union[int, str, torch.device]]]] = None, - max_memory: Optional[dict[Union[int, str], Union[int, str]]] = None, - no_split_module_classes: Optional[list[str]] = None, - offload_folder: Optional[Union[str, os.PathLike]] = None, - offload_buffers: bool = False, - dtype: Optional[Union[str, torch.dtype]] = None, - offload_state_dict: Optional[bool] = None, - skip_keys: Optional[Union[str, list[str]]] = None, - preload_module_classes: Optional[list[str]] = None, - force_hooks: bool = False, - strict: bool = False, - full_state_dict: bool = True, - broadcast_from_rank0: bool = False, -): - """ - Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are - loaded and adds the various hooks that will make this model run properly (even if split across devices). - - Args: - model (`torch.nn.Module`): The model in which we want to load a checkpoint. - checkpoint (`str` or `os.PathLike`): - The folder checkpoint to load. It can be: - - a path to a file containing a whole model state dict - - a path to a `.json` file containing the index to a sharded checkpoint - - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint. - device_map (`Dict[str, Union[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. - - To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For more - information about each option see [here](../concept_guides/big_model_inference#designing-a-device-map). - Defaults to None, which means [`dispatch_model`] will not be called. - max_memory (`Dict`, *optional*): - A dictionary device identifier to maximum memory. Will default to the maximum memory available for each GPU - and the available CPU RAM if unset. - no_split_module_classes (`List[str]`, *optional*): - A list of layer class names that should never be split across device (for instance any layer that has a - residual connection). - offload_folder (`str` or `os.PathLike`, *optional*): - If the `device_map` contains any value `"disk"`, the folder where we will offload weights. - offload_buffers (`bool`, *optional*, defaults to `False`): - In the layers that are offloaded on the CPU or the hard drive, whether or not to offload the buffers as - well as the parameters. - dtype (`str` or `torch.dtype`, *optional*): - If provided, the weights will be converted to that type when loaded. - offload_state_dict (`bool`, *optional*): - If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if - the weight of the CPU state dict + the biggest shard does not fit. Will default to `True` if the device map - picked contains `"disk"` values. - skip_keys (`str` or `List[str]`, *optional*): - A list of keys to ignore when moving inputs or outputs between devices. - preload_module_classes (`List[str]`, *optional*): - A list of classes whose instances should load all their weights (even in the submodules) at the beginning - of the forward. This should only be used for classes that have submodules which are registered but not - called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, - `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. - force_hooks (`bool`, *optional*, defaults to `False`): - Whether or not to force device hooks to be attached to the model even if all layers are dispatched to a - single device. - strict (`bool`, *optional*, defaults to `False`): - Whether to strictly enforce that the keys in the checkpoint state_dict match the keys of the model's - state_dict. - full_state_dict (`bool`, *optional*, defaults to `True`): if this is set to `True`, all the tensors in the - loaded state_dict will be gathered. No ShardedTensor and DTensor will be in the loaded state_dict. - broadcast_from_rank0 (`False`, *optional*, defaults to `False`): when the option is `True`, a distributed - `ProcessGroup` must be initialized. rank0 should receive a full state_dict and will broadcast the tensors - in the state_dict one by one to other ranks. Other ranks will receive the tensors and shard (if applicable) - according to the local shards in the model. - - Example: - - ```python - >>> from accelerate import init_empty_weights, load_checkpoint_and_dispatch - >>> from huggingface_hub import hf_hub_download - >>> from transformers import AutoConfig, AutoModelForCausalLM - - >>> # Download the Weights - >>> checkpoint = "EleutherAI/gpt-j-6B" - >>> weights_location = hf_hub_download(checkpoint, "pytorch_model.bin") - - >>> # Create a model and initialize it with empty weights - >>> config = AutoConfig.from_pretrained(checkpoint) - >>> with init_empty_weights(): - ... model = AutoModelForCausalLM.from_config(config) - - >>> # Load the checkpoint and dispatch it to the right devices - >>> model = load_checkpoint_and_dispatch( - ... model, weights_location, device_map="auto", no_split_module_classes=["GPTJBlock"] - ... ) - ``` - """ - if isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]: - raise ValueError( - "If passing a string for `device_map`, please choose 'auto', 'balanced', 'balanced_low_0' or 'sequential'." - ) - if isinstance(device_map, str): - if device_map != "sequential": - max_memory = get_balanced_memory( - model, - max_memory=max_memory, - no_split_module_classes=no_split_module_classes, - dtype=dtype, - low_zero=(device_map == "balanced_low_0"), - ) - device_map = infer_auto_device_map( - model, - max_memory=max_memory, - no_split_module_classes=no_split_module_classes, - dtype=dtype, - offload_buffers=offload_buffers, - ) - if offload_state_dict is None and device_map is not None and "disk" in device_map.values(): - offload_state_dict = True - load_checkpoint_in_model( - model, - checkpoint, - device_map=device_map, - offload_folder=offload_folder, - dtype=dtype, - offload_state_dict=offload_state_dict, - offload_buffers=offload_buffers, - strict=strict, - full_state_dict=full_state_dict, - broadcast_from_rank0=broadcast_from_rank0, - ) - if device_map is None: - return model - return dispatch_model( - model, - device_map=device_map, - offload_dir=offload_folder, - offload_buffers=offload_buffers, - skip_keys=skip_keys, - preload_module_classes=preload_module_classes, - force_hooks=force_hooks, - ) - - -def attach_layerwise_casting_hooks( - module: torch.nn.Module, - storage_dtype: torch.dtype, - compute_dtype: torch.dtype, - skip_modules_pattern: Union[str, tuple[str, ...]] = None, - skip_modules_classes: Optional[tuple[type[torch.nn.Module], ...]] = None, - non_blocking: bool = False, -) -> None: - r""" - Applies layerwise casting to a given module. The module expected here is a PyTorch `nn.Module`. This is helpful for - reducing memory requirements when one doesn't want to fully quantize a model. Model params can be kept in say, - `torch.float8_e4m3fn` and upcasted to a higher precision like `torch.bfloat16` during forward pass and downcasted - back to `torch.float8_e4m3fn` to realize memory savings. - - Args: - module (`torch.nn.Module`): - The module whose leaf modules will be cast to a high precision dtype for computation, and to a low - precision dtype for storage. - storage_dtype (`torch.dtype`): - 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 `None`): - A list of patterns to match the names of the modules to skip during the layerwise casting process. 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`): - 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. - - Example: - - ```python - >>> from accelerate.hooks import attach_layerwise_casting_hooks - >>> from transformers import AutoModelForCausalLM - >>> import torch - - >>> # Model - >>> checkpoint = "EleutherAI/gpt-j-6B" - >>> model = AutoModelForCausalLM.from_pretrained(checkpoint) - - >>> # Attach hooks and perform inference - >>> attach_layerwise_casting_hooks(model, storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16) - >>> with torch.no_grad(): - ... model(...) - ``` - - Users can also pass modules they want to avoid from getting downcasted. - - ```py - >>> attach_layerwise_casting_hooks( - ... model, storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16, skip_modules_pattern=["norm"] - ... ) - ``` - """ - _attach_layerwise_casting_hooks( - module, storage_dtype, compute_dtype, skip_modules_pattern, skip_modules_classes, non_blocking - ) - - -def _attach_layerwise_casting_hooks( - module: torch.nn.Module, - storage_dtype: torch.dtype, - compute_dtype: torch.dtype, - skip_modules_pattern: Union[str, tuple[str, ...]] = None, - skip_modules_classes: Optional[tuple[type[torch.nn.Module], ...]] = None, - non_blocking: bool = False, - _prefix: str = "", -): - should_skip = (skip_modules_classes is not None and isinstance(module, skip_modules_classes)) or ( - skip_modules_pattern is not None and any(re.search(pattern, _prefix) for pattern in skip_modules_pattern) - ) - if should_skip: - logger.debug(f'Skipping layerwise casting for layer "{_prefix}"') - return - - if isinstance(module, SUPPORTED_PYTORCH_LAYERS_FOR_UPCASTING): - logger.debug(f'Applying layerwise casting to layer "{_prefix}"') - add_hook_to_module( - module, - LayerwiseCastingHook(storage_dtype=storage_dtype, compute_dtype=compute_dtype, non_blocking=non_blocking), - append=True, - ) - return - - for name, submodule in module.named_children(): - layer_name = f"{_prefix}.{name}" if _prefix else name - _attach_layerwise_casting_hooks( - submodule, - storage_dtype, - compute_dtype, - skip_modules_pattern, - skip_modules_classes, - non_blocking, - _prefix=layer_name, - ) diff --git a/venv/lib/python3.10/site-packages/accelerate/checkpointing.py b/venv/lib/python3.10/site-packages/accelerate/checkpointing.py deleted file mode 100644 index 425dd3fecd323dc2cf161ff5b427662b82ee823b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/checkpointing.py +++ /dev/null @@ -1,330 +0,0 @@ -# Copyright 2022 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 random -from pathlib import Path - -import numpy as np -import torch -from safetensors.torch import load_model - -from .utils import ( - MODEL_NAME, - OPTIMIZER_NAME, - RNG_STATE_NAME, - SAFE_MODEL_NAME, - SAFE_WEIGHTS_NAME, - SAMPLER_NAME, - SCALER_NAME, - SCHEDULER_NAME, - WEIGHTS_NAME, - get_pretty_name, - is_cuda_available, - is_hpu_available, - is_mlu_available, - is_musa_available, - is_sdaa_available, - is_torch_version, - is_torch_xla_available, - is_xpu_available, - load, - save, -) - - -if is_torch_version(">=", "2.4.0"): - from torch.amp import GradScaler -else: - from torch.cuda.amp import GradScaler - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - -from .logging import get_logger -from .state import PartialState - - -logger = get_logger(__name__) - - -def save_accelerator_state( - output_dir: str, - model_states: list[dict], - optimizers: list, - schedulers: list, - dataloaders: list, - process_index: int, - step: int, - scaler: GradScaler = None, - save_on_each_node: bool = False, - safe_serialization: bool = True, -): - """ - Saves the current states of the models, optimizers, scaler, and RNG generators to a given directory. - - - - If `safe_serialization` is `True`, models will be saved with `safetensors` while the rest are saved using native - `pickle`. - - - - Args: - output_dir (`str` or `os.PathLike`): - The name of the folder to save all relevant weights and states. - model_states (`List[torch.nn.Module]`): - A list of model states - optimizers (`List[torch.optim.Optimizer]`): - A list of optimizer instances - schedulers (`List[torch.optim.lr_scheduler._LRScheduler]`): - A list of learning rate schedulers - dataloaders (`List[torch.utils.data.DataLoader]`): - A list of dataloader instances to save their sampler states - process_index (`int`): - The current process index in the Accelerator state - step (`int`): - The current step in the internal step tracker - scaler (`torch.amp.GradScaler`, *optional*): - An optional gradient scaler instance to save; - save_on_each_node (`bool`, *optional*): - Whether to save on every node, or only the main node. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). - """ - output_dir = Path(output_dir) - # Model states - for i, state in enumerate(model_states): - weights_name = WEIGHTS_NAME if not safe_serialization else SAFE_WEIGHTS_NAME - if i > 0: - weights_name = weights_name.replace(".", f"_{i}.") - output_model_file = output_dir.joinpath(weights_name) - save(state, output_model_file, save_on_each_node=save_on_each_node, safe_serialization=safe_serialization) - logger.info(f"Model weights saved in {output_model_file}") - # Optimizer states - for i, opt in enumerate(optimizers): - state = opt.state_dict() - optimizer_name = f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin" - output_optimizer_file = output_dir.joinpath(optimizer_name) - save(state, output_optimizer_file, save_on_each_node=save_on_each_node, safe_serialization=False) - logger.info(f"Optimizer state saved in {output_optimizer_file}") - # Scheduler states - for i, scheduler in enumerate(schedulers): - state = scheduler.state_dict() - scheduler_name = f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin" - output_scheduler_file = output_dir.joinpath(scheduler_name) - save(state, output_scheduler_file, save_on_each_node=save_on_each_node, safe_serialization=False) - logger.info(f"Scheduler state saved in {output_scheduler_file}") - # DataLoader states - for i, dataloader in enumerate(dataloaders): - sampler_name = f"{SAMPLER_NAME}.bin" if i == 0 else f"{SAMPLER_NAME}_{i}.bin" - output_sampler_file = output_dir.joinpath(sampler_name) - # Only save if we have our custom sampler - from .data_loader import IterableDatasetShard, SeedableRandomSampler - - if isinstance(dataloader.dataset, IterableDatasetShard): - sampler = dataloader.get_sampler() - if isinstance(sampler, SeedableRandomSampler): - save(sampler, output_sampler_file, save_on_each_node=save_on_each_node, safe_serialization=False) - if getattr(dataloader, "use_stateful_dataloader", False): - dataloader_state_dict_name = "dl_state_dict.bin" if i == 0 else f"dl_state_dict_{i}.bin" - output_dataloader_state_dict_file = output_dir.joinpath(dataloader_state_dict_name) - state_dict = dataloader.state_dict() - torch.save(state_dict, output_dataloader_state_dict_file) - logger.info(f"Sampler state for dataloader {i} saved in {output_sampler_file}") - - # GradScaler state - if scaler is not None: - state = scaler.state_dict() - output_scaler_file = output_dir.joinpath(SCALER_NAME) - torch.save(state, output_scaler_file) - logger.info(f"Gradient scaler state saved in {output_scaler_file}") - # Random number generator states - states = {} - states_name = f"{RNG_STATE_NAME}_{process_index}.pkl" - states["step"] = step - states["random_state"] = random.getstate() - states["numpy_random_seed"] = np.random.get_state() - states["torch_manual_seed"] = torch.get_rng_state() - if is_xpu_available(): - states["torch_xpu_manual_seed"] = torch.xpu.get_rng_state_all() - if is_mlu_available(): - states["torch_mlu_manual_seed"] = torch.mlu.get_rng_state_all() - elif is_sdaa_available(): - states["torch_sdaa_manual_seed"] = torch.sdaa.get_rng_state_all() - elif is_musa_available(): - states["torch_musa_manual_seed"] = torch.musa.get_rng_state_all() - if is_hpu_available(): - states["torch_hpu_manual_seed"] = torch.hpu.get_rng_state_all() - if is_cuda_available(): - states["torch_cuda_manual_seed"] = torch.cuda.get_rng_state_all() - if is_torch_xla_available(): - states["xm_seed"] = xm.get_rng_state() - output_states_file = output_dir.joinpath(states_name) - torch.save(states, output_states_file) - logger.info(f"Random states saved in {output_states_file}") - return output_dir - - -def load_accelerator_state( - input_dir, - models, - optimizers, - schedulers, - dataloaders, - process_index, - scaler=None, - map_location=None, - load_kwargs=None, - **load_model_func_kwargs, -): - """ - Loads states of the models, optimizers, scaler, and RNG generators from a given directory. - - Args: - input_dir (`str` or `os.PathLike`): - The name of the folder to load all relevant weights and states. - models (`List[torch.nn.Module]`): - A list of model instances - optimizers (`List[torch.optim.Optimizer]`): - A list of optimizer instances - schedulers (`List[torch.optim.lr_scheduler._LRScheduler]`): - A list of learning rate schedulers - process_index (`int`): - The current process index in the Accelerator state - scaler (`torch.amp.GradScaler`, *optional*): - An optional *GradScaler* instance to load - map_location (`str`, *optional*): - What device to load the optimizer state onto. Should be one of either "cpu" or "on_device". - load_kwargs (`dict`, *optional*): - Additional arguments that can be passed to the `load` function. - load_model_func_kwargs (`dict`, *optional*): - Additional arguments that can be passed to the model's `load_state_dict` method. - - Returns: - `dict`: Contains the `Accelerator` attributes to override while loading the state. - """ - # stores the `Accelerator` attributes to override - override_attributes = dict() - if map_location not in [None, "cpu", "on_device"]: - raise TypeError( - "Unsupported optimizer map location passed, please choose one of `None`, `'cpu'`, or `'on_device'`" - ) - if map_location is None: - map_location = "cpu" - elif map_location == "on_device": - map_location = PartialState().device - - if load_kwargs is None: - load_kwargs = {} - - input_dir = Path(input_dir) - # Model states - for i, model in enumerate(models): - ending = f"_{i}" if i > 0 else "" - input_model_file = input_dir.joinpath(f"{SAFE_MODEL_NAME}{ending}.safetensors") - if input_model_file.exists(): - load_model(model, input_model_file, device=str(map_location), **load_model_func_kwargs) - else: - # Load with torch - input_model_file = input_dir.joinpath(f"{MODEL_NAME}{ending}.bin") - state_dict = load(input_model_file, map_location=map_location) - model.load_state_dict(state_dict, **load_model_func_kwargs) - logger.info("All model weights loaded successfully") - - # Optimizer states - for i, opt in enumerate(optimizers): - optimizer_name = f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin" - input_optimizer_file = input_dir.joinpath(optimizer_name) - optimizer_state = load(input_optimizer_file, map_location=map_location, **load_kwargs) - optimizers[i].load_state_dict(optimizer_state) - logger.info("All optimizer states loaded successfully") - - # Scheduler states - for i, scheduler in enumerate(schedulers): - scheduler_name = f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin" - input_scheduler_file = input_dir.joinpath(scheduler_name) - scheduler_state = load(input_scheduler_file, **load_kwargs) - scheduler.load_state_dict(scheduler_state) - logger.info("All scheduler states loaded successfully") - - for i, dataloader in enumerate(dataloaders): - sampler_name = f"{SAMPLER_NAME}.bin" if i == 0 else f"{SAMPLER_NAME}_{i}.bin" - input_sampler_file = input_dir.joinpath(sampler_name) - # Only load if we have our custom sampler - from .data_loader import IterableDatasetShard, SeedableRandomSampler - - if isinstance(dataloader.dataset, IterableDatasetShard): - sampler = dataloader.get_sampler() - if isinstance(sampler, SeedableRandomSampler): - sampler = dataloader.set_sampler(load(input_sampler_file)) - if getattr(dataloader, "use_stateful_dataloader", False): - dataloader_state_dict_name = "dl_state_dict.bin" if i == 0 else f"dl_state_dict_{i}.bin" - input_dataloader_state_dict_file = input_dir.joinpath(dataloader_state_dict_name) - if input_dataloader_state_dict_file.exists(): - state_dict = load(input_dataloader_state_dict_file, **load_kwargs) - dataloader.load_state_dict(state_dict) - logger.info("All dataloader sampler states loaded successfully") - - # GradScaler state - if scaler is not None: - input_scaler_file = input_dir.joinpath(SCALER_NAME) - scaler_state = load(input_scaler_file) - scaler.load_state_dict(scaler_state) - logger.info("GradScaler state loaded successfully") - - # Random states - try: - states = load(input_dir.joinpath(f"{RNG_STATE_NAME}_{process_index}.pkl")) - if "step" in states: - override_attributes["step"] = states["step"] - random.setstate(states["random_state"]) - np.random.set_state(states["numpy_random_seed"]) - torch.set_rng_state(states["torch_manual_seed"]) - if is_xpu_available(): - torch.xpu.set_rng_state_all(states["torch_xpu_manual_seed"]) - if is_mlu_available(): - torch.mlu.set_rng_state_all(states["torch_mlu_manual_seed"]) - elif is_sdaa_available(): - torch.sdaa.set_rng_state_all(states["torch_sdaa_manual_seed"]) - elif is_musa_available(): - torch.musa.set_rng_state_all(states["torch_musa_manual_seed"]) - else: - torch.cuda.set_rng_state_all(states["torch_cuda_manual_seed"]) - if is_torch_xla_available(): - xm.set_rng_state(states["xm_seed"]) - logger.info("All random states loaded successfully") - except Exception: - logger.info("Could not load random states") - - return override_attributes - - -def save_custom_state(obj, path, index: int = 0, save_on_each_node: bool = False): - """ - Saves the state of `obj` to `{path}/custom_checkpoint_{index}.pkl` - """ - # Should this be the right way to get a qual_name type value from `obj`? - save_location = Path(path) / f"custom_checkpoint_{index}.pkl" - logger.info(f"Saving the state of {get_pretty_name(obj)} to {save_location}") - save(obj.state_dict(), save_location, save_on_each_node=save_on_each_node) - - -def load_custom_state(obj, path, index: int = 0): - """ - Loads the state of `obj` at `{path}/custom_checkpoint_{index}.pkl`. Will always set `weights_only=False` when - loading the state. - """ - load_location = f"{path}/custom_checkpoint_{index}.pkl" - logger.info(f"Loading the state of {get_pretty_name(obj)} from {load_location}") - obj.load_state_dict(load(load_location, map_location="cpu", weights_only=False)) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__init__.py b/venv/lib/python3.10/site-packages/accelerate/commands/__init__.py deleted file mode 100644 index c9cbe26c257b515f657c05e1996d517e69613972..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright 2020 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. diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 45fa86fa8994780078de93b850a8dfee0613eff1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/accelerate_cli.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/accelerate_cli.cpython-310.pyc deleted file mode 100644 index 37abae224e0272336f45b3083b58fa739deb8cfa..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/accelerate_cli.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/env.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/env.cpython-310.pyc deleted file mode 100644 index e3eb339d899d1abf541b7500c6eb52e1ad798db6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/env.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/estimate.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/estimate.cpython-310.pyc deleted file mode 100644 index ebfc97b2269681b3c6c494111bcbfac28aa705b6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/estimate.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/launch.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/launch.cpython-310.pyc deleted file mode 100644 index bb44c46afb04c5dc089f7faecb39498958f80576..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/launch.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/merge.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/merge.cpython-310.pyc deleted file mode 100644 index 87b67b3194bd1cf4242e9b9e1210a2308201ade5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/merge.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/test.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/test.cpython-310.pyc deleted file mode 100644 index b70afeb31a99e87a6440d04e0da8d1b290c9df37..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/test.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/to_fsdp2.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/to_fsdp2.cpython-310.pyc deleted file mode 100644 index 22f32ac4eb5b3f5950f147aa698b5daff37735e4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/to_fsdp2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/tpu.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/tpu.cpython-310.pyc deleted file mode 100644 index 0817c015174cea05b383cb659e5ccf88c54c7915..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/tpu.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index 781957c07e932b17206d33e799902bf35ef6f7a9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/accelerate_cli.py b/venv/lib/python3.10/site-packages/accelerate/commands/accelerate_cli.py deleted file mode 100644 index b878c8debd874e1418b946775b11568c7487ad72..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/accelerate_cli.py +++ /dev/null @@ -1,54 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2021 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 accelerate.commands.config import get_config_parser -from accelerate.commands.env import env_command_parser -from accelerate.commands.estimate import estimate_command_parser -from accelerate.commands.launch import launch_command_parser -from accelerate.commands.merge import merge_command_parser -from accelerate.commands.test import test_command_parser -from accelerate.commands.to_fsdp2 import to_fsdp2_command_parser -from accelerate.commands.tpu import tpu_command_parser -from accelerate.commands.utils import CustomArgumentParser - - -def main(): - parser = CustomArgumentParser("Accelerate CLI tool", usage="accelerate []", allow_abbrev=False) - subparsers = parser.add_subparsers(help="accelerate command helpers") - - # Register commands - get_config_parser(subparsers=subparsers) - estimate_command_parser(subparsers=subparsers) - env_command_parser(subparsers=subparsers) - launch_command_parser(subparsers=subparsers) - merge_command_parser(subparsers=subparsers) - tpu_command_parser(subparsers=subparsers) - test_command_parser(subparsers=subparsers) - to_fsdp2_command_parser(subparsers=subparsers) - - # Let's go - args = parser.parse_args() - - if not hasattr(args, "func"): - parser.print_help() - exit(1) - - # Run - args.func(args) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__init__.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/__init__.py deleted file mode 100644 index 649a15888cccd070b3d4ca9a600457c6ad59d4d3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/config/__init__.py +++ /dev/null @@ -1,52 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2021 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 argparse - -from .config import config_command_parser -from .config_args import default_config_file, load_config_from_file # noqa: F401 -from .default import default_command_parser -from .update import update_command_parser - - -def get_config_parser(subparsers=None): - parent_parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False) - # The main config parser - config_parser = config_command_parser(subparsers) - # The subparser to add commands to - subcommands = config_parser.add_subparsers(title="subcommands", dest="subcommand") - - # Then add other parsers with the parent parser - default_command_parser(subcommands, parents=[parent_parser]) - update_command_parser(subcommands, parents=[parent_parser]) - - return config_parser - - -def main(): - config_parser = get_config_parser() - args = config_parser.parse_args() - - if not hasattr(args, "func"): - config_parser.print_help() - exit(1) - - # Run - args.func(args) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 397787cf94c392eaa76d515f66acffda3ff0c573..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/cluster.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/cluster.cpython-310.pyc deleted file mode 100644 index ef96576540ccc60ad331d36ff449e0d42b3a72d6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/cluster.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config.cpython-310.pyc deleted file mode 100644 index bbc85f930efff617226439f8bbbedb316b9d81ca..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config_args.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config_args.cpython-310.pyc deleted file mode 100644 index 7c812228dc8ea3f8906ed38aeefb38bf87631617..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config_args.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config_utils.cpython-310.pyc deleted file mode 100644 index 9e63e09bf1cb824846eaef542816047db9f3c6c7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/config_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/default.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/default.cpython-310.pyc deleted file mode 100644 index f24657f24b3b20d6219e720b856163017b00fa4e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/default.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/sagemaker.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/sagemaker.cpython-310.pyc deleted file mode 100644 index 9b77e6c9816c00bd642ab092430188f370610cce..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/sagemaker.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/update.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/update.cpython-310.pyc deleted file mode 100644 index 32abad6203f10d9ca7356c26cab7cdb8fa84aa53..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/config/__pycache__/update.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/cluster.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/cluster.py deleted file mode 100644 index 420fa5495ba7d3df064b7d97973db4f181d044ea..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/config/cluster.py +++ /dev/null @@ -1,869 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2021 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 os - -from ...utils import ( - ComputeEnvironment, - DistributedType, - is_deepspeed_available, - is_fp8_available, - is_hpu_available, - is_mlu_available, - is_mps_available, - is_msamp_available, - is_musa_available, - is_npu_available, - is_sdaa_available, - is_transformer_engine_available, - is_transformers_available, - is_xpu_available, -) -from ...utils.constants import ( - DEEPSPEED_MULTINODE_LAUNCHERS, - FSDP2_STATE_DICT_TYPE, - FSDP_AUTO_WRAP_POLICY, - FSDP_BACKWARD_PREFETCH, - FSDP_SHARDING_STRATEGY, - FSDP_STATE_DICT_TYPE, - TORCH_DYNAMO_MODES, -) -from .config_args import ClusterConfig -from .config_utils import ( - DYNAMO_BACKENDS, - _ask_field, - _ask_options, - _convert_distributed_mode, - _convert_dynamo_backend, - _convert_fp8_backend, - _convert_mixed_precision, - _convert_yes_no_to_bool, -) - - -def get_cluster_input(): - distributed_type = _ask_options( - "Which type of machine are you using?", - [ - "No distributed training", - "multi-CPU", - "multi-XPU", - "multi-HPU", - "multi-GPU", - "multi-NPU", - "multi-MLU", - "multi-SDAA", - "multi-MUSA", - "TPU", - ], - _convert_distributed_mode, - ) - - machine_rank = 0 - num_machines = 1 - num_processes = 1 - gpu_ids = None - main_process_ip = None - main_process_port = None - rdzv_backend = "static" - same_network = True - debug = False - - if distributed_type in [ - DistributedType.MULTI_GPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_NPU, - DistributedType.MULTI_XPU, - DistributedType.MULTI_CPU, - DistributedType.MULTI_HPU, - ]: - num_machines = _ask_field( - "How many different machines will you use (use more than 1 for multi-node training)? [1]: ", - int, - default=1, - ) - if num_machines > 1: - machine_rank = _ask_options( - "What is the rank of this machine?", - list(range(num_machines)), - int, - ) - main_process_ip = _ask_field( - "What is the IP address of the machine that will host the main process? ", - ) - main_process_port = _ask_field( - "What is the port you will use to communicate with the main process? ", - int, - ) - same_network = _ask_field( - "Are all the machines on the same local network? Answer `no` if nodes are on the cloud and/or on different network hosts [YES/no]: ", - _convert_yes_no_to_bool, - default=True, - error_message="Please enter yes or no.", - ) - if not same_network: - rdzv_backend = _ask_field( - "What rendezvous backend will you use? ('static', 'c10d', ...): ", default="static" - ) - debug = _ask_field( - "Should distributed operations be checked while running for errors? This can avoid timeout issues but will be slower. [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - - if distributed_type == DistributedType.NO: - use_cpu = _ask_field( - "Do you want to run your training on CPU only (even if a GPU / Apple Silicon / Ascend NPU device is available)? [yes/NO]:", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - elif distributed_type == DistributedType.MULTI_CPU: - use_cpu = True - else: - use_cpu = False - - ipex_config = {} - mpirun_config = {} - if use_cpu or is_xpu_available(): - ipex_config["ipex"] = _ask_field( - "Do you want to use Intel PyTorch Extension (IPEX) to speed up training on CPU/XPU? [yes/NO]:", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - - if use_cpu: - if distributed_type == DistributedType.MULTI_CPU: - use_mpirun = _ask_field( - "Do you want accelerate to launch mpirun? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if use_mpirun: - mpirun_hostfile = _ask_field( - "Please enter the path to the hostfile to use with mpirun [~/hostfile]: ", - str, - default="~/hostfile", - ) - mpirun_config["mpirun_hostfile"] = os.path.expanduser(mpirun_hostfile.strip()) - mpirun_config["mpirun_ccl"] = _ask_field("Enter the number of oneCCL worker threads [1]: ", default=1) - - dynamo_config = {} - use_dynamo = _ask_field( - "Do you wish to optimize your script with torch dynamo?[yes/NO]:", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if use_dynamo: - prefix = "dynamo_" - dynamo_config[prefix + "backend"] = _ask_options( - "Which dynamo backend would you like to use?", - [x.lower() for x in DYNAMO_BACKENDS], - _convert_dynamo_backend, - default=2, - ) - use_custom_options = _ask_field( - "Do you want to customize the defaults sent to torch.compile? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - - if use_custom_options: - dynamo_config[prefix + "mode"] = _ask_options( - "Which mode do you want to use?", - TORCH_DYNAMO_MODES, - lambda x: TORCH_DYNAMO_MODES[int(x)], - default=0, - ) - dynamo_config[prefix + "use_fullgraph"] = _ask_field( - "Do you want the fullgraph mode or it is ok to break model into several subgraphs? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - dynamo_config[prefix + "use_dynamic"] = _ask_field( - "Do you want to enable dynamic shape tracing? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - dynamo_config[prefix + "use_regional_compilation"] = _ask_field( - "Do you want to enable regional compilation? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - - use_mps = not use_cpu and is_mps_available() - deepspeed_config = {} - if ( - distributed_type - in [ - DistributedType.MULTI_GPU, - DistributedType.MULTI_XPU, - DistributedType.MULTI_HPU, - DistributedType.MULTI_NPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.NO, - ] - and not use_mps - ): - use_deepspeed = _ask_field( - "Do you want to use DeepSpeed? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if use_deepspeed: - distributed_type = DistributedType.DEEPSPEED - assert is_deepspeed_available(), ( - "DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source" - ) - - if distributed_type == DistributedType.DEEPSPEED: - use_deepspeed_config = _ask_field( - "Do you want to specify a json file to a DeepSpeed config? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if use_deepspeed_config: - deepspeed_config["deepspeed_config_file"] = _ask_field( - "Please enter the path to the json DeepSpeed config file: ", - str, - default="none", - ) - else: - deepspeed_config["zero_stage"] = _ask_options( - "What should be your DeepSpeed's ZeRO optimization stage?", - [0, 1, 2, 3], - int, - default=2, - ) - - deepspeed_devices = ["none", "cpu", "nvme"] - if deepspeed_config["zero_stage"] >= 2: - deepspeed_config["offload_optimizer_device"] = _ask_options( - "Where to offload optimizer states?", deepspeed_devices, lambda x: deepspeed_devices[int(x)] - ) - deepspeed_config["offload_param_device"] = _ask_options( - "Where to offload parameters?", deepspeed_devices, lambda x: deepspeed_devices[int(x)] - ) - if deepspeed_config["offload_param_device"] == "nvme": - deepspeed_config["offload_param_nvme_path"] = _ask_field( - "Nvme Path to offload parameters?", - str, - default="/nvme", - ) - if deepspeed_config["offload_optimizer_device"] == "nvme": - deepspeed_config["offload_optimizer_nvme_path"] = _ask_field( - "Nvme Path to offload optimizer states?", - str, - default="/nvme", - ) - deepspeed_config["gradient_accumulation_steps"] = _ask_field( - "How many gradient accumulation steps you're passing in your script? [1]: ", - int, - default=1, - ) - use_gradient_clipping = _ask_field( - "Do you want to use gradient clipping? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if use_gradient_clipping: - deepspeed_config["gradient_clipping"] = _ask_field( - "What is the gradient clipping value? [1.0]: ", - float, - default=1.0, - ) - if deepspeed_config["zero_stage"] == 3: - deepspeed_config["zero3_save_16bit_model"] = _ask_field( - "Do you want to save 16-bit model weights when using ZeRO Stage-3? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - deepspeed_config["zero3_init_flag"] = _ask_field( - "Do you want to enable `deepspeed.zero.Init` when using ZeRO Stage-3 for constructing massive models? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if deepspeed_config["zero3_init_flag"]: - if not is_transformers_available(): - raise Exception( - "When `zero3_init_flag` is set, it requires Transformers to be installed. " - "Please run `pip3 install transformers`." - ) - use_moe = _ask_field( - "Do you want to enable Mixture-of-Experts training (MoE)? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if use_moe: - deepspeed_config["deepspeed_moe_layer_cls_names"] = _ask_field( - "Specify the comma-separated list of transformers MoE layer class names (case-sensitive), e.g : " - " `MixtralSparseMoeBlock`, `Qwen2MoeSparseMoeBlock`, `JetMoEAttention,JetMoEBlock` ... : ", - str, - ) - - if num_machines > 1: - launcher_query = "Which Type of launcher do you want to use?" - deepspeed_config["deepspeed_multinode_launcher"] = _ask_options( - launcher_query, - DEEPSPEED_MULTINODE_LAUNCHERS, - lambda x: DEEPSPEED_MULTINODE_LAUNCHERS[int(x)], - ) - - if deepspeed_config["deepspeed_multinode_launcher"] != DEEPSPEED_MULTINODE_LAUNCHERS[1]: - deepspeed_config["deepspeed_hostfile"] = _ask_field( - "DeepSpeed configures multi-node compute resources with hostfile. " - "Each row is of the format `hostname slots=[num_gpus]`, e.g., `localhost slots=2`; " - "for more information please refer official [documentation]" - "(https://www.deepspeed.ai/getting-started/#resource-configuration-multi-node). " - "Please specify the location of hostfile: ", - str, - ) - - is_exclusion_filter = _ask_field( - "Do you want to specify exclusion filter string? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if is_exclusion_filter: - deepspeed_config["deepspeed_exclusion_filter"] = _ask_field( - "DeepSpeed exclusion filter string: ", - str, - ) - - is_inclusion_filter = _ask_field( - "Do you want to specify inclusion filter string? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if is_inclusion_filter: - deepspeed_config["deepspeed_inclusion_filter"] = _ask_field( - "DeepSpeed inclusion filter string: ", - str, - ) - - fsdp_config = {} - - if distributed_type in [ - DistributedType.MULTI_GPU, - DistributedType.MULTI_NPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_XPU, - DistributedType.MULTI_HPU, - ]: - use_fsdp = _ask_field( - "Do you want to use FullyShardedDataParallel? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if use_fsdp: - distributed_type = DistributedType.FSDP - if distributed_type == DistributedType.FSDP: - fsdp_config["fsdp_version"] = _ask_options( - "What should be your FSDP version? [2]: ", - [1, 2], - lambda x: int(x) + 1, - default=1, - ) - fsdp_version = fsdp_config["fsdp_version"] # extract to a variable to simplify usage later - - if fsdp_version == 1: - sharding_strategy_query = "What should be your sharding strategy?" - fsdp_config["fsdp_reshard_after_forward"] = _ask_options( - sharding_strategy_query, - FSDP_SHARDING_STRATEGY, - lambda x: FSDP_SHARDING_STRATEGY[int(x)], - ) - else: - fsdp_config["fsdp_reshard_after_forward"] = _ask_field( - "Do you want to enable resharding after forward? [YES/no]: ", - _convert_yes_no_to_bool, - default=True, - error_message="Please enter yes or no.", - ) - - fsdp_config["fsdp_offload_params"] = _ask_field( - "Do you want to offload parameters and gradients to CPU? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - - fsdp_wrap_query = "What should be your auto wrap policy?" - fsdp_config["fsdp_auto_wrap_policy"] = _ask_options( - fsdp_wrap_query, - FSDP_AUTO_WRAP_POLICY, - lambda x: FSDP_AUTO_WRAP_POLICY[int(x)], - ) - if fsdp_config["fsdp_auto_wrap_policy"] == FSDP_AUTO_WRAP_POLICY[0]: - use_no_split_modules = _ask_field( - "Do you want to use the model's `_no_split_modules` to wrap. Only applicable for 🤗 Transformers [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if not use_no_split_modules: - fsdp_config["fsdp_transformer_layer_cls_to_wrap"] = _ask_field( - "Specify the comma-separated list of transformer layer class names (case-sensitive) to wrap ,e.g, :" - "`BertLayer`, `GPTJBlock`, `T5Block`, `BertLayer,BertEmbeddings,BertSelfOutput` ...? : ", - str, - ) - elif fsdp_config["fsdp_auto_wrap_policy"] == FSDP_AUTO_WRAP_POLICY[1]: - fsdp_config["fsdp_min_num_params"] = _ask_field( - "What should be your FSDP's minimum number of parameters for Default Auto Wrapping Policy? [1e8]: ", - int, - default=100000000, - ) - # Removed in FSDP2, ask for user input for FSDP1 - if fsdp_version == 1: - fsdp_backward_prefetch_query = "What should be your FSDP's backward prefetch policy?" - fsdp_config["fsdp_backward_prefetch"] = _ask_options( - fsdp_backward_prefetch_query, - FSDP_BACKWARD_PREFETCH, - lambda x: FSDP_BACKWARD_PREFETCH[int(x)], - ) - - fsdp_state_dict_type_query = "What should be your FSDP's state dict type?" - fsdp_config["fsdp_state_dict_type"] = _ask_options( - fsdp_state_dict_type_query, - FSDP_STATE_DICT_TYPE if fsdp_version == 1 else FSDP2_STATE_DICT_TYPE, - lambda x: FSDP_STATE_DICT_TYPE[int(x)] if fsdp_version == 1 else FSDP2_STATE_DICT_TYPE[int(x)], - default=0, - ) - # Not implemented in FSDP2, ask for user input for FSDP1 - if fsdp_version == 1: - fsdp_config["fsdp_forward_prefetch"] = _ask_field( - "Do you want to enable FSDP's forward prefetch policy? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - # Obsolete in FSDP2, ask for user input for FSDP1 - if fsdp_version == 1: - fsdp_config["fsdp_use_orig_params"] = _ask_field( - "Do you want to enable FSDP's `use_orig_params` feature? [YES/no]: ", - _convert_yes_no_to_bool, - default=True, - error_message="Please enter yes or no.", - ) - fsdp_config["fsdp_cpu_ram_efficient_loading"] = _ask_field( - "Do you want to enable CPU RAM efficient model loading? Only applicable for 🤗 Transformers models. [YES/no]: ", - _convert_yes_no_to_bool, - default=True, - error_message="Please enter yes or no.", - ) - # Obsolete in FSDP2, ask for user input for FSDP1 - if fsdp_version == 1: - if fsdp_config["fsdp_cpu_ram_efficient_loading"]: - fsdp_config["fsdp_sync_module_states"] = True - else: - fsdp_config["fsdp_sync_module_states"] = _ask_field( - "Do you want each individually wrapped FSDP unit to broadcast module parameters from rank 0 at the start? [YES/no]: ", - _convert_yes_no_to_bool, - default=True, - error_message="Please enter yes or no.", - ) - fsdp_config["fsdp_activation_checkpointing"] = _ask_field( - "Do you want to enable FSDP activation checkpointing? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - - megatron_lm_config = {} - if distributed_type in [DistributedType.MULTI_GPU]: - use_megatron_lm = _ask_field( - "Do you want to use Megatron-LM ? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if use_megatron_lm: - distributed_type = DistributedType.MEGATRON_LM - if distributed_type == DistributedType.MEGATRON_LM: - prefix = "megatron_lm_" - megatron_lm_config[prefix + "tp_degree"] = _ask_field( - "What is the Tensor Parallelism degree/size? [1]:", - int, - default=1, - error_message="Please enter an integer.", - ) - if megatron_lm_config[prefix + "tp_degree"] > 1: - megatron_lm_config[prefix + "sequence_parallelism"] = _ask_field( - "Do you want to enable Sequence Parallelism? [YES/no]: ", - _convert_yes_no_to_bool, - default=True, - error_message="Please enter yes or no.", - ) - - megatron_lm_config[prefix + "pp_degree"] = _ask_field( - "What is the Pipeline Parallelism degree/size? [1]:", - int, - default=1, - error_message="Please enter an integer.", - ) - if megatron_lm_config[prefix + "pp_degree"] > 1: - megatron_lm_config[prefix + "num_micro_batches"] = _ask_field( - "What is the number of micro-batches? [1]:", - int, - default=1, - error_message="Please enter an integer.", - ) - - megatron_lm_config[prefix + "recompute_activations"] = _ask_field( - "Do you want to enable selective activation recomputation? [YES/no]: ", - _convert_yes_no_to_bool, - default=True, - error_message="Please enter yes or no.", - ) - - megatron_lm_config[prefix + "use_distributed_optimizer"] = _ask_field( - "Do you want to use distributed optimizer " - "which shards optimizer state and gradients across data parallel ranks? [YES/no]: ", - _convert_yes_no_to_bool, - default=True, - error_message="Please enter yes or no.", - ) - - megatron_lm_config[prefix + "gradient_clipping"] = _ask_field( - "What is the gradient clipping value based on global L2 Norm (0 to disable)? [1.0]: ", - float, - default=1.0, - ) - # TPU specific defaults - tpu_commands = None - tpu_command_file = None - tpu_downcast_bf16 = "no" - tpu_env = [] - tpu_name = None - tpu_vm = None - tpu_zone = None - tpu_use_sudo = False - tpu_use_cluster = False - - if distributed_type in [ - DistributedType.MULTI_CPU, - DistributedType.MULTI_XPU, - DistributedType.MULTI_HPU, - DistributedType.MULTI_GPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_NPU, - DistributedType.XLA, - ]: - machine_type = str(distributed_type).split(".")[1].replace("MULTI_", "") - if machine_type == "TPU": - machine_type += " cores" - elif machine_type == "CPU": - machine_type = "processes" - else: - machine_type += "(s)" - num_processes = _ask_field( - f"How many {machine_type} should be used for distributed training? [1]:", - int, - default=1, - error_message="Please enter an integer.", - ) - elif distributed_type in [DistributedType.FSDP, DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM]: - num_processes = _ask_field( - "How many GPU(s) should be used for distributed training? [1]:", - int, - default=1, - error_message="Please enter an integer.", - ) - else: - num_processes = 1 - - if (distributed_type == DistributedType.MULTI_GPU) and (num_machines == 1) and (num_processes == 1): - raise ValueError( - f"Specified distributed type {distributed_type} but only using 1 GPU on a single machine. Please select `No distributed training` for the type of machine you are using." - ) - - if ( - distributed_type - in [ - DistributedType.MULTI_GPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_NPU, - DistributedType.MULTI_XPU, - DistributedType.MULTI_HPU, - DistributedType.NO, - ] - and not use_cpu - and not use_mps - ): - if is_npu_available(): - machine_type = "NPU(s)" - elif is_mlu_available(): - machine_type = "MLU(s)" - elif is_sdaa_available(): - machine_type = "SDAA(s)" - elif is_musa_available(): - machine_type = "MUSA(s)" - elif is_xpu_available(): - machine_type = "XPU(s)" - elif is_hpu_available(): - machine_type = "HPU(s)" - else: - machine_type = "GPU(s)" - gpu_ids = _ask_field( - f"What {machine_type} (by id) should be used for training on this machine as a comma-separated list? [all]:", - default="all", - ) - - # CPU affinity is only supported on NVIDIA hardware for now - enable_cpu_affinity = False - if distributed_type in (DistributedType.NO, DistributedType.MULTI_GPU) and not use_cpu and not use_mps: - enable_cpu_affinity = _ask_field( - "Would you like to enable numa efficiency? (Currently only supported on NVIDIA hardware). [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - - fp8_config = None - if distributed_type == DistributedType.XLA: - mixed_precision = "no" - main_training_function = _ask_field( - "What is the name of the function in your script that should be launched in all parallel scripts? [main]: ", - default="main", - ) - tpu_use_cluster = _ask_field( - "Are you using a TPU cluster? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if tpu_use_cluster: - tpu_name = _ask_field( - "What is the name of your TPU cluster? ", - default=None, - error_message="Please enter the name of your TPU cluster.", - ) - tpu_zone = _ask_field( - "What is the zone of your TPU cluster? ", - default=None, - error_message="Please enter the zone of your TPU cluster.", - ) - tpu_use_sudo = _ask_field( - "To run a python script in a TPU pod, should `sudo` be used? [yes/NO]: ", - default=False, - error_message="Please enter yes or no.", - ) - run_commands = _ask_field( - "Do you have code you wish to run on startup in each pod? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if run_commands: - use_command_file = _ask_field( - "Is this code located in a bash script? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if use_command_file: - tpu_command_file = _ask_field( - "What is the path to your bash script? ", - default=None, - error_message="Please enter the path to your bash script.", - ) - tpu_command_file = os.path.abspath(tpu_command_file) - else: - print("Please enter each command separately you wish to run on startup in each pod.") - tpu_commands = [] - another_command = True - while another_command: - tpu_commands.append( - _ask_field( - "Please enter a single command to be ran ", - default=None, - error_message="Please enter the commands you wish to run on startup in each pod as a single string.", - ) - ) - another_command = _ask_field( - "Do you wish to add another command? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - tpu_vm = _ask_field( - "If not using an instance group, what are the names of the Compute VM instances to be used, separated by a comma: ", - default="", - ).split(",") - tpu_env = _ask_field( - "What environment variables do you wish to set in each pod, separated by a comma: ", - default="", - ).split(",") - - else: - main_training_function = "main" - if distributed_type == DistributedType.DEEPSPEED and use_deepspeed_config: - mixed_precision = None - else: - mixed_precision = _ask_options( - "Do you wish to use mixed precision?", - ["no", "fp16", "bf16", "fp8"], - _convert_mixed_precision, - ) - if mixed_precision == "fp8": - if not is_fp8_available(): - raise ValueError("FP8 (either Transformer Engine or MSAMP) is not installed on this machine.") - fp8_config = {} - fp8_config["backend"] = _ask_options( - "Which FP8 backend do you want to use?", - ["te", "msamp"], - _convert_fp8_backend, - ) - if fp8_config["backend"] == "TE": - if not is_transformer_engine_available(): - raise ValueError("TransformersEngine was selected, but it is not installed on this machine.") - fp8_config["use_autocast_during_eval"] = _ask_field( - "Do you want to use FP8 autocast during eval mode? Generally better metrics are found when this is disabled [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - ) - fp8_config["margin"] = _ask_field( - "What margin should be used for gradient scaling? [0]: ", - int, - default=0, - ) - fp8_config["interval"] = _ask_field( - "What interval should be used for for how often the scaling factor is recomputed? [1]: ", - int, - default=1, - ) - fp8_config["fp8_format"] = _ask_options( - "Which weight format should be used?", - ["HYBRID", "E4M3", "E5M2"], - lambda i: ["HYBRID", "E4M3", "E5M2"][i], - default=0, - ) - fp8_config["amax_history_length"] = _ask_field( - "What length of history should be used for the amax scaling factor computation? [1024]: ", - int, - default=1024, - ) - fp8_config["amax_compute_algorithm"] = _ask_options( - "Which algorithm should be used for the amax scaling factor computation?", - ["max", "most_recent"], - lambda x: "max" if x == 0 else "most_recent", - default=0, - ) - fp8_config["override_linear_precision"] = _ask_field( - "Do you want to to execute `fprop`, `dgrad`, and `wgrad` GEMMS in higher precision? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - ) - if fp8_config["override_linear_precision"]: - fprop = _ask_field( - "Should `fprop` be executed in higher precision? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - ) - dgrad = _ask_field( - "Should `dgrad` be executed in higher precision? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - ) - wgrad = _ask_field( - "Should `wgrad` be executed in higher precision? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - ) - fp8_config["override_linear_precision"] = (fprop, dgrad, wgrad) - else: - fp8_config["override_linear_precision"] = (False, False, False) - - elif fp8_config["backend"] == "MSAMP": - if not is_msamp_available(): - raise ValueError("MSAMP was selected, but it is not installed on this machine.") - fp8_config["optimization_level"] = _ask_options( - "Which optimization level should be used?", - ["O1", "O2"], - lambda x: "O1" if x == 0 else "O2", - default=1, - ) - - if use_dynamo and mixed_precision == "no" and not use_cpu: - print( - "Torch dynamo used without mixed precision requires TF32 to be efficient. Accelerate will enable it by default when launching your scripts." - ) - - if distributed_type == DistributedType.XLA and mixed_precision == "bf16": - tpu_downcast_bf16 = _ask_field( - "Should `torch.float` be cast as `bfloat16` and `torch.double` remain `float32` on TPUs?", default="no" - ) - - return ClusterConfig( - compute_environment=ComputeEnvironment.LOCAL_MACHINE, - distributed_type=distributed_type, - num_processes=num_processes, - gpu_ids=gpu_ids, - mixed_precision=mixed_precision, - downcast_bf16=tpu_downcast_bf16, - machine_rank=machine_rank, - num_machines=num_machines, - main_process_ip=main_process_ip, - main_process_port=main_process_port, - main_training_function=main_training_function, - fp8_config=fp8_config, - deepspeed_config=deepspeed_config, - fsdp_config=fsdp_config, - megatron_lm_config=megatron_lm_config, - ipex_config=ipex_config, - mpirun_config=mpirun_config, - use_cpu=use_cpu, - rdzv_backend=rdzv_backend, - same_network=same_network, - commands=tpu_commands, - command_file=tpu_command_file, - tpu_env=tpu_env, - tpu_name=tpu_name, - tpu_vm=tpu_vm, - tpu_zone=tpu_zone, - tpu_use_sudo=tpu_use_sudo, - tpu_use_cluster=tpu_use_cluster, - dynamo_config=dynamo_config, - debug=debug, - enable_cpu_affinity=enable_cpu_affinity, - ) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/config.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/config.py deleted file mode 100644 index 72414f2abe62d76bd5133f4b0ed99bf34133f6f6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/config/config.py +++ /dev/null @@ -1,89 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2021 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 argparse -import os - -from accelerate.utils import ComputeEnvironment - -from .cluster import get_cluster_input -from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file # noqa: F401 -from .config_utils import _ask_field, _ask_options, _convert_compute_environment # noqa: F401 -from .sagemaker import get_sagemaker_input - - -description = "Launches a series of prompts to create and save a `default_config.yaml` configuration file for your training system. Should always be ran first on your machine" - - -def get_user_input(): - compute_environment = _ask_options( - "In which compute environment are you running?", - ["This machine", "AWS (Amazon SageMaker)"], - _convert_compute_environment, - ) - if compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER: - config = get_sagemaker_input() - else: - config = get_cluster_input() - return config - - -def config_command_parser(subparsers=None): - if subparsers is not None: - parser = subparsers.add_parser("config", description=description) - else: - parser = argparse.ArgumentParser("Accelerate config command", description=description) - - parser.add_argument( - "--config_file", - default=None, - help=( - "The path to use to store the config file. Will default to a file named default_config.yaml in the cache " - "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have " - "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed " - "with 'huggingface'." - ), - ) - - if subparsers is not None: - parser.set_defaults(func=config_command) - return parser - - -def config_command(args): - config = get_user_input() - if args.config_file is not None: - config_file = args.config_file - else: - if not os.path.isdir(cache_dir): - os.makedirs(cache_dir) - config_file = default_yaml_config_file - - if config_file.endswith(".json"): - config.to_json_file(config_file) - else: - config.to_yaml_file(config_file) - print(f"accelerate configuration saved at {config_file}") - - -def main(): - parser = config_command_parser() - args = parser.parse_args() - config_command(args) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/config_args.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/config_args.py deleted file mode 100644 index 3907329287b2f3ac033f4f74f84bc877b91af4f9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/config/config_args.py +++ /dev/null @@ -1,252 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2021 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 -from dataclasses import dataclass -from enum import Enum -from typing import Optional, Union - -import yaml - -from ...utils import ComputeEnvironment, DistributedType, SageMakerDistributedType -from ...utils.constants import SAGEMAKER_PYTHON_VERSION, SAGEMAKER_PYTORCH_VERSION, SAGEMAKER_TRANSFORMERS_VERSION - - -hf_cache_home = os.path.expanduser( - os.environ.get("HF_HOME", os.path.join(os.environ.get("XDG_CACHE_HOME", "~/.cache"), "huggingface")) -) -cache_dir = os.path.join(hf_cache_home, "accelerate") -default_json_config_file = os.path.join(cache_dir, "default_config.yaml") -default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml") - -# For backward compatibility: the default config is the json one if it's the only existing file. -if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file): - default_config_file = default_yaml_config_file -else: - default_config_file = default_json_config_file - - -def load_config_from_file(config_file): - if config_file is not None: - if not os.path.isfile(config_file): - raise FileNotFoundError( - f"The passed configuration file `{config_file}` does not exist. " - "Please pass an existing file to `accelerate launch`, or use the default one " - "created through `accelerate config` and run `accelerate launch` " - "without the `--config_file` argument." - ) - else: - config_file = default_config_file - with open(config_file, encoding="utf-8") as f: - if config_file.endswith(".json"): - if ( - json.load(f).get("compute_environment", ComputeEnvironment.LOCAL_MACHINE) - == ComputeEnvironment.LOCAL_MACHINE - ): - config_class = ClusterConfig - else: - config_class = SageMakerConfig - return config_class.from_json_file(json_file=config_file) - else: - if ( - yaml.safe_load(f).get("compute_environment", ComputeEnvironment.LOCAL_MACHINE) - == ComputeEnvironment.LOCAL_MACHINE - ): - config_class = ClusterConfig - else: - config_class = SageMakerConfig - return config_class.from_yaml_file(yaml_file=config_file) - - -@dataclass -class BaseConfig: - compute_environment: ComputeEnvironment - distributed_type: Union[DistributedType, SageMakerDistributedType] - mixed_precision: str - use_cpu: bool - debug: bool - - def to_dict(self): - result = self.__dict__ - # For serialization, it's best to convert Enums to strings (or their underlying value type). - - def _convert_enums(value): - if isinstance(value, Enum): - return value.value - if isinstance(value, dict): - if not bool(value): - return None - for key1, value1 in value.items(): - value[key1] = _convert_enums(value1) - return value - - for key, value in result.items(): - result[key] = _convert_enums(value) - result = {k: v for k, v in result.items() if v is not None} - return result - - @staticmethod - def process_config(config_dict): - """ - Processes `config_dict` and sets default values for any missing keys - """ - if "compute_environment" not in config_dict: - config_dict["compute_environment"] = ComputeEnvironment.LOCAL_MACHINE - if "distributed_type" not in config_dict: - raise ValueError("A `distributed_type` must be specified in the config file.") - if "num_processes" not in config_dict and config_dict["distributed_type"] == DistributedType.NO: - config_dict["num_processes"] = 1 - if "mixed_precision" not in config_dict: - config_dict["mixed_precision"] = "fp16" if ("fp16" in config_dict and config_dict["fp16"]) else None - if "fp16" in config_dict: # Convert the config to the new format. - del config_dict["fp16"] - if "dynamo_backend" in config_dict: # Convert the config to the new format. - dynamo_backend = config_dict.pop("dynamo_backend") - config_dict["dynamo_config"] = {} if dynamo_backend == "NO" else {"dynamo_backend": dynamo_backend} - if "use_cpu" not in config_dict: - config_dict["use_cpu"] = False - if "debug" not in config_dict: - config_dict["debug"] = False - if "enable_cpu_affinity" not in config_dict: - config_dict["enable_cpu_affinity"] = False - return config_dict - - @classmethod - def from_json_file(cls, json_file=None): - json_file = default_json_config_file if json_file is None else json_file - with open(json_file, encoding="utf-8") as f: - config_dict = json.load(f) - config_dict = cls.process_config(config_dict) - extra_keys = sorted(set(config_dict.keys()) - set(cls.__dataclass_fields__.keys())) - if len(extra_keys) > 0: - raise ValueError( - f"The config file at {json_file} had unknown keys ({extra_keys}), please try upgrading your `accelerate`" - " version or fix (and potentially remove) these keys from your config file." - ) - - return cls(**config_dict) - - def to_json_file(self, json_file): - with open(json_file, "w", encoding="utf-8") as f: - content = json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n" - f.write(content) - - @classmethod - def from_yaml_file(cls, yaml_file=None): - yaml_file = default_yaml_config_file if yaml_file is None else yaml_file - with open(yaml_file, encoding="utf-8") as f: - config_dict = yaml.safe_load(f) - config_dict = cls.process_config(config_dict) - extra_keys = sorted(set(config_dict.keys()) - set(cls.__dataclass_fields__.keys())) - if len(extra_keys) > 0: - raise ValueError( - f"The config file at {yaml_file} had unknown keys ({extra_keys}), please try upgrading your `accelerate`" - " version or fix (and potentially remove) these keys from your config file." - ) - return cls(**config_dict) - - def to_yaml_file(self, yaml_file): - with open(yaml_file, "w", encoding="utf-8") as f: - yaml.safe_dump(self.to_dict(), f) - - def __post_init__(self): - if isinstance(self.compute_environment, str): - self.compute_environment = ComputeEnvironment(self.compute_environment) - if isinstance(self.distributed_type, str): - if self.compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER: - self.distributed_type = SageMakerDistributedType(self.distributed_type) - else: - self.distributed_type = DistributedType(self.distributed_type) - if getattr(self, "dynamo_config", None) is None: - self.dynamo_config = {} - - -@dataclass -class ClusterConfig(BaseConfig): - num_processes: int = -1 # For instance if we use SLURM and the user manually passes it in - machine_rank: int = 0 - num_machines: int = 1 - gpu_ids: Optional[str] = None - main_process_ip: Optional[str] = None - main_process_port: Optional[int] = None - rdzv_backend: Optional[str] = "static" - same_network: Optional[bool] = False - main_training_function: str = "main" - enable_cpu_affinity: bool = False - - # args for FP8 training - fp8_config: dict = None - # args for deepspeed_plugin - deepspeed_config: dict = None - # args for fsdp - fsdp_config: dict = None - # args for megatron_lm - megatron_lm_config: dict = None - # args for ipex - ipex_config: dict = None - # args for mpirun - mpirun_config: dict = None - # args for TPU - downcast_bf16: bool = False - - # args for TPU pods - tpu_name: str = None - tpu_zone: str = None - tpu_use_cluster: bool = False - tpu_use_sudo: bool = False - command_file: str = None - commands: list[str] = None - tpu_vm: list[str] = None - tpu_env: list[str] = None - - # args for dynamo - dynamo_config: dict = None - - def __post_init__(self): - if self.deepspeed_config is None: - self.deepspeed_config = {} - if self.fsdp_config is None: - self.fsdp_config = {} - if self.megatron_lm_config is None: - self.megatron_lm_config = {} - if self.ipex_config is None: - self.ipex_config = {} - if self.mpirun_config is None: - self.mpirun_config = {} - if self.fp8_config is None: - self.fp8_config = {} - return super().__post_init__() - - -@dataclass -class SageMakerConfig(BaseConfig): - ec2_instance_type: str - iam_role_name: str - image_uri: Optional[str] = None - profile: Optional[str] = None - region: str = "us-east-1" - num_machines: int = 1 - gpu_ids: str = "all" - base_job_name: str = f"accelerate-sagemaker-{num_machines}" - pytorch_version: str = SAGEMAKER_PYTORCH_VERSION - transformers_version: str = SAGEMAKER_TRANSFORMERS_VERSION - py_version: str = SAGEMAKER_PYTHON_VERSION - sagemaker_inputs_file: str = None - sagemaker_metrics_file: str = None - additional_args: dict = None - dynamo_config: dict = None - enable_cpu_affinity: bool = False diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/config_utils.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/config_utils.py deleted file mode 100644 index 24ee971383c8fdda1491e2b5795446790755ac70..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/config/config_utils.py +++ /dev/null @@ -1,122 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2021 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 argparse - -from ...utils.dataclasses import ( - ComputeEnvironment, - DistributedType, - DynamoBackend, - FP8BackendType, - PrecisionType, - SageMakerDistributedType, -) -from ..menu import BulletMenu - - -DYNAMO_BACKENDS = [ - "EAGER", - "AOT_EAGER", - "INDUCTOR", - "AOT_TS_NVFUSER", - "NVPRIMS_NVFUSER", - "CUDAGRAPHS", - "OFI", - "FX2TRT", - "ONNXRT", - "TENSORRT", - "AOT_TORCHXLA_TRACE_ONCE", - "TORHCHXLA_TRACE_ONCE", - "IPEX", - "TVM", -] - - -def _ask_field(input_text, convert_value=None, default=None, error_message=None): - ask_again = True - while ask_again: - result = input(input_text) - try: - if default is not None and len(result) == 0: - return default - return convert_value(result) if convert_value is not None else result - except Exception: - if error_message is not None: - print(error_message) - - -def _ask_options(input_text, options=[], convert_value=None, default=0): - menu = BulletMenu(input_text, options) - result = menu.run(default_choice=default) - return convert_value(result) if convert_value is not None else result - - -def _convert_compute_environment(value): - value = int(value) - return ComputeEnvironment(["LOCAL_MACHINE", "AMAZON_SAGEMAKER"][value]) - - -def _convert_distributed_mode(value): - value = int(value) - return DistributedType( - [ - "NO", - "MULTI_CPU", - "MULTI_XPU", - "MULTI_HPU", - "MULTI_GPU", - "MULTI_NPU", - "MULTI_MLU", - "MULTI_SDAA", - "MULTI_MUSA", - "XLA", - ][value] - ) - - -def _convert_dynamo_backend(value): - value = int(value) - return DynamoBackend(DYNAMO_BACKENDS[value]).value - - -def _convert_mixed_precision(value): - value = int(value) - return PrecisionType(["no", "fp16", "bf16", "fp8"][value]) - - -def _convert_sagemaker_distributed_mode(value): - value = int(value) - return SageMakerDistributedType(["NO", "DATA_PARALLEL", "MODEL_PARALLEL"][value]) - - -def _convert_fp8_backend(value): - value = int(value) - return FP8BackendType(["TE", "MSAMP"][value]) - - -def _convert_yes_no_to_bool(value): - return {"yes": True, "no": False}[value.lower()] - - -class SubcommandHelpFormatter(argparse.RawDescriptionHelpFormatter): - """ - A custom formatter that will remove the usage line from the help message for subcommands. - """ - - def _format_usage(self, usage, actions, groups, prefix): - usage = super()._format_usage(usage, actions, groups, prefix) - usage = usage.replace(" [] ", "") - return usage diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/default.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/default.py deleted file mode 100644 index f5f267f8c47cccc80ce9ef53f970c2266472c117..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/config/default.py +++ /dev/null @@ -1,163 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2021 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 pathlib import Path - -import torch - -from ...utils import ( - is_hpu_available, - is_mlu_available, - is_musa_available, - is_npu_available, - is_sdaa_available, - is_xpu_available, -) -from .config_args import ClusterConfig, default_json_config_file -from .config_utils import SubcommandHelpFormatter - - -description = "Create a default config file for Accelerate with only a few flags set." - - -def write_basic_config(mixed_precision="no", save_location: str = default_json_config_file): - """ - Creates and saves a basic cluster config to be used on a local machine with potentially multiple GPUs. Will also - set CPU if it is a CPU-only machine. - - Args: - mixed_precision (`str`, *optional*, defaults to "no"): - Mixed Precision to use. Should be one of "no", "fp16", or "bf16" - save_location (`str`, *optional*, defaults to `default_json_config_file`): - Optional custom save location. Should be passed to `--config_file` when using `accelerate launch`. Default - location is inside the huggingface cache folder (`~/.cache/huggingface`) but can be overridden by setting - the `HF_HOME` environmental variable, followed by `accelerate/default_config.yaml`. - """ - path = Path(save_location) - path.parent.mkdir(parents=True, exist_ok=True) - if path.exists(): - print( - f"Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`." - ) - return False - mixed_precision = mixed_precision.lower() - if mixed_precision not in ["no", "fp16", "bf16", "fp8"]: - raise ValueError( - f"`mixed_precision` should be one of 'no', 'fp16', 'bf16', or 'fp8'. Received {mixed_precision}" - ) - config = { - "compute_environment": "LOCAL_MACHINE", - "mixed_precision": mixed_precision, - } - if is_mlu_available(): - num_mlus = torch.mlu.device_count() - config["num_processes"] = num_mlus - config["use_cpu"] = False - if num_mlus > 1: - config["distributed_type"] = "MULTI_MLU" - else: - config["distributed_type"] = "NO" - if is_sdaa_available(): - num_sdaas = torch.sdaa.device_count() - config["num_processes"] = num_sdaas - config["use_cpu"] = False - if num_sdaas > 1: - config["distributed_type"] = "MULTI_SDAA" - else: - config["distributed_type"] = "NO" - elif is_musa_available(): - num_musas = torch.musa.device_count() - config["num_processes"] = num_musas - config["use_cpu"] = False - if num_musas > 1: - config["distributed_type"] = "MULTI_MUSA" - else: - config["distributed_type"] = "NO" - elif is_hpu_available(): - num_hpus = torch.hpu.device_count() - config["num_processes"] = num_hpus - config["use_cpu"] = False - if num_hpus > 1: - config["distributed_type"] = "MULTI_HPU" - else: - config["distributed_type"] = "NO" - elif torch.cuda.is_available(): - num_gpus = torch.cuda.device_count() - config["num_processes"] = num_gpus - config["use_cpu"] = False - if num_gpus > 1: - config["distributed_type"] = "MULTI_GPU" - else: - config["distributed_type"] = "NO" - elif is_xpu_available(): - num_xpus = torch.xpu.device_count() - config["num_processes"] = num_xpus - config["use_cpu"] = False - if num_xpus > 1: - config["distributed_type"] = "MULTI_XPU" - else: - config["distributed_type"] = "NO" - elif is_npu_available(): - num_npus = torch.npu.device_count() - config["num_processes"] = num_npus - config["use_cpu"] = False - if num_npus > 1: - config["distributed_type"] = "MULTI_NPU" - else: - config["distributed_type"] = "NO" - else: - num_xpus = 0 - config["use_cpu"] = True - config["num_processes"] = 1 - config["distributed_type"] = "NO" - config["debug"] = False - config["enable_cpu_affinity"] = False - config = ClusterConfig(**config) - config.to_json_file(path) - return path - - -def default_command_parser(parser, parents): - parser = parser.add_parser("default", parents=parents, help=description, formatter_class=SubcommandHelpFormatter) - parser.add_argument( - "--config_file", - default=default_json_config_file, - help=( - "The path to use to store the config file. Will default to a file named default_config.yaml in the cache " - "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have " - "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed " - "with 'huggingface'." - ), - dest="save_location", - ) - - parser.add_argument( - "--mixed_precision", - choices=["no", "fp16", "bf16"], - type=str, - help="Whether or not to use mixed precision training. " - "Choose between FP16 and BF16 (bfloat16) training. " - "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.", - default="no", - ) - parser.set_defaults(func=default_config_command) - return parser - - -def default_config_command(args): - config_file = write_basic_config(args.mixed_precision, args.save_location) - if config_file: - print(f"accelerate configuration saved at {config_file}") diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/sagemaker.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/sagemaker.py deleted file mode 100644 index 5092ef31fc4715f901be6c1e7bfe80c0b140d767..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/config/sagemaker.py +++ /dev/null @@ -1,274 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2021 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 - -from ...utils.constants import SAGEMAKER_PARALLEL_EC2_INSTANCES, TORCH_DYNAMO_MODES -from ...utils.dataclasses import ComputeEnvironment, SageMakerDistributedType -from ...utils.imports import is_boto3_available -from .config_args import SageMakerConfig -from .config_utils import ( - DYNAMO_BACKENDS, - _ask_field, - _ask_options, - _convert_dynamo_backend, - _convert_mixed_precision, - _convert_sagemaker_distributed_mode, - _convert_yes_no_to_bool, -) - - -if is_boto3_available(): - import boto3 # noqa: F401 - - -def _create_iam_role_for_sagemaker(role_name): - iam_client = boto3.client("iam") - - sagemaker_trust_policy = { - "Version": "2012-10-17", - "Statement": [ - {"Effect": "Allow", "Principal": {"Service": "sagemaker.amazonaws.com"}, "Action": "sts:AssumeRole"} - ], - } - try: - # create the role, associated with the chosen trust policy - iam_client.create_role( - RoleName=role_name, AssumeRolePolicyDocument=json.dumps(sagemaker_trust_policy, indent=2) - ) - policy_document = { - "Version": "2012-10-17", - "Statement": [ - { - "Effect": "Allow", - "Action": [ - "sagemaker:*", - "ecr:GetDownloadUrlForLayer", - "ecr:BatchGetImage", - "ecr:BatchCheckLayerAvailability", - "ecr:GetAuthorizationToken", - "cloudwatch:PutMetricData", - "cloudwatch:GetMetricData", - "cloudwatch:GetMetricStatistics", - "cloudwatch:ListMetrics", - "logs:CreateLogGroup", - "logs:CreateLogStream", - "logs:DescribeLogStreams", - "logs:PutLogEvents", - "logs:GetLogEvents", - "s3:CreateBucket", - "s3:ListBucket", - "s3:GetBucketLocation", - "s3:GetObject", - "s3:PutObject", - ], - "Resource": "*", - } - ], - } - # attach policy to role - iam_client.put_role_policy( - RoleName=role_name, - PolicyName=f"{role_name}_policy_permission", - PolicyDocument=json.dumps(policy_document, indent=2), - ) - except iam_client.exceptions.EntityAlreadyExistsException: - print(f"role {role_name} already exists. Using existing one") - - -def _get_iam_role_arn(role_name): - iam_client = boto3.client("iam") - return iam_client.get_role(RoleName=role_name)["Role"]["Arn"] - - -def get_sagemaker_input(): - credentials_configuration = _ask_options( - "How do you want to authorize?", - ["AWS Profile", "Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) "], - int, - ) - aws_profile = None - if credentials_configuration == 0: - aws_profile = _ask_field("Enter your AWS Profile name: [default] ", default="default") - os.environ["AWS_PROFILE"] = aws_profile - else: - print( - "Note you will need to provide AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY when you launch you training script with," - "`accelerate launch --aws_access_key_id XXX --aws_secret_access_key YYY`" - ) - aws_access_key_id = _ask_field("AWS Access Key ID: ") - os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id - - aws_secret_access_key = _ask_field("AWS Secret Access Key: ") - os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key - - aws_region = _ask_field("Enter your AWS Region: [us-east-1]", default="us-east-1") - os.environ["AWS_DEFAULT_REGION"] = aws_region - - role_management = _ask_options( - "Do you already have an IAM Role for executing Amazon SageMaker Training Jobs?", - ["Provide IAM Role name", "Create new IAM role using credentials"], - int, - ) - if role_management == 0: - iam_role_name = _ask_field("Enter your IAM role name: ") - else: - iam_role_name = "accelerate_sagemaker_execution_role" - print(f'Accelerate will create an iam role "{iam_role_name}" using the provided credentials') - _create_iam_role_for_sagemaker(iam_role_name) - - is_custom_docker_image = _ask_field( - "Do you want to use custom Docker image? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - docker_image = None - if is_custom_docker_image: - docker_image = _ask_field("Enter your Docker image: ", lambda x: str(x).lower()) - - is_sagemaker_inputs_enabled = _ask_field( - "Do you want to provide SageMaker input channels with data locations? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - sagemaker_inputs_file = None - if is_sagemaker_inputs_enabled: - sagemaker_inputs_file = _ask_field( - "Enter the path to the SageMaker inputs TSV file with columns (channel_name, data_location): ", - lambda x: str(x).lower(), - ) - - is_sagemaker_metrics_enabled = _ask_field( - "Do you want to enable SageMaker metrics? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - sagemaker_metrics_file = None - if is_sagemaker_metrics_enabled: - sagemaker_metrics_file = _ask_field( - "Enter the path to the SageMaker metrics TSV file with columns (metric_name, metric_regex): ", - lambda x: str(x).lower(), - ) - - distributed_type = _ask_options( - "What is the distributed mode?", - ["No distributed training", "Data parallelism"], - _convert_sagemaker_distributed_mode, - ) - dynamo_config = {} - use_dynamo = _ask_field( - "Do you wish to optimize your script with torch dynamo?[yes/NO]:", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - if use_dynamo: - prefix = "dynamo_" - dynamo_config[prefix + "backend"] = _ask_options( - "Which dynamo backend would you like to use?", - [x.lower() for x in DYNAMO_BACKENDS], - _convert_dynamo_backend, - default=2, - ) - use_custom_options = _ask_field( - "Do you want to customize the defaults sent to torch.compile? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - - if use_custom_options: - dynamo_config[prefix + "mode"] = _ask_options( - "Which mode do you want to use?", - TORCH_DYNAMO_MODES, - lambda x: TORCH_DYNAMO_MODES[int(x)], - default="default", - ) - dynamo_config[prefix + "use_fullgraph"] = _ask_field( - "Do you want the fullgraph mode or it is ok to break model into several subgraphs? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - dynamo_config[prefix + "use_dynamic"] = _ask_field( - "Do you want to enable dynamic shape tracing? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - dynamo_config[prefix + "use_regional_compilation"] = _ask_field( - "Do you want to enable regional compilation? [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - - ec2_instance_query = "Which EC2 instance type you want to use for your training?" - if distributed_type != SageMakerDistributedType.NO: - ec2_instance_type = _ask_options( - ec2_instance_query, SAGEMAKER_PARALLEL_EC2_INSTANCES, lambda x: SAGEMAKER_PARALLEL_EC2_INSTANCES[int(x)] - ) - else: - ec2_instance_query += "? [ml.p3.2xlarge]:" - ec2_instance_type = _ask_field(ec2_instance_query, lambda x: str(x).lower(), default="ml.p3.2xlarge") - - debug = False - if distributed_type != SageMakerDistributedType.NO: - debug = _ask_field( - "Should distributed operations be checked while running for errors? This can avoid timeout issues but will be slower. [yes/NO]: ", - _convert_yes_no_to_bool, - default=False, - error_message="Please enter yes or no.", - ) - - num_machines = 1 - if distributed_type in (SageMakerDistributedType.DATA_PARALLEL, SageMakerDistributedType.MODEL_PARALLEL): - num_machines = _ask_field( - "How many machines do you want use? [1]: ", - int, - default=1, - ) - - mixed_precision = _ask_options( - "Do you wish to use FP16 or BF16 (mixed precision)?", - ["no", "fp16", "bf16", "fp8"], - _convert_mixed_precision, - ) - - if use_dynamo and mixed_precision == "no": - print( - "Torch dynamo used without mixed precision requires TF32 to be efficient. Accelerate will enable it by default when launching your scripts." - ) - - return SageMakerConfig( - image_uri=docker_image, - compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER, - distributed_type=distributed_type, - use_cpu=False, - dynamo_config=dynamo_config, - ec2_instance_type=ec2_instance_type, - profile=aws_profile, - region=aws_region, - iam_role_name=iam_role_name, - mixed_precision=mixed_precision, - num_machines=num_machines, - sagemaker_inputs_file=sagemaker_inputs_file, - sagemaker_metrics_file=sagemaker_metrics_file, - debug=debug, - ) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/config/update.py b/venv/lib/python3.10/site-packages/accelerate/commands/config/update.py deleted file mode 100644 index 5f025594b04ada3e3a78687befc5c1bc1d236adf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/config/update.py +++ /dev/null @@ -1,63 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2022 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 pathlib import Path - -from .config_args import default_config_file, load_config_from_file -from .config_utils import SubcommandHelpFormatter - - -description = "Update an existing config file with the latest defaults while maintaining the old configuration." - - -def update_config(args): - """ - Update an existing config file with the latest defaults while maintaining the old configuration. - """ - config_file = args.config_file - if config_file is None and Path(default_config_file).exists(): - config_file = default_config_file - elif not Path(config_file).exists(): - raise ValueError(f"The passed config file located at {config_file} doesn't exist.") - config = load_config_from_file(config_file) - - if config_file.endswith(".json"): - config.to_json_file(config_file) - else: - config.to_yaml_file(config_file) - return config_file - - -def update_command_parser(parser, parents): - parser = parser.add_parser("update", parents=parents, help=description, formatter_class=SubcommandHelpFormatter) - parser.add_argument( - "--config_file", - default=None, - help=( - "The path to the config file to update. Will default to a file named default_config.yaml in the cache " - "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have " - "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed " - "with 'huggingface'." - ), - ) - - parser.set_defaults(func=update_config_command) - return parser - - -def update_config_command(args): - config_file = update_config(args) - print(f"Sucessfully updated the configuration file at {config_file}.") diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/env.py b/venv/lib/python3.10/site-packages/accelerate/commands/env.py deleted file mode 100644 index 3dd2170aea8ba48a08e31c8031c0447825ec4797..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/env.py +++ /dev/null @@ -1,131 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2022 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 argparse -import os -import platform -import subprocess - -import numpy as np -import psutil -import torch - -from accelerate import __version__ as version -from accelerate.commands.config import default_config_file, load_config_from_file - -from ..utils import is_mlu_available, is_musa_available, is_npu_available, is_sdaa_available, is_xpu_available - - -def env_command_parser(subparsers=None): - if subparsers is not None: - parser = subparsers.add_parser("env") - else: - parser = argparse.ArgumentParser("Accelerate env command") - - parser.add_argument( - "--config_file", default=None, help="The config file to use for the default values in the launching script." - ) - - if subparsers is not None: - parser.set_defaults(func=env_command) - return parser - - -def env_command(args): - pt_version = torch.__version__ - pt_cuda_available = torch.cuda.is_available() - pt_xpu_available = is_xpu_available() - pt_mlu_available = is_mlu_available() - pt_sdaa_available = is_sdaa_available() - pt_musa_available = is_musa_available() - pt_npu_available = is_npu_available() - - accelerator = "N/A" - if pt_cuda_available: - accelerator = "CUDA" - elif pt_xpu_available: - accelerator = "XPU" - elif pt_mlu_available: - accelerator = "MLU" - elif pt_sdaa_available: - accelerator = "SDAA" - elif pt_musa_available: - accelerator = "MUSA" - elif pt_npu_available: - accelerator = "NPU" - - accelerate_config = "Not found" - # Get the default from the config file. - if args.config_file is not None or os.path.isfile(default_config_file): - accelerate_config = load_config_from_file(args.config_file).to_dict() - - # if we can run which, get it - command = None - bash_location = "Not found" - if os.name == "nt": - command = ["where", "accelerate"] - elif os.name == "posix": - command = ["which", "accelerate"] - if command is not None: - bash_location = subprocess.check_output(command, text=True, stderr=subprocess.STDOUT).strip() - info = { - "`Accelerate` version": version, - "Platform": platform.platform(), - "`accelerate` bash location": bash_location, - "Python version": platform.python_version(), - "Numpy version": np.__version__, - "PyTorch version": f"{pt_version}", - "PyTorch accelerator": accelerator, - "System RAM": f"{psutil.virtual_memory().total / 1024**3:.2f} GB", - } - if pt_cuda_available: - info["GPU type"] = torch.cuda.get_device_name() - elif pt_xpu_available: - info["XPU type"] = torch.xpu.get_device_name() - elif pt_mlu_available: - info["MLU type"] = torch.mlu.get_device_name() - elif pt_sdaa_available: - info["SDAA type"] = torch.sdaa.get_device_name() - elif pt_musa_available: - info["MUSA type"] = torch.musa.get_device_name() - elif pt_npu_available: - info["CANN version"] = torch.version.cann - - print("\nCopy-and-paste the text below in your GitHub issue\n") - print("\n".join([f"- {prop}: {val}" for prop, val in info.items()])) - - print("- `Accelerate` default config:" if args.config_file is None else "- `Accelerate` config passed:") - accelerate_config_str = ( - "\n".join([f"\t- {prop}: {val}" for prop, val in accelerate_config.items()]) - if isinstance(accelerate_config, dict) - else f"\t{accelerate_config}" - ) - print(accelerate_config_str) - - info["`Accelerate` configs"] = accelerate_config - - return info - - -def main() -> int: - parser = env_command_parser() - args = parser.parse_args() - env_command(args) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/estimate.py b/venv/lib/python3.10/site-packages/accelerate/commands/estimate.py deleted file mode 100644 index 77571777c6c8e78b7729e4b29e1a98cea6e95de7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/estimate.py +++ /dev/null @@ -1,312 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2023 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 huggingface_hub import model_info -from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError - -from accelerate import init_empty_weights -from accelerate.commands.utils import CustomArgumentParser -from accelerate.utils import ( - calculate_maximum_sizes, - convert_bytes, - is_timm_available, - is_transformers_available, -) - - -if is_transformers_available(): - import transformers - from transformers import AutoConfig, AutoModel - -if is_timm_available(): - import timm - - -def verify_on_hub(repo: str, token: str = None): - "Verifies that the model is on the hub and returns the model info." - try: - return model_info(repo, token=token) - except (OSError, GatedRepoError): - return "gated" - except RepositoryNotFoundError: - return "repo" - - -def check_has_model(error): - """ - Checks what library spawned `error` when a model is not found - """ - if is_timm_available() and isinstance(error, RuntimeError) and "Unknown model" in error.args[0]: - return "timm" - elif ( - is_transformers_available() - and isinstance(error, OSError) - and "does not appear to have a file named" in error.args[0] - ): - return "transformers" - else: - return "unknown" - - -def create_empty_model(model_name: str, library_name: str, trust_remote_code: bool = False, access_token: str = None): - """ - Creates an empty model in full precision from its parent library on the `Hub` to calculate the overall memory - consumption. - - Args: - model_name (`str`): - The model name on the Hub - library_name (`str`): - The library the model has an integration with, such as `transformers`. Will be used if `model_name` has no - metadata on the Hub to determine the library. - trust_remote_code (`bool`, `optional`, defaults to `False`): - Whether or not to allow for custom models defined on the Hub in their own modeling files. This option - should only be set to `True` for repositories you trust and in which you have read the code, as it will - execute code present on the Hub on your local machine. - access_token (`str`, `optional`, defaults to `None`): - The access token to use to access private or gated models on the Hub. (for use on the Gradio app) - - Returns: - `torch.nn.Module`: The torch model that has been initialized on the `meta` device. - - """ - model_info = verify_on_hub(model_name, access_token) - # Simplified errors - if model_info == "gated": - raise GatedRepoError( - f"Repo for model `{model_name}` is gated. You must be authenticated to access it. Please run `huggingface-cli login`." - ) - elif model_info == "repo": - raise RepositoryNotFoundError( - f"Repo for model `{model_name}` does not exist on the Hub. If you are trying to access a private repo," - " make sure you are authenticated via `huggingface-cli login` and have access." - ) - if library_name is None: - library_name = getattr(model_info, "library_name", False) - if not library_name: - raise ValueError( - f"Model `{model_name}` does not have any library metadata on the Hub, please manually pass in a `--library_name` to use (such as `transformers`)" - ) - if library_name == "transformers": - if not is_transformers_available(): - raise ImportError( - f"To check `{model_name}`, `transformers` must be installed. Please install it via `pip install transformers`" - ) - print(f"Loading pretrained config for `{model_name}` from `transformers`...") - if model_info.config is None: - raise RuntimeError(f"Tried to load `{model_name}` with `transformers` but it does not have any metadata.") - - auto_map = model_info.config.get("auto_map", False) - config = AutoConfig.from_pretrained(model_name, trust_remote_code=trust_remote_code, token=access_token) - with init_empty_weights(): - # remote code could specify a specific `AutoModel` class in the `auto_map` - constructor = AutoModel - if isinstance(auto_map, dict): - value = None - for key in auto_map.keys(): - if key.startswith("AutoModelFor"): - value = key - break - if value is not None: - constructor = getattr(transformers, value) - # we need to pass the dtype, otherwise it is going to use the torch_dtype that is saved in the config - model = constructor.from_config(config, torch_dtype=torch.float32, trust_remote_code=trust_remote_code) - elif library_name == "timm": - if not is_timm_available(): - raise ImportError( - f"To check `{model_name}`, `timm` must be installed. Please install it via `pip install timm`" - ) - print(f"Loading pretrained config for `{model_name}` from `timm`...") - with init_empty_weights(): - model = timm.create_model(model_name, pretrained=False) - else: - raise ValueError( - f"Library `{library_name}` is not supported yet, please open an issue on GitHub for us to add support." - ) - return model - - -def create_ascii_table(headers: list, rows: list, title: str): - "Creates a pretty table from a list of rows, minimal version of `tabulate`." - sep_char, in_between = "│", "─" - column_widths = [] - for i in range(len(headers)): - column_values = [row[i] for row in rows] + [headers[i]] - max_column_width = max(len(value) for value in column_values) - column_widths.append(max_column_width) - - formats = [f"%{column_widths[i]}s" for i in range(len(rows[0]))] - - pattern = f"{sep_char}{sep_char.join(formats)}{sep_char}" - diff = 0 - - def make_row(left_char, middle_char, right_char): - return f"{left_char}{middle_char.join([in_between * n for n in column_widths])}{in_between * diff}{right_char}" - - separator = make_row("├", "┼", "┤") - if len(title) > sum(column_widths): - diff = abs(len(title) - len(separator)) - column_widths[-1] += diff - - # Update with diff - separator = make_row("├", "┼", "┤") - initial_rows = [ - make_row("┌", in_between, "┐"), - f"{sep_char}{title.center(len(separator) - 2)}{sep_char}", - make_row("├", "┬", "┤"), - ] - table = "\n".join(initial_rows) + "\n" - column_widths[-1] += diff - centered_line = [text.center(column_widths[i]) for i, text in enumerate(headers)] - table += f"{pattern % tuple(centered_line)}\n{separator}\n" - for i, line in enumerate(rows): - centered_line = [t.center(column_widths[i]) for i, t in enumerate(line)] - table += f"{pattern % tuple(centered_line)}\n" - table += f"└{'┴'.join([in_between * n for n in column_widths])}┘" - - return table - - -def estimate_command_parser(subparsers=None): - if subparsers is not None: - parser = subparsers.add_parser("estimate-memory") - else: - parser = CustomArgumentParser(description="Model size estimator for fitting a model onto CUDA memory.") - - parser.add_argument("model_name", type=str, help="The model name on the Hugging Face Hub.") - parser.add_argument( - "--library_name", - type=str, - help="The library the model has an integration with, such as `transformers`, needed only if this information is not stored on the Hub.", - choices=["timm", "transformers"], - ) - parser.add_argument( - "--dtypes", - type=str, - nargs="+", - default=["float32", "float16", "int8", "int4"], - help="The dtypes to use for the model, must be one (or many) of `float32`, `float16`, `int8`, and `int4`", - choices=["float32", "float16", "int8", "int4"], - ) - parser.add_argument( - "--trust_remote_code", - action="store_true", - help="""Whether or not to allow for custom models defined on the Hub in their own modeling files. This flag - should only be used for repositories you trust and in which you have read the code, as it will execute - code present on the Hub on your local machine.""", - default=False, - ) - - if subparsers is not None: - parser.set_defaults(func=estimate_command) - return parser - - -def estimate_training_usage(bytes: int, mixed_precision: str, msamp_config: str = None) -> dict: - """ - Given an amount of `bytes` and `mixed_precision`, calculates how much training memory is needed for a batch size of - 1. - - Args: - bytes (`int`): - The size of the model being trained. - mixed_precision (`str`): - The mixed precision that would be ran. - msamp_config (`str`): - The msamp config to estimate the training memory for if `mixed_precision` is set to `"fp8"`. - """ - memory_sizes = {"model": -1, "optimizer": -1, "gradients": -1, "step": -1} - fp32_size = bytes - fp16_size = bytes // 2 - - if mixed_precision == "float32": - memory_sizes["model"] = fp32_size - memory_sizes["gradients"] = fp32_size - memory_sizes["optimizer"] = fp32_size * 2 - memory_sizes["step"] = fp32_size * 4 - elif mixed_precision in ("float16", "bfloat16") or (mixed_precision == "fp8" and msamp_config is None): - # With native `TransformersEngine`, there is no memory savings with FP8 - # With mixed precision training, the model has weights stored - # in FP16 and FP32 - memory_sizes["model"] = fp32_size - # 1.5 from weight gradient + computation (GEMM) - memory_sizes["gradients"] = fp32_size + fp16_size - # 2x from optimizer states - memory_sizes["optimizer"] = fp32_size * 2 # Optimizer states - memory_sizes["step"] = memory_sizes["optimizer"] - return memory_sizes - - -def gather_data(args): - "Creates an empty model and gathers the data for the sizes" - try: - model = create_empty_model( - args.model_name, library_name=args.library_name, trust_remote_code=args.trust_remote_code - ) - except (RuntimeError, OSError) as e: - library = check_has_model(e) - if library != "unknown": - raise RuntimeError( - f"Tried to load `{args.model_name}` with `{library}` but a possible model to load was not found inside the repo." - ) - raise e - - total_size, largest_layer = calculate_maximum_sizes(model) - - data = [] - - for dtype in args.dtypes: - dtype_total_size = total_size - dtype_largest_layer = largest_layer[0] - dtype_training_size = estimate_training_usage(dtype_total_size, dtype) - if dtype == "float16": - dtype_total_size /= 2 - dtype_largest_layer /= 2 - elif dtype == "int8": - dtype_total_size /= 4 - dtype_largest_layer /= 4 - elif dtype == "int4": - dtype_total_size /= 8 - dtype_largest_layer /= 8 - data.append([dtype, dtype_largest_layer, dtype_total_size, dtype_training_size]) - return data - - -def estimate_command(args): - data = gather_data(args) - for row in data: - for i, item in enumerate(row): - if isinstance(item, (int, float)): - row[i] = convert_bytes(item) - elif isinstance(item, dict): - training_usage = max(item.values()) - row[i] = convert_bytes(training_usage) if training_usage != -1 else "N/A" - - headers = ["dtype", "Largest Layer", "Total Size", "Training using Adam"] - - title = f"Memory Usage for loading `{args.model_name}`" - table = create_ascii_table(headers, data, title) - print(table) - - -def main(): - parser = estimate_command_parser() - args = parser.parse_args() - estimate_command(args) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/launch.py b/venv/lib/python3.10/site-packages/accelerate/commands/launch.py deleted file mode 100644 index 7f4304143a3dd67c45ee9ffb9b3e195f8b874edb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/launch.py +++ /dev/null @@ -1,1209 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2021 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 argparse -import importlib -import logging -import os -import subprocess -import sys -from pathlib import Path - -import psutil -import torch - -from accelerate.commands.config import default_config_file, load_config_from_file -from accelerate.commands.config.config_args import SageMakerConfig -from accelerate.commands.config.config_utils import DYNAMO_BACKENDS -from accelerate.commands.utils import CustomArgumentParser -from accelerate.state import get_int_from_env -from accelerate.utils import ( - ComputeEnvironment, - DistributedType, - PrepareForLaunch, - _filter_args, - check_cuda_p2p_ib_support, - convert_dict_to_env_variables, - is_bf16_available, - is_deepspeed_available, - is_hpu_available, - is_mlu_available, - is_musa_available, - is_npu_available, - is_rich_available, - is_sagemaker_available, - is_sdaa_available, - is_torch_xla_available, - is_xpu_available, - patch_environment, - prepare_deepspeed_cmd_env, - prepare_multi_gpu_env, - prepare_sagemager_args_inputs, - prepare_simple_launcher_cmd_env, - prepare_tpu, - str_to_bool, -) -from accelerate.utils.constants import DEEPSPEED_MULTINODE_LAUNCHERS, TORCH_DYNAMO_MODES - - -if is_rich_available(): - from rich import get_console - from rich.logging import RichHandler - - FORMAT = "%(message)s" - logging.basicConfig(format=FORMAT, datefmt="[%X]", handlers=[RichHandler()]) - - -logger = logging.getLogger(__name__) - - -options_to_group = { - "multi_gpu": "Distributed GPUs", - "tpu": "TPU", - "use_deepspeed": "DeepSpeed Arguments", - "use_fsdp": "FSDP Arguments", - "use_megatron_lm": "Megatron-LM Arguments", - "fp8_backend": "FP8 Arguments", -} - - -def clean_option(option): - "Finds all cases of - after the first two characters and changes them to _" - if "fp8_backend" in option: - option = "--fp8_backend" - if option.startswith("--"): - return option[2:].replace("-", "_") - - -class CustomHelpFormatter(argparse.HelpFormatter): - """ - This is a custom help formatter that will hide all arguments that are not used in the command line when the help is - called. This is useful for the case where the user is using a specific platform and only wants to see the arguments - for that platform. - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.titles = [ - "Hardware Selection Arguments", - "Resource Selection Arguments", - "Training Paradigm Arguments", - "positional arguments", - "optional arguments", - ] - - def add_argument(self, action: argparse.Action): - if "accelerate" in sys.argv[0] and "launch" in sys.argv[1:]: - args = sys.argv[2:] - else: - args = sys.argv[1:] - - if len(args) > 1: - args = list(map(clean_option, args)) - used_platforms = [arg for arg in args if arg in options_to_group.keys()] - used_titles = [options_to_group[o] for o in used_platforms] - if action.container.title not in self.titles + used_titles: - action.help = argparse.SUPPRESS - elif action.container.title == "Hardware Selection Arguments": - if set(action.option_strings).isdisjoint(set(args)): - action.help = argparse.SUPPRESS - else: - action.help = action.help + " (currently selected)" - elif action.container.title == "Training Paradigm Arguments": - if set(action.option_strings).isdisjoint(set(args)): - action.help = argparse.SUPPRESS - else: - action.help = action.help + " (currently selected)" - - action.option_strings = [s for s in action.option_strings if "-" not in s[2:]] - super().add_argument(action) - - def end_section(self): - if len(self._current_section.items) < 2: - self._current_section.items = [] - self._current_section.heading = "" - super().end_section() - - -def launch_command_parser(subparsers=None): - description = "Launch a python script in a distributed scenario. Arguments can be passed in with either hyphens (`--num-processes=2`) or underscores (`--num_processes=2`)" - if subparsers is not None: - parser = subparsers.add_parser( - "launch", description=description, add_help=False, allow_abbrev=False, formatter_class=CustomHelpFormatter - ) - else: - parser = CustomArgumentParser( - "Accelerate launch command", - description=description, - add_help=False, - allow_abbrev=False, - formatter_class=CustomHelpFormatter, - ) - - parser.add_argument("-h", "--help", action="help", help="Show this help message and exit.") - - parser.add_argument( - "--config_file", - default=None, - help="The config file to use for the default values in the launching script.", - ) - parser.add_argument( - "--quiet", - "-q", - action="store_true", - help="Silence subprocess errors from the launch stack trace and only show the relevant tracebacks. (Only applicable to DeepSpeed and single-process configurations)", - ) - # Hardware selection arguments - hardware_args = parser.add_argument_group( - "Hardware Selection Arguments", "Arguments for selecting the hardware to be used." - ) - hardware_args.add_argument( - "--cpu", default=False, action="store_true", help="Whether or not to force the training on the CPU." - ) - hardware_args.add_argument( - "--multi_gpu", - default=False, - action="store_true", - help="Whether or not this should launch a distributed GPU training.", - ) - hardware_args.add_argument( - "--tpu", default=False, action="store_true", help="Whether or not this should launch a TPU training." - ) - hardware_args.add_argument( - "--ipex", - default=False, - action="store_true", - help="Whether or not this should launch a Intel PyTorch Extension (IPEX) training.", - ) - - # Resource selection arguments - resource_args = parser.add_argument_group( - "Resource Selection Arguments", "Arguments for fine-tuning how available hardware should be used." - ) - resource_args.add_argument( - "--mixed_precision", - type=str, - choices=["no", "fp16", "bf16", "fp8"], - help="Whether or not to use mixed precision training. " - "Choose between FP16 and BF16 (bfloat16) training. " - "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.", - ) - resource_args.add_argument( - "--num_processes", type=int, default=None, help="The total number of processes to be launched in parallel." - ) - resource_args.add_argument( - "--num_machines", type=int, default=None, help="The total number of machines used in this training." - ) - resource_args.add_argument( - "--num_cpu_threads_per_process", - type=int, - default=None, - help="The number of CPU threads per process. Can be tuned for optimal performance.", - ) - resource_args.add_argument( - "--enable_cpu_affinity", - default=False, - action="store_true", - help="Whether or not CPU affinity and balancing should be enabled. Currently only supported on NVIDIA hardware.", - ) - # Dynamo arguments - resource_args.add_argument( - "--dynamo_backend", - type=str, - choices=["no"] + [b.lower() for b in DYNAMO_BACKENDS], - help="Choose a backend to optimize your training with dynamo, see more at " - "https://github.com/pytorch/torchdynamo.", - ) - resource_args.add_argument( - "--dynamo_mode", - type=str, - default="default", - choices=TORCH_DYNAMO_MODES, - help="Choose a mode to optimize your training with dynamo.", - ) - resource_args.add_argument( - "--dynamo_use_fullgraph", - default=False, - action="store_true", - help="Whether to use full graph mode for dynamo or it is ok to break model into several subgraphs", - ) - resource_args.add_argument( - "--dynamo_use_dynamic", - default=False, - action="store_true", - help="Whether to enable dynamic shape tracing.", - ) - resource_args.add_argument( - "--dynamo_use_regional_compilation", - default=False, - action="store_true", - help="Whether to enable regional compilation.", - ) - - # Training Paradigm arguments - paradigm_args = parser.add_argument_group( - "Training Paradigm Arguments", "Arguments for selecting which training paradigm to be used." - ) - paradigm_args.add_argument( - "--use_deepspeed", - default=False, - action="store_true", - help="Whether to use deepspeed.", - ) - paradigm_args.add_argument( - "--use_fsdp", - default=False, - action="store_true", - help="Whether to use fsdp.", - ) - paradigm_args.add_argument( - "--use_megatron_lm", - default=False, - action="store_true", - help="Whether to use Megatron-LM.", - ) - - paradigm_args.add_argument( - "--use_xpu", - default=None, - action="store_true", - help="Whether to use IPEX plugin to speed up training on XPU specifically. This argument is deprecated and ignored, will be removed in Accelerate v1.20.", - ) - - # distributed GPU training arguments - distributed_args = parser.add_argument_group("Distributed GPUs", "Arguments related to distributed GPU training.") - distributed_args.add_argument( - "--gpu_ids", - default=None, - help="What GPUs (by id) should be used for training on this machine as a comma-separated list", - ) - distributed_args.add_argument( - "--same_network", - default=False, - action="store_true", - help="Whether all machines used for multinode training exist on the same local network.", - ) - distributed_args.add_argument( - "--machine_rank", type=int, default=None, help="The rank of the machine on which this script is launched." - ) - distributed_args.add_argument( - "--main_process_ip", type=str, default=None, help="The IP address of the machine of rank 0." - ) - distributed_args.add_argument( - "--main_process_port", - type=int, - default=None, - help="The port to use to communicate with the machine of rank 0.", - ) - distributed_args.add_argument( - "-t", - "--tee", - default="0", - type=str, - help="Tee std streams into a log file and also to console.", - ) - distributed_args.add_argument( - "--log_dir", - type=str, - default=None, - help=( - "Base directory to use for log files when using torchrun/torch.distributed.run as launcher. " - "Use with --tee to redirect std streams info log files." - ), - ) - distributed_args.add_argument( - "--role", - type=str, - default="default", - help="User-defined role for the workers.", - ) - # Rendezvous related arguments - distributed_args.add_argument( - "--rdzv_backend", - type=str, - default="static", - help="The rendezvous method to use, such as 'static' (the default) or 'c10d'", - ) - distributed_args.add_argument( - "--rdzv_conf", - type=str, - default="", - help="Additional rendezvous configuration (=,=,...).", - ) - distributed_args.add_argument( - "--max_restarts", - type=int, - default=0, - help="Maximum number of worker group restarts before failing.", - ) - distributed_args.add_argument( - "--monitor_interval", - type=float, - default=0.1, - help="Interval, in seconds, to monitor the state of workers.", - ) - parser.add_argument( - "-m", - "--module", - action="store_true", - help="Change each process to interpret the launch script as a Python module, executing with the same behavior as 'python -m'.", - ) - parser.add_argument( - "--no_python", - action="store_true", - help="Skip prepending the training script with 'python' - just execute it directly. Useful when the script is not a Python script.", - ) - - # TPU arguments - tpu_args = parser.add_argument_group("TPU", "Arguments related to TPU.") - tpu_args.add_argument( - "--tpu_cluster", - action="store_true", - dest="tpu_use_cluster", - help="Whether to use a GCP TPU pod for training.", - ) - tpu_args.add_argument( - "--no_tpu_cluster", - action="store_false", - dest="tpu_use_cluster", - help="Should not be passed explicitly, this is for internal use only.", - ) - tpu_args.add_argument( - "--tpu_use_sudo", - action="store_true", - help="Whether to use `sudo` when running the TPU training script in each pod.", - ) - tpu_args.add_argument( - "--vm", - type=str, - action="append", - help=( - "List of single Compute VM instance names. " - "If not provided we assume usage of instance groups. For TPU pods." - ), - ) - tpu_args.add_argument( - "--env", - type=str, - action="append", - help="List of environment variables to set on the Compute VM instances. For TPU pods.", - ) - tpu_args.add_argument( - "--main_training_function", - type=str, - default=None, - help="The name of the main function to be executed in your script (only for TPU training).", - ) - tpu_args.add_argument( - "--downcast_bf16", - action="store_true", - help="Whether when using bf16 precision on TPUs if both float and double tensors are cast to bfloat16 or if double tensors remain as float32.", - ) - - # DeepSpeed arguments - deepspeed_args = parser.add_argument_group("DeepSpeed Arguments", "Arguments related to DeepSpeed.") - deepspeed_args.add_argument( - "--deepspeed_config_file", - default=None, - type=str, - help="DeepSpeed config file.", - ) - deepspeed_args.add_argument( - "--zero_stage", - default=None, - type=int, - help="DeepSpeed's ZeRO optimization stage (useful only when `use_deepspeed` flag is passed). " - "If unspecified, will default to `2`.", - ) - deepspeed_args.add_argument( - "--offload_optimizer_device", - default=None, - type=str, - help="Decides where (none|cpu|nvme) to offload optimizer states (useful only when `use_deepspeed` flag is passed). " - "If unspecified, will default to 'none'.", - ) - deepspeed_args.add_argument( - "--offload_param_device", - default=None, - type=str, - help="Decides where (none|cpu|nvme) to offload parameters (useful only when `use_deepspeed` flag is passed). " - "If unspecified, will default to 'none'.", - ) - deepspeed_args.add_argument( - "--offload_optimizer_nvme_path", - default=None, - type=str, - help="Decides Nvme Path to offload optimizer states (useful only when `use_deepspeed` flag is passed). " - "If unspecified, will default to 'none'.", - ) - deepspeed_args.add_argument( - "--offload_param_nvme_path", - default=None, - type=str, - help="Decides Nvme Path to offload parameters (useful only when `use_deepspeed` flag is passed). " - "If unspecified, will default to 'none'.", - ) - deepspeed_args.add_argument( - "--gradient_accumulation_steps", - default=None, - type=int, - help="No of gradient_accumulation_steps used in your training script (useful only when `use_deepspeed` flag is passed). " - "If unspecified, will default to `1`.", - ) - deepspeed_args.add_argument( - "--gradient_clipping", - default=None, - type=float, - help="gradient clipping value used in your training script (useful only when `use_deepspeed` flag is passed). " - "If unspecified, will default to `1.0`.", - ) - deepspeed_args.add_argument( - "--zero3_init_flag", - default=None, - type=str, - help="Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models. " - "Only applicable with DeepSpeed ZeRO Stage-3. If unspecified, will default to `true`.", - ) - deepspeed_args.add_argument( - "--zero3_save_16bit_model", - default=None, - type=str, - help="Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3. " - "Only applicable with DeepSpeed ZeRO Stage-3. If unspecified, will default to `false`.", - ) - deepspeed_args.add_argument( - "--deepspeed_hostfile", - default=None, - type=str, - help="DeepSpeed hostfile for configuring multi-node compute resources.", - ) - deepspeed_args.add_argument( - "--deepspeed_exclusion_filter", - default=None, - type=str, - help="DeepSpeed exclusion filter string when using mutli-node setup.", - ) - deepspeed_args.add_argument( - "--deepspeed_inclusion_filter", - default=None, - type=str, - help="DeepSpeed inclusion filter string when using mutli-node setup.", - ) - deepspeed_args.add_argument( - "--deepspeed_multinode_launcher", - default=None, - type=str, - help="DeepSpeed multi-node launcher to use, e.g. `pdsh`, `standard`, `openmpi`, `mvapich`, `mpich`, `slurm`, `nossh` (requires DeepSpeed >= 0.14.5). If unspecified, will default to `pdsh`.", - ) - deepspeed_args.add_argument( - "--deepspeed_moe_layer_cls_names", - default=None, - type=str, - help="comma-separated list of transformer MoE layer class names (case-sensitive) to wrap ,e.g, `MixtralSparseMoeBlock`, `Qwen2MoeSparseMoeBlock`, `JetMoEAttention,JetMoEBlock` ..." - " (useful only when `use_deepspeed` flag is passed).", - ) - - # fsdp arguments - fsdp_args = parser.add_argument_group("FSDP Arguments", "Arguments related to Fully Shared Data Parallelism.") - fsdp_args.add_argument( - "--fsdp_version", - type=str, - default="1", - choices=["1", "2"], - help="FSDP version to use. (useful only when `use_fsdp` flag is passed).", - ) - fsdp_args.add_argument( - "--fsdp_offload_params", - default="false", - type=str, - help="Decides Whether (true|false) to offload parameters and gradients to CPU. (useful only when `use_fsdp` flag is passed).", - ) - fsdp_args.add_argument( - "--fsdp_min_num_params", - type=int, - default=1e8, - help="FSDP's minimum number of parameters for Default Auto Wrapping. (useful only when `use_fsdp` flag is passed).", - ) - # We enable this for backwards compatibility, throw a warning if this is set in `FullyShardedDataParallelPlugin` - fsdp_args.add_argument( - "--fsdp_sharding_strategy", - type=str, - default="FULL_SHARD", - help="FSDP's sharding strategy. (useful only when `use_fsdp` flag is passed and `fsdp_version=1`).", - ) - fsdp_args.add_argument( - "--fsdp_reshard_after_forward", - type=str, - default="true", - help="FSDP's Reshard After Forward Strategy. (useful only when `use_fsdp` flag is passed). Supports either boolean (FSDP2) or `FULL_SHARD | SHARD_GRAD_OP | NO_RESHARD` (FSDP1).", - ) - fsdp_args.add_argument( - "--fsdp_auto_wrap_policy", - type=str, - default=None, - help="FSDP's auto wrap policy. (useful only when `use_fsdp` flag is passed).", - ) - fsdp_args.add_argument( - "--fsdp_transformer_layer_cls_to_wrap", - default=None, - type=str, - help="Transformer layer class name (case-sensitive) to wrap ,e.g, `BertLayer`, `GPTJBlock`, `T5Block` .... " - "(useful only when `use_fsdp` flag is passed).", - ) - fsdp_args.add_argument( - "--fsdp_backward_prefetch", - default=None, - type=str, - help="FSDP's backward prefetch policy. (useful only when `use_fsdp` flag is passed).", - ) - fsdp_args.add_argument( - "--fsdp_state_dict_type", - default=None, - type=str, - help="FSDP's state dict type. (useful only when `use_fsdp` flag is passed).", - ) - fsdp_args.add_argument( - "--fsdp_forward_prefetch", - default="false", - type=str, - help="If True, then FSDP explicitly prefetches the next upcoming " - "all-gather while executing in the forward pass (useful only when `use_fsdp` flag is passed).", - ) - fsdp_args.add_argument( - "--fsdp_use_orig_params", - default="true", - type=str, - help="If True, allows non-uniform `requires_grad` during init, which means support for interspersed frozen and trainable paramteres." - " (useful only when `use_fsdp` flag is passed).", - ) - fsdp_args.add_argument( - "--fsdp_cpu_ram_efficient_loading", - default="true", - type=str, - help="If True, only the first process loads the pretrained model checkoint while all other processes have empty weights. " - "Only applicable for 🤗 Transformers. When using this, `--fsdp_sync_module_states` needs to True. " - "(useful only when `use_fsdp` flag is passed).", - ) - fsdp_args.add_argument( - "--fsdp_sync_module_states", - default="true", - type=str, - help="If True, each individually wrapped FSDP unit will broadcast module parameters from rank 0." - " (useful only when `use_fsdp` flag is passed).", - ) - fsdp_args.add_argument( - "--fsdp_activation_checkpointing", - default="false", - type=str, - help="Decides Whether (true|false) intermediate activations are freed during the forward pass, and a checkpoint is left as a placeholder. (useful only when `use_fsdp` flag is passed).", - ) - - # megatron_lm args - megatron_lm_args = parser.add_argument_group("Megatron-LM Arguments", "Arguments related to Megatron-LM.") - megatron_lm_args.add_argument( - "--megatron_lm_tp_degree", - type=int, - default=1, - help="Megatron-LM's Tensor Parallelism (TP) degree. (useful only when `use_megatron_lm` flag is passed).", - ) - megatron_lm_args.add_argument( - "--megatron_lm_pp_degree", - type=int, - default=1, - help="Megatron-LM's Pipeline Parallelism (PP) degree. (useful only when `use_megatron_lm` flag is passed).", - ) - megatron_lm_args.add_argument( - "--megatron_lm_num_micro_batches", - type=int, - default=None, - help="Megatron-LM's number of micro batches when PP degree > 1. (useful only when `use_megatron_lm` flag is passed).", - ) - megatron_lm_args.add_argument( - "--megatron_lm_sequence_parallelism", - default=None, - type=str, - help="Decides Whether (true|false) to enable Sequence Parallelism when TP degree > 1. " - "(useful only when `use_megatron_lm` flag is passed).", - ) - megatron_lm_args.add_argument( - "--megatron_lm_recompute_activations", - default=None, - type=str, - help="Decides Whether (true|false) to enable Selective Activation Recomputation. " - "(useful only when `use_megatron_lm` flag is passed).", - ) - megatron_lm_args.add_argument( - "--megatron_lm_use_distributed_optimizer", - default=None, - type=str, - help="Decides Whether (true|false) to use distributed optimizer " - "which shards optimizer state and gradients across Data Pralellel (DP) ranks. " - "(useful only when `use_megatron_lm` flag is passed).", - ) - megatron_lm_args.add_argument( - "--megatron_lm_gradient_clipping", - default=1.0, - type=float, - help="Megatron-LM's gradient clipping value based on global L2 Norm (0 to disable). " - "(useful only when `use_megatron_lm` flag is passed).", - ) - - # FP8 arguments - fp8_args = parser.add_argument_group( - "FP8 Arguments", "Arguments related to FP8 training (requires `--mixed_precision=fp8`)" - ) - fp8_args.add_argument( - "--fp8_backend", - type=str, - choices=["te", "msamp"], - help="Choose a backend to train with FP8 (te: TransformerEngine, msamp: MS-AMP)", - ) - fp8_args.add_argument( - "--fp8_use_autocast_during_eval", - default=False, - action="store_true", - help="Whether to use FP8 autocast during eval mode (useful only when `--fp8_backend=te` is passed). Generally better metrics are found when this is not passed.", - ) - fp8_args.add_argument( - "--fp8_margin", - type=int, - default=0, - help="The margin to use for the gradient scaling (useful only when `--fp8_backend=te` is passed).", - ) - fp8_args.add_argument( - "--fp8_interval", - type=int, - default=1, - help="The interval to use for how often the scaling factor is recomputed (useful only when `--fp8_backend=te` is passed).", - ) - fp8_args.add_argument( - "--fp8_format", - type=str, - default="HYBRID", - choices=["HYBRID", "E4M3", "E5M2"], - help="The format to use for the FP8 recipe (useful only when `--fp8_backend=te` is passed).", - ) - fp8_args.add_argument( - "--fp8_amax_history_len", - type=int, - default=1024, - help="The length of the history to use for the scaling factor computation (useful only when `--fp8_backend=te` is passed).", - ) - fp8_args.add_argument( - "--fp8_amax_compute_algo", - type=str, - default="most_recent", - choices=["max", "most_recent"], - help="The algorithm to use for the scaling factor computation. (useful only when `--fp8_backend=te` is passed).", - ) - fp8_args.add_argument( - "--fp8_override_linear_precision", - type=lambda x: tuple(map(str_to_bool, x.split(","))), - default=(False, False, False), - help="Whether or not to execute `fprop`, `dgrad`, and `wgrad` GEMMS in higher precision. Should be passed in a comma-separated string of booleans (useful only when `--fp8_backend=te` is passed).", - ) - fp8_args.add_argument( - "--fp8_opt_level", - type=str, - default="O2", - choices=["O1", "O2"], - help="What level of 8-bit collective communication should be used with MS-AMP (useful only when `--fp8_backend=msamp` is passed).", - ) - - # AWS arguments - aws_args = parser.add_argument_group("AWS Arguments", "Arguments related to AWS.") - aws_args.add_argument( - "--aws_access_key_id", - type=str, - default=None, - help="The AWS_ACCESS_KEY_ID used to launch the Amazon SageMaker training job", - ) - aws_args.add_argument( - "--aws_secret_access_key", - type=str, - default=None, - help="The AWS_SECRET_ACCESS_KEY used to launch the Amazon SageMaker training job.", - ) - parser.add_argument( - "--debug", - action="store_true", - help="Whether to print out the torch.distributed stack trace when something fails.", - ) - parser.add_argument( - "training_script", - type=str, - help=( - "The full path to the script to be launched in parallel, followed by all the arguments for the training " - "script." - ), - ) - - # MPI arguments - mpirun_args = parser.add_argument_group("MPI Arguments", "Arguments related to mpirun for Multi-CPU") - mpirun_args.add_argument( - "--mpirun_hostfile", - type=str, - default=None, - help="Location for a hostfile for using Accelerate to launch a multi-CPU training job with mpirun. This will " - "get passed to the MPI --hostfile or -f parameter, depending on which MPI program is installed.", - ) - mpirun_args.add_argument( - "--mpirun_ccl", - type=int, - default=1, - help="The number of oneCCL worker threads when using Accelerate to launch multi-CPU training with mpirun.", - ) - - # Other arguments of the training scripts - parser.add_argument("training_script_args", nargs=argparse.REMAINDER, help="Arguments of the training script.") - - if subparsers is not None: - parser.set_defaults(func=launch_command) - return parser - - -def simple_launcher(args): - cmd, current_env = prepare_simple_launcher_cmd_env(args) - - process = subprocess.Popen(cmd, env=current_env) - process.wait() - if process.returncode != 0: - if not args.quiet: - raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd) - else: - sys.exit(1) - - -def multi_gpu_launcher(args): - import torch.distributed.run as distrib_run - - current_env = prepare_multi_gpu_env(args) - if not check_cuda_p2p_ib_support(): - message = "Using RTX 4000 series which doesn't support faster communication speedups. Ensuring P2P and IB communications are disabled." - warn = False - if "NCCL_P2P_DISABLE" not in current_env: - current_env["NCCL_P2P_DISABLE"] = "1" - warn = True - if "NCCL_IB_DISABLE" not in current_env: - current_env["NCCL_IB_DISABLE"] = "1" - warn = True - if warn: - logger.warning(message) - - debug = getattr(args, "debug", False) - args = _filter_args( - args, - distrib_run.get_args_parser(), - ["--training_script", args.training_script, "--training_script_args", args.training_script_args], - ) - - with patch_environment(**current_env): - try: - distrib_run.run(args) - except Exception: - if is_rich_available() and debug: - console = get_console() - console.print("\n[bold red]Using --debug, `torch.distributed` Stack Trace:[/bold red]") - console.print_exception(suppress=[__file__], show_locals=False) - else: - raise - - -def deepspeed_launcher(args): - import torch.distributed.run as distrib_run - - if not is_deepspeed_available(): - raise ImportError("DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source.") - else: - from deepspeed.launcher.runner import DEEPSPEED_ENVIRONMENT_NAME - - cmd, current_env = prepare_deepspeed_cmd_env(args) - if not check_cuda_p2p_ib_support(): - message = "Using RTX 4000 series which doesn't support faster communication speedups. Ensuring P2P and IB communications are disabled." - warn = False - if "NCCL_P2P_DISABLE" not in current_env: - current_env["NCCL_P2P_DISABLE"] = "1" - warn = True - if "NCCL_IB_DISABLE" not in current_env: - current_env["NCCL_IB_DISABLE"] = "1" - warn = True - if warn: - logger.warning(message) - - if args.num_machines > 1 and args.deepspeed_multinode_launcher != DEEPSPEED_MULTINODE_LAUNCHERS[1]: - with open(DEEPSPEED_ENVIRONMENT_NAME, "a") as f: - valid_env_items = convert_dict_to_env_variables(current_env) - if len(valid_env_items) > 1: - f.writelines(valid_env_items) - - process = subprocess.Popen(cmd, env=current_env) - process.wait() - if process.returncode != 0: - if not args.quiet: - raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd) - else: - sys.exit(1) - else: - debug = getattr(args, "debug", False) - args = _filter_args( - args, - distrib_run.get_args_parser(), - ["--training_script", args.training_script, "--training_script_args", args.training_script_args], - ) - with patch_environment(**current_env): - try: - distrib_run.run(args) - except Exception: - if is_rich_available() and debug: - console = get_console() - console.print("\n[bold red]Using --debug, `torch.distributed` Stack Trace:[/bold red]") - console.print_exception(suppress=[__file__], show_locals=False) - else: - raise - - -def tpu_launcher(args): - import torch_xla.distributed.xla_multiprocessing as xmp - - if args.no_python: - raise ValueError("--no_python cannot be used with TPU launcher") - - args, current_env = prepare_tpu(args, {}) - - if args.module: - mod_name = args.training_script - else: - # Import training_script as a module - script_path = Path(args.training_script) - sys.path.append(str(script_path.parent.resolve())) - mod_name = script_path.stem - - mod = importlib.import_module(mod_name) - if not hasattr(mod, args.main_training_function): - raise ValueError( - f"Your training script should have a function named {args.main_training_function}, or you should pass a " - "different value to `--main_training_function`." - ) - - # Patch sys.argv - sys.argv = [mod.__file__] + args.training_script_args - - main_function = getattr(mod, args.main_training_function) - with patch_environment(**current_env): - xmp.spawn(PrepareForLaunch(main_function), args=()) - - -def tpu_pod_launcher(args): - from torch_xla.distributed import xla_dist - - current_env = {} - args, current_env = prepare_tpu(args, current_env, True) - debug = getattr(args, "debug", False) - - training_script = args.training_script - training_script_args = args.training_script_args - new_args = _filter_args( - args, xla_dist.get_args_parser(), ["--tpu", args.tpu_name, "--positional", "", "--restart-tpuvm-pod-server"] - ) - - if args.tpu_use_sudo: - new_cmd = ["sudo"] - else: - new_cmd = [] - - new_cmd += [ - "accelerate-launch", - "--tpu", - "--no_tpu_cluster", - "--num_machines", - "1", - "--mixed_precision", - "no", - "--dynamo_backend", - "no", - "--num_processes", - str(args.num_processes), - "--main_training_function", - str(args.main_training_function), - training_script, - ] + training_script_args - - new_args.positional = new_cmd - bad_flags = "" - for arg in vars(new_args): - if arg.startswith("docker_"): - value = getattr(new_args, arg) - if value != "" and value is not None: - bad_flags += f'{arg}="{value}"\n' - if bad_flags != "": - raise ValueError( - f"Docker containers are not supported for TPU pod launcher currently, please remove the following flags:\n{bad_flags}" - ) - new_args.env = [f"{k}={v}" for k, v in current_env.items()] - new_args.env.append("ACCELERATE_IN_TPU_POD=1") - try: - xla_dist.resolve_and_execute(new_args) - except Exception: - if is_rich_available() and debug: - console = get_console() - console.print("\n[bold red]Using --debug, `torch_xla.xla_dist` Stack Trace:[/bold red]") - console.print_exception(suppress=[__file__], show_locals=False) - else: - raise - - -def sagemaker_launcher(sagemaker_config: SageMakerConfig, args): - if not is_sagemaker_available(): - raise ImportError( - "Please install sagemaker to be able to launch training on Amazon SageMaker with `pip install accelerate[sagemaker]`" - ) - if args.module or args.no_python: - raise ValueError( - "SageMaker requires a python training script file and cannot be used with --module or --no_python" - ) - - from sagemaker.huggingface import HuggingFace - - args, sagemaker_inputs = prepare_sagemager_args_inputs(sagemaker_config, args) - - huggingface_estimator = HuggingFace(**args) - - huggingface_estimator.fit(inputs=sagemaker_inputs) - print(f"You can find your model data at: {huggingface_estimator.model_data}") - - -def _validate_launch_command(args): - # Sanity checks - if sum([args.multi_gpu, args.cpu, args.tpu, args.use_deepspeed, args.use_fsdp]) > 1: - raise ValueError( - "You can only use one of `--cpu`, `--multi_gpu`, `--tpu`, `--use_deepspeed`, `--use_fsdp` at a time." - ) - if args.multi_gpu and (args.num_processes is not None) and (args.num_processes < 2): - raise ValueError("You need to use at least 2 processes to use `--multi_gpu`.") - - defaults = None - warned = [] - mp_from_config_flag = False - # Get the default from the config file. - if args.config_file is not None or os.path.isfile(default_config_file) and not args.cpu: - defaults = load_config_from_file(args.config_file) - if ( - not args.multi_gpu - and not args.tpu - and not args.tpu_use_cluster - and not args.use_deepspeed - and not args.use_fsdp - and not args.use_megatron_lm - ): - args.use_deepspeed = defaults.distributed_type == DistributedType.DEEPSPEED - args.multi_gpu = ( - True - if defaults.distributed_type - in ( - DistributedType.MULTI_GPU, - DistributedType.MULTI_NPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_XPU, - DistributedType.MULTI_HPU, - ) - else False - ) - args.tpu = defaults.distributed_type == DistributedType.XLA - args.use_fsdp = defaults.distributed_type == DistributedType.FSDP - args.use_megatron_lm = defaults.distributed_type == DistributedType.MEGATRON_LM - args.tpu_use_cluster = defaults.tpu_use_cluster if args.tpu else False - if args.gpu_ids is None: - if defaults.gpu_ids is not None: - args.gpu_ids = defaults.gpu_ids - else: - args.gpu_ids = "all" - - if args.multi_gpu and args.num_machines is None: - args.num_machines = defaults.num_machines - - if len(args.gpu_ids.split(",")) < 2 and (args.gpu_ids != "all") and args.multi_gpu and args.num_machines <= 1: - raise ValueError( - "Less than two GPU ids were configured and tried to run on on multiple GPUs. " - "Please ensure at least two are specified for `--gpu_ids`, or use `--gpu_ids='all'`." - ) - if defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE: - # Update args with the defaults - for name, attr in defaults.__dict__.items(): - if isinstance(attr, dict): - # Copy defaults.somedict.somearg to args.somearg and - # defaults.fsdp_config.x to args.fsdp_x - for key, value in attr.items(): - if name == "fsdp_config" and not key.startswith("fsdp"): - key = "fsdp_" + key - elif name == "fp8_config" and not key.startswith("fp8"): - key = "fp8_" + key - if hasattr(args, "nondefault") and key not in args.nondefault: - setattr(args, key, value) - elif ( - name not in ["compute_environment", "mixed_precision", "distributed_type"] - and getattr(args, name, None) is None - ): - # Those args are handled separately - setattr(args, name, attr) - if not args.debug: - args.debug = defaults.debug - - if not args.mixed_precision: - if defaults.mixed_precision is None: - args.mixed_precision = "no" - else: - args.mixed_precision = defaults.mixed_precision - mp_from_config_flag = True - else: - native_amp = is_bf16_available(True) - if ( - args.mixed_precision == "bf16" - and not native_amp - and not (args.tpu and is_torch_xla_available(check_is_tpu=True)) - ): - raise ValueError("bf16 mixed precision requires PyTorch >= 1.10 and a supported device.") - - # Silently set the default here - if args.dynamo_backend is None: - args.dynamo_backend = "no" - if args.num_processes == -1: - raise ValueError("You need to manually pass in `--num_processes` using this config yaml.") - else: - if args.num_processes is None: - if is_xpu_available(): - args.num_processes = torch.xpu.device_count() - elif is_mlu_available(): - args.num_processes = torch.mlu.device_count() - elif is_sdaa_available(): - args.num_processes = torch.sdaa.device_count() - elif is_musa_available(): - args.num_processes = torch.musa.device_count() - elif is_npu_available(): - args.num_processes = torch.npu.device_count() - elif is_hpu_available(): - args.num_processes = torch.hpu.device_count() - else: - args.num_processes = torch.cuda.device_count() - warned.append(f"\t`--num_processes` was set to a value of `{args.num_processes}`") - if args.debug is None: - args.debug = False - if ( - not args.multi_gpu - and args.num_processes > 1 - and ( - (is_xpu_available() and torch.xpu.device_count() > 1) - or (is_npu_available() and torch.npu.device_count() > 1) - or (is_hpu_available() and torch.hpu.device_count() > 1) - or (is_mlu_available() and torch.mlu.device_count() > 1) - or (is_sdaa_available() and torch.sdaa.device_count() > 1) - or (is_musa_available() and torch.musa.device_count() > 1) - or (torch.cuda.is_available() and torch.cuda.device_count() > 1) - ) - ): - warned.append( - "\t\tMore than one GPU was found, enabling multi-GPU training.\n" - "\t\tIf this was unintended please pass in `--num_processes=1`." - ) - args.multi_gpu = True - if args.num_machines is None: - warned.append("\t`--num_machines` was set to a value of `1`") - args.num_machines = 1 - if args.mixed_precision is None: - warned.append("\t`--mixed_precision` was set to a value of `'no'`") - args.mixed_precision = "no" - if not hasattr(args, "use_cpu"): - args.use_cpu = args.cpu - if args.dynamo_backend is None: - warned.append("\t`--dynamo_backend` was set to a value of `'no'`") - args.dynamo_backend = "no" - if args.debug: - logger.debug("Running script in debug mode, expect distributed operations to be slightly slower.") - - is_aws_env_disabled = defaults is None or ( - defaults is not None and defaults.compute_environment != ComputeEnvironment.AMAZON_SAGEMAKER - ) - if is_aws_env_disabled and args.num_cpu_threads_per_process is None: - args.num_cpu_threads_per_process = get_int_from_env(["OMP_NUM_THREADS"], 1) - if args.use_cpu and args.num_processes >= 1 and get_int_from_env(["OMP_NUM_THREADS"], 0) == 0: - local_size = get_int_from_env( - ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], - max(int(args.num_processes / args.num_machines), 1), - ) - threads_per_process = int(psutil.cpu_count(logical=False) / local_size) - if threads_per_process > 1: - args.num_cpu_threads_per_process = threads_per_process - warned.append( - f"\t`--num_cpu_threads_per_process` was set to `{args.num_cpu_threads_per_process}` to improve out-of-box performance when training on CPUs" - ) - - if args.ipex is not None: - logger.warning( - "ipex flag is deprecated, will be removed in Accelerate v1.10. " - "From 2.7.0, PyTorch has all needed optimizations for Intel CPU and XPU." - ) - - if args.use_xpu is not None: - logger.warning( - "use_xpu is deprecated and ignored, will be removed in Accelerate v1.20. " - "XPU is a PyTorch native citizen now, we don't need extra argument to enable it any more." - ) - - if any(warned): - message = "The following values were not passed to `accelerate launch` and had defaults used instead:\n" - message += "\n".join(warned) - message += ( - "\nTo avoid this warning pass in values for each of the problematic parameters or run `accelerate config`." - ) - logger.warning(message) - return args, defaults, mp_from_config_flag - - -def launch_command(args): - args, defaults, mp_from_config_flag = _validate_launch_command(args) - # Use the proper launcher - if args.use_deepspeed and not args.cpu: - args.deepspeed_fields_from_accelerate_config = list(defaults.deepspeed_config.keys()) if defaults else [] - if mp_from_config_flag: - args.deepspeed_fields_from_accelerate_config.append("mixed_precision") - args.deepspeed_fields_from_accelerate_config = ",".join(args.deepspeed_fields_from_accelerate_config) - deepspeed_launcher(args) - elif args.use_fsdp and not args.cpu: - multi_gpu_launcher(args) - elif args.use_megatron_lm and not args.cpu: - multi_gpu_launcher(args) - elif args.multi_gpu and not args.cpu: - multi_gpu_launcher(args) - elif args.tpu and not args.cpu: - if args.tpu_use_cluster: - tpu_pod_launcher(args) - else: - tpu_launcher(args) - elif defaults is not None and defaults.compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER: - sagemaker_launcher(defaults, args) - else: - simple_launcher(args) - - -def main(): - parser = launch_command_parser() - args = parser.parse_args() - launch_command(args) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__init__.py b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__init__.py deleted file mode 100644 index c2c851cc0b192ab8207d3fa68d7409868c84354c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright 2022 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 .selection_menu import BulletMenu diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index c565b5bf6119aca3772d5bfbf7413927cbeb7594..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/cursor.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/cursor.cpython-310.pyc deleted file mode 100644 index 11616afac2865b4ba21ef276ea052b3085d50725..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/cursor.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/helpers.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/helpers.cpython-310.pyc deleted file mode 100644 index 375f54ea3b9b20839c4506226c7ed1ac366f3220..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/helpers.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/input.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/input.cpython-310.pyc deleted file mode 100644 index a90a63e8fe6a7dd0852c9070aa7d3041aca0c473..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/input.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/keymap.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/keymap.cpython-310.pyc deleted file mode 100644 index 4ccbd09290b2217c4694527eb1373bff993a99d9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/keymap.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/selection_menu.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/selection_menu.cpython-310.pyc deleted file mode 100644 index 3abda192fb278fbc69b88e573a2a5a4d78a7117b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/commands/menu/__pycache__/selection_menu.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/cursor.py b/venv/lib/python3.10/site-packages/accelerate/commands/menu/cursor.py deleted file mode 100644 index c1f0bb7b68025ae4fe0c2c76c095eb36b4e64f2c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/menu/cursor.py +++ /dev/null @@ -1,65 +0,0 @@ -# Copyright 2022 The HuggingFace Team and Brian Chao. 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. - -""" -A utility for showing and hiding the terminal cursor on Windows and Linux, based on https://github.com/bchao1/bullet -""" - -import os -import sys -from contextlib import contextmanager - - -# Windows only -if os.name == "nt": - import ctypes - import msvcrt # noqa - - class CursorInfo(ctypes.Structure): - # _fields is a specific attr expected by ctypes - _fields_ = [("size", ctypes.c_int), ("visible", ctypes.c_byte)] - - -def hide_cursor(): - if os.name == "nt": - ci = CursorInfo() - handle = ctypes.windll.kernel32.GetStdHandle(-11) - ctypes.windll.kernel32.GetConsoleCursorInfo(handle, ctypes.byref(ci)) - ci.visible = False - ctypes.windll.kernel32.SetConsoleCursorInfo(handle, ctypes.byref(ci)) - elif os.name == "posix": - sys.stdout.write("\033[?25l") - sys.stdout.flush() - - -def show_cursor(): - if os.name == "nt": - ci = CursorInfo() - handle = ctypes.windll.kernel32.GetStdHandle(-11) - ctypes.windll.kernel32.GetConsoleCursorInfo(handle, ctypes.byref(ci)) - ci.visible = True - ctypes.windll.kernel32.SetConsoleCursorInfo(handle, ctypes.byref(ci)) - elif os.name == "posix": - sys.stdout.write("\033[?25h") - sys.stdout.flush() - - -@contextmanager -def hide(): - "Context manager to hide the terminal cursor" - try: - hide_cursor() - yield - finally: - show_cursor() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/helpers.py b/venv/lib/python3.10/site-packages/accelerate/commands/menu/helpers.py deleted file mode 100644 index de46f37ddcf4591167e3e01791391e4b1729034f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/menu/helpers.py +++ /dev/null @@ -1,59 +0,0 @@ -# Copyright 2022 The HuggingFace Team and Brian Chao. 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. - -""" -A variety of helper functions and constants when dealing with terminal menu choices, based on -https://github.com/bchao1/bullet -""" - -import enum -import shutil -import sys - - -TERMINAL_WIDTH, _ = shutil.get_terminal_size() - -CURSOR_TO_CHAR = {"UP": "A", "DOWN": "B", "RIGHT": "C", "LEFT": "D"} - - -class Direction(enum.Enum): - UP = 0 - DOWN = 1 - - -def forceWrite(content, end=""): - sys.stdout.write(str(content) + end) - sys.stdout.flush() - - -def writeColor(content, color, end=""): - forceWrite(f"\u001b[{color}m{content}\u001b[0m", end) - - -def reset_cursor(): - forceWrite("\r") - - -def move_cursor(num_lines: int, direction: str): - forceWrite(f"\033[{num_lines}{CURSOR_TO_CHAR[direction.upper()]}") - - -def clear_line(): - forceWrite(" " * TERMINAL_WIDTH) - reset_cursor() - - -def linebreak(): - reset_cursor() - forceWrite("-" * TERMINAL_WIDTH) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/input.py b/venv/lib/python3.10/site-packages/accelerate/commands/menu/input.py deleted file mode 100644 index f1270eaece9d4243e7282dcb31166feeeb9bdfc1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/menu/input.py +++ /dev/null @@ -1,84 +0,0 @@ -# Copyright 2022 The HuggingFace Team and Brian Chao. 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. - -""" -This file contains utilities for handling input from the user and registering specific keys to specific functions, -based on https://github.com/bchao1/bullet -""" - -from .keymap import KEYMAP, get_character - - -def mark(key: str): - """ - Mark the function with the key code so it can be handled in the register - """ - - def decorator(func): - handle = getattr(func, "handle_key", []) - handle += [key] - func.handle_key = handle - return func - - return decorator - - -def mark_multiple(*keys: list[str]): - """ - Mark the function with the key codes so it can be handled in the register - """ - - def decorator(func): - handle = getattr(func, "handle_key", []) - handle += keys - func.handle_key = handle - return func - - return decorator - - -class KeyHandler(type): - """ - Metaclass that adds the key handlers to the class - """ - - def __new__(cls, name, bases, attrs): - new_cls = super().__new__(cls, name, bases, attrs) - if not hasattr(new_cls, "key_handler"): - new_cls.key_handler = {} - new_cls.handle_input = KeyHandler.handle_input - - for value in attrs.values(): - handled_keys = getattr(value, "handle_key", []) - for key in handled_keys: - new_cls.key_handler[key] = value - return new_cls - - @staticmethod - def handle_input(cls): - "Finds and returns the selected character if it exists in the handler" - char = get_character() - if char != KEYMAP["undefined"]: - char = ord(char) - handler = cls.key_handler.get(char) - if handler: - cls.current_selection = char - return handler(cls) - else: - return None - - -def register(cls): - """Adds KeyHandler metaclass to the class""" - return KeyHandler(cls.__name__, cls.__bases__, cls.__dict__.copy()) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/keymap.py b/venv/lib/python3.10/site-packages/accelerate/commands/menu/keymap.py deleted file mode 100644 index 787db12860fe21c6786dda69c34fcccab114f2f8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/menu/keymap.py +++ /dev/null @@ -1,133 +0,0 @@ -# Copyright 2022 The HuggingFace Team and Brian Chao. 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. - -""" -Utilities relating to parsing raw characters from the keyboard, based on https://github.com/bchao1/bullet -""" - -import os -import string -import sys - - -ARROW_KEY_FLAG = 1 << 8 - -KEYMAP = { - "tab": ord("\t"), - "newline": ord("\r"), - "esc": 27, - "up": 65 + ARROW_KEY_FLAG, - "down": 66 + ARROW_KEY_FLAG, - "right": 67 + ARROW_KEY_FLAG, - "left": 68 + ARROW_KEY_FLAG, - "mod_int": 91, - "undefined": sys.maxsize, - "interrupt": 3, - "insert": 50, - "delete": 51, - "pg_up": 53, - "pg_down": 54, -} - -KEYMAP["arrow_begin"] = KEYMAP["up"] -KEYMAP["arrow_end"] = KEYMAP["left"] - -if sys.platform == "win32": - WIN_CH_BUFFER = [] - WIN_KEYMAP = { - b"\xe0H": KEYMAP["up"] - ARROW_KEY_FLAG, - b"\x00H": KEYMAP["up"] - ARROW_KEY_FLAG, - b"\xe0P": KEYMAP["down"] - ARROW_KEY_FLAG, - b"\x00P": KEYMAP["down"] - ARROW_KEY_FLAG, - b"\xe0M": KEYMAP["right"] - ARROW_KEY_FLAG, - b"\x00M": KEYMAP["right"] - ARROW_KEY_FLAG, - b"\xe0K": KEYMAP["left"] - ARROW_KEY_FLAG, - b"\x00K": KEYMAP["left"] - ARROW_KEY_FLAG, - } - -for i in range(10): - KEYMAP[str(i)] = ord(str(i)) - - -def get_raw_chars(): - "Gets raw characters from inputs" - if os.name == "nt": - import msvcrt - - encoding = "mbcs" - # Flush the keyboard buffer - while msvcrt.kbhit(): - msvcrt.getch() - if len(WIN_CH_BUFFER) == 0: - # Read the keystroke - ch = msvcrt.getch() - - # If it is a prefix char, get second part - if ch in (b"\x00", b"\xe0"): - ch2 = ch + msvcrt.getch() - # Translate actual Win chars to bullet char types - try: - chx = chr(WIN_KEYMAP[ch2]) - WIN_CH_BUFFER.append(chr(KEYMAP["mod_int"])) - WIN_CH_BUFFER.append(chx) - if ord(chx) in ( - KEYMAP["insert"] - 1 << 9, - KEYMAP["delete"] - 1 << 9, - KEYMAP["pg_up"] - 1 << 9, - KEYMAP["pg_down"] - 1 << 9, - ): - WIN_CH_BUFFER.append(chr(126)) - ch = chr(KEYMAP["esc"]) - except KeyError: - ch = ch2[1] - else: - ch = ch.decode(encoding) - else: - ch = WIN_CH_BUFFER.pop(0) - elif os.name == "posix": - import termios - import tty - - fd = sys.stdin.fileno() - old_settings = termios.tcgetattr(fd) - try: - tty.setraw(fd) - ch = sys.stdin.read(1) - finally: - termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) - return ch - - -def get_character(): - "Gets a character from the keyboard and returns the key code" - char = get_raw_chars() - if ord(char) in [KEYMAP["interrupt"], KEYMAP["newline"]]: - return char - - elif ord(char) == KEYMAP["esc"]: - combo = get_raw_chars() - if ord(combo) == KEYMAP["mod_int"]: - key = get_raw_chars() - if ord(key) >= KEYMAP["arrow_begin"] - ARROW_KEY_FLAG and ord(key) <= KEYMAP["arrow_end"] - ARROW_KEY_FLAG: - return chr(ord(key) + ARROW_KEY_FLAG) - else: - return KEYMAP["undefined"] - else: - return get_raw_chars() - - else: - if char in string.printable: - return char - else: - return KEYMAP["undefined"] diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/menu/selection_menu.py b/venv/lib/python3.10/site-packages/accelerate/commands/menu/selection_menu.py deleted file mode 100644 index ee9a771a54ef666ee46b67ae6c75fb957d49efdd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/menu/selection_menu.py +++ /dev/null @@ -1,144 +0,0 @@ -# Copyright 2022 The HuggingFace Team and Brian Chao. 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. - -""" -Main driver for the selection menu, based on https://github.com/bchao1/bullet -""" - -import builtins -import sys - -from ...utils.imports import _is_package_available -from . import cursor, input -from .helpers import Direction, clear_line, forceWrite, linebreak, move_cursor, reset_cursor, writeColor -from .keymap import KEYMAP - - -in_colab = False -try: - in_colab = _is_package_available("google.colab") -except ModuleNotFoundError: - pass - - -@input.register -class BulletMenu: - """ - A CLI menu to select a choice from a list of choices using the keyboard. - """ - - def __init__(self, prompt: str = None, choices: list = []): - self.position = 0 - self.choices = choices - self.prompt = prompt - if sys.platform == "win32": - self.arrow_char = "*" - else: - self.arrow_char = "➔ " - - def write_choice(self, index, end: str = ""): - if sys.platform != "win32": - writeColor(self.choices[index], 32, end) - else: - forceWrite(self.choices[index], end) - - def print_choice(self, index: int): - "Prints the choice at the given index" - if index == self.position: - forceWrite(f" {self.arrow_char} ") - self.write_choice(index) - else: - forceWrite(f" {self.choices[index]}") - reset_cursor() - - def move_direction(self, direction: Direction, num_spaces: int = 1): - "Should not be directly called, used to move a direction of either up or down" - old_position = self.position - if direction == Direction.DOWN: - if self.position + 1 >= len(self.choices): - return - self.position += num_spaces - else: - if self.position - 1 < 0: - return - self.position -= num_spaces - clear_line() - self.print_choice(old_position) - move_cursor(num_spaces, direction.name) - self.print_choice(self.position) - - @input.mark(KEYMAP["up"]) - def move_up(self): - self.move_direction(Direction.UP) - - @input.mark(KEYMAP["down"]) - def move_down(self): - self.move_direction(Direction.DOWN) - - @input.mark(KEYMAP["newline"]) - def select(self): - move_cursor(len(self.choices) - self.position, "DOWN") - return self.position - - @input.mark(KEYMAP["interrupt"]) - def interrupt(self): - move_cursor(len(self.choices) - self.position, "DOWN") - raise KeyboardInterrupt - - @input.mark_multiple(*[KEYMAP[str(number)] for number in range(10)]) - def select_row(self): - index = int(chr(self.current_selection)) - movement = index - self.position - if index == self.position: - return - if index < len(self.choices): - if self.position > index: - self.move_direction(Direction.UP, -movement) - elif self.position < index: - self.move_direction(Direction.DOWN, movement) - else: - return - else: - return - - def run(self, default_choice: int = 0): - "Start the menu and return the selected choice" - if self.prompt: - linebreak() - forceWrite(self.prompt, "\n") - if in_colab: - forceWrite("Please input a choice index (starting from 0), and press enter", "\n") - else: - forceWrite("Please select a choice using the arrow or number keys, and selecting with enter", "\n") - self.position = default_choice - for i in range(len(self.choices)): - self.print_choice(i) - forceWrite("\n") - move_cursor(len(self.choices) - self.position, "UP") - with cursor.hide(): - while True: - if in_colab: - try: - choice = int(builtins.input()) - except ValueError: - choice = default_choice - else: - choice = self.handle_input() - if choice is not None: - reset_cursor() - for _ in range(len(self.choices) + 1): - move_cursor(1, "UP") - clear_line() - self.write_choice(choice, "\n") - return choice diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/merge.py b/venv/lib/python3.10/site-packages/accelerate/commands/merge.py deleted file mode 100644 index 475b53b5bbb71b959057126f8667d7f61eb9d0e1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/merge.py +++ /dev/null @@ -1,69 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2024 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 accelerate.commands.utils import CustomArgumentParser -from accelerate.utils import merge_fsdp_weights - - -description = """Utility to merge the weights from multiple FSDP checkpoints into a single combined checkpoint. Should be used if -`SHARDED_STATE_DICT` was used for the model. Weights will be saved to `{output_path}`. - -This is a CPU-bound process and requires enough RAM to load the entire model state dict.""" - - -def merge_command(args): - merge_fsdp_weights( - args.checkpoint_directory, args.output_path, not args.unsafe_serialization, args.remove_checkpoint_dir - ) - - -def merge_command_parser(subparsers=None): - if subparsers is not None: - parser = subparsers.add_parser("merge-weights", description=description) - else: - parser = CustomArgumentParser(description=description) - - parser.add_argument("checkpoint_directory", type=str, help="A directory containing sharded weights saved by FSDP.") - parser.add_argument( - "output_path", - type=str, - help="The path to save the merged weights. Defaults to the current directory. ", - ) - parser.add_argument( - "--unsafe_serialization", - action="store_true", - default=False, - help="Whether to save the merged weights as `.bin` rather than `.safetensors` (not recommended).", - ) - parser.add_argument( - "--remove_checkpoint_dir", - action="store_true", - help="Whether to remove the checkpoint directory after merging.", - default=False, - ) - - if subparsers is not None: - parser.set_defaults(func=merge_command) - return parser - - -def main(): - parser = merge_command_parser() - args = parser.parse_args() - merge_command(args) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/test.py b/venv/lib/python3.10/site-packages/accelerate/commands/test.py deleted file mode 100644 index a0d2f7bcf14727aa13e3438f4cd6e6f140f5bb2f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/test.py +++ /dev/null @@ -1,65 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2021 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 argparse - -from accelerate.test_utils import execute_subprocess_async, path_in_accelerate_package - - -def test_command_parser(subparsers=None): - if subparsers is not None: - parser = subparsers.add_parser("test") - else: - parser = argparse.ArgumentParser("Accelerate test command") - - parser.add_argument( - "--config_file", - default=None, - help=( - "The path to use to store the config file. Will default to a file named default_config.yaml in the cache " - "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have " - "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed " - "with 'huggingface'." - ), - ) - - if subparsers is not None: - parser.set_defaults(func=test_command) - return parser - - -def test_command(args): - script_name = path_in_accelerate_package("test_utils", "scripts", "test_script.py") - - if args.config_file is None: - test_args = [script_name] - else: - test_args = f"--config_file={args.config_file} {script_name}".split() - - cmd = ["accelerate-launch"] + test_args - result = execute_subprocess_async(cmd) - if result.returncode == 0: - print("Test is a success! You are ready for your distributed training!") - - -def main(): - parser = test_command_parser() - args = parser.parse_args() - test_command(args) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/to_fsdp2.py b/venv/lib/python3.10/site-packages/accelerate/commands/to_fsdp2.py deleted file mode 100644 index 443407cd983dcb31711b75c2a6337f9a7af24584..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/to_fsdp2.py +++ /dev/null @@ -1,172 +0,0 @@ -#!/usr/bin/env python - -# 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 enum -import logging -from pathlib import Path - -import yaml - -from accelerate.commands.utils import CustomArgumentParser - - -class ConversionStatus(enum.Enum): - NOT_YET_IMPLEMENTED = 0 - REMOVED = -1 - - -ARGUMENT_KEY_MAPPING = { - # New keys in FSDP2 - "fsdp_version": "fsdp_version", - "fsdp_reshard_after_forward": "fsdp_reshard_after_forward", - # https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md - # https://huggingface.co/docs/accelerate/en/usage_guides/fsdp - "fsdp_auto_wrap_policy": "fsdp_auto_wrap_policy", - "fsdp_backward_prefetch": ConversionStatus.REMOVED, - "fsdp_forward_prefetch": ConversionStatus.NOT_YET_IMPLEMENTED, - "fsdp_cpu_ram_efficient_loading": "fsdp_cpu_ram_efficient_loading", - "fsdp_offload_params": "fsdp_offload_params", - "fsdp_sharding_strategy": "fsdp_reshard_after_forward", - "fsdp_state_dict_type": "fsdp_state_dict_type", - "fsdp_sync_module_states": ConversionStatus.REMOVED, - "fsdp_transformer_layer_cls_to_wrap": "fsdp_transformer_layer_cls_to_wrap", - "fsdp_min_num_params": "fsdp_min_num_params", - "fsdp_use_orig_params": ConversionStatus.REMOVED, - "fsdp_activation_checkpointing": "fsdp_activation_checkpointing", -} - -ARGUMENT_VALUE_MAPPING = { - "fsdp_sharding_strategy": { - "FULL_SHARD": True, - "SHARD_GRAD_OP": False, - "HYBRID_SHARD": True, - "HYBRID_SHARD_ZERO2": False, - "NO_SHARD": False, - }, - "fsdp_reshard_after_forward": { # Needed to convert newly created configs using FSDP1 to FSDP2 - "FULL_SHARD": True, - "SHARD_GRAD_OP": False, - "HYBRID_SHARD": True, - "HYBRID_SHARD_ZERO2": False, - "NO_SHARD": False, - }, -} - -logger = logging.getLogger(__name__) - - -def _validate_to_fsdp2_args(args): - if not Path(args.config_file).exists(): - raise FileNotFoundError(f"Config file {args.config_file} not found") - - if not args.overwrite and args.output_file is None: - raise ValueError("If --overwrite is not set, --output_file must be provided") - - if not args.overwrite and Path(args.output_file).exists(): - raise FileExistsError(f"Output file {args.output_file} already exists and --overwrite is not set") - - -def convert_config_to_fsdp2(config: dict) -> dict: - fsdp_config = config.get("fsdp_config", {}) - - if not fsdp_config: - logger.info("No FSDP config found in the config file, skipping conversion...") - return config - - new_fsdp_config = {} - - if fsdp_config.get("fsdp_version", 1) == 2: - logger.warning("Config already specfies FSDP2, skipping conversion...") - logger.warning( - "If the config doesn't use new argument names, change `fsdp_version` to `1` and rerun the command." - ) - return config - - for key, value in fsdp_config.items(): - conversion_status = ARGUMENT_KEY_MAPPING.get(key, None) - if isinstance(conversion_status, ConversionStatus) or conversion_status is None: - conversion_status = key - new_fsdp_config[conversion_status] = value - continue - - if conversion_status == ConversionStatus.REMOVED: - logger.warning(f"Argument {key} has been removed in FSDP2, skipping this key...") - continue - - if conversion_status == ConversionStatus.NOT_YET_IMPLEMENTED: - logger.warning(f"Argument {key} is not yet implemented in FSDP2, skipping this key...") - continue - - if conversion_status is None: - logger.warning(f"Argument {key} is not being converted, skipping this key...") - new_fsdp_config[key] = value - else: - if key in ARGUMENT_VALUE_MAPPING: - value = ARGUMENT_VALUE_MAPPING[key].get(value, value) - new_fsdp_config[ARGUMENT_KEY_MAPPING[key]] = value - - new_fsdp_config["fsdp_version"] = 2 - config["fsdp_config"] = new_fsdp_config - return config - - -def to_fsdp2_command_parser(subparsers=None): - description = "Convert an Accelerate config from FSDP1 to FSDP2" - - if subparsers is not None: - parser = subparsers.add_parser("to-fsdp2", description=description) - else: - parser = CustomArgumentParser(description=description) - - parser.add_argument("--config_file", type=str, help="The config file to convert to FSDP2", required=True) - parser.add_argument( - "--overwrite", - action="store_true", - help="Overwrite the config file if it exists", - default=False, - ) - parser.add_argument( - "--output_file", - type=str, - help="The path to the output file to write the converted config to. If not provided, the input file will be overwritten (if --overwrite is set)", - default=None, - ) - if subparsers is not None: - parser.set_defaults(func=to_fsdp2_command) - - return parser - - -def load_config(config_file: str) -> dict: - with open(config_file) as f: - config = yaml.safe_load(f) - if not config: - raise ValueError("Config file is empty") - - return config - - -def to_fsdp2_command(args): - _validate_to_fsdp2_args(args) - config = load_config(args.config_file) - - if args.overwrite and args.output_file is None: - args.output_file = args.config_file - - new_config = convert_config_to_fsdp2(config) - - with open(args.output_file, "w") as f: - yaml.dump(new_config, f) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/tpu.py b/venv/lib/python3.10/site-packages/accelerate/commands/tpu.py deleted file mode 100644 index fc0f07bf8697bfdb6484d3bf817f2e18b1313b00..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/tpu.py +++ /dev/null @@ -1,157 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2022 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 argparse -import os -import subprocess - -from packaging.version import Version, parse - -from accelerate.commands.config.config_args import default_config_file, load_config_from_file - - -_description = "Run commands across TPU VMs for initial setup before running `accelerate launch`." - - -def tpu_command_parser(subparsers=None): - if subparsers is not None: - parser = subparsers.add_parser("tpu-config", description=_description) - else: - parser = argparse.ArgumentParser("Accelerate tpu-config command", description=_description) - # Core arguments - config_args = parser.add_argument_group( - "Config Arguments", "Arguments that can be configured through `accelerate config`." - ) - config_args.add_argument( - "--config_file", - type=str, - default=None, - help="Path to the config file to use for accelerate.", - ) - config_args.add_argument( - "--tpu_name", - default=None, - help="The name of the TPU to use. If not specified, will use the TPU specified in the config file.", - ) - config_args.add_argument( - "--tpu_zone", - default=None, - help="The zone of the TPU to use. If not specified, will use the zone specified in the config file.", - ) - pod_args = parser.add_argument_group("TPU Arguments", "Arguments for options ran inside the TPU.") - pod_args.add_argument( - "--use_alpha", - action="store_true", - help="Whether to use `gcloud alpha` when running the TPU training script instead of `gcloud`.", - ) - pod_args.add_argument( - "--command_file", - default=None, - help="The path to the file containing the commands to run on the pod on startup.", - ) - pod_args.add_argument( - "--command", - action="append", - nargs="+", - help="A command to run on the pod. Can be passed multiple times.", - ) - pod_args.add_argument( - "--install_accelerate", - action="store_true", - help="Whether to install accelerate on the pod. Defaults to False.", - ) - pod_args.add_argument( - "--accelerate_version", - default="latest", - help="The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify 'dev' to install from GitHub.", - ) - pod_args.add_argument( - "--debug", action="store_true", help="If set, will print the command that would be run instead of running it." - ) - - if subparsers is not None: - parser.set_defaults(func=tpu_command_launcher) - return parser - - -def tpu_command_launcher(args): - defaults = None - - # Get the default from the config file if it exists. - if args.config_file is not None or os.path.isfile(default_config_file): - defaults = load_config_from_file(args.config_file) - if not args.command_file and defaults.command_file is not None and not args.command: - args.command_file = defaults.command_file - if not args.command and defaults.commands is not None: - args.command = defaults.commands - if not args.tpu_name: - args.tpu_name = defaults.tpu_name - if not args.tpu_zone: - args.tpu_zone = defaults.tpu_zone - if args.accelerate_version == "dev": - args.accelerate_version = "git+https://github.com/huggingface/accelerate.git" - elif args.accelerate_version == "latest": - args.accelerate_version = "accelerate -U" - elif isinstance(parse(args.accelerate_version), Version): - args.accelerate_version = f"accelerate=={args.accelerate_version}" - - if not args.command_file and not args.command: - raise ValueError("You must specify either a command file or a command to run on the pod.") - - if args.command_file: - with open(args.command_file) as f: - args.command = [f.read().splitlines()] - - # To turn list of lists into list of strings - if isinstance(args.command[0], list): - args.command = [line for cmd in args.command for line in cmd] - # Default to the shared folder and install accelerate - new_cmd = ["cd /usr/share"] - if args.install_accelerate: - new_cmd += [f"pip install {args.accelerate_version}"] - new_cmd += args.command - args.command = "; ".join(new_cmd) - - # Then send it to gcloud - # Eventually try to use google-api-core to do this instead of subprocess - cmd = ["gcloud"] - if args.use_alpha: - cmd += ["alpha"] - cmd += [ - "compute", - "tpus", - "tpu-vm", - "ssh", - args.tpu_name, - "--zone", - args.tpu_zone, - "--command", - args.command, - "--worker", - "all", - ] - if args.debug: - print(f"Running {' '.join(cmd)}") - return - subprocess.run(cmd) - print("Successfully setup pod.") - - -def main(): - parser = tpu_command_parser() - args = parser.parse_args() - - tpu_command_launcher(args) diff --git a/venv/lib/python3.10/site-packages/accelerate/commands/utils.py b/venv/lib/python3.10/site-packages/accelerate/commands/utils.py deleted file mode 100644 index 326f37d7f93de2417e4171e5ffe91193fb97225c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/commands/utils.py +++ /dev/null @@ -1,123 +0,0 @@ -# Copyright 2024 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 argparse - - -class _StoreAction(argparse.Action): - """ - Custom action that allows for `-` or `_` to be passed in for an argument. - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - new_option_strings = [] - for option_string in self.option_strings: - new_option_strings.append(option_string) - if "_" in option_string[2:]: - # Add `-` version to the option string - new_option_strings.append(option_string.replace("_", "-")) - self.option_strings = new_option_strings - - def __call__(self, parser, namespace, values, option_string=None): - setattr(namespace, self.dest, values) - if not hasattr(namespace, "nondefault"): - namespace.nondefault = set() - namespace.nondefault.add(self.dest) - - -class _StoreConstAction(_StoreAction): - """ - Same as `argparse._StoreConstAction` but uses the custom `_StoreAction`. - """ - - def __init__(self, option_strings, dest, const, default=None, required=False, help=None): - super().__init__( - option_strings=option_strings, - dest=dest, - nargs=0, - const=const, - default=default, - required=required, - help=help, - ) - - def __call__(self, parser, namespace, values, option_string=None): - super().__call__(parser, namespace, self.const, option_string) - - -class _StoreTrueAction(_StoreConstAction): - """ - Same as `argparse._StoreTrueAction` but uses the custom `_StoreConstAction`. - """ - - def __init__( - self, - option_strings, - dest, - default=None, - required=False, - help=None, - ): - super().__init__( - option_strings=option_strings, dest=dest, const=True, default=default, required=required, help=help - ) - - -class CustomArgumentGroup(argparse._ArgumentGroup): - """ - Custom argument group that allows for the use of `-` or `_` in arguments passed and overrides the help for each - when applicable. - """ - - def _add_action(self, action): - args = vars(action) - if isinstance(action, argparse._StoreTrueAction): - action = _StoreTrueAction( - args["option_strings"], args["dest"], args["default"], args["required"], args["help"] - ) - elif isinstance(action, argparse._StoreConstAction): - action = _StoreConstAction( - args["option_strings"], - args["dest"], - args["const"], - args["default"], - args["required"], - args["help"], - ) - elif isinstance(action, argparse._StoreAction): - action = _StoreAction(**args) - action = super()._add_action(action) - return action - - -class CustomArgumentParser(argparse.ArgumentParser): - """ - Custom argument parser that allows for the use of `-` or `_` in arguments passed and overrides the help for each - when applicable. - """ - - def add_argument(self, *args, **kwargs): - if "action" in kwargs: - # Translate action -> class - if kwargs["action"] == "store_true": - kwargs["action"] = _StoreTrueAction - else: - kwargs["action"] = _StoreAction - super().add_argument(*args, **kwargs) - - def add_argument_group(self, *args, **kwargs): - group = CustomArgumentGroup(self, *args, **kwargs) - self._action_groups.append(group) - return group diff --git a/venv/lib/python3.10/site-packages/accelerate/data_loader.py b/venv/lib/python3.10/site-packages/accelerate/data_loader.py deleted file mode 100644 index ef80eae224bebf35664ad288fbd92dcac04baccb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/data_loader.py +++ /dev/null @@ -1,1446 +0,0 @@ -# Copyright 2021 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 math -from contextlib import suppress -from typing import Callable, Optional, Union - -import torch -from packaging import version -from torch.utils.data import BatchSampler, DataLoader, IterableDataset, RandomSampler - -from .logging import get_logger -from .state import DistributedType, GradientState, PartialState, is_torch_xla_available -from .utils import ( - RNGType, - broadcast, - broadcast_object_list, - compare_versions, - concatenate, - find_batch_size, - get_data_structure, - initialize_tensors, - is_datasets_available, - is_torch_version, - is_torchdata_stateful_dataloader_available, - send_to_device, - slice_tensors, - synchronize_rng_states, -) - - -logger = get_logger(__name__) - -# kwargs of the DataLoader in min version 2.0 -_PYTORCH_DATALOADER_KWARGS = { - "batch_size": 1, - "shuffle": False, - "sampler": None, - "batch_sampler": None, - "num_workers": 0, - "collate_fn": None, - "pin_memory": False, - "drop_last": False, - "timeout": 0, - "worker_init_fn": None, - "multiprocessing_context": None, - "generator": None, - "prefetch_factor": 2, - "persistent_workers": False, - "pin_memory_device": "", -} - -# kwargs added after by version -_PYTORCH_DATALOADER_ADDITIONAL_KWARGS = {"2.6.0": {"in_order": True}} - -for v, additional_kwargs in _PYTORCH_DATALOADER_ADDITIONAL_KWARGS.items(): - if is_torch_version(">=", v): - _PYTORCH_DATALOADER_KWARGS.update(additional_kwargs) - - -class SeedableRandomSampler(RandomSampler): - """ - Same as a random sampler, except that in `__iter__` a seed can be used. - - Needed specifically in distributed cases, when the random generator for each GPU needs to start from the same seed - and be fully reproducable on multiple iterations. - - If a custom `generator` is passed, it will rely on its initial seed as well as the current iteration it is on - (stored in `self.epoch`). - """ - - def __init__(self, *args, **kwargs): - data_seed = kwargs.pop("data_seed", None) - super().__init__(*args, **kwargs) - - self.initial_seed = data_seed if data_seed is not None else torch.random.initial_seed() - self.epoch = 0 - - def __iter__(self): - if self.generator is None: - self.generator = torch.Generator( - device=torch.get_default_device() if hasattr(torch, "get_default_device") else "cpu" - ) - self.generator.manual_seed(self.initial_seed) - - # Allow `self.epoch` to modify the seed of the generator - seed = self.epoch + self.initial_seed - # print("Setting seed at epoch", self.epoch, seed) - self.generator.manual_seed(seed) - yield from super().__iter__() - self.set_epoch(self.epoch + 1) - - def set_epoch(self, epoch: int): - "Sets the current iteration of the sampler." - self.epoch = epoch - - -class BatchSamplerShard(BatchSampler): - """ - Wraps a PyTorch `BatchSampler` to generate batches for one of the processes only. Instances of this class will - always yield a number of batches that is a round multiple of `num_processes` and that all have the same size. - Depending on the value of the `drop_last` attribute of the batch sampler passed, it will either stop the iteration - at the first batch that would be too small / not present on all processes or loop with indices from the beginning. - - Args: - batch_sampler (`torch.utils.data.sampler.BatchSampler`): - The batch sampler to split in several shards. - num_processes (`int`, *optional*, defaults to 1): - The number of processes running concurrently. - process_index (`int`, *optional*, defaults to 0): - The index of the current process. - split_batches (`bool`, *optional*, defaults to `False`): - Whether the shards should be created by splitting a batch to give a piece of it on each process, or by - yielding different full batches on each process. - - On two processes with a sampler of `[[0, 1, 2, 3], [4, 5, 6, 7]]`, this will result in: - - - the sampler on process 0 to yield `[0, 1, 2, 3]` and the sampler on process 1 to yield `[4, 5, 6, 7]` if - this argument is set to `False`. - - the sampler on process 0 to yield `[0, 1]` then `[4, 5]` and the sampler on process 1 to yield `[2, 3]` - then `[6, 7]` if this argument is set to `True`. - even_batches (`bool`, *optional*, defaults to `True`): - Whether or not to loop back at the beginning of the sampler when the number of samples is not a round - multiple of (original batch size / number of processes). - - - - `BatchSampler`s with varying batch sizes are not enabled by default. To enable this behaviour, set `even_batches` - equal to `False` - - """ - - def __init__( - self, - batch_sampler: BatchSampler, - num_processes: int = 1, - process_index: int = 0, - split_batches: bool = False, - even_batches: bool = True, - ): - if split_batches and batch_sampler.batch_size % num_processes != 0: - raise ValueError( - f"To use `BatchSamplerShard` in `split_batches` mode, the batch size ({batch_sampler.batch_size}) " - f"needs to be a round multiple of the number of processes ({num_processes})." - ) - self.batch_sampler = batch_sampler - self.num_processes = num_processes - self.process_index = process_index - self.split_batches = split_batches - self.even_batches = even_batches - self.batch_size = getattr(batch_sampler, "batch_size", None) - self.drop_last = getattr(batch_sampler, "drop_last", False) - if self.batch_size is None and self.even_batches: - raise ValueError( - "You need to use `even_batches=False` when the batch sampler has no batch size. If you " - "are not calling this method directly, set `accelerator.even_batches=False` instead." - ) - - @property - def total_length(self): - return len(self.batch_sampler) - - def __len__(self): - if self.split_batches: - # Split batches does not change the length of the batch sampler - return len(self.batch_sampler) - if len(self.batch_sampler) % self.num_processes == 0: - # If the length is a round multiple of the number of processes, it's easy. - return len(self.batch_sampler) // self.num_processes - length = len(self.batch_sampler) // self.num_processes - if self.drop_last: - # Same if we drop the remainder. - return length - elif self.even_batches: - # When we even batches we always get +1 - return length + 1 - else: - # Otherwise it depends on the process index. - return length + 1 if self.process_index < len(self.batch_sampler) % self.num_processes else length - - def __iter__(self): - return self._iter_with_split() if self.split_batches else self._iter_with_no_split() - - def _iter_with_split(self): - initial_data = [] - batch_length = self.batch_sampler.batch_size // self.num_processes - for idx, batch in enumerate(self.batch_sampler): - if idx == 0: - initial_data = batch - if len(batch) == self.batch_size: - # If the batch is full, we yield the part of it this process is responsible of. - yield batch[batch_length * self.process_index : batch_length * (self.process_index + 1)] - - # If drop_last is True of the last batch was full, iteration is over, otherwise... - if not self.drop_last and len(initial_data) > 0 and len(batch) < self.batch_size: - if not self.even_batches: - if len(batch) > batch_length * self.process_index: - yield batch[batch_length * self.process_index : batch_length * (self.process_index + 1)] - else: - # For degenerate cases where the dataset has less than num_process * batch_size samples - while len(initial_data) < self.batch_size: - initial_data += initial_data - batch = batch + initial_data - yield batch[batch_length * self.process_index : batch_length * (self.process_index + 1)] - - def _iter_with_no_split(self): - initial_data = [] - batch_to_yield = [] - for idx, batch in enumerate(self.batch_sampler): - # We gather the initial indices in case we need to circle back at the end. - if not self.drop_last and idx < self.num_processes: - initial_data += batch - # We identify the batch to yield but wait until we ar sure every process gets a full batch before actually - # yielding it. - if idx % self.num_processes == self.process_index: - batch_to_yield = batch - if idx % self.num_processes == self.num_processes - 1 and ( - self.batch_size is None or len(batch) == self.batch_size - ): - yield batch_to_yield - batch_to_yield = [] - - # If drop_last is True, iteration is over, otherwise... - if not self.drop_last and len(initial_data) > 0: - if not self.even_batches: - if len(batch_to_yield) > 0: - yield batch_to_yield - else: - # ... we yield the complete batch we had saved before if it has the proper length - if len(batch_to_yield) == self.batch_size: - yield batch_to_yield - - # For degenerate cases where the dataset has less than num_process * batch_size samples - while len(initial_data) < self.num_processes * self.batch_size: - initial_data += initial_data - - # If the last batch seen was of the proper size, it has been yielded by its process so we move to the next - if len(batch) == self.batch_size: - batch = [] - idx += 1 - - # Make sure we yield a multiple of self.num_processes batches - cycle_index = 0 - while idx % self.num_processes != 0 or len(batch) > 0: - end_index = cycle_index + self.batch_size - len(batch) - batch += initial_data[cycle_index:end_index] - if idx % self.num_processes == self.process_index: - yield batch - cycle_index = end_index - batch = [] - idx += 1 - - -class IterableDatasetShard(IterableDataset): - """ - Wraps a PyTorch `IterableDataset` to generate samples for one of the processes only. Instances of this class will - always yield a number of samples that is a round multiple of the actual batch size (depending of the value of - `split_batches`, this is either `batch_size` or `batch_size x num_processes`). Depending on the value of the - `drop_last` attribute of the batch sampler passed, it will either stop the iteration at the first batch that would - be too small or loop with indices from the beginning. - - Args: - dataset (`torch.utils.data.dataset.IterableDataset`): - The batch sampler to split in several shards. - batch_size (`int`, *optional*, defaults to 1): - The size of the batches per shard (if `split_batches=False`) or the size of the batches (if - `split_batches=True`). - drop_last (`bool`, *optional*, defaults to `False`): - Whether or not to drop the last incomplete batch or complete the last batches by using the samples from the - beginning. - num_processes (`int`, *optional*, defaults to 1): - The number of processes running concurrently. - process_index (`int`, *optional*, defaults to 0): - The index of the current process. - split_batches (`bool`, *optional*, defaults to `False`): - Whether the shards should be created by splitting a batch to give a piece of it on each process, or by - yielding different full batches on each process. - - On two processes with an iterable dataset yielding of `[0, 1, 2, 3, 4, 5, 6, 7]`, this will result in: - - - the shard on process 0 to yield `[0, 1, 2, 3]` and the shard on process 1 to yield `[4, 5, 6, 7]` if this - argument is set to `False`. - - the shard on process 0 to yield `[0, 1, 4, 5]` and the sampler on process 1 to yield `[2, 3, 6, 7]` if - this argument is set to `True`. - """ - - def __init__( - self, - dataset: IterableDataset, - batch_size: int = 1, - drop_last: bool = False, - num_processes: int = 1, - process_index: int = 0, - split_batches: bool = False, - ): - if split_batches and batch_size > 1 and batch_size % num_processes != 0: - raise ValueError( - f"To use `IterableDatasetShard` in `split_batches` mode, the batch size ({batch_size}) " - f"needs to be a round multiple of the number of processes ({num_processes})." - ) - self.dataset = dataset - self.batch_size = batch_size - self.drop_last = drop_last - self.num_processes = num_processes - self.process_index = process_index - self.split_batches = split_batches - - def set_epoch(self, epoch): - self.epoch = epoch - if hasattr(self.dataset, "set_epoch"): - self.dataset.set_epoch(epoch) - - def __len__(self): - # We will just raise the downstream error if the underlying dataset is not sized - if self.drop_last: - return (len(self.dataset) // (self.batch_size * self.num_processes)) * self.batch_size - else: - return math.ceil(len(self.dataset) / (self.batch_size * self.num_processes)) * self.batch_size - - def __iter__(self): - if ( - not hasattr(self.dataset, "set_epoch") - and hasattr(self.dataset, "generator") - and isinstance(self.dataset.generator, torch.Generator) - ): - self.dataset.generator.manual_seed(self.epoch) - real_batch_size = self.batch_size if self.split_batches else (self.batch_size * self.num_processes) - process_batch_size = (self.batch_size // self.num_processes) if self.split_batches else self.batch_size - process_slice = range(self.process_index * process_batch_size, (self.process_index + 1) * process_batch_size) - - first_batch = None - current_batch = [] - for element in self.dataset: - current_batch.append(element) - # Wait to have a full batch before yielding elements. - if len(current_batch) == real_batch_size: - for i in process_slice: - yield current_batch[i] - if first_batch is None: - first_batch = current_batch.copy() - current_batch = [] - - # Finished if drop_last is True, otherwise complete the last batch with elements from the beginning. - if not self.drop_last and len(current_batch) > 0: - if first_batch is None: - first_batch = current_batch.copy() - while len(current_batch) < real_batch_size: - current_batch += first_batch - for i in process_slice: - yield current_batch[i] - - -class DataLoaderStateMixin: - """ - Mixin class that adds a state to a `DataLoader` to keep track of the status inside the dataloader such as at the - end of the iteration, the number of items in the dataset in the last batch relative to the batch size, and other - useful information that might be needed. - - **Available attributes:** - - - **end_of_dataloader** (`bool`) -- Whether at the last iteration or batch - - **remainder** (`int`) -- The number of items that are remaining in the last batch, relative to the total - batch size - - - - Inheriters of this class should ensure that the class creates a `GradientState()` instance, stored in - `self.gradient_state`. - - - - """ - - def __init_subclass__(cls, **kwargs): - cls.end_of_dataloader = False - cls.remainder = -1 - - def reset(self): - self.end_of_dataloader = False - self.remainder = -1 - - def begin(self): - "Prepares the gradient state for the current dataloader" - self.reset() - with suppress(Exception): - if not self._drop_last: - length = getattr(self.dataset, "total_dataset_length", len(self.dataset)) - self.remainder = length % self.total_batch_size - self.gradient_state._add_dataloader(self) - - def end(self): - "Cleans up the gradient state after exiting the dataloader" - self.gradient_state._remove_dataloader(self) - - -class DataLoaderAdapter: - """ - A class which wraps around a PyTorch `DataLoader` (or variants of it) to be used with the `Accelerator`. For - compatability reasons, this class inherits from the class it wraps around, so it can be used as a drop-in. - """ - - def __init__(self, dataset, use_stateful_dataloader=False, batch_sampler=None, **kwargs): - self.use_stateful_dataloader = use_stateful_dataloader - if is_torchdata_stateful_dataloader_available(): - from torchdata.stateful_dataloader import StatefulDataLoader - - if use_stateful_dataloader and not is_torchdata_stateful_dataloader_available(): - raise ImportError( - "StatefulDataLoader is not available. Please install torchdata version 0.8.0 or higher to use it." - ) - if use_stateful_dataloader: - torchdata_version = version.parse(importlib.metadata.version("torchdata")) - if ( - "in_order" in kwargs - and compare_versions(torchdata_version, "<", "0.11") - and is_torch_version(">=", "2.6.0") - ): - kwargs.pop("in_order") - self.base_dataloader = StatefulDataLoader(dataset, batch_sampler=batch_sampler, **kwargs) - else: - self.base_dataloader = DataLoader(dataset, batch_sampler=batch_sampler, **kwargs) - - if hasattr(self.base_dataloader, "state_dict"): - self.dl_state_dict = self.base_dataloader.state_dict() - - def __getattr__(self, name): - # Avoid infinite recursion if we try to access a nonexistent base_dataloader attribute. - if name == "base_dataloader": - raise AttributeError() - # Delegate attribute access to the internal dataloader - return getattr(self.base_dataloader, name) - - def state_dict(self): - return self.dl_state_dict - - def load_state_dict(self, state_dict): - self.base_dataloader.load_state_dict(state_dict) - - @property - def __class__(self): - """ - In order to maintain backwards compatability with other code, we need to ensure `isinstance(obj, DataLoader)` - returs true. This is because some downstream code assumes that the `DataLoader` is the base class of the - object. - """ - return self.base_dataloader.__class__ - - def __len__(self): - return len(self.base_dataloader) - - def adjust_state_dict_for_prefetch(self): - """ - Adjusts the state dict for prefetching. Natively, this will adjust all of the iters yielded keys in - `self.dl_state_dict` by a factor of `num_processes - 1`, however if a custom correction is needed, this can be - overridden. - - This should modify `self.dl_state_dict` directly - """ - # The state dict will be off by a factor of `n-1` batch too many during DDP, - # so we need to adjust it here - if PartialState().distributed_type != DistributedType.NO: - factor = PartialState().num_processes - 1 - if self.dl_state_dict["_sampler_iter_yielded"] > 0: - self.dl_state_dict["_sampler_iter_yielded"] -= factor - if self.dl_state_dict["_num_yielded"] > 0: - self.dl_state_dict["_num_yielded"] -= factor - if self.dl_state_dict["_index_sampler_state"] is not None: - if ( - "samples_yielded" in self.dl_state_dict["_index_sampler_state"] - and self.dl_state_dict["_index_sampler_state"]["samples_yielded"] > 0 - ): - self.dl_state_dict["_index_sampler_state"]["samples_yielded"] -= self.batch_size * factor - - def _update_state_dict(self): - # The state_dict of the underlying base_dataloader may be ahead of what is currently being yielded. - # E.g. the implementation of DataLoaderShard involves having an underlying iterator 1 element ahead of - # what it wants to yield. - # - # _update_state_dict is called to snapshot the state_dict that would properly recover the DataLoaderAdapter. - if hasattr(self.base_dataloader, "state_dict"): - self.dl_state_dict = self.base_dataloader.state_dict() - # Potentially modify the state_dict to adjust for prefetching - self.adjust_state_dict_for_prefetch() - # Then tag if we are at the end of the dataloader - self.dl_state_dict["_iterator_finished"] = self.end_of_dataloader - - -class DataLoaderShard(DataLoaderAdapter, DataLoaderStateMixin): - """ - Subclass of `DataLoaderAdapter` that will deal with device placement and current distributed setup. - - Args: - dataset (`torch.utils.data.dataset.Dataset`): - The dataset to use to build this dataloader. - device (`torch.device`, *optional*): - If passed, the device to put all batches on. - rng_types (list of `str` or [`~utils.RNGType`]): - The list of random number generators to synchronize at the beginning of each iteration. Should be one or - several of: - - - `"torch"`: the base torch random number generator - - `"cuda"`: the CUDA random number generator (GPU only) - - `"xla"`: the XLA random number generator (TPU only) - - `"generator"`: an optional `torch.Generator` - synchronized_generator (`torch.Generator`, *optional*): - A random number generator to keep synchronized across processes. - skip_batches (`int`, *optional*, defaults to 0): - The number of batches to skip at the beginning. - use_stateful_dataloader (`bool`, *optional*, defaults to `False`): - Whether to have this class adapt `StatefulDataLoader` from `torchdata` instead of the regular `DataLoader`. - **kwargs (additional keyword arguments, *optional*): - All other keyword arguments to pass to the regular `DataLoader` initialization. - - **Available attributes:** - - - **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes. - Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total - number of processes - - - **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes. - """ - - def __init__( - self, - dataset, - device=None, - rng_types=None, - synchronized_generator=None, - skip_batches=0, - use_stateful_dataloader=False, - _drop_last: bool = False, - _non_blocking: bool = False, - torch_device_mesh=None, - **kwargs, - ): - super().__init__(dataset, use_stateful_dataloader=use_stateful_dataloader, **kwargs) - self.device = device - self.rng_types = rng_types - self.synchronized_generator = synchronized_generator - self.skip_batches = skip_batches - self.gradient_state = GradientState() - self._drop_last = _drop_last - self._non_blocking = _non_blocking - self.iteration = 0 - - def __iter__(self): - if self.rng_types is not None: - synchronize_rng_states(self.rng_types, self.synchronized_generator) - self.begin() - - self.set_epoch(self.iteration) - dataloader_iter = self.base_dataloader.__iter__() - # We iterate one batch ahead to check when we are at the end - try: - current_batch = next(dataloader_iter) - except StopIteration: - self.end() - return - - batch_index = 0 - while True: - try: - # But we still move it to the device so it is done before `StopIteration` is reached - if self.device is not None: - current_batch = send_to_device(current_batch, self.device, non_blocking=self._non_blocking) - self._update_state_dict() - next_batch = next(dataloader_iter) - if batch_index >= self.skip_batches: - yield current_batch - batch_index += 1 - current_batch = next_batch - except StopIteration: - self.end_of_dataloader = True - self._update_state_dict() - if batch_index >= self.skip_batches: - yield current_batch - break - - self.iteration += 1 - self.end() - - def __reduce__(self): - """ - Define the `__reduce__` method to ensure a `DataLoaderShard` can be pickled and unpickled. This needs to be - explicitly defined since default pickling behavior is broken by `DataLoaderAdapter` messing with its - `__class__` member. - """ - args = super().__reduce__() - return (DataLoaderShard, *args[1:]) - - def set_epoch(self, epoch: int): - # In case it is manually passed in, the user can set it to what they like - if self.iteration != epoch: - self.iteration = epoch - if hasattr(self.batch_sampler, "set_epoch"): - self.batch_sampler.set_epoch(epoch) - if hasattr(self.batch_sampler, "sampler") and hasattr(self.batch_sampler.sampler, "set_epoch"): - self.batch_sampler.sampler.set_epoch(epoch) - if ( - hasattr(self.batch_sampler, "batch_sampler") - and hasattr(self.batch_sampler.batch_sampler, "sampler") - and hasattr(self.batch_sampler.batch_sampler.sampler, "set_epoch") - ): - self.batch_sampler.batch_sampler.sampler.set_epoch(epoch) - # We support if a custom `Dataset` implementation has `set_epoch` - # or in general HF datasets `Datasets` - elif hasattr(self.dataset, "set_epoch"): - self.dataset.set_epoch(epoch) - - @property - def total_batch_size(self): - batch_sampler = self.sampler if isinstance(self.sampler, BatchSampler) else self.batch_sampler - return ( - batch_sampler.batch_size - if getattr(batch_sampler, "split_batches", False) - else (batch_sampler.batch_size * getattr(batch_sampler, "num_processes", 1)) - ) - - @property - def total_dataset_length(self): - if hasattr(self.dataset, "total_length"): - return self.dataset.total_length - else: - return len(self.dataset) - - def get_sampler(self): - return get_sampler(self) - - def set_sampler(self, sampler): - sampler_is_batch_sampler = isinstance(self.sampler, BatchSampler) - if sampler_is_batch_sampler: - self.sampler.sampler = sampler - else: - self.batch_sampler.sampler = sampler - if hasattr(self.batch_sampler, "batch_sampler"): - self.batch_sampler.batch_sampler.sampler = sampler - - -if is_torch_xla_available(): - import torch_xla.distributed.parallel_loader as xpl - - class MpDeviceLoaderWrapper(xpl.MpDeviceLoader): - """ - Wrapper for the xpl.MpDeviceLoader class that knows the total batch size. - - XLA preloading threads will all call DataLoaderShard's __iter__(). Remove rng_types from DataLoaderShard to - prevent it from using the XLA device in the preloading threads, and synchronize the RNG once from the main - thread only. - - **Available attributes:** - - - **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes. - Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total - number of processes - - - **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes. - """ - - def __init__(self, dataloader: DataLoaderShard, device: torch.device): - super().__init__(dataloader, device) - self._rng_types = self._loader.rng_types - self._loader.rng_types = None - self.device = device - - def __iter__(self): - if self._rng_types is not None: - synchronize_rng_states(self._rng_types, self._loader.synchronized_generator) - - return super().__iter__() - - def set_epoch(self, epoch: int): - if hasattr(self.dataloader, "set_epoch"): - self.dataloader.set_epoch(epoch) - - @property - def total_batch_size(self): - return self._loader.total_batch_size - - @property - def total_dataset_length(self): - return self._loader.total_dataset_length - - @property - def batch_sampler(self): - return self._loader.batch_sampler - - @property - def dataloader(self): - return self._loader - - -class DataLoaderDispatcher(DataLoaderAdapter, DataLoaderStateMixin): - """ - Subclass of `DataLoaderAdapter` that will iterate and preprocess on process 0 only, then dispatch on each process - their part of the batch. - - Args: - split_batches (`bool`, *optional*, defaults to `False`): - Whether the resulting `DataLoader` should split the batches of the original data loader across devices or - yield full batches (in which case it will yield batches starting at the `process_index`-th and advancing of - `num_processes` batches at each iteration). Another way to see this is that the observed batch size will be - the same as the initial `dataloader` if this option is set to `True`, the batch size of the initial - `dataloader` multiplied by `num_processes` otherwise. Setting this option to `True` requires that the batch - size of the `dataloader` is a round multiple of `batch_size`. - skip_batches (`int`, *optional*, defaults to 0): - The number of batches to skip at the beginning of an iteration. - use_stateful_dataloader (`bool`, *optional*, defaults to `False`): - Whether to have this class adapt `StatefulDataLoader` from `torchdata` instead of the regular `DataLoader`. - - **Available attributes:** - - - **total_batch_size** (`int`) -- Total batch size of the dataloader across all processes. - Equal to the original batch size when `split_batches=True`; otherwise the original batch size * the total - number of processes - - - **total_dataset_length** (`int`) -- Total length of the inner dataset across all processes. - """ - - def __init__( - self, - dataset, - split_batches: bool = False, - skip_batches=0, - use_stateful_dataloader=False, - _drop_last: bool = False, - _non_blocking: bool = False, - slice_fn=None, - torch_device_mesh=None, - **kwargs, - ): - shuffle = False - from torch.utils.data.datapipes.iter.combinatorics import ShufflerIterDataPipe - - # We need to save the shuffling state of the DataPipe - if isinstance(dataset, ShufflerIterDataPipe): - shuffle = dataset._shuffle_enabled - super().__init__(dataset, use_stateful_dataloader=use_stateful_dataloader, **kwargs) - self.split_batches = split_batches - if shuffle: - torch.utils.data.graph_settings.apply_shuffle_settings(dataset, shuffle=shuffle) - - self.gradient_state = GradientState() - self.state = PartialState() - self._drop_last = _drop_last - self._non_blocking = _non_blocking - self.skip_batches = skip_batches - self.torch_device_mesh = torch_device_mesh - - self.slice_fn = slice_tensors if slice_fn is None else slice_fn - self.iteration = 0 - - # if a device mesh is provided extract each dimension (dp, fsdp, tp) - # device mesh may hold any number of dimensions, however, - # below code is for targetted support for dp, fsdp and tp - - # device mesh will be used only if there is tp involved - # or any multi-dimensional parallelism involving tp - # (dp, tp) (fsdp, tp) (dp, fsdp, tp) - # otherwise the default behavour not using device mesh should be sufficient - # since multi dimensional parallelism devoid of tp would anyway need - # different batches for each process irrespective of dp or fsdp - self.submesh_tp = None - self.submesh_dp = None - self.submesh_fsdp = None - if self.torch_device_mesh and "tp" in self.torch_device_mesh.mesh_dim_names: - self.submesh_tp = self.torch_device_mesh["tp"] - if "dp" in self.torch_device_mesh.mesh_dim_names: - self.submesh_dp = self.torch_device_mesh["dp"] - if "fsdp" in self.torch_device_mesh.mesh_dim_names: - self.submesh_fsdp = self.torch_device_mesh["fsdp"] - if self.submesh_tp and (self.submesh_dp or self.submesh_fsdp): - raise ValueError("TP + (DP/FSDP) is not yet supported in dispatch mode") - - def _fetch_batches(self, iterator): - batches, batch = None, None - # On process 0, we gather the batch to dispatch. - if self.state.process_index == 0: - # Procedure to support TP only is simpler - # since we want to dispatch the same batch of samples across all ranks - # this removes complexity of handling multiple tp rank groups when TP + DP - # combination is involved. - - try: - # for TP case avoid using split_batches - # since it would mean that the dataloader should be spilling out - # duplicates of batches. - if self.split_batches: - # One batch of the main iterator is dispatched and split. - if self.submesh_tp: - logger.warning( - "Use of split_batches for TP would need the dataloader to produce duplicate batches," - "otherwise, use dispatch_batches=True instead." - ) - self._update_state_dict() - batch = next(iterator) - else: - # num_processes batches of the main iterator are concatenated then dispatched and split. - # We add the batches one by one so we have the remainder available when drop_last=False. - batches = [] - if self.submesh_tp: - # when tp, extract single batch and then replicate - self._update_state_dict() - batch = next(iterator) - batches = [batch] * self.state.num_processes - else: - for _ in range(self.state.num_processes): - self._update_state_dict() - batches.append(next(iterator)) - try: - batch = concatenate(batches, dim=0) - except RuntimeError as e: - raise RuntimeError( - "You can't use batches of different size with `dispatch_batches=True` or when using an `IterableDataset`." - "either pass `dispatch_batches=False` and have each process fetch its own batch " - " or pass `split_batches=True`. By doing so, the main process will fetch a full batch and " - "slice it into `num_processes` batches for each process." - ) from e - # In both cases, we need to get the structure of the batch that we will broadcast on other - # processes to initialize the tensors with the right shape. - # data_structure, stop_iteration - batch_info = [get_data_structure(batch), False] - except StopIteration: - batch_info = [None, True] - else: - batch_info = [None, self._stop_iteration] - # This is inplace, so after this instruction, every process has the same `batch_info` as process 0. - broadcast_object_list(batch_info) - self._stop_iteration = batch_info[1] - if self._stop_iteration: - # If drop_last is False and split_batches is False, we may have a remainder to take care of. - if not self.split_batches and not self._drop_last: - if self.state.process_index == 0 and len(batches) > 0: - batch = concatenate(batches, dim=0) - batch_info = [get_data_structure(batch), False] - else: - batch_info = [None, True] - broadcast_object_list(batch_info) - return batch, batch_info - - def __iter__(self): - self.begin() - self.set_epoch(self.iteration) - main_iterator = None - if is_torch_version(">=", "2.0.1"): - # NOTE PyTorch DataLoader adds forward compatibilities for DataPipes, which broadcasts - # shared seed to all dist processes. Thus, we need to create iterator for all dist processes. - # But, we only iterate through the DataLoader on process 0. - main_iterator = self.base_dataloader.__iter__() - elif self.state.process_index == 0: - main_iterator = self.base_dataloader.__iter__() - stop_iteration = False - self._stop_iteration = False - first_batch = None - next_batch, next_batch_info = self._fetch_batches(main_iterator) - batch_index = 0 - while not stop_iteration: - batch, batch_info = next_batch, next_batch_info - - if self.state.process_index != 0: - # Initialize tensors on other processes than process 0. - batch = initialize_tensors(batch_info[0]) - batch = send_to_device(batch, self.state.device, non_blocking=self._non_blocking) - # Broadcast the batch before splitting it. - batch = broadcast(batch, from_process=0) - - if not self._drop_last and first_batch is None: - # We keep at least num processes elements of the first batch to be able to complete the last batch - first_batch = self.slice_fn( - batch, - slice(0, self.state.num_processes), - process_index=self.state.process_index, - num_processes=self.state.num_processes, - ) - - if batch is None: - raise ValueError( - f"Batch does not contain any data (`{batch}`). At the end of all iterable data available before expected stop iteration." - ) - - observed_batch_size = find_batch_size(batch) - batch_size = observed_batch_size // self.state.num_processes - - stop_iteration = self._stop_iteration - if not stop_iteration: - # We may still be at the end of the dataloader without knowing it yet: if there is nothing left in - # the dataloader since the number of batches is a round multiple of the number of processes. - next_batch, next_batch_info = self._fetch_batches(main_iterator) - # next_batch_info[0] is None when there are no more batches, otherwise we still need to process them. - if self._stop_iteration and next_batch_info[0] is None: - stop_iteration = True - - if not self._drop_last and stop_iteration and observed_batch_size % self.state.num_processes != 0: - # If the last batch is not complete, let's add the first batch to it. - batch = concatenate([batch, first_batch], dim=0) - # Batch size computation above is wrong, it's off by 1 so we fix it. - batch_size += 1 - - data_slice = slice(self.state.process_index * batch_size, (self.state.process_index + 1) * batch_size) - batch = self.slice_fn( - batch, - data_slice, - process_index=self.state.process_index, - num_processes=self.state.num_processes, - ) - - if stop_iteration: - self.end_of_dataloader = True - self._update_state_dict() - self.remainder = observed_batch_size - if batch_index >= self.skip_batches: - yield batch - batch_index += 1 - self.iteration += 1 - self.end() - - def set_epoch(self, epoch: int): - # In case it is manually passed in, the user can set it to what they like - if self.iteration != epoch: - self.iteration = epoch - if hasattr(self.batch_sampler, "sampler") and hasattr(self.batch_sampler.sampler, "set_epoch"): - self.batch_sampler.sampler.set_epoch(epoch) - elif hasattr(self.dataset, "set_epoch"): - self.dataset.set_epoch(epoch) - - def __len__(self): - whole_length = len(self.base_dataloader) - if self.split_batches: - return whole_length - elif self._drop_last: - return whole_length // self.state.num_processes - else: - return math.ceil(whole_length / self.state.num_processes) - - def __reduce__(self): - """ - Define the `__reduce__` method to ensure a `DataLoaderDispatcher` can be pickled and unpickled. This needs to - be explicitly defined since default pickling behavior is broken by `DataLoaderAdapter` messing with its - `__class__` member. - """ - args = super().__reduce__() - return (DataLoaderDispatcher, *args[1:]) - - @property - def total_batch_size(self): - return ( - self.dataset.batch_size if self.split_batches else (self.dataset.batch_size * self.dataset.num_processes) - ) - - @property - def total_dataset_length(self): - return len(self.dataset) - - def get_sampler(self): - return get_sampler(self) - - def set_sampler(self, sampler): - sampler_is_batch_sampler = isinstance(self.sampler, BatchSampler) - if sampler_is_batch_sampler: - self.sampler.sampler = sampler - else: - self.batch_sampler.sampler = sampler - if hasattr(self.batch_sampler, "batch_sampler"): - self.batch_sampler.batch_sampler.sampler = sampler - - -def get_sampler(dataloader): - """ - Get the sampler associated to the dataloader - - Args: - dataloader (`torch.utils.data.dataloader.DataLoader`): - The data loader to split across several devices. - Returns: - `torch.utils.data.Sampler`: The sampler associated to the dataloader - """ - sampler_is_batch_sampler = isinstance(dataloader.sampler, BatchSampler) - if sampler_is_batch_sampler: - sampler = getattr(dataloader.sampler, "sampler", None) - else: - sampler = getattr(dataloader.batch_sampler, "sampler", None) - return sampler - - -def prepare_data_loader( - dataloader: DataLoader, - device: Optional[torch.device] = None, - num_processes: Optional[int] = None, - process_index: Optional[int] = None, - split_batches: bool = False, - put_on_device: bool = False, - rng_types: Optional[list[Union[str, RNGType]]] = None, - dispatch_batches: Optional[bool] = None, - even_batches: bool = True, - slice_fn_for_dispatch: Optional[Callable] = None, - use_seedable_sampler: bool = False, - data_seed: Optional[int] = None, - non_blocking: bool = False, - use_stateful_dataloader: bool = False, - torch_device_mesh=None, -) -> DataLoader: - """ - Wraps a PyTorch `DataLoader` to generate batches for one of the processes only. - - Depending on the value of the `drop_last` attribute of the `dataloader` passed, it will either stop the iteration - at the first batch that would be too small / not present on all processes or loop with indices from the beginning. - - Args: - dataloader (`torch.utils.data.dataloader.DataLoader`): - The data loader to split across several devices. - device (`torch.device`): - The target device for the returned `DataLoader`. - num_processes (`int`, *optional*): - The number of processes running concurrently. Will default to the value given by [`~state.PartialState`]. - process_index (`int`, *optional*): - The index of the current process. Will default to the value given by [`~state.PartialState`]. - split_batches (`bool`, *optional*, defaults to `False`): - Whether the resulting `DataLoader` should split the batches of the original data loader across devices or - yield full batches (in which case it will yield batches starting at the `process_index`-th and advancing of - `num_processes` batches at each iteration). - - Another way to see this is that the observed batch size will be the same as the initial `dataloader` if - this option is set to `True`, the batch size of the initial `dataloader` multiplied by `num_processes` - otherwise. - - Setting this option to `True` requires that the batch size of the `dataloader` is a round multiple of - `batch_size`. - put_on_device (`bool`, *optional*, defaults to `False`): - Whether or not to put the batches on `device` (only works if the batches are nested list, tuples or - dictionaries of tensors). - rng_types (list of `str` or [`~utils.RNGType`]): - The list of random number generators to synchronize at the beginning of each iteration. Should be one or - several of: - - - `"torch"`: the base torch random number generator - - `"cuda"`: the CUDA random number generator (GPU only) - - `"xla"`: the XLA random number generator (TPU only) - - `"generator"`: the `torch.Generator` of the sampler (or batch sampler if there is no sampler in your - dataloader) or of the iterable dataset (if it exists) if the underlying dataset is of that type. - - dispatch_batches (`bool`, *optional*): - If set to `True`, the dataloader prepared is only iterated through on the main process and then the batches - are split and broadcast to each process. Will default to `True` when the underlying dataset is an - `IterableDataset`, `False` otherwise. - even_batches (`bool`, *optional*, defaults to `True`): - If set to `True`, in cases where the total batch size across all processes does not exactly divide the - dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among - all workers. - slice_fn_for_dispatch (`Callable`, *optional*`): - If passed, this function will be used to slice tensors across `num_processes`. Will default to - [`~utils.slice_tensors`]. This argument is used only when `dispatch_batches` is set to `True` and will be - ignored otherwise. - use_seedable_sampler (`bool`, *optional*, defaults to `False`): - Whether to use the [`~data_loader.SeedableRandomSampler`] instead of a `RandomSampler` for better - reproducability. Comes at a cost of potentially different performances due to different shuffling - algorithms but ensures results will be the *exact* same. Should be paired with `set_seed()` at every - `self.set_epoch` - data_seed (`int`, *optional*, defaults to `None`): - The seed to use for the underlying generator when using `use_seedable_sampler`. If `None`, the generator - will use the current default seed from torch. - non_blocking (`bool`, *optional*, defaults to `False`): - If set to `True`, dataloader will utilize non-blocking host-to-device transfers. If the dataloader has - `pin_memory` set to `True`, this will help to increase overlap between data transfer and computations. - use_stateful_dataloader (`bool`, *optional*, defaults to `False`): - "If set to true, the dataloader prepared by the Accelerator will be backed by " - "[torchdata.StatefulDataLoader](https://github.com/pytorch/data/tree/main/torchdata/stateful_dataloader). - This requires `torchdata` version 0.8.0 or higher that supports StatefulDataLoader to be installed." - torch_device_mesh (`torch.distributed.DeviceMesh`, *optional*, defaults to `None`): - PyTorch device mesh. - - - Returns: - `torch.utils.data.dataloader.DataLoader`: A new data loader that will yield the portion of the batches - - - - `BatchSampler`s with varying batch sizes are not enabled by default. To enable this behaviour, set `even_batches` - equal to `False` - - - """ - if dispatch_batches is None: - if not put_on_device: - dispatch_batches = False - else: - dispatch_batches = isinstance(dataloader.dataset, IterableDataset) - - if dispatch_batches and not put_on_device: - raise ValueError("Using `dispatch_batches=True` requires `put_on_device=True`.") - # Grab defaults from PartialState - state = PartialState() - if num_processes is None: - num_processes = state.num_processes - - if process_index is None: - process_index = state.process_index - - if torch_device_mesh: - if state.distributed_type == DistributedType.DEEPSPEED: - # In DeepSpeed, the optimizer sharing level in DP is determined by the config file. - # Only considers "dp" and "tp". - # Given a device mesh (dp, tp) = (2, 3): - # - From the data parallel perspective, ranks should be structured as: 0 0 0 1 1 1 - # - Processes with the same DP rank will receive the same batch. - if "tp" in torch_device_mesh.mesh_dim_names: - submesh_tp_size = torch_device_mesh["tp"].size() - process_index = process_index // submesh_tp_size - num_processes = num_processes // submesh_tp_size - else: - # when device mesh is used, specifically with TP - # then there is need to update process_index and num_processes - # to bring in the effect of generating same batch across TP ranks - # and different batch across FSDP and DP ranks. - # Example: - # if device mesh is (dp,fsdp,tp) = (2, 2, 3) - # ranks would range from 0...11 - # from data angle ranks should look like 0 0 0 1 1 1 2 2 2 3 3 3 - # processes with same ranks/ids would receive the same batch - submesh_fsdp_size = 1 - submesh_dp_size = 1 - submesh_tp_size = 1 - if "tp" in torch_device_mesh.mesh_dim_names: - submesh_tp_size = torch_device_mesh["tp"].size() - if "dp" in torch_device_mesh.mesh_dim_names: - submesh_dp_size = torch_device_mesh["dp"].size() - if "fsdp" in torch_device_mesh.mesh_dim_names: - submesh_fsdp_size = torch_device_mesh["fsdp"].size() - process_index = process_index // submesh_tp_size - num_processes = submesh_fsdp_size * submesh_dp_size - - # Sanity check - if split_batches: - if dataloader.batch_size is not None: - batch_size_for_check = dataloader.batch_size - else: - # For custom batch_sampler - if hasattr(dataloader.batch_sampler, "batch_size"): - batch_size_for_check = dataloader.batch_sampler.batch_size - else: - raise ValueError( - "In order to use `split_batches==True` you must have a `batch_size` attribute either in the passed " - "`dataloader` or `dataloader.batch_sampler` objects, and it has to return a natural number. " - "Your `dataloader.batch_size` is None and `dataloader.batch_sampler` " - f"(`{type(dataloader.batch_sampler)}`) does not have the `batch_size` attribute set." - ) - - if batch_size_for_check > 1 and batch_size_for_check % num_processes != 0: - raise ValueError( - f"To use a `DataLoader` in `split_batches` mode, the batch size ({dataloader.batch_size}) " - f"needs to be a round multiple of the number of processes ({num_processes})." - ) - - new_dataset = dataloader.dataset - # Iterable dataset doesn't like batch_sampler, but data_loader creates a default one for it - new_batch_sampler = dataloader.batch_sampler if not isinstance(new_dataset, IterableDataset) else None - sampler_is_batch_sampler = isinstance(dataloader.sampler, BatchSampler) - synchronized_generator = None - - sampler = get_sampler(dataloader) - if isinstance(sampler, RandomSampler) and use_seedable_sampler: - # When iterating through the dataloader during distributed processes - # we want to ensure that on each process we are iterating through the same - # samples in the same order if a seed is set. This requires a tweak - # to the `torch.utils.data.RandomSampler` class (if used). - sampler = SeedableRandomSampler( - data_source=sampler.data_source, - replacement=sampler.replacement, - num_samples=sampler._num_samples, - generator=getattr( - sampler, - "generator", - torch.Generator(device=torch.get_default_device() if hasattr(torch, "get_default_device") else "cpu"), - ), - data_seed=data_seed, - ) - - if isinstance(dataloader.sampler, RandomSampler) and state.distributed_type == DistributedType.XLA: - # isinstance(dataloader.sampler, RandomSampler) indicates the original dataloader has `shuffle` enabled. - generator = torch.Generator( - device=torch.get_default_device() if hasattr(torch, "get_default_device") else "cpu" - ) - seed = int(torch.empty((), dtype=torch.int64).random_().item()) - generator.manual_seed(seed) - dataloader.generator = generator - dataloader.sampler.generator = generator - # No change if no multiprocess - if (num_processes != 1 or state.distributed_type == DistributedType.MEGATRON_LM) and not dispatch_batches: - if is_datasets_available(): - from datasets import IterableDataset as DatasetsIterableDataset - if ( - is_datasets_available() - and isinstance(new_dataset, DatasetsIterableDataset) - and not split_batches - and new_dataset.n_shards > num_processes - ): - new_dataset = new_dataset.shard(num_shards=num_processes, index=process_index) - elif isinstance(new_dataset, IterableDataset): - if getattr(dataloader.dataset, "generator", None) is not None: - synchronized_generator = dataloader.dataset.generator - new_dataset = IterableDatasetShard( - new_dataset, - batch_size=dataloader.batch_size, - drop_last=dataloader.drop_last, - num_processes=num_processes, - process_index=process_index, - split_batches=split_batches, - ) - else: - if not use_seedable_sampler and hasattr(sampler, "generator"): - if sampler.generator is None: - sampler.generator = torch.Generator( - device=torch.get_default_device() if hasattr(torch, "get_default_device") else "cpu" - ) - seed = int(torch.empty((), dtype=torch.int64).random_().item()) - sampler.generator.manual_seed(seed) - synchronized_generator = sampler.generator - batch_sampler = dataloader.sampler if sampler_is_batch_sampler else dataloader.batch_sampler - new_batch_sampler = BatchSamplerShard( - batch_sampler, - num_processes=num_processes, - process_index=process_index, - split_batches=split_batches, - even_batches=even_batches, - ) - - # We ignore all of those since they are all dealt with by our new_batch_sampler - ignore_kwargs = [ - "batch_size", - "shuffle", - "sampler", - "batch_sampler", - "drop_last", - ] - - if rng_types is not None and synchronized_generator is None and "generator" in rng_types: - rng_types.remove("generator") - - kwargs = { - k: getattr(dataloader, k, _PYTORCH_DATALOADER_KWARGS[k]) - for k in _PYTORCH_DATALOADER_KWARGS - if k not in ignore_kwargs - } - - # Need to provide batch_size as batch_sampler is None for Iterable dataset - if new_batch_sampler is None: - kwargs["drop_last"] = dataloader.drop_last - kwargs["batch_size"] = ( - dataloader.batch_size // num_processes if split_batches and not dispatch_batches else dataloader.batch_size - ) - if dispatch_batches: - kwargs.pop("generator") - dataloader = DataLoaderDispatcher( - new_dataset, - split_batches=split_batches, - batch_sampler=new_batch_sampler, - _drop_last=dataloader.drop_last, - _non_blocking=non_blocking, - slice_fn=slice_fn_for_dispatch, - use_stateful_dataloader=use_stateful_dataloader, - torch_device_mesh=torch_device_mesh, - **kwargs, - ) - elif sampler_is_batch_sampler: - dataloader = DataLoaderShard( - new_dataset, - device=device if put_on_device and state.distributed_type != DistributedType.XLA else None, - sampler=new_batch_sampler, - batch_size=dataloader.batch_size, - rng_types=rng_types, - _drop_last=dataloader.drop_last, - _non_blocking=non_blocking, - synchronized_generator=synchronized_generator, - use_stateful_dataloader=use_stateful_dataloader, - **kwargs, - ) - else: - dataloader = DataLoaderShard( - new_dataset, - device=device if put_on_device and state.distributed_type != DistributedType.XLA else None, - batch_sampler=new_batch_sampler, - rng_types=rng_types, - synchronized_generator=synchronized_generator, - _drop_last=dataloader.drop_last, - _non_blocking=non_blocking, - use_stateful_dataloader=use_stateful_dataloader, - **kwargs, - ) - - if isinstance(sampler, SeedableRandomSampler) and use_seedable_sampler: - dataloader.set_sampler(sampler) - if state.distributed_type == DistributedType.XLA: - return MpDeviceLoaderWrapper(dataloader, device) - return dataloader - - -class SkipBatchSampler(BatchSampler): - """ - A `torch.utils.data.BatchSampler` that skips the first `n` batches of another `torch.utils.data.BatchSampler`. - Should not be used if the original dataloader is a `StatefulDataLoader`. - """ - - def __init__(self, batch_sampler, skip_batches=0): - self.batch_sampler = batch_sampler - self.skip_batches = skip_batches - - def __iter__(self): - for index, samples in enumerate(self.batch_sampler): - if index >= self.skip_batches: - yield samples - - @property - def total_length(self): - return len(self.batch_sampler) - - def __len__(self): - return len(self.batch_sampler) - self.skip_batches - - -class SkipDataLoader(DataLoaderAdapter, DataLoaderStateMixin): - """ - Subclass of a PyTorch `DataLoader` that will skip the first batches. Generally it's preferable to use - `skip_first_batches`/`torchdata.StatefulDataLoader` instead of this class. - - Args: - dataset (`torch.utils.data.dataset.Dataset`): - The dataset to use to build this dataloader. - skip_batches (`int`, *optional*, defaults to 0): - The number of batches to skip at the beginning. - kwargs: - All other keyword arguments to pass to the regular `DataLoader` initialization. - """ - - def __init__(self, dataset, skip_batches=0, use_stateful_dataloader=False, **kwargs): - super().__init__(dataset, use_stateful_dataloader=use_stateful_dataloader, **kwargs) - self.skip_batches = skip_batches - self.gradient_state = GradientState() - - def __iter__(self): - self.begin() - for index, batch in enumerate(self.base_dataloader.__iter__()): - if index >= self.skip_batches: - self._update_state_dict() - yield batch - self.end() - - def __len__(self): - return len(self.base_dataloader) - self.skip_batches - - def __reduce__(self): - """ - Define the `__reduce__` method to ensure a `SkipDataLoader` can be pickled and unpickled. This needs to be - explicitly defined since default pickling behavior is broken by `DataLoaderAdapter` messing with its - `__class__` member. - """ - args = super().__reduce__() - return (SkipDataLoader, *args[1:]) - - -def skip_first_batches(dataloader, num_batches=0): - """ - Creates a `torch.utils.data.DataLoader` that will efficiently skip the first `num_batches`. Should not be used if - the original dataloader is a `StatefulDataLoader`. - """ - state = PartialState() - if state.distributed_type == DistributedType.XLA: - device = dataloader.device - dataloader = dataloader.dataloader - - dataset = dataloader.dataset - sampler_is_batch_sampler = False - if isinstance(dataset, IterableDataset): - new_batch_sampler = None - else: - sampler_is_batch_sampler = isinstance(dataloader.sampler, BatchSampler) - batch_sampler = dataloader.sampler if sampler_is_batch_sampler else dataloader.batch_sampler - new_batch_sampler = SkipBatchSampler(batch_sampler, skip_batches=num_batches) - - # We ignore all of those since they are all dealt with by our new_batch_sampler - ignore_kwargs = [ - "batch_size", - "shuffle", - "sampler", - "batch_sampler", - "drop_last", - ] - - kwargs = { - k: getattr(dataloader, k, _PYTORCH_DATALOADER_KWARGS[k]) - for k in _PYTORCH_DATALOADER_KWARGS - if k not in ignore_kwargs - } - - # Need to provide batch_size as batch_sampler is None for Iterable dataset - if new_batch_sampler is None: - kwargs["drop_last"] = dataloader.drop_last - kwargs["batch_size"] = dataloader.batch_size - - if isinstance(dataloader, DataLoaderDispatcher): - if new_batch_sampler is None: - # Need to manually skip batches in the dataloader - kwargs["skip_batches"] = num_batches - dataloader = DataLoaderDispatcher( - dataset, - split_batches=dataloader.split_batches, - batch_sampler=new_batch_sampler, - _drop_last=dataloader._drop_last, - **kwargs, - ) - elif isinstance(dataloader, DataLoaderShard): - if new_batch_sampler is None: - # Need to manually skip batches in the dataloader - kwargs["skip_batches"] = num_batches - elif sampler_is_batch_sampler: - kwargs["sampler"] = new_batch_sampler - kwargs["batch_size"] = dataloader.batch_size - else: - kwargs["batch_sampler"] = new_batch_sampler - dataloader = DataLoaderShard( - dataset, - device=dataloader.device, - rng_types=dataloader.rng_types, - synchronized_generator=dataloader.synchronized_generator, - **kwargs, - ) - else: - if new_batch_sampler is None: - # Need to manually skip batches in the dataloader - dataloader = SkipDataLoader(dataset, skip_batches=num_batches, **kwargs) - else: - dataloader = DataLoader(dataset, batch_sampler=new_batch_sampler, **kwargs) - - if state.distributed_type == DistributedType.XLA: - dataloader = MpDeviceLoaderWrapper(dataloader, device) - - return dataloader diff --git a/venv/lib/python3.10/site-packages/accelerate/hooks.py b/venv/lib/python3.10/site-packages/accelerate/hooks.py deleted file mode 100644 index 55b6f2efb8b7f24ec38a4da37f61992fd20519de..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/hooks.py +++ /dev/null @@ -1,765 +0,0 @@ -# Copyright 2022 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 functools -from collections.abc import Mapping -from typing import Optional, Union - -import torch -import torch.nn as nn - -from .state import PartialState -from .utils import ( - PrefixedDataset, - find_device, - named_module_tensors, - send_to_device, - set_module_tensor_to_device, -) -from .utils.imports import ( - is_mlu_available, - is_musa_available, - is_npu_available, -) -from .utils.memory import clear_device_cache -from .utils.modeling import get_non_persistent_buffers -from .utils.other import recursive_getattr - - -_accelerate_added_attributes = ["to", "cuda", "npu", "xpu", "mlu", "sdaa", "musa"] - - -class ModelHook: - """ - A hook that contains callbacks to be executed just before and after the forward method of a model. The difference - with PyTorch existing hooks is that they get passed along the kwargs. - - Class attribute: - - **no_grad** (`bool`, *optional*, defaults to `False`) -- Whether or not to execute the actual forward pass under - the `torch.no_grad()` context manager. - """ - - no_grad = False - - def init_hook(self, module): - """ - To be executed when the hook is attached to the module. - - Args: - module (`torch.nn.Module`): The module attached to this hook. - """ - return module - - def pre_forward(self, module, *args, **kwargs): - """ - To be 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]`): The positional arguments passed to the module. - kwargs (`Dict[Str, Any]`): The keyword arguments passed to the module. - - Returns: - `Tuple[Tuple[Any], Dict[Str, Any]]`: A tuple with the treated `args` and `kwargs`. - """ - return args, kwargs - - def post_forward(self, module, output): - """ - To be executed just after the forward method of the model. - - Args: - module (`torch.nn.Module`): The module whose forward pass been executed just before this event. - output (`Any`): The output of the module. - - Returns: - `Any`: The processed `output`. - """ - return output - - def detach_hook(self, module): - """ - To be executed when the hook is detached from a module. - - Args: - module (`torch.nn.Module`): The module detached from this hook. - """ - return module - - -class SequentialHook(ModelHook): - """ - A hook that can contain several hooks and iterates through them at each event. - """ - - def __init__(self, *hooks): - self.hooks = hooks - - def init_hook(self, module): - for hook in self.hooks: - module = hook.init_hook(module) - return module - - def pre_forward(self, module, *args, **kwargs): - for hook in self.hooks: - args, kwargs = hook.pre_forward(module, *args, **kwargs) - return args, kwargs - - def post_forward(self, module, output): - for hook in self.hooks: - output = hook.post_forward(module, output) - return output - - def detach_hook(self, module): - for hook in self.hooks: - module = hook.detach_hook(module) - return module - - -def add_hook_to_module(module: nn.Module, hook: ModelHook, append: bool = False): - """ - Adds a hook to a given module. This will rewrite the `forward` method of the module to include the hook, to remove - this behavior and restore the original `forward` method, use `remove_hook_from_module`. - - - - If the module already contains a hook, this will replace it with the new hook passed by default. To chain two hooks - together, pass `append=True`, so it chains the current and new hook into an instance of the `SequentialHook` class. - - - - Args: - module (`torch.nn.Module`): - The module to attach a hook to. - hook (`ModelHook`): - The hook to attach. - append (`bool`, *optional*, defaults to `False`): - Whether the hook should be chained with an existing one (if module already contains a hook) or not. - - Returns: - `torch.nn.Module`: The same module, with the hook attached (the module is modified in place, so the result can - be discarded). - """ - if append and (getattr(module, "_hf_hook", None) is not None): - old_hook = module._hf_hook - remove_hook_from_module(module) - hook = SequentialHook(old_hook, hook) - - if hasattr(module, "_hf_hook") and hasattr(module, "_old_forward"): - # If we already put some hook on this module, we replace it with the new one. - old_forward = module._old_forward - else: - old_forward = module.forward - module._old_forward = old_forward - - module = hook.init_hook(module) - module._hf_hook = hook - - def new_forward(module, *args, **kwargs): - args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs) - if module._hf_hook.no_grad: - with torch.no_grad(): - output = module._old_forward(*args, **kwargs) - else: - output = module._old_forward(*args, **kwargs) - return module._hf_hook.post_forward(module, output) - - # Overriding a GraphModuleImpl forward freezes the forward call and later modifications on the graph will fail. - # Reference: https://pytorch.slack.com/archives/C3PDTEV8E/p1705929610405409 - if "GraphModuleImpl" in str(type(module)): - module.__class__.forward = functools.update_wrapper(functools.partial(new_forward, module), old_forward) - else: - module.forward = functools.update_wrapper(functools.partial(new_forward, module), old_forward) - - return module - - -def remove_hook_from_module(module: nn.Module, recurse=False): - """ - Removes any hook attached to a module via `add_hook_to_module`. - - Args: - module (`torch.nn.Module`): The module to attach a hook to. - recurse (`bool`, **optional**): Whether to remove the hooks recursively - - Returns: - `torch.nn.Module`: The same module, with the hook detached (the module is modified in place, so the result can - be discarded). - """ - - if hasattr(module, "_hf_hook"): - module._hf_hook.detach_hook(module) - delattr(module, "_hf_hook") - - if hasattr(module, "_old_forward"): - # Overriding a GraphModuleImpl forward freezes the forward call and later modifications on the graph will fail. - # Reference: https://pytorch.slack.com/archives/C3PDTEV8E/p1705929610405409 - if "GraphModuleImpl" in str(type(module)): - module.__class__.forward = module._old_forward - else: - module.forward = module._old_forward - delattr(module, "_old_forward") - - # Remove accelerate added warning hooks from dispatch_model - for attr in _accelerate_added_attributes: - module.__dict__.pop(attr, None) - - if recurse: - for child in module.children(): - remove_hook_from_module(child, recurse) - - return module - - -class AlignDevicesHook(ModelHook): - """ - A generic `ModelHook` that ensures inputs and model weights are on the same device for the forward pass of the - associated module, potentially offloading the weights after the forward pass. - - Args: - execution_device (`torch.device`, *optional*): - The device on which inputs and model weights should be placed before the forward pass. - offload (`bool`, *optional*, defaults to `False`): - Whether or not the weights should be offloaded after the forward pass. - io_same_device (`bool`, *optional*, defaults to `False`): - Whether or not the output should be placed on the same device as the input was. - weights_map (`Mapping[str, torch.Tensor]`, *optional*): - When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values. - offload_buffers (`bool`, *optional*, defaults to `False`): - Whether or not to include the associated module's buffers when offloading. - place_submodules (`bool`, *optional*, defaults to `False`): - Whether to place the submodules on `execution_device` during the `init_hook` event. - """ - - def __init__( - self, - execution_device: Optional[Union[int, str, torch.device]] = None, - offload: bool = False, - io_same_device: bool = False, - weights_map: Optional[Mapping] = None, - offload_buffers: bool = False, - place_submodules: bool = False, - skip_keys: Optional[Union[str, list[str]]] = None, - tied_params_map: Optional[dict[int, dict[torch.device, torch.Tensor]]] = None, - ): - self.execution_device = execution_device - self.offload = offload - self.io_same_device = io_same_device - self.weights_map = weights_map - self.offload_buffers = offload_buffers - self.place_submodules = place_submodules - self.skip_keys = skip_keys - - # Will contain the input device when `io_same_device=True`. - self.input_device = None - self.param_original_devices = {} - self.buffer_original_devices = {} - self.tied_params_names = set() - - # The hook pre_forward/post_forward need to have knowledge of this dictionary, as with offloading we want to avoid duplicating memory - # for tied weights already loaded on the target execution device. - self.tied_params_map = tied_params_map - - def __repr__(self): - return ( - f"AlignDevicesHook(execution_device={self.execution_device}, offload={self.offload}, " - f"io_same_device={self.io_same_device}, offload_buffers={self.offload_buffers}, " - f"place_submodules={self.place_submodules}, skip_keys={repr(self.skip_keys)})" - ) - - def init_hook(self, module): - # In case the AlignDevicesHook is on meta device, ignore tied weights as data_ptr() is then always zero. - if self.execution_device == "meta" or self.execution_device == torch.device("meta"): - self.tied_params_map = None - - if not self.offload and self.execution_device is not None: - for name, _ in named_module_tensors(module, recurse=self.place_submodules): - set_module_tensor_to_device(module, name, self.execution_device, tied_params_map=self.tied_params_map) - elif self.offload: - self.original_devices = { - name: param.device for name, param in named_module_tensors(module, recurse=self.place_submodules) - } - if self.weights_map is None: - self.weights_map = { - name: param.to("cpu") - for name, param in named_module_tensors( - module, include_buffers=self.offload_buffers, recurse=self.place_submodules - ) - } - for name, _ in named_module_tensors( - module, include_buffers=self.offload_buffers, recurse=self.place_submodules, remove_non_persistent=True - ): - # When using disk offloading, we can not rely on `weights_map[name].data_ptr()` as the reference pointer, - # as we have no guarantee that safetensors' `file.get_tensor()` will always give the same pointer. - # As we have no reliable way to track the shared data pointer of tied weights in this case, we use tied_params_names: List[str] - # to add on the fly pointers to `tied_params_map` in the pre_forward call. - if ( - self.tied_params_map is not None - and recursive_getattr(module, name).data_ptr() in self.tied_params_map - ): - self.tied_params_names.add(name) - - set_module_tensor_to_device(module, name, "meta") - - if not self.offload_buffers and self.execution_device is not None: - for name, _ in module.named_buffers(recurse=self.place_submodules): - set_module_tensor_to_device( - module, name, self.execution_device, tied_params_map=self.tied_params_map - ) - elif self.offload_buffers and self.execution_device is not None: - for name in get_non_persistent_buffers(module, recurse=self.place_submodules): - set_module_tensor_to_device( - module, name, self.execution_device, tied_params_map=self.tied_params_map - ) - - return module - - def pre_forward(self, module, *args, **kwargs): - if self.io_same_device: - self.input_device = find_device([args, kwargs]) - if self.offload: - self.tied_pointers_to_remove = set() - - for name, _ in named_module_tensors( - module, - include_buffers=self.offload_buffers, - recurse=self.place_submodules, - remove_non_persistent=True, - ): - fp16_statistics = None - value = self.weights_map[name] - if "weight" in name and name.replace("weight", "SCB") in self.weights_map.keys(): - if value.dtype == torch.int8: - fp16_statistics = self.weights_map[name.replace("weight", "SCB")] - - # In case we are using offloading with tied weights, we need to keep track of the offloaded weights - # that are loaded on device at this point, as we will need to remove them as well from the dictionary - # self.tied_params_map in order to allow to free memory. - if name in self.tied_params_names and value.data_ptr() not in self.tied_params_map: - self.tied_params_map[value.data_ptr()] = {} - - if ( - value is not None - and self.tied_params_map is not None - and value.data_ptr() in self.tied_params_map - and self.execution_device not in self.tied_params_map[value.data_ptr()] - ): - self.tied_pointers_to_remove.add((value.data_ptr(), self.execution_device)) - - set_module_tensor_to_device( - module, - name, - self.execution_device, - value=value, - fp16_statistics=fp16_statistics, - tied_params_map=self.tied_params_map, - ) - - return send_to_device(args, self.execution_device), send_to_device( - kwargs, self.execution_device, skip_keys=self.skip_keys - ) - - def post_forward(self, module, output): - if self.offload: - for name, _ in named_module_tensors( - module, - include_buffers=self.offload_buffers, - recurse=self.place_submodules, - remove_non_persistent=True, - ): - set_module_tensor_to_device(module, name, "meta") - if type(module).__name__ == "Linear8bitLt": - module.state.SCB = None - module.state.CxB = None - - # We may have loaded tied weights into self.tied_params_map (avoiding to load them several times in e.g. submodules): remove them from - # this dictionary to allow the garbage collector to do its job. - for value_pointer, device in self.tied_pointers_to_remove: - if isinstance(device, int): - if is_npu_available(): - device = f"npu:{device}" - elif is_mlu_available(): - device = f"mlu:{device}" - elif is_musa_available(): - device = f"musa:{device}" - if device in self.tied_params_map[value_pointer]: - del self.tied_params_map[value_pointer][device] - self.tied_pointers_to_remove = set() - if self.io_same_device and self.input_device is not None: - output = send_to_device(output, self.input_device, skip_keys=self.skip_keys) - - return output - - def detach_hook(self, module): - if self.offload: - for name, device in self.original_devices.items(): - if device != torch.device("meta"): - set_module_tensor_to_device(module, name, device, value=self.weights_map.get(name, None)) - return module - - -def attach_execution_device_hook( - module: torch.nn.Module, - execution_device: Union[int, str, torch.device], - skip_keys: Optional[Union[str, list[str]]] = None, - preload_module_classes: Optional[list[str]] = None, - tied_params_map: Optional[dict[int, dict[torch.device, torch.Tensor]]] = None, -): - """ - Recursively attaches `AlignDevicesHook` to all submodules of a given model to make sure they have the right - execution device - - Args: - module (`torch.nn.Module`): - The module where we want to attach the hooks. - execution_device (`int`, `str` or `torch.device`): - The device on which inputs and model weights should be placed before the forward pass. - skip_keys (`str` or `List[str]`, *optional*): - A list of keys to ignore when moving inputs or outputs between devices. - preload_module_classes (`List[str]`, *optional*): - A list of classes whose instances should load all their weights (even in the submodules) at the beginning - of the forward. This should only be used for classes that have submodules which are registered but not - called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, - `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. - tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], *optional*, defaults to `None`): - A map of data pointers to dictionaries of devices to already dispatched tied weights. For a given execution - device, this parameter is useful to reuse the first available pointer of a shared weight for all others, - instead of duplicating memory. - """ - if not hasattr(module, "_hf_hook") and len(module.state_dict()) > 0: - add_hook_to_module( - module, - AlignDevicesHook(execution_device, skip_keys=skip_keys, tied_params_map=tied_params_map), - ) - - # Break the recursion if we get to a preload module. - if preload_module_classes is not None and module.__class__.__name__ in preload_module_classes: - return - - for child in module.children(): - attach_execution_device_hook( - child, - execution_device, - skip_keys=skip_keys, - preload_module_classes=preload_module_classes, - tied_params_map=tied_params_map, - ) - - -def attach_align_device_hook( - module: torch.nn.Module, - execution_device: Optional[torch.device] = None, - offload: bool = False, - weights_map: Optional[Mapping] = None, - offload_buffers: bool = False, - module_name: str = "", - skip_keys: Optional[Union[str, list[str]]] = None, - preload_module_classes: Optional[list[str]] = None, - tied_params_map: Optional[dict[int, dict[torch.device, torch.Tensor]]] = None, -): - """ - Recursively attaches `AlignDevicesHook` to all submodules of a given model that have direct parameters and/or - buffers. - - Args: - module (`torch.nn.Module`): - The module where we want to attach the hooks. - execution_device (`torch.device`, *optional*): - The device on which inputs and model weights should be placed before the forward pass. - offload (`bool`, *optional*, defaults to `False`): - Whether or not the weights should be offloaded after the forward pass. - weights_map (`Mapping[str, torch.Tensor]`, *optional*): - When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values. - offload_buffers (`bool`, *optional*, defaults to `False`): - Whether or not to include the associated module's buffers when offloading. - module_name (`str`, *optional*, defaults to `""`): - The name of the module. - skip_keys (`str` or `List[str]`, *optional*): - A list of keys to ignore when moving inputs or outputs between devices. - preload_module_classes (`List[str]`, *optional*): - A list of classes whose instances should load all their weights (even in the submodules) at the beginning - of the forward. This should only be used for classes that have submodules which are registered but not - called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, - `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. - tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], *optional*, defaults to `None`): - A map of data pointers to dictionaries of devices to already dispatched tied weights. For a given execution - device, this parameter is useful to reuse the first available pointer of a shared weight for all others, - instead of duplicating memory. - """ - # Attach the hook on this module if it has any direct tensor. - directs = named_module_tensors(module) - full_offload = ( - offload and preload_module_classes is not None and module.__class__.__name__ in preload_module_classes - ) - - if len(list(directs)) > 0 or full_offload: - if weights_map is not None: - prefix = f"{module_name}." if len(module_name) > 0 else "" - prefixed_weights_map = PrefixedDataset(weights_map, prefix) - else: - prefixed_weights_map = None - hook = AlignDevicesHook( - execution_device=execution_device, - offload=offload, - weights_map=prefixed_weights_map, - offload_buffers=offload_buffers, - place_submodules=full_offload, - skip_keys=skip_keys, - tied_params_map=tied_params_map, - ) - add_hook_to_module(module, hook, append=True) - - # We stop the recursion in case we hit the full offload. - if full_offload: - return - - # Recurse on all children of the module. - for child_name, child in module.named_children(): - child_name = f"{module_name}.{child_name}" if len(module_name) > 0 else child_name - attach_align_device_hook( - child, - execution_device=execution_device, - offload=offload, - weights_map=weights_map, - offload_buffers=offload_buffers, - module_name=child_name, - preload_module_classes=preload_module_classes, - skip_keys=skip_keys, - tied_params_map=tied_params_map, - ) - - -def remove_hook_from_submodules(module: nn.Module): - """ - Recursively removes all hooks attached on the submodules of a given model. - - Args: - module (`torch.nn.Module`): The module on which to remove all hooks. - """ - remove_hook_from_module(module) - for child in module.children(): - remove_hook_from_submodules(child) - - -def attach_align_device_hook_on_blocks( - module: nn.Module, - execution_device: Optional[Union[torch.device, dict[str, torch.device]]] = None, - offload: Union[bool, dict[str, bool]] = False, - weights_map: Mapping = None, - offload_buffers: bool = False, - module_name: str = "", - skip_keys: Optional[Union[str, list[str]]] = None, - preload_module_classes: Optional[list[str]] = None, - tied_params_map: Optional[dict[int, dict[torch.device, torch.Tensor]]] = None, -): - """ - Attaches `AlignDevicesHook` to all blocks of a given model as needed. - - Args: - module (`torch.nn.Module`): - The module where we want to attach the hooks. - execution_device (`torch.device` or `Dict[str, torch.device]`, *optional*): - The device on which inputs and model weights should be placed before the forward pass. It can be one device - for the whole module, or a dictionary mapping module name to device. - offload (`bool`, *optional*, defaults to `False`): - Whether or not the weights should be offloaded after the forward pass. It can be one boolean for the whole - module, or a dictionary mapping module name to boolean. - weights_map (`Mapping[str, torch.Tensor]`, *optional*): - When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values. - offload_buffers (`bool`, *optional*, defaults to `False`): - Whether or not to include the associated module's buffers when offloading. - module_name (`str`, *optional*, defaults to `""`): - The name of the module. - skip_keys (`str` or `List[str]`, *optional*): - A list of keys to ignore when moving inputs or outputs between devices. - preload_module_classes (`List[str]`, *optional*): - A list of classes whose instances should load all their weights (even in the submodules) at the beginning - of the forward. This should only be used for classes that have submodules which are registered but not - called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, - `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. - tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], *optional*, defaults to `None`): - A map of data pointers to dictionaries of devices to already dispatched tied weights. For a given execution - device, this parameter is useful to reuse the first available pointer of a shared weight for all others, - instead of duplicating memory. - """ - # If one device and one offload, we've got one hook. - if not isinstance(execution_device, Mapping) and not isinstance(offload, dict): - if not offload: - hook = AlignDevicesHook( - execution_device=execution_device, - io_same_device=True, - skip_keys=skip_keys, - place_submodules=True, - tied_params_map=tied_params_map, - ) - add_hook_to_module(module, hook) - else: - attach_align_device_hook( - module, - execution_device=execution_device, - offload=True, - weights_map=weights_map, - offload_buffers=offload_buffers, - module_name=module_name, - skip_keys=skip_keys, - tied_params_map=tied_params_map, - ) - return - - if not isinstance(execution_device, Mapping): - execution_device = {key: execution_device for key in offload.keys()} - if not isinstance(offload, Mapping): - offload = {key: offload for key in execution_device.keys()} - - if module_name in execution_device and module_name in offload and not offload[module_name]: - hook = AlignDevicesHook( - execution_device=execution_device[module_name], - offload_buffers=offload_buffers, - io_same_device=(module_name == ""), - place_submodules=True, - skip_keys=skip_keys, - tied_params_map=tied_params_map, - ) - add_hook_to_module(module, hook) - attach_execution_device_hook( - module, execution_device[module_name], skip_keys=skip_keys, tied_params_map=tied_params_map - ) - elif module_name in execution_device and module_name in offload: - attach_align_device_hook( - module, - execution_device=execution_device[module_name], - offload=True, - weights_map=weights_map, - offload_buffers=offload_buffers, - module_name=module_name, - skip_keys=skip_keys, - preload_module_classes=preload_module_classes, - tied_params_map=tied_params_map, - ) - if not hasattr(module, "_hf_hook"): - hook = AlignDevicesHook( - execution_device=execution_device[module_name], - io_same_device=(module_name == ""), - skip_keys=skip_keys, - tied_params_map=tied_params_map, - ) - add_hook_to_module(module, hook) - attach_execution_device_hook( - module, - execution_device[module_name], - preload_module_classes=preload_module_classes, - skip_keys=skip_keys, - tied_params_map=tied_params_map, - ) - elif module_name == "": - hook = AlignDevicesHook( - execution_device=execution_device.get(""), - io_same_device=True, - skip_keys=skip_keys, - tied_params_map=tied_params_map, - ) - add_hook_to_module(module, hook) - - for child_name, child in module.named_children(): - child_name = f"{module_name}.{child_name}" if len(module_name) > 0 else child_name - attach_align_device_hook_on_blocks( - child, - execution_device=execution_device, - offload=offload, - weights_map=weights_map, - offload_buffers=offload_buffers, - module_name=child_name, - preload_module_classes=preload_module_classes, - skip_keys=skip_keys, - tied_params_map=tied_params_map, - ) - - -class CpuOffload(ModelHook): - """ - Offloads a model on the CPU until its forward pass is called. The model will not be offloaded back to the CPU after - the forward, the user needs to call the `init_hook` method again for this. - - Args: - execution_device(`str`, `int` or `torch.device`, *optional*): - The device on which the model should be executed. Will default to the MPS device if it's available, then - GPU 0 if there is a GPU, and finally to the CPU. - prev_module_hook (`UserCpuOffloadHook`, *optional*): - The hook sent back by [`cpu_offload_with_hook`] for a previous model in the pipeline you are running. If - passed, its offload method will be called just before the forward of the model to which this hook is - attached. - """ - - def __init__( - self, - execution_device: Optional[Union[str, int, torch.device]] = None, - prev_module_hook: Optional["UserCpuOffloadHook"] = None, - ): - self.prev_module_hook = prev_module_hook - - self.execution_device = execution_device if execution_device is not None else PartialState().default_device - - def init_hook(self, module): - return module.to("cpu") - - def pre_forward(self, module, *args, **kwargs): - if self.prev_module_hook is not None: - self.prev_module_hook.offload() - clear_device_cache() - module.to(self.execution_device) - return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device) - - -class UserCpuOffloadHook: - """ - A simple hook grouping a model and a `ModelHook`, which provides easy APIs for to call the init method of the hook - or remove it entirely. - """ - - def __init__(self, model, hook): - self.model = model - self.hook = hook - - def offload(self): - self.hook.init_hook(self.model) - - def remove(self): - remove_hook_from_module(self.model) - - -class LayerwiseCastingHook(ModelHook): - r""" - A hook that casts the weights of a module to a high precision dtype for computation, and to a low precision dtype - for storage. This process may lead to quality loss in the output, but can significantly reduce the memory - footprint. - """ - - _is_stateful = False - - def __init__(self, storage_dtype: torch.dtype, compute_dtype: torch.dtype, non_blocking: bool) -> None: - self.storage_dtype = storage_dtype - self.compute_dtype = compute_dtype - self.non_blocking = non_blocking - - def init_hook(self, module: torch.nn.Module): - module.to(dtype=self.storage_dtype, non_blocking=self.non_blocking) - return module - - def pre_forward(self, module: torch.nn.Module, *args, **kwargs): - module.to(dtype=self.compute_dtype, non_blocking=self.non_blocking) - return args, kwargs - - def post_forward(self, module: torch.nn.Module, output): - module.to(dtype=self.storage_dtype, non_blocking=self.non_blocking) - return output diff --git a/venv/lib/python3.10/site-packages/accelerate/inference.py b/venv/lib/python3.10/site-packages/accelerate/inference.py deleted file mode 100644 index d1659718b0b3020955f04060aa6f0cc7ef66d7ff..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/inference.py +++ /dev/null @@ -1,184 +0,0 @@ -# Copyright 2024 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 types import MethodType -from typing import Any, Optional, Union - -from .state import PartialState -from .utils import ( - calculate_maximum_sizes, - convert_bytes, - copy_tensor_to_devices, - ignorant_find_batch_size, - infer_auto_device_map, - is_pippy_available, - pad_input_tensors, - send_to_device, -) - - -def generate_device_map(model, num_processes: int = 1, no_split_module_classes=None, max_memory: dict = None): - """ - Calculates the device map for `model` with an offset for PiPPy - """ - if num_processes == 1: - return infer_auto_device_map(model, no_split_module_classes=no_split_module_classes, clean_result=False) - if max_memory is None: - model_size, shared = calculate_maximum_sizes(model) - - # Split into `n` chunks for each GPU - memory = (model_size + shared[0]) / num_processes - memory = convert_bytes(memory) - value, ending = memory.split(" ") - - # Add a chunk to deal with potential extra shared memory instances - memory = math.ceil(float(value)) * 1.1 - memory = f"{memory} {ending}" - max_memory = {i: memory for i in range(num_processes)} - device_map = infer_auto_device_map( - model, - max_memory=max_memory, - no_split_module_classes=no_split_module_classes, - clean_result=False, - ) - return device_map - - -def find_pippy_batch_size(args, kwargs): - found_batch_size = None - if args is not None: - for arg in args: - found_batch_size = ignorant_find_batch_size(arg) - if found_batch_size is not None: - break - if kwargs is not None and found_batch_size is None: - for kwarg in kwargs.values(): - found_batch_size = ignorant_find_batch_size(kwarg) - if found_batch_size is not None: - break - return found_batch_size - - -def build_pipeline(model, split_points, args, kwargs, num_chunks): - """ - Attaches the split points to the model based on `self.device_map` and generates a `PipelineStage`. Requires passing - in needed `args` and `kwargs` as the model needs on the CPU. - - Users can pass in custom `num_chunks` as an optional hyper-parameter. By default will use - `AcceleratorState.num_processes` - """ - # Note: We import here to reduce import time from general modules, and isolate outside dependencies - from torch.distributed.pipelining import ScheduleGPipe, SplitPoint, pipeline - - # We need to annotate the split points in the model for PiPPy - state = PartialState() - split_spec = {split_point: SplitPoint.BEGINNING for split_point in split_points} - pipe = pipeline( - model, - mb_args=args, - mb_kwargs=kwargs, - split_spec=split_spec, - ) - stage = pipe.build_stage(state.local_process_index, device=state.device) - schedule = ScheduleGPipe(stage, num_chunks) - - return schedule - - -def pippy_forward(forward, num_chunks, gather_output, *args, **kwargs): - state = PartialState() - output = None - - if state.num_processes == 1: - output = forward(*args, **kwargs) - elif state.is_local_main_process: - found_batch_size = find_pippy_batch_size(args, kwargs) - if found_batch_size is None: - raise ValueError("Could not find batch size from args or kwargs") - else: - if found_batch_size != num_chunks: - args = pad_input_tensors(args, found_batch_size, num_chunks) - kwargs = pad_input_tensors(kwargs, found_batch_size, num_chunks) - forward(*args, **kwargs) - elif state.is_last_process: - output = forward() - else: - forward() - if gather_output: - # Each node will get a copy of the full output which is only on the last GPU - output = copy_tensor_to_devices(output) - return output - - -def prepare_pippy( - model, - split_points: Optional[Union[str, list[str]]] = "auto", - no_split_module_classes: Optional[list[str]] = None, - example_args: Optional[tuple[Any]] = (), - example_kwargs: Optional[dict[str, Any]] = None, - num_chunks: Optional[int] = None, - gather_output: Optional[bool] = False, -): - """ - Wraps `model` for pipeline parallel inference. - - Args: - model (`torch.nn.Module`): - A model we want to split for pipeline-parallel inference - split_points (`str` or `List[str]`, defaults to 'auto'): - How to generate the split points and chunk the model across each GPU. 'auto' will find the best balanced - split given any model. Should be a list of layer names in the model to split by otherwise. - no_split_module_classes (`List[str]`): - A list of class names for layers we don't want to be split. - example_args (tuple of model inputs): - The expected inputs for the model that uses order-based inputs for a *single process*. Recommended to use - this method if possible. - example_kwargs (dict of model inputs) - The expected inputs for the model that uses dictionary-based inputs for a *single process*. This is a - *highly* limiting structure that requires the same keys be present at *all* inference calls. Not - recommended unless the prior condition is true for all cases. - num_chunks (`int`, defaults to the number of available GPUs): - The number of different stages the Pipeline will have. By default it will assign one chunk per GPU, but - this can be tuned and played with. In general one should have num_chunks >= num_gpus. - gather_output (`bool`, defaults to `False`): - If `True`, the output from the last GPU (which holds the true outputs) is sent across to all GPUs. - """ - if not is_pippy_available(): - raise ImportError("Using `torch.distributed.pipelining` requires PyTorch 2.4.0 or later.") - state = PartialState() - example_args = send_to_device(example_args, "cpu") - example_kwargs = send_to_device(example_kwargs, "cpu") - if num_chunks is None: - num_chunks = state.num_processes - if split_points == "auto": - device_map = generate_device_map(model, num_chunks, no_split_module_classes=no_split_module_classes) - split_points = [] - for i in range(1, num_chunks): - split_points.append(next(k for k, v in device_map.items() if v == i)) - model.hf_split_points = split_points - stage = build_pipeline(model, split_points, example_args, example_kwargs, num_chunks) - model._original_forward = model.forward - model._original_call = model.__call__ - model.pippy_stage = stage - model.hf_split_points = split_points - - def forward(*args, **kwargs): - return pippy_forward(stage.step, num_chunks, gather_output, *args, **kwargs) - - # To act like a decorator so that it can be popped when doing `extract_model_from_parallel` - # Note: creates an infinite recursion loop with `generate` - model_forward = MethodType(forward, model) - forward.__wrapped__ = model_forward - model.forward = forward - return model diff --git a/venv/lib/python3.10/site-packages/accelerate/launchers.py b/venv/lib/python3.10/site-packages/accelerate/launchers.py deleted file mode 100644 index 532c0ad4450bb22870efa5a3d46cedbfbca5a28b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/launchers.py +++ /dev/null @@ -1,306 +0,0 @@ -# Copyright 2021 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 os -import sys -import tempfile - -import torch - -from .state import AcceleratorState, PartialState -from .utils import ( - PrecisionType, - PrepareForLaunch, - are_libraries_initialized, - check_cuda_p2p_ib_support, - get_gpu_info, - is_mps_available, - is_torch_version, - patch_environment, -) -from .utils.constants import ELASTIC_LOG_LINE_PREFIX_TEMPLATE_PYTORCH_VERSION - - -def test_launch(): - "Verify a `PartialState` can be initialized." - _ = PartialState() - - -def notebook_launcher( - function, - args=(), - num_processes=None, - mixed_precision="no", - use_port="29500", - master_addr="127.0.0.1", - node_rank=0, - num_nodes=1, - rdzv_backend="static", - rdzv_endpoint="", - rdzv_conf=None, - rdzv_id="none", - max_restarts=0, - monitor_interval=0.1, - log_line_prefix_template=None, -): - """ - Launches a training function, using several processes or multiple nodes if it's possible in the current environment - (TPU with multiple cores for instance). - - - - To use this function absolutely zero calls to a device must be made in the notebook session before calling. If any - have been made, you will need to restart the notebook and make sure no cells use any device capability. - - Setting `ACCELERATE_DEBUG_MODE="1"` in your environment will run a test before truly launching to ensure that none - of those calls have been made. - - - - Args: - function (`Callable`): - The training function to execute. If it accepts arguments, the first argument should be the index of the - process run. - args (`Tuple`): - Tuple of arguments to pass to the function (it will receive `*args`). - num_processes (`int`, *optional*): - The number of processes to use for training. Will default to 8 in Colab/Kaggle if a TPU is available, to - the number of devices available otherwise. - mixed_precision (`str`, *optional*, defaults to `"no"`): - If `fp16` or `bf16`, will use mixed precision training on multi-device. - use_port (`str`, *optional*, defaults to `"29500"`): - The port to use to communicate between processes when launching a multi-device training. - master_addr (`str`, *optional*, defaults to `"127.0.0.1"`): - The address to use for communication between processes. - node_rank (`int`, *optional*, defaults to 0): - The rank of the current node. - num_nodes (`int`, *optional*, defaults to 1): - The number of nodes to use for training. - rdzv_backend (`str`, *optional*, defaults to `"static"`): - The rendezvous method to use, such as 'static' (the default) or 'c10d' - rdzv_endpoint (`str`, *optional*, defaults to `""`): - The endpoint of the rdzv sync. storage. - rdzv_conf (`Dict`, *optional*, defaults to `None`): - Additional rendezvous configuration. - rdzv_id (`str`, *optional*, defaults to `"none"`): - The unique run id of the job. - max_restarts (`int`, *optional*, defaults to 0): - The maximum amount of restarts that elastic agent will conduct on workers before failure. - monitor_interval (`float`, *optional*, defaults to 0.1): - The interval in seconds that is used by the elastic_agent as a period of monitoring workers. - log_line_prefix_template (`str`, *optional*, defaults to `None`): - The prefix template for elastic launch logging. Available from PyTorch 2.2.0. - - Example: - - ```python - # Assume this is defined in a Jupyter Notebook on an instance with two devices - from accelerate import notebook_launcher - - - def train(*args): - # Your training function here - ... - - - notebook_launcher(train, args=(arg1, arg2), num_processes=2, mixed_precision="fp16") - ``` - """ - # Are we in a google colab or a Kaggle Kernel? - in_colab = False - in_kaggle = False - if any(key.startswith("KAGGLE") for key in os.environ.keys()): - in_kaggle = True - elif "IPython" in sys.modules: - in_colab = "google.colab" in str(sys.modules["IPython"].get_ipython()) - - try: - mixed_precision = PrecisionType(mixed_precision.lower()) - except ValueError: - raise ValueError( - f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}." - ) - - if (in_colab or in_kaggle) and ( - (os.environ.get("TPU_NAME", None) is not None) or (os.environ.get("PJRT_DEVICE", "") == "TPU") - ): - # TPU launch - import torch_xla.distributed.xla_multiprocessing as xmp - - if len(AcceleratorState._shared_state) > 0: - raise ValueError( - "To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside " - "your training function. Restart your notebook and make sure no cells initializes an " - "`Accelerator`." - ) - - launcher = PrepareForLaunch(function, distributed_type="XLA") - print("Launching a training on TPU cores.") - xmp.spawn(launcher, args=args, start_method="fork") - elif in_colab and get_gpu_info()[1] < 2: - # No need for a distributed launch otherwise as it's either CPU or one GPU. - if torch.cuda.is_available(): - print("Launching training on one GPU.") - else: - print("Launching training on one CPU.") - function(*args) - else: - if num_processes is None: - raise ValueError( - "You have to specify the number of devices you would like to use, add `num_processes=...` to your call." - ) - if node_rank >= num_nodes: - raise ValueError("The node_rank must be less than the number of nodes.") - if num_processes > 1: - # Multi-device launch - from torch.distributed.launcher.api import LaunchConfig, elastic_launch - from torch.multiprocessing import start_processes - from torch.multiprocessing.spawn import ProcessRaisedException - - if len(AcceleratorState._shared_state) > 0: - raise ValueError( - "To launch a multi-device training from your notebook, the `Accelerator` should only be initialized " - "inside your training function. Restart your notebook and make sure no cells initializes an " - "`Accelerator`." - ) - # Check for specific libraries known to initialize device that users constantly use - problematic_imports = are_libraries_initialized("bitsandbytes") - if len(problematic_imports) > 0: - err = ( - "Could not start distributed process. Libraries known to initialize device upon import have been " - "imported already. Please keep these imports inside your training function to try and help with this:" - ) - for lib_name in problematic_imports: - err += f"\n\t* `{lib_name}`" - raise RuntimeError(err) - - patched_env = dict( - nproc=num_processes, - node_rank=node_rank, - world_size=num_nodes * num_processes, - master_addr=master_addr, - master_port=use_port, - mixed_precision=mixed_precision, - ) - - # Check for CUDA P2P and IB issues - if not check_cuda_p2p_ib_support(): - patched_env["nccl_p2p_disable"] = "1" - patched_env["nccl_ib_disable"] = "1" - - # torch.distributed will expect a few environment variable to be here. We set the ones common to each - # process here (the other ones will be set be the launcher). - with patch_environment(**patched_env): - # First dummy launch - device_type = torch.accelerator.current_accelerator().type if hasattr(torch, "accelerator") else "cuda" - distributed_type = "MULTI_XPU" if device_type == "xpu" else "MULTI_GPU" - if os.environ.get("ACCELERATE_DEBUG_MODE", "false").lower() == "true": - launcher = PrepareForLaunch(test_launch, distributed_type=distributed_type) - try: - start_processes(launcher, args=(), nprocs=num_processes, start_method="fork") - except ProcessRaisedException as e: - err = "An issue was found when verifying a stable environment for the notebook launcher." - if f"Cannot re-initialize {device_type.upper()} in forked subprocess" in e.args[0]: - raise RuntimeError( - f"{err}" - "This likely stems from an outside import causing issues once the `notebook_launcher()` is called. " - "Please review your imports and test them when running the `notebook_launcher()` to identify " - f"which one is problematic and causing {device_type.upper()} to be initialized." - ) from e - else: - raise RuntimeError(f"{err} The following error was raised: {e}") from e - # Now the actual launch - launcher = PrepareForLaunch(function, distributed_type=distributed_type) - print(f"Launching training on {num_processes} {device_type.upper()}s.") - try: - if rdzv_conf is None: - rdzv_conf = {} - if rdzv_backend == "static": - rdzv_conf["rank"] = node_rank - if not rdzv_endpoint: - rdzv_endpoint = f"{master_addr}:{use_port}" - launch_config_kwargs = dict( - min_nodes=num_nodes, - max_nodes=num_nodes, - nproc_per_node=num_processes, - run_id=rdzv_id, - rdzv_endpoint=rdzv_endpoint, - rdzv_backend=rdzv_backend, - rdzv_configs=rdzv_conf, - max_restarts=max_restarts, - monitor_interval=monitor_interval, - start_method="fork", - ) - if is_torch_version(">=", ELASTIC_LOG_LINE_PREFIX_TEMPLATE_PYTORCH_VERSION): - launch_config_kwargs["log_line_prefix_template"] = log_line_prefix_template - elastic_launch(config=LaunchConfig(**launch_config_kwargs), entrypoint=function)(*args) - except ProcessRaisedException as e: - if f"Cannot re-initialize {device_type.upper()} in forked subprocess" in e.args[0]: - raise RuntimeError( - f"{device_type.upper()} has been initialized before the `notebook_launcher` could create a forked subprocess. " - "This likely stems from an outside import causing issues once the `notebook_launcher()` is called. " - "Please review your imports and test them when running the `notebook_launcher()` to identify " - f"which one is problematic and causing {device_type.upper()} to be initialized." - ) from e - else: - raise RuntimeError(f"An issue was found when launching the training: {e}") from e - - else: - # No need for a distributed launch otherwise as it's either CPU, GPU, XPU or MPS. - if is_mps_available(): - os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" - print("Launching training on MPS.") - elif torch.cuda.is_available(): - print("Launching training on one GPU.") - elif torch.xpu.is_available(): - print("Launching training on one XPU.") - else: - print("Launching training on CPU.") - function(*args) - - -def debug_launcher(function, args=(), num_processes=2): - """ - Launches a training function using several processes on CPU for debugging purposes. - - - - This function is provided for internal testing and debugging, but it's not intended for real trainings. It will - only use the CPU. - - - - Args: - function (`Callable`): - The training function to execute. - args (`Tuple`): - Tuple of arguments to pass to the function (it will receive `*args`). - num_processes (`int`, *optional*, defaults to 2): - The number of processes to use for training. - """ - from torch.multiprocessing import start_processes - - with tempfile.NamedTemporaryFile() as tmp_file: - # torch.distributed will expect a few environment variable to be here. We set the ones common to each - # process here (the other ones will be set be the launcher). - with patch_environment( - world_size=num_processes, - master_addr="127.0.0.1", - master_port="29500", - accelerate_mixed_precision="no", - accelerate_debug_rdv_file=tmp_file.name, - accelerate_use_cpu="yes", - ): - launcher = PrepareForLaunch(function, debug=True) - start_processes(launcher, args=args, nprocs=num_processes, start_method="fork") diff --git a/venv/lib/python3.10/site-packages/accelerate/local_sgd.py b/venv/lib/python3.10/site-packages/accelerate/local_sgd.py deleted file mode 100644 index 40b198d46aa24610d8eaf46424474d7c1c9881a1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/local_sgd.py +++ /dev/null @@ -1,106 +0,0 @@ -# Copyright 2023 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 accelerate import Accelerator, DistributedType - - -class LocalSGD: - """ - A helper class to support local SGD on top of Accelerator. It simply runs a given number of updates independently - on each device, and averages model weights every K synchronization step. - - It should be used only in the multi-GPU (or multi-CPU) setup without extensions such as DeepSpeed. In particular, - this is a simple implementation that cannot support scenarios such as model parallelism. - - - Although we are not aware of the true origins of this simple approach, the idea of local SGD is quite old and goes - back to at least: - - Zhang, J., De Sa, C., Mitliagkas, I., & Ré, C. (2016). [Parallel SGD: When does averaging help?. arXiv preprint - arXiv:1606.07365.](https://arxiv.org/abs/1606.07365) - - We credit the term Local SGD to the following paper (but there might be earlier references we are not aware of). - - Stich, Sebastian Urban. ["Local SGD Converges Fast and Communicates Little." ICLR 2019-International Conference on - Learning Representations. No. CONF. 2019.](https://arxiv.org/abs/1805.09767) - - """ - - def __enter__(self): - if self.enabled: - self.model_sync_obj = self.model.no_sync() - self.model_sync_obj.__enter__() - - return self - - def __exit__(self, type, value, tb): - if self.enabled: - # Average all models on exit - self._sync_and_avg_model_params() - self.model_sync_obj.__exit__(type, value, tb) - - def __init__(self, accelerator: Accelerator, model: torch.nn.Module, local_sgd_steps: int, enabled: bool = True): - """ - Constructor. - - Args: - model (`torch.nn.Module): - The model whose parameters we need to average. - accelerator (`Accelerator`): - Accelerator object. - local_sgd_steps (`int`): - A number of local SGD steps (before model parameters are synchronized). - enabled (`bool): - Local SGD is disabled if this parameter set to `False`. - """ - if accelerator.distributed_type not in [ - DistributedType.NO, - DistributedType.MULTI_CPU, - DistributedType.MULTI_GPU, - DistributedType.MULTI_XPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_HPU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_NPU, - ]: - raise NotImplementedError("LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)") - self.enabled = enabled and accelerator.distributed_type != DistributedType.NO - self.num_steps = 0 - if self.enabled: - self.accelerator = accelerator - self.model = model - self.local_sgd_steps = local_sgd_steps - - def step(self): - """ - This function makes a "step" and synchronizes model parameters if necessary. - """ - self.num_steps += 1 - if not self.enabled: - return - - if self.num_steps % self.local_sgd_steps == 0: - self._sync_and_avg_model_params() - - def _sync_and_avg_model_params(self): - """ - Synchronize + Average model parameters across all GPUs - """ - - self.accelerator.wait_for_everyone() - with self.accelerator.autocast(): - for param in self.model.parameters(): - param.data = self.accelerator.reduce(param.data, reduction="mean") diff --git a/venv/lib/python3.10/site-packages/accelerate/logging.py b/venv/lib/python3.10/site-packages/accelerate/logging.py deleted file mode 100644 index 1615bc313b74d4f01166384435b6c499ff616f49..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/logging.py +++ /dev/null @@ -1,125 +0,0 @@ -# Copyright 2022 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 functools -import logging -import os - -from .state import PartialState - - -class MultiProcessAdapter(logging.LoggerAdapter): - """ - An adapter to assist with logging in multiprocess. - - `log` takes in an additional `main_process_only` kwarg, which dictates whether it should be called on all processes - or only the main executed one. Default is `main_process_only=True`. - - Does not require an `Accelerator` object to be created first. - """ - - @staticmethod - def _should_log(main_process_only): - "Check if log should be performed" - state = PartialState() - return not main_process_only or (main_process_only and state.is_main_process) - - def log(self, level, msg, *args, **kwargs): - """ - Delegates logger call after checking if we should log. - - Accepts a new kwarg of `main_process_only`, which will dictate whether it will be logged across all processes - or only the main executed one. Default is `True` if not passed - - Also accepts "in_order", which if `True` makes the processes log one by one, in order. This is much easier to - read, but comes at the cost of sometimes needing to wait for the other processes. Default is `False` to not - break with the previous behavior. - - `in_order` is ignored if `main_process_only` is passed. - """ - if PartialState._shared_state == {}: - raise RuntimeError( - "You must initialize the accelerate state by calling either `PartialState()` or `Accelerator()` before using the logging utility." - ) - main_process_only = kwargs.pop("main_process_only", True) - in_order = kwargs.pop("in_order", False) - # set `stacklevel` to exclude ourself in `Logger.findCaller()` while respecting user's choice - kwargs.setdefault("stacklevel", 2) - - if self.isEnabledFor(level): - if self._should_log(main_process_only): - msg, kwargs = self.process(msg, kwargs) - self.logger.log(level, msg, *args, **kwargs) - - elif in_order: - state = PartialState() - for i in range(state.num_processes): - if i == state.process_index: - msg, kwargs = self.process(msg, kwargs) - self.logger.log(level, msg, *args, **kwargs) - state.wait_for_everyone() - - @functools.lru_cache(None) - def warning_once(self, *args, **kwargs): - """ - This method is identical to `logger.warning()`, but will emit the warning with the same message only once - - Note: The cache is for the function arguments, so 2 different callers using the same arguments will hit the - cache. The assumption here is that all warning messages are unique across the code. If they aren't then need to - switch to another type of cache that includes the caller frame information in the hashing function. - """ - self.warning(*args, **kwargs) - - -def get_logger(name: str, log_level: str = None): - """ - Returns a `logging.Logger` for `name` that can handle multiprocessing. - - If a log should be called on all processes, pass `main_process_only=False` If a log should be called on all - processes and in order, also pass `in_order=True` - - Args: - name (`str`): - The name for the logger, such as `__file__` - log_level (`str`, *optional*): - The log level to use. If not passed, will default to the `LOG_LEVEL` environment variable, or `INFO` if not - - Example: - - ```python - >>> from accelerate.logging import get_logger - >>> from accelerate import Accelerator - - >>> logger = get_logger(__name__) - - >>> accelerator = Accelerator() - >>> logger.info("My log", main_process_only=False) - >>> logger.debug("My log", main_process_only=True) - - >>> logger = get_logger(__name__, log_level="DEBUG") - >>> logger.info("My log") - >>> logger.debug("My second log") - - >>> array = ["a", "b", "c", "d"] - >>> letter_at_rank = array[accelerator.process_index] - >>> logger.info(letter_at_rank, in_order=True) - ``` - """ - if log_level is None: - log_level = os.environ.get("ACCELERATE_LOG_LEVEL", None) - logger = logging.getLogger(name) - if log_level is not None: - logger.setLevel(log_level.upper()) - logger.root.setLevel(log_level.upper()) - return MultiProcessAdapter(logger, {}) diff --git a/venv/lib/python3.10/site-packages/accelerate/memory_utils.py b/venv/lib/python3.10/site-packages/accelerate/memory_utils.py deleted file mode 100644 index fa2e2c8b9d7d0064c3e5e282737a7ad6919bde29..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/memory_utils.py +++ /dev/null @@ -1,22 +0,0 @@ -# Copyright 2022 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 warnings - - -warnings.warn( - "memory_utils has been reorganized to utils.memory. Import `find_executable_batchsize` from the main `__init__`: " - "`from accelerate import find_executable_batch_size` to avoid this warning.", - FutureWarning, -) diff --git a/venv/lib/python3.10/site-packages/accelerate/optimizer.py b/venv/lib/python3.10/site-packages/accelerate/optimizer.py deleted file mode 100644 index c1f8faa15436fb5d1d02ae4d691a4fcac44cd236..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/optimizer.py +++ /dev/null @@ -1,213 +0,0 @@ -# Copyright 2021 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 torch - -from .state import AcceleratorState, GradientState -from .utils import DistributedType, honor_type, is_lomo_available, is_torch_xla_available - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - import torch_xla.runtime as xr - - -def move_to_device(state, device): - if isinstance(state, (list, tuple)): - return honor_type(state, (move_to_device(t, device) for t in state)) - elif isinstance(state, dict): - return type(state)({k: move_to_device(v, device) for k, v in state.items()}) - elif isinstance(state, torch.Tensor): - return state.to(device) - return state - - -class AcceleratedOptimizer(torch.optim.Optimizer): - """ - Internal wrapper around a torch optimizer. - - Conditionally will perform `step` and `zero_grad` if gradients should be synchronized when performing gradient - accumulation. - - Args: - optimizer (`torch.optim.optimizer.Optimizer`): - The optimizer to wrap. - device_placement (`bool`, *optional*, defaults to `True`): - Whether or not the optimizer should handle device placement. If so, it will place the state dictionary of - `optimizer` on the right device. - scaler (`torch.amp.GradScaler` or `torch.cuda.amp.GradScaler`, *optional*): - The scaler to use in the step function if training with mixed precision. - """ - - def __init__(self, optimizer, device_placement=True, scaler=None): - self.optimizer = optimizer - self.scaler = scaler - self.accelerator_state = AcceleratorState() - self.gradient_state = GradientState() - self.device_placement = device_placement - self._is_overflow = False - - if self.scaler is not None: - self._accelerate_step_called = False - self._optimizer_original_step_method = self.optimizer.step - self._optimizer_patched_step_method = patch_optimizer_step(self, self.optimizer.step) - - # Handle device placement - if device_placement: - state_dict = self.optimizer.state_dict() - if self.accelerator_state.distributed_type == DistributedType.XLA: - xm.send_cpu_data_to_device(state_dict, self.accelerator_state.device) - else: - state_dict = move_to_device(state_dict, self.accelerator_state.device) - self.optimizer.load_state_dict(state_dict) - - @property - def state(self): - return self.optimizer.state - - @state.setter - def state(self, state): - self.optimizer.state = state - - @property - def param_groups(self): - return self.optimizer.param_groups - - @param_groups.setter - def param_groups(self, param_groups): - self.optimizer.param_groups = param_groups - - @property - def defaults(self): - return self.optimizer.defaults - - @defaults.setter - def defaults(self, defaults): - self.optimizer.defaults = defaults - - def add_param_group(self, param_group): - self.optimizer.add_param_group(param_group) - - def load_state_dict(self, state_dict): - if self.accelerator_state.distributed_type == DistributedType.XLA and self.device_placement: - xm.send_cpu_data_to_device(state_dict, self.accelerator_state.device) - self.optimizer.load_state_dict(state_dict) - - def state_dict(self): - return self.optimizer.state_dict() - - def zero_grad(self, set_to_none=None): - if self.gradient_state.sync_gradients: - accept_arg = "set_to_none" in inspect.signature(self.optimizer.zero_grad).parameters - if accept_arg: - if set_to_none is None: - set_to_none = True - self.optimizer.zero_grad(set_to_none=set_to_none) - else: - if set_to_none is not None: - raise ValueError("`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer.") - self.optimizer.zero_grad() - - def train(self): - """ - Sets the optimizer to "train" mode. Useful for optimizers like `schedule_free` - """ - if hasattr(self.optimizer, "train") and callable(self.optimizer.train): - self.optimizer.train() - elif ( - hasattr(self.optimizer, "optimizer") - and hasattr(self.optimizer.optimizer, "train") - and callable(self.optimizer.optimizer.train) - ): - # the deepspeed optimizer further wraps the optimizer - self.optimizer.optimizer.train() - - def eval(self): - """ - Sets the optimizer to "eval" mode. Useful for optimizers like `schedule_free` - """ - if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval): - self.optimizer.eval() - - def step(self, closure=None): - if is_lomo_available(): - from lomo_optim import AdaLomo, Lomo - - if ( - not self.gradient_state.is_xla_gradients_synced - and self.accelerator_state.distributed_type == DistributedType.XLA - ): - gradients = xm._fetch_gradients(self.optimizer) - xm.all_reduce("sum", gradients, scale=1.0 / xr.world_size()) - self.gradient_state.is_xla_gradients_synced = True - - if is_lomo_available(): - # `step` should be a no-op for LOMO optimizers. - if isinstance(self.optimizer, (Lomo, AdaLomo)): - return - - if self.gradient_state.sync_gradients: - if self.scaler is not None: - self.optimizer.step = self._optimizer_patched_step_method - - self.scaler.step(self.optimizer, closure) - self.scaler.update() - - if not self._accelerate_step_called: - # If the optimizer step was skipped, gradient overflow was detected. - self._is_overflow = True - else: - self._is_overflow = False - # Reset the step method to the original one - self.optimizer.step = self._optimizer_original_step_method - # Reset the indicator - self._accelerate_step_called = False - else: - self.optimizer.step(closure) - if self.accelerator_state.distributed_type == DistributedType.XLA: - self.gradient_state.is_xla_gradients_synced = False - - def _switch_parameters(self, parameters_map): - for param_group in self.optimizer.param_groups: - param_group["params"] = [parameters_map.get(p, p) for p in param_group["params"]] - - @property - def step_was_skipped(self): - """Whether or not the optimizer step was skipped.""" - return self._is_overflow - - def __getstate__(self): - _ignored_keys = [ - "_accelerate_step_called", - "_optimizer_original_step_method", - "_optimizer_patched_step_method", - ] - return {k: v for k, v in self.__dict__.items() if k not in _ignored_keys} - - def __setstate__(self, state): - self.__dict__.update(state) - if self.scaler is not None: - self._accelerate_step_called = False - self._optimizer_original_step_method = self.optimizer.step - self._optimizer_patched_step_method = patch_optimizer_step(self, self.optimizer.step) - - -def patch_optimizer_step(accelerated_optimizer: AcceleratedOptimizer, method): - def patched_step(*args, **kwargs): - accelerated_optimizer._accelerate_step_called = True - return method(*args, **kwargs) - - return patched_step diff --git a/venv/lib/python3.10/site-packages/accelerate/scheduler.py b/venv/lib/python3.10/site-packages/accelerate/scheduler.py deleted file mode 100644 index 1fa8a13f238afd7b908ee8e8cb8e0620f48d4ff8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/scheduler.py +++ /dev/null @@ -1,98 +0,0 @@ -# Copyright 2022 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. - -# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation - -import warnings - -from .state import AcceleratorState, GradientState - - -warnings.filterwarnings("ignore", category=UserWarning, module="torch.optim.lr_scheduler") - - -class AcceleratedScheduler: - """ - A wrapper around a learning rate scheduler that will only step when the optimizer(s) have a training step. Useful - to avoid making a scheduler step too fast when gradients went overflow and there was no training step (in mixed - precision training) - - When performing gradient accumulation scheduler lengths should not be changed accordingly, Accelerate will always - step the scheduler to account for it. - - Args: - scheduler (`torch.optim.lr_scheduler._LRScheduler`): - The scheduler to wrap. - optimizers (one or a list of `torch.optim.Optimizer`): - The optimizers used. - step_with_optimizer (`bool`, *optional*, defaults to `True`): - Whether or not the scheduler should be stepped at each optimizer step. - split_batches (`bool`, *optional*, defaults to `False`): - Whether or not the dataloaders split one batch across the different processes (so batch size is the same - regardless of the number of processes) or create batches on each process (so batch size is the original - batch size multiplied by the number of processes). - """ - - def __init__(self, scheduler, optimizers, step_with_optimizer: bool = True, split_batches: bool = False): - self.scheduler = scheduler - self.optimizers = optimizers if isinstance(optimizers, (list, tuple)) else [optimizers] - self.split_batches = split_batches - self.step_with_optimizer = step_with_optimizer - self.gradient_state = GradientState() - - def step(self, *args, **kwargs): - if not self.step_with_optimizer: - # No link between scheduler and optimizer -> just step - self.scheduler.step(*args, **kwargs) - return - - # Otherwise, first make sure the optimizer was stepped. - if not self.gradient_state.sync_gradients: - if self.gradient_state.adjust_scheduler: - self.scheduler._step_count += 1 - return - - for opt in self.optimizers: - if opt.step_was_skipped: - return - if self.split_batches: - # Split batches -> the training dataloader batch size is not changed so one step per training step - self.scheduler.step(*args, **kwargs) - else: - # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do - # num_processes steps per training step - num_processes = AcceleratorState().num_processes - for _ in range(num_processes): - # Special case when using OneCycle and `drop_last` was not used - if hasattr(self.scheduler, "total_steps"): - if self.scheduler._step_count <= self.scheduler.total_steps: - self.scheduler.step(*args, **kwargs) - else: - self.scheduler.step(*args, **kwargs) - - # Passthroughs - def get_last_lr(self): - return self.scheduler.get_last_lr() - - def state_dict(self): - return self.scheduler.state_dict() - - def load_state_dict(self, state_dict): - self.scheduler.load_state_dict(state_dict) - - def get_lr(self): - return self.scheduler.get_lr() - - def print_lr(self, *args, **kwargs): - return self.scheduler.print_lr(*args, **kwargs) diff --git a/venv/lib/python3.10/site-packages/accelerate/state.py b/venv/lib/python3.10/site-packages/accelerate/state.py deleted file mode 100644 index d59305a498a63b27829947357f2a12b02f817a12..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/state.py +++ /dev/null @@ -1,1347 +0,0 @@ -# Copyright 2021 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 logging -import os -import threading -import warnings -import weakref -from contextlib import contextmanager -from functools import partial -from typing import Any, Callable - -import torch - -from .utils import ( - DistributedType, - DynamoBackend, - GradientAccumulationPlugin, - check_cuda_fp8_capability, - check_cuda_p2p_ib_support, - deepspeed_required, - get_cpu_distributed_information, - get_int_from_env, - is_ccl_available, - is_datasets_available, - is_deepspeed_available, - is_fp8_available, - is_habana_gaudi1, - is_hpu_available, - is_ipex_available, - is_mlu_available, - is_mps_available, - is_musa_available, - is_npu_available, - is_sdaa_available, - is_torch_xla_available, - is_xccl_available, - is_xpu_available, - parse_choice_from_env, - parse_flag_from_env, - set_numa_affinity, -) -from .utils.dataclasses import SageMakerDistributedType - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - import torch_xla.runtime as xr - -if is_mlu_available(check_device=False): - import torch_mlu # noqa: F401 - -if is_sdaa_available(check_device=False): - import torch_sdaa # noqa: F401 - -if is_musa_available(check_device=False): - import torch_musa # noqa: F401 - -if is_npu_available(check_device=False): - import torch_npu # noqa: F401 - - -logger = logging.getLogger(__name__) - - -def is_initialized() -> bool: - """ - Checks if the `AcceleratorState` has been initialized from `Accelerator`. Same as `AcceleratorState.initialized`, - but works as a module method. - """ - return AcceleratorState._shared_state != {} - - -# Lambda function that does nothing -def do_nothing(*args, **kwargs): - return None - - -class ThreadLocalSharedDict(threading.local): - """ - Descriptor that holds a dict shared between instances of a class in the same thread. - - Note: Descriptors have slightly different semantics than just a dict field on its own. - `PartialState(...)._shared_state` and `PartialState._shared_state` (instance vs class) give the same value: the - underlying _storage dict. Likewise, `PartialState(...)._shared_state = {...}` overrides the _storage dict inside - the descriptor as you would expect. However, `PartialState._shared_state = {}` actually replaces the descriptor - object with a dict instead Thus, you should modify the _storage dict in-place (e.g. `_shared_state.clear()`). - - See Python documentation for an explanation of descriptors: https://docs.python.org/3/howto/descriptor.html - - This is required for using PyTorch/XLA with PJRT in multithreaded mode (required for TPU v2 and v3). - - See https://github.com/pytorch/xla/blob/r2.0/docs/pjrt.md#multithreading-on-tpu-v2v3 - """ - - def __init__(self, thread_local: bool = False): - self._storage = {} - - def __get__(self, obj, objtype=None): - return self._storage - - def __set__(self, obj, value): - self._storage = value - - -# Prefer global shared dictionary, except when using TPU. -SharedDict = dict if not is_torch_xla_available() else ThreadLocalSharedDict - - -# Inspired by Alex Martelli's 'Borg'. -class PartialState: - """ - Singleton class that has information about the current training environment and functions to help with process - control. Designed to be used when only process control and device execution states are needed. Does *not* need to - be initialized from `Accelerator`. - - Args: - cpu (`bool`, *optional*): - Whether or not to force the script to execute on CPU. Will ignore any accelerators available if set to - `True` and force the execution on the CPU. - kwargs (additional keyword arguments, *optional*): - Additional keyword arguments to pass to the relevant `init_process_group` function. Valid `kwargs` can be - found in [`utils.InitProcessGroupKwargs`]. See the example section for detailed usage. - - **Available attributes:** - - - **device** (`torch.device`) -- The device to use. - - **distributed_type** ([`~accelerate.state.DistributedType`]) -- The type of distributed environment currently - in use. - - **local_process_index** (`int`) -- The index of the current process on the current server. - - **mixed_precision** (`str`) -- Whether or not the current script will use mixed precision, and if so the type - of mixed precision being performed. (Choose from 'no','fp16','bf16 or 'fp8'). - - **num_processes** (`int`) -- The number of processes currently launched in parallel. - - **process_index** (`int`) -- The index of the current process. - - **is_last_process** (`bool`) -- Whether or not the current process is the last one. - - **is_main_process** (`bool`) -- Whether or not the current process is the main one. - - **is_local_main_process** (`bool`) -- Whether or not the current process is the main one on the local node. - - **debug** (`bool`) -- Whether or not the current script is being run in debug mode. - - Example: - ```python - from accelerate.utils import InitProcessGroupKwargs - - # To include `InitProcessGroupKwargs`, init then call `.to_kwargs()` - kwargs = InitProcessGroupKwargs(...).to_kwargs() - state = PartialState(**kwargs) - ``` - """ - - _shared_state = SharedDict() - _known_attrs = [ - "_cpu", - "_mixed_precision", - "_shared_state", - "backend", - "debug", - "device", - "distributed_type", - "fork_launched", - "local_process_index", - "num_processes", - "process_index", - ] - - def __init__(self, cpu: bool = False, **kwargs): - self.__dict__ = self._shared_state - if not self.initialized: - self._cpu = cpu - self.backend = None - env_device = os.environ.get("ACCELERATE_TORCH_DEVICE", None) - self.device = torch.device(env_device) if env_device is not None else None - self.debug = parse_flag_from_env("ACCELERATE_DEBUG_MODE") - use_sagemaker_dp = kwargs.pop("_use_sagemaker_dp", None) - dist_information = None - if use_sagemaker_dp is None: - use_sagemaker_dp = ( - os.environ.get("ACCELERATE_USE_SAGEMAKER", "false") == "true" - and os.environ.get("ACCELERATE_SAGEMAKER_DISTRIBUTED_TYPE") != SageMakerDistributedType.NO - ) - - # Sets up self.backend + imports - original_backend = kwargs.pop("backend", None) - backend, distributed_type = self._prepare_backend(cpu, use_sagemaker_dp, original_backend) - if original_backend is not None and backend != original_backend: - raise ValueError(f"Your assigned backend {original_backend} is not avaliable, please use {backend}") - self.backend = backend - self.distributed_type = distributed_type - use_deepspeed = False - if not cpu and self.backend != "xla": - if int(os.environ.get("LOCAL_RANK", -1)) != -1: - # Deal with spawning deepspeed - if os.environ.get("ACCELERATE_USE_DEEPSPEED", "false") == "true": - if not is_deepspeed_available(): - raise ImportError( - "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source" - ) - from deepspeed import comm as dist - - if not dist.is_initialized(): - if self.backend == "tccl": - local_rank = os.environ.get("LOCAL_RANK", -1) - torch.sdaa.set_device(f"sdaa:{local_rank}") - dist.init_distributed(dist_backend=self.backend, auto_mpi_discovery=False, **kwargs) - # We need to flag to `use_deepspeed` to be True to override `distributed_type` later - use_deepspeed = True - # Deal with all other backends but XPU and CPU, that gets handled special later - elif ( - self.distributed_type not in (DistributedType.MULTI_XPU, DistributedType.MULTI_CPU) - and not torch.distributed.is_initialized() - ): - if self.backend == "tccl": - local_rank = os.environ.get("LOCAL_RANK", -1) - torch.sdaa.set_device(f"sdaa:{local_rank}") - if ( - self.backend == "nccl" - and os.environ.get("ACCELERATE_USE_FSDP", "false") == "true" - and ( - os.environ.get("FSDP_OFFLOAD_PARAMS", "false") == "true" - or os.environ.get("FSDP_STATE_DICT_TYPE", "SHARDED_STATE_DICT") == "FULL_STATE_DICT" - ) - ): - self.backend = "cuda:nccl,cpu:gloo" - torch.distributed.init_process_group(backend=self.backend, **kwargs) - - # XPU and CPU require special env configs to be set - if self.distributed_type in (DistributedType.MULTI_XPU, DistributedType.MULTI_CPU): - dist_information = get_cpu_distributed_information() - os.environ["RANK"] = str(dist_information.rank) - os.environ["WORLD_SIZE"] = str(dist_information.world_size) - os.environ["LOCAL_RANK"] = str(dist_information.local_rank) - os.environ["LOCAL_WORLD_SIZE"] = str(dist_information.local_world_size) - if not os.environ.get("MASTER_PORT", None): - os.environ["MASTER_PORT"] = "29500" - if ( - not os.environ.get("MASTER_ADDR", None) - and dist_information.local_world_size != dist_information.world_size - and self.backend != "mpi" - ): - raise ValueError( - "Tried to launch on distributed with multinode, but `MASTER_ADDR` env was not set, " - "please try exporting rank 0's hostname as `MASTER_ADDR`" - ) - kwargs["rank"] = dist_information.rank - kwargs["world_size"] = dist_information.world_size - - if ( - self.distributed_type == DistributedType.MULTI_CPU - and get_int_from_env(["OMP_NUM_THREADS"], 0) == 0 - ): - import psutil - - num_cpu_threads_per_process = int( - psutil.cpu_count(logical=False) / dist_information.local_world_size - ) - if num_cpu_threads_per_process == 0: - num_cpu_threads_per_process = 1 - torch.set_num_threads(num_cpu_threads_per_process) - warnings.warn( - f"OMP_NUM_THREADS/MKL_NUM_THREADS unset, we set it at {num_cpu_threads_per_process} to improve oob" - " performance." - ) - - if not torch.distributed.is_initialized(): - torch.distributed.init_process_group(backend=self.backend, **kwargs) - - # No backend == no distributed training - if self.backend is None: - self.distributed_type = DistributedType.NO - self.num_processes = 1 - self.process_index = 0 - self.local_process_index = 0 - elif self.backend == "xla": - # XLA needs device setting first for `set_replication` - self.set_device() - xm.set_replication(self.device, xm.get_xla_supported_devices()) - self.num_processes = xr.world_size() - self.process_index = xr.global_ordinal() - if is_torch_xla_available(check_is_tpu=True): - self.local_process_index = xm.get_local_ordinal() - else: - self.local_process_index = int(os.environ.get("LOCAL_RANK", -1)) - else: - self.num_processes = torch.distributed.get_world_size() - self.process_index = torch.distributed.get_rank() - self.local_process_index = ( - int(os.environ.get("LOCAL_RANK", -1)) if dist_information is None else dist_information.local_rank - ) - self.set_device() - # Now we can change to deepseed - if use_deepspeed: - self.distributed_type = DistributedType.DEEPSPEED - - # Set CPU affinity if enabled - if parse_flag_from_env("ACCELERATE_CPU_AFFINITY", False): - set_numa_affinity(self.local_process_index) - - # Check for old RTX 4000's that can't use P2P or IB and are on old drivers - if self.device.type == "cuda" and not check_cuda_p2p_ib_support(): - if "NCCL_P2P_DISABLE" not in os.environ or "NCCL_IB_DISABLE" not in os.environ: - raise NotImplementedError( - "Using RTX 4000 series doesn't support faster communication broadband via P2P or IB. " - 'Please set `NCCL_P2P_DISABLE="1"` and `NCCL_IB_DISABLE="1" or use `accelerate launch` which ' - "will do this automatically." - ) - - # Important: This should be the *only* code outside of `self.initialized!` - self.fork_launched = parse_flag_from_env("FORK_LAUNCHED", 0) - - def __repr__(self) -> str: - return ( - f"Distributed environment: {self.distributed_type}{(' Backend: ' + self.backend) if self.backend else ''}\n" - f"Num processes: {self.num_processes}\n" - f"Process index: {self.process_index}\n" - f"Local process index: {self.local_process_index}\n" - f"Device: {self.device}\n" - ) - - @staticmethod - def _reset_state(): - "Resets `_shared_state`, is used internally and should not be called" - PartialState._shared_state.clear() - - @property - def initialized(self) -> bool: - "Returns whether the `PartialState` has been initialized" - return self._shared_state != {} - - @property - def use_distributed(self): - """ - Whether the Accelerator is configured for distributed training - """ - return self.distributed_type != DistributedType.NO and self.num_processes > 1 - - @property - def is_last_process(self) -> bool: - "Returns whether the current process is the last one" - return self.process_index == self.num_processes - 1 - - @property - def is_main_process(self) -> bool: - "Returns whether the current process is the main process" - return ( - self.process_index == 0 if self.distributed_type != DistributedType.MEGATRON_LM else self.is_last_process - ) - - @property - def is_local_main_process(self) -> bool: - "Returns whether the current process is the main process on the local node" - return ( - self.local_process_index == 0 - if self.distributed_type != DistributedType.MEGATRON_LM - else self.is_last_process - ) - - def wait_for_everyone(self): - """ - Will stop the execution of the current process until every other process has reached that point (so this does - nothing when the script is only run in one process). Useful to do before saving a model. - - Example: - - ```python - >>> # Assuming two GPU processes - >>> import time - >>> from accelerate.state import PartialState - - >>> state = PartialState() - >>> if state.is_main_process: - ... time.sleep(2) - >>> else: - ... print("I'm waiting for the main process to finish its sleep...") - >>> state.wait_for_everyone() - >>> # Should print on every process at the same time - >>> print("Everyone is here") - ``` - """ - if self.distributed_type in ( - DistributedType.MULTI_GPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_NPU, - DistributedType.MULTI_XPU, - DistributedType.MULTI_CPU, - DistributedType.MULTI_HPU, - DistributedType.DEEPSPEED, - DistributedType.FSDP, - ): - torch.distributed.barrier() - elif self.distributed_type == DistributedType.XLA: - xm.rendezvous("accelerate.utils.wait_for_everyone") - - def _goes_first(self, is_main: bool): - if not is_main: - self.wait_for_everyone() - - yield - - if is_main: - self.wait_for_everyone() - - @contextmanager - def split_between_processes(self, inputs: list | tuple | dict | torch.Tensor, apply_padding: bool = False): - """ - Splits `input` between `self.num_processes` quickly and can be then used on that process. Useful when doing - distributed inference, such as with different prompts. - - Note that when using a `dict`, all keys need to have the same number of elements. - - Args: - inputs (`list`, `tuple`, `torch.Tensor`, `dict` of `list`/`tuple`/`torch.Tensor`, or `datasets.Dataset`): - The input to split between processes. - apply_padding (`bool`, `optional`, defaults to `False`): - Whether to apply padding by repeating the last element of the input so that all processes have the same - number of elements. Useful when trying to perform actions such as `gather()` on the outputs or passing - in less inputs than there are processes. If so, just remember to drop the padded elements afterwards. - - - Example: - - ```python - # Assume there are two processes - from accelerate import PartialState - - state = PartialState() - with state.split_between_processes(["A", "B", "C"]) as inputs: - print(inputs) - # Process 0 - ["A", "B"] - # Process 1 - ["C"] - - with state.split_between_processes(["A", "B", "C"], apply_padding=True) as inputs: - print(inputs) - # Process 0 - ["A", "B"] - # Process 1 - ["C", "C"] - ``` - """ - if self.num_processes == 1: - yield inputs - return - length = len(inputs) - # Nested dictionary of any types - if isinstance(inputs, dict): - length = len(inputs[list(inputs.keys())[0]]) - if not all(len(v) == length for v in inputs.values()): - raise ValueError("All values in the dictionary must have the same length") - num_samples_per_process, num_extras = divmod(length, self.num_processes) - start_index = self.process_index * num_samples_per_process + min(self.process_index, num_extras) - end_index = start_index + num_samples_per_process + (1 if self.process_index < num_extras else 0) - - def _split_values(inputs, start_index, end_index): - if isinstance(inputs, (list, tuple, torch.Tensor)): - if start_index >= len(inputs): - result = inputs[-1:] - else: - result = inputs[start_index:end_index] - if apply_padding: - if isinstance(result, torch.Tensor): - from accelerate.utils import pad_across_processes, send_to_device - - # The tensor needs to be on the device before we can pad it - tensorized_result = send_to_device(result, self.device) - result = pad_across_processes(tensorized_result, pad_index=inputs[-1]) - else: - result += [result[-1]] * (num_samples_per_process + (1 if num_extras > 0 else 0) - len(result)) - return result - elif isinstance(inputs, dict): - for key in inputs.keys(): - inputs[key] = _split_values(inputs[key], start_index, end_index) - return inputs - else: - if is_datasets_available(): - from datasets import Dataset - - if isinstance(inputs, Dataset): - if start_index >= len(inputs): - start_index = len(inputs) - 1 - if end_index > len(inputs): - end_index = len(inputs) - result_idcs = list(range(start_index, end_index)) - if apply_padding: - result_idcs += [end_index - 1] * ( - num_samples_per_process + (1 if num_extras > 0 else 0) - len(result_idcs) - ) - return inputs.select(result_idcs) - return inputs - - yield _split_values(inputs, start_index, end_index) - - @contextmanager - def main_process_first(self): - """ - Lets the main process go first inside a with block. - - The other processes will enter the with block after the main process exits. - - Example: - - ```python - >>> from accelerate import Accelerator - - >>> accelerator = Accelerator() - >>> with accelerator.main_process_first(): - ... # This will be printed first by process 0 then in a seemingly - ... # random order by the other processes. - ... print(f"This will be printed by process {accelerator.process_index}") - ``` - """ - yield from self._goes_first(self.is_main_process) - - @contextmanager - def local_main_process_first(self): - """ - Lets the local main process go inside a with block. - - The other processes will enter the with block after the main process exits. - - Example: - - ```python - >>> from accelerate.state import PartialState - - >>> state = PartialState() - >>> with state.local_main_process_first(): - ... # This will be printed first by local process 0 then in a seemingly - ... # random order by the other processes. - ... print(f"This will be printed by process {state.local_process_index}") - ``` - """ - yield from self._goes_first(self.is_local_main_process) - - def on_main_process(self, function: Callable[..., Any] = None): - """ - Decorator that only runs the decorated function on the main process. - - Args: - function (`Callable`): The function to decorate. - - Example: - - ```python - >>> from accelerate.state import PartialState - - >>> state = PartialState() - - - >>> @state.on_main_process - ... def print_something(): - ... print("This will be printed by process 0 only.") - - - >>> print_something() - "This will be printed by process 0 only" - ``` - """ - if not self.initialized: - raise ValueError("The `PartialState` or `Accelerator` must be initialized before calling this function.") - if self.is_main_process or not self.use_distributed: - return function - return do_nothing - - def on_local_main_process(self, function: Callable[..., Any] = None): - """ - Decorator that only runs the decorated function on the local main process. - - Args: - function (`Callable`): The function to decorate. - - Example: - ```python - # Assume we have 2 servers with 4 processes each. - from accelerate.state import PartialState - - state = PartialState() - - - @state.on_local_main_process - def print_something(): - print("This will be printed by process 0 only on each server.") - - - print_something() - # On server 1: - "This will be printed by process 0 only" - # On server 2: - "This will be printed by process 0 only" - ``` - """ - if self.is_local_main_process or not self.use_distributed: - return function - return do_nothing - - def on_last_process(self, function: Callable[..., Any]): - """ - Decorator that only runs the decorated function on the last process. - - Args: - function (`Callable`): The function to decorate. - - Example: - ```python - # Assume we have 4 processes. - from accelerate.state import PartialState - - state = PartialState() - - - @state.on_last_process - def print_something(): - print(f"Printed on process {state.process_index}") - - - print_something() - "Printed on process 3" - ``` - """ - if self.is_last_process or not self.use_distributed: - return function - return do_nothing - - def on_process(self, function: Callable[..., Any] = None, process_index: int = None): - """ - Decorator that only runs the decorated function on the process with the given index. - - Args: - function (`Callable`, `optional`): - The function to decorate. - process_index (`int`, `optional`): - The index of the process on which to run the function. - - Example: - ```python - # Assume we have 4 processes. - from accelerate.state import PartialState - - state = PartialState() - - - @state.on_process(process_index=2) - def print_something(): - print(f"Printed on process {state.process_index}") - - - print_something() - "Printed on process 2" - ``` - """ - if function is None: - return partial(self.on_process, process_index=process_index) - if (self.process_index == process_index) or (not self.use_distributed): - return function - return do_nothing - - def on_local_process(self, function: Callable[..., Any] = None, local_process_index: int = None): - """ - Decorator that only runs the decorated function on the process with the given index on the current node. - - Args: - function (`Callable`, *optional*): - The function to decorate. - local_process_index (`int`, *optional*): - The index of the local process on which to run the function. - - Example: - ```python - # Assume we have 2 servers with 4 processes each. - from accelerate import Accelerator - - accelerator = Accelerator() - - - @accelerator.on_local_process(local_process_index=2) - def print_something(): - print(f"Printed on process {accelerator.local_process_index}") - - - print_something() - # On server 1: - "Printed on process 2" - # On server 2: - "Printed on process 2" - ``` - """ - if function is None: - return partial(self.on_local_process, local_process_index=local_process_index) - if (self.local_process_index == local_process_index) or (not self.use_distributed): - return function - return do_nothing - - def print(self, *args, **kwargs): - if self.is_local_main_process: - print(*args, **kwargs) - - @property - def default_device(self) -> torch.device: - """ - Returns the default device which is: - - MPS if `torch.backends.mps.is_available()` and `torch.backends.mps.is_built()` both return True. - - CUDA if `torch.cuda.is_available()` - - MLU if `is_mlu_available()` - - SDAA if `is_sdaa_available()` - - MUSA if `is_musa_available()` - - NPU if `is_npu_available()` - - HPU if `is_hpu_available()` - - CPU otherwise - """ - if is_mps_available(): - os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" - return torch.device("mps") - elif is_mlu_available(): - return torch.device("mlu") - elif is_sdaa_available(): - return torch.device("sdaa") - elif is_musa_available(): - return torch.device("musa") - # NPU should be checked before CUDA when using `transfer_to_npu` - # See issue #3020: https://github.com/huggingface/accelerate/issues/3020 - elif is_npu_available(): - return torch.device("npu") - elif is_hpu_available(): - return torch.device("hpu") - elif torch.cuda.is_available(): - return torch.device("cuda") - elif is_xpu_available(): - return torch.device("xpu") - else: - return torch.device("cpu") - - def _prepare_backend( - self, cpu: bool = False, sagemaker_dp=False, backend: str = None - ) -> tuple[str, DistributedType]: - "Prepares any imports needed before initializing the distributed backend and sets `self.backend` properly" - distributed_type = None - if sagemaker_dp: - import smdistributed.dataparallel.torch.torch_smddp # noqa - - backend = "smddp" - distributed_type = DistributedType.MULTI_GPU - elif is_torch_xla_available(): - backend = "xla" - distributed_type = DistributedType.XLA - - elif int(os.environ.get("LOCAL_RANK", -1)) != -1 and not cpu: - if is_mlu_available(): - backend = "cncl" - distributed_type = DistributedType.MULTI_MLU - if is_sdaa_available(): - backend = "tccl" - distributed_type = DistributedType.MULTI_SDAA - elif is_musa_available(): - backend = "mccl" - distributed_type = DistributedType.MULTI_MUSA - # NPU should be checked before CUDA when using `transfer_to_npu` - # See issue #3020: https://github.com/huggingface/accelerate/issues/3020 - elif is_npu_available(): - backend = "hccl" - distributed_type = DistributedType.MULTI_NPU - elif is_hpu_available(init_hccl=True): - if backend is None: - backend = "hccl" - distributed_type = DistributedType.MULTI_HPU - elif torch.cuda.is_available(): - if backend is None: - backend = "nccl" - distributed_type = DistributedType.MULTI_GPU - elif is_xpu_available() and is_xccl_available(): - if backend is None: - backend = "xccl" - distributed_type = DistributedType.MULTI_XPU - - if distributed_type is None and ( - int(os.environ.get("LOCAL_RANK", -1)) != -1 - or get_int_from_env(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"], 1) > 1 - ): - if not cpu and is_xpu_available(): - distributed_type = DistributedType.MULTI_XPU - else: - distributed_type = DistributedType.MULTI_CPU - - if ( - backend in (None, "ccl") - and is_ccl_available() - and (get_int_from_env(["CCL_WORKER_COUNT"], 0) > 0 or distributed_type == DistributedType.MULTI_XPU) - ): - import oneccl_bindings_for_pytorch # noqa: F401 - - backend = "ccl" - elif backend in (None, "mpi") and torch.distributed.is_mpi_available(): - backend = "mpi" - else: - backend = "gloo" - if distributed_type is None: - distributed_type = DistributedType.NO - - return backend, distributed_type - - def set_device(self): - """ - Sets the device in `self.device` to the current distributed environment. - """ - if self.device is not None: - return - if self.distributed_type == DistributedType.NO: - self.device = torch.device("cpu") if self._cpu else self.default_device - return - device = str(self.distributed_type).split(".")[-1].replace("MULTI_", "").lower() - if device not in ("cpu", "gpu", "mlu", "musa", "npu", "xpu", "xla", "hpu", "sdaa"): - raise ValueError( - f"Can't set device for {self.distributed_type} ({device}), verify we should be calling `_set_device()` for it!" - ) - if device == "xla": - self.device = xm.xla_device() - elif device == "hpu": - self.device = torch.device("hpu", torch.hpu.current_device()) - else: - if device == "gpu": - device = "cuda" - device_module = getattr(torch, device) - device_index = self.local_process_index % device_module.device_count() - self.device = torch.device(device, device_index) - device_module.set_device(self.device) - - def destroy_process_group(self, group=None): - """ - Destroys the process group. If one is not specified, the default process group is destroyed. - """ - if self.fork_launched and group is None: - return - # needed when using torch.distributed.init_process_group - if torch.distributed.is_initialized(): - torch.distributed.destroy_process_group(group) - - def __getattr__(self, name: str): - # By this point we know that no attributes of `self` contain `name`, - # so we just modify the error message - if name in self._known_attrs: - raise AttributeError( - f"`PartialState` object has no attribute `{name}`. " - "This happens if `PartialState._reset_state()` was called and " - "an `Accelerator` or `PartialState` was not reinitialized." - ) - # Raise a typical AttributeError - raise AttributeError(f"'PartialState' object has no attribute '{name}'") - - -class AcceleratorState: - """ - Singleton class that has information about the current training environment. - - **Available attributes:** - - - **device** (`torch.device`) -- The device to use. - - **distributed_type** ([`~accelerate.state.DistributedType`]) -- The type of distributed environment currently - in use. - - **initialized** (`bool`) -- Whether or not the `AcceleratorState` has been initialized from `Accelerator`. - - **local_process_index** (`int`) -- The index of the current process on the current server. - - **mixed_precision** (`str`) -- Whether or not the current script will use mixed precision, and if so the type - of mixed precision being performed. (Choose from 'no','fp16','bf16 or 'fp8'). - - **num_processes** (`int`) -- The number of processes currently launched in parallel. - - **process_index** (`int`) -- The index of the current process. - - **is_last_process** (`bool`) -- Whether or not the current process is the last one. - - **is_main_process** (`bool`) -- Whether or not the current process is the main one. - - **is_local_main_process** (`bool`) -- Whether or not the current process is the main one on the local node. - - **debug** (`bool`) -- Whether or not the current script is being run in debug mode. - """ - - _shared_state = SharedDict() - _known_attrs = PartialState._known_attrs + [ - "deepspeed_plugin", - "use_ipex", - "fsdp_plugin", - "megatron_lm_plugin", - "dynamo_plugin", - ] - - def __init__( - self, - mixed_precision: str = None, - cpu: bool = False, - dynamo_plugin=None, - deepspeed_plugin=None, - fsdp_plugin=None, - torch_tp_plugin=None, - megatron_lm_plugin=None, - _from_accelerator: bool = False, - **kwargs, - ): - self.__dict__ = self._shared_state - if parse_flag_from_env("ACCELERATE_USE_CPU"): - cpu = True - if PartialState._shared_state == {}: - PartialState(cpu, **kwargs) - self.__dict__.update(PartialState._shared_state) - self._check_initialized(mixed_precision, cpu) - if not self.initialized: - self.deepspeed_plugins = None - self.use_ipex = None - self.torch_tp_plugin = torch_tp_plugin - mixed_precision = ( - parse_choice_from_env("ACCELERATE_MIXED_PRECISION", "no") - if mixed_precision is None - else mixed_precision.lower() - ) - if mixed_precision == "fp8": - # this is confusing, why is is_fp8_available only checks for library availability ? - if not is_fp8_available(): - raise ValueError( - "Using `fp8` precision requires `transformer_engine` or `MS-AMP` to be installed." - ) - elif torch.cuda.is_available() and not check_cuda_fp8_capability(): - logger.warning( - f"The current device has compute capability of {torch.cuda.get_device_capability()} which is " - "insufficient for FP8 mixed precision training (requires a GPU Hopper/Ada Lovelace " - "or higher, compute capability of 8.9 or higher). Will use FP16 instead." - ) - mixed_precision = "fp16" - elif is_habana_gaudi1(): - logger.warning( - "The current HPU device is Gaudi1 which does not support FP8 mixed precision training (requires " - "Gaudi2 or higher). Will use BF16 instead." - ) - mixed_precision = "bf16" - - self.dynamo_plugin = dynamo_plugin - if not _from_accelerator: - raise ValueError( - "Please make sure to properly initialize your accelerator via `accelerator = Accelerator()` " - "before using any functionality from the `accelerate` library." - ) - # deepspeed handles mixed_precision using deepspeed_config - self._mixed_precision = "no" if self.distributed_type == DistributedType.DEEPSPEED else mixed_precision - if self.distributed_type == DistributedType.XLA and is_torch_xla_available(check_is_tpu=True): - if mixed_precision == "bf16": - if os.environ.get("ACCELERATE_DOWNCAST_BF16"): - os.environ["XLA_USE_BF16"] = str(0) - os.environ["XLA_DOWNCAST_BF16"] = str(1) - self.downcast_bfloat = True - else: - os.environ["XLA_USE_BF16"] = str(1) - os.environ["XLA_DOWNCAST_BF16"] = str(0) - self.downcast_bfloat = False - elif os.environ.get("ACCELERATE_USE_DEEPSPEED", "false") == "true" and not cpu: - self.distributed_type = DistributedType.DEEPSPEED - if not isinstance(deepspeed_plugin, dict): - deepspeed_plugin.set_mixed_precision(mixed_precision) - deepspeed_plugin.select(_from_accelerator_state=True) - else: - for plugin in deepspeed_plugin.values(): - plugin.set_mixed_precision(mixed_precision) - # The first plugin passed in is always the active one - first_plugin = next(iter(deepspeed_plugin.values())) - first_plugin.select(_from_accelerator_state=True) - self.deepspeed_plugins = deepspeed_plugin - elif self.distributed_type in [ - DistributedType.MULTI_GPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_NPU, - DistributedType.MULTI_XPU, - DistributedType.MULTI_HPU, - ]: - if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true" or fsdp_plugin is not None: - self.distributed_type = DistributedType.FSDP - if self._mixed_precision != "no": - fsdp_plugin.set_mixed_precision(self._mixed_precision) - self.fsdp_plugin = fsdp_plugin - if os.environ.get("ACCELERATE_USE_MEGATRON_LM", "false") == "true" and self.distributed_type not in [ - DistributedType.MULTI_XPU, - ]: - self.distributed_type = DistributedType.MEGATRON_LM - megatron_lm_plugin.set_mixed_precision(self._mixed_precision) - self.megatron_lm_plugin = megatron_lm_plugin - if self.torch_tp_plugin is not None: - self.distributed_type = DistributedType.TP - elif self.distributed_type in [DistributedType.MULTI_CPU, DistributedType.MULTI_XPU, DistributedType.NO]: - if is_ipex_available(): - # check if user disables it explicitly - self.use_ipex = parse_flag_from_env("ACCELERATE_USE_IPEX", default=True) - else: - self.use_ipex = False - if ( - self.dynamo_plugin.backend != DynamoBackend.NO - and self._mixed_precision == "no" - and self.device.type == "cuda" - ): - torch.backends.cuda.matmul.allow_tf32 = True - if ( - self.dynamo_plugin.backend != DynamoBackend.NO - and self._mixed_precision == "no" - and self.device.type == "musa" - ): - torch.backends.musa.matmul.allow_tf32 = True - PartialState._shared_state["distributed_type"] = self.distributed_type - - @property - def initialized(self) -> bool: - return self._shared_state != PartialState._shared_state - - def __repr__(self): - repr = PartialState().__repr__() + f"\nMixed precision type: {self.mixed_precision}\n" - if self.distributed_type == DistributedType.DEEPSPEED: - repr += f"ds_config: {self.deepspeed_plugin.deepspeed_config}\n" - return repr - - def _check_initialized(self, mixed_precision=None, cpu=None): - "Checks if a modification is trying to be made and the `AcceleratorState` has already been initialized" - if self.initialized: - err = "AcceleratorState has already been initialized and cannot be changed, restart your runtime completely and pass `{flag}` to `Accelerator()`." - if cpu and self.device.type != "cpu": - raise ValueError(err.format(flag="cpu=True")) - if ( - mixed_precision is not None - and mixed_precision != self._mixed_precision - and self.distributed_type != DistributedType.DEEPSPEED - ): - raise ValueError(err.format(flag=f"mixed_precision='{mixed_precision}'")) - - @property - def mixed_precision(self): - if self.distributed_type == DistributedType.DEEPSPEED: - config = self.deepspeed_plugin.deepspeed_config - if config.get("fp16", {}).get("enabled", False): - mixed_precision = "fp16" - elif config.get("bf16", {}).get("enabled", False): - mixed_precision = "bf16" - else: - mixed_precision = "no" - else: - mixed_precision = self._mixed_precision - return mixed_precision - - @staticmethod - def _reset_state(reset_partial_state: bool = False): - "Resets `_shared_state`, is used internally and should not be called" - AcceleratorState._shared_state.clear() - if reset_partial_state: - PartialState._reset_state() - - def destroy_process_group(self, group=None): - """ - Destroys the process group. If one is not specified, the default process group is destroyed. - - If `self.fork_lauched` is `True` and `group` is `None`, nothing happens. - """ - PartialState().destroy_process_group(group) - - @property - def fork_launched(self): - return PartialState().fork_launched - - @property - def use_distributed(self): - """ - Whether the Accelerator is configured for distributed training - """ - return PartialState().use_distributed - - @property - def is_fsdp2(self) -> bool: - return self.distributed_type == DistributedType.FSDP and self.fsdp_plugin.fsdp_version == 2 - - @property - def is_last_process(self) -> bool: - "Returns whether the current process is the last one" - return PartialState().is_last_process - - @property - def is_main_process(self) -> bool: - "Returns whether the current process is the main process" - return PartialState().is_main_process - - @property - def is_local_main_process(self) -> bool: - "Returns whether the current process is the main process on the local node" - return PartialState().is_local_main_process - - def wait_for_everyone(self): - PartialState().wait_for_everyone() - - @contextmanager - def split_between_processes(self, inputs: list | tuple | dict | torch.Tensor, apply_padding: bool = False): - """ - Splits `input` between `self.num_processes` quickly and can be then used on that process. Useful when doing - distributed inference, such as with different prompts. - - Note that when using a `dict`, all keys need to have the same number of elements. - - Args: - inputs (`list`, `tuple`, `torch.Tensor`, or `dict` of `list`/`tuple`/`torch.Tensor`): - The input to split between processes. - apply_padding (`bool`, `optional`, defaults to `False`): - Whether to apply padding by repeating the last element of the input so that all processes have the same - number of elements. Useful when trying to perform actions such as `gather()` on the outputs or passing - in less inputs than there are processes. If so, just remember to drop the padded elements afterwards. - - - Example: - - ```python - # Assume there are two processes - from accelerate.state import AcceleratorState - - state = AcceleratorState() - with state.split_between_processes(["A", "B", "C"]) as inputs: - print(inputs) - # Process 0 - ["A", "B"] - # Process 1 - ["C"] - - with state.split_between_processes(["A", "B", "C"], apply_padding=True) as inputs: - print(inputs) - # Process 0 - ["A", "B"] - # Process 1 - ["C", "C"] - ``` - """ - with PartialState().split_between_processes(inputs, apply_padding=apply_padding) as inputs: - yield inputs - - @contextmanager - def main_process_first(self): - """ - Lets the main process go first inside a with block. - - The other processes will enter the with block after the main process exits. - """ - with PartialState().main_process_first(): - yield - - @contextmanager - def local_main_process_first(self): - """ - Lets the local main process go inside a with block. - - The other processes will enter the with block after the main process exits. - """ - with PartialState().local_main_process_first(): - yield - - @property - def deepspeed_plugin(self): - """ - Returns the currently active DeepSpeedPlugin. - - If not using deepspeed, returns `None`. - """ - # To maintain original behavior, return None if not using deepspeed. - if self.distributed_type != DistributedType.DEEPSPEED: - return None - from accelerate.utils.deepspeed import get_active_deepspeed_plugin - - return get_active_deepspeed_plugin(self) - - @deepspeed_required - def get_deepspeed_plugin(self, name: str): - """ - Returns the DeepSpeedPlugin with the given plugin_key. - """ - return self.deepspeed_plugins[name] - - @deepspeed_required - def select_deepspeed_plugin(self, name: str = None): - """ - Activates the DeepSpeedPlugin with the given `name`, and will disable all other plugins. - """ - for key, plugin in self.deepspeed_plugins.items(): - if key != name: - plugin._unselect() - self.deepspeed_plugins[name].select(_from_accelerator_state=True) - - def print(self, *args, **kwargs): - PartialState().print(*args, **kwargs) - - def __getattr__(self, name: str): - # By this point we know that no attributes of `self` contain `name`, - # so we just modify the error message - if name in self._known_attrs: - raise AttributeError( - f"`AcceleratorState` object has no attribute `{name}`. " - "This happens if `AcceleratorState._reset_state()` was called and " - "an `Accelerator` or `PartialState` was not reinitialized." - ) - # Raise a typical AttributeError - raise AttributeError(f"'AcceleratorState' object has no attribute '{name}'") - - -class GradientState: - """ - Singleton class that has information related to gradient synchronization for gradient accumulation - - **Available attributes:** - - - **end_of_dataloader** (`bool`) -- Whether we have reached the end the current dataloader - - **remainder** (`int`) -- The number of extra samples that were added from padding the dataloader - - **sync_gradients** (`bool`) -- Whether the gradients should be synced across all devices - - **active_dataloader** (`Optional[DataLoader]`) -- The dataloader that is currently being iterated over - - **dataloader_references** (`List[Optional[DataLoader]]`) -- A list of references to the dataloaders that are - being iterated over - - **num_steps** (`int`) -- The number of steps to accumulate over - - **adjust_scheduler** (`bool`) -- Whether the scheduler should be adjusted to account for the gradient - accumulation - - **sync_with_dataloader** (`bool`) -- Whether the gradients should be synced at the end of the dataloader - iteration and the number of total steps reset - - **is_xla_gradients_synced** (`bool`) -- Whether the XLA gradients have been synchronized. It is initialized - as false. Once gradients have been reduced before the optimizer step, this flag is set to true. Subsequently, - after each step, the flag is reset to false. FSDP will always synchronize the gradients, hence - is_xla_gradients_synced is always true. - """ - - _shared_state = SharedDict() - - def __init__(self, gradient_accumulation_plugin: GradientAccumulationPlugin | None = None): - self.__dict__ = self._shared_state - if not self.initialized: - self.sync_gradients = True - self._dataloader_references_ref = [None] - self.plugin_kwargs = ( - gradient_accumulation_plugin.to_kwargs() if gradient_accumulation_plugin is not None else {} - ) - self._is_xla_gradients_synced = False - - # Plugin args are different and can be updated - if gradient_accumulation_plugin is not None and self.plugin_kwargs != gradient_accumulation_plugin.to_kwargs(): - self.plugin_kwargs = gradient_accumulation_plugin.to_kwargs() - - @property - def num_steps(self) -> int: - "Returns the number of steps to accumulate over" - return self.plugin_kwargs.get("num_steps", 1) - - @property - def adjust_scheduler(self) -> bool: - "Returns whether the scheduler should be adjusted" - return self.plugin_kwargs.get("adjust_scheduler", False) - - @property - def sync_with_dataloader(self) -> bool: - "Returns whether the gradients should be synced at the end of the dataloader iteration and the number of total steps reset" - return self.plugin_kwargs.get("sync_with_dataloader", True) - - @property - def initialized(self) -> bool: - "Returns whether the `GradientState` has been initialized" - return GradientState._shared_state != {} - - @property - def end_of_dataloader(self) -> bool: - "Returns whether we have reached the end of the current dataloader" - if not self.in_dataloader: - return False - return self.active_dataloader.end_of_dataloader - - @property - def remainder(self) -> int: - "Returns the number of extra samples that were added from padding the dataloader" - if not self.in_dataloader: - return -1 - return self.active_dataloader.remainder - - def __repr__(self): - return ( - f"Sync Gradients: {self.sync_gradients}\n" - f"At end of current dataloader: {self.end_of_dataloader}\n" - f"Extra samples added: {self.remainder}\n" - f"Gradient accumulation plugin: {self.plugin_kwargs}\n" - ) - - @property - def is_xla_gradients_synced(self): - "Returns the value of is_xla_gradients_synced. FSDP will always synchronize the gradients, hence is_xla_gradients_synced is always true." - if parse_flag_from_env("ACCELERATE_USE_FSDP", default=False): - return True - return self._is_xla_gradients_synced - - @is_xla_gradients_synced.setter - def is_xla_gradients_synced(self, is_synced): - "Set the _is_xla_gradients_synced attribute." - self._is_xla_gradients_synced = is_synced - - def _set_sync_gradients(self, sync_gradients): - "Private function that sets whether gradients should be synchronized. Users should not have to call this." - self.sync_gradients = sync_gradients - # Allow grad-sync to automatically work on TPUs - if ( - self.sync_gradients - and is_torch_xla_available(check_is_tpu=True) - and PartialState().distributed_type == DistributedType.XLA - ): - xm.mark_step() - - def _add_dataloader(self, dataloader): - "Private function that adds a dataloader to `self.dataloader_references` and sets `in_dataloader` to `True`. Users should not have to call this." - # We explicitly use assignment to ensure that the property setter is triggered, which is required for garbage collection. - # Avoid using self.dataloader_references.append as it will not trigger the setter. - self.dataloader_references += [dataloader] - - def _remove_dataloader(self, dataloader): - "Private function that removes a dataloader from `self.dataloader_references` and sets `in_dataloader` to `False` if there are no more dataloaders. Users should not have to call this." - # We explicitly use assignment to ensure that the property setter is triggered. - self.dataloader_references = [ - dataloader_ref for dataloader_ref in self.dataloader_references if dataloader_ref != dataloader - ] - - @property - def active_dataloader(self): - return self.dataloader_references[-1] - - @property - def dataloader_references(self): - # We use a property getter and setter with weakrefs to avoid circular references that prevent garbage collection - return [reference() if reference is not None else reference for reference in self._dataloader_references_ref] - - @dataloader_references.setter - def dataloader_references(self, references): - self._dataloader_references_ref = [ - weakref.ref(dataloader) if dataloader is not None else dataloader for dataloader in references - ] - - @property - def in_dataloader(self) -> bool: - "Returns whether the current process is in a dataloader" - return self.active_dataloader is not None - - @staticmethod - def _reset_state(): - "Resets `_shared_state`, is used internally and should not be called" - GradientState._shared_state.clear() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/__init__.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/__init__.py deleted file mode 100644 index 14f2f4f5a459fef9e9862ab32ce42c36e2ed573e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/__init__.py +++ /dev/null @@ -1,65 +0,0 @@ -# Copyright 2020 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 .testing import ( - DEFAULT_LAUNCH_COMMAND, - are_the_same_tensors, - assert_exception, - capture_call_output, - device_count, - execute_subprocess_async, - get_launch_command, - get_torch_dist_unique_port, - memory_allocated_func, - path_in_accelerate_package, - pytest_xdist_worker_id, - require_bnb, - require_cpu, - require_cuda, - require_cuda_or_hpu, - require_cuda_or_xpu, - require_fp8, - require_fp16, - require_huggingface_suite, - require_mlu, - require_mps, - require_multi_device, - require_multi_gpu, - require_multi_gpu_or_xpu, - require_multi_xpu, - require_musa, - require_non_cpu, - require_non_hpu, - require_non_torch_xla, - require_non_xpu, - require_npu, - require_pippy, - require_sdaa, - require_single_device, - require_single_gpu, - require_single_xpu, - require_torch_min_version, - require_torchao, - require_torchvision, - require_tpu, - require_transformer_engine, - require_xpu, - run_first, - skip, - slow, - torch_device, -) -from .training import RegressionDataset, RegressionModel, RegressionModel4XPU - - -from .scripts import test_script, test_sync, test_ops # isort: skip diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 21d646ac8dc0ab78491bf6b26f696fdacb6206f3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/__pycache__/examples.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/__pycache__/examples.cpython-310.pyc deleted file mode 100644 index fd174b1f6074872799e59b0d67d4e0ff99e13147..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/__pycache__/examples.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/__pycache__/testing.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/__pycache__/testing.cpython-310.pyc deleted file mode 100644 index 822804acdbf1694faad036c6745387efce82fde9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/__pycache__/testing.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/__pycache__/training.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/__pycache__/training.cpython-310.pyc deleted file mode 100644 index 0a37724d0380df51bc06b43e7d5db0de52538108..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/__pycache__/training.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/examples.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/examples.py deleted file mode 100644 index 79d09eedbc263dcab49519bd7e15d08c7cf15ed2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/examples.py +++ /dev/null @@ -1,145 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2022 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. -""" -A collection of utilities for comparing `examples/complete_*_example.py` scripts with the capabilities inside of each -`examples/by_feature` example. `compare_against_test` is the main function that should be used when testing, while the -others are used to either get the code that matters, or to preprocess them (such as stripping comments) -""" - -import os - - -def get_function_contents_by_name(lines: list[str], name: str): - """ - Extracts a function from `lines` of segmented source code with the name `name`. - - Args: - lines (`List[str]`): - Source code of a script separated by line. - name (`str`): - The name of the function to extract. Should be either `training_function` or `main` - """ - if name != "training_function" and name != "main": - raise ValueError(f"Incorrect function name passed: {name}, choose either 'main' or 'training_function'") - good_lines, found_start = [], False - for line in lines: - if not found_start and f"def {name}" in line: - found_start = True - good_lines.append(line) - continue - if found_start: - if name == "training_function" and "def main" in line: - return good_lines - if name == "main" and "if __name__" in line: - return good_lines - good_lines.append(line) - - -def clean_lines(lines: list[str]): - """ - Filters `lines` and removes any entries that start with a comment ('#') or is just a newline ('\n') - - Args: - lines (`List[str]`): - Source code of a script separated by line. - """ - return [line for line in lines if not line.lstrip().startswith("#") and line != "\n"] - - -def compare_against_test(base_filename: str, feature_filename: str, parser_only: bool, secondary_filename: str = None): - """ - Tests whether the additional code inside of `feature_filename` was implemented in `base_filename`. This should be - used when testing to see if `complete_*_.py` examples have all of the implementations from each of the - `examples/by_feature/*` scripts. - - It utilizes `nlp_example.py` to extract out all of the repeated training code, so that only the new additional code - is examined and checked. If something *other* than `nlp_example.py` should be used, such as `cv_example.py` for the - `complete_cv_example.py` script, it should be passed in for the `secondary_filename` parameter. - - Args: - base_filename (`str` or `os.PathLike`): - The filepath of a single "complete" example script to test, such as `examples/complete_cv_example.py` - feature_filename (`str` or `os.PathLike`): - The filepath of a single feature example script. The contents of this script are checked to see if they - exist in `base_filename` - parser_only (`bool`): - Whether to compare only the `main()` sections in both files, or to compare the contents of - `training_loop()` - secondary_filename (`str`, *optional*): - A potential secondary filepath that should be included in the check. This function extracts the base - functionalities off of "examples/nlp_example.py", so if `base_filename` is a script other than - `complete_nlp_example.py`, the template script should be included here. Such as `examples/cv_example.py` - """ - with open(base_filename) as f: - base_file_contents = f.readlines() - with open(os.path.abspath(os.path.join("examples", "nlp_example.py"))) as f: - full_file_contents = f.readlines() - with open(feature_filename) as f: - feature_file_contents = f.readlines() - if secondary_filename is not None: - with open(secondary_filename) as f: - secondary_file_contents = f.readlines() - - # This is our base, we remove all the code from here in our `full_filename` and `feature_filename` to find the new content - if parser_only: - base_file_func = clean_lines(get_function_contents_by_name(base_file_contents, "main")) - full_file_func = clean_lines(get_function_contents_by_name(full_file_contents, "main")) - feature_file_func = clean_lines(get_function_contents_by_name(feature_file_contents, "main")) - if secondary_filename is not None: - secondary_file_func = clean_lines(get_function_contents_by_name(secondary_file_contents, "main")) - else: - base_file_func = clean_lines(get_function_contents_by_name(base_file_contents, "training_function")) - full_file_func = clean_lines(get_function_contents_by_name(full_file_contents, "training_function")) - feature_file_func = clean_lines(get_function_contents_by_name(feature_file_contents, "training_function")) - if secondary_filename is not None: - secondary_file_func = clean_lines( - get_function_contents_by_name(secondary_file_contents, "training_function") - ) - - _dl_line = "train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n" - - # Specific code in our script that differs from the full version, aka what is new - new_feature_code = [] - passed_idxs = [] # We keep track of the idxs just in case it's a repeated statement - it = iter(feature_file_func) - for i in range(len(feature_file_func) - 1): - if i not in passed_idxs: - line = next(it) - if (line not in full_file_func) and (line.lstrip() != _dl_line): - if "TESTING_MOCKED_DATALOADERS" not in line: - new_feature_code.append(line) - passed_idxs.append(i) - else: - # Skip over the `config['num_epochs'] = 2` statement - _ = next(it) - - # Extract out just the new parts from the full_file_training_func - new_full_example_parts = [] - passed_idxs = [] # We keep track of the idxs just in case it's a repeated statement - for i, line in enumerate(base_file_func): - if i not in passed_idxs: - if (line not in full_file_func) and (line.lstrip() != _dl_line): - if "TESTING_MOCKED_DATALOADERS" not in line: - new_full_example_parts.append(line) - passed_idxs.append(i) - - # Finally, get the overall diff - diff_from_example = [line for line in new_feature_code if line not in new_full_example_parts] - if secondary_filename is not None: - diff_from_two = [line for line in full_file_contents if line not in secondary_file_func] - diff_from_example = [line for line in diff_from_example if line not in diff_from_two] - - return diff_from_example diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__init__.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__init__.py deleted file mode 100644 index c9cbe26c257b515f657c05e1996d517e69613972..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright 2020 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. diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index faeec3e0ed0ecfd33f1a41a92a7046a0b2f72048..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_cli.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_cli.cpython-310.pyc deleted file mode 100644 index f7cce7bfc91448b93175c668e04d6eff0451e946..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_cli.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_ddp_comm_hook.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_ddp_comm_hook.cpython-310.pyc deleted file mode 100644 index aab3548b518620e89472f2e5091825f376202b2d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_ddp_comm_hook.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_distributed_data_loop.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_distributed_data_loop.cpython-310.pyc deleted file mode 100644 index 1bdaaeec24bc57c5296873e44cc2488b8ebe62b3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_distributed_data_loop.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_merge_weights.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_merge_weights.cpython-310.pyc deleted file mode 100644 index f48657998ec1bd2268a14e8022290c016fd68cc4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_merge_weights.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_notebook.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_notebook.cpython-310.pyc deleted file mode 100644 index cbb131cfef520fe587e00e93f62a2b43901cd15d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_notebook.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_ops.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_ops.cpython-310.pyc deleted file mode 100644 index eb411711185f10489b76cec72707b0ee9853dcb5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_ops.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_script.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_script.cpython-310.pyc deleted file mode 100644 index 77ef923ede049881409a5d236b46d1ca84922ef8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_script.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_sync.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_sync.cpython-310.pyc deleted file mode 100644 index e87eef8607120864449fa4f6f8527c3bd2ffe35e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/__pycache__/test_sync.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__init__.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__init__.py deleted file mode 100644 index c9cbe26c257b515f657c05e1996d517e69613972..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright 2020 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. diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 4c3bc486ef902410251a6c0fecb247a42c9801d2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_checkpointing.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_checkpointing.cpython-310.pyc deleted file mode 100644 index 083a957f2d83c8049f0292eee8f9fec04a37aaf4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_checkpointing.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_ds_multiple_model.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_ds_multiple_model.cpython-310.pyc deleted file mode 100644 index c7520b0bf1e4aae11ed29f8b5becc760eb7483ef..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_ds_multiple_model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_metrics.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_metrics.cpython-310.pyc deleted file mode 100644 index 52729838b0bad19cc293ca2604a260a90937b65e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_metrics.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_peak_memory_usage.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_peak_memory_usage.cpython-310.pyc deleted file mode 100644 index 7cc060478d6f3917a92e4a89f1d126e651e608f2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_peak_memory_usage.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_performance.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_performance.cpython-310.pyc deleted file mode 100644 index aa6c2eb4453a0b039a90e71e48d110f7d814760d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_performance.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_pippy.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_pippy.cpython-310.pyc deleted file mode 100644 index a530a1f570d585e532f9320ce9a1e12d54c870ec..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_pippy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_zero3_integration.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_zero3_integration.cpython-310.pyc deleted file mode 100644 index 17c32f7f7e22597140a684a86aba7f373ec0854a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/__pycache__/test_zero3_integration.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_checkpointing.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_checkpointing.py deleted file mode 100644 index 6a1553898ec3d55e64822c204ddf7e705069ce8a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_checkpointing.py +++ /dev/null @@ -1,269 +0,0 @@ -# Copyright 2022 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 json -import os - -import evaluate -import torch -from datasets import load_dataset -from torch.optim import AdamW -from torch.utils.data import DataLoader -from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed - -from accelerate import Accelerator, DistributedType -from accelerate.utils.deepspeed import DummyOptim, DummyScheduler - - -MAX_GPU_BATCH_SIZE = 16 -EVAL_BATCH_SIZE = 32 - - -def get_dataloaders(accelerator: Accelerator, batch_size: int = 16, model_name: str = "bert-base-cased"): - """ - Creates a set of `DataLoader`s for the `glue` dataset. - - Args: - accelerator (`Accelerator`): - An `Accelerator` object - batch_size (`int`, *optional*): - The batch size for the train and validation DataLoaders. - model_name (`str`, *optional*): - """ - tokenizer = AutoTokenizer.from_pretrained(model_name) - datasets = load_dataset("glue", "mrpc") - - def tokenize_function(examples): - # max_length=None => use the model max length (it's actually the default) - outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) - return outputs - - # Apply the method we just defined to all the examples in all the splits of the dataset - tokenized_datasets = datasets.map( - tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False - ) - - # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the - # transformers library - tokenized_datasets = tokenized_datasets.rename_column("label", "labels") - - def collate_fn(examples): - # On TPU it's best to pad everything to the same length or training will be very slow. - if accelerator.distributed_type == DistributedType.XLA: - return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt") - return tokenizer.pad(examples, padding="longest", return_tensors="pt") - - # Instantiate dataloaders. - train_dataloader = DataLoader( - tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size - ) - eval_dataloader = DataLoader( - tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE - ) - - return train_dataloader, eval_dataloader - - -def evaluation_loop(accelerator, model, eval_dataloader, metric): - model.eval() - samples_seen = 0 - for step, batch in enumerate(eval_dataloader): - # We could avoid this line since we set the accelerator with `device_placement=True`. - batch.to(accelerator.device) - with torch.no_grad(): - outputs = model(**batch) - predictions = outputs.logits.argmax(dim=-1) - # It is slightly faster to call this once, than multiple times - predictions, references = accelerator.gather( - (predictions, batch["labels"]) - ) # If we are in a multiprocess environment, the last batch has duplicates - if accelerator.use_distributed: - if step == len(eval_dataloader) - 1: - predictions = predictions[: len(eval_dataloader.dataset) - samples_seen] - references = references[: len(eval_dataloader.dataset) - samples_seen] - else: - samples_seen += references.shape[0] - metric.add_batch( - predictions=predictions, - references=references, - ) - - eval_metric = metric.compute() - return eval_metric["accuracy"] - - -def training_function(config, args): - # Initialize accelerator - accelerator = Accelerator() - - # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs - lr = config["lr"] - num_epochs = int(config["num_epochs"]) - seed = int(config["seed"]) - batch_size = int(config["batch_size"]) - model_name = args.model_name_or_path - - set_seed(seed) - train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size, model_name) - - # Instantiate the model (we build the model here so that the seed also control new weights initialization) - model = AutoModelForSequenceClassification.from_pretrained(model_name, return_dict=True) - - # Instantiate optimizer - optimizer_cls = ( - AdamW - if accelerator.state.deepspeed_plugin is None - or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config - else DummyOptim - ) - optimizer = optimizer_cls(params=model.parameters(), lr=lr) - - if accelerator.state.deepspeed_plugin is not None: - gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[ - "gradient_accumulation_steps" - ] - else: - gradient_accumulation_steps = 1 - max_training_steps = (len(train_dataloader) * num_epochs) // gradient_accumulation_steps - - # Instantiate scheduler - if ( - accelerator.state.deepspeed_plugin is None - or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config - ): - lr_scheduler = get_linear_schedule_with_warmup( - optimizer=optimizer, - num_warmup_steps=0, - num_training_steps=max_training_steps, - ) - else: - lr_scheduler = DummyScheduler(optimizer, total_num_steps=max_training_steps, warmup_num_steps=0) - - # Prepare everything - # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the - # prepare method. - model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare( - model, optimizer, train_dataloader, eval_dataloader, lr_scheduler - ) - - # We need to keep track of how many total steps we have iterated over - overall_step = 0 - # We also need to keep track of the stating epoch so files are named properly - starting_epoch = 0 - metric = evaluate.load("glue", "mrpc") - ending_epoch = num_epochs - - if args.partial_train_epoch is not None: - ending_epoch = args.partial_train_epoch - - if args.resume_from_checkpoint: - accelerator.load_state(args.resume_from_checkpoint) - epoch_string = args.resume_from_checkpoint.split("epoch_")[1] - state_epoch_num = "" - for char in epoch_string: - if char.isdigit(): - state_epoch_num += char - else: - break - starting_epoch = int(state_epoch_num) + 1 - accuracy = evaluation_loop(accelerator, model, eval_dataloader, metric) - accelerator.print("resumed checkpoint performance:", accuracy) - accelerator.print("resumed checkpoint's scheduler's lr:", lr_scheduler.get_lr()[0]) - accelerator.print("resumed optimizers's lr:", optimizer.param_groups[0]["lr"]) - with open(os.path.join(args.output_dir, f"state_{starting_epoch - 1}.json")) as f: - resumed_state = json.load(f) - assert resumed_state["accuracy"] == accuracy, "Accuracy mismatch, loading from checkpoint failed" - assert resumed_state["lr"] == lr_scheduler.get_lr()[0], ( - "Scheduler learning rate mismatch, loading from checkpoint failed" - ) - assert resumed_state["optimizer_lr"] == optimizer.param_groups[0]["lr"], ( - "Optimizer learning rate mismatch, loading from checkpoint failed" - ) - assert resumed_state["epoch"] == starting_epoch - 1, "Epoch mismatch, loading from checkpoint failed" - return - - # Now we train the model - state = {} - for epoch in range(starting_epoch, ending_epoch): - model.train() - for step, batch in enumerate(train_dataloader): - outputs = model(**batch) - loss = outputs.loss - loss = loss / gradient_accumulation_steps - accelerator.backward(loss) - if step % gradient_accumulation_steps == 0: - optimizer.step() - lr_scheduler.step() - optimizer.zero_grad() - - overall_step += 1 - output_dir = f"epoch_{epoch}" - output_dir = os.path.join(args.output_dir, output_dir) - accelerator.save_state(output_dir) - accuracy = evaluation_loop(accelerator, model, eval_dataloader, metric) - state["accuracy"] = accuracy - state["lr"] = lr_scheduler.get_lr()[0] - state["optimizer_lr"] = optimizer.param_groups[0]["lr"] - state["epoch"] = epoch - state["step"] = overall_step - accelerator.print(f"epoch {epoch}:", state) - - accelerator.wait_for_everyone() - if accelerator.is_main_process: - with open(os.path.join(args.output_dir, f"state_{epoch}.json"), "w") as f: - json.dump(state, f) - accelerator.end_training() - - -def main(): - parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.") - parser.add_argument( - "--model_name_or_path", - type=str, - default="bert-base-cased", - help="Path to pretrained model or model identifier from huggingface.co/models.", - required=False, - ) - parser.add_argument( - "--output_dir", - type=str, - default=".", - help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.", - ) - parser.add_argument( - "--resume_from_checkpoint", - type=str, - default=None, - help="If the training should continue from a checkpoint folder.", - ) - parser.add_argument( - "--partial_train_epoch", - type=int, - default=None, - help="If passed, the training will stop after this number of epochs.", - ) - parser.add_argument( - "--num_epochs", - type=int, - default=2, - help="Number of train epochs.", - ) - args = parser.parse_args() - config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16} - - training_function(config, args) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_ds_multiple_model.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_ds_multiple_model.py deleted file mode 100644 index 3729ecf4c72190bf865d620b6941206ab904818c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_ds_multiple_model.py +++ /dev/null @@ -1,332 +0,0 @@ -# Copyright 2024 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. - -""" -Test script for verifying multiple models can be utilized with Accelerate + DeepSpeed: - -Scenario 1: One model is training, another model is being used for inference/logits to impact training in some form. -Scenario 2: Two models are training simultaneously, which means two optimizers, etc. -""" - -import argparse -from pathlib import Path - -import evaluate -import torch -from datasets import load_dataset -from torch.optim import AdamW -from torch.utils.data import DataLoader -from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup - -from accelerate import Accelerator, DeepSpeedPlugin, DistributedType -from accelerate.state import AcceleratorState -from accelerate.utils.deepspeed import get_active_deepspeed_plugin - - -MAX_GPU_BATCH_SIZE = 16 -EVAL_BATCH_SIZE = 32 - - -class NoiseModel(torch.nn.Module): - def __init__(self, noise_factor=0.1): - super().__init__() - self.noise_factor = torch.nn.Parameter(torch.tensor(noise_factor, dtype=torch.float32)) - - def forward(self, loss): - return loss * self.noise_factor - - -def get_dataloaders(accelerator: Accelerator, batch_size: int = 16, model_name: str = "bert-base-cased"): - """ - Creates a set of `DataLoader`s for the `glue` dataset. - - Args: - accelerator (`Accelerator`): - An `Accelerator` object - batch_size (`int`, *optional*): - The batch size for the train and validation DataLoaders. - model_name (`str`, *optional*): - """ - tokenizer = AutoTokenizer.from_pretrained(model_name) - datasets = load_dataset("glue", "mrpc") - - def tokenize_function(examples): - # max_length=None => use the model max length (it's actually the default) - outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) - return outputs - - # Apply the method we just defined to all the examples in all the splits of the dataset - tokenized_datasets = datasets.map( - tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False - ) - - # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the - # transformers library - tokenized_datasets = tokenized_datasets.rename_column("label", "labels") - - def collate_fn(examples): - # On TPU it's best to pad everything to the same length or training will be very slow. - if accelerator.distributed_type == DistributedType.XLA: - return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt") - return tokenizer.pad(examples, padding="longest", return_tensors="pt") - - # Instantiate dataloaders. - train_dataloader = DataLoader( - tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size - ) - eval_dataloader = DataLoader( - tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE - ) - - return train_dataloader, eval_dataloader - - -test_file_path = __file__ -path = Path(test_file_path).resolve() -test_file_dir_str = str(path.parent.parent.parent.parent.parent.parent) - -# Create our DS plugins -# We use custom schedulers and optimizers, hence `model_only` -ds_config_file = dict( - zero2=f"{test_file_dir_str}/tests/deepspeed/ds_config_zero2_model_only.json", - zero3=f"{test_file_dir_str}/tests/deepspeed/ds_config_zero3_model_only.json", -) - - -def single_model_training(config, args): - # Training a single model, we have a `noise` model that is untrainable used to inject some noise into the training process - num_epochs = config["num_epochs"] - zero2_plugin = DeepSpeedPlugin(hf_ds_config=ds_config_file["zero2"]) - zero3_plugin = DeepSpeedPlugin(hf_ds_config=ds_config_file["zero3"]) - - deepspeed_plugins = {"training": zero2_plugin, "inference": zero3_plugin} - - # Initialize accelerator - accelerator = Accelerator( - deepspeed_plugins=deepspeed_plugins, - mixed_precision="bf16", - ) - - # Initialize model under zero2 plugin - assert get_active_deepspeed_plugin(accelerator.state) is zero2_plugin - train_model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path) - train_dataloader, eval_dataloader = get_dataloaders( - accelerator, batch_size=config["batch_size"], model_name=args.model_name_or_path - ) - max_training_steps = len(train_dataloader) * config["num_epochs"] - optimizer = AdamW(train_model.parameters(), lr=config["lr"]) - lr_scheduler = get_linear_schedule_with_warmup( - optimizer, num_warmup_steps=0, num_training_steps=max_training_steps - ) - - train_dataloader, eval_dataloader, train_model, optimizer, lr_scheduler = accelerator.prepare( - train_dataloader, eval_dataloader, train_model, optimizer, lr_scheduler - ) - - # Now prepare the model under zero3 plugin - accelerator.state.select_deepspeed_plugin("inference") - assert get_active_deepspeed_plugin(accelerator.state) is zero3_plugin - inference_model = NoiseModel() - inference_model = accelerator.prepare(inference_model) - inference_model.eval() - - # Run training loop - accelerator.state.select_deepspeed_plugin("training") - # We also need to keep track of the stating epoch so files are named properly - starting_epoch = 0 - - # Now we train the model - best_performance = 0 - metric = evaluate.load("glue", "mrpc") - performance_metric = {} - for epoch in range(starting_epoch, num_epochs): - train_model.train() - inference_model.train() - for step, batch in enumerate(train_dataloader): - with accelerator.accumulate(train_model): - outputs_1 = train_model(**batch) - with torch.no_grad(): - outputs_2 = inference_model(outputs_1.loss) - # Combine the losses - loss = outputs_1.loss + outputs_2 - accelerator.backward(loss) - optimizer.step() - lr_scheduler.step() - optimizer.zero_grad() - - train_model.eval() - for step, batch in enumerate(eval_dataloader): - with torch.no_grad(): - outputs = train_model(**batch) - predictions = outputs.logits.argmax(dim=-1) - # It is slightly faster to call this once, than multiple times - predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"])) - metric.add_batch( - predictions=predictions, - references=references, - ) - - eval_metric = metric.compute() - # Use accelerator.print to print only on the main process. - accelerator.print(f"epoch {epoch}:", eval_metric) - performance_metric[f"epoch-{epoch}"] = eval_metric["accuracy"] - - if best_performance < eval_metric["accuracy"]: - best_performance = eval_metric["accuracy"] - assert best_performance > performance_metric["epoch-0"] - - -def multiple_model_training(config, args): - # This will essentially be like a k-fold model, but one model is Zero-2 and another model is Zero-3 - num_epochs = config["num_epochs"] - zero2_plugin = DeepSpeedPlugin(hf_ds_config=ds_config_file["zero2"]) - zero3_plugin = DeepSpeedPlugin(hf_ds_config=ds_config_file["zero3"]) - - deepspeed_plugins = {"zero2": zero2_plugin, "zero3": zero3_plugin} - - # Initialize accelerator - zero2_accelerator = Accelerator( - deepspeed_plugins=deepspeed_plugins, - mixed_precision="bf16", - ) - - # Since an `AcceleratorState` has already been made, we can just reuse it here - zero3_accelerator = Accelerator() - - # Initialize model under zero2 plugin - assert get_active_deepspeed_plugin(zero2_accelerator.state) is zero2_plugin - zero2_model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path) - train_dataloader, eval_dataloader = get_dataloaders( - zero2_accelerator, batch_size=config["batch_size"], model_name=args.model_name_or_path - ) - max_training_steps = len(train_dataloader) * config["num_epochs"] - zero2_optimizer = AdamW(zero2_model.parameters(), lr=config["lr"]) - zero2_lr_scheduler = get_linear_schedule_with_warmup( - zero2_optimizer, num_warmup_steps=0, num_training_steps=max_training_steps - ) - - train_dataloader, eval_dataloader, zero2_model, zero2_optimizer, zero2_lr_scheduler = zero2_accelerator.prepare( - train_dataloader, eval_dataloader, zero2_model, zero2_optimizer, zero2_lr_scheduler - ) - assert zero2_accelerator.deepspeed_engine_wrapped.engine is zero2_model - - # now do Zero3 - zero3_accelerator.state.select_deepspeed_plugin("zero3") - zero3_plugin.deepspeed_config["train_micro_batch_size_per_gpu"] = zero2_plugin.deepspeed_config[ - "train_micro_batch_size_per_gpu" - ] - assert get_active_deepspeed_plugin(zero3_accelerator.state) is zero3_plugin - zero3_model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path) - zero3_optimizer = AdamW(zero3_model.parameters(), lr=config["lr"]) - zero3_lr_scheduler = get_linear_schedule_with_warmup( - zero3_optimizer, num_warmup_steps=0, num_training_steps=max_training_steps - ) - zero3_model, zero3_optimizer, zero3_lr_scheduler = zero3_accelerator.prepare( - zero3_model, zero3_optimizer, zero3_lr_scheduler - ) - assert zero3_accelerator.deepspeed_engine_wrapped.engine is zero3_model - - # Run training loop - starting_epoch = 0 - - # Now we train the model - best_performance_a = 0 - best_performance_b = 0 - metric_a = evaluate.load("glue", "mrpc") - metric_b = evaluate.load("glue", "mrpc") - performance_metric_a = {} - performance_metric_b = {} - for epoch in range(starting_epoch, num_epochs): - zero2_model.train() - zero3_model.train() - for step, batch in enumerate(train_dataloader): - with zero2_accelerator.accumulate(zero2_model, zero3_model): - outputs_1 = zero2_model(**batch) - zero2_accelerator.backward(outputs_1.loss) - zero2_optimizer.step() - zero2_lr_scheduler.step() - zero2_optimizer.zero_grad() - outputs_2 = zero3_model(**batch) - zero3_accelerator.backward(outputs_2.loss) - zero3_optimizer.step() - zero3_lr_scheduler.step() - zero3_optimizer.zero_grad() - - zero2_model.eval() - zero3_model.eval() - for step, batch in enumerate(eval_dataloader): - with torch.no_grad(): - logits_a = zero2_model(**batch).logits - logits_b = zero3_model(**batch).logits - # Combine the logits from both models - predictions_a = logits_a.argmax(dim=-1) - predictions_b = logits_b.argmax(dim=-1) - # It is slightly faster to call this once, than multiple times - predictions_a, predictions_b, references = zero2_accelerator.gather_for_metrics( - (predictions_a, predictions_b, batch["labels"]) - ) - metric_a.add_batch( - predictions=predictions_a, - references=references, - ) - metric_b.add_batch( - predictions=predictions_b, - references=references, - ) - - eval_metric_a = metric_a.compute() - eval_metric_b = metric_b.compute() - # Use accelerator.print to print only on the main process. - zero2_accelerator.print(f"epoch {epoch}:", eval_metric_a, eval_metric_b) - performance_metric_a[f"epoch-{epoch}"] = eval_metric_a["accuracy"] - performance_metric_b[f"epoch-{epoch}"] = eval_metric_b["accuracy"] - - if best_performance_a < eval_metric_a["accuracy"]: - best_performance_a = eval_metric_a["accuracy"] - if best_performance_b < eval_metric_b["accuracy"]: - best_performance_b = eval_metric_b["accuracy"] - assert best_performance_a > performance_metric_a["epoch-0"] - assert best_performance_b > performance_metric_b["epoch-0"] - - -def main(): - parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.") - parser.add_argument( - "--model_name_or_path", - type=str, - default="bert-base-cased", - help="Path to pretrained model or model identifier from huggingface.co/models.", - required=False, - ) - parser.add_argument( - "--performance_lower_bound", - type=float, - default=None, - help="Optional lower bound for the performance metric. If set, the training will throw error when the performance metric drops below this value.", - ) - parser.add_argument( - "--num_epochs", - type=int, - default=2, - help="Number of train epochs.", - ) - args = parser.parse_args() - config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16} - single_model_training(config, args) - AcceleratorState._reset_state(True) - multiple_model_training(config, args) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_metrics.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_metrics.py deleted file mode 100644 index d1bfe351509148ebc48067584e9d61b93e7210a6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_metrics.py +++ /dev/null @@ -1,307 +0,0 @@ -# Copyright 2022 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 logging -import math -import os -from copy import deepcopy - -import datasets -import evaluate -import torch -import transformers -from datasets import load_dataset -from torch.utils.data import DataLoader, IterableDataset -from transformers import AutoModelForSequenceClassification, AutoTokenizer - -from accelerate import Accelerator, DataLoaderConfiguration, DistributedType -from accelerate.data_loader import DataLoaderDispatcher -from accelerate.test_utils import RegressionDataset, RegressionModel, torch_device -from accelerate.utils import is_torch_xla_available, set_seed - - -os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true" - - -class ListHandler(logging.Handler): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.logs = [] - - def emit(self, record): - self.logs.append(record) - - -def get_basic_setup(accelerator, num_samples=82, batch_size=16): - "Returns everything needed to perform basic training" - set_seed(42) - model = RegressionModel() - ddp_model = deepcopy(model) - dset = RegressionDataset(length=num_samples) - dataloader = DataLoader(dset, batch_size=batch_size) - model.to(accelerator.device) - ddp_model, dataloader = accelerator.prepare(ddp_model, dataloader) - return model, ddp_model, dataloader - - -def get_dataloader(accelerator: Accelerator, use_longest=False): - tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/mrpc-bert-base-cased") - dataset = load_dataset("glue", "mrpc", split="validation") - - def tokenize_function(examples): - outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) - return outputs - - with accelerator.main_process_first(): - tokenized_datasets = dataset.map( - tokenize_function, - batched=True, - remove_columns=["idx", "sentence1", "sentence2"], - ) - - tokenized_datasets = tokenized_datasets.rename_column("label", "labels") - - def collate_fn(examples): - if use_longest: - return tokenizer.pad(examples, padding="longest", return_tensors="pt") - return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt") - - return DataLoader(tokenized_datasets, shuffle=False, collate_fn=collate_fn, batch_size=16) - - -def get_mrpc_setup(dispatch_batches, split_batches): - dataloader_config = DataLoaderConfiguration(dispatch_batches=dispatch_batches, split_batches=split_batches) - accelerator = Accelerator(dataloader_config=dataloader_config) - dataloader = get_dataloader(accelerator, not dispatch_batches) - model = AutoModelForSequenceClassification.from_pretrained( - "hf-internal-testing/mrpc-bert-base-cased", return_dict=True - ) - ddp_model, ddp_dataloader = accelerator.prepare(model, dataloader) - return { - "ddp": [ddp_model, ddp_dataloader, torch_device], - "no": [model, dataloader, accelerator.device], - }, accelerator - - -def generate_predictions(model, dataloader, accelerator): - logits_and_targets = [] - for batch in dataloader: - input, target = batch.values() - with torch.no_grad(): - logit = model(input) - logit, target = accelerator.gather_for_metrics((logit, target)) - logits_and_targets.append((logit, target)) - logits, targs = [], [] - for logit, targ in logits_and_targets: - logits.append(logit) - targs.append(targ) - logits, targs = torch.cat(logits), torch.cat(targs) - return logits, targs - - -def test_torch_metrics( - accelerator: Accelerator, num_samples=82, dispatch_batches=False, split_batches=False, batch_size=16 -): - _, ddp_model, dataloader = get_basic_setup(accelerator, num_samples, batch_size) - logits, _ = generate_predictions(ddp_model, dataloader, accelerator) - assert len(logits) == num_samples, ( - f"Unexpected number of inputs:\n Expected: {num_samples}\n Actual: {len(logits)}" - ) - - -def test_mrpc(dispatch_batches: bool = False, split_batches: bool = False): - metric = evaluate.load("glue", "mrpc") - setup, accelerator = get_mrpc_setup(dispatch_batches, split_batches) - # First do baseline - model, dataloader, device = setup["no"] - model.to(device) - model.eval() - for batch in dataloader: - batch.to(device) - with torch.inference_mode(): - outputs = model(**batch) - preds = outputs.logits.argmax(dim=-1) - metric.add_batch(predictions=preds, references=batch["labels"]) - baseline = metric.compute() - - # Then do distributed - model, dataloader, device = setup["ddp"] - model.eval() - for batch in dataloader: - with torch.inference_mode(): - outputs = model(**batch) - preds = outputs.logits.argmax(dim=-1) - references = batch["labels"] - preds, references = accelerator.gather_for_metrics((preds, references)) - metric.add_batch(predictions=preds, references=references) - distributed = metric.compute() - - for key in "accuracy f1".split(): - assert math.isclose(baseline[key], distributed[key]), ( - f"Baseline and Distributed are not the same for key {key}:\n\tBaseline: {baseline[key]}\n\tDistributed: {distributed[key]}\n" - ) - - -def test_gather_for_metrics_with_non_tensor_objects_iterable_dataset(): - class DummyIterableDataset(IterableDataset): - def __init__(self, data): - self.data = data - - def __len__(self): - return len(self.data) - - def __iter__(self): - yield from self.data - - iterable_dataset = DummyIterableDataset([n for n in range(30)]) - dataloader = DataLoader(iterable_dataset, batch_size=4) - accelerator = Accelerator() - prepared_dataloader = accelerator.prepare(dataloader) - - if accelerator.is_main_process: - logger = logging.root.manager.loggerDict["accelerate.accelerator"] - list_handler = ListHandler() - logger.addHandler(list_handler) - - batches_for_metrics = [] - for batch in prepared_dataloader: - batches_for_metrics.append(accelerator.gather_for_metrics(batch)) - - assert torch.cat(batches_for_metrics).size(0) == 30 - - if accelerator.is_main_process: - assert len(list_handler.logs) == 0 - logger.removeHandler(list_handler) - - -def test_gather_for_metrics_with_iterable_dataset(): - class DummyIterableDataset(IterableDataset): - def __init__(self, data): - self.data = data - - def __len__(self): - return len(self.data) - - def __iter__(self): - yield from self.data - - iterable_dataset = DummyIterableDataset(torch.as_tensor(range(30))) - dataloader = DataLoader(iterable_dataset, batch_size=4) - - accelerator = Accelerator() - prepared_dataloader = accelerator.prepare(dataloader) - - assert isinstance(prepared_dataloader, DataLoaderDispatcher) - - if accelerator.is_main_process: - logger = logging.root.manager.loggerDict["accelerate.accelerator"] - list_handler = ListHandler() - logger.addHandler(list_handler) - - batches_for_metrics = [] - for batch in prepared_dataloader: - batches_for_metrics.append(accelerator.gather_for_metrics(batch)) - - assert torch.cat(batches_for_metrics).size(0) == 30 - - if accelerator.is_main_process: - assert len(list_handler.logs) == 0 - - logger.removeHandler(list_handler) - - -def test_gather_for_metrics_drop_last(): - accelerator = Accelerator() - per_device_batch_size = 5 - num_items = (10 * accelerator.num_processes) + 1 - dataloader = DataLoader(range(num_items), batch_size=per_device_batch_size, drop_last=True) - dataloader = accelerator.prepare(dataloader) - - iterator = iter(dataloader) - next(iterator) # Skip first batch tensor([0, 1, 2, 3, 4], device='cuda:0') - batch = next(iterator) - gathered_items = accelerator.gather_for_metrics(batch) - - # Should return a full set of complete batches from each GPU - num_expected_items = per_device_batch_size * accelerator.num_processes - assert gathered_items.size(0) == (num_expected_items), ( - f"Expected number of items: {num_expected_items}, Actual: {gathered_items.size(0)}" - ) - - -def main(): - dataloader_config = DataLoaderConfiguration(split_batches=False, dispatch_batches=False) - accelerator = Accelerator(dataloader_config=dataloader_config) - if accelerator.is_local_main_process: - datasets.utils.logging.set_verbosity_warning() - transformers.utils.logging.set_verbosity_warning() - else: - datasets.utils.logging.set_verbosity_error() - transformers.utils.logging.set_verbosity_error() - # TorchXLA does not support batch dispatching. 'put_on_device' is always False for - # TorchXLA, which can cause a value error in 'prepare_data_loader' function. - dispatch_batches_options = [False] if accelerator.state.distributed_type == DistributedType.XLA else [True, False] - - # Temporarily close this test for TorchXLA due to the 'Cannot set version_counter for - # inference tensor' error in inference mode. Reopen it after TorchXLA fixes this bug. - # These are a bit slower so they should only be ran on the GPU or TPU - if accelerator.device.type != "cpu" and not is_torch_xla_available(): - if accelerator.is_local_main_process: - print("**Testing gather_for_metrics**") - for split_batches in [True, False]: - for dispatch_batches in dispatch_batches_options: - if accelerator.is_local_main_process: - print(f"With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`") - test_mrpc(dispatch_batches, split_batches) - accelerator.state._reset_state() - print("test_gather_for_metrics_with_iterable_dataset") - test_gather_for_metrics_with_iterable_dataset() - print("test gather_for_metrics_with_non_tensor_objects_iterable_dataset") - test_gather_for_metrics_with_non_tensor_objects_iterable_dataset() - - # MpDeviceLoader in TorchXLA is an asynchronous loader that preloads several batches into cache. - # This can cause the 'end_of_dataloader' of DataLoaderStateMixin to be set earlier than intended. - # Skip this test when TorchXLA is enabled. - if accelerator.state.distributed_type != DistributedType.XLA: - if accelerator.is_local_main_process: - print("**Test torch metrics**") - for split_batches in [True, False]: - for dispatch_batches in dispatch_batches_options: - dataloader_config = DataLoaderConfiguration( - split_batches=split_batches, dispatch_batches=dispatch_batches - ) - accelerator = Accelerator(dataloader_config=dataloader_config) - if accelerator.is_local_main_process: - print(f"With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`, length=99") - test_torch_metrics(accelerator, 99) - accelerator.state._reset_state() - if accelerator.is_local_main_process: - print("**Test last batch is not dropped when perfectly divisible**") - accelerator = Accelerator() - test_torch_metrics(accelerator, 512) - accelerator.state._reset_state() - if accelerator.is_local_main_process: - print("**Test that `drop_last` is taken into account**") - test_gather_for_metrics_drop_last() - accelerator.end_training() - accelerator.state._reset_state() - - -def _mp_fn(index): - # For xla_spawn (TPUs) - main() - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_peak_memory_usage.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_peak_memory_usage.py deleted file mode 100644 index 723e5497656020516bf072cf4112f61f59c2e5cb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_peak_memory_usage.py +++ /dev/null @@ -1,314 +0,0 @@ -# Copyright 2022 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 gc -import json -import os - -import torch -from datasets import load_dataset -from torch.optim import AdamW -from torch.utils.data import DataLoader -from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed - -from accelerate import Accelerator, DistributedType -from accelerate.utils import ( - is_hpu_available, - is_mlu_available, - is_musa_available, - is_npu_available, - is_sdaa_available, - is_xpu_available, -) -from accelerate.utils.deepspeed import DummyOptim, DummyScheduler - - -MAX_GPU_BATCH_SIZE = 16 -EVAL_BATCH_SIZE = 32 - - -# Converting Bytes to Megabytes -def b2mb(x): - return int(x / 2**20) - - -# This context manager is used to track the peak memory usage of the process -class TorchTracemalloc: - def __enter__(self): - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.reset_max_memory_allocated() # reset the peak gauge to zero - self.begin = torch.cuda.memory_allocated() - elif is_mlu_available(): - torch.mlu.empty_cache() - torch.mlu.reset_max_memory_allocated() # reset the peak gauge to zero - self.begin = torch.mlu.memory_allocated() - elif is_sdaa_available(): - torch.sdaa.empty_cache() - torch.sdaa.reset_max_memory_allocated() # reset the peak gauge to zero - self.begin = torch.sdaa.memory_allocated() - elif is_musa_available(): - torch.musa.empty_cache() - torch.musa.reset_max_memory_allocated() # reset the peak gauge to zero - self.begin = torch.musa.memory_allocated() - elif is_npu_available(): - torch.npu.empty_cache() - torch.npu.reset_max_memory_allocated() # reset the peak gauge to zero - self.begin = torch.npu.memory_allocated() - elif is_xpu_available(): - torch.xpu.empty_cache() - torch.xpu.reset_max_memory_allocated() # reset the peak gauge to zero - self.begin = torch.xpu.memory_allocated() - elif is_hpu_available(): - # torch.hpu.empty_cache() # not available on hpu as it reserves all device memory for the current process - torch.hpu.reset_peak_memory_stats() # reset the peak gauge to zero - self.begin = torch.hpu.memory_allocated() - return self - - def __exit__(self, *exc): - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - self.end = torch.cuda.memory_allocated() - self.peak = torch.cuda.max_memory_allocated() - elif is_mlu_available(): - torch.mlu.empty_cache() - self.end = torch.mlu.memory_allocated() - self.begin = torch.mlu.max_memory_allocated() - elif is_sdaa_available(): - torch.sdaa.empty_cache() - self.end = torch.sdaa.memory_allocated() - self.begin = torch.sdaa.max_memory_allocated() - elif is_musa_available(): - torch.musa.empty_cache() - self.end = torch.musa.memory_allocated() - self.begin = torch.musa.max_memory_allocated() - elif is_npu_available(): - torch.npu.empty_cache() - self.end = torch.npu.memory_allocated() - self.peak = torch.npu.max_memory_allocated() - elif is_xpu_available(): - torch.xpu.empty_cache() - self.end = torch.xpu.memory_allocated() - self.peak = torch.xpu.max_memory_allocated() - elif is_hpu_available(): - # torch.hpu.empty_cache() # not available on hpu as it reserves all device memory for the current process - self.end = torch.hpu.memory_allocated() - self.peak = torch.hpu.max_memory_allocated() - self.used = b2mb(self.end - self.begin) - self.peaked = b2mb(self.peak - self.begin) - # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}") - - -def get_dataloaders( - accelerator: Accelerator, - batch_size: int = 16, - model_name: str = "bert-base-cased", - n_train: int = 320, - n_val: int = 160, -): - """ - Creates a set of `DataLoader`s for the `glue` dataset. - - Args: - accelerator (`Accelerator`): - An `Accelerator` object - batch_size (`int`, *optional*): - The batch size for the train and validation DataLoaders. - model_name (`str`, *optional*): - The name of the model to use. - n_train (`int`, *optional*): - The number of training examples to use. - n_val (`int`, *optional*): - The number of validation examples to use. - """ - tokenizer = AutoTokenizer.from_pretrained(model_name) - datasets = load_dataset( - "glue", "mrpc", split={"train": f"train[:{n_train}]", "validation": f"validation[:{n_val}]"} - ) - - def tokenize_function(examples): - # max_length=None => use the model max length (it's actually the default) - outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) - return outputs - - # Apply the method we just defined to all the examples in all the splits of the dataset - tokenized_datasets = datasets.map( - tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False - ) - - # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the - # transformers library - tokenized_datasets = tokenized_datasets.rename_column("label", "labels") - - def collate_fn(examples): - # On TPU it's best to pad everything to the same length or training will be very slow. - if accelerator.distributed_type == DistributedType.XLA: - return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt") - return tokenizer.pad(examples, padding="longest", return_tensors="pt") - - # Instantiate dataloaders. - train_dataloader = DataLoader( - tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size - ) - eval_dataloader = DataLoader( - tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE - ) - - return train_dataloader, eval_dataloader - - -def training_function(config, args): - # Initialize accelerator - accelerator = Accelerator() - - # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs - lr = config["lr"] - num_epochs = int(config["num_epochs"]) - seed = int(config["seed"]) - batch_size = int(config["batch_size"]) - model_name = args.model_name_or_path - - set_seed(seed) - train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size, model_name, args.n_train, args.n_val) - - # Instantiate the model (we build the model here so that the seed also control new weights initialization) - model = AutoModelForSequenceClassification.from_pretrained(model_name, return_dict=True) - - # Instantiate optimizer - optimizer_cls = ( - AdamW - if accelerator.state.deepspeed_plugin is None - or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config - else DummyOptim - ) - optimizer = optimizer_cls(params=model.parameters(), lr=lr) - - if accelerator.state.deepspeed_plugin is not None: - gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[ - "gradient_accumulation_steps" - ] - else: - gradient_accumulation_steps = 1 - max_training_steps = (len(train_dataloader) * num_epochs) // gradient_accumulation_steps - - # Instantiate scheduler - if ( - accelerator.state.deepspeed_plugin is None - or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config - ): - lr_scheduler = get_linear_schedule_with_warmup( - optimizer=optimizer, - num_warmup_steps=0, - num_training_steps=max_training_steps, - ) - else: - lr_scheduler = DummyScheduler(optimizer, total_num_steps=max_training_steps, warmup_num_steps=0) - - # Prepare everything - # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the - # prepare method. - model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare( - model, optimizer, train_dataloader, eval_dataloader, lr_scheduler - ) - - # We need to keep track of how many total steps we have iterated over - overall_step = 0 - # We also need to keep track of the stating epoch so files are named properly - starting_epoch = 0 - - # Now we train the model - train_total_peak_memory = {} - for epoch in range(starting_epoch, num_epochs): - with TorchTracemalloc() as tracemalloc: - model.train() - for step, batch in enumerate(train_dataloader): - outputs = model(**batch) - loss = outputs.loss - loss = loss / gradient_accumulation_steps - accelerator.backward(loss) - if step % gradient_accumulation_steps == 0: - optimizer.step() - lr_scheduler.step() - optimizer.zero_grad() - - overall_step += 1 - - # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage - accelerator.print(f"Memory before entering the train : {b2mb(tracemalloc.begin)}") - accelerator.print(f"Memory consumed at the end of the train (end-begin): {tracemalloc.used}") - accelerator.print(f"Peak Memory consumed during the train (max-begin): {tracemalloc.peaked}") - accelerator.print( - f"Total Peak Memory consumed during the train (max): {tracemalloc.peaked + b2mb(tracemalloc.begin)}" - ) - train_total_peak_memory[f"epoch-{epoch}"] = tracemalloc.peaked + b2mb(tracemalloc.begin) - if args.peak_memory_upper_bound is not None: - assert train_total_peak_memory[f"epoch-{epoch}"] <= args.peak_memory_upper_bound, ( - "Peak memory usage exceeded the upper bound" - ) - - accelerator.wait_for_everyone() - if accelerator.is_main_process: - with open(os.path.join(args.output_dir, "peak_memory_utilization.json"), "w") as f: - json.dump(train_total_peak_memory, f) - accelerator.end_training() - - -def main(): - parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.") - parser.add_argument( - "--model_name_or_path", - type=str, - default="bert-base-cased", - help="Path to pretrained model or model identifier from huggingface.co/models.", - required=False, - ) - parser.add_argument( - "--output_dir", - type=str, - default=".", - help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.", - ) - parser.add_argument( - "--peak_memory_upper_bound", - type=float, - default=None, - help="The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value.", - ) - parser.add_argument( - "--n_train", - type=int, - default=320, - help="Number of training examples to use.", - ) - parser.add_argument( - "--n_val", - type=int, - default=160, - help="Number of validation examples to use.", - ) - parser.add_argument( - "--num_epochs", - type=int, - default=1, - help="Number of train epochs.", - ) - args = parser.parse_args() - config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16} - training_function(config, args) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_performance.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_performance.py deleted file mode 100644 index d4396ec4d6f7ce23db3678930c2de22af39681ca..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_performance.py +++ /dev/null @@ -1,298 +0,0 @@ -# Copyright 2022 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 json -import os -from contextlib import nullcontext -from pathlib import Path - -import evaluate -import torch -from datasets import load_dataset -from torch.optim import AdamW -from torch.utils.data import DataLoader -from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup - -from accelerate import Accelerator, DistributedType -from accelerate.utils import SAFE_WEIGHTS_NAME, TorchTensorParallelPlugin, set_seed -from accelerate.utils.deepspeed import DummyOptim, DummyScheduler - - -MAX_GPU_BATCH_SIZE = 16 -EVAL_BATCH_SIZE = 32 - - -def get_dataloaders(accelerator: Accelerator, batch_size: int = 16, model_name: str = "bert-base-cased"): - """ - Creates a set of `DataLoader`s for the `glue` dataset. - - Args: - accelerator (`Accelerator`): - An `Accelerator` object - batch_size (`int`, *optional*): - The batch size for the train and validation DataLoaders. - model_name (`str`, *optional*): - """ - tokenizer = AutoTokenizer.from_pretrained(model_name) - - datasets = load_dataset("glue", "mrpc") - - def tokenize_function(examples): - # max_length=None => use the model max length (it's actually the default) - outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) - return outputs - - # Apply the method we just defined to all the examples in all the splits of the dataset - tokenized_datasets = datasets.map( - tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False - ) - - # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the - # transformers library - tokenized_datasets = tokenized_datasets.rename_column("label", "labels") - - def collate_fn(examples): - # On TPU it's best to pad everything to the same length or training will be very slow. - if accelerator.distributed_type == DistributedType.XLA: - return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt") - return tokenizer.pad(examples, padding="longest", return_tensors="pt") - - # Instantiate dataloaders. - train_dataloader = DataLoader( - tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size - ) - eval_dataloader = DataLoader( - tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE - ) - - return train_dataloader, eval_dataloader - - -def training_function(config, args): - accelerator_kwargs = {} - # need this for DeepSpeed tests as `args.tp_size` would be None and `torch.distributed.init_device_mesh` would fail - if args.tp_size is not None: - accelerator_kwargs["torch_tp_plugin"] = TorchTensorParallelPlugin(tp_size=args.tp_size) - - # Initialize accelerator - accelerator = Accelerator(**accelerator_kwargs) - - # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs - lr = config["lr"] - num_epochs = int(config["num_epochs"]) - seed = int(config["seed"]) - batch_size = int(config["batch_size"]) - model_name = args.model_name_or_path - - set_seed(seed) - train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size, model_name) - - # Add TP related kwargs if provided - model_kwargs = {} - if args.tp_plan is not None: - model_kwargs["tp_plan"] = args.tp_plan - if args.tp_size is not None: - model_kwargs["tp_size"] = args.tp_size - - # Instantiate the model (we build the model here so that the seed also control new weights initialization) - model = AutoModelForSequenceClassification.from_pretrained(model_name, return_dict=True, **model_kwargs) - - if args.add_pad_token: - if model.config.pad_token_id is None: - model.config.pad_token_id = 0 - - # Instantiate optimizer - optimizer_cls = ( - AdamW - if accelerator.state.deepspeed_plugin is None - or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config - else DummyOptim - ) - optimizer = optimizer_cls(params=model.parameters(), lr=lr) - - max_training_steps = len(train_dataloader) * num_epochs - - # Instantiate scheduler - linear_decay_scheduler = False - if ( - accelerator.state.deepspeed_plugin is None - or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config - ): - lr_scheduler = get_linear_schedule_with_warmup( - optimizer=optimizer, - num_warmup_steps=0, - num_training_steps=max_training_steps, - ) - linear_decay_scheduler = True - else: - lr_scheduler = DummyScheduler(optimizer, total_num_steps=max_training_steps, warmup_num_steps=0) - - # Prepare everything - # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the - # prepare method. - model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare( - model, optimizer, train_dataloader, eval_dataloader, lr_scheduler - ) - - # We also need to keep track of the stating epoch so files are named properly - starting_epoch = 0 - - # Now we train the model - metric = evaluate.load("glue", "mrpc") - best_performance = 0 - performance_metric = {} - expected_lr_after_first_optim_step = lr * ( - 1 - 1 / (max_training_steps / accelerator.num_processes / accelerator.gradient_accumulation_steps) - ) - lr_scheduler_check_completed = False - for epoch in range(starting_epoch, num_epochs): - model.train() - for step, batch in enumerate(train_dataloader): - with accelerator.accumulate(model): - outputs = model(**batch) - loss = outputs.loss - accelerator.backward(loss) - context = nullcontext - if args.tp_plan is not None: - from torch.distributed._tensor.experimental import implicit_replication - - context = implicit_replication - with context(): - optimizer.step() - lr_scheduler.step() - optimizer.zero_grad() - - # assert the learning rate after first optimizer step - if ( - accelerator.sync_gradients - and not lr_scheduler_check_completed - and linear_decay_scheduler - and accelerator.state.mixed_precision == "no" - ): - assert lr_scheduler.get_last_lr()[0] == expected_lr_after_first_optim_step, ( - f"Wrong lr found at second step, expected {expected_lr_after_first_optim_step}, got {lr_scheduler.get_last_lr()[0]}" - ) - lr_scheduler_check_completed = True - - model.eval() - samples_seen = 0 - for step, batch in enumerate(eval_dataloader): - # We could avoid this line since we set the accelerator with `device_placement=True`. - batch.to(accelerator.device) - with torch.no_grad(): - outputs = model(**batch) - predictions = outputs.logits.argmax(dim=-1) - # It is slightly faster to call this once, than multiple times - predictions, references = accelerator.gather( - (predictions, batch["labels"]) - ) # If we are in a multiprocess environment, the last batch has duplicates - if accelerator.use_distributed: - if step == len(eval_dataloader) - 1: - predictions = predictions[: len(eval_dataloader.dataset) - samples_seen] - references = references[: len(eval_dataloader.dataset) - samples_seen] - else: - samples_seen += references.shape[0] - metric.add_batch( - predictions=predictions, - references=references, - ) - - eval_metric = metric.compute() - # Use accelerator.print to print only on the main process. - accelerator.print(f"epoch {epoch}:", eval_metric) - performance_metric[f"epoch-{epoch}"] = eval_metric["accuracy"] - - if best_performance < eval_metric["accuracy"]: - best_performance = eval_metric["accuracy"] - - # check that the LR is 0 - if linear_decay_scheduler and accelerator.state.mixed_precision == "no": - assert lr_scheduler.get_last_lr()[0] == 0, ( - f"Wrong lr found at last step, expected 0, got {lr_scheduler.get_last_lr()[0]}" - ) - - if args.performance_lower_bound is not None: - assert args.performance_lower_bound <= best_performance, ( - f"Best performance metric {best_performance} is lower than the lower bound {args.performance_lower_bound}" - ) - - accelerator.wait_for_everyone() - if accelerator.is_main_process: - with open(os.path.join(args.output_dir, "all_results.json"), "w") as f: - json.dump(performance_metric, f) - - # TODO: skip saving of the model test for TP until the feature lands - if args.tp_plan is None: - # Finally try saving the model - accelerator.save_model(model, args.output_dir) - accelerator.wait_for_everyone() - if args.tp_plan is None: - assert Path(args.output_dir, SAFE_WEIGHTS_NAME).exists(), ( - "Model was not saved when calling `Accelerator.save_model`" - ) - accelerator.end_training() - - -def main(): - parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.") - parser.add_argument( - "--model_name_or_path", - type=str, - default="bert-base-cased", - help="Path to pretrained model or model identifier from huggingface.co/models.", - required=False, - ) - parser.add_argument( - "--output_dir", - type=str, - default=".", - help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.", - ) - parser.add_argument( - "--performance_lower_bound", - type=float, - default=None, - help="Optional lower bound for the performance metric. If set, the training will throw error when the performance metric drops below this value.", - ) - parser.add_argument( - "--num_epochs", - type=int, - default=3, - help="Number of train epochs.", - ) - parser.add_argument( - "--add_pad_token", - type=bool, - default=False, - help="To add pad token if not exists.", - ) - parser.add_argument( - "--tp_plan", - type=str, - default=None, - help="pass 'auto' to use TP", - ) - parser.add_argument( - "--tp_size", - type=int, - default=None, - help="TP size to be used to shard the model", - ) - args = parser.parse_args() - config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16} - training_function(config, args) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_pippy.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_pippy.py deleted file mode 100644 index 1dbd86c46b4a0c12df8ea4d736c7cd1e03f81813..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_pippy.py +++ /dev/null @@ -1,117 +0,0 @@ -# Copyright 2024 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 torch -from transformers import ( - BertConfig, - BertForMaskedLM, - GPT2Config, - GPT2ForSequenceClassification, -) - -from accelerate import PartialState -from accelerate.inference import prepare_pippy -from accelerate.test_utils import torch_device -from accelerate.utils import DistributedType, set_seed - - -model_to_config = { - "bert": (BertForMaskedLM, BertConfig, 512), - "gpt2": (GPT2ForSequenceClassification, GPT2Config, 1024), -} - - -def get_model_and_data_for_text(model_name, device, num_processes: int = 2): - initializer, config, seq_len = model_to_config[model_name] - config_args = {} - # Eventually needed for batch inference tests on gpt-2 when bs != 1 - # if model_name == "gpt2": - # config_args["pad_token_id"] = 0 - model_config = config(**config_args) - model = initializer(model_config) - kwargs = dict(low=0, high=model_config.vocab_size, device=device, dtype=torch.int64, requires_grad=False) - trace_input = torch.randint(size=(1, seq_len), **kwargs) - inference_inputs = torch.randint(size=(num_processes, seq_len), **kwargs) - return model, trace_input, inference_inputs - - -def test_bert(batch_size: int = 2): - set_seed(42) - state = PartialState() - model, trace_input, inference_inputs = get_model_and_data_for_text("bert", "cpu", batch_size) - model = prepare_pippy(model, example_args=(trace_input,), no_split_module_classes=model._no_split_modules) - # For inference args need to be a tuple - inputs = inference_inputs.to(torch_device) - with torch.no_grad(): - output = model(inputs) - # Zach: Check that we just grab the real outputs we need at the end - if not state.is_last_process: - assert output is None, "Output was not generated on just the last process!" - else: - assert output is not None, "Output was not generated in the last process!" - - -def test_gpt2(batch_size: int = 2): - set_seed(42) - state = PartialState() - model, trace_input, inference_inputs = get_model_and_data_for_text("gpt2", "cpu", batch_size) - model = prepare_pippy(model, example_args=(trace_input,), no_split_module_classes=model._no_split_modules) - # For inference args need to be a tuple - inputs = inference_inputs.to(torch_device) - with torch.no_grad(): - output = model(inputs) - # Zach: Check that we just grab the real outputs we need at the end - if not state.is_last_process: - assert output is None, "Output was not generated on just the last process!" - else: - assert output is not None, "Output was not generated in the last process!" - - -# Currently disabled, enable again once PyTorch pippy interface can trace a resnet34 -# def test_resnet(batch_size: int = 2): -# set_seed(42) -# state = PartialState() -# model = resnet34() -# input_tensor = torch.rand(1, 3, 224, 224) -# model = prepare_pippy( -# model, -# example_args=(input_tensor,), -# ) -# inference_inputs = torch.rand(batch_size, 3, 224, 224) -# inputs = send_to_device(inference_inputs, torch_device) -# with torch.no_grad(): -# output = model(inputs) -# # Zach: Check that we just grab the real outputs we need at the end -# if not state.is_last_process: -# assert output is None, "Output was not generated on just the last process!" -# else: -# assert output is not None, "Output was not generated in the last process!" - - -if __name__ == "__main__": - state = PartialState() - state.print("Testing pippy integration...") - try: - if state.distributed_type in [DistributedType.MULTI_GPU, DistributedType.MULTI_XPU, DistributedType.MULTI_HPU]: - state.print("Testing GPT2...") - test_gpt2() - # Issue: When modifying the tokenizer for batch GPT2 inference, there's an issue - # due to references - # NameError: cannot access free variable 'chunk_args_list' where it is not associated with a value in enclosing scope - # test_gpt2(3) - state.print("Testing BERT...") - test_bert() - else: - print("Less than two GPUs found, not running tests!") - finally: - state.destroy_process_group() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_zero3_integration.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_zero3_integration.py deleted file mode 100644 index f5352b19980288115d2442229620105b8440d03d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/external_deps/test_zero3_integration.py +++ /dev/null @@ -1,59 +0,0 @@ -# Copyright 2024 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.distributed - -from accelerate.test_utils import require_huggingface_suite, torch_device -from accelerate.utils import is_transformers_available - - -if is_transformers_available(): - from transformers import AutoModel, TrainingArguments - - -GPT2_TINY = "sshleifer/tiny-gpt2" - - -@require_huggingface_suite -def init_torch_dist_then_launch_deepspeed(): - if torch_device == "xpu": - backend = "ccl" - elif torch_device == "hpu": - backend = "hccl" - else: - backend = "nccl" - - torch.distributed.init_process_group(backend=backend) - deepspeed_config = { - "zero_optimization": { - "stage": 3, - }, - "train_batch_size": "auto", - "train_micro_batch_size_per_gpu": "auto", - } - train_args = TrainingArguments( - output_dir="./", - deepspeed=deepspeed_config, - ) - model = AutoModel.from_pretrained(GPT2_TINY) - assert train_args is not None - assert model is not None - - -def main(): - init_torch_dist_then_launch_deepspeed() - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_cli.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_cli.py deleted file mode 100644 index fc9dd1d36e8f2949d6e3cebb8ce65efbf0b9e5e4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_cli.py +++ /dev/null @@ -1,32 +0,0 @@ -# Copyright 2022 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 accelerate.utils import is_xpu_available - - -def main(): - accelerator_type = "GPU" - num_accelerators = 0 - if torch.cuda.is_available(): - num_accelerators = torch.cuda.device_count() - accelerator_type = "GPU" - elif is_xpu_available(): - num_accelerators = torch.xpu.device_count() - accelerator_type = "XPU" - print(f"Successfully ran on {num_accelerators} {accelerator_type}s") - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_ddp_comm_hook.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_ddp_comm_hook.py deleted file mode 100644 index 0db5844e026d1c035670e518a8f81d33136ea665..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_ddp_comm_hook.py +++ /dev/null @@ -1,85 +0,0 @@ -# Copyright 2022 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 accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs, PartialState -from accelerate.utils import is_hpu_available - - -class MockModel(torch.nn.Module): - def __init__(self): - super().__init__() - torch.manual_seed(0) - self.p = torch.nn.Parameter(torch.randn(40, 20)) - - def forward(self, x, rank): - return self.p * (x ** (1 + rank)) - - -def _run_and_get_grads(model, rank): - torch.manual_seed(2024) - input = torch.randn(40, 20) - output = model(input, rank) - output.mean().backward() - param = next(model.parameters()) - return param.grad - - -def test_ddp_comm_hook(comm_hook, comm_wrapper, comm_state_option): - ddp_kwargs = DistributedDataParallelKwargs( - comm_hook=comm_hook, - comm_wrapper=comm_wrapper, - comm_state_option=comm_state_option, - ) - accelerator = Accelerator(kwargs_handlers=[ddp_kwargs]) - - model = accelerator.prepare(MockModel()) - hook_grads = _run_and_get_grads(model, accelerator.local_process_index) - - reference_model = torch.nn.parallel.DistributedDataParallel( - MockModel().to(accelerator.device), - device_ids=[accelerator.local_process_index], - output_device=accelerator.local_process_index, - ) - reference_grads = _run_and_get_grads(reference_model, accelerator.local_process_index) - - torch.testing.assert_close(hook_grads, reference_grads, rtol=1e-2, atol=1e-2) - - -def main(): - for comm_hook, comm_wrapper, comm_state_option in [ - (DDPCommunicationHookType.NO, DDPCommunicationHookType.NO, {}), - (DDPCommunicationHookType.FP16, DDPCommunicationHookType.NO, {}), - (DDPCommunicationHookType.BF16, DDPCommunicationHookType.NO, {}), - (DDPCommunicationHookType.POWER_SGD, DDPCommunicationHookType.NO, {}), - (DDPCommunicationHookType.POWER_SGD, DDPCommunicationHookType.FP16, {}), - (DDPCommunicationHookType.POWER_SGD, DDPCommunicationHookType.BF16, {}), - (DDPCommunicationHookType.POWER_SGD, DDPCommunicationHookType.NO, {"matrix_approximation_rank": 2}), - (DDPCommunicationHookType.BATCHED_POWER_SGD, DDPCommunicationHookType.NO, {}), - (DDPCommunicationHookType.BATCHED_POWER_SGD, DDPCommunicationHookType.FP16, {}), - (DDPCommunicationHookType.BATCHED_POWER_SGD, DDPCommunicationHookType.BF16, {}), - ]: - if is_hpu_available(): - HPU_UNSUPPORTED_COMM_HOOKS = {DDPCommunicationHookType.FP16, DDPCommunicationHookType.BF16} - if comm_hook in HPU_UNSUPPORTED_COMM_HOOKS or comm_wrapper in HPU_UNSUPPORTED_COMM_HOOKS: - print(f"Skipping test DDP comm hook: {comm_hook}, comm wrapper: {comm_wrapper} on HPU") - continue - - print(f"Test DDP comm hook: {comm_hook}, comm wrapper: {comm_wrapper}") - test_ddp_comm_hook(comm_hook, comm_wrapper, comm_state_option) - PartialState().destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_distributed_data_loop.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_distributed_data_loop.py deleted file mode 100644 index 08cbbeb844bcc3189c99d328580c251aeafc2052..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_distributed_data_loop.py +++ /dev/null @@ -1,410 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2021 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 pickle -import tempfile -import warnings -from unittest.mock import Mock - -import torch -from torch.utils.data import ( - BatchSampler, - DataLoader, - Dataset, - IterableDataset, - RandomSampler, - TensorDataset, - default_collate, -) - -from accelerate.accelerator import Accelerator, DataLoaderConfiguration -from accelerate.utils.dataclasses import DistributedType - - -NUM_ELEMENTS = 22 -NUM_WORKERS = 4 -BATCH_SIZE = 4 - - -class DummyDataset(Dataset): - def __len__(self): - return NUM_ELEMENTS - - def __getitem__(self, index): - squeeze = False - - if isinstance(index, int): - index = [index] - squeeze = True - elif isinstance(index, slice): - index = list(range(*index.indices(self.size))) - else: - index = list(index) - - batch = [{"index": i, "label": i % 2, "random_augmentation": torch.rand(1).item()} for i in index] - - if squeeze: - batch = batch[0] - - return batch - - -class DummyIterableDataset(IterableDataset): - def __init__(self, data): - self.data = data - - def __iter__(self): - yield from self.data - - -def create_accelerator(even_batches=True): - dataloader_config = DataLoaderConfiguration(even_batches=even_batches) - accelerator = Accelerator(dataloader_config=dataloader_config) - assert accelerator.num_processes == 2, "this script expects that two GPUs are available" - return accelerator - - -def create_dataloader( - accelerator: Accelerator, dataset_size: int, batch_size: int, iterable: bool = False, shuffle: bool = False -): - """ - Create a simple DataLoader to use during the test cases - """ - values = torch.as_tensor(range(dataset_size)) - if shuffle: - values = values[torch.randperm(values.size(0))] - if iterable: - dataset = DummyIterableDataset(values) - else: - dataset = TensorDataset(torch.as_tensor(range(dataset_size))) - - dl = DataLoader(dataset, batch_size=batch_size) - dl = accelerator.prepare(dl) - - return dl - - -def verify_dataloader_batch_sizes( - accelerator: Accelerator, - dataset_size: int, - batch_size: int, - process_0_expected_batch_sizes: list[int], - process_1_expected_batch_sizes: list[int], -): - """ - A helper function for verifying the batch sizes coming from a prepared dataloader in each process - """ - dl = create_dataloader(accelerator=accelerator, dataset_size=dataset_size, batch_size=batch_size) - - batch_sizes = [len(batch[0]) for batch in dl] - - if accelerator.process_index == 0: - assert batch_sizes == process_0_expected_batch_sizes - elif accelerator.process_index == 1: - assert batch_sizes == process_1_expected_batch_sizes - - -def test_default_ensures_even_batch_sizes(): - accelerator = create_accelerator() - - # without padding, we would expect a different number of batches - verify_dataloader_batch_sizes( - accelerator, - dataset_size=3, - batch_size=1, - process_0_expected_batch_sizes=[1, 1], - process_1_expected_batch_sizes=[1, 1], - ) - - # without padding, we would expect the same number of batches, but different sizes - verify_dataloader_batch_sizes( - accelerator, - dataset_size=7, - batch_size=2, - process_0_expected_batch_sizes=[2, 2], - process_1_expected_batch_sizes=[2, 2], - ) - - -def test_can_disable_even_batches(): - accelerator = create_accelerator(even_batches=False) - - verify_dataloader_batch_sizes( - accelerator, - dataset_size=3, - batch_size=1, - process_0_expected_batch_sizes=[1, 1], - process_1_expected_batch_sizes=[1], - ) - - verify_dataloader_batch_sizes( - accelerator, - dataset_size=7, - batch_size=2, - process_0_expected_batch_sizes=[2, 2], - process_1_expected_batch_sizes=[2, 1], - ) - - -def test_can_join_uneven_inputs(): - accelerator = create_accelerator(even_batches=False) - - model = torch.nn.Linear(1, 1) - ddp_model = accelerator.prepare(model) - - dl = create_dataloader(accelerator, dataset_size=3, batch_size=1) - - batch_idxs = [] - with accelerator.join_uneven_inputs([ddp_model]): - for batch_idx, batch in enumerate(dl): - output = ddp_model(batch[0].float()) - loss = output.sum() - loss.backward() - batch_idxs.append(batch_idx) - - accelerator.wait_for_everyone() - - if accelerator.process_index == 0: - assert batch_idxs == [0, 1] - elif accelerator.process_index == 1: - assert batch_idxs == [0] - - -def test_join_raises_warning_for_non_ddp_distributed(accelerator): - with warnings.catch_warnings(record=True) as w: - with accelerator.join_uneven_inputs([Mock()]): - pass - - assert issubclass(w[-1].category, UserWarning) - assert "only supported for multi-GPU" in str(w[-1].message) - - -def test_join_can_override_even_batches(): - default_even_batches = True - overridden_even_batches = False - accelerator = create_accelerator(even_batches=default_even_batches) - model = torch.nn.Linear(1, 1) - ddp_model = accelerator.prepare(model) - train_dl = create_dataloader(accelerator, dataset_size=3, batch_size=1) - valid_dl = create_dataloader(accelerator, dataset_size=3, batch_size=1) - - with accelerator.join_uneven_inputs([ddp_model], even_batches=overridden_even_batches): - train_dl_overridden_value = train_dl.batch_sampler.even_batches - valid_dl_overridden_value = valid_dl.batch_sampler.even_batches - - assert train_dl_overridden_value == overridden_even_batches - assert valid_dl_overridden_value == overridden_even_batches - assert train_dl.batch_sampler.even_batches == default_even_batches - assert valid_dl.batch_sampler.even_batches == default_even_batches - - -def test_join_can_override_for_mixed_type_dataloaders(): - default_even_batches = True - overridden_even_batches = False - accelerator = create_accelerator(even_batches=default_even_batches) - model = torch.nn.Linear(1, 1) - ddp_model = accelerator.prepare(model) - create_dataloader(accelerator, dataset_size=3, batch_size=1, iterable=True) - batch_dl = create_dataloader(accelerator, dataset_size=3, batch_size=1) - - with warnings.catch_warnings(): - warnings.filterwarnings("ignore") - try: - with accelerator.join_uneven_inputs([ddp_model], even_batches=overridden_even_batches): - batch_dl_overridden_value = batch_dl.batch_sampler.even_batches - except AttributeError: - # ensure attribute error is not raised when processing iterable dl - raise AssertionError - - assert batch_dl_overridden_value == overridden_even_batches - assert batch_dl.batch_sampler.even_batches == default_even_batches - - -def test_join_raises_warning_for_iterable_when_overriding_even_batches(): - accelerator = create_accelerator() - model = torch.nn.Linear(1, 1) - ddp_model = accelerator.prepare(model) - create_dataloader(accelerator, dataset_size=3, batch_size=1, iterable=True) - - with warnings.catch_warnings(record=True) as w: - with accelerator.join_uneven_inputs([ddp_model], even_batches=False): - pass - - assert issubclass(w[-1].category, UserWarning) - assert "only supported for map-style datasets" in str(w[-1].message) - - -def test_pickle_accelerator(): - accelerator = create_accelerator() - data_loader = create_dataloader(accelerator, dataset_size=32, batch_size=4) - _ = accelerator.prepare(data_loader) - pickled_accelerator = pickle.dumps(accelerator) - unpickled_accelerator = pickle.loads(pickled_accelerator) - # TODO: Maybe this should be implemented as __eq__ for AcceleratorState? - assert accelerator.state.__dict__ == unpickled_accelerator.state.__dict__ - - -def test_data_loader(data_loader, accelerator): - # Prepare the DataLoader - data_loader = accelerator.prepare(data_loader) - - all_examples = [] - for i, batch in enumerate(data_loader): - index, _ = accelerator.gather_for_metrics((batch["index"], batch["label"])) - all_examples.extend(index.detach().cpu().numpy().tolist()) - - # Sort the examples - sorted_all_examples = sorted(all_examples) - - # Check if all elements are present in the sorted list of iterated samples - assert len(set(sorted_all_examples)) == NUM_ELEMENTS, ( - "Not all the dataset elements have been iterated in an epoch due to duplication of samples across processes." - ) - - -def test_stateful_dataloader(accelerator): - """ - Tests that a stateful dataloader can be iterated over, saved after a few batches using `load_state_dict`, and then - resumed from the saved state. - - The result should be the same as the rest of the data that iterated over after saving. - """ - old_dataloader_config = accelerator.dataloader_config - try: - accelerator.dataloader_config = DataLoaderConfiguration(use_stateful_dataloader=True) - prepared_dl = create_dataloader( - accelerator, dataset_size=32 * accelerator.num_processes, batch_size=4, iterable=True, shuffle=True - ) - untrained_batches = [] - # Calculate what step that will be - total_batches = 32 * accelerator.num_processes // (4 * accelerator.num_processes) - last_batch_num = total_batches - 1 - for step, batch in enumerate(prepared_dl): - # Step just before - if step == last_batch_num - 1: - state_dict = prepared_dl.state_dict() - if step >= last_batch_num: - # Otherwise grab the "unseen" batches - untrained_batches.append(batch) - not_skipped_batches = accelerator.gather(untrained_batches) - prepared_dl.load_state_dict(state_dict) - resumed_batches = [] - for batch in prepared_dl: - resumed_batches.append(batch) - resumed_batches = accelerator.gather(resumed_batches) - for b1, b2 in zip(not_skipped_batches, resumed_batches): - for v1, v2 in zip(b1, b2): - assert torch.equal(v1, v2), f"Batch {b1} and {b2} are not equal" - finally: - accelerator.dataloader_config = old_dataloader_config - - -def test_stateful_dataloader_save_state(accelerator): - """ - Tests that a stateful dataloader can be iterated over, saved after a few batches using `Accelerator.save_state`, - and then resumed from the saved state. - - The result should be the same as the rest of the data that iterated over after saving. - """ - old_dataloader_config = accelerator.dataloader_config - try: - with tempfile.TemporaryDirectory() as tmpdir: - accelerator.dataloader_config = DataLoaderConfiguration(use_stateful_dataloader=True) - prepared_dl = create_dataloader( - accelerator, dataset_size=32 * accelerator.num_processes, batch_size=4, iterable=True, shuffle=True - ) - untrained_batches = [] - # Calculate what step that will be - total_batches = 32 * accelerator.num_processes // (4 * accelerator.num_processes) - last_batch_num = total_batches - 1 - for step, batch in enumerate(prepared_dl): - # Step just before - if step == last_batch_num - 1: - accelerator.save_state(tmpdir) - if step >= last_batch_num: - # Otherwise grab the "unseen" batches - untrained_batches.append(batch) - not_skipped_batches = accelerator.gather(untrained_batches) - accelerator.load_state(tmpdir) - resumed_batches = [] - for batch in prepared_dl: - resumed_batches.append(batch) - resumed_batches = accelerator.gather(resumed_batches) - for b1, b2 in zip(not_skipped_batches, resumed_batches): - for v1, v2 in zip(b1, b2): - assert torch.equal(v1, v2), f"Batch {b1} and {b2} are not equal" - finally: - accelerator.dataloader_config = old_dataloader_config - - -def main(): - accelerator = create_accelerator() - torch.manual_seed(accelerator.process_index) - - accelerator.print("Test that even_batches variable ensures uniform batches across processes") - test_default_ensures_even_batch_sizes() - - accelerator.print("Run tests with even_batches disabled") - test_can_disable_even_batches() - - accelerator.print("Test joining uneven inputs") - test_can_join_uneven_inputs() - - accelerator.print("Test overriding even_batches when joining uneven inputs") - test_join_can_override_even_batches() - - accelerator.print("Test overriding even_batches for mixed dataloader types") - test_join_can_override_for_mixed_type_dataloaders() - - accelerator.print("Test overriding even_batches raises a warning for iterable dataloaders") - test_join_raises_warning_for_iterable_when_overriding_even_batches() - - accelerator.print("Test join with non DDP distributed raises warning") - original_state = accelerator.state.distributed_type - accelerator.state.distributed_type = DistributedType.FSDP - test_join_raises_warning_for_non_ddp_distributed(accelerator) - accelerator.state.distributed_type = original_state - - accelerator.print("Test pickling an accelerator") - test_pickle_accelerator() - - dataset = DummyDataset() - - accelerator.print("Test DataLoader with shuffle=False") - loader = DataLoader(dataset, shuffle=False, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS) - test_data_loader(loader, accelerator) - - accelerator.print("Test DataLoader with shuffle=True") - loader = DataLoader(dataset, shuffle=True, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS) - test_data_loader(loader, accelerator) - - accelerator.print("Test DataLoader with batch_sampler") - sampler = BatchSampler(RandomSampler(dataset), batch_size=BATCH_SIZE, drop_last=False) - loader = DataLoader(dataset, batch_sampler=sampler, num_workers=NUM_WORKERS) - test_data_loader(loader, accelerator) - - accelerator.print("Test DataLoader with sampler as an instance of `BatchSampler`") - sampler = BatchSampler(RandomSampler(dataset), batch_size=BATCH_SIZE, drop_last=False) - loader = DataLoader(dataset, sampler=sampler, batch_size=None, collate_fn=default_collate, num_workers=NUM_WORKERS) - test_data_loader(loader, accelerator) - test_stateful_dataloader(accelerator) - test_stateful_dataloader_save_state(accelerator) - - accelerator.end_training() - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_merge_weights.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_merge_weights.py deleted file mode 100644 index f280c8fa17919367001389a8efcd78faaa041f0c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_merge_weights.py +++ /dev/null @@ -1,158 +0,0 @@ -# Copyright 2024 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 logging -import shutil -from pathlib import Path - -import torch -from safetensors.torch import load_file -from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy, StateDictType -from torch.utils.data import DataLoader - -from accelerate import Accelerator, FullyShardedDataParallelPlugin -from accelerate.commands.merge import merge_command, merge_command_parser -from accelerate.state import AcceleratorState -from accelerate.test_utils import torch_device -from accelerate.test_utils.training import RegressionDataset -from accelerate.utils import merge_fsdp_weights, patch_environment, save_fsdp_model - - -logging.basicConfig(level=logging.INFO) - -parser = merge_command_parser() - - -class TinyModel(torch.nn.Module): - def __init__(self): - super().__init__() - self.linear1 = torch.nn.Linear(16, 16) - self.activation = torch.nn.ReLU() - self.linear2 = torch.nn.Linear(16, 16) - self.softmax = torch.nn.Softmax() - - def forward(self, x): - return self.linear2(self.activation(self.linear1(x))) - - -def setup(): - if AcceleratorState._shared_state != {}: - AcceleratorState()._reset_state() - plugin = FullyShardedDataParallelPlugin( - sharding_strategy=ShardingStrategy.FULL_SHARD, state_dict_type=StateDictType.SHARDED_STATE_DICT - ) - model = TinyModel() - with patch_environment(fsdp_auto_wrap_policy="SIZE_BASED_WRAP"): - plugin.set_auto_wrap_policy(model) - accelerator = Accelerator(fsdp_plugin=plugin) - model = accelerator.prepare(model) - return model, plugin, accelerator - - -def mock_training(accelerator, model): - train_set = RegressionDataset(length=128, seed=42) - train_dl = DataLoader(train_set, batch_size=16, shuffle=False) - optimizer = torch.optim.SGD(model.parameters(), lr=0.1) - - train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer) - for _ in range(3): - for batch in train_dl: - model.zero_grad() - output = model(batch["x"]) - loss = torch.nn.functional.mse_loss(output, batch["y"]) - accelerator.backward(loss) - optimizer.step() - return model - - -def check_weights(operation, state_1, state_2): - for weight_1, weight_2 in zip(state_1.values(), state_2.values()): - if operation == "same": - assert torch.allclose(weight_1, weight_2) - else: - assert not torch.allclose(weight_1, weight_2) - - -def check_safetensors_weights(path, model): - safe_state_dict = load_file(path / "model.safetensors") - safe_loaded_model = TinyModel().to(torch_device) - check_weights("diff", model.state_dict(), safe_loaded_model.state_dict()) - safe_loaded_model.load_state_dict(safe_state_dict) - check_weights("same", model.state_dict(), safe_loaded_model.state_dict()) - - -def check_pytorch_weights(path, model): - nonsafe_state_dict = torch.load(path / "pytorch_model.bin", weights_only=True) - nonsafe_loaded_model = TinyModel().to(torch_device) - check_weights("diff", model.state_dict(), nonsafe_loaded_model.state_dict()) - nonsafe_loaded_model.load_state_dict(nonsafe_state_dict) - check_weights("same", model.state_dict(), nonsafe_loaded_model.state_dict()) - - -def test_merge_weights_safetensors(model, path): - # Should now be saved at `path/merged.safetensors` - merge_fsdp_weights(path / "pytorch_model_fsdp_0", path, safe_serialization=True) - check_safetensors_weights(path, model) - - -def test_merge_weights_command_safetensors(model, path): - args = parser.parse_args([str(path / "pytorch_model_fsdp_0"), str(path)]) - merge_command(args) - check_safetensors_weights(path, model) - - -def test_merge_weights_pytorch(model, path): - # Should now be saved at `path/merged.bin` - merge_fsdp_weights(path / "pytorch_model_fsdp_0", path, safe_serialization=False) - check_pytorch_weights(path, model) - - -def test_merge_weights_command_pytorch(model, path): - args = parser.parse_args([str(path / "pytorch_model_fsdp_0"), str(path), "--unsafe_serialization"]) - merge_command(args) - check_pytorch_weights(path, model) - - -if __name__ == "__main__": - # Note this test requires at least two accelerators! - model, plugin, accelerator = setup() - if accelerator.num_processes > 1: - try: - # Initial setup for things - out_path = Path("test_merge_weights_fsdp_weights") - if not out_path.exists(): - out_path.mkdir(parents=True, exist_ok=True) - - # Train briefly once weights aren't the baseline - model = mock_training(accelerator, model) - accelerator.wait_for_everyone() - - gc.collect() # Needed for some lingering refs after training - save_fsdp_model(plugin, accelerator, model, out_path) - accelerator.wait_for_everyone() - - # Finally we can test - test_merge_weights_safetensors(model, out_path) - test_merge_weights_command_safetensors(model, out_path) - test_merge_weights_pytorch(model, out_path) - test_merge_weights_command_pytorch(model, out_path) - except Exception: - raise - finally: - # Cleanup in case of any failures - if accelerator.is_main_process: - shutil.rmtree(out_path) - accelerator.wait_for_everyone() - accelerator.end_training() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_notebook.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_notebook.py deleted file mode 100644 index 267c11b50b22250e781f94e3643b8895cc6aeb02..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_notebook.py +++ /dev/null @@ -1,118 +0,0 @@ -# Copyright 2022 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. -""" -Test file to ensure that in general certain situational setups for notebooks work. -""" - -import os -import time -from multiprocessing import Queue - -from pytest import mark, raises -from torch.distributed.elastic.multiprocessing.errors import ChildFailedError - -from accelerate import PartialState, notebook_launcher -from accelerate.test_utils import require_bnb -from accelerate.utils import is_bnb_available - - -def basic_function(): - # Just prints the PartialState - print(f"PartialState:\n{PartialState()}") - - -def tough_nut_function(queue: Queue): - if queue.empty(): - return - trial = queue.get() - if trial > 0: - queue.put(trial - 1) - raise RuntimeError("The nut hasn't cracked yet! Try again.") - - print(f"PartialState:\n{PartialState()}") - - -def bipolar_sleep_function(sleep_sec: int): - state = PartialState() - if state.process_index % 2 == 0: - raise RuntimeError("I'm an even process. I don't like to sleep.") - else: - time.sleep(sleep_sec) - - -NUM_PROCESSES = int(os.environ.get("ACCELERATE_NUM_PROCESSES", 1)) - - -def test_can_initialize(): - notebook_launcher(basic_function, (), num_processes=NUM_PROCESSES) - - -@mark.skipif(NUM_PROCESSES < 2, reason="Need at least 2 processes to test static rendezvous backends") -def test_static_rdzv_backend(): - notebook_launcher(basic_function, (), num_processes=NUM_PROCESSES, rdzv_backend="static") - - -@mark.skipif(NUM_PROCESSES < 2, reason="Need at least 2 processes to test c10d rendezvous backends") -def test_c10d_rdzv_backend(): - notebook_launcher(basic_function, (), num_processes=NUM_PROCESSES, rdzv_backend="c10d") - - -@mark.skipif(NUM_PROCESSES < 2, reason="Need at least 2 processes to test fault tolerance") -def test_fault_tolerant(max_restarts: int = 3): - queue = Queue() - queue.put(max_restarts) - notebook_launcher(tough_nut_function, (queue,), num_processes=NUM_PROCESSES, max_restarts=max_restarts) - - -@mark.skipif(NUM_PROCESSES < 2, reason="Need at least 2 processes to test monitoring") -def test_monitoring(monitor_interval: float = 0.01, sleep_sec: int = 100): - start_time = time.time() - with raises(ChildFailedError, match="I'm an even process. I don't like to sleep."): - notebook_launcher( - bipolar_sleep_function, - (sleep_sec,), - num_processes=NUM_PROCESSES, - monitor_interval=monitor_interval, - ) - assert time.time() - start_time < sleep_sec, "Monitoring did not stop the process in time." - - -@require_bnb -def test_problematic_imports(): - with raises(RuntimeError, match="Please keep these imports"): - import bitsandbytes as bnb # noqa: F401 - - notebook_launcher(basic_function, (), num_processes=NUM_PROCESSES) - - -def main(): - print("Test basic notebook can be ran") - test_can_initialize() - print("Test static rendezvous backend") - test_static_rdzv_backend() - print("Test c10d rendezvous backend") - test_c10d_rdzv_backend() - print("Test fault tolerant") - test_fault_tolerant() - print("Test monitoring") - test_monitoring() - if is_bnb_available(): - print("Test problematic imports (bnb)") - test_problematic_imports() - if NUM_PROCESSES > 1: - PartialState().destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_ops.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_ops.py deleted file mode 100644 index f8f535d7b25a7bda527901787261591364545c09..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_ops.py +++ /dev/null @@ -1,181 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2023 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 accelerate import PartialState -from accelerate.test_utils.testing import assert_exception -from accelerate.utils.dataclasses import DistributedType -from accelerate.utils.operations import ( - DistributedOperationException, - broadcast, - copy_tensor_to_devices, - gather, - gather_object, - pad_across_processes, - reduce, -) - - -def create_tensor(state): - return (torch.arange(state.num_processes) + 1.0 + (state.num_processes * state.process_index)).to(state.device) - - -def test_gather(state): - tensor = create_tensor(state) - gathered_tensor = gather(tensor) - assert gathered_tensor.tolist() == list(range(1, state.num_processes**2 + 1)) - - -def test_gather_object(state): - # Gather objects in TorchXLA is not supported. - if state.distributed_type == DistributedType.XLA: - return - obj = [state.process_index] - gathered_obj = gather_object(obj) - assert len(gathered_obj) == state.num_processes, f"{gathered_obj}, {len(gathered_obj)} != {state.num_processes}" - assert gathered_obj == list(range(state.num_processes)), f"{gathered_obj} != {list(range(state.num_processes))}" - - -def test_gather_non_contigous(state): - # Skip this test because the 'is_contiguous' function of XLA tensor always returns True. - if state.distributed_type == DistributedType.XLA: - return - - # Create a non-contiguous tensor (enforce non-contiguity after device memory allocation) - tensor = torch.arange(12, device=state.device).view(4, 3).t() - assert not tensor.is_contiguous() - # Shouldn't error out - _ = gather(tensor) - - -def test_broadcast(state): - tensor = create_tensor(state) - broadcasted_tensor = broadcast(tensor) - assert broadcasted_tensor.shape == torch.Size([state.num_processes]) - assert broadcasted_tensor.tolist() == list(range(1, state.num_processes + 1)) - - -def test_pad_across_processes(state): - # We need to pad the tensor with one more element if we are the main process - # to ensure that we can pad - if state.is_main_process: - tensor = torch.arange(state.num_processes + 1).to(state.device) - else: - tensor = torch.arange(state.num_processes).to(state.device) - padded_tensor = pad_across_processes(tensor) - assert padded_tensor.shape == torch.Size([state.num_processes + 1]) - if not state.is_main_process: - assert padded_tensor.tolist() == list(range(0, state.num_processes)) + [0] - - -def test_reduce_sum(state): - # For now runs on only two processes - if state.num_processes != 2: - return - tensor = create_tensor(state) - reduced_tensor = reduce(tensor, "sum") - truth_tensor = torch.tensor([4.0, 6]).to(state.device) - assert torch.allclose(reduced_tensor, truth_tensor), f"{reduced_tensor} != {truth_tensor}" - - -def test_reduce_mean(state): - # For now runs on only two processes - if state.num_processes != 2: - return - tensor = create_tensor(state) - reduced_tensor = reduce(tensor, "mean") - truth_tensor = torch.tensor([2.0, 3]).to(state.device) - assert torch.allclose(reduced_tensor, truth_tensor), f"{reduced_tensor} != {truth_tensor}" - - -def test_op_checker(state): - # Must be in a distributed state, and gathering is currently not supported in TorchXLA. - if state.distributed_type in [DistributedType.NO, DistributedType.XLA]: - return - state.debug = True - # `pad_across_processes` - if state.process_index == 0: - data = {"tensor": torch.tensor([[0.0, 1, 2, 3, 4]]).to(state.device)} - else: - data = {"tensor": torch.tensor([[[0.0, 1, 2, 3, 4, 5]]]).to(state.device)} - - with assert_exception(DistributedOperationException): - pad_across_processes(data, dim=0) - - # `reduce` - if state.process_index == 0: - data = {"tensor": torch.tensor([[0.0, 1, 2, 3, 4]]).to(state.device)} - else: - data = {"tensor": torch.tensor([[[0.0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]]).to(state.device)} - - with assert_exception(DistributedOperationException): - reduce(data) - - # `broadcast` - if state.process_index == 0: - data = {"tensor": torch.tensor([[0.0, 1, 2, 3, 4]]).to(state.device)} - else: - data = {"tensor": torch.tensor([[[0.0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]]).to(state.device)} - - with assert_exception(DistributedOperationException): - broadcast(data) - - state.debug = False - - -def test_copy_tensor_to_devices(state): - if state.distributed_type not in [DistributedType.MULTI_GPU, DistributedType.XLA]: - return - if state.is_main_process: - tensor = torch.tensor([1, 2, 3], dtype=torch.int).to(state.device) - else: - tensor = None - tensor = copy_tensor_to_devices(tensor) - assert torch.allclose(tensor, torch.tensor([1, 2, 3], dtype=torch.int, device=state.device)) - - -def _mp_fn(index): - # For xla_spawn (TPUs) - main() - - -def main(): - state = PartialState() - state.print(f"State: {state}") - state.print("testing gather") - test_gather(state) - state.print("testing gather_object") - test_gather_object(state) - state.print("testing gather non-contigous") - test_gather_non_contigous(state) - state.print("testing broadcast") - test_broadcast(state) - state.print("testing pad_across_processes") - test_pad_across_processes(state) - state.print("testing reduce_sum") - test_reduce_sum(state) - state.print("testing reduce_mean") - test_reduce_mean(state) - state.print("testing op_checker") - test_op_checker(state) - state.print("testing sending tensors across devices") - test_copy_tensor_to_devices(state) - state.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_script.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_script.py deleted file mode 100644 index fa0e3b601e38057e254067984042197098efdfd6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_script.py +++ /dev/null @@ -1,952 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2021 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 contextlib -import io -import math -import time -from copy import deepcopy -from pathlib import Path - -import numpy as np -import torch -from torch.utils.data import DataLoader, Dataset - -from accelerate import Accelerator -from accelerate.data_loader import SeedableRandomSampler, prepare_data_loader -from accelerate.state import AcceleratorState -from accelerate.test_utils import RegressionDataset, are_the_same_tensors -from accelerate.utils import ( - DataLoaderConfiguration, - DistributedType, - gather, - gather_object, - is_bf16_available, - is_datasets_available, - is_fp16_available, - is_hpu_available, - is_ipex_available, - is_pytest_available, - is_xpu_available, - set_seed, - synchronize_rng_states, -) - - -# TODO: remove RegressionModel4XPU once ccl support empty buffer in broadcasting. -if is_xpu_available(): - from accelerate.test_utils import RegressionModel4XPU as RegressionModel -else: - from accelerate.test_utils import RegressionModel - -if is_hpu_available(): - ATOL = 1e-3 - RTOL = 1e-3 -else: - ATOL = 1e-6 - RTOL = 1e-6 - - -def generate_baseline_dataloader(train_set, generator, batch_size, use_seedable_sampler=False): - "Creates a dataloader that can also use the `SeedableRandomSampler`" - if use_seedable_sampler: - # The SeedableRandomSampler is needed during distributed setups - # for full reproducibility across processes with the `DataLoader` - sampler = SeedableRandomSampler( - generator=generator, - data_source=train_set, - num_samples=len(train_set), - ) - return DataLoader(train_set, batch_size=batch_size, sampler=sampler) - else: - return DataLoader(train_set, batch_size=batch_size, shuffle=True, generator=generator) - - -def print_main(state): - print(f"Printing from the main process {state.process_index}") - - -def print_local_main(state): - print(f"Printing from the local main process {state.local_process_index}") - - -def print_last(state): - print(f"Printing from the last process {state.process_index}") - - -def print_on(state, process_idx): - print(f"Printing from process {process_idx}: {state.process_index}") - - -def process_execution_check(): - accelerator = Accelerator() - num_processes = accelerator.num_processes - # Test main_process_first context manager - path = Path("check_main_process_first.txt") - with accelerator.main_process_first(): - if accelerator.is_main_process: - time.sleep(0.1) # ensure main process takes longest - with open(path, "a+") as f: - f.write("Currently in the main process\n") - else: - with open(path, "a+") as f: - f.write("Now on another process\n") - accelerator.wait_for_everyone() - - if accelerator.is_main_process: - with open(path) as f: - text = "".join(f.readlines()) - try: - assert text.startswith("Currently in the main process\n"), "Main process was not first" - if num_processes > 1: - assert text.endswith("Now on another process\n"), "Main process was not first" - assert text.count("Now on another process\n") == accelerator.num_processes - 1, ( - f"Only wrote to file {text.count('Now on another process') + 1} times, not {accelerator.num_processes}" - ) - except AssertionError: - path.unlink() - raise - - if accelerator.is_main_process and path.exists(): - path.unlink() - accelerator.wait_for_everyone() - # Test the decorators - f = io.StringIO() - with contextlib.redirect_stdout(f): - accelerator.on_main_process(print_main)(accelerator.state) - result = f.getvalue().rstrip() - if accelerator.is_main_process: - assert result == "Printing from the main process 0", f"{result} != Printing from the main process 0" - else: - assert f.getvalue().rstrip() == "", f'{result} != ""' - f.truncate(0) - f.seek(0) - - with contextlib.redirect_stdout(f): - accelerator.on_local_main_process(print_local_main)(accelerator.state) - if accelerator.is_local_main_process: - assert f.getvalue().rstrip() == "Printing from the local main process 0" - else: - assert f.getvalue().rstrip() == "" - f.truncate(0) - f.seek(0) - - with contextlib.redirect_stdout(f): - accelerator.on_last_process(print_last)(accelerator.state) - if accelerator.is_last_process: - assert f.getvalue().rstrip() == f"Printing from the last process {accelerator.state.num_processes - 1}" - else: - assert f.getvalue().rstrip() == "" - f.truncate(0) - f.seek(0) - - for process_idx in range(num_processes): - with contextlib.redirect_stdout(f): - accelerator.on_process(print_on, process_index=process_idx)(accelerator.state, process_idx) - if accelerator.process_index == process_idx: - assert f.getvalue().rstrip() == f"Printing from process {process_idx}: {accelerator.process_index}" - else: - assert f.getvalue().rstrip() == "" - f.truncate(0) - f.seek(0) - - -def init_state_check(): - # Test we can instantiate this twice in a row. - state = AcceleratorState() - if state.local_process_index == 0: - print("Testing, testing. 1, 2, 3.") - print(state) - - -def rng_sync_check(): - state = AcceleratorState() - synchronize_rng_states(["torch"]) - assert are_the_same_tensors(torch.get_rng_state()), "RNG states improperly synchronized on CPU." - if state.distributed_type == DistributedType.MULTI_GPU: - synchronize_rng_states(["cuda"]) - assert are_the_same_tensors(torch.cuda.get_rng_state()), "RNG states improperly synchronized on GPU." - elif state.distributed_type == DistributedType.MULTI_XPU: - synchronize_rng_states(["xpu"]) - assert are_the_same_tensors(torch.xpu.get_rng_state()), "RNG states improperly synchronized on XPU." - generator = torch.Generator() - synchronize_rng_states(["generator"], generator=generator) - assert are_the_same_tensors(generator.get_state()), "RNG states improperly synchronized in generator." - - if state.local_process_index == 0: - print("All rng are properly synched.") - - -def dl_preparation_check(): - state = AcceleratorState() - length = 32 * state.num_processes - - dl = DataLoader(range(length), batch_size=8) - dl = prepare_data_loader(dl, state.device, state.num_processes, state.process_index, put_on_device=True) - result = [] - for batch in dl: - result.append(gather(batch)) - result = torch.cat(result) - - assert torch.equal(result.cpu(), torch.arange(0, length).long()), "Wrong non-shuffled dataloader result." - - dl = DataLoader(range(length), batch_size=8) - dl = prepare_data_loader( - dl, - state.device, - state.num_processes, - state.process_index, - put_on_device=True, - split_batches=True, - ) - result = [] - for batch in dl: - result.append(gather(batch)) - result = torch.cat(result) - assert torch.equal(result.cpu(), torch.arange(0, length).long()), "Wrong non-shuffled dataloader result." - - if state.process_index == 0: - print("Non-shuffled dataloader passing.") - - dl = DataLoader(range(length), batch_size=8, shuffle=True) - dl = prepare_data_loader(dl, state.device, state.num_processes, state.process_index, put_on_device=True) - result = [] - for batch in dl: - result.append(gather(batch)) - result = torch.cat(result).tolist() - result.sort() - assert result == list(range(length)), "Wrong shuffled dataloader result." - - dl = DataLoader(range(length), batch_size=8, shuffle=True) - dl = prepare_data_loader( - dl, - state.device, - state.num_processes, - state.process_index, - put_on_device=True, - split_batches=True, - ) - result = [] - for batch in dl: - result.append(gather(batch)) - result = torch.cat(result).tolist() - result.sort() - assert result == list(range(length)), "Wrong shuffled dataloader result." - - if state.local_process_index == 0: - print("Shuffled dataloader passing.") - - -def central_dl_preparation_check(): - state = AcceleratorState() - length = 32 * state.num_processes - - dl = DataLoader(range(length), batch_size=8) - dl = prepare_data_loader( - dl, state.device, state.num_processes, state.process_index, put_on_device=True, dispatch_batches=True - ) - result = [] - for batch in dl: - result.append(gather(batch)) - result = torch.cat(result) - assert torch.equal(result.cpu(), torch.arange(0, length).long()), "Wrong non-shuffled dataloader result." - - dl = DataLoader(range(length), batch_size=8) - dl = prepare_data_loader( - dl, - state.device, - state.num_processes, - state.process_index, - put_on_device=True, - split_batches=True, - dispatch_batches=True, - ) - result = [] - for batch in dl: - result.append(gather(batch)) - result = torch.cat(result) - assert torch.equal(result.cpu(), torch.arange(0, length).long()), "Wrong non-shuffled dataloader result." - - if state.process_index == 0: - print("Non-shuffled central dataloader passing.") - - dl = DataLoader(range(length), batch_size=8, shuffle=True) - dl = prepare_data_loader( - dl, state.device, state.num_processes, state.process_index, put_on_device=True, dispatch_batches=True - ) - result = [] - for batch in dl: - result.append(gather(batch)) - result = torch.cat(result).tolist() - result.sort() - assert result == list(range(length)), "Wrong shuffled dataloader result." - - dl = DataLoader(range(length), batch_size=8, shuffle=True) - dl = prepare_data_loader( - dl, - state.device, - state.num_processes, - state.process_index, - put_on_device=True, - split_batches=True, - dispatch_batches=True, - ) - result = [] - for batch in dl: - result.append(gather(batch)) - result = torch.cat(result).tolist() - result.sort() - assert result == list(range(length)), "Wrong shuffled dataloader result." - - if state.local_process_index == 0: - print("Shuffled central dataloader passing.") - - -def custom_sampler_check(): - state = AcceleratorState() - - class CustomDataset(Dataset): - def __init__(self, data): - self.data = data - - def __len__(self): - return len(self.data) - - def __getitem__(self, index): - return self.data[index] - - class CustomBatchSampler: - def __init__(self, dataset_length: int, batch_size: int, shuffle: bool = True): - self.batch_size = batch_size - self.data_index = np.arange(dataset_length) - self.shuffle = shuffle - - def __iter__(self): - num_batches = len(self) - if self.shuffle: - index = np.random.permutation(self.data_index) - else: - index = self.data_index - output = np.array_split(index, num_batches) - yield from output - - def __len__(self): - return math.ceil(len(self.data_index) / self.batch_size) - - dataset = CustomDataset(range(32 * state.num_processes)) - sampler = CustomBatchSampler(len(dataset), batch_size=8) - dl = DataLoader(dataset, batch_sampler=sampler) - dl = prepare_data_loader(dl, state.device, state.num_processes, state.process_index) - # We need just ensure that `dl.batch_sampler` (or `dl.batch_sampler.batch_sampler` is indeed the old batch sampler - if hasattr(dl.batch_sampler, "batch_sampler"): - assert isinstance(dl.batch_sampler.batch_sampler, CustomBatchSampler), ( - "Custom sampler was changed after calling `prepare_data_loader`" - ) - else: - assert isinstance(dl.batch_sampler, CustomBatchSampler), ( - "Custom sampler was changed after calling `prepare_data_loader`" - ) - - -def check_seedable_sampler(): - # Set seed - set_seed(42) - train_set = RegressionDataset(length=10, seed=42) - train_dl = DataLoader(train_set, batch_size=2, shuffle=True) - - config = DataLoaderConfiguration(use_seedable_sampler=True) - accelerator = Accelerator(dataloader_config=config) - train_dl = accelerator.prepare(train_dl) - original_items = [] - for _ in range(3): - for batch in train_dl: - original_items.append(batch["x"]) - original_items = torch.cat(original_items) - - # Set seed again and the epoch - set_seed(42) - train_dl.set_epoch(0) - new_items = [] - for _ in range(3): - for batch in train_dl: - new_items.append(batch["x"]) - new_items = torch.cat(new_items) - assert torch.allclose(original_items, new_items), "Did not obtain the same items with the same seed and epoch." - - -def check_seedable_sampler_in_batch_sampler_shard(): - set_seed(42) - - config = DataLoaderConfiguration(use_seedable_sampler=True) - accelerator = Accelerator(dataloader_config=config) - assert accelerator.num_processes > 1, "This test requires more than one process." - - dataloader = DataLoader(list(range(10)), batch_size=1, shuffle=True) - prepared_data_loader = prepare_data_loader( - dataloader=dataloader, - use_seedable_sampler=True, - ) - - target_sampler = prepared_data_loader.batch_sampler.batch_sampler.sampler - assert isinstance(target_sampler, SeedableRandomSampler), ( - "Sampler in BatchSamplerShard is not SeedableRandomSampler." - ) - - -def check_seedable_sampler_with_data_seed(): - # Set seed - set_seed(42) - data_seed = 42 - train_set = RegressionDataset(length=10, seed=42) - train_dl = DataLoader(train_set, batch_size=2, shuffle=True) - - config = DataLoaderConfiguration(use_seedable_sampler=True, data_seed=data_seed) - accelerator = Accelerator(dataloader_config=config) - prepared_dl = accelerator.prepare(train_dl) - original_items = [] - for _ in range(3): - for batch in prepared_dl: - original_items.append(batch["x"]) - original_items = torch.cat(original_items) - - # Set new data seed - config.data_seed = 43 - accelerator = Accelerator(dataloader_config=config) - prepared_dl = accelerator.prepare(train_dl) - new_items = [] - for _ in range(3): - for batch in prepared_dl: - new_items.append(batch["x"]) - new_items = torch.cat(new_items) - assert not torch.allclose(original_items, new_items), "Obtained the same items with different data seed." - - -def mock_training(length, batch_size, generator, use_seedable_sampler=False): - set_seed(42) - generator.manual_seed(42) - train_set = RegressionDataset(length=length, seed=42) - - train_dl = generate_baseline_dataloader(train_set, generator, batch_size, use_seedable_sampler) - model = RegressionModel() - optimizer = torch.optim.SGD(model.parameters(), lr=0.1) - for epoch in range(3): - for batch in train_dl: - model.zero_grad() - output = model(batch["x"]) - loss = torch.nn.functional.mse_loss(output, batch["y"]) - loss.backward() - optimizer.step() - return train_set, model - - -def training_check(use_seedable_sampler=False): - state = AcceleratorState() - generator = torch.Generator() - batch_size = 8 - length = batch_size * 4 * state.num_processes - - train_set, old_model = mock_training(length, batch_size * state.num_processes, generator, use_seedable_sampler) - assert are_the_same_tensors(old_model.a), "Did not obtain the same model on both processes." - assert are_the_same_tensors(old_model.b), "Did not obtain the same model on both processes." - - accelerator = Accelerator() - train_dl = generate_baseline_dataloader(train_set, generator, batch_size, use_seedable_sampler) - model = RegressionModel() - optimizer = torch.optim.SGD(model.parameters(), lr=0.1) - - train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer) - set_seed(42) - generator.manual_seed(42) - for _ in range(3): - for batch in train_dl: - model.zero_grad() - output = model(batch["x"]) - loss = torch.nn.functional.mse_loss(output, batch["y"]) - accelerator.backward(loss) - optimizer.step() - - model = accelerator.unwrap_model(model).cpu() - torch.testing.assert_close( - old_model.a, - model.a, - atol=ATOL, - rtol=RTOL, - msg=lambda msg: f"Did not obtain the same model on CPU or distributed training.\n{msg}", - ) - torch.testing.assert_close( - old_model.b, - model.b, - atol=ATOL, - rtol=RTOL, - msg=lambda msg: f"Did not obtain the same model on CPU or distributed training.\n{msg}", - ) - - accelerator.print("Training yielded the same results on one CPU or distributed setup with no batch split.") - - dataloader_config = DataLoaderConfiguration(split_batches=True, use_seedable_sampler=use_seedable_sampler) - accelerator = Accelerator(dataloader_config=dataloader_config) - train_dl = generate_baseline_dataloader( - train_set, generator, batch_size * state.num_processes, use_seedable_sampler - ) - model = RegressionModel() - optimizer = torch.optim.SGD(model.parameters(), lr=0.1) - - train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer) - set_seed(42) - generator.manual_seed(42) - for _ in range(3): - for batch in train_dl: - model.zero_grad() - output = model(batch["x"]) - loss = torch.nn.functional.mse_loss(output, batch["y"]) - accelerator.backward(loss) - optimizer.step() - - model = accelerator.unwrap_model(model).cpu() - torch.testing.assert_close( - old_model.a, - model.a, - atol=ATOL, - rtol=RTOL, - msg=lambda msg: f"Did not obtain the same model on CPU or distributed training.\n{msg}", - ) - torch.testing.assert_close( - old_model.b, - model.b, - atol=ATOL, - rtol=RTOL, - msg=lambda msg: f"Did not obtain the same model on CPU or distributed training.\n{msg}", - ) - - accelerator.print("Training yielded the same results on one CPU or distributed setup with batch split.") - - # FP32 wrapper check - if torch.cuda.is_available(): - # Mostly a test that model.forward will have autocast when running unwrap_model(model, keep_fp32_wrapper=True) - print("Keep fp32 wrapper check.") - AcceleratorState._reset_state() - accelerator = Accelerator(mixed_precision="fp16") - - model = torch.nn.Linear(2, 4) - model = accelerator.prepare(model) - model_with_fp32_wrapper = accelerator.unwrap_model(model, keep_fp32_wrapper=True) - - # Run forward with fp16 as input. - # When the model is with mixed precision wrapper, no error will be raised. - input_tensor = torch.Tensor([1, 2]).to(dtype=torch.float16, device=accelerator.device) - output = model_with_fp32_wrapper(input_tensor) - - # BF16 support - if is_bf16_available(): - # Mostly a test that BF16 doesn't crash as the operation inside the model is not converted to BF16 - print("BF16 training check.") - AcceleratorState._reset_state() - dataloader_config = DataLoaderConfiguration(use_seedable_sampler=use_seedable_sampler) - accelerator = Accelerator(mixed_precision="bf16", dataloader_config=dataloader_config) - train_dl = generate_baseline_dataloader(train_set, generator, batch_size, use_seedable_sampler) - model = RegressionModel() - optimizer = torch.optim.SGD(model.parameters(), lr=0.1) - - train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer) - set_seed(42) - generator.manual_seed(42) - for _ in range(3): - for batch in train_dl: - model.zero_grad() - output = model(batch["x"]) - loss = torch.nn.functional.mse_loss(output, batch["y"]) - accelerator.backward(loss) - optimizer.step() - - model = accelerator.unwrap_model(model).cpu() - torch.testing.assert_close( - old_model.a, - model.a, - atol=ATOL, - rtol=RTOL, - msg=lambda msg: f"Did not obtain the same model on CPU or distributed training.\n{msg}", - ) - torch.testing.assert_close( - old_model.b, - model.b, - atol=ATOL, - rtol=RTOL, - msg=lambda msg: f"Did not obtain the same model on CPU or distributed training.\n{msg}", - ) - - # FP16 support (HPU fp16 model seems to be off by 10% from the CPU, which is a lot of numerical error) - if is_fp16_available() and not is_hpu_available(): - # Mostly a test that FP16 doesn't crash as the operation inside the model is not converted to FP16 - print("FP16 training check.") - AcceleratorState._reset_state() - dataloader_config = DataLoaderConfiguration(use_seedable_sampler=use_seedable_sampler) - accelerator = Accelerator(mixed_precision="fp16", dataloader_config=dataloader_config) - train_dl = generate_baseline_dataloader(train_set, generator, batch_size, use_seedable_sampler) - model = RegressionModel() - optimizer = torch.optim.SGD(model.parameters(), lr=0.1) - - train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer) - set_seed(42) - generator.manual_seed(42) - for _ in range(3): - for batch in train_dl: - model.zero_grad() - output = model(batch["x"]) - loss = torch.nn.functional.mse_loss(output, batch["y"]) - accelerator.backward(loss) - optimizer.step() - - model = accelerator.unwrap_model(model).cpu() - torch.testing.assert_close( - old_model.a, - model.a, - atol=ATOL, - rtol=RTOL, - msg=lambda msg: f"Did not obtain the same model on CPU or distributed training.\n{msg}", - ) - torch.testing.assert_close( - old_model.b, - model.b, - atol=ATOL, - rtol=RTOL, - msg=lambda msg: f"Did not obtain the same model on CPU or distributed training.\n{msg}", - ) - - # IPEX CPU tests - if is_ipex_available(): - print("ipex BF16 training check.") - AcceleratorState._reset_state() - dataloader_config = DataLoaderConfiguration(use_seedable_sampler=use_seedable_sampler) - accelerator = Accelerator(mixed_precision="bf16", cpu=True, dataloader_config=dataloader_config) - train_dl = generate_baseline_dataloader(train_set, generator, batch_size, use_seedable_sampler) - model = RegressionModel() - optimizer = torch.optim.SGD(model.parameters(), lr=0.1) - - train_dl, model, optimizer = accelerator.prepare(train_dl, model, optimizer) - set_seed(42) - generator.manual_seed(42) - for _ in range(3): - for batch in train_dl: - model.zero_grad() - output = model(batch["x"]) - loss = torch.nn.functional.mse_loss(output, batch["y"]) - accelerator.backward(loss) - optimizer.step() - - model = accelerator.unwrap_model(model).cpu() - torch.testing.assert_close( - old_model.a, - model.a, - atol=ATOL, - rtol=RTOL, - msg=lambda msg: f"Did not obtain the same model on CPU or distributed training.\n{msg}", - ) - torch.testing.assert_close( - old_model.b, - model.b, - atol=ATOL, - rtol=RTOL, - msg=lambda msg: f"Did not obtain the same model on CPU or distributed training.\n{msg}", - ) - - -def test_split_between_processes_dataset(datasets_Dataset): - state = AcceleratorState() - data = datasets_Dataset.from_list([dict(k=v) for v in range(2 * state.num_processes)]) - with state.split_between_processes(data, apply_padding=False) as results: - assert len(results) == 2, ( - f"Each process did not have two items. Process index: {state.process_index}; Length: {len(results)}" - ) - - data = datasets_Dataset.from_list([dict(k=v) for v in range(2 * state.num_processes - 1)]) - with state.split_between_processes(data, apply_padding=False) as results: - if state.is_last_process: - assert len(results) == 1, ( - f"Last process did not receive a single item. Process index: {state.process_index}; Length: {len(results)}" - ) - else: - assert len(results) == 2, ( - f"One of the intermediate processes did not receive two items. Process index: {state.process_index}; Length: {len(results)}" - ) - state.wait_for_everyone() - - odd_data = datasets_Dataset.from_list([dict(k=v) for v in range(2 * state.num_processes - 1)]) - even_data = datasets_Dataset.from_list([dict(k=v) for v in range(2 * state.num_processes)]) - - for data in [odd_data, even_data]: - expected_output = data["k"] - - with state.split_between_processes(data, apply_padding=True) as results: - if state.num_processes == 1: - assert len(results) == len(data), ( - f"Single process did not receive all items. Process index: {state.process_index}; Length: {len(results)}" - ) - else: - assert len(results) == 2, ( - f"Each process did not have two items. Process index: {state.process_index}; Length: {len(results)}" - ) - - results_per_process = [] - for result in results: - results_per_process.append(result) - - state.wait_for_everyone() - - gathered_results = gather_object(results_per_process) - output = [r["k"] for r in gathered_results[: len(data)]] - - assert expected_output == output, f"Gathered results is incorrect. Expected: {expected_output}; Got: {output}" - - -def test_split_between_processes_list(): - state = AcceleratorState() - data = list(range(0, 2 * state.num_processes)) - with state.split_between_processes(data) as results: - assert len(results) == 2, ( - f"Each process did not have two items. Process index: {state.process_index}; Length: {len(results)}" - ) - state.wait_for_everyone() - - even_data = list(range(0, (2 * state.num_processes))) - odd_data = list(range(0, (2 * state.num_processes) - 1)) - for data in [odd_data, even_data]: - expected_output = data - - with state.split_between_processes(data, apply_padding=True) as results: - num_samples_per_device = math.ceil(len(data) / state.num_processes) - # Test all processes gets the correct number of item(s) - assert len(results) == num_samples_per_device, ( - f"Process {state.device} did not get the correct number of item(s). Process index: {state.process_index}; Length: {len(results)}" - ) - - results_per_process = [] - for result in results: - results_per_process.append(result) - - state.wait_for_everyone() - - gathered_results = gather_object(results_per_process) - output = gathered_results[: len(data)] - - assert expected_output == output, f"Gathered results is incorrect. Expected: {expected_output}; Got: {output}" - - -def test_split_between_processes_nested_dict(): - state = AcceleratorState() - a = [1, 2, 3, 4, 5, 6, 7, 8] - b = ["a", "b", "c", "d", "e", "f", "g", "h"] - c = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]) - if state.num_processes in (1, 2, 4): - data = {"a": a, "b": b, "c": c} - data_copy = deepcopy(data) - with state.split_between_processes(data) as results: - if state.process_index == 0: - assert results["a"] == data_copy["a"][: 8 // state.num_processes] - elif state.num_processes == 2: - assert results["a"] == data_copy["a"][4:] - elif state.process_index == 3: - # We return a list each time - assert results["a"] == data_copy["a"][-2:], f"Expected: {data_copy['a'][-2]}, Actual: {results['a']}" - if state.process_index == 0: - assert results["b"] == data_copy["b"][: 8 // state.num_processes] - elif state.num_processes == 2: - assert results["b"] == data_copy["b"][4:] - elif state.process_index == 3: - assert results["b"] == data_copy["b"][-2:] - if state.process_index == 0: - assert torch.allclose(results["c"], data_copy["c"][: 8 // state.num_processes]), ( - f"Did not obtain expected values on process 0, expected `{data['c'][: 8 // state.num_processes]}`, received: {results['c']}" - ) - elif state.num_processes == 2: - assert torch.allclose(results["c"], data_copy["c"][4:]), ( - f"Did not obtain expected values on process 2, expected `{data['c'][4:]}`, received: {results['c']}" - ) - elif state.process_index == 3: - assert torch.allclose(results["c"], data_copy["c"][-2:]), ( - f"Did not obtain expected values on process 4, expected `{data['c'][-2:]}`, received: {results['c']}" - ) - - state.wait_for_everyone() - - -def test_split_between_processes_tensor(): - state = AcceleratorState() - if state.num_processes > 1: - data = torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]]).to(state.device) - with state.split_between_processes(data) as results: - if state.process_index == 0: - expected = torch.tensor([[0, 1, 2, 3]]).to(state.device) - else: - expected = torch.tensor([[4, 5, 6, 7]]).to(state.device) - torch.testing.assert_close(results, expected) - state.wait_for_everyone() - - even_data = torch.tensor([[i] for i in range(2 * state.num_processes)]).to(state.device) - odd_data = torch.tensor([[i] for i in range(2 * state.num_processes - 1)]).to(state.device) - for data in [even_data, odd_data]: - expected_output = [torch.tensor(i) for i in data.tolist()] - - with state.split_between_processes(data, apply_padding=True) as results: - num_samples_per_device = math.ceil(len(data) / state.num_processes) - assert len(results) == num_samples_per_device, ( - f"Process {state.device} did not get the correct number of item(s). Process index: {state.process_index}; Length: {len(results)}" - ) - results_per_process = [] - for result in results: - results_per_process.append(result.to("cpu")) - - state.wait_for_everyone() - - gathered_results = gather_object(results_per_process) - output = gathered_results[: len(data)] - - assert expected_output == output, f"Gathered results is incorrect. Expected: {expected_output}; Got: {output}" - - -def test_split_between_processes_evenly(): - state = AcceleratorState() - if state.num_processes in (1, 2, 4, 8): - data = list(range(17)) - num_samples_per_process = len(data) // state.num_processes - num_extras = len(data) % state.num_processes - with state.split_between_processes(data) as results: - if state.process_index < num_extras: - assert len(results) == num_samples_per_process + 1, ( - f"Each Process should have even elements. Expected: {num_samples_per_process + 1}, Actual: {len(results)}" - ) - else: - assert len(results) == num_samples_per_process, ( - f"Each Process should have even elements. Expected: {num_samples_per_process}, Actual: {len(results)}" - ) - state.wait_for_everyone() - - -def test_trigger(): - accelerator = Accelerator() - # should start with being false - assert accelerator.check_trigger() is False - - # set a breakpoint on the main process - if accelerator.is_main_process: - accelerator.set_trigger() - - # check it's been activated across all processes - # calls `all_reduce` and triggers a sync - assert accelerator.check_trigger() is True - - # check it's been reset after the sync - assert accelerator.check_trigger() is False - - -def test_reinstantiated_state(): - import pytest - - AcceleratorState._reset_state() - simple_model = torch.nn.Linear(1, 1) - # First define an accelerator - accelerator = Accelerator() - # Then call `reset_state`, breaking the state existing in the accelerator - AcceleratorState._reset_state() - # Now try and prepare a simple model, should raise the custom error early - with pytest.raises(AttributeError) as cm: - accelerator.prepare(simple_model) - assert "`AcceleratorState` object has no attribute" in str(cm.value.args[0]) - assert "This happens if `AcceleratorState._reset_state()`" in str(cm.value.args[0]) - - -def main(): - accelerator = Accelerator() - state = accelerator.state - if state.local_process_index == 0: - print("**Initialization**") - init_state_check() - state.wait_for_everyone() - - if state.distributed_type == DistributedType.MULTI_GPU: - num_processes_per_node = torch.cuda.device_count() - else: - num_processes_per_node = state.num_processes - - # We only run this test on non-multinode - if num_processes_per_node == state.num_processes: - if state.process_index == 0: - print("\n**Test process execution**") - process_execution_check() - - if state.process_index == 0: - print("\n**Test split between processes as a list**") - test_split_between_processes_list() - - if state.process_index == 0: - print("\n**Test split between processes as a dict**") - test_split_between_processes_nested_dict() - - if state.process_index == 0: - print("\n**Test split between processes as a tensor**") - test_split_between_processes_tensor() - - if state.process_index == 0: - print("\n**Test split between processes evenly**") - test_split_between_processes_evenly() - - if state.process_index == 0: - print("\n**Test split between processes as a datasets.Dataset**") - if is_datasets_available(): - from datasets import Dataset as datasets_Dataset - - test_split_between_processes_dataset(datasets_Dataset) - else: - print("Skipped because Hugging Face datasets is not available") - - if state.local_process_index == 0: - print("\n**Test random number generator synchronization**") - rng_sync_check() - - if state.local_process_index == 0: - print("\n**DataLoader integration test**") - dl_preparation_check() - if state.distributed_type != DistributedType.XLA: - central_dl_preparation_check() - custom_sampler_check() - check_seedable_sampler() - check_seedable_sampler_with_data_seed() - - if state.num_processes > 1: - check_seedable_sampler_in_batch_sampler_shard() - - # Trainings are not exactly the same in DeepSpeed and CPU mode - if state.distributed_type == DistributedType.DEEPSPEED: - return - - if state.local_process_index == 0: - print("\n**Training integration test**") - training_check(use_seedable_sampler=False) - training_check(use_seedable_sampler=True) - - if state.local_process_index == 0: - print("\n**Breakpoint trigger test**") - test_trigger() - - if is_pytest_available(): - if state.local_process_index == 0: - print("\n**Test reinstantiated state**") - test_reinstantiated_state() - - state.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_sync.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_sync.py deleted file mode 100644 index 44e1ecc1d59c5691f284282fb8cd2259c8a70658..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/scripts/test_sync.py +++ /dev/null @@ -1,410 +0,0 @@ -# Copyright 2022 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 copy import deepcopy - -import torch -import torch.nn.functional as F -from torch.optim import AdamW -from torch.optim.lr_scheduler import LambdaLR -from torch.utils.data import DataLoader - -from accelerate.accelerator import Accelerator, DataLoaderConfiguration, GradientAccumulationPlugin -from accelerate.state import GradientState -from accelerate.test_utils import RegressionDataset, RegressionModel -from accelerate.utils import DistributedType, set_seed - - -def check_model_parameters(model_a, model_b, did_step, iteration, **kwargs): - for param, grad_param in zip(model_a.parameters(), model_b.parameters()): - if not param.requires_grad: - continue - if not did_step: - # Grads should not be in sync - assert torch.allclose(param.grad, grad_param.grad, **kwargs) is False, ( - f"Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})" - ) - else: - # Grads should be in sync - assert torch.allclose(param.grad, grad_param.grad, **kwargs) is True, ( - f"Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})" - ) - - -def step_model(model, input, target, accelerator, do_backward=True): - model.train() - output = model(input) - loss = F.mse_loss(output, target.to(output.device)) - if not do_backward: - loss /= accelerator.gradient_accumulation_steps - loss.backward() - else: - accelerator.backward(loss) - - -def get_training_setup(accelerator, sched=False): - "Returns everything needed to perform basic training" - set_seed(42) - model = RegressionModel() - ddp_model = deepcopy(model) - dset = RegressionDataset(length=80) - dataloader = DataLoader(dset, batch_size=16) - model.to(accelerator.device) - if sched: - opt = AdamW(params=model.parameters(), lr=1e-3) - ddp_opt = AdamW(params=ddp_model.parameters(), lr=1e-3) - sched = LambdaLR(opt, lr_lambda=lambda epoch: epoch**0.65) - ddp_sched = LambdaLR(ddp_opt, lr_lambda=lambda epoch: epoch**0.65) - # Make a copy of `model` - if sched: - ddp_model, ddp_opt, ddp_sched, dataloader = accelerator.prepare(ddp_model, ddp_opt, ddp_sched, dataloader) - else: - ddp_model, dataloader = accelerator.prepare(ddp_model, dataloader) - if sched: - return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched) - return model, ddp_model, dataloader - - -def test_noop_sync(accelerator): - # Test when on a single CPU or GPU that the context manager does nothing - model, ddp_model, dataloader = get_training_setup(accelerator) - # Use a single batch - ddp_input, ddp_target = next(iter(dataloader)).values() - for iteration in range(3): - # Gather the distributed inputs and targs for the base model - input, target = accelerator.gather((ddp_input, ddp_target)) - input, target = input.to(accelerator.device), target.to(accelerator.device) - # Perform our initial ground truth step in non "DDP" - step_model(model, input, target, accelerator) - # Do "gradient accumulation" (noop) - if iteration % 2 == 0: - # Accumulate grads locally - with accelerator.no_sync(ddp_model): - step_model(ddp_model, ddp_input, ddp_target, accelerator) - else: - # Sync grads - step_model(ddp_model, ddp_input, ddp_target, accelerator) - - # Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync - check_model_parameters(model, ddp_model, True, iteration) - for param, ddp_param in zip(model.parameters(), ddp_model.parameters()): - if not param.requires_grad: - continue - assert torch.allclose(param.grad, ddp_param.grad), ( - f"Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})" - ) - - # Shuffle ddp_input on each iteration - torch.manual_seed(1337 + iteration) - ddp_input = ddp_input[torch.randperm(len(ddp_input))] - - -def test_distributed_sync(accelerator): - # Test on distributed setup that context manager behaves properly - model, ddp_model, dataloader = get_training_setup(accelerator) - # Use a single batch - ddp_input, ddp_target = next(iter(dataloader)).values() - for iteration in range(3): - # Gather the distributed inputs and targs for the base model - input, target = accelerator.gather((ddp_input, ddp_target)) - input, target = input.to(accelerator.device), target.to(accelerator.device) - # Perform our initial ground truth step in non "DDP" - step_model(model, input, target, accelerator) - # Do "gradient accumulation" (noop) - if iteration % 2 == 0: - # Accumulate grads locally - with accelerator.no_sync(ddp_model): - step_model(ddp_model, ddp_input, ddp_target, accelerator) - else: - # Sync grads - step_model(ddp_model, ddp_input, ddp_target, accelerator) - - # DDP model and model should only be in sync when not (iteration % 2 == 0) - for param, ddp_param in zip(model.parameters(), ddp_model.parameters()): - if not param.requires_grad: - continue - if iteration % 2 == 0: - # Grads should not be in sync - assert torch.allclose(param.grad, ddp_param.grad) is False, ( - f"Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})" - ) - else: - # Grads should be in sync - assert torch.allclose(param.grad, ddp_param.grad) is True, ( - f"Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})" - ) - - # Shuffle ddp_input on each iteration - torch.manual_seed(1337 + iteration) - ddp_input = ddp_input[torch.randperm(len(ddp_input))] - - -def test_distributed_sync_multiple_fwd(accelerator): - # Test on distributed setup that context manager behaves properly when used with multiple forwards followed by multiple backwards - model, ddp_model, dataloader = get_training_setup(accelerator) - # Do multiple forwards - losses = [] - num_iterations = 3 - for iteration in range(num_iterations): - ddp_input, ddp_target = next(iter(dataloader)).values() - - # Gather the distributed inputs and targs for the base model - input, target = accelerator.gather((ddp_input, ddp_target)) - input, target = input.to(accelerator.device), target.to(accelerator.device) - - # Perform our initial ground truth step in non "DDP" - step_model(model, input, target, accelerator) - - # Accumulate grads locally - with accelerator.no_sync(ddp_model): - ddp_output = ddp_model(ddp_input) - loss = F.mse_loss(ddp_output, ddp_target.to(ddp_output.device)) - losses.append(loss) - - # Do multiple backwards and sync only at the last backward - for iteration in range(num_iterations): - loss = losses[iteration] - - if iteration < num_iterations - 1: - # Accumulate grads locally - accelerator.backward(loss) - - # DDP model and model should only be in sync after last backward - for param, ddp_param in zip(model.parameters(), ddp_model.parameters()): - if not param.requires_grad: - continue - # Grads should not be in sync - assert torch.allclose(param.grad, ddp_param.grad) is False, ( - f"Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})" - ) - - else: - # Sync grads if last backward - with accelerator.trigger_sync_in_backward(ddp_model): - accelerator.backward(loss) - - # DDP model and model should only be in sync after last backward - for param, ddp_param in zip(model.parameters(), ddp_model.parameters()): - if not param.requires_grad: - continue - # Grads should be in sync - assert torch.allclose(param.grad, ddp_param.grad) is True, ( - f"Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})" - ) - - -def test_gradient_accumulation(split_batches=False, dispatch_batches=False, sync_each_batch=False): - gradient_accumulation_plugin = GradientAccumulationPlugin(num_steps=2, sync_each_batch=sync_each_batch) - dataloader_config = DataLoaderConfiguration(split_batches=split_batches, dispatch_batches=dispatch_batches) - accelerator = Accelerator( - dataloader_config=dataloader_config, - gradient_accumulation_plugin=gradient_accumulation_plugin, - ) - # Test that context manager behaves properly - model, ddp_model, dataloader = get_training_setup(accelerator) - for iteration, batch in enumerate(dataloader): - ddp_input, ddp_target = batch.values() - # Gather the distributed inputs and targs for the base model - input, target = accelerator.gather((ddp_input, ddp_target)) - input, target = input.to(accelerator.device), target.to(accelerator.device) - # Perform our initial ground truth step in non "DDP" - step_model(model, input, target, accelerator, False) - # Do "gradient accumulation" (noop) - with accelerator.accumulate(ddp_model): - step_model(ddp_model, ddp_input, ddp_target, accelerator) - - # DDP model and model should only be in sync when not (iteration % 2 == 0) - for param, ddp_param in zip(model.parameters(), ddp_model.parameters()): - if not param.requires_grad: - continue - if ((iteration + 1) % 2 == 0) or (iteration == len(dataloader) - 1) or sync_each_batch: - # Grads should be in sync - assert torch.allclose(param.grad, ddp_param.grad) is True, ( - f"Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})" - ) - else: - # Grads should not be in sync - assert torch.allclose(param.grad, ddp_param.grad) is False, ( - f"Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})" - ) - - # Shuffle ddp_input on each iteration - torch.manual_seed(1337 + iteration) - ddp_input = ddp_input[torch.randperm(len(ddp_input))] - GradientState._reset_state() - - -def test_gradient_accumulation_with_opt_and_scheduler( - split_batches=False, dispatch_batches=False, sync_each_batch=False -): - gradient_accumulation_plugin = GradientAccumulationPlugin(num_steps=2, sync_each_batch=sync_each_batch) - dataloader_config = DataLoaderConfiguration(split_batches=split_batches, dispatch_batches=dispatch_batches) - accelerator = Accelerator( - dataloader_config=dataloader_config, - gradient_accumulation_plugin=gradient_accumulation_plugin, - ) - # Test that context manager behaves properly - model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched = get_training_setup(accelerator, True) - for iteration, batch in enumerate(dataloader): - ddp_input, ddp_target = batch.values() - # Gather the distributed inputs and targs for the base model - input, target = accelerator.gather((ddp_input, ddp_target)) - input, target = input.to(accelerator.device), target.to(accelerator.device) - # Perform our initial ground truth step in non "DDP" - model.train() - ddp_model.train() - step_model(model, input, target, accelerator, False) - opt.step() - - if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(dataloader)): - if split_batches: - sched.step() - else: - for _ in range(accelerator.num_processes): - sched.step() - - # Perform gradient accumulation under wrapper - with accelerator.accumulate(ddp_model): - step_model(ddp_model, ddp_input, ddp_target, accelerator) - ddp_opt.step() - ddp_sched.step() - - # Learning rates should be the same - assert opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"], ( - f"Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]['lr']}\nDDP opt: {ddp_opt.param_groups[0]['lr']}\n" - ) - did_step = (((iteration + 1) % 2) == 0) or ((iteration + 1) == len(dataloader)) - if accelerator.num_processes > 1: - check_model_parameters( - model, - ddp_model, - did_step or sync_each_batch, # syncs at each grad_accum interval of if sync_each_batch==True - iteration, - rtol=1e-3, # needs a relative tolerance due to roundoff errors - ) - - if did_step: - opt.zero_grad() # flush gradients every accum step - ddp_opt.zero_grad() - - # Shuffle ddp_input on each iteration - torch.manual_seed(1337 + iteration) - GradientState._reset_state() - - -def test_dataloader_break(): - accelerator = Accelerator() - first_dset = RegressionDataset(length=80) - first_dataloader = DataLoader(first_dset, batch_size=16) - second_dset = RegressionDataset(length=96) - second_dataloader = DataLoader(second_dset, batch_size=16) - first_dataloader, second_dataloader = accelerator.prepare(first_dataloader, second_dataloader) - - assert accelerator.gradient_state.active_dataloader is None - for iteration, _ in enumerate(first_dataloader): - assert id(accelerator.gradient_state.active_dataloader) == id(first_dataloader) - if iteration < len(first_dataloader) - 1: - assert not accelerator.gradient_state.end_of_dataloader - if iteration == 1: - for batch_num, _ in enumerate(second_dataloader): - assert id(accelerator.gradient_state.active_dataloader) == id(second_dataloader) - if batch_num < len(second_dataloader) - 1: - assert not accelerator.gradient_state.end_of_dataloader - else: - assert accelerator.gradient_state.end_of_dataloader - else: - assert accelerator.gradient_state.end_of_dataloader - assert accelerator.gradient_state.active_dataloader is None - - -def main(): - accelerator = Accelerator() - state = accelerator.state - if state.local_process_index == 0: - print("**Test `accumulate` gradient accumulation with dataloader break**") - if state.distributed_type != DistributedType.XLA: - test_dataloader_break() - if state.distributed_type == DistributedType.NO: - if state.local_process_index == 0: - print("**Test NOOP `no_sync` context manager**") - test_noop_sync(accelerator) - if state.distributed_type in ( - DistributedType.MULTI_GPU, - DistributedType.MULTI_NPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_CPU, - DistributedType.MULTI_HPU, - ): - if state.local_process_index == 0: - print("**Test Distributed `no_sync` context manager**") - test_distributed_sync(accelerator) - if state.local_process_index == 0: - print("**Test Distributed `no_sync` context manager with multiple forwards**") - test_distributed_sync_multiple_fwd(accelerator) - if state.distributed_type in ( - DistributedType.MULTI_GPU, - DistributedType.MULTI_NPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_HPU, - ): - for split_batch in [True, False]: - for dispatch_batches in [True, False]: - for sync_each_batch in [True, False]: - if state.local_process_index == 0: - print( - "**Test `accumulate` gradient accumulation, ", - f"`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}` and `sync_each_batch={sync_each_batch}`**", - ) - test_gradient_accumulation(split_batch, dispatch_batches, sync_each_batch) - - # Currently will break on torch 2.0 +, need to investigate why - if state.local_process_index == 0: - print( - "**Test `accumulate` gradient accumulation with optimizer and scheduler, ", - "`split_batches=False`, `dispatch_batches=False`, `sync_each_batch=False`**", - ) - test_gradient_accumulation_with_opt_and_scheduler() - if state.distributed_type in ( - DistributedType.MULTI_GPU, - DistributedType.MULTI_NPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_HPU, - ): - for split_batch in [True, False]: - for dispatch_batches in [True, False]: - for sync_each_batch in [True, False]: - if not split_batch and not dispatch_batches and not sync_each_batch: - continue - if state.local_process_index == 0: - print( - "**Test `accumulate` gradient accumulation with optimizer and scheduler, ", - f"`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}` and `sync_each_batch={sync_each_batch}`**", - ) - test_gradient_accumulation_with_opt_and_scheduler(split_batch, dispatch_batches, sync_each_batch) - state.destroy_process_group() - - -def _mp_fn(index): - # For xla_spawn (TPUs) - main() - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/testing.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/testing.py deleted file mode 100644 index e964399784779fb9bf3d34269d08fd76d4e96319..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/testing.py +++ /dev/null @@ -1,870 +0,0 @@ -# Copyright 2021 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 asyncio -import inspect -import io -import os -import re -import shutil -import subprocess -import sys -import tempfile -import unittest -from contextlib import contextmanager -from functools import partial -from pathlib import Path -from typing import Union -from unittest import mock - -import torch - -import accelerate - -from ..state import AcceleratorState -from ..utils import ( - check_cuda_fp8_capability, - compare_versions, - gather, - is_aim_available, - is_bnb_available, - is_clearml_available, - is_comet_ml_available, - is_cuda_available, - is_datasets_available, - is_deepspeed_available, - is_dvclive_available, - is_fp8_available, - is_fp16_available, - is_habana_gaudi1, - is_hpu_available, - is_import_timer_available, - is_matplotlib_available, - is_mlflow_available, - is_mlu_available, - is_mps_available, - is_musa_available, - is_npu_available, - is_pandas_available, - is_pippy_available, - is_pytest_available, - is_schedulefree_available, - is_sdaa_available, - is_swanlab_available, - is_tensorboard_available, - is_timm_available, - is_torch_version, - is_torch_xla_available, - is_torchao_available, - is_torchdata_stateful_dataloader_available, - is_torchvision_available, - is_trackio_available, - is_transformer_engine_available, - is_transformers_available, - is_triton_available, - is_wandb_available, - is_xpu_available, - str_to_bool, -) - - -def get_backend(): - if is_torch_xla_available(): - return "xla", torch.cuda.device_count(), torch.cuda.memory_allocated - elif is_cuda_available(): - return "cuda", torch.cuda.device_count(), torch.cuda.memory_allocated - elif is_mps_available(min_version="2.0"): - return "mps", 1, torch.mps.current_allocated_memory - elif is_mps_available(): - return "mps", 1, lambda: 0 - elif is_mlu_available(): - return "mlu", torch.mlu.device_count(), torch.mlu.memory_allocated - elif is_sdaa_available(): - return "sdaa", torch.sdaa.device_count(), torch.sdaa.memory_allocated - elif is_musa_available(): - return "musa", torch.musa.device_count(), torch.musa.memory_allocated - elif is_npu_available(): - return "npu", torch.npu.device_count(), torch.npu.memory_allocated - elif is_xpu_available(): - return "xpu", torch.xpu.device_count(), torch.xpu.memory_allocated - elif is_hpu_available(): - return "hpu", torch.hpu.device_count(), torch.hpu.memory_allocated - else: - return "cpu", 1, lambda: 0 - - -torch_device, device_count, memory_allocated_func = get_backend() - - -def get_launch_command(**kwargs) -> list: - """ - Wraps around `kwargs` to help simplify launching from `subprocess`. - - Example: - ```python - # returns ['accelerate', 'launch', '--num_processes=2', '--device_count=2'] - get_launch_command(num_processes=2, device_count=2) - ``` - """ - command = ["accelerate", "launch"] - for k, v in kwargs.items(): - if isinstance(v, bool) and v: - command.append(f"--{k}") - elif v is not None: - command.append(f"--{k}={v}") - return command - - -DEFAULT_LAUNCH_COMMAND = get_launch_command(num_processes=device_count, monitor_interval=0.1) - - -def parse_flag_from_env(key, default=False): - try: - value = os.environ[key] - except KeyError: - # KEY isn't set, default to `default`. - _value = default - else: - # KEY is set, convert it to True or False. - try: - _value = str_to_bool(value) - except ValueError: - # More values are supported, but let's keep the message simple. - raise ValueError(f"If set, {key} must be yes or no.") - return _value - - -_run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False) - - -def skip(test_case): - "Decorator that skips a test unconditionally" - return unittest.skip("Test was skipped")(test_case) - - -def slow(test_case): - """ - Decorator marking a test as slow. Slow tests are skipped by default. Set the RUN_SLOW environment variable to a - truthy value to run them. - """ - return unittest.skipUnless(_run_slow_tests, "test is slow")(test_case) - - -def require_cpu(test_case): - """ - Decorator marking a test that must be only ran on the CPU. These tests are skipped when a GPU is available. - """ - return unittest.skipUnless(torch_device == "cpu", "test requires only a CPU")(test_case) - - -def require_non_cpu(test_case): - """ - Decorator marking a test that requires a hardware accelerator backend. These tests are skipped when there are no - hardware accelerator available. - """ - return unittest.skipUnless(torch_device != "cpu", "test requires a GPU")(test_case) - - -def require_cuda(test_case): - """ - Decorator marking a test that requires CUDA. These tests are skipped when there are no GPU available or when - TorchXLA is available. - """ - return unittest.skipUnless(is_cuda_available() and not is_torch_xla_available(), "test requires a GPU")(test_case) - - -def require_cuda_or_hpu(test_case): - """ - Decorator marking a test that requires CUDA or HPU. These tests are skipped when there are no GPU available or when - TorchXLA is available. - """ - return unittest.skipUnless( - (is_cuda_available() and not is_torch_xla_available()) or is_hpu_available(), "test requires a GPU or HPU" - )(test_case) - - -def require_xpu(test_case): - """ - Decorator marking a test that requires XPU. These tests are skipped when there are no XPU available. - """ - return unittest.skipUnless(is_xpu_available(), "test requires a XPU")(test_case) - - -def require_cuda_or_xpu(test_case): - """ - Decorator marking a test that requires CUDA or XPU. These tests are skipped when there are no GPU available or when - TorchXLA is available. - """ - cuda_condition = is_cuda_available() and not is_torch_xla_available() - xpu_condition = is_xpu_available() - return unittest.skipUnless(cuda_condition or xpu_condition, "test requires a CUDA GPU or XPU")(test_case) - - -def require_non_xpu(test_case): - """ - Decorator marking a test that should be skipped for XPU. - """ - return unittest.skipUnless(torch_device != "xpu", "test requires a non-XPU")(test_case) - - -def require_non_hpu(test_case): - """ - Decorator marking a test that should be skipped for HPU. - """ - return unittest.skipUnless(torch_device != "hpu", "test requires a non-HPU")(test_case) - - -def require_fp16(test_case): - """ - Decorator marking a test that requires FP16. These tests are skipped when FP16 is not supported. - """ - - return unittest.skipUnless(is_fp16_available(), "test requires FP16 support")(test_case) - - -def require_fp8(test_case): - """ - Decorator marking a test that requires FP8. These tests are skipped when FP8 is not supported. - """ - - # is_fp8_available only checks for libraries - # ideally it should check for device capability as well - fp8_is_available = is_fp8_available() - - if torch.cuda.is_available() and not check_cuda_fp8_capability(): - fp8_is_available = False - - if is_hpu_available() and is_habana_gaudi1(): - fp8_is_available = False - - return unittest.skipUnless(fp8_is_available, "test requires FP8 support")(test_case) - - -def require_fsdp2(test_case): - return unittest.skipUnless(is_torch_version(">=", "2.5.0"), "test requires FSDP2 (torch >= 2.5.0)")(test_case) - - -def require_mlu(test_case): - """ - Decorator marking a test that requires MLU. These tests are skipped when there are no MLU available. - """ - return unittest.skipUnless(is_mlu_available(), "test require a MLU")(test_case) - - -def require_sdaa(test_case): - """ - Decorator marking a test that requires SDAA. These tests are skipped when there are no SDAA available. - """ - return unittest.skipUnless(is_sdaa_available(), "test require a SDAA")(test_case) - - -def require_musa(test_case): - """ - Decorator marking a test that requires MUSA. These tests are skipped when there are no MUSA available. - """ - return unittest.skipUnless(is_musa_available(), "test require a MUSA")(test_case) - - -def require_npu(test_case): - """ - Decorator marking a test that requires NPU. These tests are skipped when there are no NPU available. - """ - return unittest.skipUnless(is_npu_available(), "test require a NPU")(test_case) - - -def require_mps(test_case): - """ - Decorator marking a test that requires MPS backend. These tests are skipped when torch doesn't support `mps` - backend. - """ - return unittest.skipUnless(is_mps_available(), "test requires a `mps` backend support in `torch`")(test_case) - - -def require_huggingface_suite(test_case): - """ - Decorator marking a test that requires transformers and datasets. These tests are skipped when they are not. - """ - return unittest.skipUnless( - is_transformers_available() and is_datasets_available(), - "test requires the Hugging Face suite", - )(test_case) - - -def require_transformers(test_case): - """ - Decorator marking a test that requires transformers. These tests are skipped when they are not. - """ - return unittest.skipUnless(is_transformers_available(), "test requires the transformers library")(test_case) - - -def require_timm(test_case): - """ - Decorator marking a test that requires timm. These tests are skipped when they are not. - """ - return unittest.skipUnless(is_timm_available(), "test requires the timm library")(test_case) - - -def require_torchvision(test_case): - """ - Decorator marking a test that requires torchvision. These tests are skipped when they are not. - """ - return unittest.skipUnless(is_torchvision_available(), "test requires the torchvision library")(test_case) - - -def require_triton(test_case): - """ - Decorator marking a test that requires triton. These tests are skipped when they are not. - """ - return unittest.skipUnless(is_triton_available(), "test requires the triton library")(test_case) - - -def require_schedulefree(test_case): - """ - Decorator marking a test that requires schedulefree. These tests are skipped when they are not. - """ - return unittest.skipUnless(is_schedulefree_available(), "test requires the schedulefree library")(test_case) - - -def require_bnb(test_case): - """ - Decorator marking a test that requires bitsandbytes. These tests are skipped when they are not. - """ - return unittest.skipUnless(is_bnb_available(), "test requires the bitsandbytes library")(test_case) - - -def require_tpu(test_case): - """ - Decorator marking a test that requires TPUs. These tests are skipped when there are no TPUs available. - """ - return unittest.skipUnless(is_torch_xla_available(check_is_tpu=True), "test requires TPU")(test_case) - - -def require_non_torch_xla(test_case): - """ - Decorator marking a test as requiring an environment without TorchXLA. These tests are skipped when TorchXLA is - available. - """ - return unittest.skipUnless(not is_torch_xla_available(), "test requires an env without TorchXLA")(test_case) - - -def require_single_device(test_case): - """ - Decorator marking a test that requires a single device. These tests are skipped when there is no hardware - accelerator available or number of devices is more than one. - """ - return unittest.skipUnless( - torch_device != "cpu" and device_count == 1, "test requires a single device accelerator" - )(test_case) - - -def require_single_gpu(test_case): - """ - Decorator marking a test that requires CUDA on a single GPU. These tests are skipped when there are no GPU - available or number of GPUs is more than one. - """ - return unittest.skipUnless(torch.cuda.device_count() == 1, "test requires a GPU")(test_case) - - -def require_single_xpu(test_case): - """ - Decorator marking a test that requires CUDA on a single XPU. These tests are skipped when there are no XPU - available or number of xPUs is more than one. - """ - return unittest.skipUnless(torch.xpu.device_count() == 1, "test requires a XPU")(test_case) - - -def require_multi_device(test_case): - """ - Decorator marking a test that requires a multi-device setup. These tests are skipped on a machine without multiple - devices. - """ - return unittest.skipUnless(device_count > 1, "test requires multiple hardware accelerators")(test_case) - - -def require_multi_gpu(test_case): - """ - Decorator marking a test that requires a multi-GPU setup. These tests are skipped on a machine without multiple - GPUs. - """ - return unittest.skipUnless(torch.cuda.device_count() > 1, "test requires multiple GPUs")(test_case) - - -def require_multi_xpu(test_case): - """ - Decorator marking a test that requires a multi-XPU setup. These tests are skipped on a machine without multiple - XPUs. - """ - return unittest.skipUnless(torch.xpu.device_count() > 1, "test requires multiple XPUs")(test_case) - - -def require_multi_gpu_or_xpu(test_case): - """ - Decorator marking a test that requires a multi-GPU setup. These tests are skipped on a machine without multiple - GPUs or XPUs. - """ - return unittest.skipUnless( - (is_cuda_available() or is_xpu_available()) and device_count > 1, "test requires multiple GPUs or XPUs" - )(test_case) - - -def require_deepspeed(test_case): - """ - Decorator marking a test that requires DeepSpeed installed. These tests are skipped when DeepSpeed isn't installed - """ - return unittest.skipUnless(is_deepspeed_available(), "test requires DeepSpeed")(test_case) - - -def require_tp(test_case): - """ - Decorator marking a test that requires TP installed. These tests are skipped when TP isn't installed - """ - return unittest.skipUnless( - is_torch_version(">=", "2.3.0") and compare_versions("transformers", ">=", "4.52.0"), - "test requires torch version >= 2.3.0 and transformers version >= 4.52.0", - )(test_case) - - -def require_torch_min_version(test_case=None, version=None): - """ - Decorator marking that a test requires a particular torch version to be tested. These tests are skipped when an - installed torch version is less than the required one. - """ - if test_case is None: - return partial(require_torch_min_version, version=version) - return unittest.skipUnless(is_torch_version(">=", version), f"test requires torch version >= {version}")(test_case) - - -def require_tensorboard(test_case): - """ - Decorator marking a test that requires tensorboard installed. These tests are skipped when tensorboard isn't - installed - """ - return unittest.skipUnless(is_tensorboard_available(), "test requires Tensorboard")(test_case) - - -def require_wandb(test_case): - """ - Decorator marking a test that requires wandb installed. These tests are skipped when wandb isn't installed - """ - return unittest.skipUnless(is_wandb_available(), "test requires wandb")(test_case) - - -def require_trackio(test_case): - """ - Decorator marking a test that requires trackio installed. These tests are skipped when trackio isn't installed - """ - return unittest.skipUnless(is_trackio_available(), "test requires trackio")(test_case) - - -def require_comet_ml(test_case): - """ - Decorator marking a test that requires comet_ml installed. These tests are skipped when comet_ml isn't installed - """ - return unittest.skipUnless(is_comet_ml_available(), "test requires comet_ml")(test_case) - - -def require_aim(test_case): - """ - Decorator marking a test that requires aim installed. These tests are skipped when aim isn't installed - """ - return unittest.skipUnless(is_aim_available(), "test requires aim")(test_case) - - -def require_clearml(test_case): - """ - Decorator marking a test that requires clearml installed. These tests are skipped when clearml isn't installed - """ - return unittest.skipUnless(is_clearml_available(), "test requires clearml")(test_case) - - -def require_dvclive(test_case): - """ - Decorator marking a test that requires dvclive installed. These tests are skipped when dvclive isn't installed - """ - return unittest.skipUnless(is_dvclive_available(), "test requires dvclive")(test_case) - - -def require_swanlab(test_case): - """ - Decorator marking a test that requires swanlab installed. These tests are skipped when swanlab isn't installed - """ - return unittest.skipUnless(is_swanlab_available(), "test requires swanlab")(test_case) - - -def require_pandas(test_case): - """ - Decorator marking a test that requires pandas installed. These tests are skipped when pandas isn't installed - """ - return unittest.skipUnless(is_pandas_available(), "test requires pandas")(test_case) - - -def require_mlflow(test_case): - """ - Decorator marking a test that requires mlflow installed. These tests are skipped when mlflow isn't installed - """ - return unittest.skipUnless(is_mlflow_available(), "test requires mlflow")(test_case) - - -def require_pippy(test_case): - """ - Decorator marking a test that requires pippy installed. These tests are skipped when pippy isn't installed It is - also checked if the test is running on a Gaudi1 device which doesn't support pippy. - """ - return unittest.skipUnless(is_pippy_available() and not is_habana_gaudi1(), "test requires pippy")(test_case) - - -def require_import_timer(test_case): - """ - Decorator marking a test that requires tuna interpreter installed. These tests are skipped when tuna isn't - installed - """ - return unittest.skipUnless(is_import_timer_available(), "test requires tuna interpreter")(test_case) - - -def require_transformer_engine(test_case): - """ - Decorator marking a test that requires transformers engine installed. These tests are skipped when transformers - engine isn't installed - """ - return unittest.skipUnless(is_transformer_engine_available(), "test requires transformers engine")(test_case) - - -def require_torchao(test_case): - """ - Decorator marking a test that requires torchao installed. These tests are skipped when torchao isn't installed - """ - return unittest.skipUnless(is_torchao_available(), "test requires torchao")(test_case) - - -def require_matplotlib(test_case): - """ - Decorator marking a test that requires matplotlib installed. These tests are skipped when matplotlib isn't - installed - """ - return unittest.skipUnless(is_matplotlib_available(), "test requires matplotlib")(test_case) - - -_atleast_one_tracker_available = ( - any([is_wandb_available(), is_tensorboard_available(), is_trackio_available(), is_swanlab_available()]) - and not is_comet_ml_available() -) - - -def require_trackers(test_case): - """ - Decorator marking that a test requires at least one tracking library installed. These tests are skipped when none - are installed - """ - return unittest.skipUnless( - _atleast_one_tracker_available, - "test requires at least one tracker to be available and for `comet_ml` to not be installed", - )(test_case) - - -def require_torchdata_stateful_dataloader(test_case): - """ - Decorator marking a test that requires torchdata.stateful_dataloader. - - These tests are skipped when torchdata with stateful_dataloader module isn't installed. - - """ - return unittest.skipUnless( - is_torchdata_stateful_dataloader_available(), "test requires torchdata.stateful_dataloader" - )(test_case) - - -def run_first(test_case): - """ - Decorator marking a test with order(1). When pytest-order plugin is installed, tests marked with this decorator are - garanteed to run first. - - This is especially useful in some test settings like on a Gaudi instance where a Gaudi device can only be used by a - single process at a time. So we make sure all tests that run in a subprocess are launched first, to avoid device - allocation conflicts. - - If pytest is not installed, test will be returned as is. - """ - - if is_pytest_available(): - import pytest - - return pytest.mark.order(1)(test_case) - return test_case - - -class TempDirTestCase(unittest.TestCase): - """ - A TestCase class that keeps a single `tempfile.TemporaryDirectory` open for the duration of the class, wipes its - data at the start of a test, and then destroyes it at the end of the TestCase. - - Useful for when a class or API requires a single constant folder throughout it's use, such as Weights and Biases - - The temporary directory location will be stored in `self.tmpdir` - """ - - clear_on_setup = True - - @classmethod - def setUpClass(cls): - "Creates a `tempfile.TemporaryDirectory` and stores it in `cls.tmpdir`" - cls.tmpdir = Path(tempfile.mkdtemp()) - - @classmethod - def tearDownClass(cls): - "Remove `cls.tmpdir` after test suite has finished" - if os.path.exists(cls.tmpdir): - shutil.rmtree(cls.tmpdir) - - def setUp(self): - "Destroy all contents in `self.tmpdir`, but not `self.tmpdir`" - if self.clear_on_setup: - for path in self.tmpdir.glob("**/*"): - if path.is_file(): - path.unlink() - elif path.is_dir(): - shutil.rmtree(path) - - -class AccelerateTestCase(unittest.TestCase): - """ - A TestCase class that will reset the accelerator state at the end of every test. Every test that checks or utilizes - the `AcceleratorState` class should inherit from this to avoid silent failures due to state being shared between - tests. - """ - - def tearDown(self): - super().tearDown() - # Reset the state of the AcceleratorState singleton. - AcceleratorState._reset_state(True) - - -class MockingTestCase(unittest.TestCase): - """ - A TestCase class designed to dynamically add various mockers that should be used in every test, mimicking the - behavior of a class-wide mock when defining one normally will not do. - - Useful when a mock requires specific information available only initialized after `TestCase.setUpClass`, such as - setting an environment variable with that information. - - The `add_mocks` function should be ran at the end of a `TestCase`'s `setUp` function, after a call to - `super().setUp()` such as: - ```python - def setUp(self): - super().setUp() - mocks = mock.patch.dict(os.environ, {"SOME_ENV_VAR", "SOME_VALUE"}) - self.add_mocks(mocks) - ``` - """ - - def add_mocks(self, mocks: Union[mock.Mock, list[mock.Mock]]): - """ - Add custom mocks for tests that should be repeated on each test. Should be called during - `MockingTestCase.setUp`, after `super().setUp()`. - - Args: - mocks (`mock.Mock` or list of `mock.Mock`): - Mocks that should be added to the `TestCase` after `TestCase.setUpClass` has been run - """ - self.mocks = mocks if isinstance(mocks, (tuple, list)) else [mocks] - for m in self.mocks: - m.start() - self.addCleanup(m.stop) - - -def are_the_same_tensors(tensor): - state = AcceleratorState() - tensor = tensor[None].clone().to(state.device) - tensors = gather(tensor).cpu() - tensor = tensor[0].cpu() - for i in range(tensors.shape[0]): - if not torch.equal(tensors[i], tensor): - return False - return True - - -class _RunOutput: - def __init__(self, returncode, stdout, stderr): - self.returncode = returncode - self.stdout = stdout - self.stderr = stderr - - -async def _read_stream(stream, callback): - while True: - line = await stream.readline() - if line: - callback(line) - else: - break - - -async def _stream_subprocess(cmd, env=None, stdin=None, timeout=None, quiet=False, echo=False) -> _RunOutput: - if echo: - print("\nRunning: ", " ".join(cmd)) - - p = await asyncio.create_subprocess_exec( - cmd[0], - *cmd[1:], - stdin=stdin, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - env=env, - ) - - # note: there is a warning for a possible deadlock when using `wait` with huge amounts of data in the pipe - # https://docs.python.org/3/library/asyncio-subprocess.html#asyncio.asyncio.subprocess.Process.wait - # - # If it starts hanging, will need to switch to the following code. The problem is that no data - # will be seen until it's done and if it hangs for example there will be no debug info. - # out, err = await p.communicate() - # return _RunOutput(p.returncode, out, err) - - out = [] - err = [] - - def tee(line, sink, pipe, label=""): - line = line.decode("utf-8").rstrip() - sink.append(line) - if not quiet: - print(label, line, file=pipe) - - # XXX: the timeout doesn't seem to make any difference here - await asyncio.wait( - [ - asyncio.create_task(_read_stream(p.stdout, lambda l: tee(l, out, sys.stdout, label="stdout:"))), - asyncio.create_task(_read_stream(p.stderr, lambda l: tee(l, err, sys.stderr, label="stderr:"))), - ], - timeout=timeout, - ) - return _RunOutput(await p.wait(), out, err) - - -def execute_subprocess_async(cmd: list, env=None, stdin=None, timeout=180, quiet=False, echo=True) -> _RunOutput: - # Cast every path in `cmd` to a string - for i, c in enumerate(cmd): - if isinstance(c, Path): - cmd[i] = str(c) - loop = asyncio.get_event_loop() - result = loop.run_until_complete( - _stream_subprocess(cmd, env=env, stdin=stdin, timeout=timeout, quiet=quiet, echo=echo) - ) - - cmd_str = " ".join(cmd) - if result.returncode > 0: - stderr = "\n".join(result.stderr) - raise RuntimeError( - f"'{cmd_str}' failed with returncode {result.returncode}\n\n" - f"The combined stderr from workers follows:\n{stderr}" - ) - - return result - - -def pytest_xdist_worker_id(): - """ - Returns an int value of worker's numerical id under `pytest-xdist`'s concurrent workers `pytest -n N` regime, or 0 - if `-n 1` or `pytest-xdist` isn't being used. - """ - worker = os.environ.get("PYTEST_XDIST_WORKER", "gw0") - worker = re.sub(r"^gw", "", worker, 0, re.M) - return int(worker) - - -def get_torch_dist_unique_port(): - """ - Returns a port number that can be fed to `torch.distributed.launch`'s `--master_port` argument. - - Under `pytest-xdist` it adds a delta number based on a worker id so that concurrent tests don't try to use the same - port at once. - """ - port = 29500 - uniq_delta = pytest_xdist_worker_id() - return port + uniq_delta - - -class SubprocessCallException(Exception): - pass - - -def run_command(command: list[str], return_stdout=False, env=None): - """ - Runs `command` with `subprocess.check_output` and will potentially return the `stdout`. Will also properly capture - if an error occurred while running `command` - """ - # Cast every path in `command` to a string - for i, c in enumerate(command): - if isinstance(c, Path): - command[i] = str(c) - if env is None: - env = os.environ.copy() - try: - output = subprocess.check_output(command, stderr=subprocess.STDOUT, env=env) - if return_stdout: - if hasattr(output, "decode"): - output = output.decode("utf-8") - return output - except subprocess.CalledProcessError as e: - raise SubprocessCallException( - f"Command `{' '.join(command)}` failed with the following error:\n\n{e.output.decode()}" - ) from e - - -def path_in_accelerate_package(*components: str) -> Path: - """ - Get a path within the `accelerate` package's directory. - - Args: - *components: Components of the path to join after the package directory. - - Returns: - `Path`: The path to the requested file or directory. - """ - - accelerate_package_dir = Path(inspect.getfile(accelerate)).parent - return accelerate_package_dir.joinpath(*components) - - -@contextmanager -def assert_exception(exception_class: Exception, msg: str = None) -> bool: - """ - Context manager to assert that the right `Exception` class was raised. - - If `msg` is provided, will check that the message is contained in the raised exception. - """ - was_ran = False - try: - yield - was_ran = True - except Exception as e: - assert isinstance(e, exception_class), f"Expected exception of type {exception_class} but got {type(e)}" - if msg is not None: - assert msg in str(e), f"Expected message '{msg}' to be in exception but got '{str(e)}'" - if was_ran: - raise AssertionError(f"Expected exception of type {exception_class} but ran without issue.") - - -def capture_call_output(func, *args, **kwargs): - """ - Takes in a `func` with `args` and `kwargs` and returns the captured stdout as a string - """ - captured_output = io.StringIO() - original_stdout = sys.stdout - try: - sys.stdout = captured_output - func(*args, **kwargs) - except Exception as e: - raise e - finally: - sys.stdout = original_stdout - return captured_output.getvalue() diff --git a/venv/lib/python3.10/site-packages/accelerate/test_utils/training.py b/venv/lib/python3.10/site-packages/accelerate/test_utils/training.py deleted file mode 100644 index e71896c1f98bf47093d772cc77d7f27ba30c7ee8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/test_utils/training.py +++ /dev/null @@ -1,162 +0,0 @@ -# Copyright 2021 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 numpy as np -import torch -from torch.utils.data import DataLoader - -from accelerate.utils.dataclasses import DistributedType - - -class RegressionDataset: - def __init__(self, a=2, b=3, length=64, seed=None): - rng = np.random.default_rng(seed) - self.length = length - self.x = rng.normal(size=(length,)).astype(np.float32) - self.y = a * self.x + b + rng.normal(scale=0.1, size=(length,)).astype(np.float32) - - def __len__(self): - return self.length - - def __getitem__(self, i): - return {"x": self.x[i], "y": self.y[i]} - - -class RegressionModel4XPU(torch.nn.Module): - def __init__(self, a=0, b=0, double_output=False): - super().__init__() - self.a = torch.nn.Parameter(torch.tensor([2, 3]).float()) - self.b = torch.nn.Parameter(torch.tensor([2, 3]).float()) - self.first_batch = True - - def forward(self, x=None): - if self.first_batch: - print(f"Model dtype: {self.a.dtype}, {self.b.dtype}. Input dtype: {x.dtype}") - self.first_batch = False - return x * self.a[0] + self.b[0] - - -class RegressionModel(torch.nn.Module): - def __init__(self, a=0, b=0, double_output=False): - super().__init__() - self.a = torch.nn.Parameter(torch.tensor(a).float()) - self.b = torch.nn.Parameter(torch.tensor(b).float()) - self.first_batch = True - - def forward(self, x=None): - if self.first_batch: - print(f"Model dtype: {self.a.dtype}, {self.b.dtype}. Input dtype: {x.dtype}") - self.first_batch = False - return x * self.a + self.b - - -def mocked_dataloaders(accelerator, batch_size: int = 16): - from datasets import load_dataset - from transformers import AutoTokenizer - - tokenizer = AutoTokenizer.from_pretrained("bert-base-cased") - data_files = {"train": "tests/test_samples/MRPC/train.csv", "validation": "tests/test_samples/MRPC/dev.csv"} - datasets = load_dataset("csv", data_files=data_files) - label_list = datasets["train"].unique("label") - - label_to_id = {v: i for i, v in enumerate(label_list)} - - def tokenize_function(examples): - # max_length=None => use the model max length (it's actually the default) - outputs = tokenizer( - examples["sentence1"], examples["sentence2"], truncation=True, max_length=None, padding="max_length" - ) - if "label" in examples: - outputs["labels"] = [label_to_id[l] for l in examples["label"]] - return outputs - - # Apply the method we just defined to all the examples in all the splits of the dataset - tokenized_datasets = datasets.map( - tokenize_function, - batched=True, - remove_columns=["sentence1", "sentence2", "label"], - ) - - def collate_fn(examples): - # On TPU it's best to pad everything to the same length or training will be very slow. - if accelerator.distributed_type == DistributedType.XLA: - return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt") - return tokenizer.pad(examples, padding="longest", return_tensors="pt") - - # Instantiate dataloaders. - train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=2) - eval_dataloader = DataLoader(tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=1) - - return train_dataloader, eval_dataloader - - -def mocked_dataloaders_for_autoregressive_models(accelerator, batch_size: int = 16): - from datasets import load_dataset - from transformers import AutoTokenizer - - tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM-360M") - tokenizer.pad_token = tokenizer.eos_token - - data_files = {"train": "tests/test_samples/MRPC/train.csv", "validation": "tests/test_samples/MRPC/dev.csv"} - datasets = load_dataset("csv", data_files=data_files) - - def tokenize_function(examples): - # max_length=None => use the model max length (it's actually the default) - outputs = tokenizer(examples["sentence1"], truncation=True, max_length=None, return_attention_mask=False) - return outputs - - # Apply the method we just defined to all the examples in all the splits of the dataset - # starting with the main process first: - with accelerator.main_process_first(): - tokenized_datasets = datasets.map( - tokenize_function, - batched=True, - remove_columns=["sentence1", "sentence2", "label"], - ) - - def collate_fn(examples): - # On TPU it's best to pad everything to the same length or training will be very slow. - max_length = ( - 128 - if accelerator.distributed_type == DistributedType.XLA - else max([len(e["input_ids"]) for e in examples]) - ) - # When using mixed precision we want round multiples of 8/16 - if accelerator.mixed_precision == "fp8": - pad_to_multiple_of = 16 - elif accelerator.mixed_precision != "no": - pad_to_multiple_of = 8 - else: - pad_to_multiple_of = None - - batch = tokenizer.pad( - examples, - padding="max_length", - max_length=max_length + 1, - pad_to_multiple_of=pad_to_multiple_of, - return_tensors="pt", - ) - - batch["labels"] = batch["input_ids"][:, 1:] - batch["input_ids"] = batch["input_ids"][:, :-1] - - batch["labels"] = torch.where(batch["labels"] == tokenizer.pad_token_id, -100, batch["labels"]) - - return batch - - # Instantiate dataloaders. - train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=False, collate_fn=collate_fn, batch_size=2) - eval_dataloader = DataLoader(tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=1) - - return train_dataloader, eval_dataloader diff --git a/venv/lib/python3.10/site-packages/accelerate/tracking.py b/venv/lib/python3.10/site-packages/accelerate/tracking.py deleted file mode 100644 index 29e0aece74eedabd6c7417c11609eb24aad8df50..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/tracking.py +++ /dev/null @@ -1,1326 +0,0 @@ -# Copyright 2022 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. - -# Expectation: -# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`} - -import json -import os -import time -from functools import wraps -from typing import Any, Optional, Union - -import yaml -from packaging import version - -from .logging import get_logger -from .state import PartialState -from .utils import ( - LoggerType, - compare_versions, - is_aim_available, - is_clearml_available, - is_comet_ml_available, - is_dvclive_available, - is_mlflow_available, - is_swanlab_available, - is_tensorboard_available, - is_trackio_available, - is_wandb_available, - listify, -) - - -_available_trackers = [] - -if is_tensorboard_available(): - _available_trackers.append(LoggerType.TENSORBOARD) - -if is_wandb_available(): - _available_trackers.append(LoggerType.WANDB) - -if is_comet_ml_available(): - _available_trackers.append(LoggerType.COMETML) - -if is_aim_available(): - _available_trackers.append(LoggerType.AIM) - -if is_mlflow_available(): - _available_trackers.append(LoggerType.MLFLOW) - -if is_clearml_available(): - _available_trackers.append(LoggerType.CLEARML) - -if is_dvclive_available(): - _available_trackers.append(LoggerType.DVCLIVE) - -if is_swanlab_available(): - _available_trackers.append(LoggerType.SWANLAB) - -if is_trackio_available(): - _available_trackers.append(LoggerType.TRACKIO) - -logger = get_logger(__name__) - - -def on_main_process(function): - """ - Decorator to selectively run the decorated function on the main process only based on the `main_process_only` - attribute in a class. - - Checks at function execution rather than initialization time, not triggering the initialization of the - `PartialState`. - """ - - @wraps(function) - def execute_on_main_process(self, *args, **kwargs): - if getattr(self, "main_process_only", False): - return PartialState().on_main_process(function)(self, *args, **kwargs) - else: - return function(self, *args, **kwargs) - - return execute_on_main_process - - -def get_available_trackers(): - "Returns a list of all supported available trackers in the system" - return _available_trackers - - -class GeneralTracker: - """ - A base Tracker class to be used for all logging integration implementations. - - Each function should take in `**kwargs` that will automatically be passed in from a base dictionary provided to - [`Accelerator`]. - - Should implement `name`, `requires_logging_directory`, and `tracker` properties such that: - - `name` (`str`): String representation of the tracker class name, such as "TensorBoard" `requires_logging_directory` - (`bool`): Whether the logger requires a directory to store their logs. `tracker` (`object`): Should return internal - tracking mechanism used by a tracker class (such as the `run` for wandb) - - Implementations can also include a `main_process_only` (`bool`) attribute to toggle if relevent logging, init, and - other functions should occur on the main process or across all processes (by default will use `True`) - """ - - main_process_only = True - - def __init__(self, _blank=False): - if not _blank: - err = "" - if not hasattr(self, "name"): - err += "`name`" - if not hasattr(self, "requires_logging_directory"): - if len(err) > 0: - err += ", " - err += "`requires_logging_directory`" - - # as tracker is a @property that relies on post-init - if "tracker" not in dir(self): - if len(err) > 0: - err += ", " - err += "`tracker`" - if len(err) > 0: - raise NotImplementedError( - f"The implementation for this tracker class is missing the following " - f"required attributes. Please define them in the class definition: " - f"{err}" - ) - - def start(self): - """ - Lazy initialization of the tracker inside Accelerator to avoid initializing PartialState before - InitProcessGroupKwargs. - """ - pass - - def store_init_configuration(self, values: dict): - """ - Logs `values` as hyperparameters for the run. Implementations should use the experiment configuration - functionality of a tracking API. - - Args: - values (Dictionary `str` to `bool`, `str`, `float` or `int`): - Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, - `str`, `float`, `int`, or `None`. - """ - pass - - def log(self, values: dict, step: Optional[int], **kwargs): - """ - Logs `values` to the current run. Base `log` implementations of a tracking API should go in here, along with - special behavior for the `step parameter. - - Args: - values (Dictionary `str` to `str`, `float`, or `int`): - Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`. - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - """ - pass - - def finish(self): - """ - Should run any finalizing functions within the tracking API. If the API should not have one, just don't - overwrite that method. - """ - pass - - -class TensorBoardTracker(GeneralTracker): - """ - A `Tracker` class that supports `tensorboard`. Should be initialized at the start of your script. - - Args: - run_name (`str`): - The name of the experiment run - logging_dir (`str`, `os.PathLike`): - Location for TensorBoard logs to be stored. - **kwargs (additional keyword arguments, *optional*): - Additional key word arguments passed along to the `tensorboard.SummaryWriter.__init__` method. - """ - - name = "tensorboard" - requires_logging_directory = True - - def __init__(self, run_name: str, logging_dir: Union[str, os.PathLike], **kwargs): - super().__init__() - self.run_name = run_name - self.logging_dir_param = logging_dir - self.init_kwargs = kwargs - - @on_main_process - def start(self): - try: - from torch.utils import tensorboard - except ModuleNotFoundError: - import tensorboardX as tensorboard - self.logging_dir = os.path.join(self.logging_dir_param, self.run_name) - self.writer = tensorboard.SummaryWriter(self.logging_dir, **self.init_kwargs) - logger.debug(f"Initialized TensorBoard project {self.run_name} logging to {self.logging_dir}") - logger.debug( - "Make sure to log any initial configurations with `self.store_init_configuration` before training!" - ) - - @property - def tracker(self): - return self.writer - - @on_main_process - def store_init_configuration(self, values: dict): - """ - Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. Stores the - hyperparameters in a yaml file for future use. - - Args: - values (Dictionary `str` to `bool`, `str`, `float` or `int`): - Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, - `str`, `float`, `int`, or `None`. - """ - self.writer.add_hparams(values, metric_dict={}) - self.writer.flush() - project_run_name = time.time() - dir_name = os.path.join(self.logging_dir, str(project_run_name)) - os.makedirs(dir_name, exist_ok=True) - with open(os.path.join(dir_name, "hparams.yml"), "w") as outfile: - try: - yaml.dump(values, outfile) - except yaml.representer.RepresenterError: - logger.error("Serialization to store hyperparameters failed") - raise - logger.debug("Stored initial configuration hyperparameters to TensorBoard and hparams yaml file") - - @on_main_process - def log(self, values: dict, step: Optional[int] = None, **kwargs): - """ - Logs `values` to the current run. - - Args: - values (Dictionary `str` to `str`, `float`, `int` or `dict` of `str` to `float`/`int`): - Values to be logged as key-value pairs. The values need to have type `str`, `float`, `int` or `dict` of - `str` to `float`/`int`. - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - kwargs: - Additional key word arguments passed along to either `SummaryWriter.add_scaler`, - `SummaryWriter.add_text`, or `SummaryWriter.add_scalers` method based on the contents of `values`. - """ - values = listify(values) - for k, v in values.items(): - if isinstance(v, (int, float)): - self.writer.add_scalar(k, v, global_step=step, **kwargs) - elif isinstance(v, str): - self.writer.add_text(k, v, global_step=step, **kwargs) - elif isinstance(v, dict): - self.writer.add_scalars(k, v, global_step=step, **kwargs) - self.writer.flush() - logger.debug("Successfully logged to TensorBoard") - - @on_main_process - def log_images(self, values: dict, step: Optional[int], **kwargs): - """ - Logs `images` to the current run. - - Args: - values (Dictionary `str` to `List` of `np.ndarray` or `PIL.Image`): - Values to be logged as key-value pairs. The values need to have type `List` of `np.ndarray` or - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - kwargs: - Additional key word arguments passed along to the `SummaryWriter.add_image` method. - """ - for k, v in values.items(): - self.writer.add_images(k, v, global_step=step, **kwargs) - logger.debug("Successfully logged images to TensorBoard") - - @on_main_process - def finish(self): - """ - Closes `TensorBoard` writer - """ - self.writer.close() - logger.debug("TensorBoard writer closed") - - -class WandBTracker(GeneralTracker): - """ - A `Tracker` class that supports `wandb`. Should be initialized at the start of your script. - - Args: - run_name (`str`): - The name of the experiment run. - **kwargs (additional keyword arguments, *optional*): - Additional key word arguments passed along to the `wandb.init` method. - """ - - name = "wandb" - requires_logging_directory = False - main_process_only = False - - def __init__(self, run_name: str, **kwargs): - super().__init__() - self.run_name = run_name - self.init_kwargs = kwargs - - @on_main_process - def start(self): - import wandb - - self.run = wandb.init(project=self.run_name, **self.init_kwargs) - logger.debug(f"Initialized WandB project {self.run_name}") - logger.debug( - "Make sure to log any initial configurations with `self.store_init_configuration` before training!" - ) - - @property - def tracker(self): - return self.run - - @on_main_process - def store_init_configuration(self, values: dict): - """ - Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. - - Args: - values (Dictionary `str` to `bool`, `str`, `float` or `int`): - Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, - `str`, `float`, `int`, or `None`. - """ - import wandb - - if os.environ.get("WANDB_MODE") == "offline": - # In offline mode, restart wandb with config included - if hasattr(self, "run") and self.run: - self.run.finish() - - init_kwargs = self.init_kwargs.copy() - init_kwargs["config"] = values - self.run = wandb.init(project=self.run_name, **init_kwargs) - else: - wandb.config.update(values, allow_val_change=True) - logger.debug("Stored initial configuration hyperparameters to WandB") - - @on_main_process - def log(self, values: dict, step: Optional[int] = None, **kwargs): - """ - Logs `values` to the current run. - - Args: - values (Dictionary `str` to `str`, `float`, `int` or `dict` of `str` to `float`/`int`): - Values to be logged as key-value pairs. The values need to have type `str`, `float`, `int` or `dict` of - `str` to `float`/`int`. - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - kwargs: - Additional key word arguments passed along to the `wandb.log` method. - """ - self.run.log(values, step=step, **kwargs) - logger.debug("Successfully logged to WandB") - - @on_main_process - def log_images(self, values: dict, step: Optional[int] = None, **kwargs): - """ - Logs `images` to the current run. - - Args: - values (Dictionary `str` to `List` of `np.ndarray` or `PIL.Image`): - Values to be logged as key-value pairs. The values need to have type `List` of `np.ndarray` or - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - kwargs: - Additional key word arguments passed along to the `wandb.log` method. - """ - import wandb - - for k, v in values.items(): - self.log({k: [wandb.Image(image) for image in v]}, step=step, **kwargs) - logger.debug("Successfully logged images to WandB") - - @on_main_process - def log_table( - self, - table_name: str, - columns: list[str] = None, - data: list[list[Any]] = None, - dataframe: Any = None, - step: Optional[int] = None, - **kwargs, - ): - """ - Log a Table containing any object type (text, image, audio, video, molecule, html, etc). Can be defined either - with `columns` and `data` or with `dataframe`. - - Args: - table_name (`str`): - The name to give to the logged table on the wandb workspace - columns (list of `str`, *optional*): - The name of the columns on the table - data (List of List of Any data type, *optional*): - The data to be logged in the table - dataframe (Any data type, *optional*): - The data to be logged in the table - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - """ - import wandb - - values = {table_name: wandb.Table(columns=columns, data=data, dataframe=dataframe)} - self.log(values, step=step, **kwargs) - - @on_main_process - def finish(self): - """ - Closes `wandb` writer - """ - self.run.finish() - logger.debug("WandB run closed") - - -class TrackioTracker(GeneralTracker): - """ - A `Tracker` class that supports `trackio`. Should be initialized at the start of your script. - - Args: - run_name (`str`): - The name of the experiment run. Will be used as the `project` name when instantiating trackio. - **kwargs (additional keyword arguments, *optional*): - Additional key word arguments passed along to the `trackio.init` method. Refer to this - [init](https://github.com/gradio-app/trackio/blob/814809552310468b13f84f33764f1369b4e5136c/trackio/__init__.py#L22) - to see all supported key word arguments. - """ - - name = "trackio" - requires_logging_directory = False - main_process_only = False - - def __init__(self, run_name: str, **kwargs): - super().__init__() - self.run_name = run_name - self.init_kwargs = kwargs - - @on_main_process - def start(self): - import trackio - - self.run = trackio.init(project=self.run_name, **self.init_kwargs) - logger.debug(f"Initialized trackio project {self.run_name}") - logger.debug( - "Make sure to log any initial configurations with `self.store_init_configuration` before training!" - ) - - @property - def tracker(self): - return self.run - - @on_main_process - def store_init_configuration(self, values: dict): - """ - Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. - - Args: - values (Dictionary `str` to `bool`, `str`, `float` or `int`): - Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, - `str`, `float`, `int`, or `None`. - """ - import trackio - - trackio.config.update(values, allow_val_change=True) - logger.debug("Stored initial configuration hyperparameters to trackio") - - @on_main_process - def log(self, values: dict, step: Optional[int] = None, **kwargs): - """ - Logs `values` to the current run. - - Args: - values (Dictionary `str` to `str`, `float`, `int` or `dict` of `str` to `float`/`int`): - Values to be logged as key-value pairs. The values need to have type `str`, `float`, `int` or `dict` of - `str` to `float`/`int`. - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - kwargs: - Additional key word arguments passed along to the `trackio.log` method. - """ - self.run.log(values, **kwargs) - logger.debug("Successfully logged to trackio") - - @on_main_process - def finish(self): - """ - Closes `trackio` run - """ - self.run.finish() - logger.debug("trackio run closed") - - -class CometMLTracker(GeneralTracker): - """ - A `Tracker` class that supports `comet_ml`. Should be initialized at the start of your script. - - API keys must be stored in a Comet config file. - - Note: - For `comet_ml` versions < 3.41.0, additional keyword arguments are passed to `comet_ml.Experiment` instead: - https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/Experiment/#comet_ml.Experiment.__init__ - - Args: - run_name (`str`): - The name of the experiment run. - **kwargs (additional keyword arguments, *optional*): - Additional key word arguments passed along to the `comet_ml.start` method: - https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/start/ - """ - - name = "comet_ml" - requires_logging_directory = False - - def __init__(self, run_name: str, **kwargs): - super().__init__() - self.run_name = run_name - self.init_kwargs = kwargs - - @on_main_process - def start(self): - import comet_ml - - comet_version = version.parse(comet_ml.__version__) - if compare_versions(comet_version, ">=", "3.41.0"): - self.writer = comet_ml.start(project_name=self.run_name, **self.init_kwargs) - else: - logger.info("Update `comet_ml` (>=3.41.0) for experiment reuse and offline support.") - self.writer = comet_ml.Experiment(project_name=self.run_name, **self.init_kwargs) - - logger.debug(f"Initialized CometML project {self.run_name}") - logger.debug( - "Make sure to log any initial configurations with `self.store_init_configuration` before training!" - ) - - @property - def tracker(self): - return self.writer - - @on_main_process - def store_init_configuration(self, values: dict): - """ - Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. - - Args: - values (Dictionary `str` to `bool`, `str`, `float` or `int`): - Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, - `str`, `float`, `int`, or `None`. - """ - self.writer.log_parameters(values) - logger.debug("Stored initial configuration hyperparameters to Comet") - - @on_main_process - def log(self, values: dict, step: Optional[int] = None, **kwargs): - """ - Logs `values` to the current run. - - Args: - values (Dictionary `str` to `str`, `float`, `int` or `dict` of `str` to `float`/`int`): - Values to be logged as key-value pairs. The values need to have type `str`, `float`, `int` or `dict` of - `str` to `float`/`int`. - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - kwargs: - Additional key word arguments passed along to either `Experiment.log_metric`, `Experiment.log_other`, - or `Experiment.log_metrics` method based on the contents of `values`. - """ - if step is not None: - self.writer.set_step(step) - for k, v in values.items(): - if isinstance(v, (int, float)): - self.writer.log_metric(k, v, step=step, **kwargs) - elif isinstance(v, str): - self.writer.log_other(k, v, **kwargs) - elif isinstance(v, dict): - self.writer.log_metrics(v, step=step, **kwargs) - logger.debug("Successfully logged to Comet") - - @on_main_process - def finish(self): - """ - Flush `comet-ml` writer - """ - self.writer.end() - logger.debug("Comet run flushed") - - -class AimTracker(GeneralTracker): - """ - A `Tracker` class that supports `aim`. Should be initialized at the start of your script. - - Args: - run_name (`str`): - The name of the experiment run. - **kwargs (additional keyword arguments, *optional*): - Additional key word arguments passed along to the `Run.__init__` method. - """ - - name = "aim" - requires_logging_directory = True - - def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs): - super().__init__() - self.run_name = run_name - self.aim_repo_path = logging_dir - self.init_kwargs = kwargs - - @on_main_process - def start(self): - from aim import Run - - self.writer = Run(repo=self.aim_repo_path, **self.init_kwargs) - self.writer.name = self.run_name - logger.debug(f"Initialized Aim project {self.run_name}") - logger.debug( - "Make sure to log any initial configurations with `self.store_init_configuration` before training!" - ) - - @property - def tracker(self): - return self.writer - - @on_main_process - def store_init_configuration(self, values: dict): - """ - Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. - - Args: - values (`dict`): - Values to be stored as initial hyperparameters as key-value pairs. - """ - self.writer["hparams"] = values - - @on_main_process - def log(self, values: dict, step: Optional[int], **kwargs): - """ - Logs `values` to the current run. - - Args: - values (`dict`): - Values to be logged as key-value pairs. - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - kwargs: - Additional key word arguments passed along to the `Run.track` method. - """ - # Note: replace this with the dictionary support when merged - for key, value in values.items(): - self.writer.track(value, name=key, step=step, **kwargs) - - @on_main_process - def log_images(self, values: dict, step: Optional[int] = None, kwargs: Optional[dict[str, dict]] = None): - """ - Logs `images` to the current run. - - Args: - values (`Dict[str, Union[np.ndarray, PIL.Image, Tuple[np.ndarray, str], Tuple[PIL.Image, str]]]`): - Values to be logged as key-value pairs. The values need to have type `np.ndarray` or PIL.Image. If a - tuple is provided, the first element should be the image and the second element should be the caption. - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - kwargs (`Dict[str, dict]`): - Additional key word arguments passed along to the `Run.Image` and `Run.track` method specified by the - keys `aim_image` and `track`, respectively. - """ - import aim - - aim_image_kw = {} - track_kw = {} - - if kwargs is not None: - aim_image_kw = kwargs.get("aim_image", {}) - track_kw = kwargs.get("track", {}) - - for key, value in values.items(): - if isinstance(value, tuple): - img, caption = value - else: - img, caption = value, "" - aim_image = aim.Image(img, caption=caption, **aim_image_kw) - self.writer.track(aim_image, name=key, step=step, **track_kw) - - @on_main_process - def finish(self): - """ - Closes `aim` writer - """ - self.writer.close() - - -class MLflowTracker(GeneralTracker): - """ - A `Tracker` class that supports `mlflow`. Should be initialized at the start of your script. - - Args: - experiment_name (`str`, *optional*): - Name of the experiment. Environment variable MLFLOW_EXPERIMENT_NAME has priority over this argument. - logging_dir (`str` or `os.PathLike`, defaults to `"."`): - Location for mlflow logs to be stored. - run_id (`str`, *optional*): - If specified, get the run with the specified UUID and log parameters and metrics under that run. The run’s - end time is unset and its status is set to running, but the run’s other attributes (source_version, - source_type, etc.) are not changed. Environment variable MLFLOW_RUN_ID has priority over this argument. - tags (`Dict[str, str]`, *optional*): - An optional `dict` of `str` keys and values, or a `str` dump from a `dict`, to set as tags on the run. If a - run is being resumed, these tags are set on the resumed run. If a new run is being created, these tags are - set on the new run. Environment variable MLFLOW_TAGS has priority over this argument. - nested_run (`bool`, *optional*, defaults to `False`): - Controls whether run is nested in parent run. True creates a nested run. Environment variable - MLFLOW_NESTED_RUN has priority over this argument. - run_name (`str`, *optional*): - Name of new run (stored as a mlflow.runName tag). Used only when `run_id` is unspecified. - description (`str`, *optional*): - An optional string that populates the description box of the run. If a run is being resumed, the - description is set on the resumed run. If a new run is being created, the description is set on the new - run. - """ - - name = "mlflow" - requires_logging_directory = False - - def __init__( - self, - experiment_name: str = None, - logging_dir: Optional[Union[str, os.PathLike]] = None, - run_id: Optional[str] = None, - tags: Optional[Union[dict[str, Any], str]] = None, - nested_run: Optional[bool] = False, - run_name: Optional[str] = None, - description: Optional[str] = None, - ): - experiment_name = os.environ.get("MLFLOW_EXPERIMENT_NAME", experiment_name) - run_id = os.environ.get("MLFLOW_RUN_ID", run_id) - tags = os.environ.get("MLFLOW_TAGS", tags) - if isinstance(tags, str): - tags = json.loads(tags) - - nested_run = os.environ.get("MLFLOW_NESTED_RUN", nested_run) - - self.experiment_name = experiment_name - self.logging_dir = logging_dir - self.run_id = run_id - self.tags = tags - self.nested_run = nested_run - self.run_name = run_name - self.description = description - - @on_main_process - def start(self): - import mlflow - - exps = mlflow.search_experiments(filter_string=f"name = '{self.experiment_name}'") - if len(exps) > 0: - if len(exps) > 1: - logger.warning("Multiple experiments with the same name found. Using first one.") - experiment_id = exps[0].experiment_id - else: - experiment_id = mlflow.create_experiment( - name=self.experiment_name, - artifact_location=self.logging_dir, - tags=self.tags, - ) - - self.active_run = mlflow.start_run( - run_id=self.run_id, - experiment_id=experiment_id, - run_name=self.run_name, - nested=self.nested_run, - tags=self.tags, - description=self.description, - ) - - logger.debug(f"Initialized mlflow experiment {self.experiment_name}") - logger.debug( - "Make sure to log any initial configurations with `self.store_init_configuration` before training!" - ) - - @property - def tracker(self): - return self.active_run - - @on_main_process - def store_init_configuration(self, values: dict): - """ - Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. - - Args: - values (`dict`): - Values to be stored as initial hyperparameters as key-value pairs. - """ - import mlflow - - for name, value in list(values.items()): - # internally, all values are converted to str in MLflow - if len(str(value)) > mlflow.utils.validation.MAX_PARAM_VAL_LENGTH: - logger.warning_once( - f'Accelerate is attempting to log a value of "{value}" for key "{name}" as a parameter. MLflow\'s' - f" log_param() only accepts values no longer than {mlflow.utils.validation.MAX_PARAM_VAL_LENGTH} characters so we dropped this attribute." - ) - del values[name] - - values_list = list(values.items()) - - # MLflow cannot log more than 100 values in one go, so we have to split it - for i in range(0, len(values_list), mlflow.utils.validation.MAX_PARAMS_TAGS_PER_BATCH): - mlflow.log_params(dict(values_list[i : i + mlflow.utils.validation.MAX_PARAMS_TAGS_PER_BATCH])) - - logger.debug("Stored initial configuration hyperparameters to MLflow") - - @on_main_process - def log(self, values: dict, step: Optional[int]): - """ - Logs `values` to the current run. - - Args: - values (`dict`): - Values to be logged as key-value pairs. - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - """ - metrics = {} - for k, v in values.items(): - if isinstance(v, (int, float)): - metrics[k] = v - else: - logger.warning_once( - f'MLflowTracker is attempting to log a value of "{v}" of type {type(v)} for key "{k}" as a metric. ' - "MLflow's log_metric() only accepts float and int types so we dropped this attribute." - ) - import mlflow - - mlflow.log_metrics(metrics, step=step) - logger.debug("Successfully logged to mlflow") - - @on_main_process - def log_figure(self, figure: Any, artifact_file: str, **save_kwargs): - """ - Logs an figure to the current run. - - Args: - figure (Any): - The figure to be logged. - artifact_file (`str`, *optional*): - The run-relative artifact file path in posixpath format to which the image is saved. - If not provided, the image is saved to a default location. - **kwargs: - Additional keyword arguments passed to the underlying mlflow.log_image function. - """ - import mlflow - - mlflow.log_figure(figure=figure, artifact_file=artifact_file, **save_kwargs) - logger.debug("Successfully logged image to mlflow") - - @on_main_process - def log_artifacts(self, local_dir: str, artifact_path: Optional[str] = None): - """ - Logs an artifacts (all content of a dir) to the current run. - - local_dir (`str`): - Path to the directory to be logged as an artifact. - artifact_path (`str`, *optional*): - Directory within the run's artifact directory where the artifact will be logged. If omitted, the - artifact will be logged to the root of the run's artifact directory. The run step. If included, the - artifact will be affiliated with this step. - """ - import mlflow - - mlflow.log_artifacts(local_dir=local_dir, artifact_path=artifact_path) - logger.debug("Successfully logged artofact to mlflow") - - @on_main_process - def log_artifact(self, local_path: str, artifact_path: Optional[str] = None): - """ - Logs an artifact (file) to the current run. - - local_path (`str`): - Path to the file to be logged as an artifact. - artifact_path (`str`, *optional*): - Directory within the run's artifact directory where the artifact will be logged. If omitted, the - artifact will be logged to the root of the run's artifact directory. The run step. If included, the - artifact will be affiliated with this step. - """ - import mlflow - - mlflow.log_artifact(local_path=local_path, artifact_path=artifact_path) - logger.debug("Successfully logged artofact to mlflow") - - @on_main_process - def finish(self): - """ - End the active MLflow run. - """ - import mlflow - - mlflow.end_run() - - -class ClearMLTracker(GeneralTracker): - """ - A `Tracker` class that supports `clearml`. Should be initialized at the start of your script. - - Args: - run_name (`str`, *optional*): - Name of the experiment. Environment variables `CLEARML_PROJECT` and `CLEARML_TASK` have priority over this - argument. - **kwargs (additional keyword arguments, *optional*): - Kwargs passed along to the `Task.__init__` method. - """ - - name = "clearml" - requires_logging_directory = False - - def __init__(self, run_name: str = None, **kwargs): - super().__init__() - self.user_provided_run_name = run_name - self._initialized_externally = False - self.init_kwargs = kwargs - - @on_main_process - def start(self): - from clearml import Task - - current_task = Task.current_task() - if current_task: - self._initialized_externally = True - self.task = current_task - return - - task_init_args = {**self.init_kwargs} - task_init_args.setdefault("project_name", os.environ.get("CLEARML_PROJECT", self.user_provided_run_name)) - task_init_args.setdefault("task_name", os.environ.get("CLEARML_TASK", self.user_provided_run_name)) - self.task = Task.init(**task_init_args) - - @property - def tracker(self): - return self.task - - @on_main_process - def store_init_configuration(self, values: dict): - """ - Connect configuration dictionary to the Task object. Should be run at the beginning of your experiment. - - Args: - values (`dict`): - Values to be stored as initial hyperparameters as key-value pairs. - """ - return self.task.connect_configuration(values) - - @on_main_process - def log(self, values: dict[str, Union[int, float]], step: Optional[int] = None, **kwargs): - """ - Logs `values` dictionary to the current run. The dictionary keys must be strings. The dictionary values must be - ints or floats - - Args: - values (`Dict[str, Union[int, float]]`): - Values to be logged as key-value pairs. If the key starts with 'eval_'/'test_'/'train_', the value will - be reported under the 'eval'/'test'/'train' series and the respective prefix will be removed. - Otherwise, the value will be reported under the 'train' series, and no prefix will be removed. - step (`int`, *optional*): - If specified, the values will be reported as scalars, with the iteration number equal to `step`. - Otherwise they will be reported as single values. - kwargs: - Additional key word arguments passed along to the `clearml.Logger.report_single_value` or - `clearml.Logger.report_scalar` methods. - """ - clearml_logger = self.task.get_logger() - for k, v in values.items(): - if not isinstance(v, (int, float)): - logger.warning_once( - "Accelerator is attempting to log a value of " - f'"{v}" of type {type(v)} for key "{k}" as a scalar. ' - "This invocation of ClearML logger's report_scalar() " - "is incorrect so we dropped this attribute." - ) - continue - if step is None: - clearml_logger.report_single_value(name=k, value=v, **kwargs) - continue - title, series = ClearMLTracker._get_title_series(k) - clearml_logger.report_scalar(title=title, series=series, value=v, iteration=step, **kwargs) - - @on_main_process - def log_images(self, values: dict, step: Optional[int] = None, **kwargs): - """ - Logs `images` to the current run. - - Args: - values (`Dict[str, List[Union[np.ndarray, PIL.Image]]`): - Values to be logged as key-value pairs. The values need to have type `List` of `np.ndarray` or - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - kwargs: - Additional key word arguments passed along to the `clearml.Logger.report_image` method. - """ - clearml_logger = self.task.get_logger() - for k, v in values.items(): - title, series = ClearMLTracker._get_title_series(k) - clearml_logger.report_image(title=title, series=series, iteration=step, image=v, **kwargs) - - @on_main_process - def log_table( - self, - table_name: str, - columns: list[str] = None, - data: list[list[Any]] = None, - dataframe: Any = None, - step: Optional[int] = None, - **kwargs, - ): - """ - Log a Table to the task. Can be defined eitherwith `columns` and `data` or with `dataframe`. - - Args: - table_name (`str`): - The name of the table - columns (list of `str`, *optional*): - The name of the columns on the table - data (List of List of Any data type, *optional*): - The data to be logged in the table. If `columns` is not specified, then the first entry in data will be - the name of the columns of the table - dataframe (Any data type, *optional*): - The data to be logged in the table - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - kwargs: - Additional key word arguments passed along to the `clearml.Logger.report_table` method. - """ - to_report = dataframe - if dataframe is None: - if data is None: - raise ValueError( - "`ClearMLTracker.log_table` requires that `data` to be supplied if `dataframe` is `None`" - ) - to_report = [columns] + data if columns else data - title, series = ClearMLTracker._get_title_series(table_name) - self.task.get_logger().report_table(title=title, series=series, table_plot=to_report, iteration=step, **kwargs) - - @on_main_process - def finish(self): - """ - Close the ClearML task. If the task was initialized externally (e.g. by manually calling `Task.init`), this - function is a noop - """ - if self.task and not self._initialized_externally: - self.task.close() - - @staticmethod - def _get_title_series(name): - for prefix in ["eval", "test", "train"]: - if name.startswith(prefix + "_"): - return name[len(prefix) + 1 :], prefix - return name, "train" - - -class DVCLiveTracker(GeneralTracker): - """ - A `Tracker` class that supports `dvclive`. Should be initialized at the start of your script. - - Args: - run_name (`str`, *optional*): - Ignored for dvclive. See `kwargs` instead. - kwargs: - Additional key word arguments passed along to [`dvclive.Live()`](https://dvc.org/doc/dvclive/live). - - Example: - - ```py - from accelerate import Accelerator - - accelerator = Accelerator(log_with="dvclive") - accelerator.init_trackers(project_name="my_project", init_kwargs={"dvclive": {"dir": "my_directory"}}) - ``` - """ - - name = "dvclive" - requires_logging_directory = False - - def __init__(self, run_name: Optional[str] = None, live: Optional[Any] = None, **kwargs): - super().__init__() - self.live = live - self.init_kwargs = kwargs - - @on_main_process - def start(self): - from dvclive import Live - - self.live = self.live if self.live is not None else Live(**self.init_kwargs) - - @property - def tracker(self): - return self.live - - @on_main_process - def store_init_configuration(self, values: dict): - """ - Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. Stores the - hyperparameters in a yaml file for future use. - - Args: - values (Dictionary `str` to `bool`, `str`, `float`, `int`, or a List or Dict of those types): - Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, - `str`, `float`, or `int`. - """ - self.live.log_params(values) - - @on_main_process - def log(self, values: dict, step: Optional[int] = None, **kwargs): - """ - Logs `values` to the current run. - - Args: - values (Dictionary `str` to `str`, `float`, or `int`): - Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`. - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - kwargs: - Additional key word arguments passed along to `dvclive.Live.log_metric()`. - """ - from dvclive.plots import Metric - - if step is not None: - self.live.step = step - for k, v in values.items(): - if Metric.could_log(v): - self.live.log_metric(k, v, **kwargs) - else: - logger.warning_once( - "Accelerator attempted to log a value of " - f'"{v}" of type {type(v)} for key "{k}" as a scalar. ' - "This invocation of DVCLive's Live.log_metric() " - "is incorrect so we dropped this attribute." - ) - self.live.next_step() - - @on_main_process - def finish(self): - """ - Closes `dvclive.Live()`. - """ - self.live.end() - - -class SwanLabTracker(GeneralTracker): - """ - A `Tracker` class that supports `swanlab`. Should be initialized at the start of your script. - - Args: - run_name (`str`): - The name of the experiment run. - **kwargs (additional keyword arguments, *optional*): - Additional key word arguments passed along to the `swanlab.init` method. - """ - - name = "swanlab" - requires_logging_directory = False - main_process_only = False - - def __init__(self, run_name: str, **kwargs): - super().__init__() - self.run_name = run_name - self.init_kwargs = kwargs - - @on_main_process - def start(self): - import swanlab - - self.run = swanlab.init(project=self.run_name, **self.init_kwargs) - swanlab.config["FRAMEWORK"] = "🤗Accelerate" # add accelerate logo in config - logger.debug(f"Initialized SwanLab project {self.run_name}") - logger.debug( - "Make sure to log any initial configurations with `self.store_init_configuration` before training!" - ) - - @property - def tracker(self): - return self.run - - @on_main_process - def store_init_configuration(self, values: dict): - """ - Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. - - Args: - values (Dictionary `str` to `bool`, `str`, `float` or `int`): - Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`, - `str`, `float`, `int`, or `None`. - """ - import swanlab - - swanlab.config.update(values, allow_val_change=True) - logger.debug("Stored initial configuration hyperparameters to SwanLab") - - @on_main_process - def log(self, values: dict, step: Optional[int] = None, **kwargs): - """ - Logs `values` to the current run. - - Args: - data : Dict[str, DataType] - Data must be a dict. The key must be a string with 0-9, a-z, A-Z, " ", "_", "-", "/". The value must be a - `float`, `float convertible object`, `int` or `swanlab.data.BaseType`. - step : int, optional - The step number of the current data, if not provided, it will be automatically incremented. - If step is duplicated, the data will be ignored. - kwargs: - Additional key word arguments passed along to the `swanlab.log` method. Likes: - print_to_console : bool, optional - Whether to print the data to the console, the default is False. - """ - self.run.log(values, step=step, **kwargs) - logger.debug("Successfully logged to SwanLab") - - @on_main_process - def log_images(self, values: dict, step: Optional[int] = None, **kwargs): - """ - Logs `images` to the current run. - - Args: - values (Dictionary `str` to `List` of `np.ndarray` or `PIL.Image`): - Values to be logged as key-value pairs. The values need to have type `List` of `np.ndarray` or - step (`int`, *optional*): - The run step. If included, the log will be affiliated with this step. - kwargs: - Additional key word arguments passed along to the `swanlab.log` method. Likes: - print_to_console : bool, optional - Whether to print the data to the console, the default is False. - """ - import swanlab - - for k, v in values.items(): - self.log({k: [swanlab.Image(image) for image in v]}, step=step, **kwargs) - logger.debug("Successfully logged images to SwanLab") - - @on_main_process - def finish(self): - """ - Closes `swanlab` writer - """ - self.run.finish() - logger.debug("SwanLab run closed") - - -LOGGER_TYPE_TO_CLASS = { - "aim": AimTracker, - "comet_ml": CometMLTracker, - "mlflow": MLflowTracker, - "tensorboard": TensorBoardTracker, - "wandb": WandBTracker, - "clearml": ClearMLTracker, - "dvclive": DVCLiveTracker, - "swanlab": SwanLabTracker, - "trackio": TrackioTracker, -} - - -def filter_trackers( - log_with: list[Union[str, LoggerType, GeneralTracker]], - logging_dir: Union[str, os.PathLike] = None, -): - """ - Takes in a list of potential tracker types and checks that: - - The tracker wanted is available in that environment - - Filters out repeats of tracker types - - If `all` is in `log_with`, will return all trackers in the environment - - If a tracker requires a `logging_dir`, ensures that `logging_dir` is not `None` - - Args: - log_with (list of `str`, [`~utils.LoggerType`] or [`~tracking.GeneralTracker`], *optional*): - A list of loggers to be setup for experiment tracking. Should be one or several of: - - - `"all"` - - `"tensorboard"` - - `"wandb"` - - `"trackio"` - - `"aim"` - - `"comet_ml"` - - `"mlflow"` - - `"dvclive"` - - `"swanlab"` - If `"all"` is selected, will pick up all available trackers in the environment and initialize them. Can - also accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`. - logging_dir (`str`, `os.PathLike`, *optional*): - A path to a directory for storing logs of locally-compatible loggers. - """ - loggers = [] - if log_with is not None: - if not isinstance(log_with, (list, tuple)): - log_with = [log_with] - if "all" in log_with or LoggerType.ALL in log_with: - loggers = [o for o in log_with if issubclass(type(o), GeneralTracker)] + get_available_trackers() - else: - for log_type in log_with: - if log_type not in LoggerType and not issubclass(type(log_type), GeneralTracker): - raise ValueError(f"Unsupported logging capability: {log_type}. Choose between {LoggerType.list()}") - if issubclass(type(log_type), GeneralTracker): - loggers.append(log_type) - else: - log_type = LoggerType(log_type) - if log_type not in loggers: - if log_type in get_available_trackers(): - tracker_init = LOGGER_TYPE_TO_CLASS[str(log_type)] - if tracker_init.requires_logging_directory: - if logging_dir is None: - raise ValueError( - f"Logging with `{log_type}` requires a `logging_dir` to be passed in." - ) - loggers.append(log_type) - else: - logger.debug(f"Tried adding logger {log_type}, but package is unavailable in the system.") - - return loggers diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__init__.py b/venv/lib/python3.10/site-packages/accelerate/utils/__init__.py deleted file mode 100644 index 34500779793601dcc4245909ed2e24864eb17f02..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/__init__.py +++ /dev/null @@ -1,300 +0,0 @@ -# Copyright 2022 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 .ao import convert_model_to_fp8_ao, filter_first_and_last_linear_layers, has_ao_layers -from .constants import ( - MITA_PROFILING_AVAILABLE_PYTORCH_VERSION, - MODEL_NAME, - OPTIMIZER_NAME, - PROFILE_PATTERN_NAME, - RNG_STATE_NAME, - SAFE_MODEL_NAME, - SAFE_WEIGHTS_INDEX_NAME, - SAFE_WEIGHTS_NAME, - SAFE_WEIGHTS_PATTERN_NAME, - SAMPLER_NAME, - SCALER_NAME, - SCHEDULER_NAME, - TORCH_DISTRIBUTED_OPERATION_TYPES, - TORCH_LAUNCH_PARAMS, - WEIGHTS_INDEX_NAME, - WEIGHTS_NAME, - WEIGHTS_PATTERN_NAME, - XPU_PROFILING_AVAILABLE_PYTORCH_VERSION, -) -from .dataclasses import ( - AORecipeKwargs, - AutocastKwargs, - BnbQuantizationConfig, - ComputeEnvironment, - CustomDtype, - DataLoaderConfiguration, - DDPCommunicationHookType, - DeepSpeedPlugin, - DistributedDataParallelKwargs, - DistributedType, - DynamoBackend, - FP8RecipeKwargs, - FullyShardedDataParallelPlugin, - GradientAccumulationPlugin, - GradScalerKwargs, - InitProcessGroupKwargs, - KwargsHandler, - LoggerType, - MegatronLMPlugin, - MSAMPRecipeKwargs, - PrecisionType, - ProfileKwargs, - ProjectConfiguration, - RNGType, - SageMakerDistributedType, - TensorInformation, - TERecipeKwargs, - TorchDynamoPlugin, - TorchTensorParallelPlugin, - add_model_config_to_megatron_parser, -) -from .environment import ( - are_libraries_initialized, - check_cuda_fp8_capability, - check_cuda_p2p_ib_support, - clear_environment, - convert_dict_to_env_variables, - get_cpu_distributed_information, - get_gpu_info, - get_int_from_env, - parse_choice_from_env, - parse_flag_from_env, - patch_environment, - purge_accelerate_environment, - set_numa_affinity, - str_to_bool, -) -from .imports import ( - deepspeed_required, - get_ccl_version, - is_4bit_bnb_available, - is_8bit_bnb_available, - is_aim_available, - is_bf16_available, - is_bitsandbytes_multi_backend_available, - is_bnb_available, - is_boto3_available, - is_ccl_available, - is_clearml_available, - is_comet_ml_available, - is_cuda_available, - is_datasets_available, - is_deepspeed_available, - is_dvclive_available, - is_fp8_available, - is_fp16_available, - is_habana_gaudi1, - is_hpu_available, - is_import_timer_available, - is_ipex_available, - is_lomo_available, - is_matplotlib_available, - is_megatron_lm_available, - is_mlflow_available, - is_mlu_available, - is_mps_available, - is_msamp_available, - is_musa_available, - is_npu_available, - is_pandas_available, - is_peft_available, - is_pippy_available, - is_pynvml_available, - is_pytest_available, - is_rich_available, - is_sagemaker_available, - is_schedulefree_available, - is_sdaa_available, - is_swanlab_available, - is_tensorboard_available, - is_timm_available, - is_torch_xla_available, - is_torchao_available, - is_torchdata_available, - is_torchdata_stateful_dataloader_available, - is_torchvision_available, - is_trackio_available, - is_transformer_engine_available, - is_transformers_available, - is_triton_available, - is_wandb_available, - is_weights_only_available, - is_xccl_available, - is_xpu_available, - torchao_required, -) -from .modeling import ( - align_module_device, - calculate_maximum_sizes, - check_device_map, - check_tied_parameters_in_config, - check_tied_parameters_on_same_device, - compute_module_sizes, - convert_file_size_to_int, - dtype_byte_size, - find_tied_parameters, - get_balanced_memory, - get_grad_scaler, - get_max_layer_size, - get_max_memory, - get_mixed_precision_context_manager, - has_offloaded_params, - id_tensor_storage, - infer_auto_device_map, - is_peft_model, - load_checkpoint_in_model, - load_offloaded_weights, - load_state_dict, - named_module_tensors, - retie_parameters, - set_module_tensor_to_device, -) -from .offload import ( - OffloadedWeightsLoader, - PrefixedDataset, - extract_submodules_state_dict, - load_offloaded_weight, - offload_state_dict, - offload_weight, - save_offload_index, -) -from .operations import ( - CannotPadNestedTensorWarning, - GatheredParameters, - broadcast, - broadcast_object_list, - concatenate, - convert_outputs_to_fp32, - convert_to_fp32, - copy_tensor_to_devices, - find_batch_size, - find_device, - gather, - gather_object, - get_data_structure, - honor_type, - ignorant_find_batch_size, - initialize_tensors, - is_namedtuple, - is_tensor_information, - is_torch_tensor, - listify, - pad_across_processes, - pad_input_tensors, - recursively_apply, - reduce, - send_to_device, - slice_tensors, -) -from .versions import compare_versions, is_torch_version - - -if is_deepspeed_available(): - from .deepspeed import ( - DeepSpeedEngineWrapper, - DeepSpeedOptimizerWrapper, - DeepSpeedSchedulerWrapper, - DummyOptim, - DummyScheduler, - HfDeepSpeedConfig, - get_active_deepspeed_plugin, - map_pytorch_optim_to_deepspeed, - ) - -from .bnb import has_4bit_bnb_layers, load_and_quantize_model -from .fsdp_utils import ( - disable_fsdp_ram_efficient_loading, - enable_fsdp_ram_efficient_loading, - ensure_weights_retied, - fsdp2_apply_ac, - fsdp2_canonicalize_names, - fsdp2_load_full_state_dict, - fsdp2_prepare_model, - fsdp2_switch_optimizer_parameters, - get_fsdp2_grad_scaler, - load_fsdp_model, - load_fsdp_optimizer, - merge_fsdp_weights, - save_fsdp_model, - save_fsdp_optimizer, -) -from .launch import ( - PrepareForLaunch, - _filter_args, - prepare_deepspeed_cmd_env, - prepare_multi_gpu_env, - prepare_sagemager_args_inputs, - prepare_simple_launcher_cmd_env, - prepare_tpu, -) - -# For docs -from .megatron_lm import ( - AbstractTrainStep, - BertTrainStep, - GPTTrainStep, - MegatronLMDummyDataLoader, - MegatronLMDummyScheduler, - T5TrainStep, - avg_losses_across_data_parallel_group, -) - - -if is_megatron_lm_available(): - from .megatron_lm import ( - MegatronEngine, - MegatronLMOptimizerWrapper, - MegatronLMSchedulerWrapper, - gather_across_data_parallel_groups, - ) - from .megatron_lm import initialize as megatron_lm_initialize - from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader - from .megatron_lm import prepare_model_optimizer_scheduler as megatron_lm_prepare_model_optimizer_scheduler - from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer - from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler -from .memory import find_executable_batch_size, release_memory -from .other import ( - check_os_kernel, - clean_state_dict_for_safetensors, - compile_regions, - compile_regions_deepspeed, - convert_bytes, - extract_model_from_parallel, - get_module_children_bottom_up, - get_pretty_name, - has_compiled_regions, - is_compiled_module, - is_port_in_use, - load, - merge_dicts, - model_has_dtensor, - recursive_getattr, - save, - wait_for_everyone, - write_basic_config, -) -from .random import set_seed, synchronize_rng_state, synchronize_rng_states -from .torch_xla import install_xla -from .tqdm import tqdm -from .transformer_engine import ( - apply_fp8_autowrap, - contextual_fp8_autocast, - convert_model, - has_transformer_engine_layers, -) diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 773307d5972562b529d145d6842a1b1d37bb5337..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/ao.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/ao.cpython-310.pyc deleted file mode 100644 index 410df5b47973818519909db6eecf04d5572f3812..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/ao.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/bnb.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/bnb.cpython-310.pyc deleted file mode 100644 index 89aa419804ae7b30f619c91b11d147252eb5b5d7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/bnb.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/constants.cpython-310.pyc deleted file mode 100644 index 0119472d895e9183a3b244672e7fde8b1a38e14e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/constants.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/dataclasses.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/dataclasses.cpython-310.pyc deleted file mode 100644 index 26bdee290c7cdf9083c41db53f2d5bb0507b6b29..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/dataclasses.cpython-310.pyc +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ffac3ba389e9e46fb79d38e8cdbc3c6c1d5fff78aacb3ec44e3e2c3259bcddd7 -size 103780 diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/deepspeed.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/deepspeed.cpython-310.pyc deleted file mode 100644 index 789efb35cdbd7961cda39da11a6aa5bb6e178eb1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/deepspeed.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/environment.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/environment.cpython-310.pyc deleted file mode 100644 index 2147f91ea5d972f9df9d7ba9948cdbd255d81090..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/environment.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/fsdp_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/fsdp_utils.cpython-310.pyc deleted file mode 100644 index 4207fb9ccea9f868e812802401fdebd0b8ce7d3c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/fsdp_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/imports.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/imports.cpython-310.pyc deleted file mode 100644 index 3dedc382d207f62d6becf560a2a585bccd96de5b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/imports.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/launch.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/launch.cpython-310.pyc deleted file mode 100644 index 5b0f9db883a1ea25c0a5cdea3972a986207dd8b0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/launch.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/megatron_lm.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/megatron_lm.cpython-310.pyc deleted file mode 100644 index 38d579955bf21c1464279871f9bd8db399708c0f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/megatron_lm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/memory.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/memory.cpython-310.pyc deleted file mode 100644 index db745c4b54ce1a0c11e9c55192b99c015312f64e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/memory.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/modeling.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/modeling.cpython-310.pyc deleted file mode 100644 index 3e42e2dd857f777e134fb5b74e0609a023939fc4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/modeling.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/offload.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/offload.cpython-310.pyc deleted file mode 100644 index 73f2b40941c5e63ebdd3f9246e3ce4b42cf25a88..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/offload.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/operations.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/operations.cpython-310.pyc deleted file mode 100644 index fa241d17271707b870edb5fa027be6ebb6e6b480..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/operations.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/other.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/other.cpython-310.pyc deleted file mode 100644 index cbef1ce5c0ca2dcfab8dae64ad8ff205bdbd9cdc..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/other.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/random.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/random.cpython-310.pyc deleted file mode 100644 index 962ef12444a3d631ecb8506156d47164b5003592..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/random.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/rich.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/rich.cpython-310.pyc deleted file mode 100644 index ffa477e80a4193c53a417b44b285312717a737a1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/rich.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/torch_xla.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/torch_xla.cpython-310.pyc deleted file mode 100644 index ab3ac43e81310b0d82be2002be92d6eea1c3d758..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/torch_xla.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/tqdm.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/tqdm.cpython-310.pyc deleted file mode 100644 index fbbda629ceacc3a9cd5a18670f302ffcdb25e013..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/tqdm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/transformer_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/transformer_engine.cpython-310.pyc deleted file mode 100644 index 59914ef8727a7223fee5f1ab50227fe2b659e519..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/transformer_engine.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/versions.cpython-310.pyc b/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/versions.cpython-310.pyc deleted file mode 100644 index 32242e60c7e0b4aa30300be84acf0d52c90feb58..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/accelerate/utils/__pycache__/versions.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/ao.py b/venv/lib/python3.10/site-packages/accelerate/utils/ao.py deleted file mode 100644 index 73155615b768a01f709bb8e0857617ba58a6ec83..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/ao.py +++ /dev/null @@ -1,140 +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. - -""" -Needed utilities for torchao FP8 training. -""" - -from functools import partial -from typing import TYPE_CHECKING, Callable, Optional - -import torch - -from .imports import is_torchao_available, torchao_required - - -if TYPE_CHECKING: - if is_torchao_available(): - from torchao.float8.float8_linear import Float8LinearConfig - - -def find_first_last_linear_layers(model: torch.nn.Module): - """ - Finds the first and last linear layer names in a model. - - This is needed during FP8 to avoid issues with instability by keeping the first and last layers unquantized. - - Ref: https://x.com/xariusrke/status/1826669142604141052 - """ - first_linear, last_linear = None, None - for name, module in model.named_modules(): - if isinstance(module, torch.nn.Linear): - if first_linear is None: - first_linear = name - last_linear = name - return first_linear, last_linear - - -def filter_linear_layers(module, fqn: str, layers_to_filter: list[str]) -> bool: - """ - A function which will check if `module` is: - - a `torch.nn.Linear` layer - - has in_features and out_features divisible by 16 - - is not part of `layers_to_filter` - - Args: - module (`torch.nn.Module`): - The module to check. - fqn (`str`): - The fully qualified name of the layer. - layers_to_filter (`List[str]`): - The list of layers to filter. - """ - if isinstance(module, torch.nn.Linear): - if module.in_features % 16 != 0 or module.out_features % 16 != 0: - return False - if fqn in layers_to_filter: - return False - return True - - -def filter_first_and_last_linear_layers(module, fqn: str) -> bool: - """ - A filter function which will filter out all linear layers except the first and last. - - - - For stability reasons, we skip the first and last linear layers Otherwise can lead to the model not training or - converging properly - - - - Args: - module (`torch.nn.Module`): - The module to check. - fqn (`str`): - The fully qualified name of the layer. - """ - first_linear, last_linear = find_first_last_linear_layers(module) - return filter_linear_layers(module, fqn, layers_to_filter=[first_linear, last_linear]) - - -@torchao_required -def has_ao_layers(model: torch.nn.Module): - from torchao.float8.float8_linear import Float8Linear - - for name, module in model.named_modules(): - if isinstance(module, Float8Linear): - return True - return False - - -@torchao_required -def convert_model_to_fp8_ao( - model: torch.nn.Module, - config: Optional["Float8LinearConfig"] = None, - module_filter_func: Optional[Callable] = filter_first_and_last_linear_layers, -): - """ - Converts all `nn.Linear` layers in the model (except the first and last) to torchao's `Float8Linear` layer inplace. - - Args: - model (`torch.nn.Module`): - The model to convert. - config (`torchao.float8.Float8LinearConfig`, *optional*): - The configuration for the FP8 training. Recommended to utilize - `torchao.float8.recipe_name_to_linear_config` to generate this. In general, the default config should be - sufficient (what is passed when set to `None`). - module_filter_func (`Callable`, *optional*, defaults to `filter_linear_layers`): - Optional function that must take in a module and layer name, and returns a boolean indicating whether the - module should be converted to FP8. Defaults to `filter_linear_layers`. See it for an example. - - Example: - - ```python - from accelerate.utils.ao import convert_model_to_fp8_ao - - model = MyModel() - model.to("cuda") - convert_to_float8_training(model) - - model.train() - ``` - """ - from torchao.float8 import convert_to_float8_training - - first_linear, last_linear = find_first_last_linear_layers(model) - if module_filter_func is None: - module_filter_func = partial(filter_linear_layers, layers_to_filter=[first_linear, last_linear]) - convert_to_float8_training(model, module_filter_fn=module_filter_func, config=config) diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/bnb.py b/venv/lib/python3.10/site-packages/accelerate/utils/bnb.py deleted file mode 100644 index af4aa541233d21193aa5b302ee37c0b57dae4926..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/bnb.py +++ /dev/null @@ -1,469 +0,0 @@ -# Copyright 2023 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 logging -import os -from copy import deepcopy -from typing import Optional, Union - -import torch -import torch.nn as nn - -from accelerate.utils.imports import ( - is_4bit_bnb_available, - is_8bit_bnb_available, -) - -from ..big_modeling import dispatch_model, init_empty_weights -from .dataclasses import BnbQuantizationConfig -from .modeling import ( - find_tied_parameters, - get_balanced_memory, - infer_auto_device_map, - load_checkpoint_in_model, - offload_weight, - set_module_tensor_to_device, -) - - -logger = logging.getLogger(__name__) - - -def load_and_quantize_model( - model: torch.nn.Module, - bnb_quantization_config: BnbQuantizationConfig, - weights_location: Union[str, os.PathLike] = None, - device_map: Optional[dict[str, Union[int, str, torch.device]]] = None, - no_split_module_classes: Optional[list[str]] = None, - max_memory: Optional[dict[Union[int, str], Union[int, str]]] = None, - offload_folder: Optional[Union[str, os.PathLike]] = None, - offload_state_dict: bool = False, -): - """ - This function will quantize the input model with the associated config passed in `bnb_quantization_config`. If the - model is in the meta device, we will load and dispatch the weights according to the `device_map` passed. If the - model is already loaded, we will quantize the model and put the model on the GPU, - - Args: - model (`torch.nn.Module`): - Input model. The model can be already loaded or on the meta device - bnb_quantization_config (`BnbQuantizationConfig`): - The bitsandbytes quantization parameters - weights_location (`str` or `os.PathLike`): - The folder weights_location to load. It can be: - - a path to a file containing a whole model state dict - - a path to a `.json` file containing the index to a sharded checkpoint - - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint. - - a path to a folder containing a unique pytorch_model.bin file. - device_map (`Dict[str, Union[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. - no_split_module_classes (`List[str]`, *optional*): - A list of layer class names that should never be split across device (for instance any layer that has a - residual connection). - max_memory (`Dict`, *optional*): - A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset. - offload_folder (`str` or `os.PathLike`, *optional*): - If the `device_map` contains any value `"disk"`, the folder where we will offload weights. - offload_state_dict (`bool`, *optional*, defaults to `False`): - If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if - the weight of the CPU state dict + the biggest shard does not fit. - - Returns: - `torch.nn.Module`: The quantized model - """ - - load_in_4bit = bnb_quantization_config.load_in_4bit - load_in_8bit = bnb_quantization_config.load_in_8bit - - if load_in_8bit and not is_8bit_bnb_available(): - raise ImportError( - "You have a version of `bitsandbytes` that is not compatible with 8bit quantization," - " make sure you have the latest version of `bitsandbytes` installed." - ) - if load_in_4bit and not is_4bit_bnb_available(): - raise ValueError( - "You have a version of `bitsandbytes` that is not compatible with 4bit quantization," - "make sure you have the latest version of `bitsandbytes` installed." - ) - - modules_on_cpu = [] - # custom device map - if isinstance(device_map, dict) and len(device_map.keys()) > 1: - modules_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]] - - # We keep some modules such as the lm_head in their original dtype for numerical stability reasons - if bnb_quantization_config.skip_modules is None: - bnb_quantization_config.skip_modules = get_keys_to_not_convert(model) - - # add cpu modules to skip modules only for 4-bit modules - if load_in_4bit: - bnb_quantization_config.skip_modules.extend(modules_on_cpu) - modules_to_not_convert = bnb_quantization_config.skip_modules - - # We add the modules we want to keep in full precision - if bnb_quantization_config.keep_in_fp32_modules is None: - bnb_quantization_config.keep_in_fp32_modules = [] - keep_in_fp32_modules = bnb_quantization_config.keep_in_fp32_modules - modules_to_not_convert.extend(keep_in_fp32_modules) - - # compatibility with peft - model.is_loaded_in_4bit = load_in_4bit - model.is_loaded_in_8bit = load_in_8bit - - model_device = get_parameter_device(model) - if model_device.type != "meta": - # quantization of an already loaded model - logger.warning( - "It is not recommended to quantize a loaded model. " - "The model should be instantiated under the `init_empty_weights` context manager." - ) - model = replace_with_bnb_layers(model, bnb_quantization_config, modules_to_not_convert=modules_to_not_convert) - # convert param to the right dtype - dtype = bnb_quantization_config.torch_dtype - for name, param in model.state_dict().items(): - if any(module_to_keep_in_fp32 in name for module_to_keep_in_fp32 in keep_in_fp32_modules): - param.to(torch.float32) - if param.dtype != torch.float32: - name = name.replace(".weight", "").replace(".bias", "") - param = getattr(model, name, None) - if param is not None: - param.to(torch.float32) - elif torch.is_floating_point(param): - param.to(dtype) - if model_device.type == "cuda": - model.cuda(torch.cuda.current_device()) - torch.cuda.empty_cache() - elif torch.cuda.is_available(): - model.to(torch.cuda.current_device()) - elif torch.xpu.is_available(): - model.to(torch.xpu.current_device()) - else: - raise RuntimeError("No GPU or Intel XPU found. A GPU or Intel XPU is needed for quantization.") - logger.info( - f"The model device type is {model_device.type}. However, gpu or intel xpu is needed for quantization." - "We move the model to it." - ) - return model - - elif weights_location is None: - raise RuntimeError( - f"`weights_location` needs to be the folder path containing the weights of the model, but we found {weights_location} " - ) - - else: - with init_empty_weights(): - model = replace_with_bnb_layers( - model, bnb_quantization_config, modules_to_not_convert=modules_to_not_convert - ) - device_map = get_quantized_model_device_map( - model, - bnb_quantization_config, - device_map, - max_memory=max_memory, - no_split_module_classes=no_split_module_classes, - ) - if offload_state_dict is None and device_map is not None and "disk" in device_map.values(): - offload_state_dict = True - - offload = any(x in list(device_map.values()) for x in ["cpu", "disk"]) - - load_checkpoint_in_model( - model, - weights_location, - device_map, - dtype=bnb_quantization_config.torch_dtype, - offload_folder=offload_folder, - offload_state_dict=offload_state_dict, - keep_in_fp32_modules=bnb_quantization_config.keep_in_fp32_modules, - offload_8bit_bnb=load_in_8bit and offload, - ) - return dispatch_model(model, device_map=device_map, offload_dir=offload_folder) - - -def get_quantized_model_device_map( - model, bnb_quantization_config, device_map=None, max_memory=None, no_split_module_classes=None -): - if device_map is None: - if torch.cuda.is_available(): - device_map = {"": torch.cuda.current_device()} - elif torch.xpu.is_available(): - device_map = {"": torch.xpu.current_device()} - else: - raise RuntimeError("No GPU found. A GPU is needed for quantization.") - logger.info("The device_map was not initialized.Setting device_map to `{'':torch.cuda.current_device()}`.") - - if isinstance(device_map, str): - if device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]: - raise ValueError( - "If passing a string for `device_map`, please choose 'auto', 'balanced', 'balanced_low_0' or " - "'sequential'." - ) - - special_dtypes = {} - special_dtypes.update( - { - name: bnb_quantization_config.torch_dtype - for name, _ in model.named_parameters() - if any(m in name for m in bnb_quantization_config.skip_modules) - } - ) - special_dtypes.update( - { - name: torch.float32 - for name, _ in model.named_parameters() - if any(m in name for m in bnb_quantization_config.keep_in_fp32_modules) - } - ) - - kwargs = {} - kwargs["special_dtypes"] = special_dtypes - kwargs["no_split_module_classes"] = no_split_module_classes - kwargs["dtype"] = bnb_quantization_config.target_dtype - - # get max_memory for each device. - if device_map != "sequential": - max_memory = get_balanced_memory( - model, - low_zero=(device_map == "balanced_low_0"), - max_memory=max_memory, - **kwargs, - ) - - kwargs["max_memory"] = max_memory - device_map = infer_auto_device_map(model, **kwargs) - - if isinstance(device_map, dict): - # check if don't have any quantized module on the cpu - modules_not_to_convert = bnb_quantization_config.skip_modules + bnb_quantization_config.keep_in_fp32_modules - - device_map_without_some_modules = { - key: device_map[key] for key in device_map.keys() if key not in modules_not_to_convert - } - for device in ["cpu", "disk"]: - if device in device_map_without_some_modules.values(): - if bnb_quantization_config.load_in_4bit: - raise ValueError( - """ - Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit - the quantized model. If you want to dispatch the model on the CPU or the disk while keeping - these modules in `torch_dtype`, you need to pass a custom `device_map` to - `load_and_quantize_model`. Check - https://huggingface.co/docs/accelerate/main/en/usage_guides/quantization#offload-modules-to-cpu-and-disk - for more details. - """ - ) - else: - logger.info( - "Some modules are are offloaded to the CPU or the disk. Note that these modules will be converted to 8-bit" - ) - del device_map_without_some_modules - return device_map - - -def replace_with_bnb_layers(model, bnb_quantization_config, modules_to_not_convert=None, current_key_name=None): - """ - A helper function to replace all `torch.nn.Linear` modules by `bnb.nn.Linear8bit` modules or by `bnb.nn.Linear4bit` - modules from the `bitsandbytes`library. The function will be run recursively and replace `torch.nn.Linear` modules. - - Parameters: - model (`torch.nn.Module`): - Input model or `torch.nn.Module` as the function is run recursively. - modules_to_not_convert (`List[str]`): - Names of the modules to not quantize convert. In practice we keep the `lm_head` in full precision for - numerical stability reasons. - 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. - """ - - if modules_to_not_convert is None: - modules_to_not_convert = [] - - model, has_been_replaced = _replace_with_bnb_layers( - model, bnb_quantization_config, modules_to_not_convert, current_key_name - ) - if not has_been_replaced: - logger.warning( - "You are loading your model in 8bit or 4bit but no linear modules were found in your model." - " this can happen for some architectures such as gpt2 that uses Conv1D instead of Linear layers." - " Please double check your model architecture, or submit an issue on github if you think this is" - " a bug." - ) - return model - - -def _replace_with_bnb_layers( - model, - bnb_quantization_config, - modules_to_not_convert=None, - current_key_name=None, -): - """ - Private method that wraps the recursion for module replacement. - - Returns the converted model and a boolean that indicates if the conversion has been successfull or not. - """ - # bitsandbytes will initialize CUDA on import, so it needs to be imported lazily - import bitsandbytes as bnb - - has_been_replaced = False - for name, module in model.named_children(): - if current_key_name is None: - current_key_name = [] - current_key_name.append(name) - if isinstance(module, nn.Linear) and name not in modules_to_not_convert: - # Check if the current key is not in the `modules_to_not_convert` - current_key_name_str = ".".join(current_key_name) - proceed = True - for key in modules_to_not_convert: - if ( - (key in current_key_name_str) and (key + "." in current_key_name_str) - ) or key == current_key_name_str: - proceed = False - break - if proceed: - # Load bnb module with empty weight and replace ``nn.Linear` module - if bnb_quantization_config.load_in_8bit: - bnb_module = bnb.nn.Linear8bitLt( - module.in_features, - module.out_features, - module.bias is not None, - has_fp16_weights=False, - threshold=bnb_quantization_config.llm_int8_threshold, - ) - elif bnb_quantization_config.load_in_4bit: - bnb_module = bnb.nn.Linear4bit( - module.in_features, - module.out_features, - module.bias is not None, - bnb_quantization_config.bnb_4bit_compute_dtype, - compress_statistics=bnb_quantization_config.bnb_4bit_use_double_quant, - quant_type=bnb_quantization_config.bnb_4bit_quant_type, - ) - else: - raise ValueError("load_in_8bit and load_in_4bit can't be both False") - bnb_module.weight.data = module.weight.data - if module.bias is not None: - bnb_module.bias.data = module.bias.data - bnb_module.requires_grad_(False) - setattr(model, name, bnb_module) - has_been_replaced = True - if len(list(module.children())) > 0: - _, _has_been_replaced = _replace_with_bnb_layers( - module, bnb_quantization_config, modules_to_not_convert, current_key_name - ) - has_been_replaced = has_been_replaced | _has_been_replaced - # Remove the last key for recursion - current_key_name.pop(-1) - return model, has_been_replaced - - -def get_keys_to_not_convert(model): - r""" - An utility function to get the key of the module to keep in full precision if any For example for CausalLM modules - we may want to keep the lm_head in full precision for numerical stability reasons. For other architectures, we want - to keep the tied weights of the model. The function will return a list of the keys of the modules to not convert in - int8. - - Parameters: - model (`torch.nn.Module`): - Input model - """ - # Create a copy of the model - with init_empty_weights(): - tied_model = deepcopy(model) # this has 0 cost since it is done inside `init_empty_weights` context manager` - - tied_params = find_tied_parameters(tied_model) - # For compatibility with Accelerate < 0.18 - if isinstance(tied_params, dict): - tied_keys = sum(list(tied_params.values()), []) + list(tied_params.keys()) - else: - tied_keys = sum(tied_params, []) - has_tied_params = len(tied_keys) > 0 - - # Check if it is a base model - is_base_model = False - if hasattr(model, "base_model_prefix"): - is_base_model = not hasattr(model, model.base_model_prefix) - - # Ignore this for base models (BertModel, GPT2Model, etc.) - if (not has_tied_params) and is_base_model: - return [] - - # otherwise they have an attached head - list_modules = list(model.named_children()) - list_last_module = [list_modules[-1][0]] - - # add last module together with tied weights - intersection = set(list_last_module) - set(tied_keys) - list_untouched = list(set(tied_keys)) + list(intersection) - - # remove ".weight" from the keys - names_to_remove = [".weight", ".bias"] - filtered_module_names = [] - for name in list_untouched: - for name_to_remove in names_to_remove: - if name_to_remove in name: - name = name.replace(name_to_remove, "") - filtered_module_names.append(name) - - return filtered_module_names - - -def has_4bit_bnb_layers(model): - """Check if we have `bnb.nn.Linear4bit` or `bnb.nn.Linear8bitLt` layers inside our model""" - # bitsandbytes will initialize CUDA on import, so it needs to be imported lazily - import bitsandbytes as bnb - - for m in model.modules(): - if isinstance(m, bnb.nn.Linear4bit): - return True - return False - - -def get_parameter_device(parameter: nn.Module): - return next(parameter.parameters()).device - - -def quantize_and_offload_8bit(model, param, param_name, new_dtype, offload_folder, offload_index, fp16_statistics): - # if it is not quantized, we quantize and offload the quantized weights and the SCB stats - if fp16_statistics is None: - set_module_tensor_to_device(model, param_name, 0, dtype=new_dtype, value=param) - tensor_name = param_name - module = model - if "." in tensor_name: - splits = tensor_name.split(".") - for split in splits[:-1]: - new_module = getattr(module, split) - if new_module is None: - raise ValueError(f"{module} has no attribute {split}.") - module = new_module - tensor_name = splits[-1] - # offload weights - module._parameters[tensor_name].requires_grad = False - offload_weight(module._parameters[tensor_name], param_name, offload_folder, index=offload_index) - if hasattr(module._parameters[tensor_name], "SCB"): - offload_weight( - module._parameters[tensor_name].SCB, - param_name.replace("weight", "SCB"), - offload_folder, - index=offload_index, - ) - else: - offload_weight(param, param_name, offload_folder, index=offload_index) - offload_weight(fp16_statistics, param_name.replace("weight", "SCB"), offload_folder, index=offload_index) - - set_module_tensor_to_device(model, param_name, "meta", dtype=new_dtype, value=torch.empty(*param.size())) diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/constants.py b/venv/lib/python3.10/site-packages/accelerate/utils/constants.py deleted file mode 100644 index b841bdbc26f0bc5f25076a19188fc6c18505f1bb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/constants.py +++ /dev/null @@ -1,103 +0,0 @@ -# Copyright 2022 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 operator as op - -import torch - - -SCALER_NAME = "scaler.pt" -MODEL_NAME = "pytorch_model" -SAFE_MODEL_NAME = "model" -RNG_STATE_NAME = "random_states" -OPTIMIZER_NAME = "optimizer" -SCHEDULER_NAME = "scheduler" -SAMPLER_NAME = "sampler" -PROFILE_PATTERN_NAME = "profile_{suffix}.json" -WEIGHTS_NAME = f"{MODEL_NAME}.bin" -WEIGHTS_PATTERN_NAME = "pytorch_model{suffix}.bin" -WEIGHTS_INDEX_NAME = f"{WEIGHTS_NAME}.index.json" -SAFE_WEIGHTS_NAME = f"{SAFE_MODEL_NAME}.safetensors" -SAFE_WEIGHTS_PATTERN_NAME = "model{suffix}.safetensors" -SAFE_WEIGHTS_INDEX_NAME = f"{SAFE_WEIGHTS_NAME}.index.json" -SAGEMAKER_PYTORCH_VERSION = "1.10.2" -SAGEMAKER_PYTHON_VERSION = "py38" -SAGEMAKER_TRANSFORMERS_VERSION = "4.17.0" -SAGEMAKER_PARALLEL_EC2_INSTANCES = ["ml.p3.16xlarge", "ml.p3dn.24xlarge", "ml.p4dn.24xlarge"] -FSDP_SHARDING_STRATEGY = ["FULL_SHARD", "SHARD_GRAD_OP", "NO_SHARD", "HYBRID_SHARD", "HYBRID_SHARD_ZERO2"] -FSDP_AUTO_WRAP_POLICY = ["TRANSFORMER_BASED_WRAP", "SIZE_BASED_WRAP", "NO_WRAP"] -FSDP_BACKWARD_PREFETCH = ["BACKWARD_PRE", "BACKWARD_POST", "NO_PREFETCH"] -FSDP_STATE_DICT_TYPE = ["FULL_STATE_DICT", "LOCAL_STATE_DICT", "SHARDED_STATE_DICT"] -FSDP2_STATE_DICT_TYPE = ["SHARDED_STATE_DICT", "FULL_STATE_DICT"] -FSDP_PYTORCH_VERSION = ( - "2.1.0.a0+32f93b1" # Technically should be 2.1.0, but MS-AMP uses this specific prerelease in their Docker image. -) -FSDP2_PYTORCH_VERSION = "2.6.0" -FSDP_MODEL_NAME = "pytorch_model_fsdp" -DEEPSPEED_MULTINODE_LAUNCHERS = ["pdsh", "standard", "openmpi", "mvapich", "mpich", "nossh", "slurm"] -TORCH_DYNAMO_MODES = ["default", "reduce-overhead", "max-autotune"] -ELASTIC_LOG_LINE_PREFIX_TEMPLATE_PYTORCH_VERSION = "2.2.0" -XPU_PROFILING_AVAILABLE_PYTORCH_VERSION = "2.4.0" -MITA_PROFILING_AVAILABLE_PYTORCH_VERSION = "2.1.0" -BETA_TP_AVAILABLE_PYTORCH_VERSION = "2.3.0" -BETA_TP_AVAILABLE_TRANSFORMERS_VERSION = "4.52.0" - -STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt} - -# These are the args for `torch.distributed.launch` for pytorch < 1.9 -TORCH_LAUNCH_PARAMS = [ - "nnodes", - "nproc_per_node", - "rdzv_backend", - "rdzv_endpoint", - "rdzv_id", - "rdzv_conf", - "standalone", - "max_restarts", - "monitor_interval", - "start_method", - "role", - "module", - "m", - "no_python", - "run_path", - "log_dir", - "r", - "redirects", - "t", - "tee", - "node_rank", - "master_addr", - "master_port", -] - -CUDA_DISTRIBUTED_TYPES = ["DEEPSPEED", "MULTI_GPU", "FSDP", "MEGATRON_LM", "TP"] -TORCH_DISTRIBUTED_OPERATION_TYPES = CUDA_DISTRIBUTED_TYPES + [ - "MULTI_NPU", - "MULTI_MLU", - "MULTI_SDAA", - "MULTI_MUSA", - "MULTI_XPU", - "MULTI_CPU", - "MULTI_HPU", -] -SUPPORTED_PYTORCH_LAYERS_FOR_UPCASTING = ( - torch.nn.Conv1d, - torch.nn.Conv2d, - torch.nn.Conv3d, - torch.nn.ConvTranspose1d, - torch.nn.ConvTranspose2d, - torch.nn.ConvTranspose3d, - torch.nn.Linear, -) diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/dataclasses.py b/venv/lib/python3.10/site-packages/accelerate/utils/dataclasses.py deleted file mode 100644 index 91ca4b987fa680b6af36fb256b4e2814c9f78345..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/dataclasses.py +++ /dev/null @@ -1,2833 +0,0 @@ -# Copyright 2022 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. - -""" -General namespace and dataclass related classes -""" - -import argparse -import copy -import enum -import functools -import logging -import os -import warnings -from collections.abc import Iterable -from contextlib import contextmanager -from dataclasses import dataclass, field -from datetime import timedelta -from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union, get_args - -import torch - -from .constants import ( - BETA_TP_AVAILABLE_PYTORCH_VERSION, - FSDP2_PYTORCH_VERSION, - FSDP_AUTO_WRAP_POLICY, - FSDP_BACKWARD_PREFETCH, - FSDP_SHARDING_STRATEGY, - MITA_PROFILING_AVAILABLE_PYTORCH_VERSION, - XPU_PROFILING_AVAILABLE_PYTORCH_VERSION, -) -from .environment import parse_flag_from_env, str_to_bool -from .imports import ( - is_cuda_available, - is_hpu_available, - is_mlu_available, - is_msamp_available, - is_musa_available, - is_npu_available, - is_transformer_engine_available, - is_xpu_available, -) -from .versions import compare_versions, is_torch_version - - -if TYPE_CHECKING: - # Mock imports for type checking - from torchao.float8 import Float8LinearConfig - -logger = logging.getLogger(__name__) - - -class KwargsHandler: - """ - Internal mixin that implements a `to_kwargs()` method for a dataclass. - """ - - def to_dict(self): - return copy.deepcopy(self.__dict__) - - def to_kwargs(self): - """ - Returns a dictionary containing the attributes with values different from the default of this class. - """ - # import clear_environment here to avoid circular import problem - from .environment import clear_environment - - with clear_environment(): - default_dict = self.__class__().to_dict() - this_dict = self.to_dict() - return {k: v for k, v in this_dict.items() if default_dict[k] != v} - - -class EnumWithContains(enum.EnumMeta): - "A metaclass that adds the ability to check if `self` contains an item with the `in` operator" - - def __contains__(cls, item): - try: - cls(item) - except ValueError: - return False - return True - - -class BaseEnum(enum.Enum, metaclass=EnumWithContains): - "An enum class that can get the value of an item with `str(Enum.key)`" - - def __str__(self): - return self.value - - @classmethod - def list(cls): - "Method to list all the possible items in `cls`" - return list(map(str, cls)) - - -@dataclass -class AutocastKwargs(KwargsHandler): - """ - Use this object in your [`Accelerator`] to customize how `torch.autocast` behaves. Please refer to the - documentation of this [context manager](https://pytorch.org/docs/stable/amp.html#torch.autocast) for more - information on each argument. - - Example: - - ```python - from accelerate import Accelerator - from accelerate.utils import AutocastKwargs - - kwargs = AutocastKwargs(cache_enabled=True) - accelerator = Accelerator(kwargs_handlers=[kwargs]) - ``` - """ - - enabled: bool = True - cache_enabled: bool = None - - -class DDPCommunicationHookType(BaseEnum): - """ - Represents a type of communication hook used in DDP. - - Values: - - - **NO** -- no communication hook - - **FP16** -- DDP communication hook to compress the gradients in FP16 - - **BF16** -- DDP communication hook to compress the gradients in BF16 - - **POWER_SGD** -- DDP communication hook to use PowerSGD - - **BATCHED_POWER_SGD** -- DDP communication hook to use batched PowerSGD - """ - - NO = "no" - FP16 = "fp16" - BF16 = "bf16" - POWER_SGD = "power_sgd" - BATCHED_POWER_SGD = "batched_power_sgd" - - -@dataclass -class DistributedDataParallelKwargs(KwargsHandler): - """ - Use this object in your [`Accelerator`] to customize how your model is wrapped in a - `torch.nn.parallel.DistributedDataParallel`. Please refer to the documentation of this - [wrapper](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) for more - information on each argument. - - - - `gradient_as_bucket_view` is only available in PyTorch 1.7.0 and later versions. - - `static_graph` is only available in PyTorch 1.11.0 and later versions. - - - - Example: - - ```python - from accelerate import Accelerator - from accelerate.utils import DistributedDataParallelKwargs - - kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) - accelerator = Accelerator(kwargs_handlers=[kwargs]) - ``` - """ - - dim: int = 0 - broadcast_buffers: bool = True - bucket_cap_mb: int = 25 - find_unused_parameters: bool = False - check_reduction: bool = False - gradient_as_bucket_view: bool = False - static_graph: bool = False - - comm_hook: DDPCommunicationHookType = DDPCommunicationHookType.NO - comm_wrapper: Literal[ - DDPCommunicationHookType.NO, DDPCommunicationHookType.FP16, DDPCommunicationHookType.BF16 - ] = DDPCommunicationHookType.NO - comm_state_option: dict = field(default_factory=dict) - - def to_dict(self, ignore_keys=("comm_hook", "comm_wrapper", "comm_state_option")): - return {k: v for k, v in super().to_dict().items() if k not in ignore_keys} - - def register_comm_hook(self, model): - from torch.distributed.algorithms.ddp_comm_hooks import default_hooks, powerSGD_hook - - hook_map: dict[DDPCommunicationHookType, Callable] = { - DDPCommunicationHookType.FP16: default_hooks.fp16_compress_hook, - DDPCommunicationHookType.BF16: default_hooks.bf16_compress_hook, - DDPCommunicationHookType.POWER_SGD: powerSGD_hook.powerSGD_hook, - DDPCommunicationHookType.BATCHED_POWER_SGD: powerSGD_hook.batched_powerSGD_hook, - } - - wrapper_map: dict[DDPCommunicationHookType, Callable] = { - DDPCommunicationHookType.FP16: default_hooks.fp16_compress_wrapper, - DDPCommunicationHookType.BF16: default_hooks.bf16_compress_wrapper, - } - - hook: Optional[Callable] = hook_map.get(self.comm_hook) - wrapper: Optional[Callable] = wrapper_map.get(self.comm_wrapper) - - if hook and wrapper: - hook = wrapper(hook) - - if hook: - state = ( - powerSGD_hook.PowerSGDState(None, **self.comm_state_option) - if self.comm_hook in (DDPCommunicationHookType.POWER_SGD, DDPCommunicationHookType.BATCHED_POWER_SGD) - else None - ) - model.register_comm_hook( - state=state, - hook=hook, - ) - - -@dataclass -class GradScalerKwargs(KwargsHandler): - """ - Use this object in your [`Accelerator`] to customize the behavior of mixed precision, specifically how the - `torch.amp.GradScaler` or `torch.cuda.amp.GradScaler` used is created. Please refer to the documentation of this - [scaler](https://pytorch.org/docs/stable/amp.html?highlight=gradscaler) for more information on each argument. - - - - `torch.cuda.amp.GradScaler` is only available in PyTorch 1.5.0 and later versions, and `torch.amp.GradScaler` is - only available in PyTorch 2.4.0 and later versions. - - - - Example: - - ```python - from accelerate import Accelerator - from accelerate.utils import GradScalerKwargs - - kwargs = GradScalerKwargs(backoff_factor=0.25) - accelerator = Accelerator(kwargs_handlers=[kwargs]) - ``` - """ - - init_scale: float = 65536.0 - growth_factor: float = 2.0 - backoff_factor: float = 0.5 - growth_interval: int = 2000 - enabled: bool = True - - -@dataclass -class InitProcessGroupKwargs(KwargsHandler): - """ - Use this object in your [`Accelerator`] to customize the initialization of the distributed processes. Please refer - to the documentation of this - [method](https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) for more - information on each argument. - - Note: If `timeout` is set to `None`, the default will be based upon how `backend` is set. - - ```python - from datetime import timedelta - from accelerate import Accelerator - from accelerate.utils import InitProcessGroupKwargs - - kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=800)) - accelerator = Accelerator(kwargs_handlers=[kwargs]) - ``` - """ - - backend: Optional[str] = "nccl" - init_method: Optional[str] = None - timeout: Optional[timedelta] = None - - def __post_init__(self): - if self.timeout is None: - seconds = 1800 if self.backend != "nccl" else 600 - self.timeout = timedelta(seconds=seconds) - - -# Literals -Backend = Literal["MSAMP", "TE"] -OptLevel = Literal["O1", "O2"] -FP8Format = Literal["HYBRID", "E4M3", "E5M2"] -AmaxComputeAlgorithm = Literal["max", "most_recent"] - - -# FP8 training recipe kwargs -@dataclass -class AORecipeKwargs(KwargsHandler): - """ - Use this object in your [`Accelerator`] to customize the initialization of the recipe for FP8 mixed precision - training with `torchao` FP8. - - Args: - config (`torchao.float8.Float8LinearConfig`, *optional*, default to `None`): - The configuration for the FP8 training. In general, the default config should be sufficient. - module_filter_func (`Callable`, *optional*, default to `None`): - Optional function that must take in a module and layer name, and returns a boolean indicating whether the - module should be converted to FP8. Defaults to `accelerate.utils.ao.filter_linear_layers`. See it for an - example. - """ - - config: Optional["Float8LinearConfig"] = None - module_filter_func: Optional[Callable] = None - - -@dataclass -class TERecipeKwargs(KwargsHandler): - """ - Use this object in your [`Accelerator`] to customize the initialization of the recipe for FP8 mixed precision - training with `transformer-engine`. - - - - For more information on the args, please refer to the API - [documentation](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/common.html). - - - - ```python - from accelerate import Accelerator - from accelerate.utils import TERecipeKwargs - - kwargs = TERecipeKwargs(fp8_format="HYBRID") - accelerator = Accelerator(mixed_precision="fp8", kwargs_handlers=[kwargs]) - ``` - - Args: - use_autocast_during_eval (`bool`, *optional*, default to `False`): - Whether to use FP8 autocast during eval mode. Generally better metrics are found when this is `False`. - margin (`int`, *optional*, default to 0): - The margin to use for the gradient scaling. - interval (`int`, *optional*, default to 1): - The interval to use for how often the scaling factor is recomputed. - fp8_format (`str`, *optional*, default to "HYBRID"): - The format to use for the FP8 recipe. Must be one of `HYBRID`, `E4M3` or `E5M2`. (Generally `HYBRID` for - training, `E4M3` or `E5M2` for evaluation) - amax_history_len (`int`, *optional*, default to 1024): - The length of the history to use for the scaling factor computation - amax_compute_algo (`str`, *optional*, default to "most_recent"): - The algorithm to use for the scaling factor computation. Must be one of `max` or `most_recent`. - override_linear_precision (`tuple` of three `bool`, *optional*, default to `(False, False, False)`): - Whether or not to execute `fprop`, `dgrad`, and `wgrad` GEMMS in higher precision. - """ - - use_autocast_during_eval: bool = None - margin: int = None - interval: int = None - fp8_format: FP8Format = None - amax_history_len: int = None - amax_compute_algo: AmaxComputeAlgorithm = None - override_linear_precision: tuple[bool, bool, bool] = None - - def __post_init__(self): - env_prefix = "ACCELERATE_FP8_" - if not is_transformer_engine_available(): - raise ImportError("TransformerEngine is not available. Please install it or use a different backend.") - if self.use_autocast_during_eval is None: - self.use_autocast_during_eval = parse_flag_from_env(env_prefix + "USE_AUTOCAST_DURING_EVAL") - if self.margin is None: - self.margin = int(os.environ.get(env_prefix + "MARGIN", 0)) - if self.interval is None: - self.interval = int(os.environ.get(env_prefix + "INTERVAL", 1)) - if self.fp8_format is None: - self.fp8_format = os.environ.get(env_prefix + "FORMAT", "HYBRID") - self.fp8_format = self.fp8_format.upper() - if self.fp8_format not in get_args(FP8Format): - raise ValueError(f"`fp8_format` must be one of {' or '.join(get_args(FP8Format))}.") - if self.amax_compute_algo is None: - self.amax_compute_algo = os.environ.get(env_prefix + "AMAX_COMPUTE_ALGO", "most_recent") - self.amax_compute_algo = self.amax_compute_algo.lower() - if self.amax_compute_algo not in get_args(AmaxComputeAlgorithm): - raise ValueError(f"`amax_compute_algo` must be one of {' or '.join(get_args(AmaxComputeAlgorithm))}") - if self.amax_history_len is None: - self.amax_history_len = int(os.environ.get(env_prefix + "AMAX_HISTORY_LEN", 1024)) - if self.override_linear_precision is None: - fprop = parse_flag_from_env(env_prefix + "OVERRIDE_FPROP") - dgrad = parse_flag_from_env(env_prefix + "OVERRIDE_DGRAD") - wgrad = parse_flag_from_env(env_prefix + "OVERRIDE_WGRAD") - self.override_linear_precision = (fprop, dgrad, wgrad) - - -@dataclass -class MSAMPRecipeKwargs(KwargsHandler): - """ - Use this object in your [`Accelerator`] to customize the initialization of the recipe for FP8 mixed precision - training with `ms-amp`. - """ - - opt_level: OptLevel = None - - def __post_init__(self): - env_prefix = "ACCELERATE_FP8_" - if self.opt_level is None: - self.opt_level = os.environ.get(env_prefix + "OPT_LEVEL", "O2") - if self.opt_level not in get_args(OptLevel): - raise ValueError(f"`opt_level` must be one of {' or '.join(get_args(OptLevel))}") - - -@dataclass -class FP8RecipeKwargs(TERecipeKwargs, MSAMPRecipeKwargs): - """ - Deprecated. Please use one of the proper FP8 recipe kwargs classes such as `TERecipeKwargs` or `MSAMPRecipeKwargs` - instead. - """ - - backend: Backend = None - - def __post_init__(self): - env_prefix = "ACCELERATE_FP8_" - warnings.warn( - "FP8RecipeKwargs is deprecated and will be removed in Accelerate v2.0.0. " - "Please use one of the proper FP8 recipe kwargs classes such as TERecipeKwargs or MSAMPRecipeKwargs instead.", - FutureWarning, - ) - default_backend = "msamp" if is_msamp_available() else "te" - if self.backend is None: - self.backend = os.environ.get(env_prefix + "BACKEND", default_backend) - self.backend = self.backend.upper() - if self.backend not in get_args(Backend): - raise ValueError("`backend` must be 'MSAMP' or 'TE' (TransformerEngine) to use `FP8RecipeKwargs`.") - super().__post_init__() - - -# Literal -ProfilerActivity = Literal["cpu", "xpu", "mtia", "cuda", "hpu"] - - -@dataclass -class ProfileKwargs(KwargsHandler): - """ - Use this object in your [`Accelerator`] to customize the initialization of the profiler. Please refer to the - documentation of this [context manager](https://pytorch.org/docs/stable/profiler.html#torch.profiler.profile) for - more information on each argument. - - - - `torch.profiler` is only available in PyTorch 1.8.1 and later versions. - - - - Example: - - ```python - from accelerate import Accelerator - from accelerate.utils import ProfileKwargs - - kwargs = ProfileKwargs(activities=["cpu", "cuda"]) - accelerator = Accelerator(kwargs_handlers=[kwargs]) - ``` - - Args: - activities (`List[str]`, *optional*, default to `None`): - The list of activity groups to use in profiling. Must be one of `"cpu"`, `"xpu"`, `"mtia"`, "hpu" or - `"cuda"`. - schedule_option (`Dict[str, int]`, *optional*, default to `None`): - The schedule option to use for the profiler. Available keys are `wait`, `warmup`, `active`, `repeat` and - `skip_first`. The profiler will skip the first `skip_first` steps, then wait for `wait` steps, then do the - warmup for the next `warmup` steps, then do the active recording for the next `active` steps and then - repeat the cycle starting with `wait` steps. The optional number of cycles is specified with the `repeat` - parameter, the zero value means that the cycles will continue until the profiling is finished. - on_trace_ready (`Callable`, *optional*, default to `None`): - Callable that is called at each step when schedule returns `ProfilerAction.RECORD_AND_SAVE` during the - profiling. - record_shapes (`bool`, *optional*, default to `False`): - Save information about operator’s input shapes. - profile_memory (`bool`, *optional*, default to `False`): - Track tensor memory allocation/deallocation - with_stack (`bool`, *optional*, default to `False`): - Record source information (file and line number) for the ops. - with_flops (`bool`, *optional*, default to `False`): - Use formula to estimate the FLOPS of specific operators - with_modules (`bool`, *optional*, default to `False`): - Record module hierarchy (including function names) corresponding to the callstack of the op. - output_trace_dir (`str`, *optional*, default to `None`): - Exports the collected trace in Chrome JSON format. Chrome use 'chrome://tracing' view json file. Defaults - to None, which means profiling does not store json files. - """ - - activities: Optional[list[ProfilerActivity]] = None - schedule_option: Optional[dict[str, int]] = None - on_trace_ready: Optional[Callable] = None - record_shapes: bool = False - profile_memory: bool = False - with_stack: bool = False - with_flops: bool = False - with_modules: bool = False - output_trace_dir: Optional[str] = None - - def _get_profiler_activity(self, activity: ProfilerActivity) -> torch.profiler.ProfilerActivity: - """Get the profiler activity from the string. - - Args: - activity (str): The profiler activity name. - - Returns: - torch.profiler.ProfilerActivity: The profiler activity. - """ - - profiler_activity_map: dict[str, torch.profiler.ProfilerActivity] = { - "cpu": torch.profiler.ProfilerActivity.CPU, - "cuda": torch.profiler.ProfilerActivity.CUDA, - } - - if is_hpu_available(): - profiler_activity_map["hpu"] = torch.profiler.ProfilerActivity.HPU - - if is_torch_version(">=", XPU_PROFILING_AVAILABLE_PYTORCH_VERSION): - if torch.xpu.is_available(): - profiler_activity_map["xpu"] = torch.profiler.ProfilerActivity.XPU - - if is_torch_version(">=", MITA_PROFILING_AVAILABLE_PYTORCH_VERSION): - if torch.mtia.is_available(): - profiler_activity_map["mtia"] = torch.profiler.ProfilerActivity.MTIA - - if activity not in profiler_activity_map: - raise ValueError(f"Invalid profiler activity: {activity}. Must be one of {list(profiler_activity_map)}.") - return profiler_activity_map[activity] - - def build(self) -> torch.profiler.profile: - """ - Build a profiler object with the current configuration. - - Returns: - torch.profiler.profile: The profiler object. - """ - activities: Optional[list[ProfilerActivity]] = None - if self.activities is not None: - activities = [self._get_profiler_activity(activity) for activity in self.activities] - schedule: Optional[torch.profiler.schedule] = None - if self.schedule_option is not None: - schedule = torch.profiler.schedule(**self.schedule_option) - - return torch.profiler.profile( - activities=activities, - schedule=schedule, - on_trace_ready=self.on_trace_ready, - record_shapes=self.record_shapes, - profile_memory=self.profile_memory, - with_stack=self.with_stack, - with_flops=self.with_flops, - with_modules=self.with_modules, - ) - - -class DistributedType(str, enum.Enum): - """ - Represents a type of distributed environment. - - Values: - - - **NO** -- Not a distributed environment, just a single process. - - **MULTI_CPU** -- Distributed on multiple CPU nodes. - - **MULTI_GPU** -- Distributed on multiple GPUs. - - **MULTI_MLU** -- Distributed on multiple MLUs. - - **MULTI_SDAA** -- Distributed on multiple SDAAs. - - **MULTI_MUSA** -- Distributed on multiple MUSAs. - - **MULTI_NPU** -- Distributed on multiple NPUs. - - **MULTI_XPU** -- Distributed on multiple XPUs. - - **MULTI_HPU** -- Distributed on multiple HPUs. - - **DEEPSPEED** -- Using DeepSpeed. - - **XLA** -- Using TorchXLA. - """ - - # Subclassing str as well as Enum allows the `DistributedType` to be JSON-serializable out of the box. - NO = "NO" - MULTI_CPU = "MULTI_CPU" - MULTI_GPU = "MULTI_GPU" - MULTI_NPU = "MULTI_NPU" - MULTI_MLU = "MULTI_MLU" - MULTI_SDAA = "MULTI_SDAA" - MULTI_MUSA = "MULTI_MUSA" - MULTI_XPU = "MULTI_XPU" - DEEPSPEED = "DEEPSPEED" - FSDP = "FSDP" - TP = "TP" - XLA = "XLA" - MEGATRON_LM = "MEGATRON_LM" - MULTI_HPU = "MULTI_HPU" - - -class SageMakerDistributedType(str, enum.Enum): - """ - Represents a type of distributed environment. - - Values: - - - **NO** -- Not a distributed environment, just a single process. - - **DATA_PARALLEL** -- using sagemaker distributed data parallelism. - - **MODEL_PARALLEL** -- using sagemaker distributed model parallelism. - """ - - # Subclassing str as well as Enum allows the `SageMakerDistributedType` to be JSON-serializable out of the box. - NO = "NO" - DATA_PARALLEL = "DATA_PARALLEL" - MODEL_PARALLEL = "MODEL_PARALLEL" - - -class FP8BackendType(str, enum.Enum): - """ - Represents the backend used for FP8. - - Values: - - - **TE** -- using TransformerEngine. - - **MSAMP** -- using msamp. - """ - - # Subclassing str as well as Enum allows the `FP8BackendType` to be JSON-serializable out of the box. - NO = "NO" - TE = "TE" - MSAMP = "MSAMP" - AO = "AO" - - -class ComputeEnvironment(str, enum.Enum): - """ - Represents a type of the compute environment. - - Values: - - - **LOCAL_MACHINE** -- private/custom cluster hardware. - - **AMAZON_SAGEMAKER** -- Amazon SageMaker as compute environment. - """ - - # Subclassing str as well as Enum allows the `ComputeEnvironment` to be JSON-serializable out of the box. - LOCAL_MACHINE = "LOCAL_MACHINE" - AMAZON_SAGEMAKER = "AMAZON_SAGEMAKER" - - -class DynamoBackend(str, BaseEnum): - """ - Represents a dynamo backend (see https://pytorch.org/docs/stable/torch.compiler.html). - - Values: - - - **NO** -- Do not use torch dynamo. - - **EAGER** -- Uses PyTorch to run the extracted GraphModule. This is quite useful in debugging TorchDynamo - issues. - - **AOT_EAGER** -- Uses AotAutograd with no compiler, i.e, just using PyTorch eager for the AotAutograd's - extracted forward and backward graphs. This is useful for debugging, and unlikely to give speedups. - - **INDUCTOR** -- Uses TorchInductor backend with AotAutograd and cudagraphs by leveraging codegened Triton - kernels. [Read - more](https://dev-discuss.pytorch.org/t/torchinductor-a-pytorch-native-compiler-with-define-by-run-ir-and-symbolic-shapes/747) - - **AOT_TS_NVFUSER** -- nvFuser with AotAutograd/TorchScript. [Read - more](https://dev-discuss.pytorch.org/t/tracing-with-primitives-update-1-nvfuser-and-its-primitives/593) - - **NVPRIMS_NVFUSER** -- nvFuser with PrimTorch. [Read - more](https://dev-discuss.pytorch.org/t/tracing-with-primitives-update-1-nvfuser-and-its-primitives/593) - - **CUDAGRAPHS** -- cudagraphs with AotAutograd. [Read more](https://github.com/pytorch/torchdynamo/pull/757) - - **OFI** -- Uses Torchscript optimize_for_inference. Inference only. [Read - more](https://pytorch.org/docs/stable/generated/torch.jit.optimize_for_inference.html) - - **FX2TRT** -- Uses Nvidia TensorRT for inference optimizations. Inference only. [Read - more](https://github.com/pytorch/TensorRT/blob/master/docsrc/tutorials/getting_started_with_fx_path.rst) - - **ONNXRT** -- Uses ONNXRT for inference on CPU/GPU. Inference only. [Read more](https://onnxruntime.ai/) - - **TENSORRT** -- Uses ONNXRT to run TensorRT for inference optimizations. [Read - more](https://github.com/onnx/onnx-tensorrt) - - **AOT_TORCHXLA_TRACE_ONCE** -- Uses Pytorch/XLA with TorchDynamo optimization, for training. [Read - more](https://github.com/pytorch/xla/blob/r2.0/docs/dynamo.md) - - **TORCHXLA_TRACE_ONCE** -- Uses Pytorch/XLA with TorchDynamo optimization, for inference. [Read - more](https://github.com/pytorch/xla/blob/r2.0/docs/dynamo.md) - - **IPEX** -- Uses IPEX for inference on CPU. Inference only. [Read - more](https://github.com/intel/intel-extension-for-pytorch). - - **TVM** -- Uses Apach TVM for inference optimizations. [Read more](https://tvm.apache.org/) - - **HPU_BACKEND** -- Uses HPU backend for inference optimizations. - - """ - - # Subclassing str as well as Enum allows the `SageMakerDistributedType` to be JSON-serializable out of the box. - NO = "NO" - EAGER = "EAGER" - AOT_EAGER = "AOT_EAGER" - INDUCTOR = "INDUCTOR" - AOT_TS_NVFUSER = "AOT_TS_NVFUSER" - NVPRIMS_NVFUSER = "NVPRIMS_NVFUSER" - CUDAGRAPHS = "CUDAGRAPHS" - OFI = "OFI" - FX2TRT = "FX2TRT" - ONNXRT = "ONNXRT" - TENSORRT = "TENSORRT" - AOT_TORCHXLA_TRACE_ONCE = "AOT_TORCHXLA_TRACE_ONCE" - TORCHXLA_TRACE_ONCE = "TORCHXLA_TRACE_ONCE" - IPEX = "IPEX" - TVM = "TVM" - HPU_BACKEND = "HPU_BACKEND" - - -class LoggerType(BaseEnum): - """Represents a type of supported experiment tracker - - Values: - - - **ALL** -- all available trackers in the environment that are supported - - **TENSORBOARD** -- TensorBoard as an experiment tracker - - **WANDB** -- wandb as an experiment tracker - - **TRACKIO** -- trackio as an experiment tracker - - **COMETML** -- comet_ml as an experiment tracker - - **MLFLOW** -- mlflow as an experiment tracker - - **CLEARML** -- clearml as an experiment tracker - - **DVCLIVE** -- dvclive as an experiment tracker - - **SWANLAB** -- swanlab as an experiment tracker - """ - - ALL = "all" - AIM = "aim" - TENSORBOARD = "tensorboard" - WANDB = "wandb" - TRACKIO = "trackio" - COMETML = "comet_ml" - MLFLOW = "mlflow" - CLEARML = "clearml" - DVCLIVE = "dvclive" - SWANLAB = "swanlab" - - -class PrecisionType(str, BaseEnum): - """Represents a type of precision used on floating point values - - Values: - - - **NO** -- using full precision (FP32) - - **FP16** -- using half precision - - **BF16** -- using brain floating point precision - """ - - NO = "no" - FP8 = "fp8" - FP16 = "fp16" - BF16 = "bf16" - - -class RNGType(BaseEnum): - TORCH = "torch" - CUDA = "cuda" - MLU = "mlu" - SDAA = "sdaa" - MUSA = "musa" - NPU = "npu" - XLA = "xla" - XPU = "xpu" - HPU = "hpu" - GENERATOR = "generator" - - -class CustomDtype(enum.Enum): - r""" - An enum that contains multiple custom dtypes that can be used for `infer_auto_device_map`. - """ - - FP8 = "fp8" - INT4 = "int4" - INT2 = "int2" - - -# data classes - - -@dataclass -class TensorInformation: - shape: torch.Size - dtype: torch.dtype - - -@dataclass -class DataLoaderConfiguration: - """ - Configuration for dataloader-related items when calling `accelerator.prepare`. - - Args: - split_batches (`bool`, defaults to `False`): - Whether or not the accelerator should split the batches yielded by the dataloaders across the devices. If - `True`, the actual batch size used will be the same on any kind of distributed processes, but it must be a - round multiple of `num_processes` you are using. If `False`, actual batch size used will be the one set in - your script multiplied by the number of processes. - dispatch_batches (`bool`, defaults to `None`): - If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process - and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose - underlying dataset is an `IterableDataset`, `False` otherwise. - even_batches (`bool`, defaults to `True`): - If set to `True`, in cases where the total batch size across all processes does not exactly divide the - dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among - all workers. - use_seedable_sampler (`bool`, defaults to `False`): - Whether or not use a fully seedable random sampler ([`data_loader.SeedableRandomSampler`]). Ensures - training results are fully reproducable using a different sampling technique. While seed-to-seed results - may differ, on average the differences are neglible when using multiple different seeds to compare. Should - also be ran with [`~utils.set_seed`] for the best results. - data_seed (`int`, defaults to `None`): - The seed to use for the underlying generator when using `use_seedable_sampler`. If `None`, the generator - will use the current default seed from torch. - non_blocking (`bool`, defaults to `False`): - If set to `True`, the dataloader prepared by the Accelerator will utilize non-blocking host-to-device - transfers, allowing for better overlap between dataloader communication and computation. Recommended that - the prepared dataloader has `pin_memory` set to `True` to work properly. - use_stateful_dataloader (`bool`, defaults to `False`): - If set to `True`, the dataloader prepared by the Accelerator will be backed by - [torchdata.StatefulDataLoader](https://github.com/pytorch/data/tree/main/torchdata/stateful_dataloader). - This requires `torchdata` version 0.8.0 or higher that supports StatefulDataLoader to be installed. - """ - - split_batches: bool = field( - default=False, - metadata={ - "help": "Whether or not the accelerator should split the batches yielded by the dataloaders across the devices. If" - " `True` the actual batch size used will be the same on any kind of distributed processes, but it must be a" - " round multiple of the `num_processes` you are using. If `False`, actual batch size used will be the one set" - " in your script multiplied by the number of processes." - }, - ) - dispatch_batches: bool = field( - default=None, - metadata={ - "help": "If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process" - " and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose" - " underlying dataset is an `IterableDataset`, `False` otherwise." - }, - ) - even_batches: bool = field( - default=True, - metadata={ - "help": "If set to `True`, in cases where the total batch size across all processes does not exactly divide the" - " dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among" - " all workers." - }, - ) - use_seedable_sampler: bool = field( - default=False, - metadata={ - "help": "Whether or not use a fully seedable random sampler ([`data_loader.SeedableRandomSampler`])." - "Ensures training results are fully reproducable using a different sampling technique. " - "While seed-to-seed results may differ, on average the differences are neglible when using" - "multiple different seeds to compare. Should also be ran with [`~utils.set_seed`] for the best results." - }, - ) - data_seed: int = field( - default=None, - metadata={ - "help": "The seed to use for the underlying generator when using `use_seedable_sampler`. If `None`, the generator" - " will use the current default seed from torch." - }, - ) - non_blocking: bool = field( - default=False, - metadata={ - "help": "If set to `True`, the dataloader prepared by the Accelerator will utilize non-blocking host-to-device" - " transfers, allowing for better overlap between dataloader communication and computation. Recommended that the" - " prepared dataloader has `pin_memory` set to `True` to work properly." - }, - ) - use_stateful_dataloader: bool = field( - default=False, - metadata={ - "help": "If set to `True`, the dataloader prepared by the Accelerator will be backed by " - "[torchdata.StatefulDataLoader](https://github.com/pytorch/data/tree/main/torchdata/stateful_dataloader). This requires `torchdata` version 0.8.0 or higher that supports StatefulDataLoader to be installed." - }, - ) - - -@dataclass -class ProjectConfiguration: - """ - Configuration for the Accelerator object based on inner-project needs. - - Args: - project_dir (`str`, defaults to `None`): - A path to a directory for storing data. - logging_dir (`str`, defaults to `None`): - A path to a directory for storing logs of locally-compatible loggers. If None, defaults to `project_dir`. - automatic_checkpoint_naming (`bool`, defaults to `False`): - Whether saved states should be automatically iteratively named. - total_limit (`int`, defaults to `None`): - The maximum number of total saved states to keep. - iteration (`int`, defaults to `0`): - The current save iteration. - save_on_each_node (`bool`, defaults to `False`): - When doing multi-node distributed training, whether to save models and checkpoints on each node, or only on - the main one. - """ - - project_dir: str = field(default=None, metadata={"help": "A path to a directory for storing data."}) - logging_dir: str = field( - default=None, - metadata={ - "help": "A path to a directory for storing logs of locally-compatible loggers. If None, defaults to `project_dir`." - }, - ) - automatic_checkpoint_naming: bool = field( - default=False, - metadata={"help": "Whether saved states should be automatically iteratively named."}, - ) - - total_limit: int = field( - default=None, - metadata={"help": "The maximum number of total saved states to keep."}, - ) - - iteration: int = field( - default=0, - metadata={"help": "The current save iteration."}, - ) - - save_on_each_node: bool = field( - default=False, - metadata={ - "help": ( - "When doing multi-node distributed training, whether to save models and checkpoints on each node, or" - " only on the main one" - ) - }, - ) - - def set_directories(self, project_dir: str = None): - "Sets `self.project_dir` and `self.logging_dir` to the appropriate values." - self.project_dir = project_dir - if self.logging_dir is None: - self.logging_dir = project_dir - - def __post_init__(self): - self.set_directories(self.project_dir) - - -@dataclass -class GradientAccumulationPlugin(KwargsHandler): - """ - A plugin to configure gradient accumulation behavior. You can only pass one of `gradient_accumulation_plugin` or - `gradient_accumulation_steps` to [`Accelerator`]. Passing both raises an error. - - Parameters: - num_steps (`int`): - The number of steps to accumulate gradients for. - adjust_scheduler (`bool`, *optional*, defaults to `True`): - Whether to adjust the scheduler steps to account for the number of steps being accumulated. Should be - `True` if the used scheduler was not adjusted for gradient accumulation. - sync_with_dataloader (`bool`, *optional*, defaults to `True`): - Whether to synchronize setting the gradients when at the end of the dataloader. - sync_each_batch (`bool`, *optional*): - Whether to synchronize setting the gradients at each data batch. Seting to `True` may reduce memory - requirements when using gradient accumulation with distributed training, at expense of speed. - - Example: - - ```python - from accelerate.utils import GradientAccumulationPlugin - - gradient_accumulation_plugin = GradientAccumulationPlugin(num_steps=2) - accelerator = Accelerator(gradient_accumulation_plugin=gradient_accumulation_plugin) - ``` - """ - - num_steps: int = field(default=None, metadata={"help": "The number of steps to accumulate gradients for."}) - adjust_scheduler: bool = field( - default=True, - metadata={ - "help": "Whether to adjust the scheduler steps to account for the number of steps being accumulated. Should be `True` if the used scheduler was not adjusted for gradient accumulation." - }, - ) - sync_with_dataloader: bool = field( - default=True, - metadata={ - "help": "Whether to synchronize setting the gradients when at the end of the dataloader. Should only be set to `False` if you know what you're doing." - }, - ) - sync_each_batch: bool = field( - default=False, - metadata={ - "help": "Whether to synchronize setting the gradients at each data batch. Setting to `True` may reduce memory requirements when using gradient accumulation with distributed training, at expense of speed." - }, - ) - - -@dataclass -class TorchDynamoPlugin(KwargsHandler): - """ - This plugin is used to compile a model with PyTorch 2.0 - - Args: - backend (`DynamoBackend`, defaults to `None`): - A valid Dynamo backend. See https://pytorch.org/docs/stable/torch.compiler.html for more details. - mode (`str`, defaults to `None`): - Possible options are 'default', 'reduce-overhead' or 'max-autotune'. - fullgraph (`bool`, defaults to `None`): - Whether it is ok to break model into several subgraphs. - dynamic (`bool`, defaults to `None`): - Whether to use dynamic shape for tracing. - options (`Any`, defaults to `None`): - A dictionary of options to pass to the backend. - disable (`bool`, defaults to `False`): - Turn torch.compile() into a no-op for testing - use_regional_compilation (`bool`, defaults to `None`): - Use it to reduce the cold start compilation time of torch.compile() by targeting repeated blocks of the - same class and compiling them sequentially to hit the compiler's cache. For example, in `GPT2LMHeadModel`, - the repeated block/class is `GPT2Block`, and can be accessed as `model.transformer.h[0]`. The rest of the - model (e.g model.lm_head) is compiled separately. - """ - - backend: DynamoBackend = field( - default=None, - metadata={"help": f"Possible options are {[b.value.lower() for b in DynamoBackend]}"}, - ) - mode: str = field( - default=None, metadata={"help": "Possible options are 'default', 'reduce-overhead' or 'max-autotune'"} - ) - fullgraph: bool = field(default=None, metadata={"help": "Whether it is ok to break model into several subgraphs"}) - dynamic: bool = field(default=None, metadata={"help": "Whether to use dynamic shape for tracing"}) - options: Any = field(default=None, metadata={"help": "A dictionary of options to pass to the backend."}) - disable: bool = field(default=False, metadata={"help": "Turn torch.compile() into a no-op for testing"}) - - use_regional_compilation: bool = field( - default=None, - metadata={ - "help": ( - # https://pytorch.org/tutorials/recipes/regional_compilation.html - "Use it to reduce the cold start compilation time of torch.compile() by targeting repeated " - "blocks of the same class and compiling them sequentially to hit the compiler's cache. For " - "example, in `GPT2LMHeadModel`, the repeated block/class is `GPT2Block`, and can be accessed " - "as `model.transformer.h[0]`. The rest of the model (e.g model.lm_head) is compiled separately." - ) - }, - ) - - def __post_init__(self): - prefix = "ACCELERATE_DYNAMO_" - if self.backend is None: - self.backend = os.environ.get(prefix + "BACKEND", "no") - self.backend = DynamoBackend(self.backend.upper()) - - if self.mode is None: - self.mode = os.environ.get(prefix + "MODE", "default") - if self.fullgraph is None: - self.fullgraph = str_to_bool(os.environ.get(prefix + "USE_FULLGRAPH", "False")) == 1 - if self.use_regional_compilation is None: - self.use_regional_compilation = ( - str_to_bool(os.environ.get(prefix + "USE_REGIONAL_COMPILATION", "False")) == 1 - ) - - if self.dynamic is None and os.environ.get(prefix + "USE_DYNAMIC", None) is not None: - self.dynamic = str_to_bool(os.environ.get(prefix + "USE_DYNAMIC", "False")) == 1 - - def to_dict(self): - dynamo_config = copy.deepcopy(self.__dict__) - dynamo_config["backend"] = dynamo_config["backend"].value.lower() - return dynamo_config - - def to_kwargs(self): - kwargs = super().to_kwargs() - kwargs.pop("use_regional_compilation", None) - return kwargs - - -@dataclass -class DeepSpeedPlugin: - """ - This plugin is used to integrate DeepSpeed. - - Args: - hf_ds_config (`Any`, defaults to `None`): - Path to DeepSpeed config file or dict or an object of class `accelerate.utils.deepspeed.HfDeepSpeedConfig`. - gradient_accumulation_steps (`int`, defaults to `None`): - Number of steps to accumulate gradients before updating optimizer states. If not set, will use the value - from the `Accelerator` directly. - gradient_clipping (`float`, defaults to `None`): - Enable gradient clipping with value. - zero_stage (`int`, defaults to `None`): - Possible options are 0, 1, 2, 3. Default will be taken from environment variable. - is_train_batch_min (`bool`, defaults to `True`): - If both train & eval dataloaders are specified, this will decide the `train_batch_size`. - offload_optimizer_device (`str`, defaults to `None`): - Possible options are none|cpu|nvme. Only applicable with ZeRO Stages 2 and 3. - offload_param_device (`str`, defaults to `None`): - Possible options are none|cpu|nvme. Only applicable with ZeRO Stage 3. - offload_optimizer_nvme_path (`str`, defaults to `None`): - Possible options are /nvme|/local_nvme. Only applicable with ZeRO Stage 3. - offload_param_nvme_path (`str`, defaults to `None`): - Possible options are /nvme|/local_nvme. Only applicable with ZeRO Stage 3. - zero3_init_flag (`bool`, defaults to `None`): - Flag to indicate whether to save 16-bit model. Only applicable with ZeRO Stage-3. - zero3_save_16bit_model (`bool`, defaults to `None`): - Flag to indicate whether to save 16-bit model. Only applicable with ZeRO Stage-3. - transformer_moe_cls_names (`str`, defaults to `None`): - Comma-separated list of Transformers MoE layer class names (case-sensitive). For example, - `MixtralSparseMoeBlock`, `Qwen2MoeSparseMoeBlock`, `JetMoEAttention`, `JetMoEBlock`, etc. - enable_msamp (`bool`, defaults to `None`): - Flag to indicate whether to enable MS-AMP backend for FP8 training. - msasmp_opt_level (`Optional[Literal["O1", "O2"]]`, defaults to `None`): - Optimization level for MS-AMP (defaults to 'O1'). Only applicable if `enable_msamp` is True. Should be one - of ['O1' or 'O2']. - """ - - hf_ds_config: Any = field( - default=None, - metadata={ - "help": "path to DeepSpeed config file or dict or an object of class `accelerate.utils.deepspeed.HfDeepSpeedConfig`." - }, - ) - gradient_accumulation_steps: int = field( - default=None, - metadata={ - "help": "Number of steps to accumulate gradients before updating optimizer states. If not set, will use the value from the `Accelerator` directly." - }, - ) - gradient_clipping: float = field(default=None, metadata={"help": "Enable gradient clipping with value"}) - zero_stage: int = field( - default=None, - metadata={"help": "Possible options are 0,1,2,3; Default will be taken from environment variable"}, - ) - is_train_batch_min: bool = field( - default=True, - metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"}, - ) - offload_optimizer_device: str = field( - default=None, - metadata={"help": "Possible options are none|cpu|nvme. Only applicable with ZeRO Stages 2 and 3."}, - ) - offload_param_device: str = field( - default=None, - metadata={"help": "Possible options are none|cpu|nvme. Only applicable with ZeRO Stage 3."}, - ) - offload_optimizer_nvme_path: str = field( - default=None, - metadata={"help": "Possible options are /nvme|/local_nvme. Only applicable with ZeRO Stage 3."}, - ) - offload_param_nvme_path: str = field( - default=None, - metadata={"help": "Possible options are /nvme|/local_nvme. Only applicable with ZeRO Stage 3."}, - ) - zero3_init_flag: bool = field( - default=None, - metadata={ - "help": "Flag to indicate whether to enable `deepspeed.zero.Init` for constructing massive models." - "Only applicable with ZeRO Stage-3." - }, - ) - zero3_save_16bit_model: bool = field( - default=None, - metadata={"help": "Flag to indicate whether to save 16-bit model. Only applicable with ZeRO Stage-3."}, - ) - transformer_moe_cls_names: str = field( - default=None, - metadata={ - "help": "comma-separated list of transformers MoE layer class names (case-sensitive), e.g : " - " `MixtralSparseMoeBlock`, `Qwen2MoeSparseMoeBlock`, `JetMoEAttention,JetMoEBlock` ..." - }, - ) - enable_msamp: bool = field( - default=None, - metadata={"help": "Flag to indicate whether to enable MS-AMP backend for FP8 training."}, - ) - msamp_opt_level: Optional[Literal["O1", "O2"]] = field( - default=None, - metadata={ - "help": "Optimization level for MS-AMP (defaults to 'O1'). Only applicable if `enable_msamp` is True. Should be one of ['O1' or 'O2']." - }, - ) - - def __post_init__(self): - from .deepspeed import HfDeepSpeedConfig - - if self.gradient_accumulation_steps is None: - gas = os.environ.get("ACCELERATE_GRADIENT_ACCUMULATION_STEPS", "auto") - self.gradient_accumulation_steps = int(gas) if gas.isdigit() else gas - - if self.gradient_clipping is None: - gradient_clipping = os.environ.get("ACCELERATE_GRADIENT_CLIPPING", "auto") - self.gradient_clipping = gradient_clipping if gradient_clipping == "auto" else float(gradient_clipping) - - if self.zero_stage is None: - self.zero_stage = int(os.environ.get("ACCELERATE_DEEPSPEED_ZERO_STAGE", 2)) - - if self.offload_optimizer_device is None: - self.offload_optimizer_device = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none") - - if self.offload_param_device is None: - self.offload_param_device = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE", "none") - - if self.offload_optimizer_nvme_path is None: - self.offload_optimizer_nvme_path = os.environ.get( - "ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_NVME_PATH", "none" - ) - - if self.offload_param_nvme_path is None: - self.offload_param_nvme_path = os.environ.get("ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_NVME_PATH", "none") - - if self.zero3_save_16bit_model is None: - self.zero3_save_16bit_model = ( - os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", "false") == "true" - ) - if self.enable_msamp is None: - self.enable_msamp = os.environ.get("ACCELERATE_FP8_BACKEND", None) == "MSAMP" - - if self.msamp_opt_level is None: - self.msamp_opt_level = os.environ.get("ACCELERATE_FP8_OPT_LEVEL", "O1") - - if self.hf_ds_config is None: - self.hf_ds_config = os.environ.get("ACCELERATE_DEEPSPEED_CONFIG_FILE", "none") - if ( - isinstance(self.hf_ds_config, dict) - or (isinstance(self.hf_ds_config, str) and self.hf_ds_config != "none") - or isinstance(self.hf_ds_config, HfDeepSpeedConfig) - ): - if not isinstance(self.hf_ds_config, HfDeepSpeedConfig): - self.hf_ds_config = HfDeepSpeedConfig(self.hf_ds_config) - if "gradient_accumulation_steps" not in self.hf_ds_config.config: - self.hf_ds_config.config["gradient_accumulation_steps"] = 1 - if "zero_optimization" not in self.hf_ds_config.config: - raise ValueError("Please specify the ZeRO optimization config in the DeepSpeed config.") - - self._deepspeed_config_checks() - plugin_to_config_mapping = { - "gradient_accumulation_steps": "gradient_accumulation_steps", - "gradient_clipping": "gradient_clipping", - "zero_stage": "zero_optimization.stage", - "offload_optimizer_device": "zero_optimization.offload_optimizer.device", - "offload_param_device": "zero_optimization.offload_param.device", - "offload_param_nvme_path": "zero_optimization.offload_param.nvme_path", - "offload_optimizer_nvme_path": "zero_optimization.offload_optimizer.nvme_path", - "zero3_save_16bit_model": "zero_optimization.stage3_gather_16bit_weights_on_model_save", - } - kwargs = {v: getattr(self, k) for k, v in plugin_to_config_mapping.items() if getattr(self, k) is not None} - for key in kwargs.keys(): - self.fill_match(key, **kwargs, must_match=False) - self.hf_ds_config.set_stage_and_offload() - - # filling the missing values in the class attributes from the DeepSpeed config - # when using the DeepSpeed config file. - for key, value in plugin_to_config_mapping.items(): - config_value = self.hf_ds_config.get_value(value) - if config_value is not None and config_value != "auto": - setattr(self, key, config_value) - else: - config = { - "train_batch_size": "auto", - "train_micro_batch_size_per_gpu": "auto", - "gradient_accumulation_steps": self.gradient_accumulation_steps, - "zero_optimization": { - "stage": self.zero_stage, - "offload_optimizer": { - "device": self.offload_optimizer_device, - "nvme_path": self.offload_optimizer_nvme_path - if self.offload_optimizer_device == "nvme" - else None, - }, - "offload_param": { - "device": self.offload_param_device, - "nvme_path": self.offload_param_nvme_path if self.offload_param_device == "nvme" else None, - }, - "stage3_gather_16bit_weights_on_model_save": self.zero3_save_16bit_model, - }, - } - if self.gradient_clipping: - config["gradient_clipping"] = self.gradient_clipping - self.hf_ds_config = HfDeepSpeedConfig(config) - - self.deepspeed_config = self.hf_ds_config.config - self.deepspeed_config["steps_per_print"] = float("inf") # this will stop deepspeed from logging @ stdout - if self.zero3_init_flag is None: - self.zero3_init_flag = ( - str_to_bool(os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", str(self.hf_ds_config.is_zero3()))) == 1 - ) - if self.zero3_init_flag and not self.hf_ds_config.is_zero3(): - warnings.warn("DeepSpeed Zero3 Init flag is only applicable for ZeRO Stage 3. Setting it to False.") - self.zero3_init_flag = False - # NOTE: Set to False by default, will be set to `True` automatically if it's the first plugin passed - # to the `Accelerator`'s `deepspeed_plugin` param, *or* `AcceleratorState().enable_deepspeed_plugin(plugin_key)` is manually called - self._set_selected(False) - - # Ignore if it's already set - if self.enable_msamp and "msamp" not in self.deepspeed_config: - if self.zero_stage == 3: - raise NotImplementedError( - "MS-AMP is not supported for ZeRO Stage 3. Please use ZeRO Stage 0, 1, or 2 instead." - ) - if self.msamp_opt_level not in ["O1", "O2"]: - raise ValueError("Invalid optimization level for MS-AMP. Please use one of ['O1' or'O2'].") - self.deepspeed_config["msamp"] = {"enabled": True, "opt_level": self.msamp_opt_level} - - def fill_match(self, ds_key_long, mismatches=None, must_match=True, **kwargs): - mismatches = [] if mismatches is None else mismatches - config, ds_key = self.hf_ds_config.find_config_node(ds_key_long) - if config is None: - return - - if config.get(ds_key) == "auto": - if ds_key_long in kwargs: - config[ds_key] = kwargs[ds_key_long] - return - else: - raise ValueError( - f"`{ds_key_long}` not found in kwargs. " - f"Please specify `{ds_key_long}` without `auto` (set to correct value) in the DeepSpeed config file or " - "pass it in kwargs." - ) - - if not must_match: - return - - ds_val = config.get(ds_key) - if ds_val is not None and ds_key_long in kwargs: - if ds_val != kwargs[ds_key_long]: - mismatches.append(f"- ds {ds_key_long}={ds_val} vs arg {ds_key_long}={kwargs[ds_key_long]}") - - def is_auto(self, ds_key_long): - val = self.hf_ds_config.get_value(ds_key_long) - if val is None: - return False - else: - return val == "auto" - - def get_value(self, ds_key_long, default=None): - return self.hf_ds_config.get_value(ds_key_long, default) - - def deepspeed_config_process(self, prefix="", mismatches=None, config=None, must_match=True, **kwargs): - """Process the DeepSpeed config with the values from the kwargs.""" - mismatches = [] if mismatches is None else mismatches - if config is None: - config = self.deepspeed_config - for key, value in config.items(): - if isinstance(value, dict): - self.deepspeed_config_process( - prefix=prefix + key + ".", mismatches=mismatches, config=value, must_match=must_match, **kwargs - ) - else: - self.fill_match(prefix + key, mismatches, must_match=must_match, **kwargs) - if len(mismatches) > 0 and prefix == "": - mismatches_msg = "\n".join(mismatches) - raise ValueError( - "Please correct the following DeepSpeed config values that mismatch kwargs " - f" values:\n{mismatches_msg}\nThe easiest method is to set these DeepSpeed config values to 'auto'." - ) - - def set_mixed_precision(self, mixed_precision): - ds_config = self.deepspeed_config - kwargs = { - "fp16.enabled": mixed_precision == "fp16", - # When training in fp8, we still rely on bf16 autocast for the core mixed precision - "bf16.enabled": mixed_precision in ("bf16", "fp8"), - } - if mixed_precision == "fp16": - if "fp16" not in ds_config: - ds_config["fp16"] = {"enabled": True, "auto_cast": True} - elif mixed_precision in ("bf16", "fp8"): - if "bf16" not in ds_config: - ds_config["bf16"] = {"enabled": True} - - if mixed_precision == "fp8" and self.enable_msamp: - if "msamp" not in ds_config: - ds_config["msamp"] = {"enabled": True, "opt_level": self.msamp_opt_level} - - if mixed_precision != "no": - diff_dtype = "bf16" if mixed_precision == "fp16" else "fp16" - if str(ds_config.get(diff_dtype, {}).get("enabled", "False")).lower() == "true": - raise ValueError( - f"`--mixed_precision` arg cannot be set to `{mixed_precision}` when `{diff_dtype}` is set in the DeepSpeed config file." - ) - for dtype in ["fp16", "bf16"]: - if dtype not in ds_config: - ds_config[dtype] = {"enabled": False} - self.fill_match("fp16.enabled", must_match=False, **kwargs) - self.fill_match("bf16.enabled", must_match=False, **kwargs) - - def set_deepspeed_weakref(self): - from .imports import is_transformers_available - - ds_config = copy.deepcopy(self.deepspeed_config) - if self.zero3_init_flag: - if not is_transformers_available(): - raise Exception( - "When `zero3_init_flag` is set, it requires Transformers to be installed. " - "Please run `pip install transformers`." - ) - if "gradient_accumulation_steps" not in ds_config or ds_config["gradient_accumulation_steps"] == "auto": - ds_config["gradient_accumulation_steps"] = 1 - if "train_micro_batch_size_per_gpu" not in ds_config or ds_config["train_micro_batch_size_per_gpu"] == "auto": - ds_config["train_micro_batch_size_per_gpu"] = 1 - if ds_config.get("train_batch_size", None) == "auto": - del ds_config["train_batch_size"] - - if compare_versions("transformers", "<", "4.46"): - from transformers.deepspeed import HfDeepSpeedConfig, unset_hf_deepspeed_config - else: - from transformers.integrations import HfDeepSpeedConfig, unset_hf_deepspeed_config - - unset_hf_deepspeed_config() - self.dschf = HfDeepSpeedConfig(ds_config) # keep this object alive # noqa - - def is_zero3_init_enabled(self): - return self.zero3_init_flag - - @contextmanager - def zero3_init_context_manager(self, enable=False): - old = self.zero3_init_flag - if old == enable: - yield - else: - self.zero3_init_flag = enable - self.dschf = None - self.set_deepspeed_weakref() - yield - self.zero3_init_flag = old - self.dschf = None - self.set_deepspeed_weakref() - - def _deepspeed_config_checks(self): - env_variable_names_to_ignore = [ - "ACCELERATE_GRADIENT_ACCUMULATION_STEPS", - "ACCELERATE_GRADIENT_CLIPPING", - "ACCELERATE_DEEPSPEED_ZERO_STAGE", - "ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", - "ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE", - "ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_NVME_PATH", - "ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_NVME_PATH", - "ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", - "ACCELERATE_MIXED_PRECISION", - ] - env_variable_names_to_ignore = [ - name.replace("ACCELERATE_", "").replace("DEEPSPEED_", "").lower() for name in env_variable_names_to_ignore - ] - - deepspeed_fields_from_accelerate_config = os.environ.get("ACCELERATE_CONFIG_DS_FIELDS", "").split(",") - - if any(name in env_variable_names_to_ignore for name in deepspeed_fields_from_accelerate_config): - raise ValueError( - f"When using `deepspeed_config_file`, the following accelerate config variables will be ignored: {env_variable_names_to_ignore}.\n" - "Please specify them appropriately in the DeepSpeed config file.\n" - "If you are using an accelerate config file, remove others config variables mentioned in the above specified list.\n" - "The easiest method is to create a new config following the questionnaire via `accelerate config`.\n" - "It will only ask for the necessary config variables when using `deepspeed_config_file`." - ) - - def set_moe_leaf_modules(self, model): - if self.transformer_moe_cls_names is None: - self.transformer_moe_cls_names = os.environ.get("ACCELERATE_DEEPSPEED_MOE_LAYER_CLS_NAMES", None) - if self.transformer_moe_cls_names is not None: - if compare_versions("deepspeed", "<", "0.14.0"): - raise ImportError("DeepSpeed version must be >= 0.14.0 to use MOE support. Please update DeepSpeed.") - from deepspeed.utils import set_z3_leaf_modules - - class_names = self.transformer_moe_cls_names.split(",") - transformer_moe_cls = [] - for layer_class in class_names: - transformer_cls = get_module_class_from_name(model, layer_class) - if transformer_cls is None: - raise Exception( - f"Could not find a transformer layer class called '{layer_class}' to wrap in the model." - ) - else: - transformer_moe_cls.append(transformer_cls) - set_z3_leaf_modules(model, transformer_moe_cls) # z3_leaf - - def select(self, _from_accelerator_state: bool = False): - """ - Sets the HfDeepSpeedWeakref to use the current deepspeed plugin configuration - """ - if not _from_accelerator_state: - raise ValueError( - "A `DeepSpeedPlugin` object must be enabled manually by calling `AcceleratorState().enable_deepspeed_plugin(plugin_key)`." - ) - self.set_deepspeed_weakref() - self._set_selected(True) - - def _unselect(self): - self._set_selected(False) - - def _set_selected(self, value: bool): - """ - Private setter for the 'enabled' attribute. - """ - self._selected = value - - @property - def selected(self): - return self._selected - - @selected.setter - def selected(self, value): - raise NotImplementedError( - "'enabled' can only be set through calling 'AcceleratorState().enable_deepspeed_plugin(key)'." - ) - - -@dataclass -class FullyShardedDataParallelPlugin: - """ - This plugin is used to enable fully sharded data parallelism. - - Args: - fsdp_version (`int`, defaults to `1`): - The version of FSDP to use. Defaults to 1. If set to 2, launcher expects the config to be converted to - FSDP2 format. - sharding_strategy (`Union[str, torch.distributed.fsdp.ShardingStrategy]`, defaults to `'FULL_SHARD'`): - Sharding strategy to use. Should be either a `str` or an instance of - `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy`. Is deprecated in favor of - `reshard_after_forward`. - reshard_after_forward (`Union[str, torch.distributed.fsdp.ShardingStrategy, bool]`, defaults to `'FULL_SHARD'` for `fsdp_version=1` and `True` for `fsdp_version=2`): - Sharding strategy to use. Should be a bool if `fsdp_version` is set to 2 else a `str` or an instance of - `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy`. - backward_prefetch (`Union[str, torch.distributed.fsdp.BackwardPrefetch]`, defaults to `'NO_PREFETCH'`): - Backward prefetch strategy to use. Should be either a `str` or an instance of - `torch.distributed.fsdp.fully_sharded_data_parallel.BackwardPrefetch`. - mixed_precision_policy (`Optional[Union[dict, torch.distributed.fsdp.MixedPrecision, torch.distributed.fsdp.MixedPrecisionPolicy]]`, defaults to `None`): - A config to enable mixed precision training with FullyShardedDataParallel. If passing in a `dict`, it - should have the following keys: `param_dtype`, `reduce_dtype`, and `buffer_dtype`, can be an instance of - `torch.distributed.fsdp.MixedPrecisionPolicy` if `fsdp_version` is set to 2. - auto_wrap_policy (`Optional(Union[Callable, Literal["transformer_based_wrap", "size_based_wrap", "no_wrap"]]), defaults to `NO_WRAP`): - A callable or string specifying a policy to recursively wrap layers with FSDP. If a string, it must be one - of `transformer_based_wrap`, `size_based_wrap`, or `no_wrap`. See - `torch.distributed.fsdp.wrap.size_based_wrap_policy` for a direction on what it should look like. - cpu_offload (`Union[bool, torch.distributed.fsdp.CPUOffload, torch.distributed.fsdp.CPUOffloadPolicy]`, defaults to `False`): - Whether to offload parameters to CPU. Should be either a `bool` or an instance of - `torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffload` or - `torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffloadPolicy` if `fsdp_version` is set to 2. - ignored_modules (`Optional[Iterable[torch.nn.Module]]`, defaults to `None`): - A list of modules to ignore when wrapping with FSDP. - state_dict_type (`Union[str, torch.distributed.fsdp.StateDictType]`, defaults to `'FULL_STATE_DICT'`): - State dict type to use. If a string, it must be one of `full_state_dict`, `local_state_dict`, or - `sharded_state_dict`. - state_dict_config (`Optional[Union[torch.distributed.fsdp.FullStateDictConfig, torch.distributed.fsdp.ShardedStateDictConfig]`, defaults to `None`): - State dict config to use. Is determined based on the `state_dict_type` if not passed in. - optim_state_dict_config (`Optional[Union[torch.distributed.fsdp.FullOptimStateDictConfig, torch.distributed.fsdp.ShardedOptimStateDictConfig]`, defaults to `None`): - Optim state dict config to use. Is determined based on the `state_dict_type` if not passed in. - limit_all_gathers (`bool`, defaults to `True`): - Whether to have FSDP explicitly synchronizes the CPU thread to prevent too many in-flight all-gathers. This - bool only affects the sharded strategies that schedule all-gathers. Enabling this can help lower the number - of CUDA malloc retries. - use_orig_params (`bool`, defaults to `False`): - Whether to use the original parameters for the optimizer. - param_init_fn (`Optional[Callable[[torch.nn.Module], None]`, defaults to `None`): - A `Callable[torch.nn.Module] -> None` that specifies how modules that are currently on the meta device - should be initialized onto an actual device. Only applicable when `sync_module_states` is `True`. By - default is a `lambda` which calls `to_empty` on the module. - sync_module_states (`bool`, defaults to `False`): - Whether each individually wrapped FSDP unit should broadcast module parameters from rank 0 to ensure they - are the same across all ranks after initialization. Defaults to `False` unless `cpu_ram_efficient_loading` - is `True`, then will be forcibly enabled. - forward_prefetch (`bool`, defaults to `False`): - Whether to have FSDP explicitly prefetches the next upcoming all-gather while executing in the forward - pass. only use with Static graphs. - activation_checkpointing (`bool`, defaults to `False`): - A technique to reduce memory usage by clearing activations of certain layers and recomputing them during a - backward pass. Effectively, this trades extra computation time for reduced memory usage. - cpu_ram_efficient_loading (`bool`, defaults to `None`): - If True, only the first process loads the pretrained model checkoint while all other processes have empty - weights. Only applicable for Transformers. When using this, `sync_module_states` needs to be `True`. - transformer_cls_names_to_wrap (`Optional[List[str]]`, defaults to `None`): - A list of transformer layer class names to wrap. Only applicable when `auto_wrap_policy` is - `transformer_based_wrap`. - min_num_params (`Optional[int]`, defaults to `None`): - The minimum number of parameters a module must have to be wrapped. Only applicable when `auto_wrap_policy` - is `size_based_wrap`. - """ - - fsdp_version: int = field( - default=None, - metadata={ - "help": "The version of FSDP to use. Defaults to 1. If set to 2, launcher expects the config to be converted to FSDP2 format." - }, - ) - - sharding_strategy: Union[str, "torch.distributed.fsdp.ShardingStrategy"] = field( - default=None, - metadata={ - "help": "Sharding strategy to use. Should be either a `str` or an instance of `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy`. Defaults to 'FULL_SHARD'. Is deprecated in favor of `reshard_after_forward` " - }, - ) - - reshard_after_forward: Union[str, "torch.distributed.fsdp.ShardingStrategy", bool] = field( - default=None, - metadata={ - "help": "Sharding strategy to use. Should be a bool if `fsdp_version` is set to 2 else a `str` or an instance of `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy`. Defaults to 'FULL_SHARD'" - }, - ) - backward_prefetch: Optional[Union[str, "torch.distributed.fsdp.BackwardPrefetch"]] = field( - default=None, - metadata={ - "help": "Backward prefetch strategy to use. Should be either a `str` or an instance of `torch.distributed.fsdp.fully_sharded_data_parallel.BackwardPrefetch`. Defaults to 'NO_PREFETCH'. This becomes obsolete in FSDP2." - }, - ) - mixed_precision_policy: Optional[ - Union[dict, "torch.distributed.fsdp.MixedPrecision", "torch.distributed.fsdp.MixedPrecisionPolicy"] - ] = field( - default=None, - metadata={ - "help": "A config to enable mixed precision training with FullyShardedDataParallel. " - "If passing in a `dict`, it should have the following keys: `param_dtype`, `reduce_dtype`, and `buffer_dtype`." - "Can also be an instance of `torch.distributed.fsdp.MixedPrecisionPolicy` if `fsdp_version` is set to 2." - }, - ) - auto_wrap_policy: Optional[Union[Callable, Literal["transformer_based_wrap", "size_based_wrap", "no_wrap"]]] = ( - field( - default=None, - metadata={ - "help": "A callable or string specifying a policy to recursively wrap layers with FSDP. If a string, it must be one of `transformer_based_wrap`, `size_based_wrap`, or `no_wrap`. " - "Defaults to `NO_WRAP`. See `torch.distributed.fsdp.wrap.size_based_wrap_policy` for a direction on what it should look like" - }, - ) - ) - cpu_offload: Union[bool, "torch.distributed.fsdp.CPUOffload", "torch.distributed.fsdp.CPUOffloadPolicy"] = field( - default=None, - metadata={ - "help": "Whether to offload parameters to CPU. Should be either a `bool` or an instance of `torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffload` or `torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffloadPolicy` if `fsdp_version` is set to 2. Defaults to `False`" - }, - ) - ignored_modules: Optional[Iterable[torch.nn.Module]] = field( - default=None, - metadata={"help": "A list of modules to ignore when wrapping with FSDP."}, - ) - - state_dict_type: Union[str, "torch.distributed.fsdp.StateDictType"] = field( - default=None, - metadata={ - "help": "State dict type to use. If a string, it must be one of `full_state_dict`, `local_state_dict`, or `sharded_state_dict`. Defaults to `FULL_STATE_DICT`" - }, - ) - state_dict_config: Optional[ - Union[ - "torch.distributed.fsdp.FullStateDictConfig", - "torch.distributed.fsdp.ShardedStateDictConfig", - ] - ] = field( - default=None, - metadata={"help": "State dict config to use. Is determined based on the `state_dict_type` if not passed in."}, - ) - optim_state_dict_config: Optional[ - Union["torch.distributed.fsdp.FullOptimStateDictConfig", "torch.distributed.fsdp.ShardedOptimStateDictConfig"] - ] = field( - default=None, - metadata={ - "help": "Optim state dict config to use. Is determined based on the `state_dict_type` if not passed in." - }, - ) - limit_all_gathers: bool = field( - default=True, - metadata={ - "help": "Whether to have FSDP explicitly synchronizes the CPU thread to prevent " - "too many in-flight all-gathers. This bool only affects the sharded strategies that schedule all-gathers. " - "Enabling this can help lower the number of CUDA malloc retries." - }, - ) - use_orig_params: Optional[bool] = field( - default=None, - metadata={ - "help": "Whether to use the original parameters for the optimizer. Defaults to `False`. This becomes obsolete in FSDP2." - }, - ) - param_init_fn: Optional[Callable[[torch.nn.Module], None]] = field( - default=None, - metadata={ - "help": "A Callable[torch.nn.Module] -> None that specifies how modules " - "that are currently on the meta device should be initialized onto an actual device. " - "Only applicable when `sync_module_states` is `True`. By default is a `lambda` which calls `to_empty` on the module." - }, - ) - sync_module_states: Optional[bool] = field( - default=None, - metadata={ - "help": "Whether each individually wrapped FSDP unit should broadcast module parameters from rank 0 " - "to ensure they are the same across all ranks after initialization. Defaults to `False` unless " - "`cpu_ram_efficient_loading` is `True`, then will be forcibly enabled. This becomes obsolete in FSDP2." - }, - ) - forward_prefetch: bool = field( - default=None, - metadata={ - "help": "Whether to have FSDP explicitly prefetches the next upcoming " - "all-gather while executing in the forward pass. only use with Static graphs. Defaults to `False`" - }, - ) - activation_checkpointing: bool = field( - default=None, - metadata={ - "help": "A technique to reduce memory usage by clearing activations of " - "certain layers and recomputing them during a backward pass. Effectively, this trades extra computation time " - "for reduced memory usage. Defaults to `False`" - }, - ) - cpu_ram_efficient_loading: bool = field( - default=None, - metadata={ - "help": "If True, only the first process loads the pretrained model checkoint while all other processes have empty weights. " - "Only applicable for 🤗 Transformers. When using this, `sync_module_states` needs to be `True`. Defaults to `False`." - }, - ) - transformer_cls_names_to_wrap: Optional[list[str]] = field( - default=None, - metadata={ - "help": "A list of transformer layer class names to wrap. Only applicable when `auto_wrap_policy` is `transformer_based_wrap`." - }, - ) - min_num_params: Optional[int] = field( - default=None, - metadata={ - "help": "The minimum number of parameters a module must have to be wrapped. Only applicable when `auto_wrap_policy` is `size_based_wrap`." - }, - ) - - def __post_init__(self): - from torch.distributed.fsdp import ( - BackwardPrefetch, - ShardingStrategy, - ) - - _fsdp2_warnings = set() - - env_prefix = "FSDP_" - # Strategy: By default we should always assume that values are passed in, else we check the environment variables - if self.fsdp_version is None: - self.fsdp_version = int(os.environ.get(env_prefix + "VERSION", "1")) - - if self.fsdp_version == 2: - if not is_torch_version(">=", FSDP2_PYTORCH_VERSION): - raise ImportError(f"FSDP2 requires PyTorch >= {FSDP2_PYTORCH_VERSION}") - - if self.sharding_strategy is not None: - # We cannot properly detect all of the cases, as by default `args.fsdp_sharding_strategy` is set to `fully_shard` - # Therefore we issue a warning only if the user has explicitly set it inside their plugin - _fsdp2_warnings.add( - "sharding_strategy is deprecated in favor of reshard_after_forward. " - "This will be removed in a future version of Accelerate." - ) - if self.fsdp_version == 1: - if self.sharding_strategy is None: - self.sharding_strategy = os.environ.get(env_prefix + "SHARDING_STRATEGY", "FULL_SHARD") - if isinstance(self.sharding_strategy, str): - if self.sharding_strategy.upper() in FSDP_SHARDING_STRATEGY: - self.sharding_strategy = FSDP_SHARDING_STRATEGY.index(self.sharding_strategy.upper()) + 1 - if isinstance(self.sharding_strategy, int) or self.sharding_strategy.isdigit(): - self.sharding_strategy = ShardingStrategy(int(self.sharding_strategy)) - else: - self.sharding_strategy = ShardingStrategy[self.sharding_strategy.upper()] - - # Fallback to `reshard_after_forward` in FSDP1 if `sharding_strategy` is not set - if self.reshard_after_forward is None and self.sharding_strategy is None: - reshard_after_forward = os.environ.get( - env_prefix + "RESHARD_AFTER_FORWARD", "true" if self.fsdp_version == 2 else "FULL_SHARD" - ) - if self.fsdp_version == 2: - self.reshard_after_forward = str_to_bool(reshard_after_forward.lower(), to_bool=True) - else: - self.reshard_after_forward = reshard_after_forward - if isinstance(self.reshard_after_forward, str): - if self.fsdp_version == 2: - self.reshard_after_forward = str_to_bool(self.reshard_after_forward.lower(), to_bool=True) - else: - # We need to remap based on custom enum values for user readability - if self.reshard_after_forward.upper() in FSDP_SHARDING_STRATEGY: - self.reshard_after_forward = FSDP_SHARDING_STRATEGY.index(self.reshard_after_forward.upper()) + 1 - if isinstance(self.reshard_after_forward, int) or self.reshard_after_forward.isdigit(): - self.reshard_after_forward = ShardingStrategy(int(self.reshard_after_forward)) - else: - self.reshard_after_forward = ShardingStrategy[self.reshard_after_forward.upper()] - - if self.fsdp_version == 2 and not isinstance(self.reshard_after_forward, bool): - raise ValueError( - f"reshard_after_forward set to {self.reshard_after_forward}. This is not supported with FSDP2, please set to a `bool`" - ) - if self.fsdp_version == 1 and isinstance(self.reshard_after_forward, bool): - raise ValueError( - f"reshard_after_forward set to {self.reshard_after_forward}. This is not supported with FSDP1, please set to a `str` or an instance of `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy`" - ) - - if self.cpu_offload is None: - self.cpu_offload = str_to_bool(os.environ.get(env_prefix + "OFFLOAD_PARAMS", "False")) == 1 - - self.set_cpu_offload() # abstracted away to hide imports due to version checks - self.validate_cpu_offload() - - if self.backward_prefetch is None: - self.backward_prefetch = os.environ.get(env_prefix + "BACKWARD_PREFETCH", None) - if isinstance(self.backward_prefetch, str) and self.backward_prefetch.upper() == "NO_PREFETCH": - self.backward_prefetch = None - if self.backward_prefetch is not None and not isinstance(self.backward_prefetch, BackwardPrefetch): - if isinstance(self.backward_prefetch, str) and self.backward_prefetch.upper() in FSDP_BACKWARD_PREFETCH: - self.backward_prefetch = FSDP_BACKWARD_PREFETCH.index(self.backward_prefetch.upper()) + 1 - if isinstance(self.backward_prefetch, int) or self.backward_prefetch.isdigit(): - self.backward_prefetch = BackwardPrefetch(int(self.backward_prefetch)) - else: - self.backward_prefetch = BackwardPrefetch[self.backward_prefetch.upper()] - if self.fsdp_version == 2 and self.backward_prefetch is not None: - _fsdp2_warnings.add("backward_prefetch is not supported in FSDP2. Setting backward prefetch to None.") - self.backward_prefetch = None - - self.set_state_dict_type() - - if self.auto_wrap_policy is None: - self.auto_wrap_policy = os.environ.get(env_prefix + "AUTO_WRAP_POLICY", "NO_WRAP") - if isinstance(self.auto_wrap_policy, str): - if self.auto_wrap_policy.upper() not in FSDP_AUTO_WRAP_POLICY: - raise ValueError( - f"Invalid auto wrap policy: {self.auto_wrap_policy}. Must be one of {FSDP_AUTO_WRAP_POLICY}" - ) - from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy - - if self.auto_wrap_policy.upper() == "TRANSFORMER_BASED_WRAP": - self.auto_wrap_policy = transformer_auto_wrap_policy - if self.transformer_cls_names_to_wrap is None: - self.transformer_cls_names_to_wrap = os.environ.get(env_prefix + "TRANSFORMER_CLS_TO_WRAP", None) - if isinstance(self.transformer_cls_names_to_wrap, str): - self.transformer_cls_names_to_wrap = self.transformer_cls_names_to_wrap.split(",") - elif self.auto_wrap_policy.upper() == "SIZE_BASED_WRAP": - self.auto_wrap_policy = size_based_auto_wrap_policy - if self.min_num_params is None: - self.min_num_params = int(os.environ.get(env_prefix + "MIN_NUM_PARAMS", 0)) - elif not isinstance(self.min_num_params, int): - raise ValueError( - f"`min_num_params` must be an integer. Got {self.min_num_params} of type {type(self.min_num_params)}" - ) - elif self.auto_wrap_policy.upper() == "NO_WRAP": - self.auto_wrap_policy = None - - if self.use_orig_params is None and self.fsdp_version == 1: - self.use_orig_params = str_to_bool(os.environ.get(env_prefix + "USE_ORIG_PARAMS", "False")) == 1 - if self.fsdp_version == 2 and self.use_orig_params is not None: - _fsdp2_warnings.add("use_orig_params is obsolete in FSDP2, as FSDP2 always uses the original parameters.") - self.use_orig_params = None - - if self.sync_module_states is None and self.fsdp_version == 1: - self.sync_module_states = str_to_bool(os.environ.get(env_prefix + "SYNC_MODULE_STATES", "False")) == 1 - if self.fsdp_version == 2 and self.sync_module_states is not None: - _fsdp2_warnings.add( - "sync_module_states is obsolete in FSDP2, as it is not needed anymore." - "Setting sync_module_states to None." - ) - self.sync_module_states = None - - if self.forward_prefetch is None and self.fsdp_version == 1: - self.forward_prefetch = str_to_bool(os.environ.get(env_prefix + "FORWARD_PREFETCH", "False")) == 1 - if self.fsdp_version == 2 and self.forward_prefetch is not None: - raise ValueError("forward_prefetch is not yet implemented in FSDP2, set to None or use `fsdp_version=1`") - - if self.activation_checkpointing is None: - self.activation_checkpointing = ( - str_to_bool(os.environ.get(env_prefix + "ACTIVATION_CHECKPOINTING", "False")) == 1 - ) - - if self.cpu_ram_efficient_loading is None: - self.cpu_ram_efficient_loading = ( - str_to_bool(os.environ.get(env_prefix + "CPU_RAM_EFFICIENT_LOADING", "False")) == 1 - ) - # There's no need to specify sync_module_states in FSDP2 - if self.fsdp_version == 1 and self.cpu_ram_efficient_loading and not self.sync_module_states: - warnings.warn( - "sync_module_states cannot be False since efficient cpu ram loading enabled. " - "Setting sync_module_states to True." - ) - self.sync_module_states = True - - if self.cpu_ram_efficient_loading != bool( - str_to_bool(os.environ.get(env_prefix + "CPU_RAM_EFFICIENT_LOADING", "False")) - ): - env_var = env_prefix + "CPU_RAM_EFFICIENT_LOADING" - warnings.warn( - f"The `cpu_ram_efficient_loading` flag for `FullyShardedDataParallelPlugin` does not match the environment variable {env_var}. " - "Setting environment variable to match `cpu_ram_efficient_loading`." - ) - os.environ[env_var] = str(self.cpu_ram_efficient_loading) - - if isinstance(self.mixed_precision_policy, dict): - self.set_mixed_precision(self.mixed_precision_policy) - if self.mixed_precision_policy is not None: - self.validate_mixed_precision_policy() - - if self.sync_module_states: - if is_npu_available(): - device = torch.npu.current_device() - elif is_mlu_available(): - device = torch.mlu.current_device() - elif is_musa_available(): - device = torch.musa.current_device() - elif is_cuda_available(): - device = torch.cuda.current_device() - elif is_xpu_available(): - device = torch.xpu.current_device() - elif is_hpu_available(): - device = torch.hpu.current_device() - else: - raise RuntimeError( - "There are currently no available devices found, must be one of 'XPU', 'CUDA', 'MLU', 'NPU', 'MUSA', or 'HPU'." - ) - # Create a function that will be used to initialize the parameters of the model - # when using `sync_module_states` - self.param_init_fn = lambda x: x.to_empty(device=device, recurse=False) - - # Single warning for all deprecation warnings due to FSDP2 conversion - if _fsdp2_warnings: - logger.warning("Multiple deprecation warnings due to FSDP2 conversion:\n".join(_fsdp2_warnings)) - - def set_state_dict_type(self, state_dict_type=None): - """ - Set the state dict config based on the `StateDictType`. - """ - from torch.distributed.fsdp.fully_sharded_data_parallel import ( - FullOptimStateDictConfig, - FullStateDictConfig, - ShardedOptimStateDictConfig, - ShardedStateDictConfig, - StateDictType, - ) - - # Override the state_dict_type if provided, typical use case: - # user trains with sharded, but final save is with full - if state_dict_type is not None: - self.state_dict_type = state_dict_type - - if self.state_dict_type is None: - self.state_dict_type = os.environ.get( - "FSDP_STATE_DICT_TYPE", "FULL_STATE_DICT" if self.fsdp_version == 1 else "SHARDED_STATE_DICT" - ) - if isinstance(self.state_dict_type, str): - if self.state_dict_type.isdigit(): - self.state_dict_type = StateDictType(int(self.state_dict_type)) - else: - self.state_dict_type = StateDictType[self.state_dict_type.upper()] - - if self.state_dict_type == StateDictType.FULL_STATE_DICT: - if self.state_dict_config is None: - self.state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) - if self.optim_state_dict_config is None: - self.optim_state_dict_config = FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True) - elif self.state_dict_type == StateDictType.SHARDED_STATE_DICT: - if self.state_dict_config is None: - self.state_dict_config = ShardedStateDictConfig(offload_to_cpu=True) - if self.optim_state_dict_config is None: - self.optim_state_dict_config = ShardedOptimStateDictConfig(offload_to_cpu=True) - - if self.fsdp_version == 2 and self.state_dict_type == StateDictType.LOCAL_STATE_DICT: - raise ValueError( - "FSDP2 does not support LOCAL_STATE_DICT. " - "Please set `fsdp_state_dict_type` to `SHARDED_STATE_DICT` or `FULL_STATE_DICT`." - ) - - def set_auto_wrap_policy(self, model): - """ - Given `model`, creates an `auto_wrap_policy` baesd on the passed in policy and if we can use the - `transformer_cls_to_wrap` - """ - from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy - - # First base off of `_no_split_modules` - no_split_modules = getattr(model, "_no_split_modules", None) - default_transformer_cls_names_to_wrap = list(no_split_modules) if no_split_modules is not None else [] - if self.auto_wrap_policy == transformer_auto_wrap_policy: - if self.transformer_cls_names_to_wrap is None: - self.transformer_cls_names_to_wrap = default_transformer_cls_names_to_wrap - transformer_cls_to_wrap = set() - for layer_class in self.transformer_cls_names_to_wrap: - transformer_cls = get_module_class_from_name(model, layer_class) - if transformer_cls is None: - raise ValueError(f"Could not find the transformer layer class {layer_class} in the model.") - transformer_cls_to_wrap.add(transformer_cls) - # Finally we set the auto_wrap_policy to a callable - self.auto_wrap_policy = functools.partial( - self.auto_wrap_policy, transformer_layer_cls=transformer_cls_to_wrap - ) - - elif self.auto_wrap_policy == size_based_auto_wrap_policy: - # If zero, we silently ignore it. - if self.min_num_params > 0: - self.auto_wrap_policy = functools.partial(self.auto_wrap_policy, min_num_params=self.min_num_params) - else: - self.auto_wrap_policy = None - - def set_mixed_precision(self, mixed_precision, buffer_autocast=False, override=False): - "Sets the mixed precision policy for FSDP" - mixed_precision_mapping = { - "fp8": torch.bfloat16, - "fp16": torch.float16, - "bf16": torch.bfloat16, - "fp32": torch.float32, - } - dtype = mixed_precision - if isinstance(mixed_precision, str): - dtype = mixed_precision_mapping.get(mixed_precision, None) - if dtype is None: - raise ValueError( - f"Invalid mixed precision: {mixed_precision}. Must be one of {list(mixed_precision_mapping.keys())}" - ) - elif isinstance(mixed_precision, torch.dtype) and mixed_precision not in mixed_precision_mapping.values(): - raise ValueError( - f"Invalid mixed precision: {mixed_precision}. Must be one of {list(mixed_precision_mapping.values())}" - ) - - buffer_type = torch.float32 if buffer_autocast else dtype - - if self.fsdp_version == 1: - from torch.distributed.fsdp import MixedPrecision - elif self.fsdp_version == 2: - from torch.distributed.fsdp import MixedPrecisionPolicy as MixedPrecision - - if override or self.mixed_precision_policy is None: - dtype_args = {"param_dtype": dtype, "reduce_dtype": dtype} - if self.fsdp_version == 1: - dtype_args["buffer_dtype"] = buffer_type - else: - dtype_args["output_dtype"] = dtype - # TODO(s1ro1): `cast_forward_inputs` for FSDP2? - self.mixed_precision_policy = MixedPrecision(**dtype_args) - elif isinstance(self.mixed_precision_policy, dict): - # Check for incompatible types - valid_keys = ["param_dtype", "reduce_dtype"] + ( - ["buffer_dtype"] if self.fsdp_version == 1 else ["output_dtype"] - ) - missing_keys = [k for k in valid_keys if k not in self.mixed_precision_policy] - invalid_values = [ - k for k, v in self.mixed_precision_policy.items() if v not in mixed_precision_mapping.values() - ] - if missing_keys or invalid_values: - raise ValueError( - f"Invalid mixed precision policy: {self.mixed_precision_policy}. " - f"Must be a `dict` with keys {valid_keys}." - f"Values must be one of {list(mixed_precision_mapping.values())}" - ) - self.mixed_precision_policy = MixedPrecision(**self.mixed_precision_policy) - - def validate_mixed_precision_policy(self): - """ - Validates the mixed precision policy, abstracted away to not bring in the imports if not needed. - """ - if self.fsdp_version == 2: - from torch.distributed.fsdp import MixedPrecisionPolicy as MixedPrecision - else: - from torch.distributed.fsdp import MixedPrecision - - if not isinstance(self.mixed_precision_policy, MixedPrecision): - required_type = ( - "`torch.distributed.fsdp.MixedPrecisionPolicy`" - if self.fsdp_version == 2 - else "`torch.distributed.fsdp.MixedPrecision`" - ) - raise ValueError(f"mixed_precision_policy must be an instance of {required_type}.") - - def set_cpu_offload(self): - if self.fsdp_version == 2: - from torch.distributed.fsdp import CPUOffloadPolicy, OffloadPolicy - else: - from torch.distributed.fsdp import CPUOffload - - if isinstance(self.cpu_offload, bool): - if self.fsdp_version == 2: - if not self.cpu_offload: - self.cpu_offload = OffloadPolicy() - else: - self.cpu_offload = CPUOffloadPolicy() - else: - self.cpu_offload = CPUOffload(offload_params=self.cpu_offload) - - def validate_cpu_offload(self): - if self.fsdp_version == 2: - from torch.distributed.fsdp import OffloadPolicy - else: - from torch.distributed.fsdp import CPUOffload - - if self.fsdp_version == 2 and not isinstance(self.cpu_offload, OffloadPolicy): - raise ValueError( - f"`cpu_offload` must be an instance of `torch.distributed.fsdp.OffloadPolicy` in FSDP2, got {self.cpu_offload}" - ) - if self.fsdp_version == 1 and not isinstance(self.cpu_offload, CPUOffload): - raise ValueError( - f"`cpu_offload` must be an instance of `torch.distributed.fsdp.CPUOffload` in FSDP1, got {self.cpu_offload}" - ) - - -@dataclass -class TorchTensorParallelPlugin: - """ - This plugin is used to enable tensor parallelism using PyTorch >= 2.0. - """ - - tp_size: int = field( - default=1, - metadata={"help": "tensor parallel size will be used in the device mesh preparation"}, - ) - - # torch_device_mesh is fo type "torch.distributed.DeviceMesh" - torch_device_mesh: Optional["torch.distributed.DeviceMesh"] = field(default=None) - - def __post_init__(self): - if not isinstance(self.tp_size, int): - raise ValueError(f"`tp_size` set to {self.tp_size}, please set to an `int`.") - - if self.tp_size <= 1: - raise ValueError("`tp_size` must be greater than 1.") - - if is_torch_version("<", BETA_TP_AVAILABLE_PYTORCH_VERSION): - raise ValueError( - f"Minimum PyTorch version {BETA_TP_AVAILABLE_PYTORCH_VERSION} needed to use tensor parallel." - ) - from torch.distributed.device_mesh import init_device_mesh - - # support for other devices has to be investigated - if is_hpu_available(init_hccl=True): - device = "hpu" - elif is_xpu_available(): - device = "xpu" - else: - device = "cuda" - - mesh_dim_name = "tp" - - # device mesh is not used for model sharding - # it is only used for preparing data loader - self.torch_device_mesh = init_device_mesh(device, (self.tp_size,), mesh_dim_names=(mesh_dim_name,)) - - -@dataclass -class MegatronLMPlugin: - """ - Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective - activation recomputation and optimized fused kernels. - - Args: - tp_degree (`int`, defaults to `None`): - Tensor parallelism degree. - pp_degree (`int`, defaults to `None`): - Pipeline parallelism degree. - num_micro_batches (`int`, defaults to `None`): - Number of micro-batches. - gradient_clipping (`float`, defaults to `None`): - Gradient clipping value based on global L2 Norm (0 to disable). - sequence_parallelism (`bool`, defaults to `None`): - Enable sequence parallelism. - recompute_activations (`bool`, defaults to `None`): - Enable selective activation recomputation. - use_distributed_optimizr (`bool`, defaults to `None`): - Enable distributed optimizer. - pipeline_model_parallel_split_rank (`int`, defaults to `None`): - Rank where encoder and decoder should be split. - num_layers_per_virtual_pipeline_stage (`int`, defaults to `None`): - Number of layers per virtual pipeline stage. - is_train_batch_min (`str`, defaults to `True`): - If both tran & eval dataloaders are specified, this will decide the `micro_batch_size`. - train_iters (`int`, defaults to `None`): - Total number of samples to train over all training runs. Note that either train-iters or train-samples - should be provided when using `MegatronLMDummyScheduler`. - train_samples (`int`, defaults to `None`): - Total number of samples to train over all training runs. Note that either train-iters or train-samples - should be provided when using `MegatronLMDummyScheduler`. - weight_decay_incr_style (`str`, defaults to `'constant'`): - Weight decay increment function. choices=["constant", "linear", "cosine"]. - start_weight_decay (`float`, defaults to `None`): - Initial weight decay coefficient for L2 regularization. - end_weight_decay (`float`, defaults to `None`): - End of run weight decay coefficient for L2 regularization. - lr_decay_style (`str`, defaults to `'linear'`): - Learning rate decay function. choices=['constant', 'linear', 'cosine']. - lr_decay_iters (`int`, defaults to `None`): - Number of iterations for learning rate decay. If None defaults to `train_iters`. - lr_decay_samples (`int`, defaults to `None`): - Number of samples for learning rate decay. If None defaults to `train_samples`. - lr_warmup_iters (`int`, defaults to `None`): - Number of iterations to linearly warmup learning rate over. - lr_warmup_samples (`int`, defaults to `None`): - Number of samples to linearly warmup learning rate over. - lr_warmup_fraction (`float`, defaults to `None`): - Fraction of lr-warmup-(iters/samples) to linearly warmup learning rate over. - min_lr (`float`, defaults to `0`): - Minumum value for learning rate. The scheduler clip values below this threshold. - consumed_samples (`List`, defaults to `None`): - Number of samples consumed in the same order as the dataloaders to `accelerator.prepare` call. - no_wd_decay_cond (`Optional`, defaults to `None`): - Condition to disable weight decay. - scale_lr_cond (`Optional`, defaults to `None`): - Condition to scale learning rate. - lr_mult (`float`, defaults to `1.0`): - Learning rate multiplier. - megatron_dataset_flag (`bool`, defaults to `False`): - Whether the format of dataset follows Megatron-LM Indexed/Cached/MemoryMapped format. - seq_length (`int`, defaults to `None`): - Maximum sequence length to process. - encoder_seq_length (`int`, defaults to `None`): - Maximum sequence length to process for the encoder. - decoder_seq_length (`int`, defaults to `None`): - Maximum sequence length to process for the decoder. - tensorboard_dir (`str`, defaults to `None`): - Path to save tensorboard logs. - set_all_logging_options (`bool`, defaults to `False`): - Whether to set all logging options. - eval_iters (`int`, defaults to `100`): - Number of iterations to run for evaluation validation/test for. - eval_interval (`int`, defaults to `1000`): - Interval between running evaluation on validation set. - return_logits (`bool`, defaults to `False`): - Whether to return logits from the model. - custom_train_step_class (`Optional`, defaults to `None`): - Custom train step class. - custom_train_step_kwargs (`Optional`, defaults to `None`): - Custom train step kwargs. - custom_model_provider_function (`Optional`, defaults to `None`): - Custom model provider function. - custom_prepare_model_function (`Optional`, defaults to `None`): - Custom prepare model function. - custom_megatron_datasets_provider_function (`Optional`, defaults to `None`): - Custom megatron train_valid_test datasets provider function. - custom_get_batch_function (`Optional`, defaults to `None`): - Custom get batch function. - custom_loss_function (`Optional`, defaults to `None`): - Custom loss function. - other_megatron_args (`Optional`, defaults to `None`): - Other Megatron-LM arguments. Please refer Megatron-LM. - """ - - tp_degree: int = field(default=None, metadata={"help": "tensor parallelism degree."}) - pp_degree: int = field(default=None, metadata={"help": "pipeline parallelism degree."}) - num_micro_batches: int = field(default=None, metadata={"help": "number of micro-batches."}) - gradient_clipping: float = field( - default=None, metadata={"help": "gradient clipping value based on global L2 Norm (0 to disable)"} - ) - sequence_parallelism: bool = field( - default=None, - metadata={"help": "enable sequence parallelism"}, - ) - recompute_activations: bool = field( - default=None, - metadata={"help": "enable selective activation recomputation"}, - ) - use_distributed_optimizer: bool = field( - default=None, - metadata={"help": "enable distributed optimizer"}, - ) - pipeline_model_parallel_split_rank: int = field( - default=None, metadata={"help": "Rank where encoder and decoder should be split."} - ) - num_layers_per_virtual_pipeline_stage: int = field( - default=None, metadata={"help": "Number of layers per virtual pipeline stage."} - ) - is_train_batch_min: str = field( - default=True, - metadata={"help": "If both train & eval dataloaders are specified, this will decide the micro_batch_size"}, - ) - train_iters: int = field( - default=None, - metadata={ - "help": "Total number of iterations to train over all training runs. " - "Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler`" - }, - ) - train_samples: int = field( - default=None, - metadata={ - "help": "Total number of samples to train over all training runs. " - "Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler`" - }, - ) - weight_decay_incr_style: str = field( - default="constant", - metadata={"help": 'Weight decay increment function. choices=["constant", "linear", "cosine"]. '}, - ) - start_weight_decay: float = field( - default=None, - metadata={"help": "Initial weight decay coefficient for L2 regularization."}, - ) - end_weight_decay: float = field( - default=None, - metadata={"help": "End of run weight decay coefficient for L2 regularization."}, - ) - lr_decay_style: str = field( - default="linear", - metadata={"help": "Learning rate decay function. choices=['constant', 'linear', 'cosine']."}, - ) - lr_decay_iters: int = field( - default=None, - metadata={"help": "Number of iterations for learning rate decay. If None defaults to `train_iters`."}, - ) - lr_decay_samples: int = field( - default=None, - metadata={"help": "Number of samples for learning rate decay. If None defaults to `train_samples`."}, - ) - lr_warmup_iters: int = field( - default=None, - metadata={"help": "number of iterations to linearly warmup learning rate over."}, - ) - lr_warmup_samples: int = field( - default=None, - metadata={"help": "number of samples to linearly warmup learning rate over."}, - ) - lr_warmup_fraction: float = field( - default=None, - metadata={"help": "fraction of lr-warmup-(iters/samples) to linearly warmup learning rate over."}, - ) - min_lr: float = field( - default=0, - metadata={"help": "Minumum value for learning rate. The scheduler clip values below this threshold."}, - ) - consumed_samples: list[int] = field( - default=None, - metadata={ - "help": "Number of samples consumed in the same order as the dataloaders to `accelerator.prepare` call." - }, - ) - no_wd_decay_cond: Optional[Callable] = field(default=None, metadata={"help": "Condition to disable weight decay."}) - scale_lr_cond: Optional[Callable] = field(default=None, metadata={"help": "Condition to scale learning rate."}) - lr_mult: float = field(default=1.0, metadata={"help": "Learning rate multiplier."}) - megatron_dataset_flag: bool = field( - default=False, - metadata={"help": "Whether the format of dataset follows Megatron-LM Indexed/Cached/MemoryMapped format."}, - ) - seq_length: int = field( - default=None, - metadata={"help": "Maximum sequence length to process."}, - ) - encoder_seq_length: int = field( - default=None, - metadata={"help": "Maximum sequence length to process for the encoder."}, - ) - decoder_seq_length: int = field( - default=None, - metadata={"help": "Maximum sequence length to process for the decoder."}, - ) - tensorboard_dir: str = field( - default=None, - metadata={"help": "Path to save tensorboard logs."}, - ) - set_all_logging_options: bool = field( - default=False, - metadata={"help": "Whether to set all logging options."}, - ) - eval_iters: int = field( - default=100, metadata={"help": "Number of iterations to run for evaluation validation/test for."} - ) - eval_interval: int = field( - default=1000, metadata={"help": "Interval between running evaluation on validation set."} - ) - return_logits: bool = field( - default=False, - metadata={"help": "Whether to return logits from the model."}, - ) - - # custom train step args - custom_train_step_class: Optional[Any] = field( - default=None, - metadata={"help": "Custom train step class."}, - ) - custom_train_step_kwargs: Optional[dict[str, Any]] = field( - default=None, - metadata={"help": "Custom train step kwargs."}, - ) - - # custom model args - custom_model_provider_function: Optional[Callable] = field( - default=None, - metadata={"help": "Custom model provider function."}, - ) - custom_prepare_model_function: Optional[Callable] = field( - default=None, - metadata={"help": "Custom prepare model function."}, - ) - custom_megatron_datasets_provider_function: Optional[Callable] = field( - default=None, - metadata={"help": "Custom megatron train_valid_test datasets provider function."}, - ) - custom_get_batch_function: Optional[Callable] = field( - default=None, - metadata={"help": "Custom get batch function."}, - ) - custom_loss_function: Optional[Callable] = field( - default=None, - metadata={"help": "Custom loss function."}, - ) - - # remaining args such as enabling Alibi/ROPE positional embeddings, - # wandb logging, Multi-Query Attention, etc. - other_megatron_args: Optional[dict[str, Any]] = field( - default=None, - metadata={"help": "Other Megatron-LM arguments. Please refer Megatron-LM"}, - ) - - def __post_init__(self): - prefix = "MEGATRON_LM_" - if self.tp_degree is None: - self.tp_degree = int(os.environ.get(prefix + "TP_DEGREE", 1)) - if self.pp_degree is None: - self.pp_degree = int(os.environ.get(prefix + "PP_DEGREE", 1)) - if self.num_micro_batches is None: - self.num_micro_batches = int(os.environ.get(prefix + "NUM_MICRO_BATCHES", 1)) - if self.gradient_clipping is None: - self.gradient_clipping = float(os.environ.get(prefix + "GRADIENT_CLIPPING", 1.0)) - if self.recompute_activations is None: - self.recompute_activations = str_to_bool(os.environ.get(prefix + "RECOMPUTE_ACTIVATIONS", "False")) == 1 - if self.use_distributed_optimizer is None: - self.use_distributed_optimizer = ( - str_to_bool(os.environ.get(prefix + "USE_DISTRIBUTED_OPTIMIZER", "False")) == 1 - ) - if self.sequence_parallelism is None: - self.sequence_parallelism = str_to_bool(os.environ.get(prefix + "SEQUENCE_PARALLELISM", "False")) == 1 - - if self.pp_degree > 1 or self.use_distributed_optimizer: - self.DDP_impl = "local" - else: - self.DDP_impl = "torch" - - if self.consumed_samples is not None: - if len(self.consumed_samples) == 1: - self.consumed_samples.extend([0, 0]) - elif len(self.consumed_samples) == 2: - self.consumed_samples.append(0) - - self.megatron_lm_default_args = { - "tensor_model_parallel_size": self.tp_degree, - "pipeline_model_parallel_size": self.pp_degree, - "pipeline_model_parallel_split_rank": self.pipeline_model_parallel_split_rank, - "num_layers_per_virtual_pipeline_stage": self.num_layers_per_virtual_pipeline_stage, - "DDP_impl": self.DDP_impl, - "use_distributed_optimizer": self.use_distributed_optimizer, - "sequence_parallel": self.sequence_parallelism, - "clip_grad": self.gradient_clipping, - "num_micro_batches": self.num_micro_batches, - "consumed_samples": self.consumed_samples, - "no_wd_decay_cond": self.no_wd_decay_cond, - "scale_lr_cond": self.scale_lr_cond, - "lr_mult": self.lr_mult, - "megatron_dataset_flag": self.megatron_dataset_flag, - "eval_iters": self.eval_iters, - "eval_interval": self.eval_interval, - } - if self.recompute_activations: - self.megatron_lm_default_args["recompute_granularity"] = "selective" - if self.tensorboard_dir is not None: - self.megatron_lm_default_args["tensorboard_dir"] = self.tensorboard_dir - if self.set_all_logging_options: - self.set_tensorboard_logging_options() - if self.other_megatron_args is not None: - self.megatron_lm_default_args.update(self.other_megatron_args) - - def set_network_size_args(self, model, batch_data=None): - model_config_type = model.config.model_type.lower() - for model_type in MODEL_CONFIGS_TO_MEGATRON_PARSERS.keys(): - if model_type in model_config_type: - MODEL_CONFIGS_TO_MEGATRON_PARSERS[model_type](self, model, batch_data) - return - raise ValueError( - f"Accelerate Megatron-LM integration not supports {model_config_type} model. " - "You can add your own model config parser." - ) - - def set_mixed_precision(self, mixed_precision): - if mixed_precision == "fp16": - self.megatron_lm_default_args["fp16"] = True - elif mixed_precision == "bf16": - self.megatron_lm_default_args["bf16"] = True - self.DDP_impl = "local" - self.megatron_lm_default_args["DDP_impl"] = self.DDP_impl - - def set_training_args(self, micro_batch_size, dp_degree): - self.data_parallel_size = dp_degree - self.micro_batch_size = micro_batch_size - self.global_batch_size = dp_degree * micro_batch_size * self.num_micro_batches - self.megatron_lm_default_args["data_parallel_size"] = self.data_parallel_size - self.megatron_lm_default_args["micro_batch_size"] = self.micro_batch_size - self.megatron_lm_default_args["global_batch_size"] = self.global_batch_size - - def set_optimizer_type(self, optimizer): - optimizer_name = optimizer.__class__.__name__.lower() - if "adam" in optimizer_name: - self.megatron_lm_default_args["optimizer"] = "adam" - self.megatron_lm_default_args["adam_beta1"] = optimizer.defaults["betas"][0] - self.megatron_lm_default_args["adam_beta2"] = optimizer.defaults["betas"][1] - self.megatron_lm_default_args["adam_eps"] = optimizer.defaults["eps"] - elif "sgd" in optimizer_name: - self.megatron_lm_default_args["optimizer"] = "sgd" - self.megatron_lm_default_args["sgd_momentum"] = optimizer.defaults["momentum"] - else: - raise ValueError(f"Optimizer {optimizer_name} is not supported by Megatron-LM") - - self.megatron_lm_default_args["lr"] = optimizer.defaults["lr"] - self.megatron_lm_default_args["weight_decay"] = optimizer.defaults["weight_decay"] - - def set_scheduler_args(self, scheduler): - if self.train_iters is None: - self.train_iters = scheduler.total_num_steps // self.megatron_lm_default_args["data_parallel_size"] - if self.train_samples is not None: - self.train_samples = None - warnings.warn( - "Ignoring `train_samples` as `train_iters` based on scheduler is being used for training." - ) - if self.lr_warmup_iters is None: - self.lr_warmup_iters = scheduler.warmup_num_steps // self.megatron_lm_default_args["data_parallel_size"] - if self.lr_warmup_samples is not None: - warnings.warn( - "Ignoring `lr_warmup_samples` as `lr_warmup_iters` based on scheduler is being used for training." - ) - self.lr_warmup_samples = 0 - - self.megatron_lm_default_args["train_iters"] = self.train_iters - self.megatron_lm_default_args["lr_warmup_iters"] = self.lr_warmup_iters - self.megatron_lm_default_args["train_samples"] = self.train_samples - self.megatron_lm_default_args["lr_warmup_samples"] = self.lr_warmup_samples - self.megatron_lm_default_args["lr_decay_iters"] = self.lr_decay_iters - self.megatron_lm_default_args["lr_decay_samples"] = self.lr_decay_samples - self.megatron_lm_default_args["lr_warmup_fraction"] = self.lr_warmup_fraction - self.megatron_lm_default_args["lr_decay_style"] = self.lr_decay_style - self.megatron_lm_default_args["weight_decay_incr_style"] = self.weight_decay_incr_style - self.megatron_lm_default_args["start_weight_decay"] = self.start_weight_decay - self.megatron_lm_default_args["end_weight_decay"] = self.end_weight_decay - self.megatron_lm_default_args["min_lr"] = self.min_lr - - def set_tensorboard_logging_options(self): - from megatron.training.arguments import _add_logging_args - - parser = argparse.ArgumentParser() - parser = _add_logging_args(parser) - logging_args = parser.parse_known_args() - self.dataset_args = vars(logging_args[0]) - for key, value in self.dataset_args.items(): - if key.startswith("log_"): - self.megatron_lm_default_args[key] = True - elif key.startswith("no_log_"): - self.megatron_lm_default_args[key.replace("no_", "")] = True - - -MODEL_CONFIGS_TO_MEGATRON_PARSERS = {} - - -def add_model_config_to_megatron_parser(model_type: str): - def add_model_config_parser_helper(func): - @functools.wraps(func) - def wrapper(*args, **kwargs): - return func(*args, **kwargs) - - MODEL_CONFIGS_TO_MEGATRON_PARSERS[model_type] = func - return wrapper - - return add_model_config_parser_helper - - -@add_model_config_to_megatron_parser("megatron-bert") -def parse_bert_config(megatron_lm_plugin, model, batch_data): - model_type_name = "bert" - num_layers = model.config.num_hidden_layers - hidden_size = model.config.hidden_size - num_attention_heads = model.config.num_attention_heads - max_position_embeddings = model.config.max_position_embeddings - num_labels = model.config.num_labels - orig_vocab_size = model.config.vocab_size - pretraining_flag = False - if "maskedlm" in model.__class__.__name__.lower(): - pretraining_flag = True - if megatron_lm_plugin.seq_length is not None: - if megatron_lm_plugin.encoder_seq_length is not None: - warnings.warn("Both `seq_length` and `encoder_seq_length` are set. Using `encoder_seq_length`.") - megatron_lm_plugin.seq_length = megatron_lm_plugin.encoder_seq_length - elif megatron_lm_plugin.encoder_seq_length is not None: - megatron_lm_plugin.seq_length = megatron_lm_plugin.encoder_seq_length - elif batch_data is not None: - megatron_lm_plugin.seq_length = batch_data["input_ids"].shape[1] - else: - megatron_lm_plugin.seq_length = max_position_embeddings - megatron_lm_plugin.megatron_lm_default_args["seq_length"] = megatron_lm_plugin.seq_length - megatron_lm_plugin.megatron_lm_default_args["model_type_name"] = model_type_name - megatron_lm_plugin.megatron_lm_default_args["num_layers"] = num_layers - megatron_lm_plugin.megatron_lm_default_args["hidden_size"] = hidden_size - megatron_lm_plugin.megatron_lm_default_args["num_attention_heads"] = num_attention_heads - megatron_lm_plugin.megatron_lm_default_args["max_position_embeddings"] = max_position_embeddings - megatron_lm_plugin.megatron_lm_default_args["pretraining_flag"] = pretraining_flag - megatron_lm_plugin.megatron_lm_default_args["orig_vocab_size"] = orig_vocab_size - megatron_lm_plugin.megatron_lm_default_args["model_return_dict"] = model.config.return_dict - megatron_lm_plugin.megatron_lm_default_args["num_labels"] = num_labels - - -@add_model_config_to_megatron_parser("gpt2") -def parse_gpt2_config(megatron_lm_plugin, model, batch_data): - model_type_name = "gpt" - num_layers = model.config.n_layer - hidden_size = model.config.n_embd - num_attention_heads = model.config.n_head - max_position_embeddings = model.config.n_positions - orig_vocab_size = model.config.vocab_size - pretraining_flag = True - if megatron_lm_plugin.seq_length is not None: - if megatron_lm_plugin.decoder_seq_length is not None: - warnings.warn("Both `seq_length` and `decoder_seq_length` are set. Using `decoder_seq_length`.") - megatron_lm_plugin.seq_length = megatron_lm_plugin.decoder_seq_length - elif megatron_lm_plugin.decoder_seq_length is not None: - megatron_lm_plugin.seq_length = megatron_lm_plugin.decoder_seq_length - elif batch_data is not None: - megatron_lm_plugin.seq_length = batch_data["input_ids"].shape[1] - else: - megatron_lm_plugin.seq_length = max_position_embeddings - megatron_lm_plugin.megatron_lm_default_args["seq_length"] = megatron_lm_plugin.seq_length - megatron_lm_plugin.megatron_lm_default_args["return_logits"] = megatron_lm_plugin.return_logits - megatron_lm_plugin.megatron_lm_default_args["tokenizer_type"] = "GPT2BPETokenizer" - megatron_lm_plugin.megatron_lm_default_args["model_type_name"] = model_type_name - megatron_lm_plugin.megatron_lm_default_args["num_layers"] = num_layers - megatron_lm_plugin.megatron_lm_default_args["hidden_size"] = hidden_size - megatron_lm_plugin.megatron_lm_default_args["num_attention_heads"] = num_attention_heads - megatron_lm_plugin.megatron_lm_default_args["max_position_embeddings"] = max_position_embeddings - megatron_lm_plugin.megatron_lm_default_args["pretraining_flag"] = pretraining_flag - megatron_lm_plugin.megatron_lm_default_args["orig_vocab_size"] = orig_vocab_size - megatron_lm_plugin.megatron_lm_default_args["model_return_dict"] = model.config.return_dict - - -@add_model_config_to_megatron_parser("t5") -def parse_t5_config(megatron_lm_plugin, model, batch_data): - model_type_name = "t5" - num_layers = model.config.num_layers - hidden_size = model.config.d_model - num_attention_heads = model.config.num_heads - max_position_embeddings = model.config.n_positions if hasattr(model.config, "n_positions") else 1024 - orig_vocab_size = model.config.vocab_size - pretraining_flag = True - if megatron_lm_plugin.encoder_seq_length is None: - if batch_data is not None: - megatron_lm_plugin.encoder_seq_length = batch_data["input_ids"].shape[1] - else: - megatron_lm_plugin.encoder_seq_length = max_position_embeddings - if megatron_lm_plugin.decoder_seq_length is None: - if batch_data is not None: - megatron_lm_plugin.decoder_seq_length = batch_data["labels"].shape[1] - else: - megatron_lm_plugin.decoder_seq_length = max_position_embeddings - megatron_lm_plugin.megatron_lm_default_args["encoder_seq_length"] = megatron_lm_plugin.encoder_seq_length - megatron_lm_plugin.megatron_lm_default_args["decoder_seq_length"] = megatron_lm_plugin.decoder_seq_length - megatron_lm_plugin.megatron_lm_default_args["model_type_name"] = model_type_name - megatron_lm_plugin.megatron_lm_default_args["num_layers"] = num_layers - megatron_lm_plugin.megatron_lm_default_args["hidden_size"] = hidden_size - megatron_lm_plugin.megatron_lm_default_args["num_attention_heads"] = num_attention_heads - megatron_lm_plugin.megatron_lm_default_args["max_position_embeddings"] = max_position_embeddings - megatron_lm_plugin.megatron_lm_default_args["pretraining_flag"] = pretraining_flag - megatron_lm_plugin.megatron_lm_default_args["orig_vocab_size"] = orig_vocab_size - megatron_lm_plugin.megatron_lm_default_args["model_return_dict"] = model.config.return_dict - - -@add_model_config_to_megatron_parser("llama") -def parse_llama_config(megatron_lm_plugin, model, batch_data): - model_type_name = "gpt" - num_layers = model.config.num_hidden_layers - pretraining_flag = True - hidden_size = model.config.hidden_size - num_attention_heads = model.config.num_attention_heads - orig_vocab_size = model.config.vocab_size - - max_position_embeddings = model.config.max_position_embeddings - seq_length = getattr(model.config, "max_sequence_length", None) - if megatron_lm_plugin.seq_length is None: - if seq_length is not None: - megatron_lm_plugin.seq_length = seq_length - elif megatron_lm_plugin.decoder_seq_length is not None: - megatron_lm_plugin.seq_length = megatron_lm_plugin.decoder_seq_length - elif batch_data is not None: - megatron_lm_plugin.seq_length = batch_data["input_ids"].shape[1] - else: - megatron_lm_plugin.seq_length = max_position_embeddings - - megatron_lm_plugin.megatron_lm_default_args["return_logits"] = megatron_lm_plugin.return_logits - megatron_lm_plugin.megatron_lm_default_args["tokenizer_type"] = "Llama2Tokenizer" - megatron_lm_plugin.megatron_lm_default_args["model_type_name"] = model_type_name - megatron_lm_plugin.megatron_lm_default_args["num_layers"] = num_layers - megatron_lm_plugin.megatron_lm_default_args["pretraining_flag"] = pretraining_flag - megatron_lm_plugin.megatron_lm_default_args["hidden_size"] = hidden_size - megatron_lm_plugin.megatron_lm_default_args["num_attention_heads"] = num_attention_heads - megatron_lm_plugin.megatron_lm_default_args["orig_vocab_size"] = orig_vocab_size - megatron_lm_plugin.megatron_lm_default_args["max_position_embeddings"] = max_position_embeddings - megatron_lm_plugin.megatron_lm_default_args["seq_length"] = megatron_lm_plugin.seq_length - megatron_lm_plugin.megatron_lm_default_args["model_return_dict"] = model.config.return_dict - - -@dataclass -class BnbQuantizationConfig: - """ - A plugin to enable BitsAndBytes 4bit and 8bit quantization - - Args: - load_in_8bit (`bool`, defaults to `False`): - Enable 8bit quantization. - llm_int8_threshold (`float`, defaults to `6.0`): - Value of the outliner threshold. Only relevant when `load_in_8bit=True`. - load_in_4_bit (`bool`, defaults to `False`): - Enable 4bit quantization. - bnb_4bit_quant_type (`str`, defaults to `fp4`): - Set the quantization data type in the `bnb.nn.Linear4Bit` layers. Options are {'fp4','np4'}. - bnb_4bit_use_double_quant (`bool`, defaults to `False`): - Enable nested quantization where the quantization constants from the first quantization are quantized - again. - bnb_4bit_compute_dtype (`bool`, defaults to `fp16`): - This sets the computational type which might be different than the input time. For example, inputs might be - fp32, but computation can be set to bf16 for speedups. Options are {'fp32','fp16','bf16'}. - torch_dtype (`torch.dtype`, defaults to `None`): - This sets the dtype of the remaining non quantized layers. `bitsandbytes` library suggests to set the value - to `torch.float16` for 8 bit model and use the same dtype as the compute dtype for 4 bit model. - skip_modules (`List[str]`, defaults to `None`): - An explicit list of the modules that we don't quantize. The dtype of these modules will be `torch_dtype`. - keep_in_fp32_modules (`List`, defaults to `None`): - An explicit list of the modules that we don't quantize. We keep them in `torch.float32`. - """ - - load_in_8bit: bool = field(default=False, metadata={"help": "enable 8bit quantization."}) - - llm_int8_threshold: float = field( - default=6.0, metadata={"help": "value of the outliner threshold. only relevant when load_in_8bit=True"} - ) - - load_in_4bit: bool = field(default=False, metadata={"help": "enable 4bit quantization."}) - - bnb_4bit_quant_type: str = field( - default="fp4", - metadata={ - "help": "set the quantization data type in the `bnb.nn.Linear4Bit` layers. Options are {'fp4','nf4'}." - }, - ) - - bnb_4bit_use_double_quant: bool = field( - default=False, - metadata={ - "help": "enable nested quantization where the quantization constants from the first quantization are quantized again." - }, - ) - - bnb_4bit_compute_dtype: str = field( - default="fp16", - metadata={ - "help": "This sets the computational type which might be different than the input time. For example, inputs might be " - "fp32, but computation can be set to bf16 for speedups. Options are {'fp32','fp16','bf16'}." - }, - ) - - torch_dtype: torch.dtype = field( - default=None, - metadata={ - "help": "this sets the dtype of the remaining non quantized layers. `bitsandbytes` library suggests to set the value" - "to `torch.float16` for 8 bit model and use the same dtype as the compute dtype for 4 bit model " - }, - ) - - skip_modules: list[str] = field( - default=None, - metadata={ - "help": "an explicit list of the modules that we don't quantize. The dtype of these modules will be `torch_dtype`." - }, - ) - - keep_in_fp32_modules: list[str] = field( - default=None, - metadata={"help": "an explicit list of the modules that we don't quantize. We keep them in `torch.float32`."}, - ) - - def __post_init__(self): - """ - Safety checker that arguments are correct - also replaces some NoneType arguments with their default values. - """ - if not isinstance(self.load_in_8bit, bool): - raise ValueError("load_in_8bit must be a boolean") - - if not isinstance(self.load_in_4bit, bool): - raise ValueError("load_in_4bit must be a boolean") - - if self.load_in_4bit and self.load_in_8bit: - raise ValueError("load_in_4bit and load_in_8bit can't be both True") - - if not self.load_in_4bit and not self.load_in_8bit: - raise ValueError("load_in_4bit and load_in_8bit can't be both False") - - if not isinstance(self.llm_int8_threshold, (int, float)): - raise ValueError("llm_int8_threshold must be a float or an int") - - if not isinstance(self.bnb_4bit_quant_type, str): - raise ValueError("bnb_4bit_quant_type must be a string") - elif self.bnb_4bit_quant_type not in ["fp4", "nf4"]: - raise ValueError(f"bnb_4bit_quant_type must be in ['fp4','nf4'] but found {self.bnb_4bit_quant_type}") - - if not isinstance(self.bnb_4bit_use_double_quant, bool): - raise ValueError("bnb_4bit_use_double_quant must be a boolean") - - if isinstance(self.bnb_4bit_compute_dtype, str): - if self.bnb_4bit_compute_dtype == "fp32": - self.bnb_4bit_compute_dtype = torch.float32 - elif self.bnb_4bit_compute_dtype == "fp16": - self.bnb_4bit_compute_dtype = torch.float16 - elif self.bnb_4bit_compute_dtype == "bf16": - self.bnb_4bit_compute_dtype = torch.bfloat16 - else: - raise ValueError( - f"bnb_4bit_compute_dtype must be in ['fp32','fp16','bf16'] but found {self.bnb_4bit_compute_dtype}" - ) - elif not isinstance(self.bnb_4bit_compute_dtype, torch.dtype): - raise ValueError("bnb_4bit_compute_dtype must be a string or a torch.dtype") - - if self.skip_modules is not None and not isinstance(self.skip_modules, list): - raise ValueError("skip_modules must be a list of strings") - - if self.keep_in_fp32_modules is not None and not isinstance(self.keep_in_fp32_modules, list): - raise ValueError("keep_in_fp_32_modules must be a list of strings") - - if self.load_in_4bit: - self.target_dtype = CustomDtype.INT4 - - if self.load_in_8bit: - self.target_dtype = torch.int8 - - if self.load_in_4bit and self.llm_int8_threshold != 6.0: - warnings.warn("llm_int8_threshold can only be used for model loaded in 8bit") - - if isinstance(self.torch_dtype, str): - if self.torch_dtype == "fp32": - self.torch_dtype = torch.float32 - elif self.torch_dtype == "fp16": - self.torch_dtype = torch.float16 - elif self.torch_dtype == "bf16": - self.torch_dtype = torch.bfloat16 - else: - raise ValueError(f"torch_dtype must be in ['fp32','fp16','bf16'] but found {self.torch_dtype}") - if self.load_in_8bit and self.torch_dtype is None: - self.torch_dtype = torch.float16 - - if self.load_in_4bit and self.torch_dtype is None: - self.torch_dtype = self.bnb_4bit_compute_dtype - - if not isinstance(self.torch_dtype, torch.dtype): - raise ValueError("torch_dtype must be a torch.dtype") - - -def get_module_class_from_name(module, name): - """ - Gets a class from a module by its name. - - Args: - module (`torch.nn.Module`): The module to get the class from. - name (`str`): The name of the class. - """ - modules_children = list(module.children()) - if module.__class__.__name__ == name: - return module.__class__ - elif len(modules_children) == 0: - return - else: - for child_module in modules_children: - module_class = get_module_class_from_name(child_module, name) - if module_class is not None: - return module_class diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/deepspeed.py b/venv/lib/python3.10/site-packages/accelerate/utils/deepspeed.py deleted file mode 100644 index 22db891c63d9bd48691acd87a15a206c270017a9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/deepspeed.py +++ /dev/null @@ -1,385 +0,0 @@ -# Copyright 2021 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 base64 -import json -import os -from copy import deepcopy - -from torch import optim - -from ..optimizer import AcceleratedOptimizer -from ..scheduler import AcceleratedScheduler -from .dataclasses import DistributedType -from .imports import is_bnb_available -from .versions import compare_versions - - -def map_pytorch_optim_to_deepspeed(optimizer): - """ - Args: - optimizer: torch.optim.Optimizer - - Returns the DeepSeedCPUOptimizer (deepspeed.ops) version of the optimizer. - """ - - defaults = {k: v for k, v in optimizer.defaults.items() if k in ["lr", "weight_decay"]} - - # Select the DeepSpeedCPUOptimizer based on the original optimizer class. - # DeepSpeedCPUAdam is the default - from deepspeed.ops.adam import DeepSpeedCPUAdam - - optimizer_class = DeepSpeedCPUAdam - - # For DeepSpeedCPUAdam (adamw_mode) - if compare_versions("deepspeed", ">=", "0.3.1"): - defaults["adamw_mode"] = False - is_adaw = isinstance(optimizer, optim.AdamW) - - if is_bnb_available() and not is_adaw: - import bitsandbytes.optim as bnb_opt - - if isinstance(optimizer, (bnb_opt.AdamW, bnb_opt.AdamW32bit)): - try: - is_adaw = optimizer.optim_bits == 32 - except AttributeError: - is_adaw = optimizer.args.optim_bits == 32 - else: - is_adaw = False - - if is_adaw: - defaults["adamw_mode"] = True - - # For DeepSpeedCPUAdagrad - if compare_versions("deepspeed", ">=", "0.5.5"): - # Check if the optimizer is PyTorch's Adagrad. - is_ada = isinstance(optimizer, optim.Adagrad) - # If not, and bitsandbytes is available, - # # check if the optimizer is the 32-bit bitsandbytes Adagrad. - if is_bnb_available() and not is_ada: - import bitsandbytes.optim as bnb_opt - - if isinstance(optimizer, (bnb_opt.Adagrad, bnb_opt.Adagrad32bit)): - try: - is_ada = optimizer.optim_bits == 32 - except AttributeError: - is_ada = optimizer.args.optim_bits == 32 - if is_ada: - from deepspeed.ops.adagrad import DeepSpeedCPUAdagrad - - optimizer_class = DeepSpeedCPUAdagrad - - # For DeepSpeedCPULion - if is_bnb_available(min_version="0.38.0") and compare_versions("deepspeed", ">=", "0.11.0"): - from bitsandbytes.optim import Lion, Lion32bit - - if isinstance(optimizer, (Lion, Lion32bit)): - try: - is_bnb_32bits = optimizer.optim_bits == 32 - except AttributeError: - is_bnb_32bits = optimizer.args.optim_bits == 32 - if is_bnb_32bits: - from deepspeed.ops.lion import DeepSpeedCPULion - - optimizer_class = DeepSpeedCPULion - - return optimizer_class(optimizer.param_groups, **defaults) - - -def get_active_deepspeed_plugin(state): - """ - Returns the currently active DeepSpeedPlugin. - - Raises: - ValueError: If DeepSpeed was not enabled and this function is called. - """ - if state.distributed_type != DistributedType.DEEPSPEED: - raise ValueError( - "Couldn't retrieve the active `DeepSpeedPlugin` as none were enabled. " - "Please make sure that either `Accelerator` is configured for `deepspeed` " - "or make sure that the desired `DeepSpeedPlugin` has been enabled (`AcceleratorState().select_deepspeed_plugin(name)`) " - "before calling this function." - ) - if not isinstance(state.deepspeed_plugins, dict): - return state.deepspeed_plugins - return next(plugin for plugin in state.deepspeed_plugins.values() if plugin.selected) - - -class HfDeepSpeedConfig: - """ - This object contains a DeepSpeed configuration dictionary and can be quickly queried for things like zero stage. - - A `weakref` of this object is stored in the module's globals to be able to access the config from areas where - things like the Trainer object is not available (e.g. `from_pretrained` and `_get_resized_embeddings`). Therefore - it's important that this object remains alive while the program is still running. - - [`Trainer`] uses the `HfTrainerDeepSpeedConfig` subclass instead. That subclass has logic to sync the configuration - with values of [`TrainingArguments`] by replacing special placeholder values: `"auto"`. Without this special logic - the DeepSpeed configuration is not modified in any way. - - Args: - config_file_or_dict (`Union[str, Dict]`): path to DeepSpeed config file or dict. - - """ - - def __init__(self, config_file_or_dict): - if isinstance(config_file_or_dict, dict): - # Don't modify user's data should they want to reuse it (e.g. in tests), because once we - # modified it, it will not be accepted here again, since `auto` values would have been overridden - config = deepcopy(config_file_or_dict) - elif os.path.exists(config_file_or_dict): - with open(config_file_or_dict, encoding="utf-8") as f: - config = json.load(f) - else: - try: - try: - # First try parsing as JSON directly - config = json.loads(config_file_or_dict) - except json.JSONDecodeError: - # If that fails, try base64 decoding - config_decoded = base64.urlsafe_b64decode(config_file_or_dict).decode("utf-8") - config = json.loads(config_decoded) - except (UnicodeDecodeError, AttributeError, ValueError): - raise ValueError( - f"Expected a string path to an existing deepspeed config, or a dictionary, or a base64 encoded string. Received: {config_file_or_dict}" - ) - - self.config = config - - self.set_stage_and_offload() - - def set_stage_and_offload(self): - # zero stage - this is done as early as possible, before model is created, to allow - # ``is_deepspeed_zero3_enabled`` query and getting to the early deepspeed config object - # during ``zero.Init()`` which needs to know the dtype, and some other hparams. - self._stage = self.get_value("zero_optimization.stage", -1) - - # offload - self._offload = False - if self.is_zero2() or self.is_zero3(): - offload_devices_valid = set(["cpu", "nvme"]) - offload_devices = set( - [ - self.get_value("zero_optimization.offload_optimizer.device"), - self.get_value("zero_optimization.offload_param.device"), - ] - ) - if len(offload_devices & offload_devices_valid) > 0: - self._offload = True - - def find_config_node(self, ds_key_long): - config = self.config - - # find the config node of interest if it exists - nodes = ds_key_long.split(".") - ds_key = nodes.pop() - for node in nodes: - config = config.get(node) - if config is None: - return None, ds_key - - return config, ds_key - - def get_value(self, ds_key_long, default=None): - """ - Returns the set value or `default` if no value is set - """ - config, ds_key = self.find_config_node(ds_key_long) - if config is None: - return default - return config.get(ds_key, default) - - def del_config_sub_tree(self, ds_key_long, must_exist=False): - """ - Deletes a sub-section of the config file if it's found. - - Unless `must_exist` is `True` the section doesn't have to exist. - """ - config = self.config - - # find the config node of interest if it exists - nodes = ds_key_long.split(".") - for node in nodes: - parent_config = config - config = config.get(node) - if config is None: - if must_exist: - raise ValueError(f"Can't find {ds_key_long} entry in the config: {self.config}") - else: - return - - # if found remove it - if parent_config is not None: - parent_config.pop(node) - - def is_true(self, ds_key_long): - """ - Returns `True`/``False` only if the value is set, always `False` otherwise. So use this method to ask the very - specific question of whether the value is set to `True` (and it's not set to `False`` or isn't set). - - """ - value = self.get_value(ds_key_long) - return False if value is None else bool(value) - - def is_false(self, ds_key_long): - """ - Returns `True`/``False` only if the value is set, always `False` otherwise. So use this method to ask the very - specific question of whether the value is set to `False` (and it's not set to `True`` or isn't set). - """ - value = self.get_value(ds_key_long) - return False if value is None else not bool(value) - - def is_zero2(self): - return self._stage == 2 - - def is_zero3(self): - return self._stage == 3 - - def is_offload(self): - return self._offload - - -class DeepSpeedEngineWrapper: - """ - Internal wrapper for deepspeed.runtime.engine.DeepSpeedEngine. This is used to follow conventional training loop. - - Args: - engine (deepspeed.runtime.engine.DeepSpeedEngine): deepspeed engine to wrap - """ - - def __init__(self, engine): - self.engine = engine - - def backward(self, loss, sync_gradients=True, **kwargs): - # Set gradient accumulation boundary based on Accelerate's sync_gradients state - # This tells DeepSpeed whether this is the final micro-batch before gradient sync - self.engine.set_gradient_accumulation_boundary(is_boundary=sync_gradients) - - # runs backpropagation and handles mixed precision - self.engine.backward(loss, **kwargs) - - # Only perform step and related operations at gradient accumulation boundaries - if sync_gradients: - # Deepspeed's `engine.step` performs the following operations: - # - gradient accumulation check - # - gradient clipping - # - optimizer step - # - zero grad - # - checking overflow - # - lr_scheduler step (only if engine.lr_scheduler is not None) - self.engine.step() - # and this plugin overrides the above calls with no-ops when Accelerate runs under - # Deepspeed, but allows normal functionality for non-Deepspeed cases thus enabling a simple - # training loop that works transparently under many training regimes. - - def get_global_grad_norm(self): - """Get the global gradient norm from DeepSpeed engine.""" - grad_norm = self.engine.get_global_grad_norm() - # Convert to scalar if it's a tensor - if hasattr(grad_norm, "item"): - return grad_norm.item() - return grad_norm - - -class DeepSpeedOptimizerWrapper(AcceleratedOptimizer): - """ - Internal wrapper around a deepspeed optimizer. - - Args: - optimizer (`torch.optim.optimizer.Optimizer`): - The optimizer to wrap. - """ - - def __init__(self, optimizer): - super().__init__(optimizer, device_placement=False, scaler=None) - self.__has_overflow__ = hasattr(self.optimizer, "overflow") - - def zero_grad(self, set_to_none=None): - pass # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed - - def step(self): - pass # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed - - @property - def step_was_skipped(self): - """Whether or not the optimizer step was done, or skipped because of gradient overflow.""" - if self.__has_overflow__: - return self.optimizer.overflow - return False - - -class DeepSpeedSchedulerWrapper(AcceleratedScheduler): - """ - Internal wrapper around a deepspeed scheduler. - - Args: - scheduler (`torch.optim.lr_scheduler.LambdaLR`): - The scheduler to wrap. - optimizers (one or a list of `torch.optim.Optimizer`): - """ - - def __init__(self, scheduler, optimizers): - super().__init__(scheduler, optimizers) - - def step(self): - pass # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed - - -class DummyOptim: - """ - Dummy optimizer presents model parameters or param groups, this is primarily used to follow conventional training - loop when optimizer config is specified in the deepspeed config file. - - Args: - lr (float): - Learning rate. - params (iterable): iterable of parameters to optimize or dicts defining - parameter groups - weight_decay (float): - Weight decay. - **kwargs (additional keyword arguments, *optional*): - Other arguments. - """ - - def __init__(self, params, lr=0.001, weight_decay=0, **kwargs): - self.params = params - self.lr = lr - self.weight_decay = weight_decay - self.kwargs = kwargs - - -class DummyScheduler: - """ - Dummy scheduler presents model parameters or param groups, this is primarily used to follow conventional training - loop when scheduler config is specified in the deepspeed config file. - - Args: - optimizer (`torch.optim.optimizer.Optimizer`): - The optimizer to wrap. - total_num_steps (int, *optional*): - Total number of steps. - warmup_num_steps (int, *optional*): - Number of steps for warmup. - lr_scheduler_callable (callable, *optional*): - A callable function that creates an LR Scheduler. It accepts only one argument `optimizer`. - **kwargs (additional keyword arguments, *optional*): - Other arguments. - """ - - def __init__(self, optimizer, total_num_steps=None, warmup_num_steps=0, lr_scheduler_callable=None, **kwargs): - self.optimizer = optimizer - self.total_num_steps = total_num_steps - self.warmup_num_steps = warmup_num_steps - self.lr_scheduler_callable = lr_scheduler_callable - self.kwargs = kwargs diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/environment.py b/venv/lib/python3.10/site-packages/accelerate/utils/environment.py deleted file mode 100644 index c913702ef30290ed3d19c08a71213663297f32de..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/environment.py +++ /dev/null @@ -1,421 +0,0 @@ -# Copyright 2022 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 logging -import math -import os -import platform -import subprocess -import sys -from contextlib import contextmanager -from dataclasses import dataclass, field -from functools import lru_cache, wraps -from shutil import which -from typing import Optional, Union - -import torch -from packaging.version import parse - - -logger = logging.getLogger(__name__) - - -def convert_dict_to_env_variables(current_env: dict): - """ - Verifies that all keys and values in `current_env` do not contain illegal keys or values, and returns a list of - strings as the result. - - Example: - ```python - >>> from accelerate.utils.environment import verify_env - - >>> env = {"ACCELERATE_DEBUG_MODE": "1", "BAD_ENV_NAME": ">> valid_env_items = verify_env(env) - >>> print(valid_env_items) - ["ACCELERATE_DEBUG_MODE=1\n", "OTHER_ENV=2\n"] - ``` - """ - forbidden_chars = [";", "\n", "<", ">", " "] - valid_env_items = [] - for key, value in current_env.items(): - if all(char not in (key + value) for char in forbidden_chars) and len(key) >= 1 and len(value) >= 1: - valid_env_items.append(f"{key}={value}\n") - else: - logger.warning(f"WARNING: Skipping {key}={value} as it contains forbidden characters or missing values.") - return valid_env_items - - -def str_to_bool(value, to_bool: bool = False) -> Union[int, bool]: - """ - Converts a string representation of truth to `True` (1) or `False` (0). - - True values are `y`, `yes`, `t`, `true`, `on`, and `1`; False value are `n`, `no`, `f`, `false`, `off`, and `0`; - """ - value = value.lower() - if value in ("y", "yes", "t", "true", "on", "1"): - return 1 if not to_bool else True - elif value in ("n", "no", "f", "false", "off", "0"): - return 0 if not to_bool else False - else: - raise ValueError(f"invalid truth value {value}") - - -def get_int_from_env(env_keys, default): - """Returns the first positive env value found in the `env_keys` list or the default.""" - for e in env_keys: - val = int(os.environ.get(e, -1)) - if val >= 0: - return val - return default - - -def parse_flag_from_env(key, default=False): - """Returns truthy value for `key` from the env if available else the default.""" - value = os.environ.get(key, str(default)) - return str_to_bool(value) == 1 # As its name indicates `str_to_bool` actually returns an int... - - -def parse_choice_from_env(key, default="no"): - value = os.environ.get(key, str(default)) - return value - - -def are_libraries_initialized(*library_names: str) -> list[str]: - """ - Checks if any of `library_names` are imported in the environment. Will return any names that are. - """ - return [lib_name for lib_name in library_names if lib_name in sys.modules.keys()] - - -def _nvidia_smi(): - """ - Returns the right nvidia-smi command based on the system. - """ - if platform.system() == "Windows": - # If platform is Windows and nvidia-smi can't be found in path - # try from systemd drive with default installation path - command = which("nvidia-smi") - if command is None: - command = f"{os.environ['systemdrive']}\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe" - else: - command = "nvidia-smi" - return command - - -def get_gpu_info(): - """ - Gets GPU count and names using `nvidia-smi` instead of torch to not initialize CUDA. - - Largely based on the `gputil` library. - """ - # Returns as list of `n` GPUs and their names - output = subprocess.check_output( - [_nvidia_smi(), "--query-gpu=count,name", "--format=csv,noheader"], universal_newlines=True - ) - output = output.strip() - gpus = output.split(os.linesep) - # Get names from output - gpu_count = len(gpus) - gpu_names = [gpu.split(",")[1].strip() for gpu in gpus] - return gpu_names, gpu_count - - -def get_driver_version(): - """ - Returns the driver version - - In the case of multiple GPUs, will return the first. - """ - output = subprocess.check_output( - [_nvidia_smi(), "--query-gpu=driver_version", "--format=csv,noheader"], universal_newlines=True - ) - output = output.strip() - return output.split(os.linesep)[0] - - -def check_cuda_p2p_ib_support(): - """ - Checks if the devices being used have issues with P2P and IB communications, namely any consumer GPU hardware after - the 3090. - - Noteably uses `nvidia-smi` instead of torch to not initialize CUDA. - """ - try: - device_names, device_count = get_gpu_info() - # As new consumer GPUs get released, add them to `unsupported_devices`` - unsupported_devices = {"RTX 40"} - if device_count > 1: - if any( - unsupported_device in device_name - for device_name in device_names - for unsupported_device in unsupported_devices - ): - # Check if they have the right driver version - acceptable_driver_version = "550.40.07" - current_driver_version = get_driver_version() - if parse(current_driver_version) < parse(acceptable_driver_version): - return False - return True - except Exception: - pass - return True - - -@lru_cache -def check_cuda_fp8_capability(): - """ - Checks if the current GPU available supports FP8. - - Notably might initialize `torch.cuda` to check. - """ - - try: - # try to get the compute capability from nvidia-smi - output = subprocess.check_output( - [_nvidia_smi(), "--query-gpu=compute_capability", "--format=csv,noheader"], universal_newlines=True - ) - output = output.strip() - # we take the first GPU's compute capability - compute_capability = tuple(map(int, output.split(os.linesep)[0].split("."))) - except Exception: - compute_capability = torch.cuda.get_device_capability() - - return compute_capability >= (8, 9) - - -@dataclass -class CPUInformation: - """ - Stores information about the CPU in a distributed environment. It contains the following attributes: - - rank: The rank of the current process. - - world_size: The total number of processes in the world. - - local_rank: The rank of the current process on the local node. - - local_world_size: The total number of processes on the local node. - """ - - rank: int = field(default=0, metadata={"help": "The rank of the current process."}) - world_size: int = field(default=1, metadata={"help": "The total number of processes in the world."}) - local_rank: int = field(default=0, metadata={"help": "The rank of the current process on the local node."}) - local_world_size: int = field(default=1, metadata={"help": "The total number of processes on the local node."}) - - -def get_cpu_distributed_information() -> CPUInformation: - """ - Returns various information about the environment in relation to CPU distributed training as a `CPUInformation` - dataclass. - """ - information = {} - information["rank"] = get_int_from_env(["RANK", "PMI_RANK", "OMPI_COMM_WORLD_RANK", "MV2_COMM_WORLD_RANK"], 0) - information["world_size"] = get_int_from_env( - ["WORLD_SIZE", "PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE"], 1 - ) - information["local_rank"] = get_int_from_env( - ["LOCAL_RANK", "MPI_LOCALRANKID", "OMPI_COMM_WORLD_LOCAL_RANK", "MV2_COMM_WORLD_LOCAL_RANK"], 0 - ) - information["local_world_size"] = get_int_from_env( - ["LOCAL_WORLD_SIZE", "MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], - 1, - ) - return CPUInformation(**information) - - -def override_numa_affinity(local_process_index: int, verbose: Optional[bool] = None) -> None: - """ - Overrides whatever NUMA affinity is set for the current process. This is very taxing and requires recalculating the - affinity to set, ideally you should use `utils.environment.set_numa_affinity` instead. - - Args: - local_process_index (int): - The index of the current process on the current server. - verbose (bool, *optional*): - Whether to log out the assignment of each CPU. If `ACCELERATE_DEBUG_MODE` is enabled, will default to True. - """ - if verbose is None: - verbose = parse_flag_from_env("ACCELERATE_DEBUG_MODE", False) - if torch.cuda.is_available(): - from accelerate.utils import is_pynvml_available - - if not is_pynvml_available(): - raise ImportError( - "To set CPU affinity on CUDA GPUs the `pynvml` package must be available. (`pip install pynvml`)" - ) - import pynvml as nvml - - # The below code is based on https://github.com/NVIDIA/DeepLearningExamples/blob/master/TensorFlow2/LanguageModeling/BERT/gpu_affinity.py - nvml.nvmlInit() - num_elements = math.ceil(os.cpu_count() / 64) - handle = nvml.nvmlDeviceGetHandleByIndex(local_process_index) - affinity_string = "" - for j in nvml.nvmlDeviceGetCpuAffinity(handle, num_elements): - # assume nvml returns list of 64 bit ints - affinity_string = f"{j:064b}{affinity_string}" - affinity_list = [int(x) for x in affinity_string] - affinity_list.reverse() # so core 0 is the 0th element - affinity_to_set = [i for i, e in enumerate(affinity_list) if e != 0] - os.sched_setaffinity(0, affinity_to_set) - if verbose: - cpu_cores = os.sched_getaffinity(0) - logger.info(f"Assigning {len(cpu_cores)} cpu cores to process {local_process_index}: {cpu_cores}") - - -@lru_cache -def set_numa_affinity(local_process_index: int, verbose: Optional[bool] = None) -> None: - """ - Assigns the current process to a specific NUMA node. Ideally most efficient when having at least 2 cpus per node. - - This result is cached between calls. If you want to override it, please use - `accelerate.utils.environment.override_numa_afifnity`. - - Args: - local_process_index (int): - The index of the current process on the current server. - verbose (bool, *optional*): - Whether to print the new cpu cores assignment for each process. If `ACCELERATE_DEBUG_MODE` is enabled, will - default to True. - """ - override_numa_affinity(local_process_index=local_process_index, verbose=verbose) - - -@contextmanager -def clear_environment(): - """ - A context manager that will temporarily clear environment variables. - - When this context exits, the previous environment variables will be back. - - Example: - - ```python - >>> import os - >>> from accelerate.utils import clear_environment - - >>> os.environ["FOO"] = "bar" - >>> with clear_environment(): - ... print(os.environ) - ... os.environ["FOO"] = "new_bar" - ... print(os.environ["FOO"]) - {} - new_bar - - >>> print(os.environ["FOO"]) - bar - ``` - """ - _old_os_environ = os.environ.copy() - os.environ.clear() - - try: - yield - finally: - os.environ.clear() # clear any added keys, - os.environ.update(_old_os_environ) # then restore previous environment - - -@contextmanager -def patch_environment(**kwargs): - """ - A context manager that will add each keyword argument passed to `os.environ` and remove them when exiting. - - Will convert the values in `kwargs` to strings and upper-case all the keys. - - Example: - - ```python - >>> import os - >>> from accelerate.utils import patch_environment - - >>> with patch_environment(FOO="bar"): - ... print(os.environ["FOO"]) # prints "bar" - >>> print(os.environ["FOO"]) # raises KeyError - ``` - """ - existing_vars = {} - for key, value in kwargs.items(): - key = key.upper() - if key in os.environ: - existing_vars[key] = os.environ[key] - os.environ[key] = str(value) - - try: - yield - finally: - for key in kwargs: - key = key.upper() - if key in existing_vars: - # restore previous value - os.environ[key] = existing_vars[key] - else: - os.environ.pop(key, None) - - -def purge_accelerate_environment(func_or_cls): - """Decorator to clean up accelerate environment variables set by the decorated class or function. - - In some circumstances, calling certain classes or functions can result in accelerate env vars being set and not - being cleaned up afterwards. As an example, when calling: - - TrainingArguments(fp16=True, ...) - - The following env var will be set: - - ACCELERATE_MIXED_PRECISION=fp16 - - This can affect subsequent code, since the env var takes precedence over TrainingArguments(fp16=False). This is - especially relevant for unit testing, where we want to avoid the individual tests to have side effects on one - another. Decorate the unit test function or whole class with this decorator to ensure that after each test, the env - vars are cleaned up. This works for both unittest.TestCase and normal classes (pytest); it also works when - decorating the parent class. - - """ - prefix = "ACCELERATE_" - - @contextmanager - def env_var_context(): - # Store existing accelerate env vars - existing_vars = {k: v for k, v in os.environ.items() if k.startswith(prefix)} - try: - yield - finally: - # Restore original env vars or remove new ones - for key in [k for k in os.environ if k.startswith(prefix)]: - if key in existing_vars: - os.environ[key] = existing_vars[key] - else: - os.environ.pop(key, None) - - def wrap_function(func): - @wraps(func) - def wrapper(*args, **kwargs): - with env_var_context(): - return func(*args, **kwargs) - - wrapper._accelerate_is_purged_environment_wrapped = True - return wrapper - - if not isinstance(func_or_cls, type): - return wrap_function(func_or_cls) - - # Handle classes by wrapping test methods - def wrap_test_methods(test_class_instance): - for name in dir(test_class_instance): - if name.startswith("test"): - method = getattr(test_class_instance, name) - if callable(method) and not hasattr(method, "_accelerate_is_purged_environment_wrapped"): - setattr(test_class_instance, name, wrap_function(method)) - return test_class_instance - - # Handle inheritance - wrap_test_methods(func_or_cls) - func_or_cls.__init_subclass__ = classmethod(lambda cls, **kw: wrap_test_methods(cls)) - return func_or_cls diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/fsdp_utils.py b/venv/lib/python3.10/site-packages/accelerate/utils/fsdp_utils.py deleted file mode 100644 index a4526a26bf4b2924e0820bd192d3c1d47a72b001..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/fsdp_utils.py +++ /dev/null @@ -1,786 +0,0 @@ -# Copyright 2023 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 functools -import os -import shutil -import warnings -from collections import defaultdict -from contextlib import nullcontext -from pathlib import Path -from typing import Callable - -import torch - -from ..logging import get_logger -from .constants import FSDP_MODEL_NAME, OPTIMIZER_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_NAME -from .dataclasses import get_module_class_from_name -from .modeling import get_non_persistent_buffers, is_peft_model -from .other import get_module_children_bottom_up, is_compiled_module, save -from .versions import is_torch_version - - -logger = get_logger(__name__) - - -def enable_fsdp_ram_efficient_loading(): - """ - Enables RAM efficient loading of Hugging Face models for FSDP in the environment. - """ - # Sets values for `transformers.modeling_utils.is_fsdp_enabled` - if "ACCELERATE_USE_FSDP" not in os.environ: - os.environ["ACCELERATE_USE_FSDP"] = "True" - os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] = "True" - - -def disable_fsdp_ram_efficient_loading(): - """ - Disables RAM efficient loading of Hugging Face models for FSDP in the environment. - """ - os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] = "False" - - -def _get_model_state_dict(model, adapter_only=False, sd_options=None): - if adapter_only and is_peft_model(model): - from peft import get_peft_model_state_dict - - return get_peft_model_state_dict(model, adapter_name=model.active_adapter) - - # Invariant: `sd_options` is not None only for FSDP2 - if sd_options is not None: - from torch.distributed.checkpoint.state_dict import get_model_state_dict - - return get_model_state_dict(model, options=sd_options) - else: - return model.state_dict() - - -def _set_model_state_dict(model, state_dict, adapter_only=False, sd_options=None): - if adapter_only and is_peft_model(model): - from peft import set_peft_model_state_dict - - return set_peft_model_state_dict(model, state_dict, adapter_name=model.active_adapter) - - # Invariant: `sd_options` is not None only for FSDP2 - if sd_options is not None: - from torch.distributed.checkpoint.state_dict import set_model_state_dict - - return set_model_state_dict(model, state_dict, options=sd_options) - else: - return model.load_state_dict(state_dict) - - -def _prepare_sd_options(fsdp_plugin): - sd_options = None - - # we use this only for FSDP2, as it requires torch >= 2.6.0 and this api requires torch >= 2.2.0 - if fsdp_plugin.fsdp_version == 2: - from torch.distributed.checkpoint.state_dict import StateDictOptions - from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType - - sd_options = StateDictOptions( - full_state_dict=fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT, - cpu_offload=getattr(fsdp_plugin.state_dict_config, "offload_to_cpu", False), - broadcast_from_rank0=getattr(fsdp_plugin.state_dict_config, "rank0_only", False), - ) - - return sd_options - - -def save_fsdp_model(fsdp_plugin, accelerator, model, output_dir, model_index=0, adapter_only=False): - # Note: We import here to reduce import time from general modules, and isolate outside dependencies - import torch.distributed.checkpoint as dist_cp - from torch.distributed.checkpoint.default_planner import DefaultSavePlanner - from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP - from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType - - os.makedirs(output_dir, exist_ok=True) - if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT: - # FSDP raises error when single GPU is used with `offload_to_cpu=True` for FULL_STATE_DICT - # so, only enable it when num_processes>1 - is_multi_process = accelerator.num_processes > 1 - fsdp_plugin.state_dict_config.offload_to_cpu = is_multi_process - fsdp_plugin.state_dict_config.rank0_only = is_multi_process - - ctx = ( - FSDP.state_dict_type( - model, fsdp_plugin.state_dict_type, fsdp_plugin.state_dict_config, fsdp_plugin.optim_state_dict_config - ) - if fsdp_plugin.fsdp_version == 1 - else nullcontext() - ) - sd_options = _prepare_sd_options(fsdp_plugin) - - with ctx: - state_dict = _get_model_state_dict(model, adapter_only=adapter_only, sd_options=sd_options) - if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT: - weights_name = f"{FSDP_MODEL_NAME}.bin" if model_index == 0 else f"{FSDP_MODEL_NAME}_{model_index}.bin" - output_model_file = os.path.join(output_dir, weights_name) - if accelerator.process_index == 0: - logger.info(f"Saving model to {output_model_file}") - torch.save(state_dict, output_model_file) - logger.info(f"Model saved to {output_model_file}") - # Invariant: `LOCAL_STATE_DICT` is never possible with `FSDP2` - elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT: - weights_name = ( - f"{FSDP_MODEL_NAME}_rank{accelerator.process_index}.bin" - if model_index == 0 - else f"{FSDP_MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin" - ) - output_model_file = os.path.join(output_dir, weights_name) - logger.info(f"Saving model to {output_model_file}") - torch.save(state_dict, output_model_file) - logger.info(f"Model saved to {output_model_file}") - elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT: - ckpt_dir = os.path.join(output_dir, f"{FSDP_MODEL_NAME}_{model_index}") - os.makedirs(ckpt_dir, exist_ok=True) - logger.info(f"Saving model to {ckpt_dir}") - state_dict = {"model": state_dict} - - dist_cp.save( - state_dict=state_dict, - storage_writer=dist_cp.FileSystemWriter(ckpt_dir), - planner=DefaultSavePlanner(), - ) - logger.info(f"Model saved to {ckpt_dir}") - - -def load_fsdp_model(fsdp_plugin, accelerator, model, input_dir, model_index=0, adapter_only=False): - # Note: We import here to reduce import time from general modules, and isolate outside dependencies - import torch.distributed.checkpoint as dist_cp - from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner - from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP - from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType - - accelerator.wait_for_everyone() - if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT: - # FSDP raises error when single GPU is used with `offload_to_cpu=True` for FULL_STATE_DICT - # so, only enable it when num_processes>1 - is_multi_process = accelerator.num_processes > 1 - fsdp_plugin.state_dict_config.offload_to_cpu = is_multi_process - fsdp_plugin.state_dict_config.rank0_only = is_multi_process - - ctx = ( - FSDP.state_dict_type( - model, fsdp_plugin.state_dict_type, fsdp_plugin.state_dict_config, fsdp_plugin.optim_state_dict_config - ) - if fsdp_plugin.fsdp_version == 1 - else nullcontext() - ) - sd_options = _prepare_sd_options(fsdp_plugin) - with ctx: - if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT: - if type(model) is not FSDP and accelerator.process_index != 0 and not accelerator.is_fsdp2: - if not fsdp_plugin.sync_module_states and fsdp_plugin.fsdp_version == 1: - raise ValueError( - "Set the `sync_module_states` flag to `True` so that model states are synced across processes when " - "initializing FSDP object" - ) - return - weights_name = f"{FSDP_MODEL_NAME}.bin" if model_index == 0 else f"{FSDP_MODEL_NAME}_{model_index}.bin" - input_model_file = os.path.join(input_dir, weights_name) - logger.info(f"Loading model from {input_model_file}") - # we want an empty state dict for FSDP2 as we use `broadcast_from_rank0` - load_model = not accelerator.is_fsdp2 or accelerator.is_main_process - if load_model: - state_dict = torch.load(input_model_file, weights_only=True) - else: - state_dict = {} - logger.info(f"Model loaded from {input_model_file}") - elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT: - weights_name = ( - f"{FSDP_MODEL_NAME}_rank{accelerator.process_index}.bin" - if model_index == 0 - else f"{FSDP_MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin" - ) - input_model_file = os.path.join(input_dir, weights_name) - logger.info(f"Loading model from {input_model_file}") - state_dict = torch.load(input_model_file, weights_only=True) - logger.info(f"Model loaded from {input_model_file}") - elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT: - ckpt_dir = ( - os.path.join(input_dir, f"{FSDP_MODEL_NAME}_{model_index}") - if f"{FSDP_MODEL_NAME}" not in input_dir - else input_dir - ) - logger.info(f"Loading model from {ckpt_dir}") - state_dict = {"model": _get_model_state_dict(model, adapter_only=adapter_only, sd_options=sd_options)} - dist_cp.load( - state_dict=state_dict, - storage_reader=dist_cp.FileSystemReader(ckpt_dir), - planner=DefaultLoadPlanner(), - ) - state_dict = state_dict["model"] - logger.info(f"Model loaded from {ckpt_dir}") - - load_result = _set_model_state_dict(model, state_dict, adapter_only=adapter_only, sd_options=sd_options) - return load_result - - -def save_fsdp_optimizer(fsdp_plugin, accelerator, optimizer, model, output_dir, optimizer_index=0): - # Note: We import here to reduce import time from general modules, and isolate outside dependencies - import torch.distributed.checkpoint as dist_cp - from torch.distributed.checkpoint.default_planner import DefaultSavePlanner - from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP - from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType - - os.makedirs(output_dir, exist_ok=True) - - ctx = ( - FSDP.state_dict_type( - model, fsdp_plugin.state_dict_type, fsdp_plugin.state_dict_config, fsdp_plugin.optim_state_dict_config - ) - if fsdp_plugin.fsdp_version == 1 - else nullcontext() - ) - - sd_options = _prepare_sd_options(fsdp_plugin) - - with ctx: - if fsdp_plugin.fsdp_version == 2: - from torch.distributed.checkpoint.state_dict import get_optimizer_state_dict - - optim_state = get_optimizer_state_dict(model, optimizer, options=sd_options) - else: - optim_state = FSDP.optim_state_dict(model, optimizer) - - if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT: - if accelerator.process_index == 0: - optim_state_name = ( - f"{OPTIMIZER_NAME}.bin" if optimizer_index == 0 else f"{OPTIMIZER_NAME}_{optimizer_index}.bin" - ) - output_optimizer_file = os.path.join(output_dir, optim_state_name) - logger.info(f"Saving Optimizer state to {output_optimizer_file}") - torch.save(optim_state, output_optimizer_file) - logger.info(f"Optimizer state saved in {output_optimizer_file}") - else: - ckpt_dir = os.path.join(output_dir, f"{OPTIMIZER_NAME}_{optimizer_index}") - os.makedirs(ckpt_dir, exist_ok=True) - logger.info(f"Saving Optimizer state to {ckpt_dir}") - dist_cp.save( - state_dict={"optimizer": optim_state}, - storage_writer=dist_cp.FileSystemWriter(ckpt_dir), - planner=DefaultSavePlanner(), - ) - logger.info(f"Optimizer state saved in {ckpt_dir}") - - -def load_fsdp_optimizer(fsdp_plugin, accelerator, optimizer, model, input_dir, optimizer_index=0, adapter_only=False): - # Note: We import here to reduce import time from general modules, and isolate outside dependencies - import torch.distributed.checkpoint as dist_cp - from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP - from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType - - accelerator.wait_for_everyone() - ctx = ( - FSDP.state_dict_type( - model, fsdp_plugin.state_dict_type, fsdp_plugin.state_dict_config, fsdp_plugin.optim_state_dict_config - ) - if fsdp_plugin.fsdp_version == 1 - else nullcontext() - ) - sd_options = _prepare_sd_options(fsdp_plugin) - with ctx: - if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT: - optim_state = None - if accelerator.process_index == 0 or not fsdp_plugin.optim_state_dict_config.rank0_only: - optimizer_name = ( - f"{OPTIMIZER_NAME}.bin" if optimizer_index == 0 else f"{OPTIMIZER_NAME}_{optimizer_index}.bin" - ) - input_optimizer_file = os.path.join(input_dir, optimizer_name) - logger.info(f"Loading Optimizer state from {input_optimizer_file}") - optim_state = torch.load(input_optimizer_file, weights_only=True) - logger.info(f"Optimizer state loaded from {input_optimizer_file}") - else: - ckpt_dir = ( - os.path.join(input_dir, f"{OPTIMIZER_NAME}_{optimizer_index}") - if f"{OPTIMIZER_NAME}" not in input_dir - else input_dir - ) - logger.info(f"Loading Optimizer from {ckpt_dir}") - optim_state = {"optimizer": optimizer.state_dict()} - dist_cp.load( - optim_state, - checkpoint_id=ckpt_dir, - storage_reader=dist_cp.FileSystemReader(ckpt_dir), - ) - optim_state = optim_state["optimizer"] - logger.info(f"Optimizer loaded from {ckpt_dir}") - - if fsdp_plugin.fsdp_version == 1: - flattened_osd = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=optim_state) - optimizer.load_state_dict(flattened_osd) - else: - from torch.distributed.checkpoint.state_dict import set_optimizer_state_dict - - set_optimizer_state_dict(model, optimizer, optim_state, options=sd_options) - - -def _distributed_checkpoint_to_merged_weights(checkpoint_dir: str, save_path: str, safe_serialization: bool = True): - """ - Passthrough to `torch.distributed.checkpoint.format_utils.dcp_to_torch_save` - - Will save under `save_path` as either `model.safetensors` or `pytorch_model.bin`. - """ - # Note: We import here to reduce import time from general modules, and isolate outside dependencies - import torch.distributed.checkpoint as dist_cp - import torch.distributed.checkpoint.format_utils as dist_cp_format_utils - - state_dict = {} - save_path = Path(save_path) - save_path.mkdir(exist_ok=True) - dist_cp_format_utils._load_state_dict( - state_dict, - storage_reader=dist_cp.FileSystemReader(checkpoint_dir), - planner=dist_cp_format_utils._EmptyStateDictLoadPlanner(), - no_dist=True, - ) - save_path = save_path / SAFE_WEIGHTS_NAME if safe_serialization else save_path / WEIGHTS_NAME - - # To handle if state is a dict like {model: {...}} - if len(state_dict.keys()) == 1: - state_dict = state_dict[list(state_dict)[0]] - save(state_dict, save_path, safe_serialization=safe_serialization) - return save_path - - -def merge_fsdp_weights( - checkpoint_dir: str, output_path: str, safe_serialization: bool = True, remove_checkpoint_dir: bool = False -): - """ - Merge the weights from sharded FSDP model checkpoints into a single combined checkpoint. Should be used if - `SHARDED_STATE_DICT` was used for the model. Weights will be saved to `{output_path}/model.safetensors` if - `safe_serialization` else `pytorch_model.bin`. - - Note: this is a CPU-bound process. - - Args: - checkpoint_dir (`str`): - The directory containing the FSDP checkpoints (can be either the model or optimizer). - output_path (`str`): - The path to save the merged checkpoint. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the merged weights with safetensors (recommended). - remove_checkpoint_dir (`bool`, *optional*, defaults to `False`): - Whether to remove the checkpoint directory after merging. - """ - checkpoint_dir = Path(checkpoint_dir) - from accelerate.state import PartialState - - if not is_torch_version(">=", "2.3.0"): - raise ValueError("`merge_fsdp_weights` requires PyTorch >= 2.3.0`") - - # Verify that the checkpoint directory exists - if not checkpoint_dir.exists(): - model_path_exists = (checkpoint_dir / "pytorch_model_fsdp_0").exists() - optimizer_path_exists = (checkpoint_dir / "optimizer_0").exists() - err = f"Tried to load from {checkpoint_dir} but couldn't find a valid metadata file." - if model_path_exists and optimizer_path_exists: - err += " However, potential model and optimizer checkpoint directories exist." - err += f"Please pass in either {checkpoint_dir}/pytorch_model_fsdp_0 or {checkpoint_dir}/optimizer_0" - err += "instead." - elif model_path_exists: - err += " However, a potential model checkpoint directory exists." - err += f"Please try passing in {checkpoint_dir}/pytorch_model_fsdp_0 instead." - elif optimizer_path_exists: - err += " However, a potential optimizer checkpoint directory exists." - err += f"Please try passing in {checkpoint_dir}/optimizer_0 instead." - raise ValueError(err) - - # To setup `save` to work - state = PartialState() - if state.is_main_process: - logger.info(f"Merging FSDP weights from {checkpoint_dir}") - save_path = _distributed_checkpoint_to_merged_weights(checkpoint_dir, output_path, safe_serialization) - logger.info(f"Successfully merged FSDP weights and saved to {save_path}") - if remove_checkpoint_dir: - logger.info(f"Removing old checkpoint directory {checkpoint_dir}") - shutil.rmtree(checkpoint_dir) - state.wait_for_everyone() - - -def ensure_weights_retied(param_init_fn, model: torch.nn.Module, device: torch.device): - _tied_names = getattr(model, "_tied_weights_keys", None) - if not _tied_names: - # if no tied names just passthrough - return param_init_fn - - # get map of parameter instances to params. - # - needed for replacement later - _tied_params = {} - for name in _tied_names: - name = name.split(".") - name, param_name = ".".join(name[:-1]), name[-1] - mod = model.get_submodule(name) - param = getattr(mod, param_name) - - _tied_params[id(param)] = None # placeholder for the param first - - # build param_init_fn for the case with tied params - def param_init_fn_tied_param(module: torch.nn.Module): - # track which params to tie - # - usually only 1, but for completeness consider > 1 - params_to_tie = defaultdict(list) - for n, param in module.named_parameters(recurse=False): - if id(param) in _tied_params: - params_to_tie[id(param)].append(n) - - # call the param init fn, which potentially re-allocates the - # parameters - module = param_init_fn(module) - - # search the parameters again and tie them up again - for id_key, _param_names in params_to_tie.items(): - for param_name in _param_names: - param = _tied_params[id_key] - if param is None: - # everything will be tied to the first time the - # param is observed - _tied_params[id_key] = getattr(module, param_name) - else: - setattr(module, param_name, param) # tie - - return module - - return param_init_fn_tied_param - - -def fsdp2_load_full_state_dict(accelerator, model: torch.nn.Module, full_sd: dict): - """ - Loads the full state dict (could be only on rank 0) into the sharded model. This is done by broadcasting the - parameters from rank 0 to all other ranks. This function modifies the model in-place. - - Args: - accelerator (`Accelerator`): The accelerator instance - model (`torch.nn.Module`): - The model to load the state dict into, expected to be on meta device or a VRAM spike can occur - full_sd (`dict`): The full state dict to load, can only be on rank 0 - """ - import torch.distributed as dist - from torch.distributed.tensor import distribute_tensor - - # Model was previously copied to meta device - meta_sharded_sd = model.state_dict() - sharded_sd = {} - - # Rank 0 distributes the full state dict to other ranks - def _infer_parameter_dtype(model, param_name, empty_param): - try: - old_param = model.get_parameter_or_buffer(param_name) - except AttributeError: - # Need this for LORA, as there some params are not *parameters* of sorts - base_param_name, local_param_name = param_name.rsplit(".", 1) - submodule = model.get_submodule(base_param_name) - old_param = getattr(submodule, local_param_name) - - is_torch_e4m3fn_available = hasattr(torch, "float8_e4m3fn") - casting_dtype = None - is_param_float8_e4m3fn = is_torch_e4m3fn_available and empty_param.dtype == torch.float8_e4m3fn - - if empty_param.dtype.is_floating_point and not is_param_float8_e4m3fn: - casting_dtype = old_param.dtype - - return old_param is not None and old_param.is_contiguous(), casting_dtype - - def _cast_and_contiguous(tensor, to_contiguous, dtype): - if dtype is not None: - tensor = tensor.to(dtype=dtype) - if to_contiguous: - tensor = tensor.contiguous() - return tensor - - if accelerator.is_main_process: - for (param_name, full_param), sharded_param in zip(full_sd.items(), meta_sharded_sd.values()): - device_mesh = sharded_param.device_mesh - full_param = full_param.detach().to(device_mesh.device_type) - dist.broadcast(full_param, src=0, group=device_mesh.get_group()) - sharded_tensor = distribute_tensor(full_param, device_mesh, sharded_param.placements) - to_contiguous, casting_dtype = _infer_parameter_dtype( - model, - param_name, - full_param, - ) - sharded_tensor = _cast_and_contiguous(sharded_tensor, to_contiguous, casting_dtype) - sharded_sd[param_name] = sharded_tensor - # We need this else to have a matching `broadcast` for all of the ranks, else we deadlock - else: - for param_name, sharded_param in meta_sharded_sd.items(): - device_mesh = sharded_param.device_mesh - full_tensor = torch.empty(sharded_param.size(), device=device_mesh.device_type, dtype=sharded_param.dtype) - dist.broadcast(full_tensor, src=0, group=device_mesh.get_group()) - sharded_tensor = distribute_tensor(full_tensor, device_mesh, sharded_param.placements) - to_contiguous, casting_dtype = _infer_parameter_dtype( - model, - param_name, - full_tensor, - ) - sharded_tensor = _cast_and_contiguous(sharded_tensor, to_contiguous, casting_dtype) - sharded_sd[param_name] = sharded_tensor - - # we set `assign=True` because our params are on meta device - model.load_state_dict(sharded_sd, assign=True) - return model - - -def fsdp2_switch_optimizer_parameters(optimizer: torch.optim.Optimizer, mapping: dict): - """ - Switches the parameters of the optimizer to new ones (sharded parameters in usual case). This function modifies the - optimizer in-place. - - Args: - optimizer (`torch.optim.Optimizer`): Optimizer instance which contains the original model parameters - mapping (`dict`): Mapping from the original parameter (specified by `data_ptr`) to the sharded parameter - - Raises: - KeyError: - If a parameter in the optimizer couldn't be switched to its sharded version. This should never happen and - indicates a bug. If we kept the original params instead of raising, the training wouldn't be numerically - correct and weights wouldn't get updated. - """ - try: - for param_group in optimizer.param_groups: - param_group["params"] = [mapping[p.data_ptr] for p in param_group["params"]] - except KeyError: - # This shouldn't ever happen, but we want to fail here else training wouldn't be numerically correct - # This basically means that we're missing a mapping from the original parameter to the sharded parameter - raise KeyError( - "A parameter in the optimizer couldn't be switched to its sharded version. This breaks the training. Please raise an issue on GitHub." - ) - - -def fsdp2_apply_ac(accelerator, model: torch.nn.Module): - """ - Applies the activation checkpointing to the model. - - Args: - accelerator (`Accelerator`): The accelerator instance - model (`torch.nn.Module`): The model to apply the activation checkpointing to - - Returns: - `torch.nn.Module`: The model with the activation checkpointing applied - """ - - from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( - checkpoint_wrapper, - ) - - auto_wrap_policy_func = fsdp2_prepare_auto_wrap_policy(accelerator.state.fsdp_plugin, model) - - for layer_name, layer in get_module_children_bottom_up(model, return_fqns=True)[:-1]: - if len(layer_name.split(".")) > 1: - parent_name, child_name = layer_name.rsplit(".", 1) - else: - parent_name = None - child_name = layer_name - - parent_module = model.get_submodule(parent_name) if parent_name else model - if auto_wrap_policy_func(parent_module): - layer = checkpoint_wrapper(layer, preserve_rng_state=False) - parent_module.register_module(child_name, layer) - - return model - - -def fsdp2_prepare_model(accelerator, model: torch.nn.Module) -> torch.nn.Module: - """Prepares the model for FSDP2 in-place. Also returns the model to avoid misuse of the original model. - - Args: - accelerator (`Accelerator`): The accelerator instance - model (`torch.nn.Module`): The model to prepare - - Returns: - `torch.nn.Module`: Prepared model - """ - from torch.distributed.fsdp import FSDPModule, MixedPrecisionPolicy, fully_shard - - is_type_fsdp = isinstance(model, FSDPModule) or ( - is_compiled_module(model) and isinstance(model._orig_mod, FSDPModule) - ) - if is_type_fsdp: - return model - - fsdp2_plugin = accelerator.state.fsdp_plugin - - fsdp2_plugin.set_auto_wrap_policy(model) - - original_sd = model.state_dict() - - fsdp2_kwargs = { - "reshard_after_forward": fsdp2_plugin.reshard_after_forward, - "offload_policy": fsdp2_plugin.cpu_offload, - # `fully_shard` doesn't accept `None` in case of `MixedPrecisionPolicy` - "mp_policy": fsdp2_plugin.mixed_precision_policy or MixedPrecisionPolicy(), - } - - model_has_params4bit = False - for name, param in model.named_parameters(): - # this is a temporary fix whereby loading models with bnb params cannot be moved from - # GPU to a meta device due with FSDP2 because torch operations don't return the original class type - # bypassing the move to meta will still cause the VRAM spike, but at least it still will load - if param.__class__.__name__ == "Params4bit": - model_has_params4bit = True - break - - if fsdp2_plugin.cpu_ram_efficient_loading and not model_has_params4bit: - # Context: `fully_shard` moves the model to GPU if it was on CPU, however it can also be on `meta` and then it stays there even after `fully_shard` - # For this reason, we need to move the model to `meta` device, as then sharding happens on `meta` device - # If we kept the model on CPU (`cpu_ram_efficient_loading` has model be on CPU on all ranks, though non-main ranks only have `torch.emtpy`), `fully_shard` would move it to GPU - # Afterwards, when we call `fsdp2_load_full_state_dict`, us creating the state_dict would result into briefly having two copies of model state_dict on the GPU -> VRAM spike - - # We need to keep the original non-persistent buffers, as those MAY not be in the state_dict, resulting in them staying on meta device - # Also, these buffers aren't getting sharded by default - # We get the FQNs of all non-persistent buffers, to re-register them after - non_persistent_buffer_fqns = get_non_persistent_buffers(model, recurse=True, fqns=True) - original_non_persistent_buffers = copy.deepcopy( - {k: v for k, v in model.named_buffers() if k in non_persistent_buffer_fqns} - ) - # We move the model to meta device, as then sharding happens on meta device - model = model.to(torch.device("meta")) - # We need to re-tie the weights, not exactly sure why, but if we don't do this, reference to `lm_head/embed_tokens` stay hanging -> more VRAM usage - # We assume `transformers` models have a `tie_weights` method if they support it - if hasattr(model, "tie_weights"): - model.tie_weights() - - auto_wrap_policy_func = fsdp2_prepare_auto_wrap_policy(fsdp2_plugin, model) - if auto_wrap_policy_func is not None: - # We skip the model itself, as that one is always wrapped - for module in get_module_children_bottom_up(model)[:-1]: - if auto_wrap_policy_func(module) and not isinstance(module, FSDPModule): - fully_shard(module, **fsdp2_kwargs) - - if not isinstance(model, FSDPModule): - fully_shard(model, **fsdp2_kwargs) - - if fsdp2_plugin.cpu_ram_efficient_loading: - # If `cpu_ram_efficient_loading` is enabled, only rank 0 loads the weights - # Other ranks have an empty model on `meta` device, so we need to distribute the weights properly - fsdp2_load_full_state_dict(accelerator, model, original_sd) - - if fsdp2_plugin.cpu_ram_efficient_loading and not model_has_params4bit: - # We re-register the buffers, as they may not be in the state_dict - for fqn, buffer_tensor in original_non_persistent_buffers.items(): - buffer_tensor = buffer_tensor.to(accelerator.device) - - if "." in fqn: - parent_fqn, local_buffer_name = fqn.rsplit(".", 1) - parent_module = model.get_submodule(parent_fqn) - else: - local_buffer_name = fqn - parent_module = model - - parent_module.register_buffer(local_buffer_name, buffer_tensor, persistent=False) - - # We need to tie the weights again, as call to `load_full_state_dict` breaks the tie - # Needs to be called both here and above - # removing this call makes the have slightly different loss - # removing the call above leads to extra memory usage as explained in the comment above - if hasattr(model, "tie_weights"): - model.tie_weights() - - # There is no `dtype` attribution for nn.Module - # Set it to None if it doesn't exist and do the upcast always - model_dtype = getattr(model, "dtype", None) - if accelerator.mixed_precision != "no" and (model_dtype is None or model_dtype != torch.float32): - # We upcast the model according to `deepspeed`'s implementation - # More info about this can be found in `accelerator.py:prepare_model`s FSDP1 section - model = model.to(torch.float32) - if accelerator.is_main_process: - # TODO(siro1): Add a warning for each parameter that was upcasted - warnings.warn( - "FSDP upcast of low precision parameters to fp32 (since mixed_precision != 'no') may affect the precision of model checkpoints." - ) - return model - - -def fsdp2_prepare_auto_wrap_policy(fsdp2_plugin, model: torch.nn.Module) -> Callable[[torch.nn.Module], bool]: - """Prepares the auto wrap policy based on its type, done to mimic the behaviour of FSDP1 auto wrap policy. - - Args: - fsdp2_plugin (`FullyShardedDataParallelPlugin`): - Instance of `FullyShardedDataParallelPlugin` containing the configuration options - auto_wrap_policy_type (`str`): - Either `transformer` or `size` - model (`torch.nn.Module`): - The model to wrap - - Returns: - `Callable[[torch.nn.Module], bool]`: - The auto wrap policy function to be applied to the model - """ - from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy - - fn = fsdp2_plugin.auto_wrap_policy - - if isinstance(fn, functools.partial): - fn = fn.func - - if fn is transformer_auto_wrap_policy: - no_split_modules = getattr(model, "_no_split_modules", None) - if no_split_modules is None: - no_split_modules = [] - transformer_cls_names_to_wrap = list(no_split_modules) - if fsdp2_plugin.transformer_cls_names_to_wrap is not None: - transformer_cls_names_to_wrap = fsdp2_plugin.transformer_cls_names_to_wrap - transformer_cls_to_wrap = set() - - for layer_class in transformer_cls_names_to_wrap: - transformer_cls = get_module_class_from_name(model, layer_class) - if transformer_cls is None: - raise ValueError(f"Could not find the transformer layer class {layer_class} in the model.") - transformer_cls_to_wrap.add(transformer_cls) - - def policy(module: torch.nn.Module) -> bool: - if fsdp2_plugin.transformer_cls_names_to_wrap is None: - return False - return isinstance(module, tuple(transformer_cls_to_wrap)) - - elif fn is size_based_auto_wrap_policy: - - def policy(module: torch.nn.Module) -> bool: - module_num_params = sum(p.numel() for p in module.parameters()) - return module_num_params > fsdp2_plugin.min_num_params - else: - return None - - return policy - - -def get_fsdp2_grad_scaler(**kwargs): - """ - Returns a `GradScaler` for FSDP2, as the current implementation of `get_grad_scaler` doesn't accept other args. We - need this as current `get_grad_scaler` accepts only `distributed_type` as arg, which doesn't differentiate between - FSDP1 and FSDP2 - """ - from torch.amp.grad_scaler import GradScaler - - return GradScaler(**kwargs) - - -def fsdp2_canonicalize_names(named_params: dict) -> dict: - """Removes parameter name modifiers in order to map them back to their original names. - - See huggingface/accelerate#3554 for more context. - - Args: - named_params (`dict`): The named parameters dictionary to canonicalize. - - Returns: - `dict`: The canonicalized named parameters dictionary - """ - named_params = {k.replace("._checkpoint_wrapped_module", ""): v for k, v in named_params.items()} - named_params = { - k.replace("_orig_mod.", "") if k.startswith("_orig_mod.") else k: v for k, v in named_params.items() - } - named_params = {k.replace("._orig_mod", ""): v for k, v in named_params.items()} - return named_params diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/imports.py b/venv/lib/python3.10/site-packages/accelerate/utils/imports.py deleted file mode 100644 index a919827994824cdbc745acac83f03e501c8688f9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/imports.py +++ /dev/null @@ -1,551 +0,0 @@ -# Copyright 2022 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 importlib.metadata -import os -import sys -import warnings -from functools import lru_cache, wraps - -import torch -from packaging import version -from packaging.version import parse - -from .environment import parse_flag_from_env, patch_environment, str_to_bool -from .versions import compare_versions, is_torch_version - - -# Try to run Torch native job in an environment with TorchXLA installed by setting this value to 0. -USE_TORCH_XLA = parse_flag_from_env("USE_TORCH_XLA", default=True) - -_torch_xla_available = False -if USE_TORCH_XLA: - try: - import torch_xla.core.xla_model as xm # noqa: F401 - import torch_xla.runtime - - _torch_xla_available = True - except ImportError: - pass - -# Keep it for is_tpu_available. It will be removed along with is_tpu_available. -_tpu_available = _torch_xla_available - -# Cache this result has it's a C FFI call which can be pretty time-consuming -_torch_distributed_available = torch.distributed.is_available() - - -def _is_package_available(pkg_name, metadata_name=None): - # Check we're not importing a "pkg_name" directory somewhere but the actual library by trying to grab the version - package_exists = importlib.util.find_spec(pkg_name) is not None - if package_exists: - try: - # Some libraries have different names in the metadata - _ = importlib.metadata.metadata(pkg_name if metadata_name is None else metadata_name) - return True - except importlib.metadata.PackageNotFoundError: - return False - - -def is_torch_distributed_available() -> bool: - return _torch_distributed_available - - -def is_xccl_available(): - if is_torch_version(">=", "2.7.0"): - return torch.distributed.distributed_c10d.is_xccl_available() - if is_ipex_available(): - return False - return False - - -def is_ccl_available(): - try: - pass - except ImportError: - print( - "Intel(R) oneCCL Bindings for PyTorch* is required to run DDP on Intel(R) XPUs, but it is not" - " detected. If you see \"ValueError: Invalid backend: 'ccl'\" error, please install Intel(R) oneCCL" - " Bindings for PyTorch*." - ) - return importlib.util.find_spec("oneccl_bindings_for_pytorch") is not None - - -def get_ccl_version(): - return importlib.metadata.version("oneccl_bind_pt") - - -def is_import_timer_available(): - return _is_package_available("import_timer") - - -def is_pynvml_available(): - return _is_package_available("pynvml") or _is_package_available("pynvml", "nvidia-ml-py") - - -def is_pytest_available(): - return _is_package_available("pytest") - - -def is_msamp_available(): - return _is_package_available("msamp", "ms-amp") - - -def is_schedulefree_available(): - return _is_package_available("schedulefree") - - -def is_transformer_engine_available(): - if is_hpu_available(): - return _is_package_available("intel_transformer_engine", "intel-transformer-engine") - else: - return _is_package_available("transformer_engine", "transformer-engine") - - -def is_lomo_available(): - return _is_package_available("lomo_optim") - - -def is_cuda_available(): - """ - Checks if `cuda` is available via an `nvml-based` check which won't trigger the drivers and leave cuda - uninitialized. - """ - with patch_environment(PYTORCH_NVML_BASED_CUDA_CHECK="1"): - available = torch.cuda.is_available() - - return available - - -@lru_cache -def is_torch_xla_available(check_is_tpu=False, check_is_gpu=False): - """ - Check if `torch_xla` is available. To train a native pytorch job in an environment with torch xla installed, set - the USE_TORCH_XLA to false. - """ - assert not (check_is_tpu and check_is_gpu), "The check_is_tpu and check_is_gpu cannot both be true." - - if not _torch_xla_available: - return False - elif check_is_gpu: - return torch_xla.runtime.device_type() in ["GPU", "CUDA"] - elif check_is_tpu: - return torch_xla.runtime.device_type() == "TPU" - - return True - - -def is_torchao_available(): - package_exists = _is_package_available("torchao") - if package_exists: - torchao_version = version.parse(importlib.metadata.version("torchao")) - return compare_versions(torchao_version, ">=", "0.6.1") - return False - - -def is_deepspeed_available(): - return _is_package_available("deepspeed") - - -def is_pippy_available(): - return is_torch_version(">=", "2.4.0") - - -def is_bf16_available(ignore_tpu=False): - "Checks if bf16 is supported, optionally ignoring the TPU" - if is_torch_xla_available(check_is_tpu=True): - return not ignore_tpu - if is_cuda_available(): - return torch.cuda.is_bf16_supported() - if is_mlu_available(): - return torch.mlu.is_bf16_supported() - if is_xpu_available(): - return torch.xpu.is_bf16_supported() - if is_mps_available(): - return False - return True - - -def is_fp16_available(): - "Checks if fp16 is supported" - if is_habana_gaudi1(): - return False - - return True - - -def is_fp8_available(): - "Checks if fp8 is supported" - return is_msamp_available() or is_transformer_engine_available() or is_torchao_available() - - -def is_4bit_bnb_available(): - package_exists = _is_package_available("bitsandbytes") - if package_exists: - bnb_version = version.parse(importlib.metadata.version("bitsandbytes")) - return compare_versions(bnb_version, ">=", "0.39.0") - return False - - -def is_8bit_bnb_available(): - package_exists = _is_package_available("bitsandbytes") - if package_exists: - bnb_version = version.parse(importlib.metadata.version("bitsandbytes")) - return compare_versions(bnb_version, ">=", "0.37.2") - return False - - -def is_bnb_available(min_version=None): - package_exists = _is_package_available("bitsandbytes") - if package_exists and min_version is not None: - bnb_version = version.parse(importlib.metadata.version("bitsandbytes")) - return compare_versions(bnb_version, ">=", min_version) - else: - return package_exists - - -def is_bitsandbytes_multi_backend_available(): - if not is_bnb_available(): - return False - import bitsandbytes as bnb - - return "multi_backend" in getattr(bnb, "features", set()) - - -def is_torchvision_available(): - return _is_package_available("torchvision") - - -def is_megatron_lm_available(): - if str_to_bool(os.environ.get("ACCELERATE_USE_MEGATRON_LM", "False")) == 1: - if importlib.util.find_spec("megatron") is not None: - try: - megatron_version = parse(importlib.metadata.version("megatron-core")) - if compare_versions(megatron_version, ">=", "0.8.0"): - return importlib.util.find_spec(".training", "megatron") - except Exception as e: - warnings.warn(f"Parse Megatron version failed. Exception:{e}") - return False - - -def is_transformers_available(): - return _is_package_available("transformers") - - -def is_datasets_available(): - return _is_package_available("datasets") - - -def is_peft_available(): - return _is_package_available("peft") - - -def is_timm_available(): - return _is_package_available("timm") - - -def is_triton_available(): - if is_xpu_available(): - return _is_package_available("triton", "pytorch-triton-xpu") - return _is_package_available("triton") - - -def is_aim_available(): - package_exists = _is_package_available("aim") - if package_exists: - aim_version = version.parse(importlib.metadata.version("aim")) - return compare_versions(aim_version, "<", "4.0.0") - return False - - -def is_tensorboard_available(): - return _is_package_available("tensorboard") or _is_package_available("tensorboardX") - - -def is_wandb_available(): - return _is_package_available("wandb") - - -def is_comet_ml_available(): - return _is_package_available("comet_ml") - - -def is_swanlab_available(): - return _is_package_available("swanlab") - - -def is_trackio_available(): - return sys.version_info >= (3, 10) and _is_package_available("trackio") - - -def is_boto3_available(): - return _is_package_available("boto3") - - -def is_rich_available(): - if _is_package_available("rich"): - return parse_flag_from_env("ACCELERATE_ENABLE_RICH", False) - return False - - -def is_sagemaker_available(): - return _is_package_available("sagemaker") - - -def is_tqdm_available(): - return _is_package_available("tqdm") - - -def is_clearml_available(): - return _is_package_available("clearml") - - -def is_pandas_available(): - return _is_package_available("pandas") - - -def is_matplotlib_available(): - return _is_package_available("matplotlib") - - -def is_mlflow_available(): - if _is_package_available("mlflow"): - return True - - if importlib.util.find_spec("mlflow") is not None: - try: - _ = importlib.metadata.metadata("mlflow-skinny") - return True - except importlib.metadata.PackageNotFoundError: - return False - return False - - -def is_mps_available(min_version="1.12"): - "Checks if MPS device is available. The minimum version required is 1.12." - # With torch 1.12, you can use torch.backends.mps - # With torch 2.0.0, you can use torch.mps - return is_torch_version(">=", min_version) and torch.backends.mps.is_available() and torch.backends.mps.is_built() - - -def is_ipex_available(): - "Checks if ipex is installed." - - def get_major_and_minor_from_version(full_version): - return str(version.parse(full_version).major) + "." + str(version.parse(full_version).minor) - - _torch_version = importlib.metadata.version("torch") - if importlib.util.find_spec("intel_extension_for_pytorch") is None: - return False - _ipex_version = "N/A" - try: - _ipex_version = importlib.metadata.version("intel_extension_for_pytorch") - except importlib.metadata.PackageNotFoundError: - return False - torch_major_and_minor = get_major_and_minor_from_version(_torch_version) - ipex_major_and_minor = get_major_and_minor_from_version(_ipex_version) - if torch_major_and_minor != ipex_major_and_minor: - warnings.warn( - f"Intel Extension for PyTorch {ipex_major_and_minor} needs to work with PyTorch {ipex_major_and_minor}.*," - f" but PyTorch {_torch_version} is found. Please switch to the matching version and run again." - ) - return False - return True - - -@lru_cache -def is_mlu_available(check_device=False): - """ - Checks if `mlu` is available via an `cndev-based` check which won't trigger the drivers and leave mlu - uninitialized. - """ - if importlib.util.find_spec("torch_mlu") is None: - return False - - import torch_mlu # noqa: F401 - - with patch_environment(PYTORCH_CNDEV_BASED_MLU_CHECK="1"): - available = torch.mlu.is_available() - - return available - - -@lru_cache -def is_musa_available(check_device=False): - "Checks if `torch_musa` is installed and potentially if a MUSA is in the environment" - if importlib.util.find_spec("torch_musa") is None: - return False - - import torch_musa # noqa: F401 - - if check_device: - try: - # Will raise a RuntimeError if no MUSA is found - _ = torch.musa.device_count() - return torch.musa.is_available() - except RuntimeError: - return False - return hasattr(torch, "musa") and torch.musa.is_available() - - -@lru_cache -def is_npu_available(check_device=False): - "Checks if `torch_npu` is installed and potentially if a NPU is in the environment" - if importlib.util.find_spec("torch_npu") is None: - return False - - import torch_npu # noqa: F401 - - if check_device: - try: - # Will raise a RuntimeError if no NPU is found - _ = torch.npu.device_count() - return torch.npu.is_available() - except RuntimeError: - return False - return hasattr(torch, "npu") and torch.npu.is_available() - - -@lru_cache -def is_sdaa_available(check_device=False): - "Checks if `torch_sdaa` is installed and potentially if a SDAA is in the environment" - if importlib.util.find_spec("torch_sdaa") is None: - return False - - import torch_sdaa # noqa: F401 - - if check_device: - try: - # Will raise a RuntimeError if no NPU is found - _ = torch.sdaa.device_count() - return torch.sdaa.is_available() - except RuntimeError: - return False - return hasattr(torch, "sdaa") and torch.sdaa.is_available() - - -@lru_cache -def is_hpu_available(init_hccl=False): - "Checks if `torch.hpu` is installed and potentially if a HPU is in the environment" - if ( - importlib.util.find_spec("habana_frameworks") is None - or importlib.util.find_spec("habana_frameworks.torch") is None - ): - return False - - import habana_frameworks.torch # noqa: F401 - - if init_hccl: - import habana_frameworks.torch.distributed.hccl as hccl # noqa: F401 - - return hasattr(torch, "hpu") and torch.hpu.is_available() - - -def is_habana_gaudi1(): - if is_hpu_available(): - import habana_frameworks.torch.utils.experimental as htexp # noqa: F401 - - if htexp._get_device_type() == htexp.synDeviceType.synDeviceGaudi: - return True - - return False - - -@lru_cache -def is_xpu_available(check_device=False): - """ - Checks if XPU acceleration is available either via `intel_extension_for_pytorch` or via stock PyTorch (>=2.4) and - potentially if a XPU is in the environment - """ - - if is_ipex_available(): - import intel_extension_for_pytorch # noqa: F401 - else: - if is_torch_version("<=", "2.3"): - return False - - if check_device: - try: - # Will raise a RuntimeError if no XPU is found - _ = torch.xpu.device_count() - return torch.xpu.is_available() - except RuntimeError: - return False - return hasattr(torch, "xpu") and torch.xpu.is_available() - - -def is_dvclive_available(): - return _is_package_available("dvclive") - - -def is_torchdata_available(): - return _is_package_available("torchdata") - - -# TODO: Remove this function once stateful_dataloader is a stable feature in torchdata. -def is_torchdata_stateful_dataloader_available(): - package_exists = _is_package_available("torchdata") - if package_exists: - torchdata_version = version.parse(importlib.metadata.version("torchdata")) - return compare_versions(torchdata_version, ">=", "0.8.0") - return False - - -def torchao_required(func): - """ - A decorator that ensures the decorated function is only called when torchao is available. - """ - - @wraps(func) - def wrapper(*args, **kwargs): - if not is_torchao_available(): - raise ImportError( - "`torchao` is not available, please install it before calling this function via `pip install torchao`." - ) - return func(*args, **kwargs) - - return wrapper - - -# TODO: Rework this into `utils.deepspeed` and migrate the "core" chunks into `accelerate.deepspeed` -def deepspeed_required(func): - """ - A decorator that ensures the decorated function is only called when deepspeed is enabled. - """ - - @wraps(func) - def wrapper(*args, **kwargs): - from accelerate.state import AcceleratorState - from accelerate.utils.dataclasses import DistributedType - - if AcceleratorState._shared_state != {} and AcceleratorState().distributed_type != DistributedType.DEEPSPEED: - raise ValueError( - "DeepSpeed is not enabled, please make sure that an `Accelerator` is configured for `deepspeed` " - "before calling this function." - ) - return func(*args, **kwargs) - - return wrapper - - -def is_weights_only_available(): - # Weights only with allowlist was added in 2.4.0 - # ref: https://github.com/pytorch/pytorch/pull/124331 - return is_torch_version(">=", "2.4.0") - - -def is_numpy_available(min_version="1.25.0"): - numpy_version = parse(importlib.metadata.version("numpy")) - return compare_versions(numpy_version, ">=", min_version) diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/launch.py b/venv/lib/python3.10/site-packages/accelerate/utils/launch.py deleted file mode 100644 index 0befeee898027678b6b6c75e2d72c9f5b33cdbfc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/launch.py +++ /dev/null @@ -1,744 +0,0 @@ -# Copyright 2022 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 argparse -import os -import subprocess -import sys -import warnings -from ast import literal_eval -from shutil import which -from typing import Any - -import torch - -from ..commands.config.config_args import SageMakerConfig -from ..utils import ( - DynamoBackend, - PrecisionType, - is_ccl_available, - is_fp8_available, - is_hpu_available, - is_ipex_available, - is_mlu_available, - is_musa_available, - is_npu_available, - is_sdaa_available, - is_torch_xla_available, - is_xpu_available, -) -from ..utils.constants import DEEPSPEED_MULTINODE_LAUNCHERS -from ..utils.other import get_free_port, is_port_in_use, merge_dicts -from ..utils.versions import compare_versions -from .dataclasses import DistributedType, SageMakerDistributedType - - -def _filter_args(args, parser, default_args=[]): - """ - Filters out all `accelerate` specific args - """ - new_args, _ = parser.parse_known_args(default_args) - for key, value in vars(args).items(): - if key in vars(new_args).keys(): - setattr(new_args, key, value) - return new_args - - -def _get_mpirun_args(): - """ - Determines the executable and argument names for mpirun, based on the type of install. The supported MPI programs - are: OpenMPI, Intel MPI, or MVAPICH. - - Returns: Program name and arg names for hostfile, num processes, and processes per node - """ - # Find the MPI program name - mpi_apps = [x for x in ["mpirun", "mpiexec"] if which(x)] - - if len(mpi_apps) == 0: - raise OSError("mpirun or mpiexec were not found. Ensure that Intel MPI, Open MPI, or MVAPICH are installed.") - - # Call the app with the --version flag to determine which MPI app is installed - mpi_app = mpi_apps[0] - mpirun_version = subprocess.check_output([mpi_app, "--version"]) - - if b"Open MPI" in mpirun_version: - return mpi_app, "--hostfile", "-n", "--npernode", "--bind-to" - else: - # Intel MPI and MVAPICH both use the same arg names - return mpi_app, "-f", "-n", "-ppn", "" - - -def setup_fp8_env(args: argparse.Namespace, current_env: dict[str, str]): - """ - Setup the FP8 environment variables. - """ - prefix = "ACCELERATE_" - for arg in vars(args): - if arg.startswith("fp8_"): - value = getattr(args, arg) - if value is not None: - if arg == "fp8_override_linear_precision": - current_env[prefix + "FP8_OVERRIDE_FPROP"] = str(value[0]) - current_env[prefix + "FP8_OVERRIDE_DGRAD"] = str(value[1]) - current_env[prefix + "FP8_OVERRIDE_WGRAD"] = str(value[2]) - else: - current_env[f"{prefix}{arg.upper()}"] = str(getattr(args, arg)) - return current_env - - -def prepare_simple_launcher_cmd_env(args: argparse.Namespace) -> tuple[list[str], dict[str, str]]: - """ - Prepares and returns the command list and an environment with the correct simple launcher environment variables. - """ - cmd = [] - if args.no_python and args.module: - raise ValueError("--module and --no_python cannot be used together") - - num_processes = getattr(args, "num_processes", None) - num_machines = args.num_machines - if args.mpirun_hostfile is not None: - mpi_app_name, hostfile_arg, num_proc_arg, proc_per_node_arg, bind_to_arg = _get_mpirun_args() - bind_to = getattr(args, "bind-to", "socket") - nproc_per_node = str(num_processes // num_machines) if num_processes and num_machines else "1" - cmd += [ - mpi_app_name, - hostfile_arg, - args.mpirun_hostfile, - proc_per_node_arg, - nproc_per_node, - ] - if num_processes: - cmd += [num_proc_arg, str(num_processes)] - if bind_to_arg: - cmd += [bind_to_arg, bind_to] - if not args.no_python: - cmd.append(sys.executable) - if args.module: - cmd.append("-m") - cmd.append(args.training_script) - cmd.extend(args.training_script_args) - - current_env = os.environ.copy() - current_env["ACCELERATE_USE_CPU"] = str(args.cpu or args.use_cpu) - if args.debug: - current_env["ACCELERATE_DEBUG_MODE"] = "true" - if args.gpu_ids != "all" and args.gpu_ids is not None: - if is_xpu_available(): - current_env["ZE_AFFINITY_MASK"] = args.gpu_ids - elif is_mlu_available(): - current_env["MLU_VISIBLE_DEVICES"] = args.gpu_ids - elif is_sdaa_available(): - current_env["SDAA_VISIBLE_DEVICES"] = args.gpu_ids - elif is_musa_available(): - current_env["MUSA_VISIBLE_DEVICES"] = args.gpu_ids - elif is_npu_available(): - current_env["ASCEND_RT_VISIBLE_DEVICES"] = args.gpu_ids - elif is_hpu_available(): - current_env["HABANA_VISIBLE_MODULES"] = args.gpu_ids - else: - current_env["CUDA_VISIBLE_DEVICES"] = args.gpu_ids - if num_machines > 1: - assert args.main_process_ip is not None, ( - "When using multiple machines, you need to specify the main process IP." - ) - assert args.main_process_port is not None, ( - "When using multiple machines, you need to specify the main process port." - ) - - ccl_worker_count = getattr(args, "mpirun_ccl", 0) if is_ccl_available() else 0 - if (num_processes is not None and num_processes > 1) or num_machines > 1: - current_env["MASTER_ADDR"] = args.main_process_ip if args.main_process_ip is not None else "127.0.0.1" - current_env["MASTER_PORT"] = str(args.main_process_port) if args.main_process_port is not None else "29500" - current_env["CCL_WORKER_COUNT"] = str(ccl_worker_count) - if current_env["ACCELERATE_USE_CPU"]: - current_env["KMP_AFFINITY"] = "granularity=fine,compact,1,0" - current_env["KMP_BLOCKTIME"] = str(1) - - try: - mixed_precision = PrecisionType(args.mixed_precision.lower()) - except ValueError: - raise ValueError( - f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}." - ) - - current_env["ACCELERATE_MIXED_PRECISION"] = str(mixed_precision) - if args.mixed_precision.lower() == "fp8": - if not is_fp8_available(): - raise RuntimeError( - "FP8 is not available on this machine. Please ensure that either Transformer Engine, MSAMP or torchao is installed." - ) - current_env = setup_fp8_env(args, current_env) - - try: - dynamo_backend = DynamoBackend(args.dynamo_backend.upper()) - except ValueError: - raise ValueError( - f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DynamoBackend.list()}." - ) - current_env["ACCELERATE_DYNAMO_BACKEND"] = dynamo_backend.value - current_env["ACCELERATE_DYNAMO_MODE"] = args.dynamo_mode - current_env["ACCELERATE_DYNAMO_USE_FULLGRAPH"] = str(args.dynamo_use_fullgraph) - current_env["ACCELERATE_DYNAMO_USE_DYNAMIC"] = str(args.dynamo_use_dynamic) - current_env["ACCELERATE_DYNAMO_USE_REGIONAL_COMPILATION"] = str(args.dynamo_use_regional_compilation) - - current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process) - if is_ipex_available(): - current_env["ACCELERATE_USE_IPEX"] = str(args.ipex).lower() - if args.enable_cpu_affinity: - current_env["ACCELERATE_CPU_AFFINITY"] = "1" - return cmd, current_env - - -def prepare_multi_gpu_env(args: argparse.Namespace) -> dict[str, str]: - """ - Prepares and returns an environment with the correct multi-GPU environment variables. - """ - # get free port and update configurations - if args.main_process_port == 0: - args.main_process_port = get_free_port() - - elif args.main_process_port is None: - args.main_process_port = 29500 - - num_processes = args.num_processes - num_machines = args.num_machines - main_process_ip = args.main_process_ip - main_process_port = args.main_process_port - if num_machines > 1: - args.nproc_per_node = str(num_processes // num_machines) - args.nnodes = str(num_machines) - args.node_rank = int(args.machine_rank) - if getattr(args, "same_network", False): - args.master_addr = str(main_process_ip) - args.master_port = str(main_process_port) - else: - args.rdzv_endpoint = f"{main_process_ip}:{main_process_port}" - else: - args.nproc_per_node = str(num_processes) - if main_process_port is not None: - args.master_port = str(main_process_port) - - # only need to check port availability in main process, in case we have to start multiple launchers on the same machine - # for some reasons like splitting log files. - need_port_check = num_machines <= 1 or int(args.machine_rank) == 0 - if need_port_check and is_port_in_use(main_process_port): - if num_machines <= 1: - args.standalone = True - warnings.warn( - f"Port `{main_process_port}` is already in use. " - "Accelerate will attempt to launch in a standalone-like mode by finding an open port automatically for this session. " - "If this current attempt fails, or for more control in future runs, please specify a different port " - "(e.g., `--main_process_port `) or use `--main_process_port 0` for automatic selection " - "in your launch command or Accelerate config file." - ) - else: - raise ConnectionError( - f"Tried to launch distributed communication on port `{main_process_port}`, but another process is utilizing it. " - "Please specify a different port (such as using the `--main_process_port` flag or specifying a different `main_process_port` in your config file)" - " and rerun your script. To automatically use the next open port (on a single node), you can set this to `0`." - ) - - if args.module and args.no_python: - raise ValueError("--module and --no_python cannot be used together") - elif args.module: - args.module = True - elif args.no_python: - args.no_python = True - - current_env = os.environ.copy() - if args.debug: - current_env["ACCELERATE_DEBUG_MODE"] = "true" - gpu_ids = getattr(args, "gpu_ids", "all") - if gpu_ids != "all" and args.gpu_ids is not None: - if is_xpu_available(): - current_env["ZE_AFFINITY_MASK"] = gpu_ids - elif is_mlu_available(): - current_env["MLU_VISIBLE_DEVICES"] = gpu_ids - elif is_sdaa_available(): - current_env["SDAA_VISIBLE_DEVICES"] = gpu_ids - elif is_musa_available(): - current_env["MUSA_VISIBLE_DEVICES"] = gpu_ids - elif is_npu_available(): - current_env["ASCEND_RT_VISIBLE_DEVICES"] = gpu_ids - elif is_hpu_available(): - current_env["HABANA_VISIBLE_MODULES"] = gpu_ids - else: - current_env["CUDA_VISIBLE_DEVICES"] = gpu_ids - mixed_precision = args.mixed_precision.lower() - try: - mixed_precision = PrecisionType(mixed_precision) - except ValueError: - raise ValueError(f"Unknown mixed_precision mode: {mixed_precision}. Choose between {PrecisionType.list()}.") - - current_env["ACCELERATE_MIXED_PRECISION"] = str(mixed_precision) - if args.mixed_precision.lower() == "fp8": - if not is_fp8_available(): - raise RuntimeError( - "FP8 is not available on this machine. Please ensure that either Transformer Engine, MSAMP or torchao is installed." - ) - current_env = setup_fp8_env(args, current_env) - - try: - dynamo_backend = DynamoBackend(args.dynamo_backend.upper()) - except ValueError: - raise ValueError( - f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DynamoBackend.list()}." - ) - current_env["ACCELERATE_DYNAMO_BACKEND"] = dynamo_backend.value - current_env["ACCELERATE_DYNAMO_MODE"] = args.dynamo_mode - current_env["ACCELERATE_DYNAMO_USE_FULLGRAPH"] = str(args.dynamo_use_fullgraph) - current_env["ACCELERATE_DYNAMO_USE_DYNAMIC"] = str(args.dynamo_use_dynamic) - current_env["ACCELERATE_DYNAMO_USE_REGIONAL_COMPILATION"] = str(args.dynamo_use_regional_compilation) - - if args.use_fsdp: - current_env["ACCELERATE_USE_FSDP"] = "true" - if args.fsdp_cpu_ram_efficient_loading and not args.fsdp_sync_module_states: - raise ValueError("When using `--fsdp_cpu_ram_efficient_loading` set `--fsdp_sync_module_states` to `True`") - - current_env["FSDP_VERSION"] = str(args.fsdp_version) if hasattr(args, "fsdp_version") else "1" - - # For backwards compatibility, we support this in launched scripts, - # however, we do not ask users for this in `accelerate config` CLI - current_env["FSDP_SHARDING_STRATEGY"] = str(args.fsdp_sharding_strategy) - - current_env["FSDP_RESHARD_AFTER_FORWARD"] = str(args.fsdp_reshard_after_forward).lower() - current_env["FSDP_OFFLOAD_PARAMS"] = str(args.fsdp_offload_params).lower() - current_env["FSDP_MIN_NUM_PARAMS"] = str(args.fsdp_min_num_params) - if args.fsdp_auto_wrap_policy is not None: - current_env["FSDP_AUTO_WRAP_POLICY"] = str(args.fsdp_auto_wrap_policy) - if args.fsdp_transformer_layer_cls_to_wrap is not None: - current_env["FSDP_TRANSFORMER_CLS_TO_WRAP"] = str(args.fsdp_transformer_layer_cls_to_wrap) - if args.fsdp_backward_prefetch is not None: - current_env["FSDP_BACKWARD_PREFETCH"] = str(args.fsdp_backward_prefetch) - if args.fsdp_state_dict_type is not None: - current_env["FSDP_STATE_DICT_TYPE"] = str(args.fsdp_state_dict_type) - current_env["FSDP_FORWARD_PREFETCH"] = str(args.fsdp_forward_prefetch).lower() - current_env["FSDP_USE_ORIG_PARAMS"] = str(args.fsdp_use_orig_params).lower() - current_env["FSDP_CPU_RAM_EFFICIENT_LOADING"] = str(args.fsdp_cpu_ram_efficient_loading).lower() - current_env["FSDP_SYNC_MODULE_STATES"] = str(args.fsdp_sync_module_states).lower() - current_env["FSDP_ACTIVATION_CHECKPOINTING"] = str(args.fsdp_activation_checkpointing).lower() - - if args.use_megatron_lm: - prefix = "MEGATRON_LM_" - current_env["ACCELERATE_USE_MEGATRON_LM"] = "true" - current_env[prefix + "TP_DEGREE"] = str(args.megatron_lm_tp_degree) - current_env[prefix + "PP_DEGREE"] = str(args.megatron_lm_pp_degree) - current_env[prefix + "GRADIENT_CLIPPING"] = str(args.megatron_lm_gradient_clipping) - if args.megatron_lm_num_micro_batches is not None: - current_env[prefix + "NUM_MICRO_BATCHES"] = str(args.megatron_lm_num_micro_batches) - if args.megatron_lm_sequence_parallelism is not None: - current_env[prefix + "SEQUENCE_PARALLELISM"] = str(args.megatron_lm_sequence_parallelism) - if args.megatron_lm_recompute_activations is not None: - current_env[prefix + "RECOMPUTE_ACTIVATIONS"] = str(args.megatron_lm_recompute_activations) - if args.megatron_lm_use_distributed_optimizer is not None: - current_env[prefix + "USE_DISTRIBUTED_OPTIMIZER"] = str(args.megatron_lm_use_distributed_optimizer) - - current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process) - if args.enable_cpu_affinity: - current_env["ACCELERATE_CPU_AFFINITY"] = "1" - return current_env - - -def prepare_deepspeed_cmd_env(args: argparse.Namespace) -> tuple[list[str], dict[str, str]]: - """ - Prepares and returns the command list and an environment with the correct DeepSpeed environment variables. - """ - # get free port and update configurations - if args.main_process_port == 0: - args.main_process_port = get_free_port() - - elif args.main_process_port is None: - args.main_process_port = 29500 - - num_processes = args.num_processes - num_machines = args.num_machines - main_process_ip = args.main_process_ip - main_process_port = args.main_process_port - cmd = None - - # make sure launcher is not None - if args.deepspeed_multinode_launcher is None: - # set to default pdsh - args.deepspeed_multinode_launcher = DEEPSPEED_MULTINODE_LAUNCHERS[0] - - if num_machines > 1 and args.deepspeed_multinode_launcher != DEEPSPEED_MULTINODE_LAUNCHERS[1]: - cmd = ["deepspeed"] - cmd.extend(["--hostfile", str(args.deepspeed_hostfile)]) - if args.deepspeed_multinode_launcher == "nossh": - if compare_versions("deepspeed", "<", "0.14.5"): - raise ValueError("nossh launcher requires DeepSpeed >= 0.14.5") - cmd.extend(["--node_rank", str(args.machine_rank), "--no_ssh"]) - else: - cmd.extend(["--no_local_rank", "--launcher", str(args.deepspeed_multinode_launcher)]) - if args.deepspeed_exclusion_filter is not None: - cmd.extend( - [ - "--exclude", - str(args.deepspeed_exclusion_filter), - ] - ) - elif args.deepspeed_inclusion_filter is not None: - cmd.extend( - [ - "--include", - str(args.deepspeed_inclusion_filter), - ] - ) - else: - cmd.extend(["--num_gpus", str(args.num_processes // args.num_machines)]) - if main_process_ip: - cmd.extend(["--master_addr", str(main_process_ip)]) - cmd.extend(["--master_port", str(main_process_port)]) - if args.module and args.no_python: - raise ValueError("--module and --no_python cannot be used together") - elif args.module: - cmd.append("--module") - elif args.no_python: - cmd.append("--no_python") - cmd.append(args.training_script) - cmd.extend(args.training_script_args) - elif num_machines > 1 and args.deepspeed_multinode_launcher == DEEPSPEED_MULTINODE_LAUNCHERS[1]: - args.nproc_per_node = str(num_processes // num_machines) - args.nnodes = str(num_machines) - args.node_rank = int(args.machine_rank) - if getattr(args, "same_network", False): - args.master_addr = str(main_process_ip) - args.master_port = str(main_process_port) - else: - args.rdzv_endpoint = f"{main_process_ip}:{main_process_port}" - else: - args.nproc_per_node = str(num_processes) - if main_process_port is not None: - args.master_port = str(main_process_port) - - # only need to check port availability in main process, in case we have to start multiple launchers on the same machine - # for some reasons like splitting log files. - need_port_check = num_machines <= 1 or int(args.machine_rank) == 0 - if need_port_check and is_port_in_use(main_process_port): - if num_machines <= 1: - args.standalone = True - warnings.warn( - f"Port `{main_process_port}` is already in use. " - "Accelerate will attempt to launch in a standalone-like mode by finding an open port automatically for this session. " - "If this current attempt fails, or for more control in future runs, please specify a different port " - "(e.g., `--main_process_port `) or use `--main_process_port 0` for automatic selection " - "in your launch command or Accelerate config file." - ) - else: - raise ConnectionError( - f"Tried to launch distributed communication on port `{main_process_port}`, but another process is utilizing it. " - "Please specify a different port (such as using the `--main_process_port` flag or specifying a different `main_process_port` in your config file)" - " and rerun your script. To automatically use the next open port (on a single node), you can set this to `0`." - ) - - if args.module and args.no_python: - raise ValueError("--module and --no_python cannot be used together") - elif args.module: - args.module = True - elif args.no_python: - args.no_python = True - - current_env = os.environ.copy() - if args.debug: - current_env["ACCELERATE_DEBUG_MODE"] = "true" - gpu_ids = getattr(args, "gpu_ids", "all") - if gpu_ids != "all" and args.gpu_ids is not None: - if is_xpu_available(): - current_env["ZE_AFFINITY_MASK"] = gpu_ids - elif is_mlu_available(): - current_env["MLU_VISIBLE_DEVICES"] = gpu_ids - elif is_sdaa_available(): - current_env["SDAA_VISIBLE_DEVICES"] = gpu_ids - elif is_musa_available(): - current_env["MUSA_VISIBLE_DEVICES"] = gpu_ids - elif is_npu_available(): - current_env["ASCEND_RT_VISIBLE_DEVICES"] = gpu_ids - elif is_hpu_available(): - current_env["HABANA_VISIBLE_MODULES"] = gpu_ids - else: - current_env["CUDA_VISIBLE_DEVICES"] = gpu_ids - try: - mixed_precision = PrecisionType(args.mixed_precision.lower()) - except ValueError: - raise ValueError( - f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}." - ) - - current_env["PYTHONPATH"] = env_var_path_add("PYTHONPATH", os.path.abspath(".")) - current_env["ACCELERATE_MIXED_PRECISION"] = str(mixed_precision) - if args.mixed_precision.lower() == "fp8": - if not is_fp8_available(): - raise RuntimeError( - "FP8 is not available on this machine. Please ensure that either Transformer Engine, MSAMP or torchao is installed." - ) - current_env = setup_fp8_env(args, current_env) - current_env["ACCELERATE_CONFIG_DS_FIELDS"] = str(args.deepspeed_fields_from_accelerate_config).lower() - current_env["ACCELERATE_USE_DEEPSPEED"] = "true" - if args.zero_stage is not None: - current_env["ACCELERATE_DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage) - if args.gradient_accumulation_steps is not None: - current_env["ACCELERATE_GRADIENT_ACCUMULATION_STEPS"] = str(args.gradient_accumulation_steps) - if args.gradient_clipping is not None: - current_env["ACCELERATE_GRADIENT_CLIPPING"] = str(args.gradient_clipping).lower() - if args.offload_optimizer_device is not None: - current_env["ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE"] = str(args.offload_optimizer_device).lower() - if args.offload_param_device is not None: - current_env["ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE"] = str(args.offload_param_device).lower() - if args.zero3_init_flag is not None: - current_env["ACCELERATE_DEEPSPEED_ZERO3_INIT"] = str(args.zero3_init_flag).lower() - if args.zero3_save_16bit_model is not None: - current_env["ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL"] = str(args.zero3_save_16bit_model).lower() - if args.deepspeed_config_file is not None: - current_env["ACCELERATE_DEEPSPEED_CONFIG_FILE"] = str(args.deepspeed_config_file) - if args.enable_cpu_affinity: - current_env["ACCELERATE_CPU_AFFINITY"] = "1" - if args.deepspeed_moe_layer_cls_names is not None: - current_env["ACCELERATE_DEEPSPEED_MOE_LAYER_CLS_NAMES"] = str(args.deepspeed_moe_layer_cls_names) - return cmd, current_env - - -def prepare_tpu( - args: argparse.Namespace, current_env: dict[str, str], pod: bool = False -) -> tuple[argparse.Namespace, dict[str, str]]: - """ - Prepares and returns an environment with the correct TPU environment variables. - """ - if args.mixed_precision == "bf16" and is_torch_xla_available(check_is_tpu=True): - if args.downcast_bf16: - current_env["XLA_DOWNCAST_BF16"] = "1" - else: - current_env["XLA_USE_BF16"] = "1" - if args.debug: - current_env["ACCELERATE_DEBUG_MODE"] = "true" - if pod: - # Take explicit args and set them up for XLA - args.vm = args.tpu_vm - args.tpu = args.tpu_name - return args, current_env - - -def _convert_nargs_to_dict(nargs: list[str]) -> dict[str, str]: - if len(nargs) < 0: - return {} - # helper function to infer type for argsparser - - def _infer_type(s): - try: - s = float(s) - - if s // 1 == s: - return int(s) - return s - except ValueError: - return s - - parser = argparse.ArgumentParser() - _, unknown = parser.parse_known_args(nargs) - for index, argument in enumerate(unknown): - if argument.startswith(("-", "--")): - action = None - if index + 1 < len(unknown): # checks if next index would be in list - if unknown[index + 1].startswith(("-", "--")): # checks if next element is an key - # raise an error if element is store_true or store_false - raise ValueError( - "SageMaker doesn’t support argparse actions for `store_true` or `store_false`. Please define explicit types" - ) - else: # raise an error if last element is store_true or store_false - raise ValueError( - "SageMaker doesn’t support argparse actions for `store_true` or `store_false`. Please define explicit types" - ) - # adds argument to parser based on action_store true - if action is None: - parser.add_argument(argument, type=_infer_type) - else: - parser.add_argument(argument, action=action) - - return { - key: (literal_eval(value) if value in ("True", "False") else value) - for key, value in parser.parse_args(nargs).__dict__.items() - } - - -def prepare_sagemager_args_inputs( - sagemaker_config: SageMakerConfig, args: argparse.Namespace -) -> tuple[argparse.Namespace, dict[str, Any]]: - # configure environment - print("Configuring Amazon SageMaker environment") - os.environ["AWS_DEFAULT_REGION"] = sagemaker_config.region - - # configure credentials - if sagemaker_config.profile is not None: - os.environ["AWS_PROFILE"] = sagemaker_config.profile - elif args.aws_access_key_id is not None and args.aws_secret_access_key is not None: - os.environ["AWS_ACCESS_KEY_ID"] = args.aws_access_key_id - os.environ["AWS_SECRET_ACCESS_KEY"] = args.aws_secret_access_key - else: - raise OSError("You need to provide an aws_access_key_id and aws_secret_access_key when not using aws_profile") - - # extract needed arguments - source_dir = os.path.dirname(args.training_script) - if not source_dir: # checks if string is empty - source_dir = "." - entry_point = os.path.basename(args.training_script) - if not entry_point.endswith(".py"): - raise ValueError(f'Your training script should be a python script and not "{entry_point}"') - - print("Converting Arguments to Hyperparameters") - hyperparameters = _convert_nargs_to_dict(args.training_script_args) - - try: - mixed_precision = PrecisionType(args.mixed_precision.lower()) - except ValueError: - raise ValueError( - f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}." - ) - - try: - dynamo_backend = DynamoBackend(args.dynamo_backend.upper()) - except ValueError: - raise ValueError( - f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DynamoBackend.list()}." - ) - - # Environment variables to be set for use during training job - environment = { - "ACCELERATE_USE_SAGEMAKER": "true", - "ACCELERATE_MIXED_PRECISION": str(mixed_precision), - "ACCELERATE_DYNAMO_BACKEND": dynamo_backend.value, - "ACCELERATE_DYNAMO_MODE": args.dynamo_mode, - "ACCELERATE_DYNAMO_USE_FULLGRAPH": str(args.dynamo_use_fullgraph), - "ACCELERATE_DYNAMO_USE_DYNAMIC": str(args.dynamo_use_dynamic), - "ACCELERATE_DYNAMO_USE_REGIONAL_COMPILATION": str(args.dynamo_use_regional_compilation), - "ACCELERATE_SAGEMAKER_DISTRIBUTED_TYPE": sagemaker_config.distributed_type.value, - } - if args.mixed_precision.lower() == "fp8": - if not is_fp8_available(): - raise RuntimeError( - "FP8 is not available on this machine. Please ensure that either Transformer Engine, MSAMP or torchao is installed." - ) - environment = setup_fp8_env(args, environment) - # configure distribution set up - distribution = None - if sagemaker_config.distributed_type == SageMakerDistributedType.DATA_PARALLEL: - distribution = {"smdistributed": {"dataparallel": {"enabled": True}}} - - # configure sagemaker inputs - sagemaker_inputs = None - if sagemaker_config.sagemaker_inputs_file is not None: - print(f"Loading SageMaker Inputs from {sagemaker_config.sagemaker_inputs_file} file") - sagemaker_inputs = {} - with open(sagemaker_config.sagemaker_inputs_file) as file: - for i, line in enumerate(file): - if i == 0: - continue - l = line.split("\t") - sagemaker_inputs[l[0]] = l[1].strip() - print(f"Loaded SageMaker Inputs: {sagemaker_inputs}") - - # configure sagemaker metrics - sagemaker_metrics = None - if sagemaker_config.sagemaker_metrics_file is not None: - print(f"Loading SageMaker Metrics from {sagemaker_config.sagemaker_metrics_file} file") - sagemaker_metrics = [] - with open(sagemaker_config.sagemaker_metrics_file) as file: - for i, line in enumerate(file): - if i == 0: - continue - l = line.split("\t") - metric_dict = { - "Name": l[0], - "Regex": l[1].strip(), - } - sagemaker_metrics.append(metric_dict) - print(f"Loaded SageMaker Metrics: {sagemaker_metrics}") - - # configure session - print("Creating Estimator") - args = { - "image_uri": sagemaker_config.image_uri, - "entry_point": entry_point, - "source_dir": source_dir, - "role": sagemaker_config.iam_role_name, - "transformers_version": sagemaker_config.transformers_version, - "pytorch_version": sagemaker_config.pytorch_version, - "py_version": sagemaker_config.py_version, - "base_job_name": sagemaker_config.base_job_name, - "instance_count": sagemaker_config.num_machines, - "instance_type": sagemaker_config.ec2_instance_type, - "debugger_hook_config": False, - "distribution": distribution, - "hyperparameters": hyperparameters, - "environment": environment, - "metric_definitions": sagemaker_metrics, - } - - if sagemaker_config.additional_args is not None: - args = merge_dicts(sagemaker_config.additional_args, args) - return args, sagemaker_inputs - - -def env_var_path_add(env_var_name, path_to_add): - """ - Extends a path-based environment variable's value with a new path and returns the updated value. It's up to the - caller to set it in os.environ. - """ - paths = [p for p in os.environ.get(env_var_name, "").split(":") if len(p) > 0] - paths.append(str(path_to_add)) - return ":".join(paths) - - -class PrepareForLaunch: - """ - Prepare a function that will launched in a distributed setup. - - Args: - launcher (`Callable`): - The function to launch. - distributed_type ([`~state.DistributedType`]): - The distributed type to prepare for. - debug (`bool`, *optional*, defaults to `False`): - Whether or not this is a debug launch. - """ - - def __init__(self, launcher, distributed_type="NO", debug=False): - self.launcher = launcher - self.distributed_type = DistributedType(distributed_type) - self.debug = debug - - def __call__(self, index, *args): - if self.debug: - world_size = int(os.environ.get("WORLD_SIZE")) - rdv_file = os.environ.get("ACCELERATE_DEBUG_RDV_FILE") - torch.distributed.init_process_group( - "gloo", - rank=index, - store=torch.distributed.FileStore(rdv_file, world_size), - world_size=world_size, - ) - elif self.distributed_type in ( - DistributedType.MULTI_GPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_NPU, - DistributedType.MULTI_XPU, - DistributedType.MULTI_CPU, - ): - # Prepare the environment for torch.distributed - os.environ["LOCAL_RANK"] = str(index) - nproc = int(os.environ.get("NPROC", 1)) - node_rank = int(os.environ.get("NODE_RANK", 0)) - os.environ["RANK"] = str(nproc * node_rank + index) - - os.environ["FORK_LAUNCHED"] = str(1) - self.launcher(*args) diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/megatron_lm.py b/venv/lib/python3.10/site-packages/accelerate/utils/megatron_lm.py deleted file mode 100644 index 9f6d87b2edb13299e0b6b34dd32851f020415242..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/megatron_lm.py +++ /dev/null @@ -1,1424 +0,0 @@ -# Copyright 2022 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 argparse -import math -import os -from abc import ABC -from functools import partial - -import torch -import torch.nn.functional as F -from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss -from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP - -from ..optimizer import AcceleratedOptimizer -from ..scheduler import AcceleratedScheduler -from .imports import is_megatron_lm_available -from .operations import recursively_apply, send_to_device - - -if is_megatron_lm_available(): - from megatron.core import mpu, tensor_parallel - from megatron.core.distributed import DistributedDataParallel as LocalDDP - from megatron.core.distributed import finalize_model_grads - from megatron.core.enums import ModelType - from megatron.core.num_microbatches_calculator import get_num_microbatches - from megatron.core.optimizer import get_megatron_optimizer - from megatron.core.parallel_state import get_tensor_model_parallel_group, get_tensor_model_parallel_src_rank - from megatron.core.pipeline_parallel import get_forward_backward_func - from megatron.core.utils import get_model_config - from megatron.inference.text_generation.communication import broadcast_int_list, broadcast_tensor - from megatron.inference.text_generation.generation import ( - beam_search_and_return_on_first_stage, - generate_tokens_probs_and_return_on_first_stage, - ) - from megatron.legacy.data.dataset_utils import build_train_valid_test_datasets - from megatron.legacy.model import BertModel, Float16Module, GPTModel, T5Model - from megatron.legacy.model.classification import Classification - from megatron.training import ( - get_args, - get_tensorboard_writer, - get_tokenizer, - print_rank_last, - ) - from megatron.training.arguments import ( - _add_data_args, - _add_validation_args, - core_transformer_config_from_args, - parse_args, - validate_args, - ) - from megatron.training.checkpointing import load_args_from_checkpoint, load_checkpoint, save_checkpoint - from megatron.training.global_vars import set_global_variables - from megatron.training.initialize import ( - _compile_dependencies, - _init_autoresume, - _initialize_distributed, - _set_random_seed, - set_jit_fusion_options, - write_args_to_tensorboard, - ) - from megatron.training.tokenizer.tokenizer import _vocab_size_with_padding - from megatron.training.training import ( - build_train_valid_test_data_iterators, - get_optimizer_param_scheduler, - num_floating_point_operations, - setup_model_and_optimizer, - train_step, - training_log, - ) - from megatron.training.utils import ( - average_losses_across_data_parallel_group, - calc_params_l2_norm, - get_ltor_masks_and_position_ids, - unwrap_model, - ) - - -# model utilities -def model_provider_func(pre_process=True, post_process=True, add_encoder=True, add_decoder=True): - """Build the model.""" - args = get_args() - mode = "pre-training" if args.pretraining_flag else "fine-tuning" - if args.rank == 0: - print(f"Building {args.model_type_name} model in the {mode} mode.") - print( - "The Megatron LM model weights are initialized at random in `accelerator.prepare`. " - "Please use `accelerator.load_checkpoint` to load a pre-trained checkpoint matching the distributed setup." - ) - config = core_transformer_config_from_args(args) - if args.model_type_name == "bert": - if args.pretraining_flag: - num_tokentypes = 2 if args.bert_binary_head else 0 - model = BertModel( - config=config, - num_tokentypes=num_tokentypes, - add_binary_head=args.bert_binary_head, - parallel_output=True, - pre_process=pre_process, - post_process=post_process, - ) - else: - model = Classification( - config=config, - num_classes=args.num_labels, - num_tokentypes=2, - pre_process=pre_process, - post_process=post_process, - ) - elif args.model_type_name == "gpt": - model = GPTModel( - config=config, - num_tokentypes=0, - parallel_output=True, - pre_process=pre_process, - post_process=post_process, - ) - elif args.model_type_name == "t5": - model = T5Model( - config=config, - num_tokentypes=0, - parallel_output=True, - pre_process=pre_process, - post_process=post_process, - add_encoder=add_encoder, - add_decoder=add_decoder, - ) - else: - raise ValueError(f"Unsupported model type: {args.model_type_name}") - return model - - -def prepare_model_optimizer_scheduler(accelerator): - accelerator.print("Preparing model optimizer scheduler") - args = get_args() - if accelerator.state.megatron_lm_plugin.custom_prepare_model_function is not None: - if accelerator.state.megatron_lm_plugin.custom_model_provider_function is None: - raise ValueError( - "You must provide a `custom_model_provider_function` when using a `custom_prepare_model_function`." - ) - custom_model_provider_func = accelerator.state.megatron_lm_plugin.custom_model_provider_function - model = accelerator.state.megatron_lm_plugin.custom_prepare_model_function(custom_model_provider_func) - optimizer = prepare_optimizer(accelerator, model) - scheduler = prepare_scheduler(accelerator, optimizer, scheduler=None) - else: - model_type = ModelType.encoder_or_decoder - if args.model_type_name == "t5": - model_type = ModelType.encoder_and_decoder - model_provider_func_ = model_provider_func - if accelerator.state.megatron_lm_plugin.custom_model_provider_function is not None: - model_provider_func_ = accelerator.state.megatron_lm_plugin.custom_model_provider_function - (model, optimizer, scheduler) = setup_model_and_optimizer( - model_provider_func_, - model_type, - no_wd_decay_cond=args.no_wd_decay_cond, - scale_lr_cond=args.scale_lr_cond, - lr_mult=args.lr_mult, - ) - args.model_len = len(model) - return model, optimizer, scheduler - - -# dataloader utilities -class MegatronLMDummyDataLoader: - """ - Dummy dataloader presents model parameters or param groups, this is primarily used to follow conventional training - - Args: - **dataset_kwargs: Megatron data arguments. - """ - - def __init__(self, **dataset_kwargs): - parser = argparse.ArgumentParser() - parser = _add_data_args(parser) - parser = _add_validation_args(parser) - data_args = parser.parse_known_args() - self.dataset_args = vars(data_args[0]) - self.dataset_args.update(dataset_kwargs) - self.dataset_args["megatron_dataset_flag"] = True - - def set_megatron_data_args(self): - args = get_args() - for key, value in self.dataset_args.items(): - old_value = getattr(args, key, "") - if old_value != value: - print( - f"WARNING: MegatronLMDummyDataLoader overriding arguments for {key}:{old_value} with {key}:{value}" - ) - setattr(args, key, value) - - def get_train_valid_test_datasets_provider(self, accelerator): - def train_valid_test_datasets_provider(train_val_test_num_samples): - """Build train, valid, and test datasets.""" - args = get_args() - dataset_args = { - "data_prefix": args.data_path if isinstance(args.data_path, (list, tuple)) else [args.data_path], - "splits_string": args.split, - "train_valid_test_num_samples": train_val_test_num_samples, - "seed": args.seed, - } - if args.model_type_name == "bert": - dataset_args.update( - { - "max_seq_length": args.seq_length, - "binary_head": args.bert_binary_head, - } - ) - elif args.model_type_name == "gpt": - dataset_args.update( - { - "max_seq_length": args.seq_length, - } - ) - elif args.model_type_name == "t5": - dataset_args.update( - { - "max_seq_length": args.encoder_seq_length, - "max_seq_length_dec": args.decoder_seq_length, - "dataset_type": "t5", - } - ) - else: - raise ValueError(f"Unsupported model type: {args.model_type_name}") - train_ds, valid_ds, test_ds = build_train_valid_test_datasets(**dataset_args) - return train_ds, valid_ds, test_ds - - if accelerator.state.megatron_lm_plugin.custom_megatron_datasets_provider_function is not None: - return accelerator.state.megatron_lm_plugin.custom_megatron_datasets_provider_function - try: - args = get_args() - # Use '--no-use-pep517 -e' to pip install nvidia's megatron from source - if args.model_type_name == "bert": - from pretrain_bert import train_valid_test_datasets_provider - - train_valid_test_datasets_provider.is_distributed = True - return train_valid_test_datasets_provider - elif args.model_type_name == "gpt": - from pretrain_gpt import train_valid_test_datasets_provider - - train_valid_test_datasets_provider.is_distributed = True - return train_valid_test_datasets_provider - elif args.model_type_name == "t5": - from pretrain_t5 import train_valid_test_datasets_provider - - train_valid_test_datasets_provider.is_distributed = True - return train_valid_test_datasets_provider - except ImportError: - pass - return train_valid_test_datasets_provider - - def build_train_valid_test_data_iterators(self, accelerator): - args = get_args() - - train_valid_test_dataset_provider = self.get_train_valid_test_datasets_provider(accelerator) - if args.virtual_pipeline_model_parallel_size is not None: - train_data_iterator = [] - valid_data_iterator = [] - test_data_iterator = [] - for i in range(getattr(args, "model_len", 0)): - mpu.set_virtual_pipeline_model_parallel_rank(i) - iterators = build_train_valid_test_data_iterators(train_valid_test_dataset_provider) - train_data_iterator.append(iterators[0]) - valid_data_iterator.append(iterators[1]) - test_data_iterator.append(iterators[2]) - else: - train_data_iterator, valid_data_iterator, test_data_iterator = build_train_valid_test_data_iterators( - train_valid_test_dataset_provider - ) - - return train_data_iterator, valid_data_iterator, test_data_iterator - - -def _handle_megatron_data_iterator(accelerator, data_iterator): - class DummyMegatronDataloader: - def __iter__(self): - return self - - def __next__(self): - return {} - - is_data_iterator_empty = data_iterator is None - is_src_data_iterator_empty = torch.tensor(is_data_iterator_empty, dtype=torch.bool, device=accelerator.device) - torch.distributed.broadcast( - is_src_data_iterator_empty, get_tensor_model_parallel_src_rank(), group=get_tensor_model_parallel_group() - ) - if not is_src_data_iterator_empty and is_data_iterator_empty: - return DummyMegatronDataloader() - return data_iterator - - -def prepare_data_loader(accelerator, dataloader): - accelerator.print("Preparing dataloader") - args = get_args() - if not args.megatron_dataset_flag: - from ..data_loader import _PYTORCH_DATALOADER_KWARGS, prepare_data_loader - - micro_batch_size = args.micro_batch_size * args.num_micro_batches - kwargs = {k: getattr(dataloader, k, _PYTORCH_DATALOADER_KWARGS[k]) for k in _PYTORCH_DATALOADER_KWARGS} - if kwargs["batch_size"] is None: - if isinstance(kwargs["sampler"], torch.utils.data.BatchSampler): - kwargs["sampler"].batch_size = micro_batch_size - else: - del kwargs["sampler"] - del kwargs["shuffle"] - del kwargs["batch_size"] - kwargs["batch_sampler"].batch_size = micro_batch_size - else: - del kwargs["batch_sampler"] - kwargs["batch_size"] = micro_batch_size - - dataloader = torch.utils.data.DataLoader(dataloader.dataset, **kwargs) - # split_batches: - # Megatron only needs to fetch different data between different dp groups, - # and does not need to split the data within the dp group. - return prepare_data_loader( - dataloader, - accelerator.device, - num_processes=mpu.get_data_parallel_world_size(), - process_index=mpu.get_data_parallel_rank(), - split_batches=False, - put_on_device=True, - rng_types=accelerator.rng_types.copy(), - dispatch_batches=accelerator.dispatch_batches, - ) - else: - if args.consumed_samples is not None: - ( - args.consumed_train_samples, - args.consumed_valid_samples, - args.consumed_test_samples, - ) = args.consumed_samples - else: - args.consumed_train_samples, args.consumed_valid_samples, args.consumed_test_samples = 0, 0, 0 - args.micro_batch_size = args.micro_batch_size * args.num_micro_batches - # In order to be compatible with data in transform format, - # it needs to increase the size of mbs first, - # and then split the large batch data into some mbs. - ( - train_data_iterator, - valid_data_iterator, - test_data_iterator, - ) = dataloader.build_train_valid_test_data_iterators(accelerator) - args.micro_batch_size = args.micro_batch_size // args.num_micro_batches - - train_data_iterator = _handle_megatron_data_iterator( - accelerator=accelerator, data_iterator=train_data_iterator - ) - valid_data_iterator = _handle_megatron_data_iterator( - accelerator=accelerator, data_iterator=valid_data_iterator - ) - test_data_iterator = _handle_megatron_data_iterator(accelerator=accelerator, data_iterator=test_data_iterator) - - return train_data_iterator, valid_data_iterator, test_data_iterator - - -# optimizer utilities -class MegatronLMOptimizerWrapper(AcceleratedOptimizer): - def __init__(self, optimizer): - super().__init__(optimizer, device_placement=False, scaler=None) - - def zero_grad(self, set_to_none=None): - pass # `model(**batch)` is doing that automatically. Therefore, its implementation is not needed - - def step(self): - pass # `model(**batch)` is doing that automatically. Therefore, its implementation is not needed - - @property - def step_was_skipped(self): - """Whether or not the optimizer step was done, or skipped because of gradient overflow.""" - return self.optimizer.skipped_iter - - -def prepare_optimizer(accelerator, model): - accelerator.print("Preparing optimizer") - args = get_args() - return get_megatron_optimizer(model, args.no_wd_decay_cond, args.scale_lr_cond, args.lr_mult) - - -# scheduler utilities -class MegatronLMDummyScheduler: - """ - Dummy scheduler presents model parameters or param groups, this is primarily used to follow conventional training - loop when scheduler config is specified in the deepspeed config file. - - Args: - optimizer (`torch.optim.optimizer.Optimizer`): - The optimizer to wrap. - total_num_steps (int): - Total number of steps. - warmup_num_steps (int): - Number of steps for warmup. - **kwargs (additional keyword arguments, *optional*): - Other arguments. - """ - - def __init__(self, optimizer, total_num_steps=None, warmup_num_steps=0, **kwargs): - self.optimizer = optimizer - self.total_num_steps = total_num_steps - self.warmup_num_steps = warmup_num_steps - self.kwargs = kwargs - - -class MegatronLMSchedulerWrapper(AcceleratedScheduler): - def __init__(self, scheduler, optimizers): - super().__init__(scheduler, optimizers) - - def step(self, *args, **kwargs): - return # `model(**batch)` is doing that automatically. Therefore, its implementation is not needed - - -def prepare_scheduler(accelerator, optimizer, scheduler): - accelerator.print("Preparing scheduler") - scheduler = get_optimizer_param_scheduler(optimizer) - return scheduler - - -class AbstractTrainStep(ABC): - """Abstract class for batching, forward pass and loss handler.""" - - def __init__(self, name): - super().__init__() - self.name = name - - def get_batch_func(self, accelerator, megatron_dataset_flag): - pass - - def get_forward_step_func(self): - pass - - def get_loss_func(self, accelerator): - pass - - -class BertTrainStep(AbstractTrainStep): - """ - Bert train step class. - - Args: - args (`argparse.Namespace`): Megatron-LM arguments. - """ - - def __init__(self, accelerator, args): - super().__init__("BertTrainStep") - self.get_batch = self.get_batch_func(accelerator, args.megatron_dataset_flag) - self.loss_func = self.get_loss_func(accelerator, args.pretraining_flag, args.num_labels) - self.forward_step = self.get_forward_step_func(args.pretraining_flag, args.bert_binary_head) - if not args.model_return_dict: - self.model_output_class = None - else: - from transformers.modeling_outputs import SequenceClassifierOutput - - self.model_output_class = SequenceClassifierOutput - - def get_batch_func(self, accelerator, megatron_dataset_flag): - def get_batch_megatron(data_iterator): - """Build the batch.""" - - # Items and their type. - keys = ["text", "types", "labels", "is_random", "loss_mask", "padding_mask"] - datatype = torch.int64 - - # Broadcast data. - if data_iterator is not None: - data = next(data_iterator) - else: - data = None - data_b = tensor_parallel.broadcast_data(keys, data, datatype) - - # Unpack. - tokens = data_b["text"].long() - types = data_b["types"].long() - sentence_order = data_b["is_random"].long() - loss_mask = data_b["loss_mask"].float() - lm_labels = data_b["labels"].long() - padding_mask = data_b["padding_mask"].long() - - return tokens, types, sentence_order, loss_mask, lm_labels, padding_mask - - def get_batch_transformer(data_iterator): - """Build the batch.""" - data = next(data_iterator) - data = send_to_device(data, torch.cuda.current_device()) - - # Unpack. - tokens = data["input_ids"].long() - padding_mask = data["attention_mask"].long() - if "token_type_ids" in data: - types = data["token_type_ids"].long() - else: - types = None - if "labels" in data: - lm_labels = data["labels"].long() - loss_mask = (data["labels"] != -100).to(torch.float) - else: - lm_labels = None - loss_mask = None - if "next_sentence_label" in data: - sentence_order = data["next_sentence_label"].long() - else: - sentence_order = None - - return tokens, types, sentence_order, loss_mask, lm_labels, padding_mask - - if accelerator.state.megatron_lm_plugin.custom_get_batch_function is not None: - return accelerator.state.megatron_lm_plugin.custom_get_batch_function - if megatron_dataset_flag: - try: - # Use '--no-use-pep517 -e' to pip install nvidia's megatron from source - from pretrain_bert import get_batch - - return get_batch - except ImportError: - pass - return get_batch_megatron - else: - return get_batch_transformer - - def get_loss_func(self, accelerator, pretraining_flag, num_labels): - def loss_func_pretrain(loss_mask, sentence_order, output_tensor): - lm_loss_, sop_logits = output_tensor - - lm_loss_ = lm_loss_.float() - loss_mask = loss_mask.float() - lm_loss = torch.sum(lm_loss_.view(-1) * loss_mask.reshape(-1)) / loss_mask.sum() - - if sop_logits is not None: - sop_loss = F.cross_entropy(sop_logits.view(-1, 2).float(), sentence_order.view(-1), ignore_index=-1) - sop_loss = sop_loss.float() - loss = lm_loss + sop_loss - averaged_losses = average_losses_across_data_parallel_group([lm_loss, sop_loss]) - return loss, {"lm loss": averaged_losses[0], "sop loss": averaged_losses[1]} - - else: - loss = lm_loss - averaged_losses = average_losses_across_data_parallel_group([lm_loss]) - return loss, {"lm loss": averaged_losses[0]} - - def loss_func_finetune(labels, logits): - if num_labels == 1: - # We are doing regression - loss_fct = MSELoss() - loss = loss_fct(logits.view(-1), labels.view(-1)) - elif self.num_labels > 1 and (labels.dtype in (torch.long, torch.int)): - loss_fct = CrossEntropyLoss() - loss = loss_fct(logits.view(-1, num_labels), labels.view(-1)) - else: - loss_fct = BCEWithLogitsLoss() - loss = loss_fct(logits, labels) - averaged_losses = average_losses_across_data_parallel_group([loss]) - return loss, {"loss": averaged_losses[0]} - - if accelerator.state.megatron_lm_plugin.custom_loss_function is not None: - return accelerator.state.megatron_lm_plugin.custom_loss_function - if pretraining_flag: - return loss_func_pretrain - else: - return loss_func_finetune - - def get_forward_step_func(self, pretraining_flag, bert_binary_head): - def forward_step(data_iterator, model): - """Forward step.""" - tokens, types, sentence_order, loss_mask, labels, padding_mask = self.get_batch(data_iterator) - if not bert_binary_head: - types = None - # Forward pass through the model. - if pretraining_flag: - output_tensor = model(tokens, padding_mask, tokentype_ids=types, lm_labels=labels) - return output_tensor, partial(self.loss_func, loss_mask, sentence_order) - else: - logits = model(tokens, padding_mask, tokentype_ids=types) - return logits, partial(self.loss_func, labels) - - return forward_step - - -class GPTTrainStep(AbstractTrainStep): - """ - GPT train step class. - - Args: - args (`argparse.Namespace`): Megatron-LM arguments. - """ - - def __init__(self, accelerator, args): - super().__init__("GPTTrainStep") - self.get_batch = self.get_batch_func(accelerator, args.megatron_dataset_flag) - self.loss_func = self.get_loss_func(accelerator) - self.forward_step = self.get_forward_step_func() - self.eod_token = args.padded_vocab_size - 1 - if args.vocab_file is not None: - tokenizer = get_tokenizer() - self.eod_token = tokenizer.eod - self.reset_position_ids = args.reset_position_ids - self.reset_attention_mask = args.reset_attention_mask - self.eod_mask_loss = args.eod_mask_loss - if not args.model_return_dict: - self.model_output_class = None - else: - from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions - - self.model_output_class = CausalLMOutputWithCrossAttentions - - def get_batch_func(self, accelerator, megatron_dataset_flag): - def get_batch_megatron(data_iterator): - """Generate a batch""" - # Items and their type. - keys = ["text"] - datatype = torch.int64 - - # Broadcast data. - if data_iterator is not None: - data = next(data_iterator) - else: - data = None - data_b = tensor_parallel.broadcast_data(keys, data, datatype) - - # Unpack. - tokens_ = data_b["text"].long() - labels = tokens_[:, 1:].contiguous() - tokens = tokens_[:, :-1].contiguous() - - # Get the masks and position ids. - attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids( - tokens, self.eod_token, self.reset_position_ids, self.reset_attention_mask, self.eod_mask_loss - ) - - return tokens, labels, loss_mask, attention_mask, position_ids - - def get_batch_transformer(data_iterator): - data = next(data_iterator) - data = {"input_ids": data["input_ids"]} - data = send_to_device(data, torch.cuda.current_device()) - - tokens_ = data["input_ids"].long() - padding = torch.zeros((tokens_.shape[0], 1), dtype=tokens_.dtype, device=tokens_.device) + self.eod_token - tokens_ = torch.concat([tokens_, padding], dim=1) - labels = tokens_[:, 1:].contiguous() - tokens = tokens_[:, :-1].contiguous() - # Get the masks and position ids. - attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids( - tokens, self.eod_token, self.reset_position_ids, self.reset_attention_mask, True - ) - return tokens, labels, loss_mask, attention_mask, position_ids - - if accelerator.state.megatron_lm_plugin.custom_get_batch_function is not None: - return accelerator.state.megatron_lm_plugin.custom_get_batch_function - if megatron_dataset_flag: - try: - # Use '--no-use-pep517 -e' to pip install nvidia's megatron from source - from pretrain_gpt import get_batch - - return get_batch - except ImportError: - pass - return get_batch_megatron - else: - return get_batch_transformer - - def get_loss_func(self, accelerator): - args = get_args() - - def loss_func(loss_mask, output_tensor): - if args.return_logits: - losses, logits = output_tensor - else: - losses = output_tensor - losses = losses.float() - loss_mask = loss_mask.view(-1).float() - if args.context_parallel_size > 1: - loss = torch.cat([torch.sum(losses.view(-1) * loss_mask).view(1), loss_mask.sum().view(1)]) - torch.distributed.all_reduce(loss, group=mpu.get_context_parallel_group()) - loss = loss[0] / loss[1] - else: - loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum() - - # Check individual rank losses are not NaN prior to DP all-reduce. - if args.check_for_nan_in_loss_and_grad: - global_rank = torch.distributed.get_rank() - assert not loss.isnan(), ( - f"Rank {global_rank}: found NaN in local forward loss calculation. " - f"Device: {torch.cuda.current_device()}, node: {os.uname()[1]}" - ) - - # Reduce loss for logging. - averaged_loss = average_losses_across_data_parallel_group([loss]) - - output_dict = {"lm loss": averaged_loss[0]} - if args.return_logits: - output_dict.update({"logits": logits}) - return loss, output_dict - - if accelerator.state.megatron_lm_plugin.custom_loss_function is not None: - return accelerator.state.megatron_lm_plugin.custom_loss_function - return loss_func - - def get_forward_step_func(self): - def forward_step(data_iterator, model): - """Forward step.""" - # Get the batch. - tokens, labels, loss_mask, attention_mask, position_ids = self.get_batch(data_iterator) - output_tensor = model(tokens, position_ids, attention_mask, labels=labels) - - return output_tensor, partial(self.loss_func, loss_mask) - - return forward_step - - -class T5TrainStep(AbstractTrainStep): - """ - T5 train step class. - - Args: - args (`argparse.Namespace`): Megatron-LM arguments. - """ - - def __init__(self, accelerator, args): - super().__init__("T5TrainStep") - self.get_batch = self.get_batch_func(accelerator, args.megatron_dataset_flag) - self.loss_func = self.get_loss_func(accelerator) - self.forward_step = self.get_forward_step_func() - if not args.model_return_dict: - self.model_output_class = None - else: - from transformers.modeling_outputs import Seq2SeqLMOutput - - self.model_output_class = Seq2SeqLMOutput - - @staticmethod - def attn_mask_postprocess(attention_mask): - # We create a 3D attention mask from a 2D tensor mask. - # [b, 1, s] - attention_mask_b1s = attention_mask.unsqueeze(1) - # [b, s, 1] - attention_mask_bs1 = attention_mask.unsqueeze(2) - # [b, s, s] - attention_mask_bss = attention_mask_b1s * attention_mask_bs1 - # Convert attention mask to binary: - extended_attention_mask = attention_mask_bss < 0.5 - return extended_attention_mask - - @staticmethod - def get_decoder_mask(seq_length, device): - attention_mask = torch.tril(torch.ones((1, seq_length, seq_length), device=device)) - attention_mask = attention_mask < 0.5 - return attention_mask - - @staticmethod - def get_enc_dec_mask(attention_mask, dec_seq_length, device): - batch_size, _ = attention_mask.shape - # We create a 3D attention mask from a 2D tensor mask. - # [b, 1, s] - attention_mask_b1s = attention_mask.unsqueeze(1) - # [b, s, 1] - attention_mask_bs1 = torch.ones((batch_size, dec_seq_length, 1), device=device) - attention_mask_bss = attention_mask_bs1 * attention_mask_b1s - extended_attention_mask = attention_mask_bss < 0.5 - return extended_attention_mask - - def get_batch_func(self, accelerator, megatron_dataset_flag): - def get_batch_megatron(data_iterator): - """Build the batch.""" - - keys = ["text_enc", "text_dec", "labels", "loss_mask", "enc_mask", "dec_mask", "enc_dec_mask"] - datatype = torch.int64 - - # Broadcast data. - if data_iterator is not None: - data = next(data_iterator) - else: - data = None - data_b = tensor_parallel.broadcast_data(keys, data, datatype) - - # Unpack. - tokens_enc = data_b["text_enc"].long() - tokens_dec = data_b["text_dec"].long() - labels = data_b["labels"].long() - loss_mask = data_b["loss_mask"].float() - - enc_mask = data_b["enc_mask"] < 0.5 - dec_mask = data_b["dec_mask"] < 0.5 - enc_dec_mask = data_b["enc_dec_mask"] < 0.5 - - return tokens_enc, tokens_dec, loss_mask, labels, enc_mask, dec_mask, enc_dec_mask - - def get_batch_transformer(data_iterator): - """Build the batch.""" - data = next(data_iterator) - data = send_to_device(data, torch.cuda.current_device()) - - tokens_enc = data["input_ids"].long() - labels = data["labels"].long() - loss_mask = (labels != -100).to(torch.float) - if "decoder_input_ids" in data: - tokens_dec = data["decoder_input_ids"].long() - else: - tokens_dec = labels.new_zeros(labels.shape, device=labels.device, dtype=torch.long) - tokens_dec[..., 1:] = labels[..., :-1].clone() - tokens_dec[..., 0] = 0 - tokens_dec.masked_fill_(tokens_dec == -100, 0) - enc_mask = T5TrainStep.attn_mask_postprocess(data["attention_mask"].long()) - dec_mask = T5TrainStep.get_decoder_mask(tokens_dec.shape[1], tokens_dec.device) - enc_dec_mask = T5TrainStep.get_enc_dec_mask( - data["attention_mask"].long(), tokens_dec.shape[1], tokens_dec.device - ) - - return tokens_enc, tokens_dec, loss_mask, labels, enc_mask, dec_mask, enc_dec_mask - - if accelerator.state.megatron_lm_plugin.custom_get_batch_function is not None: - return accelerator.state.megatron_lm_plugin.custom_get_batch_function - if megatron_dataset_flag: - try: - # Use '--no-use-pep517 -e' to pip install nvidia's megatron from source - from pretrain_t5 import get_batch - - return get_batch - except ImportError: - pass - return get_batch_megatron - else: - return get_batch_transformer - - def get_loss_func(self, accelerator): - def loss_func(loss_mask, output_tensor): - lm_loss_ = output_tensor.float() - lm_loss = torch.sum(lm_loss_.view(-1) * loss_mask.reshape(-1)) / loss_mask.sum() - - loss = lm_loss - averaged_losses = average_losses_across_data_parallel_group([lm_loss]) - - return loss, {"lm loss": averaged_losses[0]} - - if accelerator.state.megatron_lm_plugin.custom_loss_function is not None: - return accelerator.state.megatron_lm_plugin.custom_loss_function - return loss_func - - def get_forward_step_func(self): - def forward_step(data_iterator, model): - """Forward step.""" - # Get the batch. - tokens_enc, tokens_dec, loss_mask, lm_labels, enc_mask, dec_mask, enc_dec_mask = self.get_batch( - data_iterator - ) - # Forward model lm_labels - output_tensor = model( - tokens_enc, tokens_dec, enc_mask, dec_mask, enc_dec_mask, tokentype_ids=None, lm_labels=lm_labels - ) - - return output_tensor, partial(self.loss_func, loss_mask) - - return forward_step - - -def finish_mpu_init(): - # torch.distributed initialization - args = get_args() - # Pytorch distributed. - _initialize_distributed() - - # Random seeds for reproducibility. - if args.rank == 0: - print(f"> setting random seeds to {args.seed} ...") - _set_random_seed(args.seed, args.data_parallel_random_init) - - -# intialize megatron setup -def initialize(accelerator, extra_args_provider=None, args_defaults={}): - accelerator.print("Initializing Megatron-LM") - assert torch.cuda.is_available(), "Megatron requires CUDA." - - # Parse arguments - args = parse_args(extra_args_provider, ignore_unknown_args=True) - - # Set defaults - for key, value in args_defaults.items(): - if getattr(args, key, None) is not None: - if args.rank == 0: - print( - f"WARNING: overriding default arguments for {key}:{getattr(args, key)} with {key}:{value}", - flush=True, - ) - setattr(args, key, value) - - if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False): - assert args.load is not None, "--use-checkpoints-args requires --load argument" - load_args_from_checkpoint(args) - - validate_args(args) - - # set global args, build tokenizer, and set adlr-autoresume, - # tensorboard-writer, and timers. - set_global_variables(args) - - # Megatron's MPU is the master. Complete initialization right away. - finish_mpu_init() - - # Autoresume. - _init_autoresume() - - # Compile dependencies. - _compile_dependencies() - - # Set pytorch JIT layer fusion options and warmup JIT functions. - set_jit_fusion_options() - args = get_args() - if getattr(args, "padded_vocab_size", None) is None: - args.padded_vocab_size = _vocab_size_with_padding(args.orig_vocab_size, args) - if args.model_type_name == "bert" and args.pretraining_flag and args.num_labels == 2: - args.bert_binary_head = True - else: - args.bert_binary_head = False - args.iteration = 0 - - -class MegatronEngine(torch.nn.Module): - """ - Megatron-LM model wrapper - - Args: - accelerator (:class:`~accelerate.Accelerator`): The accelerator object to use. - model: Megatron-LM model - optimizer: Megatron-LM optimizer - lr_scheduler: Megatron-LM lr scheduler - """ - - def __init__(self, accelerator, model, optimizer, scheduler): - super().__init__() - self.module = model - self.base_model = model[0] - self.optimizer = optimizer - self.scheduler = scheduler - args = get_args() - if accelerator.state.megatron_lm_plugin.custom_train_step_class is not None: - self.train_step_handler = accelerator.state.megatron_lm_plugin.custom_train_step_class( - args, **accelerator.state.megatron_lm_plugin.custom_train_step_kwargs - ) - elif args.model_type_name == "bert": - self.train_step_handler = BertTrainStep(accelerator, args) - elif args.model_type_name == "gpt": - self.train_step_handler = GPTTrainStep(accelerator, args) - elif args.model_type_name == "t5": - self.train_step_handler = T5TrainStep(accelerator, args) - else: - raise ValueError(f"Unsupported model type: {args.model_type_name}") - self.optimizer.skipped_iter = False - - # Tracking loss. - self.total_loss_dict = {} - self.eval_total_loss_dict = {} - self.iteration = 0 - self.report_memory_flag = True - self.num_floating_point_operations_so_far = 0 - self.module_config = None - if args.tensorboard_dir is not None: - write_args_to_tensorboard() - - def get_module_config(self): - args = get_args() - config = get_model_config(self.module[0]) - # Setup some training config params - config.grad_scale_func = self.optimizer.scale_loss - if isinstance(self.module[0], LocalDDP) and args.overlap_grad_reduce: - assert config.no_sync_func is None, ( - "When overlap_grad_reduce is True, config.no_sync_func must be None; " - "a custom no_sync_func is not supported when overlapping grad-reduce" - ) - config.no_sync_func = [model_chunk.no_sync for model_chunk in self.module] - if len(self.module) == 1: - config.no_sync_func = config.no_sync_func[0] - if args.delay_grad_reduce: - config.grad_sync_func = [model_chunk.start_grad_sync for model_chunk in self.module] - if len(self.module) == 1: - config.grad_sync_func = config.grad_sync_func[0] - if args.overlap_param_gather and args.delay_param_gather: - config.param_sync_func = [ - lambda x: self.optimizer.finish_param_sync(model_index, x) for model_index in range(len(self.module)) - ] - if len(self.module) == 1: - config.param_sync_func = config.param_sync_func[0] - config.finalize_model_grads_func = finalize_model_grads - return config - - def train(self): - for model_module in self.module: - model_module.train() - - if self.module_config is None: - self.module_config = self.get_module_config() - - self.log_eval_results() - - def eval(self): - for model_module in self.module: - model_module.eval() - - if self.module_config is None: - self.module_config = self.get_module_config() - - def get_batch_data_iterator(self, batch_data): - args = get_args() - data_chunks = [] - if len(batch_data) > 0: - if args.num_micro_batches > 1: - for i in range(0, args.num_micro_batches): - data_chunks.append( - { - k: v[i * args.micro_batch_size : (i + 1) * args.micro_batch_size] - for k, v in batch_data.items() - } - ) - else: - data_chunks = [batch_data] - - if len(self.module) > 1: - batch_data_iterator = ( - [iter(data_chunks) for _ in range(len(self.module))] - if len(batch_data) > 0 - else [None] * len(self.module) - ) - else: - batch_data_iterator = iter(data_chunks) if len(batch_data) > 0 else None - return batch_data_iterator - - def train_step(self, **batch_data): - """ - Training step for Megatron-LM - - Args: - batch_data (:obj:`dict`): The batch data to train on. - """ - - batch_data_iterator = self.get_batch_data_iterator(batch_data) - - loss_reduced, skipped_iter, grad_norm, num_zeros_in_grad = train_step( - forward_step_func=self.train_step_handler.forward_step, - data_iterator=batch_data_iterator, - model=self.module, - optimizer=self.optimizer, - opt_param_scheduler=self.scheduler, - config=self.module_config, - ) - - self.optimizer.skipped_iter = skipped_iter == 1 - - return loss_reduced, skipped_iter, grad_norm, num_zeros_in_grad - - def eval_step(self, **batch_data): - """ - Evaluation step for Megatron-LM - - Args: - batch_data (:obj:`dict`): The batch data to evaluate on. - """ - - args = get_args() - batch_data_iterator = self.get_batch_data_iterator(batch_data) - forward_backward_func = get_forward_backward_func() - loss_dicts = forward_backward_func( - forward_step_func=self.train_step_handler.forward_step, - data_iterator=batch_data_iterator, - model=self.module, - num_microbatches=get_num_microbatches(), - seq_length=args.seq_length, - micro_batch_size=args.micro_batch_size, - forward_only=True, - ) - # Empty unused memory - if args.empty_unused_memory_level >= 1: - torch.cuda.empty_cache() - - args.consumed_valid_samples += ( - mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches() - ) - - if mpu.is_pipeline_last_stage(ignore_virtual=True): - # Average loss across microbatches. - loss_reduced = {} - for key in loss_dicts[0]: - losses_reduced_for_key = [x[key] for x in loss_dicts] - if len(losses_reduced_for_key[0].shape) == 0: - loss_reduced[key] = sum(losses_reduced_for_key) / len(losses_reduced_for_key) - else: - loss_reduced[key] = torch.concat(losses_reduced_for_key) - return loss_reduced - return {} - - def forward(self, **batch_data): - # During training, we use train_step() - # model(**batch_data) performs following operations by delegating it to `self.train_step`: - # 1. Prepare **batch_data for Tendor, Pipeline and Model Parallelism - # 2. Set grad to zero. - # 3. forward pass and backward pass using Pipeline Parallelism - # 4. Empty unused memory. - # 5. Reduce gradients. - # 6. Update parameters. - # 7. Gather params when using Distributed Optimizer (Data Parallelism). - # 8. Update learning rate if scheduler is specified. - # 9. Empty unused memory. - # 10. Average loss across microbatches and across DP ranks. - # - # During evaluation, we use eval_step() - args = get_args() - if self.module[0].training: - loss_dict, skipped_iter, grad_norm, num_zeros_in_grad = self.train_step(**batch_data) - self.iteration += 1 - batch_size = mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches() - args.consumed_train_samples += batch_size - self.num_floating_point_operations_so_far += num_floating_point_operations(args, batch_size) - if args.tensorboard_dir is not None: - # Logging. - loss_scale = self.optimizer.get_loss_scale().item() - params_norm = None - if args.log_params_norm: - params_norm = calc_params_l2_norm(self.model) - self.report_memory_flag = training_log( - loss_dict, - self.total_loss_dict, - self.optimizer.param_groups[0]["lr"], - self.iteration, - loss_scale, - self.report_memory_flag, - skipped_iter, - grad_norm, - params_norm, - num_zeros_in_grad, - ) - else: - loss_dict = self.eval_step(**batch_data) - if args.tensorboard_dir is not None: - for key in loss_dict: - self.eval_total_loss_dict[key] = ( - self.eval_total_loss_dict.get(key, torch.cuda.FloatTensor([0.0])) + loss_dict[key] - ) - self.eval_total_loss_dict[key + "_num_iters"] = self.eval_total_loss_dict.get( - key + "_num_iters", torch.cuda.FloatTensor([0.0]) - ) + torch.cuda.FloatTensor([1.0]) - - loss = torch.tensor(0.0, device=torch.cuda.current_device()) - for key in loss_dict: - if len(loss_dict[key].shape) == 0: - loss += loss_dict[key] - - logits = None - if "logits" in loss_dict: - logits = loss_dict["logits"] - if self.train_step_handler.model_output_class is not None: - return self.train_step_handler.model_output_class(loss=loss, logits=logits) - return loss - - def log_eval_results(self): - args = get_args() - if args.tensorboard_dir is None or self.iteration == 0: - return - args = get_args() - writer = get_tensorboard_writer() - string = f"validation loss at iteration {self.iteration} | " - for key in self.eval_total_loss_dict: - if key.endswith("_num_iters"): - continue - value = self.eval_total_loss_dict[key] / self.eval_total_loss_dict[key + "_num_iters"] - string += f"{key} value: {value} | " - ppl = math.exp(min(20, value.item())) - if args.pretraining_flag: - string += f"{key} PPL: {ppl} | " - if writer: - writer.add_scalar(f"{key} validation", value.item(), self.iteration) - if args.pretraining_flag: - writer.add_scalar(f"{key} validation ppl", ppl, self.iteration) - - length = len(string) + 1 - print_rank_last("-" * length) - print_rank_last(string) - print_rank_last("-" * length) - self.eval_total_loss_dict = {} - - def save_checkpoint(self, output_dir): - self.log_eval_results() - args = get_args() - args.save = output_dir - torch.distributed.barrier() - save_checkpoint( - self.iteration, - self.module, - self.optimizer, - self.scheduler, - num_floating_point_operations_so_far=self.num_floating_point_operations_so_far, - ) - torch.distributed.barrier() - - def load_checkpoint(self, input_dir): - args = get_args() - args.load = input_dir - args.consumed_train_samples = 0 - args.consumed_valid_samples = 0 - torch.distributed.barrier() - iteration, num_floating_point_operations_so_far = load_checkpoint(self.module, self.optimizer, self.scheduler) - torch.distributed.barrier() - self.iteration = iteration - self.num_floating_point_operations_so_far = num_floating_point_operations_so_far - if args.fp16 and self.iteration == 0: - self.optimizer.reload_model_params() - - def megatron_generate( - self, - inputs, - attention_mask=None, - max_length=None, - max_new_tokens=None, - num_beams=None, - temperature=None, - top_k=None, - top_p=None, - length_penalty=None, - **kwargs, - ): - """ - Generate method for GPT2 model. This method is used for inference. Supports both greedy and beam search along - with sampling. Refer the Megatron-LM repo for more details - - Args: - inputs (torch.Tensor): input ids - attention_mask (torch.Tensor, optional): attention mask. Defaults to None. - max_length (int, optional): max length of the generated sequence. Defaults to None. - Either this or max_new_tokens should be provided. - max_new_tokens (int, optional): max number of tokens to be generated. Defaults to None. - Either this or max_length should be provided. - num_beams (int, optional): number of beams to use for beam search. Defaults to None. - temperature (float, optional): temperature for sampling. Defaults to 1.0. - top_k (int, optional): top k tokens to consider for sampling. Defaults to 0.0. - top_p (float, optional): tokens in top p probability are considered for sampling. Defaults to 0.0. - length_penalty (float, optional): length penalty for beam search. Defaults to None. - kwargs: additional key-value arguments - """ - - # checking if required arguments are passed - args = get_args() - if args.model_type_name != "gpt": - raise NotImplementedError("Generate method is not implemented for this model") - - if args.data_parallel_size > 1: - raise ValueError("Generate method requires data parallelism to be 1") - - if args.sequence_parallel: - raise ValueError("Generate method requires sequence parallelism to be False") - - if args.recompute_granularity is not None: - raise ValueError("Checkpoint activations cannot be set for inference") - - if args.vocab_file is None: - raise ValueError("Vocab file is required for inference") - - # Prepare inputs - if max_length is None and max_new_tokens is None: - raise ValueError("`max_length` or `max_new_tokens` are required for inference") - - if temperature is None: - temperature = 1.0 - elif not (0.0 < temperature <= 100.0): - raise ValueError("temperature must be a positive number less than or equal to 100.0") - - if top_k is None: - top_k = 0 - elif not (0 <= top_k <= 1000): - raise ValueError("top_k must be a positive number less than or equal to 1000") - - if top_p is None: - top_p = 0.0 - elif top_p > 0.0 and top_k > 0.0: - raise ValueError("top_p and top_k sampling cannot be set together") - else: - if not (0.0 <= top_p <= 1.0): - raise ValueError("top_p must be less than or equal to 1.0") - - top_p_decay = kwargs.get("top_p_decay", 0.0) - if not (0.0 <= top_p_decay <= 1.0): - raise ValueError("top_p_decay must be less than or equal to 1.0") - - top_p_bound = kwargs.get("top_p_bound", 0.0) - if not (0.0 <= top_p_bound <= 1.0): - raise ValueError("top_p_bound must be less than or equal to 1.0") - - add_BOS = kwargs.get("add_BOS", False) - if not (isinstance(add_BOS, bool)): - raise ValueError("add_BOS must be a boolean") - - beam_width = num_beams - if beam_width is not None: - if not isinstance(beam_width, int): - raise ValueError("beam_width must be an integer") - if beam_width < 1: - raise ValueError("beam_width must be greater than 0") - if inputs.shape[0] > 1: - return "When doing beam_search, batch size must be 1" - - tokenizer = get_tokenizer() - - stop_token = kwargs.get("stop_token", tokenizer.eod) - if stop_token is not None: - if not isinstance(stop_token, int): - raise ValueError("stop_token must be an integer") - - if length_penalty is None: - length_penalty = 1.0 - - sizes_list = None - prompts_tokens_tensor = None - prompts_length_tensor = None - if torch.distributed.get_rank() == 0: - # Get the prompts length. - if attention_mask is None: - prompts_length_tensor = torch.cuda.LongTensor([inputs.shape[1]] * inputs.shape[0]) - else: - prompts_length_tensor = attention_mask.sum(axis=-1).cuda() - - if max_new_tokens is None: - max_new_tokens = max_length - inputs.shape[1] - if max_new_tokens <= 0: - raise ValueError("max_new_tokens must be greater than 0") - - if add_BOS: - max_length = max_new_tokens + inputs.shape[1] + 1 - # making sure that `max_length` is a multiple of 4 to leverage fused kernels - max_length = 4 * math.ceil(max_length / 4) - max_new_tokens = max_length - (inputs.shape[1] + 1) - padding = torch.cuda.LongTensor([[tokenizer.eod] * max_new_tokens] * inputs.shape[0]) - prompts_tokens_tensor = torch.concat( - [torch.unsqueeze(padding[:, 0], axis=-1), inputs.cuda(), padding], axis=-1 - ) - else: - # making sure that `max_length` is a multiple of 4 to leverage fused kernels - max_length = max_new_tokens + inputs.shape[1] - max_length = 4 * math.ceil(max_length / 4) - max_new_tokens = max_length - inputs.shape[1] - padding = torch.cuda.LongTensor([[tokenizer.eod] * max_new_tokens] * inputs.shape[0]) - prompts_tokens_tensor = torch.concat([inputs.cuda(), padding], axis=-1) - - # We need the sizes of these tensors for the boradcast - sizes_list = [ - prompts_tokens_tensor.size(0), # Batch size - prompts_tokens_tensor.size(1), - ] # Sequence length - - # First, broadcast the sizes. - sizes_tensor = broadcast_int_list(2, int_list=sizes_list, rank=0) - - # Now that we have the sizes, we can boradcast the tokens - # and length tensors. - sizes = sizes_tensor.tolist() - context_tokens_tensor = broadcast_tensor(sizes, torch.int64, tensor=prompts_tokens_tensor, rank=0) - context_length_tensor = broadcast_tensor(sizes[0], torch.int64, tensor=prompts_length_tensor, rank=0) - - # Run the inference - random_seed = kwargs.get("random_seed", 0) - torch.random.manual_seed(random_seed) - unwrapped_model = unwrap_model(self.base_model, (torchDDP, LocalDDP, Float16Module)) - if beam_width is not None: - tokens, _ = beam_search_and_return_on_first_stage( - unwrapped_model, - context_tokens_tensor, - context_length_tensor, - beam_width, - stop_token=stop_token, - num_return_gen=1, - length_penalty=length_penalty, - ) - else: - tokens, _, _ = generate_tokens_probs_and_return_on_first_stage( - unwrapped_model, - context_tokens_tensor, - context_length_tensor, - return_output_log_probs=False, - top_k=top_k, - top_p=top_p, - top_p_decay=top_p_decay, - top_p_bound=top_p_bound, - temperature=temperature, - use_eod_token_for_early_termination=True, - ) - return tokens - - -# other utilities -def avg_losses_across_data_parallel_group(losses): - """ - Average losses across data parallel group. - - Args: - losses (List[Tensor]): List of losses to average across data parallel group. - """ - - return average_losses_across_data_parallel_group(losses) - - -def gather_across_data_parallel_groups(tensor): - """ - Recursively gather tensor in a nested list/tuple/dictionary of tensors from data parallel ranks. - - Args: - tensor (nested list/tuple/dictionary of `torch.Tensor`): - The data to gather across data parallel ranks. - - """ - - def _gpu_gather_one(tensor): - if tensor.ndim == 0: - tensor = tensor.clone()[None] - output_tensors = [ - torch.empty_like(tensor) - for _ in range(torch.distributed.get_world_size(group=mpu.get_data_parallel_group())) - ] - torch.distributed.all_gather(output_tensors, tensor, group=mpu.get_data_parallel_group()) - return torch.cat(output_tensors, dim=0) - - return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True) diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/memory.py b/venv/lib/python3.10/site-packages/accelerate/utils/memory.py deleted file mode 100644 index e25649168ab85f189c9b115898611840858389fe..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/memory.py +++ /dev/null @@ -1,207 +0,0 @@ -# Copyright 2022 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. - -""" -A collection of utilities for ensuring that training can always occur. Heavily influenced by the -[toma](https://github.com/BlackHC/toma) library. -""" - -import functools -import gc -import importlib -import inspect -import warnings - -import torch -from packaging import version - -from .imports import ( - is_cuda_available, - is_hpu_available, - is_ipex_available, - is_mlu_available, - is_mps_available, - is_musa_available, - is_npu_available, - is_sdaa_available, - is_xpu_available, -) -from .versions import compare_versions - - -def clear_device_cache(garbage_collection=False): - """ - Clears the device cache by calling `torch.{backend}.empty_cache`. Can also run `gc.collect()`, but do note that - this is a *considerable* slowdown and should be used sparingly. - """ - if garbage_collection: - gc.collect() - - if is_xpu_available(): - torch.xpu.empty_cache() - elif is_mlu_available(): - torch.mlu.empty_cache() - elif is_sdaa_available(): - torch.sdaa.empty_cache() - elif is_musa_available(): - torch.musa.empty_cache() - elif is_npu_available(): - torch.npu.empty_cache() - elif is_mps_available(min_version="2.0"): - torch.mps.empty_cache() - elif is_cuda_available(): - torch.cuda.empty_cache() - elif is_hpu_available(): - # torch.hpu.empty_cache() # not available on hpu as it reserves all device memory for the current process - pass - - -def release_memory(*objects): - """ - Releases memory from `objects` by setting them to `None` and calls `gc.collect()` and `torch.cuda.empty_cache()`. - Returned objects should be reassigned to the same variables. - - Args: - objects (`Iterable`): - An iterable of objects - Returns: - A list of `None` objects to replace `objects` - - Example: - - ```python - >>> import torch - >>> from accelerate.utils import release_memory - - >>> a = torch.ones(1000, 1000).cuda() - >>> b = torch.ones(1000, 1000).cuda() - >>> a, b = release_memory(a, b) - ``` - """ - if not isinstance(objects, list): - objects = list(objects) - for i in range(len(objects)): - objects[i] = None - clear_device_cache(garbage_collection=True) - return objects - - -def should_reduce_batch_size(exception: Exception) -> bool: - """ - Checks if `exception` relates to CUDA out-of-memory, XPU out-of-memory, CUDNN not supported, or CPU out-of-memory - - Args: - exception (`Exception`): - An exception - """ - _statements = [ - " out of memory.", # OOM for CUDA, HIP, XPU - "cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.", # CUDNN SNAFU - "DefaultCPUAllocator: can't allocate memory", # CPU OOM - "FATAL ERROR :: MODULE:PT_DEVMEM Allocation failed", # HPU OOM - ] - if isinstance(exception, RuntimeError) and len(exception.args) == 1: - return any(err in exception.args[0] for err in _statements) - return False - - -def find_executable_batch_size( - function: callable = None, starting_batch_size: int = 128, reduce_batch_size_fn: callable = None -): - """ - A basic decorator that will try to execute `function`. If it fails from exceptions related to out-of-memory or - CUDNN, the batch size is multiplied by 0.9 and passed to `function` - - `function` must take in a `batch_size` parameter as its first argument. - - Args: - function (`callable`, *optional*): - A function to wrap - starting_batch_size (`int`, *optional*): - The batch size to try and fit into memory - - Example: - - ```python - >>> from accelerate.utils import find_executable_batch_size - - - >>> @find_executable_batch_size(starting_batch_size=128) - ... def train(batch_size, model, optimizer): - ... ... - - - >>> train(model, optimizer) - ``` - """ - if function is None: - return functools.partial(find_executable_batch_size, starting_batch_size=starting_batch_size) - - batch_size = starting_batch_size - if reduce_batch_size_fn is None: - - def reduce_batch_size_fn(): - nonlocal batch_size - batch_size = int(batch_size * 0.9) - return batch_size - - def decorator(*args, **kwargs): - nonlocal batch_size - clear_device_cache(garbage_collection=True) - params = list(inspect.signature(function).parameters.keys()) - # Guard against user error - if len(params) < (len(args) + 1): - arg_str = ", ".join([f"{arg}={value}" for arg, value in zip(params[1:], args[1:])]) - raise TypeError( - f"Batch size was passed into `{function.__name__}` as the first argument when called." - f"Remove this as the decorator already does so: `{function.__name__}({arg_str})`" - ) - while True: - if batch_size == 0: - raise RuntimeError("No executable batch size found, reached zero.") - try: - return function(batch_size, *args, **kwargs) - except Exception as e: - if should_reduce_batch_size(e): - clear_device_cache(garbage_collection=True) - batch_size = reduce_batch_size_fn() - else: - raise - - return decorator - - -def get_xpu_available_memory(device_index: int): - if version.parse(torch.__version__).release >= version.parse("2.6").release: - # torch.xpu.mem_get_info API is available starting from PyTorch 2.6 - # It further requires PyTorch built with the SYCL runtime which supports API - # to query available device memory. If not available, exception will be - # raised. Version of SYCL runtime used to build PyTorch is being reported - # with print(torch.version.xpu) and corresponds to the version of Intel DPC++ - # SYCL compiler. First version to support required feature is 20250001. - try: - return torch.xpu.mem_get_info(device_index)[0] - except Exception: - pass - elif is_ipex_available(): - ipex_version = version.parse(importlib.metadata.version("intel_extension_for_pytorch")) - if compare_versions(ipex_version, ">=", "2.5"): - from intel_extension_for_pytorch.xpu import mem_get_info - - return mem_get_info(device_index)[0] - - warnings.warn( - "The XPU `mem_get_info` API is available in IPEX version >=2.5 or PyTorch >=2.6. The current returned available memory is incorrect. Please consider upgrading your IPEX or PyTorch version." - ) - return torch.xpu.max_memory_allocated(device_index) diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/modeling.py b/venv/lib/python3.10/site-packages/accelerate/utils/modeling.py deleted file mode 100644 index 842b0c4a7c7ed0904edf536c20777f8e6b3094f6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/modeling.py +++ /dev/null @@ -1,2199 +0,0 @@ -# Copyright 2022 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 contextlib -import gc -import inspect -import json -import logging -import os -import re -import shutil -import tempfile -import warnings -from collections import OrderedDict, defaultdict -from typing import Optional, Union - -import torch -from torch import distributed as dist -from torch import nn - -from ..state import AcceleratorState -from .constants import SAFE_WEIGHTS_NAME, WEIGHTS_NAME -from .dataclasses import AutocastKwargs, CustomDtype, DistributedType -from .imports import ( - is_hpu_available, - is_mlu_available, - is_mps_available, - is_musa_available, - is_npu_available, - is_peft_available, - is_sdaa_available, - is_torch_xla_available, - is_xpu_available, -) -from .memory import clear_device_cache, get_xpu_available_memory -from .offload import load_offloaded_weight, offload_weight, save_offload_index -from .tqdm import is_tqdm_available, tqdm -from .versions import is_torch_version - - -if is_npu_available(check_device=False): - import torch_npu # noqa: F401 - -if is_mlu_available(check_device=False): - import torch_mlu # noqa: F401 - -if is_sdaa_available(check_device=False): - import torch_sdaa # noqa: F401 - -if is_musa_available(check_device=False): - import torch_musa # noqa: F401 - -from safetensors import safe_open -from safetensors.torch import load_file as safe_load_file - - -WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json" - -logger = logging.getLogger(__name__) - - -def is_peft_model(model): - from .other import extract_model_from_parallel - - if is_peft_available(): - from peft import PeftModel - - return is_peft_available() and isinstance(extract_model_from_parallel(model), PeftModel) - - -def check_device_same(first_device, second_device): - """ - Utility method to check if two `torch` devices are similar. When dealing with CUDA devices, torch throws `False` - for `torch.device("cuda") == torch.device("cuda:0")` whereas they should be the same - - Args: - first_device (`torch.device`): - First device to check - second_device (`torch.device`): - Second device to check - """ - if first_device.type != second_device.type: - return False - - if first_device.type != "cpu" and first_device.index is None: - # In case the first_device is a cuda device and have - # the index attribute set to `None`, default it to `0` - first_device = torch.device(first_device.type, index=0) - - if second_device.type != "cpu" and second_device.index is None: - # In case the second_device is a cuda device and have - # the index attribute set to `None`, default it to `0` - second_device = torch.device(second_device.type, index=0) - - return first_device == second_device - - -def convert_file_size_to_int(size: Union[int, str]): - """ - Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes). - - Args: - size (`int` or `str`): The size to convert. Will be directly returned if an `int`. - - Example: - - ```py - >>> convert_file_size_to_int("1MiB") - 1048576 - ``` - """ - mem_size = -1 - err_msg = ( - f"`size` {size} is not in a valid format. Use an integer for bytes, or a string with an unit (like '5.0GB')." - ) - try: - if isinstance(size, int): - mem_size = size - elif size.upper().endswith("GIB"): - mem_size = int(float(size[:-3]) * (2**30)) - elif size.upper().endswith("MIB"): - mem_size = int(float(size[:-3]) * (2**20)) - elif size.upper().endswith("KIB"): - mem_size = int(float(size[:-3]) * (2**10)) - elif size.upper().endswith("GB"): - int_size = int(float(size[:-2]) * (10**9)) - mem_size = int_size // 8 if size.endswith("b") else int_size - elif size.upper().endswith("MB"): - int_size = int(float(size[:-2]) * (10**6)) - mem_size = int_size // 8 if size.endswith("b") else int_size - elif size.upper().endswith("KB"): - int_size = int(float(size[:-2]) * (10**3)) - mem_size = int_size // 8 if size.endswith("b") else int_size - except ValueError: - raise ValueError(err_msg) - - if mem_size < 0: - raise ValueError(err_msg) - return mem_size - - -def dtype_byte_size(dtype: torch.dtype): - """ - Returns the size (in bytes) occupied by one parameter of type `dtype`. - - Example: - - ```py - >>> dtype_byte_size(torch.float32) - 4 - ``` - """ - if dtype == torch.bool: - return 1 / 8 - elif dtype == CustomDtype.INT2: - return 1 / 4 - elif dtype == CustomDtype.INT4: - return 1 / 2 - elif dtype == CustomDtype.FP8: - return 1 - elif is_torch_version(">=", "2.1.0") and dtype in [torch.float8_e4m3fn, torch.float8_e5m2]: - return 1 - bit_search = re.search(r"[^\d](\d+)$", str(dtype)) - if bit_search is None: - raise ValueError(f"`dtype` is not a valid dtype: {dtype}.") - bit_size = int(bit_search.groups()[0]) - return bit_size // 8 - - -def id_tensor_storage(tensor: torch.Tensor) -> tuple[torch.device, int, int]: - """ - Unique identifier to a tensor storage. Multiple different tensors can share the same underlying storage. For - example, "meta" tensors all share the same storage, and thus their identifier will all be equal. This identifier is - guaranteed to be unique and constant for this tensor's storage during its lifetime. Two tensor storages with - non-overlapping lifetimes may have the same id. - """ - _SIZE = { - torch.int64: 8, - torch.float32: 4, - torch.int32: 4, - torch.bfloat16: 2, - torch.float16: 2, - torch.int16: 2, - torch.uint8: 1, - torch.int8: 1, - torch.bool: 1, - torch.float64: 8, - } - try: - storage_ptr = tensor.untyped_storage().data_ptr() - storage_size = tensor.untyped_storage().nbytes() - except Exception: - try: - # Fallback for torch==1.10 - storage_ptr = tensor.storage().data_ptr() - storage_size = tensor.storage().size() * _SIZE[tensor.dtype] - except NotImplementedError: - # Fallback for meta storage - storage_ptr = 0 - # On torch >=2.0 this is the tensor size - storage_size = tensor.nelement() * _SIZE[tensor.dtype] - - return tensor.device, storage_ptr, storage_size - - -def set_module_tensor_to_device( - module: nn.Module, - tensor_name: str, - device: Union[int, str, torch.device], - value: Optional[torch.Tensor] = None, - dtype: Optional[Union[str, torch.dtype]] = None, - fp16_statistics: Optional[torch.HalfTensor] = None, - tied_params_map: Optional[dict[int, dict[torch.device, torch.Tensor]]] = None, - non_blocking: bool = False, - clear_cache: bool = True, -): - """ - A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing - `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function). - - Args: - module (`torch.nn.Module`): - The module in which the tensor we want to move lives. - tensor_name (`str`): - The full name of the parameter/buffer. - device (`int`, `str` or `torch.device`): - The device on which to set the tensor. - value (`torch.Tensor`, *optional*): - The value of the tensor (useful when going from the meta device to any other device). - dtype (`torch.dtype`, *optional*): - If passed along the value of the parameter will be cast to this `dtype`. Otherwise, `value` will be cast to - the dtype of the existing parameter in the model. - fp16_statistics (`torch.HalfTensor`, *optional*): - The list of fp16 statistics to set on the module, used for 8 bit model serialization. - tied_params_map (Dict[int, Dict[torch.device, torch.Tensor]], *optional*, defaults to `None`): - A map of current data pointers to dictionaries of devices to already dispatched tied weights. For a given - execution device, this parameter is useful to reuse the first available pointer of a shared weight on the - device for all others, instead of duplicating memory. - non_blocking (`bool`, *optional*, defaults to `False`): - If `True`, the device transfer will be asynchronous with respect to the host, if possible. - clear_cache (`bool`, *optional*, defaults to `True`): - Whether or not to clear the device cache after setting the tensor on the device. - """ - # Recurse if needed - if "." in tensor_name: - splits = tensor_name.split(".") - for split in splits[:-1]: - new_module = getattr(module, split) - if new_module is None: - raise ValueError(f"{module} has no attribute {split}.") - module = new_module - tensor_name = splits[-1] - - if tensor_name not in module._parameters and tensor_name not in module._buffers: - raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.") - is_buffer = tensor_name in module._buffers - old_value = getattr(module, tensor_name) - - # Treat the case where old_value (or a custom `value`, typically offloaded to RAM/disk) belongs to a tied group, and one of the weight - # in the tied group has already been dispatched to the device, by avoiding reallocating memory on the device and just copying the pointer. - if ( - value is not None - and tied_params_map is not None - and value.data_ptr() in tied_params_map - and device in tied_params_map[value.data_ptr()] - ): - module._parameters[tensor_name] = tied_params_map[value.data_ptr()][device] - return - elif ( - tied_params_map is not None - and old_value.data_ptr() in tied_params_map - and device in tied_params_map[old_value.data_ptr()] - ): - module._parameters[tensor_name] = tied_params_map[old_value.data_ptr()][device] - return - - if old_value.device == torch.device("meta") and device not in ["meta", torch.device("meta")] and value is None: - raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.") - - param = module._parameters[tensor_name] if tensor_name in module._parameters else None - param_cls = type(param) - - if value is not None: - # We can expect mismatches when using bnb 4bit since Params4bit will reshape and pack the weights. - # In other cases, we want to make sure we're not loading checkpoints that do not match the config. - if old_value.shape != value.shape and param_cls.__name__ != "Params4bit": - raise ValueError( - f'Trying to set a tensor of shape {value.shape} in "{tensor_name}" (which has shape {old_value.shape}), this looks incorrect.' - ) - - if dtype is None: - # For compatibility with PyTorch load_state_dict which converts state dict dtype to existing dtype in model - value = value.to(old_value.dtype, non_blocking=non_blocking) - elif not str(value.dtype).startswith(("torch.uint", "torch.int", "torch.bool")): - value = value.to(dtype, non_blocking=non_blocking) - - device_quantization = None - with torch.no_grad(): - # leave it on cpu first before moving them to cuda - # # fix the case where the device is meta, we don't want to put it on cpu because there is no data =0 - if ( - param is not None - and param.device.type not in ("cuda", "xpu") - and torch.device(device).type in ("cuda", "xpu") - and param_cls.__name__ in ["Int8Params", "FP4Params", "Params4bit"] - ): - device_quantization = device - device = "cpu" - # `torch.Tensor.to()` is not supported by `torch_npu` (see this [issue](https://github.com/Ascend/pytorch/issues/16)). - if isinstance(device, int): - if is_npu_available(): - device = f"npu:{device}" - elif is_mlu_available(): - device = f"mlu:{device}" - elif is_sdaa_available(): - device = f"sdaa:{device}" - elif is_musa_available(): - device = f"musa:{device}" - elif is_hpu_available(): - device = "hpu" - if "xpu" in str(device) and not is_xpu_available(): - raise ValueError(f'{device} is not available, you should use device="cpu" instead') - if value is None: - new_value = old_value.to(device, non_blocking=non_blocking) - if dtype is not None and device in ["meta", torch.device("meta")]: - if not str(old_value.dtype).startswith(("torch.uint", "torch.int", "torch.bool")): - new_value = new_value.to(dtype, non_blocking=non_blocking) - - if not is_buffer: - module._parameters[tensor_name] = param_cls(new_value, requires_grad=old_value.requires_grad) - elif isinstance(value, torch.Tensor): - new_value = value.to(device, non_blocking=non_blocking) - else: - new_value = torch.tensor(value, device=device) - if device_quantization is not None: - device = device_quantization - if is_buffer: - module._buffers[tensor_name] = new_value - elif value is not None or not check_device_same(torch.device(device), module._parameters[tensor_name].device): - param_cls = type(module._parameters[tensor_name]) - kwargs = module._parameters[tensor_name].__dict__ - if param_cls.__name__ in ["Int8Params", "FP4Params", "Params4bit"]: - if param_cls.__name__ == "Int8Params" and new_value.dtype == torch.float32: - # downcast to fp16 if any - needed for 8bit serialization - new_value = new_value.to(torch.float16, non_blocking=non_blocking) - # quantize module that are going to stay on the cpu so that we offload quantized weights - if device == "cpu" and param_cls.__name__ == "Int8Params": - new_value = param_cls(new_value, requires_grad=old_value.requires_grad, **kwargs).to(0).to("cpu") - new_value.CB = new_value.CB.to("cpu") - new_value.SCB = new_value.SCB.to("cpu") - else: - new_value = param_cls(new_value, requires_grad=old_value.requires_grad, **kwargs).to( - device, non_blocking=non_blocking - ) - elif param_cls.__name__ in ["QTensor", "QBitsTensor"]: - new_value = torch.nn.Parameter(new_value, requires_grad=old_value.requires_grad).to( - device, non_blocking=non_blocking - ) - elif param_cls.__name__ in ["AffineQuantizedTensor"]: - new_value = new_value.to(device, non_blocking=non_blocking) - else: - new_value = param_cls(new_value, requires_grad=old_value.requires_grad).to( - device, non_blocking=non_blocking - ) - - module._parameters[tensor_name] = new_value - if fp16_statistics is not None: - module._parameters[tensor_name].SCB = fp16_statistics.to(device, non_blocking=non_blocking) - del fp16_statistics - # as we put the weight to meta, it doesn't have SCB attr anymore. make sure that it is not a meta weight - if ( - module.__class__.__name__ == "Linear8bitLt" - and getattr(module.weight, "SCB", None) is None - and str(module.weight.device) != "meta" - ): - # quantize only if necessary - device_index = torch.device(device).index if torch.device(device).type == "cuda" else None - if not getattr(module.weight, "SCB", None) and device_index is not None: - if module.bias is not None and module.bias.device.type != "meta": - # if a bias exists, we need to wait until the bias is set on the correct device - module = module.cuda(device_index) - elif module.bias is None: - # if no bias exists, we can quantize right away - module = module.cuda(device_index) - elif ( - module.__class__.__name__ == "Linear4bit" - and getattr(module.weight, "quant_state", None) is None - and str(module.weight.device) != "meta" - ): - # quantize only if necessary - device_index = torch.device(device).index if torch.device(device).type == "cuda" else None - if not getattr(module.weight, "quant_state", None) and device_index is not None: - module.weight = module.weight.cuda(device_index) - - # clean pre and post forward hook - if clear_cache and device != "cpu": - clear_device_cache() - - # When handling tied weights, we update tied_params_map to keep track of the tied weights that have already been allocated on the device in - # order to avoid duplicating memory, see above. - if ( - tied_params_map is not None - and old_value.data_ptr() in tied_params_map - and device not in tied_params_map[old_value.data_ptr()] - ): - tied_params_map[old_value.data_ptr()][device] = new_value - elif ( - value is not None - and tied_params_map is not None - and value.data_ptr() in tied_params_map - and device not in tied_params_map[value.data_ptr()] - ): - tied_params_map[value.data_ptr()][device] = new_value - - -def named_module_tensors( - module: nn.Module, include_buffers: bool = True, recurse: bool = False, remove_non_persistent: bool = False -): - """ - A helper function that gathers all the tensors (parameters + buffers) of a given module. If `include_buffers=True` - it's the same as doing `module.named_parameters(recurse=recurse) + module.named_buffers(recurse=recurse)`. - - Args: - module (`torch.nn.Module`): - The module we want the tensors on. - include_buffer (`bool`, *optional*, defaults to `True`): - Whether or not to include the buffers in the result. - recurse (`bool`, *optional`, defaults to `False`): - Whether or not to go look in every submodule or just return the direct parameters and buffers. - remove_non_persistent (`bool`, *optional*, defaults to `False`): - Whether or not to remove the non persistent buffer from the buffers. Useful only when include_buffers = - True - """ - yield from module.named_parameters(recurse=recurse) - - if include_buffers: - non_persistent_buffers = set() - if remove_non_persistent: - non_persistent_buffers = get_non_persistent_buffers(module, recurse=recurse) - for named_buffer in module.named_buffers(recurse=recurse): - name, _ = named_buffer - if name not in non_persistent_buffers: - yield named_buffer - - -def get_non_persistent_buffers(module: nn.Module, recurse: bool = False, fqns: bool = False): - """ - Gather all non persistent buffers of a given modules into a set - - Args: - module (`nn.Module`): - The module we want the non persistent buffers on. - recurse (`bool`, *optional*, defaults to `False`): - Whether or not to go look in every submodule or just return the direct non persistent buffers. - fqns (`bool`, *optional*, defaults to `False`): - Whether or not to return the fully-qualified names of the non persistent buffers. - """ - - non_persistent_buffers_set = module._non_persistent_buffers_set - if recurse: - for n, m in module.named_modules(): - if fqns: - non_persistent_buffers_set |= {n + "." + b for b in m._non_persistent_buffers_set} - else: - non_persistent_buffers_set |= m._non_persistent_buffers_set - - return non_persistent_buffers_set - - -class FindTiedParametersResult(list): - """ - This is a subclass of a list to handle backward compatibility for Transformers. Do not rely on the fact this is not - a list or on the `values` method as in the future this will be removed. - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def values(self): - warnings.warn( - "The 'values' method of FindTiedParametersResult is deprecated and will be removed in Accelerate v1.3.0. ", - FutureWarning, - ) - return sum([x[1:] for x in self], []) - - -def check_tied_parameters_in_config(model: nn.Module): - """ - Check if there is any indication in the given model that some weights should be tied. - - Args: - model (`torch.nn.Module`): The model to inspect - - Returns: - bool: True if the model needs to have tied weights - """ - - # based on model.tie_weights() method - has_tied_word_embedding = False - has_tied_encoder_decoder = False - has_tied_module = False - - if "PreTrainedModel" in [c.__name__ for c in inspect.getmro(model.__class__)]: - has_tied_word_embedding = False - model_decoder_config = None - if hasattr(model, "config"): - model_decoder_config = ( - model.config.get_text_config(decoder=True) - if hasattr(model.config, "get_text_config") - else model.config - ) - has_tied_word_embedding = ( - model_decoder_config is not None - and getattr(model_decoder_config, "tie_word_embeddings", False) - and model.get_output_embeddings() - ) - - has_tied_encoder_decoder = ( - hasattr(model, "config") - and getattr(model.config, "is_encoder_decoder", False) - and getattr(model.config, "tie_encoder_decoder", False) - ) - has_tied_module = any(hasattr(module, "_tie_weights") for module in model.modules()) - return any([has_tied_word_embedding, has_tied_encoder_decoder, has_tied_module]) - - -def _get_param_device(param, device_map): - if param in device_map: - return device_map[param] - parent_param = ".".join(param.split(".")[:-1]) - if parent_param == param: - raise ValueError(f"The `device_map` does not contain the module {param}.") - else: - return _get_param_device(parent_param, device_map) - - -def check_tied_parameters_on_same_device(tied_params, device_map): - """ - Check if tied parameters are on the same device - - Args: - tied_params (`List[List[str]]`): - A list of lists of parameter names being all tied together. - - device_map (`Dict[str, Union[int, str, torch.device]]`): - A map that specifies where each submodule should go. - - """ - for tie_param in tied_params: - tie_param_devices = {} - for param in tie_param: - tie_param_devices[param] = _get_param_device(param, device_map) - if len(set(tie_param_devices.values())) > 1: - logger.warning( - f"Tied parameters are on different devices: {tie_param_devices}. " - "Please modify your custom device map or set `device_map='auto'`. " - ) - - -def find_tied_parameters(model: torch.nn.Module, **kwargs): - """ - Find the tied parameters in a given model. - - - - The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore - them. - - - - Args: - model (`torch.nn.Module`): The model to inspect. - - Returns: - List[List[str]]: A list of lists of parameter names being all tied together. - - Example: - - ```py - >>> from collections import OrderedDict - >>> import torch.nn as nn - - >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))])) - >>> model.linear2.weight = model.linear1.weight - >>> find_tied_parameters(model) - [['linear1.weight', 'linear2.weight']] - ``` - """ - - # get ALL model parameters and their names - all_named_parameters = {name: param for name, param in model.named_parameters(remove_duplicate=False)} - - # get ONLY unique named parameters, - # if parameter is tied and have multiple names, it will be included only once - no_duplicate_named_parameters = {name: param for name, param in model.named_parameters(remove_duplicate=True)} - - # the difference of the two sets will give us the tied parameters - tied_param_names = set(all_named_parameters.keys()) - set(no_duplicate_named_parameters.keys()) - - # 'tied_param_names' contains the names of parameters that are tied in the model, but we do not know - # which names refer to the same parameter. To identify this, we need to group them together. - tied_param_groups = {} - for tied_param_name in tied_param_names: - tied_param = all_named_parameters[tied_param_name] - for param_name, param in no_duplicate_named_parameters.items(): - # compare if parameters are the same, if so, group their names together - if param is tied_param: - if param_name not in tied_param_groups: - tied_param_groups[param_name] = [] - tied_param_groups[param_name].append(tied_param_name) - - return FindTiedParametersResult([sorted([weight] + list(set(tied))) for weight, tied in tied_param_groups.items()]) - - -def retie_parameters(model, tied_params): - """ - Reties tied parameters in a given model if the link was broken (for instance when adding hooks). - - Args: - model (`torch.nn.Module`): - The model in which to retie parameters. - tied_params (`List[List[str]]`): - A mapping parameter name to tied parameter name as obtained by `find_tied_parameters`. - """ - for tied_group in tied_params: - param_to_tie = None - # two loops : the first one to set param_to_tie , the second one to change the values of tied_group - for param_name in tied_group: - module = model - splits = param_name.split(".") - for split in splits[:-1]: - module = getattr(module, split) - param = getattr(module, splits[-1]) - if param_to_tie is None and param.device != torch.device("meta"): - param_to_tie = param - break - if param_to_tie is not None: - for param_name in tied_group: - module = model - splits = param_name.split(".") - for split in splits[:-1]: - module = getattr(module, split) - setattr(module, splits[-1], param_to_tie) - - -def _get_proper_dtype(dtype: Union[str, torch.device]) -> torch.dtype: - """ - Just does torch.dtype(dtype) if necessary. - """ - if isinstance(dtype, str): - # We accept "torch.float16" or just "float16" - dtype = dtype.replace("torch.", "") - dtype = getattr(torch, dtype) - return dtype - - -def compute_module_sizes( - model: nn.Module, - dtype: Optional[Union[str, torch.device]] = None, - special_dtypes: Optional[dict[str, Union[str, torch.device]]] = None, - buffers_only: bool = False, -): - """ - Compute the size of each submodule of a given model. - """ - 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 = [] - - if not buffers_only: - module_list = named_module_tensors(model, recurse=True) - else: - module_list = model.named_buffers(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 compute_module_total_buffer_size( - model: nn.Module, - dtype: Optional[Union[str, torch.device]] = None, - special_dtypes: Optional[dict[str, Union[str, torch.device]]] = None, -): - """ - Compute the total size of buffers in each submodule of a given model. - """ - module_sizes = compute_module_sizes(model, dtype=dtype, special_dtypes=special_dtypes, buffers_only=True) - return module_sizes.get("", 0) - - -def get_max_layer_size( - modules: list[tuple[str, torch.nn.Module]], module_sizes: dict[str, int], no_split_module_classes: list[str] -): - """ - Utility function that will scan a list of named modules and return the maximum size used by one full layer. The - definition of a layer being: - - a module with no direct children (just parameters and buffers) - - a module whose class name is in the list `no_split_module_classes` - - Args: - modules (`List[Tuple[str, torch.nn.Module]]`): - The list of named modules where we want to determine the maximum layer size. - module_sizes (`Dict[str, int]`): - A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`). - no_split_module_classes (`List[str]`): - A list of class names for layers we don't want to be split. - - Returns: - `Tuple[int, List[str]]`: The maximum size of a layer with the list of layer names realizing that maximum size. - """ - max_size = 0 - layer_names = [] - modules_to_treat = modules.copy() - while len(modules_to_treat) > 0: - module_name, module = modules_to_treat.pop(0) - modules_children = list(module.named_children()) if isinstance(module, torch.nn.Module) else [] - if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes: - # No splitting this one so we compare to the max_size - size = module_sizes[module_name] - if size > max_size: - max_size = size - layer_names = [module_name] - elif size == max_size: - layer_names.append(module_name) - else: - modules_to_treat = [(f"{module_name}.{n}", v) for n, v in modules_children] + modules_to_treat - return max_size, layer_names - - -def get_max_memory(max_memory: Optional[dict[Union[int, str], Union[int, str]]] = None): - """ - Get the maximum memory available if nothing is passed, converts string to int otherwise. - """ - import psutil - - if max_memory is None: - max_memory = {} - # Make sure CUDA is initialized on each GPU to have the right memory info. - if is_npu_available(): - for i in range(torch.npu.device_count()): - try: - _ = torch.tensor(0, device=torch.device("npu", i)) - max_memory[i] = torch.npu.mem_get_info(i)[0] - except Exception: - logger.info(f"Device {i} seems unavailable, Proceeding to check subsequent devices.") - continue - elif is_mlu_available(): - for i in range(torch.mlu.device_count()): - try: - _ = torch.tensor(0, device=torch.device("mlu", i)) - max_memory[i] = torch.mlu.mem_get_info(i)[0] - except Exception: - logger.info(f"Device {i} seems unavailable, Proceeding to check subsequent devices.") - continue - elif is_sdaa_available(): - for i in range(torch.sdaa.device_count()): - try: - _ = torch.tensor(0, device=torch.device("sdaa", i)) - max_memory[i] = torch.sdaa.mem_get_info(i)[0] - except Exception: - logger.info(f"Device {i} seems unavailable, Proceeding to check subsequent devices.") - continue - elif is_musa_available(): - for i in range(torch.musa.device_count()): - try: - _ = torch.tensor(0, device=torch.device("musa", i)) - max_memory[i] = torch.musa.mem_get_info(i)[0] - except Exception: - logger.info(f"Device {i} seems unavailable, Proceeding to check subsequent devices.") - continue - elif is_xpu_available(): - for i in range(torch.xpu.device_count()): - try: - _ = torch.tensor(0, device=torch.device("xpu", i)) - max_memory[i] = get_xpu_available_memory(i) - except Exception: - logger.info(f"Device {i} seems unavailable, Proceeding to check subsequent devices.") - continue - elif is_hpu_available(): - for i in range(torch.hpu.device_count()): - try: - _ = torch.tensor(0, device=torch.device("hpu", i)) - max_memory[i] = torch.hpu.mem_get_info(i)[0] - except Exception: - logger.info(f"Device {i} seems unavailable, Proceeding to check subsequent devices.") - continue - else: - for i in range(torch.cuda.device_count()): - try: - _ = torch.tensor([0], device=i) - max_memory[i] = torch.cuda.mem_get_info(i)[0] - except Exception: - logger.info(f"Device {i} seems unavailable, Proceeding to check subsequent devices.") - continue - # allocate everything in the mps device as the RAM is shared - if is_mps_available(): - max_memory["mps"] = psutil.virtual_memory().available - else: - max_memory["cpu"] = psutil.virtual_memory().available - return max_memory - - for key in max_memory: - if isinstance(max_memory[key], str): - max_memory[key] = convert_file_size_to_int(max_memory[key]) - - # Need to sort the device by type to make sure that we allocate the gpu first. - # As gpu/npu/xpu are represented by int, we need to sort them first. - gpu_devices = [k for k in max_memory.keys() if isinstance(k, int)] - gpu_devices.sort() - # check if gpu/npu/xpu devices are available and if not, throw a warning - if is_npu_available(): - num_devices = torch.npu.device_count() - elif is_mlu_available(): - num_devices = torch.mlu.device_count() - elif is_sdaa_available(): - num_devices = torch.sdaa.device_count() - elif is_musa_available(): - num_devices = torch.musa.device_count() - elif is_xpu_available(): - num_devices = torch.xpu.device_count() - elif is_hpu_available(): - num_devices = torch.hpu.device_count() - else: - num_devices = torch.cuda.device_count() - for device in gpu_devices: - if device >= num_devices or device < 0: - logger.warning(f"Device {device} is not available, available devices are {list(range(num_devices))}") - # Add the other devices in the preset order if they are available - all_devices = gpu_devices + [k for k in ["mps", "cpu", "disk"] if k in max_memory.keys()] - # Raise an error if a device is not recognized - for k in max_memory.keys(): - if k not in all_devices: - raise ValueError( - f"Device {k} is not recognized, available devices are integers(for GPU/XPU), 'mps', 'cpu' and 'disk'" - ) - max_memory = {k: max_memory[k] for k in all_devices} - - return max_memory - - -def clean_device_map(device_map: dict[str, Union[int, str, torch.device]], module_name: str = ""): - """ - Cleans a device_map by grouping all submodules that go on the same device together. - """ - # Get the value of the current module and if there is only one split across several keys, regroup it. - prefix = "" if module_name == "" else f"{module_name}." - values = [v for k, v in device_map.items() if k.startswith(prefix)] - if len(set(values)) == 1 and len(values) > 1: - for k in [k for k in device_map if k.startswith(prefix)]: - del device_map[k] - device_map[module_name] = values[0] - - # Recurse over the children - children_modules = [k for k in device_map.keys() if k.startswith(prefix) and len(k) > len(module_name)] - idx = len(module_name.split(".")) + 1 if len(module_name) > 0 else 1 - children_modules = set(".".join(k.split(".")[:idx]) for k in children_modules) - for child in children_modules: - clean_device_map(device_map, module_name=child) - - return device_map - - -def load_offloaded_weights(model, index, offload_folder): - """ - Loads the weights from the offload folder into the model. - - Args: - model (`torch.nn.Module`): - The model to load the weights into. - index (`dict`): - A dictionary containing the parameter name and its metadata for each parameter that was offloaded from the - model. - offload_folder (`str`): - The folder where the offloaded weights are stored. - """ - if index is None or len(index) == 0: - # Nothing to do - return - for param_name, metadata in index.items(): - if "SCB" in param_name: - continue - fp16_statistics = None - if "weight" in param_name and param_name.replace("weight", "SCB") in index.keys(): - weight_name = param_name.replace("weight", "SCB") - fp16_statistics = load_offloaded_weight( - os.path.join(offload_folder, f"{weight_name}.dat"), index[weight_name] - ) - tensor_file = os.path.join(offload_folder, f"{param_name}.dat") - weight = load_offloaded_weight(tensor_file, metadata) - set_module_tensor_to_device(model, param_name, "cpu", value=weight, fp16_statistics=fp16_statistics) - - -def get_module_leaves(module_sizes): - module_children = {} - for module in module_sizes: - if module == "" or "." not in module: - continue - parent = module.rsplit(".", 1)[0] - module_children[parent] = module_children.get(parent, 0) + 1 - leaves = [module for module in module_sizes if module_children.get(module, 0) == 0 and module != ""] - return leaves - - -def get_balanced_memory( - model: nn.Module, - max_memory: Optional[dict[Union[int, str], Union[int, str]]] = None, - no_split_module_classes: Optional[list[str]] = None, - dtype: Optional[Union[str, torch.dtype]] = None, - special_dtypes: Optional[dict[str, Union[str, torch.device]]] = None, - low_zero: bool = False, -): - """ - Compute a `max_memory` dictionary for [`infer_auto_device_map`] that will balance the use of each available GPU. - - - - All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the - meta device (as it would if initialized within the `init_empty_weights` context manager). - - - - Args: - model (`torch.nn.Module`): - The model to analyze. - max_memory (`Dict`, *optional*): - A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset. - Example: `max_memory={0: "1GB"}`. - no_split_module_classes (`List[str]`, *optional*): - A list of layer class names that should never be split across device (for instance any layer that has a - residual connection). - dtype (`str` or `torch.dtype`, *optional*): - If provided, the weights will be converted to that type when loaded. - special_dtypes (`Dict[str, Union[str, torch.device]]`, *optional*): - If provided, special dtypes to consider for some specific weights (will override dtype used as default for - all weights). - low_zero (`bool`, *optional*): - Minimizes the number of weights on GPU 0, which is convenient when it's used for other operations (like the - Transformers generate function). - """ - # Get default / clean up max_memory - user_not_set_max_memory = max_memory is None - max_memory = get_max_memory(max_memory) - - if is_npu_available(): - expected_device_type = "npu" - elif is_mlu_available(): - expected_device_type = "mlu" - elif is_sdaa_available(): - expected_device_type = "sdaa" - elif is_musa_available(): - expected_device_type = "musa" - elif is_xpu_available(): - expected_device_type = "xpu" - elif is_hpu_available(): - expected_device_type = "hpu" - elif is_mps_available(): - expected_device_type = "mps" - else: - expected_device_type = "cuda" - num_devices = len([d for d in max_memory if torch.device(d).type == expected_device_type and max_memory[d] > 0]) - - if num_devices == 0: - return max_memory - - if num_devices == 1: - # We cannot do low_zero on just one GPU, but we will still reserve some memory for the buffer - low_zero = False - # If user just asked us to handle memory usage, we should avoid OOM - if user_not_set_max_memory: - for key in max_memory.keys(): - if isinstance(key, int): - max_memory[key] *= 0.9 # 90% is a good compromise - logger.info( - f"We will use 90% of the memory on device {key} for storing the model, and 10% for the buffer to avoid OOM. " - "You can set `max_memory` in to a higher value to use more memory (at your own risk)." - ) - break # only one device - - module_sizes = compute_module_sizes(model, dtype=dtype, special_dtypes=special_dtypes) - per_gpu = module_sizes[""] // (num_devices - 1 if low_zero else num_devices) - - # We can't just set the memory to model_size // num_devices as it will end being too small: each GPU will get - # slightly less layers and some layers will end up offload at the end. So this function computes a buffer size to - # add which is the biggest of: - # - the size of no split block (if applicable) - # - the mean of the layer sizes - if no_split_module_classes is None: - no_split_module_classes = [] - elif not isinstance(no_split_module_classes, (list, tuple)): - no_split_module_classes = [no_split_module_classes] - - # Identify the size of the no_split_block modules - if len(no_split_module_classes) > 0: - no_split_children = {} - for name, size in module_sizes.items(): - if name == "": - continue - submodule = model - for submodule_name in name.split("."): - submodule = getattr(submodule, submodule_name) - class_name = submodule.__class__.__name__ - if class_name in no_split_module_classes and class_name not in no_split_children: - no_split_children[class_name] = size - - if set(no_split_children.keys()) == set(no_split_module_classes): - break - buffer = max(no_split_children.values()) if len(no_split_children) > 0 else 0 - else: - buffer = 0 - - # Compute mean of final modules. In the first dict of module sizes, leaves are the parameters - leaves = get_module_leaves(module_sizes) - module_sizes = {n: v for n, v in module_sizes.items() if n not in leaves} - # Once removed, leaves are the final modules. - leaves = get_module_leaves(module_sizes) - mean_leaves = int(sum([module_sizes[n] for n in leaves]) / max(len(leaves), 1)) - buffer = int(1.25 * max(buffer, mean_leaves)) - per_gpu += buffer - - # Sorted list of GPUs id (we may have some gpu ids not included in the our max_memory list - let's ignore them) - gpus_idx_list = list( - sorted( - device_id for device_id, device_mem in max_memory.items() if isinstance(device_id, int) and device_mem > 0 - ) - ) - # The last device is left with max_memory just in case the buffer is not enough. - for idx in gpus_idx_list[:-1]: - max_memory[idx] = min(max_memory[0] if low_zero and idx == 0 else per_gpu, max_memory[idx]) - - if low_zero: - min_zero = max(0, module_sizes[""] - sum([max_memory[i] for i in range(1, num_devices)])) - max_memory[0] = min(min_zero, max_memory[0]) - - return max_memory - - -def calculate_maximum_sizes(model: torch.nn.Module): - "Computes the total size of the model and its largest layer" - sizes = compute_module_sizes(model) - # `transformers` models store this information for us - no_split_modules = getattr(model, "_no_split_modules", None) - if no_split_modules is None: - no_split_modules = [] - - modules_to_treat = ( - list(model.named_parameters(recurse=False)) - + list(model.named_children()) - + list(model.named_buffers(recurse=False)) - ) - largest_layer = get_max_layer_size(modules_to_treat, sizes, no_split_modules) - total_size = sizes[""] - return total_size, largest_layer - - -def _init_infer_auto_device_map( - model: nn.Module, - max_memory: Optional[dict[Union[int, str], Union[int, str]]] = None, - no_split_module_classes: Optional[list[str]] = None, - dtype: Optional[Union[str, torch.dtype]] = None, - special_dtypes: Optional[dict[str, Union[str, torch.device]]] = None, -) -> tuple[ - list[Union[int, str]], - dict[Union[int, str], Union[int, str]], - list[Union[int, str]], - list[int], - dict[str, int], - list[list[str]], - list[str], - list[tuple[str, nn.Module]], -]: - """ - Initialize variables required for computing the device map for model allocation. - """ - max_memory = get_max_memory(max_memory) - if no_split_module_classes is None: - no_split_module_classes = [] - elif not isinstance(no_split_module_classes, (list, tuple)): - no_split_module_classes = [no_split_module_classes] - - devices = list(max_memory.keys()) - if "disk" not in devices: - devices.append("disk") - gpus = [device for device in devices if device not in ["cpu", "disk"]] - - # Devices that need to keep space for a potential offloaded layer. - if "mps" in gpus: - main_devices = ["mps"] - elif len(gpus) > 0: - main_devices = [gpus[0], "cpu"] - else: - main_devices = ["cpu"] - - module_sizes = compute_module_sizes(model, dtype=dtype, special_dtypes=special_dtypes) - tied_parameters = find_tied_parameters(model) - if check_tied_parameters_in_config(model) and len(tied_parameters) == 0: - logger.warning( - "The model weights are not tied. Please use the `tie_weights` method before using the `infer_auto_device` function." - ) - - # Direct submodules and parameters - modules_to_treat = ( - list(model.named_parameters(recurse=False)) - + list(model.named_children()) - + list(model.named_buffers(recurse=False)) - ) - - return ( - devices, - max_memory, - main_devices, - gpus, - module_sizes, - tied_parameters, - no_split_module_classes, - modules_to_treat, - ) - - -def get_module_size_with_ties( - tied_params, - module_size, - module_sizes, - modules_to_treat, -) -> tuple[int, list[str], list[nn.Module]]: - """ - Calculate the total size of a module, including its tied parameters. - - Args: - tied_params (`List[str]`): The list of tied parameters. - module_size (`int`): The size of the module without tied parameters. - module_sizes (`Dict[str, int]`): A dictionary mapping each layer name to its size. - modules_to_treat (`List[Tuple[str, nn.Module]]`): The list of named modules to treat. - - Returns: - `Tuple[int, List[str], List[nn.Module]]`: The total size of the module, the names of the tied modules, and the - tied modules. - """ - if len(tied_params) < 1: - return module_size, [], [] - tied_module_names = [] - tied_modules = [] - - for tied_param in tied_params: - tied_module_index = [i for i, (n, _) in enumerate(modules_to_treat) if tied_param.startswith(n + ".")][0] - tied_module_names.append(modules_to_treat[tied_module_index][0]) - tied_modules.append(modules_to_treat[tied_module_index][1]) - - module_size_with_ties = module_size - for tied_param, tied_module_name in zip(tied_params, tied_module_names): - module_size_with_ties += module_sizes[tied_module_name] - module_sizes[tied_param] - - return module_size_with_ties, tied_module_names, tied_modules - - -def fallback_allocate( - modules: list[tuple[str, nn.Module]], - module_sizes: dict[str, int], - size_limit: Union[int, str], - no_split_module_classes: Optional[list[str]] = None, - tied_parameters: Optional[list[list[str]]] = None, -) -> tuple[Optional[str], Optional[nn.Module], list[tuple[str, nn.Module]]]: - """ - Find a module that fits in the size limit using BFS and return it with its name and the remaining modules. - - Args: - modules (`List[Tuple[str, nn.Module]]`): - The list of named modules to search in. - module_sizes (`Dict[str, int]`): - A dictionary mapping each layer name to its size (as generated by `compute_module_sizes`). - size_limit (`Union[int, str]`): - The maximum size a module can have. - no_split_module_classes (`Optional[List[str]]`, *optional*): - A list of class names for layers we don't want to be split. - tied_parameters (`Optional[List[List[str]]`, *optional*): - A list of lists of parameter names being all tied together. - - Returns: - `Tuple[Optional[str], Optional[nn.Module], List[Tuple[str, nn.Module]]]`: A tuple containing: - - The name of the module that fits within the size limit. - - The module itself. - - The list of remaining modules after the found module is removed. - """ - try: - size_limit = convert_file_size_to_int(size_limit) - except ValueError: - return None, None, modules - - if no_split_module_classes is None: - no_split_module_classes = [] - - if tied_parameters is None: - tied_parameters = [] - - modules_to_search = modules.copy() - module_found = False - - while modules_to_search: - name, module = modules_to_search.pop(0) - - tied_param_groups = [ - tied_group - for tied_group in tied_parameters - if any(name + "." in k + "." for k in tied_group) and not all(name + "." in k + "." for k in tied_group) - ] - - tied_params = sum( - [[p for p in tied_group if name + "." not in p + "."] for tied_group in tied_param_groups], [] - ) - - module_size_with_ties, _, _ = get_module_size_with_ties( - tied_params, module_sizes[name], module_sizes, modules_to_search - ) - - # If the module fits in the size limit, we found it. - if module_size_with_ties <= size_limit: - module_found = True - break - - # The module is too big, we need to split it if possible. - modules_children = ( - [] - if isinstance(module, nn.Parameter) or isinstance(module, torch.Tensor) - else list(module.named_children()) - ) - - # Split fails, move to the next module - if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes: - continue - - # split is possible, add the children to the list of modules to search - modules_children = list(module.named_parameters(recurse=False)) + modules_children - modules_to_search = [(f"{name}.{n}", v) for n, v in modules_children] + modules_to_search - - if not module_found: - return None, None, modules - - # Prepare the module list for removal of the found module - current_names = [n for n, _ in modules] - dot_idx = [i for i, c in enumerate(name) if c == "."] - - for dot_index in dot_idx: - parent_name = name[:dot_index] - if parent_name in current_names: - parent_module_idx = current_names.index(parent_name) - _, parent_module = modules[parent_module_idx] - module_children = list(parent_module.named_parameters(recurse=False)) + list( - parent_module.named_children() - ) - modules = ( - modules[:parent_module_idx] - + [(f"{parent_name}.{n}", v) for n, v in module_children] - + modules[parent_module_idx + 1 :] - ) - current_names = [n for n, _ in modules] - - # Now the target module should be directly in the list - target_idx = current_names.index(name) - name, module = modules.pop(target_idx) - - return name, module, modules - - -def infer_auto_device_map( - model: nn.Module, - max_memory: Optional[dict[Union[int, str], Union[int, str]]] = None, - no_split_module_classes: Optional[list[str]] = None, - dtype: Optional[Union[str, torch.dtype]] = None, - special_dtypes: Optional[dict[str, Union[str, torch.dtype]]] = None, - verbose: bool = False, - clean_result: bool = True, - offload_buffers: bool = False, - fallback_allocation: bool = False, -): - """ - Compute a device map for a given model giving priority to GPUs, then offload on CPU and finally offload to disk, - such that: - - we don't exceed the memory available of any of the GPU. - - if offload to the CPU is needed, there is always room left on GPU 0 to put back the layer offloaded on CPU that - has the largest size. - - if offload to the CPU is needed,we don't exceed the RAM available on the CPU. - - if offload to the disk is needed, there is always room left on the CPU to put back the layer offloaded on disk - that has the largest size. - - - - All computation is done analyzing sizes and dtypes of the model parameters. As a result, the model can be on the - meta device (as it would if initialized within the `init_empty_weights` context manager). - - - - Args: - model (`torch.nn.Module`): - The model to analyze. - max_memory (`Dict`, *optional*): - A dictionary device identifier to maximum memory. Will default to the maximum memory available if unset. - Example: `max_memory={0: "1GB"}`. - no_split_module_classes (`List[str]`, *optional*): - A list of layer class names that should never be split across device (for instance any layer that has a - residual connection). - dtype (`str` or `torch.dtype`, *optional*): - If provided, the weights will be converted to that type when loaded. - special_dtypes (`Dict[str, Union[str, torch.device]]`, *optional*): - If provided, special dtypes to consider for some specific weights (will override dtype used as default for - all weights). - verbose (`bool`, *optional*, defaults to `False`): - Whether or not to provide debugging statements as the function builds the device_map. - clean_result (`bool`, *optional*, defaults to `True`): - Clean the resulting device_map by grouping all submodules that go on the same device together. - offload_buffers (`bool`, *optional*, defaults to `False`): - In the layers that are offloaded on the CPU or the hard drive, whether or not to offload the buffers as - well as the parameters. - fallback_allocation (`bool`, *optional*, defaults to `False`): - When regular allocation fails, try to allocate a module that fits in the size limit using BFS. - """ - - # Initialize the variables - ( - devices, - max_memory, - main_devices, - gpus, - module_sizes, - tied_parameters, - no_split_module_classes, - modules_to_treat, - ) = _init_infer_auto_device_map(model, max_memory, no_split_module_classes, dtype, special_dtypes) - - device_map = OrderedDict() - current_device = 0 - device_memory_used = {device: 0 for device in devices} - device_buffer_sizes = {} - device_minimum_assignment_memory = {} - - # Initialize maximum largest layer, to know which space to keep in memory - max_layer_size, max_layer_names = get_max_layer_size(modules_to_treat, module_sizes, no_split_module_classes) - - # Ready ? This is going to be a bit messy. - while len(modules_to_treat) > 0: - name, module = modules_to_treat.pop(0) - if verbose: - print(f"\nTreating module {name}.") - # Max size in the remaining layers may have changed since we took one, so we maybe update it. - max_layer_names = [n for n in max_layer_names if n != name and not n.startswith(name + ".")] - if len(max_layer_names) == 0: - max_layer_size, max_layer_names = get_max_layer_size( - [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)], - module_sizes, - no_split_module_classes, - ) - # Assess size needed - module_size = module_sizes[name] - - # We keep relevant tied parameters only: one of the tied parameters in the group is inside the current module - # and the other is not. - # Note: If we are currently processing the name `compute.weight`, an other parameter named - # e.g. `compute.weight_submodule.parameter` - # needs to be considered outside the current module, hence the check with additional dots. - tied_param_groups = [ - tied_group - for tied_group in tied_parameters - if any(name + "." in k + "." for k in tied_group) and not all(name + "." in k + "." for k in tied_group) - ] - - if verbose and len(tied_param_groups) > 0: - print(f" Found the relevant tied param groups {tied_param_groups}") - - # Then we keep track of all the parameters that are tied to the current module, but not in the current module - tied_params = sum( - [[p for p in tied_group if name + "." not in p + "."] for tied_group in tied_param_groups], [] - ) - - if verbose and len(tied_params) > 0: - print(f" So those parameters need to be taken into account {tied_params}") - - device = devices[current_device] - current_max_size = max_memory[device] if device != "disk" else None - current_memory_reserved = 0 - # Reduce max size available by the largest layer. - if devices[current_device] in main_devices: - current_max_size = current_max_size - max_layer_size - current_memory_reserved = max_layer_size - - module_size_with_ties, tied_module_names, tied_modules = get_module_size_with_ties( - tied_params, module_size, module_sizes, modules_to_treat - ) - - # The module and its tied modules fit on the current device. - if current_max_size is None or device_memory_used[device] + module_size_with_ties <= current_max_size: - if verbose: - output = f"Putting {name}" - - if tied_module_names: - output += f" and {tied_module_names}" - else: - output += f" (size={module_size})" - - if current_max_size is not None: - output += f" (available={current_max_size - device_memory_used[device]})" - - output += f" on {device}." - print(output) - - device_memory_used[device] += module_size_with_ties - - # Assign the primary module to the device. - device_map[name] = device - - # Assign tied modules if any. - for tied_module_name in tied_module_names: - if tied_module_name in [m[0] for m in modules_to_treat]: - # Find the index of the tied module in the list - tied_module_index = next(i for i, (n, _) in enumerate(modules_to_treat) if n == tied_module_name) - # Remove the tied module from the list to prevent reprocessing - modules_to_treat.pop(tied_module_index) - - # Assign the tied module to the device - device_map[tied_module_name] = device - - # Buffer Handling - if not offload_buffers and isinstance(module, nn.Module): - # Compute the total buffer size for the module - current_buffer_size = compute_module_total_buffer_size( - module, dtype=dtype, special_dtypes=special_dtypes - ) - # Update the buffer size on the device - device_buffer_sizes[device] = device_buffer_sizes.get(device, 0) + current_buffer_size - - continue - - # The current module itself fits, so we try to split the tied modules. - if len(tied_params) > 0 and device_memory_used[device] + module_size <= current_max_size: - # can we split one of the tied modules to make it smaller or do we need to go on the next device? - if verbose: - print( - f"Not enough space on {devices[current_device]} to put {name} and {tied_module_names} (space " - f"available {current_max_size - device_memory_used[device]}, needed size {module_size_with_ties})." - ) - split_happened = False - for tied_module_name, tied_module in zip(tied_module_names, tied_modules): - tied_module_children = list(tied_module.named_children()) - if len(tied_module_children) == 0 or tied_module.__class__.__name__ in no_split_module_classes: - # can't break this one. - continue - - if verbose: - print(f"Splitting {tied_module_name}.") - tied_module_children = list(tied_module.named_parameters(recurse=False)) + tied_module_children - tied_module_children = [(f"{tied_module_name}.{n}", v) for n, v in tied_module_children] - tied_module_index = [i for i, (n, _) in enumerate(modules_to_treat) if n == tied_module_name][0] - - modules_to_treat = ( - [(name, module)] - + modules_to_treat[:tied_module_index] - + tied_module_children - + modules_to_treat[tied_module_index + 1 :] - ) - # Update the max layer size. - max_layer_size, max_layer_names = get_max_layer_size( - [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)], - module_sizes, - no_split_module_classes, - ) - split_happened = True - break - - if split_happened: - continue - - # If the tied module is not split, we go to the next device - if verbose: - print("None of the tied module can be split, going to the next device.") - - # The current module itself doesn't fit, so we have to split it or go to the next device. - if device_memory_used[device] + module_size >= current_max_size: - # Split or not split? - modules_children = ( - [] - if isinstance(module, nn.Parameter) or isinstance(module, torch.Tensor) - else list(module.named_children()) - ) - if verbose: - print( - f"Not enough space on {devices[current_device]} to put {name} (space available " - f"{current_max_size - device_memory_used[device]}, module size {module_size})." - ) - if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes: - # -> no split, we go to the next device - if verbose: - print("This module cannot be split, going to the next device.") - - else: - # -> split, we replace the module studied by its children + parameters - if verbose: - print(f"Splitting {name}.") - modules_children = list(module.named_parameters(recurse=False)) + modules_children - modules_to_treat = [(f"{name}.{n}", v) for n, v in modules_children] + modules_to_treat - # Update the max layer size. - max_layer_size, max_layer_names = get_max_layer_size( - [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)], - module_sizes, - no_split_module_classes, - ) - continue - - # If no module is assigned to the current device, we attempt to allocate a fallback module - # if fallback_allocation is enabled. - if device_memory_used[device] == 0 and fallback_allocation and device != "disk": - # We try to allocate a module that fits in the size limit using BFS. - # Recompute the current max size as we need to consider the current module as well. - current_max_size = max_memory[device] - max(max_layer_size, module_size_with_ties) - - fallback_module_name, fallback_module, remaining_modules = fallback_allocate( - modules_to_treat, - module_sizes, - current_max_size - device_memory_used[device], - no_split_module_classes, - tied_parameters, - ) - # use the next iteration to put the fallback module on the next device to avoid code duplication - if fallback_module is not None: - modules_to_treat = [(fallback_module_name, fallback_module)] + [(name, module)] + remaining_modules - continue - - if device_memory_used[device] == 0: - device_minimum_assignment_memory[device] = module_size_with_ties + current_memory_reserved - - # Neither the current module nor any tied modules can be split, so we move to the next device. - device_memory_used[device] = device_memory_used[device] + current_memory_reserved - current_device += 1 - modules_to_treat = [(name, module)] + modules_to_treat - - device_memory_used = {device: mem for device, mem in device_memory_used.items() if mem > 0} - - if clean_result: - device_map = clean_device_map(device_map) - - non_gpu_buffer_size = device_buffer_sizes.get("cpu", 0) + device_buffer_sizes.get("disk", 0) - if non_gpu_buffer_size > 0 and not offload_buffers: - is_buffer_fit_any_gpu = False - for gpu_device, gpu_max_memory in max_memory.items(): - if gpu_device == "cpu" or gpu_device == "disk": - continue - - if not is_buffer_fit_any_gpu: - gpu_memory_used = device_memory_used.get(gpu_device, 0) - - if gpu_max_memory >= non_gpu_buffer_size + gpu_memory_used: - is_buffer_fit_any_gpu = True - - if len(gpus) > 0 and not is_buffer_fit_any_gpu: - warnings.warn( - f"Current model requires {non_gpu_buffer_size} bytes of buffer for offloaded layers, which seems does " - f"not fit any GPU's remaining memory. If you are experiencing a OOM later, please consider using " - f"offload_buffers=True." - ) - - if device_minimum_assignment_memory: - devices_info = "\n".join( - f" - {device}: {mem} bytes required" for device, mem in device_minimum_assignment_memory.items() - ) - logger.info( - f"Based on the current allocation process, no modules could be assigned to the following devices due to " - f"insufficient memory:\n" - f"{devices_info}\n" - f"These minimum requirements are specific to this allocation attempt and may vary. Consider increasing " - f"the available memory for these devices to at least the specified minimum, or adjusting the model config." - ) - return device_map - - -def check_device_map(model: nn.Module, device_map: dict[str, Union[int, str, torch.device]]): - """ - Checks a device map covers everything in a given model. - - Args: - model (`torch.nn.Module`): The model to check the device map against. - device_map (`Dict[str, Union[int, str, torch.device]]`): The device map to check. - """ - all_module_names = dict(model.named_modules()) - invalid_keys = [k for k in device_map if k != "" and k not in all_module_names] - - if invalid_keys: - warnings.warn( - f"The following device_map keys do not match any submodules in the model: {invalid_keys}", UserWarning - ) - - all_model_tensors = [name for name, _ in model.state_dict().items()] - for module_name in device_map.keys(): - if module_name == "": - all_model_tensors.clear() - break - else: - all_model_tensors = [ - name - for name in all_model_tensors - if not name == module_name and not name.startswith(module_name + ".") - ] - if len(all_model_tensors) > 0: - non_covered_params = ", ".join(all_model_tensors) - raise ValueError( - f"The device_map provided does not give any device for the following parameters: {non_covered_params}" - ) - - -def load_state_dict(checkpoint_file, device_map=None): - """ - Load a checkpoint from a given file. If the checkpoint is in the safetensors format and a device map is passed, the - weights can be fast-loaded directly on the GPU. - - Args: - checkpoint_file (`str`): The path to the checkpoint to load. - device_map (`Dict[str, Union[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. - """ - if checkpoint_file.endswith(".safetensors"): - with safe_open(checkpoint_file, framework="pt") as f: - metadata = f.metadata() - weight_names = f.keys() - - if metadata is None: - logger.warning( - f"The safetensors archive passed at {checkpoint_file} does not contain metadata. " - "Make sure to save your model with the `save_pretrained` method. Defaulting to 'pt' metadata." - ) - metadata = {"format": "pt"} - - if metadata.get("format") not in ["pt", "tf", "flax"]: - raise OSError( - f"The safetensors archive passed at {checkpoint_file} does not contain the valid metadata. Make sure " - "you save your model with the `save_pretrained` method." - ) - elif metadata["format"] != "pt": - raise ValueError(f"The checkpoint passed was saved with {metadata['format']}, we need a the pt format.") - if device_map is None: - return safe_load_file(checkpoint_file) - else: - # if we only have one device we can load everything directly - if len(set(device_map.values())) == 1: - device = list(device_map.values())[0] - target_device = device - if isinstance(device, int): - if is_npu_available(): - target_device = f"npu:{device}" - elif is_hpu_available(): - target_device = "hpu" - - return safe_load_file(checkpoint_file, device=target_device) - - devices = list(set(device_map.values()) - {"disk"}) - # cpu device should always exist as fallback option - if "cpu" not in devices: - devices.append("cpu") - - # For each device, get the weights that go there - device_weights = {device: [] for device in devices} - for module_name, device in device_map.items(): - if device in devices: - device_weights[device].extend( - [k for k in weight_names if k == module_name or k.startswith(module_name + ".")] - ) - - # all weights that haven't defined a device should be loaded on CPU - device_weights["cpu"].extend([k for k in weight_names if k not in sum(device_weights.values(), [])]) - tensors = {} - if is_tqdm_available(): - progress_bar = tqdm( - main_process_only=False, - total=sum([len(device_weights[device]) for device in devices]), - unit="w", - smoothing=0, - leave=False, - ) - else: - progress_bar = None - for device in devices: - target_device = device - if isinstance(device, int): - if is_npu_available(): - target_device = f"npu:{device}" - elif is_hpu_available(): - target_device = "hpu" - - with safe_open(checkpoint_file, framework="pt", device=target_device) as f: - for key in device_weights[device]: - if progress_bar is not None: - progress_bar.set_postfix(dev=device, refresh=False) - progress_bar.set_description(key) - tensors[key] = f.get_tensor(key) - if progress_bar is not None: - progress_bar.update() - if progress_bar is not None: - progress_bar.close() - - return tensors - else: - return torch.load(checkpoint_file, map_location=torch.device("cpu"), weights_only=True) - - -def get_state_dict_offloaded_model(model: nn.Module): - """ - Returns the state dictionary for an offloaded model via iterative onloading - - Args: - model (`torch.nn.Module`): - The offloaded model we want to save - """ - - state_dict = {} - placeholders = set() - for name, module in model.named_modules(): - if name == "": - continue - - try: - with align_module_device(module, "cpu"): - module_state_dict = module.state_dict() - except MemoryError: - raise MemoryError("Offloaded module must fit in CPU memory to call save_model!") from None - - for key in module_state_dict: - # ignore placeholder parameters that are still on the meta device - if module_state_dict[key].device == torch.device("meta"): - placeholders.add(name + f".{key}") - continue - params = module_state_dict[key] - state_dict[name + f".{key}"] = params.to("cpu") # move buffers to cpu - for key in placeholders.copy(): - if key in state_dict: - placeholders.remove(key) - if placeholders: - logger.warning(f"The following tensors were not saved because they were still on meta device: {placeholders}") - - return state_dict - - -def get_state_dict_from_offload( - module: nn.Module, - module_name: str, - state_dict: dict[str, Union[str, torch.tensor]], - device_to_put_offload: Union[int, str, torch.device] = "cpu", -): - """ - Retrieve the state dictionary (with parameters) from an offloaded module and load into a specified device (defaults - to cpu). - - Args: - module: (`torch.nn.Module`): - The module we want to retrieve a state dictionary from - module_name: (`str`): - The name of the module of interest - state_dict (`Dict[str, Union[int, str, torch.device]]`): - Dictionary of {module names: parameters} - device_to_put_offload (`Union[int, str, torch.device]`): - Device to load offloaded parameters into, defaults to the cpu. - """ - - root = module_name[: module_name.rfind(".")] # module name without .weight or .bias - - # do not move parameters if the module is not offloaded - if not has_offloaded_params(module): - device_to_put_offload = None - - # assign the device to which the offloaded parameters will be sent - with align_module_device(module, device_to_put_offload): - for m_key, params in module.state_dict().items(): - if (root + f".{m_key}") in state_dict: - state_dict[root + f".{m_key}"] = params - - return state_dict - - -def load_checkpoint_in_model( - model: nn.Module, - checkpoint: Union[str, os.PathLike], - device_map: Optional[dict[str, Union[int, str, torch.device]]] = None, - offload_folder: Optional[Union[str, os.PathLike]] = None, - dtype: Optional[Union[str, torch.dtype]] = None, - offload_state_dict: bool = False, - offload_buffers: bool = False, - keep_in_fp32_modules: list[str] = None, - offload_8bit_bnb: bool = False, - strict: bool = False, - full_state_dict: bool = True, - broadcast_from_rank0: bool = False, -): - """ - Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are - loaded. - - - - Once loaded across devices, you still need to call [`dispatch_model`] on your model to make it able to run. To - group the checkpoint loading and dispatch in one single call, use [`load_checkpoint_and_dispatch`]. - - - - Args: - model (`torch.nn.Module`): - The model in which we want to load a checkpoint. - checkpoint (`str` or `os.PathLike`): - The folder checkpoint to load. It can be: - - a path to a file containing a whole model state dict - - a path to a `.json` file containing the index to a sharded checkpoint - - a path to a folder containing a unique `.index.json` file and the shards of a checkpoint. - - a path to a folder containing a unique pytorch_model.bin or a model.safetensors file. - device_map (`Dict[str, Union[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. - offload_folder (`str` or `os.PathLike`, *optional*): - If the `device_map` contains any value `"disk"`, the folder where we will offload weights. - dtype (`str` or `torch.dtype`, *optional*): - If provided, the weights will be converted to that type when loaded. - offload_state_dict (`bool`, *optional*, defaults to `False`): - If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if - the weight of the CPU state dict + the biggest shard does not fit. - offload_buffers (`bool`, *optional*, defaults to `False`): - Whether or not to include the buffers in the weights offloaded to disk. - keep_in_fp32_modules(`List[str]`, *optional*): - A list of the modules that we keep in `torch.float32` dtype. - offload_8bit_bnb (`bool`, *optional*): - Whether or not to enable offload of 8-bit modules on cpu/disk. - strict (`bool`, *optional*, defaults to `False`): - Whether to strictly enforce that the keys in the checkpoint state_dict match the keys of the model's - state_dict. - full_state_dict (`bool`, *optional*, defaults to `True`): if this is set to `True`, all the tensors in the - loaded state_dict will be gathered. No ShardedTensor and DTensor will be in the loaded state_dict. - broadcast_from_rank0 (`False`, *optional*, defaults to `False`): when the option is `True`, a distributed - `ProcessGroup` must be initialized. rank0 should receive a full state_dict and will broadcast the tensors - in the state_dict one by one to other ranks. Other ranks will receive the tensors and shard (if applicable) - according to the local shards in the model. - - """ - if offload_8bit_bnb: - from .bnb import quantize_and_offload_8bit - - tied_params = find_tied_parameters(model) - - if check_tied_parameters_in_config(model) and len(tied_params) == 0: - logger.warning( - "The model weights are not tied. Please use the `tie_weights` method before using the `infer_auto_device` function." - ) - if device_map is not None: - check_tied_parameters_on_same_device(tied_params, device_map) - - if offload_folder is None and device_map is not None and "disk" in device_map.values(): - raise ValueError( - "At least one of the model submodule will be offloaded to disk, please pass along an `offload_folder`." - ) - elif offload_folder is not None and device_map is not None and "disk" in device_map.values(): - os.makedirs(offload_folder, exist_ok=True) - - if isinstance(dtype, str): - # We accept "torch.float16" or just "float16" - dtype = dtype.replace("torch.", "") - dtype = getattr(torch, dtype) - - checkpoint_files = None - index_filename = None - if os.path.isfile(checkpoint): - if str(checkpoint).endswith(".json"): - index_filename = checkpoint - else: - checkpoint_files = [checkpoint] - elif os.path.isdir(checkpoint): - # check if the whole state dict is present - potential_state_bin = [f for f in os.listdir(checkpoint) if f == WEIGHTS_NAME] - potential_state_safetensor = [f for f in os.listdir(checkpoint) if f == SAFE_WEIGHTS_NAME] - if len(potential_state_bin) == 1: - checkpoint_files = [os.path.join(checkpoint, potential_state_bin[0])] - elif len(potential_state_safetensor) == 1: - checkpoint_files = [os.path.join(checkpoint, potential_state_safetensor[0])] - else: - # otherwise check for sharded checkpoints - potential_index = [f for f in os.listdir(checkpoint) if f.endswith(".index.json")] - if len(potential_index) == 0: - raise ValueError( - f"{checkpoint} is not a folder containing a `.index.json` file or a {WEIGHTS_NAME} or a {SAFE_WEIGHTS_NAME} file" - ) - elif len(potential_index) == 1: - index_filename = os.path.join(checkpoint, potential_index[0]) - else: - raise ValueError( - f"{checkpoint} containing more than one `.index.json` file, delete the irrelevant ones." - ) - else: - raise ValueError( - "`checkpoint` should be the path to a file containing a whole state dict, or the index of a sharded " - f"checkpoint, or a folder containing a sharded checkpoint or the whole state dict, but got {checkpoint}." - ) - - if index_filename is not None: - checkpoint_folder = os.path.split(index_filename)[0] - with open(index_filename) as f: - index = json.loads(f.read()) - - if "weight_map" in index: - index = index["weight_map"] - checkpoint_files = sorted(list(set(index.values()))) - checkpoint_files = [os.path.join(checkpoint_folder, f) for f in checkpoint_files] - - # Logic for missing/unexepected keys goes here. - - offload_index = {} - if offload_state_dict: - state_dict_folder = tempfile.mkdtemp() - state_dict_index = {} - - unexpected_keys = set() - model_keys = set(model.state_dict().keys()) - buffer_names = [name for name, _ in model.named_buffers()] - model_devices = {t.device for t in model.state_dict().values() if isinstance(t, torch.Tensor)} - model_physical_devices = model_devices - {torch.device("meta")} - for checkpoint_file in checkpoint_files: - if device_map is None: - # exception for multi-device loading was made for the meta device in torch v2.7.0 - # https://github.com/pytorch/pytorch/blob/v2.6.0/torch/distributed/checkpoint/state_dict.py#L557-L563 - # https://github.com/pytorch/pytorch/blob/v2.7.0-rc2/torch/distributed/checkpoint/state_dict.py#L575-L587 - if is_torch_version(">=", "2.2.0") and ( - (is_torch_version(">=", "2.7.0") and len(model_physical_devices) <= 1) or len(model_devices) <= 1 - ): - from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict - - broadcast_from_rank0 &= is_torch_version(">=", "2.4.0") - loaded_checkpoint = ( - load_state_dict(checkpoint_file, device_map=device_map) - if not broadcast_from_rank0 or dist.get_rank() == 0 - else {} - ) - set_model_state_dict( - model, - loaded_checkpoint, - options=StateDictOptions( - full_state_dict=full_state_dict, - strict=strict, - **({"broadcast_from_rank0": broadcast_from_rank0} if is_torch_version(">=", "2.4.0") else {}), - ), - ) - else: - loaded_checkpoint = load_state_dict(checkpoint_file, device_map=device_map) - model.load_state_dict(loaded_checkpoint, strict=strict) - - unexpected_keys.update(set(loaded_checkpoint.keys()) - model_keys) - else: - loaded_checkpoint = load_state_dict(checkpoint_file, device_map=device_map) - - for param_name, param in loaded_checkpoint.items(): - # skip SCB parameter (for 8-bit serialization) - if "SCB" in param_name: - continue - - if param_name not in model_keys: - unexpected_keys.add(param_name) - if not strict: - continue # Skip loading this parameter. - - module_name = param_name - - while len(module_name) > 0 and module_name not in device_map: - module_name = ".".join(module_name.split(".")[:-1]) - if module_name == "" and "" not in device_map: - # TODO: group all errors and raise at the end. - raise ValueError(f"{param_name} doesn't have any device set.") - param_device = device_map[module_name] - new_dtype = dtype - if dtype is not None and torch.is_floating_point(param): - if keep_in_fp32_modules is not None and dtype == torch.float16: - proceed = False - for key in keep_in_fp32_modules: - if ((key in param_name) and (key + "." in param_name)) or key == param_name: - proceed = True - break - if proceed: - new_dtype = torch.float32 - - if "weight" in param_name and param_name.replace("weight", "SCB") in loaded_checkpoint.keys(): - if param.dtype == torch.int8: - fp16_statistics = loaded_checkpoint[param_name.replace("weight", "SCB")] - else: - fp16_statistics = None - - if param_device == "disk": - if offload_buffers or param_name not in buffer_names: - if new_dtype is None: - new_dtype = param.dtype - if offload_8bit_bnb: - quantize_and_offload_8bit( - model, param, param_name, new_dtype, offload_folder, offload_index, fp16_statistics - ) - continue - else: - set_module_tensor_to_device(model, param_name, "meta", dtype=new_dtype) - offload_weight(param, param_name, offload_folder, index=offload_index) - elif param_device == "cpu" and offload_state_dict: - if new_dtype is None: - new_dtype = param.dtype - if offload_8bit_bnb: - quantize_and_offload_8bit( - model, param, param_name, new_dtype, state_dict_folder, state_dict_index, fp16_statistics - ) - else: - set_module_tensor_to_device(model, param_name, "meta", dtype=new_dtype) - offload_weight(param, param_name, state_dict_folder, index=state_dict_index) - else: - set_module_tensor_to_device( - model, - param_name, - param_device, - value=param, - dtype=new_dtype, - fp16_statistics=fp16_statistics, - ) - - # Force Python to clean up. - del loaded_checkpoint - gc.collect() - - if not strict and len(unexpected_keys) > 0: - logger.warning( - f"Some weights of the model checkpoint at {checkpoint} were not used when" - f" initializing {model.__class__.__name__}: {unexpected_keys}. This may or may not be an issue - make sure that the checkpoint does not have unnecessary parameters, or that the model definition correctly corresponds to the checkpoint." - ) - - save_offload_index(offload_index, offload_folder) - - # Load back offloaded state dict on CPU - if offload_state_dict: - load_offloaded_weights(model, state_dict_index, state_dict_folder) - shutil.rmtree(state_dict_folder) - - retie_parameters(model, tied_params) - - -def get_mixed_precision_context_manager(native_amp: bool = False, autocast_kwargs: AutocastKwargs = None): - """ - Return a context manager for autocasting mixed precision - - Args: - native_amp (`bool`, *optional*, defaults to False): - Whether mixed precision is actually enabled. - cache_enabled (`bool`, *optional*, defaults to True): - Whether the weight cache inside autocast should be enabled. - """ - state = AcceleratorState() - if autocast_kwargs is None: - autocast_kwargs = {} - else: - autocast_kwargs = autocast_kwargs.to_kwargs() - if native_amp: - device_type = ( - "cuda" - if (state.distributed_type == DistributedType.XLA and is_torch_xla_available(check_is_gpu=True)) - else state.device.type - ) - if state.mixed_precision == "fp16": - return torch.autocast(device_type=device_type, dtype=torch.float16, **autocast_kwargs) - elif state.mixed_precision in ["bf16", "fp8"] and state.distributed_type in [ - DistributedType.NO, - DistributedType.MULTI_CPU, - DistributedType.MULTI_GPU, - DistributedType.MULTI_MLU, - DistributedType.MULTI_SDAA, - DistributedType.MULTI_MUSA, - DistributedType.MULTI_NPU, - DistributedType.MULTI_XPU, - DistributedType.MULTI_HPU, - DistributedType.FSDP, - DistributedType.XLA, - DistributedType.TP, - ]: - return torch.autocast(device_type=device_type, dtype=torch.bfloat16, **autocast_kwargs) - else: - return torch.autocast(device_type=device_type, **autocast_kwargs) - else: - return contextlib.nullcontext() - - -def get_grad_scaler(distributed_type: DistributedType = None, **kwargs): - """ - A generic helper which will initialize the correct `GradScaler` implementation based on the environment and return - it. - - Args: - distributed_type (`DistributedType`, *optional*, defaults to None): - The type of distributed environment. - kwargs: - Additional arguments for the utilized `GradScaler` constructor. - """ - if distributed_type == DistributedType.FSDP: - from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler - - return ShardedGradScaler(**kwargs) - if is_torch_xla_available(check_is_gpu=True): - import torch_xla.amp as xamp - - return xamp.GradScaler(**kwargs) - elif is_mlu_available(): - return torch.mlu.amp.GradScaler(**kwargs) - elif is_sdaa_available(): - return torch.sdaa.amp.GradScaler(**kwargs) - elif is_musa_available(): - return torch.musa.amp.GradScaler(**kwargs) - elif is_npu_available(): - return torch.npu.amp.GradScaler(**kwargs) - elif is_hpu_available(): - return torch.amp.GradScaler("hpu", **kwargs) - elif is_xpu_available(): - return torch.amp.GradScaler("xpu", **kwargs) - else: - if is_torch_version(">=", "2.3"): - return torch.amp.GradScaler("cuda", **kwargs) - else: - return torch.cuda.amp.GradScaler(**kwargs) - - -def has_offloaded_params(module: torch.nn.Module) -> bool: - """ - Checks if a module has offloaded parameters by checking if the given module has a AlignDevicesHook attached with - offloading enabled - - Args: - module (`torch.nn.Module`): The module to check for an offload hook. - - Returns: - bool: `True` if the module has an offload hook and offloading is enabled, `False` otherwise. - """ - from ..hooks import AlignDevicesHook # avoid circular import - - return hasattr(module, "_hf_hook") and isinstance(module._hf_hook, AlignDevicesHook) and module._hf_hook.offload - - -@contextlib.contextmanager -def align_module_device(module: torch.nn.Module, execution_device: Optional[torch.device] = None): - """ - Context manager that moves a module's parameters to the specified execution device. - - Args: - module (`torch.nn.Module`): - Module with parameters to align. - execution_device (`torch.device`, *optional*): - If provided, overrides the module's execution device within the context. Otherwise, use hook execution - device or pass - """ - if has_offloaded_params(module): - if execution_device is not None: - original_device = module._hf_hook.execution_device - module._hf_hook.execution_device = execution_device - - try: - module._hf_hook.pre_forward(module) - yield - finally: - module._hf_hook.post_forward(module, None) - if execution_device is not None: - module._hf_hook.execution_device = original_device - - elif execution_device is not None: - devices = {name: param.device for name, param in module.named_parameters(recurse=False)} - try: - for name in devices: - set_module_tensor_to_device(module, name, execution_device) - yield - finally: - for name, device in devices.items(): - set_module_tensor_to_device(module, name, device) - - else: - yield diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/offload.py b/venv/lib/python3.10/site-packages/accelerate/utils/offload.py deleted file mode 100644 index d8bff7dc6ad41ffc7f14555261d115d51b76ccec..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/offload.py +++ /dev/null @@ -1,213 +0,0 @@ -# Copyright 2022 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 -from collections.abc import Mapping -from typing import Optional, Union - -import numpy as np -import torch -from safetensors import safe_open - - -def offload_weight(weight, weight_name, offload_folder, index=None): - dtype = None - # Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16. - if str(weight.dtype) == "torch.bfloat16": - # Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s. - weight = weight.view(torch.int16) - dtype = "bfloat16" - array = weight.cpu().numpy() - tensor_file = os.path.join(offload_folder, f"{weight_name}.dat") - if index is not None: - if dtype is None: - dtype = str(array.dtype) - index[weight_name] = {"dtype": dtype, "shape": list(array.shape)} - if array.ndim == 0: - array = array[None] - file_array = np.memmap(tensor_file, dtype=array.dtype, mode="w+", shape=array.shape) - file_array[:] = array[:] - file_array.flush() - return index - - -def load_offloaded_weight(weight_file, weight_info): - shape = tuple(weight_info["shape"]) - if shape == (): - # NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor - shape = (1,) - - dtype = weight_info["dtype"] - if dtype == "bfloat16": - # NumPy does not support bfloat16 so this was saved as a int16 - dtype = "int16" - - weight = np.memmap(weight_file, dtype=dtype, shape=shape, mode="r") - - if len(weight_info["shape"]) == 0: - weight = weight[0] - weight = torch.tensor(weight) - if weight_info["dtype"] == "bfloat16": - weight = weight.view(torch.bfloat16) - - return weight - - -def save_offload_index(index, offload_folder): - if index is None or len(index) == 0: - # Nothing to save - return - - offload_index_file = os.path.join(offload_folder, "index.json") - if os.path.isfile(offload_index_file): - with open(offload_index_file, encoding="utf-8") as f: - current_index = json.load(f) - else: - current_index = {} - current_index.update(index) - - with open(offload_index_file, "w", encoding="utf-8") as f: - json.dump(current_index, f, indent=2) - - -def offload_state_dict(save_dir: Union[str, os.PathLike], state_dict: dict[str, torch.Tensor]): - """ - Offload a state dict in a given folder. - - Args: - save_dir (`str` or `os.PathLike`): - The directory in which to offload the state dict. - state_dict (`Dict[str, torch.Tensor]`): - The dictionary of tensors to offload. - """ - os.makedirs(save_dir, exist_ok=True) - index = {} - for name, parameter in state_dict.items(): - index = offload_weight(parameter, name, save_dir, index=index) - - # Update index - save_offload_index(index, save_dir) - - -class PrefixedDataset(Mapping): - """ - Will access keys in a given dataset by adding a prefix. - - Args: - dataset (`Mapping`): Any map with string keys. - prefix (`str`): A prefix to add when trying to access any element in the underlying dataset. - """ - - def __init__(self, dataset: Mapping, prefix: str): - self.dataset = dataset - self.prefix = prefix - - def __getitem__(self, key): - return self.dataset[f"{self.prefix}{key}"] - - def __iter__(self): - return iter([key for key in self.dataset if key.startswith(self.prefix)]) - - def __len__(self): - return len(self.dataset) - - -class OffloadedWeightsLoader(Mapping): - """ - A collection that loads weights stored in a given state dict or memory-mapped on disk. - - Args: - state_dict (`Dict[str, torch.Tensor]`, *optional*): - A dictionary parameter name to tensor. - save_folder (`str` or `os.PathLike`, *optional*): - The directory in which the weights are stored (by `offload_state_dict` for instance). - index (`Dict`, *optional*): - A dictionary from weight name to their information (`dtype`/ `shape` or safetensors filename). Will default - to the index saved in `save_folder`. - """ - - def __init__( - self, - state_dict: dict[str, torch.Tensor] = None, - save_folder: Optional[Union[str, os.PathLike]] = None, - index: Mapping = None, - device=None, - ): - if state_dict is None and save_folder is None and index is None: - raise ValueError("Need either a `state_dict`, a `save_folder` or an `index` containing offloaded weights.") - - self.state_dict = {} if state_dict is None else state_dict - self.save_folder = save_folder - if index is None and save_folder is not None: - with open(os.path.join(save_folder, "index.json")) as f: - index = json.load(f) - self.index = {} if index is None else index - self.all_keys = list(self.state_dict.keys()) - self.all_keys.extend([key for key in self.index if key not in self.all_keys]) - self.device = device - - def __getitem__(self, key: str): - # State dict gets priority - if key in self.state_dict: - return self.state_dict[key] - weight_info = self.index[key] - if weight_info.get("safetensors_file") is not None: - device = "cpu" if self.device is None else self.device - tensor = None - try: - with safe_open(weight_info["safetensors_file"], framework="pt", device=device) as f: - tensor = f.get_tensor(weight_info.get("weight_name", key)) - except TypeError: - # if failed to get_tensor on the device, such as bf16 on mps, try to load it on CPU first - with safe_open(weight_info["safetensors_file"], framework="pt", device="cpu") as f: - tensor = f.get_tensor(weight_info.get("weight_name", key)) - - if "dtype" in weight_info: - tensor = tensor.to(getattr(torch, weight_info["dtype"])) - - if tensor.device != torch.device(device): - tensor = tensor.to(device) - return tensor - - weight_file = os.path.join(self.save_folder, f"{key}.dat") - return load_offloaded_weight(weight_file, weight_info) - - def __iter__(self): - return iter(self.all_keys) - - def __len__(self): - return len(self.all_keys) - - -def extract_submodules_state_dict(state_dict: dict[str, torch.Tensor], submodule_names: list[str]): - """ - Extract the sub state-dict corresponding to a list of given submodules. - - Args: - state_dict (`Dict[str, torch.Tensor]`): The state dict to extract from. - submodule_names (`List[str]`): The list of submodule names we want to extract. - """ - result = {} - for module_name in submodule_names: - # We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the - # submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance) - result.update( - { - key: param - for key, param in state_dict.items() - if key == module_name or key.startswith(module_name + ".") - } - ) - return result diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/operations.py b/venv/lib/python3.10/site-packages/accelerate/utils/operations.py deleted file mode 100644 index 088c0f6efd2a8d7e64a778d70450dd0966984446..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/operations.py +++ /dev/null @@ -1,866 +0,0 @@ -# Copyright 2022 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. -""" -A set of basic tensor ops compatible with tpu, gpu, and multigpu -""" - -import pickle -import warnings -from collections.abc import Mapping -from contextlib import contextmanager, nullcontext -from functools import update_wrapper, wraps -from typing import Any - -import torch - -from ..state import AcceleratorState, PartialState -from .constants import TORCH_DISTRIBUTED_OPERATION_TYPES -from .dataclasses import DistributedType, TensorInformation -from .imports import ( - is_npu_available, - is_torch_distributed_available, - is_torch_xla_available, -) - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - -if is_torch_distributed_available(): - from torch.distributed import ReduceOp - - -def is_torch_tensor(tensor): - return isinstance(tensor, torch.Tensor) - - -def is_torch_xpu_tensor(tensor): - return isinstance( - tensor, - torch.xpu.FloatTensor, - torch.xpu.ByteTensor, - torch.xpu.IntTensor, - torch.xpu.LongTensor, - torch.xpu.HalfTensor, - torch.xpu.DoubleTensor, - torch.xpu.BFloat16Tensor, - ) - - -def is_tensor_information(tensor_info): - return isinstance(tensor_info, TensorInformation) - - -def is_namedtuple(data): - """ - Checks if `data` is a `namedtuple` or not. Can have false positives, but only if a user is trying to mimic a - `namedtuple` perfectly. - """ - return isinstance(data, tuple) and hasattr(data, "_asdict") and hasattr(data, "_fields") - - -def honor_type(obj, generator): - """ - Cast a generator to the same type as obj (list, tuple, or namedtuple) - """ - # Some objects may not be able to instantiate from a generator directly - if is_namedtuple(obj): - return type(obj)(*list(generator)) - else: - return type(obj)(generator) - - -def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_other_type=False, **kwargs): - """ - Recursively apply a function on a data structure that is a nested list/tuple/dictionary of a given base type. - - Args: - func (`callable`): - The function to recursively apply. - data (nested list/tuple/dictionary of `main_type`): - The data on which to apply `func` - *args: - Positional arguments that will be passed to `func` when applied on the unpacked data. - main_type (`type`, *optional*, defaults to `torch.Tensor`): - The base type of the objects to which apply `func`. - error_on_other_type (`bool`, *optional*, defaults to `False`): - Whether to return an error or not if after unpacking `data`, we get on an object that is not of type - `main_type`. If `False`, the function will leave objects of types different than `main_type` unchanged. - **kwargs (additional keyword arguments, *optional*): - Keyword arguments that will be passed to `func` when applied on the unpacked data. - - Returns: - The same data structure as `data` with `func` applied to every object of type `main_type`. - """ - if isinstance(data, (tuple, list)): - return honor_type( - data, - ( - recursively_apply( - func, o, *args, test_type=test_type, error_on_other_type=error_on_other_type, **kwargs - ) - for o in data - ), - ) - elif isinstance(data, Mapping): - return type(data)( - { - k: recursively_apply( - func, v, *args, test_type=test_type, error_on_other_type=error_on_other_type, **kwargs - ) - for k, v in data.items() - } - ) - elif test_type(data): - return func(data, *args, **kwargs) - elif error_on_other_type: - raise TypeError( - f"Unsupported types ({type(data)}) passed to `{func.__name__}`. Only nested list/tuple/dicts of " - f"objects that are valid for `{test_type.__name__}` should be passed." - ) - return data - - -def send_to_device(tensor, device, non_blocking=False, skip_keys=None): - """ - Recursively sends the elements in a nested list/tuple/dictionary of tensors to a given device. - - Args: - tensor (nested list/tuple/dictionary of `torch.Tensor`): - The data to send to a given device. - device (`torch.device`): - The device to send the data to. - - Returns: - The same data structure as `tensor` with all tensors sent to the proper device. - """ - if is_torch_tensor(tensor) or hasattr(tensor, "to"): - # `torch.Tensor.to("npu")` could not find context when called for the first time (see this [issue](https://gitee.com/ascend/pytorch/issues/I8KECW?from=project-issue)). - if device == "npu": - device = "npu:0" - try: - return tensor.to(device, non_blocking=non_blocking) - except TypeError: # .to() doesn't accept non_blocking as kwarg - return tensor.to(device) - except AssertionError as error: - # `torch.Tensor.to()` is not supported by `torch_npu` (see this [issue](https://github.com/Ascend/pytorch/issues/16)). - # This call is inside the try-block since is_npu_available is not supported by torch.compile. - if is_npu_available(): - if isinstance(device, int): - device = f"npu:{device}" - else: - raise error - try: - return tensor.to(device, non_blocking=non_blocking) - except TypeError: # .to() doesn't accept non_blocking as kwarg - return tensor.to(device) - elif isinstance(tensor, (tuple, list)): - return honor_type( - tensor, (send_to_device(t, device, non_blocking=non_blocking, skip_keys=skip_keys) for t in tensor) - ) - elif isinstance(tensor, Mapping): - if isinstance(skip_keys, str): - skip_keys = [skip_keys] - elif skip_keys is None: - skip_keys = [] - return type(tensor)( - { - k: t if k in skip_keys else send_to_device(t, device, non_blocking=non_blocking, skip_keys=skip_keys) - for k, t in tensor.items() - } - ) - else: - return tensor - - -def get_data_structure(data): - """ - Recursively gathers the information needed to rebuild a nested list/tuple/dictionary of tensors. - - Args: - data (nested list/tuple/dictionary of `torch.Tensor`): - The data to send to analyze. - - Returns: - The same data structure as `data` with [`~utils.TensorInformation`] instead of tensors. - """ - - def _get_data_structure(tensor): - return TensorInformation(shape=tensor.shape, dtype=tensor.dtype) - - return recursively_apply(_get_data_structure, data) - - -def get_shape(data): - """ - Recursively gathers the shape of a nested list/tuple/dictionary of tensors as a list. - - Args: - data (nested list/tuple/dictionary of `torch.Tensor`): - The data to send to analyze. - - Returns: - The same data structure as `data` with lists of tensor shapes instead of tensors. - """ - - def _get_shape(tensor): - return list(tensor.shape) - - return recursively_apply(_get_shape, data) - - -def initialize_tensors(data_structure): - """ - Recursively initializes tensors from a nested list/tuple/dictionary of [`~utils.TensorInformation`]. - - Returns: - The same data structure as `data` with tensors instead of [`~utils.TensorInformation`]. - """ - - def _initialize_tensor(tensor_info): - return torch.empty(*tensor_info.shape, dtype=tensor_info.dtype) - - return recursively_apply(_initialize_tensor, data_structure, test_type=is_tensor_information) - - -def find_batch_size(data): - """ - Recursively finds the batch size in a nested list/tuple/dictionary of lists of tensors. - - Args: - data (nested list/tuple/dictionary of `torch.Tensor`): The data from which to find the batch size. - - Returns: - `int`: The batch size. - """ - if isinstance(data, (tuple, list, Mapping)) and (len(data) == 0): - raise ValueError(f"Cannot find the batch size from empty {type(data)}.") - - if isinstance(data, (tuple, list)): - return find_batch_size(data[0]) - elif isinstance(data, Mapping): - for k in data.keys(): - return find_batch_size(data[k]) - elif not isinstance(data, torch.Tensor): - raise TypeError(f"Can only find the batch size of tensors but got {type(data)}.") - return data.shape[0] - - -def ignorant_find_batch_size(data): - """ - Same as [`utils.operations.find_batch_size`] except will ignore if `ValueError` and `TypeErrors` are raised - - Args: - data (nested list/tuple/dictionary of `torch.Tensor`): The data from which to find the batch size. - - Returns: - `int`: The batch size. - """ - try: - return find_batch_size(data) - except (ValueError, TypeError): - pass - return None - - -def listify(data): - """ - Recursively finds tensors in a nested list/tuple/dictionary and converts them to a list of numbers. - - Args: - data (nested list/tuple/dictionary of `torch.Tensor`): The data from which to convert to regular numbers. - - Returns: - The same data structure as `data` with lists of numbers instead of `torch.Tensor`. - """ - - def _convert_to_list(tensor): - tensor = tensor.detach().cpu() - if tensor.dtype == torch.bfloat16: - # As of Numpy 1.21.4, NumPy does not support bfloat16 (see - # https://github.com/numpy/numpy/blob/a47ecdea856986cd60eabbd53265c2ca5916ad5d/doc/source/user/basics.types.rst ). - # Until Numpy adds bfloat16, we must convert float32. - tensor = tensor.to(torch.float32) - return tensor.tolist() - - return recursively_apply(_convert_to_list, data) - - -def _tpu_gather(tensor): - def _tpu_gather_one(tensor): - if tensor.ndim == 0: - tensor = tensor.clone()[None] - - # Can only gather contiguous tensors - if not tensor.is_contiguous(): - tensor = tensor.contiguous() - return xm.all_gather(tensor) - - res = recursively_apply(_tpu_gather_one, tensor, error_on_other_type=True) - xm.mark_step() - return res - - -def _gpu_gather(tensor): - state = PartialState() - gather_op = torch.distributed.all_gather_into_tensor - - # FIXME: the below 2 lines are added to work-aound a bug related to INT64 collectives in oneCCL. Remove them once pytorch-2.9 is released. - if state.device.type == "xpu": - torch.xpu.synchronize() - - def _gpu_gather_one(tensor): - if tensor.ndim == 0: - tensor = tensor.clone()[None] - - # Can only gather contiguous tensors - if not tensor.is_contiguous(): - tensor = tensor.contiguous() - - if state.backend is not None and state.backend != "gloo": - # We use `empty` as `all_gather_into_tensor` slightly - # differs from `all_gather` for better efficiency, - # and we rely on the number of items in the tensor - # rather than its direct shape - output_tensors = torch.empty( - state.num_processes * tensor.numel(), - dtype=tensor.dtype, - device=state.device, - ) - gather_op(output_tensors, tensor) - return output_tensors.view(-1, *tensor.size()[1:]) - else: - # a backend of `None` is always CPU - # also gloo does not support `all_gather_into_tensor`, - # which will result in a larger memory overhead for the op - output_tensors = [torch.empty_like(tensor) for _ in range(state.num_processes)] - torch.distributed.all_gather(output_tensors, tensor) - return torch.cat(output_tensors, dim=0) - - return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True) - - -class DistributedOperationException(Exception): - """ - An exception class for distributed operations. Raised if the operation cannot be performed due to the shape of the - tensors. - """ - - pass - - -def verify_operation(function): - """ - Verifies that `tensor` is the same shape across all processes. Only ran if `PartialState().debug` is `True`. - """ - - @wraps(function) - def wrapper(*args, **kwargs): - if PartialState().distributed_type == DistributedType.NO or not PartialState().debug: - return function(*args, **kwargs) - operation = f"{function.__module__}.{function.__name__}" - if "tensor" in kwargs: - tensor = kwargs["tensor"] - else: - tensor = args[0] - if PartialState().device.type != find_device(tensor).type: - raise DistributedOperationException( - f"One or more of the tensors passed to {operation} were not on the {tensor.device.type} while the `Accelerator` is configured for {PartialState().device.type}. " - f"Please move it to the {PartialState().device.type} before calling {operation}." - ) - shapes = get_shape(tensor) - output = gather_object([shapes]) - if output[0] is not None: - are_same = output.count(output[0]) == len(output) - if not are_same: - process_shape_str = "\n - ".join([f"Process {i}: {shape}" for i, shape in enumerate(output)]) - raise DistributedOperationException( - f"Cannot apply desired operation due to shape mismatches. " - "All shapes across devices must be valid." - f"\n\nOperation: `{operation}`\nInput shapes:\n - {process_shape_str}" - ) - return function(*args, **kwargs) - - return wrapper - - -def chained_operation(function): - """ - Checks that `verify_operation` failed and if so reports a more helpful error chaining the existing - `DistributedOperationException`. - """ - - @wraps(function) - def wrapper(*args, **kwargs): - try: - return function(*args, **kwargs) - except DistributedOperationException as e: - operation = f"{function.__module__}.{function.__name__}" - raise DistributedOperationException( - f"Error found while calling `{operation}`. Please see the earlier error for more details." - ) from e - - return wrapper - - -@verify_operation -def gather(tensor): - """ - Recursively gather tensor in a nested list/tuple/dictionary of tensors from all devices. - - Args: - tensor (nested list/tuple/dictionary of `torch.Tensor`): - The data to gather. - - Returns: - The same data structure as `tensor` with all tensors sent to the proper device. - """ - if PartialState().distributed_type == DistributedType.XLA: - return _tpu_gather(tensor) - elif PartialState().distributed_type in TORCH_DISTRIBUTED_OPERATION_TYPES: - return _gpu_gather(tensor) - else: - return tensor - - -def _gpu_gather_object(object: Any): - output_objects = [None for _ in range(PartialState().num_processes)] - torch.distributed.all_gather_object(output_objects, object) - # all_gather_object returns a list of lists, so we need to flatten it - return [x for y in output_objects for x in y] - - -def gather_object(object: Any): - """ - Recursively gather object in a nested list/tuple/dictionary of objects from all devices. - - Args: - object (nested list/tuple/dictionary of picklable object): - The data to gather. - - Returns: - The same data structure as `object` with all the objects sent to every device. - """ - if PartialState().distributed_type == DistributedType.XLA: - raise NotImplementedError("gather objects in TPU is not supported") - elif PartialState().distributed_type in TORCH_DISTRIBUTED_OPERATION_TYPES: - return _gpu_gather_object(object) - else: - return object - - -def _gpu_broadcast(data, src=0): - def _gpu_broadcast_one(tensor, src=0): - torch.distributed.broadcast(tensor, src=src) - return tensor - - return recursively_apply(_gpu_broadcast_one, data, error_on_other_type=True, src=src) - - -def _tpu_broadcast(tensor, src=0, name="broadcast tensor"): - if isinstance(tensor, (list, tuple)): - return honor_type(tensor, (_tpu_broadcast(t, name=f"{name}_{i}") for i, t in enumerate(tensor))) - elif isinstance(tensor, Mapping): - return type(tensor)({k: _tpu_broadcast(v, name=f"{name}_{k}") for k, v in tensor.items()}) - return xm.mesh_reduce(name, tensor, lambda x: x[src]) - - -TENSOR_TYPE_TO_INT = { - torch.float: 1, - torch.double: 2, - torch.half: 3, - torch.bfloat16: 4, - torch.uint8: 5, - torch.int8: 6, - torch.int16: 7, - torch.int32: 8, - torch.int64: 9, - torch.bool: 10, -} - -TENSOR_INT_TO_DTYPE = {v: k for k, v in TENSOR_TYPE_TO_INT.items()} - - -def gather_tensor_shape(tensor): - """ - Grabs the shape of `tensor` only available on one process and returns a tensor of its shape - """ - # Allocate 80 bytes to store the shape - max_tensor_dimension = 2**20 - state = PartialState() - base_tensor = torch.empty(max_tensor_dimension, dtype=torch.int, device=state.device) - - # Since PyTorch can't just send a tensor to another GPU without - # knowing its size, we store the size of the tensor with data - # in an allocation - if tensor is not None: - shape = tensor.shape - tensor_dtype = TENSOR_TYPE_TO_INT[tensor.dtype] - base_tensor[: len(shape) + 1] = torch.tensor(list(shape) + [tensor_dtype], dtype=int) - # Perform a reduction to copy the size data onto all GPUs - base_tensor = reduce(base_tensor, reduction="sum") - base_tensor = base_tensor[base_tensor.nonzero()] - # The last non-zero data contains the coded dtype the source tensor is - dtype = int(base_tensor[-1:][0]) - base_tensor = base_tensor[:-1] - return base_tensor, dtype - - -def copy_tensor_to_devices(tensor=None) -> torch.Tensor: - """ - Copys a tensor that only exists on a single device and broadcasts it to other devices. Differs from `broadcast` as - each worker doesn't need to know its shape when used (and tensor can be `None`) - - Args: - tensor (`torch.tensor`): - The tensor that should be sent to all devices. Must only have it be defined on a single device, the rest - should be `None`. - """ - state = PartialState() - shape, dtype = gather_tensor_shape(tensor) - if tensor is None: - tensor = torch.zeros(shape, dtype=TENSOR_INT_TO_DTYPE[dtype]).to(state.device) - return reduce(tensor, reduction="sum") - - -@verify_operation -def broadcast(tensor, from_process: int = 0): - """ - Recursively broadcast tensor in a nested list/tuple/dictionary of tensors to all devices. - - Args: - tensor (nested list/tuple/dictionary of `torch.Tensor`): - The data to gather. - from_process (`int`, *optional*, defaults to 0): - The process from which to send the data - - Returns: - The same data structure as `tensor` with all tensors broadcasted to the proper device. - """ - if PartialState().distributed_type == DistributedType.XLA: - return _tpu_broadcast(tensor, src=from_process, name="accelerate.utils.broadcast") - elif PartialState().distributed_type in TORCH_DISTRIBUTED_OPERATION_TYPES: - return _gpu_broadcast(tensor, src=from_process) - else: - return tensor - - -def broadcast_object_list(object_list, from_process: int = 0): - """ - Broadcast a list of picklable objects form one process to the others. - - Args: - object_list (list of picklable objects): - The list of objects to broadcast. This list will be modified inplace. - from_process (`int`, *optional*, defaults to 0): - The process from which to send the data. - - Returns: - The same list containing the objects from process 0. - """ - if PartialState().distributed_type == DistributedType.XLA: - for i, obj in enumerate(object_list): - object_list[i] = xm.mesh_reduce("accelerate.utils.broadcast_object_list", obj, lambda x: x[from_process]) - elif PartialState().distributed_type in TORCH_DISTRIBUTED_OPERATION_TYPES: - torch.distributed.broadcast_object_list(object_list, src=from_process) - return object_list - - -def slice_tensors(data, tensor_slice, process_index=None, num_processes=None): - """ - Recursively takes a slice in a nested list/tuple/dictionary of tensors. - - Args: - data (nested list/tuple/dictionary of `torch.Tensor`): - The data to slice. - tensor_slice (`slice`): - The slice to take. - - Returns: - The same data structure as `data` with all the tensors slices. - """ - - def _slice_tensor(tensor, tensor_slice): - return tensor[tensor_slice] - - return recursively_apply(_slice_tensor, data, tensor_slice) - - -def concatenate(data, dim=0): - """ - Recursively concatenate the tensors in a nested list/tuple/dictionary of lists of tensors with the same shape. - - Args: - data (nested list/tuple/dictionary of lists of tensors `torch.Tensor`): - The data to concatenate. - dim (`int`, *optional*, defaults to 0): - The dimension on which to concatenate. - - Returns: - The same data structure as `data` with all the tensors concatenated. - """ - if isinstance(data[0], (tuple, list)): - return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0])))) - elif isinstance(data[0], Mapping): - return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()}) - elif not isinstance(data[0], torch.Tensor): - raise TypeError(f"Can only concatenate tensors but got {type(data[0])}") - return torch.cat(data, dim=dim) - - -class CannotPadNestedTensorWarning(UserWarning): - pass - - -@chained_operation -def pad_across_processes(tensor, dim=0, pad_index=0, pad_first=False): - """ - Recursively pad the tensors in a nested list/tuple/dictionary of tensors from all devices to the same size so they - can safely be gathered. - - Args: - tensor (nested list/tuple/dictionary of `torch.Tensor`): - The data to gather. - dim (`int`, *optional*, defaults to 0): - The dimension on which to pad. - pad_index (`int`, *optional*, defaults to 0): - The value with which to pad. - pad_first (`bool`, *optional*, defaults to `False`): - Whether to pad at the beginning or the end. - """ - - def _pad_across_processes(tensor, dim=0, pad_index=0, pad_first=False): - if getattr(tensor, "is_nested", False): - warnings.warn( - "Cannot pad nested tensors without more information. Leaving unprocessed.", - CannotPadNestedTensorWarning, - ) - return tensor - if dim >= len(tensor.shape) or dim < -len(tensor.shape): - return tensor - # Convert negative dimensions to non-negative - if dim < 0: - dim += len(tensor.shape) - - # Gather all sizes - size = torch.tensor(tensor.shape, device=tensor.device)[None] - sizes = gather(size).cpu() - # Then pad to the maximum size - max_size = max(s[dim] for s in sizes) - if max_size == tensor.shape[dim]: - return tensor - - old_size = tensor.shape - new_size = list(old_size) - new_size[dim] = max_size - new_tensor = tensor.new_zeros(tuple(new_size)) + pad_index - if pad_first: - indices = tuple( - slice(max_size - old_size[dim], max_size) if i == dim else slice(None) for i in range(len(new_size)) - ) - else: - indices = tuple(slice(0, old_size[dim]) if i == dim else slice(None) for i in range(len(new_size))) - new_tensor[indices] = tensor - return new_tensor - - return recursively_apply( - _pad_across_processes, tensor, error_on_other_type=True, dim=dim, pad_index=pad_index, pad_first=pad_first - ) - - -def pad_input_tensors(tensor, batch_size, num_processes, dim=0): - """ - Takes a `tensor` of arbitrary size and pads it so that it can work given `num_processes` needed dimensions. - - New tensors are just the last input repeated. - - E.g.: - Tensor: ([3,4,4]) Num processes: 4 Expected result shape: ([4,4,4]) - - """ - - def _pad_input_tensors(tensor, batch_size, num_processes, dim=0): - remainder = batch_size // num_processes - last_inputs = batch_size - (remainder * num_processes) - if batch_size // num_processes == 0: - to_pad = num_processes - batch_size - else: - to_pad = num_processes - (batch_size // num_processes) - # In the rare case that `to_pad` is negative, - # we need to pad the last inputs - the found `to_pad` - if last_inputs > to_pad & to_pad < 1: - to_pad = last_inputs - to_pad - old_size = tensor.shape - new_size = list(old_size) - new_size[0] = batch_size + to_pad - new_tensor = tensor.new_zeros(tuple(new_size)) - indices = tuple(slice(0, old_size[dim]) if i == dim else slice(None) for i in range(len(new_size))) - new_tensor[indices] = tensor - return new_tensor - - return recursively_apply( - _pad_input_tensors, - tensor, - error_on_other_type=True, - batch_size=batch_size, - num_processes=num_processes, - dim=dim, - ) - - -@verify_operation -def reduce(tensor, reduction="mean", scale=1.0): - """ - Recursively reduce the tensors in a nested list/tuple/dictionary of lists of tensors across all processes by the - mean of a given operation. - - Args: - tensor (nested list/tuple/dictionary of `torch.Tensor`): - The data to reduce. - reduction (`str`, *optional*, defaults to `"mean"`): - A reduction method. Can be of "mean", "sum", or "none" - scale (`float`, *optional*): - A default scaling value to be applied after the reduce, only valied on XLA. - - Returns: - The same data structure as `data` with all the tensors reduced. - """ - - def _reduce_across_processes(tensor, reduction="mean", scale=1.0): - state = PartialState() - cloned_tensor = tensor.clone() - if state.distributed_type == DistributedType.NO: - return cloned_tensor - if state.distributed_type == DistributedType.XLA: - # Some processes may have different HLO graphs than other - # processes, for example in the breakpoint API - # accelerator.set_trigger(). Use mark_step to make HLOs - # the same on all processes. - xm.mark_step() - xm.all_reduce(xm.REDUCE_SUM, [cloned_tensor], scale) - xm.mark_step() - elif state.distributed_type.value in TORCH_DISTRIBUTED_OPERATION_TYPES: - torch.distributed.all_reduce(cloned_tensor, ReduceOp.SUM) - if reduction == "mean": - cloned_tensor /= state.num_processes - return cloned_tensor - - return recursively_apply( - _reduce_across_processes, tensor, error_on_other_type=True, reduction=reduction, scale=scale - ) - - -def convert_to_fp32(tensor): - """ - Recursively converts the elements nested list/tuple/dictionary of tensors in FP16/BF16 precision to FP32. - - Args: - tensor (nested list/tuple/dictionary of `torch.Tensor`): - The data to convert from FP16/BF16 to FP32. - - Returns: - The same data structure as `tensor` with all tensors that were in FP16/BF16 precision converted to FP32. - """ - - def _convert_to_fp32(tensor): - return tensor.float() - - def _is_fp16_bf16_tensor(tensor): - return (is_torch_tensor(tensor) or hasattr(tensor, "dtype")) and tensor.dtype in ( - torch.float16, - torch.bfloat16, - ) - - return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_bf16_tensor) - - -class ConvertOutputsToFp32: - """ - Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16 - precision will be convert back to FP32. - - Args: - model_forward (`Callable`): - The function which outputs we want to treat. - - Returns: - The same function as `model_forward` but with converted outputs. - """ - - def __init__(self, model_forward): - self.model_forward = model_forward - update_wrapper(self, model_forward) - - def __call__(self, *args, **kwargs): - return convert_to_fp32(self.model_forward(*args, **kwargs)) - - def __getstate__(self): - raise pickle.PicklingError( - "Cannot pickle a prepared model with automatic mixed precision, please unwrap the model with `Accelerator.unwrap_model(model)` before pickling it." - ) - - -def convert_outputs_to_fp32(model_forward): - model_forward = ConvertOutputsToFp32(model_forward) - - def forward(*args, **kwargs): - return model_forward(*args, **kwargs) - - # To act like a decorator so that it can be popped when doing `extract_model_from_parallel` - forward.__wrapped__ = model_forward - - return forward - - -def find_device(data): - """ - Finds the device on which a nested dict/list/tuple of tensors lies (assuming they are all on the same device). - - Args: - (nested list/tuple/dictionary of `torch.Tensor`): The data we want to know the device of. - """ - if isinstance(data, Mapping): - for obj in data.values(): - device = find_device(obj) - if device is not None: - return device - elif isinstance(data, (tuple, list)): - for obj in data: - device = find_device(obj) - if device is not None: - return device - elif isinstance(data, torch.Tensor): - return data.device - - -@contextmanager -def GatheredParameters(params, modifier_rank=None, fwd_module=None, enabled=True): - """ - Wrapper around `deepspeed.runtime.zero.GatheredParameters`, but if Zero-3 is not enabled, will be a no-op context - manager. - """ - # We need to use the `AcceleratorState` here since it has access to the deepspeed plugin - if AcceleratorState().distributed_type != DistributedType.DEEPSPEED or ( - AcceleratorState().deepspeed_plugin is not None - and not AcceleratorState().deepspeed_plugin.is_zero3_init_enabled() - ): - gather_param_context = nullcontext() - else: - import deepspeed - - gather_param_context = deepspeed.zero.GatheredParameters( - params, modifier_rank=modifier_rank, fwd_module=fwd_module, enabled=enabled - ) - with gather_param_context: - yield diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/other.py b/venv/lib/python3.10/site-packages/accelerate/utils/other.py deleted file mode 100644 index adc0ed9cbfe66fa6e5fbcfd4da623d279c837a61..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/other.py +++ /dev/null @@ -1,560 +0,0 @@ -# Copyright 2022 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 collections -import platform -import re -import socket -from codecs import encode -from collections import OrderedDict -from functools import partial, reduce -from types import MethodType - -import numpy as np -import torch -from packaging.version import Version -from safetensors.torch import save_file as safe_save_file - -from ..commands.config.default import write_basic_config # noqa: F401 -from ..logging import get_logger -from ..state import PartialState -from .constants import FSDP_PYTORCH_VERSION -from .dataclasses import DistributedType -from .imports import ( - is_deepspeed_available, - is_numpy_available, - is_torch_distributed_available, - is_torch_xla_available, - is_weights_only_available, -) -from .modeling import id_tensor_storage -from .transformer_engine import convert_model -from .versions import is_torch_version - - -logger = get_logger(__name__) - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - - -def is_compiled_module(module: torch.nn.Module) -> bool: - """ - Check whether the module was compiled with torch.compile() - """ - if not hasattr(torch, "_dynamo"): - return False - - return isinstance(module, torch._dynamo.eval_frame.OptimizedModule) - - -def has_compiled_regions(module: torch.nn.Module) -> bool: - """ - Check whether the module has submodules that were compiled with `torch.compile()`. - """ - if not hasattr(torch, "_dynamo"): - return False - - if module._modules: - for submodule in module.modules(): - if isinstance(submodule, torch._dynamo.eval_frame.OptimizedModule): - return True - - return False - - -def is_repeated_blocks(module: torch.nn.Module) -> bool: - """ - Check whether the module is a repeated block, i.e. `torch.nn.ModuleList` with all children of the same class. This - is useful to determine whether we should apply regional compilation to the module. - """ - - return isinstance(module, torch.nn.ModuleList) and all(isinstance(m, module[0].__class__) for m in module) - - -def has_repeated_blocks(module: torch.nn.Module) -> bool: - """ - Check whether the module has repeated blocks, i.e. `torch.nn.ModuleList` with all children of the same class, at - any level of the module hierarchy. This is useful to determine whether we should apply regional compilation to the - module. - """ - if module._modules: - for submodule in module.modules(): - if is_repeated_blocks(submodule): - return True - - return False - - -def compile_regions(module: torch.nn.Module, **compile_kwargs) -> torch.nn.Module: - """ - Performs regional compilation where we target repeated blocks of the same class and compile them sequentially to - hit the compiler's cache. For example, in `GPT2LMHeadModel`, the repeated block/class is `GPT2Block`, and can be - accessed as `model.transformer.h[0]`. The rest of the model (e.g. model.lm_head) is compiled separately. - - This allows us to speed up the compilation overhead / cold start of models like LLMs and Transformers in general. - See https://pytorch.org/tutorials/recipes/regional_compilation.html for more details. - - Args: - module (`torch.nn.Module`): - The model to compile. - **compile_kwargs: - Additional keyword arguments to pass to `torch.compile()`. - - Returns: - `torch.nn.Module`: A new instance of the model with some compiled regions. - - Example: - ```python - >>> from accelerate.utils import compile_regions - >>> from transformers import AutoModelForCausalLM - - >>> model = AutoModelForCausalLM.from_pretrained("gpt2") - >>> compiled_model = compile_regions(model, mode="reduce-overhead") - >>> compiled_model.transformer.h[0] - OptimizedModule( - (_orig_mod): GPT2Block( - (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True) - (attn): GPT2Attention( - (c_attn): Conv1D(nf=2304, nx=768) - (c_proj): Conv1D(nf=768, nx=768) - (attn_dropout): Dropout(p=0.1, inplace=False) - (resid_dropout): Dropout(p=0.1, inplace=False) - ) - (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True) - (mlp): GPT2MLP( - (c_fc): Conv1D(nf=3072, nx=768) - (c_proj): Conv1D(nf=768, nx=3072) - (act): NewGELUActivation() - (dropout): Dropout(p=0.1, inplace=False) - ) - ) - ) - ``` - """ - - def _compile_regions(module: torch.nn.Module, **compile_kwargs) -> torch.nn.Module: - if is_repeated_blocks(module): - new_module = torch.nn.ModuleList() - for submodule in module: - new_module.append(torch.compile(submodule, **compile_kwargs)) - elif has_repeated_blocks(module): - new_module = module.__class__.__new__(module.__class__) - new_module.__dict__.update(module.__dict__) - new_module._modules = {} - for name, submodule in module.named_children(): - new_module.add_module(name, _compile_regions(submodule, **compile_kwargs)) - else: - new_module = torch.compile(module, **compile_kwargs) - - return new_module - - new_module = _compile_regions(module, **compile_kwargs) - - if "_orig_mod" not in new_module.__dict__: - # Keeps a reference to the original module to decompile/unwrap it later - new_module.__dict__["_orig_mod"] = module - - return new_module - - -def compile_regions_deepspeed(module: torch.nn.Module, **compile_kwargs): - """ - Performs regional compilation the same way as `compile_regions`, but specifically for `DeepSpeedEngine.module`. - Since the model is wrapped in a `DeepSpeedEngine` and has many added hooks, offloaded parameters, etc that - `torch.compile(...)` interferes with, version of trgional compilation uses the inplace `module.compile()` method - instead. - - Args: - module (`torch.nn.Module`): - The model to compile. - **compile_kwargs: - Additional keyword arguments to pass to `module.compile()`. - """ - - if is_repeated_blocks(module): - for submodule in module: - submodule.compile(**compile_kwargs) - elif has_repeated_blocks(module): - for child in module.children(): - compile_regions_deepspeed(child, **compile_kwargs) - else: # leaf node - module.compile(**compile_kwargs) - - -def model_has_dtensor(model: torch.nn.Module) -> bool: - """ - Check if the model has DTensor parameters. - - Args: - model (`torch.nn.Module`): - The model to check. - - Returns: - `bool`: Whether the model has DTensor parameters. - """ - if is_torch_version(">=", "2.5.0"): - from torch.distributed.tensor import DTensor - else: - # from torch 2.0.0 (oldest supported accelerate torch version), DTensor is in torch.distributed._tensor - from torch.distributed._tensor import DTensor - - return any(isinstance(p, DTensor) for p in model.parameters()) - - -def extract_model_from_parallel( - model, keep_fp32_wrapper: bool = True, keep_torch_compile: bool = True, recursive: bool = False -): - """ - Extract a model from its distributed containers. - - Args: - model (`torch.nn.Module`): - The model to extract. - keep_fp32_wrapper (`bool`, *optional*): - Whether to remove mixed precision hooks from the model. - keep_torch_compile (`bool`, *optional*): - Whether to unwrap compiled model. - recursive (`bool`, *optional*, defaults to `False`): - Whether to recursively extract all cases of `module.module` from `model` as well as unwrap child sublayers - recursively, not just the top-level distributed containers. - - Returns: - `torch.nn.Module`: The extracted model. - """ - options = (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel) - - is_compiled = is_compiled_module(model) - has_compiled = has_compiled_regions(model) - - if is_compiled: - compiled_model = model - model = model._orig_mod - elif has_compiled: - compiled_model = model - model = model.__dict__["_orig_mod"] - - if is_deepspeed_available(): - from deepspeed import DeepSpeedEngine - - options += (DeepSpeedEngine,) - - if is_torch_version(">=", FSDP_PYTORCH_VERSION) and is_torch_distributed_available(): - from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP - - options += (FSDP,) - - while isinstance(model, options): - model = model.module - - if recursive: - # This is needed in cases such as using FSDPv2 on XLA - def _recursive_unwrap(module): - # Wrapped modules are standardly wrapped as `module`, similar to the cases earlier - # with DDP, DataParallel, DeepSpeed, and FSDP - if hasattr(module, "module"): - unwrapped_module = _recursive_unwrap(module.module) - else: - unwrapped_module = module - # Next unwrap child sublayers recursively - for name, child in unwrapped_module.named_children(): - setattr(unwrapped_module, name, _recursive_unwrap(child)) - return unwrapped_module - - # Start with top-level - model = _recursive_unwrap(model) - - if not keep_fp32_wrapper: - forward = model.forward - original_forward = model.__dict__.pop("_original_forward", None) - if original_forward is not None: - while hasattr(forward, "__wrapped__"): - forward = forward.__wrapped__ - if forward == original_forward: - break - model.forward = MethodType(forward, model) - if getattr(model, "_converted_to_transformer_engine", False): - convert_model(model, to_transformer_engine=False) - - if keep_torch_compile: - if is_compiled: - compiled_model._orig_mod = model - model = compiled_model - elif has_compiled: - compiled_model.__dict__["_orig_mod"] = model - model = compiled_model - - return model - - -def wait_for_everyone(): - """ - Introduces a blocking point in the script, making sure all processes have reached this point before continuing. - - - - Make sure all processes will reach this instruction otherwise one of your processes will hang forever. - - - """ - PartialState().wait_for_everyone() - - -def clean_state_dict_for_safetensors(state_dict: dict): - """ - Cleans the state dictionary from a model and removes tensor aliasing if present. - - Args: - state_dict (`dict`): - The state dictionary from a model - """ - ptrs = collections.defaultdict(list) - # When bnb serialization is used, weights in state dict can be strings - for name, tensor in state_dict.items(): - if not isinstance(tensor, str): - ptrs[id_tensor_storage(tensor)].append(name) - - # These are all pointers of tensors with shared memory - shared_ptrs = {ptr: names for ptr, names in ptrs.items() if len(names) > 1} - warn_names = set() - for names in shared_ptrs.values(): - # When not all duplicates have been cleaned, we still remove those keys but put a clear warning. - # If the link between tensors was done at runtime then `from_pretrained` will not get - # the key back leading to random tensor. A proper warning will be shown - # during reload (if applicable), but since the file is not necessarily compatible with - # the config, better show a proper warning. - found_names = [name for name in names if name in state_dict] - warn_names.update(found_names[1:]) - for name in found_names[1:]: - del state_dict[name] - if len(warn_names) > 0: - logger.warning( - f"Removed shared tensor {warn_names} while saving. This should be OK, but check by verifying that you don't receive any warning while reloading", - ) - state_dict = {k: v.contiguous() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()} - return state_dict - - -def save(obj, f, save_on_each_node: bool = False, safe_serialization: bool = False): - """ - Save the data to disk. Use in place of `torch.save()`. - - Args: - obj: - The data to save - f: - The file (or file-like object) to use to save the data - save_on_each_node (`bool`, *optional*, defaults to `False`): - Whether to only save on the global main process - safe_serialization (`bool`, *optional*, defaults to `False`): - Whether to save `obj` using `safetensors` or the traditional PyTorch way (that uses `pickle`). - """ - # When TorchXLA is enabled, it's necessary to transfer all data to the CPU before saving. - # Another issue arises with `id_tensor_storage`, which treats all XLA tensors as identical. - # If tensors remain on XLA, calling `clean_state_dict_for_safetensors` will result in only - # one XLA tensor remaining. - if PartialState().distributed_type == DistributedType.XLA: - obj = xm._maybe_convert_to_cpu(obj) - # Check if it's a model and remove duplicates - if safe_serialization: - save_func = partial(safe_save_file, metadata={"format": "pt"}) - if isinstance(obj, OrderedDict): - obj = clean_state_dict_for_safetensors(obj) - else: - save_func = torch.save - - if PartialState().is_main_process and not save_on_each_node: - save_func(obj, f) - elif PartialState().is_local_main_process and save_on_each_node: - save_func(obj, f) - - -# The following are considered "safe" globals to reconstruct various types of objects when using `weights_only=True` -# These should be added and then removed after loading in the file -np_core = np._core if is_numpy_available("2.0.0") else np.core -TORCH_SAFE_GLOBALS = [ - # numpy arrays are just numbers, not objects, so we can reconstruct them safely - np_core.multiarray._reconstruct, - np.ndarray, - # The following are needed for the RNG states - encode, - np.dtype, -] - -if is_numpy_available("1.25.0"): - TORCH_SAFE_GLOBALS.append(np.dtypes.UInt32DType) - - -def load(f, map_location=None, **kwargs): - """ - Compatible drop-in replacement of `torch.load()` which allows for `weights_only` to be used if `torch` version is - 2.4.0 or higher. Otherwise will ignore the kwarg. - - Will also add (and then remove) an exception for numpy arrays - - Args: - f: - The file (or file-like object) to use to load the data - map_location: - a function, `torch.device`, string or a dict specifying how to remap storage locations - **kwargs: - Additional keyword arguments to pass to `torch.load()`. - """ - try: - if is_weights_only_available(): - old_safe_globals = torch.serialization.get_safe_globals() - if "weights_only" not in kwargs: - kwargs["weights_only"] = True - torch.serialization.add_safe_globals(TORCH_SAFE_GLOBALS) - else: - kwargs.pop("weights_only", None) - loaded_obj = torch.load(f, map_location=map_location, **kwargs) - finally: - if is_weights_only_available(): - torch.serialization.clear_safe_globals() - if old_safe_globals: - torch.serialization.add_safe_globals(old_safe_globals) - return loaded_obj - - -def get_pretty_name(obj): - """ - Gets a pretty name from `obj`. - """ - if not hasattr(obj, "__qualname__") and not hasattr(obj, "__name__"): - obj = getattr(obj, "__class__", obj) - if hasattr(obj, "__qualname__"): - return obj.__qualname__ - if hasattr(obj, "__name__"): - return obj.__name__ - return str(obj) - - -def merge_dicts(source, destination): - """ - Recursively merges two dictionaries. - - Args: - source (`dict`): The dictionary to merge into `destination`. - destination (`dict`): The dictionary to merge `source` into. - """ - for key, value in source.items(): - if isinstance(value, dict): - node = destination.setdefault(key, {}) - merge_dicts(value, node) - else: - destination[key] = value - - return destination - - -def is_port_in_use(port: int = None) -> bool: - """ - Checks if a port is in use on `localhost`. Useful for checking if multiple `accelerate launch` commands have been - run and need to see if the port is already in use. - """ - if port is None: - port = 29500 - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - return s.connect_ex(("localhost", port)) == 0 - - -def get_free_port() -> int: - """ - Gets a free port on `localhost`. Useful for automatic port selection when port 0 is specified in distributed - training scenarios. - - Returns: - int: An available port number - """ - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("", 0)) # bind to port 0 for OS to assign a free port - return s.getsockname()[1] - - -def convert_bytes(size): - "Converts `size` from bytes to the largest possible unit" - for x in ["bytes", "KB", "MB", "GB", "TB"]: - if size < 1024.0: - return f"{round(size, 2)} {x}" - size /= 1024.0 - - return f"{round(size, 2)} PB" - - -def check_os_kernel(): - """Warns if the kernel version is below the recommended minimum on Linux.""" - # see issue #1929 - info = platform.uname() - system = info.system - if system != "Linux": - return - - _, version, *_ = re.split(r"(\d+\.\d+\.\d+)", info.release) - min_version = "5.5.0" - if Version(version) < Version(min_version): - msg = ( - f"Detected kernel version {version}, which is below the recommended minimum of {min_version}; this can " - "cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher." - ) - logger.warning(msg, main_process_only=True) - - -def recursive_getattr(obj, attr: str): - """ - Recursive `getattr`. - - Args: - obj: - A class instance holding the attribute. - attr (`str`): - The attribute that is to be retrieved, e.g. 'attribute1.attribute2'. - """ - - def _getattr(obj, attr): - return getattr(obj, attr) - - return reduce(_getattr, [obj] + attr.split(".")) - - -def get_module_children_bottom_up(model: torch.nn.Module, return_fqns: bool = False) -> list[torch.nn.Module]: - """Traverse the model in bottom-up order and return the children modules in that order. - - Args: - model (`torch.nn.Module`): the model to get the children of - - Returns: - `list[torch.nn.Module]`: a list of children modules of `model` in bottom-up order. The last element is the - `model` itself. - """ - top = model if not return_fqns else ("", model) - stack = [top] - ordered_modules = [] - while stack: - current_module = stack.pop() - if return_fqns: - current_module_name, current_module = current_module - for name, attr in current_module.named_children(): - if isinstance(attr, torch.nn.Module): - if return_fqns: - child_name = current_module_name + "." + name if current_module_name else name - stack.append((child_name, attr)) - else: - stack.append(attr) - if return_fqns: - ordered_modules.append((current_module_name, current_module)) - else: - ordered_modules.append(current_module) - return ordered_modules[::-1] diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/random.py b/venv/lib/python3.10/site-packages/accelerate/utils/random.py deleted file mode 100644 index 9dceb598cacc1c1d17b198e0a7c19789ab3b9f39..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/random.py +++ /dev/null @@ -1,156 +0,0 @@ -# Copyright 2022 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 random -from typing import Optional, Union - -import numpy as np -import torch - -from ..state import AcceleratorState -from .constants import CUDA_DISTRIBUTED_TYPES -from .dataclasses import DistributedType, RNGType -from .imports import ( - is_hpu_available, - is_mlu_available, - is_musa_available, - is_npu_available, - is_sdaa_available, - is_torch_xla_available, - is_xpu_available, -) - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - - -def set_seed(seed: int, device_specific: bool = False, deterministic: bool = False): - """ - Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`. - - Args: - seed (`int`): - The seed to set. - device_specific (`bool`, *optional*, defaults to `False`): - Whether to differ the seed on each device slightly with `self.process_index`. - deterministic (`bool`, *optional*, defaults to `False`): - Whether to use deterministic algorithms where available. Can slow down training. - """ - if device_specific: - seed += AcceleratorState().process_index - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if is_xpu_available(): - torch.xpu.manual_seed_all(seed) - elif is_npu_available(): - torch.npu.manual_seed_all(seed) - elif is_mlu_available(): - torch.mlu.manual_seed_all(seed) - elif is_sdaa_available(): - torch.sdaa.manual_seed_all(seed) - elif is_musa_available(): - torch.musa.manual_seed_all(seed) - elif is_hpu_available(): - torch.hpu.manual_seed_all(seed) - else: - torch.cuda.manual_seed_all(seed) - # ^^ safe to call this function even if cuda is not available - if is_torch_xla_available(): - xm.set_rng_state(seed) - - if deterministic: - torch.use_deterministic_algorithms(True) - - -def synchronize_rng_state(rng_type: Optional[RNGType] = None, generator: Optional[torch.Generator] = None): - # Get the proper rng state - if rng_type == RNGType.TORCH: - rng_state = torch.get_rng_state() - elif rng_type == RNGType.CUDA: - rng_state = torch.cuda.get_rng_state() - elif rng_type == RNGType.XLA: - assert is_torch_xla_available(), "Can't synchronize XLA seeds as torch_xla is unavailable." - rng_state = torch.tensor(xm.get_rng_state()) - elif rng_type == RNGType.NPU: - assert is_npu_available(), "Can't synchronize NPU seeds on an environment without NPUs." - rng_state = torch.npu.get_rng_state() - elif rng_type == RNGType.MLU: - assert is_mlu_available(), "Can't synchronize MLU seeds on an environment without MLUs." - rng_state = torch.mlu.get_rng_state() - elif rng_type == RNGType.SDAA: - assert is_sdaa_available(), "Can't synchronize SDAA seeds on an environment without SDAAs." - rng_state = torch.sdaa.get_rng_state() - elif rng_type == RNGType.MUSA: - assert is_musa_available(), "Can't synchronize MUSA seeds on an environment without MUSAs." - rng_state = torch.musa.get_rng_state() - elif rng_type == RNGType.XPU: - assert is_xpu_available(), "Can't synchronize XPU seeds on an environment without XPUs." - rng_state = torch.xpu.get_rng_state() - elif rng_type == RNGType.HPU: - assert is_hpu_available(), "Can't synchronize HPU seeds on an environment without HPUs." - rng_state = torch.hpu.get_rng_state() - elif rng_type == RNGType.GENERATOR: - assert generator is not None, "Need a generator to synchronize its seed." - rng_state = generator.get_state() - - # Broadcast the rng state from device 0 to other devices - state = AcceleratorState() - if state.distributed_type == DistributedType.XLA: - rng_state = rng_state.to(xm.xla_device()) - xm.collective_broadcast([rng_state]) - xm.mark_step() - rng_state = rng_state.cpu() - elif ( - state.distributed_type in CUDA_DISTRIBUTED_TYPES - or state.distributed_type == DistributedType.MULTI_MLU - or state.distributed_type == DistributedType.MULTI_SDAA - or state.distributed_type == DistributedType.MULTI_MUSA - or state.distributed_type == DistributedType.MULTI_NPU - or state.distributed_type == DistributedType.MULTI_XPU - or state.distributed_type == DistributedType.MULTI_HPU - ): - rng_state = rng_state.to(state.device) - torch.distributed.broadcast(rng_state, 0) - rng_state = rng_state.cpu() - elif state.distributed_type == DistributedType.MULTI_CPU: - torch.distributed.broadcast(rng_state, 0) - - # Set the broadcast rng state - if rng_type == RNGType.TORCH: - torch.set_rng_state(rng_state) - elif rng_type == RNGType.CUDA: - torch.cuda.set_rng_state(rng_state) - elif rng_type == RNGType.NPU: - torch.npu.set_rng_state(rng_state) - elif rng_type == RNGType.MLU: - torch.mlu.set_rng_state(rng_state) - elif rng_type == RNGType.SDAA: - torch.sdaa.set_rng_state(rng_state) - elif rng_type == RNGType.MUSA: - torch.musa.set_rng_state(rng_state) - elif rng_type == RNGType.XPU: - torch.xpu.set_rng_state(rng_state) - elif rng_state == RNGType.HPU: - torch.hpu.set_rng_state(rng_state) - elif rng_type == RNGType.XLA: - xm.set_rng_state(rng_state.item()) - elif rng_type == RNGType.GENERATOR: - generator.set_state(rng_state) - - -def synchronize_rng_states(rng_types: list[Union[str, RNGType]], generator: Optional[torch.Generator] = None): - for rng_type in rng_types: - synchronize_rng_state(RNGType(rng_type), generator=generator) diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/rich.py b/venv/lib/python3.10/site-packages/accelerate/utils/rich.py deleted file mode 100644 index 2d48661b7fcef92ef1168b74cc275c6d3ccc67a1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/rich.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright 2022 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 .imports import is_rich_available - - -if is_rich_available(): - from rich.traceback import install - - install(show_locals=False) - -else: - raise ModuleNotFoundError("To use the rich extension, install rich with `pip install rich`") diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/torch_xla.py b/venv/lib/python3.10/site-packages/accelerate/utils/torch_xla.py deleted file mode 100644 index 140133926c2f88d39c70f5a9f46a08f88bed36da..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/torch_xla.py +++ /dev/null @@ -1,51 +0,0 @@ -# Copyright 2022 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.metadata -import subprocess -import sys - - -def install_xla(upgrade: bool = False): - """ - Helper function to install appropriate xla wheels based on the `torch` version in Google Colaboratory. - - Args: - upgrade (`bool`, *optional*, defaults to `False`): - Whether to upgrade `torch` and install the latest `torch_xla` wheels. - - Example: - - ```python - >>> from accelerate.utils import install_xla - - >>> install_xla(upgrade=True) - ``` - """ - in_colab = False - if "IPython" in sys.modules: - in_colab = "google.colab" in str(sys.modules["IPython"].get_ipython()) - - if in_colab: - if upgrade: - torch_install_cmd = ["pip", "install", "-U", "torch"] - subprocess.run(torch_install_cmd, check=True) - # get the current version of torch - torch_version = importlib.metadata.version("torch") - torch_version_trunc = torch_version[: torch_version.rindex(".")] - xla_wheel = f"https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-{torch_version_trunc}-cp37-cp37m-linux_x86_64.whl" - xla_install_cmd = ["pip", "install", xla_wheel] - subprocess.run(xla_install_cmd, check=True) - else: - raise RuntimeError("`install_xla` utility works only on google colab.") diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/tqdm.py b/venv/lib/python3.10/site-packages/accelerate/utils/tqdm.py deleted file mode 100644 index 2d4873c1573eb2ee7392162f440a76d4f07cd8ce..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/tqdm.py +++ /dev/null @@ -1,43 +0,0 @@ -# Copyright 2022 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 .imports import is_tqdm_available - - -if is_tqdm_available(): - from tqdm.auto import tqdm as _tqdm - -from ..state import PartialState - - -def tqdm(*args, main_process_only: bool = True, **kwargs): - """ - Wrapper around `tqdm.tqdm` that optionally displays only on the main process. - - Args: - main_process_only (`bool`, *optional*): - Whether to display the progress bar only on the main process - """ - if not is_tqdm_available(): - raise ImportError("Accelerate's `tqdm` module requires `tqdm` to be installed. Please run `pip install tqdm`.") - if len(args) > 0 and isinstance(args[0], bool): - raise ValueError( - "Passing `True` or `False` as the first argument to Accelerate's `tqdm` wrapper is unsupported. " - "Please use the `main_process_only` keyword argument instead." - ) - disable = kwargs.pop("disable", False) - if main_process_only and not disable: - disable = PartialState().local_process_index != 0 - return _tqdm(*args, **kwargs, disable=disable) diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/transformer_engine.py b/venv/lib/python3.10/site-packages/accelerate/utils/transformer_engine.py deleted file mode 100644 index 337213453b3d15f610f41ac71c5afe15e034b38c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/transformer_engine.py +++ /dev/null @@ -1,164 +0,0 @@ -# Copyright 2022 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 types import MethodType - -import torch.nn as nn - -from .imports import is_hpu_available, is_transformer_engine_available -from .operations import GatheredParameters - - -# Do not import `transformer_engine` at package level to avoid potential issues - - -def convert_model(model, to_transformer_engine=True, _convert_linear=True, _convert_ln=True): - """ - Recursively converts the linear and layernorm layers of a model to their `transformers_engine` counterpart. - """ - if not is_transformer_engine_available(): - raise ImportError("Using `convert_model` requires transformer_engine to be installed.") - - if is_hpu_available(): - import intel_transformer_engine as te - - if not hasattr(te, "LayerNorm"): - # HPU does not have a LayerNorm implementation in TE - te.LayerNorm = nn.LayerNorm - else: - import transformer_engine.pytorch as te - - for name, module in model.named_children(): - if isinstance(module, nn.Linear) and to_transformer_engine and _convert_linear: - has_bias = module.bias is not None - params_to_gather = [module.weight] - if has_bias: - params_to_gather.append(module.bias) - - with GatheredParameters(params_to_gather, modifier_rank=0): - if any(p % 16 != 0 for p in module.weight.shape): - return - te_module = te.Linear( - module.in_features, module.out_features, bias=has_bias, params_dtype=module.weight.dtype - ) - te_module.weight.copy_(module.weight) - if has_bias: - te_module.bias.copy_(module.bias) - - setattr(model, name, te_module) - # Note: @xrsrke (Phuc) found that te.LayerNorm doesn't have any real memory savings or speedups over nn.LayerNorm - elif isinstance(module, nn.LayerNorm) and to_transformer_engine and _convert_ln: - with GatheredParameters([module.weight, module.bias], modifier_rank=0): - te_module = te.LayerNorm(module.normalized_shape[0], eps=module.eps, params_dtype=module.weight.dtype) - te_module.weight.copy_(module.weight) - te_module.bias.copy_(module.bias) - - setattr(model, name, te_module) - elif isinstance(module, te.Linear) and not to_transformer_engine and _convert_linear: - has_bias = module.bias is not None - new_module = nn.Linear( - module.in_features, module.out_features, bias=has_bias, params_dtype=module.weight.dtype - ) - new_module.weight.copy_(module.weight) - if has_bias: - new_module.bias.copy_(module.bias) - - setattr(model, name, new_module) - elif isinstance(module, te.LayerNorm) and not to_transformer_engine and _convert_ln: - new_module = nn.LayerNorm(module.normalized_shape[0], eps=module.eps, params_dtype=module.weight.dtype) - new_module.weight.copy_(module.weight) - new_module.bias.copy_(module.bias) - - setattr(model, name, new_module) - else: - convert_model( - module, - to_transformer_engine=to_transformer_engine, - _convert_linear=_convert_linear, - _convert_ln=_convert_ln, - ) - - -def has_transformer_engine_layers(model): - """ - Returns whether a given model has some `transformer_engine` layer or not. - """ - if not is_transformer_engine_available(): - raise ImportError("Using `has_transformer_engine_layers` requires transformer_engine to be installed.") - - if is_hpu_available(): - import intel_transformer_engine as te - - module_cls_to_check = te.Linear - else: - import transformer_engine.pytorch as te - - module_cls_to_check = (te.LayerNorm, te.Linear, te.TransformerLayer) - - for m in model.modules(): - if isinstance(m, module_cls_to_check): - return True - - return False - - -def contextual_fp8_autocast(model_forward, fp8_recipe, use_during_eval=False): - """ - Wrapper for a model's forward method to apply FP8 autocast. Is context aware, meaning that by default it will - disable FP8 autocast during eval mode, which is generally better for more accurate metrics. - """ - if not is_transformer_engine_available(): - raise ImportError("Using `contextual_fp8_autocast` requires transformer_engine to be installed.") - - if is_hpu_available(): - from intel_transformer_engine import fp8_autocast - else: - from transformer_engine.pytorch import fp8_autocast - - def forward(self, *args, **kwargs): - enabled = use_during_eval or self.training - with fp8_autocast(enabled=enabled, fp8_recipe=fp8_recipe): - return model_forward(*args, **kwargs) - - # To act like a decorator so that it can be popped when doing `extract_model_from_parallel` - forward.__wrapped__ = model_forward - - return forward - - -def apply_fp8_autowrap(model, fp8_recipe_handler): - """ - Applies FP8 context manager to the model's forward method - """ - if not is_transformer_engine_available(): - raise ImportError("Using `apply_fp8_autowrap` requires transformer_engine to be installed.") - - if is_hpu_available(): - import intel_transformer_engine.recipe as te_recipe - else: - import transformer_engine.common.recipe as te_recipe - - kwargs = fp8_recipe_handler.to_kwargs() if fp8_recipe_handler is not None else {} - if "fp8_format" in kwargs: - kwargs["fp8_format"] = getattr(te_recipe.Format, kwargs["fp8_format"]) - use_during_eval = kwargs.pop("use_autocast_during_eval", False) - fp8_recipe = te_recipe.DelayedScaling(**kwargs) - new_forward = contextual_fp8_autocast(model.forward, fp8_recipe, use_during_eval) - - if hasattr(model.forward, "__func__"): - model.forward = MethodType(new_forward, model) - else: - model.forward = new_forward - - return model diff --git a/venv/lib/python3.10/site-packages/accelerate/utils/versions.py b/venv/lib/python3.10/site-packages/accelerate/utils/versions.py deleted file mode 100644 index 985c918f0e057bacc70c372f6906071bb73db577..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/accelerate/utils/versions.py +++ /dev/null @@ -1,56 +0,0 @@ -# Copyright 2022 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.metadata -from typing import Union - -from packaging.version import Version, parse - -from .constants import STR_OPERATION_TO_FUNC - - -torch_version = parse(importlib.metadata.version("torch")) - - -def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str): - """ - Compares a library version to some requirement using a given operation. - - Args: - library_or_version (`str` or `packaging.version.Version`): - A library name or a version to check. - operation (`str`): - A string representation of an operator, such as `">"` or `"<="`. - requirement_version (`str`): - The version to compare the library version against - """ - if operation not in STR_OPERATION_TO_FUNC.keys(): - raise ValueError(f"`operation` must be one of {list(STR_OPERATION_TO_FUNC.keys())}, received {operation}") - operation = STR_OPERATION_TO_FUNC[operation] - if isinstance(library_or_version, str): - library_or_version = parse(importlib.metadata.version(library_or_version)) - return operation(library_or_version, parse(requirement_version)) - - -def is_torch_version(operation: str, version: str): - """ - Compares the current PyTorch version to a given reference with an operation. - - Args: - operation (`str`): - A string representation of an operator, such as `">"` or `"<="` - version (`str`): - A string version of PyTorch - """ - return compare_versions(torch_version, operation, version) diff --git a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/LICENSE b/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/LICENSE deleted file mode 100644 index 4362b49151d7b34ef83b3067a8f9c9f877d72a0e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/LICENSE +++ /dev/null @@ -1,502 +0,0 @@ - GNU LESSER GENERAL PUBLIC LICENSE - Version 2.1, February 1999 - - Copyright (C) 1991, 1999 Free Software Foundation, Inc. - 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA - Everyone is permitted to copy and distribute verbatim copies - of this license document, but changing it is not allowed. - -[This is the first released version of the Lesser GPL. It also counts - as the successor of the GNU Library Public License, version 2, hence - the version number 2.1.] - - Preamble - - The licenses for most software are designed to take away your -freedom to share and change it. By contrast, the GNU General Public -Licenses are intended to guarantee your freedom to share and change -free software--to make sure the software is free for all its users. - - This license, the Lesser General Public License, applies to some -specially designated software packages--typically libraries--of the -Free Software Foundation and other authors who decide to use it. You -can use it too, but we suggest you first think carefully about whether -this license or the ordinary General Public License is the better -strategy to use in any particular case, based on the explanations below. - - When we speak of free software, we are referring to freedom of use, -not price. Our General Public Licenses are designed to make sure that -you have the freedom to distribute copies of free software (and charge -for this service if you wish); that you receive source code or can get -it if you want it; that you can change the software and use pieces of -it in new free programs; and that you are informed that you can do -these things. - - To protect your rights, we need to make restrictions that forbid -distributors to deny you these rights or to ask you to surrender these -rights. These restrictions translate to certain responsibilities for -you if you distribute copies of the library or if you modify it. - - For example, if you distribute copies of the library, whether gratis -or for a fee, you must give the recipients all the rights that we gave -you. You must make sure that they, too, receive or can get the source -code. If you link other code with the library, you must provide -complete object files to the recipients, so that they can relink them -with the library after making changes to the library and recompiling -it. And you must show them these terms so they know their rights. - - We protect your rights with a two-step method: (1) we copyright the -library, and (2) we offer you this license, which gives you legal -permission to copy, distribute and/or modify the library. - - To protect each distributor, we want to make it very clear that -there is no warranty for the free library. Also, if the library is -modified by someone else and passed on, the recipients should know -that what they have is not the original version, so that the original -author's reputation will not be affected by problems that might be -introduced by others. - - Finally, software patents pose a constant threat to the existence of -any free program. We wish to make sure that a company cannot -effectively restrict the users of a free program by obtaining a -restrictive license from a patent holder. Therefore, we insist that -any patent license obtained for a version of the library must be -consistent with the full freedom of use specified in this license. - - Most GNU software, including some libraries, is covered by the -ordinary GNU General Public License. This license, the GNU Lesser -General Public License, applies to certain designated libraries, and -is quite different from the ordinary General Public License. We use -this license for certain libraries in order to permit linking those -libraries into non-free programs. - - When a program is linked with a library, whether statically or using -a shared library, the combination of the two is legally speaking a -combined work, a derivative of the original library. The ordinary -General Public License therefore permits such linking only if the -entire combination fits its criteria of freedom. The Lesser General -Public License permits more lax criteria for linking other code with -the library. - - We call this license the "Lesser" General Public License because it -does Less to protect the user's freedom than the ordinary General -Public License. It also provides other free software developers Less -of an advantage over competing non-free programs. These disadvantages -are the reason we use the ordinary General Public License for many -libraries. However, the Lesser license provides advantages in certain -special circumstances. - - For example, on rare occasions, there may be a special need to -encourage the widest possible use of a certain library, so that it becomes -a de-facto standard. To achieve this, non-free programs must be -allowed to use the library. A more frequent case is that a free -library does the same job as widely used non-free libraries. In this -case, there is little to gain by limiting the free library to free -software only, so we use the Lesser General Public License. - - In other cases, permission to use a particular library in non-free -programs enables a greater number of people to use a large body of -free software. For example, permission to use the GNU C Library in -non-free programs enables many more people to use the whole GNU -operating system, as well as its variant, the GNU/Linux operating -system. - - Although the Lesser General Public License is Less protective of the -users' freedom, it does ensure that the user of a program that is -linked with the Library has the freedom and the wherewithal to run -that program using a modified version of the Library. - - The precise terms and conditions for copying, distribution and -modification follow. Pay close attention to the difference between a -"work based on the library" and a "work that uses the library". The -former contains code derived from the library, whereas the latter must -be combined with the library in order to run. - - GNU LESSER GENERAL PUBLIC LICENSE - TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION - - 0. This License Agreement applies to any software library or other -program which contains a notice placed by the copyright holder or -other authorized party saying it may be distributed under the terms of -this Lesser General Public License (also called "this License"). -Each licensee is addressed as "you". - - A "library" means a collection of software functions and/or data -prepared so as to be conveniently linked with application programs -(which use some of those functions and data) to form executables. - - The "Library", below, refers to any such software library or work -which has been distributed under these terms. A "work based on the -Library" means either the Library or any derivative work under -copyright law: that is to say, a work containing the Library or a -portion of it, either verbatim or with modifications and/or translated -straightforwardly into another language. (Hereinafter, translation is -included without limitation in the term "modification".) - - "Source code" for a work means the preferred form of the work for -making modifications to it. For a library, complete source code means -all the source code for all modules it contains, plus any associated -interface definition files, plus the scripts used to control compilation -and installation of the library. - - Activities other than copying, distribution and modification are not -covered by this License; they are outside its scope. The act of -running a program using the Library is not restricted, and output from -such a program is covered only if its contents constitute a work based -on the Library (independent of the use of the Library in a tool for -writing it). Whether that is true depends on what the Library does -and what the program that uses the Library does. - - 1. You may copy and distribute verbatim copies of the Library's -complete source code as you receive it, in any medium, provided that -you conspicuously and appropriately publish on each copy an -appropriate copyright notice and disclaimer of warranty; keep intact -all the notices that refer to this License and to the absence of any -warranty; and distribute a copy of this License along with the -Library. - - You may charge a fee for the physical act of transferring a copy, -and you may at your option offer warranty protection in exchange for a -fee. - - 2. You may modify your copy or copies of the Library or any portion -of it, thus forming a work based on the Library, and copy and -distribute such modifications or work under the terms of Section 1 -above, provided that you also meet all of these conditions: - - a) The modified work must itself be a software library. - - b) You must cause the files modified to carry prominent notices - stating that you changed the files and the date of any change. - - c) You must cause the whole of the work to be licensed at no - charge to all third parties under the terms of this License. - - d) If a facility in the modified Library refers to a function or a - table of data to be supplied by an application program that uses - the facility, other than as an argument passed when the facility - is invoked, then you must make a good faith effort to ensure that, - in the event an application does not supply such function or - table, the facility still operates, and performs whatever part of - its purpose remains meaningful. - - (For example, a function in a library to compute square roots has - a purpose that is entirely well-defined independent of the - application. Therefore, Subsection 2d requires that any - application-supplied function or table used by this function must - be optional: if the application does not supply it, the square - root function must still compute square roots.) - -These requirements apply to the modified work as a whole. If -identifiable sections of that work are not derived from the Library, -and can be reasonably considered independent and separate works in -themselves, then this License, and its terms, do not apply to those -sections when you distribute them as separate works. But when you -distribute the same sections as part of a whole which is a work based -on the Library, the distribution of the whole must be on the terms of -this License, whose permissions for other licensees extend to the -entire whole, and thus to each and every part regardless of who wrote -it. - -Thus, it is not the intent of this section to claim rights or contest -your rights to work written entirely by you; rather, the intent is to -exercise the right to control the distribution of derivative or -collective works based on the Library. - -In addition, mere aggregation of another work not based on the Library -with the Library (or with a work based on the Library) on a volume of -a storage or distribution medium does not bring the other work under -the scope of this License. - - 3. You may opt to apply the terms of the ordinary GNU General Public -License instead of this License to a given copy of the Library. To do -this, you must alter all the notices that refer to this License, so -that they refer to the ordinary GNU General Public License, version 2, -instead of to this License. (If a newer version than version 2 of the -ordinary GNU General Public License has appeared, then you can specify -that version instead if you wish.) Do not make any other change in -these notices. - - Once this change is made in a given copy, it is irreversible for -that copy, so the ordinary GNU General Public License applies to all -subsequent copies and derivative works made from that copy. - - This option is useful when you wish to copy part of the code of -the Library into a program that is not a library. - - 4. You may copy and distribute the Library (or a portion or -derivative of it, under Section 2) in object code or executable form -under the terms of Sections 1 and 2 above provided that you accompany -it with the complete corresponding machine-readable source code, which -must be distributed under the terms of Sections 1 and 2 above on a -medium customarily used for software interchange. - - If distribution of object code is made by offering access to copy -from a designated place, then offering equivalent access to copy the -source code from the same place satisfies the requirement to -distribute the source code, even though third parties are not -compelled to copy the source along with the object code. - - 5. A program that contains no derivative of any portion of the -Library, but is designed to work with the Library by being compiled or -linked with it, is called a "work that uses the Library". Such a -work, in isolation, is not a derivative work of the Library, and -therefore falls outside the scope of this License. - - However, linking a "work that uses the Library" with the Library -creates an executable that is a derivative of the Library (because it -contains portions of the Library), rather than a "work that uses the -library". The executable is therefore covered by this License. -Section 6 states terms for distribution of such executables. - - When a "work that uses the Library" uses material from a header file -that is part of the Library, the object code for the work may be a -derivative work of the Library even though the source code is not. -Whether this is true is especially significant if the work can be -linked without the Library, or if the work is itself a library. The -threshold for this to be true is not precisely defined by law. - - If such an object file uses only numerical parameters, data -structure layouts and accessors, and small macros and small inline -functions (ten lines or less in length), then the use of the object -file is unrestricted, regardless of whether it is legally a derivative -work. (Executables containing this object code plus portions of the -Library will still fall under Section 6.) - - Otherwise, if the work is a derivative of the Library, you may -distribute the object code for the work under the terms of Section 6. -Any executables containing that work also fall under Section 6, -whether or not they are linked directly with the Library itself. - - 6. As an exception to the Sections above, you may also combine or -link a "work that uses the Library" with the Library to produce a -work containing portions of the Library, and distribute that work -under terms of your choice, provided that the terms permit -modification of the work for the customer's own use and reverse -engineering for debugging such modifications. - - You must give prominent notice with each copy of the work that the -Library is used in it and that the Library and its use are covered by -this License. You must supply a copy of this License. If the work -during execution displays copyright notices, you must include the -copyright notice for the Library among them, as well as a reference -directing the user to the copy of this License. Also, you must do one -of these things: - - a) Accompany the work with the complete corresponding - machine-readable source code for the Library including whatever - changes were used in the work (which must be distributed under - Sections 1 and 2 above); and, if the work is an executable linked - with the Library, with the complete machine-readable "work that - uses the Library", as object code and/or source code, so that the - user can modify the Library and then relink to produce a modified - executable containing the modified Library. (It is understood - that the user who changes the contents of definitions files in the - Library will not necessarily be able to recompile the application - to use the modified definitions.) - - b) Use a suitable shared library mechanism for linking with the - Library. A suitable mechanism is one that (1) uses at run time a - copy of the library already present on the user's computer system, - rather than copying library functions into the executable, and (2) - will operate properly with a modified version of the library, if - the user installs one, as long as the modified version is - interface-compatible with the version that the work was made with. - - c) Accompany the work with a written offer, valid for at - least three years, to give the same user the materials - specified in Subsection 6a, above, for a charge no more - than the cost of performing this distribution. - - d) If distribution of the work is made by offering access to copy - from a designated place, offer equivalent access to copy the above - specified materials from the same place. - - e) Verify that the user has already received a copy of these - materials or that you have already sent this user a copy. - - For an executable, the required form of the "work that uses the -Library" must include any data and utility programs needed for -reproducing the executable from it. However, as a special exception, -the materials to be distributed need not include anything that is -normally distributed (in either source or binary form) with the major -components (compiler, kernel, and so on) of the operating system on -which the executable runs, unless that component itself accompanies -the executable. - - It may happen that this requirement contradicts the license -restrictions of other proprietary libraries that do not normally -accompany the operating system. Such a contradiction means you cannot -use both them and the Library together in an executable that you -distribute. - - 7. You may place library facilities that are a work based on the -Library side-by-side in a single library together with other library -facilities not covered by this License, and distribute such a combined -library, provided that the separate distribution of the work based on -the Library and of the other library facilities is otherwise -permitted, and provided that you do these two things: - - a) Accompany the combined library with a copy of the same work - based on the Library, uncombined with any other library - facilities. This must be distributed under the terms of the - Sections above. - - b) Give prominent notice with the combined library of the fact - that part of it is a work based on the Library, and explaining - where to find the accompanying uncombined form of the same work. - - 8. You may not copy, modify, sublicense, link with, or distribute -the Library except as expressly provided under this License. Any -attempt otherwise to copy, modify, sublicense, link with, or -distribute the Library is void, and will automatically terminate your -rights under this License. However, parties who have received copies, -or rights, from you under this License will not have their licenses -terminated so long as such parties remain in full compliance. - - 9. You are not required to accept this License, since you have not -signed it. However, nothing else grants you permission to modify or -distribute the Library or its derivative works. These actions are -prohibited by law if you do not accept this License. Therefore, by -modifying or distributing the Library (or any work based on the -Library), you indicate your acceptance of this License to do so, and -all its terms and conditions for copying, distributing or modifying -the Library or works based on it. - - 10. Each time you redistribute the Library (or any work based on the -Library), the recipient automatically receives a license from the -original licensor to copy, distribute, link with or modify the Library -subject to these terms and conditions. You may not impose any further -restrictions on the recipients' exercise of the rights granted herein. -You are not responsible for enforcing compliance by third parties with -this License. - - 11. If, as a consequence of a court judgment or allegation of patent -infringement or for any other reason (not limited to patent issues), -conditions are imposed on you (whether by court order, agreement or -otherwise) that contradict the conditions of this License, they do not -excuse you from the conditions of this License. If you cannot -distribute so as to satisfy simultaneously your obligations under this -License and any other pertinent obligations, then as a consequence you -may not distribute the Library at all. For example, if a patent -license would not permit royalty-free redistribution of the Library by -all those who receive copies directly or indirectly through you, then -the only way you could satisfy both it and this License would be to -refrain entirely from distribution of the Library. - -If any portion of this section is held invalid or unenforceable under any -particular circumstance, the balance of the section is intended to apply, -and the section as a whole is intended to apply in other circumstances. - -It is not the purpose of this section to induce you to infringe any -patents or other property right claims or to contest validity of any -such claims; this section has the sole purpose of protecting the -integrity of the free software distribution system which is -implemented by public license practices. Many people have made -generous contributions to the wide range of software distributed -through that system in reliance on consistent application of that -system; it is up to the author/donor to decide if he or she is willing -to distribute software through any other system and a licensee cannot -impose that choice. - -This section is intended to make thoroughly clear what is believed to -be a consequence of the rest of this License. - - 12. If the distribution and/or use of the Library is restricted in -certain countries either by patents or by copyrighted interfaces, the -original copyright holder who places the Library under this License may add -an explicit geographical distribution limitation excluding those countries, -so that distribution is permitted only in or among countries not thus -excluded. In such case, this License incorporates the limitation as if -written in the body of this License. - - 13. The Free Software Foundation may publish revised and/or new -versions of the Lesser General Public License from time to time. -Such new versions will be similar in spirit to the present version, -but may differ in detail to address new problems or concerns. - -Each version is given a distinguishing version number. If the Library -specifies a version number of this License which applies to it and -"any later version", you have the option of following the terms and -conditions either of that version or of any later version published by -the Free Software Foundation. If the Library does not specify a -license version number, you may choose any version ever published by -the Free Software Foundation. - - 14. If you wish to incorporate parts of the Library into other free -programs whose distribution conditions are incompatible with these, -write to the author to ask for permission. For software which is -copyrighted by the Free Software Foundation, write to the Free -Software Foundation; we sometimes make exceptions for this. Our -decision will be guided by the two goals of preserving the free status -of all derivatives of our free software and of promoting the sharing -and reuse of software generally. - - NO WARRANTY - - 15. BECAUSE THE LIBRARY IS LICENSED FREE OF CHARGE, THERE IS NO -WARRANTY FOR THE LIBRARY, TO THE EXTENT PERMITTED BY APPLICABLE LAW. -EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR -OTHER PARTIES PROVIDE THE LIBRARY "AS IS" WITHOUT WARRANTY OF ANY -KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE -IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR -PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE -LIBRARY IS WITH YOU. SHOULD THE LIBRARY PROVE DEFECTIVE, YOU ASSUME -THE COST OF ALL NECESSARY SERVICING, REPAIR OR CORRECTION. - - 16. IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN -WRITING WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MAY MODIFY -AND/OR REDISTRIBUTE THE LIBRARY AS PERMITTED ABOVE, BE LIABLE TO YOU -FOR DAMAGES, INCLUDING ANY GENERAL, SPECIAL, INCIDENTAL OR -CONSEQUENTIAL DAMAGES ARISING OUT OF THE USE OR INABILITY TO USE THE -LIBRARY (INCLUDING BUT NOT LIMITED TO LOSS OF DATA OR DATA BEING -RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD PARTIES OR A -FAILURE OF THE LIBRARY TO OPERATE WITH ANY OTHER SOFTWARE), EVEN IF -SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH -DAMAGES. - - END OF TERMS AND CONDITIONS - - How to Apply These Terms to Your New Libraries - - If you develop a new library, and you want it to be of the greatest -possible use to the public, we recommend making it free software that -everyone can redistribute and change. You can do so by permitting -redistribution under these terms (or, alternatively, under the terms of the -ordinary General Public License). - - To apply these terms, attach the following notices to the library. It is -safest to attach them to the start of each source file to most effectively -convey the exclusion of warranty; and each file should have at least the -"copyright" line and a pointer to where the full notice is found. - - - Copyright (C) - - This library is free software; you can redistribute it and/or - modify it under the terms of the GNU Lesser General Public - License as published by the Free Software Foundation; either - version 2.1 of the License, or (at your option) any later version. - - This library is distributed in the hope that it will be useful, - but WITHOUT ANY WARRANTY; without even the implied warranty of - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU - Lesser General Public License for more details. - - You should have received a copy of the GNU Lesser General Public - License along with this library; if not, write to the Free Software - Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA - -Also add information on how to contact you by electronic and paper mail. - -You should also get your employer (if you work as a programmer) or your -school, if any, to sign a "copyright disclaimer" for the library, if -necessary. Here is a sample; alter the names: - - Yoyodyne, Inc., hereby disclaims all copyright interest in the - library `Frob' (a library for tweaking knobs) written by James Random Hacker. - - , 1 April 1990 - Ty Coon, President of Vice - -That's all there is to it! diff --git a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/METADATA b/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/METADATA deleted file mode 100644 index 3ec4d797a2de96620acc6c38a67ea3e819bf5f46..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/METADATA +++ /dev/null @@ -1,97 +0,0 @@ -Metadata-Version: 2.1 -Name: chardet -Version: 5.2.0 -Summary: Universal encoding detector for Python 3 -Home-page: https://github.com/chardet/chardet -Author: Mark Pilgrim -Author-email: mark@diveintomark.org -Maintainer: Daniel Blanchard -Maintainer-email: dan.blanchard@gmail.com -License: LGPL -Project-URL: Documentation, https://chardet.readthedocs.io/ -Project-URL: GitHub Project, https://github.com/chardet/chardet -Project-URL: Issue Tracker, https://github.com/chardet/chardet/issues -Keywords: encoding,i18n,xml -Classifier: Development Status :: 5 - Production/Stable -Classifier: Intended Audience :: Developers -Classifier: License :: OSI Approved :: GNU Lesser General Public License v2 or later (LGPLv2+) -Classifier: Operating System :: OS Independent -Classifier: Programming Language :: Python -Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3.7 -Classifier: Programming Language :: Python :: 3.8 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Programming Language :: Python :: 3.10 -Classifier: Programming Language :: Python :: 3.11 -Classifier: Programming Language :: Python :: Implementation :: CPython -Classifier: Programming Language :: Python :: Implementation :: PyPy -Classifier: Topic :: Software Development :: Libraries :: Python Modules -Classifier: Topic :: Text Processing :: Linguistic -Requires-Python: >=3.7 -License-File: LICENSE - -Chardet: The Universal Character Encoding Detector --------------------------------------------------- - -.. image:: https://img.shields.io/travis/chardet/chardet/stable.svg - :alt: Build status - :target: https://travis-ci.org/chardet/chardet - -.. image:: https://img.shields.io/coveralls/chardet/chardet/stable.svg - :target: https://coveralls.io/r/chardet/chardet - -.. image:: https://img.shields.io/pypi/v/chardet.svg - :target: https://warehouse.python.org/project/chardet/ - :alt: Latest version on PyPI - -.. image:: https://img.shields.io/pypi/l/chardet.svg - :alt: License - - -Detects - - ASCII, UTF-8, UTF-16 (2 variants), UTF-32 (4 variants) - - Big5, GB2312, EUC-TW, HZ-GB-2312, ISO-2022-CN (Traditional and Simplified Chinese) - - EUC-JP, SHIFT_JIS, CP932, ISO-2022-JP (Japanese) - - EUC-KR, ISO-2022-KR, Johab (Korean) - - KOI8-R, MacCyrillic, IBM855, IBM866, ISO-8859-5, windows-1251 (Cyrillic) - - ISO-8859-5, windows-1251 (Bulgarian) - - ISO-8859-1, windows-1252, MacRoman (Western European languages) - - ISO-8859-7, windows-1253 (Greek) - - ISO-8859-8, windows-1255 (Visual and Logical Hebrew) - - TIS-620 (Thai) - -.. note:: - Our ISO-8859-2 and windows-1250 (Hungarian) probers have been temporarily - disabled until we can retrain the models. - -Requires Python 3.7+. - -Installation ------------- - -Install from `PyPI `_:: - - pip install chardet - -Documentation -------------- - -For users, docs are now available at https://chardet.readthedocs.io/. - -Command-line Tool ------------------ - -chardet comes with a command-line script which reports on the encodings of one -or more files:: - - % chardetect somefile someotherfile - somefile: windows-1252 with confidence 0.5 - someotherfile: ascii with confidence 1.0 - -About ------ - -This is a continuation of Mark Pilgrim's excellent original chardet port from C, and `Ian Cordasco `_'s -`charade `_ Python 3-compatible fork. - -:maintainer: Dan Blanchard diff --git a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/RECORD b/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/RECORD deleted file mode 100644 index 57627fd2e33dc87f83905ba29a56153b2a1e1eaf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/RECORD +++ /dev/null @@ -1,108 +0,0 @@ -../../../bin/chardetect,sha256=D8Tk5mkski9J5qXuc_Nvb3Lyg5tUxOJcBO1YUWPzrrU,239 -chardet-5.2.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -chardet-5.2.0.dist-info/LICENSE,sha256=3GJlINzVOiL3J68-5Cx3DlbJemT-OtsGN5nYqwMv5VE,26530 -chardet-5.2.0.dist-info/METADATA,sha256=PAr2NQ6hQWpjyFnwlI7MoxHt2S_6oRiUsucOKMNhzGw,3418 -chardet-5.2.0.dist-info/RECORD,, -chardet-5.2.0.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -chardet-5.2.0.dist-info/WHEEL,sha256=AtBG6SXL3KF_v0NxLf0ehyVOh0cold-JbJYXNGorC6Q,92 -chardet-5.2.0.dist-info/entry_points.txt,sha256=_cdvYc4jyY68GYfsQAAthNMxO-yodcGkvNC1xOEsLmI,59 -chardet-5.2.0.dist-info/top_level.txt,sha256=AowzBbZy4x8EirABDdJSLJZMkJ_53iIag8xfKR6D7kI,8 -chardet/__init__.py,sha256=57R-HSxj0PWmILMN0GFmUNqEMfrEVSamXyjD-W6_fbs,4797 -chardet/__main__.py,sha256=puNj2o_QfBRKElEkiVp1zEIL1gGYD2o-JuXLFlqHDC4,123 -chardet/__pycache__/__init__.cpython-310.pyc,, -chardet/__pycache__/__main__.cpython-310.pyc,, -chardet/__pycache__/big5freq.cpython-310.pyc,, -chardet/__pycache__/big5prober.cpython-310.pyc,, -chardet/__pycache__/chardistribution.cpython-310.pyc,, -chardet/__pycache__/charsetgroupprober.cpython-310.pyc,, -chardet/__pycache__/charsetprober.cpython-310.pyc,, -chardet/__pycache__/codingstatemachine.cpython-310.pyc,, -chardet/__pycache__/codingstatemachinedict.cpython-310.pyc,, -chardet/__pycache__/cp949prober.cpython-310.pyc,, -chardet/__pycache__/enums.cpython-310.pyc,, -chardet/__pycache__/escprober.cpython-310.pyc,, -chardet/__pycache__/escsm.cpython-310.pyc,, -chardet/__pycache__/eucjpprober.cpython-310.pyc,, -chardet/__pycache__/euckrfreq.cpython-310.pyc,, -chardet/__pycache__/euckrprober.cpython-310.pyc,, -chardet/__pycache__/euctwfreq.cpython-310.pyc,, -chardet/__pycache__/euctwprober.cpython-310.pyc,, -chardet/__pycache__/gb2312freq.cpython-310.pyc,, -chardet/__pycache__/gb2312prober.cpython-310.pyc,, -chardet/__pycache__/hebrewprober.cpython-310.pyc,, -chardet/__pycache__/jisfreq.cpython-310.pyc,, -chardet/__pycache__/johabfreq.cpython-310.pyc,, -chardet/__pycache__/johabprober.cpython-310.pyc,, -chardet/__pycache__/jpcntx.cpython-310.pyc,, -chardet/__pycache__/langbulgarianmodel.cpython-310.pyc,, -chardet/__pycache__/langgreekmodel.cpython-310.pyc,, -chardet/__pycache__/langhebrewmodel.cpython-310.pyc,, -chardet/__pycache__/langhungarianmodel.cpython-310.pyc,, -chardet/__pycache__/langrussianmodel.cpython-310.pyc,, -chardet/__pycache__/langthaimodel.cpython-310.pyc,, -chardet/__pycache__/langturkishmodel.cpython-310.pyc,, -chardet/__pycache__/latin1prober.cpython-310.pyc,, -chardet/__pycache__/macromanprober.cpython-310.pyc,, -chardet/__pycache__/mbcharsetprober.cpython-310.pyc,, -chardet/__pycache__/mbcsgroupprober.cpython-310.pyc,, -chardet/__pycache__/mbcssm.cpython-310.pyc,, -chardet/__pycache__/resultdict.cpython-310.pyc,, -chardet/__pycache__/sbcharsetprober.cpython-310.pyc,, -chardet/__pycache__/sbcsgroupprober.cpython-310.pyc,, -chardet/__pycache__/sjisprober.cpython-310.pyc,, -chardet/__pycache__/universaldetector.cpython-310.pyc,, -chardet/__pycache__/utf1632prober.cpython-310.pyc,, -chardet/__pycache__/utf8prober.cpython-310.pyc,, -chardet/__pycache__/version.cpython-310.pyc,, -chardet/big5freq.py,sha256=ltcfP-3PjlNHCoo5e4a7C4z-2DhBTXRfY6jbMbB7P30,31274 -chardet/big5prober.py,sha256=lPMfwCX6v2AaPgvFh_cSWZcgLDbWiFCHLZ_p9RQ9uxE,1763 -chardet/chardistribution.py,sha256=13B8XUG4oXDuLdXvfbIWwLFeR-ZU21AqTS1zcdON8bU,10032 -chardet/charsetgroupprober.py,sha256=UKK3SaIZB2PCdKSIS0gnvMtLR9JJX62M-fZJu3OlWyg,3915 -chardet/charsetprober.py,sha256=L3t8_wIOov8em-vZWOcbkdsrwe43N6_gqNh5pH7WPd4,5420 -chardet/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -chardet/cli/__pycache__/__init__.cpython-310.pyc,, -chardet/cli/__pycache__/chardetect.cpython-310.pyc,, -chardet/cli/chardetect.py,sha256=zibMVg5RpKb-ME9_7EYG4ZM2Sf07NHcQzZ12U-rYJho,3242 -chardet/codingstatemachine.py,sha256=K7k69sw3jY5DmTXoSJQVsUtFIQKYPQVOSJJhBuGv_yE,3732 -chardet/codingstatemachinedict.py,sha256=0GY3Hi2qIZvDrOOJ3AtqppM1RsYxr_66ER4EHjuMiMc,542 -chardet/cp949prober.py,sha256=0jKRV7fECuWI16rNnks0ZECKA1iZYCIEaP8A1ZvjUSI,1860 -chardet/enums.py,sha256=TzECiZoCKNMqgwU76cPCeKWFBqaWvAdLMev5_bCkhY8,1683 -chardet/escprober.py,sha256=Kho48X65xE0scFylIdeJjM2bcbvRvv0h0WUbMWrJD3A,4006 -chardet/escsm.py,sha256=AqyXpA2FQFD7k-buBty_7itGEYkhmVa8X09NLRul3QM,12176 -chardet/eucjpprober.py,sha256=5KYaM9fsxkRYzw1b5k0fL-j_-ezIw-ij9r97a9MHxLY,3934 -chardet/euckrfreq.py,sha256=3mHuRvXfsq_QcQysDQFb8qSudvTiol71C6Ic2w57tKM,13566 -chardet/euckrprober.py,sha256=hiFT6wM174GIwRvqDsIcuOc-dDsq2uPKMKbyV8-1Xnc,1753 -chardet/euctwfreq.py,sha256=2alILE1Lh5eqiFJZjzRkMQXolNJRHY5oBQd-vmZYFFM,36913 -chardet/euctwprober.py,sha256=NxbpNdBtU0VFI0bKfGfDkpP7S2_8_6FlO87dVH0ogws,1753 -chardet/gb2312freq.py,sha256=49OrdXzD-HXqwavkqjo8Z7gvs58hONNzDhAyMENNkvY,20735 -chardet/gb2312prober.py,sha256=KPEBueaSLSvBpFeINMu0D6TgHcR90e5PaQawifzF4o0,1759 -chardet/hebrewprober.py,sha256=96T_Lj_OmW-fK7JrSHojYjyG3fsGgbzkoTNleZ3kfYE,14537 -chardet/jisfreq.py,sha256=mm8tfrwqhpOd3wzZKS4NJqkYBQVcDfTM2JiQ5aW932E,25796 -chardet/johabfreq.py,sha256=dBpOYG34GRX6SL8k_LbS9rxZPMjLjoMlgZ03Pz5Hmqc,42498 -chardet/johabprober.py,sha256=O1Qw9nVzRnun7vZp4UZM7wvJSv9W941mEU9uDMnY3DU,1752 -chardet/jpcntx.py,sha256=uhHrYWkLxE_rF5OkHKInm0HUsrjgKHHVQvtt3UcvotA,27055 -chardet/langbulgarianmodel.py,sha256=bGoRpxBYtrbSHa6mX6PkEA26v30pWmhDjemhdxmkew8,104550 -chardet/langgreekmodel.py,sha256=3wMlEzQ8oU2MbrL2xN8lkuOB0dCMLBhW6heekxusoc0,98472 -chardet/langhebrewmodel.py,sha256=ZUTqusxMvR_earWPs5w-rH10xoe5sPjd9FLMu1DUIvE,98184 -chardet/langhungarianmodel.py,sha256=N-YtC2EiswyS7XsUicCPRycrIzRNj47Y048odp9qOoo,101351 -chardet/langrussianmodel.py,sha256=6v7RcZKGj0VH0864BHzizKNceAYbHvGts2p00ifC7w4,128023 -chardet/langthaimodel.py,sha256=Mr673U9U8rkQFfUDtLP01pp-0TOsl2o6sb75YEjvpcs,102762 -chardet/langturkishmodel.py,sha256=LkXCjWhGUEzqKXvfasHN0SFBigwKJ3xeWNVZ0EyI0kA,95360 -chardet/latin1prober.py,sha256=p15EEmFbmQUwbKLC7lOJVGHEZwcG45ubEZYTGu01J5g,5380 -chardet/macromanprober.py,sha256=9anfzmY6TBfUPDyBDOdY07kqmTHpZ1tK0jL-p1JWcOY,6077 -chardet/mbcharsetprober.py,sha256=Wr04WNI4F3X_VxEverNG-H25g7u-MDDKlNt-JGj-_uU,3715 -chardet/mbcsgroupprober.py,sha256=iRpaNBjV0DNwYPu_z6TiHgRpwYahiM7ztI_4kZ4Uz9A,2131 -chardet/mbcssm.py,sha256=hUtPvDYgWDaA2dWdgLsshbwRfm3Q5YRlRogdmeRUNQw,30391 -chardet/metadata/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -chardet/metadata/__pycache__/__init__.cpython-310.pyc,, -chardet/metadata/__pycache__/languages.cpython-310.pyc,, -chardet/metadata/languages.py,sha256=FhvBIdZFxRQ-dTwkb_0madRKgVBCaUMQz9I5xqjE5iQ,13560 -chardet/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -chardet/resultdict.py,sha256=ez4FRvN5KaSosJeJ2WzUyKdDdg35HDy_SSLPXKCdt5M,402 -chardet/sbcharsetprober.py,sha256=-nd3F90i7GpXLjehLVHqVBE0KlWzGvQUPETLBNn4o6U,6400 -chardet/sbcsgroupprober.py,sha256=gcgI0fOfgw_3YTClpbra_MNxwyEyJ3eUXraoLHYb59E,4137 -chardet/sjisprober.py,sha256=aqQufMzRw46ZpFlzmYaYeT2-nzmKb-hmcrApppJ862k,4007 -chardet/universaldetector.py,sha256=xYBrg4x0dd9WnT8qclfADVD9ondrUNkqPmvte1pa520,14848 -chardet/utf1632prober.py,sha256=pw1epGdMj1hDGiCu1AHqqzOEfjX8MVdiW7O1BlT8-eQ,8505 -chardet/utf8prober.py,sha256=8m08Ub5490H4jQ6LYXvFysGtgKoKsHUd2zH_i8_TnVw,2812 -chardet/version.py,sha256=jp8ePp1zC63YxruGcHSuKxtf3-fF1LYAMUZD2eDWYok,244 diff --git a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/REQUESTED deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/WHEEL b/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/WHEEL deleted file mode 100644 index d272f6ed555bf206d2a9572524bfa3c0b500fe8d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/WHEEL +++ /dev/null @@ -1,5 +0,0 @@ -Wheel-Version: 1.0 -Generator: bdist_wheel (0.41.0) -Root-Is-Purelib: true -Tag: py3-none-any - diff --git a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/entry_points.txt b/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/entry_points.txt deleted file mode 100644 index c36a5e3ee479856d913cae1521b4bd2122587325..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/entry_points.txt +++ /dev/null @@ -1,2 +0,0 @@ -[console_scripts] -chardetect = chardet.cli.chardetect:main diff --git a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/top_level.txt deleted file mode 100644 index 79236f25cda563eb57cd7b5838256d9f3fdf18ab..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet-5.2.0.dist-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -chardet diff --git a/venv/lib/python3.10/site-packages/chardet/__init__.py b/venv/lib/python3.10/site-packages/chardet/__init__.py deleted file mode 100644 index fe581623d89d67a49eb43f3c3e88f3f450257707..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/__init__.py +++ /dev/null @@ -1,115 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from typing import List, Union - -from .charsetgroupprober import CharSetGroupProber -from .charsetprober import CharSetProber -from .enums import InputState -from .resultdict import ResultDict -from .universaldetector import UniversalDetector -from .version import VERSION, __version__ - -__all__ = ["UniversalDetector", "detect", "detect_all", "__version__", "VERSION"] - - -def detect( - byte_str: Union[bytes, bytearray], should_rename_legacy: bool = False -) -> ResultDict: - """ - Detect the encoding of the given byte string. - - :param byte_str: The byte sequence to examine. - :type byte_str: ``bytes`` or ``bytearray`` - :param should_rename_legacy: Should we rename legacy encodings - to their more modern equivalents? - :type should_rename_legacy: ``bool`` - """ - if not isinstance(byte_str, bytearray): - if not isinstance(byte_str, bytes): - raise TypeError( - f"Expected object of type bytes or bytearray, got: {type(byte_str)}" - ) - byte_str = bytearray(byte_str) - detector = UniversalDetector(should_rename_legacy=should_rename_legacy) - detector.feed(byte_str) - return detector.close() - - -def detect_all( - byte_str: Union[bytes, bytearray], - ignore_threshold: bool = False, - should_rename_legacy: bool = False, -) -> List[ResultDict]: - """ - Detect all the possible encodings of the given byte string. - - :param byte_str: The byte sequence to examine. - :type byte_str: ``bytes`` or ``bytearray`` - :param ignore_threshold: Include encodings that are below - ``UniversalDetector.MINIMUM_THRESHOLD`` - in results. - :type ignore_threshold: ``bool`` - :param should_rename_legacy: Should we rename legacy encodings - to their more modern equivalents? - :type should_rename_legacy: ``bool`` - """ - if not isinstance(byte_str, bytearray): - if not isinstance(byte_str, bytes): - raise TypeError( - f"Expected object of type bytes or bytearray, got: {type(byte_str)}" - ) - byte_str = bytearray(byte_str) - - detector = UniversalDetector(should_rename_legacy=should_rename_legacy) - detector.feed(byte_str) - detector.close() - - if detector.input_state == InputState.HIGH_BYTE: - results: List[ResultDict] = [] - probers: List[CharSetProber] = [] - for prober in detector.charset_probers: - if isinstance(prober, CharSetGroupProber): - probers.extend(p for p in prober.probers) - else: - probers.append(prober) - for prober in probers: - if ignore_threshold or prober.get_confidence() > detector.MINIMUM_THRESHOLD: - charset_name = prober.charset_name or "" - lower_charset_name = charset_name.lower() - # Use Windows encoding name instead of ISO-8859 if we saw any - # extra Windows-specific bytes - if lower_charset_name.startswith("iso-8859") and detector.has_win_bytes: - charset_name = detector.ISO_WIN_MAP.get( - lower_charset_name, charset_name - ) - # Rename legacy encodings with superset encodings if asked - if should_rename_legacy: - charset_name = detector.LEGACY_MAP.get( - charset_name.lower(), charset_name - ) - results.append( - { - "encoding": charset_name, - "confidence": prober.get_confidence(), - "language": prober.language, - } - ) - if len(results) > 0: - return sorted(results, key=lambda result: -result["confidence"]) - - return [detector.result] diff --git a/venv/lib/python3.10/site-packages/chardet/__main__.py b/venv/lib/python3.10/site-packages/chardet/__main__.py deleted file mode 100644 index c19b0d2d7a3d8fa0d914ca4031e88241fd63a90f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/__main__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Wrapper so people can run python -m chardet""" - -from .cli.chardetect import main - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 7ebc8a831331d357f3295cb8ea6fe186c2deb05f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/__main__.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/__main__.cpython-310.pyc deleted file mode 100644 index 1b1e02abaf19fc34c76fd1c70085cb80223433eb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/__main__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/big5freq.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/big5freq.cpython-310.pyc deleted file mode 100644 index bc05119bd89b170cb6d69a883a731cd4c5a00ab9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/big5freq.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/big5prober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/big5prober.cpython-310.pyc deleted file mode 100644 index 6fe7c45a9db228bca070746c4b30246ab76879ef..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/big5prober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/chardistribution.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/chardistribution.cpython-310.pyc deleted file mode 100644 index e8764ff8f304846a4ff91a87582a4f5d52adbd55..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/chardistribution.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/charsetgroupprober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/charsetgroupprober.cpython-310.pyc deleted file mode 100644 index b72f8fe2aaa5269d50e1fe9ad586b021efb0e05f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/charsetgroupprober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/charsetprober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/charsetprober.cpython-310.pyc deleted file mode 100644 index a76676cbd6d8e337bea2fc7b5cbcf3ada8727458..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/charsetprober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/codingstatemachine.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/codingstatemachine.cpython-310.pyc deleted file mode 100644 index fbfa2d17009b81de57425c6396eb1ae4e5d49062..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/codingstatemachine.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/codingstatemachinedict.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/codingstatemachinedict.cpython-310.pyc deleted file mode 100644 index 5dfd30b81cd0189a66fdf15a0d7029ca21b9195a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/codingstatemachinedict.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/cp949prober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/cp949prober.cpython-310.pyc deleted file mode 100644 index db4a778014e7e1fa7d939379a6d34852a7863f41..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/cp949prober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/enums.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/enums.cpython-310.pyc deleted file mode 100644 index 3c142bde90be92a0f882e7c716f0ba81fe1546e7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/enums.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/escprober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/escprober.cpython-310.pyc deleted file mode 100644 index 698546bfc09fae2be11a0f22b43c8de0c2088a09..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/escprober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/escsm.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/escsm.cpython-310.pyc deleted file mode 100644 index 9a4ac4563dda29519ac76a994eb48226b295debb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/escsm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/eucjpprober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/eucjpprober.cpython-310.pyc deleted file mode 100644 index e1b52626d96fbd7f6f1ffeeb84161939f63d3aba..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/eucjpprober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/euckrfreq.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/euckrfreq.cpython-310.pyc deleted file mode 100644 index 607ba290d7569555f3feca2cfe8aa9fc7b56a51b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/euckrfreq.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/euckrprober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/euckrprober.cpython-310.pyc deleted file mode 100644 index 537b729c1a468b064e5dbe1cb7c9f23c56d5ed55..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/euckrprober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/euctwfreq.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/euctwfreq.cpython-310.pyc deleted file mode 100644 index e41364e50009903e3ddc3f6bd579a3de31c36681..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/euctwfreq.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/euctwprober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/euctwprober.cpython-310.pyc deleted file mode 100644 index f592f7532a37c43ae25f72e4024da8065b11dd23..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/euctwprober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/gb2312freq.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/gb2312freq.cpython-310.pyc deleted file mode 100644 index fd1c406370c9ef189c34dd1b8a20691a828d1acd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/gb2312freq.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/gb2312prober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/gb2312prober.cpython-310.pyc deleted file mode 100644 index de89a32cb7358880f543fa4c8e2f7e1707c8d8cc..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/gb2312prober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/hebrewprober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/hebrewprober.cpython-310.pyc deleted file mode 100644 index d08f3e6c312dfbd9a04a65d0febb5a292ce0192b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/hebrewprober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/jisfreq.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/jisfreq.cpython-310.pyc deleted file mode 100644 index d82517a692e6b28618042fb98f19fe3d9d247944..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/jisfreq.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/johabfreq.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/johabfreq.cpython-310.pyc deleted file mode 100644 index 5be12ab2067520848637faca1ce65a54838ad74e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/__pycache__/johabfreq.cpython-310.pyc +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:102bbbed9ec33e9bc16aa3fe9e54c5448673808d0bb6d4180647091d8e8e418c -size 138732 diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/johabprober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/johabprober.cpython-310.pyc deleted file mode 100644 index 9e5683b75bf78f1ba838092f1830f181022bc377..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/johabprober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/jpcntx.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/jpcntx.cpython-310.pyc deleted file mode 100644 index 830e04b678bde80afbb08e7d335ffdde76cfbb17..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/jpcntx.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/langbulgarianmodel.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/langbulgarianmodel.cpython-310.pyc deleted file mode 100644 index e7ce9ec2ec8268f7665d460a52f4932574be204d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/langbulgarianmodel.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/langgreekmodel.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/langgreekmodel.cpython-310.pyc deleted file mode 100644 index b59aa911201f5b16ed6b6dbe10fe480007ec318a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/langgreekmodel.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/langhebrewmodel.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/langhebrewmodel.cpython-310.pyc deleted file mode 100644 index 2b61d1df5dc901353a215ac0396bfddd3507e401..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/langhebrewmodel.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/langhungarianmodel.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/langhungarianmodel.cpython-310.pyc deleted file mode 100644 index becf0822a263650e4877c002749debb2fd6e8415..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/langhungarianmodel.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/langrussianmodel.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/langrussianmodel.cpython-310.pyc deleted file mode 100644 index 9bdfa3d41b7d096a7813e79070963a8350b2498c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/langrussianmodel.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/langthaimodel.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/langthaimodel.cpython-310.pyc deleted file mode 100644 index 481acf5398137fcb35cef113f53746dc4ffc3e0d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/langthaimodel.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/langturkishmodel.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/langturkishmodel.cpython-310.pyc deleted file mode 100644 index 3f71f7ffaf0f2ab2e22ba3f9a6a4148e72f8255c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/langturkishmodel.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/latin1prober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/latin1prober.cpython-310.pyc deleted file mode 100644 index 70079f7b04a2ce5068011cbfc3a414eb7be729d7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/latin1prober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/macromanprober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/macromanprober.cpython-310.pyc deleted file mode 100644 index 766d461820fa6fd0bd8cfb8401f263bbad38a6f5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/macromanprober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/mbcharsetprober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/mbcharsetprober.cpython-310.pyc deleted file mode 100644 index 56a13f08404472b2371d574f2536c4b55a4ed1d7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/mbcharsetprober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/mbcsgroupprober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/mbcsgroupprober.cpython-310.pyc deleted file mode 100644 index 31c156a907ef043f01c8ac4a4d1ba3f85fea9f05..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/mbcsgroupprober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/mbcssm.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/mbcssm.cpython-310.pyc deleted file mode 100644 index 0b5c8e4b74cc120f32267da0add512b2f0b16f99..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/mbcssm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/resultdict.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/resultdict.cpython-310.pyc deleted file mode 100644 index e5f9e9d4537f92c4086b204d1236ae8174623ba2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/resultdict.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/sbcharsetprober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/sbcharsetprober.cpython-310.pyc deleted file mode 100644 index ceb867142e79bf0f11c71a39b30cb0e7f0ea2602..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/sbcharsetprober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/sbcsgroupprober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/sbcsgroupprober.cpython-310.pyc deleted file mode 100644 index 6cd95becc50bc6a790e0f91939e52619c6442c6f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/sbcsgroupprober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/sjisprober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/sjisprober.cpython-310.pyc deleted file mode 100644 index 99b0e2065ca45301e3982fc65c2780b2b98f7c1f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/sjisprober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/universaldetector.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/universaldetector.cpython-310.pyc deleted file mode 100644 index e3b0d2baaa5791507ea2f5b215acb0dfe76ddf25..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/universaldetector.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/utf1632prober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/utf1632prober.cpython-310.pyc deleted file mode 100644 index 55d9795a193bd8484e16a886e9f82b40e65915e0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/utf1632prober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/utf8prober.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/utf8prober.cpython-310.pyc deleted file mode 100644 index c2f56916338a440756ccc3555b5853d8c808fd55..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/utf8prober.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/__pycache__/version.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/__pycache__/version.cpython-310.pyc deleted file mode 100644 index be2eb4a0ae007fe7b05a3e5b932f5b565259c873..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/__pycache__/version.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/big5freq.py b/venv/lib/python3.10/site-packages/chardet/big5freq.py deleted file mode 100644 index 87d9f972edde20d1f8e391b8010703242a8de977..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/big5freq.py +++ /dev/null @@ -1,386 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Communicator client code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -# Big5 frequency table -# by Taiwan's Mandarin Promotion Council -# -# -# 128 --> 0.42261 -# 256 --> 0.57851 -# 512 --> 0.74851 -# 1024 --> 0.89384 -# 2048 --> 0.97583 -# -# Ideal Distribution Ratio = 0.74851/(1-0.74851) =2.98 -# Random Distribution Ration = 512/(5401-512)=0.105 -# -# Typical Distribution Ratio about 25% of Ideal one, still much higher than RDR - -BIG5_TYPICAL_DISTRIBUTION_RATIO = 0.75 - -# Char to FreqOrder table -BIG5_TABLE_SIZE = 5376 -# fmt: off -BIG5_CHAR_TO_FREQ_ORDER = ( - 1,1801,1506, 255,1431, 198, 9, 82, 6,5008, 177, 202,3681,1256,2821, 110, # 16 -3814, 33,3274, 261, 76, 44,2114, 16,2946,2187,1176, 659,3971, 26,3451,2653, # 32 -1198,3972,3350,4202, 410,2215, 302, 590, 361,1964, 8, 204, 58,4510,5009,1932, # 48 - 63,5010,5011, 317,1614, 75, 222, 159,4203,2417,1480,5012,3555,3091, 224,2822, # 64 -3682, 3, 10,3973,1471, 29,2787,1135,2866,1940, 873, 130,3275,1123, 312,5013, # 80 -4511,2052, 507, 252, 682,5014, 142,1915, 124, 206,2947, 34,3556,3204, 64, 604, # 96 -5015,2501,1977,1978, 155,1991, 645, 641,1606,5016,3452, 337, 72, 406,5017, 80, # 112 - 630, 238,3205,1509, 263, 939,1092,2654, 756,1440,1094,3453, 449, 69,2987, 591, # 128 - 179,2096, 471, 115,2035,1844, 60, 50,2988, 134, 806,1869, 734,2036,3454, 180, # 144 - 995,1607, 156, 537,2907, 688,5018, 319,1305, 779,2145, 514,2379, 298,4512, 359, # 160 -2502, 90,2716,1338, 663, 11, 906,1099,2553, 20,2441, 182, 532,1716,5019, 732, # 176 -1376,4204,1311,1420,3206, 25,2317,1056, 113, 399, 382,1950, 242,3455,2474, 529, # 192 -3276, 475,1447,3683,5020, 117, 21, 656, 810,1297,2300,2334,3557,5021, 126,4205, # 208 - 706, 456, 150, 613,4513, 71,1118,2037,4206, 145,3092, 85, 835, 486,2115,1246, # 224 -1426, 428, 727,1285,1015, 800, 106, 623, 303,1281,5022,2128,2359, 347,3815, 221, # 240 -3558,3135,5023,1956,1153,4207, 83, 296,1199,3093, 192, 624, 93,5024, 822,1898, # 256 -2823,3136, 795,2065, 991,1554,1542,1592, 27, 43,2867, 859, 139,1456, 860,4514, # 272 - 437, 712,3974, 164,2397,3137, 695, 211,3037,2097, 195,3975,1608,3559,3560,3684, # 288 -3976, 234, 811,2989,2098,3977,2233,1441,3561,1615,2380, 668,2077,1638, 305, 228, # 304 -1664,4515, 467, 415,5025, 262,2099,1593, 239, 108, 300, 200,1033, 512,1247,2078, # 320 -5026,5027,2176,3207,3685,2682, 593, 845,1062,3277, 88,1723,2038,3978,1951, 212, # 336 - 266, 152, 149, 468,1899,4208,4516, 77, 187,5028,3038, 37, 5,2990,5029,3979, # 352 -5030,5031, 39,2524,4517,2908,3208,2079, 55, 148, 74,4518, 545, 483,1474,1029, # 368 -1665, 217,1870,1531,3138,1104,2655,4209, 24, 172,3562, 900,3980,3563,3564,4519, # 384 - 32,1408,2824,1312, 329, 487,2360,2251,2717, 784,2683, 4,3039,3351,1427,1789, # 400 - 188, 109, 499,5032,3686,1717,1790, 888,1217,3040,4520,5033,3565,5034,3352,1520, # 416 -3687,3981, 196,1034, 775,5035,5036, 929,1816, 249, 439, 38,5037,1063,5038, 794, # 432 -3982,1435,2301, 46, 178,3278,2066,5039,2381,5040, 214,1709,4521, 804, 35, 707, # 448 - 324,3688,1601,2554, 140, 459,4210,5041,5042,1365, 839, 272, 978,2262,2580,3456, # 464 -2129,1363,3689,1423, 697, 100,3094, 48, 70,1231, 495,3139,2196,5043,1294,5044, # 480 -2080, 462, 586,1042,3279, 853, 256, 988, 185,2382,3457,1698, 434,1084,5045,3458, # 496 - 314,2625,2788,4522,2335,2336, 569,2285, 637,1817,2525, 757,1162,1879,1616,3459, # 512 - 287,1577,2116, 768,4523,1671,2868,3566,2526,1321,3816, 909,2418,5046,4211, 933, # 528 -3817,4212,2053,2361,1222,4524, 765,2419,1322, 786,4525,5047,1920,1462,1677,2909, # 544 -1699,5048,4526,1424,2442,3140,3690,2600,3353,1775,1941,3460,3983,4213, 309,1369, # 560 -1130,2825, 364,2234,1653,1299,3984,3567,3985,3986,2656, 525,1085,3041, 902,2001, # 576 -1475, 964,4527, 421,1845,1415,1057,2286, 940,1364,3141, 376,4528,4529,1381, 7, # 592 -2527, 983,2383, 336,1710,2684,1846, 321,3461, 559,1131,3042,2752,1809,1132,1313, # 608 - 265,1481,1858,5049, 352,1203,2826,3280, 167,1089, 420,2827, 776, 792,1724,3568, # 624 -4214,2443,3281,5050,4215,5051, 446, 229, 333,2753, 901,3818,1200,1557,4530,2657, # 640 -1921, 395,2754,2685,3819,4216,1836, 125, 916,3209,2626,4531,5052,5053,3820,5054, # 656 -5055,5056,4532,3142,3691,1133,2555,1757,3462,1510,2318,1409,3569,5057,2146, 438, # 672 -2601,2910,2384,3354,1068, 958,3043, 461, 311,2869,2686,4217,1916,3210,4218,1979, # 688 - 383, 750,2755,2627,4219, 274, 539, 385,1278,1442,5058,1154,1965, 384, 561, 210, # 704 - 98,1295,2556,3570,5059,1711,2420,1482,3463,3987,2911,1257, 129,5060,3821, 642, # 720 - 523,2789,2790,2658,5061, 141,2235,1333, 68, 176, 441, 876, 907,4220, 603,2602, # 736 - 710, 171,3464, 404, 549, 18,3143,2398,1410,3692,1666,5062,3571,4533,2912,4534, # 752 -5063,2991, 368,5064, 146, 366, 99, 871,3693,1543, 748, 807,1586,1185, 22,2263, # 768 - 379,3822,3211,5065,3212, 505,1942,2628,1992,1382,2319,5066, 380,2362, 218, 702, # 784 -1818,1248,3465,3044,3572,3355,3282,5067,2992,3694, 930,3283,3823,5068, 59,5069, # 800 - 585, 601,4221, 497,3466,1112,1314,4535,1802,5070,1223,1472,2177,5071, 749,1837, # 816 - 690,1900,3824,1773,3988,1476, 429,1043,1791,2236,2117, 917,4222, 447,1086,1629, # 832 -5072, 556,5073,5074,2021,1654, 844,1090, 105, 550, 966,1758,2828,1008,1783, 686, # 848 -1095,5075,2287, 793,1602,5076,3573,2603,4536,4223,2948,2302,4537,3825, 980,2503, # 864 - 544, 353, 527,4538, 908,2687,2913,5077, 381,2629,1943,1348,5078,1341,1252, 560, # 880 -3095,5079,3467,2870,5080,2054, 973, 886,2081, 143,4539,5081,5082, 157,3989, 496, # 896 -4224, 57, 840, 540,2039,4540,4541,3468,2118,1445, 970,2264,1748,1966,2082,4225, # 912 -3144,1234,1776,3284,2829,3695, 773,1206,2130,1066,2040,1326,3990,1738,1725,4226, # 928 - 279,3145, 51,1544,2604, 423,1578,2131,2067, 173,4542,1880,5083,5084,1583, 264, # 944 - 610,3696,4543,2444, 280, 154,5085,5086,5087,1739, 338,1282,3096, 693,2871,1411, # 960 -1074,3826,2445,5088,4544,5089,5090,1240, 952,2399,5091,2914,1538,2688, 685,1483, # 976 -4227,2475,1436, 953,4228,2055,4545, 671,2400, 79,4229,2446,3285, 608, 567,2689, # 992 -3469,4230,4231,1691, 393,1261,1792,2401,5092,4546,5093,5094,5095,5096,1383,1672, # 1008 -3827,3213,1464, 522,1119, 661,1150, 216, 675,4547,3991,1432,3574, 609,4548,2690, # 1024 -2402,5097,5098,5099,4232,3045, 0,5100,2476, 315, 231,2447, 301,3356,4549,2385, # 1040 -5101, 233,4233,3697,1819,4550,4551,5102, 96,1777,1315,2083,5103, 257,5104,1810, # 1056 -3698,2718,1139,1820,4234,2022,1124,2164,2791,1778,2659,5105,3097, 363,1655,3214, # 1072 -5106,2993,5107,5108,5109,3992,1567,3993, 718, 103,3215, 849,1443, 341,3357,2949, # 1088 -1484,5110,1712, 127, 67, 339,4235,2403, 679,1412, 821,5111,5112, 834, 738, 351, # 1104 -2994,2147, 846, 235,1497,1881, 418,1993,3828,2719, 186,1100,2148,2756,3575,1545, # 1120 -1355,2950,2872,1377, 583,3994,4236,2581,2995,5113,1298,3699,1078,2557,3700,2363, # 1136 - 78,3829,3830, 267,1289,2100,2002,1594,4237, 348, 369,1274,2197,2178,1838,4552, # 1152 -1821,2830,3701,2757,2288,2003,4553,2951,2758, 144,3358, 882,4554,3995,2759,3470, # 1168 -4555,2915,5114,4238,1726, 320,5115,3996,3046, 788,2996,5116,2831,1774,1327,2873, # 1184 -3997,2832,5117,1306,4556,2004,1700,3831,3576,2364,2660, 787,2023, 506, 824,3702, # 1200 - 534, 323,4557,1044,3359,2024,1901, 946,3471,5118,1779,1500,1678,5119,1882,4558, # 1216 - 165, 243,4559,3703,2528, 123, 683,4239, 764,4560, 36,3998,1793, 589,2916, 816, # 1232 - 626,1667,3047,2237,1639,1555,1622,3832,3999,5120,4000,2874,1370,1228,1933, 891, # 1248 -2084,2917, 304,4240,5121, 292,2997,2720,3577, 691,2101,4241,1115,4561, 118, 662, # 1264 -5122, 611,1156, 854,2386,1316,2875, 2, 386, 515,2918,5123,5124,3286, 868,2238, # 1280 -1486, 855,2661, 785,2216,3048,5125,1040,3216,3578,5126,3146, 448,5127,1525,5128, # 1296 -2165,4562,5129,3833,5130,4242,2833,3579,3147, 503, 818,4001,3148,1568, 814, 676, # 1312 -1444, 306,1749,5131,3834,1416,1030, 197,1428, 805,2834,1501,4563,5132,5133,5134, # 1328 -1994,5135,4564,5136,5137,2198, 13,2792,3704,2998,3149,1229,1917,5138,3835,2132, # 1344 -5139,4243,4565,2404,3580,5140,2217,1511,1727,1120,5141,5142, 646,3836,2448, 307, # 1360 -5143,5144,1595,3217,5145,5146,5147,3705,1113,1356,4002,1465,2529,2530,5148, 519, # 1376 -5149, 128,2133, 92,2289,1980,5150,4003,1512, 342,3150,2199,5151,2793,2218,1981, # 1392 -3360,4244, 290,1656,1317, 789, 827,2365,5152,3837,4566, 562, 581,4004,5153, 401, # 1408 -4567,2252, 94,4568,5154,1399,2794,5155,1463,2025,4569,3218,1944,5156, 828,1105, # 1424 -4245,1262,1394,5157,4246, 605,4570,5158,1784,2876,5159,2835, 819,2102, 578,2200, # 1440 -2952,5160,1502, 436,3287,4247,3288,2836,4005,2919,3472,3473,5161,2721,2320,5162, # 1456 -5163,2337,2068, 23,4571, 193, 826,3838,2103, 699,1630,4248,3098, 390,1794,1064, # 1472 -3581,5164,1579,3099,3100,1400,5165,4249,1839,1640,2877,5166,4572,4573, 137,4250, # 1488 - 598,3101,1967, 780, 104, 974,2953,5167, 278, 899, 253, 402, 572, 504, 493,1339, # 1504 -5168,4006,1275,4574,2582,2558,5169,3706,3049,3102,2253, 565,1334,2722, 863, 41, # 1520 -5170,5171,4575,5172,1657,2338, 19, 463,2760,4251, 606,5173,2999,3289,1087,2085, # 1536 -1323,2662,3000,5174,1631,1623,1750,4252,2691,5175,2878, 791,2723,2663,2339, 232, # 1552 -2421,5176,3001,1498,5177,2664,2630, 755,1366,3707,3290,3151,2026,1609, 119,1918, # 1568 -3474, 862,1026,4253,5178,4007,3839,4576,4008,4577,2265,1952,2477,5179,1125, 817, # 1584 -4254,4255,4009,1513,1766,2041,1487,4256,3050,3291,2837,3840,3152,5180,5181,1507, # 1600 -5182,2692, 733, 40,1632,1106,2879, 345,4257, 841,2531, 230,4578,3002,1847,3292, # 1616 -3475,5183,1263, 986,3476,5184, 735, 879, 254,1137, 857, 622,1300,1180,1388,1562, # 1632 -4010,4011,2954, 967,2761,2665,1349, 592,2134,1692,3361,3003,1995,4258,1679,4012, # 1648 -1902,2188,5185, 739,3708,2724,1296,1290,5186,4259,2201,2202,1922,1563,2605,2559, # 1664 -1871,2762,3004,5187, 435,5188, 343,1108, 596, 17,1751,4579,2239,3477,3709,5189, # 1680 -4580, 294,3582,2955,1693, 477, 979, 281,2042,3583, 643,2043,3710,2631,2795,2266, # 1696 -1031,2340,2135,2303,3584,4581, 367,1249,2560,5190,3585,5191,4582,1283,3362,2005, # 1712 - 240,1762,3363,4583,4584, 836,1069,3153, 474,5192,2149,2532, 268,3586,5193,3219, # 1728 -1521,1284,5194,1658,1546,4260,5195,3587,3588,5196,4261,3364,2693,1685,4262, 961, # 1744 -1673,2632, 190,2006,2203,3841,4585,4586,5197, 570,2504,3711,1490,5198,4587,2633, # 1760 -3293,1957,4588, 584,1514, 396,1045,1945,5199,4589,1968,2449,5200,5201,4590,4013, # 1776 - 619,5202,3154,3294, 215,2007,2796,2561,3220,4591,3221,4592, 763,4263,3842,4593, # 1792 -5203,5204,1958,1767,2956,3365,3712,1174, 452,1477,4594,3366,3155,5205,2838,1253, # 1808 -2387,2189,1091,2290,4264, 492,5206, 638,1169,1825,2136,1752,4014, 648, 926,1021, # 1824 -1324,4595, 520,4596, 997, 847,1007, 892,4597,3843,2267,1872,3713,2405,1785,4598, # 1840 -1953,2957,3103,3222,1728,4265,2044,3714,4599,2008,1701,3156,1551, 30,2268,4266, # 1856 -5207,2027,4600,3589,5208, 501,5209,4267, 594,3478,2166,1822,3590,3479,3591,3223, # 1872 - 829,2839,4268,5210,1680,3157,1225,4269,5211,3295,4601,4270,3158,2341,5212,4602, # 1888 -4271,5213,4015,4016,5214,1848,2388,2606,3367,5215,4603, 374,4017, 652,4272,4273, # 1904 - 375,1140, 798,5216,5217,5218,2366,4604,2269, 546,1659, 138,3051,2450,4605,5219, # 1920 -2254, 612,1849, 910, 796,3844,1740,1371, 825,3845,3846,5220,2920,2562,5221, 692, # 1936 - 444,3052,2634, 801,4606,4274,5222,1491, 244,1053,3053,4275,4276, 340,5223,4018, # 1952 -1041,3005, 293,1168, 87,1357,5224,1539, 959,5225,2240, 721, 694,4277,3847, 219, # 1968 -1478, 644,1417,3368,2666,1413,1401,1335,1389,4019,5226,5227,3006,2367,3159,1826, # 1984 - 730,1515, 184,2840, 66,4607,5228,1660,2958, 246,3369, 378,1457, 226,3480, 975, # 2000 -4020,2959,1264,3592, 674, 696,5229, 163,5230,1141,2422,2167, 713,3593,3370,4608, # 2016 -4021,5231,5232,1186, 15,5233,1079,1070,5234,1522,3224,3594, 276,1050,2725, 758, # 2032 -1126, 653,2960,3296,5235,2342, 889,3595,4022,3104,3007, 903,1250,4609,4023,3481, # 2048 -3596,1342,1681,1718, 766,3297, 286, 89,2961,3715,5236,1713,5237,2607,3371,3008, # 2064 -5238,2962,2219,3225,2880,5239,4610,2505,2533, 181, 387,1075,4024, 731,2190,3372, # 2080 -5240,3298, 310, 313,3482,2304, 770,4278, 54,3054, 189,4611,3105,3848,4025,5241, # 2096 -1230,1617,1850, 355,3597,4279,4612,3373, 111,4280,3716,1350,3160,3483,3055,4281, # 2112 -2150,3299,3598,5242,2797,4026,4027,3009, 722,2009,5243,1071, 247,1207,2343,2478, # 2128 -1378,4613,2010, 864,1437,1214,4614, 373,3849,1142,2220, 667,4615, 442,2763,2563, # 2144 -3850,4028,1969,4282,3300,1840, 837, 170,1107, 934,1336,1883,5244,5245,2119,4283, # 2160 -2841, 743,1569,5246,4616,4284, 582,2389,1418,3484,5247,1803,5248, 357,1395,1729, # 2176 -3717,3301,2423,1564,2241,5249,3106,3851,1633,4617,1114,2086,4285,1532,5250, 482, # 2192 -2451,4618,5251,5252,1492, 833,1466,5253,2726,3599,1641,2842,5254,1526,1272,3718, # 2208 -4286,1686,1795, 416,2564,1903,1954,1804,5255,3852,2798,3853,1159,2321,5256,2881, # 2224 -4619,1610,1584,3056,2424,2764, 443,3302,1163,3161,5257,5258,4029,5259,4287,2506, # 2240 -3057,4620,4030,3162,2104,1647,3600,2011,1873,4288,5260,4289, 431,3485,5261, 250, # 2256 - 97, 81,4290,5262,1648,1851,1558, 160, 848,5263, 866, 740,1694,5264,2204,2843, # 2272 -3226,4291,4621,3719,1687, 950,2479, 426, 469,3227,3720,3721,4031,5265,5266,1188, # 2288 - 424,1996, 861,3601,4292,3854,2205,2694, 168,1235,3602,4293,5267,2087,1674,4622, # 2304 -3374,3303, 220,2565,1009,5268,3855, 670,3010, 332,1208, 717,5269,5270,3603,2452, # 2320 -4032,3375,5271, 513,5272,1209,2882,3376,3163,4623,1080,5273,5274,5275,5276,2534, # 2336 -3722,3604, 815,1587,4033,4034,5277,3605,3486,3856,1254,4624,1328,3058,1390,4035, # 2352 -1741,4036,3857,4037,5278, 236,3858,2453,3304,5279,5280,3723,3859,1273,3860,4625, # 2368 -5281, 308,5282,4626, 245,4627,1852,2480,1307,2583, 430, 715,2137,2454,5283, 270, # 2384 - 199,2883,4038,5284,3606,2727,1753, 761,1754, 725,1661,1841,4628,3487,3724,5285, # 2400 -5286, 587, 14,3305, 227,2608, 326, 480,2270, 943,2765,3607, 291, 650,1884,5287, # 2416 -1702,1226, 102,1547, 62,3488, 904,4629,3489,1164,4294,5288,5289,1224,1548,2766, # 2432 - 391, 498,1493,5290,1386,1419,5291,2056,1177,4630, 813, 880,1081,2368, 566,1145, # 2448 -4631,2291,1001,1035,2566,2609,2242, 394,1286,5292,5293,2069,5294, 86,1494,1730, # 2464 -4039, 491,1588, 745, 897,2963, 843,3377,4040,2767,2884,3306,1768, 998,2221,2070, # 2480 - 397,1827,1195,1970,3725,3011,3378, 284,5295,3861,2507,2138,2120,1904,5296,4041, # 2496 -2151,4042,4295,1036,3490,1905, 114,2567,4296, 209,1527,5297,5298,2964,2844,2635, # 2512 -2390,2728,3164, 812,2568,5299,3307,5300,1559, 737,1885,3726,1210, 885, 28,2695, # 2528 -3608,3862,5301,4297,1004,1780,4632,5302, 346,1982,2222,2696,4633,3863,1742, 797, # 2544 -1642,4043,1934,1072,1384,2152, 896,4044,3308,3727,3228,2885,3609,5303,2569,1959, # 2560 -4634,2455,1786,5304,5305,5306,4045,4298,1005,1308,3728,4299,2729,4635,4636,1528, # 2576 -2610, 161,1178,4300,1983, 987,4637,1101,4301, 631,4046,1157,3229,2425,1343,1241, # 2592 -1016,2243,2570, 372, 877,2344,2508,1160, 555,1935, 911,4047,5307, 466,1170, 169, # 2608 -1051,2921,2697,3729,2481,3012,1182,2012,2571,1251,2636,5308, 992,2345,3491,1540, # 2624 -2730,1201,2071,2406,1997,2482,5309,4638, 528,1923,2191,1503,1874,1570,2369,3379, # 2640 -3309,5310, 557,1073,5311,1828,3492,2088,2271,3165,3059,3107, 767,3108,2799,4639, # 2656 -1006,4302,4640,2346,1267,2179,3730,3230, 778,4048,3231,2731,1597,2667,5312,4641, # 2672 -5313,3493,5314,5315,5316,3310,2698,1433,3311, 131, 95,1504,4049, 723,4303,3166, # 2688 -1842,3610,2768,2192,4050,2028,2105,3731,5317,3013,4051,1218,5318,3380,3232,4052, # 2704 -4304,2584, 248,1634,3864, 912,5319,2845,3732,3060,3865, 654, 53,5320,3014,5321, # 2720 -1688,4642, 777,3494,1032,4053,1425,5322, 191, 820,2121,2846, 971,4643, 931,3233, # 2736 - 135, 664, 783,3866,1998, 772,2922,1936,4054,3867,4644,2923,3234, 282,2732, 640, # 2752 -1372,3495,1127, 922, 325,3381,5323,5324, 711,2045,5325,5326,4055,2223,2800,1937, # 2768 -4056,3382,2224,2255,3868,2305,5327,4645,3869,1258,3312,4057,3235,2139,2965,4058, # 2784 -4059,5328,2225, 258,3236,4646, 101,1227,5329,3313,1755,5330,1391,3314,5331,2924, # 2800 -2057, 893,5332,5333,5334,1402,4305,2347,5335,5336,3237,3611,5337,5338, 878,1325, # 2816 -1781,2801,4647, 259,1385,2585, 744,1183,2272,4648,5339,4060,2509,5340, 684,1024, # 2832 -4306,5341, 472,3612,3496,1165,3315,4061,4062, 322,2153, 881, 455,1695,1152,1340, # 2848 - 660, 554,2154,4649,1058,4650,4307, 830,1065,3383,4063,4651,1924,5342,1703,1919, # 2864 -5343, 932,2273, 122,5344,4652, 947, 677,5345,3870,2637, 297,1906,1925,2274,4653, # 2880 -2322,3316,5346,5347,4308,5348,4309, 84,4310, 112, 989,5349, 547,1059,4064, 701, # 2896 -3613,1019,5350,4311,5351,3497, 942, 639, 457,2306,2456, 993,2966, 407, 851, 494, # 2912 -4654,3384, 927,5352,1237,5353,2426,3385, 573,4312, 680, 921,2925,1279,1875, 285, # 2928 - 790,1448,1984, 719,2168,5354,5355,4655,4065,4066,1649,5356,1541, 563,5357,1077, # 2944 -5358,3386,3061,3498, 511,3015,4067,4068,3733,4069,1268,2572,3387,3238,4656,4657, # 2960 -5359, 535,1048,1276,1189,2926,2029,3167,1438,1373,2847,2967,1134,2013,5360,4313, # 2976 -1238,2586,3109,1259,5361, 700,5362,2968,3168,3734,4314,5363,4315,1146,1876,1907, # 2992 -4658,2611,4070, 781,2427, 132,1589, 203, 147, 273,2802,2407, 898,1787,2155,4071, # 3008 -4072,5364,3871,2803,5365,5366,4659,4660,5367,3239,5368,1635,3872, 965,5369,1805, # 3024 -2699,1516,3614,1121,1082,1329,3317,4073,1449,3873, 65,1128,2848,2927,2769,1590, # 3040 -3874,5370,5371, 12,2668, 45, 976,2587,3169,4661, 517,2535,1013,1037,3240,5372, # 3056 -3875,2849,5373,3876,5374,3499,5375,2612, 614,1999,2323,3877,3110,2733,2638,5376, # 3072 -2588,4316, 599,1269,5377,1811,3735,5378,2700,3111, 759,1060, 489,1806,3388,3318, # 3088 -1358,5379,5380,2391,1387,1215,2639,2256, 490,5381,5382,4317,1759,2392,2348,5383, # 3104 -4662,3878,1908,4074,2640,1807,3241,4663,3500,3319,2770,2349, 874,5384,5385,3501, # 3120 -3736,1859, 91,2928,3737,3062,3879,4664,5386,3170,4075,2669,5387,3502,1202,1403, # 3136 -3880,2969,2536,1517,2510,4665,3503,2511,5388,4666,5389,2701,1886,1495,1731,4076, # 3152 -2370,4667,5390,2030,5391,5392,4077,2702,1216, 237,2589,4318,2324,4078,3881,4668, # 3168 -4669,2703,3615,3504, 445,4670,5393,5394,5395,5396,2771, 61,4079,3738,1823,4080, # 3184 -5397, 687,2046, 935, 925, 405,2670, 703,1096,1860,2734,4671,4081,1877,1367,2704, # 3200 -3389, 918,2106,1782,2483, 334,3320,1611,1093,4672, 564,3171,3505,3739,3390, 945, # 3216 -2641,2058,4673,5398,1926, 872,4319,5399,3506,2705,3112, 349,4320,3740,4082,4674, # 3232 -3882,4321,3741,2156,4083,4675,4676,4322,4677,2408,2047, 782,4084, 400, 251,4323, # 3248 -1624,5400,5401, 277,3742, 299,1265, 476,1191,3883,2122,4324,4325,1109, 205,5402, # 3264 -2590,1000,2157,3616,1861,5403,5404,5405,4678,5406,4679,2573, 107,2484,2158,4085, # 3280 -3507,3172,5407,1533, 541,1301, 158, 753,4326,2886,3617,5408,1696, 370,1088,4327, # 3296 -4680,3618, 579, 327, 440, 162,2244, 269,1938,1374,3508, 968,3063, 56,1396,3113, # 3312 -2107,3321,3391,5409,1927,2159,4681,3016,5410,3619,5411,5412,3743,4682,2485,5413, # 3328 -2804,5414,1650,4683,5415,2613,5416,5417,4086,2671,3392,1149,3393,4087,3884,4088, # 3344 -5418,1076, 49,5419, 951,3242,3322,3323, 450,2850, 920,5420,1812,2805,2371,4328, # 3360 -1909,1138,2372,3885,3509,5421,3243,4684,1910,1147,1518,2428,4685,3886,5422,4686, # 3376 -2393,2614, 260,1796,3244,5423,5424,3887,3324, 708,5425,3620,1704,5426,3621,1351, # 3392 -1618,3394,3017,1887, 944,4329,3395,4330,3064,3396,4331,5427,3744, 422, 413,1714, # 3408 -3325, 500,2059,2350,4332,2486,5428,1344,1911, 954,5429,1668,5430,5431,4089,2409, # 3424 -4333,3622,3888,4334,5432,2307,1318,2512,3114, 133,3115,2887,4687, 629, 31,2851, # 3440 -2706,3889,4688, 850, 949,4689,4090,2970,1732,2089,4335,1496,1853,5433,4091, 620, # 3456 -3245, 981,1242,3745,3397,1619,3746,1643,3326,2140,2457,1971,1719,3510,2169,5434, # 3472 -3246,5435,5436,3398,1829,5437,1277,4690,1565,2048,5438,1636,3623,3116,5439, 869, # 3488 -2852, 655,3890,3891,3117,4092,3018,3892,1310,3624,4691,5440,5441,5442,1733, 558, # 3504 -4692,3747, 335,1549,3065,1756,4336,3748,1946,3511,1830,1291,1192, 470,2735,2108, # 3520 -2806, 913,1054,4093,5443,1027,5444,3066,4094,4693, 982,2672,3399,3173,3512,3247, # 3536 -3248,1947,2807,5445, 571,4694,5446,1831,5447,3625,2591,1523,2429,5448,2090, 984, # 3552 -4695,3749,1960,5449,3750, 852, 923,2808,3513,3751, 969,1519, 999,2049,2325,1705, # 3568 -5450,3118, 615,1662, 151, 597,4095,2410,2326,1049, 275,4696,3752,4337, 568,3753, # 3584 -3626,2487,4338,3754,5451,2430,2275, 409,3249,5452,1566,2888,3514,1002, 769,2853, # 3600 - 194,2091,3174,3755,2226,3327,4339, 628,1505,5453,5454,1763,2180,3019,4096, 521, # 3616 -1161,2592,1788,2206,2411,4697,4097,1625,4340,4341, 412, 42,3119, 464,5455,2642, # 3632 -4698,3400,1760,1571,2889,3515,2537,1219,2207,3893,2643,2141,2373,4699,4700,3328, # 3648 -1651,3401,3627,5456,5457,3628,2488,3516,5458,3756,5459,5460,2276,2092, 460,5461, # 3664 -4701,5462,3020, 962, 588,3629, 289,3250,2644,1116, 52,5463,3067,1797,5464,5465, # 3680 -5466,1467,5467,1598,1143,3757,4342,1985,1734,1067,4702,1280,3402, 465,4703,1572, # 3696 - 510,5468,1928,2245,1813,1644,3630,5469,4704,3758,5470,5471,2673,1573,1534,5472, # 3712 -5473, 536,1808,1761,3517,3894,3175,2645,5474,5475,5476,4705,3518,2929,1912,2809, # 3728 -5477,3329,1122, 377,3251,5478, 360,5479,5480,4343,1529, 551,5481,2060,3759,1769, # 3744 -2431,5482,2930,4344,3330,3120,2327,2109,2031,4706,1404, 136,1468,1479, 672,1171, # 3760 -3252,2308, 271,3176,5483,2772,5484,2050, 678,2736, 865,1948,4707,5485,2014,4098, # 3776 -2971,5486,2737,2227,1397,3068,3760,4708,4709,1735,2931,3403,3631,5487,3895, 509, # 3792 -2854,2458,2890,3896,5488,5489,3177,3178,4710,4345,2538,4711,2309,1166,1010, 552, # 3808 - 681,1888,5490,5491,2972,2973,4099,1287,1596,1862,3179, 358, 453, 736, 175, 478, # 3824 -1117, 905,1167,1097,5492,1854,1530,5493,1706,5494,2181,3519,2292,3761,3520,3632, # 3840 -4346,2093,4347,5495,3404,1193,2489,4348,1458,2193,2208,1863,1889,1421,3331,2932, # 3856 -3069,2182,3521, 595,2123,5496,4100,5497,5498,4349,1707,2646, 223,3762,1359, 751, # 3872 -3121, 183,3522,5499,2810,3021, 419,2374, 633, 704,3897,2394, 241,5500,5501,5502, # 3888 - 838,3022,3763,2277,2773,2459,3898,1939,2051,4101,1309,3122,2246,1181,5503,1136, # 3904 -2209,3899,2375,1446,4350,2310,4712,5504,5505,4351,1055,2615, 484,3764,5506,4102, # 3920 - 625,4352,2278,3405,1499,4353,4103,5507,4104,4354,3253,2279,2280,3523,5508,5509, # 3936 -2774, 808,2616,3765,3406,4105,4355,3123,2539, 526,3407,3900,4356, 955,5510,1620, # 3952 -4357,2647,2432,5511,1429,3766,1669,1832, 994, 928,5512,3633,1260,5513,5514,5515, # 3968 -1949,2293, 741,2933,1626,4358,2738,2460, 867,1184, 362,3408,1392,5516,5517,4106, # 3984 -4359,1770,1736,3254,2934,4713,4714,1929,2707,1459,1158,5518,3070,3409,2891,1292, # 4000 -1930,2513,2855,3767,1986,1187,2072,2015,2617,4360,5519,2574,2514,2170,3768,2490, # 4016 -3332,5520,3769,4715,5521,5522, 666,1003,3023,1022,3634,4361,5523,4716,1814,2257, # 4032 - 574,3901,1603, 295,1535, 705,3902,4362, 283, 858, 417,5524,5525,3255,4717,4718, # 4048 -3071,1220,1890,1046,2281,2461,4107,1393,1599, 689,2575, 388,4363,5526,2491, 802, # 4064 -5527,2811,3903,2061,1405,2258,5528,4719,3904,2110,1052,1345,3256,1585,5529, 809, # 4080 -5530,5531,5532, 575,2739,3524, 956,1552,1469,1144,2328,5533,2329,1560,2462,3635, # 4096 -3257,4108, 616,2210,4364,3180,2183,2294,5534,1833,5535,3525,4720,5536,1319,3770, # 4112 -3771,1211,3636,1023,3258,1293,2812,5537,5538,5539,3905, 607,2311,3906, 762,2892, # 4128 -1439,4365,1360,4721,1485,3072,5540,4722,1038,4366,1450,2062,2648,4367,1379,4723, # 4144 -2593,5541,5542,4368,1352,1414,2330,2935,1172,5543,5544,3907,3908,4724,1798,1451, # 4160 -5545,5546,5547,5548,2936,4109,4110,2492,2351, 411,4111,4112,3637,3333,3124,4725, # 4176 -1561,2674,1452,4113,1375,5549,5550, 47,2974, 316,5551,1406,1591,2937,3181,5552, # 4192 -1025,2142,3125,3182, 354,2740, 884,2228,4369,2412, 508,3772, 726,3638, 996,2433, # 4208 -3639, 729,5553, 392,2194,1453,4114,4726,3773,5554,5555,2463,3640,2618,1675,2813, # 4224 - 919,2352,2975,2353,1270,4727,4115, 73,5556,5557, 647,5558,3259,2856,2259,1550, # 4240 -1346,3024,5559,1332, 883,3526,5560,5561,5562,5563,3334,2775,5564,1212, 831,1347, # 4256 -4370,4728,2331,3909,1864,3073, 720,3910,4729,4730,3911,5565,4371,5566,5567,4731, # 4272 -5568,5569,1799,4732,3774,2619,4733,3641,1645,2376,4734,5570,2938, 669,2211,2675, # 4288 -2434,5571,2893,5572,5573,1028,3260,5574,4372,2413,5575,2260,1353,5576,5577,4735, # 4304 -3183, 518,5578,4116,5579,4373,1961,5580,2143,4374,5581,5582,3025,2354,2355,3912, # 4320 - 516,1834,1454,4117,2708,4375,4736,2229,2620,1972,1129,3642,5583,2776,5584,2976, # 4336 -1422, 577,1470,3026,1524,3410,5585,5586, 432,4376,3074,3527,5587,2594,1455,2515, # 4352 -2230,1973,1175,5588,1020,2741,4118,3528,4737,5589,2742,5590,1743,1361,3075,3529, # 4368 -2649,4119,4377,4738,2295, 895, 924,4378,2171, 331,2247,3076, 166,1627,3077,1098, # 4384 -5591,1232,2894,2231,3411,4739, 657, 403,1196,2377, 542,3775,3412,1600,4379,3530, # 4400 -5592,4740,2777,3261, 576, 530,1362,4741,4742,2540,2676,3776,4120,5593, 842,3913, # 4416 -5594,2814,2032,1014,4121, 213,2709,3413, 665, 621,4380,5595,3777,2939,2435,5596, # 4432 -2436,3335,3643,3414,4743,4381,2541,4382,4744,3644,1682,4383,3531,1380,5597, 724, # 4448 -2282, 600,1670,5598,1337,1233,4745,3126,2248,5599,1621,4746,5600, 651,4384,5601, # 4464 -1612,4385,2621,5602,2857,5603,2743,2312,3078,5604, 716,2464,3079, 174,1255,2710, # 4480 -4122,3645, 548,1320,1398, 728,4123,1574,5605,1891,1197,3080,4124,5606,3081,3082, # 4496 -3778,3646,3779, 747,5607, 635,4386,4747,5608,5609,5610,4387,5611,5612,4748,5613, # 4512 -3415,4749,2437, 451,5614,3780,2542,2073,4388,2744,4389,4125,5615,1764,4750,5616, # 4528 -4390, 350,4751,2283,2395,2493,5617,4391,4126,2249,1434,4127, 488,4752, 458,4392, # 4544 -4128,3781, 771,1330,2396,3914,2576,3184,2160,2414,1553,2677,3185,4393,5618,2494, # 4560 -2895,2622,1720,2711,4394,3416,4753,5619,2543,4395,5620,3262,4396,2778,5621,2016, # 4576 -2745,5622,1155,1017,3782,3915,5623,3336,2313, 201,1865,4397,1430,5624,4129,5625, # 4592 -5626,5627,5628,5629,4398,1604,5630, 414,1866, 371,2595,4754,4755,3532,2017,3127, # 4608 -4756,1708, 960,4399, 887, 389,2172,1536,1663,1721,5631,2232,4130,2356,2940,1580, # 4624 -5632,5633,1744,4757,2544,4758,4759,5634,4760,5635,2074,5636,4761,3647,3417,2896, # 4640 -4400,5637,4401,2650,3418,2815, 673,2712,2465, 709,3533,4131,3648,4402,5638,1148, # 4656 - 502, 634,5639,5640,1204,4762,3649,1575,4763,2623,3783,5641,3784,3128, 948,3263, # 4672 - 121,1745,3916,1110,5642,4403,3083,2516,3027,4132,3785,1151,1771,3917,1488,4133, # 4688 -1987,5643,2438,3534,5644,5645,2094,5646,4404,3918,1213,1407,2816, 531,2746,2545, # 4704 -3264,1011,1537,4764,2779,4405,3129,1061,5647,3786,3787,1867,2897,5648,2018, 120, # 4720 -4406,4407,2063,3650,3265,2314,3919,2678,3419,1955,4765,4134,5649,3535,1047,2713, # 4736 -1266,5650,1368,4766,2858, 649,3420,3920,2546,2747,1102,2859,2679,5651,5652,2000, # 4752 -5653,1111,3651,2977,5654,2495,3921,3652,2817,1855,3421,3788,5655,5656,3422,2415, # 4768 -2898,3337,3266,3653,5657,2577,5658,3654,2818,4135,1460, 856,5659,3655,5660,2899, # 4784 -2978,5661,2900,3922,5662,4408, 632,2517, 875,3923,1697,3924,2296,5663,5664,4767, # 4800 -3028,1239, 580,4768,4409,5665, 914, 936,2075,1190,4136,1039,2124,5666,5667,5668, # 4816 -5669,3423,1473,5670,1354,4410,3925,4769,2173,3084,4137, 915,3338,4411,4412,3339, # 4832 -1605,1835,5671,2748, 398,3656,4413,3926,4138, 328,1913,2860,4139,3927,1331,4414, # 4848 -3029, 937,4415,5672,3657,4140,4141,3424,2161,4770,3425, 524, 742, 538,3085,1012, # 4864 -5673,5674,3928,2466,5675, 658,1103, 225,3929,5676,5677,4771,5678,4772,5679,3267, # 4880 -1243,5680,4142, 963,2250,4773,5681,2714,3658,3186,5682,5683,2596,2332,5684,4774, # 4896 -5685,5686,5687,3536, 957,3426,2547,2033,1931,2941,2467, 870,2019,3659,1746,2780, # 4912 -2781,2439,2468,5688,3930,5689,3789,3130,3790,3537,3427,3791,5690,1179,3086,5691, # 4928 -3187,2378,4416,3792,2548,3188,3131,2749,4143,5692,3428,1556,2549,2297, 977,2901, # 4944 -2034,4144,1205,3429,5693,1765,3430,3189,2125,1271, 714,1689,4775,3538,5694,2333, # 4960 -3931, 533,4417,3660,2184, 617,5695,2469,3340,3539,2315,5696,5697,3190,5698,5699, # 4976 -3932,1988, 618, 427,2651,3540,3431,5700,5701,1244,1690,5702,2819,4418,4776,5703, # 4992 -3541,4777,5704,2284,1576, 473,3661,4419,3432, 972,5705,3662,5706,3087,5707,5708, # 5008 -4778,4779,5709,3793,4145,4146,5710, 153,4780, 356,5711,1892,2902,4420,2144, 408, # 5024 - 803,2357,5712,3933,5713,4421,1646,2578,2518,4781,4782,3934,5714,3935,4422,5715, # 5040 -2416,3433, 752,5716,5717,1962,3341,2979,5718, 746,3030,2470,4783,4423,3794, 698, # 5056 -4784,1893,4424,3663,2550,4785,3664,3936,5719,3191,3434,5720,1824,1302,4147,2715, # 5072 -3937,1974,4425,5721,4426,3192, 823,1303,1288,1236,2861,3542,4148,3435, 774,3938, # 5088 -5722,1581,4786,1304,2862,3939,4787,5723,2440,2162,1083,3268,4427,4149,4428, 344, # 5104 -1173, 288,2316, 454,1683,5724,5725,1461,4788,4150,2597,5726,5727,4789, 985, 894, # 5120 -5728,3436,3193,5729,1914,2942,3795,1989,5730,2111,1975,5731,4151,5732,2579,1194, # 5136 - 425,5733,4790,3194,1245,3796,4429,5734,5735,2863,5736, 636,4791,1856,3940, 760, # 5152 -1800,5737,4430,2212,1508,4792,4152,1894,1684,2298,5738,5739,4793,4431,4432,2213, # 5168 - 479,5740,5741, 832,5742,4153,2496,5743,2980,2497,3797, 990,3132, 627,1815,2652, # 5184 -4433,1582,4434,2126,2112,3543,4794,5744, 799,4435,3195,5745,4795,2113,1737,3031, # 5200 -1018, 543, 754,4436,3342,1676,4796,4797,4154,4798,1489,5746,3544,5747,2624,2903, # 5216 -4155,5748,5749,2981,5750,5751,5752,5753,3196,4799,4800,2185,1722,5754,3269,3270, # 5232 -1843,3665,1715, 481, 365,1976,1857,5755,5756,1963,2498,4801,5757,2127,3666,3271, # 5248 - 433,1895,2064,2076,5758, 602,2750,5759,5760,5761,5762,5763,3032,1628,3437,5764, # 5264 -3197,4802,4156,2904,4803,2519,5765,2551,2782,5766,5767,5768,3343,4804,2905,5769, # 5280 -4805,5770,2864,4806,4807,1221,2982,4157,2520,5771,5772,5773,1868,1990,5774,5775, # 5296 -5776,1896,5777,5778,4808,1897,4158, 318,5779,2095,4159,4437,5780,5781, 485,5782, # 5312 - 938,3941, 553,2680, 116,5783,3942,3667,5784,3545,2681,2783,3438,3344,2820,5785, # 5328 -3668,2943,4160,1747,2944,2983,5786,5787, 207,5788,4809,5789,4810,2521,5790,3033, # 5344 - 890,3669,3943,5791,1878,3798,3439,5792,2186,2358,3440,1652,5793,5794,5795, 941, # 5360 -2299, 208,3546,4161,2020, 330,4438,3944,2906,2499,3799,4439,4811,5796,5797,5798, # 5376 -) -# fmt: on diff --git a/venv/lib/python3.10/site-packages/chardet/big5prober.py b/venv/lib/python3.10/site-packages/chardet/big5prober.py deleted file mode 100644 index ef09c60e327a0122e32f95f2f10a826a033c573c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/big5prober.py +++ /dev/null @@ -1,47 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Communicator client code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from .chardistribution import Big5DistributionAnalysis -from .codingstatemachine import CodingStateMachine -from .mbcharsetprober import MultiByteCharSetProber -from .mbcssm import BIG5_SM_MODEL - - -class Big5Prober(MultiByteCharSetProber): - def __init__(self) -> None: - super().__init__() - self.coding_sm = CodingStateMachine(BIG5_SM_MODEL) - self.distribution_analyzer = Big5DistributionAnalysis() - self.reset() - - @property - def charset_name(self) -> str: - return "Big5" - - @property - def language(self) -> str: - return "Chinese" diff --git a/venv/lib/python3.10/site-packages/chardet/chardistribution.py b/venv/lib/python3.10/site-packages/chardet/chardistribution.py deleted file mode 100644 index 176cb996408e6681a88722783919efc0e9dafb29..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/chardistribution.py +++ /dev/null @@ -1,261 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Communicator client code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from typing import Tuple, Union - -from .big5freq import ( - BIG5_CHAR_TO_FREQ_ORDER, - BIG5_TABLE_SIZE, - BIG5_TYPICAL_DISTRIBUTION_RATIO, -) -from .euckrfreq import ( - EUCKR_CHAR_TO_FREQ_ORDER, - EUCKR_TABLE_SIZE, - EUCKR_TYPICAL_DISTRIBUTION_RATIO, -) -from .euctwfreq import ( - EUCTW_CHAR_TO_FREQ_ORDER, - EUCTW_TABLE_SIZE, - EUCTW_TYPICAL_DISTRIBUTION_RATIO, -) -from .gb2312freq import ( - GB2312_CHAR_TO_FREQ_ORDER, - GB2312_TABLE_SIZE, - GB2312_TYPICAL_DISTRIBUTION_RATIO, -) -from .jisfreq import ( - JIS_CHAR_TO_FREQ_ORDER, - JIS_TABLE_SIZE, - JIS_TYPICAL_DISTRIBUTION_RATIO, -) -from .johabfreq import JOHAB_TO_EUCKR_ORDER_TABLE - - -class CharDistributionAnalysis: - ENOUGH_DATA_THRESHOLD = 1024 - SURE_YES = 0.99 - SURE_NO = 0.01 - MINIMUM_DATA_THRESHOLD = 3 - - def __init__(self) -> None: - # Mapping table to get frequency order from char order (get from - # GetOrder()) - self._char_to_freq_order: Tuple[int, ...] = tuple() - self._table_size = 0 # Size of above table - # This is a constant value which varies from language to language, - # used in calculating confidence. See - # http://www.mozilla.org/projects/intl/UniversalCharsetDetection.html - # for further detail. - self.typical_distribution_ratio = 0.0 - self._done = False - self._total_chars = 0 - self._freq_chars = 0 - self.reset() - - def reset(self) -> None: - """reset analyser, clear any state""" - # If this flag is set to True, detection is done and conclusion has - # been made - self._done = False - self._total_chars = 0 # Total characters encountered - # The number of characters whose frequency order is less than 512 - self._freq_chars = 0 - - def feed(self, char: Union[bytes, bytearray], char_len: int) -> None: - """feed a character with known length""" - if char_len == 2: - # we only care about 2-bytes character in our distribution analysis - order = self.get_order(char) - else: - order = -1 - if order >= 0: - self._total_chars += 1 - # order is valid - if order < self._table_size: - if 512 > self._char_to_freq_order[order]: - self._freq_chars += 1 - - def get_confidence(self) -> float: - """return confidence based on existing data""" - # if we didn't receive any character in our consideration range, - # return negative answer - if self._total_chars <= 0 or self._freq_chars <= self.MINIMUM_DATA_THRESHOLD: - return self.SURE_NO - - if self._total_chars != self._freq_chars: - r = self._freq_chars / ( - (self._total_chars - self._freq_chars) * self.typical_distribution_ratio - ) - if r < self.SURE_YES: - return r - - # normalize confidence (we don't want to be 100% sure) - return self.SURE_YES - - def got_enough_data(self) -> bool: - # It is not necessary to receive all data to draw conclusion. - # For charset detection, certain amount of data is enough - return self._total_chars > self.ENOUGH_DATA_THRESHOLD - - def get_order(self, _: Union[bytes, bytearray]) -> int: - # We do not handle characters based on the original encoding string, - # but convert this encoding string to a number, here called order. - # This allows multiple encodings of a language to share one frequency - # table. - return -1 - - -class EUCTWDistributionAnalysis(CharDistributionAnalysis): - def __init__(self) -> None: - super().__init__() - self._char_to_freq_order = EUCTW_CHAR_TO_FREQ_ORDER - self._table_size = EUCTW_TABLE_SIZE - self.typical_distribution_ratio = EUCTW_TYPICAL_DISTRIBUTION_RATIO - - def get_order(self, byte_str: Union[bytes, bytearray]) -> int: - # for euc-TW encoding, we are interested - # first byte range: 0xc4 -- 0xfe - # second byte range: 0xa1 -- 0xfe - # no validation needed here. State machine has done that - first_char = byte_str[0] - if first_char >= 0xC4: - return 94 * (first_char - 0xC4) + byte_str[1] - 0xA1 - return -1 - - -class EUCKRDistributionAnalysis(CharDistributionAnalysis): - def __init__(self) -> None: - super().__init__() - self._char_to_freq_order = EUCKR_CHAR_TO_FREQ_ORDER - self._table_size = EUCKR_TABLE_SIZE - self.typical_distribution_ratio = EUCKR_TYPICAL_DISTRIBUTION_RATIO - - def get_order(self, byte_str: Union[bytes, bytearray]) -> int: - # for euc-KR encoding, we are interested - # first byte range: 0xb0 -- 0xfe - # second byte range: 0xa1 -- 0xfe - # no validation needed here. State machine has done that - first_char = byte_str[0] - if first_char >= 0xB0: - return 94 * (first_char - 0xB0) + byte_str[1] - 0xA1 - return -1 - - -class JOHABDistributionAnalysis(CharDistributionAnalysis): - def __init__(self) -> None: - super().__init__() - self._char_to_freq_order = EUCKR_CHAR_TO_FREQ_ORDER - self._table_size = EUCKR_TABLE_SIZE - self.typical_distribution_ratio = EUCKR_TYPICAL_DISTRIBUTION_RATIO - - def get_order(self, byte_str: Union[bytes, bytearray]) -> int: - first_char = byte_str[0] - if 0x88 <= first_char < 0xD4: - code = first_char * 256 + byte_str[1] - return JOHAB_TO_EUCKR_ORDER_TABLE.get(code, -1) - return -1 - - -class GB2312DistributionAnalysis(CharDistributionAnalysis): - def __init__(self) -> None: - super().__init__() - self._char_to_freq_order = GB2312_CHAR_TO_FREQ_ORDER - self._table_size = GB2312_TABLE_SIZE - self.typical_distribution_ratio = GB2312_TYPICAL_DISTRIBUTION_RATIO - - def get_order(self, byte_str: Union[bytes, bytearray]) -> int: - # for GB2312 encoding, we are interested - # first byte range: 0xb0 -- 0xfe - # second byte range: 0xa1 -- 0xfe - # no validation needed here. State machine has done that - first_char, second_char = byte_str[0], byte_str[1] - if (first_char >= 0xB0) and (second_char >= 0xA1): - return 94 * (first_char - 0xB0) + second_char - 0xA1 - return -1 - - -class Big5DistributionAnalysis(CharDistributionAnalysis): - def __init__(self) -> None: - super().__init__() - self._char_to_freq_order = BIG5_CHAR_TO_FREQ_ORDER - self._table_size = BIG5_TABLE_SIZE - self.typical_distribution_ratio = BIG5_TYPICAL_DISTRIBUTION_RATIO - - def get_order(self, byte_str: Union[bytes, bytearray]) -> int: - # for big5 encoding, we are interested - # first byte range: 0xa4 -- 0xfe - # second byte range: 0x40 -- 0x7e , 0xa1 -- 0xfe - # no validation needed here. State machine has done that - first_char, second_char = byte_str[0], byte_str[1] - if first_char >= 0xA4: - if second_char >= 0xA1: - return 157 * (first_char - 0xA4) + second_char - 0xA1 + 63 - return 157 * (first_char - 0xA4) + second_char - 0x40 - return -1 - - -class SJISDistributionAnalysis(CharDistributionAnalysis): - def __init__(self) -> None: - super().__init__() - self._char_to_freq_order = JIS_CHAR_TO_FREQ_ORDER - self._table_size = JIS_TABLE_SIZE - self.typical_distribution_ratio = JIS_TYPICAL_DISTRIBUTION_RATIO - - def get_order(self, byte_str: Union[bytes, bytearray]) -> int: - # for sjis encoding, we are interested - # first byte range: 0x81 -- 0x9f , 0xe0 -- 0xfe - # second byte range: 0x40 -- 0x7e, 0x81 -- oxfe - # no validation needed here. State machine has done that - first_char, second_char = byte_str[0], byte_str[1] - if 0x81 <= first_char <= 0x9F: - order = 188 * (first_char - 0x81) - elif 0xE0 <= first_char <= 0xEF: - order = 188 * (first_char - 0xE0 + 31) - else: - return -1 - order = order + second_char - 0x40 - if second_char > 0x7F: - order = -1 - return order - - -class EUCJPDistributionAnalysis(CharDistributionAnalysis): - def __init__(self) -> None: - super().__init__() - self._char_to_freq_order = JIS_CHAR_TO_FREQ_ORDER - self._table_size = JIS_TABLE_SIZE - self.typical_distribution_ratio = JIS_TYPICAL_DISTRIBUTION_RATIO - - def get_order(self, byte_str: Union[bytes, bytearray]) -> int: - # for euc-JP encoding, we are interested - # first byte range: 0xa0 -- 0xfe - # second byte range: 0xa1 -- 0xfe - # no validation needed here. State machine has done that - char = byte_str[0] - if char >= 0xA0: - return 94 * (char - 0xA1) + byte_str[1] - 0xA1 - return -1 diff --git a/venv/lib/python3.10/site-packages/chardet/charsetgroupprober.py b/venv/lib/python3.10/site-packages/chardet/charsetgroupprober.py deleted file mode 100644 index 6def56b4a75f67000ed8181ae2d2c40eefb645fb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/charsetgroupprober.py +++ /dev/null @@ -1,106 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Communicator client code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from typing import List, Optional, Union - -from .charsetprober import CharSetProber -from .enums import LanguageFilter, ProbingState - - -class CharSetGroupProber(CharSetProber): - def __init__(self, lang_filter: LanguageFilter = LanguageFilter.NONE) -> None: - super().__init__(lang_filter=lang_filter) - self._active_num = 0 - self.probers: List[CharSetProber] = [] - self._best_guess_prober: Optional[CharSetProber] = None - - def reset(self) -> None: - super().reset() - self._active_num = 0 - for prober in self.probers: - prober.reset() - prober.active = True - self._active_num += 1 - self._best_guess_prober = None - - @property - def charset_name(self) -> Optional[str]: - if not self._best_guess_prober: - self.get_confidence() - if not self._best_guess_prober: - return None - return self._best_guess_prober.charset_name - - @property - def language(self) -> Optional[str]: - if not self._best_guess_prober: - self.get_confidence() - if not self._best_guess_prober: - return None - return self._best_guess_prober.language - - def feed(self, byte_str: Union[bytes, bytearray]) -> ProbingState: - for prober in self.probers: - if not prober.active: - continue - state = prober.feed(byte_str) - if not state: - continue - if state == ProbingState.FOUND_IT: - self._best_guess_prober = prober - self._state = ProbingState.FOUND_IT - return self.state - if state == ProbingState.NOT_ME: - prober.active = False - self._active_num -= 1 - if self._active_num <= 0: - self._state = ProbingState.NOT_ME - return self.state - return self.state - - def get_confidence(self) -> float: - state = self.state - if state == ProbingState.FOUND_IT: - return 0.99 - if state == ProbingState.NOT_ME: - return 0.01 - best_conf = 0.0 - self._best_guess_prober = None - for prober in self.probers: - if not prober.active: - self.logger.debug("%s not active", prober.charset_name) - continue - conf = prober.get_confidence() - self.logger.debug( - "%s %s confidence = %s", prober.charset_name, prober.language, conf - ) - if best_conf < conf: - best_conf = conf - self._best_guess_prober = prober - if not self._best_guess_prober: - return 0.0 - return best_conf diff --git a/venv/lib/python3.10/site-packages/chardet/charsetprober.py b/venv/lib/python3.10/site-packages/chardet/charsetprober.py deleted file mode 100644 index a103ca11356606402c03b320a4fcdb8635051623..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/charsetprober.py +++ /dev/null @@ -1,147 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Universal charset detector code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 2001 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# Shy Shalom - original C code -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -import logging -import re -from typing import Optional, Union - -from .enums import LanguageFilter, ProbingState - -INTERNATIONAL_WORDS_PATTERN = re.compile( - b"[a-zA-Z]*[\x80-\xFF]+[a-zA-Z]*[^a-zA-Z\x80-\xFF]?" -) - - -class CharSetProber: - - SHORTCUT_THRESHOLD = 0.95 - - def __init__(self, lang_filter: LanguageFilter = LanguageFilter.NONE) -> None: - self._state = ProbingState.DETECTING - self.active = True - self.lang_filter = lang_filter - self.logger = logging.getLogger(__name__) - - def reset(self) -> None: - self._state = ProbingState.DETECTING - - @property - def charset_name(self) -> Optional[str]: - return None - - @property - def language(self) -> Optional[str]: - raise NotImplementedError - - def feed(self, byte_str: Union[bytes, bytearray]) -> ProbingState: - raise NotImplementedError - - @property - def state(self) -> ProbingState: - return self._state - - def get_confidence(self) -> float: - return 0.0 - - @staticmethod - def filter_high_byte_only(buf: Union[bytes, bytearray]) -> bytes: - buf = re.sub(b"([\x00-\x7F])+", b" ", buf) - return buf - - @staticmethod - def filter_international_words(buf: Union[bytes, bytearray]) -> bytearray: - """ - We define three types of bytes: - alphabet: english alphabets [a-zA-Z] - international: international characters [\x80-\xFF] - marker: everything else [^a-zA-Z\x80-\xFF] - The input buffer can be thought to contain a series of words delimited - by markers. This function works to filter all words that contain at - least one international character. All contiguous sequences of markers - are replaced by a single space ascii character. - This filter applies to all scripts which do not use English characters. - """ - filtered = bytearray() - - # This regex expression filters out only words that have at-least one - # international character. The word may include one marker character at - # the end. - words = INTERNATIONAL_WORDS_PATTERN.findall(buf) - - for word in words: - filtered.extend(word[:-1]) - - # If the last character in the word is a marker, replace it with a - # space as markers shouldn't affect our analysis (they are used - # similarly across all languages and may thus have similar - # frequencies). - last_char = word[-1:] - if not last_char.isalpha() and last_char < b"\x80": - last_char = b" " - filtered.extend(last_char) - - return filtered - - @staticmethod - def remove_xml_tags(buf: Union[bytes, bytearray]) -> bytes: - """ - Returns a copy of ``buf`` that retains only the sequences of English - alphabet and high byte characters that are not between <> characters. - This filter can be applied to all scripts which contain both English - characters and extended ASCII characters, but is currently only used by - ``Latin1Prober``. - """ - filtered = bytearray() - in_tag = False - prev = 0 - buf = memoryview(buf).cast("c") - - for curr, buf_char in enumerate(buf): - # Check if we're coming out of or entering an XML tag - - # https://github.com/python/typeshed/issues/8182 - if buf_char == b">": # type: ignore[comparison-overlap] - prev = curr + 1 - in_tag = False - # https://github.com/python/typeshed/issues/8182 - elif buf_char == b"<": # type: ignore[comparison-overlap] - if curr > prev and not in_tag: - # Keep everything after last non-extended-ASCII, - # non-alphabetic character - filtered.extend(buf[prev:curr]) - # Output a space to delimit stretch we kept - filtered.extend(b" ") - in_tag = True - - # If we're not in a tag... - if not in_tag: - # Keep everything after last non-extended-ASCII, non-alphabetic - # character - filtered.extend(buf[prev:]) - - return filtered diff --git a/venv/lib/python3.10/site-packages/chardet/cli/__init__.py b/venv/lib/python3.10/site-packages/chardet/cli/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/chardet/cli/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/cli/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 35c51973dc8f510981efb7307d4c8d158f6d5d33..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/cli/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/cli/__pycache__/chardetect.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/cli/__pycache__/chardetect.cpython-310.pyc deleted file mode 100644 index 91c0d1acae75a3df1ca46d3fda85d888ddcf5f16..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/cli/__pycache__/chardetect.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/cli/chardetect.py b/venv/lib/python3.10/site-packages/chardet/cli/chardetect.py deleted file mode 100644 index 43f6e144f677a113b5362dcbdfb75db4f41c2b2f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/cli/chardetect.py +++ /dev/null @@ -1,112 +0,0 @@ -""" -Script which takes one or more file paths and reports on their detected -encodings - -Example:: - - % chardetect somefile someotherfile - somefile: windows-1252 with confidence 0.5 - someotherfile: ascii with confidence 1.0 - -If no paths are provided, it takes its input from stdin. - -""" - - -import argparse -import sys -from typing import Iterable, List, Optional - -from .. import __version__ -from ..universaldetector import UniversalDetector - - -def description_of( - lines: Iterable[bytes], - name: str = "stdin", - minimal: bool = False, - should_rename_legacy: bool = False, -) -> Optional[str]: - """ - Return a string describing the probable encoding of a file or - list of strings. - - :param lines: The lines to get the encoding of. - :type lines: Iterable of bytes - :param name: Name of file or collection of lines - :type name: str - :param should_rename_legacy: Should we rename legacy encodings to - their more modern equivalents? - :type should_rename_legacy: ``bool`` - """ - u = UniversalDetector(should_rename_legacy=should_rename_legacy) - for line in lines: - line = bytearray(line) - u.feed(line) - # shortcut out of the loop to save reading further - particularly useful if we read a BOM. - if u.done: - break - u.close() - result = u.result - if minimal: - return result["encoding"] - if result["encoding"]: - return f'{name}: {result["encoding"]} with confidence {result["confidence"]}' - return f"{name}: no result" - - -def main(argv: Optional[List[str]] = None) -> None: - """ - Handles command line arguments and gets things started. - - :param argv: List of arguments, as if specified on the command-line. - If None, ``sys.argv[1:]`` is used instead. - :type argv: list of str - """ - # Get command line arguments - parser = argparse.ArgumentParser( - description=( - "Takes one or more file paths and reports their detected encodings" - ) - ) - parser.add_argument( - "input", - help="File whose encoding we would like to determine. (default: stdin)", - type=argparse.FileType("rb"), - nargs="*", - default=[sys.stdin.buffer], - ) - parser.add_argument( - "--minimal", - help="Print only the encoding to standard output", - action="store_true", - ) - parser.add_argument( - "-l", - "--legacy", - help="Rename legacy encodings to more modern ones.", - action="store_true", - ) - parser.add_argument( - "--version", action="version", version=f"%(prog)s {__version__}" - ) - args = parser.parse_args(argv) - - for f in args.input: - if f.isatty(): - print( - "You are running chardetect interactively. Press " - "CTRL-D twice at the start of a blank line to signal the " - "end of your input. If you want help, run chardetect " - "--help\n", - file=sys.stderr, - ) - print( - description_of( - f, f.name, minimal=args.minimal, should_rename_legacy=args.legacy - ) - ) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/chardet/codingstatemachine.py b/venv/lib/python3.10/site-packages/chardet/codingstatemachine.py deleted file mode 100644 index 8ed4a8773b8404c2705aa8728e5fd692362ba168..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/codingstatemachine.py +++ /dev/null @@ -1,90 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is mozilla.org code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -import logging - -from .codingstatemachinedict import CodingStateMachineDict -from .enums import MachineState - - -class CodingStateMachine: - """ - A state machine to verify a byte sequence for a particular encoding. For - each byte the detector receives, it will feed that byte to every active - state machine available, one byte at a time. The state machine changes its - state based on its previous state and the byte it receives. There are 3 - states in a state machine that are of interest to an auto-detector: - - START state: This is the state to start with, or a legal byte sequence - (i.e. a valid code point) for character has been identified. - - ME state: This indicates that the state machine identified a byte sequence - that is specific to the charset it is designed for and that - there is no other possible encoding which can contain this byte - sequence. This will to lead to an immediate positive answer for - the detector. - - ERROR state: This indicates the state machine identified an illegal byte - sequence for that encoding. This will lead to an immediate - negative answer for this encoding. Detector will exclude this - encoding from consideration from here on. - """ - - def __init__(self, sm: CodingStateMachineDict) -> None: - self._model = sm - self._curr_byte_pos = 0 - self._curr_char_len = 0 - self._curr_state = MachineState.START - self.active = True - self.logger = logging.getLogger(__name__) - self.reset() - - def reset(self) -> None: - self._curr_state = MachineState.START - - def next_state(self, c: int) -> int: - # for each byte we get its class - # if it is first byte, we also get byte length - byte_class = self._model["class_table"][c] - if self._curr_state == MachineState.START: - self._curr_byte_pos = 0 - self._curr_char_len = self._model["char_len_table"][byte_class] - # from byte's class and state_table, we get its next state - curr_state = self._curr_state * self._model["class_factor"] + byte_class - self._curr_state = self._model["state_table"][curr_state] - self._curr_byte_pos += 1 - return self._curr_state - - def get_current_charlen(self) -> int: - return self._curr_char_len - - def get_coding_state_machine(self) -> str: - return self._model["name"] - - @property - def language(self) -> str: - return self._model["language"] diff --git a/venv/lib/python3.10/site-packages/chardet/codingstatemachinedict.py b/venv/lib/python3.10/site-packages/chardet/codingstatemachinedict.py deleted file mode 100644 index 7a3c4c7e3fe16e91225a87cbc58b8bbd798f9cc1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/codingstatemachinedict.py +++ /dev/null @@ -1,19 +0,0 @@ -from typing import TYPE_CHECKING, Tuple - -if TYPE_CHECKING: - # TypedDict was introduced in Python 3.8. - # - # TODO: Remove the else block and TYPE_CHECKING check when dropping support - # for Python 3.7. - from typing import TypedDict - - class CodingStateMachineDict(TypedDict, total=False): - class_table: Tuple[int, ...] - class_factor: int - state_table: Tuple[int, ...] - char_len_table: Tuple[int, ...] - name: str - language: str # Optional key - -else: - CodingStateMachineDict = dict diff --git a/venv/lib/python3.10/site-packages/chardet/cp949prober.py b/venv/lib/python3.10/site-packages/chardet/cp949prober.py deleted file mode 100644 index fa7307ed8985ad7e318660da0066440f890d1624..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/cp949prober.py +++ /dev/null @@ -1,49 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is mozilla.org code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from .chardistribution import EUCKRDistributionAnalysis -from .codingstatemachine import CodingStateMachine -from .mbcharsetprober import MultiByteCharSetProber -from .mbcssm import CP949_SM_MODEL - - -class CP949Prober(MultiByteCharSetProber): - def __init__(self) -> None: - super().__init__() - self.coding_sm = CodingStateMachine(CP949_SM_MODEL) - # NOTE: CP949 is a superset of EUC-KR, so the distribution should be - # not different. - self.distribution_analyzer = EUCKRDistributionAnalysis() - self.reset() - - @property - def charset_name(self) -> str: - return "CP949" - - @property - def language(self) -> str: - return "Korean" diff --git a/venv/lib/python3.10/site-packages/chardet/enums.py b/venv/lib/python3.10/site-packages/chardet/enums.py deleted file mode 100644 index 5e3e198233698f2b007489dd299cecb87d971067..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/enums.py +++ /dev/null @@ -1,85 +0,0 @@ -""" -All of the Enums that are used throughout the chardet package. - -:author: Dan Blanchard (dan.blanchard@gmail.com) -""" - -from enum import Enum, Flag - - -class InputState: - """ - This enum represents the different states a universal detector can be in. - """ - - PURE_ASCII = 0 - ESC_ASCII = 1 - HIGH_BYTE = 2 - - -class LanguageFilter(Flag): - """ - This enum represents the different language filters we can apply to a - ``UniversalDetector``. - """ - - NONE = 0x00 - CHINESE_SIMPLIFIED = 0x01 - CHINESE_TRADITIONAL = 0x02 - JAPANESE = 0x04 - KOREAN = 0x08 - NON_CJK = 0x10 - ALL = 0x1F - CHINESE = CHINESE_SIMPLIFIED | CHINESE_TRADITIONAL - CJK = CHINESE | JAPANESE | KOREAN - - -class ProbingState(Enum): - """ - This enum represents the different states a prober can be in. - """ - - DETECTING = 0 - FOUND_IT = 1 - NOT_ME = 2 - - -class MachineState: - """ - This enum represents the different states a state machine can be in. - """ - - START = 0 - ERROR = 1 - ITS_ME = 2 - - -class SequenceLikelihood: - """ - This enum represents the likelihood of a character following the previous one. - """ - - NEGATIVE = 0 - UNLIKELY = 1 - LIKELY = 2 - POSITIVE = 3 - - @classmethod - def get_num_categories(cls) -> int: - """:returns: The number of likelihood categories in the enum.""" - return 4 - - -class CharacterCategory: - """ - This enum represents the different categories language models for - ``SingleByteCharsetProber`` put characters into. - - Anything less than CONTROL is considered a letter. - """ - - UNDEFINED = 255 - LINE_BREAK = 254 - SYMBOL = 253 - DIGIT = 252 - CONTROL = 251 diff --git a/venv/lib/python3.10/site-packages/chardet/escprober.py b/venv/lib/python3.10/site-packages/chardet/escprober.py deleted file mode 100644 index fd713830d36cabc6a0fb4ab4e8cf426a84decdc6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/escprober.py +++ /dev/null @@ -1,102 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is mozilla.org code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from typing import Optional, Union - -from .charsetprober import CharSetProber -from .codingstatemachine import CodingStateMachine -from .enums import LanguageFilter, MachineState, ProbingState -from .escsm import ( - HZ_SM_MODEL, - ISO2022CN_SM_MODEL, - ISO2022JP_SM_MODEL, - ISO2022KR_SM_MODEL, -) - - -class EscCharSetProber(CharSetProber): - """ - This CharSetProber uses a "code scheme" approach for detecting encodings, - whereby easily recognizable escape or shift sequences are relied on to - identify these encodings. - """ - - def __init__(self, lang_filter: LanguageFilter = LanguageFilter.NONE) -> None: - super().__init__(lang_filter=lang_filter) - self.coding_sm = [] - if self.lang_filter & LanguageFilter.CHINESE_SIMPLIFIED: - self.coding_sm.append(CodingStateMachine(HZ_SM_MODEL)) - self.coding_sm.append(CodingStateMachine(ISO2022CN_SM_MODEL)) - if self.lang_filter & LanguageFilter.JAPANESE: - self.coding_sm.append(CodingStateMachine(ISO2022JP_SM_MODEL)) - if self.lang_filter & LanguageFilter.KOREAN: - self.coding_sm.append(CodingStateMachine(ISO2022KR_SM_MODEL)) - self.active_sm_count = 0 - self._detected_charset: Optional[str] = None - self._detected_language: Optional[str] = None - self._state = ProbingState.DETECTING - self.reset() - - def reset(self) -> None: - super().reset() - for coding_sm in self.coding_sm: - coding_sm.active = True - coding_sm.reset() - self.active_sm_count = len(self.coding_sm) - self._detected_charset = None - self._detected_language = None - - @property - def charset_name(self) -> Optional[str]: - return self._detected_charset - - @property - def language(self) -> Optional[str]: - return self._detected_language - - def get_confidence(self) -> float: - return 0.99 if self._detected_charset else 0.00 - - def feed(self, byte_str: Union[bytes, bytearray]) -> ProbingState: - for c in byte_str: - for coding_sm in self.coding_sm: - if not coding_sm.active: - continue - coding_state = coding_sm.next_state(c) - if coding_state == MachineState.ERROR: - coding_sm.active = False - self.active_sm_count -= 1 - if self.active_sm_count <= 0: - self._state = ProbingState.NOT_ME - return self.state - elif coding_state == MachineState.ITS_ME: - self._state = ProbingState.FOUND_IT - self._detected_charset = coding_sm.get_coding_state_machine() - self._detected_language = coding_sm.language - return self.state - - return self.state diff --git a/venv/lib/python3.10/site-packages/chardet/escsm.py b/venv/lib/python3.10/site-packages/chardet/escsm.py deleted file mode 100644 index 11d4adf771f3f90bb5f1cc11043599b48e955c22..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/escsm.py +++ /dev/null @@ -1,261 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is mozilla.org code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from .codingstatemachinedict import CodingStateMachineDict -from .enums import MachineState - -# fmt: off -HZ_CLS = ( - 1, 0, 0, 0, 0, 0, 0, 0, # 00 - 07 - 0, 0, 0, 0, 0, 0, 0, 0, # 08 - 0f - 0, 0, 0, 0, 0, 0, 0, 0, # 10 - 17 - 0, 0, 0, 1, 0, 0, 0, 0, # 18 - 1f - 0, 0, 0, 0, 0, 0, 0, 0, # 20 - 27 - 0, 0, 0, 0, 0, 0, 0, 0, # 28 - 2f - 0, 0, 0, 0, 0, 0, 0, 0, # 30 - 37 - 0, 0, 0, 0, 0, 0, 0, 0, # 38 - 3f - 0, 0, 0, 0, 0, 0, 0, 0, # 40 - 47 - 0, 0, 0, 0, 0, 0, 0, 0, # 48 - 4f - 0, 0, 0, 0, 0, 0, 0, 0, # 50 - 57 - 0, 0, 0, 0, 0, 0, 0, 0, # 58 - 5f - 0, 0, 0, 0, 0, 0, 0, 0, # 60 - 67 - 0, 0, 0, 0, 0, 0, 0, 0, # 68 - 6f - 0, 0, 0, 0, 0, 0, 0, 0, # 70 - 77 - 0, 0, 0, 4, 0, 5, 2, 0, # 78 - 7f - 1, 1, 1, 1, 1, 1, 1, 1, # 80 - 87 - 1, 1, 1, 1, 1, 1, 1, 1, # 88 - 8f - 1, 1, 1, 1, 1, 1, 1, 1, # 90 - 97 - 1, 1, 1, 1, 1, 1, 1, 1, # 98 - 9f - 1, 1, 1, 1, 1, 1, 1, 1, # a0 - a7 - 1, 1, 1, 1, 1, 1, 1, 1, # a8 - af - 1, 1, 1, 1, 1, 1, 1, 1, # b0 - b7 - 1, 1, 1, 1, 1, 1, 1, 1, # b8 - bf - 1, 1, 1, 1, 1, 1, 1, 1, # c0 - c7 - 1, 1, 1, 1, 1, 1, 1, 1, # c8 - cf - 1, 1, 1, 1, 1, 1, 1, 1, # d0 - d7 - 1, 1, 1, 1, 1, 1, 1, 1, # d8 - df - 1, 1, 1, 1, 1, 1, 1, 1, # e0 - e7 - 1, 1, 1, 1, 1, 1, 1, 1, # e8 - ef - 1, 1, 1, 1, 1, 1, 1, 1, # f0 - f7 - 1, 1, 1, 1, 1, 1, 1, 1, # f8 - ff -) - -HZ_ST = ( -MachineState.START, MachineState.ERROR, 3, MachineState.START, MachineState.START, MachineState.START, MachineState.ERROR, MachineState.ERROR, # 00-07 -MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, # 08-0f -MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ERROR, MachineState.ERROR, MachineState.START, MachineState.START, 4, MachineState.ERROR, # 10-17 - 5, MachineState.ERROR, 6, MachineState.ERROR, 5, 5, 4, MachineState.ERROR, # 18-1f - 4, MachineState.ERROR, 4, 4, 4, MachineState.ERROR, 4, MachineState.ERROR, # 20-27 - 4, MachineState.ITS_ME, MachineState.START, MachineState.START, MachineState.START, MachineState.START, MachineState.START, MachineState.START, # 28-2f -) -# fmt: on - -HZ_CHAR_LEN_TABLE = (0, 0, 0, 0, 0, 0) - -HZ_SM_MODEL: CodingStateMachineDict = { - "class_table": HZ_CLS, - "class_factor": 6, - "state_table": HZ_ST, - "char_len_table": HZ_CHAR_LEN_TABLE, - "name": "HZ-GB-2312", - "language": "Chinese", -} - -# fmt: off -ISO2022CN_CLS = ( - 2, 0, 0, 0, 0, 0, 0, 0, # 00 - 07 - 0, 0, 0, 0, 0, 0, 0, 0, # 08 - 0f - 0, 0, 0, 0, 0, 0, 0, 0, # 10 - 17 - 0, 0, 0, 1, 0, 0, 0, 0, # 18 - 1f - 0, 0, 0, 0, 0, 0, 0, 0, # 20 - 27 - 0, 3, 0, 0, 0, 0, 0, 0, # 28 - 2f - 0, 0, 0, 0, 0, 0, 0, 0, # 30 - 37 - 0, 0, 0, 0, 0, 0, 0, 0, # 38 - 3f - 0, 0, 0, 4, 0, 0, 0, 0, # 40 - 47 - 0, 0, 0, 0, 0, 0, 0, 0, # 48 - 4f - 0, 0, 0, 0, 0, 0, 0, 0, # 50 - 57 - 0, 0, 0, 0, 0, 0, 0, 0, # 58 - 5f - 0, 0, 0, 0, 0, 0, 0, 0, # 60 - 67 - 0, 0, 0, 0, 0, 0, 0, 0, # 68 - 6f - 0, 0, 0, 0, 0, 0, 0, 0, # 70 - 77 - 0, 0, 0, 0, 0, 0, 0, 0, # 78 - 7f - 2, 2, 2, 2, 2, 2, 2, 2, # 80 - 87 - 2, 2, 2, 2, 2, 2, 2, 2, # 88 - 8f - 2, 2, 2, 2, 2, 2, 2, 2, # 90 - 97 - 2, 2, 2, 2, 2, 2, 2, 2, # 98 - 9f - 2, 2, 2, 2, 2, 2, 2, 2, # a0 - a7 - 2, 2, 2, 2, 2, 2, 2, 2, # a8 - af - 2, 2, 2, 2, 2, 2, 2, 2, # b0 - b7 - 2, 2, 2, 2, 2, 2, 2, 2, # b8 - bf - 2, 2, 2, 2, 2, 2, 2, 2, # c0 - c7 - 2, 2, 2, 2, 2, 2, 2, 2, # c8 - cf - 2, 2, 2, 2, 2, 2, 2, 2, # d0 - d7 - 2, 2, 2, 2, 2, 2, 2, 2, # d8 - df - 2, 2, 2, 2, 2, 2, 2, 2, # e0 - e7 - 2, 2, 2, 2, 2, 2, 2, 2, # e8 - ef - 2, 2, 2, 2, 2, 2, 2, 2, # f0 - f7 - 2, 2, 2, 2, 2, 2, 2, 2, # f8 - ff -) - -ISO2022CN_ST = ( - MachineState.START, 3, MachineState.ERROR, MachineState.START, MachineState.START, MachineState.START, MachineState.START, MachineState.START, # 00-07 - MachineState.START, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, # 08-0f - MachineState.ERROR, MachineState.ERROR, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, # 10-17 - MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, 4, MachineState.ERROR, # 18-1f - MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ITS_ME, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, # 20-27 - 5, 6, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, # 28-2f - MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ITS_ME, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, # 30-37 - MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ITS_ME, MachineState.ERROR, MachineState.START, # 38-3f -) -# fmt: on - -ISO2022CN_CHAR_LEN_TABLE = (0, 0, 0, 0, 0, 0, 0, 0, 0) - -ISO2022CN_SM_MODEL: CodingStateMachineDict = { - "class_table": ISO2022CN_CLS, - "class_factor": 9, - "state_table": ISO2022CN_ST, - "char_len_table": ISO2022CN_CHAR_LEN_TABLE, - "name": "ISO-2022-CN", - "language": "Chinese", -} - -# fmt: off -ISO2022JP_CLS = ( - 2, 0, 0, 0, 0, 0, 0, 0, # 00 - 07 - 0, 0, 0, 0, 0, 0, 2, 2, # 08 - 0f - 0, 0, 0, 0, 0, 0, 0, 0, # 10 - 17 - 0, 0, 0, 1, 0, 0, 0, 0, # 18 - 1f - 0, 0, 0, 0, 7, 0, 0, 0, # 20 - 27 - 3, 0, 0, 0, 0, 0, 0, 0, # 28 - 2f - 0, 0, 0, 0, 0, 0, 0, 0, # 30 - 37 - 0, 0, 0, 0, 0, 0, 0, 0, # 38 - 3f - 6, 0, 4, 0, 8, 0, 0, 0, # 40 - 47 - 0, 9, 5, 0, 0, 0, 0, 0, # 48 - 4f - 0, 0, 0, 0, 0, 0, 0, 0, # 50 - 57 - 0, 0, 0, 0, 0, 0, 0, 0, # 58 - 5f - 0, 0, 0, 0, 0, 0, 0, 0, # 60 - 67 - 0, 0, 0, 0, 0, 0, 0, 0, # 68 - 6f - 0, 0, 0, 0, 0, 0, 0, 0, # 70 - 77 - 0, 0, 0, 0, 0, 0, 0, 0, # 78 - 7f - 2, 2, 2, 2, 2, 2, 2, 2, # 80 - 87 - 2, 2, 2, 2, 2, 2, 2, 2, # 88 - 8f - 2, 2, 2, 2, 2, 2, 2, 2, # 90 - 97 - 2, 2, 2, 2, 2, 2, 2, 2, # 98 - 9f - 2, 2, 2, 2, 2, 2, 2, 2, # a0 - a7 - 2, 2, 2, 2, 2, 2, 2, 2, # a8 - af - 2, 2, 2, 2, 2, 2, 2, 2, # b0 - b7 - 2, 2, 2, 2, 2, 2, 2, 2, # b8 - bf - 2, 2, 2, 2, 2, 2, 2, 2, # c0 - c7 - 2, 2, 2, 2, 2, 2, 2, 2, # c8 - cf - 2, 2, 2, 2, 2, 2, 2, 2, # d0 - d7 - 2, 2, 2, 2, 2, 2, 2, 2, # d8 - df - 2, 2, 2, 2, 2, 2, 2, 2, # e0 - e7 - 2, 2, 2, 2, 2, 2, 2, 2, # e8 - ef - 2, 2, 2, 2, 2, 2, 2, 2, # f0 - f7 - 2, 2, 2, 2, 2, 2, 2, 2, # f8 - ff -) - -ISO2022JP_ST = ( - MachineState.START, 3, MachineState.ERROR, MachineState.START, MachineState.START, MachineState.START, MachineState.START, MachineState.START, # 00-07 - MachineState.START, MachineState.START, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, # 08-0f - MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, # 10-17 - MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ERROR, MachineState.ERROR, # 18-1f - MachineState.ERROR, 5, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, 4, MachineState.ERROR, MachineState.ERROR, # 20-27 - MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, 6, MachineState.ITS_ME, MachineState.ERROR, MachineState.ITS_ME, MachineState.ERROR, # 28-2f - MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ITS_ME, MachineState.ITS_ME, # 30-37 - MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ITS_ME, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, # 38-3f - MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ITS_ME, MachineState.ERROR, MachineState.START, MachineState.START, # 40-47 -) -# fmt: on - -ISO2022JP_CHAR_LEN_TABLE = (0, 0, 0, 0, 0, 0, 0, 0, 0, 0) - -ISO2022JP_SM_MODEL: CodingStateMachineDict = { - "class_table": ISO2022JP_CLS, - "class_factor": 10, - "state_table": ISO2022JP_ST, - "char_len_table": ISO2022JP_CHAR_LEN_TABLE, - "name": "ISO-2022-JP", - "language": "Japanese", -} - -# fmt: off -ISO2022KR_CLS = ( - 2, 0, 0, 0, 0, 0, 0, 0, # 00 - 07 - 0, 0, 0, 0, 0, 0, 0, 0, # 08 - 0f - 0, 0, 0, 0, 0, 0, 0, 0, # 10 - 17 - 0, 0, 0, 1, 0, 0, 0, 0, # 18 - 1f - 0, 0, 0, 0, 3, 0, 0, 0, # 20 - 27 - 0, 4, 0, 0, 0, 0, 0, 0, # 28 - 2f - 0, 0, 0, 0, 0, 0, 0, 0, # 30 - 37 - 0, 0, 0, 0, 0, 0, 0, 0, # 38 - 3f - 0, 0, 0, 5, 0, 0, 0, 0, # 40 - 47 - 0, 0, 0, 0, 0, 0, 0, 0, # 48 - 4f - 0, 0, 0, 0, 0, 0, 0, 0, # 50 - 57 - 0, 0, 0, 0, 0, 0, 0, 0, # 58 - 5f - 0, 0, 0, 0, 0, 0, 0, 0, # 60 - 67 - 0, 0, 0, 0, 0, 0, 0, 0, # 68 - 6f - 0, 0, 0, 0, 0, 0, 0, 0, # 70 - 77 - 0, 0, 0, 0, 0, 0, 0, 0, # 78 - 7f - 2, 2, 2, 2, 2, 2, 2, 2, # 80 - 87 - 2, 2, 2, 2, 2, 2, 2, 2, # 88 - 8f - 2, 2, 2, 2, 2, 2, 2, 2, # 90 - 97 - 2, 2, 2, 2, 2, 2, 2, 2, # 98 - 9f - 2, 2, 2, 2, 2, 2, 2, 2, # a0 - a7 - 2, 2, 2, 2, 2, 2, 2, 2, # a8 - af - 2, 2, 2, 2, 2, 2, 2, 2, # b0 - b7 - 2, 2, 2, 2, 2, 2, 2, 2, # b8 - bf - 2, 2, 2, 2, 2, 2, 2, 2, # c0 - c7 - 2, 2, 2, 2, 2, 2, 2, 2, # c8 - cf - 2, 2, 2, 2, 2, 2, 2, 2, # d0 - d7 - 2, 2, 2, 2, 2, 2, 2, 2, # d8 - df - 2, 2, 2, 2, 2, 2, 2, 2, # e0 - e7 - 2, 2, 2, 2, 2, 2, 2, 2, # e8 - ef - 2, 2, 2, 2, 2, 2, 2, 2, # f0 - f7 - 2, 2, 2, 2, 2, 2, 2, 2, # f8 - ff -) - -ISO2022KR_ST = ( - MachineState.START, 3, MachineState.ERROR, MachineState.START, MachineState.START, MachineState.START, MachineState.ERROR, MachineState.ERROR, # 00-07 - MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ITS_ME, # 08-0f - MachineState.ITS_ME, MachineState.ITS_ME, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, 4, MachineState.ERROR, MachineState.ERROR, # 10-17 - MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, 5, MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, # 18-1f - MachineState.ERROR, MachineState.ERROR, MachineState.ERROR, MachineState.ITS_ME, MachineState.START, MachineState.START, MachineState.START, MachineState.START, # 20-27 -) -# fmt: on - -ISO2022KR_CHAR_LEN_TABLE = (0, 0, 0, 0, 0, 0) - -ISO2022KR_SM_MODEL: CodingStateMachineDict = { - "class_table": ISO2022KR_CLS, - "class_factor": 6, - "state_table": ISO2022KR_ST, - "char_len_table": ISO2022KR_CHAR_LEN_TABLE, - "name": "ISO-2022-KR", - "language": "Korean", -} diff --git a/venv/lib/python3.10/site-packages/chardet/eucjpprober.py b/venv/lib/python3.10/site-packages/chardet/eucjpprober.py deleted file mode 100644 index 39487f4098d7c2068b67d7d3dd85b61848974a23..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/eucjpprober.py +++ /dev/null @@ -1,102 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is mozilla.org code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from typing import Union - -from .chardistribution import EUCJPDistributionAnalysis -from .codingstatemachine import CodingStateMachine -from .enums import MachineState, ProbingState -from .jpcntx import EUCJPContextAnalysis -from .mbcharsetprober import MultiByteCharSetProber -from .mbcssm import EUCJP_SM_MODEL - - -class EUCJPProber(MultiByteCharSetProber): - def __init__(self) -> None: - super().__init__() - self.coding_sm = CodingStateMachine(EUCJP_SM_MODEL) - self.distribution_analyzer = EUCJPDistributionAnalysis() - self.context_analyzer = EUCJPContextAnalysis() - self.reset() - - def reset(self) -> None: - super().reset() - self.context_analyzer.reset() - - @property - def charset_name(self) -> str: - return "EUC-JP" - - @property - def language(self) -> str: - return "Japanese" - - def feed(self, byte_str: Union[bytes, bytearray]) -> ProbingState: - assert self.coding_sm is not None - assert self.distribution_analyzer is not None - - for i, byte in enumerate(byte_str): - # PY3K: byte_str is a byte array, so byte is an int, not a byte - coding_state = self.coding_sm.next_state(byte) - if coding_state == MachineState.ERROR: - self.logger.debug( - "%s %s prober hit error at byte %s", - self.charset_name, - self.language, - i, - ) - self._state = ProbingState.NOT_ME - break - if coding_state == MachineState.ITS_ME: - self._state = ProbingState.FOUND_IT - break - if coding_state == MachineState.START: - char_len = self.coding_sm.get_current_charlen() - if i == 0: - self._last_char[1] = byte - self.context_analyzer.feed(self._last_char, char_len) - self.distribution_analyzer.feed(self._last_char, char_len) - else: - self.context_analyzer.feed(byte_str[i - 1 : i + 1], char_len) - self.distribution_analyzer.feed(byte_str[i - 1 : i + 1], char_len) - - self._last_char[0] = byte_str[-1] - - if self.state == ProbingState.DETECTING: - if self.context_analyzer.got_enough_data() and ( - self.get_confidence() > self.SHORTCUT_THRESHOLD - ): - self._state = ProbingState.FOUND_IT - - return self.state - - def get_confidence(self) -> float: - assert self.distribution_analyzer is not None - - context_conf = self.context_analyzer.get_confidence() - distrib_conf = self.distribution_analyzer.get_confidence() - return max(context_conf, distrib_conf) diff --git a/venv/lib/python3.10/site-packages/chardet/euckrfreq.py b/venv/lib/python3.10/site-packages/chardet/euckrfreq.py deleted file mode 100644 index 7dc3b10387d1c3d2da8b4e27e917ee2a85086e0c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/euckrfreq.py +++ /dev/null @@ -1,196 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Communicator client code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -# Sampling from about 20M text materials include literature and computer technology - -# 128 --> 0.79 -# 256 --> 0.92 -# 512 --> 0.986 -# 1024 --> 0.99944 -# 2048 --> 0.99999 -# -# Idea Distribution Ratio = 0.98653 / (1-0.98653) = 73.24 -# Random Distribution Ration = 512 / (2350-512) = 0.279. -# -# Typical Distribution Ratio - -EUCKR_TYPICAL_DISTRIBUTION_RATIO = 6.0 - -EUCKR_TABLE_SIZE = 2352 - -# Char to FreqOrder table , -# fmt: off -EUCKR_CHAR_TO_FREQ_ORDER = ( - 13, 130, 120,1396, 481,1719,1720, 328, 609, 212,1721, 707, 400, 299,1722, 87, -1397,1723, 104, 536,1117,1203,1724,1267, 685,1268, 508,1725,1726,1727,1728,1398, -1399,1729,1730,1731, 141, 621, 326,1057, 368,1732, 267, 488, 20,1733,1269,1734, - 945,1400,1735, 47, 904,1270,1736,1737, 773, 248,1738, 409, 313, 786, 429,1739, - 116, 987, 813,1401, 683, 75,1204, 145,1740,1741,1742,1743, 16, 847, 667, 622, - 708,1744,1745,1746, 966, 787, 304, 129,1747, 60, 820, 123, 676,1748,1749,1750, -1751, 617,1752, 626,1753,1754,1755,1756, 653,1757,1758,1759,1760,1761,1762, 856, - 344,1763,1764,1765,1766, 89, 401, 418, 806, 905, 848,1767,1768,1769, 946,1205, - 709,1770,1118,1771, 241,1772,1773,1774,1271,1775, 569,1776, 999,1777,1778,1779, -1780, 337, 751,1058, 28, 628, 254,1781, 177, 906, 270, 349, 891,1079,1782, 19, -1783, 379,1784, 315,1785, 629, 754,1402, 559,1786, 636, 203,1206,1787, 710, 567, -1788, 935, 814,1789,1790,1207, 766, 528,1791,1792,1208,1793,1794,1795,1796,1797, -1403,1798,1799, 533,1059,1404,1405,1156,1406, 936, 884,1080,1800, 351,1801,1802, -1803,1804,1805, 801,1806,1807,1808,1119,1809,1157, 714, 474,1407,1810, 298, 899, - 885,1811,1120, 802,1158,1812, 892,1813,1814,1408, 659,1815,1816,1121,1817,1818, -1819,1820,1821,1822, 319,1823, 594, 545,1824, 815, 937,1209,1825,1826, 573,1409, -1022,1827,1210,1828,1829,1830,1831,1832,1833, 556, 722, 807,1122,1060,1834, 697, -1835, 900, 557, 715,1836,1410, 540,1411, 752,1159, 294, 597,1211, 976, 803, 770, -1412,1837,1838, 39, 794,1413, 358,1839, 371, 925,1840, 453, 661, 788, 531, 723, - 544,1023,1081, 869, 91,1841, 392, 430, 790, 602,1414, 677,1082, 457,1415,1416, -1842,1843, 475, 327,1024,1417, 795, 121,1844, 733, 403,1418,1845,1846,1847, 300, - 119, 711,1212, 627,1848,1272, 207,1849,1850, 796,1213, 382,1851, 519,1852,1083, - 893,1853,1854,1855, 367, 809, 487, 671,1856, 663,1857,1858, 956, 471, 306, 857, -1859,1860,1160,1084,1861,1862,1863,1864,1865,1061,1866,1867,1868,1869,1870,1871, - 282, 96, 574,1872, 502,1085,1873,1214,1874, 907,1875,1876, 827, 977,1419,1420, -1421, 268,1877,1422,1878,1879,1880, 308,1881, 2, 537,1882,1883,1215,1884,1885, - 127, 791,1886,1273,1423,1887, 34, 336, 404, 643,1888, 571, 654, 894, 840,1889, - 0, 886,1274, 122, 575, 260, 908, 938,1890,1275, 410, 316,1891,1892, 100,1893, -1894,1123, 48,1161,1124,1025,1895, 633, 901,1276,1896,1897, 115, 816,1898, 317, -1899, 694,1900, 909, 734,1424, 572, 866,1425, 691, 85, 524,1010, 543, 394, 841, -1901,1902,1903,1026,1904,1905,1906,1907,1908,1909, 30, 451, 651, 988, 310,1910, -1911,1426, 810,1216, 93,1912,1913,1277,1217,1914, 858, 759, 45, 58, 181, 610, - 269,1915,1916, 131,1062, 551, 443,1000, 821,1427, 957, 895,1086,1917,1918, 375, -1919, 359,1920, 687,1921, 822,1922, 293,1923,1924, 40, 662, 118, 692, 29, 939, - 887, 640, 482, 174,1925, 69,1162, 728,1428, 910,1926,1278,1218,1279, 386, 870, - 217, 854,1163, 823,1927,1928,1929,1930, 834,1931, 78,1932, 859,1933,1063,1934, -1935,1936,1937, 438,1164, 208, 595,1938,1939,1940,1941,1219,1125,1942, 280, 888, -1429,1430,1220,1431,1943,1944,1945,1946,1947,1280, 150, 510,1432,1948,1949,1950, -1951,1952,1953,1954,1011,1087,1955,1433,1043,1956, 881,1957, 614, 958,1064,1065, -1221,1958, 638,1001, 860, 967, 896,1434, 989, 492, 553,1281,1165,1959,1282,1002, -1283,1222,1960,1961,1962,1963, 36, 383, 228, 753, 247, 454,1964, 876, 678,1965, -1966,1284, 126, 464, 490, 835, 136, 672, 529, 940,1088,1435, 473,1967,1968, 467, - 50, 390, 227, 587, 279, 378, 598, 792, 968, 240, 151, 160, 849, 882,1126,1285, - 639,1044, 133, 140, 288, 360, 811, 563,1027, 561, 142, 523,1969,1970,1971, 7, - 103, 296, 439, 407, 506, 634, 990,1972,1973,1974,1975, 645,1976,1977,1978,1979, -1980,1981, 236,1982,1436,1983,1984,1089, 192, 828, 618, 518,1166, 333,1127,1985, - 818,1223,1986,1987,1988,1989,1990,1991,1992,1993, 342,1128,1286, 746, 842,1994, -1995, 560, 223,1287, 98, 8, 189, 650, 978,1288,1996,1437,1997, 17, 345, 250, - 423, 277, 234, 512, 226, 97, 289, 42, 167,1998, 201,1999,2000, 843, 836, 824, - 532, 338, 783,1090, 182, 576, 436,1438,1439, 527, 500,2001, 947, 889,2002,2003, -2004,2005, 262, 600, 314, 447,2006, 547,2007, 693, 738,1129,2008, 71,1440, 745, - 619, 688,2009, 829,2010,2011, 147,2012, 33, 948,2013,2014, 74, 224,2015, 61, - 191, 918, 399, 637,2016,1028,1130, 257, 902,2017,2018,2019,2020,2021,2022,2023, -2024,2025,2026, 837,2027,2028,2029,2030, 179, 874, 591, 52, 724, 246,2031,2032, -2033,2034,1167, 969,2035,1289, 630, 605, 911,1091,1168,2036,2037,2038,1441, 912, -2039, 623,2040,2041, 253,1169,1290,2042,1442, 146, 620, 611, 577, 433,2043,1224, - 719,1170, 959, 440, 437, 534, 84, 388, 480,1131, 159, 220, 198, 679,2044,1012, - 819,1066,1443, 113,1225, 194, 318,1003,1029,2045,2046,2047,2048,1067,2049,2050, -2051,2052,2053, 59, 913, 112,2054, 632,2055, 455, 144, 739,1291,2056, 273, 681, - 499,2057, 448,2058,2059, 760,2060,2061, 970, 384, 169, 245,1132,2062,2063, 414, -1444,2064,2065, 41, 235,2066, 157, 252, 877, 568, 919, 789, 580,2067, 725,2068, -2069,1292,2070,2071,1445,2072,1446,2073,2074, 55, 588, 66,1447, 271,1092,2075, -1226,2076, 960,1013, 372,2077,2078,2079,2080,2081,1293,2082,2083,2084,2085, 850, -2086,2087,2088,2089,2090, 186,2091,1068, 180,2092,2093,2094, 109,1227, 522, 606, -2095, 867,1448,1093, 991,1171, 926, 353,1133,2096, 581,2097,2098,2099,1294,1449, -1450,2100, 596,1172,1014,1228,2101,1451,1295,1173,1229,2102,2103,1296,1134,1452, - 949,1135,2104,2105,1094,1453,1454,1455,2106,1095,2107,2108,2109,2110,2111,2112, -2113,2114,2115,2116,2117, 804,2118,2119,1230,1231, 805,1456, 405,1136,2120,2121, -2122,2123,2124, 720, 701,1297, 992,1457, 927,1004,2125,2126,2127,2128,2129,2130, - 22, 417,2131, 303,2132, 385,2133, 971, 520, 513,2134,1174, 73,1096, 231, 274, - 962,1458, 673,2135,1459,2136, 152,1137,2137,2138,2139,2140,1005,1138,1460,1139, -2141,2142,2143,2144, 11, 374, 844,2145, 154,1232, 46,1461,2146, 838, 830, 721, -1233, 106,2147, 90, 428, 462, 578, 566,1175, 352,2148,2149, 538,1234, 124,1298, -2150,1462, 761, 565,2151, 686,2152, 649,2153, 72, 173,2154, 460, 415,2155,1463, -2156,1235, 305,2157,2158,2159,2160,2161,2162, 579,2163,2164,2165,2166,2167, 747, -2168,2169,2170,2171,1464, 669,2172,2173,2174,2175,2176,1465,2177, 23, 530, 285, -2178, 335, 729,2179, 397,2180,2181,2182,1030,2183,2184, 698,2185,2186, 325,2187, -2188, 369,2189, 799,1097,1015, 348,2190,1069, 680,2191, 851,1466,2192,2193, 10, -2194, 613, 424,2195, 979, 108, 449, 589, 27, 172, 81,1031, 80, 774, 281, 350, -1032, 525, 301, 582,1176,2196, 674,1045,2197,2198,1467, 730, 762,2199,2200,2201, -2202,1468,2203, 993,2204,2205, 266,1070, 963,1140,2206,2207,2208, 664,1098, 972, -2209,2210,2211,1177,1469,1470, 871,2212,2213,2214,2215,2216,1471,2217,2218,2219, -2220,2221,2222,2223,2224,2225,2226,2227,1472,1236,2228,2229,2230,2231,2232,2233, -2234,2235,1299,2236,2237, 200,2238, 477, 373,2239,2240, 731, 825, 777,2241,2242, -2243, 521, 486, 548,2244,2245,2246,1473,1300, 53, 549, 137, 875, 76, 158,2247, -1301,1474, 469, 396,1016, 278, 712,2248, 321, 442, 503, 767, 744, 941,1237,1178, -1475,2249, 82, 178,1141,1179, 973,2250,1302,2251, 297,2252,2253, 570,2254,2255, -2256, 18, 450, 206,2257, 290, 292,1142,2258, 511, 162, 99, 346, 164, 735,2259, -1476,1477, 4, 554, 343, 798,1099,2260,1100,2261, 43, 171,1303, 139, 215,2262, -2263, 717, 775,2264,1033, 322, 216,2265, 831,2266, 149,2267,1304,2268,2269, 702, -1238, 135, 845, 347, 309,2270, 484,2271, 878, 655, 238,1006,1478,2272, 67,2273, - 295,2274,2275, 461,2276, 478, 942, 412,2277,1034,2278,2279,2280, 265,2281, 541, -2282,2283,2284,2285,2286, 70, 852,1071,2287,2288,2289,2290, 21, 56, 509, 117, - 432,2291,2292, 331, 980, 552,1101, 148, 284, 105, 393,1180,1239, 755,2293, 187, -2294,1046,1479,2295, 340,2296, 63,1047, 230,2297,2298,1305, 763,1306, 101, 800, - 808, 494,2299,2300,2301, 903,2302, 37,1072, 14, 5,2303, 79, 675,2304, 312, -2305,2306,2307,2308,2309,1480, 6,1307,2310,2311,2312, 1, 470, 35, 24, 229, -2313, 695, 210, 86, 778, 15, 784, 592, 779, 32, 77, 855, 964,2314, 259,2315, - 501, 380,2316,2317, 83, 981, 153, 689,1308,1481,1482,1483,2318,2319, 716,1484, -2320,2321,2322,2323,2324,2325,1485,2326,2327, 128, 57, 68, 261,1048, 211, 170, -1240, 31,2328, 51, 435, 742,2329,2330,2331, 635,2332, 264, 456,2333,2334,2335, - 425,2336,1486, 143, 507, 263, 943,2337, 363, 920,1487, 256,1488,1102, 243, 601, -1489,2338,2339,2340,2341,2342,2343,2344, 861,2345,2346,2347,2348,2349,2350, 395, -2351,1490,1491, 62, 535, 166, 225,2352,2353, 668, 419,1241, 138, 604, 928,2354, -1181,2355,1492,1493,2356,2357,2358,1143,2359, 696,2360, 387, 307,1309, 682, 476, -2361,2362, 332, 12, 222, 156,2363, 232,2364, 641, 276, 656, 517,1494,1495,1035, - 416, 736,1496,2365,1017, 586,2366,2367,2368,1497,2369, 242,2370,2371,2372,1498, -2373, 965, 713,2374,2375,2376,2377, 740, 982,1499, 944,1500,1007,2378,2379,1310, -1501,2380,2381,2382, 785, 329,2383,2384,1502,2385,2386,2387, 932,2388,1503,2389, -2390,2391,2392,1242,2393,2394,2395,2396,2397, 994, 950,2398,2399,2400,2401,1504, -1311,2402,2403,2404,2405,1049, 749,2406,2407, 853, 718,1144,1312,2408,1182,1505, -2409,2410, 255, 516, 479, 564, 550, 214,1506,1507,1313, 413, 239, 444, 339,1145, -1036,1508,1509,1314,1037,1510,1315,2411,1511,2412,2413,2414, 176, 703, 497, 624, - 593, 921, 302,2415, 341, 165,1103,1512,2416,1513,2417,2418,2419, 376,2420, 700, -2421,2422,2423, 258, 768,1316,2424,1183,2425, 995, 608,2426,2427,2428,2429, 221, -2430,2431,2432,2433,2434,2435,2436,2437, 195, 323, 726, 188, 897, 983,1317, 377, - 644,1050, 879,2438, 452,2439,2440,2441,2442,2443,2444, 914,2445,2446,2447,2448, - 915, 489,2449,1514,1184,2450,2451, 515, 64, 427, 495,2452, 583,2453, 483, 485, -1038, 562, 213,1515, 748, 666,2454,2455,2456,2457, 334,2458, 780, 996,1008, 705, -1243,2459,2460,2461,2462,2463, 114,2464, 493,1146, 366, 163,1516, 961,1104,2465, - 291,2466,1318,1105,2467,1517, 365,2468, 355, 951,1244,2469,1319,2470, 631,2471, -2472, 218,1320, 364, 320, 756,1518,1519,1321,1520,1322,2473,2474,2475,2476, 997, -2477,2478,2479,2480, 665,1185,2481, 916,1521,2482,2483,2484, 584, 684,2485,2486, - 797,2487,1051,1186,2488,2489,2490,1522,2491,2492, 370,2493,1039,1187, 65,2494, - 434, 205, 463,1188,2495, 125, 812, 391, 402, 826, 699, 286, 398, 155, 781, 771, - 585,2496, 590, 505,1073,2497, 599, 244, 219, 917,1018, 952, 646,1523,2498,1323, -2499,2500, 49, 984, 354, 741,2501, 625,2502,1324,2503,1019, 190, 357, 757, 491, - 95, 782, 868,2504,2505,2506,2507,2508,2509, 134,1524,1074, 422,1525, 898,2510, - 161,2511,2512,2513,2514, 769,2515,1526,2516,2517, 411,1325,2518, 472,1527,2519, -2520,2521,2522,2523,2524, 985,2525,2526,2527,2528,2529,2530, 764,2531,1245,2532, -2533, 25, 204, 311,2534, 496,2535,1052,2536,2537,2538,2539,2540,2541,2542, 199, - 704, 504, 468, 758, 657,1528, 196, 44, 839,1246, 272, 750,2543, 765, 862,2544, -2545,1326,2546, 132, 615, 933,2547, 732,2548,2549,2550,1189,1529,2551, 283,1247, -1053, 607, 929,2552,2553,2554, 930, 183, 872, 616,1040,1147,2555,1148,1020, 441, - 249,1075,2556,2557,2558, 466, 743,2559,2560,2561, 92, 514, 426, 420, 526,2562, -2563,2564,2565,2566,2567,2568, 185,2569,2570,2571,2572, 776,1530, 658,2573, 362, -2574, 361, 922,1076, 793,2575,2576,2577,2578,2579,2580,1531, 251,2581,2582,2583, -2584,1532, 54, 612, 237,1327,2585,2586, 275, 408, 647, 111,2587,1533,1106, 465, - 3, 458, 9, 38,2588, 107, 110, 890, 209, 26, 737, 498,2589,1534,2590, 431, - 202, 88,1535, 356, 287,1107, 660,1149,2591, 381,1536, 986,1150, 445,1248,1151, - 974,2592,2593, 846,2594, 446, 953, 184,1249,1250, 727,2595, 923, 193, 883,2596, -2597,2598, 102, 324, 539, 817,2599, 421,1041,2600, 832,2601, 94, 175, 197, 406, -2602, 459,2603,2604,2605,2606,2607, 330, 555,2608,2609,2610, 706,1108, 389,2611, -2612,2613,2614, 233,2615, 833, 558, 931, 954,1251,2616,2617,1537, 546,2618,2619, -1009,2620,2621,2622,1538, 690,1328,2623, 955,2624,1539,2625,2626, 772,2627,2628, -2629,2630,2631, 924, 648, 863, 603,2632,2633, 934,1540, 864, 865,2634, 642,1042, - 670,1190,2635,2636,2637,2638, 168,2639, 652, 873, 542,1054,1541,2640,2641,2642, # 512, 256 -) -# fmt: on diff --git a/venv/lib/python3.10/site-packages/chardet/euckrprober.py b/venv/lib/python3.10/site-packages/chardet/euckrprober.py deleted file mode 100644 index 1fc5de0462cd9a09472cece4087cafe699da4fa7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/euckrprober.py +++ /dev/null @@ -1,47 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is mozilla.org code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from .chardistribution import EUCKRDistributionAnalysis -from .codingstatemachine import CodingStateMachine -from .mbcharsetprober import MultiByteCharSetProber -from .mbcssm import EUCKR_SM_MODEL - - -class EUCKRProber(MultiByteCharSetProber): - def __init__(self) -> None: - super().__init__() - self.coding_sm = CodingStateMachine(EUCKR_SM_MODEL) - self.distribution_analyzer = EUCKRDistributionAnalysis() - self.reset() - - @property - def charset_name(self) -> str: - return "EUC-KR" - - @property - def language(self) -> str: - return "Korean" diff --git a/venv/lib/python3.10/site-packages/chardet/euctwfreq.py b/venv/lib/python3.10/site-packages/chardet/euctwfreq.py deleted file mode 100644 index 4900ccc160a1dbf4de3a01c234735c21dd4417d6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/euctwfreq.py +++ /dev/null @@ -1,388 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Communicator client code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -# EUCTW frequency table -# Converted from big5 work -# by Taiwan's Mandarin Promotion Council -# - -# 128 --> 0.42261 -# 256 --> 0.57851 -# 512 --> 0.74851 -# 1024 --> 0.89384 -# 2048 --> 0.97583 -# -# Idea Distribution Ratio = 0.74851/(1-0.74851) =2.98 -# Random Distribution Ration = 512/(5401-512)=0.105 -# -# Typical Distribution Ratio about 25% of Ideal one, still much higher than RDR - -EUCTW_TYPICAL_DISTRIBUTION_RATIO = 0.75 - -# Char to FreqOrder table -EUCTW_TABLE_SIZE = 5376 - -# fmt: off -EUCTW_CHAR_TO_FREQ_ORDER = ( - 1, 1800, 1506, 255, 1431, 198, 9, 82, 6, 7310, 177, 202, 3615, 1256, 2808, 110, # 2742 - 3735, 33, 3241, 261, 76, 44, 2113, 16, 2931, 2184, 1176, 659, 3868, 26, 3404, 2643, # 2758 - 1198, 3869, 3313, 4060, 410, 2211, 302, 590, 361, 1963, 8, 204, 58, 4296, 7311, 1931, # 2774 - 63, 7312, 7313, 317, 1614, 75, 222, 159, 4061, 2412, 1480, 7314, 3500, 3068, 224, 2809, # 2790 - 3616, 3, 10, 3870, 1471, 29, 2774, 1135, 2852, 1939, 873, 130, 3242, 1123, 312, 7315, # 2806 - 4297, 2051, 507, 252, 682, 7316, 142, 1914, 124, 206, 2932, 34, 3501, 3173, 64, 604, # 2822 - 7317, 2494, 1976, 1977, 155, 1990, 645, 641, 1606, 7318, 3405, 337, 72, 406, 7319, 80, # 2838 - 630, 238, 3174, 1509, 263, 939, 1092, 2644, 756, 1440, 1094, 3406, 449, 69, 2969, 591, # 2854 - 179, 2095, 471, 115, 2034, 1843, 60, 50, 2970, 134, 806, 1868, 734, 2035, 3407, 180, # 2870 - 995, 1607, 156, 537, 2893, 688, 7320, 319, 1305, 779, 2144, 514, 2374, 298, 4298, 359, # 2886 - 2495, 90, 2707, 1338, 663, 11, 906, 1099, 2545, 20, 2436, 182, 532, 1716, 7321, 732, # 2902 - 1376, 4062, 1311, 1420, 3175, 25, 2312, 1056, 113, 399, 382, 1949, 242, 3408, 2467, 529, # 2918 - 3243, 475, 1447, 3617, 7322, 117, 21, 656, 810, 1297, 2295, 2329, 3502, 7323, 126, 4063, # 2934 - 706, 456, 150, 613, 4299, 71, 1118, 2036, 4064, 145, 3069, 85, 835, 486, 2114, 1246, # 2950 - 1426, 428, 727, 1285, 1015, 800, 106, 623, 303, 1281, 7324, 2127, 2354, 347, 3736, 221, # 2966 - 3503, 3110, 7325, 1955, 1153, 4065, 83, 296, 1199, 3070, 192, 624, 93, 7326, 822, 1897, # 2982 - 2810, 3111, 795, 2064, 991, 1554, 1542, 1592, 27, 43, 2853, 859, 139, 1456, 860, 4300, # 2998 - 437, 712, 3871, 164, 2392, 3112, 695, 211, 3017, 2096, 195, 3872, 1608, 3504, 3505, 3618, # 3014 - 3873, 234, 811, 2971, 2097, 3874, 2229, 1441, 3506, 1615, 2375, 668, 2076, 1638, 305, 228, # 3030 - 1664, 4301, 467, 415, 7327, 262, 2098, 1593, 239, 108, 300, 200, 1033, 512, 1247, 2077, # 3046 - 7328, 7329, 2173, 3176, 3619, 2673, 593, 845, 1062, 3244, 88, 1723, 2037, 3875, 1950, 212, # 3062 - 266, 152, 149, 468, 1898, 4066, 4302, 77, 187, 7330, 3018, 37, 5, 2972, 7331, 3876, # 3078 - 7332, 7333, 39, 2517, 4303, 2894, 3177, 2078, 55, 148, 74, 4304, 545, 483, 1474, 1029, # 3094 - 1665, 217, 1869, 1531, 3113, 1104, 2645, 4067, 24, 172, 3507, 900, 3877, 3508, 3509, 4305, # 3110 - 32, 1408, 2811, 1312, 329, 487, 2355, 2247, 2708, 784, 2674, 4, 3019, 3314, 1427, 1788, # 3126 - 188, 109, 499, 7334, 3620, 1717, 1789, 888, 1217, 3020, 4306, 7335, 3510, 7336, 3315, 1520, # 3142 - 3621, 3878, 196, 1034, 775, 7337, 7338, 929, 1815, 249, 439, 38, 7339, 1063, 7340, 794, # 3158 - 3879, 1435, 2296, 46, 178, 3245, 2065, 7341, 2376, 7342, 214, 1709, 4307, 804, 35, 707, # 3174 - 324, 3622, 1601, 2546, 140, 459, 4068, 7343, 7344, 1365, 839, 272, 978, 2257, 2572, 3409, # 3190 - 2128, 1363, 3623, 1423, 697, 100, 3071, 48, 70, 1231, 495, 3114, 2193, 7345, 1294, 7346, # 3206 - 2079, 462, 586, 1042, 3246, 853, 256, 988, 185, 2377, 3410, 1698, 434, 1084, 7347, 3411, # 3222 - 314, 2615, 2775, 4308, 2330, 2331, 569, 2280, 637, 1816, 2518, 757, 1162, 1878, 1616, 3412, # 3238 - 287, 1577, 2115, 768, 4309, 1671, 2854, 3511, 2519, 1321, 3737, 909, 2413, 7348, 4069, 933, # 3254 - 3738, 7349, 2052, 2356, 1222, 4310, 765, 2414, 1322, 786, 4311, 7350, 1919, 1462, 1677, 2895, # 3270 - 1699, 7351, 4312, 1424, 2437, 3115, 3624, 2590, 3316, 1774, 1940, 3413, 3880, 4070, 309, 1369, # 3286 - 1130, 2812, 364, 2230, 1653, 1299, 3881, 3512, 3882, 3883, 2646, 525, 1085, 3021, 902, 2000, # 3302 - 1475, 964, 4313, 421, 1844, 1415, 1057, 2281, 940, 1364, 3116, 376, 4314, 4315, 1381, 7, # 3318 - 2520, 983, 2378, 336, 1710, 2675, 1845, 321, 3414, 559, 1131, 3022, 2742, 1808, 1132, 1313, # 3334 - 265, 1481, 1857, 7352, 352, 1203, 2813, 3247, 167, 1089, 420, 2814, 776, 792, 1724, 3513, # 3350 - 4071, 2438, 3248, 7353, 4072, 7354, 446, 229, 333, 2743, 901, 3739, 1200, 1557, 4316, 2647, # 3366 - 1920, 395, 2744, 2676, 3740, 4073, 1835, 125, 916, 3178, 2616, 4317, 7355, 7356, 3741, 7357, # 3382 - 7358, 7359, 4318, 3117, 3625, 1133, 2547, 1757, 3415, 1510, 2313, 1409, 3514, 7360, 2145, 438, # 3398 - 2591, 2896, 2379, 3317, 1068, 958, 3023, 461, 311, 2855, 2677, 4074, 1915, 3179, 4075, 1978, # 3414 - 383, 750, 2745, 2617, 4076, 274, 539, 385, 1278, 1442, 7361, 1154, 1964, 384, 561, 210, # 3430 - 98, 1295, 2548, 3515, 7362, 1711, 2415, 1482, 3416, 3884, 2897, 1257, 129, 7363, 3742, 642, # 3446 - 523, 2776, 2777, 2648, 7364, 141, 2231, 1333, 68, 176, 441, 876, 907, 4077, 603, 2592, # 3462 - 710, 171, 3417, 404, 549, 18, 3118, 2393, 1410, 3626, 1666, 7365, 3516, 4319, 2898, 4320, # 3478 - 7366, 2973, 368, 7367, 146, 366, 99, 871, 3627, 1543, 748, 807, 1586, 1185, 22, 2258, # 3494 - 379, 3743, 3180, 7368, 3181, 505, 1941, 2618, 1991, 1382, 2314, 7369, 380, 2357, 218, 702, # 3510 - 1817, 1248, 3418, 3024, 3517, 3318, 3249, 7370, 2974, 3628, 930, 3250, 3744, 7371, 59, 7372, # 3526 - 585, 601, 4078, 497, 3419, 1112, 1314, 4321, 1801, 7373, 1223, 1472, 2174, 7374, 749, 1836, # 3542 - 690, 1899, 3745, 1772, 3885, 1476, 429, 1043, 1790, 2232, 2116, 917, 4079, 447, 1086, 1629, # 3558 - 7375, 556, 7376, 7377, 2020, 1654, 844, 1090, 105, 550, 966, 1758, 2815, 1008, 1782, 686, # 3574 - 1095, 7378, 2282, 793, 1602, 7379, 3518, 2593, 4322, 4080, 2933, 2297, 4323, 3746, 980, 2496, # 3590 - 544, 353, 527, 4324, 908, 2678, 2899, 7380, 381, 2619, 1942, 1348, 7381, 1341, 1252, 560, # 3606 - 3072, 7382, 3420, 2856, 7383, 2053, 973, 886, 2080, 143, 4325, 7384, 7385, 157, 3886, 496, # 3622 - 4081, 57, 840, 540, 2038, 4326, 4327, 3421, 2117, 1445, 970, 2259, 1748, 1965, 2081, 4082, # 3638 - 3119, 1234, 1775, 3251, 2816, 3629, 773, 1206, 2129, 1066, 2039, 1326, 3887, 1738, 1725, 4083, # 3654 - 279, 3120, 51, 1544, 2594, 423, 1578, 2130, 2066, 173, 4328, 1879, 7386, 7387, 1583, 264, # 3670 - 610, 3630, 4329, 2439, 280, 154, 7388, 7389, 7390, 1739, 338, 1282, 3073, 693, 2857, 1411, # 3686 - 1074, 3747, 2440, 7391, 4330, 7392, 7393, 1240, 952, 2394, 7394, 2900, 1538, 2679, 685, 1483, # 3702 - 4084, 2468, 1436, 953, 4085, 2054, 4331, 671, 2395, 79, 4086, 2441, 3252, 608, 567, 2680, # 3718 - 3422, 4087, 4088, 1691, 393, 1261, 1791, 2396, 7395, 4332, 7396, 7397, 7398, 7399, 1383, 1672, # 3734 - 3748, 3182, 1464, 522, 1119, 661, 1150, 216, 675, 4333, 3888, 1432, 3519, 609, 4334, 2681, # 3750 - 2397, 7400, 7401, 7402, 4089, 3025, 0, 7403, 2469, 315, 231, 2442, 301, 3319, 4335, 2380, # 3766 - 7404, 233, 4090, 3631, 1818, 4336, 4337, 7405, 96, 1776, 1315, 2082, 7406, 257, 7407, 1809, # 3782 - 3632, 2709, 1139, 1819, 4091, 2021, 1124, 2163, 2778, 1777, 2649, 7408, 3074, 363, 1655, 3183, # 3798 - 7409, 2975, 7410, 7411, 7412, 3889, 1567, 3890, 718, 103, 3184, 849, 1443, 341, 3320, 2934, # 3814 - 1484, 7413, 1712, 127, 67, 339, 4092, 2398, 679, 1412, 821, 7414, 7415, 834, 738, 351, # 3830 - 2976, 2146, 846, 235, 1497, 1880, 418, 1992, 3749, 2710, 186, 1100, 2147, 2746, 3520, 1545, # 3846 - 1355, 2935, 2858, 1377, 583, 3891, 4093, 2573, 2977, 7416, 1298, 3633, 1078, 2549, 3634, 2358, # 3862 - 78, 3750, 3751, 267, 1289, 2099, 2001, 1594, 4094, 348, 369, 1274, 2194, 2175, 1837, 4338, # 3878 - 1820, 2817, 3635, 2747, 2283, 2002, 4339, 2936, 2748, 144, 3321, 882, 4340, 3892, 2749, 3423, # 3894 - 4341, 2901, 7417, 4095, 1726, 320, 7418, 3893, 3026, 788, 2978, 7419, 2818, 1773, 1327, 2859, # 3910 - 3894, 2819, 7420, 1306, 4342, 2003, 1700, 3752, 3521, 2359, 2650, 787, 2022, 506, 824, 3636, # 3926 - 534, 323, 4343, 1044, 3322, 2023, 1900, 946, 3424, 7421, 1778, 1500, 1678, 7422, 1881, 4344, # 3942 - 165, 243, 4345, 3637, 2521, 123, 683, 4096, 764, 4346, 36, 3895, 1792, 589, 2902, 816, # 3958 - 626, 1667, 3027, 2233, 1639, 1555, 1622, 3753, 3896, 7423, 3897, 2860, 1370, 1228, 1932, 891, # 3974 - 2083, 2903, 304, 4097, 7424, 292, 2979, 2711, 3522, 691, 2100, 4098, 1115, 4347, 118, 662, # 3990 - 7425, 611, 1156, 854, 2381, 1316, 2861, 2, 386, 515, 2904, 7426, 7427, 3253, 868, 2234, # 4006 - 1486, 855, 2651, 785, 2212, 3028, 7428, 1040, 3185, 3523, 7429, 3121, 448, 7430, 1525, 7431, # 4022 - 2164, 4348, 7432, 3754, 7433, 4099, 2820, 3524, 3122, 503, 818, 3898, 3123, 1568, 814, 676, # 4038 - 1444, 306, 1749, 7434, 3755, 1416, 1030, 197, 1428, 805, 2821, 1501, 4349, 7435, 7436, 7437, # 4054 - 1993, 7438, 4350, 7439, 7440, 2195, 13, 2779, 3638, 2980, 3124, 1229, 1916, 7441, 3756, 2131, # 4070 - 7442, 4100, 4351, 2399, 3525, 7443, 2213, 1511, 1727, 1120, 7444, 7445, 646, 3757, 2443, 307, # 4086 - 7446, 7447, 1595, 3186, 7448, 7449, 7450, 3639, 1113, 1356, 3899, 1465, 2522, 2523, 7451, 519, # 4102 - 7452, 128, 2132, 92, 2284, 1979, 7453, 3900, 1512, 342, 3125, 2196, 7454, 2780, 2214, 1980, # 4118 - 3323, 7455, 290, 1656, 1317, 789, 827, 2360, 7456, 3758, 4352, 562, 581, 3901, 7457, 401, # 4134 - 4353, 2248, 94, 4354, 1399, 2781, 7458, 1463, 2024, 4355, 3187, 1943, 7459, 828, 1105, 4101, # 4150 - 1262, 1394, 7460, 4102, 605, 4356, 7461, 1783, 2862, 7462, 2822, 819, 2101, 578, 2197, 2937, # 4166 - 7463, 1502, 436, 3254, 4103, 3255, 2823, 3902, 2905, 3425, 3426, 7464, 2712, 2315, 7465, 7466, # 4182 - 2332, 2067, 23, 4357, 193, 826, 3759, 2102, 699, 1630, 4104, 3075, 390, 1793, 1064, 3526, # 4198 - 7467, 1579, 3076, 3077, 1400, 7468, 4105, 1838, 1640, 2863, 7469, 4358, 4359, 137, 4106, 598, # 4214 - 3078, 1966, 780, 104, 974, 2938, 7470, 278, 899, 253, 402, 572, 504, 493, 1339, 7471, # 4230 - 3903, 1275, 4360, 2574, 2550, 7472, 3640, 3029, 3079, 2249, 565, 1334, 2713, 863, 41, 7473, # 4246 - 7474, 4361, 7475, 1657, 2333, 19, 463, 2750, 4107, 606, 7476, 2981, 3256, 1087, 2084, 1323, # 4262 - 2652, 2982, 7477, 1631, 1623, 1750, 4108, 2682, 7478, 2864, 791, 2714, 2653, 2334, 232, 2416, # 4278 - 7479, 2983, 1498, 7480, 2654, 2620, 755, 1366, 3641, 3257, 3126, 2025, 1609, 119, 1917, 3427, # 4294 - 862, 1026, 4109, 7481, 3904, 3760, 4362, 3905, 4363, 2260, 1951, 2470, 7482, 1125, 817, 4110, # 4310 - 4111, 3906, 1513, 1766, 2040, 1487, 4112, 3030, 3258, 2824, 3761, 3127, 7483, 7484, 1507, 7485, # 4326 - 2683, 733, 40, 1632, 1106, 2865, 345, 4113, 841, 2524, 230, 4364, 2984, 1846, 3259, 3428, # 4342 - 7486, 1263, 986, 3429, 7487, 735, 879, 254, 1137, 857, 622, 1300, 1180, 1388, 1562, 3907, # 4358 - 3908, 2939, 967, 2751, 2655, 1349, 592, 2133, 1692, 3324, 2985, 1994, 4114, 1679, 3909, 1901, # 4374 - 2185, 7488, 739, 3642, 2715, 1296, 1290, 7489, 4115, 2198, 2199, 1921, 1563, 2595, 2551, 1870, # 4390 - 2752, 2986, 7490, 435, 7491, 343, 1108, 596, 17, 1751, 4365, 2235, 3430, 3643, 7492, 4366, # 4406 - 294, 3527, 2940, 1693, 477, 979, 281, 2041, 3528, 643, 2042, 3644, 2621, 2782, 2261, 1031, # 4422 - 2335, 2134, 2298, 3529, 4367, 367, 1249, 2552, 7493, 3530, 7494, 4368, 1283, 3325, 2004, 240, # 4438 - 1762, 3326, 4369, 4370, 836, 1069, 3128, 474, 7495, 2148, 2525, 268, 3531, 7496, 3188, 1521, # 4454 - 1284, 7497, 1658, 1546, 4116, 7498, 3532, 3533, 7499, 4117, 3327, 2684, 1685, 4118, 961, 1673, # 4470 - 2622, 190, 2005, 2200, 3762, 4371, 4372, 7500, 570, 2497, 3645, 1490, 7501, 4373, 2623, 3260, # 4486 - 1956, 4374, 584, 1514, 396, 1045, 1944, 7502, 4375, 1967, 2444, 7503, 7504, 4376, 3910, 619, # 4502 - 7505, 3129, 3261, 215, 2006, 2783, 2553, 3189, 4377, 3190, 4378, 763, 4119, 3763, 4379, 7506, # 4518 - 7507, 1957, 1767, 2941, 3328, 3646, 1174, 452, 1477, 4380, 3329, 3130, 7508, 2825, 1253, 2382, # 4534 - 2186, 1091, 2285, 4120, 492, 7509, 638, 1169, 1824, 2135, 1752, 3911, 648, 926, 1021, 1324, # 4550 - 4381, 520, 4382, 997, 847, 1007, 892, 4383, 3764, 2262, 1871, 3647, 7510, 2400, 1784, 4384, # 4566 - 1952, 2942, 3080, 3191, 1728, 4121, 2043, 3648, 4385, 2007, 1701, 3131, 1551, 30, 2263, 4122, # 4582 - 7511, 2026, 4386, 3534, 7512, 501, 7513, 4123, 594, 3431, 2165, 1821, 3535, 3432, 3536, 3192, # 4598 - 829, 2826, 4124, 7514, 1680, 3132, 1225, 4125, 7515, 3262, 4387, 4126, 3133, 2336, 7516, 4388, # 4614 - 4127, 7517, 3912, 3913, 7518, 1847, 2383, 2596, 3330, 7519, 4389, 374, 3914, 652, 4128, 4129, # 4630 - 375, 1140, 798, 7520, 7521, 7522, 2361, 4390, 2264, 546, 1659, 138, 3031, 2445, 4391, 7523, # 4646 - 2250, 612, 1848, 910, 796, 3765, 1740, 1371, 825, 3766, 3767, 7524, 2906, 2554, 7525, 692, # 4662 - 444, 3032, 2624, 801, 4392, 4130, 7526, 1491, 244, 1053, 3033, 4131, 4132, 340, 7527, 3915, # 4678 - 1041, 2987, 293, 1168, 87, 1357, 7528, 1539, 959, 7529, 2236, 721, 694, 4133, 3768, 219, # 4694 - 1478, 644, 1417, 3331, 2656, 1413, 1401, 1335, 1389, 3916, 7530, 7531, 2988, 2362, 3134, 1825, # 4710 - 730, 1515, 184, 2827, 66, 4393, 7532, 1660, 2943, 246, 3332, 378, 1457, 226, 3433, 975, # 4726 - 3917, 2944, 1264, 3537, 674, 696, 7533, 163, 7534, 1141, 2417, 2166, 713, 3538, 3333, 4394, # 4742 - 3918, 7535, 7536, 1186, 15, 7537, 1079, 1070, 7538, 1522, 3193, 3539, 276, 1050, 2716, 758, # 4758 - 1126, 653, 2945, 3263, 7539, 2337, 889, 3540, 3919, 3081, 2989, 903, 1250, 4395, 3920, 3434, # 4774 - 3541, 1342, 1681, 1718, 766, 3264, 286, 89, 2946, 3649, 7540, 1713, 7541, 2597, 3334, 2990, # 4790 - 7542, 2947, 2215, 3194, 2866, 7543, 4396, 2498, 2526, 181, 387, 1075, 3921, 731, 2187, 3335, # 4806 - 7544, 3265, 310, 313, 3435, 2299, 770, 4134, 54, 3034, 189, 4397, 3082, 3769, 3922, 7545, # 4822 - 1230, 1617, 1849, 355, 3542, 4135, 4398, 3336, 111, 4136, 3650, 1350, 3135, 3436, 3035, 4137, # 4838 - 2149, 3266, 3543, 7546, 2784, 3923, 3924, 2991, 722, 2008, 7547, 1071, 247, 1207, 2338, 2471, # 4854 - 1378, 4399, 2009, 864, 1437, 1214, 4400, 373, 3770, 1142, 2216, 667, 4401, 442, 2753, 2555, # 4870 - 3771, 3925, 1968, 4138, 3267, 1839, 837, 170, 1107, 934, 1336, 1882, 7548, 7549, 2118, 4139, # 4886 - 2828, 743, 1569, 7550, 4402, 4140, 582, 2384, 1418, 3437, 7551, 1802, 7552, 357, 1395, 1729, # 4902 - 3651, 3268, 2418, 1564, 2237, 7553, 3083, 3772, 1633, 4403, 1114, 2085, 4141, 1532, 7554, 482, # 4918 - 2446, 4404, 7555, 7556, 1492, 833, 1466, 7557, 2717, 3544, 1641, 2829, 7558, 1526, 1272, 3652, # 4934 - 4142, 1686, 1794, 416, 2556, 1902, 1953, 1803, 7559, 3773, 2785, 3774, 1159, 2316, 7560, 2867, # 4950 - 4405, 1610, 1584, 3036, 2419, 2754, 443, 3269, 1163, 3136, 7561, 7562, 3926, 7563, 4143, 2499, # 4966 - 3037, 4406, 3927, 3137, 2103, 1647, 3545, 2010, 1872, 4144, 7564, 4145, 431, 3438, 7565, 250, # 4982 - 97, 81, 4146, 7566, 1648, 1850, 1558, 160, 848, 7567, 866, 740, 1694, 7568, 2201, 2830, # 4998 - 3195, 4147, 4407, 3653, 1687, 950, 2472, 426, 469, 3196, 3654, 3655, 3928, 7569, 7570, 1188, # 5014 - 424, 1995, 861, 3546, 4148, 3775, 2202, 2685, 168, 1235, 3547, 4149, 7571, 2086, 1674, 4408, # 5030 - 3337, 3270, 220, 2557, 1009, 7572, 3776, 670, 2992, 332, 1208, 717, 7573, 7574, 3548, 2447, # 5046 - 3929, 3338, 7575, 513, 7576, 1209, 2868, 3339, 3138, 4409, 1080, 7577, 7578, 7579, 7580, 2527, # 5062 - 3656, 3549, 815, 1587, 3930, 3931, 7581, 3550, 3439, 3777, 1254, 4410, 1328, 3038, 1390, 3932, # 5078 - 1741, 3933, 3778, 3934, 7582, 236, 3779, 2448, 3271, 7583, 7584, 3657, 3780, 1273, 3781, 4411, # 5094 - 7585, 308, 7586, 4412, 245, 4413, 1851, 2473, 1307, 2575, 430, 715, 2136, 2449, 7587, 270, # 5110 - 199, 2869, 3935, 7588, 3551, 2718, 1753, 761, 1754, 725, 1661, 1840, 4414, 3440, 3658, 7589, # 5126 - 7590, 587, 14, 3272, 227, 2598, 326, 480, 2265, 943, 2755, 3552, 291, 650, 1883, 7591, # 5142 - 1702, 1226, 102, 1547, 62, 3441, 904, 4415, 3442, 1164, 4150, 7592, 7593, 1224, 1548, 2756, # 5158 - 391, 498, 1493, 7594, 1386, 1419, 7595, 2055, 1177, 4416, 813, 880, 1081, 2363, 566, 1145, # 5174 - 4417, 2286, 1001, 1035, 2558, 2599, 2238, 394, 1286, 7596, 7597, 2068, 7598, 86, 1494, 1730, # 5190 - 3936, 491, 1588, 745, 897, 2948, 843, 3340, 3937, 2757, 2870, 3273, 1768, 998, 2217, 2069, # 5206 - 397, 1826, 1195, 1969, 3659, 2993, 3341, 284, 7599, 3782, 2500, 2137, 2119, 1903, 7600, 3938, # 5222 - 2150, 3939, 4151, 1036, 3443, 1904, 114, 2559, 4152, 209, 1527, 7601, 7602, 2949, 2831, 2625, # 5238 - 2385, 2719, 3139, 812, 2560, 7603, 3274, 7604, 1559, 737, 1884, 3660, 1210, 885, 28, 2686, # 5254 - 3553, 3783, 7605, 4153, 1004, 1779, 4418, 7606, 346, 1981, 2218, 2687, 4419, 3784, 1742, 797, # 5270 - 1642, 3940, 1933, 1072, 1384, 2151, 896, 3941, 3275, 3661, 3197, 2871, 3554, 7607, 2561, 1958, # 5286 - 4420, 2450, 1785, 7608, 7609, 7610, 3942, 4154, 1005, 1308, 3662, 4155, 2720, 4421, 4422, 1528, # 5302 - 2600, 161, 1178, 4156, 1982, 987, 4423, 1101, 4157, 631, 3943, 1157, 3198, 2420, 1343, 1241, # 5318 - 1016, 2239, 2562, 372, 877, 2339, 2501, 1160, 555, 1934, 911, 3944, 7611, 466, 1170, 169, # 5334 - 1051, 2907, 2688, 3663, 2474, 2994, 1182, 2011, 2563, 1251, 2626, 7612, 992, 2340, 3444, 1540, # 5350 - 2721, 1201, 2070, 2401, 1996, 2475, 7613, 4424, 528, 1922, 2188, 1503, 1873, 1570, 2364, 3342, # 5366 - 3276, 7614, 557, 1073, 7615, 1827, 3445, 2087, 2266, 3140, 3039, 3084, 767, 3085, 2786, 4425, # 5382 - 1006, 4158, 4426, 2341, 1267, 2176, 3664, 3199, 778, 3945, 3200, 2722, 1597, 2657, 7616, 4427, # 5398 - 7617, 3446, 7618, 7619, 7620, 3277, 2689, 1433, 3278, 131, 95, 1504, 3946, 723, 4159, 3141, # 5414 - 1841, 3555, 2758, 2189, 3947, 2027, 2104, 3665, 7621, 2995, 3948, 1218, 7622, 3343, 3201, 3949, # 5430 - 4160, 2576, 248, 1634, 3785, 912, 7623, 2832, 3666, 3040, 3786, 654, 53, 7624, 2996, 7625, # 5446 - 1688, 4428, 777, 3447, 1032, 3950, 1425, 7626, 191, 820, 2120, 2833, 971, 4429, 931, 3202, # 5462 - 135, 664, 783, 3787, 1997, 772, 2908, 1935, 3951, 3788, 4430, 2909, 3203, 282, 2723, 640, # 5478 - 1372, 3448, 1127, 922, 325, 3344, 7627, 7628, 711, 2044, 7629, 7630, 3952, 2219, 2787, 1936, # 5494 - 3953, 3345, 2220, 2251, 3789, 2300, 7631, 4431, 3790, 1258, 3279, 3954, 3204, 2138, 2950, 3955, # 5510 - 3956, 7632, 2221, 258, 3205, 4432, 101, 1227, 7633, 3280, 1755, 7634, 1391, 3281, 7635, 2910, # 5526 - 2056, 893, 7636, 7637, 7638, 1402, 4161, 2342, 7639, 7640, 3206, 3556, 7641, 7642, 878, 1325, # 5542 - 1780, 2788, 4433, 259, 1385, 2577, 744, 1183, 2267, 4434, 7643, 3957, 2502, 7644, 684, 1024, # 5558 - 4162, 7645, 472, 3557, 3449, 1165, 3282, 3958, 3959, 322, 2152, 881, 455, 1695, 1152, 1340, # 5574 - 660, 554, 2153, 4435, 1058, 4436, 4163, 830, 1065, 3346, 3960, 4437, 1923, 7646, 1703, 1918, # 5590 - 7647, 932, 2268, 122, 7648, 4438, 947, 677, 7649, 3791, 2627, 297, 1905, 1924, 2269, 4439, # 5606 - 2317, 3283, 7650, 7651, 4164, 7652, 4165, 84, 4166, 112, 989, 7653, 547, 1059, 3961, 701, # 5622 - 3558, 1019, 7654, 4167, 7655, 3450, 942, 639, 457, 2301, 2451, 993, 2951, 407, 851, 494, # 5638 - 4440, 3347, 927, 7656, 1237, 7657, 2421, 3348, 573, 4168, 680, 921, 2911, 1279, 1874, 285, # 5654 - 790, 1448, 1983, 719, 2167, 7658, 7659, 4441, 3962, 3963, 1649, 7660, 1541, 563, 7661, 1077, # 5670 - 7662, 3349, 3041, 3451, 511, 2997, 3964, 3965, 3667, 3966, 1268, 2564, 3350, 3207, 4442, 4443, # 5686 - 7663, 535, 1048, 1276, 1189, 2912, 2028, 3142, 1438, 1373, 2834, 2952, 1134, 2012, 7664, 4169, # 5702 - 1238, 2578, 3086, 1259, 7665, 700, 7666, 2953, 3143, 3668, 4170, 7667, 4171, 1146, 1875, 1906, # 5718 - 4444, 2601, 3967, 781, 2422, 132, 1589, 203, 147, 273, 2789, 2402, 898, 1786, 2154, 3968, # 5734 - 3969, 7668, 3792, 2790, 7669, 7670, 4445, 4446, 7671, 3208, 7672, 1635, 3793, 965, 7673, 1804, # 5750 - 2690, 1516, 3559, 1121, 1082, 1329, 3284, 3970, 1449, 3794, 65, 1128, 2835, 2913, 2759, 1590, # 5766 - 3795, 7674, 7675, 12, 2658, 45, 976, 2579, 3144, 4447, 517, 2528, 1013, 1037, 3209, 7676, # 5782 - 3796, 2836, 7677, 3797, 7678, 3452, 7679, 2602, 614, 1998, 2318, 3798, 3087, 2724, 2628, 7680, # 5798 - 2580, 4172, 599, 1269, 7681, 1810, 3669, 7682, 2691, 3088, 759, 1060, 489, 1805, 3351, 3285, # 5814 - 1358, 7683, 7684, 2386, 1387, 1215, 2629, 2252, 490, 7685, 7686, 4173, 1759, 2387, 2343, 7687, # 5830 - 4448, 3799, 1907, 3971, 2630, 1806, 3210, 4449, 3453, 3286, 2760, 2344, 874, 7688, 7689, 3454, # 5846 - 3670, 1858, 91, 2914, 3671, 3042, 3800, 4450, 7690, 3145, 3972, 2659, 7691, 3455, 1202, 1403, # 5862 - 3801, 2954, 2529, 1517, 2503, 4451, 3456, 2504, 7692, 4452, 7693, 2692, 1885, 1495, 1731, 3973, # 5878 - 2365, 4453, 7694, 2029, 7695, 7696, 3974, 2693, 1216, 237, 2581, 4174, 2319, 3975, 3802, 4454, # 5894 - 4455, 2694, 3560, 3457, 445, 4456, 7697, 7698, 7699, 7700, 2761, 61, 3976, 3672, 1822, 3977, # 5910 - 7701, 687, 2045, 935, 925, 405, 2660, 703, 1096, 1859, 2725, 4457, 3978, 1876, 1367, 2695, # 5926 - 3352, 918, 2105, 1781, 2476, 334, 3287, 1611, 1093, 4458, 564, 3146, 3458, 3673, 3353, 945, # 5942 - 2631, 2057, 4459, 7702, 1925, 872, 4175, 7703, 3459, 2696, 3089, 349, 4176, 3674, 3979, 4460, # 5958 - 3803, 4177, 3675, 2155, 3980, 4461, 4462, 4178, 4463, 2403, 2046, 782, 3981, 400, 251, 4179, # 5974 - 1624, 7704, 7705, 277, 3676, 299, 1265, 476, 1191, 3804, 2121, 4180, 4181, 1109, 205, 7706, # 5990 - 2582, 1000, 2156, 3561, 1860, 7707, 7708, 7709, 4464, 7710, 4465, 2565, 107, 2477, 2157, 3982, # 6006 - 3460, 3147, 7711, 1533, 541, 1301, 158, 753, 4182, 2872, 3562, 7712, 1696, 370, 1088, 4183, # 6022 - 4466, 3563, 579, 327, 440, 162, 2240, 269, 1937, 1374, 3461, 968, 3043, 56, 1396, 3090, # 6038 - 2106, 3288, 3354, 7713, 1926, 2158, 4467, 2998, 7714, 3564, 7715, 7716, 3677, 4468, 2478, 7717, # 6054 - 2791, 7718, 1650, 4469, 7719, 2603, 7720, 7721, 3983, 2661, 3355, 1149, 3356, 3984, 3805, 3985, # 6070 - 7722, 1076, 49, 7723, 951, 3211, 3289, 3290, 450, 2837, 920, 7724, 1811, 2792, 2366, 4184, # 6086 - 1908, 1138, 2367, 3806, 3462, 7725, 3212, 4470, 1909, 1147, 1518, 2423, 4471, 3807, 7726, 4472, # 6102 - 2388, 2604, 260, 1795, 3213, 7727, 7728, 3808, 3291, 708, 7729, 3565, 1704, 7730, 3566, 1351, # 6118 - 1618, 3357, 2999, 1886, 944, 4185, 3358, 4186, 3044, 3359, 4187, 7731, 3678, 422, 413, 1714, # 6134 - 3292, 500, 2058, 2345, 4188, 2479, 7732, 1344, 1910, 954, 7733, 1668, 7734, 7735, 3986, 2404, # 6150 - 4189, 3567, 3809, 4190, 7736, 2302, 1318, 2505, 3091, 133, 3092, 2873, 4473, 629, 31, 2838, # 6166 - 2697, 3810, 4474, 850, 949, 4475, 3987, 2955, 1732, 2088, 4191, 1496, 1852, 7737, 3988, 620, # 6182 - 3214, 981, 1242, 3679, 3360, 1619, 3680, 1643, 3293, 2139, 2452, 1970, 1719, 3463, 2168, 7738, # 6198 - 3215, 7739, 7740, 3361, 1828, 7741, 1277, 4476, 1565, 2047, 7742, 1636, 3568, 3093, 7743, 869, # 6214 - 2839, 655, 3811, 3812, 3094, 3989, 3000, 3813, 1310, 3569, 4477, 7744, 7745, 7746, 1733, 558, # 6230 - 4478, 3681, 335, 1549, 3045, 1756, 4192, 3682, 1945, 3464, 1829, 1291, 1192, 470, 2726, 2107, # 6246 - 2793, 913, 1054, 3990, 7747, 1027, 7748, 3046, 3991, 4479, 982, 2662, 3362, 3148, 3465, 3216, # 6262 - 3217, 1946, 2794, 7749, 571, 4480, 7750, 1830, 7751, 3570, 2583, 1523, 2424, 7752, 2089, 984, # 6278 - 4481, 3683, 1959, 7753, 3684, 852, 923, 2795, 3466, 3685, 969, 1519, 999, 2048, 2320, 1705, # 6294 - 7754, 3095, 615, 1662, 151, 597, 3992, 2405, 2321, 1049, 275, 4482, 3686, 4193, 568, 3687, # 6310 - 3571, 2480, 4194, 3688, 7755, 2425, 2270, 409, 3218, 7756, 1566, 2874, 3467, 1002, 769, 2840, # 6326 - 194, 2090, 3149, 3689, 2222, 3294, 4195, 628, 1505, 7757, 7758, 1763, 2177, 3001, 3993, 521, # 6342 - 1161, 2584, 1787, 2203, 2406, 4483, 3994, 1625, 4196, 4197, 412, 42, 3096, 464, 7759, 2632, # 6358 - 4484, 3363, 1760, 1571, 2875, 3468, 2530, 1219, 2204, 3814, 2633, 2140, 2368, 4485, 4486, 3295, # 6374 - 1651, 3364, 3572, 7760, 7761, 3573, 2481, 3469, 7762, 3690, 7763, 7764, 2271, 2091, 460, 7765, # 6390 - 4487, 7766, 3002, 962, 588, 3574, 289, 3219, 2634, 1116, 52, 7767, 3047, 1796, 7768, 7769, # 6406 - 7770, 1467, 7771, 1598, 1143, 3691, 4198, 1984, 1734, 1067, 4488, 1280, 3365, 465, 4489, 1572, # 6422 - 510, 7772, 1927, 2241, 1812, 1644, 3575, 7773, 4490, 3692, 7774, 7775, 2663, 1573, 1534, 7776, # 6438 - 7777, 4199, 536, 1807, 1761, 3470, 3815, 3150, 2635, 7778, 7779, 7780, 4491, 3471, 2915, 1911, # 6454 - 2796, 7781, 3296, 1122, 377, 3220, 7782, 360, 7783, 7784, 4200, 1529, 551, 7785, 2059, 3693, # 6470 - 1769, 2426, 7786, 2916, 4201, 3297, 3097, 2322, 2108, 2030, 4492, 1404, 136, 1468, 1479, 672, # 6486 - 1171, 3221, 2303, 271, 3151, 7787, 2762, 7788, 2049, 678, 2727, 865, 1947, 4493, 7789, 2013, # 6502 - 3995, 2956, 7790, 2728, 2223, 1397, 3048, 3694, 4494, 4495, 1735, 2917, 3366, 3576, 7791, 3816, # 6518 - 509, 2841, 2453, 2876, 3817, 7792, 7793, 3152, 3153, 4496, 4202, 2531, 4497, 2304, 1166, 1010, # 6534 - 552, 681, 1887, 7794, 7795, 2957, 2958, 3996, 1287, 1596, 1861, 3154, 358, 453, 736, 175, # 6550 - 478, 1117, 905, 1167, 1097, 7796, 1853, 1530, 7797, 1706, 7798, 2178, 3472, 2287, 3695, 3473, # 6566 - 3577, 4203, 2092, 4204, 7799, 3367, 1193, 2482, 4205, 1458, 2190, 2205, 1862, 1888, 1421, 3298, # 6582 - 2918, 3049, 2179, 3474, 595, 2122, 7800, 3997, 7801, 7802, 4206, 1707, 2636, 223, 3696, 1359, # 6598 - 751, 3098, 183, 3475, 7803, 2797, 3003, 419, 2369, 633, 704, 3818, 2389, 241, 7804, 7805, # 6614 - 7806, 838, 3004, 3697, 2272, 2763, 2454, 3819, 1938, 2050, 3998, 1309, 3099, 2242, 1181, 7807, # 6630 - 1136, 2206, 3820, 2370, 1446, 4207, 2305, 4498, 7808, 7809, 4208, 1055, 2605, 484, 3698, 7810, # 6646 - 3999, 625, 4209, 2273, 3368, 1499, 4210, 4000, 7811, 4001, 4211, 3222, 2274, 2275, 3476, 7812, # 6662 - 7813, 2764, 808, 2606, 3699, 3369, 4002, 4212, 3100, 2532, 526, 3370, 3821, 4213, 955, 7814, # 6678 - 1620, 4214, 2637, 2427, 7815, 1429, 3700, 1669, 1831, 994, 928, 7816, 3578, 1260, 7817, 7818, # 6694 - 7819, 1948, 2288, 741, 2919, 1626, 4215, 2729, 2455, 867, 1184, 362, 3371, 1392, 7820, 7821, # 6710 - 4003, 4216, 1770, 1736, 3223, 2920, 4499, 4500, 1928, 2698, 1459, 1158, 7822, 3050, 3372, 2877, # 6726 - 1292, 1929, 2506, 2842, 3701, 1985, 1187, 2071, 2014, 2607, 4217, 7823, 2566, 2507, 2169, 3702, # 6742 - 2483, 3299, 7824, 3703, 4501, 7825, 7826, 666, 1003, 3005, 1022, 3579, 4218, 7827, 4502, 1813, # 6758 - 2253, 574, 3822, 1603, 295, 1535, 705, 3823, 4219, 283, 858, 417, 7828, 7829, 3224, 4503, # 6774 - 4504, 3051, 1220, 1889, 1046, 2276, 2456, 4004, 1393, 1599, 689, 2567, 388, 4220, 7830, 2484, # 6790 - 802, 7831, 2798, 3824, 2060, 1405, 2254, 7832, 4505, 3825, 2109, 1052, 1345, 3225, 1585, 7833, # 6806 - 809, 7834, 7835, 7836, 575, 2730, 3477, 956, 1552, 1469, 1144, 2323, 7837, 2324, 1560, 2457, # 6822 - 3580, 3226, 4005, 616, 2207, 3155, 2180, 2289, 7838, 1832, 7839, 3478, 4506, 7840, 1319, 3704, # 6838 - 3705, 1211, 3581, 1023, 3227, 1293, 2799, 7841, 7842, 7843, 3826, 607, 2306, 3827, 762, 2878, # 6854 - 1439, 4221, 1360, 7844, 1485, 3052, 7845, 4507, 1038, 4222, 1450, 2061, 2638, 4223, 1379, 4508, # 6870 - 2585, 7846, 7847, 4224, 1352, 1414, 2325, 2921, 1172, 7848, 7849, 3828, 3829, 7850, 1797, 1451, # 6886 - 7851, 7852, 7853, 7854, 2922, 4006, 4007, 2485, 2346, 411, 4008, 4009, 3582, 3300, 3101, 4509, # 6902 - 1561, 2664, 1452, 4010, 1375, 7855, 7856, 47, 2959, 316, 7857, 1406, 1591, 2923, 3156, 7858, # 6918 - 1025, 2141, 3102, 3157, 354, 2731, 884, 2224, 4225, 2407, 508, 3706, 726, 3583, 996, 2428, # 6934 - 3584, 729, 7859, 392, 2191, 1453, 4011, 4510, 3707, 7860, 7861, 2458, 3585, 2608, 1675, 2800, # 6950 - 919, 2347, 2960, 2348, 1270, 4511, 4012, 73, 7862, 7863, 647, 7864, 3228, 2843, 2255, 1550, # 6966 - 1346, 3006, 7865, 1332, 883, 3479, 7866, 7867, 7868, 7869, 3301, 2765, 7870, 1212, 831, 1347, # 6982 - 4226, 4512, 2326, 3830, 1863, 3053, 720, 3831, 4513, 4514, 3832, 7871, 4227, 7872, 7873, 4515, # 6998 - 7874, 7875, 1798, 4516, 3708, 2609, 4517, 3586, 1645, 2371, 7876, 7877, 2924, 669, 2208, 2665, # 7014 - 2429, 7878, 2879, 7879, 7880, 1028, 3229, 7881, 4228, 2408, 7882, 2256, 1353, 7883, 7884, 4518, # 7030 - 3158, 518, 7885, 4013, 7886, 4229, 1960, 7887, 2142, 4230, 7888, 7889, 3007, 2349, 2350, 3833, # 7046 - 516, 1833, 1454, 4014, 2699, 4231, 4519, 2225, 2610, 1971, 1129, 3587, 7890, 2766, 7891, 2961, # 7062 - 1422, 577, 1470, 3008, 1524, 3373, 7892, 7893, 432, 4232, 3054, 3480, 7894, 2586, 1455, 2508, # 7078 - 2226, 1972, 1175, 7895, 1020, 2732, 4015, 3481, 4520, 7896, 2733, 7897, 1743, 1361, 3055, 3482, # 7094 - 2639, 4016, 4233, 4521, 2290, 895, 924, 4234, 2170, 331, 2243, 3056, 166, 1627, 3057, 1098, # 7110 - 7898, 1232, 2880, 2227, 3374, 4522, 657, 403, 1196, 2372, 542, 3709, 3375, 1600, 4235, 3483, # 7126 - 7899, 4523, 2767, 3230, 576, 530, 1362, 7900, 4524, 2533, 2666, 3710, 4017, 7901, 842, 3834, # 7142 - 7902, 2801, 2031, 1014, 4018, 213, 2700, 3376, 665, 621, 4236, 7903, 3711, 2925, 2430, 7904, # 7158 - 2431, 3302, 3588, 3377, 7905, 4237, 2534, 4238, 4525, 3589, 1682, 4239, 3484, 1380, 7906, 724, # 7174 - 2277, 600, 1670, 7907, 1337, 1233, 4526, 3103, 2244, 7908, 1621, 4527, 7909, 651, 4240, 7910, # 7190 - 1612, 4241, 2611, 7911, 2844, 7912, 2734, 2307, 3058, 7913, 716, 2459, 3059, 174, 1255, 2701, # 7206 - 4019, 3590, 548, 1320, 1398, 728, 4020, 1574, 7914, 1890, 1197, 3060, 4021, 7915, 3061, 3062, # 7222 - 3712, 3591, 3713, 747, 7916, 635, 4242, 4528, 7917, 7918, 7919, 4243, 7920, 7921, 4529, 7922, # 7238 - 3378, 4530, 2432, 451, 7923, 3714, 2535, 2072, 4244, 2735, 4245, 4022, 7924, 1764, 4531, 7925, # 7254 - 4246, 350, 7926, 2278, 2390, 2486, 7927, 4247, 4023, 2245, 1434, 4024, 488, 4532, 458, 4248, # 7270 - 4025, 3715, 771, 1330, 2391, 3835, 2568, 3159, 2159, 2409, 1553, 2667, 3160, 4249, 7928, 2487, # 7286 - 2881, 2612, 1720, 2702, 4250, 3379, 4533, 7929, 2536, 4251, 7930, 3231, 4252, 2768, 7931, 2015, # 7302 - 2736, 7932, 1155, 1017, 3716, 3836, 7933, 3303, 2308, 201, 1864, 4253, 1430, 7934, 4026, 7935, # 7318 - 7936, 7937, 7938, 7939, 4254, 1604, 7940, 414, 1865, 371, 2587, 4534, 4535, 3485, 2016, 3104, # 7334 - 4536, 1708, 960, 4255, 887, 389, 2171, 1536, 1663, 1721, 7941, 2228, 4027, 2351, 2926, 1580, # 7350 - 7942, 7943, 7944, 1744, 7945, 2537, 4537, 4538, 7946, 4539, 7947, 2073, 7948, 7949, 3592, 3380, # 7366 - 2882, 4256, 7950, 4257, 2640, 3381, 2802, 673, 2703, 2460, 709, 3486, 4028, 3593, 4258, 7951, # 7382 - 1148, 502, 634, 7952, 7953, 1204, 4540, 3594, 1575, 4541, 2613, 3717, 7954, 3718, 3105, 948, # 7398 - 3232, 121, 1745, 3837, 1110, 7955, 4259, 3063, 2509, 3009, 4029, 3719, 1151, 1771, 3838, 1488, # 7414 - 4030, 1986, 7956, 2433, 3487, 7957, 7958, 2093, 7959, 4260, 3839, 1213, 1407, 2803, 531, 2737, # 7430 - 2538, 3233, 1011, 1537, 7960, 2769, 4261, 3106, 1061, 7961, 3720, 3721, 1866, 2883, 7962, 2017, # 7446 - 120, 4262, 4263, 2062, 3595, 3234, 2309, 3840, 2668, 3382, 1954, 4542, 7963, 7964, 3488, 1047, # 7462 - 2704, 1266, 7965, 1368, 4543, 2845, 649, 3383, 3841, 2539, 2738, 1102, 2846, 2669, 7966, 7967, # 7478 - 1999, 7968, 1111, 3596, 2962, 7969, 2488, 3842, 3597, 2804, 1854, 3384, 3722, 7970, 7971, 3385, # 7494 - 2410, 2884, 3304, 3235, 3598, 7972, 2569, 7973, 3599, 2805, 4031, 1460, 856, 7974, 3600, 7975, # 7510 - 2885, 2963, 7976, 2886, 3843, 7977, 4264, 632, 2510, 875, 3844, 1697, 3845, 2291, 7978, 7979, # 7526 - 4544, 3010, 1239, 580, 4545, 4265, 7980, 914, 936, 2074, 1190, 4032, 1039, 2123, 7981, 7982, # 7542 - 7983, 3386, 1473, 7984, 1354, 4266, 3846, 7985, 2172, 3064, 4033, 915, 3305, 4267, 4268, 3306, # 7558 - 1605, 1834, 7986, 2739, 398, 3601, 4269, 3847, 4034, 328, 1912, 2847, 4035, 3848, 1331, 4270, # 7574 - 3011, 937, 4271, 7987, 3602, 4036, 4037, 3387, 2160, 4546, 3388, 524, 742, 538, 3065, 1012, # 7590 - 7988, 7989, 3849, 2461, 7990, 658, 1103, 225, 3850, 7991, 7992, 4547, 7993, 4548, 7994, 3236, # 7606 - 1243, 7995, 4038, 963, 2246, 4549, 7996, 2705, 3603, 3161, 7997, 7998, 2588, 2327, 7999, 4550, # 7622 - 8000, 8001, 8002, 3489, 3307, 957, 3389, 2540, 2032, 1930, 2927, 2462, 870, 2018, 3604, 1746, # 7638 - 2770, 2771, 2434, 2463, 8003, 3851, 8004, 3723, 3107, 3724, 3490, 3390, 3725, 8005, 1179, 3066, # 7654 - 8006, 3162, 2373, 4272, 3726, 2541, 3163, 3108, 2740, 4039, 8007, 3391, 1556, 2542, 2292, 977, # 7670 - 2887, 2033, 4040, 1205, 3392, 8008, 1765, 3393, 3164, 2124, 1271, 1689, 714, 4551, 3491, 8009, # 7686 - 2328, 3852, 533, 4273, 3605, 2181, 617, 8010, 2464, 3308, 3492, 2310, 8011, 8012, 3165, 8013, # 7702 - 8014, 3853, 1987, 618, 427, 2641, 3493, 3394, 8015, 8016, 1244, 1690, 8017, 2806, 4274, 4552, # 7718 - 8018, 3494, 8019, 8020, 2279, 1576, 473, 3606, 4275, 3395, 972, 8021, 3607, 8022, 3067, 8023, # 7734 - 8024, 4553, 4554, 8025, 3727, 4041, 4042, 8026, 153, 4555, 356, 8027, 1891, 2888, 4276, 2143, # 7750 - 408, 803, 2352, 8028, 3854, 8029, 4277, 1646, 2570, 2511, 4556, 4557, 3855, 8030, 3856, 4278, # 7766 - 8031, 2411, 3396, 752, 8032, 8033, 1961, 2964, 8034, 746, 3012, 2465, 8035, 4279, 3728, 698, # 7782 - 4558, 1892, 4280, 3608, 2543, 4559, 3609, 3857, 8036, 3166, 3397, 8037, 1823, 1302, 4043, 2706, # 7798 - 3858, 1973, 4281, 8038, 4282, 3167, 823, 1303, 1288, 1236, 2848, 3495, 4044, 3398, 774, 3859, # 7814 - 8039, 1581, 4560, 1304, 2849, 3860, 4561, 8040, 2435, 2161, 1083, 3237, 4283, 4045, 4284, 344, # 7830 - 1173, 288, 2311, 454, 1683, 8041, 8042, 1461, 4562, 4046, 2589, 8043, 8044, 4563, 985, 894, # 7846 - 8045, 3399, 3168, 8046, 1913, 2928, 3729, 1988, 8047, 2110, 1974, 8048, 4047, 8049, 2571, 1194, # 7862 - 425, 8050, 4564, 3169, 1245, 3730, 4285, 8051, 8052, 2850, 8053, 636, 4565, 1855, 3861, 760, # 7878 - 1799, 8054, 4286, 2209, 1508, 4566, 4048, 1893, 1684, 2293, 8055, 8056, 8057, 4287, 4288, 2210, # 7894 - 479, 8058, 8059, 832, 8060, 4049, 2489, 8061, 2965, 2490, 3731, 990, 3109, 627, 1814, 2642, # 7910 - 4289, 1582, 4290, 2125, 2111, 3496, 4567, 8062, 799, 4291, 3170, 8063, 4568, 2112, 1737, 3013, # 7926 - 1018, 543, 754, 4292, 3309, 1676, 4569, 4570, 4050, 8064, 1489, 8065, 3497, 8066, 2614, 2889, # 7942 - 4051, 8067, 8068, 2966, 8069, 8070, 8071, 8072, 3171, 4571, 4572, 2182, 1722, 8073, 3238, 3239, # 7958 - 1842, 3610, 1715, 481, 365, 1975, 1856, 8074, 8075, 1962, 2491, 4573, 8076, 2126, 3611, 3240, # 7974 - 433, 1894, 2063, 2075, 8077, 602, 2741, 8078, 8079, 8080, 8081, 8082, 3014, 1628, 3400, 8083, # 7990 - 3172, 4574, 4052, 2890, 4575, 2512, 8084, 2544, 2772, 8085, 8086, 8087, 3310, 4576, 2891, 8088, # 8006 - 4577, 8089, 2851, 4578, 4579, 1221, 2967, 4053, 2513, 8090, 8091, 8092, 1867, 1989, 8093, 8094, # 8022 - 8095, 1895, 8096, 8097, 4580, 1896, 4054, 318, 8098, 2094, 4055, 4293, 8099, 8100, 485, 8101, # 8038 - 938, 3862, 553, 2670, 116, 8102, 3863, 3612, 8103, 3498, 2671, 2773, 3401, 3311, 2807, 8104, # 8054 - 3613, 2929, 4056, 1747, 2930, 2968, 8105, 8106, 207, 8107, 8108, 2672, 4581, 2514, 8109, 3015, # 8070 - 890, 3614, 3864, 8110, 1877, 3732, 3402, 8111, 2183, 2353, 3403, 1652, 8112, 8113, 8114, 941, # 8086 - 2294, 208, 3499, 4057, 2019, 330, 4294, 3865, 2892, 2492, 3733, 4295, 8115, 8116, 8117, 8118, # 8102 -) -# fmt: on diff --git a/venv/lib/python3.10/site-packages/chardet/euctwprober.py b/venv/lib/python3.10/site-packages/chardet/euctwprober.py deleted file mode 100644 index a37ab18995822ad6b3372d56366becdccf9a4c26..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/euctwprober.py +++ /dev/null @@ -1,47 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is mozilla.org code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from .chardistribution import EUCTWDistributionAnalysis -from .codingstatemachine import CodingStateMachine -from .mbcharsetprober import MultiByteCharSetProber -from .mbcssm import EUCTW_SM_MODEL - - -class EUCTWProber(MultiByteCharSetProber): - def __init__(self) -> None: - super().__init__() - self.coding_sm = CodingStateMachine(EUCTW_SM_MODEL) - self.distribution_analyzer = EUCTWDistributionAnalysis() - self.reset() - - @property - def charset_name(self) -> str: - return "EUC-TW" - - @property - def language(self) -> str: - return "Taiwan" diff --git a/venv/lib/python3.10/site-packages/chardet/gb2312freq.py b/venv/lib/python3.10/site-packages/chardet/gb2312freq.py deleted file mode 100644 index b32bfc74213d93d434f1f3a47cb5d7d0bf4863d3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/gb2312freq.py +++ /dev/null @@ -1,284 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Communicator client code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -# GB2312 most frequently used character table -# -# Char to FreqOrder table , from hz6763 - -# 512 --> 0.79 -- 0.79 -# 1024 --> 0.92 -- 0.13 -# 2048 --> 0.98 -- 0.06 -# 6768 --> 1.00 -- 0.02 -# -# Ideal Distribution Ratio = 0.79135/(1-0.79135) = 3.79 -# Random Distribution Ration = 512 / (3755 - 512) = 0.157 -# -# Typical Distribution Ratio about 25% of Ideal one, still much higher that RDR - -GB2312_TYPICAL_DISTRIBUTION_RATIO = 0.9 - -GB2312_TABLE_SIZE = 3760 - -# fmt: off -GB2312_CHAR_TO_FREQ_ORDER = ( -1671, 749,1443,2364,3924,3807,2330,3921,1704,3463,2691,1511,1515, 572,3191,2205, -2361, 224,2558, 479,1711, 963,3162, 440,4060,1905,2966,2947,3580,2647,3961,3842, -2204, 869,4207, 970,2678,5626,2944,2956,1479,4048, 514,3595, 588,1346,2820,3409, - 249,4088,1746,1873,2047,1774, 581,1813, 358,1174,3590,1014,1561,4844,2245, 670, -1636,3112, 889,1286, 953, 556,2327,3060,1290,3141, 613, 185,3477,1367, 850,3820, -1715,2428,2642,2303,2732,3041,2562,2648,3566,3946,1349, 388,3098,2091,1360,3585, - 152,1687,1539, 738,1559, 59,1232,2925,2267,1388,1249,1741,1679,2960, 151,1566, -1125,1352,4271, 924,4296, 385,3166,4459, 310,1245,2850, 70,3285,2729,3534,3575, -2398,3298,3466,1960,2265, 217,3647, 864,1909,2084,4401,2773,1010,3269,5152, 853, -3051,3121,1244,4251,1895, 364,1499,1540,2313,1180,3655,2268, 562, 715,2417,3061, - 544, 336,3768,2380,1752,4075, 950, 280,2425,4382, 183,2759,3272, 333,4297,2155, -1688,2356,1444,1039,4540, 736,1177,3349,2443,2368,2144,2225, 565, 196,1482,3406, - 927,1335,4147, 692, 878,1311,1653,3911,3622,1378,4200,1840,2969,3149,2126,1816, -2534,1546,2393,2760, 737,2494, 13, 447, 245,2747, 38,2765,2129,2589,1079, 606, - 360, 471,3755,2890, 404, 848, 699,1785,1236, 370,2221,1023,3746,2074,2026,2023, -2388,1581,2119, 812,1141,3091,2536,1519, 804,2053, 406,1596,1090, 784, 548,4414, -1806,2264,2936,1100, 343,4114,5096, 622,3358, 743,3668,1510,1626,5020,3567,2513, -3195,4115,5627,2489,2991, 24,2065,2697,1087,2719, 48,1634, 315, 68, 985,2052, - 198,2239,1347,1107,1439, 597,2366,2172, 871,3307, 919,2487,2790,1867, 236,2570, -1413,3794, 906,3365,3381,1701,1982,1818,1524,2924,1205, 616,2586,2072,2004, 575, - 253,3099, 32,1365,1182, 197,1714,2454,1201, 554,3388,3224,2748, 756,2587, 250, -2567,1507,1517,3529,1922,2761,2337,3416,1961,1677,2452,2238,3153, 615, 911,1506, -1474,2495,1265,1906,2749,3756,3280,2161, 898,2714,1759,3450,2243,2444, 563, 26, -3286,2266,3769,3344,2707,3677, 611,1402, 531,1028,2871,4548,1375, 261,2948, 835, -1190,4134, 353, 840,2684,1900,3082,1435,2109,1207,1674, 329,1872,2781,4055,2686, -2104, 608,3318,2423,2957,2768,1108,3739,3512,3271,3985,2203,1771,3520,1418,2054, -1681,1153, 225,1627,2929, 162,2050,2511,3687,1954, 124,1859,2431,1684,3032,2894, - 585,4805,3969,2869,2704,2088,2032,2095,3656,2635,4362,2209, 256, 518,2042,2105, -3777,3657, 643,2298,1148,1779, 190, 989,3544, 414, 11,2135,2063,2979,1471, 403, -3678, 126, 770,1563, 671,2499,3216,2877, 600,1179, 307,2805,4937,1268,1297,2694, - 252,4032,1448,1494,1331,1394, 127,2256, 222,1647,1035,1481,3056,1915,1048, 873, -3651, 210, 33,1608,2516, 200,1520, 415, 102, 0,3389,1287, 817, 91,3299,2940, - 836,1814, 549,2197,1396,1669,2987,3582,2297,2848,4528,1070, 687, 20,1819, 121, -1552,1364,1461,1968,2617,3540,2824,2083, 177, 948,4938,2291, 110,4549,2066, 648, -3359,1755,2110,2114,4642,4845,1693,3937,3308,1257,1869,2123, 208,1804,3159,2992, -2531,2549,3361,2418,1350,2347,2800,2568,1291,2036,2680, 72, 842,1990, 212,1233, -1154,1586, 75,2027,3410,4900,1823,1337,2710,2676, 728,2810,1522,3026,4995, 157, - 755,1050,4022, 710, 785,1936,2194,2085,1406,2777,2400, 150,1250,4049,1206, 807, -1910, 534, 529,3309,1721,1660, 274, 39,2827, 661,2670,1578, 925,3248,3815,1094, -4278,4901,4252, 41,1150,3747,2572,2227,4501,3658,4902,3813,3357,3617,2884,2258, - 887, 538,4187,3199,1294,2439,3042,2329,2343,2497,1255, 107, 543,1527, 521,3478, -3568, 194,5062, 15, 961,3870,1241,1192,2664, 66,5215,3260,2111,1295,1127,2152, -3805,4135, 901,1164,1976, 398,1278, 530,1460, 748, 904,1054,1966,1426, 53,2909, - 509, 523,2279,1534, 536,1019, 239,1685, 460,2353, 673,1065,2401,3600,4298,2272, -1272,2363, 284,1753,3679,4064,1695, 81, 815,2677,2757,2731,1386, 859, 500,4221, -2190,2566, 757,1006,2519,2068,1166,1455, 337,2654,3203,1863,1682,1914,3025,1252, -1409,1366, 847, 714,2834,2038,3209, 964,2970,1901, 885,2553,1078,1756,3049, 301, -1572,3326, 688,2130,1996,2429,1805,1648,2930,3421,2750,3652,3088, 262,1158,1254, - 389,1641,1812, 526,1719, 923,2073,1073,1902, 468, 489,4625,1140, 857,2375,3070, -3319,2863, 380, 116,1328,2693,1161,2244, 273,1212,1884,2769,3011,1775,1142, 461, -3066,1200,2147,2212, 790, 702,2695,4222,1601,1058, 434,2338,5153,3640, 67,2360, -4099,2502, 618,3472,1329, 416,1132, 830,2782,1807,2653,3211,3510,1662, 192,2124, - 296,3979,1739,1611,3684, 23, 118, 324, 446,1239,1225, 293,2520,3814,3795,2535, -3116, 17,1074, 467,2692,2201, 387,2922, 45,1326,3055,1645,3659,2817, 958, 243, -1903,2320,1339,2825,1784,3289, 356, 576, 865,2315,2381,3377,3916,1088,3122,1713, -1655, 935, 628,4689,1034,1327, 441, 800, 720, 894,1979,2183,1528,5289,2702,1071, -4046,3572,2399,1571,3281, 79, 761,1103, 327, 134, 758,1899,1371,1615, 879, 442, - 215,2605,2579, 173,2048,2485,1057,2975,3317,1097,2253,3801,4263,1403,1650,2946, - 814,4968,3487,1548,2644,1567,1285, 2, 295,2636, 97, 946,3576, 832, 141,4257, -3273, 760,3821,3521,3156,2607, 949,1024,1733,1516,1803,1920,2125,2283,2665,3180, -1501,2064,3560,2171,1592, 803,3518,1416, 732,3897,4258,1363,1362,2458, 119,1427, - 602,1525,2608,1605,1639,3175, 694,3064, 10, 465, 76,2000,4846,4208, 444,3781, -1619,3353,2206,1273,3796, 740,2483, 320,1723,2377,3660,2619,1359,1137,1762,1724, -2345,2842,1850,1862, 912, 821,1866, 612,2625,1735,2573,3369,1093, 844, 89, 937, - 930,1424,3564,2413,2972,1004,3046,3019,2011, 711,3171,1452,4178, 428, 801,1943, - 432, 445,2811, 206,4136,1472, 730, 349, 73, 397,2802,2547, 998,1637,1167, 789, - 396,3217, 154,1218, 716,1120,1780,2819,4826,1931,3334,3762,2139,1215,2627, 552, -3664,3628,3232,1405,2383,3111,1356,2652,3577,3320,3101,1703, 640,1045,1370,1246, -4996, 371,1575,2436,1621,2210, 984,4033,1734,2638, 16,4529, 663,2755,3255,1451, -3917,2257,1253,1955,2234,1263,2951, 214,1229, 617, 485, 359,1831,1969, 473,2310, - 750,2058, 165, 80,2864,2419, 361,4344,2416,2479,1134, 796,3726,1266,2943, 860, -2715, 938, 390,2734,1313,1384, 248, 202, 877,1064,2854, 522,3907, 279,1602, 297, -2357, 395,3740, 137,2075, 944,4089,2584,1267,3802, 62,1533,2285, 178, 176, 780, -2440, 201,3707, 590, 478,1560,4354,2117,1075, 30, 74,4643,4004,1635,1441,2745, - 776,2596, 238,1077,1692,1912,2844, 605, 499,1742,3947, 241,3053, 980,1749, 936, -2640,4511,2582, 515,1543,2162,5322,2892,2993, 890,2148,1924, 665,1827,3581,1032, - 968,3163, 339,1044,1896, 270, 583,1791,1720,4367,1194,3488,3669, 43,2523,1657, - 163,2167, 290,1209,1622,3378, 550, 634,2508,2510, 695,2634,2384,2512,1476,1414, - 220,1469,2341,2138,2852,3183,2900,4939,2865,3502,1211,3680, 854,3227,1299,2976, -3172, 186,2998,1459, 443,1067,3251,1495, 321,1932,3054, 909, 753,1410,1828, 436, -2441,1119,1587,3164,2186,1258, 227, 231,1425,1890,3200,3942, 247, 959, 725,5254, -2741, 577,2158,2079, 929, 120, 174, 838,2813, 591,1115, 417,2024, 40,3240,1536, -1037, 291,4151,2354, 632,1298,2406,2500,3535,1825,1846,3451, 205,1171, 345,4238, - 18,1163, 811, 685,2208,1217, 425,1312,1508,1175,4308,2552,1033, 587,1381,3059, -2984,3482, 340,1316,4023,3972, 792,3176, 519, 777,4690, 918, 933,4130,2981,3741, - 90,3360,2911,2200,5184,4550, 609,3079,2030, 272,3379,2736, 363,3881,1130,1447, - 286, 779, 357,1169,3350,3137,1630,1220,2687,2391, 747,1277,3688,2618,2682,2601, -1156,3196,5290,4034,3102,1689,3596,3128, 874, 219,2783, 798, 508,1843,2461, 269, -1658,1776,1392,1913,2983,3287,2866,2159,2372, 829,4076, 46,4253,2873,1889,1894, - 915,1834,1631,2181,2318, 298, 664,2818,3555,2735, 954,3228,3117, 527,3511,2173, - 681,2712,3033,2247,2346,3467,1652, 155,2164,3382, 113,1994, 450, 899, 494, 994, -1237,2958,1875,2336,1926,3727, 545,1577,1550, 633,3473, 204,1305,3072,2410,1956, -2471, 707,2134, 841,2195,2196,2663,3843,1026,4940, 990,3252,4997, 368,1092, 437, -3212,3258,1933,1829, 675,2977,2893, 412, 943,3723,4644,3294,3283,2230,2373,5154, -2389,2241,2661,2323,1404,2524, 593, 787, 677,3008,1275,2059, 438,2709,2609,2240, -2269,2246,1446, 36,1568,1373,3892,1574,2301,1456,3962, 693,2276,5216,2035,1143, -2720,1919,1797,1811,2763,4137,2597,1830,1699,1488,1198,2090, 424,1694, 312,3634, -3390,4179,3335,2252,1214, 561,1059,3243,2295,2561, 975,5155,2321,2751,3772, 472, -1537,3282,3398,1047,2077,2348,2878,1323,3340,3076, 690,2906, 51, 369, 170,3541, -1060,2187,2688,3670,2541,1083,1683, 928,3918, 459, 109,4427, 599,3744,4286, 143, -2101,2730,2490, 82,1588,3036,2121, 281,1860, 477,4035,1238,2812,3020,2716,3312, -1530,2188,2055,1317, 843, 636,1808,1173,3495, 649, 181,1002, 147,3641,1159,2414, -3750,2289,2795, 813,3123,2610,1136,4368, 5,3391,4541,2174, 420, 429,1728, 754, -1228,2115,2219, 347,2223,2733, 735,1518,3003,2355,3134,1764,3948,3329,1888,2424, -1001,1234,1972,3321,3363,1672,1021,1450,1584, 226, 765, 655,2526,3404,3244,2302, -3665, 731, 594,2184, 319,1576, 621, 658,2656,4299,2099,3864,1279,2071,2598,2739, - 795,3086,3699,3908,1707,2352,2402,1382,3136,2475,1465,4847,3496,3865,1085,3004, -2591,1084, 213,2287,1963,3565,2250, 822, 793,4574,3187,1772,1789,3050, 595,1484, -1959,2770,1080,2650, 456, 422,2996, 940,3322,4328,4345,3092,2742, 965,2784, 739, -4124, 952,1358,2498,2949,2565, 332,2698,2378, 660,2260,2473,4194,3856,2919, 535, -1260,2651,1208,1428,1300,1949,1303,2942, 433,2455,2450,1251,1946, 614,1269, 641, -1306,1810,2737,3078,2912, 564,2365,1419,1415,1497,4460,2367,2185,1379,3005,1307, -3218,2175,1897,3063, 682,1157,4040,4005,1712,1160,1941,1399, 394, 402,2952,1573, -1151,2986,2404, 862, 299,2033,1489,3006, 346, 171,2886,3401,1726,2932, 168,2533, - 47,2507,1030,3735,1145,3370,1395,1318,1579,3609,4560,2857,4116,1457,2529,1965, - 504,1036,2690,2988,2405, 745,5871, 849,2397,2056,3081, 863,2359,3857,2096, 99, -1397,1769,2300,4428,1643,3455,1978,1757,3718,1440, 35,4879,3742,1296,4228,2280, - 160,5063,1599,2013, 166, 520,3479,1646,3345,3012, 490,1937,1545,1264,2182,2505, -1096,1188,1369,1436,2421,1667,2792,2460,1270,2122, 727,3167,2143, 806,1706,1012, -1800,3037, 960,2218,1882, 805, 139,2456,1139,1521, 851,1052,3093,3089, 342,2039, - 744,5097,1468,1502,1585,2087, 223, 939, 326,2140,2577, 892,2481,1623,4077, 982, -3708, 135,2131, 87,2503,3114,2326,1106, 876,1616, 547,2997,2831,2093,3441,4530, -4314, 9,3256,4229,4148, 659,1462,1986,1710,2046,2913,2231,4090,4880,5255,3392, -3274,1368,3689,4645,1477, 705,3384,3635,1068,1529,2941,1458,3782,1509, 100,1656, -2548, 718,2339, 408,1590,2780,3548,1838,4117,3719,1345,3530, 717,3442,2778,3220, -2898,1892,4590,3614,3371,2043,1998,1224,3483, 891, 635, 584,2559,3355, 733,1766, -1729,1172,3789,1891,2307, 781,2982,2271,1957,1580,5773,2633,2005,4195,3097,1535, -3213,1189,1934,5693,3262, 586,3118,1324,1598, 517,1564,2217,1868,1893,4445,3728, -2703,3139,1526,1787,1992,3882,2875,1549,1199,1056,2224,1904,2711,5098,4287, 338, -1993,3129,3489,2689,1809,2815,1997, 957,1855,3898,2550,3275,3057,1105,1319, 627, -1505,1911,1883,3526, 698,3629,3456,1833,1431, 746, 77,1261,2017,2296,1977,1885, - 125,1334,1600, 525,1798,1109,2222,1470,1945, 559,2236,1186,3443,2476,1929,1411, -2411,3135,1777,3372,2621,1841,1613,3229, 668,1430,1839,2643,2916, 195,1989,2671, -2358,1387, 629,3205,2293,5256,4439, 123,1310, 888,1879,4300,3021,3605,1003,1162, -3192,2910,2010, 140,2395,2859, 55,1082,2012,2901, 662, 419,2081,1438, 680,2774, -4654,3912,1620,1731,1625,5035,4065,2328, 512,1344, 802,5443,2163,2311,2537, 524, -3399, 98,1155,2103,1918,2606,3925,2816,1393,2465,1504,3773,2177,3963,1478,4346, - 180,1113,4655,3461,2028,1698, 833,2696,1235,1322,1594,4408,3623,3013,3225,2040, -3022, 541,2881, 607,3632,2029,1665,1219, 639,1385,1686,1099,2803,3231,1938,3188, -2858, 427, 676,2772,1168,2025, 454,3253,2486,3556, 230,1950, 580, 791,1991,1280, -1086,1974,2034, 630, 257,3338,2788,4903,1017, 86,4790, 966,2789,1995,1696,1131, - 259,3095,4188,1308, 179,1463,5257, 289,4107,1248, 42,3413,1725,2288, 896,1947, - 774,4474,4254, 604,3430,4264, 392,2514,2588, 452, 237,1408,3018, 988,4531,1970, -3034,3310, 540,2370,1562,1288,2990, 502,4765,1147, 4,1853,2708, 207, 294,2814, -4078,2902,2509, 684, 34,3105,3532,2551, 644, 709,2801,2344, 573,1727,3573,3557, -2021,1081,3100,4315,2100,3681, 199,2263,1837,2385, 146,3484,1195,2776,3949, 997, -1939,3973,1008,1091,1202,1962,1847,1149,4209,5444,1076, 493, 117,5400,2521, 972, -1490,2934,1796,4542,2374,1512,2933,2657, 413,2888,1135,2762,2314,2156,1355,2369, - 766,2007,2527,2170,3124,2491,2593,2632,4757,2437, 234,3125,3591,1898,1750,1376, -1942,3468,3138, 570,2127,2145,3276,4131, 962, 132,1445,4196, 19, 941,3624,3480, -3366,1973,1374,4461,3431,2629, 283,2415,2275, 808,2887,3620,2112,2563,1353,3610, - 955,1089,3103,1053, 96, 88,4097, 823,3808,1583, 399, 292,4091,3313, 421,1128, - 642,4006, 903,2539,1877,2082, 596, 29,4066,1790, 722,2157, 130, 995,1569, 769, -1485, 464, 513,2213, 288,1923,1101,2453,4316, 133, 486,2445, 50, 625, 487,2207, - 57, 423, 481,2962, 159,3729,1558, 491, 303, 482, 501, 240,2837, 112,3648,2392, -1783, 362, 8,3433,3422, 610,2793,3277,1390,1284,1654, 21,3823, 734, 367, 623, - 193, 287, 374,1009,1483, 816, 476, 313,2255,2340,1262,2150,2899,1146,2581, 782, -2116,1659,2018,1880, 255,3586,3314,1110,2867,2137,2564, 986,2767,5185,2006, 650, - 158, 926, 762, 881,3157,2717,2362,3587, 306,3690,3245,1542,3077,2427,1691,2478, -2118,2985,3490,2438, 539,2305, 983, 129,1754, 355,4201,2386, 827,2923, 104,1773, -2838,2771, 411,2905,3919, 376, 767, 122,1114, 828,2422,1817,3506, 266,3460,1007, -1609,4998, 945,2612,4429,2274, 726,1247,1964,2914,2199,2070,4002,4108, 657,3323, -1422, 579, 455,2764,4737,1222,2895,1670, 824,1223,1487,2525, 558, 861,3080, 598, -2659,2515,1967, 752,2583,2376,2214,4180, 977, 704,2464,4999,2622,4109,1210,2961, - 819,1541, 142,2284, 44, 418, 457,1126,3730,4347,4626,1644,1876,3671,1864, 302, -1063,5694, 624, 723,1984,3745,1314,1676,2488,1610,1449,3558,3569,2166,2098, 409, -1011,2325,3704,2306, 818,1732,1383,1824,1844,3757, 999,2705,3497,1216,1423,2683, -2426,2954,2501,2726,2229,1475,2554,5064,1971,1794,1666,2014,1343, 783, 724, 191, -2434,1354,2220,5065,1763,2752,2472,4152, 131, 175,2885,3434, 92,1466,4920,2616, -3871,3872,3866, 128,1551,1632, 669,1854,3682,4691,4125,1230, 188,2973,3290,1302, -1213, 560,3266, 917, 763,3909,3249,1760, 868,1958, 764,1782,2097, 145,2277,3774, -4462, 64,1491,3062, 971,2132,3606,2442, 221,1226,1617, 218, 323,1185,3207,3147, - 571, 619,1473,1005,1744,2281, 449,1887,2396,3685, 275, 375,3816,1743,3844,3731, - 845,1983,2350,4210,1377, 773, 967,3499,3052,3743,2725,4007,1697,1022,3943,1464, -3264,2855,2722,1952,1029,2839,2467, 84,4383,2215, 820,1391,2015,2448,3672, 377, -1948,2168, 797,2545,3536,2578,2645, 94,2874,1678, 405,1259,3071, 771, 546,1315, - 470,1243,3083, 895,2468, 981, 969,2037, 846,4181, 653,1276,2928, 14,2594, 557, -3007,2474, 156, 902,1338,1740,2574, 537,2518, 973,2282,2216,2433,1928, 138,2903, -1293,2631,1612, 646,3457, 839,2935, 111, 496,2191,2847, 589,3186, 149,3994,2060, -4031,2641,4067,3145,1870, 37,3597,2136,1025,2051,3009,3383,3549,1121,1016,3261, -1301, 251,2446,2599,2153, 872,3246, 637, 334,3705, 831, 884, 921,3065,3140,4092, -2198,1944, 246,2964, 108,2045,1152,1921,2308,1031, 203,3173,4170,1907,3890, 810, -1401,2003,1690, 506, 647,1242,2828,1761,1649,3208,2249,1589,3709,2931,5156,1708, - 498, 666,2613, 834,3817,1231, 184,2851,1124, 883,3197,2261,3710,1765,1553,2658, -1178,2639,2351, 93,1193, 942,2538,2141,4402, 235,1821, 870,1591,2192,1709,1871, -3341,1618,4126,2595,2334, 603, 651, 69, 701, 268,2662,3411,2555,1380,1606, 503, - 448, 254,2371,2646, 574,1187,2309,1770, 322,2235,1292,1801, 305, 566,1133, 229, -2067,2057, 706, 167, 483,2002,2672,3295,1820,3561,3067, 316, 378,2746,3452,1112, - 136,1981, 507,1651,2917,1117, 285,4591, 182,2580,3522,1304, 335,3303,1835,2504, -1795,1792,2248, 674,1018,2106,2449,1857,2292,2845, 976,3047,1781,2600,2727,1389, -1281, 52,3152, 153, 265,3950, 672,3485,3951,4463, 430,1183, 365, 278,2169, 27, -1407,1336,2304, 209,1340,1730,2202,1852,2403,2883, 979,1737,1062, 631,2829,2542, -3876,2592, 825,2086,2226,3048,3625, 352,1417,3724, 542, 991, 431,1351,3938,1861, -2294, 826,1361,2927,3142,3503,1738, 463,2462,2723, 582,1916,1595,2808, 400,3845, -3891,2868,3621,2254, 58,2492,1123, 910,2160,2614,1372,1603,1196,1072,3385,1700, -3267,1980, 696, 480,2430, 920, 799,1570,2920,1951,2041,4047,2540,1321,4223,2469, -3562,2228,1271,2602, 401,2833,3351,2575,5157, 907,2312,1256, 410, 263,3507,1582, - 996, 678,1849,2316,1480, 908,3545,2237, 703,2322, 667,1826,2849,1531,2604,2999, -2407,3146,2151,2630,1786,3711, 469,3542, 497,3899,2409, 858, 837,4446,3393,1274, - 786, 620,1845,2001,3311, 484, 308,3367,1204,1815,3691,2332,1532,2557,1842,2020, -2724,1927,2333,4440, 567, 22,1673,2728,4475,1987,1858,1144,1597, 101,1832,3601, - 12, 974,3783,4391, 951,1412, 1,3720, 453,4608,4041, 528,1041,1027,3230,2628, -1129, 875,1051,3291,1203,2262,1069,2860,2799,2149,2615,3278, 144,1758,3040, 31, - 475,1680, 366,2685,3184, 311,1642,4008,2466,5036,1593,1493,2809, 216,1420,1668, - 233, 304,2128,3284, 232,1429,1768,1040,2008,3407,2740,2967,2543, 242,2133, 778, -1565,2022,2620, 505,2189,2756,1098,2273, 372,1614, 708, 553,2846,2094,2278, 169, -3626,2835,4161, 228,2674,3165, 809,1454,1309, 466,1705,1095, 900,3423, 880,2667, -3751,5258,2317,3109,2571,4317,2766,1503,1342, 866,4447,1118, 63,2076, 314,1881, -1348,1061, 172, 978,3515,1747, 532, 511,3970, 6, 601, 905,2699,3300,1751, 276, -1467,3725,2668, 65,4239,2544,2779,2556,1604, 578,2451,1802, 992,2331,2624,1320, -3446, 713,1513,1013, 103,2786,2447,1661, 886,1702, 916, 654,3574,2031,1556, 751, -2178,2821,2179,1498,1538,2176, 271, 914,2251,2080,1325, 638,1953,2937,3877,2432, -2754, 95,3265,1716, 260,1227,4083, 775, 106,1357,3254, 426,1607, 555,2480, 772, -1985, 244,2546, 474, 495,1046,2611,1851,2061, 71,2089,1675,2590, 742,3758,2843, -3222,1433, 267,2180,2576,2826,2233,2092,3913,2435, 956,1745,3075, 856,2113,1116, - 451, 3,1988,2896,1398, 993,2463,1878,2049,1341,2718,2721,2870,2108, 712,2904, -4363,2753,2324, 277,2872,2349,2649, 384, 987, 435, 691,3000, 922, 164,3939, 652, -1500,1184,4153,2482,3373,2165,4848,2335,3775,3508,3154,2806,2830,1554,2102,1664, -2530,1434,2408, 893,1547,2623,3447,2832,2242,2532,3169,2856,3223,2078, 49,3770, -3469, 462, 318, 656,2259,3250,3069, 679,1629,2758, 344,1138,1104,3120,1836,1283, -3115,2154,1437,4448, 934, 759,1999, 794,2862,1038, 533,2560,1722,2342, 855,2626, -1197,1663,4476,3127, 85,4240,2528, 25,1111,1181,3673, 407,3470,4561,2679,2713, - 768,1925,2841,3986,1544,1165, 932, 373,1240,2146,1930,2673, 721,4766, 354,4333, - 391,2963, 187, 61,3364,1442,1102, 330,1940,1767, 341,3809,4118, 393,2496,2062, -2211, 105, 331, 300, 439, 913,1332, 626, 379,3304,1557, 328, 689,3952, 309,1555, - 931, 317,2517,3027, 325, 569, 686,2107,3084, 60,1042,1333,2794, 264,3177,4014, -1628, 258,3712, 7,4464,1176,1043,1778, 683, 114,1975, 78,1492, 383,1886, 510, - 386, 645,5291,2891,2069,3305,4138,3867,2939,2603,2493,1935,1066,1848,3588,1015, -1282,1289,4609, 697,1453,3044,2666,3611,1856,2412, 54, 719,1330, 568,3778,2459, -1748, 788, 492, 551,1191,1000, 488,3394,3763, 282,1799, 348,2016,1523,3155,2390, -1049, 382,2019,1788,1170, 729,2968,3523, 897,3926,2785,2938,3292, 350,2319,3238, -1718,1717,2655,3453,3143,4465, 161,2889,2980,2009,1421, 56,1908,1640,2387,2232, -1917,1874,2477,4921, 148, 83,3438, 592,4245,2882,1822,1055, 741, 115,1496,1624, - 381,1638,4592,1020, 516,3214, 458, 947,4575,1432, 211,1514,2926,1865,2142, 189, - 852,1221,1400,1486, 882,2299,4036, 351, 28,1122, 700,6479,6480,6481,6482,6483, #last 512 -) -# fmt: on diff --git a/venv/lib/python3.10/site-packages/chardet/gb2312prober.py b/venv/lib/python3.10/site-packages/chardet/gb2312prober.py deleted file mode 100644 index d423e7311e2fbd9a014de808c107e96ad11c66e5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/gb2312prober.py +++ /dev/null @@ -1,47 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is mozilla.org code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from .chardistribution import GB2312DistributionAnalysis -from .codingstatemachine import CodingStateMachine -from .mbcharsetprober import MultiByteCharSetProber -from .mbcssm import GB2312_SM_MODEL - - -class GB2312Prober(MultiByteCharSetProber): - def __init__(self) -> None: - super().__init__() - self.coding_sm = CodingStateMachine(GB2312_SM_MODEL) - self.distribution_analyzer = GB2312DistributionAnalysis() - self.reset() - - @property - def charset_name(self) -> str: - return "GB2312" - - @property - def language(self) -> str: - return "Chinese" diff --git a/venv/lib/python3.10/site-packages/chardet/hebrewprober.py b/venv/lib/python3.10/site-packages/chardet/hebrewprober.py deleted file mode 100644 index 785d0057bcc0ea74a4b8d65ab7a0de78474bf892..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/hebrewprober.py +++ /dev/null @@ -1,316 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Universal charset detector code. -# -# The Initial Developer of the Original Code is -# Shy Shalom -# Portions created by the Initial Developer are Copyright (C) 2005 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from typing import Optional, Union - -from .charsetprober import CharSetProber -from .enums import ProbingState -from .sbcharsetprober import SingleByteCharSetProber - -# This prober doesn't actually recognize a language or a charset. -# It is a helper prober for the use of the Hebrew model probers - -### General ideas of the Hebrew charset recognition ### -# -# Four main charsets exist in Hebrew: -# "ISO-8859-8" - Visual Hebrew -# "windows-1255" - Logical Hebrew -# "ISO-8859-8-I" - Logical Hebrew -# "x-mac-hebrew" - ?? Logical Hebrew ?? -# -# Both "ISO" charsets use a completely identical set of code points, whereas -# "windows-1255" and "x-mac-hebrew" are two different proper supersets of -# these code points. windows-1255 defines additional characters in the range -# 0x80-0x9F as some misc punctuation marks as well as some Hebrew-specific -# diacritics and additional 'Yiddish' ligature letters in the range 0xc0-0xd6. -# x-mac-hebrew defines similar additional code points but with a different -# mapping. -# -# As far as an average Hebrew text with no diacritics is concerned, all four -# charsets are identical with respect to code points. Meaning that for the -# main Hebrew alphabet, all four map the same values to all 27 Hebrew letters -# (including final letters). -# -# The dominant difference between these charsets is their directionality. -# "Visual" directionality means that the text is ordered as if the renderer is -# not aware of a BIDI rendering algorithm. The renderer sees the text and -# draws it from left to right. The text itself when ordered naturally is read -# backwards. A buffer of Visual Hebrew generally looks like so: -# "[last word of first line spelled backwards] [whole line ordered backwards -# and spelled backwards] [first word of first line spelled backwards] -# [end of line] [last word of second line] ... etc' " -# adding punctuation marks, numbers and English text to visual text is -# naturally also "visual" and from left to right. -# -# "Logical" directionality means the text is ordered "naturally" according to -# the order it is read. It is the responsibility of the renderer to display -# the text from right to left. A BIDI algorithm is used to place general -# punctuation marks, numbers and English text in the text. -# -# Texts in x-mac-hebrew are almost impossible to find on the Internet. From -# what little evidence I could find, it seems that its general directionality -# is Logical. -# -# To sum up all of the above, the Hebrew probing mechanism knows about two -# charsets: -# Visual Hebrew - "ISO-8859-8" - backwards text - Words and sentences are -# backwards while line order is natural. For charset recognition purposes -# the line order is unimportant (In fact, for this implementation, even -# word order is unimportant). -# Logical Hebrew - "windows-1255" - normal, naturally ordered text. -# -# "ISO-8859-8-I" is a subset of windows-1255 and doesn't need to be -# specifically identified. -# "x-mac-hebrew" is also identified as windows-1255. A text in x-mac-hebrew -# that contain special punctuation marks or diacritics is displayed with -# some unconverted characters showing as question marks. This problem might -# be corrected using another model prober for x-mac-hebrew. Due to the fact -# that x-mac-hebrew texts are so rare, writing another model prober isn't -# worth the effort and performance hit. -# -#### The Prober #### -# -# The prober is divided between two SBCharSetProbers and a HebrewProber, -# all of which are managed, created, fed data, inquired and deleted by the -# SBCSGroupProber. The two SBCharSetProbers identify that the text is in -# fact some kind of Hebrew, Logical or Visual. The final decision about which -# one is it is made by the HebrewProber by combining final-letter scores -# with the scores of the two SBCharSetProbers to produce a final answer. -# -# The SBCSGroupProber is responsible for stripping the original text of HTML -# tags, English characters, numbers, low-ASCII punctuation characters, spaces -# and new lines. It reduces any sequence of such characters to a single space. -# The buffer fed to each prober in the SBCS group prober is pure text in -# high-ASCII. -# The two SBCharSetProbers (model probers) share the same language model: -# Win1255Model. -# The first SBCharSetProber uses the model normally as any other -# SBCharSetProber does, to recognize windows-1255, upon which this model was -# built. The second SBCharSetProber is told to make the pair-of-letter -# lookup in the language model backwards. This in practice exactly simulates -# a visual Hebrew model using the windows-1255 logical Hebrew model. -# -# The HebrewProber is not using any language model. All it does is look for -# final-letter evidence suggesting the text is either logical Hebrew or visual -# Hebrew. Disjointed from the model probers, the results of the HebrewProber -# alone are meaningless. HebrewProber always returns 0.00 as confidence -# since it never identifies a charset by itself. Instead, the pointer to the -# HebrewProber is passed to the model probers as a helper "Name Prober". -# When the Group prober receives a positive identification from any prober, -# it asks for the name of the charset identified. If the prober queried is a -# Hebrew model prober, the model prober forwards the call to the -# HebrewProber to make the final decision. In the HebrewProber, the -# decision is made according to the final-letters scores maintained and Both -# model probers scores. The answer is returned in the form of the name of the -# charset identified, either "windows-1255" or "ISO-8859-8". - - -class HebrewProber(CharSetProber): - SPACE = 0x20 - # windows-1255 / ISO-8859-8 code points of interest - FINAL_KAF = 0xEA - NORMAL_KAF = 0xEB - FINAL_MEM = 0xED - NORMAL_MEM = 0xEE - FINAL_NUN = 0xEF - NORMAL_NUN = 0xF0 - FINAL_PE = 0xF3 - NORMAL_PE = 0xF4 - FINAL_TSADI = 0xF5 - NORMAL_TSADI = 0xF6 - - # Minimum Visual vs Logical final letter score difference. - # If the difference is below this, don't rely solely on the final letter score - # distance. - MIN_FINAL_CHAR_DISTANCE = 5 - - # Minimum Visual vs Logical model score difference. - # If the difference is below this, don't rely at all on the model score - # distance. - MIN_MODEL_DISTANCE = 0.01 - - VISUAL_HEBREW_NAME = "ISO-8859-8" - LOGICAL_HEBREW_NAME = "windows-1255" - - def __init__(self) -> None: - super().__init__() - self._final_char_logical_score = 0 - self._final_char_visual_score = 0 - self._prev = self.SPACE - self._before_prev = self.SPACE - self._logical_prober: Optional[SingleByteCharSetProber] = None - self._visual_prober: Optional[SingleByteCharSetProber] = None - self.reset() - - def reset(self) -> None: - self._final_char_logical_score = 0 - self._final_char_visual_score = 0 - # The two last characters seen in the previous buffer, - # mPrev and mBeforePrev are initialized to space in order to simulate - # a word delimiter at the beginning of the data - self._prev = self.SPACE - self._before_prev = self.SPACE - # These probers are owned by the group prober. - - def set_model_probers( - self, - logical_prober: SingleByteCharSetProber, - visual_prober: SingleByteCharSetProber, - ) -> None: - self._logical_prober = logical_prober - self._visual_prober = visual_prober - - def is_final(self, c: int) -> bool: - return c in [ - self.FINAL_KAF, - self.FINAL_MEM, - self.FINAL_NUN, - self.FINAL_PE, - self.FINAL_TSADI, - ] - - def is_non_final(self, c: int) -> bool: - # The normal Tsadi is not a good Non-Final letter due to words like - # 'lechotet' (to chat) containing an apostrophe after the tsadi. This - # apostrophe is converted to a space in FilterWithoutEnglishLetters - # causing the Non-Final tsadi to appear at an end of a word even - # though this is not the case in the original text. - # The letters Pe and Kaf rarely display a related behavior of not being - # a good Non-Final letter. Words like 'Pop', 'Winamp' and 'Mubarak' - # for example legally end with a Non-Final Pe or Kaf. However, the - # benefit of these letters as Non-Final letters outweighs the damage - # since these words are quite rare. - return c in [self.NORMAL_KAF, self.NORMAL_MEM, self.NORMAL_NUN, self.NORMAL_PE] - - def feed(self, byte_str: Union[bytes, bytearray]) -> ProbingState: - # Final letter analysis for logical-visual decision. - # Look for evidence that the received buffer is either logical Hebrew - # or visual Hebrew. - # The following cases are checked: - # 1) A word longer than 1 letter, ending with a final letter. This is - # an indication that the text is laid out "naturally" since the - # final letter really appears at the end. +1 for logical score. - # 2) A word longer than 1 letter, ending with a Non-Final letter. In - # normal Hebrew, words ending with Kaf, Mem, Nun, Pe or Tsadi, - # should not end with the Non-Final form of that letter. Exceptions - # to this rule are mentioned above in isNonFinal(). This is an - # indication that the text is laid out backwards. +1 for visual - # score - # 3) A word longer than 1 letter, starting with a final letter. Final - # letters should not appear at the beginning of a word. This is an - # indication that the text is laid out backwards. +1 for visual - # score. - # - # The visual score and logical score are accumulated throughout the - # text and are finally checked against each other in GetCharSetName(). - # No checking for final letters in the middle of words is done since - # that case is not an indication for either Logical or Visual text. - # - # We automatically filter out all 7-bit characters (replace them with - # spaces) so the word boundary detection works properly. [MAP] - - if self.state == ProbingState.NOT_ME: - # Both model probers say it's not them. No reason to continue. - return ProbingState.NOT_ME - - byte_str = self.filter_high_byte_only(byte_str) - - for cur in byte_str: - if cur == self.SPACE: - # We stand on a space - a word just ended - if self._before_prev != self.SPACE: - # next-to-last char was not a space so self._prev is not a - # 1 letter word - if self.is_final(self._prev): - # case (1) [-2:not space][-1:final letter][cur:space] - self._final_char_logical_score += 1 - elif self.is_non_final(self._prev): - # case (2) [-2:not space][-1:Non-Final letter][ - # cur:space] - self._final_char_visual_score += 1 - else: - # Not standing on a space - if ( - (self._before_prev == self.SPACE) - and (self.is_final(self._prev)) - and (cur != self.SPACE) - ): - # case (3) [-2:space][-1:final letter][cur:not space] - self._final_char_visual_score += 1 - self._before_prev = self._prev - self._prev = cur - - # Forever detecting, till the end or until both model probers return - # ProbingState.NOT_ME (handled above) - return ProbingState.DETECTING - - @property - def charset_name(self) -> str: - assert self._logical_prober is not None - assert self._visual_prober is not None - - # Make the decision: is it Logical or Visual? - # If the final letter score distance is dominant enough, rely on it. - finalsub = self._final_char_logical_score - self._final_char_visual_score - if finalsub >= self.MIN_FINAL_CHAR_DISTANCE: - return self.LOGICAL_HEBREW_NAME - if finalsub <= -self.MIN_FINAL_CHAR_DISTANCE: - return self.VISUAL_HEBREW_NAME - - # It's not dominant enough, try to rely on the model scores instead. - modelsub = ( - self._logical_prober.get_confidence() - self._visual_prober.get_confidence() - ) - if modelsub > self.MIN_MODEL_DISTANCE: - return self.LOGICAL_HEBREW_NAME - if modelsub < -self.MIN_MODEL_DISTANCE: - return self.VISUAL_HEBREW_NAME - - # Still no good, back to final letter distance, maybe it'll save the - # day. - if finalsub < 0.0: - return self.VISUAL_HEBREW_NAME - - # (finalsub > 0 - Logical) or (don't know what to do) default to - # Logical. - return self.LOGICAL_HEBREW_NAME - - @property - def language(self) -> str: - return "Hebrew" - - @property - def state(self) -> ProbingState: - assert self._logical_prober is not None - assert self._visual_prober is not None - - # Remain active as long as any of the model probers are active. - if (self._logical_prober.state == ProbingState.NOT_ME) and ( - self._visual_prober.state == ProbingState.NOT_ME - ): - return ProbingState.NOT_ME - return ProbingState.DETECTING diff --git a/venv/lib/python3.10/site-packages/chardet/jisfreq.py b/venv/lib/python3.10/site-packages/chardet/jisfreq.py deleted file mode 100644 index 3293576e012a1c931b5e89ebc065c67b65941084..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/jisfreq.py +++ /dev/null @@ -1,325 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Communicator client code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -# Sampling from about 20M text materials include literature and computer technology -# -# Japanese frequency table, applied to both S-JIS and EUC-JP -# They are sorted in order. - -# 128 --> 0.77094 -# 256 --> 0.85710 -# 512 --> 0.92635 -# 1024 --> 0.97130 -# 2048 --> 0.99431 -# -# Ideal Distribution Ratio = 0.92635 / (1-0.92635) = 12.58 -# Random Distribution Ration = 512 / (2965+62+83+86-512) = 0.191 -# -# Typical Distribution Ratio, 25% of IDR - -JIS_TYPICAL_DISTRIBUTION_RATIO = 3.0 - -# Char to FreqOrder table , -JIS_TABLE_SIZE = 4368 - -# fmt: off -JIS_CHAR_TO_FREQ_ORDER = ( - 40, 1, 6, 182, 152, 180, 295,2127, 285, 381,3295,4304,3068,4606,3165,3510, # 16 -3511,1822,2785,4607,1193,2226,5070,4608, 171,2996,1247, 18, 179,5071, 856,1661, # 32 -1262,5072, 619, 127,3431,3512,3230,1899,1700, 232, 228,1294,1298, 284, 283,2041, # 48 -2042,1061,1062, 48, 49, 44, 45, 433, 434,1040,1041, 996, 787,2997,1255,4305, # 64 -2108,4609,1684,1648,5073,5074,5075,5076,5077,5078,3687,5079,4610,5080,3927,3928, # 80 -5081,3296,3432, 290,2285,1471,2187,5082,2580,2825,1303,2140,1739,1445,2691,3375, # 96 -1691,3297,4306,4307,4611, 452,3376,1182,2713,3688,3069,4308,5083,5084,5085,5086, # 112 -5087,5088,5089,5090,5091,5092,5093,5094,5095,5096,5097,5098,5099,5100,5101,5102, # 128 -5103,5104,5105,5106,5107,5108,5109,5110,5111,5112,4097,5113,5114,5115,5116,5117, # 144 -5118,5119,5120,5121,5122,5123,5124,5125,5126,5127,5128,5129,5130,5131,5132,5133, # 160 -5134,5135,5136,5137,5138,5139,5140,5141,5142,5143,5144,5145,5146,5147,5148,5149, # 176 -5150,5151,5152,4612,5153,5154,5155,5156,5157,5158,5159,5160,5161,5162,5163,5164, # 192 -5165,5166,5167,5168,5169,5170,5171,5172,5173,5174,5175,1472, 598, 618, 820,1205, # 208 -1309,1412,1858,1307,1692,5176,5177,5178,5179,5180,5181,5182,1142,1452,1234,1172, # 224 -1875,2043,2149,1793,1382,2973, 925,2404,1067,1241, 960,1377,2935,1491, 919,1217, # 240 -1865,2030,1406,1499,2749,4098,5183,5184,5185,5186,5187,5188,2561,4099,3117,1804, # 256 -2049,3689,4309,3513,1663,5189,3166,3118,3298,1587,1561,3433,5190,3119,1625,2998, # 272 -3299,4613,1766,3690,2786,4614,5191,5192,5193,5194,2161, 26,3377, 2,3929, 20, # 288 -3691, 47,4100, 50, 17, 16, 35, 268, 27, 243, 42, 155, 24, 154, 29, 184, # 304 - 4, 91, 14, 92, 53, 396, 33, 289, 9, 37, 64, 620, 21, 39, 321, 5, # 320 - 12, 11, 52, 13, 3, 208, 138, 0, 7, 60, 526, 141, 151,1069, 181, 275, # 336 -1591, 83, 132,1475, 126, 331, 829, 15, 69, 160, 59, 22, 157, 55,1079, 312, # 352 - 109, 38, 23, 25, 10, 19, 79,5195, 61, 382,1124, 8, 30,5196,5197,5198, # 368 -5199,5200,5201,5202,5203,5204,5205,5206, 89, 62, 74, 34,2416, 112, 139, 196, # 384 - 271, 149, 84, 607, 131, 765, 46, 88, 153, 683, 76, 874, 101, 258, 57, 80, # 400 - 32, 364, 121,1508, 169,1547, 68, 235, 145,2999, 41, 360,3027, 70, 63, 31, # 416 - 43, 259, 262,1383, 99, 533, 194, 66, 93, 846, 217, 192, 56, 106, 58, 565, # 432 - 280, 272, 311, 256, 146, 82, 308, 71, 100, 128, 214, 655, 110, 261, 104,1140, # 448 - 54, 51, 36, 87, 67,3070, 185,2618,2936,2020, 28,1066,2390,2059,5207,5208, # 464 -5209,5210,5211,5212,5213,5214,5215,5216,4615,5217,5218,5219,5220,5221,5222,5223, # 480 -5224,5225,5226,5227,5228,5229,5230,5231,5232,5233,5234,5235,5236,3514,5237,5238, # 496 -5239,5240,5241,5242,5243,5244,2297,2031,4616,4310,3692,5245,3071,5246,3598,5247, # 512 -4617,3231,3515,5248,4101,4311,4618,3808,4312,4102,5249,4103,4104,3599,5250,5251, # 528 -5252,5253,5254,5255,5256,5257,5258,5259,5260,5261,5262,5263,5264,5265,5266,5267, # 544 -5268,5269,5270,5271,5272,5273,5274,5275,5276,5277,5278,5279,5280,5281,5282,5283, # 560 -5284,5285,5286,5287,5288,5289,5290,5291,5292,5293,5294,5295,5296,5297,5298,5299, # 576 -5300,5301,5302,5303,5304,5305,5306,5307,5308,5309,5310,5311,5312,5313,5314,5315, # 592 -5316,5317,5318,5319,5320,5321,5322,5323,5324,5325,5326,5327,5328,5329,5330,5331, # 608 -5332,5333,5334,5335,5336,5337,5338,5339,5340,5341,5342,5343,5344,5345,5346,5347, # 624 -5348,5349,5350,5351,5352,5353,5354,5355,5356,5357,5358,5359,5360,5361,5362,5363, # 640 -5364,5365,5366,5367,5368,5369,5370,5371,5372,5373,5374,5375,5376,5377,5378,5379, # 656 -5380,5381, 363, 642,2787,2878,2788,2789,2316,3232,2317,3434,2011, 165,1942,3930, # 672 -3931,3932,3933,5382,4619,5383,4620,5384,5385,5386,5387,5388,5389,5390,5391,5392, # 688 -5393,5394,5395,5396,5397,5398,5399,5400,5401,5402,5403,5404,5405,5406,5407,5408, # 704 -5409,5410,5411,5412,5413,5414,5415,5416,5417,5418,5419,5420,5421,5422,5423,5424, # 720 -5425,5426,5427,5428,5429,5430,5431,5432,5433,5434,5435,5436,5437,5438,5439,5440, # 736 -5441,5442,5443,5444,5445,5446,5447,5448,5449,5450,5451,5452,5453,5454,5455,5456, # 752 -5457,5458,5459,5460,5461,5462,5463,5464,5465,5466,5467,5468,5469,5470,5471,5472, # 768 -5473,5474,5475,5476,5477,5478,5479,5480,5481,5482,5483,5484,5485,5486,5487,5488, # 784 -5489,5490,5491,5492,5493,5494,5495,5496,5497,5498,5499,5500,5501,5502,5503,5504, # 800 -5505,5506,5507,5508,5509,5510,5511,5512,5513,5514,5515,5516,5517,5518,5519,5520, # 816 -5521,5522,5523,5524,5525,5526,5527,5528,5529,5530,5531,5532,5533,5534,5535,5536, # 832 -5537,5538,5539,5540,5541,5542,5543,5544,5545,5546,5547,5548,5549,5550,5551,5552, # 848 -5553,5554,5555,5556,5557,5558,5559,5560,5561,5562,5563,5564,5565,5566,5567,5568, # 864 -5569,5570,5571,5572,5573,5574,5575,5576,5577,5578,5579,5580,5581,5582,5583,5584, # 880 -5585,5586,5587,5588,5589,5590,5591,5592,5593,5594,5595,5596,5597,5598,5599,5600, # 896 -5601,5602,5603,5604,5605,5606,5607,5608,5609,5610,5611,5612,5613,5614,5615,5616, # 912 -5617,5618,5619,5620,5621,5622,5623,5624,5625,5626,5627,5628,5629,5630,5631,5632, # 928 -5633,5634,5635,5636,5637,5638,5639,5640,5641,5642,5643,5644,5645,5646,5647,5648, # 944 -5649,5650,5651,5652,5653,5654,5655,5656,5657,5658,5659,5660,5661,5662,5663,5664, # 960 -5665,5666,5667,5668,5669,5670,5671,5672,5673,5674,5675,5676,5677,5678,5679,5680, # 976 -5681,5682,5683,5684,5685,5686,5687,5688,5689,5690,5691,5692,5693,5694,5695,5696, # 992 -5697,5698,5699,5700,5701,5702,5703,5704,5705,5706,5707,5708,5709,5710,5711,5712, # 1008 -5713,5714,5715,5716,5717,5718,5719,5720,5721,5722,5723,5724,5725,5726,5727,5728, # 1024 -5729,5730,5731,5732,5733,5734,5735,5736,5737,5738,5739,5740,5741,5742,5743,5744, # 1040 -5745,5746,5747,5748,5749,5750,5751,5752,5753,5754,5755,5756,5757,5758,5759,5760, # 1056 -5761,5762,5763,5764,5765,5766,5767,5768,5769,5770,5771,5772,5773,5774,5775,5776, # 1072 -5777,5778,5779,5780,5781,5782,5783,5784,5785,5786,5787,5788,5789,5790,5791,5792, # 1088 -5793,5794,5795,5796,5797,5798,5799,5800,5801,5802,5803,5804,5805,5806,5807,5808, # 1104 -5809,5810,5811,5812,5813,5814,5815,5816,5817,5818,5819,5820,5821,5822,5823,5824, # 1120 -5825,5826,5827,5828,5829,5830,5831,5832,5833,5834,5835,5836,5837,5838,5839,5840, # 1136 -5841,5842,5843,5844,5845,5846,5847,5848,5849,5850,5851,5852,5853,5854,5855,5856, # 1152 -5857,5858,5859,5860,5861,5862,5863,5864,5865,5866,5867,5868,5869,5870,5871,5872, # 1168 -5873,5874,5875,5876,5877,5878,5879,5880,5881,5882,5883,5884,5885,5886,5887,5888, # 1184 -5889,5890,5891,5892,5893,5894,5895,5896,5897,5898,5899,5900,5901,5902,5903,5904, # 1200 -5905,5906,5907,5908,5909,5910,5911,5912,5913,5914,5915,5916,5917,5918,5919,5920, # 1216 -5921,5922,5923,5924,5925,5926,5927,5928,5929,5930,5931,5932,5933,5934,5935,5936, # 1232 -5937,5938,5939,5940,5941,5942,5943,5944,5945,5946,5947,5948,5949,5950,5951,5952, # 1248 -5953,5954,5955,5956,5957,5958,5959,5960,5961,5962,5963,5964,5965,5966,5967,5968, # 1264 -5969,5970,5971,5972,5973,5974,5975,5976,5977,5978,5979,5980,5981,5982,5983,5984, # 1280 -5985,5986,5987,5988,5989,5990,5991,5992,5993,5994,5995,5996,5997,5998,5999,6000, # 1296 -6001,6002,6003,6004,6005,6006,6007,6008,6009,6010,6011,6012,6013,6014,6015,6016, # 1312 -6017,6018,6019,6020,6021,6022,6023,6024,6025,6026,6027,6028,6029,6030,6031,6032, # 1328 -6033,6034,6035,6036,6037,6038,6039,6040,6041,6042,6043,6044,6045,6046,6047,6048, # 1344 -6049,6050,6051,6052,6053,6054,6055,6056,6057,6058,6059,6060,6061,6062,6063,6064, # 1360 -6065,6066,6067,6068,6069,6070,6071,6072,6073,6074,6075,6076,6077,6078,6079,6080, # 1376 -6081,6082,6083,6084,6085,6086,6087,6088,6089,6090,6091,6092,6093,6094,6095,6096, # 1392 -6097,6098,6099,6100,6101,6102,6103,6104,6105,6106,6107,6108,6109,6110,6111,6112, # 1408 -6113,6114,2044,2060,4621, 997,1235, 473,1186,4622, 920,3378,6115,6116, 379,1108, # 1424 -4313,2657,2735,3934,6117,3809, 636,3233, 573,1026,3693,3435,2974,3300,2298,4105, # 1440 - 854,2937,2463, 393,2581,2417, 539, 752,1280,2750,2480, 140,1161, 440, 708,1569, # 1456 - 665,2497,1746,1291,1523,3000, 164,1603, 847,1331, 537,1997, 486, 508,1693,2418, # 1472 -1970,2227, 878,1220, 299,1030, 969, 652,2751, 624,1137,3301,2619, 65,3302,2045, # 1488 -1761,1859,3120,1930,3694,3516, 663,1767, 852, 835,3695, 269, 767,2826,2339,1305, # 1504 - 896,1150, 770,1616,6118, 506,1502,2075,1012,2519, 775,2520,2975,2340,2938,4314, # 1520 -3028,2086,1224,1943,2286,6119,3072,4315,2240,1273,1987,3935,1557, 175, 597, 985, # 1536 -3517,2419,2521,1416,3029, 585, 938,1931,1007,1052,1932,1685,6120,3379,4316,4623, # 1552 - 804, 599,3121,1333,2128,2539,1159,1554,2032,3810, 687,2033,2904, 952, 675,1467, # 1568 -3436,6121,2241,1096,1786,2440,1543,1924, 980,1813,2228, 781,2692,1879, 728,1918, # 1584 -3696,4624, 548,1950,4625,1809,1088,1356,3303,2522,1944, 502, 972, 373, 513,2827, # 1600 - 586,2377,2391,1003,1976,1631,6122,2464,1084, 648,1776,4626,2141, 324, 962,2012, # 1616 -2177,2076,1384, 742,2178,1448,1173,1810, 222, 102, 301, 445, 125,2420, 662,2498, # 1632 - 277, 200,1476,1165,1068, 224,2562,1378,1446, 450,1880, 659, 791, 582,4627,2939, # 1648 -3936,1516,1274, 555,2099,3697,1020,1389,1526,3380,1762,1723,1787,2229, 412,2114, # 1664 -1900,2392,3518, 512,2597, 427,1925,2341,3122,1653,1686,2465,2499, 697, 330, 273, # 1680 - 380,2162, 951, 832, 780, 991,1301,3073, 965,2270,3519, 668,2523,2636,1286, 535, # 1696 -1407, 518, 671, 957,2658,2378, 267, 611,2197,3030,6123, 248,2299, 967,1799,2356, # 1712 - 850,1418,3437,1876,1256,1480,2828,1718,6124,6125,1755,1664,2405,6126,4628,2879, # 1728 -2829, 499,2179, 676,4629, 557,2329,2214,2090, 325,3234, 464, 811,3001, 992,2342, # 1744 -2481,1232,1469, 303,2242, 466,1070,2163, 603,1777,2091,4630,2752,4631,2714, 322, # 1760 -2659,1964,1768, 481,2188,1463,2330,2857,3600,2092,3031,2421,4632,2318,2070,1849, # 1776 -2598,4633,1302,2254,1668,1701,2422,3811,2905,3032,3123,2046,4106,1763,1694,4634, # 1792 -1604, 943,1724,1454, 917, 868,2215,1169,2940, 552,1145,1800,1228,1823,1955, 316, # 1808 -1080,2510, 361,1807,2830,4107,2660,3381,1346,1423,1134,4108,6127, 541,1263,1229, # 1824 -1148,2540, 545, 465,1833,2880,3438,1901,3074,2482, 816,3937, 713,1788,2500, 122, # 1840 -1575, 195,1451,2501,1111,6128, 859, 374,1225,2243,2483,4317, 390,1033,3439,3075, # 1856 -2524,1687, 266, 793,1440,2599, 946, 779, 802, 507, 897,1081, 528,2189,1292, 711, # 1872 -1866,1725,1167,1640, 753, 398,2661,1053, 246, 348,4318, 137,1024,3440,1600,2077, # 1888 -2129, 825,4319, 698, 238, 521, 187,2300,1157,2423,1641,1605,1464,1610,1097,2541, # 1904 -1260,1436, 759,2255,1814,2150, 705,3235, 409,2563,3304, 561,3033,2005,2564, 726, # 1920 -1956,2343,3698,4109, 949,3812,3813,3520,1669, 653,1379,2525, 881,2198, 632,2256, # 1936 -1027, 778,1074, 733,1957, 514,1481,2466, 554,2180, 702,3938,1606,1017,1398,6129, # 1952 -1380,3521, 921, 993,1313, 594, 449,1489,1617,1166, 768,1426,1360, 495,1794,3601, # 1968 -1177,3602,1170,4320,2344, 476, 425,3167,4635,3168,1424, 401,2662,1171,3382,1998, # 1984 -1089,4110, 477,3169, 474,6130,1909, 596,2831,1842, 494, 693,1051,1028,1207,3076, # 2000 - 606,2115, 727,2790,1473,1115, 743,3522, 630, 805,1532,4321,2021, 366,1057, 838, # 2016 - 684,1114,2142,4322,2050,1492,1892,1808,2271,3814,2424,1971,1447,1373,3305,1090, # 2032 -1536,3939,3523,3306,1455,2199, 336, 369,2331,1035, 584,2393, 902, 718,2600,6131, # 2048 -2753, 463,2151,1149,1611,2467, 715,1308,3124,1268, 343,1413,3236,1517,1347,2663, # 2064 -2093,3940,2022,1131,1553,2100,2941,1427,3441,2942,1323,2484,6132,1980, 872,2368, # 2080 -2441,2943, 320,2369,2116,1082, 679,1933,3941,2791,3815, 625,1143,2023, 422,2200, # 2096 -3816,6133, 730,1695, 356,2257,1626,2301,2858,2637,1627,1778, 937, 883,2906,2693, # 2112 -3002,1769,1086, 400,1063,1325,3307,2792,4111,3077, 456,2345,1046, 747,6134,1524, # 2128 - 884,1094,3383,1474,2164,1059, 974,1688,2181,2258,1047, 345,1665,1187, 358, 875, # 2144 -3170, 305, 660,3524,2190,1334,1135,3171,1540,1649,2542,1527, 927, 968,2793, 885, # 2160 -1972,1850, 482, 500,2638,1218,1109,1085,2543,1654,2034, 876, 78,2287,1482,1277, # 2176 - 861,1675,1083,1779, 724,2754, 454, 397,1132,1612,2332, 893, 672,1237, 257,2259, # 2192 -2370, 135,3384, 337,2244, 547, 352, 340, 709,2485,1400, 788,1138,2511, 540, 772, # 2208 -1682,2260,2272,2544,2013,1843,1902,4636,1999,1562,2288,4637,2201,1403,1533, 407, # 2224 - 576,3308,1254,2071, 978,3385, 170, 136,1201,3125,2664,3172,2394, 213, 912, 873, # 2240 -3603,1713,2202, 699,3604,3699, 813,3442, 493, 531,1054, 468,2907,1483, 304, 281, # 2256 -4112,1726,1252,2094, 339,2319,2130,2639, 756,1563,2944, 748, 571,2976,1588,2425, # 2272 -2715,1851,1460,2426,1528,1392,1973,3237, 288,3309, 685,3386, 296, 892,2716,2216, # 2288 -1570,2245, 722,1747,2217, 905,3238,1103,6135,1893,1441,1965, 251,1805,2371,3700, # 2304 -2601,1919,1078, 75,2182,1509,1592,1270,2640,4638,2152,6136,3310,3817, 524, 706, # 2320 -1075, 292,3818,1756,2602, 317, 98,3173,3605,3525,1844,2218,3819,2502, 814, 567, # 2336 - 385,2908,1534,6137, 534,1642,3239, 797,6138,1670,1529, 953,4323, 188,1071, 538, # 2352 - 178, 729,3240,2109,1226,1374,2000,2357,2977, 731,2468,1116,2014,2051,6139,1261, # 2368 -1593, 803,2859,2736,3443, 556, 682, 823,1541,6140,1369,2289,1706,2794, 845, 462, # 2384 -2603,2665,1361, 387, 162,2358,1740, 739,1770,1720,1304,1401,3241,1049, 627,1571, # 2400 -2427,3526,1877,3942,1852,1500, 431,1910,1503, 677, 297,2795, 286,1433,1038,1198, # 2416 -2290,1133,1596,4113,4639,2469,1510,1484,3943,6141,2442, 108, 712,4640,2372, 866, # 2432 -3701,2755,3242,1348, 834,1945,1408,3527,2395,3243,1811, 824, 994,1179,2110,1548, # 2448 -1453, 790,3003, 690,4324,4325,2832,2909,3820,1860,3821, 225,1748, 310, 346,1780, # 2464 -2470, 821,1993,2717,2796, 828, 877,3528,2860,2471,1702,2165,2910,2486,1789, 453, # 2480 - 359,2291,1676, 73,1164,1461,1127,3311, 421, 604, 314,1037, 589, 116,2487, 737, # 2496 - 837,1180, 111, 244, 735,6142,2261,1861,1362, 986, 523, 418, 581,2666,3822, 103, # 2512 - 855, 503,1414,1867,2488,1091, 657,1597, 979, 605,1316,4641,1021,2443,2078,2001, # 2528 -1209, 96, 587,2166,1032, 260,1072,2153, 173, 94, 226,3244, 819,2006,4642,4114, # 2544 -2203, 231,1744, 782, 97,2667, 786,3387, 887, 391, 442,2219,4326,1425,6143,2694, # 2560 - 633,1544,1202, 483,2015, 592,2052,1958,2472,1655, 419, 129,4327,3444,3312,1714, # 2576 -1257,3078,4328,1518,1098, 865,1310,1019,1885,1512,1734, 469,2444, 148, 773, 436, # 2592 -1815,1868,1128,1055,4329,1245,2756,3445,2154,1934,1039,4643, 579,1238, 932,2320, # 2608 - 353, 205, 801, 115,2428, 944,2321,1881, 399,2565,1211, 678, 766,3944, 335,2101, # 2624 -1459,1781,1402,3945,2737,2131,1010, 844, 981,1326,1013, 550,1816,1545,2620,1335, # 2640 -1008, 371,2881, 936,1419,1613,3529,1456,1395,2273,1834,2604,1317,2738,2503, 416, # 2656 -1643,4330, 806,1126, 229, 591,3946,1314,1981,1576,1837,1666, 347,1790, 977,3313, # 2672 - 764,2861,1853, 688,2429,1920,1462, 77, 595, 415,2002,3034, 798,1192,4115,6144, # 2688 -2978,4331,3035,2695,2582,2072,2566, 430,2430,1727, 842,1396,3947,3702, 613, 377, # 2704 - 278, 236,1417,3388,3314,3174, 757,1869, 107,3530,6145,1194, 623,2262, 207,1253, # 2720 -2167,3446,3948, 492,1117,1935, 536,1838,2757,1246,4332, 696,2095,2406,1393,1572, # 2736 -3175,1782, 583, 190, 253,1390,2230, 830,3126,3389, 934,3245,1703,1749,2979,1870, # 2752 -2545,1656,2204, 869,2346,4116,3176,1817, 496,1764,4644, 942,1504, 404,1903,1122, # 2768 -1580,3606,2945,1022, 515, 372,1735, 955,2431,3036,6146,2797,1110,2302,2798, 617, # 2784 -6147, 441, 762,1771,3447,3607,3608,1904, 840,3037, 86, 939,1385, 572,1370,2445, # 2800 -1336, 114,3703, 898, 294, 203,3315, 703,1583,2274, 429, 961,4333,1854,1951,3390, # 2816 -2373,3704,4334,1318,1381, 966,1911,2322,1006,1155, 309, 989, 458,2718,1795,1372, # 2832 -1203, 252,1689,1363,3177, 517,1936, 168,1490, 562, 193,3823,1042,4117,1835, 551, # 2848 - 470,4645, 395, 489,3448,1871,1465,2583,2641, 417,1493, 279,1295, 511,1236,1119, # 2864 - 72,1231,1982,1812,3004, 871,1564, 984,3449,1667,2696,2096,4646,2347,2833,1673, # 2880 -3609, 695,3246,2668, 807,1183,4647, 890, 388,2333,1801,1457,2911,1765,1477,1031, # 2896 -3316,3317,1278,3391,2799,2292,2526, 163,3450,4335,2669,1404,1802,6148,2323,2407, # 2912 -1584,1728,1494,1824,1269, 298, 909,3318,1034,1632, 375, 776,1683,2061, 291, 210, # 2928 -1123, 809,1249,1002,2642,3038, 206,1011,2132, 144, 975, 882,1565, 342, 667, 754, # 2944 -1442,2143,1299,2303,2062, 447, 626,2205,1221,2739,2912,1144,1214,2206,2584, 760, # 2960 -1715, 614, 950,1281,2670,2621, 810, 577,1287,2546,4648, 242,2168, 250,2643, 691, # 2976 - 123,2644, 647, 313,1029, 689,1357,2946,1650, 216, 771,1339,1306, 808,2063, 549, # 2992 - 913,1371,2913,2914,6149,1466,1092,1174,1196,1311,2605,2396,1783,1796,3079, 406, # 3008 -2671,2117,3949,4649, 487,1825,2220,6150,2915, 448,2348,1073,6151,2397,1707, 130, # 3024 - 900,1598, 329, 176,1959,2527,1620,6152,2275,4336,3319,1983,2191,3705,3610,2155, # 3040 -3706,1912,1513,1614,6153,1988, 646, 392,2304,1589,3320,3039,1826,1239,1352,1340, # 3056 -2916, 505,2567,1709,1437,2408,2547, 906,6154,2672, 384,1458,1594,1100,1329, 710, # 3072 - 423,3531,2064,2231,2622,1989,2673,1087,1882, 333, 841,3005,1296,2882,2379, 580, # 3088 -1937,1827,1293,2585, 601, 574, 249,1772,4118,2079,1120, 645, 901,1176,1690, 795, # 3104 -2207, 478,1434, 516,1190,1530, 761,2080, 930,1264, 355, 435,1552, 644,1791, 987, # 3120 - 220,1364,1163,1121,1538, 306,2169,1327,1222, 546,2645, 218, 241, 610,1704,3321, # 3136 -1984,1839,1966,2528, 451,6155,2586,3707,2568, 907,3178, 254,2947, 186,1845,4650, # 3152 - 745, 432,1757, 428,1633, 888,2246,2221,2489,3611,2118,1258,1265, 956,3127,1784, # 3168 -4337,2490, 319, 510, 119, 457,3612, 274,2035,2007,4651,1409,3128, 970,2758, 590, # 3184 -2800, 661,2247,4652,2008,3950,1420,1549,3080,3322,3951,1651,1375,2111, 485,2491, # 3200 -1429,1156,6156,2548,2183,1495, 831,1840,2529,2446, 501,1657, 307,1894,3247,1341, # 3216 - 666, 899,2156,1539,2549,1559, 886, 349,2208,3081,2305,1736,3824,2170,2759,1014, # 3232 -1913,1386, 542,1397,2948, 490, 368, 716, 362, 159, 282,2569,1129,1658,1288,1750, # 3248 -2674, 276, 649,2016, 751,1496, 658,1818,1284,1862,2209,2087,2512,3451, 622,2834, # 3264 - 376, 117,1060,2053,1208,1721,1101,1443, 247,1250,3179,1792,3952,2760,2398,3953, # 3280 -6157,2144,3708, 446,2432,1151,2570,3452,2447,2761,2835,1210,2448,3082, 424,2222, # 3296 -1251,2449,2119,2836, 504,1581,4338, 602, 817, 857,3825,2349,2306, 357,3826,1470, # 3312 -1883,2883, 255, 958, 929,2917,3248, 302,4653,1050,1271,1751,2307,1952,1430,2697, # 3328 -2719,2359, 354,3180, 777, 158,2036,4339,1659,4340,4654,2308,2949,2248,1146,2232, # 3344 -3532,2720,1696,2623,3827,6158,3129,1550,2698,1485,1297,1428, 637, 931,2721,2145, # 3360 - 914,2550,2587, 81,2450, 612, 827,2646,1242,4655,1118,2884, 472,1855,3181,3533, # 3376 -3534, 569,1353,2699,1244,1758,2588,4119,2009,2762,2171,3709,1312,1531,6159,1152, # 3392 -1938, 134,1830, 471,3710,2276,1112,1535,3323,3453,3535, 982,1337,2950, 488, 826, # 3408 - 674,1058,1628,4120,2017, 522,2399, 211, 568,1367,3454, 350, 293,1872,1139,3249, # 3424 -1399,1946,3006,1300,2360,3324, 588, 736,6160,2606, 744, 669,3536,3828,6161,1358, # 3440 - 199, 723, 848, 933, 851,1939,1505,1514,1338,1618,1831,4656,1634,3613, 443,2740, # 3456 -3829, 717,1947, 491,1914,6162,2551,1542,4121,1025,6163,1099,1223, 198,3040,2722, # 3472 - 370, 410,1905,2589, 998,1248,3182,2380, 519,1449,4122,1710, 947, 928,1153,4341, # 3488 -2277, 344,2624,1511, 615, 105, 161,1212,1076,1960,3130,2054,1926,1175,1906,2473, # 3504 - 414,1873,2801,6164,2309, 315,1319,3325, 318,2018,2146,2157, 963, 631, 223,4342, # 3520 -4343,2675, 479,3711,1197,2625,3712,2676,2361,6165,4344,4123,6166,2451,3183,1886, # 3536 -2184,1674,1330,1711,1635,1506, 799, 219,3250,3083,3954,1677,3713,3326,2081,3614, # 3552 -1652,2073,4657,1147,3041,1752, 643,1961, 147,1974,3955,6167,1716,2037, 918,3007, # 3568 -1994, 120,1537, 118, 609,3184,4345, 740,3455,1219, 332,1615,3830,6168,1621,2980, # 3584 -1582, 783, 212, 553,2350,3714,1349,2433,2082,4124, 889,6169,2310,1275,1410, 973, # 3600 - 166,1320,3456,1797,1215,3185,2885,1846,2590,2763,4658, 629, 822,3008, 763, 940, # 3616 -1990,2862, 439,2409,1566,1240,1622, 926,1282,1907,2764, 654,2210,1607, 327,1130, # 3632 -3956,1678,1623,6170,2434,2192, 686, 608,3831,3715, 903,3957,3042,6171,2741,1522, # 3648 -1915,1105,1555,2552,1359, 323,3251,4346,3457, 738,1354,2553,2311,2334,1828,2003, # 3664 -3832,1753,2351,1227,6172,1887,4125,1478,6173,2410,1874,1712,1847, 520,1204,2607, # 3680 - 264,4659, 836,2677,2102, 600,4660,3833,2278,3084,6174,4347,3615,1342, 640, 532, # 3696 - 543,2608,1888,2400,2591,1009,4348,1497, 341,1737,3616,2723,1394, 529,3252,1321, # 3712 - 983,4661,1515,2120, 971,2592, 924, 287,1662,3186,4349,2700,4350,1519, 908,1948, # 3728 -2452, 156, 796,1629,1486,2223,2055, 694,4126,1259,1036,3392,1213,2249,2742,1889, # 3744 -1230,3958,1015, 910, 408, 559,3617,4662, 746, 725, 935,4663,3959,3009,1289, 563, # 3760 - 867,4664,3960,1567,2981,2038,2626, 988,2263,2381,4351, 143,2374, 704,1895,6175, # 3776 -1188,3716,2088, 673,3085,2362,4352, 484,1608,1921,2765,2918, 215, 904,3618,3537, # 3792 - 894, 509, 976,3043,2701,3961,4353,2837,2982, 498,6176,6177,1102,3538,1332,3393, # 3808 -1487,1636,1637, 233, 245,3962, 383, 650, 995,3044, 460,1520,1206,2352, 749,3327, # 3824 - 530, 700, 389,1438,1560,1773,3963,2264, 719,2951,2724,3834, 870,1832,1644,1000, # 3840 - 839,2474,3717, 197,1630,3394, 365,2886,3964,1285,2133, 734, 922, 818,1106, 732, # 3856 - 480,2083,1774,3458, 923,2279,1350, 221,3086, 85,2233,2234,3835,1585,3010,2147, # 3872 -1387,1705,2382,1619,2475, 133, 239,2802,1991,1016,2084,2383, 411,2838,1113, 651, # 3888 -1985,1160,3328, 990,1863,3087,1048,1276,2647, 265,2627,1599,3253,2056, 150, 638, # 3904 -2019, 656, 853, 326,1479, 680,1439,4354,1001,1759, 413,3459,3395,2492,1431, 459, # 3920 -4355,1125,3329,2265,1953,1450,2065,2863, 849, 351,2678,3131,3254,3255,1104,1577, # 3936 - 227,1351,1645,2453,2193,1421,2887, 812,2121, 634, 95,2435, 201,2312,4665,1646, # 3952 -1671,2743,1601,2554,2702,2648,2280,1315,1366,2089,3132,1573,3718,3965,1729,1189, # 3968 - 328,2679,1077,1940,1136, 558,1283, 964,1195, 621,2074,1199,1743,3460,3619,1896, # 3984 -1916,1890,3836,2952,1154,2112,1064, 862, 378,3011,2066,2113,2803,1568,2839,6178, # 4000 -3088,2919,1941,1660,2004,1992,2194, 142, 707,1590,1708,1624,1922,1023,1836,1233, # 4016 -1004,2313, 789, 741,3620,6179,1609,2411,1200,4127,3719,3720,4666,2057,3721, 593, # 4032 -2840, 367,2920,1878,6180,3461,1521, 628,1168, 692,2211,2649, 300, 720,2067,2571, # 4048 -2953,3396, 959,2504,3966,3539,3462,1977, 701,6181, 954,1043, 800, 681, 183,3722, # 4064 -1803,1730,3540,4128,2103, 815,2314, 174, 467, 230,2454,1093,2134, 755,3541,3397, # 4080 -1141,1162,6182,1738,2039, 270,3256,2513,1005,1647,2185,3837, 858,1679,1897,1719, # 4096 -2954,2324,1806, 402, 670, 167,4129,1498,2158,2104, 750,6183, 915, 189,1680,1551, # 4112 - 455,4356,1501,2455, 405,1095,2955, 338,1586,1266,1819, 570, 641,1324, 237,1556, # 4128 -2650,1388,3723,6184,1368,2384,1343,1978,3089,2436, 879,3724, 792,1191, 758,3012, # 4144 -1411,2135,1322,4357, 240,4667,1848,3725,1574,6185, 420,3045,1546,1391, 714,4358, # 4160 -1967, 941,1864, 863, 664, 426, 560,1731,2680,1785,2864,1949,2363, 403,3330,1415, # 4176 -1279,2136,1697,2335, 204, 721,2097,3838, 90,6186,2085,2505, 191,3967, 124,2148, # 4192 -1376,1798,1178,1107,1898,1405, 860,4359,1243,1272,2375,2983,1558,2456,1638, 113, # 4208 -3621, 578,1923,2609, 880, 386,4130, 784,2186,2266,1422,2956,2172,1722, 497, 263, # 4224 -2514,1267,2412,2610, 177,2703,3542, 774,1927,1344, 616,1432,1595,1018, 172,4360, # 4240 -2325, 911,4361, 438,1468,3622, 794,3968,2024,2173,1681,1829,2957, 945, 895,3090, # 4256 - 575,2212,2476, 475,2401,2681, 785,2744,1745,2293,2555,1975,3133,2865, 394,4668, # 4272 -3839, 635,4131, 639, 202,1507,2195,2766,1345,1435,2572,3726,1908,1184,1181,2457, # 4288 -3727,3134,4362, 843,2611, 437, 916,4669, 234, 769,1884,3046,3047,3623, 833,6187, # 4304 -1639,2250,2402,1355,1185,2010,2047, 999, 525,1732,1290,1488,2612, 948,1578,3728, # 4320 -2413,2477,1216,2725,2159, 334,3840,1328,3624,2921,1525,4132, 564,1056, 891,4363, # 4336 -1444,1698,2385,2251,3729,1365,2281,2235,1717,6188, 864,3841,2515, 444, 527,2767, # 4352 -2922,3625, 544, 461,6189, 566, 209,2437,3398,2098,1065,2068,3331,3626,3257,2137, # 4368 #last 512 -) -# fmt: on diff --git a/venv/lib/python3.10/site-packages/chardet/johabfreq.py b/venv/lib/python3.10/site-packages/chardet/johabfreq.py deleted file mode 100644 index c12969990d73388f61a6ab98fb4ee8f0f5cbc44f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/johabfreq.py +++ /dev/null @@ -1,2382 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Communicator client code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -# The frequency data itself is the same as euc-kr. -# This is just a mapping table to euc-kr. - -JOHAB_TO_EUCKR_ORDER_TABLE = { - 0x8861: 0, - 0x8862: 1, - 0x8865: 2, - 0x8868: 3, - 0x8869: 4, - 0x886A: 5, - 0x886B: 6, - 0x8871: 7, - 0x8873: 8, - 0x8874: 9, - 0x8875: 10, - 0x8876: 11, - 0x8877: 12, - 0x8878: 13, - 0x8879: 14, - 0x887B: 15, - 0x887C: 16, - 0x887D: 17, - 0x8881: 18, - 0x8882: 19, - 0x8885: 20, - 0x8889: 21, - 0x8891: 22, - 0x8893: 23, - 0x8895: 24, - 0x8896: 25, - 0x8897: 26, - 0x88A1: 27, - 0x88A2: 28, - 0x88A5: 29, - 0x88A9: 30, - 0x88B5: 31, - 0x88B7: 32, - 0x88C1: 33, - 0x88C5: 34, - 0x88C9: 35, - 0x88E1: 36, - 0x88E2: 37, - 0x88E5: 38, - 0x88E8: 39, - 0x88E9: 40, - 0x88EB: 41, - 0x88F1: 42, - 0x88F3: 43, - 0x88F5: 44, - 0x88F6: 45, - 0x88F7: 46, - 0x88F8: 47, - 0x88FB: 48, - 0x88FC: 49, - 0x88FD: 50, - 0x8941: 51, - 0x8945: 52, - 0x8949: 53, - 0x8951: 54, - 0x8953: 55, - 0x8955: 56, - 0x8956: 57, - 0x8957: 58, - 0x8961: 59, - 0x8962: 60, - 0x8963: 61, - 0x8965: 62, - 0x8968: 63, - 0x8969: 64, - 0x8971: 65, - 0x8973: 66, - 0x8975: 67, - 0x8976: 68, - 0x8977: 69, - 0x897B: 70, - 0x8981: 71, - 0x8985: 72, - 0x8989: 73, - 0x8993: 74, - 0x8995: 75, - 0x89A1: 76, - 0x89A2: 77, - 0x89A5: 78, - 0x89A8: 79, - 0x89A9: 80, - 0x89AB: 81, - 0x89AD: 82, - 0x89B0: 83, - 0x89B1: 84, - 0x89B3: 85, - 0x89B5: 86, - 0x89B7: 87, - 0x89B8: 88, - 0x89C1: 89, - 0x89C2: 90, - 0x89C5: 91, - 0x89C9: 92, - 0x89CB: 93, - 0x89D1: 94, - 0x89D3: 95, - 0x89D5: 96, - 0x89D7: 97, - 0x89E1: 98, - 0x89E5: 99, - 0x89E9: 100, - 0x89F3: 101, - 0x89F6: 102, - 0x89F7: 103, - 0x8A41: 104, - 0x8A42: 105, - 0x8A45: 106, - 0x8A49: 107, - 0x8A51: 108, - 0x8A53: 109, - 0x8A55: 110, - 0x8A57: 111, - 0x8A61: 112, - 0x8A65: 113, - 0x8A69: 114, - 0x8A73: 115, - 0x8A75: 116, - 0x8A81: 117, - 0x8A82: 118, - 0x8A85: 119, - 0x8A88: 120, - 0x8A89: 121, - 0x8A8A: 122, - 0x8A8B: 123, - 0x8A90: 124, - 0x8A91: 125, - 0x8A93: 126, - 0x8A95: 127, - 0x8A97: 128, - 0x8A98: 129, - 0x8AA1: 130, - 0x8AA2: 131, - 0x8AA5: 132, - 0x8AA9: 133, - 0x8AB6: 134, - 0x8AB7: 135, - 0x8AC1: 136, - 0x8AD5: 137, - 0x8AE1: 138, - 0x8AE2: 139, - 0x8AE5: 140, - 0x8AE9: 141, - 0x8AF1: 142, - 0x8AF3: 143, - 0x8AF5: 144, - 0x8B41: 145, - 0x8B45: 146, - 0x8B49: 147, - 0x8B61: 148, - 0x8B62: 149, - 0x8B65: 150, - 0x8B68: 151, - 0x8B69: 152, - 0x8B6A: 153, - 0x8B71: 154, - 0x8B73: 155, - 0x8B75: 156, - 0x8B77: 157, - 0x8B81: 158, - 0x8BA1: 159, - 0x8BA2: 160, - 0x8BA5: 161, - 0x8BA8: 162, - 0x8BA9: 163, - 0x8BAB: 164, - 0x8BB1: 165, - 0x8BB3: 166, - 0x8BB5: 167, - 0x8BB7: 168, - 0x8BB8: 169, - 0x8BBC: 170, - 0x8C61: 171, - 0x8C62: 172, - 0x8C63: 173, - 0x8C65: 174, - 0x8C69: 175, - 0x8C6B: 176, - 0x8C71: 177, - 0x8C73: 178, - 0x8C75: 179, - 0x8C76: 180, - 0x8C77: 181, - 0x8C7B: 182, - 0x8C81: 183, - 0x8C82: 184, - 0x8C85: 185, - 0x8C89: 186, - 0x8C91: 187, - 0x8C93: 188, - 0x8C95: 189, - 0x8C96: 190, - 0x8C97: 191, - 0x8CA1: 192, - 0x8CA2: 193, - 0x8CA9: 194, - 0x8CE1: 195, - 0x8CE2: 196, - 0x8CE3: 197, - 0x8CE5: 198, - 0x8CE9: 199, - 0x8CF1: 200, - 0x8CF3: 201, - 0x8CF5: 202, - 0x8CF6: 203, - 0x8CF7: 204, - 0x8D41: 205, - 0x8D42: 206, - 0x8D45: 207, - 0x8D51: 208, - 0x8D55: 209, - 0x8D57: 210, - 0x8D61: 211, - 0x8D65: 212, - 0x8D69: 213, - 0x8D75: 214, - 0x8D76: 215, - 0x8D7B: 216, - 0x8D81: 217, - 0x8DA1: 218, - 0x8DA2: 219, - 0x8DA5: 220, - 0x8DA7: 221, - 0x8DA9: 222, - 0x8DB1: 223, - 0x8DB3: 224, - 0x8DB5: 225, - 0x8DB7: 226, - 0x8DB8: 227, - 0x8DB9: 228, - 0x8DC1: 229, - 0x8DC2: 230, - 0x8DC9: 231, - 0x8DD6: 232, - 0x8DD7: 233, - 0x8DE1: 234, - 0x8DE2: 235, - 0x8DF7: 236, - 0x8E41: 237, - 0x8E45: 238, - 0x8E49: 239, - 0x8E51: 240, - 0x8E53: 241, - 0x8E57: 242, - 0x8E61: 243, - 0x8E81: 244, - 0x8E82: 245, - 0x8E85: 246, - 0x8E89: 247, - 0x8E90: 248, - 0x8E91: 249, - 0x8E93: 250, - 0x8E95: 251, - 0x8E97: 252, - 0x8E98: 253, - 0x8EA1: 254, - 0x8EA9: 255, - 0x8EB6: 256, - 0x8EB7: 257, - 0x8EC1: 258, - 0x8EC2: 259, - 0x8EC5: 260, - 0x8EC9: 261, - 0x8ED1: 262, - 0x8ED3: 263, - 0x8ED6: 264, - 0x8EE1: 265, - 0x8EE5: 266, - 0x8EE9: 267, - 0x8EF1: 268, - 0x8EF3: 269, - 0x8F41: 270, - 0x8F61: 271, - 0x8F62: 272, - 0x8F65: 273, - 0x8F67: 274, - 0x8F69: 275, - 0x8F6B: 276, - 0x8F70: 277, - 0x8F71: 278, - 0x8F73: 279, - 0x8F75: 280, - 0x8F77: 281, - 0x8F7B: 282, - 0x8FA1: 283, - 0x8FA2: 284, - 0x8FA5: 285, - 0x8FA9: 286, - 0x8FB1: 287, - 0x8FB3: 288, - 0x8FB5: 289, - 0x8FB7: 290, - 0x9061: 291, - 0x9062: 292, - 0x9063: 293, - 0x9065: 294, - 0x9068: 295, - 0x9069: 296, - 0x906A: 297, - 0x906B: 298, - 0x9071: 299, - 0x9073: 300, - 0x9075: 301, - 0x9076: 302, - 0x9077: 303, - 0x9078: 304, - 0x9079: 305, - 0x907B: 306, - 0x907D: 307, - 0x9081: 308, - 0x9082: 309, - 0x9085: 310, - 0x9089: 311, - 0x9091: 312, - 0x9093: 313, - 0x9095: 314, - 0x9096: 315, - 0x9097: 316, - 0x90A1: 317, - 0x90A2: 318, - 0x90A5: 319, - 0x90A9: 320, - 0x90B1: 321, - 0x90B7: 322, - 0x90E1: 323, - 0x90E2: 324, - 0x90E4: 325, - 0x90E5: 326, - 0x90E9: 327, - 0x90EB: 328, - 0x90EC: 329, - 0x90F1: 330, - 0x90F3: 331, - 0x90F5: 332, - 0x90F6: 333, - 0x90F7: 334, - 0x90FD: 335, - 0x9141: 336, - 0x9142: 337, - 0x9145: 338, - 0x9149: 339, - 0x9151: 340, - 0x9153: 341, - 0x9155: 342, - 0x9156: 343, - 0x9157: 344, - 0x9161: 345, - 0x9162: 346, - 0x9165: 347, - 0x9169: 348, - 0x9171: 349, - 0x9173: 350, - 0x9176: 351, - 0x9177: 352, - 0x917A: 353, - 0x9181: 354, - 0x9185: 355, - 0x91A1: 356, - 0x91A2: 357, - 0x91A5: 358, - 0x91A9: 359, - 0x91AB: 360, - 0x91B1: 361, - 0x91B3: 362, - 0x91B5: 363, - 0x91B7: 364, - 0x91BC: 365, - 0x91BD: 366, - 0x91C1: 367, - 0x91C5: 368, - 0x91C9: 369, - 0x91D6: 370, - 0x9241: 371, - 0x9245: 372, - 0x9249: 373, - 0x9251: 374, - 0x9253: 375, - 0x9255: 376, - 0x9261: 377, - 0x9262: 378, - 0x9265: 379, - 0x9269: 380, - 0x9273: 381, - 0x9275: 382, - 0x9277: 383, - 0x9281: 384, - 0x9282: 385, - 0x9285: 386, - 0x9288: 387, - 0x9289: 388, - 0x9291: 389, - 0x9293: 390, - 0x9295: 391, - 0x9297: 392, - 0x92A1: 393, - 0x92B6: 394, - 0x92C1: 395, - 0x92E1: 396, - 0x92E5: 397, - 0x92E9: 398, - 0x92F1: 399, - 0x92F3: 400, - 0x9341: 401, - 0x9342: 402, - 0x9349: 403, - 0x9351: 404, - 0x9353: 405, - 0x9357: 406, - 0x9361: 407, - 0x9362: 408, - 0x9365: 409, - 0x9369: 410, - 0x936A: 411, - 0x936B: 412, - 0x9371: 413, - 0x9373: 414, - 0x9375: 415, - 0x9377: 416, - 0x9378: 417, - 0x937C: 418, - 0x9381: 419, - 0x9385: 420, - 0x9389: 421, - 0x93A1: 422, - 0x93A2: 423, - 0x93A5: 424, - 0x93A9: 425, - 0x93AB: 426, - 0x93B1: 427, - 0x93B3: 428, - 0x93B5: 429, - 0x93B7: 430, - 0x93BC: 431, - 0x9461: 432, - 0x9462: 433, - 0x9463: 434, - 0x9465: 435, - 0x9468: 436, - 0x9469: 437, - 0x946A: 438, - 0x946B: 439, - 0x946C: 440, - 0x9470: 441, - 0x9471: 442, - 0x9473: 443, - 0x9475: 444, - 0x9476: 445, - 0x9477: 446, - 0x9478: 447, - 0x9479: 448, - 0x947D: 449, - 0x9481: 450, - 0x9482: 451, - 0x9485: 452, - 0x9489: 453, - 0x9491: 454, - 0x9493: 455, - 0x9495: 456, - 0x9496: 457, - 0x9497: 458, - 0x94A1: 459, - 0x94E1: 460, - 0x94E2: 461, - 0x94E3: 462, - 0x94E5: 463, - 0x94E8: 464, - 0x94E9: 465, - 0x94EB: 466, - 0x94EC: 467, - 0x94F1: 468, - 0x94F3: 469, - 0x94F5: 470, - 0x94F7: 471, - 0x94F9: 472, - 0x94FC: 473, - 0x9541: 474, - 0x9542: 475, - 0x9545: 476, - 0x9549: 477, - 0x9551: 478, - 0x9553: 479, - 0x9555: 480, - 0x9556: 481, - 0x9557: 482, - 0x9561: 483, - 0x9565: 484, - 0x9569: 485, - 0x9576: 486, - 0x9577: 487, - 0x9581: 488, - 0x9585: 489, - 0x95A1: 490, - 0x95A2: 491, - 0x95A5: 492, - 0x95A8: 493, - 0x95A9: 494, - 0x95AB: 495, - 0x95AD: 496, - 0x95B1: 497, - 0x95B3: 498, - 0x95B5: 499, - 0x95B7: 500, - 0x95B9: 501, - 0x95BB: 502, - 0x95C1: 503, - 0x95C5: 504, - 0x95C9: 505, - 0x95E1: 506, - 0x95F6: 507, - 0x9641: 508, - 0x9645: 509, - 0x9649: 510, - 0x9651: 511, - 0x9653: 512, - 0x9655: 513, - 0x9661: 514, - 0x9681: 515, - 0x9682: 516, - 0x9685: 517, - 0x9689: 518, - 0x9691: 519, - 0x9693: 520, - 0x9695: 521, - 0x9697: 522, - 0x96A1: 523, - 0x96B6: 524, - 0x96C1: 525, - 0x96D7: 526, - 0x96E1: 527, - 0x96E5: 528, - 0x96E9: 529, - 0x96F3: 530, - 0x96F5: 531, - 0x96F7: 532, - 0x9741: 533, - 0x9745: 534, - 0x9749: 535, - 0x9751: 536, - 0x9757: 537, - 0x9761: 538, - 0x9762: 539, - 0x9765: 540, - 0x9768: 541, - 0x9769: 542, - 0x976B: 543, - 0x9771: 544, - 0x9773: 545, - 0x9775: 546, - 0x9777: 547, - 0x9781: 548, - 0x97A1: 549, - 0x97A2: 550, - 0x97A5: 551, - 0x97A8: 552, - 0x97A9: 553, - 0x97B1: 554, - 0x97B3: 555, - 0x97B5: 556, - 0x97B6: 557, - 0x97B7: 558, - 0x97B8: 559, - 0x9861: 560, - 0x9862: 561, - 0x9865: 562, - 0x9869: 563, - 0x9871: 564, - 0x9873: 565, - 0x9875: 566, - 0x9876: 567, - 0x9877: 568, - 0x987D: 569, - 0x9881: 570, - 0x9882: 571, - 0x9885: 572, - 0x9889: 573, - 0x9891: 574, - 0x9893: 575, - 0x9895: 576, - 0x9896: 577, - 0x9897: 578, - 0x98E1: 579, - 0x98E2: 580, - 0x98E5: 581, - 0x98E9: 582, - 0x98EB: 583, - 0x98EC: 584, - 0x98F1: 585, - 0x98F3: 586, - 0x98F5: 587, - 0x98F6: 588, - 0x98F7: 589, - 0x98FD: 590, - 0x9941: 591, - 0x9942: 592, - 0x9945: 593, - 0x9949: 594, - 0x9951: 595, - 0x9953: 596, - 0x9955: 597, - 0x9956: 598, - 0x9957: 599, - 0x9961: 600, - 0x9976: 601, - 0x99A1: 602, - 0x99A2: 603, - 0x99A5: 604, - 0x99A9: 605, - 0x99B7: 606, - 0x99C1: 607, - 0x99C9: 608, - 0x99E1: 609, - 0x9A41: 610, - 0x9A45: 611, - 0x9A81: 612, - 0x9A82: 613, - 0x9A85: 614, - 0x9A89: 615, - 0x9A90: 616, - 0x9A91: 617, - 0x9A97: 618, - 0x9AC1: 619, - 0x9AE1: 620, - 0x9AE5: 621, - 0x9AE9: 622, - 0x9AF1: 623, - 0x9AF3: 624, - 0x9AF7: 625, - 0x9B61: 626, - 0x9B62: 627, - 0x9B65: 628, - 0x9B68: 629, - 0x9B69: 630, - 0x9B71: 631, - 0x9B73: 632, - 0x9B75: 633, - 0x9B81: 634, - 0x9B85: 635, - 0x9B89: 636, - 0x9B91: 637, - 0x9B93: 638, - 0x9BA1: 639, - 0x9BA5: 640, - 0x9BA9: 641, - 0x9BB1: 642, - 0x9BB3: 643, - 0x9BB5: 644, - 0x9BB7: 645, - 0x9C61: 646, - 0x9C62: 647, - 0x9C65: 648, - 0x9C69: 649, - 0x9C71: 650, - 0x9C73: 651, - 0x9C75: 652, - 0x9C76: 653, - 0x9C77: 654, - 0x9C78: 655, - 0x9C7C: 656, - 0x9C7D: 657, - 0x9C81: 658, - 0x9C82: 659, - 0x9C85: 660, - 0x9C89: 661, - 0x9C91: 662, - 0x9C93: 663, - 0x9C95: 664, - 0x9C96: 665, - 0x9C97: 666, - 0x9CA1: 667, - 0x9CA2: 668, - 0x9CA5: 669, - 0x9CB5: 670, - 0x9CB7: 671, - 0x9CE1: 672, - 0x9CE2: 673, - 0x9CE5: 674, - 0x9CE9: 675, - 0x9CF1: 676, - 0x9CF3: 677, - 0x9CF5: 678, - 0x9CF6: 679, - 0x9CF7: 680, - 0x9CFD: 681, - 0x9D41: 682, - 0x9D42: 683, - 0x9D45: 684, - 0x9D49: 685, - 0x9D51: 686, - 0x9D53: 687, - 0x9D55: 688, - 0x9D57: 689, - 0x9D61: 690, - 0x9D62: 691, - 0x9D65: 692, - 0x9D69: 693, - 0x9D71: 694, - 0x9D73: 695, - 0x9D75: 696, - 0x9D76: 697, - 0x9D77: 698, - 0x9D81: 699, - 0x9D85: 700, - 0x9D93: 701, - 0x9D95: 702, - 0x9DA1: 703, - 0x9DA2: 704, - 0x9DA5: 705, - 0x9DA9: 706, - 0x9DB1: 707, - 0x9DB3: 708, - 0x9DB5: 709, - 0x9DB7: 710, - 0x9DC1: 711, - 0x9DC5: 712, - 0x9DD7: 713, - 0x9DF6: 714, - 0x9E41: 715, - 0x9E45: 716, - 0x9E49: 717, - 0x9E51: 718, - 0x9E53: 719, - 0x9E55: 720, - 0x9E57: 721, - 0x9E61: 722, - 0x9E65: 723, - 0x9E69: 724, - 0x9E73: 725, - 0x9E75: 726, - 0x9E77: 727, - 0x9E81: 728, - 0x9E82: 729, - 0x9E85: 730, - 0x9E89: 731, - 0x9E91: 732, - 0x9E93: 733, - 0x9E95: 734, - 0x9E97: 735, - 0x9EA1: 736, - 0x9EB6: 737, - 0x9EC1: 738, - 0x9EE1: 739, - 0x9EE2: 740, - 0x9EE5: 741, - 0x9EE9: 742, - 0x9EF1: 743, - 0x9EF5: 744, - 0x9EF7: 745, - 0x9F41: 746, - 0x9F42: 747, - 0x9F45: 748, - 0x9F49: 749, - 0x9F51: 750, - 0x9F53: 751, - 0x9F55: 752, - 0x9F57: 753, - 0x9F61: 754, - 0x9F62: 755, - 0x9F65: 756, - 0x9F69: 757, - 0x9F71: 758, - 0x9F73: 759, - 0x9F75: 760, - 0x9F77: 761, - 0x9F78: 762, - 0x9F7B: 763, - 0x9F7C: 764, - 0x9FA1: 765, - 0x9FA2: 766, - 0x9FA5: 767, - 0x9FA9: 768, - 0x9FB1: 769, - 0x9FB3: 770, - 0x9FB5: 771, - 0x9FB7: 772, - 0xA061: 773, - 0xA062: 774, - 0xA065: 775, - 0xA067: 776, - 0xA068: 777, - 0xA069: 778, - 0xA06A: 779, - 0xA06B: 780, - 0xA071: 781, - 0xA073: 782, - 0xA075: 783, - 0xA077: 784, - 0xA078: 785, - 0xA07B: 786, - 0xA07D: 787, - 0xA081: 788, - 0xA082: 789, - 0xA085: 790, - 0xA089: 791, - 0xA091: 792, - 0xA093: 793, - 0xA095: 794, - 0xA096: 795, - 0xA097: 796, - 0xA098: 797, - 0xA0A1: 798, - 0xA0A2: 799, - 0xA0A9: 800, - 0xA0B7: 801, - 0xA0E1: 802, - 0xA0E2: 803, - 0xA0E5: 804, - 0xA0E9: 805, - 0xA0EB: 806, - 0xA0F1: 807, - 0xA0F3: 808, - 0xA0F5: 809, - 0xA0F7: 810, - 0xA0F8: 811, - 0xA0FD: 812, - 0xA141: 813, - 0xA142: 814, - 0xA145: 815, - 0xA149: 816, - 0xA151: 817, - 0xA153: 818, - 0xA155: 819, - 0xA156: 820, - 0xA157: 821, - 0xA161: 822, - 0xA162: 823, - 0xA165: 824, - 0xA169: 825, - 0xA175: 826, - 0xA176: 827, - 0xA177: 828, - 0xA179: 829, - 0xA181: 830, - 0xA1A1: 831, - 0xA1A2: 832, - 0xA1A4: 833, - 0xA1A5: 834, - 0xA1A9: 835, - 0xA1AB: 836, - 0xA1B1: 837, - 0xA1B3: 838, - 0xA1B5: 839, - 0xA1B7: 840, - 0xA1C1: 841, - 0xA1C5: 842, - 0xA1D6: 843, - 0xA1D7: 844, - 0xA241: 845, - 0xA245: 846, - 0xA249: 847, - 0xA253: 848, - 0xA255: 849, - 0xA257: 850, - 0xA261: 851, - 0xA265: 852, - 0xA269: 853, - 0xA273: 854, - 0xA275: 855, - 0xA281: 856, - 0xA282: 857, - 0xA283: 858, - 0xA285: 859, - 0xA288: 860, - 0xA289: 861, - 0xA28A: 862, - 0xA28B: 863, - 0xA291: 864, - 0xA293: 865, - 0xA295: 866, - 0xA297: 867, - 0xA29B: 868, - 0xA29D: 869, - 0xA2A1: 870, - 0xA2A5: 871, - 0xA2A9: 872, - 0xA2B3: 873, - 0xA2B5: 874, - 0xA2C1: 875, - 0xA2E1: 876, - 0xA2E5: 877, - 0xA2E9: 878, - 0xA341: 879, - 0xA345: 880, - 0xA349: 881, - 0xA351: 882, - 0xA355: 883, - 0xA361: 884, - 0xA365: 885, - 0xA369: 886, - 0xA371: 887, - 0xA375: 888, - 0xA3A1: 889, - 0xA3A2: 890, - 0xA3A5: 891, - 0xA3A8: 892, - 0xA3A9: 893, - 0xA3AB: 894, - 0xA3B1: 895, - 0xA3B3: 896, - 0xA3B5: 897, - 0xA3B6: 898, - 0xA3B7: 899, - 0xA3B9: 900, - 0xA3BB: 901, - 0xA461: 902, - 0xA462: 903, - 0xA463: 904, - 0xA464: 905, - 0xA465: 906, - 0xA468: 907, - 0xA469: 908, - 0xA46A: 909, - 0xA46B: 910, - 0xA46C: 911, - 0xA471: 912, - 0xA473: 913, - 0xA475: 914, - 0xA477: 915, - 0xA47B: 916, - 0xA481: 917, - 0xA482: 918, - 0xA485: 919, - 0xA489: 920, - 0xA491: 921, - 0xA493: 922, - 0xA495: 923, - 0xA496: 924, - 0xA497: 925, - 0xA49B: 926, - 0xA4A1: 927, - 0xA4A2: 928, - 0xA4A5: 929, - 0xA4B3: 930, - 0xA4E1: 931, - 0xA4E2: 932, - 0xA4E5: 933, - 0xA4E8: 934, - 0xA4E9: 935, - 0xA4EB: 936, - 0xA4F1: 937, - 0xA4F3: 938, - 0xA4F5: 939, - 0xA4F7: 940, - 0xA4F8: 941, - 0xA541: 942, - 0xA542: 943, - 0xA545: 944, - 0xA548: 945, - 0xA549: 946, - 0xA551: 947, - 0xA553: 948, - 0xA555: 949, - 0xA556: 950, - 0xA557: 951, - 0xA561: 952, - 0xA562: 953, - 0xA565: 954, - 0xA569: 955, - 0xA573: 956, - 0xA575: 957, - 0xA576: 958, - 0xA577: 959, - 0xA57B: 960, - 0xA581: 961, - 0xA585: 962, - 0xA5A1: 963, - 0xA5A2: 964, - 0xA5A3: 965, - 0xA5A5: 966, - 0xA5A9: 967, - 0xA5B1: 968, - 0xA5B3: 969, - 0xA5B5: 970, - 0xA5B7: 971, - 0xA5C1: 972, - 0xA5C5: 973, - 0xA5D6: 974, - 0xA5E1: 975, - 0xA5F6: 976, - 0xA641: 977, - 0xA642: 978, - 0xA645: 979, - 0xA649: 980, - 0xA651: 981, - 0xA653: 982, - 0xA661: 983, - 0xA665: 984, - 0xA681: 985, - 0xA682: 986, - 0xA685: 987, - 0xA688: 988, - 0xA689: 989, - 0xA68A: 990, - 0xA68B: 991, - 0xA691: 992, - 0xA693: 993, - 0xA695: 994, - 0xA697: 995, - 0xA69B: 996, - 0xA69C: 997, - 0xA6A1: 998, - 0xA6A9: 999, - 0xA6B6: 1000, - 0xA6C1: 1001, - 0xA6E1: 1002, - 0xA6E2: 1003, - 0xA6E5: 1004, - 0xA6E9: 1005, - 0xA6F7: 1006, - 0xA741: 1007, - 0xA745: 1008, - 0xA749: 1009, - 0xA751: 1010, - 0xA755: 1011, - 0xA757: 1012, - 0xA761: 1013, - 0xA762: 1014, - 0xA765: 1015, - 0xA769: 1016, - 0xA771: 1017, - 0xA773: 1018, - 0xA775: 1019, - 0xA7A1: 1020, - 0xA7A2: 1021, - 0xA7A5: 1022, - 0xA7A9: 1023, - 0xA7AB: 1024, - 0xA7B1: 1025, - 0xA7B3: 1026, - 0xA7B5: 1027, - 0xA7B7: 1028, - 0xA7B8: 1029, - 0xA7B9: 1030, - 0xA861: 1031, - 0xA862: 1032, - 0xA865: 1033, - 0xA869: 1034, - 0xA86B: 1035, - 0xA871: 1036, - 0xA873: 1037, - 0xA875: 1038, - 0xA876: 1039, - 0xA877: 1040, - 0xA87D: 1041, - 0xA881: 1042, - 0xA882: 1043, - 0xA885: 1044, - 0xA889: 1045, - 0xA891: 1046, - 0xA893: 1047, - 0xA895: 1048, - 0xA896: 1049, - 0xA897: 1050, - 0xA8A1: 1051, - 0xA8A2: 1052, - 0xA8B1: 1053, - 0xA8E1: 1054, - 0xA8E2: 1055, - 0xA8E5: 1056, - 0xA8E8: 1057, - 0xA8E9: 1058, - 0xA8F1: 1059, - 0xA8F5: 1060, - 0xA8F6: 1061, - 0xA8F7: 1062, - 0xA941: 1063, - 0xA957: 1064, - 0xA961: 1065, - 0xA962: 1066, - 0xA971: 1067, - 0xA973: 1068, - 0xA975: 1069, - 0xA976: 1070, - 0xA977: 1071, - 0xA9A1: 1072, - 0xA9A2: 1073, - 0xA9A5: 1074, - 0xA9A9: 1075, - 0xA9B1: 1076, - 0xA9B3: 1077, - 0xA9B7: 1078, - 0xAA41: 1079, - 0xAA61: 1080, - 0xAA77: 1081, - 0xAA81: 1082, - 0xAA82: 1083, - 0xAA85: 1084, - 0xAA89: 1085, - 0xAA91: 1086, - 0xAA95: 1087, - 0xAA97: 1088, - 0xAB41: 1089, - 0xAB57: 1090, - 0xAB61: 1091, - 0xAB65: 1092, - 0xAB69: 1093, - 0xAB71: 1094, - 0xAB73: 1095, - 0xABA1: 1096, - 0xABA2: 1097, - 0xABA5: 1098, - 0xABA9: 1099, - 0xABB1: 1100, - 0xABB3: 1101, - 0xABB5: 1102, - 0xABB7: 1103, - 0xAC61: 1104, - 0xAC62: 1105, - 0xAC64: 1106, - 0xAC65: 1107, - 0xAC68: 1108, - 0xAC69: 1109, - 0xAC6A: 1110, - 0xAC6B: 1111, - 0xAC71: 1112, - 0xAC73: 1113, - 0xAC75: 1114, - 0xAC76: 1115, - 0xAC77: 1116, - 0xAC7B: 1117, - 0xAC81: 1118, - 0xAC82: 1119, - 0xAC85: 1120, - 0xAC89: 1121, - 0xAC91: 1122, - 0xAC93: 1123, - 0xAC95: 1124, - 0xAC96: 1125, - 0xAC97: 1126, - 0xACA1: 1127, - 0xACA2: 1128, - 0xACA5: 1129, - 0xACA9: 1130, - 0xACB1: 1131, - 0xACB3: 1132, - 0xACB5: 1133, - 0xACB7: 1134, - 0xACC1: 1135, - 0xACC5: 1136, - 0xACC9: 1137, - 0xACD1: 1138, - 0xACD7: 1139, - 0xACE1: 1140, - 0xACE2: 1141, - 0xACE3: 1142, - 0xACE4: 1143, - 0xACE5: 1144, - 0xACE8: 1145, - 0xACE9: 1146, - 0xACEB: 1147, - 0xACEC: 1148, - 0xACF1: 1149, - 0xACF3: 1150, - 0xACF5: 1151, - 0xACF6: 1152, - 0xACF7: 1153, - 0xACFC: 1154, - 0xAD41: 1155, - 0xAD42: 1156, - 0xAD45: 1157, - 0xAD49: 1158, - 0xAD51: 1159, - 0xAD53: 1160, - 0xAD55: 1161, - 0xAD56: 1162, - 0xAD57: 1163, - 0xAD61: 1164, - 0xAD62: 1165, - 0xAD65: 1166, - 0xAD69: 1167, - 0xAD71: 1168, - 0xAD73: 1169, - 0xAD75: 1170, - 0xAD76: 1171, - 0xAD77: 1172, - 0xAD81: 1173, - 0xAD85: 1174, - 0xAD89: 1175, - 0xAD97: 1176, - 0xADA1: 1177, - 0xADA2: 1178, - 0xADA3: 1179, - 0xADA5: 1180, - 0xADA9: 1181, - 0xADAB: 1182, - 0xADB1: 1183, - 0xADB3: 1184, - 0xADB5: 1185, - 0xADB7: 1186, - 0xADBB: 1187, - 0xADC1: 1188, - 0xADC2: 1189, - 0xADC5: 1190, - 0xADC9: 1191, - 0xADD7: 1192, - 0xADE1: 1193, - 0xADE5: 1194, - 0xADE9: 1195, - 0xADF1: 1196, - 0xADF5: 1197, - 0xADF6: 1198, - 0xAE41: 1199, - 0xAE45: 1200, - 0xAE49: 1201, - 0xAE51: 1202, - 0xAE53: 1203, - 0xAE55: 1204, - 0xAE61: 1205, - 0xAE62: 1206, - 0xAE65: 1207, - 0xAE69: 1208, - 0xAE71: 1209, - 0xAE73: 1210, - 0xAE75: 1211, - 0xAE77: 1212, - 0xAE81: 1213, - 0xAE82: 1214, - 0xAE85: 1215, - 0xAE88: 1216, - 0xAE89: 1217, - 0xAE91: 1218, - 0xAE93: 1219, - 0xAE95: 1220, - 0xAE97: 1221, - 0xAE99: 1222, - 0xAE9B: 1223, - 0xAE9C: 1224, - 0xAEA1: 1225, - 0xAEB6: 1226, - 0xAEC1: 1227, - 0xAEC2: 1228, - 0xAEC5: 1229, - 0xAEC9: 1230, - 0xAED1: 1231, - 0xAED7: 1232, - 0xAEE1: 1233, - 0xAEE2: 1234, - 0xAEE5: 1235, - 0xAEE9: 1236, - 0xAEF1: 1237, - 0xAEF3: 1238, - 0xAEF5: 1239, - 0xAEF7: 1240, - 0xAF41: 1241, - 0xAF42: 1242, - 0xAF49: 1243, - 0xAF51: 1244, - 0xAF55: 1245, - 0xAF57: 1246, - 0xAF61: 1247, - 0xAF62: 1248, - 0xAF65: 1249, - 0xAF69: 1250, - 0xAF6A: 1251, - 0xAF71: 1252, - 0xAF73: 1253, - 0xAF75: 1254, - 0xAF77: 1255, - 0xAFA1: 1256, - 0xAFA2: 1257, - 0xAFA5: 1258, - 0xAFA8: 1259, - 0xAFA9: 1260, - 0xAFB0: 1261, - 0xAFB1: 1262, - 0xAFB3: 1263, - 0xAFB5: 1264, - 0xAFB7: 1265, - 0xAFBC: 1266, - 0xB061: 1267, - 0xB062: 1268, - 0xB064: 1269, - 0xB065: 1270, - 0xB069: 1271, - 0xB071: 1272, - 0xB073: 1273, - 0xB076: 1274, - 0xB077: 1275, - 0xB07D: 1276, - 0xB081: 1277, - 0xB082: 1278, - 0xB085: 1279, - 0xB089: 1280, - 0xB091: 1281, - 0xB093: 1282, - 0xB096: 1283, - 0xB097: 1284, - 0xB0B7: 1285, - 0xB0E1: 1286, - 0xB0E2: 1287, - 0xB0E5: 1288, - 0xB0E9: 1289, - 0xB0EB: 1290, - 0xB0F1: 1291, - 0xB0F3: 1292, - 0xB0F6: 1293, - 0xB0F7: 1294, - 0xB141: 1295, - 0xB145: 1296, - 0xB149: 1297, - 0xB185: 1298, - 0xB1A1: 1299, - 0xB1A2: 1300, - 0xB1A5: 1301, - 0xB1A8: 1302, - 0xB1A9: 1303, - 0xB1AB: 1304, - 0xB1B1: 1305, - 0xB1B3: 1306, - 0xB1B7: 1307, - 0xB1C1: 1308, - 0xB1C2: 1309, - 0xB1C5: 1310, - 0xB1D6: 1311, - 0xB1E1: 1312, - 0xB1F6: 1313, - 0xB241: 1314, - 0xB245: 1315, - 0xB249: 1316, - 0xB251: 1317, - 0xB253: 1318, - 0xB261: 1319, - 0xB281: 1320, - 0xB282: 1321, - 0xB285: 1322, - 0xB289: 1323, - 0xB291: 1324, - 0xB293: 1325, - 0xB297: 1326, - 0xB2A1: 1327, - 0xB2B6: 1328, - 0xB2C1: 1329, - 0xB2E1: 1330, - 0xB2E5: 1331, - 0xB357: 1332, - 0xB361: 1333, - 0xB362: 1334, - 0xB365: 1335, - 0xB369: 1336, - 0xB36B: 1337, - 0xB370: 1338, - 0xB371: 1339, - 0xB373: 1340, - 0xB381: 1341, - 0xB385: 1342, - 0xB389: 1343, - 0xB391: 1344, - 0xB3A1: 1345, - 0xB3A2: 1346, - 0xB3A5: 1347, - 0xB3A9: 1348, - 0xB3B1: 1349, - 0xB3B3: 1350, - 0xB3B5: 1351, - 0xB3B7: 1352, - 0xB461: 1353, - 0xB462: 1354, - 0xB465: 1355, - 0xB466: 1356, - 0xB467: 1357, - 0xB469: 1358, - 0xB46A: 1359, - 0xB46B: 1360, - 0xB470: 1361, - 0xB471: 1362, - 0xB473: 1363, - 0xB475: 1364, - 0xB476: 1365, - 0xB477: 1366, - 0xB47B: 1367, - 0xB47C: 1368, - 0xB481: 1369, - 0xB482: 1370, - 0xB485: 1371, - 0xB489: 1372, - 0xB491: 1373, - 0xB493: 1374, - 0xB495: 1375, - 0xB496: 1376, - 0xB497: 1377, - 0xB4A1: 1378, - 0xB4A2: 1379, - 0xB4A5: 1380, - 0xB4A9: 1381, - 0xB4AC: 1382, - 0xB4B1: 1383, - 0xB4B3: 1384, - 0xB4B5: 1385, - 0xB4B7: 1386, - 0xB4BB: 1387, - 0xB4BD: 1388, - 0xB4C1: 1389, - 0xB4C5: 1390, - 0xB4C9: 1391, - 0xB4D3: 1392, - 0xB4E1: 1393, - 0xB4E2: 1394, - 0xB4E5: 1395, - 0xB4E6: 1396, - 0xB4E8: 1397, - 0xB4E9: 1398, - 0xB4EA: 1399, - 0xB4EB: 1400, - 0xB4F1: 1401, - 0xB4F3: 1402, - 0xB4F4: 1403, - 0xB4F5: 1404, - 0xB4F6: 1405, - 0xB4F7: 1406, - 0xB4F8: 1407, - 0xB4FA: 1408, - 0xB4FC: 1409, - 0xB541: 1410, - 0xB542: 1411, - 0xB545: 1412, - 0xB549: 1413, - 0xB551: 1414, - 0xB553: 1415, - 0xB555: 1416, - 0xB557: 1417, - 0xB561: 1418, - 0xB562: 1419, - 0xB563: 1420, - 0xB565: 1421, - 0xB569: 1422, - 0xB56B: 1423, - 0xB56C: 1424, - 0xB571: 1425, - 0xB573: 1426, - 0xB574: 1427, - 0xB575: 1428, - 0xB576: 1429, - 0xB577: 1430, - 0xB57B: 1431, - 0xB57C: 1432, - 0xB57D: 1433, - 0xB581: 1434, - 0xB585: 1435, - 0xB589: 1436, - 0xB591: 1437, - 0xB593: 1438, - 0xB595: 1439, - 0xB596: 1440, - 0xB5A1: 1441, - 0xB5A2: 1442, - 0xB5A5: 1443, - 0xB5A9: 1444, - 0xB5AA: 1445, - 0xB5AB: 1446, - 0xB5AD: 1447, - 0xB5B0: 1448, - 0xB5B1: 1449, - 0xB5B3: 1450, - 0xB5B5: 1451, - 0xB5B7: 1452, - 0xB5B9: 1453, - 0xB5C1: 1454, - 0xB5C2: 1455, - 0xB5C5: 1456, - 0xB5C9: 1457, - 0xB5D1: 1458, - 0xB5D3: 1459, - 0xB5D5: 1460, - 0xB5D6: 1461, - 0xB5D7: 1462, - 0xB5E1: 1463, - 0xB5E2: 1464, - 0xB5E5: 1465, - 0xB5F1: 1466, - 0xB5F5: 1467, - 0xB5F7: 1468, - 0xB641: 1469, - 0xB642: 1470, - 0xB645: 1471, - 0xB649: 1472, - 0xB651: 1473, - 0xB653: 1474, - 0xB655: 1475, - 0xB657: 1476, - 0xB661: 1477, - 0xB662: 1478, - 0xB665: 1479, - 0xB669: 1480, - 0xB671: 1481, - 0xB673: 1482, - 0xB675: 1483, - 0xB677: 1484, - 0xB681: 1485, - 0xB682: 1486, - 0xB685: 1487, - 0xB689: 1488, - 0xB68A: 1489, - 0xB68B: 1490, - 0xB691: 1491, - 0xB693: 1492, - 0xB695: 1493, - 0xB697: 1494, - 0xB6A1: 1495, - 0xB6A2: 1496, - 0xB6A5: 1497, - 0xB6A9: 1498, - 0xB6B1: 1499, - 0xB6B3: 1500, - 0xB6B6: 1501, - 0xB6B7: 1502, - 0xB6C1: 1503, - 0xB6C2: 1504, - 0xB6C5: 1505, - 0xB6C9: 1506, - 0xB6D1: 1507, - 0xB6D3: 1508, - 0xB6D7: 1509, - 0xB6E1: 1510, - 0xB6E2: 1511, - 0xB6E5: 1512, - 0xB6E9: 1513, - 0xB6F1: 1514, - 0xB6F3: 1515, - 0xB6F5: 1516, - 0xB6F7: 1517, - 0xB741: 1518, - 0xB742: 1519, - 0xB745: 1520, - 0xB749: 1521, - 0xB751: 1522, - 0xB753: 1523, - 0xB755: 1524, - 0xB757: 1525, - 0xB759: 1526, - 0xB761: 1527, - 0xB762: 1528, - 0xB765: 1529, - 0xB769: 1530, - 0xB76F: 1531, - 0xB771: 1532, - 0xB773: 1533, - 0xB775: 1534, - 0xB777: 1535, - 0xB778: 1536, - 0xB779: 1537, - 0xB77A: 1538, - 0xB77B: 1539, - 0xB77C: 1540, - 0xB77D: 1541, - 0xB781: 1542, - 0xB785: 1543, - 0xB789: 1544, - 0xB791: 1545, - 0xB795: 1546, - 0xB7A1: 1547, - 0xB7A2: 1548, - 0xB7A5: 1549, - 0xB7A9: 1550, - 0xB7AA: 1551, - 0xB7AB: 1552, - 0xB7B0: 1553, - 0xB7B1: 1554, - 0xB7B3: 1555, - 0xB7B5: 1556, - 0xB7B6: 1557, - 0xB7B7: 1558, - 0xB7B8: 1559, - 0xB7BC: 1560, - 0xB861: 1561, - 0xB862: 1562, - 0xB865: 1563, - 0xB867: 1564, - 0xB868: 1565, - 0xB869: 1566, - 0xB86B: 1567, - 0xB871: 1568, - 0xB873: 1569, - 0xB875: 1570, - 0xB876: 1571, - 0xB877: 1572, - 0xB878: 1573, - 0xB881: 1574, - 0xB882: 1575, - 0xB885: 1576, - 0xB889: 1577, - 0xB891: 1578, - 0xB893: 1579, - 0xB895: 1580, - 0xB896: 1581, - 0xB897: 1582, - 0xB8A1: 1583, - 0xB8A2: 1584, - 0xB8A5: 1585, - 0xB8A7: 1586, - 0xB8A9: 1587, - 0xB8B1: 1588, - 0xB8B7: 1589, - 0xB8C1: 1590, - 0xB8C5: 1591, - 0xB8C9: 1592, - 0xB8E1: 1593, - 0xB8E2: 1594, - 0xB8E5: 1595, - 0xB8E9: 1596, - 0xB8EB: 1597, - 0xB8F1: 1598, - 0xB8F3: 1599, - 0xB8F5: 1600, - 0xB8F7: 1601, - 0xB8F8: 1602, - 0xB941: 1603, - 0xB942: 1604, - 0xB945: 1605, - 0xB949: 1606, - 0xB951: 1607, - 0xB953: 1608, - 0xB955: 1609, - 0xB957: 1610, - 0xB961: 1611, - 0xB965: 1612, - 0xB969: 1613, - 0xB971: 1614, - 0xB973: 1615, - 0xB976: 1616, - 0xB977: 1617, - 0xB981: 1618, - 0xB9A1: 1619, - 0xB9A2: 1620, - 0xB9A5: 1621, - 0xB9A9: 1622, - 0xB9AB: 1623, - 0xB9B1: 1624, - 0xB9B3: 1625, - 0xB9B5: 1626, - 0xB9B7: 1627, - 0xB9B8: 1628, - 0xB9B9: 1629, - 0xB9BD: 1630, - 0xB9C1: 1631, - 0xB9C2: 1632, - 0xB9C9: 1633, - 0xB9D3: 1634, - 0xB9D5: 1635, - 0xB9D7: 1636, - 0xB9E1: 1637, - 0xB9F6: 1638, - 0xB9F7: 1639, - 0xBA41: 1640, - 0xBA45: 1641, - 0xBA49: 1642, - 0xBA51: 1643, - 0xBA53: 1644, - 0xBA55: 1645, - 0xBA57: 1646, - 0xBA61: 1647, - 0xBA62: 1648, - 0xBA65: 1649, - 0xBA77: 1650, - 0xBA81: 1651, - 0xBA82: 1652, - 0xBA85: 1653, - 0xBA89: 1654, - 0xBA8A: 1655, - 0xBA8B: 1656, - 0xBA91: 1657, - 0xBA93: 1658, - 0xBA95: 1659, - 0xBA97: 1660, - 0xBAA1: 1661, - 0xBAB6: 1662, - 0xBAC1: 1663, - 0xBAE1: 1664, - 0xBAE2: 1665, - 0xBAE5: 1666, - 0xBAE9: 1667, - 0xBAF1: 1668, - 0xBAF3: 1669, - 0xBAF5: 1670, - 0xBB41: 1671, - 0xBB45: 1672, - 0xBB49: 1673, - 0xBB51: 1674, - 0xBB61: 1675, - 0xBB62: 1676, - 0xBB65: 1677, - 0xBB69: 1678, - 0xBB71: 1679, - 0xBB73: 1680, - 0xBB75: 1681, - 0xBB77: 1682, - 0xBBA1: 1683, - 0xBBA2: 1684, - 0xBBA5: 1685, - 0xBBA8: 1686, - 0xBBA9: 1687, - 0xBBAB: 1688, - 0xBBB1: 1689, - 0xBBB3: 1690, - 0xBBB5: 1691, - 0xBBB7: 1692, - 0xBBB8: 1693, - 0xBBBB: 1694, - 0xBBBC: 1695, - 0xBC61: 1696, - 0xBC62: 1697, - 0xBC65: 1698, - 0xBC67: 1699, - 0xBC69: 1700, - 0xBC6C: 1701, - 0xBC71: 1702, - 0xBC73: 1703, - 0xBC75: 1704, - 0xBC76: 1705, - 0xBC77: 1706, - 0xBC81: 1707, - 0xBC82: 1708, - 0xBC85: 1709, - 0xBC89: 1710, - 0xBC91: 1711, - 0xBC93: 1712, - 0xBC95: 1713, - 0xBC96: 1714, - 0xBC97: 1715, - 0xBCA1: 1716, - 0xBCA5: 1717, - 0xBCB7: 1718, - 0xBCE1: 1719, - 0xBCE2: 1720, - 0xBCE5: 1721, - 0xBCE9: 1722, - 0xBCF1: 1723, - 0xBCF3: 1724, - 0xBCF5: 1725, - 0xBCF6: 1726, - 0xBCF7: 1727, - 0xBD41: 1728, - 0xBD57: 1729, - 0xBD61: 1730, - 0xBD76: 1731, - 0xBDA1: 1732, - 0xBDA2: 1733, - 0xBDA5: 1734, - 0xBDA9: 1735, - 0xBDB1: 1736, - 0xBDB3: 1737, - 0xBDB5: 1738, - 0xBDB7: 1739, - 0xBDB9: 1740, - 0xBDC1: 1741, - 0xBDC2: 1742, - 0xBDC9: 1743, - 0xBDD6: 1744, - 0xBDE1: 1745, - 0xBDF6: 1746, - 0xBE41: 1747, - 0xBE45: 1748, - 0xBE49: 1749, - 0xBE51: 1750, - 0xBE53: 1751, - 0xBE77: 1752, - 0xBE81: 1753, - 0xBE82: 1754, - 0xBE85: 1755, - 0xBE89: 1756, - 0xBE91: 1757, - 0xBE93: 1758, - 0xBE97: 1759, - 0xBEA1: 1760, - 0xBEB6: 1761, - 0xBEB7: 1762, - 0xBEE1: 1763, - 0xBF41: 1764, - 0xBF61: 1765, - 0xBF71: 1766, - 0xBF75: 1767, - 0xBF77: 1768, - 0xBFA1: 1769, - 0xBFA2: 1770, - 0xBFA5: 1771, - 0xBFA9: 1772, - 0xBFB1: 1773, - 0xBFB3: 1774, - 0xBFB7: 1775, - 0xBFB8: 1776, - 0xBFBD: 1777, - 0xC061: 1778, - 0xC062: 1779, - 0xC065: 1780, - 0xC067: 1781, - 0xC069: 1782, - 0xC071: 1783, - 0xC073: 1784, - 0xC075: 1785, - 0xC076: 1786, - 0xC077: 1787, - 0xC078: 1788, - 0xC081: 1789, - 0xC082: 1790, - 0xC085: 1791, - 0xC089: 1792, - 0xC091: 1793, - 0xC093: 1794, - 0xC095: 1795, - 0xC096: 1796, - 0xC097: 1797, - 0xC0A1: 1798, - 0xC0A5: 1799, - 0xC0A7: 1800, - 0xC0A9: 1801, - 0xC0B1: 1802, - 0xC0B7: 1803, - 0xC0E1: 1804, - 0xC0E2: 1805, - 0xC0E5: 1806, - 0xC0E9: 1807, - 0xC0F1: 1808, - 0xC0F3: 1809, - 0xC0F5: 1810, - 0xC0F6: 1811, - 0xC0F7: 1812, - 0xC141: 1813, - 0xC142: 1814, - 0xC145: 1815, - 0xC149: 1816, - 0xC151: 1817, - 0xC153: 1818, - 0xC155: 1819, - 0xC157: 1820, - 0xC161: 1821, - 0xC165: 1822, - 0xC176: 1823, - 0xC181: 1824, - 0xC185: 1825, - 0xC197: 1826, - 0xC1A1: 1827, - 0xC1A2: 1828, - 0xC1A5: 1829, - 0xC1A9: 1830, - 0xC1B1: 1831, - 0xC1B3: 1832, - 0xC1B5: 1833, - 0xC1B7: 1834, - 0xC1C1: 1835, - 0xC1C5: 1836, - 0xC1C9: 1837, - 0xC1D7: 1838, - 0xC241: 1839, - 0xC245: 1840, - 0xC249: 1841, - 0xC251: 1842, - 0xC253: 1843, - 0xC255: 1844, - 0xC257: 1845, - 0xC261: 1846, - 0xC271: 1847, - 0xC281: 1848, - 0xC282: 1849, - 0xC285: 1850, - 0xC289: 1851, - 0xC291: 1852, - 0xC293: 1853, - 0xC295: 1854, - 0xC297: 1855, - 0xC2A1: 1856, - 0xC2B6: 1857, - 0xC2C1: 1858, - 0xC2C5: 1859, - 0xC2E1: 1860, - 0xC2E5: 1861, - 0xC2E9: 1862, - 0xC2F1: 1863, - 0xC2F3: 1864, - 0xC2F5: 1865, - 0xC2F7: 1866, - 0xC341: 1867, - 0xC345: 1868, - 0xC349: 1869, - 0xC351: 1870, - 0xC357: 1871, - 0xC361: 1872, - 0xC362: 1873, - 0xC365: 1874, - 0xC369: 1875, - 0xC371: 1876, - 0xC373: 1877, - 0xC375: 1878, - 0xC377: 1879, - 0xC3A1: 1880, - 0xC3A2: 1881, - 0xC3A5: 1882, - 0xC3A8: 1883, - 0xC3A9: 1884, - 0xC3AA: 1885, - 0xC3B1: 1886, - 0xC3B3: 1887, - 0xC3B5: 1888, - 0xC3B7: 1889, - 0xC461: 1890, - 0xC462: 1891, - 0xC465: 1892, - 0xC469: 1893, - 0xC471: 1894, - 0xC473: 1895, - 0xC475: 1896, - 0xC477: 1897, - 0xC481: 1898, - 0xC482: 1899, - 0xC485: 1900, - 0xC489: 1901, - 0xC491: 1902, - 0xC493: 1903, - 0xC495: 1904, - 0xC496: 1905, - 0xC497: 1906, - 0xC4A1: 1907, - 0xC4A2: 1908, - 0xC4B7: 1909, - 0xC4E1: 1910, - 0xC4E2: 1911, - 0xC4E5: 1912, - 0xC4E8: 1913, - 0xC4E9: 1914, - 0xC4F1: 1915, - 0xC4F3: 1916, - 0xC4F5: 1917, - 0xC4F6: 1918, - 0xC4F7: 1919, - 0xC541: 1920, - 0xC542: 1921, - 0xC545: 1922, - 0xC549: 1923, - 0xC551: 1924, - 0xC553: 1925, - 0xC555: 1926, - 0xC557: 1927, - 0xC561: 1928, - 0xC565: 1929, - 0xC569: 1930, - 0xC571: 1931, - 0xC573: 1932, - 0xC575: 1933, - 0xC576: 1934, - 0xC577: 1935, - 0xC581: 1936, - 0xC5A1: 1937, - 0xC5A2: 1938, - 0xC5A5: 1939, - 0xC5A9: 1940, - 0xC5B1: 1941, - 0xC5B3: 1942, - 0xC5B5: 1943, - 0xC5B7: 1944, - 0xC5C1: 1945, - 0xC5C2: 1946, - 0xC5C5: 1947, - 0xC5C9: 1948, - 0xC5D1: 1949, - 0xC5D7: 1950, - 0xC5E1: 1951, - 0xC5F7: 1952, - 0xC641: 1953, - 0xC649: 1954, - 0xC661: 1955, - 0xC681: 1956, - 0xC682: 1957, - 0xC685: 1958, - 0xC689: 1959, - 0xC691: 1960, - 0xC693: 1961, - 0xC695: 1962, - 0xC697: 1963, - 0xC6A1: 1964, - 0xC6A5: 1965, - 0xC6A9: 1966, - 0xC6B7: 1967, - 0xC6C1: 1968, - 0xC6D7: 1969, - 0xC6E1: 1970, - 0xC6E2: 1971, - 0xC6E5: 1972, - 0xC6E9: 1973, - 0xC6F1: 1974, - 0xC6F3: 1975, - 0xC6F5: 1976, - 0xC6F7: 1977, - 0xC741: 1978, - 0xC745: 1979, - 0xC749: 1980, - 0xC751: 1981, - 0xC761: 1982, - 0xC762: 1983, - 0xC765: 1984, - 0xC769: 1985, - 0xC771: 1986, - 0xC773: 1987, - 0xC777: 1988, - 0xC7A1: 1989, - 0xC7A2: 1990, - 0xC7A5: 1991, - 0xC7A9: 1992, - 0xC7B1: 1993, - 0xC7B3: 1994, - 0xC7B5: 1995, - 0xC7B7: 1996, - 0xC861: 1997, - 0xC862: 1998, - 0xC865: 1999, - 0xC869: 2000, - 0xC86A: 2001, - 0xC871: 2002, - 0xC873: 2003, - 0xC875: 2004, - 0xC876: 2005, - 0xC877: 2006, - 0xC881: 2007, - 0xC882: 2008, - 0xC885: 2009, - 0xC889: 2010, - 0xC891: 2011, - 0xC893: 2012, - 0xC895: 2013, - 0xC896: 2014, - 0xC897: 2015, - 0xC8A1: 2016, - 0xC8B7: 2017, - 0xC8E1: 2018, - 0xC8E2: 2019, - 0xC8E5: 2020, - 0xC8E9: 2021, - 0xC8EB: 2022, - 0xC8F1: 2023, - 0xC8F3: 2024, - 0xC8F5: 2025, - 0xC8F6: 2026, - 0xC8F7: 2027, - 0xC941: 2028, - 0xC942: 2029, - 0xC945: 2030, - 0xC949: 2031, - 0xC951: 2032, - 0xC953: 2033, - 0xC955: 2034, - 0xC957: 2035, - 0xC961: 2036, - 0xC965: 2037, - 0xC976: 2038, - 0xC981: 2039, - 0xC985: 2040, - 0xC9A1: 2041, - 0xC9A2: 2042, - 0xC9A5: 2043, - 0xC9A9: 2044, - 0xC9B1: 2045, - 0xC9B3: 2046, - 0xC9B5: 2047, - 0xC9B7: 2048, - 0xC9BC: 2049, - 0xC9C1: 2050, - 0xC9C5: 2051, - 0xC9E1: 2052, - 0xCA41: 2053, - 0xCA45: 2054, - 0xCA55: 2055, - 0xCA57: 2056, - 0xCA61: 2057, - 0xCA81: 2058, - 0xCA82: 2059, - 0xCA85: 2060, - 0xCA89: 2061, - 0xCA91: 2062, - 0xCA93: 2063, - 0xCA95: 2064, - 0xCA97: 2065, - 0xCAA1: 2066, - 0xCAB6: 2067, - 0xCAC1: 2068, - 0xCAE1: 2069, - 0xCAE2: 2070, - 0xCAE5: 2071, - 0xCAE9: 2072, - 0xCAF1: 2073, - 0xCAF3: 2074, - 0xCAF7: 2075, - 0xCB41: 2076, - 0xCB45: 2077, - 0xCB49: 2078, - 0xCB51: 2079, - 0xCB57: 2080, - 0xCB61: 2081, - 0xCB62: 2082, - 0xCB65: 2083, - 0xCB68: 2084, - 0xCB69: 2085, - 0xCB6B: 2086, - 0xCB71: 2087, - 0xCB73: 2088, - 0xCB75: 2089, - 0xCB81: 2090, - 0xCB85: 2091, - 0xCB89: 2092, - 0xCB91: 2093, - 0xCB93: 2094, - 0xCBA1: 2095, - 0xCBA2: 2096, - 0xCBA5: 2097, - 0xCBA9: 2098, - 0xCBB1: 2099, - 0xCBB3: 2100, - 0xCBB5: 2101, - 0xCBB7: 2102, - 0xCC61: 2103, - 0xCC62: 2104, - 0xCC63: 2105, - 0xCC65: 2106, - 0xCC69: 2107, - 0xCC6B: 2108, - 0xCC71: 2109, - 0xCC73: 2110, - 0xCC75: 2111, - 0xCC76: 2112, - 0xCC77: 2113, - 0xCC7B: 2114, - 0xCC81: 2115, - 0xCC82: 2116, - 0xCC85: 2117, - 0xCC89: 2118, - 0xCC91: 2119, - 0xCC93: 2120, - 0xCC95: 2121, - 0xCC96: 2122, - 0xCC97: 2123, - 0xCCA1: 2124, - 0xCCA2: 2125, - 0xCCE1: 2126, - 0xCCE2: 2127, - 0xCCE5: 2128, - 0xCCE9: 2129, - 0xCCF1: 2130, - 0xCCF3: 2131, - 0xCCF5: 2132, - 0xCCF6: 2133, - 0xCCF7: 2134, - 0xCD41: 2135, - 0xCD42: 2136, - 0xCD45: 2137, - 0xCD49: 2138, - 0xCD51: 2139, - 0xCD53: 2140, - 0xCD55: 2141, - 0xCD57: 2142, - 0xCD61: 2143, - 0xCD65: 2144, - 0xCD69: 2145, - 0xCD71: 2146, - 0xCD73: 2147, - 0xCD76: 2148, - 0xCD77: 2149, - 0xCD81: 2150, - 0xCD89: 2151, - 0xCD93: 2152, - 0xCD95: 2153, - 0xCDA1: 2154, - 0xCDA2: 2155, - 0xCDA5: 2156, - 0xCDA9: 2157, - 0xCDB1: 2158, - 0xCDB3: 2159, - 0xCDB5: 2160, - 0xCDB7: 2161, - 0xCDC1: 2162, - 0xCDD7: 2163, - 0xCE41: 2164, - 0xCE45: 2165, - 0xCE61: 2166, - 0xCE65: 2167, - 0xCE69: 2168, - 0xCE73: 2169, - 0xCE75: 2170, - 0xCE81: 2171, - 0xCE82: 2172, - 0xCE85: 2173, - 0xCE88: 2174, - 0xCE89: 2175, - 0xCE8B: 2176, - 0xCE91: 2177, - 0xCE93: 2178, - 0xCE95: 2179, - 0xCE97: 2180, - 0xCEA1: 2181, - 0xCEB7: 2182, - 0xCEE1: 2183, - 0xCEE5: 2184, - 0xCEE9: 2185, - 0xCEF1: 2186, - 0xCEF5: 2187, - 0xCF41: 2188, - 0xCF45: 2189, - 0xCF49: 2190, - 0xCF51: 2191, - 0xCF55: 2192, - 0xCF57: 2193, - 0xCF61: 2194, - 0xCF65: 2195, - 0xCF69: 2196, - 0xCF71: 2197, - 0xCF73: 2198, - 0xCF75: 2199, - 0xCFA1: 2200, - 0xCFA2: 2201, - 0xCFA5: 2202, - 0xCFA9: 2203, - 0xCFB1: 2204, - 0xCFB3: 2205, - 0xCFB5: 2206, - 0xCFB7: 2207, - 0xD061: 2208, - 0xD062: 2209, - 0xD065: 2210, - 0xD069: 2211, - 0xD06E: 2212, - 0xD071: 2213, - 0xD073: 2214, - 0xD075: 2215, - 0xD077: 2216, - 0xD081: 2217, - 0xD082: 2218, - 0xD085: 2219, - 0xD089: 2220, - 0xD091: 2221, - 0xD093: 2222, - 0xD095: 2223, - 0xD096: 2224, - 0xD097: 2225, - 0xD0A1: 2226, - 0xD0B7: 2227, - 0xD0E1: 2228, - 0xD0E2: 2229, - 0xD0E5: 2230, - 0xD0E9: 2231, - 0xD0EB: 2232, - 0xD0F1: 2233, - 0xD0F3: 2234, - 0xD0F5: 2235, - 0xD0F7: 2236, - 0xD141: 2237, - 0xD142: 2238, - 0xD145: 2239, - 0xD149: 2240, - 0xD151: 2241, - 0xD153: 2242, - 0xD155: 2243, - 0xD157: 2244, - 0xD161: 2245, - 0xD162: 2246, - 0xD165: 2247, - 0xD169: 2248, - 0xD171: 2249, - 0xD173: 2250, - 0xD175: 2251, - 0xD176: 2252, - 0xD177: 2253, - 0xD181: 2254, - 0xD185: 2255, - 0xD189: 2256, - 0xD193: 2257, - 0xD1A1: 2258, - 0xD1A2: 2259, - 0xD1A5: 2260, - 0xD1A9: 2261, - 0xD1AE: 2262, - 0xD1B1: 2263, - 0xD1B3: 2264, - 0xD1B5: 2265, - 0xD1B7: 2266, - 0xD1BB: 2267, - 0xD1C1: 2268, - 0xD1C2: 2269, - 0xD1C5: 2270, - 0xD1C9: 2271, - 0xD1D5: 2272, - 0xD1D7: 2273, - 0xD1E1: 2274, - 0xD1E2: 2275, - 0xD1E5: 2276, - 0xD1F5: 2277, - 0xD1F7: 2278, - 0xD241: 2279, - 0xD242: 2280, - 0xD245: 2281, - 0xD249: 2282, - 0xD253: 2283, - 0xD255: 2284, - 0xD257: 2285, - 0xD261: 2286, - 0xD265: 2287, - 0xD269: 2288, - 0xD273: 2289, - 0xD275: 2290, - 0xD281: 2291, - 0xD282: 2292, - 0xD285: 2293, - 0xD289: 2294, - 0xD28E: 2295, - 0xD291: 2296, - 0xD295: 2297, - 0xD297: 2298, - 0xD2A1: 2299, - 0xD2A5: 2300, - 0xD2A9: 2301, - 0xD2B1: 2302, - 0xD2B7: 2303, - 0xD2C1: 2304, - 0xD2C2: 2305, - 0xD2C5: 2306, - 0xD2C9: 2307, - 0xD2D7: 2308, - 0xD2E1: 2309, - 0xD2E2: 2310, - 0xD2E5: 2311, - 0xD2E9: 2312, - 0xD2F1: 2313, - 0xD2F3: 2314, - 0xD2F5: 2315, - 0xD2F7: 2316, - 0xD341: 2317, - 0xD342: 2318, - 0xD345: 2319, - 0xD349: 2320, - 0xD351: 2321, - 0xD355: 2322, - 0xD357: 2323, - 0xD361: 2324, - 0xD362: 2325, - 0xD365: 2326, - 0xD367: 2327, - 0xD368: 2328, - 0xD369: 2329, - 0xD36A: 2330, - 0xD371: 2331, - 0xD373: 2332, - 0xD375: 2333, - 0xD377: 2334, - 0xD37B: 2335, - 0xD381: 2336, - 0xD385: 2337, - 0xD389: 2338, - 0xD391: 2339, - 0xD393: 2340, - 0xD397: 2341, - 0xD3A1: 2342, - 0xD3A2: 2343, - 0xD3A5: 2344, - 0xD3A9: 2345, - 0xD3B1: 2346, - 0xD3B3: 2347, - 0xD3B5: 2348, - 0xD3B7: 2349, -} diff --git a/venv/lib/python3.10/site-packages/chardet/johabprober.py b/venv/lib/python3.10/site-packages/chardet/johabprober.py deleted file mode 100644 index d7364ba61eca930aa1c868abe3b322cceb995a6b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/johabprober.py +++ /dev/null @@ -1,47 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is mozilla.org code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from .chardistribution import JOHABDistributionAnalysis -from .codingstatemachine import CodingStateMachine -from .mbcharsetprober import MultiByteCharSetProber -from .mbcssm import JOHAB_SM_MODEL - - -class JOHABProber(MultiByteCharSetProber): - def __init__(self) -> None: - super().__init__() - self.coding_sm = CodingStateMachine(JOHAB_SM_MODEL) - self.distribution_analyzer = JOHABDistributionAnalysis() - self.reset() - - @property - def charset_name(self) -> str: - return "Johab" - - @property - def language(self) -> str: - return "Korean" diff --git a/venv/lib/python3.10/site-packages/chardet/jpcntx.py b/venv/lib/python3.10/site-packages/chardet/jpcntx.py deleted file mode 100644 index 2f53bdda09e92da38e31cac1a6d415f4670137f7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/jpcntx.py +++ /dev/null @@ -1,238 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Communicator client code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from typing import List, Tuple, Union - -# This is hiragana 2-char sequence table, the number in each cell represents its frequency category -# fmt: off -jp2_char_context = ( - (0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1), - (2, 4, 0, 4, 0, 3, 0, 4, 0, 3, 4, 4, 4, 2, 4, 3, 3, 4, 3, 2, 3, 3, 4, 2, 3, 3, 3, 2, 4, 1, 4, 3, 3, 1, 5, 4, 3, 4, 3, 4, 3, 5, 3, 0, 3, 5, 4, 2, 0, 3, 1, 0, 3, 3, 0, 3, 3, 0, 1, 1, 0, 4, 3, 0, 3, 3, 0, 4, 0, 2, 0, 3, 5, 5, 5, 5, 4, 0, 4, 1, 0, 3, 4), - (0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2), - (0, 4, 0, 5, 0, 5, 0, 4, 0, 4, 5, 4, 4, 3, 5, 3, 5, 1, 5, 3, 4, 3, 4, 4, 3, 4, 3, 3, 4, 3, 5, 4, 4, 3, 5, 5, 3, 5, 5, 5, 3, 5, 5, 3, 4, 5, 5, 3, 1, 3, 2, 0, 3, 4, 0, 4, 2, 0, 4, 2, 1, 5, 3, 2, 3, 5, 0, 4, 0, 2, 0, 5, 4, 4, 5, 4, 5, 0, 4, 0, 0, 4, 4), - (0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), - (0, 3, 0, 4, 0, 3, 0, 3, 0, 4, 5, 4, 3, 3, 3, 3, 4, 3, 5, 4, 4, 3, 5, 4, 4, 3, 4, 3, 4, 4, 4, 4, 5, 3, 4, 4, 3, 4, 5, 5, 4, 5, 5, 1, 4, 5, 4, 3, 0, 3, 3, 1, 3, 3, 0, 4, 4, 0, 3, 3, 1, 5, 3, 3, 3, 5, 0, 4, 0, 3, 0, 4, 4, 3, 4, 3, 3, 0, 4, 1, 1, 3, 4), - (0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), - (0, 4, 0, 3, 0, 3, 0, 4, 0, 3, 4, 4, 3, 2, 2, 1, 2, 1, 3, 1, 3, 3, 3, 3, 3, 4, 3, 1, 3, 3, 5, 3, 3, 0, 4, 3, 0, 5, 4, 3, 3, 5, 4, 4, 3, 4, 4, 5, 0, 1, 2, 0, 1, 2, 0, 2, 2, 0, 1, 0, 0, 5, 2, 2, 1, 4, 0, 3, 0, 1, 0, 4, 4, 3, 5, 4, 3, 0, 2, 1, 0, 4, 3), - (0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), - (0, 3, 0, 5, 0, 4, 0, 2, 1, 4, 4, 2, 4, 1, 4, 2, 4, 2, 4, 3, 3, 3, 4, 3, 3, 3, 3, 1, 4, 2, 3, 3, 3, 1, 4, 4, 1, 1, 1, 4, 3, 3, 2, 0, 2, 4, 3, 2, 0, 3, 3, 0, 3, 1, 1, 0, 0, 0, 3, 3, 0, 4, 2, 2, 3, 4, 0, 4, 0, 3, 0, 4, 4, 5, 3, 4, 4, 0, 3, 0, 0, 1, 4), - (1, 4, 0, 4, 0, 4, 0, 4, 0, 3, 5, 4, 4, 3, 4, 3, 5, 4, 3, 3, 4, 3, 5, 4, 4, 4, 4, 3, 4, 2, 4, 3, 3, 1, 5, 4, 3, 2, 4, 5, 4, 5, 5, 4, 4, 5, 4, 4, 0, 3, 2, 2, 3, 3, 0, 4, 3, 1, 3, 2, 1, 4, 3, 3, 4, 5, 0, 3, 0, 2, 0, 4, 5, 5, 4, 5, 4, 0, 4, 0, 0, 5, 4), - (0, 5, 0, 5, 0, 4, 0, 3, 0, 4, 4, 3, 4, 3, 3, 3, 4, 0, 4, 4, 4, 3, 4, 3, 4, 3, 3, 1, 4, 2, 4, 3, 4, 0, 5, 4, 1, 4, 5, 4, 4, 5, 3, 2, 4, 3, 4, 3, 2, 4, 1, 3, 3, 3, 2, 3, 2, 0, 4, 3, 3, 4, 3, 3, 3, 4, 0, 4, 0, 3, 0, 4, 5, 4, 4, 4, 3, 0, 4, 1, 0, 1, 3), - (0, 3, 1, 4, 0, 3, 0, 2, 0, 3, 4, 4, 3, 1, 4, 2, 3, 3, 4, 3, 4, 3, 4, 3, 4, 4, 3, 2, 3, 1, 5, 4, 4, 1, 4, 4, 3, 5, 4, 4, 3, 5, 5, 4, 3, 4, 4, 3, 1, 2, 3, 1, 2, 2, 0, 3, 2, 0, 3, 1, 0, 5, 3, 3, 3, 4, 3, 3, 3, 3, 4, 4, 4, 4, 5, 4, 2, 0, 3, 3, 2, 4, 3), - (0, 2, 0, 3, 0, 1, 0, 1, 0, 0, 3, 2, 0, 0, 2, 0, 1, 0, 2, 1, 3, 3, 3, 1, 2, 3, 1, 0, 1, 0, 4, 2, 1, 1, 3, 3, 0, 4, 3, 3, 1, 4, 3, 3, 0, 3, 3, 2, 0, 0, 0, 0, 1, 0, 0, 2, 0, 0, 0, 0, 0, 4, 1, 0, 2, 3, 2, 2, 2, 1, 3, 3, 3, 4, 4, 3, 2, 0, 3, 1, 0, 3, 3), - (0, 4, 0, 4, 0, 3, 0, 3, 0, 4, 4, 4, 3, 3, 3, 3, 3, 3, 4, 3, 4, 2, 4, 3, 4, 3, 3, 2, 4, 3, 4, 5, 4, 1, 4, 5, 3, 5, 4, 5, 3, 5, 4, 0, 3, 5, 5, 3, 1, 3, 3, 2, 2, 3, 0, 3, 4, 1, 3, 3, 2, 4, 3, 3, 3, 4, 0, 4, 0, 3, 0, 4, 5, 4, 4, 5, 3, 0, 4, 1, 0, 3, 4), - (0, 2, 0, 3, 0, 3, 0, 0, 0, 2, 2, 2, 1, 0, 1, 0, 0, 0, 3, 0, 3, 0, 3, 0, 1, 3, 1, 0, 3, 1, 3, 3, 3, 1, 3, 3, 3, 0, 1, 3, 1, 3, 4, 0, 0, 3, 1, 1, 0, 3, 2, 0, 0, 0, 0, 1, 3, 0, 1, 0, 0, 3, 3, 2, 0, 3, 0, 0, 0, 0, 0, 3, 4, 3, 4, 3, 3, 0, 3, 0, 0, 2, 3), - (2, 3, 0, 3, 0, 2, 0, 1, 0, 3, 3, 4, 3, 1, 3, 1, 1, 1, 3, 1, 4, 3, 4, 3, 3, 3, 0, 0, 3, 1, 5, 4, 3, 1, 4, 3, 2, 5, 5, 4, 4, 4, 4, 3, 3, 4, 4, 4, 0, 2, 1, 1, 3, 2, 0, 1, 2, 0, 0, 1, 0, 4, 1, 3, 3, 3, 0, 3, 0, 1, 0, 4, 4, 4, 5, 5, 3, 0, 2, 0, 0, 4, 4), - (0, 2, 0, 1, 0, 3, 1, 3, 0, 2, 3, 3, 3, 0, 3, 1, 0, 0, 3, 0, 3, 2, 3, 1, 3, 2, 1, 1, 0, 0, 4, 2, 1, 0, 2, 3, 1, 4, 3, 2, 0, 4, 4, 3, 1, 3, 1, 3, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 4, 1, 1, 1, 2, 0, 3, 0, 0, 0, 3, 4, 2, 4, 3, 2, 0, 1, 0, 0, 3, 3), - (0, 1, 0, 4, 0, 5, 0, 4, 0, 2, 4, 4, 2, 3, 3, 2, 3, 3, 5, 3, 3, 3, 4, 3, 4, 2, 3, 0, 4, 3, 3, 3, 4, 1, 4, 3, 2, 1, 5, 5, 3, 4, 5, 1, 3, 5, 4, 2, 0, 3, 3, 0, 1, 3, 0, 4, 2, 0, 1, 3, 1, 4, 3, 3, 3, 3, 0, 3, 0, 1, 0, 3, 4, 4, 4, 5, 5, 0, 3, 0, 1, 4, 5), - (0, 2, 0, 3, 0, 3, 0, 0, 0, 2, 3, 1, 3, 0, 4, 0, 1, 1, 3, 0, 3, 4, 3, 2, 3, 1, 0, 3, 3, 2, 3, 1, 3, 0, 2, 3, 0, 2, 1, 4, 1, 2, 2, 0, 0, 3, 3, 0, 0, 2, 0, 0, 0, 1, 0, 0, 0, 0, 2, 2, 0, 3, 2, 1, 3, 3, 0, 2, 0, 2, 0, 0, 3, 3, 1, 2, 4, 0, 3, 0, 2, 2, 3), - (2, 4, 0, 5, 0, 4, 0, 4, 0, 2, 4, 4, 4, 3, 4, 3, 3, 3, 1, 2, 4, 3, 4, 3, 4, 4, 5, 0, 3, 3, 3, 3, 2, 0, 4, 3, 1, 4, 3, 4, 1, 4, 4, 3, 3, 4, 4, 3, 1, 2, 3, 0, 4, 2, 0, 4, 1, 0, 3, 3, 0, 4, 3, 3, 3, 4, 0, 4, 0, 2, 0, 3, 5, 3, 4, 5, 2, 0, 3, 0, 0, 4, 5), - (0, 3, 0, 4, 0, 1, 0, 1, 0, 1, 3, 2, 2, 1, 3, 0, 3, 0, 2, 0, 2, 0, 3, 0, 2, 0, 0, 0, 1, 0, 1, 1, 0, 0, 3, 1, 0, 0, 0, 4, 0, 3, 1, 0, 2, 1, 3, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 4, 2, 2, 3, 1, 0, 3, 0, 0, 0, 1, 4, 4, 4, 3, 0, 0, 4, 0, 0, 1, 4), - (1, 4, 1, 5, 0, 3, 0, 3, 0, 4, 5, 4, 4, 3, 5, 3, 3, 4, 4, 3, 4, 1, 3, 3, 3, 3, 2, 1, 4, 1, 5, 4, 3, 1, 4, 4, 3, 5, 4, 4, 3, 5, 4, 3, 3, 4, 4, 4, 0, 3, 3, 1, 2, 3, 0, 3, 1, 0, 3, 3, 0, 5, 4, 4, 4, 4, 4, 4, 3, 3, 5, 4, 4, 3, 3, 5, 4, 0, 3, 2, 0, 4, 4), - (0, 2, 0, 3, 0, 1, 0, 0, 0, 1, 3, 3, 3, 2, 4, 1, 3, 0, 3, 1, 3, 0, 2, 2, 1, 1, 0, 0, 2, 0, 4, 3, 1, 0, 4, 3, 0, 4, 4, 4, 1, 4, 3, 1, 1, 3, 3, 1, 0, 2, 0, 0, 1, 3, 0, 0, 0, 0, 2, 0, 0, 4, 3, 2, 4, 3, 5, 4, 3, 3, 3, 4, 3, 3, 4, 3, 3, 0, 2, 1, 0, 3, 3), - (0, 2, 0, 4, 0, 3, 0, 2, 0, 2, 5, 5, 3, 4, 4, 4, 4, 1, 4, 3, 3, 0, 4, 3, 4, 3, 1, 3, 3, 2, 4, 3, 0, 3, 4, 3, 0, 3, 4, 4, 2, 4, 4, 0, 4, 5, 3, 3, 2, 2, 1, 1, 1, 2, 0, 1, 5, 0, 3, 3, 2, 4, 3, 3, 3, 4, 0, 3, 0, 2, 0, 4, 4, 3, 5, 5, 0, 0, 3, 0, 2, 3, 3), - (0, 3, 0, 4, 0, 3, 0, 1, 0, 3, 4, 3, 3, 1, 3, 3, 3, 0, 3, 1, 3, 0, 4, 3, 3, 1, 1, 0, 3, 0, 3, 3, 0, 0, 4, 4, 0, 1, 5, 4, 3, 3, 5, 0, 3, 3, 4, 3, 0, 2, 0, 1, 1, 1, 0, 1, 3, 0, 1, 2, 1, 3, 3, 2, 3, 3, 0, 3, 0, 1, 0, 1, 3, 3, 4, 4, 1, 0, 1, 2, 2, 1, 3), - (0, 1, 0, 4, 0, 4, 0, 3, 0, 1, 3, 3, 3, 2, 3, 1, 1, 0, 3, 0, 3, 3, 4, 3, 2, 4, 2, 0, 1, 0, 4, 3, 2, 0, 4, 3, 0, 5, 3, 3, 2, 4, 4, 4, 3, 3, 3, 4, 0, 1, 3, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 4, 2, 3, 3, 3, 0, 3, 0, 0, 0, 4, 4, 4, 5, 3, 2, 0, 3, 3, 0, 3, 5), - (0, 2, 0, 3, 0, 0, 0, 3, 0, 1, 3, 0, 2, 0, 0, 0, 1, 0, 3, 1, 1, 3, 3, 0, 0, 3, 0, 0, 3, 0, 2, 3, 1, 0, 3, 1, 0, 3, 3, 2, 0, 4, 2, 2, 0, 2, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 1, 2, 0, 1, 0, 1, 0, 0, 0, 1, 3, 1, 2, 0, 0, 0, 1, 0, 0, 1, 4), - (0, 3, 0, 3, 0, 5, 0, 1, 0, 2, 4, 3, 1, 3, 3, 2, 1, 1, 5, 2, 1, 0, 5, 1, 2, 0, 0, 0, 3, 3, 2, 2, 3, 2, 4, 3, 0, 0, 3, 3, 1, 3, 3, 0, 2, 5, 3, 4, 0, 3, 3, 0, 1, 2, 0, 2, 2, 0, 3, 2, 0, 2, 2, 3, 3, 3, 0, 2, 0, 1, 0, 3, 4, 4, 2, 5, 4, 0, 3, 0, 0, 3, 5), - (0, 3, 0, 3, 0, 3, 0, 1, 0, 3, 3, 3, 3, 0, 3, 0, 2, 0, 2, 1, 1, 0, 2, 0, 1, 0, 0, 0, 2, 1, 0, 0, 1, 0, 3, 2, 0, 0, 3, 3, 1, 2, 3, 1, 0, 3, 3, 0, 0, 1, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 2, 3, 1, 2, 3, 0, 3, 0, 1, 0, 3, 2, 1, 0, 4, 3, 0, 1, 1, 0, 3, 3), - (0, 4, 0, 5, 0, 3, 0, 3, 0, 4, 5, 5, 4, 3, 5, 3, 4, 3, 5, 3, 3, 2, 5, 3, 4, 4, 4, 3, 4, 3, 4, 5, 5, 3, 4, 4, 3, 4, 4, 5, 4, 4, 4, 3, 4, 5, 5, 4, 2, 3, 4, 2, 3, 4, 0, 3, 3, 1, 4, 3, 2, 4, 3, 3, 5, 5, 0, 3, 0, 3, 0, 5, 5, 5, 5, 4, 4, 0, 4, 0, 1, 4, 4), - (0, 4, 0, 4, 0, 3, 0, 3, 0, 3, 5, 4, 4, 2, 3, 2, 5, 1, 3, 2, 5, 1, 4, 2, 3, 2, 3, 3, 4, 3, 3, 3, 3, 2, 5, 4, 1, 3, 3, 5, 3, 4, 4, 0, 4, 4, 3, 1, 1, 3, 1, 0, 2, 3, 0, 2, 3, 0, 3, 0, 0, 4, 3, 1, 3, 4, 0, 3, 0, 2, 0, 4, 4, 4, 3, 4, 5, 0, 4, 0, 0, 3, 4), - (0, 3, 0, 3, 0, 3, 1, 2, 0, 3, 4, 4, 3, 3, 3, 0, 2, 2, 4, 3, 3, 1, 3, 3, 3, 1, 1, 0, 3, 1, 4, 3, 2, 3, 4, 4, 2, 4, 4, 4, 3, 4, 4, 3, 2, 4, 4, 3, 1, 3, 3, 1, 3, 3, 0, 4, 1, 0, 2, 2, 1, 4, 3, 2, 3, 3, 5, 4, 3, 3, 5, 4, 4, 3, 3, 0, 4, 0, 3, 2, 2, 4, 4), - (0, 2, 0, 1, 0, 0, 0, 0, 0, 1, 2, 1, 3, 0, 0, 0, 0, 0, 2, 0, 1, 2, 1, 0, 0, 1, 0, 0, 0, 0, 3, 0, 0, 1, 0, 1, 1, 3, 1, 0, 0, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 2, 2, 0, 3, 4, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 1), - (0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 4, 0, 4, 1, 4, 0, 3, 0, 4, 0, 3, 0, 4, 0, 3, 0, 3, 0, 4, 1, 5, 1, 4, 0, 0, 3, 0, 5, 0, 5, 2, 0, 1, 0, 0, 0, 2, 1, 4, 0, 1, 3, 0, 0, 3, 0, 0, 3, 1, 1, 4, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0), - (1, 4, 0, 5, 0, 3, 0, 2, 0, 3, 5, 4, 4, 3, 4, 3, 5, 3, 4, 3, 3, 0, 4, 3, 3, 3, 3, 3, 3, 2, 4, 4, 3, 1, 3, 4, 4, 5, 4, 4, 3, 4, 4, 1, 3, 5, 4, 3, 3, 3, 1, 2, 2, 3, 3, 1, 3, 1, 3, 3, 3, 5, 3, 3, 4, 5, 0, 3, 0, 3, 0, 3, 4, 3, 4, 4, 3, 0, 3, 0, 2, 4, 3), - (0, 1, 0, 4, 0, 0, 0, 0, 0, 1, 4, 0, 4, 1, 4, 2, 4, 0, 3, 0, 1, 0, 1, 0, 0, 0, 0, 0, 2, 0, 3, 1, 1, 1, 0, 3, 0, 0, 0, 1, 2, 1, 0, 0, 1, 1, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 3, 0, 0, 0, 0, 3, 2, 0, 2, 2, 0, 1, 0, 0, 0, 2, 3, 2, 3, 3, 0, 0, 0, 0, 2, 1, 0), - (0, 5, 1, 5, 0, 3, 0, 3, 0, 5, 4, 4, 5, 1, 5, 3, 3, 0, 4, 3, 4, 3, 5, 3, 4, 3, 3, 2, 4, 3, 4, 3, 3, 0, 3, 3, 1, 4, 4, 3, 4, 4, 4, 3, 4, 5, 5, 3, 2, 3, 1, 1, 3, 3, 1, 3, 1, 1, 3, 3, 2, 4, 5, 3, 3, 5, 0, 4, 0, 3, 0, 4, 4, 3, 5, 3, 3, 0, 3, 4, 0, 4, 3), - (0, 5, 0, 5, 0, 3, 0, 2, 0, 4, 4, 3, 5, 2, 4, 3, 3, 3, 4, 4, 4, 3, 5, 3, 5, 3, 3, 1, 4, 0, 4, 3, 3, 0, 3, 3, 0, 4, 4, 4, 4, 5, 4, 3, 3, 5, 5, 3, 2, 3, 1, 2, 3, 2, 0, 1, 0, 0, 3, 2, 2, 4, 4, 3, 1, 5, 0, 4, 0, 3, 0, 4, 3, 1, 3, 2, 1, 0, 3, 3, 0, 3, 3), - (0, 4, 0, 5, 0, 5, 0, 4, 0, 4, 5, 5, 5, 3, 4, 3, 3, 2, 5, 4, 4, 3, 5, 3, 5, 3, 4, 0, 4, 3, 4, 4, 3, 2, 4, 4, 3, 4, 5, 4, 4, 5, 5, 0, 3, 5, 5, 4, 1, 3, 3, 2, 3, 3, 1, 3, 1, 0, 4, 3, 1, 4, 4, 3, 4, 5, 0, 4, 0, 2, 0, 4, 3, 4, 4, 3, 3, 0, 4, 0, 0, 5, 5), - (0, 4, 0, 4, 0, 5, 0, 1, 1, 3, 3, 4, 4, 3, 4, 1, 3, 0, 5, 1, 3, 0, 3, 1, 3, 1, 1, 0, 3, 0, 3, 3, 4, 0, 4, 3, 0, 4, 4, 4, 3, 4, 4, 0, 3, 5, 4, 1, 0, 3, 0, 0, 2, 3, 0, 3, 1, 0, 3, 1, 0, 3, 2, 1, 3, 5, 0, 3, 0, 1, 0, 3, 2, 3, 3, 4, 4, 0, 2, 2, 0, 4, 4), - (2, 4, 0, 5, 0, 4, 0, 3, 0, 4, 5, 5, 4, 3, 5, 3, 5, 3, 5, 3, 5, 2, 5, 3, 4, 3, 3, 4, 3, 4, 5, 3, 2, 1, 5, 4, 3, 2, 3, 4, 5, 3, 4, 1, 2, 5, 4, 3, 0, 3, 3, 0, 3, 2, 0, 2, 3, 0, 4, 1, 0, 3, 4, 3, 3, 5, 0, 3, 0, 1, 0, 4, 5, 5, 5, 4, 3, 0, 4, 2, 0, 3, 5), - (0, 5, 0, 4, 0, 4, 0, 2, 0, 5, 4, 3, 4, 3, 4, 3, 3, 3, 4, 3, 4, 2, 5, 3, 5, 3, 4, 1, 4, 3, 4, 4, 4, 0, 3, 5, 0, 4, 4, 4, 4, 5, 3, 1, 3, 4, 5, 3, 3, 3, 3, 3, 3, 3, 0, 2, 2, 0, 3, 3, 2, 4, 3, 3, 3, 5, 3, 4, 1, 3, 3, 5, 3, 2, 0, 0, 0, 0, 4, 3, 1, 3, 3), - (0, 1, 0, 3, 0, 3, 0, 1, 0, 1, 3, 3, 3, 2, 3, 3, 3, 0, 3, 0, 0, 0, 3, 1, 3, 0, 0, 0, 2, 2, 2, 3, 0, 0, 3, 2, 0, 1, 2, 4, 1, 3, 3, 0, 0, 3, 3, 3, 0, 1, 0, 0, 2, 1, 0, 0, 3, 0, 3, 1, 0, 3, 0, 0, 1, 3, 0, 2, 0, 1, 0, 3, 3, 1, 3, 3, 0, 0, 1, 1, 0, 3, 3), - (0, 2, 0, 3, 0, 2, 1, 4, 0, 2, 2, 3, 1, 1, 3, 1, 1, 0, 2, 0, 3, 1, 2, 3, 1, 3, 0, 0, 1, 0, 4, 3, 2, 3, 3, 3, 1, 4, 2, 3, 3, 3, 3, 1, 0, 3, 1, 4, 0, 1, 1, 0, 1, 2, 0, 1, 1, 0, 1, 1, 0, 3, 1, 3, 2, 2, 0, 1, 0, 0, 0, 2, 3, 3, 3, 1, 0, 0, 0, 0, 0, 2, 3), - (0, 5, 0, 4, 0, 5, 0, 2, 0, 4, 5, 5, 3, 3, 4, 3, 3, 1, 5, 4, 4, 2, 4, 4, 4, 3, 4, 2, 4, 3, 5, 5, 4, 3, 3, 4, 3, 3, 5, 5, 4, 5, 5, 1, 3, 4, 5, 3, 1, 4, 3, 1, 3, 3, 0, 3, 3, 1, 4, 3, 1, 4, 5, 3, 3, 5, 0, 4, 0, 3, 0, 5, 3, 3, 1, 4, 3, 0, 4, 0, 1, 5, 3), - (0, 5, 0, 5, 0, 4, 0, 2, 0, 4, 4, 3, 4, 3, 3, 3, 3, 3, 5, 4, 4, 4, 4, 4, 4, 5, 3, 3, 5, 2, 4, 4, 4, 3, 4, 4, 3, 3, 4, 4, 5, 5, 3, 3, 4, 3, 4, 3, 3, 4, 3, 3, 3, 3, 1, 2, 2, 1, 4, 3, 3, 5, 4, 4, 3, 4, 0, 4, 0, 3, 0, 4, 4, 4, 4, 4, 1, 0, 4, 2, 0, 2, 4), - (0, 4, 0, 4, 0, 3, 0, 1, 0, 3, 5, 2, 3, 0, 3, 0, 2, 1, 4, 2, 3, 3, 4, 1, 4, 3, 3, 2, 4, 1, 3, 3, 3, 0, 3, 3, 0, 0, 3, 3, 3, 5, 3, 3, 3, 3, 3, 2, 0, 2, 0, 0, 2, 0, 0, 2, 0, 0, 1, 0, 0, 3, 1, 2, 2, 3, 0, 3, 0, 2, 0, 4, 4, 3, 3, 4, 1, 0, 3, 0, 0, 2, 4), - (0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 2, 0, 0, 0, 0, 0, 1, 0, 2, 0, 1, 0, 0, 0, 0, 0, 3, 1, 3, 0, 3, 2, 0, 0, 0, 1, 0, 3, 2, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 4, 0, 2, 0, 0, 0, 0, 0, 0, 2), - (0, 2, 1, 3, 0, 2, 0, 2, 0, 3, 3, 3, 3, 1, 3, 1, 3, 3, 3, 3, 3, 3, 4, 2, 2, 1, 2, 1, 4, 0, 4, 3, 1, 3, 3, 3, 2, 4, 3, 5, 4, 3, 3, 3, 3, 3, 3, 3, 0, 1, 3, 0, 2, 0, 0, 1, 0, 0, 1, 0, 0, 4, 2, 0, 2, 3, 0, 3, 3, 0, 3, 3, 4, 2, 3, 1, 4, 0, 1, 2, 0, 2, 3), - (0, 3, 0, 3, 0, 1, 0, 3, 0, 2, 3, 3, 3, 0, 3, 1, 2, 0, 3, 3, 2, 3, 3, 2, 3, 2, 3, 1, 3, 0, 4, 3, 2, 0, 3, 3, 1, 4, 3, 3, 2, 3, 4, 3, 1, 3, 3, 1, 1, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 4, 1, 1, 0, 3, 0, 3, 1, 0, 2, 3, 3, 3, 3, 3, 1, 0, 0, 2, 0, 3, 3), - (0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 2, 0, 3, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 3, 0, 3, 0, 3, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 2, 0, 2, 3, 0, 0, 0, 0, 0, 0, 0, 0, 3), - (0, 2, 0, 3, 1, 3, 0, 3, 0, 2, 3, 3, 3, 1, 3, 1, 3, 1, 3, 1, 3, 3, 3, 1, 3, 0, 2, 3, 1, 1, 4, 3, 3, 2, 3, 3, 1, 2, 2, 4, 1, 3, 3, 0, 1, 4, 2, 3, 0, 1, 3, 0, 3, 0, 0, 1, 3, 0, 2, 0, 0, 3, 3, 2, 1, 3, 0, 3, 0, 2, 0, 3, 4, 4, 4, 3, 1, 0, 3, 0, 0, 3, 3), - (0, 2, 0, 1, 0, 2, 0, 0, 0, 1, 3, 2, 2, 1, 3, 0, 1, 1, 3, 0, 3, 2, 3, 1, 2, 0, 2, 0, 1, 1, 3, 3, 3, 0, 3, 3, 1, 1, 2, 3, 2, 3, 3, 1, 2, 3, 2, 0, 0, 1, 0, 0, 0, 0, 0, 0, 3, 0, 1, 0, 0, 2, 1, 2, 1, 3, 0, 3, 0, 0, 0, 3, 4, 4, 4, 3, 2, 0, 2, 0, 0, 2, 4), - (0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 3, 1, 0, 0, 0, 0, 0, 0, 0, 3), - (0, 3, 0, 3, 0, 2, 0, 3, 0, 3, 3, 3, 2, 3, 2, 2, 2, 0, 3, 1, 3, 3, 3, 2, 3, 3, 0, 0, 3, 0, 3, 2, 2, 0, 2, 3, 1, 4, 3, 4, 3, 3, 2, 3, 1, 5, 4, 4, 0, 3, 1, 2, 1, 3, 0, 3, 1, 1, 2, 0, 2, 3, 1, 3, 1, 3, 0, 3, 0, 1, 0, 3, 3, 4, 4, 2, 1, 0, 2, 1, 0, 2, 4), - (0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 4, 2, 5, 1, 4, 0, 2, 0, 2, 1, 3, 1, 4, 0, 2, 1, 0, 0, 2, 1, 4, 1, 1, 0, 3, 3, 0, 5, 1, 3, 2, 3, 3, 1, 0, 3, 2, 3, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 4, 0, 1, 0, 3, 0, 2, 0, 1, 0, 3, 3, 3, 4, 3, 3, 0, 0, 0, 0, 2, 3), - (0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 2, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 1, 0, 0, 1, 0, 0, 0, 0, 0, 3), - (0, 1, 0, 3, 0, 4, 0, 3, 0, 2, 4, 3, 1, 0, 3, 2, 2, 1, 3, 1, 2, 2, 3, 1, 1, 1, 2, 1, 3, 0, 1, 2, 0, 1, 3, 2, 1, 3, 0, 5, 5, 1, 0, 0, 1, 3, 2, 1, 0, 3, 0, 0, 1, 0, 0, 0, 0, 0, 3, 4, 0, 1, 1, 1, 3, 2, 0, 2, 0, 1, 0, 2, 3, 3, 1, 2, 3, 0, 1, 0, 1, 0, 4), - (0, 0, 0, 1, 0, 3, 0, 3, 0, 2, 2, 1, 0, 0, 4, 0, 3, 0, 3, 1, 3, 0, 3, 0, 3, 0, 1, 0, 3, 0, 3, 1, 3, 0, 3, 3, 0, 0, 1, 2, 1, 1, 1, 0, 1, 2, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 1, 2, 0, 0, 2, 0, 0, 0, 0, 2, 3, 3, 3, 3, 0, 0, 0, 0, 1, 4), - (0, 0, 0, 3, 0, 3, 0, 0, 0, 0, 3, 1, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 3, 0, 2, 0, 2, 3, 0, 0, 2, 2, 3, 1, 2, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 2, 0, 0, 0, 0, 2, 3), - (2, 4, 0, 5, 0, 5, 0, 4, 0, 3, 4, 3, 3, 3, 4, 3, 3, 3, 4, 3, 4, 4, 5, 4, 5, 5, 5, 2, 3, 0, 5, 5, 4, 1, 5, 4, 3, 1, 5, 4, 3, 4, 4, 3, 3, 4, 3, 3, 0, 3, 2, 0, 2, 3, 0, 3, 0, 0, 3, 3, 0, 5, 3, 2, 3, 3, 0, 3, 0, 3, 0, 3, 4, 5, 4, 5, 3, 0, 4, 3, 0, 3, 4), - (0, 3, 0, 3, 0, 3, 0, 3, 0, 3, 3, 4, 3, 2, 3, 2, 3, 0, 4, 3, 3, 3, 3, 3, 3, 3, 3, 0, 3, 2, 4, 3, 3, 1, 3, 4, 3, 4, 4, 4, 3, 4, 4, 3, 2, 4, 4, 1, 0, 2, 0, 0, 1, 1, 0, 2, 0, 0, 3, 1, 0, 5, 3, 2, 1, 3, 0, 3, 0, 1, 2, 4, 3, 2, 4, 3, 3, 0, 3, 2, 0, 4, 4), - (0, 3, 0, 3, 0, 1, 0, 0, 0, 1, 4, 3, 3, 2, 3, 1, 3, 1, 4, 2, 3, 2, 4, 2, 3, 4, 3, 0, 2, 2, 3, 3, 3, 0, 3, 3, 3, 0, 3, 4, 1, 3, 3, 0, 3, 4, 3, 3, 0, 1, 1, 0, 1, 0, 0, 0, 4, 0, 3, 0, 0, 3, 1, 2, 1, 3, 0, 4, 0, 1, 0, 4, 3, 3, 4, 3, 3, 0, 2, 0, 0, 3, 3), - (0, 3, 0, 4, 0, 1, 0, 3, 0, 3, 4, 3, 3, 0, 3, 3, 3, 1, 3, 1, 3, 3, 4, 3, 3, 3, 0, 0, 3, 1, 5, 3, 3, 1, 3, 3, 2, 5, 4, 3, 3, 4, 5, 3, 2, 5, 3, 4, 0, 1, 0, 0, 0, 0, 0, 2, 0, 0, 1, 1, 0, 4, 2, 2, 1, 3, 0, 3, 0, 2, 0, 4, 4, 3, 5, 3, 2, 0, 1, 1, 0, 3, 4), - (0, 5, 0, 4, 0, 5, 0, 2, 0, 4, 4, 3, 3, 2, 3, 3, 3, 1, 4, 3, 4, 1, 5, 3, 4, 3, 4, 0, 4, 2, 4, 3, 4, 1, 5, 4, 0, 4, 4, 4, 4, 5, 4, 1, 3, 5, 4, 2, 1, 4, 1, 1, 3, 2, 0, 3, 1, 0, 3, 2, 1, 4, 3, 3, 3, 4, 0, 4, 0, 3, 0, 4, 4, 4, 3, 3, 3, 0, 4, 2, 0, 3, 4), - (1, 4, 0, 4, 0, 3, 0, 1, 0, 3, 3, 3, 1, 1, 3, 3, 2, 2, 3, 3, 1, 0, 3, 2, 2, 1, 2, 0, 3, 1, 2, 1, 2, 0, 3, 2, 0, 2, 2, 3, 3, 4, 3, 0, 3, 3, 1, 2, 0, 1, 1, 3, 1, 2, 0, 0, 3, 0, 1, 1, 0, 3, 2, 2, 3, 3, 0, 3, 0, 0, 0, 2, 3, 3, 4, 3, 3, 0, 1, 0, 0, 1, 4), - (0, 4, 0, 4, 0, 4, 0, 0, 0, 3, 4, 4, 3, 1, 4, 2, 3, 2, 3, 3, 3, 1, 4, 3, 4, 0, 3, 0, 4, 2, 3, 3, 2, 2, 5, 4, 2, 1, 3, 4, 3, 4, 3, 1, 3, 3, 4, 2, 0, 2, 1, 0, 3, 3, 0, 0, 2, 0, 3, 1, 0, 4, 4, 3, 4, 3, 0, 4, 0, 1, 0, 2, 4, 4, 4, 4, 4, 0, 3, 2, 0, 3, 3), - (0, 0, 0, 1, 0, 4, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 3, 2, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 2), - (0, 2, 0, 3, 0, 4, 0, 4, 0, 1, 3, 3, 3, 0, 4, 0, 2, 1, 2, 1, 1, 1, 2, 0, 3, 1, 1, 0, 1, 0, 3, 1, 0, 0, 3, 3, 2, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 2, 0, 2, 2, 0, 3, 1, 0, 0, 1, 0, 1, 1, 0, 1, 2, 0, 3, 0, 0, 0, 0, 1, 0, 0, 3, 3, 4, 3, 1, 0, 1, 0, 3, 0, 2), - (0, 0, 0, 3, 0, 5, 0, 0, 0, 0, 1, 0, 2, 0, 3, 1, 0, 1, 3, 0, 0, 0, 2, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0, 0, 4, 0, 0, 0, 2, 3, 0, 1, 4, 1, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 3, 0, 0, 0, 0, 0, 3), - (0, 2, 0, 5, 0, 5, 0, 1, 0, 2, 4, 3, 3, 2, 5, 1, 3, 2, 3, 3, 3, 0, 4, 1, 2, 0, 3, 0, 4, 0, 2, 2, 1, 1, 5, 3, 0, 0, 1, 4, 2, 3, 2, 0, 3, 3, 3, 2, 0, 2, 4, 1, 1, 2, 0, 1, 1, 0, 3, 1, 0, 1, 3, 1, 2, 3, 0, 2, 0, 0, 0, 1, 3, 5, 4, 4, 4, 0, 3, 0, 0, 1, 3), - (0, 4, 0, 5, 0, 4, 0, 4, 0, 4, 5, 4, 3, 3, 4, 3, 3, 3, 4, 3, 4, 4, 5, 3, 4, 5, 4, 2, 4, 2, 3, 4, 3, 1, 4, 4, 1, 3, 5, 4, 4, 5, 5, 4, 4, 5, 5, 5, 2, 3, 3, 1, 4, 3, 1, 3, 3, 0, 3, 3, 1, 4, 3, 4, 4, 4, 0, 3, 0, 4, 0, 3, 3, 4, 4, 5, 0, 0, 4, 3, 0, 4, 5), - (0, 4, 0, 4, 0, 3, 0, 3, 0, 3, 4, 4, 4, 3, 3, 2, 4, 3, 4, 3, 4, 3, 5, 3, 4, 3, 2, 1, 4, 2, 4, 4, 3, 1, 3, 4, 2, 4, 5, 5, 3, 4, 5, 4, 1, 5, 4, 3, 0, 3, 2, 2, 3, 2, 1, 3, 1, 0, 3, 3, 3, 5, 3, 3, 3, 5, 4, 4, 2, 3, 3, 4, 3, 3, 3, 2, 1, 0, 3, 2, 1, 4, 3), - (0, 4, 0, 5, 0, 4, 0, 3, 0, 3, 5, 5, 3, 2, 4, 3, 4, 0, 5, 4, 4, 1, 4, 4, 4, 3, 3, 3, 4, 3, 5, 5, 2, 3, 3, 4, 1, 2, 5, 5, 3, 5, 5, 2, 3, 5, 5, 4, 0, 3, 2, 0, 3, 3, 1, 1, 5, 1, 4, 1, 0, 4, 3, 2, 3, 5, 0, 4, 0, 3, 0, 5, 4, 3, 4, 3, 0, 0, 4, 1, 0, 4, 4), - (1, 3, 0, 4, 0, 2, 0, 2, 0, 2, 5, 5, 3, 3, 3, 3, 3, 0, 4, 2, 3, 4, 4, 4, 3, 4, 0, 0, 3, 4, 5, 4, 3, 3, 3, 3, 2, 5, 5, 4, 5, 5, 5, 4, 3, 5, 5, 5, 1, 3, 1, 0, 1, 0, 0, 3, 2, 0, 4, 2, 0, 5, 2, 3, 2, 4, 1, 3, 0, 3, 0, 4, 5, 4, 5, 4, 3, 0, 4, 2, 0, 5, 4), - (0, 3, 0, 4, 0, 5, 0, 3, 0, 3, 4, 4, 3, 2, 3, 2, 3, 3, 3, 3, 3, 2, 4, 3, 3, 2, 2, 0, 3, 3, 3, 3, 3, 1, 3, 3, 3, 0, 4, 4, 3, 4, 4, 1, 1, 4, 4, 2, 0, 3, 1, 0, 1, 1, 0, 4, 1, 0, 2, 3, 1, 3, 3, 1, 3, 4, 0, 3, 0, 1, 0, 3, 1, 3, 0, 0, 1, 0, 2, 0, 0, 4, 4), - (0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), - (0, 3, 0, 3, 0, 2, 0, 3, 0, 1, 5, 4, 3, 3, 3, 1, 4, 2, 1, 2, 3, 4, 4, 2, 4, 4, 5, 0, 3, 1, 4, 3, 4, 0, 4, 3, 3, 3, 2, 3, 2, 5, 3, 4, 3, 2, 2, 3, 0, 0, 3, 0, 2, 1, 0, 1, 2, 0, 0, 0, 0, 2, 1, 1, 3, 1, 0, 2, 0, 4, 0, 3, 4, 4, 4, 5, 2, 0, 2, 0, 0, 1, 3), - (0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 0, 0, 0, 4, 2, 1, 1, 0, 1, 0, 3, 2, 0, 0, 3, 1, 1, 1, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 1, 0, 0, 0, 2, 0, 0, 0, 1, 4, 0, 4, 2, 1, 0, 0, 0, 0, 0, 1), - (0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 3, 1, 0, 0, 0, 2, 0, 2, 1, 0, 0, 1, 2, 1, 0, 1, 1, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 3, 1, 0, 0, 0, 0, 0, 1, 0, 0, 2, 1, 0, 0, 0, 0, 0, 0, 0, 0, 2), - (0, 4, 0, 4, 0, 4, 0, 3, 0, 4, 4, 3, 4, 2, 4, 3, 2, 0, 4, 4, 4, 3, 5, 3, 5, 3, 3, 2, 4, 2, 4, 3, 4, 3, 1, 4, 0, 2, 3, 4, 4, 4, 3, 3, 3, 4, 4, 4, 3, 4, 1, 3, 4, 3, 2, 1, 2, 1, 3, 3, 3, 4, 4, 3, 3, 5, 0, 4, 0, 3, 0, 4, 3, 3, 3, 2, 1, 0, 3, 0, 0, 3, 3), - (0, 4, 0, 3, 0, 3, 0, 3, 0, 3, 5, 5, 3, 3, 3, 3, 4, 3, 4, 3, 3, 3, 4, 4, 4, 3, 3, 3, 3, 4, 3, 5, 3, 3, 1, 3, 2, 4, 5, 5, 5, 5, 4, 3, 4, 5, 5, 3, 2, 2, 3, 3, 3, 3, 2, 3, 3, 1, 2, 3, 2, 4, 3, 3, 3, 4, 0, 4, 0, 2, 0, 4, 3, 2, 2, 1, 2, 0, 3, 0, 0, 4, 1), -) -# fmt: on - - -class JapaneseContextAnalysis: - NUM_OF_CATEGORY = 6 - DONT_KNOW = -1 - ENOUGH_REL_THRESHOLD = 100 - MAX_REL_THRESHOLD = 1000 - MINIMUM_DATA_THRESHOLD = 4 - - def __init__(self) -> None: - self._total_rel = 0 - self._rel_sample: List[int] = [] - self._need_to_skip_char_num = 0 - self._last_char_order = -1 - self._done = False - self.reset() - - def reset(self) -> None: - self._total_rel = 0 # total sequence received - # category counters, each integer counts sequence in its category - self._rel_sample = [0] * self.NUM_OF_CATEGORY - # if last byte in current buffer is not the last byte of a character, - # we need to know how many bytes to skip in next buffer - self._need_to_skip_char_num = 0 - self._last_char_order = -1 # The order of previous char - # If this flag is set to True, detection is done and conclusion has - # been made - self._done = False - - def feed(self, byte_str: Union[bytes, bytearray], num_bytes: int) -> None: - if self._done: - return - - # The buffer we got is byte oriented, and a character may span in more than one - # buffers. In case the last one or two byte in last buffer is not - # complete, we record how many byte needed to complete that character - # and skip these bytes here. We can choose to record those bytes as - # well and analyse the character once it is complete, but since a - # character will not make much difference, by simply skipping - # this character will simply our logic and improve performance. - i = self._need_to_skip_char_num - while i < num_bytes: - order, char_len = self.get_order(byte_str[i : i + 2]) - i += char_len - if i > num_bytes: - self._need_to_skip_char_num = i - num_bytes - self._last_char_order = -1 - else: - if (order != -1) and (self._last_char_order != -1): - self._total_rel += 1 - if self._total_rel > self.MAX_REL_THRESHOLD: - self._done = True - break - self._rel_sample[ - jp2_char_context[self._last_char_order][order] - ] += 1 - self._last_char_order = order - - def got_enough_data(self) -> bool: - return self._total_rel > self.ENOUGH_REL_THRESHOLD - - def get_confidence(self) -> float: - # This is just one way to calculate confidence. It works well for me. - if self._total_rel > self.MINIMUM_DATA_THRESHOLD: - return (self._total_rel - self._rel_sample[0]) / self._total_rel - return self.DONT_KNOW - - def get_order(self, _: Union[bytes, bytearray]) -> Tuple[int, int]: - return -1, 1 - - -class SJISContextAnalysis(JapaneseContextAnalysis): - def __init__(self) -> None: - super().__init__() - self._charset_name = "SHIFT_JIS" - - @property - def charset_name(self) -> str: - return self._charset_name - - def get_order(self, byte_str: Union[bytes, bytearray]) -> Tuple[int, int]: - if not byte_str: - return -1, 1 - # find out current char's byte length - first_char = byte_str[0] - if (0x81 <= first_char <= 0x9F) or (0xE0 <= first_char <= 0xFC): - char_len = 2 - if (first_char == 0x87) or (0xFA <= first_char <= 0xFC): - self._charset_name = "CP932" - else: - char_len = 1 - - # return its order if it is hiragana - if len(byte_str) > 1: - second_char = byte_str[1] - if (first_char == 202) and (0x9F <= second_char <= 0xF1): - return second_char - 0x9F, char_len - - return -1, char_len - - -class EUCJPContextAnalysis(JapaneseContextAnalysis): - def get_order(self, byte_str: Union[bytes, bytearray]) -> Tuple[int, int]: - if not byte_str: - return -1, 1 - # find out current char's byte length - first_char = byte_str[0] - if (first_char == 0x8E) or (0xA1 <= first_char <= 0xFE): - char_len = 2 - elif first_char == 0x8F: - char_len = 3 - else: - char_len = 1 - - # return its order if it is hiragana - if len(byte_str) > 1: - second_char = byte_str[1] - if (first_char == 0xA4) and (0xA1 <= second_char <= 0xF3): - return second_char - 0xA1, char_len - - return -1, char_len diff --git a/venv/lib/python3.10/site-packages/chardet/langbulgarianmodel.py b/venv/lib/python3.10/site-packages/chardet/langbulgarianmodel.py deleted file mode 100644 index 2f771bb8170c1ea374bbfb5156c6d4949ed6e59a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/langbulgarianmodel.py +++ /dev/null @@ -1,4649 +0,0 @@ -from chardet.sbcharsetprober import SingleByteCharSetModel - -# 3: Positive -# 2: Likely -# 1: Unlikely -# 0: Negative - -BULGARIAN_LANG_MODEL = { - 63: { # 'e' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 0, # 'а' - 18: 1, # 'б' - 9: 1, # 'в' - 20: 1, # 'г' - 11: 1, # 'д' - 3: 1, # 'е' - 23: 1, # 'ж' - 15: 1, # 'з' - 2: 0, # 'и' - 26: 1, # 'й' - 12: 1, # 'к' - 10: 1, # 'л' - 14: 1, # 'м' - 6: 1, # 'н' - 4: 1, # 'о' - 13: 1, # 'п' - 7: 1, # 'р' - 8: 1, # 'с' - 5: 1, # 'т' - 19: 0, # 'у' - 29: 1, # 'ф' - 25: 1, # 'х' - 22: 0, # 'ц' - 21: 1, # 'ч' - 27: 1, # 'ш' - 24: 1, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 45: { # '\xad' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 1, # 'Б' - 35: 1, # 'В' - 43: 0, # 'Г' - 37: 1, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 1, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 0, # 'Л' - 38: 1, # 'М' - 36: 0, # 'Н' - 41: 1, # 'О' - 30: 1, # 'П' - 39: 1, # 'Р' - 28: 1, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 1, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 0, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 0, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 0, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 0, # 'л' - 14: 0, # 'м' - 6: 0, # 'н' - 4: 0, # 'о' - 13: 0, # 'п' - 7: 0, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 0, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 31: { # 'А' - 63: 0, # 'e' - 45: 1, # '\xad' - 31: 1, # 'А' - 32: 1, # 'Б' - 35: 2, # 'В' - 43: 1, # 'Г' - 37: 2, # 'Д' - 44: 2, # 'Е' - 55: 1, # 'Ж' - 47: 2, # 'З' - 40: 1, # 'И' - 59: 1, # 'Й' - 33: 1, # 'К' - 46: 2, # 'Л' - 38: 1, # 'М' - 36: 2, # 'Н' - 41: 1, # 'О' - 30: 2, # 'П' - 39: 2, # 'Р' - 28: 2, # 'С' - 34: 2, # 'Т' - 51: 1, # 'У' - 48: 2, # 'Ф' - 49: 1, # 'Х' - 53: 1, # 'Ц' - 50: 1, # 'Ч' - 54: 1, # 'Ш' - 57: 2, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 1, # 'Я' - 1: 1, # 'а' - 18: 2, # 'б' - 9: 2, # 'в' - 20: 2, # 'г' - 11: 2, # 'д' - 3: 1, # 'е' - 23: 1, # 'ж' - 15: 2, # 'з' - 2: 0, # 'и' - 26: 2, # 'й' - 12: 2, # 'к' - 10: 3, # 'л' - 14: 2, # 'м' - 6: 3, # 'н' - 4: 0, # 'о' - 13: 2, # 'п' - 7: 2, # 'р' - 8: 2, # 'с' - 5: 2, # 'т' - 19: 1, # 'у' - 29: 2, # 'ф' - 25: 1, # 'х' - 22: 1, # 'ц' - 21: 1, # 'ч' - 27: 1, # 'ш' - 24: 0, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 32: { # 'Б' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 2, # 'А' - 32: 2, # 'Б' - 35: 1, # 'В' - 43: 1, # 'Г' - 37: 2, # 'Д' - 44: 1, # 'Е' - 55: 1, # 'Ж' - 47: 2, # 'З' - 40: 1, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 1, # 'М' - 36: 2, # 'Н' - 41: 2, # 'О' - 30: 1, # 'П' - 39: 1, # 'Р' - 28: 2, # 'С' - 34: 2, # 'Т' - 51: 1, # 'У' - 48: 2, # 'Ф' - 49: 1, # 'Х' - 53: 1, # 'Ц' - 50: 1, # 'Ч' - 54: 0, # 'Ш' - 57: 1, # 'Щ' - 61: 2, # 'Ъ' - 60: 1, # 'Ю' - 56: 1, # 'Я' - 1: 3, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 1, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 2, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 2, # 'л' - 14: 0, # 'м' - 6: 0, # 'н' - 4: 3, # 'о' - 13: 0, # 'п' - 7: 2, # 'р' - 8: 1, # 'с' - 5: 0, # 'т' - 19: 2, # 'у' - 29: 0, # 'ф' - 25: 1, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 3, # 'ъ' - 52: 1, # 'ь' - 42: 1, # 'ю' - 16: 2, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 35: { # 'В' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 2, # 'А' - 32: 1, # 'Б' - 35: 1, # 'В' - 43: 0, # 'Г' - 37: 1, # 'Д' - 44: 2, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 2, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 1, # 'М' - 36: 1, # 'Н' - 41: 1, # 'О' - 30: 1, # 'П' - 39: 2, # 'Р' - 28: 2, # 'С' - 34: 1, # 'Т' - 51: 1, # 'У' - 48: 2, # 'Ф' - 49: 0, # 'Х' - 53: 1, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 1, # 'Ъ' - 60: 1, # 'Ю' - 56: 2, # 'Я' - 1: 3, # 'а' - 18: 1, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 1, # 'д' - 3: 3, # 'е' - 23: 1, # 'ж' - 15: 2, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 1, # 'к' - 10: 2, # 'л' - 14: 1, # 'м' - 6: 2, # 'н' - 4: 2, # 'о' - 13: 1, # 'п' - 7: 2, # 'р' - 8: 2, # 'с' - 5: 2, # 'т' - 19: 1, # 'у' - 29: 0, # 'ф' - 25: 1, # 'х' - 22: 0, # 'ц' - 21: 2, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 2, # 'ъ' - 52: 1, # 'ь' - 42: 1, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 43: { # 'Г' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 2, # 'А' - 32: 1, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 1, # 'Д' - 44: 2, # 'Е' - 55: 0, # 'Ж' - 47: 1, # 'З' - 40: 1, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 0, # 'М' - 36: 1, # 'Н' - 41: 1, # 'О' - 30: 0, # 'П' - 39: 1, # 'Р' - 28: 1, # 'С' - 34: 0, # 'Т' - 51: 1, # 'У' - 48: 1, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 1, # 'Щ' - 61: 1, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 2, # 'а' - 18: 1, # 'б' - 9: 1, # 'в' - 20: 0, # 'г' - 11: 1, # 'д' - 3: 3, # 'е' - 23: 1, # 'ж' - 15: 0, # 'з' - 2: 2, # 'и' - 26: 0, # 'й' - 12: 1, # 'к' - 10: 2, # 'л' - 14: 1, # 'м' - 6: 1, # 'н' - 4: 2, # 'о' - 13: 0, # 'п' - 7: 2, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 2, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 1, # 'щ' - 17: 2, # 'ъ' - 52: 1, # 'ь' - 42: 1, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 37: { # 'Д' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 2, # 'А' - 32: 1, # 'Б' - 35: 2, # 'В' - 43: 1, # 'Г' - 37: 2, # 'Д' - 44: 2, # 'Е' - 55: 2, # 'Ж' - 47: 1, # 'З' - 40: 2, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 1, # 'М' - 36: 1, # 'Н' - 41: 2, # 'О' - 30: 2, # 'П' - 39: 1, # 'Р' - 28: 2, # 'С' - 34: 1, # 'Т' - 51: 1, # 'У' - 48: 1, # 'Ф' - 49: 0, # 'Х' - 53: 1, # 'Ц' - 50: 1, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 1, # 'Ъ' - 60: 1, # 'Ю' - 56: 1, # 'Я' - 1: 3, # 'а' - 18: 0, # 'б' - 9: 2, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 3, # 'е' - 23: 3, # 'ж' - 15: 1, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 1, # 'л' - 14: 1, # 'м' - 6: 2, # 'н' - 4: 3, # 'о' - 13: 0, # 'п' - 7: 2, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 2, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 2, # 'ъ' - 52: 1, # 'ь' - 42: 2, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 44: { # 'Е' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 1, # 'А' - 32: 1, # 'Б' - 35: 2, # 'В' - 43: 1, # 'Г' - 37: 1, # 'Д' - 44: 1, # 'Е' - 55: 1, # 'Ж' - 47: 1, # 'З' - 40: 1, # 'И' - 59: 1, # 'Й' - 33: 2, # 'К' - 46: 2, # 'Л' - 38: 1, # 'М' - 36: 2, # 'Н' - 41: 2, # 'О' - 30: 1, # 'П' - 39: 2, # 'Р' - 28: 2, # 'С' - 34: 2, # 'Т' - 51: 1, # 'У' - 48: 2, # 'Ф' - 49: 1, # 'Х' - 53: 2, # 'Ц' - 50: 1, # 'Ч' - 54: 1, # 'Ш' - 57: 1, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 1, # 'Я' - 1: 0, # 'а' - 18: 1, # 'б' - 9: 2, # 'в' - 20: 1, # 'г' - 11: 2, # 'д' - 3: 0, # 'е' - 23: 1, # 'ж' - 15: 1, # 'з' - 2: 0, # 'и' - 26: 1, # 'й' - 12: 2, # 'к' - 10: 2, # 'л' - 14: 2, # 'м' - 6: 2, # 'н' - 4: 0, # 'о' - 13: 1, # 'п' - 7: 2, # 'р' - 8: 2, # 'с' - 5: 1, # 'т' - 19: 1, # 'у' - 29: 1, # 'ф' - 25: 1, # 'х' - 22: 0, # 'ц' - 21: 1, # 'ч' - 27: 1, # 'ш' - 24: 1, # 'щ' - 17: 1, # 'ъ' - 52: 0, # 'ь' - 42: 1, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 55: { # 'Ж' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 1, # 'А' - 32: 0, # 'Б' - 35: 1, # 'В' - 43: 0, # 'Г' - 37: 1, # 'Д' - 44: 1, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 1, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 1, # 'Н' - 41: 1, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 1, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 2, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 1, # 'д' - 3: 2, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 2, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 0, # 'л' - 14: 0, # 'м' - 6: 0, # 'н' - 4: 2, # 'о' - 13: 1, # 'п' - 7: 1, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 1, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 1, # 'ъ' - 52: 1, # 'ь' - 42: 1, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 47: { # 'З' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 2, # 'А' - 32: 1, # 'Б' - 35: 1, # 'В' - 43: 1, # 'Г' - 37: 1, # 'Д' - 44: 1, # 'Е' - 55: 0, # 'Ж' - 47: 1, # 'З' - 40: 1, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 1, # 'М' - 36: 2, # 'Н' - 41: 1, # 'О' - 30: 1, # 'П' - 39: 1, # 'Р' - 28: 1, # 'С' - 34: 1, # 'Т' - 51: 1, # 'У' - 48: 0, # 'Ф' - 49: 1, # 'Х' - 53: 1, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 1, # 'Ъ' - 60: 0, # 'Ю' - 56: 1, # 'Я' - 1: 3, # 'а' - 18: 1, # 'б' - 9: 2, # 'в' - 20: 1, # 'г' - 11: 2, # 'д' - 3: 2, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 1, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 2, # 'л' - 14: 1, # 'м' - 6: 1, # 'н' - 4: 1, # 'о' - 13: 0, # 'п' - 7: 1, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 1, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 1, # 'ъ' - 52: 0, # 'ь' - 42: 1, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 40: { # 'И' - 63: 0, # 'e' - 45: 1, # '\xad' - 31: 1, # 'А' - 32: 1, # 'Б' - 35: 1, # 'В' - 43: 1, # 'Г' - 37: 1, # 'Д' - 44: 2, # 'Е' - 55: 1, # 'Ж' - 47: 2, # 'З' - 40: 1, # 'И' - 59: 1, # 'Й' - 33: 2, # 'К' - 46: 2, # 'Л' - 38: 2, # 'М' - 36: 2, # 'Н' - 41: 1, # 'О' - 30: 1, # 'П' - 39: 2, # 'Р' - 28: 2, # 'С' - 34: 2, # 'Т' - 51: 0, # 'У' - 48: 1, # 'Ф' - 49: 1, # 'Х' - 53: 1, # 'Ц' - 50: 1, # 'Ч' - 54: 1, # 'Ш' - 57: 1, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 2, # 'Я' - 1: 1, # 'а' - 18: 1, # 'б' - 9: 3, # 'в' - 20: 2, # 'г' - 11: 1, # 'д' - 3: 1, # 'е' - 23: 0, # 'ж' - 15: 3, # 'з' - 2: 0, # 'и' - 26: 1, # 'й' - 12: 1, # 'к' - 10: 2, # 'л' - 14: 2, # 'м' - 6: 2, # 'н' - 4: 0, # 'о' - 13: 1, # 'п' - 7: 2, # 'р' - 8: 2, # 'с' - 5: 2, # 'т' - 19: 0, # 'у' - 29: 1, # 'ф' - 25: 1, # 'х' - 22: 1, # 'ц' - 21: 1, # 'ч' - 27: 1, # 'ш' - 24: 1, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 59: { # 'Й' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 1, # 'Д' - 44: 1, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 1, # 'М' - 36: 1, # 'Н' - 41: 1, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 1, # 'С' - 34: 1, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 1, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 1, # 'Я' - 1: 0, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 1, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 0, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 0, # 'л' - 14: 0, # 'м' - 6: 0, # 'н' - 4: 2, # 'о' - 13: 0, # 'п' - 7: 0, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 0, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 1, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 33: { # 'К' - 63: 0, # 'e' - 45: 1, # '\xad' - 31: 2, # 'А' - 32: 1, # 'Б' - 35: 1, # 'В' - 43: 1, # 'Г' - 37: 1, # 'Д' - 44: 1, # 'Е' - 55: 0, # 'Ж' - 47: 1, # 'З' - 40: 2, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 0, # 'М' - 36: 2, # 'Н' - 41: 2, # 'О' - 30: 2, # 'П' - 39: 1, # 'Р' - 28: 2, # 'С' - 34: 1, # 'Т' - 51: 1, # 'У' - 48: 1, # 'Ф' - 49: 1, # 'Х' - 53: 1, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 1, # 'Ъ' - 60: 1, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 0, # 'б' - 9: 1, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 2, # 'е' - 23: 1, # 'ж' - 15: 0, # 'з' - 2: 2, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 2, # 'л' - 14: 1, # 'м' - 6: 2, # 'н' - 4: 3, # 'о' - 13: 0, # 'п' - 7: 3, # 'р' - 8: 1, # 'с' - 5: 0, # 'т' - 19: 2, # 'у' - 29: 0, # 'ф' - 25: 1, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 1, # 'ш' - 24: 0, # 'щ' - 17: 2, # 'ъ' - 52: 1, # 'ь' - 42: 2, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 46: { # 'Л' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 2, # 'А' - 32: 1, # 'Б' - 35: 1, # 'В' - 43: 2, # 'Г' - 37: 1, # 'Д' - 44: 2, # 'Е' - 55: 0, # 'Ж' - 47: 1, # 'З' - 40: 2, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 0, # 'М' - 36: 1, # 'Н' - 41: 2, # 'О' - 30: 1, # 'П' - 39: 0, # 'Р' - 28: 1, # 'С' - 34: 1, # 'Т' - 51: 1, # 'У' - 48: 0, # 'Ф' - 49: 1, # 'Х' - 53: 1, # 'Ц' - 50: 1, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 1, # 'Ъ' - 60: 1, # 'Ю' - 56: 1, # 'Я' - 1: 2, # 'а' - 18: 0, # 'б' - 9: 1, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 2, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 0, # 'л' - 14: 0, # 'м' - 6: 0, # 'н' - 4: 2, # 'о' - 13: 0, # 'п' - 7: 0, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 2, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 1, # 'ъ' - 52: 1, # 'ь' - 42: 2, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 38: { # 'М' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 2, # 'А' - 32: 1, # 'Б' - 35: 2, # 'В' - 43: 0, # 'Г' - 37: 1, # 'Д' - 44: 1, # 'Е' - 55: 0, # 'Ж' - 47: 1, # 'З' - 40: 2, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 1, # 'М' - 36: 1, # 'Н' - 41: 2, # 'О' - 30: 1, # 'П' - 39: 1, # 'Р' - 28: 2, # 'С' - 34: 1, # 'Т' - 51: 1, # 'У' - 48: 1, # 'Ф' - 49: 0, # 'Х' - 53: 1, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 1, # 'Ъ' - 60: 0, # 'Ю' - 56: 1, # 'Я' - 1: 3, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 2, # 'л' - 14: 0, # 'м' - 6: 2, # 'н' - 4: 3, # 'о' - 13: 0, # 'п' - 7: 1, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 2, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 2, # 'ъ' - 52: 1, # 'ь' - 42: 2, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 36: { # 'Н' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 2, # 'А' - 32: 2, # 'Б' - 35: 1, # 'В' - 43: 1, # 'Г' - 37: 2, # 'Д' - 44: 2, # 'Е' - 55: 1, # 'Ж' - 47: 1, # 'З' - 40: 2, # 'И' - 59: 1, # 'Й' - 33: 2, # 'К' - 46: 1, # 'Л' - 38: 1, # 'М' - 36: 1, # 'Н' - 41: 2, # 'О' - 30: 1, # 'П' - 39: 1, # 'Р' - 28: 2, # 'С' - 34: 2, # 'Т' - 51: 1, # 'У' - 48: 1, # 'Ф' - 49: 1, # 'Х' - 53: 1, # 'Ц' - 50: 1, # 'Ч' - 54: 1, # 'Ш' - 57: 0, # 'Щ' - 61: 1, # 'Ъ' - 60: 1, # 'Ю' - 56: 1, # 'Я' - 1: 3, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 1, # 'г' - 11: 0, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 0, # 'л' - 14: 0, # 'м' - 6: 0, # 'н' - 4: 3, # 'о' - 13: 0, # 'п' - 7: 0, # 'р' - 8: 0, # 'с' - 5: 1, # 'т' - 19: 1, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 1, # 'ш' - 24: 0, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 2, # 'ю' - 16: 2, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 41: { # 'О' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 1, # 'А' - 32: 1, # 'Б' - 35: 2, # 'В' - 43: 1, # 'Г' - 37: 2, # 'Д' - 44: 1, # 'Е' - 55: 1, # 'Ж' - 47: 1, # 'З' - 40: 1, # 'И' - 59: 1, # 'Й' - 33: 2, # 'К' - 46: 2, # 'Л' - 38: 2, # 'М' - 36: 2, # 'Н' - 41: 2, # 'О' - 30: 1, # 'П' - 39: 2, # 'Р' - 28: 2, # 'С' - 34: 2, # 'Т' - 51: 1, # 'У' - 48: 1, # 'Ф' - 49: 1, # 'Х' - 53: 0, # 'Ц' - 50: 1, # 'Ч' - 54: 1, # 'Ш' - 57: 1, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 1, # 'Я' - 1: 1, # 'а' - 18: 2, # 'б' - 9: 2, # 'в' - 20: 2, # 'г' - 11: 1, # 'д' - 3: 1, # 'е' - 23: 1, # 'ж' - 15: 1, # 'з' - 2: 0, # 'и' - 26: 1, # 'й' - 12: 2, # 'к' - 10: 2, # 'л' - 14: 1, # 'м' - 6: 1, # 'н' - 4: 0, # 'о' - 13: 2, # 'п' - 7: 2, # 'р' - 8: 2, # 'с' - 5: 3, # 'т' - 19: 1, # 'у' - 29: 1, # 'ф' - 25: 1, # 'х' - 22: 1, # 'ц' - 21: 2, # 'ч' - 27: 0, # 'ш' - 24: 2, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 30: { # 'П' - 63: 0, # 'e' - 45: 1, # '\xad' - 31: 2, # 'А' - 32: 1, # 'Б' - 35: 1, # 'В' - 43: 1, # 'Г' - 37: 1, # 'Д' - 44: 1, # 'Е' - 55: 0, # 'Ж' - 47: 1, # 'З' - 40: 2, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 1, # 'М' - 36: 1, # 'Н' - 41: 2, # 'О' - 30: 2, # 'П' - 39: 2, # 'Р' - 28: 2, # 'С' - 34: 1, # 'Т' - 51: 2, # 'У' - 48: 1, # 'Ф' - 49: 0, # 'Х' - 53: 1, # 'Ц' - 50: 1, # 'Ч' - 54: 1, # 'Ш' - 57: 0, # 'Щ' - 61: 1, # 'Ъ' - 60: 1, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 2, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 2, # 'и' - 26: 0, # 'й' - 12: 1, # 'к' - 10: 3, # 'л' - 14: 0, # 'м' - 6: 1, # 'н' - 4: 3, # 'о' - 13: 0, # 'п' - 7: 3, # 'р' - 8: 1, # 'с' - 5: 1, # 'т' - 19: 2, # 'у' - 29: 1, # 'ф' - 25: 1, # 'х' - 22: 0, # 'ц' - 21: 1, # 'ч' - 27: 1, # 'ш' - 24: 0, # 'щ' - 17: 2, # 'ъ' - 52: 1, # 'ь' - 42: 1, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 39: { # 'Р' - 63: 0, # 'e' - 45: 1, # '\xad' - 31: 2, # 'А' - 32: 1, # 'Б' - 35: 1, # 'В' - 43: 2, # 'Г' - 37: 2, # 'Д' - 44: 2, # 'Е' - 55: 0, # 'Ж' - 47: 1, # 'З' - 40: 2, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 0, # 'Л' - 38: 1, # 'М' - 36: 1, # 'Н' - 41: 2, # 'О' - 30: 2, # 'П' - 39: 1, # 'Р' - 28: 1, # 'С' - 34: 1, # 'Т' - 51: 1, # 'У' - 48: 1, # 'Ф' - 49: 1, # 'Х' - 53: 1, # 'Ц' - 50: 1, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 1, # 'Ъ' - 60: 1, # 'Ю' - 56: 1, # 'Я' - 1: 3, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 2, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 2, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 0, # 'л' - 14: 0, # 'м' - 6: 1, # 'н' - 4: 3, # 'о' - 13: 0, # 'п' - 7: 0, # 'р' - 8: 1, # 'с' - 5: 0, # 'т' - 19: 3, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 1, # 'ъ' - 52: 0, # 'ь' - 42: 1, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 28: { # 'С' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 3, # 'А' - 32: 2, # 'Б' - 35: 2, # 'В' - 43: 1, # 'Г' - 37: 2, # 'Д' - 44: 2, # 'Е' - 55: 1, # 'Ж' - 47: 1, # 'З' - 40: 2, # 'И' - 59: 0, # 'Й' - 33: 2, # 'К' - 46: 1, # 'Л' - 38: 1, # 'М' - 36: 1, # 'Н' - 41: 2, # 'О' - 30: 2, # 'П' - 39: 1, # 'Р' - 28: 2, # 'С' - 34: 2, # 'Т' - 51: 1, # 'У' - 48: 1, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 1, # 'Ъ' - 60: 1, # 'Ю' - 56: 1, # 'Я' - 1: 3, # 'а' - 18: 1, # 'б' - 9: 2, # 'в' - 20: 1, # 'г' - 11: 1, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 2, # 'к' - 10: 3, # 'л' - 14: 2, # 'м' - 6: 1, # 'н' - 4: 3, # 'о' - 13: 3, # 'п' - 7: 2, # 'р' - 8: 0, # 'с' - 5: 3, # 'т' - 19: 2, # 'у' - 29: 2, # 'ф' - 25: 1, # 'х' - 22: 1, # 'ц' - 21: 1, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 3, # 'ъ' - 52: 1, # 'ь' - 42: 1, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 34: { # 'Т' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 2, # 'А' - 32: 2, # 'Б' - 35: 1, # 'В' - 43: 0, # 'Г' - 37: 1, # 'Д' - 44: 2, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 2, # 'И' - 59: 0, # 'Й' - 33: 2, # 'К' - 46: 1, # 'Л' - 38: 1, # 'М' - 36: 1, # 'Н' - 41: 2, # 'О' - 30: 1, # 'П' - 39: 2, # 'Р' - 28: 2, # 'С' - 34: 1, # 'Т' - 51: 1, # 'У' - 48: 1, # 'Ф' - 49: 0, # 'Х' - 53: 1, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 1, # 'Ъ' - 60: 0, # 'Ю' - 56: 1, # 'Я' - 1: 3, # 'а' - 18: 1, # 'б' - 9: 1, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 2, # 'и' - 26: 0, # 'й' - 12: 1, # 'к' - 10: 1, # 'л' - 14: 0, # 'м' - 6: 0, # 'н' - 4: 3, # 'о' - 13: 0, # 'п' - 7: 3, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 2, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 2, # 'ъ' - 52: 0, # 'ь' - 42: 1, # 'ю' - 16: 2, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 51: { # 'У' - 63: 0, # 'e' - 45: 1, # '\xad' - 31: 1, # 'А' - 32: 1, # 'Б' - 35: 1, # 'В' - 43: 1, # 'Г' - 37: 1, # 'Д' - 44: 2, # 'Е' - 55: 1, # 'Ж' - 47: 1, # 'З' - 40: 1, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 1, # 'М' - 36: 1, # 'Н' - 41: 0, # 'О' - 30: 1, # 'П' - 39: 1, # 'Р' - 28: 1, # 'С' - 34: 2, # 'Т' - 51: 0, # 'У' - 48: 1, # 'Ф' - 49: 1, # 'Х' - 53: 1, # 'Ц' - 50: 1, # 'Ч' - 54: 1, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 1, # 'а' - 18: 1, # 'б' - 9: 2, # 'в' - 20: 1, # 'г' - 11: 1, # 'д' - 3: 2, # 'е' - 23: 1, # 'ж' - 15: 1, # 'з' - 2: 2, # 'и' - 26: 1, # 'й' - 12: 2, # 'к' - 10: 1, # 'л' - 14: 1, # 'м' - 6: 2, # 'н' - 4: 2, # 'о' - 13: 1, # 'п' - 7: 1, # 'р' - 8: 2, # 'с' - 5: 1, # 'т' - 19: 1, # 'у' - 29: 0, # 'ф' - 25: 1, # 'х' - 22: 0, # 'ц' - 21: 2, # 'ч' - 27: 1, # 'ш' - 24: 0, # 'щ' - 17: 1, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 48: { # 'Ф' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 2, # 'А' - 32: 1, # 'Б' - 35: 1, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 1, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 2, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 0, # 'М' - 36: 1, # 'Н' - 41: 1, # 'О' - 30: 2, # 'П' - 39: 1, # 'Р' - 28: 2, # 'С' - 34: 1, # 'Т' - 51: 1, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 2, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 2, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 2, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 2, # 'л' - 14: 0, # 'м' - 6: 0, # 'н' - 4: 2, # 'о' - 13: 0, # 'п' - 7: 2, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 1, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 1, # 'ъ' - 52: 1, # 'ь' - 42: 1, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 49: { # 'Х' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 1, # 'А' - 32: 0, # 'Б' - 35: 1, # 'В' - 43: 1, # 'Г' - 37: 1, # 'Д' - 44: 1, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 1, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 1, # 'Л' - 38: 1, # 'М' - 36: 1, # 'Н' - 41: 1, # 'О' - 30: 1, # 'П' - 39: 1, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 1, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 2, # 'а' - 18: 0, # 'б' - 9: 1, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 2, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 2, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 1, # 'л' - 14: 1, # 'м' - 6: 0, # 'н' - 4: 2, # 'о' - 13: 0, # 'п' - 7: 2, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 2, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 2, # 'ъ' - 52: 1, # 'ь' - 42: 1, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 53: { # 'Ц' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 1, # 'А' - 32: 0, # 'Б' - 35: 1, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 1, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 2, # 'И' - 59: 0, # 'Й' - 33: 2, # 'К' - 46: 1, # 'Л' - 38: 1, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 1, # 'Р' - 28: 2, # 'С' - 34: 0, # 'Т' - 51: 1, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 2, # 'а' - 18: 0, # 'б' - 9: 2, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 2, # 'е' - 23: 0, # 'ж' - 15: 1, # 'з' - 2: 2, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 0, # 'л' - 14: 0, # 'м' - 6: 0, # 'н' - 4: 1, # 'о' - 13: 0, # 'п' - 7: 1, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 1, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 1, # 'ъ' - 52: 0, # 'ь' - 42: 1, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 50: { # 'Ч' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 2, # 'А' - 32: 1, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 1, # 'Е' - 55: 0, # 'Ж' - 47: 1, # 'З' - 40: 1, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 0, # 'М' - 36: 1, # 'Н' - 41: 1, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 1, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 2, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 3, # 'е' - 23: 1, # 'ж' - 15: 0, # 'з' - 2: 2, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 1, # 'л' - 14: 0, # 'м' - 6: 0, # 'н' - 4: 2, # 'о' - 13: 0, # 'п' - 7: 1, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 2, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 1, # 'ъ' - 52: 1, # 'ь' - 42: 0, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 54: { # 'Ш' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 1, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 1, # 'Е' - 55: 0, # 'Ж' - 47: 1, # 'З' - 40: 1, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 1, # 'Н' - 41: 1, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 1, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 2, # 'а' - 18: 0, # 'б' - 9: 2, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 2, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 2, # 'и' - 26: 0, # 'й' - 12: 1, # 'к' - 10: 1, # 'л' - 14: 1, # 'м' - 6: 1, # 'н' - 4: 2, # 'о' - 13: 1, # 'п' - 7: 1, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 2, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 1, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 1, # 'ъ' - 52: 1, # 'ь' - 42: 0, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 57: { # 'Щ' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 1, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 1, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 1, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 1, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 2, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 2, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 1, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 0, # 'л' - 14: 0, # 'м' - 6: 0, # 'н' - 4: 1, # 'о' - 13: 0, # 'п' - 7: 1, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 1, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 1, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 61: { # 'Ъ' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 1, # 'Б' - 35: 1, # 'В' - 43: 0, # 'Г' - 37: 1, # 'Д' - 44: 0, # 'Е' - 55: 1, # 'Ж' - 47: 1, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 2, # 'Л' - 38: 1, # 'М' - 36: 1, # 'Н' - 41: 0, # 'О' - 30: 1, # 'П' - 39: 2, # 'Р' - 28: 1, # 'С' - 34: 1, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 1, # 'Х' - 53: 1, # 'Ц' - 50: 1, # 'Ч' - 54: 1, # 'Ш' - 57: 1, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 0, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 0, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 0, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 1, # 'л' - 14: 0, # 'м' - 6: 1, # 'н' - 4: 0, # 'о' - 13: 0, # 'п' - 7: 1, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 0, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 60: { # 'Ю' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 1, # 'А' - 32: 1, # 'Б' - 35: 0, # 'В' - 43: 1, # 'Г' - 37: 1, # 'Д' - 44: 0, # 'Е' - 55: 1, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 0, # 'М' - 36: 1, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 1, # 'Р' - 28: 1, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 0, # 'а' - 18: 1, # 'б' - 9: 1, # 'в' - 20: 2, # 'г' - 11: 1, # 'д' - 3: 0, # 'е' - 23: 2, # 'ж' - 15: 1, # 'з' - 2: 1, # 'и' - 26: 0, # 'й' - 12: 1, # 'к' - 10: 1, # 'л' - 14: 1, # 'м' - 6: 1, # 'н' - 4: 0, # 'о' - 13: 1, # 'п' - 7: 1, # 'р' - 8: 1, # 'с' - 5: 1, # 'т' - 19: 0, # 'у' - 29: 0, # 'ф' - 25: 1, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 56: { # 'Я' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 1, # 'Б' - 35: 1, # 'В' - 43: 1, # 'Г' - 37: 1, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 1, # 'Л' - 38: 1, # 'М' - 36: 1, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 1, # 'С' - 34: 2, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 0, # 'а' - 18: 1, # 'б' - 9: 1, # 'в' - 20: 1, # 'г' - 11: 1, # 'д' - 3: 0, # 'е' - 23: 0, # 'ж' - 15: 1, # 'з' - 2: 1, # 'и' - 26: 1, # 'й' - 12: 1, # 'к' - 10: 1, # 'л' - 14: 2, # 'м' - 6: 2, # 'н' - 4: 0, # 'о' - 13: 2, # 'п' - 7: 1, # 'р' - 8: 1, # 'с' - 5: 1, # 'т' - 19: 0, # 'у' - 29: 0, # 'ф' - 25: 1, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 1, # 'ш' - 24: 0, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 1, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 1: { # 'а' - 63: 1, # 'e' - 45: 1, # '\xad' - 31: 1, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 1, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 1, # 'а' - 18: 3, # 'б' - 9: 3, # 'в' - 20: 3, # 'г' - 11: 3, # 'д' - 3: 3, # 'е' - 23: 3, # 'ж' - 15: 3, # 'з' - 2: 3, # 'и' - 26: 3, # 'й' - 12: 3, # 'к' - 10: 3, # 'л' - 14: 3, # 'м' - 6: 3, # 'н' - 4: 2, # 'о' - 13: 3, # 'п' - 7: 3, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 3, # 'у' - 29: 3, # 'ф' - 25: 3, # 'х' - 22: 3, # 'ц' - 21: 3, # 'ч' - 27: 3, # 'ш' - 24: 3, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 1, # 'ю' - 16: 3, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 18: { # 'б' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 0, # 'б' - 9: 3, # 'в' - 20: 1, # 'г' - 11: 2, # 'д' - 3: 3, # 'е' - 23: 1, # 'ж' - 15: 1, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 1, # 'к' - 10: 3, # 'л' - 14: 2, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 1, # 'п' - 7: 3, # 'р' - 8: 3, # 'с' - 5: 0, # 'т' - 19: 3, # 'у' - 29: 0, # 'ф' - 25: 2, # 'х' - 22: 1, # 'ц' - 21: 1, # 'ч' - 27: 1, # 'ш' - 24: 3, # 'щ' - 17: 3, # 'ъ' - 52: 1, # 'ь' - 42: 2, # 'ю' - 16: 3, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 9: { # 'в' - 63: 1, # 'e' - 45: 1, # '\xad' - 31: 0, # 'А' - 32: 1, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 1, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 1, # 'б' - 9: 0, # 'в' - 20: 2, # 'г' - 11: 3, # 'д' - 3: 3, # 'е' - 23: 1, # 'ж' - 15: 3, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 3, # 'к' - 10: 3, # 'л' - 14: 2, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 2, # 'п' - 7: 3, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 2, # 'у' - 29: 0, # 'ф' - 25: 2, # 'х' - 22: 2, # 'ц' - 21: 3, # 'ч' - 27: 2, # 'ш' - 24: 1, # 'щ' - 17: 3, # 'ъ' - 52: 1, # 'ь' - 42: 2, # 'ю' - 16: 3, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 20: { # 'г' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 1, # 'б' - 9: 2, # 'в' - 20: 1, # 'г' - 11: 2, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 1, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 1, # 'к' - 10: 3, # 'л' - 14: 1, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 1, # 'п' - 7: 3, # 'р' - 8: 2, # 'с' - 5: 2, # 'т' - 19: 3, # 'у' - 29: 1, # 'ф' - 25: 1, # 'х' - 22: 0, # 'ц' - 21: 1, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 3, # 'ъ' - 52: 1, # 'ь' - 42: 1, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 11: { # 'д' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 2, # 'б' - 9: 3, # 'в' - 20: 2, # 'г' - 11: 2, # 'д' - 3: 3, # 'е' - 23: 3, # 'ж' - 15: 2, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 3, # 'к' - 10: 3, # 'л' - 14: 3, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 3, # 'п' - 7: 3, # 'р' - 8: 3, # 'с' - 5: 1, # 'т' - 19: 3, # 'у' - 29: 1, # 'ф' - 25: 2, # 'х' - 22: 2, # 'ц' - 21: 2, # 'ч' - 27: 1, # 'ш' - 24: 1, # 'щ' - 17: 3, # 'ъ' - 52: 1, # 'ь' - 42: 1, # 'ю' - 16: 3, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 3: { # 'е' - 63: 0, # 'e' - 45: 1, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 2, # 'а' - 18: 3, # 'б' - 9: 3, # 'в' - 20: 3, # 'г' - 11: 3, # 'д' - 3: 2, # 'е' - 23: 3, # 'ж' - 15: 3, # 'з' - 2: 2, # 'и' - 26: 3, # 'й' - 12: 3, # 'к' - 10: 3, # 'л' - 14: 3, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 3, # 'п' - 7: 3, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 2, # 'у' - 29: 3, # 'ф' - 25: 3, # 'х' - 22: 3, # 'ц' - 21: 3, # 'ч' - 27: 3, # 'ш' - 24: 3, # 'щ' - 17: 1, # 'ъ' - 52: 0, # 'ь' - 42: 1, # 'ю' - 16: 3, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 23: { # 'ж' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 3, # 'б' - 9: 2, # 'в' - 20: 1, # 'г' - 11: 3, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 2, # 'к' - 10: 1, # 'л' - 14: 1, # 'м' - 6: 3, # 'н' - 4: 2, # 'о' - 13: 1, # 'п' - 7: 1, # 'р' - 8: 1, # 'с' - 5: 1, # 'т' - 19: 2, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 1, # 'ц' - 21: 1, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 2, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 15: { # 'з' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 3, # 'б' - 9: 3, # 'в' - 20: 3, # 'г' - 11: 3, # 'д' - 3: 3, # 'е' - 23: 1, # 'ж' - 15: 1, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 3, # 'к' - 10: 3, # 'л' - 14: 3, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 3, # 'п' - 7: 3, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 3, # 'у' - 29: 1, # 'ф' - 25: 2, # 'х' - 22: 2, # 'ц' - 21: 2, # 'ч' - 27: 2, # 'ш' - 24: 1, # 'щ' - 17: 2, # 'ъ' - 52: 1, # 'ь' - 42: 1, # 'ю' - 16: 2, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 2: { # 'и' - 63: 1, # 'e' - 45: 1, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 1, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 1, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 1, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 1, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 3, # 'б' - 9: 3, # 'в' - 20: 3, # 'г' - 11: 3, # 'д' - 3: 3, # 'е' - 23: 3, # 'ж' - 15: 3, # 'з' - 2: 3, # 'и' - 26: 3, # 'й' - 12: 3, # 'к' - 10: 3, # 'л' - 14: 3, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 3, # 'п' - 7: 3, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 2, # 'у' - 29: 3, # 'ф' - 25: 3, # 'х' - 22: 3, # 'ц' - 21: 3, # 'ч' - 27: 3, # 'ш' - 24: 3, # 'щ' - 17: 2, # 'ъ' - 52: 0, # 'ь' - 42: 1, # 'ю' - 16: 3, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 26: { # 'й' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 1, # 'а' - 18: 2, # 'б' - 9: 2, # 'в' - 20: 1, # 'г' - 11: 2, # 'д' - 3: 2, # 'е' - 23: 0, # 'ж' - 15: 2, # 'з' - 2: 1, # 'и' - 26: 0, # 'й' - 12: 3, # 'к' - 10: 2, # 'л' - 14: 2, # 'м' - 6: 3, # 'н' - 4: 2, # 'о' - 13: 1, # 'п' - 7: 2, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 1, # 'у' - 29: 2, # 'ф' - 25: 1, # 'х' - 22: 2, # 'ц' - 21: 2, # 'ч' - 27: 1, # 'ш' - 24: 1, # 'щ' - 17: 1, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 12: { # 'к' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 1, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 1, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 1, # 'б' - 9: 3, # 'в' - 20: 2, # 'г' - 11: 1, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 2, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 1, # 'к' - 10: 3, # 'л' - 14: 2, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 1, # 'п' - 7: 3, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 3, # 'у' - 29: 1, # 'ф' - 25: 1, # 'х' - 22: 3, # 'ц' - 21: 2, # 'ч' - 27: 1, # 'ш' - 24: 0, # 'щ' - 17: 3, # 'ъ' - 52: 1, # 'ь' - 42: 2, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 10: { # 'л' - 63: 1, # 'e' - 45: 1, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 1, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 3, # 'б' - 9: 3, # 'в' - 20: 3, # 'г' - 11: 2, # 'д' - 3: 3, # 'е' - 23: 3, # 'ж' - 15: 2, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 3, # 'к' - 10: 1, # 'л' - 14: 2, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 2, # 'п' - 7: 2, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 3, # 'у' - 29: 2, # 'ф' - 25: 2, # 'х' - 22: 2, # 'ц' - 21: 2, # 'ч' - 27: 2, # 'ш' - 24: 1, # 'щ' - 17: 3, # 'ъ' - 52: 2, # 'ь' - 42: 3, # 'ю' - 16: 3, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 14: { # 'м' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 1, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 3, # 'б' - 9: 3, # 'в' - 20: 1, # 'г' - 11: 1, # 'д' - 3: 3, # 'е' - 23: 1, # 'ж' - 15: 1, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 2, # 'к' - 10: 3, # 'л' - 14: 1, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 3, # 'п' - 7: 2, # 'р' - 8: 2, # 'с' - 5: 1, # 'т' - 19: 3, # 'у' - 29: 2, # 'ф' - 25: 1, # 'х' - 22: 2, # 'ц' - 21: 2, # 'ч' - 27: 2, # 'ш' - 24: 1, # 'щ' - 17: 3, # 'ъ' - 52: 1, # 'ь' - 42: 2, # 'ю' - 16: 3, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 6: { # 'н' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 1, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 2, # 'б' - 9: 2, # 'в' - 20: 3, # 'г' - 11: 3, # 'д' - 3: 3, # 'е' - 23: 2, # 'ж' - 15: 2, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 3, # 'к' - 10: 2, # 'л' - 14: 1, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 1, # 'п' - 7: 2, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 3, # 'у' - 29: 3, # 'ф' - 25: 2, # 'х' - 22: 3, # 'ц' - 21: 3, # 'ч' - 27: 2, # 'ш' - 24: 1, # 'щ' - 17: 3, # 'ъ' - 52: 2, # 'ь' - 42: 2, # 'ю' - 16: 3, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 4: { # 'о' - 63: 0, # 'e' - 45: 1, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 2, # 'а' - 18: 3, # 'б' - 9: 3, # 'в' - 20: 3, # 'г' - 11: 3, # 'д' - 3: 3, # 'е' - 23: 3, # 'ж' - 15: 3, # 'з' - 2: 3, # 'и' - 26: 3, # 'й' - 12: 3, # 'к' - 10: 3, # 'л' - 14: 3, # 'м' - 6: 3, # 'н' - 4: 2, # 'о' - 13: 3, # 'п' - 7: 3, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 2, # 'у' - 29: 3, # 'ф' - 25: 3, # 'х' - 22: 3, # 'ц' - 21: 3, # 'ч' - 27: 3, # 'ш' - 24: 3, # 'щ' - 17: 1, # 'ъ' - 52: 0, # 'ь' - 42: 1, # 'ю' - 16: 3, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 13: { # 'п' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 1, # 'б' - 9: 2, # 'в' - 20: 1, # 'г' - 11: 1, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 1, # 'з' - 2: 3, # 'и' - 26: 1, # 'й' - 12: 2, # 'к' - 10: 3, # 'л' - 14: 1, # 'м' - 6: 2, # 'н' - 4: 3, # 'о' - 13: 1, # 'п' - 7: 3, # 'р' - 8: 2, # 'с' - 5: 2, # 'т' - 19: 3, # 'у' - 29: 1, # 'ф' - 25: 1, # 'х' - 22: 2, # 'ц' - 21: 2, # 'ч' - 27: 1, # 'ш' - 24: 1, # 'щ' - 17: 3, # 'ъ' - 52: 1, # 'ь' - 42: 2, # 'ю' - 16: 2, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 7: { # 'р' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 3, # 'б' - 9: 3, # 'в' - 20: 3, # 'г' - 11: 3, # 'д' - 3: 3, # 'е' - 23: 3, # 'ж' - 15: 2, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 3, # 'к' - 10: 3, # 'л' - 14: 3, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 2, # 'п' - 7: 1, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 3, # 'у' - 29: 2, # 'ф' - 25: 3, # 'х' - 22: 3, # 'ц' - 21: 2, # 'ч' - 27: 3, # 'ш' - 24: 1, # 'щ' - 17: 3, # 'ъ' - 52: 1, # 'ь' - 42: 2, # 'ю' - 16: 3, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 8: { # 'с' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 2, # 'б' - 9: 3, # 'в' - 20: 2, # 'г' - 11: 2, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 1, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 3, # 'к' - 10: 3, # 'л' - 14: 3, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 3, # 'п' - 7: 3, # 'р' - 8: 1, # 'с' - 5: 3, # 'т' - 19: 3, # 'у' - 29: 2, # 'ф' - 25: 2, # 'х' - 22: 2, # 'ц' - 21: 2, # 'ч' - 27: 2, # 'ш' - 24: 0, # 'щ' - 17: 3, # 'ъ' - 52: 2, # 'ь' - 42: 2, # 'ю' - 16: 3, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 5: { # 'т' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 3, # 'б' - 9: 3, # 'в' - 20: 2, # 'г' - 11: 2, # 'д' - 3: 3, # 'е' - 23: 1, # 'ж' - 15: 1, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 3, # 'к' - 10: 3, # 'л' - 14: 2, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 2, # 'п' - 7: 3, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 3, # 'у' - 29: 1, # 'ф' - 25: 2, # 'х' - 22: 2, # 'ц' - 21: 2, # 'ч' - 27: 1, # 'ш' - 24: 1, # 'щ' - 17: 3, # 'ъ' - 52: 2, # 'ь' - 42: 2, # 'ю' - 16: 3, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 19: { # 'у' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 3, # 'б' - 9: 3, # 'в' - 20: 3, # 'г' - 11: 3, # 'д' - 3: 2, # 'е' - 23: 3, # 'ж' - 15: 3, # 'з' - 2: 2, # 'и' - 26: 2, # 'й' - 12: 3, # 'к' - 10: 3, # 'л' - 14: 3, # 'м' - 6: 3, # 'н' - 4: 2, # 'о' - 13: 3, # 'п' - 7: 3, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 1, # 'у' - 29: 2, # 'ф' - 25: 2, # 'х' - 22: 2, # 'ц' - 21: 3, # 'ч' - 27: 3, # 'ш' - 24: 2, # 'щ' - 17: 1, # 'ъ' - 52: 0, # 'ь' - 42: 1, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 29: { # 'ф' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 1, # 'б' - 9: 1, # 'в' - 20: 1, # 'г' - 11: 0, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 2, # 'к' - 10: 2, # 'л' - 14: 1, # 'м' - 6: 1, # 'н' - 4: 3, # 'о' - 13: 0, # 'п' - 7: 2, # 'р' - 8: 2, # 'с' - 5: 2, # 'т' - 19: 2, # 'у' - 29: 0, # 'ф' - 25: 1, # 'х' - 22: 0, # 'ц' - 21: 1, # 'ч' - 27: 1, # 'ш' - 24: 0, # 'щ' - 17: 2, # 'ъ' - 52: 2, # 'ь' - 42: 1, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 25: { # 'х' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 1, # 'б' - 9: 3, # 'в' - 20: 0, # 'г' - 11: 1, # 'д' - 3: 2, # 'е' - 23: 0, # 'ж' - 15: 1, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 1, # 'к' - 10: 2, # 'л' - 14: 2, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 1, # 'п' - 7: 3, # 'р' - 8: 1, # 'с' - 5: 2, # 'т' - 19: 3, # 'у' - 29: 0, # 'ф' - 25: 1, # 'х' - 22: 0, # 'ц' - 21: 1, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 2, # 'ъ' - 52: 0, # 'ь' - 42: 1, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 22: { # 'ц' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 1, # 'б' - 9: 2, # 'в' - 20: 1, # 'г' - 11: 1, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 1, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 2, # 'к' - 10: 1, # 'л' - 14: 1, # 'м' - 6: 1, # 'н' - 4: 2, # 'о' - 13: 1, # 'п' - 7: 1, # 'р' - 8: 1, # 'с' - 5: 1, # 'т' - 19: 2, # 'у' - 29: 1, # 'ф' - 25: 1, # 'х' - 22: 1, # 'ц' - 21: 1, # 'ч' - 27: 1, # 'ш' - 24: 1, # 'щ' - 17: 2, # 'ъ' - 52: 1, # 'ь' - 42: 0, # 'ю' - 16: 2, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 21: { # 'ч' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 1, # 'б' - 9: 3, # 'в' - 20: 1, # 'г' - 11: 0, # 'д' - 3: 3, # 'е' - 23: 1, # 'ж' - 15: 0, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 3, # 'к' - 10: 2, # 'л' - 14: 2, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 0, # 'п' - 7: 2, # 'р' - 8: 0, # 'с' - 5: 2, # 'т' - 19: 3, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 1, # 'ш' - 24: 0, # 'щ' - 17: 2, # 'ъ' - 52: 0, # 'ь' - 42: 1, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 27: { # 'ш' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 0, # 'б' - 9: 2, # 'в' - 20: 0, # 'г' - 11: 1, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 3, # 'к' - 10: 2, # 'л' - 14: 1, # 'м' - 6: 3, # 'н' - 4: 2, # 'о' - 13: 2, # 'п' - 7: 1, # 'р' - 8: 0, # 'с' - 5: 1, # 'т' - 19: 2, # 'у' - 29: 1, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 1, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 2, # 'ъ' - 52: 1, # 'ь' - 42: 1, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 24: { # 'щ' - 63: 1, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 3, # 'а' - 18: 0, # 'б' - 9: 1, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 3, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 3, # 'и' - 26: 0, # 'й' - 12: 1, # 'к' - 10: 0, # 'л' - 14: 0, # 'м' - 6: 2, # 'н' - 4: 3, # 'о' - 13: 0, # 'п' - 7: 1, # 'р' - 8: 0, # 'с' - 5: 2, # 'т' - 19: 3, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 1, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 1, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 2, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 17: { # 'ъ' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 1, # 'а' - 18: 3, # 'б' - 9: 3, # 'в' - 20: 3, # 'г' - 11: 3, # 'д' - 3: 2, # 'е' - 23: 3, # 'ж' - 15: 3, # 'з' - 2: 1, # 'и' - 26: 2, # 'й' - 12: 3, # 'к' - 10: 3, # 'л' - 14: 3, # 'м' - 6: 3, # 'н' - 4: 3, # 'о' - 13: 3, # 'п' - 7: 3, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 1, # 'у' - 29: 1, # 'ф' - 25: 2, # 'х' - 22: 2, # 'ц' - 21: 3, # 'ч' - 27: 2, # 'ш' - 24: 3, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 2, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 52: { # 'ь' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 0, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 1, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 0, # 'и' - 26: 0, # 'й' - 12: 1, # 'к' - 10: 0, # 'л' - 14: 0, # 'м' - 6: 1, # 'н' - 4: 3, # 'о' - 13: 0, # 'п' - 7: 0, # 'р' - 8: 0, # 'с' - 5: 1, # 'т' - 19: 0, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 1, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 1, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 42: { # 'ю' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 1, # 'а' - 18: 2, # 'б' - 9: 1, # 'в' - 20: 2, # 'г' - 11: 2, # 'д' - 3: 1, # 'е' - 23: 2, # 'ж' - 15: 2, # 'з' - 2: 1, # 'и' - 26: 1, # 'й' - 12: 2, # 'к' - 10: 2, # 'л' - 14: 2, # 'м' - 6: 2, # 'н' - 4: 1, # 'о' - 13: 1, # 'п' - 7: 2, # 'р' - 8: 2, # 'с' - 5: 2, # 'т' - 19: 1, # 'у' - 29: 1, # 'ф' - 25: 1, # 'х' - 22: 2, # 'ц' - 21: 3, # 'ч' - 27: 1, # 'ш' - 24: 1, # 'щ' - 17: 1, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 16: { # 'я' - 63: 0, # 'e' - 45: 1, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 0, # 'а' - 18: 3, # 'б' - 9: 3, # 'в' - 20: 2, # 'г' - 11: 3, # 'д' - 3: 2, # 'е' - 23: 1, # 'ж' - 15: 2, # 'з' - 2: 1, # 'и' - 26: 2, # 'й' - 12: 3, # 'к' - 10: 3, # 'л' - 14: 3, # 'м' - 6: 3, # 'н' - 4: 1, # 'о' - 13: 2, # 'п' - 7: 2, # 'р' - 8: 3, # 'с' - 5: 3, # 'т' - 19: 1, # 'у' - 29: 1, # 'ф' - 25: 3, # 'х' - 22: 2, # 'ц' - 21: 1, # 'ч' - 27: 1, # 'ш' - 24: 2, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 1, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 58: { # 'є' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 0, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 0, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 0, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 0, # 'л' - 14: 0, # 'м' - 6: 0, # 'н' - 4: 0, # 'о' - 13: 0, # 'п' - 7: 0, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 0, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, - 62: { # '№' - 63: 0, # 'e' - 45: 0, # '\xad' - 31: 0, # 'А' - 32: 0, # 'Б' - 35: 0, # 'В' - 43: 0, # 'Г' - 37: 0, # 'Д' - 44: 0, # 'Е' - 55: 0, # 'Ж' - 47: 0, # 'З' - 40: 0, # 'И' - 59: 0, # 'Й' - 33: 0, # 'К' - 46: 0, # 'Л' - 38: 0, # 'М' - 36: 0, # 'Н' - 41: 0, # 'О' - 30: 0, # 'П' - 39: 0, # 'Р' - 28: 0, # 'С' - 34: 0, # 'Т' - 51: 0, # 'У' - 48: 0, # 'Ф' - 49: 0, # 'Х' - 53: 0, # 'Ц' - 50: 0, # 'Ч' - 54: 0, # 'Ш' - 57: 0, # 'Щ' - 61: 0, # 'Ъ' - 60: 0, # 'Ю' - 56: 0, # 'Я' - 1: 0, # 'а' - 18: 0, # 'б' - 9: 0, # 'в' - 20: 0, # 'г' - 11: 0, # 'д' - 3: 0, # 'е' - 23: 0, # 'ж' - 15: 0, # 'з' - 2: 0, # 'и' - 26: 0, # 'й' - 12: 0, # 'к' - 10: 0, # 'л' - 14: 0, # 'м' - 6: 0, # 'н' - 4: 0, # 'о' - 13: 0, # 'п' - 7: 0, # 'р' - 8: 0, # 'с' - 5: 0, # 'т' - 19: 0, # 'у' - 29: 0, # 'ф' - 25: 0, # 'х' - 22: 0, # 'ц' - 21: 0, # 'ч' - 27: 0, # 'ш' - 24: 0, # 'щ' - 17: 0, # 'ъ' - 52: 0, # 'ь' - 42: 0, # 'ю' - 16: 0, # 'я' - 58: 0, # 'є' - 62: 0, # '№' - }, -} - -# 255: Undefined characters that did not exist in training text -# 254: Carriage/Return -# 253: symbol (punctuation) that does not belong to word -# 252: 0 - 9 -# 251: Control characters - -# Character Mapping Table(s): -ISO_8859_5_BULGARIAN_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 254, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 254, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 253, # ' ' - 33: 253, # '!' - 34: 253, # '"' - 35: 253, # '#' - 36: 253, # '$' - 37: 253, # '%' - 38: 253, # '&' - 39: 253, # "'" - 40: 253, # '(' - 41: 253, # ')' - 42: 253, # '*' - 43: 253, # '+' - 44: 253, # ',' - 45: 253, # '-' - 46: 253, # '.' - 47: 253, # '/' - 48: 252, # '0' - 49: 252, # '1' - 50: 252, # '2' - 51: 252, # '3' - 52: 252, # '4' - 53: 252, # '5' - 54: 252, # '6' - 55: 252, # '7' - 56: 252, # '8' - 57: 252, # '9' - 58: 253, # ':' - 59: 253, # ';' - 60: 253, # '<' - 61: 253, # '=' - 62: 253, # '>' - 63: 253, # '?' - 64: 253, # '@' - 65: 77, # 'A' - 66: 90, # 'B' - 67: 99, # 'C' - 68: 100, # 'D' - 69: 72, # 'E' - 70: 109, # 'F' - 71: 107, # 'G' - 72: 101, # 'H' - 73: 79, # 'I' - 74: 185, # 'J' - 75: 81, # 'K' - 76: 102, # 'L' - 77: 76, # 'M' - 78: 94, # 'N' - 79: 82, # 'O' - 80: 110, # 'P' - 81: 186, # 'Q' - 82: 108, # 'R' - 83: 91, # 'S' - 84: 74, # 'T' - 85: 119, # 'U' - 86: 84, # 'V' - 87: 96, # 'W' - 88: 111, # 'X' - 89: 187, # 'Y' - 90: 115, # 'Z' - 91: 253, # '[' - 92: 253, # '\\' - 93: 253, # ']' - 94: 253, # '^' - 95: 253, # '_' - 96: 253, # '`' - 97: 65, # 'a' - 98: 69, # 'b' - 99: 70, # 'c' - 100: 66, # 'd' - 101: 63, # 'e' - 102: 68, # 'f' - 103: 112, # 'g' - 104: 103, # 'h' - 105: 92, # 'i' - 106: 194, # 'j' - 107: 104, # 'k' - 108: 95, # 'l' - 109: 86, # 'm' - 110: 87, # 'n' - 111: 71, # 'o' - 112: 116, # 'p' - 113: 195, # 'q' - 114: 85, # 'r' - 115: 93, # 's' - 116: 97, # 't' - 117: 113, # 'u' - 118: 196, # 'v' - 119: 197, # 'w' - 120: 198, # 'x' - 121: 199, # 'y' - 122: 200, # 'z' - 123: 253, # '{' - 124: 253, # '|' - 125: 253, # '}' - 126: 253, # '~' - 127: 253, # '\x7f' - 128: 194, # '\x80' - 129: 195, # '\x81' - 130: 196, # '\x82' - 131: 197, # '\x83' - 132: 198, # '\x84' - 133: 199, # '\x85' - 134: 200, # '\x86' - 135: 201, # '\x87' - 136: 202, # '\x88' - 137: 203, # '\x89' - 138: 204, # '\x8a' - 139: 205, # '\x8b' - 140: 206, # '\x8c' - 141: 207, # '\x8d' - 142: 208, # '\x8e' - 143: 209, # '\x8f' - 144: 210, # '\x90' - 145: 211, # '\x91' - 146: 212, # '\x92' - 147: 213, # '\x93' - 148: 214, # '\x94' - 149: 215, # '\x95' - 150: 216, # '\x96' - 151: 217, # '\x97' - 152: 218, # '\x98' - 153: 219, # '\x99' - 154: 220, # '\x9a' - 155: 221, # '\x9b' - 156: 222, # '\x9c' - 157: 223, # '\x9d' - 158: 224, # '\x9e' - 159: 225, # '\x9f' - 160: 81, # '\xa0' - 161: 226, # 'Ё' - 162: 227, # 'Ђ' - 163: 228, # 'Ѓ' - 164: 229, # 'Є' - 165: 230, # 'Ѕ' - 166: 105, # 'І' - 167: 231, # 'Ї' - 168: 232, # 'Ј' - 169: 233, # 'Љ' - 170: 234, # 'Њ' - 171: 235, # 'Ћ' - 172: 236, # 'Ќ' - 173: 45, # '\xad' - 174: 237, # 'Ў' - 175: 238, # 'Џ' - 176: 31, # 'А' - 177: 32, # 'Б' - 178: 35, # 'В' - 179: 43, # 'Г' - 180: 37, # 'Д' - 181: 44, # 'Е' - 182: 55, # 'Ж' - 183: 47, # 'З' - 184: 40, # 'И' - 185: 59, # 'Й' - 186: 33, # 'К' - 187: 46, # 'Л' - 188: 38, # 'М' - 189: 36, # 'Н' - 190: 41, # 'О' - 191: 30, # 'П' - 192: 39, # 'Р' - 193: 28, # 'С' - 194: 34, # 'Т' - 195: 51, # 'У' - 196: 48, # 'Ф' - 197: 49, # 'Х' - 198: 53, # 'Ц' - 199: 50, # 'Ч' - 200: 54, # 'Ш' - 201: 57, # 'Щ' - 202: 61, # 'Ъ' - 203: 239, # 'Ы' - 204: 67, # 'Ь' - 205: 240, # 'Э' - 206: 60, # 'Ю' - 207: 56, # 'Я' - 208: 1, # 'а' - 209: 18, # 'б' - 210: 9, # 'в' - 211: 20, # 'г' - 212: 11, # 'д' - 213: 3, # 'е' - 214: 23, # 'ж' - 215: 15, # 'з' - 216: 2, # 'и' - 217: 26, # 'й' - 218: 12, # 'к' - 219: 10, # 'л' - 220: 14, # 'м' - 221: 6, # 'н' - 222: 4, # 'о' - 223: 13, # 'п' - 224: 7, # 'р' - 225: 8, # 'с' - 226: 5, # 'т' - 227: 19, # 'у' - 228: 29, # 'ф' - 229: 25, # 'х' - 230: 22, # 'ц' - 231: 21, # 'ч' - 232: 27, # 'ш' - 233: 24, # 'щ' - 234: 17, # 'ъ' - 235: 75, # 'ы' - 236: 52, # 'ь' - 237: 241, # 'э' - 238: 42, # 'ю' - 239: 16, # 'я' - 240: 62, # '№' - 241: 242, # 'ё' - 242: 243, # 'ђ' - 243: 244, # 'ѓ' - 244: 58, # 'є' - 245: 245, # 'ѕ' - 246: 98, # 'і' - 247: 246, # 'ї' - 248: 247, # 'ј' - 249: 248, # 'љ' - 250: 249, # 'њ' - 251: 250, # 'ћ' - 252: 251, # 'ќ' - 253: 91, # '§' - 254: 252, # 'ў' - 255: 253, # 'џ' -} - -ISO_8859_5_BULGARIAN_MODEL = SingleByteCharSetModel( - charset_name="ISO-8859-5", - language="Bulgarian", - char_to_order_map=ISO_8859_5_BULGARIAN_CHAR_TO_ORDER, - language_model=BULGARIAN_LANG_MODEL, - typical_positive_ratio=0.969392, - keep_ascii_letters=False, - alphabet="АБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЬЮЯабвгдежзийклмнопрстуфхцчшщъьюя", -) - -WINDOWS_1251_BULGARIAN_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 254, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 254, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 253, # ' ' - 33: 253, # '!' - 34: 253, # '"' - 35: 253, # '#' - 36: 253, # '$' - 37: 253, # '%' - 38: 253, # '&' - 39: 253, # "'" - 40: 253, # '(' - 41: 253, # ')' - 42: 253, # '*' - 43: 253, # '+' - 44: 253, # ',' - 45: 253, # '-' - 46: 253, # '.' - 47: 253, # '/' - 48: 252, # '0' - 49: 252, # '1' - 50: 252, # '2' - 51: 252, # '3' - 52: 252, # '4' - 53: 252, # '5' - 54: 252, # '6' - 55: 252, # '7' - 56: 252, # '8' - 57: 252, # '9' - 58: 253, # ':' - 59: 253, # ';' - 60: 253, # '<' - 61: 253, # '=' - 62: 253, # '>' - 63: 253, # '?' - 64: 253, # '@' - 65: 77, # 'A' - 66: 90, # 'B' - 67: 99, # 'C' - 68: 100, # 'D' - 69: 72, # 'E' - 70: 109, # 'F' - 71: 107, # 'G' - 72: 101, # 'H' - 73: 79, # 'I' - 74: 185, # 'J' - 75: 81, # 'K' - 76: 102, # 'L' - 77: 76, # 'M' - 78: 94, # 'N' - 79: 82, # 'O' - 80: 110, # 'P' - 81: 186, # 'Q' - 82: 108, # 'R' - 83: 91, # 'S' - 84: 74, # 'T' - 85: 119, # 'U' - 86: 84, # 'V' - 87: 96, # 'W' - 88: 111, # 'X' - 89: 187, # 'Y' - 90: 115, # 'Z' - 91: 253, # '[' - 92: 253, # '\\' - 93: 253, # ']' - 94: 253, # '^' - 95: 253, # '_' - 96: 253, # '`' - 97: 65, # 'a' - 98: 69, # 'b' - 99: 70, # 'c' - 100: 66, # 'd' - 101: 63, # 'e' - 102: 68, # 'f' - 103: 112, # 'g' - 104: 103, # 'h' - 105: 92, # 'i' - 106: 194, # 'j' - 107: 104, # 'k' - 108: 95, # 'l' - 109: 86, # 'm' - 110: 87, # 'n' - 111: 71, # 'o' - 112: 116, # 'p' - 113: 195, # 'q' - 114: 85, # 'r' - 115: 93, # 's' - 116: 97, # 't' - 117: 113, # 'u' - 118: 196, # 'v' - 119: 197, # 'w' - 120: 198, # 'x' - 121: 199, # 'y' - 122: 200, # 'z' - 123: 253, # '{' - 124: 253, # '|' - 125: 253, # '}' - 126: 253, # '~' - 127: 253, # '\x7f' - 128: 206, # 'Ђ' - 129: 207, # 'Ѓ' - 130: 208, # '‚' - 131: 209, # 'ѓ' - 132: 210, # '„' - 133: 211, # '…' - 134: 212, # '†' - 135: 213, # '‡' - 136: 120, # '€' - 137: 214, # '‰' - 138: 215, # 'Љ' - 139: 216, # '‹' - 140: 217, # 'Њ' - 141: 218, # 'Ќ' - 142: 219, # 'Ћ' - 143: 220, # 'Џ' - 144: 221, # 'ђ' - 145: 78, # '‘' - 146: 64, # '’' - 147: 83, # '“' - 148: 121, # '”' - 149: 98, # '•' - 150: 117, # '–' - 151: 105, # '—' - 152: 222, # None - 153: 223, # '™' - 154: 224, # 'љ' - 155: 225, # '›' - 156: 226, # 'њ' - 157: 227, # 'ќ' - 158: 228, # 'ћ' - 159: 229, # 'џ' - 160: 88, # '\xa0' - 161: 230, # 'Ў' - 162: 231, # 'ў' - 163: 232, # 'Ј' - 164: 233, # '¤' - 165: 122, # 'Ґ' - 166: 89, # '¦' - 167: 106, # '§' - 168: 234, # 'Ё' - 169: 235, # '©' - 170: 236, # 'Є' - 171: 237, # '«' - 172: 238, # '¬' - 173: 45, # '\xad' - 174: 239, # '®' - 175: 240, # 'Ї' - 176: 73, # '°' - 177: 80, # '±' - 178: 118, # 'І' - 179: 114, # 'і' - 180: 241, # 'ґ' - 181: 242, # 'µ' - 182: 243, # '¶' - 183: 244, # '·' - 184: 245, # 'ё' - 185: 62, # '№' - 186: 58, # 'є' - 187: 246, # '»' - 188: 247, # 'ј' - 189: 248, # 'Ѕ' - 190: 249, # 'ѕ' - 191: 250, # 'ї' - 192: 31, # 'А' - 193: 32, # 'Б' - 194: 35, # 'В' - 195: 43, # 'Г' - 196: 37, # 'Д' - 197: 44, # 'Е' - 198: 55, # 'Ж' - 199: 47, # 'З' - 200: 40, # 'И' - 201: 59, # 'Й' - 202: 33, # 'К' - 203: 46, # 'Л' - 204: 38, # 'М' - 205: 36, # 'Н' - 206: 41, # 'О' - 207: 30, # 'П' - 208: 39, # 'Р' - 209: 28, # 'С' - 210: 34, # 'Т' - 211: 51, # 'У' - 212: 48, # 'Ф' - 213: 49, # 'Х' - 214: 53, # 'Ц' - 215: 50, # 'Ч' - 216: 54, # 'Ш' - 217: 57, # 'Щ' - 218: 61, # 'Ъ' - 219: 251, # 'Ы' - 220: 67, # 'Ь' - 221: 252, # 'Э' - 222: 60, # 'Ю' - 223: 56, # 'Я' - 224: 1, # 'а' - 225: 18, # 'б' - 226: 9, # 'в' - 227: 20, # 'г' - 228: 11, # 'д' - 229: 3, # 'е' - 230: 23, # 'ж' - 231: 15, # 'з' - 232: 2, # 'и' - 233: 26, # 'й' - 234: 12, # 'к' - 235: 10, # 'л' - 236: 14, # 'м' - 237: 6, # 'н' - 238: 4, # 'о' - 239: 13, # 'п' - 240: 7, # 'р' - 241: 8, # 'с' - 242: 5, # 'т' - 243: 19, # 'у' - 244: 29, # 'ф' - 245: 25, # 'х' - 246: 22, # 'ц' - 247: 21, # 'ч' - 248: 27, # 'ш' - 249: 24, # 'щ' - 250: 17, # 'ъ' - 251: 75, # 'ы' - 252: 52, # 'ь' - 253: 253, # 'э' - 254: 42, # 'ю' - 255: 16, # 'я' -} - -WINDOWS_1251_BULGARIAN_MODEL = SingleByteCharSetModel( - charset_name="windows-1251", - language="Bulgarian", - char_to_order_map=WINDOWS_1251_BULGARIAN_CHAR_TO_ORDER, - language_model=BULGARIAN_LANG_MODEL, - typical_positive_ratio=0.969392, - keep_ascii_letters=False, - alphabet="АБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЬЮЯабвгдежзийклмнопрстуфхцчшщъьюя", -) diff --git a/venv/lib/python3.10/site-packages/chardet/langgreekmodel.py b/venv/lib/python3.10/site-packages/chardet/langgreekmodel.py deleted file mode 100644 index 0471d8bb1891415468a7394e34b2fef58ad6e0b0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/langgreekmodel.py +++ /dev/null @@ -1,4397 +0,0 @@ -from chardet.sbcharsetprober import SingleByteCharSetModel - -# 3: Positive -# 2: Likely -# 1: Unlikely -# 0: Negative - -GREEK_LANG_MODEL = { - 60: { # 'e' - 60: 2, # 'e' - 55: 1, # 'o' - 58: 2, # 't' - 36: 1, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 1, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 0, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 55: { # 'o' - 60: 0, # 'e' - 55: 2, # 'o' - 58: 2, # 't' - 36: 1, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 1, # 'ν' - 30: 0, # 'ξ' - 4: 0, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 1, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 58: { # 't' - 60: 2, # 'e' - 55: 1, # 'o' - 58: 1, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 2, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 1, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 36: { # '·' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 0, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 61: { # 'Ά' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 1, # 'γ' - 21: 2, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 2, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 0, # 'ο' - 9: 1, # 'π' - 8: 2, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 46: { # 'Έ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 2, # 'β' - 20: 2, # 'γ' - 21: 0, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 2, # 'κ' - 16: 2, # 'λ' - 10: 0, # 'μ' - 6: 3, # 'ν' - 30: 2, # 'ξ' - 4: 0, # 'ο' - 9: 2, # 'π' - 8: 2, # 'ρ' - 14: 0, # 'ς' - 7: 1, # 'σ' - 2: 2, # 'τ' - 12: 0, # 'υ' - 28: 2, # 'φ' - 23: 3, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 54: { # 'Ό' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 2, # 'λ' - 10: 2, # 'μ' - 6: 2, # 'ν' - 30: 0, # 'ξ' - 4: 0, # 'ο' - 9: 2, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 2, # 'σ' - 2: 3, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 2, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 31: { # 'Α' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 2, # 'Β' - 43: 2, # 'Γ' - 41: 1, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 2, # 'Θ' - 47: 2, # 'Ι' - 44: 2, # 'Κ' - 53: 2, # 'Λ' - 38: 2, # 'Μ' - 49: 2, # 'Ν' - 59: 1, # 'Ξ' - 39: 0, # 'Ο' - 35: 2, # 'Π' - 48: 2, # 'Ρ' - 37: 2, # 'Σ' - 33: 2, # 'Τ' - 45: 2, # 'Υ' - 56: 2, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 2, # 'γ' - 21: 0, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 1, # 'θ' - 5: 0, # 'ι' - 11: 2, # 'κ' - 16: 3, # 'λ' - 10: 2, # 'μ' - 6: 3, # 'ν' - 30: 2, # 'ξ' - 4: 0, # 'ο' - 9: 3, # 'π' - 8: 3, # 'ρ' - 14: 2, # 'ς' - 7: 2, # 'σ' - 2: 0, # 'τ' - 12: 3, # 'υ' - 28: 2, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 2, # 'ύ' - 27: 0, # 'ώ' - }, - 51: { # 'Β' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 2, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 1, # 'Ε' - 40: 1, # 'Η' - 52: 0, # 'Θ' - 47: 1, # 'Ι' - 44: 0, # 'Κ' - 53: 1, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 2, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 2, # 'ά' - 18: 2, # 'έ' - 22: 2, # 'ή' - 15: 0, # 'ί' - 1: 2, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 2, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 2, # 'ι' - 11: 0, # 'κ' - 16: 2, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 2, # 'ο' - 9: 0, # 'π' - 8: 2, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 43: { # 'Γ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 1, # 'Α' - 51: 0, # 'Β' - 43: 2, # 'Γ' - 41: 0, # 'Δ' - 34: 2, # 'Ε' - 40: 1, # 'Η' - 52: 0, # 'Θ' - 47: 2, # 'Ι' - 44: 1, # 'Κ' - 53: 1, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 1, # 'Ο' - 35: 0, # 'Π' - 48: 2, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 2, # 'Υ' - 56: 0, # 'Φ' - 50: 1, # 'Χ' - 57: 2, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 2, # 'ί' - 1: 2, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 2, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 3, # 'ι' - 11: 0, # 'κ' - 16: 2, # 'λ' - 10: 0, # 'μ' - 6: 2, # 'ν' - 30: 0, # 'ξ' - 4: 0, # 'ο' - 9: 0, # 'π' - 8: 2, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 41: { # 'Δ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 2, # 'Ε' - 40: 2, # 'Η' - 52: 0, # 'Θ' - 47: 2, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 2, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 2, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 2, # 'ή' - 15: 2, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 2, # 'η' - 25: 0, # 'θ' - 5: 3, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 2, # 'ο' - 9: 0, # 'π' - 8: 2, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 2, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 2, # 'ω' - 19: 1, # 'ό' - 26: 2, # 'ύ' - 27: 2, # 'ώ' - }, - 34: { # 'Ε' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 2, # 'Α' - 51: 0, # 'Β' - 43: 2, # 'Γ' - 41: 2, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 2, # 'Ι' - 44: 2, # 'Κ' - 53: 2, # 'Λ' - 38: 2, # 'Μ' - 49: 2, # 'Ν' - 59: 1, # 'Ξ' - 39: 0, # 'Ο' - 35: 2, # 'Π' - 48: 2, # 'Ρ' - 37: 2, # 'Σ' - 33: 2, # 'Τ' - 45: 2, # 'Υ' - 56: 0, # 'Φ' - 50: 2, # 'Χ' - 57: 2, # 'Ω' - 17: 3, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 3, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 3, # 'γ' - 21: 2, # 'δ' - 3: 1, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 1, # 'θ' - 5: 2, # 'ι' - 11: 3, # 'κ' - 16: 3, # 'λ' - 10: 2, # 'μ' - 6: 3, # 'ν' - 30: 2, # 'ξ' - 4: 0, # 'ο' - 9: 3, # 'π' - 8: 2, # 'ρ' - 14: 0, # 'ς' - 7: 2, # 'σ' - 2: 2, # 'τ' - 12: 2, # 'υ' - 28: 2, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 1, # 'ύ' - 27: 0, # 'ώ' - }, - 40: { # 'Η' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 1, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 2, # 'Θ' - 47: 0, # 'Ι' - 44: 2, # 'Κ' - 53: 0, # 'Λ' - 38: 2, # 'Μ' - 49: 2, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 2, # 'Π' - 48: 2, # 'Ρ' - 37: 2, # 'Σ' - 33: 2, # 'Τ' - 45: 1, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 2, # 'λ' - 10: 0, # 'μ' - 6: 1, # 'ν' - 30: 0, # 'ξ' - 4: 0, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 1, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 52: { # 'Θ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 2, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 2, # 'Ε' - 40: 2, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 2, # 'Ο' - 35: 0, # 'Π' - 48: 1, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 1, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 2, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 2, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 0, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 2, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 2, # 'ύ' - 27: 0, # 'ώ' - }, - 47: { # 'Ι' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 2, # 'Α' - 51: 1, # 'Β' - 43: 1, # 'Γ' - 41: 2, # 'Δ' - 34: 2, # 'Ε' - 40: 2, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 2, # 'Κ' - 53: 2, # 'Λ' - 38: 2, # 'Μ' - 49: 2, # 'Ν' - 59: 0, # 'Ξ' - 39: 2, # 'Ο' - 35: 0, # 'Π' - 48: 2, # 'Ρ' - 37: 2, # 'Σ' - 33: 2, # 'Τ' - 45: 0, # 'Υ' - 56: 2, # 'Φ' - 50: 0, # 'Χ' - 57: 2, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 2, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 2, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 1, # 'ν' - 30: 0, # 'ξ' - 4: 2, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 2, # 'σ' - 2: 1, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 1, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 44: { # 'Κ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 2, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 1, # 'Δ' - 34: 2, # 'Ε' - 40: 2, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 1, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 2, # 'Ο' - 35: 0, # 'Π' - 48: 2, # 'Ρ' - 37: 0, # 'Σ' - 33: 1, # 'Τ' - 45: 2, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 1, # 'Ω' - 17: 3, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 2, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 2, # 'ι' - 11: 0, # 'κ' - 16: 2, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 2, # 'ο' - 9: 0, # 'π' - 8: 2, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 2, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 2, # 'ό' - 26: 2, # 'ύ' - 27: 2, # 'ώ' - }, - 53: { # 'Λ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 2, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 2, # 'Ε' - 40: 2, # 'Η' - 52: 0, # 'Θ' - 47: 2, # 'Ι' - 44: 0, # 'Κ' - 53: 2, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 2, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 2, # 'Σ' - 33: 0, # 'Τ' - 45: 2, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 2, # 'Ω' - 17: 2, # 'ά' - 18: 2, # 'έ' - 22: 0, # 'ή' - 15: 2, # 'ί' - 1: 2, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 2, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 1, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 2, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 2, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 2, # 'ό' - 26: 2, # 'ύ' - 27: 0, # 'ώ' - }, - 38: { # 'Μ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 2, # 'Α' - 51: 2, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 2, # 'Ε' - 40: 2, # 'Η' - 52: 0, # 'Θ' - 47: 2, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 2, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 2, # 'Ο' - 35: 2, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 2, # 'ά' - 18: 2, # 'έ' - 22: 2, # 'ή' - 15: 2, # 'ί' - 1: 2, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 2, # 'η' - 25: 0, # 'θ' - 5: 3, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 2, # 'ο' - 9: 3, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 2, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 2, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 49: { # 'Ν' - 60: 2, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 2, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 2, # 'Ε' - 40: 2, # 'Η' - 52: 0, # 'Θ' - 47: 2, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 2, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 2, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 2, # 'Ω' - 17: 0, # 'ά' - 18: 2, # 'έ' - 22: 0, # 'ή' - 15: 2, # 'ί' - 1: 2, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 1, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 2, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 1, # 'ω' - 19: 2, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 59: { # 'Ξ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 1, # 'Ε' - 40: 1, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 1, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 2, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 2, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 2, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 0, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 39: { # 'Ο' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 1, # 'Β' - 43: 2, # 'Γ' - 41: 2, # 'Δ' - 34: 2, # 'Ε' - 40: 1, # 'Η' - 52: 2, # 'Θ' - 47: 2, # 'Ι' - 44: 2, # 'Κ' - 53: 2, # 'Λ' - 38: 2, # 'Μ' - 49: 2, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 2, # 'Π' - 48: 2, # 'Ρ' - 37: 2, # 'Σ' - 33: 2, # 'Τ' - 45: 2, # 'Υ' - 56: 2, # 'Φ' - 50: 2, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 2, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 3, # 'ι' - 11: 2, # 'κ' - 16: 2, # 'λ' - 10: 2, # 'μ' - 6: 2, # 'ν' - 30: 0, # 'ξ' - 4: 0, # 'ο' - 9: 2, # 'π' - 8: 2, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 2, # 'τ' - 12: 2, # 'υ' - 28: 1, # 'φ' - 23: 1, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 2, # 'ύ' - 27: 0, # 'ώ' - }, - 35: { # 'Π' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 2, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 2, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 2, # 'Ι' - 44: 0, # 'Κ' - 53: 2, # 'Λ' - 38: 1, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 2, # 'Ο' - 35: 0, # 'Π' - 48: 2, # 'Ρ' - 37: 0, # 'Σ' - 33: 1, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 1, # 'Χ' - 57: 2, # 'Ω' - 17: 2, # 'ά' - 18: 1, # 'έ' - 22: 1, # 'ή' - 15: 2, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 2, # 'η' - 25: 0, # 'θ' - 5: 2, # 'ι' - 11: 0, # 'κ' - 16: 2, # 'λ' - 10: 0, # 'μ' - 6: 2, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 0, # 'π' - 8: 3, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 2, # 'υ' - 28: 0, # 'φ' - 23: 2, # 'χ' - 42: 0, # 'ψ' - 24: 2, # 'ω' - 19: 2, # 'ό' - 26: 0, # 'ύ' - 27: 3, # 'ώ' - }, - 48: { # 'Ρ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 2, # 'Α' - 51: 0, # 'Β' - 43: 1, # 'Γ' - 41: 1, # 'Δ' - 34: 2, # 'Ε' - 40: 2, # 'Η' - 52: 0, # 'Θ' - 47: 2, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 2, # 'Ν' - 59: 0, # 'Ξ' - 39: 2, # 'Ο' - 35: 0, # 'Π' - 48: 2, # 'Ρ' - 37: 0, # 'Σ' - 33: 1, # 'Τ' - 45: 1, # 'Υ' - 56: 0, # 'Φ' - 50: 1, # 'Χ' - 57: 1, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 2, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 1, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 3, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 2, # 'ω' - 19: 0, # 'ό' - 26: 2, # 'ύ' - 27: 0, # 'ώ' - }, - 37: { # 'Σ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 2, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 1, # 'Δ' - 34: 2, # 'Ε' - 40: 2, # 'Η' - 52: 0, # 'Θ' - 47: 2, # 'Ι' - 44: 2, # 'Κ' - 53: 0, # 'Λ' - 38: 2, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 2, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 2, # 'Σ' - 33: 2, # 'Τ' - 45: 2, # 'Υ' - 56: 0, # 'Φ' - 50: 2, # 'Χ' - 57: 2, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 2, # 'ή' - 15: 2, # 'ί' - 1: 2, # 'α' - 29: 2, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 3, # 'η' - 25: 0, # 'θ' - 5: 2, # 'ι' - 11: 2, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 2, # 'ο' - 9: 2, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 3, # 'τ' - 12: 3, # 'υ' - 28: 0, # 'φ' - 23: 2, # 'χ' - 42: 0, # 'ψ' - 24: 2, # 'ω' - 19: 0, # 'ό' - 26: 2, # 'ύ' - 27: 2, # 'ώ' - }, - 33: { # 'Τ' - 60: 0, # 'e' - 55: 1, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 2, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 2, # 'Ε' - 40: 2, # 'Η' - 52: 0, # 'Θ' - 47: 2, # 'Ι' - 44: 2, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 2, # 'Ο' - 35: 0, # 'Π' - 48: 2, # 'Ρ' - 37: 0, # 'Σ' - 33: 1, # 'Τ' - 45: 1, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 2, # 'Ω' - 17: 2, # 'ά' - 18: 2, # 'έ' - 22: 0, # 'ή' - 15: 2, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 2, # 'ε' - 32: 0, # 'ζ' - 13: 2, # 'η' - 25: 0, # 'θ' - 5: 2, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 2, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 0, # 'π' - 8: 2, # 'ρ' - 14: 0, # 'ς' - 7: 2, # 'σ' - 2: 0, # 'τ' - 12: 2, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 2, # 'ό' - 26: 2, # 'ύ' - 27: 3, # 'ώ' - }, - 45: { # 'Υ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 2, # 'Γ' - 41: 0, # 'Δ' - 34: 1, # 'Ε' - 40: 2, # 'Η' - 52: 2, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 1, # 'Λ' - 38: 2, # 'Μ' - 49: 2, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 2, # 'Π' - 48: 1, # 'Ρ' - 37: 2, # 'Σ' - 33: 2, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 1, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 2, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 0, # 'ο' - 9: 3, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 56: { # 'Φ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 1, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 1, # 'Η' - 52: 0, # 'Θ' - 47: 2, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 2, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 2, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 2, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 2, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 2, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 2, # 'τ' - 12: 2, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 1, # 'ύ' - 27: 1, # 'ώ' - }, - 50: { # 'Χ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 1, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 2, # 'Ε' - 40: 2, # 'Η' - 52: 0, # 'Θ' - 47: 2, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 1, # 'Ν' - 59: 0, # 'Ξ' - 39: 1, # 'Ο' - 35: 0, # 'Π' - 48: 2, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 1, # 'Χ' - 57: 1, # 'Ω' - 17: 2, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 2, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 2, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 2, # 'ο' - 9: 0, # 'π' - 8: 3, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 2, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 2, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 57: { # 'Ω' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 1, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 1, # 'Λ' - 38: 0, # 'Μ' - 49: 2, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 2, # 'Ρ' - 37: 2, # 'Σ' - 33: 2, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 0, # 'ο' - 9: 0, # 'π' - 8: 2, # 'ρ' - 14: 2, # 'ς' - 7: 2, # 'σ' - 2: 0, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 1, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 17: { # 'ά' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 2, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 3, # 'β' - 20: 3, # 'γ' - 21: 3, # 'δ' - 3: 3, # 'ε' - 32: 3, # 'ζ' - 13: 0, # 'η' - 25: 3, # 'θ' - 5: 2, # 'ι' - 11: 3, # 'κ' - 16: 3, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 3, # 'ξ' - 4: 0, # 'ο' - 9: 3, # 'π' - 8: 3, # 'ρ' - 14: 3, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 0, # 'υ' - 28: 3, # 'φ' - 23: 3, # 'χ' - 42: 3, # 'ψ' - 24: 2, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 18: { # 'έ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 3, # 'α' - 29: 2, # 'β' - 20: 3, # 'γ' - 21: 2, # 'δ' - 3: 3, # 'ε' - 32: 2, # 'ζ' - 13: 0, # 'η' - 25: 3, # 'θ' - 5: 0, # 'ι' - 11: 3, # 'κ' - 16: 3, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 3, # 'ξ' - 4: 3, # 'ο' - 9: 3, # 'π' - 8: 3, # 'ρ' - 14: 3, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 0, # 'υ' - 28: 3, # 'φ' - 23: 3, # 'χ' - 42: 3, # 'ψ' - 24: 2, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 22: { # 'ή' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 1, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 3, # 'γ' - 21: 3, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 3, # 'θ' - 5: 0, # 'ι' - 11: 3, # 'κ' - 16: 2, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 2, # 'ξ' - 4: 0, # 'ο' - 9: 3, # 'π' - 8: 3, # 'ρ' - 14: 3, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 0, # 'υ' - 28: 2, # 'φ' - 23: 3, # 'χ' - 42: 2, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 15: { # 'ί' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 3, # 'α' - 29: 2, # 'β' - 20: 3, # 'γ' - 21: 3, # 'δ' - 3: 3, # 'ε' - 32: 3, # 'ζ' - 13: 3, # 'η' - 25: 3, # 'θ' - 5: 0, # 'ι' - 11: 3, # 'κ' - 16: 3, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 3, # 'ξ' - 4: 3, # 'ο' - 9: 3, # 'π' - 8: 3, # 'ρ' - 14: 3, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 0, # 'υ' - 28: 1, # 'φ' - 23: 3, # 'χ' - 42: 2, # 'ψ' - 24: 3, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 1: { # 'α' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 2, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 2, # 'έ' - 22: 0, # 'ή' - 15: 3, # 'ί' - 1: 0, # 'α' - 29: 3, # 'β' - 20: 3, # 'γ' - 21: 3, # 'δ' - 3: 2, # 'ε' - 32: 3, # 'ζ' - 13: 1, # 'η' - 25: 3, # 'θ' - 5: 3, # 'ι' - 11: 3, # 'κ' - 16: 3, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 3, # 'ξ' - 4: 2, # 'ο' - 9: 3, # 'π' - 8: 3, # 'ρ' - 14: 3, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 3, # 'υ' - 28: 3, # 'φ' - 23: 3, # 'χ' - 42: 2, # 'ψ' - 24: 0, # 'ω' - 19: 2, # 'ό' - 26: 2, # 'ύ' - 27: 0, # 'ώ' - }, - 29: { # 'β' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 3, # 'ά' - 18: 2, # 'έ' - 22: 3, # 'ή' - 15: 2, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 2, # 'γ' - 21: 2, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 2, # 'η' - 25: 0, # 'θ' - 5: 3, # 'ι' - 11: 0, # 'κ' - 16: 3, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 0, # 'π' - 8: 3, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 2, # 'ω' - 19: 2, # 'ό' - 26: 2, # 'ύ' - 27: 2, # 'ώ' - }, - 20: { # 'γ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 3, # 'ά' - 18: 3, # 'έ' - 22: 3, # 'ή' - 15: 3, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 3, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 3, # 'η' - 25: 0, # 'θ' - 5: 3, # 'ι' - 11: 3, # 'κ' - 16: 3, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 3, # 'ξ' - 4: 3, # 'ο' - 9: 0, # 'π' - 8: 3, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 2, # 'υ' - 28: 0, # 'φ' - 23: 3, # 'χ' - 42: 0, # 'ψ' - 24: 3, # 'ω' - 19: 3, # 'ό' - 26: 2, # 'ύ' - 27: 3, # 'ώ' - }, - 21: { # 'δ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 2, # 'ά' - 18: 3, # 'έ' - 22: 3, # 'ή' - 15: 3, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 3, # 'η' - 25: 0, # 'θ' - 5: 3, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 0, # 'π' - 8: 3, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 3, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 3, # 'ω' - 19: 3, # 'ό' - 26: 3, # 'ύ' - 27: 3, # 'ώ' - }, - 3: { # 'ε' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 2, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 3, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 3, # 'ί' - 1: 2, # 'α' - 29: 3, # 'β' - 20: 3, # 'γ' - 21: 3, # 'δ' - 3: 2, # 'ε' - 32: 2, # 'ζ' - 13: 0, # 'η' - 25: 3, # 'θ' - 5: 3, # 'ι' - 11: 3, # 'κ' - 16: 3, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 3, # 'ξ' - 4: 2, # 'ο' - 9: 3, # 'π' - 8: 3, # 'ρ' - 14: 3, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 3, # 'υ' - 28: 3, # 'φ' - 23: 3, # 'χ' - 42: 2, # 'ψ' - 24: 3, # 'ω' - 19: 2, # 'ό' - 26: 3, # 'ύ' - 27: 2, # 'ώ' - }, - 32: { # 'ζ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 2, # 'ά' - 18: 2, # 'έ' - 22: 2, # 'ή' - 15: 2, # 'ί' - 1: 2, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 3, # 'η' - 25: 0, # 'θ' - 5: 2, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 1, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 3, # 'ω' - 19: 2, # 'ό' - 26: 0, # 'ύ' - 27: 2, # 'ώ' - }, - 13: { # 'η' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 2, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 3, # 'γ' - 21: 2, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 3, # 'θ' - 5: 0, # 'ι' - 11: 3, # 'κ' - 16: 3, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 2, # 'ξ' - 4: 0, # 'ο' - 9: 2, # 'π' - 8: 3, # 'ρ' - 14: 3, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 0, # 'υ' - 28: 2, # 'φ' - 23: 3, # 'χ' - 42: 2, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 25: { # 'θ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 2, # 'ά' - 18: 3, # 'έ' - 22: 3, # 'ή' - 15: 2, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 3, # 'η' - 25: 0, # 'θ' - 5: 3, # 'ι' - 11: 0, # 'κ' - 16: 1, # 'λ' - 10: 3, # 'μ' - 6: 2, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 0, # 'π' - 8: 3, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 3, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 3, # 'ω' - 19: 3, # 'ό' - 26: 3, # 'ύ' - 27: 3, # 'ώ' - }, - 5: { # 'ι' - 60: 0, # 'e' - 55: 1, # 'o' - 58: 0, # 't' - 36: 2, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 1, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 3, # 'ά' - 18: 3, # 'έ' - 22: 3, # 'ή' - 15: 0, # 'ί' - 1: 3, # 'α' - 29: 3, # 'β' - 20: 3, # 'γ' - 21: 3, # 'δ' - 3: 3, # 'ε' - 32: 2, # 'ζ' - 13: 3, # 'η' - 25: 3, # 'θ' - 5: 0, # 'ι' - 11: 3, # 'κ' - 16: 3, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 3, # 'ξ' - 4: 3, # 'ο' - 9: 3, # 'π' - 8: 3, # 'ρ' - 14: 3, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 0, # 'υ' - 28: 2, # 'φ' - 23: 3, # 'χ' - 42: 2, # 'ψ' - 24: 3, # 'ω' - 19: 3, # 'ό' - 26: 0, # 'ύ' - 27: 3, # 'ώ' - }, - 11: { # 'κ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 3, # 'ά' - 18: 3, # 'έ' - 22: 3, # 'ή' - 15: 3, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 3, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 3, # 'η' - 25: 2, # 'θ' - 5: 3, # 'ι' - 11: 3, # 'κ' - 16: 3, # 'λ' - 10: 3, # 'μ' - 6: 2, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 2, # 'π' - 8: 3, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 3, # 'τ' - 12: 3, # 'υ' - 28: 2, # 'φ' - 23: 2, # 'χ' - 42: 0, # 'ψ' - 24: 3, # 'ω' - 19: 3, # 'ό' - 26: 3, # 'ύ' - 27: 3, # 'ώ' - }, - 16: { # 'λ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 3, # 'ά' - 18: 3, # 'έ' - 22: 3, # 'ή' - 15: 3, # 'ί' - 1: 3, # 'α' - 29: 1, # 'β' - 20: 2, # 'γ' - 21: 1, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 3, # 'η' - 25: 2, # 'θ' - 5: 3, # 'ι' - 11: 2, # 'κ' - 16: 3, # 'λ' - 10: 2, # 'μ' - 6: 2, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 3, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 3, # 'τ' - 12: 3, # 'υ' - 28: 2, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 3, # 'ω' - 19: 3, # 'ό' - 26: 3, # 'ύ' - 27: 3, # 'ώ' - }, - 10: { # 'μ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 1, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 3, # 'ά' - 18: 3, # 'έ' - 22: 3, # 'ή' - 15: 3, # 'ί' - 1: 3, # 'α' - 29: 3, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 3, # 'η' - 25: 0, # 'θ' - 5: 3, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 3, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 2, # 'υ' - 28: 3, # 'φ' - 23: 0, # 'χ' - 42: 2, # 'ψ' - 24: 3, # 'ω' - 19: 3, # 'ό' - 26: 2, # 'ύ' - 27: 2, # 'ώ' - }, - 6: { # 'ν' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 2, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 3, # 'ά' - 18: 3, # 'έ' - 22: 3, # 'ή' - 15: 3, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 3, # 'δ' - 3: 3, # 'ε' - 32: 2, # 'ζ' - 13: 3, # 'η' - 25: 3, # 'θ' - 5: 3, # 'ι' - 11: 0, # 'κ' - 16: 1, # 'λ' - 10: 0, # 'μ' - 6: 2, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 3, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 3, # 'ω' - 19: 3, # 'ό' - 26: 3, # 'ύ' - 27: 3, # 'ώ' - }, - 30: { # 'ξ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 2, # 'ά' - 18: 3, # 'έ' - 22: 3, # 'ή' - 15: 2, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 3, # 'η' - 25: 0, # 'θ' - 5: 2, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 3, # 'τ' - 12: 2, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 3, # 'ω' - 19: 2, # 'ό' - 26: 3, # 'ύ' - 27: 1, # 'ώ' - }, - 4: { # 'ο' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 2, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 2, # 'έ' - 22: 3, # 'ή' - 15: 3, # 'ί' - 1: 2, # 'α' - 29: 3, # 'β' - 20: 3, # 'γ' - 21: 3, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 3, # 'η' - 25: 3, # 'θ' - 5: 3, # 'ι' - 11: 3, # 'κ' - 16: 3, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 2, # 'ξ' - 4: 2, # 'ο' - 9: 3, # 'π' - 8: 3, # 'ρ' - 14: 3, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 3, # 'υ' - 28: 3, # 'φ' - 23: 3, # 'χ' - 42: 2, # 'ψ' - 24: 2, # 'ω' - 19: 1, # 'ό' - 26: 3, # 'ύ' - 27: 2, # 'ώ' - }, - 9: { # 'π' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 3, # 'ά' - 18: 3, # 'έ' - 22: 3, # 'ή' - 15: 3, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 3, # 'η' - 25: 0, # 'θ' - 5: 3, # 'ι' - 11: 0, # 'κ' - 16: 3, # 'λ' - 10: 0, # 'μ' - 6: 2, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 0, # 'π' - 8: 3, # 'ρ' - 14: 2, # 'ς' - 7: 0, # 'σ' - 2: 3, # 'τ' - 12: 3, # 'υ' - 28: 0, # 'φ' - 23: 2, # 'χ' - 42: 0, # 'ψ' - 24: 3, # 'ω' - 19: 3, # 'ό' - 26: 2, # 'ύ' - 27: 3, # 'ώ' - }, - 8: { # 'ρ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 3, # 'ά' - 18: 3, # 'έ' - 22: 3, # 'ή' - 15: 3, # 'ί' - 1: 3, # 'α' - 29: 2, # 'β' - 20: 3, # 'γ' - 21: 2, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 3, # 'η' - 25: 3, # 'θ' - 5: 3, # 'ι' - 11: 3, # 'κ' - 16: 1, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 2, # 'ξ' - 4: 3, # 'ο' - 9: 2, # 'π' - 8: 2, # 'ρ' - 14: 0, # 'ς' - 7: 2, # 'σ' - 2: 3, # 'τ' - 12: 3, # 'υ' - 28: 3, # 'φ' - 23: 3, # 'χ' - 42: 0, # 'ψ' - 24: 3, # 'ω' - 19: 3, # 'ό' - 26: 3, # 'ύ' - 27: 3, # 'ώ' - }, - 14: { # 'ς' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 2, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 0, # 'θ' - 5: 0, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 0, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 0, # 'τ' - 12: 0, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 7: { # 'σ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 2, # 'ά' - 18: 2, # 'έ' - 22: 3, # 'ή' - 15: 3, # 'ί' - 1: 3, # 'α' - 29: 3, # 'β' - 20: 0, # 'γ' - 21: 2, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 3, # 'η' - 25: 3, # 'θ' - 5: 3, # 'ι' - 11: 3, # 'κ' - 16: 2, # 'λ' - 10: 3, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 3, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 3, # 'υ' - 28: 3, # 'φ' - 23: 3, # 'χ' - 42: 0, # 'ψ' - 24: 3, # 'ω' - 19: 3, # 'ό' - 26: 3, # 'ύ' - 27: 2, # 'ώ' - }, - 2: { # 'τ' - 60: 0, # 'e' - 55: 2, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 3, # 'ά' - 18: 3, # 'έ' - 22: 3, # 'ή' - 15: 3, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 2, # 'ζ' - 13: 3, # 'η' - 25: 0, # 'θ' - 5: 3, # 'ι' - 11: 2, # 'κ' - 16: 2, # 'λ' - 10: 3, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 0, # 'π' - 8: 3, # 'ρ' - 14: 0, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 3, # 'υ' - 28: 2, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 3, # 'ω' - 19: 3, # 'ό' - 26: 3, # 'ύ' - 27: 3, # 'ώ' - }, - 12: { # 'υ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 2, # 'ά' - 18: 2, # 'έ' - 22: 3, # 'ή' - 15: 2, # 'ί' - 1: 3, # 'α' - 29: 2, # 'β' - 20: 3, # 'γ' - 21: 2, # 'δ' - 3: 2, # 'ε' - 32: 2, # 'ζ' - 13: 2, # 'η' - 25: 3, # 'θ' - 5: 2, # 'ι' - 11: 3, # 'κ' - 16: 3, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 3, # 'ξ' - 4: 3, # 'ο' - 9: 3, # 'π' - 8: 3, # 'ρ' - 14: 3, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 0, # 'υ' - 28: 2, # 'φ' - 23: 3, # 'χ' - 42: 2, # 'ψ' - 24: 2, # 'ω' - 19: 2, # 'ό' - 26: 0, # 'ύ' - 27: 2, # 'ώ' - }, - 28: { # 'φ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 3, # 'ά' - 18: 3, # 'έ' - 22: 3, # 'ή' - 15: 3, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 2, # 'η' - 25: 2, # 'θ' - 5: 3, # 'ι' - 11: 0, # 'κ' - 16: 2, # 'λ' - 10: 0, # 'μ' - 6: 1, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 0, # 'π' - 8: 3, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 3, # 'τ' - 12: 3, # 'υ' - 28: 1, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 3, # 'ω' - 19: 3, # 'ό' - 26: 2, # 'ύ' - 27: 2, # 'ώ' - }, - 23: { # 'χ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 3, # 'ά' - 18: 2, # 'έ' - 22: 3, # 'ή' - 15: 3, # 'ί' - 1: 3, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 2, # 'η' - 25: 2, # 'θ' - 5: 3, # 'ι' - 11: 0, # 'κ' - 16: 2, # 'λ' - 10: 2, # 'μ' - 6: 3, # 'ν' - 30: 0, # 'ξ' - 4: 3, # 'ο' - 9: 0, # 'π' - 8: 3, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 3, # 'τ' - 12: 3, # 'υ' - 28: 0, # 'φ' - 23: 2, # 'χ' - 42: 0, # 'ψ' - 24: 3, # 'ω' - 19: 3, # 'ό' - 26: 3, # 'ύ' - 27: 3, # 'ώ' - }, - 42: { # 'ψ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 2, # 'ά' - 18: 2, # 'έ' - 22: 1, # 'ή' - 15: 2, # 'ί' - 1: 2, # 'α' - 29: 0, # 'β' - 20: 0, # 'γ' - 21: 0, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 3, # 'η' - 25: 0, # 'θ' - 5: 2, # 'ι' - 11: 0, # 'κ' - 16: 0, # 'λ' - 10: 0, # 'μ' - 6: 0, # 'ν' - 30: 0, # 'ξ' - 4: 2, # 'ο' - 9: 0, # 'π' - 8: 0, # 'ρ' - 14: 0, # 'ς' - 7: 0, # 'σ' - 2: 2, # 'τ' - 12: 1, # 'υ' - 28: 0, # 'φ' - 23: 0, # 'χ' - 42: 0, # 'ψ' - 24: 2, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 24: { # 'ω' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 1, # 'ά' - 18: 0, # 'έ' - 22: 2, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 2, # 'β' - 20: 3, # 'γ' - 21: 2, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 0, # 'η' - 25: 3, # 'θ' - 5: 2, # 'ι' - 11: 0, # 'κ' - 16: 2, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 0, # 'ξ' - 4: 0, # 'ο' - 9: 3, # 'π' - 8: 3, # 'ρ' - 14: 3, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 0, # 'υ' - 28: 2, # 'φ' - 23: 2, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 19: { # 'ό' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 3, # 'β' - 20: 3, # 'γ' - 21: 3, # 'δ' - 3: 1, # 'ε' - 32: 2, # 'ζ' - 13: 2, # 'η' - 25: 2, # 'θ' - 5: 2, # 'ι' - 11: 3, # 'κ' - 16: 3, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 1, # 'ξ' - 4: 2, # 'ο' - 9: 3, # 'π' - 8: 3, # 'ρ' - 14: 3, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 0, # 'υ' - 28: 2, # 'φ' - 23: 3, # 'χ' - 42: 2, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 26: { # 'ύ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 2, # 'α' - 29: 2, # 'β' - 20: 2, # 'γ' - 21: 1, # 'δ' - 3: 3, # 'ε' - 32: 0, # 'ζ' - 13: 2, # 'η' - 25: 3, # 'θ' - 5: 0, # 'ι' - 11: 3, # 'κ' - 16: 3, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 2, # 'ξ' - 4: 3, # 'ο' - 9: 3, # 'π' - 8: 3, # 'ρ' - 14: 3, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 0, # 'υ' - 28: 2, # 'φ' - 23: 2, # 'χ' - 42: 2, # 'ψ' - 24: 2, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, - 27: { # 'ώ' - 60: 0, # 'e' - 55: 0, # 'o' - 58: 0, # 't' - 36: 0, # '·' - 61: 0, # 'Ά' - 46: 0, # 'Έ' - 54: 0, # 'Ό' - 31: 0, # 'Α' - 51: 0, # 'Β' - 43: 0, # 'Γ' - 41: 0, # 'Δ' - 34: 0, # 'Ε' - 40: 0, # 'Η' - 52: 0, # 'Θ' - 47: 0, # 'Ι' - 44: 0, # 'Κ' - 53: 0, # 'Λ' - 38: 0, # 'Μ' - 49: 0, # 'Ν' - 59: 0, # 'Ξ' - 39: 0, # 'Ο' - 35: 0, # 'Π' - 48: 0, # 'Ρ' - 37: 0, # 'Σ' - 33: 0, # 'Τ' - 45: 0, # 'Υ' - 56: 0, # 'Φ' - 50: 0, # 'Χ' - 57: 0, # 'Ω' - 17: 0, # 'ά' - 18: 0, # 'έ' - 22: 0, # 'ή' - 15: 0, # 'ί' - 1: 0, # 'α' - 29: 1, # 'β' - 20: 0, # 'γ' - 21: 3, # 'δ' - 3: 0, # 'ε' - 32: 0, # 'ζ' - 13: 1, # 'η' - 25: 2, # 'θ' - 5: 2, # 'ι' - 11: 0, # 'κ' - 16: 2, # 'λ' - 10: 3, # 'μ' - 6: 3, # 'ν' - 30: 1, # 'ξ' - 4: 0, # 'ο' - 9: 2, # 'π' - 8: 3, # 'ρ' - 14: 3, # 'ς' - 7: 3, # 'σ' - 2: 3, # 'τ' - 12: 0, # 'υ' - 28: 1, # 'φ' - 23: 1, # 'χ' - 42: 0, # 'ψ' - 24: 0, # 'ω' - 19: 0, # 'ό' - 26: 0, # 'ύ' - 27: 0, # 'ώ' - }, -} - -# 255: Undefined characters that did not exist in training text -# 254: Carriage/Return -# 253: symbol (punctuation) that does not belong to word -# 252: 0 - 9 -# 251: Control characters - -# Character Mapping Table(s): -WINDOWS_1253_GREEK_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 254, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 254, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 253, # ' ' - 33: 253, # '!' - 34: 253, # '"' - 35: 253, # '#' - 36: 253, # '$' - 37: 253, # '%' - 38: 253, # '&' - 39: 253, # "'" - 40: 253, # '(' - 41: 253, # ')' - 42: 253, # '*' - 43: 253, # '+' - 44: 253, # ',' - 45: 253, # '-' - 46: 253, # '.' - 47: 253, # '/' - 48: 252, # '0' - 49: 252, # '1' - 50: 252, # '2' - 51: 252, # '3' - 52: 252, # '4' - 53: 252, # '5' - 54: 252, # '6' - 55: 252, # '7' - 56: 252, # '8' - 57: 252, # '9' - 58: 253, # ':' - 59: 253, # ';' - 60: 253, # '<' - 61: 253, # '=' - 62: 253, # '>' - 63: 253, # '?' - 64: 253, # '@' - 65: 82, # 'A' - 66: 100, # 'B' - 67: 104, # 'C' - 68: 94, # 'D' - 69: 98, # 'E' - 70: 101, # 'F' - 71: 116, # 'G' - 72: 102, # 'H' - 73: 111, # 'I' - 74: 187, # 'J' - 75: 117, # 'K' - 76: 92, # 'L' - 77: 88, # 'M' - 78: 113, # 'N' - 79: 85, # 'O' - 80: 79, # 'P' - 81: 118, # 'Q' - 82: 105, # 'R' - 83: 83, # 'S' - 84: 67, # 'T' - 85: 114, # 'U' - 86: 119, # 'V' - 87: 95, # 'W' - 88: 99, # 'X' - 89: 109, # 'Y' - 90: 188, # 'Z' - 91: 253, # '[' - 92: 253, # '\\' - 93: 253, # ']' - 94: 253, # '^' - 95: 253, # '_' - 96: 253, # '`' - 97: 72, # 'a' - 98: 70, # 'b' - 99: 80, # 'c' - 100: 81, # 'd' - 101: 60, # 'e' - 102: 96, # 'f' - 103: 93, # 'g' - 104: 89, # 'h' - 105: 68, # 'i' - 106: 120, # 'j' - 107: 97, # 'k' - 108: 77, # 'l' - 109: 86, # 'm' - 110: 69, # 'n' - 111: 55, # 'o' - 112: 78, # 'p' - 113: 115, # 'q' - 114: 65, # 'r' - 115: 66, # 's' - 116: 58, # 't' - 117: 76, # 'u' - 118: 106, # 'v' - 119: 103, # 'w' - 120: 87, # 'x' - 121: 107, # 'y' - 122: 112, # 'z' - 123: 253, # '{' - 124: 253, # '|' - 125: 253, # '}' - 126: 253, # '~' - 127: 253, # '\x7f' - 128: 255, # '€' - 129: 255, # None - 130: 255, # '‚' - 131: 255, # 'ƒ' - 132: 255, # '„' - 133: 255, # '…' - 134: 255, # '†' - 135: 255, # '‡' - 136: 255, # None - 137: 255, # '‰' - 138: 255, # None - 139: 255, # '‹' - 140: 255, # None - 141: 255, # None - 142: 255, # None - 143: 255, # None - 144: 255, # None - 145: 255, # '‘' - 146: 255, # '’' - 147: 255, # '“' - 148: 255, # '”' - 149: 255, # '•' - 150: 255, # '–' - 151: 255, # '—' - 152: 255, # None - 153: 255, # '™' - 154: 255, # None - 155: 255, # '›' - 156: 255, # None - 157: 255, # None - 158: 255, # None - 159: 255, # None - 160: 253, # '\xa0' - 161: 233, # '΅' - 162: 61, # 'Ά' - 163: 253, # '£' - 164: 253, # '¤' - 165: 253, # '¥' - 166: 253, # '¦' - 167: 253, # '§' - 168: 253, # '¨' - 169: 253, # '©' - 170: 253, # None - 171: 253, # '«' - 172: 253, # '¬' - 173: 74, # '\xad' - 174: 253, # '®' - 175: 253, # '―' - 176: 253, # '°' - 177: 253, # '±' - 178: 253, # '²' - 179: 253, # '³' - 180: 247, # '΄' - 181: 253, # 'µ' - 182: 253, # '¶' - 183: 36, # '·' - 184: 46, # 'Έ' - 185: 71, # 'Ή' - 186: 73, # 'Ί' - 187: 253, # '»' - 188: 54, # 'Ό' - 189: 253, # '½' - 190: 108, # 'Ύ' - 191: 123, # 'Ώ' - 192: 110, # 'ΐ' - 193: 31, # 'Α' - 194: 51, # 'Β' - 195: 43, # 'Γ' - 196: 41, # 'Δ' - 197: 34, # 'Ε' - 198: 91, # 'Ζ' - 199: 40, # 'Η' - 200: 52, # 'Θ' - 201: 47, # 'Ι' - 202: 44, # 'Κ' - 203: 53, # 'Λ' - 204: 38, # 'Μ' - 205: 49, # 'Ν' - 206: 59, # 'Ξ' - 207: 39, # 'Ο' - 208: 35, # 'Π' - 209: 48, # 'Ρ' - 210: 250, # None - 211: 37, # 'Σ' - 212: 33, # 'Τ' - 213: 45, # 'Υ' - 214: 56, # 'Φ' - 215: 50, # 'Χ' - 216: 84, # 'Ψ' - 217: 57, # 'Ω' - 218: 120, # 'Ϊ' - 219: 121, # 'Ϋ' - 220: 17, # 'ά' - 221: 18, # 'έ' - 222: 22, # 'ή' - 223: 15, # 'ί' - 224: 124, # 'ΰ' - 225: 1, # 'α' - 226: 29, # 'β' - 227: 20, # 'γ' - 228: 21, # 'δ' - 229: 3, # 'ε' - 230: 32, # 'ζ' - 231: 13, # 'η' - 232: 25, # 'θ' - 233: 5, # 'ι' - 234: 11, # 'κ' - 235: 16, # 'λ' - 236: 10, # 'μ' - 237: 6, # 'ν' - 238: 30, # 'ξ' - 239: 4, # 'ο' - 240: 9, # 'π' - 241: 8, # 'ρ' - 242: 14, # 'ς' - 243: 7, # 'σ' - 244: 2, # 'τ' - 245: 12, # 'υ' - 246: 28, # 'φ' - 247: 23, # 'χ' - 248: 42, # 'ψ' - 249: 24, # 'ω' - 250: 64, # 'ϊ' - 251: 75, # 'ϋ' - 252: 19, # 'ό' - 253: 26, # 'ύ' - 254: 27, # 'ώ' - 255: 253, # None -} - -WINDOWS_1253_GREEK_MODEL = SingleByteCharSetModel( - charset_name="windows-1253", - language="Greek", - char_to_order_map=WINDOWS_1253_GREEK_CHAR_TO_ORDER, - language_model=GREEK_LANG_MODEL, - typical_positive_ratio=0.982851, - keep_ascii_letters=False, - alphabet="ΆΈΉΊΌΎΏΑΒΓΔΕΖΗΘΙΚΛΜΝΞΟΠΡΣΤΥΦΧΨΩάέήίαβγδεζηθικλμνξοπρςστυφχψωόύώ", -) - -ISO_8859_7_GREEK_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 254, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 254, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 253, # ' ' - 33: 253, # '!' - 34: 253, # '"' - 35: 253, # '#' - 36: 253, # '$' - 37: 253, # '%' - 38: 253, # '&' - 39: 253, # "'" - 40: 253, # '(' - 41: 253, # ')' - 42: 253, # '*' - 43: 253, # '+' - 44: 253, # ',' - 45: 253, # '-' - 46: 253, # '.' - 47: 253, # '/' - 48: 252, # '0' - 49: 252, # '1' - 50: 252, # '2' - 51: 252, # '3' - 52: 252, # '4' - 53: 252, # '5' - 54: 252, # '6' - 55: 252, # '7' - 56: 252, # '8' - 57: 252, # '9' - 58: 253, # ':' - 59: 253, # ';' - 60: 253, # '<' - 61: 253, # '=' - 62: 253, # '>' - 63: 253, # '?' - 64: 253, # '@' - 65: 82, # 'A' - 66: 100, # 'B' - 67: 104, # 'C' - 68: 94, # 'D' - 69: 98, # 'E' - 70: 101, # 'F' - 71: 116, # 'G' - 72: 102, # 'H' - 73: 111, # 'I' - 74: 187, # 'J' - 75: 117, # 'K' - 76: 92, # 'L' - 77: 88, # 'M' - 78: 113, # 'N' - 79: 85, # 'O' - 80: 79, # 'P' - 81: 118, # 'Q' - 82: 105, # 'R' - 83: 83, # 'S' - 84: 67, # 'T' - 85: 114, # 'U' - 86: 119, # 'V' - 87: 95, # 'W' - 88: 99, # 'X' - 89: 109, # 'Y' - 90: 188, # 'Z' - 91: 253, # '[' - 92: 253, # '\\' - 93: 253, # ']' - 94: 253, # '^' - 95: 253, # '_' - 96: 253, # '`' - 97: 72, # 'a' - 98: 70, # 'b' - 99: 80, # 'c' - 100: 81, # 'd' - 101: 60, # 'e' - 102: 96, # 'f' - 103: 93, # 'g' - 104: 89, # 'h' - 105: 68, # 'i' - 106: 120, # 'j' - 107: 97, # 'k' - 108: 77, # 'l' - 109: 86, # 'm' - 110: 69, # 'n' - 111: 55, # 'o' - 112: 78, # 'p' - 113: 115, # 'q' - 114: 65, # 'r' - 115: 66, # 's' - 116: 58, # 't' - 117: 76, # 'u' - 118: 106, # 'v' - 119: 103, # 'w' - 120: 87, # 'x' - 121: 107, # 'y' - 122: 112, # 'z' - 123: 253, # '{' - 124: 253, # '|' - 125: 253, # '}' - 126: 253, # '~' - 127: 253, # '\x7f' - 128: 255, # '\x80' - 129: 255, # '\x81' - 130: 255, # '\x82' - 131: 255, # '\x83' - 132: 255, # '\x84' - 133: 255, # '\x85' - 134: 255, # '\x86' - 135: 255, # '\x87' - 136: 255, # '\x88' - 137: 255, # '\x89' - 138: 255, # '\x8a' - 139: 255, # '\x8b' - 140: 255, # '\x8c' - 141: 255, # '\x8d' - 142: 255, # '\x8e' - 143: 255, # '\x8f' - 144: 255, # '\x90' - 145: 255, # '\x91' - 146: 255, # '\x92' - 147: 255, # '\x93' - 148: 255, # '\x94' - 149: 255, # '\x95' - 150: 255, # '\x96' - 151: 255, # '\x97' - 152: 255, # '\x98' - 153: 255, # '\x99' - 154: 255, # '\x9a' - 155: 255, # '\x9b' - 156: 255, # '\x9c' - 157: 255, # '\x9d' - 158: 255, # '\x9e' - 159: 255, # '\x9f' - 160: 253, # '\xa0' - 161: 233, # '‘' - 162: 90, # '’' - 163: 253, # '£' - 164: 253, # '€' - 165: 253, # '₯' - 166: 253, # '¦' - 167: 253, # '§' - 168: 253, # '¨' - 169: 253, # '©' - 170: 253, # 'ͺ' - 171: 253, # '«' - 172: 253, # '¬' - 173: 74, # '\xad' - 174: 253, # None - 175: 253, # '―' - 176: 253, # '°' - 177: 253, # '±' - 178: 253, # '²' - 179: 253, # '³' - 180: 247, # '΄' - 181: 248, # '΅' - 182: 61, # 'Ά' - 183: 36, # '·' - 184: 46, # 'Έ' - 185: 71, # 'Ή' - 186: 73, # 'Ί' - 187: 253, # '»' - 188: 54, # 'Ό' - 189: 253, # '½' - 190: 108, # 'Ύ' - 191: 123, # 'Ώ' - 192: 110, # 'ΐ' - 193: 31, # 'Α' - 194: 51, # 'Β' - 195: 43, # 'Γ' - 196: 41, # 'Δ' - 197: 34, # 'Ε' - 198: 91, # 'Ζ' - 199: 40, # 'Η' - 200: 52, # 'Θ' - 201: 47, # 'Ι' - 202: 44, # 'Κ' - 203: 53, # 'Λ' - 204: 38, # 'Μ' - 205: 49, # 'Ν' - 206: 59, # 'Ξ' - 207: 39, # 'Ο' - 208: 35, # 'Π' - 209: 48, # 'Ρ' - 210: 250, # None - 211: 37, # 'Σ' - 212: 33, # 'Τ' - 213: 45, # 'Υ' - 214: 56, # 'Φ' - 215: 50, # 'Χ' - 216: 84, # 'Ψ' - 217: 57, # 'Ω' - 218: 120, # 'Ϊ' - 219: 121, # 'Ϋ' - 220: 17, # 'ά' - 221: 18, # 'έ' - 222: 22, # 'ή' - 223: 15, # 'ί' - 224: 124, # 'ΰ' - 225: 1, # 'α' - 226: 29, # 'β' - 227: 20, # 'γ' - 228: 21, # 'δ' - 229: 3, # 'ε' - 230: 32, # 'ζ' - 231: 13, # 'η' - 232: 25, # 'θ' - 233: 5, # 'ι' - 234: 11, # 'κ' - 235: 16, # 'λ' - 236: 10, # 'μ' - 237: 6, # 'ν' - 238: 30, # 'ξ' - 239: 4, # 'ο' - 240: 9, # 'π' - 241: 8, # 'ρ' - 242: 14, # 'ς' - 243: 7, # 'σ' - 244: 2, # 'τ' - 245: 12, # 'υ' - 246: 28, # 'φ' - 247: 23, # 'χ' - 248: 42, # 'ψ' - 249: 24, # 'ω' - 250: 64, # 'ϊ' - 251: 75, # 'ϋ' - 252: 19, # 'ό' - 253: 26, # 'ύ' - 254: 27, # 'ώ' - 255: 253, # None -} - -ISO_8859_7_GREEK_MODEL = SingleByteCharSetModel( - charset_name="ISO-8859-7", - language="Greek", - char_to_order_map=ISO_8859_7_GREEK_CHAR_TO_ORDER, - language_model=GREEK_LANG_MODEL, - typical_positive_ratio=0.982851, - keep_ascii_letters=False, - alphabet="ΆΈΉΊΌΎΏΑΒΓΔΕΖΗΘΙΚΛΜΝΞΟΠΡΣΤΥΦΧΨΩάέήίαβγδεζηθικλμνξοπρςστυφχψωόύώ", -) diff --git a/venv/lib/python3.10/site-packages/chardet/langhebrewmodel.py b/venv/lib/python3.10/site-packages/chardet/langhebrewmodel.py deleted file mode 100644 index 86b3c5e64558ea9cd8ed4a4b603a8799fb2312e6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/langhebrewmodel.py +++ /dev/null @@ -1,4380 +0,0 @@ -from chardet.sbcharsetprober import SingleByteCharSetModel - -# 3: Positive -# 2: Likely -# 1: Unlikely -# 0: Negative - -HEBREW_LANG_MODEL = { - 50: { # 'a' - 50: 0, # 'a' - 60: 1, # 'c' - 61: 1, # 'd' - 42: 1, # 'e' - 53: 1, # 'i' - 56: 2, # 'l' - 54: 2, # 'n' - 49: 0, # 'o' - 51: 2, # 'r' - 43: 1, # 's' - 44: 2, # 't' - 63: 1, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 1, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 1, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 1, # 'ק' - 7: 0, # 'ר' - 10: 1, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 1, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 60: { # 'c' - 50: 1, # 'a' - 60: 1, # 'c' - 61: 0, # 'd' - 42: 1, # 'e' - 53: 1, # 'i' - 56: 1, # 'l' - 54: 0, # 'n' - 49: 1, # 'o' - 51: 1, # 'r' - 43: 1, # 's' - 44: 2, # 't' - 63: 1, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 1, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 1, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 1, # 'מ' - 23: 0, # 'ן' - 12: 1, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 61: { # 'd' - 50: 1, # 'a' - 60: 0, # 'c' - 61: 1, # 'd' - 42: 1, # 'e' - 53: 1, # 'i' - 56: 1, # 'l' - 54: 1, # 'n' - 49: 2, # 'o' - 51: 1, # 'r' - 43: 1, # 's' - 44: 0, # 't' - 63: 1, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 1, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 1, # '–' - 52: 1, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 42: { # 'e' - 50: 1, # 'a' - 60: 1, # 'c' - 61: 2, # 'd' - 42: 1, # 'e' - 53: 1, # 'i' - 56: 2, # 'l' - 54: 2, # 'n' - 49: 1, # 'o' - 51: 2, # 'r' - 43: 2, # 's' - 44: 2, # 't' - 63: 1, # 'u' - 34: 1, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 1, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 1, # '–' - 52: 2, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 53: { # 'i' - 50: 1, # 'a' - 60: 2, # 'c' - 61: 1, # 'd' - 42: 1, # 'e' - 53: 0, # 'i' - 56: 1, # 'l' - 54: 2, # 'n' - 49: 2, # 'o' - 51: 1, # 'r' - 43: 2, # 's' - 44: 2, # 't' - 63: 1, # 'u' - 34: 0, # '\xa0' - 55: 1, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 1, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 56: { # 'l' - 50: 1, # 'a' - 60: 1, # 'c' - 61: 1, # 'd' - 42: 2, # 'e' - 53: 2, # 'i' - 56: 2, # 'l' - 54: 1, # 'n' - 49: 1, # 'o' - 51: 0, # 'r' - 43: 1, # 's' - 44: 1, # 't' - 63: 1, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 1, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 54: { # 'n' - 50: 1, # 'a' - 60: 1, # 'c' - 61: 1, # 'd' - 42: 1, # 'e' - 53: 1, # 'i' - 56: 1, # 'l' - 54: 1, # 'n' - 49: 1, # 'o' - 51: 0, # 'r' - 43: 1, # 's' - 44: 2, # 't' - 63: 1, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 1, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 2, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 49: { # 'o' - 50: 1, # 'a' - 60: 1, # 'c' - 61: 1, # 'd' - 42: 1, # 'e' - 53: 1, # 'i' - 56: 1, # 'l' - 54: 2, # 'n' - 49: 1, # 'o' - 51: 2, # 'r' - 43: 1, # 's' - 44: 1, # 't' - 63: 1, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 1, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 51: { # 'r' - 50: 2, # 'a' - 60: 1, # 'c' - 61: 1, # 'd' - 42: 2, # 'e' - 53: 1, # 'i' - 56: 1, # 'l' - 54: 1, # 'n' - 49: 2, # 'o' - 51: 1, # 'r' - 43: 1, # 's' - 44: 1, # 't' - 63: 1, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 2, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 43: { # 's' - 50: 1, # 'a' - 60: 1, # 'c' - 61: 0, # 'd' - 42: 2, # 'e' - 53: 1, # 'i' - 56: 1, # 'l' - 54: 1, # 'n' - 49: 1, # 'o' - 51: 1, # 'r' - 43: 1, # 's' - 44: 2, # 't' - 63: 1, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 1, # '’' - 47: 0, # '“' - 46: 2, # '”' - 58: 0, # '†' - 40: 2, # '…' - }, - 44: { # 't' - 50: 1, # 'a' - 60: 1, # 'c' - 61: 0, # 'd' - 42: 2, # 'e' - 53: 2, # 'i' - 56: 1, # 'l' - 54: 0, # 'n' - 49: 1, # 'o' - 51: 1, # 'r' - 43: 1, # 's' - 44: 1, # 't' - 63: 1, # 'u' - 34: 1, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 2, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 63: { # 'u' - 50: 1, # 'a' - 60: 1, # 'c' - 61: 1, # 'd' - 42: 1, # 'e' - 53: 1, # 'i' - 56: 1, # 'l' - 54: 1, # 'n' - 49: 0, # 'o' - 51: 1, # 'r' - 43: 2, # 's' - 44: 1, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 1, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 34: { # '\xa0' - 50: 1, # 'a' - 60: 0, # 'c' - 61: 1, # 'd' - 42: 0, # 'e' - 53: 1, # 'i' - 56: 0, # 'l' - 54: 1, # 'n' - 49: 1, # 'o' - 51: 0, # 'r' - 43: 1, # 's' - 44: 1, # 't' - 63: 0, # 'u' - 34: 2, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 2, # 'א' - 8: 1, # 'ב' - 20: 1, # 'ג' - 16: 1, # 'ד' - 3: 1, # 'ה' - 2: 1, # 'ו' - 24: 1, # 'ז' - 14: 1, # 'ח' - 22: 1, # 'ט' - 1: 2, # 'י' - 25: 0, # 'ך' - 15: 1, # 'כ' - 4: 1, # 'ל' - 11: 0, # 'ם' - 6: 2, # 'מ' - 23: 0, # 'ן' - 12: 1, # 'נ' - 19: 1, # 'ס' - 13: 1, # 'ע' - 26: 0, # 'ף' - 18: 1, # 'פ' - 27: 0, # 'ץ' - 21: 1, # 'צ' - 17: 1, # 'ק' - 7: 1, # 'ר' - 10: 1, # 'ש' - 5: 1, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 55: { # '´' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 1, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 1, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 1, # 'ה' - 2: 1, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 2, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 1, # 'ל' - 11: 0, # 'ם' - 6: 1, # 'מ' - 23: 1, # 'ן' - 12: 1, # 'נ' - 19: 1, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 1, # 'ר' - 10: 1, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 48: { # '¼' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 1, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 1, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 1, # 'כ' - 4: 1, # 'ל' - 11: 0, # 'ם' - 6: 1, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 39: { # '½' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 1, # 'כ' - 4: 1, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 1, # 'צ' - 17: 1, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 57: { # '¾' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 30: { # 'ְ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 1, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 1, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 2, # 'א' - 8: 2, # 'ב' - 20: 2, # 'ג' - 16: 2, # 'ד' - 3: 2, # 'ה' - 2: 2, # 'ו' - 24: 2, # 'ז' - 14: 2, # 'ח' - 22: 2, # 'ט' - 1: 2, # 'י' - 25: 2, # 'ך' - 15: 2, # 'כ' - 4: 2, # 'ל' - 11: 1, # 'ם' - 6: 2, # 'מ' - 23: 0, # 'ן' - 12: 2, # 'נ' - 19: 2, # 'ס' - 13: 2, # 'ע' - 26: 0, # 'ף' - 18: 2, # 'פ' - 27: 0, # 'ץ' - 21: 2, # 'צ' - 17: 2, # 'ק' - 7: 2, # 'ר' - 10: 2, # 'ש' - 5: 2, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 59: { # 'ֱ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 1, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 1, # 'ב' - 20: 1, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 0, # 'ו' - 24: 1, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 1, # 'י' - 25: 0, # 'ך' - 15: 1, # 'כ' - 4: 2, # 'ל' - 11: 0, # 'ם' - 6: 2, # 'מ' - 23: 0, # 'ן' - 12: 1, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 1, # 'ר' - 10: 1, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 41: { # 'ֲ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 2, # 'ב' - 20: 1, # 'ג' - 16: 2, # 'ד' - 3: 1, # 'ה' - 2: 1, # 'ו' - 24: 1, # 'ז' - 14: 1, # 'ח' - 22: 1, # 'ט' - 1: 1, # 'י' - 25: 1, # 'ך' - 15: 1, # 'כ' - 4: 2, # 'ל' - 11: 0, # 'ם' - 6: 2, # 'מ' - 23: 0, # 'ן' - 12: 2, # 'נ' - 19: 1, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 1, # 'פ' - 27: 0, # 'ץ' - 21: 2, # 'צ' - 17: 1, # 'ק' - 7: 2, # 'ר' - 10: 2, # 'ש' - 5: 1, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 33: { # 'ִ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 1, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 1, # 'ִ' - 37: 0, # 'ֵ' - 36: 1, # 'ֶ' - 31: 0, # 'ַ' - 29: 1, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 1, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 1, # 'א' - 8: 2, # 'ב' - 20: 2, # 'ג' - 16: 2, # 'ד' - 3: 1, # 'ה' - 2: 1, # 'ו' - 24: 2, # 'ז' - 14: 1, # 'ח' - 22: 1, # 'ט' - 1: 3, # 'י' - 25: 1, # 'ך' - 15: 2, # 'כ' - 4: 2, # 'ל' - 11: 2, # 'ם' - 6: 2, # 'מ' - 23: 2, # 'ן' - 12: 2, # 'נ' - 19: 2, # 'ס' - 13: 1, # 'ע' - 26: 0, # 'ף' - 18: 2, # 'פ' - 27: 1, # 'ץ' - 21: 2, # 'צ' - 17: 2, # 'ק' - 7: 2, # 'ר' - 10: 2, # 'ש' - 5: 2, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 37: { # 'ֵ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 1, # 'ֶ' - 31: 1, # 'ַ' - 29: 1, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 2, # 'א' - 8: 2, # 'ב' - 20: 1, # 'ג' - 16: 2, # 'ד' - 3: 2, # 'ה' - 2: 1, # 'ו' - 24: 1, # 'ז' - 14: 2, # 'ח' - 22: 1, # 'ט' - 1: 3, # 'י' - 25: 2, # 'ך' - 15: 1, # 'כ' - 4: 2, # 'ל' - 11: 2, # 'ם' - 6: 1, # 'מ' - 23: 2, # 'ן' - 12: 2, # 'נ' - 19: 1, # 'ס' - 13: 2, # 'ע' - 26: 1, # 'ף' - 18: 1, # 'פ' - 27: 1, # 'ץ' - 21: 1, # 'צ' - 17: 1, # 'ק' - 7: 2, # 'ר' - 10: 2, # 'ש' - 5: 2, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 36: { # 'ֶ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 1, # 'ֶ' - 31: 1, # 'ַ' - 29: 1, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 2, # 'א' - 8: 2, # 'ב' - 20: 1, # 'ג' - 16: 2, # 'ד' - 3: 2, # 'ה' - 2: 1, # 'ו' - 24: 1, # 'ז' - 14: 2, # 'ח' - 22: 1, # 'ט' - 1: 2, # 'י' - 25: 2, # 'ך' - 15: 1, # 'כ' - 4: 2, # 'ל' - 11: 2, # 'ם' - 6: 2, # 'מ' - 23: 2, # 'ן' - 12: 2, # 'נ' - 19: 2, # 'ס' - 13: 1, # 'ע' - 26: 1, # 'ף' - 18: 1, # 'פ' - 27: 2, # 'ץ' - 21: 1, # 'צ' - 17: 1, # 'ק' - 7: 2, # 'ר' - 10: 2, # 'ש' - 5: 2, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 31: { # 'ַ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 1, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 1, # 'ֶ' - 31: 0, # 'ַ' - 29: 2, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 2, # 'א' - 8: 2, # 'ב' - 20: 2, # 'ג' - 16: 2, # 'ד' - 3: 2, # 'ה' - 2: 1, # 'ו' - 24: 2, # 'ז' - 14: 2, # 'ח' - 22: 2, # 'ט' - 1: 3, # 'י' - 25: 1, # 'ך' - 15: 2, # 'כ' - 4: 2, # 'ל' - 11: 2, # 'ם' - 6: 2, # 'מ' - 23: 2, # 'ן' - 12: 2, # 'נ' - 19: 2, # 'ס' - 13: 2, # 'ע' - 26: 2, # 'ף' - 18: 2, # 'פ' - 27: 1, # 'ץ' - 21: 2, # 'צ' - 17: 2, # 'ק' - 7: 2, # 'ר' - 10: 2, # 'ש' - 5: 2, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 29: { # 'ָ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 1, # 'ַ' - 29: 2, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 1, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 2, # 'א' - 8: 2, # 'ב' - 20: 2, # 'ג' - 16: 2, # 'ד' - 3: 3, # 'ה' - 2: 2, # 'ו' - 24: 2, # 'ז' - 14: 2, # 'ח' - 22: 1, # 'ט' - 1: 2, # 'י' - 25: 2, # 'ך' - 15: 2, # 'כ' - 4: 2, # 'ל' - 11: 2, # 'ם' - 6: 2, # 'מ' - 23: 2, # 'ן' - 12: 2, # 'נ' - 19: 1, # 'ס' - 13: 2, # 'ע' - 26: 1, # 'ף' - 18: 2, # 'פ' - 27: 1, # 'ץ' - 21: 2, # 'צ' - 17: 2, # 'ק' - 7: 2, # 'ר' - 10: 2, # 'ש' - 5: 2, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 35: { # 'ֹ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 1, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 2, # 'א' - 8: 2, # 'ב' - 20: 1, # 'ג' - 16: 2, # 'ד' - 3: 2, # 'ה' - 2: 1, # 'ו' - 24: 1, # 'ז' - 14: 1, # 'ח' - 22: 1, # 'ט' - 1: 1, # 'י' - 25: 1, # 'ך' - 15: 2, # 'כ' - 4: 2, # 'ל' - 11: 2, # 'ם' - 6: 2, # 'מ' - 23: 2, # 'ן' - 12: 2, # 'נ' - 19: 2, # 'ס' - 13: 2, # 'ע' - 26: 1, # 'ף' - 18: 2, # 'פ' - 27: 1, # 'ץ' - 21: 2, # 'צ' - 17: 2, # 'ק' - 7: 2, # 'ר' - 10: 2, # 'ש' - 5: 2, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 62: { # 'ֻ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 1, # 'ב' - 20: 1, # 'ג' - 16: 1, # 'ד' - 3: 1, # 'ה' - 2: 1, # 'ו' - 24: 1, # 'ז' - 14: 1, # 'ח' - 22: 0, # 'ט' - 1: 1, # 'י' - 25: 0, # 'ך' - 15: 1, # 'כ' - 4: 2, # 'ל' - 11: 1, # 'ם' - 6: 1, # 'מ' - 23: 1, # 'ן' - 12: 1, # 'נ' - 19: 1, # 'ס' - 13: 1, # 'ע' - 26: 0, # 'ף' - 18: 1, # 'פ' - 27: 0, # 'ץ' - 21: 1, # 'צ' - 17: 1, # 'ק' - 7: 1, # 'ר' - 10: 1, # 'ש' - 5: 1, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 28: { # 'ּ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 3, # 'ְ' - 59: 0, # 'ֱ' - 41: 1, # 'ֲ' - 33: 3, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 3, # 'ַ' - 29: 3, # 'ָ' - 35: 2, # 'ֹ' - 62: 1, # 'ֻ' - 28: 0, # 'ּ' - 38: 2, # 'ׁ' - 45: 1, # 'ׂ' - 9: 2, # 'א' - 8: 2, # 'ב' - 20: 1, # 'ג' - 16: 2, # 'ד' - 3: 1, # 'ה' - 2: 2, # 'ו' - 24: 1, # 'ז' - 14: 1, # 'ח' - 22: 1, # 'ט' - 1: 2, # 'י' - 25: 2, # 'ך' - 15: 2, # 'כ' - 4: 2, # 'ל' - 11: 1, # 'ם' - 6: 2, # 'מ' - 23: 1, # 'ן' - 12: 2, # 'נ' - 19: 1, # 'ס' - 13: 2, # 'ע' - 26: 1, # 'ף' - 18: 1, # 'פ' - 27: 1, # 'ץ' - 21: 1, # 'צ' - 17: 1, # 'ק' - 7: 2, # 'ר' - 10: 2, # 'ש' - 5: 2, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 38: { # 'ׁ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 1, # 'ֹ' - 62: 1, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 2, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 1, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 1, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 45: { # 'ׂ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 1, # 'ֵ' - 36: 2, # 'ֶ' - 31: 1, # 'ַ' - 29: 2, # 'ָ' - 35: 1, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 1, # 'א' - 8: 0, # 'ב' - 20: 1, # 'ג' - 16: 0, # 'ד' - 3: 1, # 'ה' - 2: 2, # 'ו' - 24: 0, # 'ז' - 14: 1, # 'ח' - 22: 0, # 'ט' - 1: 1, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 1, # 'ם' - 6: 1, # 'מ' - 23: 0, # 'ן' - 12: 1, # 'נ' - 19: 0, # 'ס' - 13: 1, # 'ע' - 26: 0, # 'ף' - 18: 1, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 1, # 'ר' - 10: 0, # 'ש' - 5: 1, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 9: { # 'א' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 1, # '´' - 48: 1, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 2, # 'ֱ' - 41: 2, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 2, # 'ֹ' - 62: 1, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 2, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 3, # 'ז' - 14: 3, # 'ח' - 22: 3, # 'ט' - 1: 3, # 'י' - 25: 3, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 3, # 'ן' - 12: 3, # 'נ' - 19: 3, # 'ס' - 13: 2, # 'ע' - 26: 3, # 'ף' - 18: 3, # 'פ' - 27: 1, # 'ץ' - 21: 3, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 8: { # 'ב' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 1, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 1, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 2, # 'ֹ' - 62: 1, # 'ֻ' - 28: 3, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 3, # 'ז' - 14: 3, # 'ח' - 22: 3, # 'ט' - 1: 3, # 'י' - 25: 2, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 2, # 'ם' - 6: 3, # 'מ' - 23: 3, # 'ן' - 12: 3, # 'נ' - 19: 3, # 'ס' - 13: 3, # 'ע' - 26: 1, # 'ף' - 18: 3, # 'פ' - 27: 2, # 'ץ' - 21: 3, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 1, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 20: { # 'ג' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 2, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 1, # 'ִ' - 37: 1, # 'ֵ' - 36: 1, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 1, # 'ֹ' - 62: 0, # 'ֻ' - 28: 2, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 2, # 'א' - 8: 3, # 'ב' - 20: 2, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 3, # 'ז' - 14: 2, # 'ח' - 22: 2, # 'ט' - 1: 3, # 'י' - 25: 1, # 'ך' - 15: 1, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 3, # 'ן' - 12: 3, # 'נ' - 19: 2, # 'ס' - 13: 3, # 'ע' - 26: 2, # 'ף' - 18: 2, # 'פ' - 27: 1, # 'ץ' - 21: 1, # 'צ' - 17: 1, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 1, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 16: { # 'ד' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 2, # 'ֹ' - 62: 1, # 'ֻ' - 28: 2, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 1, # 'ז' - 14: 2, # 'ח' - 22: 2, # 'ט' - 1: 3, # 'י' - 25: 2, # 'ך' - 15: 2, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 2, # 'ן' - 12: 3, # 'נ' - 19: 2, # 'ס' - 13: 3, # 'ע' - 26: 2, # 'ף' - 18: 3, # 'פ' - 27: 0, # 'ץ' - 21: 2, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 3: { # 'ה' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 1, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 0, # '´' - 48: 1, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 1, # 'ְ' - 59: 1, # 'ֱ' - 41: 2, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 3, # 'ַ' - 29: 2, # 'ָ' - 35: 1, # 'ֹ' - 62: 1, # 'ֻ' - 28: 2, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 3, # 'ז' - 14: 3, # 'ח' - 22: 3, # 'ט' - 1: 3, # 'י' - 25: 1, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 3, # 'ן' - 12: 3, # 'נ' - 19: 3, # 'ס' - 13: 3, # 'ע' - 26: 0, # 'ף' - 18: 3, # 'פ' - 27: 1, # 'ץ' - 21: 3, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 1, # '–' - 52: 1, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 2, # '…' - }, - 2: { # 'ו' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 1, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 1, # '´' - 48: 1, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 1, # 'ֵ' - 36: 1, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 3, # 'ֹ' - 62: 0, # 'ֻ' - 28: 3, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 3, # 'ז' - 14: 3, # 'ח' - 22: 3, # 'ט' - 1: 3, # 'י' - 25: 3, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 3, # 'ן' - 12: 3, # 'נ' - 19: 3, # 'ס' - 13: 3, # 'ע' - 26: 3, # 'ף' - 18: 3, # 'פ' - 27: 3, # 'ץ' - 21: 3, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 1, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 2, # '…' - }, - 24: { # 'ז' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 1, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 1, # 'ֲ' - 33: 1, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 1, # 'ֹ' - 62: 1, # 'ֻ' - 28: 2, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 2, # 'ב' - 20: 2, # 'ג' - 16: 2, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 2, # 'ז' - 14: 2, # 'ח' - 22: 1, # 'ט' - 1: 3, # 'י' - 25: 1, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 2, # 'ם' - 6: 3, # 'מ' - 23: 2, # 'ן' - 12: 2, # 'נ' - 19: 1, # 'ס' - 13: 2, # 'ע' - 26: 1, # 'ף' - 18: 1, # 'פ' - 27: 0, # 'ץ' - 21: 2, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 1, # 'ש' - 5: 2, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 14: { # 'ח' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 1, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 1, # 'ֱ' - 41: 2, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 2, # 'ֹ' - 62: 1, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 2, # 'א' - 8: 3, # 'ב' - 20: 2, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 3, # 'ז' - 14: 2, # 'ח' - 22: 2, # 'ט' - 1: 3, # 'י' - 25: 1, # 'ך' - 15: 2, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 2, # 'ן' - 12: 3, # 'נ' - 19: 3, # 'ס' - 13: 1, # 'ע' - 26: 2, # 'ף' - 18: 2, # 'פ' - 27: 2, # 'ץ' - 21: 3, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 1, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 22: { # 'ט' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 1, # 'ֵ' - 36: 1, # 'ֶ' - 31: 2, # 'ַ' - 29: 1, # 'ָ' - 35: 1, # 'ֹ' - 62: 1, # 'ֻ' - 28: 1, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 1, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 2, # 'ז' - 14: 3, # 'ח' - 22: 2, # 'ט' - 1: 3, # 'י' - 25: 1, # 'ך' - 15: 2, # 'כ' - 4: 3, # 'ל' - 11: 2, # 'ם' - 6: 2, # 'מ' - 23: 2, # 'ן' - 12: 3, # 'נ' - 19: 2, # 'ס' - 13: 3, # 'ע' - 26: 2, # 'ף' - 18: 3, # 'פ' - 27: 1, # 'ץ' - 21: 2, # 'צ' - 17: 2, # 'ק' - 7: 3, # 'ר' - 10: 2, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 1: { # 'י' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 1, # '´' - 48: 1, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 1, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 2, # 'ֹ' - 62: 1, # 'ֻ' - 28: 2, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 3, # 'ז' - 14: 3, # 'ח' - 22: 3, # 'ט' - 1: 3, # 'י' - 25: 3, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 3, # 'ן' - 12: 3, # 'נ' - 19: 3, # 'ס' - 13: 3, # 'ע' - 26: 3, # 'ף' - 18: 3, # 'פ' - 27: 3, # 'ץ' - 21: 3, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 1, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 2, # '…' - }, - 25: { # 'ך' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 2, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 1, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 1, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 1, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 1, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 1, # 'ל' - 11: 0, # 'ם' - 6: 1, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 1, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 15: { # 'כ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 1, # 'ֹ' - 62: 1, # 'ֻ' - 28: 3, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 2, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 3, # 'ז' - 14: 3, # 'ח' - 22: 2, # 'ט' - 1: 3, # 'י' - 25: 3, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 3, # 'ן' - 12: 3, # 'נ' - 19: 3, # 'ס' - 13: 2, # 'ע' - 26: 3, # 'ף' - 18: 3, # 'פ' - 27: 1, # 'ץ' - 21: 2, # 'צ' - 17: 2, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 4: { # 'ל' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 1, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 3, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 2, # 'ֹ' - 62: 1, # 'ֻ' - 28: 2, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 3, # 'ז' - 14: 3, # 'ח' - 22: 3, # 'ט' - 1: 3, # 'י' - 25: 3, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 2, # 'ן' - 12: 3, # 'נ' - 19: 3, # 'ס' - 13: 3, # 'ע' - 26: 2, # 'ף' - 18: 3, # 'פ' - 27: 2, # 'ץ' - 21: 3, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 1, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 11: { # 'ם' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 1, # 'א' - 8: 1, # 'ב' - 20: 1, # 'ג' - 16: 0, # 'ד' - 3: 1, # 'ה' - 2: 1, # 'ו' - 24: 1, # 'ז' - 14: 1, # 'ח' - 22: 0, # 'ט' - 1: 1, # 'י' - 25: 0, # 'ך' - 15: 1, # 'כ' - 4: 1, # 'ל' - 11: 1, # 'ם' - 6: 1, # 'מ' - 23: 0, # 'ן' - 12: 1, # 'נ' - 19: 0, # 'ס' - 13: 1, # 'ע' - 26: 0, # 'ף' - 18: 1, # 'פ' - 27: 1, # 'ץ' - 21: 1, # 'צ' - 17: 1, # 'ק' - 7: 1, # 'ר' - 10: 1, # 'ש' - 5: 1, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 2, # '…' - }, - 6: { # 'מ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 1, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 2, # 'ֹ' - 62: 1, # 'ֻ' - 28: 2, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 3, # 'ז' - 14: 3, # 'ח' - 22: 3, # 'ט' - 1: 3, # 'י' - 25: 2, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 3, # 'ן' - 12: 3, # 'נ' - 19: 3, # 'ס' - 13: 3, # 'ע' - 26: 0, # 'ף' - 18: 3, # 'פ' - 27: 2, # 'ץ' - 21: 3, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 23: { # 'ן' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 0, # '´' - 48: 1, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 1, # 'א' - 8: 1, # 'ב' - 20: 1, # 'ג' - 16: 1, # 'ד' - 3: 1, # 'ה' - 2: 1, # 'ו' - 24: 0, # 'ז' - 14: 1, # 'ח' - 22: 1, # 'ט' - 1: 1, # 'י' - 25: 0, # 'ך' - 15: 1, # 'כ' - 4: 1, # 'ל' - 11: 1, # 'ם' - 6: 1, # 'מ' - 23: 0, # 'ן' - 12: 1, # 'נ' - 19: 1, # 'ס' - 13: 1, # 'ע' - 26: 1, # 'ף' - 18: 1, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 1, # 'ק' - 7: 1, # 'ר' - 10: 1, # 'ש' - 5: 1, # 'ת' - 32: 1, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 2, # '…' - }, - 12: { # 'נ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 1, # 'ֹ' - 62: 1, # 'ֻ' - 28: 2, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 3, # 'ז' - 14: 3, # 'ח' - 22: 3, # 'ט' - 1: 3, # 'י' - 25: 2, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 3, # 'ן' - 12: 3, # 'נ' - 19: 3, # 'ס' - 13: 3, # 'ע' - 26: 2, # 'ף' - 18: 3, # 'פ' - 27: 2, # 'ץ' - 21: 3, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 19: { # 'ס' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 1, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 1, # 'ֵ' - 36: 2, # 'ֶ' - 31: 2, # 'ַ' - 29: 1, # 'ָ' - 35: 1, # 'ֹ' - 62: 2, # 'ֻ' - 28: 2, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 2, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 1, # 'ז' - 14: 3, # 'ח' - 22: 3, # 'ט' - 1: 3, # 'י' - 25: 2, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 2, # 'ם' - 6: 3, # 'מ' - 23: 2, # 'ן' - 12: 3, # 'נ' - 19: 2, # 'ס' - 13: 3, # 'ע' - 26: 3, # 'ף' - 18: 3, # 'פ' - 27: 0, # 'ץ' - 21: 2, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 1, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 13: { # 'ע' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 1, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 1, # 'ְ' - 59: 1, # 'ֱ' - 41: 2, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 2, # 'ֹ' - 62: 1, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 2, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 3, # 'ז' - 14: 1, # 'ח' - 22: 3, # 'ט' - 1: 3, # 'י' - 25: 2, # 'ך' - 15: 2, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 2, # 'ן' - 12: 3, # 'נ' - 19: 3, # 'ס' - 13: 2, # 'ע' - 26: 1, # 'ף' - 18: 2, # 'פ' - 27: 2, # 'ץ' - 21: 3, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 26: { # 'ף' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 1, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 1, # 'ו' - 24: 0, # 'ז' - 14: 1, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 1, # 'כ' - 4: 1, # 'ל' - 11: 0, # 'ם' - 6: 1, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 1, # 'ס' - 13: 0, # 'ע' - 26: 1, # 'ף' - 18: 1, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 1, # 'ק' - 7: 1, # 'ר' - 10: 1, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 18: { # 'פ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 1, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 1, # 'ֵ' - 36: 2, # 'ֶ' - 31: 1, # 'ַ' - 29: 2, # 'ָ' - 35: 1, # 'ֹ' - 62: 1, # 'ֻ' - 28: 2, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 2, # 'ב' - 20: 3, # 'ג' - 16: 2, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 2, # 'ז' - 14: 3, # 'ח' - 22: 3, # 'ט' - 1: 3, # 'י' - 25: 2, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 2, # 'ם' - 6: 2, # 'מ' - 23: 3, # 'ן' - 12: 3, # 'נ' - 19: 3, # 'ס' - 13: 3, # 'ע' - 26: 2, # 'ף' - 18: 2, # 'פ' - 27: 2, # 'ץ' - 21: 3, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 27: { # 'ץ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 1, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 1, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 1, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 1, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 1, # 'ר' - 10: 0, # 'ש' - 5: 1, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 21: { # 'צ' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 1, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 1, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 1, # 'ֹ' - 62: 1, # 'ֻ' - 28: 2, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 2, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 1, # 'ז' - 14: 3, # 'ח' - 22: 2, # 'ט' - 1: 3, # 'י' - 25: 1, # 'ך' - 15: 1, # 'כ' - 4: 3, # 'ל' - 11: 2, # 'ם' - 6: 3, # 'מ' - 23: 2, # 'ן' - 12: 3, # 'נ' - 19: 1, # 'ס' - 13: 3, # 'ע' - 26: 2, # 'ף' - 18: 3, # 'פ' - 27: 2, # 'ץ' - 21: 2, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 0, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 17: { # 'ק' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 1, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 1, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 2, # 'ֹ' - 62: 1, # 'ֻ' - 28: 2, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 2, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 2, # 'ז' - 14: 3, # 'ח' - 22: 3, # 'ט' - 1: 3, # 'י' - 25: 1, # 'ך' - 15: 1, # 'כ' - 4: 3, # 'ל' - 11: 2, # 'ם' - 6: 3, # 'מ' - 23: 2, # 'ן' - 12: 3, # 'נ' - 19: 3, # 'ס' - 13: 3, # 'ע' - 26: 2, # 'ף' - 18: 3, # 'פ' - 27: 2, # 'ץ' - 21: 3, # 'צ' - 17: 2, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 1, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 7: { # 'ר' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 2, # '´' - 48: 1, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 1, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 2, # 'ֹ' - 62: 1, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 3, # 'ז' - 14: 3, # 'ח' - 22: 3, # 'ט' - 1: 3, # 'י' - 25: 3, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 3, # 'ן' - 12: 3, # 'נ' - 19: 3, # 'ס' - 13: 3, # 'ע' - 26: 2, # 'ף' - 18: 3, # 'פ' - 27: 3, # 'ץ' - 21: 3, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 2, # '…' - }, - 10: { # 'ש' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 1, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 1, # 'ִ' - 37: 1, # 'ֵ' - 36: 1, # 'ֶ' - 31: 1, # 'ַ' - 29: 1, # 'ָ' - 35: 1, # 'ֹ' - 62: 1, # 'ֻ' - 28: 2, # 'ּ' - 38: 3, # 'ׁ' - 45: 2, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 3, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 2, # 'ז' - 14: 3, # 'ח' - 22: 3, # 'ט' - 1: 3, # 'י' - 25: 3, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 2, # 'ן' - 12: 3, # 'נ' - 19: 2, # 'ס' - 13: 3, # 'ע' - 26: 2, # 'ף' - 18: 3, # 'פ' - 27: 1, # 'ץ' - 21: 2, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 1, # '…' - }, - 5: { # 'ת' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 1, # '\xa0' - 55: 0, # '´' - 48: 1, # '¼' - 39: 1, # '½' - 57: 0, # '¾' - 30: 2, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 2, # 'ִ' - 37: 2, # 'ֵ' - 36: 2, # 'ֶ' - 31: 2, # 'ַ' - 29: 2, # 'ָ' - 35: 1, # 'ֹ' - 62: 1, # 'ֻ' - 28: 2, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 3, # 'א' - 8: 3, # 'ב' - 20: 3, # 'ג' - 16: 2, # 'ד' - 3: 3, # 'ה' - 2: 3, # 'ו' - 24: 2, # 'ז' - 14: 3, # 'ח' - 22: 2, # 'ט' - 1: 3, # 'י' - 25: 2, # 'ך' - 15: 3, # 'כ' - 4: 3, # 'ל' - 11: 3, # 'ם' - 6: 3, # 'מ' - 23: 3, # 'ן' - 12: 3, # 'נ' - 19: 2, # 'ס' - 13: 3, # 'ע' - 26: 2, # 'ף' - 18: 3, # 'פ' - 27: 1, # 'ץ' - 21: 2, # 'צ' - 17: 3, # 'ק' - 7: 3, # 'ר' - 10: 3, # 'ש' - 5: 3, # 'ת' - 32: 1, # '–' - 52: 1, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 2, # '…' - }, - 32: { # '–' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 1, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 1, # 'א' - 8: 1, # 'ב' - 20: 1, # 'ג' - 16: 1, # 'ד' - 3: 1, # 'ה' - 2: 1, # 'ו' - 24: 0, # 'ז' - 14: 1, # 'ח' - 22: 0, # 'ט' - 1: 1, # 'י' - 25: 0, # 'ך' - 15: 1, # 'כ' - 4: 1, # 'ל' - 11: 0, # 'ם' - 6: 1, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 1, # 'ס' - 13: 1, # 'ע' - 26: 0, # 'ף' - 18: 1, # 'פ' - 27: 0, # 'ץ' - 21: 1, # 'צ' - 17: 0, # 'ק' - 7: 1, # 'ר' - 10: 1, # 'ש' - 5: 1, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 52: { # '’' - 50: 1, # 'a' - 60: 0, # 'c' - 61: 1, # 'd' - 42: 1, # 'e' - 53: 1, # 'i' - 56: 1, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 1, # 'r' - 43: 2, # 's' - 44: 2, # 't' - 63: 1, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 1, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 1, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 1, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 47: { # '“' - 50: 1, # 'a' - 60: 1, # 'c' - 61: 1, # 'd' - 42: 1, # 'e' - 53: 1, # 'i' - 56: 1, # 'l' - 54: 1, # 'n' - 49: 1, # 'o' - 51: 1, # 'r' - 43: 1, # 's' - 44: 1, # 't' - 63: 1, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 2, # 'א' - 8: 1, # 'ב' - 20: 1, # 'ג' - 16: 1, # 'ד' - 3: 1, # 'ה' - 2: 1, # 'ו' - 24: 1, # 'ז' - 14: 1, # 'ח' - 22: 1, # 'ט' - 1: 1, # 'י' - 25: 0, # 'ך' - 15: 1, # 'כ' - 4: 1, # 'ל' - 11: 0, # 'ם' - 6: 1, # 'מ' - 23: 0, # 'ן' - 12: 1, # 'נ' - 19: 1, # 'ס' - 13: 1, # 'ע' - 26: 0, # 'ף' - 18: 1, # 'פ' - 27: 0, # 'ץ' - 21: 1, # 'צ' - 17: 1, # 'ק' - 7: 1, # 'ר' - 10: 1, # 'ש' - 5: 1, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 46: { # '”' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 1, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 1, # 'א' - 8: 1, # 'ב' - 20: 1, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 1, # 'י' - 25: 0, # 'ך' - 15: 1, # 'כ' - 4: 1, # 'ל' - 11: 0, # 'ם' - 6: 1, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 1, # 'צ' - 17: 0, # 'ק' - 7: 1, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 0, # '†' - 40: 0, # '…' - }, - 58: { # '†' - 50: 0, # 'a' - 60: 0, # 'c' - 61: 0, # 'd' - 42: 0, # 'e' - 53: 0, # 'i' - 56: 0, # 'l' - 54: 0, # 'n' - 49: 0, # 'o' - 51: 0, # 'r' - 43: 0, # 's' - 44: 0, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 0, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 0, # 'ה' - 2: 0, # 'ו' - 24: 0, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 0, # 'י' - 25: 0, # 'ך' - 15: 0, # 'כ' - 4: 0, # 'ל' - 11: 0, # 'ם' - 6: 0, # 'מ' - 23: 0, # 'ן' - 12: 0, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 0, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 0, # 'ר' - 10: 0, # 'ש' - 5: 0, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 0, # '”' - 58: 2, # '†' - 40: 0, # '…' - }, - 40: { # '…' - 50: 1, # 'a' - 60: 1, # 'c' - 61: 1, # 'd' - 42: 1, # 'e' - 53: 1, # 'i' - 56: 0, # 'l' - 54: 1, # 'n' - 49: 0, # 'o' - 51: 1, # 'r' - 43: 1, # 's' - 44: 1, # 't' - 63: 0, # 'u' - 34: 0, # '\xa0' - 55: 0, # '´' - 48: 0, # '¼' - 39: 0, # '½' - 57: 0, # '¾' - 30: 0, # 'ְ' - 59: 0, # 'ֱ' - 41: 0, # 'ֲ' - 33: 0, # 'ִ' - 37: 0, # 'ֵ' - 36: 0, # 'ֶ' - 31: 0, # 'ַ' - 29: 0, # 'ָ' - 35: 0, # 'ֹ' - 62: 0, # 'ֻ' - 28: 0, # 'ּ' - 38: 0, # 'ׁ' - 45: 0, # 'ׂ' - 9: 1, # 'א' - 8: 0, # 'ב' - 20: 0, # 'ג' - 16: 0, # 'ד' - 3: 1, # 'ה' - 2: 1, # 'ו' - 24: 1, # 'ז' - 14: 0, # 'ח' - 22: 0, # 'ט' - 1: 1, # 'י' - 25: 0, # 'ך' - 15: 1, # 'כ' - 4: 1, # 'ל' - 11: 0, # 'ם' - 6: 1, # 'מ' - 23: 0, # 'ן' - 12: 1, # 'נ' - 19: 0, # 'ס' - 13: 0, # 'ע' - 26: 0, # 'ף' - 18: 1, # 'פ' - 27: 0, # 'ץ' - 21: 0, # 'צ' - 17: 0, # 'ק' - 7: 1, # 'ר' - 10: 1, # 'ש' - 5: 1, # 'ת' - 32: 0, # '–' - 52: 0, # '’' - 47: 0, # '“' - 46: 1, # '”' - 58: 0, # '†' - 40: 2, # '…' - }, -} - -# 255: Undefined characters that did not exist in training text -# 254: Carriage/Return -# 253: symbol (punctuation) that does not belong to word -# 252: 0 - 9 -# 251: Control characters - -# Character Mapping Table(s): -WINDOWS_1255_HEBREW_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 254, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 254, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 253, # ' ' - 33: 253, # '!' - 34: 253, # '"' - 35: 253, # '#' - 36: 253, # '$' - 37: 253, # '%' - 38: 253, # '&' - 39: 253, # "'" - 40: 253, # '(' - 41: 253, # ')' - 42: 253, # '*' - 43: 253, # '+' - 44: 253, # ',' - 45: 253, # '-' - 46: 253, # '.' - 47: 253, # '/' - 48: 252, # '0' - 49: 252, # '1' - 50: 252, # '2' - 51: 252, # '3' - 52: 252, # '4' - 53: 252, # '5' - 54: 252, # '6' - 55: 252, # '7' - 56: 252, # '8' - 57: 252, # '9' - 58: 253, # ':' - 59: 253, # ';' - 60: 253, # '<' - 61: 253, # '=' - 62: 253, # '>' - 63: 253, # '?' - 64: 253, # '@' - 65: 69, # 'A' - 66: 91, # 'B' - 67: 79, # 'C' - 68: 80, # 'D' - 69: 92, # 'E' - 70: 89, # 'F' - 71: 97, # 'G' - 72: 90, # 'H' - 73: 68, # 'I' - 74: 111, # 'J' - 75: 112, # 'K' - 76: 82, # 'L' - 77: 73, # 'M' - 78: 95, # 'N' - 79: 85, # 'O' - 80: 78, # 'P' - 81: 121, # 'Q' - 82: 86, # 'R' - 83: 71, # 'S' - 84: 67, # 'T' - 85: 102, # 'U' - 86: 107, # 'V' - 87: 84, # 'W' - 88: 114, # 'X' - 89: 103, # 'Y' - 90: 115, # 'Z' - 91: 253, # '[' - 92: 253, # '\\' - 93: 253, # ']' - 94: 253, # '^' - 95: 253, # '_' - 96: 253, # '`' - 97: 50, # 'a' - 98: 74, # 'b' - 99: 60, # 'c' - 100: 61, # 'd' - 101: 42, # 'e' - 102: 76, # 'f' - 103: 70, # 'g' - 104: 64, # 'h' - 105: 53, # 'i' - 106: 105, # 'j' - 107: 93, # 'k' - 108: 56, # 'l' - 109: 65, # 'm' - 110: 54, # 'n' - 111: 49, # 'o' - 112: 66, # 'p' - 113: 110, # 'q' - 114: 51, # 'r' - 115: 43, # 's' - 116: 44, # 't' - 117: 63, # 'u' - 118: 81, # 'v' - 119: 77, # 'w' - 120: 98, # 'x' - 121: 75, # 'y' - 122: 108, # 'z' - 123: 253, # '{' - 124: 253, # '|' - 125: 253, # '}' - 126: 253, # '~' - 127: 253, # '\x7f' - 128: 124, # '€' - 129: 202, # None - 130: 203, # '‚' - 131: 204, # 'ƒ' - 132: 205, # '„' - 133: 40, # '…' - 134: 58, # '†' - 135: 206, # '‡' - 136: 207, # 'ˆ' - 137: 208, # '‰' - 138: 209, # None - 139: 210, # '‹' - 140: 211, # None - 141: 212, # None - 142: 213, # None - 143: 214, # None - 144: 215, # None - 145: 83, # '‘' - 146: 52, # '’' - 147: 47, # '“' - 148: 46, # '”' - 149: 72, # '•' - 150: 32, # '–' - 151: 94, # '—' - 152: 216, # '˜' - 153: 113, # '™' - 154: 217, # None - 155: 109, # '›' - 156: 218, # None - 157: 219, # None - 158: 220, # None - 159: 221, # None - 160: 34, # '\xa0' - 161: 116, # '¡' - 162: 222, # '¢' - 163: 118, # '£' - 164: 100, # '₪' - 165: 223, # '¥' - 166: 224, # '¦' - 167: 117, # '§' - 168: 119, # '¨' - 169: 104, # '©' - 170: 125, # '×' - 171: 225, # '«' - 172: 226, # '¬' - 173: 87, # '\xad' - 174: 99, # '®' - 175: 227, # '¯' - 176: 106, # '°' - 177: 122, # '±' - 178: 123, # '²' - 179: 228, # '³' - 180: 55, # '´' - 181: 229, # 'µ' - 182: 230, # '¶' - 183: 101, # '·' - 184: 231, # '¸' - 185: 232, # '¹' - 186: 120, # '÷' - 187: 233, # '»' - 188: 48, # '¼' - 189: 39, # '½' - 190: 57, # '¾' - 191: 234, # '¿' - 192: 30, # 'ְ' - 193: 59, # 'ֱ' - 194: 41, # 'ֲ' - 195: 88, # 'ֳ' - 196: 33, # 'ִ' - 197: 37, # 'ֵ' - 198: 36, # 'ֶ' - 199: 31, # 'ַ' - 200: 29, # 'ָ' - 201: 35, # 'ֹ' - 202: 235, # None - 203: 62, # 'ֻ' - 204: 28, # 'ּ' - 205: 236, # 'ֽ' - 206: 126, # '־' - 207: 237, # 'ֿ' - 208: 238, # '׀' - 209: 38, # 'ׁ' - 210: 45, # 'ׂ' - 211: 239, # '׃' - 212: 240, # 'װ' - 213: 241, # 'ױ' - 214: 242, # 'ײ' - 215: 243, # '׳' - 216: 127, # '״' - 217: 244, # None - 218: 245, # None - 219: 246, # None - 220: 247, # None - 221: 248, # None - 222: 249, # None - 223: 250, # None - 224: 9, # 'א' - 225: 8, # 'ב' - 226: 20, # 'ג' - 227: 16, # 'ד' - 228: 3, # 'ה' - 229: 2, # 'ו' - 230: 24, # 'ז' - 231: 14, # 'ח' - 232: 22, # 'ט' - 233: 1, # 'י' - 234: 25, # 'ך' - 235: 15, # 'כ' - 236: 4, # 'ל' - 237: 11, # 'ם' - 238: 6, # 'מ' - 239: 23, # 'ן' - 240: 12, # 'נ' - 241: 19, # 'ס' - 242: 13, # 'ע' - 243: 26, # 'ף' - 244: 18, # 'פ' - 245: 27, # 'ץ' - 246: 21, # 'צ' - 247: 17, # 'ק' - 248: 7, # 'ר' - 249: 10, # 'ש' - 250: 5, # 'ת' - 251: 251, # None - 252: 252, # None - 253: 128, # '\u200e' - 254: 96, # '\u200f' - 255: 253, # None -} - -WINDOWS_1255_HEBREW_MODEL = SingleByteCharSetModel( - charset_name="windows-1255", - language="Hebrew", - char_to_order_map=WINDOWS_1255_HEBREW_CHAR_TO_ORDER, - language_model=HEBREW_LANG_MODEL, - typical_positive_ratio=0.984004, - keep_ascii_letters=False, - alphabet="אבגדהוזחטיךכלםמןנסעףפץצקרשתװױײ", -) diff --git a/venv/lib/python3.10/site-packages/chardet/langhungarianmodel.py b/venv/lib/python3.10/site-packages/chardet/langhungarianmodel.py deleted file mode 100644 index bd6630a0513447bb56e1ffbed7aa07e173f62f5b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/langhungarianmodel.py +++ /dev/null @@ -1,4649 +0,0 @@ -from chardet.sbcharsetprober import SingleByteCharSetModel - -# 3: Positive -# 2: Likely -# 1: Unlikely -# 0: Negative - -HUNGARIAN_LANG_MODEL = { - 28: { # 'A' - 28: 0, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 2, # 'D' - 32: 1, # 'E' - 50: 1, # 'F' - 49: 2, # 'G' - 38: 1, # 'H' - 39: 2, # 'I' - 53: 1, # 'J' - 36: 2, # 'K' - 41: 2, # 'L' - 34: 1, # 'M' - 35: 2, # 'N' - 47: 1, # 'O' - 46: 2, # 'P' - 43: 2, # 'R' - 33: 2, # 'S' - 37: 2, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 1, # 'Y' - 52: 2, # 'Z' - 2: 0, # 'a' - 18: 1, # 'b' - 26: 1, # 'c' - 17: 2, # 'd' - 1: 1, # 'e' - 27: 1, # 'f' - 12: 1, # 'g' - 20: 1, # 'h' - 9: 1, # 'i' - 22: 1, # 'j' - 7: 2, # 'k' - 6: 2, # 'l' - 13: 2, # 'm' - 4: 2, # 'n' - 8: 0, # 'o' - 23: 2, # 'p' - 10: 2, # 'r' - 5: 1, # 's' - 3: 1, # 't' - 21: 1, # 'u' - 19: 1, # 'v' - 62: 1, # 'x' - 16: 0, # 'y' - 11: 3, # 'z' - 51: 1, # 'Á' - 44: 0, # 'É' - 61: 1, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 40: { # 'B' - 28: 2, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 2, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 1, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 0, # 'M' - 35: 1, # 'N' - 47: 2, # 'O' - 46: 0, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 2, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 3, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 0, # 'h' - 9: 2, # 'i' - 22: 1, # 'j' - 7: 0, # 'k' - 6: 1, # 'l' - 13: 0, # 'm' - 4: 0, # 'n' - 8: 2, # 'o' - 23: 1, # 'p' - 10: 2, # 'r' - 5: 0, # 's' - 3: 0, # 't' - 21: 3, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 0, # 'z' - 51: 1, # 'Á' - 44: 1, # 'É' - 61: 1, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 2, # 'á' - 15: 2, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 1, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 1, # 'ő' - 56: 1, # 'ű' - }, - 54: { # 'C' - 28: 1, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 1, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 1, # 'H' - 39: 2, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 0, # 'N' - 47: 1, # 'O' - 46: 1, # 'P' - 43: 1, # 'R' - 33: 2, # 'S' - 37: 1, # 'T' - 57: 1, # 'U' - 48: 0, # 'V' - 55: 1, # 'Y' - 52: 1, # 'Z' - 2: 2, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 1, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 1, # 'h' - 9: 1, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 1, # 'l' - 13: 0, # 'm' - 4: 0, # 'n' - 8: 2, # 'o' - 23: 0, # 'p' - 10: 1, # 'r' - 5: 3, # 's' - 3: 0, # 't' - 21: 1, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 1, # 'z' - 51: 1, # 'Á' - 44: 1, # 'É' - 61: 1, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 1, # 'á' - 15: 1, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 45: { # 'D' - 28: 2, # 'A' - 40: 1, # 'B' - 54: 0, # 'C' - 45: 1, # 'D' - 32: 2, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 2, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 0, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 2, # 'O' - 46: 0, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 1, # 'Y' - 52: 1, # 'Z' - 2: 2, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 3, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 0, # 'h' - 9: 1, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 0, # 'l' - 13: 0, # 'm' - 4: 0, # 'n' - 8: 1, # 'o' - 23: 0, # 'p' - 10: 2, # 'r' - 5: 0, # 's' - 3: 0, # 't' - 21: 2, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 1, # 'z' - 51: 1, # 'Á' - 44: 1, # 'É' - 61: 1, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 1, # 'á' - 15: 1, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 1, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 1, # 'ő' - 56: 0, # 'ű' - }, - 32: { # 'E' - 28: 1, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 1, # 'E' - 50: 1, # 'F' - 49: 2, # 'G' - 38: 1, # 'H' - 39: 1, # 'I' - 53: 1, # 'J' - 36: 2, # 'K' - 41: 2, # 'L' - 34: 2, # 'M' - 35: 2, # 'N' - 47: 1, # 'O' - 46: 1, # 'P' - 43: 2, # 'R' - 33: 2, # 'S' - 37: 2, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 1, # 'Y' - 52: 1, # 'Z' - 2: 1, # 'a' - 18: 1, # 'b' - 26: 1, # 'c' - 17: 2, # 'd' - 1: 1, # 'e' - 27: 1, # 'f' - 12: 3, # 'g' - 20: 1, # 'h' - 9: 1, # 'i' - 22: 1, # 'j' - 7: 1, # 'k' - 6: 2, # 'l' - 13: 2, # 'm' - 4: 2, # 'n' - 8: 0, # 'o' - 23: 1, # 'p' - 10: 2, # 'r' - 5: 2, # 's' - 3: 1, # 't' - 21: 2, # 'u' - 19: 1, # 'v' - 62: 1, # 'x' - 16: 0, # 'y' - 11: 3, # 'z' - 51: 1, # 'Á' - 44: 1, # 'É' - 61: 0, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 0, # 'Ú' - 63: 1, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 1, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 50: { # 'F' - 28: 1, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 1, # 'E' - 50: 1, # 'F' - 49: 0, # 'G' - 38: 1, # 'H' - 39: 1, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 1, # 'O' - 46: 0, # 'P' - 43: 1, # 'R' - 33: 0, # 'S' - 37: 1, # 'T' - 57: 1, # 'U' - 48: 0, # 'V' - 55: 1, # 'Y' - 52: 0, # 'Z' - 2: 2, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 2, # 'e' - 27: 1, # 'f' - 12: 0, # 'g' - 20: 0, # 'h' - 9: 2, # 'i' - 22: 1, # 'j' - 7: 0, # 'k' - 6: 1, # 'l' - 13: 0, # 'm' - 4: 0, # 'n' - 8: 2, # 'o' - 23: 0, # 'p' - 10: 2, # 'r' - 5: 0, # 's' - 3: 0, # 't' - 21: 1, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 0, # 'z' - 51: 1, # 'Á' - 44: 1, # 'É' - 61: 0, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 0, # 'Ú' - 63: 1, # 'Ü' - 14: 1, # 'á' - 15: 1, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 2, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 1, # 'ő' - 56: 1, # 'ű' - }, - 49: { # 'G' - 28: 2, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 2, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 1, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 1, # 'O' - 46: 1, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 2, # 'Y' - 52: 1, # 'Z' - 2: 2, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 2, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 0, # 'h' - 9: 1, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 1, # 'l' - 13: 0, # 'm' - 4: 0, # 'n' - 8: 2, # 'o' - 23: 0, # 'p' - 10: 2, # 'r' - 5: 0, # 's' - 3: 0, # 't' - 21: 1, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 2, # 'y' - 11: 0, # 'z' - 51: 1, # 'Á' - 44: 1, # 'É' - 61: 1, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 1, # 'á' - 15: 1, # 'é' - 30: 0, # 'í' - 25: 1, # 'ó' - 24: 1, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 1, # 'ő' - 56: 0, # 'ű' - }, - 38: { # 'H' - 28: 2, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 0, # 'D' - 32: 1, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 1, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 1, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 1, # 'O' - 46: 0, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 1, # 'U' - 48: 0, # 'V' - 55: 1, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 2, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 0, # 'h' - 9: 2, # 'i' - 22: 1, # 'j' - 7: 0, # 'k' - 6: 1, # 'l' - 13: 1, # 'm' - 4: 0, # 'n' - 8: 3, # 'o' - 23: 0, # 'p' - 10: 1, # 'r' - 5: 0, # 's' - 3: 0, # 't' - 21: 2, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 0, # 'z' - 51: 2, # 'Á' - 44: 2, # 'É' - 61: 1, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 2, # 'á' - 15: 1, # 'é' - 30: 2, # 'í' - 25: 1, # 'ó' - 24: 1, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 1, # 'ő' - 56: 1, # 'ű' - }, - 39: { # 'I' - 28: 2, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 1, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 2, # 'I' - 53: 1, # 'J' - 36: 2, # 'K' - 41: 2, # 'L' - 34: 1, # 'M' - 35: 2, # 'N' - 47: 1, # 'O' - 46: 1, # 'P' - 43: 1, # 'R' - 33: 2, # 'S' - 37: 1, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 0, # 'Y' - 52: 2, # 'Z' - 2: 0, # 'a' - 18: 1, # 'b' - 26: 1, # 'c' - 17: 2, # 'd' - 1: 0, # 'e' - 27: 1, # 'f' - 12: 2, # 'g' - 20: 1, # 'h' - 9: 0, # 'i' - 22: 1, # 'j' - 7: 1, # 'k' - 6: 2, # 'l' - 13: 2, # 'm' - 4: 1, # 'n' - 8: 0, # 'o' - 23: 1, # 'p' - 10: 2, # 'r' - 5: 2, # 's' - 3: 2, # 't' - 21: 0, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 1, # 'z' - 51: 1, # 'Á' - 44: 1, # 'É' - 61: 0, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 53: { # 'J' - 28: 2, # 'A' - 40: 0, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 2, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 1, # 'H' - 39: 1, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 1, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 1, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 1, # 'Z' - 2: 2, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 2, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 0, # 'h' - 9: 1, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 0, # 'l' - 13: 0, # 'm' - 4: 0, # 'n' - 8: 1, # 'o' - 23: 0, # 'p' - 10: 0, # 'r' - 5: 0, # 's' - 3: 0, # 't' - 21: 2, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 0, # 'z' - 51: 1, # 'Á' - 44: 1, # 'É' - 61: 0, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 2, # 'á' - 15: 1, # 'é' - 30: 0, # 'í' - 25: 2, # 'ó' - 24: 2, # 'ö' - 31: 1, # 'ú' - 29: 0, # 'ü' - 42: 1, # 'ő' - 56: 0, # 'ű' - }, - 36: { # 'K' - 28: 2, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 2, # 'E' - 50: 1, # 'F' - 49: 0, # 'G' - 38: 1, # 'H' - 39: 2, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 2, # 'O' - 46: 0, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 1, # 'Y' - 52: 0, # 'Z' - 2: 2, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 2, # 'e' - 27: 1, # 'f' - 12: 0, # 'g' - 20: 1, # 'h' - 9: 3, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 1, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 8: 2, # 'o' - 23: 0, # 'p' - 10: 2, # 'r' - 5: 0, # 's' - 3: 0, # 't' - 21: 1, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 0, # 'z' - 51: 1, # 'Á' - 44: 1, # 'É' - 61: 1, # 'Í' - 58: 1, # 'Ó' - 59: 2, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 2, # 'á' - 15: 2, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 2, # 'ö' - 31: 1, # 'ú' - 29: 2, # 'ü' - 42: 1, # 'ő' - 56: 0, # 'ű' - }, - 41: { # 'L' - 28: 2, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 2, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 2, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 2, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 2, # 'O' - 46: 0, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 2, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 1, # 'Y' - 52: 1, # 'Z' - 2: 2, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 3, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 0, # 'h' - 9: 2, # 'i' - 22: 1, # 'j' - 7: 0, # 'k' - 6: 1, # 'l' - 13: 0, # 'm' - 4: 0, # 'n' - 8: 2, # 'o' - 23: 0, # 'p' - 10: 0, # 'r' - 5: 0, # 's' - 3: 0, # 't' - 21: 2, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 0, # 'z' - 51: 2, # 'Á' - 44: 1, # 'É' - 61: 1, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 2, # 'á' - 15: 1, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 1, # 'ö' - 31: 0, # 'ú' - 29: 1, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 34: { # 'M' - 28: 2, # 'A' - 40: 1, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 2, # 'E' - 50: 1, # 'F' - 49: 0, # 'G' - 38: 1, # 'H' - 39: 2, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 1, # 'O' - 46: 1, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 1, # 'Y' - 52: 1, # 'Z' - 2: 3, # 'a' - 18: 0, # 'b' - 26: 1, # 'c' - 17: 0, # 'd' - 1: 3, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 0, # 'h' - 9: 3, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 0, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 8: 3, # 'o' - 23: 0, # 'p' - 10: 1, # 'r' - 5: 0, # 's' - 3: 0, # 't' - 21: 2, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 0, # 'z' - 51: 2, # 'Á' - 44: 1, # 'É' - 61: 1, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 2, # 'á' - 15: 2, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 1, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 0, # 'ő' - 56: 1, # 'ű' - }, - 35: { # 'N' - 28: 2, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 2, # 'D' - 32: 2, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 1, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 1, # 'O' - 46: 1, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 2, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 2, # 'Y' - 52: 1, # 'Z' - 2: 3, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 3, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 0, # 'h' - 9: 2, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 0, # 'l' - 13: 0, # 'm' - 4: 1, # 'n' - 8: 2, # 'o' - 23: 0, # 'p' - 10: 0, # 'r' - 5: 0, # 's' - 3: 0, # 't' - 21: 1, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 2, # 'y' - 11: 0, # 'z' - 51: 1, # 'Á' - 44: 1, # 'É' - 61: 1, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 1, # 'á' - 15: 2, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 1, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 1, # 'ő' - 56: 0, # 'ű' - }, - 47: { # 'O' - 28: 1, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 1, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 1, # 'I' - 53: 1, # 'J' - 36: 2, # 'K' - 41: 2, # 'L' - 34: 2, # 'M' - 35: 2, # 'N' - 47: 1, # 'O' - 46: 1, # 'P' - 43: 2, # 'R' - 33: 2, # 'S' - 37: 2, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 1, # 'Y' - 52: 1, # 'Z' - 2: 0, # 'a' - 18: 1, # 'b' - 26: 1, # 'c' - 17: 1, # 'd' - 1: 1, # 'e' - 27: 1, # 'f' - 12: 1, # 'g' - 20: 1, # 'h' - 9: 1, # 'i' - 22: 1, # 'j' - 7: 2, # 'k' - 6: 2, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 8: 1, # 'o' - 23: 1, # 'p' - 10: 2, # 'r' - 5: 1, # 's' - 3: 2, # 't' - 21: 1, # 'u' - 19: 0, # 'v' - 62: 1, # 'x' - 16: 0, # 'y' - 11: 1, # 'z' - 51: 1, # 'Á' - 44: 1, # 'É' - 61: 0, # 'Í' - 58: 1, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 46: { # 'P' - 28: 1, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 1, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 1, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 0, # 'M' - 35: 1, # 'N' - 47: 1, # 'O' - 46: 1, # 'P' - 43: 2, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 0, # 'Y' - 52: 1, # 'Z' - 2: 2, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 2, # 'e' - 27: 1, # 'f' - 12: 0, # 'g' - 20: 1, # 'h' - 9: 2, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 1, # 'l' - 13: 0, # 'm' - 4: 1, # 'n' - 8: 2, # 'o' - 23: 0, # 'p' - 10: 2, # 'r' - 5: 1, # 's' - 3: 0, # 't' - 21: 1, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 0, # 'z' - 51: 2, # 'Á' - 44: 1, # 'É' - 61: 1, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 0, # 'Ú' - 63: 1, # 'Ü' - 14: 3, # 'á' - 15: 2, # 'é' - 30: 0, # 'í' - 25: 1, # 'ó' - 24: 1, # 'ö' - 31: 0, # 'ú' - 29: 1, # 'ü' - 42: 1, # 'ő' - 56: 0, # 'ű' - }, - 43: { # 'R' - 28: 2, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 2, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 2, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 2, # 'O' - 46: 1, # 'P' - 43: 1, # 'R' - 33: 2, # 'S' - 37: 2, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 1, # 'Y' - 52: 1, # 'Z' - 2: 2, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 2, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 1, # 'h' - 9: 2, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 0, # 'l' - 13: 0, # 'm' - 4: 0, # 'n' - 8: 2, # 'o' - 23: 0, # 'p' - 10: 0, # 'r' - 5: 0, # 's' - 3: 0, # 't' - 21: 1, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 0, # 'z' - 51: 2, # 'Á' - 44: 1, # 'É' - 61: 1, # 'Í' - 58: 2, # 'Ó' - 59: 1, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 2, # 'á' - 15: 2, # 'é' - 30: 1, # 'í' - 25: 2, # 'ó' - 24: 1, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 33: { # 'S' - 28: 2, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 2, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 2, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 2, # 'O' - 46: 1, # 'P' - 43: 1, # 'R' - 33: 2, # 'S' - 37: 2, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 1, # 'Y' - 52: 3, # 'Z' - 2: 2, # 'a' - 18: 0, # 'b' - 26: 1, # 'c' - 17: 0, # 'd' - 1: 2, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 1, # 'h' - 9: 2, # 'i' - 22: 0, # 'j' - 7: 1, # 'k' - 6: 1, # 'l' - 13: 1, # 'm' - 4: 0, # 'n' - 8: 2, # 'o' - 23: 1, # 'p' - 10: 0, # 'r' - 5: 0, # 's' - 3: 1, # 't' - 21: 1, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 3, # 'z' - 51: 2, # 'Á' - 44: 1, # 'É' - 61: 1, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 2, # 'á' - 15: 1, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 1, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 1, # 'ő' - 56: 1, # 'ű' - }, - 37: { # 'T' - 28: 2, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 2, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 2, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 2, # 'O' - 46: 1, # 'P' - 43: 2, # 'R' - 33: 1, # 'S' - 37: 2, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 1, # 'Y' - 52: 1, # 'Z' - 2: 2, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 2, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 1, # 'h' - 9: 2, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 0, # 'l' - 13: 0, # 'm' - 4: 0, # 'n' - 8: 2, # 'o' - 23: 0, # 'p' - 10: 1, # 'r' - 5: 1, # 's' - 3: 0, # 't' - 21: 2, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 1, # 'z' - 51: 2, # 'Á' - 44: 2, # 'É' - 61: 1, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 2, # 'á' - 15: 1, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 2, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 1, # 'ő' - 56: 1, # 'ű' - }, - 57: { # 'U' - 28: 1, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 1, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 1, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 1, # 'O' - 46: 1, # 'P' - 43: 1, # 'R' - 33: 2, # 'S' - 37: 1, # 'T' - 57: 0, # 'U' - 48: 1, # 'V' - 55: 0, # 'Y' - 52: 1, # 'Z' - 2: 0, # 'a' - 18: 1, # 'b' - 26: 1, # 'c' - 17: 1, # 'd' - 1: 1, # 'e' - 27: 0, # 'f' - 12: 2, # 'g' - 20: 0, # 'h' - 9: 0, # 'i' - 22: 1, # 'j' - 7: 1, # 'k' - 6: 1, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 8: 0, # 'o' - 23: 1, # 'p' - 10: 1, # 'r' - 5: 1, # 's' - 3: 1, # 't' - 21: 0, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 1, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 1, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 48: { # 'V' - 28: 2, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 1, # 'D' - 32: 2, # 'E' - 50: 1, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 2, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 0, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 1, # 'O' - 46: 1, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 1, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 2, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 0, # 'h' - 9: 2, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 1, # 'l' - 13: 0, # 'm' - 4: 0, # 'n' - 8: 2, # 'o' - 23: 0, # 'p' - 10: 0, # 'r' - 5: 0, # 's' - 3: 0, # 't' - 21: 1, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 0, # 'z' - 51: 2, # 'Á' - 44: 2, # 'É' - 61: 1, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 0, # 'Ú' - 63: 1, # 'Ü' - 14: 2, # 'á' - 15: 2, # 'é' - 30: 1, # 'í' - 25: 0, # 'ó' - 24: 1, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 55: { # 'Y' - 28: 2, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 2, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 1, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 1, # 'O' - 46: 1, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 0, # 'Y' - 52: 2, # 'Z' - 2: 1, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 1, # 'd' - 1: 1, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 0, # 'h' - 9: 0, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 0, # 'l' - 13: 0, # 'm' - 4: 0, # 'n' - 8: 1, # 'o' - 23: 1, # 'p' - 10: 0, # 'r' - 5: 0, # 's' - 3: 0, # 't' - 21: 0, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 0, # 'z' - 51: 1, # 'Á' - 44: 1, # 'É' - 61: 1, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 52: { # 'Z' - 28: 2, # 'A' - 40: 1, # 'B' - 54: 0, # 'C' - 45: 1, # 'D' - 32: 2, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 2, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 2, # 'O' - 46: 1, # 'P' - 43: 1, # 'R' - 33: 2, # 'S' - 37: 1, # 'T' - 57: 1, # 'U' - 48: 1, # 'V' - 55: 1, # 'Y' - 52: 1, # 'Z' - 2: 1, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 1, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 0, # 'h' - 9: 1, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 0, # 'l' - 13: 0, # 'm' - 4: 1, # 'n' - 8: 1, # 'o' - 23: 0, # 'p' - 10: 1, # 'r' - 5: 2, # 's' - 3: 0, # 't' - 21: 1, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 0, # 'z' - 51: 2, # 'Á' - 44: 1, # 'É' - 61: 1, # 'Í' - 58: 1, # 'Ó' - 59: 1, # 'Ö' - 60: 1, # 'Ú' - 63: 1, # 'Ü' - 14: 1, # 'á' - 15: 1, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 1, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 2: { # 'a' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 1, # 'a' - 18: 3, # 'b' - 26: 3, # 'c' - 17: 3, # 'd' - 1: 2, # 'e' - 27: 2, # 'f' - 12: 3, # 'g' - 20: 3, # 'h' - 9: 3, # 'i' - 22: 3, # 'j' - 7: 3, # 'k' - 6: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 8: 2, # 'o' - 23: 3, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 3, # 'u' - 19: 3, # 'v' - 62: 1, # 'x' - 16: 2, # 'y' - 11: 3, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 1, # 'á' - 15: 1, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 1, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 18: { # 'b' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 3, # 'b' - 26: 1, # 'c' - 17: 1, # 'd' - 1: 3, # 'e' - 27: 1, # 'f' - 12: 1, # 'g' - 20: 1, # 'h' - 9: 3, # 'i' - 22: 2, # 'j' - 7: 2, # 'k' - 6: 2, # 'l' - 13: 1, # 'm' - 4: 2, # 'n' - 8: 3, # 'o' - 23: 1, # 'p' - 10: 3, # 'r' - 5: 2, # 's' - 3: 1, # 't' - 21: 3, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 1, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 2, # 'í' - 25: 3, # 'ó' - 24: 2, # 'ö' - 31: 2, # 'ú' - 29: 2, # 'ü' - 42: 2, # 'ő' - 56: 1, # 'ű' - }, - 26: { # 'c' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 1, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 1, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 2, # 'a' - 18: 1, # 'b' - 26: 2, # 'c' - 17: 1, # 'd' - 1: 3, # 'e' - 27: 1, # 'f' - 12: 1, # 'g' - 20: 3, # 'h' - 9: 3, # 'i' - 22: 1, # 'j' - 7: 2, # 'k' - 6: 1, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 8: 3, # 'o' - 23: 1, # 'p' - 10: 2, # 'r' - 5: 3, # 's' - 3: 2, # 't' - 21: 2, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 2, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 2, # 'á' - 15: 2, # 'é' - 30: 2, # 'í' - 25: 1, # 'ó' - 24: 1, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 17: { # 'd' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 2, # 'b' - 26: 1, # 'c' - 17: 2, # 'd' - 1: 3, # 'e' - 27: 1, # 'f' - 12: 1, # 'g' - 20: 2, # 'h' - 9: 3, # 'i' - 22: 3, # 'j' - 7: 2, # 'k' - 6: 1, # 'l' - 13: 2, # 'm' - 4: 3, # 'n' - 8: 3, # 'o' - 23: 1, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 3, # 'u' - 19: 3, # 'v' - 62: 0, # 'x' - 16: 2, # 'y' - 11: 2, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 3, # 'í' - 25: 3, # 'ó' - 24: 3, # 'ö' - 31: 2, # 'ú' - 29: 2, # 'ü' - 42: 2, # 'ő' - 56: 1, # 'ű' - }, - 1: { # 'e' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 2, # 'a' - 18: 3, # 'b' - 26: 3, # 'c' - 17: 3, # 'd' - 1: 2, # 'e' - 27: 3, # 'f' - 12: 3, # 'g' - 20: 3, # 'h' - 9: 3, # 'i' - 22: 3, # 'j' - 7: 3, # 'k' - 6: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 8: 2, # 'o' - 23: 3, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 2, # 'u' - 19: 3, # 'v' - 62: 2, # 'x' - 16: 2, # 'y' - 11: 3, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 1, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 1, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 27: { # 'f' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 1, # 'b' - 26: 1, # 'c' - 17: 1, # 'd' - 1: 3, # 'e' - 27: 2, # 'f' - 12: 1, # 'g' - 20: 1, # 'h' - 9: 3, # 'i' - 22: 2, # 'j' - 7: 1, # 'k' - 6: 1, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 8: 3, # 'o' - 23: 0, # 'p' - 10: 3, # 'r' - 5: 1, # 's' - 3: 1, # 't' - 21: 2, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 0, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 3, # 'ö' - 31: 1, # 'ú' - 29: 2, # 'ü' - 42: 1, # 'ő' - 56: 1, # 'ű' - }, - 12: { # 'g' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 3, # 'b' - 26: 2, # 'c' - 17: 2, # 'd' - 1: 3, # 'e' - 27: 2, # 'f' - 12: 3, # 'g' - 20: 3, # 'h' - 9: 3, # 'i' - 22: 3, # 'j' - 7: 2, # 'k' - 6: 3, # 'l' - 13: 2, # 'm' - 4: 3, # 'n' - 8: 3, # 'o' - 23: 1, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 3, # 'u' - 19: 3, # 'v' - 62: 0, # 'x' - 16: 3, # 'y' - 11: 2, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 2, # 'í' - 25: 3, # 'ó' - 24: 2, # 'ö' - 31: 2, # 'ú' - 29: 2, # 'ü' - 42: 2, # 'ő' - 56: 1, # 'ű' - }, - 20: { # 'h' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 1, # 'b' - 26: 1, # 'c' - 17: 0, # 'd' - 1: 3, # 'e' - 27: 0, # 'f' - 12: 1, # 'g' - 20: 2, # 'h' - 9: 3, # 'i' - 22: 1, # 'j' - 7: 1, # 'k' - 6: 1, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 8: 3, # 'o' - 23: 0, # 'p' - 10: 1, # 'r' - 5: 2, # 's' - 3: 1, # 't' - 21: 3, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 2, # 'y' - 11: 0, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 3, # 'í' - 25: 2, # 'ó' - 24: 2, # 'ö' - 31: 2, # 'ú' - 29: 1, # 'ü' - 42: 1, # 'ő' - 56: 1, # 'ű' - }, - 9: { # 'i' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 3, # 'b' - 26: 3, # 'c' - 17: 3, # 'd' - 1: 3, # 'e' - 27: 3, # 'f' - 12: 3, # 'g' - 20: 3, # 'h' - 9: 2, # 'i' - 22: 2, # 'j' - 7: 3, # 'k' - 6: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 8: 2, # 'o' - 23: 2, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 3, # 'u' - 19: 3, # 'v' - 62: 1, # 'x' - 16: 1, # 'y' - 11: 3, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 2, # 'é' - 30: 1, # 'í' - 25: 3, # 'ó' - 24: 1, # 'ö' - 31: 2, # 'ú' - 29: 1, # 'ü' - 42: 0, # 'ő' - 56: 1, # 'ű' - }, - 22: { # 'j' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 2, # 'b' - 26: 1, # 'c' - 17: 3, # 'd' - 1: 3, # 'e' - 27: 1, # 'f' - 12: 1, # 'g' - 20: 2, # 'h' - 9: 1, # 'i' - 22: 2, # 'j' - 7: 2, # 'k' - 6: 2, # 'l' - 13: 1, # 'm' - 4: 2, # 'n' - 8: 3, # 'o' - 23: 1, # 'p' - 10: 2, # 'r' - 5: 2, # 's' - 3: 3, # 't' - 21: 3, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 2, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 1, # 'í' - 25: 3, # 'ó' - 24: 3, # 'ö' - 31: 3, # 'ú' - 29: 2, # 'ü' - 42: 1, # 'ő' - 56: 1, # 'ű' - }, - 7: { # 'k' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 3, # 'b' - 26: 2, # 'c' - 17: 1, # 'd' - 1: 3, # 'e' - 27: 1, # 'f' - 12: 1, # 'g' - 20: 2, # 'h' - 9: 3, # 'i' - 22: 2, # 'j' - 7: 3, # 'k' - 6: 3, # 'l' - 13: 1, # 'm' - 4: 3, # 'n' - 8: 3, # 'o' - 23: 1, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 3, # 'u' - 19: 2, # 'v' - 62: 0, # 'x' - 16: 2, # 'y' - 11: 1, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 3, # 'í' - 25: 2, # 'ó' - 24: 3, # 'ö' - 31: 1, # 'ú' - 29: 3, # 'ü' - 42: 1, # 'ő' - 56: 1, # 'ű' - }, - 6: { # 'l' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 1, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 1, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 2, # 'b' - 26: 3, # 'c' - 17: 3, # 'd' - 1: 3, # 'e' - 27: 3, # 'f' - 12: 3, # 'g' - 20: 3, # 'h' - 9: 3, # 'i' - 22: 3, # 'j' - 7: 3, # 'k' - 6: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 8: 3, # 'o' - 23: 2, # 'p' - 10: 2, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 3, # 'u' - 19: 3, # 'v' - 62: 0, # 'x' - 16: 3, # 'y' - 11: 2, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 3, # 'í' - 25: 3, # 'ó' - 24: 3, # 'ö' - 31: 2, # 'ú' - 29: 2, # 'ü' - 42: 3, # 'ő' - 56: 1, # 'ű' - }, - 13: { # 'm' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 3, # 'b' - 26: 2, # 'c' - 17: 1, # 'd' - 1: 3, # 'e' - 27: 1, # 'f' - 12: 1, # 'g' - 20: 2, # 'h' - 9: 3, # 'i' - 22: 2, # 'j' - 7: 1, # 'k' - 6: 3, # 'l' - 13: 3, # 'm' - 4: 2, # 'n' - 8: 3, # 'o' - 23: 3, # 'p' - 10: 2, # 'r' - 5: 2, # 's' - 3: 2, # 't' - 21: 3, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 2, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 2, # 'í' - 25: 2, # 'ó' - 24: 2, # 'ö' - 31: 2, # 'ú' - 29: 2, # 'ü' - 42: 1, # 'ő' - 56: 2, # 'ű' - }, - 4: { # 'n' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 3, # 'b' - 26: 3, # 'c' - 17: 3, # 'd' - 1: 3, # 'e' - 27: 2, # 'f' - 12: 3, # 'g' - 20: 3, # 'h' - 9: 3, # 'i' - 22: 2, # 'j' - 7: 3, # 'k' - 6: 2, # 'l' - 13: 2, # 'm' - 4: 3, # 'n' - 8: 3, # 'o' - 23: 2, # 'p' - 10: 2, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 3, # 'u' - 19: 2, # 'v' - 62: 1, # 'x' - 16: 3, # 'y' - 11: 3, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 2, # 'í' - 25: 2, # 'ó' - 24: 3, # 'ö' - 31: 2, # 'ú' - 29: 3, # 'ü' - 42: 2, # 'ő' - 56: 1, # 'ű' - }, - 8: { # 'o' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 1, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 2, # 'a' - 18: 3, # 'b' - 26: 3, # 'c' - 17: 3, # 'd' - 1: 2, # 'e' - 27: 2, # 'f' - 12: 3, # 'g' - 20: 3, # 'h' - 9: 2, # 'i' - 22: 2, # 'j' - 7: 3, # 'k' - 6: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 8: 1, # 'o' - 23: 3, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 2, # 'u' - 19: 3, # 'v' - 62: 1, # 'x' - 16: 1, # 'y' - 11: 3, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 1, # 'á' - 15: 2, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 1, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 23: { # 'p' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 1, # 'b' - 26: 2, # 'c' - 17: 1, # 'd' - 1: 3, # 'e' - 27: 1, # 'f' - 12: 1, # 'g' - 20: 2, # 'h' - 9: 3, # 'i' - 22: 2, # 'j' - 7: 2, # 'k' - 6: 3, # 'l' - 13: 1, # 'm' - 4: 2, # 'n' - 8: 3, # 'o' - 23: 3, # 'p' - 10: 3, # 'r' - 5: 2, # 's' - 3: 2, # 't' - 21: 3, # 'u' - 19: 2, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 2, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 2, # 'í' - 25: 2, # 'ó' - 24: 2, # 'ö' - 31: 1, # 'ú' - 29: 2, # 'ü' - 42: 1, # 'ő' - 56: 1, # 'ű' - }, - 10: { # 'r' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 3, # 'b' - 26: 3, # 'c' - 17: 3, # 'd' - 1: 3, # 'e' - 27: 2, # 'f' - 12: 3, # 'g' - 20: 2, # 'h' - 9: 3, # 'i' - 22: 3, # 'j' - 7: 3, # 'k' - 6: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 8: 3, # 'o' - 23: 2, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 3, # 'u' - 19: 3, # 'v' - 62: 1, # 'x' - 16: 2, # 'y' - 11: 3, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 2, # 'í' - 25: 3, # 'ó' - 24: 3, # 'ö' - 31: 3, # 'ú' - 29: 3, # 'ü' - 42: 2, # 'ő' - 56: 2, # 'ű' - }, - 5: { # 's' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 3, # 'b' - 26: 2, # 'c' - 17: 2, # 'd' - 1: 3, # 'e' - 27: 2, # 'f' - 12: 2, # 'g' - 20: 2, # 'h' - 9: 3, # 'i' - 22: 1, # 'j' - 7: 3, # 'k' - 6: 2, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 8: 3, # 'o' - 23: 2, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 3, # 'u' - 19: 2, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 3, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 3, # 'í' - 25: 3, # 'ó' - 24: 3, # 'ö' - 31: 3, # 'ú' - 29: 3, # 'ü' - 42: 2, # 'ő' - 56: 1, # 'ű' - }, - 3: { # 't' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 3, # 'b' - 26: 2, # 'c' - 17: 1, # 'd' - 1: 3, # 'e' - 27: 2, # 'f' - 12: 1, # 'g' - 20: 3, # 'h' - 9: 3, # 'i' - 22: 3, # 'j' - 7: 3, # 'k' - 6: 3, # 'l' - 13: 2, # 'm' - 4: 3, # 'n' - 8: 3, # 'o' - 23: 1, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 3, # 'u' - 19: 3, # 'v' - 62: 0, # 'x' - 16: 3, # 'y' - 11: 1, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 2, # 'í' - 25: 3, # 'ó' - 24: 3, # 'ö' - 31: 3, # 'ú' - 29: 3, # 'ü' - 42: 3, # 'ő' - 56: 2, # 'ű' - }, - 21: { # 'u' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 1, # 'a' - 18: 2, # 'b' - 26: 2, # 'c' - 17: 3, # 'd' - 1: 2, # 'e' - 27: 1, # 'f' - 12: 3, # 'g' - 20: 2, # 'h' - 9: 2, # 'i' - 22: 2, # 'j' - 7: 3, # 'k' - 6: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 8: 1, # 'o' - 23: 2, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 1, # 'u' - 19: 3, # 'v' - 62: 1, # 'x' - 16: 1, # 'y' - 11: 2, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 2, # 'á' - 15: 1, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 0, # 'ö' - 31: 1, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 19: { # 'v' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 2, # 'b' - 26: 1, # 'c' - 17: 1, # 'd' - 1: 3, # 'e' - 27: 1, # 'f' - 12: 1, # 'g' - 20: 1, # 'h' - 9: 3, # 'i' - 22: 1, # 'j' - 7: 1, # 'k' - 6: 1, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 8: 3, # 'o' - 23: 1, # 'p' - 10: 1, # 'r' - 5: 2, # 's' - 3: 2, # 't' - 21: 2, # 'u' - 19: 2, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 1, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 2, # 'í' - 25: 2, # 'ó' - 24: 2, # 'ö' - 31: 1, # 'ú' - 29: 2, # 'ü' - 42: 1, # 'ő' - 56: 1, # 'ű' - }, - 62: { # 'x' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 1, # 'a' - 18: 1, # 'b' - 26: 1, # 'c' - 17: 0, # 'd' - 1: 1, # 'e' - 27: 1, # 'f' - 12: 0, # 'g' - 20: 0, # 'h' - 9: 1, # 'i' - 22: 0, # 'j' - 7: 1, # 'k' - 6: 1, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 8: 1, # 'o' - 23: 1, # 'p' - 10: 1, # 'r' - 5: 1, # 's' - 3: 1, # 't' - 21: 1, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 0, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 1, # 'á' - 15: 1, # 'é' - 30: 1, # 'í' - 25: 1, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 16: { # 'y' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 2, # 'b' - 26: 1, # 'c' - 17: 1, # 'd' - 1: 3, # 'e' - 27: 2, # 'f' - 12: 2, # 'g' - 20: 2, # 'h' - 9: 3, # 'i' - 22: 2, # 'j' - 7: 2, # 'k' - 6: 2, # 'l' - 13: 2, # 'm' - 4: 3, # 'n' - 8: 3, # 'o' - 23: 2, # 'p' - 10: 2, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 3, # 'u' - 19: 3, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 2, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 2, # 'í' - 25: 2, # 'ó' - 24: 3, # 'ö' - 31: 2, # 'ú' - 29: 2, # 'ü' - 42: 1, # 'ő' - 56: 2, # 'ű' - }, - 11: { # 'z' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 3, # 'a' - 18: 2, # 'b' - 26: 1, # 'c' - 17: 3, # 'd' - 1: 3, # 'e' - 27: 1, # 'f' - 12: 2, # 'g' - 20: 2, # 'h' - 9: 3, # 'i' - 22: 1, # 'j' - 7: 3, # 'k' - 6: 2, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 8: 3, # 'o' - 23: 1, # 'p' - 10: 2, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 3, # 'u' - 19: 2, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 3, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 3, # 'á' - 15: 3, # 'é' - 30: 3, # 'í' - 25: 3, # 'ó' - 24: 3, # 'ö' - 31: 2, # 'ú' - 29: 3, # 'ü' - 42: 2, # 'ő' - 56: 1, # 'ű' - }, - 51: { # 'Á' - 28: 0, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 0, # 'E' - 50: 1, # 'F' - 49: 2, # 'G' - 38: 1, # 'H' - 39: 1, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 2, # 'L' - 34: 1, # 'M' - 35: 2, # 'N' - 47: 0, # 'O' - 46: 1, # 'P' - 43: 2, # 'R' - 33: 2, # 'S' - 37: 1, # 'T' - 57: 0, # 'U' - 48: 1, # 'V' - 55: 0, # 'Y' - 52: 1, # 'Z' - 2: 0, # 'a' - 18: 1, # 'b' - 26: 1, # 'c' - 17: 1, # 'd' - 1: 0, # 'e' - 27: 0, # 'f' - 12: 1, # 'g' - 20: 1, # 'h' - 9: 0, # 'i' - 22: 1, # 'j' - 7: 1, # 'k' - 6: 2, # 'l' - 13: 2, # 'm' - 4: 0, # 'n' - 8: 0, # 'o' - 23: 1, # 'p' - 10: 1, # 'r' - 5: 1, # 's' - 3: 1, # 't' - 21: 0, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 1, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 1, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 44: { # 'É' - 28: 0, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 1, # 'E' - 50: 0, # 'F' - 49: 2, # 'G' - 38: 1, # 'H' - 39: 1, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 2, # 'L' - 34: 1, # 'M' - 35: 2, # 'N' - 47: 0, # 'O' - 46: 1, # 'P' - 43: 2, # 'R' - 33: 2, # 'S' - 37: 2, # 'T' - 57: 0, # 'U' - 48: 1, # 'V' - 55: 0, # 'Y' - 52: 1, # 'Z' - 2: 0, # 'a' - 18: 1, # 'b' - 26: 1, # 'c' - 17: 1, # 'd' - 1: 0, # 'e' - 27: 0, # 'f' - 12: 1, # 'g' - 20: 1, # 'h' - 9: 0, # 'i' - 22: 1, # 'j' - 7: 1, # 'k' - 6: 2, # 'l' - 13: 1, # 'm' - 4: 2, # 'n' - 8: 0, # 'o' - 23: 1, # 'p' - 10: 2, # 'r' - 5: 3, # 's' - 3: 1, # 't' - 21: 0, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 0, # 'z' - 51: 0, # 'Á' - 44: 1, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 61: { # 'Í' - 28: 0, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 0, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 1, # 'J' - 36: 0, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 0, # 'O' - 46: 1, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 0, # 'U' - 48: 1, # 'V' - 55: 0, # 'Y' - 52: 1, # 'Z' - 2: 0, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 0, # 'e' - 27: 0, # 'f' - 12: 2, # 'g' - 20: 0, # 'h' - 9: 0, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 0, # 'l' - 13: 1, # 'm' - 4: 0, # 'n' - 8: 0, # 'o' - 23: 0, # 'p' - 10: 1, # 'r' - 5: 0, # 's' - 3: 1, # 't' - 21: 0, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 1, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 58: { # 'Ó' - 28: 1, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 0, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 1, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 2, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 0, # 'O' - 46: 1, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 0, # 'U' - 48: 1, # 'V' - 55: 0, # 'Y' - 52: 1, # 'Z' - 2: 0, # 'a' - 18: 1, # 'b' - 26: 1, # 'c' - 17: 1, # 'd' - 1: 0, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 2, # 'h' - 9: 0, # 'i' - 22: 0, # 'j' - 7: 1, # 'k' - 6: 1, # 'l' - 13: 0, # 'm' - 4: 1, # 'n' - 8: 0, # 'o' - 23: 1, # 'p' - 10: 1, # 'r' - 5: 1, # 's' - 3: 0, # 't' - 21: 0, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 1, # 'z' - 51: 0, # 'Á' - 44: 1, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 59: { # 'Ö' - 28: 0, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 0, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 0, # 'O' - 46: 1, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 0, # 'U' - 48: 1, # 'V' - 55: 0, # 'Y' - 52: 1, # 'Z' - 2: 0, # 'a' - 18: 0, # 'b' - 26: 1, # 'c' - 17: 1, # 'd' - 1: 0, # 'e' - 27: 0, # 'f' - 12: 0, # 'g' - 20: 0, # 'h' - 9: 0, # 'i' - 22: 0, # 'j' - 7: 1, # 'k' - 6: 1, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 8: 0, # 'o' - 23: 0, # 'p' - 10: 2, # 'r' - 5: 1, # 's' - 3: 1, # 't' - 21: 0, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 1, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 60: { # 'Ú' - 28: 0, # 'A' - 40: 1, # 'B' - 54: 1, # 'C' - 45: 1, # 'D' - 32: 0, # 'E' - 50: 1, # 'F' - 49: 1, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 0, # 'U' - 48: 1, # 'V' - 55: 0, # 'Y' - 52: 1, # 'Z' - 2: 0, # 'a' - 18: 0, # 'b' - 26: 0, # 'c' - 17: 0, # 'd' - 1: 0, # 'e' - 27: 0, # 'f' - 12: 2, # 'g' - 20: 0, # 'h' - 9: 0, # 'i' - 22: 2, # 'j' - 7: 0, # 'k' - 6: 0, # 'l' - 13: 0, # 'm' - 4: 1, # 'n' - 8: 0, # 'o' - 23: 0, # 'p' - 10: 1, # 'r' - 5: 1, # 's' - 3: 1, # 't' - 21: 0, # 'u' - 19: 0, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 0, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 63: { # 'Ü' - 28: 0, # 'A' - 40: 1, # 'B' - 54: 0, # 'C' - 45: 1, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 1, # 'G' - 38: 1, # 'H' - 39: 0, # 'I' - 53: 1, # 'J' - 36: 1, # 'K' - 41: 1, # 'L' - 34: 1, # 'M' - 35: 1, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 1, # 'R' - 33: 1, # 'S' - 37: 1, # 'T' - 57: 0, # 'U' - 48: 1, # 'V' - 55: 0, # 'Y' - 52: 1, # 'Z' - 2: 0, # 'a' - 18: 1, # 'b' - 26: 0, # 'c' - 17: 1, # 'd' - 1: 0, # 'e' - 27: 0, # 'f' - 12: 1, # 'g' - 20: 0, # 'h' - 9: 0, # 'i' - 22: 0, # 'j' - 7: 0, # 'k' - 6: 1, # 'l' - 13: 0, # 'm' - 4: 1, # 'n' - 8: 0, # 'o' - 23: 0, # 'p' - 10: 1, # 'r' - 5: 1, # 's' - 3: 1, # 't' - 21: 0, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 1, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 14: { # 'á' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 1, # 'a' - 18: 3, # 'b' - 26: 3, # 'c' - 17: 3, # 'd' - 1: 1, # 'e' - 27: 2, # 'f' - 12: 3, # 'g' - 20: 2, # 'h' - 9: 2, # 'i' - 22: 3, # 'j' - 7: 3, # 'k' - 6: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 8: 1, # 'o' - 23: 2, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 2, # 'u' - 19: 3, # 'v' - 62: 0, # 'x' - 16: 1, # 'y' - 11: 3, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 1, # 'á' - 15: 2, # 'é' - 30: 1, # 'í' - 25: 0, # 'ó' - 24: 1, # 'ö' - 31: 0, # 'ú' - 29: 1, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 15: { # 'é' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 1, # 'a' - 18: 3, # 'b' - 26: 2, # 'c' - 17: 3, # 'd' - 1: 1, # 'e' - 27: 1, # 'f' - 12: 3, # 'g' - 20: 3, # 'h' - 9: 2, # 'i' - 22: 2, # 'j' - 7: 3, # 'k' - 6: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 8: 1, # 'o' - 23: 3, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 0, # 'u' - 19: 3, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 3, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 1, # 'á' - 15: 1, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 1, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 30: { # 'í' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 0, # 'a' - 18: 1, # 'b' - 26: 2, # 'c' - 17: 1, # 'd' - 1: 0, # 'e' - 27: 1, # 'f' - 12: 3, # 'g' - 20: 0, # 'h' - 9: 0, # 'i' - 22: 1, # 'j' - 7: 1, # 'k' - 6: 2, # 'l' - 13: 2, # 'm' - 4: 3, # 'n' - 8: 0, # 'o' - 23: 1, # 'p' - 10: 3, # 'r' - 5: 2, # 's' - 3: 3, # 't' - 21: 0, # 'u' - 19: 3, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 2, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 25: { # 'ó' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 2, # 'a' - 18: 3, # 'b' - 26: 2, # 'c' - 17: 3, # 'd' - 1: 1, # 'e' - 27: 2, # 'f' - 12: 2, # 'g' - 20: 2, # 'h' - 9: 2, # 'i' - 22: 2, # 'j' - 7: 3, # 'k' - 6: 3, # 'l' - 13: 2, # 'm' - 4: 3, # 'n' - 8: 1, # 'o' - 23: 2, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 1, # 'u' - 19: 2, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 3, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 1, # 'á' - 15: 1, # 'é' - 30: 1, # 'í' - 25: 0, # 'ó' - 24: 1, # 'ö' - 31: 1, # 'ú' - 29: 1, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 24: { # 'ö' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 0, # 'a' - 18: 3, # 'b' - 26: 1, # 'c' - 17: 2, # 'd' - 1: 0, # 'e' - 27: 1, # 'f' - 12: 2, # 'g' - 20: 1, # 'h' - 9: 0, # 'i' - 22: 1, # 'j' - 7: 3, # 'k' - 6: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 8: 0, # 'o' - 23: 2, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 3, # 't' - 21: 0, # 'u' - 19: 3, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 3, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 31: { # 'ú' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 1, # 'a' - 18: 1, # 'b' - 26: 2, # 'c' - 17: 1, # 'd' - 1: 1, # 'e' - 27: 2, # 'f' - 12: 3, # 'g' - 20: 1, # 'h' - 9: 1, # 'i' - 22: 3, # 'j' - 7: 1, # 'k' - 6: 3, # 'l' - 13: 1, # 'm' - 4: 2, # 'n' - 8: 0, # 'o' - 23: 1, # 'p' - 10: 3, # 'r' - 5: 3, # 's' - 3: 2, # 't' - 21: 1, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 2, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 1, # 'á' - 15: 1, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 29: { # 'ü' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 1, # 'a' - 18: 1, # 'b' - 26: 1, # 'c' - 17: 2, # 'd' - 1: 1, # 'e' - 27: 1, # 'f' - 12: 3, # 'g' - 20: 2, # 'h' - 9: 1, # 'i' - 22: 1, # 'j' - 7: 3, # 'k' - 6: 3, # 'l' - 13: 1, # 'm' - 4: 3, # 'n' - 8: 0, # 'o' - 23: 1, # 'p' - 10: 2, # 'r' - 5: 2, # 's' - 3: 2, # 't' - 21: 0, # 'u' - 19: 2, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 2, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 1, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 42: { # 'ő' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 1, # 'a' - 18: 2, # 'b' - 26: 1, # 'c' - 17: 2, # 'd' - 1: 1, # 'e' - 27: 1, # 'f' - 12: 1, # 'g' - 20: 1, # 'h' - 9: 1, # 'i' - 22: 1, # 'j' - 7: 2, # 'k' - 6: 3, # 'l' - 13: 1, # 'm' - 4: 2, # 'n' - 8: 1, # 'o' - 23: 1, # 'p' - 10: 2, # 'r' - 5: 2, # 's' - 3: 2, # 't' - 21: 1, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 2, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 1, # 'é' - 30: 1, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 1, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, - 56: { # 'ű' - 28: 0, # 'A' - 40: 0, # 'B' - 54: 0, # 'C' - 45: 0, # 'D' - 32: 0, # 'E' - 50: 0, # 'F' - 49: 0, # 'G' - 38: 0, # 'H' - 39: 0, # 'I' - 53: 0, # 'J' - 36: 0, # 'K' - 41: 0, # 'L' - 34: 0, # 'M' - 35: 0, # 'N' - 47: 0, # 'O' - 46: 0, # 'P' - 43: 0, # 'R' - 33: 0, # 'S' - 37: 0, # 'T' - 57: 0, # 'U' - 48: 0, # 'V' - 55: 0, # 'Y' - 52: 0, # 'Z' - 2: 1, # 'a' - 18: 1, # 'b' - 26: 0, # 'c' - 17: 1, # 'd' - 1: 1, # 'e' - 27: 1, # 'f' - 12: 1, # 'g' - 20: 1, # 'h' - 9: 1, # 'i' - 22: 1, # 'j' - 7: 1, # 'k' - 6: 1, # 'l' - 13: 0, # 'm' - 4: 2, # 'n' - 8: 0, # 'o' - 23: 0, # 'p' - 10: 1, # 'r' - 5: 1, # 's' - 3: 1, # 't' - 21: 0, # 'u' - 19: 1, # 'v' - 62: 0, # 'x' - 16: 0, # 'y' - 11: 2, # 'z' - 51: 0, # 'Á' - 44: 0, # 'É' - 61: 0, # 'Í' - 58: 0, # 'Ó' - 59: 0, # 'Ö' - 60: 0, # 'Ú' - 63: 0, # 'Ü' - 14: 0, # 'á' - 15: 0, # 'é' - 30: 0, # 'í' - 25: 0, # 'ó' - 24: 0, # 'ö' - 31: 0, # 'ú' - 29: 0, # 'ü' - 42: 0, # 'ő' - 56: 0, # 'ű' - }, -} - -# 255: Undefined characters that did not exist in training text -# 254: Carriage/Return -# 253: symbol (punctuation) that does not belong to word -# 252: 0 - 9 -# 251: Control characters - -# Character Mapping Table(s): -WINDOWS_1250_HUNGARIAN_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 254, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 254, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 253, # ' ' - 33: 253, # '!' - 34: 253, # '"' - 35: 253, # '#' - 36: 253, # '$' - 37: 253, # '%' - 38: 253, # '&' - 39: 253, # "'" - 40: 253, # '(' - 41: 253, # ')' - 42: 253, # '*' - 43: 253, # '+' - 44: 253, # ',' - 45: 253, # '-' - 46: 253, # '.' - 47: 253, # '/' - 48: 252, # '0' - 49: 252, # '1' - 50: 252, # '2' - 51: 252, # '3' - 52: 252, # '4' - 53: 252, # '5' - 54: 252, # '6' - 55: 252, # '7' - 56: 252, # '8' - 57: 252, # '9' - 58: 253, # ':' - 59: 253, # ';' - 60: 253, # '<' - 61: 253, # '=' - 62: 253, # '>' - 63: 253, # '?' - 64: 253, # '@' - 65: 28, # 'A' - 66: 40, # 'B' - 67: 54, # 'C' - 68: 45, # 'D' - 69: 32, # 'E' - 70: 50, # 'F' - 71: 49, # 'G' - 72: 38, # 'H' - 73: 39, # 'I' - 74: 53, # 'J' - 75: 36, # 'K' - 76: 41, # 'L' - 77: 34, # 'M' - 78: 35, # 'N' - 79: 47, # 'O' - 80: 46, # 'P' - 81: 72, # 'Q' - 82: 43, # 'R' - 83: 33, # 'S' - 84: 37, # 'T' - 85: 57, # 'U' - 86: 48, # 'V' - 87: 64, # 'W' - 88: 68, # 'X' - 89: 55, # 'Y' - 90: 52, # 'Z' - 91: 253, # '[' - 92: 253, # '\\' - 93: 253, # ']' - 94: 253, # '^' - 95: 253, # '_' - 96: 253, # '`' - 97: 2, # 'a' - 98: 18, # 'b' - 99: 26, # 'c' - 100: 17, # 'd' - 101: 1, # 'e' - 102: 27, # 'f' - 103: 12, # 'g' - 104: 20, # 'h' - 105: 9, # 'i' - 106: 22, # 'j' - 107: 7, # 'k' - 108: 6, # 'l' - 109: 13, # 'm' - 110: 4, # 'n' - 111: 8, # 'o' - 112: 23, # 'p' - 113: 67, # 'q' - 114: 10, # 'r' - 115: 5, # 's' - 116: 3, # 't' - 117: 21, # 'u' - 118: 19, # 'v' - 119: 65, # 'w' - 120: 62, # 'x' - 121: 16, # 'y' - 122: 11, # 'z' - 123: 253, # '{' - 124: 253, # '|' - 125: 253, # '}' - 126: 253, # '~' - 127: 253, # '\x7f' - 128: 161, # '€' - 129: 162, # None - 130: 163, # '‚' - 131: 164, # None - 132: 165, # '„' - 133: 166, # '…' - 134: 167, # '†' - 135: 168, # '‡' - 136: 169, # None - 137: 170, # '‰' - 138: 171, # 'Š' - 139: 172, # '‹' - 140: 173, # 'Ś' - 141: 174, # 'Ť' - 142: 175, # 'Ž' - 143: 176, # 'Ź' - 144: 177, # None - 145: 178, # '‘' - 146: 179, # '’' - 147: 180, # '“' - 148: 78, # '”' - 149: 181, # '•' - 150: 69, # '–' - 151: 182, # '—' - 152: 183, # None - 153: 184, # '™' - 154: 185, # 'š' - 155: 186, # '›' - 156: 187, # 'ś' - 157: 188, # 'ť' - 158: 189, # 'ž' - 159: 190, # 'ź' - 160: 191, # '\xa0' - 161: 192, # 'ˇ' - 162: 193, # '˘' - 163: 194, # 'Ł' - 164: 195, # '¤' - 165: 196, # 'Ą' - 166: 197, # '¦' - 167: 76, # '§' - 168: 198, # '¨' - 169: 199, # '©' - 170: 200, # 'Ş' - 171: 201, # '«' - 172: 202, # '¬' - 173: 203, # '\xad' - 174: 204, # '®' - 175: 205, # 'Ż' - 176: 81, # '°' - 177: 206, # '±' - 178: 207, # '˛' - 179: 208, # 'ł' - 180: 209, # '´' - 181: 210, # 'µ' - 182: 211, # '¶' - 183: 212, # '·' - 184: 213, # '¸' - 185: 214, # 'ą' - 186: 215, # 'ş' - 187: 216, # '»' - 188: 217, # 'Ľ' - 189: 218, # '˝' - 190: 219, # 'ľ' - 191: 220, # 'ż' - 192: 221, # 'Ŕ' - 193: 51, # 'Á' - 194: 83, # 'Â' - 195: 222, # 'Ă' - 196: 80, # 'Ä' - 197: 223, # 'Ĺ' - 198: 224, # 'Ć' - 199: 225, # 'Ç' - 200: 226, # 'Č' - 201: 44, # 'É' - 202: 227, # 'Ę' - 203: 228, # 'Ë' - 204: 229, # 'Ě' - 205: 61, # 'Í' - 206: 230, # 'Î' - 207: 231, # 'Ď' - 208: 232, # 'Đ' - 209: 233, # 'Ń' - 210: 234, # 'Ň' - 211: 58, # 'Ó' - 212: 235, # 'Ô' - 213: 66, # 'Ő' - 214: 59, # 'Ö' - 215: 236, # '×' - 216: 237, # 'Ř' - 217: 238, # 'Ů' - 218: 60, # 'Ú' - 219: 70, # 'Ű' - 220: 63, # 'Ü' - 221: 239, # 'Ý' - 222: 240, # 'Ţ' - 223: 241, # 'ß' - 224: 84, # 'ŕ' - 225: 14, # 'á' - 226: 75, # 'â' - 227: 242, # 'ă' - 228: 71, # 'ä' - 229: 82, # 'ĺ' - 230: 243, # 'ć' - 231: 73, # 'ç' - 232: 244, # 'č' - 233: 15, # 'é' - 234: 85, # 'ę' - 235: 79, # 'ë' - 236: 86, # 'ě' - 237: 30, # 'í' - 238: 77, # 'î' - 239: 87, # 'ď' - 240: 245, # 'đ' - 241: 246, # 'ń' - 242: 247, # 'ň' - 243: 25, # 'ó' - 244: 74, # 'ô' - 245: 42, # 'ő' - 246: 24, # 'ö' - 247: 248, # '÷' - 248: 249, # 'ř' - 249: 250, # 'ů' - 250: 31, # 'ú' - 251: 56, # 'ű' - 252: 29, # 'ü' - 253: 251, # 'ý' - 254: 252, # 'ţ' - 255: 253, # '˙' -} - -WINDOWS_1250_HUNGARIAN_MODEL = SingleByteCharSetModel( - charset_name="windows-1250", - language="Hungarian", - char_to_order_map=WINDOWS_1250_HUNGARIAN_CHAR_TO_ORDER, - language_model=HUNGARIAN_LANG_MODEL, - typical_positive_ratio=0.947368, - keep_ascii_letters=True, - alphabet="ABCDEFGHIJKLMNOPRSTUVZabcdefghijklmnoprstuvzÁÉÍÓÖÚÜáéíóöúüŐőŰű", -) - -ISO_8859_2_HUNGARIAN_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 254, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 254, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 253, # ' ' - 33: 253, # '!' - 34: 253, # '"' - 35: 253, # '#' - 36: 253, # '$' - 37: 253, # '%' - 38: 253, # '&' - 39: 253, # "'" - 40: 253, # '(' - 41: 253, # ')' - 42: 253, # '*' - 43: 253, # '+' - 44: 253, # ',' - 45: 253, # '-' - 46: 253, # '.' - 47: 253, # '/' - 48: 252, # '0' - 49: 252, # '1' - 50: 252, # '2' - 51: 252, # '3' - 52: 252, # '4' - 53: 252, # '5' - 54: 252, # '6' - 55: 252, # '7' - 56: 252, # '8' - 57: 252, # '9' - 58: 253, # ':' - 59: 253, # ';' - 60: 253, # '<' - 61: 253, # '=' - 62: 253, # '>' - 63: 253, # '?' - 64: 253, # '@' - 65: 28, # 'A' - 66: 40, # 'B' - 67: 54, # 'C' - 68: 45, # 'D' - 69: 32, # 'E' - 70: 50, # 'F' - 71: 49, # 'G' - 72: 38, # 'H' - 73: 39, # 'I' - 74: 53, # 'J' - 75: 36, # 'K' - 76: 41, # 'L' - 77: 34, # 'M' - 78: 35, # 'N' - 79: 47, # 'O' - 80: 46, # 'P' - 81: 71, # 'Q' - 82: 43, # 'R' - 83: 33, # 'S' - 84: 37, # 'T' - 85: 57, # 'U' - 86: 48, # 'V' - 87: 64, # 'W' - 88: 68, # 'X' - 89: 55, # 'Y' - 90: 52, # 'Z' - 91: 253, # '[' - 92: 253, # '\\' - 93: 253, # ']' - 94: 253, # '^' - 95: 253, # '_' - 96: 253, # '`' - 97: 2, # 'a' - 98: 18, # 'b' - 99: 26, # 'c' - 100: 17, # 'd' - 101: 1, # 'e' - 102: 27, # 'f' - 103: 12, # 'g' - 104: 20, # 'h' - 105: 9, # 'i' - 106: 22, # 'j' - 107: 7, # 'k' - 108: 6, # 'l' - 109: 13, # 'm' - 110: 4, # 'n' - 111: 8, # 'o' - 112: 23, # 'p' - 113: 67, # 'q' - 114: 10, # 'r' - 115: 5, # 's' - 116: 3, # 't' - 117: 21, # 'u' - 118: 19, # 'v' - 119: 65, # 'w' - 120: 62, # 'x' - 121: 16, # 'y' - 122: 11, # 'z' - 123: 253, # '{' - 124: 253, # '|' - 125: 253, # '}' - 126: 253, # '~' - 127: 253, # '\x7f' - 128: 159, # '\x80' - 129: 160, # '\x81' - 130: 161, # '\x82' - 131: 162, # '\x83' - 132: 163, # '\x84' - 133: 164, # '\x85' - 134: 165, # '\x86' - 135: 166, # '\x87' - 136: 167, # '\x88' - 137: 168, # '\x89' - 138: 169, # '\x8a' - 139: 170, # '\x8b' - 140: 171, # '\x8c' - 141: 172, # '\x8d' - 142: 173, # '\x8e' - 143: 174, # '\x8f' - 144: 175, # '\x90' - 145: 176, # '\x91' - 146: 177, # '\x92' - 147: 178, # '\x93' - 148: 179, # '\x94' - 149: 180, # '\x95' - 150: 181, # '\x96' - 151: 182, # '\x97' - 152: 183, # '\x98' - 153: 184, # '\x99' - 154: 185, # '\x9a' - 155: 186, # '\x9b' - 156: 187, # '\x9c' - 157: 188, # '\x9d' - 158: 189, # '\x9e' - 159: 190, # '\x9f' - 160: 191, # '\xa0' - 161: 192, # 'Ą' - 162: 193, # '˘' - 163: 194, # 'Ł' - 164: 195, # '¤' - 165: 196, # 'Ľ' - 166: 197, # 'Ś' - 167: 75, # '§' - 168: 198, # '¨' - 169: 199, # 'Š' - 170: 200, # 'Ş' - 171: 201, # 'Ť' - 172: 202, # 'Ź' - 173: 203, # '\xad' - 174: 204, # 'Ž' - 175: 205, # 'Ż' - 176: 79, # '°' - 177: 206, # 'ą' - 178: 207, # '˛' - 179: 208, # 'ł' - 180: 209, # '´' - 181: 210, # 'ľ' - 182: 211, # 'ś' - 183: 212, # 'ˇ' - 184: 213, # '¸' - 185: 214, # 'š' - 186: 215, # 'ş' - 187: 216, # 'ť' - 188: 217, # 'ź' - 189: 218, # '˝' - 190: 219, # 'ž' - 191: 220, # 'ż' - 192: 221, # 'Ŕ' - 193: 51, # 'Á' - 194: 81, # 'Â' - 195: 222, # 'Ă' - 196: 78, # 'Ä' - 197: 223, # 'Ĺ' - 198: 224, # 'Ć' - 199: 225, # 'Ç' - 200: 226, # 'Č' - 201: 44, # 'É' - 202: 227, # 'Ę' - 203: 228, # 'Ë' - 204: 229, # 'Ě' - 205: 61, # 'Í' - 206: 230, # 'Î' - 207: 231, # 'Ď' - 208: 232, # 'Đ' - 209: 233, # 'Ń' - 210: 234, # 'Ň' - 211: 58, # 'Ó' - 212: 235, # 'Ô' - 213: 66, # 'Ő' - 214: 59, # 'Ö' - 215: 236, # '×' - 216: 237, # 'Ř' - 217: 238, # 'Ů' - 218: 60, # 'Ú' - 219: 69, # 'Ű' - 220: 63, # 'Ü' - 221: 239, # 'Ý' - 222: 240, # 'Ţ' - 223: 241, # 'ß' - 224: 82, # 'ŕ' - 225: 14, # 'á' - 226: 74, # 'â' - 227: 242, # 'ă' - 228: 70, # 'ä' - 229: 80, # 'ĺ' - 230: 243, # 'ć' - 231: 72, # 'ç' - 232: 244, # 'č' - 233: 15, # 'é' - 234: 83, # 'ę' - 235: 77, # 'ë' - 236: 84, # 'ě' - 237: 30, # 'í' - 238: 76, # 'î' - 239: 85, # 'ď' - 240: 245, # 'đ' - 241: 246, # 'ń' - 242: 247, # 'ň' - 243: 25, # 'ó' - 244: 73, # 'ô' - 245: 42, # 'ő' - 246: 24, # 'ö' - 247: 248, # '÷' - 248: 249, # 'ř' - 249: 250, # 'ů' - 250: 31, # 'ú' - 251: 56, # 'ű' - 252: 29, # 'ü' - 253: 251, # 'ý' - 254: 252, # 'ţ' - 255: 253, # '˙' -} - -ISO_8859_2_HUNGARIAN_MODEL = SingleByteCharSetModel( - charset_name="ISO-8859-2", - language="Hungarian", - char_to_order_map=ISO_8859_2_HUNGARIAN_CHAR_TO_ORDER, - language_model=HUNGARIAN_LANG_MODEL, - typical_positive_ratio=0.947368, - keep_ascii_letters=True, - alphabet="ABCDEFGHIJKLMNOPRSTUVZabcdefghijklmnoprstuvzÁÉÍÓÖÚÜáéíóöúüŐőŰű", -) diff --git a/venv/lib/python3.10/site-packages/chardet/langrussianmodel.py b/venv/lib/python3.10/site-packages/chardet/langrussianmodel.py deleted file mode 100644 index 0d5b178446d243dbd98bd1e9ebb0a89e19d94258..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/langrussianmodel.py +++ /dev/null @@ -1,5725 +0,0 @@ -from chardet.sbcharsetprober import SingleByteCharSetModel - -# 3: Positive -# 2: Likely -# 1: Unlikely -# 0: Negative - -RUSSIAN_LANG_MODEL = { - 37: { # 'А' - 37: 0, # 'А' - 44: 1, # 'Б' - 33: 1, # 'В' - 46: 1, # 'Г' - 41: 1, # 'Д' - 48: 1, # 'Е' - 56: 1, # 'Ж' - 51: 1, # 'З' - 42: 1, # 'И' - 60: 1, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 2, # 'Н' - 34: 1, # 'О' - 35: 1, # 'П' - 45: 1, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 1, # 'У' - 53: 1, # 'Ф' - 55: 1, # 'Х' - 58: 1, # 'Ц' - 50: 1, # 'Ч' - 57: 1, # 'Ш' - 63: 1, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 1, # 'Ю' - 43: 1, # 'Я' - 3: 1, # 'а' - 21: 2, # 'б' - 10: 2, # 'в' - 19: 2, # 'г' - 13: 2, # 'д' - 2: 0, # 'е' - 24: 1, # 'ж' - 20: 1, # 'з' - 4: 0, # 'и' - 23: 1, # 'й' - 11: 2, # 'к' - 8: 3, # 'л' - 12: 2, # 'м' - 5: 2, # 'н' - 1: 0, # 'о' - 15: 2, # 'п' - 9: 2, # 'р' - 7: 2, # 'с' - 6: 2, # 'т' - 14: 2, # 'у' - 39: 2, # 'ф' - 26: 2, # 'х' - 28: 0, # 'ц' - 22: 1, # 'ч' - 25: 2, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 1, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 44: { # 'Б' - 37: 1, # 'А' - 44: 0, # 'Б' - 33: 1, # 'В' - 46: 1, # 'Г' - 41: 0, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 0, # 'П' - 45: 1, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 1, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 1, # 'Я' - 3: 2, # 'а' - 21: 0, # 'б' - 10: 0, # 'в' - 19: 0, # 'г' - 13: 1, # 'д' - 2: 3, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 2, # 'л' - 12: 0, # 'м' - 5: 0, # 'н' - 1: 3, # 'о' - 15: 0, # 'п' - 9: 2, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 2, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 2, # 'ы' - 17: 1, # 'ь' - 30: 2, # 'э' - 27: 1, # 'ю' - 16: 1, # 'я' - }, - 33: { # 'В' - 37: 2, # 'А' - 44: 0, # 'Б' - 33: 1, # 'В' - 46: 0, # 'Г' - 41: 1, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 1, # 'П' - 45: 1, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 1, # 'Ш' - 63: 0, # 'Щ' - 62: 1, # 'Ы' - 61: 1, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 1, # 'Я' - 3: 2, # 'а' - 21: 1, # 'б' - 10: 1, # 'в' - 19: 1, # 'г' - 13: 2, # 'д' - 2: 3, # 'е' - 24: 0, # 'ж' - 20: 2, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 1, # 'к' - 8: 2, # 'л' - 12: 2, # 'м' - 5: 2, # 'н' - 1: 3, # 'о' - 15: 2, # 'п' - 9: 2, # 'р' - 7: 3, # 'с' - 6: 2, # 'т' - 14: 2, # 'у' - 39: 0, # 'ф' - 26: 1, # 'х' - 28: 1, # 'ц' - 22: 2, # 'ч' - 25: 1, # 'ш' - 29: 0, # 'щ' - 54: 1, # 'ъ' - 18: 3, # 'ы' - 17: 1, # 'ь' - 30: 2, # 'э' - 27: 0, # 'ю' - 16: 1, # 'я' - }, - 46: { # 'Г' - 37: 1, # 'А' - 44: 1, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 1, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 1, # 'П' - 45: 1, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 2, # 'а' - 21: 0, # 'б' - 10: 1, # 'в' - 19: 0, # 'г' - 13: 2, # 'д' - 2: 2, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 2, # 'л' - 12: 1, # 'м' - 5: 1, # 'н' - 1: 3, # 'о' - 15: 0, # 'п' - 9: 2, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 2, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 1, # 'ь' - 30: 1, # 'э' - 27: 1, # 'ю' - 16: 0, # 'я' - }, - 41: { # 'Д' - 37: 1, # 'А' - 44: 0, # 'Б' - 33: 1, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 2, # 'Е' - 56: 1, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 0, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 0, # 'П' - 45: 1, # 'Р' - 32: 1, # 'С' - 40: 0, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 1, # 'Ц' - 50: 1, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 1, # 'Ы' - 61: 1, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 1, # 'Я' - 3: 3, # 'а' - 21: 0, # 'б' - 10: 2, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 2, # 'е' - 24: 3, # 'ж' - 20: 1, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 2, # 'л' - 12: 1, # 'м' - 5: 1, # 'н' - 1: 3, # 'о' - 15: 0, # 'п' - 9: 2, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 2, # 'у' - 39: 0, # 'ф' - 26: 1, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 1, # 'ы' - 17: 1, # 'ь' - 30: 2, # 'э' - 27: 1, # 'ю' - 16: 1, # 'я' - }, - 48: { # 'Е' - 37: 1, # 'А' - 44: 1, # 'Б' - 33: 1, # 'В' - 46: 1, # 'Г' - 41: 1, # 'Д' - 48: 1, # 'Е' - 56: 1, # 'Ж' - 51: 1, # 'З' - 42: 1, # 'И' - 60: 1, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 2, # 'Н' - 34: 1, # 'О' - 35: 1, # 'П' - 45: 2, # 'Р' - 32: 2, # 'С' - 40: 1, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 1, # 'Х' - 58: 1, # 'Ц' - 50: 1, # 'Ч' - 57: 1, # 'Ш' - 63: 1, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 1, # 'Я' - 3: 0, # 'а' - 21: 0, # 'б' - 10: 2, # 'в' - 19: 2, # 'г' - 13: 2, # 'д' - 2: 2, # 'е' - 24: 1, # 'ж' - 20: 1, # 'з' - 4: 0, # 'и' - 23: 2, # 'й' - 11: 1, # 'к' - 8: 2, # 'л' - 12: 2, # 'м' - 5: 1, # 'н' - 1: 0, # 'о' - 15: 1, # 'п' - 9: 1, # 'р' - 7: 3, # 'с' - 6: 0, # 'т' - 14: 0, # 'у' - 39: 1, # 'ф' - 26: 1, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 1, # 'ш' - 29: 2, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 1, # 'ю' - 16: 0, # 'я' - }, - 56: { # 'Ж' - 37: 1, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 1, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 1, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 2, # 'а' - 21: 1, # 'б' - 10: 0, # 'в' - 19: 1, # 'г' - 13: 1, # 'д' - 2: 2, # 'е' - 24: 1, # 'ж' - 20: 0, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 0, # 'л' - 12: 1, # 'м' - 5: 0, # 'н' - 1: 2, # 'о' - 15: 0, # 'п' - 9: 1, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 2, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 2, # 'ю' - 16: 0, # 'я' - }, - 51: { # 'З' - 37: 1, # 'А' - 44: 0, # 'Б' - 33: 1, # 'В' - 46: 1, # 'Г' - 41: 1, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 0, # 'П' - 45: 1, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 1, # 'Ы' - 61: 1, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 1, # 'б' - 10: 2, # 'в' - 19: 0, # 'г' - 13: 2, # 'д' - 2: 2, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 1, # 'л' - 12: 1, # 'м' - 5: 2, # 'н' - 1: 2, # 'о' - 15: 0, # 'п' - 9: 1, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 1, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 1, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 0, # 'ю' - 16: 1, # 'я' - }, - 42: { # 'И' - 37: 1, # 'А' - 44: 1, # 'Б' - 33: 1, # 'В' - 46: 1, # 'Г' - 41: 1, # 'Д' - 48: 2, # 'Е' - 56: 1, # 'Ж' - 51: 1, # 'З' - 42: 1, # 'И' - 60: 1, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 1, # 'П' - 45: 1, # 'Р' - 32: 2, # 'С' - 40: 1, # 'Т' - 52: 0, # 'У' - 53: 1, # 'Ф' - 55: 1, # 'Х' - 58: 1, # 'Ц' - 50: 1, # 'Ч' - 57: 0, # 'Ш' - 63: 1, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 1, # 'Ю' - 43: 1, # 'Я' - 3: 1, # 'а' - 21: 2, # 'б' - 10: 2, # 'в' - 19: 2, # 'г' - 13: 2, # 'д' - 2: 2, # 'е' - 24: 0, # 'ж' - 20: 2, # 'з' - 4: 1, # 'и' - 23: 0, # 'й' - 11: 1, # 'к' - 8: 2, # 'л' - 12: 2, # 'м' - 5: 2, # 'н' - 1: 1, # 'о' - 15: 1, # 'п' - 9: 2, # 'р' - 7: 2, # 'с' - 6: 2, # 'т' - 14: 1, # 'у' - 39: 1, # 'ф' - 26: 2, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 1, # 'ш' - 29: 1, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 1, # 'ю' - 16: 0, # 'я' - }, - 60: { # 'Й' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 1, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 0, # 'М' - 31: 1, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 1, # 'Х' - 58: 1, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 0, # 'а' - 21: 0, # 'б' - 10: 0, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 1, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 0, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 0, # 'л' - 12: 0, # 'м' - 5: 0, # 'н' - 1: 2, # 'о' - 15: 0, # 'п' - 9: 0, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 0, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 36: { # 'К' - 37: 2, # 'А' - 44: 0, # 'Б' - 33: 1, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 1, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 1, # 'Л' - 38: 0, # 'М' - 31: 1, # 'Н' - 34: 2, # 'О' - 35: 1, # 'П' - 45: 1, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 1, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 0, # 'б' - 10: 1, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 2, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 2, # 'л' - 12: 0, # 'м' - 5: 1, # 'н' - 1: 3, # 'о' - 15: 0, # 'п' - 9: 2, # 'р' - 7: 2, # 'с' - 6: 2, # 'т' - 14: 2, # 'у' - 39: 0, # 'ф' - 26: 1, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 1, # 'ы' - 17: 1, # 'ь' - 30: 2, # 'э' - 27: 1, # 'ю' - 16: 0, # 'я' - }, - 49: { # 'Л' - 37: 2, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 1, # 'Г' - 41: 0, # 'Д' - 48: 1, # 'Е' - 56: 1, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 0, # 'Н' - 34: 1, # 'О' - 35: 1, # 'П' - 45: 0, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 1, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 1, # 'Ы' - 61: 1, # 'Ь' - 47: 0, # 'Э' - 59: 1, # 'Ю' - 43: 1, # 'Я' - 3: 2, # 'а' - 21: 0, # 'б' - 10: 0, # 'в' - 19: 1, # 'г' - 13: 0, # 'д' - 2: 2, # 'е' - 24: 1, # 'ж' - 20: 0, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 1, # 'л' - 12: 0, # 'м' - 5: 1, # 'н' - 1: 2, # 'о' - 15: 0, # 'п' - 9: 0, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 2, # 'у' - 39: 0, # 'ф' - 26: 1, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 1, # 'ы' - 17: 1, # 'ь' - 30: 2, # 'э' - 27: 2, # 'ю' - 16: 1, # 'я' - }, - 38: { # 'М' - 37: 1, # 'А' - 44: 1, # 'Б' - 33: 1, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 1, # 'П' - 45: 1, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 1, # 'У' - 53: 1, # 'Ф' - 55: 1, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 1, # 'Ы' - 61: 0, # 'Ь' - 47: 1, # 'Э' - 59: 0, # 'Ю' - 43: 1, # 'Я' - 3: 3, # 'а' - 21: 0, # 'б' - 10: 0, # 'в' - 19: 1, # 'г' - 13: 0, # 'д' - 2: 2, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 1, # 'л' - 12: 1, # 'м' - 5: 2, # 'н' - 1: 3, # 'о' - 15: 0, # 'п' - 9: 1, # 'р' - 7: 1, # 'с' - 6: 0, # 'т' - 14: 2, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 3, # 'ы' - 17: 1, # 'ь' - 30: 2, # 'э' - 27: 1, # 'ю' - 16: 1, # 'я' - }, - 31: { # 'Н' - 37: 2, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 1, # 'Г' - 41: 1, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 1, # 'З' - 42: 2, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 0, # 'П' - 45: 1, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 1, # 'У' - 53: 1, # 'Ф' - 55: 1, # 'Х' - 58: 1, # 'Ц' - 50: 1, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 1, # 'Ы' - 61: 1, # 'Ь' - 47: 1, # 'Э' - 59: 0, # 'Ю' - 43: 1, # 'Я' - 3: 3, # 'а' - 21: 0, # 'б' - 10: 0, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 3, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 0, # 'л' - 12: 0, # 'м' - 5: 0, # 'н' - 1: 3, # 'о' - 15: 0, # 'п' - 9: 1, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 3, # 'у' - 39: 0, # 'ф' - 26: 1, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 1, # 'ы' - 17: 2, # 'ь' - 30: 1, # 'э' - 27: 1, # 'ю' - 16: 1, # 'я' - }, - 34: { # 'О' - 37: 0, # 'А' - 44: 1, # 'Б' - 33: 1, # 'В' - 46: 1, # 'Г' - 41: 2, # 'Д' - 48: 1, # 'Е' - 56: 1, # 'Ж' - 51: 1, # 'З' - 42: 1, # 'И' - 60: 1, # 'Й' - 36: 1, # 'К' - 49: 2, # 'Л' - 38: 1, # 'М' - 31: 2, # 'Н' - 34: 1, # 'О' - 35: 1, # 'П' - 45: 2, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 1, # 'У' - 53: 1, # 'Ф' - 55: 1, # 'Х' - 58: 0, # 'Ц' - 50: 1, # 'Ч' - 57: 1, # 'Ш' - 63: 1, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 1, # 'Я' - 3: 1, # 'а' - 21: 2, # 'б' - 10: 1, # 'в' - 19: 2, # 'г' - 13: 2, # 'д' - 2: 0, # 'е' - 24: 1, # 'ж' - 20: 1, # 'з' - 4: 0, # 'и' - 23: 1, # 'й' - 11: 2, # 'к' - 8: 2, # 'л' - 12: 1, # 'м' - 5: 3, # 'н' - 1: 0, # 'о' - 15: 2, # 'п' - 9: 2, # 'р' - 7: 2, # 'с' - 6: 2, # 'т' - 14: 1, # 'у' - 39: 1, # 'ф' - 26: 2, # 'х' - 28: 1, # 'ц' - 22: 2, # 'ч' - 25: 2, # 'ш' - 29: 1, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 35: { # 'П' - 37: 1, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 1, # 'Л' - 38: 0, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 1, # 'П' - 45: 2, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 1, # 'Ы' - 61: 1, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 1, # 'Я' - 3: 2, # 'а' - 21: 0, # 'б' - 10: 0, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 2, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 2, # 'л' - 12: 0, # 'м' - 5: 1, # 'н' - 1: 3, # 'о' - 15: 0, # 'п' - 9: 3, # 'р' - 7: 1, # 'с' - 6: 1, # 'т' - 14: 2, # 'у' - 39: 1, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 1, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 1, # 'ы' - 17: 2, # 'ь' - 30: 1, # 'э' - 27: 0, # 'ю' - 16: 2, # 'я' - }, - 45: { # 'Р' - 37: 2, # 'А' - 44: 1, # 'Б' - 33: 1, # 'В' - 46: 1, # 'Г' - 41: 1, # 'Д' - 48: 2, # 'Е' - 56: 1, # 'Ж' - 51: 0, # 'З' - 42: 2, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 2, # 'О' - 35: 0, # 'П' - 45: 1, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 1, # 'Х' - 58: 1, # 'Ц' - 50: 1, # 'Ч' - 57: 1, # 'Ш' - 63: 0, # 'Щ' - 62: 1, # 'Ы' - 61: 1, # 'Ь' - 47: 1, # 'Э' - 59: 1, # 'Ю' - 43: 1, # 'Я' - 3: 3, # 'а' - 21: 0, # 'б' - 10: 1, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 2, # 'е' - 24: 1, # 'ж' - 20: 0, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 0, # 'л' - 12: 0, # 'м' - 5: 0, # 'н' - 1: 3, # 'о' - 15: 0, # 'п' - 9: 1, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 2, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 2, # 'ы' - 17: 0, # 'ь' - 30: 1, # 'э' - 27: 1, # 'ю' - 16: 2, # 'я' - }, - 32: { # 'С' - 37: 1, # 'А' - 44: 1, # 'Б' - 33: 1, # 'В' - 46: 1, # 'Г' - 41: 1, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 1, # 'П' - 45: 1, # 'Р' - 32: 1, # 'С' - 40: 2, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 1, # 'Х' - 58: 1, # 'Ц' - 50: 1, # 'Ч' - 57: 1, # 'Ш' - 63: 0, # 'Щ' - 62: 1, # 'Ы' - 61: 1, # 'Ь' - 47: 1, # 'Э' - 59: 1, # 'Ю' - 43: 1, # 'Я' - 3: 2, # 'а' - 21: 1, # 'б' - 10: 2, # 'в' - 19: 1, # 'г' - 13: 2, # 'д' - 2: 3, # 'е' - 24: 1, # 'ж' - 20: 1, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 2, # 'к' - 8: 2, # 'л' - 12: 2, # 'м' - 5: 2, # 'н' - 1: 2, # 'о' - 15: 2, # 'п' - 9: 2, # 'р' - 7: 1, # 'с' - 6: 3, # 'т' - 14: 2, # 'у' - 39: 1, # 'ф' - 26: 1, # 'х' - 28: 1, # 'ц' - 22: 1, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 1, # 'ъ' - 18: 1, # 'ы' - 17: 1, # 'ь' - 30: 2, # 'э' - 27: 1, # 'ю' - 16: 1, # 'я' - }, - 40: { # 'Т' - 37: 1, # 'А' - 44: 0, # 'Б' - 33: 1, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 2, # 'О' - 35: 0, # 'П' - 45: 1, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 1, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 1, # 'Ы' - 61: 1, # 'Ь' - 47: 1, # 'Э' - 59: 1, # 'Ю' - 43: 1, # 'Я' - 3: 3, # 'а' - 21: 1, # 'б' - 10: 2, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 3, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 1, # 'к' - 8: 1, # 'л' - 12: 0, # 'м' - 5: 0, # 'н' - 1: 3, # 'о' - 15: 0, # 'п' - 9: 2, # 'р' - 7: 1, # 'с' - 6: 0, # 'т' - 14: 2, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 1, # 'щ' - 54: 0, # 'ъ' - 18: 3, # 'ы' - 17: 1, # 'ь' - 30: 2, # 'э' - 27: 1, # 'ю' - 16: 1, # 'я' - }, - 52: { # 'У' - 37: 1, # 'А' - 44: 1, # 'Б' - 33: 1, # 'В' - 46: 1, # 'Г' - 41: 1, # 'Д' - 48: 1, # 'Е' - 56: 1, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 1, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 1, # 'П' - 45: 1, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 1, # 'Х' - 58: 0, # 'Ц' - 50: 1, # 'Ч' - 57: 1, # 'Ш' - 63: 1, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 1, # 'Ю' - 43: 0, # 'Я' - 3: 1, # 'а' - 21: 2, # 'б' - 10: 2, # 'в' - 19: 1, # 'г' - 13: 2, # 'д' - 2: 1, # 'е' - 24: 2, # 'ж' - 20: 2, # 'з' - 4: 2, # 'и' - 23: 1, # 'й' - 11: 1, # 'к' - 8: 2, # 'л' - 12: 2, # 'м' - 5: 1, # 'н' - 1: 2, # 'о' - 15: 1, # 'п' - 9: 2, # 'р' - 7: 2, # 'с' - 6: 2, # 'т' - 14: 0, # 'у' - 39: 1, # 'ф' - 26: 1, # 'х' - 28: 1, # 'ц' - 22: 2, # 'ч' - 25: 1, # 'ш' - 29: 1, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 2, # 'э' - 27: 1, # 'ю' - 16: 0, # 'я' - }, - 53: { # 'Ф' - 37: 1, # 'А' - 44: 1, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 1, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 1, # 'О' - 35: 0, # 'П' - 45: 1, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 2, # 'а' - 21: 0, # 'б' - 10: 0, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 2, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 2, # 'л' - 12: 0, # 'м' - 5: 0, # 'н' - 1: 2, # 'о' - 15: 0, # 'п' - 9: 2, # 'р' - 7: 0, # 'с' - 6: 1, # 'т' - 14: 2, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 1, # 'ь' - 30: 2, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 55: { # 'Х' - 37: 1, # 'А' - 44: 0, # 'Б' - 33: 1, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 2, # 'а' - 21: 0, # 'б' - 10: 2, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 2, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 2, # 'л' - 12: 1, # 'м' - 5: 0, # 'н' - 1: 2, # 'о' - 15: 0, # 'п' - 9: 2, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 1, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 1, # 'ь' - 30: 1, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 58: { # 'Ц' - 37: 1, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 1, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 1, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 1, # 'а' - 21: 0, # 'б' - 10: 1, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 2, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 0, # 'л' - 12: 0, # 'м' - 5: 0, # 'н' - 1: 0, # 'о' - 15: 0, # 'п' - 9: 0, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 1, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 1, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 1, # 'ю' - 16: 0, # 'я' - }, - 50: { # 'Ч' - 37: 1, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 1, # 'Н' - 34: 0, # 'О' - 35: 1, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 1, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 1, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 2, # 'а' - 21: 0, # 'б' - 10: 0, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 2, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 1, # 'л' - 12: 0, # 'м' - 5: 0, # 'н' - 1: 1, # 'о' - 15: 0, # 'п' - 9: 1, # 'р' - 7: 0, # 'с' - 6: 3, # 'т' - 14: 2, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 1, # 'ь' - 30: 0, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 57: { # 'Ш' - 37: 1, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 0, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 1, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 2, # 'а' - 21: 0, # 'б' - 10: 1, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 2, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 1, # 'и' - 23: 0, # 'й' - 11: 1, # 'к' - 8: 2, # 'л' - 12: 1, # 'м' - 5: 1, # 'н' - 1: 2, # 'о' - 15: 2, # 'п' - 9: 1, # 'р' - 7: 0, # 'с' - 6: 2, # 'т' - 14: 2, # 'у' - 39: 0, # 'ф' - 26: 1, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 1, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 1, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 63: { # 'Щ' - 37: 1, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 1, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 1, # 'а' - 21: 0, # 'б' - 10: 0, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 1, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 1, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 0, # 'л' - 12: 0, # 'м' - 5: 0, # 'н' - 1: 1, # 'о' - 15: 0, # 'п' - 9: 0, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 1, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 62: { # 'Ы' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 1, # 'В' - 46: 1, # 'Г' - 41: 0, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 1, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 0, # 'О' - 35: 1, # 'П' - 45: 1, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 1, # 'Х' - 58: 1, # 'Ц' - 50: 0, # 'Ч' - 57: 1, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 0, # 'а' - 21: 0, # 'б' - 10: 0, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 0, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 0, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 0, # 'л' - 12: 0, # 'м' - 5: 0, # 'н' - 1: 0, # 'о' - 15: 0, # 'п' - 9: 0, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 0, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 61: { # 'Ь' - 37: 0, # 'А' - 44: 1, # 'Б' - 33: 1, # 'В' - 46: 0, # 'Г' - 41: 1, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 0, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 1, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 1, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 1, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 1, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 1, # 'Ю' - 43: 1, # 'Я' - 3: 0, # 'а' - 21: 0, # 'б' - 10: 0, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 0, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 0, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 0, # 'л' - 12: 0, # 'м' - 5: 0, # 'н' - 1: 0, # 'о' - 15: 0, # 'п' - 9: 0, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 0, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 47: { # 'Э' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 1, # 'В' - 46: 0, # 'Г' - 41: 1, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 1, # 'Й' - 36: 1, # 'К' - 49: 1, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 0, # 'О' - 35: 1, # 'П' - 45: 1, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 1, # 'а' - 21: 1, # 'б' - 10: 2, # 'в' - 19: 1, # 'г' - 13: 2, # 'д' - 2: 0, # 'е' - 24: 1, # 'ж' - 20: 0, # 'з' - 4: 0, # 'и' - 23: 2, # 'й' - 11: 2, # 'к' - 8: 2, # 'л' - 12: 2, # 'м' - 5: 2, # 'н' - 1: 0, # 'о' - 15: 1, # 'п' - 9: 2, # 'р' - 7: 1, # 'с' - 6: 3, # 'т' - 14: 1, # 'у' - 39: 1, # 'ф' - 26: 1, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 1, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 59: { # 'Ю' - 37: 1, # 'А' - 44: 1, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 1, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 1, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 1, # 'Р' - 32: 0, # 'С' - 40: 1, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 1, # 'Ч' - 57: 0, # 'Ш' - 63: 1, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 0, # 'а' - 21: 1, # 'б' - 10: 0, # 'в' - 19: 1, # 'г' - 13: 1, # 'д' - 2: 0, # 'е' - 24: 1, # 'ж' - 20: 0, # 'з' - 4: 0, # 'и' - 23: 0, # 'й' - 11: 1, # 'к' - 8: 2, # 'л' - 12: 1, # 'м' - 5: 2, # 'н' - 1: 0, # 'о' - 15: 1, # 'п' - 9: 1, # 'р' - 7: 1, # 'с' - 6: 0, # 'т' - 14: 0, # 'у' - 39: 0, # 'ф' - 26: 1, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 43: { # 'Я' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 1, # 'В' - 46: 1, # 'Г' - 41: 0, # 'Д' - 48: 1, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 1, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 1, # 'С' - 40: 1, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 1, # 'Х' - 58: 0, # 'Ц' - 50: 1, # 'Ч' - 57: 0, # 'Ш' - 63: 1, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 1, # 'Ю' - 43: 1, # 'Я' - 3: 0, # 'а' - 21: 1, # 'б' - 10: 1, # 'в' - 19: 1, # 'г' - 13: 1, # 'д' - 2: 0, # 'е' - 24: 0, # 'ж' - 20: 1, # 'з' - 4: 0, # 'и' - 23: 1, # 'й' - 11: 1, # 'к' - 8: 1, # 'л' - 12: 1, # 'м' - 5: 2, # 'н' - 1: 0, # 'о' - 15: 1, # 'п' - 9: 1, # 'р' - 7: 1, # 'с' - 6: 0, # 'т' - 14: 0, # 'у' - 39: 0, # 'ф' - 26: 1, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 1, # 'ш' - 29: 1, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 3: { # 'а' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 1, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 1, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 2, # 'а' - 21: 3, # 'б' - 10: 3, # 'в' - 19: 3, # 'г' - 13: 3, # 'д' - 2: 3, # 'е' - 24: 3, # 'ж' - 20: 3, # 'з' - 4: 3, # 'и' - 23: 3, # 'й' - 11: 3, # 'к' - 8: 3, # 'л' - 12: 3, # 'м' - 5: 3, # 'н' - 1: 2, # 'о' - 15: 3, # 'п' - 9: 3, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 3, # 'у' - 39: 2, # 'ф' - 26: 3, # 'х' - 28: 3, # 'ц' - 22: 3, # 'ч' - 25: 3, # 'ш' - 29: 3, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 2, # 'э' - 27: 3, # 'ю' - 16: 3, # 'я' - }, - 21: { # 'б' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 1, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 2, # 'б' - 10: 2, # 'в' - 19: 1, # 'г' - 13: 2, # 'д' - 2: 3, # 'е' - 24: 2, # 'ж' - 20: 1, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 2, # 'к' - 8: 3, # 'л' - 12: 2, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 1, # 'п' - 9: 3, # 'р' - 7: 3, # 'с' - 6: 2, # 'т' - 14: 3, # 'у' - 39: 0, # 'ф' - 26: 2, # 'х' - 28: 1, # 'ц' - 22: 1, # 'ч' - 25: 2, # 'ш' - 29: 3, # 'щ' - 54: 2, # 'ъ' - 18: 3, # 'ы' - 17: 2, # 'ь' - 30: 1, # 'э' - 27: 2, # 'ю' - 16: 3, # 'я' - }, - 10: { # 'в' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 2, # 'б' - 10: 2, # 'в' - 19: 2, # 'г' - 13: 3, # 'д' - 2: 3, # 'е' - 24: 1, # 'ж' - 20: 3, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 3, # 'к' - 8: 3, # 'л' - 12: 2, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 3, # 'п' - 9: 3, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 3, # 'у' - 39: 1, # 'ф' - 26: 2, # 'х' - 28: 2, # 'ц' - 22: 2, # 'ч' - 25: 3, # 'ш' - 29: 2, # 'щ' - 54: 2, # 'ъ' - 18: 3, # 'ы' - 17: 3, # 'ь' - 30: 1, # 'э' - 27: 1, # 'ю' - 16: 3, # 'я' - }, - 19: { # 'г' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 1, # 'б' - 10: 2, # 'в' - 19: 1, # 'г' - 13: 3, # 'д' - 2: 3, # 'е' - 24: 0, # 'ж' - 20: 1, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 2, # 'к' - 8: 3, # 'л' - 12: 2, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 0, # 'п' - 9: 3, # 'р' - 7: 2, # 'с' - 6: 2, # 'т' - 14: 3, # 'у' - 39: 1, # 'ф' - 26: 1, # 'х' - 28: 1, # 'ц' - 22: 2, # 'ч' - 25: 1, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 1, # 'ы' - 17: 1, # 'ь' - 30: 1, # 'э' - 27: 1, # 'ю' - 16: 0, # 'я' - }, - 13: { # 'д' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 2, # 'б' - 10: 3, # 'в' - 19: 2, # 'г' - 13: 2, # 'д' - 2: 3, # 'е' - 24: 2, # 'ж' - 20: 2, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 3, # 'к' - 8: 3, # 'л' - 12: 2, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 2, # 'п' - 9: 3, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 3, # 'у' - 39: 1, # 'ф' - 26: 2, # 'х' - 28: 3, # 'ц' - 22: 2, # 'ч' - 25: 2, # 'ш' - 29: 1, # 'щ' - 54: 2, # 'ъ' - 18: 3, # 'ы' - 17: 3, # 'ь' - 30: 1, # 'э' - 27: 2, # 'ю' - 16: 3, # 'я' - }, - 2: { # 'е' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 2, # 'а' - 21: 3, # 'б' - 10: 3, # 'в' - 19: 3, # 'г' - 13: 3, # 'д' - 2: 3, # 'е' - 24: 3, # 'ж' - 20: 3, # 'з' - 4: 2, # 'и' - 23: 3, # 'й' - 11: 3, # 'к' - 8: 3, # 'л' - 12: 3, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 3, # 'п' - 9: 3, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 2, # 'у' - 39: 2, # 'ф' - 26: 3, # 'х' - 28: 3, # 'ц' - 22: 3, # 'ч' - 25: 3, # 'ш' - 29: 3, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 1, # 'э' - 27: 2, # 'ю' - 16: 3, # 'я' - }, - 24: { # 'ж' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 2, # 'б' - 10: 1, # 'в' - 19: 2, # 'г' - 13: 3, # 'д' - 2: 3, # 'е' - 24: 2, # 'ж' - 20: 1, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 2, # 'к' - 8: 2, # 'л' - 12: 1, # 'м' - 5: 3, # 'н' - 1: 2, # 'о' - 15: 1, # 'п' - 9: 2, # 'р' - 7: 2, # 'с' - 6: 1, # 'т' - 14: 3, # 'у' - 39: 1, # 'ф' - 26: 0, # 'х' - 28: 1, # 'ц' - 22: 2, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 1, # 'ы' - 17: 2, # 'ь' - 30: 1, # 'э' - 27: 1, # 'ю' - 16: 1, # 'я' - }, - 20: { # 'з' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 3, # 'б' - 10: 3, # 'в' - 19: 3, # 'г' - 13: 3, # 'д' - 2: 3, # 'е' - 24: 2, # 'ж' - 20: 2, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 3, # 'к' - 8: 3, # 'л' - 12: 3, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 0, # 'п' - 9: 3, # 'р' - 7: 2, # 'с' - 6: 2, # 'т' - 14: 3, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 1, # 'ц' - 22: 2, # 'ч' - 25: 1, # 'ш' - 29: 0, # 'щ' - 54: 2, # 'ъ' - 18: 3, # 'ы' - 17: 2, # 'ь' - 30: 1, # 'э' - 27: 1, # 'ю' - 16: 3, # 'я' - }, - 4: { # 'и' - 37: 1, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 1, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 3, # 'б' - 10: 3, # 'в' - 19: 3, # 'г' - 13: 3, # 'д' - 2: 3, # 'е' - 24: 3, # 'ж' - 20: 3, # 'з' - 4: 3, # 'и' - 23: 3, # 'й' - 11: 3, # 'к' - 8: 3, # 'л' - 12: 3, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 3, # 'п' - 9: 3, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 2, # 'у' - 39: 2, # 'ф' - 26: 3, # 'х' - 28: 3, # 'ц' - 22: 3, # 'ч' - 25: 3, # 'ш' - 29: 3, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 2, # 'э' - 27: 3, # 'ю' - 16: 3, # 'я' - }, - 23: { # 'й' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 1, # 'а' - 21: 1, # 'б' - 10: 1, # 'в' - 19: 2, # 'г' - 13: 3, # 'д' - 2: 2, # 'е' - 24: 0, # 'ж' - 20: 2, # 'з' - 4: 1, # 'и' - 23: 0, # 'й' - 11: 2, # 'к' - 8: 2, # 'л' - 12: 2, # 'м' - 5: 3, # 'н' - 1: 2, # 'о' - 15: 1, # 'п' - 9: 2, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 1, # 'у' - 39: 2, # 'ф' - 26: 1, # 'х' - 28: 2, # 'ц' - 22: 3, # 'ч' - 25: 2, # 'ш' - 29: 1, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 1, # 'э' - 27: 1, # 'ю' - 16: 2, # 'я' - }, - 11: { # 'к' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 1, # 'б' - 10: 3, # 'в' - 19: 1, # 'г' - 13: 1, # 'д' - 2: 3, # 'е' - 24: 2, # 'ж' - 20: 2, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 2, # 'к' - 8: 3, # 'л' - 12: 1, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 0, # 'п' - 9: 3, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 3, # 'у' - 39: 1, # 'ф' - 26: 2, # 'х' - 28: 2, # 'ц' - 22: 1, # 'ч' - 25: 2, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 1, # 'ы' - 17: 1, # 'ь' - 30: 1, # 'э' - 27: 1, # 'ю' - 16: 1, # 'я' - }, - 8: { # 'л' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 2, # 'б' - 10: 2, # 'в' - 19: 3, # 'г' - 13: 2, # 'д' - 2: 3, # 'е' - 24: 3, # 'ж' - 20: 2, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 3, # 'к' - 8: 3, # 'л' - 12: 2, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 2, # 'п' - 9: 1, # 'р' - 7: 3, # 'с' - 6: 2, # 'т' - 14: 3, # 'у' - 39: 2, # 'ф' - 26: 2, # 'х' - 28: 1, # 'ц' - 22: 3, # 'ч' - 25: 2, # 'ш' - 29: 1, # 'щ' - 54: 0, # 'ъ' - 18: 3, # 'ы' - 17: 3, # 'ь' - 30: 1, # 'э' - 27: 3, # 'ю' - 16: 3, # 'я' - }, - 12: { # 'м' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 2, # 'б' - 10: 2, # 'в' - 19: 2, # 'г' - 13: 1, # 'д' - 2: 3, # 'е' - 24: 1, # 'ж' - 20: 1, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 2, # 'к' - 8: 3, # 'л' - 12: 2, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 2, # 'п' - 9: 2, # 'р' - 7: 3, # 'с' - 6: 2, # 'т' - 14: 3, # 'у' - 39: 2, # 'ф' - 26: 2, # 'х' - 28: 2, # 'ц' - 22: 2, # 'ч' - 25: 1, # 'ш' - 29: 1, # 'щ' - 54: 0, # 'ъ' - 18: 3, # 'ы' - 17: 2, # 'ь' - 30: 2, # 'э' - 27: 1, # 'ю' - 16: 3, # 'я' - }, - 5: { # 'н' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 2, # 'б' - 10: 2, # 'в' - 19: 3, # 'г' - 13: 3, # 'д' - 2: 3, # 'е' - 24: 2, # 'ж' - 20: 2, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 3, # 'к' - 8: 2, # 'л' - 12: 1, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 1, # 'п' - 9: 2, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 3, # 'у' - 39: 2, # 'ф' - 26: 2, # 'х' - 28: 3, # 'ц' - 22: 3, # 'ч' - 25: 2, # 'ш' - 29: 2, # 'щ' - 54: 1, # 'ъ' - 18: 3, # 'ы' - 17: 3, # 'ь' - 30: 1, # 'э' - 27: 3, # 'ю' - 16: 3, # 'я' - }, - 1: { # 'о' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 2, # 'а' - 21: 3, # 'б' - 10: 3, # 'в' - 19: 3, # 'г' - 13: 3, # 'д' - 2: 3, # 'е' - 24: 3, # 'ж' - 20: 3, # 'з' - 4: 3, # 'и' - 23: 3, # 'й' - 11: 3, # 'к' - 8: 3, # 'л' - 12: 3, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 3, # 'п' - 9: 3, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 2, # 'у' - 39: 2, # 'ф' - 26: 3, # 'х' - 28: 2, # 'ц' - 22: 3, # 'ч' - 25: 3, # 'ш' - 29: 3, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 2, # 'э' - 27: 3, # 'ю' - 16: 3, # 'я' - }, - 15: { # 'п' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 1, # 'б' - 10: 0, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 3, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 2, # 'к' - 8: 3, # 'л' - 12: 1, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 2, # 'п' - 9: 3, # 'р' - 7: 2, # 'с' - 6: 2, # 'т' - 14: 3, # 'у' - 39: 1, # 'ф' - 26: 0, # 'х' - 28: 2, # 'ц' - 22: 2, # 'ч' - 25: 1, # 'ш' - 29: 1, # 'щ' - 54: 0, # 'ъ' - 18: 3, # 'ы' - 17: 2, # 'ь' - 30: 1, # 'э' - 27: 1, # 'ю' - 16: 3, # 'я' - }, - 9: { # 'р' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 2, # 'б' - 10: 3, # 'в' - 19: 3, # 'г' - 13: 3, # 'д' - 2: 3, # 'е' - 24: 3, # 'ж' - 20: 2, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 3, # 'к' - 8: 2, # 'л' - 12: 3, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 2, # 'п' - 9: 2, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 3, # 'у' - 39: 2, # 'ф' - 26: 3, # 'х' - 28: 2, # 'ц' - 22: 2, # 'ч' - 25: 3, # 'ш' - 29: 2, # 'щ' - 54: 0, # 'ъ' - 18: 3, # 'ы' - 17: 3, # 'ь' - 30: 2, # 'э' - 27: 2, # 'ю' - 16: 3, # 'я' - }, - 7: { # 'с' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 1, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 2, # 'б' - 10: 3, # 'в' - 19: 2, # 'г' - 13: 3, # 'д' - 2: 3, # 'е' - 24: 2, # 'ж' - 20: 2, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 3, # 'к' - 8: 3, # 'л' - 12: 3, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 3, # 'п' - 9: 3, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 3, # 'у' - 39: 2, # 'ф' - 26: 3, # 'х' - 28: 2, # 'ц' - 22: 3, # 'ч' - 25: 2, # 'ш' - 29: 1, # 'щ' - 54: 2, # 'ъ' - 18: 3, # 'ы' - 17: 3, # 'ь' - 30: 2, # 'э' - 27: 3, # 'ю' - 16: 3, # 'я' - }, - 6: { # 'т' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 2, # 'б' - 10: 3, # 'в' - 19: 2, # 'г' - 13: 2, # 'д' - 2: 3, # 'е' - 24: 1, # 'ж' - 20: 1, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 3, # 'к' - 8: 3, # 'л' - 12: 2, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 2, # 'п' - 9: 3, # 'р' - 7: 3, # 'с' - 6: 2, # 'т' - 14: 3, # 'у' - 39: 2, # 'ф' - 26: 2, # 'х' - 28: 2, # 'ц' - 22: 2, # 'ч' - 25: 2, # 'ш' - 29: 2, # 'щ' - 54: 2, # 'ъ' - 18: 3, # 'ы' - 17: 3, # 'ь' - 30: 2, # 'э' - 27: 2, # 'ю' - 16: 3, # 'я' - }, - 14: { # 'у' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 2, # 'а' - 21: 3, # 'б' - 10: 3, # 'в' - 19: 3, # 'г' - 13: 3, # 'д' - 2: 3, # 'е' - 24: 3, # 'ж' - 20: 3, # 'з' - 4: 2, # 'и' - 23: 2, # 'й' - 11: 3, # 'к' - 8: 3, # 'л' - 12: 3, # 'м' - 5: 3, # 'н' - 1: 2, # 'о' - 15: 3, # 'п' - 9: 3, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 1, # 'у' - 39: 2, # 'ф' - 26: 3, # 'х' - 28: 2, # 'ц' - 22: 3, # 'ч' - 25: 3, # 'ш' - 29: 3, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 2, # 'э' - 27: 3, # 'ю' - 16: 2, # 'я' - }, - 39: { # 'ф' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 1, # 'б' - 10: 0, # 'в' - 19: 1, # 'г' - 13: 0, # 'д' - 2: 3, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 1, # 'к' - 8: 2, # 'л' - 12: 1, # 'м' - 5: 1, # 'н' - 1: 3, # 'о' - 15: 1, # 'п' - 9: 2, # 'р' - 7: 2, # 'с' - 6: 2, # 'т' - 14: 2, # 'у' - 39: 2, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 1, # 'ч' - 25: 1, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 2, # 'ы' - 17: 1, # 'ь' - 30: 2, # 'э' - 27: 1, # 'ю' - 16: 1, # 'я' - }, - 26: { # 'х' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 0, # 'б' - 10: 3, # 'в' - 19: 1, # 'г' - 13: 1, # 'д' - 2: 2, # 'е' - 24: 0, # 'ж' - 20: 1, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 1, # 'к' - 8: 2, # 'л' - 12: 2, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 1, # 'п' - 9: 3, # 'р' - 7: 2, # 'с' - 6: 2, # 'т' - 14: 2, # 'у' - 39: 1, # 'ф' - 26: 1, # 'х' - 28: 1, # 'ц' - 22: 1, # 'ч' - 25: 2, # 'ш' - 29: 0, # 'щ' - 54: 1, # 'ъ' - 18: 0, # 'ы' - 17: 1, # 'ь' - 30: 1, # 'э' - 27: 1, # 'ю' - 16: 0, # 'я' - }, - 28: { # 'ц' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 1, # 'б' - 10: 2, # 'в' - 19: 1, # 'г' - 13: 1, # 'д' - 2: 3, # 'е' - 24: 0, # 'ж' - 20: 1, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 2, # 'к' - 8: 1, # 'л' - 12: 1, # 'м' - 5: 1, # 'н' - 1: 3, # 'о' - 15: 0, # 'п' - 9: 1, # 'р' - 7: 0, # 'с' - 6: 1, # 'т' - 14: 3, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 1, # 'ц' - 22: 0, # 'ч' - 25: 1, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 3, # 'ы' - 17: 1, # 'ь' - 30: 0, # 'э' - 27: 1, # 'ю' - 16: 0, # 'я' - }, - 22: { # 'ч' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 1, # 'б' - 10: 1, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 3, # 'е' - 24: 1, # 'ж' - 20: 0, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 3, # 'к' - 8: 2, # 'л' - 12: 1, # 'м' - 5: 3, # 'н' - 1: 2, # 'о' - 15: 0, # 'п' - 9: 2, # 'р' - 7: 1, # 'с' - 6: 3, # 'т' - 14: 3, # 'у' - 39: 1, # 'ф' - 26: 1, # 'х' - 28: 0, # 'ц' - 22: 1, # 'ч' - 25: 2, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 3, # 'ь' - 30: 0, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 25: { # 'ш' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 1, # 'б' - 10: 2, # 'в' - 19: 1, # 'г' - 13: 0, # 'д' - 2: 3, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 3, # 'к' - 8: 3, # 'л' - 12: 2, # 'м' - 5: 3, # 'н' - 1: 3, # 'о' - 15: 2, # 'п' - 9: 2, # 'р' - 7: 1, # 'с' - 6: 2, # 'т' - 14: 3, # 'у' - 39: 2, # 'ф' - 26: 1, # 'х' - 28: 1, # 'ц' - 22: 1, # 'ч' - 25: 1, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 3, # 'ь' - 30: 1, # 'э' - 27: 1, # 'ю' - 16: 0, # 'я' - }, - 29: { # 'щ' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 3, # 'а' - 21: 0, # 'б' - 10: 1, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 3, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 3, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 0, # 'л' - 12: 1, # 'м' - 5: 2, # 'н' - 1: 1, # 'о' - 15: 0, # 'п' - 9: 2, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 2, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 2, # 'ь' - 30: 0, # 'э' - 27: 0, # 'ю' - 16: 0, # 'я' - }, - 54: { # 'ъ' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 0, # 'а' - 21: 0, # 'б' - 10: 0, # 'в' - 19: 0, # 'г' - 13: 0, # 'д' - 2: 2, # 'е' - 24: 0, # 'ж' - 20: 0, # 'з' - 4: 0, # 'и' - 23: 0, # 'й' - 11: 0, # 'к' - 8: 0, # 'л' - 12: 0, # 'м' - 5: 0, # 'н' - 1: 0, # 'о' - 15: 0, # 'п' - 9: 0, # 'р' - 7: 0, # 'с' - 6: 0, # 'т' - 14: 0, # 'у' - 39: 0, # 'ф' - 26: 0, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 0, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 1, # 'ю' - 16: 2, # 'я' - }, - 18: { # 'ы' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 0, # 'а' - 21: 3, # 'б' - 10: 3, # 'в' - 19: 2, # 'г' - 13: 2, # 'д' - 2: 3, # 'е' - 24: 2, # 'ж' - 20: 2, # 'з' - 4: 2, # 'и' - 23: 3, # 'й' - 11: 3, # 'к' - 8: 3, # 'л' - 12: 3, # 'м' - 5: 3, # 'н' - 1: 1, # 'о' - 15: 3, # 'п' - 9: 3, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 1, # 'у' - 39: 0, # 'ф' - 26: 3, # 'х' - 28: 2, # 'ц' - 22: 3, # 'ч' - 25: 3, # 'ш' - 29: 2, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 0, # 'ю' - 16: 2, # 'я' - }, - 17: { # 'ь' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 0, # 'а' - 21: 2, # 'б' - 10: 2, # 'в' - 19: 2, # 'г' - 13: 2, # 'д' - 2: 3, # 'е' - 24: 1, # 'ж' - 20: 3, # 'з' - 4: 2, # 'и' - 23: 0, # 'й' - 11: 3, # 'к' - 8: 0, # 'л' - 12: 3, # 'м' - 5: 3, # 'н' - 1: 2, # 'о' - 15: 2, # 'п' - 9: 1, # 'р' - 7: 3, # 'с' - 6: 2, # 'т' - 14: 0, # 'у' - 39: 2, # 'ф' - 26: 1, # 'х' - 28: 2, # 'ц' - 22: 2, # 'ч' - 25: 3, # 'ш' - 29: 2, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 1, # 'э' - 27: 3, # 'ю' - 16: 3, # 'я' - }, - 30: { # 'э' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 1, # 'М' - 31: 1, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 1, # 'Р' - 32: 1, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 1, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 0, # 'а' - 21: 1, # 'б' - 10: 1, # 'в' - 19: 1, # 'г' - 13: 2, # 'д' - 2: 1, # 'е' - 24: 0, # 'ж' - 20: 1, # 'з' - 4: 0, # 'и' - 23: 2, # 'й' - 11: 2, # 'к' - 8: 2, # 'л' - 12: 2, # 'м' - 5: 2, # 'н' - 1: 0, # 'о' - 15: 2, # 'п' - 9: 2, # 'р' - 7: 2, # 'с' - 6: 3, # 'т' - 14: 1, # 'у' - 39: 2, # 'ф' - 26: 1, # 'х' - 28: 0, # 'ц' - 22: 0, # 'ч' - 25: 1, # 'ш' - 29: 0, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 1, # 'э' - 27: 1, # 'ю' - 16: 1, # 'я' - }, - 27: { # 'ю' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 2, # 'а' - 21: 3, # 'б' - 10: 1, # 'в' - 19: 2, # 'г' - 13: 3, # 'д' - 2: 1, # 'е' - 24: 2, # 'ж' - 20: 2, # 'з' - 4: 1, # 'и' - 23: 1, # 'й' - 11: 2, # 'к' - 8: 2, # 'л' - 12: 2, # 'м' - 5: 2, # 'н' - 1: 1, # 'о' - 15: 2, # 'п' - 9: 2, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 0, # 'у' - 39: 1, # 'ф' - 26: 2, # 'х' - 28: 2, # 'ц' - 22: 2, # 'ч' - 25: 2, # 'ш' - 29: 3, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 1, # 'э' - 27: 2, # 'ю' - 16: 1, # 'я' - }, - 16: { # 'я' - 37: 0, # 'А' - 44: 0, # 'Б' - 33: 0, # 'В' - 46: 0, # 'Г' - 41: 0, # 'Д' - 48: 0, # 'Е' - 56: 0, # 'Ж' - 51: 0, # 'З' - 42: 0, # 'И' - 60: 0, # 'Й' - 36: 0, # 'К' - 49: 0, # 'Л' - 38: 0, # 'М' - 31: 0, # 'Н' - 34: 0, # 'О' - 35: 0, # 'П' - 45: 0, # 'Р' - 32: 0, # 'С' - 40: 0, # 'Т' - 52: 0, # 'У' - 53: 0, # 'Ф' - 55: 0, # 'Х' - 58: 0, # 'Ц' - 50: 0, # 'Ч' - 57: 0, # 'Ш' - 63: 0, # 'Щ' - 62: 0, # 'Ы' - 61: 0, # 'Ь' - 47: 0, # 'Э' - 59: 0, # 'Ю' - 43: 0, # 'Я' - 3: 0, # 'а' - 21: 2, # 'б' - 10: 3, # 'в' - 19: 2, # 'г' - 13: 3, # 'д' - 2: 3, # 'е' - 24: 3, # 'ж' - 20: 3, # 'з' - 4: 2, # 'и' - 23: 2, # 'й' - 11: 3, # 'к' - 8: 3, # 'л' - 12: 3, # 'м' - 5: 3, # 'н' - 1: 0, # 'о' - 15: 2, # 'п' - 9: 2, # 'р' - 7: 3, # 'с' - 6: 3, # 'т' - 14: 1, # 'у' - 39: 1, # 'ф' - 26: 3, # 'х' - 28: 2, # 'ц' - 22: 2, # 'ч' - 25: 2, # 'ш' - 29: 3, # 'щ' - 54: 0, # 'ъ' - 18: 0, # 'ы' - 17: 0, # 'ь' - 30: 0, # 'э' - 27: 2, # 'ю' - 16: 2, # 'я' - }, -} - -# 255: Undefined characters that did not exist in training text -# 254: Carriage/Return -# 253: symbol (punctuation) that does not belong to word -# 252: 0 - 9 -# 251: Control characters - -# Character Mapping Table(s): -IBM866_RUSSIAN_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 254, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 254, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 253, # ' ' - 33: 253, # '!' - 34: 253, # '"' - 35: 253, # '#' - 36: 253, # '$' - 37: 253, # '%' - 38: 253, # '&' - 39: 253, # "'" - 40: 253, # '(' - 41: 253, # ')' - 42: 253, # '*' - 43: 253, # '+' - 44: 253, # ',' - 45: 253, # '-' - 46: 253, # '.' - 47: 253, # '/' - 48: 252, # '0' - 49: 252, # '1' - 50: 252, # '2' - 51: 252, # '3' - 52: 252, # '4' - 53: 252, # '5' - 54: 252, # '6' - 55: 252, # '7' - 56: 252, # '8' - 57: 252, # '9' - 58: 253, # ':' - 59: 253, # ';' - 60: 253, # '<' - 61: 253, # '=' - 62: 253, # '>' - 63: 253, # '?' - 64: 253, # '@' - 65: 142, # 'A' - 66: 143, # 'B' - 67: 144, # 'C' - 68: 145, # 'D' - 69: 146, # 'E' - 70: 147, # 'F' - 71: 148, # 'G' - 72: 149, # 'H' - 73: 150, # 'I' - 74: 151, # 'J' - 75: 152, # 'K' - 76: 74, # 'L' - 77: 153, # 'M' - 78: 75, # 'N' - 79: 154, # 'O' - 80: 155, # 'P' - 81: 156, # 'Q' - 82: 157, # 'R' - 83: 158, # 'S' - 84: 159, # 'T' - 85: 160, # 'U' - 86: 161, # 'V' - 87: 162, # 'W' - 88: 163, # 'X' - 89: 164, # 'Y' - 90: 165, # 'Z' - 91: 253, # '[' - 92: 253, # '\\' - 93: 253, # ']' - 94: 253, # '^' - 95: 253, # '_' - 96: 253, # '`' - 97: 71, # 'a' - 98: 172, # 'b' - 99: 66, # 'c' - 100: 173, # 'd' - 101: 65, # 'e' - 102: 174, # 'f' - 103: 76, # 'g' - 104: 175, # 'h' - 105: 64, # 'i' - 106: 176, # 'j' - 107: 177, # 'k' - 108: 77, # 'l' - 109: 72, # 'm' - 110: 178, # 'n' - 111: 69, # 'o' - 112: 67, # 'p' - 113: 179, # 'q' - 114: 78, # 'r' - 115: 73, # 's' - 116: 180, # 't' - 117: 181, # 'u' - 118: 79, # 'v' - 119: 182, # 'w' - 120: 183, # 'x' - 121: 184, # 'y' - 122: 185, # 'z' - 123: 253, # '{' - 124: 253, # '|' - 125: 253, # '}' - 126: 253, # '~' - 127: 253, # '\x7f' - 128: 37, # 'А' - 129: 44, # 'Б' - 130: 33, # 'В' - 131: 46, # 'Г' - 132: 41, # 'Д' - 133: 48, # 'Е' - 134: 56, # 'Ж' - 135: 51, # 'З' - 136: 42, # 'И' - 137: 60, # 'Й' - 138: 36, # 'К' - 139: 49, # 'Л' - 140: 38, # 'М' - 141: 31, # 'Н' - 142: 34, # 'О' - 143: 35, # 'П' - 144: 45, # 'Р' - 145: 32, # 'С' - 146: 40, # 'Т' - 147: 52, # 'У' - 148: 53, # 'Ф' - 149: 55, # 'Х' - 150: 58, # 'Ц' - 151: 50, # 'Ч' - 152: 57, # 'Ш' - 153: 63, # 'Щ' - 154: 70, # 'Ъ' - 155: 62, # 'Ы' - 156: 61, # 'Ь' - 157: 47, # 'Э' - 158: 59, # 'Ю' - 159: 43, # 'Я' - 160: 3, # 'а' - 161: 21, # 'б' - 162: 10, # 'в' - 163: 19, # 'г' - 164: 13, # 'д' - 165: 2, # 'е' - 166: 24, # 'ж' - 167: 20, # 'з' - 168: 4, # 'и' - 169: 23, # 'й' - 170: 11, # 'к' - 171: 8, # 'л' - 172: 12, # 'м' - 173: 5, # 'н' - 174: 1, # 'о' - 175: 15, # 'п' - 176: 191, # '░' - 177: 192, # '▒' - 178: 193, # '▓' - 179: 194, # '│' - 180: 195, # '┤' - 181: 196, # '╡' - 182: 197, # '╢' - 183: 198, # '╖' - 184: 199, # '╕' - 185: 200, # '╣' - 186: 201, # '║' - 187: 202, # '╗' - 188: 203, # '╝' - 189: 204, # '╜' - 190: 205, # '╛' - 191: 206, # '┐' - 192: 207, # '└' - 193: 208, # '┴' - 194: 209, # '┬' - 195: 210, # '├' - 196: 211, # '─' - 197: 212, # '┼' - 198: 213, # '╞' - 199: 214, # '╟' - 200: 215, # '╚' - 201: 216, # '╔' - 202: 217, # '╩' - 203: 218, # '╦' - 204: 219, # '╠' - 205: 220, # '═' - 206: 221, # '╬' - 207: 222, # '╧' - 208: 223, # '╨' - 209: 224, # '╤' - 210: 225, # '╥' - 211: 226, # '╙' - 212: 227, # '╘' - 213: 228, # '╒' - 214: 229, # '╓' - 215: 230, # '╫' - 216: 231, # '╪' - 217: 232, # '┘' - 218: 233, # '┌' - 219: 234, # '█' - 220: 235, # '▄' - 221: 236, # '▌' - 222: 237, # '▐' - 223: 238, # '▀' - 224: 9, # 'р' - 225: 7, # 'с' - 226: 6, # 'т' - 227: 14, # 'у' - 228: 39, # 'ф' - 229: 26, # 'х' - 230: 28, # 'ц' - 231: 22, # 'ч' - 232: 25, # 'ш' - 233: 29, # 'щ' - 234: 54, # 'ъ' - 235: 18, # 'ы' - 236: 17, # 'ь' - 237: 30, # 'э' - 238: 27, # 'ю' - 239: 16, # 'я' - 240: 239, # 'Ё' - 241: 68, # 'ё' - 242: 240, # 'Є' - 243: 241, # 'є' - 244: 242, # 'Ї' - 245: 243, # 'ї' - 246: 244, # 'Ў' - 247: 245, # 'ў' - 248: 246, # '°' - 249: 247, # '∙' - 250: 248, # '·' - 251: 249, # '√' - 252: 250, # '№' - 253: 251, # '¤' - 254: 252, # '■' - 255: 255, # '\xa0' -} - -IBM866_RUSSIAN_MODEL = SingleByteCharSetModel( - charset_name="IBM866", - language="Russian", - char_to_order_map=IBM866_RUSSIAN_CHAR_TO_ORDER, - language_model=RUSSIAN_LANG_MODEL, - typical_positive_ratio=0.976601, - keep_ascii_letters=False, - alphabet="ЁАБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯабвгдежзийклмнопрстуфхцчшщъыьэюяё", -) - -WINDOWS_1251_RUSSIAN_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 254, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 254, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 253, # ' ' - 33: 253, # '!' - 34: 253, # '"' - 35: 253, # '#' - 36: 253, # '$' - 37: 253, # '%' - 38: 253, # '&' - 39: 253, # "'" - 40: 253, # '(' - 41: 253, # ')' - 42: 253, # '*' - 43: 253, # '+' - 44: 253, # ',' - 45: 253, # '-' - 46: 253, # '.' - 47: 253, # '/' - 48: 252, # '0' - 49: 252, # '1' - 50: 252, # '2' - 51: 252, # '3' - 52: 252, # '4' - 53: 252, # '5' - 54: 252, # '6' - 55: 252, # '7' - 56: 252, # '8' - 57: 252, # '9' - 58: 253, # ':' - 59: 253, # ';' - 60: 253, # '<' - 61: 253, # '=' - 62: 253, # '>' - 63: 253, # '?' - 64: 253, # '@' - 65: 142, # 'A' - 66: 143, # 'B' - 67: 144, # 'C' - 68: 145, # 'D' - 69: 146, # 'E' - 70: 147, # 'F' - 71: 148, # 'G' - 72: 149, # 'H' - 73: 150, # 'I' - 74: 151, # 'J' - 75: 152, # 'K' - 76: 74, # 'L' - 77: 153, # 'M' - 78: 75, # 'N' - 79: 154, # 'O' - 80: 155, # 'P' - 81: 156, # 'Q' - 82: 157, # 'R' - 83: 158, # 'S' - 84: 159, # 'T' - 85: 160, # 'U' - 86: 161, # 'V' - 87: 162, # 'W' - 88: 163, # 'X' - 89: 164, # 'Y' - 90: 165, # 'Z' - 91: 253, # '[' - 92: 253, # '\\' - 93: 253, # ']' - 94: 253, # '^' - 95: 253, # '_' - 96: 253, # '`' - 97: 71, # 'a' - 98: 172, # 'b' - 99: 66, # 'c' - 100: 173, # 'd' - 101: 65, # 'e' - 102: 174, # 'f' - 103: 76, # 'g' - 104: 175, # 'h' - 105: 64, # 'i' - 106: 176, # 'j' - 107: 177, # 'k' - 108: 77, # 'l' - 109: 72, # 'm' - 110: 178, # 'n' - 111: 69, # 'o' - 112: 67, # 'p' - 113: 179, # 'q' - 114: 78, # 'r' - 115: 73, # 's' - 116: 180, # 't' - 117: 181, # 'u' - 118: 79, # 'v' - 119: 182, # 'w' - 120: 183, # 'x' - 121: 184, # 'y' - 122: 185, # 'z' - 123: 253, # '{' - 124: 253, # '|' - 125: 253, # '}' - 126: 253, # '~' - 127: 253, # '\x7f' - 128: 191, # 'Ђ' - 129: 192, # 'Ѓ' - 130: 193, # '‚' - 131: 194, # 'ѓ' - 132: 195, # '„' - 133: 196, # '…' - 134: 197, # '†' - 135: 198, # '‡' - 136: 199, # '€' - 137: 200, # '‰' - 138: 201, # 'Љ' - 139: 202, # '‹' - 140: 203, # 'Њ' - 141: 204, # 'Ќ' - 142: 205, # 'Ћ' - 143: 206, # 'Џ' - 144: 207, # 'ђ' - 145: 208, # '‘' - 146: 209, # '’' - 147: 210, # '“' - 148: 211, # '”' - 149: 212, # '•' - 150: 213, # '–' - 151: 214, # '—' - 152: 215, # None - 153: 216, # '™' - 154: 217, # 'љ' - 155: 218, # '›' - 156: 219, # 'њ' - 157: 220, # 'ќ' - 158: 221, # 'ћ' - 159: 222, # 'џ' - 160: 223, # '\xa0' - 161: 224, # 'Ў' - 162: 225, # 'ў' - 163: 226, # 'Ј' - 164: 227, # '¤' - 165: 228, # 'Ґ' - 166: 229, # '¦' - 167: 230, # '§' - 168: 231, # 'Ё' - 169: 232, # '©' - 170: 233, # 'Є' - 171: 234, # '«' - 172: 235, # '¬' - 173: 236, # '\xad' - 174: 237, # '®' - 175: 238, # 'Ї' - 176: 239, # '°' - 177: 240, # '±' - 178: 241, # 'І' - 179: 242, # 'і' - 180: 243, # 'ґ' - 181: 244, # 'µ' - 182: 245, # '¶' - 183: 246, # '·' - 184: 68, # 'ё' - 185: 247, # '№' - 186: 248, # 'є' - 187: 249, # '»' - 188: 250, # 'ј' - 189: 251, # 'Ѕ' - 190: 252, # 'ѕ' - 191: 253, # 'ї' - 192: 37, # 'А' - 193: 44, # 'Б' - 194: 33, # 'В' - 195: 46, # 'Г' - 196: 41, # 'Д' - 197: 48, # 'Е' - 198: 56, # 'Ж' - 199: 51, # 'З' - 200: 42, # 'И' - 201: 60, # 'Й' - 202: 36, # 'К' - 203: 49, # 'Л' - 204: 38, # 'М' - 205: 31, # 'Н' - 206: 34, # 'О' - 207: 35, # 'П' - 208: 45, # 'Р' - 209: 32, # 'С' - 210: 40, # 'Т' - 211: 52, # 'У' - 212: 53, # 'Ф' - 213: 55, # 'Х' - 214: 58, # 'Ц' - 215: 50, # 'Ч' - 216: 57, # 'Ш' - 217: 63, # 'Щ' - 218: 70, # 'Ъ' - 219: 62, # 'Ы' - 220: 61, # 'Ь' - 221: 47, # 'Э' - 222: 59, # 'Ю' - 223: 43, # 'Я' - 224: 3, # 'а' - 225: 21, # 'б' - 226: 10, # 'в' - 227: 19, # 'г' - 228: 13, # 'д' - 229: 2, # 'е' - 230: 24, # 'ж' - 231: 20, # 'з' - 232: 4, # 'и' - 233: 23, # 'й' - 234: 11, # 'к' - 235: 8, # 'л' - 236: 12, # 'м' - 237: 5, # 'н' - 238: 1, # 'о' - 239: 15, # 'п' - 240: 9, # 'р' - 241: 7, # 'с' - 242: 6, # 'т' - 243: 14, # 'у' - 244: 39, # 'ф' - 245: 26, # 'х' - 246: 28, # 'ц' - 247: 22, # 'ч' - 248: 25, # 'ш' - 249: 29, # 'щ' - 250: 54, # 'ъ' - 251: 18, # 'ы' - 252: 17, # 'ь' - 253: 30, # 'э' - 254: 27, # 'ю' - 255: 16, # 'я' -} - -WINDOWS_1251_RUSSIAN_MODEL = SingleByteCharSetModel( - charset_name="windows-1251", - language="Russian", - char_to_order_map=WINDOWS_1251_RUSSIAN_CHAR_TO_ORDER, - language_model=RUSSIAN_LANG_MODEL, - typical_positive_ratio=0.976601, - keep_ascii_letters=False, - alphabet="ЁАБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯабвгдежзийклмнопрстуфхцчшщъыьэюяё", -) - -IBM855_RUSSIAN_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 254, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 254, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 253, # ' ' - 33: 253, # '!' - 34: 253, # '"' - 35: 253, # '#' - 36: 253, # '$' - 37: 253, # '%' - 38: 253, # '&' - 39: 253, # "'" - 40: 253, # '(' - 41: 253, # ')' - 42: 253, # '*' - 43: 253, # '+' - 44: 253, # ',' - 45: 253, # '-' - 46: 253, # '.' - 47: 253, # '/' - 48: 252, # '0' - 49: 252, # '1' - 50: 252, # '2' - 51: 252, # '3' - 52: 252, # '4' - 53: 252, # '5' - 54: 252, # '6' - 55: 252, # '7' - 56: 252, # '8' - 57: 252, # '9' - 58: 253, # ':' - 59: 253, # ';' - 60: 253, # '<' - 61: 253, # '=' - 62: 253, # '>' - 63: 253, # '?' - 64: 253, # '@' - 65: 142, # 'A' - 66: 143, # 'B' - 67: 144, # 'C' - 68: 145, # 'D' - 69: 146, # 'E' - 70: 147, # 'F' - 71: 148, # 'G' - 72: 149, # 'H' - 73: 150, # 'I' - 74: 151, # 'J' - 75: 152, # 'K' - 76: 74, # 'L' - 77: 153, # 'M' - 78: 75, # 'N' - 79: 154, # 'O' - 80: 155, # 'P' - 81: 156, # 'Q' - 82: 157, # 'R' - 83: 158, # 'S' - 84: 159, # 'T' - 85: 160, # 'U' - 86: 161, # 'V' - 87: 162, # 'W' - 88: 163, # 'X' - 89: 164, # 'Y' - 90: 165, # 'Z' - 91: 253, # '[' - 92: 253, # '\\' - 93: 253, # ']' - 94: 253, # '^' - 95: 253, # '_' - 96: 253, # '`' - 97: 71, # 'a' - 98: 172, # 'b' - 99: 66, # 'c' - 100: 173, # 'd' - 101: 65, # 'e' - 102: 174, # 'f' - 103: 76, # 'g' - 104: 175, # 'h' - 105: 64, # 'i' - 106: 176, # 'j' - 107: 177, # 'k' - 108: 77, # 'l' - 109: 72, # 'm' - 110: 178, # 'n' - 111: 69, # 'o' - 112: 67, # 'p' - 113: 179, # 'q' - 114: 78, # 'r' - 115: 73, # 's' - 116: 180, # 't' - 117: 181, # 'u' - 118: 79, # 'v' - 119: 182, # 'w' - 120: 183, # 'x' - 121: 184, # 'y' - 122: 185, # 'z' - 123: 253, # '{' - 124: 253, # '|' - 125: 253, # '}' - 126: 253, # '~' - 127: 253, # '\x7f' - 128: 191, # 'ђ' - 129: 192, # 'Ђ' - 130: 193, # 'ѓ' - 131: 194, # 'Ѓ' - 132: 68, # 'ё' - 133: 195, # 'Ё' - 134: 196, # 'є' - 135: 197, # 'Є' - 136: 198, # 'ѕ' - 137: 199, # 'Ѕ' - 138: 200, # 'і' - 139: 201, # 'І' - 140: 202, # 'ї' - 141: 203, # 'Ї' - 142: 204, # 'ј' - 143: 205, # 'Ј' - 144: 206, # 'љ' - 145: 207, # 'Љ' - 146: 208, # 'њ' - 147: 209, # 'Њ' - 148: 210, # 'ћ' - 149: 211, # 'Ћ' - 150: 212, # 'ќ' - 151: 213, # 'Ќ' - 152: 214, # 'ў' - 153: 215, # 'Ў' - 154: 216, # 'џ' - 155: 217, # 'Џ' - 156: 27, # 'ю' - 157: 59, # 'Ю' - 158: 54, # 'ъ' - 159: 70, # 'Ъ' - 160: 3, # 'а' - 161: 37, # 'А' - 162: 21, # 'б' - 163: 44, # 'Б' - 164: 28, # 'ц' - 165: 58, # 'Ц' - 166: 13, # 'д' - 167: 41, # 'Д' - 168: 2, # 'е' - 169: 48, # 'Е' - 170: 39, # 'ф' - 171: 53, # 'Ф' - 172: 19, # 'г' - 173: 46, # 'Г' - 174: 218, # '«' - 175: 219, # '»' - 176: 220, # '░' - 177: 221, # '▒' - 178: 222, # '▓' - 179: 223, # '│' - 180: 224, # '┤' - 181: 26, # 'х' - 182: 55, # 'Х' - 183: 4, # 'и' - 184: 42, # 'И' - 185: 225, # '╣' - 186: 226, # '║' - 187: 227, # '╗' - 188: 228, # '╝' - 189: 23, # 'й' - 190: 60, # 'Й' - 191: 229, # '┐' - 192: 230, # '└' - 193: 231, # '┴' - 194: 232, # '┬' - 195: 233, # '├' - 196: 234, # '─' - 197: 235, # '┼' - 198: 11, # 'к' - 199: 36, # 'К' - 200: 236, # '╚' - 201: 237, # '╔' - 202: 238, # '╩' - 203: 239, # '╦' - 204: 240, # '╠' - 205: 241, # '═' - 206: 242, # '╬' - 207: 243, # '¤' - 208: 8, # 'л' - 209: 49, # 'Л' - 210: 12, # 'м' - 211: 38, # 'М' - 212: 5, # 'н' - 213: 31, # 'Н' - 214: 1, # 'о' - 215: 34, # 'О' - 216: 15, # 'п' - 217: 244, # '┘' - 218: 245, # '┌' - 219: 246, # '█' - 220: 247, # '▄' - 221: 35, # 'П' - 222: 16, # 'я' - 223: 248, # '▀' - 224: 43, # 'Я' - 225: 9, # 'р' - 226: 45, # 'Р' - 227: 7, # 'с' - 228: 32, # 'С' - 229: 6, # 'т' - 230: 40, # 'Т' - 231: 14, # 'у' - 232: 52, # 'У' - 233: 24, # 'ж' - 234: 56, # 'Ж' - 235: 10, # 'в' - 236: 33, # 'В' - 237: 17, # 'ь' - 238: 61, # 'Ь' - 239: 249, # '№' - 240: 250, # '\xad' - 241: 18, # 'ы' - 242: 62, # 'Ы' - 243: 20, # 'з' - 244: 51, # 'З' - 245: 25, # 'ш' - 246: 57, # 'Ш' - 247: 30, # 'э' - 248: 47, # 'Э' - 249: 29, # 'щ' - 250: 63, # 'Щ' - 251: 22, # 'ч' - 252: 50, # 'Ч' - 253: 251, # '§' - 254: 252, # '■' - 255: 255, # '\xa0' -} - -IBM855_RUSSIAN_MODEL = SingleByteCharSetModel( - charset_name="IBM855", - language="Russian", - char_to_order_map=IBM855_RUSSIAN_CHAR_TO_ORDER, - language_model=RUSSIAN_LANG_MODEL, - typical_positive_ratio=0.976601, - keep_ascii_letters=False, - alphabet="ЁАБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯабвгдежзийклмнопрстуфхцчшщъыьэюяё", -) - -KOI8_R_RUSSIAN_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 254, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 254, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 253, # ' ' - 33: 253, # '!' - 34: 253, # '"' - 35: 253, # '#' - 36: 253, # '$' - 37: 253, # '%' - 38: 253, # '&' - 39: 253, # "'" - 40: 253, # '(' - 41: 253, # ')' - 42: 253, # '*' - 43: 253, # '+' - 44: 253, # ',' - 45: 253, # '-' - 46: 253, # '.' - 47: 253, # '/' - 48: 252, # '0' - 49: 252, # '1' - 50: 252, # '2' - 51: 252, # '3' - 52: 252, # '4' - 53: 252, # '5' - 54: 252, # '6' - 55: 252, # '7' - 56: 252, # '8' - 57: 252, # '9' - 58: 253, # ':' - 59: 253, # ';' - 60: 253, # '<' - 61: 253, # '=' - 62: 253, # '>' - 63: 253, # '?' - 64: 253, # '@' - 65: 142, # 'A' - 66: 143, # 'B' - 67: 144, # 'C' - 68: 145, # 'D' - 69: 146, # 'E' - 70: 147, # 'F' - 71: 148, # 'G' - 72: 149, # 'H' - 73: 150, # 'I' - 74: 151, # 'J' - 75: 152, # 'K' - 76: 74, # 'L' - 77: 153, # 'M' - 78: 75, # 'N' - 79: 154, # 'O' - 80: 155, # 'P' - 81: 156, # 'Q' - 82: 157, # 'R' - 83: 158, # 'S' - 84: 159, # 'T' - 85: 160, # 'U' - 86: 161, # 'V' - 87: 162, # 'W' - 88: 163, # 'X' - 89: 164, # 'Y' - 90: 165, # 'Z' - 91: 253, # '[' - 92: 253, # '\\' - 93: 253, # ']' - 94: 253, # '^' - 95: 253, # '_' - 96: 253, # '`' - 97: 71, # 'a' - 98: 172, # 'b' - 99: 66, # 'c' - 100: 173, # 'd' - 101: 65, # 'e' - 102: 174, # 'f' - 103: 76, # 'g' - 104: 175, # 'h' - 105: 64, # 'i' - 106: 176, # 'j' - 107: 177, # 'k' - 108: 77, # 'l' - 109: 72, # 'm' - 110: 178, # 'n' - 111: 69, # 'o' - 112: 67, # 'p' - 113: 179, # 'q' - 114: 78, # 'r' - 115: 73, # 's' - 116: 180, # 't' - 117: 181, # 'u' - 118: 79, # 'v' - 119: 182, # 'w' - 120: 183, # 'x' - 121: 184, # 'y' - 122: 185, # 'z' - 123: 253, # '{' - 124: 253, # '|' - 125: 253, # '}' - 126: 253, # '~' - 127: 253, # '\x7f' - 128: 191, # '─' - 129: 192, # '│' - 130: 193, # '┌' - 131: 194, # '┐' - 132: 195, # '└' - 133: 196, # '┘' - 134: 197, # '├' - 135: 198, # '┤' - 136: 199, # '┬' - 137: 200, # '┴' - 138: 201, # '┼' - 139: 202, # '▀' - 140: 203, # '▄' - 141: 204, # '█' - 142: 205, # '▌' - 143: 206, # '▐' - 144: 207, # '░' - 145: 208, # '▒' - 146: 209, # '▓' - 147: 210, # '⌠' - 148: 211, # '■' - 149: 212, # '∙' - 150: 213, # '√' - 151: 214, # '≈' - 152: 215, # '≤' - 153: 216, # '≥' - 154: 217, # '\xa0' - 155: 218, # '⌡' - 156: 219, # '°' - 157: 220, # '²' - 158: 221, # '·' - 159: 222, # '÷' - 160: 223, # '═' - 161: 224, # '║' - 162: 225, # '╒' - 163: 68, # 'ё' - 164: 226, # '╓' - 165: 227, # '╔' - 166: 228, # '╕' - 167: 229, # '╖' - 168: 230, # '╗' - 169: 231, # '╘' - 170: 232, # '╙' - 171: 233, # '╚' - 172: 234, # '╛' - 173: 235, # '╜' - 174: 236, # '╝' - 175: 237, # '╞' - 176: 238, # '╟' - 177: 239, # '╠' - 178: 240, # '╡' - 179: 241, # 'Ё' - 180: 242, # '╢' - 181: 243, # '╣' - 182: 244, # '╤' - 183: 245, # '╥' - 184: 246, # '╦' - 185: 247, # '╧' - 186: 248, # '╨' - 187: 249, # '╩' - 188: 250, # '╪' - 189: 251, # '╫' - 190: 252, # '╬' - 191: 253, # '©' - 192: 27, # 'ю' - 193: 3, # 'а' - 194: 21, # 'б' - 195: 28, # 'ц' - 196: 13, # 'д' - 197: 2, # 'е' - 198: 39, # 'ф' - 199: 19, # 'г' - 200: 26, # 'х' - 201: 4, # 'и' - 202: 23, # 'й' - 203: 11, # 'к' - 204: 8, # 'л' - 205: 12, # 'м' - 206: 5, # 'н' - 207: 1, # 'о' - 208: 15, # 'п' - 209: 16, # 'я' - 210: 9, # 'р' - 211: 7, # 'с' - 212: 6, # 'т' - 213: 14, # 'у' - 214: 24, # 'ж' - 215: 10, # 'в' - 216: 17, # 'ь' - 217: 18, # 'ы' - 218: 20, # 'з' - 219: 25, # 'ш' - 220: 30, # 'э' - 221: 29, # 'щ' - 222: 22, # 'ч' - 223: 54, # 'ъ' - 224: 59, # 'Ю' - 225: 37, # 'А' - 226: 44, # 'Б' - 227: 58, # 'Ц' - 228: 41, # 'Д' - 229: 48, # 'Е' - 230: 53, # 'Ф' - 231: 46, # 'Г' - 232: 55, # 'Х' - 233: 42, # 'И' - 234: 60, # 'Й' - 235: 36, # 'К' - 236: 49, # 'Л' - 237: 38, # 'М' - 238: 31, # 'Н' - 239: 34, # 'О' - 240: 35, # 'П' - 241: 43, # 'Я' - 242: 45, # 'Р' - 243: 32, # 'С' - 244: 40, # 'Т' - 245: 52, # 'У' - 246: 56, # 'Ж' - 247: 33, # 'В' - 248: 61, # 'Ь' - 249: 62, # 'Ы' - 250: 51, # 'З' - 251: 57, # 'Ш' - 252: 47, # 'Э' - 253: 63, # 'Щ' - 254: 50, # 'Ч' - 255: 70, # 'Ъ' -} - -KOI8_R_RUSSIAN_MODEL = SingleByteCharSetModel( - charset_name="KOI8-R", - language="Russian", - char_to_order_map=KOI8_R_RUSSIAN_CHAR_TO_ORDER, - language_model=RUSSIAN_LANG_MODEL, - typical_positive_ratio=0.976601, - keep_ascii_letters=False, - alphabet="ЁАБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯабвгдежзийклмнопрстуфхцчшщъыьэюяё", -) - -MACCYRILLIC_RUSSIAN_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 254, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 254, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 253, # ' ' - 33: 253, # '!' - 34: 253, # '"' - 35: 253, # '#' - 36: 253, # '$' - 37: 253, # '%' - 38: 253, # '&' - 39: 253, # "'" - 40: 253, # '(' - 41: 253, # ')' - 42: 253, # '*' - 43: 253, # '+' - 44: 253, # ',' - 45: 253, # '-' - 46: 253, # '.' - 47: 253, # '/' - 48: 252, # '0' - 49: 252, # '1' - 50: 252, # '2' - 51: 252, # '3' - 52: 252, # '4' - 53: 252, # '5' - 54: 252, # '6' - 55: 252, # '7' - 56: 252, # '8' - 57: 252, # '9' - 58: 253, # ':' - 59: 253, # ';' - 60: 253, # '<' - 61: 253, # '=' - 62: 253, # '>' - 63: 253, # '?' - 64: 253, # '@' - 65: 142, # 'A' - 66: 143, # 'B' - 67: 144, # 'C' - 68: 145, # 'D' - 69: 146, # 'E' - 70: 147, # 'F' - 71: 148, # 'G' - 72: 149, # 'H' - 73: 150, # 'I' - 74: 151, # 'J' - 75: 152, # 'K' - 76: 74, # 'L' - 77: 153, # 'M' - 78: 75, # 'N' - 79: 154, # 'O' - 80: 155, # 'P' - 81: 156, # 'Q' - 82: 157, # 'R' - 83: 158, # 'S' - 84: 159, # 'T' - 85: 160, # 'U' - 86: 161, # 'V' - 87: 162, # 'W' - 88: 163, # 'X' - 89: 164, # 'Y' - 90: 165, # 'Z' - 91: 253, # '[' - 92: 253, # '\\' - 93: 253, # ']' - 94: 253, # '^' - 95: 253, # '_' - 96: 253, # '`' - 97: 71, # 'a' - 98: 172, # 'b' - 99: 66, # 'c' - 100: 173, # 'd' - 101: 65, # 'e' - 102: 174, # 'f' - 103: 76, # 'g' - 104: 175, # 'h' - 105: 64, # 'i' - 106: 176, # 'j' - 107: 177, # 'k' - 108: 77, # 'l' - 109: 72, # 'm' - 110: 178, # 'n' - 111: 69, # 'o' - 112: 67, # 'p' - 113: 179, # 'q' - 114: 78, # 'r' - 115: 73, # 's' - 116: 180, # 't' - 117: 181, # 'u' - 118: 79, # 'v' - 119: 182, # 'w' - 120: 183, # 'x' - 121: 184, # 'y' - 122: 185, # 'z' - 123: 253, # '{' - 124: 253, # '|' - 125: 253, # '}' - 126: 253, # '~' - 127: 253, # '\x7f' - 128: 37, # 'А' - 129: 44, # 'Б' - 130: 33, # 'В' - 131: 46, # 'Г' - 132: 41, # 'Д' - 133: 48, # 'Е' - 134: 56, # 'Ж' - 135: 51, # 'З' - 136: 42, # 'И' - 137: 60, # 'Й' - 138: 36, # 'К' - 139: 49, # 'Л' - 140: 38, # 'М' - 141: 31, # 'Н' - 142: 34, # 'О' - 143: 35, # 'П' - 144: 45, # 'Р' - 145: 32, # 'С' - 146: 40, # 'Т' - 147: 52, # 'У' - 148: 53, # 'Ф' - 149: 55, # 'Х' - 150: 58, # 'Ц' - 151: 50, # 'Ч' - 152: 57, # 'Ш' - 153: 63, # 'Щ' - 154: 70, # 'Ъ' - 155: 62, # 'Ы' - 156: 61, # 'Ь' - 157: 47, # 'Э' - 158: 59, # 'Ю' - 159: 43, # 'Я' - 160: 191, # '†' - 161: 192, # '°' - 162: 193, # 'Ґ' - 163: 194, # '£' - 164: 195, # '§' - 165: 196, # '•' - 166: 197, # '¶' - 167: 198, # 'І' - 168: 199, # '®' - 169: 200, # '©' - 170: 201, # '™' - 171: 202, # 'Ђ' - 172: 203, # 'ђ' - 173: 204, # '≠' - 174: 205, # 'Ѓ' - 175: 206, # 'ѓ' - 176: 207, # '∞' - 177: 208, # '±' - 178: 209, # '≤' - 179: 210, # '≥' - 180: 211, # 'і' - 181: 212, # 'µ' - 182: 213, # 'ґ' - 183: 214, # 'Ј' - 184: 215, # 'Є' - 185: 216, # 'є' - 186: 217, # 'Ї' - 187: 218, # 'ї' - 188: 219, # 'Љ' - 189: 220, # 'љ' - 190: 221, # 'Њ' - 191: 222, # 'њ' - 192: 223, # 'ј' - 193: 224, # 'Ѕ' - 194: 225, # '¬' - 195: 226, # '√' - 196: 227, # 'ƒ' - 197: 228, # '≈' - 198: 229, # '∆' - 199: 230, # '«' - 200: 231, # '»' - 201: 232, # '…' - 202: 233, # '\xa0' - 203: 234, # 'Ћ' - 204: 235, # 'ћ' - 205: 236, # 'Ќ' - 206: 237, # 'ќ' - 207: 238, # 'ѕ' - 208: 239, # '–' - 209: 240, # '—' - 210: 241, # '“' - 211: 242, # '”' - 212: 243, # '‘' - 213: 244, # '’' - 214: 245, # '÷' - 215: 246, # '„' - 216: 247, # 'Ў' - 217: 248, # 'ў' - 218: 249, # 'Џ' - 219: 250, # 'џ' - 220: 251, # '№' - 221: 252, # 'Ё' - 222: 68, # 'ё' - 223: 16, # 'я' - 224: 3, # 'а' - 225: 21, # 'б' - 226: 10, # 'в' - 227: 19, # 'г' - 228: 13, # 'д' - 229: 2, # 'е' - 230: 24, # 'ж' - 231: 20, # 'з' - 232: 4, # 'и' - 233: 23, # 'й' - 234: 11, # 'к' - 235: 8, # 'л' - 236: 12, # 'м' - 237: 5, # 'н' - 238: 1, # 'о' - 239: 15, # 'п' - 240: 9, # 'р' - 241: 7, # 'с' - 242: 6, # 'т' - 243: 14, # 'у' - 244: 39, # 'ф' - 245: 26, # 'х' - 246: 28, # 'ц' - 247: 22, # 'ч' - 248: 25, # 'ш' - 249: 29, # 'щ' - 250: 54, # 'ъ' - 251: 18, # 'ы' - 252: 17, # 'ь' - 253: 30, # 'э' - 254: 27, # 'ю' - 255: 255, # '€' -} - -MACCYRILLIC_RUSSIAN_MODEL = SingleByteCharSetModel( - charset_name="MacCyrillic", - language="Russian", - char_to_order_map=MACCYRILLIC_RUSSIAN_CHAR_TO_ORDER, - language_model=RUSSIAN_LANG_MODEL, - typical_positive_ratio=0.976601, - keep_ascii_letters=False, - alphabet="ЁАБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯабвгдежзийклмнопрстуфхцчшщъыьэюяё", -) - -ISO_8859_5_RUSSIAN_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 254, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 254, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 253, # ' ' - 33: 253, # '!' - 34: 253, # '"' - 35: 253, # '#' - 36: 253, # '$' - 37: 253, # '%' - 38: 253, # '&' - 39: 253, # "'" - 40: 253, # '(' - 41: 253, # ')' - 42: 253, # '*' - 43: 253, # '+' - 44: 253, # ',' - 45: 253, # '-' - 46: 253, # '.' - 47: 253, # '/' - 48: 252, # '0' - 49: 252, # '1' - 50: 252, # '2' - 51: 252, # '3' - 52: 252, # '4' - 53: 252, # '5' - 54: 252, # '6' - 55: 252, # '7' - 56: 252, # '8' - 57: 252, # '9' - 58: 253, # ':' - 59: 253, # ';' - 60: 253, # '<' - 61: 253, # '=' - 62: 253, # '>' - 63: 253, # '?' - 64: 253, # '@' - 65: 142, # 'A' - 66: 143, # 'B' - 67: 144, # 'C' - 68: 145, # 'D' - 69: 146, # 'E' - 70: 147, # 'F' - 71: 148, # 'G' - 72: 149, # 'H' - 73: 150, # 'I' - 74: 151, # 'J' - 75: 152, # 'K' - 76: 74, # 'L' - 77: 153, # 'M' - 78: 75, # 'N' - 79: 154, # 'O' - 80: 155, # 'P' - 81: 156, # 'Q' - 82: 157, # 'R' - 83: 158, # 'S' - 84: 159, # 'T' - 85: 160, # 'U' - 86: 161, # 'V' - 87: 162, # 'W' - 88: 163, # 'X' - 89: 164, # 'Y' - 90: 165, # 'Z' - 91: 253, # '[' - 92: 253, # '\\' - 93: 253, # ']' - 94: 253, # '^' - 95: 253, # '_' - 96: 253, # '`' - 97: 71, # 'a' - 98: 172, # 'b' - 99: 66, # 'c' - 100: 173, # 'd' - 101: 65, # 'e' - 102: 174, # 'f' - 103: 76, # 'g' - 104: 175, # 'h' - 105: 64, # 'i' - 106: 176, # 'j' - 107: 177, # 'k' - 108: 77, # 'l' - 109: 72, # 'm' - 110: 178, # 'n' - 111: 69, # 'o' - 112: 67, # 'p' - 113: 179, # 'q' - 114: 78, # 'r' - 115: 73, # 's' - 116: 180, # 't' - 117: 181, # 'u' - 118: 79, # 'v' - 119: 182, # 'w' - 120: 183, # 'x' - 121: 184, # 'y' - 122: 185, # 'z' - 123: 253, # '{' - 124: 253, # '|' - 125: 253, # '}' - 126: 253, # '~' - 127: 253, # '\x7f' - 128: 191, # '\x80' - 129: 192, # '\x81' - 130: 193, # '\x82' - 131: 194, # '\x83' - 132: 195, # '\x84' - 133: 196, # '\x85' - 134: 197, # '\x86' - 135: 198, # '\x87' - 136: 199, # '\x88' - 137: 200, # '\x89' - 138: 201, # '\x8a' - 139: 202, # '\x8b' - 140: 203, # '\x8c' - 141: 204, # '\x8d' - 142: 205, # '\x8e' - 143: 206, # '\x8f' - 144: 207, # '\x90' - 145: 208, # '\x91' - 146: 209, # '\x92' - 147: 210, # '\x93' - 148: 211, # '\x94' - 149: 212, # '\x95' - 150: 213, # '\x96' - 151: 214, # '\x97' - 152: 215, # '\x98' - 153: 216, # '\x99' - 154: 217, # '\x9a' - 155: 218, # '\x9b' - 156: 219, # '\x9c' - 157: 220, # '\x9d' - 158: 221, # '\x9e' - 159: 222, # '\x9f' - 160: 223, # '\xa0' - 161: 224, # 'Ё' - 162: 225, # 'Ђ' - 163: 226, # 'Ѓ' - 164: 227, # 'Є' - 165: 228, # 'Ѕ' - 166: 229, # 'І' - 167: 230, # 'Ї' - 168: 231, # 'Ј' - 169: 232, # 'Љ' - 170: 233, # 'Њ' - 171: 234, # 'Ћ' - 172: 235, # 'Ќ' - 173: 236, # '\xad' - 174: 237, # 'Ў' - 175: 238, # 'Џ' - 176: 37, # 'А' - 177: 44, # 'Б' - 178: 33, # 'В' - 179: 46, # 'Г' - 180: 41, # 'Д' - 181: 48, # 'Е' - 182: 56, # 'Ж' - 183: 51, # 'З' - 184: 42, # 'И' - 185: 60, # 'Й' - 186: 36, # 'К' - 187: 49, # 'Л' - 188: 38, # 'М' - 189: 31, # 'Н' - 190: 34, # 'О' - 191: 35, # 'П' - 192: 45, # 'Р' - 193: 32, # 'С' - 194: 40, # 'Т' - 195: 52, # 'У' - 196: 53, # 'Ф' - 197: 55, # 'Х' - 198: 58, # 'Ц' - 199: 50, # 'Ч' - 200: 57, # 'Ш' - 201: 63, # 'Щ' - 202: 70, # 'Ъ' - 203: 62, # 'Ы' - 204: 61, # 'Ь' - 205: 47, # 'Э' - 206: 59, # 'Ю' - 207: 43, # 'Я' - 208: 3, # 'а' - 209: 21, # 'б' - 210: 10, # 'в' - 211: 19, # 'г' - 212: 13, # 'д' - 213: 2, # 'е' - 214: 24, # 'ж' - 215: 20, # 'з' - 216: 4, # 'и' - 217: 23, # 'й' - 218: 11, # 'к' - 219: 8, # 'л' - 220: 12, # 'м' - 221: 5, # 'н' - 222: 1, # 'о' - 223: 15, # 'п' - 224: 9, # 'р' - 225: 7, # 'с' - 226: 6, # 'т' - 227: 14, # 'у' - 228: 39, # 'ф' - 229: 26, # 'х' - 230: 28, # 'ц' - 231: 22, # 'ч' - 232: 25, # 'ш' - 233: 29, # 'щ' - 234: 54, # 'ъ' - 235: 18, # 'ы' - 236: 17, # 'ь' - 237: 30, # 'э' - 238: 27, # 'ю' - 239: 16, # 'я' - 240: 239, # '№' - 241: 68, # 'ё' - 242: 240, # 'ђ' - 243: 241, # 'ѓ' - 244: 242, # 'є' - 245: 243, # 'ѕ' - 246: 244, # 'і' - 247: 245, # 'ї' - 248: 246, # 'ј' - 249: 247, # 'љ' - 250: 248, # 'њ' - 251: 249, # 'ћ' - 252: 250, # 'ќ' - 253: 251, # '§' - 254: 252, # 'ў' - 255: 255, # 'џ' -} - -ISO_8859_5_RUSSIAN_MODEL = SingleByteCharSetModel( - charset_name="ISO-8859-5", - language="Russian", - char_to_order_map=ISO_8859_5_RUSSIAN_CHAR_TO_ORDER, - language_model=RUSSIAN_LANG_MODEL, - typical_positive_ratio=0.976601, - keep_ascii_letters=False, - alphabet="ЁАБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯабвгдежзийклмнопрстуфхцчшщъыьэюяё", -) diff --git a/venv/lib/python3.10/site-packages/chardet/langthaimodel.py b/venv/lib/python3.10/site-packages/chardet/langthaimodel.py deleted file mode 100644 index 883fdb1eafea7c7122846bdd478da800375a55bd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/langthaimodel.py +++ /dev/null @@ -1,4380 +0,0 @@ -from chardet.sbcharsetprober import SingleByteCharSetModel - -# 3: Positive -# 2: Likely -# 1: Unlikely -# 0: Negative - -THAI_LANG_MODEL = { - 5: { # 'ก' - 5: 2, # 'ก' - 30: 2, # 'ข' - 24: 2, # 'ค' - 8: 2, # 'ง' - 26: 2, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 1, # 'ซ' - 47: 0, # 'ญ' - 58: 3, # 'ฎ' - 57: 2, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 2, # 'ณ' - 20: 2, # 'ด' - 19: 3, # 'ต' - 44: 0, # 'ถ' - 14: 2, # 'ท' - 48: 0, # 'ธ' - 3: 2, # 'น' - 17: 1, # 'บ' - 25: 2, # 'ป' - 39: 1, # 'ผ' - 62: 1, # 'ฝ' - 31: 1, # 'พ' - 54: 0, # 'ฟ' - 45: 1, # 'ภ' - 9: 2, # 'ม' - 16: 1, # 'ย' - 2: 3, # 'ร' - 61: 2, # 'ฤ' - 15: 3, # 'ล' - 12: 3, # 'ว' - 42: 2, # 'ศ' - 46: 3, # 'ษ' - 18: 2, # 'ส' - 21: 2, # 'ห' - 4: 3, # 'อ' - 63: 1, # 'ฯ' - 22: 2, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 3, # 'ำ' - 23: 3, # 'ิ' - 13: 3, # 'ี' - 40: 0, # 'ึ' - 27: 2, # 'ื' - 32: 2, # 'ุ' - 35: 1, # 'ู' - 11: 2, # 'เ' - 28: 2, # 'แ' - 41: 1, # 'โ' - 29: 1, # 'ใ' - 33: 2, # 'ไ' - 50: 1, # 'ๆ' - 37: 3, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 2, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 30: { # 'ข' - 5: 1, # 'ก' - 30: 0, # 'ข' - 24: 1, # 'ค' - 8: 1, # 'ง' - 26: 1, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 2, # 'ณ' - 20: 0, # 'ด' - 19: 2, # 'ต' - 44: 0, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 2, # 'น' - 17: 1, # 'บ' - 25: 1, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 2, # 'ย' - 2: 1, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 2, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 1, # 'ส' - 21: 1, # 'ห' - 4: 3, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 2, # 'ี' - 40: 3, # 'ึ' - 27: 1, # 'ื' - 32: 1, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 1, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 1, # '็' - 6: 2, # '่' - 7: 3, # '้' - 38: 1, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 24: { # 'ค' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 2, # 'ค' - 8: 2, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 2, # 'ณ' - 20: 2, # 'ด' - 19: 2, # 'ต' - 44: 0, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 0, # 'บ' - 25: 1, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 2, # 'ม' - 16: 2, # 'ย' - 2: 3, # 'ร' - 61: 0, # 'ฤ' - 15: 3, # 'ล' - 12: 3, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 1, # 'ส' - 21: 0, # 'ห' - 4: 2, # 'อ' - 63: 0, # 'ฯ' - 22: 2, # 'ะ' - 10: 3, # 'ั' - 1: 2, # 'า' - 36: 3, # 'ำ' - 23: 3, # 'ิ' - 13: 2, # 'ี' - 40: 0, # 'ึ' - 27: 3, # 'ื' - 32: 3, # 'ุ' - 35: 2, # 'ู' - 11: 1, # 'เ' - 28: 0, # 'แ' - 41: 3, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 1, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 3, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 8: { # 'ง' - 5: 3, # 'ก' - 30: 2, # 'ข' - 24: 3, # 'ค' - 8: 2, # 'ง' - 26: 2, # 'จ' - 52: 1, # 'ฉ' - 34: 2, # 'ช' - 51: 1, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 2, # 'ด' - 19: 2, # 'ต' - 44: 1, # 'ถ' - 14: 3, # 'ท' - 48: 1, # 'ธ' - 3: 3, # 'น' - 17: 2, # 'บ' - 25: 2, # 'ป' - 39: 2, # 'ผ' - 62: 1, # 'ฝ' - 31: 2, # 'พ' - 54: 0, # 'ฟ' - 45: 1, # 'ภ' - 9: 2, # 'ม' - 16: 1, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 2, # 'ว' - 42: 2, # 'ศ' - 46: 1, # 'ษ' - 18: 3, # 'ส' - 21: 3, # 'ห' - 4: 2, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 1, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 2, # 'ิ' - 13: 1, # 'ี' - 40: 0, # 'ึ' - 27: 1, # 'ื' - 32: 1, # 'ุ' - 35: 0, # 'ู' - 11: 3, # 'เ' - 28: 2, # 'แ' - 41: 1, # 'โ' - 29: 2, # 'ใ' - 33: 2, # 'ไ' - 50: 3, # 'ๆ' - 37: 0, # '็' - 6: 2, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 26: { # 'จ' - 5: 2, # 'ก' - 30: 1, # 'ข' - 24: 0, # 'ค' - 8: 2, # 'ง' - 26: 3, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 2, # 'ด' - 19: 1, # 'ต' - 44: 1, # 'ถ' - 14: 2, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 1, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 1, # 'ม' - 16: 1, # 'ย' - 2: 3, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 1, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 2, # 'ส' - 21: 1, # 'ห' - 4: 2, # 'อ' - 63: 0, # 'ฯ' - 22: 3, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 3, # 'ำ' - 23: 2, # 'ิ' - 13: 1, # 'ี' - 40: 3, # 'ึ' - 27: 1, # 'ื' - 32: 3, # 'ุ' - 35: 2, # 'ู' - 11: 1, # 'เ' - 28: 1, # 'แ' - 41: 0, # 'โ' - 29: 1, # 'ใ' - 33: 1, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 2, # '่' - 7: 2, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 52: { # 'ฉ' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 3, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 3, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 1, # 'ม' - 16: 1, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 1, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 1, # 'ะ' - 10: 1, # 'ั' - 1: 1, # 'า' - 36: 0, # 'ำ' - 23: 1, # 'ิ' - 13: 1, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 1, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 34: { # 'ช' - 5: 1, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 1, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 1, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 2, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 2, # 'ม' - 16: 1, # 'ย' - 2: 1, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 1, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 2, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 2, # 'ั' - 1: 3, # 'า' - 36: 1, # 'ำ' - 23: 3, # 'ิ' - 13: 2, # 'ี' - 40: 0, # 'ึ' - 27: 3, # 'ื' - 32: 3, # 'ุ' - 35: 1, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 1, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 51: { # 'ซ' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 1, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 1, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 1, # 'ส' - 21: 0, # 'ห' - 4: 2, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 1, # 'ั' - 1: 1, # 'า' - 36: 0, # 'ำ' - 23: 1, # 'ิ' - 13: 2, # 'ี' - 40: 3, # 'ึ' - 27: 2, # 'ื' - 32: 1, # 'ุ' - 35: 1, # 'ู' - 11: 1, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 1, # '็' - 6: 1, # '่' - 7: 2, # '้' - 38: 1, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 47: { # 'ญ' - 5: 1, # 'ก' - 30: 1, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 0, # 'ซ' - 47: 3, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 1, # 'บ' - 25: 1, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 1, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 1, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 1, # 'ส' - 21: 2, # 'ห' - 4: 1, # 'อ' - 63: 0, # 'ฯ' - 22: 1, # 'ะ' - 10: 2, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 1, # 'ิ' - 13: 1, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 1, # 'เ' - 28: 1, # 'แ' - 41: 0, # 'โ' - 29: 1, # 'ใ' - 33: 0, # 'ไ' - 50: 1, # 'ๆ' - 37: 0, # '็' - 6: 2, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 58: { # 'ฎ' - 5: 2, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 1, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 1, # 'ิ' - 13: 2, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 57: { # 'ฏ' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 3, # 'ิ' - 13: 1, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 49: { # 'ฐ' - 5: 1, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 2, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 2, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 1, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 1, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 1, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 53: { # 'ฑ' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 2, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 3, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 55: { # 'ฒ' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 1, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 43: { # 'ณ' - 5: 1, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 3, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 3, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 1, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 1, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 1, # 'ส' - 21: 1, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 3, # 'ะ' - 10: 0, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 1, # 'ิ' - 13: 2, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 1, # 'เ' - 28: 1, # 'แ' - 41: 0, # 'โ' - 29: 1, # 'ใ' - 33: 1, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 3, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 20: { # 'ด' - 5: 2, # 'ก' - 30: 2, # 'ข' - 24: 2, # 'ค' - 8: 3, # 'ง' - 26: 2, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 1, # 'ด' - 19: 2, # 'ต' - 44: 1, # 'ถ' - 14: 2, # 'ท' - 48: 0, # 'ธ' - 3: 1, # 'น' - 17: 1, # 'บ' - 25: 1, # 'ป' - 39: 1, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 0, # 'ฟ' - 45: 1, # 'ภ' - 9: 2, # 'ม' - 16: 3, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 2, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 2, # 'ส' - 21: 2, # 'ห' - 4: 1, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 3, # 'ั' - 1: 2, # 'า' - 36: 2, # 'ำ' - 23: 3, # 'ิ' - 13: 3, # 'ี' - 40: 1, # 'ึ' - 27: 2, # 'ื' - 32: 3, # 'ุ' - 35: 2, # 'ู' - 11: 2, # 'เ' - 28: 2, # 'แ' - 41: 1, # 'โ' - 29: 2, # 'ใ' - 33: 2, # 'ไ' - 50: 2, # 'ๆ' - 37: 2, # '็' - 6: 1, # '่' - 7: 3, # '้' - 38: 1, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 19: { # 'ต' - 5: 2, # 'ก' - 30: 1, # 'ข' - 24: 1, # 'ค' - 8: 0, # 'ง' - 26: 1, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 1, # 'ด' - 19: 1, # 'ต' - 44: 2, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 2, # 'น' - 17: 1, # 'บ' - 25: 1, # 'ป' - 39: 1, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 0, # 'ฟ' - 45: 2, # 'ภ' - 9: 1, # 'ม' - 16: 1, # 'ย' - 2: 3, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 1, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 3, # 'ส' - 21: 0, # 'ห' - 4: 3, # 'อ' - 63: 1, # 'ฯ' - 22: 2, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 2, # 'ำ' - 23: 3, # 'ิ' - 13: 2, # 'ี' - 40: 1, # 'ึ' - 27: 1, # 'ื' - 32: 3, # 'ุ' - 35: 2, # 'ู' - 11: 1, # 'เ' - 28: 1, # 'แ' - 41: 1, # 'โ' - 29: 1, # 'ใ' - 33: 1, # 'ไ' - 50: 0, # 'ๆ' - 37: 2, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 2, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 44: { # 'ถ' - 5: 1, # 'ก' - 30: 0, # 'ข' - 24: 1, # 'ค' - 8: 0, # 'ง' - 26: 1, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 1, # 'ต' - 44: 0, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 1, # 'น' - 17: 2, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 1, # 'ร' - 61: 0, # 'ฤ' - 15: 1, # 'ล' - 12: 1, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 1, # 'ส' - 21: 0, # 'ห' - 4: 1, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 2, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 2, # 'ิ' - 13: 1, # 'ี' - 40: 3, # 'ึ' - 27: 2, # 'ื' - 32: 2, # 'ุ' - 35: 3, # 'ู' - 11: 1, # 'เ' - 28: 1, # 'แ' - 41: 0, # 'โ' - 29: 1, # 'ใ' - 33: 1, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 2, # '่' - 7: 3, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 14: { # 'ท' - 5: 1, # 'ก' - 30: 1, # 'ข' - 24: 3, # 'ค' - 8: 1, # 'ง' - 26: 1, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 2, # 'ด' - 19: 1, # 'ต' - 44: 0, # 'ถ' - 14: 1, # 'ท' - 48: 3, # 'ธ' - 3: 3, # 'น' - 17: 2, # 'บ' - 25: 2, # 'ป' - 39: 1, # 'ผ' - 62: 0, # 'ฝ' - 31: 2, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 1, # 'ม' - 16: 3, # 'ย' - 2: 3, # 'ร' - 61: 1, # 'ฤ' - 15: 1, # 'ล' - 12: 2, # 'ว' - 42: 3, # 'ศ' - 46: 1, # 'ษ' - 18: 1, # 'ส' - 21: 0, # 'ห' - 4: 2, # 'อ' - 63: 0, # 'ฯ' - 22: 2, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 3, # 'ำ' - 23: 2, # 'ิ' - 13: 3, # 'ี' - 40: 2, # 'ึ' - 27: 1, # 'ื' - 32: 3, # 'ุ' - 35: 1, # 'ู' - 11: 0, # 'เ' - 28: 1, # 'แ' - 41: 0, # 'โ' - 29: 1, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 1, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 2, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 48: { # 'ธ' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 1, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 1, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 2, # 'า' - 36: 0, # 'ำ' - 23: 3, # 'ิ' - 13: 3, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 2, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 3, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 3: { # 'น' - 5: 3, # 'ก' - 30: 2, # 'ข' - 24: 3, # 'ค' - 8: 1, # 'ง' - 26: 2, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 1, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 1, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 3, # 'ด' - 19: 3, # 'ต' - 44: 2, # 'ถ' - 14: 3, # 'ท' - 48: 3, # 'ธ' - 3: 2, # 'น' - 17: 2, # 'บ' - 25: 2, # 'ป' - 39: 2, # 'ผ' - 62: 0, # 'ฝ' - 31: 2, # 'พ' - 54: 1, # 'ฟ' - 45: 1, # 'ภ' - 9: 2, # 'ม' - 16: 2, # 'ย' - 2: 2, # 'ร' - 61: 1, # 'ฤ' - 15: 2, # 'ล' - 12: 3, # 'ว' - 42: 1, # 'ศ' - 46: 0, # 'ษ' - 18: 2, # 'ส' - 21: 2, # 'ห' - 4: 3, # 'อ' - 63: 1, # 'ฯ' - 22: 2, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 3, # 'ำ' - 23: 3, # 'ิ' - 13: 3, # 'ี' - 40: 3, # 'ึ' - 27: 3, # 'ื' - 32: 3, # 'ุ' - 35: 2, # 'ู' - 11: 3, # 'เ' - 28: 2, # 'แ' - 41: 3, # 'โ' - 29: 3, # 'ใ' - 33: 3, # 'ไ' - 50: 2, # 'ๆ' - 37: 1, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 2, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 17: { # 'บ' - 5: 3, # 'ก' - 30: 2, # 'ข' - 24: 2, # 'ค' - 8: 1, # 'ง' - 26: 1, # 'จ' - 52: 1, # 'ฉ' - 34: 1, # 'ช' - 51: 1, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 1, # 'ด' - 19: 2, # 'ต' - 44: 1, # 'ถ' - 14: 3, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 3, # 'บ' - 25: 2, # 'ป' - 39: 2, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 1, # 'ฟ' - 45: 1, # 'ภ' - 9: 1, # 'ม' - 16: 0, # 'ย' - 2: 3, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 3, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 2, # 'ส' - 21: 2, # 'ห' - 4: 2, # 'อ' - 63: 1, # 'ฯ' - 22: 0, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 2, # 'ำ' - 23: 2, # 'ิ' - 13: 2, # 'ี' - 40: 0, # 'ึ' - 27: 2, # 'ื' - 32: 3, # 'ุ' - 35: 2, # 'ู' - 11: 2, # 'เ' - 28: 2, # 'แ' - 41: 1, # 'โ' - 29: 2, # 'ใ' - 33: 2, # 'ไ' - 50: 0, # 'ๆ' - 37: 1, # '็' - 6: 2, # '่' - 7: 2, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 25: { # 'ป' - 5: 2, # 'ก' - 30: 0, # 'ข' - 24: 1, # 'ค' - 8: 0, # 'ง' - 26: 1, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 1, # 'ซ' - 47: 0, # 'ญ' - 58: 1, # 'ฎ' - 57: 3, # 'ฏ' - 49: 1, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 1, # 'ด' - 19: 1, # 'ต' - 44: 1, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 2, # 'น' - 17: 0, # 'บ' - 25: 1, # 'ป' - 39: 1, # 'ผ' - 62: 1, # 'ฝ' - 31: 1, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 1, # 'ม' - 16: 0, # 'ย' - 2: 3, # 'ร' - 61: 0, # 'ฤ' - 15: 3, # 'ล' - 12: 1, # 'ว' - 42: 0, # 'ศ' - 46: 1, # 'ษ' - 18: 2, # 'ส' - 21: 1, # 'ห' - 4: 2, # 'อ' - 63: 0, # 'ฯ' - 22: 1, # 'ะ' - 10: 3, # 'ั' - 1: 1, # 'า' - 36: 0, # 'ำ' - 23: 2, # 'ิ' - 13: 3, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 1, # 'ุ' - 35: 0, # 'ู' - 11: 1, # 'เ' - 28: 2, # 'แ' - 41: 0, # 'โ' - 29: 1, # 'ใ' - 33: 2, # 'ไ' - 50: 0, # 'ๆ' - 37: 3, # '็' - 6: 1, # '่' - 7: 2, # '้' - 38: 1, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 39: { # 'ผ' - 5: 1, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 1, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 2, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 1, # 'ม' - 16: 2, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 3, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 1, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 1, # 'ะ' - 10: 1, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 2, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 1, # 'ื' - 32: 0, # 'ุ' - 35: 3, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 3, # '่' - 7: 1, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 62: { # 'ฝ' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 1, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 1, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 1, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 1, # 'ี' - 40: 2, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 2, # '่' - 7: 1, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 31: { # 'พ' - 5: 1, # 'ก' - 30: 1, # 'ข' - 24: 1, # 'ค' - 8: 1, # 'ง' - 26: 1, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 1, # 'ณ' - 20: 1, # 'ด' - 19: 1, # 'ต' - 44: 0, # 'ถ' - 14: 2, # 'ท' - 48: 1, # 'ธ' - 3: 3, # 'น' - 17: 2, # 'บ' - 25: 0, # 'ป' - 39: 1, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 1, # 'ม' - 16: 2, # 'ย' - 2: 3, # 'ร' - 61: 2, # 'ฤ' - 15: 2, # 'ล' - 12: 2, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 1, # 'ส' - 21: 1, # 'ห' - 4: 2, # 'อ' - 63: 1, # 'ฯ' - 22: 0, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 3, # 'ิ' - 13: 2, # 'ี' - 40: 1, # 'ึ' - 27: 3, # 'ื' - 32: 1, # 'ุ' - 35: 2, # 'ู' - 11: 1, # 'เ' - 28: 1, # 'แ' - 41: 0, # 'โ' - 29: 1, # 'ใ' - 33: 1, # 'ไ' - 50: 0, # 'ๆ' - 37: 1, # '็' - 6: 0, # '่' - 7: 1, # '้' - 38: 3, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 54: { # 'ฟ' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 1, # 'ต' - 44: 0, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 2, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 1, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 1, # 'ส' - 21: 0, # 'ห' - 4: 1, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 2, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 1, # 'ิ' - 13: 1, # 'ี' - 40: 0, # 'ึ' - 27: 1, # 'ื' - 32: 1, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 1, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 2, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 45: { # 'ภ' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 1, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 3, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 1, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 1, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 2, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 1, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 9: { # 'ม' - 5: 2, # 'ก' - 30: 2, # 'ข' - 24: 2, # 'ค' - 8: 2, # 'ง' - 26: 2, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 1, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 1, # 'ณ' - 20: 2, # 'ด' - 19: 2, # 'ต' - 44: 1, # 'ถ' - 14: 2, # 'ท' - 48: 1, # 'ธ' - 3: 3, # 'น' - 17: 2, # 'บ' - 25: 2, # 'ป' - 39: 1, # 'ผ' - 62: 0, # 'ฝ' - 31: 3, # 'พ' - 54: 0, # 'ฟ' - 45: 1, # 'ภ' - 9: 2, # 'ม' - 16: 1, # 'ย' - 2: 2, # 'ร' - 61: 2, # 'ฤ' - 15: 2, # 'ล' - 12: 2, # 'ว' - 42: 1, # 'ศ' - 46: 1, # 'ษ' - 18: 3, # 'ส' - 21: 3, # 'ห' - 4: 3, # 'อ' - 63: 0, # 'ฯ' - 22: 1, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 3, # 'ิ' - 13: 3, # 'ี' - 40: 0, # 'ึ' - 27: 3, # 'ื' - 32: 3, # 'ุ' - 35: 3, # 'ู' - 11: 2, # 'เ' - 28: 2, # 'แ' - 41: 2, # 'โ' - 29: 2, # 'ใ' - 33: 2, # 'ไ' - 50: 1, # 'ๆ' - 37: 1, # '็' - 6: 3, # '่' - 7: 2, # '้' - 38: 1, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 16: { # 'ย' - 5: 3, # 'ก' - 30: 1, # 'ข' - 24: 2, # 'ค' - 8: 3, # 'ง' - 26: 2, # 'จ' - 52: 0, # 'ฉ' - 34: 2, # 'ช' - 51: 0, # 'ซ' - 47: 2, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 2, # 'ด' - 19: 2, # 'ต' - 44: 1, # 'ถ' - 14: 2, # 'ท' - 48: 1, # 'ธ' - 3: 3, # 'น' - 17: 3, # 'บ' - 25: 1, # 'ป' - 39: 1, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 0, # 'ฟ' - 45: 1, # 'ภ' - 9: 2, # 'ม' - 16: 0, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 1, # 'ล' - 12: 3, # 'ว' - 42: 1, # 'ศ' - 46: 0, # 'ษ' - 18: 2, # 'ส' - 21: 1, # 'ห' - 4: 2, # 'อ' - 63: 0, # 'ฯ' - 22: 2, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 2, # 'ิ' - 13: 3, # 'ี' - 40: 1, # 'ึ' - 27: 2, # 'ื' - 32: 2, # 'ุ' - 35: 3, # 'ู' - 11: 2, # 'เ' - 28: 1, # 'แ' - 41: 1, # 'โ' - 29: 2, # 'ใ' - 33: 2, # 'ไ' - 50: 2, # 'ๆ' - 37: 1, # '็' - 6: 3, # '่' - 7: 2, # '้' - 38: 3, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 2: { # 'ร' - 5: 3, # 'ก' - 30: 2, # 'ข' - 24: 2, # 'ค' - 8: 3, # 'ง' - 26: 2, # 'จ' - 52: 0, # 'ฉ' - 34: 2, # 'ช' - 51: 1, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 3, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 3, # 'ณ' - 20: 2, # 'ด' - 19: 2, # 'ต' - 44: 3, # 'ถ' - 14: 3, # 'ท' - 48: 1, # 'ธ' - 3: 2, # 'น' - 17: 2, # 'บ' - 25: 3, # 'ป' - 39: 2, # 'ผ' - 62: 1, # 'ฝ' - 31: 2, # 'พ' - 54: 1, # 'ฟ' - 45: 1, # 'ภ' - 9: 3, # 'ม' - 16: 2, # 'ย' - 2: 3, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 3, # 'ว' - 42: 2, # 'ศ' - 46: 2, # 'ษ' - 18: 2, # 'ส' - 21: 2, # 'ห' - 4: 3, # 'อ' - 63: 1, # 'ฯ' - 22: 3, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 3, # 'ิ' - 13: 3, # 'ี' - 40: 2, # 'ึ' - 27: 3, # 'ื' - 32: 3, # 'ุ' - 35: 3, # 'ู' - 11: 3, # 'เ' - 28: 3, # 'แ' - 41: 1, # 'โ' - 29: 2, # 'ใ' - 33: 1, # 'ไ' - 50: 0, # 'ๆ' - 37: 3, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 3, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 61: { # 'ฤ' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 2, # 'ต' - 44: 0, # 'ถ' - 14: 2, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 1, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 2, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 15: { # 'ล' - 5: 2, # 'ก' - 30: 3, # 'ข' - 24: 1, # 'ค' - 8: 3, # 'ง' - 26: 1, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 2, # 'ด' - 19: 2, # 'ต' - 44: 1, # 'ถ' - 14: 2, # 'ท' - 48: 0, # 'ธ' - 3: 1, # 'น' - 17: 2, # 'บ' - 25: 2, # 'ป' - 39: 1, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 1, # 'ภ' - 9: 1, # 'ม' - 16: 3, # 'ย' - 2: 1, # 'ร' - 61: 0, # 'ฤ' - 15: 1, # 'ล' - 12: 1, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 2, # 'ส' - 21: 1, # 'ห' - 4: 3, # 'อ' - 63: 2, # 'ฯ' - 22: 3, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 2, # 'ำ' - 23: 3, # 'ิ' - 13: 3, # 'ี' - 40: 2, # 'ึ' - 27: 3, # 'ื' - 32: 2, # 'ุ' - 35: 3, # 'ู' - 11: 2, # 'เ' - 28: 1, # 'แ' - 41: 1, # 'โ' - 29: 2, # 'ใ' - 33: 1, # 'ไ' - 50: 0, # 'ๆ' - 37: 2, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 2, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 12: { # 'ว' - 5: 3, # 'ก' - 30: 2, # 'ข' - 24: 1, # 'ค' - 8: 3, # 'ง' - 26: 2, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 1, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 1, # 'ณ' - 20: 2, # 'ด' - 19: 1, # 'ต' - 44: 1, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 2, # 'บ' - 25: 1, # 'ป' - 39: 1, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 1, # 'ฟ' - 45: 0, # 'ภ' - 9: 3, # 'ม' - 16: 3, # 'ย' - 2: 3, # 'ร' - 61: 0, # 'ฤ' - 15: 3, # 'ล' - 12: 1, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 2, # 'ส' - 21: 2, # 'ห' - 4: 2, # 'อ' - 63: 0, # 'ฯ' - 22: 2, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 3, # 'ิ' - 13: 2, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 2, # 'ุ' - 35: 0, # 'ู' - 11: 3, # 'เ' - 28: 2, # 'แ' - 41: 1, # 'โ' - 29: 1, # 'ใ' - 33: 2, # 'ไ' - 50: 1, # 'ๆ' - 37: 0, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 1, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 42: { # 'ศ' - 5: 1, # 'ก' - 30: 0, # 'ข' - 24: 1, # 'ค' - 8: 0, # 'ง' - 26: 1, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 1, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 1, # 'ต' - 44: 0, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 2, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 2, # 'ว' - 42: 1, # 'ศ' - 46: 2, # 'ษ' - 18: 1, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 2, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 2, # 'ิ' - 13: 0, # 'ี' - 40: 3, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 2, # 'ู' - 11: 0, # 'เ' - 28: 1, # 'แ' - 41: 0, # 'โ' - 29: 1, # 'ใ' - 33: 1, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 1, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 46: { # 'ษ' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 2, # 'ฎ' - 57: 1, # 'ฏ' - 49: 2, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 3, # 'ณ' - 20: 0, # 'ด' - 19: 1, # 'ต' - 44: 0, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 1, # 'ภ' - 9: 1, # 'ม' - 16: 2, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 1, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 2, # 'ะ' - 10: 2, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 1, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 1, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 2, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 18: { # 'ส' - 5: 2, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 2, # 'ง' - 26: 1, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 3, # 'ด' - 19: 3, # 'ต' - 44: 3, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 2, # 'บ' - 25: 1, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 2, # 'ภ' - 9: 3, # 'ม' - 16: 1, # 'ย' - 2: 3, # 'ร' - 61: 0, # 'ฤ' - 15: 1, # 'ล' - 12: 2, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 2, # 'ห' - 4: 3, # 'อ' - 63: 0, # 'ฯ' - 22: 2, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 3, # 'ำ' - 23: 3, # 'ิ' - 13: 3, # 'ี' - 40: 2, # 'ึ' - 27: 3, # 'ื' - 32: 3, # 'ุ' - 35: 3, # 'ู' - 11: 2, # 'เ' - 28: 0, # 'แ' - 41: 1, # 'โ' - 29: 0, # 'ใ' - 33: 1, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 3, # '่' - 7: 1, # '้' - 38: 2, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 21: { # 'ห' - 5: 3, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 1, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 2, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 1, # 'ด' - 19: 3, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 0, # 'บ' - 25: 1, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 3, # 'ม' - 16: 2, # 'ย' - 2: 3, # 'ร' - 61: 0, # 'ฤ' - 15: 3, # 'ล' - 12: 2, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 3, # 'อ' - 63: 0, # 'ฯ' - 22: 1, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 0, # 'ำ' - 23: 1, # 'ิ' - 13: 1, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 1, # 'ุ' - 35: 1, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 3, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 2, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 4: { # 'อ' - 5: 3, # 'ก' - 30: 1, # 'ข' - 24: 2, # 'ค' - 8: 3, # 'ง' - 26: 1, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 3, # 'ด' - 19: 2, # 'ต' - 44: 1, # 'ถ' - 14: 2, # 'ท' - 48: 1, # 'ธ' - 3: 3, # 'น' - 17: 3, # 'บ' - 25: 1, # 'ป' - 39: 1, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 1, # 'ฟ' - 45: 1, # 'ภ' - 9: 3, # 'ม' - 16: 3, # 'ย' - 2: 3, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 2, # 'ว' - 42: 1, # 'ศ' - 46: 0, # 'ษ' - 18: 2, # 'ส' - 21: 2, # 'ห' - 4: 3, # 'อ' - 63: 0, # 'ฯ' - 22: 2, # 'ะ' - 10: 3, # 'ั' - 1: 3, # 'า' - 36: 2, # 'ำ' - 23: 2, # 'ิ' - 13: 3, # 'ี' - 40: 0, # 'ึ' - 27: 3, # 'ื' - 32: 3, # 'ุ' - 35: 0, # 'ู' - 11: 3, # 'เ' - 28: 1, # 'แ' - 41: 1, # 'โ' - 29: 2, # 'ใ' - 33: 2, # 'ไ' - 50: 1, # 'ๆ' - 37: 1, # '็' - 6: 2, # '่' - 7: 2, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 63: { # 'ฯ' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 22: { # 'ะ' - 5: 3, # 'ก' - 30: 1, # 'ข' - 24: 2, # 'ค' - 8: 1, # 'ง' - 26: 2, # 'จ' - 52: 0, # 'ฉ' - 34: 3, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 3, # 'ด' - 19: 3, # 'ต' - 44: 1, # 'ถ' - 14: 3, # 'ท' - 48: 1, # 'ธ' - 3: 2, # 'น' - 17: 3, # 'บ' - 25: 2, # 'ป' - 39: 1, # 'ผ' - 62: 0, # 'ฝ' - 31: 2, # 'พ' - 54: 0, # 'ฟ' - 45: 1, # 'ภ' - 9: 3, # 'ม' - 16: 2, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 2, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 3, # 'ส' - 21: 3, # 'ห' - 4: 2, # 'อ' - 63: 1, # 'ฯ' - 22: 1, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 3, # 'เ' - 28: 2, # 'แ' - 41: 1, # 'โ' - 29: 2, # 'ใ' - 33: 2, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 10: { # 'ั' - 5: 3, # 'ก' - 30: 0, # 'ข' - 24: 1, # 'ค' - 8: 3, # 'ง' - 26: 3, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 0, # 'ซ' - 47: 3, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 2, # 'ฐ' - 53: 0, # 'ฑ' - 55: 3, # 'ฒ' - 43: 3, # 'ณ' - 20: 3, # 'ด' - 19: 3, # 'ต' - 44: 0, # 'ถ' - 14: 2, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 3, # 'บ' - 25: 1, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 2, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 3, # 'ม' - 16: 3, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 3, # 'ว' - 42: 2, # 'ศ' - 46: 0, # 'ษ' - 18: 3, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 1: { # 'า' - 5: 3, # 'ก' - 30: 2, # 'ข' - 24: 3, # 'ค' - 8: 3, # 'ง' - 26: 3, # 'จ' - 52: 0, # 'ฉ' - 34: 3, # 'ช' - 51: 1, # 'ซ' - 47: 2, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 3, # 'ณ' - 20: 3, # 'ด' - 19: 3, # 'ต' - 44: 1, # 'ถ' - 14: 3, # 'ท' - 48: 2, # 'ธ' - 3: 3, # 'น' - 17: 3, # 'บ' - 25: 2, # 'ป' - 39: 1, # 'ผ' - 62: 1, # 'ฝ' - 31: 3, # 'พ' - 54: 1, # 'ฟ' - 45: 1, # 'ภ' - 9: 3, # 'ม' - 16: 3, # 'ย' - 2: 3, # 'ร' - 61: 0, # 'ฤ' - 15: 3, # 'ล' - 12: 3, # 'ว' - 42: 2, # 'ศ' - 46: 3, # 'ษ' - 18: 3, # 'ส' - 21: 3, # 'ห' - 4: 2, # 'อ' - 63: 1, # 'ฯ' - 22: 3, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 3, # 'เ' - 28: 2, # 'แ' - 41: 1, # 'โ' - 29: 2, # 'ใ' - 33: 2, # 'ไ' - 50: 1, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 36: { # 'ำ' - 5: 2, # 'ก' - 30: 1, # 'ข' - 24: 3, # 'ค' - 8: 2, # 'ง' - 26: 1, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 1, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 1, # 'ด' - 19: 1, # 'ต' - 44: 1, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 1, # 'บ' - 25: 1, # 'ป' - 39: 1, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 0, # 'ฟ' - 45: 1, # 'ภ' - 9: 1, # 'ม' - 16: 0, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 1, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 1, # 'ส' - 21: 3, # 'ห' - 4: 1, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 3, # 'เ' - 28: 2, # 'แ' - 41: 1, # 'โ' - 29: 2, # 'ใ' - 33: 2, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 23: { # 'ิ' - 5: 3, # 'ก' - 30: 1, # 'ข' - 24: 2, # 'ค' - 8: 3, # 'ง' - 26: 3, # 'จ' - 52: 0, # 'ฉ' - 34: 3, # 'ช' - 51: 0, # 'ซ' - 47: 2, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 3, # 'ด' - 19: 3, # 'ต' - 44: 1, # 'ถ' - 14: 3, # 'ท' - 48: 3, # 'ธ' - 3: 3, # 'น' - 17: 3, # 'บ' - 25: 2, # 'ป' - 39: 2, # 'ผ' - 62: 0, # 'ฝ' - 31: 3, # 'พ' - 54: 1, # 'ฟ' - 45: 2, # 'ภ' - 9: 3, # 'ม' - 16: 2, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 3, # 'ว' - 42: 3, # 'ศ' - 46: 2, # 'ษ' - 18: 2, # 'ส' - 21: 3, # 'ห' - 4: 1, # 'อ' - 63: 1, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 3, # 'เ' - 28: 1, # 'แ' - 41: 1, # 'โ' - 29: 1, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 3, # '่' - 7: 2, # '้' - 38: 2, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 13: { # 'ี' - 5: 3, # 'ก' - 30: 2, # 'ข' - 24: 2, # 'ค' - 8: 0, # 'ง' - 26: 1, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 2, # 'ด' - 19: 1, # 'ต' - 44: 0, # 'ถ' - 14: 2, # 'ท' - 48: 0, # 'ธ' - 3: 1, # 'น' - 17: 2, # 'บ' - 25: 2, # 'ป' - 39: 1, # 'ผ' - 62: 0, # 'ฝ' - 31: 2, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 2, # 'ม' - 16: 3, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 1, # 'ล' - 12: 2, # 'ว' - 42: 1, # 'ศ' - 46: 0, # 'ษ' - 18: 2, # 'ส' - 21: 1, # 'ห' - 4: 2, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 2, # 'เ' - 28: 2, # 'แ' - 41: 1, # 'โ' - 29: 1, # 'ใ' - 33: 1, # 'ไ' - 50: 1, # 'ๆ' - 37: 0, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 40: { # 'ึ' - 5: 3, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 3, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 1, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 1, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 27: { # 'ื' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 1, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 2, # 'น' - 17: 3, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 2, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 3, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 32: { # 'ุ' - 5: 3, # 'ก' - 30: 2, # 'ข' - 24: 3, # 'ค' - 8: 3, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 2, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 1, # 'ฒ' - 43: 3, # 'ณ' - 20: 3, # 'ด' - 19: 3, # 'ต' - 44: 1, # 'ถ' - 14: 2, # 'ท' - 48: 1, # 'ธ' - 3: 2, # 'น' - 17: 2, # 'บ' - 25: 2, # 'ป' - 39: 2, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 0, # 'ฟ' - 45: 1, # 'ภ' - 9: 3, # 'ม' - 16: 1, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 1, # 'ว' - 42: 1, # 'ศ' - 46: 2, # 'ษ' - 18: 1, # 'ส' - 21: 1, # 'ห' - 4: 1, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 1, # 'เ' - 28: 0, # 'แ' - 41: 1, # 'โ' - 29: 0, # 'ใ' - 33: 1, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 3, # '่' - 7: 2, # '้' - 38: 1, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 35: { # 'ู' - 5: 3, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 2, # 'ง' - 26: 1, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 2, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 1, # 'ณ' - 20: 2, # 'ด' - 19: 2, # 'ต' - 44: 0, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 2, # 'น' - 17: 0, # 'บ' - 25: 3, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 2, # 'ม' - 16: 0, # 'ย' - 2: 1, # 'ร' - 61: 0, # 'ฤ' - 15: 3, # 'ล' - 12: 1, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 1, # 'เ' - 28: 1, # 'แ' - 41: 1, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 3, # '่' - 7: 3, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 11: { # 'เ' - 5: 3, # 'ก' - 30: 3, # 'ข' - 24: 3, # 'ค' - 8: 2, # 'ง' - 26: 3, # 'จ' - 52: 3, # 'ฉ' - 34: 3, # 'ช' - 51: 2, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 1, # 'ณ' - 20: 3, # 'ด' - 19: 3, # 'ต' - 44: 1, # 'ถ' - 14: 3, # 'ท' - 48: 1, # 'ธ' - 3: 3, # 'น' - 17: 3, # 'บ' - 25: 3, # 'ป' - 39: 2, # 'ผ' - 62: 1, # 'ฝ' - 31: 3, # 'พ' - 54: 1, # 'ฟ' - 45: 3, # 'ภ' - 9: 3, # 'ม' - 16: 2, # 'ย' - 2: 3, # 'ร' - 61: 0, # 'ฤ' - 15: 3, # 'ล' - 12: 3, # 'ว' - 42: 2, # 'ศ' - 46: 0, # 'ษ' - 18: 3, # 'ส' - 21: 3, # 'ห' - 4: 3, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 28: { # 'แ' - 5: 3, # 'ก' - 30: 2, # 'ข' - 24: 2, # 'ค' - 8: 1, # 'ง' - 26: 2, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 2, # 'ด' - 19: 3, # 'ต' - 44: 2, # 'ถ' - 14: 3, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 3, # 'บ' - 25: 2, # 'ป' - 39: 3, # 'ผ' - 62: 0, # 'ฝ' - 31: 2, # 'พ' - 54: 2, # 'ฟ' - 45: 0, # 'ภ' - 9: 2, # 'ม' - 16: 2, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 3, # 'ล' - 12: 2, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 3, # 'ส' - 21: 3, # 'ห' - 4: 1, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 41: { # 'โ' - 5: 2, # 'ก' - 30: 1, # 'ข' - 24: 2, # 'ค' - 8: 0, # 'ง' - 26: 1, # 'จ' - 52: 1, # 'ฉ' - 34: 1, # 'ช' - 51: 1, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 3, # 'ด' - 19: 2, # 'ต' - 44: 0, # 'ถ' - 14: 2, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 1, # 'บ' - 25: 3, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 1, # 'ฟ' - 45: 1, # 'ภ' - 9: 1, # 'ม' - 16: 2, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 3, # 'ล' - 12: 0, # 'ว' - 42: 1, # 'ศ' - 46: 0, # 'ษ' - 18: 2, # 'ส' - 21: 0, # 'ห' - 4: 2, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 29: { # 'ใ' - 5: 2, # 'ก' - 30: 0, # 'ข' - 24: 1, # 'ค' - 8: 0, # 'ง' - 26: 3, # 'จ' - 52: 0, # 'ฉ' - 34: 3, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 3, # 'ด' - 19: 1, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 2, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 1, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 3, # 'ส' - 21: 3, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 33: { # 'ไ' - 5: 1, # 'ก' - 30: 2, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 1, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 3, # 'ด' - 19: 1, # 'ต' - 44: 0, # 'ถ' - 14: 3, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 1, # 'บ' - 25: 3, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 2, # 'ฟ' - 45: 0, # 'ภ' - 9: 3, # 'ม' - 16: 0, # 'ย' - 2: 3, # 'ร' - 61: 0, # 'ฤ' - 15: 1, # 'ล' - 12: 3, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 1, # 'ส' - 21: 2, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 50: { # 'ๆ' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 37: { # '็' - 5: 2, # 'ก' - 30: 1, # 'ข' - 24: 2, # 'ค' - 8: 2, # 'ง' - 26: 3, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 1, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 1, # 'ด' - 19: 2, # 'ต' - 44: 0, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 3, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 2, # 'ม' - 16: 1, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 2, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 1, # 'ส' - 21: 0, # 'ห' - 4: 1, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 1, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 1, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 6: { # '่' - 5: 2, # 'ก' - 30: 1, # 'ข' - 24: 2, # 'ค' - 8: 3, # 'ง' - 26: 2, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 1, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 1, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 1, # 'ด' - 19: 2, # 'ต' - 44: 1, # 'ถ' - 14: 2, # 'ท' - 48: 1, # 'ธ' - 3: 3, # 'น' - 17: 1, # 'บ' - 25: 2, # 'ป' - 39: 2, # 'ผ' - 62: 1, # 'ฝ' - 31: 1, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 3, # 'ม' - 16: 3, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 2, # 'ล' - 12: 3, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 2, # 'ส' - 21: 1, # 'ห' - 4: 3, # 'อ' - 63: 0, # 'ฯ' - 22: 1, # 'ะ' - 10: 0, # 'ั' - 1: 3, # 'า' - 36: 2, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 3, # 'เ' - 28: 2, # 'แ' - 41: 1, # 'โ' - 29: 2, # 'ใ' - 33: 2, # 'ไ' - 50: 1, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 7: { # '้' - 5: 2, # 'ก' - 30: 1, # 'ข' - 24: 2, # 'ค' - 8: 3, # 'ง' - 26: 2, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 1, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 1, # 'ด' - 19: 2, # 'ต' - 44: 1, # 'ถ' - 14: 2, # 'ท' - 48: 0, # 'ธ' - 3: 3, # 'น' - 17: 2, # 'บ' - 25: 2, # 'ป' - 39: 2, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 1, # 'ฟ' - 45: 0, # 'ภ' - 9: 3, # 'ม' - 16: 2, # 'ย' - 2: 2, # 'ร' - 61: 0, # 'ฤ' - 15: 1, # 'ล' - 12: 3, # 'ว' - 42: 1, # 'ศ' - 46: 0, # 'ษ' - 18: 2, # 'ส' - 21: 2, # 'ห' - 4: 3, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 3, # 'า' - 36: 2, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 2, # 'เ' - 28: 2, # 'แ' - 41: 1, # 'โ' - 29: 2, # 'ใ' - 33: 2, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 38: { # '์' - 5: 2, # 'ก' - 30: 1, # 'ข' - 24: 1, # 'ค' - 8: 0, # 'ง' - 26: 1, # 'จ' - 52: 0, # 'ฉ' - 34: 1, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 2, # 'ด' - 19: 1, # 'ต' - 44: 1, # 'ถ' - 14: 1, # 'ท' - 48: 0, # 'ธ' - 3: 1, # 'น' - 17: 1, # 'บ' - 25: 1, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 1, # 'พ' - 54: 1, # 'ฟ' - 45: 0, # 'ภ' - 9: 2, # 'ม' - 16: 0, # 'ย' - 2: 1, # 'ร' - 61: 1, # 'ฤ' - 15: 1, # 'ล' - 12: 1, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 1, # 'ส' - 21: 1, # 'ห' - 4: 2, # 'อ' - 63: 1, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 2, # 'เ' - 28: 2, # 'แ' - 41: 1, # 'โ' - 29: 1, # 'ใ' - 33: 1, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 0, # '๑' - 59: 0, # '๒' - 60: 0, # '๕' - }, - 56: { # '๑' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 2, # '๑' - 59: 1, # '๒' - 60: 1, # '๕' - }, - 59: { # '๒' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 1, # '๑' - 59: 1, # '๒' - 60: 3, # '๕' - }, - 60: { # '๕' - 5: 0, # 'ก' - 30: 0, # 'ข' - 24: 0, # 'ค' - 8: 0, # 'ง' - 26: 0, # 'จ' - 52: 0, # 'ฉ' - 34: 0, # 'ช' - 51: 0, # 'ซ' - 47: 0, # 'ญ' - 58: 0, # 'ฎ' - 57: 0, # 'ฏ' - 49: 0, # 'ฐ' - 53: 0, # 'ฑ' - 55: 0, # 'ฒ' - 43: 0, # 'ณ' - 20: 0, # 'ด' - 19: 0, # 'ต' - 44: 0, # 'ถ' - 14: 0, # 'ท' - 48: 0, # 'ธ' - 3: 0, # 'น' - 17: 0, # 'บ' - 25: 0, # 'ป' - 39: 0, # 'ผ' - 62: 0, # 'ฝ' - 31: 0, # 'พ' - 54: 0, # 'ฟ' - 45: 0, # 'ภ' - 9: 0, # 'ม' - 16: 0, # 'ย' - 2: 0, # 'ร' - 61: 0, # 'ฤ' - 15: 0, # 'ล' - 12: 0, # 'ว' - 42: 0, # 'ศ' - 46: 0, # 'ษ' - 18: 0, # 'ส' - 21: 0, # 'ห' - 4: 0, # 'อ' - 63: 0, # 'ฯ' - 22: 0, # 'ะ' - 10: 0, # 'ั' - 1: 0, # 'า' - 36: 0, # 'ำ' - 23: 0, # 'ิ' - 13: 0, # 'ี' - 40: 0, # 'ึ' - 27: 0, # 'ื' - 32: 0, # 'ุ' - 35: 0, # 'ู' - 11: 0, # 'เ' - 28: 0, # 'แ' - 41: 0, # 'โ' - 29: 0, # 'ใ' - 33: 0, # 'ไ' - 50: 0, # 'ๆ' - 37: 0, # '็' - 6: 0, # '่' - 7: 0, # '้' - 38: 0, # '์' - 56: 2, # '๑' - 59: 1, # '๒' - 60: 0, # '๕' - }, -} - -# 255: Undefined characters that did not exist in training text -# 254: Carriage/Return -# 253: symbol (punctuation) that does not belong to word -# 252: 0 - 9 -# 251: Control characters - -# Character Mapping Table(s): -TIS_620_THAI_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 254, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 254, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 253, # ' ' - 33: 253, # '!' - 34: 253, # '"' - 35: 253, # '#' - 36: 253, # '$' - 37: 253, # '%' - 38: 253, # '&' - 39: 253, # "'" - 40: 253, # '(' - 41: 253, # ')' - 42: 253, # '*' - 43: 253, # '+' - 44: 253, # ',' - 45: 253, # '-' - 46: 253, # '.' - 47: 253, # '/' - 48: 252, # '0' - 49: 252, # '1' - 50: 252, # '2' - 51: 252, # '3' - 52: 252, # '4' - 53: 252, # '5' - 54: 252, # '6' - 55: 252, # '7' - 56: 252, # '8' - 57: 252, # '9' - 58: 253, # ':' - 59: 253, # ';' - 60: 253, # '<' - 61: 253, # '=' - 62: 253, # '>' - 63: 253, # '?' - 64: 253, # '@' - 65: 182, # 'A' - 66: 106, # 'B' - 67: 107, # 'C' - 68: 100, # 'D' - 69: 183, # 'E' - 70: 184, # 'F' - 71: 185, # 'G' - 72: 101, # 'H' - 73: 94, # 'I' - 74: 186, # 'J' - 75: 187, # 'K' - 76: 108, # 'L' - 77: 109, # 'M' - 78: 110, # 'N' - 79: 111, # 'O' - 80: 188, # 'P' - 81: 189, # 'Q' - 82: 190, # 'R' - 83: 89, # 'S' - 84: 95, # 'T' - 85: 112, # 'U' - 86: 113, # 'V' - 87: 191, # 'W' - 88: 192, # 'X' - 89: 193, # 'Y' - 90: 194, # 'Z' - 91: 253, # '[' - 92: 253, # '\\' - 93: 253, # ']' - 94: 253, # '^' - 95: 253, # '_' - 96: 253, # '`' - 97: 64, # 'a' - 98: 72, # 'b' - 99: 73, # 'c' - 100: 114, # 'd' - 101: 74, # 'e' - 102: 115, # 'f' - 103: 116, # 'g' - 104: 102, # 'h' - 105: 81, # 'i' - 106: 201, # 'j' - 107: 117, # 'k' - 108: 90, # 'l' - 109: 103, # 'm' - 110: 78, # 'n' - 111: 82, # 'o' - 112: 96, # 'p' - 113: 202, # 'q' - 114: 91, # 'r' - 115: 79, # 's' - 116: 84, # 't' - 117: 104, # 'u' - 118: 105, # 'v' - 119: 97, # 'w' - 120: 98, # 'x' - 121: 92, # 'y' - 122: 203, # 'z' - 123: 253, # '{' - 124: 253, # '|' - 125: 253, # '}' - 126: 253, # '~' - 127: 253, # '\x7f' - 128: 209, # '\x80' - 129: 210, # '\x81' - 130: 211, # '\x82' - 131: 212, # '\x83' - 132: 213, # '\x84' - 133: 88, # '\x85' - 134: 214, # '\x86' - 135: 215, # '\x87' - 136: 216, # '\x88' - 137: 217, # '\x89' - 138: 218, # '\x8a' - 139: 219, # '\x8b' - 140: 220, # '\x8c' - 141: 118, # '\x8d' - 142: 221, # '\x8e' - 143: 222, # '\x8f' - 144: 223, # '\x90' - 145: 224, # '\x91' - 146: 99, # '\x92' - 147: 85, # '\x93' - 148: 83, # '\x94' - 149: 225, # '\x95' - 150: 226, # '\x96' - 151: 227, # '\x97' - 152: 228, # '\x98' - 153: 229, # '\x99' - 154: 230, # '\x9a' - 155: 231, # '\x9b' - 156: 232, # '\x9c' - 157: 233, # '\x9d' - 158: 234, # '\x9e' - 159: 235, # '\x9f' - 160: 236, # None - 161: 5, # 'ก' - 162: 30, # 'ข' - 163: 237, # 'ฃ' - 164: 24, # 'ค' - 165: 238, # 'ฅ' - 166: 75, # 'ฆ' - 167: 8, # 'ง' - 168: 26, # 'จ' - 169: 52, # 'ฉ' - 170: 34, # 'ช' - 171: 51, # 'ซ' - 172: 119, # 'ฌ' - 173: 47, # 'ญ' - 174: 58, # 'ฎ' - 175: 57, # 'ฏ' - 176: 49, # 'ฐ' - 177: 53, # 'ฑ' - 178: 55, # 'ฒ' - 179: 43, # 'ณ' - 180: 20, # 'ด' - 181: 19, # 'ต' - 182: 44, # 'ถ' - 183: 14, # 'ท' - 184: 48, # 'ธ' - 185: 3, # 'น' - 186: 17, # 'บ' - 187: 25, # 'ป' - 188: 39, # 'ผ' - 189: 62, # 'ฝ' - 190: 31, # 'พ' - 191: 54, # 'ฟ' - 192: 45, # 'ภ' - 193: 9, # 'ม' - 194: 16, # 'ย' - 195: 2, # 'ร' - 196: 61, # 'ฤ' - 197: 15, # 'ล' - 198: 239, # 'ฦ' - 199: 12, # 'ว' - 200: 42, # 'ศ' - 201: 46, # 'ษ' - 202: 18, # 'ส' - 203: 21, # 'ห' - 204: 76, # 'ฬ' - 205: 4, # 'อ' - 206: 66, # 'ฮ' - 207: 63, # 'ฯ' - 208: 22, # 'ะ' - 209: 10, # 'ั' - 210: 1, # 'า' - 211: 36, # 'ำ' - 212: 23, # 'ิ' - 213: 13, # 'ี' - 214: 40, # 'ึ' - 215: 27, # 'ื' - 216: 32, # 'ุ' - 217: 35, # 'ู' - 218: 86, # 'ฺ' - 219: 240, # None - 220: 241, # None - 221: 242, # None - 222: 243, # None - 223: 244, # '฿' - 224: 11, # 'เ' - 225: 28, # 'แ' - 226: 41, # 'โ' - 227: 29, # 'ใ' - 228: 33, # 'ไ' - 229: 245, # 'ๅ' - 230: 50, # 'ๆ' - 231: 37, # '็' - 232: 6, # '่' - 233: 7, # '้' - 234: 67, # '๊' - 235: 77, # '๋' - 236: 38, # '์' - 237: 93, # 'ํ' - 238: 246, # '๎' - 239: 247, # '๏' - 240: 68, # '๐' - 241: 56, # '๑' - 242: 59, # '๒' - 243: 65, # '๓' - 244: 69, # '๔' - 245: 60, # '๕' - 246: 70, # '๖' - 247: 80, # '๗' - 248: 71, # '๘' - 249: 87, # '๙' - 250: 248, # '๚' - 251: 249, # '๛' - 252: 250, # None - 253: 251, # None - 254: 252, # None - 255: 253, # None -} - -TIS_620_THAI_MODEL = SingleByteCharSetModel( - charset_name="TIS-620", - language="Thai", - char_to_order_map=TIS_620_THAI_CHAR_TO_ORDER, - language_model=THAI_LANG_MODEL, - typical_positive_ratio=0.926386, - keep_ascii_letters=False, - alphabet="กขฃคฅฆงจฉชซฌญฎฏฐฑฒณดตถทธนบปผฝพฟภมยรฤลฦวศษสหฬอฮฯะัาำิีึืฺุู฿เแโใไๅๆ็่้๊๋์ํ๎๏๐๑๒๓๔๕๖๗๘๙๚๛", -) diff --git a/venv/lib/python3.10/site-packages/chardet/langturkishmodel.py b/venv/lib/python3.10/site-packages/chardet/langturkishmodel.py deleted file mode 100644 index 64c94336cb2827db8e28ebe246c132779d390b89..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/langturkishmodel.py +++ /dev/null @@ -1,4380 +0,0 @@ -from chardet.sbcharsetprober import SingleByteCharSetModel - -# 3: Positive -# 2: Likely -# 1: Unlikely -# 0: Negative - -TURKISH_LANG_MODEL = { - 23: { # 'A' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 0, # 'c' - 12: 2, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 1, # 'g' - 25: 1, # 'h' - 3: 1, # 'i' - 24: 0, # 'j' - 10: 2, # 'k' - 5: 1, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 15: 0, # 'o' - 26: 0, # 'p' - 7: 1, # 'r' - 8: 1, # 's' - 9: 1, # 't' - 14: 1, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 3, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 0, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 37: { # 'B' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 2, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 2, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 1, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 1, # 'P' - 44: 0, # 'R' - 35: 1, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 1, # 'Y' - 56: 0, # 'Z' - 1: 2, # 'a' - 21: 0, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 0, # 'j' - 10: 0, # 'k' - 5: 0, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 15: 0, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 2, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 1, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 1, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 0, # 'ı' - 40: 1, # 'Ş' - 19: 1, # 'ş' - }, - 47: { # 'C' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 1, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 1, # 'L' - 20: 0, # 'M' - 46: 1, # 'N' - 42: 0, # 'O' - 48: 1, # 'P' - 44: 1, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 1, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 2, # 'j' - 10: 1, # 'k' - 5: 2, # 'l' - 13: 2, # 'm' - 4: 2, # 'n' - 15: 1, # 'o' - 26: 0, # 'p' - 7: 2, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 3, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 1, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 1, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 39: { # 'D' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 1, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 1, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 1, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 2, # 'a' - 21: 0, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 2, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 0, # 'j' - 10: 0, # 'k' - 5: 1, # 'l' - 13: 3, # 'm' - 4: 0, # 'n' - 15: 1, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 1, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 1, # 'z' - 63: 0, # '·' - 54: 1, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 1, # 'ğ' - 41: 0, # 'İ' - 6: 1, # 'ı' - 40: 1, # 'Ş' - 19: 0, # 'ş' - }, - 29: { # 'E' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 1, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 1, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 0, # 'c' - 12: 2, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 1, # 'g' - 25: 0, # 'h' - 3: 1, # 'i' - 24: 1, # 'j' - 10: 0, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 1, # 's' - 9: 1, # 't' - 14: 1, # 'u' - 32: 1, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 2, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 52: { # 'F' - 23: 0, # 'A' - 37: 1, # 'B' - 47: 1, # 'C' - 39: 1, # 'D' - 29: 1, # 'E' - 52: 2, # 'F' - 36: 0, # 'G' - 45: 2, # 'H' - 53: 1, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 1, # 'M' - 46: 1, # 'N' - 42: 1, # 'O' - 48: 2, # 'P' - 44: 1, # 'R' - 35: 1, # 'S' - 31: 1, # 'T' - 51: 1, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 2, # 'Y' - 56: 0, # 'Z' - 1: 0, # 'a' - 21: 1, # 'b' - 28: 1, # 'c' - 12: 1, # 'd' - 2: 0, # 'e' - 18: 1, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 2, # 'i' - 24: 1, # 'j' - 10: 0, # 'k' - 5: 0, # 'l' - 13: 1, # 'm' - 4: 2, # 'n' - 15: 1, # 'o' - 26: 0, # 'p' - 7: 2, # 'r' - 8: 1, # 's' - 9: 1, # 't' - 14: 1, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 1, # 'y' - 22: 1, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 1, # 'Ö' - 55: 2, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 2, # 'ö' - 17: 0, # 'ü' - 30: 1, # 'ğ' - 41: 1, # 'İ' - 6: 2, # 'ı' - 40: 0, # 'Ş' - 19: 2, # 'ş' - }, - 36: { # 'G' - 23: 1, # 'A' - 37: 0, # 'B' - 47: 1, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 1, # 'F' - 36: 2, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 2, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 2, # 'N' - 42: 1, # 'O' - 48: 1, # 'P' - 44: 1, # 'R' - 35: 1, # 'S' - 31: 0, # 'T' - 51: 1, # 'U' - 38: 2, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 1, # 'c' - 12: 0, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 1, # 'j' - 10: 1, # 'k' - 5: 0, # 'l' - 13: 3, # 'm' - 4: 2, # 'n' - 15: 0, # 'o' - 26: 1, # 'p' - 7: 0, # 'r' - 8: 1, # 's' - 9: 1, # 't' - 14: 3, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 1, # 'x' - 11: 0, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 1, # 'Ç' - 50: 2, # 'Ö' - 55: 0, # 'Ü' - 59: 1, # 'â' - 33: 2, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 1, # 'ğ' - 41: 1, # 'İ' - 6: 2, # 'ı' - 40: 2, # 'Ş' - 19: 1, # 'ş' - }, - 45: { # 'H' - 23: 0, # 'A' - 37: 1, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 2, # 'F' - 36: 2, # 'G' - 45: 1, # 'H' - 53: 1, # 'I' - 60: 0, # 'J' - 16: 2, # 'K' - 49: 1, # 'L' - 20: 0, # 'M' - 46: 1, # 'N' - 42: 1, # 'O' - 48: 1, # 'P' - 44: 0, # 'R' - 35: 2, # 'S' - 31: 0, # 'T' - 51: 1, # 'U' - 38: 2, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 2, # 'i' - 24: 0, # 'j' - 10: 1, # 'k' - 5: 0, # 'l' - 13: 2, # 'm' - 4: 0, # 'n' - 15: 1, # 'o' - 26: 1, # 'p' - 7: 1, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 3, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 1, # 'Ç' - 50: 1, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 0, # 'ü' - 30: 2, # 'ğ' - 41: 1, # 'İ' - 6: 0, # 'ı' - 40: 2, # 'Ş' - 19: 1, # 'ş' - }, - 53: { # 'I' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 1, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 2, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 1, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 2, # 'a' - 21: 0, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 2, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 0, # 'j' - 10: 0, # 'k' - 5: 2, # 'l' - 13: 2, # 'm' - 4: 0, # 'n' - 15: 0, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 2, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 1, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 2, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 0, # 'ı' - 40: 1, # 'Ş' - 19: 1, # 'ş' - }, - 60: { # 'J' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 1, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 0, # 'a' - 21: 1, # 'b' - 28: 0, # 'c' - 12: 1, # 'd' - 2: 0, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 1, # 'i' - 24: 0, # 'j' - 10: 0, # 'k' - 5: 0, # 'l' - 13: 0, # 'm' - 4: 1, # 'n' - 15: 0, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 1, # 's' - 9: 0, # 't' - 14: 0, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 0, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 16: { # 'K' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 3, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 2, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 2, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 2, # 'a' - 21: 3, # 'b' - 28: 0, # 'c' - 12: 3, # 'd' - 2: 1, # 'e' - 18: 3, # 'f' - 27: 3, # 'g' - 25: 3, # 'h' - 3: 3, # 'i' - 24: 2, # 'j' - 10: 3, # 'k' - 5: 0, # 'l' - 13: 0, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 1, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 0, # 'u' - 32: 3, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 2, # 'y' - 22: 1, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 2, # 'ü' - 30: 0, # 'ğ' - 41: 1, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 49: { # 'L' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 2, # 'E' - 52: 0, # 'F' - 36: 1, # 'G' - 45: 1, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 1, # 'M' - 46: 0, # 'N' - 42: 2, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 1, # 'Y' - 56: 0, # 'Z' - 1: 0, # 'a' - 21: 3, # 'b' - 28: 0, # 'c' - 12: 2, # 'd' - 2: 0, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 2, # 'i' - 24: 0, # 'j' - 10: 1, # 'k' - 5: 0, # 'l' - 13: 0, # 'm' - 4: 2, # 'n' - 15: 1, # 'o' - 26: 1, # 'p' - 7: 1, # 'r' - 8: 1, # 's' - 9: 1, # 't' - 14: 0, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 2, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 2, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 1, # 'ü' - 30: 1, # 'ğ' - 41: 0, # 'İ' - 6: 2, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 20: { # 'M' - 23: 1, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 1, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 2, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 1, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 2, # 'b' - 28: 0, # 'c' - 12: 3, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 1, # 'g' - 25: 1, # 'h' - 3: 2, # 'i' - 24: 2, # 'j' - 10: 2, # 'k' - 5: 2, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 1, # 'p' - 7: 3, # 'r' - 8: 0, # 's' - 9: 2, # 't' - 14: 3, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 2, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 3, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 46: { # 'N' - 23: 0, # 'A' - 37: 1, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 1, # 'F' - 36: 1, # 'G' - 45: 1, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 2, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 1, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 1, # 'R' - 35: 1, # 'S' - 31: 0, # 'T' - 51: 1, # 'U' - 38: 2, # 'V' - 62: 0, # 'W' - 43: 1, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 1, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 2, # 'j' - 10: 1, # 'k' - 5: 1, # 'l' - 13: 3, # 'm' - 4: 2, # 'n' - 15: 1, # 'o' - 26: 1, # 'p' - 7: 1, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 3, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 1, # 'x' - 11: 1, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 1, # 'Ç' - 50: 1, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 1, # 'İ' - 6: 2, # 'ı' - 40: 1, # 'Ş' - 19: 1, # 'ş' - }, - 42: { # 'O' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 1, # 'F' - 36: 0, # 'G' - 45: 1, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 2, # 'K' - 49: 1, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 2, # 'P' - 44: 1, # 'R' - 35: 1, # 'S' - 31: 0, # 'T' - 51: 1, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 2, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 0, # 'j' - 10: 0, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 0, # 'n' - 15: 1, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 2, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 2, # 'Ç' - 50: 1, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 2, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 0, # 'ü' - 30: 1, # 'ğ' - 41: 2, # 'İ' - 6: 1, # 'ı' - 40: 1, # 'Ş' - 19: 1, # 'ş' - }, - 48: { # 'P' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 2, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 2, # 'F' - 36: 1, # 'G' - 45: 1, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 2, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 1, # 'N' - 42: 1, # 'O' - 48: 1, # 'P' - 44: 0, # 'R' - 35: 1, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 2, # 'a' - 21: 0, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 0, # 'j' - 10: 1, # 'k' - 5: 0, # 'l' - 13: 2, # 'm' - 4: 0, # 'n' - 15: 2, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 2, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 2, # 'x' - 11: 0, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 1, # 'Ç' - 50: 2, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 2, # 'ö' - 17: 0, # 'ü' - 30: 1, # 'ğ' - 41: 1, # 'İ' - 6: 0, # 'ı' - 40: 2, # 'Ş' - 19: 1, # 'ş' - }, - 44: { # 'R' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 1, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 1, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 1, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 1, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 1, # 'b' - 28: 1, # 'c' - 12: 0, # 'd' - 2: 2, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 0, # 'j' - 10: 1, # 'k' - 5: 2, # 'l' - 13: 2, # 'm' - 4: 0, # 'n' - 15: 1, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 2, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 1, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 1, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 1, # 'ü' - 30: 1, # 'ğ' - 41: 0, # 'İ' - 6: 2, # 'ı' - 40: 1, # 'Ş' - 19: 1, # 'ş' - }, - 35: { # 'S' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 1, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 1, # 'F' - 36: 1, # 'G' - 45: 1, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 3, # 'K' - 49: 1, # 'L' - 20: 1, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 1, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 1, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 1, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 0, # 'j' - 10: 1, # 'k' - 5: 1, # 'l' - 13: 2, # 'm' - 4: 1, # 'n' - 15: 0, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 1, # 't' - 14: 2, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 1, # 'z' - 63: 0, # '·' - 54: 2, # 'Ç' - 50: 2, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 3, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 2, # 'Ş' - 19: 1, # 'ş' - }, - 31: { # 'T' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 1, # 'J' - 16: 2, # 'K' - 49: 0, # 'L' - 20: 1, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 2, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 2, # 'b' - 28: 0, # 'c' - 12: 1, # 'd' - 2: 3, # 'e' - 18: 2, # 'f' - 27: 2, # 'g' - 25: 0, # 'h' - 3: 1, # 'i' - 24: 1, # 'j' - 10: 2, # 'k' - 5: 2, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 2, # 'p' - 7: 2, # 'r' - 8: 0, # 's' - 9: 2, # 't' - 14: 2, # 'u' - 32: 1, # 'v' - 57: 1, # 'w' - 58: 1, # 'x' - 11: 2, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 1, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 51: { # 'U' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 1, # 'F' - 36: 1, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 1, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 1, # 'N' - 42: 0, # 'O' - 48: 1, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 1, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 1, # 'c' - 12: 0, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 2, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 0, # 'j' - 10: 1, # 'k' - 5: 1, # 'l' - 13: 3, # 'm' - 4: 2, # 'n' - 15: 0, # 'o' - 26: 1, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 2, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 1, # 'Ç' - 50: 1, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 1, # 'ğ' - 41: 1, # 'İ' - 6: 2, # 'ı' - 40: 0, # 'Ş' - 19: 1, # 'ş' - }, - 38: { # 'V' - 23: 1, # 'A' - 37: 1, # 'B' - 47: 1, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 2, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 3, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 1, # 'P' - 44: 1, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 1, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 0, # 'j' - 10: 0, # 'k' - 5: 2, # 'l' - 13: 2, # 'm' - 4: 0, # 'n' - 15: 2, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 1, # 't' - 14: 3, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 1, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 1, # 'Ç' - 50: 1, # 'Ö' - 55: 0, # 'Ü' - 59: 1, # 'â' - 33: 2, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 0, # 'ü' - 30: 1, # 'ğ' - 41: 1, # 'İ' - 6: 3, # 'ı' - 40: 2, # 'Ş' - 19: 1, # 'ş' - }, - 62: { # 'W' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 0, # 'a' - 21: 0, # 'b' - 28: 0, # 'c' - 12: 0, # 'd' - 2: 0, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 0, # 'j' - 10: 0, # 'k' - 5: 0, # 'l' - 13: 0, # 'm' - 4: 0, # 'n' - 15: 0, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 0, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 0, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 43: { # 'Y' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 1, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 2, # 'F' - 36: 0, # 'G' - 45: 1, # 'H' - 53: 1, # 'I' - 60: 0, # 'J' - 16: 2, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 2, # 'N' - 42: 0, # 'O' - 48: 2, # 'P' - 44: 1, # 'R' - 35: 1, # 'S' - 31: 0, # 'T' - 51: 1, # 'U' - 38: 2, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 2, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 1, # 'j' - 10: 1, # 'k' - 5: 1, # 'l' - 13: 3, # 'm' - 4: 0, # 'n' - 15: 2, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 3, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 1, # 'x' - 11: 0, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 1, # 'Ç' - 50: 2, # 'Ö' - 55: 1, # 'Ü' - 59: 1, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 0, # 'ü' - 30: 1, # 'ğ' - 41: 1, # 'İ' - 6: 0, # 'ı' - 40: 2, # 'Ş' - 19: 1, # 'ş' - }, - 56: { # 'Z' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 2, # 'Z' - 1: 2, # 'a' - 21: 1, # 'b' - 28: 0, # 'c' - 12: 0, # 'd' - 2: 2, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 2, # 'i' - 24: 1, # 'j' - 10: 0, # 'k' - 5: 0, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 15: 0, # 'o' - 26: 0, # 'p' - 7: 1, # 'r' - 8: 1, # 's' - 9: 0, # 't' - 14: 2, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 1, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 1, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 1: { # 'a' - 23: 3, # 'A' - 37: 0, # 'B' - 47: 1, # 'C' - 39: 0, # 'D' - 29: 3, # 'E' - 52: 0, # 'F' - 36: 1, # 'G' - 45: 1, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 3, # 'M' - 46: 1, # 'N' - 42: 0, # 'O' - 48: 1, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 3, # 'T' - 51: 0, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 2, # 'Z' - 1: 2, # 'a' - 21: 3, # 'b' - 28: 0, # 'c' - 12: 3, # 'd' - 2: 2, # 'e' - 18: 3, # 'f' - 27: 3, # 'g' - 25: 3, # 'h' - 3: 3, # 'i' - 24: 3, # 'j' - 10: 3, # 'k' - 5: 0, # 'l' - 13: 2, # 'm' - 4: 3, # 'n' - 15: 1, # 'o' - 26: 3, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 3, # 'u' - 32: 3, # 'v' - 57: 2, # 'w' - 58: 0, # 'x' - 11: 3, # 'y' - 22: 0, # 'z' - 63: 1, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 1, # 'î' - 34: 1, # 'ö' - 17: 3, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 1, # 'ş' - }, - 21: { # 'b' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 1, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 1, # 'J' - 16: 2, # 'K' - 49: 0, # 'L' - 20: 2, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 1, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 1, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 2, # 'b' - 28: 0, # 'c' - 12: 3, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 3, # 'g' - 25: 1, # 'h' - 3: 3, # 'i' - 24: 2, # 'j' - 10: 3, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 3, # 'p' - 7: 1, # 'r' - 8: 2, # 's' - 9: 2, # 't' - 14: 2, # 'u' - 32: 1, # 'v' - 57: 0, # 'w' - 58: 1, # 'x' - 11: 3, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 1, # 'ğ' - 41: 0, # 'İ' - 6: 2, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 28: { # 'c' - 23: 0, # 'A' - 37: 1, # 'B' - 47: 1, # 'C' - 39: 1, # 'D' - 29: 2, # 'E' - 52: 0, # 'F' - 36: 2, # 'G' - 45: 2, # 'H' - 53: 1, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 2, # 'M' - 46: 1, # 'N' - 42: 1, # 'O' - 48: 2, # 'P' - 44: 1, # 'R' - 35: 1, # 'S' - 31: 2, # 'T' - 51: 2, # 'U' - 38: 2, # 'V' - 62: 0, # 'W' - 43: 3, # 'Y' - 56: 0, # 'Z' - 1: 1, # 'a' - 21: 1, # 'b' - 28: 2, # 'c' - 12: 2, # 'd' - 2: 1, # 'e' - 18: 1, # 'f' - 27: 2, # 'g' - 25: 2, # 'h' - 3: 3, # 'i' - 24: 1, # 'j' - 10: 3, # 'k' - 5: 0, # 'l' - 13: 2, # 'm' - 4: 3, # 'n' - 15: 2, # 'o' - 26: 2, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 1, # 'u' - 32: 0, # 'v' - 57: 1, # 'w' - 58: 0, # 'x' - 11: 2, # 'y' - 22: 1, # 'z' - 63: 1, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 1, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 1, # 'î' - 34: 2, # 'ö' - 17: 2, # 'ü' - 30: 2, # 'ğ' - 41: 1, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 2, # 'ş' - }, - 12: { # 'd' - 23: 1, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 2, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 3, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 1, # 'S' - 31: 1, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 2, # 'b' - 28: 1, # 'c' - 12: 3, # 'd' - 2: 3, # 'e' - 18: 1, # 'f' - 27: 3, # 'g' - 25: 3, # 'h' - 3: 2, # 'i' - 24: 3, # 'j' - 10: 2, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 1, # 'o' - 26: 2, # 'p' - 7: 3, # 'r' - 8: 2, # 's' - 9: 2, # 't' - 14: 3, # 'u' - 32: 1, # 'v' - 57: 0, # 'w' - 58: 1, # 'x' - 11: 3, # 'y' - 22: 1, # 'z' - 63: 1, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 1, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 2, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 2: { # 'e' - 23: 2, # 'A' - 37: 0, # 'B' - 47: 2, # 'C' - 39: 0, # 'D' - 29: 3, # 'E' - 52: 1, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 1, # 'K' - 49: 0, # 'L' - 20: 3, # 'M' - 46: 1, # 'N' - 42: 0, # 'O' - 48: 1, # 'P' - 44: 1, # 'R' - 35: 0, # 'S' - 31: 3, # 'T' - 51: 0, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 1, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 3, # 'b' - 28: 0, # 'c' - 12: 3, # 'd' - 2: 2, # 'e' - 18: 3, # 'f' - 27: 3, # 'g' - 25: 3, # 'h' - 3: 3, # 'i' - 24: 3, # 'j' - 10: 3, # 'k' - 5: 0, # 'l' - 13: 2, # 'm' - 4: 3, # 'n' - 15: 1, # 'o' - 26: 3, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 3, # 'u' - 32: 3, # 'v' - 57: 2, # 'w' - 58: 0, # 'x' - 11: 3, # 'y' - 22: 1, # 'z' - 63: 1, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 3, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 18: { # 'f' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 2, # 'K' - 49: 0, # 'L' - 20: 2, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 2, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 1, # 'b' - 28: 0, # 'c' - 12: 3, # 'd' - 2: 3, # 'e' - 18: 2, # 'f' - 27: 1, # 'g' - 25: 1, # 'h' - 3: 1, # 'i' - 24: 1, # 'j' - 10: 1, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 2, # 'p' - 7: 1, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 1, # 'u' - 32: 2, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 1, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 1, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 1, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 27: { # 'g' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 1, # 'S' - 31: 1, # 'T' - 51: 0, # 'U' - 38: 2, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 1, # 'b' - 28: 0, # 'c' - 12: 1, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 2, # 'g' - 25: 1, # 'h' - 3: 2, # 'i' - 24: 3, # 'j' - 10: 2, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 2, # 'n' - 15: 0, # 'o' - 26: 1, # 'p' - 7: 2, # 'r' - 8: 2, # 's' - 9: 3, # 't' - 14: 3, # 'u' - 32: 1, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 1, # 'y' - 22: 0, # 'z' - 63: 1, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 2, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 25: { # 'h' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 2, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 0, # 'c' - 12: 2, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 1, # 'g' - 25: 2, # 'h' - 3: 2, # 'i' - 24: 3, # 'j' - 10: 3, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 1, # 'o' - 26: 1, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 2, # 't' - 14: 3, # 'u' - 32: 2, # 'v' - 57: 1, # 'w' - 58: 0, # 'x' - 11: 1, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 3: { # 'i' - 23: 2, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 1, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 3, # 'M' - 46: 0, # 'N' - 42: 1, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 1, # 'S' - 31: 2, # 'T' - 51: 0, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 2, # 'b' - 28: 0, # 'c' - 12: 3, # 'd' - 2: 3, # 'e' - 18: 2, # 'f' - 27: 3, # 'g' - 25: 1, # 'h' - 3: 3, # 'i' - 24: 2, # 'j' - 10: 3, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 1, # 'o' - 26: 3, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 3, # 'u' - 32: 2, # 'v' - 57: 1, # 'w' - 58: 1, # 'x' - 11: 3, # 'y' - 22: 1, # 'z' - 63: 1, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 1, # 'Ü' - 59: 0, # 'â' - 33: 2, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 3, # 'ü' - 30: 0, # 'ğ' - 41: 1, # 'İ' - 6: 2, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 24: { # 'j' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 1, # 'J' - 16: 2, # 'K' - 49: 0, # 'L' - 20: 2, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 1, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 1, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 1, # 'Z' - 1: 3, # 'a' - 21: 1, # 'b' - 28: 1, # 'c' - 12: 3, # 'd' - 2: 3, # 'e' - 18: 2, # 'f' - 27: 1, # 'g' - 25: 1, # 'h' - 3: 2, # 'i' - 24: 1, # 'j' - 10: 2, # 'k' - 5: 2, # 'l' - 13: 3, # 'm' - 4: 2, # 'n' - 15: 0, # 'o' - 26: 1, # 'p' - 7: 2, # 'r' - 8: 3, # 's' - 9: 2, # 't' - 14: 3, # 'u' - 32: 2, # 'v' - 57: 0, # 'w' - 58: 2, # 'x' - 11: 1, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 1, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 10: { # 'k' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 2, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 3, # 'T' - 51: 0, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 1, # 'Z' - 1: 3, # 'a' - 21: 2, # 'b' - 28: 0, # 'c' - 12: 2, # 'd' - 2: 3, # 'e' - 18: 1, # 'f' - 27: 2, # 'g' - 25: 2, # 'h' - 3: 3, # 'i' - 24: 2, # 'j' - 10: 2, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 3, # 'p' - 7: 2, # 'r' - 8: 2, # 's' - 9: 2, # 't' - 14: 3, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 1, # 'x' - 11: 3, # 'y' - 22: 0, # 'z' - 63: 1, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 3, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 3, # 'ü' - 30: 1, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 1, # 'ş' - }, - 5: { # 'l' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 3, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 2, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 1, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 0, # 'a' - 21: 3, # 'b' - 28: 0, # 'c' - 12: 3, # 'd' - 2: 1, # 'e' - 18: 3, # 'f' - 27: 3, # 'g' - 25: 2, # 'h' - 3: 3, # 'i' - 24: 2, # 'j' - 10: 3, # 'k' - 5: 1, # 'l' - 13: 1, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 2, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 2, # 'u' - 32: 2, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 3, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 2, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 13: { # 'm' - 23: 1, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 3, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 3, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 3, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 1, # 'Y' - 56: 0, # 'Z' - 1: 2, # 'a' - 21: 3, # 'b' - 28: 0, # 'c' - 12: 3, # 'd' - 2: 2, # 'e' - 18: 3, # 'f' - 27: 3, # 'g' - 25: 3, # 'h' - 3: 3, # 'i' - 24: 3, # 'j' - 10: 3, # 'k' - 5: 0, # 'l' - 13: 2, # 'm' - 4: 3, # 'n' - 15: 1, # 'o' - 26: 2, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 2, # 'u' - 32: 2, # 'v' - 57: 1, # 'w' - 58: 0, # 'x' - 11: 3, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 3, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 1, # 'ş' - }, - 4: { # 'n' - 23: 1, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 1, # 'H' - 53: 0, # 'I' - 60: 2, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 3, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 2, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 2, # 'b' - 28: 1, # 'c' - 12: 3, # 'd' - 2: 3, # 'e' - 18: 1, # 'f' - 27: 2, # 'g' - 25: 3, # 'h' - 3: 2, # 'i' - 24: 2, # 'j' - 10: 3, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 1, # 'o' - 26: 3, # 'p' - 7: 2, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 3, # 'u' - 32: 2, # 'v' - 57: 0, # 'w' - 58: 2, # 'x' - 11: 3, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 2, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 1, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 15: { # 'o' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 1, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 2, # 'F' - 36: 1, # 'G' - 45: 1, # 'H' - 53: 1, # 'I' - 60: 0, # 'J' - 16: 3, # 'K' - 49: 2, # 'L' - 20: 0, # 'M' - 46: 2, # 'N' - 42: 1, # 'O' - 48: 2, # 'P' - 44: 1, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 1, # 'i' - 24: 2, # 'j' - 10: 1, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 2, # 'n' - 15: 2, # 'o' - 26: 0, # 'p' - 7: 1, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 3, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 2, # 'x' - 11: 0, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 1, # 'Ç' - 50: 2, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 3, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 0, # 'ü' - 30: 2, # 'ğ' - 41: 2, # 'İ' - 6: 3, # 'ı' - 40: 2, # 'Ş' - 19: 2, # 'ş' - }, - 26: { # 'p' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 1, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 1, # 'b' - 28: 0, # 'c' - 12: 1, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 1, # 'g' - 25: 1, # 'h' - 3: 2, # 'i' - 24: 3, # 'j' - 10: 1, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 2, # 'n' - 15: 0, # 'o' - 26: 2, # 'p' - 7: 2, # 'r' - 8: 1, # 's' - 9: 1, # 't' - 14: 3, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 1, # 'x' - 11: 1, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 3, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 1, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 7: { # 'r' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 1, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 2, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 2, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 2, # 'T' - 51: 1, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 1, # 'Z' - 1: 3, # 'a' - 21: 1, # 'b' - 28: 0, # 'c' - 12: 3, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 2, # 'g' - 25: 3, # 'h' - 3: 2, # 'i' - 24: 2, # 'j' - 10: 3, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 2, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 3, # 'u' - 32: 2, # 'v' - 57: 0, # 'w' - 58: 1, # 'x' - 11: 2, # 'y' - 22: 0, # 'z' - 63: 1, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 2, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 3, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 2, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 8: { # 's' - 23: 1, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 1, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 3, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 2, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 1, # 'Z' - 1: 3, # 'a' - 21: 2, # 'b' - 28: 1, # 'c' - 12: 3, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 2, # 'g' - 25: 2, # 'h' - 3: 2, # 'i' - 24: 3, # 'j' - 10: 3, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 3, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 3, # 'u' - 32: 2, # 'v' - 57: 0, # 'w' - 58: 1, # 'x' - 11: 2, # 'y' - 22: 1, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 2, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 2, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 1, # 'ş' - }, - 9: { # 't' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 1, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 2, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 2, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 1, # 'Z' - 1: 3, # 'a' - 21: 3, # 'b' - 28: 0, # 'c' - 12: 3, # 'd' - 2: 3, # 'e' - 18: 2, # 'f' - 27: 2, # 'g' - 25: 2, # 'h' - 3: 2, # 'i' - 24: 2, # 'j' - 10: 3, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 2, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 3, # 'u' - 32: 3, # 'v' - 57: 0, # 'w' - 58: 2, # 'x' - 11: 2, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 3, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 2, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 14: { # 'u' - 23: 3, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 3, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 1, # 'H' - 53: 0, # 'I' - 60: 1, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 3, # 'M' - 46: 2, # 'N' - 42: 0, # 'O' - 48: 1, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 3, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 1, # 'Y' - 56: 2, # 'Z' - 1: 2, # 'a' - 21: 3, # 'b' - 28: 0, # 'c' - 12: 3, # 'd' - 2: 2, # 'e' - 18: 2, # 'f' - 27: 3, # 'g' - 25: 3, # 'h' - 3: 3, # 'i' - 24: 2, # 'j' - 10: 3, # 'k' - 5: 0, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 3, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 3, # 'u' - 32: 2, # 'v' - 57: 2, # 'w' - 58: 0, # 'x' - 11: 3, # 'y' - 22: 0, # 'z' - 63: 1, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 3, # 'ü' - 30: 1, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 32: { # 'v' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 1, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 0, # 'c' - 12: 3, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 1, # 'j' - 10: 1, # 'k' - 5: 3, # 'l' - 13: 2, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 1, # 'p' - 7: 1, # 'r' - 8: 2, # 's' - 9: 3, # 't' - 14: 3, # 'u' - 32: 1, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 2, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 1, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 57: { # 'w' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 1, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 1, # 'a' - 21: 0, # 'b' - 28: 0, # 'c' - 12: 0, # 'd' - 2: 2, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 1, # 'h' - 3: 0, # 'i' - 24: 0, # 'j' - 10: 1, # 'k' - 5: 0, # 'l' - 13: 0, # 'm' - 4: 1, # 'n' - 15: 0, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 1, # 's' - 9: 0, # 't' - 14: 1, # 'u' - 32: 0, # 'v' - 57: 2, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 0, # 'z' - 63: 1, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 1, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 0, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 58: { # 'x' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 1, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 1, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 1, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 0, # 'a' - 21: 1, # 'b' - 28: 0, # 'c' - 12: 2, # 'd' - 2: 1, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 2, # 'i' - 24: 2, # 'j' - 10: 1, # 'k' - 5: 0, # 'l' - 13: 0, # 'm' - 4: 2, # 'n' - 15: 0, # 'o' - 26: 0, # 'p' - 7: 1, # 'r' - 8: 2, # 's' - 9: 1, # 't' - 14: 0, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 2, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 1, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 2, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 11: { # 'y' - 23: 1, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 1, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 1, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 1, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 1, # 'Y' - 56: 1, # 'Z' - 1: 3, # 'a' - 21: 1, # 'b' - 28: 0, # 'c' - 12: 2, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 2, # 'g' - 25: 2, # 'h' - 3: 2, # 'i' - 24: 1, # 'j' - 10: 2, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 1, # 'p' - 7: 2, # 'r' - 8: 1, # 's' - 9: 2, # 't' - 14: 3, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 1, # 'x' - 11: 3, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 3, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 2, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 22: { # 'z' - 23: 2, # 'A' - 37: 2, # 'B' - 47: 1, # 'C' - 39: 2, # 'D' - 29: 3, # 'E' - 52: 1, # 'F' - 36: 2, # 'G' - 45: 2, # 'H' - 53: 1, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 3, # 'M' - 46: 2, # 'N' - 42: 2, # 'O' - 48: 2, # 'P' - 44: 1, # 'R' - 35: 1, # 'S' - 31: 3, # 'T' - 51: 2, # 'U' - 38: 2, # 'V' - 62: 0, # 'W' - 43: 2, # 'Y' - 56: 1, # 'Z' - 1: 1, # 'a' - 21: 2, # 'b' - 28: 1, # 'c' - 12: 2, # 'd' - 2: 2, # 'e' - 18: 3, # 'f' - 27: 2, # 'g' - 25: 2, # 'h' - 3: 3, # 'i' - 24: 2, # 'j' - 10: 3, # 'k' - 5: 0, # 'l' - 13: 2, # 'm' - 4: 3, # 'n' - 15: 2, # 'o' - 26: 2, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 0, # 'u' - 32: 2, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 3, # 'y' - 22: 2, # 'z' - 63: 1, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 2, # 'Ü' - 59: 1, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 2, # 'ö' - 17: 2, # 'ü' - 30: 2, # 'ğ' - 41: 1, # 'İ' - 6: 3, # 'ı' - 40: 1, # 'Ş' - 19: 2, # 'ş' - }, - 63: { # '·' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 0, # 'a' - 21: 0, # 'b' - 28: 0, # 'c' - 12: 0, # 'd' - 2: 1, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 0, # 'j' - 10: 0, # 'k' - 5: 0, # 'l' - 13: 2, # 'm' - 4: 0, # 'n' - 15: 0, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 2, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 0, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 54: { # 'Ç' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 1, # 'C' - 39: 1, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 1, # 'G' - 45: 1, # 'H' - 53: 1, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 1, # 'O' - 48: 1, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 1, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 2, # 'Y' - 56: 0, # 'Z' - 1: 0, # 'a' - 21: 1, # 'b' - 28: 0, # 'c' - 12: 1, # 'd' - 2: 0, # 'e' - 18: 0, # 'f' - 27: 1, # 'g' - 25: 0, # 'h' - 3: 3, # 'i' - 24: 0, # 'j' - 10: 1, # 'k' - 5: 0, # 'l' - 13: 0, # 'm' - 4: 2, # 'n' - 15: 1, # 'o' - 26: 0, # 'p' - 7: 2, # 'r' - 8: 0, # 's' - 9: 1, # 't' - 14: 0, # 'u' - 32: 2, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 2, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 2, # 'ı' - 40: 0, # 'Ş' - 19: 1, # 'ş' - }, - 50: { # 'Ö' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 1, # 'C' - 39: 1, # 'D' - 29: 2, # 'E' - 52: 0, # 'F' - 36: 1, # 'G' - 45: 2, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 1, # 'M' - 46: 1, # 'N' - 42: 2, # 'O' - 48: 2, # 'P' - 44: 1, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 1, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 2, # 'Y' - 56: 0, # 'Z' - 1: 0, # 'a' - 21: 2, # 'b' - 28: 1, # 'c' - 12: 2, # 'd' - 2: 0, # 'e' - 18: 1, # 'f' - 27: 1, # 'g' - 25: 1, # 'h' - 3: 2, # 'i' - 24: 0, # 'j' - 10: 2, # 'k' - 5: 0, # 'l' - 13: 0, # 'm' - 4: 3, # 'n' - 15: 2, # 'o' - 26: 2, # 'p' - 7: 3, # 'r' - 8: 1, # 's' - 9: 2, # 't' - 14: 0, # 'u' - 32: 1, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 1, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 2, # 'ö' - 17: 2, # 'ü' - 30: 1, # 'ğ' - 41: 0, # 'İ' - 6: 2, # 'ı' - 40: 0, # 'Ş' - 19: 1, # 'ş' - }, - 55: { # 'Ü' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 2, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 1, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 1, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 2, # 'a' - 21: 0, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 2, # 'e' - 18: 0, # 'f' - 27: 1, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 0, # 'j' - 10: 0, # 'k' - 5: 1, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 15: 0, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 1, # 't' - 14: 2, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 1, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 1, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 0, # 'ü' - 30: 1, # 'ğ' - 41: 1, # 'İ' - 6: 0, # 'ı' - 40: 0, # 'Ş' - 19: 1, # 'ş' - }, - 59: { # 'â' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 1, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 1, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 2, # 'a' - 21: 0, # 'b' - 28: 0, # 'c' - 12: 0, # 'd' - 2: 2, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 0, # 'j' - 10: 0, # 'k' - 5: 0, # 'l' - 13: 2, # 'm' - 4: 0, # 'n' - 15: 1, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 2, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 1, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 1, # 'ı' - 40: 1, # 'Ş' - 19: 0, # 'ş' - }, - 33: { # 'ç' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 3, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 1, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 2, # 'T' - 51: 0, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 0, # 'Z' - 1: 0, # 'a' - 21: 3, # 'b' - 28: 0, # 'c' - 12: 2, # 'd' - 2: 0, # 'e' - 18: 2, # 'f' - 27: 1, # 'g' - 25: 3, # 'h' - 3: 3, # 'i' - 24: 0, # 'j' - 10: 3, # 'k' - 5: 0, # 'l' - 13: 0, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 1, # 'p' - 7: 3, # 'r' - 8: 2, # 's' - 9: 3, # 't' - 14: 0, # 'u' - 32: 2, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 2, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 1, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 61: { # 'î' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 0, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 0, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 1, # 'Z' - 1: 2, # 'a' - 21: 0, # 'b' - 28: 0, # 'c' - 12: 0, # 'd' - 2: 2, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 1, # 'j' - 10: 0, # 'k' - 5: 0, # 'l' - 13: 1, # 'm' - 4: 1, # 'n' - 15: 0, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 1, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 1, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 1, # 'î' - 34: 0, # 'ö' - 17: 0, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 1, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 34: { # 'ö' - 23: 0, # 'A' - 37: 1, # 'B' - 47: 1, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 2, # 'F' - 36: 1, # 'G' - 45: 1, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 3, # 'K' - 49: 1, # 'L' - 20: 0, # 'M' - 46: 1, # 'N' - 42: 1, # 'O' - 48: 2, # 'P' - 44: 1, # 'R' - 35: 1, # 'S' - 31: 1, # 'T' - 51: 1, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 1, # 'Z' - 1: 3, # 'a' - 21: 1, # 'b' - 28: 2, # 'c' - 12: 1, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 2, # 'g' - 25: 2, # 'h' - 3: 1, # 'i' - 24: 2, # 'j' - 10: 1, # 'k' - 5: 2, # 'l' - 13: 3, # 'm' - 4: 2, # 'n' - 15: 2, # 'o' - 26: 0, # 'p' - 7: 0, # 'r' - 8: 3, # 's' - 9: 1, # 't' - 14: 3, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 1, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 1, # 'Ç' - 50: 2, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 2, # 'ç' - 61: 0, # 'î' - 34: 2, # 'ö' - 17: 0, # 'ü' - 30: 2, # 'ğ' - 41: 1, # 'İ' - 6: 1, # 'ı' - 40: 2, # 'Ş' - 19: 1, # 'ş' - }, - 17: { # 'ü' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 1, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 0, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 1, # 'J' - 16: 1, # 'K' - 49: 0, # 'L' - 20: 1, # 'M' - 46: 0, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 1, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 0, # 'Y' - 56: 1, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 0, # 'c' - 12: 1, # 'd' - 2: 3, # 'e' - 18: 1, # 'f' - 27: 2, # 'g' - 25: 0, # 'h' - 3: 1, # 'i' - 24: 1, # 'j' - 10: 2, # 'k' - 5: 3, # 'l' - 13: 2, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 2, # 'p' - 7: 2, # 'r' - 8: 3, # 's' - 9: 2, # 't' - 14: 3, # 'u' - 32: 1, # 'v' - 57: 1, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 2, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 2, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 30: { # 'ğ' - 23: 0, # 'A' - 37: 2, # 'B' - 47: 1, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 2, # 'F' - 36: 1, # 'G' - 45: 0, # 'H' - 53: 1, # 'I' - 60: 0, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 1, # 'M' - 46: 2, # 'N' - 42: 2, # 'O' - 48: 1, # 'P' - 44: 1, # 'R' - 35: 0, # 'S' - 31: 1, # 'T' - 51: 0, # 'U' - 38: 2, # 'V' - 62: 0, # 'W' - 43: 2, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 0, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 2, # 'e' - 18: 0, # 'f' - 27: 0, # 'g' - 25: 0, # 'h' - 3: 0, # 'i' - 24: 3, # 'j' - 10: 1, # 'k' - 5: 2, # 'l' - 13: 3, # 'm' - 4: 0, # 'n' - 15: 1, # 'o' - 26: 0, # 'p' - 7: 1, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 3, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 2, # 'Ç' - 50: 2, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 0, # 'î' - 34: 2, # 'ö' - 17: 0, # 'ü' - 30: 1, # 'ğ' - 41: 2, # 'İ' - 6: 2, # 'ı' - 40: 2, # 'Ş' - 19: 1, # 'ş' - }, - 41: { # 'İ' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 1, # 'C' - 39: 1, # 'D' - 29: 1, # 'E' - 52: 0, # 'F' - 36: 2, # 'G' - 45: 2, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 2, # 'M' - 46: 1, # 'N' - 42: 1, # 'O' - 48: 2, # 'P' - 44: 0, # 'R' - 35: 1, # 'S' - 31: 1, # 'T' - 51: 1, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 2, # 'Y' - 56: 0, # 'Z' - 1: 1, # 'a' - 21: 2, # 'b' - 28: 1, # 'c' - 12: 2, # 'd' - 2: 1, # 'e' - 18: 0, # 'f' - 27: 3, # 'g' - 25: 2, # 'h' - 3: 2, # 'i' - 24: 2, # 'j' - 10: 2, # 'k' - 5: 0, # 'l' - 13: 1, # 'm' - 4: 3, # 'n' - 15: 1, # 'o' - 26: 1, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 2, # 't' - 14: 0, # 'u' - 32: 0, # 'v' - 57: 1, # 'w' - 58: 0, # 'x' - 11: 2, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 1, # 'Ü' - 59: 1, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 1, # 'ö' - 17: 1, # 'ü' - 30: 2, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 1, # 'ş' - }, - 6: { # 'ı' - 23: 2, # 'A' - 37: 0, # 'B' - 47: 0, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 0, # 'F' - 36: 1, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 2, # 'J' - 16: 3, # 'K' - 49: 0, # 'L' - 20: 3, # 'M' - 46: 1, # 'N' - 42: 0, # 'O' - 48: 0, # 'P' - 44: 0, # 'R' - 35: 0, # 'S' - 31: 2, # 'T' - 51: 0, # 'U' - 38: 0, # 'V' - 62: 0, # 'W' - 43: 2, # 'Y' - 56: 1, # 'Z' - 1: 3, # 'a' - 21: 2, # 'b' - 28: 1, # 'c' - 12: 3, # 'd' - 2: 3, # 'e' - 18: 3, # 'f' - 27: 3, # 'g' - 25: 2, # 'h' - 3: 3, # 'i' - 24: 3, # 'j' - 10: 3, # 'k' - 5: 3, # 'l' - 13: 3, # 'm' - 4: 3, # 'n' - 15: 0, # 'o' - 26: 3, # 'p' - 7: 3, # 'r' - 8: 3, # 's' - 9: 3, # 't' - 14: 3, # 'u' - 32: 3, # 'v' - 57: 1, # 'w' - 58: 1, # 'x' - 11: 3, # 'y' - 22: 0, # 'z' - 63: 1, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 2, # 'ç' - 61: 0, # 'î' - 34: 0, # 'ö' - 17: 3, # 'ü' - 30: 0, # 'ğ' - 41: 0, # 'İ' - 6: 3, # 'ı' - 40: 0, # 'Ş' - 19: 0, # 'ş' - }, - 40: { # 'Ş' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 1, # 'C' - 39: 1, # 'D' - 29: 1, # 'E' - 52: 0, # 'F' - 36: 1, # 'G' - 45: 2, # 'H' - 53: 1, # 'I' - 60: 0, # 'J' - 16: 0, # 'K' - 49: 0, # 'L' - 20: 2, # 'M' - 46: 1, # 'N' - 42: 1, # 'O' - 48: 2, # 'P' - 44: 2, # 'R' - 35: 1, # 'S' - 31: 1, # 'T' - 51: 0, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 2, # 'Y' - 56: 1, # 'Z' - 1: 0, # 'a' - 21: 2, # 'b' - 28: 0, # 'c' - 12: 2, # 'd' - 2: 0, # 'e' - 18: 3, # 'f' - 27: 0, # 'g' - 25: 2, # 'h' - 3: 3, # 'i' - 24: 2, # 'j' - 10: 1, # 'k' - 5: 0, # 'l' - 13: 1, # 'm' - 4: 3, # 'n' - 15: 2, # 'o' - 26: 0, # 'p' - 7: 3, # 'r' - 8: 2, # 's' - 9: 2, # 't' - 14: 1, # 'u' - 32: 3, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 2, # 'y' - 22: 0, # 'z' - 63: 0, # '·' - 54: 0, # 'Ç' - 50: 0, # 'Ö' - 55: 1, # 'Ü' - 59: 0, # 'â' - 33: 0, # 'ç' - 61: 0, # 'î' - 34: 2, # 'ö' - 17: 1, # 'ü' - 30: 2, # 'ğ' - 41: 0, # 'İ' - 6: 2, # 'ı' - 40: 1, # 'Ş' - 19: 2, # 'ş' - }, - 19: { # 'ş' - 23: 0, # 'A' - 37: 0, # 'B' - 47: 1, # 'C' - 39: 0, # 'D' - 29: 0, # 'E' - 52: 2, # 'F' - 36: 1, # 'G' - 45: 0, # 'H' - 53: 0, # 'I' - 60: 0, # 'J' - 16: 3, # 'K' - 49: 2, # 'L' - 20: 0, # 'M' - 46: 1, # 'N' - 42: 1, # 'O' - 48: 1, # 'P' - 44: 1, # 'R' - 35: 1, # 'S' - 31: 0, # 'T' - 51: 1, # 'U' - 38: 1, # 'V' - 62: 0, # 'W' - 43: 1, # 'Y' - 56: 0, # 'Z' - 1: 3, # 'a' - 21: 1, # 'b' - 28: 2, # 'c' - 12: 0, # 'd' - 2: 3, # 'e' - 18: 0, # 'f' - 27: 2, # 'g' - 25: 1, # 'h' - 3: 1, # 'i' - 24: 0, # 'j' - 10: 2, # 'k' - 5: 2, # 'l' - 13: 3, # 'm' - 4: 0, # 'n' - 15: 0, # 'o' - 26: 1, # 'p' - 7: 3, # 'r' - 8: 0, # 's' - 9: 0, # 't' - 14: 3, # 'u' - 32: 0, # 'v' - 57: 0, # 'w' - 58: 0, # 'x' - 11: 0, # 'y' - 22: 2, # 'z' - 63: 0, # '·' - 54: 1, # 'Ç' - 50: 2, # 'Ö' - 55: 0, # 'Ü' - 59: 0, # 'â' - 33: 1, # 'ç' - 61: 1, # 'î' - 34: 2, # 'ö' - 17: 0, # 'ü' - 30: 1, # 'ğ' - 41: 1, # 'İ' - 6: 1, # 'ı' - 40: 1, # 'Ş' - 19: 1, # 'ş' - }, -} - -# 255: Undefined characters that did not exist in training text -# 254: Carriage/Return -# 253: symbol (punctuation) that does not belong to word -# 252: 0 - 9 -# 251: Control characters - -# Character Mapping Table(s): -ISO_8859_9_TURKISH_CHAR_TO_ORDER = { - 0: 255, # '\x00' - 1: 255, # '\x01' - 2: 255, # '\x02' - 3: 255, # '\x03' - 4: 255, # '\x04' - 5: 255, # '\x05' - 6: 255, # '\x06' - 7: 255, # '\x07' - 8: 255, # '\x08' - 9: 255, # '\t' - 10: 255, # '\n' - 11: 255, # '\x0b' - 12: 255, # '\x0c' - 13: 255, # '\r' - 14: 255, # '\x0e' - 15: 255, # '\x0f' - 16: 255, # '\x10' - 17: 255, # '\x11' - 18: 255, # '\x12' - 19: 255, # '\x13' - 20: 255, # '\x14' - 21: 255, # '\x15' - 22: 255, # '\x16' - 23: 255, # '\x17' - 24: 255, # '\x18' - 25: 255, # '\x19' - 26: 255, # '\x1a' - 27: 255, # '\x1b' - 28: 255, # '\x1c' - 29: 255, # '\x1d' - 30: 255, # '\x1e' - 31: 255, # '\x1f' - 32: 255, # ' ' - 33: 255, # '!' - 34: 255, # '"' - 35: 255, # '#' - 36: 255, # '$' - 37: 255, # '%' - 38: 255, # '&' - 39: 255, # "'" - 40: 255, # '(' - 41: 255, # ')' - 42: 255, # '*' - 43: 255, # '+' - 44: 255, # ',' - 45: 255, # '-' - 46: 255, # '.' - 47: 255, # '/' - 48: 255, # '0' - 49: 255, # '1' - 50: 255, # '2' - 51: 255, # '3' - 52: 255, # '4' - 53: 255, # '5' - 54: 255, # '6' - 55: 255, # '7' - 56: 255, # '8' - 57: 255, # '9' - 58: 255, # ':' - 59: 255, # ';' - 60: 255, # '<' - 61: 255, # '=' - 62: 255, # '>' - 63: 255, # '?' - 64: 255, # '@' - 65: 23, # 'A' - 66: 37, # 'B' - 67: 47, # 'C' - 68: 39, # 'D' - 69: 29, # 'E' - 70: 52, # 'F' - 71: 36, # 'G' - 72: 45, # 'H' - 73: 53, # 'I' - 74: 60, # 'J' - 75: 16, # 'K' - 76: 49, # 'L' - 77: 20, # 'M' - 78: 46, # 'N' - 79: 42, # 'O' - 80: 48, # 'P' - 81: 69, # 'Q' - 82: 44, # 'R' - 83: 35, # 'S' - 84: 31, # 'T' - 85: 51, # 'U' - 86: 38, # 'V' - 87: 62, # 'W' - 88: 65, # 'X' - 89: 43, # 'Y' - 90: 56, # 'Z' - 91: 255, # '[' - 92: 255, # '\\' - 93: 255, # ']' - 94: 255, # '^' - 95: 255, # '_' - 96: 255, # '`' - 97: 1, # 'a' - 98: 21, # 'b' - 99: 28, # 'c' - 100: 12, # 'd' - 101: 2, # 'e' - 102: 18, # 'f' - 103: 27, # 'g' - 104: 25, # 'h' - 105: 3, # 'i' - 106: 24, # 'j' - 107: 10, # 'k' - 108: 5, # 'l' - 109: 13, # 'm' - 110: 4, # 'n' - 111: 15, # 'o' - 112: 26, # 'p' - 113: 64, # 'q' - 114: 7, # 'r' - 115: 8, # 's' - 116: 9, # 't' - 117: 14, # 'u' - 118: 32, # 'v' - 119: 57, # 'w' - 120: 58, # 'x' - 121: 11, # 'y' - 122: 22, # 'z' - 123: 255, # '{' - 124: 255, # '|' - 125: 255, # '}' - 126: 255, # '~' - 127: 255, # '\x7f' - 128: 180, # '\x80' - 129: 179, # '\x81' - 130: 178, # '\x82' - 131: 177, # '\x83' - 132: 176, # '\x84' - 133: 175, # '\x85' - 134: 174, # '\x86' - 135: 173, # '\x87' - 136: 172, # '\x88' - 137: 171, # '\x89' - 138: 170, # '\x8a' - 139: 169, # '\x8b' - 140: 168, # '\x8c' - 141: 167, # '\x8d' - 142: 166, # '\x8e' - 143: 165, # '\x8f' - 144: 164, # '\x90' - 145: 163, # '\x91' - 146: 162, # '\x92' - 147: 161, # '\x93' - 148: 160, # '\x94' - 149: 159, # '\x95' - 150: 101, # '\x96' - 151: 158, # '\x97' - 152: 157, # '\x98' - 153: 156, # '\x99' - 154: 155, # '\x9a' - 155: 154, # '\x9b' - 156: 153, # '\x9c' - 157: 152, # '\x9d' - 158: 151, # '\x9e' - 159: 106, # '\x9f' - 160: 150, # '\xa0' - 161: 149, # '¡' - 162: 148, # '¢' - 163: 147, # '£' - 164: 146, # '¤' - 165: 145, # '¥' - 166: 144, # '¦' - 167: 100, # '§' - 168: 143, # '¨' - 169: 142, # '©' - 170: 141, # 'ª' - 171: 140, # '«' - 172: 139, # '¬' - 173: 138, # '\xad' - 174: 137, # '®' - 175: 136, # '¯' - 176: 94, # '°' - 177: 80, # '±' - 178: 93, # '²' - 179: 135, # '³' - 180: 105, # '´' - 181: 134, # 'µ' - 182: 133, # '¶' - 183: 63, # '·' - 184: 132, # '¸' - 185: 131, # '¹' - 186: 130, # 'º' - 187: 129, # '»' - 188: 128, # '¼' - 189: 127, # '½' - 190: 126, # '¾' - 191: 125, # '¿' - 192: 124, # 'À' - 193: 104, # 'Á' - 194: 73, # 'Â' - 195: 99, # 'Ã' - 196: 79, # 'Ä' - 197: 85, # 'Å' - 198: 123, # 'Æ' - 199: 54, # 'Ç' - 200: 122, # 'È' - 201: 98, # 'É' - 202: 92, # 'Ê' - 203: 121, # 'Ë' - 204: 120, # 'Ì' - 205: 91, # 'Í' - 206: 103, # 'Î' - 207: 119, # 'Ï' - 208: 68, # 'Ğ' - 209: 118, # 'Ñ' - 210: 117, # 'Ò' - 211: 97, # 'Ó' - 212: 116, # 'Ô' - 213: 115, # 'Õ' - 214: 50, # 'Ö' - 215: 90, # '×' - 216: 114, # 'Ø' - 217: 113, # 'Ù' - 218: 112, # 'Ú' - 219: 111, # 'Û' - 220: 55, # 'Ü' - 221: 41, # 'İ' - 222: 40, # 'Ş' - 223: 86, # 'ß' - 224: 89, # 'à' - 225: 70, # 'á' - 226: 59, # 'â' - 227: 78, # 'ã' - 228: 71, # 'ä' - 229: 82, # 'å' - 230: 88, # 'æ' - 231: 33, # 'ç' - 232: 77, # 'è' - 233: 66, # 'é' - 234: 84, # 'ê' - 235: 83, # 'ë' - 236: 110, # 'ì' - 237: 75, # 'í' - 238: 61, # 'î' - 239: 96, # 'ï' - 240: 30, # 'ğ' - 241: 67, # 'ñ' - 242: 109, # 'ò' - 243: 74, # 'ó' - 244: 87, # 'ô' - 245: 102, # 'õ' - 246: 34, # 'ö' - 247: 95, # '÷' - 248: 81, # 'ø' - 249: 108, # 'ù' - 250: 76, # 'ú' - 251: 72, # 'û' - 252: 17, # 'ü' - 253: 6, # 'ı' - 254: 19, # 'ş' - 255: 107, # 'ÿ' -} - -ISO_8859_9_TURKISH_MODEL = SingleByteCharSetModel( - charset_name="ISO-8859-9", - language="Turkish", - char_to_order_map=ISO_8859_9_TURKISH_CHAR_TO_ORDER, - language_model=TURKISH_LANG_MODEL, - typical_positive_ratio=0.97029, - keep_ascii_letters=True, - alphabet="ABCDEFGHIJKLMNOPRSTUVYZabcdefghijklmnoprstuvyzÂÇÎÖÛÜâçîöûüĞğİıŞş", -) diff --git a/venv/lib/python3.10/site-packages/chardet/latin1prober.py b/venv/lib/python3.10/site-packages/chardet/latin1prober.py deleted file mode 100644 index 59a01d91b87d4282bede38ade7cc78c0f7552d0e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/latin1prober.py +++ /dev/null @@ -1,147 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Universal charset detector code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 2001 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# Shy Shalom - original C code -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from typing import List, Union - -from .charsetprober import CharSetProber -from .enums import ProbingState - -FREQ_CAT_NUM = 4 - -UDF = 0 # undefined -OTH = 1 # other -ASC = 2 # ascii capital letter -ASS = 3 # ascii small letter -ACV = 4 # accent capital vowel -ACO = 5 # accent capital other -ASV = 6 # accent small vowel -ASO = 7 # accent small other -CLASS_NUM = 8 # total classes - -# fmt: off -Latin1_CharToClass = ( - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 00 - 07 - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 08 - 0F - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 10 - 17 - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 18 - 1F - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 20 - 27 - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 28 - 2F - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 30 - 37 - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 38 - 3F - OTH, ASC, ASC, ASC, ASC, ASC, ASC, ASC, # 40 - 47 - ASC, ASC, ASC, ASC, ASC, ASC, ASC, ASC, # 48 - 4F - ASC, ASC, ASC, ASC, ASC, ASC, ASC, ASC, # 50 - 57 - ASC, ASC, ASC, OTH, OTH, OTH, OTH, OTH, # 58 - 5F - OTH, ASS, ASS, ASS, ASS, ASS, ASS, ASS, # 60 - 67 - ASS, ASS, ASS, ASS, ASS, ASS, ASS, ASS, # 68 - 6F - ASS, ASS, ASS, ASS, ASS, ASS, ASS, ASS, # 70 - 77 - ASS, ASS, ASS, OTH, OTH, OTH, OTH, OTH, # 78 - 7F - OTH, UDF, OTH, ASO, OTH, OTH, OTH, OTH, # 80 - 87 - OTH, OTH, ACO, OTH, ACO, UDF, ACO, UDF, # 88 - 8F - UDF, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 90 - 97 - OTH, OTH, ASO, OTH, ASO, UDF, ASO, ACO, # 98 - 9F - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # A0 - A7 - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # A8 - AF - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # B0 - B7 - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # B8 - BF - ACV, ACV, ACV, ACV, ACV, ACV, ACO, ACO, # C0 - C7 - ACV, ACV, ACV, ACV, ACV, ACV, ACV, ACV, # C8 - CF - ACO, ACO, ACV, ACV, ACV, ACV, ACV, OTH, # D0 - D7 - ACV, ACV, ACV, ACV, ACV, ACO, ACO, ACO, # D8 - DF - ASV, ASV, ASV, ASV, ASV, ASV, ASO, ASO, # E0 - E7 - ASV, ASV, ASV, ASV, ASV, ASV, ASV, ASV, # E8 - EF - ASO, ASO, ASV, ASV, ASV, ASV, ASV, OTH, # F0 - F7 - ASV, ASV, ASV, ASV, ASV, ASO, ASO, ASO, # F8 - FF -) - -# 0 : illegal -# 1 : very unlikely -# 2 : normal -# 3 : very likely -Latin1ClassModel = ( -# UDF OTH ASC ASS ACV ACO ASV ASO - 0, 0, 0, 0, 0, 0, 0, 0, # UDF - 0, 3, 3, 3, 3, 3, 3, 3, # OTH - 0, 3, 3, 3, 3, 3, 3, 3, # ASC - 0, 3, 3, 3, 1, 1, 3, 3, # ASS - 0, 3, 3, 3, 1, 2, 1, 2, # ACV - 0, 3, 3, 3, 3, 3, 3, 3, # ACO - 0, 3, 1, 3, 1, 1, 1, 3, # ASV - 0, 3, 1, 3, 1, 1, 3, 3, # ASO -) -# fmt: on - - -class Latin1Prober(CharSetProber): - def __init__(self) -> None: - super().__init__() - self._last_char_class = OTH - self._freq_counter: List[int] = [] - self.reset() - - def reset(self) -> None: - self._last_char_class = OTH - self._freq_counter = [0] * FREQ_CAT_NUM - super().reset() - - @property - def charset_name(self) -> str: - return "ISO-8859-1" - - @property - def language(self) -> str: - return "" - - def feed(self, byte_str: Union[bytes, bytearray]) -> ProbingState: - byte_str = self.remove_xml_tags(byte_str) - for c in byte_str: - char_class = Latin1_CharToClass[c] - freq = Latin1ClassModel[(self._last_char_class * CLASS_NUM) + char_class] - if freq == 0: - self._state = ProbingState.NOT_ME - break - self._freq_counter[freq] += 1 - self._last_char_class = char_class - - return self.state - - def get_confidence(self) -> float: - if self.state == ProbingState.NOT_ME: - return 0.01 - - total = sum(self._freq_counter) - confidence = ( - 0.0 - if total < 0.01 - else (self._freq_counter[3] - self._freq_counter[1] * 20.0) / total - ) - confidence = max(confidence, 0.0) - # lower the confidence of latin1 so that other more accurate - # detector can take priority. - confidence *= 0.73 - return confidence diff --git a/venv/lib/python3.10/site-packages/chardet/macromanprober.py b/venv/lib/python3.10/site-packages/chardet/macromanprober.py deleted file mode 100644 index 1425d10ecaa59a9e49b73cea2b8b4747de73f6b5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/macromanprober.py +++ /dev/null @@ -1,162 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# This code was modified from latin1prober.py by Rob Speer . -# The Original Code is Mozilla Universal charset detector code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 2001 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Rob Speer - adapt to MacRoman encoding -# Mark Pilgrim - port to Python -# Shy Shalom - original C code -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from typing import List, Union - -from .charsetprober import CharSetProber -from .enums import ProbingState - -FREQ_CAT_NUM = 4 - -UDF = 0 # undefined -OTH = 1 # other -ASC = 2 # ascii capital letter -ASS = 3 # ascii small letter -ACV = 4 # accent capital vowel -ACO = 5 # accent capital other -ASV = 6 # accent small vowel -ASO = 7 # accent small other -ODD = 8 # character that is unlikely to appear -CLASS_NUM = 9 # total classes - -# The change from Latin1 is that we explicitly look for extended characters -# that are infrequently-occurring symbols, and consider them to always be -# improbable. This should let MacRoman get out of the way of more likely -# encodings in most situations. - -# fmt: off -MacRoman_CharToClass = ( - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 00 - 07 - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 08 - 0F - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 10 - 17 - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 18 - 1F - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 20 - 27 - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 28 - 2F - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 30 - 37 - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # 38 - 3F - OTH, ASC, ASC, ASC, ASC, ASC, ASC, ASC, # 40 - 47 - ASC, ASC, ASC, ASC, ASC, ASC, ASC, ASC, # 48 - 4F - ASC, ASC, ASC, ASC, ASC, ASC, ASC, ASC, # 50 - 57 - ASC, ASC, ASC, OTH, OTH, OTH, OTH, OTH, # 58 - 5F - OTH, ASS, ASS, ASS, ASS, ASS, ASS, ASS, # 60 - 67 - ASS, ASS, ASS, ASS, ASS, ASS, ASS, ASS, # 68 - 6F - ASS, ASS, ASS, ASS, ASS, ASS, ASS, ASS, # 70 - 77 - ASS, ASS, ASS, OTH, OTH, OTH, OTH, OTH, # 78 - 7F - ACV, ACV, ACO, ACV, ACO, ACV, ACV, ASV, # 80 - 87 - ASV, ASV, ASV, ASV, ASV, ASO, ASV, ASV, # 88 - 8F - ASV, ASV, ASV, ASV, ASV, ASV, ASO, ASV, # 90 - 97 - ASV, ASV, ASV, ASV, ASV, ASV, ASV, ASV, # 98 - 9F - OTH, OTH, OTH, OTH, OTH, OTH, OTH, ASO, # A0 - A7 - OTH, OTH, ODD, ODD, OTH, OTH, ACV, ACV, # A8 - AF - OTH, OTH, OTH, OTH, OTH, OTH, OTH, OTH, # B0 - B7 - OTH, OTH, OTH, OTH, OTH, OTH, ASV, ASV, # B8 - BF - OTH, OTH, ODD, OTH, ODD, OTH, OTH, OTH, # C0 - C7 - OTH, OTH, OTH, ACV, ACV, ACV, ACV, ASV, # C8 - CF - OTH, OTH, OTH, OTH, OTH, OTH, OTH, ODD, # D0 - D7 - ASV, ACV, ODD, OTH, OTH, OTH, OTH, OTH, # D8 - DF - OTH, OTH, OTH, OTH, OTH, ACV, ACV, ACV, # E0 - E7 - ACV, ACV, ACV, ACV, ACV, ACV, ACV, ACV, # E8 - EF - ODD, ACV, ACV, ACV, ACV, ASV, ODD, ODD, # F0 - F7 - ODD, ODD, ODD, ODD, ODD, ODD, ODD, ODD, # F8 - FF -) - -# 0 : illegal -# 1 : very unlikely -# 2 : normal -# 3 : very likely -MacRomanClassModel = ( -# UDF OTH ASC ASS ACV ACO ASV ASO ODD - 0, 0, 0, 0, 0, 0, 0, 0, 0, # UDF - 0, 3, 3, 3, 3, 3, 3, 3, 1, # OTH - 0, 3, 3, 3, 3, 3, 3, 3, 1, # ASC - 0, 3, 3, 3, 1, 1, 3, 3, 1, # ASS - 0, 3, 3, 3, 1, 2, 1, 2, 1, # ACV - 0, 3, 3, 3, 3, 3, 3, 3, 1, # ACO - 0, 3, 1, 3, 1, 1, 1, 3, 1, # ASV - 0, 3, 1, 3, 1, 1, 3, 3, 1, # ASO - 0, 1, 1, 1, 1, 1, 1, 1, 1, # ODD -) -# fmt: on - - -class MacRomanProber(CharSetProber): - def __init__(self) -> None: - super().__init__() - self._last_char_class = OTH - self._freq_counter: List[int] = [] - self.reset() - - def reset(self) -> None: - self._last_char_class = OTH - self._freq_counter = [0] * FREQ_CAT_NUM - - # express the prior that MacRoman is a somewhat rare encoding; - # this can be done by starting out in a slightly improbable state - # that must be overcome - self._freq_counter[2] = 10 - - super().reset() - - @property - def charset_name(self) -> str: - return "MacRoman" - - @property - def language(self) -> str: - return "" - - def feed(self, byte_str: Union[bytes, bytearray]) -> ProbingState: - byte_str = self.remove_xml_tags(byte_str) - for c in byte_str: - char_class = MacRoman_CharToClass[c] - freq = MacRomanClassModel[(self._last_char_class * CLASS_NUM) + char_class] - if freq == 0: - self._state = ProbingState.NOT_ME - break - self._freq_counter[freq] += 1 - self._last_char_class = char_class - - return self.state - - def get_confidence(self) -> float: - if self.state == ProbingState.NOT_ME: - return 0.01 - - total = sum(self._freq_counter) - confidence = ( - 0.0 - if total < 0.01 - else (self._freq_counter[3] - self._freq_counter[1] * 20.0) / total - ) - confidence = max(confidence, 0.0) - # lower the confidence of MacRoman so that other more accurate - # detector can take priority. - confidence *= 0.73 - return confidence diff --git a/venv/lib/python3.10/site-packages/chardet/mbcharsetprober.py b/venv/lib/python3.10/site-packages/chardet/mbcharsetprober.py deleted file mode 100644 index 666307e8fe0608c69f2b6578a49794e1e20a139a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/mbcharsetprober.py +++ /dev/null @@ -1,95 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Universal charset detector code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 2001 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# Shy Shalom - original C code -# Proofpoint, Inc. -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from typing import Optional, Union - -from .chardistribution import CharDistributionAnalysis -from .charsetprober import CharSetProber -from .codingstatemachine import CodingStateMachine -from .enums import LanguageFilter, MachineState, ProbingState - - -class MultiByteCharSetProber(CharSetProber): - """ - MultiByteCharSetProber - """ - - def __init__(self, lang_filter: LanguageFilter = LanguageFilter.NONE) -> None: - super().__init__(lang_filter=lang_filter) - self.distribution_analyzer: Optional[CharDistributionAnalysis] = None - self.coding_sm: Optional[CodingStateMachine] = None - self._last_char = bytearray(b"\0\0") - - def reset(self) -> None: - super().reset() - if self.coding_sm: - self.coding_sm.reset() - if self.distribution_analyzer: - self.distribution_analyzer.reset() - self._last_char = bytearray(b"\0\0") - - def feed(self, byte_str: Union[bytes, bytearray]) -> ProbingState: - assert self.coding_sm is not None - assert self.distribution_analyzer is not None - - for i, byte in enumerate(byte_str): - coding_state = self.coding_sm.next_state(byte) - if coding_state == MachineState.ERROR: - self.logger.debug( - "%s %s prober hit error at byte %s", - self.charset_name, - self.language, - i, - ) - self._state = ProbingState.NOT_ME - break - if coding_state == MachineState.ITS_ME: - self._state = ProbingState.FOUND_IT - break - if coding_state == MachineState.START: - char_len = self.coding_sm.get_current_charlen() - if i == 0: - self._last_char[1] = byte - self.distribution_analyzer.feed(self._last_char, char_len) - else: - self.distribution_analyzer.feed(byte_str[i - 1 : i + 1], char_len) - - self._last_char[0] = byte_str[-1] - - if self.state == ProbingState.DETECTING: - if self.distribution_analyzer.got_enough_data() and ( - self.get_confidence() > self.SHORTCUT_THRESHOLD - ): - self._state = ProbingState.FOUND_IT - - return self.state - - def get_confidence(self) -> float: - assert self.distribution_analyzer is not None - return self.distribution_analyzer.get_confidence() diff --git a/venv/lib/python3.10/site-packages/chardet/mbcsgroupprober.py b/venv/lib/python3.10/site-packages/chardet/mbcsgroupprober.py deleted file mode 100644 index 6cb9cc7b3bc751fbb5a54ba06eaaf953bf14ed8d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/mbcsgroupprober.py +++ /dev/null @@ -1,57 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Universal charset detector code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 2001 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# Shy Shalom - original C code -# Proofpoint, Inc. -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from .big5prober import Big5Prober -from .charsetgroupprober import CharSetGroupProber -from .cp949prober import CP949Prober -from .enums import LanguageFilter -from .eucjpprober import EUCJPProber -from .euckrprober import EUCKRProber -from .euctwprober import EUCTWProber -from .gb2312prober import GB2312Prober -from .johabprober import JOHABProber -from .sjisprober import SJISProber -from .utf8prober import UTF8Prober - - -class MBCSGroupProber(CharSetGroupProber): - def __init__(self, lang_filter: LanguageFilter = LanguageFilter.NONE) -> None: - super().__init__(lang_filter=lang_filter) - self.probers = [ - UTF8Prober(), - SJISProber(), - EUCJPProber(), - GB2312Prober(), - EUCKRProber(), - CP949Prober(), - Big5Prober(), - EUCTWProber(), - JOHABProber(), - ] - self.reset() diff --git a/venv/lib/python3.10/site-packages/chardet/mbcssm.py b/venv/lib/python3.10/site-packages/chardet/mbcssm.py deleted file mode 100644 index 7bbe97e6665356327814e2b797ffcc5724974a46..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/mbcssm.py +++ /dev/null @@ -1,661 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is mozilla.org code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from .codingstatemachinedict import CodingStateMachineDict -from .enums import MachineState - -# BIG5 - -# fmt: off -BIG5_CLS = ( - 1, 1, 1, 1, 1, 1, 1, 1, # 00 - 07 #allow 0x00 as legal value - 1, 1, 1, 1, 1, 1, 0, 0, # 08 - 0f - 1, 1, 1, 1, 1, 1, 1, 1, # 10 - 17 - 1, 1, 1, 0, 1, 1, 1, 1, # 18 - 1f - 1, 1, 1, 1, 1, 1, 1, 1, # 20 - 27 - 1, 1, 1, 1, 1, 1, 1, 1, # 28 - 2f - 1, 1, 1, 1, 1, 1, 1, 1, # 30 - 37 - 1, 1, 1, 1, 1, 1, 1, 1, # 38 - 3f - 2, 2, 2, 2, 2, 2, 2, 2, # 40 - 47 - 2, 2, 2, 2, 2, 2, 2, 2, # 48 - 4f - 2, 2, 2, 2, 2, 2, 2, 2, # 50 - 57 - 2, 2, 2, 2, 2, 2, 2, 2, # 58 - 5f - 2, 2, 2, 2, 2, 2, 2, 2, # 60 - 67 - 2, 2, 2, 2, 2, 2, 2, 2, # 68 - 6f - 2, 2, 2, 2, 2, 2, 2, 2, # 70 - 77 - 2, 2, 2, 2, 2, 2, 2, 1, # 78 - 7f - 4, 4, 4, 4, 4, 4, 4, 4, # 80 - 87 - 4, 4, 4, 4, 4, 4, 4, 4, # 88 - 8f - 4, 4, 4, 4, 4, 4, 4, 4, # 90 - 97 - 4, 4, 4, 4, 4, 4, 4, 4, # 98 - 9f - 4, 3, 3, 3, 3, 3, 3, 3, # a0 - a7 - 3, 3, 3, 3, 3, 3, 3, 3, # a8 - af - 3, 3, 3, 3, 3, 3, 3, 3, # b0 - b7 - 3, 3, 3, 3, 3, 3, 3, 3, # b8 - bf - 3, 3, 3, 3, 3, 3, 3, 3, # c0 - c7 - 3, 3, 3, 3, 3, 3, 3, 3, # c8 - cf - 3, 3, 3, 3, 3, 3, 3, 3, # d0 - d7 - 3, 3, 3, 3, 3, 3, 3, 3, # d8 - df - 3, 3, 3, 3, 3, 3, 3, 3, # e0 - e7 - 3, 3, 3, 3, 3, 3, 3, 3, # e8 - ef - 3, 3, 3, 3, 3, 3, 3, 3, # f0 - f7 - 3, 3, 3, 3, 3, 3, 3, 0 # f8 - ff -) - -BIG5_ST = ( - MachineState.ERROR,MachineState.START,MachineState.START, 3,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#00-07 - MachineState.ERROR,MachineState.ERROR,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ERROR,#08-0f - MachineState.ERROR,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START#10-17 -) -# fmt: on - -BIG5_CHAR_LEN_TABLE = (0, 1, 1, 2, 0) - -BIG5_SM_MODEL: CodingStateMachineDict = { - "class_table": BIG5_CLS, - "class_factor": 5, - "state_table": BIG5_ST, - "char_len_table": BIG5_CHAR_LEN_TABLE, - "name": "Big5", -} - -# CP949 -# fmt: off -CP949_CLS = ( - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, # 00 - 0f - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, # 10 - 1f - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, # 20 - 2f - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, # 30 - 3f - 1, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, # 40 - 4f - 4, 4, 5, 5, 5, 5, 5, 5, 5, 5, 5, 1, 1, 1, 1, 1, # 50 - 5f - 1, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, # 60 - 6f - 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 1, 1, 1, 1, 1, # 70 - 7f - 0, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, # 80 - 8f - 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, # 90 - 9f - 6, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 8, 8, 8, # a0 - af - 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, # b0 - bf - 7, 7, 7, 7, 7, 7, 9, 2, 2, 3, 2, 2, 2, 2, 2, 2, # c0 - cf - 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, # d0 - df - 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, # e0 - ef - 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 0, # f0 - ff -) - -CP949_ST = ( -#cls= 0 1 2 3 4 5 6 7 8 9 # previous state = - MachineState.ERROR,MachineState.START, 3,MachineState.ERROR,MachineState.START,MachineState.START, 4, 5,MachineState.ERROR, 6, # MachineState.START - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR, # MachineState.ERROR - MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME, # MachineState.ITS_ME - MachineState.ERROR,MachineState.ERROR,MachineState.START,MachineState.START,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.START,MachineState.START,MachineState.START, # 3 - MachineState.ERROR,MachineState.ERROR,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START, # 4 - MachineState.ERROR,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START, # 5 - MachineState.ERROR,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.ERROR,MachineState.ERROR,MachineState.START,MachineState.START,MachineState.START, # 6 -) -# fmt: on - -CP949_CHAR_LEN_TABLE = (0, 1, 2, 0, 1, 1, 2, 2, 0, 2) - -CP949_SM_MODEL: CodingStateMachineDict = { - "class_table": CP949_CLS, - "class_factor": 10, - "state_table": CP949_ST, - "char_len_table": CP949_CHAR_LEN_TABLE, - "name": "CP949", -} - -# EUC-JP -# fmt: off -EUCJP_CLS = ( - 4, 4, 4, 4, 4, 4, 4, 4, # 00 - 07 - 4, 4, 4, 4, 4, 4, 5, 5, # 08 - 0f - 4, 4, 4, 4, 4, 4, 4, 4, # 10 - 17 - 4, 4, 4, 5, 4, 4, 4, 4, # 18 - 1f - 4, 4, 4, 4, 4, 4, 4, 4, # 20 - 27 - 4, 4, 4, 4, 4, 4, 4, 4, # 28 - 2f - 4, 4, 4, 4, 4, 4, 4, 4, # 30 - 37 - 4, 4, 4, 4, 4, 4, 4, 4, # 38 - 3f - 4, 4, 4, 4, 4, 4, 4, 4, # 40 - 47 - 4, 4, 4, 4, 4, 4, 4, 4, # 48 - 4f - 4, 4, 4, 4, 4, 4, 4, 4, # 50 - 57 - 4, 4, 4, 4, 4, 4, 4, 4, # 58 - 5f - 4, 4, 4, 4, 4, 4, 4, 4, # 60 - 67 - 4, 4, 4, 4, 4, 4, 4, 4, # 68 - 6f - 4, 4, 4, 4, 4, 4, 4, 4, # 70 - 77 - 4, 4, 4, 4, 4, 4, 4, 4, # 78 - 7f - 5, 5, 5, 5, 5, 5, 5, 5, # 80 - 87 - 5, 5, 5, 5, 5, 5, 1, 3, # 88 - 8f - 5, 5, 5, 5, 5, 5, 5, 5, # 90 - 97 - 5, 5, 5, 5, 5, 5, 5, 5, # 98 - 9f - 5, 2, 2, 2, 2, 2, 2, 2, # a0 - a7 - 2, 2, 2, 2, 2, 2, 2, 2, # a8 - af - 2, 2, 2, 2, 2, 2, 2, 2, # b0 - b7 - 2, 2, 2, 2, 2, 2, 2, 2, # b8 - bf - 2, 2, 2, 2, 2, 2, 2, 2, # c0 - c7 - 2, 2, 2, 2, 2, 2, 2, 2, # c8 - cf - 2, 2, 2, 2, 2, 2, 2, 2, # d0 - d7 - 2, 2, 2, 2, 2, 2, 2, 2, # d8 - df - 0, 0, 0, 0, 0, 0, 0, 0, # e0 - e7 - 0, 0, 0, 0, 0, 0, 0, 0, # e8 - ef - 0, 0, 0, 0, 0, 0, 0, 0, # f0 - f7 - 0, 0, 0, 0, 0, 0, 0, 5 # f8 - ff -) - -EUCJP_ST = ( - 3, 4, 3, 5,MachineState.START,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#00-07 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,#08-0f - MachineState.ITS_ME,MachineState.ITS_ME,MachineState.START,MachineState.ERROR,MachineState.START,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#10-17 - MachineState.ERROR,MachineState.ERROR,MachineState.START,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR, 3,MachineState.ERROR,#18-1f - 3,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.START,MachineState.START,MachineState.START,MachineState.START#20-27 -) -# fmt: on - -EUCJP_CHAR_LEN_TABLE = (2, 2, 2, 3, 1, 0) - -EUCJP_SM_MODEL: CodingStateMachineDict = { - "class_table": EUCJP_CLS, - "class_factor": 6, - "state_table": EUCJP_ST, - "char_len_table": EUCJP_CHAR_LEN_TABLE, - "name": "EUC-JP", -} - -# EUC-KR -# fmt: off -EUCKR_CLS = ( - 1, 1, 1, 1, 1, 1, 1, 1, # 00 - 07 - 1, 1, 1, 1, 1, 1, 0, 0, # 08 - 0f - 1, 1, 1, 1, 1, 1, 1, 1, # 10 - 17 - 1, 1, 1, 0, 1, 1, 1, 1, # 18 - 1f - 1, 1, 1, 1, 1, 1, 1, 1, # 20 - 27 - 1, 1, 1, 1, 1, 1, 1, 1, # 28 - 2f - 1, 1, 1, 1, 1, 1, 1, 1, # 30 - 37 - 1, 1, 1, 1, 1, 1, 1, 1, # 38 - 3f - 1, 1, 1, 1, 1, 1, 1, 1, # 40 - 47 - 1, 1, 1, 1, 1, 1, 1, 1, # 48 - 4f - 1, 1, 1, 1, 1, 1, 1, 1, # 50 - 57 - 1, 1, 1, 1, 1, 1, 1, 1, # 58 - 5f - 1, 1, 1, 1, 1, 1, 1, 1, # 60 - 67 - 1, 1, 1, 1, 1, 1, 1, 1, # 68 - 6f - 1, 1, 1, 1, 1, 1, 1, 1, # 70 - 77 - 1, 1, 1, 1, 1, 1, 1, 1, # 78 - 7f - 0, 0, 0, 0, 0, 0, 0, 0, # 80 - 87 - 0, 0, 0, 0, 0, 0, 0, 0, # 88 - 8f - 0, 0, 0, 0, 0, 0, 0, 0, # 90 - 97 - 0, 0, 0, 0, 0, 0, 0, 0, # 98 - 9f - 0, 2, 2, 2, 2, 2, 2, 2, # a0 - a7 - 2, 2, 2, 2, 2, 3, 3, 3, # a8 - af - 2, 2, 2, 2, 2, 2, 2, 2, # b0 - b7 - 2, 2, 2, 2, 2, 2, 2, 2, # b8 - bf - 2, 2, 2, 2, 2, 2, 2, 2, # c0 - c7 - 2, 3, 2, 2, 2, 2, 2, 2, # c8 - cf - 2, 2, 2, 2, 2, 2, 2, 2, # d0 - d7 - 2, 2, 2, 2, 2, 2, 2, 2, # d8 - df - 2, 2, 2, 2, 2, 2, 2, 2, # e0 - e7 - 2, 2, 2, 2, 2, 2, 2, 2, # e8 - ef - 2, 2, 2, 2, 2, 2, 2, 2, # f0 - f7 - 2, 2, 2, 2, 2, 2, 2, 0 # f8 - ff -) - -EUCKR_ST = ( - MachineState.ERROR,MachineState.START, 3,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#00-07 - MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ERROR,MachineState.ERROR,MachineState.START,MachineState.START #08-0f -) -# fmt: on - -EUCKR_CHAR_LEN_TABLE = (0, 1, 2, 0) - -EUCKR_SM_MODEL: CodingStateMachineDict = { - "class_table": EUCKR_CLS, - "class_factor": 4, - "state_table": EUCKR_ST, - "char_len_table": EUCKR_CHAR_LEN_TABLE, - "name": "EUC-KR", -} - -# JOHAB -# fmt: off -JOHAB_CLS = ( - 4,4,4,4,4,4,4,4, # 00 - 07 - 4,4,4,4,4,4,0,0, # 08 - 0f - 4,4,4,4,4,4,4,4, # 10 - 17 - 4,4,4,0,4,4,4,4, # 18 - 1f - 4,4,4,4,4,4,4,4, # 20 - 27 - 4,4,4,4,4,4,4,4, # 28 - 2f - 4,3,3,3,3,3,3,3, # 30 - 37 - 3,3,3,3,3,3,3,3, # 38 - 3f - 3,1,1,1,1,1,1,1, # 40 - 47 - 1,1,1,1,1,1,1,1, # 48 - 4f - 1,1,1,1,1,1,1,1, # 50 - 57 - 1,1,1,1,1,1,1,1, # 58 - 5f - 1,1,1,1,1,1,1,1, # 60 - 67 - 1,1,1,1,1,1,1,1, # 68 - 6f - 1,1,1,1,1,1,1,1, # 70 - 77 - 1,1,1,1,1,1,1,2, # 78 - 7f - 6,6,6,6,8,8,8,8, # 80 - 87 - 8,8,8,8,8,8,8,8, # 88 - 8f - 8,7,7,7,7,7,7,7, # 90 - 97 - 7,7,7,7,7,7,7,7, # 98 - 9f - 7,7,7,7,7,7,7,7, # a0 - a7 - 7,7,7,7,7,7,7,7, # a8 - af - 7,7,7,7,7,7,7,7, # b0 - b7 - 7,7,7,7,7,7,7,7, # b8 - bf - 7,7,7,7,7,7,7,7, # c0 - c7 - 7,7,7,7,7,7,7,7, # c8 - cf - 7,7,7,7,5,5,5,5, # d0 - d7 - 5,9,9,9,9,9,9,5, # d8 - df - 9,9,9,9,9,9,9,9, # e0 - e7 - 9,9,9,9,9,9,9,9, # e8 - ef - 9,9,9,9,9,9,9,9, # f0 - f7 - 9,9,5,5,5,5,5,0 # f8 - ff -) - -JOHAB_ST = ( -# cls = 0 1 2 3 4 5 6 7 8 9 - MachineState.ERROR ,MachineState.START ,MachineState.START ,MachineState.START ,MachineState.START ,MachineState.ERROR ,MachineState.ERROR ,3 ,3 ,4 , # MachineState.START - MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME, # MachineState.ITS_ME - MachineState.ERROR ,MachineState.ERROR ,MachineState.ERROR ,MachineState.ERROR ,MachineState.ERROR ,MachineState.ERROR ,MachineState.ERROR ,MachineState.ERROR ,MachineState.ERROR ,MachineState.ERROR , # MachineState.ERROR - MachineState.ERROR ,MachineState.START ,MachineState.START ,MachineState.ERROR ,MachineState.ERROR ,MachineState.START ,MachineState.START ,MachineState.START ,MachineState.START ,MachineState.START , # 3 - MachineState.ERROR ,MachineState.START ,MachineState.ERROR ,MachineState.START ,MachineState.ERROR ,MachineState.START ,MachineState.ERROR ,MachineState.START ,MachineState.ERROR ,MachineState.START , # 4 -) -# fmt: on - -JOHAB_CHAR_LEN_TABLE = (0, 1, 1, 1, 1, 0, 0, 2, 2, 2) - -JOHAB_SM_MODEL: CodingStateMachineDict = { - "class_table": JOHAB_CLS, - "class_factor": 10, - "state_table": JOHAB_ST, - "char_len_table": JOHAB_CHAR_LEN_TABLE, - "name": "Johab", -} - -# EUC-TW -# fmt: off -EUCTW_CLS = ( - 2, 2, 2, 2, 2, 2, 2, 2, # 00 - 07 - 2, 2, 2, 2, 2, 2, 0, 0, # 08 - 0f - 2, 2, 2, 2, 2, 2, 2, 2, # 10 - 17 - 2, 2, 2, 0, 2, 2, 2, 2, # 18 - 1f - 2, 2, 2, 2, 2, 2, 2, 2, # 20 - 27 - 2, 2, 2, 2, 2, 2, 2, 2, # 28 - 2f - 2, 2, 2, 2, 2, 2, 2, 2, # 30 - 37 - 2, 2, 2, 2, 2, 2, 2, 2, # 38 - 3f - 2, 2, 2, 2, 2, 2, 2, 2, # 40 - 47 - 2, 2, 2, 2, 2, 2, 2, 2, # 48 - 4f - 2, 2, 2, 2, 2, 2, 2, 2, # 50 - 57 - 2, 2, 2, 2, 2, 2, 2, 2, # 58 - 5f - 2, 2, 2, 2, 2, 2, 2, 2, # 60 - 67 - 2, 2, 2, 2, 2, 2, 2, 2, # 68 - 6f - 2, 2, 2, 2, 2, 2, 2, 2, # 70 - 77 - 2, 2, 2, 2, 2, 2, 2, 2, # 78 - 7f - 0, 0, 0, 0, 0, 0, 0, 0, # 80 - 87 - 0, 0, 0, 0, 0, 0, 6, 0, # 88 - 8f - 0, 0, 0, 0, 0, 0, 0, 0, # 90 - 97 - 0, 0, 0, 0, 0, 0, 0, 0, # 98 - 9f - 0, 3, 4, 4, 4, 4, 4, 4, # a0 - a7 - 5, 5, 1, 1, 1, 1, 1, 1, # a8 - af - 1, 1, 1, 1, 1, 1, 1, 1, # b0 - b7 - 1, 1, 1, 1, 1, 1, 1, 1, # b8 - bf - 1, 1, 3, 1, 3, 3, 3, 3, # c0 - c7 - 3, 3, 3, 3, 3, 3, 3, 3, # c8 - cf - 3, 3, 3, 3, 3, 3, 3, 3, # d0 - d7 - 3, 3, 3, 3, 3, 3, 3, 3, # d8 - df - 3, 3, 3, 3, 3, 3, 3, 3, # e0 - e7 - 3, 3, 3, 3, 3, 3, 3, 3, # e8 - ef - 3, 3, 3, 3, 3, 3, 3, 3, # f0 - f7 - 3, 3, 3, 3, 3, 3, 3, 0 # f8 - ff -) - -EUCTW_ST = ( - MachineState.ERROR,MachineState.ERROR,MachineState.START, 3, 3, 3, 4,MachineState.ERROR,#00-07 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ITS_ME,MachineState.ITS_ME,#08-0f - MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ERROR,MachineState.START,MachineState.ERROR,#10-17 - MachineState.START,MachineState.START,MachineState.START,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#18-1f - 5,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.START,MachineState.ERROR,MachineState.START,MachineState.START,#20-27 - MachineState.START,MachineState.ERROR,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START #28-2f -) -# fmt: on - -EUCTW_CHAR_LEN_TABLE = (0, 0, 1, 2, 2, 2, 3) - -EUCTW_SM_MODEL: CodingStateMachineDict = { - "class_table": EUCTW_CLS, - "class_factor": 7, - "state_table": EUCTW_ST, - "char_len_table": EUCTW_CHAR_LEN_TABLE, - "name": "x-euc-tw", -} - -# GB2312 -# fmt: off -GB2312_CLS = ( - 1, 1, 1, 1, 1, 1, 1, 1, # 00 - 07 - 1, 1, 1, 1, 1, 1, 0, 0, # 08 - 0f - 1, 1, 1, 1, 1, 1, 1, 1, # 10 - 17 - 1, 1, 1, 0, 1, 1, 1, 1, # 18 - 1f - 1, 1, 1, 1, 1, 1, 1, 1, # 20 - 27 - 1, 1, 1, 1, 1, 1, 1, 1, # 28 - 2f - 3, 3, 3, 3, 3, 3, 3, 3, # 30 - 37 - 3, 3, 1, 1, 1, 1, 1, 1, # 38 - 3f - 2, 2, 2, 2, 2, 2, 2, 2, # 40 - 47 - 2, 2, 2, 2, 2, 2, 2, 2, # 48 - 4f - 2, 2, 2, 2, 2, 2, 2, 2, # 50 - 57 - 2, 2, 2, 2, 2, 2, 2, 2, # 58 - 5f - 2, 2, 2, 2, 2, 2, 2, 2, # 60 - 67 - 2, 2, 2, 2, 2, 2, 2, 2, # 68 - 6f - 2, 2, 2, 2, 2, 2, 2, 2, # 70 - 77 - 2, 2, 2, 2, 2, 2, 2, 4, # 78 - 7f - 5, 6, 6, 6, 6, 6, 6, 6, # 80 - 87 - 6, 6, 6, 6, 6, 6, 6, 6, # 88 - 8f - 6, 6, 6, 6, 6, 6, 6, 6, # 90 - 97 - 6, 6, 6, 6, 6, 6, 6, 6, # 98 - 9f - 6, 6, 6, 6, 6, 6, 6, 6, # a0 - a7 - 6, 6, 6, 6, 6, 6, 6, 6, # a8 - af - 6, 6, 6, 6, 6, 6, 6, 6, # b0 - b7 - 6, 6, 6, 6, 6, 6, 6, 6, # b8 - bf - 6, 6, 6, 6, 6, 6, 6, 6, # c0 - c7 - 6, 6, 6, 6, 6, 6, 6, 6, # c8 - cf - 6, 6, 6, 6, 6, 6, 6, 6, # d0 - d7 - 6, 6, 6, 6, 6, 6, 6, 6, # d8 - df - 6, 6, 6, 6, 6, 6, 6, 6, # e0 - e7 - 6, 6, 6, 6, 6, 6, 6, 6, # e8 - ef - 6, 6, 6, 6, 6, 6, 6, 6, # f0 - f7 - 6, 6, 6, 6, 6, 6, 6, 0 # f8 - ff -) - -GB2312_ST = ( - MachineState.ERROR,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START, 3,MachineState.ERROR,#00-07 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ITS_ME,MachineState.ITS_ME,#08-0f - MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ERROR,MachineState.ERROR,MachineState.START,#10-17 - 4,MachineState.ERROR,MachineState.START,MachineState.START,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#18-1f - MachineState.ERROR,MachineState.ERROR, 5,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ITS_ME,MachineState.ERROR,#20-27 - MachineState.ERROR,MachineState.ERROR,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.START #28-2f -) -# fmt: on - -# To be accurate, the length of class 6 can be either 2 or 4. -# But it is not necessary to discriminate between the two since -# it is used for frequency analysis only, and we are validating -# each code range there as well. So it is safe to set it to be -# 2 here. -GB2312_CHAR_LEN_TABLE = (0, 1, 1, 1, 1, 1, 2) - -GB2312_SM_MODEL: CodingStateMachineDict = { - "class_table": GB2312_CLS, - "class_factor": 7, - "state_table": GB2312_ST, - "char_len_table": GB2312_CHAR_LEN_TABLE, - "name": "GB2312", -} - -# Shift_JIS -# fmt: off -SJIS_CLS = ( - 1, 1, 1, 1, 1, 1, 1, 1, # 00 - 07 - 1, 1, 1, 1, 1, 1, 0, 0, # 08 - 0f - 1, 1, 1, 1, 1, 1, 1, 1, # 10 - 17 - 1, 1, 1, 0, 1, 1, 1, 1, # 18 - 1f - 1, 1, 1, 1, 1, 1, 1, 1, # 20 - 27 - 1, 1, 1, 1, 1, 1, 1, 1, # 28 - 2f - 1, 1, 1, 1, 1, 1, 1, 1, # 30 - 37 - 1, 1, 1, 1, 1, 1, 1, 1, # 38 - 3f - 2, 2, 2, 2, 2, 2, 2, 2, # 40 - 47 - 2, 2, 2, 2, 2, 2, 2, 2, # 48 - 4f - 2, 2, 2, 2, 2, 2, 2, 2, # 50 - 57 - 2, 2, 2, 2, 2, 2, 2, 2, # 58 - 5f - 2, 2, 2, 2, 2, 2, 2, 2, # 60 - 67 - 2, 2, 2, 2, 2, 2, 2, 2, # 68 - 6f - 2, 2, 2, 2, 2, 2, 2, 2, # 70 - 77 - 2, 2, 2, 2, 2, 2, 2, 1, # 78 - 7f - 3, 3, 3, 3, 3, 2, 2, 3, # 80 - 87 - 3, 3, 3, 3, 3, 3, 3, 3, # 88 - 8f - 3, 3, 3, 3, 3, 3, 3, 3, # 90 - 97 - 3, 3, 3, 3, 3, 3, 3, 3, # 98 - 9f - #0xa0 is illegal in sjis encoding, but some pages does - #contain such byte. We need to be more error forgiven. - 2, 2, 2, 2, 2, 2, 2, 2, # a0 - a7 - 2, 2, 2, 2, 2, 2, 2, 2, # a8 - af - 2, 2, 2, 2, 2, 2, 2, 2, # b0 - b7 - 2, 2, 2, 2, 2, 2, 2, 2, # b8 - bf - 2, 2, 2, 2, 2, 2, 2, 2, # c0 - c7 - 2, 2, 2, 2, 2, 2, 2, 2, # c8 - cf - 2, 2, 2, 2, 2, 2, 2, 2, # d0 - d7 - 2, 2, 2, 2, 2, 2, 2, 2, # d8 - df - 3, 3, 3, 3, 3, 3, 3, 3, # e0 - e7 - 3, 3, 3, 3, 3, 4, 4, 4, # e8 - ef - 3, 3, 3, 3, 3, 3, 3, 3, # f0 - f7 - 3, 3, 3, 3, 3, 0, 0, 0, # f8 - ff -) - -SJIS_ST = ( - MachineState.ERROR,MachineState.START,MachineState.START, 3,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#00-07 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,#08-0f - MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ERROR,MachineState.ERROR,MachineState.START,MachineState.START,MachineState.START,MachineState.START #10-17 -) -# fmt: on - -SJIS_CHAR_LEN_TABLE = (0, 1, 1, 2, 0, 0) - -SJIS_SM_MODEL: CodingStateMachineDict = { - "class_table": SJIS_CLS, - "class_factor": 6, - "state_table": SJIS_ST, - "char_len_table": SJIS_CHAR_LEN_TABLE, - "name": "Shift_JIS", -} - -# UCS2-BE -# fmt: off -UCS2BE_CLS = ( - 0, 0, 0, 0, 0, 0, 0, 0, # 00 - 07 - 0, 0, 1, 0, 0, 2, 0, 0, # 08 - 0f - 0, 0, 0, 0, 0, 0, 0, 0, # 10 - 17 - 0, 0, 0, 3, 0, 0, 0, 0, # 18 - 1f - 0, 0, 0, 0, 0, 0, 0, 0, # 20 - 27 - 0, 3, 3, 3, 3, 3, 0, 0, # 28 - 2f - 0, 0, 0, 0, 0, 0, 0, 0, # 30 - 37 - 0, 0, 0, 0, 0, 0, 0, 0, # 38 - 3f - 0, 0, 0, 0, 0, 0, 0, 0, # 40 - 47 - 0, 0, 0, 0, 0, 0, 0, 0, # 48 - 4f - 0, 0, 0, 0, 0, 0, 0, 0, # 50 - 57 - 0, 0, 0, 0, 0, 0, 0, 0, # 58 - 5f - 0, 0, 0, 0, 0, 0, 0, 0, # 60 - 67 - 0, 0, 0, 0, 0, 0, 0, 0, # 68 - 6f - 0, 0, 0, 0, 0, 0, 0, 0, # 70 - 77 - 0, 0, 0, 0, 0, 0, 0, 0, # 78 - 7f - 0, 0, 0, 0, 0, 0, 0, 0, # 80 - 87 - 0, 0, 0, 0, 0, 0, 0, 0, # 88 - 8f - 0, 0, 0, 0, 0, 0, 0, 0, # 90 - 97 - 0, 0, 0, 0, 0, 0, 0, 0, # 98 - 9f - 0, 0, 0, 0, 0, 0, 0, 0, # a0 - a7 - 0, 0, 0, 0, 0, 0, 0, 0, # a8 - af - 0, 0, 0, 0, 0, 0, 0, 0, # b0 - b7 - 0, 0, 0, 0, 0, 0, 0, 0, # b8 - bf - 0, 0, 0, 0, 0, 0, 0, 0, # c0 - c7 - 0, 0, 0, 0, 0, 0, 0, 0, # c8 - cf - 0, 0, 0, 0, 0, 0, 0, 0, # d0 - d7 - 0, 0, 0, 0, 0, 0, 0, 0, # d8 - df - 0, 0, 0, 0, 0, 0, 0, 0, # e0 - e7 - 0, 0, 0, 0, 0, 0, 0, 0, # e8 - ef - 0, 0, 0, 0, 0, 0, 0, 0, # f0 - f7 - 0, 0, 0, 0, 0, 0, 4, 5 # f8 - ff -) - -UCS2BE_ST = ( - 5, 7, 7,MachineState.ERROR, 4, 3,MachineState.ERROR,MachineState.ERROR,#00-07 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,#08-0f - MachineState.ITS_ME,MachineState.ITS_ME, 6, 6, 6, 6,MachineState.ERROR,MachineState.ERROR,#10-17 - 6, 6, 6, 6, 6,MachineState.ITS_ME, 6, 6,#18-1f - 6, 6, 6, 6, 5, 7, 7,MachineState.ERROR,#20-27 - 5, 8, 6, 6,MachineState.ERROR, 6, 6, 6,#28-2f - 6, 6, 6, 6,MachineState.ERROR,MachineState.ERROR,MachineState.START,MachineState.START #30-37 -) -# fmt: on - -UCS2BE_CHAR_LEN_TABLE = (2, 2, 2, 0, 2, 2) - -UCS2BE_SM_MODEL: CodingStateMachineDict = { - "class_table": UCS2BE_CLS, - "class_factor": 6, - "state_table": UCS2BE_ST, - "char_len_table": UCS2BE_CHAR_LEN_TABLE, - "name": "UTF-16BE", -} - -# UCS2-LE -# fmt: off -UCS2LE_CLS = ( - 0, 0, 0, 0, 0, 0, 0, 0, # 00 - 07 - 0, 0, 1, 0, 0, 2, 0, 0, # 08 - 0f - 0, 0, 0, 0, 0, 0, 0, 0, # 10 - 17 - 0, 0, 0, 3, 0, 0, 0, 0, # 18 - 1f - 0, 0, 0, 0, 0, 0, 0, 0, # 20 - 27 - 0, 3, 3, 3, 3, 3, 0, 0, # 28 - 2f - 0, 0, 0, 0, 0, 0, 0, 0, # 30 - 37 - 0, 0, 0, 0, 0, 0, 0, 0, # 38 - 3f - 0, 0, 0, 0, 0, 0, 0, 0, # 40 - 47 - 0, 0, 0, 0, 0, 0, 0, 0, # 48 - 4f - 0, 0, 0, 0, 0, 0, 0, 0, # 50 - 57 - 0, 0, 0, 0, 0, 0, 0, 0, # 58 - 5f - 0, 0, 0, 0, 0, 0, 0, 0, # 60 - 67 - 0, 0, 0, 0, 0, 0, 0, 0, # 68 - 6f - 0, 0, 0, 0, 0, 0, 0, 0, # 70 - 77 - 0, 0, 0, 0, 0, 0, 0, 0, # 78 - 7f - 0, 0, 0, 0, 0, 0, 0, 0, # 80 - 87 - 0, 0, 0, 0, 0, 0, 0, 0, # 88 - 8f - 0, 0, 0, 0, 0, 0, 0, 0, # 90 - 97 - 0, 0, 0, 0, 0, 0, 0, 0, # 98 - 9f - 0, 0, 0, 0, 0, 0, 0, 0, # a0 - a7 - 0, 0, 0, 0, 0, 0, 0, 0, # a8 - af - 0, 0, 0, 0, 0, 0, 0, 0, # b0 - b7 - 0, 0, 0, 0, 0, 0, 0, 0, # b8 - bf - 0, 0, 0, 0, 0, 0, 0, 0, # c0 - c7 - 0, 0, 0, 0, 0, 0, 0, 0, # c8 - cf - 0, 0, 0, 0, 0, 0, 0, 0, # d0 - d7 - 0, 0, 0, 0, 0, 0, 0, 0, # d8 - df - 0, 0, 0, 0, 0, 0, 0, 0, # e0 - e7 - 0, 0, 0, 0, 0, 0, 0, 0, # e8 - ef - 0, 0, 0, 0, 0, 0, 0, 0, # f0 - f7 - 0, 0, 0, 0, 0, 0, 4, 5 # f8 - ff -) - -UCS2LE_ST = ( - 6, 6, 7, 6, 4, 3,MachineState.ERROR,MachineState.ERROR,#00-07 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,#08-0f - MachineState.ITS_ME,MachineState.ITS_ME, 5, 5, 5,MachineState.ERROR,MachineState.ITS_ME,MachineState.ERROR,#10-17 - 5, 5, 5,MachineState.ERROR, 5,MachineState.ERROR, 6, 6,#18-1f - 7, 6, 8, 8, 5, 5, 5,MachineState.ERROR,#20-27 - 5, 5, 5,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR, 5, 5,#28-2f - 5, 5, 5,MachineState.ERROR, 5,MachineState.ERROR,MachineState.START,MachineState.START #30-37 -) -# fmt: on - -UCS2LE_CHAR_LEN_TABLE = (2, 2, 2, 2, 2, 2) - -UCS2LE_SM_MODEL: CodingStateMachineDict = { - "class_table": UCS2LE_CLS, - "class_factor": 6, - "state_table": UCS2LE_ST, - "char_len_table": UCS2LE_CHAR_LEN_TABLE, - "name": "UTF-16LE", -} - -# UTF-8 -# fmt: off -UTF8_CLS = ( - 1, 1, 1, 1, 1, 1, 1, 1, # 00 - 07 #allow 0x00 as a legal value - 1, 1, 1, 1, 1, 1, 0, 0, # 08 - 0f - 1, 1, 1, 1, 1, 1, 1, 1, # 10 - 17 - 1, 1, 1, 0, 1, 1, 1, 1, # 18 - 1f - 1, 1, 1, 1, 1, 1, 1, 1, # 20 - 27 - 1, 1, 1, 1, 1, 1, 1, 1, # 28 - 2f - 1, 1, 1, 1, 1, 1, 1, 1, # 30 - 37 - 1, 1, 1, 1, 1, 1, 1, 1, # 38 - 3f - 1, 1, 1, 1, 1, 1, 1, 1, # 40 - 47 - 1, 1, 1, 1, 1, 1, 1, 1, # 48 - 4f - 1, 1, 1, 1, 1, 1, 1, 1, # 50 - 57 - 1, 1, 1, 1, 1, 1, 1, 1, # 58 - 5f - 1, 1, 1, 1, 1, 1, 1, 1, # 60 - 67 - 1, 1, 1, 1, 1, 1, 1, 1, # 68 - 6f - 1, 1, 1, 1, 1, 1, 1, 1, # 70 - 77 - 1, 1, 1, 1, 1, 1, 1, 1, # 78 - 7f - 2, 2, 2, 2, 3, 3, 3, 3, # 80 - 87 - 4, 4, 4, 4, 4, 4, 4, 4, # 88 - 8f - 4, 4, 4, 4, 4, 4, 4, 4, # 90 - 97 - 4, 4, 4, 4, 4, 4, 4, 4, # 98 - 9f - 5, 5, 5, 5, 5, 5, 5, 5, # a0 - a7 - 5, 5, 5, 5, 5, 5, 5, 5, # a8 - af - 5, 5, 5, 5, 5, 5, 5, 5, # b0 - b7 - 5, 5, 5, 5, 5, 5, 5, 5, # b8 - bf - 0, 0, 6, 6, 6, 6, 6, 6, # c0 - c7 - 6, 6, 6, 6, 6, 6, 6, 6, # c8 - cf - 6, 6, 6, 6, 6, 6, 6, 6, # d0 - d7 - 6, 6, 6, 6, 6, 6, 6, 6, # d8 - df - 7, 8, 8, 8, 8, 8, 8, 8, # e0 - e7 - 8, 8, 8, 8, 8, 9, 8, 8, # e8 - ef - 10, 11, 11, 11, 11, 11, 11, 11, # f0 - f7 - 12, 13, 13, 13, 14, 15, 0, 0 # f8 - ff -) - -UTF8_ST = ( - MachineState.ERROR,MachineState.START,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR, 12, 10,#00-07 - 9, 11, 8, 7, 6, 5, 4, 3,#08-0f - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#10-17 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#18-1f - MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,#20-27 - MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,MachineState.ITS_ME,#28-2f - MachineState.ERROR,MachineState.ERROR, 5, 5, 5, 5,MachineState.ERROR,MachineState.ERROR,#30-37 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#38-3f - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR, 5, 5, 5,MachineState.ERROR,MachineState.ERROR,#40-47 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#48-4f - MachineState.ERROR,MachineState.ERROR, 7, 7, 7, 7,MachineState.ERROR,MachineState.ERROR,#50-57 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#58-5f - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR, 7, 7,MachineState.ERROR,MachineState.ERROR,#60-67 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#68-6f - MachineState.ERROR,MachineState.ERROR, 9, 9, 9, 9,MachineState.ERROR,MachineState.ERROR,#70-77 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#78-7f - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR, 9,MachineState.ERROR,MachineState.ERROR,#80-87 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#88-8f - MachineState.ERROR,MachineState.ERROR, 12, 12, 12, 12,MachineState.ERROR,MachineState.ERROR,#90-97 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#98-9f - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR, 12,MachineState.ERROR,MachineState.ERROR,#a0-a7 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#a8-af - MachineState.ERROR,MachineState.ERROR, 12, 12, 12,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#b0-b7 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,#b8-bf - MachineState.ERROR,MachineState.ERROR,MachineState.START,MachineState.START,MachineState.START,MachineState.START,MachineState.ERROR,MachineState.ERROR,#c0-c7 - MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR,MachineState.ERROR #c8-cf -) -# fmt: on - -UTF8_CHAR_LEN_TABLE = (0, 1, 0, 0, 0, 0, 2, 3, 3, 3, 4, 4, 5, 5, 6, 6) - -UTF8_SM_MODEL: CodingStateMachineDict = { - "class_table": UTF8_CLS, - "class_factor": 16, - "state_table": UTF8_ST, - "char_len_table": UTF8_CHAR_LEN_TABLE, - "name": "UTF-8", -} diff --git a/venv/lib/python3.10/site-packages/chardet/metadata/__init__.py b/venv/lib/python3.10/site-packages/chardet/metadata/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/chardet/metadata/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/metadata/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 5b2bd9f202b47ed341b19f5344ae24d5e7cd3b37..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/metadata/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/metadata/__pycache__/languages.cpython-310.pyc b/venv/lib/python3.10/site-packages/chardet/metadata/__pycache__/languages.cpython-310.pyc deleted file mode 100644 index 69615d514f8156d59d25b51e162a012d76c1111b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/chardet/metadata/__pycache__/languages.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/chardet/metadata/languages.py b/venv/lib/python3.10/site-packages/chardet/metadata/languages.py deleted file mode 100644 index eb40c5f0c8526208d434d762855d23079dc68b36..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/metadata/languages.py +++ /dev/null @@ -1,352 +0,0 @@ -""" -Metadata about languages used by our model training code for our -SingleByteCharSetProbers. Could be used for other things in the future. - -This code is based on the language metadata from the uchardet project. -""" - -from string import ascii_letters -from typing import List, Optional - -# TODO: Add Ukrainian (KOI8-U) - - -class Language: - """Metadata about a language useful for training models - - :ivar name: The human name for the language, in English. - :type name: str - :ivar iso_code: 2-letter ISO 639-1 if possible, 3-letter ISO code otherwise, - or use another catalog as a last resort. - :type iso_code: str - :ivar use_ascii: Whether or not ASCII letters should be included in trained - models. - :type use_ascii: bool - :ivar charsets: The charsets we want to support and create data for. - :type charsets: list of str - :ivar alphabet: The characters in the language's alphabet. If `use_ascii` is - `True`, you only need to add those not in the ASCII set. - :type alphabet: str - :ivar wiki_start_pages: The Wikipedia pages to start from if we're crawling - Wikipedia for training data. - :type wiki_start_pages: list of str - """ - - def __init__( - self, - name: Optional[str] = None, - iso_code: Optional[str] = None, - use_ascii: bool = True, - charsets: Optional[List[str]] = None, - alphabet: Optional[str] = None, - wiki_start_pages: Optional[List[str]] = None, - ) -> None: - super().__init__() - self.name = name - self.iso_code = iso_code - self.use_ascii = use_ascii - self.charsets = charsets - if self.use_ascii: - if alphabet: - alphabet += ascii_letters - else: - alphabet = ascii_letters - elif not alphabet: - raise ValueError("Must supply alphabet if use_ascii is False") - self.alphabet = "".join(sorted(set(alphabet))) if alphabet else None - self.wiki_start_pages = wiki_start_pages - - def __repr__(self) -> str: - param_str = ", ".join( - f"{k}={v!r}" for k, v in self.__dict__.items() if not k.startswith("_") - ) - return f"{self.__class__.__name__}({param_str})" - - -LANGUAGES = { - "Arabic": Language( - name="Arabic", - iso_code="ar", - use_ascii=False, - # We only support encodings that use isolated - # forms, because the current recommendation is - # that the rendering system handles presentation - # forms. This means we purposefully skip IBM864. - charsets=["ISO-8859-6", "WINDOWS-1256", "CP720", "CP864"], - alphabet="ءآأؤإئابةتثجحخدذرزسشصضطظعغػؼؽؾؿـفقكلمنهوىيًٌٍَُِّ", - wiki_start_pages=["الصفحة_الرئيسية"], - ), - "Belarusian": Language( - name="Belarusian", - iso_code="be", - use_ascii=False, - charsets=["ISO-8859-5", "WINDOWS-1251", "IBM866", "MacCyrillic"], - alphabet="АБВГДЕЁЖЗІЙКЛМНОПРСТУЎФХЦЧШЫЬЭЮЯабвгдеёжзійклмнопрстуўфхцчшыьэюяʼ", - wiki_start_pages=["Галоўная_старонка"], - ), - "Bulgarian": Language( - name="Bulgarian", - iso_code="bg", - use_ascii=False, - charsets=["ISO-8859-5", "WINDOWS-1251", "IBM855"], - alphabet="АБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЬЮЯабвгдежзийклмнопрстуфхцчшщъьюя", - wiki_start_pages=["Начална_страница"], - ), - "Czech": Language( - name="Czech", - iso_code="cz", - use_ascii=True, - charsets=["ISO-8859-2", "WINDOWS-1250"], - alphabet="áčďéěíňóřšťúůýžÁČĎÉĚÍŇÓŘŠŤÚŮÝŽ", - wiki_start_pages=["Hlavní_strana"], - ), - "Danish": Language( - name="Danish", - iso_code="da", - use_ascii=True, - charsets=["ISO-8859-1", "ISO-8859-15", "WINDOWS-1252", "MacRoman"], - alphabet="æøåÆØÅ", - wiki_start_pages=["Forside"], - ), - "German": Language( - name="German", - iso_code="de", - use_ascii=True, - charsets=["ISO-8859-1", "ISO-8859-15", "WINDOWS-1252", "MacRoman"], - alphabet="äöüßẞÄÖÜ", - wiki_start_pages=["Wikipedia:Hauptseite"], - ), - "Greek": Language( - name="Greek", - iso_code="el", - use_ascii=False, - charsets=["ISO-8859-7", "WINDOWS-1253"], - alphabet="αβγδεζηθικλμνξοπρσςτυφχψωάέήίόύώΑΒΓΔΕΖΗΘΙΚΛΜΝΞΟΠΡΣΣΤΥΦΧΨΩΆΈΉΊΌΎΏ", - wiki_start_pages=["Πύλη:Κύρια"], - ), - "English": Language( - name="English", - iso_code="en", - use_ascii=True, - charsets=["ISO-8859-1", "WINDOWS-1252", "MacRoman"], - wiki_start_pages=["Main_Page"], - ), - "Esperanto": Language( - name="Esperanto", - iso_code="eo", - # Q, W, X, and Y not used at all - use_ascii=False, - charsets=["ISO-8859-3"], - alphabet="abcĉdefgĝhĥijĵklmnoprsŝtuŭvzABCĈDEFGĜHĤIJĴKLMNOPRSŜTUŬVZ", - wiki_start_pages=["Vikipedio:Ĉefpaĝo"], - ), - "Spanish": Language( - name="Spanish", - iso_code="es", - use_ascii=True, - charsets=["ISO-8859-1", "ISO-8859-15", "WINDOWS-1252", "MacRoman"], - alphabet="ñáéíóúüÑÁÉÍÓÚÜ", - wiki_start_pages=["Wikipedia:Portada"], - ), - "Estonian": Language( - name="Estonian", - iso_code="et", - use_ascii=False, - charsets=["ISO-8859-4", "ISO-8859-13", "WINDOWS-1257"], - # C, F, Š, Q, W, X, Y, Z, Ž are only for - # loanwords - alphabet="ABDEGHIJKLMNOPRSTUVÕÄÖÜabdeghijklmnoprstuvõäöü", - wiki_start_pages=["Esileht"], - ), - "Finnish": Language( - name="Finnish", - iso_code="fi", - use_ascii=True, - charsets=["ISO-8859-1", "ISO-8859-15", "WINDOWS-1252", "MacRoman"], - alphabet="ÅÄÖŠŽåäöšž", - wiki_start_pages=["Wikipedia:Etusivu"], - ), - "French": Language( - name="French", - iso_code="fr", - use_ascii=True, - charsets=["ISO-8859-1", "ISO-8859-15", "WINDOWS-1252", "MacRoman"], - alphabet="œàâçèéîïùûêŒÀÂÇÈÉÎÏÙÛÊ", - wiki_start_pages=["Wikipédia:Accueil_principal", "Bœuf (animal)"], - ), - "Hebrew": Language( - name="Hebrew", - iso_code="he", - use_ascii=False, - charsets=["ISO-8859-8", "WINDOWS-1255"], - alphabet="אבגדהוזחטיךכלםמןנסעףפץצקרשתװױײ", - wiki_start_pages=["עמוד_ראשי"], - ), - "Croatian": Language( - name="Croatian", - iso_code="hr", - # Q, W, X, Y are only used for foreign words. - use_ascii=False, - charsets=["ISO-8859-2", "WINDOWS-1250"], - alphabet="abcčćdđefghijklmnoprsštuvzžABCČĆDĐEFGHIJKLMNOPRSŠTUVZŽ", - wiki_start_pages=["Glavna_stranica"], - ), - "Hungarian": Language( - name="Hungarian", - iso_code="hu", - # Q, W, X, Y are only used for foreign words. - use_ascii=False, - charsets=["ISO-8859-2", "WINDOWS-1250"], - alphabet="abcdefghijklmnoprstuvzáéíóöőúüűABCDEFGHIJKLMNOPRSTUVZÁÉÍÓÖŐÚÜŰ", - wiki_start_pages=["Kezdőlap"], - ), - "Italian": Language( - name="Italian", - iso_code="it", - use_ascii=True, - charsets=["ISO-8859-1", "ISO-8859-15", "WINDOWS-1252", "MacRoman"], - alphabet="ÀÈÉÌÒÓÙàèéìòóù", - wiki_start_pages=["Pagina_principale"], - ), - "Lithuanian": Language( - name="Lithuanian", - iso_code="lt", - use_ascii=False, - charsets=["ISO-8859-13", "WINDOWS-1257", "ISO-8859-4"], - # Q, W, and X not used at all - alphabet="AĄBCČDEĘĖFGHIĮYJKLMNOPRSŠTUŲŪVZŽaąbcčdeęėfghiįyjklmnoprsštuųūvzž", - wiki_start_pages=["Pagrindinis_puslapis"], - ), - "Latvian": Language( - name="Latvian", - iso_code="lv", - use_ascii=False, - charsets=["ISO-8859-13", "WINDOWS-1257", "ISO-8859-4"], - # Q, W, X, Y are only for loanwords - alphabet="AĀBCČDEĒFGĢHIĪJKĶLĻMNŅOPRSŠTUŪVZŽaābcčdeēfgģhiījkķlļmnņoprsštuūvzž", - wiki_start_pages=["Sākumlapa"], - ), - "Macedonian": Language( - name="Macedonian", - iso_code="mk", - use_ascii=False, - charsets=["ISO-8859-5", "WINDOWS-1251", "MacCyrillic", "IBM855"], - alphabet="АБВГДЃЕЖЗЅИЈКЛЉМНЊОПРСТЌУФХЦЧЏШабвгдѓежзѕијклљмнњопрстќуфхцчџш", - wiki_start_pages=["Главна_страница"], - ), - "Dutch": Language( - name="Dutch", - iso_code="nl", - use_ascii=True, - charsets=["ISO-8859-1", "WINDOWS-1252", "MacRoman"], - wiki_start_pages=["Hoofdpagina"], - ), - "Polish": Language( - name="Polish", - iso_code="pl", - # Q and X are only used for foreign words. - use_ascii=False, - charsets=["ISO-8859-2", "WINDOWS-1250"], - alphabet="AĄBCĆDEĘFGHIJKLŁMNŃOÓPRSŚTUWYZŹŻaąbcćdeęfghijklłmnńoóprsśtuwyzźż", - wiki_start_pages=["Wikipedia:Strona_główna"], - ), - "Portuguese": Language( - name="Portuguese", - iso_code="pt", - use_ascii=True, - charsets=["ISO-8859-1", "ISO-8859-15", "WINDOWS-1252", "MacRoman"], - alphabet="ÁÂÃÀÇÉÊÍÓÔÕÚáâãàçéêíóôõú", - wiki_start_pages=["Wikipédia:Página_principal"], - ), - "Romanian": Language( - name="Romanian", - iso_code="ro", - use_ascii=True, - charsets=["ISO-8859-2", "WINDOWS-1250"], - alphabet="ăâîșțĂÂÎȘȚ", - wiki_start_pages=["Pagina_principală"], - ), - "Russian": Language( - name="Russian", - iso_code="ru", - use_ascii=False, - charsets=[ - "ISO-8859-5", - "WINDOWS-1251", - "KOI8-R", - "MacCyrillic", - "IBM866", - "IBM855", - ], - alphabet="абвгдеёжзийклмнопрстуфхцчшщъыьэюяАБВГДЕЁЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯ", - wiki_start_pages=["Заглавная_страница"], - ), - "Slovak": Language( - name="Slovak", - iso_code="sk", - use_ascii=True, - charsets=["ISO-8859-2", "WINDOWS-1250"], - alphabet="áäčďéíĺľňóôŕšťúýžÁÄČĎÉÍĹĽŇÓÔŔŠŤÚÝŽ", - wiki_start_pages=["Hlavná_stránka"], - ), - "Slovene": Language( - name="Slovene", - iso_code="sl", - # Q, W, X, Y are only used for foreign words. - use_ascii=False, - charsets=["ISO-8859-2", "WINDOWS-1250"], - alphabet="abcčdefghijklmnoprsštuvzžABCČDEFGHIJKLMNOPRSŠTUVZŽ", - wiki_start_pages=["Glavna_stran"], - ), - # Serbian can be written in both Latin and Cyrillic, but there's no - # simple way to get the Latin alphabet pages from Wikipedia through - # the API, so for now we just support Cyrillic. - "Serbian": Language( - name="Serbian", - iso_code="sr", - alphabet="АБВГДЂЕЖЗИЈКЛЉМНЊОПРСТЋУФХЦЧЏШабвгдђежзијклљмнњопрстћуфхцчџш", - charsets=["ISO-8859-5", "WINDOWS-1251", "MacCyrillic", "IBM855"], - wiki_start_pages=["Главна_страна"], - ), - "Thai": Language( - name="Thai", - iso_code="th", - use_ascii=False, - charsets=["ISO-8859-11", "TIS-620", "CP874"], - alphabet="กขฃคฅฆงจฉชซฌญฎฏฐฑฒณดตถทธนบปผฝพฟภมยรฤลฦวศษสหฬอฮฯะัาำิีึืฺุู฿เแโใไๅๆ็่้๊๋์ํ๎๏๐๑๒๓๔๕๖๗๘๙๚๛", - wiki_start_pages=["หน้าหลัก"], - ), - "Turkish": Language( - name="Turkish", - iso_code="tr", - # Q, W, and X are not used by Turkish - use_ascii=False, - charsets=["ISO-8859-3", "ISO-8859-9", "WINDOWS-1254"], - alphabet="abcçdefgğhıijklmnoöprsştuüvyzâîûABCÇDEFGĞHIİJKLMNOÖPRSŞTUÜVYZÂÎÛ", - wiki_start_pages=["Ana_Sayfa"], - ), - "Vietnamese": Language( - name="Vietnamese", - iso_code="vi", - use_ascii=False, - # Windows-1258 is the only common 8-bit - # Vietnamese encoding supported by Python. - # From Wikipedia: - # For systems that lack support for Unicode, - # dozens of 8-bit Vietnamese code pages are - # available.[1] The most common are VISCII - # (TCVN 5712:1993), VPS, and Windows-1258.[3] - # Where ASCII is required, such as when - # ensuring readability in plain text e-mail, - # Vietnamese letters are often encoded - # according to Vietnamese Quoted-Readable - # (VIQR) or VSCII Mnemonic (VSCII-MNEM),[4] - # though usage of either variable-width - # scheme has declined dramatically following - # the adoption of Unicode on the World Wide - # Web. - charsets=["WINDOWS-1258"], - alphabet="aăâbcdđeêghiklmnoôơpqrstuưvxyAĂÂBCDĐEÊGHIKLMNOÔƠPQRSTUƯVXY", - wiki_start_pages=["Chữ_Quốc_ngữ"], - ), -} diff --git a/venv/lib/python3.10/site-packages/chardet/py.typed b/venv/lib/python3.10/site-packages/chardet/py.typed deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/chardet/resultdict.py b/venv/lib/python3.10/site-packages/chardet/resultdict.py deleted file mode 100644 index 7d36e64c467ca8d9cadc88ab03da71faf1aa8abb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/resultdict.py +++ /dev/null @@ -1,16 +0,0 @@ -from typing import TYPE_CHECKING, Optional - -if TYPE_CHECKING: - # TypedDict was introduced in Python 3.8. - # - # TODO: Remove the else block and TYPE_CHECKING check when dropping support - # for Python 3.7. - from typing import TypedDict - - class ResultDict(TypedDict): - encoding: Optional[str] - confidence: float - language: Optional[str] - -else: - ResultDict = dict diff --git a/venv/lib/python3.10/site-packages/chardet/sbcharsetprober.py b/venv/lib/python3.10/site-packages/chardet/sbcharsetprober.py deleted file mode 100644 index 0ffbcdd2c3e21b68566c88a3f05239447489df84..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/sbcharsetprober.py +++ /dev/null @@ -1,162 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Universal charset detector code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 2001 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# Shy Shalom - original C code -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from typing import Dict, List, NamedTuple, Optional, Union - -from .charsetprober import CharSetProber -from .enums import CharacterCategory, ProbingState, SequenceLikelihood - - -class SingleByteCharSetModel(NamedTuple): - charset_name: str - language: str - char_to_order_map: Dict[int, int] - language_model: Dict[int, Dict[int, int]] - typical_positive_ratio: float - keep_ascii_letters: bool - alphabet: str - - -class SingleByteCharSetProber(CharSetProber): - SAMPLE_SIZE = 64 - SB_ENOUGH_REL_THRESHOLD = 1024 # 0.25 * SAMPLE_SIZE^2 - POSITIVE_SHORTCUT_THRESHOLD = 0.95 - NEGATIVE_SHORTCUT_THRESHOLD = 0.05 - - def __init__( - self, - model: SingleByteCharSetModel, - is_reversed: bool = False, - name_prober: Optional[CharSetProber] = None, - ) -> None: - super().__init__() - self._model = model - # TRUE if we need to reverse every pair in the model lookup - self._reversed = is_reversed - # Optional auxiliary prober for name decision - self._name_prober = name_prober - self._last_order = 255 - self._seq_counters: List[int] = [] - self._total_seqs = 0 - self._total_char = 0 - self._control_char = 0 - self._freq_char = 0 - self.reset() - - def reset(self) -> None: - super().reset() - # char order of last character - self._last_order = 255 - self._seq_counters = [0] * SequenceLikelihood.get_num_categories() - self._total_seqs = 0 - self._total_char = 0 - self._control_char = 0 - # characters that fall in our sampling range - self._freq_char = 0 - - @property - def charset_name(self) -> Optional[str]: - if self._name_prober: - return self._name_prober.charset_name - return self._model.charset_name - - @property - def language(self) -> Optional[str]: - if self._name_prober: - return self._name_prober.language - return self._model.language - - def feed(self, byte_str: Union[bytes, bytearray]) -> ProbingState: - # TODO: Make filter_international_words keep things in self.alphabet - if not self._model.keep_ascii_letters: - byte_str = self.filter_international_words(byte_str) - else: - byte_str = self.remove_xml_tags(byte_str) - if not byte_str: - return self.state - char_to_order_map = self._model.char_to_order_map - language_model = self._model.language_model - for char in byte_str: - order = char_to_order_map.get(char, CharacterCategory.UNDEFINED) - # XXX: This was SYMBOL_CAT_ORDER before, with a value of 250, but - # CharacterCategory.SYMBOL is actually 253, so we use CONTROL - # to make it closer to the original intent. The only difference - # is whether or not we count digits and control characters for - # _total_char purposes. - if order < CharacterCategory.CONTROL: - self._total_char += 1 - if order < self.SAMPLE_SIZE: - self._freq_char += 1 - if self._last_order < self.SAMPLE_SIZE: - self._total_seqs += 1 - if not self._reversed: - lm_cat = language_model[self._last_order][order] - else: - lm_cat = language_model[order][self._last_order] - self._seq_counters[lm_cat] += 1 - self._last_order = order - - charset_name = self._model.charset_name - if self.state == ProbingState.DETECTING: - if self._total_seqs > self.SB_ENOUGH_REL_THRESHOLD: - confidence = self.get_confidence() - if confidence > self.POSITIVE_SHORTCUT_THRESHOLD: - self.logger.debug( - "%s confidence = %s, we have a winner", charset_name, confidence - ) - self._state = ProbingState.FOUND_IT - elif confidence < self.NEGATIVE_SHORTCUT_THRESHOLD: - self.logger.debug( - "%s confidence = %s, below negative shortcut threshold %s", - charset_name, - confidence, - self.NEGATIVE_SHORTCUT_THRESHOLD, - ) - self._state = ProbingState.NOT_ME - - return self.state - - def get_confidence(self) -> float: - r = 0.01 - if self._total_seqs > 0: - r = ( - ( - self._seq_counters[SequenceLikelihood.POSITIVE] - + 0.25 * self._seq_counters[SequenceLikelihood.LIKELY] - ) - / self._total_seqs - / self._model.typical_positive_ratio - ) - # The more control characters (proportionnaly to the size - # of the text), the less confident we become in the current - # charset. - r = r * (self._total_char - self._control_char) / self._total_char - r = r * self._freq_char / self._total_char - if r >= 1.0: - r = 0.99 - return r diff --git a/venv/lib/python3.10/site-packages/chardet/sbcsgroupprober.py b/venv/lib/python3.10/site-packages/chardet/sbcsgroupprober.py deleted file mode 100644 index 890ae8465c5b0ad2a5f99464fe5f5c0be49809f1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/sbcsgroupprober.py +++ /dev/null @@ -1,88 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Universal charset detector code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 2001 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# Shy Shalom - original C code -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from .charsetgroupprober import CharSetGroupProber -from .hebrewprober import HebrewProber -from .langbulgarianmodel import ISO_8859_5_BULGARIAN_MODEL, WINDOWS_1251_BULGARIAN_MODEL -from .langgreekmodel import ISO_8859_7_GREEK_MODEL, WINDOWS_1253_GREEK_MODEL -from .langhebrewmodel import WINDOWS_1255_HEBREW_MODEL - -# from .langhungarianmodel import (ISO_8859_2_HUNGARIAN_MODEL, -# WINDOWS_1250_HUNGARIAN_MODEL) -from .langrussianmodel import ( - IBM855_RUSSIAN_MODEL, - IBM866_RUSSIAN_MODEL, - ISO_8859_5_RUSSIAN_MODEL, - KOI8_R_RUSSIAN_MODEL, - MACCYRILLIC_RUSSIAN_MODEL, - WINDOWS_1251_RUSSIAN_MODEL, -) -from .langthaimodel import TIS_620_THAI_MODEL -from .langturkishmodel import ISO_8859_9_TURKISH_MODEL -from .sbcharsetprober import SingleByteCharSetProber - - -class SBCSGroupProber(CharSetGroupProber): - def __init__(self) -> None: - super().__init__() - hebrew_prober = HebrewProber() - logical_hebrew_prober = SingleByteCharSetProber( - WINDOWS_1255_HEBREW_MODEL, is_reversed=False, name_prober=hebrew_prober - ) - # TODO: See if using ISO-8859-8 Hebrew model works better here, since - # it's actually the visual one - visual_hebrew_prober = SingleByteCharSetProber( - WINDOWS_1255_HEBREW_MODEL, is_reversed=True, name_prober=hebrew_prober - ) - hebrew_prober.set_model_probers(logical_hebrew_prober, visual_hebrew_prober) - # TODO: ORDER MATTERS HERE. I changed the order vs what was in master - # and several tests failed that did not before. Some thought - # should be put into the ordering, and we should consider making - # order not matter here, because that is very counter-intuitive. - self.probers = [ - SingleByteCharSetProber(WINDOWS_1251_RUSSIAN_MODEL), - SingleByteCharSetProber(KOI8_R_RUSSIAN_MODEL), - SingleByteCharSetProber(ISO_8859_5_RUSSIAN_MODEL), - SingleByteCharSetProber(MACCYRILLIC_RUSSIAN_MODEL), - SingleByteCharSetProber(IBM866_RUSSIAN_MODEL), - SingleByteCharSetProber(IBM855_RUSSIAN_MODEL), - SingleByteCharSetProber(ISO_8859_7_GREEK_MODEL), - SingleByteCharSetProber(WINDOWS_1253_GREEK_MODEL), - SingleByteCharSetProber(ISO_8859_5_BULGARIAN_MODEL), - SingleByteCharSetProber(WINDOWS_1251_BULGARIAN_MODEL), - # TODO: Restore Hungarian encodings (iso-8859-2 and windows-1250) - # after we retrain model. - # SingleByteCharSetProber(ISO_8859_2_HUNGARIAN_MODEL), - # SingleByteCharSetProber(WINDOWS_1250_HUNGARIAN_MODEL), - SingleByteCharSetProber(TIS_620_THAI_MODEL), - SingleByteCharSetProber(ISO_8859_9_TURKISH_MODEL), - hebrew_prober, - logical_hebrew_prober, - visual_hebrew_prober, - ] - self.reset() diff --git a/venv/lib/python3.10/site-packages/chardet/sjisprober.py b/venv/lib/python3.10/site-packages/chardet/sjisprober.py deleted file mode 100644 index 91df077961b6310b8e1c708b74003d5343bff6a8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/sjisprober.py +++ /dev/null @@ -1,105 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is mozilla.org code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from typing import Union - -from .chardistribution import SJISDistributionAnalysis -from .codingstatemachine import CodingStateMachine -from .enums import MachineState, ProbingState -from .jpcntx import SJISContextAnalysis -from .mbcharsetprober import MultiByteCharSetProber -from .mbcssm import SJIS_SM_MODEL - - -class SJISProber(MultiByteCharSetProber): - def __init__(self) -> None: - super().__init__() - self.coding_sm = CodingStateMachine(SJIS_SM_MODEL) - self.distribution_analyzer = SJISDistributionAnalysis() - self.context_analyzer = SJISContextAnalysis() - self.reset() - - def reset(self) -> None: - super().reset() - self.context_analyzer.reset() - - @property - def charset_name(self) -> str: - return self.context_analyzer.charset_name - - @property - def language(self) -> str: - return "Japanese" - - def feed(self, byte_str: Union[bytes, bytearray]) -> ProbingState: - assert self.coding_sm is not None - assert self.distribution_analyzer is not None - - for i, byte in enumerate(byte_str): - coding_state = self.coding_sm.next_state(byte) - if coding_state == MachineState.ERROR: - self.logger.debug( - "%s %s prober hit error at byte %s", - self.charset_name, - self.language, - i, - ) - self._state = ProbingState.NOT_ME - break - if coding_state == MachineState.ITS_ME: - self._state = ProbingState.FOUND_IT - break - if coding_state == MachineState.START: - char_len = self.coding_sm.get_current_charlen() - if i == 0: - self._last_char[1] = byte - self.context_analyzer.feed( - self._last_char[2 - char_len :], char_len - ) - self.distribution_analyzer.feed(self._last_char, char_len) - else: - self.context_analyzer.feed( - byte_str[i + 1 - char_len : i + 3 - char_len], char_len - ) - self.distribution_analyzer.feed(byte_str[i - 1 : i + 1], char_len) - - self._last_char[0] = byte_str[-1] - - if self.state == ProbingState.DETECTING: - if self.context_analyzer.got_enough_data() and ( - self.get_confidence() > self.SHORTCUT_THRESHOLD - ): - self._state = ProbingState.FOUND_IT - - return self.state - - def get_confidence(self) -> float: - assert self.distribution_analyzer is not None - - context_conf = self.context_analyzer.get_confidence() - distrib_conf = self.distribution_analyzer.get_confidence() - return max(context_conf, distrib_conf) diff --git a/venv/lib/python3.10/site-packages/chardet/universaldetector.py b/venv/lib/python3.10/site-packages/chardet/universaldetector.py deleted file mode 100644 index 30c441dc28ee327076a850b1d3c88a9a2c8f04f0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/universaldetector.py +++ /dev/null @@ -1,362 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is Mozilla Universal charset detector code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 2001 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# Shy Shalom - original C code -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### -""" -Module containing the UniversalDetector detector class, which is the primary -class a user of ``chardet`` should use. - -:author: Mark Pilgrim (initial port to Python) -:author: Shy Shalom (original C code) -:author: Dan Blanchard (major refactoring for 3.0) -:author: Ian Cordasco -""" - - -import codecs -import logging -import re -from typing import List, Optional, Union - -from .charsetgroupprober import CharSetGroupProber -from .charsetprober import CharSetProber -from .enums import InputState, LanguageFilter, ProbingState -from .escprober import EscCharSetProber -from .latin1prober import Latin1Prober -from .macromanprober import MacRomanProber -from .mbcsgroupprober import MBCSGroupProber -from .resultdict import ResultDict -from .sbcsgroupprober import SBCSGroupProber -from .utf1632prober import UTF1632Prober - - -class UniversalDetector: - """ - The ``UniversalDetector`` class underlies the ``chardet.detect`` function - and coordinates all of the different charset probers. - - To get a ``dict`` containing an encoding and its confidence, you can simply - run: - - .. code:: - - u = UniversalDetector() - u.feed(some_bytes) - u.close() - detected = u.result - - """ - - MINIMUM_THRESHOLD = 0.20 - HIGH_BYTE_DETECTOR = re.compile(b"[\x80-\xFF]") - ESC_DETECTOR = re.compile(b"(\033|~{)") - WIN_BYTE_DETECTOR = re.compile(b"[\x80-\x9F]") - ISO_WIN_MAP = { - "iso-8859-1": "Windows-1252", - "iso-8859-2": "Windows-1250", - "iso-8859-5": "Windows-1251", - "iso-8859-6": "Windows-1256", - "iso-8859-7": "Windows-1253", - "iso-8859-8": "Windows-1255", - "iso-8859-9": "Windows-1254", - "iso-8859-13": "Windows-1257", - } - # Based on https://encoding.spec.whatwg.org/#names-and-labels - # but altered to match Python names for encodings and remove mappings - # that break tests. - LEGACY_MAP = { - "ascii": "Windows-1252", - "iso-8859-1": "Windows-1252", - "tis-620": "ISO-8859-11", - "iso-8859-9": "Windows-1254", - "gb2312": "GB18030", - "euc-kr": "CP949", - "utf-16le": "UTF-16", - } - - def __init__( - self, - lang_filter: LanguageFilter = LanguageFilter.ALL, - should_rename_legacy: bool = False, - ) -> None: - self._esc_charset_prober: Optional[EscCharSetProber] = None - self._utf1632_prober: Optional[UTF1632Prober] = None - self._charset_probers: List[CharSetProber] = [] - self.result: ResultDict = { - "encoding": None, - "confidence": 0.0, - "language": None, - } - self.done = False - self._got_data = False - self._input_state = InputState.PURE_ASCII - self._last_char = b"" - self.lang_filter = lang_filter - self.logger = logging.getLogger(__name__) - self._has_win_bytes = False - self.should_rename_legacy = should_rename_legacy - self.reset() - - @property - def input_state(self) -> int: - return self._input_state - - @property - def has_win_bytes(self) -> bool: - return self._has_win_bytes - - @property - def charset_probers(self) -> List[CharSetProber]: - return self._charset_probers - - def reset(self) -> None: - """ - Reset the UniversalDetector and all of its probers back to their - initial states. This is called by ``__init__``, so you only need to - call this directly in between analyses of different documents. - """ - self.result = {"encoding": None, "confidence": 0.0, "language": None} - self.done = False - self._got_data = False - self._has_win_bytes = False - self._input_state = InputState.PURE_ASCII - self._last_char = b"" - if self._esc_charset_prober: - self._esc_charset_prober.reset() - if self._utf1632_prober: - self._utf1632_prober.reset() - for prober in self._charset_probers: - prober.reset() - - def feed(self, byte_str: Union[bytes, bytearray]) -> None: - """ - Takes a chunk of a document and feeds it through all of the relevant - charset probers. - - After calling ``feed``, you can check the value of the ``done`` - attribute to see if you need to continue feeding the - ``UniversalDetector`` more data, or if it has made a prediction - (in the ``result`` attribute). - - .. note:: - You should always call ``close`` when you're done feeding in your - document if ``done`` is not already ``True``. - """ - if self.done: - return - - if not byte_str: - return - - if not isinstance(byte_str, bytearray): - byte_str = bytearray(byte_str) - - # First check for known BOMs, since these are guaranteed to be correct - if not self._got_data: - # If the data starts with BOM, we know it is UTF - if byte_str.startswith(codecs.BOM_UTF8): - # EF BB BF UTF-8 with BOM - self.result = { - "encoding": "UTF-8-SIG", - "confidence": 1.0, - "language": "", - } - elif byte_str.startswith((codecs.BOM_UTF32_LE, codecs.BOM_UTF32_BE)): - # FF FE 00 00 UTF-32, little-endian BOM - # 00 00 FE FF UTF-32, big-endian BOM - self.result = {"encoding": "UTF-32", "confidence": 1.0, "language": ""} - elif byte_str.startswith(b"\xFE\xFF\x00\x00"): - # FE FF 00 00 UCS-4, unusual octet order BOM (3412) - self.result = { - # TODO: This encoding is not supported by Python. Should remove? - "encoding": "X-ISO-10646-UCS-4-3412", - "confidence": 1.0, - "language": "", - } - elif byte_str.startswith(b"\x00\x00\xFF\xFE"): - # 00 00 FF FE UCS-4, unusual octet order BOM (2143) - self.result = { - # TODO: This encoding is not supported by Python. Should remove? - "encoding": "X-ISO-10646-UCS-4-2143", - "confidence": 1.0, - "language": "", - } - elif byte_str.startswith((codecs.BOM_LE, codecs.BOM_BE)): - # FF FE UTF-16, little endian BOM - # FE FF UTF-16, big endian BOM - self.result = {"encoding": "UTF-16", "confidence": 1.0, "language": ""} - - self._got_data = True - if self.result["encoding"] is not None: - self.done = True - return - - # If none of those matched and we've only see ASCII so far, check - # for high bytes and escape sequences - if self._input_state == InputState.PURE_ASCII: - if self.HIGH_BYTE_DETECTOR.search(byte_str): - self._input_state = InputState.HIGH_BYTE - elif ( - self._input_state == InputState.PURE_ASCII - and self.ESC_DETECTOR.search(self._last_char + byte_str) - ): - self._input_state = InputState.ESC_ASCII - - self._last_char = byte_str[-1:] - - # next we will look to see if it is appears to be either a UTF-16 or - # UTF-32 encoding - if not self._utf1632_prober: - self._utf1632_prober = UTF1632Prober() - - if self._utf1632_prober.state == ProbingState.DETECTING: - if self._utf1632_prober.feed(byte_str) == ProbingState.FOUND_IT: - self.result = { - "encoding": self._utf1632_prober.charset_name, - "confidence": self._utf1632_prober.get_confidence(), - "language": "", - } - self.done = True - return - - # If we've seen escape sequences, use the EscCharSetProber, which - # uses a simple state machine to check for known escape sequences in - # HZ and ISO-2022 encodings, since those are the only encodings that - # use such sequences. - if self._input_state == InputState.ESC_ASCII: - if not self._esc_charset_prober: - self._esc_charset_prober = EscCharSetProber(self.lang_filter) - if self._esc_charset_prober.feed(byte_str) == ProbingState.FOUND_IT: - self.result = { - "encoding": self._esc_charset_prober.charset_name, - "confidence": self._esc_charset_prober.get_confidence(), - "language": self._esc_charset_prober.language, - } - self.done = True - # If we've seen high bytes (i.e., those with values greater than 127), - # we need to do more complicated checks using all our multi-byte and - # single-byte probers that are left. The single-byte probers - # use character bigram distributions to determine the encoding, whereas - # the multi-byte probers use a combination of character unigram and - # bigram distributions. - elif self._input_state == InputState.HIGH_BYTE: - if not self._charset_probers: - self._charset_probers = [MBCSGroupProber(self.lang_filter)] - # If we're checking non-CJK encodings, use single-byte prober - if self.lang_filter & LanguageFilter.NON_CJK: - self._charset_probers.append(SBCSGroupProber()) - self._charset_probers.append(Latin1Prober()) - self._charset_probers.append(MacRomanProber()) - for prober in self._charset_probers: - if prober.feed(byte_str) == ProbingState.FOUND_IT: - self.result = { - "encoding": prober.charset_name, - "confidence": prober.get_confidence(), - "language": prober.language, - } - self.done = True - break - if self.WIN_BYTE_DETECTOR.search(byte_str): - self._has_win_bytes = True - - def close(self) -> ResultDict: - """ - Stop analyzing the current document and come up with a final - prediction. - - :returns: The ``result`` attribute, a ``dict`` with the keys - `encoding`, `confidence`, and `language`. - """ - # Don't bother with checks if we're already done - if self.done: - return self.result - self.done = True - - if not self._got_data: - self.logger.debug("no data received!") - - # Default to ASCII if it is all we've seen so far - elif self._input_state == InputState.PURE_ASCII: - self.result = {"encoding": "ascii", "confidence": 1.0, "language": ""} - - # If we have seen non-ASCII, return the best that met MINIMUM_THRESHOLD - elif self._input_state == InputState.HIGH_BYTE: - prober_confidence = None - max_prober_confidence = 0.0 - max_prober = None - for prober in self._charset_probers: - if not prober: - continue - prober_confidence = prober.get_confidence() - if prober_confidence > max_prober_confidence: - max_prober_confidence = prober_confidence - max_prober = prober - if max_prober and (max_prober_confidence > self.MINIMUM_THRESHOLD): - charset_name = max_prober.charset_name - assert charset_name is not None - lower_charset_name = charset_name.lower() - confidence = max_prober.get_confidence() - # Use Windows encoding name instead of ISO-8859 if we saw any - # extra Windows-specific bytes - if lower_charset_name.startswith("iso-8859"): - if self._has_win_bytes: - charset_name = self.ISO_WIN_MAP.get( - lower_charset_name, charset_name - ) - # Rename legacy encodings with superset encodings if asked - if self.should_rename_legacy: - charset_name = self.LEGACY_MAP.get( - (charset_name or "").lower(), charset_name - ) - self.result = { - "encoding": charset_name, - "confidence": confidence, - "language": max_prober.language, - } - - # Log all prober confidences if none met MINIMUM_THRESHOLD - if self.logger.getEffectiveLevel() <= logging.DEBUG: - if self.result["encoding"] is None: - self.logger.debug("no probers hit minimum threshold") - for group_prober in self._charset_probers: - if not group_prober: - continue - if isinstance(group_prober, CharSetGroupProber): - for prober in group_prober.probers: - self.logger.debug( - "%s %s confidence = %s", - prober.charset_name, - prober.language, - prober.get_confidence(), - ) - else: - self.logger.debug( - "%s %s confidence = %s", - group_prober.charset_name, - group_prober.language, - group_prober.get_confidence(), - ) - return self.result diff --git a/venv/lib/python3.10/site-packages/chardet/utf1632prober.py b/venv/lib/python3.10/site-packages/chardet/utf1632prober.py deleted file mode 100644 index 6bdec63d6867928bf73a7e513f60cee8f49ca050..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/utf1632prober.py +++ /dev/null @@ -1,225 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# -# Contributor(s): -# Jason Zavaglia -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### -from typing import List, Union - -from .charsetprober import CharSetProber -from .enums import ProbingState - - -class UTF1632Prober(CharSetProber): - """ - This class simply looks for occurrences of zero bytes, and infers - whether the file is UTF16 or UTF32 (low-endian or big-endian) - For instance, files looking like ( \0 \0 \0 [nonzero] )+ - have a good probability to be UTF32BE. Files looking like ( \0 [nonzero] )+ - may be guessed to be UTF16BE, and inversely for little-endian varieties. - """ - - # how many logical characters to scan before feeling confident of prediction - MIN_CHARS_FOR_DETECTION = 20 - # a fixed constant ratio of expected zeros or non-zeros in modulo-position. - EXPECTED_RATIO = 0.94 - - def __init__(self) -> None: - super().__init__() - self.position = 0 - self.zeros_at_mod = [0] * 4 - self.nonzeros_at_mod = [0] * 4 - self._state = ProbingState.DETECTING - self.quad = [0, 0, 0, 0] - self.invalid_utf16be = False - self.invalid_utf16le = False - self.invalid_utf32be = False - self.invalid_utf32le = False - self.first_half_surrogate_pair_detected_16be = False - self.first_half_surrogate_pair_detected_16le = False - self.reset() - - def reset(self) -> None: - super().reset() - self.position = 0 - self.zeros_at_mod = [0] * 4 - self.nonzeros_at_mod = [0] * 4 - self._state = ProbingState.DETECTING - self.invalid_utf16be = False - self.invalid_utf16le = False - self.invalid_utf32be = False - self.invalid_utf32le = False - self.first_half_surrogate_pair_detected_16be = False - self.first_half_surrogate_pair_detected_16le = False - self.quad = [0, 0, 0, 0] - - @property - def charset_name(self) -> str: - if self.is_likely_utf32be(): - return "utf-32be" - if self.is_likely_utf32le(): - return "utf-32le" - if self.is_likely_utf16be(): - return "utf-16be" - if self.is_likely_utf16le(): - return "utf-16le" - # default to something valid - return "utf-16" - - @property - def language(self) -> str: - return "" - - def approx_32bit_chars(self) -> float: - return max(1.0, self.position / 4.0) - - def approx_16bit_chars(self) -> float: - return max(1.0, self.position / 2.0) - - def is_likely_utf32be(self) -> bool: - approx_chars = self.approx_32bit_chars() - return approx_chars >= self.MIN_CHARS_FOR_DETECTION and ( - self.zeros_at_mod[0] / approx_chars > self.EXPECTED_RATIO - and self.zeros_at_mod[1] / approx_chars > self.EXPECTED_RATIO - and self.zeros_at_mod[2] / approx_chars > self.EXPECTED_RATIO - and self.nonzeros_at_mod[3] / approx_chars > self.EXPECTED_RATIO - and not self.invalid_utf32be - ) - - def is_likely_utf32le(self) -> bool: - approx_chars = self.approx_32bit_chars() - return approx_chars >= self.MIN_CHARS_FOR_DETECTION and ( - self.nonzeros_at_mod[0] / approx_chars > self.EXPECTED_RATIO - and self.zeros_at_mod[1] / approx_chars > self.EXPECTED_RATIO - and self.zeros_at_mod[2] / approx_chars > self.EXPECTED_RATIO - and self.zeros_at_mod[3] / approx_chars > self.EXPECTED_RATIO - and not self.invalid_utf32le - ) - - def is_likely_utf16be(self) -> bool: - approx_chars = self.approx_16bit_chars() - return approx_chars >= self.MIN_CHARS_FOR_DETECTION and ( - (self.nonzeros_at_mod[1] + self.nonzeros_at_mod[3]) / approx_chars - > self.EXPECTED_RATIO - and (self.zeros_at_mod[0] + self.zeros_at_mod[2]) / approx_chars - > self.EXPECTED_RATIO - and not self.invalid_utf16be - ) - - def is_likely_utf16le(self) -> bool: - approx_chars = self.approx_16bit_chars() - return approx_chars >= self.MIN_CHARS_FOR_DETECTION and ( - (self.nonzeros_at_mod[0] + self.nonzeros_at_mod[2]) / approx_chars - > self.EXPECTED_RATIO - and (self.zeros_at_mod[1] + self.zeros_at_mod[3]) / approx_chars - > self.EXPECTED_RATIO - and not self.invalid_utf16le - ) - - def validate_utf32_characters(self, quad: List[int]) -> None: - """ - Validate if the quad of bytes is valid UTF-32. - - UTF-32 is valid in the range 0x00000000 - 0x0010FFFF - excluding 0x0000D800 - 0x0000DFFF - - https://en.wikipedia.org/wiki/UTF-32 - """ - if ( - quad[0] != 0 - or quad[1] > 0x10 - or (quad[0] == 0 and quad[1] == 0 and 0xD8 <= quad[2] <= 0xDF) - ): - self.invalid_utf32be = True - if ( - quad[3] != 0 - or quad[2] > 0x10 - or (quad[3] == 0 and quad[2] == 0 and 0xD8 <= quad[1] <= 0xDF) - ): - self.invalid_utf32le = True - - def validate_utf16_characters(self, pair: List[int]) -> None: - """ - Validate if the pair of bytes is valid UTF-16. - - UTF-16 is valid in the range 0x0000 - 0xFFFF excluding 0xD800 - 0xFFFF - with an exception for surrogate pairs, which must be in the range - 0xD800-0xDBFF followed by 0xDC00-0xDFFF - - https://en.wikipedia.org/wiki/UTF-16 - """ - if not self.first_half_surrogate_pair_detected_16be: - if 0xD8 <= pair[0] <= 0xDB: - self.first_half_surrogate_pair_detected_16be = True - elif 0xDC <= pair[0] <= 0xDF: - self.invalid_utf16be = True - else: - if 0xDC <= pair[0] <= 0xDF: - self.first_half_surrogate_pair_detected_16be = False - else: - self.invalid_utf16be = True - - if not self.first_half_surrogate_pair_detected_16le: - if 0xD8 <= pair[1] <= 0xDB: - self.first_half_surrogate_pair_detected_16le = True - elif 0xDC <= pair[1] <= 0xDF: - self.invalid_utf16le = True - else: - if 0xDC <= pair[1] <= 0xDF: - self.first_half_surrogate_pair_detected_16le = False - else: - self.invalid_utf16le = True - - def feed(self, byte_str: Union[bytes, bytearray]) -> ProbingState: - for c in byte_str: - mod4 = self.position % 4 - self.quad[mod4] = c - if mod4 == 3: - self.validate_utf32_characters(self.quad) - self.validate_utf16_characters(self.quad[0:2]) - self.validate_utf16_characters(self.quad[2:4]) - if c == 0: - self.zeros_at_mod[mod4] += 1 - else: - self.nonzeros_at_mod[mod4] += 1 - self.position += 1 - return self.state - - @property - def state(self) -> ProbingState: - if self._state in {ProbingState.NOT_ME, ProbingState.FOUND_IT}: - # terminal, decided states - return self._state - if self.get_confidence() > 0.80: - self._state = ProbingState.FOUND_IT - elif self.position > 4 * 1024: - # if we get to 4kb into the file, and we can't conclude it's UTF, - # let's give up - self._state = ProbingState.NOT_ME - return self._state - - def get_confidence(self) -> float: - return ( - 0.85 - if ( - self.is_likely_utf16le() - or self.is_likely_utf16be() - or self.is_likely_utf32le() - or self.is_likely_utf32be() - ) - else 0.00 - ) diff --git a/venv/lib/python3.10/site-packages/chardet/utf8prober.py b/venv/lib/python3.10/site-packages/chardet/utf8prober.py deleted file mode 100644 index d96354d97c2195320d0acc1717a5876eafbea2af..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/utf8prober.py +++ /dev/null @@ -1,82 +0,0 @@ -######################## BEGIN LICENSE BLOCK ######################## -# The Original Code is mozilla.org code. -# -# The Initial Developer of the Original Code is -# Netscape Communications Corporation. -# Portions created by the Initial Developer are Copyright (C) 1998 -# the Initial Developer. All Rights Reserved. -# -# Contributor(s): -# Mark Pilgrim - port to Python -# -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA -# 02110-1301 USA -######################### END LICENSE BLOCK ######################### - -from typing import Union - -from .charsetprober import CharSetProber -from .codingstatemachine import CodingStateMachine -from .enums import MachineState, ProbingState -from .mbcssm import UTF8_SM_MODEL - - -class UTF8Prober(CharSetProber): - ONE_CHAR_PROB = 0.5 - - def __init__(self) -> None: - super().__init__() - self.coding_sm = CodingStateMachine(UTF8_SM_MODEL) - self._num_mb_chars = 0 - self.reset() - - def reset(self) -> None: - super().reset() - self.coding_sm.reset() - self._num_mb_chars = 0 - - @property - def charset_name(self) -> str: - return "utf-8" - - @property - def language(self) -> str: - return "" - - def feed(self, byte_str: Union[bytes, bytearray]) -> ProbingState: - for c in byte_str: - coding_state = self.coding_sm.next_state(c) - if coding_state == MachineState.ERROR: - self._state = ProbingState.NOT_ME - break - if coding_state == MachineState.ITS_ME: - self._state = ProbingState.FOUND_IT - break - if coding_state == MachineState.START: - if self.coding_sm.get_current_charlen() >= 2: - self._num_mb_chars += 1 - - if self.state == ProbingState.DETECTING: - if self.get_confidence() > self.SHORTCUT_THRESHOLD: - self._state = ProbingState.FOUND_IT - - return self.state - - def get_confidence(self) -> float: - unlike = 0.99 - if self._num_mb_chars < 6: - unlike *= self.ONE_CHAR_PROB**self._num_mb_chars - return 1.0 - unlike - return unlike diff --git a/venv/lib/python3.10/site-packages/chardet/version.py b/venv/lib/python3.10/site-packages/chardet/version.py deleted file mode 100644 index 19dd01e030156d04171bac727474595fba305864..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/chardet/version.py +++ /dev/null @@ -1,9 +0,0 @@ -""" -This module exists only to simplify retrieving the version number of chardet -from within setuptools and from chardet subpackages. - -:author: Dan Blanchard (dan.blanchard@gmail.com) -""" - -__version__ = "5.2.0" -VERSION = __version__.split(".") diff --git a/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/METADATA b/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/METADATA deleted file mode 100644 index e6c05af4a78cabdf5689f3f6f7fd57304d8a0eb1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/METADATA +++ /dev/null @@ -1,82 +0,0 @@ -Metadata-Version: 2.4 -Name: click -Version: 8.2.1 -Summary: Composable command line interface toolkit -Maintainer-email: Pallets -Requires-Python: >=3.10 -Description-Content-Type: text/markdown -License-Expression: BSD-3-Clause -Classifier: Development Status :: 5 - Production/Stable -Classifier: Intended Audience :: Developers -Classifier: Operating System :: OS Independent -Classifier: Programming Language :: Python -Classifier: Typing :: Typed -License-File: LICENSE.txt -Requires-Dist: colorama; platform_system == 'Windows' -Project-URL: Changes, https://click.palletsprojects.com/page/changes/ -Project-URL: Chat, https://discord.gg/pallets -Project-URL: Documentation, https://click.palletsprojects.com/ -Project-URL: Donate, https://palletsprojects.com/donate -Project-URL: Source, https://github.com/pallets/click/ - -# $ click_ - -Click is a Python package for creating beautiful command line interfaces -in a composable way with as little code as necessary. It's the "Command -Line Interface Creation Kit". It's highly configurable but comes with -sensible defaults out of the box. - -It aims to make the process of writing command line tools quick and fun -while also preventing any frustration caused by the inability to -implement an intended CLI API. - -Click in three points: - -- Arbitrary nesting of commands -- Automatic help page generation -- Supports lazy loading of subcommands at runtime - - -## A Simple Example - -```python -import click - -@click.command() -@click.option("--count", default=1, help="Number of greetings.") -@click.option("--name", prompt="Your name", help="The person to greet.") -def hello(count, name): - """Simple program that greets NAME for a total of COUNT times.""" - for _ in range(count): - click.echo(f"Hello, {name}!") - -if __name__ == '__main__': - hello() -``` - -``` -$ python hello.py --count=3 -Your name: Click -Hello, Click! -Hello, Click! -Hello, Click! -``` - - -## Donate - -The Pallets organization develops and supports Click and other popular -packages. In order to grow the community of contributors and users, and -allow the maintainers to devote more time to the projects, [please -donate today][]. - -[please donate today]: https://palletsprojects.com/donate - -## Contributing - -See our [detailed contributing documentation][contrib] for many ways to -contribute, including reporting issues, requesting features, asking or answering -questions, and making PRs. - -[contrib]: https://palletsprojects.com/contributing/ - diff --git a/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/RECORD b/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/RECORD deleted file mode 100644 index 621f456cf0d0bd3be7c24c3f098842cf5370131b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/RECORD +++ /dev/null @@ -1,38 +0,0 @@ -click-8.2.1.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -click-8.2.1.dist-info/METADATA,sha256=dI1MbhHTLoKD2tNCCGnx9rK2gok23HDNylFeLKdLSik,2471 -click-8.2.1.dist-info/RECORD,, -click-8.2.1.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82 -click-8.2.1.dist-info/licenses/LICENSE.txt,sha256=morRBqOU6FO_4h9C9OctWSgZoigF2ZG18ydQKSkrZY0,1475 -click/__init__.py,sha256=6YyS1aeyknZ0LYweWozNZy0A9nZ_11wmYIhv3cbQrYo,4473 -click/__pycache__/__init__.cpython-310.pyc,, -click/__pycache__/_compat.cpython-310.pyc,, -click/__pycache__/_termui_impl.cpython-310.pyc,, -click/__pycache__/_textwrap.cpython-310.pyc,, -click/__pycache__/_winconsole.cpython-310.pyc,, -click/__pycache__/core.cpython-310.pyc,, -click/__pycache__/decorators.cpython-310.pyc,, -click/__pycache__/exceptions.cpython-310.pyc,, -click/__pycache__/formatting.cpython-310.pyc,, -click/__pycache__/globals.cpython-310.pyc,, -click/__pycache__/parser.cpython-310.pyc,, -click/__pycache__/shell_completion.cpython-310.pyc,, -click/__pycache__/termui.cpython-310.pyc,, -click/__pycache__/testing.cpython-310.pyc,, -click/__pycache__/types.cpython-310.pyc,, -click/__pycache__/utils.cpython-310.pyc,, -click/_compat.py,sha256=v3xBZkFbvA1BXPRkFfBJc6-pIwPI7345m-kQEnpVAs4,18693 -click/_termui_impl.py,sha256=ASXhLi9IQIc0Js9KQSS-3-SLZcPet3VqysBf9WgbbpI,26712 -click/_textwrap.py,sha256=BOae0RQ6vg3FkNgSJyOoGzG1meGMxJ_ukWVZKx_v-0o,1400 -click/_winconsole.py,sha256=_vxUuUaxwBhoR0vUWCNuHY8VUefiMdCIyU2SXPqoF-A,8465 -click/core.py,sha256=gUhpNS9cFBGdEXXdisGVG-eRvGf49RTyFagxulqwdFw,117343 -click/decorators.py,sha256=5P7abhJtAQYp_KHgjUvhMv464ERwOzrv2enNknlwHyQ,18461 -click/exceptions.py,sha256=1rdtXgHJ1b3OjGkN-UpXB9t_HCBihJvh_DtpmLmwn9s,9891 -click/formatting.py,sha256=Bhqx4QXdKQ9W4WKknIwj5KPKFmtduGOuGq1yw_THLZ8,9726 -click/globals.py,sha256=gM-Nh6A4M0HB_SgkaF5M4ncGGMDHc_flHXu9_oh4GEU,1923 -click/parser.py,sha256=nU1Ah2p11q29ul1vNdU9swPo_PUuKrxU6YXToi71q1c,18979 -click/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -click/shell_completion.py,sha256=CQSGdjgun4ORbOZrXP0CVhEtPx4knsufOkRsDiK64cM,19857 -click/termui.py,sha256=vAYrKC2a7f_NfEIhAThEVYfa__ib5XQbTSCGtJlABRA,30847 -click/testing.py,sha256=2eLdAaCJCGToP5Tw-XN8JjrDb3wbJIfARxg3d0crW5M,18702 -click/types.py,sha256=KBTRxN28cR1VZ5mb9iJX98MQSw_p9SGzljqfEI8z5Tw,38389 -click/utils.py,sha256=b1Mm-usEDBHtEwcPltPIn3zSK4nw2KTp5GC7_oSTlLo,20245 diff --git a/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/WHEEL b/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/WHEEL deleted file mode 100644 index d8b9936dad9ab2513fa6979f411560d3b6b57e37..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/WHEEL +++ /dev/null @@ -1,4 +0,0 @@ -Wheel-Version: 1.0 -Generator: flit 3.12.0 -Root-Is-Purelib: true -Tag: py3-none-any diff --git a/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/licenses/LICENSE.txt b/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/licenses/LICENSE.txt deleted file mode 100644 index d12a849186982399c537c5b9a8fd77bf2edd5eab..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click-8.2.1.dist-info/licenses/LICENSE.txt +++ /dev/null @@ -1,28 +0,0 @@ -Copyright 2014 Pallets - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are -met: - -1. Redistributions of source code must retain the above copyright - notice, this list of conditions and the following disclaimer. - -2. Redistributions in binary form must reproduce the above copyright - notice, this list of conditions and the following disclaimer in the - documentation and/or other materials provided with the distribution. - -3. Neither the name of the copyright holder nor the names of its - contributors may be used to endorse or promote products derived from - this software without specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS -"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT -LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A -PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT -HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, -SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED -TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR -PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF -LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING -NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS -SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/venv/lib/python3.10/site-packages/click/__init__.py b/venv/lib/python3.10/site-packages/click/__init__.py deleted file mode 100644 index 1aa547c57af68954b3bcd6012873881030e53d91..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/__init__.py +++ /dev/null @@ -1,123 +0,0 @@ -""" -Click is a simple Python module inspired by the stdlib optparse to make -writing command line scripts fun. Unlike other modules, it's based -around a simple API that does not come with too much magic and is -composable. -""" - -from __future__ import annotations - -from .core import Argument as Argument -from .core import Command as Command -from .core import CommandCollection as CommandCollection -from .core import Context as Context -from .core import Group as Group -from .core import Option as Option -from .core import Parameter as Parameter -from .decorators import argument as argument -from .decorators import command as command -from .decorators import confirmation_option as confirmation_option -from .decorators import group as group -from .decorators import help_option as help_option -from .decorators import make_pass_decorator as make_pass_decorator -from .decorators import option as option -from .decorators import pass_context as pass_context -from .decorators import pass_obj as pass_obj -from .decorators import password_option as password_option -from .decorators import version_option as version_option -from .exceptions import Abort as Abort -from .exceptions import BadArgumentUsage as BadArgumentUsage -from .exceptions import BadOptionUsage as BadOptionUsage -from .exceptions import BadParameter as BadParameter -from .exceptions import ClickException as ClickException -from .exceptions import FileError as FileError -from .exceptions import MissingParameter as MissingParameter -from .exceptions import NoSuchOption as NoSuchOption -from .exceptions import UsageError as UsageError -from .formatting import HelpFormatter as HelpFormatter -from .formatting import wrap_text as wrap_text -from .globals import get_current_context as get_current_context -from .termui import clear as clear -from .termui import confirm as confirm -from .termui import echo_via_pager as echo_via_pager -from .termui import edit as edit -from .termui import getchar as getchar -from .termui import launch as launch -from .termui import pause as pause -from .termui import progressbar as progressbar -from .termui import prompt as prompt -from .termui import secho as secho -from .termui import style as style -from .termui import unstyle as unstyle -from .types import BOOL as BOOL -from .types import Choice as Choice -from .types import DateTime as DateTime -from .types import File as File -from .types import FLOAT as FLOAT -from .types import FloatRange as FloatRange -from .types import INT as INT -from .types import IntRange as IntRange -from .types import ParamType as ParamType -from .types import Path as Path -from .types import STRING as STRING -from .types import Tuple as Tuple -from .types import UNPROCESSED as UNPROCESSED -from .types import UUID as UUID -from .utils import echo as echo -from .utils import format_filename as format_filename -from .utils import get_app_dir as get_app_dir -from .utils import get_binary_stream as get_binary_stream -from .utils import get_text_stream as get_text_stream -from .utils import open_file as open_file - - -def __getattr__(name: str) -> object: - import warnings - - if name == "BaseCommand": - from .core import _BaseCommand - - warnings.warn( - "'BaseCommand' is deprecated and will be removed in Click 9.0. Use" - " 'Command' instead.", - DeprecationWarning, - stacklevel=2, - ) - return _BaseCommand - - if name == "MultiCommand": - from .core import _MultiCommand - - warnings.warn( - "'MultiCommand' is deprecated and will be removed in Click 9.0. Use" - " 'Group' instead.", - DeprecationWarning, - stacklevel=2, - ) - return _MultiCommand - - if name == "OptionParser": - from .parser import _OptionParser - - warnings.warn( - "'OptionParser' is deprecated and will be removed in Click 9.0. The" - " old parser is available in 'optparse'.", - DeprecationWarning, - stacklevel=2, - ) - return _OptionParser - - if name == "__version__": - import importlib.metadata - import warnings - - warnings.warn( - "The '__version__' attribute is deprecated and will be removed in" - " Click 9.1. Use feature detection or" - " 'importlib.metadata.version(\"click\")' instead.", - DeprecationWarning, - stacklevel=2, - ) - return importlib.metadata.version("click") - - raise AttributeError(name) diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 56deca6c2a4964d8c09a5187b958ee3361a1208b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/_compat.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/_compat.cpython-310.pyc deleted file mode 100644 index 0b1b56337bd94a9401b3141197ad029e90725a80..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/_compat.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/_termui_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/_termui_impl.cpython-310.pyc deleted file mode 100644 index 57b5d1b1f8b5ca16d9cb968af09c385a7de6d418..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/_termui_impl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/_textwrap.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/_textwrap.cpython-310.pyc deleted file mode 100644 index 750a8e860ca315e87677c9fc740d7fcc353fec51..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/_textwrap.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/_winconsole.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/_winconsole.cpython-310.pyc deleted file mode 100644 index a29968f590de5ae73c38b28ccb3a53b4b9aa1867..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/_winconsole.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/core.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/core.cpython-310.pyc deleted file mode 100644 index 8e40144d2f524c664a4863d8ae09d207cf5c4e30..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/core.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/decorators.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/decorators.cpython-310.pyc deleted file mode 100644 index 721600162edbd092d703a536222b7d01edcfe425..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/decorators.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/exceptions.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/exceptions.cpython-310.pyc deleted file mode 100644 index 4b3cbbe5513b06433f2553ee08384e6a00b5f93e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/exceptions.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/formatting.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/formatting.cpython-310.pyc deleted file mode 100644 index 7174c4810af49dc44860873d371102d619464a16..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/formatting.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/globals.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/globals.cpython-310.pyc deleted file mode 100644 index 3de18fcfeb1caad83a5d1603efb76b12cf6a1f3f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/globals.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/parser.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/parser.cpython-310.pyc deleted file mode 100644 index 3c518a98136bd2cc49205777046e15a07a4448ff..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/parser.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/shell_completion.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/shell_completion.cpython-310.pyc deleted file mode 100644 index f9417b6e8413df3d7dcad595ef6773b0a616c30c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/shell_completion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/termui.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/termui.cpython-310.pyc deleted file mode 100644 index e8983688b073551d7e698b7b2c8fb7f4c7d26a6e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/termui.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/testing.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/testing.cpython-310.pyc deleted file mode 100644 index 1873cafa1bcbbb6151529f81facbe91fbf4e1035..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/testing.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/types.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/types.cpython-310.pyc deleted file mode 100644 index 32e8a2412bb9832452be9d81bffd5991a3217fb5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/types.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/click/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index 5819807d449d6abf9402307b124a3636e5594fc8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/click/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/click/_compat.py b/venv/lib/python3.10/site-packages/click/_compat.py deleted file mode 100644 index f2726b93afe2ed69ba20d8e709c9cbd22b9d2339..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/_compat.py +++ /dev/null @@ -1,622 +0,0 @@ -from __future__ import annotations - -import codecs -import collections.abc as cabc -import io -import os -import re -import sys -import typing as t -from types import TracebackType -from weakref import WeakKeyDictionary - -CYGWIN = sys.platform.startswith("cygwin") -WIN = sys.platform.startswith("win") -auto_wrap_for_ansi: t.Callable[[t.TextIO], t.TextIO] | None = None -_ansi_re = re.compile(r"\033\[[;?0-9]*[a-zA-Z]") - - -def _make_text_stream( - stream: t.BinaryIO, - encoding: str | None, - errors: str | None, - force_readable: bool = False, - force_writable: bool = False, -) -> t.TextIO: - if encoding is None: - encoding = get_best_encoding(stream) - if errors is None: - errors = "replace" - return _NonClosingTextIOWrapper( - stream, - encoding, - errors, - line_buffering=True, - force_readable=force_readable, - force_writable=force_writable, - ) - - -def is_ascii_encoding(encoding: str) -> bool: - """Checks if a given encoding is ascii.""" - try: - return codecs.lookup(encoding).name == "ascii" - except LookupError: - return False - - -def get_best_encoding(stream: t.IO[t.Any]) -> str: - """Returns the default stream encoding if not found.""" - rv = getattr(stream, "encoding", None) or sys.getdefaultencoding() - if is_ascii_encoding(rv): - return "utf-8" - return rv - - -class _NonClosingTextIOWrapper(io.TextIOWrapper): - def __init__( - self, - stream: t.BinaryIO, - encoding: str | None, - errors: str | None, - force_readable: bool = False, - force_writable: bool = False, - **extra: t.Any, - ) -> None: - self._stream = stream = t.cast( - t.BinaryIO, _FixupStream(stream, force_readable, force_writable) - ) - super().__init__(stream, encoding, errors, **extra) - - def __del__(self) -> None: - try: - self.detach() - except Exception: - pass - - def isatty(self) -> bool: - # https://bitbucket.org/pypy/pypy/issue/1803 - return self._stream.isatty() - - -class _FixupStream: - """The new io interface needs more from streams than streams - traditionally implement. As such, this fix-up code is necessary in - some circumstances. - - The forcing of readable and writable flags are there because some tools - put badly patched objects on sys (one such offender are certain version - of jupyter notebook). - """ - - def __init__( - self, - stream: t.BinaryIO, - force_readable: bool = False, - force_writable: bool = False, - ): - self._stream = stream - self._force_readable = force_readable - self._force_writable = force_writable - - def __getattr__(self, name: str) -> t.Any: - return getattr(self._stream, name) - - def read1(self, size: int) -> bytes: - f = getattr(self._stream, "read1", None) - - if f is not None: - return t.cast(bytes, f(size)) - - return self._stream.read(size) - - def readable(self) -> bool: - if self._force_readable: - return True - x = getattr(self._stream, "readable", None) - if x is not None: - return t.cast(bool, x()) - try: - self._stream.read(0) - except Exception: - return False - return True - - def writable(self) -> bool: - if self._force_writable: - return True - x = getattr(self._stream, "writable", None) - if x is not None: - return t.cast(bool, x()) - try: - self._stream.write(b"") - except Exception: - try: - self._stream.write(b"") - except Exception: - return False - return True - - def seekable(self) -> bool: - x = getattr(self._stream, "seekable", None) - if x is not None: - return t.cast(bool, x()) - try: - self._stream.seek(self._stream.tell()) - except Exception: - return False - return True - - -def _is_binary_reader(stream: t.IO[t.Any], default: bool = False) -> bool: - try: - return isinstance(stream.read(0), bytes) - except Exception: - return default - # This happens in some cases where the stream was already - # closed. In this case, we assume the default. - - -def _is_binary_writer(stream: t.IO[t.Any], default: bool = False) -> bool: - try: - stream.write(b"") - except Exception: - try: - stream.write("") - return False - except Exception: - pass - return default - return True - - -def _find_binary_reader(stream: t.IO[t.Any]) -> t.BinaryIO | None: - # We need to figure out if the given stream is already binary. - # This can happen because the official docs recommend detaching - # the streams to get binary streams. Some code might do this, so - # we need to deal with this case explicitly. - if _is_binary_reader(stream, False): - return t.cast(t.BinaryIO, stream) - - buf = getattr(stream, "buffer", None) - - # Same situation here; this time we assume that the buffer is - # actually binary in case it's closed. - if buf is not None and _is_binary_reader(buf, True): - return t.cast(t.BinaryIO, buf) - - return None - - -def _find_binary_writer(stream: t.IO[t.Any]) -> t.BinaryIO | None: - # We need to figure out if the given stream is already binary. - # This can happen because the official docs recommend detaching - # the streams to get binary streams. Some code might do this, so - # we need to deal with this case explicitly. - if _is_binary_writer(stream, False): - return t.cast(t.BinaryIO, stream) - - buf = getattr(stream, "buffer", None) - - # Same situation here; this time we assume that the buffer is - # actually binary in case it's closed. - if buf is not None and _is_binary_writer(buf, True): - return t.cast(t.BinaryIO, buf) - - return None - - -def _stream_is_misconfigured(stream: t.TextIO) -> bool: - """A stream is misconfigured if its encoding is ASCII.""" - # If the stream does not have an encoding set, we assume it's set - # to ASCII. This appears to happen in certain unittest - # environments. It's not quite clear what the correct behavior is - # but this at least will force Click to recover somehow. - return is_ascii_encoding(getattr(stream, "encoding", None) or "ascii") - - -def _is_compat_stream_attr(stream: t.TextIO, attr: str, value: str | None) -> bool: - """A stream attribute is compatible if it is equal to the - desired value or the desired value is unset and the attribute - has a value. - """ - stream_value = getattr(stream, attr, None) - return stream_value == value or (value is None and stream_value is not None) - - -def _is_compatible_text_stream( - stream: t.TextIO, encoding: str | None, errors: str | None -) -> bool: - """Check if a stream's encoding and errors attributes are - compatible with the desired values. - """ - return _is_compat_stream_attr( - stream, "encoding", encoding - ) and _is_compat_stream_attr(stream, "errors", errors) - - -def _force_correct_text_stream( - text_stream: t.IO[t.Any], - encoding: str | None, - errors: str | None, - is_binary: t.Callable[[t.IO[t.Any], bool], bool], - find_binary: t.Callable[[t.IO[t.Any]], t.BinaryIO | None], - force_readable: bool = False, - force_writable: bool = False, -) -> t.TextIO: - if is_binary(text_stream, False): - binary_reader = t.cast(t.BinaryIO, text_stream) - else: - text_stream = t.cast(t.TextIO, text_stream) - # If the stream looks compatible, and won't default to a - # misconfigured ascii encoding, return it as-is. - if _is_compatible_text_stream(text_stream, encoding, errors) and not ( - encoding is None and _stream_is_misconfigured(text_stream) - ): - return text_stream - - # Otherwise, get the underlying binary reader. - possible_binary_reader = find_binary(text_stream) - - # If that's not possible, silently use the original reader - # and get mojibake instead of exceptions. - if possible_binary_reader is None: - return text_stream - - binary_reader = possible_binary_reader - - # Default errors to replace instead of strict in order to get - # something that works. - if errors is None: - errors = "replace" - - # Wrap the binary stream in a text stream with the correct - # encoding parameters. - return _make_text_stream( - binary_reader, - encoding, - errors, - force_readable=force_readable, - force_writable=force_writable, - ) - - -def _force_correct_text_reader( - text_reader: t.IO[t.Any], - encoding: str | None, - errors: str | None, - force_readable: bool = False, -) -> t.TextIO: - return _force_correct_text_stream( - text_reader, - encoding, - errors, - _is_binary_reader, - _find_binary_reader, - force_readable=force_readable, - ) - - -def _force_correct_text_writer( - text_writer: t.IO[t.Any], - encoding: str | None, - errors: str | None, - force_writable: bool = False, -) -> t.TextIO: - return _force_correct_text_stream( - text_writer, - encoding, - errors, - _is_binary_writer, - _find_binary_writer, - force_writable=force_writable, - ) - - -def get_binary_stdin() -> t.BinaryIO: - reader = _find_binary_reader(sys.stdin) - if reader is None: - raise RuntimeError("Was not able to determine binary stream for sys.stdin.") - return reader - - -def get_binary_stdout() -> t.BinaryIO: - writer = _find_binary_writer(sys.stdout) - if writer is None: - raise RuntimeError("Was not able to determine binary stream for sys.stdout.") - return writer - - -def get_binary_stderr() -> t.BinaryIO: - writer = _find_binary_writer(sys.stderr) - if writer is None: - raise RuntimeError("Was not able to determine binary stream for sys.stderr.") - return writer - - -def get_text_stdin(encoding: str | None = None, errors: str | None = None) -> t.TextIO: - rv = _get_windows_console_stream(sys.stdin, encoding, errors) - if rv is not None: - return rv - return _force_correct_text_reader(sys.stdin, encoding, errors, force_readable=True) - - -def get_text_stdout(encoding: str | None = None, errors: str | None = None) -> t.TextIO: - rv = _get_windows_console_stream(sys.stdout, encoding, errors) - if rv is not None: - return rv - return _force_correct_text_writer(sys.stdout, encoding, errors, force_writable=True) - - -def get_text_stderr(encoding: str | None = None, errors: str | None = None) -> t.TextIO: - rv = _get_windows_console_stream(sys.stderr, encoding, errors) - if rv is not None: - return rv - return _force_correct_text_writer(sys.stderr, encoding, errors, force_writable=True) - - -def _wrap_io_open( - file: str | os.PathLike[str] | int, - mode: str, - encoding: str | None, - errors: str | None, -) -> t.IO[t.Any]: - """Handles not passing ``encoding`` and ``errors`` in binary mode.""" - if "b" in mode: - return open(file, mode) - - return open(file, mode, encoding=encoding, errors=errors) - - -def open_stream( - filename: str | os.PathLike[str], - mode: str = "r", - encoding: str | None = None, - errors: str | None = "strict", - atomic: bool = False, -) -> tuple[t.IO[t.Any], bool]: - binary = "b" in mode - filename = os.fspath(filename) - - # Standard streams first. These are simple because they ignore the - # atomic flag. Use fsdecode to handle Path("-"). - if os.fsdecode(filename) == "-": - if any(m in mode for m in ["w", "a", "x"]): - if binary: - return get_binary_stdout(), False - return get_text_stdout(encoding=encoding, errors=errors), False - if binary: - return get_binary_stdin(), False - return get_text_stdin(encoding=encoding, errors=errors), False - - # Non-atomic writes directly go out through the regular open functions. - if not atomic: - return _wrap_io_open(filename, mode, encoding, errors), True - - # Some usability stuff for atomic writes - if "a" in mode: - raise ValueError( - "Appending to an existing file is not supported, because that" - " would involve an expensive `copy`-operation to a temporary" - " file. Open the file in normal `w`-mode and copy explicitly" - " if that's what you're after." - ) - if "x" in mode: - raise ValueError("Use the `overwrite`-parameter instead.") - if "w" not in mode: - raise ValueError("Atomic writes only make sense with `w`-mode.") - - # Atomic writes are more complicated. They work by opening a file - # as a proxy in the same folder and then using the fdopen - # functionality to wrap it in a Python file. Then we wrap it in an - # atomic file that moves the file over on close. - import errno - import random - - try: - perm: int | None = os.stat(filename).st_mode - except OSError: - perm = None - - flags = os.O_RDWR | os.O_CREAT | os.O_EXCL - - if binary: - flags |= getattr(os, "O_BINARY", 0) - - while True: - tmp_filename = os.path.join( - os.path.dirname(filename), - f".__atomic-write{random.randrange(1 << 32):08x}", - ) - try: - fd = os.open(tmp_filename, flags, 0o666 if perm is None else perm) - break - except OSError as e: - if e.errno == errno.EEXIST or ( - os.name == "nt" - and e.errno == errno.EACCES - and os.path.isdir(e.filename) - and os.access(e.filename, os.W_OK) - ): - continue - raise - - if perm is not None: - os.chmod(tmp_filename, perm) # in case perm includes bits in umask - - f = _wrap_io_open(fd, mode, encoding, errors) - af = _AtomicFile(f, tmp_filename, os.path.realpath(filename)) - return t.cast(t.IO[t.Any], af), True - - -class _AtomicFile: - def __init__(self, f: t.IO[t.Any], tmp_filename: str, real_filename: str) -> None: - self._f = f - self._tmp_filename = tmp_filename - self._real_filename = real_filename - self.closed = False - - @property - def name(self) -> str: - return self._real_filename - - def close(self, delete: bool = False) -> None: - if self.closed: - return - self._f.close() - os.replace(self._tmp_filename, self._real_filename) - self.closed = True - - def __getattr__(self, name: str) -> t.Any: - return getattr(self._f, name) - - def __enter__(self) -> _AtomicFile: - return self - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc_value: BaseException | None, - tb: TracebackType | None, - ) -> None: - self.close(delete=exc_type is not None) - - def __repr__(self) -> str: - return repr(self._f) - - -def strip_ansi(value: str) -> str: - return _ansi_re.sub("", value) - - -def _is_jupyter_kernel_output(stream: t.IO[t.Any]) -> bool: - while isinstance(stream, (_FixupStream, _NonClosingTextIOWrapper)): - stream = stream._stream - - return stream.__class__.__module__.startswith("ipykernel.") - - -def should_strip_ansi( - stream: t.IO[t.Any] | None = None, color: bool | None = None -) -> bool: - if color is None: - if stream is None: - stream = sys.stdin - return not isatty(stream) and not _is_jupyter_kernel_output(stream) - return not color - - -# On Windows, wrap the output streams with colorama to support ANSI -# color codes. -# NOTE: double check is needed so mypy does not analyze this on Linux -if sys.platform.startswith("win") and WIN: - from ._winconsole import _get_windows_console_stream - - def _get_argv_encoding() -> str: - import locale - - return locale.getpreferredencoding() - - _ansi_stream_wrappers: cabc.MutableMapping[t.TextIO, t.TextIO] = WeakKeyDictionary() - - def auto_wrap_for_ansi(stream: t.TextIO, color: bool | None = None) -> t.TextIO: - """Support ANSI color and style codes on Windows by wrapping a - stream with colorama. - """ - try: - cached = _ansi_stream_wrappers.get(stream) - except Exception: - cached = None - - if cached is not None: - return cached - - import colorama - - strip = should_strip_ansi(stream, color) - ansi_wrapper = colorama.AnsiToWin32(stream, strip=strip) - rv = t.cast(t.TextIO, ansi_wrapper.stream) - _write = rv.write - - def _safe_write(s: str) -> int: - try: - return _write(s) - except BaseException: - ansi_wrapper.reset_all() - raise - - rv.write = _safe_write # type: ignore[method-assign] - - try: - _ansi_stream_wrappers[stream] = rv - except Exception: - pass - - return rv - -else: - - def _get_argv_encoding() -> str: - return getattr(sys.stdin, "encoding", None) or sys.getfilesystemencoding() - - def _get_windows_console_stream( - f: t.TextIO, encoding: str | None, errors: str | None - ) -> t.TextIO | None: - return None - - -def term_len(x: str) -> int: - return len(strip_ansi(x)) - - -def isatty(stream: t.IO[t.Any]) -> bool: - try: - return stream.isatty() - except Exception: - return False - - -def _make_cached_stream_func( - src_func: t.Callable[[], t.TextIO | None], - wrapper_func: t.Callable[[], t.TextIO], -) -> t.Callable[[], t.TextIO | None]: - cache: cabc.MutableMapping[t.TextIO, t.TextIO] = WeakKeyDictionary() - - def func() -> t.TextIO | None: - stream = src_func() - - if stream is None: - return None - - try: - rv = cache.get(stream) - except Exception: - rv = None - if rv is not None: - return rv - rv = wrapper_func() - try: - cache[stream] = rv - except Exception: - pass - return rv - - return func - - -_default_text_stdin = _make_cached_stream_func(lambda: sys.stdin, get_text_stdin) -_default_text_stdout = _make_cached_stream_func(lambda: sys.stdout, get_text_stdout) -_default_text_stderr = _make_cached_stream_func(lambda: sys.stderr, get_text_stderr) - - -binary_streams: cabc.Mapping[str, t.Callable[[], t.BinaryIO]] = { - "stdin": get_binary_stdin, - "stdout": get_binary_stdout, - "stderr": get_binary_stderr, -} - -text_streams: cabc.Mapping[str, t.Callable[[str | None, str | None], t.TextIO]] = { - "stdin": get_text_stdin, - "stdout": get_text_stdout, - "stderr": get_text_stderr, -} diff --git a/venv/lib/python3.10/site-packages/click/_termui_impl.py b/venv/lib/python3.10/site-packages/click/_termui_impl.py deleted file mode 100644 index 51fd9bf358f0ee8dedceb88efa8c6a60dc40d4a8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/_termui_impl.py +++ /dev/null @@ -1,839 +0,0 @@ -""" -This module contains implementations for the termui module. To keep the -import time of Click down, some infrequently used functionality is -placed in this module and only imported as needed. -""" - -from __future__ import annotations - -import collections.abc as cabc -import contextlib -import math -import os -import shlex -import sys -import time -import typing as t -from gettext import gettext as _ -from io import StringIO -from pathlib import Path -from shutil import which -from types import TracebackType - -from ._compat import _default_text_stdout -from ._compat import CYGWIN -from ._compat import get_best_encoding -from ._compat import isatty -from ._compat import open_stream -from ._compat import strip_ansi -from ._compat import term_len -from ._compat import WIN -from .exceptions import ClickException -from .utils import echo - -V = t.TypeVar("V") - -if os.name == "nt": - BEFORE_BAR = "\r" - AFTER_BAR = "\n" -else: - BEFORE_BAR = "\r\033[?25l" - AFTER_BAR = "\033[?25h\n" - - -class ProgressBar(t.Generic[V]): - def __init__( - self, - iterable: cabc.Iterable[V] | None, - length: int | None = None, - fill_char: str = "#", - empty_char: str = " ", - bar_template: str = "%(bar)s", - info_sep: str = " ", - hidden: bool = False, - show_eta: bool = True, - show_percent: bool | None = None, - show_pos: bool = False, - item_show_func: t.Callable[[V | None], str | None] | None = None, - label: str | None = None, - file: t.TextIO | None = None, - color: bool | None = None, - update_min_steps: int = 1, - width: int = 30, - ) -> None: - self.fill_char = fill_char - self.empty_char = empty_char - self.bar_template = bar_template - self.info_sep = info_sep - self.hidden = hidden - self.show_eta = show_eta - self.show_percent = show_percent - self.show_pos = show_pos - self.item_show_func = item_show_func - self.label: str = label or "" - - if file is None: - file = _default_text_stdout() - - # There are no standard streams attached to write to. For example, - # pythonw on Windows. - if file is None: - file = StringIO() - - self.file = file - self.color = color - self.update_min_steps = update_min_steps - self._completed_intervals = 0 - self.width: int = width - self.autowidth: bool = width == 0 - - if length is None: - from operator import length_hint - - length = length_hint(iterable, -1) - - if length == -1: - length = None - if iterable is None: - if length is None: - raise TypeError("iterable or length is required") - iterable = t.cast("cabc.Iterable[V]", range(length)) - self.iter: cabc.Iterable[V] = iter(iterable) - self.length = length - self.pos: int = 0 - self.avg: list[float] = [] - self.last_eta: float - self.start: float - self.start = self.last_eta = time.time() - self.eta_known: bool = False - self.finished: bool = False - self.max_width: int | None = None - self.entered: bool = False - self.current_item: V | None = None - self._is_atty = isatty(self.file) - self._last_line: str | None = None - - def __enter__(self) -> ProgressBar[V]: - self.entered = True - self.render_progress() - return self - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc_value: BaseException | None, - tb: TracebackType | None, - ) -> None: - self.render_finish() - - def __iter__(self) -> cabc.Iterator[V]: - if not self.entered: - raise RuntimeError("You need to use progress bars in a with block.") - self.render_progress() - return self.generator() - - def __next__(self) -> V: - # Iteration is defined in terms of a generator function, - # returned by iter(self); use that to define next(). This works - # because `self.iter` is an iterable consumed by that generator, - # so it is re-entry safe. Calling `next(self.generator())` - # twice works and does "what you want". - return next(iter(self)) - - def render_finish(self) -> None: - if self.hidden or not self._is_atty: - return - self.file.write(AFTER_BAR) - self.file.flush() - - @property - def pct(self) -> float: - if self.finished: - return 1.0 - return min(self.pos / (float(self.length or 1) or 1), 1.0) - - @property - def time_per_iteration(self) -> float: - if not self.avg: - return 0.0 - return sum(self.avg) / float(len(self.avg)) - - @property - def eta(self) -> float: - if self.length is not None and not self.finished: - return self.time_per_iteration * (self.length - self.pos) - return 0.0 - - def format_eta(self) -> str: - if self.eta_known: - t = int(self.eta) - seconds = t % 60 - t //= 60 - minutes = t % 60 - t //= 60 - hours = t % 24 - t //= 24 - if t > 0: - return f"{t}d {hours:02}:{minutes:02}:{seconds:02}" - else: - return f"{hours:02}:{minutes:02}:{seconds:02}" - return "" - - def format_pos(self) -> str: - pos = str(self.pos) - if self.length is not None: - pos += f"/{self.length}" - return pos - - def format_pct(self) -> str: - return f"{int(self.pct * 100): 4}%"[1:] - - def format_bar(self) -> str: - if self.length is not None: - bar_length = int(self.pct * self.width) - bar = self.fill_char * bar_length - bar += self.empty_char * (self.width - bar_length) - elif self.finished: - bar = self.fill_char * self.width - else: - chars = list(self.empty_char * (self.width or 1)) - if self.time_per_iteration != 0: - chars[ - int( - (math.cos(self.pos * self.time_per_iteration) / 2.0 + 0.5) - * self.width - ) - ] = self.fill_char - bar = "".join(chars) - return bar - - def format_progress_line(self) -> str: - show_percent = self.show_percent - - info_bits = [] - if self.length is not None and show_percent is None: - show_percent = not self.show_pos - - if self.show_pos: - info_bits.append(self.format_pos()) - if show_percent: - info_bits.append(self.format_pct()) - if self.show_eta and self.eta_known and not self.finished: - info_bits.append(self.format_eta()) - if self.item_show_func is not None: - item_info = self.item_show_func(self.current_item) - if item_info is not None: - info_bits.append(item_info) - - return ( - self.bar_template - % { - "label": self.label, - "bar": self.format_bar(), - "info": self.info_sep.join(info_bits), - } - ).rstrip() - - def render_progress(self) -> None: - import shutil - - if self.hidden: - return - - if not self._is_atty: - # Only output the label once if the output is not a TTY. - if self._last_line != self.label: - self._last_line = self.label - echo(self.label, file=self.file, color=self.color) - return - - buf = [] - # Update width in case the terminal has been resized - if self.autowidth: - old_width = self.width - self.width = 0 - clutter_length = term_len(self.format_progress_line()) - new_width = max(0, shutil.get_terminal_size().columns - clutter_length) - if new_width < old_width and self.max_width is not None: - buf.append(BEFORE_BAR) - buf.append(" " * self.max_width) - self.max_width = new_width - self.width = new_width - - clear_width = self.width - if self.max_width is not None: - clear_width = self.max_width - - buf.append(BEFORE_BAR) - line = self.format_progress_line() - line_len = term_len(line) - if self.max_width is None or self.max_width < line_len: - self.max_width = line_len - - buf.append(line) - buf.append(" " * (clear_width - line_len)) - line = "".join(buf) - # Render the line only if it changed. - - if line != self._last_line: - self._last_line = line - echo(line, file=self.file, color=self.color, nl=False) - self.file.flush() - - def make_step(self, n_steps: int) -> None: - self.pos += n_steps - if self.length is not None and self.pos >= self.length: - self.finished = True - - if (time.time() - self.last_eta) < 1.0: - return - - self.last_eta = time.time() - - # self.avg is a rolling list of length <= 7 of steps where steps are - # defined as time elapsed divided by the total progress through - # self.length. - if self.pos: - step = (time.time() - self.start) / self.pos - else: - step = time.time() - self.start - - self.avg = self.avg[-6:] + [step] - - self.eta_known = self.length is not None - - def update(self, n_steps: int, current_item: V | None = None) -> None: - """Update the progress bar by advancing a specified number of - steps, and optionally set the ``current_item`` for this new - position. - - :param n_steps: Number of steps to advance. - :param current_item: Optional item to set as ``current_item`` - for the updated position. - - .. versionchanged:: 8.0 - Added the ``current_item`` optional parameter. - - .. versionchanged:: 8.0 - Only render when the number of steps meets the - ``update_min_steps`` threshold. - """ - if current_item is not None: - self.current_item = current_item - - self._completed_intervals += n_steps - - if self._completed_intervals >= self.update_min_steps: - self.make_step(self._completed_intervals) - self.render_progress() - self._completed_intervals = 0 - - def finish(self) -> None: - self.eta_known = False - self.current_item = None - self.finished = True - - def generator(self) -> cabc.Iterator[V]: - """Return a generator which yields the items added to the bar - during construction, and updates the progress bar *after* the - yielded block returns. - """ - # WARNING: the iterator interface for `ProgressBar` relies on - # this and only works because this is a simple generator which - # doesn't create or manage additional state. If this function - # changes, the impact should be evaluated both against - # `iter(bar)` and `next(bar)`. `next()` in particular may call - # `self.generator()` repeatedly, and this must remain safe in - # order for that interface to work. - if not self.entered: - raise RuntimeError("You need to use progress bars in a with block.") - - if not self._is_atty: - yield from self.iter - else: - for rv in self.iter: - self.current_item = rv - - # This allows show_item_func to be updated before the - # item is processed. Only trigger at the beginning of - # the update interval. - if self._completed_intervals == 0: - self.render_progress() - - yield rv - self.update(1) - - self.finish() - self.render_progress() - - -def pager(generator: cabc.Iterable[str], color: bool | None = None) -> None: - """Decide what method to use for paging through text.""" - stdout = _default_text_stdout() - - # There are no standard streams attached to write to. For example, - # pythonw on Windows. - if stdout is None: - stdout = StringIO() - - if not isatty(sys.stdin) or not isatty(stdout): - return _nullpager(stdout, generator, color) - - # Split and normalize the pager command into parts. - pager_cmd_parts = shlex.split(os.environ.get("PAGER", ""), posix=False) - if pager_cmd_parts: - if WIN: - if _tempfilepager(generator, pager_cmd_parts, color): - return - elif _pipepager(generator, pager_cmd_parts, color): - return - - if os.environ.get("TERM") in ("dumb", "emacs"): - return _nullpager(stdout, generator, color) - if (WIN or sys.platform.startswith("os2")) and _tempfilepager( - generator, ["more"], color - ): - return - if _pipepager(generator, ["less"], color): - return - - import tempfile - - fd, filename = tempfile.mkstemp() - os.close(fd) - try: - if _pipepager(generator, ["more"], color): - return - return _nullpager(stdout, generator, color) - finally: - os.unlink(filename) - - -def _pipepager( - generator: cabc.Iterable[str], cmd_parts: list[str], color: bool | None -) -> bool: - """Page through text by feeding it to another program. Invoking a - pager through this might support colors. - - Returns `True` if the command was found, `False` otherwise and thus another - pager should be attempted. - """ - # Split the command into the invoked CLI and its parameters. - if not cmd_parts: - return False - cmd = cmd_parts[0] - cmd_params = cmd_parts[1:] - - cmd_filepath = which(cmd) - if not cmd_filepath: - return False - # Resolves symlinks and produces a normalized absolute path string. - cmd_path = Path(cmd_filepath).resolve() - cmd_name = cmd_path.name - - import subprocess - - # Make a local copy of the environment to not affect the global one. - env = dict(os.environ) - - # If we're piping to less and the user hasn't decided on colors, we enable - # them by default we find the -R flag in the command line arguments. - if color is None and cmd_name == "less": - less_flags = f"{os.environ.get('LESS', '')}{' '.join(cmd_params)}" - if not less_flags: - env["LESS"] = "-R" - color = True - elif "r" in less_flags or "R" in less_flags: - color = True - - c = subprocess.Popen( - [str(cmd_path)] + cmd_params, - shell=True, - stdin=subprocess.PIPE, - env=env, - errors="replace", - text=True, - ) - assert c.stdin is not None - try: - for text in generator: - if not color: - text = strip_ansi(text) - - c.stdin.write(text) - except BrokenPipeError: - # In case the pager exited unexpectedly, ignore the broken pipe error. - pass - except Exception as e: - # In case there is an exception we want to close the pager immediately - # and let the caller handle it. - # Otherwise the pager will keep running, and the user may not notice - # the error message, or worse yet it may leave the terminal in a broken state. - c.terminate() - raise e - finally: - # We must close stdin and wait for the pager to exit before we continue - try: - c.stdin.close() - # Close implies flush, so it might throw a BrokenPipeError if the pager - # process exited already. - except BrokenPipeError: - pass - - # Less doesn't respect ^C, but catches it for its own UI purposes (aborting - # search or other commands inside less). - # - # That means when the user hits ^C, the parent process (click) terminates, - # but less is still alive, paging the output and messing up the terminal. - # - # If the user wants to make the pager exit on ^C, they should set - # `LESS='-K'`. It's not our decision to make. - while True: - try: - c.wait() - except KeyboardInterrupt: - pass - else: - break - - return True - - -def _tempfilepager( - generator: cabc.Iterable[str], cmd_parts: list[str], color: bool | None -) -> bool: - """Page through text by invoking a program on a temporary file. - - Returns `True` if the command was found, `False` otherwise and thus another - pager should be attempted. - """ - # Split the command into the invoked CLI and its parameters. - if not cmd_parts: - return False - cmd = cmd_parts[0] - - cmd_filepath = which(cmd) - if not cmd_filepath: - return False - # Resolves symlinks and produces a normalized absolute path string. - cmd_path = Path(cmd_filepath).resolve() - - import subprocess - import tempfile - - fd, filename = tempfile.mkstemp() - # TODO: This never terminates if the passed generator never terminates. - text = "".join(generator) - if not color: - text = strip_ansi(text) - encoding = get_best_encoding(sys.stdout) - with open_stream(filename, "wb")[0] as f: - f.write(text.encode(encoding)) - try: - subprocess.call([str(cmd_path), filename]) - except OSError: - # Command not found - pass - finally: - os.close(fd) - os.unlink(filename) - - return True - - -def _nullpager( - stream: t.TextIO, generator: cabc.Iterable[str], color: bool | None -) -> None: - """Simply print unformatted text. This is the ultimate fallback.""" - for text in generator: - if not color: - text = strip_ansi(text) - stream.write(text) - - -class Editor: - def __init__( - self, - editor: str | None = None, - env: cabc.Mapping[str, str] | None = None, - require_save: bool = True, - extension: str = ".txt", - ) -> None: - self.editor = editor - self.env = env - self.require_save = require_save - self.extension = extension - - def get_editor(self) -> str: - if self.editor is not None: - return self.editor - for key in "VISUAL", "EDITOR": - rv = os.environ.get(key) - if rv: - return rv - if WIN: - return "notepad" - for editor in "sensible-editor", "vim", "nano": - if which(editor) is not None: - return editor - return "vi" - - def edit_files(self, filenames: cabc.Iterable[str]) -> None: - import subprocess - - editor = self.get_editor() - environ: dict[str, str] | None = None - - if self.env: - environ = os.environ.copy() - environ.update(self.env) - - exc_filename = " ".join(f'"{filename}"' for filename in filenames) - - try: - c = subprocess.Popen( - args=f"{editor} {exc_filename}", env=environ, shell=True - ) - exit_code = c.wait() - if exit_code != 0: - raise ClickException( - _("{editor}: Editing failed").format(editor=editor) - ) - except OSError as e: - raise ClickException( - _("{editor}: Editing failed: {e}").format(editor=editor, e=e) - ) from e - - @t.overload - def edit(self, text: bytes | bytearray) -> bytes | None: ... - - # We cannot know whether or not the type expected is str or bytes when None - # is passed, so str is returned as that was what was done before. - @t.overload - def edit(self, text: str | None) -> str | None: ... - - def edit(self, text: str | bytes | bytearray | None) -> str | bytes | None: - import tempfile - - if text is None: - data = b"" - elif isinstance(text, (bytes, bytearray)): - data = text - else: - if text and not text.endswith("\n"): - text += "\n" - - if WIN: - data = text.replace("\n", "\r\n").encode("utf-8-sig") - else: - data = text.encode("utf-8") - - fd, name = tempfile.mkstemp(prefix="editor-", suffix=self.extension) - f: t.BinaryIO - - try: - with os.fdopen(fd, "wb") as f: - f.write(data) - - # If the filesystem resolution is 1 second, like Mac OS - # 10.12 Extended, or 2 seconds, like FAT32, and the editor - # closes very fast, require_save can fail. Set the modified - # time to be 2 seconds in the past to work around this. - os.utime(name, (os.path.getatime(name), os.path.getmtime(name) - 2)) - # Depending on the resolution, the exact value might not be - # recorded, so get the new recorded value. - timestamp = os.path.getmtime(name) - - self.edit_files((name,)) - - if self.require_save and os.path.getmtime(name) == timestamp: - return None - - with open(name, "rb") as f: - rv = f.read() - - if isinstance(text, (bytes, bytearray)): - return rv - - return rv.decode("utf-8-sig").replace("\r\n", "\n") - finally: - os.unlink(name) - - -def open_url(url: str, wait: bool = False, locate: bool = False) -> int: - import subprocess - - def _unquote_file(url: str) -> str: - from urllib.parse import unquote - - if url.startswith("file://"): - url = unquote(url[7:]) - - return url - - if sys.platform == "darwin": - args = ["open"] - if wait: - args.append("-W") - if locate: - args.append("-R") - args.append(_unquote_file(url)) - null = open("/dev/null", "w") - try: - return subprocess.Popen(args, stderr=null).wait() - finally: - null.close() - elif WIN: - if locate: - url = _unquote_file(url) - args = ["explorer", f"/select,{url}"] - else: - args = ["start"] - if wait: - args.append("/WAIT") - args.append("") - args.append(url) - try: - return subprocess.call(args) - except OSError: - # Command not found - return 127 - elif CYGWIN: - if locate: - url = _unquote_file(url) - args = ["cygstart", os.path.dirname(url)] - else: - args = ["cygstart"] - if wait: - args.append("-w") - args.append(url) - try: - return subprocess.call(args) - except OSError: - # Command not found - return 127 - - try: - if locate: - url = os.path.dirname(_unquote_file(url)) or "." - else: - url = _unquote_file(url) - c = subprocess.Popen(["xdg-open", url]) - if wait: - return c.wait() - return 0 - except OSError: - if url.startswith(("http://", "https://")) and not locate and not wait: - import webbrowser - - webbrowser.open(url) - return 0 - return 1 - - -def _translate_ch_to_exc(ch: str) -> None: - if ch == "\x03": - raise KeyboardInterrupt() - - if ch == "\x04" and not WIN: # Unix-like, Ctrl+D - raise EOFError() - - if ch == "\x1a" and WIN: # Windows, Ctrl+Z - raise EOFError() - - return None - - -if sys.platform == "win32": - import msvcrt - - @contextlib.contextmanager - def raw_terminal() -> cabc.Iterator[int]: - yield -1 - - def getchar(echo: bool) -> str: - # The function `getch` will return a bytes object corresponding to - # the pressed character. Since Windows 10 build 1803, it will also - # return \x00 when called a second time after pressing a regular key. - # - # `getwch` does not share this probably-bugged behavior. Moreover, it - # returns a Unicode object by default, which is what we want. - # - # Either of these functions will return \x00 or \xe0 to indicate - # a special key, and you need to call the same function again to get - # the "rest" of the code. The fun part is that \u00e0 is - # "latin small letter a with grave", so if you type that on a French - # keyboard, you _also_ get a \xe0. - # E.g., consider the Up arrow. This returns \xe0 and then \x48. The - # resulting Unicode string reads as "a with grave" + "capital H". - # This is indistinguishable from when the user actually types - # "a with grave" and then "capital H". - # - # When \xe0 is returned, we assume it's part of a special-key sequence - # and call `getwch` again, but that means that when the user types - # the \u00e0 character, `getchar` doesn't return until a second - # character is typed. - # The alternative is returning immediately, but that would mess up - # cross-platform handling of arrow keys and others that start with - # \xe0. Another option is using `getch`, but then we can't reliably - # read non-ASCII characters, because return values of `getch` are - # limited to the current 8-bit codepage. - # - # Anyway, Click doesn't claim to do this Right(tm), and using `getwch` - # is doing the right thing in more situations than with `getch`. - - if echo: - func = t.cast(t.Callable[[], str], msvcrt.getwche) - else: - func = t.cast(t.Callable[[], str], msvcrt.getwch) - - rv = func() - - if rv in ("\x00", "\xe0"): - # \x00 and \xe0 are control characters that indicate special key, - # see above. - rv += func() - - _translate_ch_to_exc(rv) - return rv - -else: - import termios - import tty - - @contextlib.contextmanager - def raw_terminal() -> cabc.Iterator[int]: - f: t.TextIO | None - fd: int - - if not isatty(sys.stdin): - f = open("/dev/tty") - fd = f.fileno() - else: - fd = sys.stdin.fileno() - f = None - - try: - old_settings = termios.tcgetattr(fd) - - try: - tty.setraw(fd) - yield fd - finally: - termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) - sys.stdout.flush() - - if f is not None: - f.close() - except termios.error: - pass - - def getchar(echo: bool) -> str: - with raw_terminal() as fd: - ch = os.read(fd, 32).decode(get_best_encoding(sys.stdin), "replace") - - if echo and isatty(sys.stdout): - sys.stdout.write(ch) - - _translate_ch_to_exc(ch) - return ch diff --git a/venv/lib/python3.10/site-packages/click/_textwrap.py b/venv/lib/python3.10/site-packages/click/_textwrap.py deleted file mode 100644 index 97fbee3dc6ce372600e702d38d9139bcf6a40bec..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/_textwrap.py +++ /dev/null @@ -1,51 +0,0 @@ -from __future__ import annotations - -import collections.abc as cabc -import textwrap -from contextlib import contextmanager - - -class TextWrapper(textwrap.TextWrapper): - def _handle_long_word( - self, - reversed_chunks: list[str], - cur_line: list[str], - cur_len: int, - width: int, - ) -> None: - space_left = max(width - cur_len, 1) - - if self.break_long_words: - last = reversed_chunks[-1] - cut = last[:space_left] - res = last[space_left:] - cur_line.append(cut) - reversed_chunks[-1] = res - elif not cur_line: - cur_line.append(reversed_chunks.pop()) - - @contextmanager - def extra_indent(self, indent: str) -> cabc.Iterator[None]: - old_initial_indent = self.initial_indent - old_subsequent_indent = self.subsequent_indent - self.initial_indent += indent - self.subsequent_indent += indent - - try: - yield - finally: - self.initial_indent = old_initial_indent - self.subsequent_indent = old_subsequent_indent - - def indent_only(self, text: str) -> str: - rv = [] - - for idx, line in enumerate(text.splitlines()): - indent = self.initial_indent - - if idx > 0: - indent = self.subsequent_indent - - rv.append(f"{indent}{line}") - - return "\n".join(rv) diff --git a/venv/lib/python3.10/site-packages/click/_winconsole.py b/venv/lib/python3.10/site-packages/click/_winconsole.py deleted file mode 100644 index e56c7c6ae7ee7e577db6ace659cd58c0c652074c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/_winconsole.py +++ /dev/null @@ -1,296 +0,0 @@ -# This module is based on the excellent work by Adam Bartoš who -# provided a lot of what went into the implementation here in -# the discussion to issue1602 in the Python bug tracker. -# -# There are some general differences in regards to how this works -# compared to the original patches as we do not need to patch -# the entire interpreter but just work in our little world of -# echo and prompt. -from __future__ import annotations - -import collections.abc as cabc -import io -import sys -import time -import typing as t -from ctypes import Array -from ctypes import byref -from ctypes import c_char -from ctypes import c_char_p -from ctypes import c_int -from ctypes import c_ssize_t -from ctypes import c_ulong -from ctypes import c_void_p -from ctypes import POINTER -from ctypes import py_object -from ctypes import Structure -from ctypes.wintypes import DWORD -from ctypes.wintypes import HANDLE -from ctypes.wintypes import LPCWSTR -from ctypes.wintypes import LPWSTR - -from ._compat import _NonClosingTextIOWrapper - -assert sys.platform == "win32" -import msvcrt # noqa: E402 -from ctypes import windll # noqa: E402 -from ctypes import WINFUNCTYPE # noqa: E402 - -c_ssize_p = POINTER(c_ssize_t) - -kernel32 = windll.kernel32 -GetStdHandle = kernel32.GetStdHandle -ReadConsoleW = kernel32.ReadConsoleW -WriteConsoleW = kernel32.WriteConsoleW -GetConsoleMode = kernel32.GetConsoleMode -GetLastError = kernel32.GetLastError -GetCommandLineW = WINFUNCTYPE(LPWSTR)(("GetCommandLineW", windll.kernel32)) -CommandLineToArgvW = WINFUNCTYPE(POINTER(LPWSTR), LPCWSTR, POINTER(c_int))( - ("CommandLineToArgvW", windll.shell32) -) -LocalFree = WINFUNCTYPE(c_void_p, c_void_p)(("LocalFree", windll.kernel32)) - -STDIN_HANDLE = GetStdHandle(-10) -STDOUT_HANDLE = GetStdHandle(-11) -STDERR_HANDLE = GetStdHandle(-12) - -PyBUF_SIMPLE = 0 -PyBUF_WRITABLE = 1 - -ERROR_SUCCESS = 0 -ERROR_NOT_ENOUGH_MEMORY = 8 -ERROR_OPERATION_ABORTED = 995 - -STDIN_FILENO = 0 -STDOUT_FILENO = 1 -STDERR_FILENO = 2 - -EOF = b"\x1a" -MAX_BYTES_WRITTEN = 32767 - -if t.TYPE_CHECKING: - try: - # Using `typing_extensions.Buffer` instead of `collections.abc` - # on Windows for some reason does not have `Sized` implemented. - from collections.abc import Buffer # type: ignore - except ImportError: - from typing_extensions import Buffer - -try: - from ctypes import pythonapi -except ImportError: - # On PyPy we cannot get buffers so our ability to operate here is - # severely limited. - get_buffer = None -else: - - class Py_buffer(Structure): - _fields_ = [ # noqa: RUF012 - ("buf", c_void_p), - ("obj", py_object), - ("len", c_ssize_t), - ("itemsize", c_ssize_t), - ("readonly", c_int), - ("ndim", c_int), - ("format", c_char_p), - ("shape", c_ssize_p), - ("strides", c_ssize_p), - ("suboffsets", c_ssize_p), - ("internal", c_void_p), - ] - - PyObject_GetBuffer = pythonapi.PyObject_GetBuffer - PyBuffer_Release = pythonapi.PyBuffer_Release - - def get_buffer(obj: Buffer, writable: bool = False) -> Array[c_char]: - buf = Py_buffer() - flags: int = PyBUF_WRITABLE if writable else PyBUF_SIMPLE - PyObject_GetBuffer(py_object(obj), byref(buf), flags) - - try: - buffer_type = c_char * buf.len - out: Array[c_char] = buffer_type.from_address(buf.buf) - return out - finally: - PyBuffer_Release(byref(buf)) - - -class _WindowsConsoleRawIOBase(io.RawIOBase): - def __init__(self, handle: int | None) -> None: - self.handle = handle - - def isatty(self) -> t.Literal[True]: - super().isatty() - return True - - -class _WindowsConsoleReader(_WindowsConsoleRawIOBase): - def readable(self) -> t.Literal[True]: - return True - - def readinto(self, b: Buffer) -> int: - bytes_to_be_read = len(b) - if not bytes_to_be_read: - return 0 - elif bytes_to_be_read % 2: - raise ValueError( - "cannot read odd number of bytes from UTF-16-LE encoded console" - ) - - buffer = get_buffer(b, writable=True) - code_units_to_be_read = bytes_to_be_read // 2 - code_units_read = c_ulong() - - rv = ReadConsoleW( - HANDLE(self.handle), - buffer, - code_units_to_be_read, - byref(code_units_read), - None, - ) - if GetLastError() == ERROR_OPERATION_ABORTED: - # wait for KeyboardInterrupt - time.sleep(0.1) - if not rv: - raise OSError(f"Windows error: {GetLastError()}") - - if buffer[0] == EOF: - return 0 - return 2 * code_units_read.value - - -class _WindowsConsoleWriter(_WindowsConsoleRawIOBase): - def writable(self) -> t.Literal[True]: - return True - - @staticmethod - def _get_error_message(errno: int) -> str: - if errno == ERROR_SUCCESS: - return "ERROR_SUCCESS" - elif errno == ERROR_NOT_ENOUGH_MEMORY: - return "ERROR_NOT_ENOUGH_MEMORY" - return f"Windows error {errno}" - - def write(self, b: Buffer) -> int: - bytes_to_be_written = len(b) - buf = get_buffer(b) - code_units_to_be_written = min(bytes_to_be_written, MAX_BYTES_WRITTEN) // 2 - code_units_written = c_ulong() - - WriteConsoleW( - HANDLE(self.handle), - buf, - code_units_to_be_written, - byref(code_units_written), - None, - ) - bytes_written = 2 * code_units_written.value - - if bytes_written == 0 and bytes_to_be_written > 0: - raise OSError(self._get_error_message(GetLastError())) - return bytes_written - - -class ConsoleStream: - def __init__(self, text_stream: t.TextIO, byte_stream: t.BinaryIO) -> None: - self._text_stream = text_stream - self.buffer = byte_stream - - @property - def name(self) -> str: - return self.buffer.name - - def write(self, x: t.AnyStr) -> int: - if isinstance(x, str): - return self._text_stream.write(x) - try: - self.flush() - except Exception: - pass - return self.buffer.write(x) - - def writelines(self, lines: cabc.Iterable[t.AnyStr]) -> None: - for line in lines: - self.write(line) - - def __getattr__(self, name: str) -> t.Any: - return getattr(self._text_stream, name) - - def isatty(self) -> bool: - return self.buffer.isatty() - - def __repr__(self) -> str: - return f"" - - -def _get_text_stdin(buffer_stream: t.BinaryIO) -> t.TextIO: - text_stream = _NonClosingTextIOWrapper( - io.BufferedReader(_WindowsConsoleReader(STDIN_HANDLE)), - "utf-16-le", - "strict", - line_buffering=True, - ) - return t.cast(t.TextIO, ConsoleStream(text_stream, buffer_stream)) - - -def _get_text_stdout(buffer_stream: t.BinaryIO) -> t.TextIO: - text_stream = _NonClosingTextIOWrapper( - io.BufferedWriter(_WindowsConsoleWriter(STDOUT_HANDLE)), - "utf-16-le", - "strict", - line_buffering=True, - ) - return t.cast(t.TextIO, ConsoleStream(text_stream, buffer_stream)) - - -def _get_text_stderr(buffer_stream: t.BinaryIO) -> t.TextIO: - text_stream = _NonClosingTextIOWrapper( - io.BufferedWriter(_WindowsConsoleWriter(STDERR_HANDLE)), - "utf-16-le", - "strict", - line_buffering=True, - ) - return t.cast(t.TextIO, ConsoleStream(text_stream, buffer_stream)) - - -_stream_factories: cabc.Mapping[int, t.Callable[[t.BinaryIO], t.TextIO]] = { - 0: _get_text_stdin, - 1: _get_text_stdout, - 2: _get_text_stderr, -} - - -def _is_console(f: t.TextIO) -> bool: - if not hasattr(f, "fileno"): - return False - - try: - fileno = f.fileno() - except (OSError, io.UnsupportedOperation): - return False - - handle = msvcrt.get_osfhandle(fileno) - return bool(GetConsoleMode(handle, byref(DWORD()))) - - -def _get_windows_console_stream( - f: t.TextIO, encoding: str | None, errors: str | None -) -> t.TextIO | None: - if ( - get_buffer is None - or encoding not in {"utf-16-le", None} - or errors not in {"strict", None} - or not _is_console(f) - ): - return None - - func = _stream_factories.get(f.fileno()) - if func is None: - return None - - b = getattr(f, "buffer", None) - - if b is None: - return None - - return func(b) diff --git a/venv/lib/python3.10/site-packages/click/core.py b/venv/lib/python3.10/site-packages/click/core.py deleted file mode 100644 index f57ada62b195fa533387b5d06b366b17edb7d2f4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/core.py +++ /dev/null @@ -1,3135 +0,0 @@ -from __future__ import annotations - -import collections.abc as cabc -import enum -import errno -import inspect -import os -import sys -import typing as t -from collections import abc -from collections import Counter -from contextlib import AbstractContextManager -from contextlib import contextmanager -from contextlib import ExitStack -from functools import update_wrapper -from gettext import gettext as _ -from gettext import ngettext -from itertools import repeat -from types import TracebackType - -from . import types -from .exceptions import Abort -from .exceptions import BadParameter -from .exceptions import ClickException -from .exceptions import Exit -from .exceptions import MissingParameter -from .exceptions import NoArgsIsHelpError -from .exceptions import UsageError -from .formatting import HelpFormatter -from .formatting import join_options -from .globals import pop_context -from .globals import push_context -from .parser import _flag_needs_value -from .parser import _OptionParser -from .parser import _split_opt -from .termui import confirm -from .termui import prompt -from .termui import style -from .utils import _detect_program_name -from .utils import _expand_args -from .utils import echo -from .utils import make_default_short_help -from .utils import make_str -from .utils import PacifyFlushWrapper - -if t.TYPE_CHECKING: - from .shell_completion import CompletionItem - -F = t.TypeVar("F", bound="t.Callable[..., t.Any]") -V = t.TypeVar("V") - - -def _complete_visible_commands( - ctx: Context, incomplete: str -) -> cabc.Iterator[tuple[str, Command]]: - """List all the subcommands of a group that start with the - incomplete value and aren't hidden. - - :param ctx: Invocation context for the group. - :param incomplete: Value being completed. May be empty. - """ - multi = t.cast(Group, ctx.command) - - for name in multi.list_commands(ctx): - if name.startswith(incomplete): - command = multi.get_command(ctx, name) - - if command is not None and not command.hidden: - yield name, command - - -def _check_nested_chain( - base_command: Group, cmd_name: str, cmd: Command, register: bool = False -) -> None: - if not base_command.chain or not isinstance(cmd, Group): - return - - if register: - message = ( - f"It is not possible to add the group {cmd_name!r} to another" - f" group {base_command.name!r} that is in chain mode." - ) - else: - message = ( - f"Found the group {cmd_name!r} as subcommand to another group " - f" {base_command.name!r} that is in chain mode. This is not supported." - ) - - raise RuntimeError(message) - - -def batch(iterable: cabc.Iterable[V], batch_size: int) -> list[tuple[V, ...]]: - return list(zip(*repeat(iter(iterable), batch_size), strict=False)) - - -@contextmanager -def augment_usage_errors( - ctx: Context, param: Parameter | None = None -) -> cabc.Iterator[None]: - """Context manager that attaches extra information to exceptions.""" - try: - yield - except BadParameter as e: - if e.ctx is None: - e.ctx = ctx - if param is not None and e.param is None: - e.param = param - raise - except UsageError as e: - if e.ctx is None: - e.ctx = ctx - raise - - -def iter_params_for_processing( - invocation_order: cabc.Sequence[Parameter], - declaration_order: cabc.Sequence[Parameter], -) -> list[Parameter]: - """Returns all declared parameters in the order they should be processed. - - The declared parameters are re-shuffled depending on the order in which - they were invoked, as well as the eagerness of each parameters. - - The invocation order takes precedence over the declaration order. I.e. the - order in which the user provided them to the CLI is respected. - - This behavior and its effect on callback evaluation is detailed at: - https://click.palletsprojects.com/en/stable/advanced/#callback-evaluation-order - """ - - def sort_key(item: Parameter) -> tuple[bool, float]: - try: - idx: float = invocation_order.index(item) - except ValueError: - idx = float("inf") - - return not item.is_eager, idx - - return sorted(declaration_order, key=sort_key) - - -class ParameterSource(enum.Enum): - """This is an :class:`~enum.Enum` that indicates the source of a - parameter's value. - - Use :meth:`click.Context.get_parameter_source` to get the - source for a parameter by name. - - .. versionchanged:: 8.0 - Use :class:`~enum.Enum` and drop the ``validate`` method. - - .. versionchanged:: 8.0 - Added the ``PROMPT`` value. - """ - - COMMANDLINE = enum.auto() - """The value was provided by the command line args.""" - ENVIRONMENT = enum.auto() - """The value was provided with an environment variable.""" - DEFAULT = enum.auto() - """Used the default specified by the parameter.""" - DEFAULT_MAP = enum.auto() - """Used a default provided by :attr:`Context.default_map`.""" - PROMPT = enum.auto() - """Used a prompt to confirm a default or provide a value.""" - - -class Context: - """The context is a special internal object that holds state relevant - for the script execution at every single level. It's normally invisible - to commands unless they opt-in to getting access to it. - - The context is useful as it can pass internal objects around and can - control special execution features such as reading data from - environment variables. - - A context can be used as context manager in which case it will call - :meth:`close` on teardown. - - :param command: the command class for this context. - :param parent: the parent context. - :param info_name: the info name for this invocation. Generally this - is the most descriptive name for the script or - command. For the toplevel script it is usually - the name of the script, for commands below it it's - the name of the script. - :param obj: an arbitrary object of user data. - :param auto_envvar_prefix: the prefix to use for automatic environment - variables. If this is `None` then reading - from environment variables is disabled. This - does not affect manually set environment - variables which are always read. - :param default_map: a dictionary (like object) with default values - for parameters. - :param terminal_width: the width of the terminal. The default is - inherit from parent context. If no context - defines the terminal width then auto - detection will be applied. - :param max_content_width: the maximum width for content rendered by - Click (this currently only affects help - pages). This defaults to 80 characters if - not overridden. In other words: even if the - terminal is larger than that, Click will not - format things wider than 80 characters by - default. In addition to that, formatters might - add some safety mapping on the right. - :param resilient_parsing: if this flag is enabled then Click will - parse without any interactivity or callback - invocation. Default values will also be - ignored. This is useful for implementing - things such as completion support. - :param allow_extra_args: if this is set to `True` then extra arguments - at the end will not raise an error and will be - kept on the context. The default is to inherit - from the command. - :param allow_interspersed_args: if this is set to `False` then options - and arguments cannot be mixed. The - default is to inherit from the command. - :param ignore_unknown_options: instructs click to ignore options it does - not know and keeps them for later - processing. - :param help_option_names: optionally a list of strings that define how - the default help parameter is named. The - default is ``['--help']``. - :param token_normalize_func: an optional function that is used to - normalize tokens (options, choices, - etc.). This for instance can be used to - implement case insensitive behavior. - :param color: controls if the terminal supports ANSI colors or not. The - default is autodetection. This is only needed if ANSI - codes are used in texts that Click prints which is by - default not the case. This for instance would affect - help output. - :param show_default: Show the default value for commands. If this - value is not set, it defaults to the value from the parent - context. ``Command.show_default`` overrides this default for the - specific command. - - .. versionchanged:: 8.2 - The ``protected_args`` attribute is deprecated and will be removed in - Click 9.0. ``args`` will contain remaining unparsed tokens. - - .. versionchanged:: 8.1 - The ``show_default`` parameter is overridden by - ``Command.show_default``, instead of the other way around. - - .. versionchanged:: 8.0 - The ``show_default`` parameter defaults to the value from the - parent context. - - .. versionchanged:: 7.1 - Added the ``show_default`` parameter. - - .. versionchanged:: 4.0 - Added the ``color``, ``ignore_unknown_options``, and - ``max_content_width`` parameters. - - .. versionchanged:: 3.0 - Added the ``allow_extra_args`` and ``allow_interspersed_args`` - parameters. - - .. versionchanged:: 2.0 - Added the ``resilient_parsing``, ``help_option_names``, and - ``token_normalize_func`` parameters. - """ - - #: The formatter class to create with :meth:`make_formatter`. - #: - #: .. versionadded:: 8.0 - formatter_class: type[HelpFormatter] = HelpFormatter - - def __init__( - self, - command: Command, - parent: Context | None = None, - info_name: str | None = None, - obj: t.Any | None = None, - auto_envvar_prefix: str | None = None, - default_map: cabc.MutableMapping[str, t.Any] | None = None, - terminal_width: int | None = None, - max_content_width: int | None = None, - resilient_parsing: bool = False, - allow_extra_args: bool | None = None, - allow_interspersed_args: bool | None = None, - ignore_unknown_options: bool | None = None, - help_option_names: list[str] | None = None, - token_normalize_func: t.Callable[[str], str] | None = None, - color: bool | None = None, - show_default: bool | None = None, - ) -> None: - #: the parent context or `None` if none exists. - self.parent = parent - #: the :class:`Command` for this context. - self.command = command - #: the descriptive information name - self.info_name = info_name - #: Map of parameter names to their parsed values. Parameters - #: with ``expose_value=False`` are not stored. - self.params: dict[str, t.Any] = {} - #: the leftover arguments. - self.args: list[str] = [] - #: protected arguments. These are arguments that are prepended - #: to `args` when certain parsing scenarios are encountered but - #: must be never propagated to another arguments. This is used - #: to implement nested parsing. - self._protected_args: list[str] = [] - #: the collected prefixes of the command's options. - self._opt_prefixes: set[str] = set(parent._opt_prefixes) if parent else set() - - if obj is None and parent is not None: - obj = parent.obj - - #: the user object stored. - self.obj: t.Any = obj - self._meta: dict[str, t.Any] = getattr(parent, "meta", {}) - - #: A dictionary (-like object) with defaults for parameters. - if ( - default_map is None - and info_name is not None - and parent is not None - and parent.default_map is not None - ): - default_map = parent.default_map.get(info_name) - - self.default_map: cabc.MutableMapping[str, t.Any] | None = default_map - - #: This flag indicates if a subcommand is going to be executed. A - #: group callback can use this information to figure out if it's - #: being executed directly or because the execution flow passes - #: onwards to a subcommand. By default it's None, but it can be - #: the name of the subcommand to execute. - #: - #: If chaining is enabled this will be set to ``'*'`` in case - #: any commands are executed. It is however not possible to - #: figure out which ones. If you require this knowledge you - #: should use a :func:`result_callback`. - self.invoked_subcommand: str | None = None - - if terminal_width is None and parent is not None: - terminal_width = parent.terminal_width - - #: The width of the terminal (None is autodetection). - self.terminal_width: int | None = terminal_width - - if max_content_width is None and parent is not None: - max_content_width = parent.max_content_width - - #: The maximum width of formatted content (None implies a sensible - #: default which is 80 for most things). - self.max_content_width: int | None = max_content_width - - if allow_extra_args is None: - allow_extra_args = command.allow_extra_args - - #: Indicates if the context allows extra args or if it should - #: fail on parsing. - #: - #: .. versionadded:: 3.0 - self.allow_extra_args = allow_extra_args - - if allow_interspersed_args is None: - allow_interspersed_args = command.allow_interspersed_args - - #: Indicates if the context allows mixing of arguments and - #: options or not. - #: - #: .. versionadded:: 3.0 - self.allow_interspersed_args: bool = allow_interspersed_args - - if ignore_unknown_options is None: - ignore_unknown_options = command.ignore_unknown_options - - #: Instructs click to ignore options that a command does not - #: understand and will store it on the context for later - #: processing. This is primarily useful for situations where you - #: want to call into external programs. Generally this pattern is - #: strongly discouraged because it's not possibly to losslessly - #: forward all arguments. - #: - #: .. versionadded:: 4.0 - self.ignore_unknown_options: bool = ignore_unknown_options - - if help_option_names is None: - if parent is not None: - help_option_names = parent.help_option_names - else: - help_option_names = ["--help"] - - #: The names for the help options. - self.help_option_names: list[str] = help_option_names - - if token_normalize_func is None and parent is not None: - token_normalize_func = parent.token_normalize_func - - #: An optional normalization function for tokens. This is - #: options, choices, commands etc. - self.token_normalize_func: t.Callable[[str], str] | None = token_normalize_func - - #: Indicates if resilient parsing is enabled. In that case Click - #: will do its best to not cause any failures and default values - #: will be ignored. Useful for completion. - self.resilient_parsing: bool = resilient_parsing - - # If there is no envvar prefix yet, but the parent has one and - # the command on this level has a name, we can expand the envvar - # prefix automatically. - if auto_envvar_prefix is None: - if ( - parent is not None - and parent.auto_envvar_prefix is not None - and self.info_name is not None - ): - auto_envvar_prefix = ( - f"{parent.auto_envvar_prefix}_{self.info_name.upper()}" - ) - else: - auto_envvar_prefix = auto_envvar_prefix.upper() - - if auto_envvar_prefix is not None: - auto_envvar_prefix = auto_envvar_prefix.replace("-", "_") - - self.auto_envvar_prefix: str | None = auto_envvar_prefix - - if color is None and parent is not None: - color = parent.color - - #: Controls if styling output is wanted or not. - self.color: bool | None = color - - if show_default is None and parent is not None: - show_default = parent.show_default - - #: Show option default values when formatting help text. - self.show_default: bool | None = show_default - - self._close_callbacks: list[t.Callable[[], t.Any]] = [] - self._depth = 0 - self._parameter_source: dict[str, ParameterSource] = {} - self._exit_stack = ExitStack() - - @property - def protected_args(self) -> list[str]: - import warnings - - warnings.warn( - "'protected_args' is deprecated and will be removed in Click 9.0." - " 'args' will contain remaining unparsed tokens.", - DeprecationWarning, - stacklevel=2, - ) - return self._protected_args - - def to_info_dict(self) -> dict[str, t.Any]: - """Gather information that could be useful for a tool generating - user-facing documentation. This traverses the entire CLI - structure. - - .. code-block:: python - - with Context(cli) as ctx: - info = ctx.to_info_dict() - - .. versionadded:: 8.0 - """ - return { - "command": self.command.to_info_dict(self), - "info_name": self.info_name, - "allow_extra_args": self.allow_extra_args, - "allow_interspersed_args": self.allow_interspersed_args, - "ignore_unknown_options": self.ignore_unknown_options, - "auto_envvar_prefix": self.auto_envvar_prefix, - } - - def __enter__(self) -> Context: - self._depth += 1 - push_context(self) - return self - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc_value: BaseException | None, - tb: TracebackType | None, - ) -> None: - self._depth -= 1 - if self._depth == 0: - self.close() - pop_context() - - @contextmanager - def scope(self, cleanup: bool = True) -> cabc.Iterator[Context]: - """This helper method can be used with the context object to promote - it to the current thread local (see :func:`get_current_context`). - The default behavior of this is to invoke the cleanup functions which - can be disabled by setting `cleanup` to `False`. The cleanup - functions are typically used for things such as closing file handles. - - If the cleanup is intended the context object can also be directly - used as a context manager. - - Example usage:: - - with ctx.scope(): - assert get_current_context() is ctx - - This is equivalent:: - - with ctx: - assert get_current_context() is ctx - - .. versionadded:: 5.0 - - :param cleanup: controls if the cleanup functions should be run or - not. The default is to run these functions. In - some situations the context only wants to be - temporarily pushed in which case this can be disabled. - Nested pushes automatically defer the cleanup. - """ - if not cleanup: - self._depth += 1 - try: - with self as rv: - yield rv - finally: - if not cleanup: - self._depth -= 1 - - @property - def meta(self) -> dict[str, t.Any]: - """This is a dictionary which is shared with all the contexts - that are nested. It exists so that click utilities can store some - state here if they need to. It is however the responsibility of - that code to manage this dictionary well. - - The keys are supposed to be unique dotted strings. For instance - module paths are a good choice for it. What is stored in there is - irrelevant for the operation of click. However what is important is - that code that places data here adheres to the general semantics of - the system. - - Example usage:: - - LANG_KEY = f'{__name__}.lang' - - def set_language(value): - ctx = get_current_context() - ctx.meta[LANG_KEY] = value - - def get_language(): - return get_current_context().meta.get(LANG_KEY, 'en_US') - - .. versionadded:: 5.0 - """ - return self._meta - - def make_formatter(self) -> HelpFormatter: - """Creates the :class:`~click.HelpFormatter` for the help and - usage output. - - To quickly customize the formatter class used without overriding - this method, set the :attr:`formatter_class` attribute. - - .. versionchanged:: 8.0 - Added the :attr:`formatter_class` attribute. - """ - return self.formatter_class( - width=self.terminal_width, max_width=self.max_content_width - ) - - def with_resource(self, context_manager: AbstractContextManager[V]) -> V: - """Register a resource as if it were used in a ``with`` - statement. The resource will be cleaned up when the context is - popped. - - Uses :meth:`contextlib.ExitStack.enter_context`. It calls the - resource's ``__enter__()`` method and returns the result. When - the context is popped, it closes the stack, which calls the - resource's ``__exit__()`` method. - - To register a cleanup function for something that isn't a - context manager, use :meth:`call_on_close`. Or use something - from :mod:`contextlib` to turn it into a context manager first. - - .. code-block:: python - - @click.group() - @click.option("--name") - @click.pass_context - def cli(ctx): - ctx.obj = ctx.with_resource(connect_db(name)) - - :param context_manager: The context manager to enter. - :return: Whatever ``context_manager.__enter__()`` returns. - - .. versionadded:: 8.0 - """ - return self._exit_stack.enter_context(context_manager) - - def call_on_close(self, f: t.Callable[..., t.Any]) -> t.Callable[..., t.Any]: - """Register a function to be called when the context tears down. - - This can be used to close resources opened during the script - execution. Resources that support Python's context manager - protocol which would be used in a ``with`` statement should be - registered with :meth:`with_resource` instead. - - :param f: The function to execute on teardown. - """ - return self._exit_stack.callback(f) - - def close(self) -> None: - """Invoke all close callbacks registered with - :meth:`call_on_close`, and exit all context managers entered - with :meth:`with_resource`. - """ - self._exit_stack.close() - # In case the context is reused, create a new exit stack. - self._exit_stack = ExitStack() - - @property - def command_path(self) -> str: - """The computed command path. This is used for the ``usage`` - information on the help page. It's automatically created by - combining the info names of the chain of contexts to the root. - """ - rv = "" - if self.info_name is not None: - rv = self.info_name - if self.parent is not None: - parent_command_path = [self.parent.command_path] - - if isinstance(self.parent.command, Command): - for param in self.parent.command.get_params(self): - parent_command_path.extend(param.get_usage_pieces(self)) - - rv = f"{' '.join(parent_command_path)} {rv}" - return rv.lstrip() - - def find_root(self) -> Context: - """Finds the outermost context.""" - node = self - while node.parent is not None: - node = node.parent - return node - - def find_object(self, object_type: type[V]) -> V | None: - """Finds the closest object of a given type.""" - node: Context | None = self - - while node is not None: - if isinstance(node.obj, object_type): - return node.obj - - node = node.parent - - return None - - def ensure_object(self, object_type: type[V]) -> V: - """Like :meth:`find_object` but sets the innermost object to a - new instance of `object_type` if it does not exist. - """ - rv = self.find_object(object_type) - if rv is None: - self.obj = rv = object_type() - return rv - - @t.overload - def lookup_default( - self, name: str, call: t.Literal[True] = True - ) -> t.Any | None: ... - - @t.overload - def lookup_default( - self, name: str, call: t.Literal[False] = ... - ) -> t.Any | t.Callable[[], t.Any] | None: ... - - def lookup_default(self, name: str, call: bool = True) -> t.Any | None: - """Get the default for a parameter from :attr:`default_map`. - - :param name: Name of the parameter. - :param call: If the default is a callable, call it. Disable to - return the callable instead. - - .. versionchanged:: 8.0 - Added the ``call`` parameter. - """ - if self.default_map is not None: - value = self.default_map.get(name) - - if call and callable(value): - return value() - - return value - - return None - - def fail(self, message: str) -> t.NoReturn: - """Aborts the execution of the program with a specific error - message. - - :param message: the error message to fail with. - """ - raise UsageError(message, self) - - def abort(self) -> t.NoReturn: - """Aborts the script.""" - raise Abort() - - def exit(self, code: int = 0) -> t.NoReturn: - """Exits the application with a given exit code. - - .. versionchanged:: 8.2 - Callbacks and context managers registered with :meth:`call_on_close` - and :meth:`with_resource` are closed before exiting. - """ - self.close() - raise Exit(code) - - def get_usage(self) -> str: - """Helper method to get formatted usage string for the current - context and command. - """ - return self.command.get_usage(self) - - def get_help(self) -> str: - """Helper method to get formatted help page for the current - context and command. - """ - return self.command.get_help(self) - - def _make_sub_context(self, command: Command) -> Context: - """Create a new context of the same type as this context, but - for a new command. - - :meta private: - """ - return type(self)(command, info_name=command.name, parent=self) - - @t.overload - def invoke( - self, callback: t.Callable[..., V], /, *args: t.Any, **kwargs: t.Any - ) -> V: ... - - @t.overload - def invoke(self, callback: Command, /, *args: t.Any, **kwargs: t.Any) -> t.Any: ... - - def invoke( - self, callback: Command | t.Callable[..., V], /, *args: t.Any, **kwargs: t.Any - ) -> t.Any | V: - """Invokes a command callback in exactly the way it expects. There - are two ways to invoke this method: - - 1. the first argument can be a callback and all other arguments and - keyword arguments are forwarded directly to the function. - 2. the first argument is a click command object. In that case all - arguments are forwarded as well but proper click parameters - (options and click arguments) must be keyword arguments and Click - will fill in defaults. - - .. versionchanged:: 8.0 - All ``kwargs`` are tracked in :attr:`params` so they will be - passed if :meth:`forward` is called at multiple levels. - - .. versionchanged:: 3.2 - A new context is created, and missing arguments use default values. - """ - if isinstance(callback, Command): - other_cmd = callback - - if other_cmd.callback is None: - raise TypeError( - "The given command does not have a callback that can be invoked." - ) - else: - callback = t.cast("t.Callable[..., V]", other_cmd.callback) - - ctx = self._make_sub_context(other_cmd) - - for param in other_cmd.params: - if param.name not in kwargs and param.expose_value: - kwargs[param.name] = param.type_cast_value( # type: ignore - ctx, param.get_default(ctx) - ) - - # Track all kwargs as params, so that forward() will pass - # them on in subsequent calls. - ctx.params.update(kwargs) - else: - ctx = self - - with augment_usage_errors(self): - with ctx: - return callback(*args, **kwargs) - - def forward(self, cmd: Command, /, *args: t.Any, **kwargs: t.Any) -> t.Any: - """Similar to :meth:`invoke` but fills in default keyword - arguments from the current context if the other command expects - it. This cannot invoke callbacks directly, only other commands. - - .. versionchanged:: 8.0 - All ``kwargs`` are tracked in :attr:`params` so they will be - passed if ``forward`` is called at multiple levels. - """ - # Can only forward to other commands, not direct callbacks. - if not isinstance(cmd, Command): - raise TypeError("Callback is not a command.") - - for param in self.params: - if param not in kwargs: - kwargs[param] = self.params[param] - - return self.invoke(cmd, *args, **kwargs) - - def set_parameter_source(self, name: str, source: ParameterSource) -> None: - """Set the source of a parameter. This indicates the location - from which the value of the parameter was obtained. - - :param name: The name of the parameter. - :param source: A member of :class:`~click.core.ParameterSource`. - """ - self._parameter_source[name] = source - - def get_parameter_source(self, name: str) -> ParameterSource | None: - """Get the source of a parameter. This indicates the location - from which the value of the parameter was obtained. - - This can be useful for determining when a user specified a value - on the command line that is the same as the default value. It - will be :attr:`~click.core.ParameterSource.DEFAULT` only if the - value was actually taken from the default. - - :param name: The name of the parameter. - :rtype: ParameterSource - - .. versionchanged:: 8.0 - Returns ``None`` if the parameter was not provided from any - source. - """ - return self._parameter_source.get(name) - - -class Command: - """Commands are the basic building block of command line interfaces in - Click. A basic command handles command line parsing and might dispatch - more parsing to commands nested below it. - - :param name: the name of the command to use unless a group overrides it. - :param context_settings: an optional dictionary with defaults that are - passed to the context object. - :param callback: the callback to invoke. This is optional. - :param params: the parameters to register with this command. This can - be either :class:`Option` or :class:`Argument` objects. - :param help: the help string to use for this command. - :param epilog: like the help string but it's printed at the end of the - help page after everything else. - :param short_help: the short help to use for this command. This is - shown on the command listing of the parent command. - :param add_help_option: by default each command registers a ``--help`` - option. This can be disabled by this parameter. - :param no_args_is_help: this controls what happens if no arguments are - provided. This option is disabled by default. - If enabled this will add ``--help`` as argument - if no arguments are passed - :param hidden: hide this command from help outputs. - :param deprecated: If ``True`` or non-empty string, issues a message - indicating that the command is deprecated and highlights - its deprecation in --help. The message can be customized - by using a string as the value. - - .. versionchanged:: 8.2 - This is the base class for all commands, not ``BaseCommand``. - ``deprecated`` can be set to a string as well to customize the - deprecation message. - - .. versionchanged:: 8.1 - ``help``, ``epilog``, and ``short_help`` are stored unprocessed, - all formatting is done when outputting help text, not at init, - and is done even if not using the ``@command`` decorator. - - .. versionchanged:: 8.0 - Added a ``repr`` showing the command name. - - .. versionchanged:: 7.1 - Added the ``no_args_is_help`` parameter. - - .. versionchanged:: 2.0 - Added the ``context_settings`` parameter. - """ - - #: The context class to create with :meth:`make_context`. - #: - #: .. versionadded:: 8.0 - context_class: type[Context] = Context - - #: the default for the :attr:`Context.allow_extra_args` flag. - allow_extra_args = False - - #: the default for the :attr:`Context.allow_interspersed_args` flag. - allow_interspersed_args = True - - #: the default for the :attr:`Context.ignore_unknown_options` flag. - ignore_unknown_options = False - - def __init__( - self, - name: str | None, - context_settings: cabc.MutableMapping[str, t.Any] | None = None, - callback: t.Callable[..., t.Any] | None = None, - params: list[Parameter] | None = None, - help: str | None = None, - epilog: str | None = None, - short_help: str | None = None, - options_metavar: str | None = "[OPTIONS]", - add_help_option: bool = True, - no_args_is_help: bool = False, - hidden: bool = False, - deprecated: bool | str = False, - ) -> None: - #: the name the command thinks it has. Upon registering a command - #: on a :class:`Group` the group will default the command name - #: with this information. You should instead use the - #: :class:`Context`\'s :attr:`~Context.info_name` attribute. - self.name = name - - if context_settings is None: - context_settings = {} - - #: an optional dictionary with defaults passed to the context. - self.context_settings: cabc.MutableMapping[str, t.Any] = context_settings - - #: the callback to execute when the command fires. This might be - #: `None` in which case nothing happens. - self.callback = callback - #: the list of parameters for this command in the order they - #: should show up in the help page and execute. Eager parameters - #: will automatically be handled before non eager ones. - self.params: list[Parameter] = params or [] - self.help = help - self.epilog = epilog - self.options_metavar = options_metavar - self.short_help = short_help - self.add_help_option = add_help_option - self._help_option = None - self.no_args_is_help = no_args_is_help - self.hidden = hidden - self.deprecated = deprecated - - def to_info_dict(self, ctx: Context) -> dict[str, t.Any]: - return { - "name": self.name, - "params": [param.to_info_dict() for param in self.get_params(ctx)], - "help": self.help, - "epilog": self.epilog, - "short_help": self.short_help, - "hidden": self.hidden, - "deprecated": self.deprecated, - } - - def __repr__(self) -> str: - return f"<{self.__class__.__name__} {self.name}>" - - def get_usage(self, ctx: Context) -> str: - """Formats the usage line into a string and returns it. - - Calls :meth:`format_usage` internally. - """ - formatter = ctx.make_formatter() - self.format_usage(ctx, formatter) - return formatter.getvalue().rstrip("\n") - - def get_params(self, ctx: Context) -> list[Parameter]: - params = self.params - help_option = self.get_help_option(ctx) - - if help_option is not None: - params = [*params, help_option] - - if __debug__: - import warnings - - opts = [opt for param in params for opt in param.opts] - opts_counter = Counter(opts) - duplicate_opts = (opt for opt, count in opts_counter.items() if count > 1) - - for duplicate_opt in duplicate_opts: - warnings.warn( - ( - f"The parameter {duplicate_opt} is used more than once. " - "Remove its duplicate as parameters should be unique." - ), - stacklevel=3, - ) - - return params - - def format_usage(self, ctx: Context, formatter: HelpFormatter) -> None: - """Writes the usage line into the formatter. - - This is a low-level method called by :meth:`get_usage`. - """ - pieces = self.collect_usage_pieces(ctx) - formatter.write_usage(ctx.command_path, " ".join(pieces)) - - def collect_usage_pieces(self, ctx: Context) -> list[str]: - """Returns all the pieces that go into the usage line and returns - it as a list of strings. - """ - rv = [self.options_metavar] if self.options_metavar else [] - - for param in self.get_params(ctx): - rv.extend(param.get_usage_pieces(ctx)) - - return rv - - def get_help_option_names(self, ctx: Context) -> list[str]: - """Returns the names for the help option.""" - all_names = set(ctx.help_option_names) - for param in self.params: - all_names.difference_update(param.opts) - all_names.difference_update(param.secondary_opts) - return list(all_names) - - def get_help_option(self, ctx: Context) -> Option | None: - """Returns the help option object. - - Skipped if :attr:`add_help_option` is ``False``. - - .. versionchanged:: 8.1.8 - The help option is now cached to avoid creating it multiple times. - """ - help_option_names = self.get_help_option_names(ctx) - - if not help_option_names or not self.add_help_option: - return None - - # Cache the help option object in private _help_option attribute to - # avoid creating it multiple times. Not doing this will break the - # callback odering by iter_params_for_processing(), which relies on - # object comparison. - if self._help_option is None: - # Avoid circular import. - from .decorators import help_option - - # Apply help_option decorator and pop resulting option - help_option(*help_option_names)(self) - self._help_option = self.params.pop() # type: ignore[assignment] - - return self._help_option - - def make_parser(self, ctx: Context) -> _OptionParser: - """Creates the underlying option parser for this command.""" - parser = _OptionParser(ctx) - for param in self.get_params(ctx): - param.add_to_parser(parser, ctx) - return parser - - def get_help(self, ctx: Context) -> str: - """Formats the help into a string and returns it. - - Calls :meth:`format_help` internally. - """ - formatter = ctx.make_formatter() - self.format_help(ctx, formatter) - return formatter.getvalue().rstrip("\n") - - def get_short_help_str(self, limit: int = 45) -> str: - """Gets short help for the command or makes it by shortening the - long help string. - """ - if self.short_help: - text = inspect.cleandoc(self.short_help) - elif self.help: - text = make_default_short_help(self.help, limit) - else: - text = "" - - if self.deprecated: - deprecated_message = ( - f"(DEPRECATED: {self.deprecated})" - if isinstance(self.deprecated, str) - else "(DEPRECATED)" - ) - text = _("{text} {deprecated_message}").format( - text=text, deprecated_message=deprecated_message - ) - - return text.strip() - - def format_help(self, ctx: Context, formatter: HelpFormatter) -> None: - """Writes the help into the formatter if it exists. - - This is a low-level method called by :meth:`get_help`. - - This calls the following methods: - - - :meth:`format_usage` - - :meth:`format_help_text` - - :meth:`format_options` - - :meth:`format_epilog` - """ - self.format_usage(ctx, formatter) - self.format_help_text(ctx, formatter) - self.format_options(ctx, formatter) - self.format_epilog(ctx, formatter) - - def format_help_text(self, ctx: Context, formatter: HelpFormatter) -> None: - """Writes the help text to the formatter if it exists.""" - if self.help is not None: - # truncate the help text to the first form feed - text = inspect.cleandoc(self.help).partition("\f")[0] - else: - text = "" - - if self.deprecated: - deprecated_message = ( - f"(DEPRECATED: {self.deprecated})" - if isinstance(self.deprecated, str) - else "(DEPRECATED)" - ) - text = _("{text} {deprecated_message}").format( - text=text, deprecated_message=deprecated_message - ) - - if text: - formatter.write_paragraph() - - with formatter.indentation(): - formatter.write_text(text) - - def format_options(self, ctx: Context, formatter: HelpFormatter) -> None: - """Writes all the options into the formatter if they exist.""" - opts = [] - for param in self.get_params(ctx): - rv = param.get_help_record(ctx) - if rv is not None: - opts.append(rv) - - if opts: - with formatter.section(_("Options")): - formatter.write_dl(opts) - - def format_epilog(self, ctx: Context, formatter: HelpFormatter) -> None: - """Writes the epilog into the formatter if it exists.""" - if self.epilog: - epilog = inspect.cleandoc(self.epilog) - formatter.write_paragraph() - - with formatter.indentation(): - formatter.write_text(epilog) - - def make_context( - self, - info_name: str | None, - args: list[str], - parent: Context | None = None, - **extra: t.Any, - ) -> Context: - """This function when given an info name and arguments will kick - off the parsing and create a new :class:`Context`. It does not - invoke the actual command callback though. - - To quickly customize the context class used without overriding - this method, set the :attr:`context_class` attribute. - - :param info_name: the info name for this invocation. Generally this - is the most descriptive name for the script or - command. For the toplevel script it's usually - the name of the script, for commands below it's - the name of the command. - :param args: the arguments to parse as list of strings. - :param parent: the parent context if available. - :param extra: extra keyword arguments forwarded to the context - constructor. - - .. versionchanged:: 8.0 - Added the :attr:`context_class` attribute. - """ - for key, value in self.context_settings.items(): - if key not in extra: - extra[key] = value - - ctx = self.context_class(self, info_name=info_name, parent=parent, **extra) - - with ctx.scope(cleanup=False): - self.parse_args(ctx, args) - return ctx - - def parse_args(self, ctx: Context, args: list[str]) -> list[str]: - if not args and self.no_args_is_help and not ctx.resilient_parsing: - raise NoArgsIsHelpError(ctx) - - parser = self.make_parser(ctx) - opts, args, param_order = parser.parse_args(args=args) - - for param in iter_params_for_processing(param_order, self.get_params(ctx)): - value, args = param.handle_parse_result(ctx, opts, args) - - if args and not ctx.allow_extra_args and not ctx.resilient_parsing: - ctx.fail( - ngettext( - "Got unexpected extra argument ({args})", - "Got unexpected extra arguments ({args})", - len(args), - ).format(args=" ".join(map(str, args))) - ) - - ctx.args = args - ctx._opt_prefixes.update(parser._opt_prefixes) - return args - - def invoke(self, ctx: Context) -> t.Any: - """Given a context, this invokes the attached callback (if it exists) - in the right way. - """ - if self.deprecated: - extra_message = ( - f" {self.deprecated}" if isinstance(self.deprecated, str) else "" - ) - message = _( - "DeprecationWarning: The command {name!r} is deprecated.{extra_message}" - ).format(name=self.name, extra_message=extra_message) - echo(style(message, fg="red"), err=True) - - if self.callback is not None: - return ctx.invoke(self.callback, **ctx.params) - - def shell_complete(self, ctx: Context, incomplete: str) -> list[CompletionItem]: - """Return a list of completions for the incomplete value. Looks - at the names of options and chained multi-commands. - - Any command could be part of a chained multi-command, so sibling - commands are valid at any point during command completion. - - :param ctx: Invocation context for this command. - :param incomplete: Value being completed. May be empty. - - .. versionadded:: 8.0 - """ - from click.shell_completion import CompletionItem - - results: list[CompletionItem] = [] - - if incomplete and not incomplete[0].isalnum(): - for param in self.get_params(ctx): - if ( - not isinstance(param, Option) - or param.hidden - or ( - not param.multiple - and ctx.get_parameter_source(param.name) # type: ignore - is ParameterSource.COMMANDLINE - ) - ): - continue - - results.extend( - CompletionItem(name, help=param.help) - for name in [*param.opts, *param.secondary_opts] - if name.startswith(incomplete) - ) - - while ctx.parent is not None: - ctx = ctx.parent - - if isinstance(ctx.command, Group) and ctx.command.chain: - results.extend( - CompletionItem(name, help=command.get_short_help_str()) - for name, command in _complete_visible_commands(ctx, incomplete) - if name not in ctx._protected_args - ) - - return results - - @t.overload - def main( - self, - args: cabc.Sequence[str] | None = None, - prog_name: str | None = None, - complete_var: str | None = None, - standalone_mode: t.Literal[True] = True, - **extra: t.Any, - ) -> t.NoReturn: ... - - @t.overload - def main( - self, - args: cabc.Sequence[str] | None = None, - prog_name: str | None = None, - complete_var: str | None = None, - standalone_mode: bool = ..., - **extra: t.Any, - ) -> t.Any: ... - - def main( - self, - args: cabc.Sequence[str] | None = None, - prog_name: str | None = None, - complete_var: str | None = None, - standalone_mode: bool = True, - windows_expand_args: bool = True, - **extra: t.Any, - ) -> t.Any: - """This is the way to invoke a script with all the bells and - whistles as a command line application. This will always terminate - the application after a call. If this is not wanted, ``SystemExit`` - needs to be caught. - - This method is also available by directly calling the instance of - a :class:`Command`. - - :param args: the arguments that should be used for parsing. If not - provided, ``sys.argv[1:]`` is used. - :param prog_name: the program name that should be used. By default - the program name is constructed by taking the file - name from ``sys.argv[0]``. - :param complete_var: the environment variable that controls the - bash completion support. The default is - ``"__COMPLETE"`` with prog_name in - uppercase. - :param standalone_mode: the default behavior is to invoke the script - in standalone mode. Click will then - handle exceptions and convert them into - error messages and the function will never - return but shut down the interpreter. If - this is set to `False` they will be - propagated to the caller and the return - value of this function is the return value - of :meth:`invoke`. - :param windows_expand_args: Expand glob patterns, user dir, and - env vars in command line args on Windows. - :param extra: extra keyword arguments are forwarded to the context - constructor. See :class:`Context` for more information. - - .. versionchanged:: 8.0.1 - Added the ``windows_expand_args`` parameter to allow - disabling command line arg expansion on Windows. - - .. versionchanged:: 8.0 - When taking arguments from ``sys.argv`` on Windows, glob - patterns, user dir, and env vars are expanded. - - .. versionchanged:: 3.0 - Added the ``standalone_mode`` parameter. - """ - if args is None: - args = sys.argv[1:] - - if os.name == "nt" and windows_expand_args: - args = _expand_args(args) - else: - args = list(args) - - if prog_name is None: - prog_name = _detect_program_name() - - # Process shell completion requests and exit early. - self._main_shell_completion(extra, prog_name, complete_var) - - try: - try: - with self.make_context(prog_name, args, **extra) as ctx: - rv = self.invoke(ctx) - if not standalone_mode: - return rv - # it's not safe to `ctx.exit(rv)` here! - # note that `rv` may actually contain data like "1" which - # has obvious effects - # more subtle case: `rv=[None, None]` can come out of - # chained commands which all returned `None` -- so it's not - # even always obvious that `rv` indicates success/failure - # by its truthiness/falsiness - ctx.exit() - except (EOFError, KeyboardInterrupt) as e: - echo(file=sys.stderr) - raise Abort() from e - except ClickException as e: - if not standalone_mode: - raise - e.show() - sys.exit(e.exit_code) - except OSError as e: - if e.errno == errno.EPIPE: - sys.stdout = t.cast(t.TextIO, PacifyFlushWrapper(sys.stdout)) - sys.stderr = t.cast(t.TextIO, PacifyFlushWrapper(sys.stderr)) - sys.exit(1) - else: - raise - except Exit as e: - if standalone_mode: - sys.exit(e.exit_code) - else: - # in non-standalone mode, return the exit code - # note that this is only reached if `self.invoke` above raises - # an Exit explicitly -- thus bypassing the check there which - # would return its result - # the results of non-standalone execution may therefore be - # somewhat ambiguous: if there are codepaths which lead to - # `ctx.exit(1)` and to `return 1`, the caller won't be able to - # tell the difference between the two - return e.exit_code - except Abort: - if not standalone_mode: - raise - echo(_("Aborted!"), file=sys.stderr) - sys.exit(1) - - def _main_shell_completion( - self, - ctx_args: cabc.MutableMapping[str, t.Any], - prog_name: str, - complete_var: str | None = None, - ) -> None: - """Check if the shell is asking for tab completion, process - that, then exit early. Called from :meth:`main` before the - program is invoked. - - :param prog_name: Name of the executable in the shell. - :param complete_var: Name of the environment variable that holds - the completion instruction. Defaults to - ``_{PROG_NAME}_COMPLETE``. - - .. versionchanged:: 8.2.0 - Dots (``.``) in ``prog_name`` are replaced with underscores (``_``). - """ - if complete_var is None: - complete_name = prog_name.replace("-", "_").replace(".", "_") - complete_var = f"_{complete_name}_COMPLETE".upper() - - instruction = os.environ.get(complete_var) - - if not instruction: - return - - from .shell_completion import shell_complete - - rv = shell_complete(self, ctx_args, prog_name, complete_var, instruction) - sys.exit(rv) - - def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any: - """Alias for :meth:`main`.""" - return self.main(*args, **kwargs) - - -class _FakeSubclassCheck(type): - def __subclasscheck__(cls, subclass: type) -> bool: - return issubclass(subclass, cls.__bases__[0]) - - def __instancecheck__(cls, instance: t.Any) -> bool: - return isinstance(instance, cls.__bases__[0]) - - -class _BaseCommand(Command, metaclass=_FakeSubclassCheck): - """ - .. deprecated:: 8.2 - Will be removed in Click 9.0. Use ``Command`` instead. - """ - - -class Group(Command): - """A group is a command that nests other commands (or more groups). - - :param name: The name of the group command. - :param commands: Map names to :class:`Command` objects. Can be a list, which - will use :attr:`Command.name` as the keys. - :param invoke_without_command: Invoke the group's callback even if a - subcommand is not given. - :param no_args_is_help: If no arguments are given, show the group's help and - exit. Defaults to the opposite of ``invoke_without_command``. - :param subcommand_metavar: How to represent the subcommand argument in help. - The default will represent whether ``chain`` is set or not. - :param chain: Allow passing more than one subcommand argument. After parsing - a command's arguments, if any arguments remain another command will be - matched, and so on. - :param result_callback: A function to call after the group's and - subcommand's callbacks. The value returned by the subcommand is passed. - If ``chain`` is enabled, the value will be a list of values returned by - all the commands. If ``invoke_without_command`` is enabled, the value - will be the value returned by the group's callback, or an empty list if - ``chain`` is enabled. - :param kwargs: Other arguments passed to :class:`Command`. - - .. versionchanged:: 8.0 - The ``commands`` argument can be a list of command objects. - - .. versionchanged:: 8.2 - Merged with and replaces the ``MultiCommand`` base class. - """ - - allow_extra_args = True - allow_interspersed_args = False - - #: If set, this is used by the group's :meth:`command` decorator - #: as the default :class:`Command` class. This is useful to make all - #: subcommands use a custom command class. - #: - #: .. versionadded:: 8.0 - command_class: type[Command] | None = None - - #: If set, this is used by the group's :meth:`group` decorator - #: as the default :class:`Group` class. This is useful to make all - #: subgroups use a custom group class. - #: - #: If set to the special value :class:`type` (literally - #: ``group_class = type``), this group's class will be used as the - #: default class. This makes a custom group class continue to make - #: custom groups. - #: - #: .. versionadded:: 8.0 - group_class: type[Group] | type[type] | None = None - # Literal[type] isn't valid, so use Type[type] - - def __init__( - self, - name: str | None = None, - commands: cabc.MutableMapping[str, Command] - | cabc.Sequence[Command] - | None = None, - invoke_without_command: bool = False, - no_args_is_help: bool | None = None, - subcommand_metavar: str | None = None, - chain: bool = False, - result_callback: t.Callable[..., t.Any] | None = None, - **kwargs: t.Any, - ) -> None: - super().__init__(name, **kwargs) - - if commands is None: - commands = {} - elif isinstance(commands, abc.Sequence): - commands = {c.name: c for c in commands if c.name is not None} - - #: The registered subcommands by their exported names. - self.commands: cabc.MutableMapping[str, Command] = commands - - if no_args_is_help is None: - no_args_is_help = not invoke_without_command - - self.no_args_is_help = no_args_is_help - self.invoke_without_command = invoke_without_command - - if subcommand_metavar is None: - if chain: - subcommand_metavar = "COMMAND1 [ARGS]... [COMMAND2 [ARGS]...]..." - else: - subcommand_metavar = "COMMAND [ARGS]..." - - self.subcommand_metavar = subcommand_metavar - self.chain = chain - # The result callback that is stored. This can be set or - # overridden with the :func:`result_callback` decorator. - self._result_callback = result_callback - - if self.chain: - for param in self.params: - if isinstance(param, Argument) and not param.required: - raise RuntimeError( - "A group in chain mode cannot have optional arguments." - ) - - def to_info_dict(self, ctx: Context) -> dict[str, t.Any]: - info_dict = super().to_info_dict(ctx) - commands = {} - - for name in self.list_commands(ctx): - command = self.get_command(ctx, name) - - if command is None: - continue - - sub_ctx = ctx._make_sub_context(command) - - with sub_ctx.scope(cleanup=False): - commands[name] = command.to_info_dict(sub_ctx) - - info_dict.update(commands=commands, chain=self.chain) - return info_dict - - def add_command(self, cmd: Command, name: str | None = None) -> None: - """Registers another :class:`Command` with this group. If the name - is not provided, the name of the command is used. - """ - name = name or cmd.name - if name is None: - raise TypeError("Command has no name.") - _check_nested_chain(self, name, cmd, register=True) - self.commands[name] = cmd - - @t.overload - def command(self, __func: t.Callable[..., t.Any]) -> Command: ... - - @t.overload - def command( - self, *args: t.Any, **kwargs: t.Any - ) -> t.Callable[[t.Callable[..., t.Any]], Command]: ... - - def command( - self, *args: t.Any, **kwargs: t.Any - ) -> t.Callable[[t.Callable[..., t.Any]], Command] | Command: - """A shortcut decorator for declaring and attaching a command to - the group. This takes the same arguments as :func:`command` and - immediately registers the created command with this group by - calling :meth:`add_command`. - - To customize the command class used, set the - :attr:`command_class` attribute. - - .. versionchanged:: 8.1 - This decorator can be applied without parentheses. - - .. versionchanged:: 8.0 - Added the :attr:`command_class` attribute. - """ - from .decorators import command - - func: t.Callable[..., t.Any] | None = None - - if args and callable(args[0]): - assert len(args) == 1 and not kwargs, ( - "Use 'command(**kwargs)(callable)' to provide arguments." - ) - (func,) = args - args = () - - if self.command_class and kwargs.get("cls") is None: - kwargs["cls"] = self.command_class - - def decorator(f: t.Callable[..., t.Any]) -> Command: - cmd: Command = command(*args, **kwargs)(f) - self.add_command(cmd) - return cmd - - if func is not None: - return decorator(func) - - return decorator - - @t.overload - def group(self, __func: t.Callable[..., t.Any]) -> Group: ... - - @t.overload - def group( - self, *args: t.Any, **kwargs: t.Any - ) -> t.Callable[[t.Callable[..., t.Any]], Group]: ... - - def group( - self, *args: t.Any, **kwargs: t.Any - ) -> t.Callable[[t.Callable[..., t.Any]], Group] | Group: - """A shortcut decorator for declaring and attaching a group to - the group. This takes the same arguments as :func:`group` and - immediately registers the created group with this group by - calling :meth:`add_command`. - - To customize the group class used, set the :attr:`group_class` - attribute. - - .. versionchanged:: 8.1 - This decorator can be applied without parentheses. - - .. versionchanged:: 8.0 - Added the :attr:`group_class` attribute. - """ - from .decorators import group - - func: t.Callable[..., t.Any] | None = None - - if args and callable(args[0]): - assert len(args) == 1 and not kwargs, ( - "Use 'group(**kwargs)(callable)' to provide arguments." - ) - (func,) = args - args = () - - if self.group_class is not None and kwargs.get("cls") is None: - if self.group_class is type: - kwargs["cls"] = type(self) - else: - kwargs["cls"] = self.group_class - - def decorator(f: t.Callable[..., t.Any]) -> Group: - cmd: Group = group(*args, **kwargs)(f) - self.add_command(cmd) - return cmd - - if func is not None: - return decorator(func) - - return decorator - - def result_callback(self, replace: bool = False) -> t.Callable[[F], F]: - """Adds a result callback to the command. By default if a - result callback is already registered this will chain them but - this can be disabled with the `replace` parameter. The result - callback is invoked with the return value of the subcommand - (or the list of return values from all subcommands if chaining - is enabled) as well as the parameters as they would be passed - to the main callback. - - Example:: - - @click.group() - @click.option('-i', '--input', default=23) - def cli(input): - return 42 - - @cli.result_callback() - def process_result(result, input): - return result + input - - :param replace: if set to `True` an already existing result - callback will be removed. - - .. versionchanged:: 8.0 - Renamed from ``resultcallback``. - - .. versionadded:: 3.0 - """ - - def decorator(f: F) -> F: - old_callback = self._result_callback - - if old_callback is None or replace: - self._result_callback = f - return f - - def function(value: t.Any, /, *args: t.Any, **kwargs: t.Any) -> t.Any: - inner = old_callback(value, *args, **kwargs) - return f(inner, *args, **kwargs) - - self._result_callback = rv = update_wrapper(t.cast(F, function), f) - return rv # type: ignore[return-value] - - return decorator - - def get_command(self, ctx: Context, cmd_name: str) -> Command | None: - """Given a context and a command name, this returns a :class:`Command` - object if it exists or returns ``None``. - """ - return self.commands.get(cmd_name) - - def list_commands(self, ctx: Context) -> list[str]: - """Returns a list of subcommand names in the order they should appear.""" - return sorted(self.commands) - - def collect_usage_pieces(self, ctx: Context) -> list[str]: - rv = super().collect_usage_pieces(ctx) - rv.append(self.subcommand_metavar) - return rv - - def format_options(self, ctx: Context, formatter: HelpFormatter) -> None: - super().format_options(ctx, formatter) - self.format_commands(ctx, formatter) - - def format_commands(self, ctx: Context, formatter: HelpFormatter) -> None: - """Extra format methods for multi methods that adds all the commands - after the options. - """ - commands = [] - for subcommand in self.list_commands(ctx): - cmd = self.get_command(ctx, subcommand) - # What is this, the tool lied about a command. Ignore it - if cmd is None: - continue - if cmd.hidden: - continue - - commands.append((subcommand, cmd)) - - # allow for 3 times the default spacing - if len(commands): - limit = formatter.width - 6 - max(len(cmd[0]) for cmd in commands) - - rows = [] - for subcommand, cmd in commands: - help = cmd.get_short_help_str(limit) - rows.append((subcommand, help)) - - if rows: - with formatter.section(_("Commands")): - formatter.write_dl(rows) - - def parse_args(self, ctx: Context, args: list[str]) -> list[str]: - if not args and self.no_args_is_help and not ctx.resilient_parsing: - raise NoArgsIsHelpError(ctx) - - rest = super().parse_args(ctx, args) - - if self.chain: - ctx._protected_args = rest - ctx.args = [] - elif rest: - ctx._protected_args, ctx.args = rest[:1], rest[1:] - - return ctx.args - - def invoke(self, ctx: Context) -> t.Any: - def _process_result(value: t.Any) -> t.Any: - if self._result_callback is not None: - value = ctx.invoke(self._result_callback, value, **ctx.params) - return value - - if not ctx._protected_args: - if self.invoke_without_command: - # No subcommand was invoked, so the result callback is - # invoked with the group return value for regular - # groups, or an empty list for chained groups. - with ctx: - rv = super().invoke(ctx) - return _process_result([] if self.chain else rv) - ctx.fail(_("Missing command.")) - - # Fetch args back out - args = [*ctx._protected_args, *ctx.args] - ctx.args = [] - ctx._protected_args = [] - - # If we're not in chain mode, we only allow the invocation of a - # single command but we also inform the current context about the - # name of the command to invoke. - if not self.chain: - # Make sure the context is entered so we do not clean up - # resources until the result processor has worked. - with ctx: - cmd_name, cmd, args = self.resolve_command(ctx, args) - assert cmd is not None - ctx.invoked_subcommand = cmd_name - super().invoke(ctx) - sub_ctx = cmd.make_context(cmd_name, args, parent=ctx) - with sub_ctx: - return _process_result(sub_ctx.command.invoke(sub_ctx)) - - # In chain mode we create the contexts step by step, but after the - # base command has been invoked. Because at that point we do not - # know the subcommands yet, the invoked subcommand attribute is - # set to ``*`` to inform the command that subcommands are executed - # but nothing else. - with ctx: - ctx.invoked_subcommand = "*" if args else None - super().invoke(ctx) - - # Otherwise we make every single context and invoke them in a - # chain. In that case the return value to the result processor - # is the list of all invoked subcommand's results. - contexts = [] - while args: - cmd_name, cmd, args = self.resolve_command(ctx, args) - assert cmd is not None - sub_ctx = cmd.make_context( - cmd_name, - args, - parent=ctx, - allow_extra_args=True, - allow_interspersed_args=False, - ) - contexts.append(sub_ctx) - args, sub_ctx.args = sub_ctx.args, [] - - rv = [] - for sub_ctx in contexts: - with sub_ctx: - rv.append(sub_ctx.command.invoke(sub_ctx)) - return _process_result(rv) - - def resolve_command( - self, ctx: Context, args: list[str] - ) -> tuple[str | None, Command | None, list[str]]: - cmd_name = make_str(args[0]) - original_cmd_name = cmd_name - - # Get the command - cmd = self.get_command(ctx, cmd_name) - - # If we can't find the command but there is a normalization - # function available, we try with that one. - if cmd is None and ctx.token_normalize_func is not None: - cmd_name = ctx.token_normalize_func(cmd_name) - cmd = self.get_command(ctx, cmd_name) - - # If we don't find the command we want to show an error message - # to the user that it was not provided. However, there is - # something else we should do: if the first argument looks like - # an option we want to kick off parsing again for arguments to - # resolve things like --help which now should go to the main - # place. - if cmd is None and not ctx.resilient_parsing: - if _split_opt(cmd_name)[0]: - self.parse_args(ctx, args) - ctx.fail(_("No such command {name!r}.").format(name=original_cmd_name)) - return cmd_name if cmd else None, cmd, args[1:] - - def shell_complete(self, ctx: Context, incomplete: str) -> list[CompletionItem]: - """Return a list of completions for the incomplete value. Looks - at the names of options, subcommands, and chained - multi-commands. - - :param ctx: Invocation context for this command. - :param incomplete: Value being completed. May be empty. - - .. versionadded:: 8.0 - """ - from click.shell_completion import CompletionItem - - results = [ - CompletionItem(name, help=command.get_short_help_str()) - for name, command in _complete_visible_commands(ctx, incomplete) - ] - results.extend(super().shell_complete(ctx, incomplete)) - return results - - -class _MultiCommand(Group, metaclass=_FakeSubclassCheck): - """ - .. deprecated:: 8.2 - Will be removed in Click 9.0. Use ``Group`` instead. - """ - - -class CommandCollection(Group): - """A :class:`Group` that looks up subcommands on other groups. If a command - is not found on this group, each registered source is checked in order. - Parameters on a source are not added to this group, and a source's callback - is not invoked when invoking its commands. In other words, this "flattens" - commands in many groups into this one group. - - :param name: The name of the group command. - :param sources: A list of :class:`Group` objects to look up commands from. - :param kwargs: Other arguments passed to :class:`Group`. - - .. versionchanged:: 8.2 - This is a subclass of ``Group``. Commands are looked up first on this - group, then each of its sources. - """ - - def __init__( - self, - name: str | None = None, - sources: list[Group] | None = None, - **kwargs: t.Any, - ) -> None: - super().__init__(name, **kwargs) - #: The list of registered groups. - self.sources: list[Group] = sources or [] - - def add_source(self, group: Group) -> None: - """Add a group as a source of commands.""" - self.sources.append(group) - - def get_command(self, ctx: Context, cmd_name: str) -> Command | None: - rv = super().get_command(ctx, cmd_name) - - if rv is not None: - return rv - - for source in self.sources: - rv = source.get_command(ctx, cmd_name) - - if rv is not None: - if self.chain: - _check_nested_chain(self, cmd_name, rv) - - return rv - - return None - - def list_commands(self, ctx: Context) -> list[str]: - rv: set[str] = set(super().list_commands(ctx)) - - for source in self.sources: - rv.update(source.list_commands(ctx)) - - return sorted(rv) - - -def _check_iter(value: t.Any) -> cabc.Iterator[t.Any]: - """Check if the value is iterable but not a string. Raises a type - error, or return an iterator over the value. - """ - if isinstance(value, str): - raise TypeError - - return iter(value) - - -class Parameter: - r"""A parameter to a command comes in two versions: they are either - :class:`Option`\s or :class:`Argument`\s. Other subclasses are currently - not supported by design as some of the internals for parsing are - intentionally not finalized. - - Some settings are supported by both options and arguments. - - :param param_decls: the parameter declarations for this option or - argument. This is a list of flags or argument - names. - :param type: the type that should be used. Either a :class:`ParamType` - or a Python type. The latter is converted into the former - automatically if supported. - :param required: controls if this is optional or not. - :param default: the default value if omitted. This can also be a callable, - in which case it's invoked when the default is needed - without any arguments. - :param callback: A function to further process or validate the value - after type conversion. It is called as ``f(ctx, param, value)`` - and must return the value. It is called for all sources, - including prompts. - :param nargs: the number of arguments to match. If not ``1`` the return - value is a tuple instead of single value. The default for - nargs is ``1`` (except if the type is a tuple, then it's - the arity of the tuple). If ``nargs=-1``, all remaining - parameters are collected. - :param metavar: how the value is represented in the help page. - :param expose_value: if this is `True` then the value is passed onwards - to the command callback and stored on the context, - otherwise it's skipped. - :param is_eager: eager values are processed before non eager ones. This - should not be set for arguments or it will inverse the - order of processing. - :param envvar: a string or list of strings that are environment variables - that should be checked. - :param shell_complete: A function that returns custom shell - completions. Used instead of the param's type completion if - given. Takes ``ctx, param, incomplete`` and must return a list - of :class:`~click.shell_completion.CompletionItem` or a list of - strings. - :param deprecated: If ``True`` or non-empty string, issues a message - indicating that the argument is deprecated and highlights - its deprecation in --help. The message can be customized - by using a string as the value. A deprecated parameter - cannot be required, a ValueError will be raised otherwise. - - .. versionchanged:: 8.2.0 - Introduction of ``deprecated``. - - .. versionchanged:: 8.2 - Adding duplicate parameter names to a :class:`~click.core.Command` will - result in a ``UserWarning`` being shown. - - .. versionchanged:: 8.2 - Adding duplicate parameter names to a :class:`~click.core.Command` will - result in a ``UserWarning`` being shown. - - .. versionchanged:: 8.0 - ``process_value`` validates required parameters and bounded - ``nargs``, and invokes the parameter callback before returning - the value. This allows the callback to validate prompts. - ``full_process_value`` is removed. - - .. versionchanged:: 8.0 - ``autocompletion`` is renamed to ``shell_complete`` and has new - semantics described above. The old name is deprecated and will - be removed in 8.1, until then it will be wrapped to match the - new requirements. - - .. versionchanged:: 8.0 - For ``multiple=True, nargs>1``, the default must be a list of - tuples. - - .. versionchanged:: 8.0 - Setting a default is no longer required for ``nargs>1``, it will - default to ``None``. ``multiple=True`` or ``nargs=-1`` will - default to ``()``. - - .. versionchanged:: 7.1 - Empty environment variables are ignored rather than taking the - empty string value. This makes it possible for scripts to clear - variables if they can't unset them. - - .. versionchanged:: 2.0 - Changed signature for parameter callback to also be passed the - parameter. The old callback format will still work, but it will - raise a warning to give you a chance to migrate the code easier. - """ - - param_type_name = "parameter" - - def __init__( - self, - param_decls: cabc.Sequence[str] | None = None, - type: types.ParamType | t.Any | None = None, - required: bool = False, - default: t.Any | t.Callable[[], t.Any] | None = None, - callback: t.Callable[[Context, Parameter, t.Any], t.Any] | None = None, - nargs: int | None = None, - multiple: bool = False, - metavar: str | None = None, - expose_value: bool = True, - is_eager: bool = False, - envvar: str | cabc.Sequence[str] | None = None, - shell_complete: t.Callable[ - [Context, Parameter, str], list[CompletionItem] | list[str] - ] - | None = None, - deprecated: bool | str = False, - ) -> None: - self.name: str | None - self.opts: list[str] - self.secondary_opts: list[str] - self.name, self.opts, self.secondary_opts = self._parse_decls( - param_decls or (), expose_value - ) - self.type: types.ParamType = types.convert_type(type, default) - - # Default nargs to what the type tells us if we have that - # information available. - if nargs is None: - if self.type.is_composite: - nargs = self.type.arity - else: - nargs = 1 - - self.required = required - self.callback = callback - self.nargs = nargs - self.multiple = multiple - self.expose_value = expose_value - self.default = default - self.is_eager = is_eager - self.metavar = metavar - self.envvar = envvar - self._custom_shell_complete = shell_complete - self.deprecated = deprecated - - if __debug__: - if self.type.is_composite and nargs != self.type.arity: - raise ValueError( - f"'nargs' must be {self.type.arity} (or None) for" - f" type {self.type!r}, but it was {nargs}." - ) - - # Skip no default or callable default. - check_default = default if not callable(default) else None - - if check_default is not None: - if multiple: - try: - # Only check the first value against nargs. - check_default = next(_check_iter(check_default), None) - except TypeError: - raise ValueError( - "'default' must be a list when 'multiple' is true." - ) from None - - # Can be None for multiple with empty default. - if nargs != 1 and check_default is not None: - try: - _check_iter(check_default) - except TypeError: - if multiple: - message = ( - "'default' must be a list of lists when 'multiple' is" - " true and 'nargs' != 1." - ) - else: - message = "'default' must be a list when 'nargs' != 1." - - raise ValueError(message) from None - - if nargs > 1 and len(check_default) != nargs: - subject = "item length" if multiple else "length" - raise ValueError( - f"'default' {subject} must match nargs={nargs}." - ) - - if required and deprecated: - raise ValueError( - f"The {self.param_type_name} '{self.human_readable_name}' " - "is deprecated and still required. A deprecated " - f"{self.param_type_name} cannot be required." - ) - - def to_info_dict(self) -> dict[str, t.Any]: - """Gather information that could be useful for a tool generating - user-facing documentation. - - Use :meth:`click.Context.to_info_dict` to traverse the entire - CLI structure. - - .. versionadded:: 8.0 - """ - return { - "name": self.name, - "param_type_name": self.param_type_name, - "opts": self.opts, - "secondary_opts": self.secondary_opts, - "type": self.type.to_info_dict(), - "required": self.required, - "nargs": self.nargs, - "multiple": self.multiple, - "default": self.default, - "envvar": self.envvar, - } - - def __repr__(self) -> str: - return f"<{self.__class__.__name__} {self.name}>" - - def _parse_decls( - self, decls: cabc.Sequence[str], expose_value: bool - ) -> tuple[str | None, list[str], list[str]]: - raise NotImplementedError() - - @property - def human_readable_name(self) -> str: - """Returns the human readable name of this parameter. This is the - same as the name for options, but the metavar for arguments. - """ - return self.name # type: ignore - - def make_metavar(self, ctx: Context) -> str: - if self.metavar is not None: - return self.metavar - - metavar = self.type.get_metavar(param=self, ctx=ctx) - - if metavar is None: - metavar = self.type.name.upper() - - if self.nargs != 1: - metavar += "..." - - return metavar - - @t.overload - def get_default( - self, ctx: Context, call: t.Literal[True] = True - ) -> t.Any | None: ... - - @t.overload - def get_default( - self, ctx: Context, call: bool = ... - ) -> t.Any | t.Callable[[], t.Any] | None: ... - - def get_default( - self, ctx: Context, call: bool = True - ) -> t.Any | t.Callable[[], t.Any] | None: - """Get the default for the parameter. Tries - :meth:`Context.lookup_default` first, then the local default. - - :param ctx: Current context. - :param call: If the default is a callable, call it. Disable to - return the callable instead. - - .. versionchanged:: 8.0.2 - Type casting is no longer performed when getting a default. - - .. versionchanged:: 8.0.1 - Type casting can fail in resilient parsing mode. Invalid - defaults will not prevent showing help text. - - .. versionchanged:: 8.0 - Looks at ``ctx.default_map`` first. - - .. versionchanged:: 8.0 - Added the ``call`` parameter. - """ - value = ctx.lookup_default(self.name, call=False) # type: ignore - - if value is None: - value = self.default - - if call and callable(value): - value = value() - - return value - - def add_to_parser(self, parser: _OptionParser, ctx: Context) -> None: - raise NotImplementedError() - - def consume_value( - self, ctx: Context, opts: cabc.Mapping[str, t.Any] - ) -> tuple[t.Any, ParameterSource]: - value = opts.get(self.name) # type: ignore - source = ParameterSource.COMMANDLINE - - if value is None: - value = self.value_from_envvar(ctx) - source = ParameterSource.ENVIRONMENT - - if value is None: - value = ctx.lookup_default(self.name) # type: ignore - source = ParameterSource.DEFAULT_MAP - - if value is None: - value = self.get_default(ctx) - source = ParameterSource.DEFAULT - - return value, source - - def type_cast_value(self, ctx: Context, value: t.Any) -> t.Any: - """Convert and validate a value against the option's - :attr:`type`, :attr:`multiple`, and :attr:`nargs`. - """ - if value is None: - return () if self.multiple or self.nargs == -1 else None - - def check_iter(value: t.Any) -> cabc.Iterator[t.Any]: - try: - return _check_iter(value) - except TypeError: - # This should only happen when passing in args manually, - # the parser should construct an iterable when parsing - # the command line. - raise BadParameter( - _("Value must be an iterable."), ctx=ctx, param=self - ) from None - - if self.nargs == 1 or self.type.is_composite: - - def convert(value: t.Any) -> t.Any: - return self.type(value, param=self, ctx=ctx) - - elif self.nargs == -1: - - def convert(value: t.Any) -> t.Any: # tuple[t.Any, ...] - return tuple(self.type(x, self, ctx) for x in check_iter(value)) - - else: # nargs > 1 - - def convert(value: t.Any) -> t.Any: # tuple[t.Any, ...] - value = tuple(check_iter(value)) - - if len(value) != self.nargs: - raise BadParameter( - ngettext( - "Takes {nargs} values but 1 was given.", - "Takes {nargs} values but {len} were given.", - len(value), - ).format(nargs=self.nargs, len=len(value)), - ctx=ctx, - param=self, - ) - - return tuple(self.type(x, self, ctx) for x in value) - - if self.multiple: - return tuple(convert(x) for x in check_iter(value)) - - return convert(value) - - def value_is_missing(self, value: t.Any) -> bool: - if value is None: - return True - - if (self.nargs != 1 or self.multiple) and value == (): - return True - - return False - - def process_value(self, ctx: Context, value: t.Any) -> t.Any: - value = self.type_cast_value(ctx, value) - - if self.required and self.value_is_missing(value): - raise MissingParameter(ctx=ctx, param=self) - - if self.callback is not None: - value = self.callback(ctx, self, value) - - return value - - def resolve_envvar_value(self, ctx: Context) -> str | None: - if self.envvar is None: - return None - - if isinstance(self.envvar, str): - rv = os.environ.get(self.envvar) - - if rv: - return rv - else: - for envvar in self.envvar: - rv = os.environ.get(envvar) - - if rv: - return rv - - return None - - def value_from_envvar(self, ctx: Context) -> t.Any | None: - rv: t.Any | None = self.resolve_envvar_value(ctx) - - if rv is not None and self.nargs != 1: - rv = self.type.split_envvar_value(rv) - - return rv - - def handle_parse_result( - self, ctx: Context, opts: cabc.Mapping[str, t.Any], args: list[str] - ) -> tuple[t.Any, list[str]]: - with augment_usage_errors(ctx, param=self): - value, source = self.consume_value(ctx, opts) - - if ( - self.deprecated - and value is not None - and source - not in ( - ParameterSource.DEFAULT, - ParameterSource.DEFAULT_MAP, - ) - ): - extra_message = ( - f" {self.deprecated}" if isinstance(self.deprecated, str) else "" - ) - message = _( - "DeprecationWarning: The {param_type} {name!r} is deprecated." - "{extra_message}" - ).format( - param_type=self.param_type_name, - name=self.human_readable_name, - extra_message=extra_message, - ) - echo(style(message, fg="red"), err=True) - - ctx.set_parameter_source(self.name, source) # type: ignore - - try: - value = self.process_value(ctx, value) - except Exception: - if not ctx.resilient_parsing: - raise - - value = None - - if self.expose_value: - ctx.params[self.name] = value # type: ignore - - return value, args - - def get_help_record(self, ctx: Context) -> tuple[str, str] | None: - pass - - def get_usage_pieces(self, ctx: Context) -> list[str]: - return [] - - def get_error_hint(self, ctx: Context) -> str: - """Get a stringified version of the param for use in error messages to - indicate which param caused the error. - """ - hint_list = self.opts or [self.human_readable_name] - return " / ".join(f"'{x}'" for x in hint_list) - - def shell_complete(self, ctx: Context, incomplete: str) -> list[CompletionItem]: - """Return a list of completions for the incomplete value. If a - ``shell_complete`` function was given during init, it is used. - Otherwise, the :attr:`type` - :meth:`~click.types.ParamType.shell_complete` function is used. - - :param ctx: Invocation context for this command. - :param incomplete: Value being completed. May be empty. - - .. versionadded:: 8.0 - """ - if self._custom_shell_complete is not None: - results = self._custom_shell_complete(ctx, self, incomplete) - - if results and isinstance(results[0], str): - from click.shell_completion import CompletionItem - - results = [CompletionItem(c) for c in results] - - return t.cast("list[CompletionItem]", results) - - return self.type.shell_complete(ctx, self, incomplete) - - -class Option(Parameter): - """Options are usually optional values on the command line and - have some extra features that arguments don't have. - - All other parameters are passed onwards to the parameter constructor. - - :param show_default: Show the default value for this option in its - help text. Values are not shown by default, unless - :attr:`Context.show_default` is ``True``. If this value is a - string, it shows that string in parentheses instead of the - actual value. This is particularly useful for dynamic options. - For single option boolean flags, the default remains hidden if - its value is ``False``. - :param show_envvar: Controls if an environment variable should be - shown on the help page and error messages. - Normally, environment variables are not shown. - :param prompt: If set to ``True`` or a non empty string then the - user will be prompted for input. If set to ``True`` the prompt - will be the option name capitalized. A deprecated option cannot be - prompted. - :param confirmation_prompt: Prompt a second time to confirm the - value if it was prompted for. Can be set to a string instead of - ``True`` to customize the message. - :param prompt_required: If set to ``False``, the user will be - prompted for input only when the option was specified as a flag - without a value. - :param hide_input: If this is ``True`` then the input on the prompt - will be hidden from the user. This is useful for password input. - :param is_flag: forces this option to act as a flag. The default is - auto detection. - :param flag_value: which value should be used for this flag if it's - enabled. This is set to a boolean automatically if - the option string contains a slash to mark two options. - :param multiple: if this is set to `True` then the argument is accepted - multiple times and recorded. This is similar to ``nargs`` - in how it works but supports arbitrary number of - arguments. - :param count: this flag makes an option increment an integer. - :param allow_from_autoenv: if this is enabled then the value of this - parameter will be pulled from an environment - variable in case a prefix is defined on the - context. - :param help: the help string. - :param hidden: hide this option from help outputs. - :param attrs: Other command arguments described in :class:`Parameter`. - - .. versionchanged:: 8.2 - ``envvar`` used with ``flag_value`` will always use the ``flag_value``, - previously it would use the value of the environment variable. - - .. versionchanged:: 8.1 - Help text indentation is cleaned here instead of only in the - ``@option`` decorator. - - .. versionchanged:: 8.1 - The ``show_default`` parameter overrides - ``Context.show_default``. - - .. versionchanged:: 8.1 - The default of a single option boolean flag is not shown if the - default value is ``False``. - - .. versionchanged:: 8.0.1 - ``type`` is detected from ``flag_value`` if given. - """ - - param_type_name = "option" - - def __init__( - self, - param_decls: cabc.Sequence[str] | None = None, - show_default: bool | str | None = None, - prompt: bool | str = False, - confirmation_prompt: bool | str = False, - prompt_required: bool = True, - hide_input: bool = False, - is_flag: bool | None = None, - flag_value: t.Any | None = None, - multiple: bool = False, - count: bool = False, - allow_from_autoenv: bool = True, - type: types.ParamType | t.Any | None = None, - help: str | None = None, - hidden: bool = False, - show_choices: bool = True, - show_envvar: bool = False, - deprecated: bool | str = False, - **attrs: t.Any, - ) -> None: - if help: - help = inspect.cleandoc(help) - - default_is_missing = "default" not in attrs - super().__init__( - param_decls, type=type, multiple=multiple, deprecated=deprecated, **attrs - ) - - if prompt is True: - if self.name is None: - raise TypeError("'name' is required with 'prompt=True'.") - - prompt_text: str | None = self.name.replace("_", " ").capitalize() - elif prompt is False: - prompt_text = None - else: - prompt_text = prompt - - if deprecated: - deprecated_message = ( - f"(DEPRECATED: {deprecated})" - if isinstance(deprecated, str) - else "(DEPRECATED)" - ) - help = help + deprecated_message if help is not None else deprecated_message - - self.prompt = prompt_text - self.confirmation_prompt = confirmation_prompt - self.prompt_required = prompt_required - self.hide_input = hide_input - self.hidden = hidden - - # If prompt is enabled but not required, then the option can be - # used as a flag to indicate using prompt or flag_value. - self._flag_needs_value = self.prompt is not None and not self.prompt_required - - if is_flag is None: - if flag_value is not None: - # Implicitly a flag because flag_value was set. - is_flag = True - elif self._flag_needs_value: - # Not a flag, but when used as a flag it shows a prompt. - is_flag = False - else: - # Implicitly a flag because flag options were given. - is_flag = bool(self.secondary_opts) - elif is_flag is False and not self._flag_needs_value: - # Not a flag, and prompt is not enabled, can be used as a - # flag if flag_value is set. - self._flag_needs_value = flag_value is not None - - self.default: t.Any | t.Callable[[], t.Any] - - if is_flag and default_is_missing and not self.required: - if multiple: - self.default = () - else: - self.default = False - - if is_flag and flag_value is None: - flag_value = not self.default - - self.type: types.ParamType - if is_flag and type is None: - # Re-guess the type from the flag value instead of the - # default. - self.type = types.convert_type(None, flag_value) - - self.is_flag: bool = is_flag - self.is_bool_flag: bool = is_flag and isinstance(self.type, types.BoolParamType) - self.flag_value: t.Any = flag_value - - # Counting - self.count = count - if count: - if type is None: - self.type = types.IntRange(min=0) - if default_is_missing: - self.default = 0 - - self.allow_from_autoenv = allow_from_autoenv - self.help = help - self.show_default = show_default - self.show_choices = show_choices - self.show_envvar = show_envvar - - if __debug__: - if deprecated and prompt: - raise ValueError("`deprecated` options cannot use `prompt`.") - - if self.nargs == -1: - raise TypeError("nargs=-1 is not supported for options.") - - if self.prompt and self.is_flag and not self.is_bool_flag: - raise TypeError("'prompt' is not valid for non-boolean flag.") - - if not self.is_bool_flag and self.secondary_opts: - raise TypeError("Secondary flag is not valid for non-boolean flag.") - - if self.is_bool_flag and self.hide_input and self.prompt is not None: - raise TypeError( - "'prompt' with 'hide_input' is not valid for boolean flag." - ) - - if self.count: - if self.multiple: - raise TypeError("'count' is not valid with 'multiple'.") - - if self.is_flag: - raise TypeError("'count' is not valid with 'is_flag'.") - - def to_info_dict(self) -> dict[str, t.Any]: - info_dict = super().to_info_dict() - info_dict.update( - help=self.help, - prompt=self.prompt, - is_flag=self.is_flag, - flag_value=self.flag_value, - count=self.count, - hidden=self.hidden, - ) - return info_dict - - def get_error_hint(self, ctx: Context) -> str: - result = super().get_error_hint(ctx) - if self.show_envvar: - result += f" (env var: '{self.envvar}')" - return result - - def _parse_decls( - self, decls: cabc.Sequence[str], expose_value: bool - ) -> tuple[str | None, list[str], list[str]]: - opts = [] - secondary_opts = [] - name = None - possible_names = [] - - for decl in decls: - if decl.isidentifier(): - if name is not None: - raise TypeError(f"Name '{name}' defined twice") - name = decl - else: - split_char = ";" if decl[:1] == "/" else "/" - if split_char in decl: - first, second = decl.split(split_char, 1) - first = first.rstrip() - if first: - possible_names.append(_split_opt(first)) - opts.append(first) - second = second.lstrip() - if second: - secondary_opts.append(second.lstrip()) - if first == second: - raise ValueError( - f"Boolean option {decl!r} cannot use the" - " same flag for true/false." - ) - else: - possible_names.append(_split_opt(decl)) - opts.append(decl) - - if name is None and possible_names: - possible_names.sort(key=lambda x: -len(x[0])) # group long options first - name = possible_names[0][1].replace("-", "_").lower() - if not name.isidentifier(): - name = None - - if name is None: - if not expose_value: - return None, opts, secondary_opts - raise TypeError( - f"Could not determine name for option with declarations {decls!r}" - ) - - if not opts and not secondary_opts: - raise TypeError( - f"No options defined but a name was passed ({name})." - " Did you mean to declare an argument instead? Did" - f" you mean to pass '--{name}'?" - ) - - return name, opts, secondary_opts - - def add_to_parser(self, parser: _OptionParser, ctx: Context) -> None: - if self.multiple: - action = "append" - elif self.count: - action = "count" - else: - action = "store" - - if self.is_flag: - action = f"{action}_const" - - if self.is_bool_flag and self.secondary_opts: - parser.add_option( - obj=self, opts=self.opts, dest=self.name, action=action, const=True - ) - parser.add_option( - obj=self, - opts=self.secondary_opts, - dest=self.name, - action=action, - const=False, - ) - else: - parser.add_option( - obj=self, - opts=self.opts, - dest=self.name, - action=action, - const=self.flag_value, - ) - else: - parser.add_option( - obj=self, - opts=self.opts, - dest=self.name, - action=action, - nargs=self.nargs, - ) - - def get_help_record(self, ctx: Context) -> tuple[str, str] | None: - if self.hidden: - return None - - any_prefix_is_slash = False - - def _write_opts(opts: cabc.Sequence[str]) -> str: - nonlocal any_prefix_is_slash - - rv, any_slashes = join_options(opts) - - if any_slashes: - any_prefix_is_slash = True - - if not self.is_flag and not self.count: - rv += f" {self.make_metavar(ctx=ctx)}" - - return rv - - rv = [_write_opts(self.opts)] - - if self.secondary_opts: - rv.append(_write_opts(self.secondary_opts)) - - help = self.help or "" - - extra = self.get_help_extra(ctx) - extra_items = [] - if "envvars" in extra: - extra_items.append( - _("env var: {var}").format(var=", ".join(extra["envvars"])) - ) - if "default" in extra: - extra_items.append(_("default: {default}").format(default=extra["default"])) - if "range" in extra: - extra_items.append(extra["range"]) - if "required" in extra: - extra_items.append(_(extra["required"])) - - if extra_items: - extra_str = "; ".join(extra_items) - help = f"{help} [{extra_str}]" if help else f"[{extra_str}]" - - return ("; " if any_prefix_is_slash else " / ").join(rv), help - - def get_help_extra(self, ctx: Context) -> types.OptionHelpExtra: - extra: types.OptionHelpExtra = {} - - if self.show_envvar: - envvar = self.envvar - - if envvar is None: - if ( - self.allow_from_autoenv - and ctx.auto_envvar_prefix is not None - and self.name is not None - ): - envvar = f"{ctx.auto_envvar_prefix}_{self.name.upper()}" - - if envvar is not None: - if isinstance(envvar, str): - extra["envvars"] = (envvar,) - else: - extra["envvars"] = tuple(str(d) for d in envvar) - - # Temporarily enable resilient parsing to avoid type casting - # failing for the default. Might be possible to extend this to - # help formatting in general. - resilient = ctx.resilient_parsing - ctx.resilient_parsing = True - - try: - default_value = self.get_default(ctx, call=False) - finally: - ctx.resilient_parsing = resilient - - show_default = False - show_default_is_str = False - - if self.show_default is not None: - if isinstance(self.show_default, str): - show_default_is_str = show_default = True - else: - show_default = self.show_default - elif ctx.show_default is not None: - show_default = ctx.show_default - - if show_default_is_str or (show_default and (default_value is not None)): - if show_default_is_str: - default_string = f"({self.show_default})" - elif isinstance(default_value, (list, tuple)): - default_string = ", ".join(str(d) for d in default_value) - elif inspect.isfunction(default_value): - default_string = _("(dynamic)") - elif self.is_bool_flag and self.secondary_opts: - # For boolean flags that have distinct True/False opts, - # use the opt without prefix instead of the value. - default_string = _split_opt( - (self.opts if default_value else self.secondary_opts)[0] - )[1] - elif self.is_bool_flag and not self.secondary_opts and not default_value: - default_string = "" - elif default_value == "": - default_string = '""' - else: - default_string = str(default_value) - - if default_string: - extra["default"] = default_string - - if ( - isinstance(self.type, types._NumberRangeBase) - # skip count with default range type - and not (self.count and self.type.min == 0 and self.type.max is None) - ): - range_str = self.type._describe_range() - - if range_str: - extra["range"] = range_str - - if self.required: - extra["required"] = "required" - - return extra - - @t.overload - def get_default( - self, ctx: Context, call: t.Literal[True] = True - ) -> t.Any | None: ... - - @t.overload - def get_default( - self, ctx: Context, call: bool = ... - ) -> t.Any | t.Callable[[], t.Any] | None: ... - - def get_default( - self, ctx: Context, call: bool = True - ) -> t.Any | t.Callable[[], t.Any] | None: - # If we're a non boolean flag our default is more complex because - # we need to look at all flags in the same group to figure out - # if we're the default one in which case we return the flag - # value as default. - if self.is_flag and not self.is_bool_flag: - for param in ctx.command.params: - if param.name == self.name and param.default: - return t.cast(Option, param).flag_value - - return None - - return super().get_default(ctx, call=call) - - def prompt_for_value(self, ctx: Context) -> t.Any: - """This is an alternative flow that can be activated in the full - value processing if a value does not exist. It will prompt the - user until a valid value exists and then returns the processed - value as result. - """ - assert self.prompt is not None - - # Calculate the default before prompting anything to be stable. - default = self.get_default(ctx) - - # If this is a prompt for a flag we need to handle this - # differently. - if self.is_bool_flag: - return confirm(self.prompt, default) - - # If show_default is set to True/False, provide this to `prompt` as well. For - # non-bool values of `show_default`, we use `prompt`'s default behavior - prompt_kwargs: t.Any = {} - if isinstance(self.show_default, bool): - prompt_kwargs["show_default"] = self.show_default - - return prompt( - self.prompt, - default=default, - type=self.type, - hide_input=self.hide_input, - show_choices=self.show_choices, - confirmation_prompt=self.confirmation_prompt, - value_proc=lambda x: self.process_value(ctx, x), - **prompt_kwargs, - ) - - def resolve_envvar_value(self, ctx: Context) -> str | None: - rv = super().resolve_envvar_value(ctx) - - if rv is not None: - if self.is_flag and self.flag_value: - return str(self.flag_value) - return rv - - if ( - self.allow_from_autoenv - and ctx.auto_envvar_prefix is not None - and self.name is not None - ): - envvar = f"{ctx.auto_envvar_prefix}_{self.name.upper()}" - rv = os.environ.get(envvar) - - if rv: - return rv - - return None - - def value_from_envvar(self, ctx: Context) -> t.Any | None: - rv: t.Any | None = self.resolve_envvar_value(ctx) - - if rv is None: - return None - - value_depth = (self.nargs != 1) + bool(self.multiple) - - if value_depth > 0: - rv = self.type.split_envvar_value(rv) - - if self.multiple and self.nargs != 1: - rv = batch(rv, self.nargs) - - return rv - - def consume_value( - self, ctx: Context, opts: cabc.Mapping[str, Parameter] - ) -> tuple[t.Any, ParameterSource]: - value, source = super().consume_value(ctx, opts) - - # The parser will emit a sentinel value if the option can be - # given as a flag without a value. This is different from None - # to distinguish from the flag not being given at all. - if value is _flag_needs_value: - if self.prompt is not None and not ctx.resilient_parsing: - value = self.prompt_for_value(ctx) - source = ParameterSource.PROMPT - else: - value = self.flag_value - source = ParameterSource.COMMANDLINE - - elif ( - self.multiple - and value is not None - and any(v is _flag_needs_value for v in value) - ): - value = [self.flag_value if v is _flag_needs_value else v for v in value] - source = ParameterSource.COMMANDLINE - - # The value wasn't set, or used the param's default, prompt if - # prompting is enabled. - elif ( - source in {None, ParameterSource.DEFAULT} - and self.prompt is not None - and (self.required or self.prompt_required) - and not ctx.resilient_parsing - ): - value = self.prompt_for_value(ctx) - source = ParameterSource.PROMPT - - return value, source - - -class Argument(Parameter): - """Arguments are positional parameters to a command. They generally - provide fewer features than options but can have infinite ``nargs`` - and are required by default. - - All parameters are passed onwards to the constructor of :class:`Parameter`. - """ - - param_type_name = "argument" - - def __init__( - self, - param_decls: cabc.Sequence[str], - required: bool | None = None, - **attrs: t.Any, - ) -> None: - if required is None: - if attrs.get("default") is not None: - required = False - else: - required = attrs.get("nargs", 1) > 0 - - if "multiple" in attrs: - raise TypeError("__init__() got an unexpected keyword argument 'multiple'.") - - super().__init__(param_decls, required=required, **attrs) - - if __debug__: - if self.default is not None and self.nargs == -1: - raise TypeError("'default' is not supported for nargs=-1.") - - @property - def human_readable_name(self) -> str: - if self.metavar is not None: - return self.metavar - return self.name.upper() # type: ignore - - def make_metavar(self, ctx: Context) -> str: - if self.metavar is not None: - return self.metavar - var = self.type.get_metavar(param=self, ctx=ctx) - if not var: - var = self.name.upper() # type: ignore - if self.deprecated: - var += "!" - if not self.required: - var = f"[{var}]" - if self.nargs != 1: - var += "..." - return var - - def _parse_decls( - self, decls: cabc.Sequence[str], expose_value: bool - ) -> tuple[str | None, list[str], list[str]]: - if not decls: - if not expose_value: - return None, [], [] - raise TypeError("Argument is marked as exposed, but does not have a name.") - if len(decls) == 1: - name = arg = decls[0] - name = name.replace("-", "_").lower() - else: - raise TypeError( - "Arguments take exactly one parameter declaration, got" - f" {len(decls)}: {decls}." - ) - return name, [arg], [] - - def get_usage_pieces(self, ctx: Context) -> list[str]: - return [self.make_metavar(ctx)] - - def get_error_hint(self, ctx: Context) -> str: - return f"'{self.make_metavar(ctx)}'" - - def add_to_parser(self, parser: _OptionParser, ctx: Context) -> None: - parser.add_argument(dest=self.name, nargs=self.nargs, obj=self) - - -def __getattr__(name: str) -> object: - import warnings - - if name == "BaseCommand": - warnings.warn( - "'BaseCommand' is deprecated and will be removed in Click 9.0. Use" - " 'Command' instead.", - DeprecationWarning, - stacklevel=2, - ) - return _BaseCommand - - if name == "MultiCommand": - warnings.warn( - "'MultiCommand' is deprecated and will be removed in Click 9.0. Use" - " 'Group' instead.", - DeprecationWarning, - stacklevel=2, - ) - return _MultiCommand - - raise AttributeError(name) diff --git a/venv/lib/python3.10/site-packages/click/decorators.py b/venv/lib/python3.10/site-packages/click/decorators.py deleted file mode 100644 index 21f4c342245fa93a7946b17d7661e0d2fa6a0c29..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/decorators.py +++ /dev/null @@ -1,551 +0,0 @@ -from __future__ import annotations - -import inspect -import typing as t -from functools import update_wrapper -from gettext import gettext as _ - -from .core import Argument -from .core import Command -from .core import Context -from .core import Group -from .core import Option -from .core import Parameter -from .globals import get_current_context -from .utils import echo - -if t.TYPE_CHECKING: - import typing_extensions as te - - P = te.ParamSpec("P") - -R = t.TypeVar("R") -T = t.TypeVar("T") -_AnyCallable = t.Callable[..., t.Any] -FC = t.TypeVar("FC", bound="_AnyCallable | Command") - - -def pass_context(f: t.Callable[te.Concatenate[Context, P], R]) -> t.Callable[P, R]: - """Marks a callback as wanting to receive the current context - object as first argument. - """ - - def new_func(*args: P.args, **kwargs: P.kwargs) -> R: - return f(get_current_context(), *args, **kwargs) - - return update_wrapper(new_func, f) - - -def pass_obj(f: t.Callable[te.Concatenate[T, P], R]) -> t.Callable[P, R]: - """Similar to :func:`pass_context`, but only pass the object on the - context onwards (:attr:`Context.obj`). This is useful if that object - represents the state of a nested system. - """ - - def new_func(*args: P.args, **kwargs: P.kwargs) -> R: - return f(get_current_context().obj, *args, **kwargs) - - return update_wrapper(new_func, f) - - -def make_pass_decorator( - object_type: type[T], ensure: bool = False -) -> t.Callable[[t.Callable[te.Concatenate[T, P], R]], t.Callable[P, R]]: - """Given an object type this creates a decorator that will work - similar to :func:`pass_obj` but instead of passing the object of the - current context, it will find the innermost context of type - :func:`object_type`. - - This generates a decorator that works roughly like this:: - - from functools import update_wrapper - - def decorator(f): - @pass_context - def new_func(ctx, *args, **kwargs): - obj = ctx.find_object(object_type) - return ctx.invoke(f, obj, *args, **kwargs) - return update_wrapper(new_func, f) - return decorator - - :param object_type: the type of the object to pass. - :param ensure: if set to `True`, a new object will be created and - remembered on the context if it's not there yet. - """ - - def decorator(f: t.Callable[te.Concatenate[T, P], R]) -> t.Callable[P, R]: - def new_func(*args: P.args, **kwargs: P.kwargs) -> R: - ctx = get_current_context() - - obj: T | None - if ensure: - obj = ctx.ensure_object(object_type) - else: - obj = ctx.find_object(object_type) - - if obj is None: - raise RuntimeError( - "Managed to invoke callback without a context" - f" object of type {object_type.__name__!r}" - " existing." - ) - - return ctx.invoke(f, obj, *args, **kwargs) - - return update_wrapper(new_func, f) - - return decorator - - -def pass_meta_key( - key: str, *, doc_description: str | None = None -) -> t.Callable[[t.Callable[te.Concatenate[T, P], R]], t.Callable[P, R]]: - """Create a decorator that passes a key from - :attr:`click.Context.meta` as the first argument to the decorated - function. - - :param key: Key in ``Context.meta`` to pass. - :param doc_description: Description of the object being passed, - inserted into the decorator's docstring. Defaults to "the 'key' - key from Context.meta". - - .. versionadded:: 8.0 - """ - - def decorator(f: t.Callable[te.Concatenate[T, P], R]) -> t.Callable[P, R]: - def new_func(*args: P.args, **kwargs: P.kwargs) -> R: - ctx = get_current_context() - obj = ctx.meta[key] - return ctx.invoke(f, obj, *args, **kwargs) - - return update_wrapper(new_func, f) - - if doc_description is None: - doc_description = f"the {key!r} key from :attr:`click.Context.meta`" - - decorator.__doc__ = ( - f"Decorator that passes {doc_description} as the first argument" - " to the decorated function." - ) - return decorator - - -CmdType = t.TypeVar("CmdType", bound=Command) - - -# variant: no call, directly as decorator for a function. -@t.overload -def command(name: _AnyCallable) -> Command: ... - - -# variant: with positional name and with positional or keyword cls argument: -# @command(namearg, CommandCls, ...) or @command(namearg, cls=CommandCls, ...) -@t.overload -def command( - name: str | None, - cls: type[CmdType], - **attrs: t.Any, -) -> t.Callable[[_AnyCallable], CmdType]: ... - - -# variant: name omitted, cls _must_ be a keyword argument, @command(cls=CommandCls, ...) -@t.overload -def command( - name: None = None, - *, - cls: type[CmdType], - **attrs: t.Any, -) -> t.Callable[[_AnyCallable], CmdType]: ... - - -# variant: with optional string name, no cls argument provided. -@t.overload -def command( - name: str | None = ..., cls: None = None, **attrs: t.Any -) -> t.Callable[[_AnyCallable], Command]: ... - - -def command( - name: str | _AnyCallable | None = None, - cls: type[CmdType] | None = None, - **attrs: t.Any, -) -> Command | t.Callable[[_AnyCallable], Command | CmdType]: - r"""Creates a new :class:`Command` and uses the decorated function as - callback. This will also automatically attach all decorated - :func:`option`\s and :func:`argument`\s as parameters to the command. - - The name of the command defaults to the name of the function, converted to - lowercase, with underscores ``_`` replaced by dashes ``-``, and the suffixes - ``_command``, ``_cmd``, ``_group``, and ``_grp`` are removed. For example, - ``init_data_command`` becomes ``init-data``. - - All keyword arguments are forwarded to the underlying command class. - For the ``params`` argument, any decorated params are appended to - the end of the list. - - Once decorated the function turns into a :class:`Command` instance - that can be invoked as a command line utility or be attached to a - command :class:`Group`. - - :param name: The name of the command. Defaults to modifying the function's - name as described above. - :param cls: The command class to create. Defaults to :class:`Command`. - - .. versionchanged:: 8.2 - The suffixes ``_command``, ``_cmd``, ``_group``, and ``_grp`` are - removed when generating the name. - - .. versionchanged:: 8.1 - This decorator can be applied without parentheses. - - .. versionchanged:: 8.1 - The ``params`` argument can be used. Decorated params are - appended to the end of the list. - """ - - func: t.Callable[[_AnyCallable], t.Any] | None = None - - if callable(name): - func = name - name = None - assert cls is None, "Use 'command(cls=cls)(callable)' to specify a class." - assert not attrs, "Use 'command(**kwargs)(callable)' to provide arguments." - - if cls is None: - cls = t.cast("type[CmdType]", Command) - - def decorator(f: _AnyCallable) -> CmdType: - if isinstance(f, Command): - raise TypeError("Attempted to convert a callback into a command twice.") - - attr_params = attrs.pop("params", None) - params = attr_params if attr_params is not None else [] - - try: - decorator_params = f.__click_params__ # type: ignore - except AttributeError: - pass - else: - del f.__click_params__ # type: ignore - params.extend(reversed(decorator_params)) - - if attrs.get("help") is None: - attrs["help"] = f.__doc__ - - if t.TYPE_CHECKING: - assert cls is not None - assert not callable(name) - - if name is not None: - cmd_name = name - else: - cmd_name = f.__name__.lower().replace("_", "-") - cmd_left, sep, suffix = cmd_name.rpartition("-") - - if sep and suffix in {"command", "cmd", "group", "grp"}: - cmd_name = cmd_left - - cmd = cls(name=cmd_name, callback=f, params=params, **attrs) - cmd.__doc__ = f.__doc__ - return cmd - - if func is not None: - return decorator(func) - - return decorator - - -GrpType = t.TypeVar("GrpType", bound=Group) - - -# variant: no call, directly as decorator for a function. -@t.overload -def group(name: _AnyCallable) -> Group: ... - - -# variant: with positional name and with positional or keyword cls argument: -# @group(namearg, GroupCls, ...) or @group(namearg, cls=GroupCls, ...) -@t.overload -def group( - name: str | None, - cls: type[GrpType], - **attrs: t.Any, -) -> t.Callable[[_AnyCallable], GrpType]: ... - - -# variant: name omitted, cls _must_ be a keyword argument, @group(cmd=GroupCls, ...) -@t.overload -def group( - name: None = None, - *, - cls: type[GrpType], - **attrs: t.Any, -) -> t.Callable[[_AnyCallable], GrpType]: ... - - -# variant: with optional string name, no cls argument provided. -@t.overload -def group( - name: str | None = ..., cls: None = None, **attrs: t.Any -) -> t.Callable[[_AnyCallable], Group]: ... - - -def group( - name: str | _AnyCallable | None = None, - cls: type[GrpType] | None = None, - **attrs: t.Any, -) -> Group | t.Callable[[_AnyCallable], Group | GrpType]: - """Creates a new :class:`Group` with a function as callback. This - works otherwise the same as :func:`command` just that the `cls` - parameter is set to :class:`Group`. - - .. versionchanged:: 8.1 - This decorator can be applied without parentheses. - """ - if cls is None: - cls = t.cast("type[GrpType]", Group) - - if callable(name): - return command(cls=cls, **attrs)(name) - - return command(name, cls, **attrs) - - -def _param_memo(f: t.Callable[..., t.Any], param: Parameter) -> None: - if isinstance(f, Command): - f.params.append(param) - else: - if not hasattr(f, "__click_params__"): - f.__click_params__ = [] # type: ignore - - f.__click_params__.append(param) # type: ignore - - -def argument( - *param_decls: str, cls: type[Argument] | None = None, **attrs: t.Any -) -> t.Callable[[FC], FC]: - """Attaches an argument to the command. All positional arguments are - passed as parameter declarations to :class:`Argument`; all keyword - arguments are forwarded unchanged (except ``cls``). - This is equivalent to creating an :class:`Argument` instance manually - and attaching it to the :attr:`Command.params` list. - - For the default argument class, refer to :class:`Argument` and - :class:`Parameter` for descriptions of parameters. - - :param cls: the argument class to instantiate. This defaults to - :class:`Argument`. - :param param_decls: Passed as positional arguments to the constructor of - ``cls``. - :param attrs: Passed as keyword arguments to the constructor of ``cls``. - """ - if cls is None: - cls = Argument - - def decorator(f: FC) -> FC: - _param_memo(f, cls(param_decls, **attrs)) - return f - - return decorator - - -def option( - *param_decls: str, cls: type[Option] | None = None, **attrs: t.Any -) -> t.Callable[[FC], FC]: - """Attaches an option to the command. All positional arguments are - passed as parameter declarations to :class:`Option`; all keyword - arguments are forwarded unchanged (except ``cls``). - This is equivalent to creating an :class:`Option` instance manually - and attaching it to the :attr:`Command.params` list. - - For the default option class, refer to :class:`Option` and - :class:`Parameter` for descriptions of parameters. - - :param cls: the option class to instantiate. This defaults to - :class:`Option`. - :param param_decls: Passed as positional arguments to the constructor of - ``cls``. - :param attrs: Passed as keyword arguments to the constructor of ``cls``. - """ - if cls is None: - cls = Option - - def decorator(f: FC) -> FC: - _param_memo(f, cls(param_decls, **attrs)) - return f - - return decorator - - -def confirmation_option(*param_decls: str, **kwargs: t.Any) -> t.Callable[[FC], FC]: - """Add a ``--yes`` option which shows a prompt before continuing if - not passed. If the prompt is declined, the program will exit. - - :param param_decls: One or more option names. Defaults to the single - value ``"--yes"``. - :param kwargs: Extra arguments are passed to :func:`option`. - """ - - def callback(ctx: Context, param: Parameter, value: bool) -> None: - if not value: - ctx.abort() - - if not param_decls: - param_decls = ("--yes",) - - kwargs.setdefault("is_flag", True) - kwargs.setdefault("callback", callback) - kwargs.setdefault("expose_value", False) - kwargs.setdefault("prompt", "Do you want to continue?") - kwargs.setdefault("help", "Confirm the action without prompting.") - return option(*param_decls, **kwargs) - - -def password_option(*param_decls: str, **kwargs: t.Any) -> t.Callable[[FC], FC]: - """Add a ``--password`` option which prompts for a password, hiding - input and asking to enter the value again for confirmation. - - :param param_decls: One or more option names. Defaults to the single - value ``"--password"``. - :param kwargs: Extra arguments are passed to :func:`option`. - """ - if not param_decls: - param_decls = ("--password",) - - kwargs.setdefault("prompt", True) - kwargs.setdefault("confirmation_prompt", True) - kwargs.setdefault("hide_input", True) - return option(*param_decls, **kwargs) - - -def version_option( - version: str | None = None, - *param_decls: str, - package_name: str | None = None, - prog_name: str | None = None, - message: str | None = None, - **kwargs: t.Any, -) -> t.Callable[[FC], FC]: - """Add a ``--version`` option which immediately prints the version - number and exits the program. - - If ``version`` is not provided, Click will try to detect it using - :func:`importlib.metadata.version` to get the version for the - ``package_name``. - - If ``package_name`` is not provided, Click will try to detect it by - inspecting the stack frames. This will be used to detect the - version, so it must match the name of the installed package. - - :param version: The version number to show. If not provided, Click - will try to detect it. - :param param_decls: One or more option names. Defaults to the single - value ``"--version"``. - :param package_name: The package name to detect the version from. If - not provided, Click will try to detect it. - :param prog_name: The name of the CLI to show in the message. If not - provided, it will be detected from the command. - :param message: The message to show. The values ``%(prog)s``, - ``%(package)s``, and ``%(version)s`` are available. Defaults to - ``"%(prog)s, version %(version)s"``. - :param kwargs: Extra arguments are passed to :func:`option`. - :raise RuntimeError: ``version`` could not be detected. - - .. versionchanged:: 8.0 - Add the ``package_name`` parameter, and the ``%(package)s`` - value for messages. - - .. versionchanged:: 8.0 - Use :mod:`importlib.metadata` instead of ``pkg_resources``. The - version is detected based on the package name, not the entry - point name. The Python package name must match the installed - package name, or be passed with ``package_name=``. - """ - if message is None: - message = _("%(prog)s, version %(version)s") - - if version is None and package_name is None: - frame = inspect.currentframe() - f_back = frame.f_back if frame is not None else None - f_globals = f_back.f_globals if f_back is not None else None - # break reference cycle - # https://docs.python.org/3/library/inspect.html#the-interpreter-stack - del frame - - if f_globals is not None: - package_name = f_globals.get("__name__") - - if package_name == "__main__": - package_name = f_globals.get("__package__") - - if package_name: - package_name = package_name.partition(".")[0] - - def callback(ctx: Context, param: Parameter, value: bool) -> None: - if not value or ctx.resilient_parsing: - return - - nonlocal prog_name - nonlocal version - - if prog_name is None: - prog_name = ctx.find_root().info_name - - if version is None and package_name is not None: - import importlib.metadata - - try: - version = importlib.metadata.version(package_name) - except importlib.metadata.PackageNotFoundError: - raise RuntimeError( - f"{package_name!r} is not installed. Try passing" - " 'package_name' instead." - ) from None - - if version is None: - raise RuntimeError( - f"Could not determine the version for {package_name!r} automatically." - ) - - echo( - message % {"prog": prog_name, "package": package_name, "version": version}, - color=ctx.color, - ) - ctx.exit() - - if not param_decls: - param_decls = ("--version",) - - kwargs.setdefault("is_flag", True) - kwargs.setdefault("expose_value", False) - kwargs.setdefault("is_eager", True) - kwargs.setdefault("help", _("Show the version and exit.")) - kwargs["callback"] = callback - return option(*param_decls, **kwargs) - - -def help_option(*param_decls: str, **kwargs: t.Any) -> t.Callable[[FC], FC]: - """Pre-configured ``--help`` option which immediately prints the help page - and exits the program. - - :param param_decls: One or more option names. Defaults to the single - value ``"--help"``. - :param kwargs: Extra arguments are passed to :func:`option`. - """ - - def show_help(ctx: Context, param: Parameter, value: bool) -> None: - """Callback that print the help page on ```` and exits.""" - if value and not ctx.resilient_parsing: - echo(ctx.get_help(), color=ctx.color) - ctx.exit() - - if not param_decls: - param_decls = ("--help",) - - kwargs.setdefault("is_flag", True) - kwargs.setdefault("expose_value", False) - kwargs.setdefault("is_eager", True) - kwargs.setdefault("help", _("Show this message and exit.")) - kwargs.setdefault("callback", show_help) - - return option(*param_decls, **kwargs) diff --git a/venv/lib/python3.10/site-packages/click/exceptions.py b/venv/lib/python3.10/site-packages/click/exceptions.py deleted file mode 100644 index f141a832e1ab678709998c3bbafcdea6561b9661..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/exceptions.py +++ /dev/null @@ -1,308 +0,0 @@ -from __future__ import annotations - -import collections.abc as cabc -import typing as t -from gettext import gettext as _ -from gettext import ngettext - -from ._compat import get_text_stderr -from .globals import resolve_color_default -from .utils import echo -from .utils import format_filename - -if t.TYPE_CHECKING: - from .core import Command - from .core import Context - from .core import Parameter - - -def _join_param_hints(param_hint: cabc.Sequence[str] | str | None) -> str | None: - if param_hint is not None and not isinstance(param_hint, str): - return " / ".join(repr(x) for x in param_hint) - - return param_hint - - -class ClickException(Exception): - """An exception that Click can handle and show to the user.""" - - #: The exit code for this exception. - exit_code = 1 - - def __init__(self, message: str) -> None: - super().__init__(message) - # The context will be removed by the time we print the message, so cache - # the color settings here to be used later on (in `show`) - self.show_color: bool | None = resolve_color_default() - self.message = message - - def format_message(self) -> str: - return self.message - - def __str__(self) -> str: - return self.message - - def show(self, file: t.IO[t.Any] | None = None) -> None: - if file is None: - file = get_text_stderr() - - echo( - _("Error: {message}").format(message=self.format_message()), - file=file, - color=self.show_color, - ) - - -class UsageError(ClickException): - """An internal exception that signals a usage error. This typically - aborts any further handling. - - :param message: the error message to display. - :param ctx: optionally the context that caused this error. Click will - fill in the context automatically in some situations. - """ - - exit_code = 2 - - def __init__(self, message: str, ctx: Context | None = None) -> None: - super().__init__(message) - self.ctx = ctx - self.cmd: Command | None = self.ctx.command if self.ctx else None - - def show(self, file: t.IO[t.Any] | None = None) -> None: - if file is None: - file = get_text_stderr() - color = None - hint = "" - if ( - self.ctx is not None - and self.ctx.command.get_help_option(self.ctx) is not None - ): - hint = _("Try '{command} {option}' for help.").format( - command=self.ctx.command_path, option=self.ctx.help_option_names[0] - ) - hint = f"{hint}\n" - if self.ctx is not None: - color = self.ctx.color - echo(f"{self.ctx.get_usage()}\n{hint}", file=file, color=color) - echo( - _("Error: {message}").format(message=self.format_message()), - file=file, - color=color, - ) - - -class BadParameter(UsageError): - """An exception that formats out a standardized error message for a - bad parameter. This is useful when thrown from a callback or type as - Click will attach contextual information to it (for instance, which - parameter it is). - - .. versionadded:: 2.0 - - :param param: the parameter object that caused this error. This can - be left out, and Click will attach this info itself - if possible. - :param param_hint: a string that shows up as parameter name. This - can be used as alternative to `param` in cases - where custom validation should happen. If it is - a string it's used as such, if it's a list then - each item is quoted and separated. - """ - - def __init__( - self, - message: str, - ctx: Context | None = None, - param: Parameter | None = None, - param_hint: str | None = None, - ) -> None: - super().__init__(message, ctx) - self.param = param - self.param_hint = param_hint - - def format_message(self) -> str: - if self.param_hint is not None: - param_hint = self.param_hint - elif self.param is not None: - param_hint = self.param.get_error_hint(self.ctx) # type: ignore - else: - return _("Invalid value: {message}").format(message=self.message) - - return _("Invalid value for {param_hint}: {message}").format( - param_hint=_join_param_hints(param_hint), message=self.message - ) - - -class MissingParameter(BadParameter): - """Raised if click required an option or argument but it was not - provided when invoking the script. - - .. versionadded:: 4.0 - - :param param_type: a string that indicates the type of the parameter. - The default is to inherit the parameter type from - the given `param`. Valid values are ``'parameter'``, - ``'option'`` or ``'argument'``. - """ - - def __init__( - self, - message: str | None = None, - ctx: Context | None = None, - param: Parameter | None = None, - param_hint: str | None = None, - param_type: str | None = None, - ) -> None: - super().__init__(message or "", ctx, param, param_hint) - self.param_type = param_type - - def format_message(self) -> str: - if self.param_hint is not None: - param_hint: str | None = self.param_hint - elif self.param is not None: - param_hint = self.param.get_error_hint(self.ctx) # type: ignore - else: - param_hint = None - - param_hint = _join_param_hints(param_hint) - param_hint = f" {param_hint}" if param_hint else "" - - param_type = self.param_type - if param_type is None and self.param is not None: - param_type = self.param.param_type_name - - msg = self.message - if self.param is not None: - msg_extra = self.param.type.get_missing_message( - param=self.param, ctx=self.ctx - ) - if msg_extra: - if msg: - msg += f". {msg_extra}" - else: - msg = msg_extra - - msg = f" {msg}" if msg else "" - - # Translate param_type for known types. - if param_type == "argument": - missing = _("Missing argument") - elif param_type == "option": - missing = _("Missing option") - elif param_type == "parameter": - missing = _("Missing parameter") - else: - missing = _("Missing {param_type}").format(param_type=param_type) - - return f"{missing}{param_hint}.{msg}" - - def __str__(self) -> str: - if not self.message: - param_name = self.param.name if self.param else None - return _("Missing parameter: {param_name}").format(param_name=param_name) - else: - return self.message - - -class NoSuchOption(UsageError): - """Raised if click attempted to handle an option that does not - exist. - - .. versionadded:: 4.0 - """ - - def __init__( - self, - option_name: str, - message: str | None = None, - possibilities: cabc.Sequence[str] | None = None, - ctx: Context | None = None, - ) -> None: - if message is None: - message = _("No such option: {name}").format(name=option_name) - - super().__init__(message, ctx) - self.option_name = option_name - self.possibilities = possibilities - - def format_message(self) -> str: - if not self.possibilities: - return self.message - - possibility_str = ", ".join(sorted(self.possibilities)) - suggest = ngettext( - "Did you mean {possibility}?", - "(Possible options: {possibilities})", - len(self.possibilities), - ).format(possibility=possibility_str, possibilities=possibility_str) - return f"{self.message} {suggest}" - - -class BadOptionUsage(UsageError): - """Raised if an option is generally supplied but the use of the option - was incorrect. This is for instance raised if the number of arguments - for an option is not correct. - - .. versionadded:: 4.0 - - :param option_name: the name of the option being used incorrectly. - """ - - def __init__( - self, option_name: str, message: str, ctx: Context | None = None - ) -> None: - super().__init__(message, ctx) - self.option_name = option_name - - -class BadArgumentUsage(UsageError): - """Raised if an argument is generally supplied but the use of the argument - was incorrect. This is for instance raised if the number of values - for an argument is not correct. - - .. versionadded:: 6.0 - """ - - -class NoArgsIsHelpError(UsageError): - def __init__(self, ctx: Context) -> None: - self.ctx: Context - super().__init__(ctx.get_help(), ctx=ctx) - - def show(self, file: t.IO[t.Any] | None = None) -> None: - echo(self.format_message(), file=file, err=True, color=self.ctx.color) - - -class FileError(ClickException): - """Raised if a file cannot be opened.""" - - def __init__(self, filename: str, hint: str | None = None) -> None: - if hint is None: - hint = _("unknown error") - - super().__init__(hint) - self.ui_filename: str = format_filename(filename) - self.filename = filename - - def format_message(self) -> str: - return _("Could not open file {filename!r}: {message}").format( - filename=self.ui_filename, message=self.message - ) - - -class Abort(RuntimeError): - """An internal signalling exception that signals Click to abort.""" - - -class Exit(RuntimeError): - """An exception that indicates that the application should exit with some - status code. - - :param code: the status code to exit with. - """ - - __slots__ = ("exit_code",) - - def __init__(self, code: int = 0) -> None: - self.exit_code: int = code diff --git a/venv/lib/python3.10/site-packages/click/formatting.py b/venv/lib/python3.10/site-packages/click/formatting.py deleted file mode 100644 index 9891f8809feff24a529ce47a860589c129835322..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/formatting.py +++ /dev/null @@ -1,301 +0,0 @@ -from __future__ import annotations - -import collections.abc as cabc -from contextlib import contextmanager -from gettext import gettext as _ - -from ._compat import term_len -from .parser import _split_opt - -# Can force a width. This is used by the test system -FORCED_WIDTH: int | None = None - - -def measure_table(rows: cabc.Iterable[tuple[str, str]]) -> tuple[int, ...]: - widths: dict[int, int] = {} - - for row in rows: - for idx, col in enumerate(row): - widths[idx] = max(widths.get(idx, 0), term_len(col)) - - return tuple(y for x, y in sorted(widths.items())) - - -def iter_rows( - rows: cabc.Iterable[tuple[str, str]], col_count: int -) -> cabc.Iterator[tuple[str, ...]]: - for row in rows: - yield row + ("",) * (col_count - len(row)) - - -def wrap_text( - text: str, - width: int = 78, - initial_indent: str = "", - subsequent_indent: str = "", - preserve_paragraphs: bool = False, -) -> str: - """A helper function that intelligently wraps text. By default, it - assumes that it operates on a single paragraph of text but if the - `preserve_paragraphs` parameter is provided it will intelligently - handle paragraphs (defined by two empty lines). - - If paragraphs are handled, a paragraph can be prefixed with an empty - line containing the ``\\b`` character (``\\x08``) to indicate that - no rewrapping should happen in that block. - - :param text: the text that should be rewrapped. - :param width: the maximum width for the text. - :param initial_indent: the initial indent that should be placed on the - first line as a string. - :param subsequent_indent: the indent string that should be placed on - each consecutive line. - :param preserve_paragraphs: if this flag is set then the wrapping will - intelligently handle paragraphs. - """ - from ._textwrap import TextWrapper - - text = text.expandtabs() - wrapper = TextWrapper( - width, - initial_indent=initial_indent, - subsequent_indent=subsequent_indent, - replace_whitespace=False, - ) - if not preserve_paragraphs: - return wrapper.fill(text) - - p: list[tuple[int, bool, str]] = [] - buf: list[str] = [] - indent = None - - def _flush_par() -> None: - if not buf: - return - if buf[0].strip() == "\b": - p.append((indent or 0, True, "\n".join(buf[1:]))) - else: - p.append((indent or 0, False, " ".join(buf))) - del buf[:] - - for line in text.splitlines(): - if not line: - _flush_par() - indent = None - else: - if indent is None: - orig_len = term_len(line) - line = line.lstrip() - indent = orig_len - term_len(line) - buf.append(line) - _flush_par() - - rv = [] - for indent, raw, text in p: - with wrapper.extra_indent(" " * indent): - if raw: - rv.append(wrapper.indent_only(text)) - else: - rv.append(wrapper.fill(text)) - - return "\n\n".join(rv) - - -class HelpFormatter: - """This class helps with formatting text-based help pages. It's - usually just needed for very special internal cases, but it's also - exposed so that developers can write their own fancy outputs. - - At present, it always writes into memory. - - :param indent_increment: the additional increment for each level. - :param width: the width for the text. This defaults to the terminal - width clamped to a maximum of 78. - """ - - def __init__( - self, - indent_increment: int = 2, - width: int | None = None, - max_width: int | None = None, - ) -> None: - import shutil - - self.indent_increment = indent_increment - if max_width is None: - max_width = 80 - if width is None: - width = FORCED_WIDTH - if width is None: - width = max(min(shutil.get_terminal_size().columns, max_width) - 2, 50) - self.width = width - self.current_indent: int = 0 - self.buffer: list[str] = [] - - def write(self, string: str) -> None: - """Writes a unicode string into the internal buffer.""" - self.buffer.append(string) - - def indent(self) -> None: - """Increases the indentation.""" - self.current_indent += self.indent_increment - - def dedent(self) -> None: - """Decreases the indentation.""" - self.current_indent -= self.indent_increment - - def write_usage(self, prog: str, args: str = "", prefix: str | None = None) -> None: - """Writes a usage line into the buffer. - - :param prog: the program name. - :param args: whitespace separated list of arguments. - :param prefix: The prefix for the first line. Defaults to - ``"Usage: "``. - """ - if prefix is None: - prefix = f"{_('Usage:')} " - - usage_prefix = f"{prefix:>{self.current_indent}}{prog} " - text_width = self.width - self.current_indent - - if text_width >= (term_len(usage_prefix) + 20): - # The arguments will fit to the right of the prefix. - indent = " " * term_len(usage_prefix) - self.write( - wrap_text( - args, - text_width, - initial_indent=usage_prefix, - subsequent_indent=indent, - ) - ) - else: - # The prefix is too long, put the arguments on the next line. - self.write(usage_prefix) - self.write("\n") - indent = " " * (max(self.current_indent, term_len(prefix)) + 4) - self.write( - wrap_text( - args, text_width, initial_indent=indent, subsequent_indent=indent - ) - ) - - self.write("\n") - - def write_heading(self, heading: str) -> None: - """Writes a heading into the buffer.""" - self.write(f"{'':>{self.current_indent}}{heading}:\n") - - def write_paragraph(self) -> None: - """Writes a paragraph into the buffer.""" - if self.buffer: - self.write("\n") - - def write_text(self, text: str) -> None: - """Writes re-indented text into the buffer. This rewraps and - preserves paragraphs. - """ - indent = " " * self.current_indent - self.write( - wrap_text( - text, - self.width, - initial_indent=indent, - subsequent_indent=indent, - preserve_paragraphs=True, - ) - ) - self.write("\n") - - def write_dl( - self, - rows: cabc.Sequence[tuple[str, str]], - col_max: int = 30, - col_spacing: int = 2, - ) -> None: - """Writes a definition list into the buffer. This is how options - and commands are usually formatted. - - :param rows: a list of two item tuples for the terms and values. - :param col_max: the maximum width of the first column. - :param col_spacing: the number of spaces between the first and - second column. - """ - rows = list(rows) - widths = measure_table(rows) - if len(widths) != 2: - raise TypeError("Expected two columns for definition list") - - first_col = min(widths[0], col_max) + col_spacing - - for first, second in iter_rows(rows, len(widths)): - self.write(f"{'':>{self.current_indent}}{first}") - if not second: - self.write("\n") - continue - if term_len(first) <= first_col - col_spacing: - self.write(" " * (first_col - term_len(first))) - else: - self.write("\n") - self.write(" " * (first_col + self.current_indent)) - - text_width = max(self.width - first_col - 2, 10) - wrapped_text = wrap_text(second, text_width, preserve_paragraphs=True) - lines = wrapped_text.splitlines() - - if lines: - self.write(f"{lines[0]}\n") - - for line in lines[1:]: - self.write(f"{'':>{first_col + self.current_indent}}{line}\n") - else: - self.write("\n") - - @contextmanager - def section(self, name: str) -> cabc.Iterator[None]: - """Helpful context manager that writes a paragraph, a heading, - and the indents. - - :param name: the section name that is written as heading. - """ - self.write_paragraph() - self.write_heading(name) - self.indent() - try: - yield - finally: - self.dedent() - - @contextmanager - def indentation(self) -> cabc.Iterator[None]: - """A context manager that increases the indentation.""" - self.indent() - try: - yield - finally: - self.dedent() - - def getvalue(self) -> str: - """Returns the buffer contents.""" - return "".join(self.buffer) - - -def join_options(options: cabc.Sequence[str]) -> tuple[str, bool]: - """Given a list of option strings this joins them in the most appropriate - way and returns them in the form ``(formatted_string, - any_prefix_is_slash)`` where the second item in the tuple is a flag that - indicates if any of the option prefixes was a slash. - """ - rv = [] - any_prefix_is_slash = False - - for opt in options: - prefix = _split_opt(opt)[0] - - if prefix == "/": - any_prefix_is_slash = True - - rv.append((len(prefix), opt)) - - rv.sort(key=lambda x: x[0]) - return ", ".join(x[1] for x in rv), any_prefix_is_slash diff --git a/venv/lib/python3.10/site-packages/click/globals.py b/venv/lib/python3.10/site-packages/click/globals.py deleted file mode 100644 index a2f91723d21cefdf658be327174e1f8ccdc20fcd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/globals.py +++ /dev/null @@ -1,67 +0,0 @@ -from __future__ import annotations - -import typing as t -from threading import local - -if t.TYPE_CHECKING: - from .core import Context - -_local = local() - - -@t.overload -def get_current_context(silent: t.Literal[False] = False) -> Context: ... - - -@t.overload -def get_current_context(silent: bool = ...) -> Context | None: ... - - -def get_current_context(silent: bool = False) -> Context | None: - """Returns the current click context. This can be used as a way to - access the current context object from anywhere. This is a more implicit - alternative to the :func:`pass_context` decorator. This function is - primarily useful for helpers such as :func:`echo` which might be - interested in changing its behavior based on the current context. - - To push the current context, :meth:`Context.scope` can be used. - - .. versionadded:: 5.0 - - :param silent: if set to `True` the return value is `None` if no context - is available. The default behavior is to raise a - :exc:`RuntimeError`. - """ - try: - return t.cast("Context", _local.stack[-1]) - except (AttributeError, IndexError) as e: - if not silent: - raise RuntimeError("There is no active click context.") from e - - return None - - -def push_context(ctx: Context) -> None: - """Pushes a new context to the current stack.""" - _local.__dict__.setdefault("stack", []).append(ctx) - - -def pop_context() -> None: - """Removes the top level from the stack.""" - _local.stack.pop() - - -def resolve_color_default(color: bool | None = None) -> bool | None: - """Internal helper to get the default value of the color flag. If a - value is passed it's returned unchanged, otherwise it's looked up from - the current context. - """ - if color is not None: - return color - - ctx = get_current_context(silent=True) - - if ctx is not None: - return ctx.color - - return None diff --git a/venv/lib/python3.10/site-packages/click/parser.py b/venv/lib/python3.10/site-packages/click/parser.py deleted file mode 100644 index a8b7d2634a89df3c9e0b7a1f85dfd4bace60a608..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/parser.py +++ /dev/null @@ -1,532 +0,0 @@ -""" -This module started out as largely a copy paste from the stdlib's -optparse module with the features removed that we do not need from -optparse because we implement them in Click on a higher level (for -instance type handling, help formatting and a lot more). - -The plan is to remove more and more from here over time. - -The reason this is a different module and not optparse from the stdlib -is that there are differences in 2.x and 3.x about the error messages -generated and optparse in the stdlib uses gettext for no good reason -and might cause us issues. - -Click uses parts of optparse written by Gregory P. Ward and maintained -by the Python Software Foundation. This is limited to code in parser.py. - -Copyright 2001-2006 Gregory P. Ward. All rights reserved. -Copyright 2002-2006 Python Software Foundation. All rights reserved. -""" - -# This code uses parts of optparse written by Gregory P. Ward and -# maintained by the Python Software Foundation. -# Copyright 2001-2006 Gregory P. Ward -# Copyright 2002-2006 Python Software Foundation -from __future__ import annotations - -import collections.abc as cabc -import typing as t -from collections import deque -from gettext import gettext as _ -from gettext import ngettext - -from .exceptions import BadArgumentUsage -from .exceptions import BadOptionUsage -from .exceptions import NoSuchOption -from .exceptions import UsageError - -if t.TYPE_CHECKING: - from .core import Argument as CoreArgument - from .core import Context - from .core import Option as CoreOption - from .core import Parameter as CoreParameter - -V = t.TypeVar("V") - -# Sentinel value that indicates an option was passed as a flag without a -# value but is not a flag option. Option.consume_value uses this to -# prompt or use the flag_value. -_flag_needs_value = object() - - -def _unpack_args( - args: cabc.Sequence[str], nargs_spec: cabc.Sequence[int] -) -> tuple[cabc.Sequence[str | cabc.Sequence[str | None] | None], list[str]]: - """Given an iterable of arguments and an iterable of nargs specifications, - it returns a tuple with all the unpacked arguments at the first index - and all remaining arguments as the second. - - The nargs specification is the number of arguments that should be consumed - or `-1` to indicate that this position should eat up all the remainders. - - Missing items are filled with `None`. - """ - args = deque(args) - nargs_spec = deque(nargs_spec) - rv: list[str | tuple[str | None, ...] | None] = [] - spos: int | None = None - - def _fetch(c: deque[V]) -> V | None: - try: - if spos is None: - return c.popleft() - else: - return c.pop() - except IndexError: - return None - - while nargs_spec: - nargs = _fetch(nargs_spec) - - if nargs is None: - continue - - if nargs == 1: - rv.append(_fetch(args)) - elif nargs > 1: - x = [_fetch(args) for _ in range(nargs)] - - # If we're reversed, we're pulling in the arguments in reverse, - # so we need to turn them around. - if spos is not None: - x.reverse() - - rv.append(tuple(x)) - elif nargs < 0: - if spos is not None: - raise TypeError("Cannot have two nargs < 0") - - spos = len(rv) - rv.append(None) - - # spos is the position of the wildcard (star). If it's not `None`, - # we fill it with the remainder. - if spos is not None: - rv[spos] = tuple(args) - args = [] - rv[spos + 1 :] = reversed(rv[spos + 1 :]) - - return tuple(rv), list(args) - - -def _split_opt(opt: str) -> tuple[str, str]: - first = opt[:1] - if first.isalnum(): - return "", opt - if opt[1:2] == first: - return opt[:2], opt[2:] - return first, opt[1:] - - -def _normalize_opt(opt: str, ctx: Context | None) -> str: - if ctx is None or ctx.token_normalize_func is None: - return opt - prefix, opt = _split_opt(opt) - return f"{prefix}{ctx.token_normalize_func(opt)}" - - -class _Option: - def __init__( - self, - obj: CoreOption, - opts: cabc.Sequence[str], - dest: str | None, - action: str | None = None, - nargs: int = 1, - const: t.Any | None = None, - ): - self._short_opts = [] - self._long_opts = [] - self.prefixes: set[str] = set() - - for opt in opts: - prefix, value = _split_opt(opt) - if not prefix: - raise ValueError(f"Invalid start character for option ({opt})") - self.prefixes.add(prefix[0]) - if len(prefix) == 1 and len(value) == 1: - self._short_opts.append(opt) - else: - self._long_opts.append(opt) - self.prefixes.add(prefix) - - if action is None: - action = "store" - - self.dest = dest - self.action = action - self.nargs = nargs - self.const = const - self.obj = obj - - @property - def takes_value(self) -> bool: - return self.action in ("store", "append") - - def process(self, value: t.Any, state: _ParsingState) -> None: - if self.action == "store": - state.opts[self.dest] = value # type: ignore - elif self.action == "store_const": - state.opts[self.dest] = self.const # type: ignore - elif self.action == "append": - state.opts.setdefault(self.dest, []).append(value) # type: ignore - elif self.action == "append_const": - state.opts.setdefault(self.dest, []).append(self.const) # type: ignore - elif self.action == "count": - state.opts[self.dest] = state.opts.get(self.dest, 0) + 1 # type: ignore - else: - raise ValueError(f"unknown action '{self.action}'") - state.order.append(self.obj) - - -class _Argument: - def __init__(self, obj: CoreArgument, dest: str | None, nargs: int = 1): - self.dest = dest - self.nargs = nargs - self.obj = obj - - def process( - self, - value: str | cabc.Sequence[str | None] | None, - state: _ParsingState, - ) -> None: - if self.nargs > 1: - assert value is not None - holes = sum(1 for x in value if x is None) - if holes == len(value): - value = None - elif holes != 0: - raise BadArgumentUsage( - _("Argument {name!r} takes {nargs} values.").format( - name=self.dest, nargs=self.nargs - ) - ) - - if self.nargs == -1 and self.obj.envvar is not None and value == (): - # Replace empty tuple with None so that a value from the - # environment may be tried. - value = None - - state.opts[self.dest] = value # type: ignore - state.order.append(self.obj) - - -class _ParsingState: - def __init__(self, rargs: list[str]) -> None: - self.opts: dict[str, t.Any] = {} - self.largs: list[str] = [] - self.rargs = rargs - self.order: list[CoreParameter] = [] - - -class _OptionParser: - """The option parser is an internal class that is ultimately used to - parse options and arguments. It's modelled after optparse and brings - a similar but vastly simplified API. It should generally not be used - directly as the high level Click classes wrap it for you. - - It's not nearly as extensible as optparse or argparse as it does not - implement features that are implemented on a higher level (such as - types or defaults). - - :param ctx: optionally the :class:`~click.Context` where this parser - should go with. - - .. deprecated:: 8.2 - Will be removed in Click 9.0. - """ - - def __init__(self, ctx: Context | None = None) -> None: - #: The :class:`~click.Context` for this parser. This might be - #: `None` for some advanced use cases. - self.ctx = ctx - #: This controls how the parser deals with interspersed arguments. - #: If this is set to `False`, the parser will stop on the first - #: non-option. Click uses this to implement nested subcommands - #: safely. - self.allow_interspersed_args: bool = True - #: This tells the parser how to deal with unknown options. By - #: default it will error out (which is sensible), but there is a - #: second mode where it will ignore it and continue processing - #: after shifting all the unknown options into the resulting args. - self.ignore_unknown_options: bool = False - - if ctx is not None: - self.allow_interspersed_args = ctx.allow_interspersed_args - self.ignore_unknown_options = ctx.ignore_unknown_options - - self._short_opt: dict[str, _Option] = {} - self._long_opt: dict[str, _Option] = {} - self._opt_prefixes = {"-", "--"} - self._args: list[_Argument] = [] - - def add_option( - self, - obj: CoreOption, - opts: cabc.Sequence[str], - dest: str | None, - action: str | None = None, - nargs: int = 1, - const: t.Any | None = None, - ) -> None: - """Adds a new option named `dest` to the parser. The destination - is not inferred (unlike with optparse) and needs to be explicitly - provided. Action can be any of ``store``, ``store_const``, - ``append``, ``append_const`` or ``count``. - - The `obj` can be used to identify the option in the order list - that is returned from the parser. - """ - opts = [_normalize_opt(opt, self.ctx) for opt in opts] - option = _Option(obj, opts, dest, action=action, nargs=nargs, const=const) - self._opt_prefixes.update(option.prefixes) - for opt in option._short_opts: - self._short_opt[opt] = option - for opt in option._long_opts: - self._long_opt[opt] = option - - def add_argument(self, obj: CoreArgument, dest: str | None, nargs: int = 1) -> None: - """Adds a positional argument named `dest` to the parser. - - The `obj` can be used to identify the option in the order list - that is returned from the parser. - """ - self._args.append(_Argument(obj, dest=dest, nargs=nargs)) - - def parse_args( - self, args: list[str] - ) -> tuple[dict[str, t.Any], list[str], list[CoreParameter]]: - """Parses positional arguments and returns ``(values, args, order)`` - for the parsed options and arguments as well as the leftover - arguments if there are any. The order is a list of objects as they - appear on the command line. If arguments appear multiple times they - will be memorized multiple times as well. - """ - state = _ParsingState(args) - try: - self._process_args_for_options(state) - self._process_args_for_args(state) - except UsageError: - if self.ctx is None or not self.ctx.resilient_parsing: - raise - return state.opts, state.largs, state.order - - def _process_args_for_args(self, state: _ParsingState) -> None: - pargs, args = _unpack_args( - state.largs + state.rargs, [x.nargs for x in self._args] - ) - - for idx, arg in enumerate(self._args): - arg.process(pargs[idx], state) - - state.largs = args - state.rargs = [] - - def _process_args_for_options(self, state: _ParsingState) -> None: - while state.rargs: - arg = state.rargs.pop(0) - arglen = len(arg) - # Double dashes always handled explicitly regardless of what - # prefixes are valid. - if arg == "--": - return - elif arg[:1] in self._opt_prefixes and arglen > 1: - self._process_opts(arg, state) - elif self.allow_interspersed_args: - state.largs.append(arg) - else: - state.rargs.insert(0, arg) - return - - # Say this is the original argument list: - # [arg0, arg1, ..., arg(i-1), arg(i), arg(i+1), ..., arg(N-1)] - # ^ - # (we are about to process arg(i)). - # - # Then rargs is [arg(i), ..., arg(N-1)] and largs is a *subset* of - # [arg0, ..., arg(i-1)] (any options and their arguments will have - # been removed from largs). - # - # The while loop will usually consume 1 or more arguments per pass. - # If it consumes 1 (eg. arg is an option that takes no arguments), - # then after _process_arg() is done the situation is: - # - # largs = subset of [arg0, ..., arg(i)] - # rargs = [arg(i+1), ..., arg(N-1)] - # - # If allow_interspersed_args is false, largs will always be - # *empty* -- still a subset of [arg0, ..., arg(i-1)], but - # not a very interesting subset! - - def _match_long_opt( - self, opt: str, explicit_value: str | None, state: _ParsingState - ) -> None: - if opt not in self._long_opt: - from difflib import get_close_matches - - possibilities = get_close_matches(opt, self._long_opt) - raise NoSuchOption(opt, possibilities=possibilities, ctx=self.ctx) - - option = self._long_opt[opt] - if option.takes_value: - # At this point it's safe to modify rargs by injecting the - # explicit value, because no exception is raised in this - # branch. This means that the inserted value will be fully - # consumed. - if explicit_value is not None: - state.rargs.insert(0, explicit_value) - - value = self._get_value_from_state(opt, option, state) - - elif explicit_value is not None: - raise BadOptionUsage( - opt, _("Option {name!r} does not take a value.").format(name=opt) - ) - - else: - value = None - - option.process(value, state) - - def _match_short_opt(self, arg: str, state: _ParsingState) -> None: - stop = False - i = 1 - prefix = arg[0] - unknown_options = [] - - for ch in arg[1:]: - opt = _normalize_opt(f"{prefix}{ch}", self.ctx) - option = self._short_opt.get(opt) - i += 1 - - if not option: - if self.ignore_unknown_options: - unknown_options.append(ch) - continue - raise NoSuchOption(opt, ctx=self.ctx) - if option.takes_value: - # Any characters left in arg? Pretend they're the - # next arg, and stop consuming characters of arg. - if i < len(arg): - state.rargs.insert(0, arg[i:]) - stop = True - - value = self._get_value_from_state(opt, option, state) - - else: - value = None - - option.process(value, state) - - if stop: - break - - # If we got any unknown options we recombine the string of the - # remaining options and re-attach the prefix, then report that - # to the state as new larg. This way there is basic combinatorics - # that can be achieved while still ignoring unknown arguments. - if self.ignore_unknown_options and unknown_options: - state.largs.append(f"{prefix}{''.join(unknown_options)}") - - def _get_value_from_state( - self, option_name: str, option: _Option, state: _ParsingState - ) -> t.Any: - nargs = option.nargs - - if len(state.rargs) < nargs: - if option.obj._flag_needs_value: - # Option allows omitting the value. - value = _flag_needs_value - else: - raise BadOptionUsage( - option_name, - ngettext( - "Option {name!r} requires an argument.", - "Option {name!r} requires {nargs} arguments.", - nargs, - ).format(name=option_name, nargs=nargs), - ) - elif nargs == 1: - next_rarg = state.rargs[0] - - if ( - option.obj._flag_needs_value - and isinstance(next_rarg, str) - and next_rarg[:1] in self._opt_prefixes - and len(next_rarg) > 1 - ): - # The next arg looks like the start of an option, don't - # use it as the value if omitting the value is allowed. - value = _flag_needs_value - else: - value = state.rargs.pop(0) - else: - value = tuple(state.rargs[:nargs]) - del state.rargs[:nargs] - - return value - - def _process_opts(self, arg: str, state: _ParsingState) -> None: - explicit_value = None - # Long option handling happens in two parts. The first part is - # supporting explicitly attached values. In any case, we will try - # to long match the option first. - if "=" in arg: - long_opt, explicit_value = arg.split("=", 1) - else: - long_opt = arg - norm_long_opt = _normalize_opt(long_opt, self.ctx) - - # At this point we will match the (assumed) long option through - # the long option matching code. Note that this allows options - # like "-foo" to be matched as long options. - try: - self._match_long_opt(norm_long_opt, explicit_value, state) - except NoSuchOption: - # At this point the long option matching failed, and we need - # to try with short options. However there is a special rule - # which says, that if we have a two character options prefix - # (applies to "--foo" for instance), we do not dispatch to the - # short option code and will instead raise the no option - # error. - if arg[:2] not in self._opt_prefixes: - self._match_short_opt(arg, state) - return - - if not self.ignore_unknown_options: - raise - - state.largs.append(arg) - - -def __getattr__(name: str) -> object: - import warnings - - if name in { - "OptionParser", - "Argument", - "Option", - "split_opt", - "normalize_opt", - "ParsingState", - }: - warnings.warn( - f"'parser.{name}' is deprecated and will be removed in Click 9.0." - " The old parser is available in 'optparse'.", - DeprecationWarning, - stacklevel=2, - ) - return globals()[f"_{name}"] - - if name == "split_arg_string": - from .shell_completion import split_arg_string - - warnings.warn( - "Importing 'parser.split_arg_string' is deprecated, it will only be" - " available in 'shell_completion' in Click 9.0.", - DeprecationWarning, - stacklevel=2, - ) - return split_arg_string - - raise AttributeError(name) diff --git a/venv/lib/python3.10/site-packages/click/py.typed b/venv/lib/python3.10/site-packages/click/py.typed deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/click/shell_completion.py b/venv/lib/python3.10/site-packages/click/shell_completion.py deleted file mode 100644 index 6c39d5eb3889000b499d37262cc5f7de8b7a443e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/shell_completion.py +++ /dev/null @@ -1,644 +0,0 @@ -from __future__ import annotations - -import collections.abc as cabc -import os -import re -import typing as t -from gettext import gettext as _ - -from .core import Argument -from .core import Command -from .core import Context -from .core import Group -from .core import Option -from .core import Parameter -from .core import ParameterSource -from .utils import echo - - -def shell_complete( - cli: Command, - ctx_args: cabc.MutableMapping[str, t.Any], - prog_name: str, - complete_var: str, - instruction: str, -) -> int: - """Perform shell completion for the given CLI program. - - :param cli: Command being called. - :param ctx_args: Extra arguments to pass to - ``cli.make_context``. - :param prog_name: Name of the executable in the shell. - :param complete_var: Name of the environment variable that holds - the completion instruction. - :param instruction: Value of ``complete_var`` with the completion - instruction and shell, in the form ``instruction_shell``. - :return: Status code to exit with. - """ - shell, _, instruction = instruction.partition("_") - comp_cls = get_completion_class(shell) - - if comp_cls is None: - return 1 - - comp = comp_cls(cli, ctx_args, prog_name, complete_var) - - if instruction == "source": - echo(comp.source()) - return 0 - - if instruction == "complete": - echo(comp.complete()) - return 0 - - return 1 - - -class CompletionItem: - """Represents a completion value and metadata about the value. The - default metadata is ``type`` to indicate special shell handling, - and ``help`` if a shell supports showing a help string next to the - value. - - Arbitrary parameters can be passed when creating the object, and - accessed using ``item.attr``. If an attribute wasn't passed, - accessing it returns ``None``. - - :param value: The completion suggestion. - :param type: Tells the shell script to provide special completion - support for the type. Click uses ``"dir"`` and ``"file"``. - :param help: String shown next to the value if supported. - :param kwargs: Arbitrary metadata. The built-in implementations - don't use this, but custom type completions paired with custom - shell support could use it. - """ - - __slots__ = ("value", "type", "help", "_info") - - def __init__( - self, - value: t.Any, - type: str = "plain", - help: str | None = None, - **kwargs: t.Any, - ) -> None: - self.value: t.Any = value - self.type: str = type - self.help: str | None = help - self._info = kwargs - - def __getattr__(self, name: str) -> t.Any: - return self._info.get(name) - - -# Only Bash >= 4.4 has the nosort option. -_SOURCE_BASH = """\ -%(complete_func)s() { - local IFS=$'\\n' - local response - - response=$(env COMP_WORDS="${COMP_WORDS[*]}" COMP_CWORD=$COMP_CWORD \ -%(complete_var)s=bash_complete $1) - - for completion in $response; do - IFS=',' read type value <<< "$completion" - - if [[ $type == 'dir' ]]; then - COMPREPLY=() - compopt -o dirnames - elif [[ $type == 'file' ]]; then - COMPREPLY=() - compopt -o default - elif [[ $type == 'plain' ]]; then - COMPREPLY+=($value) - fi - done - - return 0 -} - -%(complete_func)s_setup() { - complete -o nosort -F %(complete_func)s %(prog_name)s -} - -%(complete_func)s_setup; -""" - -_SOURCE_ZSH = """\ -#compdef %(prog_name)s - -%(complete_func)s() { - local -a completions - local -a completions_with_descriptions - local -a response - (( ! $+commands[%(prog_name)s] )) && return 1 - - response=("${(@f)$(env COMP_WORDS="${words[*]}" COMP_CWORD=$((CURRENT-1)) \ -%(complete_var)s=zsh_complete %(prog_name)s)}") - - for type key descr in ${response}; do - if [[ "$type" == "plain" ]]; then - if [[ "$descr" == "_" ]]; then - completions+=("$key") - else - completions_with_descriptions+=("$key":"$descr") - fi - elif [[ "$type" == "dir" ]]; then - _path_files -/ - elif [[ "$type" == "file" ]]; then - _path_files -f - fi - done - - if [ -n "$completions_with_descriptions" ]; then - _describe -V unsorted completions_with_descriptions -U - fi - - if [ -n "$completions" ]; then - compadd -U -V unsorted -a completions - fi -} - -if [[ $zsh_eval_context[-1] == loadautofunc ]]; then - # autoload from fpath, call function directly - %(complete_func)s "$@" -else - # eval/source/. command, register function for later - compdef %(complete_func)s %(prog_name)s -fi -""" - -_SOURCE_FISH = """\ -function %(complete_func)s; - set -l response (env %(complete_var)s=fish_complete COMP_WORDS=(commandline -cp) \ -COMP_CWORD=(commandline -t) %(prog_name)s); - - for completion in $response; - set -l metadata (string split "," $completion); - - if test $metadata[1] = "dir"; - __fish_complete_directories $metadata[2]; - else if test $metadata[1] = "file"; - __fish_complete_path $metadata[2]; - else if test $metadata[1] = "plain"; - echo $metadata[2]; - end; - end; -end; - -complete --no-files --command %(prog_name)s --arguments \ -"(%(complete_func)s)"; -""" - - -class ShellComplete: - """Base class for providing shell completion support. A subclass for - a given shell will override attributes and methods to implement the - completion instructions (``source`` and ``complete``). - - :param cli: Command being called. - :param prog_name: Name of the executable in the shell. - :param complete_var: Name of the environment variable that holds - the completion instruction. - - .. versionadded:: 8.0 - """ - - name: t.ClassVar[str] - """Name to register the shell as with :func:`add_completion_class`. - This is used in completion instructions (``{name}_source`` and - ``{name}_complete``). - """ - - source_template: t.ClassVar[str] - """Completion script template formatted by :meth:`source`. This must - be provided by subclasses. - """ - - def __init__( - self, - cli: Command, - ctx_args: cabc.MutableMapping[str, t.Any], - prog_name: str, - complete_var: str, - ) -> None: - self.cli = cli - self.ctx_args = ctx_args - self.prog_name = prog_name - self.complete_var = complete_var - - @property - def func_name(self) -> str: - """The name of the shell function defined by the completion - script. - """ - safe_name = re.sub(r"\W*", "", self.prog_name.replace("-", "_"), flags=re.ASCII) - return f"_{safe_name}_completion" - - def source_vars(self) -> dict[str, t.Any]: - """Vars for formatting :attr:`source_template`. - - By default this provides ``complete_func``, ``complete_var``, - and ``prog_name``. - """ - return { - "complete_func": self.func_name, - "complete_var": self.complete_var, - "prog_name": self.prog_name, - } - - def source(self) -> str: - """Produce the shell script that defines the completion - function. By default this ``%``-style formats - :attr:`source_template` with the dict returned by - :meth:`source_vars`. - """ - return self.source_template % self.source_vars() - - def get_completion_args(self) -> tuple[list[str], str]: - """Use the env vars defined by the shell script to return a - tuple of ``args, incomplete``. This must be implemented by - subclasses. - """ - raise NotImplementedError - - def get_completions(self, args: list[str], incomplete: str) -> list[CompletionItem]: - """Determine the context and last complete command or parameter - from the complete args. Call that object's ``shell_complete`` - method to get the completions for the incomplete value. - - :param args: List of complete args before the incomplete value. - :param incomplete: Value being completed. May be empty. - """ - ctx = _resolve_context(self.cli, self.ctx_args, self.prog_name, args) - obj, incomplete = _resolve_incomplete(ctx, args, incomplete) - return obj.shell_complete(ctx, incomplete) - - def format_completion(self, item: CompletionItem) -> str: - """Format a completion item into the form recognized by the - shell script. This must be implemented by subclasses. - - :param item: Completion item to format. - """ - raise NotImplementedError - - def complete(self) -> str: - """Produce the completion data to send back to the shell. - - By default this calls :meth:`get_completion_args`, gets the - completions, then calls :meth:`format_completion` for each - completion. - """ - args, incomplete = self.get_completion_args() - completions = self.get_completions(args, incomplete) - out = [self.format_completion(item) for item in completions] - return "\n".join(out) - - -class BashComplete(ShellComplete): - """Shell completion for Bash.""" - - name = "bash" - source_template = _SOURCE_BASH - - @staticmethod - def _check_version() -> None: - import shutil - import subprocess - - bash_exe = shutil.which("bash") - - if bash_exe is None: - match = None - else: - output = subprocess.run( - [bash_exe, "--norc", "-c", 'echo "${BASH_VERSION}"'], - stdout=subprocess.PIPE, - ) - match = re.search(r"^(\d+)\.(\d+)\.\d+", output.stdout.decode()) - - if match is not None: - major, minor = match.groups() - - if major < "4" or major == "4" and minor < "4": - echo( - _( - "Shell completion is not supported for Bash" - " versions older than 4.4." - ), - err=True, - ) - else: - echo( - _("Couldn't detect Bash version, shell completion is not supported."), - err=True, - ) - - def source(self) -> str: - self._check_version() - return super().source() - - def get_completion_args(self) -> tuple[list[str], str]: - cwords = split_arg_string(os.environ["COMP_WORDS"]) - cword = int(os.environ["COMP_CWORD"]) - args = cwords[1:cword] - - try: - incomplete = cwords[cword] - except IndexError: - incomplete = "" - - return args, incomplete - - def format_completion(self, item: CompletionItem) -> str: - return f"{item.type},{item.value}" - - -class ZshComplete(ShellComplete): - """Shell completion for Zsh.""" - - name = "zsh" - source_template = _SOURCE_ZSH - - def get_completion_args(self) -> tuple[list[str], str]: - cwords = split_arg_string(os.environ["COMP_WORDS"]) - cword = int(os.environ["COMP_CWORD"]) - args = cwords[1:cword] - - try: - incomplete = cwords[cword] - except IndexError: - incomplete = "" - - return args, incomplete - - def format_completion(self, item: CompletionItem) -> str: - return f"{item.type}\n{item.value}\n{item.help if item.help else '_'}" - - -class FishComplete(ShellComplete): - """Shell completion for Fish.""" - - name = "fish" - source_template = _SOURCE_FISH - - def get_completion_args(self) -> tuple[list[str], str]: - cwords = split_arg_string(os.environ["COMP_WORDS"]) - incomplete = os.environ["COMP_CWORD"] - args = cwords[1:] - - # Fish stores the partial word in both COMP_WORDS and - # COMP_CWORD, remove it from complete args. - if incomplete and args and args[-1] == incomplete: - args.pop() - - return args, incomplete - - def format_completion(self, item: CompletionItem) -> str: - if item.help: - return f"{item.type},{item.value}\t{item.help}" - - return f"{item.type},{item.value}" - - -ShellCompleteType = t.TypeVar("ShellCompleteType", bound="type[ShellComplete]") - - -_available_shells: dict[str, type[ShellComplete]] = { - "bash": BashComplete, - "fish": FishComplete, - "zsh": ZshComplete, -} - - -def add_completion_class( - cls: ShellCompleteType, name: str | None = None -) -> ShellCompleteType: - """Register a :class:`ShellComplete` subclass under the given name. - The name will be provided by the completion instruction environment - variable during completion. - - :param cls: The completion class that will handle completion for the - shell. - :param name: Name to register the class under. Defaults to the - class's ``name`` attribute. - """ - if name is None: - name = cls.name - - _available_shells[name] = cls - - return cls - - -def get_completion_class(shell: str) -> type[ShellComplete] | None: - """Look up a registered :class:`ShellComplete` subclass by the name - provided by the completion instruction environment variable. If the - name isn't registered, returns ``None``. - - :param shell: Name the class is registered under. - """ - return _available_shells.get(shell) - - -def split_arg_string(string: str) -> list[str]: - """Split an argument string as with :func:`shlex.split`, but don't - fail if the string is incomplete. Ignores a missing closing quote or - incomplete escape sequence and uses the partial token as-is. - - .. code-block:: python - - split_arg_string("example 'my file") - ["example", "my file"] - - split_arg_string("example my\\") - ["example", "my"] - - :param string: String to split. - - .. versionchanged:: 8.2 - Moved to ``shell_completion`` from ``parser``. - """ - import shlex - - lex = shlex.shlex(string, posix=True) - lex.whitespace_split = True - lex.commenters = "" - out = [] - - try: - for token in lex: - out.append(token) - except ValueError: - # Raised when end-of-string is reached in an invalid state. Use - # the partial token as-is. The quote or escape character is in - # lex.state, not lex.token. - out.append(lex.token) - - return out - - -def _is_incomplete_argument(ctx: Context, param: Parameter) -> bool: - """Determine if the given parameter is an argument that can still - accept values. - - :param ctx: Invocation context for the command represented by the - parsed complete args. - :param param: Argument object being checked. - """ - if not isinstance(param, Argument): - return False - - assert param.name is not None - # Will be None if expose_value is False. - value = ctx.params.get(param.name) - return ( - param.nargs == -1 - or ctx.get_parameter_source(param.name) is not ParameterSource.COMMANDLINE - or ( - param.nargs > 1 - and isinstance(value, (tuple, list)) - and len(value) < param.nargs - ) - ) - - -def _start_of_option(ctx: Context, value: str) -> bool: - """Check if the value looks like the start of an option.""" - if not value: - return False - - c = value[0] - return c in ctx._opt_prefixes - - -def _is_incomplete_option(ctx: Context, args: list[str], param: Parameter) -> bool: - """Determine if the given parameter is an option that needs a value. - - :param args: List of complete args before the incomplete value. - :param param: Option object being checked. - """ - if not isinstance(param, Option): - return False - - if param.is_flag or param.count: - return False - - last_option = None - - for index, arg in enumerate(reversed(args)): - if index + 1 > param.nargs: - break - - if _start_of_option(ctx, arg): - last_option = arg - - return last_option is not None and last_option in param.opts - - -def _resolve_context( - cli: Command, - ctx_args: cabc.MutableMapping[str, t.Any], - prog_name: str, - args: list[str], -) -> Context: - """Produce the context hierarchy starting with the command and - traversing the complete arguments. This only follows the commands, - it doesn't trigger input prompts or callbacks. - - :param cli: Command being called. - :param prog_name: Name of the executable in the shell. - :param args: List of complete args before the incomplete value. - """ - ctx_args["resilient_parsing"] = True - with cli.make_context(prog_name, args.copy(), **ctx_args) as ctx: - args = ctx._protected_args + ctx.args - - while args: - command = ctx.command - - if isinstance(command, Group): - if not command.chain: - name, cmd, args = command.resolve_command(ctx, args) - - if cmd is None: - return ctx - - with cmd.make_context( - name, args, parent=ctx, resilient_parsing=True - ) as sub_ctx: - ctx = sub_ctx - args = ctx._protected_args + ctx.args - else: - sub_ctx = ctx - - while args: - name, cmd, args = command.resolve_command(ctx, args) - - if cmd is None: - return ctx - - with cmd.make_context( - name, - args, - parent=ctx, - allow_extra_args=True, - allow_interspersed_args=False, - resilient_parsing=True, - ) as sub_sub_ctx: - sub_ctx = sub_sub_ctx - args = sub_ctx.args - - ctx = sub_ctx - args = [*sub_ctx._protected_args, *sub_ctx.args] - else: - break - - return ctx - - -def _resolve_incomplete( - ctx: Context, args: list[str], incomplete: str -) -> tuple[Command | Parameter, str]: - """Find the Click object that will handle the completion of the - incomplete value. Return the object and the incomplete value. - - :param ctx: Invocation context for the command represented by - the parsed complete args. - :param args: List of complete args before the incomplete value. - :param incomplete: Value being completed. May be empty. - """ - # Different shells treat an "=" between a long option name and - # value differently. Might keep the value joined, return the "=" - # as a separate item, or return the split name and value. Always - # split and discard the "=" to make completion easier. - if incomplete == "=": - incomplete = "" - elif "=" in incomplete and _start_of_option(ctx, incomplete): - name, _, incomplete = incomplete.partition("=") - args.append(name) - - # The "--" marker tells Click to stop treating values as options - # even if they start with the option character. If it hasn't been - # given and the incomplete arg looks like an option, the current - # command will provide option name completions. - if "--" not in args and _start_of_option(ctx, incomplete): - return ctx.command, incomplete - - params = ctx.command.get_params(ctx) - - # If the last complete arg is an option name with an incomplete - # value, the option will provide value completions. - for param in params: - if _is_incomplete_option(ctx, args, param): - return param, incomplete - - # It's not an option name or value. The first argument without a - # parsed value will provide value completions. - for param in params: - if _is_incomplete_argument(ctx, param): - return param, incomplete - - # There were no unparsed arguments, the command may be a group that - # will provide command name completions. - return ctx.command, incomplete diff --git a/venv/lib/python3.10/site-packages/click/termui.py b/venv/lib/python3.10/site-packages/click/termui.py deleted file mode 100644 index dcbb22216c4bac521eee02fbac97e3b67f302076..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/termui.py +++ /dev/null @@ -1,877 +0,0 @@ -from __future__ import annotations - -import collections.abc as cabc -import inspect -import io -import itertools -import sys -import typing as t -from contextlib import AbstractContextManager -from gettext import gettext as _ - -from ._compat import isatty -from ._compat import strip_ansi -from .exceptions import Abort -from .exceptions import UsageError -from .globals import resolve_color_default -from .types import Choice -from .types import convert_type -from .types import ParamType -from .utils import echo -from .utils import LazyFile - -if t.TYPE_CHECKING: - from ._termui_impl import ProgressBar - -V = t.TypeVar("V") - -# The prompt functions to use. The doc tools currently override these -# functions to customize how they work. -visible_prompt_func: t.Callable[[str], str] = input - -_ansi_colors = { - "black": 30, - "red": 31, - "green": 32, - "yellow": 33, - "blue": 34, - "magenta": 35, - "cyan": 36, - "white": 37, - "reset": 39, - "bright_black": 90, - "bright_red": 91, - "bright_green": 92, - "bright_yellow": 93, - "bright_blue": 94, - "bright_magenta": 95, - "bright_cyan": 96, - "bright_white": 97, -} -_ansi_reset_all = "\033[0m" - - -def hidden_prompt_func(prompt: str) -> str: - import getpass - - return getpass.getpass(prompt) - - -def _build_prompt( - text: str, - suffix: str, - show_default: bool = False, - default: t.Any | None = None, - show_choices: bool = True, - type: ParamType | None = None, -) -> str: - prompt = text - if type is not None and show_choices and isinstance(type, Choice): - prompt += f" ({', '.join(map(str, type.choices))})" - if default is not None and show_default: - prompt = f"{prompt} [{_format_default(default)}]" - return f"{prompt}{suffix}" - - -def _format_default(default: t.Any) -> t.Any: - if isinstance(default, (io.IOBase, LazyFile)) and hasattr(default, "name"): - return default.name - - return default - - -def prompt( - text: str, - default: t.Any | None = None, - hide_input: bool = False, - confirmation_prompt: bool | str = False, - type: ParamType | t.Any | None = None, - value_proc: t.Callable[[str], t.Any] | None = None, - prompt_suffix: str = ": ", - show_default: bool = True, - err: bool = False, - show_choices: bool = True, -) -> t.Any: - """Prompts a user for input. This is a convenience function that can - be used to prompt a user for input later. - - If the user aborts the input by sending an interrupt signal, this - function will catch it and raise a :exc:`Abort` exception. - - :param text: the text to show for the prompt. - :param default: the default value to use if no input happens. If this - is not given it will prompt until it's aborted. - :param hide_input: if this is set to true then the input value will - be hidden. - :param confirmation_prompt: Prompt a second time to confirm the - value. Can be set to a string instead of ``True`` to customize - the message. - :param type: the type to use to check the value against. - :param value_proc: if this parameter is provided it's a function that - is invoked instead of the type conversion to - convert a value. - :param prompt_suffix: a suffix that should be added to the prompt. - :param show_default: shows or hides the default value in the prompt. - :param err: if set to true the file defaults to ``stderr`` instead of - ``stdout``, the same as with echo. - :param show_choices: Show or hide choices if the passed type is a Choice. - For example if type is a Choice of either day or week, - show_choices is true and text is "Group by" then the - prompt will be "Group by (day, week): ". - - .. versionadded:: 8.0 - ``confirmation_prompt`` can be a custom string. - - .. versionadded:: 7.0 - Added the ``show_choices`` parameter. - - .. versionadded:: 6.0 - Added unicode support for cmd.exe on Windows. - - .. versionadded:: 4.0 - Added the `err` parameter. - - """ - - def prompt_func(text: str) -> str: - f = hidden_prompt_func if hide_input else visible_prompt_func - try: - # Write the prompt separately so that we get nice - # coloring through colorama on Windows - echo(text.rstrip(" "), nl=False, err=err) - # Echo a space to stdout to work around an issue where - # readline causes backspace to clear the whole line. - return f(" ") - except (KeyboardInterrupt, EOFError): - # getpass doesn't print a newline if the user aborts input with ^C. - # Allegedly this behavior is inherited from getpass(3). - # A doc bug has been filed at https://bugs.python.org/issue24711 - if hide_input: - echo(None, err=err) - raise Abort() from None - - if value_proc is None: - value_proc = convert_type(type, default) - - prompt = _build_prompt( - text, prompt_suffix, show_default, default, show_choices, type - ) - - if confirmation_prompt: - if confirmation_prompt is True: - confirmation_prompt = _("Repeat for confirmation") - - confirmation_prompt = _build_prompt(confirmation_prompt, prompt_suffix) - - while True: - while True: - value = prompt_func(prompt) - if value: - break - elif default is not None: - value = default - break - try: - result = value_proc(value) - except UsageError as e: - if hide_input: - echo(_("Error: The value you entered was invalid."), err=err) - else: - echo(_("Error: {e.message}").format(e=e), err=err) - continue - if not confirmation_prompt: - return result - while True: - value2 = prompt_func(confirmation_prompt) - is_empty = not value and not value2 - if value2 or is_empty: - break - if value == value2: - return result - echo(_("Error: The two entered values do not match."), err=err) - - -def confirm( - text: str, - default: bool | None = False, - abort: bool = False, - prompt_suffix: str = ": ", - show_default: bool = True, - err: bool = False, -) -> bool: - """Prompts for confirmation (yes/no question). - - If the user aborts the input by sending a interrupt signal this - function will catch it and raise a :exc:`Abort` exception. - - :param text: the question to ask. - :param default: The default value to use when no input is given. If - ``None``, repeat until input is given. - :param abort: if this is set to `True` a negative answer aborts the - exception by raising :exc:`Abort`. - :param prompt_suffix: a suffix that should be added to the prompt. - :param show_default: shows or hides the default value in the prompt. - :param err: if set to true the file defaults to ``stderr`` instead of - ``stdout``, the same as with echo. - - .. versionchanged:: 8.0 - Repeat until input is given if ``default`` is ``None``. - - .. versionadded:: 4.0 - Added the ``err`` parameter. - """ - prompt = _build_prompt( - text, - prompt_suffix, - show_default, - "y/n" if default is None else ("Y/n" if default else "y/N"), - ) - - while True: - try: - # Write the prompt separately so that we get nice - # coloring through colorama on Windows - echo(prompt.rstrip(" "), nl=False, err=err) - # Echo a space to stdout to work around an issue where - # readline causes backspace to clear the whole line. - value = visible_prompt_func(" ").lower().strip() - except (KeyboardInterrupt, EOFError): - raise Abort() from None - if value in ("y", "yes"): - rv = True - elif value in ("n", "no"): - rv = False - elif default is not None and value == "": - rv = default - else: - echo(_("Error: invalid input"), err=err) - continue - break - if abort and not rv: - raise Abort() - return rv - - -def echo_via_pager( - text_or_generator: cabc.Iterable[str] | t.Callable[[], cabc.Iterable[str]] | str, - color: bool | None = None, -) -> None: - """This function takes a text and shows it via an environment specific - pager on stdout. - - .. versionchanged:: 3.0 - Added the `color` flag. - - :param text_or_generator: the text to page, or alternatively, a - generator emitting the text to page. - :param color: controls if the pager supports ANSI colors or not. The - default is autodetection. - """ - color = resolve_color_default(color) - - if inspect.isgeneratorfunction(text_or_generator): - i = t.cast("t.Callable[[], cabc.Iterable[str]]", text_or_generator)() - elif isinstance(text_or_generator, str): - i = [text_or_generator] - else: - i = iter(t.cast("cabc.Iterable[str]", text_or_generator)) - - # convert every element of i to a text type if necessary - text_generator = (el if isinstance(el, str) else str(el) for el in i) - - from ._termui_impl import pager - - return pager(itertools.chain(text_generator, "\n"), color) - - -@t.overload -def progressbar( - *, - length: int, - label: str | None = None, - hidden: bool = False, - show_eta: bool = True, - show_percent: bool | None = None, - show_pos: bool = False, - fill_char: str = "#", - empty_char: str = "-", - bar_template: str = "%(label)s [%(bar)s] %(info)s", - info_sep: str = " ", - width: int = 36, - file: t.TextIO | None = None, - color: bool | None = None, - update_min_steps: int = 1, -) -> ProgressBar[int]: ... - - -@t.overload -def progressbar( - iterable: cabc.Iterable[V] | None = None, - length: int | None = None, - label: str | None = None, - hidden: bool = False, - show_eta: bool = True, - show_percent: bool | None = None, - show_pos: bool = False, - item_show_func: t.Callable[[V | None], str | None] | None = None, - fill_char: str = "#", - empty_char: str = "-", - bar_template: str = "%(label)s [%(bar)s] %(info)s", - info_sep: str = " ", - width: int = 36, - file: t.TextIO | None = None, - color: bool | None = None, - update_min_steps: int = 1, -) -> ProgressBar[V]: ... - - -def progressbar( - iterable: cabc.Iterable[V] | None = None, - length: int | None = None, - label: str | None = None, - hidden: bool = False, - show_eta: bool = True, - show_percent: bool | None = None, - show_pos: bool = False, - item_show_func: t.Callable[[V | None], str | None] | None = None, - fill_char: str = "#", - empty_char: str = "-", - bar_template: str = "%(label)s [%(bar)s] %(info)s", - info_sep: str = " ", - width: int = 36, - file: t.TextIO | None = None, - color: bool | None = None, - update_min_steps: int = 1, -) -> ProgressBar[V]: - """This function creates an iterable context manager that can be used - to iterate over something while showing a progress bar. It will - either iterate over the `iterable` or `length` items (that are counted - up). While iteration happens, this function will print a rendered - progress bar to the given `file` (defaults to stdout) and will attempt - to calculate remaining time and more. By default, this progress bar - will not be rendered if the file is not a terminal. - - The context manager creates the progress bar. When the context - manager is entered the progress bar is already created. With every - iteration over the progress bar, the iterable passed to the bar is - advanced and the bar is updated. When the context manager exits, - a newline is printed and the progress bar is finalized on screen. - - Note: The progress bar is currently designed for use cases where the - total progress can be expected to take at least several seconds. - Because of this, the ProgressBar class object won't display - progress that is considered too fast, and progress where the time - between steps is less than a second. - - No printing must happen or the progress bar will be unintentionally - destroyed. - - Example usage:: - - with progressbar(items) as bar: - for item in bar: - do_something_with(item) - - Alternatively, if no iterable is specified, one can manually update the - progress bar through the `update()` method instead of directly - iterating over the progress bar. The update method accepts the number - of steps to increment the bar with:: - - with progressbar(length=chunks.total_bytes) as bar: - for chunk in chunks: - process_chunk(chunk) - bar.update(chunks.bytes) - - The ``update()`` method also takes an optional value specifying the - ``current_item`` at the new position. This is useful when used - together with ``item_show_func`` to customize the output for each - manual step:: - - with click.progressbar( - length=total_size, - label='Unzipping archive', - item_show_func=lambda a: a.filename - ) as bar: - for archive in zip_file: - archive.extract() - bar.update(archive.size, archive) - - :param iterable: an iterable to iterate over. If not provided the length - is required. - :param length: the number of items to iterate over. By default the - progressbar will attempt to ask the iterator about its - length, which might or might not work. If an iterable is - also provided this parameter can be used to override the - length. If an iterable is not provided the progress bar - will iterate over a range of that length. - :param label: the label to show next to the progress bar. - :param hidden: hide the progressbar. Defaults to ``False``. When no tty is - detected, it will only print the progressbar label. Setting this to - ``False`` also disables that. - :param show_eta: enables or disables the estimated time display. This is - automatically disabled if the length cannot be - determined. - :param show_percent: enables or disables the percentage display. The - default is `True` if the iterable has a length or - `False` if not. - :param show_pos: enables or disables the absolute position display. The - default is `False`. - :param item_show_func: A function called with the current item which - can return a string to show next to the progress bar. If the - function returns ``None`` nothing is shown. The current item can - be ``None``, such as when entering and exiting the bar. - :param fill_char: the character to use to show the filled part of the - progress bar. - :param empty_char: the character to use to show the non-filled part of - the progress bar. - :param bar_template: the format string to use as template for the bar. - The parameters in it are ``label`` for the label, - ``bar`` for the progress bar and ``info`` for the - info section. - :param info_sep: the separator between multiple info items (eta etc.) - :param width: the width of the progress bar in characters, 0 means full - terminal width - :param file: The file to write to. If this is not a terminal then - only the label is printed. - :param color: controls if the terminal supports ANSI colors or not. The - default is autodetection. This is only needed if ANSI - codes are included anywhere in the progress bar output - which is not the case by default. - :param update_min_steps: Render only when this many updates have - completed. This allows tuning for very fast iterators. - - .. versionadded:: 8.2 - The ``hidden`` argument. - - .. versionchanged:: 8.0 - Output is shown even if execution time is less than 0.5 seconds. - - .. versionchanged:: 8.0 - ``item_show_func`` shows the current item, not the previous one. - - .. versionchanged:: 8.0 - Labels are echoed if the output is not a TTY. Reverts a change - in 7.0 that removed all output. - - .. versionadded:: 8.0 - The ``update_min_steps`` parameter. - - .. versionadded:: 4.0 - The ``color`` parameter and ``update`` method. - - .. versionadded:: 2.0 - """ - from ._termui_impl import ProgressBar - - color = resolve_color_default(color) - return ProgressBar( - iterable=iterable, - length=length, - hidden=hidden, - show_eta=show_eta, - show_percent=show_percent, - show_pos=show_pos, - item_show_func=item_show_func, - fill_char=fill_char, - empty_char=empty_char, - bar_template=bar_template, - info_sep=info_sep, - file=file, - label=label, - width=width, - color=color, - update_min_steps=update_min_steps, - ) - - -def clear() -> None: - """Clears the terminal screen. This will have the effect of clearing - the whole visible space of the terminal and moving the cursor to the - top left. This does not do anything if not connected to a terminal. - - .. versionadded:: 2.0 - """ - if not isatty(sys.stdout): - return - - # ANSI escape \033[2J clears the screen, \033[1;1H moves the cursor - echo("\033[2J\033[1;1H", nl=False) - - -def _interpret_color(color: int | tuple[int, int, int] | str, offset: int = 0) -> str: - if isinstance(color, int): - return f"{38 + offset};5;{color:d}" - - if isinstance(color, (tuple, list)): - r, g, b = color - return f"{38 + offset};2;{r:d};{g:d};{b:d}" - - return str(_ansi_colors[color] + offset) - - -def style( - text: t.Any, - fg: int | tuple[int, int, int] | str | None = None, - bg: int | tuple[int, int, int] | str | None = None, - bold: bool | None = None, - dim: bool | None = None, - underline: bool | None = None, - overline: bool | None = None, - italic: bool | None = None, - blink: bool | None = None, - reverse: bool | None = None, - strikethrough: bool | None = None, - reset: bool = True, -) -> str: - """Styles a text with ANSI styles and returns the new string. By - default the styling is self contained which means that at the end - of the string a reset code is issued. This can be prevented by - passing ``reset=False``. - - Examples:: - - click.echo(click.style('Hello World!', fg='green')) - click.echo(click.style('ATTENTION!', blink=True)) - click.echo(click.style('Some things', reverse=True, fg='cyan')) - click.echo(click.style('More colors', fg=(255, 12, 128), bg=117)) - - Supported color names: - - * ``black`` (might be a gray) - * ``red`` - * ``green`` - * ``yellow`` (might be an orange) - * ``blue`` - * ``magenta`` - * ``cyan`` - * ``white`` (might be light gray) - * ``bright_black`` - * ``bright_red`` - * ``bright_green`` - * ``bright_yellow`` - * ``bright_blue`` - * ``bright_magenta`` - * ``bright_cyan`` - * ``bright_white`` - * ``reset`` (reset the color code only) - - If the terminal supports it, color may also be specified as: - - - An integer in the interval [0, 255]. The terminal must support - 8-bit/256-color mode. - - An RGB tuple of three integers in [0, 255]. The terminal must - support 24-bit/true-color mode. - - See https://en.wikipedia.org/wiki/ANSI_color and - https://gist.github.com/XVilka/8346728 for more information. - - :param text: the string to style with ansi codes. - :param fg: if provided this will become the foreground color. - :param bg: if provided this will become the background color. - :param bold: if provided this will enable or disable bold mode. - :param dim: if provided this will enable or disable dim mode. This is - badly supported. - :param underline: if provided this will enable or disable underline. - :param overline: if provided this will enable or disable overline. - :param italic: if provided this will enable or disable italic. - :param blink: if provided this will enable or disable blinking. - :param reverse: if provided this will enable or disable inverse - rendering (foreground becomes background and the - other way round). - :param strikethrough: if provided this will enable or disable - striking through text. - :param reset: by default a reset-all code is added at the end of the - string which means that styles do not carry over. This - can be disabled to compose styles. - - .. versionchanged:: 8.0 - A non-string ``message`` is converted to a string. - - .. versionchanged:: 8.0 - Added support for 256 and RGB color codes. - - .. versionchanged:: 8.0 - Added the ``strikethrough``, ``italic``, and ``overline`` - parameters. - - .. versionchanged:: 7.0 - Added support for bright colors. - - .. versionadded:: 2.0 - """ - if not isinstance(text, str): - text = str(text) - - bits = [] - - if fg: - try: - bits.append(f"\033[{_interpret_color(fg)}m") - except KeyError: - raise TypeError(f"Unknown color {fg!r}") from None - - if bg: - try: - bits.append(f"\033[{_interpret_color(bg, 10)}m") - except KeyError: - raise TypeError(f"Unknown color {bg!r}") from None - - if bold is not None: - bits.append(f"\033[{1 if bold else 22}m") - if dim is not None: - bits.append(f"\033[{2 if dim else 22}m") - if underline is not None: - bits.append(f"\033[{4 if underline else 24}m") - if overline is not None: - bits.append(f"\033[{53 if overline else 55}m") - if italic is not None: - bits.append(f"\033[{3 if italic else 23}m") - if blink is not None: - bits.append(f"\033[{5 if blink else 25}m") - if reverse is not None: - bits.append(f"\033[{7 if reverse else 27}m") - if strikethrough is not None: - bits.append(f"\033[{9 if strikethrough else 29}m") - bits.append(text) - if reset: - bits.append(_ansi_reset_all) - return "".join(bits) - - -def unstyle(text: str) -> str: - """Removes ANSI styling information from a string. Usually it's not - necessary to use this function as Click's echo function will - automatically remove styling if necessary. - - .. versionadded:: 2.0 - - :param text: the text to remove style information from. - """ - return strip_ansi(text) - - -def secho( - message: t.Any | None = None, - file: t.IO[t.AnyStr] | None = None, - nl: bool = True, - err: bool = False, - color: bool | None = None, - **styles: t.Any, -) -> None: - """This function combines :func:`echo` and :func:`style` into one - call. As such the following two calls are the same:: - - click.secho('Hello World!', fg='green') - click.echo(click.style('Hello World!', fg='green')) - - All keyword arguments are forwarded to the underlying functions - depending on which one they go with. - - Non-string types will be converted to :class:`str`. However, - :class:`bytes` are passed directly to :meth:`echo` without applying - style. If you want to style bytes that represent text, call - :meth:`bytes.decode` first. - - .. versionchanged:: 8.0 - A non-string ``message`` is converted to a string. Bytes are - passed through without style applied. - - .. versionadded:: 2.0 - """ - if message is not None and not isinstance(message, (bytes, bytearray)): - message = style(message, **styles) - - return echo(message, file=file, nl=nl, err=err, color=color) - - -@t.overload -def edit( - text: bytes | bytearray, - editor: str | None = None, - env: cabc.Mapping[str, str] | None = None, - require_save: bool = False, - extension: str = ".txt", -) -> bytes | None: ... - - -@t.overload -def edit( - text: str, - editor: str | None = None, - env: cabc.Mapping[str, str] | None = None, - require_save: bool = True, - extension: str = ".txt", -) -> str | None: ... - - -@t.overload -def edit( - text: None = None, - editor: str | None = None, - env: cabc.Mapping[str, str] | None = None, - require_save: bool = True, - extension: str = ".txt", - filename: str | cabc.Iterable[str] | None = None, -) -> None: ... - - -def edit( - text: str | bytes | bytearray | None = None, - editor: str | None = None, - env: cabc.Mapping[str, str] | None = None, - require_save: bool = True, - extension: str = ".txt", - filename: str | cabc.Iterable[str] | None = None, -) -> str | bytes | bytearray | None: - r"""Edits the given text in the defined editor. If an editor is given - (should be the full path to the executable but the regular operating - system search path is used for finding the executable) it overrides - the detected editor. Optionally, some environment variables can be - used. If the editor is closed without changes, `None` is returned. In - case a file is edited directly the return value is always `None` and - `require_save` and `extension` are ignored. - - If the editor cannot be opened a :exc:`UsageError` is raised. - - Note for Windows: to simplify cross-platform usage, the newlines are - automatically converted from POSIX to Windows and vice versa. As such, - the message here will have ``\n`` as newline markers. - - :param text: the text to edit. - :param editor: optionally the editor to use. Defaults to automatic - detection. - :param env: environment variables to forward to the editor. - :param require_save: if this is true, then not saving in the editor - will make the return value become `None`. - :param extension: the extension to tell the editor about. This defaults - to `.txt` but changing this might change syntax - highlighting. - :param filename: if provided it will edit this file instead of the - provided text contents. It will not use a temporary - file as an indirection in that case. If the editor supports - editing multiple files at once, a sequence of files may be - passed as well. Invoke `click.file` once per file instead - if multiple files cannot be managed at once or editing the - files serially is desired. - - .. versionchanged:: 8.2.0 - ``filename`` now accepts any ``Iterable[str]`` in addition to a ``str`` - if the ``editor`` supports editing multiple files at once. - - """ - from ._termui_impl import Editor - - ed = Editor(editor=editor, env=env, require_save=require_save, extension=extension) - - if filename is None: - return ed.edit(text) - - if isinstance(filename, str): - filename = (filename,) - - ed.edit_files(filenames=filename) - return None - - -def launch(url: str, wait: bool = False, locate: bool = False) -> int: - """This function launches the given URL (or filename) in the default - viewer application for this file type. If this is an executable, it - might launch the executable in a new session. The return value is - the exit code of the launched application. Usually, ``0`` indicates - success. - - Examples:: - - click.launch('https://click.palletsprojects.com/') - click.launch('/my/downloaded/file', locate=True) - - .. versionadded:: 2.0 - - :param url: URL or filename of the thing to launch. - :param wait: Wait for the program to exit before returning. This - only works if the launched program blocks. In particular, - ``xdg-open`` on Linux does not block. - :param locate: if this is set to `True` then instead of launching the - application associated with the URL it will attempt to - launch a file manager with the file located. This - might have weird effects if the URL does not point to - the filesystem. - """ - from ._termui_impl import open_url - - return open_url(url, wait=wait, locate=locate) - - -# If this is provided, getchar() calls into this instead. This is used -# for unittesting purposes. -_getchar: t.Callable[[bool], str] | None = None - - -def getchar(echo: bool = False) -> str: - """Fetches a single character from the terminal and returns it. This - will always return a unicode character and under certain rare - circumstances this might return more than one character. The - situations which more than one character is returned is when for - whatever reason multiple characters end up in the terminal buffer or - standard input was not actually a terminal. - - Note that this will always read from the terminal, even if something - is piped into the standard input. - - Note for Windows: in rare cases when typing non-ASCII characters, this - function might wait for a second character and then return both at once. - This is because certain Unicode characters look like special-key markers. - - .. versionadded:: 2.0 - - :param echo: if set to `True`, the character read will also show up on - the terminal. The default is to not show it. - """ - global _getchar - - if _getchar is None: - from ._termui_impl import getchar as f - - _getchar = f - - return _getchar(echo) - - -def raw_terminal() -> AbstractContextManager[int]: - from ._termui_impl import raw_terminal as f - - return f() - - -def pause(info: str | None = None, err: bool = False) -> None: - """This command stops execution and waits for the user to press any - key to continue. This is similar to the Windows batch "pause" - command. If the program is not run through a terminal, this command - will instead do nothing. - - .. versionadded:: 2.0 - - .. versionadded:: 4.0 - Added the `err` parameter. - - :param info: The message to print before pausing. Defaults to - ``"Press any key to continue..."``. - :param err: if set to message goes to ``stderr`` instead of - ``stdout``, the same as with echo. - """ - if not isatty(sys.stdin) or not isatty(sys.stdout): - return - - if info is None: - info = _("Press any key to continue...") - - try: - if info: - echo(info, nl=False, err=err) - try: - getchar() - except (KeyboardInterrupt, EOFError): - pass - finally: - if info: - echo(err=err) diff --git a/venv/lib/python3.10/site-packages/click/testing.py b/venv/lib/python3.10/site-packages/click/testing.py deleted file mode 100644 index 7c0e8741e2d07f54294a82792c776d08c639add4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/testing.py +++ /dev/null @@ -1,565 +0,0 @@ -from __future__ import annotations - -import collections.abc as cabc -import contextlib -import io -import os -import shlex -import shutil -import sys -import tempfile -import typing as t -from types import TracebackType - -from . import _compat -from . import formatting -from . import termui -from . import utils -from ._compat import _find_binary_reader - -if t.TYPE_CHECKING: - from _typeshed import ReadableBuffer - - from .core import Command - - -class EchoingStdin: - def __init__(self, input: t.BinaryIO, output: t.BinaryIO) -> None: - self._input = input - self._output = output - self._paused = False - - def __getattr__(self, x: str) -> t.Any: - return getattr(self._input, x) - - def _echo(self, rv: bytes) -> bytes: - if not self._paused: - self._output.write(rv) - - return rv - - def read(self, n: int = -1) -> bytes: - return self._echo(self._input.read(n)) - - def read1(self, n: int = -1) -> bytes: - return self._echo(self._input.read1(n)) # type: ignore - - def readline(self, n: int = -1) -> bytes: - return self._echo(self._input.readline(n)) - - def readlines(self) -> list[bytes]: - return [self._echo(x) for x in self._input.readlines()] - - def __iter__(self) -> cabc.Iterator[bytes]: - return iter(self._echo(x) for x in self._input) - - def __repr__(self) -> str: - return repr(self._input) - - -@contextlib.contextmanager -def _pause_echo(stream: EchoingStdin | None) -> cabc.Iterator[None]: - if stream is None: - yield - else: - stream._paused = True - yield - stream._paused = False - - -class BytesIOCopy(io.BytesIO): - """Patch ``io.BytesIO`` to let the written stream be copied to another. - - .. versionadded:: 8.2 - """ - - def __init__(self, copy_to: io.BytesIO) -> None: - super().__init__() - self.copy_to = copy_to - - def flush(self) -> None: - super().flush() - self.copy_to.flush() - - def write(self, b: ReadableBuffer) -> int: - self.copy_to.write(b) - return super().write(b) - - -class StreamMixer: - """Mixes `` and `` streams. - - The result is available in the ``output`` attribute. - - .. versionadded:: 8.2 - """ - - def __init__(self) -> None: - self.output: io.BytesIO = io.BytesIO() - self.stdout: io.BytesIO = BytesIOCopy(copy_to=self.output) - self.stderr: io.BytesIO = BytesIOCopy(copy_to=self.output) - - -class _NamedTextIOWrapper(io.TextIOWrapper): - def __init__( - self, buffer: t.BinaryIO, name: str, mode: str, **kwargs: t.Any - ) -> None: - super().__init__(buffer, **kwargs) - self._name = name - self._mode = mode - - @property - def name(self) -> str: - return self._name - - @property - def mode(self) -> str: - return self._mode - - def __next__(self) -> str: # type: ignore - try: - line = super().__next__() - except StopIteration as e: - raise EOFError() from e - return line - - -def make_input_stream( - input: str | bytes | t.IO[t.Any] | None, charset: str -) -> t.BinaryIO: - # Is already an input stream. - if hasattr(input, "read"): - rv = _find_binary_reader(t.cast("t.IO[t.Any]", input)) - - if rv is not None: - return rv - - raise TypeError("Could not find binary reader for input stream.") - - if input is None: - input = b"" - elif isinstance(input, str): - input = input.encode(charset) - - return io.BytesIO(input) - - -class Result: - """Holds the captured result of an invoked CLI script. - - :param runner: The runner that created the result - :param stdout_bytes: The standard output as bytes. - :param stderr_bytes: The standard error as bytes. - :param output_bytes: A mix of ``stdout_bytes`` and ``stderr_bytes``, as the - user would see it in its terminal. - :param return_value: The value returned from the invoked command. - :param exit_code: The exit code as integer. - :param exception: The exception that happened if one did. - :param exc_info: Exception information (exception type, exception instance, - traceback type). - - .. versionchanged:: 8.2 - ``stderr_bytes`` no longer optional, ``output_bytes`` introduced and - ``mix_stderr`` has been removed. - - .. versionadded:: 8.0 - Added ``return_value``. - """ - - def __init__( - self, - runner: CliRunner, - stdout_bytes: bytes, - stderr_bytes: bytes, - output_bytes: bytes, - return_value: t.Any, - exit_code: int, - exception: BaseException | None, - exc_info: tuple[type[BaseException], BaseException, TracebackType] - | None = None, - ): - self.runner = runner - self.stdout_bytes = stdout_bytes - self.stderr_bytes = stderr_bytes - self.output_bytes = output_bytes - self.return_value = return_value - self.exit_code = exit_code - self.exception = exception - self.exc_info = exc_info - - @property - def output(self) -> str: - """The terminal output as unicode string, as the user would see it. - - .. versionchanged:: 8.2 - No longer a proxy for ``self.stdout``. Now has its own independent stream - that is mixing `` and ``, in the order they were written. - """ - return self.output_bytes.decode(self.runner.charset, "replace").replace( - "\r\n", "\n" - ) - - @property - def stdout(self) -> str: - """The standard output as unicode string.""" - return self.stdout_bytes.decode(self.runner.charset, "replace").replace( - "\r\n", "\n" - ) - - @property - def stderr(self) -> str: - """The standard error as unicode string. - - .. versionchanged:: 8.2 - No longer raise an exception, always returns the `` string. - """ - return self.stderr_bytes.decode(self.runner.charset, "replace").replace( - "\r\n", "\n" - ) - - def __repr__(self) -> str: - exc_str = repr(self.exception) if self.exception else "okay" - return f"<{type(self).__name__} {exc_str}>" - - -class CliRunner: - """The CLI runner provides functionality to invoke a Click command line - script for unittesting purposes in a isolated environment. This only - works in single-threaded systems without any concurrency as it changes the - global interpreter state. - - :param charset: the character set for the input and output data. - :param env: a dictionary with environment variables for overriding. - :param echo_stdin: if this is set to `True`, then reading from `` writes - to ``. This is useful for showing examples in - some circumstances. Note that regular prompts - will automatically echo the input. - :param catch_exceptions: Whether to catch any exceptions other than - ``SystemExit`` when running :meth:`~CliRunner.invoke`. - - .. versionchanged:: 8.2 - Added the ``catch_exceptions`` parameter. - - .. versionchanged:: 8.2 - ``mix_stderr`` parameter has been removed. - """ - - def __init__( - self, - charset: str = "utf-8", - env: cabc.Mapping[str, str | None] | None = None, - echo_stdin: bool = False, - catch_exceptions: bool = True, - ) -> None: - self.charset = charset - self.env: cabc.Mapping[str, str | None] = env or {} - self.echo_stdin = echo_stdin - self.catch_exceptions = catch_exceptions - - def get_default_prog_name(self, cli: Command) -> str: - """Given a command object it will return the default program name - for it. The default is the `name` attribute or ``"root"`` if not - set. - """ - return cli.name or "root" - - def make_env( - self, overrides: cabc.Mapping[str, str | None] | None = None - ) -> cabc.Mapping[str, str | None]: - """Returns the environment overrides for invoking a script.""" - rv = dict(self.env) - if overrides: - rv.update(overrides) - return rv - - @contextlib.contextmanager - def isolation( - self, - input: str | bytes | t.IO[t.Any] | None = None, - env: cabc.Mapping[str, str | None] | None = None, - color: bool = False, - ) -> cabc.Iterator[tuple[io.BytesIO, io.BytesIO, io.BytesIO]]: - """A context manager that sets up the isolation for invoking of a - command line tool. This sets up `` with the given input data - and `os.environ` with the overrides from the given dictionary. - This also rebinds some internals in Click to be mocked (like the - prompt functionality). - - This is automatically done in the :meth:`invoke` method. - - :param input: the input stream to put into `sys.stdin`. - :param env: the environment overrides as dictionary. - :param color: whether the output should contain color codes. The - application can still override this explicitly. - - .. versionadded:: 8.2 - An additional output stream is returned, which is a mix of - `` and `` streams. - - .. versionchanged:: 8.2 - Always returns the `` stream. - - .. versionchanged:: 8.0 - `` is opened with ``errors="backslashreplace"`` - instead of the default ``"strict"``. - - .. versionchanged:: 4.0 - Added the ``color`` parameter. - """ - bytes_input = make_input_stream(input, self.charset) - echo_input = None - - old_stdin = sys.stdin - old_stdout = sys.stdout - old_stderr = sys.stderr - old_forced_width = formatting.FORCED_WIDTH - formatting.FORCED_WIDTH = 80 - - env = self.make_env(env) - - stream_mixer = StreamMixer() - - if self.echo_stdin: - bytes_input = echo_input = t.cast( - t.BinaryIO, EchoingStdin(bytes_input, stream_mixer.stdout) - ) - - sys.stdin = text_input = _NamedTextIOWrapper( - bytes_input, encoding=self.charset, name="", mode="r" - ) - - if self.echo_stdin: - # Force unbuffered reads, otherwise TextIOWrapper reads a - # large chunk which is echoed early. - text_input._CHUNK_SIZE = 1 # type: ignore - - sys.stdout = _NamedTextIOWrapper( - stream_mixer.stdout, encoding=self.charset, name="", mode="w" - ) - - sys.stderr = _NamedTextIOWrapper( - stream_mixer.stderr, - encoding=self.charset, - name="", - mode="w", - errors="backslashreplace", - ) - - @_pause_echo(echo_input) # type: ignore - def visible_input(prompt: str | None = None) -> str: - sys.stdout.write(prompt or "") - val = next(text_input).rstrip("\r\n") - sys.stdout.write(f"{val}\n") - sys.stdout.flush() - return val - - @_pause_echo(echo_input) # type: ignore - def hidden_input(prompt: str | None = None) -> str: - sys.stdout.write(f"{prompt or ''}\n") - sys.stdout.flush() - return next(text_input).rstrip("\r\n") - - @_pause_echo(echo_input) # type: ignore - def _getchar(echo: bool) -> str: - char = sys.stdin.read(1) - - if echo: - sys.stdout.write(char) - - sys.stdout.flush() - return char - - default_color = color - - def should_strip_ansi( - stream: t.IO[t.Any] | None = None, color: bool | None = None - ) -> bool: - if color is None: - return not default_color - return not color - - old_visible_prompt_func = termui.visible_prompt_func - old_hidden_prompt_func = termui.hidden_prompt_func - old__getchar_func = termui._getchar - old_should_strip_ansi = utils.should_strip_ansi # type: ignore - old__compat_should_strip_ansi = _compat.should_strip_ansi - termui.visible_prompt_func = visible_input - termui.hidden_prompt_func = hidden_input - termui._getchar = _getchar - utils.should_strip_ansi = should_strip_ansi # type: ignore - _compat.should_strip_ansi = should_strip_ansi - - old_env = {} - try: - for key, value in env.items(): - old_env[key] = os.environ.get(key) - if value is None: - try: - del os.environ[key] - except Exception: - pass - else: - os.environ[key] = value - yield (stream_mixer.stdout, stream_mixer.stderr, stream_mixer.output) - finally: - for key, value in old_env.items(): - if value is None: - try: - del os.environ[key] - except Exception: - pass - else: - os.environ[key] = value - sys.stdout = old_stdout - sys.stderr = old_stderr - sys.stdin = old_stdin - termui.visible_prompt_func = old_visible_prompt_func - termui.hidden_prompt_func = old_hidden_prompt_func - termui._getchar = old__getchar_func - utils.should_strip_ansi = old_should_strip_ansi # type: ignore - _compat.should_strip_ansi = old__compat_should_strip_ansi - formatting.FORCED_WIDTH = old_forced_width - - def invoke( - self, - cli: Command, - args: str | cabc.Sequence[str] | None = None, - input: str | bytes | t.IO[t.Any] | None = None, - env: cabc.Mapping[str, str | None] | None = None, - catch_exceptions: bool | None = None, - color: bool = False, - **extra: t.Any, - ) -> Result: - """Invokes a command in an isolated environment. The arguments are - forwarded directly to the command line script, the `extra` keyword - arguments are passed to the :meth:`~clickpkg.Command.main` function of - the command. - - This returns a :class:`Result` object. - - :param cli: the command to invoke - :param args: the arguments to invoke. It may be given as an iterable - or a string. When given as string it will be interpreted - as a Unix shell command. More details at - :func:`shlex.split`. - :param input: the input data for `sys.stdin`. - :param env: the environment overrides. - :param catch_exceptions: Whether to catch any other exceptions than - ``SystemExit``. If :data:`None`, the value - from :class:`CliRunner` is used. - :param extra: the keyword arguments to pass to :meth:`main`. - :param color: whether the output should contain color codes. The - application can still override this explicitly. - - .. versionadded:: 8.2 - The result object has the ``output_bytes`` attribute with - the mix of ``stdout_bytes`` and ``stderr_bytes``, as the user would - see it in its terminal. - - .. versionchanged:: 8.2 - The result object always returns the ``stderr_bytes`` stream. - - .. versionchanged:: 8.0 - The result object has the ``return_value`` attribute with - the value returned from the invoked command. - - .. versionchanged:: 4.0 - Added the ``color`` parameter. - - .. versionchanged:: 3.0 - Added the ``catch_exceptions`` parameter. - - .. versionchanged:: 3.0 - The result object has the ``exc_info`` attribute with the - traceback if available. - """ - exc_info = None - if catch_exceptions is None: - catch_exceptions = self.catch_exceptions - - with self.isolation(input=input, env=env, color=color) as outstreams: - return_value = None - exception: BaseException | None = None - exit_code = 0 - - if isinstance(args, str): - args = shlex.split(args) - - try: - prog_name = extra.pop("prog_name") - except KeyError: - prog_name = self.get_default_prog_name(cli) - - try: - return_value = cli.main(args=args or (), prog_name=prog_name, **extra) - except SystemExit as e: - exc_info = sys.exc_info() - e_code = t.cast("int | t.Any | None", e.code) - - if e_code is None: - e_code = 0 - - if e_code != 0: - exception = e - - if not isinstance(e_code, int): - sys.stdout.write(str(e_code)) - sys.stdout.write("\n") - e_code = 1 - - exit_code = e_code - - except Exception as e: - if not catch_exceptions: - raise - exception = e - exit_code = 1 - exc_info = sys.exc_info() - finally: - sys.stdout.flush() - sys.stderr.flush() - stdout = outstreams[0].getvalue() - stderr = outstreams[1].getvalue() - output = outstreams[2].getvalue() - - return Result( - runner=self, - stdout_bytes=stdout, - stderr_bytes=stderr, - output_bytes=output, - return_value=return_value, - exit_code=exit_code, - exception=exception, - exc_info=exc_info, # type: ignore - ) - - @contextlib.contextmanager - def isolated_filesystem( - self, temp_dir: str | os.PathLike[str] | None = None - ) -> cabc.Iterator[str]: - """A context manager that creates a temporary directory and - changes the current working directory to it. This isolates tests - that affect the contents of the CWD to prevent them from - interfering with each other. - - :param temp_dir: Create the temporary directory under this - directory. If given, the created directory is not removed - when exiting. - - .. versionchanged:: 8.0 - Added the ``temp_dir`` parameter. - """ - cwd = os.getcwd() - dt = tempfile.mkdtemp(dir=temp_dir) - os.chdir(dt) - - try: - yield dt - finally: - os.chdir(cwd) - - if temp_dir is None: - try: - shutil.rmtree(dt) - except OSError: - pass diff --git a/venv/lib/python3.10/site-packages/click/types.py b/venv/lib/python3.10/site-packages/click/types.py deleted file mode 100644 index 684cb3b1e6107bada194d18968de69a7f383a2a0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/types.py +++ /dev/null @@ -1,1165 +0,0 @@ -from __future__ import annotations - -import collections.abc as cabc -import enum -import os -import stat -import sys -import typing as t -from datetime import datetime -from gettext import gettext as _ -from gettext import ngettext - -from ._compat import _get_argv_encoding -from ._compat import open_stream -from .exceptions import BadParameter -from .utils import format_filename -from .utils import LazyFile -from .utils import safecall - -if t.TYPE_CHECKING: - import typing_extensions as te - - from .core import Context - from .core import Parameter - from .shell_completion import CompletionItem - -ParamTypeValue = t.TypeVar("ParamTypeValue") - - -class ParamType: - """Represents the type of a parameter. Validates and converts values - from the command line or Python into the correct type. - - To implement a custom type, subclass and implement at least the - following: - - - The :attr:`name` class attribute must be set. - - Calling an instance of the type with ``None`` must return - ``None``. This is already implemented by default. - - :meth:`convert` must convert string values to the correct type. - - :meth:`convert` must accept values that are already the correct - type. - - It must be able to convert a value if the ``ctx`` and ``param`` - arguments are ``None``. This can occur when converting prompt - input. - """ - - is_composite: t.ClassVar[bool] = False - arity: t.ClassVar[int] = 1 - - #: the descriptive name of this type - name: str - - #: if a list of this type is expected and the value is pulled from a - #: string environment variable, this is what splits it up. `None` - #: means any whitespace. For all parameters the general rule is that - #: whitespace splits them up. The exception are paths and files which - #: are split by ``os.path.pathsep`` by default (":" on Unix and ";" on - #: Windows). - envvar_list_splitter: t.ClassVar[str | None] = None - - def to_info_dict(self) -> dict[str, t.Any]: - """Gather information that could be useful for a tool generating - user-facing documentation. - - Use :meth:`click.Context.to_info_dict` to traverse the entire - CLI structure. - - .. versionadded:: 8.0 - """ - # The class name without the "ParamType" suffix. - param_type = type(self).__name__.partition("ParamType")[0] - param_type = param_type.partition("ParameterType")[0] - - # Custom subclasses might not remember to set a name. - if hasattr(self, "name"): - name = self.name - else: - name = param_type - - return {"param_type": param_type, "name": name} - - def __call__( - self, - value: t.Any, - param: Parameter | None = None, - ctx: Context | None = None, - ) -> t.Any: - if value is not None: - return self.convert(value, param, ctx) - - def get_metavar(self, param: Parameter, ctx: Context) -> str | None: - """Returns the metavar default for this param if it provides one.""" - - def get_missing_message(self, param: Parameter, ctx: Context | None) -> str | None: - """Optionally might return extra information about a missing - parameter. - - .. versionadded:: 2.0 - """ - - def convert( - self, value: t.Any, param: Parameter | None, ctx: Context | None - ) -> t.Any: - """Convert the value to the correct type. This is not called if - the value is ``None`` (the missing value). - - This must accept string values from the command line, as well as - values that are already the correct type. It may also convert - other compatible types. - - The ``param`` and ``ctx`` arguments may be ``None`` in certain - situations, such as when converting prompt input. - - If the value cannot be converted, call :meth:`fail` with a - descriptive message. - - :param value: The value to convert. - :param param: The parameter that is using this type to convert - its value. May be ``None``. - :param ctx: The current context that arrived at this value. May - be ``None``. - """ - return value - - def split_envvar_value(self, rv: str) -> cabc.Sequence[str]: - """Given a value from an environment variable this splits it up - into small chunks depending on the defined envvar list splitter. - - If the splitter is set to `None`, which means that whitespace splits, - then leading and trailing whitespace is ignored. Otherwise, leading - and trailing splitters usually lead to empty items being included. - """ - return (rv or "").split(self.envvar_list_splitter) - - def fail( - self, - message: str, - param: Parameter | None = None, - ctx: Context | None = None, - ) -> t.NoReturn: - """Helper method to fail with an invalid value message.""" - raise BadParameter(message, ctx=ctx, param=param) - - def shell_complete( - self, ctx: Context, param: Parameter, incomplete: str - ) -> list[CompletionItem]: - """Return a list of - :class:`~click.shell_completion.CompletionItem` objects for the - incomplete value. Most types do not provide completions, but - some do, and this allows custom types to provide custom - completions as well. - - :param ctx: Invocation context for this command. - :param param: The parameter that is requesting completion. - :param incomplete: Value being completed. May be empty. - - .. versionadded:: 8.0 - """ - return [] - - -class CompositeParamType(ParamType): - is_composite = True - - @property - def arity(self) -> int: # type: ignore - raise NotImplementedError() - - -class FuncParamType(ParamType): - def __init__(self, func: t.Callable[[t.Any], t.Any]) -> None: - self.name: str = func.__name__ - self.func = func - - def to_info_dict(self) -> dict[str, t.Any]: - info_dict = super().to_info_dict() - info_dict["func"] = self.func - return info_dict - - def convert( - self, value: t.Any, param: Parameter | None, ctx: Context | None - ) -> t.Any: - try: - return self.func(value) - except ValueError: - try: - value = str(value) - except UnicodeError: - value = value.decode("utf-8", "replace") - - self.fail(value, param, ctx) - - -class UnprocessedParamType(ParamType): - name = "text" - - def convert( - self, value: t.Any, param: Parameter | None, ctx: Context | None - ) -> t.Any: - return value - - def __repr__(self) -> str: - return "UNPROCESSED" - - -class StringParamType(ParamType): - name = "text" - - def convert( - self, value: t.Any, param: Parameter | None, ctx: Context | None - ) -> t.Any: - if isinstance(value, bytes): - enc = _get_argv_encoding() - try: - value = value.decode(enc) - except UnicodeError: - fs_enc = sys.getfilesystemencoding() - if fs_enc != enc: - try: - value = value.decode(fs_enc) - except UnicodeError: - value = value.decode("utf-8", "replace") - else: - value = value.decode("utf-8", "replace") - return value - return str(value) - - def __repr__(self) -> str: - return "STRING" - - -class Choice(ParamType, t.Generic[ParamTypeValue]): - """The choice type allows a value to be checked against a fixed set - of supported values. - - You may pass any iterable value which will be converted to a tuple - and thus will only be iterated once. - - The resulting value will always be one of the originally passed choices. - See :meth:`normalize_choice` for more info on the mapping of strings - to choices. See :ref:`choice-opts` for an example. - - :param case_sensitive: Set to false to make choices case - insensitive. Defaults to true. - - .. versionchanged:: 8.2.0 - Non-``str`` ``choices`` are now supported. It can additionally be any - iterable. Before you were not recommended to pass anything but a list or - tuple. - - .. versionadded:: 8.2.0 - Choice normalization can be overridden via :meth:`normalize_choice`. - """ - - name = "choice" - - def __init__( - self, choices: cabc.Iterable[ParamTypeValue], case_sensitive: bool = True - ) -> None: - self.choices: cabc.Sequence[ParamTypeValue] = tuple(choices) - self.case_sensitive = case_sensitive - - def to_info_dict(self) -> dict[str, t.Any]: - info_dict = super().to_info_dict() - info_dict["choices"] = self.choices - info_dict["case_sensitive"] = self.case_sensitive - return info_dict - - def _normalized_mapping( - self, ctx: Context | None = None - ) -> cabc.Mapping[ParamTypeValue, str]: - """ - Returns mapping where keys are the original choices and the values are - the normalized values that are accepted via the command line. - - This is a simple wrapper around :meth:`normalize_choice`, use that - instead which is supported. - """ - return { - choice: self.normalize_choice( - choice=choice, - ctx=ctx, - ) - for choice in self.choices - } - - def normalize_choice(self, choice: ParamTypeValue, ctx: Context | None) -> str: - """ - Normalize a choice value, used to map a passed string to a choice. - Each choice must have a unique normalized value. - - By default uses :meth:`Context.token_normalize_func` and if not case - sensitive, convert it to a casefolded value. - - .. versionadded:: 8.2.0 - """ - normed_value = choice.name if isinstance(choice, enum.Enum) else str(choice) - - if ctx is not None and ctx.token_normalize_func is not None: - normed_value = ctx.token_normalize_func(normed_value) - - if not self.case_sensitive: - normed_value = normed_value.casefold() - - return normed_value - - def get_metavar(self, param: Parameter, ctx: Context) -> str | None: - if param.param_type_name == "option" and not param.show_choices: # type: ignore - choice_metavars = [ - convert_type(type(choice)).name.upper() for choice in self.choices - ] - choices_str = "|".join([*dict.fromkeys(choice_metavars)]) - else: - choices_str = "|".join( - [str(i) for i in self._normalized_mapping(ctx=ctx).values()] - ) - - # Use curly braces to indicate a required argument. - if param.required and param.param_type_name == "argument": - return f"{{{choices_str}}}" - - # Use square braces to indicate an option or optional argument. - return f"[{choices_str}]" - - def get_missing_message(self, param: Parameter, ctx: Context | None) -> str: - """ - Message shown when no choice is passed. - - .. versionchanged:: 8.2.0 Added ``ctx`` argument. - """ - return _("Choose from:\n\t{choices}").format( - choices=",\n\t".join(self._normalized_mapping(ctx=ctx).values()) - ) - - def convert( - self, value: t.Any, param: Parameter | None, ctx: Context | None - ) -> ParamTypeValue: - """ - For a given value from the parser, normalize it and find its - matching normalized value in the list of choices. Then return the - matched "original" choice. - """ - normed_value = self.normalize_choice(choice=value, ctx=ctx) - normalized_mapping = self._normalized_mapping(ctx=ctx) - - try: - return next( - original - for original, normalized in normalized_mapping.items() - if normalized == normed_value - ) - except StopIteration: - self.fail( - self.get_invalid_choice_message(value=value, ctx=ctx), - param=param, - ctx=ctx, - ) - - def get_invalid_choice_message(self, value: t.Any, ctx: Context | None) -> str: - """Get the error message when the given choice is invalid. - - :param value: The invalid value. - - .. versionadded:: 8.2 - """ - choices_str = ", ".join(map(repr, self._normalized_mapping(ctx=ctx).values())) - return ngettext( - "{value!r} is not {choice}.", - "{value!r} is not one of {choices}.", - len(self.choices), - ).format(value=value, choice=choices_str, choices=choices_str) - - def __repr__(self) -> str: - return f"Choice({list(self.choices)})" - - def shell_complete( - self, ctx: Context, param: Parameter, incomplete: str - ) -> list[CompletionItem]: - """Complete choices that start with the incomplete value. - - :param ctx: Invocation context for this command. - :param param: The parameter that is requesting completion. - :param incomplete: Value being completed. May be empty. - - .. versionadded:: 8.0 - """ - from click.shell_completion import CompletionItem - - str_choices = map(str, self.choices) - - if self.case_sensitive: - matched = (c for c in str_choices if c.startswith(incomplete)) - else: - incomplete = incomplete.lower() - matched = (c for c in str_choices if c.lower().startswith(incomplete)) - - return [CompletionItem(c) for c in matched] - - -class DateTime(ParamType): - """The DateTime type converts date strings into `datetime` objects. - - The format strings which are checked are configurable, but default to some - common (non-timezone aware) ISO 8601 formats. - - When specifying *DateTime* formats, you should only pass a list or a tuple. - Other iterables, like generators, may lead to surprising results. - - The format strings are processed using ``datetime.strptime``, and this - consequently defines the format strings which are allowed. - - Parsing is tried using each format, in order, and the first format which - parses successfully is used. - - :param formats: A list or tuple of date format strings, in the order in - which they should be tried. Defaults to - ``'%Y-%m-%d'``, ``'%Y-%m-%dT%H:%M:%S'``, - ``'%Y-%m-%d %H:%M:%S'``. - """ - - name = "datetime" - - def __init__(self, formats: cabc.Sequence[str] | None = None): - self.formats: cabc.Sequence[str] = formats or [ - "%Y-%m-%d", - "%Y-%m-%dT%H:%M:%S", - "%Y-%m-%d %H:%M:%S", - ] - - def to_info_dict(self) -> dict[str, t.Any]: - info_dict = super().to_info_dict() - info_dict["formats"] = self.formats - return info_dict - - def get_metavar(self, param: Parameter, ctx: Context) -> str | None: - return f"[{'|'.join(self.formats)}]" - - def _try_to_convert_date(self, value: t.Any, format: str) -> datetime | None: - try: - return datetime.strptime(value, format) - except ValueError: - return None - - def convert( - self, value: t.Any, param: Parameter | None, ctx: Context | None - ) -> t.Any: - if isinstance(value, datetime): - return value - - for format in self.formats: - converted = self._try_to_convert_date(value, format) - - if converted is not None: - return converted - - formats_str = ", ".join(map(repr, self.formats)) - self.fail( - ngettext( - "{value!r} does not match the format {format}.", - "{value!r} does not match the formats {formats}.", - len(self.formats), - ).format(value=value, format=formats_str, formats=formats_str), - param, - ctx, - ) - - def __repr__(self) -> str: - return "DateTime" - - -class _NumberParamTypeBase(ParamType): - _number_class: t.ClassVar[type[t.Any]] - - def convert( - self, value: t.Any, param: Parameter | None, ctx: Context | None - ) -> t.Any: - try: - return self._number_class(value) - except ValueError: - self.fail( - _("{value!r} is not a valid {number_type}.").format( - value=value, number_type=self.name - ), - param, - ctx, - ) - - -class _NumberRangeBase(_NumberParamTypeBase): - def __init__( - self, - min: float | None = None, - max: float | None = None, - min_open: bool = False, - max_open: bool = False, - clamp: bool = False, - ) -> None: - self.min = min - self.max = max - self.min_open = min_open - self.max_open = max_open - self.clamp = clamp - - def to_info_dict(self) -> dict[str, t.Any]: - info_dict = super().to_info_dict() - info_dict.update( - min=self.min, - max=self.max, - min_open=self.min_open, - max_open=self.max_open, - clamp=self.clamp, - ) - return info_dict - - def convert( - self, value: t.Any, param: Parameter | None, ctx: Context | None - ) -> t.Any: - import operator - - rv = super().convert(value, param, ctx) - lt_min: bool = self.min is not None and ( - operator.le if self.min_open else operator.lt - )(rv, self.min) - gt_max: bool = self.max is not None and ( - operator.ge if self.max_open else operator.gt - )(rv, self.max) - - if self.clamp: - if lt_min: - return self._clamp(self.min, 1, self.min_open) # type: ignore - - if gt_max: - return self._clamp(self.max, -1, self.max_open) # type: ignore - - if lt_min or gt_max: - self.fail( - _("{value} is not in the range {range}.").format( - value=rv, range=self._describe_range() - ), - param, - ctx, - ) - - return rv - - def _clamp(self, bound: float, dir: t.Literal[1, -1], open: bool) -> float: - """Find the valid value to clamp to bound in the given - direction. - - :param bound: The boundary value. - :param dir: 1 or -1 indicating the direction to move. - :param open: If true, the range does not include the bound. - """ - raise NotImplementedError - - def _describe_range(self) -> str: - """Describe the range for use in help text.""" - if self.min is None: - op = "<" if self.max_open else "<=" - return f"x{op}{self.max}" - - if self.max is None: - op = ">" if self.min_open else ">=" - return f"x{op}{self.min}" - - lop = "<" if self.min_open else "<=" - rop = "<" if self.max_open else "<=" - return f"{self.min}{lop}x{rop}{self.max}" - - def __repr__(self) -> str: - clamp = " clamped" if self.clamp else "" - return f"<{type(self).__name__} {self._describe_range()}{clamp}>" - - -class IntParamType(_NumberParamTypeBase): - name = "integer" - _number_class = int - - def __repr__(self) -> str: - return "INT" - - -class IntRange(_NumberRangeBase, IntParamType): - """Restrict an :data:`click.INT` value to a range of accepted - values. See :ref:`ranges`. - - If ``min`` or ``max`` are not passed, any value is accepted in that - direction. If ``min_open`` or ``max_open`` are enabled, the - corresponding boundary is not included in the range. - - If ``clamp`` is enabled, a value outside the range is clamped to the - boundary instead of failing. - - .. versionchanged:: 8.0 - Added the ``min_open`` and ``max_open`` parameters. - """ - - name = "integer range" - - def _clamp( # type: ignore - self, bound: int, dir: t.Literal[1, -1], open: bool - ) -> int: - if not open: - return bound - - return bound + dir - - -class FloatParamType(_NumberParamTypeBase): - name = "float" - _number_class = float - - def __repr__(self) -> str: - return "FLOAT" - - -class FloatRange(_NumberRangeBase, FloatParamType): - """Restrict a :data:`click.FLOAT` value to a range of accepted - values. See :ref:`ranges`. - - If ``min`` or ``max`` are not passed, any value is accepted in that - direction. If ``min_open`` or ``max_open`` are enabled, the - corresponding boundary is not included in the range. - - If ``clamp`` is enabled, a value outside the range is clamped to the - boundary instead of failing. This is not supported if either - boundary is marked ``open``. - - .. versionchanged:: 8.0 - Added the ``min_open`` and ``max_open`` parameters. - """ - - name = "float range" - - def __init__( - self, - min: float | None = None, - max: float | None = None, - min_open: bool = False, - max_open: bool = False, - clamp: bool = False, - ) -> None: - super().__init__( - min=min, max=max, min_open=min_open, max_open=max_open, clamp=clamp - ) - - if (min_open or max_open) and clamp: - raise TypeError("Clamping is not supported for open bounds.") - - def _clamp(self, bound: float, dir: t.Literal[1, -1], open: bool) -> float: - if not open: - return bound - - # Could use math.nextafter here, but clamping an - # open float range doesn't seem to be particularly useful. It's - # left up to the user to write a callback to do it if needed. - raise RuntimeError("Clamping is not supported for open bounds.") - - -class BoolParamType(ParamType): - name = "boolean" - - def convert( - self, value: t.Any, param: Parameter | None, ctx: Context | None - ) -> t.Any: - if value in {False, True}: - return bool(value) - - norm = value.strip().lower() - - if norm in {"1", "true", "t", "yes", "y", "on"}: - return True - - if norm in {"0", "false", "f", "no", "n", "off"}: - return False - - self.fail( - _("{value!r} is not a valid boolean.").format(value=value), param, ctx - ) - - def __repr__(self) -> str: - return "BOOL" - - -class UUIDParameterType(ParamType): - name = "uuid" - - def convert( - self, value: t.Any, param: Parameter | None, ctx: Context | None - ) -> t.Any: - import uuid - - if isinstance(value, uuid.UUID): - return value - - value = value.strip() - - try: - return uuid.UUID(value) - except ValueError: - self.fail( - _("{value!r} is not a valid UUID.").format(value=value), param, ctx - ) - - def __repr__(self) -> str: - return "UUID" - - -class File(ParamType): - """Declares a parameter to be a file for reading or writing. The file - is automatically closed once the context tears down (after the command - finished working). - - Files can be opened for reading or writing. The special value ``-`` - indicates stdin or stdout depending on the mode. - - By default, the file is opened for reading text data, but it can also be - opened in binary mode or for writing. The encoding parameter can be used - to force a specific encoding. - - The `lazy` flag controls if the file should be opened immediately or upon - first IO. The default is to be non-lazy for standard input and output - streams as well as files opened for reading, `lazy` otherwise. When opening a - file lazily for reading, it is still opened temporarily for validation, but - will not be held open until first IO. lazy is mainly useful when opening - for writing to avoid creating the file until it is needed. - - Files can also be opened atomically in which case all writes go into a - separate file in the same folder and upon completion the file will - be moved over to the original location. This is useful if a file - regularly read by other users is modified. - - See :ref:`file-args` for more information. - - .. versionchanged:: 2.0 - Added the ``atomic`` parameter. - """ - - name = "filename" - envvar_list_splitter: t.ClassVar[str] = os.path.pathsep - - def __init__( - self, - mode: str = "r", - encoding: str | None = None, - errors: str | None = "strict", - lazy: bool | None = None, - atomic: bool = False, - ) -> None: - self.mode = mode - self.encoding = encoding - self.errors = errors - self.lazy = lazy - self.atomic = atomic - - def to_info_dict(self) -> dict[str, t.Any]: - info_dict = super().to_info_dict() - info_dict.update(mode=self.mode, encoding=self.encoding) - return info_dict - - def resolve_lazy_flag(self, value: str | os.PathLike[str]) -> bool: - if self.lazy is not None: - return self.lazy - if os.fspath(value) == "-": - return False - elif "w" in self.mode: - return True - return False - - def convert( - self, - value: str | os.PathLike[str] | t.IO[t.Any], - param: Parameter | None, - ctx: Context | None, - ) -> t.IO[t.Any]: - if _is_file_like(value): - return value - - value = t.cast("str | os.PathLike[str]", value) - - try: - lazy = self.resolve_lazy_flag(value) - - if lazy: - lf = LazyFile( - value, self.mode, self.encoding, self.errors, atomic=self.atomic - ) - - if ctx is not None: - ctx.call_on_close(lf.close_intelligently) - - return t.cast("t.IO[t.Any]", lf) - - f, should_close = open_stream( - value, self.mode, self.encoding, self.errors, atomic=self.atomic - ) - - # If a context is provided, we automatically close the file - # at the end of the context execution (or flush out). If a - # context does not exist, it's the caller's responsibility to - # properly close the file. This for instance happens when the - # type is used with prompts. - if ctx is not None: - if should_close: - ctx.call_on_close(safecall(f.close)) - else: - ctx.call_on_close(safecall(f.flush)) - - return f - except OSError as e: - self.fail(f"'{format_filename(value)}': {e.strerror}", param, ctx) - - def shell_complete( - self, ctx: Context, param: Parameter, incomplete: str - ) -> list[CompletionItem]: - """Return a special completion marker that tells the completion - system to use the shell to provide file path completions. - - :param ctx: Invocation context for this command. - :param param: The parameter that is requesting completion. - :param incomplete: Value being completed. May be empty. - - .. versionadded:: 8.0 - """ - from click.shell_completion import CompletionItem - - return [CompletionItem(incomplete, type="file")] - - -def _is_file_like(value: t.Any) -> te.TypeGuard[t.IO[t.Any]]: - return hasattr(value, "read") or hasattr(value, "write") - - -class Path(ParamType): - """The ``Path`` type is similar to the :class:`File` type, but - returns the filename instead of an open file. Various checks can be - enabled to validate the type of file and permissions. - - :param exists: The file or directory needs to exist for the value to - be valid. If this is not set to ``True``, and the file does not - exist, then all further checks are silently skipped. - :param file_okay: Allow a file as a value. - :param dir_okay: Allow a directory as a value. - :param readable: if true, a readable check is performed. - :param writable: if true, a writable check is performed. - :param executable: if true, an executable check is performed. - :param resolve_path: Make the value absolute and resolve any - symlinks. A ``~`` is not expanded, as this is supposed to be - done by the shell only. - :param allow_dash: Allow a single dash as a value, which indicates - a standard stream (but does not open it). Use - :func:`~click.open_file` to handle opening this value. - :param path_type: Convert the incoming path value to this type. If - ``None``, keep Python's default, which is ``str``. Useful to - convert to :class:`pathlib.Path`. - - .. versionchanged:: 8.1 - Added the ``executable`` parameter. - - .. versionchanged:: 8.0 - Allow passing ``path_type=pathlib.Path``. - - .. versionchanged:: 6.0 - Added the ``allow_dash`` parameter. - """ - - envvar_list_splitter: t.ClassVar[str] = os.path.pathsep - - def __init__( - self, - exists: bool = False, - file_okay: bool = True, - dir_okay: bool = True, - writable: bool = False, - readable: bool = True, - resolve_path: bool = False, - allow_dash: bool = False, - path_type: type[t.Any] | None = None, - executable: bool = False, - ): - self.exists = exists - self.file_okay = file_okay - self.dir_okay = dir_okay - self.readable = readable - self.writable = writable - self.executable = executable - self.resolve_path = resolve_path - self.allow_dash = allow_dash - self.type = path_type - - if self.file_okay and not self.dir_okay: - self.name: str = _("file") - elif self.dir_okay and not self.file_okay: - self.name = _("directory") - else: - self.name = _("path") - - def to_info_dict(self) -> dict[str, t.Any]: - info_dict = super().to_info_dict() - info_dict.update( - exists=self.exists, - file_okay=self.file_okay, - dir_okay=self.dir_okay, - writable=self.writable, - readable=self.readable, - allow_dash=self.allow_dash, - ) - return info_dict - - def coerce_path_result( - self, value: str | os.PathLike[str] - ) -> str | bytes | os.PathLike[str]: - if self.type is not None and not isinstance(value, self.type): - if self.type is str: - return os.fsdecode(value) - elif self.type is bytes: - return os.fsencode(value) - else: - return t.cast("os.PathLike[str]", self.type(value)) - - return value - - def convert( - self, - value: str | os.PathLike[str], - param: Parameter | None, - ctx: Context | None, - ) -> str | bytes | os.PathLike[str]: - rv = value - - is_dash = self.file_okay and self.allow_dash and rv in (b"-", "-") - - if not is_dash: - if self.resolve_path: - rv = os.path.realpath(rv) - - try: - st = os.stat(rv) - except OSError: - if not self.exists: - return self.coerce_path_result(rv) - self.fail( - _("{name} {filename!r} does not exist.").format( - name=self.name.title(), filename=format_filename(value) - ), - param, - ctx, - ) - - if not self.file_okay and stat.S_ISREG(st.st_mode): - self.fail( - _("{name} {filename!r} is a file.").format( - name=self.name.title(), filename=format_filename(value) - ), - param, - ctx, - ) - if not self.dir_okay and stat.S_ISDIR(st.st_mode): - self.fail( - _("{name} {filename!r} is a directory.").format( - name=self.name.title(), filename=format_filename(value) - ), - param, - ctx, - ) - - if self.readable and not os.access(rv, os.R_OK): - self.fail( - _("{name} {filename!r} is not readable.").format( - name=self.name.title(), filename=format_filename(value) - ), - param, - ctx, - ) - - if self.writable and not os.access(rv, os.W_OK): - self.fail( - _("{name} {filename!r} is not writable.").format( - name=self.name.title(), filename=format_filename(value) - ), - param, - ctx, - ) - - if self.executable and not os.access(value, os.X_OK): - self.fail( - _("{name} {filename!r} is not executable.").format( - name=self.name.title(), filename=format_filename(value) - ), - param, - ctx, - ) - - return self.coerce_path_result(rv) - - def shell_complete( - self, ctx: Context, param: Parameter, incomplete: str - ) -> list[CompletionItem]: - """Return a special completion marker that tells the completion - system to use the shell to provide path completions for only - directories or any paths. - - :param ctx: Invocation context for this command. - :param param: The parameter that is requesting completion. - :param incomplete: Value being completed. May be empty. - - .. versionadded:: 8.0 - """ - from click.shell_completion import CompletionItem - - type = "dir" if self.dir_okay and not self.file_okay else "file" - return [CompletionItem(incomplete, type=type)] - - -class Tuple(CompositeParamType): - """The default behavior of Click is to apply a type on a value directly. - This works well in most cases, except for when `nargs` is set to a fixed - count and different types should be used for different items. In this - case the :class:`Tuple` type can be used. This type can only be used - if `nargs` is set to a fixed number. - - For more information see :ref:`tuple-type`. - - This can be selected by using a Python tuple literal as a type. - - :param types: a list of types that should be used for the tuple items. - """ - - def __init__(self, types: cabc.Sequence[type[t.Any] | ParamType]) -> None: - self.types: cabc.Sequence[ParamType] = [convert_type(ty) for ty in types] - - def to_info_dict(self) -> dict[str, t.Any]: - info_dict = super().to_info_dict() - info_dict["types"] = [t.to_info_dict() for t in self.types] - return info_dict - - @property - def name(self) -> str: # type: ignore - return f"<{' '.join(ty.name for ty in self.types)}>" - - @property - def arity(self) -> int: # type: ignore - return len(self.types) - - def convert( - self, value: t.Any, param: Parameter | None, ctx: Context | None - ) -> t.Any: - len_type = len(self.types) - len_value = len(value) - - if len_value != len_type: - self.fail( - ngettext( - "{len_type} values are required, but {len_value} was given.", - "{len_type} values are required, but {len_value} were given.", - len_value, - ).format(len_type=len_type, len_value=len_value), - param=param, - ctx=ctx, - ) - - return tuple( - ty(x, param, ctx) for ty, x in zip(self.types, value, strict=False) - ) - - -def convert_type(ty: t.Any | None, default: t.Any | None = None) -> ParamType: - """Find the most appropriate :class:`ParamType` for the given Python - type. If the type isn't provided, it can be inferred from a default - value. - """ - guessed_type = False - - if ty is None and default is not None: - if isinstance(default, (tuple, list)): - # If the default is empty, ty will remain None and will - # return STRING. - if default: - item = default[0] - - # A tuple of tuples needs to detect the inner types. - # Can't call convert recursively because that would - # incorrectly unwind the tuple to a single type. - if isinstance(item, (tuple, list)): - ty = tuple(map(type, item)) - else: - ty = type(item) - else: - ty = type(default) - - guessed_type = True - - if isinstance(ty, tuple): - return Tuple(ty) - - if isinstance(ty, ParamType): - return ty - - if ty is str or ty is None: - return STRING - - if ty is int: - return INT - - if ty is float: - return FLOAT - - if ty is bool: - return BOOL - - if guessed_type: - return STRING - - if __debug__: - try: - if issubclass(ty, ParamType): - raise AssertionError( - f"Attempted to use an uninstantiated parameter type ({ty})." - ) - except TypeError: - # ty is an instance (correct), so issubclass fails. - pass - - return FuncParamType(ty) - - -#: A dummy parameter type that just does nothing. From a user's -#: perspective this appears to just be the same as `STRING` but -#: internally no string conversion takes place if the input was bytes. -#: This is usually useful when working with file paths as they can -#: appear in bytes and unicode. -#: -#: For path related uses the :class:`Path` type is a better choice but -#: there are situations where an unprocessed type is useful which is why -#: it is is provided. -#: -#: .. versionadded:: 4.0 -UNPROCESSED = UnprocessedParamType() - -#: A unicode string parameter type which is the implicit default. This -#: can also be selected by using ``str`` as type. -STRING = StringParamType() - -#: An integer parameter. This can also be selected by using ``int`` as -#: type. -INT = IntParamType() - -#: A floating point value parameter. This can also be selected by using -#: ``float`` as type. -FLOAT = FloatParamType() - -#: A boolean parameter. This is the default for boolean flags. This can -#: also be selected by using ``bool`` as a type. -BOOL = BoolParamType() - -#: A UUID parameter. -UUID = UUIDParameterType() - - -class OptionHelpExtra(t.TypedDict, total=False): - envvars: tuple[str, ...] - default: str - range: str - required: str diff --git a/venv/lib/python3.10/site-packages/click/utils.py b/venv/lib/python3.10/site-packages/click/utils.py deleted file mode 100644 index ab2fe5889cee75d559092442696f9b23bf20ebfb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/click/utils.py +++ /dev/null @@ -1,627 +0,0 @@ -from __future__ import annotations - -import collections.abc as cabc -import os -import re -import sys -import typing as t -from functools import update_wrapper -from types import ModuleType -from types import TracebackType - -from ._compat import _default_text_stderr -from ._compat import _default_text_stdout -from ._compat import _find_binary_writer -from ._compat import auto_wrap_for_ansi -from ._compat import binary_streams -from ._compat import open_stream -from ._compat import should_strip_ansi -from ._compat import strip_ansi -from ._compat import text_streams -from ._compat import WIN -from .globals import resolve_color_default - -if t.TYPE_CHECKING: - import typing_extensions as te - - P = te.ParamSpec("P") - -R = t.TypeVar("R") - - -def _posixify(name: str) -> str: - return "-".join(name.split()).lower() - - -def safecall(func: t.Callable[P, R]) -> t.Callable[P, R | None]: - """Wraps a function so that it swallows exceptions.""" - - def wrapper(*args: P.args, **kwargs: P.kwargs) -> R | None: - try: - return func(*args, **kwargs) - except Exception: - pass - return None - - return update_wrapper(wrapper, func) - - -def make_str(value: t.Any) -> str: - """Converts a value into a valid string.""" - if isinstance(value, bytes): - try: - return value.decode(sys.getfilesystemencoding()) - except UnicodeError: - return value.decode("utf-8", "replace") - return str(value) - - -def make_default_short_help(help: str, max_length: int = 45) -> str: - """Returns a condensed version of help string.""" - # Consider only the first paragraph. - paragraph_end = help.find("\n\n") - - if paragraph_end != -1: - help = help[:paragraph_end] - - # Collapse newlines, tabs, and spaces. - words = help.split() - - if not words: - return "" - - # The first paragraph started with a "no rewrap" marker, ignore it. - if words[0] == "\b": - words = words[1:] - - total_length = 0 - last_index = len(words) - 1 - - for i, word in enumerate(words): - total_length += len(word) + (i > 0) - - if total_length > max_length: # too long, truncate - break - - if word[-1] == ".": # sentence end, truncate without "..." - return " ".join(words[: i + 1]) - - if total_length == max_length and i != last_index: - break # not at sentence end, truncate with "..." - else: - return " ".join(words) # no truncation needed - - # Account for the length of the suffix. - total_length += len("...") - - # remove words until the length is short enough - while i > 0: - total_length -= len(words[i]) + (i > 0) - - if total_length <= max_length: - break - - i -= 1 - - return " ".join(words[:i]) + "..." - - -class LazyFile: - """A lazy file works like a regular file but it does not fully open - the file but it does perform some basic checks early to see if the - filename parameter does make sense. This is useful for safely opening - files for writing. - """ - - def __init__( - self, - filename: str | os.PathLike[str], - mode: str = "r", - encoding: str | None = None, - errors: str | None = "strict", - atomic: bool = False, - ): - self.name: str = os.fspath(filename) - self.mode = mode - self.encoding = encoding - self.errors = errors - self.atomic = atomic - self._f: t.IO[t.Any] | None - self.should_close: bool - - if self.name == "-": - self._f, self.should_close = open_stream(filename, mode, encoding, errors) - else: - if "r" in mode: - # Open and close the file in case we're opening it for - # reading so that we can catch at least some errors in - # some cases early. - open(filename, mode).close() - self._f = None - self.should_close = True - - def __getattr__(self, name: str) -> t.Any: - return getattr(self.open(), name) - - def __repr__(self) -> str: - if self._f is not None: - return repr(self._f) - return f"" - - def open(self) -> t.IO[t.Any]: - """Opens the file if it's not yet open. This call might fail with - a :exc:`FileError`. Not handling this error will produce an error - that Click shows. - """ - if self._f is not None: - return self._f - try: - rv, self.should_close = open_stream( - self.name, self.mode, self.encoding, self.errors, atomic=self.atomic - ) - except OSError as e: - from .exceptions import FileError - - raise FileError(self.name, hint=e.strerror) from e - self._f = rv - return rv - - def close(self) -> None: - """Closes the underlying file, no matter what.""" - if self._f is not None: - self._f.close() - - def close_intelligently(self) -> None: - """This function only closes the file if it was opened by the lazy - file wrapper. For instance this will never close stdin. - """ - if self.should_close: - self.close() - - def __enter__(self) -> LazyFile: - return self - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc_value: BaseException | None, - tb: TracebackType | None, - ) -> None: - self.close_intelligently() - - def __iter__(self) -> cabc.Iterator[t.AnyStr]: - self.open() - return iter(self._f) # type: ignore - - -class KeepOpenFile: - def __init__(self, file: t.IO[t.Any]) -> None: - self._file: t.IO[t.Any] = file - - def __getattr__(self, name: str) -> t.Any: - return getattr(self._file, name) - - def __enter__(self) -> KeepOpenFile: - return self - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc_value: BaseException | None, - tb: TracebackType | None, - ) -> None: - pass - - def __repr__(self) -> str: - return repr(self._file) - - def __iter__(self) -> cabc.Iterator[t.AnyStr]: - return iter(self._file) - - -def echo( - message: t.Any | None = None, - file: t.IO[t.Any] | None = None, - nl: bool = True, - err: bool = False, - color: bool | None = None, -) -> None: - """Print a message and newline to stdout or a file. This should be - used instead of :func:`print` because it provides better support - for different data, files, and environments. - - Compared to :func:`print`, this does the following: - - - Ensures that the output encoding is not misconfigured on Linux. - - Supports Unicode in the Windows console. - - Supports writing to binary outputs, and supports writing bytes - to text outputs. - - Supports colors and styles on Windows. - - Removes ANSI color and style codes if the output does not look - like an interactive terminal. - - Always flushes the output. - - :param message: The string or bytes to output. Other objects are - converted to strings. - :param file: The file to write to. Defaults to ``stdout``. - :param err: Write to ``stderr`` instead of ``stdout``. - :param nl: Print a newline after the message. Enabled by default. - :param color: Force showing or hiding colors and other styles. By - default Click will remove color if the output does not look like - an interactive terminal. - - .. versionchanged:: 6.0 - Support Unicode output on the Windows console. Click does not - modify ``sys.stdout``, so ``sys.stdout.write()`` and ``print()`` - will still not support Unicode. - - .. versionchanged:: 4.0 - Added the ``color`` parameter. - - .. versionadded:: 3.0 - Added the ``err`` parameter. - - .. versionchanged:: 2.0 - Support colors on Windows if colorama is installed. - """ - if file is None: - if err: - file = _default_text_stderr() - else: - file = _default_text_stdout() - - # There are no standard streams attached to write to. For example, - # pythonw on Windows. - if file is None: - return - - # Convert non bytes/text into the native string type. - if message is not None and not isinstance(message, (str, bytes, bytearray)): - out: str | bytes | None = str(message) - else: - out = message - - if nl: - out = out or "" - if isinstance(out, str): - out += "\n" - else: - out += b"\n" - - if not out: - file.flush() - return - - # If there is a message and the value looks like bytes, we manually - # need to find the binary stream and write the message in there. - # This is done separately so that most stream types will work as you - # would expect. Eg: you can write to StringIO for other cases. - if isinstance(out, (bytes, bytearray)): - binary_file = _find_binary_writer(file) - - if binary_file is not None: - file.flush() - binary_file.write(out) - binary_file.flush() - return - - # ANSI style code support. For no message or bytes, nothing happens. - # When outputting to a file instead of a terminal, strip codes. - else: - color = resolve_color_default(color) - - if should_strip_ansi(file, color): - out = strip_ansi(out) - elif WIN: - if auto_wrap_for_ansi is not None: - file = auto_wrap_for_ansi(file, color) # type: ignore - elif not color: - out = strip_ansi(out) - - file.write(out) # type: ignore - file.flush() - - -def get_binary_stream(name: t.Literal["stdin", "stdout", "stderr"]) -> t.BinaryIO: - """Returns a system stream for byte processing. - - :param name: the name of the stream to open. Valid names are ``'stdin'``, - ``'stdout'`` and ``'stderr'`` - """ - opener = binary_streams.get(name) - if opener is None: - raise TypeError(f"Unknown standard stream '{name}'") - return opener() - - -def get_text_stream( - name: t.Literal["stdin", "stdout", "stderr"], - encoding: str | None = None, - errors: str | None = "strict", -) -> t.TextIO: - """Returns a system stream for text processing. This usually returns - a wrapped stream around a binary stream returned from - :func:`get_binary_stream` but it also can take shortcuts for already - correctly configured streams. - - :param name: the name of the stream to open. Valid names are ``'stdin'``, - ``'stdout'`` and ``'stderr'`` - :param encoding: overrides the detected default encoding. - :param errors: overrides the default error mode. - """ - opener = text_streams.get(name) - if opener is None: - raise TypeError(f"Unknown standard stream '{name}'") - return opener(encoding, errors) - - -def open_file( - filename: str | os.PathLike[str], - mode: str = "r", - encoding: str | None = None, - errors: str | None = "strict", - lazy: bool = False, - atomic: bool = False, -) -> t.IO[t.Any]: - """Open a file, with extra behavior to handle ``'-'`` to indicate - a standard stream, lazy open on write, and atomic write. Similar to - the behavior of the :class:`~click.File` param type. - - If ``'-'`` is given to open ``stdout`` or ``stdin``, the stream is - wrapped so that using it in a context manager will not close it. - This makes it possible to use the function without accidentally - closing a standard stream: - - .. code-block:: python - - with open_file(filename) as f: - ... - - :param filename: The name or Path of the file to open, or ``'-'`` for - ``stdin``/``stdout``. - :param mode: The mode in which to open the file. - :param encoding: The encoding to decode or encode a file opened in - text mode. - :param errors: The error handling mode. - :param lazy: Wait to open the file until it is accessed. For read - mode, the file is temporarily opened to raise access errors - early, then closed until it is read again. - :param atomic: Write to a temporary file and replace the given file - on close. - - .. versionadded:: 3.0 - """ - if lazy: - return t.cast( - "t.IO[t.Any]", LazyFile(filename, mode, encoding, errors, atomic=atomic) - ) - - f, should_close = open_stream(filename, mode, encoding, errors, atomic=atomic) - - if not should_close: - f = t.cast("t.IO[t.Any]", KeepOpenFile(f)) - - return f - - -def format_filename( - filename: str | bytes | os.PathLike[str] | os.PathLike[bytes], - shorten: bool = False, -) -> str: - """Format a filename as a string for display. Ensures the filename can be - displayed by replacing any invalid bytes or surrogate escapes in the name - with the replacement character ``�``. - - Invalid bytes or surrogate escapes will raise an error when written to a - stream with ``errors="strict"``. This will typically happen with ``stdout`` - when the locale is something like ``en_GB.UTF-8``. - - Many scenarios *are* safe to write surrogates though, due to PEP 538 and - PEP 540, including: - - - Writing to ``stderr``, which uses ``errors="backslashreplace"``. - - The system has ``LANG=C.UTF-8``, ``C``, or ``POSIX``. Python opens - stdout and stderr with ``errors="surrogateescape"``. - - None of ``LANG/LC_*`` are set. Python assumes ``LANG=C.UTF-8``. - - Python is started in UTF-8 mode with ``PYTHONUTF8=1`` or ``-X utf8``. - Python opens stdout and stderr with ``errors="surrogateescape"``. - - :param filename: formats a filename for UI display. This will also convert - the filename into unicode without failing. - :param shorten: this optionally shortens the filename to strip of the - path that leads up to it. - """ - if shorten: - filename = os.path.basename(filename) - else: - filename = os.fspath(filename) - - if isinstance(filename, bytes): - filename = filename.decode(sys.getfilesystemencoding(), "replace") - else: - filename = filename.encode("utf-8", "surrogateescape").decode( - "utf-8", "replace" - ) - - return filename - - -def get_app_dir(app_name: str, roaming: bool = True, force_posix: bool = False) -> str: - r"""Returns the config folder for the application. The default behavior - is to return whatever is most appropriate for the operating system. - - To give you an idea, for an app called ``"Foo Bar"``, something like - the following folders could be returned: - - Mac OS X: - ``~/Library/Application Support/Foo Bar`` - Mac OS X (POSIX): - ``~/.foo-bar`` - Unix: - ``~/.config/foo-bar`` - Unix (POSIX): - ``~/.foo-bar`` - Windows (roaming): - ``C:\Users\\AppData\Roaming\Foo Bar`` - Windows (not roaming): - ``C:\Users\\AppData\Local\Foo Bar`` - - .. versionadded:: 2.0 - - :param app_name: the application name. This should be properly capitalized - and can contain whitespace. - :param roaming: controls if the folder should be roaming or not on Windows. - Has no effect otherwise. - :param force_posix: if this is set to `True` then on any POSIX system the - folder will be stored in the home folder with a leading - dot instead of the XDG config home or darwin's - application support folder. - """ - if WIN: - key = "APPDATA" if roaming else "LOCALAPPDATA" - folder = os.environ.get(key) - if folder is None: - folder = os.path.expanduser("~") - return os.path.join(folder, app_name) - if force_posix: - return os.path.join(os.path.expanduser(f"~/.{_posixify(app_name)}")) - if sys.platform == "darwin": - return os.path.join( - os.path.expanduser("~/Library/Application Support"), app_name - ) - return os.path.join( - os.environ.get("XDG_CONFIG_HOME", os.path.expanduser("~/.config")), - _posixify(app_name), - ) - - -class PacifyFlushWrapper: - """This wrapper is used to catch and suppress BrokenPipeErrors resulting - from ``.flush()`` being called on broken pipe during the shutdown/final-GC - of the Python interpreter. Notably ``.flush()`` is always called on - ``sys.stdout`` and ``sys.stderr``. So as to have minimal impact on any - other cleanup code, and the case where the underlying file is not a broken - pipe, all calls and attributes are proxied. - """ - - def __init__(self, wrapped: t.IO[t.Any]) -> None: - self.wrapped = wrapped - - def flush(self) -> None: - try: - self.wrapped.flush() - except OSError as e: - import errno - - if e.errno != errno.EPIPE: - raise - - def __getattr__(self, attr: str) -> t.Any: - return getattr(self.wrapped, attr) - - -def _detect_program_name( - path: str | None = None, _main: ModuleType | None = None -) -> str: - """Determine the command used to run the program, for use in help - text. If a file or entry point was executed, the file name is - returned. If ``python -m`` was used to execute a module or package, - ``python -m name`` is returned. - - This doesn't try to be too precise, the goal is to give a concise - name for help text. Files are only shown as their name without the - path. ``python`` is only shown for modules, and the full path to - ``sys.executable`` is not shown. - - :param path: The Python file being executed. Python puts this in - ``sys.argv[0]``, which is used by default. - :param _main: The ``__main__`` module. This should only be passed - during internal testing. - - .. versionadded:: 8.0 - Based on command args detection in the Werkzeug reloader. - - :meta private: - """ - if _main is None: - _main = sys.modules["__main__"] - - if not path: - path = sys.argv[0] - - # The value of __package__ indicates how Python was called. It may - # not exist if a setuptools script is installed as an egg. It may be - # set incorrectly for entry points created with pip on Windows. - # It is set to "" inside a Shiv or PEX zipapp. - if getattr(_main, "__package__", None) in {None, ""} or ( - os.name == "nt" - and _main.__package__ == "" - and not os.path.exists(path) - and os.path.exists(f"{path}.exe") - ): - # Executed a file, like "python app.py". - return os.path.basename(path) - - # Executed a module, like "python -m example". - # Rewritten by Python from "-m script" to "/path/to/script.py". - # Need to look at main module to determine how it was executed. - py_module = t.cast(str, _main.__package__) - name = os.path.splitext(os.path.basename(path))[0] - - # A submodule like "example.cli". - if name != "__main__": - py_module = f"{py_module}.{name}" - - return f"python -m {py_module.lstrip('.')}" - - -def _expand_args( - args: cabc.Iterable[str], - *, - user: bool = True, - env: bool = True, - glob_recursive: bool = True, -) -> list[str]: - """Simulate Unix shell expansion with Python functions. - - See :func:`glob.glob`, :func:`os.path.expanduser`, and - :func:`os.path.expandvars`. - - This is intended for use on Windows, where the shell does not do any - expansion. It may not exactly match what a Unix shell would do. - - :param args: List of command line arguments to expand. - :param user: Expand user home directory. - :param env: Expand environment variables. - :param glob_recursive: ``**`` matches directories recursively. - - .. versionchanged:: 8.1 - Invalid glob patterns are treated as empty expansions rather - than raising an error. - - .. versionadded:: 8.0 - - :meta private: - """ - from glob import glob - - out = [] - - for arg in args: - if user: - arg = os.path.expanduser(arg) - - if env: - arg = os.path.expandvars(arg) - - try: - matches = glob(arg, recursive=glob_recursive) - except re.error: - matches = [] - - if not matches: - out.append(arg) - else: - out.extend(matches) - - return out diff --git a/venv/lib/python3.10/site-packages/coloredlogs.pth b/venv/lib/python3.10/site-packages/coloredlogs.pth deleted file mode 100644 index 27ad5ca9eb2dbb5d758cf6bdcf78a905eb6a1b7b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/coloredlogs.pth +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:dda83a855986efa5cd87f0248b0199c0086eb0e8e7fece7d6741959c5ce39536 -size 147 diff --git a/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/METADATA b/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/METADATA deleted file mode 100644 index 8f69ac5cbfc5ab4488fadd1f14a972fc593272e8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/METADATA +++ /dev/null @@ -1,139 +0,0 @@ -Metadata-Version: 2.4 -Name: cryptography -Version: 45.0.5 -Classifier: Development Status :: 5 - Production/Stable -Classifier: Intended Audience :: Developers -Classifier: Natural Language :: English -Classifier: Operating System :: MacOS :: MacOS X -Classifier: Operating System :: POSIX -Classifier: Operating System :: POSIX :: BSD -Classifier: Operating System :: POSIX :: Linux -Classifier: Operating System :: Microsoft :: Windows -Classifier: Programming Language :: Python -Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3 :: Only -Classifier: Programming Language :: Python :: 3.7 -Classifier: Programming Language :: Python :: 3.8 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Programming Language :: Python :: 3.10 -Classifier: Programming Language :: Python :: 3.11 -Classifier: Programming Language :: Python :: 3.12 -Classifier: Programming Language :: Python :: 3.13 -Classifier: Programming Language :: Python :: Implementation :: CPython -Classifier: Programming Language :: Python :: Implementation :: PyPy -Classifier: Topic :: Security :: Cryptography -Requires-Dist: cffi>=1.14 ; platform_python_implementation != 'PyPy' -Requires-Dist: bcrypt>=3.1.5 ; extra == 'ssh' -Requires-Dist: nox>=2024.4.15 ; extra == 'nox' -Requires-Dist: nox[uv]>=2024.3.2 ; python_full_version >= '3.8' and extra == 'nox' -Requires-Dist: cryptography-vectors==45.0.5 ; extra == 'test' -Requires-Dist: pytest>=7.4.0 ; extra == 'test' -Requires-Dist: pytest-benchmark>=4.0 ; extra == 'test' -Requires-Dist: pytest-cov>=2.10.1 ; extra == 'test' -Requires-Dist: pytest-xdist>=3.5.0 ; extra == 'test' -Requires-Dist: pretend>=0.7 ; extra == 'test' -Requires-Dist: certifi>=2024 ; extra == 'test' -Requires-Dist: pytest-randomly ; extra == 'test-randomorder' -Requires-Dist: sphinx>=5.3.0 ; extra == 'docs' -Requires-Dist: sphinx-rtd-theme>=3.0.0 ; python_full_version >= '3.8' and extra == 'docs' -Requires-Dist: sphinx-inline-tabs ; python_full_version >= '3.8' and extra == 'docs' -Requires-Dist: pyenchant>=3 ; extra == 'docstest' -Requires-Dist: readme-renderer>=30.0 ; extra == 'docstest' -Requires-Dist: sphinxcontrib-spelling>=7.3.1 ; extra == 'docstest' -Requires-Dist: build>=1.0.0 ; extra == 'sdist' -Requires-Dist: ruff>=0.3.6 ; extra == 'pep8test' -Requires-Dist: mypy>=1.4 ; extra == 'pep8test' -Requires-Dist: check-sdist ; python_full_version >= '3.8' and extra == 'pep8test' -Requires-Dist: click>=8.0.1 ; extra == 'pep8test' -Provides-Extra: ssh -Provides-Extra: nox -Provides-Extra: test -Provides-Extra: test-randomorder -Provides-Extra: docs -Provides-Extra: docstest -Provides-Extra: sdist -Provides-Extra: pep8test -License-File: LICENSE -License-File: LICENSE.APACHE -License-File: LICENSE.BSD -Summary: cryptography is a package which provides cryptographic recipes and primitives to Python developers. -Author: The cryptography developers -Author-email: The Python Cryptographic Authority and individual contributors -License: Apache-2.0 OR BSD-3-Clause -Requires-Python: >=3.7, !=3.9.0, !=3.9.1 -Description-Content-Type: text/x-rst; charset=UTF-8 -Project-URL: homepage, https://github.com/pyca/cryptography -Project-URL: documentation, https://cryptography.io/ -Project-URL: source, https://github.com/pyca/cryptography/ -Project-URL: issues, https://github.com/pyca/cryptography/issues -Project-URL: changelog, https://cryptography.io/en/latest/changelog/ - -pyca/cryptography -================= - -.. image:: https://img.shields.io/pypi/v/cryptography.svg - :target: https://pypi.org/project/cryptography/ - :alt: Latest Version - -.. image:: https://readthedocs.org/projects/cryptography/badge/?version=latest - :target: https://cryptography.io - :alt: Latest Docs - -.. image:: https://github.com/pyca/cryptography/workflows/CI/badge.svg?branch=main - :target: https://github.com/pyca/cryptography/actions?query=workflow%3ACI+branch%3Amain - - -``cryptography`` is a package which provides cryptographic recipes and -primitives to Python developers. Our goal is for it to be your "cryptographic -standard library". It supports Python 3.7+ and PyPy3 7.3.11+. - -``cryptography`` includes both high level recipes and low level interfaces to -common cryptographic algorithms such as symmetric ciphers, message digests, and -key derivation functions. For example, to encrypt something with -``cryptography``'s high level symmetric encryption recipe: - -.. code-block:: pycon - - >>> from cryptography.fernet import Fernet - >>> # Put this somewhere safe! - >>> key = Fernet.generate_key() - >>> f = Fernet(key) - >>> token = f.encrypt(b"A really secret message. Not for prying eyes.") - >>> token - b'...' - >>> f.decrypt(token) - b'A really secret message. Not for prying eyes.' - -You can find more information in the `documentation`_. - -You can install ``cryptography`` with: - -.. code-block:: console - - $ pip install cryptography - -For full details see `the installation documentation`_. - -Discussion -~~~~~~~~~~ - -If you run into bugs, you can file them in our `issue tracker`_. - -We maintain a `cryptography-dev`_ mailing list for development discussion. - -You can also join ``#pyca`` on ``irc.libera.chat`` to ask questions or get -involved. - -Security -~~~~~~~~ - -Need to report a security issue? Please consult our `security reporting`_ -documentation. - - -.. _`documentation`: https://cryptography.io/ -.. _`the installation documentation`: https://cryptography.io/en/latest/installation/ -.. _`issue tracker`: https://github.com/pyca/cryptography/issues -.. _`cryptography-dev`: https://mail.python.org/mailman/listinfo/cryptography-dev -.. _`security reporting`: https://cryptography.io/en/latest/security/ - diff --git a/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/RECORD b/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/RECORD deleted file mode 100644 index 45e51fc85f515db8127c2254faeebffc4633b702..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/RECORD +++ /dev/null @@ -1,175 +0,0 @@ -cryptography-45.0.5.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -cryptography-45.0.5.dist-info/METADATA,sha256=Dn2Ls7PhdM88VKgRTBFDskPTSNeGHm7SezAcdIPYdKc,5690 -cryptography-45.0.5.dist-info/RECORD,, -cryptography-45.0.5.dist-info/WHEEL,sha256=VqNc6WOTGHZndw6FKXcyBqGvpOjDTrsbnnH79sUpcvY,107 -cryptography-45.0.5.dist-info/licenses/LICENSE,sha256=Pgx8CRqUi4JTO6mP18u0BDLW8amsv4X1ki0vmak65rs,197 -cryptography-45.0.5.dist-info/licenses/LICENSE.APACHE,sha256=qsc7MUj20dcRHbyjIJn2jSbGRMaBOuHk8F9leaomY_4,11360 -cryptography-45.0.5.dist-info/licenses/LICENSE.BSD,sha256=YCxMdILeZHndLpeTzaJ15eY9dz2s0eymiSMqtwCPtPs,1532 -cryptography/__about__.py,sha256=X8cr9wNd4PDeKHYUt6ppLOYb-V56WYSgLBuYKlivuBY,445 -cryptography/__init__.py,sha256=QDuQ3te_14R9iGF-2Q5yEwD1FTJrCWqpnuhlMIJBsno,762 -cryptography/__pycache__/__about__.cpython-310.pyc,, -cryptography/__pycache__/__init__.cpython-310.pyc,, -cryptography/__pycache__/exceptions.cpython-310.pyc,, -cryptography/__pycache__/fernet.cpython-310.pyc,, -cryptography/__pycache__/utils.cpython-310.pyc,, -cryptography/exceptions.py,sha256=835EWILc2fwxw-gyFMriciC2SqhViETB10LBSytnDIc,1087 -cryptography/fernet.py,sha256=3Cvxkh0KJSbX8HbnCHu4wfCW7U0GgfUA3v_qQ8a8iWc,6963 -cryptography/hazmat/__init__.py,sha256=5IwrLWrVp0AjEr_4FdWG_V057NSJGY_W4egNNsuct0g,455 -cryptography/hazmat/__pycache__/__init__.cpython-310.pyc,, -cryptography/hazmat/__pycache__/_oid.cpython-310.pyc,, -cryptography/hazmat/_oid.py,sha256=7-aiQLWCYT7FDrQ0sUWBNYkN5ADEGTZwLQE_ZCksPxs,16795 -cryptography/hazmat/backends/__init__.py,sha256=O5jvKFQdZnXhKeqJ-HtulaEL9Ni7mr1mDzZY5kHlYhI,361 -cryptography/hazmat/backends/__pycache__/__init__.cpython-310.pyc,, -cryptography/hazmat/backends/openssl/__init__.py,sha256=p3jmJfnCag9iE5sdMrN6VvVEu55u46xaS_IjoI0SrmA,305 -cryptography/hazmat/backends/openssl/__pycache__/__init__.cpython-310.pyc,, -cryptography/hazmat/backends/openssl/__pycache__/backend.cpython-310.pyc,, -cryptography/hazmat/backends/openssl/backend.py,sha256=ev-Prba_u1bCMk3ZWspkdOuLBbDlf_7wJnHK6xdYs7c,10329 -cryptography/hazmat/bindings/__init__.py,sha256=s9oKCQ2ycFdXoERdS1imafueSkBsL9kvbyfghaauZ9Y,180 -cryptography/hazmat/bindings/__pycache__/__init__.cpython-310.pyc,, -cryptography/hazmat/bindings/_rust.abi3.so,sha256=BtscyLV0h1_Wy8zPS9g0OAG4KOJsnqLVG1w4ZqZnr5s,12446840 -cryptography/hazmat/bindings/_rust/__init__.pyi,sha256=KhqLhXFPArPzzJ7DYO9Fl8FoXB_BagAd_r4Dm_Ze9Xo,1257 -cryptography/hazmat/bindings/_rust/_openssl.pyi,sha256=mpNJLuYLbCVrd5i33FBTmWwL_55Dw7JPkSLlSX9Q7oI,230 -cryptography/hazmat/bindings/_rust/asn1.pyi,sha256=BrGjC8J6nwuS-r3EVcdXJB8ndotfY9mbQYOfpbPG0HA,354 -cryptography/hazmat/bindings/_rust/exceptions.pyi,sha256=exXr2xw_0pB1kk93cYbM3MohbzoUkjOms1ZMUi0uQZE,640 -cryptography/hazmat/bindings/_rust/ocsp.pyi,sha256=VPVWuKHI9EMs09ZLRYAGvR0Iz0mCMmEzXAkgJHovpoM,4020 -cryptography/hazmat/bindings/_rust/openssl/__init__.pyi,sha256=iOAMDyHoNwwCSZfZzuXDr64g4GpGUeDgEN-LjXqdrBM,1522 -cryptography/hazmat/bindings/_rust/openssl/aead.pyi,sha256=4Nddw6-ynzIB3w2W86WvkGKTLlTDk_6F5l54RHCuy3E,2688 -cryptography/hazmat/bindings/_rust/openssl/ciphers.pyi,sha256=LhPzHWSXJq4grAJXn6zSvSSdV-aYIIscHDwIPlJGGPs,1315 -cryptography/hazmat/bindings/_rust/openssl/cmac.pyi,sha256=nPH0X57RYpsAkRowVpjQiHE566ThUTx7YXrsadmrmHk,564 -cryptography/hazmat/bindings/_rust/openssl/dh.pyi,sha256=Z3TC-G04-THtSdAOPLM1h2G7ml5bda1ElZUcn5wpuhk,1564 -cryptography/hazmat/bindings/_rust/openssl/dsa.pyi,sha256=qBtkgj2albt2qFcnZ9UDrhzoNhCVO7HTby5VSf1EXMI,1299 -cryptography/hazmat/bindings/_rust/openssl/ec.pyi,sha256=zJy0pRa5n-_p2dm45PxECB_-B6SVZyNKfjxFDpPqT38,1691 -cryptography/hazmat/bindings/_rust/openssl/ed25519.pyi,sha256=VXfXd5G6hUivg399R1DYdmW3eTb0EebzDTqjRC2gaRw,532 -cryptography/hazmat/bindings/_rust/openssl/ed448.pyi,sha256=Yx49lqdnjsD7bxiDV1kcaMrDktug5evi5a6zerMiy2s,514 -cryptography/hazmat/bindings/_rust/openssl/hashes.pyi,sha256=OWZvBx7xfo_HJl41Nc--DugVyCVPIprZ3HlOPTSWH9g,984 -cryptography/hazmat/bindings/_rust/openssl/hmac.pyi,sha256=BXZn7NDjL3JAbYW0SQ8pg1iyC5DbQXVhUAiwsi8DFR8,702 -cryptography/hazmat/bindings/_rust/openssl/kdf.pyi,sha256=4FGpmCR2H8SVGoe4SUX7h5wzhyvuRi8hnzKKfvnJJrw,1379 -cryptography/hazmat/bindings/_rust/openssl/keys.pyi,sha256=teIt8M6ZEMJrn4s3W0UnW0DZ-30Jd68WnSsKKG124l0,912 -cryptography/hazmat/bindings/_rust/openssl/poly1305.pyi,sha256=_SW9NtQ5FDlAbdclFtWpT4lGmxKIKHpN-4j8J2BzYfQ,585 -cryptography/hazmat/bindings/_rust/openssl/rsa.pyi,sha256=2OQCNSXkxgc-3uw1xiCCloIQTV6p9_kK79Yu0rhZgPc,1364 -cryptography/hazmat/bindings/_rust/openssl/x25519.pyi,sha256=ewn4GpQyb7zPwE-ni7GtyQgMC0A1mLuqYsSyqv6nI_s,523 -cryptography/hazmat/bindings/_rust/openssl/x448.pyi,sha256=juTZTmli8jO_5Vcufg-vHvx_tCyezmSLIh_9PU3TczI,505 -cryptography/hazmat/bindings/_rust/pkcs12.pyi,sha256=vEEd5wDiZvb8ZGFaziLCaWLzAwoG_tvPUxLQw5_uOl8,1605 -cryptography/hazmat/bindings/_rust/pkcs7.pyi,sha256=txGBJijqZshEcqra6byPNbnisIdlxzOSIHP2hl9arPs,1601 -cryptography/hazmat/bindings/_rust/test_support.pyi,sha256=PPhld-WkO743iXFPebeG0LtgK0aTzGdjcIsay1Gm5GE,757 -cryptography/hazmat/bindings/_rust/x509.pyi,sha256=WWFr7RV3n1eFYbLpmM_xwXGKYZalyYi_V07p6KLOHX8,10189 -cryptography/hazmat/bindings/openssl/__init__.py,sha256=s9oKCQ2ycFdXoERdS1imafueSkBsL9kvbyfghaauZ9Y,180 -cryptography/hazmat/bindings/openssl/__pycache__/__init__.cpython-310.pyc,, -cryptography/hazmat/bindings/openssl/__pycache__/_conditional.cpython-310.pyc,, -cryptography/hazmat/bindings/openssl/__pycache__/binding.cpython-310.pyc,, -cryptography/hazmat/bindings/openssl/_conditional.py,sha256=eKiDLdRSXTGA6mQ8jj5S2OXEB_sod6iL7cwed-yOuUw,5340 -cryptography/hazmat/bindings/openssl/binding.py,sha256=hAmRXYzQlpIAP_jj2tl-SIBJBxwtclLlr_3aIL-RkR0,4072 -cryptography/hazmat/decrepit/__init__.py,sha256=wHCbWfaefa-fk6THSw9th9fJUsStJo7245wfFBqmduA,216 -cryptography/hazmat/decrepit/__pycache__/__init__.cpython-310.pyc,, -cryptography/hazmat/decrepit/ciphers/__init__.py,sha256=wHCbWfaefa-fk6THSw9th9fJUsStJo7245wfFBqmduA,216 -cryptography/hazmat/decrepit/ciphers/__pycache__/__init__.cpython-310.pyc,, -cryptography/hazmat/decrepit/ciphers/__pycache__/algorithms.cpython-310.pyc,, -cryptography/hazmat/decrepit/ciphers/algorithms.py,sha256=YrKgHS4MfwWaMmPBYRymRRlC0phwWp9ycICFezeJPGk,2595 -cryptography/hazmat/primitives/__init__.py,sha256=s9oKCQ2ycFdXoERdS1imafueSkBsL9kvbyfghaauZ9Y,180 -cryptography/hazmat/primitives/__pycache__/__init__.cpython-310.pyc,, -cryptography/hazmat/primitives/__pycache__/_asymmetric.cpython-310.pyc,, -cryptography/hazmat/primitives/__pycache__/_cipheralgorithm.cpython-310.pyc,, -cryptography/hazmat/primitives/__pycache__/_serialization.cpython-310.pyc,, -cryptography/hazmat/primitives/__pycache__/cmac.cpython-310.pyc,, -cryptography/hazmat/primitives/__pycache__/constant_time.cpython-310.pyc,, -cryptography/hazmat/primitives/__pycache__/hashes.cpython-310.pyc,, -cryptography/hazmat/primitives/__pycache__/hmac.cpython-310.pyc,, -cryptography/hazmat/primitives/__pycache__/keywrap.cpython-310.pyc,, -cryptography/hazmat/primitives/__pycache__/padding.cpython-310.pyc,, -cryptography/hazmat/primitives/__pycache__/poly1305.cpython-310.pyc,, -cryptography/hazmat/primitives/_asymmetric.py,sha256=RhgcouUB6HTiFDBrR1LxqkMjpUxIiNvQ1r_zJjRG6qQ,532 -cryptography/hazmat/primitives/_cipheralgorithm.py,sha256=Eh3i7lwedHfi0eLSsH93PZxQKzY9I6lkK67vL4V5tOc,1522 -cryptography/hazmat/primitives/_serialization.py,sha256=chgPCSF2jxI2Cr5gB-qbWXOvOfupBh4CARS0KAhv9AM,5123 -cryptography/hazmat/primitives/asymmetric/__init__.py,sha256=s9oKCQ2ycFdXoERdS1imafueSkBsL9kvbyfghaauZ9Y,180 -cryptography/hazmat/primitives/asymmetric/__pycache__/__init__.cpython-310.pyc,, -cryptography/hazmat/primitives/asymmetric/__pycache__/dh.cpython-310.pyc,, -cryptography/hazmat/primitives/asymmetric/__pycache__/dsa.cpython-310.pyc,, -cryptography/hazmat/primitives/asymmetric/__pycache__/ec.cpython-310.pyc,, -cryptography/hazmat/primitives/asymmetric/__pycache__/ed25519.cpython-310.pyc,, -cryptography/hazmat/primitives/asymmetric/__pycache__/ed448.cpython-310.pyc,, -cryptography/hazmat/primitives/asymmetric/__pycache__/padding.cpython-310.pyc,, -cryptography/hazmat/primitives/asymmetric/__pycache__/rsa.cpython-310.pyc,, -cryptography/hazmat/primitives/asymmetric/__pycache__/types.cpython-310.pyc,, -cryptography/hazmat/primitives/asymmetric/__pycache__/utils.cpython-310.pyc,, -cryptography/hazmat/primitives/asymmetric/__pycache__/x25519.cpython-310.pyc,, -cryptography/hazmat/primitives/asymmetric/__pycache__/x448.cpython-310.pyc,, -cryptography/hazmat/primitives/asymmetric/dh.py,sha256=0v_vEFFz5pQ1QG-FkWDyvgv7IfuVZSH5Q6LyFI5A8rg,3645 -cryptography/hazmat/primitives/asymmetric/dsa.py,sha256=Ld_bbbqQFz12dObHxIkzEQzX0SWWP41RLSWkYSaKhqE,4213 -cryptography/hazmat/primitives/asymmetric/ec.py,sha256=Vf5ig2PcS3PVnsb5N49Kx1uIkFBJyhg4BWXThDz5cug,12999 -cryptography/hazmat/primitives/asymmetric/ed25519.py,sha256=jZW5cs472wXXV3eB0sE1b8w64gdazwwU0_MT5UOTiXs,3700 -cryptography/hazmat/primitives/asymmetric/ed448.py,sha256=yAetgn2f2JYf0BO8MapGzXeThsvSMG5LmUCrxVOidAA,3729 -cryptography/hazmat/primitives/asymmetric/padding.py,sha256=eZcvUqVLbe3u48SunLdeniaPlV4-k6pwBl67OW4jSy8,2885 -cryptography/hazmat/primitives/asymmetric/rsa.py,sha256=kegiZAGeb6yJISjpxpaAEpM3wGdfhJSU9RgZpVwKYwk,7967 -cryptography/hazmat/primitives/asymmetric/types.py,sha256=LnsOJym-wmPUJ7Knu_7bCNU3kIiELCd6krOaW_JU08I,2996 -cryptography/hazmat/primitives/asymmetric/utils.py,sha256=DPTs6T4F-UhwzFQTh-1fSEpQzazH2jf2xpIro3ItF4o,790 -cryptography/hazmat/primitives/asymmetric/x25519.py,sha256=_4nQeZ3yJ3Lg0RpXnaqA-1yt6vbx1F-wzLcaZHwSpeE,3613 -cryptography/hazmat/primitives/asymmetric/x448.py,sha256=WKBLtuVfJqiBRro654fGaQAlvsKbqbNkK7c4A_ZCdV0,3642 -cryptography/hazmat/primitives/ciphers/__init__.py,sha256=eyEXmjk6_CZXaOPYDr7vAYGXr29QvzgWL2-4CSolLFs,680 -cryptography/hazmat/primitives/ciphers/__pycache__/__init__.cpython-310.pyc,, -cryptography/hazmat/primitives/ciphers/__pycache__/aead.cpython-310.pyc,, -cryptography/hazmat/primitives/ciphers/__pycache__/algorithms.cpython-310.pyc,, -cryptography/hazmat/primitives/ciphers/__pycache__/base.cpython-310.pyc,, -cryptography/hazmat/primitives/ciphers/__pycache__/modes.cpython-310.pyc,, -cryptography/hazmat/primitives/ciphers/aead.py,sha256=Fzlyx7w8KYQakzDp1zWgJnIr62zgZrgVh1u2h4exB54,634 -cryptography/hazmat/primitives/ciphers/algorithms.py,sha256=S04j4NdFCBDt5KLIvbYohohRO--MCWhTia_lkVt8xD8,4542 -cryptography/hazmat/primitives/ciphers/base.py,sha256=aBC7HHBBoixebmparVr0UlODs3VD0A7B6oz_AaRjDv8,4253 -cryptography/hazmat/primitives/ciphers/modes.py,sha256=20stpwhDtbAvpH0SMf9EDHIciwmTF-JMBUOZ9bU8WiQ,8318 -cryptography/hazmat/primitives/cmac.py,sha256=sz_s6H_cYnOvx-VNWdIKhRhe3Ymp8z8J0D3CBqOX3gg,338 -cryptography/hazmat/primitives/constant_time.py,sha256=xdunWT0nf8OvKdcqUhhlFKayGp4_PgVJRU2W1wLSr_A,422 -cryptography/hazmat/primitives/hashes.py,sha256=M8BrlKB3U6DEtHvWTV5VRjpteHv1kS3Zxm_Bsk04cr8,5184 -cryptography/hazmat/primitives/hmac.py,sha256=RpB3z9z5skirCQrm7zQbtnp9pLMnAjrlTUvKqF5aDDc,423 -cryptography/hazmat/primitives/kdf/__init__.py,sha256=4XibZnrYq4hh5xBjWiIXzaYW6FKx8hPbVaa_cB9zS64,750 -cryptography/hazmat/primitives/kdf/__pycache__/__init__.cpython-310.pyc,, -cryptography/hazmat/primitives/kdf/__pycache__/argon2.cpython-310.pyc,, -cryptography/hazmat/primitives/kdf/__pycache__/concatkdf.cpython-310.pyc,, -cryptography/hazmat/primitives/kdf/__pycache__/hkdf.cpython-310.pyc,, -cryptography/hazmat/primitives/kdf/__pycache__/kbkdf.cpython-310.pyc,, -cryptography/hazmat/primitives/kdf/__pycache__/pbkdf2.cpython-310.pyc,, -cryptography/hazmat/primitives/kdf/__pycache__/scrypt.cpython-310.pyc,, -cryptography/hazmat/primitives/kdf/__pycache__/x963kdf.cpython-310.pyc,, -cryptography/hazmat/primitives/kdf/argon2.py,sha256=UFDNXG0v-rw3DqAQTB1UQAsQC2M5Ejg0k_6OCyhLKus,460 -cryptography/hazmat/primitives/kdf/concatkdf.py,sha256=Ua8KoLXXnzgsrAUmHpyKymaPt8aPRP0EHEaBz7QCQ9I,3737 -cryptography/hazmat/primitives/kdf/hkdf.py,sha256=2HlP_huUzGAy_FZWOn8L9qP9B5ke7XKmbiER7_2l__Q,3043 -cryptography/hazmat/primitives/kdf/kbkdf.py,sha256=tGH5dQGEVaKYMEbgeqlaSPpqhlDMoF0xwCxk2vuhcTc,9211 -cryptography/hazmat/primitives/kdf/pbkdf2.py,sha256=1WIwhELR0w8ztTpTu8BrFiYWmK3hUfJq08I79TxwieE,1957 -cryptography/hazmat/primitives/kdf/scrypt.py,sha256=XyWUdUUmhuI9V6TqAPOvujCSMGv1XQdg0a21IWCmO-U,590 -cryptography/hazmat/primitives/kdf/x963kdf.py,sha256=zLTcF665QFvXX2f8TS7fmBZTteXpFjKahzfjjQcCJyw,1999 -cryptography/hazmat/primitives/keywrap.py,sha256=XV4Pj2fqSeD-RqZVvY2cA3j5_7RwJSFygYuLfk2ujCo,5650 -cryptography/hazmat/primitives/padding.py,sha256=QT-U-NvV2eQGO1wVPbDiNGNSc9keRDS-ig5cQOrLz0E,1865 -cryptography/hazmat/primitives/poly1305.py,sha256=P5EPQV-RB_FJPahpg01u0Ts4S_PnAmsroxIGXbGeRRo,355 -cryptography/hazmat/primitives/serialization/__init__.py,sha256=Q7uTgDlt7n3WfsMT6jYwutC6DIg_7SEeoAm1GHZ5B5E,1705 -cryptography/hazmat/primitives/serialization/__pycache__/__init__.cpython-310.pyc,, -cryptography/hazmat/primitives/serialization/__pycache__/base.cpython-310.pyc,, -cryptography/hazmat/primitives/serialization/__pycache__/pkcs12.cpython-310.pyc,, -cryptography/hazmat/primitives/serialization/__pycache__/pkcs7.cpython-310.pyc,, -cryptography/hazmat/primitives/serialization/__pycache__/ssh.cpython-310.pyc,, -cryptography/hazmat/primitives/serialization/base.py,sha256=ikq5MJIwp_oUnjiaBco_PmQwOTYuGi-XkYUYHKy8Vo0,615 -cryptography/hazmat/primitives/serialization/pkcs12.py,sha256=mS9cFNG4afzvseoc5e1MWoY2VskfL8N8Y_OFjl67luY,5104 -cryptography/hazmat/primitives/serialization/pkcs7.py,sha256=5OR_Tkysxaprn4FegvJIfbep9rJ9wok6FLWvWwQ5-Mg,13943 -cryptography/hazmat/primitives/serialization/ssh.py,sha256=CeDSD_KmtcpQDSUg8QY3jruZDOKjjWBKcUByDOGQXnY,53693 -cryptography/hazmat/primitives/twofactor/__init__.py,sha256=tmMZGB-g4IU1r7lIFqASU019zr0uPp_wEBYcwdDCKCA,258 -cryptography/hazmat/primitives/twofactor/__pycache__/__init__.cpython-310.pyc,, -cryptography/hazmat/primitives/twofactor/__pycache__/hotp.cpython-310.pyc,, -cryptography/hazmat/primitives/twofactor/__pycache__/totp.cpython-310.pyc,, -cryptography/hazmat/primitives/twofactor/hotp.py,sha256=ivZo5BrcCGWLsqql4nZV0XXCjyGPi_iHfDFltGlOJwk,3256 -cryptography/hazmat/primitives/twofactor/totp.py,sha256=m5LPpRL00kp4zY8gTjr55Hfz9aMlPS53kHmVkSQCmdY,1652 -cryptography/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -cryptography/utils.py,sha256=r90VtamCjiba37omFqp2aa--eEQVXiQIMrvWeLT-BY4,4397 -cryptography/x509/__init__.py,sha256=xloN0swseNx-m2WFZmCA17gOoxQWqeU82UVjEdJBePQ,8257 -cryptography/x509/__pycache__/__init__.cpython-310.pyc,, -cryptography/x509/__pycache__/base.cpython-310.pyc,, -cryptography/x509/__pycache__/certificate_transparency.cpython-310.pyc,, -cryptography/x509/__pycache__/extensions.cpython-310.pyc,, -cryptography/x509/__pycache__/general_name.cpython-310.pyc,, -cryptography/x509/__pycache__/name.cpython-310.pyc,, -cryptography/x509/__pycache__/ocsp.cpython-310.pyc,, -cryptography/x509/__pycache__/oid.cpython-310.pyc,, -cryptography/x509/__pycache__/verification.cpython-310.pyc,, -cryptography/x509/base.py,sha256=OrmTw3y8B6AE_nGXQPN8x9kq-d7rDWeH13gCq6T6D6U,27997 -cryptography/x509/certificate_transparency.py,sha256=JqoOIDhlwInrYMFW6IFn77WJ0viF-PB_rlZV3vs9MYc,797 -cryptography/x509/extensions.py,sha256=QxYrqR6SF1qzR9ZraP8wDiIczlEVlAFuwDRVcltB6Tk,77724 -cryptography/x509/general_name.py,sha256=sP_rV11Qlpsk4x3XXGJY_Mv0Q_s9dtjeLckHsjpLQoQ,7836 -cryptography/x509/name.py,sha256=oi0CqY_B72CEJgNXy9XkJya26QWXfwOeGJBYfXjRYvI,15121 -cryptography/x509/ocsp.py,sha256=Yey6NdFV1MPjop24Mj_VenjEpg3kUaMopSWOK0AbeBs,12699 -cryptography/x509/oid.py,sha256=BUzgXXGVWilkBkdKPTm9R4qElE9gAGHgdYPMZAp7PJo,931 -cryptography/x509/verification.py,sha256=gR2C2c-XZQtblZhT5T5vjSKOtCb74ef2alPVmEcwFlM,958 diff --git a/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/WHEEL b/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/WHEEL deleted file mode 100644 index 81333fb421bd1f1b406398d8768d025d0a28970c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/WHEEL +++ /dev/null @@ -1,5 +0,0 @@ -Wheel-Version: 1.0 -Generator: maturin (1.8.6) -Root-Is-Purelib: false -Tag: cp37-abi3-manylinux_2_34_x86_64 - diff --git a/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/licenses/LICENSE b/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/licenses/LICENSE deleted file mode 100644 index b11f379efe1504d235b4d2d42685ba5dc6af6e9f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/licenses/LICENSE +++ /dev/null @@ -1,3 +0,0 @@ -This software is made available under the terms of *either* of the licenses -found in LICENSE.APACHE or LICENSE.BSD. Contributions to cryptography are made -under the terms of *both* these licenses. diff --git a/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/licenses/LICENSE.APACHE b/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/licenses/LICENSE.APACHE deleted file mode 100644 index 62589edd12a37dd28b6b6fed1e2d728ac9f05c8d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/licenses/LICENSE.APACHE +++ /dev/null @@ -1,202 +0,0 @@ - - Apache License - Version 2.0, January 2004 - https://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - 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 - 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 - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - 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 - - https://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. diff --git a/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/licenses/LICENSE.BSD b/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/licenses/LICENSE.BSD deleted file mode 100644 index ec1a29d34d6e419411c75523408aca72f705345c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography-45.0.5.dist-info/licenses/LICENSE.BSD +++ /dev/null @@ -1,27 +0,0 @@ -Copyright (c) Individual contributors. -All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: - - 1. Redistributions of source code must retain the above copyright notice, - this list of conditions and the following disclaimer. - - 2. Redistributions in binary form must reproduce the above copyright - notice, this list of conditions and the following disclaimer in the - documentation and/or other materials provided with the distribution. - - 3. Neither the name of PyCA Cryptography nor the names of its contributors - may be used to endorse or promote products derived from this software - without specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND -ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED -WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR -ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES -(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; -LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON -ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT -(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS -SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/venv/lib/python3.10/site-packages/cryptography/__about__.py b/venv/lib/python3.10/site-packages/cryptography/__about__.py deleted file mode 100644 index 97168d231b6276189033c9478336a5e502fe463e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/__about__.py +++ /dev/null @@ -1,17 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -__all__ = [ - "__author__", - "__copyright__", - "__version__", -] - -__version__ = "45.0.5" - - -__author__ = "The Python Cryptographic Authority and individual contributors" -__copyright__ = f"Copyright 2013-2025 {__author__}" diff --git a/venv/lib/python3.10/site-packages/cryptography/__init__.py b/venv/lib/python3.10/site-packages/cryptography/__init__.py deleted file mode 100644 index e8febfb8d554723371d1023daa23947224320c3f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import sys -import warnings - -from cryptography import utils -from cryptography.__about__ import __author__, __copyright__, __version__ - -__all__ = [ - "__author__", - "__copyright__", - "__version__", -] - -if sys.version_info[:2] == (3, 7): - warnings.warn( - "Python 3.7 is no longer supported by the Python core team " - "and support for it is deprecated in cryptography. The next release " - "of cryptography will remove support for Python 3.7.", - utils.CryptographyDeprecationWarning, - stacklevel=2, - ) diff --git a/venv/lib/python3.10/site-packages/cryptography/__pycache__/__about__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/__pycache__/__about__.cpython-310.pyc deleted file mode 100644 index 5417530d512d0b0768ba67dac6815f15ba42740f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/__pycache__/__about__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 94aeb1484c4ebe50f4d50f5a55d70e2dbc54941c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/__pycache__/exceptions.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/__pycache__/exceptions.cpython-310.pyc deleted file mode 100644 index e0213f28d43a3e9b2d33d331cd2d3925aa16ee4c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/__pycache__/exceptions.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/__pycache__/fernet.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/__pycache__/fernet.cpython-310.pyc deleted file mode 100644 index 6d39d63dfbdbc9ac84926f76c741fc904c385dca..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/__pycache__/fernet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index 01398af0d1fa667834047ef1a3fca37bb299dde6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/exceptions.py b/venv/lib/python3.10/site-packages/cryptography/exceptions.py deleted file mode 100644 index fe125ea9a763ad731ba99e4761d6623e1d2c639c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/exceptions.py +++ /dev/null @@ -1,52 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import typing - -from cryptography.hazmat.bindings._rust import exceptions as rust_exceptions - -if typing.TYPE_CHECKING: - from cryptography.hazmat.bindings._rust import openssl as rust_openssl - -_Reasons = rust_exceptions._Reasons - - -class UnsupportedAlgorithm(Exception): - def __init__(self, message: str, reason: _Reasons | None = None) -> None: - super().__init__(message) - self._reason = reason - - -class AlreadyFinalized(Exception): - pass - - -class AlreadyUpdated(Exception): - pass - - -class NotYetFinalized(Exception): - pass - - -class InvalidTag(Exception): - pass - - -class InvalidSignature(Exception): - pass - - -class InternalError(Exception): - def __init__( - self, msg: str, err_code: list[rust_openssl.OpenSSLError] - ) -> None: - super().__init__(msg) - self.err_code = err_code - - -class InvalidKey(Exception): - pass diff --git a/venv/lib/python3.10/site-packages/cryptography/fernet.py b/venv/lib/python3.10/site-packages/cryptography/fernet.py deleted file mode 100644 index c6744ae384e3497d2855b32398f8ee59b972f371..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/fernet.py +++ /dev/null @@ -1,224 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import base64 -import binascii -import os -import time -import typing -from collections.abc import Iterable - -from cryptography import utils -from cryptography.exceptions import InvalidSignature -from cryptography.hazmat.primitives import hashes, padding -from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes -from cryptography.hazmat.primitives.hmac import HMAC - - -class InvalidToken(Exception): - pass - - -_MAX_CLOCK_SKEW = 60 - - -class Fernet: - def __init__( - self, - key: bytes | str, - backend: typing.Any = None, - ) -> None: - try: - key = base64.urlsafe_b64decode(key) - except binascii.Error as exc: - raise ValueError( - "Fernet key must be 32 url-safe base64-encoded bytes." - ) from exc - if len(key) != 32: - raise ValueError( - "Fernet key must be 32 url-safe base64-encoded bytes." - ) - - self._signing_key = key[:16] - self._encryption_key = key[16:] - - @classmethod - def generate_key(cls) -> bytes: - return base64.urlsafe_b64encode(os.urandom(32)) - - def encrypt(self, data: bytes) -> bytes: - return self.encrypt_at_time(data, int(time.time())) - - def encrypt_at_time(self, data: bytes, current_time: int) -> bytes: - iv = os.urandom(16) - return self._encrypt_from_parts(data, current_time, iv) - - def _encrypt_from_parts( - self, data: bytes, current_time: int, iv: bytes - ) -> bytes: - utils._check_bytes("data", data) - - padder = padding.PKCS7(algorithms.AES.block_size).padder() - padded_data = padder.update(data) + padder.finalize() - encryptor = Cipher( - algorithms.AES(self._encryption_key), - modes.CBC(iv), - ).encryptor() - ciphertext = encryptor.update(padded_data) + encryptor.finalize() - - basic_parts = ( - b"\x80" - + current_time.to_bytes(length=8, byteorder="big") - + iv - + ciphertext - ) - - h = HMAC(self._signing_key, hashes.SHA256()) - h.update(basic_parts) - hmac = h.finalize() - return base64.urlsafe_b64encode(basic_parts + hmac) - - def decrypt(self, token: bytes | str, ttl: int | None = None) -> bytes: - timestamp, data = Fernet._get_unverified_token_data(token) - if ttl is None: - time_info = None - else: - time_info = (ttl, int(time.time())) - return self._decrypt_data(data, timestamp, time_info) - - def decrypt_at_time( - self, token: bytes | str, ttl: int, current_time: int - ) -> bytes: - if ttl is None: - raise ValueError( - "decrypt_at_time() can only be used with a non-None ttl" - ) - timestamp, data = Fernet._get_unverified_token_data(token) - return self._decrypt_data(data, timestamp, (ttl, current_time)) - - def extract_timestamp(self, token: bytes | str) -> int: - timestamp, data = Fernet._get_unverified_token_data(token) - # Verify the token was not tampered with. - self._verify_signature(data) - return timestamp - - @staticmethod - def _get_unverified_token_data(token: bytes | str) -> tuple[int, bytes]: - if not isinstance(token, (str, bytes)): - raise TypeError("token must be bytes or str") - - try: - data = base64.urlsafe_b64decode(token) - except (TypeError, binascii.Error): - raise InvalidToken - - if not data or data[0] != 0x80: - raise InvalidToken - - if len(data) < 9: - raise InvalidToken - - timestamp = int.from_bytes(data[1:9], byteorder="big") - return timestamp, data - - def _verify_signature(self, data: bytes) -> None: - h = HMAC(self._signing_key, hashes.SHA256()) - h.update(data[:-32]) - try: - h.verify(data[-32:]) - except InvalidSignature: - raise InvalidToken - - def _decrypt_data( - self, - data: bytes, - timestamp: int, - time_info: tuple[int, int] | None, - ) -> bytes: - if time_info is not None: - ttl, current_time = time_info - if timestamp + ttl < current_time: - raise InvalidToken - - if current_time + _MAX_CLOCK_SKEW < timestamp: - raise InvalidToken - - self._verify_signature(data) - - iv = data[9:25] - ciphertext = data[25:-32] - decryptor = Cipher( - algorithms.AES(self._encryption_key), modes.CBC(iv) - ).decryptor() - plaintext_padded = decryptor.update(ciphertext) - try: - plaintext_padded += decryptor.finalize() - except ValueError: - raise InvalidToken - unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder() - - unpadded = unpadder.update(plaintext_padded) - try: - unpadded += unpadder.finalize() - except ValueError: - raise InvalidToken - return unpadded - - -class MultiFernet: - def __init__(self, fernets: Iterable[Fernet]): - fernets = list(fernets) - if not fernets: - raise ValueError( - "MultiFernet requires at least one Fernet instance" - ) - self._fernets = fernets - - def encrypt(self, msg: bytes) -> bytes: - return self.encrypt_at_time(msg, int(time.time())) - - def encrypt_at_time(self, msg: bytes, current_time: int) -> bytes: - return self._fernets[0].encrypt_at_time(msg, current_time) - - def rotate(self, msg: bytes | str) -> bytes: - timestamp, data = Fernet._get_unverified_token_data(msg) - for f in self._fernets: - try: - p = f._decrypt_data(data, timestamp, None) - break - except InvalidToken: - pass - else: - raise InvalidToken - - iv = os.urandom(16) - return self._fernets[0]._encrypt_from_parts(p, timestamp, iv) - - def decrypt(self, msg: bytes | str, ttl: int | None = None) -> bytes: - for f in self._fernets: - try: - return f.decrypt(msg, ttl) - except InvalidToken: - pass - raise InvalidToken - - def decrypt_at_time( - self, msg: bytes | str, ttl: int, current_time: int - ) -> bytes: - for f in self._fernets: - try: - return f.decrypt_at_time(msg, ttl, current_time) - except InvalidToken: - pass - raise InvalidToken - - def extract_timestamp(self, msg: bytes | str) -> int: - for f in self._fernets: - try: - return f.extract_timestamp(msg) - except InvalidToken: - pass - raise InvalidToken diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/__init__.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/__init__.py deleted file mode 100644 index b9f1187011bdaa0720bc462564582393700f3d4a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -""" -Hazardous Materials - -This is a "Hazardous Materials" module. You should ONLY use it if you're -100% absolutely sure that you know what you're doing because this module -is full of land mines, dragons, and dinosaurs with laser guns. -""" diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 85a7f011d8368a2fdff838e2d0b5bc0ef00b0b62..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/__pycache__/_oid.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/__pycache__/_oid.cpython-310.pyc deleted file mode 100644 index fbcdf7745d5a88c6cc4cd2e8b70a1dc823e5c39f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/__pycache__/_oid.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/_oid.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/_oid.py deleted file mode 100644 index 249b4dc857af98bc76f5245a9512c639f5e78533..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/_oid.py +++ /dev/null @@ -1,348 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography.hazmat.bindings._rust import ( - ObjectIdentifier as ObjectIdentifier, -) -from cryptography.hazmat.primitives import hashes - - -class ExtensionOID: - SUBJECT_DIRECTORY_ATTRIBUTES = ObjectIdentifier("2.5.29.9") - SUBJECT_KEY_IDENTIFIER = ObjectIdentifier("2.5.29.14") - KEY_USAGE = ObjectIdentifier("2.5.29.15") - PRIVATE_KEY_USAGE_PERIOD = ObjectIdentifier("2.5.29.16") - SUBJECT_ALTERNATIVE_NAME = ObjectIdentifier("2.5.29.17") - ISSUER_ALTERNATIVE_NAME = ObjectIdentifier("2.5.29.18") - BASIC_CONSTRAINTS = ObjectIdentifier("2.5.29.19") - NAME_CONSTRAINTS = ObjectIdentifier("2.5.29.30") - CRL_DISTRIBUTION_POINTS = ObjectIdentifier("2.5.29.31") - CERTIFICATE_POLICIES = ObjectIdentifier("2.5.29.32") - POLICY_MAPPINGS = ObjectIdentifier("2.5.29.33") - AUTHORITY_KEY_IDENTIFIER = ObjectIdentifier("2.5.29.35") - POLICY_CONSTRAINTS = ObjectIdentifier("2.5.29.36") - EXTENDED_KEY_USAGE = ObjectIdentifier("2.5.29.37") - FRESHEST_CRL = ObjectIdentifier("2.5.29.46") - INHIBIT_ANY_POLICY = ObjectIdentifier("2.5.29.54") - ISSUING_DISTRIBUTION_POINT = ObjectIdentifier("2.5.29.28") - AUTHORITY_INFORMATION_ACCESS = ObjectIdentifier("1.3.6.1.5.5.7.1.1") - SUBJECT_INFORMATION_ACCESS = ObjectIdentifier("1.3.6.1.5.5.7.1.11") - OCSP_NO_CHECK = ObjectIdentifier("1.3.6.1.5.5.7.48.1.5") - TLS_FEATURE = ObjectIdentifier("1.3.6.1.5.5.7.1.24") - CRL_NUMBER = ObjectIdentifier("2.5.29.20") - DELTA_CRL_INDICATOR = ObjectIdentifier("2.5.29.27") - PRECERT_SIGNED_CERTIFICATE_TIMESTAMPS = ObjectIdentifier( - "1.3.6.1.4.1.11129.2.4.2" - ) - PRECERT_POISON = ObjectIdentifier("1.3.6.1.4.1.11129.2.4.3") - SIGNED_CERTIFICATE_TIMESTAMPS = ObjectIdentifier("1.3.6.1.4.1.11129.2.4.5") - MS_CERTIFICATE_TEMPLATE = ObjectIdentifier("1.3.6.1.4.1.311.21.7") - ADMISSIONS = ObjectIdentifier("1.3.36.8.3.3") - - -class OCSPExtensionOID: - NONCE = ObjectIdentifier("1.3.6.1.5.5.7.48.1.2") - ACCEPTABLE_RESPONSES = ObjectIdentifier("1.3.6.1.5.5.7.48.1.4") - - -class CRLEntryExtensionOID: - CERTIFICATE_ISSUER = ObjectIdentifier("2.5.29.29") - CRL_REASON = ObjectIdentifier("2.5.29.21") - INVALIDITY_DATE = ObjectIdentifier("2.5.29.24") - - -class NameOID: - COMMON_NAME = ObjectIdentifier("2.5.4.3") - COUNTRY_NAME = ObjectIdentifier("2.5.4.6") - LOCALITY_NAME = ObjectIdentifier("2.5.4.7") - STATE_OR_PROVINCE_NAME = ObjectIdentifier("2.5.4.8") - STREET_ADDRESS = ObjectIdentifier("2.5.4.9") - ORGANIZATION_IDENTIFIER = ObjectIdentifier("2.5.4.97") - ORGANIZATION_NAME = ObjectIdentifier("2.5.4.10") - ORGANIZATIONAL_UNIT_NAME = ObjectIdentifier("2.5.4.11") - SERIAL_NUMBER = ObjectIdentifier("2.5.4.5") - SURNAME = ObjectIdentifier("2.5.4.4") - GIVEN_NAME = ObjectIdentifier("2.5.4.42") - TITLE = ObjectIdentifier("2.5.4.12") - INITIALS = ObjectIdentifier("2.5.4.43") - GENERATION_QUALIFIER = ObjectIdentifier("2.5.4.44") - X500_UNIQUE_IDENTIFIER = ObjectIdentifier("2.5.4.45") - DN_QUALIFIER = ObjectIdentifier("2.5.4.46") - PSEUDONYM = ObjectIdentifier("2.5.4.65") - USER_ID = ObjectIdentifier("0.9.2342.19200300.100.1.1") - DOMAIN_COMPONENT = ObjectIdentifier("0.9.2342.19200300.100.1.25") - EMAIL_ADDRESS = ObjectIdentifier("1.2.840.113549.1.9.1") - JURISDICTION_COUNTRY_NAME = ObjectIdentifier("1.3.6.1.4.1.311.60.2.1.3") - JURISDICTION_LOCALITY_NAME = ObjectIdentifier("1.3.6.1.4.1.311.60.2.1.1") - JURISDICTION_STATE_OR_PROVINCE_NAME = ObjectIdentifier( - "1.3.6.1.4.1.311.60.2.1.2" - ) - BUSINESS_CATEGORY = ObjectIdentifier("2.5.4.15") - POSTAL_ADDRESS = ObjectIdentifier("2.5.4.16") - POSTAL_CODE = ObjectIdentifier("2.5.4.17") - INN = ObjectIdentifier("1.2.643.3.131.1.1") - OGRN = ObjectIdentifier("1.2.643.100.1") - SNILS = ObjectIdentifier("1.2.643.100.3") - UNSTRUCTURED_NAME = ObjectIdentifier("1.2.840.113549.1.9.2") - - -class SignatureAlgorithmOID: - RSA_WITH_MD5 = ObjectIdentifier("1.2.840.113549.1.1.4") - RSA_WITH_SHA1 = ObjectIdentifier("1.2.840.113549.1.1.5") - # This is an alternate OID for RSA with SHA1 that is occasionally seen - _RSA_WITH_SHA1 = ObjectIdentifier("1.3.14.3.2.29") - RSA_WITH_SHA224 = ObjectIdentifier("1.2.840.113549.1.1.14") - RSA_WITH_SHA256 = ObjectIdentifier("1.2.840.113549.1.1.11") - RSA_WITH_SHA384 = ObjectIdentifier("1.2.840.113549.1.1.12") - RSA_WITH_SHA512 = ObjectIdentifier("1.2.840.113549.1.1.13") - RSA_WITH_SHA3_224 = ObjectIdentifier("2.16.840.1.101.3.4.3.13") - RSA_WITH_SHA3_256 = ObjectIdentifier("2.16.840.1.101.3.4.3.14") - RSA_WITH_SHA3_384 = ObjectIdentifier("2.16.840.1.101.3.4.3.15") - RSA_WITH_SHA3_512 = ObjectIdentifier("2.16.840.1.101.3.4.3.16") - RSASSA_PSS = ObjectIdentifier("1.2.840.113549.1.1.10") - ECDSA_WITH_SHA1 = ObjectIdentifier("1.2.840.10045.4.1") - ECDSA_WITH_SHA224 = ObjectIdentifier("1.2.840.10045.4.3.1") - ECDSA_WITH_SHA256 = ObjectIdentifier("1.2.840.10045.4.3.2") - ECDSA_WITH_SHA384 = ObjectIdentifier("1.2.840.10045.4.3.3") - ECDSA_WITH_SHA512 = ObjectIdentifier("1.2.840.10045.4.3.4") - ECDSA_WITH_SHA3_224 = ObjectIdentifier("2.16.840.1.101.3.4.3.9") - ECDSA_WITH_SHA3_256 = ObjectIdentifier("2.16.840.1.101.3.4.3.10") - ECDSA_WITH_SHA3_384 = ObjectIdentifier("2.16.840.1.101.3.4.3.11") - ECDSA_WITH_SHA3_512 = ObjectIdentifier("2.16.840.1.101.3.4.3.12") - DSA_WITH_SHA1 = ObjectIdentifier("1.2.840.10040.4.3") - DSA_WITH_SHA224 = ObjectIdentifier("2.16.840.1.101.3.4.3.1") - DSA_WITH_SHA256 = ObjectIdentifier("2.16.840.1.101.3.4.3.2") - DSA_WITH_SHA384 = ObjectIdentifier("2.16.840.1.101.3.4.3.3") - DSA_WITH_SHA512 = ObjectIdentifier("2.16.840.1.101.3.4.3.4") - ED25519 = ObjectIdentifier("1.3.101.112") - ED448 = ObjectIdentifier("1.3.101.113") - GOSTR3411_94_WITH_3410_2001 = ObjectIdentifier("1.2.643.2.2.3") - GOSTR3410_2012_WITH_3411_2012_256 = ObjectIdentifier("1.2.643.7.1.1.3.2") - GOSTR3410_2012_WITH_3411_2012_512 = ObjectIdentifier("1.2.643.7.1.1.3.3") - - -_SIG_OIDS_TO_HASH: dict[ObjectIdentifier, hashes.HashAlgorithm | None] = { - SignatureAlgorithmOID.RSA_WITH_MD5: hashes.MD5(), - SignatureAlgorithmOID.RSA_WITH_SHA1: hashes.SHA1(), - SignatureAlgorithmOID._RSA_WITH_SHA1: hashes.SHA1(), - SignatureAlgorithmOID.RSA_WITH_SHA224: hashes.SHA224(), - SignatureAlgorithmOID.RSA_WITH_SHA256: hashes.SHA256(), - SignatureAlgorithmOID.RSA_WITH_SHA384: hashes.SHA384(), - SignatureAlgorithmOID.RSA_WITH_SHA512: hashes.SHA512(), - SignatureAlgorithmOID.RSA_WITH_SHA3_224: hashes.SHA3_224(), - SignatureAlgorithmOID.RSA_WITH_SHA3_256: hashes.SHA3_256(), - SignatureAlgorithmOID.RSA_WITH_SHA3_384: hashes.SHA3_384(), - SignatureAlgorithmOID.RSA_WITH_SHA3_512: hashes.SHA3_512(), - SignatureAlgorithmOID.ECDSA_WITH_SHA1: hashes.SHA1(), - SignatureAlgorithmOID.ECDSA_WITH_SHA224: hashes.SHA224(), - SignatureAlgorithmOID.ECDSA_WITH_SHA256: hashes.SHA256(), - SignatureAlgorithmOID.ECDSA_WITH_SHA384: hashes.SHA384(), - SignatureAlgorithmOID.ECDSA_WITH_SHA512: hashes.SHA512(), - SignatureAlgorithmOID.ECDSA_WITH_SHA3_224: hashes.SHA3_224(), - SignatureAlgorithmOID.ECDSA_WITH_SHA3_256: hashes.SHA3_256(), - SignatureAlgorithmOID.ECDSA_WITH_SHA3_384: hashes.SHA3_384(), - SignatureAlgorithmOID.ECDSA_WITH_SHA3_512: hashes.SHA3_512(), - SignatureAlgorithmOID.DSA_WITH_SHA1: hashes.SHA1(), - SignatureAlgorithmOID.DSA_WITH_SHA224: hashes.SHA224(), - SignatureAlgorithmOID.DSA_WITH_SHA256: hashes.SHA256(), - SignatureAlgorithmOID.ED25519: None, - SignatureAlgorithmOID.ED448: None, - SignatureAlgorithmOID.GOSTR3411_94_WITH_3410_2001: None, - SignatureAlgorithmOID.GOSTR3410_2012_WITH_3411_2012_256: None, - SignatureAlgorithmOID.GOSTR3410_2012_WITH_3411_2012_512: None, -} - - -class HashAlgorithmOID: - SHA1 = ObjectIdentifier("1.3.14.3.2.26") - SHA224 = ObjectIdentifier("2.16.840.1.101.3.4.2.4") - SHA256 = ObjectIdentifier("2.16.840.1.101.3.4.2.1") - SHA384 = ObjectIdentifier("2.16.840.1.101.3.4.2.2") - SHA512 = ObjectIdentifier("2.16.840.1.101.3.4.2.3") - SHA3_224 = ObjectIdentifier("1.3.6.1.4.1.37476.3.2.1.99.7.224") - SHA3_256 = ObjectIdentifier("1.3.6.1.4.1.37476.3.2.1.99.7.256") - SHA3_384 = ObjectIdentifier("1.3.6.1.4.1.37476.3.2.1.99.7.384") - SHA3_512 = ObjectIdentifier("1.3.6.1.4.1.37476.3.2.1.99.7.512") - - -class PublicKeyAlgorithmOID: - DSA = ObjectIdentifier("1.2.840.10040.4.1") - EC_PUBLIC_KEY = ObjectIdentifier("1.2.840.10045.2.1") - RSAES_PKCS1_v1_5 = ObjectIdentifier("1.2.840.113549.1.1.1") - RSASSA_PSS = ObjectIdentifier("1.2.840.113549.1.1.10") - X25519 = ObjectIdentifier("1.3.101.110") - X448 = ObjectIdentifier("1.3.101.111") - ED25519 = ObjectIdentifier("1.3.101.112") - ED448 = ObjectIdentifier("1.3.101.113") - - -class ExtendedKeyUsageOID: - SERVER_AUTH = ObjectIdentifier("1.3.6.1.5.5.7.3.1") - CLIENT_AUTH = ObjectIdentifier("1.3.6.1.5.5.7.3.2") - CODE_SIGNING = ObjectIdentifier("1.3.6.1.5.5.7.3.3") - EMAIL_PROTECTION = ObjectIdentifier("1.3.6.1.5.5.7.3.4") - TIME_STAMPING = ObjectIdentifier("1.3.6.1.5.5.7.3.8") - OCSP_SIGNING = ObjectIdentifier("1.3.6.1.5.5.7.3.9") - ANY_EXTENDED_KEY_USAGE = ObjectIdentifier("2.5.29.37.0") - SMARTCARD_LOGON = ObjectIdentifier("1.3.6.1.4.1.311.20.2.2") - KERBEROS_PKINIT_KDC = ObjectIdentifier("1.3.6.1.5.2.3.5") - IPSEC_IKE = ObjectIdentifier("1.3.6.1.5.5.7.3.17") - BUNDLE_SECURITY = ObjectIdentifier("1.3.6.1.5.5.7.3.35") - CERTIFICATE_TRANSPARENCY = ObjectIdentifier("1.3.6.1.4.1.11129.2.4.4") - - -class OtherNameFormOID: - PERMANENT_IDENTIFIER = ObjectIdentifier("1.3.6.1.5.5.7.8.3") - HW_MODULE_NAME = ObjectIdentifier("1.3.6.1.5.5.7.8.4") - DNS_SRV = ObjectIdentifier("1.3.6.1.5.5.7.8.7") - NAI_REALM = ObjectIdentifier("1.3.6.1.5.5.7.8.8") - SMTP_UTF8_MAILBOX = ObjectIdentifier("1.3.6.1.5.5.7.8.9") - ACP_NODE_NAME = ObjectIdentifier("1.3.6.1.5.5.7.8.10") - BUNDLE_EID = ObjectIdentifier("1.3.6.1.5.5.7.8.11") - - -class AuthorityInformationAccessOID: - CA_ISSUERS = ObjectIdentifier("1.3.6.1.5.5.7.48.2") - OCSP = ObjectIdentifier("1.3.6.1.5.5.7.48.1") - - -class SubjectInformationAccessOID: - CA_REPOSITORY = ObjectIdentifier("1.3.6.1.5.5.7.48.5") - - -class CertificatePoliciesOID: - CPS_QUALIFIER = ObjectIdentifier("1.3.6.1.5.5.7.2.1") - CPS_USER_NOTICE = ObjectIdentifier("1.3.6.1.5.5.7.2.2") - ANY_POLICY = ObjectIdentifier("2.5.29.32.0") - - -class AttributeOID: - CHALLENGE_PASSWORD = ObjectIdentifier("1.2.840.113549.1.9.7") - UNSTRUCTURED_NAME = ObjectIdentifier("1.2.840.113549.1.9.2") - - -_OID_NAMES = { - NameOID.COMMON_NAME: "commonName", - NameOID.COUNTRY_NAME: "countryName", - NameOID.LOCALITY_NAME: "localityName", - NameOID.STATE_OR_PROVINCE_NAME: "stateOrProvinceName", - NameOID.STREET_ADDRESS: "streetAddress", - NameOID.ORGANIZATION_NAME: "organizationName", - NameOID.ORGANIZATIONAL_UNIT_NAME: "organizationalUnitName", - NameOID.SERIAL_NUMBER: "serialNumber", - NameOID.SURNAME: "surname", - NameOID.GIVEN_NAME: "givenName", - NameOID.TITLE: "title", - NameOID.GENERATION_QUALIFIER: "generationQualifier", - NameOID.X500_UNIQUE_IDENTIFIER: "x500UniqueIdentifier", - NameOID.DN_QUALIFIER: "dnQualifier", - NameOID.PSEUDONYM: "pseudonym", - NameOID.USER_ID: "userID", - NameOID.DOMAIN_COMPONENT: "domainComponent", - NameOID.EMAIL_ADDRESS: "emailAddress", - NameOID.JURISDICTION_COUNTRY_NAME: "jurisdictionCountryName", - NameOID.JURISDICTION_LOCALITY_NAME: "jurisdictionLocalityName", - NameOID.JURISDICTION_STATE_OR_PROVINCE_NAME: ( - "jurisdictionStateOrProvinceName" - ), - NameOID.BUSINESS_CATEGORY: "businessCategory", - NameOID.POSTAL_ADDRESS: "postalAddress", - NameOID.POSTAL_CODE: "postalCode", - NameOID.INN: "INN", - NameOID.OGRN: "OGRN", - NameOID.SNILS: "SNILS", - NameOID.UNSTRUCTURED_NAME: "unstructuredName", - SignatureAlgorithmOID.RSA_WITH_MD5: "md5WithRSAEncryption", - SignatureAlgorithmOID.RSA_WITH_SHA1: "sha1WithRSAEncryption", - SignatureAlgorithmOID.RSA_WITH_SHA224: "sha224WithRSAEncryption", - SignatureAlgorithmOID.RSA_WITH_SHA256: "sha256WithRSAEncryption", - SignatureAlgorithmOID.RSA_WITH_SHA384: "sha384WithRSAEncryption", - SignatureAlgorithmOID.RSA_WITH_SHA512: "sha512WithRSAEncryption", - SignatureAlgorithmOID.RSASSA_PSS: "rsassaPss", - SignatureAlgorithmOID.ECDSA_WITH_SHA1: "ecdsa-with-SHA1", - SignatureAlgorithmOID.ECDSA_WITH_SHA224: "ecdsa-with-SHA224", - SignatureAlgorithmOID.ECDSA_WITH_SHA256: "ecdsa-with-SHA256", - SignatureAlgorithmOID.ECDSA_WITH_SHA384: "ecdsa-with-SHA384", - SignatureAlgorithmOID.ECDSA_WITH_SHA512: "ecdsa-with-SHA512", - SignatureAlgorithmOID.DSA_WITH_SHA1: "dsa-with-sha1", - SignatureAlgorithmOID.DSA_WITH_SHA224: "dsa-with-sha224", - SignatureAlgorithmOID.DSA_WITH_SHA256: "dsa-with-sha256", - SignatureAlgorithmOID.ED25519: "ed25519", - SignatureAlgorithmOID.ED448: "ed448", - SignatureAlgorithmOID.GOSTR3411_94_WITH_3410_2001: ( - "GOST R 34.11-94 with GOST R 34.10-2001" - ), - SignatureAlgorithmOID.GOSTR3410_2012_WITH_3411_2012_256: ( - "GOST R 34.10-2012 with GOST R 34.11-2012 (256 bit)" - ), - SignatureAlgorithmOID.GOSTR3410_2012_WITH_3411_2012_512: ( - "GOST R 34.10-2012 with GOST R 34.11-2012 (512 bit)" - ), - HashAlgorithmOID.SHA1: "sha1", - HashAlgorithmOID.SHA224: "sha224", - HashAlgorithmOID.SHA256: "sha256", - HashAlgorithmOID.SHA384: "sha384", - HashAlgorithmOID.SHA512: "sha512", - HashAlgorithmOID.SHA3_224: "sha3_224", - HashAlgorithmOID.SHA3_256: "sha3_256", - HashAlgorithmOID.SHA3_384: "sha3_384", - HashAlgorithmOID.SHA3_512: "sha3_512", - PublicKeyAlgorithmOID.DSA: "dsaEncryption", - PublicKeyAlgorithmOID.EC_PUBLIC_KEY: "id-ecPublicKey", - PublicKeyAlgorithmOID.RSAES_PKCS1_v1_5: "rsaEncryption", - PublicKeyAlgorithmOID.X25519: "X25519", - PublicKeyAlgorithmOID.X448: "X448", - ExtendedKeyUsageOID.SERVER_AUTH: "serverAuth", - ExtendedKeyUsageOID.CLIENT_AUTH: "clientAuth", - ExtendedKeyUsageOID.CODE_SIGNING: "codeSigning", - ExtendedKeyUsageOID.EMAIL_PROTECTION: "emailProtection", - ExtendedKeyUsageOID.TIME_STAMPING: "timeStamping", - ExtendedKeyUsageOID.OCSP_SIGNING: "OCSPSigning", - ExtendedKeyUsageOID.SMARTCARD_LOGON: "msSmartcardLogin", - ExtendedKeyUsageOID.KERBEROS_PKINIT_KDC: "pkInitKDC", - ExtensionOID.SUBJECT_DIRECTORY_ATTRIBUTES: "subjectDirectoryAttributes", - ExtensionOID.SUBJECT_KEY_IDENTIFIER: "subjectKeyIdentifier", - ExtensionOID.KEY_USAGE: "keyUsage", - ExtensionOID.PRIVATE_KEY_USAGE_PERIOD: "privateKeyUsagePeriod", - ExtensionOID.SUBJECT_ALTERNATIVE_NAME: "subjectAltName", - ExtensionOID.ISSUER_ALTERNATIVE_NAME: "issuerAltName", - ExtensionOID.BASIC_CONSTRAINTS: "basicConstraints", - ExtensionOID.PRECERT_SIGNED_CERTIFICATE_TIMESTAMPS: ( - "signedCertificateTimestampList" - ), - ExtensionOID.SIGNED_CERTIFICATE_TIMESTAMPS: ( - "signedCertificateTimestampList" - ), - ExtensionOID.PRECERT_POISON: "ctPoison", - ExtensionOID.MS_CERTIFICATE_TEMPLATE: "msCertificateTemplate", - ExtensionOID.ADMISSIONS: "Admissions", - CRLEntryExtensionOID.CRL_REASON: "cRLReason", - CRLEntryExtensionOID.INVALIDITY_DATE: "invalidityDate", - CRLEntryExtensionOID.CERTIFICATE_ISSUER: "certificateIssuer", - ExtensionOID.NAME_CONSTRAINTS: "nameConstraints", - ExtensionOID.CRL_DISTRIBUTION_POINTS: "cRLDistributionPoints", - ExtensionOID.CERTIFICATE_POLICIES: "certificatePolicies", - ExtensionOID.POLICY_MAPPINGS: "policyMappings", - ExtensionOID.AUTHORITY_KEY_IDENTIFIER: "authorityKeyIdentifier", - ExtensionOID.POLICY_CONSTRAINTS: "policyConstraints", - ExtensionOID.EXTENDED_KEY_USAGE: "extendedKeyUsage", - ExtensionOID.FRESHEST_CRL: "freshestCRL", - ExtensionOID.INHIBIT_ANY_POLICY: "inhibitAnyPolicy", - ExtensionOID.ISSUING_DISTRIBUTION_POINT: "issuingDistributionPoint", - ExtensionOID.AUTHORITY_INFORMATION_ACCESS: "authorityInfoAccess", - ExtensionOID.SUBJECT_INFORMATION_ACCESS: "subjectInfoAccess", - ExtensionOID.OCSP_NO_CHECK: "OCSPNoCheck", - ExtensionOID.CRL_NUMBER: "cRLNumber", - ExtensionOID.DELTA_CRL_INDICATOR: "deltaCRLIndicator", - ExtensionOID.TLS_FEATURE: "TLSFeature", - AuthorityInformationAccessOID.OCSP: "OCSP", - AuthorityInformationAccessOID.CA_ISSUERS: "caIssuers", - SubjectInformationAccessOID.CA_REPOSITORY: "caRepository", - CertificatePoliciesOID.CPS_QUALIFIER: "id-qt-cps", - CertificatePoliciesOID.CPS_USER_NOTICE: "id-qt-unotice", - OCSPExtensionOID.NONCE: "OCSPNonce", - AttributeOID.CHALLENGE_PASSWORD: "challengePassword", -} diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/__init__.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/__init__.py deleted file mode 100644 index b4400aa037451d10d6bcadd455dcf98f9b72e221..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from typing import Any - - -def default_backend() -> Any: - from cryptography.hazmat.backends.openssl.backend import backend - - return backend diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index c4edc20f2779d1df1431607f5f8c5a2c35823805..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/openssl/__init__.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/openssl/__init__.py deleted file mode 100644 index 51b04476cbb7f4a98051f2fc55bcc1aa35e22d05..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/openssl/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography.hazmat.backends.openssl.backend import backend - -__all__ = ["backend"] diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/openssl/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/openssl/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index d679f655e0e60a498849a4fe438074e90a716244..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/openssl/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/openssl/__pycache__/backend.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/openssl/__pycache__/backend.cpython-310.pyc deleted file mode 100644 index bdc408c3a0a23e3c6b3d5fdbb7effa1b4c1d8f30..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/openssl/__pycache__/backend.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/openssl/backend.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/openssl/backend.py deleted file mode 100644 index 361011a86c440b5dc06ea546d983838d1e570467..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/backends/openssl/backend.py +++ /dev/null @@ -1,308 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.hazmat.bindings.openssl import binding -from cryptography.hazmat.primitives import hashes -from cryptography.hazmat.primitives._asymmetric import AsymmetricPadding -from cryptography.hazmat.primitives.asymmetric import ec -from cryptography.hazmat.primitives.asymmetric import utils as asym_utils -from cryptography.hazmat.primitives.asymmetric.padding import ( - MGF1, - OAEP, - PSS, - PKCS1v15, -) -from cryptography.hazmat.primitives.ciphers import ( - CipherAlgorithm, -) -from cryptography.hazmat.primitives.ciphers.algorithms import ( - AES, -) -from cryptography.hazmat.primitives.ciphers.modes import ( - CBC, - Mode, -) - - -class Backend: - """ - OpenSSL API binding interfaces. - """ - - name = "openssl" - - # TripleDES encryption is disallowed/deprecated throughout 2023 in - # FIPS 140-3. To keep it simple we denylist any use of TripleDES (TDEA). - _fips_ciphers = (AES,) - # Sometimes SHA1 is still permissible. That logic is contained - # within the various *_supported methods. - _fips_hashes = ( - hashes.SHA224, - hashes.SHA256, - hashes.SHA384, - hashes.SHA512, - hashes.SHA512_224, - hashes.SHA512_256, - hashes.SHA3_224, - hashes.SHA3_256, - hashes.SHA3_384, - hashes.SHA3_512, - hashes.SHAKE128, - hashes.SHAKE256, - ) - _fips_ecdh_curves = ( - ec.SECP224R1, - ec.SECP256R1, - ec.SECP384R1, - ec.SECP521R1, - ) - _fips_rsa_min_key_size = 2048 - _fips_rsa_min_public_exponent = 65537 - _fips_dsa_min_modulus = 1 << 2048 - _fips_dh_min_key_size = 2048 - _fips_dh_min_modulus = 1 << _fips_dh_min_key_size - - def __init__(self) -> None: - self._binding = binding.Binding() - self._ffi = self._binding.ffi - self._lib = self._binding.lib - self._fips_enabled = rust_openssl.is_fips_enabled() - - def __repr__(self) -> str: - return ( - f"" - ) - - def openssl_assert(self, ok: bool) -> None: - return binding._openssl_assert(ok) - - def _enable_fips(self) -> None: - # This function enables FIPS mode for OpenSSL 3.0.0 on installs that - # have the FIPS provider installed properly. - rust_openssl.enable_fips(rust_openssl._providers) - assert rust_openssl.is_fips_enabled() - self._fips_enabled = rust_openssl.is_fips_enabled() - - def openssl_version_text(self) -> str: - """ - Friendly string name of the loaded OpenSSL library. This is not - necessarily the same version as it was compiled against. - - Example: OpenSSL 3.2.1 30 Jan 2024 - """ - return rust_openssl.openssl_version_text() - - def openssl_version_number(self) -> int: - return rust_openssl.openssl_version() - - def hash_supported(self, algorithm: hashes.HashAlgorithm) -> bool: - if self._fips_enabled and not isinstance(algorithm, self._fips_hashes): - return False - - return rust_openssl.hashes.hash_supported(algorithm) - - def signature_hash_supported( - self, algorithm: hashes.HashAlgorithm - ) -> bool: - # Dedicated check for hashing algorithm use in message digest for - # signatures, e.g. RSA PKCS#1 v1.5 SHA1 (sha1WithRSAEncryption). - if self._fips_enabled and isinstance(algorithm, hashes.SHA1): - return False - return self.hash_supported(algorithm) - - def scrypt_supported(self) -> bool: - if self._fips_enabled: - return False - else: - return hasattr(rust_openssl.kdf.Scrypt, "derive") - - def argon2_supported(self) -> bool: - if self._fips_enabled: - return False - else: - return hasattr(rust_openssl.kdf.Argon2id, "derive") - - def hmac_supported(self, algorithm: hashes.HashAlgorithm) -> bool: - # FIPS mode still allows SHA1 for HMAC - if self._fips_enabled and isinstance(algorithm, hashes.SHA1): - return True - if rust_openssl.CRYPTOGRAPHY_IS_AWSLC: - return isinstance( - algorithm, - ( - hashes.SHA1, - hashes.SHA224, - hashes.SHA256, - hashes.SHA384, - hashes.SHA512, - hashes.SHA512_224, - hashes.SHA512_256, - ), - ) - return self.hash_supported(algorithm) - - def cipher_supported(self, cipher: CipherAlgorithm, mode: Mode) -> bool: - if self._fips_enabled: - # FIPS mode requires AES. TripleDES is disallowed/deprecated in - # FIPS 140-3. - if not isinstance(cipher, self._fips_ciphers): - return False - - return rust_openssl.ciphers.cipher_supported(cipher, mode) - - def pbkdf2_hmac_supported(self, algorithm: hashes.HashAlgorithm) -> bool: - return self.hmac_supported(algorithm) - - def _consume_errors(self) -> list[rust_openssl.OpenSSLError]: - return rust_openssl.capture_error_stack() - - def _oaep_hash_supported(self, algorithm: hashes.HashAlgorithm) -> bool: - if self._fips_enabled and isinstance(algorithm, hashes.SHA1): - return False - - return isinstance( - algorithm, - ( - hashes.SHA1, - hashes.SHA224, - hashes.SHA256, - hashes.SHA384, - hashes.SHA512, - ), - ) - - def rsa_padding_supported(self, padding: AsymmetricPadding) -> bool: - if isinstance(padding, PKCS1v15): - return True - elif isinstance(padding, PSS) and isinstance(padding._mgf, MGF1): - # FIPS 186-4 only allows salt length == digest length for PSS - # It is technically acceptable to set an explicit salt length - # equal to the digest length and this will incorrectly fail, but - # since we don't do that in the tests and this method is - # private, we'll ignore that until we need to do otherwise. - if ( - self._fips_enabled - and padding._salt_length != PSS.DIGEST_LENGTH - ): - return False - return self.hash_supported(padding._mgf._algorithm) - elif isinstance(padding, OAEP) and isinstance(padding._mgf, MGF1): - return self._oaep_hash_supported( - padding._mgf._algorithm - ) and self._oaep_hash_supported(padding._algorithm) - else: - return False - - def rsa_encryption_supported(self, padding: AsymmetricPadding) -> bool: - if self._fips_enabled and isinstance(padding, PKCS1v15): - return False - else: - return self.rsa_padding_supported(padding) - - def dsa_supported(self) -> bool: - return ( - not rust_openssl.CRYPTOGRAPHY_IS_BORINGSSL - and not self._fips_enabled - ) - - def dsa_hash_supported(self, algorithm: hashes.HashAlgorithm) -> bool: - if not self.dsa_supported(): - return False - return self.signature_hash_supported(algorithm) - - def cmac_algorithm_supported(self, algorithm) -> bool: - return self.cipher_supported( - algorithm, CBC(b"\x00" * algorithm.block_size) - ) - - def elliptic_curve_supported(self, curve: ec.EllipticCurve) -> bool: - if self._fips_enabled and not isinstance( - curve, self._fips_ecdh_curves - ): - return False - - return rust_openssl.ec.curve_supported(curve) - - def elliptic_curve_signature_algorithm_supported( - self, - signature_algorithm: ec.EllipticCurveSignatureAlgorithm, - curve: ec.EllipticCurve, - ) -> bool: - # We only support ECDSA right now. - if not isinstance(signature_algorithm, ec.ECDSA): - return False - - return self.elliptic_curve_supported(curve) and ( - isinstance(signature_algorithm.algorithm, asym_utils.Prehashed) - or self.hash_supported(signature_algorithm.algorithm) - ) - - def elliptic_curve_exchange_algorithm_supported( - self, algorithm: ec.ECDH, curve: ec.EllipticCurve - ) -> bool: - return self.elliptic_curve_supported(curve) and isinstance( - algorithm, ec.ECDH - ) - - def dh_supported(self) -> bool: - return ( - not rust_openssl.CRYPTOGRAPHY_IS_BORINGSSL - and not rust_openssl.CRYPTOGRAPHY_IS_AWSLC - ) - - def dh_x942_serialization_supported(self) -> bool: - return self._lib.Cryptography_HAS_EVP_PKEY_DHX == 1 - - def x25519_supported(self) -> bool: - if self._fips_enabled: - return False - return True - - def x448_supported(self) -> bool: - if self._fips_enabled: - return False - return ( - not rust_openssl.CRYPTOGRAPHY_IS_LIBRESSL - and not rust_openssl.CRYPTOGRAPHY_IS_BORINGSSL - and not rust_openssl.CRYPTOGRAPHY_IS_AWSLC - ) - - def ed25519_supported(self) -> bool: - if self._fips_enabled: - return False - return True - - def ed448_supported(self) -> bool: - if self._fips_enabled: - return False - return ( - not rust_openssl.CRYPTOGRAPHY_IS_LIBRESSL - and not rust_openssl.CRYPTOGRAPHY_IS_BORINGSSL - and not rust_openssl.CRYPTOGRAPHY_IS_AWSLC - ) - - def ecdsa_deterministic_supported(self) -> bool: - return ( - rust_openssl.CRYPTOGRAPHY_OPENSSL_320_OR_GREATER - and not self._fips_enabled - ) - - def poly1305_supported(self) -> bool: - if self._fips_enabled: - return False - return True - - def pkcs7_supported(self) -> bool: - return ( - not rust_openssl.CRYPTOGRAPHY_IS_BORINGSSL - and not rust_openssl.CRYPTOGRAPHY_IS_AWSLC - ) - - -backend = Backend() diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/__init__.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/__init__.py deleted file mode 100644 index b509336233c2fafe4185a49da5909c8bbb38dfd7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index f82e3df95e5592462b2528fe2a685c4e8db3b700..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust.abi3.so b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust.abi3.so deleted file mode 100644 index 609af0780c2e8beef809173dd59272d48790a77a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:06db1cc8b574875fd6cbcccf4bd8343801b828e26c9ea2d51b5c3866a667af9b -size 12446840 diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/__init__.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/__init__.pyi deleted file mode 100644 index 2f4eef4ead80ad7c006ddf7c4b28fcd48ad6fa33..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/__init__.pyi +++ /dev/null @@ -1,37 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import typing - -from cryptography.hazmat.primitives import padding -from cryptography.utils import Buffer - -class PKCS7PaddingContext(padding.PaddingContext): - def __init__(self, block_size: int) -> None: ... - def update(self, data: Buffer) -> bytes: ... - def finalize(self) -> bytes: ... - -class ANSIX923PaddingContext(padding.PaddingContext): - def __init__(self, block_size: int) -> None: ... - def update(self, data: Buffer) -> bytes: ... - def finalize(self) -> bytes: ... - -class PKCS7UnpaddingContext(padding.PaddingContext): - def __init__(self, block_size: int) -> None: ... - def update(self, data: Buffer) -> bytes: ... - def finalize(self) -> bytes: ... - -class ANSIX923UnpaddingContext(padding.PaddingContext): - def __init__(self, block_size: int) -> None: ... - def update(self, data: Buffer) -> bytes: ... - def finalize(self) -> bytes: ... - -class ObjectIdentifier: - def __init__(self, value: str) -> None: ... - @property - def dotted_string(self) -> str: ... - @property - def _name(self) -> str: ... - -T = typing.TypeVar("T") diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/_openssl.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/_openssl.pyi deleted file mode 100644 index 80100082acd30b6bc69e5b8ab0972d77a33dc11a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/_openssl.pyi +++ /dev/null @@ -1,8 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import typing - -lib = typing.Any -ffi = typing.Any diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/asn1.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/asn1.pyi deleted file mode 100644 index 3b5f208ecf09aa661854e6983507b0e7b1ef15cb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/asn1.pyi +++ /dev/null @@ -1,7 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -def decode_dss_signature(signature: bytes) -> tuple[int, int]: ... -def encode_dss_signature(r: int, s: int) -> bytes: ... -def parse_spki_for_data(data: bytes) -> bytes: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/exceptions.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/exceptions.pyi deleted file mode 100644 index 09f46b1e817f806873c4a4bdd7b8466046b5f64d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/exceptions.pyi +++ /dev/null @@ -1,17 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -class _Reasons: - BACKEND_MISSING_INTERFACE: _Reasons - UNSUPPORTED_HASH: _Reasons - UNSUPPORTED_CIPHER: _Reasons - UNSUPPORTED_PADDING: _Reasons - UNSUPPORTED_MGF: _Reasons - UNSUPPORTED_PUBLIC_KEY_ALGORITHM: _Reasons - UNSUPPORTED_ELLIPTIC_CURVE: _Reasons - UNSUPPORTED_SERIALIZATION: _Reasons - UNSUPPORTED_X509: _Reasons - UNSUPPORTED_EXCHANGE_ALGORITHM: _Reasons - UNSUPPORTED_DIFFIE_HELLMAN: _Reasons - UNSUPPORTED_MAC: _Reasons diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/ocsp.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/ocsp.pyi deleted file mode 100644 index 103e96c1f117ed06924569fadb6a2b680fbd5456..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/ocsp.pyi +++ /dev/null @@ -1,117 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import datetime -from collections.abc import Iterator - -from cryptography import x509 -from cryptography.hazmat.primitives import hashes, serialization -from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes -from cryptography.x509 import ocsp - -class OCSPRequest: - @property - def issuer_key_hash(self) -> bytes: ... - @property - def issuer_name_hash(self) -> bytes: ... - @property - def hash_algorithm(self) -> hashes.HashAlgorithm: ... - @property - def serial_number(self) -> int: ... - def public_bytes(self, encoding: serialization.Encoding) -> bytes: ... - @property - def extensions(self) -> x509.Extensions: ... - -class OCSPResponse: - @property - def responses(self) -> Iterator[OCSPSingleResponse]: ... - @property - def response_status(self) -> ocsp.OCSPResponseStatus: ... - @property - def signature_algorithm_oid(self) -> x509.ObjectIdentifier: ... - @property - def signature_hash_algorithm( - self, - ) -> hashes.HashAlgorithm | None: ... - @property - def signature(self) -> bytes: ... - @property - def tbs_response_bytes(self) -> bytes: ... - @property - def certificates(self) -> list[x509.Certificate]: ... - @property - def responder_key_hash(self) -> bytes | None: ... - @property - def responder_name(self) -> x509.Name | None: ... - @property - def produced_at(self) -> datetime.datetime: ... - @property - def produced_at_utc(self) -> datetime.datetime: ... - @property - def certificate_status(self) -> ocsp.OCSPCertStatus: ... - @property - def revocation_time(self) -> datetime.datetime | None: ... - @property - def revocation_time_utc(self) -> datetime.datetime | None: ... - @property - def revocation_reason(self) -> x509.ReasonFlags | None: ... - @property - def this_update(self) -> datetime.datetime: ... - @property - def this_update_utc(self) -> datetime.datetime: ... - @property - def next_update(self) -> datetime.datetime | None: ... - @property - def next_update_utc(self) -> datetime.datetime | None: ... - @property - def issuer_key_hash(self) -> bytes: ... - @property - def issuer_name_hash(self) -> bytes: ... - @property - def hash_algorithm(self) -> hashes.HashAlgorithm: ... - @property - def serial_number(self) -> int: ... - @property - def extensions(self) -> x509.Extensions: ... - @property - def single_extensions(self) -> x509.Extensions: ... - def public_bytes(self, encoding: serialization.Encoding) -> bytes: ... - -class OCSPSingleResponse: - @property - def certificate_status(self) -> ocsp.OCSPCertStatus: ... - @property - def revocation_time(self) -> datetime.datetime | None: ... - @property - def revocation_time_utc(self) -> datetime.datetime | None: ... - @property - def revocation_reason(self) -> x509.ReasonFlags | None: ... - @property - def this_update(self) -> datetime.datetime: ... - @property - def this_update_utc(self) -> datetime.datetime: ... - @property - def next_update(self) -> datetime.datetime | None: ... - @property - def next_update_utc(self) -> datetime.datetime | None: ... - @property - def issuer_key_hash(self) -> bytes: ... - @property - def issuer_name_hash(self) -> bytes: ... - @property - def hash_algorithm(self) -> hashes.HashAlgorithm: ... - @property - def serial_number(self) -> int: ... - -def load_der_ocsp_request(data: bytes) -> ocsp.OCSPRequest: ... -def load_der_ocsp_response(data: bytes) -> ocsp.OCSPResponse: ... -def create_ocsp_request( - builder: ocsp.OCSPRequestBuilder, -) -> ocsp.OCSPRequest: ... -def create_ocsp_response( - status: ocsp.OCSPResponseStatus, - builder: ocsp.OCSPResponseBuilder | None, - private_key: PrivateKeyTypes | None, - hash_algorithm: hashes.HashAlgorithm | None, -) -> ocsp.OCSPResponse: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/__init__.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/__init__.pyi deleted file mode 100644 index 5fb3cb2403c0cd9b39059f0030036037f8ae0094..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/__init__.pyi +++ /dev/null @@ -1,75 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import typing - -from cryptography.hazmat.bindings._rust.openssl import ( - aead, - ciphers, - cmac, - dh, - dsa, - ec, - ed448, - ed25519, - hashes, - hmac, - kdf, - keys, - poly1305, - rsa, - x448, - x25519, -) - -__all__ = [ - "aead", - "ciphers", - "cmac", - "dh", - "dsa", - "ec", - "ed448", - "ed25519", - "hashes", - "hmac", - "kdf", - "keys", - "openssl_version", - "openssl_version_text", - "poly1305", - "raise_openssl_error", - "rsa", - "x448", - "x25519", -] - -CRYPTOGRAPHY_IS_LIBRESSL: bool -CRYPTOGRAPHY_IS_BORINGSSL: bool -CRYPTOGRAPHY_IS_AWSLC: bool -CRYPTOGRAPHY_OPENSSL_300_OR_GREATER: bool -CRYPTOGRAPHY_OPENSSL_309_OR_GREATER: bool -CRYPTOGRAPHY_OPENSSL_320_OR_GREATER: bool -CRYPTOGRAPHY_OPENSSL_330_OR_GREATER: bool -CRYPTOGRAPHY_OPENSSL_350_OR_GREATER: bool - -class Providers: ... - -_legacy_provider_loaded: bool -_providers: Providers - -def openssl_version() -> int: ... -def openssl_version_text() -> str: ... -def raise_openssl_error() -> typing.NoReturn: ... -def capture_error_stack() -> list[OpenSSLError]: ... -def is_fips_enabled() -> bool: ... -def enable_fips(providers: Providers) -> None: ... - -class OpenSSLError: - @property - def lib(self) -> int: ... - @property - def reason(self) -> int: ... - @property - def reason_text(self) -> bytes: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/aead.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/aead.pyi deleted file mode 100644 index 831fcd1469d563a87b7db8db50a1bbdb590f866f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/aead.pyi +++ /dev/null @@ -1,107 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from collections.abc import Sequence - -from cryptography.utils import Buffer - -class AESGCM: - def __init__(self, key: Buffer) -> None: ... - @staticmethod - def generate_key(bit_length: int) -> bytes: ... - def encrypt( - self, - nonce: Buffer, - data: Buffer, - associated_data: Buffer | None, - ) -> bytes: ... - def decrypt( - self, - nonce: Buffer, - data: Buffer, - associated_data: Buffer | None, - ) -> bytes: ... - -class ChaCha20Poly1305: - def __init__(self, key: Buffer) -> None: ... - @staticmethod - def generate_key() -> bytes: ... - def encrypt( - self, - nonce: Buffer, - data: Buffer, - associated_data: Buffer | None, - ) -> bytes: ... - def decrypt( - self, - nonce: Buffer, - data: Buffer, - associated_data: Buffer | None, - ) -> bytes: ... - -class AESCCM: - def __init__(self, key: Buffer, tag_length: int = 16) -> None: ... - @staticmethod - def generate_key(bit_length: int) -> bytes: ... - def encrypt( - self, - nonce: Buffer, - data: Buffer, - associated_data: Buffer | None, - ) -> bytes: ... - def decrypt( - self, - nonce: Buffer, - data: Buffer, - associated_data: Buffer | None, - ) -> bytes: ... - -class AESSIV: - def __init__(self, key: Buffer) -> None: ... - @staticmethod - def generate_key(bit_length: int) -> bytes: ... - def encrypt( - self, - data: Buffer, - associated_data: Sequence[Buffer] | None, - ) -> bytes: ... - def decrypt( - self, - data: Buffer, - associated_data: Sequence[Buffer] | None, - ) -> bytes: ... - -class AESOCB3: - def __init__(self, key: Buffer) -> None: ... - @staticmethod - def generate_key(bit_length: int) -> bytes: ... - def encrypt( - self, - nonce: Buffer, - data: Buffer, - associated_data: Buffer | None, - ) -> bytes: ... - def decrypt( - self, - nonce: Buffer, - data: Buffer, - associated_data: Buffer | None, - ) -> bytes: ... - -class AESGCMSIV: - def __init__(self, key: Buffer) -> None: ... - @staticmethod - def generate_key(bit_length: int) -> bytes: ... - def encrypt( - self, - nonce: Buffer, - data: Buffer, - associated_data: Buffer | None, - ) -> bytes: ... - def decrypt( - self, - nonce: Buffer, - data: Buffer, - associated_data: Buffer | None, - ) -> bytes: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/ciphers.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/ciphers.pyi deleted file mode 100644 index a48fb017ff562b3d26fa849ebf0c0f9dafb36aaf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/ciphers.pyi +++ /dev/null @@ -1,38 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import typing - -from cryptography.hazmat.primitives import ciphers -from cryptography.hazmat.primitives.ciphers import modes - -@typing.overload -def create_encryption_ctx( - algorithm: ciphers.CipherAlgorithm, mode: modes.ModeWithAuthenticationTag -) -> ciphers.AEADEncryptionContext: ... -@typing.overload -def create_encryption_ctx( - algorithm: ciphers.CipherAlgorithm, mode: modes.Mode | None -) -> ciphers.CipherContext: ... -@typing.overload -def create_decryption_ctx( - algorithm: ciphers.CipherAlgorithm, mode: modes.ModeWithAuthenticationTag -) -> ciphers.AEADDecryptionContext: ... -@typing.overload -def create_decryption_ctx( - algorithm: ciphers.CipherAlgorithm, mode: modes.Mode | None -) -> ciphers.CipherContext: ... -def cipher_supported( - algorithm: ciphers.CipherAlgorithm, mode: modes.Mode -) -> bool: ... -def _advance( - ctx: ciphers.AEADEncryptionContext | ciphers.AEADDecryptionContext, n: int -) -> None: ... -def _advance_aad( - ctx: ciphers.AEADEncryptionContext | ciphers.AEADDecryptionContext, n: int -) -> None: ... - -class CipherContext: ... -class AEADEncryptionContext: ... -class AEADDecryptionContext: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/cmac.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/cmac.pyi deleted file mode 100644 index 9c03508bc89bf087e24c684802eeec344d76bb9e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/cmac.pyi +++ /dev/null @@ -1,18 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import typing - -from cryptography.hazmat.primitives import ciphers - -class CMAC: - def __init__( - self, - algorithm: ciphers.BlockCipherAlgorithm, - backend: typing.Any = None, - ) -> None: ... - def update(self, data: bytes) -> None: ... - def finalize(self) -> bytes: ... - def verify(self, signature: bytes) -> None: ... - def copy(self) -> CMAC: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/dh.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/dh.pyi deleted file mode 100644 index 08733d745c3ddc994f358193dd4b79ce359583d8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/dh.pyi +++ /dev/null @@ -1,51 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import typing - -from cryptography.hazmat.primitives.asymmetric import dh - -MIN_MODULUS_SIZE: int - -class DHPrivateKey: ... -class DHPublicKey: ... -class DHParameters: ... - -class DHPrivateNumbers: - def __init__(self, x: int, public_numbers: DHPublicNumbers) -> None: ... - def private_key(self, backend: typing.Any = None) -> dh.DHPrivateKey: ... - @property - def x(self) -> int: ... - @property - def public_numbers(self) -> DHPublicNumbers: ... - -class DHPublicNumbers: - def __init__( - self, y: int, parameter_numbers: DHParameterNumbers - ) -> None: ... - def public_key(self, backend: typing.Any = None) -> dh.DHPublicKey: ... - @property - def y(self) -> int: ... - @property - def parameter_numbers(self) -> DHParameterNumbers: ... - -class DHParameterNumbers: - def __init__(self, p: int, g: int, q: int | None = None) -> None: ... - def parameters(self, backend: typing.Any = None) -> dh.DHParameters: ... - @property - def p(self) -> int: ... - @property - def g(self) -> int: ... - @property - def q(self) -> int | None: ... - -def generate_parameters( - generator: int, key_size: int, backend: typing.Any = None -) -> dh.DHParameters: ... -def from_pem_parameters( - data: bytes, backend: typing.Any = None -) -> dh.DHParameters: ... -def from_der_parameters( - data: bytes, backend: typing.Any = None -) -> dh.DHParameters: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/dsa.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/dsa.pyi deleted file mode 100644 index 0922a4c4041a90b1e884f56522f82bced398c80e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/dsa.pyi +++ /dev/null @@ -1,41 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import typing - -from cryptography.hazmat.primitives.asymmetric import dsa - -class DSAPrivateKey: ... -class DSAPublicKey: ... -class DSAParameters: ... - -class DSAPrivateNumbers: - def __init__(self, x: int, public_numbers: DSAPublicNumbers) -> None: ... - @property - def x(self) -> int: ... - @property - def public_numbers(self) -> DSAPublicNumbers: ... - def private_key(self, backend: typing.Any = None) -> dsa.DSAPrivateKey: ... - -class DSAPublicNumbers: - def __init__( - self, y: int, parameter_numbers: DSAParameterNumbers - ) -> None: ... - @property - def y(self) -> int: ... - @property - def parameter_numbers(self) -> DSAParameterNumbers: ... - def public_key(self, backend: typing.Any = None) -> dsa.DSAPublicKey: ... - -class DSAParameterNumbers: - def __init__(self, p: int, q: int, g: int) -> None: ... - @property - def p(self) -> int: ... - @property - def q(self) -> int: ... - @property - def g(self) -> int: ... - def parameters(self, backend: typing.Any = None) -> dsa.DSAParameters: ... - -def generate_parameters(key_size: int) -> dsa.DSAParameters: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/ec.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/ec.pyi deleted file mode 100644 index 5c3b7bf6e4a9bafac1423bd3936ebc347885a950..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/ec.pyi +++ /dev/null @@ -1,52 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import typing - -from cryptography.hazmat.primitives.asymmetric import ec - -class ECPrivateKey: ... -class ECPublicKey: ... - -class EllipticCurvePrivateNumbers: - def __init__( - self, private_value: int, public_numbers: EllipticCurvePublicNumbers - ) -> None: ... - def private_key( - self, backend: typing.Any = None - ) -> ec.EllipticCurvePrivateKey: ... - @property - def private_value(self) -> int: ... - @property - def public_numbers(self) -> EllipticCurvePublicNumbers: ... - -class EllipticCurvePublicNumbers: - def __init__(self, x: int, y: int, curve: ec.EllipticCurve) -> None: ... - def public_key( - self, backend: typing.Any = None - ) -> ec.EllipticCurvePublicKey: ... - @property - def x(self) -> int: ... - @property - def y(self) -> int: ... - @property - def curve(self) -> ec.EllipticCurve: ... - def __eq__(self, other: object) -> bool: ... - -def curve_supported(curve: ec.EllipticCurve) -> bool: ... -def generate_private_key( - curve: ec.EllipticCurve, backend: typing.Any = None -) -> ec.EllipticCurvePrivateKey: ... -def from_private_numbers( - numbers: ec.EllipticCurvePrivateNumbers, -) -> ec.EllipticCurvePrivateKey: ... -def from_public_numbers( - numbers: ec.EllipticCurvePublicNumbers, -) -> ec.EllipticCurvePublicKey: ... -def from_public_bytes( - curve: ec.EllipticCurve, data: bytes -) -> ec.EllipticCurvePublicKey: ... -def derive_private_key( - private_value: int, curve: ec.EllipticCurve -) -> ec.EllipticCurvePrivateKey: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/ed25519.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/ed25519.pyi deleted file mode 100644 index f85b3d1b2361832095a00ba9b67675ad365339dd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/ed25519.pyi +++ /dev/null @@ -1,13 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from cryptography.hazmat.primitives.asymmetric import ed25519 -from cryptography.utils import Buffer - -class Ed25519PrivateKey: ... -class Ed25519PublicKey: ... - -def generate_key() -> ed25519.Ed25519PrivateKey: ... -def from_private_bytes(data: Buffer) -> ed25519.Ed25519PrivateKey: ... -def from_public_bytes(data: bytes) -> ed25519.Ed25519PublicKey: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/ed448.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/ed448.pyi deleted file mode 100644 index c8ca0ecb156b9d9a7e69cf10bf71f187362f3315..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/ed448.pyi +++ /dev/null @@ -1,13 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from cryptography.hazmat.primitives.asymmetric import ed448 -from cryptography.utils import Buffer - -class Ed448PrivateKey: ... -class Ed448PublicKey: ... - -def generate_key() -> ed448.Ed448PrivateKey: ... -def from_private_bytes(data: Buffer) -> ed448.Ed448PrivateKey: ... -def from_public_bytes(data: bytes) -> ed448.Ed448PublicKey: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/hashes.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/hashes.pyi deleted file mode 100644 index 6bfd295af88964a4a09c59ecafcce3d0890ff350..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/hashes.pyi +++ /dev/null @@ -1,28 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import typing - -from cryptography.hazmat.primitives import hashes -from cryptography.utils import Buffer - -class Hash(hashes.HashContext): - def __init__( - self, algorithm: hashes.HashAlgorithm, backend: typing.Any = None - ) -> None: ... - @property - def algorithm(self) -> hashes.HashAlgorithm: ... - def update(self, data: Buffer) -> None: ... - def finalize(self) -> bytes: ... - def copy(self) -> Hash: ... - -def hash_supported(algorithm: hashes.HashAlgorithm) -> bool: ... - -class XOFHash: - def __init__(self, algorithm: hashes.ExtendableOutputFunction) -> None: ... - @property - def algorithm(self) -> hashes.ExtendableOutputFunction: ... - def update(self, data: Buffer) -> None: ... - def squeeze(self, length: int) -> bytes: ... - def copy(self) -> XOFHash: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/hmac.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/hmac.pyi deleted file mode 100644 index 3883d1b1a920b0f118a6e9c6084f9532a5484bec..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/hmac.pyi +++ /dev/null @@ -1,22 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import typing - -from cryptography.hazmat.primitives import hashes -from cryptography.utils import Buffer - -class HMAC(hashes.HashContext): - def __init__( - self, - key: Buffer, - algorithm: hashes.HashAlgorithm, - backend: typing.Any = None, - ) -> None: ... - @property - def algorithm(self) -> hashes.HashAlgorithm: ... - def update(self, data: Buffer) -> None: ... - def finalize(self) -> bytes: ... - def verify(self, signature: bytes) -> None: ... - def copy(self) -> HMAC: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/kdf.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/kdf.pyi deleted file mode 100644 index 9979e42db66103f2b6cbb4f0308547c707f4bd32..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/kdf.pyi +++ /dev/null @@ -1,49 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import typing - -from cryptography.hazmat.primitives.hashes import HashAlgorithm -from cryptography.utils import Buffer - -def derive_pbkdf2_hmac( - key_material: Buffer, - algorithm: HashAlgorithm, - salt: bytes, - iterations: int, - length: int, -) -> bytes: ... - -class Scrypt: - def __init__( - self, - salt: bytes, - length: int, - n: int, - r: int, - p: int, - backend: typing.Any = None, - ) -> None: ... - def derive(self, key_material: Buffer) -> bytes: ... - def verify(self, key_material: bytes, expected_key: bytes) -> None: ... - -class Argon2id: - def __init__( - self, - *, - salt: bytes, - length: int, - iterations: int, - lanes: int, - memory_cost: int, - ad: bytes | None = None, - secret: bytes | None = None, - ) -> None: ... - def derive(self, key_material: bytes) -> bytes: ... - def verify(self, key_material: bytes, expected_key: bytes) -> None: ... - def derive_phc_encoded(self, key_material: bytes) -> str: ... - @classmethod - def verify_phc_encoded( - cls, key_material: bytes, phc_encoded: str, secret: bytes | None = None - ) -> None: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/keys.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/keys.pyi deleted file mode 100644 index 404057e037408dd1979afc482768430b8b478646..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/keys.pyi +++ /dev/null @@ -1,34 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import typing - -from cryptography.hazmat.primitives.asymmetric.types import ( - PrivateKeyTypes, - PublicKeyTypes, -) -from cryptography.utils import Buffer - -def load_der_private_key( - data: Buffer, - password: bytes | None, - backend: typing.Any = None, - *, - unsafe_skip_rsa_key_validation: bool = False, -) -> PrivateKeyTypes: ... -def load_pem_private_key( - data: Buffer, - password: bytes | None, - backend: typing.Any = None, - *, - unsafe_skip_rsa_key_validation: bool = False, -) -> PrivateKeyTypes: ... -def load_der_public_key( - data: bytes, - backend: typing.Any = None, -) -> PublicKeyTypes: ... -def load_pem_public_key( - data: bytes, - backend: typing.Any = None, -) -> PublicKeyTypes: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/poly1305.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/poly1305.pyi deleted file mode 100644 index 45a2a39f689071b5d544d8495bfd5bd57d4ab6e4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/poly1305.pyi +++ /dev/null @@ -1,15 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from cryptography.utils import Buffer - -class Poly1305: - def __init__(self, key: Buffer) -> None: ... - @staticmethod - def generate_tag(key: Buffer, data: Buffer) -> bytes: ... - @staticmethod - def verify_tag(key: Buffer, data: Buffer, tag: bytes) -> None: ... - def update(self, data: Buffer) -> None: ... - def finalize(self) -> bytes: ... - def verify(self, tag: bytes) -> None: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/rsa.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/rsa.pyi deleted file mode 100644 index ef7752ddb79db62489bcf71df933fdefd83c0f9d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/rsa.pyi +++ /dev/null @@ -1,55 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import typing - -from cryptography.hazmat.primitives.asymmetric import rsa - -class RSAPrivateKey: ... -class RSAPublicKey: ... - -class RSAPrivateNumbers: - def __init__( - self, - p: int, - q: int, - d: int, - dmp1: int, - dmq1: int, - iqmp: int, - public_numbers: RSAPublicNumbers, - ) -> None: ... - @property - def p(self) -> int: ... - @property - def q(self) -> int: ... - @property - def d(self) -> int: ... - @property - def dmp1(self) -> int: ... - @property - def dmq1(self) -> int: ... - @property - def iqmp(self) -> int: ... - @property - def public_numbers(self) -> RSAPublicNumbers: ... - def private_key( - self, - backend: typing.Any = None, - *, - unsafe_skip_rsa_key_validation: bool = False, - ) -> rsa.RSAPrivateKey: ... - -class RSAPublicNumbers: - def __init__(self, e: int, n: int) -> None: ... - @property - def n(self) -> int: ... - @property - def e(self) -> int: ... - def public_key(self, backend: typing.Any = None) -> rsa.RSAPublicKey: ... - -def generate_private_key( - public_exponent: int, - key_size: int, -) -> rsa.RSAPrivateKey: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/x25519.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/x25519.pyi deleted file mode 100644 index 38d2adddb1017030139f2cb2411dc43101a7448e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/x25519.pyi +++ /dev/null @@ -1,13 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from cryptography.hazmat.primitives.asymmetric import x25519 -from cryptography.utils import Buffer - -class X25519PrivateKey: ... -class X25519PublicKey: ... - -def generate_key() -> x25519.X25519PrivateKey: ... -def from_private_bytes(data: Buffer) -> x25519.X25519PrivateKey: ... -def from_public_bytes(data: bytes) -> x25519.X25519PublicKey: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/x448.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/x448.pyi deleted file mode 100644 index 3ac098091af514639889486723710f851cd93e6b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/openssl/x448.pyi +++ /dev/null @@ -1,13 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from cryptography.hazmat.primitives.asymmetric import x448 -from cryptography.utils import Buffer - -class X448PrivateKey: ... -class X448PublicKey: ... - -def generate_key() -> x448.X448PrivateKey: ... -def from_private_bytes(data: Buffer) -> x448.X448PrivateKey: ... -def from_public_bytes(data: bytes) -> x448.X448PublicKey: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/pkcs12.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/pkcs12.pyi deleted file mode 100644 index b25becb6bdef7ce259dcb40f34aa9386774437ca..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/pkcs12.pyi +++ /dev/null @@ -1,52 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import typing -from collections.abc import Iterable - -from cryptography import x509 -from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes -from cryptography.hazmat.primitives.serialization import ( - KeySerializationEncryption, -) -from cryptography.hazmat.primitives.serialization.pkcs12 import ( - PKCS12KeyAndCertificates, - PKCS12PrivateKeyTypes, -) -from cryptography.utils import Buffer - -class PKCS12Certificate: - def __init__( - self, cert: x509.Certificate, friendly_name: bytes | None - ) -> None: ... - @property - def friendly_name(self) -> bytes | None: ... - @property - def certificate(self) -> x509.Certificate: ... - -def load_key_and_certificates( - data: Buffer, - password: Buffer | None, - backend: typing.Any = None, -) -> tuple[ - PrivateKeyTypes | None, - x509.Certificate | None, - list[x509.Certificate], -]: ... -def load_pkcs12( - data: bytes, - password: bytes | None, - backend: typing.Any = None, -) -> PKCS12KeyAndCertificates: ... -def serialize_java_truststore( - certs: Iterable[PKCS12Certificate], - encryption_algorithm: KeySerializationEncryption, -) -> bytes: ... -def serialize_key_and_certificates( - name: bytes | None, - key: PKCS12PrivateKeyTypes | None, - cert: x509.Certificate | None, - cas: Iterable[x509.Certificate | PKCS12Certificate] | None, - encryption_algorithm: KeySerializationEncryption, -) -> bytes: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/pkcs7.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/pkcs7.pyi deleted file mode 100644 index 358b135865a89a6675706e0fd880d3eabdb081b5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/pkcs7.pyi +++ /dev/null @@ -1,50 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from collections.abc import Iterable - -from cryptography import x509 -from cryptography.hazmat.primitives import serialization -from cryptography.hazmat.primitives.asymmetric import rsa -from cryptography.hazmat.primitives.serialization import pkcs7 - -def serialize_certificates( - certs: list[x509.Certificate], - encoding: serialization.Encoding, -) -> bytes: ... -def encrypt_and_serialize( - builder: pkcs7.PKCS7EnvelopeBuilder, - content_encryption_algorithm: pkcs7.ContentEncryptionAlgorithm, - encoding: serialization.Encoding, - options: Iterable[pkcs7.PKCS7Options], -) -> bytes: ... -def sign_and_serialize( - builder: pkcs7.PKCS7SignatureBuilder, - encoding: serialization.Encoding, - options: Iterable[pkcs7.PKCS7Options], -) -> bytes: ... -def decrypt_der( - data: bytes, - certificate: x509.Certificate, - private_key: rsa.RSAPrivateKey, - options: Iterable[pkcs7.PKCS7Options], -) -> bytes: ... -def decrypt_pem( - data: bytes, - certificate: x509.Certificate, - private_key: rsa.RSAPrivateKey, - options: Iterable[pkcs7.PKCS7Options], -) -> bytes: ... -def decrypt_smime( - data: bytes, - certificate: x509.Certificate, - private_key: rsa.RSAPrivateKey, - options: Iterable[pkcs7.PKCS7Options], -) -> bytes: ... -def load_pem_pkcs7_certificates( - data: bytes, -) -> list[x509.Certificate]: ... -def load_der_pkcs7_certificates( - data: bytes, -) -> list[x509.Certificate]: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/test_support.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/test_support.pyi deleted file mode 100644 index c6c6d0bbe12916b4299d08398493c5da8f762ae7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/test_support.pyi +++ /dev/null @@ -1,23 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from cryptography import x509 -from cryptography.hazmat.primitives import serialization -from cryptography.hazmat.primitives.serialization import pkcs7 -from cryptography.utils import Buffer - -class TestCertificate: - not_after_tag: int - not_before_tag: int - issuer_value_tags: list[int] - subject_value_tags: list[int] - -def test_parse_certificate(data: bytes) -> TestCertificate: ... -def pkcs7_verify( - encoding: serialization.Encoding, - sig: bytes, - msg: Buffer | None, - certs: list[x509.Certificate], - options: list[pkcs7.PKCS7Options], -) -> None: ... diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/x509.pyi b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/x509.pyi deleted file mode 100644 index 6ba9afb3da57c620864537acff36b7e1beadbdb3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/_rust/x509.pyi +++ /dev/null @@ -1,313 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -import datetime -import typing -from collections.abc import Iterator - -from cryptography import x509 -from cryptography.hazmat.primitives import hashes, serialization -from cryptography.hazmat.primitives.asymmetric.ec import ECDSA -from cryptography.hazmat.primitives.asymmetric.padding import PSS, PKCS1v15 -from cryptography.hazmat.primitives.asymmetric.types import ( - CertificateIssuerPublicKeyTypes, - CertificatePublicKeyTypes, - PrivateKeyTypes, -) -from cryptography.x509 import certificate_transparency - -def load_pem_x509_certificate( - data: bytes, backend: typing.Any = None -) -> x509.Certificate: ... -def load_der_x509_certificate( - data: bytes, backend: typing.Any = None -) -> x509.Certificate: ... -def load_pem_x509_certificates( - data: bytes, -) -> list[x509.Certificate]: ... -def load_pem_x509_crl( - data: bytes, backend: typing.Any = None -) -> x509.CertificateRevocationList: ... -def load_der_x509_crl( - data: bytes, backend: typing.Any = None -) -> x509.CertificateRevocationList: ... -def load_pem_x509_csr( - data: bytes, backend: typing.Any = None -) -> x509.CertificateSigningRequest: ... -def load_der_x509_csr( - data: bytes, backend: typing.Any = None -) -> x509.CertificateSigningRequest: ... -def encode_name_bytes(name: x509.Name) -> bytes: ... -def encode_extension_value(extension: x509.ExtensionType) -> bytes: ... -def create_x509_certificate( - builder: x509.CertificateBuilder, - private_key: PrivateKeyTypes, - hash_algorithm: hashes.HashAlgorithm | None, - rsa_padding: PKCS1v15 | PSS | None, - ecdsa_deterministic: bool | None, -) -> x509.Certificate: ... -def create_x509_csr( - builder: x509.CertificateSigningRequestBuilder, - private_key: PrivateKeyTypes, - hash_algorithm: hashes.HashAlgorithm | None, - rsa_padding: PKCS1v15 | PSS | None, - ecdsa_deterministic: bool | None, -) -> x509.CertificateSigningRequest: ... -def create_x509_crl( - builder: x509.CertificateRevocationListBuilder, - private_key: PrivateKeyTypes, - hash_algorithm: hashes.HashAlgorithm | None, - rsa_padding: PKCS1v15 | PSS | None, - ecdsa_deterministic: bool | None, -) -> x509.CertificateRevocationList: ... - -class Sct: - @property - def version(self) -> certificate_transparency.Version: ... - @property - def log_id(self) -> bytes: ... - @property - def timestamp(self) -> datetime.datetime: ... - @property - def entry_type(self) -> certificate_transparency.LogEntryType: ... - @property - def signature_hash_algorithm(self) -> hashes.HashAlgorithm: ... - @property - def signature_algorithm( - self, - ) -> certificate_transparency.SignatureAlgorithm: ... - @property - def signature(self) -> bytes: ... - @property - def extension_bytes(self) -> bytes: ... - -class Certificate: - def fingerprint(self, algorithm: hashes.HashAlgorithm) -> bytes: ... - @property - def serial_number(self) -> int: ... - @property - def version(self) -> x509.Version: ... - def public_key(self) -> CertificatePublicKeyTypes: ... - @property - def public_key_algorithm_oid(self) -> x509.ObjectIdentifier: ... - @property - def not_valid_before(self) -> datetime.datetime: ... - @property - def not_valid_before_utc(self) -> datetime.datetime: ... - @property - def not_valid_after(self) -> datetime.datetime: ... - @property - def not_valid_after_utc(self) -> datetime.datetime: ... - @property - def issuer(self) -> x509.Name: ... - @property - def subject(self) -> x509.Name: ... - @property - def signature_hash_algorithm( - self, - ) -> hashes.HashAlgorithm | None: ... - @property - def signature_algorithm_oid(self) -> x509.ObjectIdentifier: ... - @property - def signature_algorithm_parameters( - self, - ) -> PSS | PKCS1v15 | ECDSA | None: ... - @property - def extensions(self) -> x509.Extensions: ... - @property - def signature(self) -> bytes: ... - @property - def tbs_certificate_bytes(self) -> bytes: ... - @property - def tbs_precertificate_bytes(self) -> bytes: ... - def __eq__(self, other: object) -> bool: ... - def __hash__(self) -> int: ... - def public_bytes(self, encoding: serialization.Encoding) -> bytes: ... - def verify_directly_issued_by(self, issuer: Certificate) -> None: ... - -class RevokedCertificate: ... - -class CertificateRevocationList: - def public_bytes(self, encoding: serialization.Encoding) -> bytes: ... - def fingerprint(self, algorithm: hashes.HashAlgorithm) -> bytes: ... - def get_revoked_certificate_by_serial_number( - self, serial_number: int - ) -> x509.RevokedCertificate | None: ... - @property - def signature_hash_algorithm( - self, - ) -> hashes.HashAlgorithm | None: ... - @property - def signature_algorithm_oid(self) -> x509.ObjectIdentifier: ... - @property - def signature_algorithm_parameters( - self, - ) -> PSS | PKCS1v15 | ECDSA | None: ... - @property - def issuer(self) -> x509.Name: ... - @property - def next_update(self) -> datetime.datetime | None: ... - @property - def next_update_utc(self) -> datetime.datetime | None: ... - @property - def last_update(self) -> datetime.datetime: ... - @property - def last_update_utc(self) -> datetime.datetime: ... - @property - def extensions(self) -> x509.Extensions: ... - @property - def signature(self) -> bytes: ... - @property - def tbs_certlist_bytes(self) -> bytes: ... - def __eq__(self, other: object) -> bool: ... - def __len__(self) -> int: ... - @typing.overload - def __getitem__(self, idx: int) -> x509.RevokedCertificate: ... - @typing.overload - def __getitem__(self, idx: slice) -> list[x509.RevokedCertificate]: ... - def __iter__(self) -> Iterator[x509.RevokedCertificate]: ... - def is_signature_valid( - self, public_key: CertificateIssuerPublicKeyTypes - ) -> bool: ... - -class CertificateSigningRequest: - def __eq__(self, other: object) -> bool: ... - def __hash__(self) -> int: ... - def public_key(self) -> CertificatePublicKeyTypes: ... - @property - def subject(self) -> x509.Name: ... - @property - def signature_hash_algorithm( - self, - ) -> hashes.HashAlgorithm | None: ... - @property - def signature_algorithm_oid(self) -> x509.ObjectIdentifier: ... - @property - def signature_algorithm_parameters( - self, - ) -> PSS | PKCS1v15 | ECDSA | None: ... - @property - def extensions(self) -> x509.Extensions: ... - @property - def attributes(self) -> x509.Attributes: ... - def public_bytes(self, encoding: serialization.Encoding) -> bytes: ... - @property - def signature(self) -> bytes: ... - @property - def tbs_certrequest_bytes(self) -> bytes: ... - @property - def is_signature_valid(self) -> bool: ... - def get_attribute_for_oid(self, oid: x509.ObjectIdentifier) -> bytes: ... - -class PolicyBuilder: - def time(self, time: datetime.datetime) -> PolicyBuilder: ... - def store(self, store: Store) -> PolicyBuilder: ... - def max_chain_depth(self, max_chain_depth: int) -> PolicyBuilder: ... - def extension_policies( - self, *, ca_policy: ExtensionPolicy, ee_policy: ExtensionPolicy - ) -> PolicyBuilder: ... - def build_client_verifier(self) -> ClientVerifier: ... - def build_server_verifier( - self, subject: x509.verification.Subject - ) -> ServerVerifier: ... - -class Policy: - @property - def max_chain_depth(self) -> int: ... - @property - def subject(self) -> x509.verification.Subject | None: ... - @property - def validation_time(self) -> datetime.datetime: ... - @property - def extended_key_usage(self) -> x509.ObjectIdentifier: ... - @property - def minimum_rsa_modulus(self) -> int: ... - -class Criticality: - CRITICAL: Criticality - AGNOSTIC: Criticality - NON_CRITICAL: Criticality - -T = typing.TypeVar("T", contravariant=True, bound=x509.ExtensionType) - -MaybeExtensionValidatorCallback = typing.Callable[ - [ - Policy, - x509.Certificate, - T | None, - ], - None, -] - -PresentExtensionValidatorCallback = typing.Callable[ - [Policy, x509.Certificate, T], - None, -] - -class ExtensionPolicy: - @staticmethod - def permit_all() -> ExtensionPolicy: ... - @staticmethod - def webpki_defaults_ca() -> ExtensionPolicy: ... - @staticmethod - def webpki_defaults_ee() -> ExtensionPolicy: ... - def require_not_present( - self, extension_type: type[x509.ExtensionType] - ) -> ExtensionPolicy: ... - def may_be_present( - self, - extension_type: type[T], - criticality: Criticality, - validator: MaybeExtensionValidatorCallback[T] | None, - ) -> ExtensionPolicy: ... - def require_present( - self, - extension_type: type[T], - criticality: Criticality, - validator: PresentExtensionValidatorCallback[T] | None, - ) -> ExtensionPolicy: ... - -class VerifiedClient: - @property - def subjects(self) -> list[x509.GeneralName] | None: ... - @property - def chain(self) -> list[x509.Certificate]: ... - -class ClientVerifier: - @property - def policy(self) -> Policy: ... - @property - def validation_time(self) -> datetime.datetime: ... - @property - def max_chain_depth(self) -> int: ... - @property - def store(self) -> Store: ... - def verify( - self, - leaf: x509.Certificate, - intermediates: list[x509.Certificate], - ) -> VerifiedClient: ... - -class ServerVerifier: - @property - def policy(self) -> Policy: ... - @property - def subject(self) -> x509.verification.Subject: ... - @property - def validation_time(self) -> datetime.datetime: ... - @property - def max_chain_depth(self) -> int: ... - @property - def store(self) -> Store: ... - def verify( - self, - leaf: x509.Certificate, - intermediates: list[x509.Certificate], - ) -> list[x509.Certificate]: ... - -class Store: - def __init__(self, certs: list[x509.Certificate]) -> None: ... - -class VerificationError(Exception): - pass diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/__init__.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/__init__.py deleted file mode 100644 index b509336233c2fafe4185a49da5909c8bbb38dfd7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 1d1c4498d2c00cd149984204b103006cf02dea1c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/__pycache__/_conditional.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/__pycache__/_conditional.cpython-310.pyc deleted file mode 100644 index 85b706d899ce40ccd6f222050d3e19fb83c7a1ae..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/__pycache__/_conditional.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/__pycache__/binding.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/__pycache__/binding.cpython-310.pyc deleted file mode 100644 index 2be24117577d01ec83119d1f19f3df0434421271..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/__pycache__/binding.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/_conditional.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/_conditional.py deleted file mode 100644 index c90c916dcd08076b98c1d2bab73c7662c8f32b49..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/_conditional.py +++ /dev/null @@ -1,191 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - - -def cryptography_has_set_cert_cb() -> list[str]: - return [ - "SSL_CTX_set_cert_cb", - "SSL_set_cert_cb", - ] - - -def cryptography_has_ssl_st() -> list[str]: - return [ - "SSL_ST_BEFORE", - "SSL_ST_OK", - "SSL_ST_INIT", - "SSL_ST_RENEGOTIATE", - ] - - -def cryptography_has_tls_st() -> list[str]: - return [ - "TLS_ST_BEFORE", - "TLS_ST_OK", - ] - - -def cryptography_has_ssl_sigalgs() -> list[str]: - return [ - "SSL_CTX_set1_sigalgs_list", - ] - - -def cryptography_has_psk() -> list[str]: - return [ - "SSL_CTX_use_psk_identity_hint", - "SSL_CTX_set_psk_server_callback", - "SSL_CTX_set_psk_client_callback", - ] - - -def cryptography_has_psk_tlsv13() -> list[str]: - return [ - "SSL_CTX_set_psk_find_session_callback", - "SSL_CTX_set_psk_use_session_callback", - "Cryptography_SSL_SESSION_new", - "SSL_CIPHER_find", - "SSL_SESSION_set1_master_key", - "SSL_SESSION_set_cipher", - "SSL_SESSION_set_protocol_version", - ] - - -def cryptography_has_custom_ext() -> list[str]: - return [ - "SSL_CTX_add_client_custom_ext", - "SSL_CTX_add_server_custom_ext", - "SSL_extension_supported", - ] - - -def cryptography_has_tlsv13_functions() -> list[str]: - return [ - "SSL_CTX_set_ciphersuites", - ] - - -def cryptography_has_tlsv13_hs_functions() -> list[str]: - return [ - "SSL_VERIFY_POST_HANDSHAKE", - "SSL_verify_client_post_handshake", - "SSL_CTX_set_post_handshake_auth", - "SSL_set_post_handshake_auth", - "SSL_SESSION_get_max_early_data", - "SSL_write_early_data", - "SSL_read_early_data", - "SSL_CTX_set_max_early_data", - ] - - -def cryptography_has_engine() -> list[str]: - return [ - "ENGINE_by_id", - "ENGINE_init", - "ENGINE_finish", - "ENGINE_get_default_RAND", - "ENGINE_set_default_RAND", - "ENGINE_unregister_RAND", - "ENGINE_ctrl_cmd", - "ENGINE_free", - "ENGINE_get_name", - "ENGINE_ctrl_cmd_string", - "ENGINE_load_builtin_engines", - "ENGINE_load_private_key", - "ENGINE_load_public_key", - "SSL_CTX_set_client_cert_engine", - ] - - -def cryptography_has_verified_chain() -> list[str]: - return [ - "SSL_get0_verified_chain", - ] - - -def cryptography_has_srtp() -> list[str]: - return [ - "SSL_CTX_set_tlsext_use_srtp", - "SSL_set_tlsext_use_srtp", - "SSL_get_selected_srtp_profile", - ] - - -def cryptography_has_op_no_renegotiation() -> list[str]: - return [ - "SSL_OP_NO_RENEGOTIATION", - ] - - -def cryptography_has_dtls_get_data_mtu() -> list[str]: - return [ - "DTLS_get_data_mtu", - ] - - -def cryptography_has_ssl_cookie() -> list[str]: - return [ - "SSL_OP_COOKIE_EXCHANGE", - "DTLSv1_listen", - "SSL_CTX_set_cookie_generate_cb", - "SSL_CTX_set_cookie_verify_cb", - ] - - -def cryptography_has_prime_checks() -> list[str]: - return [ - "BN_prime_checks_for_size", - ] - - -def cryptography_has_unexpected_eof_while_reading() -> list[str]: - return ["SSL_R_UNEXPECTED_EOF_WHILE_READING"] - - -def cryptography_has_ssl_op_ignore_unexpected_eof() -> list[str]: - return [ - "SSL_OP_IGNORE_UNEXPECTED_EOF", - ] - - -def cryptography_has_get_extms_support() -> list[str]: - return ["SSL_get_extms_support"] - - -# This is a mapping of -# {condition: function-returning-names-dependent-on-that-condition} so we can -# loop over them and delete unsupported names at runtime. It will be removed -# when cffi supports #if in cdef. We use functions instead of just a dict of -# lists so we can use coverage to measure which are used. -CONDITIONAL_NAMES = { - "Cryptography_HAS_SET_CERT_CB": cryptography_has_set_cert_cb, - "Cryptography_HAS_SSL_ST": cryptography_has_ssl_st, - "Cryptography_HAS_TLS_ST": cryptography_has_tls_st, - "Cryptography_HAS_SIGALGS": cryptography_has_ssl_sigalgs, - "Cryptography_HAS_PSK": cryptography_has_psk, - "Cryptography_HAS_PSK_TLSv1_3": cryptography_has_psk_tlsv13, - "Cryptography_HAS_CUSTOM_EXT": cryptography_has_custom_ext, - "Cryptography_HAS_TLSv1_3_FUNCTIONS": cryptography_has_tlsv13_functions, - "Cryptography_HAS_TLSv1_3_HS_FUNCTIONS": ( - cryptography_has_tlsv13_hs_functions - ), - "Cryptography_HAS_ENGINE": cryptography_has_engine, - "Cryptography_HAS_VERIFIED_CHAIN": cryptography_has_verified_chain, - "Cryptography_HAS_SRTP": cryptography_has_srtp, - "Cryptography_HAS_OP_NO_RENEGOTIATION": ( - cryptography_has_op_no_renegotiation - ), - "Cryptography_HAS_DTLS_GET_DATA_MTU": cryptography_has_dtls_get_data_mtu, - "Cryptography_HAS_SSL_COOKIE": cryptography_has_ssl_cookie, - "Cryptography_HAS_PRIME_CHECKS": cryptography_has_prime_checks, - "Cryptography_HAS_UNEXPECTED_EOF_WHILE_READING": ( - cryptography_has_unexpected_eof_while_reading - ), - "Cryptography_HAS_SSL_OP_IGNORE_UNEXPECTED_EOF": ( - cryptography_has_ssl_op_ignore_unexpected_eof - ), - "Cryptography_HAS_GET_EXTMS_SUPPORT": cryptography_has_get_extms_support, -} diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/binding.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/binding.py deleted file mode 100644 index 7de2c65356bb3f73dcfa4bd2a572dab6425a4adf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/bindings/openssl/binding.py +++ /dev/null @@ -1,122 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import os -import sys -import threading -import types -import typing -import warnings -from collections.abc import Callable - -import cryptography -from cryptography.exceptions import InternalError -from cryptography.hazmat.bindings._rust import _openssl, openssl -from cryptography.hazmat.bindings.openssl._conditional import CONDITIONAL_NAMES - - -def _openssl_assert(ok: bool) -> None: - if not ok: - errors = openssl.capture_error_stack() - - raise InternalError( - "Unknown OpenSSL error. This error is commonly encountered when " - "another library is not cleaning up the OpenSSL error stack. If " - "you are using cryptography with another library that uses " - "OpenSSL try disabling it before reporting a bug. Otherwise " - "please file an issue at https://github.com/pyca/cryptography/" - "issues with information on how to reproduce " - f"this. ({errors!r})", - errors, - ) - - -def build_conditional_library( - lib: typing.Any, - conditional_names: dict[str, Callable[[], list[str]]], -) -> typing.Any: - conditional_lib = types.ModuleType("lib") - conditional_lib._original_lib = lib # type: ignore[attr-defined] - excluded_names = set() - for condition, names_cb in conditional_names.items(): - if not getattr(lib, condition): - excluded_names.update(names_cb()) - - for attr in dir(lib): - if attr not in excluded_names: - setattr(conditional_lib, attr, getattr(lib, attr)) - - return conditional_lib - - -class Binding: - """ - OpenSSL API wrapper. - """ - - lib: typing.ClassVar = None - ffi = _openssl.ffi - _lib_loaded = False - _init_lock = threading.Lock() - - def __init__(self) -> None: - self._ensure_ffi_initialized() - - @classmethod - def _ensure_ffi_initialized(cls) -> None: - with cls._init_lock: - if not cls._lib_loaded: - cls.lib = build_conditional_library( - _openssl.lib, CONDITIONAL_NAMES - ) - cls._lib_loaded = True - - @classmethod - def init_static_locks(cls) -> None: - cls._ensure_ffi_initialized() - - -def _verify_package_version(version: str) -> None: - # Occasionally we run into situations where the version of the Python - # package does not match the version of the shared object that is loaded. - # This may occur in environments where multiple versions of cryptography - # are installed and available in the python path. To avoid errors cropping - # up later this code checks that the currently imported package and the - # shared object that were loaded have the same version and raise an - # ImportError if they do not - so_package_version = _openssl.ffi.string( - _openssl.lib.CRYPTOGRAPHY_PACKAGE_VERSION - ) - if version.encode("ascii") != so_package_version: - raise ImportError( - "The version of cryptography does not match the loaded " - "shared object. This can happen if you have multiple copies of " - "cryptography installed in your Python path. Please try creating " - "a new virtual environment to resolve this issue. " - f"Loaded python version: {version}, " - f"shared object version: {so_package_version}" - ) - - _openssl_assert( - _openssl.lib.OpenSSL_version_num() == openssl.openssl_version(), - ) - - -_verify_package_version(cryptography.__version__) - -Binding.init_static_locks() - -if ( - sys.platform == "win32" - and os.environ.get("PROCESSOR_ARCHITEW6432") is not None -): - warnings.warn( - "You are using cryptography on a 32-bit Python on a 64-bit Windows " - "Operating System. Cryptography will be significantly faster if you " - "switch to using a 64-bit Python.", - UserWarning, - stacklevel=2, - ) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/__init__.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/__init__.py deleted file mode 100644 index 41d731863aa2927b7a4fb02d7e536c35c678a102..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 660f033a6d7fa6278dbc4ec1fedce2f915b2747f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/ciphers/__init__.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/ciphers/__init__.py deleted file mode 100644 index 41d731863aa2927b7a4fb02d7e536c35c678a102..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/ciphers/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/ciphers/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/ciphers/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index d76402f31586c7ae5b38c3e2d24455d656a3319f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/ciphers/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/ciphers/__pycache__/algorithms.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/ciphers/__pycache__/algorithms.cpython-310.pyc deleted file mode 100644 index 8333efcbe51a2eff819cd6b5da92eac2997ac50c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/ciphers/__pycache__/algorithms.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/ciphers/algorithms.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/ciphers/algorithms.py deleted file mode 100644 index 072a99143d6e3fd6c69e4454d61894f34e6c59b9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/decrepit/ciphers/algorithms.py +++ /dev/null @@ -1,112 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography.hazmat.primitives._cipheralgorithm import ( - BlockCipherAlgorithm, - CipherAlgorithm, - _verify_key_size, -) - - -class ARC4(CipherAlgorithm): - name = "RC4" - key_sizes = frozenset([40, 56, 64, 80, 128, 160, 192, 256]) - - def __init__(self, key: bytes): - self.key = _verify_key_size(self, key) - - @property - def key_size(self) -> int: - return len(self.key) * 8 - - -class TripleDES(BlockCipherAlgorithm): - name = "3DES" - block_size = 64 - key_sizes = frozenset([64, 128, 192]) - - def __init__(self, key: bytes): - if len(key) == 8: - key += key + key - elif len(key) == 16: - key += key[:8] - self.key = _verify_key_size(self, key) - - @property - def key_size(self) -> int: - return len(self.key) * 8 - - -# Not actually supported, marker for tests -class _DES: - key_size = 64 - - -class Blowfish(BlockCipherAlgorithm): - name = "Blowfish" - block_size = 64 - key_sizes = frozenset(range(32, 449, 8)) - - def __init__(self, key: bytes): - self.key = _verify_key_size(self, key) - - @property - def key_size(self) -> int: - return len(self.key) * 8 - - -class CAST5(BlockCipherAlgorithm): - name = "CAST5" - block_size = 64 - key_sizes = frozenset(range(40, 129, 8)) - - def __init__(self, key: bytes): - self.key = _verify_key_size(self, key) - - @property - def key_size(self) -> int: - return len(self.key) * 8 - - -class SEED(BlockCipherAlgorithm): - name = "SEED" - block_size = 128 - key_sizes = frozenset([128]) - - def __init__(self, key: bytes): - self.key = _verify_key_size(self, key) - - @property - def key_size(self) -> int: - return len(self.key) * 8 - - -class IDEA(BlockCipherAlgorithm): - name = "IDEA" - block_size = 64 - key_sizes = frozenset([128]) - - def __init__(self, key: bytes): - self.key = _verify_key_size(self, key) - - @property - def key_size(self) -> int: - return len(self.key) * 8 - - -# This class only allows RC2 with a 128-bit key. No support for -# effective key bits or other key sizes is provided. -class RC2(BlockCipherAlgorithm): - name = "RC2" - block_size = 64 - key_sizes = frozenset([128]) - - def __init__(self, key: bytes): - self.key = _verify_key_size(self, key) - - @property - def key_size(self) -> int: - return len(self.key) * 8 diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__init__.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__init__.py deleted file mode 100644 index b509336233c2fafe4185a49da5909c8bbb38dfd7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index d515ec7d35c258e6c5ead164dcbd384e4c7b91a2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/_asymmetric.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/_asymmetric.cpython-310.pyc deleted file mode 100644 index 3929e845a3115f116abde9dbd62a0062c7c286f5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/_asymmetric.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/_cipheralgorithm.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/_cipheralgorithm.cpython-310.pyc deleted file mode 100644 index d95a65cf08cccad12b583459d5176a301483d277..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/_cipheralgorithm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/_serialization.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/_serialization.cpython-310.pyc deleted file mode 100644 index 97e7e6dc6a81fcda9e451a343b98baf519211470..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/_serialization.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/cmac.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/cmac.cpython-310.pyc deleted file mode 100644 index 4f4e840b95b454b9836d595632368b55334d6c5c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/cmac.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/constant_time.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/constant_time.cpython-310.pyc deleted file mode 100644 index 29374154566924871ff0768090ad11e221b42bd5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/constant_time.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/hashes.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/hashes.cpython-310.pyc deleted file mode 100644 index 6e0bbadbef6c0473a6f6ef10054fbde30c43940c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/hashes.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/hmac.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/hmac.cpython-310.pyc deleted file mode 100644 index 90b4f55743073c149ae8cfe7b36cc848c06bb5a0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/hmac.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/keywrap.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/keywrap.cpython-310.pyc deleted file mode 100644 index 21a511f74b799fbaee821aa1f90dc2125b21122c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/keywrap.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/padding.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/padding.cpython-310.pyc deleted file mode 100644 index 73a4955ca29b766fdaf23eca45cce8ccd1c52e00..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/padding.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/poly1305.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/poly1305.cpython-310.pyc deleted file mode 100644 index d896615ca4fb7aaab7b4740c657092b14a0e277e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/__pycache__/poly1305.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/_asymmetric.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/_asymmetric.py deleted file mode 100644 index ea55ffdf1a721f8fd2de8ae67de913bc47cbf55d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/_asymmetric.py +++ /dev/null @@ -1,19 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc - -# This exists to break an import cycle. It is normally accessible from the -# asymmetric padding module. - - -class AsymmetricPadding(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def name(self) -> str: - """ - A string naming this padding (e.g. "PSS", "PKCS1"). - """ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/_cipheralgorithm.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/_cipheralgorithm.py deleted file mode 100644 index 305a9fd3179cd439b7a86355e6cd6c8251583532..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/_cipheralgorithm.py +++ /dev/null @@ -1,60 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc - -from cryptography import utils - -# This exists to break an import cycle. It is normally accessible from the -# ciphers module. - - -class CipherAlgorithm(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def name(self) -> str: - """ - A string naming this mode (e.g. "AES", "Camellia"). - """ - - @property - @abc.abstractmethod - def key_sizes(self) -> frozenset[int]: - """ - Valid key sizes for this algorithm in bits - """ - - @property - @abc.abstractmethod - def key_size(self) -> int: - """ - The size of the key being used as an integer in bits (e.g. 128, 256). - """ - - -class BlockCipherAlgorithm(CipherAlgorithm): - key: utils.Buffer - - @property - @abc.abstractmethod - def block_size(self) -> int: - """ - The size of a block as an integer in bits (e.g. 64, 128). - """ - - -def _verify_key_size( - algorithm: CipherAlgorithm, key: utils.Buffer -) -> utils.Buffer: - # Verify that the key is instance of bytes - utils._check_byteslike("key", key) - - # Verify that the key size matches the expected key size - if len(key) * 8 not in algorithm.key_sizes: - raise ValueError( - f"Invalid key size ({len(key) * 8}) for {algorithm.name}." - ) - return key diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/_serialization.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/_serialization.py deleted file mode 100644 index e998865cdd9787e2aa2b6d87d5917b87b3ef20e9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/_serialization.py +++ /dev/null @@ -1,168 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc - -from cryptography import utils -from cryptography.hazmat.primitives.hashes import HashAlgorithm - -# This exists to break an import cycle. These classes are normally accessible -# from the serialization module. - - -class PBES(utils.Enum): - PBESv1SHA1And3KeyTripleDESCBC = "PBESv1 using SHA1 and 3-Key TripleDES" - PBESv2SHA256AndAES256CBC = "PBESv2 using SHA256 PBKDF2 and AES256 CBC" - - -class Encoding(utils.Enum): - PEM = "PEM" - DER = "DER" - OpenSSH = "OpenSSH" - Raw = "Raw" - X962 = "ANSI X9.62" - SMIME = "S/MIME" - - -class PrivateFormat(utils.Enum): - PKCS8 = "PKCS8" - TraditionalOpenSSL = "TraditionalOpenSSL" - Raw = "Raw" - OpenSSH = "OpenSSH" - PKCS12 = "PKCS12" - - def encryption_builder(self) -> KeySerializationEncryptionBuilder: - if self not in (PrivateFormat.OpenSSH, PrivateFormat.PKCS12): - raise ValueError( - "encryption_builder only supported with PrivateFormat.OpenSSH" - " and PrivateFormat.PKCS12" - ) - return KeySerializationEncryptionBuilder(self) - - -class PublicFormat(utils.Enum): - SubjectPublicKeyInfo = "X.509 subjectPublicKeyInfo with PKCS#1" - PKCS1 = "Raw PKCS#1" - OpenSSH = "OpenSSH" - Raw = "Raw" - CompressedPoint = "X9.62 Compressed Point" - UncompressedPoint = "X9.62 Uncompressed Point" - - -class ParameterFormat(utils.Enum): - PKCS3 = "PKCS3" - - -class KeySerializationEncryption(metaclass=abc.ABCMeta): - pass - - -class BestAvailableEncryption(KeySerializationEncryption): - def __init__(self, password: bytes): - if not isinstance(password, bytes) or len(password) == 0: - raise ValueError("Password must be 1 or more bytes.") - - self.password = password - - -class NoEncryption(KeySerializationEncryption): - pass - - -class KeySerializationEncryptionBuilder: - def __init__( - self, - format: PrivateFormat, - *, - _kdf_rounds: int | None = None, - _hmac_hash: HashAlgorithm | None = None, - _key_cert_algorithm: PBES | None = None, - ) -> None: - self._format = format - - self._kdf_rounds = _kdf_rounds - self._hmac_hash = _hmac_hash - self._key_cert_algorithm = _key_cert_algorithm - - def kdf_rounds(self, rounds: int) -> KeySerializationEncryptionBuilder: - if self._kdf_rounds is not None: - raise ValueError("kdf_rounds already set") - - if not isinstance(rounds, int): - raise TypeError("kdf_rounds must be an integer") - - if rounds < 1: - raise ValueError("kdf_rounds must be a positive integer") - - return KeySerializationEncryptionBuilder( - self._format, - _kdf_rounds=rounds, - _hmac_hash=self._hmac_hash, - _key_cert_algorithm=self._key_cert_algorithm, - ) - - def hmac_hash( - self, algorithm: HashAlgorithm - ) -> KeySerializationEncryptionBuilder: - if self._format is not PrivateFormat.PKCS12: - raise TypeError( - "hmac_hash only supported with PrivateFormat.PKCS12" - ) - - if self._hmac_hash is not None: - raise ValueError("hmac_hash already set") - return KeySerializationEncryptionBuilder( - self._format, - _kdf_rounds=self._kdf_rounds, - _hmac_hash=algorithm, - _key_cert_algorithm=self._key_cert_algorithm, - ) - - def key_cert_algorithm( - self, algorithm: PBES - ) -> KeySerializationEncryptionBuilder: - if self._format is not PrivateFormat.PKCS12: - raise TypeError( - "key_cert_algorithm only supported with PrivateFormat.PKCS12" - ) - if self._key_cert_algorithm is not None: - raise ValueError("key_cert_algorithm already set") - return KeySerializationEncryptionBuilder( - self._format, - _kdf_rounds=self._kdf_rounds, - _hmac_hash=self._hmac_hash, - _key_cert_algorithm=algorithm, - ) - - def build(self, password: bytes) -> KeySerializationEncryption: - if not isinstance(password, bytes) or len(password) == 0: - raise ValueError("Password must be 1 or more bytes.") - - return _KeySerializationEncryption( - self._format, - password, - kdf_rounds=self._kdf_rounds, - hmac_hash=self._hmac_hash, - key_cert_algorithm=self._key_cert_algorithm, - ) - - -class _KeySerializationEncryption(KeySerializationEncryption): - def __init__( - self, - format: PrivateFormat, - password: bytes, - *, - kdf_rounds: int | None, - hmac_hash: HashAlgorithm | None, - key_cert_algorithm: PBES | None, - ): - self._format = format - self.password = password - - self._kdf_rounds = kdf_rounds - self._hmac_hash = hmac_hash - self._key_cert_algorithm = key_cert_algorithm diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__init__.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__init__.py deleted file mode 100644 index b509336233c2fafe4185a49da5909c8bbb38dfd7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 5c6bf4529d9d12a772bc5e07a1b33b1aa4858db3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/dh.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/dh.cpython-310.pyc deleted file mode 100644 index 6dfa9e8ea6cc453a04b71a1fbf637c8db9f26ab4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/dh.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/dsa.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/dsa.cpython-310.pyc deleted file mode 100644 index f47042b73b8e3b2d85f308f09f98bf46555d083f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/dsa.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/ec.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/ec.cpython-310.pyc deleted file mode 100644 index ed01248035eaa37d996ef3f35dbb3838f3e0b24f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/ec.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/ed25519.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/ed25519.cpython-310.pyc deleted file mode 100644 index 4a33e6dfd91853d7dbb0ddac305ddd3c93a3f185..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/ed25519.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/ed448.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/ed448.cpython-310.pyc deleted file mode 100644 index c799aaa51f3d76e622ccc40316fcd5ba46ac7176..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/ed448.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/padding.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/padding.cpython-310.pyc deleted file mode 100644 index 3a2c384cb2026fad6bc850911b0f6b96a085dba3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/padding.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/rsa.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/rsa.cpython-310.pyc deleted file mode 100644 index ded885869e890b57ee7afe6850d625f5243911cb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/rsa.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/types.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/types.cpython-310.pyc deleted file mode 100644 index c4bd9e7c84e7cbf315135afd14550be75009ac61..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/types.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index 72f78e637f8fd6f9b59c1a0709b057c3bf10a60b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/x25519.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/x25519.cpython-310.pyc deleted file mode 100644 index e48269f0635821bdb06ed9052b239c27f208e314..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/x25519.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/x448.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/x448.cpython-310.pyc deleted file mode 100644 index 50a5fd8e9607c36ca0abf99caa18d6be85fda7ab..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/__pycache__/x448.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/dh.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/dh.py deleted file mode 100644 index 1822e99dcda89bc16a0a9a184499c097e854102d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/dh.py +++ /dev/null @@ -1,147 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc - -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.hazmat.primitives import _serialization - -generate_parameters = rust_openssl.dh.generate_parameters - - -DHPrivateNumbers = rust_openssl.dh.DHPrivateNumbers -DHPublicNumbers = rust_openssl.dh.DHPublicNumbers -DHParameterNumbers = rust_openssl.dh.DHParameterNumbers - - -class DHParameters(metaclass=abc.ABCMeta): - @abc.abstractmethod - def generate_private_key(self) -> DHPrivateKey: - """ - Generates and returns a DHPrivateKey. - """ - - @abc.abstractmethod - def parameter_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.ParameterFormat, - ) -> bytes: - """ - Returns the parameters serialized as bytes. - """ - - @abc.abstractmethod - def parameter_numbers(self) -> DHParameterNumbers: - """ - Returns a DHParameterNumbers. - """ - - -DHParametersWithSerialization = DHParameters -DHParameters.register(rust_openssl.dh.DHParameters) - - -class DHPublicKey(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def key_size(self) -> int: - """ - The bit length of the prime modulus. - """ - - @abc.abstractmethod - def parameters(self) -> DHParameters: - """ - The DHParameters object associated with this public key. - """ - - @abc.abstractmethod - def public_numbers(self) -> DHPublicNumbers: - """ - Returns a DHPublicNumbers. - """ - - @abc.abstractmethod - def public_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PublicFormat, - ) -> bytes: - """ - Returns the key serialized as bytes. - """ - - @abc.abstractmethod - def __eq__(self, other: object) -> bool: - """ - Checks equality. - """ - - @abc.abstractmethod - def __copy__(self) -> DHPublicKey: - """ - Returns a copy. - """ - - -DHPublicKeyWithSerialization = DHPublicKey -DHPublicKey.register(rust_openssl.dh.DHPublicKey) - - -class DHPrivateKey(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def key_size(self) -> int: - """ - The bit length of the prime modulus. - """ - - @abc.abstractmethod - def public_key(self) -> DHPublicKey: - """ - The DHPublicKey associated with this private key. - """ - - @abc.abstractmethod - def parameters(self) -> DHParameters: - """ - The DHParameters object associated with this private key. - """ - - @abc.abstractmethod - def exchange(self, peer_public_key: DHPublicKey) -> bytes: - """ - Given peer's DHPublicKey, carry out the key exchange and - return shared key as bytes. - """ - - @abc.abstractmethod - def private_numbers(self) -> DHPrivateNumbers: - """ - Returns a DHPrivateNumbers. - """ - - @abc.abstractmethod - def private_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PrivateFormat, - encryption_algorithm: _serialization.KeySerializationEncryption, - ) -> bytes: - """ - Returns the key serialized as bytes. - """ - - @abc.abstractmethod - def __copy__(self) -> DHPrivateKey: - """ - Returns a copy. - """ - - -DHPrivateKeyWithSerialization = DHPrivateKey -DHPrivateKey.register(rust_openssl.dh.DHPrivateKey) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/dsa.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/dsa.py deleted file mode 100644 index 21d78ba95e036156fc20d6b096c136d0abdddda2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/dsa.py +++ /dev/null @@ -1,167 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc -import typing - -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.hazmat.primitives import _serialization, hashes -from cryptography.hazmat.primitives.asymmetric import utils as asym_utils -from cryptography.utils import Buffer - - -class DSAParameters(metaclass=abc.ABCMeta): - @abc.abstractmethod - def generate_private_key(self) -> DSAPrivateKey: - """ - Generates and returns a DSAPrivateKey. - """ - - @abc.abstractmethod - def parameter_numbers(self) -> DSAParameterNumbers: - """ - Returns a DSAParameterNumbers. - """ - - -DSAParametersWithNumbers = DSAParameters -DSAParameters.register(rust_openssl.dsa.DSAParameters) - - -class DSAPrivateKey(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def key_size(self) -> int: - """ - The bit length of the prime modulus. - """ - - @abc.abstractmethod - def public_key(self) -> DSAPublicKey: - """ - The DSAPublicKey associated with this private key. - """ - - @abc.abstractmethod - def parameters(self) -> DSAParameters: - """ - The DSAParameters object associated with this private key. - """ - - @abc.abstractmethod - def sign( - self, - data: Buffer, - algorithm: asym_utils.Prehashed | hashes.HashAlgorithm, - ) -> bytes: - """ - Signs the data - """ - - @abc.abstractmethod - def private_numbers(self) -> DSAPrivateNumbers: - """ - Returns a DSAPrivateNumbers. - """ - - @abc.abstractmethod - def private_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PrivateFormat, - encryption_algorithm: _serialization.KeySerializationEncryption, - ) -> bytes: - """ - Returns the key serialized as bytes. - """ - - @abc.abstractmethod - def __copy__(self) -> DSAPrivateKey: - """ - Returns a copy. - """ - - -DSAPrivateKeyWithSerialization = DSAPrivateKey -DSAPrivateKey.register(rust_openssl.dsa.DSAPrivateKey) - - -class DSAPublicKey(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def key_size(self) -> int: - """ - The bit length of the prime modulus. - """ - - @abc.abstractmethod - def parameters(self) -> DSAParameters: - """ - The DSAParameters object associated with this public key. - """ - - @abc.abstractmethod - def public_numbers(self) -> DSAPublicNumbers: - """ - Returns a DSAPublicNumbers. - """ - - @abc.abstractmethod - def public_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PublicFormat, - ) -> bytes: - """ - Returns the key serialized as bytes. - """ - - @abc.abstractmethod - def verify( - self, - signature: Buffer, - data: Buffer, - algorithm: asym_utils.Prehashed | hashes.HashAlgorithm, - ) -> None: - """ - Verifies the signature of the data. - """ - - @abc.abstractmethod - def __eq__(self, other: object) -> bool: - """ - Checks equality. - """ - - @abc.abstractmethod - def __copy__(self) -> DSAPublicKey: - """ - Returns a copy. - """ - - -DSAPublicKeyWithSerialization = DSAPublicKey -DSAPublicKey.register(rust_openssl.dsa.DSAPublicKey) - -DSAPrivateNumbers = rust_openssl.dsa.DSAPrivateNumbers -DSAPublicNumbers = rust_openssl.dsa.DSAPublicNumbers -DSAParameterNumbers = rust_openssl.dsa.DSAParameterNumbers - - -def generate_parameters( - key_size: int, backend: typing.Any = None -) -> DSAParameters: - if key_size not in (1024, 2048, 3072, 4096): - raise ValueError("Key size must be 1024, 2048, 3072, or 4096 bits.") - - return rust_openssl.dsa.generate_parameters(key_size) - - -def generate_private_key( - key_size: int, backend: typing.Any = None -) -> DSAPrivateKey: - parameters = generate_parameters(key_size) - return parameters.generate_private_key() diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/ec.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/ec.py deleted file mode 100644 index a13d9827f44afc42de62b7c21fd48fe2c99b1578..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/ec.py +++ /dev/null @@ -1,447 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc -import typing - -from cryptography import utils -from cryptography.exceptions import UnsupportedAlgorithm, _Reasons -from cryptography.hazmat._oid import ObjectIdentifier -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.hazmat.primitives import _serialization, hashes -from cryptography.hazmat.primitives.asymmetric import utils as asym_utils - - -class EllipticCurveOID: - SECP192R1 = ObjectIdentifier("1.2.840.10045.3.1.1") - SECP224R1 = ObjectIdentifier("1.3.132.0.33") - SECP256K1 = ObjectIdentifier("1.3.132.0.10") - SECP256R1 = ObjectIdentifier("1.2.840.10045.3.1.7") - SECP384R1 = ObjectIdentifier("1.3.132.0.34") - SECP521R1 = ObjectIdentifier("1.3.132.0.35") - BRAINPOOLP256R1 = ObjectIdentifier("1.3.36.3.3.2.8.1.1.7") - BRAINPOOLP384R1 = ObjectIdentifier("1.3.36.3.3.2.8.1.1.11") - BRAINPOOLP512R1 = ObjectIdentifier("1.3.36.3.3.2.8.1.1.13") - SECT163K1 = ObjectIdentifier("1.3.132.0.1") - SECT163R2 = ObjectIdentifier("1.3.132.0.15") - SECT233K1 = ObjectIdentifier("1.3.132.0.26") - SECT233R1 = ObjectIdentifier("1.3.132.0.27") - SECT283K1 = ObjectIdentifier("1.3.132.0.16") - SECT283R1 = ObjectIdentifier("1.3.132.0.17") - SECT409K1 = ObjectIdentifier("1.3.132.0.36") - SECT409R1 = ObjectIdentifier("1.3.132.0.37") - SECT571K1 = ObjectIdentifier("1.3.132.0.38") - SECT571R1 = ObjectIdentifier("1.3.132.0.39") - - -class EllipticCurve(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def name(self) -> str: - """ - The name of the curve. e.g. secp256r1. - """ - - @property - @abc.abstractmethod - def key_size(self) -> int: - """ - Bit size of a secret scalar for the curve. - """ - - @property - @abc.abstractmethod - def group_order(self) -> int: - """ - The order of the curve's group. - """ - - -class EllipticCurveSignatureAlgorithm(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def algorithm( - self, - ) -> asym_utils.Prehashed | hashes.HashAlgorithm: - """ - The digest algorithm used with this signature. - """ - - -class EllipticCurvePrivateKey(metaclass=abc.ABCMeta): - @abc.abstractmethod - def exchange( - self, algorithm: ECDH, peer_public_key: EllipticCurvePublicKey - ) -> bytes: - """ - Performs a key exchange operation using the provided algorithm with the - provided peer's public key. - """ - - @abc.abstractmethod - def public_key(self) -> EllipticCurvePublicKey: - """ - The EllipticCurvePublicKey for this private key. - """ - - @property - @abc.abstractmethod - def curve(self) -> EllipticCurve: - """ - The EllipticCurve that this key is on. - """ - - @property - @abc.abstractmethod - def key_size(self) -> int: - """ - Bit size of a secret scalar for the curve. - """ - - @abc.abstractmethod - def sign( - self, - data: utils.Buffer, - signature_algorithm: EllipticCurveSignatureAlgorithm, - ) -> bytes: - """ - Signs the data - """ - - @abc.abstractmethod - def private_numbers(self) -> EllipticCurvePrivateNumbers: - """ - Returns an EllipticCurvePrivateNumbers. - """ - - @abc.abstractmethod - def private_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PrivateFormat, - encryption_algorithm: _serialization.KeySerializationEncryption, - ) -> bytes: - """ - Returns the key serialized as bytes. - """ - - @abc.abstractmethod - def __copy__(self) -> EllipticCurvePrivateKey: - """ - Returns a copy. - """ - - -EllipticCurvePrivateKeyWithSerialization = EllipticCurvePrivateKey -EllipticCurvePrivateKey.register(rust_openssl.ec.ECPrivateKey) - - -class EllipticCurvePublicKey(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def curve(self) -> EllipticCurve: - """ - The EllipticCurve that this key is on. - """ - - @property - @abc.abstractmethod - def key_size(self) -> int: - """ - Bit size of a secret scalar for the curve. - """ - - @abc.abstractmethod - def public_numbers(self) -> EllipticCurvePublicNumbers: - """ - Returns an EllipticCurvePublicNumbers. - """ - - @abc.abstractmethod - def public_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PublicFormat, - ) -> bytes: - """ - Returns the key serialized as bytes. - """ - - @abc.abstractmethod - def verify( - self, - signature: utils.Buffer, - data: utils.Buffer, - signature_algorithm: EllipticCurveSignatureAlgorithm, - ) -> None: - """ - Verifies the signature of the data. - """ - - @classmethod - def from_encoded_point( - cls, curve: EllipticCurve, data: bytes - ) -> EllipticCurvePublicKey: - utils._check_bytes("data", data) - - if len(data) == 0: - raise ValueError("data must not be an empty byte string") - - if data[0] not in [0x02, 0x03, 0x04]: - raise ValueError("Unsupported elliptic curve point type") - - return rust_openssl.ec.from_public_bytes(curve, data) - - @abc.abstractmethod - def __eq__(self, other: object) -> bool: - """ - Checks equality. - """ - - @abc.abstractmethod - def __copy__(self) -> EllipticCurvePublicKey: - """ - Returns a copy. - """ - - -EllipticCurvePublicKeyWithSerialization = EllipticCurvePublicKey -EllipticCurvePublicKey.register(rust_openssl.ec.ECPublicKey) - -EllipticCurvePrivateNumbers = rust_openssl.ec.EllipticCurvePrivateNumbers -EllipticCurvePublicNumbers = rust_openssl.ec.EllipticCurvePublicNumbers - - -class SECT571R1(EllipticCurve): - name = "sect571r1" - key_size = 570 - group_order = 0x3FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFE661CE18FF55987308059B186823851EC7DD9CA1161DE93D5174D66E8382E9BB2FE84E47 # noqa: E501 - - -class SECT409R1(EllipticCurve): - name = "sect409r1" - key_size = 409 - group_order = 0x10000000000000000000000000000000000000000000000000001E2AAD6A612F33307BE5FA47C3C9E052F838164CD37D9A21173 # noqa: E501 - - -class SECT283R1(EllipticCurve): - name = "sect283r1" - key_size = 283 - group_order = 0x3FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEF90399660FC938A90165B042A7CEFADB307 # noqa: E501 - - -class SECT233R1(EllipticCurve): - name = "sect233r1" - key_size = 233 - group_order = 0x1000000000000000000000000000013E974E72F8A6922031D2603CFE0D7 - - -class SECT163R2(EllipticCurve): - name = "sect163r2" - key_size = 163 - group_order = 0x40000000000000000000292FE77E70C12A4234C33 - - -class SECT571K1(EllipticCurve): - name = "sect571k1" - key_size = 571 - group_order = 0x20000000000000000000000000000000000000000000000000000000000000000000000131850E1F19A63E4B391A8DB917F4138B630D84BE5D639381E91DEB45CFE778F637C1001 # noqa: E501 - - -class SECT409K1(EllipticCurve): - name = "sect409k1" - key_size = 409 - group_order = 0x7FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFE5F83B2D4EA20400EC4557D5ED3E3E7CA5B4B5C83B8E01E5FCF # noqa: E501 - - -class SECT283K1(EllipticCurve): - name = "sect283k1" - key_size = 283 - group_order = 0x1FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFE9AE2ED07577265DFF7F94451E061E163C61 # noqa: E501 - - -class SECT233K1(EllipticCurve): - name = "sect233k1" - key_size = 233 - group_order = 0x8000000000000000000000000000069D5BB915BCD46EFB1AD5F173ABDF - - -class SECT163K1(EllipticCurve): - name = "sect163k1" - key_size = 163 - group_order = 0x4000000000000000000020108A2E0CC0D99F8A5EF - - -class SECP521R1(EllipticCurve): - name = "secp521r1" - key_size = 521 - group_order = 0x1FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFA51868783BF2F966B7FCC0148F709A5D03BB5C9B8899C47AEBB6FB71E91386409 # noqa: E501 - - -class SECP384R1(EllipticCurve): - name = "secp384r1" - key_size = 384 - group_order = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFC7634D81F4372DDF581A0DB248B0A77AECEC196ACCC52973 # noqa: E501 - - -class SECP256R1(EllipticCurve): - name = "secp256r1" - key_size = 256 - group_order = ( - 0xFFFFFFFF00000000FFFFFFFFFFFFFFFFBCE6FAADA7179E84F3B9CAC2FC632551 - ) - - -class SECP256K1(EllipticCurve): - name = "secp256k1" - key_size = 256 - group_order = ( - 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141 - ) - - -class SECP224R1(EllipticCurve): - name = "secp224r1" - key_size = 224 - group_order = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFF16A2E0B8F03E13DD29455C5C2A3D - - -class SECP192R1(EllipticCurve): - name = "secp192r1" - key_size = 192 - group_order = 0xFFFFFFFFFFFFFFFFFFFFFFFF99DEF836146BC9B1B4D22831 - - -class BrainpoolP256R1(EllipticCurve): - name = "brainpoolP256r1" - key_size = 256 - group_order = ( - 0xA9FB57DBA1EEA9BC3E660A909D838D718C397AA3B561A6F7901E0E82974856A7 - ) - - -class BrainpoolP384R1(EllipticCurve): - name = "brainpoolP384r1" - key_size = 384 - group_order = 0x8CB91E82A3386D280F5D6F7E50E641DF152F7109ED5456B31F166E6CAC0425A7CF3AB6AF6B7FC3103B883202E9046565 # noqa: E501 - - -class BrainpoolP512R1(EllipticCurve): - name = "brainpoolP512r1" - key_size = 512 - group_order = 0xAADD9DB8DBE9C48B3FD4E6AE33C9FC07CB308DB3B3C9D20ED6639CCA70330870553E5C414CA92619418661197FAC10471DB1D381085DDADDB58796829CA90069 # noqa: E501 - - -_CURVE_TYPES: dict[str, EllipticCurve] = { - "prime192v1": SECP192R1(), - "prime256v1": SECP256R1(), - "secp192r1": SECP192R1(), - "secp224r1": SECP224R1(), - "secp256r1": SECP256R1(), - "secp384r1": SECP384R1(), - "secp521r1": SECP521R1(), - "secp256k1": SECP256K1(), - "sect163k1": SECT163K1(), - "sect233k1": SECT233K1(), - "sect283k1": SECT283K1(), - "sect409k1": SECT409K1(), - "sect571k1": SECT571K1(), - "sect163r2": SECT163R2(), - "sect233r1": SECT233R1(), - "sect283r1": SECT283R1(), - "sect409r1": SECT409R1(), - "sect571r1": SECT571R1(), - "brainpoolP256r1": BrainpoolP256R1(), - "brainpoolP384r1": BrainpoolP384R1(), - "brainpoolP512r1": BrainpoolP512R1(), -} - - -class ECDSA(EllipticCurveSignatureAlgorithm): - def __init__( - self, - algorithm: asym_utils.Prehashed | hashes.HashAlgorithm, - deterministic_signing: bool = False, - ): - from cryptography.hazmat.backends.openssl.backend import backend - - if ( - deterministic_signing - and not backend.ecdsa_deterministic_supported() - ): - raise UnsupportedAlgorithm( - "ECDSA with deterministic signature (RFC 6979) is not " - "supported by this version of OpenSSL.", - _Reasons.UNSUPPORTED_PUBLIC_KEY_ALGORITHM, - ) - self._algorithm = algorithm - self._deterministic_signing = deterministic_signing - - @property - def algorithm( - self, - ) -> asym_utils.Prehashed | hashes.HashAlgorithm: - return self._algorithm - - @property - def deterministic_signing( - self, - ) -> bool: - return self._deterministic_signing - - -generate_private_key = rust_openssl.ec.generate_private_key - - -def derive_private_key( - private_value: int, - curve: EllipticCurve, - backend: typing.Any = None, -) -> EllipticCurvePrivateKey: - if not isinstance(private_value, int): - raise TypeError("private_value must be an integer type.") - - if private_value <= 0: - raise ValueError("private_value must be a positive integer.") - - return rust_openssl.ec.derive_private_key(private_value, curve) - - -class ECDH: - pass - - -_OID_TO_CURVE = { - EllipticCurveOID.SECP192R1: SECP192R1, - EllipticCurveOID.SECP224R1: SECP224R1, - EllipticCurveOID.SECP256K1: SECP256K1, - EllipticCurveOID.SECP256R1: SECP256R1, - EllipticCurveOID.SECP384R1: SECP384R1, - EllipticCurveOID.SECP521R1: SECP521R1, - EllipticCurveOID.BRAINPOOLP256R1: BrainpoolP256R1, - EllipticCurveOID.BRAINPOOLP384R1: BrainpoolP384R1, - EllipticCurveOID.BRAINPOOLP512R1: BrainpoolP512R1, - EllipticCurveOID.SECT163K1: SECT163K1, - EllipticCurveOID.SECT163R2: SECT163R2, - EllipticCurveOID.SECT233K1: SECT233K1, - EllipticCurveOID.SECT233R1: SECT233R1, - EllipticCurveOID.SECT283K1: SECT283K1, - EllipticCurveOID.SECT283R1: SECT283R1, - EllipticCurveOID.SECT409K1: SECT409K1, - EllipticCurveOID.SECT409R1: SECT409R1, - EllipticCurveOID.SECT571K1: SECT571K1, - EllipticCurveOID.SECT571R1: SECT571R1, -} - - -def get_curve_for_oid(oid: ObjectIdentifier) -> type[EllipticCurve]: - try: - return _OID_TO_CURVE[oid] - except KeyError: - raise LookupError( - "The provided object identifier has no matching elliptic " - "curve class" - ) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/ed25519.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/ed25519.py deleted file mode 100644 index e576dc9f42f159bb37cc6dd791498587d731d9fa..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/ed25519.py +++ /dev/null @@ -1,129 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc - -from cryptography.exceptions import UnsupportedAlgorithm, _Reasons -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.hazmat.primitives import _serialization -from cryptography.utils import Buffer - - -class Ed25519PublicKey(metaclass=abc.ABCMeta): - @classmethod - def from_public_bytes(cls, data: bytes) -> Ed25519PublicKey: - from cryptography.hazmat.backends.openssl.backend import backend - - if not backend.ed25519_supported(): - raise UnsupportedAlgorithm( - "ed25519 is not supported by this version of OpenSSL.", - _Reasons.UNSUPPORTED_PUBLIC_KEY_ALGORITHM, - ) - - return rust_openssl.ed25519.from_public_bytes(data) - - @abc.abstractmethod - def public_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PublicFormat, - ) -> bytes: - """ - The serialized bytes of the public key. - """ - - @abc.abstractmethod - def public_bytes_raw(self) -> bytes: - """ - The raw bytes of the public key. - Equivalent to public_bytes(Raw, Raw). - """ - - @abc.abstractmethod - def verify(self, signature: Buffer, data: Buffer) -> None: - """ - Verify the signature. - """ - - @abc.abstractmethod - def __eq__(self, other: object) -> bool: - """ - Checks equality. - """ - - @abc.abstractmethod - def __copy__(self) -> Ed25519PublicKey: - """ - Returns a copy. - """ - - -Ed25519PublicKey.register(rust_openssl.ed25519.Ed25519PublicKey) - - -class Ed25519PrivateKey(metaclass=abc.ABCMeta): - @classmethod - def generate(cls) -> Ed25519PrivateKey: - from cryptography.hazmat.backends.openssl.backend import backend - - if not backend.ed25519_supported(): - raise UnsupportedAlgorithm( - "ed25519 is not supported by this version of OpenSSL.", - _Reasons.UNSUPPORTED_PUBLIC_KEY_ALGORITHM, - ) - - return rust_openssl.ed25519.generate_key() - - @classmethod - def from_private_bytes(cls, data: Buffer) -> Ed25519PrivateKey: - from cryptography.hazmat.backends.openssl.backend import backend - - if not backend.ed25519_supported(): - raise UnsupportedAlgorithm( - "ed25519 is not supported by this version of OpenSSL.", - _Reasons.UNSUPPORTED_PUBLIC_KEY_ALGORITHM, - ) - - return rust_openssl.ed25519.from_private_bytes(data) - - @abc.abstractmethod - def public_key(self) -> Ed25519PublicKey: - """ - The Ed25519PublicKey derived from the private key. - """ - - @abc.abstractmethod - def private_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PrivateFormat, - encryption_algorithm: _serialization.KeySerializationEncryption, - ) -> bytes: - """ - The serialized bytes of the private key. - """ - - @abc.abstractmethod - def private_bytes_raw(self) -> bytes: - """ - The raw bytes of the private key. - Equivalent to private_bytes(Raw, Raw, NoEncryption()). - """ - - @abc.abstractmethod - def sign(self, data: Buffer) -> bytes: - """ - Signs the data. - """ - - @abc.abstractmethod - def __copy__(self) -> Ed25519PrivateKey: - """ - Returns a copy. - """ - - -Ed25519PrivateKey.register(rust_openssl.ed25519.Ed25519PrivateKey) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/ed448.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/ed448.py deleted file mode 100644 index 89db20984fd6fda41b5ae26941484e681c512d31..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/ed448.py +++ /dev/null @@ -1,131 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc - -from cryptography.exceptions import UnsupportedAlgorithm, _Reasons -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.hazmat.primitives import _serialization -from cryptography.utils import Buffer - - -class Ed448PublicKey(metaclass=abc.ABCMeta): - @classmethod - def from_public_bytes(cls, data: bytes) -> Ed448PublicKey: - from cryptography.hazmat.backends.openssl.backend import backend - - if not backend.ed448_supported(): - raise UnsupportedAlgorithm( - "ed448 is not supported by this version of OpenSSL.", - _Reasons.UNSUPPORTED_PUBLIC_KEY_ALGORITHM, - ) - - return rust_openssl.ed448.from_public_bytes(data) - - @abc.abstractmethod - def public_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PublicFormat, - ) -> bytes: - """ - The serialized bytes of the public key. - """ - - @abc.abstractmethod - def public_bytes_raw(self) -> bytes: - """ - The raw bytes of the public key. - Equivalent to public_bytes(Raw, Raw). - """ - - @abc.abstractmethod - def verify(self, signature: Buffer, data: Buffer) -> None: - """ - Verify the signature. - """ - - @abc.abstractmethod - def __eq__(self, other: object) -> bool: - """ - Checks equality. - """ - - @abc.abstractmethod - def __copy__(self) -> Ed448PublicKey: - """ - Returns a copy. - """ - - -if hasattr(rust_openssl, "ed448"): - Ed448PublicKey.register(rust_openssl.ed448.Ed448PublicKey) - - -class Ed448PrivateKey(metaclass=abc.ABCMeta): - @classmethod - def generate(cls) -> Ed448PrivateKey: - from cryptography.hazmat.backends.openssl.backend import backend - - if not backend.ed448_supported(): - raise UnsupportedAlgorithm( - "ed448 is not supported by this version of OpenSSL.", - _Reasons.UNSUPPORTED_PUBLIC_KEY_ALGORITHM, - ) - - return rust_openssl.ed448.generate_key() - - @classmethod - def from_private_bytes(cls, data: Buffer) -> Ed448PrivateKey: - from cryptography.hazmat.backends.openssl.backend import backend - - if not backend.ed448_supported(): - raise UnsupportedAlgorithm( - "ed448 is not supported by this version of OpenSSL.", - _Reasons.UNSUPPORTED_PUBLIC_KEY_ALGORITHM, - ) - - return rust_openssl.ed448.from_private_bytes(data) - - @abc.abstractmethod - def public_key(self) -> Ed448PublicKey: - """ - The Ed448PublicKey derived from the private key. - """ - - @abc.abstractmethod - def sign(self, data: Buffer) -> bytes: - """ - Signs the data. - """ - - @abc.abstractmethod - def private_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PrivateFormat, - encryption_algorithm: _serialization.KeySerializationEncryption, - ) -> bytes: - """ - The serialized bytes of the private key. - """ - - @abc.abstractmethod - def private_bytes_raw(self) -> bytes: - """ - The raw bytes of the private key. - Equivalent to private_bytes(Raw, Raw, NoEncryption()). - """ - - @abc.abstractmethod - def __copy__(self) -> Ed448PrivateKey: - """ - Returns a copy. - """ - - -if hasattr(rust_openssl, "x448"): - Ed448PrivateKey.register(rust_openssl.ed448.Ed448PrivateKey) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/padding.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/padding.py deleted file mode 100644 index b4babf44f79b1fa6ed1920d55c1d641e5904d6ab..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/padding.py +++ /dev/null @@ -1,113 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc - -from cryptography.hazmat.primitives import hashes -from cryptography.hazmat.primitives._asymmetric import ( - AsymmetricPadding as AsymmetricPadding, -) -from cryptography.hazmat.primitives.asymmetric import rsa - - -class PKCS1v15(AsymmetricPadding): - name = "EMSA-PKCS1-v1_5" - - -class _MaxLength: - "Sentinel value for `MAX_LENGTH`." - - -class _Auto: - "Sentinel value for `AUTO`." - - -class _DigestLength: - "Sentinel value for `DIGEST_LENGTH`." - - -class PSS(AsymmetricPadding): - MAX_LENGTH = _MaxLength() - AUTO = _Auto() - DIGEST_LENGTH = _DigestLength() - name = "EMSA-PSS" - _salt_length: int | _MaxLength | _Auto | _DigestLength - - def __init__( - self, - mgf: MGF, - salt_length: int | _MaxLength | _Auto | _DigestLength, - ) -> None: - self._mgf = mgf - - if not isinstance( - salt_length, (int, _MaxLength, _Auto, _DigestLength) - ): - raise TypeError( - "salt_length must be an integer, MAX_LENGTH, " - "DIGEST_LENGTH, or AUTO" - ) - - if isinstance(salt_length, int) and salt_length < 0: - raise ValueError("salt_length must be zero or greater.") - - self._salt_length = salt_length - - @property - def mgf(self) -> MGF: - return self._mgf - - -class OAEP(AsymmetricPadding): - name = "EME-OAEP" - - def __init__( - self, - mgf: MGF, - algorithm: hashes.HashAlgorithm, - label: bytes | None, - ): - if not isinstance(algorithm, hashes.HashAlgorithm): - raise TypeError("Expected instance of hashes.HashAlgorithm.") - - self._mgf = mgf - self._algorithm = algorithm - self._label = label - - @property - def algorithm(self) -> hashes.HashAlgorithm: - return self._algorithm - - @property - def mgf(self) -> MGF: - return self._mgf - - -class MGF(metaclass=abc.ABCMeta): - _algorithm: hashes.HashAlgorithm - - -class MGF1(MGF): - MAX_LENGTH = _MaxLength() - - def __init__(self, algorithm: hashes.HashAlgorithm): - if not isinstance(algorithm, hashes.HashAlgorithm): - raise TypeError("Expected instance of hashes.HashAlgorithm.") - - self._algorithm = algorithm - - -def calculate_max_pss_salt_length( - key: rsa.RSAPrivateKey | rsa.RSAPublicKey, - hash_algorithm: hashes.HashAlgorithm, -) -> int: - if not isinstance(key, (rsa.RSAPrivateKey, rsa.RSAPublicKey)): - raise TypeError("key must be an RSA public or private key") - # bit length - 1 per RFC 3447 - emlen = (key.key_size + 6) // 8 - salt_length = emlen - hash_algorithm.digest_size - 2 - assert salt_length >= 0 - return salt_length diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/rsa.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/rsa.py deleted file mode 100644 index 2e192aaaa749a3ad9809dfd25752d64b5d756b1c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/rsa.py +++ /dev/null @@ -1,277 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc -import random -import typing -from math import gcd - -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.hazmat.primitives import _serialization, hashes -from cryptography.hazmat.primitives._asymmetric import AsymmetricPadding -from cryptography.hazmat.primitives.asymmetric import utils as asym_utils - - -class RSAPrivateKey(metaclass=abc.ABCMeta): - @abc.abstractmethod - def decrypt(self, ciphertext: bytes, padding: AsymmetricPadding) -> bytes: - """ - Decrypts the provided ciphertext. - """ - - @property - @abc.abstractmethod - def key_size(self) -> int: - """ - The bit length of the public modulus. - """ - - @abc.abstractmethod - def public_key(self) -> RSAPublicKey: - """ - The RSAPublicKey associated with this private key. - """ - - @abc.abstractmethod - def sign( - self, - data: bytes, - padding: AsymmetricPadding, - algorithm: asym_utils.Prehashed | hashes.HashAlgorithm, - ) -> bytes: - """ - Signs the data. - """ - - @abc.abstractmethod - def private_numbers(self) -> RSAPrivateNumbers: - """ - Returns an RSAPrivateNumbers. - """ - - @abc.abstractmethod - def private_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PrivateFormat, - encryption_algorithm: _serialization.KeySerializationEncryption, - ) -> bytes: - """ - Returns the key serialized as bytes. - """ - - @abc.abstractmethod - def __copy__(self) -> RSAPrivateKey: - """ - Returns a copy. - """ - - -RSAPrivateKeyWithSerialization = RSAPrivateKey -RSAPrivateKey.register(rust_openssl.rsa.RSAPrivateKey) - - -class RSAPublicKey(metaclass=abc.ABCMeta): - @abc.abstractmethod - def encrypt(self, plaintext: bytes, padding: AsymmetricPadding) -> bytes: - """ - Encrypts the given plaintext. - """ - - @property - @abc.abstractmethod - def key_size(self) -> int: - """ - The bit length of the public modulus. - """ - - @abc.abstractmethod - def public_numbers(self) -> RSAPublicNumbers: - """ - Returns an RSAPublicNumbers - """ - - @abc.abstractmethod - def public_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PublicFormat, - ) -> bytes: - """ - Returns the key serialized as bytes. - """ - - @abc.abstractmethod - def verify( - self, - signature: bytes, - data: bytes, - padding: AsymmetricPadding, - algorithm: asym_utils.Prehashed | hashes.HashAlgorithm, - ) -> None: - """ - Verifies the signature of the data. - """ - - @abc.abstractmethod - def recover_data_from_signature( - self, - signature: bytes, - padding: AsymmetricPadding, - algorithm: hashes.HashAlgorithm | None, - ) -> bytes: - """ - Recovers the original data from the signature. - """ - - @abc.abstractmethod - def __eq__(self, other: object) -> bool: - """ - Checks equality. - """ - - @abc.abstractmethod - def __copy__(self) -> RSAPublicKey: - """ - Returns a copy. - """ - - -RSAPublicKeyWithSerialization = RSAPublicKey -RSAPublicKey.register(rust_openssl.rsa.RSAPublicKey) - -RSAPrivateNumbers = rust_openssl.rsa.RSAPrivateNumbers -RSAPublicNumbers = rust_openssl.rsa.RSAPublicNumbers - - -def generate_private_key( - public_exponent: int, - key_size: int, - backend: typing.Any = None, -) -> RSAPrivateKey: - _verify_rsa_parameters(public_exponent, key_size) - return rust_openssl.rsa.generate_private_key(public_exponent, key_size) - - -def _verify_rsa_parameters(public_exponent: int, key_size: int) -> None: - if public_exponent not in (3, 65537): - raise ValueError( - "public_exponent must be either 3 (for legacy compatibility) or " - "65537. Almost everyone should choose 65537 here!" - ) - - if key_size < 1024: - raise ValueError("key_size must be at least 1024-bits.") - - -def _modinv(e: int, m: int) -> int: - """ - Modular Multiplicative Inverse. Returns x such that: (x*e) mod m == 1 - """ - x1, x2 = 1, 0 - a, b = e, m - while b > 0: - q, r = divmod(a, b) - xn = x1 - q * x2 - a, b, x1, x2 = b, r, x2, xn - return x1 % m - - -def rsa_crt_iqmp(p: int, q: int) -> int: - """ - Compute the CRT (q ** -1) % p value from RSA primes p and q. - """ - return _modinv(q, p) - - -def rsa_crt_dmp1(private_exponent: int, p: int) -> int: - """ - Compute the CRT private_exponent % (p - 1) value from the RSA - private_exponent (d) and p. - """ - return private_exponent % (p - 1) - - -def rsa_crt_dmq1(private_exponent: int, q: int) -> int: - """ - Compute the CRT private_exponent % (q - 1) value from the RSA - private_exponent (d) and q. - """ - return private_exponent % (q - 1) - - -def rsa_recover_private_exponent(e: int, p: int, q: int) -> int: - """ - Compute the RSA private_exponent (d) given the public exponent (e) - and the RSA primes p and q. - - This uses the Carmichael totient function to generate the - smallest possible working value of the private exponent. - """ - # This lambda_n is the Carmichael totient function. - # The original RSA paper uses the Euler totient function - # here: phi_n = (p - 1) * (q - 1) - # Either version of the private exponent will work, but the - # one generated by the older formulation may be larger - # than necessary. (lambda_n always divides phi_n) - # - # TODO: Replace with lcm(p - 1, q - 1) once the minimum - # supported Python version is >= 3.9. - lambda_n = (p - 1) * (q - 1) // gcd(p - 1, q - 1) - return _modinv(e, lambda_n) - - -# Controls the number of iterations rsa_recover_prime_factors will perform -# to obtain the prime factors. -_MAX_RECOVERY_ATTEMPTS = 500 - - -def rsa_recover_prime_factors(n: int, e: int, d: int) -> tuple[int, int]: - """ - Compute factors p and q from the private exponent d. We assume that n has - no more than two factors. This function is adapted from code in PyCrypto. - """ - # reject invalid values early - if d <= 1 or e <= 1: - raise ValueError("d, e can't be <= 1") - if 17 != pow(17, e * d, n): - raise ValueError("n, d, e don't match") - # See 8.2.2(i) in Handbook of Applied Cryptography. - ktot = d * e - 1 - # The quantity d*e-1 is a multiple of phi(n), even, - # and can be represented as t*2^s. - t = ktot - while t % 2 == 0: - t = t // 2 - # Cycle through all multiplicative inverses in Zn. - # The algorithm is non-deterministic, but there is a 50% chance - # any candidate a leads to successful factoring. - # See "Digitalized Signatures and Public Key Functions as Intractable - # as Factorization", M. Rabin, 1979 - spotted = False - tries = 0 - while not spotted and tries < _MAX_RECOVERY_ATTEMPTS: - a = random.randint(2, n - 1) - tries += 1 - k = t - # Cycle through all values a^{t*2^i}=a^k - while k < ktot: - cand = pow(a, k, n) - # Check if a^k is a non-trivial root of unity (mod n) - if cand != 1 and cand != (n - 1) and pow(cand, 2, n) == 1: - # We have found a number such that (cand-1)(cand+1)=0 (mod n). - # Either of the terms divides n. - p = gcd(cand + 1, n) - spotted = True - break - k *= 2 - if not spotted: - raise ValueError("Unable to compute factors p and q from exponent d.") - # Found ! - q, r = divmod(n, p) - assert r == 0 - p, q = sorted((p, q), reverse=True) - return (p, q) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/types.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/types.py deleted file mode 100644 index 1fe4eaf51d850c30ef7764d1c0bbf17533b8ad38..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/types.py +++ /dev/null @@ -1,111 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import typing - -from cryptography import utils -from cryptography.hazmat.primitives.asymmetric import ( - dh, - dsa, - ec, - ed448, - ed25519, - rsa, - x448, - x25519, -) - -# Every asymmetric key type -PublicKeyTypes = typing.Union[ - dh.DHPublicKey, - dsa.DSAPublicKey, - rsa.RSAPublicKey, - ec.EllipticCurvePublicKey, - ed25519.Ed25519PublicKey, - ed448.Ed448PublicKey, - x25519.X25519PublicKey, - x448.X448PublicKey, -] -PUBLIC_KEY_TYPES = PublicKeyTypes -utils.deprecated( - PUBLIC_KEY_TYPES, - __name__, - "Use PublicKeyTypes instead", - utils.DeprecatedIn40, - name="PUBLIC_KEY_TYPES", -) -# Every asymmetric key type -PrivateKeyTypes = typing.Union[ - dh.DHPrivateKey, - ed25519.Ed25519PrivateKey, - ed448.Ed448PrivateKey, - rsa.RSAPrivateKey, - dsa.DSAPrivateKey, - ec.EllipticCurvePrivateKey, - x25519.X25519PrivateKey, - x448.X448PrivateKey, -] -PRIVATE_KEY_TYPES = PrivateKeyTypes -utils.deprecated( - PRIVATE_KEY_TYPES, - __name__, - "Use PrivateKeyTypes instead", - utils.DeprecatedIn40, - name="PRIVATE_KEY_TYPES", -) -# Just the key types we allow to be used for x509 signing. This mirrors -# the certificate public key types -CertificateIssuerPrivateKeyTypes = typing.Union[ - ed25519.Ed25519PrivateKey, - ed448.Ed448PrivateKey, - rsa.RSAPrivateKey, - dsa.DSAPrivateKey, - ec.EllipticCurvePrivateKey, -] -CERTIFICATE_PRIVATE_KEY_TYPES = CertificateIssuerPrivateKeyTypes -utils.deprecated( - CERTIFICATE_PRIVATE_KEY_TYPES, - __name__, - "Use CertificateIssuerPrivateKeyTypes instead", - utils.DeprecatedIn40, - name="CERTIFICATE_PRIVATE_KEY_TYPES", -) -# Just the key types we allow to be used for x509 signing. This mirrors -# the certificate private key types -CertificateIssuerPublicKeyTypes = typing.Union[ - dsa.DSAPublicKey, - rsa.RSAPublicKey, - ec.EllipticCurvePublicKey, - ed25519.Ed25519PublicKey, - ed448.Ed448PublicKey, -] -CERTIFICATE_ISSUER_PUBLIC_KEY_TYPES = CertificateIssuerPublicKeyTypes -utils.deprecated( - CERTIFICATE_ISSUER_PUBLIC_KEY_TYPES, - __name__, - "Use CertificateIssuerPublicKeyTypes instead", - utils.DeprecatedIn40, - name="CERTIFICATE_ISSUER_PUBLIC_KEY_TYPES", -) -# This type removes DHPublicKey. x448/x25519 can be a public key -# but cannot be used in signing so they are allowed here. -CertificatePublicKeyTypes = typing.Union[ - dsa.DSAPublicKey, - rsa.RSAPublicKey, - ec.EllipticCurvePublicKey, - ed25519.Ed25519PublicKey, - ed448.Ed448PublicKey, - x25519.X25519PublicKey, - x448.X448PublicKey, -] -CERTIFICATE_PUBLIC_KEY_TYPES = CertificatePublicKeyTypes -utils.deprecated( - CERTIFICATE_PUBLIC_KEY_TYPES, - __name__, - "Use CertificatePublicKeyTypes instead", - utils.DeprecatedIn40, - name="CERTIFICATE_PUBLIC_KEY_TYPES", -) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/utils.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/utils.py deleted file mode 100644 index 826b9567b47bc704902eb959fd21376c8969f695..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/utils.py +++ /dev/null @@ -1,24 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography.hazmat.bindings._rust import asn1 -from cryptography.hazmat.primitives import hashes - -decode_dss_signature = asn1.decode_dss_signature -encode_dss_signature = asn1.encode_dss_signature - - -class Prehashed: - def __init__(self, algorithm: hashes.HashAlgorithm): - if not isinstance(algorithm, hashes.HashAlgorithm): - raise TypeError("Expected instance of HashAlgorithm.") - - self._algorithm = algorithm - self._digest_size = algorithm.digest_size - - @property - def digest_size(self) -> int: - return self._digest_size diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/x25519.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/x25519.py deleted file mode 100644 index a4993766b07691cf7795c0540a3e87ed4b2245e0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/x25519.py +++ /dev/null @@ -1,122 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc - -from cryptography.exceptions import UnsupportedAlgorithm, _Reasons -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.hazmat.primitives import _serialization -from cryptography.utils import Buffer - - -class X25519PublicKey(metaclass=abc.ABCMeta): - @classmethod - def from_public_bytes(cls, data: bytes) -> X25519PublicKey: - from cryptography.hazmat.backends.openssl.backend import backend - - if not backend.x25519_supported(): - raise UnsupportedAlgorithm( - "X25519 is not supported by this version of OpenSSL.", - _Reasons.UNSUPPORTED_EXCHANGE_ALGORITHM, - ) - - return rust_openssl.x25519.from_public_bytes(data) - - @abc.abstractmethod - def public_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PublicFormat, - ) -> bytes: - """ - The serialized bytes of the public key. - """ - - @abc.abstractmethod - def public_bytes_raw(self) -> bytes: - """ - The raw bytes of the public key. - Equivalent to public_bytes(Raw, Raw). - """ - - @abc.abstractmethod - def __eq__(self, other: object) -> bool: - """ - Checks equality. - """ - - @abc.abstractmethod - def __copy__(self) -> X25519PublicKey: - """ - Returns a copy. - """ - - -X25519PublicKey.register(rust_openssl.x25519.X25519PublicKey) - - -class X25519PrivateKey(metaclass=abc.ABCMeta): - @classmethod - def generate(cls) -> X25519PrivateKey: - from cryptography.hazmat.backends.openssl.backend import backend - - if not backend.x25519_supported(): - raise UnsupportedAlgorithm( - "X25519 is not supported by this version of OpenSSL.", - _Reasons.UNSUPPORTED_EXCHANGE_ALGORITHM, - ) - return rust_openssl.x25519.generate_key() - - @classmethod - def from_private_bytes(cls, data: Buffer) -> X25519PrivateKey: - from cryptography.hazmat.backends.openssl.backend import backend - - if not backend.x25519_supported(): - raise UnsupportedAlgorithm( - "X25519 is not supported by this version of OpenSSL.", - _Reasons.UNSUPPORTED_EXCHANGE_ALGORITHM, - ) - - return rust_openssl.x25519.from_private_bytes(data) - - @abc.abstractmethod - def public_key(self) -> X25519PublicKey: - """ - Returns the public key associated with this private key - """ - - @abc.abstractmethod - def private_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PrivateFormat, - encryption_algorithm: _serialization.KeySerializationEncryption, - ) -> bytes: - """ - The serialized bytes of the private key. - """ - - @abc.abstractmethod - def private_bytes_raw(self) -> bytes: - """ - The raw bytes of the private key. - Equivalent to private_bytes(Raw, Raw, NoEncryption()). - """ - - @abc.abstractmethod - def exchange(self, peer_public_key: X25519PublicKey) -> bytes: - """ - Performs a key exchange operation using the provided peer's public key. - """ - - @abc.abstractmethod - def __copy__(self) -> X25519PrivateKey: - """ - Returns a copy. - """ - - -X25519PrivateKey.register(rust_openssl.x25519.X25519PrivateKey) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/x448.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/x448.py deleted file mode 100644 index c6fd71ba5003ae757f9c0c842bb16ef0739a1fcf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/asymmetric/x448.py +++ /dev/null @@ -1,125 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc - -from cryptography.exceptions import UnsupportedAlgorithm, _Reasons -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.hazmat.primitives import _serialization -from cryptography.utils import Buffer - - -class X448PublicKey(metaclass=abc.ABCMeta): - @classmethod - def from_public_bytes(cls, data: bytes) -> X448PublicKey: - from cryptography.hazmat.backends.openssl.backend import backend - - if not backend.x448_supported(): - raise UnsupportedAlgorithm( - "X448 is not supported by this version of OpenSSL.", - _Reasons.UNSUPPORTED_EXCHANGE_ALGORITHM, - ) - - return rust_openssl.x448.from_public_bytes(data) - - @abc.abstractmethod - def public_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PublicFormat, - ) -> bytes: - """ - The serialized bytes of the public key. - """ - - @abc.abstractmethod - def public_bytes_raw(self) -> bytes: - """ - The raw bytes of the public key. - Equivalent to public_bytes(Raw, Raw). - """ - - @abc.abstractmethod - def __eq__(self, other: object) -> bool: - """ - Checks equality. - """ - - @abc.abstractmethod - def __copy__(self) -> X448PublicKey: - """ - Returns a copy. - """ - - -if hasattr(rust_openssl, "x448"): - X448PublicKey.register(rust_openssl.x448.X448PublicKey) - - -class X448PrivateKey(metaclass=abc.ABCMeta): - @classmethod - def generate(cls) -> X448PrivateKey: - from cryptography.hazmat.backends.openssl.backend import backend - - if not backend.x448_supported(): - raise UnsupportedAlgorithm( - "X448 is not supported by this version of OpenSSL.", - _Reasons.UNSUPPORTED_EXCHANGE_ALGORITHM, - ) - - return rust_openssl.x448.generate_key() - - @classmethod - def from_private_bytes(cls, data: Buffer) -> X448PrivateKey: - from cryptography.hazmat.backends.openssl.backend import backend - - if not backend.x448_supported(): - raise UnsupportedAlgorithm( - "X448 is not supported by this version of OpenSSL.", - _Reasons.UNSUPPORTED_EXCHANGE_ALGORITHM, - ) - - return rust_openssl.x448.from_private_bytes(data) - - @abc.abstractmethod - def public_key(self) -> X448PublicKey: - """ - Returns the public key associated with this private key - """ - - @abc.abstractmethod - def private_bytes( - self, - encoding: _serialization.Encoding, - format: _serialization.PrivateFormat, - encryption_algorithm: _serialization.KeySerializationEncryption, - ) -> bytes: - """ - The serialized bytes of the private key. - """ - - @abc.abstractmethod - def private_bytes_raw(self) -> bytes: - """ - The raw bytes of the private key. - Equivalent to private_bytes(Raw, Raw, NoEncryption()). - """ - - @abc.abstractmethod - def exchange(self, peer_public_key: X448PublicKey) -> bytes: - """ - Performs a key exchange operation using the provided peer's public key. - """ - - @abc.abstractmethod - def __copy__(self) -> X448PrivateKey: - """ - Returns a copy. - """ - - -if hasattr(rust_openssl, "x448"): - X448PrivateKey.register(rust_openssl.x448.X448PrivateKey) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__init__.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__init__.py deleted file mode 100644 index 10c15d0f5cb36f19ce23b3a4255e92ef27998c9b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__init__.py +++ /dev/null @@ -1,27 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography.hazmat.primitives._cipheralgorithm import ( - BlockCipherAlgorithm, - CipherAlgorithm, -) -from cryptography.hazmat.primitives.ciphers.base import ( - AEADCipherContext, - AEADDecryptionContext, - AEADEncryptionContext, - Cipher, - CipherContext, -) - -__all__ = [ - "AEADCipherContext", - "AEADDecryptionContext", - "AEADEncryptionContext", - "BlockCipherAlgorithm", - "Cipher", - "CipherAlgorithm", - "CipherContext", -] diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 7879c5eac1002f6b731645eb2af702c4eff1c2b4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/aead.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/aead.cpython-310.pyc deleted file mode 100644 index fa17b6aa1982d43f093397b9e685eced84827cb2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/aead.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/algorithms.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/algorithms.cpython-310.pyc deleted file mode 100644 index c59265ddbefaa86016c752caf06c523179d1a09a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/algorithms.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/base.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/base.cpython-310.pyc deleted file mode 100644 index 33f438b84465282dcf693561723fa06c2573f398..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/modes.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/modes.cpython-310.pyc deleted file mode 100644 index 33011e620c4b47fb4a63688ee89ff4b259d98f78..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/__pycache__/modes.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/aead.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/aead.py deleted file mode 100644 index c8a582d7844ddf57670ba657b69b5ac1f20af1d1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/aead.py +++ /dev/null @@ -1,23 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography.hazmat.bindings._rust import openssl as rust_openssl - -__all__ = [ - "AESCCM", - "AESGCM", - "AESGCMSIV", - "AESOCB3", - "AESSIV", - "ChaCha20Poly1305", -] - -AESGCM = rust_openssl.aead.AESGCM -ChaCha20Poly1305 = rust_openssl.aead.ChaCha20Poly1305 -AESCCM = rust_openssl.aead.AESCCM -AESSIV = rust_openssl.aead.AESSIV -AESOCB3 = rust_openssl.aead.AESOCB3 -AESGCMSIV = rust_openssl.aead.AESGCMSIV diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/algorithms.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/algorithms.py deleted file mode 100644 index 9d650045b0fe15753043878b0e17a932009c1457..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/algorithms.py +++ /dev/null @@ -1,183 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography import utils -from cryptography.hazmat.decrepit.ciphers.algorithms import ( - ARC4 as ARC4, -) -from cryptography.hazmat.decrepit.ciphers.algorithms import ( - CAST5 as CAST5, -) -from cryptography.hazmat.decrepit.ciphers.algorithms import ( - IDEA as IDEA, -) -from cryptography.hazmat.decrepit.ciphers.algorithms import ( - SEED as SEED, -) -from cryptography.hazmat.decrepit.ciphers.algorithms import ( - Blowfish as Blowfish, -) -from cryptography.hazmat.decrepit.ciphers.algorithms import ( - TripleDES as TripleDES, -) -from cryptography.hazmat.primitives._cipheralgorithm import _verify_key_size -from cryptography.hazmat.primitives.ciphers import ( - BlockCipherAlgorithm, - CipherAlgorithm, -) - - -class AES(BlockCipherAlgorithm): - name = "AES" - block_size = 128 - # 512 added to support AES-256-XTS, which uses 512-bit keys - key_sizes = frozenset([128, 192, 256, 512]) - - def __init__(self, key: utils.Buffer): - self.key = _verify_key_size(self, key) - - @property - def key_size(self) -> int: - return len(self.key) * 8 - - -class AES128(BlockCipherAlgorithm): - name = "AES" - block_size = 128 - key_sizes = frozenset([128]) - key_size = 128 - - def __init__(self, key: utils.Buffer): - self.key = _verify_key_size(self, key) - - -class AES256(BlockCipherAlgorithm): - name = "AES" - block_size = 128 - key_sizes = frozenset([256]) - key_size = 256 - - def __init__(self, key: utils.Buffer): - self.key = _verify_key_size(self, key) - - -class Camellia(BlockCipherAlgorithm): - name = "camellia" - block_size = 128 - key_sizes = frozenset([128, 192, 256]) - - def __init__(self, key: utils.Buffer): - self.key = _verify_key_size(self, key) - - @property - def key_size(self) -> int: - return len(self.key) * 8 - - -utils.deprecated( - ARC4, - __name__, - "ARC4 has been moved to " - "cryptography.hazmat.decrepit.ciphers.algorithms.ARC4 and " - "will be removed from " - "cryptography.hazmat.primitives.ciphers.algorithms in 48.0.0.", - utils.DeprecatedIn43, - name="ARC4", -) - - -utils.deprecated( - TripleDES, - __name__, - "TripleDES has been moved to " - "cryptography.hazmat.decrepit.ciphers.algorithms.TripleDES and " - "will be removed from " - "cryptography.hazmat.primitives.ciphers.algorithms in 48.0.0.", - utils.DeprecatedIn43, - name="TripleDES", -) - -utils.deprecated( - Blowfish, - __name__, - "Blowfish has been moved to " - "cryptography.hazmat.decrepit.ciphers.algorithms.Blowfish and " - "will be removed from " - "cryptography.hazmat.primitives.ciphers.algorithms in 45.0.0.", - utils.DeprecatedIn37, - name="Blowfish", -) - - -utils.deprecated( - CAST5, - __name__, - "CAST5 has been moved to " - "cryptography.hazmat.decrepit.ciphers.algorithms.CAST5 and " - "will be removed from " - "cryptography.hazmat.primitives.ciphers.algorithms in 45.0.0.", - utils.DeprecatedIn37, - name="CAST5", -) - - -utils.deprecated( - IDEA, - __name__, - "IDEA has been moved to " - "cryptography.hazmat.decrepit.ciphers.algorithms.IDEA and " - "will be removed from " - "cryptography.hazmat.primitives.ciphers.algorithms in 45.0.0.", - utils.DeprecatedIn37, - name="IDEA", -) - - -utils.deprecated( - SEED, - __name__, - "SEED has been moved to " - "cryptography.hazmat.decrepit.ciphers.algorithms.SEED and " - "will be removed from " - "cryptography.hazmat.primitives.ciphers.algorithms in 45.0.0.", - utils.DeprecatedIn37, - name="SEED", -) - - -class ChaCha20(CipherAlgorithm): - name = "ChaCha20" - key_sizes = frozenset([256]) - - def __init__(self, key: utils.Buffer, nonce: utils.Buffer): - self.key = _verify_key_size(self, key) - utils._check_byteslike("nonce", nonce) - - if len(nonce) != 16: - raise ValueError("nonce must be 128-bits (16 bytes)") - - self._nonce = nonce - - @property - def nonce(self) -> utils.Buffer: - return self._nonce - - @property - def key_size(self) -> int: - return len(self.key) * 8 - - -class SM4(BlockCipherAlgorithm): - name = "SM4" - block_size = 128 - key_sizes = frozenset([128]) - - def __init__(self, key: bytes): - self.key = _verify_key_size(self, key) - - @property - def key_size(self) -> int: - return len(self.key) * 8 diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/base.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/base.py deleted file mode 100644 index 24fceea236cb0f385f4b8c8c8e2e5772d018681b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/base.py +++ /dev/null @@ -1,146 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc -import typing - -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.hazmat.primitives._cipheralgorithm import CipherAlgorithm -from cryptography.hazmat.primitives.ciphers import modes -from cryptography.utils import Buffer - - -class CipherContext(metaclass=abc.ABCMeta): - @abc.abstractmethod - def update(self, data: Buffer) -> bytes: - """ - Processes the provided bytes through the cipher and returns the results - as bytes. - """ - - @abc.abstractmethod - def update_into(self, data: Buffer, buf: Buffer) -> int: - """ - Processes the provided bytes and writes the resulting data into the - provided buffer. Returns the number of bytes written. - """ - - @abc.abstractmethod - def finalize(self) -> bytes: - """ - Returns the results of processing the final block as bytes. - """ - - @abc.abstractmethod - def reset_nonce(self, nonce: bytes) -> None: - """ - Resets the nonce for the cipher context to the provided value. - Raises an exception if it does not support reset or if the - provided nonce does not have a valid length. - """ - - -class AEADCipherContext(CipherContext, metaclass=abc.ABCMeta): - @abc.abstractmethod - def authenticate_additional_data(self, data: Buffer) -> None: - """ - Authenticates the provided bytes. - """ - - -class AEADDecryptionContext(AEADCipherContext, metaclass=abc.ABCMeta): - @abc.abstractmethod - def finalize_with_tag(self, tag: bytes) -> bytes: - """ - Returns the results of processing the final block as bytes and allows - delayed passing of the authentication tag. - """ - - -class AEADEncryptionContext(AEADCipherContext, metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def tag(self) -> bytes: - """ - Returns tag bytes. This is only available after encryption is - finalized. - """ - - -Mode = typing.TypeVar( - "Mode", bound=typing.Optional[modes.Mode], covariant=True -) - - -class Cipher(typing.Generic[Mode]): - def __init__( - self, - algorithm: CipherAlgorithm, - mode: Mode, - backend: typing.Any = None, - ) -> None: - if not isinstance(algorithm, CipherAlgorithm): - raise TypeError("Expected interface of CipherAlgorithm.") - - if mode is not None: - # mypy needs this assert to narrow the type from our generic - # type. Maybe it won't some time in the future. - assert isinstance(mode, modes.Mode) - mode.validate_for_algorithm(algorithm) - - self.algorithm = algorithm - self.mode = mode - - @typing.overload - def encryptor( - self: Cipher[modes.ModeWithAuthenticationTag], - ) -> AEADEncryptionContext: ... - - @typing.overload - def encryptor( - self: _CIPHER_TYPE, - ) -> CipherContext: ... - - def encryptor(self): - if isinstance(self.mode, modes.ModeWithAuthenticationTag): - if self.mode.tag is not None: - raise ValueError( - "Authentication tag must be None when encrypting." - ) - - return rust_openssl.ciphers.create_encryption_ctx( - self.algorithm, self.mode - ) - - @typing.overload - def decryptor( - self: Cipher[modes.ModeWithAuthenticationTag], - ) -> AEADDecryptionContext: ... - - @typing.overload - def decryptor( - self: _CIPHER_TYPE, - ) -> CipherContext: ... - - def decryptor(self): - return rust_openssl.ciphers.create_decryption_ctx( - self.algorithm, self.mode - ) - - -_CIPHER_TYPE = Cipher[ - typing.Union[ - modes.ModeWithNonce, - modes.ModeWithTweak, - modes.ECB, - modes.ModeWithInitializationVector, - None, - ] -] - -CipherContext.register(rust_openssl.ciphers.CipherContext) -AEADEncryptionContext.register(rust_openssl.ciphers.AEADEncryptionContext) -AEADDecryptionContext.register(rust_openssl.ciphers.AEADDecryptionContext) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/modes.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/modes.py deleted file mode 100644 index 36c555c68a90e52e2eaf355ab1cf170166aeb95c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/ciphers/modes.py +++ /dev/null @@ -1,268 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc - -from cryptography import utils -from cryptography.exceptions import UnsupportedAlgorithm, _Reasons -from cryptography.hazmat.primitives._cipheralgorithm import ( - BlockCipherAlgorithm, - CipherAlgorithm, -) -from cryptography.hazmat.primitives.ciphers import algorithms - - -class Mode(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def name(self) -> str: - """ - A string naming this mode (e.g. "ECB", "CBC"). - """ - - @abc.abstractmethod - def validate_for_algorithm(self, algorithm: CipherAlgorithm) -> None: - """ - Checks that all the necessary invariants of this (mode, algorithm) - combination are met. - """ - - -class ModeWithInitializationVector(Mode, metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def initialization_vector(self) -> utils.Buffer: - """ - The value of the initialization vector for this mode as bytes. - """ - - -class ModeWithTweak(Mode, metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def tweak(self) -> utils.Buffer: - """ - The value of the tweak for this mode as bytes. - """ - - -class ModeWithNonce(Mode, metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def nonce(self) -> utils.Buffer: - """ - The value of the nonce for this mode as bytes. - """ - - -class ModeWithAuthenticationTag(Mode, metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def tag(self) -> bytes | None: - """ - The value of the tag supplied to the constructor of this mode. - """ - - -def _check_aes_key_length(self: Mode, algorithm: CipherAlgorithm) -> None: - if algorithm.key_size > 256 and algorithm.name == "AES": - raise ValueError( - "Only 128, 192, and 256 bit keys are allowed for this AES mode" - ) - - -def _check_iv_length( - self: ModeWithInitializationVector, algorithm: BlockCipherAlgorithm -) -> None: - iv_len = len(self.initialization_vector) - if iv_len * 8 != algorithm.block_size: - raise ValueError(f"Invalid IV size ({iv_len}) for {self.name}.") - - -def _check_nonce_length( - nonce: utils.Buffer, name: str, algorithm: CipherAlgorithm -) -> None: - if not isinstance(algorithm, BlockCipherAlgorithm): - raise UnsupportedAlgorithm( - f"{name} requires a block cipher algorithm", - _Reasons.UNSUPPORTED_CIPHER, - ) - if len(nonce) * 8 != algorithm.block_size: - raise ValueError(f"Invalid nonce size ({len(nonce)}) for {name}.") - - -def _check_iv_and_key_length( - self: ModeWithInitializationVector, algorithm: CipherAlgorithm -) -> None: - if not isinstance(algorithm, BlockCipherAlgorithm): - raise UnsupportedAlgorithm( - f"{self} requires a block cipher algorithm", - _Reasons.UNSUPPORTED_CIPHER, - ) - _check_aes_key_length(self, algorithm) - _check_iv_length(self, algorithm) - - -class CBC(ModeWithInitializationVector): - name = "CBC" - - def __init__(self, initialization_vector: utils.Buffer): - utils._check_byteslike("initialization_vector", initialization_vector) - self._initialization_vector = initialization_vector - - @property - def initialization_vector(self) -> utils.Buffer: - return self._initialization_vector - - validate_for_algorithm = _check_iv_and_key_length - - -class XTS(ModeWithTweak): - name = "XTS" - - def __init__(self, tweak: utils.Buffer): - utils._check_byteslike("tweak", tweak) - - if len(tweak) != 16: - raise ValueError("tweak must be 128-bits (16 bytes)") - - self._tweak = tweak - - @property - def tweak(self) -> utils.Buffer: - return self._tweak - - def validate_for_algorithm(self, algorithm: CipherAlgorithm) -> None: - if isinstance(algorithm, (algorithms.AES128, algorithms.AES256)): - raise TypeError( - "The AES128 and AES256 classes do not support XTS, please use " - "the standard AES class instead." - ) - - if algorithm.key_size not in (256, 512): - raise ValueError( - "The XTS specification requires a 256-bit key for AES-128-XTS" - " and 512-bit key for AES-256-XTS" - ) - - -class ECB(Mode): - name = "ECB" - - validate_for_algorithm = _check_aes_key_length - - -class OFB(ModeWithInitializationVector): - name = "OFB" - - def __init__(self, initialization_vector: utils.Buffer): - utils._check_byteslike("initialization_vector", initialization_vector) - self._initialization_vector = initialization_vector - - @property - def initialization_vector(self) -> utils.Buffer: - return self._initialization_vector - - validate_for_algorithm = _check_iv_and_key_length - - -class CFB(ModeWithInitializationVector): - name = "CFB" - - def __init__(self, initialization_vector: utils.Buffer): - utils._check_byteslike("initialization_vector", initialization_vector) - self._initialization_vector = initialization_vector - - @property - def initialization_vector(self) -> utils.Buffer: - return self._initialization_vector - - validate_for_algorithm = _check_iv_and_key_length - - -class CFB8(ModeWithInitializationVector): - name = "CFB8" - - def __init__(self, initialization_vector: utils.Buffer): - utils._check_byteslike("initialization_vector", initialization_vector) - self._initialization_vector = initialization_vector - - @property - def initialization_vector(self) -> utils.Buffer: - return self._initialization_vector - - validate_for_algorithm = _check_iv_and_key_length - - -class CTR(ModeWithNonce): - name = "CTR" - - def __init__(self, nonce: utils.Buffer): - utils._check_byteslike("nonce", nonce) - self._nonce = nonce - - @property - def nonce(self) -> utils.Buffer: - return self._nonce - - def validate_for_algorithm(self, algorithm: CipherAlgorithm) -> None: - _check_aes_key_length(self, algorithm) - _check_nonce_length(self.nonce, self.name, algorithm) - - -class GCM(ModeWithInitializationVector, ModeWithAuthenticationTag): - name = "GCM" - _MAX_ENCRYPTED_BYTES = (2**39 - 256) // 8 - _MAX_AAD_BYTES = (2**64) // 8 - - def __init__( - self, - initialization_vector: utils.Buffer, - tag: bytes | None = None, - min_tag_length: int = 16, - ): - # OpenSSL 3.0.0 constrains GCM IVs to [64, 1024] bits inclusive - # This is a sane limit anyway so we'll enforce it here. - utils._check_byteslike("initialization_vector", initialization_vector) - if len(initialization_vector) < 8 or len(initialization_vector) > 128: - raise ValueError( - "initialization_vector must be between 8 and 128 bytes (64 " - "and 1024 bits)." - ) - self._initialization_vector = initialization_vector - if tag is not None: - utils._check_bytes("tag", tag) - if min_tag_length < 4: - raise ValueError("min_tag_length must be >= 4") - if len(tag) < min_tag_length: - raise ValueError( - f"Authentication tag must be {min_tag_length} bytes or " - "longer." - ) - self._tag = tag - self._min_tag_length = min_tag_length - - @property - def tag(self) -> bytes | None: - return self._tag - - @property - def initialization_vector(self) -> utils.Buffer: - return self._initialization_vector - - def validate_for_algorithm(self, algorithm: CipherAlgorithm) -> None: - _check_aes_key_length(self, algorithm) - if not isinstance(algorithm, BlockCipherAlgorithm): - raise UnsupportedAlgorithm( - "GCM requires a block cipher algorithm", - _Reasons.UNSUPPORTED_CIPHER, - ) - block_size_bytes = algorithm.block_size // 8 - if self._tag is not None and len(self._tag) > block_size_bytes: - raise ValueError( - f"Authentication tag cannot be more than {block_size_bytes} " - "bytes." - ) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/cmac.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/cmac.py deleted file mode 100644 index 2c67ce2206e4c36c7ab8c4a49ba9d87f07529598..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/cmac.py +++ /dev/null @@ -1,10 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography.hazmat.bindings._rust import openssl as rust_openssl - -__all__ = ["CMAC"] -CMAC = rust_openssl.cmac.CMAC diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/constant_time.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/constant_time.py deleted file mode 100644 index 3975c7147eb92b56685423aa1c5810adcf253a23..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/constant_time.py +++ /dev/null @@ -1,14 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import hmac - - -def bytes_eq(a: bytes, b: bytes) -> bool: - if not isinstance(a, bytes) or not isinstance(b, bytes): - raise TypeError("a and b must be bytes.") - - return hmac.compare_digest(a, b) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/hashes.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/hashes.py deleted file mode 100644 index 4b55ec33dbff64bb848c8633619bbbd36362936a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/hashes.py +++ /dev/null @@ -1,246 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc - -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.utils import Buffer - -__all__ = [ - "MD5", - "SHA1", - "SHA3_224", - "SHA3_256", - "SHA3_384", - "SHA3_512", - "SHA224", - "SHA256", - "SHA384", - "SHA512", - "SHA512_224", - "SHA512_256", - "SHAKE128", - "SHAKE256", - "SM3", - "BLAKE2b", - "BLAKE2s", - "ExtendableOutputFunction", - "Hash", - "HashAlgorithm", - "HashContext", - "XOFHash", -] - - -class HashAlgorithm(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def name(self) -> str: - """ - A string naming this algorithm (e.g. "sha256", "md5"). - """ - - @property - @abc.abstractmethod - def digest_size(self) -> int: - """ - The size of the resulting digest in bytes. - """ - - @property - @abc.abstractmethod - def block_size(self) -> int | None: - """ - The internal block size of the hash function, or None if the hash - function does not use blocks internally (e.g. SHA3). - """ - - -class HashContext(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def algorithm(self) -> HashAlgorithm: - """ - A HashAlgorithm that will be used by this context. - """ - - @abc.abstractmethod - def update(self, data: Buffer) -> None: - """ - Processes the provided bytes through the hash. - """ - - @abc.abstractmethod - def finalize(self) -> bytes: - """ - Finalizes the hash context and returns the hash digest as bytes. - """ - - @abc.abstractmethod - def copy(self) -> HashContext: - """ - Return a HashContext that is a copy of the current context. - """ - - -Hash = rust_openssl.hashes.Hash -HashContext.register(Hash) - -XOFHash = rust_openssl.hashes.XOFHash - - -class ExtendableOutputFunction(metaclass=abc.ABCMeta): - """ - An interface for extendable output functions. - """ - - -class SHA1(HashAlgorithm): - name = "sha1" - digest_size = 20 - block_size = 64 - - -class SHA512_224(HashAlgorithm): # noqa: N801 - name = "sha512-224" - digest_size = 28 - block_size = 128 - - -class SHA512_256(HashAlgorithm): # noqa: N801 - name = "sha512-256" - digest_size = 32 - block_size = 128 - - -class SHA224(HashAlgorithm): - name = "sha224" - digest_size = 28 - block_size = 64 - - -class SHA256(HashAlgorithm): - name = "sha256" - digest_size = 32 - block_size = 64 - - -class SHA384(HashAlgorithm): - name = "sha384" - digest_size = 48 - block_size = 128 - - -class SHA512(HashAlgorithm): - name = "sha512" - digest_size = 64 - block_size = 128 - - -class SHA3_224(HashAlgorithm): # noqa: N801 - name = "sha3-224" - digest_size = 28 - block_size = None - - -class SHA3_256(HashAlgorithm): # noqa: N801 - name = "sha3-256" - digest_size = 32 - block_size = None - - -class SHA3_384(HashAlgorithm): # noqa: N801 - name = "sha3-384" - digest_size = 48 - block_size = None - - -class SHA3_512(HashAlgorithm): # noqa: N801 - name = "sha3-512" - digest_size = 64 - block_size = None - - -class SHAKE128(HashAlgorithm, ExtendableOutputFunction): - name = "shake128" - block_size = None - - def __init__(self, digest_size: int): - if not isinstance(digest_size, int): - raise TypeError("digest_size must be an integer") - - if digest_size < 1: - raise ValueError("digest_size must be a positive integer") - - self._digest_size = digest_size - - @property - def digest_size(self) -> int: - return self._digest_size - - -class SHAKE256(HashAlgorithm, ExtendableOutputFunction): - name = "shake256" - block_size = None - - def __init__(self, digest_size: int): - if not isinstance(digest_size, int): - raise TypeError("digest_size must be an integer") - - if digest_size < 1: - raise ValueError("digest_size must be a positive integer") - - self._digest_size = digest_size - - @property - def digest_size(self) -> int: - return self._digest_size - - -class MD5(HashAlgorithm): - name = "md5" - digest_size = 16 - block_size = 64 - - -class BLAKE2b(HashAlgorithm): - name = "blake2b" - _max_digest_size = 64 - _min_digest_size = 1 - block_size = 128 - - def __init__(self, digest_size: int): - if digest_size != 64: - raise ValueError("Digest size must be 64") - - self._digest_size = digest_size - - @property - def digest_size(self) -> int: - return self._digest_size - - -class BLAKE2s(HashAlgorithm): - name = "blake2s" - block_size = 64 - _max_digest_size = 32 - _min_digest_size = 1 - - def __init__(self, digest_size: int): - if digest_size != 32: - raise ValueError("Digest size must be 32") - - self._digest_size = digest_size - - @property - def digest_size(self) -> int: - return self._digest_size - - -class SM3(HashAlgorithm): - name = "sm3" - digest_size = 32 - block_size = 64 diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/hmac.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/hmac.py deleted file mode 100644 index a9442d59ab474ac680ff7253f2ca0e67d0e93c1a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/hmac.py +++ /dev/null @@ -1,13 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.hazmat.primitives import hashes - -__all__ = ["HMAC"] - -HMAC = rust_openssl.hmac.HMAC -hashes.HashContext.register(HMAC) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__init__.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__init__.py deleted file mode 100644 index 79bb459f01ec288d67b4f7c20fd5af436fcab4a5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc - - -class KeyDerivationFunction(metaclass=abc.ABCMeta): - @abc.abstractmethod - def derive(self, key_material: bytes) -> bytes: - """ - Deterministically generates and returns a new key based on the existing - key material. - """ - - @abc.abstractmethod - def verify(self, key_material: bytes, expected_key: bytes) -> None: - """ - Checks whether the key generated by the key material matches the - expected derived key. Raises an exception if they do not match. - """ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index bca8d7cabcb57958302f09da3b8faac199ce33e9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/argon2.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/argon2.cpython-310.pyc deleted file mode 100644 index 9b20c7cfdac019cf6d780924ce295748209c0e3d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/argon2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/concatkdf.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/concatkdf.cpython-310.pyc deleted file mode 100644 index df173da0ebb585a6b320b1cc3c92ebbf7e5a53c0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/concatkdf.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/hkdf.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/hkdf.cpython-310.pyc deleted file mode 100644 index 776789d143765decfe4dbeea9eb7e066bae1aeeb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/hkdf.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/kbkdf.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/kbkdf.cpython-310.pyc deleted file mode 100644 index 65668ea128403d9da607004fdc94e661d5ca3928..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/kbkdf.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/pbkdf2.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/pbkdf2.cpython-310.pyc deleted file mode 100644 index d13d9c0d283936c635aff918ee2df64a1a755cae..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/pbkdf2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/scrypt.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/scrypt.cpython-310.pyc deleted file mode 100644 index 38396569466b8f1e333326b89cedf4b652224ed7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/scrypt.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/x963kdf.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/x963kdf.cpython-310.pyc deleted file mode 100644 index 5af8412cbc3dbfff181f5ab8fb5de9303d6aa214..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/__pycache__/x963kdf.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/argon2.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/argon2.py deleted file mode 100644 index 405fc8dff26876404dc27279b815d4dea6508340..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/argon2.py +++ /dev/null @@ -1,13 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.hazmat.primitives.kdf import KeyDerivationFunction - -Argon2id = rust_openssl.kdf.Argon2id -KeyDerivationFunction.register(Argon2id) - -__all__ = ["Argon2id"] diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/concatkdf.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/concatkdf.py deleted file mode 100644 index 1b928415c5c11369792a3522f001a557a9d0c5cf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/concatkdf.py +++ /dev/null @@ -1,125 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import typing -from collections.abc import Callable - -from cryptography import utils -from cryptography.exceptions import AlreadyFinalized, InvalidKey -from cryptography.hazmat.primitives import constant_time, hashes, hmac -from cryptography.hazmat.primitives.kdf import KeyDerivationFunction - - -def _int_to_u32be(n: int) -> bytes: - return n.to_bytes(length=4, byteorder="big") - - -def _common_args_checks( - algorithm: hashes.HashAlgorithm, - length: int, - otherinfo: bytes | None, -) -> None: - max_length = algorithm.digest_size * (2**32 - 1) - if length > max_length: - raise ValueError(f"Cannot derive keys larger than {max_length} bits.") - if otherinfo is not None: - utils._check_bytes("otherinfo", otherinfo) - - -def _concatkdf_derive( - key_material: utils.Buffer, - length: int, - auxfn: Callable[[], hashes.HashContext], - otherinfo: bytes, -) -> bytes: - utils._check_byteslike("key_material", key_material) - output = [b""] - outlen = 0 - counter = 1 - - while length > outlen: - h = auxfn() - h.update(_int_to_u32be(counter)) - h.update(key_material) - h.update(otherinfo) - output.append(h.finalize()) - outlen += len(output[-1]) - counter += 1 - - return b"".join(output)[:length] - - -class ConcatKDFHash(KeyDerivationFunction): - def __init__( - self, - algorithm: hashes.HashAlgorithm, - length: int, - otherinfo: bytes | None, - backend: typing.Any = None, - ): - _common_args_checks(algorithm, length, otherinfo) - self._algorithm = algorithm - self._length = length - self._otherinfo: bytes = otherinfo if otherinfo is not None else b"" - - self._used = False - - def _hash(self) -> hashes.Hash: - return hashes.Hash(self._algorithm) - - def derive(self, key_material: utils.Buffer) -> bytes: - if self._used: - raise AlreadyFinalized - self._used = True - return _concatkdf_derive( - key_material, self._length, self._hash, self._otherinfo - ) - - def verify(self, key_material: bytes, expected_key: bytes) -> None: - if not constant_time.bytes_eq(self.derive(key_material), expected_key): - raise InvalidKey - - -class ConcatKDFHMAC(KeyDerivationFunction): - def __init__( - self, - algorithm: hashes.HashAlgorithm, - length: int, - salt: bytes | None, - otherinfo: bytes | None, - backend: typing.Any = None, - ): - _common_args_checks(algorithm, length, otherinfo) - self._algorithm = algorithm - self._length = length - self._otherinfo: bytes = otherinfo if otherinfo is not None else b"" - - if algorithm.block_size is None: - raise TypeError(f"{algorithm.name} is unsupported for ConcatKDF") - - if salt is None: - salt = b"\x00" * algorithm.block_size - else: - utils._check_bytes("salt", salt) - - self._salt = salt - - self._used = False - - def _hmac(self) -> hmac.HMAC: - return hmac.HMAC(self._salt, self._algorithm) - - def derive(self, key_material: utils.Buffer) -> bytes: - if self._used: - raise AlreadyFinalized - self._used = True - return _concatkdf_derive( - key_material, self._length, self._hmac, self._otherinfo - ) - - def verify(self, key_material: bytes, expected_key: bytes) -> None: - if not constant_time.bytes_eq(self.derive(key_material), expected_key): - raise InvalidKey diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/hkdf.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/hkdf.py deleted file mode 100644 index ce11c54baa317b9c4470a98172c07694a7cab48d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/hkdf.py +++ /dev/null @@ -1,101 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import typing - -from cryptography import utils -from cryptography.exceptions import AlreadyFinalized, InvalidKey -from cryptography.hazmat.primitives import constant_time, hashes, hmac -from cryptography.hazmat.primitives.kdf import KeyDerivationFunction - - -class HKDF(KeyDerivationFunction): - def __init__( - self, - algorithm: hashes.HashAlgorithm, - length: int, - salt: bytes | None, - info: bytes | None, - backend: typing.Any = None, - ): - self._algorithm = algorithm - - if salt is None: - salt = b"\x00" * self._algorithm.digest_size - else: - utils._check_bytes("salt", salt) - - self._salt = salt - - self._hkdf_expand = HKDFExpand(self._algorithm, length, info) - - def _extract(self, key_material: utils.Buffer) -> bytes: - h = hmac.HMAC(self._salt, self._algorithm) - h.update(key_material) - return h.finalize() - - def derive(self, key_material: utils.Buffer) -> bytes: - utils._check_byteslike("key_material", key_material) - return self._hkdf_expand.derive(self._extract(key_material)) - - def verify(self, key_material: bytes, expected_key: bytes) -> None: - if not constant_time.bytes_eq(self.derive(key_material), expected_key): - raise InvalidKey - - -class HKDFExpand(KeyDerivationFunction): - def __init__( - self, - algorithm: hashes.HashAlgorithm, - length: int, - info: bytes | None, - backend: typing.Any = None, - ): - self._algorithm = algorithm - - max_length = 255 * algorithm.digest_size - - if length > max_length: - raise ValueError( - f"Cannot derive keys larger than {max_length} octets." - ) - - self._length = length - - if info is None: - info = b"" - else: - utils._check_bytes("info", info) - - self._info = info - - self._used = False - - def _expand(self, key_material: utils.Buffer) -> bytes: - output = [b""] - counter = 1 - - while self._algorithm.digest_size * (len(output) - 1) < self._length: - h = hmac.HMAC(key_material, self._algorithm) - h.update(output[-1]) - h.update(self._info) - h.update(bytes([counter])) - output.append(h.finalize()) - counter += 1 - - return b"".join(output)[: self._length] - - def derive(self, key_material: utils.Buffer) -> bytes: - utils._check_byteslike("key_material", key_material) - if self._used: - raise AlreadyFinalized - - self._used = True - return self._expand(key_material) - - def verify(self, key_material: bytes, expected_key: bytes) -> None: - if not constant_time.bytes_eq(self.derive(key_material), expected_key): - raise InvalidKey diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/kbkdf.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/kbkdf.py deleted file mode 100644 index d6c3ff383bb6ece101e62e90d9d11d920ba3d613..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/kbkdf.py +++ /dev/null @@ -1,305 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import typing -from collections.abc import Callable - -from cryptography import utils -from cryptography.exceptions import ( - AlreadyFinalized, - InvalidKey, - UnsupportedAlgorithm, - _Reasons, -) -from cryptography.hazmat.primitives import ( - ciphers, - cmac, - constant_time, - hashes, - hmac, -) -from cryptography.hazmat.primitives.kdf import KeyDerivationFunction - - -class Mode(utils.Enum): - CounterMode = "ctr" - - -class CounterLocation(utils.Enum): - BeforeFixed = "before_fixed" - AfterFixed = "after_fixed" - MiddleFixed = "middle_fixed" - - -class _KBKDFDeriver: - def __init__( - self, - prf: Callable, - mode: Mode, - length: int, - rlen: int, - llen: int | None, - location: CounterLocation, - break_location: int | None, - label: bytes | None, - context: bytes | None, - fixed: bytes | None, - ): - assert callable(prf) - - if not isinstance(mode, Mode): - raise TypeError("mode must be of type Mode") - - if not isinstance(location, CounterLocation): - raise TypeError("location must be of type CounterLocation") - - if break_location is None and location is CounterLocation.MiddleFixed: - raise ValueError("Please specify a break_location") - - if ( - break_location is not None - and location != CounterLocation.MiddleFixed - ): - raise ValueError( - "break_location is ignored when location is not" - " CounterLocation.MiddleFixed" - ) - - if break_location is not None and not isinstance(break_location, int): - raise TypeError("break_location must be an integer") - - if break_location is not None and break_location < 0: - raise ValueError("break_location must be a positive integer") - - if (label or context) and fixed: - raise ValueError( - "When supplying fixed data, label and context are ignored." - ) - - if rlen is None or not self._valid_byte_length(rlen): - raise ValueError("rlen must be between 1 and 4") - - if llen is None and fixed is None: - raise ValueError("Please specify an llen") - - if llen is not None and not isinstance(llen, int): - raise TypeError("llen must be an integer") - - if llen == 0: - raise ValueError("llen must be non-zero") - - if label is None: - label = b"" - - if context is None: - context = b"" - - utils._check_bytes("label", label) - utils._check_bytes("context", context) - self._prf = prf - self._mode = mode - self._length = length - self._rlen = rlen - self._llen = llen - self._location = location - self._break_location = break_location - self._label = label - self._context = context - self._used = False - self._fixed_data = fixed - - @staticmethod - def _valid_byte_length(value: int) -> bool: - if not isinstance(value, int): - raise TypeError("value must be of type int") - - value_bin = utils.int_to_bytes(1, value) - if not 1 <= len(value_bin) <= 4: - return False - return True - - def derive( - self, key_material: utils.Buffer, prf_output_size: int - ) -> bytes: - if self._used: - raise AlreadyFinalized - - utils._check_byteslike("key_material", key_material) - self._used = True - - # inverse floor division (equivalent to ceiling) - rounds = -(-self._length // prf_output_size) - - output = [b""] - - # For counter mode, the number of iterations shall not be - # larger than 2^r-1, where r <= 32 is the binary length of the counter - # This ensures that the counter values used as an input to the - # PRF will not repeat during a particular call to the KDF function. - r_bin = utils.int_to_bytes(1, self._rlen) - if rounds > pow(2, len(r_bin) * 8) - 1: - raise ValueError("There are too many iterations.") - - fixed = self._generate_fixed_input() - - if self._location == CounterLocation.BeforeFixed: - data_before_ctr = b"" - data_after_ctr = fixed - elif self._location == CounterLocation.AfterFixed: - data_before_ctr = fixed - data_after_ctr = b"" - else: - if isinstance( - self._break_location, int - ) and self._break_location > len(fixed): - raise ValueError("break_location offset > len(fixed)") - data_before_ctr = fixed[: self._break_location] - data_after_ctr = fixed[self._break_location :] - - for i in range(1, rounds + 1): - h = self._prf(key_material) - - counter = utils.int_to_bytes(i, self._rlen) - input_data = data_before_ctr + counter + data_after_ctr - - h.update(input_data) - - output.append(h.finalize()) - - return b"".join(output)[: self._length] - - def _generate_fixed_input(self) -> bytes: - if self._fixed_data and isinstance(self._fixed_data, bytes): - return self._fixed_data - - l_val = utils.int_to_bytes(self._length * 8, self._llen) - - return b"".join([self._label, b"\x00", self._context, l_val]) - - -class KBKDFHMAC(KeyDerivationFunction): - def __init__( - self, - algorithm: hashes.HashAlgorithm, - mode: Mode, - length: int, - rlen: int, - llen: int | None, - location: CounterLocation, - label: bytes | None, - context: bytes | None, - fixed: bytes | None, - backend: typing.Any = None, - *, - break_location: int | None = None, - ): - if not isinstance(algorithm, hashes.HashAlgorithm): - raise UnsupportedAlgorithm( - "Algorithm supplied is not a supported hash algorithm.", - _Reasons.UNSUPPORTED_HASH, - ) - - from cryptography.hazmat.backends.openssl.backend import ( - backend as ossl, - ) - - if not ossl.hmac_supported(algorithm): - raise UnsupportedAlgorithm( - "Algorithm supplied is not a supported hmac algorithm.", - _Reasons.UNSUPPORTED_HASH, - ) - - self._algorithm = algorithm - - self._deriver = _KBKDFDeriver( - self._prf, - mode, - length, - rlen, - llen, - location, - break_location, - label, - context, - fixed, - ) - - def _prf(self, key_material: bytes) -> hmac.HMAC: - return hmac.HMAC(key_material, self._algorithm) - - def derive(self, key_material: utils.Buffer) -> bytes: - return self._deriver.derive(key_material, self._algorithm.digest_size) - - def verify(self, key_material: bytes, expected_key: bytes) -> None: - if not constant_time.bytes_eq(self.derive(key_material), expected_key): - raise InvalidKey - - -class KBKDFCMAC(KeyDerivationFunction): - def __init__( - self, - algorithm, - mode: Mode, - length: int, - rlen: int, - llen: int | None, - location: CounterLocation, - label: bytes | None, - context: bytes | None, - fixed: bytes | None, - backend: typing.Any = None, - *, - break_location: int | None = None, - ): - if not issubclass( - algorithm, ciphers.BlockCipherAlgorithm - ) or not issubclass(algorithm, ciphers.CipherAlgorithm): - raise UnsupportedAlgorithm( - "Algorithm supplied is not a supported cipher algorithm.", - _Reasons.UNSUPPORTED_CIPHER, - ) - - self._algorithm = algorithm - self._cipher: ciphers.BlockCipherAlgorithm | None = None - - self._deriver = _KBKDFDeriver( - self._prf, - mode, - length, - rlen, - llen, - location, - break_location, - label, - context, - fixed, - ) - - def _prf(self, _: bytes) -> cmac.CMAC: - assert self._cipher is not None - - return cmac.CMAC(self._cipher) - - def derive(self, key_material: utils.Buffer) -> bytes: - self._cipher = self._algorithm(key_material) - - assert self._cipher is not None - - from cryptography.hazmat.backends.openssl.backend import ( - backend as ossl, - ) - - if not ossl.cmac_algorithm_supported(self._cipher): - raise UnsupportedAlgorithm( - "Algorithm supplied is not a supported cipher algorithm.", - _Reasons.UNSUPPORTED_CIPHER, - ) - - return self._deriver.derive(key_material, self._cipher.block_size // 8) - - def verify(self, key_material: bytes, expected_key: bytes) -> None: - if not constant_time.bytes_eq(self.derive(key_material), expected_key): - raise InvalidKey diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/pbkdf2.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/pbkdf2.py deleted file mode 100644 index d539f1317556e71b5636eb2f896ec2676d458af8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/pbkdf2.py +++ /dev/null @@ -1,62 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import typing - -from cryptography import utils -from cryptography.exceptions import ( - AlreadyFinalized, - InvalidKey, - UnsupportedAlgorithm, - _Reasons, -) -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.hazmat.primitives import constant_time, hashes -from cryptography.hazmat.primitives.kdf import KeyDerivationFunction - - -class PBKDF2HMAC(KeyDerivationFunction): - def __init__( - self, - algorithm: hashes.HashAlgorithm, - length: int, - salt: bytes, - iterations: int, - backend: typing.Any = None, - ): - from cryptography.hazmat.backends.openssl.backend import ( - backend as ossl, - ) - - if not ossl.pbkdf2_hmac_supported(algorithm): - raise UnsupportedAlgorithm( - f"{algorithm.name} is not supported for PBKDF2.", - _Reasons.UNSUPPORTED_HASH, - ) - self._used = False - self._algorithm = algorithm - self._length = length - utils._check_bytes("salt", salt) - self._salt = salt - self._iterations = iterations - - def derive(self, key_material: utils.Buffer) -> bytes: - if self._used: - raise AlreadyFinalized("PBKDF2 instances can only be used once.") - self._used = True - - return rust_openssl.kdf.derive_pbkdf2_hmac( - key_material, - self._algorithm, - self._salt, - self._iterations, - self._length, - ) - - def verify(self, key_material: bytes, expected_key: bytes) -> None: - derived_key = self.derive(key_material) - if not constant_time.bytes_eq(derived_key, expected_key): - raise InvalidKey("Keys do not match.") diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/scrypt.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/scrypt.py deleted file mode 100644 index f791ceea371b11c0d873ff031cf3bd653cb1b2b8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/scrypt.py +++ /dev/null @@ -1,19 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import sys - -from cryptography.hazmat.bindings._rust import openssl as rust_openssl -from cryptography.hazmat.primitives.kdf import KeyDerivationFunction - -# This is used by the scrypt tests to skip tests that require more memory -# than the MEM_LIMIT -_MEM_LIMIT = sys.maxsize // 2 - -Scrypt = rust_openssl.kdf.Scrypt -KeyDerivationFunction.register(Scrypt) - -__all__ = ["Scrypt"] diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/x963kdf.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/x963kdf.py deleted file mode 100644 index 63870cdc1f5ffb6f43ae22adda0fa00fd9931862..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/kdf/x963kdf.py +++ /dev/null @@ -1,61 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import typing - -from cryptography import utils -from cryptography.exceptions import AlreadyFinalized, InvalidKey -from cryptography.hazmat.primitives import constant_time, hashes -from cryptography.hazmat.primitives.kdf import KeyDerivationFunction - - -def _int_to_u32be(n: int) -> bytes: - return n.to_bytes(length=4, byteorder="big") - - -class X963KDF(KeyDerivationFunction): - def __init__( - self, - algorithm: hashes.HashAlgorithm, - length: int, - sharedinfo: bytes | None, - backend: typing.Any = None, - ): - max_len = algorithm.digest_size * (2**32 - 1) - if length > max_len: - raise ValueError(f"Cannot derive keys larger than {max_len} bits.") - if sharedinfo is not None: - utils._check_bytes("sharedinfo", sharedinfo) - - self._algorithm = algorithm - self._length = length - self._sharedinfo = sharedinfo - self._used = False - - def derive(self, key_material: utils.Buffer) -> bytes: - if self._used: - raise AlreadyFinalized - self._used = True - utils._check_byteslike("key_material", key_material) - output = [b""] - outlen = 0 - counter = 1 - - while self._length > outlen: - h = hashes.Hash(self._algorithm) - h.update(key_material) - h.update(_int_to_u32be(counter)) - if self._sharedinfo is not None: - h.update(self._sharedinfo) - output.append(h.finalize()) - outlen += len(output[-1]) - counter += 1 - - return b"".join(output)[: self._length] - - def verify(self, key_material: bytes, expected_key: bytes) -> None: - if not constant_time.bytes_eq(self.derive(key_material), expected_key): - raise InvalidKey diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/keywrap.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/keywrap.py deleted file mode 100644 index b93d87d31cff5cecdf6513e81f8e506f2b2b6d3e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/keywrap.py +++ /dev/null @@ -1,177 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import typing - -from cryptography.hazmat.primitives.ciphers import Cipher -from cryptography.hazmat.primitives.ciphers.algorithms import AES -from cryptography.hazmat.primitives.ciphers.modes import ECB -from cryptography.hazmat.primitives.constant_time import bytes_eq - - -def _wrap_core( - wrapping_key: bytes, - a: bytes, - r: list[bytes], -) -> bytes: - # RFC 3394 Key Wrap - 2.2.1 (index method) - encryptor = Cipher(AES(wrapping_key), ECB()).encryptor() - n = len(r) - for j in range(6): - for i in range(n): - # every encryption operation is a discrete 16 byte chunk (because - # AES has a 128-bit block size) and since we're using ECB it is - # safe to reuse the encryptor for the entire operation - b = encryptor.update(a + r[i]) - a = ( - int.from_bytes(b[:8], byteorder="big") ^ ((n * j) + i + 1) - ).to_bytes(length=8, byteorder="big") - r[i] = b[-8:] - - assert encryptor.finalize() == b"" - - return a + b"".join(r) - - -def aes_key_wrap( - wrapping_key: bytes, - key_to_wrap: bytes, - backend: typing.Any = None, -) -> bytes: - if len(wrapping_key) not in [16, 24, 32]: - raise ValueError("The wrapping key must be a valid AES key length") - - if len(key_to_wrap) < 16: - raise ValueError("The key to wrap must be at least 16 bytes") - - if len(key_to_wrap) % 8 != 0: - raise ValueError("The key to wrap must be a multiple of 8 bytes") - - a = b"\xa6\xa6\xa6\xa6\xa6\xa6\xa6\xa6" - r = [key_to_wrap[i : i + 8] for i in range(0, len(key_to_wrap), 8)] - return _wrap_core(wrapping_key, a, r) - - -def _unwrap_core( - wrapping_key: bytes, - a: bytes, - r: list[bytes], -) -> tuple[bytes, list[bytes]]: - # Implement RFC 3394 Key Unwrap - 2.2.2 (index method) - decryptor = Cipher(AES(wrapping_key), ECB()).decryptor() - n = len(r) - for j in reversed(range(6)): - for i in reversed(range(n)): - atr = ( - int.from_bytes(a, byteorder="big") ^ ((n * j) + i + 1) - ).to_bytes(length=8, byteorder="big") + r[i] - # every decryption operation is a discrete 16 byte chunk so - # it is safe to reuse the decryptor for the entire operation - b = decryptor.update(atr) - a = b[:8] - r[i] = b[-8:] - - assert decryptor.finalize() == b"" - return a, r - - -def aes_key_wrap_with_padding( - wrapping_key: bytes, - key_to_wrap: bytes, - backend: typing.Any = None, -) -> bytes: - if len(wrapping_key) not in [16, 24, 32]: - raise ValueError("The wrapping key must be a valid AES key length") - - aiv = b"\xa6\x59\x59\xa6" + len(key_to_wrap).to_bytes( - length=4, byteorder="big" - ) - # pad the key to wrap if necessary - pad = (8 - (len(key_to_wrap) % 8)) % 8 - key_to_wrap = key_to_wrap + b"\x00" * pad - if len(key_to_wrap) == 8: - # RFC 5649 - 4.1 - exactly 8 octets after padding - encryptor = Cipher(AES(wrapping_key), ECB()).encryptor() - b = encryptor.update(aiv + key_to_wrap) - assert encryptor.finalize() == b"" - return b - else: - r = [key_to_wrap[i : i + 8] for i in range(0, len(key_to_wrap), 8)] - return _wrap_core(wrapping_key, aiv, r) - - -def aes_key_unwrap_with_padding( - wrapping_key: bytes, - wrapped_key: bytes, - backend: typing.Any = None, -) -> bytes: - if len(wrapped_key) < 16: - raise InvalidUnwrap("Must be at least 16 bytes") - - if len(wrapping_key) not in [16, 24, 32]: - raise ValueError("The wrapping key must be a valid AES key length") - - if len(wrapped_key) == 16: - # RFC 5649 - 4.2 - exactly two 64-bit blocks - decryptor = Cipher(AES(wrapping_key), ECB()).decryptor() - out = decryptor.update(wrapped_key) - assert decryptor.finalize() == b"" - a = out[:8] - data = out[8:] - n = 1 - else: - r = [wrapped_key[i : i + 8] for i in range(0, len(wrapped_key), 8)] - encrypted_aiv = r.pop(0) - n = len(r) - a, r = _unwrap_core(wrapping_key, encrypted_aiv, r) - data = b"".join(r) - - # 1) Check that MSB(32,A) = A65959A6. - # 2) Check that 8*(n-1) < LSB(32,A) <= 8*n. If so, let - # MLI = LSB(32,A). - # 3) Let b = (8*n)-MLI, and then check that the rightmost b octets of - # the output data are zero. - mli = int.from_bytes(a[4:], byteorder="big") - b = (8 * n) - mli - if ( - not bytes_eq(a[:4], b"\xa6\x59\x59\xa6") - or not 8 * (n - 1) < mli <= 8 * n - or (b != 0 and not bytes_eq(data[-b:], b"\x00" * b)) - ): - raise InvalidUnwrap() - - if b == 0: - return data - else: - return data[:-b] - - -def aes_key_unwrap( - wrapping_key: bytes, - wrapped_key: bytes, - backend: typing.Any = None, -) -> bytes: - if len(wrapped_key) < 24: - raise InvalidUnwrap("Must be at least 24 bytes") - - if len(wrapped_key) % 8 != 0: - raise InvalidUnwrap("The wrapped key must be a multiple of 8 bytes") - - if len(wrapping_key) not in [16, 24, 32]: - raise ValueError("The wrapping key must be a valid AES key length") - - aiv = b"\xa6\xa6\xa6\xa6\xa6\xa6\xa6\xa6" - r = [wrapped_key[i : i + 8] for i in range(0, len(wrapped_key), 8)] - a = r.pop(0) - a, r = _unwrap_core(wrapping_key, a, r) - if not bytes_eq(a, aiv): - raise InvalidUnwrap() - - return b"".join(r) - - -class InvalidUnwrap(Exception): - pass diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/padding.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/padding.py deleted file mode 100644 index f9cd1f13c469beec5512d58b3de6ee1ced4362d5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/padding.py +++ /dev/null @@ -1,69 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc - -from cryptography import utils -from cryptography.hazmat.bindings._rust import ( - ANSIX923PaddingContext, - ANSIX923UnpaddingContext, - PKCS7PaddingContext, - PKCS7UnpaddingContext, -) - - -class PaddingContext(metaclass=abc.ABCMeta): - @abc.abstractmethod - def update(self, data: utils.Buffer) -> bytes: - """ - Pads the provided bytes and returns any available data as bytes. - """ - - @abc.abstractmethod - def finalize(self) -> bytes: - """ - Finalize the padding, returns bytes. - """ - - -def _byte_padding_check(block_size: int) -> None: - if not (0 <= block_size <= 2040): - raise ValueError("block_size must be in range(0, 2041).") - - if block_size % 8 != 0: - raise ValueError("block_size must be a multiple of 8.") - - -class PKCS7: - def __init__(self, block_size: int): - _byte_padding_check(block_size) - self.block_size = block_size - - def padder(self) -> PaddingContext: - return PKCS7PaddingContext(self.block_size) - - def unpadder(self) -> PaddingContext: - return PKCS7UnpaddingContext(self.block_size) - - -PaddingContext.register(PKCS7PaddingContext) -PaddingContext.register(PKCS7UnpaddingContext) - - -class ANSIX923: - def __init__(self, block_size: int): - _byte_padding_check(block_size) - self.block_size = block_size - - def padder(self) -> PaddingContext: - return ANSIX923PaddingContext(self.block_size) - - def unpadder(self) -> PaddingContext: - return ANSIX923UnpaddingContext(self.block_size) - - -PaddingContext.register(ANSIX923PaddingContext) -PaddingContext.register(ANSIX923UnpaddingContext) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/poly1305.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/poly1305.py deleted file mode 100644 index 7f5a77a576fd258a9a634ae3c3365fc3f7d4702a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/poly1305.py +++ /dev/null @@ -1,11 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography.hazmat.bindings._rust import openssl as rust_openssl - -__all__ = ["Poly1305"] - -Poly1305 = rust_openssl.poly1305.Poly1305 diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__init__.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__init__.py deleted file mode 100644 index 62283cc70fd6bb7d29da3dd67773761a56e86f10..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__init__.py +++ /dev/null @@ -1,65 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography.hazmat.primitives._serialization import ( - BestAvailableEncryption, - Encoding, - KeySerializationEncryption, - NoEncryption, - ParameterFormat, - PrivateFormat, - PublicFormat, - _KeySerializationEncryption, -) -from cryptography.hazmat.primitives.serialization.base import ( - load_der_parameters, - load_der_private_key, - load_der_public_key, - load_pem_parameters, - load_pem_private_key, - load_pem_public_key, -) -from cryptography.hazmat.primitives.serialization.ssh import ( - SSHCertificate, - SSHCertificateBuilder, - SSHCertificateType, - SSHCertPrivateKeyTypes, - SSHCertPublicKeyTypes, - SSHPrivateKeyTypes, - SSHPublicKeyTypes, - load_ssh_private_key, - load_ssh_public_identity, - load_ssh_public_key, - ssh_key_fingerprint, -) - -__all__ = [ - "BestAvailableEncryption", - "Encoding", - "KeySerializationEncryption", - "NoEncryption", - "ParameterFormat", - "PrivateFormat", - "PublicFormat", - "SSHCertPrivateKeyTypes", - "SSHCertPublicKeyTypes", - "SSHCertificate", - "SSHCertificateBuilder", - "SSHCertificateType", - "SSHPrivateKeyTypes", - "SSHPublicKeyTypes", - "_KeySerializationEncryption", - "load_der_parameters", - "load_der_private_key", - "load_der_public_key", - "load_pem_parameters", - "load_pem_private_key", - "load_pem_public_key", - "load_ssh_private_key", - "load_ssh_public_identity", - "load_ssh_public_key", - "ssh_key_fingerprint", -] diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 8c4ddb9b236eee97ddb5e8552d737ee3ef28ed37..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/base.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/base.cpython-310.pyc deleted file mode 100644 index b792148a4052efb7b3cea1fabe6212d715d053ad..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/pkcs12.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/pkcs12.cpython-310.pyc deleted file mode 100644 index 5fd289d7afc3046e6f38a51d762ec4a39b2d74ca..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/pkcs12.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/pkcs7.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/pkcs7.cpython-310.pyc deleted file mode 100644 index aac958430b373f5ca7df1a7160447aae03bfe4b5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/pkcs7.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/ssh.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/ssh.cpython-310.pyc deleted file mode 100644 index 6ff1d39d2c39020676e2fb5874a81de07aa826a6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/__pycache__/ssh.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/base.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/base.py deleted file mode 100644 index e7c998b7f35b3e364909348f0516f51cf902ebfe..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/base.py +++ /dev/null @@ -1,14 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from cryptography.hazmat.bindings._rust import openssl as rust_openssl - -load_pem_private_key = rust_openssl.keys.load_pem_private_key -load_der_private_key = rust_openssl.keys.load_der_private_key - -load_pem_public_key = rust_openssl.keys.load_pem_public_key -load_der_public_key = rust_openssl.keys.load_der_public_key - -load_pem_parameters = rust_openssl.dh.from_pem_parameters -load_der_parameters = rust_openssl.dh.from_der_parameters diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/pkcs12.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/pkcs12.py deleted file mode 100644 index 58884ff61a79e50afe2df7ed134c33d447d02998..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/pkcs12.py +++ /dev/null @@ -1,176 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import typing -from collections.abc import Iterable - -from cryptography import x509 -from cryptography.hazmat.bindings._rust import pkcs12 as rust_pkcs12 -from cryptography.hazmat.primitives import serialization -from cryptography.hazmat.primitives._serialization import PBES as PBES -from cryptography.hazmat.primitives.asymmetric import ( - dsa, - ec, - ed448, - ed25519, - rsa, -) -from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes - -__all__ = [ - "PBES", - "PKCS12Certificate", - "PKCS12KeyAndCertificates", - "PKCS12PrivateKeyTypes", - "load_key_and_certificates", - "load_pkcs12", - "serialize_java_truststore", - "serialize_key_and_certificates", -] - -PKCS12PrivateKeyTypes = typing.Union[ - rsa.RSAPrivateKey, - dsa.DSAPrivateKey, - ec.EllipticCurvePrivateKey, - ed25519.Ed25519PrivateKey, - ed448.Ed448PrivateKey, -] - - -PKCS12Certificate = rust_pkcs12.PKCS12Certificate - - -class PKCS12KeyAndCertificates: - def __init__( - self, - key: PrivateKeyTypes | None, - cert: PKCS12Certificate | None, - additional_certs: list[PKCS12Certificate], - ): - if key is not None and not isinstance( - key, - ( - rsa.RSAPrivateKey, - dsa.DSAPrivateKey, - ec.EllipticCurvePrivateKey, - ed25519.Ed25519PrivateKey, - ed448.Ed448PrivateKey, - ), - ): - raise TypeError( - "Key must be RSA, DSA, EllipticCurve, ED25519, or ED448" - " private key, or None." - ) - if cert is not None and not isinstance(cert, PKCS12Certificate): - raise TypeError("cert must be a PKCS12Certificate object or None") - if not all( - isinstance(add_cert, PKCS12Certificate) - for add_cert in additional_certs - ): - raise TypeError( - "all values in additional_certs must be PKCS12Certificate" - " objects" - ) - self._key = key - self._cert = cert - self._additional_certs = additional_certs - - @property - def key(self) -> PrivateKeyTypes | None: - return self._key - - @property - def cert(self) -> PKCS12Certificate | None: - return self._cert - - @property - def additional_certs(self) -> list[PKCS12Certificate]: - return self._additional_certs - - def __eq__(self, other: object) -> bool: - if not isinstance(other, PKCS12KeyAndCertificates): - return NotImplemented - - return ( - self.key == other.key - and self.cert == other.cert - and self.additional_certs == other.additional_certs - ) - - def __hash__(self) -> int: - return hash((self.key, self.cert, tuple(self.additional_certs))) - - def __repr__(self) -> str: - fmt = ( - "" - ) - return fmt.format(self.key, self.cert, self.additional_certs) - - -load_key_and_certificates = rust_pkcs12.load_key_and_certificates -load_pkcs12 = rust_pkcs12.load_pkcs12 - - -_PKCS12CATypes = typing.Union[ - x509.Certificate, - PKCS12Certificate, -] - - -def serialize_java_truststore( - certs: Iterable[PKCS12Certificate], - encryption_algorithm: serialization.KeySerializationEncryption, -) -> bytes: - if not certs: - raise ValueError("You must supply at least one cert") - - if not isinstance( - encryption_algorithm, serialization.KeySerializationEncryption - ): - raise TypeError( - "Key encryption algorithm must be a " - "KeySerializationEncryption instance" - ) - - return rust_pkcs12.serialize_java_truststore(certs, encryption_algorithm) - - -def serialize_key_and_certificates( - name: bytes | None, - key: PKCS12PrivateKeyTypes | None, - cert: x509.Certificate | None, - cas: Iterable[_PKCS12CATypes] | None, - encryption_algorithm: serialization.KeySerializationEncryption, -) -> bytes: - if key is not None and not isinstance( - key, - ( - rsa.RSAPrivateKey, - dsa.DSAPrivateKey, - ec.EllipticCurvePrivateKey, - ed25519.Ed25519PrivateKey, - ed448.Ed448PrivateKey, - ), - ): - raise TypeError( - "Key must be RSA, DSA, EllipticCurve, ED25519, or ED448" - " private key, or None." - ) - - if not isinstance( - encryption_algorithm, serialization.KeySerializationEncryption - ): - raise TypeError( - "Key encryption algorithm must be a " - "KeySerializationEncryption instance" - ) - - if key is None and cert is None and not cas: - raise ValueError("You must supply at least one of key, cert, or cas") - - return rust_pkcs12.serialize_key_and_certificates( - name, key, cert, cas, encryption_algorithm - ) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/pkcs7.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/pkcs7.py deleted file mode 100644 index 456dc5b0831c0961d715c1e2697b48e59dc2ad24..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/pkcs7.py +++ /dev/null @@ -1,411 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import email.base64mime -import email.generator -import email.message -import email.policy -import io -import typing -from collections.abc import Iterable - -from cryptography import utils, x509 -from cryptography.exceptions import UnsupportedAlgorithm, _Reasons -from cryptography.hazmat.bindings._rust import pkcs7 as rust_pkcs7 -from cryptography.hazmat.primitives import hashes, serialization -from cryptography.hazmat.primitives.asymmetric import ec, padding, rsa -from cryptography.hazmat.primitives.ciphers import ( - algorithms, -) -from cryptography.utils import _check_byteslike - -load_pem_pkcs7_certificates = rust_pkcs7.load_pem_pkcs7_certificates - -load_der_pkcs7_certificates = rust_pkcs7.load_der_pkcs7_certificates - -serialize_certificates = rust_pkcs7.serialize_certificates - -PKCS7HashTypes = typing.Union[ - hashes.SHA224, - hashes.SHA256, - hashes.SHA384, - hashes.SHA512, -] - -PKCS7PrivateKeyTypes = typing.Union[ - rsa.RSAPrivateKey, ec.EllipticCurvePrivateKey -] - -ContentEncryptionAlgorithm = typing.Union[ - typing.Type[algorithms.AES128], typing.Type[algorithms.AES256] -] - - -class PKCS7Options(utils.Enum): - Text = "Add text/plain MIME type" - Binary = "Don't translate input data into canonical MIME format" - DetachedSignature = "Don't embed data in the PKCS7 structure" - NoCapabilities = "Don't embed SMIME capabilities" - NoAttributes = "Don't embed authenticatedAttributes" - NoCerts = "Don't embed signer certificate" - - -class PKCS7SignatureBuilder: - def __init__( - self, - data: utils.Buffer | None = None, - signers: list[ - tuple[ - x509.Certificate, - PKCS7PrivateKeyTypes, - PKCS7HashTypes, - padding.PSS | padding.PKCS1v15 | None, - ] - ] = [], - additional_certs: list[x509.Certificate] = [], - ): - self._data = data - self._signers = signers - self._additional_certs = additional_certs - - def set_data(self, data: utils.Buffer) -> PKCS7SignatureBuilder: - _check_byteslike("data", data) - if self._data is not None: - raise ValueError("data may only be set once") - - return PKCS7SignatureBuilder(data, self._signers) - - def add_signer( - self, - certificate: x509.Certificate, - private_key: PKCS7PrivateKeyTypes, - hash_algorithm: PKCS7HashTypes, - *, - rsa_padding: padding.PSS | padding.PKCS1v15 | None = None, - ) -> PKCS7SignatureBuilder: - if not isinstance( - hash_algorithm, - ( - hashes.SHA224, - hashes.SHA256, - hashes.SHA384, - hashes.SHA512, - ), - ): - raise TypeError( - "hash_algorithm must be one of hashes.SHA224, " - "SHA256, SHA384, or SHA512" - ) - if not isinstance(certificate, x509.Certificate): - raise TypeError("certificate must be a x509.Certificate") - - if not isinstance( - private_key, (rsa.RSAPrivateKey, ec.EllipticCurvePrivateKey) - ): - raise TypeError("Only RSA & EC keys are supported at this time.") - - if rsa_padding is not None: - if not isinstance(rsa_padding, (padding.PSS, padding.PKCS1v15)): - raise TypeError("Padding must be PSS or PKCS1v15") - if not isinstance(private_key, rsa.RSAPrivateKey): - raise TypeError("Padding is only supported for RSA keys") - - return PKCS7SignatureBuilder( - self._data, - [ - *self._signers, - (certificate, private_key, hash_algorithm, rsa_padding), - ], - ) - - def add_certificate( - self, certificate: x509.Certificate - ) -> PKCS7SignatureBuilder: - if not isinstance(certificate, x509.Certificate): - raise TypeError("certificate must be a x509.Certificate") - - return PKCS7SignatureBuilder( - self._data, self._signers, [*self._additional_certs, certificate] - ) - - def sign( - self, - encoding: serialization.Encoding, - options: Iterable[PKCS7Options], - backend: typing.Any = None, - ) -> bytes: - if len(self._signers) == 0: - raise ValueError("Must have at least one signer") - if self._data is None: - raise ValueError("You must add data to sign") - options = list(options) - if not all(isinstance(x, PKCS7Options) for x in options): - raise ValueError("options must be from the PKCS7Options enum") - if encoding not in ( - serialization.Encoding.PEM, - serialization.Encoding.DER, - serialization.Encoding.SMIME, - ): - raise ValueError( - "Must be PEM, DER, or SMIME from the Encoding enum" - ) - - # Text is a meaningless option unless it is accompanied by - # DetachedSignature - if ( - PKCS7Options.Text in options - and PKCS7Options.DetachedSignature not in options - ): - raise ValueError( - "When passing the Text option you must also pass " - "DetachedSignature" - ) - - if PKCS7Options.Text in options and encoding in ( - serialization.Encoding.DER, - serialization.Encoding.PEM, - ): - raise ValueError( - "The Text option is only available for SMIME serialization" - ) - - # No attributes implies no capabilities so we'll error if you try to - # pass both. - if ( - PKCS7Options.NoAttributes in options - and PKCS7Options.NoCapabilities in options - ): - raise ValueError( - "NoAttributes is a superset of NoCapabilities. Do not pass " - "both values." - ) - - return rust_pkcs7.sign_and_serialize(self, encoding, options) - - -class PKCS7EnvelopeBuilder: - def __init__( - self, - *, - _data: bytes | None = None, - _recipients: list[x509.Certificate] | None = None, - _content_encryption_algorithm: ContentEncryptionAlgorithm - | None = None, - ): - from cryptography.hazmat.backends.openssl.backend import ( - backend as ossl, - ) - - if not ossl.rsa_encryption_supported(padding=padding.PKCS1v15()): - raise UnsupportedAlgorithm( - "RSA with PKCS1 v1.5 padding is not supported by this version" - " of OpenSSL.", - _Reasons.UNSUPPORTED_PADDING, - ) - self._data = _data - self._recipients = _recipients if _recipients is not None else [] - self._content_encryption_algorithm = _content_encryption_algorithm - - def set_data(self, data: bytes) -> PKCS7EnvelopeBuilder: - _check_byteslike("data", data) - if self._data is not None: - raise ValueError("data may only be set once") - - return PKCS7EnvelopeBuilder( - _data=data, - _recipients=self._recipients, - _content_encryption_algorithm=self._content_encryption_algorithm, - ) - - def add_recipient( - self, - certificate: x509.Certificate, - ) -> PKCS7EnvelopeBuilder: - if not isinstance(certificate, x509.Certificate): - raise TypeError("certificate must be a x509.Certificate") - - if not isinstance(certificate.public_key(), rsa.RSAPublicKey): - raise TypeError("Only RSA keys are supported at this time.") - - return PKCS7EnvelopeBuilder( - _data=self._data, - _recipients=[ - *self._recipients, - certificate, - ], - _content_encryption_algorithm=self._content_encryption_algorithm, - ) - - def set_content_encryption_algorithm( - self, content_encryption_algorithm: ContentEncryptionAlgorithm - ) -> PKCS7EnvelopeBuilder: - if self._content_encryption_algorithm is not None: - raise ValueError("Content encryption algo may only be set once") - if content_encryption_algorithm not in { - algorithms.AES128, - algorithms.AES256, - }: - raise TypeError("Only AES128 and AES256 are supported") - - return PKCS7EnvelopeBuilder( - _data=self._data, - _recipients=self._recipients, - _content_encryption_algorithm=content_encryption_algorithm, - ) - - def encrypt( - self, - encoding: serialization.Encoding, - options: Iterable[PKCS7Options], - ) -> bytes: - if len(self._recipients) == 0: - raise ValueError("Must have at least one recipient") - if self._data is None: - raise ValueError("You must add data to encrypt") - - # The default content encryption algorithm is AES-128, which the S/MIME - # v3.2 RFC specifies as MUST support (https://datatracker.ietf.org/doc/html/rfc5751#section-2.7) - content_encryption_algorithm = ( - self._content_encryption_algorithm or algorithms.AES128 - ) - - options = list(options) - if not all(isinstance(x, PKCS7Options) for x in options): - raise ValueError("options must be from the PKCS7Options enum") - if encoding not in ( - serialization.Encoding.PEM, - serialization.Encoding.DER, - serialization.Encoding.SMIME, - ): - raise ValueError( - "Must be PEM, DER, or SMIME from the Encoding enum" - ) - - # Only allow options that make sense for encryption - if any( - opt not in [PKCS7Options.Text, PKCS7Options.Binary] - for opt in options - ): - raise ValueError( - "Only the following options are supported for encryption: " - "Text, Binary" - ) - elif PKCS7Options.Text in options and PKCS7Options.Binary in options: - # OpenSSL accepts both options at the same time, but ignores Text. - # We fail defensively to avoid unexpected outputs. - raise ValueError( - "Cannot use Binary and Text options at the same time" - ) - - return rust_pkcs7.encrypt_and_serialize( - self, content_encryption_algorithm, encoding, options - ) - - -pkcs7_decrypt_der = rust_pkcs7.decrypt_der -pkcs7_decrypt_pem = rust_pkcs7.decrypt_pem -pkcs7_decrypt_smime = rust_pkcs7.decrypt_smime - - -def _smime_signed_encode( - data: bytes, signature: bytes, micalg: str, text_mode: bool -) -> bytes: - # This function works pretty hard to replicate what OpenSSL does - # precisely. For good and for ill. - - m = email.message.Message() - m.add_header("MIME-Version", "1.0") - m.add_header( - "Content-Type", - "multipart/signed", - protocol="application/x-pkcs7-signature", - micalg=micalg, - ) - - m.preamble = "This is an S/MIME signed message\n" - - msg_part = OpenSSLMimePart() - msg_part.set_payload(data) - if text_mode: - msg_part.add_header("Content-Type", "text/plain") - m.attach(msg_part) - - sig_part = email.message.MIMEPart() - sig_part.add_header( - "Content-Type", "application/x-pkcs7-signature", name="smime.p7s" - ) - sig_part.add_header("Content-Transfer-Encoding", "base64") - sig_part.add_header( - "Content-Disposition", "attachment", filename="smime.p7s" - ) - sig_part.set_payload( - email.base64mime.body_encode(signature, maxlinelen=65) - ) - del sig_part["MIME-Version"] - m.attach(sig_part) - - fp = io.BytesIO() - g = email.generator.BytesGenerator( - fp, - maxheaderlen=0, - mangle_from_=False, - policy=m.policy.clone(linesep="\r\n"), - ) - g.flatten(m) - return fp.getvalue() - - -def _smime_enveloped_encode(data: bytes) -> bytes: - m = email.message.Message() - m.add_header("MIME-Version", "1.0") - m.add_header("Content-Disposition", "attachment", filename="smime.p7m") - m.add_header( - "Content-Type", - "application/pkcs7-mime", - smime_type="enveloped-data", - name="smime.p7m", - ) - m.add_header("Content-Transfer-Encoding", "base64") - - m.set_payload(email.base64mime.body_encode(data, maxlinelen=65)) - - return m.as_bytes(policy=m.policy.clone(linesep="\n", max_line_length=0)) - - -def _smime_enveloped_decode(data: bytes) -> bytes: - m = email.message_from_bytes(data) - if m.get_content_type() not in { - "application/x-pkcs7-mime", - "application/pkcs7-mime", - }: - raise ValueError("Not an S/MIME enveloped message") - return bytes(m.get_payload(decode=True)) - - -def _smime_remove_text_headers(data: bytes) -> bytes: - m = email.message_from_bytes(data) - # Using get() instead of get_content_type() since it has None as default, - # where the latter has "text/plain". Both methods are case-insensitive. - content_type = m.get("content-type") - if content_type is None: - raise ValueError( - "Decrypted MIME data has no 'Content-Type' header. " - "Please remove the 'Text' option to parse it manually." - ) - if "text/plain" not in content_type: - raise ValueError( - f"Decrypted MIME data content type is '{content_type}', not " - "'text/plain'. Remove the 'Text' option to parse it manually." - ) - return bytes(m.get_payload(decode=True)) - - -class OpenSSLMimePart(email.message.MIMEPart): - # A MIMEPart subclass that replicates OpenSSL's behavior of not including - # a newline if there are no headers. - def _write_headers(self, generator) -> None: - if list(self.raw_items()): - generator._write_headers(self) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/ssh.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/ssh.py deleted file mode 100644 index c08b98121f2c759d8b0665c47487c7506db53994..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/serialization/ssh.py +++ /dev/null @@ -1,1619 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import binascii -import enum -import os -import re -import typing -import warnings -from base64 import encodebytes as _base64_encode -from dataclasses import dataclass - -from cryptography import utils -from cryptography.exceptions import UnsupportedAlgorithm -from cryptography.hazmat.primitives import hashes -from cryptography.hazmat.primitives.asymmetric import ( - dsa, - ec, - ed25519, - padding, - rsa, -) -from cryptography.hazmat.primitives.asymmetric import utils as asym_utils -from cryptography.hazmat.primitives.ciphers import ( - AEADDecryptionContext, - Cipher, - algorithms, - modes, -) -from cryptography.hazmat.primitives.serialization import ( - Encoding, - KeySerializationEncryption, - NoEncryption, - PrivateFormat, - PublicFormat, - _KeySerializationEncryption, -) - -try: - from bcrypt import kdf as _bcrypt_kdf - - _bcrypt_supported = True -except ImportError: - _bcrypt_supported = False - - def _bcrypt_kdf( - password: bytes, - salt: bytes, - desired_key_bytes: int, - rounds: int, - ignore_few_rounds: bool = False, - ) -> bytes: - raise UnsupportedAlgorithm("Need bcrypt module") - - -_SSH_ED25519 = b"ssh-ed25519" -_SSH_RSA = b"ssh-rsa" -_SSH_DSA = b"ssh-dss" -_ECDSA_NISTP256 = b"ecdsa-sha2-nistp256" -_ECDSA_NISTP384 = b"ecdsa-sha2-nistp384" -_ECDSA_NISTP521 = b"ecdsa-sha2-nistp521" -_CERT_SUFFIX = b"-cert-v01@openssh.com" - -# U2F application string suffixed pubkey -_SK_SSH_ED25519 = b"sk-ssh-ed25519@openssh.com" -_SK_SSH_ECDSA_NISTP256 = b"sk-ecdsa-sha2-nistp256@openssh.com" - -# These are not key types, only algorithms, so they cannot appear -# as a public key type -_SSH_RSA_SHA256 = b"rsa-sha2-256" -_SSH_RSA_SHA512 = b"rsa-sha2-512" - -_SSH_PUBKEY_RC = re.compile(rb"\A(\S+)[ \t]+(\S+)") -_SK_MAGIC = b"openssh-key-v1\0" -_SK_START = b"-----BEGIN OPENSSH PRIVATE KEY-----" -_SK_END = b"-----END OPENSSH PRIVATE KEY-----" -_BCRYPT = b"bcrypt" -_NONE = b"none" -_DEFAULT_CIPHER = b"aes256-ctr" -_DEFAULT_ROUNDS = 16 - -# re is only way to work on bytes-like data -_PEM_RC = re.compile(_SK_START + b"(.*?)" + _SK_END, re.DOTALL) - -# padding for max blocksize -_PADDING = memoryview(bytearray(range(1, 1 + 16))) - - -@dataclass -class _SSHCipher: - alg: type[algorithms.AES] - key_len: int - mode: type[modes.CTR] | type[modes.CBC] | type[modes.GCM] - block_len: int - iv_len: int - tag_len: int | None - is_aead: bool - - -# ciphers that are actually used in key wrapping -_SSH_CIPHERS: dict[bytes, _SSHCipher] = { - b"aes256-ctr": _SSHCipher( - alg=algorithms.AES, - key_len=32, - mode=modes.CTR, - block_len=16, - iv_len=16, - tag_len=None, - is_aead=False, - ), - b"aes256-cbc": _SSHCipher( - alg=algorithms.AES, - key_len=32, - mode=modes.CBC, - block_len=16, - iv_len=16, - tag_len=None, - is_aead=False, - ), - b"aes256-gcm@openssh.com": _SSHCipher( - alg=algorithms.AES, - key_len=32, - mode=modes.GCM, - block_len=16, - iv_len=12, - tag_len=16, - is_aead=True, - ), -} - -# map local curve name to key type -_ECDSA_KEY_TYPE = { - "secp256r1": _ECDSA_NISTP256, - "secp384r1": _ECDSA_NISTP384, - "secp521r1": _ECDSA_NISTP521, -} - - -def _get_ssh_key_type(key: SSHPrivateKeyTypes | SSHPublicKeyTypes) -> bytes: - if isinstance(key, ec.EllipticCurvePrivateKey): - key_type = _ecdsa_key_type(key.public_key()) - elif isinstance(key, ec.EllipticCurvePublicKey): - key_type = _ecdsa_key_type(key) - elif isinstance(key, (rsa.RSAPrivateKey, rsa.RSAPublicKey)): - key_type = _SSH_RSA - elif isinstance(key, (dsa.DSAPrivateKey, dsa.DSAPublicKey)): - key_type = _SSH_DSA - elif isinstance( - key, (ed25519.Ed25519PrivateKey, ed25519.Ed25519PublicKey) - ): - key_type = _SSH_ED25519 - else: - raise ValueError("Unsupported key type") - - return key_type - - -def _ecdsa_key_type(public_key: ec.EllipticCurvePublicKey) -> bytes: - """Return SSH key_type and curve_name for private key.""" - curve = public_key.curve - if curve.name not in _ECDSA_KEY_TYPE: - raise ValueError( - f"Unsupported curve for ssh private key: {curve.name!r}" - ) - return _ECDSA_KEY_TYPE[curve.name] - - -def _ssh_pem_encode( - data: utils.Buffer, - prefix: bytes = _SK_START + b"\n", - suffix: bytes = _SK_END + b"\n", -) -> bytes: - return b"".join([prefix, _base64_encode(data), suffix]) - - -def _check_block_size(data: utils.Buffer, block_len: int) -> None: - """Require data to be full blocks""" - if not data or len(data) % block_len != 0: - raise ValueError("Corrupt data: missing padding") - - -def _check_empty(data: utils.Buffer) -> None: - """All data should have been parsed.""" - if data: - raise ValueError("Corrupt data: unparsed data") - - -def _init_cipher( - ciphername: bytes, - password: bytes | None, - salt: bytes, - rounds: int, -) -> Cipher[modes.CBC | modes.CTR | modes.GCM]: - """Generate key + iv and return cipher.""" - if not password: - raise TypeError( - "Key is password-protected, but password was not provided." - ) - - ciph = _SSH_CIPHERS[ciphername] - seed = _bcrypt_kdf( - password, salt, ciph.key_len + ciph.iv_len, rounds, True - ) - return Cipher( - ciph.alg(seed[: ciph.key_len]), - ciph.mode(seed[ciph.key_len :]), - ) - - -def _get_u32(data: memoryview) -> tuple[int, memoryview]: - """Uint32""" - if len(data) < 4: - raise ValueError("Invalid data") - return int.from_bytes(data[:4], byteorder="big"), data[4:] - - -def _get_u64(data: memoryview) -> tuple[int, memoryview]: - """Uint64""" - if len(data) < 8: - raise ValueError("Invalid data") - return int.from_bytes(data[:8], byteorder="big"), data[8:] - - -def _get_sshstr(data: memoryview) -> tuple[memoryview, memoryview]: - """Bytes with u32 length prefix""" - n, data = _get_u32(data) - if n > len(data): - raise ValueError("Invalid data") - return data[:n], data[n:] - - -def _get_mpint(data: memoryview) -> tuple[int, memoryview]: - """Big integer.""" - val, data = _get_sshstr(data) - if val and val[0] > 0x7F: - raise ValueError("Invalid data") - return int.from_bytes(val, "big"), data - - -def _to_mpint(val: int) -> bytes: - """Storage format for signed bigint.""" - if val < 0: - raise ValueError("negative mpint not allowed") - if not val: - return b"" - nbytes = (val.bit_length() + 8) // 8 - return utils.int_to_bytes(val, nbytes) - - -class _FragList: - """Build recursive structure without data copy.""" - - flist: list[utils.Buffer] - - def __init__(self, init: list[utils.Buffer] | None = None) -> None: - self.flist = [] - if init: - self.flist.extend(init) - - def put_raw(self, val: utils.Buffer) -> None: - """Add plain bytes""" - self.flist.append(val) - - def put_u32(self, val: int) -> None: - """Big-endian uint32""" - self.flist.append(val.to_bytes(length=4, byteorder="big")) - - def put_u64(self, val: int) -> None: - """Big-endian uint64""" - self.flist.append(val.to_bytes(length=8, byteorder="big")) - - def put_sshstr(self, val: bytes | _FragList) -> None: - """Bytes prefixed with u32 length""" - if isinstance(val, (bytes, memoryview, bytearray)): - self.put_u32(len(val)) - self.flist.append(val) - else: - self.put_u32(val.size()) - self.flist.extend(val.flist) - - def put_mpint(self, val: int) -> None: - """Big-endian bigint prefixed with u32 length""" - self.put_sshstr(_to_mpint(val)) - - def size(self) -> int: - """Current number of bytes""" - return sum(map(len, self.flist)) - - def render(self, dstbuf: memoryview, pos: int = 0) -> int: - """Write into bytearray""" - for frag in self.flist: - flen = len(frag) - start, pos = pos, pos + flen - dstbuf[start:pos] = frag - return pos - - def tobytes(self) -> bytes: - """Return as bytes""" - buf = memoryview(bytearray(self.size())) - self.render(buf) - return buf.tobytes() - - -class _SSHFormatRSA: - """Format for RSA keys. - - Public: - mpint e, n - Private: - mpint n, e, d, iqmp, p, q - """ - - def get_public( - self, data: memoryview - ) -> tuple[tuple[int, int], memoryview]: - """RSA public fields""" - e, data = _get_mpint(data) - n, data = _get_mpint(data) - return (e, n), data - - def load_public( - self, data: memoryview - ) -> tuple[rsa.RSAPublicKey, memoryview]: - """Make RSA public key from data.""" - (e, n), data = self.get_public(data) - public_numbers = rsa.RSAPublicNumbers(e, n) - public_key = public_numbers.public_key() - return public_key, data - - def load_private( - self, data: memoryview, pubfields, unsafe_skip_rsa_key_validation: bool - ) -> tuple[rsa.RSAPrivateKey, memoryview]: - """Make RSA private key from data.""" - n, data = _get_mpint(data) - e, data = _get_mpint(data) - d, data = _get_mpint(data) - iqmp, data = _get_mpint(data) - p, data = _get_mpint(data) - q, data = _get_mpint(data) - - if (e, n) != pubfields: - raise ValueError("Corrupt data: rsa field mismatch") - dmp1 = rsa.rsa_crt_dmp1(d, p) - dmq1 = rsa.rsa_crt_dmq1(d, q) - public_numbers = rsa.RSAPublicNumbers(e, n) - private_numbers = rsa.RSAPrivateNumbers( - p, q, d, dmp1, dmq1, iqmp, public_numbers - ) - private_key = private_numbers.private_key( - unsafe_skip_rsa_key_validation=unsafe_skip_rsa_key_validation - ) - return private_key, data - - def encode_public( - self, public_key: rsa.RSAPublicKey, f_pub: _FragList - ) -> None: - """Write RSA public key""" - pubn = public_key.public_numbers() - f_pub.put_mpint(pubn.e) - f_pub.put_mpint(pubn.n) - - def encode_private( - self, private_key: rsa.RSAPrivateKey, f_priv: _FragList - ) -> None: - """Write RSA private key""" - private_numbers = private_key.private_numbers() - public_numbers = private_numbers.public_numbers - - f_priv.put_mpint(public_numbers.n) - f_priv.put_mpint(public_numbers.e) - - f_priv.put_mpint(private_numbers.d) - f_priv.put_mpint(private_numbers.iqmp) - f_priv.put_mpint(private_numbers.p) - f_priv.put_mpint(private_numbers.q) - - -class _SSHFormatDSA: - """Format for DSA keys. - - Public: - mpint p, q, g, y - Private: - mpint p, q, g, y, x - """ - - def get_public(self, data: memoryview) -> tuple[tuple, memoryview]: - """DSA public fields""" - p, data = _get_mpint(data) - q, data = _get_mpint(data) - g, data = _get_mpint(data) - y, data = _get_mpint(data) - return (p, q, g, y), data - - def load_public( - self, data: memoryview - ) -> tuple[dsa.DSAPublicKey, memoryview]: - """Make DSA public key from data.""" - (p, q, g, y), data = self.get_public(data) - parameter_numbers = dsa.DSAParameterNumbers(p, q, g) - public_numbers = dsa.DSAPublicNumbers(y, parameter_numbers) - self._validate(public_numbers) - public_key = public_numbers.public_key() - return public_key, data - - def load_private( - self, data: memoryview, pubfields, unsafe_skip_rsa_key_validation: bool - ) -> tuple[dsa.DSAPrivateKey, memoryview]: - """Make DSA private key from data.""" - (p, q, g, y), data = self.get_public(data) - x, data = _get_mpint(data) - - if (p, q, g, y) != pubfields: - raise ValueError("Corrupt data: dsa field mismatch") - parameter_numbers = dsa.DSAParameterNumbers(p, q, g) - public_numbers = dsa.DSAPublicNumbers(y, parameter_numbers) - self._validate(public_numbers) - private_numbers = dsa.DSAPrivateNumbers(x, public_numbers) - private_key = private_numbers.private_key() - return private_key, data - - def encode_public( - self, public_key: dsa.DSAPublicKey, f_pub: _FragList - ) -> None: - """Write DSA public key""" - public_numbers = public_key.public_numbers() - parameter_numbers = public_numbers.parameter_numbers - self._validate(public_numbers) - - f_pub.put_mpint(parameter_numbers.p) - f_pub.put_mpint(parameter_numbers.q) - f_pub.put_mpint(parameter_numbers.g) - f_pub.put_mpint(public_numbers.y) - - def encode_private( - self, private_key: dsa.DSAPrivateKey, f_priv: _FragList - ) -> None: - """Write DSA private key""" - self.encode_public(private_key.public_key(), f_priv) - f_priv.put_mpint(private_key.private_numbers().x) - - def _validate(self, public_numbers: dsa.DSAPublicNumbers) -> None: - parameter_numbers = public_numbers.parameter_numbers - if parameter_numbers.p.bit_length() != 1024: - raise ValueError("SSH supports only 1024 bit DSA keys") - - -class _SSHFormatECDSA: - """Format for ECDSA keys. - - Public: - str curve - bytes point - Private: - str curve - bytes point - mpint secret - """ - - def __init__(self, ssh_curve_name: bytes, curve: ec.EllipticCurve): - self.ssh_curve_name = ssh_curve_name - self.curve = curve - - def get_public( - self, data: memoryview - ) -> tuple[tuple[memoryview, memoryview], memoryview]: - """ECDSA public fields""" - curve, data = _get_sshstr(data) - point, data = _get_sshstr(data) - if curve != self.ssh_curve_name: - raise ValueError("Curve name mismatch") - if point[0] != 4: - raise NotImplementedError("Need uncompressed point") - return (curve, point), data - - def load_public( - self, data: memoryview - ) -> tuple[ec.EllipticCurvePublicKey, memoryview]: - """Make ECDSA public key from data.""" - (_, point), data = self.get_public(data) - public_key = ec.EllipticCurvePublicKey.from_encoded_point( - self.curve, point.tobytes() - ) - return public_key, data - - def load_private( - self, data: memoryview, pubfields, unsafe_skip_rsa_key_validation: bool - ) -> tuple[ec.EllipticCurvePrivateKey, memoryview]: - """Make ECDSA private key from data.""" - (curve_name, point), data = self.get_public(data) - secret, data = _get_mpint(data) - - if (curve_name, point) != pubfields: - raise ValueError("Corrupt data: ecdsa field mismatch") - private_key = ec.derive_private_key(secret, self.curve) - return private_key, data - - def encode_public( - self, public_key: ec.EllipticCurvePublicKey, f_pub: _FragList - ) -> None: - """Write ECDSA public key""" - point = public_key.public_bytes( - Encoding.X962, PublicFormat.UncompressedPoint - ) - f_pub.put_sshstr(self.ssh_curve_name) - f_pub.put_sshstr(point) - - def encode_private( - self, private_key: ec.EllipticCurvePrivateKey, f_priv: _FragList - ) -> None: - """Write ECDSA private key""" - public_key = private_key.public_key() - private_numbers = private_key.private_numbers() - - self.encode_public(public_key, f_priv) - f_priv.put_mpint(private_numbers.private_value) - - -class _SSHFormatEd25519: - """Format for Ed25519 keys. - - Public: - bytes point - Private: - bytes point - bytes secret_and_point - """ - - def get_public( - self, data: memoryview - ) -> tuple[tuple[memoryview], memoryview]: - """Ed25519 public fields""" - point, data = _get_sshstr(data) - return (point,), data - - def load_public( - self, data: memoryview - ) -> tuple[ed25519.Ed25519PublicKey, memoryview]: - """Make Ed25519 public key from data.""" - (point,), data = self.get_public(data) - public_key = ed25519.Ed25519PublicKey.from_public_bytes( - point.tobytes() - ) - return public_key, data - - def load_private( - self, data: memoryview, pubfields, unsafe_skip_rsa_key_validation: bool - ) -> tuple[ed25519.Ed25519PrivateKey, memoryview]: - """Make Ed25519 private key from data.""" - (point,), data = self.get_public(data) - keypair, data = _get_sshstr(data) - - secret = keypair[:32] - point2 = keypair[32:] - if point != point2 or (point,) != pubfields: - raise ValueError("Corrupt data: ed25519 field mismatch") - private_key = ed25519.Ed25519PrivateKey.from_private_bytes(secret) - return private_key, data - - def encode_public( - self, public_key: ed25519.Ed25519PublicKey, f_pub: _FragList - ) -> None: - """Write Ed25519 public key""" - raw_public_key = public_key.public_bytes( - Encoding.Raw, PublicFormat.Raw - ) - f_pub.put_sshstr(raw_public_key) - - def encode_private( - self, private_key: ed25519.Ed25519PrivateKey, f_priv: _FragList - ) -> None: - """Write Ed25519 private key""" - public_key = private_key.public_key() - raw_private_key = private_key.private_bytes( - Encoding.Raw, PrivateFormat.Raw, NoEncryption() - ) - raw_public_key = public_key.public_bytes( - Encoding.Raw, PublicFormat.Raw - ) - f_keypair = _FragList([raw_private_key, raw_public_key]) - - self.encode_public(public_key, f_priv) - f_priv.put_sshstr(f_keypair) - - -def load_application(data) -> tuple[memoryview, memoryview]: - """ - U2F application strings - """ - application, data = _get_sshstr(data) - if not application.tobytes().startswith(b"ssh:"): - raise ValueError( - "U2F application string does not start with b'ssh:' " - f"({application})" - ) - return application, data - - -class _SSHFormatSKEd25519: - """ - The format of a sk-ssh-ed25519@openssh.com public key is: - - string "sk-ssh-ed25519@openssh.com" - string public key - string application (user-specified, but typically "ssh:") - """ - - def load_public( - self, data: memoryview - ) -> tuple[ed25519.Ed25519PublicKey, memoryview]: - """Make Ed25519 public key from data.""" - public_key, data = _lookup_kformat(_SSH_ED25519).load_public(data) - _, data = load_application(data) - return public_key, data - - def get_public(self, data: memoryview) -> typing.NoReturn: - # Confusingly `get_public` is an entry point used by private key - # loading. - raise UnsupportedAlgorithm( - "sk-ssh-ed25519 private keys cannot be loaded" - ) - - -class _SSHFormatSKECDSA: - """ - The format of a sk-ecdsa-sha2-nistp256@openssh.com public key is: - - string "sk-ecdsa-sha2-nistp256@openssh.com" - string curve name - ec_point Q - string application (user-specified, but typically "ssh:") - """ - - def load_public( - self, data: memoryview - ) -> tuple[ec.EllipticCurvePublicKey, memoryview]: - """Make ECDSA public key from data.""" - public_key, data = _lookup_kformat(_ECDSA_NISTP256).load_public(data) - _, data = load_application(data) - return public_key, data - - def get_public(self, data: memoryview) -> typing.NoReturn: - # Confusingly `get_public` is an entry point used by private key - # loading. - raise UnsupportedAlgorithm( - "sk-ecdsa-sha2-nistp256 private keys cannot be loaded" - ) - - -_KEY_FORMATS = { - _SSH_RSA: _SSHFormatRSA(), - _SSH_DSA: _SSHFormatDSA(), - _SSH_ED25519: _SSHFormatEd25519(), - _ECDSA_NISTP256: _SSHFormatECDSA(b"nistp256", ec.SECP256R1()), - _ECDSA_NISTP384: _SSHFormatECDSA(b"nistp384", ec.SECP384R1()), - _ECDSA_NISTP521: _SSHFormatECDSA(b"nistp521", ec.SECP521R1()), - _SK_SSH_ED25519: _SSHFormatSKEd25519(), - _SK_SSH_ECDSA_NISTP256: _SSHFormatSKECDSA(), -} - - -def _lookup_kformat(key_type: utils.Buffer): - """Return valid format or throw error""" - if not isinstance(key_type, bytes): - key_type = memoryview(key_type).tobytes() - if key_type in _KEY_FORMATS: - return _KEY_FORMATS[key_type] - raise UnsupportedAlgorithm(f"Unsupported key type: {key_type!r}") - - -SSHPrivateKeyTypes = typing.Union[ - ec.EllipticCurvePrivateKey, - rsa.RSAPrivateKey, - dsa.DSAPrivateKey, - ed25519.Ed25519PrivateKey, -] - - -def load_ssh_private_key( - data: utils.Buffer, - password: bytes | None, - backend: typing.Any = None, - *, - unsafe_skip_rsa_key_validation: bool = False, -) -> SSHPrivateKeyTypes: - """Load private key from OpenSSH custom encoding.""" - utils._check_byteslike("data", data) - if password is not None: - utils._check_bytes("password", password) - - m = _PEM_RC.search(data) - if not m: - raise ValueError("Not OpenSSH private key format") - p1 = m.start(1) - p2 = m.end(1) - data = binascii.a2b_base64(memoryview(data)[p1:p2]) - if not data.startswith(_SK_MAGIC): - raise ValueError("Not OpenSSH private key format") - data = memoryview(data)[len(_SK_MAGIC) :] - - # parse header - ciphername, data = _get_sshstr(data) - kdfname, data = _get_sshstr(data) - kdfoptions, data = _get_sshstr(data) - nkeys, data = _get_u32(data) - if nkeys != 1: - raise ValueError("Only one key supported") - - # load public key data - pubdata, data = _get_sshstr(data) - pub_key_type, pubdata = _get_sshstr(pubdata) - kformat = _lookup_kformat(pub_key_type) - pubfields, pubdata = kformat.get_public(pubdata) - _check_empty(pubdata) - - if ciphername != _NONE or kdfname != _NONE: - ciphername_bytes = ciphername.tobytes() - if ciphername_bytes not in _SSH_CIPHERS: - raise UnsupportedAlgorithm( - f"Unsupported cipher: {ciphername_bytes!r}" - ) - if kdfname != _BCRYPT: - raise UnsupportedAlgorithm(f"Unsupported KDF: {kdfname!r}") - blklen = _SSH_CIPHERS[ciphername_bytes].block_len - tag_len = _SSH_CIPHERS[ciphername_bytes].tag_len - # load secret data - edata, data = _get_sshstr(data) - # see https://bugzilla.mindrot.org/show_bug.cgi?id=3553 for - # information about how OpenSSH handles AEAD tags - if _SSH_CIPHERS[ciphername_bytes].is_aead: - tag = bytes(data) - if len(tag) != tag_len: - raise ValueError("Corrupt data: invalid tag length for cipher") - else: - _check_empty(data) - _check_block_size(edata, blklen) - salt, kbuf = _get_sshstr(kdfoptions) - rounds, kbuf = _get_u32(kbuf) - _check_empty(kbuf) - ciph = _init_cipher(ciphername_bytes, password, salt.tobytes(), rounds) - dec = ciph.decryptor() - edata = memoryview(dec.update(edata)) - if _SSH_CIPHERS[ciphername_bytes].is_aead: - assert isinstance(dec, AEADDecryptionContext) - _check_empty(dec.finalize_with_tag(tag)) - else: - # _check_block_size requires data to be a full block so there - # should be no output from finalize - _check_empty(dec.finalize()) - else: - if password: - raise TypeError( - "Password was given but private key is not encrypted." - ) - # load secret data - edata, data = _get_sshstr(data) - _check_empty(data) - blklen = 8 - _check_block_size(edata, blklen) - ck1, edata = _get_u32(edata) - ck2, edata = _get_u32(edata) - if ck1 != ck2: - raise ValueError("Corrupt data: broken checksum") - - # load per-key struct - key_type, edata = _get_sshstr(edata) - if key_type != pub_key_type: - raise ValueError("Corrupt data: key type mismatch") - private_key, edata = kformat.load_private( - edata, - pubfields, - unsafe_skip_rsa_key_validation=unsafe_skip_rsa_key_validation, - ) - # We don't use the comment - _, edata = _get_sshstr(edata) - - # yes, SSH does padding check *after* all other parsing is done. - # need to follow as it writes zero-byte padding too. - if edata != _PADDING[: len(edata)]: - raise ValueError("Corrupt data: invalid padding") - - if isinstance(private_key, dsa.DSAPrivateKey): - warnings.warn( - "SSH DSA keys are deprecated and will be removed in a future " - "release.", - utils.DeprecatedIn40, - stacklevel=2, - ) - - return private_key - - -def _serialize_ssh_private_key( - private_key: SSHPrivateKeyTypes, - password: bytes, - encryption_algorithm: KeySerializationEncryption, -) -> bytes: - """Serialize private key with OpenSSH custom encoding.""" - utils._check_bytes("password", password) - if isinstance(private_key, dsa.DSAPrivateKey): - warnings.warn( - "SSH DSA key support is deprecated and will be " - "removed in a future release", - utils.DeprecatedIn40, - stacklevel=4, - ) - - key_type = _get_ssh_key_type(private_key) - kformat = _lookup_kformat(key_type) - - # setup parameters - f_kdfoptions = _FragList() - if password: - ciphername = _DEFAULT_CIPHER - blklen = _SSH_CIPHERS[ciphername].block_len - kdfname = _BCRYPT - rounds = _DEFAULT_ROUNDS - if ( - isinstance(encryption_algorithm, _KeySerializationEncryption) - and encryption_algorithm._kdf_rounds is not None - ): - rounds = encryption_algorithm._kdf_rounds - salt = os.urandom(16) - f_kdfoptions.put_sshstr(salt) - f_kdfoptions.put_u32(rounds) - ciph = _init_cipher(ciphername, password, salt, rounds) - else: - ciphername = kdfname = _NONE - blklen = 8 - ciph = None - nkeys = 1 - checkval = os.urandom(4) - comment = b"" - - # encode public and private parts together - f_public_key = _FragList() - f_public_key.put_sshstr(key_type) - kformat.encode_public(private_key.public_key(), f_public_key) - - f_secrets = _FragList([checkval, checkval]) - f_secrets.put_sshstr(key_type) - kformat.encode_private(private_key, f_secrets) - f_secrets.put_sshstr(comment) - f_secrets.put_raw(_PADDING[: blklen - (f_secrets.size() % blklen)]) - - # top-level structure - f_main = _FragList() - f_main.put_raw(_SK_MAGIC) - f_main.put_sshstr(ciphername) - f_main.put_sshstr(kdfname) - f_main.put_sshstr(f_kdfoptions) - f_main.put_u32(nkeys) - f_main.put_sshstr(f_public_key) - f_main.put_sshstr(f_secrets) - - # copy result info bytearray - slen = f_secrets.size() - mlen = f_main.size() - buf = memoryview(bytearray(mlen + blklen)) - f_main.render(buf) - ofs = mlen - slen - - # encrypt in-place - if ciph is not None: - ciph.encryptor().update_into(buf[ofs:mlen], buf[ofs:]) - - return _ssh_pem_encode(buf[:mlen]) - - -SSHPublicKeyTypes = typing.Union[ - ec.EllipticCurvePublicKey, - rsa.RSAPublicKey, - dsa.DSAPublicKey, - ed25519.Ed25519PublicKey, -] - -SSHCertPublicKeyTypes = typing.Union[ - ec.EllipticCurvePublicKey, - rsa.RSAPublicKey, - ed25519.Ed25519PublicKey, -] - - -class SSHCertificateType(enum.Enum): - USER = 1 - HOST = 2 - - -class SSHCertificate: - def __init__( - self, - _nonce: memoryview, - _public_key: SSHPublicKeyTypes, - _serial: int, - _cctype: int, - _key_id: memoryview, - _valid_principals: list[bytes], - _valid_after: int, - _valid_before: int, - _critical_options: dict[bytes, bytes], - _extensions: dict[bytes, bytes], - _sig_type: memoryview, - _sig_key: memoryview, - _inner_sig_type: memoryview, - _signature: memoryview, - _tbs_cert_body: memoryview, - _cert_key_type: bytes, - _cert_body: memoryview, - ): - self._nonce = _nonce - self._public_key = _public_key - self._serial = _serial - try: - self._type = SSHCertificateType(_cctype) - except ValueError: - raise ValueError("Invalid certificate type") - self._key_id = _key_id - self._valid_principals = _valid_principals - self._valid_after = _valid_after - self._valid_before = _valid_before - self._critical_options = _critical_options - self._extensions = _extensions - self._sig_type = _sig_type - self._sig_key = _sig_key - self._inner_sig_type = _inner_sig_type - self._signature = _signature - self._cert_key_type = _cert_key_type - self._cert_body = _cert_body - self._tbs_cert_body = _tbs_cert_body - - @property - def nonce(self) -> bytes: - return bytes(self._nonce) - - def public_key(self) -> SSHCertPublicKeyTypes: - # make mypy happy until we remove DSA support entirely and - # the underlying union won't have a disallowed type - return typing.cast(SSHCertPublicKeyTypes, self._public_key) - - @property - def serial(self) -> int: - return self._serial - - @property - def type(self) -> SSHCertificateType: - return self._type - - @property - def key_id(self) -> bytes: - return bytes(self._key_id) - - @property - def valid_principals(self) -> list[bytes]: - return self._valid_principals - - @property - def valid_before(self) -> int: - return self._valid_before - - @property - def valid_after(self) -> int: - return self._valid_after - - @property - def critical_options(self) -> dict[bytes, bytes]: - return self._critical_options - - @property - def extensions(self) -> dict[bytes, bytes]: - return self._extensions - - def signature_key(self) -> SSHCertPublicKeyTypes: - sigformat = _lookup_kformat(self._sig_type) - signature_key, sigkey_rest = sigformat.load_public(self._sig_key) - _check_empty(sigkey_rest) - return signature_key - - def public_bytes(self) -> bytes: - return ( - bytes(self._cert_key_type) - + b" " - + binascii.b2a_base64(bytes(self._cert_body), newline=False) - ) - - def verify_cert_signature(self) -> None: - signature_key = self.signature_key() - if isinstance(signature_key, ed25519.Ed25519PublicKey): - signature_key.verify( - bytes(self._signature), bytes(self._tbs_cert_body) - ) - elif isinstance(signature_key, ec.EllipticCurvePublicKey): - # The signature is encoded as a pair of big-endian integers - r, data = _get_mpint(self._signature) - s, data = _get_mpint(data) - _check_empty(data) - computed_sig = asym_utils.encode_dss_signature(r, s) - hash_alg = _get_ec_hash_alg(signature_key.curve) - signature_key.verify( - computed_sig, bytes(self._tbs_cert_body), ec.ECDSA(hash_alg) - ) - else: - assert isinstance(signature_key, rsa.RSAPublicKey) - if self._inner_sig_type == _SSH_RSA: - hash_alg = hashes.SHA1() - elif self._inner_sig_type == _SSH_RSA_SHA256: - hash_alg = hashes.SHA256() - else: - assert self._inner_sig_type == _SSH_RSA_SHA512 - hash_alg = hashes.SHA512() - signature_key.verify( - bytes(self._signature), - bytes(self._tbs_cert_body), - padding.PKCS1v15(), - hash_alg, - ) - - -def _get_ec_hash_alg(curve: ec.EllipticCurve) -> hashes.HashAlgorithm: - if isinstance(curve, ec.SECP256R1): - return hashes.SHA256() - elif isinstance(curve, ec.SECP384R1): - return hashes.SHA384() - else: - assert isinstance(curve, ec.SECP521R1) - return hashes.SHA512() - - -def _load_ssh_public_identity( - data: utils.Buffer, - _legacy_dsa_allowed=False, -) -> SSHCertificate | SSHPublicKeyTypes: - utils._check_byteslike("data", data) - - m = _SSH_PUBKEY_RC.match(data) - if not m: - raise ValueError("Invalid line format") - key_type = orig_key_type = m.group(1) - key_body = m.group(2) - with_cert = False - if key_type.endswith(_CERT_SUFFIX): - with_cert = True - key_type = key_type[: -len(_CERT_SUFFIX)] - if key_type == _SSH_DSA and not _legacy_dsa_allowed: - raise UnsupportedAlgorithm( - "DSA keys aren't supported in SSH certificates" - ) - kformat = _lookup_kformat(key_type) - - try: - rest = memoryview(binascii.a2b_base64(key_body)) - except (TypeError, binascii.Error): - raise ValueError("Invalid format") - - if with_cert: - cert_body = rest - inner_key_type, rest = _get_sshstr(rest) - if inner_key_type != orig_key_type: - raise ValueError("Invalid key format") - if with_cert: - nonce, rest = _get_sshstr(rest) - public_key, rest = kformat.load_public(rest) - if with_cert: - serial, rest = _get_u64(rest) - cctype, rest = _get_u32(rest) - key_id, rest = _get_sshstr(rest) - principals, rest = _get_sshstr(rest) - valid_principals = [] - while principals: - principal, principals = _get_sshstr(principals) - valid_principals.append(bytes(principal)) - valid_after, rest = _get_u64(rest) - valid_before, rest = _get_u64(rest) - crit_options, rest = _get_sshstr(rest) - critical_options = _parse_exts_opts(crit_options) - exts, rest = _get_sshstr(rest) - extensions = _parse_exts_opts(exts) - # Get the reserved field, which is unused. - _, rest = _get_sshstr(rest) - sig_key_raw, rest = _get_sshstr(rest) - sig_type, sig_key = _get_sshstr(sig_key_raw) - if sig_type == _SSH_DSA and not _legacy_dsa_allowed: - raise UnsupportedAlgorithm( - "DSA signatures aren't supported in SSH certificates" - ) - # Get the entire cert body and subtract the signature - tbs_cert_body = cert_body[: -len(rest)] - signature_raw, rest = _get_sshstr(rest) - _check_empty(rest) - inner_sig_type, sig_rest = _get_sshstr(signature_raw) - # RSA certs can have multiple algorithm types - if ( - sig_type == _SSH_RSA - and inner_sig_type - not in [_SSH_RSA_SHA256, _SSH_RSA_SHA512, _SSH_RSA] - ) or (sig_type != _SSH_RSA and inner_sig_type != sig_type): - raise ValueError("Signature key type does not match") - signature, sig_rest = _get_sshstr(sig_rest) - _check_empty(sig_rest) - return SSHCertificate( - nonce, - public_key, - serial, - cctype, - key_id, - valid_principals, - valid_after, - valid_before, - critical_options, - extensions, - sig_type, - sig_key, - inner_sig_type, - signature, - tbs_cert_body, - orig_key_type, - cert_body, - ) - else: - _check_empty(rest) - return public_key - - -def load_ssh_public_identity( - data: bytes, -) -> SSHCertificate | SSHPublicKeyTypes: - return _load_ssh_public_identity(data) - - -def _parse_exts_opts(exts_opts: memoryview) -> dict[bytes, bytes]: - result: dict[bytes, bytes] = {} - last_name = None - while exts_opts: - name, exts_opts = _get_sshstr(exts_opts) - bname: bytes = bytes(name) - if bname in result: - raise ValueError("Duplicate name") - if last_name is not None and bname < last_name: - raise ValueError("Fields not lexically sorted") - value, exts_opts = _get_sshstr(exts_opts) - if len(value) > 0: - value, extra = _get_sshstr(value) - if len(extra) > 0: - raise ValueError("Unexpected extra data after value") - result[bname] = bytes(value) - last_name = bname - return result - - -def ssh_key_fingerprint( - key: SSHPublicKeyTypes, - hash_algorithm: hashes.MD5 | hashes.SHA256, -) -> bytes: - if not isinstance(hash_algorithm, (hashes.MD5, hashes.SHA256)): - raise TypeError("hash_algorithm must be either MD5 or SHA256") - - key_type = _get_ssh_key_type(key) - kformat = _lookup_kformat(key_type) - - f_pub = _FragList() - f_pub.put_sshstr(key_type) - kformat.encode_public(key, f_pub) - - ssh_binary_data = f_pub.tobytes() - - # Hash the binary data - hash_obj = hashes.Hash(hash_algorithm) - hash_obj.update(ssh_binary_data) - return hash_obj.finalize() - - -def load_ssh_public_key( - data: utils.Buffer, backend: typing.Any = None -) -> SSHPublicKeyTypes: - cert_or_key = _load_ssh_public_identity(data, _legacy_dsa_allowed=True) - public_key: SSHPublicKeyTypes - if isinstance(cert_or_key, SSHCertificate): - public_key = cert_or_key.public_key() - else: - public_key = cert_or_key - - if isinstance(public_key, dsa.DSAPublicKey): - warnings.warn( - "SSH DSA keys are deprecated and will be removed in a future " - "release.", - utils.DeprecatedIn40, - stacklevel=2, - ) - return public_key - - -def serialize_ssh_public_key(public_key: SSHPublicKeyTypes) -> bytes: - """One-line public key format for OpenSSH""" - if isinstance(public_key, dsa.DSAPublicKey): - warnings.warn( - "SSH DSA key support is deprecated and will be " - "removed in a future release", - utils.DeprecatedIn40, - stacklevel=4, - ) - key_type = _get_ssh_key_type(public_key) - kformat = _lookup_kformat(key_type) - - f_pub = _FragList() - f_pub.put_sshstr(key_type) - kformat.encode_public(public_key, f_pub) - - pub = binascii.b2a_base64(f_pub.tobytes()).strip() - return b"".join([key_type, b" ", pub]) - - -SSHCertPrivateKeyTypes = typing.Union[ - ec.EllipticCurvePrivateKey, - rsa.RSAPrivateKey, - ed25519.Ed25519PrivateKey, -] - - -# This is an undocumented limit enforced in the openssh codebase for sshd and -# ssh-keygen, but it is undefined in the ssh certificates spec. -_SSHKEY_CERT_MAX_PRINCIPALS = 256 - - -class SSHCertificateBuilder: - def __init__( - self, - _public_key: SSHCertPublicKeyTypes | None = None, - _serial: int | None = None, - _type: SSHCertificateType | None = None, - _key_id: bytes | None = None, - _valid_principals: list[bytes] = [], - _valid_for_all_principals: bool = False, - _valid_before: int | None = None, - _valid_after: int | None = None, - _critical_options: list[tuple[bytes, bytes]] = [], - _extensions: list[tuple[bytes, bytes]] = [], - ): - self._public_key = _public_key - self._serial = _serial - self._type = _type - self._key_id = _key_id - self._valid_principals = _valid_principals - self._valid_for_all_principals = _valid_for_all_principals - self._valid_before = _valid_before - self._valid_after = _valid_after - self._critical_options = _critical_options - self._extensions = _extensions - - def public_key( - self, public_key: SSHCertPublicKeyTypes - ) -> SSHCertificateBuilder: - if not isinstance( - public_key, - ( - ec.EllipticCurvePublicKey, - rsa.RSAPublicKey, - ed25519.Ed25519PublicKey, - ), - ): - raise TypeError("Unsupported key type") - if self._public_key is not None: - raise ValueError("public_key already set") - - return SSHCertificateBuilder( - _public_key=public_key, - _serial=self._serial, - _type=self._type, - _key_id=self._key_id, - _valid_principals=self._valid_principals, - _valid_for_all_principals=self._valid_for_all_principals, - _valid_before=self._valid_before, - _valid_after=self._valid_after, - _critical_options=self._critical_options, - _extensions=self._extensions, - ) - - def serial(self, serial: int) -> SSHCertificateBuilder: - if not isinstance(serial, int): - raise TypeError("serial must be an integer") - if not 0 <= serial < 2**64: - raise ValueError("serial must be between 0 and 2**64") - if self._serial is not None: - raise ValueError("serial already set") - - return SSHCertificateBuilder( - _public_key=self._public_key, - _serial=serial, - _type=self._type, - _key_id=self._key_id, - _valid_principals=self._valid_principals, - _valid_for_all_principals=self._valid_for_all_principals, - _valid_before=self._valid_before, - _valid_after=self._valid_after, - _critical_options=self._critical_options, - _extensions=self._extensions, - ) - - def type(self, type: SSHCertificateType) -> SSHCertificateBuilder: - if not isinstance(type, SSHCertificateType): - raise TypeError("type must be an SSHCertificateType") - if self._type is not None: - raise ValueError("type already set") - - return SSHCertificateBuilder( - _public_key=self._public_key, - _serial=self._serial, - _type=type, - _key_id=self._key_id, - _valid_principals=self._valid_principals, - _valid_for_all_principals=self._valid_for_all_principals, - _valid_before=self._valid_before, - _valid_after=self._valid_after, - _critical_options=self._critical_options, - _extensions=self._extensions, - ) - - def key_id(self, key_id: bytes) -> SSHCertificateBuilder: - if not isinstance(key_id, bytes): - raise TypeError("key_id must be bytes") - if self._key_id is not None: - raise ValueError("key_id already set") - - return SSHCertificateBuilder( - _public_key=self._public_key, - _serial=self._serial, - _type=self._type, - _key_id=key_id, - _valid_principals=self._valid_principals, - _valid_for_all_principals=self._valid_for_all_principals, - _valid_before=self._valid_before, - _valid_after=self._valid_after, - _critical_options=self._critical_options, - _extensions=self._extensions, - ) - - def valid_principals( - self, valid_principals: list[bytes] - ) -> SSHCertificateBuilder: - if self._valid_for_all_principals: - raise ValueError( - "Principals can't be set because the cert is valid " - "for all principals" - ) - if ( - not all(isinstance(x, bytes) for x in valid_principals) - or not valid_principals - ): - raise TypeError( - "principals must be a list of bytes and can't be empty" - ) - if self._valid_principals: - raise ValueError("valid_principals already set") - - if len(valid_principals) > _SSHKEY_CERT_MAX_PRINCIPALS: - raise ValueError( - "Reached or exceeded the maximum number of valid_principals" - ) - - return SSHCertificateBuilder( - _public_key=self._public_key, - _serial=self._serial, - _type=self._type, - _key_id=self._key_id, - _valid_principals=valid_principals, - _valid_for_all_principals=self._valid_for_all_principals, - _valid_before=self._valid_before, - _valid_after=self._valid_after, - _critical_options=self._critical_options, - _extensions=self._extensions, - ) - - def valid_for_all_principals(self): - if self._valid_principals: - raise ValueError( - "valid_principals already set, can't set " - "valid_for_all_principals" - ) - if self._valid_for_all_principals: - raise ValueError("valid_for_all_principals already set") - - return SSHCertificateBuilder( - _public_key=self._public_key, - _serial=self._serial, - _type=self._type, - _key_id=self._key_id, - _valid_principals=self._valid_principals, - _valid_for_all_principals=True, - _valid_before=self._valid_before, - _valid_after=self._valid_after, - _critical_options=self._critical_options, - _extensions=self._extensions, - ) - - def valid_before(self, valid_before: int | float) -> SSHCertificateBuilder: - if not isinstance(valid_before, (int, float)): - raise TypeError("valid_before must be an int or float") - valid_before = int(valid_before) - if valid_before < 0 or valid_before >= 2**64: - raise ValueError("valid_before must [0, 2**64)") - if self._valid_before is not None: - raise ValueError("valid_before already set") - - return SSHCertificateBuilder( - _public_key=self._public_key, - _serial=self._serial, - _type=self._type, - _key_id=self._key_id, - _valid_principals=self._valid_principals, - _valid_for_all_principals=self._valid_for_all_principals, - _valid_before=valid_before, - _valid_after=self._valid_after, - _critical_options=self._critical_options, - _extensions=self._extensions, - ) - - def valid_after(self, valid_after: int | float) -> SSHCertificateBuilder: - if not isinstance(valid_after, (int, float)): - raise TypeError("valid_after must be an int or float") - valid_after = int(valid_after) - if valid_after < 0 or valid_after >= 2**64: - raise ValueError("valid_after must [0, 2**64)") - if self._valid_after is not None: - raise ValueError("valid_after already set") - - return SSHCertificateBuilder( - _public_key=self._public_key, - _serial=self._serial, - _type=self._type, - _key_id=self._key_id, - _valid_principals=self._valid_principals, - _valid_for_all_principals=self._valid_for_all_principals, - _valid_before=self._valid_before, - _valid_after=valid_after, - _critical_options=self._critical_options, - _extensions=self._extensions, - ) - - def add_critical_option( - self, name: bytes, value: bytes - ) -> SSHCertificateBuilder: - if not isinstance(name, bytes) or not isinstance(value, bytes): - raise TypeError("name and value must be bytes") - # This is O(n**2) - if name in [name for name, _ in self._critical_options]: - raise ValueError("Duplicate critical option name") - - return SSHCertificateBuilder( - _public_key=self._public_key, - _serial=self._serial, - _type=self._type, - _key_id=self._key_id, - _valid_principals=self._valid_principals, - _valid_for_all_principals=self._valid_for_all_principals, - _valid_before=self._valid_before, - _valid_after=self._valid_after, - _critical_options=[*self._critical_options, (name, value)], - _extensions=self._extensions, - ) - - def add_extension( - self, name: bytes, value: bytes - ) -> SSHCertificateBuilder: - if not isinstance(name, bytes) or not isinstance(value, bytes): - raise TypeError("name and value must be bytes") - # This is O(n**2) - if name in [name for name, _ in self._extensions]: - raise ValueError("Duplicate extension name") - - return SSHCertificateBuilder( - _public_key=self._public_key, - _serial=self._serial, - _type=self._type, - _key_id=self._key_id, - _valid_principals=self._valid_principals, - _valid_for_all_principals=self._valid_for_all_principals, - _valid_before=self._valid_before, - _valid_after=self._valid_after, - _critical_options=self._critical_options, - _extensions=[*self._extensions, (name, value)], - ) - - def sign(self, private_key: SSHCertPrivateKeyTypes) -> SSHCertificate: - if not isinstance( - private_key, - ( - ec.EllipticCurvePrivateKey, - rsa.RSAPrivateKey, - ed25519.Ed25519PrivateKey, - ), - ): - raise TypeError("Unsupported private key type") - - if self._public_key is None: - raise ValueError("public_key must be set") - - # Not required - serial = 0 if self._serial is None else self._serial - - if self._type is None: - raise ValueError("type must be set") - - # Not required - key_id = b"" if self._key_id is None else self._key_id - - # A zero length list is valid, but means the certificate - # is valid for any principal of the specified type. We require - # the user to explicitly set valid_for_all_principals to get - # that behavior. - if not self._valid_principals and not self._valid_for_all_principals: - raise ValueError( - "valid_principals must be set if valid_for_all_principals " - "is False" - ) - - if self._valid_before is None: - raise ValueError("valid_before must be set") - - if self._valid_after is None: - raise ValueError("valid_after must be set") - - if self._valid_after > self._valid_before: - raise ValueError("valid_after must be earlier than valid_before") - - # lexically sort our byte strings - self._critical_options.sort(key=lambda x: x[0]) - self._extensions.sort(key=lambda x: x[0]) - - key_type = _get_ssh_key_type(self._public_key) - cert_prefix = key_type + _CERT_SUFFIX - - # Marshal the bytes to be signed - nonce = os.urandom(32) - kformat = _lookup_kformat(key_type) - f = _FragList() - f.put_sshstr(cert_prefix) - f.put_sshstr(nonce) - kformat.encode_public(self._public_key, f) - f.put_u64(serial) - f.put_u32(self._type.value) - f.put_sshstr(key_id) - fprincipals = _FragList() - for p in self._valid_principals: - fprincipals.put_sshstr(p) - f.put_sshstr(fprincipals.tobytes()) - f.put_u64(self._valid_after) - f.put_u64(self._valid_before) - fcrit = _FragList() - for name, value in self._critical_options: - fcrit.put_sshstr(name) - if len(value) > 0: - foptval = _FragList() - foptval.put_sshstr(value) - fcrit.put_sshstr(foptval.tobytes()) - else: - fcrit.put_sshstr(value) - f.put_sshstr(fcrit.tobytes()) - fext = _FragList() - for name, value in self._extensions: - fext.put_sshstr(name) - if len(value) > 0: - fextval = _FragList() - fextval.put_sshstr(value) - fext.put_sshstr(fextval.tobytes()) - else: - fext.put_sshstr(value) - f.put_sshstr(fext.tobytes()) - f.put_sshstr(b"") # RESERVED FIELD - # encode CA public key - ca_type = _get_ssh_key_type(private_key) - caformat = _lookup_kformat(ca_type) - caf = _FragList() - caf.put_sshstr(ca_type) - caformat.encode_public(private_key.public_key(), caf) - f.put_sshstr(caf.tobytes()) - # Sigs according to the rules defined for the CA's public key - # (RFC4253 section 6.6 for ssh-rsa, RFC5656 for ECDSA, - # and RFC8032 for Ed25519). - if isinstance(private_key, ed25519.Ed25519PrivateKey): - signature = private_key.sign(f.tobytes()) - fsig = _FragList() - fsig.put_sshstr(ca_type) - fsig.put_sshstr(signature) - f.put_sshstr(fsig.tobytes()) - elif isinstance(private_key, ec.EllipticCurvePrivateKey): - hash_alg = _get_ec_hash_alg(private_key.curve) - signature = private_key.sign(f.tobytes(), ec.ECDSA(hash_alg)) - r, s = asym_utils.decode_dss_signature(signature) - fsig = _FragList() - fsig.put_sshstr(ca_type) - fsigblob = _FragList() - fsigblob.put_mpint(r) - fsigblob.put_mpint(s) - fsig.put_sshstr(fsigblob.tobytes()) - f.put_sshstr(fsig.tobytes()) - - else: - assert isinstance(private_key, rsa.RSAPrivateKey) - # Just like Golang, we're going to use SHA512 for RSA - # https://cs.opensource.google/go/x/crypto/+/refs/tags/ - # v0.4.0:ssh/certs.go;l=445 - # RFC 8332 defines SHA256 and 512 as options - fsig = _FragList() - fsig.put_sshstr(_SSH_RSA_SHA512) - signature = private_key.sign( - f.tobytes(), padding.PKCS1v15(), hashes.SHA512() - ) - fsig.put_sshstr(signature) - f.put_sshstr(fsig.tobytes()) - - cert_data = binascii.b2a_base64(f.tobytes()).strip() - # load_ssh_public_identity returns a union, but this is - # guaranteed to be an SSHCertificate, so we cast to make - # mypy happy. - return typing.cast( - SSHCertificate, - load_ssh_public_identity(b"".join([cert_prefix, b" ", cert_data])), - ) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/__init__.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/__init__.py deleted file mode 100644 index c1af423004863e7db3e163f61a6baa4ddc157351..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - - -class InvalidToken(Exception): - pass diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 6a37d103667fa7f28a35f918642dc75782f346a2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/__pycache__/hotp.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/__pycache__/hotp.cpython-310.pyc deleted file mode 100644 index 0a994f61dc301a4b948e6cb4cd31a3cd3c072871..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/__pycache__/hotp.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/__pycache__/totp.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/__pycache__/totp.cpython-310.pyc deleted file mode 100644 index b7ce1720d3aaf90f5215a2ff1cec677b737d26d0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/__pycache__/totp.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/hotp.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/hotp.py deleted file mode 100644 index 21fb000499c43cb8492ce60d9a2c5ed6fab704e2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/hotp.py +++ /dev/null @@ -1,101 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import base64 -import typing -from urllib.parse import quote, urlencode - -from cryptography.hazmat.primitives import constant_time, hmac -from cryptography.hazmat.primitives.hashes import SHA1, SHA256, SHA512 -from cryptography.hazmat.primitives.twofactor import InvalidToken -from cryptography.utils import Buffer - -HOTPHashTypes = typing.Union[SHA1, SHA256, SHA512] - - -def _generate_uri( - hotp: HOTP, - type_name: str, - account_name: str, - issuer: str | None, - extra_parameters: list[tuple[str, int]], -) -> str: - parameters = [ - ("digits", hotp._length), - ("secret", base64.b32encode(hotp._key)), - ("algorithm", hotp._algorithm.name.upper()), - ] - - if issuer is not None: - parameters.append(("issuer", issuer)) - - parameters.extend(extra_parameters) - - label = ( - f"{quote(issuer)}:{quote(account_name)}" - if issuer - else quote(account_name) - ) - return f"otpauth://{type_name}/{label}?{urlencode(parameters)}" - - -class HOTP: - def __init__( - self, - key: Buffer, - length: int, - algorithm: HOTPHashTypes, - backend: typing.Any = None, - enforce_key_length: bool = True, - ) -> None: - if len(key) < 16 and enforce_key_length is True: - raise ValueError("Key length has to be at least 128 bits.") - - if not isinstance(length, int): - raise TypeError("Length parameter must be an integer type.") - - if length < 6 or length > 8: - raise ValueError("Length of HOTP has to be between 6 and 8.") - - if not isinstance(algorithm, (SHA1, SHA256, SHA512)): - raise TypeError("Algorithm must be SHA1, SHA256 or SHA512.") - - self._key = key - self._length = length - self._algorithm = algorithm - - def generate(self, counter: int) -> bytes: - if not isinstance(counter, int): - raise TypeError("Counter parameter must be an integer type.") - - truncated_value = self._dynamic_truncate(counter) - hotp = truncated_value % (10**self._length) - return "{0:0{1}}".format(hotp, self._length).encode() - - def verify(self, hotp: bytes, counter: int) -> None: - if not constant_time.bytes_eq(self.generate(counter), hotp): - raise InvalidToken("Supplied HOTP value does not match.") - - def _dynamic_truncate(self, counter: int) -> int: - ctx = hmac.HMAC(self._key, self._algorithm) - - try: - ctx.update(counter.to_bytes(length=8, byteorder="big")) - except OverflowError: - raise ValueError(f"Counter must be between 0 and {2**64 - 1}.") - - hmac_value = ctx.finalize() - - offset = hmac_value[len(hmac_value) - 1] & 0b1111 - p = hmac_value[offset : offset + 4] - return int.from_bytes(p, byteorder="big") & 0x7FFFFFFF - - def get_provisioning_uri( - self, account_name: str, counter: int, issuer: str | None - ) -> str: - return _generate_uri( - self, "hotp", account_name, issuer, [("counter", int(counter))] - ) diff --git a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/totp.py b/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/totp.py deleted file mode 100644 index 10c725cc0ab0c9fff5a978a3d5278ccce52d4e67..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/hazmat/primitives/twofactor/totp.py +++ /dev/null @@ -1,56 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import typing - -from cryptography.hazmat.primitives import constant_time -from cryptography.hazmat.primitives.twofactor import InvalidToken -from cryptography.hazmat.primitives.twofactor.hotp import ( - HOTP, - HOTPHashTypes, - _generate_uri, -) -from cryptography.utils import Buffer - - -class TOTP: - def __init__( - self, - key: Buffer, - length: int, - algorithm: HOTPHashTypes, - time_step: int, - backend: typing.Any = None, - enforce_key_length: bool = True, - ): - self._time_step = time_step - self._hotp = HOTP( - key, length, algorithm, enforce_key_length=enforce_key_length - ) - - def generate(self, time: int | float) -> bytes: - if not isinstance(time, (int, float)): - raise TypeError( - "Time parameter must be an integer type or float type." - ) - - counter = int(time / self._time_step) - return self._hotp.generate(counter) - - def verify(self, totp: bytes, time: int) -> None: - if not constant_time.bytes_eq(self.generate(time), totp): - raise InvalidToken("Supplied TOTP value does not match.") - - def get_provisioning_uri( - self, account_name: str, issuer: str | None - ) -> str: - return _generate_uri( - self._hotp, - "totp", - account_name, - issuer, - [("period", int(self._time_step))], - ) diff --git a/venv/lib/python3.10/site-packages/cryptography/py.typed b/venv/lib/python3.10/site-packages/cryptography/py.typed deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/cryptography/utils.py b/venv/lib/python3.10/site-packages/cryptography/utils.py deleted file mode 100644 index 7de7e4dfc06b52e21a108f1a2dc10e8f19f537cf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/utils.py +++ /dev/null @@ -1,139 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import enum -import sys -import types -import typing -import warnings -from collections.abc import Callable, Sequence - - -# We use a UserWarning subclass, instead of DeprecationWarning, because CPython -# decided deprecation warnings should be invisible by default. -class CryptographyDeprecationWarning(UserWarning): - pass - - -# Several APIs were deprecated with no specific end-of-life date because of the -# ubiquity of their use. They should not be removed until we agree on when that -# cycle ends. -DeprecatedIn36 = CryptographyDeprecationWarning -DeprecatedIn37 = CryptographyDeprecationWarning -DeprecatedIn40 = CryptographyDeprecationWarning -DeprecatedIn41 = CryptographyDeprecationWarning -DeprecatedIn42 = CryptographyDeprecationWarning -DeprecatedIn43 = CryptographyDeprecationWarning -DeprecatedIn45 = CryptographyDeprecationWarning - - -# If you're wondering why we don't use `Buffer`, it's because `Buffer` would -# be more accurately named: Bufferable. It means something which has an -# `__buffer__`. Which means you can't actually treat the result as a buffer -# (and do things like take a `len()`). -if sys.version_info >= (3, 9): - Buffer = typing.Union[bytes, bytearray, memoryview] -else: - Buffer = typing.ByteString - - -def _check_bytes(name: str, value: bytes) -> None: - if not isinstance(value, bytes): - raise TypeError(f"{name} must be bytes") - - -def _check_byteslike(name: str, value: Buffer) -> None: - try: - memoryview(value) - except TypeError: - raise TypeError(f"{name} must be bytes-like") - - -def int_to_bytes(integer: int, length: int | None = None) -> bytes: - if length == 0: - raise ValueError("length argument can't be 0") - return integer.to_bytes( - length or (integer.bit_length() + 7) // 8 or 1, "big" - ) - - -class InterfaceNotImplemented(Exception): - pass - - -class _DeprecatedValue: - def __init__(self, value: object, message: str, warning_class): - self.value = value - self.message = message - self.warning_class = warning_class - - -class _ModuleWithDeprecations(types.ModuleType): - def __init__(self, module: types.ModuleType): - super().__init__(module.__name__) - self.__dict__["_module"] = module - - def __getattr__(self, attr: str) -> object: - obj = getattr(self._module, attr) - if isinstance(obj, _DeprecatedValue): - warnings.warn(obj.message, obj.warning_class, stacklevel=2) - obj = obj.value - return obj - - def __setattr__(self, attr: str, value: object) -> None: - setattr(self._module, attr, value) - - def __delattr__(self, attr: str) -> None: - obj = getattr(self._module, attr) - if isinstance(obj, _DeprecatedValue): - warnings.warn(obj.message, obj.warning_class, stacklevel=2) - - delattr(self._module, attr) - - def __dir__(self) -> Sequence[str]: - return ["_module", *dir(self._module)] - - -def deprecated( - value: object, - module_name: str, - message: str, - warning_class: type[Warning], - name: str | None = None, -) -> _DeprecatedValue: - module = sys.modules[module_name] - if not isinstance(module, _ModuleWithDeprecations): - sys.modules[module_name] = module = _ModuleWithDeprecations(module) - dv = _DeprecatedValue(value, message, warning_class) - # Maintain backwards compatibility with `name is None` for pyOpenSSL. - if name is not None: - setattr(module, name, dv) - return dv - - -def cached_property(func: Callable) -> property: - cached_name = f"_cached_{func}" - sentinel = object() - - def inner(instance: object): - cache = getattr(instance, cached_name, sentinel) - if cache is not sentinel: - return cache - result = func(instance) - setattr(instance, cached_name, result) - return result - - return property(inner) - - -# Python 3.10 changed representation of enums. We use well-defined object -# representation and string representation from Python 3.9. -class Enum(enum.Enum): - def __repr__(self) -> str: - return f"<{self.__class__.__name__}.{self._name_}: {self._value_!r}>" - - def __str__(self) -> str: - return f"{self.__class__.__name__}.{self._name_}" diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/__init__.py b/venv/lib/python3.10/site-packages/cryptography/x509/__init__.py deleted file mode 100644 index 318eecc96e1df01dc3e26bb236c55cbfd77d05c1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/x509/__init__.py +++ /dev/null @@ -1,270 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography.x509 import certificate_transparency, verification -from cryptography.x509.base import ( - Attribute, - AttributeNotFound, - Attributes, - Certificate, - CertificateBuilder, - CertificateRevocationList, - CertificateRevocationListBuilder, - CertificateSigningRequest, - CertificateSigningRequestBuilder, - InvalidVersion, - RevokedCertificate, - RevokedCertificateBuilder, - Version, - load_der_x509_certificate, - load_der_x509_crl, - load_der_x509_csr, - load_pem_x509_certificate, - load_pem_x509_certificates, - load_pem_x509_crl, - load_pem_x509_csr, - random_serial_number, -) -from cryptography.x509.extensions import ( - AccessDescription, - Admission, - Admissions, - AuthorityInformationAccess, - AuthorityKeyIdentifier, - BasicConstraints, - CertificateIssuer, - CertificatePolicies, - CRLDistributionPoints, - CRLNumber, - CRLReason, - DeltaCRLIndicator, - DistributionPoint, - DuplicateExtension, - ExtendedKeyUsage, - Extension, - ExtensionNotFound, - Extensions, - ExtensionType, - FreshestCRL, - GeneralNames, - InhibitAnyPolicy, - InvalidityDate, - IssuerAlternativeName, - IssuingDistributionPoint, - KeyUsage, - MSCertificateTemplate, - NameConstraints, - NamingAuthority, - NoticeReference, - OCSPAcceptableResponses, - OCSPNoCheck, - OCSPNonce, - PolicyConstraints, - PolicyInformation, - PrecertificateSignedCertificateTimestamps, - PrecertPoison, - PrivateKeyUsagePeriod, - ProfessionInfo, - ReasonFlags, - SignedCertificateTimestamps, - SubjectAlternativeName, - SubjectInformationAccess, - SubjectKeyIdentifier, - TLSFeature, - TLSFeatureType, - UnrecognizedExtension, - UserNotice, -) -from cryptography.x509.general_name import ( - DirectoryName, - DNSName, - GeneralName, - IPAddress, - OtherName, - RegisteredID, - RFC822Name, - UniformResourceIdentifier, - UnsupportedGeneralNameType, -) -from cryptography.x509.name import ( - Name, - NameAttribute, - RelativeDistinguishedName, -) -from cryptography.x509.oid import ( - AuthorityInformationAccessOID, - CertificatePoliciesOID, - CRLEntryExtensionOID, - ExtendedKeyUsageOID, - ExtensionOID, - NameOID, - ObjectIdentifier, - PublicKeyAlgorithmOID, - SignatureAlgorithmOID, -) - -OID_AUTHORITY_INFORMATION_ACCESS = ExtensionOID.AUTHORITY_INFORMATION_ACCESS -OID_AUTHORITY_KEY_IDENTIFIER = ExtensionOID.AUTHORITY_KEY_IDENTIFIER -OID_BASIC_CONSTRAINTS = ExtensionOID.BASIC_CONSTRAINTS -OID_CERTIFICATE_POLICIES = ExtensionOID.CERTIFICATE_POLICIES -OID_CRL_DISTRIBUTION_POINTS = ExtensionOID.CRL_DISTRIBUTION_POINTS -OID_EXTENDED_KEY_USAGE = ExtensionOID.EXTENDED_KEY_USAGE -OID_FRESHEST_CRL = ExtensionOID.FRESHEST_CRL -OID_INHIBIT_ANY_POLICY = ExtensionOID.INHIBIT_ANY_POLICY -OID_ISSUER_ALTERNATIVE_NAME = ExtensionOID.ISSUER_ALTERNATIVE_NAME -OID_KEY_USAGE = ExtensionOID.KEY_USAGE -OID_PRIVATE_KEY_USAGE_PERIOD = ExtensionOID.PRIVATE_KEY_USAGE_PERIOD -OID_NAME_CONSTRAINTS = ExtensionOID.NAME_CONSTRAINTS -OID_OCSP_NO_CHECK = ExtensionOID.OCSP_NO_CHECK -OID_POLICY_CONSTRAINTS = ExtensionOID.POLICY_CONSTRAINTS -OID_POLICY_MAPPINGS = ExtensionOID.POLICY_MAPPINGS -OID_SUBJECT_ALTERNATIVE_NAME = ExtensionOID.SUBJECT_ALTERNATIVE_NAME -OID_SUBJECT_DIRECTORY_ATTRIBUTES = ExtensionOID.SUBJECT_DIRECTORY_ATTRIBUTES -OID_SUBJECT_INFORMATION_ACCESS = ExtensionOID.SUBJECT_INFORMATION_ACCESS -OID_SUBJECT_KEY_IDENTIFIER = ExtensionOID.SUBJECT_KEY_IDENTIFIER - -OID_DSA_WITH_SHA1 = SignatureAlgorithmOID.DSA_WITH_SHA1 -OID_DSA_WITH_SHA224 = SignatureAlgorithmOID.DSA_WITH_SHA224 -OID_DSA_WITH_SHA256 = SignatureAlgorithmOID.DSA_WITH_SHA256 -OID_ECDSA_WITH_SHA1 = SignatureAlgorithmOID.ECDSA_WITH_SHA1 -OID_ECDSA_WITH_SHA224 = SignatureAlgorithmOID.ECDSA_WITH_SHA224 -OID_ECDSA_WITH_SHA256 = SignatureAlgorithmOID.ECDSA_WITH_SHA256 -OID_ECDSA_WITH_SHA384 = SignatureAlgorithmOID.ECDSA_WITH_SHA384 -OID_ECDSA_WITH_SHA512 = SignatureAlgorithmOID.ECDSA_WITH_SHA512 -OID_RSA_WITH_MD5 = SignatureAlgorithmOID.RSA_WITH_MD5 -OID_RSA_WITH_SHA1 = SignatureAlgorithmOID.RSA_WITH_SHA1 -OID_RSA_WITH_SHA224 = SignatureAlgorithmOID.RSA_WITH_SHA224 -OID_RSA_WITH_SHA256 = SignatureAlgorithmOID.RSA_WITH_SHA256 -OID_RSA_WITH_SHA384 = SignatureAlgorithmOID.RSA_WITH_SHA384 -OID_RSA_WITH_SHA512 = SignatureAlgorithmOID.RSA_WITH_SHA512 -OID_RSASSA_PSS = SignatureAlgorithmOID.RSASSA_PSS - -OID_COMMON_NAME = NameOID.COMMON_NAME -OID_COUNTRY_NAME = NameOID.COUNTRY_NAME -OID_DOMAIN_COMPONENT = NameOID.DOMAIN_COMPONENT -OID_DN_QUALIFIER = NameOID.DN_QUALIFIER -OID_EMAIL_ADDRESS = NameOID.EMAIL_ADDRESS -OID_GENERATION_QUALIFIER = NameOID.GENERATION_QUALIFIER -OID_GIVEN_NAME = NameOID.GIVEN_NAME -OID_LOCALITY_NAME = NameOID.LOCALITY_NAME -OID_ORGANIZATIONAL_UNIT_NAME = NameOID.ORGANIZATIONAL_UNIT_NAME -OID_ORGANIZATION_NAME = NameOID.ORGANIZATION_NAME -OID_PSEUDONYM = NameOID.PSEUDONYM -OID_SERIAL_NUMBER = NameOID.SERIAL_NUMBER -OID_STATE_OR_PROVINCE_NAME = NameOID.STATE_OR_PROVINCE_NAME -OID_SURNAME = NameOID.SURNAME -OID_TITLE = NameOID.TITLE - -OID_CLIENT_AUTH = ExtendedKeyUsageOID.CLIENT_AUTH -OID_CODE_SIGNING = ExtendedKeyUsageOID.CODE_SIGNING -OID_EMAIL_PROTECTION = ExtendedKeyUsageOID.EMAIL_PROTECTION -OID_OCSP_SIGNING = ExtendedKeyUsageOID.OCSP_SIGNING -OID_SERVER_AUTH = ExtendedKeyUsageOID.SERVER_AUTH -OID_TIME_STAMPING = ExtendedKeyUsageOID.TIME_STAMPING - -OID_ANY_POLICY = CertificatePoliciesOID.ANY_POLICY -OID_CPS_QUALIFIER = CertificatePoliciesOID.CPS_QUALIFIER -OID_CPS_USER_NOTICE = CertificatePoliciesOID.CPS_USER_NOTICE - -OID_CERTIFICATE_ISSUER = CRLEntryExtensionOID.CERTIFICATE_ISSUER -OID_CRL_REASON = CRLEntryExtensionOID.CRL_REASON -OID_INVALIDITY_DATE = CRLEntryExtensionOID.INVALIDITY_DATE - -OID_CA_ISSUERS = AuthorityInformationAccessOID.CA_ISSUERS -OID_OCSP = AuthorityInformationAccessOID.OCSP - -__all__ = [ - "OID_CA_ISSUERS", - "OID_OCSP", - "AccessDescription", - "Admission", - "Admissions", - "Attribute", - "AttributeNotFound", - "Attributes", - "AuthorityInformationAccess", - "AuthorityKeyIdentifier", - "BasicConstraints", - "CRLDistributionPoints", - "CRLNumber", - "CRLReason", - "Certificate", - "CertificateBuilder", - "CertificateIssuer", - "CertificatePolicies", - "CertificateRevocationList", - "CertificateRevocationListBuilder", - "CertificateSigningRequest", - "CertificateSigningRequestBuilder", - "DNSName", - "DeltaCRLIndicator", - "DirectoryName", - "DistributionPoint", - "DuplicateExtension", - "ExtendedKeyUsage", - "Extension", - "ExtensionNotFound", - "ExtensionType", - "Extensions", - "FreshestCRL", - "GeneralName", - "GeneralNames", - "IPAddress", - "InhibitAnyPolicy", - "InvalidVersion", - "InvalidityDate", - "IssuerAlternativeName", - "IssuingDistributionPoint", - "KeyUsage", - "MSCertificateTemplate", - "Name", - "NameAttribute", - "NameConstraints", - "NameOID", - "NamingAuthority", - "NoticeReference", - "OCSPAcceptableResponses", - "OCSPNoCheck", - "OCSPNonce", - "ObjectIdentifier", - "OtherName", - "PolicyConstraints", - "PolicyInformation", - "PrecertPoison", - "PrecertificateSignedCertificateTimestamps", - "PrivateKeyUsagePeriod", - "ProfessionInfo", - "PublicKeyAlgorithmOID", - "RFC822Name", - "ReasonFlags", - "RegisteredID", - "RelativeDistinguishedName", - "RevokedCertificate", - "RevokedCertificateBuilder", - "SignatureAlgorithmOID", - "SignedCertificateTimestamps", - "SubjectAlternativeName", - "SubjectInformationAccess", - "SubjectKeyIdentifier", - "TLSFeature", - "TLSFeatureType", - "UniformResourceIdentifier", - "UnrecognizedExtension", - "UnsupportedGeneralNameType", - "UserNotice", - "Version", - "certificate_transparency", - "load_der_x509_certificate", - "load_der_x509_crl", - "load_der_x509_csr", - "load_pem_x509_certificate", - "load_pem_x509_certificates", - "load_pem_x509_crl", - "load_pem_x509_csr", - "random_serial_number", - "verification", - "verification", -] diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index a0bf1f38550c7e40297c3ddae07f6b96edb2a258..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/base.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/base.cpython-310.pyc deleted file mode 100644 index 70ed8466753e4b3643eaaa95fac3ae1c7d88d6ac..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/certificate_transparency.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/certificate_transparency.cpython-310.pyc deleted file mode 100644 index 78bec157bb733d83078de9771adcc0edf3d0b45f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/certificate_transparency.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/extensions.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/extensions.cpython-310.pyc deleted file mode 100644 index 4f8405cddf12e0872acfcb459c8c3933e56de278..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/extensions.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/general_name.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/general_name.cpython-310.pyc deleted file mode 100644 index 053a56f93005c3867e79518ee12b7666e9713ed1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/general_name.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/name.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/name.cpython-310.pyc deleted file mode 100644 index f476ba93b18e6c36a935bc6f4c0203c29464641a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/name.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/ocsp.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/ocsp.cpython-310.pyc deleted file mode 100644 index 040c6389a4f67ddaf3dc2168ad04e0bf2f223fe2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/ocsp.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/oid.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/oid.cpython-310.pyc deleted file mode 100644 index 25e1a265011c28241066cc6070705d41c236df96..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/oid.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/verification.cpython-310.pyc b/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/verification.cpython-310.pyc deleted file mode 100644 index 18dfd0eac1d8d194ed6f08c149d180199ffdd57d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/cryptography/x509/__pycache__/verification.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/base.py b/venv/lib/python3.10/site-packages/cryptography/x509/base.py deleted file mode 100644 index 1be612be607b4338ed30c06a33a4b5157517e11c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/x509/base.py +++ /dev/null @@ -1,848 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc -import datetime -import os -import typing -import warnings -from collections.abc import Iterable - -from cryptography import utils -from cryptography.hazmat.bindings._rust import x509 as rust_x509 -from cryptography.hazmat.primitives import hashes -from cryptography.hazmat.primitives.asymmetric import ( - dsa, - ec, - ed448, - ed25519, - padding, - rsa, - x448, - x25519, -) -from cryptography.hazmat.primitives.asymmetric.types import ( - CertificateIssuerPrivateKeyTypes, - CertificatePublicKeyTypes, -) -from cryptography.x509.extensions import ( - Extension, - Extensions, - ExtensionType, - _make_sequence_methods, -) -from cryptography.x509.name import Name, _ASN1Type -from cryptography.x509.oid import ObjectIdentifier - -_EARLIEST_UTC_TIME = datetime.datetime(1950, 1, 1) - -# This must be kept in sync with sign.rs's list of allowable types in -# identify_hash_type -_AllowedHashTypes = typing.Union[ - hashes.SHA224, - hashes.SHA256, - hashes.SHA384, - hashes.SHA512, - hashes.SHA3_224, - hashes.SHA3_256, - hashes.SHA3_384, - hashes.SHA3_512, -] - - -class AttributeNotFound(Exception): - def __init__(self, msg: str, oid: ObjectIdentifier) -> None: - super().__init__(msg) - self.oid = oid - - -def _reject_duplicate_extension( - extension: Extension[ExtensionType], - extensions: list[Extension[ExtensionType]], -) -> None: - # This is quadratic in the number of extensions - for e in extensions: - if e.oid == extension.oid: - raise ValueError("This extension has already been set.") - - -def _reject_duplicate_attribute( - oid: ObjectIdentifier, - attributes: list[tuple[ObjectIdentifier, bytes, int | None]], -) -> None: - # This is quadratic in the number of attributes - for attr_oid, _, _ in attributes: - if attr_oid == oid: - raise ValueError("This attribute has already been set.") - - -def _convert_to_naive_utc_time(time: datetime.datetime) -> datetime.datetime: - """Normalizes a datetime to a naive datetime in UTC. - - time -- datetime to normalize. Assumed to be in UTC if not timezone - aware. - """ - if time.tzinfo is not None: - offset = time.utcoffset() - offset = offset if offset else datetime.timedelta() - return time.replace(tzinfo=None) - offset - else: - return time - - -class Attribute: - def __init__( - self, - oid: ObjectIdentifier, - value: bytes, - _type: int = _ASN1Type.UTF8String.value, - ) -> None: - self._oid = oid - self._value = value - self._type = _type - - @property - def oid(self) -> ObjectIdentifier: - return self._oid - - @property - def value(self) -> bytes: - return self._value - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, Attribute): - return NotImplemented - - return ( - self.oid == other.oid - and self.value == other.value - and self._type == other._type - ) - - def __hash__(self) -> int: - return hash((self.oid, self.value, self._type)) - - -class Attributes: - def __init__( - self, - attributes: Iterable[Attribute], - ) -> None: - self._attributes = list(attributes) - - __len__, __iter__, __getitem__ = _make_sequence_methods("_attributes") - - def __repr__(self) -> str: - return f"" - - def get_attribute_for_oid(self, oid: ObjectIdentifier) -> Attribute: - for attr in self: - if attr.oid == oid: - return attr - - raise AttributeNotFound(f"No {oid} attribute was found", oid) - - -class Version(utils.Enum): - v1 = 0 - v3 = 2 - - -class InvalidVersion(Exception): - def __init__(self, msg: str, parsed_version: int) -> None: - super().__init__(msg) - self.parsed_version = parsed_version - - -Certificate = rust_x509.Certificate - - -class RevokedCertificate(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def serial_number(self) -> int: - """ - Returns the serial number of the revoked certificate. - """ - - @property - @abc.abstractmethod - def revocation_date(self) -> datetime.datetime: - """ - Returns the date of when this certificate was revoked. - """ - - @property - @abc.abstractmethod - def revocation_date_utc(self) -> datetime.datetime: - """ - Returns the date of when this certificate was revoked as a non-naive - UTC datetime. - """ - - @property - @abc.abstractmethod - def extensions(self) -> Extensions: - """ - Returns an Extensions object containing a list of Revoked extensions. - """ - - -# Runtime isinstance checks need this since the rust class is not a subclass. -RevokedCertificate.register(rust_x509.RevokedCertificate) - - -class _RawRevokedCertificate(RevokedCertificate): - def __init__( - self, - serial_number: int, - revocation_date: datetime.datetime, - extensions: Extensions, - ): - self._serial_number = serial_number - self._revocation_date = revocation_date - self._extensions = extensions - - @property - def serial_number(self) -> int: - return self._serial_number - - @property - def revocation_date(self) -> datetime.datetime: - warnings.warn( - "Properties that return a naïve datetime object have been " - "deprecated. Please switch to revocation_date_utc.", - utils.DeprecatedIn42, - stacklevel=2, - ) - return self._revocation_date - - @property - def revocation_date_utc(self) -> datetime.datetime: - return self._revocation_date.replace(tzinfo=datetime.timezone.utc) - - @property - def extensions(self) -> Extensions: - return self._extensions - - -CertificateRevocationList = rust_x509.CertificateRevocationList -CertificateSigningRequest = rust_x509.CertificateSigningRequest - - -load_pem_x509_certificate = rust_x509.load_pem_x509_certificate -load_der_x509_certificate = rust_x509.load_der_x509_certificate - -load_pem_x509_certificates = rust_x509.load_pem_x509_certificates - -load_pem_x509_csr = rust_x509.load_pem_x509_csr -load_der_x509_csr = rust_x509.load_der_x509_csr - -load_pem_x509_crl = rust_x509.load_pem_x509_crl -load_der_x509_crl = rust_x509.load_der_x509_crl - - -class CertificateSigningRequestBuilder: - def __init__( - self, - subject_name: Name | None = None, - extensions: list[Extension[ExtensionType]] = [], - attributes: list[tuple[ObjectIdentifier, bytes, int | None]] = [], - ): - """ - Creates an empty X.509 certificate request (v1). - """ - self._subject_name = subject_name - self._extensions = extensions - self._attributes = attributes - - def subject_name(self, name: Name) -> CertificateSigningRequestBuilder: - """ - Sets the certificate requestor's distinguished name. - """ - if not isinstance(name, Name): - raise TypeError("Expecting x509.Name object.") - if self._subject_name is not None: - raise ValueError("The subject name may only be set once.") - return CertificateSigningRequestBuilder( - name, self._extensions, self._attributes - ) - - def add_extension( - self, extval: ExtensionType, critical: bool - ) -> CertificateSigningRequestBuilder: - """ - Adds an X.509 extension to the certificate request. - """ - if not isinstance(extval, ExtensionType): - raise TypeError("extension must be an ExtensionType") - - extension = Extension(extval.oid, critical, extval) - _reject_duplicate_extension(extension, self._extensions) - - return CertificateSigningRequestBuilder( - self._subject_name, - [*self._extensions, extension], - self._attributes, - ) - - def add_attribute( - self, - oid: ObjectIdentifier, - value: bytes, - *, - _tag: _ASN1Type | None = None, - ) -> CertificateSigningRequestBuilder: - """ - Adds an X.509 attribute with an OID and associated value. - """ - if not isinstance(oid, ObjectIdentifier): - raise TypeError("oid must be an ObjectIdentifier") - - if not isinstance(value, bytes): - raise TypeError("value must be bytes") - - if _tag is not None and not isinstance(_tag, _ASN1Type): - raise TypeError("tag must be _ASN1Type") - - _reject_duplicate_attribute(oid, self._attributes) - - if _tag is not None: - tag = _tag.value - else: - tag = None - - return CertificateSigningRequestBuilder( - self._subject_name, - self._extensions, - [*self._attributes, (oid, value, tag)], - ) - - def sign( - self, - private_key: CertificateIssuerPrivateKeyTypes, - algorithm: _AllowedHashTypes | None, - backend: typing.Any = None, - *, - rsa_padding: padding.PSS | padding.PKCS1v15 | None = None, - ecdsa_deterministic: bool | None = None, - ) -> CertificateSigningRequest: - """ - Signs the request using the requestor's private key. - """ - if self._subject_name is None: - raise ValueError("A CertificateSigningRequest must have a subject") - - if rsa_padding is not None: - if not isinstance(rsa_padding, (padding.PSS, padding.PKCS1v15)): - raise TypeError("Padding must be PSS or PKCS1v15") - if not isinstance(private_key, rsa.RSAPrivateKey): - raise TypeError("Padding is only supported for RSA keys") - - if ecdsa_deterministic is not None: - if not isinstance(private_key, ec.EllipticCurvePrivateKey): - raise TypeError( - "Deterministic ECDSA is only supported for EC keys" - ) - - return rust_x509.create_x509_csr( - self, - private_key, - algorithm, - rsa_padding, - ecdsa_deterministic, - ) - - -class CertificateBuilder: - _extensions: list[Extension[ExtensionType]] - - def __init__( - self, - issuer_name: Name | None = None, - subject_name: Name | None = None, - public_key: CertificatePublicKeyTypes | None = None, - serial_number: int | None = None, - not_valid_before: datetime.datetime | None = None, - not_valid_after: datetime.datetime | None = None, - extensions: list[Extension[ExtensionType]] = [], - ) -> None: - self._version = Version.v3 - self._issuer_name = issuer_name - self._subject_name = subject_name - self._public_key = public_key - self._serial_number = serial_number - self._not_valid_before = not_valid_before - self._not_valid_after = not_valid_after - self._extensions = extensions - - def issuer_name(self, name: Name) -> CertificateBuilder: - """ - Sets the CA's distinguished name. - """ - if not isinstance(name, Name): - raise TypeError("Expecting x509.Name object.") - if self._issuer_name is not None: - raise ValueError("The issuer name may only be set once.") - return CertificateBuilder( - name, - self._subject_name, - self._public_key, - self._serial_number, - self._not_valid_before, - self._not_valid_after, - self._extensions, - ) - - def subject_name(self, name: Name) -> CertificateBuilder: - """ - Sets the requestor's distinguished name. - """ - if not isinstance(name, Name): - raise TypeError("Expecting x509.Name object.") - if self._subject_name is not None: - raise ValueError("The subject name may only be set once.") - return CertificateBuilder( - self._issuer_name, - name, - self._public_key, - self._serial_number, - self._not_valid_before, - self._not_valid_after, - self._extensions, - ) - - def public_key( - self, - key: CertificatePublicKeyTypes, - ) -> CertificateBuilder: - """ - Sets the requestor's public key (as found in the signing request). - """ - if not isinstance( - key, - ( - dsa.DSAPublicKey, - rsa.RSAPublicKey, - ec.EllipticCurvePublicKey, - ed25519.Ed25519PublicKey, - ed448.Ed448PublicKey, - x25519.X25519PublicKey, - x448.X448PublicKey, - ), - ): - raise TypeError( - "Expecting one of DSAPublicKey, RSAPublicKey," - " EllipticCurvePublicKey, Ed25519PublicKey," - " Ed448PublicKey, X25519PublicKey, or " - "X448PublicKey." - ) - if self._public_key is not None: - raise ValueError("The public key may only be set once.") - return CertificateBuilder( - self._issuer_name, - self._subject_name, - key, - self._serial_number, - self._not_valid_before, - self._not_valid_after, - self._extensions, - ) - - def serial_number(self, number: int) -> CertificateBuilder: - """ - Sets the certificate serial number. - """ - if not isinstance(number, int): - raise TypeError("Serial number must be of integral type.") - if self._serial_number is not None: - raise ValueError("The serial number may only be set once.") - if number <= 0: - raise ValueError("The serial number should be positive.") - - # ASN.1 integers are always signed, so most significant bit must be - # zero. - if number.bit_length() >= 160: # As defined in RFC 5280 - raise ValueError( - "The serial number should not be more than 159 bits." - ) - return CertificateBuilder( - self._issuer_name, - self._subject_name, - self._public_key, - number, - self._not_valid_before, - self._not_valid_after, - self._extensions, - ) - - def not_valid_before(self, time: datetime.datetime) -> CertificateBuilder: - """ - Sets the certificate activation time. - """ - if not isinstance(time, datetime.datetime): - raise TypeError("Expecting datetime object.") - if self._not_valid_before is not None: - raise ValueError("The not valid before may only be set once.") - time = _convert_to_naive_utc_time(time) - if time < _EARLIEST_UTC_TIME: - raise ValueError( - "The not valid before date must be on or after" - " 1950 January 1)." - ) - if self._not_valid_after is not None and time > self._not_valid_after: - raise ValueError( - "The not valid before date must be before the not valid after " - "date." - ) - return CertificateBuilder( - self._issuer_name, - self._subject_name, - self._public_key, - self._serial_number, - time, - self._not_valid_after, - self._extensions, - ) - - def not_valid_after(self, time: datetime.datetime) -> CertificateBuilder: - """ - Sets the certificate expiration time. - """ - if not isinstance(time, datetime.datetime): - raise TypeError("Expecting datetime object.") - if self._not_valid_after is not None: - raise ValueError("The not valid after may only be set once.") - time = _convert_to_naive_utc_time(time) - if time < _EARLIEST_UTC_TIME: - raise ValueError( - "The not valid after date must be on or after 1950 January 1." - ) - if ( - self._not_valid_before is not None - and time < self._not_valid_before - ): - raise ValueError( - "The not valid after date must be after the not valid before " - "date." - ) - return CertificateBuilder( - self._issuer_name, - self._subject_name, - self._public_key, - self._serial_number, - self._not_valid_before, - time, - self._extensions, - ) - - def add_extension( - self, extval: ExtensionType, critical: bool - ) -> CertificateBuilder: - """ - Adds an X.509 extension to the certificate. - """ - if not isinstance(extval, ExtensionType): - raise TypeError("extension must be an ExtensionType") - - extension = Extension(extval.oid, critical, extval) - _reject_duplicate_extension(extension, self._extensions) - - return CertificateBuilder( - self._issuer_name, - self._subject_name, - self._public_key, - self._serial_number, - self._not_valid_before, - self._not_valid_after, - [*self._extensions, extension], - ) - - def sign( - self, - private_key: CertificateIssuerPrivateKeyTypes, - algorithm: _AllowedHashTypes | None, - backend: typing.Any = None, - *, - rsa_padding: padding.PSS | padding.PKCS1v15 | None = None, - ecdsa_deterministic: bool | None = None, - ) -> Certificate: - """ - Signs the certificate using the CA's private key. - """ - if self._subject_name is None: - raise ValueError("A certificate must have a subject name") - - if self._issuer_name is None: - raise ValueError("A certificate must have an issuer name") - - if self._serial_number is None: - raise ValueError("A certificate must have a serial number") - - if self._not_valid_before is None: - raise ValueError("A certificate must have a not valid before time") - - if self._not_valid_after is None: - raise ValueError("A certificate must have a not valid after time") - - if self._public_key is None: - raise ValueError("A certificate must have a public key") - - if rsa_padding is not None: - if not isinstance(rsa_padding, (padding.PSS, padding.PKCS1v15)): - raise TypeError("Padding must be PSS or PKCS1v15") - if not isinstance(private_key, rsa.RSAPrivateKey): - raise TypeError("Padding is only supported for RSA keys") - - if ecdsa_deterministic is not None: - if not isinstance(private_key, ec.EllipticCurvePrivateKey): - raise TypeError( - "Deterministic ECDSA is only supported for EC keys" - ) - - return rust_x509.create_x509_certificate( - self, - private_key, - algorithm, - rsa_padding, - ecdsa_deterministic, - ) - - -class CertificateRevocationListBuilder: - _extensions: list[Extension[ExtensionType]] - _revoked_certificates: list[RevokedCertificate] - - def __init__( - self, - issuer_name: Name | None = None, - last_update: datetime.datetime | None = None, - next_update: datetime.datetime | None = None, - extensions: list[Extension[ExtensionType]] = [], - revoked_certificates: list[RevokedCertificate] = [], - ): - self._issuer_name = issuer_name - self._last_update = last_update - self._next_update = next_update - self._extensions = extensions - self._revoked_certificates = revoked_certificates - - def issuer_name( - self, issuer_name: Name - ) -> CertificateRevocationListBuilder: - if not isinstance(issuer_name, Name): - raise TypeError("Expecting x509.Name object.") - if self._issuer_name is not None: - raise ValueError("The issuer name may only be set once.") - return CertificateRevocationListBuilder( - issuer_name, - self._last_update, - self._next_update, - self._extensions, - self._revoked_certificates, - ) - - def last_update( - self, last_update: datetime.datetime - ) -> CertificateRevocationListBuilder: - if not isinstance(last_update, datetime.datetime): - raise TypeError("Expecting datetime object.") - if self._last_update is not None: - raise ValueError("Last update may only be set once.") - last_update = _convert_to_naive_utc_time(last_update) - if last_update < _EARLIEST_UTC_TIME: - raise ValueError( - "The last update date must be on or after 1950 January 1." - ) - if self._next_update is not None and last_update > self._next_update: - raise ValueError( - "The last update date must be before the next update date." - ) - return CertificateRevocationListBuilder( - self._issuer_name, - last_update, - self._next_update, - self._extensions, - self._revoked_certificates, - ) - - def next_update( - self, next_update: datetime.datetime - ) -> CertificateRevocationListBuilder: - if not isinstance(next_update, datetime.datetime): - raise TypeError("Expecting datetime object.") - if self._next_update is not None: - raise ValueError("Last update may only be set once.") - next_update = _convert_to_naive_utc_time(next_update) - if next_update < _EARLIEST_UTC_TIME: - raise ValueError( - "The last update date must be on or after 1950 January 1." - ) - if self._last_update is not None and next_update < self._last_update: - raise ValueError( - "The next update date must be after the last update date." - ) - return CertificateRevocationListBuilder( - self._issuer_name, - self._last_update, - next_update, - self._extensions, - self._revoked_certificates, - ) - - def add_extension( - self, extval: ExtensionType, critical: bool - ) -> CertificateRevocationListBuilder: - """ - Adds an X.509 extension to the certificate revocation list. - """ - if not isinstance(extval, ExtensionType): - raise TypeError("extension must be an ExtensionType") - - extension = Extension(extval.oid, critical, extval) - _reject_duplicate_extension(extension, self._extensions) - return CertificateRevocationListBuilder( - self._issuer_name, - self._last_update, - self._next_update, - [*self._extensions, extension], - self._revoked_certificates, - ) - - def add_revoked_certificate( - self, revoked_certificate: RevokedCertificate - ) -> CertificateRevocationListBuilder: - """ - Adds a revoked certificate to the CRL. - """ - if not isinstance(revoked_certificate, RevokedCertificate): - raise TypeError("Must be an instance of RevokedCertificate") - - return CertificateRevocationListBuilder( - self._issuer_name, - self._last_update, - self._next_update, - self._extensions, - [*self._revoked_certificates, revoked_certificate], - ) - - def sign( - self, - private_key: CertificateIssuerPrivateKeyTypes, - algorithm: _AllowedHashTypes | None, - backend: typing.Any = None, - *, - rsa_padding: padding.PSS | padding.PKCS1v15 | None = None, - ecdsa_deterministic: bool | None = None, - ) -> CertificateRevocationList: - if self._issuer_name is None: - raise ValueError("A CRL must have an issuer name") - - if self._last_update is None: - raise ValueError("A CRL must have a last update time") - - if self._next_update is None: - raise ValueError("A CRL must have a next update time") - - if rsa_padding is not None: - if not isinstance(rsa_padding, (padding.PSS, padding.PKCS1v15)): - raise TypeError("Padding must be PSS or PKCS1v15") - if not isinstance(private_key, rsa.RSAPrivateKey): - raise TypeError("Padding is only supported for RSA keys") - - if ecdsa_deterministic is not None: - if not isinstance(private_key, ec.EllipticCurvePrivateKey): - raise TypeError( - "Deterministic ECDSA is only supported for EC keys" - ) - - return rust_x509.create_x509_crl( - self, - private_key, - algorithm, - rsa_padding, - ecdsa_deterministic, - ) - - -class RevokedCertificateBuilder: - def __init__( - self, - serial_number: int | None = None, - revocation_date: datetime.datetime | None = None, - extensions: list[Extension[ExtensionType]] = [], - ): - self._serial_number = serial_number - self._revocation_date = revocation_date - self._extensions = extensions - - def serial_number(self, number: int) -> RevokedCertificateBuilder: - if not isinstance(number, int): - raise TypeError("Serial number must be of integral type.") - if self._serial_number is not None: - raise ValueError("The serial number may only be set once.") - if number <= 0: - raise ValueError("The serial number should be positive") - - # ASN.1 integers are always signed, so most significant bit must be - # zero. - if number.bit_length() >= 160: # As defined in RFC 5280 - raise ValueError( - "The serial number should not be more than 159 bits." - ) - return RevokedCertificateBuilder( - number, self._revocation_date, self._extensions - ) - - def revocation_date( - self, time: datetime.datetime - ) -> RevokedCertificateBuilder: - if not isinstance(time, datetime.datetime): - raise TypeError("Expecting datetime object.") - if self._revocation_date is not None: - raise ValueError("The revocation date may only be set once.") - time = _convert_to_naive_utc_time(time) - if time < _EARLIEST_UTC_TIME: - raise ValueError( - "The revocation date must be on or after 1950 January 1." - ) - return RevokedCertificateBuilder( - self._serial_number, time, self._extensions - ) - - def add_extension( - self, extval: ExtensionType, critical: bool - ) -> RevokedCertificateBuilder: - if not isinstance(extval, ExtensionType): - raise TypeError("extension must be an ExtensionType") - - extension = Extension(extval.oid, critical, extval) - _reject_duplicate_extension(extension, self._extensions) - return RevokedCertificateBuilder( - self._serial_number, - self._revocation_date, - [*self._extensions, extension], - ) - - def build(self, backend: typing.Any = None) -> RevokedCertificate: - if self._serial_number is None: - raise ValueError("A revoked certificate must have a serial number") - if self._revocation_date is None: - raise ValueError( - "A revoked certificate must have a revocation date" - ) - return _RawRevokedCertificate( - self._serial_number, - self._revocation_date, - Extensions(self._extensions), - ) - - -def random_serial_number() -> int: - return int.from_bytes(os.urandom(20), "big") >> 1 diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/certificate_transparency.py b/venv/lib/python3.10/site-packages/cryptography/x509/certificate_transparency.py deleted file mode 100644 index fb66cc6049528dd012f07bfcef765e284d0062d2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/x509/certificate_transparency.py +++ /dev/null @@ -1,35 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography import utils -from cryptography.hazmat.bindings._rust import x509 as rust_x509 - - -class LogEntryType(utils.Enum): - X509_CERTIFICATE = 0 - PRE_CERTIFICATE = 1 - - -class Version(utils.Enum): - v1 = 0 - - -class SignatureAlgorithm(utils.Enum): - """ - Signature algorithms that are valid for SCTs. - - These are exactly the same as SignatureAlgorithm in RFC 5246 (TLS 1.2). - - See: - """ - - ANONYMOUS = 0 - RSA = 1 - DSA = 2 - ECDSA = 3 - - -SignedCertificateTimestamp = rust_x509.Sct diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/extensions.py b/venv/lib/python3.10/site-packages/cryptography/x509/extensions.py deleted file mode 100644 index dfa472d38061584e086a60820fac8555bc4fa808..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/x509/extensions.py +++ /dev/null @@ -1,2528 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc -import datetime -import hashlib -import ipaddress -import typing -from collections.abc import Iterable, Iterator - -from cryptography import utils -from cryptography.hazmat.bindings._rust import asn1 -from cryptography.hazmat.bindings._rust import x509 as rust_x509 -from cryptography.hazmat.primitives import constant_time, serialization -from cryptography.hazmat.primitives.asymmetric.ec import EllipticCurvePublicKey -from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey -from cryptography.hazmat.primitives.asymmetric.types import ( - CertificateIssuerPublicKeyTypes, - CertificatePublicKeyTypes, -) -from cryptography.x509.certificate_transparency import ( - SignedCertificateTimestamp, -) -from cryptography.x509.general_name import ( - DirectoryName, - DNSName, - GeneralName, - IPAddress, - OtherName, - RegisteredID, - RFC822Name, - UniformResourceIdentifier, - _IPAddressTypes, -) -from cryptography.x509.name import Name, RelativeDistinguishedName -from cryptography.x509.oid import ( - CRLEntryExtensionOID, - ExtensionOID, - ObjectIdentifier, - OCSPExtensionOID, -) - -ExtensionTypeVar = typing.TypeVar( - "ExtensionTypeVar", bound="ExtensionType", covariant=True -) - - -def _key_identifier_from_public_key( - public_key: CertificatePublicKeyTypes, -) -> bytes: - if isinstance(public_key, RSAPublicKey): - data = public_key.public_bytes( - serialization.Encoding.DER, - serialization.PublicFormat.PKCS1, - ) - elif isinstance(public_key, EllipticCurvePublicKey): - data = public_key.public_bytes( - serialization.Encoding.X962, - serialization.PublicFormat.UncompressedPoint, - ) - else: - # This is a very slow way to do this. - serialized = public_key.public_bytes( - serialization.Encoding.DER, - serialization.PublicFormat.SubjectPublicKeyInfo, - ) - data = asn1.parse_spki_for_data(serialized) - - return hashlib.sha1(data).digest() - - -def _make_sequence_methods(field_name: str): - def len_method(self) -> int: - return len(getattr(self, field_name)) - - def iter_method(self): - return iter(getattr(self, field_name)) - - def getitem_method(self, idx): - return getattr(self, field_name)[idx] - - return len_method, iter_method, getitem_method - - -class DuplicateExtension(Exception): - def __init__(self, msg: str, oid: ObjectIdentifier) -> None: - super().__init__(msg) - self.oid = oid - - -class ExtensionNotFound(Exception): - def __init__(self, msg: str, oid: ObjectIdentifier) -> None: - super().__init__(msg) - self.oid = oid - - -class ExtensionType(metaclass=abc.ABCMeta): - oid: typing.ClassVar[ObjectIdentifier] - - def public_bytes(self) -> bytes: - """ - Serializes the extension type to DER. - """ - raise NotImplementedError( - f"public_bytes is not implemented for extension type {self!r}" - ) - - -class Extensions: - def __init__(self, extensions: Iterable[Extension[ExtensionType]]) -> None: - self._extensions = list(extensions) - - def get_extension_for_oid( - self, oid: ObjectIdentifier - ) -> Extension[ExtensionType]: - for ext in self: - if ext.oid == oid: - return ext - - raise ExtensionNotFound(f"No {oid} extension was found", oid) - - def get_extension_for_class( - self, extclass: type[ExtensionTypeVar] - ) -> Extension[ExtensionTypeVar]: - if extclass is UnrecognizedExtension: - raise TypeError( - "UnrecognizedExtension can't be used with " - "get_extension_for_class because more than one instance of the" - " class may be present." - ) - - for ext in self: - if isinstance(ext.value, extclass): - return ext - - raise ExtensionNotFound( - f"No {extclass} extension was found", extclass.oid - ) - - __len__, __iter__, __getitem__ = _make_sequence_methods("_extensions") - - def __repr__(self) -> str: - return f"" - - -class CRLNumber(ExtensionType): - oid = ExtensionOID.CRL_NUMBER - - def __init__(self, crl_number: int) -> None: - if not isinstance(crl_number, int): - raise TypeError("crl_number must be an integer") - - self._crl_number = crl_number - - def __eq__(self, other: object) -> bool: - if not isinstance(other, CRLNumber): - return NotImplemented - - return self.crl_number == other.crl_number - - def __hash__(self) -> int: - return hash(self.crl_number) - - def __repr__(self) -> str: - return f"" - - @property - def crl_number(self) -> int: - return self._crl_number - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class AuthorityKeyIdentifier(ExtensionType): - oid = ExtensionOID.AUTHORITY_KEY_IDENTIFIER - - def __init__( - self, - key_identifier: bytes | None, - authority_cert_issuer: Iterable[GeneralName] | None, - authority_cert_serial_number: int | None, - ) -> None: - if (authority_cert_issuer is None) != ( - authority_cert_serial_number is None - ): - raise ValueError( - "authority_cert_issuer and authority_cert_serial_number " - "must both be present or both None" - ) - - if authority_cert_issuer is not None: - authority_cert_issuer = list(authority_cert_issuer) - if not all( - isinstance(x, GeneralName) for x in authority_cert_issuer - ): - raise TypeError( - "authority_cert_issuer must be a list of GeneralName " - "objects" - ) - - if authority_cert_serial_number is not None and not isinstance( - authority_cert_serial_number, int - ): - raise TypeError("authority_cert_serial_number must be an integer") - - self._key_identifier = key_identifier - self._authority_cert_issuer = authority_cert_issuer - self._authority_cert_serial_number = authority_cert_serial_number - - # This takes a subset of CertificatePublicKeyTypes because an issuer - # cannot have an X25519/X448 key. This introduces some unfortunate - # asymmetry that requires typing users to explicitly - # narrow their type, but we should make this accurate and not just - # convenient. - @classmethod - def from_issuer_public_key( - cls, public_key: CertificateIssuerPublicKeyTypes - ) -> AuthorityKeyIdentifier: - digest = _key_identifier_from_public_key(public_key) - return cls( - key_identifier=digest, - authority_cert_issuer=None, - authority_cert_serial_number=None, - ) - - @classmethod - def from_issuer_subject_key_identifier( - cls, ski: SubjectKeyIdentifier - ) -> AuthorityKeyIdentifier: - return cls( - key_identifier=ski.digest, - authority_cert_issuer=None, - authority_cert_serial_number=None, - ) - - def __repr__(self) -> str: - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, AuthorityKeyIdentifier): - return NotImplemented - - return ( - self.key_identifier == other.key_identifier - and self.authority_cert_issuer == other.authority_cert_issuer - and self.authority_cert_serial_number - == other.authority_cert_serial_number - ) - - def __hash__(self) -> int: - if self.authority_cert_issuer is None: - aci = None - else: - aci = tuple(self.authority_cert_issuer) - return hash( - (self.key_identifier, aci, self.authority_cert_serial_number) - ) - - @property - def key_identifier(self) -> bytes | None: - return self._key_identifier - - @property - def authority_cert_issuer( - self, - ) -> list[GeneralName] | None: - return self._authority_cert_issuer - - @property - def authority_cert_serial_number(self) -> int | None: - return self._authority_cert_serial_number - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class SubjectKeyIdentifier(ExtensionType): - oid = ExtensionOID.SUBJECT_KEY_IDENTIFIER - - def __init__(self, digest: bytes) -> None: - self._digest = digest - - @classmethod - def from_public_key( - cls, public_key: CertificatePublicKeyTypes - ) -> SubjectKeyIdentifier: - return cls(_key_identifier_from_public_key(public_key)) - - @property - def digest(self) -> bytes: - return self._digest - - @property - def key_identifier(self) -> bytes: - return self._digest - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, SubjectKeyIdentifier): - return NotImplemented - - return constant_time.bytes_eq(self.digest, other.digest) - - def __hash__(self) -> int: - return hash(self.digest) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class AuthorityInformationAccess(ExtensionType): - oid = ExtensionOID.AUTHORITY_INFORMATION_ACCESS - - def __init__(self, descriptions: Iterable[AccessDescription]) -> None: - descriptions = list(descriptions) - if not all(isinstance(x, AccessDescription) for x in descriptions): - raise TypeError( - "Every item in the descriptions list must be an " - "AccessDescription" - ) - - self._descriptions = descriptions - - __len__, __iter__, __getitem__ = _make_sequence_methods("_descriptions") - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, AuthorityInformationAccess): - return NotImplemented - - return self._descriptions == other._descriptions - - def __hash__(self) -> int: - return hash(tuple(self._descriptions)) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class SubjectInformationAccess(ExtensionType): - oid = ExtensionOID.SUBJECT_INFORMATION_ACCESS - - def __init__(self, descriptions: Iterable[AccessDescription]) -> None: - descriptions = list(descriptions) - if not all(isinstance(x, AccessDescription) for x in descriptions): - raise TypeError( - "Every item in the descriptions list must be an " - "AccessDescription" - ) - - self._descriptions = descriptions - - __len__, __iter__, __getitem__ = _make_sequence_methods("_descriptions") - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, SubjectInformationAccess): - return NotImplemented - - return self._descriptions == other._descriptions - - def __hash__(self) -> int: - return hash(tuple(self._descriptions)) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class AccessDescription: - def __init__( - self, access_method: ObjectIdentifier, access_location: GeneralName - ) -> None: - if not isinstance(access_method, ObjectIdentifier): - raise TypeError("access_method must be an ObjectIdentifier") - - if not isinstance(access_location, GeneralName): - raise TypeError("access_location must be a GeneralName") - - self._access_method = access_method - self._access_location = access_location - - def __repr__(self) -> str: - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, AccessDescription): - return NotImplemented - - return ( - self.access_method == other.access_method - and self.access_location == other.access_location - ) - - def __hash__(self) -> int: - return hash((self.access_method, self.access_location)) - - @property - def access_method(self) -> ObjectIdentifier: - return self._access_method - - @property - def access_location(self) -> GeneralName: - return self._access_location - - -class BasicConstraints(ExtensionType): - oid = ExtensionOID.BASIC_CONSTRAINTS - - def __init__(self, ca: bool, path_length: int | None) -> None: - if not isinstance(ca, bool): - raise TypeError("ca must be a boolean value") - - if path_length is not None and not ca: - raise ValueError("path_length must be None when ca is False") - - if path_length is not None and ( - not isinstance(path_length, int) or path_length < 0 - ): - raise TypeError( - "path_length must be a non-negative integer or None" - ) - - self._ca = ca - self._path_length = path_length - - @property - def ca(self) -> bool: - return self._ca - - @property - def path_length(self) -> int | None: - return self._path_length - - def __repr__(self) -> str: - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, BasicConstraints): - return NotImplemented - - return self.ca == other.ca and self.path_length == other.path_length - - def __hash__(self) -> int: - return hash((self.ca, self.path_length)) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class DeltaCRLIndicator(ExtensionType): - oid = ExtensionOID.DELTA_CRL_INDICATOR - - def __init__(self, crl_number: int) -> None: - if not isinstance(crl_number, int): - raise TypeError("crl_number must be an integer") - - self._crl_number = crl_number - - @property - def crl_number(self) -> int: - return self._crl_number - - def __eq__(self, other: object) -> bool: - if not isinstance(other, DeltaCRLIndicator): - return NotImplemented - - return self.crl_number == other.crl_number - - def __hash__(self) -> int: - return hash(self.crl_number) - - def __repr__(self) -> str: - return f"" - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class CRLDistributionPoints(ExtensionType): - oid = ExtensionOID.CRL_DISTRIBUTION_POINTS - - def __init__( - self, distribution_points: Iterable[DistributionPoint] - ) -> None: - distribution_points = list(distribution_points) - if not all( - isinstance(x, DistributionPoint) for x in distribution_points - ): - raise TypeError( - "distribution_points must be a list of DistributionPoint " - "objects" - ) - - self._distribution_points = distribution_points - - __len__, __iter__, __getitem__ = _make_sequence_methods( - "_distribution_points" - ) - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, CRLDistributionPoints): - return NotImplemented - - return self._distribution_points == other._distribution_points - - def __hash__(self) -> int: - return hash(tuple(self._distribution_points)) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class FreshestCRL(ExtensionType): - oid = ExtensionOID.FRESHEST_CRL - - def __init__( - self, distribution_points: Iterable[DistributionPoint] - ) -> None: - distribution_points = list(distribution_points) - if not all( - isinstance(x, DistributionPoint) for x in distribution_points - ): - raise TypeError( - "distribution_points must be a list of DistributionPoint " - "objects" - ) - - self._distribution_points = distribution_points - - __len__, __iter__, __getitem__ = _make_sequence_methods( - "_distribution_points" - ) - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, FreshestCRL): - return NotImplemented - - return self._distribution_points == other._distribution_points - - def __hash__(self) -> int: - return hash(tuple(self._distribution_points)) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class DistributionPoint: - def __init__( - self, - full_name: Iterable[GeneralName] | None, - relative_name: RelativeDistinguishedName | None, - reasons: frozenset[ReasonFlags] | None, - crl_issuer: Iterable[GeneralName] | None, - ) -> None: - if full_name and relative_name: - raise ValueError( - "You cannot provide both full_name and relative_name, at " - "least one must be None." - ) - if not full_name and not relative_name and not crl_issuer: - raise ValueError( - "Either full_name, relative_name or crl_issuer must be " - "provided." - ) - - if full_name is not None: - full_name = list(full_name) - if not all(isinstance(x, GeneralName) for x in full_name): - raise TypeError( - "full_name must be a list of GeneralName objects" - ) - - if relative_name: - if not isinstance(relative_name, RelativeDistinguishedName): - raise TypeError( - "relative_name must be a RelativeDistinguishedName" - ) - - if crl_issuer is not None: - crl_issuer = list(crl_issuer) - if not all(isinstance(x, GeneralName) for x in crl_issuer): - raise TypeError( - "crl_issuer must be None or a list of general names" - ) - - if reasons and ( - not isinstance(reasons, frozenset) - or not all(isinstance(x, ReasonFlags) for x in reasons) - ): - raise TypeError("reasons must be None or frozenset of ReasonFlags") - - if reasons and ( - ReasonFlags.unspecified in reasons - or ReasonFlags.remove_from_crl in reasons - ): - raise ValueError( - "unspecified and remove_from_crl are not valid reasons in a " - "DistributionPoint" - ) - - self._full_name = full_name - self._relative_name = relative_name - self._reasons = reasons - self._crl_issuer = crl_issuer - - def __repr__(self) -> str: - return ( - "".format(self) - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, DistributionPoint): - return NotImplemented - - return ( - self.full_name == other.full_name - and self.relative_name == other.relative_name - and self.reasons == other.reasons - and self.crl_issuer == other.crl_issuer - ) - - def __hash__(self) -> int: - if self.full_name is not None: - fn: tuple[GeneralName, ...] | None = tuple(self.full_name) - else: - fn = None - - if self.crl_issuer is not None: - crl_issuer: tuple[GeneralName, ...] | None = tuple(self.crl_issuer) - else: - crl_issuer = None - - return hash((fn, self.relative_name, self.reasons, crl_issuer)) - - @property - def full_name(self) -> list[GeneralName] | None: - return self._full_name - - @property - def relative_name(self) -> RelativeDistinguishedName | None: - return self._relative_name - - @property - def reasons(self) -> frozenset[ReasonFlags] | None: - return self._reasons - - @property - def crl_issuer(self) -> list[GeneralName] | None: - return self._crl_issuer - - -class ReasonFlags(utils.Enum): - unspecified = "unspecified" - key_compromise = "keyCompromise" - ca_compromise = "cACompromise" - affiliation_changed = "affiliationChanged" - superseded = "superseded" - cessation_of_operation = "cessationOfOperation" - certificate_hold = "certificateHold" - privilege_withdrawn = "privilegeWithdrawn" - aa_compromise = "aACompromise" - remove_from_crl = "removeFromCRL" - - -# These are distribution point bit string mappings. Not to be confused with -# CRLReason reason flags bit string mappings. -# ReasonFlags ::= BIT STRING { -# unused (0), -# keyCompromise (1), -# cACompromise (2), -# affiliationChanged (3), -# superseded (4), -# cessationOfOperation (5), -# certificateHold (6), -# privilegeWithdrawn (7), -# aACompromise (8) } -_REASON_BIT_MAPPING = { - 1: ReasonFlags.key_compromise, - 2: ReasonFlags.ca_compromise, - 3: ReasonFlags.affiliation_changed, - 4: ReasonFlags.superseded, - 5: ReasonFlags.cessation_of_operation, - 6: ReasonFlags.certificate_hold, - 7: ReasonFlags.privilege_withdrawn, - 8: ReasonFlags.aa_compromise, -} - -_CRLREASONFLAGS = { - ReasonFlags.key_compromise: 1, - ReasonFlags.ca_compromise: 2, - ReasonFlags.affiliation_changed: 3, - ReasonFlags.superseded: 4, - ReasonFlags.cessation_of_operation: 5, - ReasonFlags.certificate_hold: 6, - ReasonFlags.privilege_withdrawn: 7, - ReasonFlags.aa_compromise: 8, -} - -# CRLReason ::= ENUMERATED { -# unspecified (0), -# keyCompromise (1), -# cACompromise (2), -# affiliationChanged (3), -# superseded (4), -# cessationOfOperation (5), -# certificateHold (6), -# -- value 7 is not used -# removeFromCRL (8), -# privilegeWithdrawn (9), -# aACompromise (10) } -_CRL_ENTRY_REASON_ENUM_TO_CODE = { - ReasonFlags.unspecified: 0, - ReasonFlags.key_compromise: 1, - ReasonFlags.ca_compromise: 2, - ReasonFlags.affiliation_changed: 3, - ReasonFlags.superseded: 4, - ReasonFlags.cessation_of_operation: 5, - ReasonFlags.certificate_hold: 6, - ReasonFlags.remove_from_crl: 8, - ReasonFlags.privilege_withdrawn: 9, - ReasonFlags.aa_compromise: 10, -} - - -class PolicyConstraints(ExtensionType): - oid = ExtensionOID.POLICY_CONSTRAINTS - - def __init__( - self, - require_explicit_policy: int | None, - inhibit_policy_mapping: int | None, - ) -> None: - if require_explicit_policy is not None and not isinstance( - require_explicit_policy, int - ): - raise TypeError( - "require_explicit_policy must be a non-negative integer or " - "None" - ) - - if inhibit_policy_mapping is not None and not isinstance( - inhibit_policy_mapping, int - ): - raise TypeError( - "inhibit_policy_mapping must be a non-negative integer or None" - ) - - if inhibit_policy_mapping is None and require_explicit_policy is None: - raise ValueError( - "At least one of require_explicit_policy and " - "inhibit_policy_mapping must not be None" - ) - - self._require_explicit_policy = require_explicit_policy - self._inhibit_policy_mapping = inhibit_policy_mapping - - def __repr__(self) -> str: - return ( - "".format(self) - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, PolicyConstraints): - return NotImplemented - - return ( - self.require_explicit_policy == other.require_explicit_policy - and self.inhibit_policy_mapping == other.inhibit_policy_mapping - ) - - def __hash__(self) -> int: - return hash( - (self.require_explicit_policy, self.inhibit_policy_mapping) - ) - - @property - def require_explicit_policy(self) -> int | None: - return self._require_explicit_policy - - @property - def inhibit_policy_mapping(self) -> int | None: - return self._inhibit_policy_mapping - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class CertificatePolicies(ExtensionType): - oid = ExtensionOID.CERTIFICATE_POLICIES - - def __init__(self, policies: Iterable[PolicyInformation]) -> None: - policies = list(policies) - if not all(isinstance(x, PolicyInformation) for x in policies): - raise TypeError( - "Every item in the policies list must be a PolicyInformation" - ) - - self._policies = policies - - __len__, __iter__, __getitem__ = _make_sequence_methods("_policies") - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, CertificatePolicies): - return NotImplemented - - return self._policies == other._policies - - def __hash__(self) -> int: - return hash(tuple(self._policies)) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class PolicyInformation: - def __init__( - self, - policy_identifier: ObjectIdentifier, - policy_qualifiers: Iterable[str | UserNotice] | None, - ) -> None: - if not isinstance(policy_identifier, ObjectIdentifier): - raise TypeError("policy_identifier must be an ObjectIdentifier") - - self._policy_identifier = policy_identifier - - if policy_qualifiers is not None: - policy_qualifiers = list(policy_qualifiers) - if not all( - isinstance(x, (str, UserNotice)) for x in policy_qualifiers - ): - raise TypeError( - "policy_qualifiers must be a list of strings and/or " - "UserNotice objects or None" - ) - - self._policy_qualifiers = policy_qualifiers - - def __repr__(self) -> str: - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, PolicyInformation): - return NotImplemented - - return ( - self.policy_identifier == other.policy_identifier - and self.policy_qualifiers == other.policy_qualifiers - ) - - def __hash__(self) -> int: - if self.policy_qualifiers is not None: - pq = tuple(self.policy_qualifiers) - else: - pq = None - - return hash((self.policy_identifier, pq)) - - @property - def policy_identifier(self) -> ObjectIdentifier: - return self._policy_identifier - - @property - def policy_qualifiers( - self, - ) -> list[str | UserNotice] | None: - return self._policy_qualifiers - - -class UserNotice: - def __init__( - self, - notice_reference: NoticeReference | None, - explicit_text: str | None, - ) -> None: - if notice_reference and not isinstance( - notice_reference, NoticeReference - ): - raise TypeError( - "notice_reference must be None or a NoticeReference" - ) - - self._notice_reference = notice_reference - self._explicit_text = explicit_text - - def __repr__(self) -> str: - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, UserNotice): - return NotImplemented - - return ( - self.notice_reference == other.notice_reference - and self.explicit_text == other.explicit_text - ) - - def __hash__(self) -> int: - return hash((self.notice_reference, self.explicit_text)) - - @property - def notice_reference(self) -> NoticeReference | None: - return self._notice_reference - - @property - def explicit_text(self) -> str | None: - return self._explicit_text - - -class NoticeReference: - def __init__( - self, - organization: str | None, - notice_numbers: Iterable[int], - ) -> None: - self._organization = organization - notice_numbers = list(notice_numbers) - if not all(isinstance(x, int) for x in notice_numbers): - raise TypeError("notice_numbers must be a list of integers") - - self._notice_numbers = notice_numbers - - def __repr__(self) -> str: - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, NoticeReference): - return NotImplemented - - return ( - self.organization == other.organization - and self.notice_numbers == other.notice_numbers - ) - - def __hash__(self) -> int: - return hash((self.organization, tuple(self.notice_numbers))) - - @property - def organization(self) -> str | None: - return self._organization - - @property - def notice_numbers(self) -> list[int]: - return self._notice_numbers - - -class ExtendedKeyUsage(ExtensionType): - oid = ExtensionOID.EXTENDED_KEY_USAGE - - def __init__(self, usages: Iterable[ObjectIdentifier]) -> None: - usages = list(usages) - if not all(isinstance(x, ObjectIdentifier) for x in usages): - raise TypeError( - "Every item in the usages list must be an ObjectIdentifier" - ) - - self._usages = usages - - __len__, __iter__, __getitem__ = _make_sequence_methods("_usages") - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, ExtendedKeyUsage): - return NotImplemented - - return self._usages == other._usages - - def __hash__(self) -> int: - return hash(tuple(self._usages)) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class OCSPNoCheck(ExtensionType): - oid = ExtensionOID.OCSP_NO_CHECK - - def __eq__(self, other: object) -> bool: - if not isinstance(other, OCSPNoCheck): - return NotImplemented - - return True - - def __hash__(self) -> int: - return hash(OCSPNoCheck) - - def __repr__(self) -> str: - return "" - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class PrecertPoison(ExtensionType): - oid = ExtensionOID.PRECERT_POISON - - def __eq__(self, other: object) -> bool: - if not isinstance(other, PrecertPoison): - return NotImplemented - - return True - - def __hash__(self) -> int: - return hash(PrecertPoison) - - def __repr__(self) -> str: - return "" - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class TLSFeature(ExtensionType): - oid = ExtensionOID.TLS_FEATURE - - def __init__(self, features: Iterable[TLSFeatureType]) -> None: - features = list(features) - if ( - not all(isinstance(x, TLSFeatureType) for x in features) - or len(features) == 0 - ): - raise TypeError( - "features must be a list of elements from the TLSFeatureType " - "enum" - ) - - self._features = features - - __len__, __iter__, __getitem__ = _make_sequence_methods("_features") - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, TLSFeature): - return NotImplemented - - return self._features == other._features - - def __hash__(self) -> int: - return hash(tuple(self._features)) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class TLSFeatureType(utils.Enum): - # status_request is defined in RFC 6066 and is used for what is commonly - # called OCSP Must-Staple when present in the TLS Feature extension in an - # X.509 certificate. - status_request = 5 - # status_request_v2 is defined in RFC 6961 and allows multiple OCSP - # responses to be provided. It is not currently in use by clients or - # servers. - status_request_v2 = 17 - - -_TLS_FEATURE_TYPE_TO_ENUM = {x.value: x for x in TLSFeatureType} - - -class InhibitAnyPolicy(ExtensionType): - oid = ExtensionOID.INHIBIT_ANY_POLICY - - def __init__(self, skip_certs: int) -> None: - if not isinstance(skip_certs, int): - raise TypeError("skip_certs must be an integer") - - if skip_certs < 0: - raise ValueError("skip_certs must be a non-negative integer") - - self._skip_certs = skip_certs - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, InhibitAnyPolicy): - return NotImplemented - - return self.skip_certs == other.skip_certs - - def __hash__(self) -> int: - return hash(self.skip_certs) - - @property - def skip_certs(self) -> int: - return self._skip_certs - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class KeyUsage(ExtensionType): - oid = ExtensionOID.KEY_USAGE - - def __init__( - self, - digital_signature: bool, - content_commitment: bool, - key_encipherment: bool, - data_encipherment: bool, - key_agreement: bool, - key_cert_sign: bool, - crl_sign: bool, - encipher_only: bool, - decipher_only: bool, - ) -> None: - if not key_agreement and (encipher_only or decipher_only): - raise ValueError( - "encipher_only and decipher_only can only be true when " - "key_agreement is true" - ) - - self._digital_signature = digital_signature - self._content_commitment = content_commitment - self._key_encipherment = key_encipherment - self._data_encipherment = data_encipherment - self._key_agreement = key_agreement - self._key_cert_sign = key_cert_sign - self._crl_sign = crl_sign - self._encipher_only = encipher_only - self._decipher_only = decipher_only - - @property - def digital_signature(self) -> bool: - return self._digital_signature - - @property - def content_commitment(self) -> bool: - return self._content_commitment - - @property - def key_encipherment(self) -> bool: - return self._key_encipherment - - @property - def data_encipherment(self) -> bool: - return self._data_encipherment - - @property - def key_agreement(self) -> bool: - return self._key_agreement - - @property - def key_cert_sign(self) -> bool: - return self._key_cert_sign - - @property - def crl_sign(self) -> bool: - return self._crl_sign - - @property - def encipher_only(self) -> bool: - if not self.key_agreement: - raise ValueError( - "encipher_only is undefined unless key_agreement is true" - ) - else: - return self._encipher_only - - @property - def decipher_only(self) -> bool: - if not self.key_agreement: - raise ValueError( - "decipher_only is undefined unless key_agreement is true" - ) - else: - return self._decipher_only - - def __repr__(self) -> str: - try: - encipher_only = self.encipher_only - decipher_only = self.decipher_only - except ValueError: - # Users found None confusing because even though encipher/decipher - # have no meaning unless key_agreement is true, to construct an - # instance of the class you still need to pass False. - encipher_only = False - decipher_only = False - - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, KeyUsage): - return NotImplemented - - return ( - self.digital_signature == other.digital_signature - and self.content_commitment == other.content_commitment - and self.key_encipherment == other.key_encipherment - and self.data_encipherment == other.data_encipherment - and self.key_agreement == other.key_agreement - and self.key_cert_sign == other.key_cert_sign - and self.crl_sign == other.crl_sign - and self._encipher_only == other._encipher_only - and self._decipher_only == other._decipher_only - ) - - def __hash__(self) -> int: - return hash( - ( - self.digital_signature, - self.content_commitment, - self.key_encipherment, - self.data_encipherment, - self.key_agreement, - self.key_cert_sign, - self.crl_sign, - self._encipher_only, - self._decipher_only, - ) - ) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class PrivateKeyUsagePeriod(ExtensionType): - oid = ExtensionOID.PRIVATE_KEY_USAGE_PERIOD - - def __init__( - self, - not_before: datetime.datetime | None, - not_after: datetime.datetime | None, - ) -> None: - if ( - not isinstance(not_before, datetime.datetime) - and not_before is not None - ): - raise TypeError("not_before must be a datetime.datetime or None") - - if ( - not isinstance(not_after, datetime.datetime) - and not_after is not None - ): - raise TypeError("not_after must be a datetime.datetime or None") - - if not_before is None and not_after is None: - raise ValueError( - "At least one of not_before and not_after must not be None" - ) - - if ( - not_before is not None - and not_after is not None - and not_before > not_after - ): - raise ValueError("not_before must be before not_after") - - self._not_before = not_before - self._not_after = not_after - - @property - def not_before(self) -> datetime.datetime | None: - return self._not_before - - @property - def not_after(self) -> datetime.datetime | None: - return self._not_after - - def __repr__(self) -> str: - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, PrivateKeyUsagePeriod): - return NotImplemented - - return ( - self.not_before == other.not_before - and self.not_after == other.not_after - ) - - def __hash__(self) -> int: - return hash((self.not_before, self.not_after)) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class NameConstraints(ExtensionType): - oid = ExtensionOID.NAME_CONSTRAINTS - - def __init__( - self, - permitted_subtrees: Iterable[GeneralName] | None, - excluded_subtrees: Iterable[GeneralName] | None, - ) -> None: - if permitted_subtrees is not None: - permitted_subtrees = list(permitted_subtrees) - if not permitted_subtrees: - raise ValueError( - "permitted_subtrees must be a non-empty list or None" - ) - if not all(isinstance(x, GeneralName) for x in permitted_subtrees): - raise TypeError( - "permitted_subtrees must be a list of GeneralName objects " - "or None" - ) - - self._validate_tree(permitted_subtrees) - - if excluded_subtrees is not None: - excluded_subtrees = list(excluded_subtrees) - if not excluded_subtrees: - raise ValueError( - "excluded_subtrees must be a non-empty list or None" - ) - if not all(isinstance(x, GeneralName) for x in excluded_subtrees): - raise TypeError( - "excluded_subtrees must be a list of GeneralName objects " - "or None" - ) - - self._validate_tree(excluded_subtrees) - - if permitted_subtrees is None and excluded_subtrees is None: - raise ValueError( - "At least one of permitted_subtrees and excluded_subtrees " - "must not be None" - ) - - self._permitted_subtrees = permitted_subtrees - self._excluded_subtrees = excluded_subtrees - - def __eq__(self, other: object) -> bool: - if not isinstance(other, NameConstraints): - return NotImplemented - - return ( - self.excluded_subtrees == other.excluded_subtrees - and self.permitted_subtrees == other.permitted_subtrees - ) - - def _validate_tree(self, tree: Iterable[GeneralName]) -> None: - self._validate_ip_name(tree) - self._validate_dns_name(tree) - - def _validate_ip_name(self, tree: Iterable[GeneralName]) -> None: - if any( - isinstance(name, IPAddress) - and not isinstance( - name.value, (ipaddress.IPv4Network, ipaddress.IPv6Network) - ) - for name in tree - ): - raise TypeError( - "IPAddress name constraints must be an IPv4Network or" - " IPv6Network object" - ) - - def _validate_dns_name(self, tree: Iterable[GeneralName]) -> None: - if any( - isinstance(name, DNSName) and "*" in name.value for name in tree - ): - raise ValueError( - "DNSName name constraints must not contain the '*' wildcard" - " character" - ) - - def __repr__(self) -> str: - return ( - f"" - ) - - def __hash__(self) -> int: - if self.permitted_subtrees is not None: - ps: tuple[GeneralName, ...] | None = tuple(self.permitted_subtrees) - else: - ps = None - - if self.excluded_subtrees is not None: - es: tuple[GeneralName, ...] | None = tuple(self.excluded_subtrees) - else: - es = None - - return hash((ps, es)) - - @property - def permitted_subtrees( - self, - ) -> list[GeneralName] | None: - return self._permitted_subtrees - - @property - def excluded_subtrees( - self, - ) -> list[GeneralName] | None: - return self._excluded_subtrees - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class Extension(typing.Generic[ExtensionTypeVar]): - def __init__( - self, oid: ObjectIdentifier, critical: bool, value: ExtensionTypeVar - ) -> None: - if not isinstance(oid, ObjectIdentifier): - raise TypeError( - "oid argument must be an ObjectIdentifier instance." - ) - - if not isinstance(critical, bool): - raise TypeError("critical must be a boolean value") - - self._oid = oid - self._critical = critical - self._value = value - - @property - def oid(self) -> ObjectIdentifier: - return self._oid - - @property - def critical(self) -> bool: - return self._critical - - @property - def value(self) -> ExtensionTypeVar: - return self._value - - def __repr__(self) -> str: - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, Extension): - return NotImplemented - - return ( - self.oid == other.oid - and self.critical == other.critical - and self.value == other.value - ) - - def __hash__(self) -> int: - return hash((self.oid, self.critical, self.value)) - - -class GeneralNames: - def __init__(self, general_names: Iterable[GeneralName]) -> None: - general_names = list(general_names) - if not all(isinstance(x, GeneralName) for x in general_names): - raise TypeError( - "Every item in the general_names list must be an " - "object conforming to the GeneralName interface" - ) - - self._general_names = general_names - - __len__, __iter__, __getitem__ = _make_sequence_methods("_general_names") - - @typing.overload - def get_values_for_type( - self, - type: type[DNSName] - | type[UniformResourceIdentifier] - | type[RFC822Name], - ) -> list[str]: ... - - @typing.overload - def get_values_for_type( - self, - type: type[DirectoryName], - ) -> list[Name]: ... - - @typing.overload - def get_values_for_type( - self, - type: type[RegisteredID], - ) -> list[ObjectIdentifier]: ... - - @typing.overload - def get_values_for_type( - self, type: type[IPAddress] - ) -> list[_IPAddressTypes]: ... - - @typing.overload - def get_values_for_type( - self, type: type[OtherName] - ) -> list[OtherName]: ... - - def get_values_for_type( - self, - type: type[DNSName] - | type[DirectoryName] - | type[IPAddress] - | type[OtherName] - | type[RFC822Name] - | type[RegisteredID] - | type[UniformResourceIdentifier], - ) -> ( - list[_IPAddressTypes] - | list[str] - | list[OtherName] - | list[Name] - | list[ObjectIdentifier] - ): - # Return the value of each GeneralName, except for OtherName instances - # which we return directly because it has two important properties not - # just one value. - objs = (i for i in self if isinstance(i, type)) - if type != OtherName: - return [i.value for i in objs] - return list(objs) - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, GeneralNames): - return NotImplemented - - return self._general_names == other._general_names - - def __hash__(self) -> int: - return hash(tuple(self._general_names)) - - -class SubjectAlternativeName(ExtensionType): - oid = ExtensionOID.SUBJECT_ALTERNATIVE_NAME - - def __init__(self, general_names: Iterable[GeneralName]) -> None: - self._general_names = GeneralNames(general_names) - - __len__, __iter__, __getitem__ = _make_sequence_methods("_general_names") - - @typing.overload - def get_values_for_type( - self, - type: type[DNSName] - | type[UniformResourceIdentifier] - | type[RFC822Name], - ) -> list[str]: ... - - @typing.overload - def get_values_for_type( - self, - type: type[DirectoryName], - ) -> list[Name]: ... - - @typing.overload - def get_values_for_type( - self, - type: type[RegisteredID], - ) -> list[ObjectIdentifier]: ... - - @typing.overload - def get_values_for_type( - self, type: type[IPAddress] - ) -> list[_IPAddressTypes]: ... - - @typing.overload - def get_values_for_type( - self, type: type[OtherName] - ) -> list[OtherName]: ... - - def get_values_for_type( - self, - type: type[DNSName] - | type[DirectoryName] - | type[IPAddress] - | type[OtherName] - | type[RFC822Name] - | type[RegisteredID] - | type[UniformResourceIdentifier], - ) -> ( - list[_IPAddressTypes] - | list[str] - | list[OtherName] - | list[Name] - | list[ObjectIdentifier] - ): - return self._general_names.get_values_for_type(type) - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, SubjectAlternativeName): - return NotImplemented - - return self._general_names == other._general_names - - def __hash__(self) -> int: - return hash(self._general_names) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class IssuerAlternativeName(ExtensionType): - oid = ExtensionOID.ISSUER_ALTERNATIVE_NAME - - def __init__(self, general_names: Iterable[GeneralName]) -> None: - self._general_names = GeneralNames(general_names) - - __len__, __iter__, __getitem__ = _make_sequence_methods("_general_names") - - @typing.overload - def get_values_for_type( - self, - type: type[DNSName] - | type[UniformResourceIdentifier] - | type[RFC822Name], - ) -> list[str]: ... - - @typing.overload - def get_values_for_type( - self, - type: type[DirectoryName], - ) -> list[Name]: ... - - @typing.overload - def get_values_for_type( - self, - type: type[RegisteredID], - ) -> list[ObjectIdentifier]: ... - - @typing.overload - def get_values_for_type( - self, type: type[IPAddress] - ) -> list[_IPAddressTypes]: ... - - @typing.overload - def get_values_for_type( - self, type: type[OtherName] - ) -> list[OtherName]: ... - - def get_values_for_type( - self, - type: type[DNSName] - | type[DirectoryName] - | type[IPAddress] - | type[OtherName] - | type[RFC822Name] - | type[RegisteredID] - | type[UniformResourceIdentifier], - ) -> ( - list[_IPAddressTypes] - | list[str] - | list[OtherName] - | list[Name] - | list[ObjectIdentifier] - ): - return self._general_names.get_values_for_type(type) - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, IssuerAlternativeName): - return NotImplemented - - return self._general_names == other._general_names - - def __hash__(self) -> int: - return hash(self._general_names) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class CertificateIssuer(ExtensionType): - oid = CRLEntryExtensionOID.CERTIFICATE_ISSUER - - def __init__(self, general_names: Iterable[GeneralName]) -> None: - self._general_names = GeneralNames(general_names) - - __len__, __iter__, __getitem__ = _make_sequence_methods("_general_names") - - @typing.overload - def get_values_for_type( - self, - type: type[DNSName] - | type[UniformResourceIdentifier] - | type[RFC822Name], - ) -> list[str]: ... - - @typing.overload - def get_values_for_type( - self, - type: type[DirectoryName], - ) -> list[Name]: ... - - @typing.overload - def get_values_for_type( - self, - type: type[RegisteredID], - ) -> list[ObjectIdentifier]: ... - - @typing.overload - def get_values_for_type( - self, type: type[IPAddress] - ) -> list[_IPAddressTypes]: ... - - @typing.overload - def get_values_for_type( - self, type: type[OtherName] - ) -> list[OtherName]: ... - - def get_values_for_type( - self, - type: type[DNSName] - | type[DirectoryName] - | type[IPAddress] - | type[OtherName] - | type[RFC822Name] - | type[RegisteredID] - | type[UniformResourceIdentifier], - ) -> ( - list[_IPAddressTypes] - | list[str] - | list[OtherName] - | list[Name] - | list[ObjectIdentifier] - ): - return self._general_names.get_values_for_type(type) - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, CertificateIssuer): - return NotImplemented - - return self._general_names == other._general_names - - def __hash__(self) -> int: - return hash(self._general_names) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class CRLReason(ExtensionType): - oid = CRLEntryExtensionOID.CRL_REASON - - def __init__(self, reason: ReasonFlags) -> None: - if not isinstance(reason, ReasonFlags): - raise TypeError("reason must be an element from ReasonFlags") - - self._reason = reason - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, CRLReason): - return NotImplemented - - return self.reason == other.reason - - def __hash__(self) -> int: - return hash(self.reason) - - @property - def reason(self) -> ReasonFlags: - return self._reason - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class InvalidityDate(ExtensionType): - oid = CRLEntryExtensionOID.INVALIDITY_DATE - - def __init__(self, invalidity_date: datetime.datetime) -> None: - if not isinstance(invalidity_date, datetime.datetime): - raise TypeError("invalidity_date must be a datetime.datetime") - - self._invalidity_date = invalidity_date - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, InvalidityDate): - return NotImplemented - - return self.invalidity_date == other.invalidity_date - - def __hash__(self) -> int: - return hash(self.invalidity_date) - - @property - def invalidity_date(self) -> datetime.datetime: - return self._invalidity_date - - @property - def invalidity_date_utc(self) -> datetime.datetime: - if self._invalidity_date.tzinfo is None: - return self._invalidity_date.replace(tzinfo=datetime.timezone.utc) - else: - return self._invalidity_date.astimezone(tz=datetime.timezone.utc) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class PrecertificateSignedCertificateTimestamps(ExtensionType): - oid = ExtensionOID.PRECERT_SIGNED_CERTIFICATE_TIMESTAMPS - - def __init__( - self, - signed_certificate_timestamps: Iterable[SignedCertificateTimestamp], - ) -> None: - signed_certificate_timestamps = list(signed_certificate_timestamps) - if not all( - isinstance(sct, SignedCertificateTimestamp) - for sct in signed_certificate_timestamps - ): - raise TypeError( - "Every item in the signed_certificate_timestamps list must be " - "a SignedCertificateTimestamp" - ) - self._signed_certificate_timestamps = signed_certificate_timestamps - - __len__, __iter__, __getitem__ = _make_sequence_methods( - "_signed_certificate_timestamps" - ) - - def __repr__(self) -> str: - return f"" - - def __hash__(self) -> int: - return hash(tuple(self._signed_certificate_timestamps)) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, PrecertificateSignedCertificateTimestamps): - return NotImplemented - - return ( - self._signed_certificate_timestamps - == other._signed_certificate_timestamps - ) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class SignedCertificateTimestamps(ExtensionType): - oid = ExtensionOID.SIGNED_CERTIFICATE_TIMESTAMPS - - def __init__( - self, - signed_certificate_timestamps: Iterable[SignedCertificateTimestamp], - ) -> None: - signed_certificate_timestamps = list(signed_certificate_timestamps) - if not all( - isinstance(sct, SignedCertificateTimestamp) - for sct in signed_certificate_timestamps - ): - raise TypeError( - "Every item in the signed_certificate_timestamps list must be " - "a SignedCertificateTimestamp" - ) - self._signed_certificate_timestamps = signed_certificate_timestamps - - __len__, __iter__, __getitem__ = _make_sequence_methods( - "_signed_certificate_timestamps" - ) - - def __repr__(self) -> str: - return f"" - - def __hash__(self) -> int: - return hash(tuple(self._signed_certificate_timestamps)) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, SignedCertificateTimestamps): - return NotImplemented - - return ( - self._signed_certificate_timestamps - == other._signed_certificate_timestamps - ) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class OCSPNonce(ExtensionType): - oid = OCSPExtensionOID.NONCE - - def __init__(self, nonce: bytes) -> None: - if not isinstance(nonce, bytes): - raise TypeError("nonce must be bytes") - - self._nonce = nonce - - def __eq__(self, other: object) -> bool: - if not isinstance(other, OCSPNonce): - return NotImplemented - - return self.nonce == other.nonce - - def __hash__(self) -> int: - return hash(self.nonce) - - def __repr__(self) -> str: - return f"" - - @property - def nonce(self) -> bytes: - return self._nonce - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class OCSPAcceptableResponses(ExtensionType): - oid = OCSPExtensionOID.ACCEPTABLE_RESPONSES - - def __init__(self, responses: Iterable[ObjectIdentifier]) -> None: - responses = list(responses) - if any(not isinstance(r, ObjectIdentifier) for r in responses): - raise TypeError("All responses must be ObjectIdentifiers") - - self._responses = responses - - def __eq__(self, other: object) -> bool: - if not isinstance(other, OCSPAcceptableResponses): - return NotImplemented - - return self._responses == other._responses - - def __hash__(self) -> int: - return hash(tuple(self._responses)) - - def __repr__(self) -> str: - return f"" - - def __iter__(self) -> Iterator[ObjectIdentifier]: - return iter(self._responses) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class IssuingDistributionPoint(ExtensionType): - oid = ExtensionOID.ISSUING_DISTRIBUTION_POINT - - def __init__( - self, - full_name: Iterable[GeneralName] | None, - relative_name: RelativeDistinguishedName | None, - only_contains_user_certs: bool, - only_contains_ca_certs: bool, - only_some_reasons: frozenset[ReasonFlags] | None, - indirect_crl: bool, - only_contains_attribute_certs: bool, - ) -> None: - if full_name is not None: - full_name = list(full_name) - - if only_some_reasons and ( - not isinstance(only_some_reasons, frozenset) - or not all(isinstance(x, ReasonFlags) for x in only_some_reasons) - ): - raise TypeError( - "only_some_reasons must be None or frozenset of ReasonFlags" - ) - - if only_some_reasons and ( - ReasonFlags.unspecified in only_some_reasons - or ReasonFlags.remove_from_crl in only_some_reasons - ): - raise ValueError( - "unspecified and remove_from_crl are not valid reasons in an " - "IssuingDistributionPoint" - ) - - if not ( - isinstance(only_contains_user_certs, bool) - and isinstance(only_contains_ca_certs, bool) - and isinstance(indirect_crl, bool) - and isinstance(only_contains_attribute_certs, bool) - ): - raise TypeError( - "only_contains_user_certs, only_contains_ca_certs, " - "indirect_crl and only_contains_attribute_certs " - "must all be boolean." - ) - - # Per RFC5280 Section 5.2.5, the Issuing Distribution Point extension - # in a CRL can have only one of onlyContainsUserCerts, - # onlyContainsCACerts, onlyContainsAttributeCerts set to TRUE. - crl_constraints = [ - only_contains_user_certs, - only_contains_ca_certs, - only_contains_attribute_certs, - ] - - if len([x for x in crl_constraints if x]) > 1: - raise ValueError( - "Only one of the following can be set to True: " - "only_contains_user_certs, only_contains_ca_certs, " - "only_contains_attribute_certs" - ) - - if not any( - [ - only_contains_user_certs, - only_contains_ca_certs, - indirect_crl, - only_contains_attribute_certs, - full_name, - relative_name, - only_some_reasons, - ] - ): - raise ValueError( - "Cannot create empty extension: " - "if only_contains_user_certs, only_contains_ca_certs, " - "indirect_crl, and only_contains_attribute_certs are all False" - ", then either full_name, relative_name, or only_some_reasons " - "must have a value." - ) - - self._only_contains_user_certs = only_contains_user_certs - self._only_contains_ca_certs = only_contains_ca_certs - self._indirect_crl = indirect_crl - self._only_contains_attribute_certs = only_contains_attribute_certs - self._only_some_reasons = only_some_reasons - self._full_name = full_name - self._relative_name = relative_name - - def __repr__(self) -> str: - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, IssuingDistributionPoint): - return NotImplemented - - return ( - self.full_name == other.full_name - and self.relative_name == other.relative_name - and self.only_contains_user_certs == other.only_contains_user_certs - and self.only_contains_ca_certs == other.only_contains_ca_certs - and self.only_some_reasons == other.only_some_reasons - and self.indirect_crl == other.indirect_crl - and self.only_contains_attribute_certs - == other.only_contains_attribute_certs - ) - - def __hash__(self) -> int: - return hash( - ( - self.full_name, - self.relative_name, - self.only_contains_user_certs, - self.only_contains_ca_certs, - self.only_some_reasons, - self.indirect_crl, - self.only_contains_attribute_certs, - ) - ) - - @property - def full_name(self) -> list[GeneralName] | None: - return self._full_name - - @property - def relative_name(self) -> RelativeDistinguishedName | None: - return self._relative_name - - @property - def only_contains_user_certs(self) -> bool: - return self._only_contains_user_certs - - @property - def only_contains_ca_certs(self) -> bool: - return self._only_contains_ca_certs - - @property - def only_some_reasons( - self, - ) -> frozenset[ReasonFlags] | None: - return self._only_some_reasons - - @property - def indirect_crl(self) -> bool: - return self._indirect_crl - - @property - def only_contains_attribute_certs(self) -> bool: - return self._only_contains_attribute_certs - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class MSCertificateTemplate(ExtensionType): - oid = ExtensionOID.MS_CERTIFICATE_TEMPLATE - - def __init__( - self, - template_id: ObjectIdentifier, - major_version: int | None, - minor_version: int | None, - ) -> None: - if not isinstance(template_id, ObjectIdentifier): - raise TypeError("oid must be an ObjectIdentifier") - self._template_id = template_id - if ( - major_version is not None and not isinstance(major_version, int) - ) or ( - minor_version is not None and not isinstance(minor_version, int) - ): - raise TypeError( - "major_version and minor_version must be integers or None" - ) - self._major_version = major_version - self._minor_version = minor_version - - @property - def template_id(self) -> ObjectIdentifier: - return self._template_id - - @property - def major_version(self) -> int | None: - return self._major_version - - @property - def minor_version(self) -> int | None: - return self._minor_version - - def __repr__(self) -> str: - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, MSCertificateTemplate): - return NotImplemented - - return ( - self.template_id == other.template_id - and self.major_version == other.major_version - and self.minor_version == other.minor_version - ) - - def __hash__(self) -> int: - return hash((self.template_id, self.major_version, self.minor_version)) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class NamingAuthority: - def __init__( - self, - id: ObjectIdentifier | None, - url: str | None, - text: str | None, - ) -> None: - if id is not None and not isinstance(id, ObjectIdentifier): - raise TypeError("id must be an ObjectIdentifier") - - if url is not None and not isinstance(url, str): - raise TypeError("url must be a str") - - if text is not None and not isinstance(text, str): - raise TypeError("text must be a str") - - self._id = id - self._url = url - self._text = text - - @property - def id(self) -> ObjectIdentifier | None: - return self._id - - @property - def url(self) -> str | None: - return self._url - - @property - def text(self) -> str | None: - return self._text - - def __repr__(self) -> str: - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, NamingAuthority): - return NotImplemented - - return ( - self.id == other.id - and self.url == other.url - and self.text == other.text - ) - - def __hash__(self) -> int: - return hash( - ( - self.id, - self.url, - self.text, - ) - ) - - -class ProfessionInfo: - def __init__( - self, - naming_authority: NamingAuthority | None, - profession_items: Iterable[str], - profession_oids: Iterable[ObjectIdentifier] | None, - registration_number: str | None, - add_profession_info: bytes | None, - ) -> None: - if naming_authority is not None and not isinstance( - naming_authority, NamingAuthority - ): - raise TypeError("naming_authority must be a NamingAuthority") - - profession_items = list(profession_items) - if not all(isinstance(item, str) for item in profession_items): - raise TypeError( - "Every item in the profession_items list must be a str" - ) - - if profession_oids is not None: - profession_oids = list(profession_oids) - if not all( - isinstance(oid, ObjectIdentifier) for oid in profession_oids - ): - raise TypeError( - "Every item in the profession_oids list must be an " - "ObjectIdentifier" - ) - - if registration_number is not None and not isinstance( - registration_number, str - ): - raise TypeError("registration_number must be a str") - - if add_profession_info is not None and not isinstance( - add_profession_info, bytes - ): - raise TypeError("add_profession_info must be bytes") - - self._naming_authority = naming_authority - self._profession_items = profession_items - self._profession_oids = profession_oids - self._registration_number = registration_number - self._add_profession_info = add_profession_info - - @property - def naming_authority(self) -> NamingAuthority | None: - return self._naming_authority - - @property - def profession_items(self) -> list[str]: - return self._profession_items - - @property - def profession_oids(self) -> list[ObjectIdentifier] | None: - return self._profession_oids - - @property - def registration_number(self) -> str | None: - return self._registration_number - - @property - def add_profession_info(self) -> bytes | None: - return self._add_profession_info - - def __repr__(self) -> str: - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, ProfessionInfo): - return NotImplemented - - return ( - self.naming_authority == other.naming_authority - and self.profession_items == other.profession_items - and self.profession_oids == other.profession_oids - and self.registration_number == other.registration_number - and self.add_profession_info == other.add_profession_info - ) - - def __hash__(self) -> int: - if self.profession_oids is not None: - profession_oids = tuple(self.profession_oids) - else: - profession_oids = None - return hash( - ( - self.naming_authority, - tuple(self.profession_items), - profession_oids, - self.registration_number, - self.add_profession_info, - ) - ) - - -class Admission: - def __init__( - self, - admission_authority: GeneralName | None, - naming_authority: NamingAuthority | None, - profession_infos: Iterable[ProfessionInfo], - ) -> None: - if admission_authority is not None and not isinstance( - admission_authority, GeneralName - ): - raise TypeError("admission_authority must be a GeneralName") - - if naming_authority is not None and not isinstance( - naming_authority, NamingAuthority - ): - raise TypeError("naming_authority must be a NamingAuthority") - - profession_infos = list(profession_infos) - if not all( - isinstance(info, ProfessionInfo) for info in profession_infos - ): - raise TypeError( - "Every item in the profession_infos list must be a " - "ProfessionInfo" - ) - - self._admission_authority = admission_authority - self._naming_authority = naming_authority - self._profession_infos = profession_infos - - @property - def admission_authority(self) -> GeneralName | None: - return self._admission_authority - - @property - def naming_authority(self) -> NamingAuthority | None: - return self._naming_authority - - @property - def profession_infos(self) -> list[ProfessionInfo]: - return self._profession_infos - - def __repr__(self) -> str: - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, Admission): - return NotImplemented - - return ( - self.admission_authority == other.admission_authority - and self.naming_authority == other.naming_authority - and self.profession_infos == other.profession_infos - ) - - def __hash__(self) -> int: - return hash( - ( - self.admission_authority, - self.naming_authority, - tuple(self.profession_infos), - ) - ) - - -class Admissions(ExtensionType): - oid = ExtensionOID.ADMISSIONS - - def __init__( - self, - authority: GeneralName | None, - admissions: Iterable[Admission], - ) -> None: - if authority is not None and not isinstance(authority, GeneralName): - raise TypeError("authority must be a GeneralName") - - admissions = list(admissions) - if not all( - isinstance(admission, Admission) for admission in admissions - ): - raise TypeError( - "Every item in the contents_of_admissions list must be an " - "Admission" - ) - - self._authority = authority - self._admissions = admissions - - __len__, __iter__, __getitem__ = _make_sequence_methods("_admissions") - - @property - def authority(self) -> GeneralName | None: - return self._authority - - def __repr__(self) -> str: - return ( - f"" - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, Admissions): - return NotImplemented - - return ( - self.authority == other.authority - and self._admissions == other._admissions - ) - - def __hash__(self) -> int: - return hash((self.authority, tuple(self._admissions))) - - def public_bytes(self) -> bytes: - return rust_x509.encode_extension_value(self) - - -class UnrecognizedExtension(ExtensionType): - def __init__(self, oid: ObjectIdentifier, value: bytes) -> None: - if not isinstance(oid, ObjectIdentifier): - raise TypeError("oid must be an ObjectIdentifier") - self._oid = oid - self._value = value - - @property - def oid(self) -> ObjectIdentifier: # type: ignore[override] - return self._oid - - @property - def value(self) -> bytes: - return self._value - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, UnrecognizedExtension): - return NotImplemented - - return self.oid == other.oid and self.value == other.value - - def __hash__(self) -> int: - return hash((self.oid, self.value)) - - def public_bytes(self) -> bytes: - return self.value diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/general_name.py b/venv/lib/python3.10/site-packages/cryptography/x509/general_name.py deleted file mode 100644 index 672f28759cb0ab6d60cb4e38a574095ae3e0159d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/x509/general_name.py +++ /dev/null @@ -1,281 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import abc -import ipaddress -import typing -from email.utils import parseaddr - -from cryptography.x509.name import Name -from cryptography.x509.oid import ObjectIdentifier - -_IPAddressTypes = typing.Union[ - ipaddress.IPv4Address, - ipaddress.IPv6Address, - ipaddress.IPv4Network, - ipaddress.IPv6Network, -] - - -class UnsupportedGeneralNameType(Exception): - pass - - -class GeneralName(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def value(self) -> typing.Any: - """ - Return the value of the object - """ - - -class RFC822Name(GeneralName): - def __init__(self, value: str) -> None: - if isinstance(value, str): - try: - value.encode("ascii") - except UnicodeEncodeError: - raise ValueError( - "RFC822Name values should be passed as an A-label string. " - "This means unicode characters should be encoded via " - "a library like idna." - ) - else: - raise TypeError("value must be string") - - name, address = parseaddr(value) - if name or not address: - # parseaddr has found a name (e.g. Name ) or the entire - # value is an empty string. - raise ValueError("Invalid rfc822name value") - - self._value = value - - @property - def value(self) -> str: - return self._value - - @classmethod - def _init_without_validation(cls, value: str) -> RFC822Name: - instance = cls.__new__(cls) - instance._value = value - return instance - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, RFC822Name): - return NotImplemented - - return self.value == other.value - - def __hash__(self) -> int: - return hash(self.value) - - -class DNSName(GeneralName): - def __init__(self, value: str) -> None: - if isinstance(value, str): - try: - value.encode("ascii") - except UnicodeEncodeError: - raise ValueError( - "DNSName values should be passed as an A-label string. " - "This means unicode characters should be encoded via " - "a library like idna." - ) - else: - raise TypeError("value must be string") - - self._value = value - - @property - def value(self) -> str: - return self._value - - @classmethod - def _init_without_validation(cls, value: str) -> DNSName: - instance = cls.__new__(cls) - instance._value = value - return instance - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, DNSName): - return NotImplemented - - return self.value == other.value - - def __hash__(self) -> int: - return hash(self.value) - - -class UniformResourceIdentifier(GeneralName): - def __init__(self, value: str) -> None: - if isinstance(value, str): - try: - value.encode("ascii") - except UnicodeEncodeError: - raise ValueError( - "URI values should be passed as an A-label string. " - "This means unicode characters should be encoded via " - "a library like idna." - ) - else: - raise TypeError("value must be string") - - self._value = value - - @property - def value(self) -> str: - return self._value - - @classmethod - def _init_without_validation(cls, value: str) -> UniformResourceIdentifier: - instance = cls.__new__(cls) - instance._value = value - return instance - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, UniformResourceIdentifier): - return NotImplemented - - return self.value == other.value - - def __hash__(self) -> int: - return hash(self.value) - - -class DirectoryName(GeneralName): - def __init__(self, value: Name) -> None: - if not isinstance(value, Name): - raise TypeError("value must be a Name") - - self._value = value - - @property - def value(self) -> Name: - return self._value - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, DirectoryName): - return NotImplemented - - return self.value == other.value - - def __hash__(self) -> int: - return hash(self.value) - - -class RegisteredID(GeneralName): - def __init__(self, value: ObjectIdentifier) -> None: - if not isinstance(value, ObjectIdentifier): - raise TypeError("value must be an ObjectIdentifier") - - self._value = value - - @property - def value(self) -> ObjectIdentifier: - return self._value - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, RegisteredID): - return NotImplemented - - return self.value == other.value - - def __hash__(self) -> int: - return hash(self.value) - - -class IPAddress(GeneralName): - def __init__(self, value: _IPAddressTypes) -> None: - if not isinstance( - value, - ( - ipaddress.IPv4Address, - ipaddress.IPv6Address, - ipaddress.IPv4Network, - ipaddress.IPv6Network, - ), - ): - raise TypeError( - "value must be an instance of ipaddress.IPv4Address, " - "ipaddress.IPv6Address, ipaddress.IPv4Network, or " - "ipaddress.IPv6Network" - ) - - self._value = value - - @property - def value(self) -> _IPAddressTypes: - return self._value - - def _packed(self) -> bytes: - if isinstance( - self.value, (ipaddress.IPv4Address, ipaddress.IPv6Address) - ): - return self.value.packed - else: - return ( - self.value.network_address.packed + self.value.netmask.packed - ) - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, IPAddress): - return NotImplemented - - return self.value == other.value - - def __hash__(self) -> int: - return hash(self.value) - - -class OtherName(GeneralName): - def __init__(self, type_id: ObjectIdentifier, value: bytes) -> None: - if not isinstance(type_id, ObjectIdentifier): - raise TypeError("type_id must be an ObjectIdentifier") - if not isinstance(value, bytes): - raise TypeError("value must be a binary string") - - self._type_id = type_id - self._value = value - - @property - def type_id(self) -> ObjectIdentifier: - return self._type_id - - @property - def value(self) -> bytes: - return self._value - - def __repr__(self) -> str: - return f"" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, OtherName): - return NotImplemented - - return self.type_id == other.type_id and self.value == other.value - - def __hash__(self) -> int: - return hash((self.type_id, self.value)) diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/name.py b/venv/lib/python3.10/site-packages/cryptography/x509/name.py deleted file mode 100644 index 5f827818c254f07481c7500814347a0072371772..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/x509/name.py +++ /dev/null @@ -1,477 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import binascii -import re -import sys -import typing -import warnings -from collections.abc import Iterable, Iterator - -from cryptography import utils -from cryptography.hazmat.bindings._rust import x509 as rust_x509 -from cryptography.x509.oid import NameOID, ObjectIdentifier - - -class _ASN1Type(utils.Enum): - BitString = 3 - OctetString = 4 - UTF8String = 12 - NumericString = 18 - PrintableString = 19 - T61String = 20 - IA5String = 22 - UTCTime = 23 - GeneralizedTime = 24 - VisibleString = 26 - UniversalString = 28 - BMPString = 30 - - -_ASN1_TYPE_TO_ENUM = {i.value: i for i in _ASN1Type} -_NAMEOID_DEFAULT_TYPE: dict[ObjectIdentifier, _ASN1Type] = { - NameOID.COUNTRY_NAME: _ASN1Type.PrintableString, - NameOID.JURISDICTION_COUNTRY_NAME: _ASN1Type.PrintableString, - NameOID.SERIAL_NUMBER: _ASN1Type.PrintableString, - NameOID.DN_QUALIFIER: _ASN1Type.PrintableString, - NameOID.EMAIL_ADDRESS: _ASN1Type.IA5String, - NameOID.DOMAIN_COMPONENT: _ASN1Type.IA5String, -} - -# Type alias -_OidNameMap = typing.Mapping[ObjectIdentifier, str] -_NameOidMap = typing.Mapping[str, ObjectIdentifier] - -#: Short attribute names from RFC 4514: -#: https://tools.ietf.org/html/rfc4514#page-7 -_NAMEOID_TO_NAME: _OidNameMap = { - NameOID.COMMON_NAME: "CN", - NameOID.LOCALITY_NAME: "L", - NameOID.STATE_OR_PROVINCE_NAME: "ST", - NameOID.ORGANIZATION_NAME: "O", - NameOID.ORGANIZATIONAL_UNIT_NAME: "OU", - NameOID.COUNTRY_NAME: "C", - NameOID.STREET_ADDRESS: "STREET", - NameOID.DOMAIN_COMPONENT: "DC", - NameOID.USER_ID: "UID", -} -_NAME_TO_NAMEOID = {v: k for k, v in _NAMEOID_TO_NAME.items()} - -_NAMEOID_LENGTH_LIMIT = { - NameOID.COUNTRY_NAME: (2, 2), - NameOID.JURISDICTION_COUNTRY_NAME: (2, 2), - NameOID.COMMON_NAME: (1, 64), -} - - -def _escape_dn_value(val: str | bytes) -> str: - """Escape special characters in RFC4514 Distinguished Name value.""" - - if not val: - return "" - - # RFC 4514 Section 2.4 defines the value as being the # (U+0023) character - # followed by the hexadecimal encoding of the octets. - if isinstance(val, bytes): - return "#" + binascii.hexlify(val).decode("utf8") - - # See https://tools.ietf.org/html/rfc4514#section-2.4 - val = val.replace("\\", "\\\\") - val = val.replace('"', '\\"') - val = val.replace("+", "\\+") - val = val.replace(",", "\\,") - val = val.replace(";", "\\;") - val = val.replace("<", "\\<") - val = val.replace(">", "\\>") - val = val.replace("\0", "\\00") - - if val[0] in ("#", " "): - val = "\\" + val - if val[-1] == " ": - val = val[:-1] + "\\ " - - return val - - -def _unescape_dn_value(val: str) -> str: - if not val: - return "" - - # See https://tools.ietf.org/html/rfc4514#section-3 - - # special = escaped / SPACE / SHARP / EQUALS - # escaped = DQUOTE / PLUS / COMMA / SEMI / LANGLE / RANGLE - def sub(m): - val = m.group(1) - # Regular escape - if len(val) == 1: - return val - # Hex-value scape - return chr(int(val, 16)) - - return _RFC4514NameParser._PAIR_RE.sub(sub, val) - - -NameAttributeValueType = typing.TypeVar( - "NameAttributeValueType", - typing.Union[str, bytes], - str, - bytes, - covariant=True, -) - - -class NameAttribute(typing.Generic[NameAttributeValueType]): - def __init__( - self, - oid: ObjectIdentifier, - value: NameAttributeValueType, - _type: _ASN1Type | None = None, - *, - _validate: bool = True, - ) -> None: - if not isinstance(oid, ObjectIdentifier): - raise TypeError( - "oid argument must be an ObjectIdentifier instance." - ) - if _type == _ASN1Type.BitString: - if oid != NameOID.X500_UNIQUE_IDENTIFIER: - raise TypeError( - "oid must be X500_UNIQUE_IDENTIFIER for BitString type." - ) - if not isinstance(value, bytes): - raise TypeError("value must be bytes for BitString") - else: - if not isinstance(value, str): - raise TypeError("value argument must be a str") - - length_limits = _NAMEOID_LENGTH_LIMIT.get(oid) - if length_limits is not None: - min_length, max_length = length_limits - assert isinstance(value, str) - c_len = len(value.encode("utf8")) - if c_len < min_length or c_len > max_length: - msg = ( - f"Attribute's length must be >= {min_length} and " - f"<= {max_length}, but it was {c_len}" - ) - if _validate is True: - raise ValueError(msg) - else: - warnings.warn(msg, stacklevel=2) - - # The appropriate ASN1 string type varies by OID and is defined across - # multiple RFCs including 2459, 3280, and 5280. In general UTF8String - # is preferred (2459), but 3280 and 5280 specify several OIDs with - # alternate types. This means when we see the sentinel value we need - # to look up whether the OID has a non-UTF8 type. If it does, set it - # to that. Otherwise, UTF8! - if _type is None: - _type = _NAMEOID_DEFAULT_TYPE.get(oid, _ASN1Type.UTF8String) - - if not isinstance(_type, _ASN1Type): - raise TypeError("_type must be from the _ASN1Type enum") - - self._oid = oid - self._value: NameAttributeValueType = value - self._type: _ASN1Type = _type - - @property - def oid(self) -> ObjectIdentifier: - return self._oid - - @property - def value(self) -> NameAttributeValueType: - return self._value - - @property - def rfc4514_attribute_name(self) -> str: - """ - The short attribute name (for example "CN") if available, - otherwise the OID dotted string. - """ - return _NAMEOID_TO_NAME.get(self.oid, self.oid.dotted_string) - - def rfc4514_string( - self, attr_name_overrides: _OidNameMap | None = None - ) -> str: - """ - Format as RFC4514 Distinguished Name string. - - Use short attribute name if available, otherwise fall back to OID - dotted string. - """ - attr_name = ( - attr_name_overrides.get(self.oid) if attr_name_overrides else None - ) - if attr_name is None: - attr_name = self.rfc4514_attribute_name - - return f"{attr_name}={_escape_dn_value(self.value)}" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, NameAttribute): - return NotImplemented - - return self.oid == other.oid and self.value == other.value - - def __hash__(self) -> int: - return hash((self.oid, self.value)) - - def __repr__(self) -> str: - return f"" - - -class RelativeDistinguishedName: - def __init__(self, attributes: Iterable[NameAttribute]): - attributes = list(attributes) - if not attributes: - raise ValueError("a relative distinguished name cannot be empty") - if not all(isinstance(x, NameAttribute) for x in attributes): - raise TypeError("attributes must be an iterable of NameAttribute") - - # Keep list and frozenset to preserve attribute order where it matters - self._attributes = attributes - self._attribute_set = frozenset(attributes) - - if len(self._attribute_set) != len(attributes): - raise ValueError("duplicate attributes are not allowed") - - def get_attributes_for_oid( - self, - oid: ObjectIdentifier, - ) -> list[NameAttribute[str | bytes]]: - return [i for i in self if i.oid == oid] - - def rfc4514_string( - self, attr_name_overrides: _OidNameMap | None = None - ) -> str: - """ - Format as RFC4514 Distinguished Name string. - - Within each RDN, attributes are joined by '+', although that is rarely - used in certificates. - """ - return "+".join( - attr.rfc4514_string(attr_name_overrides) - for attr in self._attributes - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, RelativeDistinguishedName): - return NotImplemented - - return self._attribute_set == other._attribute_set - - def __hash__(self) -> int: - return hash(self._attribute_set) - - def __iter__(self) -> Iterator[NameAttribute]: - return iter(self._attributes) - - def __len__(self) -> int: - return len(self._attributes) - - def __repr__(self) -> str: - return f"" - - -class Name: - @typing.overload - def __init__(self, attributes: Iterable[NameAttribute]) -> None: ... - - @typing.overload - def __init__( - self, attributes: Iterable[RelativeDistinguishedName] - ) -> None: ... - - def __init__( - self, - attributes: Iterable[NameAttribute | RelativeDistinguishedName], - ) -> None: - attributes = list(attributes) - if all(isinstance(x, NameAttribute) for x in attributes): - self._attributes = [ - RelativeDistinguishedName([typing.cast(NameAttribute, x)]) - for x in attributes - ] - elif all(isinstance(x, RelativeDistinguishedName) for x in attributes): - self._attributes = typing.cast( - typing.List[RelativeDistinguishedName], attributes - ) - else: - raise TypeError( - "attributes must be a list of NameAttribute" - " or a list RelativeDistinguishedName" - ) - - @classmethod - def from_rfc4514_string( - cls, - data: str, - attr_name_overrides: _NameOidMap | None = None, - ) -> Name: - return _RFC4514NameParser(data, attr_name_overrides or {}).parse() - - def rfc4514_string( - self, attr_name_overrides: _OidNameMap | None = None - ) -> str: - """ - Format as RFC4514 Distinguished Name string. - For example 'CN=foobar.com,O=Foo Corp,C=US' - - An X.509 name is a two-level structure: a list of sets of attributes. - Each list element is separated by ',' and within each list element, set - elements are separated by '+'. The latter is almost never used in - real world certificates. According to RFC4514 section 2.1 the - RDNSequence must be reversed when converting to string representation. - """ - return ",".join( - attr.rfc4514_string(attr_name_overrides) - for attr in reversed(self._attributes) - ) - - def get_attributes_for_oid( - self, - oid: ObjectIdentifier, - ) -> list[NameAttribute[str | bytes]]: - return [i for i in self if i.oid == oid] - - @property - def rdns(self) -> list[RelativeDistinguishedName]: - return self._attributes - - def public_bytes(self, backend: typing.Any = None) -> bytes: - return rust_x509.encode_name_bytes(self) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, Name): - return NotImplemented - - return self._attributes == other._attributes - - def __hash__(self) -> int: - # TODO: this is relatively expensive, if this looks like a bottleneck - # for you, consider optimizing! - return hash(tuple(self._attributes)) - - def __iter__(self) -> Iterator[NameAttribute]: - for rdn in self._attributes: - yield from rdn - - def __len__(self) -> int: - return sum(len(rdn) for rdn in self._attributes) - - def __repr__(self) -> str: - rdns = ",".join(attr.rfc4514_string() for attr in self._attributes) - return f"" - - -class _RFC4514NameParser: - _OID_RE = re.compile(r"(0|([1-9]\d*))(\.(0|([1-9]\d*)))+") - _DESCR_RE = re.compile(r"[a-zA-Z][a-zA-Z\d-]*") - - _PAIR = r"\\([\\ #=\"\+,;<>]|[\da-zA-Z]{2})" - _PAIR_RE = re.compile(_PAIR) - _LUTF1 = r"[\x01-\x1f\x21\x24-\x2A\x2D-\x3A\x3D\x3F-\x5B\x5D-\x7F]" - _SUTF1 = r"[\x01-\x21\x23-\x2A\x2D-\x3A\x3D\x3F-\x5B\x5D-\x7F]" - _TUTF1 = r"[\x01-\x1F\x21\x23-\x2A\x2D-\x3A\x3D\x3F-\x5B\x5D-\x7F]" - _UTFMB = rf"[\x80-{chr(sys.maxunicode)}]" - _LEADCHAR = rf"{_LUTF1}|{_UTFMB}" - _STRINGCHAR = rf"{_SUTF1}|{_UTFMB}" - _TRAILCHAR = rf"{_TUTF1}|{_UTFMB}" - _STRING_RE = re.compile( - rf""" - ( - ({_LEADCHAR}|{_PAIR}) - ( - ({_STRINGCHAR}|{_PAIR})* - ({_TRAILCHAR}|{_PAIR}) - )? - )? - """, - re.VERBOSE, - ) - _HEXSTRING_RE = re.compile(r"#([\da-zA-Z]{2})+") - - def __init__(self, data: str, attr_name_overrides: _NameOidMap) -> None: - self._data = data - self._idx = 0 - - self._attr_name_overrides = attr_name_overrides - - def _has_data(self) -> bool: - return self._idx < len(self._data) - - def _peek(self) -> str | None: - if self._has_data(): - return self._data[self._idx] - return None - - def _read_char(self, ch: str) -> None: - if self._peek() != ch: - raise ValueError - self._idx += 1 - - def _read_re(self, pat) -> str: - match = pat.match(self._data, pos=self._idx) - if match is None: - raise ValueError - val = match.group() - self._idx += len(val) - return val - - def parse(self) -> Name: - """ - Parses the `data` string and converts it to a Name. - - According to RFC4514 section 2.1 the RDNSequence must be - reversed when converting to string representation. So, when - we parse it, we need to reverse again to get the RDNs on the - correct order. - """ - - if not self._has_data(): - return Name([]) - - rdns = [self._parse_rdn()] - - while self._has_data(): - self._read_char(",") - rdns.append(self._parse_rdn()) - - return Name(reversed(rdns)) - - def _parse_rdn(self) -> RelativeDistinguishedName: - nas = [self._parse_na()] - while self._peek() == "+": - self._read_char("+") - nas.append(self._parse_na()) - - return RelativeDistinguishedName(nas) - - def _parse_na(self) -> NameAttribute: - try: - oid_value = self._read_re(self._OID_RE) - except ValueError: - name = self._read_re(self._DESCR_RE) - oid = self._attr_name_overrides.get( - name, _NAME_TO_NAMEOID.get(name) - ) - if oid is None: - raise ValueError - else: - oid = ObjectIdentifier(oid_value) - - self._read_char("=") - if self._peek() == "#": - value = self._read_re(self._HEXSTRING_RE) - value = binascii.unhexlify(value[1:]).decode() - else: - raw_value = self._read_re(self._STRING_RE) - value = _unescape_dn_value(raw_value) - - return NameAttribute(oid, value) diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/ocsp.py b/venv/lib/python3.10/site-packages/cryptography/x509/ocsp.py deleted file mode 100644 index f61ed80b82bdea361b70d676a632f173368adc26..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/x509/ocsp.py +++ /dev/null @@ -1,379 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import datetime -from collections.abc import Iterable - -from cryptography import utils, x509 -from cryptography.hazmat.bindings._rust import ocsp -from cryptography.hazmat.primitives import hashes -from cryptography.hazmat.primitives.asymmetric.types import ( - CertificateIssuerPrivateKeyTypes, -) -from cryptography.x509.base import _reject_duplicate_extension - - -class OCSPResponderEncoding(utils.Enum): - HASH = "By Hash" - NAME = "By Name" - - -class OCSPResponseStatus(utils.Enum): - SUCCESSFUL = 0 - MALFORMED_REQUEST = 1 - INTERNAL_ERROR = 2 - TRY_LATER = 3 - SIG_REQUIRED = 5 - UNAUTHORIZED = 6 - - -_ALLOWED_HASHES = ( - hashes.SHA1, - hashes.SHA224, - hashes.SHA256, - hashes.SHA384, - hashes.SHA512, -) - - -def _verify_algorithm(algorithm: hashes.HashAlgorithm) -> None: - if not isinstance(algorithm, _ALLOWED_HASHES): - raise ValueError( - "Algorithm must be SHA1, SHA224, SHA256, SHA384, or SHA512" - ) - - -class OCSPCertStatus(utils.Enum): - GOOD = 0 - REVOKED = 1 - UNKNOWN = 2 - - -class _SingleResponse: - def __init__( - self, - resp: tuple[x509.Certificate, x509.Certificate] | None, - resp_hash: tuple[bytes, bytes, int] | None, - algorithm: hashes.HashAlgorithm, - cert_status: OCSPCertStatus, - this_update: datetime.datetime, - next_update: datetime.datetime | None, - revocation_time: datetime.datetime | None, - revocation_reason: x509.ReasonFlags | None, - ): - _verify_algorithm(algorithm) - if not isinstance(this_update, datetime.datetime): - raise TypeError("this_update must be a datetime object") - if next_update is not None and not isinstance( - next_update, datetime.datetime - ): - raise TypeError("next_update must be a datetime object or None") - - self._resp = resp - self._resp_hash = resp_hash - self._algorithm = algorithm - self._this_update = this_update - self._next_update = next_update - - if not isinstance(cert_status, OCSPCertStatus): - raise TypeError( - "cert_status must be an item from the OCSPCertStatus enum" - ) - if cert_status is not OCSPCertStatus.REVOKED: - if revocation_time is not None: - raise ValueError( - "revocation_time can only be provided if the certificate " - "is revoked" - ) - if revocation_reason is not None: - raise ValueError( - "revocation_reason can only be provided if the certificate" - " is revoked" - ) - else: - if not isinstance(revocation_time, datetime.datetime): - raise TypeError("revocation_time must be a datetime object") - - if revocation_reason is not None and not isinstance( - revocation_reason, x509.ReasonFlags - ): - raise TypeError( - "revocation_reason must be an item from the ReasonFlags " - "enum or None" - ) - - self._cert_status = cert_status - self._revocation_time = revocation_time - self._revocation_reason = revocation_reason - - -OCSPRequest = ocsp.OCSPRequest -OCSPResponse = ocsp.OCSPResponse -OCSPSingleResponse = ocsp.OCSPSingleResponse - - -class OCSPRequestBuilder: - def __init__( - self, - request: tuple[ - x509.Certificate, x509.Certificate, hashes.HashAlgorithm - ] - | None = None, - request_hash: tuple[bytes, bytes, int, hashes.HashAlgorithm] - | None = None, - extensions: list[x509.Extension[x509.ExtensionType]] = [], - ) -> None: - self._request = request - self._request_hash = request_hash - self._extensions = extensions - - def add_certificate( - self, - cert: x509.Certificate, - issuer: x509.Certificate, - algorithm: hashes.HashAlgorithm, - ) -> OCSPRequestBuilder: - if self._request is not None or self._request_hash is not None: - raise ValueError("Only one certificate can be added to a request") - - _verify_algorithm(algorithm) - if not isinstance(cert, x509.Certificate) or not isinstance( - issuer, x509.Certificate - ): - raise TypeError("cert and issuer must be a Certificate") - - return OCSPRequestBuilder( - (cert, issuer, algorithm), self._request_hash, self._extensions - ) - - def add_certificate_by_hash( - self, - issuer_name_hash: bytes, - issuer_key_hash: bytes, - serial_number: int, - algorithm: hashes.HashAlgorithm, - ) -> OCSPRequestBuilder: - if self._request is not None or self._request_hash is not None: - raise ValueError("Only one certificate can be added to a request") - - if not isinstance(serial_number, int): - raise TypeError("serial_number must be an integer") - - _verify_algorithm(algorithm) - utils._check_bytes("issuer_name_hash", issuer_name_hash) - utils._check_bytes("issuer_key_hash", issuer_key_hash) - if algorithm.digest_size != len( - issuer_name_hash - ) or algorithm.digest_size != len(issuer_key_hash): - raise ValueError( - "issuer_name_hash and issuer_key_hash must be the same length " - "as the digest size of the algorithm" - ) - - return OCSPRequestBuilder( - self._request, - (issuer_name_hash, issuer_key_hash, serial_number, algorithm), - self._extensions, - ) - - def add_extension( - self, extval: x509.ExtensionType, critical: bool - ) -> OCSPRequestBuilder: - if not isinstance(extval, x509.ExtensionType): - raise TypeError("extension must be an ExtensionType") - - extension = x509.Extension(extval.oid, critical, extval) - _reject_duplicate_extension(extension, self._extensions) - - return OCSPRequestBuilder( - self._request, self._request_hash, [*self._extensions, extension] - ) - - def build(self) -> OCSPRequest: - if self._request is None and self._request_hash is None: - raise ValueError("You must add a certificate before building") - - return ocsp.create_ocsp_request(self) - - -class OCSPResponseBuilder: - def __init__( - self, - response: _SingleResponse | None = None, - responder_id: tuple[x509.Certificate, OCSPResponderEncoding] - | None = None, - certs: list[x509.Certificate] | None = None, - extensions: list[x509.Extension[x509.ExtensionType]] = [], - ): - self._response = response - self._responder_id = responder_id - self._certs = certs - self._extensions = extensions - - def add_response( - self, - cert: x509.Certificate, - issuer: x509.Certificate, - algorithm: hashes.HashAlgorithm, - cert_status: OCSPCertStatus, - this_update: datetime.datetime, - next_update: datetime.datetime | None, - revocation_time: datetime.datetime | None, - revocation_reason: x509.ReasonFlags | None, - ) -> OCSPResponseBuilder: - if self._response is not None: - raise ValueError("Only one response per OCSPResponse.") - - if not isinstance(cert, x509.Certificate) or not isinstance( - issuer, x509.Certificate - ): - raise TypeError("cert and issuer must be a Certificate") - - singleresp = _SingleResponse( - (cert, issuer), - None, - algorithm, - cert_status, - this_update, - next_update, - revocation_time, - revocation_reason, - ) - return OCSPResponseBuilder( - singleresp, - self._responder_id, - self._certs, - self._extensions, - ) - - def add_response_by_hash( - self, - issuer_name_hash: bytes, - issuer_key_hash: bytes, - serial_number: int, - algorithm: hashes.HashAlgorithm, - cert_status: OCSPCertStatus, - this_update: datetime.datetime, - next_update: datetime.datetime | None, - revocation_time: datetime.datetime | None, - revocation_reason: x509.ReasonFlags | None, - ) -> OCSPResponseBuilder: - if self._response is not None: - raise ValueError("Only one response per OCSPResponse.") - - if not isinstance(serial_number, int): - raise TypeError("serial_number must be an integer") - - utils._check_bytes("issuer_name_hash", issuer_name_hash) - utils._check_bytes("issuer_key_hash", issuer_key_hash) - _verify_algorithm(algorithm) - if algorithm.digest_size != len( - issuer_name_hash - ) or algorithm.digest_size != len(issuer_key_hash): - raise ValueError( - "issuer_name_hash and issuer_key_hash must be the same length " - "as the digest size of the algorithm" - ) - - singleresp = _SingleResponse( - None, - (issuer_name_hash, issuer_key_hash, serial_number), - algorithm, - cert_status, - this_update, - next_update, - revocation_time, - revocation_reason, - ) - return OCSPResponseBuilder( - singleresp, - self._responder_id, - self._certs, - self._extensions, - ) - - def responder_id( - self, encoding: OCSPResponderEncoding, responder_cert: x509.Certificate - ) -> OCSPResponseBuilder: - if self._responder_id is not None: - raise ValueError("responder_id can only be set once") - if not isinstance(responder_cert, x509.Certificate): - raise TypeError("responder_cert must be a Certificate") - if not isinstance(encoding, OCSPResponderEncoding): - raise TypeError( - "encoding must be an element from OCSPResponderEncoding" - ) - - return OCSPResponseBuilder( - self._response, - (responder_cert, encoding), - self._certs, - self._extensions, - ) - - def certificates( - self, certs: Iterable[x509.Certificate] - ) -> OCSPResponseBuilder: - if self._certs is not None: - raise ValueError("certificates may only be set once") - certs = list(certs) - if len(certs) == 0: - raise ValueError("certs must not be an empty list") - if not all(isinstance(x, x509.Certificate) for x in certs): - raise TypeError("certs must be a list of Certificates") - return OCSPResponseBuilder( - self._response, - self._responder_id, - certs, - self._extensions, - ) - - def add_extension( - self, extval: x509.ExtensionType, critical: bool - ) -> OCSPResponseBuilder: - if not isinstance(extval, x509.ExtensionType): - raise TypeError("extension must be an ExtensionType") - - extension = x509.Extension(extval.oid, critical, extval) - _reject_duplicate_extension(extension, self._extensions) - - return OCSPResponseBuilder( - self._response, - self._responder_id, - self._certs, - [*self._extensions, extension], - ) - - def sign( - self, - private_key: CertificateIssuerPrivateKeyTypes, - algorithm: hashes.HashAlgorithm | None, - ) -> OCSPResponse: - if self._response is None: - raise ValueError("You must add a response before signing") - if self._responder_id is None: - raise ValueError("You must add a responder_id before signing") - - return ocsp.create_ocsp_response( - OCSPResponseStatus.SUCCESSFUL, self, private_key, algorithm - ) - - @classmethod - def build_unsuccessful( - cls, response_status: OCSPResponseStatus - ) -> OCSPResponse: - if not isinstance(response_status, OCSPResponseStatus): - raise TypeError( - "response_status must be an item from OCSPResponseStatus" - ) - if response_status is OCSPResponseStatus.SUCCESSFUL: - raise ValueError("response_status cannot be SUCCESSFUL") - - return ocsp.create_ocsp_response(response_status, None, None, None) - - -load_der_ocsp_request = ocsp.load_der_ocsp_request -load_der_ocsp_response = ocsp.load_der_ocsp_response diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/oid.py b/venv/lib/python3.10/site-packages/cryptography/x509/oid.py deleted file mode 100644 index 520fc7ab018c1b62132b957d332a83abef51b39e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/x509/oid.py +++ /dev/null @@ -1,37 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -from cryptography.hazmat._oid import ( - AttributeOID, - AuthorityInformationAccessOID, - CertificatePoliciesOID, - CRLEntryExtensionOID, - ExtendedKeyUsageOID, - ExtensionOID, - NameOID, - ObjectIdentifier, - OCSPExtensionOID, - OtherNameFormOID, - PublicKeyAlgorithmOID, - SignatureAlgorithmOID, - SubjectInformationAccessOID, -) - -__all__ = [ - "AttributeOID", - "AuthorityInformationAccessOID", - "CRLEntryExtensionOID", - "CertificatePoliciesOID", - "ExtendedKeyUsageOID", - "ExtensionOID", - "NameOID", - "OCSPExtensionOID", - "ObjectIdentifier", - "OtherNameFormOID", - "PublicKeyAlgorithmOID", - "SignatureAlgorithmOID", - "SubjectInformationAccessOID", -] diff --git a/venv/lib/python3.10/site-packages/cryptography/x509/verification.py b/venv/lib/python3.10/site-packages/cryptography/x509/verification.py deleted file mode 100644 index 2db4324d9615fb0d4dce6cd4dc5a9b55f4e28745..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cryptography/x509/verification.py +++ /dev/null @@ -1,34 +0,0 @@ -# This file is dual licensed under the terms of the Apache License, Version -# 2.0, and the BSD License. See the LICENSE file in the root of this repository -# for complete details. - -from __future__ import annotations - -import typing - -from cryptography.hazmat.bindings._rust import x509 as rust_x509 -from cryptography.x509.general_name import DNSName, IPAddress - -__all__ = [ - "ClientVerifier", - "Criticality", - "ExtensionPolicy", - "Policy", - "PolicyBuilder", - "ServerVerifier", - "Store", - "Subject", - "VerificationError", - "VerifiedClient", -] - -Store = rust_x509.Store -Subject = typing.Union[DNSName, IPAddress] -VerifiedClient = rust_x509.VerifiedClient -ClientVerifier = rust_x509.ClientVerifier -ServerVerifier = rust_x509.ServerVerifier -PolicyBuilder = rust_x509.PolicyBuilder -Policy = rust_x509.Policy -ExtensionPolicy = rust_x509.ExtensionPolicy -Criticality = rust_x509.Criticality -VerificationError = rust_x509.VerificationError diff --git a/venv/lib/python3.10/site-packages/cython.py b/venv/lib/python3.10/site-packages/cython.py deleted file mode 100644 index 77f4ce04704e929f07923210cf7a55401abf3d89..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/cython.py +++ /dev/null @@ -1,29 +0,0 @@ -#!/usr/bin/env python - -# -# Cython -- Main Program, generic -# - -try: - from typing import TYPE_CHECKING -except ImportError: - TYPE_CHECKING = False - -if not TYPE_CHECKING and __name__ == '__main__': - - import os - import sys - - # Make sure we import the right Cython - cythonpath, _ = os.path.split(os.path.realpath(__file__)) - sys.path.insert(0, cythonpath) - - from Cython.Compiler.Main import main - main(command_line = 1) - -else: - # Void cython.* directives. - from Cython.Shadow import * - ## and bring in the __version__ - from Cython import __version__ - from Cython import load_ipython_extension diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/METADATA b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/METADATA deleted file mode 100644 index 592ac50132955ef28992b625cdd3d2a091e9d167..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/METADATA +++ /dev/null @@ -1,443 +0,0 @@ -Metadata-Version: 2.4 -Name: deepspeed -Version: 0.17.2 -Summary: DeepSpeed library -Home-page: http://deepspeed.ai -Author: DeepSpeed Team -Author-email: deepspeed-info@microsoft.com -License: Apache Software License 2.0 -Project-URL: Documentation, https://deepspeed.readthedocs.io -Project-URL: Source, https://github.com/deepspeedai/DeepSpeed -Classifier: Programming Language :: Python :: 3.8 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Programming Language :: Python :: 3.10 -Classifier: Programming Language :: Python :: 3.11 -Classifier: Programming Language :: Python :: 3.12 -Description-Content-Type: text/markdown -License-File: LICENSE -Requires-Dist: einops -Requires-Dist: hjson -Requires-Dist: msgpack -Requires-Dist: ninja -Requires-Dist: numpy -Requires-Dist: packaging>=20.0 -Requires-Dist: psutil -Requires-Dist: py-cpuinfo -Requires-Dist: pydantic>=2.0.0 -Requires-Dist: torch -Requires-Dist: tqdm -Provides-Extra: 1bit -Provides-Extra: 1bit-mpi -Requires-Dist: mpi4py; extra == "1bit-mpi" -Provides-Extra: readthedocs -Requires-Dist: autodoc_pydantic>=2.0.0; extra == "readthedocs" -Requires-Dist: docutils<0.18; extra == "readthedocs" -Requires-Dist: hjson; extra == "readthedocs" -Requires-Dist: packaging; extra == "readthedocs" -Requires-Dist: psutil; extra == "readthedocs" -Requires-Dist: py-cpuinfo; extra == "readthedocs" -Requires-Dist: pydantic>=2.0.0; extra == "readthedocs" -Requires-Dist: recommonmark; extra == "readthedocs" -Requires-Dist: sphinx_rtd_theme; extra == "readthedocs" -Requires-Dist: torch; extra == "readthedocs" -Requires-Dist: tqdm; extra == "readthedocs" -Provides-Extra: dev -Requires-Dist: accelerate; extra == "dev" -Requires-Dist: clang-format==18.1.3; extra == "dev" -Requires-Dist: comet_ml>=3.41.0; extra == "dev" -Requires-Dist: deepspeed-kernels; sys_platform == "linux" and extra == "dev" -Requires-Dist: docutils<0.18; extra == "dev" -Requires-Dist: future; extra == "dev" -Requires-Dist: importlib-metadata>=4; extra == "dev" -Requires-Dist: mup; extra == "dev" -Requires-Dist: pre-commit>=3.2.0; extra == "dev" -Requires-Dist: pytest<8.4.0,>=7.2.0; extra == "dev" -Requires-Dist: pytest-forked; extra == "dev" -Requires-Dist: pytest-randomly; extra == "dev" -Requires-Dist: pytest-xdist; extra == "dev" -Requires-Dist: qtorch==0.3.0; extra == "dev" -Requires-Dist: recommonmark; extra == "dev" -Requires-Dist: sphinx; extra == "dev" -Requires-Dist: sphinx-rtd-theme; extra == "dev" -Requires-Dist: tensorboard; extra == "dev" -Requires-Dist: torchvision; extra == "dev" -Requires-Dist: transformers>=4.51.3; extra == "dev" -Requires-Dist: wandb; extra == "dev" -Provides-Extra: autotuning -Requires-Dist: tabulate; extra == "autotuning" -Provides-Extra: autotuning-ml -Requires-Dist: hjson; extra == "autotuning-ml" -Requires-Dist: tabulate; extra == "autotuning-ml" -Requires-Dist: xgboost; extra == "autotuning-ml" -Provides-Extra: sparse-attn -Requires-Dist: triton==1.0.0; extra == "sparse-attn" -Provides-Extra: sparse -Requires-Dist: neural-compressor==2.1.0; extra == "sparse" -Provides-Extra: inf -Requires-Dist: google; extra == "inf" -Requires-Dist: lm-eval==0.3.0; extra == "inf" -Requires-Dist: protobuf; extra == "inf" -Requires-Dist: qtorch; extra == "inf" -Requires-Dist: safetensors; extra == "inf" -Requires-Dist: sentencepiece; extra == "inf" -Requires-Dist: transformers>=4.32.1; extra == "inf" -Provides-Extra: sd -Requires-Dist: diffusers>=0.25.0; extra == "sd" -Requires-Dist: triton>=2.1.0; extra == "sd" -Provides-Extra: triton -Requires-Dist: triton==2.1.0; extra == "triton" -Provides-Extra: deepcompile -Requires-Dist: scipy; extra == "deepcompile" -Provides-Extra: all -Requires-Dist: clang-format==18.1.3; extra == "all" -Requires-Dist: xgboost; extra == "all" -Requires-Dist: pydantic>=2.0.0; extra == "all" -Requires-Dist: comet_ml>=3.41.0; extra == "all" -Requires-Dist: deepspeed-kernels; sys_platform == "linux" and extra == "all" -Requires-Dist: sentencepiece; extra == "all" -Requires-Dist: qtorch; extra == "all" -Requires-Dist: psutil; extra == "all" -Requires-Dist: triton==2.1.0; extra == "all" -Requires-Dist: torch; extra == "all" -Requires-Dist: accelerate; extra == "all" -Requires-Dist: tabulate; extra == "all" -Requires-Dist: pytest-xdist; extra == "all" -Requires-Dist: tqdm; extra == "all" -Requires-Dist: importlib-metadata>=4; extra == "all" -Requires-Dist: py-cpuinfo; extra == "all" -Requires-Dist: pre-commit>=3.2.0; extra == "all" -Requires-Dist: triton==1.0.0; extra == "all" -Requires-Dist: protobuf; extra == "all" -Requires-Dist: lm-eval==0.3.0; extra == "all" -Requires-Dist: wandb; extra == "all" -Requires-Dist: triton>=2.1.0; extra == "all" -Requires-Dist: qtorch==0.3.0; extra == "all" -Requires-Dist: transformers>=4.51.3; extra == "all" -Requires-Dist: sphinx_rtd_theme; extra == "all" -Requires-Dist: safetensors; extra == "all" -Requires-Dist: torchvision; extra == "all" -Requires-Dist: pytest-randomly; extra == "all" -Requires-Dist: mup; extra == "all" -Requires-Dist: sphinx; extra == "all" -Requires-Dist: sphinx-rtd-theme; extra == "all" -Requires-Dist: hjson; extra == "all" -Requires-Dist: docutils<0.18; extra == "all" -Requires-Dist: future; extra == "all" -Requires-Dist: transformers>=4.32.1; extra == "all" -Requires-Dist: diffusers>=0.25.0; extra == "all" -Requires-Dist: pytest<8.4.0,>=7.2.0; extra == "all" -Requires-Dist: neural-compressor==2.1.0; extra == "all" -Requires-Dist: tensorboard; extra == "all" -Requires-Dist: recommonmark; extra == "all" -Requires-Dist: pytest-forked; extra == "all" -Requires-Dist: google; extra == "all" -Requires-Dist: scipy; extra == "all" -Requires-Dist: mpi4py; extra == "all" -Requires-Dist: autodoc_pydantic>=2.0.0; extra == "all" -Requires-Dist: packaging; extra == "all" -Dynamic: author -Dynamic: author-email -Dynamic: classifier -Dynamic: description -Dynamic: description-content-type -Dynamic: home-page -Dynamic: license -Dynamic: license-file -Dynamic: project-url -Dynamic: provides-extra -Dynamic: requires-dist -Dynamic: summary - -[![License Apache 2.0](https://badgen.net/badge/license/apache2.0/blue)](https://github.com/deepspeedai/DeepSpeed/blob/master/LICENSE) -[![PyPI version](https://badge.fury.io/py/deepspeed.svg)](https://pypi.org/project/deepspeed/) -[![Downloads](https://static.pepy.tech/badge/deepspeed)](https://pepy.tech/project/deepspeed) -[![Build](https://badgen.net/badge/build/check-status/blue)](#build-pipeline-status) -[![OpenSSF Best Practices](https://www.bestpractices.dev/projects/9530/badge)](https://www.bestpractices.dev/projects/9530) -[![Twitter](https://img.shields.io/twitter/follow/DeepSpeedAI)](https://twitter.com/intent/follow?screen_name=DeepSpeedAI) -[![Japanese Twitter](https://img.shields.io/badge/%E6%97%A5%E6%9C%AC%E8%AA%9ETwitter-%40DeepSpeedAI_JP-blue)](https://twitter.com/DeepSpeedAI_JP) -[![Chinese Zhihu](https://img.shields.io/badge/%E7%9F%A5%E4%B9%8E-%E5%BE%AE%E8%BD%AFDeepSpeed-blue)](https://www.zhihu.com/people/deepspeed) - - -
- - -
- -## Latest News - DeepSpeed empowers ChatGPT-like model training with a single click, offering 15x speedup over SOTA RLHF systems with unprecedented cost reduction at all scales; [learn how](https://github.com/deepspeedai/DeepSpeed/tree/master/blogs/deepspeed-chat). -* [2025/06] [Arctic Long Sequence Training (ALST) with DeepSpeed: Scalable And Efficient Training For Multi-Million Token Sequences](https://www.snowflake.com/en/engineering-blog/arctic-long-sequence-training-multi-million-token-ai/) -* [2025/04] [DeepCompile: Unlocking Compiler Optimization for Distributed Training](https://github.com/deepspeedai/DeepSpeed/blob/master/blogs/deepcompile/README.md) -* [2025/03] [DeepSpeed-AutoTP: Automatic Tensor Parallel Training of Hugging Face models](https://github.com/deepspeedai/DeepSpeed/blob/master/blogs/huggingface-tp/README.md) -* [2024/12] [Ulysses-Offload: Democratizing Long Context LLM Training ](https://github.com/deepspeedai/DeepSpeed/blob/master/blogs/ulysses-offload/README.md) -* [2024/12] [DeepSpeed-Domino: Communication-Free LLM Training Engine](https://github.com/deepspeedai/DeepSpeed/blob/master/blogs/deepspeed-domino/README.md) -* [2024/08] [DeepSpeed on Windows](https://github.com/deepspeedai/DeepSpeed/tree/master/blogs/windows/08-2024/README.md) [[日本語](https://github.com/deepspeedai/DeepSpeed/tree/master/blogs/windows/08-2024/japanese/README.md)] [[中文](https://github.com/deepspeedai/DeepSpeed/tree/master/blogs/windows/08-2024/chinese/README.md)] - - -
- More news - -
- ---- - -# Extreme Speed and Scale for DL Training and Inference - -***[DeepSpeed](https://www.deepspeed.ai/) enabled the world's most powerful language models (at the time of this writing) such as [MT-530B](https://www.microsoft.com/en-us/research/blog/using-deepspeed-and-megatron-to-train-megatron-turing-nlg-530b-the-worlds-largest-and-most-powerful-generative-language-model/) and [BLOOM](https://huggingface.co/blog/bloom-megatron-deepspeed)***. It is an easy-to-use deep learning optimization software suite that powers unprecedented scale and speed for both training and inference. With DeepSpeed you can: - -* Train/Inference dense or sparse models with billions or trillions of parameters -* Achieve excellent system throughput and efficiently scale to thousands of GPUs -* Train/Inference on resource constrained GPU systems -* Achieve unprecedented low latency and high throughput for inference -* Achieve extreme compression for an unparalleled inference latency and model size reduction with low costs - ---- - -# DeepSpeed's four innovation pillars - - - - -## DeepSpeed-Training - -DeepSpeed offers a confluence of system innovations, that has made large scale DL training effective, and efficient, greatly improved ease of use, and redefined the DL training landscape in terms of scale that is possible. These innovations such as ZeRO, 3D-Parallelism, DeepSpeed-MoE, ZeRO-Infinity, etc. fall under the training pillar. Learn more: [DeepSpeed-Training](https://www.deepspeed.ai/training/) - -## DeepSpeed-Inference - -DeepSpeed brings together innovations in parallelism technology such as tensor, pipeline, expert and ZeRO-parallelism, and combines them with high performance custom inference kernels, communication optimizations and heterogeneous memory technologies to enable inference at an unprecedented scale, while achieving unparalleled latency, throughput and cost reduction. This systematic composition of system technologies for inference falls under the inference pillar. Learn more: [DeepSpeed-Inference](https://www.deepspeed.ai/inference) - - -## DeepSpeed-Compression - -To further increase the inference efficiency, DeepSpeed offers easy-to-use and flexible-to-compose compression techniques for researchers and practitioners to compress their models while delivering faster speed, smaller model size, and significantly reduced compression cost. Moreover, SoTA innovations on compression like ZeroQuant and XTC are included under the compression pillar. Learn more: [DeepSpeed-Compression](https://www.deepspeed.ai/compression) - -## DeepSpeed4Science - -In line with Microsoft's mission to solve humanity's most pressing challenges, the DeepSpeed team at Microsoft is responding to this opportunity by launching a new initiative called *DeepSpeed4Science*, aiming to build unique capabilities through AI system technology innovations to help domain experts to unlock today's biggest science mysteries. Learn more: [tutorials](https://www.deepspeed.ai/deepspeed4science/) - ---- - -# DeepSpeed Software Suite - -## DeepSpeed Library - - The [DeepSpeed](https://github.com/deepspeedai/deepspeed) library (this repository) implements and packages the innovations and technologies in DeepSpeed Training, Inference and Compression Pillars into a single easy-to-use, open-sourced repository. It allows for easy composition of multitude of features within a single training, inference or compression pipeline. The DeepSpeed Library is heavily adopted by the DL community, and has been used to enable some of the most powerful models (see [DeepSpeed Adoption](#deepspeed-adoption)). - -## Model Implementations for Inference (MII) - - [Model Implementations for Inference (MII)](https://github.com/deepspeedai/deepspeed-mii) is an open-sourced repository for making low-latency and high-throughput inference accessible to all data scientists by alleviating the need to apply complex system optimization techniques themselves. Out-of-box, MII offers support for thousands of widely used DL models, optimized using DeepSpeed-Inference, that can be deployed with a few lines of code, while achieving significant latency reduction compared to their vanilla open-sourced versions. - -## DeepSpeed on Azure - - DeepSpeed users are diverse and have access to different environments. We recommend to try DeepSpeed on Azure as it is the simplest and easiest method. The recommended method to try DeepSpeed on Azure is through AzureML [recipes](https://github.com/Azure/azureml-examples/tree/main/v1/python-sdk/workflows/train/deepspeed). The job submission and data preparation scripts have been made available [here](https://github.com/deepspeedai/Megatron-DeepSpeed/tree/main/examples_deepspeed/azureml). For more details on how to use DeepSpeed on Azure, please follow the [Azure tutorial](https://www.deepspeed.ai/tutorials/azure/). - ---- - -# DeepSpeed Adoption - -DeepSpeed was an important part of Microsoft’s -[AI at Scale](https://www.microsoft.com/en-us/research/project/ai-at-scale/) -initiative to enable next-generation AI capabilities at scale, where you can find more -information [here](https://innovation.microsoft.com/en-us/exploring-ai-at-scale). - -DeepSpeed has been used to train many different large-scale models, below is a list of several examples that we are aware of (if you'd like to include your model please submit a PR): - - * [Megatron-Turing NLG (530B)](https://www.microsoft.com/en-us/research/blog/using-deepspeed-and-megatron-to-train-megatron-turing-nlg-530b-the-worlds-largest-and-most-powerful-generative-language-model/) - * [Jurassic-1 (178B)](https://uploads-ssl.webflow.com/60fd4503684b466578c0d307/61138924626a6981ee09caf6_jurassic_tech_paper.pdf) - * [BLOOM (176B)](https://huggingface.co/blog/bloom-megatron-deepspeed) - * [GLM (130B)](https://github.com/THUDM/GLM-130B) - * [xTrimoPGLM (100B)](https://www.biorxiv.org/content/10.1101/2023.07.05.547496v2) - * [YaLM (100B)](https://github.com/yandex/YaLM-100B) - * [GPT-NeoX (20B)](https://github.com/EleutherAI/gpt-neox) - * [AlexaTM (20B)](https://www.amazon.science/blog/20b-parameter-alexa-model-sets-new-marks-in-few-shot-learning) - * [Turing NLG (17B)](https://www.microsoft.com/en-us/research/blog/turing-nlg-a-17-billion-parameter-language-model-by-microsoft/) - * [METRO-LM (5.4B)](https://arxiv.org/pdf/2204.06644.pdf) - -DeepSpeed has been integrated with several different popular open-source DL frameworks such as: - -| | Documentation | -| ---------------------------------------------------------------------------------------------- | -------------------------------------------- | - | [Transformers with DeepSpeed](https://huggingface.co/docs/transformers/deepspeed) | -| | [Accelerate with DeepSpeed](https://huggingface.co/docs/accelerate/usage_guides/deepspeed) | -| | [Lightning with DeepSpeed](https://lightning.ai/docs/pytorch/stable/advanced/model_parallel.html#deepspeed) | -| | [MosaicML with DeepSpeed](https://docs.mosaicml.com/projects/composer/en/latest/trainer/using_the_trainer.html?highlight=deepspeed#deepspeed-integration) | -| | [Determined with DeepSpeed](https://docs.determined.ai/latest/training/apis-howto/deepspeed/overview.html) | -| | [MMEngine with DeepSpeed](https://mmengine.readthedocs.io/en/latest/common_usage/large_model_training.html#deepspeed) | - ---- - -# Build Pipeline Status - -| Description | Status | -| ----------- | ------ | -| NVIDIA | [![nv-torch110-p40](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch110-p40.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch110-p40.yml) [![nv-torch110-v100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch110-v100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch110-v100.yml) [![nv-torch-latest-v100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch-latest-v100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch-latest-v100.yml) [![nv-h100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-h100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-h100.yml) [![nv-inference](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-inference.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-inference.yml) [![nv-nightly](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-nightly.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-nightly.yml) | -| AMD | [![amd-mi200](https://github.com/deepspeedai/DeepSpeed/actions/workflows/amd-mi200.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/amd-mi200.yml) | -| CPU | [![torch-latest-cpu](https://github.com/deepspeedai/DeepSpeed/actions/workflows/cpu-torch-latest.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/cpu-torch-latest.yml) [![cpu-inference](https://github.com/deepspeedai/DeepSpeed/actions/workflows/cpu-inference.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/cpu-inference.yml) | -| Intel Gaudi | [![hpu-gaudi2](https://github.com/deepspeedai/DeepSpeed/actions/workflows/hpu-gaudi2.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/hpu-gaudi2.yml) | -| Intel XPU | [![xpu-max1100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/xpu-max1100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/xpu-max1100.yml) | -| PyTorch Nightly | [![nv-torch-nightly-v100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch-nightly-v100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-torch-nightly-v100.yml) | -| Integrations | [![nv-transformers-v100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-transformers-v100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-transformers-v100.yml) [![nv-lightning-v100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-lightning-v100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-lightning-v100.yml) [![nv-accelerate-v100](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-accelerate-v100.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-accelerate-v100.yml) [![nv-mii](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-mii.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-mii.yml) [![nv-ds-chat](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-ds-chat.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-ds-chat.yml) [![nv-sd](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-sd.yml/badge.svg)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/nv-sd.yml) | -| Misc | [![Formatting](https://github.com/deepspeedai/DeepSpeed/actions/workflows/formatting.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/formatting.yml) [![pages-build-deployment](https://github.com/deepspeedai/DeepSpeed/actions/workflows/pages/pages-build-deployment/badge.svg)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/pages/pages-build-deployment) [![Documentation Status](https://readthedocs.org/projects/deepspeed/badge/?version=latest)](https://deepspeed.readthedocs.io/en/latest/?badge=latest)[![python](https://github.com/deepspeedai/DeepSpeed/actions/workflows/python.yml/badge.svg?branch=master)](https://github.com/deepspeedai/DeepSpeed/actions/workflows/python.yml) | -| Huawei Ascend NPU | [![Huawei Ascend NPU](https://github.com/Ascend/Ascend-CI/actions/workflows/deepspeed.yaml/badge.svg?branch=main)](https://github.com/Ascend/Ascend-CI/actions/workflows/deepspeed.yaml) | - -# Installation - -The quickest way to get started with DeepSpeed is via pip, this will install -the latest release of DeepSpeed which is not tied to specific PyTorch or CUDA -versions. DeepSpeed includes several C++/CUDA extensions that we commonly refer -to as our 'ops'. By default, all of these extensions/ops will be built -just-in-time (JIT) using [torch's JIT C++ extension loader that relies on -ninja](https://pytorch.org/docs/stable/cpp_extension.html) to build and -dynamically link them at runtime. - -## Requirements -* [PyTorch](https://pytorch.org/) must be installed _before_ installing DeepSpeed. -* For full feature support we recommend a version of PyTorch that is >= 1.9 and ideally the latest PyTorch stable release. -* A CUDA or ROCm compiler such as [nvcc](https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/#introduction) or [hipcc](https://github.com/ROCm-Developer-Tools/HIPCC) used to compile C++/CUDA/HIP extensions. -* Specific GPUs we develop and test against are listed below, this doesn't mean your GPU will not work if it doesn't fall into this category it's just DeepSpeed is most well tested on the following: - * NVIDIA: Pascal, Volta, Ampere, and Hopper architectures - * AMD: MI100 and MI200 - -## Contributed HW support -* DeepSpeed now support various HW accelerators. - -| Contributor | Hardware | Accelerator Name | Contributor validated | Upstream validated | -|-------------|-------------------------------------|------------------| --------------------- |--------------------| -| Huawei | Huawei Ascend NPU | npu | Yes | No | -| Intel | Intel(R) Gaudi(R) 2 AI accelerator | hpu | Yes | Yes | -| Intel | Intel(R) Xeon(R) Processors | cpu | Yes | Yes | -| Intel | Intel(R) Data Center GPU Max series | xpu | Yes | Yes | -| Tecorigin | Scalable Data Analytics Accelerator | sdaa | Yes | No | - -## PyPI -We regularly push releases to [PyPI](https://pypi.org/project/deepspeed/) and encourage users to install from there in most cases. - -```bash -pip install deepspeed -``` - -After installation, you can validate your install and see which extensions/ops -your machine is compatible with via the DeepSpeed environment report. - -```bash -ds_report -``` - -If you would like to pre-install any of the DeepSpeed extensions/ops (instead -of JIT compiling) or install pre-compiled ops via PyPI please see our [advanced -installation instructions](https://www.deepspeed.ai/tutorials/advanced-install/). - -## Windows -Many DeepSpeed features are supported on Windows for both training and inference. You can read more about this in the original blog post [here](https://github.com/deepspeedai/DeepSpeed/tree/master/blogs/windows/08-2024/README.md). Among features that are currently not supported are async io (AIO) and GDS (which does not support Windows). -1. Install PyTorch, such as pytorch 2.3+cu121. -2. Install Visual C++ build tools, such as VS2022 C++ x64/x86 build tools. -3. Launch Cmd console with Administrator permissions for creating required symlink folders and ensure MSVC tools are added to your PATH or launch the Developer Command Prompt for Visual Studio 2022 with administrator permissions. -4. Run `build_win.bat` to build wheel in `dist` folder. - -# Features - -Please checkout [DeepSpeed-Training](https://www.deepspeed.ai/training), [DeepSpeed-Inference](https://www.deepspeed.ai/inference) and [DeepSpeed-Compression](https://www.deepspeed.ai/compression) pages for full set of features offered along each of these three pillars. - -# Further Reading - -All DeepSpeed documentation, tutorials, and blogs can be found on our website: [deepspeed.ai](https://www.deepspeed.ai/) - - -| | Description | -| ---------------------------------------------------------------------------------------------- | -------------------------------------------- | -| [Getting Started](https://www.deepspeed.ai/getting-started/) | First steps with DeepSpeed | -| [DeepSpeed JSON Configuration](https://www.deepspeed.ai/docs/config-json/) | Configuring DeepSpeed | -| [API Documentation](https://deepspeed.readthedocs.io/en/latest/) | Generated DeepSpeed API documentation | -| [Tutorials](https://www.deepspeed.ai/tutorials/) | Tutorials | -| [Blogs](https://www.deepspeed.ai/posts/) | Blogs | - - -# Contributing -DeepSpeed welcomes your contributions! Please see our -[contributing](CONTRIBUTING.md) guide for more details on formatting, testing, -etc.
-Thanks so much to all of our amazing contributors! - - - - - -## Contributor License Agreement -This project welcomes contributions and suggestions. Most contributions require you to -agree to a Contributor License Agreement (CLA) declaring that you have the right to, and -actually do, grant us the rights to use your contribution. For details, visit -https://cla.opensource.microsoft.com. - -When you submit a pull request, a CLA bot will automatically determine whether you need -to provide a CLA and decorate the PR appropriately (e.g., status check, comment). Simply -follow the instructions provided by the bot. You will only need to do this once across -all repos using our CLA. - -## Code of Conduct -This project has adopted the [Microsoft Open Source Code of -Conduct](https://opensource.microsoft.com/codeofconduct/). For more information see the -[Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or contact -[opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments. - -# Publications -1. Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He. (2019) ZeRO: memory optimizations toward training trillion parameter models. [arXiv:1910.02054](https://arxiv.org/abs/1910.02054) and [In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC '20)](https://dl.acm.org/doi/10.5555/3433701.3433727). -2. Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. (2020) DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters. [In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining (KDD '20, Tutorial)](https://dl.acm.org/doi/10.1145/3394486.3406703). -3. Minjia Zhang, Yuxiong He. (2020) Accelerating Training of Transformer-Based Language Models with Progressive Layer Dropping. [arXiv:2010.13369](https://arxiv.org/abs/2010.13369) and [NeurIPS 2020](https://proceedings.neurips.cc/paper/2020/hash/a1140a3d0df1c81e24ae954d935e8926-Abstract.html). -4. Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, Yuxiong He. (2021) ZeRO-Offload: Democratizing Billion-Scale Model Training. [arXiv:2101.06840](https://arxiv.org/abs/2101.06840) and [USENIX ATC 2021](https://www.usenix.org/conference/atc21/presentation/ren-jie). [[paper]](https://arxiv.org/abs/2101.06840) [[slides]](https://www.usenix.org/system/files/atc21_slides_ren-jie.pdf) [[blog]](https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/) -5. Hanlin Tang, Shaoduo Gan, Ammar Ahmad Awan, Samyam Rajbhandari, Conglong Li, Xiangru Lian, Ji Liu, Ce Zhang, Yuxiong He. (2021) 1-bit Adam: Communication Efficient Large-Scale Training with Adam's Convergence Speed. [arXiv:2102.02888](https://arxiv.org/abs/2102.02888) and [ICML 2021](http://proceedings.mlr.press/v139/tang21a.html). -6. Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, Yuxiong He. (2021) ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning. [arXiv:2104.07857](https://arxiv.org/abs/2104.07857) and [SC 2021](https://dl.acm.org/doi/abs/10.1145/3458817.3476205). [[paper]](https://arxiv.org/abs/2104.07857) [[slides]](docs/assets/files/SC21-ZeRO-Infinity.pdf) [[blog]](https://www.microsoft.com/en-us/research/blog/zero-infinity-and-deepspeed-unlocking-unprecedented-model-scale-for-deep-learning-training/) -7. Conglong Li, Ammar Ahmad Awan, Hanlin Tang, Samyam Rajbhandari, Yuxiong He. (2021) 1-bit LAMB: Communication Efficient Large-Scale Large-Batch Training with LAMB's Convergence Speed. [arXiv:2104.06069](https://arxiv.org/abs/2104.06069) and [HiPC 2022](https://hipc.org/advance-program/). -8. Conglong Li, Minjia Zhang, Yuxiong He. (2021) The Stability-Efficiency Dilemma: Investigating Sequence Length Warmup for Training GPT Models. [arXiv:2108.06084](https://arxiv.org/abs/2108.06084) and [NeurIPS 2022](https://openreview.net/forum?id=JpZ5du_Kdh). -9. Yucheng Lu, Conglong Li, Minjia Zhang, Christopher De Sa, Yuxiong He. (2022) Maximizing Communication Efficiency for Large-scale Training via 0/1 Adam. [arXiv:2202.06009](https://arxiv.org/abs/2202.06009). -10. Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, Yuxiong He. (2022) DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale [arXiv:2201.05596](https://arxiv.org/abs/2201.05596) and [ICML 2022](https://proceedings.mlr.press/v162/rajbhandari22a.html). [[pdf]](https://arxiv.org/abs/2201.05596) [[slides]](docs/assets/files/ICML-5mins.pdf) [[blog]](https://www.microsoft.com/en-us/research/blog/deepspeed-advancing-moe-inference-and-training-to-power-next-generation-ai-scale/) -11. Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, Elton Zhang, Rewon Child, Reza Yazdani Aminabadi, Julie Bernauer, Xia Song, Mohammad Shoeybi, Yuxiong He, Michael Houston, Saurabh Tiwary, Bryan Catanzaro. (2022) Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model [arXiv:2201.11990](https://arxiv.org/abs/2201.11990). -12. Xiaoxia Wu, Zhewei Yao, Minjia Zhang, Conglong Li, Yuxiong He. (2022) Extreme Compression for Pre-trained Transformers Made Simple and Efficient. [arXiv:2206.01859](https://arxiv.org/abs/2206.01859) and [NeurIPS 2022](https://openreview.net/forum?id=xNeAhc2CNAl). -13. Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, Yuxiong He. (2022) ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers. [arXiv:2206.01861](https://arxiv.org/abs/2206.01861) and [NeurIPS 2022](https://openreview.net/forum?id=f-fVCElZ-G1) [[slides]](docs/assets/files/zeroquant_series.pdf) [[blog]](https://www.microsoft.com/en-us/research/blog/deepspeed-compression-a-composable-library-for-extreme-compression-and-zero-cost-quantization/) -14. Reza Yazdani Aminabadi, Samyam Rajbhandari, Minjia Zhang, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Jeff Rasley, Shaden Smith, Olatunji Ruwase, Yuxiong He. (2022) DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale. [arXiv:2207.00032](https://arxiv.org/abs/2207.00032) and [SC 2022](https://dl.acm.org/doi/abs/10.5555/3571885.3571946). [[paper]](https://arxiv.org/abs/2207.00032) [[slides]](docs/assets/files/sc22-ds-inference.pdf) [[blog]](https://www.microsoft.com/en-us/research/blog/deepspeed-accelerating-large-scale-model-inference-and-training-via-system-optimizations-and-compression/) -15. Zhewei Yao, Xiaoxia Wu, Conglong Li, Connor Holmes, Minjia Zhang, Cheng Li, Yuxiong He. (2022) Random-LTD: Random and Layerwise Token Dropping Brings Efficient Training for Large-scale Transformers. [arXiv:2211.11586](https://arxiv.org/abs/2211.11586). -16. Conglong Li, Zhewei Yao, Xiaoxia Wu, Minjia Zhang, Yuxiong He. (2022) DeepSpeed Data Efficiency: Improving Deep Learning Model Quality and Training Efficiency via Efficient Data Sampling and Routing. [arXiv:2212.03597](https://arxiv.org/abs/2212.03597) [ENLSP2023 Workshop at NeurIPS2023](https://neurips2023-enlsp.github.io/) -17. Xiaoxia Wu, Cheng Li, Reza Yazdani Aminabadi, Zhewei Yao, Yuxiong He. (2023) Understanding INT4 Quantization for Transformer Models: Latency Speedup, Composability, and Failure Cases. [arXiv:2301.12017](https://arxiv.org/abs/2301.12017) and [ICML2023](https://icml.cc/Conferences/2023). -18. Syed Zawad, Cheng Li, Zhewei Yao, Elton Zheng, Yuxiong He, Feng Yan. (2023) DySR: Adaptive Super-Resolution via Algorithm and System Co-design. [ICLR:2023](https://openreview.net/forum?id=Pgtn4l6eKjv). -19. Sheng Shen, Zhewei Yao, Chunyuan Li, Trevor Darrell, Kurt Keutzer, Yuxiong He. (2023) Scaling Vision-Language Models with Sparse Mixture of Experts. [arXiv:2303.07226](https://arxiv.org/abs/2303.07226) and [Finding at EMNLP2023](https://2023.emnlp.org/). -20. Quentin Anthony, Ammar Ahmad Awan, Jeff Rasley, Yuxiong He, Aamir Shafi, Mustafa Abduljabbar, Hari Subramoni, Dhabaleswar Panda. (2023) MCR-DL: Mix-and-Match Communication Runtime for Deep Learning [arXiv:2303.08374](https://arxiv.org/abs/2303.08374) and will appear at IPDPS 2023. -21. Siddharth Singh, Olatunji Ruwase, Ammar Ahmad Awan, Samyam Rajbhandari, Yuxiong He, Abhinav Bhatele. (2023) A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training [arXiv:2303.06318](https://arxiv.org/abs/2303.06318) and [ICS 2023](https://dl.acm.org/doi/10.1145/3577193.3593704). -22. Guanhua Wang, Heyang Qin, Sam Ade Jacobs, Xiaoxia Wu, Connor Holmes, Zhewei Yao, Samyam Rajbhandari, Olatunji Ruwase, Feng Yan, Lei Yang, Yuxiong He. (2023) ZeRO++: Extremely Efficient Collective Communication for Giant Model Training [arXiv:2306.10209](https://arxiv.org/abs/2306.10209) and [ML for Sys Workshop at NeurIPS2023](http://mlforsystems.org/) [[blog]](https://www.microsoft.com/en-us/research/blog/deepspeed-zero-a-leap-in-speed-for-llm-and-chat-model-training-with-4x-less-communication/) -23. Zhewei Yao, Xiaoxia Wu, Cheng Li, Stephen Youn, Yuxiong He. (2023) ZeroQuant-V2: Exploring Post-training Quantization in LLMs from Comprehensive Study to Low Rank Compensation [arXiv:2303.08302](https://arxiv.org/abs/2303.08302) and [ENLSP2023 Workshop at NeurIPS2023](https://neurips2023-enlsp.github.io/) [[slides]](docs/assets/files/zeroquant_series.pdf) -24. Pareesa Ameneh Golnari, Zhewei Yao, Yuxiong He. (2023) Selective Guidance: Are All the Denoising Steps of Guided Diffusion Important? [arXiv:2305.09847](https://arxiv.org/abs/2305.09847) -25. Zhewei Yao, Reza Yazdani Aminabadi, Olatunji Ruwase, Samyam Rajbhandari, Xiaoxia Wu, Ammar Ahmad Awan, Jeff Rasley, Minjia Zhang, Conglong Li, Connor Holmes, Zhongzhu Zhou, Michael Wyatt, Molly Smith, Lev Kurilenko, Heyang Qin, Masahiro Tanaka, Shuai Che, Shuaiwen Leon Song, Yuxiong He. (2023) DeepSpeed-Chat: Easy, Fast and Affordable RLHF Training of ChatGPT-like Models at All Scales [arXiv:2308.01320](https://arxiv.org/abs/2308.01320). -26. Xiaoxia Wu, Zhewei Yao, Yuxiong He. (2023) ZeroQuant-FP: A Leap Forward in LLMs Post-Training W4A8 Quantization Using Floating-Point Formats [arXiv:2307.09782](https://arxiv.org/abs/2307.09782) and [ENLSP2023 Workshop at NeurIPS2023](https://neurips2023-enlsp.github.io/) [[slides]](docs/assets/files/zeroquant_series.pdf) -27. Zhewei Yao, Xiaoxia Wu, Conglong Li, Minjia Zhang, Heyang Qin, Olatunji Ruwase, Ammar Ahmad Awan, Samyam Rajbhandari, Yuxiong He. (2023) DeepSpeed-VisualChat: Multi-Round Multi-Image Interleave Chat via Multi-Modal Causal Attention [arXiv:2309.14327](https://arxiv.org/pdf/2309.14327.pdf) -28. Shuaiwen Leon Song, Bonnie Kruft, Minjia Zhang, Conglong Li, Shiyang Chen, Chengming Zhang, Masahiro Tanaka, Xiaoxia Wu, Jeff Rasley, Ammar Ahmad Awan, Connor Holmes, Martin Cai, Adam Ghanem, Zhongzhu Zhou, Yuxiong He, et al. (2023) DeepSpeed4Science Initiative: Enabling Large-Scale Scientific Discovery through Sophisticated AI System Technologies [arXiv:2310.04610](https://arxiv.org/abs/2310.04610) [[blog]](https://www.microsoft.com/en-us/research/blog/announcing-the-deepspeed4science-initiative-enabling-large-scale-scientific-discovery-through-sophisticated-ai-system-technologies/) -29. Zhewei Yao, Reza Yazdani Aminabadi, Stephen Youn, Xiaoxia Wu, Elton Zheng, Yuxiong He. (2023) ZeroQuant-HERO: Hardware-Enhanced Robust Optimized Post-Training Quantization Framework for W8A8 Transformers [arXiv:2310.17723](https://arxiv.org/abs/2310.17723) - -30. Xiaoxia Wu, Haojun Xia, Stephen Youn, Zhen Zheng, Shiyang Chen, Arash Bakhtiari, Michael Wyatt, Reza Yazdani Aminabadi, Yuxiong He, Olatunji Ruwase, Leon Song, Zhewei Yao (2023) ZeroQuant(4+2): Redefining LLMs Quantization with a New FP6-Centric Strategy for Diverse Generative Tasks [arXiv:2312.08583](https://arxiv.org/abs/2312.08583) - -31. Haojun Xia, Zhen Zheng, Xiaoxia Wu, Shiyang Chen, Zhewei Yao, Stephen Youn, Arash Bakhtiari, Michael Wyatt, Donglin Zhuang, Zhongzhu Zhou, Olatunji Ruwase, Yuxiong He, Shuaiwen Leon Song. (2024) FP6-LLM: Efficiently Serving Large Language Models Through FP6-Centric Algorithm-System Co-Design [arXiv:2401.14112](https://arxiv.org/abs/2401.14112) -32. Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Reza Yazdani Aminadabi, Shuaiwen Leon Song, Samyam Rajbhandari, Yuxiong He. (2024) [System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://dl.acm.org/doi/10.1145/3662158.3662806) -33. Xinyu Lian, Sam Ade Jacobs, Lev Kurilenko, Masahiro Tanaka, Stas Bekman, Olatunji Ruwase, Minjia Zhang. (2024) Universal Checkpointing: Efficient and Flexible Checkpointing for Large Scale Distributed Training [arXiv:2406.18820](https://arxiv.org/abs/2406.18820) -34. Stas Bekman, Samyam Rajbhandari, Michael Wyatt, Jeff Rasley, Tunji Ruwase, Zhewei Yao, Aurick Qiao, Yuxiong He. (2025) Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences [arXiv:2506.13996](https://arxiv.org/abs/2506.13996) - - -# Videos -1. DeepSpeed KDD 2020 Tutorial - 1. [Overview](https://www.youtube.com/watch?v=CaseqC45DNc&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=29) - 2. [ZeRO + large model training](https://www.youtube.com/watch?v=y4_bCiAsIAk&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=28) - 3. [17B T-NLG demo](https://www.youtube.com/watch?v=9V-ZbP92drg&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=27) - 4. [Fastest BERT training + RScan tuning](https://www.youtube.com/watch?v=o1K-ZG9F6u0&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=26) - 5. DeepSpeed hands on deep dive: [part 1](https://www.youtube.com/watch?v=_NOk-mBwDYg&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=92), [part 2](https://www.youtube.com/watch?v=sG6_c4VXLww&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=94), [part 3](https://www.youtube.com/watch?v=k9yPkBTayos&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=93) - 6. [FAQ](https://www.youtube.com/watch?v=nsHu6vEgPew&list=PLa85ZdUjfWS21mgibJ2vCvLziprjpKoW0&index=24) -2. Microsoft Research Webinar - * Registration is free and all videos are available on-demand. - * [ZeRO & Fastest BERT: Increasing the scale and speed of deep learning training in DeepSpeed](https://note.microsoft.com/MSR-Webinar-DeepSpeed-Registration-On-Demand.html). -3. [DeepSpeed on AzureML](https://youtu.be/yBVXR8G8Bg8) -4. [Large Model Training and Inference with DeepSpeed // Samyam Rajbhandari // LLMs in Prod Conference](https://www.youtube.com/watch?v=cntxC3g22oU) [[slides]](docs/assets/files/presentation-mlops.pdf) -5. Community Tutorials - * [DeepSpeed: All the tricks to scale to gigantic models (Mark Saroufim)](https://www.youtube.com/watch?v=pDGI668pNg0) - * [Turing-NLG, DeepSpeed and the ZeRO optimizer (Yannic Kilcher)](https://www.youtube.com/watch?v=tC01FRB0M7w) - * [Ultimate Guide To Scaling ML Models (The AI Epiphany)](https://www.youtube.com/watch?v=hc0u4avAkuM) diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/RECORD b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/RECORD deleted file mode 100644 index 0eb7d39a8a0d89ad563c4f08aefeb2e0c1511303..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/RECORD +++ /dev/null @@ -1,1580 +0,0 @@ -../../../bin/deepspeed,sha256=fVvGbO4f9l-ArUn5Ic-zdG_onz7cAAiR6zp4IMhwP0w,120 -../../../bin/deepspeed.pt,sha256=fVvGbO4f9l-ArUn5Ic-zdG_onz7cAAiR6zp4IMhwP0w,120 -../../../bin/ds,sha256=fVvGbO4f9l-ArUn5Ic-zdG_onz7cAAiR6zp4IMhwP0w,120 -../../../bin/ds_bench,sha256=NCeQ5VdNlkSZRWIC2aQcrjIhd4Wz5OEEA9P-ibHbbec,870 -../../../bin/ds_elastic,sha256=9KRAtw2socyGIPD2PyG1hyHxeLAejtiJF0RIKVoAERA,1883 -../../../bin/ds_io,sha256=SSadr7J4RwFc3drWWJ6Eyvt5fon_WwNff5mwunNAN0Q,121 -../../../bin/ds_nvme_tune,sha256=4MaRWLsmGkRMji03UQB4eK5inNdzlfjrtEhGVxTB6Ec,299 -../../../bin/ds_report,sha256=SA4zydPPRa4xRemY_nRR1yqbFjSoYKdUF7_vQuksZBU,123 -../../../bin/ds_ssh,sha256=BT6cdZ47ceJiKPK40eVf9NCXF5RLjayEwWlyZ7viWbE,680 -../../../bin/dsr,sha256=SA4zydPPRa4xRemY_nRR1yqbFjSoYKdUF7_vQuksZBU,123 -deepspeed-0.17.2.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -deepspeed-0.17.2.dist-info/METADATA,sha256=cq8vc6YI5SGW6twYehsWxaN1IauzMrlCSpIx1nsE-z4,43003 -deepspeed-0.17.2.dist-info/RECORD,, -deepspeed-0.17.2.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -deepspeed-0.17.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91 -deepspeed-0.17.2.dist-info/entry_points.txt,sha256=btWVtkuERrrqyxQojUVMJZOeaoNPz0pjiZKaJ_Q4MJ8,84 -deepspeed-0.17.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357 -deepspeed-0.17.2.dist-info/top_level.txt,sha256=VK6WJ48PiB2ke4oVIo2mLzqXYqGXdu4MPE2GP-vfvA8,10 -deepspeed/__init__.py,sha256=tcS3eXx2HvMLr38_ofg-2RC9l9k1s3AsYX-TyJEBrJU,16832 -deepspeed/__pycache__/__init__.cpython-310.pyc,, -deepspeed/__pycache__/constants.cpython-310.pyc,, -deepspeed/__pycache__/env_report.cpython-310.pyc,, -deepspeed/__pycache__/git_version_info.cpython-310.pyc,, -deepspeed/__pycache__/git_version_info_installed.cpython-310.pyc,, -deepspeed/accelerator/__init__.py,sha256=1pNnvGX9h418ekjY_Xnpb363tLo-5Szw3icQf9hRvk0,248 -deepspeed/accelerator/__pycache__/__init__.cpython-310.pyc,, -deepspeed/accelerator/__pycache__/abstract_accelerator.cpython-310.pyc,, -deepspeed/accelerator/__pycache__/cpu_accelerator.cpython-310.pyc,, -deepspeed/accelerator/__pycache__/cuda_accelerator.cpython-310.pyc,, -deepspeed/accelerator/__pycache__/hpu_accelerator.cpython-310.pyc,, -deepspeed/accelerator/__pycache__/mlu_accelerator.cpython-310.pyc,, -deepspeed/accelerator/__pycache__/mps_accelerator.cpython-310.pyc,, -deepspeed/accelerator/__pycache__/npu_accelerator.cpython-310.pyc,, -deepspeed/accelerator/__pycache__/real_accelerator.cpython-310.pyc,, -deepspeed/accelerator/__pycache__/sdaa_accelerator.cpython-310.pyc,, -deepspeed/accelerator/__pycache__/xpu_accelerator.cpython-310.pyc,, -deepspeed/accelerator/abstract_accelerator.py,sha256=mF-Km8bgg_yGqwcoPcSheKRlYoyFP4i1cFzxUHRL72o,5840 -deepspeed/accelerator/cpu_accelerator.py,sha256=Mn58Mq3eZJr5QWmQliXOY3-D5c1ne-uzmYgeiz1oexM,10553 -deepspeed/accelerator/cuda_accelerator.py,sha256=mwj2tI4736nG1XD3C3p0OqCIpUKt1SzEir1PmDaPVtA,12920 -deepspeed/accelerator/hpu_accelerator.py,sha256=U8MLbGI_qP6PnTwP2F7ibzUalB0-MxnhsJr9Ld0F3BA,10971 -deepspeed/accelerator/mlu_accelerator.py,sha256=1tUxv8mnAHOfo4AY4BxiQ4rCD8KULFKZVeNRULOQNB4,9337 -deepspeed/accelerator/mps_accelerator.py,sha256=XkqWL7P86KFp2ebtGSNF6pPsR5OPsFm1Jl5UPEcR6C4,7006 -deepspeed/accelerator/npu_accelerator.py,sha256=FwTKiHZeaWW5PngtqZYFS49w5hKFmQwhsOPR1KreX2s,9020 -deepspeed/accelerator/real_accelerator.py,sha256=CJS2JocHzQCeGJRGQqxpGdzgLJBHqzR31PDB_SaFxN0,12791 -deepspeed/accelerator/sdaa_accelerator.py,sha256=YXdjlEwIUYuocPsu-0_etZH-2ak4yCNQYAFGL-_DKgM,10813 -deepspeed/accelerator/xpu_accelerator.py,sha256=AopBKrqyTu4_GLr1LpTeH075PUQ2iparLBU56ceDng8,10936 -deepspeed/autotuning/__init__.py,sha256=y0O9XDcX76E6dmhNFmRhy9TKfl0ywJrdkXmM2JuzErU,129 -deepspeed/autotuning/__pycache__/__init__.cpython-310.pyc,, -deepspeed/autotuning/__pycache__/autotuner.cpython-310.pyc,, -deepspeed/autotuning/__pycache__/config.cpython-310.pyc,, -deepspeed/autotuning/__pycache__/constants.cpython-310.pyc,, -deepspeed/autotuning/__pycache__/scheduler.cpython-310.pyc,, -deepspeed/autotuning/__pycache__/utils.cpython-310.pyc,, -deepspeed/autotuning/autotuner.py,sha256=_ucGsohYkJ1MFsz1JuB9DJyky-brugx_IDz0mbhqeY0,54335 -deepspeed/autotuning/config.py,sha256=pg5OeryVqrg32xak1NLEWaFiN1MYujdwlSgTuAz3Xl0,4633 -deepspeed/autotuning/config_templates/template_zero0.json,sha256=hR1baTMA5HzfTD1JlXjiTYjMtB9C8CFfxIBzqZeW2Sw,48 -deepspeed/autotuning/config_templates/template_zero1.json,sha256=AAIox-1dviiYg-Z7L-02bxItSnaggAlHkR4TekXXEIw,113 -deepspeed/autotuning/config_templates/template_zero2.json,sha256=FfrPKZHzMJuHn_TVJ7UZ-2s9yJ5IzBIKgT7nLYnIHh4,237 -deepspeed/autotuning/config_templates/template_zero3.json,sha256=WvsoMMp5MDgdVLq1RX1fhKsKbGSHQtpplZzGJikZvzs,485 -deepspeed/autotuning/constants.py,sha256=sr02-8B5ORdimK5VZ1N9FK1luWadmECFLSlMnM-uIxA,5943 -deepspeed/autotuning/scheduler.py,sha256=4IjLO56K-ZXr-HK1FuSuw5-v4tUdAWRZfH8R02WNx3U,15703 -deepspeed/autotuning/tuner/__init__.py,sha256=c9ImdL2iEc89lFltZ2PPe2EOUCuB3Olq3hRP1Q-yUMo,235 -deepspeed/autotuning/tuner/__pycache__/__init__.cpython-310.pyc,, -deepspeed/autotuning/tuner/__pycache__/base_tuner.cpython-310.pyc,, -deepspeed/autotuning/tuner/__pycache__/cost_model.cpython-310.pyc,, -deepspeed/autotuning/tuner/__pycache__/index_based_tuner.cpython-310.pyc,, -deepspeed/autotuning/tuner/__pycache__/model_based_tuner.cpython-310.pyc,, -deepspeed/autotuning/tuner/__pycache__/utils.cpython-310.pyc,, -deepspeed/autotuning/tuner/base_tuner.py,sha256=psA1I4-AEfqGrGLvYYUl3EjM8jG8AlhUXz-AnrFZ7G4,2754 -deepspeed/autotuning/tuner/cost_model.py,sha256=Uu9jD65cvdUK6aozm6Du3XJQZ_W-yOyUrSWFLdOT0do,1820 -deepspeed/autotuning/tuner/index_based_tuner.py,sha256=AEkTByT3XXCyYolAqBiMquV6XHN-ntGmYKMLZKPCRYE,1158 -deepspeed/autotuning/tuner/model_based_tuner.py,sha256=f_CrbgZNONRjdMuRhrJkQbhtih6GOg4LdkXqRDRDzaM,5614 -deepspeed/autotuning/tuner/utils.py,sha256=o5nD51Z6LBylJDRnDNIiOARlBHnOgINBVjBos6ypLEg,2329 -deepspeed/autotuning/utils.py,sha256=d4ePhVxVtekB7gIEK1dIOs6bhxOpmFF8yhKPC_rNRDc,15045 -deepspeed/checkpoint/__init__.py,sha256=3qhX6swvx0-zo72xLl99eZzBr_KC7PguCVCgemYyAGw,576 -deepspeed/checkpoint/__pycache__/__init__.cpython-310.pyc,, -deepspeed/checkpoint/__pycache__/constants.cpython-310.pyc,, -deepspeed/checkpoint/__pycache__/deepspeed_checkpoint.cpython-310.pyc,, -deepspeed/checkpoint/__pycache__/ds_to_universal.cpython-310.pyc,, -deepspeed/checkpoint/__pycache__/reshape_3d_utils.cpython-310.pyc,, -deepspeed/checkpoint/__pycache__/reshape_meg_2d.cpython-310.pyc,, -deepspeed/checkpoint/__pycache__/reshape_utils.cpython-310.pyc,, -deepspeed/checkpoint/__pycache__/universal_checkpoint.cpython-310.pyc,, -deepspeed/checkpoint/__pycache__/utils.cpython-310.pyc,, -deepspeed/checkpoint/__pycache__/zero_checkpoint.cpython-310.pyc,, -deepspeed/checkpoint/constants.py,sha256=FQ55iGvLxiXitjitVfepx3-vLrPx7J-btyfLUGyPGCc,3510 -deepspeed/checkpoint/deepspeed_checkpoint.py,sha256=mrL3zQ6pvWBQIZyBtzsXTNEIt2jbyemunrweHGE-Cl8,13095 -deepspeed/checkpoint/ds_to_universal.py,sha256=IbP-9e1A6Y95I6AWLolK1OCHjJ78Ndrp9i07zaLqJN0,22788 -deepspeed/checkpoint/reshape_3d_utils.py,sha256=pmnOtdI3srRS46R9sXLh9kijXPxT3bFJ_c6lWdygbwc,4674 -deepspeed/checkpoint/reshape_meg_2d.py,sha256=o5dpF3CbrbtEt1twz0XsL3xbqCfLfEse7O2ke-WR9_s,7885 -deepspeed/checkpoint/reshape_utils.py,sha256=n3k8xjXn2--nRzcs7jo7V9lFlma7LTLXjkk9S3IgUhM,3458 -deepspeed/checkpoint/universal_checkpoint.py,sha256=EF1VbI2eeHdHXmdgYbHCYVOIf8sjSTcGqY82qjtCtmQ,6788 -deepspeed/checkpoint/utils.py,sha256=Wog7wLoR_F88UOSlLRAX8XspdjA9mF3CiJ5iOYVWh9E,2693 -deepspeed/checkpoint/zero_checkpoint.py,sha256=OsB42DLdUAeuqFbynEHaJGAo3UTbuog1UicXg8dVJAU,5441 -deepspeed/comm/__init__.py,sha256=N31DOMdAdlfSnQd64slK89dpW_y1Vc9oPVZabLArcYw,137 -deepspeed/comm/__pycache__/__init__.cpython-310.pyc,, -deepspeed/comm/__pycache__/backend.cpython-310.pyc,, -deepspeed/comm/__pycache__/ccl.cpython-310.pyc,, -deepspeed/comm/__pycache__/comm.cpython-310.pyc,, -deepspeed/comm/__pycache__/config.cpython-310.pyc,, -deepspeed/comm/__pycache__/constants.cpython-310.pyc,, -deepspeed/comm/__pycache__/reduce_op.cpython-310.pyc,, -deepspeed/comm/__pycache__/torch.cpython-310.pyc,, -deepspeed/comm/__pycache__/utils.cpython-310.pyc,, -deepspeed/comm/backend.py,sha256=TPFMx6ViU4n_NcfwWS56_Yp35PeCi-0Kh9za7UELeyA,1416 -deepspeed/comm/ccl.py,sha256=wGGn3sFUvbzblOsA4uMfKCltk6giyI2dkG6SkmMFzVc,8614 -deepspeed/comm/comm.py,sha256=XFedePFBQqfJngaz1WWhC9dLVEenUd-bHpR0M2JQWlQ,31135 -deepspeed/comm/config.py,sha256=vd3v2i0X638bgqgVENb03GddaJfZVJYV17k3dRHdmKM,728 -deepspeed/comm/constants.py,sha256=Adtnb5LCEjRpylnwgsMAFRmWgPvxb8CEEDvn4-D39lw,1298 -deepspeed/comm/reduce_op.py,sha256=A_KxQtpfUrNun6Rn8SX6iQiuoVgfYM7bTtAYc-YTENE,259 -deepspeed/comm/torch.py,sha256=IOBv9fuZjBdKK_vQvoOnHbUAzlYKnJSgzTCzKVByjs8,19481 -deepspeed/comm/utils.py,sha256=I7hJ2Jq5HEPaJDD09eEI9xV5VJZeDdag0ovTeAc3Ex8,3842 -deepspeed/compile/__init__.py,sha256=-DiusbPt_BUtHIDjxVtMiLEuFDUb6ANm671qfLuUXVY,126 -deepspeed/compile/__pycache__/__init__.cpython-310.pyc,, -deepspeed/compile/__pycache__/backend.cpython-310.pyc,, -deepspeed/compile/__pycache__/config.cpython-310.pyc,, -deepspeed/compile/__pycache__/fx.cpython-310.pyc,, -deepspeed/compile/__pycache__/graph_param.cpython-310.pyc,, -deepspeed/compile/__pycache__/inductor.cpython-310.pyc,, -deepspeed/compile/__pycache__/init_z1.cpython-310.pyc,, -deepspeed/compile/__pycache__/init_z3.cpython-310.pyc,, -deepspeed/compile/__pycache__/input_storage.cpython-310.pyc,, -deepspeed/compile/__pycache__/list_schedule.cpython-310.pyc,, -deepspeed/compile/__pycache__/partitioner.cpython-310.pyc,, -deepspeed/compile/__pycache__/patch_compiled_func.cpython-310.pyc,, -deepspeed/compile/__pycache__/patch_fake_tensor.cpython-310.pyc,, -deepspeed/compile/__pycache__/util.cpython-310.pyc,, -deepspeed/compile/backend.py,sha256=7J3IuJJwMKDuH5hg5Ot4nwJhjDbesQxGRwT9-SV3pKc,15885 -deepspeed/compile/config.py,sha256=n-0B63IyyOlK6j-bfUSfVlCxdvJHMDdjPmadKYrb2MQ,1521 -deepspeed/compile/fx.py,sha256=WY5Shuq80GsEj1r3vNd20azPZNfcbMbGkr-c6JfMLVY,4717 -deepspeed/compile/graph_param.py,sha256=KBsOiXs6KfQ_p3UAakfOoTFQfvKXaEc-zNcnRh33V3w,3802 -deepspeed/compile/inductor.py,sha256=dfQC6RSYFZbR2_IXtfkrhNl88t31QZjosp159k477q0,8933 -deepspeed/compile/init_z1.py,sha256=slw23gEhnglbkN1UciCkrgClPmY1cCIf7JHWABwf-n0,3488 -deepspeed/compile/init_z3.py,sha256=59J3fEVt0J1cS_hAjE13tj59SO4RUyE2xv72khpVJOg,3905 -deepspeed/compile/input_storage.py,sha256=u71Xm3YBa5TV7in8SvqzF7PJA4iVO0Fjyw9YlRibwfc,7835 -deepspeed/compile/list_schedule.py,sha256=S2d9OolXfOwpyUkE68QHslWUPOmGWCBc8ijx2qCAsjc,15870 -deepspeed/compile/partitioner.py,sha256=fOtJtoRgoULNvYImDDgAnYli0OH3khT3W0GI98dM2Q8,6902 -deepspeed/compile/passes/__init__.py,sha256=YPYE70D8HDgY-Wo0z8HgOF7ikgUN0F29deid37kdF8k,1470 -deepspeed/compile/passes/__pycache__/__init__.cpython-310.pyc,, -deepspeed/compile/passes/__pycache__/offload_activation.cpython-310.pyc,, -deepspeed/compile/passes/__pycache__/offload_adam_states.cpython-310.pyc,, -deepspeed/compile/passes/__pycache__/offload_parameters.cpython-310.pyc,, -deepspeed/compile/passes/__pycache__/prefetch.cpython-310.pyc,, -deepspeed/compile/passes/__pycache__/selective_gather.cpython-310.pyc,, -deepspeed/compile/passes/__pycache__/zero1_compile.cpython-310.pyc,, -deepspeed/compile/passes/__pycache__/zero3_compile.cpython-310.pyc,, -deepspeed/compile/passes/offload_activation.py,sha256=fPCEG4OtfyMdLEtVCzZYrlvkeZkAeiVX29qh7tlrt_0,3849 -deepspeed/compile/passes/offload_adam_states.py,sha256=vprSD81ZODjObxQFOpkAoawQoyD1nL_X2DVYeKX8YpY,20605 -deepspeed/compile/passes/offload_parameters.py,sha256=zOOIpmbra727IUfSb-Rihuco-PVcTF2cJMfys626ZbM,2046 -deepspeed/compile/passes/prefetch.py,sha256=6NW_QLnIw6VzV5edhmb6sonp0Wu-TSqRPh7KhMO4Ho0,7696 -deepspeed/compile/passes/selective_gather.py,sha256=icQWdK73H61VaD54xZdq-R-9hHvLF7-spofiwdeK3Jw,5466 -deepspeed/compile/passes/zero1_compile.py,sha256=UHM5rtq1JX6Rjj0asbF7Nt9EUy25pvalseM4xXd7QsY,2323 -deepspeed/compile/passes/zero3_compile.py,sha256=PzrtTW-oDyEWQb_J38F3cAzIsRvU0w100SWv-FnB8L8,7645 -deepspeed/compile/patch_compiled_func.py,sha256=ijs5Yo7IkLve0j5Hpwe3QxiiTL1AEhOrlGCW7GiAjVo,2856 -deepspeed/compile/patch_fake_tensor.py,sha256=1d-sXFu63HF2vCPqA7Wg3rVvUupeQjezrRlqeMOaYxU,2010 -deepspeed/compile/profilers/__init__.py,sha256=OSidjoyeb6NwY7vBzdvPEJHDPpS08WRUbR-m8HFLyXE,957 -deepspeed/compile/profilers/__pycache__/__init__.cpython-310.pyc,, -deepspeed/compile/profilers/__pycache__/comm_profile.cpython-310.pyc,, -deepspeed/compile/profilers/__pycache__/graph_profile.cpython-310.pyc,, -deepspeed/compile/profilers/comm_profile.py,sha256=xM3vbFB7gYPdyVqTicuEqOCB4mW89TCR81j6kRVVbew,5103 -deepspeed/compile/profilers/graph_profile.py,sha256=UDzJGK7ENwoDzpB2LepW0Aa--CHmF4I3Ow5D5rdoSrk,11243 -deepspeed/compile/util.py,sha256=Zirqj_4obWpWc-wPi_Hrn5SU7UZ11vBIiSALu4WHKCE,15586 -deepspeed/compression/__init__.py,sha256=s3fNEEJJdLa4D4rEwtraP5axoSqpxxmh6LAIvyzfUTM,243 -deepspeed/compression/__pycache__/__init__.cpython-310.pyc,, -deepspeed/compression/__pycache__/basic_layer.cpython-310.pyc,, -deepspeed/compression/__pycache__/compress.cpython-310.pyc,, -deepspeed/compression/__pycache__/config.cpython-310.pyc,, -deepspeed/compression/__pycache__/constants.cpython-310.pyc,, -deepspeed/compression/__pycache__/helper.cpython-310.pyc,, -deepspeed/compression/__pycache__/scheduler.cpython-310.pyc,, -deepspeed/compression/__pycache__/utils.cpython-310.pyc,, -deepspeed/compression/basic_layer.py,sha256=oZvNusdxqnlhM-RsaALf5Q-4m2kQLcKAUUMuIKBtXec,36047 -deepspeed/compression/compress.py,sha256=DueechYg3bDNAIn2XW0L9nv8qN81acTPwimvD4usE_8,11886 -deepspeed/compression/config.py,sha256=5TTZNLB_GDULXEf704-dbNJo0ILL47DNVWHWiMGio8Q,25067 -deepspeed/compression/constants.py,sha256=Fc2681jhRfm0ADzNjFKjhYIPSXAI29gqeQSGW0FxLoA,5569 -deepspeed/compression/helper.py,sha256=lEFpurhyU7-yKWgO-wp1RojHECkqR8qJ16KrT4Hl06U,14637 -deepspeed/compression/scheduler.py,sha256=QwnyBtRhkEkufjyhns-Riz_-GBdUePzuI8UEQJOu7Ho,8112 -deepspeed/compression/utils.py,sha256=vLWAEOQHa1w61zmMyFQ3dwxbm-5kM_S-zOI_eQhYx5E,7818 -deepspeed/constants.py,sha256=K5l2jzMTXYfqyUBbO9QIMbe_1TFMxdJ9LJDEQ6HJLRY,867 -deepspeed/elasticity/__init__.py,sha256=Cmi3gEN48VYdPSg4zLOJ1KGKP2G3Xeruz9TykgC13m4,383 -deepspeed/elasticity/__pycache__/__init__.cpython-310.pyc,, -deepspeed/elasticity/__pycache__/config.cpython-310.pyc,, -deepspeed/elasticity/__pycache__/constants.cpython-310.pyc,, -deepspeed/elasticity/__pycache__/elastic_agent.cpython-310.pyc,, -deepspeed/elasticity/__pycache__/elasticity.cpython-310.pyc,, -deepspeed/elasticity/__pycache__/utils.cpython-310.pyc,, -deepspeed/elasticity/config.py,sha256=OCmF8fP8_dL3LkKiEJDtOREnJB3NYQuCaK48qRs-eFA,4703 -deepspeed/elasticity/constants.py,sha256=5Ynz57XbIFxeFkhC3JwanR-Dw0-U7qrsqKIH9oQo-Hw,2454 -deepspeed/elasticity/elastic_agent.py,sha256=KrPIykg6Q4K6BrjkgWmj0PRkgGHNuOb5vvgNjf7XZwk,8009 -deepspeed/elasticity/elasticity.py,sha256=oGJXhEj-VVdT2iUa-ei_gHXAfhP5e_VgtUlw5N500mU,17361 -deepspeed/elasticity/utils.py,sha256=-mbP_VX-SSstZKrEp0hkFwezql9pkAEvxx3BxMO2-ic,347 -deepspeed/env_report.py,sha256=0ZVv7RaHSaQIOJKWG5dqs6OfFalx6adq3NFZwB3nVIs,7491 -deepspeed/git_version_info.py,sha256=IZlPdw3IO5LiSJQ5ECQBVNnIj9TL3ezz-EEYfSR54ag,1075 -deepspeed/git_version_info_installed.py,sha256=XabhZhUkeGWlt88HDwsiXiUqQgraSOvvYg27wpRyljU,369 -deepspeed/inference/__init__.py,sha256=ZV8FRXVTJ3IGIDnyXJgJjRlCbB2jH_KJ1hWoTS4l_y0,267 -deepspeed/inference/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/__pycache__/config.cpython-310.pyc,, -deepspeed/inference/__pycache__/engine.cpython-310.pyc,, -deepspeed/inference/config.py,sha256=zBiw9UvgAQ79FfauGJPLpXwXgtxrrBMUwUdokM-TEww,11283 -deepspeed/inference/engine.py,sha256=2ndGZLN0dHv2RDDeaWj0do5o5OLrzutfqZMPDO6DsIA,30250 -deepspeed/inference/quantization/__init__.py,sha256=4I9UpQ5vMRU5SYSF_dW9FJDEnBq4m_0SuwtVQ92lGaA,95 -deepspeed/inference/quantization/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/quantization/__pycache__/layers.cpython-310.pyc,, -deepspeed/inference/quantization/__pycache__/quantization.cpython-310.pyc,, -deepspeed/inference/quantization/__pycache__/quantization_context.cpython-310.pyc,, -deepspeed/inference/quantization/__pycache__/utils.cpython-310.pyc,, -deepspeed/inference/quantization/layers.py,sha256=s62KAh9pysdpMzV3KuaFeOdrezxLSsPcZi-Fmsj0w84,5729 -deepspeed/inference/quantization/quantization.py,sha256=-r1gbQC0AZe3umfsjla_2lp6opWrfI-0jfs6KtOaJ8o,4396 -deepspeed/inference/quantization/quantization_context.py,sha256=0IV04DTAk8PiOi07TGgABxFMkrKM4QFOqzRaT-RZO7w,514 -deepspeed/inference/quantization/utils.py,sha256=qu0HJHOSz020DT1UuMUo45c1IyARecLl_cCoMcDuf9c,11933 -deepspeed/inference/v2/__init__.py,sha256=9LrBSo2xdUl33MPzLedqmI13rnQOJJO-eT6fVsrt61k,283 -deepspeed/inference/v2/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/__pycache__/allocator.cpython-310.pyc,, -deepspeed/inference/v2/__pycache__/config_v2.cpython-310.pyc,, -deepspeed/inference/v2/__pycache__/engine_factory.cpython-310.pyc,, -deepspeed/inference/v2/__pycache__/engine_v2.cpython-310.pyc,, -deepspeed/inference/v2/__pycache__/inference_parameter.cpython-310.pyc,, -deepspeed/inference/v2/__pycache__/inference_utils.cpython-310.pyc,, -deepspeed/inference/v2/__pycache__/logging.cpython-310.pyc,, -deepspeed/inference/v2/__pycache__/scheduling_utils.cpython-310.pyc,, -deepspeed/inference/v2/allocator.py,sha256=A-yTKojaNr_O8COkDIFagTl4dd9aYJRIMtR7g8TIvJ4,1182 -deepspeed/inference/v2/checkpoint/__init__.py,sha256=6y_8HbNnkff-I3LYNyzZQ3DVDcy8CwjQAYykp_JiCw8,252 -deepspeed/inference/v2/checkpoint/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/checkpoint/__pycache__/base_engine.cpython-310.pyc,, -deepspeed/inference/v2/checkpoint/__pycache__/huggingface_engine.cpython-310.pyc,, -deepspeed/inference/v2/checkpoint/__pycache__/in_memory_engine.cpython-310.pyc,, -deepspeed/inference/v2/checkpoint/base_engine.py,sha256=9EmyYEwt1ZE861yuS-ayx9ykG2OzRi9p7_ua4hECMEs,1391 -deepspeed/inference/v2/checkpoint/huggingface_engine.py,sha256=sn6nmRLiOLEVxTCMHRExUbbIjMVavCAFyerbdxssz-M,5997 -deepspeed/inference/v2/checkpoint/in_memory_engine.py,sha256=4tEQQPAWl4BG4FJFdcgnoa8jahMQUtrXtcssTwRUSeo,1487 -deepspeed/inference/v2/config_v2.py,sha256=HT1FCeys_QnuUTNteBLQhJ8Y6Cg4WVW5B1HwAwQSUX4,1389 -deepspeed/inference/v2/engine_factory.py,sha256=ap6CwbdaOvD3IcCOCaqfhAvIRhp0lu0efZCLHTEMST8,6028 -deepspeed/inference/v2/engine_v2.py,sha256=2IiT1buZLPYL66oC8OQdpANhRSFjN-xjvnnCJh9Wqts,10491 -deepspeed/inference/v2/inference_parameter.py,sha256=rPeDJ4h58L4weNeIiFDIUrpVgDLpkQtzihUl_L-rkbg,2788 -deepspeed/inference/v2/inference_utils.py,sha256=V5gryIBXY9kVVkwPWukMMzQKbn0vH2hyDWnDhQ4jWQk,2378 -deepspeed/inference/v2/kernels/__init__.py,sha256=KHPoSEGybfJdU54HoIOCZzeIovB3EsfzaZuHv_i36V4,132 -deepspeed/inference/v2/kernels/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/__pycache__/ds_kernel.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/__init__.py,sha256=72TCKbfeUY6LbxtwVeNKljxXvxFrY-zobANphLq7nVQ,276 -deepspeed/inference/v2/kernels/core_ops/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/bias_activations/__init__.py,sha256=BpRq2zT5tKcsx4e7f8zcnDA5F6wEedmQXSI6idSEoTo,127 -deepspeed/inference/v2/kernels/core_ops/bias_activations/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/bias_activations/__pycache__/bias_activation.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/bias_activations/bias_activation.cpp,sha256=PdSuObELmrSu3lsYgfVEOO9tEEmexqO0an3_pP4nU3k,3084 -deepspeed/inference/v2/kernels/core_ops/bias_activations/bias_activation.h,sha256=q3vJjn6I7s_raM1Gy2TNnrY3nbskDDP9QyWwcDjpk6E,688 -deepspeed/inference/v2/kernels/core_ops/bias_activations/bias_activation.py,sha256=6uu2cWl65xBVWVHfZR9TfLiqrovUJOxOm50y2iZ9Y28,2437 -deepspeed/inference/v2/kernels/core_ops/bias_activations/bias_activation_cuda.cu,sha256=dq_Bwd-UGv8bK3wWA1llNGVUEXO_rAMZInl5PcfpO4Y,4820 -deepspeed/inference/v2/kernels/core_ops/blas_kernels/__init__.py,sha256=ND9hDC4WOLfDXnlqe7FySprn8RkCXvzrwQKwYG8ZFRQ,123 -deepspeed/inference/v2/kernels/core_ops/blas_kernels/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/blas_kernels/__pycache__/blas_linear.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/blas_kernels/blas.h,sha256=YtINP8OfWWiyZEx44zb1yylsXBHGBLdfxcb3xjBUNeI,5164 -deepspeed/inference/v2/kernels/core_ops/blas_kernels/blas_linear.py,sha256=AQ4eOYQ3ejZspCeBHm87Q32f_Lf_ooauZIoVobdD_G0,2023 -deepspeed/inference/v2/kernels/core_ops/blas_kernels/blas_utils.h,sha256=IGCBmqS3jASIg3Trt70BRfl2KIxtF2T8fQ2HxIOUJNs,11400 -deepspeed/inference/v2/kernels/core_ops/core_ops.cpp,sha256=4HbbHAOUmkaolemLrP_nWOdXSIbAP7KGgl1jTU1LMfI,1514 -deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/__init__.py,sha256=tUarONsuAuUxyfogdH03MC3yebDySSrUEusyIEWNjrg,174 -deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/__pycache__/cuda_fp_ln_base.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/__pycache__/cuda_ln.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/__pycache__/cuda_post_ln.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/__pycache__/cuda_pre_ln.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/cuda_fp_ln_base.py,sha256=CmnqS4_iIbPWKbzNCKEv3wnng8UnRtyhSJiEz08cfPI,1302 -deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/cuda_ln.py,sha256=bp4CvANVf8hQQzyKOIScnR1OU7kB41uBNK5ANIydCvs,854 -deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/cuda_post_ln.py,sha256=cIJ9l9ey_fAdjuX00ZjooepUgvNEg_bc2ikSiR0MWK0,990 -deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/cuda_pre_ln.py,sha256=meb0fP8vvu1lFGlMR8iEus7YKlGP6FMlad9FJ7ZPg8A,1225 -deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm.cpp,sha256=cWe1Z6kgQuPsrSTCOO1mWBtFY_CYqXHUt1HqhfiOyNs,4199 -deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm.h,sha256=SCuiMaChAEgmE0t2kwfbqaYXQ1ycU1QXDVDJSHxtQ6o,2157 -deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm_cuda.cu,sha256=NLCh_1U7k5iA5EOmmqISVMzoa47ypMYYZ5UAk_8tvMo,19730 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/__init__.py,sha256=0J6ih4FAYwgMMo0HSRWLTiBimI9hdeAFh9Ui5iwx2NU,123 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/cuda_linear/__pycache__/cuda_linear.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/cuda_linear/cuda_linear.py,sha256=_2OxKHsY9P5mKOW3sT1ZTIy72d3BGpEjt7_Kb2bNsD4,6827 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/configs.h,sha256=XPl9PmMyD-OVFksGXN-eK67a-aoxJvjxNEBIXm3q6Fs,5044 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/kernel_matmul.cuh,sha256=e92X51umAswaFGMz9ntC58TRD1W33CjVZPS2eVeZFOk,14700 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/kernel_reduction.cuh,sha256=GQix8cchlsjHeTajB6oM08OEoEC1Mcp62LEsq3XKnAU,1494 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/ptx_cp.async.cuh,sha256=kdZEfq2Akuw1rSJMdR6DQ0k23mbYhRuQ4cKffK9_wUc,2357 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/ptx_mma.cuh,sha256=_-OJfWV4dhi9or_B1orbEkxBdZCNiXnmjW5-eNHtmno,5139 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/utils_core.cuh,sha256=9RDNZ3HQFo0rBbFgaBJpM7JzxXQ4fYiUNdnlsZOfM84,10251 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/utils_gmem.cuh,sha256=pyOFBGMAlcDFZx5iSS-HPnXA0Pqwx9oCvrUFV_EkioY,3207 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/utils_paralleldequant.cuh,sha256=bqJJIaioGtHe1feD1vIjjWS1OpKGTrau1h03cCdkiaI,4057 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/weight_prepacking.h,sha256=JTrsVJdRovts8Q7bG7MlCzsVyKL_R76moU-zyakMPGI,8896 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/linear_kernels.cpp,sha256=XIFk-nXPumydZOOpupiRPrqgN6wbPRMr86u8_AAqcnM,8892 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/linear_kernels.h,sha256=i97f3j3gOIiFK4Ijoiu05U9aj0S4p2NfhBIIYe-hAbM,828 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/linear_kernels_cuda.cu,sha256=l5rmu_QnbM1VNmao97uVySwP2sOGjjF919pX18_KXwA,13947 -deepspeed/inference/v2/kernels/core_ops/cuda_linear/linear_kernels_cuda.h,sha256=ON5qqFROLBv392C3nwM0adOlKtCHsjO59yYUgGQwz_A,1669 -deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/__init__.py,sha256=S3Q1i1iwyLp6kWWJjD2kTinhJQ_ND4gKHejhuxQ8-74,171 -deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/__pycache__/rms_norm.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/__pycache__/rms_norm_base.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/__pycache__/rms_pre_norm.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm.cpp,sha256=XqvW2ObN0dkRQIsqtWRGl_JYxkFjKAi-zjFMeEZVa3M,5873 -deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm.h,sha256=FUf2hMir5TnFA6iU1Rx48DlnSYX__Yl0jkiNWARdREA,978 -deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm.py,sha256=OkC86DHLQ03zshmFNUQO9r1mZ4mWD3gN0505IvzC1NE,770 -deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm_base.py,sha256=cbZXvZSXRaD1MoTokRSRuw1anp_FXdGzg0moAxasQ1I,1311 -deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm_cuda.cu,sha256=GrWivwyCSZ8gh1zaiJoKHh8sTcq8aot5m7XKx-Auw1A,10216 -deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_pre_norm.py,sha256=qzME0FbrgV9IXbGCqIA0IP91rfFyRiTdiXwwQ39SufM,1204 -deepspeed/inference/v2/kernels/core_ops/gated_activations/__init__.py,sha256=2mIcBil301zqLsisa-B2C2EnTgQ2Zaz2QFRxPTD8eUE,128 -deepspeed/inference/v2/kernels/core_ops/gated_activations/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/gated_activations/__pycache__/gated_activation.cpython-310.pyc,, -deepspeed/inference/v2/kernels/core_ops/gated_activations/gated_activation.py,sha256=XQXwTY3PN0qWIdZMBWh8De2lEGtN8syw7ItH_uN319g,2755 -deepspeed/inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels.cpp,sha256=aog8VhJGL0F4bwUOZTR232KTgRJRavFFk3Xj_55shLk,3521 -deepspeed/inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels.h,sha256=DpAhhx6lO3TNA8eid6Tct9qfDrNFweHRypmn4serMW0,776 -deepspeed/inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels_cuda.cu,sha256=0uLVwa_H3oSJczKK61AZCbBIUZF1mts1QGIM3yKUI6k,6080 -deepspeed/inference/v2/kernels/cutlass_ops/__init__.py,sha256=hpTBFM7draPD0dwxTHvGYgAdic-ysFMG6_XEr8x6V5c,146 -deepspeed/inference/v2/kernels/cutlass_ops/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/cutlass_ops/cutlass_ops.cpp,sha256=v6-PzCFZnmxFPXJm-pooeUISCq2V0S4E4_PAvg7oA6w,483 -deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/__init__.py,sha256=leIZbZ0995YwKwql6mU54HN_DJ1iBe67Cq_GdlXSMAY,122 -deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/__pycache__/mixed_gemm.cpython-310.pyc,, -deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm.cu,sha256=fDh2z9IwBffz5AI1NltOo4tgC2a2c0bwZ-leLahnR2k,4131 -deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm.h,sha256=laUi-GFh6qbLo0Ywiz_h7TtNIxroXSFwPQKc6yBGhE8,410 -deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm.py,sha256=gZ7ljMs2f20rcL1624vIhiLg_yemxkrYYFDoX-K1M3Y,2674 -deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm_api.h,sha256=soDC6vGW7UCYrnfw1EftpmOz6BuOi27Ihyi8J2Cfw2Q,1570 -deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/__init__.py,sha256=RRc_fJh0o2_l9wVb7Czf2GqTxLzvS6Rj_zeYUGzRjjI,150 -deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/__pycache__/mixed_moe_gemm.cpython-310.pyc,, -deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/__pycache__/moe_gemm.cpython-310.pyc,, -deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/mixed_moe_gemm.py,sha256=pAaowuHqgwtCqv3tyof9jpbPxZxgLdcj1vPM2NgDcx4,2996 -deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm.cu,sha256=FRZSfEKLH3lHjvdbTbuEr1qGAi1spT6lhTLuIJ_OGuI,8594 -deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm.h,sha256=hdqzk6YOF68wO1DBVWDdtX5z-MIGHdtekijLU8z1cHM,740 -deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm.py,sha256=zX6btY3JLxP-6mgcjlNLSHlmgHJ2z3ZMH6PV3DQPfeE,2572 -deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm_api.h,sha256=RWmGN2LrcQukrsLoYaqykymCQJ1onmubdqXrHHgx2Is,1839 -deepspeed/inference/v2/kernels/cutlass_ops/shared_resources/weight_variant.h,sha256=1Mr74mw0SDFt0BsMgWiLvYXXaMWmhIgCUz1MZlWXYMs,278 -deepspeed/inference/v2/kernels/ds_kernel.py,sha256=R0n_uBCLCGAGlJa7lvep4cNF0DBxkD6TDKDFNAWpBpU,961 -deepspeed/inference/v2/kernels/includes/activation_type.h,sha256=WkFfkL5mgLyYAZimLUouThCkyA2aiF9LQlh1q206L9o,264 -deepspeed/inference/v2/kernels/includes/conversion_utils.h,sha256=augq8Zf-t46MaxWh0ryFb46_Guxud0k4kIK-h93rI8Y,12379 -deepspeed/inference/v2/kernels/includes/ds_kernel_utils.h,sha256=OKOMT46XItRnkrzqKHHK9n_ewxHh6AN-Uhv4aKixBm4,1297 -deepspeed/inference/v2/kernels/includes/memory_access_utils.h,sha256=wqpAqnfmSD5Fx15qwsnY63VCL6tYyJaChCqLxakM10w,33966 -deepspeed/inference/v2/kernels/includes/reduction_utils.h,sha256=189h93WF77CblgWonfW5eQqoTeJnb-Rt2Ov2fZ6mO-E,22352 -deepspeed/inference/v2/kernels/ragged_ops/__init__.py,sha256=8C6TomumccM_-p-iiyPyjZnDLPjTdQS8E4ydGQGNLu4,324 -deepspeed/inference/v2/kernels/ragged_ops/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/atom_builder/__init__.py,sha256=djM0DTjGkr4vvNuhRIWV-2OI7lroepRMBiHIJtrcqTA,124 -deepspeed/inference/v2/kernels/ragged_ops/atom_builder/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/atom_builder/__pycache__/atom_builder.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/atom_builder/atom_builder.cpp,sha256=AjIgvgoS7PUYGt4MW7oFu8QEcQvm3J16tyerW-dx8cI,2051 -deepspeed/inference/v2/kernels/ragged_ops/atom_builder/atom_builder.h,sha256=i232uHXwHjh2EaAF3709YNcTQ29uVDI2GQOs6GnYAsQ,732 -deepspeed/inference/v2/kernels/ragged_ops/atom_builder/atom_builder.py,sha256=QgVnamgI1qG8xjzu-9i5sdBccOuhgbTeAtufUtLNdog,1744 -deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/__init__.py,sha256=696rc8-r9nWCnfSG8UZ_fagI28Nerwq-nZCDvIIXPRs,125 -deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/__pycache__/blocked_flash.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/attention_atom.h,sha256=-ITylnAQx_20rXiWdPPldb149cGG-fNQUqOg8iHVugU,1062 -deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/blocked_flash.cpp,sha256=qaNovmlpqw7Cw_Y1ArGmH-7fB1nPQd0idvmC4nDR9pw,3912 -deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/blocked_flash.h,sha256=io7L6krERzOgiBIKNk3zwEppMDSvWtOqCY9nH35dB_U,459 -deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/blocked_flash.py,sha256=-aU8A7FbwrRqJPSHZEOef6Y3_PRndGKfCFT1c-3S1DQ,3830 -deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/flash.h,sha256=VJKtKOCufGTGt8ZGuPXfsA-dzsi87cqNTYX7pbqHKH4,1931 -deepspeed/inference/v2/kernels/ragged_ops/embed/__init__.py,sha256=y4HCIAaLuHexvTBcNppgiXqo76vHb_pAahL75kQBKY0,137 -deepspeed/inference/v2/kernels/ragged_ops/embed/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/embed/__pycache__/embed.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/embed/embed.cpp,sha256=iI-rjkQlCGnujjVzY1h8e6ccO24G7eVxwPLY5t_HKxY,4840 -deepspeed/inference/v2/kernels/ragged_ops/embed/embed.cuh,sha256=C5ZS1W9lnEDPBmWsuoJpRP-pgCYwLxqTK_OI-DQrT04,956 -deepspeed/inference/v2/kernels/ragged_ops/embed/embed.h,sha256=XbqCgqlx16bRtRLljUstzjssRJIGx7AixbbaofhILl8,700 -deepspeed/inference/v2/kernels/ragged_ops/embed/embed.py,sha256=6lxDG2ymTyx2yGZ_9jDFyvcchbtD1iwKr2icLgHtmxE,3011 -deepspeed/inference/v2/kernels/ragged_ops/embed/embed_cuda.cu,sha256=5JfvhMBcKAPXQbF4WvE3Ljmolom8ArjvkNtprx1SaKo,6055 -deepspeed/inference/v2/kernels/ragged_ops/includes/top_k_utils.h,sha256=T0LJvEtOulSPO1r0YA4jEHW4l-3XSYyS8nuI6Ok2F6E,782 -deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/__init__.py,sha256=ou-mEvP9FMj89HvxswZbtTvS_T8Gcy3zf16jPcb_Urw,208 -deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/__pycache__/blocked_kv_rotary.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/__pycache__/blocked_trained_kv_rotary.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/__pycache__/linear_blocked_kv_copy.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.cpp,sha256=qUdVHb-2UMnlH8Q6zZJLUW69T75AMwfqEZjco2EXbdY,8110 -deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.cuh,sha256=X9w1u75CiGNEOV0uaXAmjX8usEiTXFTZG7EK3ulvZDc,1676 -deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.h,sha256=5RfNKDify3eHIV6GjfqNozT6_lc2shcwNidgrQ6AJD0,2406 -deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.py,sha256=spVrPPi6tMh31Zjb3l4UooFgJ2YTFq9UJI2f50Mziuk,3142 -deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary_cuda.cu,sha256=ztJXnG_ik8nCCJQw2d3P0ScuNmTUPvCKQT1I3L28EzA,19275 -deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_trained_kv_rotary.py,sha256=FGU4a0Wcb4NtsbJiLswLnd6-M89yE800UW19v5kFp9Y,3220 -deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/linear_blocked_kv_copy.py,sha256=mAa-291Bzc06yNGZVgsYPhrpq7dVoYDkWDHFUr2i1fY,2957 -deepspeed/inference/v2/kernels/ragged_ops/logits_gather/__init__.py,sha256=-6rVFXUnZ3U4MBA1VNuKQfhcdjlWgc_xnufXpnFfGyg,125 -deepspeed/inference/v2/kernels/ragged_ops/logits_gather/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/logits_gather/__pycache__/logits_gather.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/logits_gather/logits_gather.cpp,sha256=wItdeUW_dSJzfePXbRKjYeggcfGG_l8yO7OAPXF4IeU,1800 -deepspeed/inference/v2/kernels/ragged_ops/logits_gather/logits_gather.cuh,sha256=OQlrhwuHJIwGfhbozWlM4NKhYoKBtAL-3nc0vqATdZY,624 -deepspeed/inference/v2/kernels/ragged_ops/logits_gather/logits_gather.h,sha256=bRw1hkiDAAqkP9FkNwsL1QZ0UeGSx1_pYHOY4yoTF5k,571 -deepspeed/inference/v2/kernels/ragged_ops/logits_gather/logits_gather.py,sha256=kGPLmxKKqOh2x35QWtCpeXLXxoG15kxmE_Z_-cd7Yv0,2121 -deepspeed/inference/v2/kernels/ragged_ops/logits_gather/logits_gather_cuda.cu,sha256=Dd2HJJHpPlDkt2C91667Fn1GVuc4jonsIW8s1vPKhHk,3312 -deepspeed/inference/v2/kernels/ragged_ops/moe_gather/__init__.py,sha256=K3BZ4_CDJGkNmvFCeYvIEBY2DRMWqq5CzCL7HAStX8Q,122 -deepspeed/inference/v2/kernels/ragged_ops/moe_gather/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/moe_gather/__pycache__/moe_gather.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/moe_gather/moe_gather.cpp,sha256=zWULPg_6e15LzTk5q2DJJcsgwT4zoLhTNIc0ClJf73w,2392 -deepspeed/inference/v2/kernels/ragged_ops/moe_gather/moe_gather.cuh,sha256=orang8iBcPNQ0BjFTeiBWNX09UjDq82Sh8tyye-5xyU,706 -deepspeed/inference/v2/kernels/ragged_ops/moe_gather/moe_gather.h,sha256=MfJuJ2rg_ba6KvIm1LNbTTHMBH8MYQHgd7JeYJSG9Q8,561 -deepspeed/inference/v2/kernels/ragged_ops/moe_gather/moe_gather.py,sha256=31zC0Sb3caA12JdhT7qJoawCUTkKB5PjZpFupQI6FvE,2265 -deepspeed/inference/v2/kernels/ragged_ops/moe_gather/moe_gather_cuda.cu,sha256=oy_bzQhyMfmjwiXag0u4YKjyKHaLXMDWYHlJQnQZNWM,6679 -deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/__init__.py,sha256=e3cxpdOKqIBYpxvNmMnj6Ms883dqdsl4B_fXdqTnlnM,123 -deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/__pycache__/moe_scatter.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.cpp,sha256=ywodf_RYumTot90MnTLb4o1h7HDWgiFuVTuhHwxNooE,2858 -deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.cuh,sha256=M6t_ydSLYKsekCym5m07AGDu8-U_qVG0hLuVH_nrO-A,772 -deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.h,sha256=pBdy99sow7TKoB3eV3gdPbCB4_pGNx48K-Eo8C8XjHs,701 -deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.py,sha256=7c_jFxhoJINWl57LIhFZwvlr08W8Nbdk7ULSFmC8Yqc,2504 -deepspeed/inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter_cuda.cu,sha256=US5_PjVcIqYcqR2DDKXUIJAMfdXKI0ElfbNCUmLscbw,8070 -deepspeed/inference/v2/kernels/ragged_ops/ragged_helpers/ragged_dtypes.h,sha256=OVEob4D9LFPkyxA4PscOrGV5BzBalfiMrTrkcifCvL0,940 -deepspeed/inference/v2/kernels/ragged_ops/ragged_helpers/ragged_kernel_helpers.cpp,sha256=QaWqjTz4kXXR804xXlgNBTOxtOtLSP3NcGDNHuFqqQI,1149 -deepspeed/inference/v2/kernels/ragged_ops/ragged_helpers/ragged_kernel_helpers.h,sha256=_7Xb4l1mfYbSz7i2HAZ8TqQSLJy6TBBC51OMaarJXfM,563 -deepspeed/inference/v2/kernels/ragged_ops/ragged_ops.cpp,sha256=Bnku0tlvBSyPIwOtx7-kvIGg3OPGmsgjMO72rLABEkQ,1472 -deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/__init__.py,sha256=N0qSPddtzqj_UkISW7rxEVXpdk8a5JItbXBjNAY_FBE,139 -deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/__pycache__/top_k_gating.cpython-310.pyc,, -deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.cpp,sha256=EBxWZ_Nm2ygSbFtnzI7kGMsl3nZxECXXefAVjMDDVx4,2435 -deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.cuh,sha256=1U1wNpJPmIhCiSwRi7xeRSR_zCcGnuXbFwBOz_gQplw,751 -deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.h,sha256=FdryRq5nESxX9HgzKIhnUgwfaoOzLir5O361gf-4g_M,556 -deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.py,sha256=-xgNLLmp6ea_QIMkGtuYDgt-fLfvq6wPVKgMWQpIHfA,2578 -deepspeed/inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating_cuda.cu,sha256=LCuHhCZhWYGT69EW4mq3sPYr9JafW9Ef1YIScTNCQC0,5130 -deepspeed/inference/v2/logging.py,sha256=H4Dp4vKNi42HZaNgXOAOG-Vki2oC3U3APVQwq3ay8kE,780 -deepspeed/inference/v2/model_implementations/__init__.py,sha256=g4p-wyfxcks_kDpGydBi2Sho_dI46rEfin8-xAydRd4,577 -deepspeed/inference/v2/model_implementations/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/__pycache__/flat_model_helpers.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/__pycache__/inference_model_base.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/__pycache__/inference_policy_base.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/__pycache__/inference_transformer_base.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/__pycache__/layer_container_base.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/__pycache__/parameter_base.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/common_parameters/__init__.py,sha256=_mp6NDvxO-3TM9LiiOTrMBRIweeiRjU6s6L-GpTvX9U,359 -deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/attn_output_parameters.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/embedding_parameters.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/invfreq_parameters.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/mlp_parameters.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/moe_parameters.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/norm_parameters.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/qkv_parameters.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/common_parameters/__pycache__/unembed_parameters.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/common_parameters/attn_output_parameters.py,sha256=VwOj8hRbIQQnWfoHVq-vbRK_QLt9f-uimARsTHfot4w,763 -deepspeed/inference/v2/model_implementations/common_parameters/embedding_parameters.py,sha256=k7CMnQ06UchamvQXsdeiTfxAkjpor3qZjTU0cD20hZU,605 -deepspeed/inference/v2/model_implementations/common_parameters/invfreq_parameters.py,sha256=Srz4b0mFK73OBzsw3EK3Aydj21Z1gZjLuyC8qfry6QI,401 -deepspeed/inference/v2/model_implementations/common_parameters/mlp_parameters.py,sha256=fMdVvya2UlX-DJ52QT0Mv8YKFukIAWx4F833ZoI0-IA,3363 -deepspeed/inference/v2/model_implementations/common_parameters/moe_parameters.py,sha256=hTKZ8zIERsvQ6i6zQRWe47qUcJmMAGOHfJinwpY3r04,2548 -deepspeed/inference/v2/model_implementations/common_parameters/norm_parameters.py,sha256=xTaLT5H81-Et8y9FivEaY8N9VQ0XDsvtlV4v1wnFUG4,454 -deepspeed/inference/v2/model_implementations/common_parameters/qkv_parameters.py,sha256=pxBm6EBMNLvbJTPAY9ABjdQ68-oxbY_nWb5CFGDdvdA,4086 -deepspeed/inference/v2/model_implementations/common_parameters/unembed_parameters.py,sha256=-gdVA8UMMvFsd-lSAvCbQp5mQ25j1sT5f7KC2mSVY_A,651 -deepspeed/inference/v2/model_implementations/falcon/__init__.py,sha256=IVA-W-wyht49JDQt7MTTsBMN4dHAAQxWPEULVVwkoVY,129 -deepspeed/inference/v2/model_implementations/falcon/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/falcon/__pycache__/container.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/falcon/__pycache__/model.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/falcon/__pycache__/policy.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/falcon/container.py,sha256=TNdqevl_yh4g06l80i3RJC88dExK2bfd_22NSA-yz0M,4479 -deepspeed/inference/v2/model_implementations/falcon/model.py,sha256=ELmk5fmti6cUmUqYRcsnykS746An6eEibPf0G3IQsgs,7768 -deepspeed/inference/v2/model_implementations/falcon/policy.py,sha256=oFyIvfzFEistRJCzXS-wPtFfUXDPqt9ANhPa7ywrKSQ,1342 -deepspeed/inference/v2/model_implementations/flat_model_helpers.py,sha256=XT6oMqJAfZhYLCr5KrPBxSFVW-WKnFjOi3FfEElshP4,9772 -deepspeed/inference/v2/model_implementations/inference_model_base.py,sha256=3YVtQafVIrZMI-kZadRKmQklafI5fjc38GTdTZeM3jQ,10083 -deepspeed/inference/v2/model_implementations/inference_policy_base.py,sha256=XKH3GcOtP8WAeAopwuQACflJA3cvX4F2MGybrIqcBPc,9481 -deepspeed/inference/v2/model_implementations/inference_transformer_base.py,sha256=SgXRyFYWj31DJMbDpgo06nEyAJLD-w3qKJF1xbAm9Jo,23592 -deepspeed/inference/v2/model_implementations/layer_container_base.py,sha256=qvvGjZ0aL6Ah4MDtjpowV76DKSpI8VuZyY_4hZRbB1o,15471 -deepspeed/inference/v2/model_implementations/llama_v2/__init__.py,sha256=J8aLmbt1KDVL_LOtSqaP-xACSxlgIDjeaIA9R-Grsgs,129 -deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/container.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/model.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/policy.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/llama_v2/container.py,sha256=blAioiSrdB9kH0zHFV0ZO1kWChv6-geaZW1ivEaSQx8,2664 -deepspeed/inference/v2/model_implementations/llama_v2/model.py,sha256=YmauqGYxtx2o_apo28aYvEhxi3SmYYvLm1IoEE7tGgs,7515 -deepspeed/inference/v2/model_implementations/llama_v2/policy.py,sha256=A1z_j19nj8yOURNNC29DeSyMatAjzuPUtcnIz09zZII,1148 -deepspeed/inference/v2/model_implementations/mistral/__init__.py,sha256=rzXGZBulBvo8ZNt7E5z1mSi8OlN1edNDTFlLCMCqo3M,130 -deepspeed/inference/v2/model_implementations/mistral/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/mistral/__pycache__/container.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/mistral/__pycache__/model.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/mistral/__pycache__/policy.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/mistral/container.py,sha256=0THLABZ6g4W4H_Kf2INGnyVrXHroT2uveKr-F2deP88,2784 -deepspeed/inference/v2/model_implementations/mistral/model.py,sha256=QhDoAoi2aMguJf4GJXqNyivqRKAjJLUwf5EsVgYHl38,7350 -deepspeed/inference/v2/model_implementations/mistral/policy.py,sha256=ceBkx-EOWH6UmcosmmhcE7anO5cDMaPGhLInK-x3PBA,1056 -deepspeed/inference/v2/model_implementations/mixtral/__init__.py,sha256=ywKUfNNkXZuDaMa9NT05H3MkYJQgiCWl3eM8kZMkhJE,130 -deepspeed/inference/v2/model_implementations/mixtral/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/mixtral/__pycache__/container.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/mixtral/__pycache__/model.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/mixtral/__pycache__/policy.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/mixtral/container.py,sha256=YMdy5flaitAQHdy71VBCi7n0vQBfWcPRyXTfZeABchA,1658 -deepspeed/inference/v2/model_implementations/mixtral/model.py,sha256=smAqT1nkiFmn2ipNNvJqd2BiIHHprI5O7iWKOFsKM4M,9096 -deepspeed/inference/v2/model_implementations/mixtral/policy.py,sha256=u1Am_kWG7UvEt8A5V4lTn_BjUZ4eSlnTB32R5agIYy4,1057 -deepspeed/inference/v2/model_implementations/opt/__init__.py,sha256=WeJqbtuxqlXaqif4F_4KpZnGBCwXwKEn0OwUZOcxmpo,126 -deepspeed/inference/v2/model_implementations/opt/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/opt/__pycache__/container.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/opt/__pycache__/model.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/opt/__pycache__/policy.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/opt/container.py,sha256=1jIU2Icv-lSEi_XMhaS4jGUVWFfbP-B2t2xf52eQVwY,3415 -deepspeed/inference/v2/model_implementations/opt/model.py,sha256=3FFAoq-iBo_F5ep0guXKXk6d5BnNOx0QZFiv2VILScs,7262 -deepspeed/inference/v2/model_implementations/opt/policy.py,sha256=9cqEPs_pM9gdzsBegbXJ7GHEgHfmD6n_AaRpjMH1yGE,1066 -deepspeed/inference/v2/model_implementations/parameter_base.py,sha256=FioZn7e4Gw2MUq_ov4popyXxD-dle2X0r7dYagAXmvg,9026 -deepspeed/inference/v2/model_implementations/phi/__init__.py,sha256=jKL7LDi4qrR3X3S6rcFpkls2K0wXVV6Qb3mIkiOtRKk,126 -deepspeed/inference/v2/model_implementations/phi/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/phi/__pycache__/containers.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/phi/__pycache__/model.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/phi/__pycache__/policy.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/phi/containers.py,sha256=L1S3qe2iHCv73EtWzb-IPn_dg8rxzwUn27HPXMatmCg,3197 -deepspeed/inference/v2/model_implementations/phi/model.py,sha256=b12uXjTpbHy_0EklvawD67KxwkrV-ooYSSdjNxGbz58,7251 -deepspeed/inference/v2/model_implementations/phi/policy.py,sha256=FIcCwm7nwUQOoFT1yGV_OTFd2xldUWmL_TqmpjtjUgU,1175 -deepspeed/inference/v2/model_implementations/phi3/__init__.py,sha256=97ojikdkWIzS9unY8jgler8jl3vaOBjGXpifKCtPXxY,127 -deepspeed/inference/v2/model_implementations/phi3/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/phi3/__pycache__/containers.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/phi3/__pycache__/model.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/phi3/__pycache__/policy.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/phi3/containers.py,sha256=hn-rEguXEgYbcu1GSYH7Pc4e6bOFTBQeEbGomMhmG0Y,2464 -deepspeed/inference/v2/model_implementations/phi3/model.py,sha256=0MtLieFL5uPP_eHcKdsAe9Q48ONQIWtPEYuB6OND5wQ,7296 -deepspeed/inference/v2/model_implementations/phi3/policy.py,sha256=ytHLbsKAqQN27Lku715CNU5bMg71EohF5W2-4P7v0lU,1033 -deepspeed/inference/v2/model_implementations/qwen/__init__.py,sha256=1TgzvNCP0loWbohqWRMH9NXCKKjXVVagBtatz_NCHc8,127 -deepspeed/inference/v2/model_implementations/qwen/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/qwen/__pycache__/container.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/qwen/__pycache__/model.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/qwen/__pycache__/policy.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/qwen/container.py,sha256=4fnYCNY92cDw116OyEVhmp_fuFtfeeyq4ZG4T5NHVi0,2362 -deepspeed/inference/v2/model_implementations/qwen/model.py,sha256=cCoVWHR2t_9LrLSwYE2LTCJa6nxXj9-l-Qr5q819Iq0,8112 -deepspeed/inference/v2/model_implementations/qwen/policy.py,sha256=0nP7dboH8FSnh6xvImRKK0gVWkw2iOTImwCJw73V-OA,1066 -deepspeed/inference/v2/model_implementations/qwen_v2/__init__.py,sha256=s0OcwPN9W24-8vYSQ8kztzky9jEwOvER9aMR60Xejuk,128 -deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/container.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/model.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/policy.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/qwen_v2/container.py,sha256=HpDpzng1cMtevhB-XykEKvmPPXENaDRV6BiLVYM7wjo,2817 -deepspeed/inference/v2/model_implementations/qwen_v2/model.py,sha256=EENwSZm-DZRQrv64UCLoEHop_U8xZ9yVejtboPzri_8,8029 -deepspeed/inference/v2/model_implementations/qwen_v2/policy.py,sha256=U9rMVvegwkKMqw6SoFXGZXobYTrjfpfySUryrpF0TDs,1140 -deepspeed/inference/v2/model_implementations/qwen_v2_moe/__init__.py,sha256=nxhwmjwG8giEHatkVAzPoQDEzgLgXCOcxluBG7Y4cn8,131 -deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/container.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/model.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/policy.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/qwen_v2_moe/container.py,sha256=QK62BZaTWjQACeHUJklh_qVJgRbX_N9WGhaOAtY-8dQ,4087 -deepspeed/inference/v2/model_implementations/qwen_v2_moe/model.py,sha256=bcL2sE2G5mHIPk1EZHT3vwFqNEiaBNmosSsKm1d1bjY,13516 -deepspeed/inference/v2/model_implementations/qwen_v2_moe/policy.py,sha256=k_HNyspBw8UG9mjCnCilY8khGf5dp8TI_74Du7eSwCM,1064 -deepspeed/inference/v2/model_implementations/sharding/__init__.py,sha256=M0NsMjUqO0Rj0PAPm2lyRepH5lc3uFQswNM5Gr3QAcI,247 -deepspeed/inference/v2/model_implementations/sharding/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/sharding/__pycache__/attn.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/sharding/__pycache__/attn_out.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/sharding/__pycache__/embedding.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/sharding/__pycache__/mlp.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/sharding/__pycache__/qkv.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/sharding/__pycache__/types.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/sharding/__pycache__/unembed.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/sharding/__pycache__/utils.cpython-310.pyc,, -deepspeed/inference/v2/model_implementations/sharding/attn.py,sha256=QAapGaB4Bp1k9xI53-QhhyKFoaeIPandzI1qnPT2rxE,2376 -deepspeed/inference/v2/model_implementations/sharding/attn_out.py,sha256=MWf6O0P55Am0x947474ZqGcOEuZn2nGptKcNPCxfS0c,4909 -deepspeed/inference/v2/model_implementations/sharding/embedding.py,sha256=mJgZR289D6_ODuTZm4pVsedyJEHaXH2a4ncI8_7YdG0,1270 -deepspeed/inference/v2/model_implementations/sharding/mlp.py,sha256=UIpBwlbU9xDNv1iDbIfcxqPgcAf_WvO42NxtMGyZ1QQ,2895 -deepspeed/inference/v2/model_implementations/sharding/qkv.py,sha256=pQfXbEPeqEZ9eCJwcdSvJsid5jeBJ2jptWhqVaPJSPM,7609 -deepspeed/inference/v2/model_implementations/sharding/types.py,sha256=YqDHnyDLGhy2sf37n-a0eHIIXqcQ8OjHKHQ-rcKAl8E,575 -deepspeed/inference/v2/model_implementations/sharding/unembed.py,sha256=1AmukReQ6SAfFlq2YdiQwJrrvPLVx7fhel7KKCGXXjQ,1625 -deepspeed/inference/v2/model_implementations/sharding/utils.py,sha256=wKlbKIwSyblinL1xupsHDXodws2zfWhxt-flmac1grc,5066 -deepspeed/inference/v2/modules/__init__.py,sha256=oZhP5-VwkgPMtHVFO_9MefXPnSEfcHdQeIXaRg3VOeY,193 -deepspeed/inference/v2/modules/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/modules/__pycache__/ds_module.cpython-310.pyc,, -deepspeed/inference/v2/modules/__pycache__/heuristics.cpython-310.pyc,, -deepspeed/inference/v2/modules/__pycache__/module_registry.cpython-310.pyc,, -deepspeed/inference/v2/modules/configs/__init__.py,sha256=IJ5d0k84imLnM7jh0UcZ9k3_ITPU7YI_6ZoGRpEJcqs,449 -deepspeed/inference/v2/modules/configs/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/modules/configs/__pycache__/attention_configs.cpython-310.pyc,, -deepspeed/inference/v2/modules/configs/__pycache__/embedding_config.cpython-310.pyc,, -deepspeed/inference/v2/modules/configs/__pycache__/linear_config.cpython-310.pyc,, -deepspeed/inference/v2/modules/configs/__pycache__/moe_config.cpython-310.pyc,, -deepspeed/inference/v2/modules/configs/__pycache__/norm_config.cpython-310.pyc,, -deepspeed/inference/v2/modules/configs/__pycache__/unembed_config.cpython-310.pyc,, -deepspeed/inference/v2/modules/configs/attention_configs.py,sha256=IWbmcmCc3CdL1E4FEUhrN9xU9wfbEaswUD8a0cM_D3Y,2834 -deepspeed/inference/v2/modules/configs/embedding_config.py,sha256=PuRrpOWLwtqjZRg9jumeMHuqzzsTcrbCjRxKL2L6xww,1812 -deepspeed/inference/v2/modules/configs/linear_config.py,sha256=fz9m8RIL6NVjbkq06-0WSP4-7IxIdSVQi1nRY4WQdTc,1301 -deepspeed/inference/v2/modules/configs/moe_config.py,sha256=EqX1OLq_QGt4IrSGhAJUXfVvy0MqjZk2E3toeUkb7kA,1289 -deepspeed/inference/v2/modules/configs/norm_config.py,sha256=iTKxfBxWcHs2JiiuxNZU356DGJR-a0OgDW3p5oshedk,839 -deepspeed/inference/v2/modules/configs/unembed_config.py,sha256=2E9XPh8VSRJckwng_DwbQksLMXtWZMFiXUjo5GWhzqI,822 -deepspeed/inference/v2/modules/ds_module.py,sha256=1h8X2TvYvDsUyEUjllPlziT7XOqwRxP02GBagzJw7FU,1856 -deepspeed/inference/v2/modules/heuristics.py,sha256=eRitZ_RX2yg4KOgmjiJaTEDZfH-zauMDIaf3WVQ2rGc,7978 -deepspeed/inference/v2/modules/implementations/__init__.py,sha256=0rEv4jixmtnw_sIe_Y5UKaIEcG7qwALUSzx1GNORTdk,289 -deepspeed/inference/v2/modules/implementations/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/attention/__init__.py,sha256=tNj2_fpavmUWaI5VAkUF0I7k0OxmIlZRNZZl7JIqrR4,157 -deepspeed/inference/v2/modules/implementations/attention/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/attention/__pycache__/dense_blocked_attention.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/attention/dense_blocked_attention.py,sha256=JgS7E2TzKmWylwRYWiYR-F7ldLdg8qhGr8Q9bXSPoCA,7735 -deepspeed/inference/v2/modules/implementations/embedding/__init__.py,sha256=CUiHnEi77BVJmH3aoJOanq5mTJIvEcIo3cHcbkddDHs,144 -deepspeed/inference/v2/modules/implementations/embedding/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/embedding/__pycache__/ragged_embedding.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/embedding/ragged_embedding.py,sha256=zddQlHcpjLe2_dKQUTp_5VfJwTP6L-BRk8HUHD2vPRo,2740 -deepspeed/inference/v2/modules/implementations/linear/__init__.py,sha256=FXT-bP-U4YlbMbBDrgwkpMPTHps3NiztMczMulIgmPU,203 -deepspeed/inference/v2/modules/implementations/linear/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/linear/__pycache__/blas_fp_linear.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/linear/__pycache__/quantized_linear.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/linear/blas_fp_linear.py,sha256=nHZHiLKVcppWyn0hNrQq4sHrS9UYKlHunrwL92DUVkQ,3629 -deepspeed/inference/v2/modules/implementations/linear/quantized_linear.py,sha256=5CQEpEBNq-Tzyt5dH0OK3RzZJUWjAxBieqzuC2OFJgw,7772 -deepspeed/inference/v2/modules/implementations/moe/__init__.py,sha256=XP3Ddz7eBtLvpIGum9u3gzUhl61w0W_jNfopcMI3QRE,143 -deepspeed/inference/v2/modules/implementations/moe/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/moe/__pycache__/cutlass_multi_gemm.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/moe/cutlass_multi_gemm.py,sha256=OzMYYHO7YIT5zOrAqfndOohQgJGwM4oHiPT_1mJU4EU,11126 -deepspeed/inference/v2/modules/implementations/post_norm/__init__.py,sha256=4hE4KqEcl74cOWm1kvexksZwHMK8rk-OQMzJYkdmgQo,141 -deepspeed/inference/v2/modules/implementations/post_norm/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/post_norm/__pycache__/cuda_post_ln.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/post_norm/cuda_post_ln.py,sha256=DNHGpJ2wNC8P6vyyk8M8hDpuXpA7k-ni78W3olRKO18,2086 -deepspeed/inference/v2/modules/implementations/pre_norm/__init__.py,sha256=_yJ5_1apmG5KriWHdy9klvL5H0vFzmLyWnZf1VIrmrA,184 -deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/cuda_pre_ln.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/cuda_pre_rms.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/pre_norm/cuda_pre_ln.py,sha256=VvFlm3Xc6U7VVyfXnFUeA5tsLw4TLhK_LVlKJkN4SmY,2744 -deepspeed/inference/v2/modules/implementations/pre_norm/cuda_pre_rms.py,sha256=WxZeP0B1ghhpOiMIHig4vb0ou0q7jkiLZFfbbdREXtQ,3245 -deepspeed/inference/v2/modules/implementations/unembed/__init__.py,sha256=ve_166eLGPGfIncTIvEvqoFGKi2pvtgvSLVVvTPXpHQ,140 -deepspeed/inference/v2/modules/implementations/unembed/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/unembed/__pycache__/ragged_unembed.cpython-310.pyc,, -deepspeed/inference/v2/modules/implementations/unembed/ragged_unembed.py,sha256=cAd3By6SBykoCHJI3oxwE0FXHcMP5HEF5CSrWJyBjjg,4921 -deepspeed/inference/v2/modules/interfaces/__init__.py,sha256=2djwkw0m2KAAuBV-LuWY0W0ncwIrCxDX5uwf4oc8E-U,519 -deepspeed/inference/v2/modules/interfaces/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/modules/interfaces/__pycache__/attention_base.cpython-310.pyc,, -deepspeed/inference/v2/modules/interfaces/__pycache__/embedding_base.cpython-310.pyc,, -deepspeed/inference/v2/modules/interfaces/__pycache__/linear_base.cpython-310.pyc,, -deepspeed/inference/v2/modules/interfaces/__pycache__/moe_base.cpython-310.pyc,, -deepspeed/inference/v2/modules/interfaces/__pycache__/post_norm_base.cpython-310.pyc,, -deepspeed/inference/v2/modules/interfaces/__pycache__/pre_norm_base.cpython-310.pyc,, -deepspeed/inference/v2/modules/interfaces/__pycache__/unembed_base.cpython-310.pyc,, -deepspeed/inference/v2/modules/interfaces/attention_base.py,sha256=Kl8aPbpzyhP4bF21LZ_OcEwHFoI1KFirOiukiV4YajE,3613 -deepspeed/inference/v2/modules/interfaces/embedding_base.py,sha256=3pXB9F4IlWMqFGV8ZYE3k3bG2Je4B9uvVnVS0pRzLeU,3071 -deepspeed/inference/v2/modules/interfaces/linear_base.py,sha256=8g8or7xUMtGX_3uQj6zt0YfUBS9SwSbRY02lrr3gUVE,2146 -deepspeed/inference/v2/modules/interfaces/moe_base.py,sha256=5BCoKWZplUfxMj7mOPA58b4aBGk7APM-WG-ByxUIUyA,2850 -deepspeed/inference/v2/modules/interfaces/post_norm_base.py,sha256=PktjhFD5Pi3y2zt_VEwB9L-yNVyIo4dHpJgdVwgbcVc,2123 -deepspeed/inference/v2/modules/interfaces/pre_norm_base.py,sha256=QlYoS3N4N3ZKfUJk6C9u_vCT10p4-ACVgAEydLquHnY,2106 -deepspeed/inference/v2/modules/interfaces/unembed_base.py,sha256=W5ToOG46DZwH5e1VnFXxXp_f9-EvnO2iIw6zqZkL38w,2055 -deepspeed/inference/v2/modules/module_registry.py,sha256=UyDd-5FYQwbEJ7ihPZ60GkEiQ7M6Z2XMWnEAaLZN_WQ,2073 -deepspeed/inference/v2/ragged/__init__.py,sha256=2fFLZ35zp9WgcDaQ3XaTti8Hzi99zt8EgueCaaOqZs4,418 -deepspeed/inference/v2/ragged/__pycache__/__init__.cpython-310.pyc,, -deepspeed/inference/v2/ragged/__pycache__/blocked_allocator.cpython-310.pyc,, -deepspeed/inference/v2/ragged/__pycache__/kv_cache.cpython-310.pyc,, -deepspeed/inference/v2/ragged/__pycache__/manager_configs.cpython-310.pyc,, -deepspeed/inference/v2/ragged/__pycache__/ragged_manager.cpython-310.pyc,, -deepspeed/inference/v2/ragged/__pycache__/ragged_wrapper.cpython-310.pyc,, -deepspeed/inference/v2/ragged/__pycache__/sequence_descriptor.cpython-310.pyc,, -deepspeed/inference/v2/ragged/blocked_allocator.py,sha256=8rwqli63wrLAbzOuLoOgVqLi5UlpPlyxXtyRS1gU6zM,3661 -deepspeed/inference/v2/ragged/csrc/fast_host_buffer.cu,sha256=8U10U6VZk12oGwXCniQsqduta_cEcVBFuN9Fq3LO1JQ,501 -deepspeed/inference/v2/ragged/csrc/ragged_ops.cpp,sha256=Z6na7zRDPD7PrYbZ4lBhBsGyMatZ9yUNYw_PBzRGHw8,2726 -deepspeed/inference/v2/ragged/includes/fast_host_buffer.h,sha256=N73g7w1uE1gEsyWzrEnEibGeU8nBJPDeL-SWJEWK_0M,302 -deepspeed/inference/v2/ragged/kv_cache.py,sha256=eekwPz2OMNgQJHYXcxIrW86yK0guS419GeNqXJ38wIg,8559 -deepspeed/inference/v2/ragged/manager_configs.py,sha256=_SDSRv5VC5ZdmN6VDckjHtuC4HdJK6KjH9Vec839WcM,5965 -deepspeed/inference/v2/ragged/ragged_manager.py,sha256=Mk09dkqkEj5CMA9kWvTWuQeiw7fWiO9xn_nMIs-1Hho,7334 -deepspeed/inference/v2/ragged/ragged_wrapper.py,sha256=RJJpqSVz6vsElbtfePx99TDYZ2Q_Hao7qt_2Jt6RP2s,12701 -deepspeed/inference/v2/ragged/sequence_descriptor.py,sha256=oo1kf0upitTJkbZB6ME7GY3_MoOnGABXcw1coq49ljM,10891 -deepspeed/inference/v2/scheduling_utils.py,sha256=S1MPmx8zU8qHmDd_eDLhKP6hh71N9lo_Q8kJxO1AyOY,1363 -deepspeed/io/__init__.py,sha256=948yAWx0rWW6wzxU7f2JTmwF88rNisJdS1FoPs_JNZc,249 -deepspeed/io/__pycache__/__init__.cpython-310.pyc,, -deepspeed/io/__pycache__/base_file_writer.cpython-310.pyc,, -deepspeed/io/__pycache__/base_io_buffer.cpython-310.pyc,, -deepspeed/io/__pycache__/constants.cpython-310.pyc,, -deepspeed/io/__pycache__/double_io_buffer.cpython-310.pyc,, -deepspeed/io/__pycache__/fast_file_writer.cpython-310.pyc,, -deepspeed/io/__pycache__/mock_file_writer.cpython-310.pyc,, -deepspeed/io/__pycache__/py_file_writer.cpython-310.pyc,, -deepspeed/io/__pycache__/single_io_buffer.cpython-310.pyc,, -deepspeed/io/__pycache__/utils.cpython-310.pyc,, -deepspeed/io/base_file_writer.py,sha256=nXR4OuhLiHga6-Colhu578n97LP8sEkXCwKfFKOZ5rw,1007 -deepspeed/io/base_io_buffer.py,sha256=C-FVSF4w6wc7lCZ9VXkRhwBnW75TUE-n_tTko4YplnU,1841 -deepspeed/io/constants.py,sha256=9nxuv0_QNaB_1CN9d9FXvlEPrZ2fEHFL3yuo8eoshZg,806 -deepspeed/io/double_io_buffer.py,sha256=6qRY3Je1o2onzvXIP50KZdfJfzraYzBCBlRJoC-K0xM,3028 -deepspeed/io/fast_file_writer.py,sha256=lh8RGJc9JAKKNrzKkyOh28DCpLoh_DHqUatNbC-WJIU,11049 -deepspeed/io/mock_file_writer.py,sha256=pWXxx5o7DyV1xKsAOiWoSUH9ZIYDIlvRtlS8JHTkj-4,1638 -deepspeed/io/py_file_writer.py,sha256=knJ292prjk4AdfbbSiDEK0DGoBsZU6rMG_RYqo7KcAU,906 -deepspeed/io/single_io_buffer.py,sha256=CeUIj3GO-CnfTWHA9wrqe89W_PRjOF7GjXJYuCE36i4,1324 -deepspeed/io/utils.py,sha256=D0IOzZmhByrqOBQxP5m3KvA2HGanM3tV0OxZXQyMyas,1464 -deepspeed/launcher/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/launcher/__pycache__/__init__.cpython-310.pyc,, -deepspeed/launcher/__pycache__/constants.cpython-310.pyc,, -deepspeed/launcher/__pycache__/launch.cpython-310.pyc,, -deepspeed/launcher/__pycache__/launcher_helper.cpython-310.pyc,, -deepspeed/launcher/__pycache__/multinode_runner.cpython-310.pyc,, -deepspeed/launcher/__pycache__/runner.cpython-310.pyc,, -deepspeed/launcher/constants.py,sha256=X-6bI3aC4NXSZgrMJ6LaIILrsWu_TNteNBwgKOhNrgI,375 -deepspeed/launcher/launch.py,sha256=aHutEg7kj_kY5VdH_CQLclyYOIZoh2YSRLD-wXvM_BI,14898 -deepspeed/launcher/launcher_helper.py,sha256=tpvx6ItaLbBNYnUwP8XoKcn8P0mFXRffpB1Z_uLB_U0,3931 -deepspeed/launcher/multinode_runner.py,sha256=QcTvKOnkP94gmjLrK9om9Upw8FKClAmSLyCEtKQFKjU,18797 -deepspeed/launcher/runner.py,sha256=3u_dE4Q1e0bTmFjwk85JVRRWwbRmnzYmnloS0crJS_0,26306 -deepspeed/linear/__init__.py,sha256=fwQSoiw98NsTIwp6BuhtNMykJsaJOpJWd5QAUpitfcc,238 -deepspeed/linear/__pycache__/__init__.cpython-310.pyc,, -deepspeed/linear/__pycache__/config.cpython-310.pyc,, -deepspeed/linear/__pycache__/context_manager.cpython-310.pyc,, -deepspeed/linear/__pycache__/optimized_linear.cpython-310.pyc,, -deepspeed/linear/__pycache__/quantization.cpython-310.pyc,, -deepspeed/linear/config.py,sha256=U4yr-hEialc28KiEp4qRfm-8H_SF5xsLij2DROOTGSQ,2178 -deepspeed/linear/context_manager.py,sha256=L_aADkH1e59C7LCpoS7XyxC3rJ3QBSBqXucrVOjy1ss,3574 -deepspeed/linear/optimized_linear.py,sha256=RcpH8q8QpB7hE85DqU_evn29xVxKcWxOYzvx8aiUaUQ,10474 -deepspeed/linear/quantization.py,sha256=Lpvwk_-yb3EQpsNQC1sFT9fGvlj6dPdQx5fg0Rxb0s8,6352 -deepspeed/model_implementations/__init__.py,sha256=jJCJcnbvmFc9eQxIUIb7Gd6Ysd8ijSuuk-aaqFrDkK4,220 -deepspeed/model_implementations/__pycache__/__init__.cpython-310.pyc,, -deepspeed/model_implementations/diffusers/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/model_implementations/diffusers/__pycache__/__init__.cpython-310.pyc,, -deepspeed/model_implementations/diffusers/__pycache__/unet.cpython-310.pyc,, -deepspeed/model_implementations/diffusers/__pycache__/vae.cpython-310.pyc,, -deepspeed/model_implementations/diffusers/unet.py,sha256=W2_9-c-5Tz-nymu7kK2du2rzpV8-eC8JZEpQGRhGMek,3056 -deepspeed/model_implementations/diffusers/vae.py,sha256=xbjGaYsXeJA4nu49gQaGO5GrEazCHwcTUVqPsxoOTn4,6244 -deepspeed/model_implementations/features/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/model_implementations/features/__pycache__/__init__.cpython-310.pyc,, -deepspeed/model_implementations/features/__pycache__/cuda_graph.cpython-310.pyc,, -deepspeed/model_implementations/features/cuda_graph.py,sha256=-KgILcht5qw_ayzI0CqSPGGFoqZKHG8E5CTfD8ACgw8,563 -deepspeed/model_implementations/transformers/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/model_implementations/transformers/__pycache__/__init__.cpython-310.pyc,, -deepspeed/model_implementations/transformers/__pycache__/clip_encoder.cpython-310.pyc,, -deepspeed/model_implementations/transformers/__pycache__/ds_base.cpython-310.pyc,, -deepspeed/model_implementations/transformers/__pycache__/ds_bert.cpython-310.pyc,, -deepspeed/model_implementations/transformers/__pycache__/ds_bloom.cpython-310.pyc,, -deepspeed/model_implementations/transformers/__pycache__/ds_gpt.cpython-310.pyc,, -deepspeed/model_implementations/transformers/__pycache__/ds_llama2.cpython-310.pyc,, -deepspeed/model_implementations/transformers/__pycache__/ds_megatron_gpt.cpython-310.pyc,, -deepspeed/model_implementations/transformers/__pycache__/ds_opt.cpython-310.pyc,, -deepspeed/model_implementations/transformers/__pycache__/ds_transformer.cpython-310.pyc,, -deepspeed/model_implementations/transformers/clip_encoder.py,sha256=BTc_do6u4A79uhoNwQYxKrbqSnk2C63BO6xN0Ez-fz0,3096 -deepspeed/model_implementations/transformers/ds_base.py,sha256=feBwoCtVctmdYb_KZruLVCJr7WUBsfNZtVUxEKtO5MQ,388 -deepspeed/model_implementations/transformers/ds_bert.py,sha256=fckfXCn3zf9DsvqiyVAb0pwn25GA9H4ZceaCI4gLLSI,667 -deepspeed/model_implementations/transformers/ds_bloom.py,sha256=_YHlkBOUBTT27ZRWs7x6DxKz_5gILGBAe2YNmlkCij0,669 -deepspeed/model_implementations/transformers/ds_gpt.py,sha256=PkhQU0iHFkSGRxaEb_j1UDvRWruwrOPF8MLZfAo55ZY,665 -deepspeed/model_implementations/transformers/ds_llama2.py,sha256=geGAyCXgRNzMlk0xhvqu3FFmHfFgotVXT6aVV4ZlsYo,2128 -deepspeed/model_implementations/transformers/ds_megatron_gpt.py,sha256=kVtyp2Uss9Z_MRGwgAvJI_2tfWlPgzMpywIPS0xeOIo,682 -deepspeed/model_implementations/transformers/ds_opt.py,sha256=oHS210T27R5zOtg8ik_QwV22F8BIJxKVWPiPQaDXzjY,665 -deepspeed/model_implementations/transformers/ds_transformer.py,sha256=y981LDa0stEBTe8BW_FnQ_1WxHoJSCHNt7FrLiodBzo,8659 -deepspeed/module_inject/__init__.py,sha256=AechzL-WB6riK3mtss4gUD4HD79MhxZzwmPpZzBJSSk,461 -deepspeed/module_inject/__pycache__/__init__.cpython-310.pyc,, -deepspeed/module_inject/__pycache__/auto_tp.cpython-310.pyc,, -deepspeed/module_inject/__pycache__/auto_tp_model_utils.cpython-310.pyc,, -deepspeed/module_inject/__pycache__/fusedqkv_utils.cpython-310.pyc,, -deepspeed/module_inject/__pycache__/inject.cpython-310.pyc,, -deepspeed/module_inject/__pycache__/layers.cpython-310.pyc,, -deepspeed/module_inject/__pycache__/load_checkpoint.cpython-310.pyc,, -deepspeed/module_inject/__pycache__/module_quantize.cpython-310.pyc,, -deepspeed/module_inject/__pycache__/policy.cpython-310.pyc,, -deepspeed/module_inject/__pycache__/replace_module.cpython-310.pyc,, -deepspeed/module_inject/__pycache__/replace_policy.cpython-310.pyc,, -deepspeed/module_inject/__pycache__/tp_shard.cpython-310.pyc,, -deepspeed/module_inject/__pycache__/utils.cpython-310.pyc,, -deepspeed/module_inject/auto_tp.py,sha256=KAyjZgGCmaQ5_HRgUQvjBgUikmecpnXZw9q62eYQFyw,24809 -deepspeed/module_inject/auto_tp_model_utils.py,sha256=G_tPVAaT7b422Wa8J785rEOcM742uXnoW-5fvd1cinU,5910 -deepspeed/module_inject/containers/__init__.py,sha256=AQbuSXpxEyLpLN3RoAkQeyvtioX_8ShH6C8cqzcQ2Jw,1015 -deepspeed/module_inject/containers/__pycache__/__init__.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/base.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/base_moe.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/bert.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/bloom.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/clip.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/distil_bert.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/gpt2.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/gptj.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/gptneo.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/gptneox.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/internlm.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/llama.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/llama2.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/megatron_gpt.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/megatron_gpt_moe.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/opt.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/unet.cpython-310.pyc,, -deepspeed/module_inject/containers/__pycache__/vae.cpython-310.pyc,, -deepspeed/module_inject/containers/base.py,sha256=kw-culwCQ480zVCgeY7K9LH30BSuKWuPtm2lBMh-L_o,13762 -deepspeed/module_inject/containers/base_moe.py,sha256=lfyM4AhAXmYp5dxa60XoEbcJDQVj9uB-pRlzIDR2ypI,5756 -deepspeed/module_inject/containers/bert.py,sha256=f0VMB8R5mCHxNgIVhQSYeTgVsxo1IN4BV6vzSIOgLZc,3769 -deepspeed/module_inject/containers/bloom.py,sha256=fV_vVFyjB0BB95kP7TrjwqrqruKbEwWQH6ttduzfcwg,6256 -deepspeed/module_inject/containers/clip.py,sha256=KR3s7ev3v-ccvrFhqfm2pKjApQeYeCfhm4aZZ0_Ou_A,2822 -deepspeed/module_inject/containers/distil_bert.py,sha256=j8ufi0mXYe862fLYMtY1J0Rh2uY-xzkXKDKrhK-EsgI,3188 -deepspeed/module_inject/containers/features/__init__.py,sha256=GOivP3ix552VHgGPZ70UehyLQ_YZ8TLRc0-rgTCO07c,275 -deepspeed/module_inject/containers/features/__pycache__/__init__.cpython-310.pyc,, -deepspeed/module_inject/containers/features/__pycache__/gated_mlp.cpython-310.pyc,, -deepspeed/module_inject/containers/features/__pycache__/hybrid_engine.cpython-310.pyc,, -deepspeed/module_inject/containers/features/__pycache__/hybrid_megatron.cpython-310.pyc,, -deepspeed/module_inject/containers/features/__pycache__/megatron.cpython-310.pyc,, -deepspeed/module_inject/containers/features/__pycache__/meta_tensor.cpython-310.pyc,, -deepspeed/module_inject/containers/features/__pycache__/split_qkv.cpython-310.pyc,, -deepspeed/module_inject/containers/features/gated_mlp.py,sha256=ep8fnRs-06C7oL75w5UBAaLyCb9BwAx2v3c-sBhkB2k,5397 -deepspeed/module_inject/containers/features/hybrid_engine.py,sha256=ki_k_0SstM39XfN0SNWwiSO3qzIBs2qKxxAyrGxpfUo,8374 -deepspeed/module_inject/containers/features/hybrid_megatron.py,sha256=qp04fT0ManTJwzVdABFP6OWvLQsD6pWDaJAxNJU4-ZQ,4129 -deepspeed/module_inject/containers/features/megatron.py,sha256=BrbmtBC9ZKFL4ESbIYRJLzoEnEp3QOhkj1VcoqlcEek,1200 -deepspeed/module_inject/containers/features/meta_tensor.py,sha256=oem92mPilw6_UOBzkE-7VMoi9vpA9kvA1lemDcCSPwQ,2928 -deepspeed/module_inject/containers/features/split_qkv.py,sha256=totYVBg1M6DcsdC3762Vhoq7gpGJ1BIioTfDFdFwkF0,7136 -deepspeed/module_inject/containers/gpt2.py,sha256=98yUKFAGrfyyCv5yK7V4Zr8oWf6m4qTXgGIzUbRG0kM,2221 -deepspeed/module_inject/containers/gptj.py,sha256=OJ4Z-G5sjd-92yrNDJ0D53Jz_o0n2tfkg9fS9Mll1ss,5074 -deepspeed/module_inject/containers/gptneo.py,sha256=NakhEgR8qh-U9rZ6JD6nIyY_r1mH2obkrgLat4eKpMo,5790 -deepspeed/module_inject/containers/gptneox.py,sha256=zM8hYq0e1FKrDd0n0m1oUPKKMDNpwxVmXQHm8jlv1CA,5897 -deepspeed/module_inject/containers/internlm.py,sha256=3JYddGlL-87jC35MSKvYw9XzZFHYIIYC4Un0Ek5jsL4,7761 -deepspeed/module_inject/containers/llama.py,sha256=II9F76RipDydYbh6R-KYtDekTTQvVcAAHzgizkjT5Z4,6566 -deepspeed/module_inject/containers/llama2.py,sha256=5WbtiVVAL1dVseHiO5kvx2qOwVAAvq-cKBKo0JZ7YjA,6359 -deepspeed/module_inject/containers/megatron_gpt.py,sha256=PiD58-Ul61ZJE4QAkCKCtGwjg44wMaSPOkqnWcRudyo,5417 -deepspeed/module_inject/containers/megatron_gpt_moe.py,sha256=xFOff3OOIKMz2YAFRrqLHxpkzLgB-chA4L69HXiey98,3936 -deepspeed/module_inject/containers/opt.py,sha256=exIuEYIu-KCjwPL3L-ykz1ud477q-F4WBwbXIL0Hef4,6905 -deepspeed/module_inject/containers/unet.py,sha256=fWteUHNx0S8u6C3xpiU2TLK4snlTr7Bv81yCNlVIDY8,1862 -deepspeed/module_inject/containers/vae.py,sha256=yLb5XWwZcaVOXVbUjGW1n0hUr3wrL4rqzqFkvBbD1Y4,1505 -deepspeed/module_inject/fusedqkv_utils.py,sha256=GvycFptz45D6k8eCPqJ-jJs4Vui8QzoyqK0CLVqbICs,10188 -deepspeed/module_inject/inject.py,sha256=KkpC_LLybQeiFzIhr6pz-OlKOl92XNFf3Vh7GqD7jsY,4719 -deepspeed/module_inject/layers.py,sha256=TcMmFqGcapbLQLiZUSlpBFuavpo3PB8V6TCbY7JZ5Bc,31475 -deepspeed/module_inject/load_checkpoint.py,sha256=BwyqWbY18JGZJv709sKG3_ZSUP1lGUjFnP8KYDpESgM,15434 -deepspeed/module_inject/module_quantize.py,sha256=NrOwvSfPWWOnlUVe82sSOZhay76JzqkgZ-j4NP-yEOY,3144 -deepspeed/module_inject/policy.py,sha256=kfyIzx4R4INZve94RPBPZu0jnnQKYpWtou-nRvZQ0Z0,8259 -deepspeed/module_inject/replace_module.py,sha256=_pfrNWVj286Qv4gQYglPWSiMvumKjBj2Tkam7CH4cL4,33068 -deepspeed/module_inject/replace_policy.py,sha256=jlIhG7pbXcjD8JPjWJCU-Gs_WSTG_q7K-IHmewu04pw,1119 -deepspeed/module_inject/tp_shard.py,sha256=n7n9vJYwPMuJVJjN5UaMiDwiPI-xTDOi8oaJMRNKOxg,2126 -deepspeed/module_inject/utils.py,sha256=q5gZWY7YK_-BI7ce16-G8dvU9Qgx7Qbezz9Hwz7oKS0,1995 -deepspeed/moe/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/moe/__pycache__/__init__.cpython-310.pyc,, -deepspeed/moe/__pycache__/experts.cpython-310.pyc,, -deepspeed/moe/__pycache__/layer.cpython-310.pyc,, -deepspeed/moe/__pycache__/mappings.cpython-310.pyc,, -deepspeed/moe/__pycache__/sharded_moe.cpython-310.pyc,, -deepspeed/moe/__pycache__/utils.cpython-310.pyc,, -deepspeed/moe/experts.py,sha256=7xFMeuKcBmfyVdaVKpUw5uyGzHYbtxXTiKqcUuBjrQ8,1315 -deepspeed/moe/layer.py,sha256=TayzIRGqEOakrq4TthCiBpqdwwdm5sGqISWMZG13KMs,6905 -deepspeed/moe/mappings.py,sha256=S-o5bwyGoZPmKL0i2E2TrWakyri_P7zARLgOY_4a8u8,4090 -deepspeed/moe/sharded_moe.py,sha256=7z2v3inspO0O8YSGWTll036Xxc3_TiB-vECG2wlKy9E,27332 -deepspeed/moe/utils.py,sha256=YGnqMJD26e2q3ezp-A0srHyb2sdWaHUEg1ld_I6qPgU,6699 -deepspeed/monitor/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/monitor/__pycache__/__init__.cpython-310.pyc,, -deepspeed/monitor/__pycache__/comet.cpython-310.pyc,, -deepspeed/monitor/__pycache__/config.cpython-310.pyc,, -deepspeed/monitor/__pycache__/csv_monitor.cpython-310.pyc,, -deepspeed/monitor/__pycache__/monitor.cpython-310.pyc,, -deepspeed/monitor/__pycache__/tensorboard.cpython-310.pyc,, -deepspeed/monitor/__pycache__/utils.cpython-310.pyc,, -deepspeed/monitor/__pycache__/wandb.cpython-310.pyc,, -deepspeed/monitor/comet.py,sha256=16xIXzwylfuGJyt4p3Lu46JXemCe89_2XCi_cLoWDlY,2903 -deepspeed/monitor/config.py,sha256=vdDaZ7qM0RLXIStjW6S2C4vJimCx0g0IQIi6Hv_sn4k,5258 -deepspeed/monitor/csv_monitor.py,sha256=E9bWWx7QKhex8GWOxMOLjOsrAyBmXqCkaxNUyVirSY0,2907 -deepspeed/monitor/monitor.py,sha256=d_Ke7YvSsXpXeyLmx61KaqkYpgG026sSSbMel0-lfVg,1894 -deepspeed/monitor/tensorboard.py,sha256=9crE_YD9adHS65QmRxAuSuezRQUpL_VHrJMZmPhmuiE,2227 -deepspeed/monitor/utils.py,sha256=MCyfYHyZL-F1a4jbEqG5_d2-1CvX81-gZC6JPlkIblY,1226 -deepspeed/monitor/wandb.py,sha256=G-2BnrWGjiuWG9Tg5-bzkBsjRLwlsyINAYryxOmMdvE,1128 -deepspeed/nebula/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/nebula/__pycache__/__init__.cpython-310.pyc,, -deepspeed/nebula/__pycache__/config.cpython-310.pyc,, -deepspeed/nebula/__pycache__/constants.cpython-310.pyc,, -deepspeed/nebula/config.py,sha256=THO1Mwto2utxnBWRfn0vIVWevOpT2EPFJFZkDCJF-ck,1764 -deepspeed/nebula/constants.py,sha256=0HQOkViV_lRv13vtczIhBJG_YAH876vhh0R4BqHAj68,2786 -deepspeed/nvme/__init__.py,sha256=iPQdRvD-GrL7bbRWcJtX0TFD7kz-OZVTmHDQa1mG3Z0,241 -deepspeed/nvme/__pycache__/__init__.cpython-310.pyc,, -deepspeed/nvme/__pycache__/ds_aio_args.cpython-310.pyc,, -deepspeed/nvme/__pycache__/ds_aio_basic.cpython-310.pyc,, -deepspeed/nvme/__pycache__/ds_aio_constants.cpython-310.pyc,, -deepspeed/nvme/__pycache__/ds_aio_handle.cpython-310.pyc,, -deepspeed/nvme/__pycache__/ds_aio_job.cpython-310.pyc,, -deepspeed/nvme/__pycache__/io_engine.cpython-310.pyc,, -deepspeed/nvme/__pycache__/parse_nvme_stats.cpython-310.pyc,, -deepspeed/nvme/__pycache__/perf_generate_param.cpython-310.pyc,, -deepspeed/nvme/__pycache__/perf_run_sweep.cpython-310.pyc,, -deepspeed/nvme/__pycache__/perf_sweep_utils.cpython-310.pyc,, -deepspeed/nvme/__pycache__/test_ds_aio.cpython-310.pyc,, -deepspeed/nvme/__pycache__/test_ds_aio_utils.cpython-310.pyc,, -deepspeed/nvme/__pycache__/torch_fastio_engine.cpython-310.pyc,, -deepspeed/nvme/__pycache__/torch_io.cpython-310.pyc,, -deepspeed/nvme/__pycache__/validate_async_io.cpython-310.pyc,, -deepspeed/nvme/ds_aio_args.py,sha256=gNjCXHv3VyyioJAWrMoAx4HJYmrN68VNNdq9pwly7SE,8458 -deepspeed/nvme/ds_aio_basic.py,sha256=WXImsGkd3E5qwuMOjZ6wHhCHQM7tzt7nKyg-LzeUpVg,2410 -deepspeed/nvme/ds_aio_constants.py,sha256=Af0dLTpRrw8VXmMby28zy60_tEWx61i15mpUSNzk51E,493 -deepspeed/nvme/ds_aio_handle.py,sha256=HS-sNFfUgwJn6T8o0pZjdK2gMQt4_xLzzHXu-M7PXtU,4944 -deepspeed/nvme/ds_aio_job.py,sha256=v-4U2oTjTqPZ7jjg1yIDtZOMlwfOQcOPMKCmzqTtAQQ,1276 -deepspeed/nvme/io_engine.py,sha256=nUJ1cXfLg6x3IQQ9g-zILY2I584beeigBaTHbsn2lPs,3847 -deepspeed/nvme/parse_nvme_stats.py,sha256=UjaSbCHcP0MifIngLT7mmWbf1BbkKRb-VYO8WbfBmqc,3835 -deepspeed/nvme/perf_generate_param.py,sha256=VPIleO_Ihe5JJRARe94DrXf_sWNCgC8RBPIwz7H_tYI,3289 -deepspeed/nvme/perf_run_sweep.py,sha256=PaLX4sfWHjbIgIJiG3W3t_Mgll_NfHCm6Qmp627Kwb4,9747 -deepspeed/nvme/perf_sweep_utils.py,sha256=pvreOkVAjNHwxSQ33bgXr4Yelwjzy-l9SEOcFTKGjhw,425 -deepspeed/nvme/test_ds_aio.py,sha256=0v_ZTRfro2gjANPa9CVoii1gz021Z5Eol1z9_hbLlR8,590 -deepspeed/nvme/test_ds_aio_utils.py,sha256=BORiau54UKygrVdLsqQLIq1e90ED3PsW1X6WI9DO2s8,3103 -deepspeed/nvme/torch_fastio_engine.py,sha256=zf4S1zsoE64jHUx9hrwYY-Ciy_6GSUC1O0x9rsL5DnQ,3509 -deepspeed/nvme/torch_io.py,sha256=gIq32rm0kRVYP_R-RFZphkmVVDv3TnJVYD8JjKx2gWg,2270 -deepspeed/nvme/validate_async_io.py,sha256=RsuUU_w27HuHWZAKsan5M0-pcllXRyf_H_VQcs-U4T8,302 -deepspeed/ops/__init__.py,sha256=jr72zRfZAdER7PTHnZc-9GziYpnxW80sBCyBthYWYhI,407 -deepspeed/ops/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/adagrad/__init__.py,sha256=aJHYZSC21yulH09u72ZYAe43nLgeMemqvmC9wNNh9CQ,141 -deepspeed/ops/adagrad/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/adagrad/__pycache__/cpu_adagrad.cpython-310.pyc,, -deepspeed/ops/adagrad/cpu_adagrad.py,sha256=nXZLRyBmJS4hcdt6XULXVrX7hQq09nrLFFajJTk5do8,4181 -deepspeed/ops/adam/__init__.py,sha256=ZVagnjkzHVw4akcDibUCfsArTcu8J9wsz9xKmbrJnIs,169 -deepspeed/ops/adam/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/adam/__pycache__/cpu_adam.cpython-310.pyc,, -deepspeed/ops/adam/__pycache__/fused_adam.cpython-310.pyc,, -deepspeed/ops/adam/__pycache__/multi_tensor_apply.cpython-310.pyc,, -deepspeed/ops/adam/cpu_adam.py,sha256=URu3D8FceMISyD_-49MgPYWRHJ2SS59deX7sGqnQAKQ,7520 -deepspeed/ops/adam/fused_adam.py,sha256=5_JBk_UvBFcnOL0_hPr5SVFATmmmKHFIlVSR1am-sBc,8767 -deepspeed/ops/adam/multi_tensor_apply.py,sha256=APt3UCnfw-nLjJOGWXC3izWRXNu-16TYhB-s213DpM0,429 -deepspeed/ops/aio/__init__.py,sha256=dHc8QXzdN4Cw0D7px-gbO_fzV-lKT21ySZJ67LtUJuI,136 -deepspeed/ops/aio/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/compile/__init__.py,sha256=6GIOlNEFXSRcqWtTL1JcczbQQcLh2Ab1oBVcnlnii0I,140 -deepspeed/ops/compile/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/csrc/adagrad/cpu_adagrad.cpp,sha256=FPSi-AtroLc-zPnzN1yF2OnKkb-wCpKALqN_e8wKTPI,7455 -deepspeed/ops/csrc/adam/cpu_adam.cpp,sha256=9bS2_kdvipPCUoonzJNR6CLuCjNviTPcZJlzaNqGxGY,409 -deepspeed/ops/csrc/adam/cpu_adam_impl.cpp,sha256=QycbyNXh4VGNq4FihHwBysu-_oHSWf1d-STXBgtKZFE,8468 -deepspeed/ops/csrc/adam/fused_adam_frontend.cpp,sha256=tu6pBcAvqcdh1ti6SB7fj8HnqFDQg3Rqbaa4xsQ2v7M,880 -deepspeed/ops/csrc/adam/multi_tensor_adam.cu,sha256=13ixr7zAf-vH8NzAPKYiKXPmVRgsJxILT6cj75_92j8,8417 -deepspeed/ops/csrc/adam/multi_tensor_apply.cuh,sha256=tcBmATZa_Et2VH9hnq5Jf6LJNnLaWDNrSqHdRr2yCE4,5609 -deepspeed/ops/csrc/aio/common/deepspeed_aio_common.cpp,sha256=hZ8EYNZ-Lcrll0VVnvmb-Gd7mr2CZ_P2a0F9_9bnexo,13274 -deepspeed/ops/csrc/aio/common/deepspeed_aio_common.h,sha256=8fp1mBNfXIAgE48vL5CEtH_n6PeP8TeVs3UXkm-4lXs,1358 -deepspeed/ops/csrc/aio/common/deepspeed_aio_types.cpp,sha256=vrR8UQ9EK-8EgHBUyQIMgPLxNSFtwNUQVeoClggfPnQ,2033 -deepspeed/ops/csrc/aio/common/deepspeed_aio_types.h,sha256=pnFoIKThq3P_bHoo-jxZgL86SMbXlJZGfpghK7pQqBg,1402 -deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.cpp,sha256=UihEk2VXEqQLCiT0u2ovXoQULBMGbZgSfNBwIx-csSQ,4674 -deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.h,sha256=iCd_aS9H_DqJae7-O5qS9GASyE0gNYLLwwrJT-td_GY,2176 -deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp,sha256=G23cmsy1S5u6_vOynCYJnO3Q1B4iZkvwpMEzcprcqQ8,1215 -deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.h,sha256=1n67p-zqla1BFNUPQ54Zrpq9OuWaKc7DuAUPKPDjScU,1065 -deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.cpp,sha256=pUuqwGkCkFLDtXtU5mDbgSTW7BXoFlAzFsTyIcx88GI,1394 -deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.h,sha256=fAX8as116Ecs7NObKX_LdjALHck0godLITqVkg-Byt4,885 -deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.cpp,sha256=Wye-QMpyW8AxLqqouStbdUDRzikV20ajizyskDygd0g,3433 -deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.h,sha256=lXYUW0DsMKdighnXSdKSMfmM4oZxpXCMJ71-Swtmf0c,1038 -deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.cpp,sha256=MpYmRFffPgfkcjI_4tjw-nnEFimy1wgUxFF-HCNl5K4,1850 -deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.h,sha256=oIJ0bbxcgE3XFwEMxSOx15XXk1rNbQFUNcYdHqFkQdU,861 -deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.cpp,sha256=aZysoGV-1JU2wfAKQ3qYpyTGaAFPC8NDJGLQuucwRA0,4232 -deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.h,sha256=z3O_AoPJXMHKjkwsrfrXIcIONXUNc2Yr3Pz1w3Q-99U,965 -deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp,sha256=YvU1K6pKxcMG9knW0BRkbUh3jr-_N4ySp_QhDjIyudY,903 -deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.h,sha256=WgLYd05nnvNWO492PGBPNoIP8B0X2JI3axKFFSGLntQ,624 -deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.cpp,sha256=xIqqhPWz-bqBD89HV9TItBnacCWr-JalPjqSaUodwrg,4414 -deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.h,sha256=LxCbVzeTiKRmU2P_cbRZFHIInlW6HQWszpB5s_p2l0M,1169 -deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.cpp,sha256=Frut6FvfUuJDMIPioBIJzXuuTR54Xjj6nPdBXNAclOo,12834 -deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.h,sha256=1bDbRqqOuAOHlIO91k3MytydEQM4YLArmaCakBu5tig,3975 -deepspeed/ops/csrc/aio/py_lib/py_ds_aio.cpp,sha256=KGjMfQDWaKmViwB6cmJ7uZCJBK5HANTOkufHehFwz3I,4772 -deepspeed/ops/csrc/aio/py_test/single_process_config.json,sha256=6dSFspQHi0nNbLClaIfQ_3Jzj4Nw7urGCly8ZEWnfIo,433 -deepspeed/ops/csrc/compile/deepcompile.cpp,sha256=lar-g-tVzTCXzkrUXy3bevvZ4H65Se3hIfiQF5MUeHU,5868 -deepspeed/ops/csrc/compile/init.cpp,sha256=YfdOvY9YZGyWlo1JohwrsJe5RZuF6EibZekMyphzsGI,4461 -deepspeed/ops/csrc/compile/util.cpp,sha256=v9lmBKZDKPWRWOyi7Azve5PYGa5QO89uLwBHqlwzCCI,3769 -deepspeed/ops/csrc/compile/z1.cpp,sha256=rUEojV-gw-92ygEdS3hDBK52vE1cOfEU9nfz9jqXpMA,4391 -deepspeed/ops/csrc/compile/z1.h,sha256=gn4YNfbL21uG7McpLUhbmml8R7BRLJL9FTpUiwUyZoQ,459 -deepspeed/ops/csrc/compile/z2.cpp,sha256=ACKjRQS_TzisIyohi2HXVbRi_oL5fEsWd8boNO3b7fk,3973 -deepspeed/ops/csrc/compile/z2.h,sha256=HsOGpkbegv3bhZ65vKreTMAD8G2ZS1WJ0ZT-Vx5z_GM,247 -deepspeed/ops/csrc/compile/z3.cpp,sha256=pVo-ktWIbEYygvphK7E8CpBP3H6sV-_rUgy2ftrgcTM,20140 -deepspeed/ops/csrc/compile/z3.h,sha256=eRBbP0wfO15pLQtk4abWFNluBOHNGmOXEGiPQaj1WbY,2347 -deepspeed/ops/csrc/cpu/adam/fused_adam.cpp,sha256=m45BQnPCTkrYHiYcdKNa791zqVqtbdwia4TFoXjnOK0,1433 -deepspeed/ops/csrc/cpu/comm/ccl.cpp,sha256=1U2qIchXfW7ypHD1DIG7mjPXR2levhB55Jz00Kv5Hnk,12342 -deepspeed/ops/csrc/cpu/comm/shm.cpp,sha256=bcPXkJacINgDy081QimwLTBnXXxeUexV3zb_oxStsmk,27403 -deepspeed/ops/csrc/cpu/comm/shm.h,sha256=iM952l2aMWnNP7uI44blK4b8sS80Am756zWpjhd7K3o,401 -deepspeed/ops/csrc/cpu/comm/shm_interface.cpp,sha256=C19jMwHp4HhTinkHeFMdV7NtdMrq2JwUsXIcQK8CwnI,6195 -deepspeed/ops/csrc/cpu/lion/fused_lion.cpp,sha256=-_QUZhwjkjsePW-YyZ4ishpUPCGmbgfDp4C2jNgRIA0,1230 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention.cpp,sha256=1d5msQb6IiUdDeAYpzsiKjFHEFPHNAvBXDkO-aY2W_g,2095 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_back.cu,sha256=10fKkFtrljPwiHjIAhoZJrBEcIRjdBP8JDagnFdWyOU,9971 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_cu.cu,sha256=lfsGocib3dneilmyTKnEL9fdmALP8v_UhvwApmwgpVM,7509 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_grad_bias.h,sha256=ziZZar1z41wsr6ia5HtX8aAjpAuomtq5cCgZMzYylWI,12185 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_pipelined.h,sha256=RvYmxAW4NN7kMc9pLfEmkXRO-nXMKR1BX3fhhRUG3SI,24898 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_rescale_output.h,sha256=WWnsDQeR-XHcjKeEZixxhaTkv89tpyc1jaXoSAEiWt8,9606 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_thread_apply_logsumexp.h,sha256=Q_QP5LK3AwDtt6GHxLdUg-ttVqtDkxN2gv4bXiDnlPw,6392 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma.h,sha256=rXA6J4g_TuQ1Ufm2c0Q2dMcuC1FV6X4MpkYpk7TeYp0,6270 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_base.h,sha256=zn3uylR7vb9GNLKG3UjodjXyPO_XpIuK0bRBwILm4zo,6511 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_multistage.h,sha256=WeVgJF-IIsVjd50TwDsL_oKwr2UMYUWdLEdhOiy_bfE,30281 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_pipelined.h,sha256=7MyvAD-eXwtNkfyjtd459IvbAbhW81Ys1RzFN8TasTM,15770 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/find_default_mma.h,sha256=l0wa_nFB9Vrz11bR_0lH-D4vaK-sWfuHm31Mg_Em3bc,9420 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_accum_lambda_iterator.h,sha256=G6LeFpLrmBPVVsnfiYCPjh0h4VdJ5su6jyYPLvPymqI,16280 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_from_smem.h,sha256=o_6L3rJlRtGmughlt9g0xLCzvTqBbUObl72lNxkbhNY,84216 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm_kernel_utils.h,sha256=G02ENUZevpYkNI8FwYfxdiaY8BmFsZkIZh3ieGGWpzg,11661 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/epilogue_predicated_tile_iterator.h,sha256=Vl83Qgukp3lN0EwAkZdftc1gxIfvwvFzjRw7M8wSQ9A,27084 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/make_residual_last.h,sha256=bl7iRKlZYnjqYwoST7SRHWwqygjsvYCUmLMBfbAxV7Q,4245 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_access_iterator_residual_last.h,sha256=Jg4qjZ0ukBy2ZIZ-2wxZ9u4mrIcdVAcil56e7CdCY8M,71866 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_atomic.h,sha256=RQ1sW5cWwiPveRxIs8rVuOgaXBBhsgHU5mbiLh6rawc,31326 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_residual_last.h,sha256=aZy_jvL9zakb9oeFXkoihOabhLLhZjRSbLkf_6lUwS8,71626 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/transpose_warp_iterator.h,sha256=PtMp-E5HhdBmwgRiEbE2ZjlVL-b97euA80waCNO7jYQ,2536 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/warp_iterator_from_smem.h,sha256=QvDXwcGimdRqSoaGSzN5aS3q37VSXeihlU_A6ZPee9M,10125 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_backward.h,sha256=lLlud-SKkjuOa9UcJkOUfSZRV3RQpPRf9J6qMtAODBQ,95042 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_forward.h,sha256=CTv6YV9cCkmzU7w49AqRhENGWOccHYz0d4YzmRhvVBM,47657 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/bias_broadcast.h,sha256=7yLAnvgViEcY5qls5eEJyfQDNkToG0TfPKQhm8Qi0-E,5952 -deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/tile_smem_loader.h,sha256=Dr_uSpu3dH8GAfChWggDeon9rn7F4G8BFrjm5_Xyt0E,4476 -deepspeed/ops/csrc/fp_quantizer/fp_quantize.cpp,sha256=74zRV9Yj3aoyA6rYhQmZUu0tVa4TqC7g3uR8mbytux0,5789 -deepspeed/ops/csrc/fp_quantizer/fp_quantize.cu,sha256=j17WEtr-J7xnSmOd4yt04OqKZawqPK-P7Aqdfk_omIM,25629 -deepspeed/ops/csrc/fp_quantizer/includes/context.h,sha256=n-WCCbWIdbfNBCeS8c1wdY8LNdGIkz8qwn8C9ImX_nM,1427 -deepspeed/ops/csrc/fp_quantizer/includes/fp_quantize.h,sha256=UTTOIZa18Cz0wOip7Tfn6itcCeUWySA0DJLDRZLOUY8,5979 -deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.cpp,sha256=bmF57wG0eiTrKQ1WVR3eavTPfVPBSnxnb1i9tcUKQns,5815 -deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.h,sha256=pD5329Gd2a75fzfThYDYHIb1Fs6TAeT0Br-YPjunHGg,1131 -deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_utils.h,sha256=4abARh9qrz9TE9GkOUJOxZGzavRhQCubOFRuj-G3uLs,4790 -deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.cpp,sha256=X337agbrlRfSF-jnSzCLv4z_CejhaUyjyZowl8EoC4Q,4309 -deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.h,sha256=ZookPemckgcq6mTCuStAFxq47xmaRLnKGKeV7ca3BBA,1657 -deepspeed/ops/csrc/gds/py_lib/py_ds_gds.cpp,sha256=oJb2OJ-RoTtlqHc04Yz6DCg_rG9XZsayURECIbXaqwM,5106 -deepspeed/ops/csrc/includes/StopWatch.h,sha256=DocXNohnOPWiZBqfrdI6AKV9IggYRB1Bylzxfke_i5s,1981 -deepspeed/ops/csrc/includes/Timer.h,sha256=4mYxcATDJElSy5GCj_xFDZ2AvHHmMrhJlW0jJbMP7nQ,1180 -deepspeed/ops/csrc/includes/activation_type.h,sha256=WkFfkL5mgLyYAZimLUouThCkyA2aiF9LQlh1q206L9o,264 -deepspeed/ops/csrc/includes/compat.h,sha256=5ivuTerIqiCaYu2HpRPzJcUtqMsAtRB7aOXel69HeZQ,336 -deepspeed/ops/csrc/includes/context.h,sha256=Cs6hRXBiFSjpKKQOAMCX9F0uP_Hh1IeZLloNJxHxk1E,7207 -deepspeed/ops/csrc/includes/conversion_utils.h,sha256=augq8Zf-t46MaxWh0ryFb46_Guxud0k4kIK-h93rI8Y,12379 -deepspeed/ops/csrc/includes/cpu_adagrad.h,sha256=zcXXyFLZyjNSFfiXgaoYS8ytfFtvKyKQk8hP-SrmLuY,4036 -deepspeed/ops/csrc/includes/cpu_adam.h,sha256=Xkulb32K2B35yGaYqKKeBWEOBrhGb2160UqO-D2C0ok,7365 -deepspeed/ops/csrc/includes/cpu_lion.h,sha256=mFmdjrrWqll4mKKsNGeJysPomGFKBD1e5F2zosWSAq8,5516 -deepspeed/ops/csrc/includes/cublas_wrappers.h,sha256=Gyu-fwvPndgvI1PqGjb-CaoAkeGW8YsRRBBpIjKUjLg,3840 -deepspeed/ops/csrc/includes/custom_cuda_layers.h,sha256=lqmn3AfbfSynF2BLu3YWcZZnwQxsO89LiOluP-J2efU,12964 -deepspeed/ops/csrc/includes/deepcompile.h,sha256=qR0lZVVJCLzBxU_am8OxEOQIg8s4QB7yyRMragxyFLE,21211 -deepspeed/ops/csrc/includes/dequantization_utils.h,sha256=uFpPmBIYd3gNkKXImQLidYnGDYYTXkpY9V4o-AkdYpk,7301 -deepspeed/ops/csrc/includes/dropout.h,sha256=l5JuSzyfxRsfufBKIebPwpghnmYR_E6ucCmGbLvZ1qc,2195 -deepspeed/ops/csrc/includes/ds_kernel_utils.h,sha256=OKOMT46XItRnkrzqKHHK9n_ewxHh6AN-Uhv4aKixBm4,1297 -deepspeed/ops/csrc/includes/ds_transformer_cuda.h,sha256=OwVCwN4tS1kMEeF9J6QV9GG8AvtYklaHRJrMPmji-zY,6164 -deepspeed/ops/csrc/includes/feed_forward.h,sha256=8IbZYyUWVFGMhkjve5ZdWGQc8WvYXVhNEAddT2Vvp6s,3560 -deepspeed/ops/csrc/includes/gelu.h,sha256=iCsuAdNFzhLeccvc5D7KnsAcOcvXN_QD4jlaFRPF340,1018 -deepspeed/ops/csrc/includes/gemm_test.h,sha256=X9oIVKaZpFOGU0yCw-kUz4m5XDq3V1leW5v5dMGQYnA,11546 -deepspeed/ops/csrc/includes/general_kernels.h,sha256=0ksHr78VwjPiAfYPAloHM5uMxriRfNPEZyk9OklLhhI,1507 -deepspeed/ops/csrc/includes/memory_access_utils.h,sha256=Lhpfu-OXn3TscCXK6VdWr8SEp-9rFrtnThrQGDU6aQI,34901 -deepspeed/ops/csrc/includes/normalize_layer.h,sha256=j2lJkrxENFILqdUpEZIMj1qGrFty5SNoAeUrifThvSI,7089 -deepspeed/ops/csrc/includes/quantization.h,sha256=vqhmpPNw0Um-TqEtR_KyO3Pu6eZV5lGpXbraNLIfEoQ,5889 -deepspeed/ops/csrc/includes/quantization_utils.h,sha256=ciuykcFNUW9BxCYTE_moCtN5gJYFd6dgibCleBX7NdM,17366 -deepspeed/ops/csrc/includes/quantizer.h,sha256=YIRNwwr1eZYMmXbs6AZT82shpXuWbn94P4OVMcA3tss,346 -deepspeed/ops/csrc/includes/reduction_utils.h,sha256=AeQx6X3WD4vExcIQRO9LuHdCAWa5vLBcf2QmcIDgePM,23380 -deepspeed/ops/csrc/includes/simd.h,sha256=Ti4NkJ6ciwgPqEuAYmIr0zGRNPRq0z8gK8YmsRSZDKE,10190 -deepspeed/ops/csrc/includes/softmax.h,sha256=UGZHziAfbGN6OdTIA24G5CWHhV8hAfQbq16NN2_0Wlo,1642 -deepspeed/ops/csrc/includes/strided_batch_gemm.h,sha256=vp3ZGYjoSj_xZFTvdIj6PMrn_mupI3YlA4rLFjhYE8I,7279 -deepspeed/ops/csrc/includes/type_shim.h,sha256=8VuX9zK5EsmleO7wuAq9D1eHZC7zKWaf0N97eaGg7x8,6388 -deepspeed/ops/csrc/lamb/fused_lamb_cuda.cpp,sha256=icBfzI6QpN8LqIZT4-AY91UXlg_zDqkptNZy10erbws,3995 -deepspeed/ops/csrc/lamb/fused_lamb_cuda_kernel.cu,sha256=vjhQw1pj6v19tVE2otOCSStk53Ss0R82g9o1lXfJ5vw,15291 -deepspeed/ops/csrc/lion/cpu_lion.cpp,sha256=didiEmqBezyZ5C2yv92MVbPne8WQbDgkGN5JYV-e1ww,409 -deepspeed/ops/csrc/lion/cpu_lion_impl.cpp,sha256=xCsrkojmjl4uk-QSNa5TpONYRbAQi_V68yacdwvL5vs,7402 -deepspeed/ops/csrc/lion/fused_lion_frontend.cpp,sha256=RMiEOozd9gjXPidc4Ykg-kDIJek0OOASFVwmLIfFs0M,732 -deepspeed/ops/csrc/lion/multi_tensor_apply.cuh,sha256=tMXW3UwsOhNM4J1axQbNkn4gNfmig7_vXPiBLdFCCv4,5595 -deepspeed/ops/csrc/lion/multi_tensor_lion.cu,sha256=11RptKP9XmwMuvGBrmq3P1T8PAH9IvBJirB8gInMVoQ,4377 -deepspeed/ops/csrc/quantization/dequantize.cu,sha256=iibytiJox1QcqE8b57KCLPOaD610STv2VUmKoC7RVM8,3302 -deepspeed/ops/csrc/quantization/fake_quantizer.cu,sha256=CPMZnjOjJTKFibCBeYG1WrXnwEG-0l6POuZCRsBrHBo,37581 -deepspeed/ops/csrc/quantization/pt_binding.cpp,sha256=rcfjK9PldJDktjppS0_Rc9vFh_Rd_YSauGoyj6KzzAY,17434 -deepspeed/ops/csrc/quantization/quant_reduce.cu,sha256=mQ4rRfiPRDfgGJHgt5uFlgsBMav94aZqDYeui9JLBLs,25376 -deepspeed/ops/csrc/quantization/quantize.cu,sha256=IzmI5sq3V8B3tF92i84wlouV186djjVvdISREcepYRc,6988 -deepspeed/ops/csrc/quantization/quantize_intX.cu,sha256=SIEWB2qHtFkHINoTTNukiJnXnQS06BKVH3Wjtg-im1k,8514 -deepspeed/ops/csrc/quantization/swizzled_quantize.cu,sha256=dbNOMXtGQwBI7GnGZXKfVvR9h0iRC-aDU6jSeKrzoMo,19787 -deepspeed/ops/csrc/random_ltd/gather_scatter.cu,sha256=YFiLiYlgkq1RNYS8yRNI3CKYUrv_8ht0lsvZsUcsUYg,8408 -deepspeed/ops/csrc/random_ltd/pt_binding.cpp,sha256=7NAanMKS-j_hBputc6Ltv1yH6-c7xLwhDZJUfLiH--M,9590 -deepspeed/ops/csrc/random_ltd/slice_attn_masks.cu,sha256=KqANgsPTAAHyys0cPh3023JaBF8b0kOKF-W13W39Q_4,5076 -deepspeed/ops/csrc/random_ltd/token_sort.cu,sha256=Tm3oqRHM6UF1PWMuVaghnV0NftUjv3SUEbtBejNLJNc,7044 -deepspeed/ops/csrc/sparse_attention/utils.cpp,sha256=HmNSJfE5WzsE3dYSWjOKOAxawUOcgZ2JwQpJSjvB3b4,4523 -deepspeed/ops/csrc/spatial/csrc/opt_bias_add.cu,sha256=8N-FAKh-6Fwu7byz6e_sjXpd8XCRwdFntbZzZL6H8x8,6273 -deepspeed/ops/csrc/spatial/csrc/pt_binding.cpp,sha256=K2zwA7EXdtjM05P9Bl8YV_p35TQXrt4xrADEkq44Dlo,3863 -deepspeed/ops/csrc/spatial/includes/spatial_cuda_layers.h,sha256=dPh5gJbsCh2gB4_g9PSzMb1zN6LJvAGLfD7Vw_08_mc,915 -deepspeed/ops/csrc/transformer/cublas_wrappers.cu,sha256=HntTKLKThdx3Gj_jBK6gfLD1pnbVE_-wt5Yd_JTQatk,21052 -deepspeed/ops/csrc/transformer/dropout_kernels.cu,sha256=u6U7RuMioRS_aixeRL0aGvU3r7exZ-EFeE7Etgg66IM,29835 -deepspeed/ops/csrc/transformer/ds_transformer_cuda.cpp,sha256=sztj1XsEX1rbeq2lovNbYf_jZ-ExLL9qB-TLVRP_1D4,47582 -deepspeed/ops/csrc/transformer/gelu_kernels.cu,sha256=f1emIwHWorXkuIqaB9kD93L4HjzSCcQ2JL7GlE6hh6s,12191 -deepspeed/ops/csrc/transformer/general_kernels.cu,sha256=TzoOK2TqstG0MX9swk0UptK2zBgkpYwo99PeY9j0Fns,14520 -deepspeed/ops/csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu,sha256=7z2HGEwCY1N-JutgeAYGsJJtXyvi36D8NdlpNaqRjcE,8594 -deepspeed/ops/csrc/transformer/inference/csrc/dequantize.cu,sha256=z7GzRO9XUJFb0DNPr7luWML9cJTYOZrR5wAQGeuStqY,4916 -deepspeed/ops/csrc/transformer/inference/csrc/gelu.cu,sha256=peyOpVxxboodFWAY0YB6vr3kakCcr8x0EJ4j7HiEY4U,28877 -deepspeed/ops/csrc/transformer/inference/csrc/layer_norm.cu,sha256=79gLSVFW19gsgFQNezyITbBImha7_oXJo5NfBloOhzU,20883 -deepspeed/ops/csrc/transformer/inference/csrc/pointwise_ops.cu,sha256=tJP7V172fEHn86xtO1H8lKzBS7CToDNofJBnOVhivEA,2476 -deepspeed/ops/csrc/transformer/inference/csrc/pt_binding.cpp,sha256=aIdgiEyTGJm_InXaO_eXQ4zH_2MMiRAx9u7vXEEnGkU,89443 -deepspeed/ops/csrc/transformer/inference/csrc/relu.cu,sha256=F_W6D8eeHtm_btn-FTyAV8FnLTDJ9H8mpkDBWpmxxDc,2318 -deepspeed/ops/csrc/transformer/inference/csrc/rms_norm.cu,sha256=0vHKVYfQQ8w-qLUD-ihfCd0jWDQ_N03QuQu0rimHrdU,10251 -deepspeed/ops/csrc/transformer/inference/csrc/softmax.cu,sha256=GY3hRUkY59BcQqe2BA0QCGw3tlN2UO4GQipRezSzJLM,27245 -deepspeed/ops/csrc/transformer/inference/csrc/transform.cu,sha256=6S1NPVL8Fw9y6wJ0rho2NVIXm64YxthNMKxSWTDsMPo,31887 -deepspeed/ops/csrc/transformer/inference/includes/inference_context.h,sha256=2-KQOUFuWdx1rOLR3rZtTmOpz9LUmY9Vs7a5qNybUzI,10531 -deepspeed/ops/csrc/transformer/inference/includes/inference_cublas_wrappers.h,sha256=AXxx5E4YDDpbKkAxozTHsNJQpIN7dSWr8CiNnpJnBLE,21657 -deepspeed/ops/csrc/transformer/inference/includes/inference_cuda_layers.h,sha256=TlXSYeyoKSWBaxuN_dDxrTMhutjQocPjKeGL_NPXIPo,9057 -deepspeed/ops/csrc/transformer/normalize_kernels.cu,sha256=u9fnXA77eqpoC1hri_w0IzC0xW7_OIG5Cs6F_UFnivs,74900 -deepspeed/ops/csrc/transformer/softmax_kernels.cu,sha256=GKmqpxpbAFlhM2p1SFMZ02phEpYMSGUFA2J9ukLovno,26758 -deepspeed/ops/csrc/transformer/transform_kernels.cu,sha256=Ob-Hae1q-KZmDxaji0Ci0nHyNuRYDvPFBF_yrnB8ZoE,22709 -deepspeed/ops/csrc/utils/py_ds_utils.cpp,sha256=jB0kJ0caZitZaw81e6LJ989yaCtCHP_tEFFXvXEfMLk,682 -deepspeed/ops/csrc/utils/tensor_cast.cpp,sha256=GZMlX8HcDDPI0ahjL4_Pf-xjurK1VlvI1dp_LUyhEXQ,797 -deepspeed/ops/csrc/utils/tensor_cast.h,sha256=OBCrxtt5tY0wmP6gd7SS8qlPVcOuYEba-PvC5MTsBSA,383 -deepspeed/ops/csrc/xpu/adagrad/cpu_adagrad.cpp,sha256=-LiTifCTt4Pr_4miMyP5-8W84FHhORWZn8mx-VZDR58,6814 -deepspeed/ops/csrc/xpu/adam/fused_adam_frontend.cpp,sha256=tu6pBcAvqcdh1ti6SB7fj8HnqFDQg3Rqbaa4xsQ2v7M,880 -deepspeed/ops/csrc/xpu/adam/multi_tensor_adam.dp.cpp,sha256=Du_BWrHEoYqoCfDb6qKxDeKXP69ZfdFEW1mt7ir5ah0,6595 -deepspeed/ops/csrc/xpu/adam/multi_tensor_apply.dp.hpp,sha256=__3PLu5FDG8vFF9fOwP45KL_yWTXYxKiNdLUYNyqnNI,9482 -deepspeed/ops/csrc/xpu/common/custom_cuda_kernel.dp.cpp,sha256=fSy_UoyFgDlpKWEDRtUqwrhaLMMUV1OG7eCVZfzcGKs,3525 -deepspeed/ops/csrc/xpu/includes/compat.h,sha256=5ivuTerIqiCaYu2HpRPzJcUtqMsAtRB7aOXel69HeZQ,336 -deepspeed/ops/csrc/xpu/includes/cpu_adagrad.h,sha256=7yZNV7T38Yvnv2JLfcdm4KdBuXyPdvT-7Plzva_5BbQ,3951 -deepspeed/ops/csrc/xpu/includes/cpu_adam.h,sha256=3a0twZXzUUnAhrEymYF8gKFoVW594B-I-Snjg6paK9Q,7956 -deepspeed/ops/csrc/xpu/includes/simd.h,sha256=tRwYTW1Xm20rDzsew6c7dfEPbCBd0tsjY26mx1VnVYk,6588 -deepspeed/ops/csrc/xpu/includes/type_shim.h,sha256=qA0DdBShoh_Wute7v52S18jvHB81xi2ADhYvn8EXRVg,7834 -deepspeed/ops/csrc/xpu/packbits/packing.cpp,sha256=oUzgPsI6tXEK56YWrSxAvf5ei289d36We-vmqTXKPuU,3292 -deepspeed/ops/deepspeed4science/__init__.py,sha256=LRmYuJYPcnqiyIuSCwsvqpD61wcScSmS7An8zq0UY4M,175 -deepspeed/ops/deepspeed4science/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/deepspeed4science/__pycache__/evoformer_attn.cpython-310.pyc,, -deepspeed/ops/deepspeed4science/evoformer_attn.py,sha256=50z8I1w4lJww3bRiBxq8PfEJxcacZw_qXk4DyTyWsSQ,4121 -deepspeed/ops/fp_quantizer/__init__.py,sha256=UxAI8a6mhBD9nD5dIfV6dAcO4qGlz8tMLBSEgVucflA,174 -deepspeed/ops/fp_quantizer/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm.cpython-310.pyc,, -deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm_triton.cpython-310.pyc,, -deepspeed/ops/fp_quantizer/__pycache__/quantize.cpython-310.pyc,, -deepspeed/ops/fp_quantizer/fp8_gemm.py,sha256=ZFcVM33qjxFg5utlPDtZ9GYMThcdFMgeqMAblhs8Fcs,939 -deepspeed/ops/fp_quantizer/fp8_gemm_triton.py,sha256=XDdslcyiSZUoOhPKKm7Hdzy60gqqkOdODqehj-iCYJY,7106 -deepspeed/ops/fp_quantizer/quantize.py,sha256=Cw19CwiaIk8kC_1dyjZSbMcYqVHEbdXNYybkVXN1xog,6481 -deepspeed/ops/gds/__init__.py,sha256=j3yUaB9pVgZKU8_d5i6-ypeWGZbpzKovc-r3ZpF5XLQ,132 -deepspeed/ops/gds/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/lamb/__init__.py,sha256=HG9WeYIi4tTtKZBuLuJNIVpLSK75djGxjXDkrgfreEk,130 -deepspeed/ops/lamb/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/lamb/__pycache__/fused_lamb.cpython-310.pyc,, -deepspeed/ops/lamb/fused_lamb.py,sha256=gZbjO38b517Ii1tEJwT0O81YwJ5bXtflM2uq57NFXtI,7815 -deepspeed/ops/lion/__init__.py,sha256=EFF2vXYA9hYPNuqBsLsd1MTFCPiOFfW1ZWBrlbxfbD8,169 -deepspeed/ops/lion/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/lion/__pycache__/cpu_lion.cpython-310.pyc,, -deepspeed/ops/lion/__pycache__/fused_lion.cpython-310.pyc,, -deepspeed/ops/lion/__pycache__/multi_tensor_apply.cpython-310.pyc,, -deepspeed/ops/lion/cpu_lion.py,sha256=mJkqnY3EG7YEDi71Fg2qD0HFiSFFFkkz84Vcx8xfG2Y,5306 -deepspeed/ops/lion/fused_lion.py,sha256=QoYzriRhYkaa_bDxIY_YGviCUe-BcNZvvNtFA4g5kLc,5548 -deepspeed/ops/lion/multi_tensor_apply.py,sha256=APt3UCnfw-nLjJOGWXC3izWRXNu-16TYhB-s213DpM0,429 -deepspeed/ops/op_builder/__init__.py,sha256=HJ00U1SrhiXDftTIq9k5hZp-CuRPZ5GWOyQerqgC21U,2005 -deepspeed/ops/op_builder/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/all_ops.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/async_io.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/builder.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/cpu_adagrad.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/cpu_adam.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/cpu_lion.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/dc.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/evoformer_attn.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/fp_quantizer.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/fused_adam.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/fused_lamb.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/fused_lion.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/gds.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/inference_core_ops.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/inference_cutlass_builder.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/quantizer.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/ragged_ops.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/ragged_utils.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/random_ltd.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/sparse_attn.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/spatial_inference.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/stochastic_transformer.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/transformer.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/transformer_inference.cpython-310.pyc,, -deepspeed/ops/op_builder/__pycache__/utils.cpython-310.pyc,, -deepspeed/ops/op_builder/all_ops.py,sha256=bymFEiEsBwsM_BlMOePCM-9xMGeQdk6u0ZXh2fPQo9w,1223 -deepspeed/ops/op_builder/async_io.py,sha256=EMeLJZ9GBIAjuD2AwDDsfW0GpKPCg-YrK12yqze_Y9w,4500 -deepspeed/ops/op_builder/builder.py,sha256=r3sljupH_zvXGHApweR1FuCGDEQ7hdnuePttpC6UL9E,35323 -deepspeed/ops/op_builder/cpu/__init__.py,sha256=fRbecWZV5CWo2W5bRzfweUx0C_8gjVwmj4Yb6YgIbzY,353 -deepspeed/ops/op_builder/cpu/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/op_builder/cpu/__pycache__/async_io.cpython-310.pyc,, -deepspeed/ops/op_builder/cpu/__pycache__/builder.cpython-310.pyc,, -deepspeed/ops/op_builder/cpu/__pycache__/comm.cpython-310.pyc,, -deepspeed/ops/op_builder/cpu/__pycache__/cpu_adam.cpython-310.pyc,, -deepspeed/ops/op_builder/cpu/__pycache__/fused_adam.cpython-310.pyc,, -deepspeed/ops/op_builder/cpu/__pycache__/no_impl.cpython-310.pyc,, -deepspeed/ops/op_builder/cpu/async_io.py,sha256=aY_ndccYdKQYQ7adRJfl3NfwY8Dj6yqn1OkBr42qn9g,3550 -deepspeed/ops/op_builder/cpu/builder.py,sha256=Bi1gqcUHoDb12ci-5Urb2QRopaQFAy5E5uNELEd4Hzg,1460 -deepspeed/ops/op_builder/cpu/comm.py,sha256=e_bVQ0AnLnZSpsuY4du3NgcpYjhpRoQOVqJu1NQmqW4,2114 -deepspeed/ops/op_builder/cpu/cpu_adam.py,sha256=rGtsjwhzzP9iMDIphGxjh5TLUBtUiEwtrnf5x5BbG08,625 -deepspeed/ops/op_builder/cpu/fused_adam.py,sha256=5LAQSEL2H4Bcc1b7LArqKCcF2Ztuq2uaABmRmI7m5Ss,546 -deepspeed/ops/op_builder/cpu/no_impl.py,sha256=OIZ3aVWXK9W8iRtgNL0pAtl033WV0hEoOuPP35UuFUc,616 -deepspeed/ops/op_builder/cpu_adagrad.py,sha256=fB81OcyxFAMVzINEi1uY1CneERg5KKS8IafOM_ICzCo,622 -deepspeed/ops/op_builder/cpu_adam.py,sha256=aBvkhLFpMxELFvFEpB57C2Fmlq_-3O3azRNEI8XglGo,635 -deepspeed/ops/op_builder/cpu_lion.py,sha256=8l8o0abx4TCn30ps_Iy40HCnymQ9ve225kp_uLRwpWk,635 -deepspeed/ops/op_builder/dc.py,sha256=VL5DLUUq2NbCrJh3Y2WrIcqZmgA1WKj9AVvGfVFDY1U,1157 -deepspeed/ops/op_builder/evoformer_attn.py,sha256=dmGaZQv-G74ckTri8yAMRwjdwomqIjjDeiu2zyCPNgU,3602 -deepspeed/ops/op_builder/fp_quantizer.py,sha256=WUFftK3MJApYc2IZDF7Cm0fBD2PmYCYprbdA23dKSIU,4130 -deepspeed/ops/op_builder/fused_adam.py,sha256=IC87eCjvRPJ2ETVdQUv0-d_yxMXhDJvbPML1AEU9cSQ,1044 -deepspeed/ops/op_builder/fused_lamb.py,sha256=Ojx9euSMUeeU7DVIFxE11w_fHERqYsZxNNtaGxnAAag,1216 -deepspeed/ops/op_builder/fused_lion.py,sha256=1ZJh92_chzH0Jupt8RO8GUOnrm6ACKLECZ5mTiV7WI0,1044 -deepspeed/ops/op_builder/gds.py,sha256=zJPklZuTb6tbi8rnswfmBAC862pwVKv-E6gcv-0QT0M,2195 -deepspeed/ops/op_builder/hpu/__init__.py,sha256=E7-aI0c0S1kZjNCKsjp7XImSfdziWHt3psvwHc82YHQ,370 -deepspeed/ops/op_builder/hpu/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/op_builder/hpu/__pycache__/builder.cpython-310.pyc,, -deepspeed/ops/op_builder/hpu/__pycache__/cpu_adam.cpython-310.pyc,, -deepspeed/ops/op_builder/hpu/__pycache__/fp_quantizer.cpython-310.pyc,, -deepspeed/ops/op_builder/hpu/__pycache__/fused_adam.cpython-310.pyc,, -deepspeed/ops/op_builder/hpu/__pycache__/no_impl.cpython-310.pyc,, -deepspeed/ops/op_builder/hpu/__pycache__/transformer_inference.cpython-310.pyc,, -deepspeed/ops/op_builder/hpu/builder.py,sha256=nVgIhaRZRgppzIeuYrqfvWcwpGShcK7JUDxHgosA5Cc,1391 -deepspeed/ops/op_builder/hpu/cpu_adam.py,sha256=FzHM2A3XH3dXSYy1DWA7H61swa7m4S9OLsLrbvL8Myk,681 -deepspeed/ops/op_builder/hpu/fp_quantizer.py,sha256=0278EpesoRE54UveUX1ZdlcVlqdmHkzhPk2bFg1bH-U,2947 -deepspeed/ops/op_builder/hpu/fused_adam.py,sha256=ZpP8fkXQ7KdEDtrW5XT9qQ5QlsiGQhONo5r2Mpn03SE,3498 -deepspeed/ops/op_builder/hpu/no_impl.py,sha256=cUJ87-4-I46CK3GmS-Ve52nzfMhcFayT_YZF5FWEmzY,616 -deepspeed/ops/op_builder/hpu/transformer_inference.py,sha256=dPHPLX_wdRFIbCbx4FbxeREn1AJoWOzRbO36dKmcYig,1341 -deepspeed/ops/op_builder/inference_core_ops.py,sha256=ZwfV2HtgoSUv055qWQ8xGSyjI3EUW7TiPy9Xd5DbgQ8,3835 -deepspeed/ops/op_builder/inference_cutlass_builder.py,sha256=OMvFKGBukkUTiYcNHu5CcjZ_oWNA7ifTVGlDu_wBWd0,3373 -deepspeed/ops/op_builder/mlu/__init__.py,sha256=oPAhp7_PS9EV3UCyKDFRfFifj7mJPV2YEMHlqTbVbD8,402 -deepspeed/ops/op_builder/mlu/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/op_builder/mlu/__pycache__/builder.cpython-310.pyc,, -deepspeed/ops/op_builder/mlu/__pycache__/cpu_adagrad.cpython-310.pyc,, -deepspeed/ops/op_builder/mlu/__pycache__/cpu_adam.cpython-310.pyc,, -deepspeed/ops/op_builder/mlu/__pycache__/fused_adam.cpython-310.pyc,, -deepspeed/ops/op_builder/mlu/__pycache__/no_impl.cpython-310.pyc,, -deepspeed/ops/op_builder/mlu/builder.py,sha256=QdxnbQZxPI98OtPEvDl186lgOat6nQIoIJsIlxYC1aw,1283 -deepspeed/ops/op_builder/mlu/cpu_adagrad.py,sha256=FCrdOCfiYHdfs1HSwhYgohE2x27jmOcFpE9TtqZRa48,565 -deepspeed/ops/op_builder/mlu/cpu_adam.py,sha256=i-Oj1EP5jeuU_7JHu7iqrs16BONW93S7RIq_ImUfyZQ,669 -deepspeed/ops/op_builder/mlu/fused_adam.py,sha256=o574SNUxqSYNXVzGAsRkXdoXpINpwQ6-rysW6SRDVMo,1137 -deepspeed/ops/op_builder/mlu/no_impl.py,sha256=j0r-HhbzJuCVzna9vLWVaDp2wE21AuHDHBN8yy7vUQU,799 -deepspeed/ops/op_builder/npu/__init__.py,sha256=utWarSvvEuL_LUk8-i4Kx2oCHJRemddoGeGfzHOCsTE,419 -deepspeed/ops/op_builder/npu/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/op_builder/npu/__pycache__/async_io.cpython-310.pyc,, -deepspeed/ops/op_builder/npu/__pycache__/builder.cpython-310.pyc,, -deepspeed/ops/op_builder/npu/__pycache__/cpu_adagrad.cpython-310.pyc,, -deepspeed/ops/op_builder/npu/__pycache__/cpu_adam.cpython-310.pyc,, -deepspeed/ops/op_builder/npu/__pycache__/cpu_lion.cpython-310.pyc,, -deepspeed/ops/op_builder/npu/__pycache__/fused_adam.cpython-310.pyc,, -deepspeed/ops/op_builder/npu/__pycache__/inference.cpython-310.pyc,, -deepspeed/ops/op_builder/npu/__pycache__/no_impl.cpython-310.pyc,, -deepspeed/ops/op_builder/npu/async_io.py,sha256=D6N-B3SF9bpg8y8hyHxljAdSWIcpvjhXQ3J80DJnACQ,3770 -deepspeed/ops/op_builder/npu/builder.py,sha256=k21-MqhnSCYyMpuOxbRjDWbsIcnDnA6GxWPtsmtBuC8,3200 -deepspeed/ops/op_builder/npu/cpu_adagrad.py,sha256=CM4wGBKhOHayOyzvOGX2yDIh4PgOCu6AUHXF4GKq9kE,581 -deepspeed/ops/op_builder/npu/cpu_adam.py,sha256=S5V76jfADBzNNP-CyBoTXZSWrQMLYWuos6ACygoumDE,594 -deepspeed/ops/op_builder/npu/cpu_lion.py,sha256=rDkiA8vVCMolVFWUnvN9PkrLOJlc6P_rSRNYmOUrYV8,594 -deepspeed/ops/op_builder/npu/fused_adam.py,sha256=uKjA8NKcDUqjr34ToapiXrQHD661uClj6N9equFMh_g,2142 -deepspeed/ops/op_builder/npu/inference.py,sha256=-IsUpVW106uskkegCrYzyMoed5W-RhhffawBWsXYad8,15187 -deepspeed/ops/op_builder/npu/no_impl.py,sha256=qZnbySBzMLLUrMlZV656Pdh8TGfEZKNACKdmD6uT9ho,755 -deepspeed/ops/op_builder/quantizer.py,sha256=k0kTBN9_9Ri_dyzhW_ENxYpUVhiFOeHaIdVY1W50ZPE,1037 -deepspeed/ops/op_builder/ragged_ops.py,sha256=dpipyaz2IADHNLHYC9vavFlXLEMUBJhVwWgwxagKjbI,4912 -deepspeed/ops/op_builder/ragged_utils.py,sha256=CVQGORZRAxjymG9HCuhRtnKudI2sMXd3v3MRVQhi1fQ,2727 -deepspeed/ops/op_builder/random_ltd.py,sha256=zNBXnNjckT9dZqtj4pVYHcSUu4m2govven5qh5Qp9VU,879 -deepspeed/ops/op_builder/sdaa/__init__.py,sha256=N_LPrOGu5tIJVcRRXzDm86SBtEUInPAIJLWMJWqMr1s,1811 -deepspeed/ops/op_builder/sdaa/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/op_builder/sdaa/__pycache__/builder.cpython-310.pyc,, -deepspeed/ops/op_builder/sdaa/__pycache__/cpu_adam.cpython-310.pyc,, -deepspeed/ops/op_builder/sdaa/__pycache__/fused_adam.cpython-310.pyc,, -deepspeed/ops/op_builder/sdaa/__pycache__/no_impl.cpython-310.pyc,, -deepspeed/ops/op_builder/sdaa/builder.py,sha256=59jhxPvqJLB-5qePUz0qXUryvSwskkFKKwaRPqRD0U4,2788 -deepspeed/ops/op_builder/sdaa/cpu_adam.py,sha256=DgbrhjECQzOkSyyfDlbrm1l1zxDKC9z3hJ0CrUVeXtQ,2175 -deepspeed/ops/op_builder/sdaa/fused_adam.py,sha256=4u_xou5pjkpTtiPMYhvmX5XuLS-RSQNwRAg9OCNvTlI,2647 -deepspeed/ops/op_builder/sdaa/no_impl.py,sha256=spQASQULrUtgzUL4Cr5TNlItdW_9TWCJwIvGIikLOUQ,2306 -deepspeed/ops/op_builder/sparse_attn.py,sha256=nled88Ns_XIbbbqEITlFF_ZgQujuuv3JKyZIEqV9GNc,3221 -deepspeed/ops/op_builder/spatial_inference.py,sha256=0c8zejdt2H05Q879i-kVCuUxnrsJo7ni0F532EGK5Gg,1599 -deepspeed/ops/op_builder/stochastic_transformer.py,sha256=nQPCaLrtuHVtJxMWIQp_8IT1iQ2XS9nQ1BJVwwAeIPI,565 -deepspeed/ops/op_builder/transformer.py,sha256=wyrj-FZTKPyXzwKOGPYDyic3w5hdtXh4dWJS_wfRdCU,1094 -deepspeed/ops/op_builder/transformer_inference.py,sha256=9R9XaLdkPxf89oZgSfZ9yTcJBpN7B4ic27vBDKTBHbY,2868 -deepspeed/ops/op_builder/utils.py,sha256=w2kP0Vg2PMV0RGbmy0erVVCYK4FN6V7A3YWhpSgJODo,455 -deepspeed/ops/op_builder/xpu/__init__.py,sha256=3WbwNA_gET5S9B3BiSglGVJB6325AKYgdsMXWYU1-7Y,421 -deepspeed/ops/op_builder/xpu/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/op_builder/xpu/__pycache__/async_io.cpython-310.pyc,, -deepspeed/ops/op_builder/xpu/__pycache__/builder.cpython-310.pyc,, -deepspeed/ops/op_builder/xpu/__pycache__/cpu_adagrad.cpython-310.pyc,, -deepspeed/ops/op_builder/xpu/__pycache__/cpu_adam.cpython-310.pyc,, -deepspeed/ops/op_builder/xpu/__pycache__/flash_attn.cpython-310.pyc,, -deepspeed/ops/op_builder/xpu/__pycache__/fused_adam.cpython-310.pyc,, -deepspeed/ops/op_builder/xpu/__pycache__/inference.cpython-310.pyc,, -deepspeed/ops/op_builder/xpu/__pycache__/no_impl.cpython-310.pyc,, -deepspeed/ops/op_builder/xpu/__pycache__/packbits.cpython-310.pyc,, -deepspeed/ops/op_builder/xpu/async_io.py,sha256=gFV9LPwhIHNGcSk9bisKe_Z3bCXCxsPyIBFaL-sEtc0,3738 -deepspeed/ops/op_builder/xpu/builder.py,sha256=Pul5G1bF8D5fb-RpXISLJhVoZyE5vH3GNeOA8qKW_NM,5816 -deepspeed/ops/op_builder/xpu/cpu_adagrad.py,sha256=bi9U6d3UmHRYw7qkIcRNRYTuasPYoSe38ngF71ZM1RQ,576 -deepspeed/ops/op_builder/xpu/cpu_adam.py,sha256=1mUDRDd_SV_ZFxCzbt39PajtYqEjRdzHAQPzCClx7FY,627 -deepspeed/ops/op_builder/xpu/flash_attn.py,sha256=YUkFzxNQ_c8MsriNLC1pdqjk-8op60DiDIaky8u52Jk,1615 -deepspeed/ops/op_builder/xpu/fused_adam.py,sha256=yEvt078k2gqqqDjPG_D5I5Ntuz_aAqEjYgtCgjpgAjg,701 -deepspeed/ops/op_builder/xpu/inference.py,sha256=_4-VhP2iIhY9qA_tYadGlKWWtDQ95NVPooVGcyiqzIc,1134 -deepspeed/ops/op_builder/xpu/no_impl.py,sha256=ODPPWw4qSX5ITgjXYT3LHWc-JSe92eBzAIQVKpgkZRA,752 -deepspeed/ops/op_builder/xpu/packbits.py,sha256=eDSmRqT3cFETG48aeHvSHtYg5oaenw2X7WQFW4WJcM0,626 -deepspeed/ops/quantizer/__init__.py,sha256=5IdLoKmcCi6MuZNNTGADQs88dgcCUdN50WofVXDcvpI,132 -deepspeed/ops/quantizer/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/quantizer/__pycache__/quantizer.cpython-310.pyc,, -deepspeed/ops/quantizer/quantizer.py,sha256=z3dh0MRUklHnR74y662k9H31FOMU0buQOGE9uPLcUZA,1193 -deepspeed/ops/random_ltd/__init__.py,sha256=MC02456CIFtrgIYKClskYz4kDnQ9X7zRyyKNp110l58,191 -deepspeed/ops/random_ltd/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/random_ltd/__pycache__/dropping_utils.cpython-310.pyc,, -deepspeed/ops/random_ltd/dropping_utils.py,sha256=tHMRvk6WIVDoPR58HWc0OwUNzXhfn8VYd7Xw13X-r5o,4902 -deepspeed/ops/sparse_attention/__init__.py,sha256=qhwanCYxLO9eaH1cx52E4LNe_An8Nrja_8WyyoYk42I,467 -deepspeed/ops/sparse_attention/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/sparse_attention/__pycache__/bert_sparse_self_attention.cpython-310.pyc,, -deepspeed/ops/sparse_attention/__pycache__/matmul.cpython-310.pyc,, -deepspeed/ops/sparse_attention/__pycache__/softmax.cpython-310.pyc,, -deepspeed/ops/sparse_attention/__pycache__/sparse_attention_utils.cpython-310.pyc,, -deepspeed/ops/sparse_attention/__pycache__/sparse_self_attention.cpython-310.pyc,, -deepspeed/ops/sparse_attention/__pycache__/sparsity_config.cpython-310.pyc,, -deepspeed/ops/sparse_attention/bert_sparse_self_attention.py,sha256=kQM4sl9L-jNirgCWwA9U-4EWGo-AKSzQE7d2CPcCwgw,3465 -deepspeed/ops/sparse_attention/matmul.py,sha256=RYTMwGqaYRUTR7DC60Ek7PNWUc5n9sKJNgcTpzmiZfo,32948 -deepspeed/ops/sparse_attention/softmax.py,sha256=B_q4n2TP29iPfMLTnkRqgHxJ8y5cBjsTixPOO7zrrMw,11322 -deepspeed/ops/sparse_attention/sparse_attention_utils.py,sha256=XlNgnI3CdG8Kur-Ac1O5qJVQRX_lgexTzXtfarMAHb4,12300 -deepspeed/ops/sparse_attention/sparse_self_attention.py,sha256=Mgs_4NaoBoHIghomqaqnKtjbAgYs1xWlZOlZAoSaMao,6746 -deepspeed/ops/sparse_attention/sparsity_config.py,sha256=eUCYxt5QPHWXdot-Myu206H9Y5WV8HrqxvO7LPKHHKg,42463 -deepspeed/ops/sparse_attention/trsrc/__init__.py,sha256=2G0yT0H4-bH6beifS-DoGK8lQcfZ2aIqt7IFekVRxP0,1032 -deepspeed/ops/sparse_attention/trsrc/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/sparse_attention/trsrc/matmul.tr,sha256=nM6iQGzBXyKFzHqmCDjpr8xmykwEls4Te6PWVJg22Dg,6628 -deepspeed/ops/sparse_attention/trsrc/softmax_bwd.tr,sha256=C--cFHczPdwJA7wm0s66p9f_nkzDMatqu3-iBQoAlog,1923 -deepspeed/ops/sparse_attention/trsrc/softmax_fwd.tr,sha256=t4bvel9w4ilPsn24DQgzIqGNG08Pt4pfCBFua8rj9fc,4047 -deepspeed/ops/transformer/__init__.py,sha256=ADlxWPujgq5jT_XL8yuRtTg9jdQiBK-5ydzgnNwap-o,413 -deepspeed/ops/transformer/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/transformer/__pycache__/transformer.cpython-310.pyc,, -deepspeed/ops/transformer/inference/__init__.py,sha256=ENKx17yjSD2dIqD6kPq-zxY-Xae-O84-ZPGGNeFq064,315 -deepspeed/ops/transformer/inference/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/transformer/inference/__pycache__/bias_add.cpython-310.pyc,, -deepspeed/ops/transformer/inference/__pycache__/config.cpython-310.pyc,, -deepspeed/ops/transformer/inference/__pycache__/diffusers_2d_transformer.cpython-310.pyc,, -deepspeed/ops/transformer/inference/__pycache__/diffusers_attention.cpython-310.pyc,, -deepspeed/ops/transformer/inference/__pycache__/diffusers_transformer_block.cpython-310.pyc,, -deepspeed/ops/transformer/inference/__pycache__/ds_attention.cpython-310.pyc,, -deepspeed/ops/transformer/inference/__pycache__/ds_mlp.cpython-310.pyc,, -deepspeed/ops/transformer/inference/__pycache__/moe_inference.cpython-310.pyc,, -deepspeed/ops/transformer/inference/__pycache__/triton_ops.cpython-310.pyc,, -deepspeed/ops/transformer/inference/bias_add.py,sha256=x1gk_iN4pmFNsSt6jwtBXil-3no4yDhTMbi_6yrs7gY,876 -deepspeed/ops/transformer/inference/config.py,sha256=g9IEbiLEc1kZUvJ3pKvYNQEvwljcIthhupUXpSX914c,6120 -deepspeed/ops/transformer/inference/diffusers_2d_transformer.py,sha256=JVSE9B60n6XHeg37zE26W3-jzKOhxiOtJkmkCiM1nDU,236 -deepspeed/ops/transformer/inference/diffusers_attention.py,sha256=BJa7R6fYdgeqB1gBNX38hD5CqLQdcj8R0g5m4sAp8_Q,9600 -deepspeed/ops/transformer/inference/diffusers_transformer_block.py,sha256=1WHyfxnq1-JMtCZavXr0mK0lPM3vm6QNCRAqVlxwzyc,4324 -deepspeed/ops/transformer/inference/ds_attention.py,sha256=Oiy0oW5ByvWArjqzqU3qUG3eWQmhYQRf_JBwQOyJPu4,15207 -deepspeed/ops/transformer/inference/ds_mlp.py,sha256=Lr5q6pk3MrAx8D36wXhTNDlV7MsFMkFFaHb7CiKzfzo,6294 -deepspeed/ops/transformer/inference/moe_inference.py,sha256=p3Ck2gBRtnY-km8EYnFfjhuyc63rN-ZUVdBBVs97s4A,17360 -deepspeed/ops/transformer/inference/op_binding/__init__.py,sha256=JhdbkMuKVwM62BEU4aI4OscG3GGkxFKsvnttFx9vc9k,382 -deepspeed/ops/transformer/inference/op_binding/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/base.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_add.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_gelu.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_relu.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_residual.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/einsum_sec_sm_ecm.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/gated_activation.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/gelu_gemm.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/layer_norm.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/linear.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/mlp_gemm.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/moe_res_matmul.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/pad_transform.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/pre_rms_norm.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/qkv_gemm.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/residual_add.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/rms_norm.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax_context.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_add.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_matmul.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/__pycache__/workspace.cpython-310.pyc,, -deepspeed/ops/transformer/inference/op_binding/base.py,sha256=YN1s4f-7BKDmvCyKkVW3DOaXz40Xxk8FrIeDU2BfmJA,536 -deepspeed/ops/transformer/inference/op_binding/bias_add.py,sha256=9zrXRrRgwfd21zT_jqYF0tJUFur3Dup3x4-WwD_4d_E,980 -deepspeed/ops/transformer/inference/op_binding/bias_gelu.py,sha256=dOuLD_SwIw0p0V45qQnosr82bCDBn1cXtMMzGUXzQ5g,1214 -deepspeed/ops/transformer/inference/op_binding/bias_relu.py,sha256=Z87V8pcsHv01FxuXk2qw-qBsEtG6CNJ0CjJ8MSBUQL0,1161 -deepspeed/ops/transformer/inference/op_binding/bias_residual.py,sha256=ZURDKvdWx8fHPr_2pe9V0d66C_8hbcPLAxgTqKeTClc,952 -deepspeed/ops/transformer/inference/op_binding/einsum_sec_sm_ecm.py,sha256=P5OTQgxyxc73DLQvyEQzzYsH7NQde25awdmSMV9HhK4,938 -deepspeed/ops/transformer/inference/op_binding/gated_activation.py,sha256=xLe2XuWm0-myxEkamkjwliNAode8-GiiCMtcpSB7iLg,1549 -deepspeed/ops/transformer/inference/op_binding/gelu_gemm.py,sha256=LZ3wUeuGrBG56n0jMFbv7-7X_TbzUnJQ6FyQpIGDHNc,2191 -deepspeed/ops/transformer/inference/op_binding/layer_norm.py,sha256=Fj2rFo3RV8Bqg1XpeX1hYVCvWiP2QReN_b9TUSypB60,2370 -deepspeed/ops/transformer/inference/op_binding/linear.py,sha256=vIYuAmj9CX0KERm9_IE7pdmf4XG4UPXkanBqqtlnhgc,2719 -deepspeed/ops/transformer/inference/op_binding/mlp_gemm.py,sha256=ZTwwz6ptQ2hRQw1Mp1_0bheTc8X7Wc8m7e5Clw7yKZ8,5745 -deepspeed/ops/transformer/inference/op_binding/moe_res_matmul.py,sha256=2w4xeLtx3u7MdD4jZQ7f1DPeWPIvbOaHVMSICVa9bDw,1002 -deepspeed/ops/transformer/inference/op_binding/pad_transform.py,sha256=dlxmmFGmY0xnUUCONADu2lH4mMnHH4KWvm1odvEnTAc,899 -deepspeed/ops/transformer/inference/op_binding/pre_rms_norm.py,sha256=L9j-2d55UW8cqUyq4g1bVxVHT7doW4ACcFdhJXjhs64,1009 -deepspeed/ops/transformer/inference/op_binding/qkv_gemm.py,sha256=vDrBmcbq1xnFcoQ7J6AGQS9X_2NHEaveC-6kPW3E1y0,4514 -deepspeed/ops/transformer/inference/op_binding/residual_add.py,sha256=1w8a4TEIRPtZKPg6z40PC-_dj0ssNDpVA-8JlfPgBFo,3084 -deepspeed/ops/transformer/inference/op_binding/rms_norm.py,sha256=S68R8QOSlT3qyzdmubRx0q6YkMC_qggDH1qPl_Hxmho,985 -deepspeed/ops/transformer/inference/op_binding/softmax.py,sha256=LE1zRTF5JvIN3YmAexhtmTiWdkk5zwdl8A8ZTkmfp1M,3105 -deepspeed/ops/transformer/inference/op_binding/softmax_context.py,sha256=A0L_4CG72sLDUK2n0dWdEVATqt5loB3UpqM66jHNVco,7527 -deepspeed/ops/transformer/inference/op_binding/vector_add.py,sha256=n_l4idWGoSjrFze9_9j_V2YnaDwnVJYzOoSCERnoVAI,882 -deepspeed/ops/transformer/inference/op_binding/vector_matmul.py,sha256=LhLI-FzBDCM9ivWOENA_NkvGAna-eqxESQBNVp7zu44,2640 -deepspeed/ops/transformer/inference/op_binding/workspace.py,sha256=-Of25xBArRF5u0zYYv2rK8SVrjWaQSDla0ztbaqv6Sw,9584 -deepspeed/ops/transformer/inference/triton/__init__.py,sha256=FVRc5cfjotwsFtcB7gkAm09lXFbNl01gaoet67x4B_k,350 -deepspeed/ops/transformer/inference/triton/__pycache__/__init__.cpython-310.pyc,, -deepspeed/ops/transformer/inference/triton/__pycache__/attention.cpython-310.pyc,, -deepspeed/ops/transformer/inference/triton/__pycache__/gelu.cpython-310.pyc,, -deepspeed/ops/transformer/inference/triton/__pycache__/layer_norm.cpython-310.pyc,, -deepspeed/ops/transformer/inference/triton/__pycache__/matmul_ext.cpython-310.pyc,, -deepspeed/ops/transformer/inference/triton/__pycache__/mlp.cpython-310.pyc,, -deepspeed/ops/transformer/inference/triton/__pycache__/ops.cpython-310.pyc,, -deepspeed/ops/transformer/inference/triton/__pycache__/residual_add.cpython-310.pyc,, -deepspeed/ops/transformer/inference/triton/__pycache__/softmax.cpython-310.pyc,, -deepspeed/ops/transformer/inference/triton/__pycache__/triton_matmul_kernel.cpython-310.pyc,, -deepspeed/ops/transformer/inference/triton/attention.py,sha256=iLSFq5MOkMXhUc9uQvpmoGo5fzYSIQPJttGNAE6BG5E,16576 -deepspeed/ops/transformer/inference/triton/gelu.py,sha256=OAViD-qU6B50KhjQgg-Y3leuEKP6gW2OkX7MX_VGDko,1152 -deepspeed/ops/transformer/inference/triton/layer_norm.py,sha256=q5Xt2ov4z5DonYRBnWevaUYTFwJU5qIrJc6RVDtLaPw,7512 -deepspeed/ops/transformer/inference/triton/matmul_ext.py,sha256=6020-OGlpSZv8ytcKDnctmwggKCQ9XdLcOjjgkK3R-0,16028 -deepspeed/ops/transformer/inference/triton/mlp.py,sha256=c4yRQv3GPzEKcaxMFwkms4q2eF47ewemkYcnuhu0iec,4225 -deepspeed/ops/transformer/inference/triton/ops.py,sha256=5YLnOEbqn-GydhgR26Qw2z5zZRNDDxQQ_O3oNQyng5w,3804 -deepspeed/ops/transformer/inference/triton/residual_add.py,sha256=PYcAlDC8yz00LzJ8eXsJKoMmxjD_nMM0X5PbK8w_J70,3118 -deepspeed/ops/transformer/inference/triton/softmax.py,sha256=2q_qNmffFM4l2UgsAFKrJmAVrID95KK1rcg5dEzBmco,3208 -deepspeed/ops/transformer/inference/triton/triton_matmul_kernel.py,sha256=8pa5U6ZJphYxJ3oX2cPdxdYilmuKf8w_lvt1A8-pjbg,12570 -deepspeed/ops/transformer/inference/triton_ops.py,sha256=lR82VswoMZ9b2Axbck_eF5GCOczUmvKTDAT1W8bVH1w,5487 -deepspeed/ops/transformer/transformer.py,sha256=1UMs9JNjlK3mD5PJi3iER9QARUO2_a5h0sZAY0Rmo24,20600 -deepspeed/pipe/__init__.py,sha256=ddCnO6IbTpGzlAIgfEbGg6RPCz70HNN-MwTjCpTsqLI,164 -deepspeed/pipe/__pycache__/__init__.cpython-310.pyc,, -deepspeed/profiling/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/profiling/__pycache__/__init__.cpython-310.pyc,, -deepspeed/profiling/__pycache__/config.cpython-310.pyc,, -deepspeed/profiling/__pycache__/constants.cpython-310.pyc,, -deepspeed/profiling/config.py,sha256=OzJuTtxVvm62T1iA5eXocvP8xQQXHNtzcCdqISLTTKM,1959 -deepspeed/profiling/constants.py,sha256=QaZFzozkotv71aArhiwtQJ13DNDGDU1xC0Jm0KrVH08,1243 -deepspeed/profiling/flops_profiler/__init__.py,sha256=NCZv_Ktz4sFmmfpBwGKVVbbKrI36VQHwPapvLlWUUxE,120 -deepspeed/profiling/flops_profiler/__pycache__/__init__.cpython-310.pyc,, -deepspeed/profiling/flops_profiler/__pycache__/profiler.cpython-310.pyc,, -deepspeed/profiling/flops_profiler/profiler.py,sha256=NtUZADA8U0-9PioVuNKZnOIFPnQ_-dSlxWhdguJwaXw,51818 -deepspeed/runtime/__init__.py,sha256=4I9UpQ5vMRU5SYSF_dW9FJDEnBq4m_0SuwtVQ92lGaA,95 -deepspeed/runtime/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/__pycache__/base_optimizer.cpython-310.pyc,, -deepspeed/runtime/__pycache__/bf16_optimizer.cpython-310.pyc,, -deepspeed/runtime/__pycache__/compiler.cpython-310.pyc,, -deepspeed/runtime/__pycache__/config.cpython-310.pyc,, -deepspeed/runtime/__pycache__/config_utils.cpython-310.pyc,, -deepspeed/runtime/__pycache__/constants.cpython-310.pyc,, -deepspeed/runtime/__pycache__/dataloader.cpython-310.pyc,, -deepspeed/runtime/__pycache__/eigenvalue.cpython-310.pyc,, -deepspeed/runtime/__pycache__/engine.cpython-310.pyc,, -deepspeed/runtime/__pycache__/hybrid_engine.cpython-310.pyc,, -deepspeed/runtime/__pycache__/lr_schedules.cpython-310.pyc,, -deepspeed/runtime/__pycache__/precision_config.cpython-310.pyc,, -deepspeed/runtime/__pycache__/progressive_layer_drop.cpython-310.pyc,, -deepspeed/runtime/__pycache__/quantize.cpython-310.pyc,, -deepspeed/runtime/__pycache__/sparse_tensor.cpython-310.pyc,, -deepspeed/runtime/__pycache__/state_dict_factory.cpython-310.pyc,, -deepspeed/runtime/__pycache__/torch_autocast.cpython-310.pyc,, -deepspeed/runtime/__pycache__/utils.cpython-310.pyc,, -deepspeed/runtime/__pycache__/weight_quantizer.cpython-310.pyc,, -deepspeed/runtime/activation_checkpointing/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/runtime/activation_checkpointing/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/activation_checkpointing/__pycache__/checkpointing.cpython-310.pyc,, -deepspeed/runtime/activation_checkpointing/__pycache__/config.cpython-310.pyc,, -deepspeed/runtime/activation_checkpointing/checkpointing.py,sha256=AfqMBlOkP0DYiNJyqej_9zbzQdaP3aomJ0scwxn6CZs,44313 -deepspeed/runtime/activation_checkpointing/config.py,sha256=YQkhHYT6Qr2GJl_KusLU0pUJ-rbLLKKj0Zo7LkeZkAg,3988 -deepspeed/runtime/base_optimizer.py,sha256=qqzmVfz_Xdah_QogN3Nz6hiT8NUJacyGjX1uh7TTzQ0,3475 -deepspeed/runtime/bf16_optimizer.py,sha256=hvKo9HanSTB9kClFJEiBQDBNFVa5mS89pl0z2M_ZkxE,25572 -deepspeed/runtime/checkpoint_engine/__init__.py,sha256=LaqB9nJrtf9qqyhJ0Eyi_mDvYXOkGVK8-XIuxWWBVuw,420 -deepspeed/runtime/checkpoint_engine/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/checkpoint_engine/__pycache__/checkpoint_engine.cpython-310.pyc,, -deepspeed/runtime/checkpoint_engine/__pycache__/decoupled_checkpoint_engine.cpython-310.pyc,, -deepspeed/runtime/checkpoint_engine/__pycache__/fast_checkpoint_engine.cpython-310.pyc,, -deepspeed/runtime/checkpoint_engine/__pycache__/nebula_checkpoint_engine.cpython-310.pyc,, -deepspeed/runtime/checkpoint_engine/__pycache__/torch_checkpoint_engine.cpython-310.pyc,, -deepspeed/runtime/checkpoint_engine/__pycache__/utils.cpython-310.pyc,, -deepspeed/runtime/checkpoint_engine/checkpoint_engine.py,sha256=82YbbHCTWv898UU075AemsYniypgmHLGUVLq6C-HwPU,1257 -deepspeed/runtime/checkpoint_engine/decoupled_checkpoint_engine.py,sha256=1dnj6UgtfkaOqEq3NpwwwOE_3pnv9mmRuFSKTSoC1NY,5808 -deepspeed/runtime/checkpoint_engine/fast_checkpoint_engine.py,sha256=Cr-9Px8TGIqVtrErYArLFZ84ueqKIGIXDh4KGO4njTg,1761 -deepspeed/runtime/checkpoint_engine/nebula_checkpoint_engine.py,sha256=5FDE7UUwM_xbB6StNjtynSyEfLpljyx0FwmWi4u9sME,5126 -deepspeed/runtime/checkpoint_engine/torch_checkpoint_engine.py,sha256=UDWWw1hnn9oYWS0lRa_JP9Mp80Yic0iJ2g6Y7t5SZtg,1764 -deepspeed/runtime/checkpoint_engine/utils.py,sha256=ABeLEQGo9amYK0h_jkTytjENf2ybNkyBKphmyVdQyEI,1820 -deepspeed/runtime/comm/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/runtime/comm/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/comm/__pycache__/coalesced_collectives.cpython-310.pyc,, -deepspeed/runtime/comm/__pycache__/compressed.cpython-310.pyc,, -deepspeed/runtime/comm/__pycache__/hccl.cpython-310.pyc,, -deepspeed/runtime/comm/__pycache__/mpi.cpython-310.pyc,, -deepspeed/runtime/comm/__pycache__/nccl.cpython-310.pyc,, -deepspeed/runtime/comm/coalesced_collectives.py,sha256=2tdtNA68uONSuqZull-4T_p2eao2wChCLsRUIxM0BR8,11545 -deepspeed/runtime/comm/compressed.py,sha256=yflJJlYN2Wb47h6PUOnUJfmRIS8TJdjMi3LwVAbhvGA,5661 -deepspeed/runtime/comm/hccl.py,sha256=O1XqCH-DdBbJMlTxRkR2l5OUpMQYWXLhkgyZlPwbdCk,5176 -deepspeed/runtime/comm/mpi.py,sha256=m3-XVeEL6F5ln5krGaExKcBj1Qlxg7y_XwvvVo8APcU,10076 -deepspeed/runtime/comm/nccl.py,sha256=kAvrk95McM3accR11f2Grg1spe1kwP8UmC000zemsJw,7584 -deepspeed/runtime/compiler.py,sha256=tCN_Upy5tkLFRtT6Qa0u8-EbK0RLGP6Oaej0YyM8qJ4,712 -deepspeed/runtime/compression/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/runtime/compression/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/compression/__pycache__/cupy.cpython-310.pyc,, -deepspeed/runtime/compression/cupy.py,sha256=gCQ7N_EE7TRgZmRaGRHuhGAXrtEkvfbAnJTBMGO-pXI,691 -deepspeed/runtime/config.py,sha256=tuJDGKzsXBASgi1tFUImqeGxsZ2ip3nDie_rYgLB-Zo,40536 -deepspeed/runtime/config_utils.py,sha256=Mjlms4PInPnKZct1xW-ceuulZCvhsspm1SwnuqTs_SE,8460 -deepspeed/runtime/constants.py,sha256=eizwrjKZpfT0qFgAcYBZNF4MGPexhMYgM20o7QZgkFc,15571 -deepspeed/runtime/data_pipeline/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/runtime/data_pipeline/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/__pycache__/config.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/__pycache__/constants.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/__pycache__/curriculum_scheduler.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/config.py,sha256=9GGQNeHTR-bog6aBpcxCm65jc5QO2Y3KWNMZoPfK4gs,7721 -deepspeed/runtime/data_pipeline/constants.py,sha256=x4rApIpDlX1zPViP01wMiKqGNPSiqzhzNZ4VlQ7zFp4,5702 -deepspeed/runtime/data_pipeline/curriculum_scheduler.py,sha256=g6Heo7I5L8mpyKJXSH8t1SwO2Lxk9jMfOIMbYDksV00,10025 -deepspeed/runtime/data_pipeline/data_routing/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/runtime/data_pipeline/data_routing/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/data_routing/__pycache__/basic_layer.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/data_routing/__pycache__/helper.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/data_routing/__pycache__/scheduler.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/data_routing/__pycache__/utils.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/data_routing/basic_layer.py,sha256=QvJuYJRbFTZSkKovN5Z7PVQP5o_fNC-2QgYqrWGy8ns,5638 -deepspeed/runtime/data_pipeline/data_routing/helper.py,sha256=mFWBiepSdAAejDPSaIV6sXeKpRlfXoWFGF0CUmiDACY,1282 -deepspeed/runtime/data_pipeline/data_routing/scheduler.py,sha256=8ddlbZJ3RO0btPKmurh3288Lk1CIoMxRWHcrxlVMZkY,4638 -deepspeed/runtime/data_pipeline/data_routing/utils.py,sha256=ZGrHeImPXbVhLwOMQGWowXEK5YpIGXfC8i3RHVF4NAI,955 -deepspeed/runtime/data_pipeline/data_sampling/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/runtime/data_pipeline/data_sampling/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_analyzer.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_sampler.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/data_sampling/__pycache__/indexed_dataset.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/data_sampling/__pycache__/utils.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/data_sampling/__pycache__/variable_batch_size_and_lr.cpython-310.pyc,, -deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py,sha256=eN9-LeQE66qratDmZ4vyZL1jPs2Fc-M4UM0HXG4Yg-Q,48993 -deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py,sha256=PJNw94Zc0L31DRk4DH9vkXXGvXPMqtisnbF79gFY7UY,19894 -deepspeed/runtime/data_pipeline/data_sampling/indexed_dataset.py,sha256=QQdA40qmTehZqhBDX2oGeJI8m1r8c6Jc5Hx-ygPa-j8,21449 -deepspeed/runtime/data_pipeline/data_sampling/utils.py,sha256=5QtKBV_54fKL9I7ydqF40uBaLkUpcX5LFeYVz2KY_r4,1690 -deepspeed/runtime/data_pipeline/data_sampling/variable_batch_size_and_lr.py,sha256=o-fRgO_JiJ32nQh0rrh-Zz_uGfqsfqHzPsHTkl_KpLE,24311 -deepspeed/runtime/dataloader.py,sha256=rNiiaHzYOU3hp8OAiC2g9tFWPwACJTMhJ_EVrxvP5OU,6977 -deepspeed/runtime/domino/__init__.py,sha256=4I9UpQ5vMRU5SYSF_dW9FJDEnBq4m_0SuwtVQ92lGaA,95 -deepspeed/runtime/domino/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/domino/__pycache__/async_linear.cpython-310.pyc,, -deepspeed/runtime/domino/__pycache__/transformer.cpython-310.pyc,, -deepspeed/runtime/domino/async_linear.py,sha256=4kDO5n0bLu71BuksABfK4bbi5Opp95JfeYJrOqRJuWQ,4576 -deepspeed/runtime/domino/transformer.py,sha256=xU4nTT-xy52S33rM2tuWlAjG6GxRpNEPQ2lf4MVT0Lg,25278 -deepspeed/runtime/eigenvalue.py,sha256=5SEkMHwnGPqRFyU7ymuoO6cOkRcDDcJnjsnp2wXa2JM,5872 -deepspeed/runtime/engine.py,sha256=tu30PwPILQwNeaqWJWzS83WtaH9SS7rG3jG1WfkKIg4,193515 -deepspeed/runtime/fp16/__init__.py,sha256=Dzme9x1YQb-Ru4gzsCNsv8r19zkGyqwPwPwPgZzief4,140 -deepspeed/runtime/fp16/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/fp16/__pycache__/fused_optimizer.cpython-310.pyc,, -deepspeed/runtime/fp16/__pycache__/loss_scaler.cpython-310.pyc,, -deepspeed/runtime/fp16/__pycache__/unfused_optimizer.cpython-310.pyc,, -deepspeed/runtime/fp16/fused_optimizer.py,sha256=w1ke8oaojnHEz-w8kVr_dy_3d72x1lnFt_z4UY18he4,22322 -deepspeed/runtime/fp16/loss_scaler.py,sha256=7E3yzMJEFzbKcIR2aeUx9-9GrB6_2sGv5wOX8NZ6ezY,11652 -deepspeed/runtime/fp16/onebit/__init__.py,sha256=8hnDOpk71O_EekeWgWIJ4CnbtYt67dre4flu5hoES-4,186 -deepspeed/runtime/fp16/onebit/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/fp16/onebit/__pycache__/adam.cpython-310.pyc,, -deepspeed/runtime/fp16/onebit/__pycache__/lamb.cpython-310.pyc,, -deepspeed/runtime/fp16/onebit/__pycache__/zoadam.cpython-310.pyc,, -deepspeed/runtime/fp16/onebit/adam.py,sha256=jRCUSTzBhvo1xPD5XvRxmHni8R2JXMBGL1_P7udBmpc,15694 -deepspeed/runtime/fp16/onebit/lamb.py,sha256=jMZBvuwG8GxoF4tD22ECq9jZyl8_XBcAKnDLyS0GDHA,23556 -deepspeed/runtime/fp16/onebit/zoadam.py,sha256=z3Pc2u9LyJXm5hCfQjAzj_wp72RJP8In6LYv5KyrChA,19554 -deepspeed/runtime/fp16/unfused_optimizer.py,sha256=qB2H34wcs1q05TqhyrA5wPYmC-RJduqjxswGA1rs3Ng,18029 -deepspeed/runtime/hybrid_engine.py,sha256=rfP2MpmnB-DqCiSRjwKyklcl78RtGNwFOQa-PiutlWU,20965 -deepspeed/runtime/lr_schedules.py,sha256=eDap8FWaxsp6M-VSxi07NWLKWCwEwyqbOLIQb6frQxU,39074 -deepspeed/runtime/model_checkpointing/__init__.py,sha256=9EfmW2osMT_aBTR0xOE5NadoNHh7mETjTJxf5lnz-4k,173 -deepspeed/runtime/model_checkpointing/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/model_checkpointing/__pycache__/config.cpython-310.pyc,, -deepspeed/runtime/model_checkpointing/__pycache__/constants.cpython-310.pyc,, -deepspeed/runtime/model_checkpointing/__pycache__/data_parallel_writer_factory.cpython-310.pyc,, -deepspeed/runtime/model_checkpointing/__pycache__/utils.cpython-310.pyc,, -deepspeed/runtime/model_checkpointing/__pycache__/writer_factory.cpython-310.pyc,, -deepspeed/runtime/model_checkpointing/config.py,sha256=jFcd7SQFwXj7JqWgpJHXinDN5j9X2-73U4aWd4BBFaU,2992 -deepspeed/runtime/model_checkpointing/constants.py,sha256=JJ2jAMNPTnWuFwr6JJelgayVBryCcIQhpXlxXBga0zg,2258 -deepspeed/runtime/model_checkpointing/data_parallel_writer_factory.py,sha256=hsW-v_h4VfZM8gbn57EBK-Ov4noINYpnKAku8GD9x6Y,10732 -deepspeed/runtime/model_checkpointing/utils.py,sha256=C78G1D098fb3ffbstnGkhcO8CEw2pUbGq9NLtxMSsl8,3275 -deepspeed/runtime/model_checkpointing/writer_factory.py,sha256=SaFlZX0cRvXrIVKTg_qSmOOd7NK3mYvpQx4NYdpPQz4,4930 -deepspeed/runtime/pipe/__init__.py,sha256=4Xc534VEOCSKYVc-ImMFLx5eTmPB_URVgWqweehFsVI,195 -deepspeed/runtime/pipe/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/pipe/__pycache__/engine.cpython-310.pyc,, -deepspeed/runtime/pipe/__pycache__/module.cpython-310.pyc,, -deepspeed/runtime/pipe/__pycache__/p2p.cpython-310.pyc,, -deepspeed/runtime/pipe/__pycache__/schedule.cpython-310.pyc,, -deepspeed/runtime/pipe/__pycache__/topology.cpython-310.pyc,, -deepspeed/runtime/pipe/engine.py,sha256=Rg89Sfun9GELR8Ob4OrClFyEmi8fMn6cOrOn6-WQ8eQ,63690 -deepspeed/runtime/pipe/module.py,sha256=FfLY5D8DZVGy0kEtcTs7Pqmc4U3V7Lq6TdPwZyRMYy0,30551 -deepspeed/runtime/pipe/p2p.py,sha256=t4K_Q2W2YI4K5bPOF1iqC5z7Ytyj4K1vHKA3gcSioBU,5344 -deepspeed/runtime/pipe/schedule.py,sha256=ZDzAeTPZTaYt77Wi2KSkycfuwneBGFsa7V6afQAzHo4,15546 -deepspeed/runtime/pipe/topology.py,sha256=Y_bQMiTNioGRb4twzWfsxlSkZc0iepbTbs9yvx7H-Cs,17698 -deepspeed/runtime/precision_config.py,sha256=HGZxwS8lygvliSHrtYLBl3B2eJmW37j7crRYzRt4UMQ,3612 -deepspeed/runtime/progressive_layer_drop.py,sha256=5zb3-BrMbRxYZL5lk1FtvaOknMi31xT1refDkgSMQbQ,1353 -deepspeed/runtime/quantize.py,sha256=OmimrTJV_RoVG8TLVdjcQR91kSatQCNdFkPf3U39fsk,7699 -deepspeed/runtime/sequence_parallel/__init__.py,sha256=2BSwDtJJ1v4uCmvYlXq0W17dbHPGC4v0-Xv0IfmtQb0,99 -deepspeed/runtime/sequence_parallel/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/sequence_parallel/__pycache__/parallel_state_sp.cpython-310.pyc,, -deepspeed/runtime/sequence_parallel/__pycache__/ulysses_sp.cpython-310.pyc,, -deepspeed/runtime/sequence_parallel/parallel_state_sp.py,sha256=SfNF80rWVfeByWao4GcT3wCBOibzTNqZlFi4GO8l2Fs,3677 -deepspeed/runtime/sequence_parallel/ulysses_sp.py,sha256=0n4mc4u70lhFHNlSZ5cJze3T533qnKdXHluPryErMH8,55816 -deepspeed/runtime/sparse_tensor.py,sha256=QBoplFkxSYLqE9Jfj8mN2qNLMMpVpEgpP3nuIsobpRg,2466 -deepspeed/runtime/state_dict_factory.py,sha256=Zp2mJbIySbYgp2cl5pFvLkpROOOqX9eqGnG1b3d0kD0,18177 -deepspeed/runtime/swap_tensor/__init__.py,sha256=OiMgE9mkwbOsv3gC91ENJRAHc7zlVfMXXhn4pxN7zcc,134 -deepspeed/runtime/swap_tensor/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/swap_tensor/__pycache__/aio_config.cpython-310.pyc,, -deepspeed/runtime/swap_tensor/__pycache__/async_swapper.cpython-310.pyc,, -deepspeed/runtime/swap_tensor/__pycache__/constants.cpython-310.pyc,, -deepspeed/runtime/swap_tensor/__pycache__/optimizer_utils.cpython-310.pyc,, -deepspeed/runtime/swap_tensor/__pycache__/partitioned_optimizer_swapper.cpython-310.pyc,, -deepspeed/runtime/swap_tensor/__pycache__/partitioned_param_swapper.cpython-310.pyc,, -deepspeed/runtime/swap_tensor/__pycache__/pipelined_optimizer_swapper.cpython-310.pyc,, -deepspeed/runtime/swap_tensor/__pycache__/utils.cpython-310.pyc,, -deepspeed/runtime/swap_tensor/aio_config.py,sha256=CaII3llLuEecOBceu7MuaqUtyFyzxARYE3Fspyj7DQ0,1643 -deepspeed/runtime/swap_tensor/async_swapper.py,sha256=RlwlqBbkHa3nviVT4nrYzZ4s05nPch_dNMHTZcFOdxw,6349 -deepspeed/runtime/swap_tensor/constants.py,sha256=lmGtuinGdUVZ8lEgvKKj0dVLUuiLRQTfPaEhxxWPg7E,700 -deepspeed/runtime/swap_tensor/optimizer_utils.py,sha256=MTQIkQ9zT0C6gpyS8zeB9et40z1BXFnh1e1z5FJBaZA,21403 -deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py,sha256=0-t1zOLHC_PMm5BnllC0BwjToFpYvBYYHby_ALmIOYw,10629 -deepspeed/runtime/swap_tensor/partitioned_param_swapper.py,sha256=U3Oopxi88nOvXx7wkDAt8eGnk1fuvIoV5ClhDce6QHc,18791 -deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py,sha256=AiFBvtcln-TrI3pSnaQzLTvaVTNBfiObA9_muL35R_M,11346 -deepspeed/runtime/swap_tensor/utils.py,sha256=Yrgbja_CefG1juXQVS-_KfnUu4vmP5Cw-r0p_fJ8DgU,7820 -deepspeed/runtime/tensor_parallel/__init__.py,sha256=BwFdU2Aa2rmV01DRps6I9hRTGsOiBzsTGH7oaEoeUDE,198 -deepspeed/runtime/tensor_parallel/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/tensor_parallel/__pycache__/config.cpython-310.pyc,, -deepspeed/runtime/tensor_parallel/__pycache__/tp_manager.cpython-310.pyc,, -deepspeed/runtime/tensor_parallel/config.py,sha256=rcNLwzsdTG1wInLFF1JI916VjJFYxhuoFuZZNWiZfmE,2860 -deepspeed/runtime/tensor_parallel/tp_manager.py,sha256=nPEscixH1Of-jaTsyrXzGo3qSHK0gaWDpfnVbZzbDhg,2539 -deepspeed/runtime/torch_autocast.py,sha256=ERSLVErVR1athv4vKhGPebAT8oBWzS1oQexgfCPt7wA,3638 -deepspeed/runtime/utils.py,sha256=4WyJnDEc7k87RxQco63P1PIxikNABzB62aYN7E0v34A,46087 -deepspeed/runtime/weight_quantizer.py,sha256=-aS6hTZnMcc3ckoYaCI6MiHybwy2NZo5AqEZ8rT7G_o,7027 -deepspeed/runtime/zero/__init__.py,sha256=3hNEJ-32gAoqG24c00D2lrC4WY0qFQPr0AbNSqZIazU,611 -deepspeed/runtime/zero/__pycache__/__init__.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/config.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/contiguous_memory_allocator.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/linear.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/mics.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/mics_utils.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/offload_config.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/offload_states.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/parameter_offload.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/partition_parameters.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/partitioned_param_coordinator.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/partitioned_param_profiler.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/stage3.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/stage_1_and_2.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/test.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/tiling.cpython-310.pyc,, -deepspeed/runtime/zero/__pycache__/utils.cpython-310.pyc,, -deepspeed/runtime/zero/config.py,sha256=tc0rFE0llwvM70WIyiLTc4iVWH8WZZhE0PftZp98AF8,14345 -deepspeed/runtime/zero/contiguous_memory_allocator.py,sha256=r1J-KRJqO_Aa4Rqlsb1cgtHPa4mpnxm_zrd19FzZVQU,10923 -deepspeed/runtime/zero/linear.py,sha256=vOhCaYR-E_Xg154-9VUsaF1rYJJUq6jytC_qxeSW9Zg,7998 -deepspeed/runtime/zero/mics.py,sha256=mKc4itpTLX3oQqciDiHHHmgF7cd_IRjOjgMv_5chk7E,22339 -deepspeed/runtime/zero/mics_utils.py,sha256=GWmBVESUd699i4YjCKScOwAwlwwFGs0si1dmPAeMz60,7500 -deepspeed/runtime/zero/offload_config.py,sha256=NN8h1kMNtH61GrSbnLTED2JnjzwGYL8v3IozOrSDIPI,3359 -deepspeed/runtime/zero/offload_states.py,sha256=WThDC2EhyaOArrehBNr-_VUNPc4hY5ZeTIXKsGh9Qhg,2956 -deepspeed/runtime/zero/parameter_offload.py,sha256=ObXU8MT1znWw8gg5vlIREvTzXqiGPIta4QBzVdAwFc8,26426 -deepspeed/runtime/zero/partition_parameters.py,sha256=Ob77bimPLPGXRsODOiIdnqg3jdNibK9MyVhNhd2omOk,110399 -deepspeed/runtime/zero/partitioned_param_coordinator.py,sha256=oVpLaxrhKbwsi3kvjlbJL79NSOhfmwQc4JJY5Ns9PUM,29471 -deepspeed/runtime/zero/partitioned_param_profiler.py,sha256=xTXgTnu1ew215TluuO2HvP78hZNbqPEhzJBieeOChf8,1801 -deepspeed/runtime/zero/stage3.py,sha256=88nREaCuBp62fCz_5uTRBiumtlSdLVvPK0y608QAcPU,155200 -deepspeed/runtime/zero/stage_1_and_2.py,sha256=A1t4UUHgHKBQ3KyTybM1LpvZ2qVdhtscfGWf9Sb73hE,126808 -deepspeed/runtime/zero/test.py,sha256=j7zu8Fnxm5_Wmp-aaUiQicPD1EUL2JGpODeSARFSZVs,2727 -deepspeed/runtime/zero/tiling.py,sha256=Uz6IgvJPAQ4GICXj4E509p_BIDvnczWVtBT7EjZU9rg,11727 -deepspeed/runtime/zero/utils.py,sha256=mpmn-5rjhKzd9BsXWoAf6cB90Ed4PnpynoCRgE0tBAs,5868 -deepspeed/sequence/__init__.py,sha256=4I9UpQ5vMRU5SYSF_dW9FJDEnBq4m_0SuwtVQ92lGaA,95 -deepspeed/sequence/__pycache__/__init__.cpython-310.pyc,, -deepspeed/sequence/__pycache__/cross_entropy.cpython-310.pyc,, -deepspeed/sequence/__pycache__/fpdt_layer.cpython-310.pyc,, -deepspeed/sequence/__pycache__/layer.cpython-310.pyc,, -deepspeed/sequence/cross_entropy.py,sha256=-YW4WYSHzZf_42OVu4WJuAp53GQAEnjDfIvGugAKA5k,2140 -deepspeed/sequence/fpdt_layer.py,sha256=MqPtxPXHxa5Gp6RFBIrQujHyJ_Qu8cJWH0g4gOEIGU0,56246 -deepspeed/sequence/layer.py,sha256=o-qB87cWrjJjf6152Xh0JWiCOzO-Kg6HCrAtEytHeW4,20516 -deepspeed/utils/__init__.py,sha256=n_hu2LhHCYt0EIvkNxEVJrSqyL_AhGWqQAAN6TkYPdg,1407 -deepspeed/utils/__pycache__/__init__.cpython-310.pyc,, -deepspeed/utils/__pycache__/bwc.cpython-310.pyc,, -deepspeed/utils/__pycache__/comms_logging.cpython-310.pyc,, -deepspeed/utils/__pycache__/config.cpython-310.pyc,, -deepspeed/utils/__pycache__/debug.cpython-310.pyc,, -deepspeed/utils/__pycache__/exceptions.cpython-310.pyc,, -deepspeed/utils/__pycache__/groups.cpython-310.pyc,, -deepspeed/utils/__pycache__/init_on_device.cpython-310.pyc,, -deepspeed/utils/__pycache__/logging.cpython-310.pyc,, -deepspeed/utils/__pycache__/mixed_precision_linkage.cpython-310.pyc,, -deepspeed/utils/__pycache__/numa.cpython-310.pyc,, -deepspeed/utils/__pycache__/nvtx.cpython-310.pyc,, -deepspeed/utils/__pycache__/tensor_fragment.cpython-310.pyc,, -deepspeed/utils/__pycache__/timer.cpython-310.pyc,, -deepspeed/utils/__pycache__/torch.cpython-310.pyc,, -deepspeed/utils/__pycache__/types.cpython-310.pyc,, -deepspeed/utils/__pycache__/z3_leaf_module.cpython-310.pyc,, -deepspeed/utils/__pycache__/zero_to_fp32.cpython-310.pyc,, -deepspeed/utils/bwc.py,sha256=EbdKjvtp_6XwUc3XhCyPxYfIqe9v5I23ynXGG39aAcY,3970 -deepspeed/utils/comms_logging.py,sha256=Bx8Gz6g1BPgMZ7FjtGLlqRyhKGuPPgyKMAGVdW4K_pg,7996 -deepspeed/utils/config.py,sha256=O4l-6Da_8RWBeyJ4tZpKqJPCLNtYK7eMVaqsVCbQkRU,1344 -deepspeed/utils/debug.py,sha256=ilbE4Z7cz3Ww0DR7TzP6vPQbJjyZFaq-ayO6FagSTAo,4936 -deepspeed/utils/exceptions.py,sha256=h4J_9uk3HmKG8LdiSO8DKYzOmIYrD1MGwtZD1nA-Q3g,144 -deepspeed/utils/groups.py,sha256=CNBrWWq3J9m04nnpZjaViZYNnZEtEp89LViE1LVnDog,30800 -deepspeed/utils/init_on_device.py,sha256=Q4RFeRMi7PYGBHfTKh3E63X8IZquheDWYKbT56aOuuQ,3004 -deepspeed/utils/logging.py,sha256=Khqk1omHmqqcYJgI88YW-JxY0zFnNYZ7QAfLAsixX40,5016 -deepspeed/utils/mixed_precision_linkage.py,sha256=OAzA20w_15nLUurjltJRTkGC__-aNx68KWrf5cbYfRg,2484 -deepspeed/utils/numa.py,sha256=ggGkhc8gXvuzkSqpl7JzmQ4K3H9Y44pnaq9XUAgBq-Q,7174 -deepspeed/utils/nvtx.py,sha256=XeqDmA2fdAUTH1FtFQwEr0K5Y7-gTWe9ZqLSp5qttec,746 -deepspeed/utils/tensor_fragment.py,sha256=6uEsYS1-eq8iz9iViCx8LzA231jIpSPIjSFC5QUdQ_o,17039 -deepspeed/utils/timer.py,sha256=V1iXDzsYWmuE0l1hoMKditAJMr1y8sShO5MZOfyF9r0,10742 -deepspeed/utils/torch.py,sha256=xxVsKHjO8HqeDJeffSed1TiP36IAlMn8lsakLuj-MGo,888 -deepspeed/utils/types.py,sha256=IcEwFod7RqLTRCGf7lr2sbKhFRf1an1KvtQ98nswwNw,434 -deepspeed/utils/z3_leaf_module.py,sha256=9eVQubxo8VxE-hBkrI3qxrddUKFoOBTUL4SCKCQiBe8,4758 -deepspeed/utils/zero_to_fp32.py,sha256=Rkl1ZczytKix9vGMg0EELzdJYFqUM1yB9p3xvSaK9k8,33272 diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/REQUESTED deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/WHEEL b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/WHEEL deleted file mode 100644 index e7fa31b6f3f78deb1022c1f7927f07d4d16da822..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/WHEEL +++ /dev/null @@ -1,5 +0,0 @@ -Wheel-Version: 1.0 -Generator: setuptools (80.9.0) -Root-Is-Purelib: true -Tag: py3-none-any - diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/entry_points.txt b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/entry_points.txt deleted file mode 100644 index cfa0d2ef62101b531baf7ac75575e64b2c0cb6a2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/entry_points.txt +++ /dev/null @@ -1,2 +0,0 @@ -[pytest_randomly.random_seeder] -deepspeed = deepspeed.runtime.utils:set_random_seed diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/licenses/LICENSE b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/licenses/LICENSE deleted file mode 100644 index 261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/licenses/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - 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 - 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 - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - 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. diff --git a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/top_level.txt deleted file mode 100644 index 0f2c819def28c8c4109156014fde317020d4ddb9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed-0.17.2.dist-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -deepspeed diff --git a/venv/lib/python3.10/site-packages/deepspeed/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/__init__.py deleted file mode 100644 index c7e01a0ac969325ec381d69a6a9174c3811cda52..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/__init__.py +++ /dev/null @@ -1,398 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import sys -import types -import json -from typing import Optional, Union -import torch -from torch.optim import Optimizer -from torch.optim.lr_scheduler import _LRScheduler -from packaging import version as pkg_version - -# Skip Triton import for AMD due to pytorch-triton-rocm module breaking device API in DeepSpeed -if not (hasattr(torch.version, 'hip') and torch.version.hip is not None): - try: - import triton # noqa: F401 # type: ignore - HAS_TRITON = True - except ImportError: - HAS_TRITON = False -else: - HAS_TRITON = False - -from . import ops -from . import module_inject - -from .accelerator import get_accelerator -from .constants import TORCH_DISTRIBUTED_DEFAULT_PORT -from .runtime.engine import DeepSpeedEngine, DeepSpeedOptimizerCallable, DeepSpeedSchedulerCallable -from .runtime.engine import ADAM_OPTIMIZER, LAMB_OPTIMIZER -from .runtime.hybrid_engine import DeepSpeedHybridEngine -from .runtime.pipe.engine import PipelineEngine -from .inference.engine import InferenceEngine -from .inference.config import DeepSpeedInferenceConfig -from .runtime.lr_schedules import add_tuning_arguments -from .runtime.config import DeepSpeedConfig, DeepSpeedConfigError -from .runtime.activation_checkpointing import checkpointing -from .ops.transformer import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig -from .module_inject import replace_transformer_layer, revert_transformer_layer, set_autotp_mode - -from .utils import log_dist, OnDevice, logger -from .comm.comm import init_distributed - -from .runtime import zero, domino -from .runtime.compiler import is_compile_supported - -from .pipe import PipelineModule - -from .git_version_info import version, git_hash, git_branch - - -def _parse_version(version_str): - '''Parse a version string and extract the major, minor, and patch versions.''' - ver = pkg_version.parse(version_str) - return ver.major, ver.minor, ver.micro - - -# Export version information -__version__ = version -__version_major__, __version_minor__, __version_patch__ = _parse_version(__version__) -__git_hash__ = git_hash -__git_branch__ = git_branch - -# Set to torch's distributed package or deepspeed.comm based inside DeepSpeedEngine init -dist = None - - -def initialize(args=None, - model: torch.nn.Module = None, - optimizer: Optional[Union[Optimizer, DeepSpeedOptimizerCallable]] = None, - model_parameters: Optional[torch.nn.Module] = None, - training_data: Optional[torch.utils.data.Dataset] = None, - lr_scheduler: Optional[Union[_LRScheduler, DeepSpeedSchedulerCallable]] = None, - distributed_port: int = TORCH_DISTRIBUTED_DEFAULT_PORT, - mpu=None, - dist_init_required: Optional[bool] = None, - collate_fn=None, - config=None, - mesh_param=None, - config_params=None): - """Initialize the DeepSpeed Engine. - - Arguments: - args: an object containing local_rank and deepspeed_config fields. - This is optional if `config` is passed. - - model: Required: nn.module class before apply any wrappers - - optimizer: Optional: a user defined Optimizer or Callable that returns an Optimizer object. - This overrides any optimizer definition in the DeepSpeed json config. - - model_parameters: Optional: An iterable of torch.Tensors or dicts. - Specifies what Tensors should be optimized. - - training_data: Optional: Dataset of type torch.utils.data.Dataset - - lr_scheduler: Optional: Learning Rate Scheduler Object or a Callable that takes an Optimizer and returns a Scheduler object. - The scheduler object should define a get_lr(), step(), state_dict(), and load_state_dict() methods - - distributed_port: Optional: Master node (rank 0)'s free port that needs to be used for communication during distributed training - - mpu: Optional: A model parallelism unit object that implements - get_{model,data}_parallel_{rank,group,world_size}() - - dist_init_required: Optional: None will auto-initialize torch distributed if needed, - otherwise the user can force it to be initialized or not via boolean. - - collate_fn: Optional: Merges a list of samples to form a - mini-batch of Tensor(s). Used when using batched loading from a - map-style dataset. - - config: Optional: Instead of requiring args.deepspeed_config you can pass your deepspeed config - as an argument instead, as a path or a dictionary. - - config_params: Optional: Same as `config`, kept for backwards compatibility. - - Returns: - A tuple of ``engine``, ``optimizer``, ``training_dataloader``, ``lr_scheduler`` - - * ``engine``: DeepSpeed runtime engine which wraps the client model for distributed training. - - * ``optimizer``: Wrapped optimizer if a user defined ``optimizer`` is supplied, or if - optimizer is specified in json config else ``None``. - - * ``training_dataloader``: DeepSpeed dataloader if ``training_data`` was supplied, - otherwise ``None``. - - * ``lr_scheduler``: Wrapped lr scheduler if user ``lr_scheduler`` is passed, or - if ``lr_scheduler`` specified in JSON configuration. Otherwise ``None``. - """ - log_dist("DeepSpeed info: version={}, git-hash={}, git-branch={}".format(__version__, __git_hash__, - __git_branch__), - ranks=[0]) - - # Disable zero.Init context if it's currently enabled - zero.partition_parameters.shutdown_init_context() - - assert model is not None, "deepspeed.initialize requires a model" - - global dist - from deepspeed import comm as dist - dist_backend = get_accelerator().communication_backend_name() - dist.init_distributed(dist_backend=dist_backend, - distributed_port=distributed_port, - dist_init_required=dist_init_required) - - ##TODO: combine reuse mpu as mesh device and vice versa - # Set config using config_params for backwards compat - if config is None and config_params is not None: - config = config_params - - mesh_device = None - if mesh_param: - logger.info(f"mesh_param to Initialize mesh device: {mesh_param}") - mesh_device = dist.initialize_mesh_device(mesh_param, ("data_parallel", "sequence_parallel")) - #if config file has sequence parallelize and data parallelize, then use them to initialize mesh device - elif config is not None: - if "sequence_parallel_size" in config and "data_parallel_size" in config: - logger.info(f"config to Initialize mesh device: {config}") - mesh_device = dist.initialize_mesh_device((config["data_parallel_size"], config["sequence_parallel_size"]), \ - ("data_parallel", "sequence_parallel")) - - # Check for deepscale_config for backwards compat - if hasattr(args, "deepscale_config") and args.deepscale_config is not None: - logger.warning("************ --deepscale_config is deprecated, please use --deepspeed_config ************") - if hasattr(args, "deepspeed_config"): - assert (args.deepspeed_config - is None), "Not sure how to proceed, we were given both a deepscale_config and deepspeed_config" - args.deepspeed_config = args.deepscale_config - args.deepscale_config = None - - # Check that we have only one config passed - if hasattr(args, "deepspeed_config") and args.deepspeed_config is not None: - assert config is None, "Not sure how to proceed, we were given deepspeed configs in the deepspeed arguments and deepspeed.initialize() function call" - config = args.deepspeed_config - assert config is not None, "DeepSpeed requires --deepspeed_config to specify configuration file" - if not isinstance(model, PipelineModule): - config_class = DeepSpeedConfig(config, mpu, mesh_device=mesh_device) - if config_class.hybrid_engine.enabled: - engine = DeepSpeedHybridEngine(args=args, - model=model, - optimizer=optimizer, - model_parameters=model_parameters, - training_data=training_data, - lr_scheduler=lr_scheduler, - mpu=mpu, - dist_init_required=dist_init_required, - collate_fn=collate_fn, - config=config, - config_class=config_class) - else: - engine = DeepSpeedEngine(args=args, - model=model, - optimizer=optimizer, - model_parameters=model_parameters, - training_data=training_data, - lr_scheduler=lr_scheduler, - mpu=mpu, - dist_init_required=dist_init_required, - collate_fn=collate_fn, - config=config, - mesh_device=mesh_device, - config_class=config_class) - else: - assert mpu is None, "mpu must be None with pipeline parallelism" - mpu = model.mpu() - config_class = DeepSpeedConfig(config, mpu) - engine = PipelineEngine(args=args, - model=model, - optimizer=optimizer, - model_parameters=model_parameters, - training_data=training_data, - lr_scheduler=lr_scheduler, - mpu=mpu, - dist_init_required=dist_init_required, - collate_fn=collate_fn, - config=config, - config_class=config_class) - - # Restore zero.Init context if necessary - zero.partition_parameters.restore_init_context() - - return_items = [ - engine, - engine.optimizer, - engine.training_dataloader, - engine.lr_scheduler, - ] - return tuple(return_items) - - -def _add_core_arguments(parser): - r"""Helper (internal) function to update an argument parser with an argument group of the core DeepSpeed arguments. - The core set of DeepSpeed arguments include the following: - 1) --deepspeed: boolean flag to enable DeepSpeed - 2) --deepspeed_config : path of a json configuration file to configure DeepSpeed runtime. - - This is a helper function to the public add_config_arguments() - - Arguments: - parser: argument parser - Return: - parser: Updated Parser - """ - group = parser.add_argument_group('DeepSpeed', 'DeepSpeed configurations') - - group.add_argument('--deepspeed', - default=False, - action='store_true', - help='Enable DeepSpeed (helper flag for user code, no impact on DeepSpeed backend)') - - group.add_argument('--deepspeed_config', default=None, type=str, help='DeepSpeed json configuration file.') - - group.add_argument('--deepscale', - default=False, - action='store_true', - help='Deprecated enable DeepSpeed (helper flag for user code, no impact on DeepSpeed backend)') - - group.add_argument('--deepscale_config', - default=None, - type=str, - help='Deprecated DeepSpeed json configuration file.') - - return parser - - -def add_config_arguments(parser): - r"""Update the argument parser to enabling parsing of DeepSpeed command line arguments. - The set of DeepSpeed arguments include the following: - 1) --deepspeed: boolean flag to enable DeepSpeed - 2) --deepspeed_config : path of a json configuration file to configure DeepSpeed runtime. - - Arguments: - parser: argument parser - Return: - parser: Updated Parser - """ - parser = _add_core_arguments(parser) - - return parser - - -def default_inference_config(): - """ - Return a default DeepSpeed inference configuration dictionary. - """ - return DeepSpeedInferenceConfig().dict() - - -def init_inference(model, config=None, **kwargs): - """Initialize the DeepSpeed InferenceEngine. - - Description: all four cases are valid and supported in DS init_inference() API. - - # Case 1: user provides no config and no kwargs. Default config will be used. - - .. code-block:: python - - generator.model = deepspeed.init_inference(generator.model) - string = generator("DeepSpeed is") - print(string) - - # Case 2: user provides a config and no kwargs. User supplied config will be used. - - .. code-block:: python - - generator.model = deepspeed.init_inference(generator.model, config=config) - string = generator("DeepSpeed is") - print(string) - - # Case 3: user provides no config and uses keyword arguments (kwargs) only. - - .. code-block:: python - - generator.model = deepspeed.init_inference(generator.model, - tensor_parallel={"tp_size": world_size}, - dtype=torch.half, - replace_with_kernel_inject=True) - string = generator("DeepSpeed is") - print(string) - - # Case 4: user provides config and keyword arguments (kwargs). Both config and kwargs are merged and kwargs take precedence. - - .. code-block:: python - - generator.model = deepspeed.init_inference(generator.model, config={"dtype": torch.half}, replace_with_kernel_inject=True) - string = generator("DeepSpeed is") - print(string) - - Arguments: - model: Required: original nn.module object without any wrappers - - config: Optional: instead of arguments, you can pass in a DS inference config dict or path to JSON file - - Returns: - A deepspeed.InferenceEngine wrapped model. - """ - log_dist("DeepSpeed info: version={}, git-hash={}, git-branch={}".format(__version__, __git_hash__, - __git_branch__), - ranks=[0]) - - # Load config_dict from config first - if config is None: - config = {} - if isinstance(config, str): - with open(config, "r") as f: - config_dict = json.load(f) - elif isinstance(config, dict): - config_dict = config - else: - raise ValueError(f"'config' argument expected string or dictionary, got {type(config)}") - - # Update with values from kwargs, ensuring no conflicting overlap between config and kwargs - overlap_keys = set(config_dict.keys()).intersection(kwargs.keys()) - # If there is overlap, error out if values are different - for key in overlap_keys: - if config_dict[key] != kwargs[key]: - raise ValueError(f"Conflicting argument '{key}' in 'config':{config_dict[key]} and kwargs:{kwargs[key]}") - config_dict.update(kwargs) - - ds_inference_config = DeepSpeedInferenceConfig(**config_dict) - - engine = InferenceEngine(model, config=ds_inference_config) - - return engine - - -def tp_model_init(model, tp_size, dtype, config=None, **kwargs): - """ - Initialize the model for tensor parallelism. - - Args: - model (torch.nn.Module): The model to be initialized. - tp_size (int): The tensor parallelism size. - dtype (torch.dtype): The data type to be used for the model. - - Returns: - torch.nn.Module: The initialized model with tensor parallelism. - """ - # avoid re-entry - if hasattr(model, 'ds_autotp_parsed'): - logger.warning("ds_autotp_parsed' attribute already exists in the model, re-entry is not allowed.") - return - - set_autotp_mode(training=True) - - from deepspeed.runtime.tensor_parallel import TpTrainingManager - # The expected usage here is for it to be invoked by transformers package. - - #TODO: We should provide a custom TP mapping solution without using autoTP - #as modifying the autoTP logic may be more difficult for users compared to configuring it - - model = TpTrainingManager(model=model, tp_size=tp_size, dtype=dtype).module - - setattr(model, 'ds_autotp_parsed', True) - - return model diff --git a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 6981f53b588133cc939af983dc429aa1adaf3c0f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/constants.cpython-310.pyc deleted file mode 100644 index 140c17b6a46dfe9ef070030dbea92368bcd620c2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/constants.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/env_report.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/env_report.cpython-310.pyc deleted file mode 100644 index 635373e3cc9c1fdb7db868634a31e61bef97f562..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/env_report.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/git_version_info.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/git_version_info.cpython-310.pyc deleted file mode 100644 index acf9a7a17f9b6267d68e26a26f194ea13453a1fd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/git_version_info.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/git_version_info_installed.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/__pycache__/git_version_info_installed.cpython-310.pyc deleted file mode 100644 index 8f8e1e7da2799ef0cd75e51024adef5c0fc65c06..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/__pycache__/git_version_info_installed.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/constants.py b/venv/lib/python3.10/site-packages/deepspeed/constants.py deleted file mode 100644 index 8378e8f2264fd4e99b7309dc21a9844afa19c54e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/constants.py +++ /dev/null @@ -1,25 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -from datetime import timedelta - -############################################# -# Torch distributed constants -############################################# -TORCH_DISTRIBUTED_DEFAULT_PORT = 29500 - -# Default process group wide timeout, if applicable. -# This only applies to the gloo and nccl backends -# (only if NCCL_BLOCKING_WAIT or NCCL_ASYNC_ERROR_HANDLING is set to 1). -# To make an attempt at backwards compatibility with THD, we use an -# extraordinarily high default timeout, given that THD did not have timeouts. -default_pg_timeout = timedelta(minutes=int(os.getenv("DEEPSPEED_TIMEOUT", default=30))) -INFERENCE_GENERIC_MODE = 'generic' -INFERENCE_SPECIALIZED_MODE = 'specialized' - -CROSS_RANK = "CROSS_RANK" -CROSS_SIZE = "CROSS_SIZE" -LOCAL_RANK = 'LOCAL_RANK' diff --git a/venv/lib/python3.10/site-packages/deepspeed/env_report.py b/venv/lib/python3.10/site-packages/deepspeed/env_report.py deleted file mode 100644 index 37e33b1e873972e10194ef70a7fd4470af6a7e30..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/env_report.py +++ /dev/null @@ -1,195 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import torch -import deepspeed -import subprocess -import argparse -from .ops.op_builder.all_ops import ALL_OPS -from .git_version_info import installed_ops, torch_info, accelerator_name -from deepspeed.accelerator import get_accelerator - -GREEN = '\033[92m' -RED = '\033[91m' -YELLOW = '\033[93m' -END = '\033[0m' -SUCCESS = f"{GREEN} [SUCCESS] {END}" -OKAY = f"{GREEN}[OKAY]{END}" -WARNING = f"{YELLOW}[WARNING]{END}" -FAIL = f'{RED}[FAIL]{END}' -INFO = '[INFO]' - -color_len = len(GREEN) + len(END) -okay = f"{GREEN}[OKAY]{END}" -warning = f"{YELLOW}[WARNING]{END}" - - -def op_report(verbose=True): - max_dots = 23 - max_dots2 = 11 - h = ["op name", "installed", "compatible"] - print("-" * (max_dots + max_dots2 + len(h[0]) + len(h[1]))) - print("DeepSpeed C++/CUDA extension op report") - print("-" * (max_dots + max_dots2 + len(h[0]) + len(h[1]))) - - print("NOTE: Ops not installed will be just-in-time (JIT) compiled at\n" - " runtime if needed. Op compatibility means that your system\n" - " meet the required dependencies to JIT install the op.") - - print("-" * (max_dots + max_dots2 + len(h[0]) + len(h[1]))) - print("JIT compiled ops requires ninja") - ninja_status = OKAY if ninja_installed() else FAIL - print('ninja', "." * (max_dots - 5), ninja_status) - print("-" * (max_dots + max_dots2 + len(h[0]) + len(h[1]))) - print(h[0], "." * (max_dots - len(h[0])), h[1], "." * (max_dots2 - len(h[1])), h[2]) - print("-" * (max_dots + max_dots2 + len(h[0]) + len(h[1]))) - installed = f"{GREEN}[YES]{END}" - no = f"{YELLOW}[NO]{END}" - for op_name, builder in ALL_OPS.items(): - dots = "." * (max_dots - len(op_name)) - is_compatible = OKAY if builder.is_compatible(verbose) else no - is_installed = installed if installed_ops.get(op_name, - False) and accelerator_name == get_accelerator()._name else no - dots2 = '.' * ((len(h[1]) + (max_dots2 - len(h[1]))) - (len(is_installed) - color_len)) - print(op_name, dots, is_installed, dots2, is_compatible) - print("-" * (max_dots + max_dots2 + len(h[0]) + len(h[1]))) - - -def ninja_installed(): - try: - import ninja # noqa: F401 # type: ignore - except ImportError: - return False - return True - - -def nvcc_version(): - import torch.utils.cpp_extension - cuda_home = torch.utils.cpp_extension.CUDA_HOME - if cuda_home is None: - return f"{RED} [FAIL] cannot find CUDA_HOME via torch.utils.cpp_extension.CUDA_HOME={torch.utils.cpp_extension.CUDA_HOME} {END}" - try: - output = subprocess.check_output([cuda_home + "/bin/nvcc", "-V"], universal_newlines=True) - except FileNotFoundError: - return f"{RED} [FAIL] nvcc missing {END}" - output_split = output.split() - release_idx = output_split.index("release") - release = output_split[release_idx + 1].replace(',', '').split(".") - return ".".join(release) - - -def installed_cann_path(): - if "ASCEND_HOME_PATH" in os.environ or os.path.exists(os.environ["ASCEND_HOME_PATH"]): - return os.environ["ASCEND_HOME_PATH"] - return None - - -def installed_cann_version(): - import re - ascend_path = installed_cann_path() - if ascend_path is None: - return f"CANN_HOME does not exist, unable to compile NPU op(s)" - cann_version = "" - for dirpath, _, filenames in os.walk(os.path.realpath(ascend_path)): - if cann_version: - break - install_files = [file for file in filenames if re.match(r"ascend_.*_install\.info", file)] - if install_files: - filepath = os.path.join(dirpath, install_files[0]) - with open(filepath, "r") as f: - for line in f: - if line.find("version") != -1: - cann_version = line.strip().split("=")[-1] - break - return cann_version - - -def get_shm_size(): - try: - shm_stats = os.statvfs('/dev/shm') - except (OSError, FileNotFoundError, ValueError, AttributeError): - return "UNKNOWN", None - - shm_size = shm_stats.f_frsize * shm_stats.f_blocks - shm_hbytes = human_readable_size(shm_size) - warn = [] - if shm_size < 512 * 1024**2: - warn.append( - f" {YELLOW} [WARNING] /dev/shm size might be too small, if running in docker increase to at least --shm-size='1gb' {END}" - ) - if get_accelerator().communication_backend_name() == "nccl": - warn.append( - f" {YELLOW} [WARNING] see more details about NCCL requirements: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/troubleshooting.html#sharing-data {END}" - ) - return shm_hbytes, warn - - -def human_readable_size(size): - units = ['B', 'KB', 'MB', 'GB', 'TB'] - i = 0 - while size >= 1024 and i < len(units) - 1: - size /= 1024 - i += 1 - return f'{size:.2f} {units[i]}' - - -def debug_report(): - max_dots = 33 - - report = [("torch install path", torch.__path__), ("torch version", torch.__version__), - ("deepspeed install path", deepspeed.__path__), - ("deepspeed info", f"{deepspeed.__version__}, {deepspeed.__git_hash__}, {deepspeed.__git_branch__}")] - if get_accelerator().device_name() == 'cuda': - hip_version = getattr(torch.version, "hip", None) - report.extend([("torch cuda version", torch.version.cuda), ("torch hip version", hip_version), - ("nvcc version", (None if hip_version else nvcc_version())), - ("deepspeed wheel compiled w.", f"torch {torch_info['version']}, " + - (f"hip {torch_info['hip_version']}" if hip_version else f"cuda {torch_info['cuda_version']}")) - ]) - elif get_accelerator().device_name() == 'npu': - import torch_npu - report.extend([("deepspeed wheel compiled w.", f"torch {torch_info['version']}"), - ("torch_npu install path", torch_npu.__path__), ("torch_npu version", torch_npu.__version__), - ("ascend_cann version", installed_cann_version())]) - else: - report.extend([("deepspeed wheel compiled w.", f"torch {torch_info['version']} ")]) - - report.append(("shared memory (/dev/shm) size", get_shm_size())) - - print("DeepSpeed general environment info:") - for name, value in report: - warns = [] - if isinstance(value, tuple): - value, warns = value - print(name, "." * (max_dots - len(name)), value) - if warns: - for warn in warns: - print(warn) - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--hide_operator_status', - action='store_true', - help='Suppress display of installation and compatibility statuses of DeepSpeed operators. ') - parser.add_argument('--hide_errors_and_warnings', action='store_true', help='Suppress warning and error messages.') - args = parser.parse_args() - return args - - -def main(hide_operator_status=False, hide_errors_and_warnings=False): - if not hide_operator_status: - op_report(verbose=not hide_errors_and_warnings) - debug_report() - - -def cli_main(): - args = parse_arguments() - main(hide_operator_status=args.hide_operator_status, hide_errors_and_warnings=args.hide_errors_and_warnings) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/deepspeed/git_version_info.py b/venv/lib/python3.10/site-packages/deepspeed/git_version_info.py deleted file mode 100644 index 70c536d2f78eee24c21db985c43fc9a3c017d4ee..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/git_version_info.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -try: - # This is populated by setup.py - from .git_version_info_installed import * # noqa: F401 # type: ignore -except ModuleNotFoundError: - import os - if os.path.isfile('version.txt'): - # Will be missing from checkouts that haven't been installed (e.g., readthedocs) - version = open('version.txt', 'r').read().strip() - else: - version = "0.0.0" - git_hash = '[none]' - git_branch = '[none]' - - from .ops.op_builder.all_ops import ALL_OPS - installed_ops = dict.fromkeys(ALL_OPS.keys(), False) - accelerator_name = "" - torch_info = {'version': "0.0", "cuda_version": "0.0", "hip_version": "0.0"} - -# compatible_ops list is recreated for each launch -from .ops.op_builder.all_ops import ALL_OPS - -compatible_ops = dict.fromkeys(ALL_OPS.keys(), False) -for op_name, builder in ALL_OPS.items(): - op_compatible = builder.is_compatible() - compatible_ops[op_name] = op_compatible - compatible_ops["deepspeed_not_implemented"] = False diff --git a/venv/lib/python3.10/site-packages/deepspeed/git_version_info_installed.py b/venv/lib/python3.10/site-packages/deepspeed/git_version_info_installed.py deleted file mode 100644 index c2692d4544024549a3b18539ed5266f8aebf76f2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/git_version_info_installed.py +++ /dev/null @@ -1,6 +0,0 @@ -version='0.17.2' -git_hash='unknown' -git_branch='unknown' -installed_ops={'deepspeed_not_implemented': False, 'async_io': False, 'deepspeed_ccl_comm': False, 'deepspeed_shm_comm': False, 'cpu_adam': False, 'fused_adam': False} -accelerator_name='cpu' -torch_info={'version': '0.0', 'bf16_support': False, 'cuda_version': '0.0', 'nccl_version': '0.0', 'hip_version': '0.0'} diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/__init__.py deleted file mode 100644 index cdd00fec935b826c821c8ca9d62e2711a91ea811..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -from .v2 import RaggedInferenceEngineConfig, DeepSpeedTPConfig -from .v2.engine_v2 import InferenceEngineV2 -from .v2 import build_hf_engine, build_engine_from_ds_checkpoint diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 8fc31fa34764571cdfeff7d04b4e56c12e859f7f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/__pycache__/config.cpython-310.pyc deleted file mode 100644 index 104e939e8a045c7c5fc6436d90ab442a20a570b7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/__pycache__/engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/__pycache__/engine.cpython-310.pyc deleted file mode 100644 index a02f64e465be9bee4eb116c6e8172b587cb7fc54..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/__pycache__/engine.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/config.py b/venv/lib/python3.10/site-packages/deepspeed/inference/config.py deleted file mode 100644 index 6df61f7c8841df35540230cf26874418f58cf5a3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/config.py +++ /dev/null @@ -1,323 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import deepspeed -from pydantic import Field, field_validator -from deepspeed.runtime.config_utils import DeepSpeedConfigModel -from deepspeed.runtime.zero.config import DeepSpeedZeroConfig -from typing import Dict, Union, Optional -from enum import Enum - - -class DtypeEnum(Enum): - fp16 = (torch.float16, "torch.float16", "fp16", "float16", "half") - fp32 = (torch.float32, "torch.float32", "fp32", "float32", "float") - bf16 = (torch.bfloat16, "torch.bfloat16", "bf16", "bfloat16", "bfloat") - int8 = (torch.int8, "torch.int8", "int8") - - @classmethod - def from_str(cls, value: str): - for dtype in cls: - if value in dtype.value: - return dtype - raise ValueError(f"'{value}' is not a valid DtypeEnum") - - -class MoETypeEnum(str, Enum): - residual = "residual" - standard = "standard" - - -class DeepSpeedTPConfig(DeepSpeedConfigModel): - """ Configure tensor parallelism settings """ - - enabled: bool = True - """ Turn tensor parallelism on/off. """ - - tp_size: int = 1 - """ Number of devices to split the model across using tensor parallelism. """ - - tp_grain_size: int = 64 - "Desired MLP/lm_head tp size granularity. DNN library favors tensor size in granularity of power of 2, we pick 64 as a default size." - - mpu: object = None - """ - A model parallelism unit object that implements - ``get_{model,data}_parallel_{rank,group,world_size}()``. - """ - - tp_group: object = None - - -class DeepSpeedMoEConfig(DeepSpeedConfigModel): - """ Sets parameters for MoE """ - - enabled: bool = True - ep_size: int = 1 - """ - The expert-parallelism size which is used for partitioning the experts - across the GPUs in the expert-parallel group. - """ - - moe_experts: list = Field([1], alias="num_experts") - """ The global number of experts used in an MoE layer. """ - - type: MoETypeEnum = MoETypeEnum.standard - """ - Specify the type of MoE layer. We have two types of MoE layer: 'Standard' - and 'Residual'. - """ - - ep_mp_group: object = None - ep_group: object = Field(None, alias="expert_group") - - -class QuantTypeEnum(str, Enum): - asym = "asymmetric" - sym = "symmetric" - - -class BaseQuantConfig(DeepSpeedConfigModel): - enabled: bool = True - num_bits: int = 8 - q_type: QuantTypeEnum = QuantTypeEnum.sym - q_groups: int = 1 - - -class WeightQuantConfig(BaseQuantConfig): - enabled: bool = True - quantized_initialization: Dict = {} - post_init_quant: Dict = {} - - -class ActivationQuantConfig(BaseQuantConfig): - enabled: bool = True - - -class QKVQuantConfig(DeepSpeedConfigModel): - enabled: bool = True - - -class QuantizationConfig(DeepSpeedConfigModel): - enabled: bool = True - activation: ActivationQuantConfig = ActivationQuantConfig() - weight: WeightQuantConfig = WeightQuantConfig() - qkv: QKVQuantConfig = QKVQuantConfig() - - -# todo: brainstorm on how to do ckpt loading for DS inference -class InferenceCheckpointConfig(DeepSpeedConfigModel): - checkpoint_dir: Optional[str] = None - save_mp_checkpoint_path: Optional[str] = None - base_dir: Optional[str] = None - - -class DeepSpeedInferenceConfig(DeepSpeedConfigModel): - """ Sets parameters for DeepSpeed Inference Engine. """ - - replace_with_kernel_inject: bool = Field(False, alias="kernel_inject") - """ - Set to true to inject inference kernels for models such as, Bert, GPT2, - GPT-Neo and GPT-J. Otherwise, the injection_dict provides the names of two - linear layers as a tuple: - `(attention_output projection, transformer output projection)` - """ - - dtype: torch.dtype = torch.float16 - """ - Desired model data type, will convert model to this type. - Supported target types: `torch.half`, `torch.int8`, `torch.float` - """ - - tensor_parallel: DeepSpeedTPConfig = Field({}, alias="tp") - """ - Configuration for tensor parallelism used to split the model across several - GPUs. Expects a dictionary containing values for :any:`DeepSpeedTPConfig`. - """ - - enable_cuda_graph: bool = False - """ - Use this flag for capturing the CUDA-Graph of the inference ops, so that it - can run faster using the graph replay method. - """ - - use_triton: bool = False - """ - Use this flag to use triton kernels for inference ops. - """ - - triton_autotune: bool = False - """ - Use this flag to enable triton autotuning. - Turning it on is better for performance but increase the 1st runtime for - autotuning. - """ - - zero: DeepSpeedZeroConfig = {} - """ - ZeRO configuration to use with the Inference Engine. Expects a dictionary - containing values for :any:`DeepSpeedZeroConfig`. - """ - - triangular_masking: bool = Field(True, alias="tm") - """ - Controls the type of masking for attention scores in transformer layer. - Note that the masking is application specific. - """ - - moe: Union[bool, DeepSpeedMoEConfig] = {} - """ - Specify if the type of Transformer is MoE. Expects a dictionary containing - values for :any:`DeepSpeedMoEConfig`. - """ - - keep_module_on_host: bool = False - """ - When loading checkpoints to model parameters, they are moved to the device. In very large models - this might fill the device and cause OOM. Setting this flag to true, will keep checkpoints on - host and not move them directly to the device (giving an option to quantize checkpoint data before - moving it to the device for example). - Set only for models with injection policies and auto TP. - """ - - quant: QuantizationConfig = {} - """ - NOTE: only works for int8 dtype. - Quantization settings used for quantizing your model using the MoQ. The - setting can be one element or a tuple. If one value is passed in, we - consider it as the number of groups used in quantization. A tuple is passed - in if we want to mention that there is extra-grouping for the MLP part of a - Transformer layer (e.g. (True, 8) shows we quantize the model using 8 - groups for all the network except the MLP part that we use 8 extra - grouping). Expects a dictionary containing values for - :any:`QuantizationConfig`. - """ - - #todo: refactor the following 3 into the new checkpoint_config - checkpoint: Optional[Union[str, Dict]] = None - """ - Path to deepspeed compatible checkpoint or path to JSON with load policy. - """ - - base_dir: str = "" - """ - This shows the root directory under which all the checkpoint files exists. - This can be passed through the json config too. - """ - - set_empty_params: bool = False - """ - specifying whether the inference-module is created with empty or real Tensor - """ - - save_mp_checkpoint_path: Optional[str] = None - """ - The path for which we want to save the loaded model with a checkpoint. This - feature is used for adjusting the parallelism degree to help alleviate the - model loading overhead. It does not save any new checkpoint if no path is - passed. - """ - - checkpoint_config: InferenceCheckpointConfig = Field({}, alias="ckpt_config") - """ - TODO: Add docs. Expects a dictionary containing values for - :any:`InferenceCheckpointConfig`. - """ - - return_tuple: bool = True - """ - Specify whether or not the transformer layers need to return a tuple or a - Tensor. - """ - - training_mp_size: int = 1 - """ - If loading a checkpoint this is the mp size that it was trained with, it - may be different than what the mp size that you want to use during - inference. - """ - - replace_method: str = Field( - "auto", - json_schema_extra={ - "deprecated": True, - "deprecated_msg": "This parameter is no longer needed, please remove from your call to DeepSpeed-inference" - }) - - injection_policy: Optional[Dict] = Field(None, alias="injection_dict") - """ - Dictionary mapping a client nn.Module to its corresponding injection - policy. e.g., `{BertLayer : deepspeed.inference.HFBertLayerPolicy}` - """ - - injection_policy_tuple: Optional[tuple] = None - """ TODO: Add docs """ - - config: Optional[Dict] = Field(None, alias="args") # todo: really no need for this field if we can refactor - - max_out_tokens: int = Field(1024, alias="max_tokens") - """ - This argument shows the maximum number of tokens inference-engine can work - with, including the input and output tokens. Please consider increasing it - to the required token-length required for your use-case. - """ - - min_out_tokens: int = Field(1, alias="min_tokens") - """ - This argument communicates to the runtime the minimum number of tokens you - expect you will need to generate. This will cause the runtime to error - if it unable to provide this and provide context on the memory pressure - rather than seg-faulting or providing corrupted output. - """ - - transposed_mode: bool = Field(False, alias="transposed_mode") - - mp_size: int = Field(1, json_schema_extra={"deprecated": True, "new_param": "tensor_parallel.tp_size"}) - """ - Desired model parallel size, default is 1 meaning no model parallelism. - Deprecated, please use the ``tensor_parallel` config to control model - parallelism. - """ - mpu: object = Field(None, json_schema_extra={"deprecated": True, "new_param": "tensor_parallel.mpu"}) - ep_size: int = Field(1, json_schema_extra={"deprecated": True, "new_param": "moe.ep_size"}) - ep_group: object = Field(None, - alias="expert_group", - json_schema_extra={ - "deprecated": True, - "new_param": "moe.ep_group" - }) - ep_mp_group: object = Field(None, - alias="expert_mp_group", - json_schema_extra={ - "deprecated": True, - "new_param": "moe.ep_mp_group" - }) - moe_experts: list = Field([1], json_schema_extra={"deprecated": True, "new_param": "moe.moe_experts"}) - moe_type: MoETypeEnum = Field(MoETypeEnum.standard, - json_schema_extra={ - "deprecated": True, - "new_param": "moe.type" - }) - - @field_validator("dtype", mode="before") - def validate_dtype(cls, field_value, values): - if isinstance(field_value, str): - return DtypeEnum.from_str(field_value).value[0] - if isinstance(field_value, torch.dtype): - return field_value - raise TypeError(f"Invalid type for dtype: {type(field_value)}") - - @field_validator("moe") - def moe_backward_compat(cls, field_value, values): - if isinstance(field_value, bool): - return DeepSpeedMoEConfig(moe=field_value) - return field_value - - @field_validator("use_triton") - def has_triton(cls, field_value, values): - if field_value and not deepspeed.HAS_TRITON: - raise ValueError('Triton needs to be installed to use deepspeed with triton kernels') - return field_value diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/engine.py b/venv/lib/python3.10/site-packages/deepspeed/inference/engine.py deleted file mode 100644 index 0a74d19e91f5a6bfe8206db0503ca3efd95b835a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/engine.py +++ /dev/null @@ -1,625 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import time -import os -import deepspeed -from deepspeed import comm as dist -from deepspeed.utils.logging import log_dist - -from torch.nn.modules import Module -from packaging import version as pkg_version -from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine -from deepspeed.utils.timer import SynchronizedWallClockTimer -from deepspeed.runtime.compiler import is_compile_supported -from ..runtime.state_dict_factory import SDLoaderFactory -from ..runtime.weight_quantizer import WeightQuantization -from ..module_inject import replace_transformer_layer, generic_injection -from ..comm.comm import init_distributed -from ..pipe import PipelineModule -from ..moe.utils import has_moe_layers -from ..module_inject import LinearAllreduce, LinearLayer, Normalize, ReplaceWithTensorSlicing -from deepspeed.accelerator import get_accelerator -from ..module_inject.policy import TransformerPolicy -from ..module_inject.auto_tp import AutoTP - -from ..module_inject.replace_policy import generic_policies -from ..module_inject.auto_tp_model_utils import build_bloom_alibi_tensor, build_mpt_atten_bias_tensor, build_mpt_alibi_tensor, get_alibi_mask -from ..ops.transformer.inference.ds_attention import DeepSpeedSelfAttention -from ..model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference - -DS_INFERENCE_ENABLED = False -from torch import nn - -INFERENCE_MODEL_TIMER = "model-forward-inference" - - -class InferenceEngine(Module): - inference_mp_group = None - inference_ep_group = None - expert_mp_group = None - - def __init__(self, model, config): - """ - Args: - model: torch.nn.Module - config: DeepSpeedInferenceConfig - """ - global DS_INFERENCE_ENABLED - DS_INFERENCE_ENABLED = True - - super().__init__() - if DeepSpeedTransformerInference.workspace is not None: - self.destroy() - - self.module = model - self._config = config - - self._get_model_config_generate(config) # keep for weird backward compatibility - - # patch model generate with ours if model uses it - if hasattr(self.module, "generate"): - self.generate = self._generate - - if hasattr(self.module, "config"): - TransformerPolicy.hf_model_config = self.module.config - - if config.dtype not in get_accelerator().supported_dtypes(): - raise ValueError( - f"Data type {config.dtype} is not supported by {get_accelerator().device_name()} accelerator") - - # todo: keep this self.injection_dict because we don't use to change config.injection_policy API - # todo: this will get changed when Molly's PR on auto injection dict is merged - self.injection_dict = config.injection_policy - - # todo: refactor the mp_group and mp_size related in the next refactor - self.mp_group = config.tensor_parallel.tp_group - self.mpu = config.tensor_parallel.mpu - - self.quantize_merge_count = 1 - self.quantization_scales = None - - # these are not needed in the config as we are creating them ourselves in the inference engine - self.ep_group = None # config.moe.ep_group - self.expert_mp_group = None # config.moe.ep_mp_group - - self.cuda_graph_created = False - self.checkpoint_engine = TorchCheckpointEngine() - quantization_setting = None - self._init_quantization_setting( - quantization_setting) # todo: update with the new quant config for weight quant - self.model_profile_enabled = False - self._model_times = [] - - if not self.injection_dict and config.replace_with_kernel_inject: - # This is a hack to remove the prepare_mask function on HF side for BLOOM architecture - self.remove_mask_prepare_for_bloom() - - if self.injection_dict or not config.replace_with_kernel_inject: - # This is a hack to redefine the alibi func due to TP - if config.tensor_parallel.tp_size > 1: - self.build_alibi_tensor() - self.build_attn_bias() - - if get_accelerator().device_name() == 'cuda' and config.enable_cuda_graph: - assert pkg_version.parse(torch.__version__) >= pkg_version.parse("1.10"), \ - "If you want to use cuda graph, please upgrade torch to at least v1.10" - - # convert model to intended dtype - if config.dtype: - self._convert_to_dtype(config) - - if self.mpu: - config.tensor_parallel.tp_size = dist.get_world_size(group=self.mpu.get_model_parallel_group()) - self.mp_group = self.mpu.get_model_parallel_group() - elif config.tensor_parallel.tp_size > 1: - self._create_model_parallel_group(config) - config.tensor_parallel.tp_group = self.mp_group - - if isinstance(self.module, torch.nn.Module): - moe, _ = has_moe_layers(self.module) - else: - moe = False - - if moe and dist.get_world_size() > 1: - self._create_ep_parallel_group(config.moe.moe_experts) - - # We only support three modes: 1) user specified policy for tensor-parallelism, 2) kernel injection (replace_with_kernel_inject), and 3) automatic tensor parallelism if tp_size > 1. - if self.injection_dict: - # 1. User specified Tensor Parallelism - assert not config.replace_with_kernel_inject, "Cannot use both user specified injection policy and kernel injection" - for client_module, injection_policy in self.injection_dict.items(): - - assert issubclass(client_module, - torch.nn.Module), f"{client_module} is not a subclass of torch.nn.Module" - - # construct the tuple and pass that instead of a string or dict. - if isinstance(injection_policy, str): - config.injection_policy_tuple = (injection_policy, ) - else: - config.injection_policy_tuple = injection_policy - - layer_names = [name for name, _ in self.module.named_modules()] - for policy in config.injection_policy_tuple: - if not any(name.endswith(policy) for name in layer_names): - raise ValueError(f"Injection policy layer'{policy}' not valid.") - - self._apply_injection_policy(config, client_module) - else: - if config.replace_with_kernel_inject: - # 2. DeepSpeed Kernel Injection - self._apply_injection_policy(config) - elif config.tensor_parallel.tp_size > 1: - # 3. Automatic Tensor Parallelism - parser_dict = AutoTP.tp_parser(model) - print("AutoTP: ", parser_dict) - for client_module, injection_policy in parser_dict: - if isinstance(injection_policy, str): - config.injection_policy_tuple = (injection_policy, ) - else: - config.injection_policy_tuple = injection_policy - self._apply_injection_policy(config, client_module) - - device = get_accelerator().current_device_name() - # NOTE: This check assumes a Hugging Face hierarchy for the device type i.e. module.device.type - is_meta_device = hasattr(self.module, "device") and self.module.device.type == 'meta' - if is_meta_device: - self.module.to_empty(device=device) - elif not config.keep_module_on_host: - self.module.to(device) - - if config.tensor_parallel.tp_size > 1: - _rng_state = get_accelerator().get_rng_state().to(get_accelerator().current_device_name()) - dist.broadcast(_rng_state, 0) - get_accelerator().set_rng_state(_rng_state.cpu()) - - if config.tensor_parallel.tp_size > 1: - assert not config.enable_cuda_graph, "Cuda graph is not supported for model parallelism" - - # Check if local CUDA graphs can be created in replacement modules - self.local_cuda_graph = self._local_cuda_graph_used(self.module) - self._is_compiled = False - - def destroy(self): - DeepSpeedTransformerInference.layer_id = 0 - DeepSpeedSelfAttention.num_layers = 0 - if DeepSpeedTransformerInference.workspace.is_allocated(): - DeepSpeedTransformerInference.workspace.release_workspace() - DeepSpeedTransformerInference.workspace = None - - def profile_model_time(self, use_cuda_events=True): - if not self.model_profile_enabled and not self._config.enable_cuda_graph: - self.module.register_forward_pre_hook(self._pre_forward_hook) - self.module.register_forward_hook(self._post_forward_hook) - self.model_profile_enabled = True - self.use_cuda_events = use_cuda_events - if self.use_cuda_events: - self.timers = SynchronizedWallClockTimer() - - # todo: remove this once all the config dicts are centralized from top level pydantic config - def _get_model_config_generate(self, config): - # this is being passed to replace_transformer_layer(config=self.user_model_config_dict) - self.config = getattr(self.module, 'config', None) if config.config is None else config.config - - def remove_mask_prepare_for_bloom(self): - if hasattr(self.module, 'transformer'): - if hasattr(self.module.transformer, '_prepare_attn_mask'): - self.module.transformer._prepare_attn_mask = lambda attention_mask, *args, **kwargs: attention_mask - - def build_alibi_tensor(self): - if hasattr(self.module, 'transformer'): - if hasattr(self.module.transformer, 'build_alibi_tensor'): - self.module.transformer.build_alibi_tensor = build_bloom_alibi_tensor - if hasattr(self.module.transformer, 'build_mpt_alibi_tensor'): - self.module.transformer.build_mpt_alibi_tensor_orig = self.module.transformer.build_mpt_alibi_tensor - self.module.transformer.__class__.build_mpt_alibi_tensor = build_mpt_alibi_tensor - if hasattr(self.module, 'model'): - if hasattr(self.module.model, 'get_alibi_mask'): - self.module.model.get_alibi_mask_orig = self.module.model.get_alibi_mask - self.module.model.__class__.get_alibi_mask = get_alibi_mask - - def build_attn_bias(self): - if hasattr(self.module, 'transformer'): - if hasattr(self.module.transformer, '_attn_bias'): - self.module.transformer._attn_bias_orig = self.module.transformer._attn_bias - self.module.transformer.__class__._attn_bias = build_mpt_atten_bias_tensor - - def _pre_forward_hook(self, module, *inputs, **kwargs): - if self.use_cuda_events: - self.timers(INFERENCE_MODEL_TIMER).start() - else: - get_accelerator().synchronize() - self._start = time.time() - - def _post_forward_hook(self, module, input, output): - if self.use_cuda_events: - self.timers(INFERENCE_MODEL_TIMER).stop() - elapsed_time = self.timers(INFERENCE_MODEL_TIMER).elapsed(reset=True) - else: - get_accelerator().synchronize() - self._end = time.time() - elapsed_time = (self._end - self._start) * 1e3 # convert seconds to ms - self._model_times.append(elapsed_time) - - def _create_model_parallel_group(self, config): - # Call the init process - if InferenceEngine.inference_mp_group is None: - init_distributed() - local_rank = int(os.getenv('LOCAL_RANK', '0')) - get_accelerator().set_device(local_rank) - - ranks = [i for i in range(config.tensor_parallel.tp_size)] - self.mp_group = dist.new_group(ranks) - InferenceEngine.inference_mp_group = self.mp_group - else: - self.mp_group = InferenceEngine.inference_mp_group - - def _create_ep_parallel_group(self, moe_experts): - # Call the init process - self.ep_group = {} - self.expert_mp_group = {} - moe_experts = moe_experts if type(moe_experts) is list else [moe_experts] - for e in moe_experts: - self.ep_group.update({e: None}) - self.expert_mp_group.update({e: None}) - for moe_ep_size in self.ep_group.keys(): - num_ep_groups = dist.get_world_size() // moe_ep_size - for i in range(num_ep_groups): - ep_cnt = i * moe_ep_size - size = dist.get_world_size() if moe_ep_size > dist.get_world_size() else moe_ep_size - ranks = list(range(ep_cnt, ep_cnt + size)) - _ep_group = dist.new_group(ranks) - if dist.get_rank() in ranks: - self.ep_group.update({moe_ep_size: _ep_group}) - - if dist.get_world_size() > moe_ep_size: - num_expert_mp_groups = dist.get_world_size() // num_ep_groups - expert_mp_size = dist.get_world_size() // moe_ep_size - for i in range(num_expert_mp_groups): - expert_mp_comm_ranks = [i + nr * moe_ep_size for nr in range(expert_mp_size)] - _expert_mp_group = dist.new_group(expert_mp_comm_ranks) - if dist.get_rank() in expert_mp_comm_ranks: - self.expert_mp_group.update({moe_ep_size: _expert_mp_group}) - - def _init_quantization_setting(self, quantization_setting): - self.quantize_bits = 8 - self.mlp_extra_grouping = False - self.quantize_groups = 1 - if type(quantization_setting) is tuple: - self.mlp_extra_grouping, \ - self.quantize_groups = quantization_setting - elif quantization_setting is not None: - self.quantize_groups = quantization_setting - log_dist( - f"quantize_bits = {self.quantize_bits} " - f"mlp_extra_grouping = {self.mlp_extra_grouping}, " - f"quantize_groups = {self.quantize_groups}", [0]) - - def load_model_with_checkpoint(self, r_module): - self.mp_replace = ReplaceWithTensorSlicing( - mp_group=self.mp_group, mp_size=self._config.tensor_parallel.tp_size) #, out_dim=0, in_dim=1) - error_msgs = [] - - def load(module, state_dict, prefix): - args = (state_dict, prefix, {}, True, [], [], error_msgs) - if hasattr(module, 'weight'): - if module.weight.data.is_meta: - # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here - module.weight = torch.nn.parameter.Parameter(data=torch.empty_like(module.weight.data, - device="cpu"), - requires_grad=module.weight.data.requires_grad) - if 'query_key_value' in prefix: - module.weight = self.mp_replace.strided_copy(module.weight.data, - state_dict[prefix + 'weight'], - num_splits=3) - else: - module.weight = self.mp_replace.copy(module.weight.data, state_dict[prefix + 'weight']) - else: - if module.norm.weight.data.is_meta: - # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here - module.norm.weight = torch.nn.parameter.Parameter( - data=torch.empty_like(module.norm.weight.data, device="cpu"), - requires_grad=module.norm.weight.data.requires_grad) - module.norm.weight = self.mp_replace.copy(module.norm.weight.data, state_dict[prefix + 'weight']) - if prefix + 'bias' in self.key_list: - if hasattr(module, 'norm'): - if module.norm.bias.data.is_meta: - # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here - module.norm.bias = torch.nn.parameter.Parameter( - data=torch.empty_like(module.norm.bias.data, device="cpu"), - requires_grad=module.norm.bias.data.requires_grad) - module.norm.bias = self.mp_replace.copy(module.norm.bias, state_dict[prefix + 'bias']) - else: - if module.bias.data.is_meta: - # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here - module.bias = torch.nn.parameter.Parameter(data=torch.empty_like(module.bias.data, - device="cpu"), - requires_grad=module.bias.data.requires_grad) - data = state_dict[prefix + 'bias'] - data = data.to(get_accelerator().current_device_name()) - module.bias = self.mp_replace.copy(module.bias, data) - - layer_policies = { - nn.Linear: load, - nn.Embedding: load, - nn.LayerNorm: load, - LinearLayer: load, - LinearAllreduce: load - } - - def load_module_recursive(module, prefix='', level=0): - for name, child in module.named_children(): - if child.__class__ in layer_policies: - checking_key = prefix + name + '.' - if not any(checking_key in item for item in self.key_list): - continue - if len(list(child.parameters())) > 0 and list(child.parameters())[0].numel() == 0: - if len(child.weight.ds_shape) == 1: - child = Normalize(dim=child.weight.ds_shape[-1], dtype=child.weight.dtype, eps=child.eps) - setattr(module, name, child) - load(child, self.sd, prefix + name + '.') - else: - load_module_recursive(child, prefix if level == 0 else prefix + name + '.', level + 1) - - load_module_recursive(r_module) - - embedding_weight = None - - for n, p in r_module.named_parameters(): - if "word_embeddings." in n or "embed_tokens." in n or "wte." in n: - embedding_weight = p - if embedding_weight is not None and hasattr(r_module, "lm_head") and hasattr( - r_module.lm_head, "weight") and r_module.lm_head.weight.is_meta: - r_module.lm_head.weight = embedding_weight - - def _apply_injection_policy(self, config, client_module=None): - # client_module is only passed when using the injection_dict method. - checkpoint_dir = config.checkpoint - checkpoint = SDLoaderFactory.get_sd_loader_json(checkpoint_dir, - self.checkpoint_engine) if checkpoint_dir is not None else None - - generic_injection(self.module, dtype=config.dtype, enable_cuda_graph=config.enable_cuda_graph) - - if isinstance(self.module, torch.nn.Module): - # config is our DeepSpeedInferenceConfig and self.config is the HF model config - replace_transformer_layer(client_module, self.module, checkpoint, config, self.config) - - def _get_all_ckpt_names(self, checkpoints_path, tag): - ckpt_file_pattern = self._get_ckpt_name(checkpoints_path, tag, mp_placeholder="*") - import glob - - ckpt_files = glob.glob(ckpt_file_pattern) - ckpt_files.sort() - return ckpt_files - - def _get_ckpt_name(self, checkpoints_path, tag, mp_placeholder=None): - if mp_placeholder is not None: - mp_rank_str = mp_placeholder - else: - mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() - mp_rank_str = "{:02d}".format(mp_rank) - - ckpt_name = os.path.join( - checkpoints_path, - "mp_rank_" + mp_rank_str + "_model_states.pt", - ) - return ckpt_name - - def _load_checkpoint(self, load_dir, load_module_strict=True, tag=None): - is_pipe_parallel = isinstance(self.module, PipelineModule) - if is_pipe_parallel: - raise RuntimeError('pipeline parallelism is currently not supported in inference.') - if not isinstance(load_dir, dict) and os.path.isdir(load_dir): - if tag is None: - latest_path = os.path.join(load_dir, "latest") - if os.path.isfile(latest_path): - with open(latest_path, "r") as fd: - tag = fd.read().strip() - - ckpt_list = self._get_all_ckpt_names(load_dir, tag) - sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, self.checkpoint_engine) - else: - sd_loader = SDLoaderFactory.get_sd_loader_json(load_dir, self.checkpoint_engine) - - checkpoint = sd_loader['checkpoints'] - - if type(checkpoint) is list: - self.sd = torch.load(checkpoint[0], map_location='cpu', weights_only=False) - self.key_list = list(self.sd.keys()) - - self.load_model_with_checkpoint(self.module) - - for i in range(1, len(checkpoint)): - if not dist.is_initialized() or dist.get_rank() == 0: - print(f"loading checkpoint ({i})") - self.sd = torch.load(checkpoint[i], map_location=get_accelerator().device_name(), weights_only=False) - self.key_list = list(self.sd.keys()) - self.load_model_with_checkpoint(self.module) - else: - mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() - - load_path, checkpoint, quantize_config = sd_loader.load(self._config.tensor_parallel.tp_size, - mp_rank, - is_pipe_parallel=is_pipe_parallel, - quantize=(self._config.dtype is torch.int8), - quantize_groups=self.quantize_groups, - mlp_extra_grouping=self.mlp_extra_grouping) - - self.quantization_scales, self.quantize_merge_count = quantize_config - - moe, _ = has_moe_layers(self.module) - if moe: - from deepspeed.runtime.engine import DeepSpeedEngine - old_moe_load = False - if not isinstance(checkpoint['num_experts'], list): - old_moe_load = True - DeepSpeedEngine.load_moe_state_dict(load_dir, - tag, - state_dict=checkpoint[self._choose_module_key(checkpoint)], - old_moe_load=old_moe_load, - model=self.module, - mpu=self.mpu, - checkpoint_engine=self.checkpoint_engine) - - self.module.load_state_dict(state_dict=checkpoint[self._choose_module_key(checkpoint)], - strict=load_module_strict) - - def _choose_module_key(self, sd): - assert not ('module' in sd - and 'model' in sd), "checkpoint has both 'model' and 'module' keys, not sure how to proceed" - assert 'module' in sd or 'model' in sd, "checkpoint contains neither 'model' or 'module' keys, not sure how to proceed" - if 'module' in sd: - return 'module' - elif 'model' in sd: - return 'model' - - def _convert_to_dtype(self, config): - if not isinstance(self.module, torch.nn.Module): - return - - if False: #config.dtype is torch.int8 and self.quantization_scales is None: - quantizer = WeightQuantization(mlp_extra_grouping=self.mlp_extra_grouping) - model, self.quantization_scales = quantizer.model_quantize(self.module, self.injection_dict, - self.quantize_bits, self.quantize_groups) - elif config.dtype == torch.half: - self.module.half() - elif config.dtype == torch.bfloat16: - self.module.bfloat16() - elif config.dtype == torch.float: - self.module.float() - - def _create_cuda_graph(self, *inputs, **kwargs): - # warmup to create the workspace and cublas handle - cuda_stream = get_accelerator().Stream() - cuda_stream.wait_stream(get_accelerator().current_stream()) - with get_accelerator().stream(cuda_stream): - for i in range(3): - ret = self.module(*inputs, **kwargs) - get_accelerator().current_stream().wait_stream(cuda_stream) - - # create cuda_graph and assign static_inputs and static_outputs - self._cuda_graphs = get_accelerator().create_graph() - self.static_inputs = inputs - self.static_kwargs = kwargs - - with get_accelerator().capture_to_graph(self._cuda_graphs): - self.static_output = self.module(*self.static_inputs, **self.static_kwargs) - - self.cuda_graph_created = True - - def _graph_replay(self, *inputs, **kwargs): - for i in range(len(inputs)): - if torch.is_tensor(inputs[i]): - self.static_inputs[i].copy_(inputs[i]) - for k in kwargs: - if torch.is_tensor(kwargs[k]): - self.static_kwargs[k].copy_(kwargs[k]) - get_accelerator().replay_graph(self._cuda_graphs) - return self.static_output - - def model_times(self): - assert self.model_profile_enabled, "model profiling is not enabled" - model_times = self._model_times - if self._config.enable_cuda_graph and len(self._model_times) == 0: - raise ValueError("Model times are empty and cuda graph is enabled. If " - "this is a GPT-style model this combo is not supported. If this is a " - "BERT-style model this is a bug, please report it. " - f"Model type is: {type(self.module)}") - self._model_times = [] - return model_times - - def _module_match(self, module): - for policy in generic_policies: - policy = policy() - if policy.match_replaced(module): - return True - return False - - def _local_cuda_graph_used(self, module): - if isinstance(module, torch.nn.Module): - return False - else: - sub_module_cuda_graph = False - for name in module.__dict__.keys(): - sub_module = getattr(module, name) - - if self._module_match(sub_module) and hasattr(sub_module, "enable_cuda_graph"): - sub_module_cuda_graph = True - - return sub_module_cuda_graph - - def forward(self, *inputs, **kwargs): - """Execute forward propagation - - Arguments: - *inputs: Variable length input list - **kwargs: variable length keyword arguments - """ - start = None - if self.model_profile_enabled and get_accelerator().device_name() == 'cuda' and self._config.enable_cuda_graph: - get_accelerator().synchronize() - start = time.time() - - if get_accelerator().device_name() == 'cuda' and self._config.enable_cuda_graph and not self.local_cuda_graph: - if self.cuda_graph_created: - outputs = self._graph_replay(*inputs, **kwargs) - else: - self._create_cuda_graph(*inputs, **kwargs) - outputs = self._graph_replay(*inputs, **kwargs) - - else: - outputs = self.module(*inputs, **kwargs) - - if self.model_profile_enabled and self._config.enable_cuda_graph: - get_accelerator().synchronize() - duration = (time.time() - start) * 1e3 # convert seconds to ms - self._model_times.append(duration) - - return outputs - - def _generate(self, *inputs, **kwargs): - # Reset KV-cache at the beginning of generate - if hasattr(self.module, 'reset_cache'): - self.module.reset_cache() - num_beams = 1 - if "generation_config" in kwargs: - gen_config = kwargs["generation_config"] - num_beams = getattr(gen_config, "num_beams", 1) - if "num_beams" in kwargs: - num_beams = kwargs["num_beams"] - - if num_beams > 1: - raise NotImplementedError("DeepSpeed does not support `num_beams` > 1, if this is important to you please " - "add your request to: https://github.com/deepspeedai/DeepSpeed/issues/2506") - - if ("input_ids" in kwargs) and (kwargs["input_ids"].dim() == 2): - for input_tensor in kwargs["input_ids"]: - tensor_length = input_tensor.shape[-1] - if tensor_length > self._config.max_out_tokens: - raise RuntimeError( - f"Input with size {tensor_length} exceeds maximum length of {self._config.max_out_tokens}. Please increase max_tokens in the DeepSpeed Inference Config." - ) - - return self.module.generate(*inputs, **kwargs) - - def compile(self, backend=get_accelerator().get_compile_backend(), compile_kwargs={}) -> None: - """ - Compile the module using the specified backend and kwargs. - """ - if not is_compile_supported(): - raise RuntimeError("compile is not supported in your version of PyTorch.") - - if self._is_compiled: - return - - # Avoid graph breaks - deepspeed.utils.nvtx.enable_nvtx = False - self.module.compile(backend=backend, **compile_kwargs) - self._is_compiled = True - - @property - def is_compiled(self) -> bool: - return self._is_compiled diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__init__.py deleted file mode 100644 index ac8a42da8ab3d9e4ace2a4f1d7b1d455cf7be7fb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -from .config_v2 import RaggedInferenceEngineConfig, DeepSpeedTPConfig -from .engine_v2 import InferenceEngineV2 -from .engine_factory import build_hf_engine, build_engine_from_ds_checkpoint diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 8f384da9a0d934d099e2bcd085666981d32c8554..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/allocator.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/allocator.cpython-310.pyc deleted file mode 100644 index 21766683f40705639deb3fd2bc81af0dcddc5378..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/allocator.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/config_v2.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/config_v2.cpython-310.pyc deleted file mode 100644 index d8f330c1d278086d2afcb3380435c807163fae17..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/config_v2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/engine_factory.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/engine_factory.cpython-310.pyc deleted file mode 100644 index b1cb95d907e9166a409f174322108a09c94945fa..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/engine_factory.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/engine_v2.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/engine_v2.cpython-310.pyc deleted file mode 100644 index c565b72a9519e8e05e50d8c8a7da84c4f88348d5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/engine_v2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/inference_parameter.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/inference_parameter.cpython-310.pyc deleted file mode 100644 index 303d3c1181a80e0bc221d0ad8ac273a8859b9062..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/inference_parameter.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/inference_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/inference_utils.cpython-310.pyc deleted file mode 100644 index a3282a8b0e65d831e7b26636207b20a3a5b651c8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/inference_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/logging.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/logging.cpython-310.pyc deleted file mode 100644 index a5342ceda7461eb166d971b011254743ad60ae1d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/logging.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/scheduling_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/scheduling_utils.cpython-310.pyc deleted file mode 100644 index dc8c7fcf80ed0d291e92ec30366503867aa4e77f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/__pycache__/scheduling_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/allocator.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/allocator.py deleted file mode 100644 index fcc0d94c0f825170dd89c54db73e53a2baf0c077..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/allocator.py +++ /dev/null @@ -1,42 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from functools import reduce -from typing import Iterable -from collections import defaultdict -import torch - -from deepspeed.accelerator import get_accelerator - - -class Allocator: - cache = defaultdict(dict) - - def empty_from(tensor: torch.Tensor, shape: Iterable[int]) -> torch.Tensor: - try: - return Allocator.cache[tensor][shape] - except KeyError: - shape_size = reduce(lambda x, y: x * y, shape) - if shape_size == 0: - raise ValueError("Cannot create empty tensor with size 0") - Allocator.cache[tensor][shape] = tensor.flatten()[:shape_size].view(shape) - return Allocator.cache[tensor][shape] - - -empty_from = Allocator.empty_from - - -def on_device(method) -> torch.Tensor: - """ - Wraps a method to ensure the returned tensor is on the current device. - """ - - def wrapped(self, *args, **kwargs): - tensor = method(self, *args, **kwargs) - if isinstance(tensor, torch.Tensor): - return tensor.to(get_accelerator().current_device()) - return tensor - - return wrapped diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/config_v2.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/config_v2.py deleted file mode 100644 index 325b57d8f56a8730158096a2f8109a15b662ecf7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/config_v2.py +++ /dev/null @@ -1,44 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from pydantic import Field -from typing import Optional - -from deepspeed.runtime.config_utils import DeepSpeedConfigModel -from .ragged import DSStateManagerConfig - - -class DeepSpeedTPConfig(DeepSpeedConfigModel): - """ Configure tensor parallelism settings """ - - tp_size: int = 1 - """ Number of devices to split the model across using tensor parallelism. """ - - -class QuantizationConfig(DeepSpeedConfigModel): - """ Configure tensor parallelism settings """ - - quantization_mode: Optional[str] = None - """ The quantization mode in string format. The supported modes are as follows: - - 'wf6af16', weight-only quantization with FP6 weight and FP16 activation. - """ - # TODO: may reuse the constants in deepspeed/compression/constants.py - - -class RaggedInferenceEngineConfig(DeepSpeedConfigModel): - """ Sets parameters for DeepSpeed Inference Engine. """ - - tensor_parallel: DeepSpeedTPConfig = Field({}, alias="tp") - """ - Configuration for tensor parallelism used to split the model across several - GPUs. Expects a dictionary containing values for :any:`DeepSpeedTPConfig`. - """ - - state_manager: DSStateManagerConfig = Field({}, alias="manager") - """ - Configuration for managing persistent state - """ - - quantization: QuantizationConfig = {} diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/engine_factory.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/engine_factory.py deleted file mode 100644 index 9c3188dfebb86e08aea555858b3873c14a9438aa..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/engine_factory.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import json -import logging -import os -import pickle -from packaging import version - -from .engine_v2 import InferenceEngineV2 -from .config_v2 import RaggedInferenceEngineConfig -from .checkpoint import HuggingFaceCheckpointEngine -from .logging import inference_logger -from .model_implementations import ( - OPTPolicy, - Llama2Policy, - MistralPolicy, - MixtralPolicy, - FalconPolicy, - PhiPolicy, - Phi3Policy, - QwenPolicy, - Qwen2Policy, - Qwen2MoePolicy, -) -from .model_implementations.inference_policy_base import POLICIES, InferenceV2Policy -from .model_implementations.flat_model_helpers import make_metadata_filename, ModelMetadata - - -def build_engine_from_ds_checkpoint(path: str, - engine_config: RaggedInferenceEngineConfig, - debug_level: int = logging.INFO) -> InferenceEngineV2: - """ - Creates an engine from a checkpoint saved by ``InferenceEngineV2``. - - Arguments: - path: Path to the checkpoint. This does not need to point to any files in particular, - just the directory containing the checkpoint. - engine_config: Engine configuration. See ``RaggedInferenceEngineConfig`` for details. - debug_level: Logging level to use. Unless you are actively seeing issues, the recommended - value is ``logging.INFO``. - - Returns: - Fully initialized inference engine ready to serve queries. - """ - - inference_logger(level=debug_level) - # Load metadata, for grabbing the policy name we'll have all ranks just check for - # rank 0. - metadata_filename = make_metadata_filename(path, 0, engine_config.tensor_parallel.tp_size) - metadata = json.load(open(metadata_filename, "r")) - metadata = ModelMetadata.parse_raw(metadata) - - # Get the policy - try: - policy_cls: InferenceV2Policy = POLICIES[metadata.policy] - except KeyError: - raise ValueError(f"Unknown policy {metadata.policy} for model {path}") - - # Load the model config - model_config = pickle.load(open(os.path.join(path, "ds_model_config.pkl"), "rb")) - policy = policy_cls(model_config, inf_checkpoint_path=path) - - return InferenceEngineV2(policy, engine_config) - - -def build_hf_engine(path: str, - engine_config: RaggedInferenceEngineConfig, - debug_level: int = logging.INFO) -> InferenceEngineV2: - """ - Build an InferenceV2 engine for HuggingFace models. This can accept both a HuggingFace - model name or a path to an Inference-V2 checkpoint. - - Arguments: - path: Path to the checkpoint. This does not need to point to any files in particular, - just the directory containing the checkpoint. - engine_config: Engine configuration. See ``RaggedInferenceEngineConfig`` for details. - debug_level: Logging level to use. Unless you are actively seeing issues, the recommended - value is ``logging.INFO``. - - Returns: - Fully initialized inference engine ready to serve queries. - """ - - if os.path.exists(os.path.join(path, "ds_model_config.pkl")): - return build_engine_from_ds_checkpoint(path, engine_config, debug_level=debug_level) - else: - # Set up logging - inference_logger(level=debug_level) - # get HF checkpoint engine - checkpoint_engine = HuggingFaceCheckpointEngine(path) - - # get model config from HF AutoConfig - model_config = checkpoint_engine.model_config - - # get the policy - # TODO: generalize this to other models - if model_config.model_type == "opt": - if not model_config.do_layer_norm_before: - raise ValueError( - "Detected OPT-350m model. This model is not currently supported. If this is not the 350m model, please open an issue: https://github.com/deepspeedai/DeepSpeed-MII/issues" - ) - policy = OPTPolicy(model_config, checkpoint_engine=checkpoint_engine) - elif model_config.model_type == "llama": - policy = Llama2Policy(model_config, checkpoint_engine=checkpoint_engine) - elif model_config.model_type == "mistral": - # Ensure we're using the correct version of transformers for mistral - import transformers - assert version.parse(transformers.__version__) >= version.parse("4.34.0"), \ - f"Mistral requires transformers >= 4.34.0, you have version {transformers.__version__}" - policy = MistralPolicy(model_config, checkpoint_engine=checkpoint_engine) - elif model_config.model_type == "mixtral": - # Ensure we're using the correct version of transformers for mistral - import transformers - assert version.parse(transformers.__version__) >= version.parse("4.36.1"), \ - f"Mistral requires transformers >= 4.36.1, you have version {transformers.__version__}" - policy = MixtralPolicy(model_config, checkpoint_engine=checkpoint_engine) - elif model_config.model_type == "falcon": - policy = FalconPolicy(model_config, checkpoint_engine=checkpoint_engine) - elif model_config.model_type == "phi": - policy = PhiPolicy(model_config, checkpoint_engine=checkpoint_engine) - elif model_config.model_type == "phi3": - policy = Phi3Policy(model_config, checkpoint_engine=checkpoint_engine) - elif model_config.model_type == "qwen": - policy = QwenPolicy(model_config, checkpoint_engine=checkpoint_engine) - elif model_config.model_type == "qwen2": - policy = Qwen2Policy(model_config, checkpoint_engine=checkpoint_engine) - elif model_config.model_type == "qwen2_moe": - policy = Qwen2MoePolicy(model_config, checkpoint_engine=checkpoint_engine) - else: - raise ValueError(f"Unsupported model type {model_config.model_type}") - - return InferenceEngineV2(policy, engine_config) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/engine_v2.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/engine_v2.py deleted file mode 100644 index 4a358310377f00af00d92b6e5dcd2fd5d600392d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/engine_v2.py +++ /dev/null @@ -1,268 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import json -import pickle -from typing import Iterable, Tuple - -import torch - -import deepspeed.comm as dist - -from deepspeed.accelerator import get_accelerator -from deepspeed.comm.comm import init_distributed - -from .model_implementations import InferenceV2Policy -from .logging import inference_logger -from .ragged import DSStateManager, RaggedBatchWrapper, PlaceholderSequenceDescriptor -from .scheduling_utils import SchedulingError, SchedulingResult -from .model_implementations.flat_model_helpers import make_param_filename, make_metadata_filename -from .model_implementations.inference_model_base import DSInferenceModelBase - -from .config_v2 import RaggedInferenceEngineConfig - -INFERENCE_MODEL_TIMER = "model-forward-inference" - - -class InferenceEngineV2: - - _config: RaggedInferenceEngineConfig - """ - Configuration of the inference engine. - """ - - _model: DSInferenceModelBase - """ - Inference model supporting ragged inference. - """ - - _state_manager: DSStateManager - """ - Persistent state manager for sequences and KV-cache. - """ - - @property - def free_blocks(self) -> torch.Tensor: - """ - Number of free KV blocks. This is a tensor of shape [n_kv_cache_groups] where each - element is the number of free blocks in the corresponding KV cache group. - """ - return self._state_manager.free_blocks - - @property - def n_kv_cache_groups(self) -> int: - """ - Number of KV cache groups. - """ - return self._state_manager.n_kv_cache_groups - - def model(self) -> DSInferenceModelBase: - """ - The model implementation. - """ - return self._model - - def __init__(self, policy: InferenceV2Policy, engine_config: RaggedInferenceEngineConfig) -> None: - """ - Create the Inference V2 engine. - - Arguments: - policy (InferenceV2Policy): Policy for the model implementation. This policy object - will be used to build the model and load the checkpoint associated with it. - engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine. - """ - self._config = engine_config - self._policy = policy - self._base_mp_group = self._initialize_tp_group() - - # Build model from policy - inference_logger().info("Building model...") - self._model = self._policy.build_model(self._config, self._base_mp_group) - inference_logger().info("Model built.") - - # Create state manager - self._batch = RaggedBatchWrapper(self._config.state_manager) - self._state_manager = DSStateManager(self._config.state_manager, - self._model.kv_cache_config(), - base_mp_group=self._base_mp_group) - self._model.set_state_manager(self._state_manager) - - def _initialize_tp_group(self): - """ - Implementation of our TP group initialization. - """ - init_distributed() - local_rank = int(os.getenv("LOCAL_RANK", 0)) - get_accelerator().set_device(local_rank) - - if local_rank >= self._config.tensor_parallel.tp_size: - raise RuntimeError("Local rank is greater than TP size, ensure that the TP config is correct.") - - ranks = list(range(self._config.tensor_parallel.tp_size)) - return dist.new_group(ranks=ranks) - - def put(self, - batch_uids: Iterable[int], - batch_tokens: Iterable[torch.Tensor], - do_checks: bool = True) -> torch.Tensor: - """ - Put a ragged batch onto the inference engine. This will perform one forward and return - a Tensor of the shape [len(batch_uids), *output_shape]. Logits for the non-final tokens - are not calculated. - - Arguments: - batch_uids: Iterable of uids for the batch on the host - batch_tokens: Iterable of token tensors for the batch on the host - do_checks: Check schedulability when it is set to True. You can skip this check for better performance when it has already been completed. - """ - - if do_checks: - token_lens = [len(tokens) for tokens in batch_tokens] - schedule_check = self.can_schedule(batch_uids, token_lens) - if schedule_check != SchedulingResult.Success: - raise SchedulingError(schedule_check) - - self._batch.clear() - for uid, tokens in zip(batch_uids, batch_tokens): - - host_seq_desc = self._state_manager.get_or_create_sequence(uid) - self._model.maybe_allocate_kv(host_seq_desc, tokens.numel()) - host_seq_desc.pre_forward(tokens.numel()) - - # We can disable checks since we already validated schedulability. - self._batch.insert_sequence(host_seq_desc, tokens, do_checks=do_checks) - - # Send all metadata to the device - self._batch.finalize() - - # Prep all data structures for the actual forward (in anticipation of CG in the future) - # and also to amortize some of the costs in a more straightforward way. - self._model.prepare_batch(self._batch) - - # Model implementation will pick up in the forward. - logits = self._model.forward(self._batch) - - # We return one set of logits per sequence in the batch (saves cost on unembedding) - assert logits.shape[0] == self._batch.current_sequences - - for uid in batch_uids: - host_seq_desc = self._state_manager.get_sequence(uid) - host_seq_desc.post_forward() # Updates sequence metadata. - self._model.maybe_free_kv(host_seq_desc) - - return logits - - def query(self, uid: int, max_request_tokens: int, max_request_blocks) -> Tuple[int, torch.Tensor]: - """ - Determine the number of tokens and KV blocks to reserve for a given request. Given a UID - (this UID may not be recognized by the model yet), this will return the number of tokens - and blocks to reserve for the request. - - Arguments: - uid (int): The UID of the sequence (as tracked by the scheduling entity). If - this is a new sequence (with a UID unknown to the inference engine), then - an empty placeholder is created to pass to the occupancy logic. - n_tokens (int): The number of tokens to hypothetically send. - - Returns: - Tuple[int, Optional[int]]: Tuple of free kv blocks and the number of blocks - required to schedule the sequence. - """ - seq_desc = self._state_manager.get_sequence(uid) - if seq_desc is None: - if (self._state_manager.n_tracked_sequences == self._config.state_manager.max_tracked_sequences): - return (0, 0) - seq_desc = PlaceholderSequenceDescriptor() - - req_tokens, req_blocks = self._model.get_kv_requirements(seq_desc, max_request_tokens, max_request_blocks) - - return (req_tokens, req_blocks) - - def can_schedule(self, uids: Iterable[int], lengths: Iterable[int]) -> SchedulingResult: - """ - Dry run a batch to determine if it can be scheduled. Placeholder sequences will be - created for any UIDs that are unknown to the inference engine. - - Arguments: - uids (Iterable[int]): Iterable of UIDs for the batch - lengths (Iterable[int]): Iterable of lengths for each sequence of the batch. This lengths - corresponds to the number of tokens to send in the hypothetical forward; history - tokens will be determined via UID lookup and future tokens are disregarded. - - Returns: - bool: True if the batch can be scheduled, False otherwise. - """ - - cur_seqs = self._state_manager.n_tracked_sequences - free_blocks = self._state_manager.free_blocks - req_blocks = 0 - batch_len = 0 - - if len(uids) > self._config.state_manager.max_ragged_sequence_count: - # Can only compose a batch from a limited number of sequences - return SchedulingResult.BatchSequenceLimitExceeded - - for uid, length in zip(uids, lengths): - seq_desc = self._state_manager.get_sequence(uid) - if seq_desc is None: - cur_seqs += 1 - seq_desc = PlaceholderSequenceDescriptor() - - sched_len, sched_blocks = self._model.get_kv_requirements(seq_desc, length, free_blocks) - - if sched_len != length: - # We ran out of KV cache - return SchedulingResult.KVCacheLimitExceeded - - batch_len += length - free_blocks -= sched_blocks - - if cur_seqs > self._config.state_manager.max_tracked_sequences: - # Would run out of tracking metadata - return SchedulingResult.EngineSequenceLimitExceeded - - if batch_len > self._config.state_manager.max_ragged_batch_size: - # Would exceed the maximum batch size - return SchedulingResult.BatchTokenLimitExceeded - - return SchedulingResult.Success - - def get_remaining_block_capacity(self, uid: int) -> int: - """ - Get the remaining capacity of the last block already allocated. - """ - seq_desc = self._state_manager.get_sequence(uid) - if seq_desc is None: - return 0 - return self._model.get_remaining_block_capacity(seq_desc) - - def flush(self, uid: int) -> None: - """ - Remove all state associated with a sequence from the inference engine. - - Arguments: - uid (int): The UID of the sequence to flush. - """ - self._state_manager.flush_sequence(uid) - - def serialize(self, save_path: str) -> None: - """ - Serialize the model to a file. - - Arguments: - path (str): Path to the file to serialize to. - """ - param_file_name = make_param_filename(save_path, self._model.tp_rank, self._model.tp_size) - metadata_file_name = make_metadata_filename(save_path, self._model.tp_rank, self._model.tp_size) - - # Save the flattened parameters - - torch.save(self._model.flattened_params, param_file_name) - - json.dump(self._model.flattened_param_metadata.json(), open(metadata_file_name, "w")) - - if self._model.tp_rank == 0: - pickle.dump(self._model._config, open(os.path.join(save_path, "ds_model_config.pkl"), "wb")) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/inference_parameter.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/inference_parameter.py deleted file mode 100644 index 4dcff16a4515ce37ed334c9b0b0d623eea5b2ac2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/inference_parameter.py +++ /dev/null @@ -1,89 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Dict - -import torch - -CORE_PARAM = "_ds_core_param_key" - -STR_TO_DTYPE = { - "torch.float32": torch.float32, - "torch.float64": torch.float64, - "torch.float16": torch.float16, - "torch.bfloat16": torch.bfloat16, - "torch.int64": torch.int64, - "torch.int32": torch.int32, - "torch.int16": torch.int16, - "torch.int8": torch.int8, - "torch.uint8": torch.uint8, - "torch.bool": torch.bool, -} - - -class InferenceParameter(torch.Tensor): - """ - An extension of the torch.Tensor class to support our inference focused features. One important - thing to note here is that an InferenceParam can be used a torch.Tensor, but outputs of - torch.Tensor operations will not be InferenceParams. - """ - - @staticmethod - def __new__(cls, tensor, *args, **kwargs): - new_tensor = super().__new__(cls, tensor, *args, **kwargs) - if hasattr(tensor, "_aux_attrs"): - setattr(new_tensor, "_aux_attrs", tensor.aux_attrs) - return new_tensor - - def to(self, *args, **kwargs): - new_tensor = super().to(*args, **kwargs) - if hasattr(self, "_aux_attrs"): - setattr(new_tensor, "_aux_attrs", self.aux_attrs) - try: - _ = torch.device(args[0]) - for name, attr in new_tensor.aux_attrs.items(): - new_attr = attr.to(*args, **kwargs) - setattr(new_tensor, name, new_attr) - new_tensor.aux_attrs[name] = new_attr - except: - pass - - return new_tensor - - @classmethod - def initialize(cls, core_param: torch.Tensor, **kwargs) -> 'InferenceParameter': - """ - Create the inference parameter. - """ - param = InferenceParameter(core_param) - setattr(param, "_aux_attrs", kwargs) - - for attr_name, attr in kwargs.items(): - if hasattr(param, attr_name): - raise ValueError(f"Attribute {attr_name} already exists on param.") - - if not isinstance(attr, torch.Tensor): - raise ValueError(f"Attribute {attr_name} must be a tensor.") - - setattr(param, attr_name, attr) - - return param - - @classmethod - def initialize_raw(self, **kwargs) -> 'InferenceParameter': - """ - All kwargs must be torch.Tensors and must include the core parameter. - """ - if CORE_PARAM not in kwargs: - raise ValueError(f"Must provide core parameter, with key {CORE_PARAM}.") - - return InferenceParameter.initialize(kwargs[CORE_PARAM], **kwargs) - - @property - def aux_attrs(self) -> Dict[str, torch.Tensor]: - """ - Dictionary of auxiliary attributes. - """ - return self._aux_attrs diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/inference_utils.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/inference_utils.py deleted file mode 100644 index 7b2dd4237353d85b4249faafb2d8c3051289b2ef..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/inference_utils.py +++ /dev/null @@ -1,105 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Dict - -import torch - -from enum import Enum, IntEnum - - -class NormTypeEnum(Enum): - LayerNorm: str = "layer_norm" - RMSNorm: str = "rms_norm" - - -class DtypeEnum(Enum): - # The torch dtype must always be the first value (so we return torch.dtype) - fp16 = torch.float16, "torch.float16", "fp16", "float16", "half" - fp32 = torch.float32, "torch.float32", "fp32", "float32", "float" - bf16 = torch.bfloat16, "torch.bfloat16", "bf16", "bfloat16", "bfloat" - int8 = torch.int8, "torch.int8", "int8" - - # Copied from https://stackoverflow.com/a/43210118 - # Allows us to use multiple values for each Enum index and returns first - # listed value when Enum is called - def __new__(cls, *values): - obj = object.__new__(cls) - # first value is canonical value - obj._value_ = values[0] - for other_value in values[1:]: - cls._value2member_map_[other_value] = obj - obj._all_values = values - return obj - - def __repr__(self): - return "<%s.%s: %s>" % ( - self.__class__.__name__, - self._name_, - ", ".join([repr(v) for v in self._all_values]), - ) - - -ELEM_SIZES: Dict[torch.dtype, int] = { - torch.float16: 2, - torch.bfloat16: 2, - torch.float32: 4, - torch.float64: 8, - torch.int8: 1, - torch.uint8: 1, - torch.int16: 2, - torch.int32: 4, - torch.int64: 8, - torch.bool: 1, -} - - -class ActivationType(IntEnum): - """ - Types of activations supported by DS-Inference - """ - - GELU = 0 - - RELU = 1 - - SILU = 2 - - GEGLU = 3 - - ReGLU = 4 - - SiGLU = 5 - - IDENTITY = 6 - - InvalidType = -1 - - -def is_gated(act_fn: ActivationType) -> bool: - """ - Return True if the given activation function is gated. - """ - if not isinstance(act_fn, ActivationType): - act_fn = ActivationType(act_fn) - - return act_fn in [ActivationType.GEGLU, ActivationType.ReGLU, ActivationType.SiGLU] - - -def elem_size(dtype: torch.dtype) -> int: - """ - Return size in bytes of the given dtype. - """ - try: - return ELEM_SIZES[dtype] - except KeyError: - raise ValueError("Unknown dtype size for {}".format(dtype)) - - -def ceil_div(a: int, b: int) -> int: - """ - Return ceil(a / b). - """ - return -(-a // b) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/logging.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/logging.py deleted file mode 100644 index 77afe351cbea127c2f0f5dbc9a249dd71516ca31..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/logging.py +++ /dev/null @@ -1,26 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import logging - -from deepspeed.utils.logging import LoggerFactory - -inf_logger = None - - -def inference_logger(level: int = logging.INFO) -> logging.Logger: - """ - Create the inference logger. NOTE: Logging is not cost free. On a 3960X, - there is a cost of about 6 us per call to a no-op logger, so this should - be used during setup only and not during the inference loop. - - Args: - level (int, optional): The logging level. Defaults to logging.INFO. - """ - global inf_logger - if inf_logger is None: - inf_logger = LoggerFactory.create_logger(name="DS-Inference", level=level) - inf_logger.debug("Inference logger created.") - return inf_logger diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__init__.py deleted file mode 100644 index 3483d9348c55599f8066be65b5d89ef04aade950..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .inference_model_base import DSInferenceModelBase -from .inference_transformer_base import DSTransformerModelBase, DSMoETransformerModelBase -from .inference_policy_base import InferenceV2Policy, ContainerMap -from .sharding import * - -# Model Implementations -from .llama_v2 import * -from .opt import * -from .mistral import * -from .mixtral import * -from .falcon import * -from .phi import * -from .phi3 import * -from .qwen import * -from .qwen_v2 import * -from .qwen_v2_moe import * diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index ae2885b085ee21f87569c5187bfa4e68959dc496..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/flat_model_helpers.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/flat_model_helpers.cpython-310.pyc deleted file mode 100644 index e61167e07003a220236e5b53c59d86dfb5074a54..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/flat_model_helpers.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/inference_model_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/inference_model_base.cpython-310.pyc deleted file mode 100644 index c7730e79da2155c19fb4cd0cec779db682dfa9c4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/inference_model_base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/inference_policy_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/inference_policy_base.cpython-310.pyc deleted file mode 100644 index 8f05ae250c989db8b487446679d1a9865e0a7017..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/inference_policy_base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/inference_transformer_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/inference_transformer_base.cpython-310.pyc deleted file mode 100644 index 690cbb162b6a1aa8407d05a553613fb71eaf8c1e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/inference_transformer_base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/layer_container_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/layer_container_base.cpython-310.pyc deleted file mode 100644 index 8014dcd0bb5aafc029634d968fcd6b6a76716993..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/layer_container_base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/parameter_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/parameter_base.cpython-310.pyc deleted file mode 100644 index 4cf2ff0812aa081d4f4dd238ffad605f7ce2c4a7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/__pycache__/parameter_base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/__init__.py deleted file mode 100644 index 60963011cd660fe5e43b9a90efdacec4b16651b9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .attn_output_parameters import * -from .embedding_parameters import * -from .mlp_parameters import * -from .moe_parameters import * -from .norm_parameters import * -from .qkv_parameters import * -from .unembed_parameters import * -from .invfreq_parameters import * diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/attn_output_parameters.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/attn_output_parameters.py deleted file mode 100644 index f220cf7a7125d030b05829d0615210c79e9d562d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/attn_output_parameters.py +++ /dev/null @@ -1,29 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from ...model_implementations.parameter_base import ParameterBase -""" -Common Attention Output Parameter Patterns -""" - - -class AttentionOutputParameter(ParameterBase): - """ - Attention output parameter container. - - Note: The differentiation for something like GQA for this matrix is primarily - encompassed in the sharding logic, which is currently expected to be performed by - the model implementation. - """ - - params: torch.Tensor - """ - Unsharded attention output parameter of shape [model_dim, model_dim] - """ - - def finalize(self) -> torch.Tensor: - return self.inference_model.transform_attn_out_param(self.params) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/embedding_parameters.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/embedding_parameters.py deleted file mode 100644 index 2ed34b5fd259a77e858073a32deef1f805dd1325..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/embedding_parameters.py +++ /dev/null @@ -1,26 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from ...model_implementations.parameter_base import ParameterBase -""" -Embedding containers. -""" - - -class EmbeddingParameter(ParameterBase): - """ - Embedding container. This should be safe to use for all types of embeddings (i.e. word, position, - and token type). - """ - - params: torch.Tensor - """ - Vocabulary parameter of shape [vocab_size, model_dim]. - """ - - def finalize(self) -> torch.Tensor: - return self.inference_model.transform_embedding_param(self.params) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/invfreq_parameters.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/invfreq_parameters.py deleted file mode 100644 index 163f9de81d98bc30e8fb7bbf3b0b724dd74cd0c3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/invfreq_parameters.py +++ /dev/null @@ -1,19 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from ...model_implementations.parameter_base import ParameterBase -""" -Common InvFreq Parameter Patterns -""" - - -class InvFreqParameter(ParameterBase): - - params: torch.Tensor - - def finalize(self) -> torch.Tensor: - return self.params.to(self.inference_model.activation_dtype.value) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/mlp_parameters.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/mlp_parameters.py deleted file mode 100644 index 17def1fa021fb1ce8c23b9c30ca5c27c40d67488..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/mlp_parameters.py +++ /dev/null @@ -1,99 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from ...model_implementations.parameter_base import ParameterBase -""" -MLP Parameter Containers -""" - - -class MLP1Parameter(ParameterBase): - """ - First MLP projection weight container. This performs a straight pass-through to the - model implementation for transformation. - """ - params: torch.Tensor - - def finalize(self) -> torch.Tensor: - # NOTE(cmikeh2): If we are gated but not in the format specified below, we should trigger a permutation here. - # I am not currently aware of any models that use this format (or how we should even detect it; probably should - # just be a different param entirely, but until then we'll just assume the format is correct). - return self.inference_model.transform_mlp_1_param(self.params) - - -class GatedMLPParameter(ParameterBase): - """ - Gated MLP projection container. - """ - - gate_params: torch.Tensor - """ - Weight parameter for the gating matrix. - """ - - up_params: torch.Tensor - """ - For lack of a better name, the non-gating weight parameters. - """ - - def finalize(self) -> torch.Tensor: - """ - Our gated format (this is different from InferenceV1!) is to have the gate and activated neurons - interleaved. So if we have 4 output neurons (two effective neurons) with 4 input neurons, the finalized - parameter will look like: - [g0_0, g0_1, g0_2, g0_3] - [a0_0, a0_1, a0_2, a0_3] - [g1_0, g1_1, g1_2, g1_3] - [a1_0, a1_1, a1_2, a1_3] - - As a reference, in inference v1, the format is: - [g0_0, g0_1, g0_2, g0_3] - [g1_0, g1_1, g1_2, g1_3] - [a0_0, a0_1, a0_2, a0_3] - [a1_0, a1_1, a1_2, a1_3] - """ - assert self.gate_params.shape[0] == self.up_params.shape[ - 0], "Gated MLP parameters must have the same number of neurons." - total_neurons = self.gate_params.shape[0] + self.up_params.shape[0] - - # flip the order if even with the correct tokenizer we get wrong output - #fused_param = torch.cat([self.up_params, self.gate_params], dim=-1).reshape(total_neurons, -1) - fused_param = torch.cat([self.gate_params, self.up_params], dim=-1).reshape(total_neurons, -1) - return self.inference_model.transform_mlp_1_param(fused_param) - - -class FusedGatedMLPParameter(ParameterBase): - """ - Gated MLP projection container. - """ - - params: torch.Tensor - """ - Weight parameter for the fused gating and non-gating weight parameters. - """ - - def finalize(self) -> torch.Tensor: - gate_params = self.params[:self.params.shape[0] // 2] - up_params = self.params[self.params.shape[0] // 2:] - total_neurons = gate_params.shape[0] + up_params.shape[0] - fused_param = torch.cat([gate_params, up_params], dim=-1).reshape(total_neurons, -1) - return self.inference_model.transform_mlp_1_param(fused_param) - - -class MLP2Parameter(ParameterBase): - """ - Second MLP projection weight container. This performs a straight pass-through to the - model implementation for transformation. - """ - - params: torch.Tensor - """ - Full weight parameter. - """ - - def finalize(self) -> torch.Tensor: - return self.inference_model.transform_mlp_2_param(self.params) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/moe_parameters.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/moe_parameters.py deleted file mode 100644 index 8ababf567ba9a499624c5924dd32564ad82922fb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/moe_parameters.py +++ /dev/null @@ -1,78 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from ...model_implementations.parameter_base import ParameterBase, ParamList -""" -Moe Parameters - -These parameters are compatible with any model inheriting from ``DSMoETransformerModelBase``. -""" - - -class MoEGatingWeightParameter(ParameterBase): - """ - Gating weight matrix. - """ - - params: torch.Tensor - """ - Projection matrix from the input activations to the gate logits. - """ - - def finalize(self) -> torch.Tensor: - return self.inference_model.transform_moe_gate_param(self.params) - - -class UnfusedMoEMLP1Parameter(ParameterBase): - """ - This container should be used when the experts are held in separate parameters - and need to be joined into a single group. - """ - - experts: ParamList("n_experts") # noqa: F821 - - def finalize(self) -> torch.Tensor: - stacked_experts = torch.stack([p for p in self.experts], dim=0) - return self.inference_model.transform_moe_mlp_1_param(stacked_experts) - - -class UnfusedMoEMLP2Parameter(ParameterBase): - """ - This container should be used when the experts are held in separate parameters - and need to be joined into a single group. - """ - - experts: ParamList("n_experts") # noqa: F821 - - def finalize(self) -> torch.Tensor: - stacked_experts = torch.stack([p for p in self.experts], dim=0) - return self.inference_model.transform_moe_mlp_2_param(stacked_experts) - - -class UnfusedMoEGatedMLPParameter(ParameterBase): - """ - MoE Parameter for a gated activation function in which the gating matrix is not - fused in the same parameter as the non-gating matrix. - - This is a stacked version of the ``GatedMLPParameter``. Please see that class for more - documentation on the layout of the parameters. - """ - - gating_experts: ParamList("n_experts") # noqa: F821 - - up_experts: ParamList("n_experts") # noqa: F821 - - def finalize(self) -> torch.Tensor: - transposed_experts = [] - for gate, up in zip(self.gating_experts, self.up_experts): - assert gate.shape[0] == up.shape[0], "Gated MLP parameters must have the same number of neurons." - total_neurons = gate.shape[0] + up.shape[0] - fused_expert = torch.cat([gate, up], dim=-1).reshape(total_neurons, -1) - transposed_experts.append(fused_expert) - - stacked_experts = torch.stack(transposed_experts, dim=0) - return self.inference_model.transform_moe_mlp_1_param(stacked_experts) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/norm_parameters.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/norm_parameters.py deleted file mode 100644 index 81ffcc3221df2dd2c05f9d8739a905a5ea2399a5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/norm_parameters.py +++ /dev/null @@ -1,22 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from ...model_implementations.parameter_base import ParameterBase -""" -Common Attention Output Parameter Patterns -""" - - -class NormParameter(ParameterBase): - """ - Simple normalization container. - """ - - params: torch.Tensor - - def finalize(self) -> torch.Tensor: - return self.inference_model.transform_norm_param(self.params) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/qkv_parameters.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/qkv_parameters.py deleted file mode 100644 index e240137186fe6114cb58ea66c1a25da59184d8f7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/qkv_parameters.py +++ /dev/null @@ -1,115 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from ...model_implementations.parameter_base import ParameterBase -""" -Common QKV Parameter Patterns -""" - - -class FusedQKVParameter(ParameterBase): - """ - Traditional fused QKV parameters for QKV projection. This is functionally - a direct copy. - - src_qkv_w shape: [3 * out_features, in_features] - qkv_w shape: [3 * out_features, in_features] - """ - - params: torch.Tensor - - def finalize(self) -> torch.Tensor: - return self.inference_model.transform_qkv_param(self.params) - - -class UnfusedQKVParameter(ParameterBase): - """ - QKV parameter container for unfused QKV projection. - - src_param shapes: 3 x [out_features, in_features] - dst_param shape: [3 x out_features, in_features] - """ - - q_params: torch.Tensor - - k_params: torch.Tensor - - v_params: torch.Tensor - - def finalize(self): - fused_param = torch.cat([self.q_params, self.k_params, self.v_params], dim=0) - return self.inference_model.transform_qkv_param(fused_param) - - -def megatron_qkv_reshape(param: torch.Tensor, head_size: int, n_heads: int) -> torch.Tensor: - assert param.shape[0] == 3 * n_heads * head_size - - all_heads = torch.chunk(param, chunks=3 * n_heads, dim=0) - q_heads = all_heads[::3] - k_heads = all_heads[1::3] - v_heads = all_heads[2::3] - return torch.cat([q_heads, k_heads, v_heads], dim=0) - - -class MegatronQKVParameter(ParameterBase): - """ - QKV parameter container for Megatron-style QKV projection. Megatron stores the parameter - as [n_heads, 3, head_size, in_features] whereas our inference system is built around - [3, n_heads, head_size, in_features]. This container handles the conversion. - - Note: this container expects the model implementation to implement properties for - `head_size` and `n_heads`. - - src_qkv_w shape: [3 * out_features, in_features] - qkv_w shape: [3 * out_features, in_features] - """ - - params: torch.Tensor - - def finalize(self) -> torch.Tensor: - head_size = self.inference_model.head_size - n_heads = self.inference_model.n_heads - - transposed_param = megatron_qkv_reshape(self.params, head_size, n_heads) - return self.inference_model.transform_qkv_param(transposed_param) - - -def transform_gqa_megatron(src_param: torch.Tensor, head_size: int, n_q_heads: int, n_kv_heads: int) -> torch.Tensor: - assert src_param.shape[0] == (2 * n_kv_heads + n_q_heads) * head_size - - head_ratio = n_q_heads // n_kv_heads - - # Reshape to get the groups as the leading dimension - groups_leading_view = src_param.reshape(n_kv_heads, 2 + head_ratio, head_size, -1) - q_heads = groups_leading_view[:, :head_ratio, :, :].reshape(-1, groups_leading_view.shape[-1]) - k_heads = groups_leading_view[:, head_ratio, :, :].reshape(-1, groups_leading_view.shape[-1]) - v_heads = groups_leading_view[:, head_ratio + 1, :, :].reshape(-1, groups_leading_view.shape[-1]) - # Squeeze will remove extra dimension for bias - return torch.cat([q_heads, k_heads, v_heads], dim=0).squeeze() - - -class GQAMegatronQKVParameter(ParameterBase): - """ - QKV parameter for Megatron-style QKV projection with GQA-style QKV projection. In this - storage format each of the groups is stored consecutively, so there will be multiple q_heads, - then one k head, and one v head. - - Note: this container expects the model implementation to implement properties for - `head_size`, `n_q_heads`, and `n_kv_heads`. - - src_qkv_w shape: [(2 * n_kv_heads + n_q_heads) * head_size, in_features] - qkv_w shape: [(2 * n_kv_heads + n_q_heads) * head_size, in_features] - """ - - params: torch.Tensor - - def finalize(self) -> torch.Tensor: - head_size = self.inference_model.head_size - n_q_heads = self.inference_model.n_heads_q - n_kv_heads = self.inference_model.n_heads_kv - transposed_param = transform_gqa_megatron(self.params, head_size, n_q_heads, n_kv_heads) - return self.inference_model.transform_qkv_param(transposed_param) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/unembed_parameters.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/unembed_parameters.py deleted file mode 100644 index 9f67c0ce3c27d2fc07fc735acba8a64bda8bcbf1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/common_parameters/unembed_parameters.py +++ /dev/null @@ -1,26 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from ...model_implementations.parameter_base import ParameterBase -""" -Unembedding containers. -""" - - -class UnembedParameter(ParameterBase): - """ - Unembedding parameter. This will likely be mapped to the same original weight in the model as the - embedding, but we have a different preferred sharding approach. - """ - - params: torch.Tensor - """ - Unembedding parameter of shape [vocab_size, model_dim]. - """ - - def finalize(self) -> torch.Tensor: - return self.inference_model.transform_unembed_param(self.params) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__init__.py deleted file mode 100644 index 20f37538274ccda7ab68ceb7e82c87675120b7e7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .policy import FalconPolicy diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 67930a550825f888e28a4ec1b3784b900780e382..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__pycache__/container.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__pycache__/container.cpython-310.pyc deleted file mode 100644 index 0d056163211fcef85b4889adc120d4e8be6a6ea5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__pycache__/container.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__pycache__/model.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__pycache__/model.cpython-310.pyc deleted file mode 100644 index abe0f7256f85b0cf5935a7a5f70950903e3313ac..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__pycache__/model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__pycache__/policy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__pycache__/policy.cpython-310.pyc deleted file mode 100644 index 6c08dfbcc5cf4ac685d6f22a97d5ad14f2a06765..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/__pycache__/policy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/container.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/container.py deleted file mode 100644 index caccfe1ecb00c37b08121decb250b2625c64eb9a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/container.py +++ /dev/null @@ -1,129 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Create a container object to save model-specific tensors using the policy file above. - -from ..common_parameters import * -from ..layer_container_base import LayerContainer -''' - # HF Falcon 7b model looks like this: - -FalconForCausalLM( - (transformer): FalconModel( - (word_embeddings): Embedding(65024, 4544) - (h): ModuleList( - (0-31): 32 x FalconDecoderLayer( - (self_attention): FalconAttention( - (maybe_rotary): FalconRotaryEmbedding() - (query_key_value): FalconLinear(in_features=4544, out_features=4672, bias=False) - (dense): FalconLinear(in_features=4544, out_features=4544, bias=False) - (attention_dropout): Dropout(p=0.0, inplace=False) - ) - (mlp): FalconMLP( - (dense_h_to_4h): FalconLinear(in_features=4544, out_features=18176, bias=False) - (act): GELU(approximate='none') - (dense_4h_to_h): FalconLinear(in_features=18176, out_features=4544, bias=False) - ) - (input_layernorm): LayerNorm((4544,), eps=1e-05, elementwise_affine=True) - ) - ) - (ln_f): LayerNorm((4544,), eps=1e-05, elementwise_affine=True) - ) - (lm_head): Linear(in_features=4544, out_features=65024, bias=False) -) -''' - - -class FalconTransformerContainer(LayerContainer): - """ - Transformer layer container for the Falcon model. - """ - qkv_w: FusedQKVParameter - attn_out_w: AttentionOutputParameter - mlp_1_w: MLP1Parameter - mlp_2_w: MLP2Parameter - ln_attn_gamma: NormParameter - ln_attn_beta: NormParameter - - PARAM_MAPPING = { - "self_attention.query_key_value.weight": "qkv_w.params", - "self_attention.dense.weight": "attn_out_w.params", - "mlp.dense_h_to_4h.weight": "mlp_1_w.params", - "mlp.dense_4h_to_h.weight": "mlp_2_w.params", - "input_layernorm.weight": "ln_attn_gamma.params", - "input_layernorm.bias": "ln_attn_beta.params", - } - - -class FalconNonTransformerContainer(LayerContainer): - """ - Non-Transformer layer container for the Falcon model. - """ - word_emb: EmbeddingParameter - word_unembed: UnembedParameter - final_norm_gamma: NormParameter - final_norm_beta: NormParameter - - PARAM_MAPPING = { - "transformer.word_embeddings.weight": "word_emb.params", - "transformer.ln_f.weight": "final_norm_gamma.params", - "transformer.ln_f.bias": "final_norm_beta.params", - "lm_head.weight": "word_unembed.params", - } - - -''' - # HF Falcon 40b model looks like this: - - FalconForCausalLM( - (transformer): FalconModel( - (word_embeddings): Embedding(65024, 8192) - (h): ModuleList( - (0-59): 60 x FalconDecoderLayer( - (self_attention): FalconAttention( - (maybe_rotary): FalconRotaryEmbedding() - (query_key_value): FalconLinear(in_features=8192, out_features=9216, bias=False) - (dense): FalconLinear(in_features=8192, out_features=8192, bias=False) - (attention_dropout): Dropout(p=0.0, inplace=False) - ) - (mlp): FalconMLP( - (dense_h_to_4h): FalconLinear(in_features=8192, out_features=32768, bias=False) - (act): GELU(approximate='none') - (dense_4h_to_h): FalconLinear(in_features=32768, out_features=8192, bias=False) - ) - (ln_attn): LayerNorm((8192,), eps=1e-05, elementwise_affine=True) - (ln_mlp): LayerNorm((8192,), eps=1e-05, elementwise_affine=True) - ) - ) - (ln_f): LayerNorm((8192,), eps=1e-05, elementwise_affine=True) - ) - (lm_head): Linear(in_features=8192, out_features=65024, bias=False) -) -''' - - -class FalconNewArchTransformerContainer(LayerContainer): - """ - Transformer layer container for the Falcon model. - """ - qkv_w: GQAMegatronQKVParameter - attn_out_w: AttentionOutputParameter - mlp_1_w: MLP1Parameter - mlp_2_w: MLP2Parameter - ln_attn_gamma: NormParameter - ln_attn_beta: NormParameter - ln_mlp_gamma: NormParameter - ln_mlp_beta: NormParameter - - PARAM_MAPPING = { - "self_attention.query_key_value.weight": "qkv_w.params", - "self_attention.dense.weight": "attn_out_w.params", - "mlp.dense_h_to_4h.weight": "mlp_1_w.params", - "mlp.dense_4h_to_h.weight": "mlp_2_w.params", - "ln_attn.weight": "ln_attn_gamma.params", - "ln_attn.bias": "ln_attn_beta.params", - "ln_mlp.weight": "ln_mlp_gamma.params", - "ln_mlp.bias": "ln_mlp_beta.params", - } diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/model.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/model.py deleted file mode 100644 index b2830c80b562546d138430da702b5c9882ab00b6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/model.py +++ /dev/null @@ -1,213 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Iterable, Optional, Tuple - -import torch - -import deepspeed.comm as dist - -from ...allocator import empty_from -from ...inference_utils import ActivationType, DtypeEnum -from .. import * -from ...modules.configs import * -from ...modules.interfaces import * -from ...ragged import RaggedBatchWrapper - -from .container import FalconNonTransformerContainer, FalconTransformerContainer - - -class FalconInferenceModel(DSTransformerModelBase): - """ - Inference model implementation for ragged batching for Llama-2 models. - """ - - _non_transformer: Optional[FalconNonTransformerContainer] - """ - Embed + unembed container. Specializing the type annotation. - """ - - _transformer: Optional[Iterable[FalconTransformerContainer]] - """ - Per-layer transformer container. Specializing the type annotation. - """ - """ - Properties inherited from `DSInferenceModelBase` - """ - - @property - def max_sequence_length(self) -> int: - return self._config.max_seq_length - - """ - Properties inherited from `DSTransformerModelBase` - """ - - @property - def num_layers(self) -> int: - return self._config.num_hidden_layers - - @property - def model_dim(self) -> int: - return self._config.hidden_size - - @property - def vocab_size(self) -> int: - return self._config.vocab_size - - @property - def head_size(self) -> int: - return self.model_dim // self.n_heads - - @property - def n_heads(self) -> int: - return self._config.num_attention_heads - - @property - def intermediate_dim(self) -> int: - return 4 * self._config.hidden_size - - @property - def n_heads_kv(self) -> int: - return self._config.num_kv_heads if (self._config.new_decoder_architecture - or not self._config.multi_query) else 1 - - @property - def activation_dtype(self) -> DtypeEnum: - if self._config.torch_dtype == torch.float16: - return DtypeEnum.fp16 - elif self._config.torch_dtype == torch.bfloat16: - return DtypeEnum.bf16 - else: - raise NotImplementedError("Only fp16 and bf16 are supported") - - @property - def mlp_activation_fn(self) -> ActivationType: - return ActivationType.GELU - - @property - def norm_type(self) -> NormTypeEnum: - return NormTypeEnum.LayerNorm - - @property - def positional_embedding_type(self) -> PositionalEmbeddingType: - return PositionalEmbeddingType.rotate_half - - @property - def positional_embedding_config(self) -> RotateHalfConfig: - """ - The positional embedding configuration for the model. - """ - return RotateHalfConfig() - - """ - Forward implementations - """ - - def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs the embedding lookup prior to running the transformer of the model. - - Arguments: - ragged_batch (RaggedBatchWrapper): The batch to embed. - - Returns: - torch.Tensor: The embedded batch. - """ - embed = self.embed(ragged_batch, self._non_transformer.word_emb) - - if embed.shape[-1] != self.model_dim: - raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}") - - return embed - - def _forward_transformer_layer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor, - ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead - optimization to fuse the layer norm of the next layer into the current layer. - - Arguments: - layer_idx (int): The index of the layer to execute. - residual (torch.Tensor): The residual tensor from the previous layer. - hidden_states (torch.Tensor): The hidden states from the previous layer. This is the - hidden states after pre normalization. - ragged_batch_info (RaggedBatchWrapper): The batch metadata. - """ - assert self.config.parallel_attn, "Only parallel attention implementation is supported" - - cur_params = self._transformer[layer_idx] - kv_cache = self.state_manager.get_cache(layer_idx) - - attn_ln_out = hidden_states - attn_hidden_state = self.qkv(attn_ln_out, cur_params.qkv_w, b=None) - attn_hidden_state = self.attn(attn_hidden_state, kv_cache, ragged_batch_info) - attention_output = self.attn_out(attn_hidden_state, cur_params.attn_out_w, b=None) - - if self.config.new_decoder_architecture: - residual, mlp_ln_out = self.norm(residual, - None, - gamma=cur_params.ln_mlp_gamma, - beta=cur_params.ln_mlp_beta) - else: - mlp_ln_out = hidden_states - - mlp_hidden_state = self.mlp_1(mlp_ln_out, cur_params.mlp_1_w, b=None) - mlp_output = self.mlp_2(mlp_hidden_state, cur_params.mlp_2_w, b=None) - - mlp_output.add_(attention_output) - - if self.tp_size > 1: - dist.all_reduce(mlp_output, group=self._base_mp_group) - - if layer_idx != self.num_layers - 1: - next_params = self._transformer[layer_idx + 1] - residual, mlp_output = self.norm(residual, - mlp_output, - next_params.ln_attn_gamma, - beta=next_params.ln_attn_beta) - else: - # On last layer, we just need to perform the residual add. Adding into the residual - # here is safe. - residual.add_(mlp_output) - - return residual, mlp_output - - def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs unembedding of the hidden states to logits. This will only sample the final - token of each sequence. - """ - logits = self.unembed(hidden_states, - self._non_transformer.word_unembed, - ragged_batch_info, - gamma=self._non_transformer.final_norm_gamma, - beta=self._non_transformer.final_norm_beta) - - if self.tp_size > 1: - comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1])) - full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size)) - - dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group) - - full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size)) - - return full_logits - else: - return logits - - def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor: - residual = self._forward_embed(wrapped_batch) - - residual, hidden_states = self.norm(residual, - None, - gamma=self._transformer[0].ln_attn_gamma, - beta=self._transformer[0].ln_attn_beta) - - for layer_idx in range(self.num_layers): - residual, hidden_states = self._forward_transformer_layer(layer_idx, residual, hidden_states, - wrapped_batch) - - return self._forward_unembed(residual, wrapped_batch) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/policy.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/policy.py deleted file mode 100644 index c6612090a0df41228244c193b492493824a86394..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/falcon/policy.py +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any - -from ...config_v2 import RaggedInferenceEngineConfig -from ..inference_policy_base import ContainerMap, InferenceV2Policy -from .container import FalconNonTransformerContainer, FalconTransformerContainer -from .container import FalconNewArchTransformerContainer -from .model import FalconInferenceModel - - -class FalconPolicy(InferenceV2Policy): - - def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> FalconInferenceModel: - return FalconInferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group) - - def build_container_map(self) -> ContainerMap: - map = ContainerMap() - - trans_container_cls = FalconNewArchTransformerContainer if self._model_config.new_decoder_architecture else FalconTransformerContainer - transformer_containers = [trans_container_cls(self.model) for _ in range(self.model.num_layers)] - - map.set_transformer_params(['transformer.h'], transformer_containers) - - map.set_non_transformer_params(FalconNonTransformerContainer(self.model)) - - map.set_unmapped_params( - [f'model.layers.{i}.self_attn.rotary_emb.inv_freq' for i in range(self.model.num_layers)]) - - return map diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/flat_model_helpers.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/flat_model_helpers.py deleted file mode 100644 index c5e02adaffc4d4f095c966ff31eb985bb27eca26..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/flat_model_helpers.py +++ /dev/null @@ -1,282 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Dict, Iterable, Tuple, Optional -from os import path - -import torch - -from deepspeed.accelerator import get_accelerator -from deepspeed.ops.op_builder import RaggedUtilsBuilder -from deepspeed.runtime.config_utils import DeepSpeedConfigModel -from .layer_container_base import LayerContainer -from ..inference_parameter import InferenceParameter, STR_TO_DTYPE -from ..inference_utils import elem_size - - -def pad_to_aligned_offset(offset: int, alignment: int = 256) -> int: - """ - Pad the provided offset to a well-aligned value. - """ - return ((offset + alignment - 1) // alignment) * alignment - - -class TensorMetadata(DeepSpeedConfigModel): - """ - A class to represent a tensor specification. - """ - dtype: Optional[str] = None - shape: Optional[Tuple[int, ...]] = None - strides: Optional[Tuple[int, ...]] = None - offset: int - - -class ParameterMetadata(DeepSpeedConfigModel): - """ - A class to represent a parameter specification. - """ - core_param: Optional[TensorMetadata] = None - aux_params: Dict[str, TensorMetadata] = {} - - -class LayerMetadata(DeepSpeedConfigModel): - """ - A class to represent a layer specification. - """ - params: Dict[str, ParameterMetadata] = {} - - -class ModelMetadata(DeepSpeedConfigModel): - """ - A class to represent a model specification. - """ - policy: str = "" - layers: Dict[str, LayerMetadata] = {} - - -def make_param_filename(base: str, rank: int, n_ranks: int) -> str: - """ - Make a filename for a parameter file. - - Arguments: - rank: Rank of the file. - n_ranks: Total number of ranks. - - Returns: - str: Filename. - """ - return path.join(base, f"params_rank_{rank}_of_{n_ranks}.pt") - - -def make_metadata_filename(base: str, rank: int, n_ranks: int) -> str: - """ - Make a filename for a metadata file. - - Arguments: - rank: Rank of the file. - n_ranks: Total number of ranks. - - Returns: - str: Filename. - """ - return path.join(base, f"metadata_rank_{rank}_of_{n_ranks}.json") - - -def make_model_config_filename(base: str) -> str: - """ - Make a filename for a model config file. - - Arguments: - base: Base directory. - - Returns: - str: Filename. - """ - return path.join(base, "ds_model_config.json") - - -def flatten_inference_model( - transformer_containers: Iterable[LayerContainer], - non_transformer_container: LayerContainer, - policy_name: str, -) -> Tuple[torch.Tensor, ModelMetadata]: - """ - Flatten the underlying parameters into - - Arguments: - transformer_containers: Iterable of layer containers corresponding to the transformer - parameters. - non_transformer_container: Layer container corresponding to the non-transformer parameters. - policy_name: The name of the policy class (typically accessed with `type(policy).__name__`). - - Returns: - Iterable[Any]: Flattened list of parameters. - """ - alloc_fn = RaggedUtilsBuilder().load().allocate_view_on - - total_size = 0 - metadata = ModelMetadata(policy=policy_name) - - def process_layer(layer_container: LayerContainer, l_name: str, cur_offset: int) -> int: - """ - Iterate over the parameters of a single container and collect metadata for the final - flattened buffer. - - Arguments: - layer_container: The layer container to process. - l_name: The name of the layer container to key the metadata. - cur_offset: The current offset into the flattened buffer. - - Captured Variables: - metadata: The metadata object to populate. - - Returns: - int: The updated offset into the flattened buffer. - """ - try: - _ = layer_container.is_populated - except ValueError as e: - raise ValueError(f"Layer container {l_name} is not populated.") from e - - layer_metadata = LayerMetadata() - - for p_name in layer_container.annotation_attrs: - param = getattr(layer_container, p_name) - param_metadata = ParameterMetadata() - - if param is None: - param_metadata.core_param = TensorMetadata(offset=-1) - layer_metadata.params[p_name] = param_metadata - continue - - param_metadata.core_param = TensorMetadata(dtype=str(param.dtype), - shape=param.shape, - strides=param.stride(), - offset=cur_offset) - - cur_offset += pad_to_aligned_offset(elem_size(param.dtype) * param.numel()) - - for t_name, tensor in param.aux_attrs.items(): - param_metadata.aux_params[t_name] = TensorMetadata(dtype=str(tensor.dtype), - shape=tensor.shape, - strides=tensor.stride(), - offset=cur_offset) - - cur_offset += pad_to_aligned_offset(elem_size(tensor.dtype) * tensor.numel()) - - layer_metadata.params[p_name] = param_metadata - - metadata.layers[l_name] = layer_metadata - return cur_offset - - for i, layer in enumerate(transformer_containers): - l_name = f"transformer_layer_{i}" - total_size = process_layer(layer, l_name, total_size) - - l_name = "non_transformer" - total_size = process_layer(non_transformer_container, l_name, total_size) - - buffer = torch.empty(total_size, dtype=torch.uint8, device=get_accelerator().current_device()) - - def copy_layer(layer_container: LayerContainer, l_name: str) -> None: - """ - Local method for copying from the layer container to the flattened buffer. - - Arguments: - layer_container: The layer container to copy from. - l_name: The name of the layer container to key the metadata. - - Captured Variables: - buffer: The flattened buffer to copy into. - metadata: The metadata object to populate. - """ - l_metadata = metadata.layers[l_name] - for p_name in layer_container.annotation_attrs: - p_metadata = l_metadata.params[p_name] - param = getattr(layer_container, p_name) - - if param is None: - continue - - core_param = alloc_fn(param, buffer, p_metadata.core_param.offset) - core_param.copy_(param) - - aux_params = {} - - for t_name, tensor in param.aux_attrs.items(): - t_view = alloc_fn(tensor, buffer, p_metadata.aux_params[t_name].offset) - aux_params[t_name] = t_view - t_view.copy_(tensor) - - setattr(layer_container, p_name, InferenceParameter.initialize(core_param, **aux_params)) - - for i, layer in enumerate(transformer_containers): - l_name = f"transformer_layer_{i}" - copy_layer(layer, l_name) - - l_name = "non_transformer" - copy_layer(non_transformer_container, l_name) - - return buffer, metadata - - -def restore_inference_model(buffer: torch.Tensor, metadata: ModelMetadata, - transformer_containers: Iterable[LayerContainer], - non_transformer_container: LayerContainer) -> None: - """ - Restore the model from the buffer and metadata. - - Arguments: - buffer: Buffer containing the model parameters. - metadata: Metadata for the model. - transformer_containers: Iterable of transformer layer containers. - non_transformer_container: Non-transformer layer container. - """ - alloc_fn = RaggedUtilsBuilder().load().allocate_view_like - - def restore_layer(layer_container: LayerContainer, l_name: str) -> None: - """ - Local method for restoring a layer container from a flattened buffer. This - only constructs views for the parameters onto the buffer. No data movement - is performed. - - Arguments: - layer_container: The layer container to restore. - l_name: The name of the layer container to key the metadata. - - Captured Variables: - buffer: The flattened buffer to reconstruct views on top of. - metadata: The metadata object describing the each parameter in the model. - """ - l_metadata = metadata.layers[l_name] - - for p_name in layer_container.annotation_attrs: - p_metadata = l_metadata.params[p_name] - - if p_metadata.core_param.offset == -1: - layer_container.direct_injection(p_name, None) - continue - - dummy_tensor = torch.empty([], dtype=STR_TO_DTYPE[p_metadata.core_param.dtype]) - core_param = alloc_fn(p_metadata.core_param.shape, p_metadata.core_param.strides, dummy_tensor, buffer, - p_metadata.core_param.offset) - - aux_params = {} - - for t_name, t_metadata in p_metadata.aux_params.items(): - dummy_tensor = torch.empty([], dtype=STR_TO_DTYPE[t_metadata.dtype]) - t_view = alloc_fn(t_metadata.shape, t_metadata.strides, dummy_tensor, buffer, t_metadata.offset) - - aux_params[t_name] = t_view - - restored_param = InferenceParameter.initialize(core_param, **aux_params) - layer_container.direct_injection(p_name, restored_param) - - for i, layer in enumerate(transformer_containers): - l_name = f"transformer_layer_{i}" - restore_layer(layer, l_name) - - l_name = "non_transformer" - restore_layer(non_transformer_container, l_name) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/inference_model_base.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/inference_model_base.py deleted file mode 100644 index 894a4137407e9b0c7c38f18d504cebf2977649bd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/inference_model_base.py +++ /dev/null @@ -1,272 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import ABC, abstractmethod -from typing import Iterable, Optional, Tuple, Type - -import torch - -import deepspeed.comm as dist -from ..ragged import DSStateManager, RaggedBatchWrapper -from ..ragged.manager_configs import KVCacheConfig -from ..ragged import DSSequenceDescriptor -from ..model_implementations.layer_container_base import LayerContainer -from ..config_v2 import RaggedInferenceEngineConfig -from .flat_model_helpers import ModelMetadata - -try: - from functools import cached_property -except ImportError: - - def cached_property(func): - return property(func) - - -""" -This abstract class defines the interfaces that a model implementation should implement -in order to include anything that may be called by the engine. Most models should be able -to inherit from `DSInferenceTransformerModelBase` to reduce implementation work so it is recommended -to begin there. -""" -""" -Placeholder for typing the model config, which can vary based on model implementation/ -""" -DSModelImplementationConfig = Type['DSModelImplementationConfig'] -""" -Placeholder for typing the distributed comm object. - -TODO(cmikeh2): Replace when we have a more defined API for the inference communication system. -""" -MPType = Type["MPType"] - - -class DSInferenceModelBase(torch.nn.Module, ABC): - """ - Implementation of a model for inference composable with ragged batching. - """ - - _config: DSModelImplementationConfig - """ - Model-specific configuration. No abstraction surrounds this yet. - """ - - _engine_config: RaggedInferenceEngineConfig - """ - Engine configuration. - """ - - _base_mp_group: MPType - """ - Base communication group for Tensor-parallel inference. - """ - - _non_transformer: Optional[LayerContainer] - """ - Abstract container for storing both embedding (pre-transformer) and unembedding (post-transformer) - parameters. This attribute should be None at model instantiation until the Policy sets - the model parameters. These parameters are grouped together since many model implementations - will tie the embedding and unembedding parameters together. - """ - - _transformer: Optional[Iterable[LayerContainer]] - """ - List of abstract containers (1 per layer) for storing transformer (transformer) - parameters. This attribute should be None at model instantiation until the Policy - sets the model parameters. - """ - - state_manager: Optional[DSStateManager] - """ - Since the state manager is lazy initialized, by the engine, it is not guaranteed to be present - until full initialization. - """ - - def __init__(self, config: DSModelImplementationConfig, engine_config: RaggedInferenceEngineConfig, - base_mp_group: MPType) -> None: - """ - Minimal initialization of the model. - - Arguments: - config (DSModelImplementationConfig): Model-specific configuration. No assumptions - should be made about this config that are not closely tied to the specific - model implementation. - engine_config (RaggedInferenceEngineConfig): Engine configuration. - base_mp_group (MPType): Base communication group for Tensor-parallel inference. - """ - super().__init__() - self._config = config - self._engine_config = engine_config - self._base_mp_group = base_mp_group - - # Set to None until the Policy sets the model parameters - self._non_transformer = None - self._transformer = None - self._flattened_param_buffer = None - self._flattened_param_metadata = None - - @property - def config(self) -> DSModelImplementationConfig: - """ - The model config. - """ - return self._config - - def set_parameters(self, transformer: Iterable[LayerContainer], non_transformer: LayerContainer, - flattened_param_buffer: torch.Tensor, flattened_param_metadata: ModelMetadata): - """ - Set the model parameters for the embedding, transformer, and unembedding containers. - """ - self._transformer = transformer - self._non_transformer = non_transformer - self._flattened_param_buffer = flattened_param_buffer - self._flattened_param_metadata = flattened_param_metadata - - def set_state_manager(self, state_manager: DSStateManager): - """ - Sets the state manager attribute. This is called by the inference engine after - the model is fully initialized. - """ - self.state_manager = state_manager - - @cached_property - def tp_rank(self) -> int: - """ - The rank of the current process. - - # TODO(cmikeh2): Kind of a hack right now, but this is too verbose to use at - the frequency we need. - """ - return dist.get_rank(group=self._base_mp_group) - - @cached_property - def tp_size(self) -> int: - """ - The total number of processes. - - # TODO(cmikeh2): Kind of a hack right now, but this is too verbose to use at - the frequency we need. - """ - return dist.get_world_size(group=self._base_mp_group) - - @property - def model_config(self): - """ - The model config. - """ - return self._config - - @property - def engine_config(self): - """ - The engine config. - """ - return self._engine_config - - @property - def flattened_params(self) -> Optional[torch.Tensor]: - """ - The flattened parameter buffer. - """ - return self._flattened_param_buffer - - @property - def flattened_param_metadata(self) -> Optional[ModelMetadata]: - """ - The flattened parameter metadata. - """ - return self._flattened_param_metadata - - @abstractmethod - def get_kv_requirements(self, sequence: DSSequenceDescriptor, max_new_tokens: int, - max_new_blocks: Tuple[int, ...]) -> Tuple[int, torch.Tensor]: - """ - Given a sequence and the number of new tokens in the sequence, determine the - number of new KV blocks needed to support the sequence. This method is - used to help the engine provide schedulability APIs and can be used as a helper - for ``maybe_allocate_kv``. - - Args: - sequence (DSSequenceDescriptor): The sequence for which to allocate KV-storage. - max_new_tokens (int): Maximum number of tokens to hypothetically schedule. - max_new_blocks (int): Maximum number of blocks to hypothetically allocate. - - Returns: - Tuple[int, torch.Tensor]: The tuple of number of tokens scheduled and number - of blocks allocated (per KV cache). In general, only one of these numbers will - match the corresponding input argument, but this is not guaranteed. - """ - raise NotImplementedError() - - @abstractmethod - def get_remaining_block_capacity(self, sequence: DSSequenceDescriptor) -> int: - raise NotImplementedError() - - @abstractmethod - def maybe_allocate_kv(self, sequence: DSSequenceDescriptor, n_new_tokens: int) -> None: - """ - Given a sequence and the number of new tokens in the sequence, determine - whether or not additional KV-storage is needed and allocate it if so. - - Args: - sequence (DSSequenceDescriptor): The sequence for which to allocate KV-storage. - n_new_tokens (int): The number of new tokens in the sequence. - """ - raise NotImplementedError() - - @abstractmethod - def kv_cache_config(self) -> Tuple[KVCacheConfig, ...]: - """ - Return the KV-cache configuration for this model. This should be a tuple of one or more - KVCacheConfig objects (one for each distinct cache group). - """ - raise NotImplementedError() - - @property - @abstractmethod - def max_sequence_length(self) -> int: - """ - The maximum sequence length supported by the model. - """ - ... - - def maybe_free_kv(self, sequence: DSSequenceDescriptor) -> None: - """ - After completing a forward pass, determine whether or not the there are any KV blocks - that maybe freed since they are no longer in use. - - Consider the following example: - - We have a block size of 4 and a local window size of 8. At the beginning of the forward - pass there 10 tokens had been seen and the new forward has a size of 4. This would lend - itself to the following cache structure prior to the forward: - [[0, 1, 2*, 3*] [4*, 5*, 6*, 7*] [8*, 9*, x, x] [x x x x]] - Where x's denote empty cache locations and * denote values that are needed for attention - of the next open slot. After the forward, the cache would look like the following: - [[0, 1, 2, 3] [4, 5, 6*, 7*] [8*, 9*, 10*, 11*] [12* 13* x x]] - In this case, the first block is no longer needed since it is not needed for any future - local attention windows. This function would be responsible for freeing that block. - - Default behavior assumes no local patterns that require freeing and in general should - be sufficient. - """ - pass - - @abstractmethod - def prepare_batch(self, wrapped_batch: RaggedBatchWrapper) -> None: - """ - This will be called before each forward with the intent of building forward-specific metadata - about a batch. The intent here is to build data structures like attention atoms without necessarily - needing to implement graphable kernels to do so. - - Abstract so as to force model implementations to opt out of doing anything here explicitly. - """ - raise NotImplementedError() - - def forward(wrapped_batch: RaggedBatchWrapper) -> torch.Tensor: - """ - Complete a forward pass of the model. This interface should be graphable, so it - should not rely on the ability to use python control flow. - """ - raise NotImplementedError() diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/inference_policy_base.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/inference_policy_base.py deleted file mode 100644 index 2f4266a8cb880ccf7ce97a82c708962d3b75b059..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/inference_policy_base.py +++ /dev/null @@ -1,220 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import json -from abc import ABC, ABCMeta, abstractmethod -from typing import Any, Iterable, List, Optional, Union - -import torch - -from ..config_v2 import RaggedInferenceEngineConfig -from ..checkpoint import CheckpointEngineBase -from ..logging import inference_logger -from .layer_container_base import LayerContainer -from .inference_model_base import DSInferenceModelBase -from .flat_model_helpers import ( - flatten_inference_model, - make_param_filename, - make_metadata_filename, - ModelMetadata, - restore_inference_model, -) - -POLICIES = {} - - -class ContainerMap: - - def __init__(self) -> None: - self._prefix_map = {} - self._transformer_params = None - self._non_transformer_params = None - - @property - def transformer_params(self) -> Iterable[LayerContainer]: - return self._transformer_params - - @property - def non_transformer_params(self) -> LayerContainer: - return self._non_transformer_params - - def set_transformer_params(self, prefixes: Union[str, Iterable[str]], containers: List[LayerContainer]) -> None: - if not isinstance(containers, list): - raise ValueError( - f"The transformer containers should be a list, of one container per layer, but got {type(containers)} instead." - ) - - self._transformer_prefixes = prefixes if isinstance(prefixes, list) else [prefixes] - self._transformer_params = containers - - def set_non_transformer_params(self, container: LayerContainer) -> None: - self._non_transformer_params = container - - def set_unmapped_params(self, prefixes: Union[str, Iterable[str]]) -> None: - self._unmapped_prefixes = prefixes - - def map_param(self, name, parameter) -> None: - for unmapped_prefix in self._unmapped_prefixes: - if name.startswith(unmapped_prefix): - inference_logger().debug(f"Ignoring: {name} for {unmapped_prefix}") - return - - for transformer_prefix in self._transformer_prefixes: - if name.startswith(transformer_prefix): - popped_name = name[len(transformer_prefix) + 1:] - layer_idx = popped_name.split(".")[0] - assert layer_idx.isdigit( - ), f"expected name to start w. list index but got {layer_idx} instead, name={name}" - layer_idx = int(layer_idx) - inference_logger().debug( - f"Setting: {'.'.join(popped_name.split('.')[1:])} for layer-idx={layer_idx} to {parameter.shape}") - self._transformer_params[layer_idx].set_dependency(".".join(popped_name.split(".")[1:]), parameter) - return - - try: - inference_logger().debug(f"Setting: {name} to {parameter.shape}") - self._non_transformer_params.set_dependency(name, parameter) - except ValueError: - # Catch the ValueError here from the non_transformer_params because we are knowingly - # calling it with something that may not match. This should allow us to raise a slightly more - # informative error message. - raise ValueError(f"Cannot find container for {name}, please double check the Containers/ContainerMap") - - def validate(self) -> None: - if not self._non_transformer_params.is_initialized: - raise RuntimeError("Non-transformer parameters not fully initialized after checkpoint load.") - - for layer_idx, container in enumerate(self._transformer_params): - if not container.is_initialized: - raise RuntimeError( - f"Transformer container at index {layer_idx} not fully initialized after checkpoint load.") - - -class PolicyMeta(ABCMeta): - - def __new__(cls, name, bases, dct): - new_obj = super().__new__(cls, name, bases, dct) - if name != "InferenceV2Policy": - POLICIES[name] = new_obj - return new_obj - - -class InferenceV2Policy(ABC, metaclass=PolicyMeta): - """ - The InferenceV2Policy is the base class for all inference policies. An inference policy - is responsible for instantiating the inference model and mapping the parameters from the - checkpoint engine to the model itself. - """ - - def __init__( - self, - model_config: Any, - checkpoint_engine: Optional[CheckpointEngineBase] = None, - inf_checkpoint_path: Optional[str] = None, - ) -> None: - """ - Create the Policy with sufficient context to build the model. There are two supported - model creation mechanisms. - - The first is the generalized ``checkpoint_engine`` which - will iterate over the parameters of the model and provide them to the policy. These in - turn will be sharded/transformed by the model implementation. - - The second is used to re-create a previously serialized DeepSpeed inference model. These - checkpoints should not be used across different model backend configurations. - - TODO(cmikeh2): Enforce this in code - """ - if checkpoint_engine is None and inf_checkpoint_path is None: - raise ValueError("Either checkpoint_engine or ds_checkpoint_path must be provided.") - - if checkpoint_engine is not None and inf_checkpoint_path is not None: - raise ValueError("Only one of checkpoint_engine or ds_checkpoint_path can be provided.") - - self._checkpoint_engine = checkpoint_engine - self._inf_checkpoint_path = inf_checkpoint_path - self._model_config = model_config - - def build_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> DSInferenceModelBase: - """ - Completely instantiate the inference model. This will both create the ops needed to run the - model, as well as load the model parameters via the checkpoint engine. For more context - on each of these components please see ``instantiate_model`` and ``populate_model_parameters``. - - Arguments: - engine_config: The config that has been used to instantiate the engine. This is used - to communicate to the model implementation the limits on batches (sequences/tokens) - and bound the size of intermediate buffers. - mp_group: Object to enable communication between tensor parallel ranks. - - Returns: - DSInferenceModelBase: An implementation of the inference model abstraction that will be - run by the engine. - """ - self.model = self.instantiate_model(engine_config, mp_group) - self.populate_model_parameters() - return self.model - - @abstractmethod - def instantiate_model(self, engine_config: RaggedInferenceEngineConfig) -> DSInferenceModelBase: - """ - Instantiate the inference model. Depending on the engine/model config, this could be where - different model implementations could be selected. - - Arguments: - engine_config: The config that has been used to instantiate the engine. This is used - to communicate to the model implementation the limits on batches (sequences/tokens) - and bound the size of intermediate buffers. - - Returns: - DSInferenceModelBase: An implementation of the inference model abstraction that will be - run by the engine. - """ - ... - - @abstractmethod - def build_container_map(self) -> ContainerMap: - """ - Build a dictionary representing the structure of the string prefixes leading - to the parameters to be mapped to the container. - - Returns: - ContainerMap: An instantiated mapping describing how checkpoint prefixes map - to ``LayerContainer`` instances. - """ - raise NotImplementedError() - - def populate_model_parameters(self) -> None: - """ - This model will iterate over the parameters (as provided by the checkpoint engine) and - use the container map built by ``build_container_map`` to populate the model - """ - - container_map = self.build_container_map() - - if self._checkpoint_engine is not None: - for name, parameter in self._checkpoint_engine.parameters(): - container_map.map_param(name, parameter) - - buffer, metadata = flatten_inference_model(container_map.transformer_params, - container_map.non_transformer_params, self.__class__.__name__) - else: - - buffer_path = make_param_filename(self._inf_checkpoint_path, self.model.tp_rank, self.model.tp_size) - metadata_path = make_metadata_filename(self._inf_checkpoint_path, self.model.tp_rank, self.model.tp_size) - - buffer = torch.load(buffer_path, weights_only=False) - metadata = json.load(open(metadata_path, "r")) - metadata = ModelMetadata.parse_raw(metadata) - - restore_inference_model(buffer, metadata, container_map.transformer_params, - container_map.non_transformer_params) - - container_map.validate() - - self.model.set_parameters(transformer=container_map.transformer_params, - non_transformer=container_map.non_transformer_params, - flattened_param_buffer=buffer, - flattened_param_metadata=metadata) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/inference_transformer_base.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/inference_transformer_base.py deleted file mode 100644 index fae67dc8fc2ad807df5ed4b337bd709cab20c9b5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/inference_transformer_base.py +++ /dev/null @@ -1,617 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import abstractmethod -from typing import Optional - -import torch - -from deepspeed.accelerator import get_accelerator -from ..config_v2 import RaggedInferenceEngineConfig -from ..inference_utils import ActivationType, ceil_div, is_gated -from ..model_implementations import * -from ..model_implementations.sharding import * -from ..modules.configs import ( - DSEmbeddingsConfig, - DSLinearConfig, - DSMoEConfig, - DSNormConfig, - DSSelfAttentionConfig, - DSUnembedConfig, - NormTypeEnum, - PositionalEmbeddingType, - RotateHalfConfig, -) -from ..modules import heuristics -from ..ragged import ( - DSSequenceDescriptor, - KVCacheConfig, - RaggedBatchWrapper, -) -from .inference_model_base import ( - DSInferenceModelBase, - DSModelImplementationConfig, - MPType, -) -from ..inference_parameter import InferenceParameter - -try: - from functools import cached_property -except ImportError: - - def cached_property(func): - return property(func) - - -class DSTransformerModelBase(DSInferenceModelBase): - """ - Dimensioning properties - """ - - @property - @abstractmethod - def num_layers(self) -> int: - """ - Number of the layers in the model - """ - ... - - @property - @abstractmethod - def model_dim(self) -> int: - """ - Size of embedding projection and residuals. - """ - ... - - @property - @abstractmethod - def vocab_size(self) -> int: - """ - Size of the vocabulary (including padding). - """ - ... - - @property - @abstractmethod - def head_size(self) -> int: - """ - Size of each attention head. - """ - ... - - @property - @abstractmethod - def n_heads(self) -> int: - """ - The number of query heads on the model. This should not take into account - any dimension reductions from model sharding. - """ - ... - - @property - def n_heads_q(self) -> int: - """ - Alias to n_heads. - """ - return self.n_heads - - @property - def n_heads_kv(self) -> int: - """ - The number of key and value heads on the model. For GQA or MQA, overload this attribute. - Otherwise it adopts MHA formulations and uses n_heads. This should not take into account - any dimension reductions from model sharding. - """ - return self.n_heads - - @property - @abstractmethod - def intermediate_dim(self) -> int: - """ - The size of the (unsharded) intermediate projection dim. For a gated activation function - this is the size of the input to the second MLP layer. This should not take into account - any dimension reductions from model sharding. - """ - ... - - @property - @abstractmethod - def positional_embedding_type(self) -> PositionalEmbeddingType: - """ - The type of positional embedding used by the model. - """ - ... - - """ - Architectural properties - """ - - @property - @abstractmethod - def activation_dtype(self) -> torch.dtype: - """ - The activation dtype of the model. - """ - ... - - @property - @abstractmethod - def mlp_activation_fn(self) -> ActivationType: - """ - The activation function used in the MLP. - """ - ... - - @property - @abstractmethod - def norm_type(self) -> NormTypeEnum: - """ - The type of normalization used in the model. - """ - ... - - @property - @abstractmethod - def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - """ - The positional embedding configuration for the model. - """ - ... - - """ - Derived helpers - """ - - @cached_property - def n_heads_q_local(self) -> int: - """ - Number of local heads post sharding. - """ - return get_local_heads(self.tp_rank, self.tp_size, self.n_heads_q, self.n_heads_kv)[0] - - @cached_property - def n_heads_kv_local(self) -> int: - """ - Number of local heads post sharding. - """ - return get_local_heads(self.tp_rank, self.tp_size, self.n_heads_q, self.n_heads_kv)[1] - - @property - def gated_mlp(self) -> bool: - """ - Return a boolean to determine whether the model uses a gated activation function. - """ - return is_gated(self.mlp_activation_fn) - - """ - Method implementations - """ - - def __init__(self, config: DSModelImplementationConfig, engine_config: RaggedInferenceEngineConfig, - base_mp_group: MPType) -> None: - """ - Base implementation for initialization. By default, this will initialize - the traditional components of a transformer model: - - Embedding - - QKV projection - - Self attention - - Attention output projection - - Feed forward network - - Normalization - - Unembedding - - Arguments: - config (DSModelImplementationConfig): Model-specific configuration. No assumptions - should be made about this config that are not closely tied to the specific - model implementation. - engine_config (RaggedInferenceEngineConfig): Engine configuration. - base_mp_group (MPType): Base communication group for Tensor-parallel inference. - """ - super().__init__(config, engine_config, base_mp_group) - - self.make_norm_layer() - self.make_qkv_layer() - self.make_attn_layer() - self.make_attn_out_layer() - self.make_mlp_1_layer() - self.make_mlp_2_layer() - self.make_embedding_layer() - self.make_unembedding_layer() - self._kv_cache_config = None - - ######### Embedding ######### - def make_embedding_layer(self) -> None: - """ - Performs setup and creates embedding DSModule. This will set the `self.embed` attribute. - """ - - embed_config = DSEmbeddingsConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - residual_dtype=self.activation_dtype, - embedding_dim=self.model_dim, - ) - - self.embed = heuristics.instantiate_embed(embed_config, self._engine_config) - - def transform_embedding_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Performs embedding sharding along the channels dimension. - """ - # Until we can do non-contiguous all-gather, we won't shard the embedding parameters. - param = param.to(self.activation_dtype.value) - return InferenceParameter.initialize(param) - - ######### Unembedding ######### - def make_unembedding_layer(self) -> None: - """ - Performs setup and creates an unembedding layer. This implementation assumes - normalization prior to the LM head projection. If this does not match the model's - implementation, override this method. This will set the ``self.unembed`` attribute. - """ - unembed_dim = sharded_unembed_dim(self.vocab_size, self.tp_rank, self.tp_size) - - unembed_config = DSUnembedConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - max_sequences=self._engine_config.state_manager.max_ragged_sequence_count, - dtype=self.activation_dtype, - model_dim=self.model_dim, - vocab_size=unembed_dim, - norm_type=self.norm_type, - ) - - self.unembed = heuristics.instantiate_unembed(unembed_config, self._engine_config) - - if self.tp_size > 1: - self._comm_logits = torch.empty(self.tp_size, - self._engine_config.state_manager.max_ragged_sequence_count, - unembed_dim, - device=get_accelerator().current_device(), - dtype=self.activation_dtype.value) - self._return_logits = torch.empty(self._engine_config.state_manager.max_ragged_sequence_count, - self.vocab_size, - device=get_accelerator().current_device(), - dtype=self.activation_dtype.value) - - def transform_unembed_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Performs sharding along the vocab dimension. - """ - param = shard_unembed_param(param, self.tp_rank, self.tp_size).to(self.activation_dtype.value) - return InferenceParameter.initialize(param) - - ######### QKV ######### - def make_qkv_layer(self) -> None: - """ - Instantiates the linear projection layer for the QKV linear layer. This sets the - `self.qkv` attribute. - """ - out_features = qkv_out_features(self.model_dim, self.tp_rank, self.tp_size, self.head_size, self.n_heads_q, - self.n_heads_kv) - - linear_config = DSLinearConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - in_channels=self.model_dim, - out_channels=out_features, - input_dtype=self.activation_dtype, - output_dtype=self.activation_dtype, - ) - - self.qkv = heuristics.instantiate_linear(linear_config, self._engine_config) - - def transform_qkv_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Passes a QKV parameter to the underlying implementation for any necessary - transformations. - - Args: - param (torch.Tensor): The parameter to transform. This may be either a bias or weight and should have - the shape (out_neurons, in_neurons) - """ - param = shard_qkv_param(param, self.tp_rank, self.tp_size, self.head_size, self.n_heads_q, self.n_heads_kv) - return self.qkv.transform_param(param) - - ######### Attention ######### - def make_attn_layer(self) -> None: - """ - Builds the attention layer for the model. This sets the `self.attn` attribute. - """ - softmax_scale = 1.0 / (self.head_size**0.5) - - attn_config = DSSelfAttentionConfig(max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - n_heads_q=self.n_heads_q_local, - n_heads_kv=self.n_heads_kv_local, - head_size=self.head_size, - max_sequences=self._engine_config.state_manager.max_ragged_sequence_count, - scale_factor=softmax_scale, - input_dtype=self.activation_dtype, - output_dtype=self.activation_dtype, - positional_embedding_type=self.positional_embedding_type, - positional_embedding_config=self.positional_embedding_config) - - self.attn = heuristics.instantiate_attention(attn_config, self._engine_config) - - def get_kv_requirements(self, sequence: DSSequenceDescriptor, max_new_tokens: int, - max_new_blocks: int) -> Tuple[int, int]: - """ - See ``DSInferenceModelBase.get_kv_requirements`` for documentation. - - This method assumes an autoregressive dense attention pattern. Override this method - if this does not match the model's attention pattern. - """ - total_tokens = sequence.seen_tokens + max_new_tokens - req_blocks = ceil_div(total_tokens, self.attn.kv_block_size) - block_lim = req_blocks - sequence.cur_allocated_blocks - - if block_lim <= max_new_blocks: - return max_new_tokens, block_lim - - token_capacity = (max_new_blocks + - sequence.cur_allocated_blocks) * self.attn.kv_block_size - sequence.seen_tokens - - return token_capacity, max_new_blocks - - def get_remaining_block_capacity(self, sequence: DSSequenceDescriptor) -> int: - return sequence.seen_tokens % self.attn.kv_block_size - - def maybe_allocate_kv(self, sequence: DSSequenceDescriptor, n_new_tokens: int) -> None: - """ - See ``DSInferenceModelBase.maybe_allocate_kv`` for documentation. - - This method assumes an autoregressive dense attention pattern. Override this method - if this does not match the model's attention pattern. - """ - free_block = self.state_manager.free_blocks[0] - _, n_needed_blocks = self.get_kv_requirements(sequence, n_new_tokens, free_block) - - if n_needed_blocks > 0: - new_blocks = self.state_manager.allocate_blocks(n_needed_blocks) - sequence.extend_kv_cache(new_blocks) - - def kv_cache_config(self) -> Tuple[KVCacheConfig, ...]: - """ - See ``DSInferenceModelBase.kv_cache_config`` for documentation. - - This method assumes an autoregressive dense attention pattern. Override this method - if this does not match the model's attention pattern. - """ - if self._kv_cache_config is None: - cache_shape = (self.num_layers, self.n_heads_kv_local, self.head_size) - max_blocks = ceil_div(self.max_sequence_length, self.attn.kv_block_size) - self._kv_cache_config = KVCacheConfig(block_size=self.attn.kv_block_size, - cache_shape=cache_shape, - cache_dtype=self.activation_dtype, - max_blocks_per_allocation_group=max_blocks) - return (self._kv_cache_config, ) - - def prepare_batch(self, wrapped_batch: RaggedBatchWrapper) -> None: - """ - See ``DSInferenceModelBase.prepare_batch`` for documentation. - - This method assumes an autoregressive dense attention pattern. Override this method - if this does not match the model's attention pattern. - """ - self.attn.build_atoms(wrapped_batch) - - ######### Attention output ######### - def make_attn_out_layer(self) -> None: - """ - Instantiates the linear projection layer for the attention output linear layer. This sets the - `self.attn_out` attribute. - """ - in_features = attn_out_in_features(self.model_dim, self.tp_rank, self.tp_size, self.head_size, self.n_heads_q, - self.n_heads_kv) - - linear_config = DSLinearConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - in_channels=in_features, - out_channels=self.model_dim, - input_dtype=self.activation_dtype, - output_dtype=self.activation_dtype, - ) - - self.attn_out = heuristics.instantiate_linear(linear_config, self._engine_config) - - def transform_attn_out_param(self, param: torch.Tensor) -> Optional[InferenceParameter]: - """ - Shards an attention output projection parameter and passes it to the underlying - implementation for any necessary transformations. This will return `None` for bias parameters - if they are not on TP rank 0. - - Args: - param (torch.Tensor): The parameter to transform. This may be either a bias or weight and should have - the shape (out_neurons, in_neurons). - """ - param = shard_attn_out_param(param, self.tp_rank, self.tp_size, self.head_size, self.n_heads_q, - self.n_heads_kv) - - if param is not None: - param = self.attn_out.transform_param(param) - - return param - - ######### MLP ######### - def make_mlp_1_layer(self) -> None: - """ - Instantiates the linear projection layer for the first MLP in the feedforward network. - This sets the `self.mlp_1` attribute. - """ - shard_size = sharded_intermediate_dim(self.intermediate_dim, self.tp_size, self.tp_rank) - - linear_config = DSLinearConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - in_channels=self.model_dim, - out_channels=shard_size, - activation=self.mlp_activation_fn, - input_dtype=self.activation_dtype, - output_dtype=self.activation_dtype, - ) - - self.mlp_1 = heuristics.instantiate_linear(linear_config, self._engine_config) - - def transform_mlp_1_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Shards the first MLP parameter and passes it to the underlying implementation - for any necessary transformations. - - Args: - param (torch.Tensor): The parameter to transform. This may be either a bias or weight and should have - the shape (out_neurons, in_neurons). - """ - param = shard_mlp_1_param(param, self.tp_rank, self.tp_size, gated=self.gated_mlp) - - return self.mlp_1.transform_param(param) - - def make_mlp_2_layer(self) -> None: - """ - Instantiates the linear projection layer for the second MLP in the feedforward network. - This sets the `self.mlp_2` attribute. - """ - shard_size = sharded_intermediate_dim(self.intermediate_dim, self.tp_size, self.tp_rank) - - linear_config = DSLinearConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - in_channels=shard_size, - out_channels=self.model_dim, - input_dtype=self.activation_dtype, - output_dtype=self.activation_dtype, - ) - - self.mlp_2 = heuristics.instantiate_linear(linear_config, self._engine_config) - - def transform_mlp_2_param(self, param: torch.Tensor) -> Optional[InferenceParameter]: - """ - Shards the second MLP parameter and passes it to the underlying implementation - for any necessary transformations. This will return `None` for bias parameters - if they are not on TP rank 0. - - Args: - param (torch.Tensor): The parameter to transform. This may be either a bias or weight and should have - the shape (out_neurons, in_neurons). - """ - param = shard_mlp_2_param(param, self.tp_rank, self.tp_size) - - if param is not None: - param = self.mlp_2.transform_param(param) - - return param - - ######### Norm ######### - def make_norm_layer(self) -> None: - """ - Instantiates the normalization layer for the model. This sets the `self.norm` attribute. - - TODO(cmikeh2): In the future we'll distinguish between the different norm objects, - but for now we'll just use the same one for all of them. - """ - norm_config = DSNormConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - type=self.norm_type, - channels=self.model_dim, - residual_dtype=self.activation_dtype, - input_dtype=self.activation_dtype, - output_dtype=self.activation_dtype, - ) - - self.norm = heuristics.instantiate_pre_norm(norm_config, self._engine_config) - - def transform_norm_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Passes a normalization parameter to the underlying implementation for any - necessary transformations. - - TODO(cmikeh2): In the future we'll distinguish between the different norm objects, - but for now we'll just use the same one for all of them. - - Args: - param (torch.Tensor): The parameter to transform. This may be either a bias or weight and should have - shape (model_dim,) - """ - return self.norm.transform_param(param) - - -class DSMoETransformerModelBase(DSTransformerModelBase): - - @property - def n_experts(self) -> int: - """ - Return the number of experts in the model. - """ - raise NotImplementedError("Attempted to access an unimplemented number of experts") - - @property - def n_top_k(self) -> int: - """ - Number of experts per token. - """ - raise NotImplementedError("Attempted to access an unimplemented number of experts per token") - - @property - def normalize_expert_scores(self) -> bool: - """ - Whether to normalize expert scores. If true, sum(expert_scores) = 1. - """ - raise NotImplementedError("Attempted to access an unimplemented normalization flag") - - def make_moe_layer(self) -> None: - """ - Instantiates the MoE layer for the model. This sets the `self.moe` attribute. - """ - sharded_dim = sharded_intermediate_dim(self.intermediate_dim, self.tp_size, self.tp_rank) - - moe_config = DSMoEConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - model_dim=self.model_dim, - intermediate_features=sharded_dim, - activation=self.mlp_activation_fn, - n_experts=self.n_experts, - top_k=self.n_top_k, - input_dtype=self.activation_dtype, - output_dtype=self.activation_dtype, - normalize_scores=self.normalize_expert_scores, - ) - - self.moe = heuristics.instantiate_moe(moe_config, self._engine_config) - - def transform_moe_gate_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Passes a MoE gate parameter to the underlying implementation for any necessary transformations. - - TODO(cmikeh2): This will need to be updated/overridden for expert parallelism. - """ - return self.moe.transform_gate_param(param) - - def transform_moe_mlp_1_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Shards the first MoE param and passes it to the underlying implementation. Since it's possible for an architecture - to have both MoE and non-MoE layers, this can't be overloaded on the MLP1 transform. Furthermore, since both - the MoE DSModule owns both MLP1 and MLP2, under certain sharding conditions it's not possible for the model implementation - to infer from the shape whether to perform a different transformation based on MLP1 or MLP2. This (and the below) - separations are intended to solve both these issues. - - Args: - param (torch.Tensor): The parameter to transform. This should have shape (n_experts, out_neurons, in_neurons). - """ - param = shard_mlp_1_param(param, self.tp_rank, self.tp_size, gated=self.gated_mlp, is_moe=True) - - return self.moe.transform_moe_mlp_1_param(param) - - def transform_moe_mlp_2_param(self, param: torch.Tensor) -> Optional[torch.Tensor]: - """ - Shards the second MoE param and passes it to the underlying implementation. See the above for context on why this API - exists. - - This will return `None` for expert bias params not on TP rank 0. NOTE(cmikeh2): Does it make sense to round-robin assign? - My intuition is that this will make debugging much more difficult for minimal memory reduction. - - Args: - param (torch.Tensor): The parameter to transform. This should have shape (n_experts, out_neurons, in_neurons). - """ - param = shard_mlp_2_param(param, self.tp_rank, self.tp_size, is_moe=True) - - if param is not None: - param = self.moe.transform_moe_mlp_2_param(param) - - return param diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/layer_container_base.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/layer_container_base.py deleted file mode 100644 index feb65b4a5f5d143807bddb44ee09844aeeb00141..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/layer_container_base.py +++ /dev/null @@ -1,355 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import re -from typing import Type - -import torch - -from deepspeed.accelerator import get_accelerator -from .parameter_base import ParameterBase, ParametrizedList -from ..inference_parameter import InferenceParameter - -# Currently have dependency loops for the type hints. -InferenceModel = Type["InferenceModel"] -LayerContainer = Type["LayerContainer"] # noqa: F811 - -MAPPING_KEY = "PARAM_MAPPING" -PLIST_HELPERS = "_ds_plist_strip_vals" - - -def make_finalization_callback(all_names: str): - """ - Helper method for building the finalization callback for a LayerContainer. This - is not client code and should not be used or called directly. - """ - - def finalization_callback(self, param: ParameterBase, finalized_param: torch.Tensor) -> None: - """ - Callback for when a parameter is finalized. - """ - self._finalized_params += 1 - - for name in all_names: - if getattr(self, name) is param: - setattr(self, name, finalized_param) - - return finalization_callback - - -class LayerMetaclass(type): - """ - MetaClass for the LayerContainer base class. This class will parse the annotations - of the class that correspond to `ParameterBase` and create None initializers for each - as well as a finalization callback that for when each `ParameterBase` is finalized - and should be replaced with a Tensor. - """ - - def __new__(cls, clsname, bases, attrs): - - annotations = attrs.get("__annotations__", {}) - - for base in bases: - # We'll pick up all annotations on any base classes. This will allow us to - # to use inheritance to share common parameter groups in base classes. - if hasattr(base, "__annotations__"): - annotations.update(base.__annotations__) - - if hasattr(base, MAPPING_KEY): - if MAPPING_KEY not in attrs: - # This is likely a fail state. If a parent has MAPPING KEY but the child does - # not, then we're guaranteed only a subset of the parameters will be mapped. - attrs[MAPPING_KEY] = {} - attrs[MAPPING_KEY].update(getattr(base, MAPPING_KEY)) - - all_names = [name for name, annotation in annotations.items() if issubclass(annotation, ParameterBase)] - - if MAPPING_KEY in attrs: - # If we have a mapping key at all, then we will enter the validation mode for building - # helpers for mapping and ensuring we have complete mapping. - - # First we'll build a flat list of every dependency for this layer. - all_deps = set() - for name in all_names: - parameter_deps = [ - name for name, annotation in annotations[name].__annotations__.items() - if issubclass(annotation, (torch.Tensor, ParametrizedList)) - ] - - all_deps.update([f"{name}.{dep}" for dep in parameter_deps]) - - # Create static helper for doing the string processing only once. - attrs[PLIST_HELPERS] = [] - - # Iterate over all the mappings - for src_name, target_or_targets in attrs[MAPPING_KEY].items(): - if isinstance(target_or_targets, str): - target_or_targets = [target_or_targets] - - actual_targets = [] - for target_name in target_or_targets: - base_dependency, dependency_attr = target_name.split(".") - - # Check for invalid mappings - if base_dependency not in all_names: - raise ValueError( - "Target parameter \"{}\" not found in this layer. Valid targets are {}".format( - base_dependency, all_names)) - if dependency_attr not in annotations[base_dependency].__annotations__: - # This check is not universal (see below) if a single dependency is being - # mapped to by a single row. - raise ValueError( - "Target dependency \"{}\" not found on parameter \"{}\". Valid targets are {}".format( - dependency_attr, base_dependency, annotations[base_dependency].__annotations__.keys())) - if target_name not in all_deps: - raise ValueError( - "Target dependency \"{}\" was targeted with multiple mapping rules.".format(target_name)) - - # If we've made it this far, the dependency definitely exists. - actual_targets.append(annotations[base_dependency].__annotations__[dependency_attr]) - - all_deps.remove(target_name) - - are_plists = [issubclass(target, ParametrizedList) for target in actual_targets] - if all(are_plists): - # We can do direct sets on everything but ParametrizedLists, so we'll only explicitly - # handle these here. - # TODO(cmikeh2): SPLIT, error if more than 1 - glob_count = src_name.count("*") - if glob_count > 1: - raise ValueError( - "ParametrizedList index inference can only work with a single glob: {}".format(src_name)) - elif glob_count == 0: - raise ValueError( - "Must have wildcard (*) in source name for ParametrizedList mapping: {}".format(src_name)) - - wildcard_idx = src_name.find("*") - prefix = src_name[:wildcard_idx] - suffix = src_name[wildcard_idx + 1:] - attrs[PLIST_HELPERS].append((prefix, suffix, target_or_targets)) - elif any(are_plists): - raise ValueError("Cannot mix ParametrizedLists and Tensors in a single mapping rule.") - - if len(all_deps) > 0: - raise ValueError( - "A parameter mapping was provided for {}, but the following dependencies were not mapped: {}". - format(clsname, all_deps)) - - attrs["finalization_callback"] = make_finalization_callback(all_names) - - new_obj = super().__new__(cls, clsname, bases, attrs) - - setattr(new_obj, "_n_params", len(all_names)) - setattr(new_obj, "_annotation_attrs", all_names) - - return new_obj - - def __call__(cls, *args, **kwargs): - instance = cls.__new__(cls, *args, **kwargs) - instance.__init__(*args, **kwargs) - - for name, annotation in instance.__annotations__.items(): - if issubclass(annotation, ParameterBase): - # TODO(cmikeh2): Do we want to make this a property - # It might also make sense to do this in the base class __init__ - # but since it is tied with the changes made in __new__ it feels - # to me like it should be here. - setattr(instance, name, annotation(instance.inference_model, instance)) - - return instance - - -class LayerContainer(metaclass=LayerMetaclass): # noqa: F811 - """ - Abstract base class for containing model parameters. - - This is primarily a guidance abstraction since we do not put any restrictions - on how the parameters are stored. - - To use this class, annotate the class with `ParameterBase` subclasses and give them - names. As a checkpoint is loaded into this container, the `ParameterBase` instances - will be replaced with realized Tensors as soon as each of their dependencies are met. - - To enable automatic mapping, add a static attribute `PARAM_MAPPING` to the class - definition. This should be a dictionary mapping from a source string to one or - more dependencies. - - ```python - class MyLayer(LayerContainer): - PARAM_MAPPING = { - "path.to.param.dependency", "container_param_1.dependency", - "path.to.param2.dependency", "container_param_2.dependency", - "path.to.param3.*.dependency", "container_param_3.list_dependency" - } - - ... - ``` - """ - - def __init__(self, model: InferenceModel) -> None: - """ - Initialization of the LayerContainer. This method does not need to be overridden - for any children classes. - - Args: - model (InferenceModel): Inference model that will be used to shard and transform - parameters correctly, as well as provide specific information about the model - for `ParameterizedList`s that may be part of one of the member `ParameterBase`s. - """ - self.inference_model = model - self._finalized_params = 0 - - def _initialization_checker(self, check_device: bool = True) -> bool: - """ - Returns whether or not all parameters have been initialized and transformed by - the model. Once this returns True, all the `ParameterBase` instances will be - torch.Tensors. - """ - if self._finalized_params != self.n_params: - return False - - for name in self._annotation_attrs: - tensor = getattr(self, name) - if tensor is None: - continue - elif not isinstance(tensor, InferenceParameter): - raise ValueError("Layer should be finalized, but {} ({}) is neither InferenceParameter or None".format( - name, type(tensor))) - elif check_device and tensor.device != torch.device(get_accelerator().current_device()): - raise RuntimeError("Layer should be finalized, but {} is not on device {}".format( - name, - get_accelerator().current_device())) - return True - - @property - def is_populated(self) -> bool: - """ - Returns whether or not all parameters have been populated by the checkpoint engine, but - does not validat the parameters are on the correct device. - """ - return self._initialization_checker(check_device=False) - - @property - def is_initialized(self) -> bool: - """ - Returns whether or not all parameters have been initialized and transformed by - the model and are located on the appropriate device. Once this returns True, all - the `ParameterBase` instances ``InferenceParameter``s or explicitly set to ``None``. - """ - return self._initialization_checker() - - @property - def n_params(self) -> int: - """ - The number of parameters this container holds. This is a read-only value - that is set by the metaclass. - """ - return self._n_params - - @property - def annotation_attrs(self) -> list: - return self._annotation_attrs - - @property - def mapping_params(self) -> dict: - return getattr(self.__class__, MAPPING_KEY, {}) - - @property - def plist_helpers(self) -> list: - return getattr(self.__class__, PLIST_HELPERS, []) - - def direct_injection(self, name: str, tensor: InferenceParameter) -> None: - - if name not in self._annotation_attrs: - raise ValueError(f"Cannot directly inject {name}, not a valid parameter.") - - setattr(self, name, tensor) - self._finalized_params += 1 - - def set_dependency(self, dep_name: str, dep_value: torch.Tensor) -> None: - """ - Set dependency can be used for managing dependencies when a mapping is provided - in the class definition for the layer. The dep_name here should have any prefix - for transformer layers removed (such as model.layers.*.attn.qkv.weight -> attn.qkv.weight). - - Args: - dep_name (str): The name of the dependency to set. - dep_value (torch.Tensor): The value to set the dependency to. - """ - - def get_dep_name_target(dep_name: str) -> str: - """ - Helper method for getting the target name for a dependency from the - mapping params. Tries to match exact string first, then looks for - wildcards and attempts regex matching. Will return empty string if - no match found. - """ - if dep_name in self.mapping_params: - # If we have an exact match, it's a direct mapping and we can - # immediately set the value. - return self.mapping_params[dep_name] - - matched_targets = [] - for key, target in self.mapping_params.items(): - regex_key = key.replace("*", ".*") - if re.match(regex_key, dep_name): - matched_targets.append(target) - if len(matched_targets) > 1: - raise ValueError(f"Multiple targets matched for dependency {dep_name}: {matched_targets}") - if matched_targets: - return matched_targets[0] - return "" - - if dep_name in self.mapping_params: - # If we have an exact match, it's a direct mapping and we can immediately set - # the value. - target = self.mapping_params[dep_name] - - # Convert single targets to a list for consistency - if isinstance(target, str): - target = [target] - - for target_name in target: - # Double setting doesn't set the attribute correctly, so we do a getattr then setattr - target_param_name, target_dependency_name = target_name.split(".") - target_param = getattr(self, target_param_name) - setattr(target_param, target_dependency_name, dep_value) - return - - # Otherwise we need to map to one of the parameter lists. - for prefix, suffix, dests in self.plist_helpers: - if dep_name.startswith(prefix) and dep_name.endswith(suffix): - # We have a match, so we can set the value. - target_idx = int(dep_name[len(prefix):-len(suffix)]) - - # Convert single targets to a list for consistency - if isinstance(dests, str): - dests = [dests] - - for dest in dests: - target_param_name, target_dependency_name = dest.split(".") - target_param = getattr(self, target_param_name) - target_dependency = getattr(target_param, target_dependency_name) - target_dependency[target_idx] = dep_value - return - - # TODO: Refactor this with the help of cmikeh2 - # We should be able to combine this with the wildcard matching above. - target = get_dep_name_target(dep_name) - if target: - # Convert single targets to a list for consistency - if isinstance(target, str): - target = [target] - - for target_name in target: - # Double setting doesn't set the attribute correctly, so we do a getattr then setattr - target_param_name, target_dependency_name = target_name.split(".") - target_param = getattr(self, target_param_name) - setattr(target_param, target_dependency_name, dep_value) - return - - raise ValueError( - "Could not find a mapping for dependency \"{}\". Check that it is included in the ``MAPPING_PARAMS``. See docstring for more on ``MAPPING_PARAMS``" - .format(dep_name)) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__init__.py deleted file mode 100644 index 79605a76a4c28151336040e33e5a87f6ab7ec64b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .policy import Llama2Policy diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index d1394b111b5c3e6839d915e210e844ab464818a9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/container.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/container.cpython-310.pyc deleted file mode 100644 index e20dd5439f8a9ac474c0b27a73cd1240869edb3b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/container.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/model.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/model.cpython-310.pyc deleted file mode 100644 index 8e52a56dc67f8e718193132815f8ab308e3de880..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/policy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/policy.cpython-310.pyc deleted file mode 100644 index 0f935e18757a1791c526a4765ce6a52e4049d43d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/__pycache__/policy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/container.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/container.py deleted file mode 100644 index 9de9bdb345743c413841424b3d5ed8755152832a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/container.py +++ /dev/null @@ -1,78 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Create a container object to save model-specific tensors using the policy file above. - -from ..common_parameters import * -from ..layer_container_base import LayerContainer -''' - # HF Llama model looks like this: - -LlamaForCausalLM( - (model): LlamaModel( - (embed_tokens): Embedding(32000, 4096, padding_idx=0) - (layers): ModuleList( - (0-31): 32 x LlamaDecoderLayer( - (self_attn): LlamaAttention( - (q_proj): Linear(in_features=4096, out_features=4096, bias=False) - (k_proj): Linear(in_features=4096, out_features=4096, bias=False) - (v_proj): Linear(in_features=4096, out_features=4096, bias=False) - (o_proj): Linear(in_features=4096, out_features=4096, bias=False) - (rotary_emb): LlamaRotaryEmbedding() - ) - (mlp): LlamaMLP( - (gate_proj): Linear(in_features=4096, out_features=11008, bias=False) - (up_proj): Linear(in_features=4096, out_features=11008, bias=False) - (down_proj): Linear(in_features=11008, out_features=4096, bias=False) - (act_fn): SiLUActivation() - ) - (input_layernorm): LlamaRMSNorm() - (post_attention_layernorm): LlamaRMSNorm() - ) - ) - (norm): LlamaRMSNorm() - ) - (lm_head): Linear(in_features=4096, out_features=32000, bias=False) -) -''' - - -class Llama2TransformerContainer(LayerContainer): - """ - Transformer layer container for the Llama-2 model. - """ - qkv_w: UnfusedQKVParameter - attn_out_w: AttentionOutputParameter - mlp_1_w: GatedMLPParameter - mlp_2_w: MLP2Parameter - attn_norm_gamma: NormParameter - mlp_norm_gamma: NormParameter - - PARAM_MAPPING = { - "self_attn.q_proj.weight": "qkv_w.q_params", - "self_attn.k_proj.weight": "qkv_w.k_params", - "self_attn.v_proj.weight": "qkv_w.v_params", - "self_attn.o_proj.weight": "attn_out_w.params", - "mlp.gate_proj.weight": "mlp_1_w.gate_params", - "mlp.up_proj.weight": "mlp_1_w.up_params", - "mlp.down_proj.weight": "mlp_2_w.params", - "input_layernorm.weight": "attn_norm_gamma.params", - "post_attention_layernorm.weight": "mlp_norm_gamma.params", - } - - -class Llama2NonTransformerContainer(LayerContainer): - """ - Non-Transformer layer container for the Llama-2 model. - """ - word_emb: EmbeddingParameter - word_unembed: UnembedParameter - final_norm: NormParameter - - PARAM_MAPPING = { - "model.embed_tokens.weight": "word_emb.params", - "model.norm.weight": "final_norm.params", - "lm_head.weight": "word_unembed.params", - } diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/model.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/model.py deleted file mode 100644 index a0c81f4d749e13709e44fd2a10a6912d54f04310..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/model.py +++ /dev/null @@ -1,209 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Iterable, Optional, Tuple - -import torch - -import deepspeed.comm as dist - -from ...allocator import empty_from -from ...inference_utils import ActivationType, DtypeEnum -from .. import * -from ...modules.configs import * -from ...modules.interfaces import * -from ...ragged import RaggedBatchWrapper - -from .container import Llama2NonTransformerContainer, Llama2TransformerContainer - - -class Llama2InferenceModel(DSTransformerModelBase): - """ - Inference model implementation for ragged batching for Llama-2 models. - """ - - _non_transformer: Optional[Llama2NonTransformerContainer] - """ - Embed + unembed container. Specializing the type annotation. - """ - - _transformer: Optional[Iterable[Llama2TransformerContainer]] - """ - Per-layer transformer container. Specializing the type annotation. - """ - """ - Properties ineherited from `DSInferenceModelBase` - """ - - @property - def max_sequence_length(self) -> int: - return self._config.max_seq_length - - """ - Properties ineherited from `DSTransformerModelBase` - """ - - @property - def num_layers(self) -> int: - return self._config.num_hidden_layers - - @property - def model_dim(self) -> int: - return self._config.hidden_size - - @property - def vocab_size(self) -> int: - return self._config.vocab_size - - @property - def head_size(self) -> int: - return self.model_dim // self.n_heads - - @property - def n_heads(self) -> int: - return self._config.num_attention_heads - - @property - def intermediate_dim(self) -> int: - return self._config.intermediate_size - - @property - def n_heads_kv(self) -> int: - return self._config.num_key_value_heads - - @property - def activation_dtype(self) -> DtypeEnum: - if self._config.torch_dtype == torch.float16: - return DtypeEnum.fp16 - elif self._config.torch_dtype == torch.bfloat16: - return DtypeEnum.bf16 - else: - raise NotImplementedError("Only fp16 and bf16 are supported") - - @property - def mlp_activation_fn(self) -> ActivationType: - activation = self._config.hidden_act.lower() - # llama model family is special and is always gated so force gated versions of relu, gelu, silu - if activation == "gelu": - return ActivationType.GEGLU - elif activation == "relu": - return ActivationType.ReGLU - elif activation == "gegelu": - return ActivationType.GEGLU - elif activation == "silu": - return ActivationType.SiGLU - else: - raise NotImplementedError(f"Activation {activation} not supported") - - @property - def norm_type(self) -> NormTypeEnum: - return NormTypeEnum.RMSNorm - - @property - def positional_embedding_type(self) -> PositionalEmbeddingType: - return PositionalEmbeddingType.rotate_half - - @property - def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - return RotateHalfConfig(theta_base=self._config.rope_theta) - - """ - Forward implementations - """ - - def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs the embedding lookup prior to running the transformer of the model. - - Arguments: - ragged_batch (RaggedBatchWrapper): The batch to embed. - - Returns: - torch.Tensor: The embedded batch. - """ - embed = self.embed(ragged_batch, self._non_transformer.word_emb) - - if embed.shape[-1] != self.model_dim: - raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}") - - return embed - - def _forward_transformer_layer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor, - ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead - optimization to fuse the layer norm of the next layer into the current layer. - - Arguments: - layer_idx (int): The index of the layer to execute. - residual (torch.Tensor): The residual tensor from the previous layer. - hidden_states (torch.Tensor): The hidden states from the previous layer. This is the - hidden states after pre normalization. - ragged_batch_info (RaggedBatchWrapper): The batch metadata. - """ - # TODO(cmikeh2): Distribute ragged_batch_info to all modules - - cur_params = self._transformer[layer_idx] - kv_cache = self.state_manager.get_cache(layer_idx) - - hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=None) - hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info) - hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=None) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - residual, hidden_states = self.norm(residual, hidden_states, cur_params.mlp_norm_gamma, beta=None) - - # Should be configurable in the future - hidden_states = self.mlp_1(hidden_states, cur_params.mlp_1_w, b=None) - hidden_states = self.mlp_2(hidden_states, cur_params.mlp_2_w, b=None) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - if layer_idx != self.num_layers - 1: - next_params = self._transformer[layer_idx + 1] - residual, hidden_states = self.norm(residual, hidden_states, next_params.attn_norm_gamma, beta=None) - else: - # On last layer, we just need to perform the residual add. Adding into the residual - # here is safe. - residual.add_(hidden_states) - - return residual, hidden_states - - def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs unembedding of the hidden states to logits. This will only sample the final - token of each sequence. - """ - logits = self.unembed(hidden_states, - self._non_transformer.word_unembed, - ragged_batch_info, - gamma=self._non_transformer.final_norm) - - if self.tp_size > 1: - comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1])) - full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size)) - - dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group) - - full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size)) - - return full_logits - else: - return logits - - def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor: - - residual = self._forward_embed(wrapped_batch) - - residual, hidden_states = self.norm(residual, None, self._transformer[0].attn_norm_gamma, beta=None) - - for layer_idx in range(self.num_layers): - residual, hidden_states = self._forward_transformer_layer(layer_idx, residual, hidden_states, - wrapped_batch) - - return self._forward_unembed(residual, wrapped_batch) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/policy.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/policy.py deleted file mode 100644 index bb13ab6d5bf46fc5f834f69a97d08ca008f7e9c8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/llama_v2/policy.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any - -from ...config_v2 import RaggedInferenceEngineConfig -from ..inference_policy_base import ContainerMap, InferenceV2Policy -from .container import Llama2NonTransformerContainer, Llama2TransformerContainer -from .model import Llama2InferenceModel - - -class Llama2Policy(InferenceV2Policy): - - def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> Llama2InferenceModel: - return Llama2InferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group) - - def build_container_map(self) -> ContainerMap: - map = ContainerMap() - - transformer_containers = [Llama2TransformerContainer(self.model) for _ in range(self.model.num_layers)] - - map.set_transformer_params(['model.layers'], transformer_containers) - - map.set_non_transformer_params(Llama2NonTransformerContainer(self.model)) - - map.set_unmapped_params( - [f'model.layers.{i}.self_attn.rotary_emb.inv_freq' for i in range(self.model.num_layers)]) - - return map diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__init__.py deleted file mode 100644 index 60d636693ef3ad9d82af98310ffb7206ec23a781..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .policy import MistralPolicy diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 07dc9527288dec8ab6c024c8965553b09ba629a7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__pycache__/container.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__pycache__/container.cpython-310.pyc deleted file mode 100644 index 4046fa78209835c57ced1f3de3bdb0bc2ace3d07..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__pycache__/container.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__pycache__/model.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__pycache__/model.cpython-310.pyc deleted file mode 100644 index f102e9fbdc483593b746e498455cb8a90011b194..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__pycache__/model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__pycache__/policy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__pycache__/policy.cpython-310.pyc deleted file mode 100644 index f3bf35a4075b4bc2d46d988912f3c33ed387c0cb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/__pycache__/policy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/container.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/container.py deleted file mode 100644 index b4c0956f4049bd1515eb9d95cb865468eb6740b8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/container.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Create a container object to save model-specific tensors using the policy file above. - -from deepspeed.inference.v2.model_implementations.common_parameters import * -from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer -''' - # HF Mistral model (mistralai/Mistral-7B-v0.1) looks like this: -MistralForCausalLM( - (model): MistralModel( - (embed_tokens): Embedding(32000, 4096) - (layers): ModuleList( - (0-31): 32 x MistralDecoderLayer( - (self_attn): MistralAttention( - (q_proj): Linear(in_features=4096, out_features=4096, bias=False) - (k_proj): Linear(in_features=4096, out_features=1024, bias=False) - (v_proj): Linear(in_features=4096, out_features=1024, bias=False) - (o_proj): Linear(in_features=4096, out_features=4096, bias=False) - (rotary_emb): MistralRotaryEmbedding() - ) - (mlp): MistralMLP( - (gate_proj): Linear(in_features=4096, out_features=14336, bias=False) - (up_proj): Linear(in_features=4096, out_features=14336, bias=False) - (down_proj): Linear(in_features=14336, out_features=4096, bias=False) - (act_fn): SiLUActivation() - ) - (input_layernorm): MistralRMSNorm() - (post_attention_layernorm): MistralRMSNorm() - ) - ) - (norm): MistralRMSNorm() - ) - (lm_head): Linear(in_features=4096, out_features=32000, bias=False) -) -''' - - -class MistralTransformerContainer(LayerContainer): - """ - Transformer layer container for the Mistral model. - """ - qkv_w: UnfusedQKVParameter - attn_out_w: AttentionOutputParameter - mlp_1_w: GatedMLPParameter - mlp_2_w: MLP2Parameter - attn_norm_gamma: NormParameter - mlp_norm_gamma: NormParameter - - PARAM_MAPPING = { - "self_attn.q_proj.weight": "qkv_w.q_params", - "self_attn.k_proj.weight": "qkv_w.k_params", - "self_attn.v_proj.weight": "qkv_w.v_params", - "self_attn.o_proj.weight": "attn_out_w.params", - "mlp.gate_proj.weight": "mlp_1_w.gate_params", - "mlp.up_proj.weight": "mlp_1_w.up_params", - "mlp.down_proj.weight": "mlp_2_w.params", - "input_layernorm.weight": "attn_norm_gamma.params", - "post_attention_layernorm.weight": "mlp_norm_gamma.params", - } - - -class MistralNonTransformerContainer(LayerContainer): - """ - Non-Transformer layer container for the Mistral model. - """ - word_emb: EmbeddingParameter - word_unembed: UnembedParameter - final_norm: NormParameter - - PARAM_MAPPING = { - "model.embed_tokens.weight": "word_emb.params", - "model.norm.weight": "final_norm.params", - "lm_head.weight": "word_unembed.params", - } diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/model.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/model.py deleted file mode 100644 index 318d362f1a64acc117a3e266d88afdd0b336ba56..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/model.py +++ /dev/null @@ -1,207 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Iterable, Optional, Tuple - -import torch - -import deepspeed.comm as dist - -from ...allocator import empty_from -from ...inference_utils import ActivationType, DtypeEnum -from ...model_implementations import * -from ...modules.configs import * -from ...modules.interfaces import * -from ...ragged import RaggedBatchWrapper - -from .container import MistralNonTransformerContainer, MistralTransformerContainer - - -class MistralInferenceModel(DSTransformerModelBase): - """ - Inference model implementation for ragged batching for Mistral models. - """ - - _non_transformer: Optional[MistralNonTransformerContainer] - """ - Embed + unembed container. Specializing the type annotation. - """ - - _transformer: Optional[Iterable[MistralTransformerContainer]] - """ - Per-layer transformer container. Specializing the type annotation. - """ - """ - Properties ineherited from `DSInferenceModelBase` - """ - - @property - def max_sequence_length(self) -> int: - return self._config.max_seq_length - - """ - Properties ineherited from `DSTransformerModelBase` - """ - - @property - def num_layers(self) -> int: - return self._config.num_hidden_layers - - @property - def model_dim(self) -> int: - return self._config.hidden_size - - @property - def vocab_size(self) -> int: - return self._config.vocab_size - - @property - def head_size(self) -> int: - return self.model_dim // self.n_heads - - @property - def n_heads(self) -> int: - return self._config.num_attention_heads - - @property - def intermediate_dim(self) -> int: - return self._config.intermediate_size - - @property - def n_heads_kv(self) -> int: - return self._config.num_key_value_heads - - @property - def activation_dtype(self) -> DtypeEnum: - if self._config.torch_dtype == torch.float16: - return DtypeEnum.fp16 - elif self._config.torch_dtype == torch.bfloat16: - return DtypeEnum.bf16 - else: - raise NotImplementedError("Only fp16 and bf16 are supported") - - @property - def mlp_activation_fn(self) -> ActivationType: - activation = self._config.hidden_act.lower() - if activation == "gelu": - return ActivationType.GEGLU - elif activation == "relu": - return ActivationType.ReGLU - elif activation == "gegelu": - return ActivationType.GEGLU - elif activation == "silu": - return ActivationType.SiGLU - else: - raise NotImplementedError(f"Activation {activation} not supported") - - @property - def norm_type(self) -> NormTypeEnum: - return NormTypeEnum.RMSNorm - - @property - def positional_embedding_type(self) -> PositionalEmbeddingType: - return PositionalEmbeddingType.rotate_half - - @property - def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - return RotateHalfConfig(theta_base=self._config.rope_theta) - - """ - Forward implementations - """ - - def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs the embedding lookup prior to running the transformer of the model. - - Arguments: - ragged_batch (RaggedBatchWrapper): The batch to embed. - - Returns: - torch.Tensor: The embedded batch. - """ - embed = self.embed(ragged_batch, self._non_transformer.word_emb) - - if embed.shape[-1] != self.model_dim: - raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}") - - return embed - - def _forward_transformer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor, - ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead - optimization to fuse the layer norm of the next layer into the current layer. - - Arguments: - layer_idx (int): The index of the layer to execute. - residual (torch.Tensor): The residual tensor from the previous layer. - hidden_states (torch.Tensor): The hidden states from the previous layer. This is the - hidden states after pre normalization. - ragged_batch_info (RaggedBatchWrapper): The batch metadata. - """ - # TODO(cmikeh2): Distribute ragged_batch_info to all modules - - cur_params = self._transformer[layer_idx] - kv_cache = self.state_manager.get_cache(layer_idx) - - hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=None) - hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info) - hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=None) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - residual, hidden_states = self.norm(residual, hidden_states, cur_params.mlp_norm_gamma, beta=None) - - # Should be configurable in the future - hidden_states = self.mlp_1(hidden_states, cur_params.mlp_1_w, b=None) - hidden_states = self.mlp_2(hidden_states, cur_params.mlp_2_w, b=None) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - if layer_idx != self.num_layers - 1: - next_params = self._transformer[layer_idx + 1] - residual, hidden_states = self.norm(residual, hidden_states, next_params.attn_norm_gamma, beta=None) - else: - # On last layer, we just need to perform the residual add. Adding into the residual - # here is safe. - residual.add_(hidden_states) - - return residual, hidden_states - - def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs unembedding of the hidden states to logits. This will only sample the final - token of each sequence. - """ - logits = self.unembed(hidden_states, - self._non_transformer.word_unembed, - ragged_batch_info, - gamma=self._non_transformer.final_norm) - - if self.tp_size > 1: - comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1])) - full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size)) - - dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group) - - full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size)) - - return full_logits - else: - return logits - - def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor: - - residual = self._forward_embed(wrapped_batch) - - residual, hidden_states = self.norm(residual, None, self._transformer[0].attn_norm_gamma, beta=None) - - for layer_idx in range(self.num_layers): - residual, hidden_states = self._forward_transformer(layer_idx, residual, hidden_states, wrapped_batch) - - return self._forward_unembed(residual, wrapped_batch) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/policy.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/policy.py deleted file mode 100644 index b67ec311c952df8f1de538952f4d922cbba6b2ba..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mistral/policy.py +++ /dev/null @@ -1,30 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any - -from ...config_v2 import RaggedInferenceEngineConfig -from ..inference_policy_base import ContainerMap, InferenceV2Policy -from .container import MistralNonTransformerContainer, MistralTransformerContainer -from .model import MistralInferenceModel - - -class MistralPolicy(InferenceV2Policy): - - def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> MistralInferenceModel: - return MistralInferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group) - - def build_container_map(self) -> ContainerMap: - map = ContainerMap() - - transformer_containers = [MistralTransformerContainer(self.model) for _ in range(self.model.num_layers)] - - map.set_transformer_params(['model.layers'], transformer_containers) - - map.set_non_transformer_params(MistralNonTransformerContainer(self.model)) - - map.set_unmapped_params([]) - - return map diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__init__.py deleted file mode 100644 index 2cb1aa889291d63dc93bba27494c42e11aa9139f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .policy import MixtralPolicy diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 19e3aaab4cd18cb6387749e2396fe4b93efced0f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__pycache__/container.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__pycache__/container.cpython-310.pyc deleted file mode 100644 index f071a010ec746ef57dea0d5824bd86ed0c6d5d70..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__pycache__/container.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__pycache__/model.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__pycache__/model.cpython-310.pyc deleted file mode 100644 index 26c43c1ae515fe80d5ec3ef5280d3a1e995d074f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__pycache__/model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__pycache__/policy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__pycache__/policy.cpython-310.pyc deleted file mode 100644 index 787e8d85ea3a34c2df2a9f8a13f353ec69e2c262..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/__pycache__/policy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/container.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/container.py deleted file mode 100644 index 6ec4a0552b8f53430083e045b913ef5db412f277..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/container.py +++ /dev/null @@ -1,46 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Create a container object to save model-specific tensors using the policy file above. - -from deepspeed.inference.v2.model_implementations.common_parameters import * -from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer - - -class MixtralTransformerContainer(LayerContainer): - - qkv_w: UnfusedQKVParameter - attn_out_w: AttentionOutputParameter - moe_gate: MoEGatingWeightParameter - moe_mlp_1: UnfusedMoEGatedMLPParameter - moe_mlp_2: UnfusedMoEMLP2Parameter - attn_norm_gamma: NormParameter - mlp_norm_gamma: NormParameter - - PARAM_MAPPING = { - "input_layernorm.weight": "attn_norm_gamma.params", - "post_attention_layernorm.weight": "mlp_norm_gamma.params", - "self_attn.q_proj.weight": "qkv_w.q_params", - "self_attn.k_proj.weight": "qkv_w.k_params", - "self_attn.v_proj.weight": "qkv_w.v_params", - "self_attn.o_proj.weight": "attn_out_w.params", - "block_sparse_moe.gate.weight": "moe_gate.params", - "block_sparse_moe.experts.*.w1.weight": "moe_mlp_1.gating_experts", - "block_sparse_moe.experts.*.w3.weight": "moe_mlp_1.up_experts", - "block_sparse_moe.experts.*.w2.weight": "moe_mlp_2.experts", - } - - -class MixtralNonTransformerContainer(LayerContainer): - - word_emb: EmbeddingParameter - word_unembed: UnembedParameter - final_norm: NormParameter - - PARAM_MAPPING = { - "model.embed_tokens.weight": "word_emb.params", - "lm_head.weight": "word_unembed.params", - "model.norm.weight": "final_norm.params", - } diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/model.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/model.py deleted file mode 100644 index 878cd8e31cec07f3955928e2ac49acdbe97b47b5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/model.py +++ /dev/null @@ -1,261 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Iterable, Optional, Tuple - -import torch - -import deepspeed.comm as dist - -from ...allocator import empty_from -from ...config_v2 import RaggedInferenceEngineConfig -from ...inference_utils import ActivationType, DtypeEnum -from ...model_implementations import * -from ...modules.configs import * -from ...modules.interfaces import * -from ...ragged import RaggedBatchWrapper -from ..inference_model_base import ( - DSModelImplementationConfig, - MPType, -) - -from .container import MixtralNonTransformerContainer, MixtralTransformerContainer - - -class MixtralInferenceModel(DSMoETransformerModelBase): - """ - Inference model implementation for Mixtral models. - """ - - _non_transformer: Optional[MixtralNonTransformerContainer] - """ - Embed + unembed container. Specializing the type annotation. - """ - - _transformer: Optional[Iterable[MixtralTransformerContainer]] - """ - Per-layer transformer container. Specializing the type annotation. - """ - """ - Properties ineherited from `DSInferenceModelBase` - """ - - @property - def max_sequence_length(self) -> int: - return self._config.max_position_embeddings - - """ - Properties ineherited from `DSTransformerModelBase` - """ - - @property - def num_layers(self) -> int: - return self._config.num_hidden_layers - - @property - def model_dim(self) -> int: - return self._config.hidden_size - - @property - def vocab_size(self) -> int: - return self._config.vocab_size - - @property - def head_size(self) -> int: - return self.model_dim // self.n_heads - - @property - def n_heads(self) -> int: - return self._config.num_attention_heads - - @property - def intermediate_dim(self) -> int: - return self._config.intermediate_size - - @property - def n_heads_kv(self) -> int: - return self._config.num_key_value_heads - - @property - def activation_dtype(self) -> DtypeEnum: - if self._config.torch_dtype == torch.float16: - return DtypeEnum.fp16 - elif self._config.torch_dtype == torch.bfloat16: - return DtypeEnum.bf16 - else: - raise NotImplementedError("Only fp16 and bf16 are supported") - - @property - def mlp_activation_fn(self) -> ActivationType: - activation = self._config.hidden_act.lower() - if activation == "gelu": - return ActivationType.GEGLU - elif activation == "relu": - return ActivationType.ReGLU - elif activation == "gegelu": - return ActivationType.GEGLU - elif activation == "silu": - return ActivationType.SiGLU - else: - raise NotImplementedError(f"Activation {activation} not supported") - - @property - def norm_type(self) -> NormTypeEnum: - return NormTypeEnum.RMSNorm - - @property - def positional_embedding_type(self) -> PositionalEmbeddingType: - return PositionalEmbeddingType.rotate_half - - @property - def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - """ - The positional embedding configuration for the model. - """ - return RotateHalfConfig(theta_base=self._config.rope_theta) - - """ - Inherited from `DSMoETransformerModelBase` - """ - - @property - def n_experts(self) -> int: - return self._config.num_local_experts - - @property - def n_top_k(self) -> int: - return self._config.num_experts_per_tok - - @property - def normalize_expert_scores(self) -> bool: - return True - - """ - Model implementation - """ - - def __init__(self, config: DSModelImplementationConfig, engine_config: RaggedInferenceEngineConfig, - base_mp_group: MPType) -> None: - """ - Base implementation for initialization. By default, this will initialize - the traditional components of a transformer model: - - Embedding - - QKV projection - - Self attention - - Attention output projection - - Feed forward network - - Normalization - - Unembedding - - Arguments: - config (DSModelImplementationConfig): Model-specific configuration. No assumptions - should be made about this config that are not closely tied to the specific - model implementation. - engine_config (RaggedInferenceEngineConfig): Engine configuration. - base_mp_group (MPType): Base communication group for Tensor-parallel inference. - """ - super().__init__(config, engine_config, base_mp_group) - - self.make_norm_layer() - self.make_qkv_layer() - self.make_attn_layer() - self.make_attn_out_layer() - self.make_moe_layer() - self.make_embedding_layer() - self.make_unembedding_layer() - self._kv_cache_config = None - - def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs the embedding lookup prior to running the transformer of the model. - - Arguments: - ragged_batch (RaggedBatchWrapper): The batch to embed. - - Returns: - torch.Tensor: The embedded batch. - """ - embed = self.embed(ragged_batch, self._non_transformer.word_emb) - - if embed.shape[-1] != self.model_dim: - raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}") - - return embed - - def _forward_transformer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor, - ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead - optimization to fuse the layer norm of the next layer into the current layer. - - Arguments: - layer_idx (int): The index of the layer to execute. - residual (torch.Tensor): The residual tensor from the previous layer. - hidden_states (torch.Tensor): The hidden states from the previous layer. This is the - hidden states after pre normalization. - ragged_batch_info (RaggedBatchWrapper): The batch metadata. - """ - # TODO(cmikeh2): Distribute ragged_batch_info to all modules - - cur_params = self._transformer[layer_idx] - kv_cache = self.state_manager.get_cache(layer_idx) - - hidden_states = self.qkv(hidden_states, cur_params.qkv_w) - hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info) - hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - residual, hidden_states = self.norm(residual, hidden_states, cur_params.mlp_norm_gamma) - - hidden_states = self.moe(hidden_states, ragged_batch_info, cur_params.moe_gate, cur_params.moe_mlp_1, - cur_params.moe_mlp_2) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - if layer_idx != self.num_layers - 1: - next_params = self._transformer[layer_idx + 1] - residual, hidden_states = self.norm(residual, hidden_states, next_params.attn_norm_gamma) - else: - # On last layer, we just need to perform the residual add. Adding into the residual - # here is safe. - residual.add_(hidden_states) - - return residual, hidden_states - - def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs unembedding of the hidden states to logits. This will only sample the final - token of each sequence. - """ - logits = self.unembed(hidden_states, - self._non_transformer.word_unembed, - ragged_batch_info, - gamma=self._non_transformer.final_norm) - - if self.tp_size > 1: - comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1])) - full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size)) - - dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group) - - full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size)) - - return full_logits - else: - return logits - - def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor: - - residual = self._forward_embed(wrapped_batch) - - residual, hidden_states = self.norm(residual, None, self._transformer[0].attn_norm_gamma, beta=None) - - for layer_idx in range(self.num_layers): - residual, hidden_states = self._forward_transformer(layer_idx, residual, hidden_states, wrapped_batch) - - return self._forward_unembed(residual, wrapped_batch) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/policy.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/policy.py deleted file mode 100644 index 2f0087919720d040ee53eda7723268f2635db585..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/policy.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any - -from ...config_v2 import RaggedInferenceEngineConfig -from ..inference_policy_base import ContainerMap, InferenceV2Policy -from .container import MixtralTransformerContainer, MixtralNonTransformerContainer -from .model import MixtralInferenceModel - - -class MixtralPolicy(InferenceV2Policy): - - def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> MixtralInferenceModel: - return MixtralInferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group) - - def build_container_map(self) -> ContainerMap: - - map = ContainerMap() - - transformer_containers = [MixtralTransformerContainer(self.model) for _ in range(self.model.num_layers)] - - map.set_transformer_params(['model.layers'], transformer_containers) - - map.set_non_transformer_params(MixtralNonTransformerContainer(self.model)) - - map.set_unmapped_params([]) - - return map diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__init__.py deleted file mode 100644 index c0f24d5243b820748a9ddb9a1380fd51644ae917..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .policy import OPTPolicy diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 94aa0380b971ad054de735929206248ee86adb32..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__pycache__/container.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__pycache__/container.cpython-310.pyc deleted file mode 100644 index 29d14a0e9c02bba3b103554631a08581311e6afd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__pycache__/container.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__pycache__/model.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__pycache__/model.cpython-310.pyc deleted file mode 100644 index c0e3b7e1236b4cc8ed92f5c8f10eace7f4917fee..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__pycache__/model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__pycache__/policy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__pycache__/policy.cpython-310.pyc deleted file mode 100644 index 166bd55de72279de3643bc6f8500407a4cbdd616..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/__pycache__/policy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/container.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/container.py deleted file mode 100644 index e97599ef8e50e3dc30ce5540bd4721bb055df87e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/container.py +++ /dev/null @@ -1,94 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Create a container object to save model-specific tensors using the policy file above. - -from ..common_parameters import * -from ..layer_container_base import LayerContainer -''' - # HF OPT model looks like this: - -OPTForCausalLM( - (model): OPTModel( - (decoder): OPTDecoder( - (embed_tokens): Embedding(50272, 768, padding_idx=1) - (embed_positions): OPTLearnedPositionalEmbedding(2050, 768) - (final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True) - (layers): ModuleList( - (0-11): 12 x OPTDecoderLayer( - (self_attn): OPTAttention( - (k_proj): Linear(in_features=768, out_features=768, bias=True) - (v_proj): Linear(in_features=768, out_features=768, bias=True) - (q_proj): Linear(in_features=768, out_features=768, bias=True) - (out_proj): Linear(in_features=768, out_features=768, bias=True) - ) - (activation_fn): ReLU() - (self_attn_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True) - (fc1): Linear(in_features=768, out_features=3072, bias=True) - (fc2): Linear(in_features=3072, out_features=768, bias=True) - (final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True) - ) - ) - ) - ) - (lm_head): Linear(in_features=768, out_features=50272, bias=False) -) - -''' - - -class OPTTransformerContainer(LayerContainer): - """ - Transformer layer container for the OPT model. - """ - qkv_w: UnfusedQKVParameter - qkv_b: UnfusedQKVParameter - attn_out_w: AttentionOutputParameter - attn_out_b: AttentionOutputParameter - mlp_1_w: MLP1Parameter - mlp_1_b: MLP1Parameter - mlp_2_w: MLP2Parameter - mlp_2_b: MLP2Parameter - attn_norm_beta: NormParameter - attn_norm_gamma: NormParameter - mlp_norm_beta: NormParameter - mlp_norm_gamma: NormParameter - - PARAM_MAPPING = { - "self_attn.q_proj.weight": "qkv_w.q_params", - "self_attn.q_proj.bias": "qkv_b.q_params", - "self_attn.k_proj.weight": "qkv_w.k_params", - "self_attn.k_proj.bias": "qkv_b.k_params", - "self_attn.v_proj.weight": "qkv_w.v_params", - "self_attn.v_proj.bias": "qkv_b.v_params", - "self_attn.out_proj.weight": "attn_out_w.params", - "self_attn.out_proj.bias": "attn_out_b.params", - "fc1.weight": "mlp_1_w.params", - "fc1.bias": "mlp_1_b.params", - "fc2.weight": "mlp_2_w.params", - "fc2.bias": "mlp_2_b.params", - "self_attn_layer_norm.weight": "attn_norm_gamma.params", - "self_attn_layer_norm.bias": "attn_norm_beta.params", - "final_layer_norm.weight": "mlp_norm_gamma.params", - "final_layer_norm.bias": "mlp_norm_beta.params", - } - - -class OPTNonTransformerContainer(LayerContainer): - """ - Non-Transformer layer container for the OPT model. - """ - word_emb: EmbeddingParameter - word_emb_pos: EmbeddingParameter - word_unembed: UnembedParameter - final_norm_w: NormParameter - final_norm_b: NormParameter - - PARAM_MAPPING = { - "*decoder.embed_tokens.weight": ["word_emb.params", "word_unembed.params"], - "*decoder.embed_positions.weight": "word_emb_pos.params", - "*decoder.final_layer_norm.weight": "final_norm_w.params", - "*decoder.final_layer_norm.bias": "final_norm_b.params", - } diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/model.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/model.py deleted file mode 100644 index adf011d8f1a7884f47569ee186336dcc77b355bb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/model.py +++ /dev/null @@ -1,197 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Iterable, Optional, Tuple - -import torch - -import deepspeed.comm as dist - -from ...allocator import empty_from -from ...inference_utils import ActivationType, DtypeEnum -from ...model_implementations import * -from ...modules.configs import * -from ...ragged import RaggedBatchWrapper -from .container import OPTNonTransformerContainer, OPTTransformerContainer - -from ...modules.heuristics import instantiate_embed - - -class OPTInferenceModel(DSTransformerModelBase): - """ - Inference model implementation for ragged batching for OPT models. - """ - - _non_transformer: Optional[OPTNonTransformerContainer] - """ - Embed + unembed container. Specializing the type annotation. - """ - - _transformer: Optional[Iterable[OPTTransformerContainer]] - """ - Per-layer transformer container. Specializing the type annotation. - """ - """ - Properties ineherited from `DSInferenceModelBase` - """ - - @property - def max_sequence_length(self) -> int: - return self._config.max_seq_length - - """ - Properties ineherited from `DSTransformerModelBase` - """ - - @property - def num_layers(self) -> int: - return self._config.num_hidden_layers - - @property - def model_dim(self) -> int: - return self._config.hidden_size - - @property - def vocab_size(self) -> int: - return self._config.vocab_size - - @property - def head_size(self) -> int: - return self.model_dim // self.n_heads - - @property - def n_heads(self) -> int: - return self._config.num_attention_heads - - @property - def intermediate_dim(self) -> int: - return self._config.ffn_dim - - @property - def activation_dtype(self) -> DtypeEnum: - if self._config.torch_dtype == torch.float16: - return DtypeEnum.fp16 - elif self._config.torch_dtype == torch.bfloat16: - return DtypeEnum.bf16 - else: - raise NotImplementedError("Only fp16 and bf16 are supported") - - @property - def mlp_activation_fn(self) -> ActivationType: - return ActivationType.RELU - - @property - def norm_type(self) -> NormTypeEnum: - return NormTypeEnum.LayerNorm - - @property - def positional_embedding_type(self) -> PositionalEmbeddingType: - return PositionalEmbeddingType.none - - @property - def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - return None - - """ - Overrides of ``DSTransformerModelBase`` methods - """ - - def make_embedding_layer(self) -> None: - """ - Performs setup and creates embedding DSModule. Since OPT includes trained - positional embeddings, we will override the base model implementation. - """ - - embed_config = DSEmbeddingsConfig(max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - residual_dtype=self.activation_dtype, - embedding_dim=self.model_dim, - positional_embedding=True, - positional_offset=2) - - self.embed = instantiate_embed(embed_config, self._engine_config) - - """ - Forward implementations - """ - - def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor: - embed = self.embed(ragged_batch, self._non_transformer.word_emb, self._non_transformer.word_emb_pos) - if embed.shape[-1] != self.model_dim: - raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}") - - return embed - - def _forward_transformer_layer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor, - ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]: - # TODO(cmikeh2): Distribute ragged_batch_info to all modules - - cur_params = self._transformer[layer_idx] - kv_cache = self.state_manager.get_cache(layer_idx) - - hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=cur_params.qkv_b) - hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info) - hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=cur_params.attn_out_b) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - residual, hidden_states = self.norm(residual, - hidden_states, - cur_params.mlp_norm_gamma, - beta=cur_params.mlp_norm_beta) - - # Should be configurable in the future - hidden_states = self.mlp_1(hidden_states, cur_params.mlp_1_w, b=cur_params.mlp_1_b) - hidden_states = self.mlp_2(hidden_states, cur_params.mlp_2_w, b=cur_params.mlp_2_b) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - if layer_idx != self.num_layers - 1: - next_params = self._transformer[layer_idx + 1] - residual, hidden_states = self.norm(residual, - hidden_states, - next_params.attn_norm_gamma, - beta=next_params.attn_norm_beta) - else: - # On last layer, we just need to perform the residual add. Adding into the residual - # here is safe. - residual.add_(hidden_states) - - return residual, hidden_states - - def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor: - logits = self.unembed(hidden_states, - self._non_transformer.word_unembed, - ragged_batch_info, - gamma=self._non_transformer.final_norm_w, - beta=self._non_transformer.final_norm_b) - - if self.tp_size > 1: - comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1])) - full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size)) - - dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group) - - full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size)) - - return full_logits - else: - return logits - - def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor: - - residual = self._forward_embed(wrapped_batch) - - residual, hidden_states = self.norm(residual, - None, - self._transformer[0].attn_norm_gamma, - beta=self._transformer[0].attn_norm_beta) - - for layer_idx in range(self.num_layers): - residual, hidden_states = self._forward_transformer_layer(layer_idx, residual, hidden_states, - wrapped_batch) - - return self._forward_unembed(residual, wrapped_batch) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/policy.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/policy.py deleted file mode 100644 index d57d5beb48d515fbd6d4c5a0741edf1a4867136f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/opt/policy.py +++ /dev/null @@ -1,30 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any - -from ...config_v2 import RaggedInferenceEngineConfig -from ..inference_policy_base import ContainerMap, InferenceV2Policy -from .container import OPTNonTransformerContainer, OPTTransformerContainer -from .model import OPTInferenceModel - - -class OPTPolicy(InferenceV2Policy): - - def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> OPTInferenceModel: - return OPTInferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group) - - def build_container_map(self) -> ContainerMap: - map = ContainerMap() - - transformer_containers = [OPTTransformerContainer(self.model) for _ in range(self.model.num_layers)] - - map.set_transformer_params(['model.decoder.layers', 'decoder.layers'], transformer_containers) - - map.set_non_transformer_params(OPTNonTransformerContainer(self.model)) - - map.set_unmapped_params(['lm_head.weight']) - - return map diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/parameter_base.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/parameter_base.py deleted file mode 100644 index e30eedbdcf0c6087567ce522aa085306bcf3bdd7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/parameter_base.py +++ /dev/null @@ -1,256 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import weakref -from abc import abstractmethod -from typing import Type - -import torch - -# Currently have dependency loops for the type hints. -InferenceModel = Type["InferenceModel"] -LayerContainer = Type["LayerContainer"] - -MAPPING_KEY = "PARAM_MAPPING" - - -def make_param_getter(clsname, param): - """ - Normal getter implementation for a property. - """ - - def param_getter(self): - return getattr(self, f"__{clsname}__{param}") - - return param_getter - - -def make_param_setter(clsname, param): - """ - Setter implementation that will call complete component to potentially - finalize the parameter. - """ - - def param_setter(self, value): - setattr(self, f"__{clsname}__{param}", value) - self.dtype = value.dtype - self.complete_component() - - return param_setter - - -def make_readonly_setter(): - """ - Setter implementation that will raise an error if called. - """ - - def paramlist_setter(self, value): - raise ValueError("Cannot set a ParametrizedList directly.") - - return paramlist_setter - - -class ParameterMetaclass(type): - """ - MetaClass for the ParameterBase base class. This class will parse the `src_params` - attribute and create properties for each of the dependencies. A dependency can either - be represented as a string, which is interpreted as a named Tensor, or a `ParametrizedList` - subclass. - """ - - def __new__(cls, clsname, bases, attrs): - - annotations = attrs.get("__annotations__", {}) - dependencies = { - name: annotation - for name, annotation in annotations.items() if issubclass(annotation, (torch.Tensor, ParametrizedList)) - } - n_dependencies = len(dependencies) - - # Create properties for each of our dependencies - for d_name, d_type in dependencies.items(): - if issubclass(d_type, ParametrizedList): - assert hasattr( - d_type, "count_attr" - ), "ParametrizedList must have a count_attr attribute to access on the inference module." - attrs[d_name] = property(make_param_getter(clsname, d_name), make_readonly_setter()) - else: # torch.Tensor - attrs[d_name] = property(make_param_getter(clsname, d_name), make_param_setter(clsname, d_name)) - - new_cls = super().__new__(cls, clsname, bases, attrs) - new_cls.n_dependencies = n_dependencies - - return new_cls - - def __call__(cls, *args, **kwargs): - new_obj = super().__call__(*args, **kwargs) - new_obj.__init__(*args, **kwargs) - - setattr(new_obj, "dest_param", None) - - # Initialize our dependences to None/empty `ParametrizedList`s - for name, annotation in new_obj.__annotations__.items(): - if issubclass(annotation, ParametrizedList): - #TODO(jeff): update assert with this, model implementation attribute does not align or missing wrt the ParametrizedList attributes - assert hasattr( - new_obj.inference_model, annotation.count_attr - ), f"new_obj={new_obj.__class__.__name__}, name={name}, annotation.count_attr={annotation.count_attr}" - param_list = annotation(new_obj, getattr(new_obj.inference_model, annotation.count_attr)) - setattr(new_obj, f"__{new_obj.__class__.__name__}__{name}", param_list) - else: # torch.Tensor - setattr(new_obj, f"__{new_obj.__class__.__name__}__{name}", None) - - return new_obj - - -class ParameterBase(metaclass=ParameterMetaclass): - """ - A ParameterBase allows us to consolidate tracking the dependencies of loading a parameter from - a checkpoint into a single object. This class should not be used directly, but rather subclassed - and the `src_params` attribute set to a list of strings and/or `ParametrizedList`s. - """ - - # inference_model: InferenceModel - """ - Inference model that will provide context on how to shard and transform the parameter. - """ - - #completed_components: int - """ - How many of the layer dependencies have been met. This is used to determine when the parameter - is ready to be finalized. A ParametrizedList counts as a single dependency for the purposes - of this counter. - """ - - def __init__(self, model: InferenceModel, parent_container: LayerContainer) -> None: - """ - Direct constructor. This should not be called from client code. - - Args: - model (InferenceModel): Inference model that will be used to shard and transform the - parameter in `finalize`. - parent_container (LayerContainer): The parent container that this parameter is a member - of. We will build a weakref to this container to call the finalization callback. - """ - self.inference_model = model - self.completed_components = 0 - self.parent_container = weakref.ref(parent_container) - - @abstractmethod - def finalize(self) -> torch.Tensor: - """ - Finalize the parameter after all of its source parameters have been set. This method - will be automatically called when all inputs have been set. It should return the Tensor - with all transformations performed on it. - """ - pass - - def complete_component(self) -> None: - """ - Mark a component as completed. This should be called by the relevant setter of a direct - property or a ParametrizedList. This method will automatically call `finalize` when all - dependencies have been met and then call the finalization callback on the parent container. - - Once the finalization callback has been called, the parameter will be replaced with the - `dst_param` attribute on the parent container, and this instance will be destroyed. - """ - self.completed_components += 1 - - if self.completed_components != self.n_dependencies: - return - - finalized_param = self.finalize() - self.parent_container().finalization_callback(self, finalized_param) - - -class ParametrizedList: - """ - A ParametrizedList is a list of parameters that are dependencies - of a `ParameterBase` but may vary in length depending on the model - configuration (rather than architecture). For example, a MoE layer - may have different number of experts depending on the size of the model. - - This class is used to manage these lists and provide integer indexing - of a single component rather than accessing names directly. For example, - it tends to be more natural to access the 8th expert with `experts[8]` - rather than a name like `expert_8`, especially as an attribute. - - To inherit from this class, set static variables `name` and `count_attr`. - - ```python - class MyParametrizedList(ParametrizedList): - count_attr: str = "my_list_count" - ``` - - In the above example, `my_list_count` should be an accessible attribute - of the inference model (i.e. via `self.inference_model.my_list_count`). - - NOTE: There are some APIs in which this type cannot be used as if it is - just a list of Tensors. For example, `torch.cat(param_list)` will not work. - However, you can make it compatible with a tuple wrapper: - `torch.cat(tuple(param_list))` - """ - - n_params: int - """ - Number of params this list contains. - """ - - param: ParameterBase - """ - WeakRef to the owning parameter. - """ - - def __init__(self, param: ParameterBase, n_params: int) -> None: - """ - Constructor. Should not be called from client code. - - Args: - param (ParameterBase): The owning parameter. - n_params (int): The number of parameters this list contains. This should be - """ - self.n_params = n_params - self.set_params = 0 - self.param = weakref.ref(param) - self._params = [None] * n_params - - def __getitem__(self, index): - return self._params[index] - - def __setitem__(self, index, value): - if self._params[index] is not None: - raise ValueError("Cannot set a parameter twice.") - - self._params[index] = value - self.set_params += 1 - - if self.set_params != self.n_params: - return - - self.param().complete_component() - - def __iter__(self): - return iter(self._params) - - -def ParamList(attr: str): - """ - Helper to create a subclass of ParametrizedList with the desired `count_attr`. - - In this manner, we can annotate the type of a Parameter dependency with the - following: - - ```python - class CustomParameter(ParameterBase): - dependency_list: ParamList("dependencies_count_name") - ``` - - where "dependencies_count_name" is the name of the attribute on the inference model. - """ - - class ParametrizedListInstance(ParametrizedList): - count_attr: str = attr - - return ParametrizedListInstance diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__init__.py deleted file mode 100644 index 3ab107e75a9147f30176f9e3f6bc575898d8e572..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .policy import PhiPolicy diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 58d7f7869503414c07f2df86dea35323ce4ab536..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__pycache__/containers.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__pycache__/containers.cpython-310.pyc deleted file mode 100644 index 04e1be5c039f39f085fc2555f879c43ba7667c26..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__pycache__/containers.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__pycache__/model.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__pycache__/model.cpython-310.pyc deleted file mode 100644 index 4f328d06cd5015fd1b5e33d61ecd197c88e8d7c6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__pycache__/model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__pycache__/policy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__pycache__/policy.cpython-310.pyc deleted file mode 100644 index 98bd67fbce1a82743d63a56ca4e7dad295d3f8ca..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/__pycache__/policy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/containers.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/containers.py deleted file mode 100644 index 21f07eb8c99a037243086688b551bbe3fa9ec3d6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/containers.py +++ /dev/null @@ -1,91 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Create a container object to save model-specific tensors using the policy file above. - -from ..common_parameters import * -from ..layer_container_base import LayerContainer -''' - # HF Phi-2 model looks like this: - -PhiForCausalLM( - (model): PhiModel( - (embed_tokens): Embedding(51200, 2560) - (embed_dropout): Dropout(p=0.0, inplace=False) - (layers): ModuleList( - (0-31): 32 x PhiDecoderLayer( - (self_attn): PhiAttention( - (q_proj): Linear(in_features=2560, out_features=2560, bias=True) - (k_proj): Linear(in_features=2560, out_features=2560, bias=True) - (v_proj): Linear(in_features=2560, out_features=2560, bias=True) - (dense): Linear(in_features=2560, out_features=2560, bias=True) - (rotary_emb): PhiRotaryEmbedding() - ) - (mlp): PhiMLP( - (activation_fn): NewGELUActivation() - (fc1): Linear(in_features=2560, out_features=10240, bias=True) - (fc2): Linear(in_features=10240, out_features=2560, bias=True) - ) - (input_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True) - (resid_dropout): Dropout(p=0.1, inplace=False) - ) - ) - (final_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True) - ) - (lm_head): Linear(in_features=2560, out_features=51200, bias=True) -) -''' - - -class PhiTransformerContainer(LayerContainer): - """ - Transformer layer container for the Phi model. - """ - qkv_w: UnfusedQKVParameter - qkv_b: UnfusedQKVParameter - attn_out_w: AttentionOutputParameter - attn_out_b: AttentionOutputParameter - mlp_1_w: MLP1Parameter - mlp_1_b: MLP1Parameter - mlp_2_w: MLP2Parameter - mlp_2_b: MLP2Parameter - ln_gamma: NormParameter - ln_beta: NormParameter - - PARAM_MAPPING = { - "self_attn.q_proj.weight": "qkv_w.q_params", - "self_attn.k_proj.weight": "qkv_w.k_params", - "self_attn.v_proj.weight": "qkv_w.v_params", - "self_attn.q_proj.bias": "qkv_b.q_params", - "self_attn.k_proj.bias": "qkv_b.k_params", - "self_attn.v_proj.bias": "qkv_b.v_params", - "self_attn.dense.weight": "attn_out_w.params", - "self_attn.dense.bias": "attn_out_b.params", - "mlp.fc1.weight": "mlp_1_w.params", - "mlp.fc1.bias": "mlp_1_b.params", - "mlp.fc2.weight": "mlp_2_w.params", - "mlp.fc2.bias": "mlp_2_b.params", - "input_layernorm.weight": "ln_gamma.params", - "input_layernorm.bias": "ln_beta.params", - } - - -class PhiNonTransformerContainer(LayerContainer): - """ - Non-Transformer layer container for the Phi model. - """ - word_emb: EmbeddingParameter - word_unembed_w: UnembedParameter - word_unembed_b: UnembedParameter - final_norm_gamma: NormParameter - final_norm_beta: NormParameter - - PARAM_MAPPING = { - "model.embed_tokens.weight": "word_emb.params", - "model.final_layernorm.weight": "final_norm_gamma.params", - "model.final_layernorm.bias": "final_norm_beta.params", - "lm_head.weight": "word_unembed_w.params", - "lm_head.bias": "word_unembed_b.params", - } diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/model.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/model.py deleted file mode 100644 index 2d5826810cb57bfb7fe9b7df75edf1cf947bd1f8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/model.py +++ /dev/null @@ -1,199 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Iterable, Optional, Tuple - -import torch - -import deepspeed.comm as dist - -from ...allocator import empty_from -from ...inference_utils import ActivationType, DtypeEnum -from .. import * -from ...modules.configs import * -from ...modules.interfaces import * -from ...ragged import RaggedBatchWrapper - -from .containers import PhiNonTransformerContainer, PhiTransformerContainer - - -class PhiInferenceModel(DSTransformerModelBase): - """ - Inference model implementation for ragged batching for Llama-2 models. - """ - - _non_transformer: Optional[PhiNonTransformerContainer] - """ - Embed + unembed container. Specializing the type annotation. - """ - - _transformer: Optional[Iterable[PhiTransformerContainer]] - """ - Per-layer transformer container. Specializing the type annotation. - """ - """ - Properties inherited from `DSInferenceModelBase` - """ - - @property - def max_sequence_length(self) -> int: - return self._config.max_seq_length - - """ - Properties inherited from `DSTransformerModelBase` - """ - - @property - def num_layers(self) -> int: - return self._config.num_hidden_layers - - @property - def model_dim(self) -> int: - return self._config.hidden_size - - @property - def vocab_size(self) -> int: - return self._config.vocab_size - - @property - def head_size(self) -> int: - return self.model_dim // self.n_heads - - @property - def n_heads(self) -> int: - return self._config.num_attention_heads - - @property - def intermediate_dim(self) -> int: - return self._config.intermediate_size - - @property - def n_heads_kv(self) -> int: - return self._config.num_key_value_heads - - @property - def activation_dtype(self) -> DtypeEnum: - if self._config.torch_dtype == torch.float16: - return DtypeEnum.fp16 - elif self._config.torch_dtype == torch.bfloat16: - return DtypeEnum.bf16 - else: - raise NotImplementedError("Only fp16 and bf16 are supported") - - @property - def mlp_activation_fn(self) -> ActivationType: - return ActivationType.GELU - - @property - def norm_type(self) -> NormTypeEnum: - return NormTypeEnum.LayerNorm - - @property - def positional_embedding_type(self) -> PositionalEmbeddingType: - return PositionalEmbeddingType.rotate_half - - @property - def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - rotary_dim = int(self._config.partial_rotary_factor * self.head_size) - return RotateHalfConfig(rotate_dim=rotary_dim, theta_base=self._config.rope_theta) - - """ - Forward implementations - """ - - def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs the embedding lookup prior to running the transformer of the model. - - Arguments: - ragged_batch (RaggedBatchWrapper): The batch to embed. - - Returns: - torch.Tensor: The embedded batch. - """ - embed = self.embed(ragged_batch, self._non_transformer.word_emb) - - if embed.shape[-1] != self.model_dim: - raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}") - - return embed - - def _forward_transformer_layer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor, - ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead - optimization to fuse the layer norm of the next layer into the current layer. - - Arguments: - layer_idx (int): The index of the layer to execute. - residual (torch.Tensor): The residual tensor from the previous layer. - hidden_states (torch.Tensor): The hidden states from the previous layer. This is the - hidden states after pre normalization. - ragged_batch_info (RaggedBatchWrapper): The batch metadata. - """ - cur_params = self._transformer[layer_idx] - kv_cache = self.state_manager.get_cache(layer_idx) - - attn_ln_out = hidden_states - attn_hidden_state = self.qkv(attn_ln_out, cur_params.qkv_w, b=cur_params.qkv_b) - attn_hidden_state = self.attn(attn_hidden_state, kv_cache, ragged_batch_info) - attention_output = self.attn_out(attn_hidden_state, cur_params.attn_out_w, b=cur_params.attn_out_b) - - mlp_ln_out = hidden_states - mlp_hidden_state = self.mlp_1(mlp_ln_out, cur_params.mlp_1_w, b=cur_params.mlp_1_b) - mlp_output = self.mlp_2(mlp_hidden_state, cur_params.mlp_2_w, b=cur_params.mlp_2_b) - - mlp_output.add_(attention_output) - - if self.tp_size > 1: - dist.all_reduce(mlp_output, group=self._base_mp_group) - - if layer_idx != self.num_layers - 1: - next_params = self._transformer[layer_idx + 1] - residual, mlp_output = self.norm(residual, mlp_output, next_params.ln_gamma, beta=next_params.ln_beta) - else: - # On last layer, we just need to perform the residual add. Adding into the residual - # here is safe. - residual.add_(mlp_output) - - return residual, mlp_output - - def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs unembedding of the hidden states to logits. This will only sample the final - token of each sequence. - """ - logits = self.unembed(hidden_states, - self._non_transformer.word_unembed_w, - ragged_batch_info, - bias=self._non_transformer.word_unembed_b, - gamma=self._non_transformer.final_norm_gamma, - beta=self._non_transformer.final_norm_beta) - - if self.tp_size > 1: - comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1])) - full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size)) - - dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group) - - full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size)) - - return full_logits - else: - return logits - - def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor: - residual = self._forward_embed(wrapped_batch) - - residual, hidden_states = self.norm(residual, - None, - gamma=self._transformer[0].ln_gamma, - beta=self._transformer[0].ln_beta) - - for layer_idx in range(self.num_layers): - residual, hidden_states = self._forward_transformer_layer(layer_idx, residual, hidden_states, - wrapped_batch) - - return self._forward_unembed(residual, wrapped_batch) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/policy.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/policy.py deleted file mode 100644 index 4b081a8e61bde9304c8f1bee920bb8729f1c6aef..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi/policy.py +++ /dev/null @@ -1,32 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any - -from ...config_v2 import RaggedInferenceEngineConfig -from ..inference_policy_base import ContainerMap, InferenceV2Policy -from .containers import PhiNonTransformerContainer, PhiTransformerContainer -from .model import PhiInferenceModel - - -class PhiPolicy(InferenceV2Policy): - - def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> PhiInferenceModel: - return PhiInferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group) - - def build_container_map(self) -> ContainerMap: - map = ContainerMap() - - trans_container_cls = PhiTransformerContainer - transformer_containers = [trans_container_cls(self.model) for _ in range(self.model.num_layers)] - - map.set_transformer_params(['model.layers'], transformer_containers) - - map.set_non_transformer_params(PhiNonTransformerContainer(self.model)) - - map.set_unmapped_params( - [f'model.layers.{i}.self_attn.rotary_emb.inv_freq' for i in range(self.model.num_layers)]) - - return map diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__init__.py deleted file mode 100644 index 1a4b756d210c5e336cec1053e6c0533f69afd461..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .policy import Phi3Policy diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index f2df5ad7aeff5725854a5ebfe1a56a217a635cc9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__pycache__/containers.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__pycache__/containers.cpython-310.pyc deleted file mode 100644 index a6041352b215848c76901cc962fc22e5c4865be1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__pycache__/containers.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__pycache__/model.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__pycache__/model.cpython-310.pyc deleted file mode 100644 index 266dac826d0c8d1868249309893120c64dc5a0f8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__pycache__/model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__pycache__/policy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__pycache__/policy.cpython-310.pyc deleted file mode 100644 index e0bc428845b6a7488e7270505bdd599d63e8b68f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/__pycache__/policy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/containers.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/containers.py deleted file mode 100644 index 1cb52a75ae0b8a6c26a92db651ba881509df9405..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/containers.py +++ /dev/null @@ -1,75 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Create a container object to save model-specific tensors using the policy file above. - -from ..common_parameters import * -from ..layer_container_base import LayerContainer -''' - # HF Phi-3 model looks like this: - -Phi3ForCausalLM( - (model): Phi3Model( - (embed_tokens): Embedding(32064, 3072) - (embed_dropout): Dropout(p=0.0, inplace=False) - (layers): ModuleList( - (0-31): 32 x Phi3DecoderLayer( - (self_attn): Phi3Attention( - (o_proj): Linear(in_features=3072, out_features=3072, bias=False) - (qkv_proj): Linear(in_features=3072, out_features=9216, bias=False) - (rotary_emb): Phi3RotaryEmbedding() - ) - (mlp): PhiMLP( - (gate_up_proj): Linear(in_features=3072, out_features=16384, bias=False) - (down_proj): Linear(in_features=16384, out_features=3072, bias=False) - (activation_fn): SiLU() - ) - (input_layernorm): Phi3RMSNorm((3072,), eps=1e-05) - (resid_attn_dropout): Dropout(p=0.0) - (resid_mlp_dropout): Dropout(p=0.0) - (post_attention_layernorm): Phi3RMSNorm((3072,), eps=1e-05) - ) - ) - (final_layernorm): Phi3RMSNorm((3072,), eps=1e-05) - ) - (lm_head): Linear(in_features=3072, out_features=32064, bias=False) -) -''' - - -class Phi3TransformerContainer(LayerContainer): - """ - Transformer layer container for the Phi model. - """ - qkv_w: FusedQKVParameter - attn_out_w: AttentionOutputParameter - mlp_1_w: FusedGatedMLPParameter - mlp_2_w: MLP2Parameter - attn_norm_gamma: NormParameter - mlp_norm_gamma: NormParameter - - PARAM_MAPPING = { - "self_attn.qkv_proj.weight": "qkv_w.params", - "self_attn.o_proj.weight": "attn_out_w.params", - "mlp.gate_up_proj.weight": "mlp_1_w.params", - "mlp.down_proj.weight": "mlp_2_w.params", - "input_layernorm.weight": "attn_norm_gamma.params", - "post_attention_layernorm.weight": "mlp_norm_gamma.params", - } - - -class Phi3NonTransformerContainer(LayerContainer): - """ - Non-Transformer layer container for the Phi model. - """ - word_emb: EmbeddingParameter - word_unembed_w: UnembedParameter - final_norm_gamma: NormParameter - - PARAM_MAPPING = { - "model.embed_tokens.weight": "word_emb.params", - "model.norm.weight": "final_norm_gamma.params", - "lm_head.weight": "word_unembed_w.params", - } diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/model.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/model.py deleted file mode 100644 index 507bb4fc9af1a7755b281c3df5b579c00abb98ff..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/model.py +++ /dev/null @@ -1,204 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Iterable, Optional, Tuple - -import torch - -import deepspeed.comm as dist - -from ...allocator import empty_from -from ...inference_utils import ActivationType, DtypeEnum -from .. import * -from ...modules.configs import * -from ...modules.interfaces import * -from ...ragged import RaggedBatchWrapper - -from .containers import Phi3NonTransformerContainer, Phi3TransformerContainer - - -class Phi3InferenceModel(DSTransformerModelBase): - """ - Inference model implementation for ragged batching for Llama-2 models. - """ - - _non_transformer: Optional[Phi3NonTransformerContainer] - """ - Embed + unembed container. Specializing the type annotation. - """ - - _transformer: Optional[Iterable[Phi3TransformerContainer]] - """ - Per-layer transformer container. Specializing the type annotation. - """ - """ - Properties inherited from `DSInferenceModelBase` - """ - - @property - def max_sequence_length(self) -> int: - return self._config.max_seq_length - - """ - Properties inherited from `DSTransformerModelBase` - """ - - @property - def num_layers(self) -> int: - return self._config.num_hidden_layers - - @property - def model_dim(self) -> int: - return self._config.hidden_size - - @property - def vocab_size(self) -> int: - return self._config.vocab_size - - @property - def head_size(self) -> int: - return self.model_dim // self.n_heads - - @property - def n_heads(self) -> int: - return self._config.num_attention_heads - - @property - def intermediate_dim(self) -> int: - return self._config.intermediate_size - - @property - def n_heads_kv(self) -> int: - return self._config.num_key_value_heads - - @property - def activation_dtype(self) -> DtypeEnum: - if self._config.torch_dtype == torch.float16: - return DtypeEnum.fp16 - elif self._config.torch_dtype == torch.bfloat16: - return DtypeEnum.bf16 - else: - raise NotImplementedError("Only fp16 and bf16 are supported") - - @property - def mlp_activation_fn(self) -> ActivationType: - activation = self._config.hidden_act.lower() - if activation == "gelu": - return ActivationType.GEGLU - elif activation == "relu": - return ActivationType.ReGLU - elif activation == "gegelu": - return ActivationType.GEGLU - elif activation == "silu": - return ActivationType.SiGLU - else: - raise NotImplementedError(f"Activation {activation} not supported") - - @property - def norm_type(self) -> NormTypeEnum: - return NormTypeEnum.RMSNorm - - @property - def positional_embedding_type(self) -> PositionalEmbeddingType: - return PositionalEmbeddingType.rotate_half - - @property - def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - return RotateHalfConfig(theta_base=self._config.rope_theta) - - """ - Forward implementations - """ - - def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs the embedding lookup prior to running the transformer of the model. - - Arguments: - ragged_batch (RaggedBatchWrapper): The batch to embed. - - Returns: - torch.Tensor: The embedded batch. - """ - embed = self.embed(ragged_batch, self._non_transformer.word_emb) - - if embed.shape[-1] != self.model_dim: - raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}") - - return embed - - def _forward_transformer_layer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor, - ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead - optimization to fuse the layer norm of the next layer into the current layer. - - Arguments: - layer_idx (int): The index of the layer to execute. - residual (torch.Tensor): The residual tensor from the previous layer. - hidden_states (torch.Tensor): The hidden states from the previous layer. This is the - hidden states after pre normalization. - ragged_batch_info (RaggedBatchWrapper): The batch metadata. - """ - cur_params = self._transformer[layer_idx] - kv_cache = self.state_manager.get_cache(layer_idx) - - hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=None) - hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info) - hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=None) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - residual, hidden_states = self.norm(residual, hidden_states, cur_params.mlp_norm_gamma, beta=None) - - hidden_states = self.mlp_1(hidden_states, cur_params.mlp_1_w, b=None) - hidden_states = self.mlp_2(hidden_states, cur_params.mlp_2_w, b=None) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - if layer_idx != self.num_layers - 1: - next_params = self._transformer[layer_idx + 1] - residual, hidden_states = self.norm(residual, hidden_states, next_params.attn_norm_gamma, beta=None) - else: - # On last layer, we just need to perform the residual add. Adding into the residual - # here is safe. - residual.add_(hidden_states) - - return residual, hidden_states - - def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs unembedding of the hidden states to logits. This will only sample the final - token of each sequence. - """ - logits = self.unembed(hidden_states, - self._non_transformer.word_unembed_w, - ragged_batch_info, - gamma=self._non_transformer.final_norm_gamma) - - if self.tp_size > 1: - comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1])) - full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size)) - - dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group) - - full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size)) - - return full_logits - else: - return logits - - def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor: - residual = self._forward_embed(wrapped_batch) - - residual, hidden_states = self.norm(residual, None, gamma=self._transformer[0].attn_norm_gamma, beta=None) - - for layer_idx in range(self.num_layers): - residual, hidden_states = self._forward_transformer_layer(layer_idx, residual, hidden_states, - wrapped_batch) - - return self._forward_unembed(residual, wrapped_batch) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/policy.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/policy.py deleted file mode 100644 index a1b445929053afb53e6998cd11f176b51651baa5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/phi3/policy.py +++ /dev/null @@ -1,30 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any - -from ...config_v2 import RaggedInferenceEngineConfig -from ..inference_policy_base import ContainerMap, InferenceV2Policy -from .containers import Phi3NonTransformerContainer, Phi3TransformerContainer -from .model import Phi3InferenceModel - - -class Phi3Policy(InferenceV2Policy): - - def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> Phi3InferenceModel: - return Phi3InferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group) - - def build_container_map(self) -> ContainerMap: - map = ContainerMap() - - transformer_containers = [Phi3TransformerContainer(self.model) for _ in range(self.model.num_layers)] - - map.set_transformer_params(['model.layers'], transformer_containers) - - map.set_non_transformer_params(Phi3NonTransformerContainer(self.model)) - - map.set_unmapped_params([]) - - return map diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__init__.py deleted file mode 100644 index 18206048fa299b8ced074aba9ea7f5212d3e0e3b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .policy import QwenPolicy diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 4b2d247f71cbe6e9798d34db45744144488b3c00..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__pycache__/container.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__pycache__/container.cpython-310.pyc deleted file mode 100644 index bc13a0f8760782c1ab287c7b3b8d6b184fa0f1d1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__pycache__/container.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__pycache__/model.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__pycache__/model.cpython-310.pyc deleted file mode 100644 index 138c8500c4e788639870150e846f3720e5077a31..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__pycache__/model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__pycache__/policy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__pycache__/policy.cpython-310.pyc deleted file mode 100644 index 1a887390a957ab190ddd66a863ae4f222922189e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/__pycache__/policy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/container.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/container.py deleted file mode 100644 index 313de68555b90f1dd38935de286d93f9cbd8b382..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/container.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Create a container object to save model-specific tensors using the policy file above. - -from ..common_parameters import * -from ..layer_container_base import LayerContainer -''' - # HF Qwen model looks like this: - -QWenLMHeadModel( - (transformer): QWenModel( - (wte): Embedding(151936, 4096) - (drop): Dropout(p=0.0, inplace=False) - (rotary_emb): RotaryEmbedding() - (h): ModuleList( - (0-31): 32 x QWenBlock( - (ln_1): RMSNorm() - (attn): QWenAttention( - (c_attn): Linear(in_features=4096, out_features=12288, bias=True) - (c_proj): Linear(in_features=4096, out_features=4096, bias=False) - (attn_dropout): Dropout(p=0.0, inplace=False) - ) - (ln_2): RMSNorm() - (mlp): QWenMLP( - (w1): Linear(in_features=4096, out_features=11008, bias=False) - (w2): Linear(in_features=4096, out_features=11008, bias=False) - (c_proj): Linear(in_features=11008, out_features=4096, bias=False) - ) - ) - ) - (ln_f): RMSNorm() - ) - (lm_head): Linear(in_features=4096, out_features=151936, bias=False) -) -''' - - -class QwenTransformerContainer(LayerContainer): - """ - Transformer layer container for the Qwen model. - """ - qkv_w: FusedQKVParameter - qkv_b: FusedQKVParameter - attn_out_w: AttentionOutputParameter - mlp_1_w: GatedMLPParameter - mlp_2_w: MLP2Parameter - attn_norm_gamma: NormParameter - mlp_norm_gamma: NormParameter - - PARAM_MAPPING = { - "attn.c_attn.weight": "qkv_w.params", - "attn.c_attn.bias": "qkv_b.params", - "attn.c_proj.weight": "attn_out_w.params", - "mlp.w1.weight": "mlp_1_w.up_params", - "mlp.w2.weight": "mlp_1_w.gate_params", - "mlp.c_proj.weight": "mlp_2_w.params", - "ln_1.weight": "attn_norm_gamma.params", - "ln_2.weight": "mlp_norm_gamma.params", - } - - -class QwenNonTransformerContainer(LayerContainer): - """ - Non-Transformer layer container for the Qwen model. - """ - word_emb: EmbeddingParameter - word_unembed: UnembedParameter - final_norm: NormParameter - - PARAM_MAPPING = { - "transformer.wte.weight": "word_emb.params", - "transformer.ln_f.weight": "final_norm.params", - "lm_head.weight": "word_unembed.params", - } diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/model.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/model.py deleted file mode 100644 index e867e4be67133cb512c737db3a3c45cb2294e404..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/model.py +++ /dev/null @@ -1,223 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Iterable, Optional, Tuple - -import torch - -import deepspeed.comm as dist - -from ...allocator import empty_from -from ...inference_utils import ActivationType, DtypeEnum -from .. import * -from ...modules.configs import * -from ...modules.interfaces import * -from ...modules import heuristics -from ...ragged import RaggedBatchWrapper - -from .container import QwenNonTransformerContainer, QwenTransformerContainer - - -class QwenInferenceModel(DSTransformerModelBase): - """ - Inference model implementation for ragged batching for Llama-2 models. - """ - - _non_transformer: Optional[QwenNonTransformerContainer] - """ - Embed + unembed container. Specializing the type annotation. - """ - - _transformer: Optional[Iterable[QwenTransformerContainer]] - """ - Per-layer transformer container. Specializing the type annotation. - """ - """ - Properties ineherited from `DSInferenceModelBase` - """ - - @property - def max_sequence_length(self) -> int: - return self._config.max_seq_length - - """ - Properties ineherited from `DSTransformerModelBase` - """ - - @property - def num_layers(self) -> int: - return self._config.num_hidden_layers - - @property - def model_dim(self) -> int: - return self._config.hidden_size - - @property - def vocab_size(self) -> int: - return self._config.vocab_size - - @property - def head_size(self) -> int: - return self.model_dim // self.n_heads - - @property - def n_heads(self) -> int: - return self._config.num_attention_heads - - @property - def intermediate_dim(self) -> int: - return self._config.intermediate_size // 2 - - @property - def n_heads_kv(self) -> int: - return self._config.hidden_size // self._config.kv_channels - - @property - def activation_dtype(self) -> DtypeEnum: - autoset_precision = self._config.bf16 + self._config.fp16 == 0 - if autoset_precision: - return DtypeEnum.fp16 - if self._config.fp16: - return DtypeEnum.fp16 - elif self._config.bf16: - # TODO(ZonePG): bf16 inference results may be different from huggingface bf16, - # because in rms_norm, Qwen still use float() instead of bf16 - return DtypeEnum.bf16 - else: - raise NotImplementedError("Only fp16 and bf16 are supported") - - @property - def mlp_activation_fn(self) -> ActivationType: - return ActivationType.SiGLU - - @property - def norm_type(self) -> NormTypeEnum: - return NormTypeEnum.RMSNorm - - @property - def positional_embedding_type(self) -> PositionalEmbeddingType: - return PositionalEmbeddingType.rotate_half - - @property - def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - return RotateHalfConfig(theta_base=self._config.rotary_emb_base) - - def make_norm_layer(self) -> None: - """ - Instantiates the normalization layer for the model. This sets the `self.norm` attribute. - - TODO(cmikeh2): In the future we'll distinguish between the different norm objects, - but for now we'll just use the same one for all of them. - """ - norm_config = DSNormConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - type=self.norm_type, - channels=self.model_dim, - residual_dtype=self.activation_dtype, - input_dtype=self.activation_dtype, - output_dtype=self.activation_dtype, - eps=self._config.layer_norm_epsilon, - ) - - self.norm = heuristics.instantiate_pre_norm(norm_config, self._engine_config) - - """ - Forward implementations - """ - - def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs the embedding lookup prior to running the transformer of the model. - - Arguments: - ragged_batch (RaggedBatchWrapper): The batch to embed. - - Returns: - torch.Tensor: The embedded batch. - """ - embed = self.embed(ragged_batch, self._non_transformer.word_emb) - - if embed.shape[-1] != self.model_dim: - raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}") - - return embed - - def _forward_transformer_layer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor, - ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead - optimization to fuse the layer norm of the next layer into the current layer. - - Arguments: - layer_idx (int): The index of the layer to execute. - residual (torch.Tensor): The residual tensor from the previous layer. - hidden_states (torch.Tensor): The hidden states from the previous layer. This is the - hidden states after pre normalization. - ragged_batch_info (RaggedBatchWrapper): The batch metadata. - """ - # TODO(cmikeh2): Distribute ragged_batch_info to all modules - - cur_params = self._transformer[layer_idx] - kv_cache = self.state_manager.get_cache(layer_idx) - - hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=cur_params.qkv_b) - hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info) - hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=None) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - residual, hidden_states = self.norm(residual, hidden_states, cur_params.mlp_norm_gamma, beta=None) - - # Should be configurable in the future - hidden_states = self.mlp_1(hidden_states, cur_params.mlp_1_w, b=None) - hidden_states = self.mlp_2(hidden_states, cur_params.mlp_2_w, b=None) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - if layer_idx != self.num_layers - 1: - next_params = self._transformer[layer_idx + 1] - residual, hidden_states = self.norm(residual, hidden_states, next_params.attn_norm_gamma, beta=None) - else: - # On last layer, we just need to perform the residual add. Adding into the residual - # here is safe. - residual.add_(hidden_states) - - return residual, hidden_states - - def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs unembedding of the hidden states to logits. This will only sample the final - token of each sequence. - """ - logits = self.unembed(hidden_states, - self._non_transformer.word_unembed, - ragged_batch_info, - gamma=self._non_transformer.final_norm) - - if self.tp_size > 1: - comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1])) - full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size)) - - dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group) - - full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size)) - - return full_logits - else: - return logits - - def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor: - - residual = self._forward_embed(wrapped_batch) - - residual, hidden_states = self.norm(residual, None, self._transformer[0].attn_norm_gamma, beta=None) - - for layer_idx in range(self.num_layers): - residual, hidden_states = self._forward_transformer_layer(layer_idx, residual, hidden_states, - wrapped_batch) - - return self._forward_unembed(residual, wrapped_batch) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/policy.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/policy.py deleted file mode 100644 index a9263f553621ec715cc5527e95034b842be0132b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen/policy.py +++ /dev/null @@ -1,30 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any - -from ...config_v2 import RaggedInferenceEngineConfig -from ..inference_policy_base import ContainerMap, InferenceV2Policy -from .container import QwenNonTransformerContainer, QwenTransformerContainer -from .model import QwenInferenceModel - - -class QwenPolicy(InferenceV2Policy): - - def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> QwenInferenceModel: - return QwenInferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group) - - def build_container_map(self) -> ContainerMap: - map = ContainerMap() - - transformer_containers = [QwenTransformerContainer(self.model) for _ in range(self.model.num_layers)] - - map.set_transformer_params(['transformer.h'], transformer_containers) - - map.set_non_transformer_params(QwenNonTransformerContainer(self.model)) - - map.set_unmapped_params(['transformer.rotary_emb.inv_freq']) - - return map diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__init__.py deleted file mode 100644 index 80b09757c74db181e5a7729579c89d530f65f25c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .policy import Qwen2Policy diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 90fa6a62fe8c75ba69d0dd87e357a03e698977f1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/container.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/container.cpython-310.pyc deleted file mode 100644 index dcffe68423dfdc80fb2f4cf012065910b48e82af..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/container.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/model.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/model.cpython-310.pyc deleted file mode 100644 index e5666d13f68d74f86ab36488c641b0da246ee94b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/policy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/policy.cpython-310.pyc deleted file mode 100644 index 91a3acc1602778e67f59b04b4d278c7b46f6b93c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/__pycache__/policy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/container.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/container.py deleted file mode 100644 index 6556d87d6afb2c2d36781d9af8c4a09ce035aee1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/container.py +++ /dev/null @@ -1,82 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Create a container object to save model-specific tensors using the policy file above. - -from ..common_parameters import * -from ..layer_container_base import LayerContainer -''' - # HF Qwen2 model looks like this: - -Qwen2ForCausalLM( - (model): Qwen2Model( - (embed_tokens): Embedding(151936, 1024) - (layers): ModuleList( - (0-23): 24 x Qwen2DecoderLayer( - (self_attn): Qwen2SdpaAttention( - (q_proj): Linear(in_features=1024, out_features=1024, bias=True) - (k_proj): Linear(in_features=1024, out_features=1024, bias=True) - (v_proj): Linear(in_features=1024, out_features=1024, bias=True) - (o_proj): Linear(in_features=1024, out_features=1024, bias=False) - (rotary_emb): Qwen2RotaryEmbedding() - ) - (mlp): Qwen2MLP( - (gate_proj): Linear(in_features=1024, out_features=2816, bias=False) - (up_proj): Linear(in_features=1024, out_features=2816, bias=False) - (down_proj): Linear(in_features=2816, out_features=1024, bias=False) - (act_fn): SiLU() - ) - (input_layernorm): Qwen2RMSNorm() - (post_attention_layernorm): Qwen2RMSNorm() - ) - ) - (norm): Qwen2RMSNorm() - ) - (lm_head): Linear(in_features=1024, out_features=151936, bias=False) -) -''' - - -class Qwen2TransformerContainer(LayerContainer): - """ - Transformer layer container for the Qwen2 model. - """ - qkv_w: UnfusedQKVParameter - qkv_b: UnfusedQKVParameter - attn_out_w: AttentionOutputParameter - mlp_1_w: GatedMLPParameter - mlp_2_w: MLP2Parameter - attn_norm_gamma: NormParameter - mlp_norm_gamma: NormParameter - - PARAM_MAPPING = { - "self_attn.q_proj.weight": "qkv_w.q_params", - "self_attn.k_proj.weight": "qkv_w.k_params", - "self_attn.v_proj.weight": "qkv_w.v_params", - "self_attn.q_proj.bias": "qkv_b.q_params", - "self_attn.k_proj.bias": "qkv_b.k_params", - "self_attn.v_proj.bias": "qkv_b.v_params", - "self_attn.o_proj.weight": "attn_out_w.params", - "mlp.gate_proj.weight": "mlp_1_w.gate_params", - "mlp.up_proj.weight": "mlp_1_w.up_params", - "mlp.down_proj.weight": "mlp_2_w.params", - "input_layernorm.weight": "attn_norm_gamma.params", - "post_attention_layernorm.weight": "mlp_norm_gamma.params", - } - - -class Qwen2NonTransformerContainer(LayerContainer): - """ - Non-Transformer layer container for the Qwen2 model. - """ - word_emb: EmbeddingParameter - word_unembed: UnembedParameter - final_norm: NormParameter - - PARAM_MAPPING = { - "model.embed_tokens.weight": "word_emb.params", - "model.norm.weight": "final_norm.params", - "lm_head.weight": "word_unembed.params", - } diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/model.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/model.py deleted file mode 100644 index d535462a954d4156ff32f2d3667a6a17713b4871..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/model.py +++ /dev/null @@ -1,221 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Iterable, Optional, Tuple - -import torch - -import deepspeed.comm as dist - -from ...allocator import empty_from -from ...inference_utils import ActivationType, DtypeEnum -from .. import * -from ...modules.configs import * -from ...modules.interfaces import * -from ...modules import heuristics -from ...ragged import RaggedBatchWrapper - -from .container import Qwen2NonTransformerContainer, Qwen2TransformerContainer - - -class Qwen2InferenceModel(DSTransformerModelBase): - """ - Inference model implementation for ragged batching for Llama-2 models. - """ - - _non_transformer: Optional[Qwen2NonTransformerContainer] - """ - Embed + unembed container. Specializing the type annotation. - """ - - _transformer: Optional[Iterable[Qwen2TransformerContainer]] - """ - Per-layer transformer container. Specializing the type annotation. - """ - """ - Properties ineherited from `DSInferenceModelBase` - """ - - @property - def max_sequence_length(self) -> int: - return self._config.max_seq_length - - """ - Properties ineherited from `DSTransformerModelBase` - """ - - @property - def num_layers(self) -> int: - return self._config.num_hidden_layers - - @property - def model_dim(self) -> int: - return self._config.hidden_size - - @property - def vocab_size(self) -> int: - return self._config.vocab_size - - @property - def head_size(self) -> int: - return self.model_dim // self.n_heads - - @property - def n_heads(self) -> int: - return self._config.num_attention_heads - - @property - def intermediate_dim(self) -> int: - return self._config.intermediate_size - - @property - def n_heads_kv(self) -> int: - return self._config.num_key_value_heads - - @property - def activation_dtype(self) -> DtypeEnum: - # TODO(ZonePG): bf16 inference results may be different from huggingface bf16, - # because in rms_norm, Qwen still use float() instead of bf16 - # if self._config.torch_dtype == torch.float16: - # return DtypeEnum.fp16 - # elif self._config.torch_dtype == torch.bfloat16: - # return DtypeEnum.bf16 - # else: - # raise NotImplementedError("Only fp16 and bf16 are supported") - return DtypeEnum.fp16 - - @property - def mlp_activation_fn(self) -> ActivationType: - return ActivationType.SiGLU - - @property - def norm_type(self) -> NormTypeEnum: - return NormTypeEnum.RMSNorm - - @property - def positional_embedding_type(self) -> PositionalEmbeddingType: - return PositionalEmbeddingType.rotate_half - - @property - def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - return RotateHalfConfig(theta_base=self._config.rope_theta) - - def make_norm_layer(self) -> None: - """ - Instantiates the normalization layer for the model. This sets the `self.norm` attribute. - - TODO(cmikeh2): In the future we'll distinguish between the different norm objects, - but for now we'll just use the same one for all of them. - """ - norm_config = DSNormConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - type=self.norm_type, - channels=self.model_dim, - residual_dtype=self.activation_dtype, - input_dtype=self.activation_dtype, - output_dtype=self.activation_dtype, - eps=self._config.rms_norm_eps, - ) - - self.norm = heuristics.instantiate_pre_norm(norm_config, self._engine_config) - - """ - Forward implementations - """ - - def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs the embedding lookup prior to running the transformer of the model. - - Arguments: - ragged_batch (RaggedBatchWrapper): The batch to embed. - - Returns: - torch.Tensor: The embedded batch. - """ - embed = self.embed(ragged_batch, self._non_transformer.word_emb) - - if embed.shape[-1] != self.model_dim: - raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}") - - return embed - - def _forward_transformer_layer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor, - ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead - optimization to fuse the layer norm of the next layer into the current layer. - - Arguments: - layer_idx (int): The index of the layer to execute. - residual (torch.Tensor): The residual tensor from the previous layer. - hidden_states (torch.Tensor): The hidden states from the previous layer. This is the - hidden states after pre normalization. - ragged_batch_info (RaggedBatchWrapper): The batch metadata. - """ - # TODO(cmikeh2): Distribute ragged_batch_info to all modules - - cur_params = self._transformer[layer_idx] - kv_cache = self.state_manager.get_cache(layer_idx) - - hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=cur_params.qkv_b) - hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info) - hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=None) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - residual, hidden_states = self.norm(residual, hidden_states, cur_params.mlp_norm_gamma, beta=None) - - # Should be configurable in the future - hidden_states = self.mlp_1(hidden_states, cur_params.mlp_1_w, b=None) - hidden_states = self.mlp_2(hidden_states, cur_params.mlp_2_w, b=None) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - if layer_idx != self.num_layers - 1: - next_params = self._transformer[layer_idx + 1] - residual, hidden_states = self.norm(residual, hidden_states, next_params.attn_norm_gamma, beta=None) - else: - # On last layer, we just need to perform the residual add. Adding into the residual - # here is safe. - residual.add_(hidden_states) - - return residual, hidden_states - - def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs unembedding of the hidden states to logits. This will only sample the final - token of each sequence. - """ - logits = self.unembed(hidden_states, - self._non_transformer.word_unembed, - ragged_batch_info, - gamma=self._non_transformer.final_norm) - - if self.tp_size > 1: - comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1])) - full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size)) - - dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group) - - full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size)) - - return full_logits - else: - return logits - - def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor: - - residual = self._forward_embed(wrapped_batch) - - residual, hidden_states = self.norm(residual, None, self._transformer[0].attn_norm_gamma, beta=None) - - for layer_idx in range(self.num_layers): - residual, hidden_states = self._forward_transformer_layer(layer_idx, residual, hidden_states, - wrapped_batch) - - return self._forward_unembed(residual, wrapped_batch) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/policy.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/policy.py deleted file mode 100644 index 9c5db2ba0065e0b545464c3ffe89095cb0e03148..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2/policy.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any - -from ...config_v2 import RaggedInferenceEngineConfig -from ..inference_policy_base import ContainerMap, InferenceV2Policy -from .container import Qwen2NonTransformerContainer, Qwen2TransformerContainer -from .model import Qwen2InferenceModel - - -class Qwen2Policy(InferenceV2Policy): - - def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> Qwen2InferenceModel: - return Qwen2InferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group) - - def build_container_map(self) -> ContainerMap: - map = ContainerMap() - - transformer_containers = [Qwen2TransformerContainer(self.model) for _ in range(self.model.num_layers)] - - map.set_transformer_params(['model.layers'], transformer_containers) - - map.set_non_transformer_params(Qwen2NonTransformerContainer(self.model)) - - map.set_unmapped_params( - [f'model.layers.{i}.self_attn.rotary_emb.inv_freq' for i in range(self.model.num_layers)]) - - return map diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__init__.py deleted file mode 100644 index 23e06a7700237f1e8e7c91f83c6f2986f2db8993..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .policy import Qwen2MoePolicy diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index d343f1142c7477d1ff86fa1e71eab2fb98af14aa..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/container.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/container.cpython-310.pyc deleted file mode 100644 index 3703faabccdba69c4e0ebf8e60f5995f72a21074..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/container.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/model.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/model.cpython-310.pyc deleted file mode 100644 index 11410e319c32627fa0c5be9aba95b40dd574320d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/policy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/policy.cpython-310.pyc deleted file mode 100644 index c6a54970ac4afe3cf429b141ae90fc723bd9738f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/__pycache__/policy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/container.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/container.py deleted file mode 100644 index e499379da7e368ccb9279333bd4e28ec8cff7e0d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/container.py +++ /dev/null @@ -1,103 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Create a container object to save model-specific tensors using the policy file above. - -from ..common_parameters import * -from ..layer_container_base import LayerContainer -''' - # HF Qwen2-57B-A14B model looks like this: - -Qwen2MoeForCausalLM( - (model): Qwen2MoeModel( - (embed_tokens): Embedding(151936, 3584) - (layers): ModuleList( - (0-27): 28 x Qwen2MoeDecoderLayer( - (self_attn): Qwen2MoeSdpaAttention( - (q_proj): Linear(in_features=3584, out_features=3584, bias=True) - (k_proj): Linear(in_features=3584, out_features=512, bias=True) - (v_proj): Linear(in_features=3584, out_features=512, bias=True) - (o_proj): Linear(in_features=3584, out_features=3584, bias=False) - (rotary_emb): Qwen2MoeRotaryEmbedding() - ) - (mlp): Qwen2MoeSparseMoeBlock( - (gate): Linear(in_features=3584, out_features=64, bias=False) - (experts): ModuleList( - (0-63): 64 x Qwen2MoeMLP( - (gate_proj): Linear(in_features=3584, out_features=2560, bias=False) - (up_proj): Linear(in_features=3584, out_features=2560, bias=False) - (down_proj): Linear(in_features=2560, out_features=3584, bias=False) - (act_fn): SiLU() - ) - ) - (shared_expert): Qwen2MoeMLP( - (gate_proj): Linear(in_features=3584, out_features=20480, bias=False) - (up_proj): Linear(in_features=3584, out_features=20480, bias=False) - (down_proj): Linear(in_features=20480, out_features=3584, bias=False) - (act_fn): SiLU() - ) - (shared_expert_gate): Linear(in_features=3584, out_features=1, bias=False) - ) - (input_layernorm): Qwen2MoeRMSNorm((3584,), eps=1e-06) - (post_attention_layernorm): Qwen2MoeRMSNorm((3584,), eps=1e-06) - ) - ) - (norm): Qwen2MoeRMSNorm((3584,), eps=1e-06) - ) - (lm_head): Linear(in_features=3584, out_features=151936, bias=False) -) -''' - - -class Qwen2MoeTransformerContainer(LayerContainer): - """ - Transformer layer container for the Qwen2Moe model. - """ - qkv_w: UnfusedQKVParameter - qkv_b: UnfusedQKVParameter - attn_out_w: AttentionOutputParameter - moe_gate: MoEGatingWeightParameter - moe_mlp_1: UnfusedMoEGatedMLPParameter - moe_mlp_2: UnfusedMoEMLP2Parameter - shared_moe_mlp_1: GatedMLPParameter - shared_moe_mlp_2: MLP2Parameter - shared_moe_gate: MoEGatingWeightParameter - attn_norm_gamma: NormParameter - mlp_norm_gamma: NormParameter - - PARAM_MAPPING = { - "self_attn.q_proj.weight": "qkv_w.q_params", - "self_attn.k_proj.weight": "qkv_w.k_params", - "self_attn.v_proj.weight": "qkv_w.v_params", - "self_attn.q_proj.bias": "qkv_b.q_params", - "self_attn.k_proj.bias": "qkv_b.k_params", - "self_attn.v_proj.bias": "qkv_b.v_params", - "self_attn.o_proj.weight": "attn_out_w.params", - "mlp.gate.weight": "moe_gate.params", - "mlp.experts.*.gate_proj.weight": "moe_mlp_1.gating_experts", - "mlp.experts.*.up_proj.weight": "moe_mlp_1.up_experts", - "mlp.experts.*.down_proj.weight": "moe_mlp_2.experts", - "mlp.shared_expert.gate_proj.weight": "shared_moe_mlp_1.gate_params", - "mlp.shared_expert.up_proj.weight": "shared_moe_mlp_1.up_params", - "mlp.shared_expert.down_proj.weight": "shared_moe_mlp_2.params", - "mlp.shared_expert_gate.weight": "shared_moe_gate.params", - "input_layernorm.weight": "attn_norm_gamma.params", - "post_attention_layernorm.weight": "mlp_norm_gamma.params", - } - - -class Qwen2MoeNonTransformerContainer(LayerContainer): - """ - Non-Transformer layer container for the Qwen2Moe model. - """ - word_emb: EmbeddingParameter - word_unembed: UnembedParameter - final_norm: NormParameter - - PARAM_MAPPING = { - "model.embed_tokens.weight": "word_emb.params", - "model.norm.weight": "final_norm.params", - "lm_head.weight": "word_unembed.params", - } diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/model.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/model.py deleted file mode 100644 index c7841b24e5fc79c6428422ca07618e78683cd7f1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/model.py +++ /dev/null @@ -1,359 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Iterable, Optional, Tuple - -import torch - -import deepspeed.comm as dist - -from ...allocator import empty_from -from ...config_v2 import RaggedInferenceEngineConfig -from ...inference_utils import ActivationType, DtypeEnum -from ...model_implementations import * -from ...modules.configs import * -from ...modules.interfaces import * -from ...modules import heuristics -from ...ragged import RaggedBatchWrapper -from ..inference_model_base import ( - DSModelImplementationConfig, - MPType, -) - -from .container import Qwen2MoeNonTransformerContainer, Qwen2MoeTransformerContainer - - -class Qwen2MoeInferenceModel(DSMoETransformerModelBase): - """ - Inference model implementation for Qwen2MoE models. - """ - - _non_transformer: Optional[Qwen2MoeNonTransformerContainer] - """ - Embed + unembed container. Specializing the type annotation. - """ - - _transformer: Optional[Iterable[Qwen2MoeTransformerContainer]] - """ - Per-layer transformer container. Specializing the type annotation. - """ - """ - Properties ineherited from `DSInferenceModelBase` - """ - - @property - def max_sequence_length(self) -> int: - return self._config.max_position_embeddings - - """ - Properties ineherited from `DSTransformerModelBase` - """ - - @property - def num_layers(self) -> int: - return self._config.num_hidden_layers - - @property - def model_dim(self) -> int: - return self._config.hidden_size - - @property - def vocab_size(self) -> int: - return self._config.vocab_size - - @property - def head_size(self) -> int: - return self.model_dim // self.n_heads - - @property - def n_heads(self) -> int: - return self._config.num_attention_heads - - @property - def intermediate_dim(self) -> int: - return self._config.shared_expert_intermediate_size - - @property - def n_heads_kv(self) -> int: - return self._config.num_key_value_heads - - @property - def activation_dtype(self) -> DtypeEnum: - # TODO(ZonePG): bf16 inference results may be different from huggingface bf16, - # because in rms_norm, Qwen still use float() instead of bf16 - # if self._config.torch_dtype == torch.float16: - # return DtypeEnum.fp16 - # elif self._config.torch_dtype == torch.bfloat16: - # return DtypeEnum.bf16 - # else: - # raise NotImplementedError("Only fp16 and bf16 are supported") - return DtypeEnum.fp16 - - @property - def mlp_activation_fn(self) -> ActivationType: - return ActivationType.SiGLU - - @property - def norm_type(self) -> NormTypeEnum: - return NormTypeEnum.RMSNorm - - @property - def positional_embedding_type(self) -> PositionalEmbeddingType: - return PositionalEmbeddingType.rotate_half - - @property - def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - return RotateHalfConfig(theta_base=self._config.rope_theta) - - """ - Inherited from `DSMoETransformerModelBase` - """ - - @property - def n_experts(self) -> int: - return self._config.num_experts - - @property - def n_top_k(self) -> int: - return self._config.num_experts_per_tok - - @property - def normalize_expert_scores(self) -> bool: - return self._config.norm_topk_prob - - def make_moe_layer(self) -> None: - """ - Instantiates the MoE layer for the model. This sets the `self.moe` attribute. - """ - sharded_dim = sharded_intermediate_dim(self.intermediate_dim // self.n_top_k, self.tp_size, self.tp_rank) - - moe_config = DSMoEConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - model_dim=self.model_dim, - intermediate_features=sharded_dim, - activation=self.mlp_activation_fn, - n_experts=self.n_experts, - top_k=self.n_top_k, - input_dtype=self.activation_dtype, - output_dtype=self.activation_dtype, - normalize_scores=self.normalize_expert_scores, - ) - - self.moe = heuristics.instantiate_moe(moe_config, self._engine_config) - - ######### MLP 1 ######### - def make_shared_expert_mlp_1_layer(self) -> None: - """ - Instantiates the linear projection layer for the first MLP in the feedforward network. - This sets the `self.mlp_1` attribute. - """ - shard_size = sharded_intermediate_dim(self.intermediate_dim, self.tp_size, self.tp_rank) - - linear_config = DSLinearConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - in_channels=self.model_dim, - out_channels=shard_size, - activation=self.mlp_activation_fn, - input_dtype=self.activation_dtype, - output_dtype=self.activation_dtype, - ) - - self.shared_expert_mlp_1 = heuristics.instantiate_linear(linear_config, self._engine_config) - - ######### MLP 2 ######### - def make_shared_expert_mlp_2_layer(self) -> None: - """ - Instantiates the linear projection layer for the second MLP in the feedforward network. - This sets the `self.mlp_2` attribute. - """ - shard_size = sharded_intermediate_dim(self.intermediate_dim, self.tp_size, self.tp_rank) - - linear_config = DSLinearConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - in_channels=shard_size, - out_channels=self.model_dim, - input_dtype=self.activation_dtype, - output_dtype=self.activation_dtype, - ) - - self.shared_expert_mlp_2 = heuristics.instantiate_linear(linear_config, self._engine_config) - - ######### MLP 2 ######### - def make_shared_expert_gate_layer(self) -> None: - """ - Instantiates the linear projection layer for the second MLP in the feedforward network. - This sets the `self.mlp_2` attribute. - """ - shard_size = sharded_intermediate_dim(self.model_dim, self.tp_size, self.tp_rank) - - linear_config = DSLinearConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - in_channels=shard_size, - out_channels=8, - input_dtype=self.activation_dtype, - output_dtype=self.activation_dtype, - ) - - self.shared_expert_gate = heuristics.instantiate_linear(linear_config, self._engine_config) - - def make_norm_layer(self) -> None: - """ - Instantiates the normalization layer for the model. This sets the `self.norm` attribute. - - TODO(cmikeh2): In the future we'll distinguish between the different norm objects, - but for now we'll just use the same one for all of them. - """ - norm_config = DSNormConfig( - max_tokens=self._engine_config.state_manager.max_ragged_batch_size, - type=self.norm_type, - channels=self.model_dim, - residual_dtype=self.activation_dtype, - input_dtype=self.activation_dtype, - output_dtype=self.activation_dtype, - eps=self._config.rms_norm_eps, - ) - - self.norm = heuristics.instantiate_pre_norm(norm_config, self._engine_config) - - """ - Model implementation - """ - - def __init__(self, config: DSModelImplementationConfig, engine_config: RaggedInferenceEngineConfig, - base_mp_group: MPType) -> None: - """ - Base implementation for initialization. By default, this will initialize - the traditional components of a transformer model: - - Embedding - - QKV projection - - Self attention - - Attention output projection - - Feed forward network - - Normalization - - Unembedding - - Arguments: - config (DSModelImplementationConfig): Model-specific configuration. No assumptions - should be made about this config that are not closely tied to the specific - model implementation. - engine_config (RaggedInferenceEngineConfig): Engine configuration. - base_mp_group (MPType): Base communication group for Tensor-parallel inference. - """ - super().__init__(config, engine_config, base_mp_group) - - self.make_norm_layer() - self.make_qkv_layer() - self.make_attn_layer() - self.make_attn_out_layer() - self.make_moe_layer() - self.make_shared_expert_mlp_1_layer() - self.make_shared_expert_mlp_2_layer() - self.make_shared_expert_gate_layer() - self.make_embedding_layer() - self.make_unembedding_layer() - self._kv_cache_config = None - - """ - Forward implementations - """ - - def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs the embedding lookup prior to running the transformer of the model. - - Arguments: - ragged_batch (RaggedBatchWrapper): The batch to embed. - - Returns: - torch.Tensor: The embedded batch. - """ - embed = self.embed(ragged_batch, self._non_transformer.word_emb) - - if embed.shape[-1] != self.model_dim: - raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}") - - return embed - - def _forward_transformer(self, layer_idx: int, residual: torch.Tensor, hidden_states: torch.Tensor, - ragged_batch_info: RaggedBatchWrapper) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Executes one (slightly offset) layer of the transformer. This implementation does a peak-ahead - optimization to fuse the layer norm of the next layer into the current layer. - - Arguments: - layer_idx (int): The index of the layer to execute. - residual (torch.Tensor): The residual tensor from the previous layer. - hidden_states (torch.Tensor): The hidden states from the previous layer. This is the - hidden states after pre normalization. - ragged_batch_info (RaggedBatchWrapper): The batch metadata. - """ - # TODO(cmikeh2): Distribute ragged_batch_info to all modules - - cur_params = self._transformer[layer_idx] - kv_cache = self.state_manager.get_cache(layer_idx) - - hidden_states = self.qkv(hidden_states, cur_params.qkv_w, b=cur_params.qkv_b) - hidden_states = self.attn(hidden_states, kv_cache, ragged_batch_info) - hidden_states = self.attn_out(hidden_states, cur_params.attn_out_w, b=None) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - residual, hidden_states = self.norm(residual, hidden_states, cur_params.mlp_norm_gamma, beta=None) - - shared_expert_output = self.shared_expert_mlp_1(hidden_states, cur_params.shared_moe_mlp_1, b=None) - shared_expert_output = self.shared_expert_mlp_2(shared_expert_output, cur_params.shared_moe_mlp_2, b=None) - shared_expert_gate_output = self.shared_expert_gate(hidden_states, cur_params.shared_moe_gate, b=None)[..., :1] - # shared_expert_gate_output shape[-1] is 1 - shared_expert_output.mul_(torch.sigmoid(shared_expert_gate_output)) - hidden_states = self.moe(hidden_states, ragged_batch_info, cur_params.moe_gate, cur_params.moe_mlp_1, - cur_params.moe_mlp_2) - hidden_states.add_(shared_expert_output) - - if self.tp_size > 1: - dist.all_reduce(hidden_states, group=self._base_mp_group) - - if layer_idx != self.num_layers - 1: - next_params = self._transformer[layer_idx + 1] - residual, hidden_states = self.norm(residual, hidden_states, next_params.attn_norm_gamma, beta=None) - else: - # On last layer, we just need to perform the residual add. Adding into the residual - # here is safe. - residual.add_(hidden_states) - - return residual, hidden_states - - def _forward_unembed(self, hidden_states: torch.Tensor, ragged_batch_info: RaggedBatchWrapper) -> torch.Tensor: - """ - Performs unembedding of the hidden states to logits. This will only sample the final - token of each sequence. - """ - logits = self.unembed(hidden_states, - self._non_transformer.word_unembed, - ragged_batch_info, - gamma=self._non_transformer.final_norm) - - if self.tp_size > 1: - comm_buffer = empty_from(self._comm_logits, (self.tp_size, logits.shape[0], logits.shape[1])) - full_logits = empty_from(self._return_logits, (logits.shape[0], self.vocab_size)) - - dist.all_gather_into_tensor(comm_buffer, logits, group=self._base_mp_group) - - full_logits.copy_(comm_buffer.permute(1, 0, 2).reshape(logits.shape[0], self.vocab_size)) - - return full_logits - else: - return logits - - def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor: - - residual = self._forward_embed(wrapped_batch) - - residual, hidden_states = self.norm(residual, None, self._transformer[0].attn_norm_gamma, beta=None) - - for layer_idx in range(self.num_layers): - residual, hidden_states = self._forward_transformer(layer_idx, residual, hidden_states, wrapped_batch) - - return self._forward_unembed(residual, wrapped_batch) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/policy.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/policy.py deleted file mode 100644 index 630bafe993a886cbd9d6e83b58450b711c870319..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/qwen_v2_moe/policy.py +++ /dev/null @@ -1,30 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any - -from ...config_v2 import RaggedInferenceEngineConfig -from ..inference_policy_base import ContainerMap, InferenceV2Policy -from .container import Qwen2MoeNonTransformerContainer, Qwen2MoeTransformerContainer -from .model import Qwen2MoeInferenceModel - - -class Qwen2MoePolicy(InferenceV2Policy): - - def instantiate_model(self, engine_config: RaggedInferenceEngineConfig, mp_group: Any) -> Qwen2MoeInferenceModel: - return Qwen2MoeInferenceModel(config=self._model_config, engine_config=engine_config, base_mp_group=mp_group) - - def build_container_map(self) -> ContainerMap: - map = ContainerMap() - - transformer_containers = [Qwen2MoeTransformerContainer(self.model) for _ in range(self.model.num_layers)] - - map.set_transformer_params(['model.layers'], transformer_containers) - - map.set_non_transformer_params(Qwen2MoeNonTransformerContainer(self.model)) - - map.set_unmapped_params([]) - - return map diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__init__.py deleted file mode 100644 index 63421bc1c622822f3975680bdbbcffe66b261c4d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .attn import * -from .attn_out import * -from .embedding import * -from .mlp import * -from .qkv import * -from .types import * -from .unembed import * diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 28448792503ff71f786914b20d008d445e53118d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/attn.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/attn.cpython-310.pyc deleted file mode 100644 index ee665bde9a6f160a817f50d4c3856f5787414cc8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/attn.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/attn_out.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/attn_out.cpython-310.pyc deleted file mode 100644 index 5030fc08570f9a83480b046c33eddf4fd4c0b765..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/attn_out.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/embedding.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/embedding.cpython-310.pyc deleted file mode 100644 index 1c7ecd3d3e60fa897a800c308a873ed9b329461c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/embedding.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/mlp.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/mlp.cpython-310.pyc deleted file mode 100644 index 1a9cdde745e5d0a1fcd5f15566fe8ad82774e652..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/mlp.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/qkv.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/qkv.cpython-310.pyc deleted file mode 100644 index 6b47f19436512f3d6c37d9535270488b3b5eb9d2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/qkv.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/types.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/types.cpython-310.pyc deleted file mode 100644 index 2241d51c03623aa84329e24de25c1fe77d03c065..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/types.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/unembed.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/unembed.cpython-310.pyc deleted file mode 100644 index 5b4b693d5d1b5d428954dd30dc8de980ee2e6973..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/unembed.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index 71d343e6cd451d97fbcea998658139842f2bc534..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/attn.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/attn.py deleted file mode 100644 index de8d6f6ac4c56dc591ea1dce59ef3b7e4f8b2b8b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/attn.py +++ /dev/null @@ -1,63 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Optional, Tuple - - -def get_local_heads(shard_rank: int, - num_shards: int, - n_heads_q: int, - n_heads_kv: Optional[int] = None) -> Tuple[int, int]: - """ - Helper to determine the number of local heads of a given shard. - - Args: - shard_rank (int): The rank of the shard. - num_shards (int): The total number of shards that attention is distributed over. - n_heads_q (int): The number of query heads. - n_heads_kv (int): The number of key/value heads. If not passed, it is assumed that - the number of query and key/value heads are the same. - """ - if n_heads_q < num_shards: - raise ValueError("There must be at least as many attention heads as there are shards.") - - if n_heads_kv is None or n_heads_kv == n_heads_q: - # MHA attention - base_heads = n_heads_q // num_shards - extra_heads = n_heads_q % num_shards - - if shard_rank < extra_heads: - return (base_heads + 1), (base_heads + 1) - else: - return base_heads, base_heads - else: - # GQA attention - if n_heads_q % n_heads_kv != 0: - raise ValueError("Must be an even ratio between query and key/value heads.") - - if n_heads_kv < num_shards and num_shards % n_heads_kv != 0: - raise ValueError( - "If splitting a group across multiple shards, we must be able to distribute the groups evenly.") - - if n_heads_kv >= num_shards and n_heads_kv % num_shards != 0: - raise ValueError("If parallelizing groups, must be able to evenly distribute them.") - - q_ratio = n_heads_q // n_heads_kv - - if n_heads_kv >= num_shards: - local_kv_heads = n_heads_kv // num_shards - local_q_heads = local_kv_heads * q_ratio - return local_q_heads, local_kv_heads - else: - group_sharding_size = num_shards // n_heads_kv - group_rank_idx = shard_rank % group_sharding_size - - base_heads = q_ratio // group_sharding_size - extra_heads = q_ratio % group_sharding_size - - if group_rank_idx < extra_heads: - return (base_heads + 1), 1 - else: - return base_heads, 1 diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/attn_out.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/attn_out.py deleted file mode 100644 index ce7c105531eaba2c3948b385282b2bd1c3358b5f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/attn_out.py +++ /dev/null @@ -1,111 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Optional - -import torch - -from .types import ShardingType -from .utils import shard_param, get_shard_endpoints - - -def shard_attn_out_param(param: torch.Tensor, - shard_rank: int, - num_shards: int, - head_size: int, - n_heads_q: Optional[int] = None, - n_heads_kv: Optional[int] = None) -> Optional[torch.Tensor]: - """ - Utility method for sharding an attention output parameter. - """ - if len(param.shape) == 1: - # We will do the bias addition on the 0th rank only rather than scale the parameter and - # implicitly reconstruct this in the distributed reduce. - return param if shard_rank == 0 else None - - assert n_heads_kv is None or (n_heads_q is not None - and n_heads_kv is not None), "n_heads_kv should not be passed without n_heads_q" - - mha_sharding = n_heads_kv is None or n_heads_q == n_heads_kv - - if mha_sharding: - return shard_param(param, ShardingType.INNER_DIMENSION, shard_rank, num_shards, granularity=head_size) - else: - assert param.shape[0] == head_size * n_heads_q, "GQA param shape is not correct" - - # 32 KV heads, 16 shards for example - even_kv_sharding = n_heads_kv % num_shards == 0 - - # 8 KV heads, 16 shards for example - even_kv_distribution = num_shards % n_heads_kv == 0 - - assert even_kv_sharding or even_kv_distribution, "No partitioning algorithm for this yet." - - if even_kv_sharding: - # Same as original sharding scenario - return shard_param(param, ShardingType.INNER_DIMENSION, shard_rank, num_shards, granularity=head_size) - else: - # We will first do a sharding on the KV and Q to map to the one KV shard per group of Q. - q_sharding_degree = num_shards // n_heads_kv - - kv_head = shard_rank // q_sharding_degree - - q_sharding_rank = shard_rank % q_sharding_degree - q_factor = n_heads_q // n_heads_kv - - q_chunk = param[..., q_factor * kv_head * head_size:q_factor * (kv_head + 1) * head_size] - - return shard_param(q_chunk, - ShardingType.INNER_DIMENSION, - q_sharding_rank, - q_sharding_degree, - granularity=head_size) - - -def attn_out_in_features(out_features: int, - shard_rank: int, - num_shards: int, - head_size: int, - n_heads_q: Optional[int] = None, - n_heads_kv: Optional[int] = None) -> int: - """ - Helper to calculate the expected output projection dimension of a QKV projection matrix. - - Args: - in_features (int): The model dimension. - shard_rank (int): Which rank to return the corresponding size for. - num_shards (int): The total number of shards the parameter is distributed across. - head_size (int): The size of each attention head. - n_heads_q (int): The number of query heads on the model. This only needs to be passed if the number - of query and key/value heads are different. If passed without n_heads_kv, default - MHA partitioning will be used. - n_heads_kv (int): The number of key and value heads on the model. This only needs to be passed - if the number of query and key/value heads are different. This argument cannot be passed without - also passing n_heads_q (we want to explicitly opt into GQA sharding). - """ - assert n_heads_kv is None or (n_heads_q is not None - and n_heads_kv is not None), "n_heads_kv should not be passed without n_heads_q" - - mha_sharding = n_heads_kv is None or n_heads_q == n_heads_kv - - if mha_sharding: - endpoints = get_shard_endpoints(out_features, shard_rank, num_shards, granularity=head_size) - return endpoints[1] - endpoints[0] - else: - if n_heads_kv >= num_shards: - assert n_heads_kv % num_shards == 0, "No partitioning algorithm for this yet." - n_local_groups = n_heads_kv // num_shards - group_size = n_heads_q // n_heads_kv - - return n_local_groups * head_size * group_size - else: - assert num_shards % n_heads_kv == 0, "No partitioning algorithm for this yet." - q_split_degree = num_shards // n_heads_kv - q_split_rank = shard_rank % q_split_degree - split_granularity = (n_heads_q // n_heads_kv) * head_size - - q_endpoints = get_shard_endpoints(split_granularity, q_split_rank, q_split_degree, granularity=head_size) - - return q_endpoints[1] - q_endpoints[0] diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/embedding.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/embedding.py deleted file mode 100644 index 00d335768ae69b1d897dac92c4c222fdedd1e609..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/embedding.py +++ /dev/null @@ -1,34 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from .types import ShardingType -from .utils import shard_param, get_shard_endpoints - - -def shard_embedding_param(param: torch.Tensor, shard_rank: int, num_shards: int) -> torch.Tensor: - """ - Utility method for sharding an embedding parameter. - - Args: - param (torch.Tensor): The parameter to shard. Should be of shape [vocab_size, model_dim] - shard_rank (int): Which shard of the partitioned tensor to return. - num_shards (int): The total number of shards the parameter is distributed across. - """ - return shard_param(param, ShardingType.INNER_DIMENSION, shard_rank, num_shards) - - -def sharded_embedding_dim(embedding_size: int, shard_rank: int, num_shards: int) -> int: - """ - Utility method for getting the size of the embedding dimension of a sharded embedding. - - Args: - embedding_size (int): The size of the embedding. - shard_rank (int): Which shard of the partitioned tensor to return. - num_shards (int): The total number of shards the parameter is distributed across. - """ - start_idx, end_idx = get_shard_endpoints(embedding_size, shard_rank, num_shards) - return end_idx - start_idx diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/mlp.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/mlp.py deleted file mode 100644 index 8abd0ff8622df3bd302f13292b48e9be9cae9782..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/mlp.py +++ /dev/null @@ -1,75 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Optional - -import torch - -from .types import ShardingType, DEFAULT_SHARD_GRANULARITY -from .utils import shard_param, get_shard_endpoints - - -def shard_mlp_1_param(param: torch.Tensor, - shard_rank: int, - num_shards: int, - gated: bool = False, - is_moe: bool = False) -> torch.Tensor: - """ - Utility method for sharding an MLP 1 parameter. Both biases and weights are supported, as well - as for fused weights for MoE. - - Args: - param (torch.Tensor): The parameter to shard. - shard_rank (int): Which shard of the partitioned tensor to return. - num_shards (int): The total number of shards the parameter is distributed across. - gated (bool): Whether or not the parameter is from a gated MLP. - """ - bias_dims = 2 if is_moe else 1 - - if gated: - return shard_param(param, - ShardingType.OUTER_DIMENSION, - shard_rank, - num_shards, - granularity=DEFAULT_SHARD_GRANULARITY * 2, - bias_dims=bias_dims) - else: - return shard_param(param, ShardingType.OUTER_DIMENSION, shard_rank, num_shards, bias_dims=bias_dims) - - -def shard_mlp_2_param(param: torch.Tensor, - shard_rank: int, - num_shards: int, - is_moe: bool = False) -> Optional[torch.Tensor]: - """ - Utility method for sharding an MLP 2 parameter. - - Args: - param (torch.Tensor): The parameter to shard. - shard_rank (int): Which shard of the partitioned tensor to return. - num_shards (int): The total number of shards the parameter is distributed across. - is_moe (bool): Whether or not the parameter is from a MoE model. - """ - bias_dim_size = 2 if is_moe else 1 - - if len(param.shape) == bias_dim_size: - # We will do the bias addition on the 0th rank only rather than scale the parameter and - # implicitly reconstruct this in the distributed reduce. - return param if shard_rank == 0 else None - - return shard_param(param, ShardingType.INNER_DIMENSION, shard_rank, num_shards) - - -def sharded_intermediate_dim(intermediate_size: int, num_shards: int, shard_rank: int) -> int: - """ - Utility method for getting the size of the intermediate dimension of a sharded MLP. - - Args: - intermediate_size (int): The size of the intermediate dimension. - num_shards (int): The total number of shards the parameter is distributed across. - shard_rank (int): Which shard of the partitioned tensor to return. - """ - endpoints = get_shard_endpoints(intermediate_size, shard_rank, num_shards) - return endpoints[1] - endpoints[0] diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/qkv.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/qkv.py deleted file mode 100644 index 2b6d7f40836e8ac72d479b2c2792551e38e3d379..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/qkv.py +++ /dev/null @@ -1,166 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Optional - -import torch - -from .types import ShardingType -from .utils import shard_param, get_shard_endpoints - - -def shard_qkv_param(param: torch.Tensor, - shard_rank: int, - num_shards: int, - head_size: int, - n_heads_q: Optional[int] = None, - n_heads_kv: Optional[int] = None) -> Optional[torch.Tensor]: - """ - Utility method for sharding a QKV parameter. Both biases and weights are supported. It is assumed - that the layout of the parameter is such that all Q heads, all K heads, and all V heads - are contiguous with respect to each other. - - Args: - param (torch.Tensor): The parameter to shard. - shard_rank (int): Which shard of the partitioned tensor to return. - num_shards (int): The total number of shards the parameter is distributed across. - head_size (int): The size of each head. - n_heads_q (int): The number of query heads. This only needs to be passed if the number - of query and key/value heads are different. If passed without n_heads_kv, default - MHA partitioning will be used. - n_heads_kv (int): The number of key/value heads. This only needs to be passed if the number - of query and key/value heads are different. This argument should not be passed without - n_heads_q (we want to explicitly opt into GQA sharding). - """ - if n_heads_kv is not None and n_heads_q is None: - raise ValueError("n_heads_kv should not be passed without n_heads_q") - - if n_heads_q is None: - # Guaranteed to be in MHA - if param.shape[0] // 3 % head_size != 0: - raise ValueError("MHA param shape is not correct") - n_heads_q = param.shape[0] // head_size // 3 - mha_sharding = True - else: - mha_sharding = n_heads_q == n_heads_kv - - if n_heads_q < num_shards: - raise ValueError("There must be at least as many query heads as there are shards.") - - if mha_sharding: - return shard_param(param, - ShardingType.OUTER_DIMENSION, - shard_rank, - num_shards, - num_concatenated_matrices=3, - granularity=head_size) - else: - if n_heads_q % n_heads_kv != 0: - raise ValueError("Must be an even ratio between query and key/value heads.") - - if param.shape[0] != head_size * (n_heads_q + 2 * n_heads_kv): - raise ValueError("GQA param shape is not correct") - - # 32 KV heads, 16 shards for example - if n_heads_kv >= num_shards and n_heads_kv % num_shards != 0: - raise ValueError("Currently do not support uneven partitioning of KV heads for GQA.") - - # 8 KV heads, 16 shards for example - if n_heads_kv < num_shards and num_shards % n_heads_kv != 0: - raise ValueError("Currently do not support distributing KV heads across different numbers of shards.") - else: - even_kv_sharding = n_heads_kv >= num_shards - - if param is None: - return None - - q_param = param[:head_size * n_heads_q] - kv_param = param[head_size * n_heads_q:] - - if even_kv_sharding: - # This is equivalent to the original sharding algorithm since n_heads_q = C * n_heads_kv. - # If n_heads_kv % num_shards == 0, then n_heads_q % num_shards == 0. - q_param = shard_param(q_param, ShardingType.OUTER_DIMENSION, shard_rank, num_shards, granularity=head_size) - kv_param = shard_param(kv_param, - ShardingType.OUTER_DIMENSION, - shard_rank, - num_shards, - num_concatenated_matrices=2, - granularity=head_size) - return torch.cat([q_param, kv_param], dim=0) - else: - # We will first do a sharding on the KV and Q to map to the one KV shard per group of Q. - q_sharding_degree = num_shards // n_heads_kv - - kv_head = shard_rank // q_sharding_degree - k_param = kv_param[kv_head * head_size:(kv_head + 1) * head_size] - v_param = kv_param[(n_heads_kv + kv_head) * head_size:(n_heads_kv + kv_head + 1) * head_size] - - q_sharding_rank = shard_rank % q_sharding_degree - q_factor = n_heads_q // n_heads_kv - - q_chunk = q_param[q_factor * kv_head * head_size:q_factor * (kv_head + 1) * head_size] - - q_param = shard_param(q_chunk, - ShardingType.OUTER_DIMENSION, - q_sharding_rank, - q_sharding_degree, - granularity=head_size) - - return torch.cat([q_param, k_param, v_param], dim=0) - - -def qkv_out_features(in_features: int, - shard_rank: int, - num_shards: int, - head_size: int, - n_heads_q: Optional[int] = None, - n_heads_kv: Optional[int] = None) -> int: - """ - Helper to calculate the expected output projection dimension of a QKV projection matrix. - - Args: - in_features (int): The model dimension. - shard_rank (int): Which rank to return the corresponding size for. - num_shards (int): The total number of shards the parameter is distributed across. - head_size (int): The size of each head. - n_heads_q (int): The number of query heads. This only needs to be passed if the number - of query and key/value heads are different. If passed without n_heads_kv, default - MHA partitioning will be used. - n_heads_kv (int): The number of key/value heads. This only needs to be passed if the number - of query and key/value heads are different. This argument cannot be passed without also - passing n_heads_q (we want to explicitly opt into GQA sharding). - """ - if n_heads_kv is not None and n_heads_q is None: - raise ValueError("n_heads_kv should not be passed without n_heads_q") - - mha_sharding = n_heads_kv is None or n_heads_q == n_heads_kv - - if n_heads_q is not None and in_features != head_size * n_heads_q: - raise ValueError("in_features is not consistent with n_heads_q and head_size") - - if mha_sharding: - endpoints = get_shard_endpoints(in_features, shard_rank, num_shards, granularity=head_size) - return (endpoints[1] - endpoints[0]) * 3 - else: - if n_heads_kv >= num_shards: - if n_heads_kv % num_shards != 0: - raise ValueError("The KV heads must be evenly distributed across the shards.") - - n_local_groups = n_heads_kv // num_shards - group_size = n_heads_q // n_heads_kv - - return n_local_groups * head_size * (2 + group_size) - else: - if num_shards % n_heads_kv != 0: - raise ValueError("A shared KV head must always partition across the same number of shards.") - - q_split_degree = num_shards // n_heads_kv - q_split_rank = shard_rank % q_split_degree - split_granularity = (n_heads_q // n_heads_kv) * head_size - - q_endpoints = get_shard_endpoints(split_granularity, q_split_rank, q_split_degree, granularity=head_size) - - return (q_endpoints[1] - q_endpoints[0]) + 2 * head_size diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/types.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/types.py deleted file mode 100644 index 01dce0db523a9e904052c8336fd14278c5162b2c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/types.py +++ /dev/null @@ -1,18 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from enum import Enum - -DEFAULT_SHARD_GRANULARITY = 32 - - -class ShardingType(Enum): - # Inner dimension sharding corresponds to splitting the Tensor along the K-dimension - # of a matrix multiplication. This would be used for attention_output or MLP2. - INNER_DIMENSION = 1 - - # Outer dimension sharding corresponds to splitting the Tensor along the N-dimension - # of a matrix multiplication. This would be used for the QKV and MLP1 projections. - OUTER_DIMENSION = 0 diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/unembed.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/unembed.py deleted file mode 100644 index 6cc771969ad9e195f04216e1374ea0ced6fb7065..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/unembed.py +++ /dev/null @@ -1,41 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from .types import ShardingType -from .utils import shard_param, get_shard_endpoints - - -def shard_unembed_param(param: torch.Tensor, shard_rank: int, num_shards: int) -> torch.Tensor: - """ - Utility method for sharding an unembed parameter. We shard unembeddings on the vocab dimension - with the expectation of an all-gather to produce the full results. - - TODO(cmikeh2): Really ideal would be if MII could have access to the comm and we would do - an A2A and sharded sampling. - - Args: - param (torch.Tensor): The parameter to shard. Should be of shape [vocab_size, model_dim] - shard_rank (int): Which shard of the partitioned tensor to return. - num_shards (int): The total number of shards the parameter is distributed across. - - Returns: - torch.Tensor: The sharded parameter of shape [sharded_vocab_size, model_dim] - """ - return shard_param(param, ShardingType.OUTER_DIMENSION, shard_rank, num_shards, granularity=1) - - -def sharded_unembed_dim(vocab_size: int, shard_rank: int, num_shards: int) -> int: - """ - Utility method for determining the sharded vocab size of a sharded unembed parameter. - - Args: - vocab_size (int): The size of the vocabulary. - shard_rank (int): Which shard of the partitioned tensor to return. - num_shards (int): The total number of shards the parameter is distributed across. - """ - start_idx, end_idx = get_shard_endpoints(vocab_size, shard_rank, num_shards, granularity=1) - return end_idx - start_idx diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/utils.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/utils.py deleted file mode 100644 index fd0eb51873f83834f99bf14b084cf62efbde05f3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/sharding/utils.py +++ /dev/null @@ -1,104 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Optional, Tuple - -import torch - -from .types import ShardingType, DEFAULT_SHARD_GRANULARITY - - -def get_shard_endpoints(dim_size: int, - shard_rank: int, - num_shards: int, - granularity: int = DEFAULT_SHARD_GRANULARITY) -> Tuple[int, int]: - """ - Given a dimension to shard with size dim_size, return the start and end indices of the slice - that belong to the given rank. - - The typical use of this is as an internal helper function, so see if there is a higher level - API that better suits the application. - - Args: - dim_size (int): The size of the dimension to shard. - shard_rank (int): The rank of the shard to return. - num_shards (int): Total number of shards the dimension will be distributed across. - granularity (int): The minimum alignment of the shard endpoints. This is used to support - non-even head counts as well as align dimensions to cleaner GEMM boundaries. - """ - assert dim_size % granularity == 0, "Dimension size must be divisible by granularity" - - total_chunks = dim_size // granularity - base_chunks_per_rank = total_chunks // num_shards - remainder_chunks = total_chunks % num_shards - - start_chunk_id = shard_rank * base_chunks_per_rank + min(shard_rank, remainder_chunks) - end_chunk_id = start_chunk_id + base_chunks_per_rank + (1 if shard_rank < remainder_chunks else 0) - - return start_chunk_id * granularity, end_chunk_id * granularity - - -def shard_param(param: Optional[torch.Tensor], - shard_mode: ShardingType, - shard_rank: int, - num_shards: int, - num_concatenated_matrices: int = 1, - granularity: int = 32, - bias_dims: int = 1) -> torch.Tensor: - """ - Utility for sharding a parameter. This will return the slice of the parameter that should - exist on the given shard_rank given the sharding configuration. The workflow here is - to find the minimum bounded Tensor to shard, get the slicing endpoints, and then concatenate - as needed. - - The typical use of this is as an internal helper function, so see if there is a higher level - API that better suits the application. - - Args: - param (torch.Tensor): The parameter to shard. - shard_mode (ShardingType): The type of sharding to apply. See ShardingType for more context. - shard_rank (int): The rank of the shard to return. - num_shards (int): Total number of shards the parameter will be distrbuted across. - num_concatenated_matrices (int): The number of matrices that have been concatenated together in the original - parameter. An example of this is a fused QKV projection matrix, where the `num_concatenated_matrices` - argument would be 3. - granularity (int): The minimum alignment of the shard endpoints. For attention projection matrices, this - should be set to the head size to support non-even sharding. - bias_dims (int): The number of dimensions that are considered bias dimensions. This is used to support - sharding of MoE and non-MoE biases on the same codepath. - """ - assert shard_rank < num_shards, "Shard rank must be less than num_shards" - - # Easier to hide this inside of the sharding logic than to add checks in every model - # implementation. - if param is None: - return None - - if num_shards == 1: - # Trivial case of no sharding. - return param - - if shard_mode == ShardingType.OUTER_DIMENSION: - - def get_matrices(dim_idx: int) -> torch.Tensor: - dim_size = param.size(dim_idx) // num_concatenated_matrices - start_channel_id, end_channel_id = get_shard_endpoints(dim_size, shard_rank, num_shards, granularity) - return torch.chunk(param, num_concatenated_matrices, dim=dim_idx), start_channel_id, end_channel_id - - if param.ndim == bias_dims: - # Special case for bias parameters. - matrices, start_channel_id, end_channel_id = get_matrices(dim_idx=-1) - return torch.cat([mat[..., start_channel_id:end_channel_id] for mat in matrices], dim=-1) - else: - # General case for weight parameters. This assumes MoE parameters are stored in the format of - # [num_experts, out_features, in_features] - matrices, start_channel_id, end_channel_id = get_matrices(dim_idx=-2) - return torch.cat([mat[..., start_channel_id:end_channel_id, :] for mat in matrices], dim=-2) - - elif shard_mode == ShardingType.INNER_DIMENSION: - dim_size = param.size(-1) // num_concatenated_matrices - start_channel_id, end_channel_id = get_shard_endpoints(dim_size, shard_rank, num_shards, granularity) - matrices = torch.chunk(param, num_concatenated_matrices, dim=-1) - return torch.cat([mat[..., start_channel_id:end_channel_id] for mat in matrices], dim=-1) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__init__.py deleted file mode 100644 index 917c1599de2e6124ef7ee09f73a655053f6acb7e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from . import implementations -from . import interfaces -from .module_registry import ConfigBundle diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index cbb30f5e86f371546a552eb10f37570ace59d7b0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__pycache__/ds_module.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__pycache__/ds_module.cpython-310.pyc deleted file mode 100644 index d2279dae9b97ff0f38950c3d99e401d14189df0e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__pycache__/ds_module.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__pycache__/heuristics.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__pycache__/heuristics.cpython-310.pyc deleted file mode 100644 index a8a9573161578a366630860abac65bf36eaf83c5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__pycache__/heuristics.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__pycache__/module_registry.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__pycache__/module_registry.cpython-310.pyc deleted file mode 100644 index 4dc2b69b1e6abb0048cd3db6f191f33cfb36cd6b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/__pycache__/module_registry.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__init__.py deleted file mode 100644 index 3429e69b47de32cd14346b8fd4ad1c7d0ac460c3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .attention_configs import ( - DSSelfAttentionConfig, - PositionalEmbeddingType, - MaskingType, - RotateHalfConfig, -) -from .embedding_config import DSEmbeddingsConfig -from .linear_config import DSLinearConfig -from .moe_config import DSMoEConfig -from .norm_config import DSNormConfig, NormTypeEnum -from .unembed_config import DSUnembedConfig diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 3e6b61d97e78a14567b6534b3761a8fae6f88f70..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/attention_configs.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/attention_configs.cpython-310.pyc deleted file mode 100644 index 38c66727ecc3d8910a24285ef56bf86b30e7abb2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/attention_configs.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/embedding_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/embedding_config.cpython-310.pyc deleted file mode 100644 index 3014804c58dc8b8ae04ee41be59cec3c92ffc901..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/embedding_config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/linear_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/linear_config.cpython-310.pyc deleted file mode 100644 index e6288ce8e74c9a72e719db8077fd590b5001a4d1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/linear_config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/moe_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/moe_config.cpython-310.pyc deleted file mode 100644 index 306a4ba78945bbc045695e94337bc683b03f09c6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/moe_config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/norm_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/norm_config.cpython-310.pyc deleted file mode 100644 index 5914d92583c27cd62a178205773b56c40b2e423b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/norm_config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/unembed_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/unembed_config.cpython-310.pyc deleted file mode 100644 index 89d7a9ab25961f97033f365f8d607481c5caa3c6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/__pycache__/unembed_config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/attention_configs.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/attention_configs.py deleted file mode 100644 index be6a3535024c1e2e90d89d57dca9efb454b8889f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/attention_configs.py +++ /dev/null @@ -1,110 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from enum import Enum -from typing import Dict, Optional - -from ...inference_utils import DtypeEnum -from ...modules.ds_module import DSModuleConfig -from deepspeed.runtime.config_utils import DeepSpeedConfigModel - - -class PositionalEmbeddingType(Enum): - - # No positional embeddings - none = "none" - - # Rotary positional embeddings - every half - rotate_half = "rotate_half" - - # Rotary positional embeddings - every other - rotate_every_other = "rotate_every_other" - - # Alibi - alibi = "alibi" - - -class RotateHalfConfig(DeepSpeedConfigModel): - - use_trained_freqs: bool = False - """ - Whether to use a passed `trained_freqs` tensor for the attention implementation - or to use default synthesized frequencies. - """ - - theta_base: float = 10_000.0 - """ - Base for theta. This will only be used if `use_trained_freqs` is False. - """ - - rotate_dim: Optional[int] = None - """ - How many neurons to rotate. If None, then all neurons will be rotated. Many external configs - will set this number to half the head dimension and then internally multiply by 2. To make it - more clear to understand what is happening (rotate_dim < head_dim -> then only partial rotation), - we do not do this multiplication internally. - """ - - -class MaskingType(Enum): - - # No masking - none = "none" - - # Causal masking - causal = "causal" - - # Local masking - local = "local" - - # Symmetric masking (this is a 1D tensor mask) - symmetric = "symmetric" - - # Arbitrary masking (this would correspond to a 2D tensor mask) - asymmetric = "asymmetric" - - -class DSSelfAttentionConfig(DSModuleConfig): - """ - Config class for attention. - """ - - # Number of query attention heads on this shard - n_heads_q: int - - # Number of KV attention heads on this shard - n_heads_kv: int - - # Size of each attention head - head_size: int - - # Max number of sequences that may compose a ragged batch - max_sequences: int - - # Scale factor for attention scores - scale_factor: float = 1.0 - - # Input data type - input_dtype: DtypeEnum = DtypeEnum.fp16 - - # Output data type - output_dtype: DtypeEnum = DtypeEnum.fp16 - - # Masking type - masking_type: MaskingType = MaskingType.causal - - # Masking args - masking_args: Dict = {} - - # Positional embedding type - positional_embedding_type: PositionalEmbeddingType = PositionalEmbeddingType.none - - # Positional embedding args - positional_embedding_config: Optional[RotateHalfConfig] = None - """ - To extend this for the other positional embedding types, we would need to add - new configs for each type (as necessary) and annotate this with the - Union[RotateHalfConfig, OtherConfig, ...] type. - """ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/embedding_config.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/embedding_config.py deleted file mode 100644 index 2486c5986e9531c7377f1dddbe81ebafb6ef377f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/embedding_config.py +++ /dev/null @@ -1,70 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Optional - -from ...inference_utils import DtypeEnum, NormTypeEnum -from ...modules.ds_module import DSModuleConfig -""" -Trying to define the space we need to support here right now: - -Types of embeddings I've found so far: - 1. Token embedding - 2. Position embedding - 3. Token type embedding - 4. LN - -GPTNeo: 1, 2, 3 (shared with 1) -GPTNeoX: 1 -GPTJ: 1, 3 -LLaMA: 1 -BERT: 1, 2, 3, 4 -GPT2: 1, 2, 3 (shared with 1) - -Sidebar for OPT: -OPT: 1, 2 -1 may not actually project to the actual hidden dimension according to the raw -code, but for the model configs we care about it does. -2 has a weird offset associated with it that the others do not. -""" - - -class DSEmbeddingsConfig(DSModuleConfig): - """ - Config class for DSEmbeddings. - """ - - residual_dtype: DtypeEnum = DtypeEnum.fp16 - """ - Data type the module should use for its output. - """ - - embedding_dim: int - """ - Dimensionality of the embedding projections. - """ - - positional_embedding: bool = False - """ - Whether the module should expect a positional embedding matrix. The shape of this - matrix should be of shape [max_seq_len + positional_offset, embedding_dim] - """ - - positional_offset: int = 0 - """ - Whether the linearized token IDs should be offset by a certain amount. For an example - of this, see the OPT model implementation. - """ - - use_token_type: bool = False - """ - Whether the module should expect a token type embedding matrix. - """ - - output_normalization: Optional[NormTypeEnum] = None - """ - If a the output of the embedding module should be normalized, specify here. See - ``inference.inference_utils.NormTypeEnum`` for supported values. - """ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/linear_config.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/linear_config.py deleted file mode 100644 index 40fe0773aeeee92d505a115f85c440a753491329..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/linear_config.py +++ /dev/null @@ -1,43 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from ...inference_utils import ActivationType, DtypeEnum -from ...modules.ds_module import DSModuleConfig - - -class DSLinearConfig(DSModuleConfig): - """ - Config class for DSLinearBase. - """ - - in_channels: int - """ - Number of input channels - """ - - out_channels: int - """ - Number of output channels. NOTE: If this linear layer is using a gated activation function, - the value for ``out_channels`` passed here should refer to the number of channels after - gating (i.e., the expected weight shape before transformations will be ``[out_channels * 2, in_channels]``). - """ - - activation: ActivationType = ActivationType.IDENTITY - """ - The activation function for this layer. See :class:`deepspeed.inference.inference_utils.ActivationType` for - supported activation functions. - """ - - input_dtype: DtypeEnum = DtypeEnum.fp16 - """ - The data type of the input tensor. See :class:`deepspeed.inference.inference_utils.DtypeEnum` for supported - data types. - """ - - output_dtype: DtypeEnum = DtypeEnum.fp16 - """ - The data type of the output tensor. See :class:`deepspeed.inference.inference_utils.DtypeEnum` for supported - data types. - """ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/moe_config.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/moe_config.py deleted file mode 100644 index 7bc944f55e17cf4e837c2a5db1144099dd942fdd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/moe_config.py +++ /dev/null @@ -1,56 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from ...inference_utils import ActivationType, DtypeEnum -from ...modules.ds_module import DSModuleConfig - - -class DSMoEConfig(DSModuleConfig): - """ - Config class for DSMoEBase - """ - - model_dim: int - """ - Size of input activation. - """ - - intermediate_features: int - """ - Size of intermediate activation. Specifically, this is the number of input features - in the second linear layer. Depending on the activation function, the output of the first - linear layer may have increased dimensionality. - """ - - n_experts: int - """ - Number of experts. - """ - - top_k: int = 1 - """ - top-k gating function (like top-1 or top-2) - """ - - input_dtype: DtypeEnum = DtypeEnum.fp16 - """ - Data type for the input activations. - """ - - output_dtype: DtypeEnum = DtypeEnum.fp16 - """ - Data type for the output activations. - """ - - activation: ActivationType = ActivationType.IDENTITY - """ - Activation function of the first MLP1 - """ - - normalize_scores: bool = False - """ - Whether normalization is applied to the selected scores. If true, the module - should rescale the scores such that their sum is 1.0. - """ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/norm_config.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/norm_config.py deleted file mode 100644 index 358982253756af4c4065d2e8cc53f7d0dd0b0287..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/norm_config.py +++ /dev/null @@ -1,32 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from ...inference_utils import DtypeEnum, NormTypeEnum -from ...modules.ds_module import DSModuleConfig - - -class DSNormConfig(DSModuleConfig): - """ - Config class for both DSPreLN and DSPostLN. - """ - - # Type of normalization - type: NormTypeEnum - - # Number of channels in the model embedding - channels: int - - # Data type of the residual input/outputs (we assume the residual must - # be the same data type for the entire model). - residual_dtype: DtypeEnum = DtypeEnum.fp16 - - # Data type of the hidden states input - input_dtype: DtypeEnum = DtypeEnum.fp16 - - # Data type of the hidden states output - output_dtype: DtypeEnum = DtypeEnum.fp16 - - # Epsilon value for numerical stability - eps: float = 1e-5 diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/unembed_config.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/unembed_config.py deleted file mode 100644 index ea4cc3cc99c17af9201db4774ebc0d095539dd82..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/configs/unembed_config.py +++ /dev/null @@ -1,39 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from ...inference_utils import DtypeEnum, NormTypeEnum -from ...modules.ds_module import DSModuleConfig -from typing import Optional - - -class DSUnembedConfig(DSModuleConfig): - """ - Config class for DSUnembed - """ - - dtype: DtypeEnum = DtypeEnum.fp16 - """ - Expected data type. - """ - - norm_type: Optional[NormTypeEnum] = None - """ - Whether the input to the unembed is normalized prior to the unembedding projection. - """ - - model_dim: int - """ - Model embedding size. - """ - - max_sequences: int - """ - Max sequences composing the ragged batch. - """ - - vocab_size: int - """ - Local vocab size (the full vocab size may have been sharded across model parallel ranks) - """ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/ds_module.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/ds_module.py deleted file mode 100644 index 2a6d294f32666cf46ab918ba03150d3f17ae37ff..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/ds_module.py +++ /dev/null @@ -1,62 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import ABC, abstractstaticmethod -from typing import Any, Dict, Type - -import torch - -from deepspeed.runtime.config_utils import DeepSpeedConfigModel - - -class DSModuleConfig(DeepSpeedConfigModel): - - max_tokens: int - - -class DSModuleBase(torch.nn.Module, ABC): - """ - Base class for all DeepSpeed Inference modules. This class establishes - the basic attributes of a DSModule. Only abstract functionality modules should inherit - directly from this class, not specific implementations. - """ - - @abstractstaticmethod - def name() -> str: - """ - Return a memorable, human-readable name for this module. - - This will be used as a key in custom inference configurations and should only - be implemented by the children of functionality modules. - """ - ... - - @abstractstaticmethod - def config_class() -> Type[DSModuleConfig]: - """ - Return the associated config class for this module. - - This should be implemented (along with the config class) by an abstract functionality - module. - """ - ... - - @abstractstaticmethod - def supports_config(config: DSModuleConfig) -> bool: - """ - Return whether or not this module supports the given config. - - This should be implemented by the children of functionality modules and should report - whether it would be feasible to instantiate this module with the given config. - """ - ... - - def __init__(self, config: DSModuleConfig, implementation_config: Dict[str, Any] = {}) -> None: - """ - Initialize the module with the given config. - """ - super().__init__() - self._config = config - self._implementation_config = implementation_config diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/heuristics.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/heuristics.py deleted file mode 100644 index f719e299a4b2eb99b40c58807a14d71ac2cc2319..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/heuristics.py +++ /dev/null @@ -1,195 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from ..config_v2 import RaggedInferenceEngineConfig -from ..inference_utils import NormTypeEnum - -from .module_registry import ConfigBundle -from ..modules.configs import ( - DSEmbeddingsConfig, - DSLinearConfig, - DSMoEConfig, - DSNormConfig, - DSSelfAttentionConfig, - DSUnembedConfig, -) -from ..modules.interfaces import ( - DSEmbeddingBase, - DSEmbeddingRegistry, - DSLinearBase, - DSLinearRegistry, - DSMoEBase, - DSMoERegistry, - DSPostNormBase, - DSPostNormRegistry, - DSPreNormBase, - DSPreNormRegistry, - DSSelfAttentionBase, - DSSelfAttentionRegistry, - DSUnembedBase, - DSUnembedRegistry, -) - - -def instantiate_attention(attention_config: DSSelfAttentionConfig, - engine_config: RaggedInferenceEngineConfig) -> DSSelfAttentionBase: - """ - Choose an appropriate attention implementation based on the given configurations. This - method is currently a stub, but as more implementations may be developed we can centralize - the logic for choosing between them here. - - Arguments: - attention_config (DSSelfAttentionConfig): Configuration for the attention module. - engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine. - - Returns: - An attention module implementing the given configuration. - """ - - # Currently, we only have one implementation, so we just return it. - config = ConfigBundle(name="dense_blocked_attention", config=attention_config) - return DSSelfAttentionRegistry.instantiate_config(config) - - -def instantiate_embed(embed_config: DSEmbeddingsConfig, engine_config: RaggedInferenceEngineConfig) -> DSEmbeddingBase: - """ - Choose an appropriate embedding implementation based on the given configurations. This - method is currently a stub, but as more implementations may be developed we can centralize - the logic for choosing between them here. - - Arguments: - embed_config (DSEmbeddingsConfig): Configuration for the embedding module. - engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine. - - Returns: - An embedding module implementing the given configuration. - """ - - # Currently, we only have one implementation, so we just return it. - config = ConfigBundle(name="ragged_embedding", config=embed_config) - return DSEmbeddingRegistry.instantiate_config(config) - - -def instantiate_linear(linear_config: DSLinearConfig, engine_config: RaggedInferenceEngineConfig) -> DSLinearBase: - """ - Choose an appropriate linear implementation based on the given configurations. This - method is currently a stub, but as more implementations may be developed we can centralize - the logic for choosing between them here. - - Arguments: - linear_config (DSLinearConfig): Configuration for the linear module. - engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine. - - Returns: - A linear module implementing the given configuration. - """ - - quantization_mode = engine_config.quantization.quantization_mode - if quantization_mode is None: - config = ConfigBundle(name="blas_fp_linear", config=linear_config) - else: - # Currently, we only support ``quantized_wf6af16_linear`` on NVIDIA Ampere GPUs. - if quantization_mode == "wf6af16": - import torch - if not torch.cuda.is_available(): #ignore-cuda - raise ValueError("WF6AF16 quantization is only supported on CUDA") - else: - is_rocm_pytorch = hasattr(torch.version, 'hip') and torch.version.hip is not None - if is_rocm_pytorch: - raise ValueError("WF6AF16 quantization is only supported on NVIDIA GPUs") - elif torch.cuda.get_device_properties(0).major != 8: #ignore-cuda - raise ValueError("WF6AF16 quantization is only supported on Ampere architectures") - config = ConfigBundle(name="quantized_wf6af16_linear", config=linear_config) - else: - raise ValueError(f"Unsupported quantization mode: {quantization_mode}") - return DSLinearRegistry.instantiate_config(config) - - -def instantiate_moe(moe_config: DSMoEConfig, engine_config: RaggedInferenceEngineConfig) -> DSMoEBase: - """ - Choose an appropriate MoE implementation based on the given configurations. This - method is currently a stub, but as more implementations may be developed we can centralize - the logic for choosing between them here. - - Arguments: - moe_config (DSMoEConfig): Configuration for the MoE module. - engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine. - - Returns: - A MoE module implementing the given configuration. - """ - - moe_type = "cutlass_multi_gemm_moe" - - if moe_type == "cutlass_multi_gemm_moe": - # TODO: Get this off an engine config - implementation_config = { - "weight_dtype": moe_config.input_dtype, - } - - # Currently, we only have one implementation, so we just return it. - config = ConfigBundle(name="cutlass_multi_gemm_moe", - config=moe_config, - implementation_config=implementation_config) - return DSMoERegistry.instantiate_config(config) - - -def instantiate_post_norm(norm_config: DSNormConfig, engine_config: RaggedInferenceEngineConfig) -> DSPostNormBase: - """ - Choose an appropriate post-norm implementation based on the given configurations. This - method is currently a stub, but as more implementations may be developed we can centralize - the logic for choosing between them here. - - Arguments: - norm_config (DSNormConfig): Configuration for the post-norm module. - engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine. - - Returns: - A post-norm module implementing the given configuration. - """ - - # Currently, we only have one implementation, so we just return it. - config = ConfigBundle(name="cuda_post_ln", config=norm_config) - return DSPostNormRegistry.instantiate_config(config) - - -def instantiate_pre_norm(norm_config: DSNormConfig, engine_config: RaggedInferenceEngineConfig) -> DSPreNormBase: - """ - Choose an appropriate pre-norm implementation based on the given configurations. Currently, - this will select between two CUDA implementations, one for LayerNorm and one for RMSNorm. - - Arguments: - norm_config (DSNormConfig): Configuration for the pre-norm module. - engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine. - - Returns: - A pre-norm module implementing the given configuration. - """ - if NormTypeEnum(norm_config.type) == NormTypeEnum.LayerNorm: - module_name = "cuda_pre_ln" - elif NormTypeEnum(norm_config.type) == NormTypeEnum.RMSNorm: - module_name = "cuda_pre_rms" - - config = ConfigBundle(name=module_name, config=norm_config) - return DSPreNormRegistry.instantiate_config(config) - - -def instantiate_unembed(unembed_config: DSUnembedConfig, engine_config: RaggedInferenceEngineConfig) -> DSUnembedBase: - """ - Choose an appropriate unembedding implementation based on the given configurations. This - method is currently a stub, but as more implementations may be developed we can centralize - the logic for choosing between them here. - - Arguments: - unembed_config (DSUnembedConfig): Configuration for the unembed module. - engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine. - - Returns: - An unembed module implementing the given configuration. - """ - - # Currently, we only have one implementation, so we just return it. - config = ConfigBundle(name="ragged_unembed", config=unembed_config) - return DSUnembedRegistry.instantiate_config(config) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/__init__.py deleted file mode 100644 index 1b500a9a0b5a25b17696bffcbcc5fa1ffd582de4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Imports for registering ops -from .attention import * -from .linear import * -from .post_norm import * -from .pre_norm import * -from .embedding import * -from .unembed import * -from .moe import * diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index b50cfc7cfe50aebeef0720b9e5b35af920cce863..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/attention/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/attention/__init__.py deleted file mode 100644 index 1ff593ac1368235cc03b427532154ec92b8bbe82..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/attention/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .dense_blocked_attention import DSDenseBlockedAttention diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/attention/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/attention/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index fe4cbaa4e7efcdb53bd35e8d84ef820697f8a23a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/attention/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/attention/__pycache__/dense_blocked_attention.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/attention/__pycache__/dense_blocked_attention.cpython-310.pyc deleted file mode 100644 index 63a362ca3ba109606d1f7ceb1a23cfa9c55d7159..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/attention/__pycache__/dense_blocked_attention.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/attention/dense_blocked_attention.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/attention/dense_blocked_attention.py deleted file mode 100644 index 3515b3c2b690c6346db7afba02b5790a2eb98982..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/attention/dense_blocked_attention.py +++ /dev/null @@ -1,180 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any, Dict, Optional - -import torch - -from deepspeed.accelerator import get_accelerator -from ....allocator import empty_from -from ....inference_utils import DtypeEnum -from ....kernels.ragged_ops import ( - AtomBuilder, - BlockedFlashAttn, - BlockedRotaryEmbeddings, - BlockedTrainedRotaryEmbeddings, - get_q_block_size, - get_kv_block_size, - LinearBlockedKVCopy, -) -from ....ragged import RaggedBatchWrapper, split_kv -from deepspeed.ops.op_builder import RaggedUtilsBuilder - -from ...interfaces import DSSelfAttentionBase, DSSelfAttentionRegistry -from ...configs import DSSelfAttentionConfig, PositionalEmbeddingType, MaskingType - -try: - from functools import cached_property -except ImportError: - - def cached_property(func): - return property(func) - - -@DSSelfAttentionRegistry.register_module -class DSDenseBlockedAttention(DSSelfAttentionBase): - """ - Self attention implementation for dense, blocked self attention. - """ - - @staticmethod - def name() -> str: - return 'dense_blocked_attention' - - @staticmethod - def supports_config(config: DSSelfAttentionConfig) -> bool: - - if config.input_dtype != config.output_dtype: - return False - - if DtypeEnum(config.input_dtype) not in (DtypeEnum.fp16, DtypeEnum.bf16): - return False - - if PositionalEmbeddingType(config.positional_embedding_type) not in [ - PositionalEmbeddingType.none, PositionalEmbeddingType.rotate_half - ]: - return False - - if MaskingType(config.masking_type) != MaskingType.causal: - return False - - return True - - def __init__(self, config: DSSelfAttentionConfig, implementation_config: Dict[str, Any]) -> None: - """ - Create the Attention DSModule. - - Args: - config (DSSelfAttentionConfig): The self attention config for all attention DSModules. - implementation_config (Dict[str, Any]): - There are two (dependent) potential components in the implementtion config. - - 1. `trained_freqs` - If the embedding weights for RoPE are trained, the implementation - config should contain {'trained_freqs': True}. This will mean the implementation will - expect a `trained_freqs` tensor in the `forward` method and will not synthesize the - values internally. - - 2. `theta_base` - The base value for synthesized frequencies in the rotary embeddings. - This will only be used if `trained_freqs` is False or not present in the `implementation_config`. If this is not included, the default value of 10000.0 will be used. - """ - super().__init__(config, implementation_config) - - embed_type = PositionalEmbeddingType(config.positional_embedding_type) - if embed_type == PositionalEmbeddingType.none: - self._kv_copy = LinearBlockedKVCopy(self._config.head_size, self._config.n_heads_q, - self._config.n_heads_kv, self._config.input_dtype) - elif embed_type == PositionalEmbeddingType.rotate_half: - rotary_config = config.positional_embedding_config - assert rotary_config is not None, "Rotary config must be provided if using rotate_half as Positional Embedding Type." - - if rotary_config.use_trained_freqs: - # Theta and rotary dim are effectively embedded into either the values (theta) or the shape (rotary_dim) - # of the trained_freqs tensor. - self._kv_copy = BlockedTrainedRotaryEmbeddings(self._config.head_size, self._config.n_heads_q, - self._config.n_heads_kv, self._config.input_dtype) - else: - theta_base = rotary_config.theta_base - rotary_dim = rotary_config.rotate_dim if rotary_config.rotate_dim is not None else self._config.head_size - self._kv_copy = BlockedRotaryEmbeddings(self._config.head_size, self._config.n_heads_q, - self._config.n_heads_kv, self._config.input_dtype, rotary_dim, - theta_base) - - self._softmax_scale = self._config.scale_factor - - # TODO(cmikeh2): Attention kernel gets created here. - self._attn_kernel = BlockedFlashAttn(self._config.head_size, self._config.input_dtype) - self._atom_builder = AtomBuilder() - - self.model_dim = self._config.head_size * self._config.n_heads_q - self._output = torch.empty((self._config.max_tokens, self._config.head_size * self._config.n_heads_q), - dtype=self._config.output_dtype, - device=get_accelerator().current_device()) - - # TODO(cmikeh2): Pre-allocate storage buffer for the attention atoms. - self._max_atoms = self._config.max_sequences - self._atoms = torch.empty((self._max_atoms, 8), dtype=torch.int32, device=get_accelerator().current_device()) - - alloc_func = RaggedUtilsBuilder().load().allocate_fast_host_buffer - self._atoms_shadow = alloc_func(self._atoms) - self._cur_atoms = 0 - - @cached_property - def kv_block_size(self) -> int: - """ - Return preferred granulatity for blocked KV-cache implementation. - """ - return get_kv_block_size(self._config.head_size) - - @cached_property - def q_block_size(self) -> int: - """ - Property to calculate blocking granularity for the query dimension. - This has no impact on the KV-cache structure, but will affect the - number of attention atoms associated with a batch. - """ - return get_q_block_size(self._config.head_size) - - def build_atoms(self, ragged_batch: RaggedBatchWrapper) -> None: - """ - Build the atoms for the attention kernel. - - Args: - ragged_batch (RaggedBatchWrapper): The input ids and associated ragged batch metadata. - """ - host_atoms, n_atoms = self._atom_builder(self._atoms_shadow, ragged_batch, self.q_block_size, - self.kv_block_size) - - self._cur_atoms = n_atoms - self._atoms[:n_atoms].copy_(host_atoms[:n_atoms], non_blocking=True) - - def forward(self, - q_k_v: torch.Tensor, - kv_cache: torch.Tensor, - batch: RaggedBatchWrapper, - inv_freqs: Optional[torch.Tensor] = None) -> torch.Tensor: - """ - Forward implementation. - - Args: - q_k_v (torch.Tensor): Query/Key/Value projection Tensor of shape - [n_heads, (n_heads_q + 2 * n_heads_kv) * head_size]. - kv_cache (torch.Tensor): Blocked persistent cache of shape - [2, batch, block_size, n_heads_kv, head_size]. - batch (RaggedBatchWrapper): The input ids and associated ragged batch metadata. - inv_freqs (Optional[torch.Tensor]): The inverse frequencies for the rotary embeddings if they - have been modified from synthesizable values. - """ - if inv_freqs is not None: - self._kv_copy(kv_cache, q_k_v, batch, inv_freqs) - else: - self._kv_copy(kv_cache, q_k_v, batch) - - q = q_k_v[:, :self._config.head_size * self._config.n_heads_q] - output = empty_from(self._output, q.shape) - k_cache, v_cache = split_kv(kv_cache) - - self._attn_kernel(output, q, k_cache, v_cache, self._atoms[:self._cur_atoms], self._softmax_scale) - - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/embedding/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/embedding/__init__.py deleted file mode 100644 index 5458a855abf45c3870c2208663dd5f49fb4fb22b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/embedding/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .ragged_embedding import DSRaggedEmbedding diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/embedding/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/embedding/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 095e890f8e82b748aa4f867d3dd853a370affc81..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/embedding/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/embedding/__pycache__/ragged_embedding.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/embedding/__pycache__/ragged_embedding.cpython-310.pyc deleted file mode 100644 index 274997a07deff2a7c9f6768917bc49f2b0028a97..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/embedding/__pycache__/ragged_embedding.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/embedding/ragged_embedding.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/embedding/ragged_embedding.py deleted file mode 100644 index 90cdd39d1be7f7e51da65a1a517325e585386a72..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/embedding/ragged_embedding.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any, Dict, Optional - -import torch - -from deepspeed.accelerator import get_accelerator -from ....allocator import empty_from -from ....inference_utils import DtypeEnum -from ....kernels.ragged_ops import RaggedEmbeddingKernel -from ....ragged import RaggedBatchWrapper -from ...interfaces import DSEmbeddingBase, DSEmbeddingRegistry -from ...configs import DSEmbeddingsConfig - - -@DSEmbeddingRegistry.register_module -class DSRaggedEmbedding(DSEmbeddingBase): - - @staticmethod - def name(): - return 'ragged_embedding' - - @staticmethod - def supports_config(config: DSEmbeddingsConfig) -> bool: - - if DtypeEnum(config.residual_dtype) not in [DtypeEnum.fp16, DtypeEnum.bf16, DtypeEnum.fp32]: - return False - - if config.use_token_type: - return False - - if config.output_normalization is not None: - return False - - try: - _ = RaggedEmbeddingKernel(config.residual_dtype, torch.int32, config.embedding_dim) - except ValueError: - return False - - return True - - def __init__(self, config: DSEmbeddingsConfig, implementation_config: Dict[str, Any]) -> None: - super().__init__(config, implementation_config) - - self.embed_offset = self._config.positional_offset - - # TODO(cmikeh2): How do we want to avoid the int32 vs int64 issue? - self._ragged_embed = RaggedEmbeddingKernel(self._config.residual_dtype, torch.int32, - self._config.embedding_dim) - - self._output = torch.empty((self._config.max_tokens, self._config.embedding_dim), - dtype=self._config.residual_dtype, - device=get_accelerator().current_device()) - - @property - def output(self) -> torch.Tensor: - return self._output - - def forward(self, - ragged_batch: RaggedBatchWrapper, - word_embeddings: torch.Tensor, - position_embeddings: Optional[torch.Tensor] = None) -> torch.Tensor: - """ - Parameters: - ragged_batch (RaggedBatchWrapper): The input ids and associated ragged batch metadata. - word_embeddings (torch.Tensor): The word embedding table - """ - output = empty_from(self._output, (ragged_batch.tensor_toks, self._config.embedding_dim)) - self._ragged_embed(output, - ragged_batch, - word_embeddings, - position_embed_weight=position_embeddings, - position_embed_offset=self.embed_offset) - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/__init__.py deleted file mode 100644 index 0501af54c4e6dfb89ee0db49d56e46809ae7945c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .blas_fp_linear import BlasFPLinear -from .quantized_linear import QuantizedWf6Af16Linear, fp_quantize diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 14c610642e162addc5e8a570bf95da0533427390..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/__pycache__/blas_fp_linear.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/__pycache__/blas_fp_linear.cpython-310.pyc deleted file mode 100644 index 4093d65a73adeebb97f98db2e3da2490db522ca7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/__pycache__/blas_fp_linear.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/__pycache__/quantized_linear.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/__pycache__/quantized_linear.cpython-310.pyc deleted file mode 100644 index 4f5657a006b0670000230da56cfbc62cd5b0243e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/__pycache__/quantized_linear.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/blas_fp_linear.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/blas_fp_linear.py deleted file mode 100644 index c58dab0b826b48ed5ac4b58217936a2f9de4ae29..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/blas_fp_linear.py +++ /dev/null @@ -1,103 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any, Dict, Optional - -import torch - -from deepspeed.accelerator import get_accelerator -from ....allocator import empty_from -from ....inference_utils import is_gated -from ....kernels.core_ops import ( - BlasLibLinear, - CUDABiasActivation, - CUDAGatedActivation, -) - -from ...interfaces import DSLinearBase, DSLinearRegistry -from ...configs import DSLinearConfig -from ....inference_parameter import InferenceParameter - - -@DSLinearRegistry.register_module -class BlasFPLinear(DSLinearBase): - """ - Linear DSModule based on BLAS library and standalone bias + activation kernel implementation. - """ - - @staticmethod - def name(): - return 'blas_fp_linear' - - @staticmethod - def supports_config(config: DSLinearConfig) -> bool: - if config.input_dtype != config.output_dtype: - return False - - if config.input_dtype != torch.float16 and config.input_dtype != torch.bfloat16: - return False - - if is_gated(config.activation): - try: - _ = CUDAGatedActivation(config.out_channels, config.output_dtype, config.activation) - except ValueError: - return False - else: - try: - _ = CUDABiasActivation(config.out_channels, config.output_dtype, config.activation) - except ValueError: - return False - - return True - - def __init__(self, config: DSLinearConfig, implementation_config: Dict[str, Any]) -> None: - super().__init__(config, implementation_config) - - self._linear_impl = BlasLibLinear(self._config.input_dtype) - - if is_gated(config.activation): - self._is_gated = True - self._act_fn = CUDAGatedActivation(config.out_channels, config.output_dtype, config.activation) - self._double_buffer = torch.empty((config.max_tokens, config.out_channels * 2), - dtype=config.output_dtype, - device=get_accelerator().current_device()) - else: - self._is_gated = False - self._act_fn = CUDABiasActivation(config.out_channels, config.output_dtype, config.activation) - - self._output = torch.empty((config.max_tokens, config.out_channels), - dtype=config.output_dtype, - device=get_accelerator().current_device()) - - def transform_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Converts param to same data type as input and output. - - Parameters: - param (torch.Tensor): Weight or bias tensor. - """ - param = param.to(self._config.output_dtype) - return InferenceParameter.initialize(param) - - def forward(self, hidden_states: torch.Tensor, w: torch.Tensor, b: Optional[torch.Tensor] = None) -> torch.Tensor: - - output = empty_from(self._output, (hidden_states.shape[0], self._config.out_channels)) - - if self._is_gated: - staging_output = empty_from(self._double_buffer, (hidden_states.shape[0], self._config.out_channels * 2)) - self._linear_impl(staging_output, hidden_states, w) - self._act_fn(output, staging_output, b) - else: - self._linear_impl(output, hidden_states, w) - self._act_fn(output, b) - - return output - - @property - def output(self) -> torch.Tensor: - """ - Return the padded, pre-allocated output Tensor. - """ - return self._output diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/quantized_linear.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/quantized_linear.py deleted file mode 100644 index 933cf55b2391b0132f392e770d0b747f357f8ca0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/linear/quantized_linear.py +++ /dev/null @@ -1,205 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any, Dict, Optional - -import torch - -from deepspeed.accelerator import get_accelerator -from deepspeed.ops.op_builder import InferenceCoreBuilder -from ....allocator import empty_from -from ....inference_utils import is_gated -from ....kernels.core_ops import ( - CUDAWf6Af16Linear, - CUDABiasActivation, - CUDAGatedActivation, -) - -from ...interfaces import DSLinearBase, DSLinearRegistry -from ...configs import DSLinearConfig -from ....inference_parameter import InferenceParameter - - -def fp_quantize(input: torch.FloatTensor, - num_bits: int = 6, - exp_bits: int = 3, - min_value: torch.FloatTensor = None, - max_value: torch.FloatTensor = None, - group_size: int = -1): - """ - Args: - inputs (`torch.FloatTensor`) - The input which needs to be quantized - num_bits (int, >=4) - Number of bits to use for quantization - exp_bits: - fp exp_bits - min_value/max_vlue (torch.FloatTensor) - Used for static activation quantization - group_size (int) N - The quantization block size, each N numbers has its own scaling - factor and off-site. -1 means use the last dim as the group_size - Returns: - quantized_fake_fp6 - The quantized weights, in fp16 format and contains fp6 value. - scales - Quantization scales - """ - - try: - from qtorch.quant import float_quantize - except ImportError: - raise ImportError("Please install qtorch to use this function") - - assert (min_value is None and max_value is None) or (min_value is not None and max_value is not None) - - assert input.dtype == torch.float16 - - orig_device = input.device - input = input.to(torch.float32).to(get_accelerator().current_device()) - if num_bits == 6 and exp_bits == 3: # this is default - q_range = 28 - else: - raise NotImplementedError - - man_bits = num_bits - exp_bits - 1 - input_shape = input.shape - - if group_size == -1: - group_size = input_shape[-1] - else: - # Only support per-channel quantization - raise NotImplementedError - num_groups = input.numel() // group_size - input = input.reshape(num_groups, -1) - - if min_value is None: - max_input = torch.amax(torch.abs(input), dim=-1).view(num_groups, -1) - else: - max_input = torch.max(min_value.abs(), max_value) # .view(-1) - scales = max_input / q_range # q_range + 1 - scales[scales == 0] = 1 # avoid zero scales - scaled_input = input / scales - - quantized_fake_fp6 = float_quantize(scaled_input, exp_bits, man_bits, rounding="nearest") - - quantized_fake_fp6 = quantized_fake_fp6.reshape(input_shape).contiguous().to(torch.float16).to(orig_device) - scales = scales.to(torch.float16).to(orig_device) - # Now the dequantized value is quantized_fake_fp6 * scales - - return quantized_fake_fp6, scales - - -@DSLinearRegistry.register_module -class QuantizedWf6Af16Linear(DSLinearBase): - """ - Linear DSModule for FP6 weight-only quantization kernel, where weight is FP6 - and activation is FP16. - """ - - @staticmethod - def name(): - return 'quantized_wf6af16_linear' - - @staticmethod - def supports_config(config: DSLinearConfig) -> bool: - if config.input_dtype != config.output_dtype: - return False - - # As for fp6 data items, they are packed and stored in a set of fp16 - # tensors. E.g., 8 fp6 data items are stored in 3 fp16 tensor. - if config.input_dtype != torch.float16: - return False - - if is_gated(config.activation): - try: - _ = CUDAGatedActivation(config.out_channels, config.output_dtype, config.activation) - except ValueError: - return False - else: - try: - _ = CUDABiasActivation(config.out_channels, config.output_dtype, config.activation) - except ValueError: - return False - - return True - - def __init__(self, config: DSLinearConfig, implementation_config: Dict[str, Any]) -> None: - super().__init__(config, implementation_config) - - self._linear_impl = CUDAWf6Af16Linear() - - if is_gated(config.activation): - # In the FP6 kernel implementation, the MatMul is W * A, where W is - # the weight and A is activation. M is the output channel size. - self.out_channels = self._config.out_channels * 2 - self.in_channels = self._config.in_channels - self._is_gated = True - self._act_fn = CUDAGatedActivation(config.out_channels, config.output_dtype, config.activation) - self._double_buffer = torch.empty((config.max_tokens, config.out_channels * 2), - dtype=config.output_dtype, - device=get_accelerator().current_device()) - else: - self.out_channels = self._config.out_channels - self.in_channels = self._config.in_channels - self._is_gated = False - self._act_fn = CUDABiasActivation(config.out_channels, config.output_dtype, config.activation) - - self._output = torch.empty((config.max_tokens, config.out_channels), - dtype=config.output_dtype, - device=get_accelerator().current_device()) - - self.inf_module = InferenceCoreBuilder().load() - self.inf_module.create_handle() - self.preprocess_weight = self.inf_module.preprocess_weight - - self.quantizer = fp_quantize - - def transform_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Converts param to same data type as input and output. - - Parameters: - param (torch.Tensor): Weight or bias tensor. - """ - # It expects that the quantization scales are store in the attribute `scales`. - - if param.ndim == 1: # bias, do nothing - return InferenceParameter.initialize(param) - - quantized_fake_fp6, scales = self.quantizer(param, num_bits=6, exp_bits=3) - - # This is for debugging, will delete before release. - assert (quantized_fake_fp6.dtype == torch.float16) - assert quantized_fake_fp6.shape[0] == self.out_channels - assert scales.numel() == self.out_channels - - weights_2bit, weights_4bit = self.preprocess_weight(quantized_fake_fp6) - - return InferenceParameter.initialize(weights_2bit, weights_4bit=weights_4bit, scales=scales) - - def forward(self, hidden_states: torch.Tensor, w: torch.Tensor, b: Optional[torch.Tensor] = None) -> torch.Tensor: - weights_2bit = w - weights_4bit = w.weights_4bit - scales = w.scales - output = empty_from(self._output, (hidden_states.shape[0], self._config.out_channels)) - if self._is_gated: - staging_output = empty_from(self._double_buffer, (hidden_states.shape[0], self.out_channels)) - self._linear_impl(staging_output, hidden_states, weights_2bit, weights_4bit, scales, self.out_channels, - hidden_states.shape[0], self.in_channels) - self._act_fn(output, staging_output, b) - else: - self._linear_impl(output, hidden_states, weights_2bit, weights_4bit, scales, self.out_channels, - hidden_states.shape[0], self.in_channels) - self._act_fn(output, b) - - return output - - @property - def output(self) -> torch.Tensor: - """ - Return the padded, pre-allocated output Tensor. - """ - return self._output diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/moe/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/moe/__init__.py deleted file mode 100644 index 053ad5da77460974cbc34361dcd887c5ea2851d8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/moe/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .cutlass_multi_gemm import DSMultiGemmMoE diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/moe/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/moe/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 18c4e22921022c08e7a89c47c5f4358c2a10427e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/moe/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/moe/__pycache__/cutlass_multi_gemm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/moe/__pycache__/cutlass_multi_gemm.cpython-310.pyc deleted file mode 100644 index 498ca27d6647084c92cb6c4709b9a8d4c38caff2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/moe/__pycache__/cutlass_multi_gemm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/moe/cutlass_multi_gemm.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/moe/cutlass_multi_gemm.py deleted file mode 100644 index a9b01d1233cd9af3aa9608ffedcdac13a031e411..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/moe/cutlass_multi_gemm.py +++ /dev/null @@ -1,249 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any, Dict, Optional, Tuple - -import torch - -from deepspeed.accelerator import get_accelerator -from ....allocator import empty_from -from ....inference_utils import ActivationType, is_gated -from ....kernels.core_ops import BlasLibLinear, CUDAGatedActivation -from ....kernels.ragged_ops import ( - MoEGather, - MoEScatter, - RaggedTopKGating, -) -from ....ragged import RaggedBatchWrapper - -from ...interfaces import DSMoEBase, DSMoERegistry -from ...configs import DSMoEConfig -from ....kernels.cutlass_ops import MoEGEMM -from ....inference_parameter import InferenceParameter - - -@DSMoERegistry.register_module -class DSMultiGemmMoE(DSMoEBase): - """ - MoE implementation based on the CUTLASS multi-GEMM. - """ - - @staticmethod - def name(): - return 'cutlass_multi_gemm_moe' - - @staticmethod - def supports_config(config: DSMoEConfig) -> bool: - if config.input_dtype != config.output_dtype: - return False - - if config.input_dtype != torch.float16 and config.input_dtype != torch.bfloat16: - return False - - if config.top_k != 1 and config.top_k != 2 and config.top_k != 4 and config.top_k != 8: - return False - - return True - - def __init__(self, config: DSMoEConfig, implementation_config: Dict[str, Any]) -> None: - super().__init__(config, implementation_config) - - # Convenience variables for frequently accessed items. - self.max_tokens = self._config.max_tokens - self.n_experts = self._config.n_experts - self.n_top_k = self._config.top_k - self.intermediate_dim = self._config.intermediate_features - - moe_op_act_fn = ActivationType.IDENTITY if is_gated(self._config.activation) else self._config.activation - - self._mlp_1 = MoEGEMM(fp_dtype=implementation_config['weight_dtype'], act_fn=moe_op_act_fn) - self._mlp_2 = MoEGEMM(fp_dtype=implementation_config['weight_dtype'], act_fn=ActivationType.IDENTITY) - - if is_gated(self._config.activation): - self._activation = CUDAGatedActivation(self._config.model_dim, self._config.input_dtype, - self._config.activation) - else: - self._activation = None - - self._gate_proj = BlasLibLinear(self._config.input_dtype) - self._top_1_gate = RaggedTopKGating(config.input_dtype) - self._moe_scatter = MoEScatter(config.input_dtype, config.model_dim) - self._moe_gather = MoEGather(config.input_dtype, config.model_dim, config.normalize_scores) - - self._create_buffers() - - def _create_buffers(self): - - # Gating buffers - self._logits = torch.empty((self._config.max_tokens, self.n_experts), - dtype=self._config.input_dtype, - device=get_accelerator().current_device()) - self._expert_counts = torch.empty((self.n_experts, ), - dtype=torch.int32, - device=get_accelerator().current_device()) - self._scores = torch.empty((self._config.max_tokens, self.n_top_k), - dtype=torch.float32, - device=get_accelerator().current_device()) - self._assignments = torch.empty((self._config.max_tokens, self.n_top_k), - dtype=torch.int32, - device=get_accelerator().current_device()) - self._offsets = torch.empty((self._config.max_tokens, self.n_top_k), - dtype=torch.int32, - device=get_accelerator().current_device()) - - # Scatter buffers - self._moe_input = torch.empty((self._config.max_tokens * self.n_top_k, self._config.model_dim), - dtype=self._config.input_dtype, - device=get_accelerator().current_device()) - self._expert_cumsum = torch.empty((self._config.n_experts, ), - dtype=torch.int64, - device=get_accelerator().current_device()) - self._mapped_slots = torch.empty((self._config.max_tokens, self.n_top_k), - dtype=torch.int32, - device=get_accelerator().current_device()) - - # GEMM Buffers - self._intermediate = torch.empty((self._config.max_tokens * self.n_top_k, self._config.intermediate_features), - dtype=self._config.output_dtype, - device=get_accelerator().current_device()) - if self._activation is not None: - self._gated_intermediate = torch.empty( - (self._config.max_tokens * self.n_top_k, self._config.intermediate_features * 2), - dtype=self._config.output_dtype, - device=get_accelerator().current_device()) - - self._output_unordered = torch.empty((self._config.max_tokens * self.n_top_k, self._config.model_dim), - dtype=self._config.output_dtype, - device=get_accelerator().current_device()) - - # Gather buffer - self._output = torch.empty((self._config.max_tokens, self._config.model_dim), - dtype=self._config.output_dtype, - device=get_accelerator().current_device()) - - def transform_gate_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Ensures gate param is going to match the activation data type. - """ - param = param.to(self._config.input_dtype) - return InferenceParameter.initialize(param) - - def transform_moe_mlp_1_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Converts param to same data type as input and output. - - Parameters: - param (torch.Tensor): Weight or bias tensor. - """ - param = param.to(self._config.input_dtype) - - if len(param.shape) == 3: - param = param.permute(0, 2, 1).contiguous() - return InferenceParameter.initialize(param) - - def transform_moe_mlp_2_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Converts param to same data type as input and output. - - Parameters: - param (torch.Tensor): Weight or bias tensor. - """ - param = param.to(self._config.input_dtype) - - if len(param.shape) == 3: - param = param.permute(0, 2, 1).contiguous() - return InferenceParameter.initialize(param) - - @property - def output(self) -> torch.Tensor: - return self._output - - def _gate(self, hidden_states: torch.Tensor, batch_metadata: RaggedBatchWrapper, - gate_w: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Helper function to isolate the logit for gating. This will take the hidden states - and produce the metadata + tensors for the CUTLASS ragged GEMMs. If the input has - been padded for CG, this will strip the padding for MoE. - - Parameters: - hidden_states (torch.Tensor): Hidden states tensor. Expected shape is [n_tokens, model_dim]. - batch_metadata (RaggedBatchWrapper): Batch metadata for the hidden states. - gate_w (torch.Tensor): Gate weight tensor. Expected shape is [num_experts, model_dim]. - - Returns: - Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: The MoE input, the cumsum of the offsets (for the MoE kernels themselves), the scores, and the mapped slots (to recover the original order of the tokens) - """ - - # Get views on the buffers for gating - logits = empty_from(self._logits, (hidden_states.shape[0], self._logits.shape[-1])) - scores = empty_from(self._scores, (hidden_states.shape[0], self.n_top_k)) - assignments = empty_from(self._assignments, (hidden_states.shape[0], self.n_top_k)) - offsets = empty_from(self._offsets, (hidden_states.shape[0], self.n_top_k)) - mapped_slots = empty_from(self._mapped_slots, (hidden_states.shape[0], self.n_top_k)) - moe_input = empty_from(self._moe_input, (hidden_states.shape[0] * self.n_top_k, self._moe_input.shape[-1])) - - self._gate_proj(logits, hidden_states, gate_w) - self._expert_counts.zero_() - self._top_1_gate(self._expert_counts, scores, assignments, offsets, logits, batch_metadata) - self._moe_scatter(moe_input, self._expert_cumsum, mapped_slots, hidden_states, self._expert_counts, - assignments, offsets) - - return moe_input, self._expert_cumsum, scores, mapped_slots - - def forward(self, - hidden_states: torch.Tensor, - batch_metadata: RaggedBatchWrapper, - gate_w: torch.Tensor, - mlp_1_w: torch.Tensor, - mlp_2_w: torch.Tensor, - mlp_1_b: Optional[torch.Tensor] = None, - mlp_2_b: Optional[torch.Tensor] = None) -> torch.Tensor: - """ - MoE forward pass built on top of CUTLASS multi-GEMM. - - Parameters: - hidden_states (torch.Tensor): Hidden states tensor. Expected shape is [batch, seq_len, model_dim]. - gate_w (torch.Tensor): Gate weight tensor. Expected shape is [num_experts, model_dim]. - """ - - moe_input, expert_cumsum, scores, mapped_slots = self._gate(hidden_states, batch_metadata, gate_w) - - # Get views on the buffers for GEMM - intermediate = empty_from(self._intermediate, - (hidden_states.shape[0] * self.n_top_k, self._intermediate.shape[-1])) - output_unordered = empty_from(self._output_unordered, - (hidden_states.shape[0] * self.n_top_k, self._output_unordered.shape[-1])) - output = empty_from(self._output, (hidden_states.shape[0], self._output.shape[-1])) - - if self._activation is not None: - gated_intermediate = empty_from( - self._gated_intermediate, (hidden_states.shape[0] * self.n_top_k, self._gated_intermediate.shape[-1])) - self._mlp_1( - gated_intermediate, - moe_input, - mlp_1_w, - expert_cumsum, - mlp_1_b, - ) - self._activation(intermediate, gated_intermediate) - else: - self._mlp_1( - intermediate, - moe_input, - mlp_1_w, - expert_cumsum, - mlp_1_b, - ) - - self._mlp_2( - output_unordered, - intermediate, - mlp_2_w, - expert_cumsum, - mlp_2_b, - ) - - self._moe_gather(output, output_unordered, scores, mapped_slots, self._expert_counts) - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/post_norm/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/post_norm/__init__.py deleted file mode 100644 index 653a2fe4fb5be742c873aa51e76db79ead7325df..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/post_norm/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .cuda_post_ln import DSPostLNCUDAModule diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/post_norm/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/post_norm/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 372a53d51aa77af6525f57896d64d1033e664376..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/post_norm/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/post_norm/__pycache__/cuda_post_ln.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/post_norm/__pycache__/cuda_post_ln.cpython-310.pyc deleted file mode 100644 index 44a4fec6040f58e6ebcec24e725668508d3cba5d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/post_norm/__pycache__/cuda_post_ln.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/post_norm/cuda_post_ln.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/post_norm/cuda_post_ln.py deleted file mode 100644 index 9b2af4bb90231cd3f37e3406fe94042c4006faa3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/post_norm/cuda_post_ln.py +++ /dev/null @@ -1,56 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any, Dict, Tuple - -import torch - -from deepspeed.accelerator import get_accelerator -from ...interfaces import DSPostNormBase, DSPostNormRegistry -from ...configs import DSNormConfig -from ....kernels.core_ops.cuda_layer_norm.cuda_post_ln import CUDAFPPostLN -from ....allocator import empty_from -from ....inference_parameter import InferenceParameter - - -@DSPostNormRegistry.register_module -class DSPostLNCUDAModule(DSPostNormBase): - - @staticmethod - def name(): - return 'cuda_post_ln' - - @staticmethod - def supports_config(config: DSNormConfig): - if len(set([config.residual_dtype, config.input_dtype, config.output_dtype])) != 1: - return False - - try: - _ = CUDAFPPostLN(config.channels, config.residual_dtype) - except ValueError: - return False - return True - - def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]): - super().__init__(config, implementation_config) - self._fp_post_ln = CUDAFPPostLN(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps) - - self._output = torch.empty((config.max_tokens, config.channels), - dtype=config.output_dtype, - device=get_accelerator().current_device()) - - def transform_param(self, param: torch.Tensor) -> InferenceParameter: - param = param.to(self._config.input_dtype) - return InferenceParameter.initialize(param) - - def forward(self, residual: torch.Tensor, hidden_in: torch.Tensor, gamma: torch.Tensor, - beta: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Since the CUDA FP only supports all data types being the same, we will alias the residual - with our output. - """ - self._residual_output = empty_from(self._output, residual.shape) - self._fp_post_ln(residual, residual, hidden_in, gamma, beta) - return residual, residual diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/__init__.py deleted file mode 100644 index 12605f13f955e00b598516a97a8cc7d0369f7405..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .cuda_pre_ln import DSPreLNCUDAModule -from .cuda_pre_rms import DSPreRMSCUDAModule diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index b79f937cc487e2037eebb6d16daa2541aff12943..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/cuda_pre_ln.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/cuda_pre_ln.cpython-310.pyc deleted file mode 100644 index f10c8390427df0172f9ca01f441bc1ff2df9068a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/cuda_pre_ln.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/cuda_pre_rms.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/cuda_pre_rms.cpython-310.pyc deleted file mode 100644 index 43189a8e11f211b7bb16b9c7518b09a46631de89..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/__pycache__/cuda_pre_rms.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/cuda_pre_ln.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/cuda_pre_ln.py deleted file mode 100644 index 90783ce8c9a63320fc2047b4e8dfed9522df7899..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/cuda_pre_ln.py +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any, Dict, Optional, Tuple - -import torch - -from deepspeed.accelerator import get_accelerator -from ...interfaces import DSPreNormBase, DSPreNormRegistry -from ...configs import DSNormConfig, NormTypeEnum -from ....kernels.core_ops.cuda_layer_norm.cuda_pre_ln import CUDAFPPreLN -from ....kernels.core_ops.cuda_layer_norm.cuda_ln import CUDAFPLN -from ....allocator import empty_from -from ....inference_parameter import InferenceParameter - - -@DSPreNormRegistry.register_module -class DSPreLNCUDAModule(DSPreNormBase): - - @staticmethod - def name(): - return 'cuda_pre_ln' - - @staticmethod - def supports_config(config: DSNormConfig): - type = NormTypeEnum(config.type) - if type != NormTypeEnum.LayerNorm: - return False - - if len(set([config.residual_dtype, config.input_dtype, config.output_dtype])) != 1: - return False - - try: - _ = CUDAFPPreLN(config.channels, config.residual_dtype) - except ValueError: - return False - return True - - def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]): - super().__init__(config, implementation_config) - self._fp_pre_ln = CUDAFPPreLN(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps) - self._fp_ln = CUDAFPLN(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps) - - # Buffers for the hidden output (residual is updated in-place) - self._hidden_output = torch.empty((config.max_tokens, config.channels), - dtype=config.output_dtype, - device=get_accelerator().current_device()) - - def transform_param(self, param: torch.Tensor) -> InferenceParameter: - param = param.to(self._config.input_dtype) - return InferenceParameter.initialize(param) - - def forward(self, residual: torch.Tensor, hidden_in: Optional[torch.Tensor], gamma: torch.Tensor, - beta: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Since the CUDA FP only supports all data types being the same, we will alias the residual - with our output. - - If hidden_in is None, that means we do not need to perform the residual add and will - only return the hidden output modified. - """ - hidden_out = empty_from(self._hidden_output, residual.shape) - if hidden_in is None: - self._fp_ln(hidden_out, residual, gamma, beta) - else: - self._fp_pre_ln(residual, hidden_out, residual, hidden_in, gamma, beta) - return residual, hidden_out diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/cuda_pre_rms.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/cuda_pre_rms.py deleted file mode 100644 index 986262b31b1f50c3c139e7f71aca94c13ed866c3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/pre_norm/cuda_pre_rms.py +++ /dev/null @@ -1,79 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any, Dict, Optional, Tuple - -import torch - -from deepspeed.accelerator import get_accelerator -from ...interfaces import DSPreNormBase, DSPreNormRegistry -from ...configs import DSNormConfig, NormTypeEnum -from ....kernels.core_ops import CUDARMSNorm, CUDARMSPreNorm -from ....allocator import empty_from -from ....inference_parameter import InferenceParameter - - -@DSPreNormRegistry.register_module -class DSPreRMSCUDAModule(DSPreNormBase): - - @staticmethod - def name(): - return 'cuda_pre_rms' - - @staticmethod - def supports_config(config: DSNormConfig): - type = NormTypeEnum(config.type) - if type != NormTypeEnum.RMSNorm: - return False - - if len(set([config.residual_dtype, config.input_dtype, config.output_dtype])) != 1: - return False - - try: - # Only need to check one since the support matrix for the two rms kernels is the same - _ = CUDARMSPreNorm(config.channels, config.residual_dtype) - except ValueError: - return False - return True - - def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]): - super().__init__(config, implementation_config) - self._fp_rms = CUDARMSNorm(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps) - self._fp_rms_pre = CUDARMSPreNorm(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps) - - # Buffers for both the hidden and residual outputs - self._hidden_output = torch.empty((config.max_tokens, config.channels), - dtype=config.output_dtype, - device=get_accelerator().current_device()) - self._residual_output = torch.empty((config.max_tokens, config.channels), - dtype=config.output_dtype, - device=get_accelerator().current_device()) - - def transform_param(self, param: torch.Tensor) -> InferenceParameter: - param = param.to(self._config.input_dtype) - return InferenceParameter.initialize(param) - - def forward(self, - residual: torch.Tensor, - hidden_in: Optional[torch.Tensor], - gamma: torch.Tensor, - beta: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Since the CUDA FP only supports all data types being the same, we will alias the residual - with our output. - - If hidden_in is None, that means we do not need to perform the residual add and will - only return the hidden output modified. - """ - assert beta is None, "Beta is not supported for RMSNorm" - - hidden_out = empty_from(self._hidden_output, residual.shape) - if hidden_in is None: - self._fp_rms(hidden_out, residual, gamma) - residual_out = residual - else: - residual_out = empty_from(self._residual_output, residual.shape) - self._fp_rms_pre(residual_out, hidden_out, residual, hidden_in, gamma) - return residual_out, hidden_out diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/unembed/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/unembed/__init__.py deleted file mode 100644 index 4a5fd24d518bbf29f88771478b04da73399d4a53..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/unembed/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .ragged_unembed import DSRaggedUnembed diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/unembed/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/unembed/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index fa208955658fec909ac4545552f400e78b027dc5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/unembed/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/unembed/__pycache__/ragged_unembed.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/unembed/__pycache__/ragged_unembed.cpython-310.pyc deleted file mode 100644 index ddc1503cf46087a66a488b542e414bbf7cfda34b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/unembed/__pycache__/ragged_unembed.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/unembed/ragged_unembed.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/unembed/ragged_unembed.py deleted file mode 100644 index 36130902c665c7adff419ef0ea2be7138292e193..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/implementations/unembed/ragged_unembed.py +++ /dev/null @@ -1,120 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any, Dict, Optional - -import torch - -from deepspeed.accelerator import get_accelerator -from ....allocator import empty_from -from ....inference_utils import DtypeEnum, ActivationType -from ....kernels.core_ops import CUDAFPLN, BlasLibLinear, CUDARMSNorm, CUDABiasActivation -from ....kernels.ragged_ops import RaggedLogitsGather -from ....ragged import RaggedBatchWrapper -from ...interfaces import DSUnembedBase, DSUnembedRegistry -from ...configs import DSUnembedConfig - - -@DSUnembedRegistry.register_module -class DSRaggedUnembed(DSUnembedBase): - """ - Ragged unembedding implementation. This implementation will gather only the last token - of each sequence in the ragged inflight batch and calculate the logits only for those rows. - """ - - @staticmethod - def name(): - return 'ragged_unembed' - - @staticmethod - def supports_config(config: DSUnembedConfig): - - if DtypeEnum(config.dtype) not in [DtypeEnum.fp16, DtypeEnum.bf16, DtypeEnum.fp32]: - return False - - try: - _ = RaggedLogitsGather(config.model_dim, config.dtype) - except ValueError: - return False - - if config.norm_type == 'rms_norm': - try: - _ = CUDARMSNorm(config.model_dim, config.dtype) - except ValueError: - return False - elif config.norm_type == 'layer_norm': - try: - _ = CUDAFPLN(config.model_dim, config.dtype) - except ValueError: - return False - - return True - - def __init__(self, config: DSUnembedConfig, implementation_config: Dict[str, Any]) -> None: - super().__init__(config, implementation_config) - - self._logits_gather = RaggedLogitsGather(config.model_dim, self._config.dtype) - - if self._config.norm_type == 'layer_norm': - self._norm = CUDAFPLN(self._config.model_dim, self._config.dtype) - elif self._config.norm_type == 'rms_norm': - self._norm = CUDARMSNorm(self._config.model_dim, self._config.dtype) - else: - self._norm = None - - self._linear = BlasLibLinear(self._config.dtype) - # Here the activation kernel is being used to apply bias, hence the identity activation type! - self._act_fn = CUDABiasActivation(self._config.vocab_size, self._config.dtype, ActivationType.IDENTITY) - - self._intermediate = torch.empty((self._config.max_sequences, self._config.model_dim), - dtype=self._config.dtype, - device=get_accelerator().current_device()) - - self._output = torch.empty((self._config.max_sequences, self._config.vocab_size), - dtype=self._config.dtype, - device=get_accelerator().current_device()) - - @property - def output(self) -> torch.Tensor: - return self._output - - def forward(self, - hidden_states: torch.Tensor, - vocab_embedding: torch.Tensor, - ragged_metadata: RaggedBatchWrapper, - bias: Optional[torch.Tensor] = None, - gamma: Optional[torch.Tensor] = None, - beta: Optional[torch.Tensor] = None) -> torch.Tensor: - """ - Return final model logits. - - Args: - hidden_states (torch.Tensor): The hidden states from the model. This is the output of the - final layer of the model. - vocab_embedding (torch.Tensor): The vocab embedding table. - raged_metadata (RaggedBatchWrapper): The ragged batch metadata. - gamma (Optional[torch.Tensor]): The gamma tensor for normalization. - beta (Optional[torch.Tensor]): The beta tensor for normalization. - """ - - cut_down_hidden_states = empty_from(self._intermediate, - (ragged_metadata.current_sequences, self._config.model_dim)) - self._logits_gather(cut_down_hidden_states, hidden_states, ragged_metadata) - - if self._config.norm_type == 'rms_norm': - if gamma is None: - raise ValueError('RMS Normalization enabled but gamma not provided.') - self._norm(cut_down_hidden_states, cut_down_hidden_states, gamma) - elif self._config.norm_type == 'layer_norm': - if gamma is None or beta is None: - raise ValueError('Normalization enabled but gamma and/or beta not provided.') - self._norm(cut_down_hidden_states, cut_down_hidden_states, gamma, beta) - - output = empty_from(self._output, (ragged_metadata.current_sequences, self._config.vocab_size)) - self._linear(output, cut_down_hidden_states, vocab_embedding) - if bias is not None: - self._act_fn(output, bias) - - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__init__.py deleted file mode 100644 index 13b556789e4e43c4e816319cd48e95e9e51afb0c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .attention_base import DSSelfAttentionRegistry, DSSelfAttentionBase -from .embedding_base import DSEmbeddingRegistry, DSEmbeddingBase -from .linear_base import DSLinearRegistry, DSLinearBase -from .moe_base import DSMoERegistry, DSMoEBase -from .post_norm_base import DSPostNormRegistry, DSPostNormBase -from .pre_norm_base import DSPreNormRegistry, DSPreNormBase -from .unembed_base import DSUnembedRegistry, DSUnembedBase diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index dc65fa6493332df8bc9ec12530d030a76a4ea443..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/attention_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/attention_base.cpython-310.pyc deleted file mode 100644 index 95f113a50a1221880938244676c4ae900b6b6ea7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/attention_base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/embedding_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/embedding_base.cpython-310.pyc deleted file mode 100644 index 6fc941e816b04b306956382f87dcfdc4b6fb9311..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/embedding_base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/linear_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/linear_base.cpython-310.pyc deleted file mode 100644 index fd23cbcd290f15fa4b2b0fe067a463ff30dc4038..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/linear_base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/moe_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/moe_base.cpython-310.pyc deleted file mode 100644 index 75b023637a032e68c651143860f31ecb26207512..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/moe_base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/post_norm_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/post_norm_base.cpython-310.pyc deleted file mode 100644 index b6f24b581288dd541ec579f0bd652613a3788287..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/post_norm_base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/pre_norm_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/pre_norm_base.cpython-310.pyc deleted file mode 100644 index 2e31399989bd5f52c856dc37bc9e92e913129526..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/pre_norm_base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/unembed_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/unembed_base.cpython-310.pyc deleted file mode 100644 index 43593e0a58d9f1b153aef6bc0b9c282803b362a5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/__pycache__/unembed_base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/attention_base.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/attention_base.py deleted file mode 100644 index c67dc033f92ada844c1c5d27bde5ac0cf7f3f0ac..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/attention_base.py +++ /dev/null @@ -1,97 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any, Dict, Optional, Tuple, Type - -import torch - -from ...ragged import RaggedBatchWrapper -from deepspeed.runtime.config_utils import DeepSpeedConfigModel -from ..ds_module import DSModuleBase -from ..module_registry import DSModuleRegistryBase -from ..configs import DSSelfAttentionConfig - - -class DSSelfAttentionBase(DSModuleBase): - """ - Base mixin for all attention modules. The interface represented by this module - is broadly: - - output = attention(query_key_value, - Optional[kv_cache], - Optional[attention_mask], - Optional[attention_bias]) - """ - - @staticmethod - def config_class() -> Type[DeepSpeedConfigModel]: - return DSSelfAttentionConfig - - def __init__(self, config: DSSelfAttentionConfig, implementation_config: Dict[str, Any]) -> None: - super().__init__(config, implementation_config) - - @property - def kv_block_size(self) -> int: - """ - Return preferred granulatity for blocked KV-cache implementation. - """ - raise NotImplementedError() - - @property - def q_block_size(self) -> int: - """ - Property to calculate blocking granularity for the query dimension. - This has no impact on the KV-cache structure, but will affect the - number of attention atoms associated with a batch. - """ - raise NotImplementedError() - - def build_atoms(self, ragged_batch: RaggedBatchWrapper) -> None: - """ - Build the atoms for this module. This is not a strict requirement for the class, - so this method is a no-op by default rather than abstract. - """ - pass - - def forward(self, - q_k_v: torch.Tensor, - kv_cache: torch.Tensor, - batch: RaggedBatchWrapper, - attention_mask: Optional[torch.Tensor] = None, - attention_bias: Optional[torch.Tensor] = None, - inv_freqs: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Parameters: - q_k_v (torch.Tensor): Query, key, and value tensors. Expected shape is: - [ - batch, - seq_len, - 2 * self._config.n_heads_kv + self._config.n_heads_q, - self._config.head_size - ]. - kv_cache (Optional[torch.Tensor]): Key and value cache tensor. Expected shape is - [ - 2, - batch, - kv_cache_len, - self._config.n_heads_kv, - self._config.head_size - ]. If None, cache is disabled. The `kv_cache_len` dimension does not need to - be contiguous (it should expand stride by `max_out_tokens`). - batch (RaggedBatchWrapper): Ragged batch metadata. - attention_mask (Optional[torch.Tensor]): Attention mask tensor. If None, masking is - disabled. This will defer to the config in the case of conflicting information. - This means if the config class is implying causal attention, the mask will be ignored. - attention_bias (Optional[torch.Tensor]): Attention bias tensor. If None, bias is disabled. - """ - raise NotImplementedError() - - -class DSSelfAttentionRegistry(DSModuleRegistryBase): - registry: Dict = {} - - @staticmethod - def associated_class() -> Type[DSModuleBase]: - return DSSelfAttentionBase diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/embedding_base.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/embedding_base.py deleted file mode 100644 index 1ab7e5f0b7a2455ad3e60bb12a9d262900340e45..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/embedding_base.py +++ /dev/null @@ -1,85 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import abstractmethod -from typing import Any, Dict, Optional, Type - -import torch - -from deepspeed.runtime.config_utils import DeepSpeedConfigModel -from ...ragged import RaggedBatchWrapper -from ..ds_module import DSModuleBase -from ..module_registry import DSModuleRegistryBase -from ..configs import DSEmbeddingsConfig -from ...inference_parameter import InferenceParameter - - -class DSEmbeddingBase(DSModuleBase): - """ - Base mixin for embedding modules. The interface represented by this module is: - - hidden_out = embedding(input_ids) + - position_embedding(position_ids) + - token_type_embedding(token_type_ids) - with optional normalization. - """ - - @staticmethod - def config_class() -> Type[DeepSpeedConfigModel]: - return DSEmbeddingsConfig - - def __init__(self, config: DSEmbeddingsConfig, implementation_config: Dict[str, Any]) -> None: - super().__init__(config, implementation_config) - - def transform_param(self, embed_param: torch.Tensor) -> InferenceParameter: - """ - Perform any necessary transformations on an embedding parameter. This module assumes - that all embedding parameters would require the same set of transformations. - - Parameters: - embed_param (torch.Tensor): Embedding parameter. Shape is of [vocab_size, hidden_size] - """ - raise NotImplementedError() - - @property - @abstractmethod - def output(self) -> torch.Tensor: - """ - Pre-allocated output Tensor. This currently needs to be exposed for gather operations - on the output. - - TODO(cmikeh2): This is not ideal. We need a better abstraction for this, such as giving - access to the inference comm object to the DSModule. - """ - raise NotImplementedError() - - def forward(self, - ragged_batch: RaggedBatchWrapper, - word_embeddings: torch.Tensor, - position_embeddings: Optional[torch.Tensor] = None, - token_type_ids: Optional[torch.Tensor] = None, - token_type_embeddings: Optional[torch.Tensor] = None) -> InferenceParameter: - """ - Parameters: - ragged_batch (torch.Tensor): Ragged batch of token ids + associated metadata. - word_embeddings (torch.Tensor): Word embeddings. - position_embeddings (torch.Tensor): Position embeddings. If passed, IDs will be - inferred from the ragged batch itself. - token_type_ids (torch.Tensor): Token type ids. - token_type_embeddings (torch.Tensor): Token type embeddings. - - Returns: - torch.Tensor: Hidden states. This should be the sum of the relevant - encodings for the model. - """ - raise NotImplementedError() - - -class DSEmbeddingRegistry(DSModuleRegistryBase): - registry: Dict = {} - - @staticmethod - def associated_class() -> Type[DSModuleBase]: - return DSEmbeddingBase diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/linear_base.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/linear_base.py deleted file mode 100644 index fe6ccbcd934490c94a0b656f09b1f25aef70c163..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/linear_base.py +++ /dev/null @@ -1,72 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import abstractmethod -from typing import Any, Dict, Optional, Type - -import torch - -from deepspeed.runtime.config_utils import DeepSpeedConfigModel -from ..ds_module import DSModuleBase -from ..module_registry import DSModuleRegistryBase -from ..configs import DSLinearConfig -from ...inference_parameter import InferenceParameter - - -class DSLinearBase(DSModuleBase): - """ - Base mixin for all Linear modules. The interface represented by this module - is: - - hidden_out = activation(hidden_in * weight + bias) - - The format and dtype of the weight and bias tensors are not defined and implementations - may compress as necessary. Must support a bias. - """ - - @staticmethod - def config_class() -> Type[DeepSpeedConfigModel]: - return DSLinearConfig - - def __init__(self, config: DSLinearConfig, implementation_config: Dict[str, Any]) -> None: - super().__init__(config, implementation_config) - - @abstractmethod - def transform_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Perform any necessary transformations of the parameters of this module. - - Parameters: - param (torch.Tensor): Weight or bias tensor. - """ - ... - - def forward(self, hidden_states: torch.Tensor, w: torch.Tensor, b: Optional[torch.Tensor] = None) -> torch.Tensor: - """ - Parameters: - hidden_states (torch.Tensor): Hidden states tensor. Expected shape is either - [batch, seq_len, in_channels] or [batch, in_channels]. - - Returns: - torch.Tensor: Output tensor. Tensor should have same number of dimensions as - input tensor. - """ - raise NotImplementedError() - - @property - @abstractmethod - def output(self) -> torch.Tensor: - """ - Return the padded, pre-allocated output Tensor. - """ - ... - - -class DSLinearRegistry(DSModuleRegistryBase): - registry: Dict = {} - - @staticmethod - def associated_class() -> Type[DSModuleBase]: - return DSLinearBase diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/moe_base.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/moe_base.py deleted file mode 100644 index 78bdc0700f63b9aa532fc62b66b3e9c19261061c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/moe_base.py +++ /dev/null @@ -1,91 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import abstractmethod -from typing import Any, Dict, Optional, Type - -import torch - -from deepspeed.runtime.config_utils import DeepSpeedConfigModel -from ..ds_module import DSModuleBase -from ..module_registry import DSModuleRegistryBase -from ..configs import DSMoEConfig -from ...inference_parameter import InferenceParameter - - -class DSMoEBase(DSModuleBase): - """ - Base mixing for MoE modules. The interface represented by this module is: - - expert_assignments = gate(hidden_states) - intermediate = ragged_linear(hidden_states, expert_assignments) - output = ragged_linear(intermediate, expert_assignments) - """ - - @staticmethod - def config_class() -> Type[DeepSpeedConfigModel]: - return DSMoEConfig - - def __init__(self, config: DSMoEConfig, implementation_config: Dict[str, Any]) -> None: - super().__init__(config, implementation_config) - - @abstractmethod - def transform_gate_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Perform any necessary transformations of the gate parameter. - - Args: - param (torch.Tensor): gate_w (shape: [num_experts, model_dim]) - """ - ... - - @abstractmethod - def transform_moe_mlp_1_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Perform any necessary transformations of the parameter. The specific component - being transformed should be inferred from the shape of the parameter. - - Args: - param (torch.Tensor): One of either mlp_1_w, mlp_1_b - """ - ... - - @abstractmethod - def transform_moe_mlp_2_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Perform any necessary transformations of the parameter. The specified component being - transformed should be inferred from the shape of the parameter. This interface is - separate from transform_moe_1_param because the two components may have identical - shapes. - - Args: - param (torch.Tensor): One of either mlp_2_w or mlp_2_b - """ - ... - - def forward(self, - hidden_states: torch.Tensor, - gate_w: torch.Tensor, - mlp_1_w: torch.Tensor, - mlp_2_w: torch.Tensor, - mlp_1_b: Optional[torch.Tensor] = None, - mlp_2_b: Optional[torch.Tensor] = None) -> torch.Tensor: - raise NotImplementedError() - - @property - @abstractmethod - def output(self) -> torch.Tensor: - """ - Returns the pre-allocated, padded output Tensor. - """ - ... - - -class DSMoERegistry(DSModuleRegistryBase): - registry: Dict = {} - - @staticmethod - def associated_class() -> Type[DSModuleBase]: - return DSMoEBase diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/post_norm_base.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/post_norm_base.py deleted file mode 100644 index cc80e5c94bf73e7e2173906bd220fc39ff7df610..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/post_norm_base.py +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import abstractmethod -from typing import Any, Dict, Optional, Tuple, Type - -import torch - -from deepspeed.runtime.config_utils import DeepSpeedConfigModel -from ..ds_module import DSModuleBase -from ..configs.norm_config import DSNormConfig -from ..module_registry import DSModuleRegistryBase -from ...inference_parameter import InferenceParameter - - -class DSPostNormBase(DSModuleBase): - """ - Base MixIn for all Post-Normalization modules. The interface represented by this - module is: - - residual, hidden_out = norm(residual + hidden_in) - - If residual and hidden_out are the same data type, then they may alias each other. - Furthermore, residual should be updated in-place. - """ - - @staticmethod - def config_class() -> Type[DeepSpeedConfigModel]: - return DSNormConfig - - def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]) -> None: - super().__init__(config, implementation_config) - - @abstractmethod - def transform_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Transform a gamma/beta parameter. It is assumed that both transformations are - the same. - - Parameters: - param (torch.Tensor): Gamma or beta parameter. - """ - ... - - def forward(self, - residual: torch.Tensor, - hidden_states: torch.Tensor, - gamma: torch.Tensor, - beta: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Parameters: - residual (torch.Tensor): Residual tensor. - hidden_states (torch.Tensor): Hidden states tensor. - - Returns: - (torch.Tensor, torch.Tensor): Tuple of residual and hidden states. - Hidden states may alias with residual. - """ - raise NotImplementedError() - - -class DSPostNormRegistry(DSModuleRegistryBase): - registry: Dict = {} - - @staticmethod - def associated_class() -> Type[DSModuleBase]: - return DSPostNormBase diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/pre_norm_base.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/pre_norm_base.py deleted file mode 100644 index 84f51cff6947ecf8ee8563c289aec39324f37e7b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/pre_norm_base.py +++ /dev/null @@ -1,73 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import abstractmethod -from typing import Any, Dict, Optional, Tuple, Type - -import torch - -from deepspeed.runtime.config_utils import DeepSpeedConfigModel -from ..ds_module import DSModuleBase -from ..configs.norm_config import DSNormConfig -from ..module_registry import DSModuleRegistryBase -from ...inference_parameter import InferenceParameter - - -class DSPreNormBase(DSModuleBase): - """ - Base mixin for all Pre-Normalization modules. The interface represented by this module - is: - - if hidden_in is not None: - residual_out = residual + hidden_in - else: - residual_out = residual - - hidden_out = normalize(residual_out) - return residual_out, hidden_out - - Residual should be updated in-place. - """ - - @staticmethod - def config_class() -> Type[DeepSpeedConfigModel]: - return DSNormConfig - - def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]): - super().__init__(config, implementation_config) - - @abstractmethod - def transform_param(self, param: torch.Tensor) -> InferenceParameter: - """ - Transform a gamma/beta parameter. It is assumed that both transformations are - the same. - - Parameters: - param (torch.Tensor): Gamma or beta parameter. - """ - ... - - def forward(self, - residual: torch.Tensor, - hidden_states: Optional[torch.Tensor], - gamma: torch.Tensor, - beta: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Parameters: - residual (torch.Tensor): Residual tensor. - hidden_states (torch.Tensor): Hidden states tensor. - - Returns: - (torch.Tensor, torch.Tensor): Tuple of residual and hidden states. - """ - raise NotImplementedError() - - -class DSPreNormRegistry(DSModuleRegistryBase): - registry: Dict = {} - - @staticmethod - def associated_class() -> Type[DSModuleBase]: - return DSPreNormBase diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/unembed_base.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/unembed_base.py deleted file mode 100644 index 9eca6fcde7682903ef445f4703e009b5c1b9e641..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/interfaces/unembed_base.py +++ /dev/null @@ -1,61 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Any, Dict, Optional, Type - -import torch - -from deepspeed.runtime.config_utils import DeepSpeedConfigModel -from ...ragged import RaggedBatchWrapper -from ..ds_module import DSModuleBase -from ..module_registry import DSModuleRegistryBase -from ..configs import DSUnembedConfig - - -class DSUnembedBase(DSModuleBase): - """ - Base mixin for unmebedding modules. The interface represented by this module is: - - if config.do_normalization - hidden = layer_norm(hidden) - logits = hidden @ projection - """ - - @staticmethod - def config_class() -> Type[DeepSpeedConfigModel]: - return DSUnembedConfig - - def __init__(self, config: DSUnembedConfig, implementation_config: Dict[str, Any]) -> None: - super().__init__(config, implementation_config) - - def forward(self, - hidden_states: torch.Tensor, - vocab_embedding: torch.Tensor, - ragged_metadata: RaggedBatchWrapper, - gamma: Optional[torch.Tensor] = None, - beta: Optional[torch.Tensor] = None) -> torch.Tensor: - """ - Forward interface. Gamma and beta are optional parameters passed depending on - `self.config.do_normalization`. - - Args: - hidden_states (torch.Tensor): Hidden states of shape [tokens, model_dim] - vocab_embedding (torch.Tensor): Embedding matrix of shape [vocab_size, model_dim] - ragged_metadata (RaggedBatchWrapper): Metadata for the ragged batch. - gamma (Optional[torch.Tensor]): Gamma parameter for layer norm. - beta (Optional[torch.Tensor]): Beta parameter for layer norm. - - Returns: - torch.Tensor: Unembedded hidden states of shape [n_seqs, model_dim] - """ - raise NotImplementedError() - - -class DSUnembedRegistry(DSModuleRegistryBase): - registry: Dict = {} - - @staticmethod - def associated_class() -> Type[DSModuleBase]: - return DSUnembedBase diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/module_registry.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/module_registry.py deleted file mode 100644 index e04b8d734518b2322fe3f07a38c5d01222e4ecbe..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/modules/module_registry.py +++ /dev/null @@ -1,58 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import ABC, abstractstaticmethod -from typing import Any, Dict, Type - -from deepspeed.runtime.config_utils import DeepSpeedConfigModel -from .ds_module import DSModuleBase - - -class ConfigBundle(DeepSpeedConfigModel): - """ - A config bundle is a collection of configs that are used to instantiate a model implementation. - """ - name: str - config: DeepSpeedConfigModel - implementation_config: Dict[str, Any] = {} - - -class DSModuleRegistryBase(ABC): - """ - Class holding logic for tracking the DSModule implementations of a given interface. - """ - - @classmethod - def instantiate_config(cls, config_bundle: ConfigBundle) -> DSModuleBase: - """ - Given a DSModule key, attempt to instantiate - """ - if config_bundle.name not in cls.registry: - raise KeyError(f"Unknown DSModule: {config_bundle.name}, cls.registry={cls.registry}") - - target_implementation = cls.registry[config_bundle.name] - if not target_implementation.supports_config(config_bundle.config): - raise ValueError(f"Config {config_bundle.config} is not supported by {target_implementation}") - - return cls.registry[config_bundle.name](config_bundle.config, config_bundle.implementation_config) - - @abstractstaticmethod - def associated_class() -> Type[DSModuleBase]: - """ - Return the class associated with this registry. - """ - raise NotImplementedError("Must associated a DSModule class with its registry.") - - @classmethod - def register_module(cls, child_class: DSModuleBase) -> None: - """ - Register a module with this registry. - """ - if not issubclass(child_class, cls.associated_class()): - raise TypeError( - f"Can only register subclasses of {cls.associated_class()}, {child_class} does not inherit from {cls.associated_class()}" - ) - cls.registry[child_class.name()] = child_class - return child_class diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__init__.py deleted file mode 100644 index 3af09cff4be57716e3e4295b88163db6b90b608a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .kv_cache import split_kv -from .manager_configs import ( - AllocationMode, - DSStateManagerConfig, - KVCacheConfig, - MemoryConfig, -) -from .ragged_manager import DSStateManager -from .ragged_wrapper import RaggedBatchWrapper -from .sequence_descriptor import DSSequenceDescriptor, PlaceholderSequenceDescriptor diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 8d894340ae064cde5bb72a2800adb4b733edb87b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/blocked_allocator.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/blocked_allocator.cpython-310.pyc deleted file mode 100644 index 26ce4312eb3cb941223c25c976a865ef18ccb6e5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/blocked_allocator.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/kv_cache.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/kv_cache.cpython-310.pyc deleted file mode 100644 index 5e4dd3e2c655be32f6f2b8cf33a01cbaf99fd57e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/kv_cache.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/manager_configs.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/manager_configs.cpython-310.pyc deleted file mode 100644 index 5a3caa98791d87a592f2685cb8400bb0d07a9c1b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/manager_configs.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/ragged_manager.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/ragged_manager.cpython-310.pyc deleted file mode 100644 index 5078f6630b80deea564b126216bb9e32b4674dd6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/ragged_manager.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/ragged_wrapper.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/ragged_wrapper.cpython-310.pyc deleted file mode 100644 index 1c82fb5e9b9bf1efc8a16ff0f5b908005bef5224..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/ragged_wrapper.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/sequence_descriptor.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/sequence_descriptor.cpython-310.pyc deleted file mode 100644 index 558d1aecddfb02843039e9578344f260145011b9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/__pycache__/sequence_descriptor.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/blocked_allocator.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/blocked_allocator.py deleted file mode 100644 index 7884d8cccb47a1e9a4d94759a12007fabfea90a4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/blocked_allocator.py +++ /dev/null @@ -1,105 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Iterable, Union - -import torch - - -class BlockedAllocator: - """ - Allocator class for managing which blocks are free/used in the - blocked KV-cache. This is a simple allocator that uses a linked list - to keep track of which blocks are free/used. The cost of allocation/deallocation - is O(blocks), where blocks is the number of blocks to allocate/deallocate. - - TODO(cmikeh2): Evaluate performance of this allocator and migrate - to C++ if necessary. - """ - # Number of blocks in the KV-cache(s). - _num_blocks: int - - # Array of blocks, where each element is the next block in the linked list. - _blocks: torch.Tensor - - # Index of the head of the linked list. - _head: int - - # Number of free blocks in the KV-cache. - _free_blocks: int - - def __init__(self, num_blocks: int) -> None: - """ - Initialize an allocator with `num_blocks` blocks. This requires at least - `num_blocks` * 4 bytes of host memory. - - Parameters: - num_blocks (int): The number of blocks to allocate. - """ - - if num_blocks < 1: - raise ValueError(f'Blocked KV-cache must have at least 1 block, provided {num_blocks}') - - self._num_blocks = num_blocks - self._blocks = torch.arange(1, num_blocks + 1, dtype=torch.int32, device='cpu', pin_memory=True) - self._head = 0 - self._free_blocks = num_blocks - - def allocate(self, num_blocks: int) -> torch.Tensor: - """ - Allocate a list of blocks from the associated KV-caches. This will - return `num_blocks` blocks from the KV-cache if they are available, - or raise an exception if there are not enough free blocks. - - Parameters: - num_blocks (int): The number of blocks to allocate. - - Returns: - List[int]: The list of blocks allocated. - """ - if num_blocks > self._free_blocks: - raise ValueError(f'Not enough free blocks in the KV-cache to allocate {num_blocks} blocks') - - allocated_blocks = torch.zeros(num_blocks, dtype=torch.int32) - for i in range(num_blocks): - allocated_blocks[i] = self._head - self._head = self._blocks[self._head].item() - self._blocks[allocated_blocks[i]] = -1 # Mark as used - self._free_blocks -= 1 - - return allocated_blocks - - def free(self, blocks: Union[Iterable[int], int]) -> None: - """ - Return a list of blocks to the free pool. If a single invalid block is provided (i.e., - one that is out of range of the allocator or is already free), then an exception is raised - and no blocks are freed. - - Parameters: - blocks (Union[Iterable[int], int]): The list of blocks to free. If only one block - is to be freed, this can be alone as an integer. - """ - if isinstance(blocks, int): - blocks = [blocks] - - for block in blocks: - # Parse all blocks for validity before mutating the list. - if block < 0 or block >= self._num_blocks: - raise ValueError(f'Invalid block {block} provided to free') - - if self._blocks[block] != -1: - raise ValueError(f'Block {block} is already free') - - for block in blocks: - self._blocks[block] = self._head - self._head = block - self._free_blocks += 1 - - @property - def free_blocks(self) -> int: - """ - Return the number of free blocks in the KV-cache. - """ - return self._free_blocks diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/csrc/fast_host_buffer.cu b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/csrc/fast_host_buffer.cu deleted file mode 100644 index 31347636b50c40fc74f0c477e848a9cddbab0c92..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/csrc/fast_host_buffer.cu +++ /dev/null @@ -1,18 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "ds_kernel_utils.h" -#include "fast_host_buffer.h" - -void* get_cuda_fast_buffer(int64_t size) -{ - void* buffer_ptr; - // Host allocation flags that should minimize the host -> accelerator copy latency - unsigned int alloc_flags = - cudaHostAllocPortable | cudaHostAllocMapped | cudaHostAllocWriteCombined; - - cudaHostAlloc(&buffer_ptr, size, alloc_flags); - return buffer_ptr; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/csrc/ragged_ops.cpp b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/csrc/ragged_ops.cpp deleted file mode 100644 index ce115f993c3c5676c77055773569b50ac68f0f7c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/csrc/ragged_ops.cpp +++ /dev/null @@ -1,76 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include - -#include "fast_host_buffer.h" - -/* -Similar to doing an empty_like to replicate a Tensor on the host, but will -attempt to optimize for faster host -> accelerator copies. Since this is on the critical -path for the forward pass, this should directly improve performance. -Allocates the shadow buffers for the input_ids, batch, seq and kv_ids tensors. - -Arguments: - device_mirror: A tensor on the accelerator that should be mirrored by the host. - -Returns: - A tensor on the host of the same size and datatype optimized for fast host -> accelerator -copies. -*/ -torch::Tensor allocate_fast_host_buffer(torch::Tensor device_mirror) -{ -#ifdef __HIP_PLATFORM_AMD__ - auto options = - torch::TensorOptions().device(torch::kCPU).pinned_memory(true).dtype(device_mirror.dtype()); - auto buffer = torch::empty(device_mirror.sizes(), options); -#else - - void* buffer_ptr = get_cuda_fast_buffer(device_mirror.numel() * device_mirror.element_size()); - - auto options = torch::TensorOptions().device(torch::kCPU).dtype(device_mirror.dtype()); - auto buffer = torch::from_blob(buffer_ptr, device_mirror.sizes(), options); -#endif - return buffer; -} - -torch::Tensor allocate_view_on(torch::Tensor& tensor, torch::Tensor& buffer, int64_t offset) -{ - int8_t* data = reinterpret_cast(buffer.data_ptr()); - - auto options = tensor.options().device(buffer.device()); - - return at::from_blob(data + offset, tensor.sizes(), tensor.strides(), options); -} - -torch::Tensor allocate_view_like(py::tuple shape, - py::tuple strides, - torch::Tensor& dummy_tensor, - torch::Tensor& buffer, - int64_t offset) -{ - int8_t* data = reinterpret_cast(buffer.data_ptr()); - - auto options = torch::TensorOptions().device(buffer.device()).dtype(dummy_tensor.dtype()); - - return at::from_blob(data + offset, - shape.cast>(), - strides.cast>(), - options); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("allocate_fast_host_buffer", - &allocate_fast_host_buffer, - "Allocate a host mirror of an accelerator Tensor."); - m.def("allocate_view_on", - &allocate_view_on, - "Allocate a view on a Tensor on the same device as the input Tensor."); - m.def("allocate_view_like", - &allocate_view_like, - "Allocate a view on a Tensor on the same device as the input Tensor."); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/includes/fast_host_buffer.h b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/includes/fast_host_buffer.h deleted file mode 100644 index 81f24ed8fdaadfcf1b712e7c025e6f7a2e4a95a1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/includes/fast_host_buffer.h +++ /dev/null @@ -1,14 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include "ds_kernel_utils.h" - -/* -Wrapper around cudaHostAlloc with some specific flags. Returns a pointer to the -memory region of `size` bytes. -*/ -void* get_cuda_fast_buffer(int64_t size); diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/kv_cache.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/kv_cache.py deleted file mode 100644 index ceba3190b93cdc378bc4581def235ffa99210bfd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/kv_cache.py +++ /dev/null @@ -1,208 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import operator -from functools import reduce -from typing import Any, Iterable, Optional, Tuple - -import torch - -import deepspeed.comm as dist -from deepspeed.comm.reduce_op import ReduceOp - -from deepspeed.accelerator import get_accelerator -from ..inference_utils import elem_size -from ..logging import inference_logger -from .blocked_allocator import BlockedAllocator -from .manager_configs import AllocationMode, KVCacheConfig, MemoryConfig - - -def split_kv(kv_cache: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Split a KV cache instance into its key and value components. - - Parameters: - kv_cache (torch.Tensor): The KV-cache to split. This should be a 5D tensor with the - following shape: [num_blocks, block_size, 2, num_heads, head_size] - - Returns: - Tuple[torch.Tensor, torch.Tensor]: The key and value components of the KV-cache. Both - tensors will have the shape [num_blocks, block_size, num_heads, head_size]. - """ - if kv_cache.ndim != 5: - raise ValueError(f"KV-cache must have 5 dimensions, got {kv_cache.ndim}.") - - return kv_cache[:, :, 0, :, :], kv_cache[:, :, 1, :, :] - - -class BlockedKVCache: - - _caches: Tuple[torch.Tensor, ...] - """ - Backing storage for all KV caches. This is a 6D tensor with the following shape: - (num_caches, num_blocks, block_size, 2, num_heads, head_size) - """ - - _allocators: Tuple[BlockedAllocator, ...] - """ - Block allocator for tracking cache usage. This manages the GPU cache. - """ - - _configs: Tuple[KVCacheConfig, ...] - """ - Configuration of the KV cache(s). See ``KVCacheConfig`` for more details. This enables the support - for different types/shapes of KV-caches (i.e. the alternating local and global attention in - GPT-Neo). - """ - - def __init__(self, - configs: Tuple[KVCacheConfig, ...], - memory_config: MemoryConfig, - mp_group: Optional[Any] = None, - offload: bool = False) -> None: - """ - Create a container that will maintain the storage and allocations for a set of - blocked KV-caches. - - Parameters: - config (KVCacheConfig): The configuration of the KV-cache. - slack (int): The amount of slack space to reserve in GPU memory for the cache. - enable_offload (bool): Whether to enable offloading of the cache to the host. - blocks (int): The number of blocks to pre-allocate for the cache. If this is set, - slack will be ignored. - """ - self._configs = configs - self._memory_config = memory_config - self._enable_offload = offload - - if self._enable_offload: - raise NotImplementedError("Offloading of KV-caches is not yet supported.") - - if AllocationMode(self._memory_config.mode) is AllocationMode.RESERVE: - # TODO(cmikeh2): Change the weighting based on the type of the KV-cache - - total_per_block_footprint = 0 - for config in self._configs: - per_block_footprint = reduce(operator.mul, config.cache_shape, config.block_size) - per_block_footprint *= 2 # for key and value - total_per_block_footprint += per_block_footprint * elem_size(config.cache_dtype) - - # Perform a dummy nccl call before calculating available memory, on some systems (H100) we've observed higher memory allocations from NCCL - if dist.get_world_size(group=mp_group) > 1: - dummy_tensor = torch.tensor(0, dtype=torch.int32, device=get_accelerator().current_device()) - dist.all_reduce(dummy_tensor, op=ReduceOp.MIN, group=mp_group) - - get_accelerator().empty_cache() - available_kv_memory = get_accelerator().available_memory() - self._memory_config.size - total_memory = get_accelerator().total_memory() - - inference_logger().debug( - f"Memory usage before KV-cache allocation: total_memory={total_memory}, available_kv_memory={available_kv_memory}, total_per_block_footprint={total_per_block_footprint}" - ) - - if available_kv_memory < total_per_block_footprint: - raise ValueError( - f"Insufficient memory to allocate KV-caches. Required: {total_per_block_footprint}, Available: {available_kv_memory}" - ) - - num_blocks = available_kv_memory // total_per_block_footprint - - # In a multi-process setting, we need to ensure that all processes have the same - # KV cache capacity to ensure scheduling guarantees are equivalent on all ranks. - if dist.get_world_size(group=mp_group) > 1: - reduce_tensor = torch.tensor(num_blocks, dtype=torch.int32, device=get_accelerator().current_device()) - dist.all_reduce(reduce_tensor, op=ReduceOp.MIN, group=mp_group) - num_blocks = reduce_tensor.item() - - # This is ugly but don't want the fragmentation of the 8 byte Tensor maybe - # hanging around. - del reduce_tensor - get_accelerator().empty_cache() - else: # AllocationMode.ALLOCATE - num_blocks = self._memory_config.size - - caches = [] - allocators = [] - - for cache_group_id, config in enumerate(self._configs): - num_caches = config.cache_shape[0] - num_heads = config.cache_shape[1] - head_size = config.cache_shape[2] - - alloc_shape = (num_caches, num_blocks, config.block_size, 2, num_heads, head_size) - inference_logger().info( - f"Allocating KV-cache {cache_group_id} with shape: {alloc_shape} consisting of {num_blocks} blocks.") - caches.append(torch.empty(alloc_shape, dtype=config.cache_dtype, - device=get_accelerator().current_device())) - allocators.append(BlockedAllocator(num_blocks)) - - self._caches = tuple(caches) - self._allocators = tuple(allocators) - - def reserve(self, num_blocks: int, cache_group: int = 0) -> torch.Tensor: - """ - Reserve a number of blocks from the cache. This will return a 1D tensor of - block_ids that have been marked as reserved. - - Parameters: - num_blocks (int): The number of blocks to reserve. - cache_group (int): The cache group to reserve from. Default is 0. - """ - return self._allocators[cache_group].allocate(num_blocks) - - def free(self, blocks: Iterable[int], cache_group: int = 0) -> None: - """ - Free a set of blocks from the cache. This will mark the blocks as free in the - allocator. - - Parameters: - blocks (Iterable[int]): The blocks to free. - cache_group (int): The cache group to free from. Default is 0. - """ - self._allocators[cache_group].free(blocks) - - def offload(self, blocks: Iterable[int], cache_group: int = 0) -> torch.Tensor: - """ - Offload KV-cache blocks from accelerator memory to the host. - - Parameters: - blocks (Iterable[int]): The blocks to offload. - cache_group (int): The cache group to offload from. Default is 0. - """ - raise NotImplementedError("Offloading is not yet supported.") - - def restore(self, blocks: Iterable[int], cache_group: int = 0) -> torch.Tensor: - """ - Restore KV-cache blocks from the host to accelerator memory. - - Parameters: - blocks (Iterable[int]): The blocks to restore. - cache_group (int): The cache group to restore to. Default is 0. - """ - raise NotImplementedError("Offloading is not yet supported.") - - def get_cache(self, cache_id: int, cache_group: int = 0) -> torch.Tensor: - """ - Get the tensor associated with the given cache ID. - - Parameters: - cache_id (int): The ID of the cache tensor to get. - cache_group (int): The cache group to get from. Default is 0. - """ - return self._caches[cache_group][cache_id] - - @property - def free_blocks(self) -> torch.Tensor: - """ - Return the number of free blocks in each cache - """ - return [allocator.free_blocks for allocator in self._allocators] - - @property - def num_caches(self) -> int: - """ - Return the number of caches - """ - return len(self._caches) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/manager_configs.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/manager_configs.py deleted file mode 100644 index 17283b8bc0c4707ed95ad5e91414d160b59374ea..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/manager_configs.py +++ /dev/null @@ -1,181 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from enum import Enum -from typing import Tuple - -from pydantic import PositiveInt, model_validator - -from deepspeed.runtime.config_utils import DeepSpeedConfigModel -from ..inference_utils import DtypeEnum - - -class KVCacheType(Enum): - - DENSE = "dense" - """ - Dense KV-cache. This is the default type. - """ - - LOCAL = "local" - """ - KV-cache that attends to only a local (trailing) window of tokens. - """ - - -class KVCacheConfig(DeepSpeedConfigModel): - - type: KVCacheType = KVCacheType.DENSE - """ - Type of KV-cache to use. This may inform the allocator of the expected access/retention pattern - to enable more efficient memory management. - """ - - block_size: int = 128 - """ - Number of tokens that may be contained in each cache block. - """ - - num_allocation_groups: PositiveInt = 1 - """ - Allocation groups are assumed to be able to use the same allocation block size because - the allocation granularity is the same but the number of blocks required in each group - may differ. - - As a concrete example, consider a model with alternating layers of local and global - attention (such as GPTNeo). The local attention layers do not require the same number - of cache blocks as the global layer. However, a static partitioning scheme is sub-optimal since the ratio of local to global KV-cache blocks is not constant across - the range of sequence lengths that may be encountered. - - NOTE: In theory, this functionality could be used to do per-head and per-layer - KV-cache allocation, but it is likely the allocator will struggle with managing that - many blocks. - - NOTE: This will need to be primarily understood and handled by the model implementation - itself, rather than the KV cache manager. However, I'd like to make this explicit. - """ - - cache_shape: Tuple[PositiveInt, PositiveInt, PositiveInt] - """ - The shape of the cache per token. The first dimension is the number of individual - caches, the second is the number of heads, and the third is the head size. The number - of caches argument here is per allocation group. - """ - - cache_dtype: DtypeEnum = DtypeEnum.fp16 - """ - Data type of the KV-cache. - """ - - max_blocks_per_allocation_group: PositiveInt = 64 - """ - Maximum number of blocks that can be associated with an allocation group. - """ - - -""" -The config above is a little confusing so let's use a couple of concrete examples of -usage: - -Model 1: Llama-13B with a block size of 256 - -Llama is uniform attention so we have a single allocation group. The cache shape is -(40 layers, 40 heads, 128 head size) - -```python -llama_kv_config = KVCacheConfig(block_size=256, - num_allocation_groups=1, - cache_shape=(40, 40, 128)) -``` - -Model 2: GPTNeo-2.7B with a block size of 128 - -GPTNeo has alternating local and global attention layers. We have two allocation groups. -There are 16 layers of each type with 20 heads apiece at 128 head size. - -```python -gptneo_kv_config = KVCacheConfig(num_allocation_groups=2, cache_shape=(16, 20, 128)) -``` -""" - - -class AllocationMode(Enum): - """ - Helper class to describe memory allocation strategies for the KV-cache. - """ - - RESERVE = "reserve" - """ - Reserve a small amount of memory for non-KV cache allocations. - """ - - ALLOCATE = "allocate" - """ - Allocate an explicit number of KV blocks. - """ - - -class MemoryConfig(DeepSpeedConfigModel): - - mode: AllocationMode = AllocationMode.RESERVE - - size: PositiveInt = 1_000_000_000 - """ - Parameter for each of the modes. - - If mode is RESERVE, this is the amount of memory in bytes to reserve after allocating the - KV-cache. If in a tensor-parallel regime, this amount is guaranteed to be reserved on - all devices. - - If mode is ALLOCATE, this is the number of blocks to allocate for the KV-cache. This may - require tuning for model/GPU setups. - """ - - -class DSStateManagerConfig(DeepSpeedConfigModel): - - max_tracked_sequences: PositiveInt = 2048 - """ - How many sequences this engine will track simultaneously. This limit should be greater - than the ``max_ragged_sequence_count``. - """ - - max_ragged_batch_size: PositiveInt = 768 - """ - The maximum number of tokens that can be contained in a single ragged batch. Passing - a larger value than this will raise an exception that must be handled by the runtime. - """ - - max_ragged_sequence_count: PositiveInt = 512 - """ - The maximum number of sequences that can compose a batch. This limitation is only - relevant under CUDA graphing scenarios currently, where the maximum number of blocks - is largely bound by the total number of sequences in the ragged batch. This number cannot - be larger than ``max_tracked_sequences`` or ``max_ragged_batch_size``. - """ - - max_context: PositiveInt = 8192 - """ - The maximum number of tokens (inclusive of generation) that can be contained in a single - sequence. Currently used to bound the size of the KV cache metadata. - """ - - memory_config: MemoryConfig = MemoryConfig() - """ - Directive for how to manage the creation of the KV-cache. See MemoryConfig for more - details. - """ - - offload: bool = False - """ - Enable tracking for offloading KV-cache to host memory. Currently unsupported. - """ - - @model_validator(mode="after") - def max_ragged_sequence_count_validator(self): - # If the attributes below failed their validation they won't appear in the values dict. - assert self.max_ragged_sequence_count <= self.max_tracked_sequences, "max_ragged_sequence_count must be less than max_tracked_sequences" - assert self.max_ragged_sequence_count <= self.max_ragged_batch_size, "max_ragged_sequence_count must be less than max_ragged_batch_size" - return self diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/ragged_manager.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/ragged_manager.py deleted file mode 100644 index ecc3c52a5834321255d3e475fc2fcf8eb391dd82..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/ragged_manager.py +++ /dev/null @@ -1,206 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from typing import Any, Dict, Optional, Tuple - -from deepspeed.accelerator import get_accelerator -from deepspeed.ops.op_builder import RaggedUtilsBuilder -from deepspeed.utils.logging import logger - -from .blocked_allocator import BlockedAllocator -from .kv_cache import BlockedKVCache -from .manager_configs import DSStateManagerConfig, KVCacheConfig -from .sequence_descriptor import DSSequenceDescriptor - - -class DSStateManager: - """ - Base abstract class for managing blocked KV caches. Will probably have a single - implementation for now. - """ - - _config: DSStateManagerConfig - """ - Config for state management. See DSStateManagerConfig for more details. The arguments here - should come from the engine config. - """ - - _kv_configs: Tuple[KVCacheConfig] - """ - Config for the KV cache. See KVCacheConfig for more details. These arguments should derive - from the model implementation. - """ - - _kv_cache: BlockedKVCache - """ - Persistent KV cache store. - """ - - # Container for tracking all sequences in the system. - _seqs: Dict[int, DSSequenceDescriptor] - """ - Container for tracking all sequences in the system. - - TODO(cmikeh2): Evaluate if this has any performance implications. - """ - - # Allocator for tracking sequences. - _tracking_allocator: BlockedAllocator - _all_block_ids: Tuple[torch.Tensor, ...] - _all_block_ids_shadow: Tuple[torch.Tensor, ...] - - def __init__(self, - config: DSStateManagerConfig, - kv_configs: Tuple[KVCacheConfig, ...], - base_mp_group: Optional[Any] = None) -> None: - """ - The key - - Parameters: - block_size (int): The number of tokens to allocate in each block. - """ - self._config = config - self._kv_configs = kv_configs - - # Load our helpers for host allocation. - self._ragged_utils = RaggedUtilsBuilder().load() - - # Initialize the allocator for tracking sequences (so this doesn't need to be ad-hoc). - self._tracking_allocator = BlockedAllocator(self._config.max_tracked_sequences) - - all_block_ids = [] - all_block_ids_shadow = [] - - for cache_config in self._kv_configs: - # Storage to back tracking the KV cache allocation. - ids_shape = ( - self._config.max_tracked_sequences, - cache_config.num_allocation_groups, - cache_config.max_blocks_per_allocation_group, - ) - - all_block_ids.append(torch.zeros(ids_shape, dtype=torch.int32, device=get_accelerator().current_device())) - all_block_ids_shadow.append(self._ragged_utils.allocate_fast_host_buffer(all_block_ids[-1])) - - self._all_block_ids = tuple(all_block_ids) - self._all_block_ids_shadow = tuple(all_block_ids_shadow) - - # Initialize the sequence container. - self._seqs = {} - - # Finally initialize the KV cache. - self._kv_cache = BlockedKVCache(self._kv_configs, - self._config.memory_config, - mp_group=base_mp_group, - offload=self._config.offload) - - def get_cache(self, cache_id: int, cache_group: int = 0) -> torch.Tensor: - """ - Return the Tensor associated with the given cache id in the specified cache group. - - Arguments: - cache_group (str): The KV cache group. - cache_id (int): The cache id within that group. - """ - return self._kv_cache.get_cache(cache_id, cache_group=cache_group) - - def flush_sequence(self, uid: int) -> None: - """ - Free all resources associated with the given sequence id. - """ - if uid not in self._seqs: - logger.warning(f"Attempting to flush sequence {uid} which does not exist.") - return - - seq = self._seqs[uid] - for i in range(self.n_kv_cache_groups): - self._kv_cache.free(seq.all_block_ids(cache_group=i), cache_group=i) - - self._tracking_allocator.free(seq.tracking_id) - del self._seqs[uid] - - def get_sequence(self, uid: int) -> Optional[DSSequenceDescriptor]: - """ - Get the sequence descriptor for the given sequence id. If the sequence does not exist, - then None is returned. - """ - return self._seqs.get(uid, None) - - def get_or_create_sequence(self, uid: int) -> DSSequenceDescriptor: - """ - Get the existing sequence descriptor for a given uid or initialize one if - it does not exist. NOTE: This will always return a valid sequence descriptor - if one may be allocated and should not be used from APIs that are attempting - to test the schedulability of a hypothetical batch. - """ - seq = self.get_sequence(uid) - if seq is not None: - return seq - else: - return self._create_sequence(uid) - - def _create_sequence(self, uid: int) -> DSSequenceDescriptor: - """ - Create a new sequence descriptor for the given sequence id. - """ - if uid in self._seqs: - raise ValueError(f"Sequence {uid} already exists.") - - try: - tracking_slot = self._tracking_allocator.allocate(1).item() - except ValueError: - raise RuntimeError( - f"Unable to create tracking slot for sequence {uid} since the metadata buffers are full.") - - seq_block_ids = tuple(all_block_ids[tracking_slot] for all_block_ids in self._all_block_ids) - seq_block_ids_shadow = tuple(all_block_ids_shadow[tracking_slot] - for all_block_ids_shadow in self._all_block_ids_shadow) - - self._seqs[uid] = DSSequenceDescriptor(tracking_slot, - seq_block_ids, - seq_block_ids_shadow, - max_context=self._config.max_context) - # TODO(cmikeh2): Debug call here might be unnecessary and is potentially on critical path. - logger.debug(f"Created sequence {uid} with tracking slot {tracking_slot}.") - return self._seqs[uid] - - @property - def tracked_sequences(self) -> Dict[int, DSSequenceDescriptor]: - """ - Return the tracked sequences. - """ - return self._seqs - - @property - def n_tracked_sequences(self) -> int: - """ - Return the number of sequences currently tracked. - """ - return len(self._seqs) - - @property - def kv_block_size(self) -> int: - """ - Return the block size of the KV cache. - """ - return self._kv_config.block_size - - @property - def n_kv_cache_groups(self) -> int: - """ - Return the number of KV caches. - """ - return self._kv_cache.num_caches - - @property - def free_blocks(self) -> torch.Tensor: - """ - Return the number of free blocks in the KV cache. - """ - return self._kv_cache.free_blocks - - def allocate_blocks(self, n_blocks: int, cache_group: int = 0) -> torch.Tensor: - return self._kv_cache.reserve(n_blocks, cache_group=cache_group) diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/ragged_wrapper.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/ragged_wrapper.py deleted file mode 100644 index 056ecfa2ac4094cf9380db533a1f031332815641..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/ragged_wrapper.py +++ /dev/null @@ -1,292 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Optional - -import torch - -from deepspeed.accelerator import get_accelerator -from deepspeed.ops.op_builder import RaggedUtilsBuilder - -from .sequence_descriptor import DSSequenceDescriptor -from .manager_configs import DSStateManagerConfig - - -def to_padded(original_size: int) -> int: - """ - Pad to a backend friendly granularity. - """ - - def _pad_to_mul_of_pow2(val: int, pow_2_val: int) -> int: - return val + (pow_2_val - 1) & ~(pow_2_val - 1) - - # TODO(cmikeh2): Tune this approach. This is mainly a placeholder right now. - granularity = 64 if original_size <= 512 else 128 - - return _pad_to_mul_of_pow2(original_size, granularity) - - -class RaggedBatchWrapper: - """ - Container for all the auxiliary Tensors used in the management of a ragged batch. - - For each Tensor, we maintain a shadow Tensor on the host. This Tensor is what is - directly populated when constructing the ragged batch. The shadow Tensors, when possible, - should be allocated so as to support fast host-to-accelerator copies. - """ - - # Tensors to populate the ragged batch into. - _input_ids_shadow: torch.Tensor - _input_ids: torch.Tensor - """ - Forward pass input buffer. - """ - - _batch_metadata_storage: torch.Tensor - _batch_metadata_storage_shadow: torch.Tensor - """ - Holds the number of inflight sequences and tokens for the ragged batch. - """ - - _token_to_seq_storage: torch.Tensor - _token_to_seq_storage_shadow: torch.Tensor - """ - Linear mapping for each of the tokens. Let's say we have 8 tokens in the batch, - with the sequence breakdown being [4, 1, 3]. Then, the mapping would be: - [0, 0, 0, 0, 1, 2, 2, 2] - """ - - _inflight_seq_descriptors: torch.Tensor - _inflight_seq_descriptors_shadow: torch.Tensor - """ - For each sequence in the batch, we store the start token in the batch, the number of tokens - the number of tokens in the history of this sequence, and an unused 4th reserved for alignment. - For the above example this would give: - [[0, 4, H0, X], [4, 1, H1, X], [5, 3, H2, X]] - """ - - # Holds the block ids for each sequence in the ragged batch. - _kv_ptrs: torch.Tensor - _kv_ptrs_shadow: torch.Tensor - """ - List of ptrs pointing to the GPU buffer that holds the KV-block ids for each sequence. - If there are multiple allocation groups associated with each of the sequences, then - then accessing the Nth cache will require accessing the Nth block id - """ - - def __init__(self, config: DSStateManagerConfig) -> None: - """ - Convenience wrapper around the data structures used to represent a ragged - batch for inference. Only a single `RaggedBatchWrapper` should be used per - ragged inference engine. - - The underlying data structures are implemented in `ragged_batch_descriptor.h`. - """ - self._config = config - self._input_ids = torch.zeros((self._config.max_ragged_batch_size), - dtype=torch.int64, - device=get_accelerator().current_device()) - - self._batch_metadata_storage = torch.zeros(2, dtype=torch.int32, device=get_accelerator().current_device()) - - self._token_to_seq_storage = torch.zeros((self._config.max_ragged_batch_size), - dtype=torch.int32, - device=get_accelerator().current_device()) - self._inflight_seq_descriptors = torch.zeros((self._config.max_ragged_sequence_count, 4), - dtype=torch.int32, - device=get_accelerator().current_device()) - self._kv_ptrs = torch.zeros((self._config.max_ragged_sequence_count), - dtype=torch.int64, - device=get_accelerator().current_device()) - - self._utils_module = RaggedUtilsBuilder().load() - host_alloc = self._utils_module.allocate_fast_host_buffer - - self._input_ids_shadow = host_alloc(self._input_ids) - self._batch_metadata_storage_shadow = host_alloc(self._batch_metadata_storage) - self._token_to_seq_storage_shadow = host_alloc(self._token_to_seq_storage) - self._inflight_seq_descriptors_shadow = host_alloc(self._inflight_seq_descriptors) - self._kv_ptrs_shadow = host_alloc(self._kv_ptrs) - - # Default behavior should be no padding - self._is_padded = False - - self._current_tokens = 0 - self._current_sequences = 0 - self._batch_tokens = [] - self._inflight_seq_descriptors_shadow_buf = [] - self._kv_blocks_ptr_buf = [] - self._token_to_seq_storage_shadow_buf = [] - - def clear(self) -> None: - """ - Clear the ragged batch. This will reset the number of tokens and sequences to 0. - """ - self._current_tokens = 0 - self._current_sequences = 0 - self._batch_tokens = [] - self._inflight_seq_descriptors_shadow_buf = [] - self._kv_blocks_ptr_buf = [] - self._token_to_seq_storage_shadow_buf = [] - - def insert_sequence(self, seq_descriptor: DSSequenceDescriptor, tokens: torch.Tensor, do_checks=True) -> None: - """ - Incrementally insert a sequence into the ragged batch. This will update the - metadata for the ragged batch and the sequence. - - Arguments: - seq_descriptor () - """ - if tokens.device != torch.device("cpu"): - # This doesn't really fall under schedulability, so we'll unconditionally check for it. - raise RuntimeError(f"Expected tokens to be on host but found device '{tokens.device}'") - - if do_checks and self.current_sequences == self._config.max_ragged_sequence_count: - raise RuntimeError(f"Ragged batch is full due to sequence limit: {self._config.max_ragged_sequence_count}") - - seq_tokens = tokens.numel() - - if do_checks and self.current_tokens + seq_tokens > self._config.max_ragged_batch_size: - raise RuntimeError(f"Ragged batch is full due to capacity limit: {self._config.max_ragged_batch_size})") - - # The values in _inflight_seq_descriptors_shadow_buf, _token_to_seq_storage_shadow_buf, _kv_blocks_ptr_buf, etc., - # are ultimately stored in PyTorch tensors: _inflight_seq_descriptors_shadow, _token_to_seq_storage_shadow, _kv_ptrs_shadow, etc. - # However, we found it inefficient to iterate over and substitute values into tensor slices or to use copy/fill calls for this purpose. - # Therefore, we initially store the values in Python lists or primitive data types and then copy them collectively in the finalize() method, - # instead of updating the tensors directly in each iteration. - self._batch_tokens.append(tokens) - self._inflight_seq_descriptors_shadow_buf.append(self.current_tokens) - self._inflight_seq_descriptors_shadow_buf.append(seq_tokens) - self._inflight_seq_descriptors_shadow_buf.append(seq_descriptor.seen_tokens) - self._inflight_seq_descriptors_shadow_buf.append(0) # alignment - - self._token_to_seq_storage_shadow_buf.extend([self.current_sequences] * seq_tokens) - - self._kv_blocks_ptr_buf.append(seq_descriptor.kv_blocks_ptr) - - self._current_tokens += seq_tokens - self._current_sequences += 1 - - @property - def tensor_toks(self) -> torch.Tensor: - """ - The number of tokens in the in-flight ragged batch. This will not trigger - synchronization with the device. - """ - cur_toks = self.current_tokens - if self._is_padded: - return to_padded(cur_toks) - else: - return cur_toks - - def finalize(self, padding: Optional[bool] = False) -> None: - """ - Completes construction of the ragged batch by flushing the host buffers to the device. - """ - cur_toks = self.current_tokens - - # Batch-copy the values recorded in insert_sequence() into PyTorch tensors to enhance efficiency. - self._inflight_seq_descriptors_shadow.flatten()[:len(self._inflight_seq_descriptors_shadow_buf)].copy_( - torch.tensor(self._inflight_seq_descriptors_shadow_buf)) - self._input_ids_shadow[:self.current_tokens].copy_(torch.cat(self._batch_tokens, dim=0)) - self._token_to_seq_storage_shadow[:len(self._token_to_seq_storage_shadow_buf)].copy_( - torch.tensor(self._token_to_seq_storage_shadow_buf)) - self._kv_ptrs_shadow[:len(self._kv_blocks_ptr_buf)].copy_(torch.tensor(self._kv_blocks_ptr_buf)) - self._batch_metadata_storage_shadow.copy_(torch.tensor([cur_toks, self.current_sequences])) - - if padding: - padded_toks = to_padded(cur_toks) - self._input_ids_shadow[cur_toks:padded_toks].fill_(-1) - self._token_to_seq_storage_shadow[cur_toks:padded_toks].fill_(-1) - self._is_padded = True - else: - padded_toks = cur_toks - self._is_padded = False - - current_sequences = self.current_sequences - - def _noblock_copy(dst: torch.Tensor, src: torch.Tensor) -> None: - dst.copy_(src, non_blocking=True) - - _noblock_copy(self._input_ids[:padded_toks], self._input_ids_shadow[:padded_toks]) - _noblock_copy(self._batch_metadata_storage, self._batch_metadata_storage_shadow) - _noblock_copy(self._token_to_seq_storage[:padded_toks], self._token_to_seq_storage_shadow[:padded_toks]) - _noblock_copy(self._inflight_seq_descriptors[:current_sequences], - self._inflight_seq_descriptors_shadow[:current_sequences]) - _noblock_copy(self._kv_ptrs[:current_sequences], self._kv_ptrs_shadow[:current_sequences]) - - def input_ids(self, on_device: bool = True) -> torch.Tensor: - """ - The input ids tensor for the ragged batch. If the device Tensor is requested, the Tensor - is truncated to the number of tokens in the batch. - """ - if on_device: - return self._input_ids[:self.tensor_toks] - else: - return self._input_ids_shadow - - def batch_metadata_buffer(self, on_device: bool = True) -> torch.Tensor: - """ - Buffer associated with the batch metadata tensor that can - be populated in preparation for passing a new input to the device. - """ - if on_device: - return self._batch_metadata_storage - else: - return self._batch_metadata_storage_shadow - - def tokens_to_seq(self, on_device: bool = True) -> torch.Tensor: - """ - Mapping of token to which sequence it belongs to in the ragged batch. If the device Tensor - is requested, the Tensor is truncated to the number of tokens in the batch. - """ - if on_device: - return self._token_to_seq_storage[:self.tensor_toks] - else: - return self._token_to_seq_storage_shadow - - def inflight_seq_descriptors(self, on_device: bool = True) -> torch.Tensor: - """ - Buffer associated with the metadata of each sequence in the ragged batch. If the device Tensor - is requested, the Tensor is truncated to the number of sequences in the batch. - """ - if on_device: - return self._inflight_seq_descriptors[:self.current_sequences] - else: - return self._inflight_seq_descriptors_shadow - - def kv_ptrs(self, on_device: bool = True) -> torch.Tensor: - """ - Pointer to where the list of KV ids associated with a sequence are. If the device Tensor - is requested, the Tensor is truncated to the number of sequences in the batch. - """ - if on_device: - return self._kv_ptrs[:self.current_sequences] - else: - return self._kv_ptrs_shadow - - def masks(self, on_device: bool = True) -> Optional[torch.Tensor]: - """ - Placeholder for supporting complex masks. Currently not supported. - - Models that will need this will be BERT-like, not generative. - """ - return None - - @property - def current_tokens(self) -> int: - """ - The number of tokens in the in-flight ragged batch. This will not trigger - synchronization with the device. - """ - return self._current_tokens - - @property - def current_sequences(self) -> int: - """ - The number of sequences in the in-flight ragged batch. This will not trigger - synchronization with the device. - """ - return self._current_sequences diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/sequence_descriptor.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/sequence_descriptor.py deleted file mode 100644 index 6b9f65255eec8ed0996c45e224941779777ef7f3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/ragged/sequence_descriptor.py +++ /dev/null @@ -1,280 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import List, Tuple, Union - -import torch - - -class BaseSequenceDescriptor: - - @property - def seen_tokens(self) -> int: - """ - The number of tokens for this sequence that have completed a forward pass. - """ - raise NotImplementedError() - - @property - def cur_allocated_blocks(self, cache_group: int = 0) -> int: - """ - The number of KV blocks currently allocated for this sequence. - """ - raise NotImplementedError() - - @property - def kv_blocks_ptr(self, cache_group: int = 0) -> int: - """ - The pointer to the KV blocks for this sequence. - """ - raise NotImplementedError() - - -class PlaceholderSequenceDescriptor(BaseSequenceDescriptor): - """ - The DummySequenceDescriptor is an empty object that allows us to perform schedulability - checks before formally tracking a sequence. - """ - - def __init__(self, seen_tokens=0, cur_allocated_blocks=0, kv_blocks_ptr=0) -> None: - self._seen_tokens = seen_tokens - self._cur_allocated_blocks = cur_allocated_blocks - self._kv_blocks_ptr = kv_blocks_ptr - - @property - def seen_tokens(self) -> int: - return self._seen_tokens - - @property - def cur_allocated_blocks(self, cache_group: int = 0) -> int: - return self._cur_allocated_blocks - - @property - def kv_blocks_ptr(self, cache_group: int = 0) -> int: - return self._kv_blocks_ptr - - -class DSSequenceDescriptor(BaseSequenceDescriptor): - - _seen_tokens: int - """ - Number of tokens in the sequence that have completed a forward pass. - """ - - _in_flight_tokens: int - """ - Number of tokens that have begun a forward pass but not yet completed it. - """ - - _max_context: int - """ - Maximum number of tokens this sequence may eventually include. Currently unused but - may be used in future implementations for speculative caching. - """ - - _num_allocation_groups: Tuple[int, ...] - """ - Number of unique allocation groups associated with the sequence for each cache group. - """ - - _blocks_per_allocation_group: Tuple[torch.IntTensor, ...] - """ - Number of blocks allocated for each allocation group in each cache group. - """ - - # Padded list of KV-cache IDs for the sequence. - _kv_cache_ids: Tuple[torch.Tensor, ...] - _kv_cache_ids_shadow: Tuple[torch.Tensor, ...] - """ - Padded list of KV-cache IDs for the sequence. The padded shape is [num_allocation_groups, max_blocks_per_allocation_group]. - """ - - # The location in the broader ID tensor where the KV-cache IDs for the sequence - # are stored. Used on flush. - _tracking_id: int - - def __init__(self, - tracking_id: int, - kv_cache_ids: Tuple[torch.Tensor, ...], - kv_cache_ids_shadow: Tuple[torch.Tensor, ...], - max_context: int = -1) -> None: - """ - Create the metadata to track a single sequence in the system. - - Arguments: - tracking_id (int): The slot in the tracking buffers used to track this sequence. - kv_cache_ids (Tuple[torch.Tensor, ...]): The KV-cache IDs for the sequence. The shape - of the tensor should be [num_allocation_groups, max_blocks_per_allocation_group]. - There should be one tensor per cache group. - kv_cache_ids_shadow (Tuple[torch.Tensor, ...]): The shadow tensor for the KV-cache IDs. - This tensor should be allocated on the host and should have the same shape as the - tensor provided in ``kv_cache_ids``. There should be one tensor per cache group. - max_context (int): The maximum number of tokens this sequence may eventually include. - Currently unused but may be used in future implementations for speculative caching. - """ - self._tracking_id = tracking_id - self._kv_cache_ids = kv_cache_ids - self._kv_cache_ids_shadow = kv_cache_ids_shadow - self._max_context = max_context - self._n_cache_groups = len(kv_cache_ids) - - self._seen_tokens = 0 - self._in_flight_tokens = 0 - - self._num_allocation_groups = tuple(kv_cache_ids_shadow.shape[0] - for kv_cache_ids_shadow in kv_cache_ids_shadow) - self._blocks_per_allocation_group = tuple( - torch.zeros(num_groups, dtype=torch.int32, device="cpu") for num_groups in self._num_allocation_groups) - - for cache_group, kv_cache_ids in enumerate(kv_cache_ids): - assert self._num_allocation_groups[cache_group] == kv_cache_ids.shape[0] - assert len(kv_cache_ids.shape) == 2 - - @property - def seen_tokens(self) -> int: - """ - Number of tokens in the sequence that have completed a forward pass. - """ - return self._seen_tokens - - @property - def in_flight_tokens(self) -> int: - """ - Number of tokens that have begun a forward pass but not yet completed it. - """ - return self._in_flight_tokens - - @property - def max_context(self) -> int: - """ - Maximum number of tokens for this sequence. Currently unused. - """ - return self._max_context - - @property - def tracking_id(self) -> int: - """ - Return the slot in the tracking buffers used to track this sequence. - """ - return self._tracking_id - - @property - def cur_allocated_blocks(self, cache_group: int = 0) -> int: - """ - Returns the number of blocks currently allocated for this sequence in the specified cache group. - - Arguments: - cache_group (int): The cache group to query. - """ - # Currently, there is only one allocation group. - # A shortcut is used here to bypass the overhead of sum(). - if len(self._blocks_per_allocation_group) == 1: - return self._blocks_per_allocation_group[0].item() - return self._blocks_per_allocation_group[cache_group].sum().item() - - def kv_cache_ids(self, cache_group: int = 0, on_device: bool = False) -> torch.Tensor: - """ - Returns the Tensor containing the block IDs for this sequence on the appropriate device - for the specified cache group. - - Arguments: - cache_group (int): The cache group to query. - on_device (bool): Whether or not to return the Tensor on the device or on the host. - """ - if on_device: - return self._kv_cache_ids[cache_group] - else: - return self._kv_cache_ids_shadow[cache_group] - - @property - def kv_blocks_ptr(self, cache_group: int = 0) -> int: - """ - Get the device pointer to the base of the KV-cache ids for the specified cache group and - sequence. - - Arguments: - cache_group (int): The cache group to query. - """ - return self._kv_cache_ids[cache_group].data_ptr() - - #TODO: this was previously a property but causing issues with PR-4668 need to consult w. Connor - def all_block_ids(self, cache_group: int = 0) -> torch.Tensor: - """ - Return the Tensor containing all block IDs for this sequence in the specified cache group. - - Arguments: - cache_group (int): The cache group to query. - """ - block_ids = [] - for allocation_group, num_blocks in zip(self._kv_cache_ids[cache_group], - self._blocks_per_allocation_group[cache_group]): - block_ids.append(allocation_group[:num_blocks]) - return torch.cat(block_ids) - - def pre_forward(self, num_tokens: int) -> None: - """ - Update the state of the sequence before a forward pass. - - Arguments: - num_tokens (int): The number of tokens in the sequence that will be executed during the - next forward pass of the model. - """ - self._in_flight_tokens = num_tokens - - def post_forward(self) -> None: - """ - Update the state of the sequence after a forward pass. This should be called after the forward - pass completes. NOTE: due to the asynchronous nature of the accelerator, this may be called - before the forward pass completes on the device itself. - """ - self._seen_tokens += self._in_flight_tokens - self._in_flight_tokens = 0 - - def extend_kv_cache(self, new_ids: Union[List[torch.IntTensor], torch.IntTensor], cache_group: int = 0) -> None: - """ - Extend the KV-cache for the sequence. - - Arguments: - new_ids (Union[List[torch.IntTensor], torch.IntTensor]): For each allocation group, the IDs - to add to the KV-cache. If there is only one allocation group, a single tensor can be - provided. Otherwise, a list of tensors should be provided. The tensors do not need - to have the same shape. - """ - if isinstance(new_ids, torch.Tensor): - new_ids = [new_ids] - - if len(new_ids) != self._num_allocation_groups[cache_group]: - raise ValueError( - f"Only {len(new_ids)} allocation groups provided, expected {self._num_allocation_groups[cache_group]}") - - for group_id, new_group_ids in enumerate(new_ids): - new_blocks = new_group_ids.numel() - - if new_blocks == 0: - # If we have multiple groups, it's possible to have an empty group. - continue - - shadow_alloc_group = self._kv_cache_ids_shadow[cache_group][group_id] - alloc_group = self._kv_cache_ids[cache_group][group_id] - cur_blocks = self._blocks_per_allocation_group[cache_group][group_id] - - shadow_alloc_group[cur_blocks:cur_blocks + new_blocks].copy_(new_group_ids) - alloc_group[cur_blocks:cur_blocks + new_blocks].copy_(shadow_alloc_group[cur_blocks:cur_blocks + - new_blocks], - non_blocking=True) - - self._blocks_per_allocation_group[cache_group][group_id] += new_blocks - - def free_kv_cache(self, free_ids: Union[List[torch.IntTensor], torch.IntTensor], cache_group: int = 0) -> None: - """ - Free blocks from the KV-cache for the sequence. - - Arguments: - free_ids (Union[List[torch.IntTensor], torch.IntTensor]): The ids of blocks to free - from the KV-cache. If there is only one allocation group, a single tensor can be - provided. Otherwise, a list of tensors should be provided. The tensors do not need - to have the same shape. - """ - raise NotImplementedError("Partial KV-cache freeing is not yet supported.") diff --git a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/scheduling_utils.py b/venv/lib/python3.10/site-packages/deepspeed/inference/v2/scheduling_utils.py deleted file mode 100644 index 6d3818d46675d9b00735454b3db021e770bb4947..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/inference/v2/scheduling_utils.py +++ /dev/null @@ -1,54 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from enum import Enum - - -class SchedulingResult(Enum): - - Success = 0 - """ - The proposed batch is valid and can be scheduled. - """ - - EngineSequenceLimitExceeded = 1 - """ - The proposed batch would would overflow the number of concurrent sequences the engine may support. - """ - - BatchSequenceLimitExceeded = 2 - """ - The proposed batch contains more sequences than the engine was configured - to support in a single forwardp - """ - - BatchTokenLimitExceeded = 3 - """ - The proposed batch contains more tokens than the engine was configured - to support in a single forward. - """ - - KVCacheLimitExceeded = 4 - """ - The proposed batch would require more KV cache to be allocated than the engine - currently has available. - """ - - SequenceTokenLimitExceeded = 5 - """ - The proposed batch contains a sequence that is longer than the engine/model can support. - """ - - -class SchedulingError(RuntimeError): - - result: SchedulingResult - """ - The failed result of the scheduling check. Guaranteed to not be SchedulingResult.Success. - """ - - def __init__(self, result: SchedulingResult) -> None: - self.result = result - super().__init__(f"Batch scheduling failed with result {result}") diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/io/__init__.py deleted file mode 100644 index aab8353b8e17254c09b002eb18a9d925d4778d99..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/io/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .py_file_writer import PyFileWriter -from .fast_file_writer import FastFileWriter, FastFileWriterConfig -from .mock_file_writer import MockFileWriter diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index c5f0244762e89619599c64a2c6ac69ffdd136ace..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/base_file_writer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/base_file_writer.cpython-310.pyc deleted file mode 100644 index 4a7c9d2b7e9b481279cf916ae8dce9cb5fa274c9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/base_file_writer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/base_io_buffer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/base_io_buffer.cpython-310.pyc deleted file mode 100644 index d9e8b64243d47f61aa6525dc7a8a236cb3f09ed1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/base_io_buffer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/constants.cpython-310.pyc deleted file mode 100644 index 62c6f1752023ab5c565b189fc3fb4eab076356eb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/constants.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/double_io_buffer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/double_io_buffer.cpython-310.pyc deleted file mode 100644 index d5c00b9d78227a90bc72c1c3c1c677b172858e65..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/double_io_buffer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/fast_file_writer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/fast_file_writer.cpython-310.pyc deleted file mode 100644 index dc964e8d0b5fe00dcc45758af849450c453bbedb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/fast_file_writer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/mock_file_writer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/mock_file_writer.cpython-310.pyc deleted file mode 100644 index 93c9eb5f19e53f0592d6235dc17d7aba7053d67c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/mock_file_writer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/py_file_writer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/py_file_writer.cpython-310.pyc deleted file mode 100644 index 649d237c20471bd754583a8609d1cce4c660f618..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/py_file_writer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/single_io_buffer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/single_io_buffer.cpython-310.pyc deleted file mode 100644 index e9bd48ca6adee552f97e16b40dabf996defb8566..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/single_io_buffer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index a411a60bc4b15f8e41f30ae21406c7f438b46908..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/io/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/base_file_writer.py b/venv/lib/python3.10/site-packages/deepspeed/io/base_file_writer.py deleted file mode 100644 index ef7d9148450cc04cb571d7ef0206d9356862cf57..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/io/base_file_writer.py +++ /dev/null @@ -1,43 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .constants import * - -BASE_STAT_KEYS = [ - CLOSE_COUNT_KEY, FILENO_COUNT_KEY, FLUSH_COUNT_KEY, WRITE_COUNT_KEY, WRITE_BYTES_KEY, WRITE_SEC_KEY, - WRITE_SPEED_KEY -] - - -class BaseFileWriter(object): - - def __init__(self, file_path): - self._file_path = file_path - self._stats = {k: 0 for k in BASE_STAT_KEYS} - - def close(self): - pass - - def fileno(self): - pass - - def flush(self): - pass - - def write(self, buffer): - pass - - def file_path(self): - return self._file_path - - def _incr_stats(self, key, incr=1): - self._stats[key] += incr - - def _dump_state(self): - if self._stats[WRITE_SEC_KEY] > 0: - self._stats[WRITE_SPEED_KEY] = (self._stats[WRITE_BYTES_KEY] / self._stats[WRITE_SEC_KEY] / (1024**3)) - state = self._stats - state[FILE_PATH_KEY] = self.file_path() - print(f'stats = {self._stats}') diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/base_io_buffer.py b/venv/lib/python3.10/site-packages/deepspeed/io/base_io_buffer.py deleted file mode 100644 index d06a20e0165612fb6ad20c2c231b159829e3ce39..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/io/base_io_buffer.py +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - - -class Base_IO_Buffer(object): - - def __init__(self, pinned_tensor, dnvme_handle): - assert pinned_tensor.numel() % dnvme_handle.get_alignment() == 0 - self._dnvme_handle = dnvme_handle - self._pinned_tensor = pinned_tensor - - def fill(self, src_tensor, src_offset): - pass - - def drain(self, num_bytes, fd, file_offset): - pass - - def is_empty(self): - pass - - def is_full(self): - pass - - def get_buffer(self): - pass - - def get_offset(self): - pass - - def get_aligned_num_bytes(self): - pass - - def get_unaligned_num_bytes(self): - pass - - def reset(self): - pass - - def complete_ongoing_drain(self): - pass - - def _drain(self, num_bytes, fd, file_offset, blocking=False): - assert num_bytes <= self.get_offset() - assert num_bytes % self._dnvme_handle.get_alignment() == 0 - buffer = self.get_buffer() - r = self._dnvme_handle.async_pwrite(torch.narrow(buffer, 0, 0, num_bytes), fd, file_offset) - assert 0 == r - if blocking: - assert 1 == self._dnvme_handle.wait() - - @staticmethod - def fill_buffer(src_tensor, src_offset, buffer_tensor, buffer_offset): - src_bytes = src_tensor.numel() - src_offset - assert src_bytes > 0 - - dst_bytes = buffer_tensor.numel() - buffer_offset - copy_bytes = min(src_bytes, dst_bytes) - assert (buffer_offset + copy_bytes) <= buffer_tensor.numel() - - if copy_bytes > 0: - src_slice = torch.narrow(src_tensor, 0, src_offset, copy_bytes) - dst_slice = torch.narrow(buffer_tensor, 0, buffer_offset, copy_bytes) - dst_slice.data.copy_(src_slice.data) - - return copy_bytes diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/constants.py b/venv/lib/python3.10/site-packages/deepspeed/io/constants.py deleted file mode 100644 index e402a365a0d3f3f3d005e4b9afbd5509cdf8aed3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/io/constants.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -INVALID_FD = -1 - -FILE_PATH_KEY = 'path' -FLUSH_COUNT_KEY = 'flush' -WRITE_COUNT_KEY = 'write' -CLOSE_COUNT_KEY = 'close' -FILENO_COUNT_KEY = 'fileno' -WRITE_BYTES_KEY = 'bytes' -WRITE_SEC_KEY = 'write_secs' -WRITE_SPEED_KEY = 'write_GB/s' - -AIO_WRITE_SEC_KEY = 'aio_write_secs' -AIO_WRITE_BYTES_KEY = 'aio_bytes' -AIO_SPEED_KEY = 'aio_GB/s' -SLOW_WRITE_BYTES_KEY = 'slow_bytes' -SLOW_WRITE_SEC_KEY = 'slow_write_secs' -AIO_FILL_BUFFER_SEC_KEY = 'fill_buffer_secs' -AIO_FILL_BUFFER_COUNT_KEY = 'fill_buffer_count' -AIO_FILL_BUFFER_SPEED_KEY = 'fill_buffer_GB/s' - -SAVE_STORAGE_KEY = 'save_storage' -SAVE_STORAGE_BYTES_KEY = 'save_storage_bytes' -SAVE_STORAGE_SEC_KEY = 'save_storage_secs' -STORAGE_OBJ_SIZE = 8 - -RANK_KEY = 'rank' diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/double_io_buffer.py b/venv/lib/python3.10/site-packages/deepspeed/io/double_io_buffer.py deleted file mode 100644 index 0491ff1ac93eb2c2f77dfce34fbc4d550c274863..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/io/double_io_buffer.py +++ /dev/null @@ -1,84 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from .base_io_buffer import Base_IO_Buffer - -NUM_BUFFERS = 2 -INVALID_BUFFER_INDEX = -1 - - -class Double_IO_Buffer(Base_IO_Buffer): - - def __init__(self, pinned_tensor, dnvme_handle): - super(Double_IO_Buffer, self).__init__(pinned_tensor, dnvme_handle) - assert self._pinned_tensor.numel() % (NUM_BUFFERS * self._dnvme_handle.get_alignment()) == 0 - self._buffers = self._split_buffer() - self._fill_index = 0 - self._drain_index = INVALID_BUFFER_INDEX - self._buffer_offset = 0 - - def fill(self, src_tensor, src_offset): - self._validate_buffer_index(self._fill_index) - copy_bytes = Base_IO_Buffer.fill_buffer(src_tensor, src_offset, self._buffers[self._fill_index], - self._buffer_offset) - self._buffer_offset += copy_bytes - return copy_bytes - - def drain(self, num_bytes, fd, file_offset): - self._validate_buffer_index(self._fill_index) - self.complete_ongoing_drain() - assert self._drain_index == INVALID_BUFFER_INDEX - self._drain(num_bytes, fd, file_offset, blocking=False) - self._drain_index = self._fill_index - self._fill_index = (self._fill_index + 1) % NUM_BUFFERS - self._buffer_offset = 0 - - def get_buffer(self): - self._validate_buffer_index(self._fill_index) - return self._buffers[self._fill_index] - - def get_offset(self): - self._validate_buffer_index(self._fill_index) - return self._buffer_offset - - def get_aligned_num_bytes(self): - self._validate_buffer_index(self._fill_index) - aligned_size = self._dnvme_handle.get_alignment() - return (self._buffer_offset // aligned_size) * aligned_size - - def get_unaligned_num_bytes(self): - self._validate_buffer_index(self._fill_index) - return self._buffer_offset % self._dnvme_handle.get_alignment() - - def is_full(self): - self._validate_buffer_index(self._fill_index) - return self._buffer_offset == self._buffers[self._fill_index].numel() - - def is_empty(self): - self._validate_buffer_index(self._fill_index) - return self._buffer_offset == 0 and not self._is_ongoing_drain() - - def reset(self): - self._buffer_offset = 0 - - def complete_ongoing_drain(self): - if self._is_ongoing_drain(): - self._wait_for_drain() - - def _split_buffer(self): - buffer_size = self._pinned_tensor.numel() // NUM_BUFFERS - return [torch.narrow(self._pinned_tensor, 0, (i * buffer_size), buffer_size) for i in range(NUM_BUFFERS)] - - def _validate_buffer_index(self, index): - assert index in [0, 1] - - def _wait_for_drain(self): - self._validate_buffer_index(self._drain_index) - assert 1 == self._dnvme_handle.wait() - self._drain_index = INVALID_BUFFER_INDEX - - def _is_ongoing_drain(self): - return self._drain_index != INVALID_BUFFER_INDEX diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/fast_file_writer.py b/venv/lib/python3.10/site-packages/deepspeed/io/fast_file_writer.py deleted file mode 100644 index 57b78072aa2edfd807cb7f5d1f17107cc17e2261..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/io/fast_file_writer.py +++ /dev/null @@ -1,267 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import torch -import time -from dataclasses import dataclass - -from .constants import * -from .base_file_writer import BaseFileWriter -from .single_io_buffer import Single_IO_Buffer -from .double_io_buffer import Double_IO_Buffer -from deepspeed.ops.op_builder import UtilsBuilder -from deepspeed.accelerator import get_accelerator - -from .utils import (tensor_to_bytes, bytes_to_tensor, obj_serialization_details) - -FASTIO_STAT_KEYS = [ - AIO_WRITE_SEC_KEY, - AIO_WRITE_BYTES_KEY, - AIO_SPEED_KEY, - SLOW_WRITE_BYTES_KEY, - SLOW_WRITE_SEC_KEY, - AIO_FILL_BUFFER_COUNT_KEY, - AIO_FILL_BUFFER_SEC_KEY, - AIO_FILL_BUFFER_SPEED_KEY, - SAVE_STORAGE_KEY, - SAVE_STORAGE_BYTES_KEY, -] - - -@dataclass -class FastFileWriterConfig: - dnvme_handle: object - pinned_tensor: torch.Tensor - double_buffer: bool = True - num_parallel_writers: int = 1 - writer_rank: int = 0 - global_rank: int = 0 - - -class FastFileWriter(BaseFileWriter): - - def __init__(self, file_path, config): - super(FastFileWriter, self).__init__(file_path) - self._aio_fd = os.open(self._file_path, flags=os.O_DIRECT | os.O_CREAT | os.O_WRONLY) - self._dnvme_handle = config.dnvme_handle - self._file_offset = 0 - io_buffer_type = Double_IO_Buffer if config.double_buffer else Single_IO_Buffer - self._io_buffer = io_buffer_type(config.pinned_tensor, self._dnvme_handle) - self._cast_to_byte_tensor = UtilsBuilder().load().cast_to_byte_tensor - self._get_serialization_details = obj_serialization_details() - self._num_parallel_writers = config.num_parallel_writers - self._writer_rank = config.writer_rank - self._global_rank = config.global_rank - - for k in FASTIO_STAT_KEYS: - self._stats[k] = 0 - - def write(self, buffer): - assert self._file_offset % self._dnvme_handle.get_alignment() == 0 - buffer_num_bytes = len(buffer) - num_written_bytes = self._write_from_tensor(bytes_to_tensor(buffer)) - assert buffer_num_bytes == num_written_bytes - return buffer_num_bytes - - def split_index_list(self, storage_obj_list, num_splits): - assert num_splits > 0 - split_list = [-1] * num_splits - # t[0] is data, t[1] is data_type - tensor_bytes_list = [len(t[0]) for t in storage_obj_list] - print(tensor_bytes_list) - total_bytes = sum(tensor_bytes_list) - bytes_per_group = total_bytes / num_splits - split_counter = 0 - tmp_size = 0 - for i in range(len(tensor_bytes_list)): - tmp_size += tensor_bytes_list[i] - if tmp_size > bytes_per_group: - split_list[split_counter] = i - tmp_size = 0 - split_counter += 1 - if split_list[num_splits - 1] == -1: - split_list[num_splits - 1] = len(tensor_bytes_list) - return split_list - - def save_torch_storage_object_list(self, storage_obj_list, save_size): - assert self._file_offset % self._dnvme_handle.get_alignment() == 0 - num_bytes_written = self._save_storage_list(storage_obj_list, save_size) - return num_bytes_written - - def close(self): - self._fini() - self._incr_stats(CLOSE_COUNT_KEY) - - def fileno(self): - self._incr_stats(FILENO_COUNT_KEY) - return INVALID_FD # self._aio_fd - - def flush(self): - self._incr_stats(FLUSH_COUNT_KEY) - - def __del__(self): - self._fini() - assert self._aio_fd == INVALID_FD - assert self._io_buffer.get_offset() == 0, \ - f'__del__ assert: pinned_offset {self._io_buffer.get_offset()} != 0' - assert self._file_offset == self._stats[WRITE_BYTES_KEY], \ - f'__del__ assert: file_offset != write_bytes - {self._file_offset} != {self._stats[WRITE_BYTES_KEY]}' - - def _fini(self): - if not self._io_buffer_is_empty(): - self._force_drain() - self._io_buffer.reset() - self._aio_fd = INVALID_FD - - def _fill_io_buffer(self, src_tensor, src_offset): - st = time.time() - copy_bytes = self._io_buffer.fill(src_tensor, src_offset) - self._incr_stats(AIO_FILL_BUFFER_SEC_KEY, time.time() - st) - self._incr_stats(AIO_FILL_BUFFER_COUNT_KEY) - return copy_bytes - - def _drain_io_buffer(self, num_bytes): - st = time.time() - self._io_buffer.drain(num_bytes, self._aio_fd, self._file_offset) - self._incr_stats(AIO_WRITE_SEC_KEY, time.time() - st) - self._incr_stats(AIO_WRITE_BYTES_KEY, num_bytes) - self._file_offset += num_bytes - - def _io_buffer_is_full(self): - return self._io_buffer.is_full() - - def _io_buffer_is_empty(self): - return self._io_buffer.is_empty() - - def _force_drain(self): - st = time.time() - aligned_num_bytes = self._io_buffer.get_aligned_num_bytes() - # Important to retrieve unaligned drain bytes and tensor before doing aligned drain because of the side effects. - # TODO: Need to eliminate this dependency - unaligned_num_bytes = self._io_buffer.get_unaligned_num_bytes() - unaligned_tensor = torch.narrow(self._io_buffer.get_buffer(), 0, aligned_num_bytes, unaligned_num_bytes) - - if aligned_num_bytes > 0: - self._drain_io_buffer(aligned_num_bytes) - - self._io_buffer.complete_ongoing_drain() - self._incr_stats(AIO_WRITE_SEC_KEY, time.time() - st) - - if unaligned_num_bytes > 0: - self._unaligned_drain(unaligned_tensor) - self._incr_stats(WRITE_SEC_KEY, time.time() - st) - - def _unaligned_drain(self, unaligned_tensor): - os.close(self._aio_fd) - st = time.time() - fp = open(self._file_path, 'ab') - fp.write(tensor_to_bytes(unaligned_tensor.cpu())) - fp.close() - self._file_offset += unaligned_tensor.numel() - self._incr_stats(SLOW_WRITE_SEC_KEY, time.time() - st) - self._incr_stats(SLOW_WRITE_BYTES_KEY, unaligned_tensor.numel()) - self._aio_fd = os.open(self._file_path, flags=os.O_DIRECT | os.O_WRONLY | os.O_APPEND) - - def _dump_state(self): - if self._stats[AIO_WRITE_SEC_KEY] > 0: - self._stats[AIO_SPEED_KEY] = (self._stats[AIO_WRITE_BYTES_KEY] / self._stats[AIO_WRITE_SEC_KEY] / - (1024**3)) - if self._stats[AIO_FILL_BUFFER_SEC_KEY] > 0: - self._stats[AIO_FILL_BUFFER_SPEED_KEY] = (self._stats[AIO_WRITE_BYTES_KEY] / - self._stats[AIO_FILL_BUFFER_SEC_KEY] / (1024**3)) - super()._dump_state() - - def _update_write_stats(self, num_bytes, secs_latency): - self._incr_stats(WRITE_COUNT_KEY) - self._incr_stats(WRITE_BYTES_KEY, num_bytes) - self._incr_stats(WRITE_SEC_KEY, secs_latency) - - def _write_from_tensor(self, buffer_tensor): - st = time.time() - buffer_offset = 0 - while (buffer_offset < buffer_tensor.numel()): - num_copied_bytes = self._fill_io_buffer(buffer_tensor, buffer_offset) - if self._io_buffer_is_full(): - self._drain_io_buffer(self._io_buffer.get_offset()) - buffer_offset += num_copied_bytes - - self._update_write_stats(buffer_offset, time.time() - st) - - return buffer_offset - - def _save_storage_list(self, obj_list, save_size): - byte_tensor_list, byte_tensor_nbytes = self._convert_to_byte_tensors(obj_list, save_size) - if self._num_parallel_writers > 1: - my_byte_tensor_list = self._partition_byte_tensors(byte_tensor_list, byte_tensor_nbytes, - self._num_parallel_writers, self._writer_rank) - else: - my_byte_tensor_list = byte_tensor_list - - num_object_bytes_written = 0 - for byte_tensor in my_byte_tensor_list: - num_object_bytes_written += self._write_from_tensor(byte_tensor) - - self._incr_stats(SAVE_STORAGE_KEY, len(obj_list)) - self._incr_stats(SAVE_STORAGE_BYTES_KEY, num_object_bytes_written) - return num_object_bytes_written - - # Convert list of storage objects into list of byte tensors of object and size bytes - def _convert_to_byte_tensors(self, obj_list, save_size): - tensor_list = [] - num_bytes = 0 - for storage_obj in obj_list: - details = self._get_serialization_details(storage_obj) - if save_size: - tensor_list.append( - torch.tensor( - details.size, - dtype=torch.int64, - ).to(get_accelerator().device_name())) - tensor_list.append(torch.empty(0, dtype=details.dtype, device=details.obj.device).set_(details.obj)) - num_bytes += details.nbytes - if save_size: - num_bytes += STORAGE_OBJ_SIZE * len(obj_list) - - return self._cast_to_byte_tensor(tensor_list), num_bytes - - def _partition_byte_tensors(self, byte_tensor_list, byte_tensor_nbytes, num_ranks, my_rank): - assert my_rank >= 0, f'Invalid for rank number to be negative: {my_rank}' - assert num_ranks > my_rank, f'Number of ranks {num_ranks} must be greater than rank {my_rank}' - - partition_size = int(byte_tensor_nbytes // num_ranks) - num_remainder_bytes = byte_tensor_nbytes % num_ranks - if num_remainder_bytes == 0: - partition_start = partition_size * my_rank - else: - # Spread extra bytes evenly among early ranks - if num_remainder_bytes > my_rank: - partition_size += 1 - partition_start = partition_size * my_rank - else: - # Account for allocation of extra bytes to earlier ranks - partition_start = (partition_size * my_rank) + num_remainder_bytes - - partition_end = min(partition_start + partition_size, byte_tensor_nbytes) - partition_tensor_list = [] - current_offset = 0 - for byte_tensor in byte_tensor_list: - byte_tensor_end = current_offset + byte_tensor.numel() - if current_offset < partition_end and byte_tensor_end > partition_start: - fragment_start = max(current_offset, partition_start) - fragment_end = min(byte_tensor_end, partition_end) - assert fragment_start < fragment_end, \ - f'fragment start {fragment_start} should be < fragment_end {fragment_end}' - - fragment_numel = fragment_end - fragment_start - partition_tensor_list.append(byte_tensor.narrow(0, fragment_start - current_offset, fragment_numel)) - - current_offset += byte_tensor.numel() - - actual_partition_nbytes = sum([t.numel() for t in partition_tensor_list]) - assert actual_partition_nbytes == partition_size, \ - f'Incorrect partition bytes for rank {my_rank}, expected = {partition_size} actual = {actual_partition_nbytes}' - - return partition_tensor_list diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/mock_file_writer.py b/venv/lib/python3.10/site-packages/deepspeed/io/mock_file_writer.py deleted file mode 100644 index 5957ad771974010e98c7d57b7743f1c64da14ef6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/io/mock_file_writer.py +++ /dev/null @@ -1,49 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .constants import * -from .base_file_writer import BaseFileWriter -from .utils import obj_serialization_details - - -class MockFileWriter(BaseFileWriter): - - def __init__(self, file_path): - super(MockFileWriter, self).__init__(file_path) - self._fp = open(file_path, 'wb') - self._stats[SAVE_STORAGE_KEY] = 0 - self._stats[SAVE_STORAGE_BYTES_KEY] = 0 - self._get_serialization_details = obj_serialization_details() - - def close(self): - self._incr_stats(CLOSE_COUNT_KEY) - self._fp.close() - - def fileno(self): - self._incr_stats(FILENO_COUNT_KEY) - return INVALID_FD # self._fp.fileno() - - def flush(self): - self._incr_stats(FLUSH_COUNT_KEY) - self._fp.flush() - - def write(self, buffer): - return self._write(len(buffer)) - - def save_torch_storage_object_list(self, storage_obj_list, save_size): - num_bytes = sum([self._save_torch_storage_object(obj, save_size) for obj in storage_obj_list]) - return num_bytes - - def _save_torch_storage_object(self, storage_obj, save_size): - details = self._get_serialization_details(storage_obj) - self._incr_stats(SAVE_STORAGE_KEY) - self._incr_stats(SAVE_STORAGE_BYTES_KEY, details.size) - num_written_bytes = self._write(STORAGE_OBJ_SIZE) if save_size else 0 - return num_written_bytes + self._write(details.size) - - def _write(self, num_bytes): - self._incr_stats(WRITE_COUNT_KEY) - self._incr_stats(WRITE_BYTES_KEY, num_bytes) - return num_bytes diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/py_file_writer.py b/venv/lib/python3.10/site-packages/deepspeed/io/py_file_writer.py deleted file mode 100644 index 51849cb89c8af07c59f5a6f9a7e297d58405b13a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/io/py_file_writer.py +++ /dev/null @@ -1,35 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import time -from .constants import * -from .base_file_writer import BaseFileWriter - - -class PyFileWriter(BaseFileWriter): - - def __init__(self, file_path): - super(PyFileWriter, self).__init__(file_path) - self._fp = open(file_path, 'wb') - - def close(self): - self._incr_stats(CLOSE_COUNT_KEY) - self._fp.close() - - def fileno(self): - self._incr_stats(FILENO_COUNT_KEY) - return INVALID_FD # self._fp.fileno() - - def flush(self): - self._incr_stats(FLUSH_COUNT_KEY) - self._fp.flush() - - def write(self, buffer): - st = time.time() - self._fp.write(buffer) - self._incr_stats(WRITE_SEC_KEY, time.time() - st) - self._incr_stats(WRITE_COUNT_KEY) - self._incr_stats(WRITE_BYTES_KEY, len(buffer)) - return len(buffer) diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/single_io_buffer.py b/venv/lib/python3.10/site-packages/deepspeed/io/single_io_buffer.py deleted file mode 100644 index 57fbfe38c45d4b6d54ac1ef4b07a6da17f732b91..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/io/single_io_buffer.py +++ /dev/null @@ -1,44 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .base_io_buffer import Base_IO_Buffer - - -class Single_IO_Buffer(Base_IO_Buffer): - - def __init__(self, pinned_tensor, dnvme_handle): - super(Single_IO_Buffer, self).__init__(pinned_tensor, dnvme_handle) - self._pinned_offset = 0 - - def fill(self, src_tensor, src_offset): - copy_bytes = Base_IO_Buffer.fill_buffer(src_tensor, src_offset, self._pinned_tensor, self._pinned_offset) - self._pinned_offset += copy_bytes - return copy_bytes - - def drain(self, num_bytes, fd, file_offset): - self._drain(num_bytes, fd, file_offset, blocking=True) - self._pinned_offset = 0 - - def get_buffer(self): - return self._pinned_tensor - - def get_offset(self): - return self._pinned_offset - - def get_aligned_num_bytes(self): - aligned_size = self._dnvme_handle.get_alignment() - return (self._pinned_offset // aligned_size) * aligned_size - - def get_unaligned_num_bytes(self): - return self._pinned_offset % self._dnvme_handle.get_alignment() - - def is_full(self): - return self._pinned_offset == self._pinned_tensor.numel() - - def is_empty(self): - return self._pinned_offset == 0 - - def reset(self): - self._pinned_offset = 0 diff --git a/venv/lib/python3.10/site-packages/deepspeed/io/utils.py b/venv/lib/python3.10/site-packages/deepspeed/io/utils.py deleted file mode 100644 index f811aeb84577570a2d284e3c5e181c9b003a5a0e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/io/utils.py +++ /dev/null @@ -1,56 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import numpy -import torch -from dataclasses import dataclass - - -@dataclass -class serialize_details: - obj: object - dtype: torch.dtype - size: int - nbytes: int - - -def tensor_to_bytes(tensor): - return tensor.numpy().tobytes() - - -def bytes_to_tensor(buffer): - return torch.from_numpy(numpy.array(numpy.frombuffer(buffer, dtype=numpy.uint8))) - - -def required_minimum_torch_version(major_version, minor_version): - TORCH_MAJOR = int(torch.__version__.split('.')[0]) - TORCH_MINOR = int(torch.__version__.split('.')[1]) - - if TORCH_MAJOR < major_version: - return False - - return TORCH_MAJOR > major_version or TORCH_MINOR >= minor_version - - -# torch < 1.12 -def _legacy_obj_serialization_details(storage_obj): - nbytes = storage_obj.element_size() * storage_obj.size() - return serialize_details(obj=storage_obj, dtype=storage_obj.dtype, size=nbytes, nbytes=nbytes) - - -# torch >= 1.12 -def _new_obj_serialization_details(storage_obj): - obj, dtype = storage_obj - return serialize_details(obj=obj, - dtype=dtype, - size=obj.size() // torch._utils._element_size(dtype), - nbytes=obj.size()) - - -def obj_serialization_details(): - if required_minimum_torch_version(1, 12): - return _new_obj_serialization_details - - return _legacy_obj_serialization_details diff --git a/venv/lib/python3.10/site-packages/deepspeed/launcher/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/launcher/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/launcher/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 20d2f1bb9bb9736fecbe01407fb28d50fa938ee4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/constants.cpython-310.pyc deleted file mode 100644 index dc7f464d686d69d19d258dc4da30a6899939c7b3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/constants.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/launch.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/launch.cpython-310.pyc deleted file mode 100644 index 740da94c2e524d048fc41e289113bc49bdc6b198..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/launch.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/launcher_helper.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/launcher_helper.cpython-310.pyc deleted file mode 100644 index 03f17274f7bd480d6aa3fbad1f85725190b34613..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/launcher_helper.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/multinode_runner.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/multinode_runner.cpython-310.pyc deleted file mode 100644 index 45b3eb76aa2da46a3cf8458f01c4f1b4dc9c8b3d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/multinode_runner.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/runner.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/runner.cpython-310.pyc deleted file mode 100644 index f6ae930932a7ef9c3ef5b1c823c8b1ae3c1f4d04..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/launcher/__pycache__/runner.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/launcher/constants.py b/venv/lib/python3.10/site-packages/deepspeed/launcher/constants.py deleted file mode 100644 index 366ae0b236f81a94b1a4ed2cecc01a1f3e1aa69e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/launcher/constants.py +++ /dev/null @@ -1,16 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -PDSH_LAUNCHER = 'pdsh' -PDSH_MAX_FAN_OUT = 1024 - -OPENMPI_LAUNCHER = 'openmpi' -MPICH_LAUNCHER = 'mpich' -IMPI_LAUNCHER = 'impi' -SLURM_LAUNCHER = 'slurm' -MVAPICH_LAUNCHER = 'mvapich' -MVAPICH_TMP_HOSTFILE = '/tmp/deepspeed_mvapich_hostfile' - -ELASTIC_TRAINING_ID_DEFAULT = "123456789" diff --git a/venv/lib/python3.10/site-packages/deepspeed/launcher/launch.py b/venv/lib/python3.10/site-packages/deepspeed/launcher/launch.py deleted file mode 100644 index 2abfcae80b2b128b641a6048db94939c1dbaf24f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/launcher/launch.py +++ /dev/null @@ -1,359 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -DeepSpeed launcher, this is similar to torch's distributed.launch but supports -additional features such as arbitrary gpu exclusion. - -deepspeed.launcher.launch is intended to be run on a single worker node and -will spawn several worker sub-processes depending on how many devices/ranks -are on the worker. -""" - -import sys -import subprocess -import os -import json -import base64 -import time -import signal -import psutil -from collections import defaultdict -from typing import Dict -from argparse import ArgumentParser, REMAINDER -from ..constants import TORCH_DISTRIBUTED_DEFAULT_PORT, CROSS_RANK, CROSS_SIZE -from deepspeed.accelerator import get_accelerator -from ..nebula.constants import DLTS_POD_ENV_PATH -from ..utils import logger, get_numactl_cmd -from ..elasticity import is_torch_elastic_compatible -from .constants import ELASTIC_TRAINING_ID_DEFAULT - -PID_FILE_BASEPATH = "/tmp" - - -def parse_args(): - parser = ArgumentParser(description="DeepSpeed distributed training launch" - " utility that creates multiple distributed" - " processes on a single node") - - # Optional arguments for the launch helper - parser.add_argument("--node_rank", - type=int, - default=0, - help="The rank of the node for multi-node distributed " - "training") - parser.add_argument("--master_addr", - default="127.0.0.1", - type=str, - help="Master node (rank 0)'s address, should be either" - " the IP address or the hostname of node 0, for" - " single node multi-proc training, the" - " --master_addr can simply be 127.0.0.1") - parser.add_argument("--master_port", - default=TORCH_DISTRIBUTED_DEFAULT_PORT, - type=int, - help="Master node (rank 0)'s free port that needs to " - "be used for communication during distributed " - "training") - parser.add_argument("--world_info", default="None", type=str, help="world info base64 encoded dictionary") - - parser.add_argument("--module", - action="store_true", - help="Change each process to interpret the launch " - "script as a Python module, executing with the same " - "behavior as 'python -m'.") - - parser.add_argument("--no_python", - action="store_true", - help="Skip prepending the training script with " - "'python' - just execute it directly.") - - parser.add_argument("--enable_elastic_training", action="store_true", help="Enable elastic training support.") - - parser.add_argument("--min_elastic_nodes", type=int, default=-1, help="Min number of nodes in elastic training.") - - parser.add_argument("--max_elastic_nodes", type=int, default=-1, help="Max number of nodes in elastic training.") - - parser.add_argument("--no_local_rank", - action="store_true", - help="Do not pass local_rank as an argument when calling " - "the user's training script.") - - parser.add_argument("--save_pid", - type=int, - default=0, - help="main launching process pid, for internal pid tracking") - - parser.add_argument("--enable_each_rank_log", - default="None", - type=str, - help="redirect the stdout and stderr from each rank into different log files") - - parser.add_argument("--bind_cores_to_rank", - action="store_true", - help="Bind each rank to different cores of the host. " - "This improves host efficiency especially for CPU backend") - - parser.add_argument("--bind_core_list", - type=str, - default=None, - help="List of cores to bind to with comma separated list of " - "numbers and range. i.e. 1,3-5,7 => [1,3,4,5,7]. When not " - "specified, all cores on system would be used rank binding") - - # positional - parser.add_argument("training_script", - type=str, - help="The full path to the single GPU training " - "program/script to be launched in parallel, " - "followed by all the arguments for the " - "training script") - - # rest from the training program - parser.add_argument('training_script_args', nargs=REMAINDER) - return parser.parse_args() - - -# Adapted from https://psutil.readthedocs.io/en/latest/#kill-process-tree -def terminate_process_tree(pid): - process = psutil.Process(pid) - children = process.children(recursive=True) - children.append(process) - for child in children: - try: - child.terminate() - except psutil.NoSuchProcess: - pass - gone, alive = psutil.wait_procs(children, timeout=30) - for p in alive: - p.kill() - - -def main(): - args = parse_args() - current_env = os.environ.copy() - - for k in current_env.keys(): - if "NCCL" in k: - logger.info(f"{args.node_rank} {k}={current_env[k]}") - - if args.world_info == "None": - raise ValueError("world_info can not be None") - world_info = base64.urlsafe_b64decode(args.world_info) - world_info = json.loads(world_info) - - logger.info(f"WORLD INFO DICT: {world_info}") - node_list = list(world_info.keys()) - args.nnodes = len(node_list) - local_node = node_list[args.node_rank] - local_accelerator_ids = world_info[local_node] - num_local_procs = len(local_accelerator_ids) - logger.info(f"nnodes={args.nnodes}, num_local_procs={num_local_procs}, node_rank={args.node_rank}") - - global_rank_mapping = defaultdict(list) - curr_global_rank = 0 - dist_world_size = 0 - for node_id in node_list: - gids = world_info[node_id] - dist_world_size += len(gids) - for gid in gids: - global_rank_mapping[node_id].append(curr_global_rank) - curr_global_rank += 1 - logger.info(f"global_rank_mapping={global_rank_mapping}") - logger.info(f"dist_world_size={dist_world_size}") - - get_accelerator().set_visible_devices_envs(current_env, local_accelerator_ids) - for env in get_accelerator().visible_devices_envs(): - logger.info(f"Setting {env}={current_env[env]}") - - # set PyTorch distributed related environmental variables - current_env["MASTER_ADDR"] = args.master_addr - current_env["MASTER_PORT"] = str(args.master_port) - current_env["WORLD_SIZE"] = str(dist_world_size) - current_env[CROSS_RANK] = str(args.node_rank) - current_env[CROSS_SIZE] = str(args.nnodes) - current_env["LOCAL_SIZE"] = str(num_local_procs) - - if args.save_pid: - print(f"launcher pid: {os.getpid()}") - - pid_file = None - if args.save_pid: - launcher_pid = os.getpid() - pid_file = os.path.join(PID_FILE_BASEPATH, f"{args.save_pid}.deepspeed") - assert not os.path.isfile(pid_file), "pid file exists but shouldn't" - with open(pid_file, 'w') as fd: - fd.write(f"{launcher_pid}") - - if not is_torch_elastic_compatible(): - if args.enable_elastic_training: - logger.info(f"Disabling elastic training support as \ - PyTorch version should be greater than 1.11.x") - args.enable_elastic_training = False - - if os.path.exists(DLTS_POD_ENV_PATH): - with open(DLTS_POD_ENV_PATH) as file: - lines = file.readlines() - lines = [line.rstrip() for line in lines] - for line in lines: - if line.startswith('export FC_TASKROLE_NAME') or line.startswith('export FC_TASK_INDEX'): - key_val = line.split()[1] - key, val = key_val.split('=') - current_env[key] = val - - processes = [] - cmd = [] - - if not args.enable_elastic_training: - if args.enable_each_rank_log != "None": - # prepare the log path and the file name prefix - if os.path.isfile(args.enable_each_rank_log): - raise ValueError(f"{args.enable_each_rank_log} should not be a file, it should be a directory.") - if not os.path.exists(args.enable_each_rank_log): - try: - os.makedirs(args.enable_each_rank_log) - except Exception as e: - print(e) - raise ValueError(f"unable to create directory {args.enable_each_rank_log} for each rank log.") - log_name_prefix = time.strftime("%Y%m%d%H%M%S", time.localtime()) - - for local_proc in range(0, num_local_procs): - # each process's rank - dist_rank = global_rank_mapping[local_node][local_proc] - local_rank = dist_rank % num_local_procs - current_env["RANK"] = str(dist_rank) - current_env["LOCAL_RANK"] = str(local_rank) - - # spawn the processes - cmd = [] - if args.bind_cores_to_rank: - cores_per_rank, numactl_cmd = get_numactl_cmd(args.bind_core_list, num_local_procs, local_rank) - current_env["OMP_NUM_THREADS"] = f"{cores_per_rank}" - cmd = cmd + numactl_cmd - if not args.no_python: - cmd.append(sys.executable) - cmd.append("-u") - if args.module: - cmd.append("-m") - else: - if args.module: - raise ValueError("Don't use both the '--no_python' flag" - " and the '--module' flag at the same time.") - cmd.append(args.training_script) - # A user may not want to pass local_rank as a keyword arg so we make this optional. - if not args.no_local_rank: - cmd.append(f"--local_rank={local_rank}") - cmd += args.training_script_args - - if args.enable_each_rank_log != "None": - log_file = os.path.join(args.enable_each_rank_log, f"{log_name_prefix}_rank{dist_rank}.log") - log_fd = open(log_file, 'w') - process = subprocess.Popen(cmd, env=current_env, stdout=log_fd, stderr=log_fd) - else: - process = subprocess.Popen(cmd, env=current_env) - # logs the command from processes - logger.info(f"process {process.pid} spawned with command: {cmd}") - processes.append(process) - else: - from ..elasticity import DSElasticAgent - from torch.distributed.elastic.rendezvous import RendezvousParameters - from torch.distributed.elastic.agent.server.api import WorkerSpec - import torch.distributed.elastic.rendezvous.registry as rdzv_registry - from torch.distributed.elastic.multiprocessing import Std - - if args.min_elastic_nodes == -1: - args.min_elastic_nodes = 1 - if args.max_elastic_nodes == -1: - args.max_elastic_nodes = args.nnodes - assert args.max_elastic_nodes > 0 and args.min_elastic_nodes > 0, "Max and Min nodes should be positive" - - current_env["NCCL_ASYNC_ERROR_HANDLING"] = str(1) - - # Get config and arguments - cmd = [] - if not args.no_python: - cmd = [sys.executable, "-u"] - if args.module: - cmd.append("-m") - else: - if args.module: - raise ValueError("Don't use both the '--no_python' flag" - " and the '--module' flag at the same time.") - cmd.append(args.training_script) - cmd += args.training_script_args - cmd_args = cmd[1:] - - rdzv_configs: Dict[str, str] = {'timeout': 100} - run_id = os.environ.get("ELASTIC_RUN_ID", ELASTIC_TRAINING_ID_DEFAULT) - - # Creating config for rendezvous class - rdzv_parameters = RendezvousParameters(backend='c10d', - endpoint=args.master_addr + ":" + str(args.master_port), - run_id=run_id, - min_nodes=args.min_elastic_nodes, - max_nodes=args.max_elastic_nodes, - **rdzv_configs) - - spec = WorkerSpec( - role='trainer', - local_world_size=num_local_procs, - entrypoint=cmd[0], - args=cmd[1:], - rdzv_handler=rdzv_registry.get_rendezvous_handler(rdzv_parameters), - max_restarts=100, - monitor_interval=5, - redirects=Std.from_str("0"), - tee=Std.from_str("0"), - master_addr=None, - master_port=None, - ) - agent = DSElasticAgent(spec, current_env) - agent.run() - - sig_names = {2: "SIGINT", 15: "SIGTERM"} - last_return_code = None - - def sigkill_handler(signum, frame): - for process in processes: - logger.info(f"Killing subprocess {process.pid}") - try: - terminate_process_tree(process.pid) - except Exception: - pass - if last_return_code is not None: - logger.error(f"{cmd} exits with return code = {last_return_code}") - sys.exit(last_return_code) - if signum in sig_names: - logger.info(f"Main process received {sig_names[signum]}, exiting") - if args.save_pid: - if os.path.isfile(pid_file): - os.remove(pid_file) - sys.exit(1) - - # pass SIGINT/SIGTERM to children if the parent is being terminated - signal.signal(signal.SIGINT, sigkill_handler) - signal.signal(signal.SIGTERM, sigkill_handler) - - alive_processes = set(processes) - while len(alive_processes): - finished_processes = [] - for process in alive_processes: - if process.poll() is None: - # the process is still running - continue - else: - if process.returncode != 0: - last_return_code = process.returncode # for sigkill_handler - sigkill_handler(signal.SIGTERM, None) # not coming back - else: - # exited cleanly - logger.info(f"Process {process.pid} exits successfully.") - finished_processes.append(process) - alive_processes = set(alive_processes) - set(finished_processes) - - time.sleep(1) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/deepspeed/launcher/launcher_helper.py b/venv/lib/python3.10/site-packages/deepspeed/launcher/launcher_helper.py deleted file mode 100644 index 05ce14bcc52ed9dd2c6182b30be541a1873569dc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/launcher/launcher_helper.py +++ /dev/null @@ -1,108 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import sys -import argparse -import subprocess -from deepspeed.utils import logger -from deepspeed.launcher.constants import MPICH_LAUNCHER - - -def parse_args(args=None): - parser = argparse.ArgumentParser(description="DeepSpeed launcher helper to map environment variables for" - "multi-node/multi-gpu training jobs.", - formatter_class=argparse.ArgumentDefaultsHelpFormatter) - - parser.add_argument("--launcher", - default=MPICH_LAUNCHER, - type=str, - help="(optional) choose launcher backend for multi-node " - "training. Options currently include MPICH.") - - parser.add_argument("--module", - action="store_true", - help="Change each process to interpret the launch " - "script as a Python module, executing with the same " - "behavior as 'python -m'.") - - parser.add_argument("--no_python", - action="store_true", - help="Skip prepending the training script with " - "'python' - just execute it directly.") - - parser.add_argument("user_script", type=str, help="User script to launch, followed by any required " - "arguments.") - - parser.add_argument('user_args', nargs=argparse.REMAINDER) - - parser.add_argument("--bind_cores_to_rank", - action="store_true", - help="Bind each rank to different cores of the host") - - parser.add_argument("--bind_core_list", - type=str, - default=None, - help="List of cores to bind to with comma separated list of " - "numbers and range. i.e. 1,3-5,7 => [1,3,4,5,7]. When not " - "specified, all cores on system would be used rank binding") - - return parser.parse_args(args=args) - - -def env_mapping(env, rank_name_list=None, local_rank_name_list=None): - rank = None - for rank_name in rank_name_list: - if rank_name in env: - if rank == None: - rank = env.get(rank_name) - elif rank != env.get(rank_name): - raise EnvironmentError(f"rank number doesn't match!") - if rank == None: - raise EnvironmentError(f"rank number is not in current env!") - env['RANK'] = rank - - local_rank = None - for local_rank_name in local_rank_name_list: - if local_rank_name in env: - if local_rank == None: - local_rank = env.get(local_rank_name) - elif local_rank != env.get(local_rank_name): - raise EnvironmentError(f"local_rank number doesn't match!") - if local_rank == None: - raise EnvironmentError(f"rank number is not in current env!") - env['LOCAL_RANK'] = local_rank - - return env - - -def main(args=None): - args = parse_args(args) - - env = os.environ.copy() - - args.launcher = args.launcher.lower() - if args.launcher == MPICH_LAUNCHER: - rank_name_list = ["PMIX_RANK"] + ["PMI_RANK"] - local_rank_name_list = ["PALS_LOCAL_RANKID"] + ["MPI_LOCALRANKID"] - env = env_mapping(env, rank_name_list=rank_name_list, local_rank_name_list=local_rank_name_list) - else: - raise NotImplementedError(f"Unknown launcher {args.launcher}") - - python_exec = [] - if not args.no_python: - python_exec += [sys.executable, "-u"] - if args.module: - python_exec.append("-m") - cmd = python_exec + [args.user_script] + args.user_args - - logger.info(f"launcher_helper cmd = {' '.join(cmd)}") - - result = subprocess.Popen(cmd, env=env, close_fds=False) - result.wait() - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/deepspeed/launcher/multinode_runner.py b/venv/lib/python3.10/site-packages/deepspeed/launcher/multinode_runner.py deleted file mode 100644 index 13ab07b6314d1a97f84d5deff60500af8388c80e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/launcher/multinode_runner.py +++ /dev/null @@ -1,495 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import sys -import shutil -import subprocess -import warnings -import re -from shlex import split -from abc import ABC, abstractmethod -from deepspeed.accelerator import get_accelerator -from ..utils import logger, get_numactl_cmd -from .constants import PDSH_MAX_FAN_OUT, MVAPICH_TMP_HOSTFILE - - -class MultiNodeRunner(ABC): - - def __init__(self, args, world_info_base64): - self.args = args - self.validate_args() - self.user_arguments = self.parse_user_args() - self.user_script = args.user_script - self.world_info_base64 = world_info_base64 - self.exports = {} - - @abstractmethod - def backend_exists(self): - """Return whether the corresponding backend exists""" - - @abstractmethod - def get_cmd(self, environment, active_resources): - """Return the command to execute on node""" - - def add_export(self, key, var): - var = var.strip() - if re.search(r'[^\w@%+=:,./-]', var): - var = f"\"{var}\"" - self.exports[key.strip()] = var - - def parse_user_args(self): - return self.args.user_args - - @property - def name(self): - """Return the name of the backend""" - return self.__class__.__name__ - - def validate_args(self): - """Validate self.args""" - - -class PDSHRunner(MultiNodeRunner): - - def __init__(self, args, world_info_base64): - super().__init__(args, world_info_base64) - - def backend_exists(self): - return shutil.which('pdsh') - - def parse_user_args(self): - processed_args = [] - for arg in self.args.user_args: - # With pdsh, if we are passing a string as an argument, it will get - # split on whitespace. To avoid this and support strings that - # contain '"', we do this extra processing step: - if " " in arg: - arg = '"{}"'.format(arg.replace('"', '\\"')) - processed_args.append(arg) - return processed_args - - @property - def name(self): - return "pdsh" - - def get_cmd(self, environment, active_resources): - environment['PDSH_RCMD_TYPE'] = 'ssh' - if self.args.ssh_port is not None: # only specify ssh port if it is specified - environment["PDSH_SSH_ARGS_APPEND"] = f"{environment.get('PDSH_SSH_ARGS_APPEND', '')} \ - -p {self.args.ssh_port}" - - active_workers = ",".join(active_resources.keys()) - logger.info("Running on the following workers: %s" % active_workers) - - # PDSH flags for max node fan out and specific hosts to launch on - # See https://linux.die.net/man/1/pdsh for flag details - pdsh_cmd_args = ['pdsh', '-S', '-f', str(PDSH_MAX_FAN_OUT), '-w', active_workers] + split( - self.args.launcher_args) - - exports = "" - for key, val in self.exports.items(): - exports += "export {}={}; ".format(key, val) - - # https://linux.die.net/man/1/pdsh - # %n will be replaced by pdsh command - deepspeed_launch = [ - exports, f"cd {os.path.abspath('.')};", sys.executable, "-u", "-m", "deepspeed.launcher.launch", - f'--world_info={self.world_info_base64}', "--node_rank=%n", f"--master_addr={self.args.master_addr}", - f"--master_port={self.args.master_port}" - ] - if self.args.venv_script is not None: - deepspeed_launch = [f"source {self.args.venv_script}"] + deepspeed_launch - if self.args.no_python: - deepspeed_launch.append("--no_python") - if self.args.module: - deepspeed_launch.append("--module") - if self.args.no_local_rank: - deepspeed_launch.append("--no_local_rank") - if self.args.save_pid: - deepspeed_launch += ["--save_pid", f"{os.getpid()}"] - if self.args.enable_each_rank_log: - deepspeed_launch.append(f"--enable_each_rank_log={self.args.enable_each_rank_log}") - if self.args.elastic_training: - deepspeed_launch.append("--enable_elastic_training") - deepspeed_launch.append(f"--max_elastic_nodes={self.args.max_elastic_nodes}") - deepspeed_launch.append(f"--min_elastic_nodes={self.args.min_elastic_nodes}") - - cmd_to_search = [i + "\\" for i in deepspeed_launch[2:6]] - - kill_command = pdsh_cmd_args + ["pkill -f ", " ".join(cmd_to_search)[:-2]] - return pdsh_cmd_args + deepspeed_launch + [self.user_script] + self.user_arguments, kill_command, environment - - -class OpenMPIRunner(MultiNodeRunner): - - def __init__(self, args, world_info_base64, resource_pool): - super().__init__(args, world_info_base64) - self.resource_pool = resource_pool - self.add_export('UCX_TLS', 'tcp') - - def backend_exists(self): - #TODO: if IB is available we should suggestion mvapich - return shutil.which('ompi_info') - - @property - def name(self): - return "openmpi" - - def validate_args(self): - super().validate_args() - - # Validate and set MPI environment variables - self._setup_mpi_environment() - - #TODO: Allow for include/exclude at node-level but not gpu-level - if self.args.include != "" or self.args.exclude != "": - raise ValueError(f"{self.name} backend does not support worker include/exclusion") - if self.args.num_nodes != -1 or self.args.num_gpus != -1: - raise ValueError(f"{self.name} backend does not support limiting num nodes/gpus") - - def _setup_mpi_environment(self): - """Sets up MPI-related environment variables or raises an error if they're missing.""" - - required_vars = ['OMPI_COMM_WORLD_LOCAL_RANK', 'OMPI_COMM_WORLD_RANK', 'OMPI_COMM_WORLD_SIZE'] - - # Check if all these are present - if not all(var in os.environ for var in required_vars): - raise EnvironmentError("MPI environment variables are not set. " - "Ensure you are running the script with an MPI-compatible launcher.") - - # Now safe to read all - os.environ['LOCAL_RANK'] = os.environ['OMPI_COMM_WORLD_LOCAL_RANK'] - os.environ['RANK'] = os.environ['OMPI_COMM_WORLD_RANK'] - os.environ['WORLD_SIZE'] = os.environ['OMPI_COMM_WORLD_SIZE'] - - def get_cmd(self, environment, active_resources): - total_process_count = sum(self.resource_pool.values()) - - launcher_args = split(self.args.launcher_args) - - # If btl_tcp_if_include option is provided through launcher_args, we use it. Otherwise, we add - # `--mca btl_tcp_if_include eth0` option as a default value for compatibility. - btl_tcp_opt = ['--mca', 'btl_tcp_if_include', 'eth0'] - if len(launcher_args) >= 2: - for i in range(len(launcher_args) - 1): - if launcher_args[i] in ['-mca', '--mca'] and launcher_args[i + 1] == 'btl_tcp_if_include': - btl_tcp_opt = [] - break - - mpirun_cmd = [ - 'mpirun', - '-n', - f'{total_process_count}', - '-hostfile', - f'{self.args.hostfile}', - '--mca', - 'btl', - '^openib', - ] + btl_tcp_opt + launcher_args - - export_cmd = [] - for k, v in self.exports.items(): - export_cmd += ['-x', "{}={}".format(k, v)] - - python_exec = [] - if not self.args.no_python: - python_exec = [sys.executable, "-u"] - if self.args.module: - python_exec.append("-m") - - return mpirun_cmd + export_cmd + python_exec + [self.user_script] + self.user_arguments - - -class MPICHRunner(MultiNodeRunner): - - def __init__(self, args, world_info_base64, resource_pool): - super().__init__(args, world_info_base64) - self.resource_pool = resource_pool - - def backend_exists(self): - #TODO: if IB is available we should suggestion mpich - return shutil.which('mpirun') #mpich_info - - @property - def name(self): - return "mpich" - - def validate_args(self): - super().validate_args() - #TODO: Allow for include/exclude at node-level but not gpu-level - if self.args.include != "" or self.args.exclude != "": - raise ValueError(f"{self.name} backend does not support worker include/exclusion") - - if self.args.num_nodes != -1 or self.args.num_gpus != -1: - raise ValueError(f"{self.name} backend does not support limiting num nodes/gpus") - - def get_cmd(self, environment, active_resources): - devices_per_node = self.resource_pool.values() - total_process_count = sum(devices_per_node) - process_per_node = list(devices_per_node)[0] - if not all([n == process_per_node for n in devices_per_node]): - raise ValueError("MPICH requires same number of devices per node") - - mpirun_cmd = [ - 'mpirun', - '-n', - f'{total_process_count}', - '-ppn', - f'{process_per_node}', - ] + split(self.args.launcher_args) - export_cmd = [] - - for k, v in self.exports.items(): - export_cmd += ['-genv', "{}={}".format(k, v)] - - export_cmd += ['-genv', 'MASTER_ADDR', str(self.args.master_addr)] - export_cmd += ['-genv', 'MASTER_PORT', str(self.args.master_port)] - export_cmd += ['-genv', 'WORLD_SIZE', str(total_process_count)] - export_cmd += ['-genv', 'LOCAL_SIZE', str(process_per_node)] - - export_cmd += ['-hosts'] - hosts = "" - for i, host in enumerate(self.resource_pool.keys()): - if i == 0: - hosts = f"{host}" - else: - hosts += f",{host}" - export_cmd += [hosts] - - helper_args = ["--launcher"] + [self.args.launcher] - python_exec = [] - if not self.args.no_python: - python_exec += [sys.executable, "-u"] - if self.args.module: - python_exec.append("-m") - helper_args.append("--module") - else: - helper_args.append("--no_python") - - helper_cmd = str(os.path.dirname(os.path.realpath(__file__))) + '/launcher_helper.py' - helper_cmd = [helper_cmd] + helper_args + [self.user_script] + self.user_arguments - - return mpirun_cmd + export_cmd + python_exec + helper_cmd - - -class IMPIRunner(MultiNodeRunner): - - def __init__(self, args, world_info_base64, resource_pool): - super().__init__(args, world_info_base64) - self.resource_pool = resource_pool - - def backend_exists(self): - #TODO: if IB is available we should suggestion mpich - return shutil.which('mpirun') #mpich_info - - @property - def name(self): - return "impi" - - def validate_args(self): - super().validate_args() - #TODO: Allow for include/exclude at node-level but not gpu-level - if self.args.include != "" or self.args.exclude != "": - raise ValueError(f"{self.name} backend does not support worker include/exclusion") - - if self.args.num_nodes != -1 or self.args.num_gpus != -1: - raise ValueError(f"{self.name} backend does not support limiting num nodes/gpus") - - def get_cmd(self, environment, active_resources): - devices_per_node = self.resource_pool.values() - total_process_count = sum(devices_per_node) - process_per_node = list(devices_per_node)[0] - if not all([n == process_per_node for n in devices_per_node]): - raise ValueError("Intel MPI requires same number of devices per node") - - mpirun_cmd = [ - 'mpirun', - '-ppn', - f'{process_per_node}', - ] + split(self.args.launcher_args) - export_cmd = [] - - for k, v in self.exports.items(): - export_cmd += ['-genv', f'{k}', f'{v}'] - - if self.args.bind_cores_to_rank: - cores_per_rank, _ = get_numactl_cmd(self.args.bind_core_list, process_per_node, 0) - export_cmd += ['-genv', 'OMP_NUM_THREADS', str(cores_per_rank)] - - export_cmd += ['-genv', 'MASTER_ADDR', str(self.args.master_addr)] - export_cmd += ['-genv', 'MASTER_PORT', str(self.args.master_port)] - export_cmd += ['-genv', 'WORLD_SIZE', str(total_process_count)] - export_cmd += ['-genv', 'LOCAL_SIZE', str(process_per_node)] - - # turn off IMPI core binding, use deepspeed's own core binding - export_cmd += ['-genv', 'I_MPI_PIN', '0'] - - export_cmd += ['-hosts'] - hosts = "" - for i, host in enumerate(self.resource_pool.keys()): - if i == 0: - hosts = f"{host}" - else: - hosts += f",{host}" - export_cmd += [hosts] - - per_host_cmd = [] - - for i in range(total_process_count): - local_rank = i % process_per_node - python_exec = [] - if self.args.bind_cores_to_rank: - _, numactl_cmd = get_numactl_cmd(self.args.bind_core_list, process_per_node, local_rank) - python_exec += numactl_cmd - - if not self.args.no_python: - python_exec += [sys.executable, "-u"] - if self.args.module: - python_exec.append("-m") - env_mapping = ['-env', 'RANK', str(i)] - env_mapping += ['-env', 'LOCAL_RANK', str(local_rank)] - if i == 0: - per_host_cmd = ['-n', '1'] + env_mapping + python_exec + [self.user_script] + self.user_arguments - else: - per_host_cmd = per_host_cmd + [':', '-n', '1'] + env_mapping + python_exec + [self.user_script - ] + self.user_arguments - print(mpirun_cmd + export_cmd + per_host_cmd) - return mpirun_cmd + export_cmd + per_host_cmd - - -class SlurmRunner(MultiNodeRunner): - - def __init__(self, args, world_info_base64, resource_pool): - super().__init__(args, world_info_base64) - self.resource_pool = resource_pool - - def backend_exists(self): - return shutil.which('sinfo') - - @property - def name(self): - return 'slurm' - - def get_cmd(self, environment, active_resources): - assert not getattr(self.args, 'detect_nvlink_pairs', - False), "slurm backend does not support remapping visible devices" - total_process_count = sum(self.resource_pool.values()) - srun_cmd = [ - 'srun', - '-n', - f'{total_process_count}', - ] + split(self.args.launcher_args) - - if getattr(self.args, 'slurm_comment', ''): - srun_cmd += ['--comment', self.args.slurm_comment] - - if self.args.include != "": - srun_cmd.append('--include') - srun_cmd.append(f'{self.args.include}') - if self.args.exclude != "": - srun_cmd.append('--exclude') - srun_cmd.append(f'{self.args.exclude}') - if self.args.num_nodes > 0: - srun_cmd.append('--nodes') - srun_cmd.append(f'{self.args.num_nodes}') - if self.args.num_gpus > 0: - srun_cmd.append('--gpus') - srun_cmd.append(f'{self.args.num_gpus}') - - exports = '--export=ALL' - for key, val in self.exports.items(): - exports += f",{key}={val}" - - python_exec = [sys.executable, "-u"] - command = srun_cmd + [exports] + python_exec + [self.user_script] + self.user_arguments - return command - - -class MVAPICHRunner(MultiNodeRunner): - - def __init__(self, args, world_info_base64, resource_pool): - super().__init__(args, world_info_base64) - self.resource_pool = resource_pool - - # Disable the CMA kernel module, not available on Ubuntu systems - self.add_export('MV2_SMP_USE_CMA', '0') - - # If we fail this will output more verbose logging - self.add_export('MV2_DEBUG_SHOW_BACKTRACE', '1') - - # Enabled cuda-aware communication - if get_accelerator().device_name() == 'cuda': - self.add_export('MV2_USE_CUDA', '1') - - # Support deep learning frameworks: http://hidl.cse.ohio-state.edu/userguide/horovod/ - self.add_export('MV2_SUPPORT_DL', '1') - - # Support MPI_THREAD_MULTIPLE - self.add_export('MV2_ENABLE_AFFINITY', '0') - - # Performance tuning flags for allgather - self.add_export('MV2_INTER_ALLGATHER_TUNING', '5') - self.add_export('MV2_CUDA_USE_NAIVE', '0') - - def backend_exists(self): - #TODO: if IB is available we should suggestion mvapich - mpiname_exists = shutil.which('mpiname') - exists = False - if not mpiname_exists: - warnings.warn("mpiname does not exist, mvapich is not installed properly") - else: - results = subprocess.check_output(['mpiname']) - mpiname_results = results.decode('utf-8').strip() - if "MVAPICH2-GDR" in mpiname_results: - exists = True - else: - warnings.warn(f"Expected MVAPICH2-GDR as return for mpiname but received {mpiname_results}") - return exists - - @property - def name(self): - return "mvapich" - - def validate_args(self): - super().validate_args() - #TODO: Allow for include/exclude at node-level but not gpu-level - if self.args.include != "" or self.args.exclude != "": - raise ValueError(f"{self.name} backend does not support worker include/exclusion") - if self.args.num_nodes != -1 or self.args.num_gpus != -1: - raise ValueError(f"{self.name} backend does not support limiting num nodes/gpus") - - def get_cmd(self, environment, active_resources): - devices_per_node = self.resource_pool.values() - total_process_count = sum(devices_per_node) - process_per_node = list(devices_per_node)[0] - if not all([n == process_per_node for n in devices_per_node]): - raise ValueError("mvapich requires same number of devices per node") - - with open(MVAPICH_TMP_HOSTFILE, 'w') as fd: - for host in self.resource_pool.keys(): - fd.write(f'{host}\n') - - mpirun_cmd = [ - 'mpirun', - '-np', - f'{total_process_count}', - '-ppn', - f'{process_per_node}', - '--hostfile', - f'{MVAPICH_TMP_HOSTFILE}', - ] + split(self.args.launcher_args) - - export_cmd = [] - for k, v in self.exports.items(): - export_cmd += ['-env', "{}={}".format(k, v)] - - python_exec = [] - if not self.args.no_python: - python_exec = [sys.executable, "-u"] - if self.args.module: - python_exec.append("-m") - - return mpirun_cmd + export_cmd + python_exec + [self.user_script] + self.user_arguments diff --git a/venv/lib/python3.10/site-packages/deepspeed/launcher/runner.py b/venv/lib/python3.10/site-packages/deepspeed/launcher/runner.py deleted file mode 100644 index 7c792153c8817fb207313c670a38492122a3fe7e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/launcher/runner.py +++ /dev/null @@ -1,636 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -DeepSpeed runner is the main front-end to launching multi-worker -training jobs with DeepSpeed. By default this uses pdsh to parallel -ssh into multiple worker nodes and launch all the necessary processes -per rank for training. -""" - -import os -import re -import sys -import json -import base64 -import argparse -import subprocess -import collections -from copy import deepcopy -import signal -import time -from typing import Tuple, List, Dict -from collections import defaultdict -import shlex - -from .multinode_runner import PDSHRunner, OpenMPIRunner, MVAPICHRunner, SlurmRunner, MPICHRunner, IMPIRunner -from .constants import PDSH_LAUNCHER, OPENMPI_LAUNCHER, MVAPICH_LAUNCHER, SLURM_LAUNCHER, MPICH_LAUNCHER, IMPI_LAUNCHER -from ..constants import TORCH_DISTRIBUTED_DEFAULT_PORT -from ..nebula.constants import NEBULA_EXPORT_ENVS -from ..utils import logger - -from ..autotuning import Autotuner -from deepspeed.accelerator import get_accelerator - -DLTS_HOSTFILE = "/job/hostfile" -EXPORT_ENVS = ['MLFLOW', 'PYTHON', 'MV2', 'UCX'] -EXPORT_ENVS += NEBULA_EXPORT_ENVS -DEEPSPEED_ENVIRONMENT_NAME = os.getenv("DS_ENV_FILE", ".deepspeed_env") -DEEPSPEED_ENVIRONMENT_PATHS = [os.path.expanduser("~"), '.'] -PDSH_MAX_FAN_OUT = 1024 - -# On AISC compute, each node sets environment variables independently, want to prevent -# exporting rank-0 env variables in case of heterogeneous compute. -EXCLUDE_ENVS = {'AISC_JOB_NAME': ['NCCL_IB_HCA', 'UCX_NET_DEVICES']} - - -def parse_args(args=None): - parser = argparse.ArgumentParser(description="DeepSpeed runner to help launch distributed " - "multi-node/multi-gpu training jobs.", - formatter_class=argparse.ArgumentDefaultsHelpFormatter) - - parser.add_argument("-H", - "--hostfile", - type=str, - default=DLTS_HOSTFILE, - help="Hostfile path (in MPI style) that defines the " - "resource pool available to the job (e.g., " - "worker-0 slots=4)") - - parser.add_argument("-i", - "--include", - type=str, - default="", - help='''Specify hardware resources to use during execution. - String format is - NODE_SPEC[@NODE_SPEC ...], - where - NODE_SPEC=NAME[:SLOT[,SLOT ...]]. - If :SLOT is omitted, include all slots on that host. - Example: -i "worker-0@worker-1:0,2" will use all slots - on worker-0 and slots [0, 2] on worker-1. - ''') - - parser.add_argument("-e", - "--exclude", - type=str, - default="", - help='''Specify hardware resources to NOT use during execution. - Mutually exclusive with --include. Resource formatting - is the same as --include. - Example: -e "worker-1:0" will use all available - resources except slot 0 on worker-1. - ''') - - parser.add_argument("--num_nodes", - type=int, - default=-1, - help="Total number of worker nodes to run on, this will use " - "the top N hosts from the given hostfile.") - - parser.add_argument("--min_elastic_nodes", - type=int, - default=-1, - help="Minimum number of nodes to run elastic training on. " - "Default is 1 when elastic training is enabled") - - parser.add_argument("--max_elastic_nodes", - type=int, - default=-1, - help="Maximum number of nodes to run elastic training on. " - "Default is num_nodes when elastic training is enabled") - - parser.add_argument("--num_gpus", - "--num_accelerators", - type=int, - default=-1, - help="Max number of GPUs to use on each node, will use " - "[0:N) GPU ids on each node.") - - parser.add_argument("--master_port", - default=TORCH_DISTRIBUTED_DEFAULT_PORT, - type=int, - help="(optional) Port used by PyTorch distributed for " - "communication during training.") - - parser.add_argument("--master_addr", - default="", - type=str, - help="(optional) IP address of node 0, will be " - "inferred via 'hostname -I' if not specified.") - - parser.add_argument("--node_rank", - default=-1, - type=int, - help="ID of each node in the range [0:N). " - "Only required when --no_ssh is set.") - - parser.add_argument("--launcher", - default=PDSH_LAUNCHER, - type=str, - help="(optional) choose launcher backend for multi-node " - "training. Options currently include PDSH, OpenMPI, MVAPICH, SLURM, MPICH, IMPI.") - - parser.add_argument("--launcher_args", - default="", - type=str, - help="(optional) pass launcher specific arguments as a " - "single quoted argument.") - - parser.add_argument("--module", - action="store_true", - help="Change each process to interpret the launch " - "script as a Python module, executing with the same " - "behavior as 'python -m'.") - - parser.add_argument("--no_python", - action="store_true", - help="Skip prepending the training script with " - "'python' - just execute it directly.") - - parser.add_argument("--no_local_rank", - action="store_true", - help="Do not pass local_rank as an argument when calling " - "the user's training script.") - - parser.add_argument("--no_ssh", - action="store_true", - help="Launch training independently on each node without ssh setup.") - - parser.add_argument("--no_ssh_check", - action="store_true", - help="Do not perform ssh check in multi-node launcher model") - - parser.add_argument("--force_multi", - action="store_true", - help="Force multi-node launcher mode, helps in cases where user " - "wants to launch on single remote node.") - - parser.add_argument("--save_pid", - action="store_true", - help="Save file containing launcher process id (pid) at /tmp/.ds, " - "where is the pid of the first process that invoked `deepspeed`. " - "Useful when launching deepspeed processes programmatically.") - - parser.add_argument("--enable_each_rank_log", - default="None", - type=str, - help="redirect the stdout and stderr from each rank into different log files") - - parser.add_argument("--autotuning", - default="", - choices=["tune", "run"], - type=str, - help="Run DeepSpeed autotuner to discover optimal configuration parameters " - "before running job.") - - parser.add_argument("--elastic_training", - action="store_true", - help="Enable elastic training support in DeepSpeed.") - - parser.add_argument("user_script", type=str, help="User script to launch, followed by any required " - "arguments.") - - parser.add_argument('user_args', nargs=argparse.REMAINDER) - - parser.add_argument("--bind_cores_to_rank", - action="store_true", - help="Bind each rank to different cores of the host") - - parser.add_argument("--bind_core_list", - type=str, - default=None, - help="List of cores to bind to with comma separated list of " - "numbers and range. i.e. 1,3-5,7 => [1,3,4,5,7]. When not " - "specified, all cores on system would be used rank binding") - - parser.add_argument("--ssh_port", type=int, default=None, help="SSH port to use for remote connections") - - parser.add_argument("--venv_script", - type=str, - default=None, - help="Python virtual environment activation script for job.") - - return parser.parse_args(args=args) - - -def fetch_hostfile(hostfile_path): - if not os.path.isfile(hostfile_path): - logger.warning("Unable to find hostfile, will proceed with training " - "with local resources only.") - return None - - # e.g., worker-0 slots=16 - with open(hostfile_path, 'r') as fd: - hostfile_text = fd.readlines() - - return _parse_hostfile(hostfile_text) - - -def _parse_hostfile(hostfile_lines): - # Regex matches one or more non-whitespace characters (\S+) at the start of - # the line, followed by one or more whitespace characters (\s+), followed - # by the string "slots=", followed by one or more digits (\d+). - pattern = r'^(\S+)\s+slots=(\d+)' - - resource_pool = collections.OrderedDict() - - for line in hostfile_lines: - line = line.strip() - match = re.search(pattern, line) - if line.startswith("#") or line == "": - # hostfile comment or empty line, ignore - continue - elif match: - host = match.group(1) - num_slots = int(match.group(2)) - if host in resource_pool: - logger.error(f"Bad hostfile text: {hostfile_lines}") - raise ValueError(f"Hostfile contains multiple entries for {host}, unable to proceed with launching") - resource_pool[host] = num_slots - else: - logger.error(f"Bad hostfile text: {hostfile_lines}") - raise ValueError(f"Hostfile contains a bad entry: {line}, unable to proceed with launching") - - if len(resource_pool) == 0: - logger.error(f"Bad hostfile text: {hostfile_lines}") - raise ValueError("Hostfile is empty or not formatted correctly, unable to proceed with launching.") - - return resource_pool - - -def _stable_remove_duplicates(data): - # Create a new list in the same order as original but with duplicates - # removed, should never be more than ~16 elements so simple is best - new_list = [] - for x in data: - if x not in new_list: - new_list.append(x) - return new_list - - -def parse_node_config(node_config: str) -> Tuple[str, List[int]]: - SLOT_LIST_START = ':' - SLOT_SEP = ',' - - if SLOT_LIST_START not in node_config: - return node_config, [] - - hostname, slots = node_config.split(SLOT_LIST_START) - slots = [int(x) for x in slots.split(SLOT_SEP)] - - return hostname, slots - - -def parse_node_config_list(node_config_list: List[str]) -> Dict[str, List[int]]: - NODE_SEP = '@' - - node_configs = defaultdict(list) - - for node_config in node_config_list.split(NODE_SEP): - hostname, slots = parse_node_config(node_config) - node_configs[hostname] += slots - - return {k: sorted(list(set(v))) for k, v in node_configs.items()} - - -def parse_resource_filter(host_info, include_str="", exclude_str=""): - '''Parse an inclusion or exclusion string and filter a hostfile dictionary. - - String format is NODE_SPEC[@NODE_SPEC ...], where - NODE_SPEC = NAME[:SLOT[,SLOT ...]]. - If :SLOT is omitted, include/exclude all slots on that host. - - Examples: - include_str="worker-0@worker-1:0,2" will use all slots on worker-0 and - slots [0, 2] on worker-1. - exclude_str="worker-1:0" will use all available resources except - slot 0 on worker-1. - ''' - - # Ensure include/exclude are mutually exclusive - if (include_str != "") and (exclude_str != ""): - raise ValueError('include_str and exclude_str are mutually exclusive.') - - # no-op - if (include_str == "") and (exclude_str == ""): - return host_info - - # Either build from scratch or remove items - filtered_hosts = dict() - if include_str: - parse_str = include_str - if exclude_str != "": - filtered_hosts = deepcopy(host_info) - parse_str = exclude_str - - # foreach node in the list - for hostname, slots in parse_node_config_list(parse_str).items(): - # Node can either be alone or node:slot,slot,slot - if len(slots) > 0: - # sanity checks - if hostname not in host_info: - raise ValueError(f"Hostname '{hostname}' not found in hostfile") - for slot in slots: - if slot not in host_info[hostname]: - raise ValueError(f"No slot '{slot}' specified on host '{hostname}'") - - # If include string, build the list from here - if include_str: - filtered_hosts[hostname] = slots - elif exclude_str: - for slot in slots: - logger.info(f'removing {slot} from {hostname}') - filtered_hosts[hostname].remove(slot) - - # User just specified the whole node - else: - # sanity check hostname - if hostname not in host_info: - raise ValueError(f"Hostname '{hostname}' not found in hostfile") - - if include_str: - filtered_hosts[hostname] = host_info[hostname] - elif exclude_str: - filtered_hosts[hostname] = [] - - # Post-processing to remove duplicates and empty nodes - del_keys = [] - for hostname in filtered_hosts: - # Remove duplicates - filtered_hosts[hostname] = _stable_remove_duplicates(filtered_hosts[hostname]) - # Remove empty hosts - if len(filtered_hosts[hostname]) == 0: - del_keys.append(hostname) - for name in del_keys: - del filtered_hosts[name] - - # Lastly, go over filtered_hosts and convert to a OrderedDict() to ensure - # we map ranks to nodes correctly by maintaining host_info ordering. - ordered_hosts = collections.OrderedDict() - for host in host_info: - if host in filtered_hosts: - ordered_hosts[host] = filtered_hosts[host] - - return ordered_hosts - - -def parse_inclusion_exclusion(resource_pool, inclusion, exclusion): - active_resources = collections.OrderedDict() - node_configs = parse_node_config_list(inclusion) - - for hostname, slots in resource_pool.items(): - active_resources[hostname] = node_configs[hostname] if hostname in node_configs else list(range(slots)) - - return parse_resource_filter(active_resources, include_str=inclusion, exclude_str=exclusion) - - -def encode_world_info(world_info): - world_info_json = json.dumps(world_info).encode('utf-8') - world_info_base64 = base64.urlsafe_b64encode(world_info_json).decode('utf-8') - return world_info_base64 - - -def run_autotuning(args, active_resources): - tuner = Autotuner(args, active_resources) - logger.info("[Start] Running autotuning") - - tuner.tune() - tuner.print_tuning_results() - - logger.info("[End] Running autotuning") - tuner.write_optimal_config() - - if args.autotuning == "run": - tuner.run_after_tuning() - - -def parse_num_nodes(str_num_nodes: str, elastic_training: bool): - node_list = str_num_nodes.split(":") - - if len(node_list) == 1: - min_nodes, max_nodes = int(node_list[0]), -1 - elif len(node_list) == 2 and elastic_training: - min_nodes, max_nodes = int(node_list[0]), int(node_list[1]) - elif len(node_list) == 2 and not elastic_training: - raise RuntimeError("MIN:MAX format is only supported in elastic training") - else: - raise RuntimeError("num_nodes {} is not in MIN:MAX format".format(str_num_nodes)) - - return min_nodes, max_nodes - - -def main(args=None): - args = parse_args(args) - - if args.elastic_training: - assert args.master_addr != "", "Master Addr is required when elastic training is enabled" - - resource_pool = fetch_hostfile(args.hostfile) - - # respect VISIBLE_DEVICES for a single node and no explicit resource filters - visible_devices_env = get_accelerator().visible_devices_envs()[0] - visible_devices = os.environ.get(visible_devices_env, "") - if not resource_pool and len(visible_devices): - detected_str = f"Detected VISIBLE_DEVICES={visible_devices}" - if len(args.include) or len(args.exclude) or args.num_nodes > 1 or args.num_gpus > 0: - print( - f"{detected_str} but ignoring it because one or several of --include/--exclude/--num_gpus/--num_nodes cl args were used. If you want to use CUDA_VISIBLE_DEVICES don't pass any of these arguments to deepspeed." - ) - else: - args.include = f"localhost:{visible_devices}" - print(f"{detected_str}: setting --include={args.include}") - del os.environ[visible_devices_env] - - if args.num_nodes >= 0 or args.num_gpus >= 0: - if args.include != "" or args.exclude != "": - raise ValueError("Cannot specify num_nodes/gpus with include/exclude") - - multi_node_exec = True - if not resource_pool: - resource_pool = {} - device_count = get_accelerator().device_count() - if device_count == 0: - raise RuntimeError("Unable to proceed, no GPU resources available") - resource_pool['localhost'] = device_count - args.master_addr = "127.0.0.1" - multi_node_exec = False - - if not multi_node_exec and args.num_nodes > 1: - raise ValueError("Num nodes is >1 but no extra nodes available via hostfile") - - active_resources = parse_inclusion_exclusion(resource_pool, args.include, args.exclude) - env = os.environ.copy() - - # validate that passwordless-ssh is workly properly with this hostfile - if multi_node_exec and not args.no_ssh_check and not args.no_ssh: - first_host = list(active_resources.keys())[0] - try: - ssh_check_cmd = ("ssh -o PasswordAuthentication=no " + - (f"-p {args.ssh_port} " if args.ssh_port is not None else "") + f"{first_host} hostname") - safe_ssh_cmd = shlex.split(ssh_check_cmd) - subprocess.check_call(safe_ssh_cmd, stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL) - except subprocess.CalledProcessError: - raise RuntimeError( - f"Using hostfile at {args.hostfile} but host={first_host} was not reachable via ssh. If you are running with a single node please remove {args.hostfile} or setup passwordless ssh." - ) - - if not args.master_addr: - assert multi_node_exec - first_host = list(active_resources.keys())[0] - ssh_check_cmd = "ssh " - if args.ssh_port is not None: - ssh_check_cmd += f" -p {args.ssh_port}" - ssh_check_cmd += f" {first_host} hostname -I" - hostname_cmd = shlex.split(ssh_check_cmd) - try: - result = subprocess.check_output(hostname_cmd) - except subprocess.CalledProcessError as err: - logger.error( - "Unable to detect suitable master address via 'hostname -I', please manually specify one via --master_addr" - ) - raise err - args.master_addr = result.decode('utf-8').split()[0] - if not args.master_addr: - raise RuntimeError( - "Unable to detect suitable master address via `hostname -I`, please manually specify one via --master_addr" - ) - logger.info(f"Using IP address of {args.master_addr} for node {first_host}") - - if args.autotuning != "": - run_autotuning(args, active_resources) - return - - if args.num_nodes > 0: - updated_active_resources = collections.OrderedDict() - for count, hostname in enumerate(active_resources.keys()): - if args.num_nodes == count: - break - updated_active_resources[hostname] = active_resources[hostname] - active_resources = updated_active_resources - - if args.num_gpus > 0: - updated_active_resources = collections.OrderedDict() - for hostname in active_resources.keys(): - updated_active_resources[hostname] = list(range(args.num_gpus)) - active_resources = updated_active_resources - - if args.elastic_training: - assert not args.no_local_rank, "--no_local_rank argument is not supported in Elastic training" - - if args.no_ssh: - assert (0 <= args.node_rank < - len(active_resources)), "Launching training without ssh, but --node_rank is not set correctly." - - # encode world info as base64 to make it easier to pass via command line - world_info_base64 = encode_world_info(active_resources) - - multi_node_exec = (args.force_multi or len(active_resources) > 1) and not args.no_ssh - - if not multi_node_exec: - deepspeed_launch = [ - sys.executable, "-u", "-m", "deepspeed.launcher.launch", f"--world_info={world_info_base64}", - f"--master_addr={args.master_addr}", f"--master_port={args.master_port}" - ] - if args.no_ssh: - deepspeed_launch.append(f"--node_rank={args.node_rank}") - if args.no_python: - deepspeed_launch.append("--no_python") - if args.module: - deepspeed_launch.append("--module") - if args.no_local_rank: - deepspeed_launch.append("--no_local_rank") - if args.save_pid: - deepspeed_launch += ["--save_pid", f"{os.getpid()}"] - if args.enable_each_rank_log: - deepspeed_launch.append(f"--enable_each_rank_log={args.enable_each_rank_log}") - if args.elastic_training: - deepspeed_launch.append("--enable_elastic_training") - deepspeed_launch.append(f"--max_elastic_nodes={args.max_elastic_nodes}") - deepspeed_launch.append(f"--min_elastic_nodes={args.min_elastic_nodes}") - if args.bind_cores_to_rank: - deepspeed_launch.append("--bind_cores_to_rank") - if args.bind_core_list is not None: - deepspeed_launch.append(f"--bind_core_list={args.bind_core_list}") - cmd = deepspeed_launch + [args.user_script] + args.user_args - else: - args.launcher = args.launcher.lower() - if args.launcher == PDSH_LAUNCHER: - runner = PDSHRunner(args, world_info_base64) - elif args.launcher == OPENMPI_LAUNCHER: - runner = OpenMPIRunner(args, world_info_base64, resource_pool) - elif args.launcher == MPICH_LAUNCHER: - runner = MPICHRunner(args, world_info_base64, resource_pool) - elif args.launcher == IMPI_LAUNCHER: - runner = IMPIRunner(args, world_info_base64, resource_pool) - elif args.launcher == MVAPICH_LAUNCHER: - runner = MVAPICHRunner(args, world_info_base64, resource_pool) - elif args.launcher == SLURM_LAUNCHER: - runner = SlurmRunner(args, world_info_base64, resource_pool) - else: - raise NotImplementedError(f"Unknown launcher {args.launcher}") - - if not runner.backend_exists(): - raise RuntimeError(f"launcher '{args.launcher}' not installed.") - - curr_path = os.path.abspath('.') - if 'PYTHONPATH' in env: - env['PYTHONPATH'] = curr_path + ":" + env['PYTHONPATH'] - else: - env['PYTHONPATH'] = curr_path - - excluded_vars = [] - for exclude_key, var_list in EXCLUDE_ENVS.items(): - if exclude_key in env.keys(): - # key exists in launcher env -> var list should be used - excluded_vars += var_list - - # load envs from accelerator - exports = EXPORT_ENVS + get_accelerator().export_envs() - for var in env.keys(): - if any([var.startswith(name) for name in exports]): - if not any([var == name for name in excluded_vars]): - runner.add_export(var, env[var]) - - for environ_path in DEEPSPEED_ENVIRONMENT_PATHS: - environ_file = os.path.join(environ_path, DEEPSPEED_ENVIRONMENT_NAME) - if os.path.isfile(environ_file): - logger.info(f"deepspeed_env file = {environ_file}") - with open(environ_file, 'r') as fd: - for var in fd.readlines(): - key, val = var.split('=', maxsplit=1) - runner.add_export(key, val) - - if args.launcher == PDSH_LAUNCHER: - cmd, kill_cmd, env = runner.get_cmd(env, active_resources) - else: - cmd = runner.get_cmd(env, active_resources) - - logger.info(f"cmd = {' '.join(cmd)}") - result = subprocess.Popen(cmd, env=env) - - def sigkill_handler(signum, frame): - result.send_signal(signal.SIGINT) - time.sleep(0.1) - result.send_signal(signal.SIGTERM) - result_kill = subprocess.Popen(kill_cmd, env=env) - result_kill.wait() - time.sleep(1) - sys.exit(1) - - if args.launcher == PDSH_LAUNCHER and multi_node_exec: - signal.signal(signal.SIGINT, sigkill_handler) - signal.signal(signal.SIGTERM, sigkill_handler) - - result.wait() - - # In case of failure must propagate the error-condition back to the caller (usually shell). The - # actual error and traceback should have been printed in the subprocess, so in order to avoid - # unnecessary noise we just quietly exit here with the same code as the subprocess - if result.returncode > 0: - sys.exit(result.returncode) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/deepspeed/linear/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/linear/__init__.py deleted file mode 100644 index 9931a95a0a40c449b40b9152a5b8419f66d60a10..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/linear/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .optimized_linear import OptimizedLinear -from .config import LoRAConfig, QuantizationConfig -from .context_manager import Init, init_lora diff --git a/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index a74c627359ba58ec6af8c28349b2dc07a7c6d8a9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/config.cpython-310.pyc deleted file mode 100644 index 7dcfa866ad2e357f7c5150858c5633008aa052c5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/context_manager.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/context_manager.cpython-310.pyc deleted file mode 100644 index 7227aee3792b7d822dd8f6f3d2783b87813616fc..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/context_manager.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/optimized_linear.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/optimized_linear.cpython-310.pyc deleted file mode 100644 index 14d7c9f0d3f7594adb6484f7f5da6f675b53bc2f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/optimized_linear.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/quantization.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/quantization.cpython-310.pyc deleted file mode 100644 index c397826fc2d759caf7299e99fd89d670fcbad7b3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/linear/__pycache__/quantization.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/linear/config.py b/venv/lib/python3.10/site-packages/deepspeed/linear/config.py deleted file mode 100644 index 1459704a32c5f36a827d4427de96942e3b6fa06b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/linear/config.py +++ /dev/null @@ -1,54 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from dataclasses import dataclass, field -from typing import List - -import torch - - -@dataclass -class LoRAConfig: - """ - Configuration settings for LoRAOptimizedLinear. - - Attributes: - lora_r (int): LoRA attention dimension, also known as the rank. Defaults is 64. - lora_alpha (float): LoRA scaling factor, default is 16. - base_weight_sharding (int): The degree to which the base weights are sharded, - should typically be set to the data-parallel world size to maximize the memory - reduction benefits. Defaults to 1, which means this feature is disabled. - offload (bool): offload frozen parameters to cpu when not in use - offload_ratio (float): ratio of parameters to offload to cpu when not in use - delay_lora_init (bool): initialize lora parameters at time of model init or allow manual init later - target_mods (str): target module names to apply LoRA to, defaults to llama-3.1 arch - """ - lora_r: int = 64 - lora_alpha: float = 16. - base_weight_sharding: int = 1 - offload: bool = False - offload_ratio: float = 0.0 - delay_lora_init: bool = False - target_mods: List[str] = field( - default_factory=lambda: ['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj']) - - -@dataclass -class QuantizationConfig: - """ - Configuration settings for quantization for LoRAOptimizedLinear, QuantizedLinear, - and QuantizedParameter - - Attributes: - q_bits (int): The number of bits used for quantization. Default is 8. - mantissa_bits (int): The number of bits reserved for the mantissa in fixed-point quantization. Default is 3. - group_size (int): The number of elements used for quantization. Default is 512. - q_dtype (torch.dtype): The data type to quantize to. Default is uint8. (in CUDA, buffers are allocated as - uint8, but inside the kernels the quantization is done to fp8) - """ - q_bits: int = 8 - mantissa_bits: int = 3 - group_size: int = 512 - q_dtype: torch.dtype = torch.uint8 diff --git a/venv/lib/python3.10/site-packages/deepspeed/linear/context_manager.py b/venv/lib/python3.10/site-packages/deepspeed/linear/context_manager.py deleted file mode 100644 index 204fa0fe9c1da3ed2b1c8ccad2828257da536d57..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/linear/context_manager.py +++ /dev/null @@ -1,90 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .optimized_linear import LoRAOptimizedLinear, OptimizedLinear - -import torch - -try: - import transformers -except ImportError: - transformers = None - - -def init_lora(model): - model.requires_grad_(False) - for m in model.modules(): - if isinstance(m, LoRAOptimizedLinear): - m.init_lora() - - -class Init(object): - """ - Init context wrapper similar in style to zero.Init. Allows for injecting OptimizedLinear during model - construction which will shard base weights and reduce overall memory usage during model init. Primarily - useful when initializing a model via transformers.AutoModelForCausalLM. - - Example usage: - lora_config = deepspeed.linear.LoRAConfig(..) - quant_config = deepspeed.linear.QuantizationConfig(..) - with deepspeed.linear.Init(lora_config=lora_config, quant_config=quant_config): - model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3.1-405B") - - """ - - def __init__(self, lora_config=None, quant_config=None): - self._orig_nn_linear = torch.nn.Linear - self._orig_causallm_pretrained = None - if transformers != None: - self._orig_causallm_pretrained = transformers.AutoModelForCausalLM.from_pretrained - self._orig_causallm_config = transformers.AutoModelForCausalLM.from_config - self.lora_config = lora_config - self.quant_config = quant_config - self._post_init_complete = False - - def __enter__(self): - - class OptLinearWrapper: - _orig_nn_linear = self._orig_nn_linear - _lora_config = self.lora_config - _quant_config = self.quant_config - - def __new__(self, *args, **kwargs): - self._lora_config.delay_lora_init = True - kwargs['lora_config'] = self._lora_config - kwargs['quantization_config'] = self._quant_config - kwargs['linear_cls'] = self._orig_nn_linear - return OptimizedLinear(*args, **kwargs) - - def _model_init(model): - if self.lora_config != None: - init_lora(model) - self._post_init_complete = True - return model - - # ensures non-lora params are frozen and lora weights are initialized - def from_pretrained(*args, **kwargs): - model = self._orig_causallm_pretrained(*args, **kwargs) - return _model_init(model) - - def from_config(*args, **kwargs): - model = self._orig_causallm_config(*args, **kwargs) - return _model_init(model) - - torch.nn.Linear = OptLinearWrapper - if transformers != None: - transformers.AutoModelForCausalLM.from_pretrained = from_pretrained - transformers.AutoModelForCausalLM.from_config = from_config - - def __exit__(self, *args, **kwargs): - torch.nn.Linear = self._orig_nn_linear - if not self._post_init_complete: - print('WARNING: For some reason LoRA modules are not initialized, this is usually done automatically ' - 'if using transformers via (AutoModelForCausalLM from_pretrained/from_config). ' - 'You must call `init_lora` on each module in order to use DeepSpeed LoRA, otherwise ' - 'you will error out during runtime.') - else: - transformers.AutoModelForCausalLM.from_pretrained = self._orig_causallm_pretrained - transformers.AutoModelForCausalLM.from_config = self._orig_causallm_config diff --git a/venv/lib/python3.10/site-packages/deepspeed/linear/optimized_linear.py b/venv/lib/python3.10/site-packages/deepspeed/linear/optimized_linear.py deleted file mode 100644 index 3720196aa2554b27bd235522a73447597ef7f97b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/linear/optimized_linear.py +++ /dev/null @@ -1,222 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import math -import torch.nn as nn -import torch.nn.functional as F -from dataclasses import is_dataclass -from deepspeed.accelerator import get_accelerator -import deepspeed.comm as dist - -from .config import LoRAConfig, QuantizationConfig -from .quantization import QuantizedParameter, QuantizedLinear - - -class OptimizedLinear(nn.Module): - """ - Optimized version of nn.Linear that adds features such as: - * LoRA w. base weight sharding - * FP [6,8,12] quantization - - Arguments: - input_dim: Required: size of each input sample - output_dim: Required: size of each output sample - bias: Optional: If set to False, the layer will not learn an additive bias. Default: False - lora_config: Optional: LoRAConfig defining lora features and base-weight-sharding degree - quantization_config: Optional: QuantizationConfig defining quantization features - dtype: Optional: parameter dtype, only supports bfloat16 currently - - Returns: - Returns a new nn.Module depending on the input config. Either native - torch.nn.Linear, QuantizedLinear, or the full-featured DSOptimizedLinear. - """ - - def __new__(self, - input_dim: int, - output_dim: int, - bias: bool = False, - lora_config: LoRAConfig = None, - quantization_config: QuantizationConfig = None, - device=None, - dtype=torch.bfloat16, - linear_cls=nn.Linear): - - if quantization_config is not None and not is_dataclass(quantization_config): - raise ValueError(f"Expecting QuantizationConfig but received {type(quantization_config)}") - if lora_config is not None and not is_dataclass(lora_config): - raise ValueError(f"Expecting LoRAConfig but received {type(lora_config)}") - if lora_config is None and quantization_config is None: - # Everything disabled, fall back to normal nn.Linear - self = linear_cls(input_dim, output_dim, bias=bias, dtype=dtype, device=device) - - elif lora_config: - # lora enabled, quantization may or may not be - self = LoRAOptimizedLinear(input_dim=input_dim, - output_dim=output_dim, - bias=bias, - lora_config=lora_config, - quantization_config=quantization_config, - dtype=dtype, - device=device, - linear_cls=linear_cls) - - elif quantization_config: - # only quantization enabled, no lora - self = QuantizedLinear(input_dim=input_dim, - output_dim=output_dim, - bias=bias, - quantization_config=quantization_config, - dtype=dtype) - return self - - -class LoRAOptimizedLinear(nn.Module): - - def __init__(self, - input_dim: int, - output_dim: int, - bias: bool = False, - lora_config: LoRAConfig = None, - quantization_config: QuantizationConfig = None, - device=None, - dtype=torch.bfloat16, - linear_cls=nn.Linear): - super().__init__() - self.input_dim = input_dim - self.output_dim = output_dim - self.bias = bias - self.lora_config = lora_config - self.quantization_config = quantization_config - self.device = get_accelerator().current_device_name() if device is None else device - self.linear_cls = linear_cls - self.dtype = dtype - assert self.lora_config is not None, "DSOptimizedLinear requires a LoRA config" - assert not self.bias, "bias=True is not supported by LoRAOptimizedLinear" - self.zero_shards = self.lora_config.base_weight_sharding - self.sharded_weight_size = int(float(self.input_dim) // self.zero_shards) - if self.zero_shards > 1: - assert self.zero_shards == dist.get_world_size( - ), "base weight sharding is only supported across world size" - w = torch.nn.Parameter(torch.empty(self.output_dim * self.sharded_weight_size, dtype=dtype), - requires_grad=False) - else: - w = torch.nn.Parameter(torch.empty((self.output_dim, self.input_dim), dtype=dtype), requires_grad=False) - torch.nn.init.xavier_uniform_(w.reshape(self.sharded_weight_size, self.output_dim)) - - if self.quantization_config is not None: - assert dtype == torch.bfloat16, "only bfloat16 is supported when using quantization" - self.weight = QuantizedParameter(w, quantization_config=quantization_config) - else: - self.weight = w - - self.disabled = False - self._initialized = False - if not self.lora_config.delay_lora_init: - self.init_lora() - - def disable(self): - self.disabled = True - self.weight = torch.nn.Parameter(torch.empty((self.output_dim, self.input_dim), dtype=self.dtype), - requires_grad=False) - - def init_lora(self): - if self.disabled: - return - - if self.quantization_config is not None: - # ensure quant-param wasn't stripped, in some cases transformers will do this during model init - if not isinstance(self.weight, QuantizedParameter): - self.weight = QuantizedParameter(self.weight, quantization_config=self.quantization_config) - - self._initialized = True - self.weight.requires_grad = False - - # Mark base weight to prevent broadcast and ensure proper offload behavior - self.weight.ds_optim_param = True - - self.lora_scaling_factor = self.lora_config.lora_alpha / self.lora_config.lora_r - - # Keeping lora weights in bf16 precision for ease of training. - self.lora_weight_1 = self.linear_cls(self.input_dim, - self.lora_config.lora_r, - bias=self.bias, - device=self.device, - dtype=self.dtype) - self.lora_weight_2 = self.linear_cls(self.lora_config.lora_r, - self.output_dim, - bias=self.bias, - device=self.device, - dtype=self.dtype) - - # initialize "A" with kaiming uniform and "B" with zeros following this - # https://github.com/huggingface/peft/blob/62122b5add8d6892f70c82eaef2147a6ba33b90b/src/peft/tuners/lora/layer.py#L155 - nn.init.kaiming_uniform_(self.lora_weight_1.weight, a=math.sqrt(5)) - nn.init.zeros_(self.lora_weight_2.weight) - self.lora_weight_1.weight.requires_grad = True - self.lora_weight_2.weight.requires_grad = True - - def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, - error_msgs): - if not any([target in prefix for target in self.lora_config.target_mods]): - # module does not match any target_mods, we must revert to normal nn.Linear via disable - self.disable() - return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, - unexpected_keys, error_msgs) - - if self.zero_shards > 1: - if not dist.is_initialized(): - raise RuntimeError( - "attempting to use optimized linear base weight sharding but torch-distributed is not initialized, please init first." - ) - rank = dist.get_rank() - shape_local = self.output_dim * self.sharded_weight_size - base_weight_name = f"{prefix}weight" - incoming_param = state_dict[base_weight_name] - state_dict[base_weight_name] = incoming_param.flatten().narrow(0, rank * shape_local, shape_local) - - return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, - error_msgs) - - def full_weight(self): - base_weight = self.weight - if getattr(base_weight, 'ds_offload', False): - # move to gpu so we can dequant and all-gather - assert base_weight.device == torch.device('cpu'), \ - f"expected base weight on cpu but found {base_weight.device}" - base_weight.offload(revert=True) - local_weight = base_weight.dequantized() if isinstance(base_weight, QuantizedParameter) else base_weight - base_weight.offload() - else: - local_weight = base_weight.dequantized() if isinstance(base_weight, QuantizedParameter) else base_weight - - tensor_out = torch.empty(self.output_dim * self.input_dim, - dtype=local_weight.dtype, - device=local_weight.device) - dist.all_gather_into_tensor(tensor_out, local_weight) - return tensor_out.reshape(self.output_dim, self.input_dim) - - def linear_without_F_linear(self, input, weight): - output = torch.mm(input.reshape(-1, input.shape[-1]), weight) - output = output.view(*input.shape[:-1], weight.shape[1]) - return output - - def forward(self, input_tensor): - if self.disabled: - return F.linear(input_tensor, self.weight) - assert self._initialized, "init_lora was never called, please initialize before proceeding" - - # Gather the sharded base weight - if self.zero_shards > 1: - with torch.no_grad(): - base_weight = self.full_weight() - elif self.quantization_config: - base_weight = self.weight.dequantized() - else: - base_weight = self.weight - - base_weight_output = F.linear(input_tensor, base_weight) - lora_output = self.lora_weight_2(self.lora_weight_1(input_tensor)) - return base_weight_output + self.lora_scaling_factor * lora_output diff --git a/venv/lib/python3.10/site-packages/deepspeed/linear/quantization.py b/venv/lib/python3.10/site-packages/deepspeed/linear/quantization.py deleted file mode 100644 index 2023601be281734b68c434acda1f652191f7db8d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/linear/quantization.py +++ /dev/null @@ -1,147 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import copy -import torch -import torch.nn as nn -import torch.nn.functional as F - -from typing import Optional - -from deepspeed.accelerator import get_accelerator -from deepspeed.ops.fp_quantizer import Quantizer, FP_Quantize -from .config import QuantizationConfig - - -class QuantizedParameter(nn.Parameter): - """ - Quantized parameter class that implements weight quantization. Weights - are stored in quantized form on GPUs, and can be dequantized on-the-fly when - needed by the model. The weights are actually quantized during any `.to(device)`. - - Arguments: - data (Tensor): parameter tensor. - requires_grad (bool, optional): if the parameter requires gradient. Defaults - to False and is not supported to be True. Argument provided only for interface - compatibility with torch.nn.Parameter. - quantization_config (QuantizationConfig, optional): - quantizer (Quantizer, optional): Defaults to FP_Quantize but can be any quantizer - that implements deepspeed.ops.fp_quantizer.Quantizer. This argument is also - required since the quantizer is stashed in the Parameter itself, some models - may clone the Parameter by passing an attribute __dict__. For an example, see - tests/unit/linear/test_quant_param.py::TestQuantParam::test_hf_clone - """ - - def __new__( - cls, - data: Optional[torch.Tensor] = None, - requires_grad: bool = False, # quantized weights must be frozen - quantization_config: QuantizationConfig = None, - quantizer: Quantizer = None, - ): - if requires_grad: - raise ValueError(f"requires_grad=True is not supported with QuantizedParameter") - if data is None: - data = torch.empty(0) - self = torch.Tensor._make_subclass(cls, data, requires_grad) - self.quantization_config = QuantizationConfig() if quantization_config is None else quantization_config - if quantizer is not None: - self.quantizer = quantizer - else: - # if FPQuantizerBuilder is not compatible in this env this init will fail - self.quantizer = FP_Quantize(quantization_config=self.quantization_config) - self._ensure_quantized(self) - return self - - def _ensure_quantized(self, tensor: torch.Tensor): - # If the tensor is on the accelerator and is not quantized, then quantize it in-place. - if get_accelerator().on_accelerator(tensor) and tensor.dtype != self.quantization_config.q_dtype: - with get_accelerator().stream(get_accelerator().current_stream(tensor.device)): - tensor.data = self.quantizer.quantize(tensor.data, - q_bits=self.quantization_config.q_bits, - q_mantisa_bits=self.quantization_config.mantissa_bits) - assert tensor.dtype == self.quantization_config.q_dtype - - def dequantized(self) -> torch.Tensor: - """ - Return a tensor containing the dequantized weights of this parameter. - """ - if get_accelerator().on_accelerator(self.data) and self.data.dtype == self.quantization_config.q_dtype: - with get_accelerator().stream(get_accelerator().current_stream(self.data.device)): - return self.quantizer.dequantize(self.data, - q_bits=self.quantization_config.q_bits, - q_mantisa_bits=self.quantization_config.mantissa_bits) - return self.data - - def offload(self, revert=False): - if getattr(self, 'ds_offload', False): - if revert: - self.data = self.to(get_accelerator().current_device_name()) - else: - self.data = self.to('cpu') - - def __getstate__(self): - state = self.__dict__ - state["data"] = self.data - state["quantization_config"] = self.quantization_config - state["requires_grad"] = self.requires_grad - return state - - def __setstate__(self, state): - self.quantizer = state["quantizer"] - self.quantization_config = state["quantization_config"] - self.data = state["data"] - self.requires_grad = state["requires_grad"] - - def __deepcopy__(self, memo): - new_instance = type(self).__new__(type(self)) - state = self.__getstate__() - new_instance.__setstate__(state) - new_instance.quantizer = copy.deepcopy(state["quantizer"]) - new_instance.quantization_config = copy.deepcopy(state["quantization_config"]) - new_instance.data = copy.deepcopy(state["data"]) - return new_instance - - def __copy__(self): - new_instance = type(self).__new__(type(self)) - state = self.__getstate__() - new_instance.__setstate__(state) - return new_instance - - def cuda(self, device=None, non_blocking=False): - device = "cuda" if device is None else device - self.quantizer.to(device, non_blocking=non_blocking) - return self.to(device, non_blocking=non_blocking) - - def to(self, *args, **kwargs): - """ - Move the parameter to the given device. Then, if the device is a cuda device, - quantize it. - """ - tensor = super().to(*args, **kwargs) - self.quantizer.to(*args, **kwargs) - self._ensure_quantized(tensor) - return tensor - - -class QuantizedLinear(nn.Linear): - """ - Linear layer that implements weight quantization. Parameters - are stored via `QuantizedParameter` and are dequantized on-the-fly during any - forward pass. - """ - - def __init__(self, - input_dim: int, - output_dim: int, - bias: bool = False, - quantization_config: QuantizationConfig = None, - dtype=torch.bfloat16): - super().__init__(input_dim, output_dim, bias=bias, dtype=dtype) - assert dtype == torch.bfloat16, "currently only supports bfloat16 dtype" - self.weight = QuantizedParameter(self.weight.data, quantization_config=quantization_config) - - def forward(self, input: torch.Tensor) -> torch.Tensor: - return F.linear(input, self.weight.dequantized(), self.bias) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/__init__.py deleted file mode 100644 index ad95c58f76090188b5d9c73e1e3df4eb2b2c678f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .transformers.ds_transformer import DeepSpeedTransformerInference -from .transformers.clip_encoder import DSClipEncoder diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 2d649bdff2998d339af87783ae03549774593d51..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 9137eb29d3876ddab0770673a08d0641a4069d90..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/__pycache__/unet.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/__pycache__/unet.cpython-310.pyc deleted file mode 100644 index b4ba9610b0773a9e8a14f4f4aea6e5fd31ae6b46..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/__pycache__/unet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/__pycache__/vae.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/__pycache__/vae.cpython-310.pyc deleted file mode 100644 index 87dac9308756736c195f7e8aa77cd8ef4c271674..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/__pycache__/vae.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/unet.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/unet.py deleted file mode 100644 index 8d5ddd95437a42de5cb8552c529d20e422c792f2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/unet.py +++ /dev/null @@ -1,81 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed.accelerator import get_accelerator -from ..features.cuda_graph import CUDAGraph - - -class DSUNet(CUDAGraph, torch.nn.Module): - - def __init__(self, unet, enable_cuda_graph=True): - super().__init__(enable_cuda_graph=enable_cuda_graph) - self.unet = unet - # SD pipeline accesses this attribute - self.in_channels = unet.in_channels - self.device = self.unet.device - self.dtype = self.unet.dtype - self.config = self.unet.config - self.fwd_count = 0 - self.unet.requires_grad_(requires_grad=False) - self.unet.to(memory_format=torch.channels_last) - self.cuda_graph_created = False - - def _graph_replay(self, *inputs, **kwargs): - for i in range(len(inputs)): - if torch.is_tensor(inputs[i]): - self.static_inputs[i].copy_(inputs[i]) - for k in kwargs: - if torch.is_tensor(kwargs[k]): - self.static_kwargs[k].copy_(kwargs[k]) - get_accelerator().replay_graph(self._cuda_graphs) - return self.static_output - - def forward(self, *inputs, **kwargs): - if self.enable_cuda_graph: - if self.cuda_graph_created: - outputs = self._graph_replay(*inputs, **kwargs) - else: - self._create_cuda_graph(*inputs, **kwargs) - outputs = self._graph_replay(*inputs, **kwargs) - return outputs - else: - return self._forward(*inputs, **kwargs) - - def _create_cuda_graph(self, *inputs, **kwargs): - # warmup to create the workspace and cublas handle - cuda_stream = torch.cuda.Stream() - cuda_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(cuda_stream): - for i in range(3): - ret = self._forward(*inputs, **kwargs) - torch.cuda.current_stream().wait_stream(cuda_stream) - - # create cuda_graph and assign static_inputs and static_outputs - self._cuda_graphs = get_accelerator().create_graph() - self.static_inputs = inputs - self.static_kwargs = kwargs - - with get_accelerator().capture_to_graph(self._cuda_graphs): - self.static_output = self._forward(*self.static_inputs, **self.static_kwargs) - - self.cuda_graph_created = True - - def _forward(self, - sample, - timestamp, - encoder_hidden_states, - return_dict=True, - cross_attention_kwargs=None, - timestep_cond=None, - added_cond_kwargs=None): - if cross_attention_kwargs: - return self.unet(sample, - timestamp, - encoder_hidden_states, - return_dict, - cross_attention_kwargs=cross_attention_kwargs) - else: - return self.unet(sample, timestamp, encoder_hidden_states, return_dict) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/vae.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/vae.py deleted file mode 100644 index ce50ade647a8b85c1858e9d47e1f6a1133379467..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/diffusers/vae.py +++ /dev/null @@ -1,151 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed.accelerator import get_accelerator -from ..features.cuda_graph import CUDAGraph - - -class DSVAE(CUDAGraph, torch.nn.Module): - - def __init__(self, vae, enable_cuda_graph=True): - super().__init__(enable_cuda_graph=enable_cuda_graph) - self.vae = vae - self.config = vae.config - self.device = self.vae.device - self.dtype = self.vae.dtype - self.vae.requires_grad_(requires_grad=False) - self.decoder_cuda_graph_created = False - self.encoder_cuda_graph_created = False - self.all_cuda_graph_created = False - - def _graph_replay_decoder(self, *inputs, **kwargs): - for i in range(len(inputs)): - if torch.is_tensor(inputs[i]): - self.static_decoder_inputs[i].copy_(inputs[i]) - for k in kwargs: - if torch.is_tensor(kwargs[k]): - self.static_decoder_kwargs[k].copy_(kwargs[k]) - get_accelerator().replay_graph(self._decoder_cuda_graph) - return self.static_decoder_output - - def _decode(self, x, return_dict=True, generator=None): - return self.vae.decode(x, return_dict=return_dict) - - def _create_cuda_graph_decoder(self, *inputs, **kwargs): - # warmup to create the workspace and cublas handle - cuda_stream = torch.cuda.Stream() - cuda_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(cuda_stream): - for i in range(3): - ret = self._decode(*inputs, **kwargs) - torch.cuda.current_stream().wait_stream(cuda_stream) - - # create cuda_graph and assign static_inputs and static_outputs - self._decoder_cuda_graph = get_accelerator().create_graph() - self.static_decoder_inputs = inputs - self.static_decoder_kwargs = kwargs - - with get_accelerator().capture_to_graph(self._decoder_cuda_graph): - self.static_decoder_output = self._decode(*self.static_decoder_inputs, **self.static_decoder_kwargs) - - self.decoder_cuda_graph_created = True - - def decode(self, *inputs, **kwargs): - if self.enable_cuda_graph: - if self.decoder_cuda_graph_created: - outputs = self._graph_replay_decoder(*inputs, **kwargs) - else: - self._create_cuda_graph_decoder(*inputs, **kwargs) - outputs = self._graph_replay_decoder(*inputs, **kwargs) - return outputs - else: - return self._decode(*inputs, **kwargs) - - def _graph_replay_encoder(self, *inputs, **kwargs): - for i in range(len(inputs)): - if torch.is_tensor(inputs[i]): - self.static_encoder_inputs[i].copy_(inputs[i]) - for k in kwargs: - if torch.is_tensor(kwargs[k]): - self.static_encoder_kwargs[k].copy_(kwargs[k]) - get_accelerator().replay_graph(self._encoder_cuda_graph) - return self.static_encoder_output - - def _encode(self, x, return_dict=True): - return self.vae.encode(x, return_dict=return_dict) - - def _create_cuda_graph_encoder(self, *inputs, **kwargs): - # warmup to create the workspace and cublas handle - cuda_stream = torch.cuda.Stream() - cuda_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(cuda_stream): - for i in range(3): - ret = self._encode(*inputs, **kwargs) - torch.cuda.current_stream().wait_stream(cuda_stream) - - # create cuda_graph and assign static_inputs and static_outputs - self._encoder_cuda_graph = get_accelerator().create_graph() - self.static_encoder_inputs = inputs - self.static_encoder_kwargs = kwargs - - with get_accelerator().capture_to_graph(self._encoder_cuda_graph): - self.static_encoder_output = self._encode(*self.static_encoder_inputs, **self.static_encoder_kwargs) - - self.encoder_cuda_graph_created = True - - def encode(self, *inputs, **kwargs): - if self.enable_cuda_graph: - if self.encoder_cuda_graph_created: - outputs = self._graph_replay_encoder(*inputs, **kwargs) - else: - self._create_cuda_graph_encoder(*inputs, **kwargs) - outputs = self._graph_replay_encoder(*inputs, **kwargs) - return outputs - else: - return self._encode(*inputs, **kwargs) - - def _graph_replay(self, *inputs, **kwargs): - for i in range(len(inputs)): - if torch.is_tensor(inputs[i]): - self.static_inputs[i].copy_(inputs[i]) - for k in kwargs: - if torch.is_tensor(kwargs[k]): - self.static_kwargs[k].copy_(kwargs[k]) - get_accelerator().replay_graph(self._all_cuda_graph) - return self.static_output - - def forward(self, *inputs, **kwargs): - if self.enable_cuda_graph: - if self.cuda_graph_created: - outputs = self._graph_replay(*inputs, **kwargs) - else: - self._create_cuda_graph(*inputs, **kwargs) - outputs = self._graph_replay(*inputs, **kwargs) - return outputs - else: - return self._forward(*inputs, **kwargs) - - def _create_cuda_graph(self, *inputs, **kwargs): - # warmup to create the workspace and cublas handle - cuda_stream = torch.cuda.Stream() - cuda_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(cuda_stream): - for i in range(3): - ret = self._forward(*inputs, **kwargs) - torch.cuda.current_stream().wait_stream(cuda_stream) - - # create cuda_graph and assign static_inputs and static_outputs - self._all_cuda_graph = get_accelerator().create_graph() - self.static_inputs = inputs - self.static_kwargs = kwargs - - with get_accelerator().capture_to_graph(self._all_cuda_graph): - self.static_output = self._forward(*self.static_inputs, **self.static_kwargs) - - self.all_cuda_graph_created = True - - def _forward(self, sample, timestamp, encoder_hidden_states, return_dict=True): - return self.vae(sample, timestamp, encoder_hidden_states, return_dict) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/features/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/features/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/features/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/features/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/features/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 21f786c6376f44dd821973571ead6a96031dc16e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/features/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/features/__pycache__/cuda_graph.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/features/__pycache__/cuda_graph.cpython-310.pyc deleted file mode 100644 index fba6280745e53244daafe14a8536c296e5c8b40a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/features/__pycache__/cuda_graph.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/features/cuda_graph.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/features/cuda_graph.py deleted file mode 100644 index fc3a2951665426b5c98112ce307e395993644f8b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/features/cuda_graph.py +++ /dev/null @@ -1,27 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import ABC, abstractmethod - - -class CUDAGraph(ABC): - - def __init__(self, enable_cuda_graph=False): - super().__init__() - self.enable_cuda_graph = enable_cuda_graph - - @abstractmethod - def _create_cuda_graph(self): - """ - Create CUDA graph(s) - """ - raise NotImplementedError - - @abstractmethod - def _graph_replay(self): - """ - Replay CUDA graph(s) - """ - raise NotImplementedError diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 0efcf277084e0d115aae717f1b13ed7c9943728b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/clip_encoder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/clip_encoder.cpython-310.pyc deleted file mode 100644 index f1a0a7d525b0eb9f1e8f018c6760b13937a2ae05..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/clip_encoder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_base.cpython-310.pyc deleted file mode 100644 index ad528ca9fe48e0816ae43eadbcb8a2043966348e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bert.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bert.cpython-310.pyc deleted file mode 100644 index 2625a7dba14f2a13c08d03c643eb2fe2e3c1f608..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bert.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bloom.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bloom.cpython-310.pyc deleted file mode 100644 index 6503be7fd03b02dd22333203a9286b2c897bc344..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_bloom.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_gpt.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_gpt.cpython-310.pyc deleted file mode 100644 index 95f0489319d24eb06a3a830010419912c81e29ba..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_gpt.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_llama2.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_llama2.cpython-310.pyc deleted file mode 100644 index bdced73b9855965a4c27fcf0cd01304d712c6016..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_llama2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_megatron_gpt.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_megatron_gpt.cpython-310.pyc deleted file mode 100644 index 4149d2933d447221b0c866366f92c147876c0f68..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_megatron_gpt.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_opt.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_opt.cpython-310.pyc deleted file mode 100644 index 9b56d2e8c38b202987ca339142aa50791839b3b4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_opt.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_transformer.cpython-310.pyc deleted file mode 100644 index 7e9894384691904286574b1175f95a7b4ec1a061..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/__pycache__/ds_transformer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/clip_encoder.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/clip_encoder.py deleted file mode 100644 index 848a5b48dcf1336a4c80c20a9e21388cdff9e514..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/clip_encoder.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed.accelerator import get_accelerator -from ..features.cuda_graph import CUDAGraph - - -class DSClipEncoder(CUDAGraph, torch.nn.Module): - - def __init__(self, enc, enable_cuda_graph=False): - super().__init__(enable_cuda_graph=enable_cuda_graph) - enc.text_model._build_causal_attention_mask = self._build_causal_attention_mask - self.enc = enc - self.device = self.enc.device - self.dtype = self.enc.dtype - self.cuda_graph_created = [False, False] - self.static_inputs = [None, None] - self.static_kwargs = [None, None] - self.static_output = [None, None] - self._cuda_graphs = [None, None] - self.iter = 0 - self.config = self.enc.config - - def _build_causal_attention_mask(self, bsz, seq_len, dtype): - mask = torch.empty(bsz, seq_len, seq_len, dtype=dtype, device=get_accelerator().current_device_name()) - mask.fill_(torch.tensor(torch.finfo(dtype).min)) - mask.triu_(1) - mask = mask.unsqueeze(1) - return mask - - def _graph_replay(self, *inputs, **kwargs): - for i in range(len(inputs)): - if torch.is_tensor(inputs[i]): - self.static_inputs[self.iter][i].copy_(inputs[i]) - for k in kwargs: - if torch.is_tensor(kwargs[k]): - self.static_kwargs[self.iter][k].copy_(kwargs[k]) - get_accelerator().replay_graph(self._cuda_graphs[self.iter]) - return self.static_output[self.iter] - - def forward(self, *inputs, **kwargs): - if self.enable_cuda_graph: - if self.cuda_graph_created[self.iter]: - outputs = self._graph_replay(*inputs, **kwargs) - else: - self._create_cuda_graph(*inputs, **kwargs) - outputs = self._graph_replay(*inputs, **kwargs) - self.iter = (self.iter + 1) % 2 - return outputs - else: - return self.enc(*inputs, **kwargs) - - def _create_cuda_graph(self, *inputs, **kwargs): - # warmup to create the workspace and cublas handle - cuda_stream = torch.cuda.Stream() - cuda_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(cuda_stream): - for i in range(3): - ret = self._forward(*inputs, **kwargs) - torch.cuda.current_stream().wait_stream(cuda_stream) - - # create cuda_graph and assign static_inputs and static_outputs - self._cuda_graphs[self.iter] = get_accelerator().create_graph() - self.static_inputs[self.iter] = inputs - self.static_kwargs[self.iter] = kwargs - - with get_accelerator().capture_to_graph(self._cuda_graphs[self.iter]): - self.static_output[self.iter] = self._forward(*self.static_inputs[self.iter], - **self.static_kwargs[self.iter]) - - self.cuda_graph_created[self.iter] = True - - def _forward(self, *inputs, **kwargs): - return self.enc(*inputs, **kwargs) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_base.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_base.py deleted file mode 100644 index b4b113904997fd124aecdfe51529aab6c9e1b25c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_base.py +++ /dev/null @@ -1,15 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch.nn as nn - - -class DeepSpeedTransformerBase(nn.module): - - def __init__(self): - pass - - # this would be the new clean base class that will replace DeepSpeedTransformerInference. - # we currently don't know how this will look like but keeping it here as a placeholder. diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_bert.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_bert.py deleted file mode 100644 index 13075553ec8b6d818744cc5404ef4db31283ac8e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_bert.py +++ /dev/null @@ -1,20 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference - - -class DeepSpeedBERTInference(DeepSpeedTransformerInference): - """Initialize the DeepSpeed BERT Transformer Layer. - """ - - def __init__(self, - config, - mp_group=None, - quantize_scales=None, - quantize_groups=1, - merge_count=1, - mlp_extra_grouping=False): - super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_bloom.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_bloom.py deleted file mode 100644 index c48c7ed58ae5e2746c513a1acea509bbfd129784..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_bloom.py +++ /dev/null @@ -1,20 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference - - -class DeepSpeedBloomInference(DeepSpeedTransformerInference): - """Initialize the DeepSpeed Bloom Transformer Layer. - """ - - def __init__(self, - config, - mp_group=None, - quantize_scales=None, - quantize_groups=1, - merge_count=1, - mlp_extra_grouping=False): - super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_gpt.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_gpt.py deleted file mode 100644 index 3c3baed1f6186a89ed2b852f27ad56a5ab00d9f6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_gpt.py +++ /dev/null @@ -1,20 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference - - -class DeepSpeedGPTInference(DeepSpeedTransformerInference): - """Initialize the DeepSpeed GPT Transformer Layer. - """ - - def __init__(self, - config, - mp_group=None, - quantize_scales=None, - quantize_groups=1, - merge_count=1, - mlp_extra_grouping=False): - super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_llama2.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_llama2.py deleted file mode 100644 index 325bfb4f7e181d8b718748fd5ebf2f4f1ca3af64..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_llama2.py +++ /dev/null @@ -1,58 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference - - -class DeepSpeedLlama2Inference(DeepSpeedTransformerInference): - """Initialize the DeepSpeed OPT Transformer Layer. - """ - - def __init__(self, - config, - mp_group=None, - quantize_scales=None, - quantize_groups=1, - merge_count=1, - mlp_extra_grouping=False): - super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) - - def forward(self, *args, **kwargs): - - input = args[0] - input_mask = None - get_present = True - - self.allocate_workspace(input.size()) - - # We set the prev key/value to None when there is a prompt - if input.shape[1] > 1: - self.layer_past = None - layer_past = self.layer_past - - input_type = input.dtype - - if (self.config.dtype in [torch.float16, torch.bfloat16, torch.int8]) \ - and input.dtype == torch.float: - target_dtype = torch.half if self.dtype == torch.int8 else self.dtype - input = input.to(target_dtype) - - with torch.no_grad(): - attention_output, key, value, context_outputtn_ctx, inp_norm = \ - self.attention(input, - input_mask, - None, - layer_past, - get_present, - None, None, None, - self.norm_w, - self.norm_b, - None) - self.layer_past = (key, value) - output = self.mlp(attention_output, input, inp_norm, self.attention.attn_ob) - - output = output.to(input_type) - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_megatron_gpt.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_megatron_gpt.py deleted file mode 100644 index 055ed6d27d7db4cb39cf868d5c189ef517f4a08a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_megatron_gpt.py +++ /dev/null @@ -1,20 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference - - -class DeepSpeedMegatronGPTInference(DeepSpeedTransformerInference): - """Initialize the DeepSpeed Megatron GPT Transformer Layer. - """ - - def __init__(self, - config, - mp_group=None, - quantize_scales=None, - quantize_groups=1, - merge_count=1, - mlp_extra_grouping=False): - super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_opt.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_opt.py deleted file mode 100644 index 7bc5524d71c7155fad0b40eafeca490705f3e324..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_opt.py +++ /dev/null @@ -1,20 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference - - -class DeepSpeedOPTInference(DeepSpeedTransformerInference): - """Initialize the DeepSpeed OPT Transformer Layer. - """ - - def __init__(self, - config, - mp_group=None, - quantize_scales=None, - quantize_groups=1, - merge_count=1, - mlp_extra_grouping=False): - super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping) diff --git a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_transformer.py b/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_transformer.py deleted file mode 100644 index 7e3c81b714c06b84fff1d18daf28ecccbc1db4f6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds_transformer.py +++ /dev/null @@ -1,191 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import torch.nn as nn -from deepspeed import comm as dist -from deepspeed.ops.transformer.inference.op_binding.layer_norm import LayerNormOp -from deepspeed.utils.logging import log_dist - -from deepspeed.ops.transformer.inference.ds_mlp import DeepSpeedMLP -from deepspeed.ops.transformer.inference.ds_attention import DeepSpeedSelfAttention, BloomSelfAttention -from deepspeed.ops.transformer.inference.op_binding.workspace import WorkspaceOp -from deepspeed.accelerator import get_accelerator -import deepspeed -if deepspeed.HAS_TRITON and get_accelerator().is_triton_supported(): - from deepspeed.ops.transformer.inference.triton.mlp import TritonMLP - from deepspeed.ops.transformer.inference.triton.attention import TritonSelfAttention - - -class DeepSpeedTransformerInference(nn.Module): - """Initialize the DeepSpeed Transformer Layer. - Arguments: - layer_id: The layer index starting from 0, e.g. if model has 24 transformer layers, - layer_id will be 0,1,2...23 when each layer object is instantiated - config: An object of DeepSpeedInferenceConfig - mp_group: Model parallelism group initialized on the modeling side. - quantize_scales: This argument groups all the layers' scales used for quantization - quantize_groups: Number of groups used for quantizing the model - merge_count: Shows the number of model-parallel checkpoints merged before running inference. - We use this argument to control the quantization scale for the model parameters if a bigger - quantize-grouping than 1 is used. - mlp_extra_grouping: This flag is used to show a 2x higher number of groups used for the MLP part - of a Transformer layer. We use this feature for quantization to reduce the convergence impact - for specific downstream tasks. - """ - layer_id = 0 - workspace = None - - def __init__(self, - config, - mp_group=None, - quantize_scales=None, - quantize_groups=1, - merge_count=1, - mlp_extra_grouping=False): - super(DeepSpeedTransformerInference, self).__init__() - - self.config = config - self.config.layer_id = DeepSpeedTransformerInference.layer_id - DeepSpeedTransformerInference.layer_id += 1 - - data_type = torch.half if self.config.dtype == torch.int8 else self.config.dtype - - if DeepSpeedTransformerInference.layer_id == 1: - log_dist(f"DeepSpeed-Inference config: {self.config.__dict__}", [0]) - if deepspeed.HAS_TRITON and self.config.use_triton: - log_dist(f"Injecting Triton kernels ...", [0]) - - if self.config.bigscience_bloom: - self.attention = BloomSelfAttention(self.config, mp_group, quantize_scales, quantize_groups, merge_count) - assert not self.config.use_triton - else: - if deepspeed.HAS_TRITON and self.config.use_triton: - self.attention = TritonSelfAttention(self.config) - else: - self.attention = DeepSpeedSelfAttention(self.config, mp_group, quantize_scales, quantize_groups, - merge_count) - - if deepspeed.HAS_TRITON and self.config.use_triton: - self.mlp = TritonMLP(self.config) - else: - self.mlp = DeepSpeedMLP(self.config, mp_group, quantize_scales, quantize_groups, merge_count, - mlp_extra_grouping) - - device = get_accelerator().current_device_name() # if config.bigscience_bloom else 'cpu' - if self.config.set_empty_params: - self.norm_w = None - self.norm_b = None - else: - self.norm_w = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type, device=device), - requires_grad=False) - self.norm_b = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type, device=device), - requires_grad=False) - self.layer_past = None - self.layer_norm = LayerNormOp() - if DeepSpeedTransformerInference.workspace is None: - DeepSpeedTransformerInference.workspace = WorkspaceOp(self.config) - self._should_allocate_workspace = True - - def allocate_workspace(self, size): - # Allocate memory only on first layer forward - if self.config.layer_id == 0 and self._should_allocate_workspace: - DeepSpeedTransformerInference.workspace.allocate_workspace( - self.config.hidden_size, self.config.heads, size[1], size[0], DeepSpeedTransformerInference.layer_id, - self.config.mp_size, self.config.bigscience_bloom, - dist.get_rank() if dist.is_initialized() else 0, self.config.max_out_tokens, - self.config.min_out_tokens) - self._should_allocate_workspace = False - - @classmethod - def reset_cache(cls): - if cls.workspace is not None: - cls.workspace.reset_cache() - - def forward( - self, - input=None, - input_mask=None, - attention_mask=None, - attn_mask=None, - head_mask=None, - layer_past=None, - get_key_value=False, - get_present=False, - encoder_output=None, - enc_dec_attn_mask=None, - x=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - use_cache=False, - alibi=None, - output_attentions=False, - # TODO(arashb): 'layer_head_mask' and 'past_key_value' are only added to satisfy the OPT models API. - # This needs to be redesigned later! - layer_head_mask=None, - past_key_value=None, - **kwargs): - - if x is not None: - input = x - if "hidden_states" in kwargs: - input = kwargs["hidden_states"] - - input_mask = (input_mask if attn_mask is None else attn_mask) if attention_mask is None else attention_mask - - self.allocate_workspace(input.size()) - - get_present = (get_present or get_key_value or use_cache) - input_mask = input_mask if attention_mask is None else attention_mask - - # We set the prev key/value to None when there is a prompt - if input.shape[1] > 1: - self.layer_past = None - layer_past = layer_past if layer_past is not None else self.layer_past - head_mask = layer_head_mask if layer_head_mask is not None else head_mask - - attn_mask = None - if isinstance(input, tuple): - attn_mask = input[1] - input = input[0] - input_type = input.dtype - - if (self.config.dtype in [torch.float16, torch.bfloat16, torch.int8]) \ - and input.dtype == torch.float: - target_dtype = torch.half if self.config.dtype == torch.int8 else self.config.dtype - input = input.to(target_dtype) - - with torch.no_grad(): - attention_output, key, value, context_outputtn_ctx, inp_norm = \ - self.attention(input, - input_mask, - head_mask, - layer_past, - get_present, - encoder_hidden_states, - encoder_attention_mask, - output_attentions, - self.norm_w, - self.norm_b, - alibi, - **kwargs) - - presents = (key, value) - self.layer_past = presents if layer_past is None else None - output = self.mlp(attention_output, input, inp_norm, self.attention.attn_ob) - - if not self.config.pre_layer_norm: - output = self.layer_norm(output, self.norm_w, self.norm_b, self.config.epsilon) - - output = output.to(input_type) - if get_present: - output = (output, presents) - - if self.config.return_single_tuple: - return (output, ) - elif self.config.return_tuple: - return output if type(output) is tuple else (output, attn_mask) - else: - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__init__.py deleted file mode 100644 index 9fc2f979a04b7adb8aaa638d8e2db3d6b2946d98..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .replace_module import replace_transformer_layer, revert_transformer_layer, ReplaceWithTensorSlicing, GroupQuantizer, generic_injection -from .module_quantize import quantize_transformer_layer -from .replace_policy import HFBertLayerPolicy -from .layers import LinearAllreduce, LinearLayer, EmbeddingLayer, Normalize, set_autotp_mode -from .policy import DSPolicy diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 5843b8eae3585d5acb207ab3664d3ce83f10ac14..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/auto_tp.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/auto_tp.cpython-310.pyc deleted file mode 100644 index 2aab19375809def7fc0acd0be7b1063a29f8c59e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/auto_tp.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/auto_tp_model_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/auto_tp_model_utils.cpython-310.pyc deleted file mode 100644 index c1df6bc5312408dca6dae73ff84e08b6e4752823..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/auto_tp_model_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/fusedqkv_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/fusedqkv_utils.cpython-310.pyc deleted file mode 100644 index 26f4cf99c2ae7ccbc2e532007d756b99284e5584..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/fusedqkv_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/inject.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/inject.cpython-310.pyc deleted file mode 100644 index 9139fd19fba43cfa8b41246bfac4b532280aa583..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/inject.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/layers.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/layers.cpython-310.pyc deleted file mode 100644 index eb8150805609840f2c76bd8e02aeb9950ad2d2e3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/layers.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/load_checkpoint.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/load_checkpoint.cpython-310.pyc deleted file mode 100644 index 7e122146f8b2c7f1a7555083340e2077a5c8a27e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/load_checkpoint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/module_quantize.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/module_quantize.cpython-310.pyc deleted file mode 100644 index c732953190afa13a42574e0676cfe11d66235eff..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/module_quantize.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/policy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/policy.cpython-310.pyc deleted file mode 100644 index cbeea9444375c2bfd4a2f963183fdf80e458e982..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/policy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/replace_module.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/replace_module.cpython-310.pyc deleted file mode 100644 index 9f8953f31d1a59d4115ffcf9cdd7633ed576a693..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/replace_module.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/replace_policy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/replace_policy.cpython-310.pyc deleted file mode 100644 index 51334220208926bc087fe27a70c34297d827f7f1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/replace_policy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/tp_shard.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/tp_shard.cpython-310.pyc deleted file mode 100644 index bce51fe3800c11cd8a3d7cb447b0dade3fb13246..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/tp_shard.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index 8e89180cb01ac69e5f0d646f36f59219734d2c6d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/auto_tp.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/auto_tp.py deleted file mode 100644 index 82cd9042071e7de85cffdcbdba8dad6aa94568bc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/auto_tp.py +++ /dev/null @@ -1,509 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Automatic Tensor Parallelism -import re - -from torch import nn -from .replace_policy import replace_policies -from typing import Optional -import torch -from deepspeed import comm as dist -from .layers import * -from deepspeed.accelerator import get_accelerator -from .fusedqkv_utils import require_tp_fused_qkvw -from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list -from deepspeed.utils import groups -from deepspeed.module_inject.layers import is_autotp_training_mode - - -def move(tensor, device, copy=True): - if tensor.is_meta: - return torch.empty_like(tensor, device=device) - else: - # Using new tensors help in freeing memory (after split for example) was done before by calling clone(). - # Using copy=True instead of clone() will help in case of cpu --> cpu. - # Otherwise to() will not create a new copy for the view of the full tensor, and it will not be de-referenced. - return tensor.to(device, copy=copy) - - -class ReplaceWithTensorSlicing: - - def __init__(self, mp_group=None, mp_size=1, out_dim=1, in_dim=0): - if mp_group is not None: - self.gpu_index = dist.get_rank(group=mp_group) - else: - self.gpu_index = 0 - self.out_dim = out_dim - self.in_dim = in_dim - self.mp_size = mp_size - - def merge_assert(self, dim1, dim2): - assert dim1 > dim2, \ - 'Merging tensors is not allowed here! Please use deepspeed load_checkpoint\ - for merging your checkpoints before replacing the transformer layer with\ - inference-kernels' - - def strided_copy(self, - dst: Optional[torch.Tensor], - src: Optional[torch.Tensor], - num_splits: int, - int8: bool = False, - allocate_tensor: bool = False): - if src is None: - return src - src_shape = src.shape - dst_shape = dst.shape - - outer_dim = 0 if int8 else -1 - - if allocate_tensor: - dst = torch.empty_like(dst) - - src_split = torch.split(src.data, src.shape[outer_dim] // num_splits, dim=outer_dim) - if (len(src_shape) == 2 and len(dst_shape) == 2): - if src_shape[outer_dim] == dst_shape[self.out_dim]: - try: - dst = dst.reshape(-1).data.copy_(src.data.reshape(-1)).reshape(src.shape) - except: - print(dst.shape, src.shape) - exit() - dst = torch.nn.parameter.Parameter(dst, requires_grad=False) - if hasattr(src, 'scale'): - dst.scale = src.scale - return dst - self.merge_assert(src_shape[outer_dim], dst_shape[self.out_dim]) - qkv_size = dst_shape[self.out_dim] // num_splits - qkv_split = [torch.split(src_s, qkv_size, dim=outer_dim) for src_s in src_split] - weight_split = [ - torch.cat([qkv_s[i] for qkv_s in qkv_split], axis=outer_dim) for i in range(len(qkv_split[0])) - ] - dst = dst.reshape(-1).data.copy_(weight_split[self.gpu_index].contiguous().reshape(-1)).reshape( - weight_split[self.gpu_index].shape) - else: - if src_shape[0] == dst_shape[0]: - return torch.nn.parameter.Parameter(src) - qkv_size = dst_shape[0] // num_splits - qkv_split = [torch.split(src_s, qkv_size, dim=0) for src_s in src_split] - bias_split = [torch.cat([qkv_s[i] for qkv_s in qkv_split], axis=0) for i in range(len(qkv_split[0]))] - dst.data.copy_(bias_split[self.gpu_index].contiguous()) - - dst = torch.nn.parameter.Parameter(dst, requires_grad=False) - if hasattr(src, 'scale'): - dst.scale = src.scale - return dst - - def copy(self, dst, src, int8=False, allocate_tensor=False): - if src is None: - return src - assert not dst.data.is_meta # the torch.Tensor.copy_ method used below will silently fail on meta tensors - if allocate_tensor: - dst = torch.empty_like(dst) - outer_dim = 0 if int8 else 1 - inner_dim = 1 if int8 else 0 - src_shape = src.shape - dst_shape = dst.shape - if (len(src_shape) == 2 and len(dst_shape) == 2): - - if src_shape[inner_dim] == dst_shape[self.in_dim] and src_shape[outer_dim] == dst_shape[self.out_dim]: - dst = dst.reshape(-1).data.copy_(src.data.reshape(-1)).reshape(src.shape) - else: - if src_shape[inner_dim] != dst_shape[self.in_dim]: - self.merge_assert(src_shape[inner_dim], dst_shape[self.in_dim]) - dst.data.copy_(src[:, self.gpu_index * dst_shape[self.in_dim]: (self.gpu_index + 1) * dst_shape[self.in_dim]] if inner_dim == 1 else \ - src[self.gpu_index * dst_shape[self.in_dim]: (self.gpu_index + 1) * dst_shape[self.in_dim], :]) - else: - self.merge_assert(src_shape[outer_dim], dst_shape[self.out_dim]) - dst.data.copy_(src[:, self.gpu_index * dst_shape[self.out_dim]: (self.gpu_index + 1) * dst_shape[self.out_dim]] if outer_dim == 1 else \ - src[self.gpu_index * dst_shape[self.out_dim]: (self.gpu_index + 1) * dst_shape[self.out_dim], :]) - else: - if src_shape[0] == dst_shape[0]: - dst = src if src.dtype == dst.dtype else dst.data.copy_(src) - else: - dst.data.copy_(src[self.gpu_index * dst_shape[-1]:(self.gpu_index + 1) * dst_shape[-1]]) - dst = torch.nn.parameter.Parameter(dst, requires_grad=False) - if hasattr(src, 'scale'): - dst.scale = src.scale - return dst - - -class Loading(): - - def is_load_module(module): - load_layers = [nn.Linear, nn.Embedding, nn.LayerNorm] - load_layer_names = [ - "LPLayerNorm", "SharedEmbedding", "OPTLearnedPositionalEmbedding", "LlamaRMSNorm", "FalconLinear", - "MistralRMSNorm", "T5LayerNorm", "MixtralRMSNorm", "Phi3RotaryEmbedding", "Phi3SuScaledRotaryEmbedding", - "Phi3RMSNorm", "YuanRMSNorm", "YuanRotaryEmbedding", "Phi3LongRoPEScaledRotaryEmbedding", "Qwen2RMSNorm", - "Qwen3RMSNorm", "Qwen3MoeRMSNorm", "DeepseekV2RMSNorm", "DeepseekV3RMSNorm", - "DeepseekV2YarnRotaryEmbedding", "DeepseekV3YarnRotaryEmbedding", "MoEGate" - ] - return module.__class__ in load_layers or module._get_name() in load_layer_names - - def load_buffer(module, state_dict, prefix): - for name in module._buffers.keys(): - if module._buffers[name].data.is_meta: - module._buffers[name] = torch.nn.parameter.Parameter( - data=torch.empty_like(module._buffers[name].data, device="cpu"), - requires_grad=module._buffers[name].data.requires_grad) - if prefix + name in state_dict.keys(): - module._buffers[name].data.copy_(state_dict[prefix + name]) - - def load(module, state_dict, prefix, mp_group=None): - mp_replace = ReplaceWithTensorSlicing(mp_group=mp_group) - if hasattr(module, 'weight'): - if module.weight.data.is_meta: - # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here - module.weight = torch.nn.parameter.Parameter(data=torch.empty_like(module.weight.data, device="cpu"), - requires_grad=module.weight.data.requires_grad) - if 'query_key_value' in prefix: - module.weight = mp_replace.strided_copy(module.weight.data, - state_dict[prefix + 'weight'], - num_splits=3) - else: - module.weight = mp_replace.copy(module.weight.data, state_dict[prefix + 'weight']) - else: - if hasattr(module, 'norm') and hasattr(module.norm, 'weight'): - if module.norm.weight.data.is_meta: - # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here - module.norm.weight = torch.nn.parameter.Parameter( - data=torch.empty_like(module.norm.weight.data, device="cpu"), - requires_grad=module.norm.weight.data.requires_grad) - module.norm.weight = mp_replace.copy(module.norm.weight.data, state_dict[prefix + 'weight']) - - if prefix + 'bias' in state_dict.keys(): - if hasattr(module, 'bias'): - if module.bias.data.is_meta: - # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here - module.bias = torch.nn.parameter.Parameter(data=torch.empty_like(module.bias.data, device="cpu"), - requires_grad=module.bias.data.requires_grad) - module.bias = mp_replace.copy(module.bias, state_dict[prefix + 'bias']) - else: - if hasattr(module, 'norm') and hasattr(module.norm, 'bias'): - if module.norm.bias.data.is_meta: - # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here - module.norm.bias = torch.nn.parameter.Parameter( - data=torch.empty_like(module.norm.bias.data, device="cpu"), - requires_grad=module.norm.bias.data.requires_grad) - module.norm.bias = mp_replace.copy(module.norm.bias, state_dict[prefix + 'bias']) - - -class AutoTP(): - - def __init__(self, - module, - all_reduce_linears, - prefix, - state_dict, - linear_layer_setting, - orig_layer_impl, - keep_module_on_host=False): - self.module = module - self.all_reduce_linears = all_reduce_linears - self.prefix = prefix - self.state_dict = state_dict - - self.mp_size = None - self.mp_group = None - self.linear_layer_setting = linear_layer_setting - self.orig_layer_impl = orig_layer_impl - self.linear_policies = None - self.conv_linear_layer = False - TensorParallel_Layer.set_keep_module_on_host(keep_module_on_host) - - def in_module_list(module, module_list): - for item in module_list: - if type(item).__name__ == type(module).__name__: - return True - return False - - def get_module_list(model): - mlist = [] - for child in model.children(): - if isinstance(child, nn.ModuleList): - for module in child.children(): - if not mlist: - mlist = [module] - elif not AutoTP.in_module_list(module, mlist): - mlist = mlist + [module] - else: - mlist = mlist + AutoTP.get_module_list(child) - return mlist - - def supported(model): - unsupported = ['deberta', 'flaubert', 'fsmt', 'gpt2', 'led', 'longformer', 'xlm', 'xlnet'] - model = str(model) - key = re.search(r": (.*?)Model", model) - if key is None: - key = re.search(r": (.*?)Stack", model) - if key is None: - key = re.match(r"(.*?)Model", model) - assert key is not None, "Not able to determine model policy automatically. Please provide policy." - if key.group(1).lower() in unsupported: - return False - return True - - def get_layers(parent, module): - layer_list = [] - for key, submodule in module._modules.items(): - if isinstance(submodule, nn.Linear): - layer_list = layer_list + [parent + "." + key] - elif isinstance(submodule, nn.LayerNorm) or key == 'LayerNorm' or key == 'layer_norm': - layer_list = layer_list + ["ln"] - else: - layer_list = layer_list + AutoTP.get_layers(key, submodule) - return layer_list - - def update_policy_list(policy_list, new_module, new_gems): - if len(policy_list): - for i, policy in enumerate(policy_list): - # if module already exists in policy, combine gems and remove duplicates - if policy[0] == type(new_module): - new_gems = set(new_gems + policy[1]) - policy_list[i] = tuple([type(new_module), new_gems]) - return policy_list - policy_list.append(tuple([type(new_module), new_gems])) - return policy_list - - def kernel_supported(module_list): - policy = [] - for plcy in replace_policies: - # instantiate a throw-away policy in order to populate the _orig_layer_class - _ = plcy(None) - if isinstance(plcy._orig_layer_class, list): - for orig_layer_class in plcy._orig_layer_class: - policy.append(orig_layer_class) - elif plcy._orig_layer_class is not None: - policy.append(plcy._orig_layer_class) - for child in module_list: - if child.__class__ in policy: - return True - return False - - def tp_parser(model): - policy_list = [] - module_list = [] - layer_list = [] - gem_list = [] - - module_list = AutoTP.get_module_list(model) - assert AutoTP.supported(model), "AutoTP not supported for model. Please use kernel injection since container policy for model exists." \ - if AutoTP.kernel_supported(module_list) else "AutoTP not supported for model. Please provide policy." - norm_layer_name_list = ['LayerNorm', 'layer_norm', 'ln_1', 'ln_2'] - #ln_1 , ln_2 for Qwen - for module in module_list: - for key, submodule in module._modules.items(): - if isinstance(submodule, nn.Linear): - layer_list = layer_list + ["." + key] - elif isinstance(submodule, nn.LayerNorm) or key in norm_layer_name_list: - layer_list = layer_list + ["ln"] - else: - layer_list = layer_list + AutoTP.get_layers(key, submodule) - for i, layer in enumerate(layer_list): - if layer == 'ln': - if layer_list[i - 1] != 'ln': - gem_list = gem_list + [layer_list[i - 1]] - elif 'out_proj' in layer: - gem_list = gem_list + [layer] - elif 'o_proj' in layer: - gem_list = gem_list + [layer] - elif 'down_proj' in layer: - gem_list = gem_list + [layer] - elif 'attention.dense' in layer and 'GPTNeoX' in str(model): - gem_list = gem_list + [layer] - elif 'self_attention.dense' in layer and 'falcon' in str( - type(module)): # this is a hack to get the right linear layer for this model! - gem_list = gem_list + [layer] - # Mixtral-7x8b used w2*act(w1*w3) linear. need to replace w2 to linearallreduce. - elif 'w2' in layer and 'Mixtral' in str(type(module)): - gem_list = gem_list + [layer] - elif 'self_attn.dense' in layer and 'Phi' in str(type(module)): - gem_list = gem_list + [layer] - elif 'self_attention.dense' in layer and 'ChatGLM' in str(model): - gem_list = gem_list + [layer] - elif 'dense_4h_to_h' in layer and 'ChatGLM' in str(model): - gem_list = gem_list + [layer] - - layer_list = [] - if gem_list != []: - gem_list = list(set(gem_list)) - policy_list = AutoTP.update_policy_list(policy_list, module, gem_list) - gem_list = [] - assert len(policy_list), "AutoTP not supported for model. Please use kernel injection since container policy for model exists." \ - if AutoTP.kernel_supported(module_list) else "Not able to determine model policy automatically. Please provide policy." - return policy_list - - def set_tensor_parallel_config(self, mp_size, mp_group): - - if is_autotp_training_mode(): - self.mp_group = groups.get_tensor_model_parallel_group() - self.mp_size = groups.get_tensor_model_parallel_world_size() - return - - self.mp_size = mp_size - self.mp_group = mp_group - - def _replace(self, child, name, conv_linear_layer): - # This function should clearly define the routing rules for specific layers - # and avoid any complex shard-related logic. - if getattr(child, "replaced", False) == True: - return - - weight_shape = child.weight.shape - mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group) - # For TP layer skip, e.g., MoE gate, deepseek low rank layer skip - if "mlp.gate" == name or "q_a_proj" in name or "kv_a_proj_with_mqa" in name or name == "block_sparse_moe.gate" or ( - ('mlp.shared_expert_gate' == name or 'mlp.gate' == name) and 'qwen2_moe' in str(type(self.module))): - return child - # For Yuan model - if 'Yuan' in str(self.module): - if 'v_proj' in name: - return Yuan_LinearLayer(child, self.mp_group) - - elif 'o_proj' in name: - return Yuan_LinearAllreduce(child, self.mp_group) - - # For MLP including chunk layer. - if 'gate_up_proj' in name or ('dense_h_to_4h' in name and 'GLM' in str(self.module)): - return GateUpPack_LinearLayer(child, self.mp_group) - # For Arctic model, bypass to all_reduce replacement for w2 weights - arctic_w2_all_reduce_linear = False - if 'Arctic' in str(self.module) and 'w2' in name: - arctic_w2_all_reduce_linear = True - # For MoE MLP model, e.g., deepseek and jamba - down_proj = False - if 'down_proj' in name: - down_proj = True - if name in self.all_reduce_linears or arctic_w2_all_reduce_linear or down_proj: - - setattr(child, "replaced", True) - if self.conv_linear_layer: - return Conv_LinearALlreduce(child, self.mp_group, name=name) - elif name == "lm_head" or name == 'embed_out': - return LmHeadLinearAllreduce(child, self.mp_group) - - return LinearAllreduce(child, self.mp_group, name=name) - else: - - setattr(child, "replaced", True) - if self.conv_linear_layer: - conv_LinearLayer(child, self.mp_group) - elif require_tp_fused_qkvw(name, self.mp_size): - #Check and handle fused qkv for TP - return fused_LinearLayer(child, self.mp_group, fused_module=self.module) - - return LinearLayer(child, self.mp_group, name=name) - - def _slice_embedding(self, child, name, conv_linear_layer): - if getattr(child, "replaced", False) == True: - return - mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group) - - if hasattr(child.weight, 'ds_tensor'): - data = child.weight.ds_tensor.data.split(get_shard_size_list(child.weight.shape[1], self.mp_size), dim=1) - else: - data = child.weight.data.split(get_shard_size_list(child.weight.shape[1], self.mp_size, name), dim=1) - data = data[mp_replace.gpu_index].to(get_accelerator().current_device_name()) - data = torch.nn.parameter.Parameter(data, requires_grad=False) - - new_embedding = nn.Embedding(child.weight.shape[0], get_shard_size(child.weight.shape[1], self.mp_size, name)) - new_embedding.weight.data.copy_(data) - setattr(child, "replaced", True) - return new_embedding - - def update_mp_params(self, child): - if getattr(child, "replaced", False) == True: - return - param_list = [ - "n_heads", "inner_dim", "num_heads", "num_kv", "num_attention_heads", "num_attn_heads", "all_head_size", - "embed_dim", "hidden_size", "num_key_value_heads", "num_kv_heads", "kv_n_heads", "d_model", - "num_attention_heads_per_partition", "num_multi_query_groups_per_partition", "hidden_size_per_partition" - ] - for param in param_list: - if "Yuan" in str(child) and 'embed_dim' in param_list: - param_list.remove('embed_dim') - if hasattr(child, param): - param_val = getattr(child, param) - setattr(child, param, get_shard_size(param_val, self.mp_size)) - setattr(child, "replaced", True) - - def update_linear_policies(self): - self.conv_linear_layer = False - if self.linear_layer_setting is not None: - self.linear_policies = {self.linear_layer_setting[0]: self._replace} - if len(self.linear_layer_setting) == 2: - self.linear_policies.update({self.linear_layer_setting[1]: self._slice_embedding}) - else: - import transformers - if self.orig_layer_impl is transformers.models.gpt2.modeling_gpt2.GPT2Block: - try: - self.conv_linear_layer = True - self.linear_policies = {transformers.pytorch_utils.Conv1D: self._replace} - except ImportError: - self.linear_policies = {nn.Linear: self._replace} - else: - self.linear_policies = {nn.Linear: self._replace, nn.Embedding: self._slice_embedding} - - def _replace_module(self, r_module, prev_name='', prev_class_name=''): - for name, child in r_module.named_children(): - if prev_class_name == "": - class_name = prev_name - elif prev_name == "": - class_name = prev_class_name - else: - class_name = prev_class_name + '.' + prev_name - checking_key = self.prefix + '.' + class_name + '.' + name + '.' if class_name != "" else self.prefix + '.' + name + '.' - if Loading.is_load_module(child) and self.state_dict is not None: - if any(checking_key in item for item in self.state_dict): - Loading.load(child, self.state_dict, checking_key, self.mp_group) - else: - continue - if len(child._buffers) != 0 and self.state_dict is not None: - Loading.load_buffer(child, self.state_dict, checking_key) - if child.__class__ in self.linear_policies: - setattr(r_module, name, self.linear_policies[child.__class__](child, prev_name + '.' + name, - self.conv_linear_layer)) - elif any(isinstance(child, lp) for lp in self.linear_policies): - # Added for falcon model support - # Note: isinstance will account for class inheritance, child.__class__ does not - key = None - for lp in self.linear_policies: - if isinstance(child, lp): - key = lp - break - assert key is not None - setattr(r_module, name, self.linear_policies[key](child, prev_name + '.' + name, - self.conv_linear_layer)) - else: - self.update_mp_params(child) - self._replace_module(child, name, class_name) - return r_module - - def get_model_num_kv_heads(self, config): - num_kv_heads = None - # multi_query_group_num is for chatglm2 & chatglm3 - kv_head_names = [ - 'multi_query_group_num', 'num_kv_heads', 'num_key_value_heads', 'num_attention_heads', 'n_heads', - 'attention_heads' - ] - for name in kv_head_names: - if hasattr(config, name): - num_kv_heads = getattr(config, name) - if num_kv_heads is not None: - break - return num_kv_heads - - def _replace_last_linear_module(self, r_module): - if hasattr(r_module, "lm_head"): - name = "lm_head" - child = r_module.lm_head - elif hasattr(r_module, "embed_out"): - name = "embed_out" - child = r_module.embed_out - else: - return r_module - if child.__class__ in self.linear_policies: - setattr(r_module, name, self.linear_policies[child.__class__](child, name, self.conv_linear_layer)) - return r_module diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/auto_tp_model_utils.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/auto_tp_model_utils.py deleted file mode 100644 index a71b1a54d6f6031c18899b1a5294dd5dd963e92d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/auto_tp_model_utils.py +++ /dev/null @@ -1,104 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed import comm as dist -import torch -from typing import Optional -from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list - - -def build_bloom_alibi_tensor(attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor: - """ - Link to paper: https://arxiv.org/abs/2108.12409 Alibi tensor is not causal as the original paper mentions, it - relies on a translation invariance of softmax for quick implementation: with l being a tensor, and a fixed value - `softmax(l+a) = softmax(l)`. Based on - https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L742 - TODO @thomasw21 this doesn't work as nicely due to the masking strategy, and so masking varies slightly. - - Args: - Returns tensor shaped (batch_size * num_heads, 1, max_seq_len) - attention_mask (`torch.Tensor`): - Token-wise attention mask, this should be of shape (batch_size, max_seq_len). - num_heads (`int`, *required*): - number of heads - dtype (`torch.dtype`, *optional*, default=`torch.bfloat16`): - dtype of the output tensor - """ - import math - batch_size, seq_length = attention_mask.shape - closest_power_of_2 = 2**math.floor(math.log2(num_heads)) - base = torch.tensor(2**(-(2**-(math.log2(closest_power_of_2) - 3))), - device=attention_mask.device, - dtype=torch.float32) - powers = torch.arange(1, 1 + closest_power_of_2, device=attention_mask.device, dtype=torch.int32) - slopes = torch.pow(base, powers) - - if closest_power_of_2 != num_heads: - extra_base = torch.tensor(2**(-(2**-(math.log2(2 * closest_power_of_2) - 3))), - device=attention_mask.device, - dtype=torch.float32) - num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2) - extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=attention_mask.device, dtype=torch.int32) - slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0) - - # Note: alibi will added to the attention bias that will be applied to the query, key product of attention - # => therefore alibi will have to be of shape (batch_size, num_heads, query_length, key_length) - # => here we set (batch_size=1, num_heads=num_heads, query_length=1, key_length=max_length) - # => the query_length dimension will then be broadcasted correctly - # This is more or less identical to T5's relative position bias: - # https://github.com/huggingface/transformers/blob/f681437203baa7671de3174b0fa583c349d9d5e1/src/transformers/models/t5/modeling_t5.py#L527 - arange_tensor = ((attention_mask.cumsum(dim=-1) - 1) * attention_mask)[:, None, :] - alibi = slopes[..., None] * arange_tensor - if dist.is_initialized(): - num_heads_per_rank = get_shard_size(num_heads, dist.get_world_size()) - offset = sum(get_shard_size_list(num_heads, dist.get_world_size())[0:dist.get_rank()]) - alibi = alibi.view(batch_size, num_heads, 1, seq_length) - alibi = alibi[:, offset:num_heads_per_rank + offset, :, :] - return alibi.reshape(batch_size * num_heads_per_rank, 1, seq_length).to(dtype) - else: - return alibi.reshape(batch_size * num_heads, 1, seq_length).to(dtype) - - -def get_alibi_mask(self, tensor, seq_length_with_past): - mask = self.get_alibi_mask_orig(tensor, seq_length_with_past) - if not self.training and dist.is_initialized(): - num_heads_per_rank = get_shard_size(self.n_head, dist.get_world_size()) - offset = sum(get_shard_size_list(self.n_head, dist.get_world_size())[0:dist.get_rank()]) - mask = mask[offset:num_heads_per_rank + offset, :seq_length_with_past, :seq_length_with_past] - - return mask - - -def build_mpt_atten_bias_tensor(self, - device, - dtype, - attention_mask: Optional[torch.ByteTensor] = None, - prefix_mask: Optional[torch.ByteTensor] = None, - sequence_id: Optional[torch.LongTensor] = None): - (attn_bias, attention_mask) = self._attn_bias_orig(device, - dtype, - attention_mask=attention_mask, - prefix_mask=prefix_mask, - sequence_id=sequence_id) - if dist.is_initialized(): - num_heads_per_rank = get_shard_size(self.config.n_heads, dist.get_world_size()) - offset = sum(get_shard_size_list(self.config.n_heads, dist.get_world_size())[0:dist.get_rank()]) - attn_bias = attn_bias[:, offset:num_heads_per_rank + offset, :, :] - return attn_bias, attention_mask - - -def build_mpt_alibi_tensor(self, num_heads, sequence_length, alibi_bias_max=8, device=None) -> torch.Tensor: - r""" - Link to paper: https://arxiv.org/abs/2108.12409 - Alibi tensor is not causal as the original paper mentions, it - relies on a translation invariance of softmax for quick implementation. This implementation has been copied from - the alibi implementation of MPT source code that led to slightly different results than the Bloom alibi: - https://huggingface.co/mosaicml/mpt-7b/blob/main/attention.py#L292 - """ - alibi = self.build_mpt_alibi_tensor_orig(num_heads, sequence_length, alibi_bias_max, device) - if dist.is_initialized(): - num_heads_per_rank = int(num_heads / dist.get_world_size()) - offset = dist.get_rank() * num_heads_per_rank - alibi = alibi[offset:num_heads_per_rank + offset, :, :] - return alibi diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__init__.py deleted file mode 100644 index 993d14071659303ea6702b0d5182c6d04c931905..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .bert import DS_BERTContainer, HFBertLayerPolicy -from .bloom import DS_BloomContainer, BLOOMLayerPolicy, supported_models -from .distil_bert import DS_DistilBERTContainer, HFDistilBertLayerPolicy -from .gpt2 import DS_GPT2Container, HFGPT2LayerPolicy -from .gptj import DS_GPTJContainer, HFGPTJLayerPolicy -from .gptneo import DS_GPTNEOContainer, HFGPTNEOLayerPolicy -from .gptneox import DS_GPTNEOXContainer, GPTNEOXLayerPolicy -from .llama import DS_LLAMAContainer, LLAMALayerPolicy -from .llama2 import LLAMA2LayerPolicy, DS_LLAMA2Container -from .internlm import DS_InternLMContainer, InternLMLayerPolicy -from .megatron_gpt import DS_MegatronGPTContainer, MegatronLayerPolicy -from .megatron_gpt_moe import DS_MegatronGPTMoEContainer, MegatronMoELayerPolicy -from .opt import DS_OPTContainer, HFOPTLayerPolicy -from .clip import DS_CLIPContainer, HFCLIPLayerPolicy -from .unet import UNetPolicy -from .vae import VAEPolicy diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 923e1c264af0d375b68bd6c4a180e8ae65076187..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/base.cpython-310.pyc deleted file mode 100644 index f5fbfdf8125d1e5bfb3901f1b84da00154c26961..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/base_moe.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/base_moe.cpython-310.pyc deleted file mode 100644 index 022a65b51ef6f457403191d975be4da143ca0453..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/base_moe.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/bert.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/bert.cpython-310.pyc deleted file mode 100644 index 142129d93547d351e8cf00390cb5fa3aa631f7d4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/bert.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/bloom.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/bloom.cpython-310.pyc deleted file mode 100644 index 71258bd8924ec289ee2e65d6591f38cd18dcb529..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/bloom.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/clip.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/clip.cpython-310.pyc deleted file mode 100644 index 022541e8b36ad8481f5544b59704bdca0e55c974..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/clip.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/distil_bert.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/distil_bert.cpython-310.pyc deleted file mode 100644 index 9658edcf096a448a74e6d4f8675f6fb6bb9272b9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/distil_bert.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gpt2.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gpt2.cpython-310.pyc deleted file mode 100644 index f9508feb18fece4614eca8ec96abef7d19963fbe..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gpt2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptj.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptj.cpython-310.pyc deleted file mode 100644 index b6fc1d9c3570254acfe29ed62b3a01f5b6a1d077..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptj.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptneo.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptneo.cpython-310.pyc deleted file mode 100644 index 20b61ddb249f7381aaa32e8db9530554fc788748..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptneo.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptneox.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptneox.cpython-310.pyc deleted file mode 100644 index 038b291a68bf1ce696f6f4e24c4b20fbe7c89bf6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/gptneox.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/internlm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/internlm.cpython-310.pyc deleted file mode 100644 index a81336c56abd511a04f5d463869e3df7d33bf975..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/internlm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/llama.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/llama.cpython-310.pyc deleted file mode 100644 index 4cc1f680066364d364cc77a9eb9355225475371a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/llama.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/llama2.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/llama2.cpython-310.pyc deleted file mode 100644 index ffe78d415cd1603a1d1f061561fb6417f0c4f880..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/llama2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt.cpython-310.pyc deleted file mode 100644 index 3a2f8b05af386551ed85536b7603e9816d69fbd6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt_moe.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt_moe.cpython-310.pyc deleted file mode 100644 index d11434d939c72588cd4c6167bc66075eaa35edcd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/megatron_gpt_moe.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/opt.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/opt.cpython-310.pyc deleted file mode 100644 index a8857af09d7f331995f83877647b1ddc8f68852b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/opt.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/unet.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/unet.cpython-310.pyc deleted file mode 100644 index e66279c3d92bc21329c700c01997c5420a67dfba..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/unet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/vae.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/vae.cpython-310.pyc deleted file mode 100644 index fa76a9e11111cd96f9e15cf94f5d0dc666457495..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/__pycache__/vae.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/base.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/base.py deleted file mode 100644 index 83e109167ffec669b6b1d77c2996010a70678f27..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/base.py +++ /dev/null @@ -1,322 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Create a container object to save model-specific tensors using the policy file above. -from abc import ABC - -import torch - -import deepspeed -from deepspeed.ops.transformer.inference.config import DeepSpeedInferenceConfig -from deepspeed.accelerator import get_accelerator - -# If the intermediate size attribute is set DEFAULT_INTERMEDIATE_SIZE -# it is assumed the intermediate size is 4x the embedding dimension -DEFAULT_INTERMEDIATE_SIZE = -1 - - -class BaseConvolutionContainer(ABC): - # not implemented - def __init__(self): - pass - - -class BaseTransformerContainer(ABC): - - def __init__(self, policy, config, model_config, layer_id, child): - self.policy = policy - self.config = config - self.model_config = model_config - self.layer_id = layer_id - self.child = child - - self.megatron_v2 = self.policy.is_megatron_v2 - self.scale_attention = self.policy.scale_attention - self.ckpt_load_enabled = False - - # configuration for models. todo: can this be moved to a pydantic model config? - self.hidden_size = None - self.intermediate_size = None - self.num_attention_heads = None - self.mp_size = self.config.tensor_parallel.tp_size - self.pre_layer_norm = self.model_config.do_layer_norm_before if \ - hasattr(self.model_config, 'do_layer_norm_before') else self.policy.pre_attn_norm - self.dtype = self.config.dtype - self.attn_linear_layer = self.policy.linear_layer - self.mlp_linear_layer = self.policy.linear_layer - self.return_tuple = self.config.return_tuple - self.triangular_masking = True - self.local_attention = ((self.model_config.attention_layers[self.layer_id] == "local") if hasattr( - self.model_config, 'attention_layers') else False) - self.window_size = getattr(self.model_config, "window_size", 1) - self.mlp_act_func_type = self.policy.mlp_act_func_type - self.norm_type = self.policy.norm_type - self.training_mp_size = self.config.training_mp_size - self.bigscience_bloom = False - self.max_out_tokens = self.config.max_out_tokens - self.min_out_tokens = self.config.min_out_tokens - self.scale_attn_by_inverse_layer_idx = getattr(self.config, "scale_attn_by_inverse_layer_idx", False) - self.use_mup = self.policy.use_mup - self.return_single_tuple = False - self.rotary_dim = self.get_rotary_dim() - self.mlp_after_attn = (self.rotary_dim is None or self.rotary_dim < 0) - - # Attention tensors - self.qkvw = None - self.qkvb = None - self.dense_w = None - self.dense_b = None - # MLP tensors - self._h4h_w = None - self._h4h_b = None - self._4hh_w = None - self._4hh_b = None - # LayerNorm tensors - self.attn_nw = None - self.attn_nb = None - self.input_nw = None - self.input_nb = None - - self.mp_group = None - self.use_triton = False - - # Triton - self.use_triton = config.use_triton and deepspeed.HAS_TRITON - - def create_ds_model_config(self): - self.set_hidden_heads(*self.policy.get_hidden_heads()) - assert self.num_attention_heads % self.mp_size == 0,\ - "To run the model parallel across the GPUs, the attention_heads require to be divisible by the world_size!" +\ - "This is because the attention computation is partitioned evenly among the parallel GPUs." - - self.ds_model_config = DeepSpeedInferenceConfig( - hidden_size=self.hidden_size, - intermediate_size=self.intermediate_size, - heads=self.num_attention_heads, - layer_norm_eps=self.layernorm_epsilon, - dtype=self.dtype, - pre_layer_norm=self.pre_layer_norm, - norm_type=self.norm_type, - mp_size=self.mp_size, - return_tuple=self.return_tuple, - triangular_masking=self.triangular_masking, - local_attention=self.local_attention, - window_size=self.window_size, - rotary_dim=self.rotary_dim, - mlp_after_attn=self.mlp_after_attn, - mlp_act_func_type=self.mlp_act_func_type, - training_mp_size=self.training_mp_size, - bigscience_bloom=self.bigscience_bloom, - max_out_tokens=self.max_out_tokens, - min_out_tokens=self.min_out_tokens, - scale_attn_by_inverse_layer_idx=self.scale_attn_by_inverse_layer_idx, - use_mup=self.use_mup, - return_single_tuple=self.return_single_tuple, - set_empty_params=self.config.set_empty_params, - transposed_mode=self.config.transposed_mode, - use_triton=self.use_triton, - triton_autotune=self.config.triton_autotune) - - if self.use_triton and deepspeed.HAS_TRITON: - from .bert import DS_BERTContainer - if not isinstance(self, DS_BERTContainer): - raise NotImplementedError("Triton kernels are only for BERT-like models yet") - - if not self.config.triton_autotune: - from deepspeed.ops.transformer.inference.triton.matmul_ext import fp16_matmul - fp16_matmul.skip_autotune() - - return self.ds_model_config - - def check_meta_tensor_support(self): - if hasattr(self.qkvw, 'is_meta'): - if self.qkvw.is_meta: - assert self.ckpt_load_enabled, "Meta tensors are not supported for this model currently." - else: - raise NotImplementedError("Meta tensor support is not available, please upgrade to torch 1.10+") - - def initialize_tensors(self, enable_training=False): - # Set the tensors from policy (user module) to container (DS module) - self.set_attention(*self.policy.attention(enable_training=enable_training)) - self.set_mlp(*self.policy.mlp(enable_training=enable_training)) - self.set_layernorm(*self.policy.layernorm()) - #self.check_meta_tensor_support() - - def convert_to_required_dtype(self): - # Note: converting tensors to fp16 requires that we do it in-place using self.__dict__ and not make a list/dict copy - if self.dtype in [torch.half, torch.bfloat16]: - for k, v in self.__dict__.items(): - # The list comprehension is used for MoE tensor lists - if isinstance(v, list) and all((isinstance(tensor, torch.Tensor) \ - or isinstance(tensor, torch.nn.Parameter)) for tensor in v): - self.__dict__[k] = [moe_tensor.to(self.dtype) for moe_tensor in v] - - if isinstance(v, torch.Tensor) or isinstance(v, torch.nn.Parameter): - self.__dict__[k] = v.to(self.dtype) - - def get_rotary_dim(self): - if hasattr(self.model_config, 'rotary_dim'): - return self.model_config.rotary_dim - if hasattr(self.child, 'attention') and hasattr(self.child.attention, 'rotary_ndims'): - return self.child.attention.rotary_ndims - return -1 - - def set_moe(self, moe=False): - self.moe = moe - - def set_tensor_parallel_config(self, mp_size, mp_group): - self.mp_size = mp_size - self.mp_group = mp_group - - def set_quantization_config(self, quantizer): - self.quantizer = quantizer - - def set_hidden_heads(self, hidden_size, num_attention_heads, epsilon, intermediate_size): - """ - Args: - hidden_size: embedding dimension of the model - num_attention_heads: number of attention heads in the model - epsilon: epsilon value for layer norm (same value used for all norms) - intermediate_size: Size of MLP projection. If `DEFAULT_INTERMEDIATE_SIZE` is passed - it is assumed to be `4 * hidden_size` - """ - self.hidden_size = hidden_size - if intermediate_size == DEFAULT_INTERMEDIATE_SIZE: - self.intermediate_size = 4 * hidden_size - else: - self.intermediate_size = intermediate_size - self.num_attention_heads = num_attention_heads - self.layernorm_epsilon = epsilon - - def set_attention(self, qkvw, qkvb, dense_w, dense_b): - self.qkvw = qkvw - self.qkvb = qkvb - self.dense_w = dense_w - self.dense_b = dense_b - - def set_mlp(self, _h4h_w, _h4h_b, _4hh_w, _4hh_b): - self._h4h_w = _h4h_w - self._h4h_b = _h4h_b - self._4hh_w = _4hh_w - self._4hh_b = _4hh_b - - def set_layernorm(self, attn_nw, attn_nb, input_nw, input_nb): - self.attn_nw = attn_nw - self.attn_nb = attn_nb - self.input_nw = input_nw - self.input_nb = input_nb - - def apply_weight_quantization(self): - # quantize attention weights - self.attention_quantization() - - # quantize mlp weights - self.mlp_quantization() - - def attention_quantization(self): - self.module.attention.attn_qkvw = self.quantizer.quantize(self.module.attention.attn_qkvw) - self.module.attention.attn_ow = self.quantizer.quantize(self.module.attention.attn_ow) - - def mlp_quantization(self): - self.module.mlp.inter_w = self.quantizer.quantize(self.module.mlp.inter_w) - self.module.mlp.output_w = self.quantizer.quantize(self.module.mlp.output_w) - - def apply_tensor_parallelism(self, mp_replace): - # setup the new Attention module - self.attention_qkv_mp(mp_replace) - self.attention_o_mp(mp_replace) - - # setup the new MLP module - self.mlp_inter_mp(mp_replace) - self.mlp_output_mp(mp_replace) - - # Apply weight quantization - # TODO(cmikeh2): Re-enable this once verified - #self.apply_weight_quantization() - - def attention_qkv_mp(self, mp_replace, reversed_dim=False): - self.module.attention.attn_qkvw = mp_replace.strided_copy(self.module.attention.attn_qkvw, - self.qkvw, - num_splits=3, - int8=reversed_dim) - self.module.attention.attn_qkvb = mp_replace.strided_copy(self.module.attention.attn_qkvb, - self.qkvb, - num_splits=3, - int8=reversed_dim) - - def attention_o_mp(self, mp_replace, reversed_dim=False): - self.module.attention.attn_ow = mp_replace.copy(self.module.attention.attn_ow, self.dense_w, int8=reversed_dim) - self.module.attention.attn_ob = mp_replace.copy(self.module.attention.attn_ob, - self.dense_b, - int8=reversed_dim, - allocate_tensor=reversed_dim) - - def mlp_inter_mp(self, mp_replace, reversed_dim=False): - self.module.mlp.inter_w = mp_replace.copy(self.module.mlp.inter_w, self._h4h_w, int8=reversed_dim) - self.module.mlp.inter_b = mp_replace.copy(self.module.mlp.inter_b, self._h4h_b, int8=reversed_dim) - - def mlp_output_mp(self, mp_replace, reversed_dim=False): - self.module.mlp.output_w = mp_replace.copy(self.module.mlp.output_w, self._4hh_w, int8=reversed_dim) - self.module.mlp.output_b = mp_replace.copy(self.module.mlp.output_b, - self._4hh_b, - int8=reversed_dim, - allocate_tensor=reversed_dim) - - def copy_data_to_new_module(self): - params = {'attn_nw': self.attn_nw, 'attn_nb': self.attn_nb} - for key in params: - if params[key] is None: - setattr(self.module.mlp, key, None) - else: - setattr(self.module.mlp, key, - torch.nn.parameter.Parameter(params[key].to(get_accelerator().current_device_name()))) - - params = {'norm_w': self.input_nw, 'norm_b': self.input_nb} - for key in params: - if params[key] is None: - setattr(self.module, key, None) - else: - setattr(self.module, key, - torch.nn.parameter.Parameter(params[key].to(get_accelerator().current_device_name()))) - - def transpose(self): - self.transpose_attention() - self.transpose_mlp() - - def transpose_attention(self): - if self.attn_linear_layer: - self.qkvw = self.transpose_impl(self.qkvw.data) - self.dense_w = self.transpose_impl(self.dense_w.data) - - def transpose_mlp(self): - if self.mlp_linear_layer: - self._h4h_w = self.transpose_impl(self._h4h_w.data) - self._4hh_w = self.transpose_impl(self._4hh_w.data) - - def transpose_impl(self, data): - data = data.contiguous() - data.reshape(-1).copy_(data.transpose(-1, -2).contiguous().reshape(-1)) - data = data.reshape(data.shape[-1], data.shape[-2]) - data.to(get_accelerator().current_device_name()) - return data - - def get_all_params(self): - params = [ - self.attn_nw, - self.attn_nb, - self.input_nw, - self.input_nb, - ] - - params.extend(self.get_attn_params()) - params.extend(self.get_mlp_params()) - - return params - - def get_attn_params(self): - return [self.qkvw, self.qkvb, self.dense_w, self.dense_b] - - def get_mlp_params(self): - return [self._h4h_w, self._h4h_b, self._4hh_w, self._4hh_b] diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/base_moe.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/base_moe.py deleted file mode 100644 index 4be1b849ba70da04b6b08ea011c27bbbcf96b8bf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/base_moe.py +++ /dev/null @@ -1,130 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Create a container object to save model-specific tensors using the policy file above. -from .base import * -from deepspeed import comm as dist -import deepspeed.ops.transformer as transformer_inference -from deepspeed.accelerator import get_accelerator - - -class BaseTransformerMoEContainer(BaseTransformerContainer): - - def __init__(self, **kwargs): - # Call the init function of the parent class to initialize the tensors and configs from parent class - super().__init__(**kwargs) - - self.num_experts = self.policy.get_num_experts() - self.ep_world_size = dist.get_world_size() - self.local_ep_size = 1 if self.num_experts < self.ep_world_size else self.num_experts // self.ep_world_size - - self.layer_norm_eps = self.config.layer_norm_eps if hasattr(self.config, 'layer_norm_eps') else 1e-12, - - # MoE models will have a list of mlp related tensors - self._h4h_w = [] - self._h4h_b = [] - self._4hh_w = [] - self._4hh_b = [] - - # Residual MoE needs extra parameters - self._res_h4h_w = None - self._res_h4h_b = None - self._res_4hh_w = None - self._res_4hh_b = None - self._res_coef = None - - def create_ds_model_config(self): - self.set_hidden_heads(*self.policy.get_hidden_heads()) - assert self.num_attention_heads % self.mp_size == 0,\ - "To run the model parallel across the GPUs, the attention_heads require to be divisible by the world_size!" +\ - "This is because the attention computation is partitioned evenly among the parallel GPUs." - - self.ds_model_config = transformer_inference.DeepSpeedMoEInferenceConfig( - hidden_size=self.hidden_size, - heads=self.num_attention_heads, - layer_norm_eps=self.layer_norm_eps, - fp16=self.fp16, - pre_layer_norm=self.pre_layer_norm, - mp_size=self.mp_size, - q_int8=self.quantize, - moe_experts=self.local_ep_size, - global_experts=self.num_experts, - mlp_type=self.config.moe.type, - scale_attn_by_inverse_layer_idx=self.scale_attn_by_inverse_layer_idx, - ) - - return self.ds_model_config - - def initialize_tensors(self): - # Set the tensors from policy (user module) to container (DS module) - self.set_attention(*self.policy.attention()) - self.set_mlp(self.config.moe.type) - self.set_layernorm(*self.policy.layernorm()) - - def set_mlp(self, config_moe_type): - if config_moe_type == 'standard': - self._h4h_w, self._h4h_b, \ - self._4hh_w, self._4hh_b = self.policy.mlp() - else: - self._h4h_w, self._h4h_b, self._4hh_w, \ - self._4hh_b, self._res_h4h_w, self._res_h4h_b, \ - self._res_4hh_w, self._res_4hh_b, \ - self._res_coef = self.policy.mlp(config_moe_type) - - def transpose(self): - self.transpose_attention() - self.transpose_mlp() - - if self.config.moe.type == 'residual': - self.transpose_residual() - - def transpose_mlp(self): - self._h4h_w = [self.transpose_impl(moe_w1.data) for moe_w1 in self._h4h_w] - self._4hh_w = [self.transpose_impl(moe_w1.data) for moe_w1 in self._4hh_w] - - def transpose_residual(self): - self._res_h4h_w.data = self.transpose_impl(self._res_h4h_w.data) - self._res_4hh_w.data = self.transpose_impl(self._res_4hh_w.data) - self._res_coef.data = self.transpose_impl(self._res_coef.data) - - def apply_tensor_parallelism(self, mp_replace): - # setup the new Attention module - self.attention_qkv_mp(mp_replace) - self.attention_o_mp(mp_replace) - - # quantize attention weights - self.attention_quantization() - - # setup the new MLP module - self.mlp_mp() - - def mlp_mp(self): - gpu_index = dist.get_rank() - for ep_index in range(self.local_ep_size): - # mlp inter - self.module.mlp[ep_index].inter_w.data = self._h4h_w[gpu_index * self.local_ep_size + ep_index].to( - get_accelerator().current_device_name()) - self.module.mlp[ep_index].inter_b.data = self._h4h_b[gpu_index * self.local_ep_size + ep_index].to( - get_accelerator().current_device_name()) - - # mlp output - self.module.mlp[ep_index].output_w.data = self._4hh_w[gpu_index * self.local_ep_size + ep_index].to( - get_accelerator().current_device_name()) - self.module.mlp[ep_index].output_b.data = self._4hh_b[gpu_index * self.local_ep_size + ep_index].to( - get_accelerator().current_device_name()) - - def copy_data_to_new_module(self): - self.module.attn_nw.data = self.attn_nw.to(get_accelerator().current_device_name()) - self.module.attn_nb.data = self.attn_nb.to(get_accelerator().current_device_name()) - - self.module.norm_w.data.copy_(self.input_nw.to(get_accelerator().current_device_name())) - self.module.norm_b.data.copy_(self.input_nb.to(get_accelerator().current_device_name())) - - if self.config.moe.type == 'residual': - self.module.res_mlp.inter_w.data = self._res_h4h_w.to(get_accelerator().current_device_name()) - self.module.res_mlp.inter_b.data = self._res_h4h_b.to(get_accelerator().current_device_name()) - self.module.res_mlp.output_w.data = self._res_4hh_w.to(get_accelerator().current_device_name()) - self.module.res_mlp.output_b.data = self._res_4hh_b.to(get_accelerator().current_device_name()) - self.module.res_coef.data = self._res_coef.to(get_accelerator().current_device_name()) diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/bert.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/bert.py deleted file mode 100644 index 20ae575f45144733a82b609eed21284746cb95d0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/bert.py +++ /dev/null @@ -1,93 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .base import * -from deepspeed.model_implementations.transformers.ds_bert import DeepSpeedBERTInference -import torch -from torch.nn.parameter import Parameter -from ..policy import TransformerPolicy - - -class DS_BERTContainer(BaseTransformerContainer): - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - # All model specific things should be defined here instead of the base class. - self.return_tuple = True - self.triangular_masking = False - self.use_triton = kwargs['config'].use_triton and deepspeed.HAS_TRITON - - def create_module(self, config=None): - _config = config if config is not None else self.ds_model_config - self.module = DeepSpeedBERTInference(_config, mp_group=self.mp_group) - self.module.config.scale_attention = self.scale_attention - return self.module - - -class HFBertLayerPolicy(TransformerPolicy): - - def __init__(self, client_module, inference=False): - super().__init__(inference, pre_attn_norm=False) - self.client_module = client_module - self.cuda_graph_supported = True - - if HFBertLayerPolicy._orig_layer_class is None: - try: - import transformers - HFBertLayerPolicy._orig_layer_class = [ - transformers.models.bert.modeling_bert.BertLayer, - transformers.models.roberta.modeling_roberta.RobertaLayer - ] - except: - HFBertLayerPolicy._orig_layer_class = None - - def get_hidden_heads(self): - if self.pre_attn_norm: - attention_layernorm = self.client_module.PostAttentionLayerNorm - else: - attention_layernorm = self.client_module.attention.output.LayerNorm - return self.client_module.attention.self.query.weight.shape[1], \ - self.client_module.attention.self.num_attention_heads, \ - attention_layernorm.eps, \ - DEFAULT_INTERMEDIATE_SIZE - - def attention(self, enable_training=False): - qw = self.client_module.attention.self.query.weight - qb = self.client_module.attention.self.query.bias - kw = self.client_module.attention.self.key.weight - kb = self.client_module.attention.self.key.bias - vw = self.client_module.attention.self.value.weight - vb = self.client_module.attention.self.value.bias - - qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) - qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) - - return qkvw, \ - qkvb, \ - self.client_module.attention.output.dense.weight, \ - self.client_module.attention.output.dense.bias, \ - - def mlp(self, enable_training=False): - if self.pre_attn_norm: - intermediate_ff = self.client_module.intermediate.dense_act - else: - intermediate_ff = self.client_module.intermediate.dense - - return intermediate_ff.weight, intermediate_ff.bias, \ - self.client_module.output.dense.weight, \ - self.client_module.output.dense.bias - - def layernorm(self): - if self.pre_attn_norm: - attention_layernorm = self.client_module.PostAttentionLayerNorm - transformer_layernorm = self.client_module.PreAttentionLayerNorm - else: - attention_layernorm = self.client_module.attention.output.LayerNorm - transformer_layernorm = self.client_module.output.LayerNorm - return attention_layernorm.weight, \ - attention_layernorm.bias, \ - transformer_layernorm.weight, \ - transformer_layernorm.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/bloom.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/bloom.py deleted file mode 100644 index 7a9b9ca2065bf969efba949588bc9ac4c95ec8dc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/bloom.py +++ /dev/null @@ -1,143 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .base import * -from .features.meta_tensor import MetaTensorContainer -from .features.hybrid_engine import HybridEngineContainer -from deepspeed.model_implementations.transformers.ds_bloom import DeepSpeedBloomInference -from ..policy import TransformerPolicy -from ..policy import transformer_param_names -from ..policy import maybe_copy - -from ..policy import maybe_get_lora - -supported_models = {None} - - -class DS_BloomContainer(MetaTensorContainer, HybridEngineContainer, BaseTransformerContainer): - - def __init__(self, **kwargs): - # Check transformers version, error if > 4.43.4 (breaks at 4.44.0) - from importlib.metadata import version - v_transformers = version('transformers') - vers = v_transformers.split('.') - major = int(vers[0]) - minor = int(vers[1]) - if major > 4 or (major == 4 and minor > 43): - import sys - sys.exit( - f"Transformers version {v_transformers} exceeds version 4.43.4! After transformers version 4.43.4, BLOOM inference with DeepSpeed is no longer supported." - ) - - super().__init__(**kwargs) - - # All model specific things should be defined here instead of the base class. - self.bigscience_bloom = True - self.triangular_masking = False - - def create_module(self, config=None): - _config = config if config is not None else self.ds_model_config - - self.module = DeepSpeedBloomInference(_config, mp_group=self.mp_group) - self.module.config.scale_attention = self.scale_attention - self.module.config.invert_mask = False - return self.module - - def attention_qkv_mp(self, mp_replace, reversed_dim=False): - self.module.attention.attn_qkvw = mp_replace.copy(self.module.attention.attn_qkvw, self.qkvw) - self.module.attention.attn_qkvb = mp_replace.copy(self.module.attention.attn_qkvb, self.qkvb) - - def get_lora_matched_pair(self): - """ - Necessary to implement for `HybridEngineContainer` - """ - fc1_lora, fc2_lora, qkv_lora, out_lora = self.get_lora_params() - ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (qkv_lora, self.qkvw), (out_lora, self.dense_w)] - return ret - - def set_lora_params(self): - """ - Necessary to implement for `HybridEngineContainer` - """ - self.lora_params = [ - maybe_get_lora(p) for p in [ - self.policy.client_module.mlp.dense_h_to_4h, self.policy.client_module.mlp.dense_4h_to_h, self.policy. - client_module.self_attention.query_key_value, self.policy.client_module.self_attention.dense - ] - ] - - def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): - param_names = ( - 'self_attention.query_key_value.weight', \ - 'self_attention.query_key_value.bias', \ - 'self_attention.dense.weight', \ - 'self_attention.dense.bias', \ - 'mlp.dense_h_to_4h.weight', \ - 'mlp.dense_h_to_4h.bias', \ - 'mlp.dense_4h_to_h.weight', \ - 'mlp.dense_4h_to_h.bias', \ - 'post_attention_layernorm.weight', \ - 'post_attention_layernorm.bias', \ - 'input_layernorm.weight', \ - 'input_layernorm.bias' - ) - for i in range(0, 2): - maybe_copy(module.attention, - sd, - weight_quantizer, - mp_replace, - transformer_param_names[i], - prefix + param_names[i], - qkv=True, - megatron_v2=self.policy.is_megatron_v2, - split_qkv=self.policy.split_qkv) - for i in range(2, 4): - maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i], - prefix + param_names[i]) - for i in range(4, 10): - maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i], - prefix + param_names[i]) - for i in range(10, 12): - maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i], prefix + param_names[i]) - - -class BLOOMLayerPolicy(TransformerPolicy): - _orig_layer_class = None - - def __init__(self, client_module, inference=True, use_load_prefix=True, split_qkv=False): - super().__init__(inference, linear_layer=True, use_load_prefix=use_load_prefix, split_qkv=split_qkv) - self.client_module = client_module - try: - import transformers - BLOOMLayerPolicy._orig_layer_class = transformers.models.bloom.modeling_bloom.BloomBlock - global supported_models - supported_models.update({transformers.models.bloom.modeling_bloom.BloomModel}) - except Exception as e: - print(f"WARNING! Setting BLOOMLayerPolicy._orig_layer_class to None due to Exception: {e}") - BLOOMLayerPolicy._orig_layer_class = None - - def get_hidden_heads(self): - return self.client_module.self_attention.hidden_size, \ - self.client_module.self_attention.num_heads, \ - self.client_module.input_layernorm.eps, \ - DEFAULT_INTERMEDIATE_SIZE - - def attention(self, enable_training=False): - return self.client_module.self_attention.query_key_value.weight, \ - self.client_module.self_attention.query_key_value.bias, \ - self.client_module.self_attention.dense.weight, \ - self.client_module.self_attention.dense.bias, - - def mlp(self, enable_training=False): - return self.client_module.mlp.dense_h_to_4h.weight, \ - self.client_module.mlp.dense_h_to_4h.bias, \ - self.client_module.mlp.dense_4h_to_h.weight, \ - self.client_module.mlp.dense_4h_to_h.bias - - def layernorm(self): - return self.client_module.post_attention_layernorm.weight, \ - self.client_module.post_attention_layernorm.bias, \ - self.client_module.input_layernorm.weight, \ - self.client_module.input_layernorm.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/clip.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/clip.py deleted file mode 100644 index afe4a76086d80ceec7dce15c9cfb2a9b8718b0c4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/clip.py +++ /dev/null @@ -1,73 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .base import * -from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference -import torch -from torch.nn.parameter import Parameter -from ..policy import TransformerPolicy - - -class DS_CLIPContainer(BaseTransformerContainer): - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - # All model specific things should be defined here instead of the base class. - - def create_module(self, config=None): - _config = config if config is not None else self.ds_model_config - self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) - self.module.config.scale_attention = self.scale_attention - return self.module - - -class HFCLIPLayerPolicy(TransformerPolicy): - - def __init__(self, client_module, inference=False): - super().__init__(inference, pre_attn_norm=True, scale_attention=True) - self.client_module = client_module - self.cuda_graph_supported = True - - if HFCLIPLayerPolicy._orig_layer_class is None: - try: - import transformers - HFCLIPLayerPolicy._orig_layer_class = transformers.models.clip.modeling_clip.CLIPEncoderLayer - except: - HFCLIPLayerPolicy._orig_layer_class = None - - def get_hidden_heads(self): - return self.client_module.self_attn.q_proj.weight.shape[1], \ - self.client_module.self_attn.num_heads, \ - self.client_module.layer_norm1.eps, \ - DEFAULT_INTERMEDIATE_SIZE - - def attention(self, enable_training=False): - qw = self.client_module.self_attn.q_proj.weight - qb = self.client_module.self_attn.q_proj.bias - kw = self.client_module.self_attn.k_proj.weight - kb = self.client_module.self_attn.k_proj.bias - vw = self.client_module.self_attn.v_proj.weight - vb = self.client_module.self_attn.v_proj.bias - - qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) - qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) - - return qkvw, \ - qkvb, \ - self.client_module.self_attn.out_proj.weight, \ - self.client_module.self_attn.out_proj.bias - - def mlp(self, enable_training=False): - return self.client_module.mlp.fc1.weight, \ - self.client_module.mlp.fc1.bias, \ - self.client_module.mlp.fc2.weight, \ - self.client_module.mlp.fc2.bias - - def layernorm(self): - return self.client_module.layer_norm2.weight, \ - self.client_module.layer_norm2.bias, \ - self.client_module.layer_norm1.weight, \ - self.client_module.layer_norm1.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/distil_bert.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/distil_bert.py deleted file mode 100644 index ecd0562438b5ac634cf8b4536fd3413d0f9ed9d8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/distil_bert.py +++ /dev/null @@ -1,82 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .base import * -from deepspeed.model_implementations.transformers.ds_bert import DeepSpeedBERTInference -import torch -from torch.nn.parameter import Parameter -from ..policy import TransformerPolicy - - -class DS_DistilBERTContainer(BaseTransformerContainer): - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - # All model specific things should be defined here instead of the base class. - self.triangular_masking = False - self.return_single_tuple = True - self.use_triton = kwargs['config'].use_triton and deepspeed.HAS_TRITON - - def create_module(self, config=None): - _config = config if config is not None else self.ds_model_config - self.module = DeepSpeedBERTInference(_config, mp_group=self.mp_group) - self.module.config.scale_attention = self.scale_attention - return self.module - - -class HFDistilBertLayerPolicy(TransformerPolicy): - _orig_layer_class = None - - def __init__(self, client_module, inference=False, preln=False): - super().__init__(inference) - self.client_module = client_module - self.preln = preln - self.cuda_graph_supported = True - if HFDistilBertLayerPolicy._orig_layer_class is None: - try: - import transformers - HFDistilBertLayerPolicy._orig_layer_class = [ - transformers.models.distilbert.modeling_distilbert.TransformerBlock, - ] - except: - HFDistilBertLayerPolicy._orig_layer_class = None - - def get_hidden_heads(self): - return self.client_module.attention.q_lin.weight.shape[1], \ - self.client_module.attention.n_heads, \ - self.client_module.sa_layer_norm.eps, \ - DEFAULT_INTERMEDIATE_SIZE - - def attention(self, enable_training=False): - qw = self.client_module.attention.q_lin.weight - qb = self.client_module.attention.q_lin.bias - kw = self.client_module.attention.k_lin.weight - kb = self.client_module.attention.k_lin.bias - vw = self.client_module.attention.v_lin.weight - vb = self.client_module.attention.v_lin.bias - - qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) - qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) - - return qkvw, \ - qkvb, \ - self.client_module.attention.out_lin.weight, \ - self.client_module.attention.out_lin.bias - - def mlp(self, enable_training=False): - intermediate_ff = self.client_module.ffn.lin1 - - return intermediate_ff.weight, intermediate_ff.bias, \ - self.client_module.ffn.lin2.weight, \ - self.client_module.ffn.lin2.bias - - def layernorm(self): - attention_layernorm = self.client_module.sa_layer_norm - transformer_layernorm = self.client_module.output_layer_norm - return attention_layernorm.weight, \ - attention_layernorm.bias, \ - transformer_layernorm.weight, \ - transformer_layernorm.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__init__.py deleted file mode 100644 index fc2eb2a65531e61bcea078c0035bd2ee6ee861a6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .gated_mlp import HybridGatedMLPContainer -from .megatron import MegatronContainer -from .meta_tensor import MetaTensorContainer -from .split_qkv import HybridSplitQKVContainer diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index f2e0fefb8cb1b38c9a720ca7a0b7144c35920640..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/gated_mlp.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/gated_mlp.cpython-310.pyc deleted file mode 100644 index e77e5a265dc7994eb429128f49256b1b3324c28a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/gated_mlp.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_engine.cpython-310.pyc deleted file mode 100644 index 6682782ba7778a60b6d7654a56370884d979e186..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_engine.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_megatron.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_megatron.cpython-310.pyc deleted file mode 100644 index 020a773cc8a8fcdfbd3808cabe9858109a70c93e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/hybrid_megatron.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/megatron.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/megatron.cpython-310.pyc deleted file mode 100644 index 927c08f789979738e6f26c52415f7f9a01c99f6f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/megatron.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/meta_tensor.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/meta_tensor.cpython-310.pyc deleted file mode 100644 index 50e60628f15732c3f9c079b1a887a89919da0ab8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/meta_tensor.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/split_qkv.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/split_qkv.cpython-310.pyc deleted file mode 100644 index cb0ffb5ad10ff593f4bdaf7dbd58d0c318694835..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/__pycache__/split_qkv.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/gated_mlp.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/gated_mlp.py deleted file mode 100644 index 24f0826db14ed08f373a2824f845ac6f9d7d9508..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/gated_mlp.py +++ /dev/null @@ -1,118 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import abstractmethod - -from .hybrid_engine import HybridEngineContainer - - -class HybridGatedMLPContainer(HybridEngineContainer): - """ - The HybridGatedMLPContainer supports models for which the first MLP layer - is represented with two separate weights, one for the activation function - and one for the gating function. - """ - - def set_mlp(self, _h4h_w, _h4h_b, _4hh_w, _4hh_b): - super().set_mlp(_h4h_w, _h4h_b, _4hh_w, _4hh_b) - self.set_mlp_gate() - - @abstractmethod - def set_mlp_gate(self): - """ - In `set_mlp_gate`, it is necessary to populate the following variables (where appropriate) - for the given model: - self.inter_up_w: inter up weight - self.inter_up_b: inter up bias - self.inter_gate_w: inter gate weight - self.inter_gate_b: inter gate bias - If the parameter does not exist in the original model, set the attribute to None. - """ - raise NotImplementedError("A set_mlp_gate() function must be defined in the model container \ - in order to set the unfused inter up and gate tensors.") - - def mlp_inter_mp(self, mp_replace, reversed_dim=False): - # Only need to alter behavior if we can't do the normal destructive copy - if self.module.mlp.inter_w is None: - params = [ - (self.module.mlp.inter_up_w, self.inter_up_w), - (self.module.mlp.inter_up_b, self.inter_up_b), - (self.module.mlp.inter_gate_w, self.inter_gate_w), - (self.module.mlp.inter_gate_b, self.inter_gate_b), - ] - for dst, src in params: - dst = mp_replace.copy(dst[:self.inter_up_w.shape[0] // mp_replace.mp_size], - src, - int8=reversed_dim, - allocate_tensor=reversed_dim) if src is not None else None - else: - self.module.mlp.inter_w = mp_replace.strided_copy(self.module.mlp.inter_w, - self._h4h_w, - num_splits=2, - int8=reversed_dim) - self.module.mlp.inter_b = mp_replace.strided_copy(self.module.mlp.inter_b, - self._h4h_b, - num_splits=2, - int8=reversed_dim) - - def release_mlp(self): - super().release_mlp() - gated_mlp_params = [ - (self.module.mlp.inter_up_w, self.inter_up_w), - (self.module.mlp.inter_up_b, self.inter_up_b), - (self.module.mlp.inter_gate_w, self.inter_gate_w), - (self.module.mlp.inter_gate_b, self.inter_gate_b), - ] - - self._release_params(gated_mlp_params) - - def reset_mlp(self): - self._h4h_w.data[:self.inter_up_w.shape[0]] = self.inter_up_w.data - self._h4h_w.data[self.inter_up_w.shape[0]:] = self.inter_gate_w.data - - if self.inter_up_b is not None: - self._h4h_b.data[:self.inter_up_b.shape[0]] = self.inter_up_b.data - self._h4h_b.data[self.inter_up_b.shape[0]:] = self.inter_gate_b.data - - inter_data = [self.inter_up_w.data, self.inter_gate_w.data] - if self.inter_up_b is not None: - inter_data.extend([self.inter_up_b.data, self.inter_gate_b.data]) - - self.inter_up_w.data = self._h4h_w.data[:self.inter_up_w.shape[0]] - self.inter_gate_w.data = self._h4h_w.data[self.inter_up_w.shape[0]:] - - if self.inter_up_b is not None: - self.inter_up_b.data = self._h4h_b.data[:self.inter_up_b.shape[0]] - self.inter_gate_b.data = self._h4h_b.data[self.inter_up_b.shape[0]:] - - for data in inter_data: - del data - - def set_mlp_params_wo_copy(self, Z3_enabled=False): - self.module.mlp.output_w = self._4hh_w - self.module.mlp.output_b = self._4hh_b - - if not Z3_enabled: - # In initialize_tensors, we create a fused inter projection with the appropriate shape - # and copy the up projection and gate projection into it - self.module.mlp.inter_w = self._h4h_w - self.module.mlp.inter_b = self._h4h_b - - self.inter_up_w.data = self._h4h_w[:self.inter_up_w.shape[0], :] - self.inter_gate_w.data = self._h4h_w[self.inter_up_w.shape[0]:, :] - - if self.inter_up_b is not None: - self.inter_up_b.data = self._h4h_b[:self.inter_up_w.shape[0]] if self._h4h_b is not None else None - self.inter_gate_b.data = self._h4h_b[self.inter_up_w.shape[0]:] if self._h4h_b is not None else None - else: - self.module.mlp.inter_up_w = self.inter_up_w - self.module.mlp.inter_up_b = self.inter_up_b - self.module.mlp.inter_gate_w = self.inter_gate_w - self.module.mlp.inter_gate_b = self.inter_gate_b - - def get_mlp_params(self): - params = super().get_mlp_params() - params.extend([self.inter_up_w, self.inter_up_b, self.inter_gate_w, self.inter_gate_b]) - return params diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/hybrid_engine.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/hybrid_engine.py deleted file mode 100644 index 3f702abcf319a6db997b6e2607dcba26f3440841..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/hybrid_engine.py +++ /dev/null @@ -1,212 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import ABC, abstractmethod -from typing import List, Tuple - -import torch - - -class HybridEngineContainer(ABC): - """ - This container identifies which methods need to be overridden in addition to - the base container to enable use in the RLHF pipeline. These methods are not - necessary for inference alone. - - NOTE: If you are using this feature with a container that - also inherits from `MetaTensorContainer`, ensure that `MetaTensorContainer` - is inherited before `HybridEngineContainer` in the class definition. - """ - - def initialize_tensors(self, enable_training=False): - """ - Same purposes as the base container, but also grabs the hooks for any LoRA - parameters. If it's necessary to override specific sub-components of the model, - it's best to augment the specific `set_[component]` itself rather than modifying - the `initialize_tensors` method. See the `HybridSplitQKVContainer` for an example. - """ - super().initialize_tensors(enable_training=enable_training) - self.set_lora_params() - - def transform_for_training(self): - """ - If the views on certain parameters are largely incompatible, it may be necessary to do - more substantial transformations to the parameters. This method should be overridden to - transform the inference format to what is necessary for training. - """ - pass - - def transform_for_inference(self): - """ - If the views on certain parameters are largely incompatible, it may be necessary to do - more substantial transformations to the parameters. This method should be overridden to - transform the training format to what is necessary for inference. - """ - pass - - @abstractmethod - def set_lora_params(self): - """ - If available, set the LoRA parameters for the module. An implementation - for this would iterate over all parameters of the model and use the `maybe_get_lora` helper - method to check if the parameter does in fact have any LoRA params. - """ - raise NotImplementedError("A set_lora_params() function must be defined for the relevant parameters.") - - @abstractmethod - def get_lora_matched_pair(self): - """Get the pair of lora params and its matched model parameters.""" - raise NotImplementedError("get_lora_matched_pair() must be defined for the relevant parameters.") - - def fuse_lora(self): - """Fuse the LoRA parameters for the inference mode.""" - for maybe_lora_param, param in self.get_lora_matched_pair(): - if len(maybe_lora_param) == 3: - lora_right_weight, \ - lora_left_weight, \ - lora_scaling = maybe_lora_param - param.data += lora_scaling * torch.matmul(lora_left_weight.t(), lora_right_weight.t()) - - def unfuse_lora(self): - """Unfuse the LoRA parameters for the training mode.""" - for maybe_lora_param, param in self.get_lora_matched_pair(): - if len(maybe_lora_param) == 3: - lora_right_weight, \ - lora_left_weight, \ - lora_scaling = maybe_lora_param - param.data -= lora_scaling * torch.matmul(lora_left_weight.t(), lora_right_weight.t()) - - def apply_tensor_parallelism(self, mp_replace, reversed_dim=False): - """ - Add support for reversed dim in tensor parallelism. If necessary, override - the called methods to handle partitioned weights (i.e. if qkv is split, override - the `attention_qkv_mp` method). If the model component is not split, it should - be safe to use the default implementation. - """ - # Setup the new Attention module - self.attention_qkv_mp(mp_replace, reversed_dim=reversed_dim) - self.attention_o_mp(mp_replace, reversed_dim=reversed_dim) - - # Setup the new MLP module - self.mlp_inter_mp(mp_replace, reversed_dim=reversed_dim) - self.mlp_output_mp(mp_replace, reversed_dim=reversed_dim) - - # Apply weight quantization - # TODO(cmikeh2): Re-enable this once verified - #self.apply_weight_quantization() - - def _release_params(self, param_pairs: List[Tuple[torch.Tensor, torch.Tensor]]): - """ - Helper for `release_[component]` methods. Accepts a list of tuples where the first - element is the module param that needs to be deleted, and the second is the reassignment - from the container. - """ - for module_param, container_param in param_pairs: - if module_param is not None: - del module_param - module_param = container_param - - def release_memory(self): - """ - Delete module parameters if they exist and point them back to the container. The primary - purpose of this is for TP-inference with ZeRO-3. In this scenario, we need to delete the - parameters we've created for inference to free their memory. - """ - general_params = [ - (self.module.attention.attn_ow, self.dense_w), - (self.module.attention.attn_ob, self.dense_b), - (self.module.mlp.attn_nw, self.attn_nw), - (self.module.mlp.attn_nb, self.attn_nb), - (self.module.norm_w, self.input_nw), - (self.module.norm_b, self.input_nb), - ] - - self._release_params(general_params) - - self.release_qkv() - self.release_mlp() - - def release_qkv(self): - """ - Release for QKV parameters (as well as any aliases). - """ - qkv_params = [ - (self.module.attention.attn_qkvw, self.qkvw), - (self.module.attention.attn_qkvb, self.qkvb), - ] - - self._release_params(qkv_params) - - def release_mlp(self): - """ - Release for MLP parameters (as well as any aliases). - """ - mlp_params = [ - (self.module.mlp.inter_w, self._h4h_w), - (self.module.mlp.inter_b, self._h4h_b), - (self.module.mlp.output_w, self._4hh_w), - (self.module.mlp.output_b, self._4hh_b), - ] - - self._release_params(mlp_params) - - def reset_params(self): - """ - The purpose of reset params is to get the weights from the FP16 training - copy of the model and copy to them to contiguous inference view. This only needs - to be performed when the container parameters cannot be used directly for inference. - """ - self.reset_qkv() - self.reset_mlp() - - def reset_qkv(self): - """ - Perform any necessary resets of the model parameters for the QKV components. - """ - pass - - def reset_mlp(self): - """ - Perform any necessary resets of the model parameters for the MLP components. - """ - pass - - def get_lora_params(self): - """ - Return a list of all parameters that would have LoRA for the module. - """ - if not hasattr(self, "lora_params"): - self.set_lora_params() - return self.lora_params - - def set_params_wo_copy(self, Z3_enabled=False): - """ - Rather than copying into, set the parameters directly. This is necessary to provide - an inexpensive (low-memory-overhead) view onto the FP16 forward weights. - """ - self.module.mlp.attn_nw = self.attn_nw - self.module.mlp.attn_nb = self.attn_nb - self.module.norm_w = self.input_nw - self.module.norm_b = self.input_nb - self.set_attn_params_wo_copy(Z3_enabled=Z3_enabled) - self.set_mlp_params_wo_copy(Z3_enabled=Z3_enabled) - - def set_attn_params_wo_copy(self, **kwargs): - """ - Narrower sub-method for finer grained overriding. - """ - self.module.attention.attn_ow = self.dense_w - self.module.attention.attn_ob = self.dense_b - self.module.attention.attn_qkvw = self.qkvw - self.module.attention.attn_qkvb = self.qkvb - - def set_mlp_params_wo_copy(self, **kwargs): - """ - Narrower sub-method for finer grained overriding. - """ - self.module.mlp.inter_w = self._h4h_w - self.module.mlp.inter_b = self._h4h_b - self.module.mlp.output_w = self._4hh_w - self.module.mlp.output_b = self._4hh_b diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/hybrid_megatron.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/hybrid_megatron.py deleted file mode 100644 index d40f2a6b57e86087dc13ff2b74a389c67826fa76..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/hybrid_megatron.py +++ /dev/null @@ -1,87 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from .hybrid_engine import HybridEngineContainer -from .megatron import MegatronContainer - - -class HybridMegatronContainer(MegatronContainer, HybridEngineContainer): - - def _align_qkv(self, x: torch.Tensor): - """ - Internal helper for accepting the head-contiguous weight matrix and chunking - the query, key, and value components. - """ - attention_head_size = x.shape[0] // self.num_attention_heads - new_x_shape = (self.num_attention_heads, attention_head_size) + x.size()[1:] - x_1 = x.view(*new_x_shape) - div_dim = len(x_1.size()) - 2 if len(x.shape) == 2 else -1 - (q, k, v) = torch.split(x_1, (x_1.shape[div_dim] // 3), dim=div_dim) - if len(q.shape) > 2: - x.data.copy_( - torch.cat((q.reshape(-1, q.shape[-1]), k.reshape(-1, q.shape[-1]), v.reshape(-1, q.shape[-1])), - dim=0).reshape(x.shape)) - else: - x.data.copy_(torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape)) - - def transform_for_inference(self) -> None: - """ - Overrides the HybridEngineContainer implementation. - - The alternative layout of the QKV matrix for Megatron is such that each head's Q, K, and V - are sequential in memory. This is different from the default layout in which all of the Qs - are sequential, followed by all of the Ks, and then all of the Vs. Here, we take the default - layout and transform it to the inference layout. - """ - if hasattr(self.qkvw, 'ds_id'): - from deepspeed.runtime.zero import GatheredParameters - from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus - param_list = [self.qkvw, self.qkvb] - non_active_params = [param for param in param_list if (hasattr(param, 'ds_id') and \ - param.ds_status == ZeroParamStatus.NOT_AVAILABLE)] - with GatheredParameters(non_active_params): - self._align_qkv(self.qkvw) - self._align_qkv(self.qkvb) - else: - self._align_qkv(self.qkvw) - self._align_qkv(self.qkvb) - - def _partition_qkv(self, x: torch.Tensor): - """ - Internal helper for taking contiguous QKV and partitioning it for contiguous - heads. - """ - q_k_v = torch.split(x, (x.shape[0] // 3), dim=0) - attention_head_size = q_k_v[0].shape[0] // self.num_attention_heads - new_x_shape = (self.num_attention_heads, attention_head_size) + x.size()[1:] - q, k, v = [data.view(*new_x_shape) for data in q_k_v] - if len(q.shape) > 2: - x.data.copy_(torch.cat((q, k, v), dim=-2).reshape(-1, q.shape[-1])) - else: - x.data.copy_(torch.cat((q, k, v), dim=-1).reshape(-1)) - - def transform_for_training(self): - """ - Overrides the HybridEngineContainer implementation. - - The alternative layout of the QKV matrix for Megatron is such that each head's Q, K, and V - are sequential in memory. This is different from the default layout in which all of the Qs - are sequential, followed by all of the Ks, and then all of the Vs. This function takes the inference format and reverts it back to the default format. - """ - # If parameter is distributed, handle gathering it - if hasattr(self.qkvw, 'ds_id'): - from deepspeed.runtime.zero import GatheredParameters - from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus - param_list = [self.qkvw, self.qkvb] - non_active_params = [param for param in param_list if (hasattr(param, 'ds_id') and \ - param.ds_status == ZeroParamStatus.NOT_AVAILABLE)] - with GatheredParameters(non_active_params): - self._partition_qkv(self.qkvw) - self._partition_qkv(self.qkvb) - else: - self._partition_qkv(self.qkvw) - self._partition_qkv(self.qkvb) diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/megatron.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/megatron.py deleted file mode 100644 index 4daccf7d7c8d4db06bd65243bd83a5f1fd50df2e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/megatron.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from abc import ABC - - -class MegatronContainer(ABC): - - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.megatron_v2 = self.policy.is_megatron_v2 - - def _align_qkv_transposed(self, x): - attention_head_size = x.shape[-1] // self.num_attention_heads - new_x_shape = x.size()[:-1] + (self.num_attention_heads, attention_head_size) - x_1 = x.view(*new_x_shape) - (q, k, v) = torch.split(x_1, (x_1.shape[-1] // 3), dim=(x_1.dim() - 1)) - if len(q.shape) > 2: - return torch.cat((q.reshape(q.shape[0], -1), k.reshape(q.shape[0], -1), v.reshape(q.shape[0], -1)), - dim=-1).reshape(x.shape) - else: - return torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape) - - def transpose(self): - super().transpose() - if self.megatron_v2: - self.qkvw = torch.nn.parameter.Parameter(self._align_qkv_transposed(self.qkvw).contiguous()) - self.qkvb = torch.nn.parameter.Parameter(self._align_qkv_transposed(self.qkvb).contiguous()) diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/meta_tensor.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/meta_tensor.py deleted file mode 100644 index 57b136663be370863dad81cc9d6a4a29425376b3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/meta_tensor.py +++ /dev/null @@ -1,70 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import ABC, abstractmethod -from packaging import version as pkg_version -import torch - - -class MetaTensorContainer(ABC): - """ - NOTE: If you are using this feature with a container that - also inherits from `HybridEngineContainer`, ensure that `MetaTensorContainer` - is inherited before `HybridEngineContainer` in the class definition. - """ - - def __init__(self, **kwargs): - if pkg_version.parse('1.10') > pkg_version.parse(torch.__version__): - raise NotImplementedError("Meta tensor support is not available, please upgrade to torch 1.10+") - super().__init__(**kwargs) - self.is_meta = False - self.ckpt_load_enabled = True - - def initialize_tensors(self, enable_training=False): - super().initialize_tensors(enable_training=enable_training) - self.is_meta = self.qkvw.is_meta - - def apply_tensor_parallelism(self, mp_replace, **kwargs): - if self.is_meta: - if self.qkvb is None: - self.module.attention.attn_qkvb = None - if self.dense_b is None: - self.module.attention.attn_ob = None - else: - super().apply_tensor_parallelism(mp_replace, **kwargs) - - def copy_data_to_new_module(self): - if self.is_meta: - if self.attn_nw is None: - self.module.mlp.attn_nw = self.attn_nw - self.module.mlp.attn_nb = self.attn_nb - else: - super().copy_data_to_new_module() - - def transpose(self): - if not self.is_meta: - super().transpose() - - @abstractmethod - def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): - """ - Load all the transformer parameter from the checkpoint file (sd). - In addition to the parameter names, we require two - more parameters to help read the data correctly - from the checkpoint and split the qkv heads in the - right order: - 1. `use_load_prefix` (Default: False): this specifies - whether we need to use the name of first abstraction - layer of the model for searching the parameter's name - in a checkpoint file. For more information of how this - is used please see - https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/module_inject/load_checkpoint.py - 2. `split_qkv` (Default: True): we use this flag when splitting - the qkv parameter into heads. If it is False, it means the heads - of q, k, and v are stored together and needs to split in the - DeepSpeed-Inference API. - """ - raise NotImplementedError("A load_params() function must be defined in the model container \ - when inheriting the MetaTensorContainer feature") diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/split_qkv.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/split_qkv.py deleted file mode 100644 index f4c14d4e425a7e9096bae3bf5788ab7b9ca0dcc2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/features/split_qkv.py +++ /dev/null @@ -1,159 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import abstractmethod -import torch - -from .hybrid_engine import HybridEngineContainer - - -class HybridSplitQKVContainer(HybridEngineContainer): - - def set_attention(self, qkvw, qkvb, dense_w, dense_b): - super().set_attention(qkvw, qkvb, dense_w, dense_b) - self.set_q_k_v() - - @abstractmethod - def set_q_k_v(self): - """ - In `set_q_k_v`, it is necessary to populate the following variables (where appropriate) - for the given model: - self.qw: q weight - self.qb: q bias - self.kw: k weight - self.kb: k bias - self.vw: v weight - self.vb: v bias - """ - raise NotImplementedError("A set_q_k_v() function must be defined in the model container \ - in order to set the unfused q, k, and v tensors.") - - def attention_qkv_mp(self, mp_replace, reversed_dim=False): - # Only need to alter - if self.module.attention.attn_qkvw is None: - params = [ - (self.module.attention.attn_qw, self.qw), - (self.module.attention.attn_qb, self.qb), - (self.module.attention.attn_kw, self.kw), - (self.module.attention.attn_kb, self.kb), - (self.module.attention.attn_vw, self.vw), - (self.module.attention.attn_vb, self.vb), - ] - for dst, src in params: - dst = mp_replace.copy( - dst[:self.qw.shape[0] // mp_replace.mp_size], src, int8=reversed_dim, - allocate_tensor=reversed_dim) if src is not None else None - else: - super().attention_qkv_mp(mp_replace) - - def release_qkv(self): - super().release_qkv() - split_qkv_params = [ - (self.module.attention.attn_qw, self.qw), - (self.module.attention.attn_qb, self.qb), - (self.module.attention.attn_kw, self.kw), - (self.module.attention.attn_kb, self.kb), - (self.module.attention.attn_vw, self.vw), - (self.module.attention.attn_vb, self.vb), - ] - - self._release_params(split_qkv_params) - - def reset_qkv(self): - self.qkvw.data[:self.qw.shape[0]] = self.qw.data - self.qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kw.data - self.qkvw.data[2 * self.qw.shape[0]:] = self.vw.data - - qkv_data = [self.qw.data, self.kw.data, self.vw.data] - - self.qw.data = self.qkvw.data[:self.qw.shape[0]] - self.kw.data = self.qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] - self.vw.data = self.qkvw.data[2 * self.qw.shape[0]:] - - if self.qkvb is not None: - self.qkvb.data[:self.qw.shape[0]] = self.qb.data - self.qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kb.data - self.qkvb.data[2 * self.qw.shape[0]:] = self.vb.data - - qkv_data.extend([self.qb.data, self.kb.data, self.vb.data]) - - self.qb.data = self.qkvb.data[:self.qw.shape[0]] - self.kb.data = self.qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] - self.vb.data = self.qkvb.data[2 * self.qw.shape[0]:] - - for data in qkv_data: - del data - - def reset_qkv_experimental(self): - """ - WIP - experimental and likely to be changed/improved. - Unused by keeping for now. - """ - if self.module.attention.attn_qkvw is None: - self.module.attention.attn_qkvw = torch.empty(self.qw.shape[0] * 3, - self.qw.shape[0], - dtype=self.qw.dtype, - device=self.qw.device) - self.module.attention.attn_qkvb = torch.empty(self.qw.shape[0] * 3, - dtype=self.qw.dtype, - device=self.qw.device) - self.module.attention.attn_qkvw.data[:self.qw.shape[0]] = self.qw.data - self.module.attention.attn_qkvb.data[:self.qw.shape[0]] = self.qb.data - self.module.attention.attn_qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kw.data - self.module.attention.attn_qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kb.data - self.module.attention.attn_qkvw.data[2 * self.qw.shape[0]:] = self.vw.data - self.module.attention.attn_qkvb.data[2 * self.qw.shape[0]:] = self.vb.data - - qkv_data = [self.qw.data, \ - self.qb.data, \ - self.kw.data, \ - self.kb.data, \ - self.vw.data, \ - self.vb.data] - - self.qw.data = self.module.attention.attn_qkvw.data[:self.qw.shape[0]] - self.qb.data = self.module.attention.attn_qkvb.data[:self.qw.shape[0]] - self.kw.data = self.module.attention.attn_qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] - self.kb.data = self.module.attention.attn_qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] - self.vw.data = self.module.attention.attn_qkvw.data[2 * self.qw.shape[0]:] - self.vb.data = self.module.attention.attn_qkvb.data[2 * self.qw.shape[0]:] - - for data in qkv_data: - del data - - def set_attn_params_wo_copy(self, Z3_enabled=False): - self.module.attention.attn_ow = self.dense_w - self.module.attention.attn_ob = self.dense_b - if not Z3_enabled: - # In initialize_tensors, we create a fused qkvw with the appropriate shape - # and copy the qw, qb, kw, kb, vw, vb into it - self.module.attention.attn_qkvw = self.qkvw - self.module.attention.attn_qkvb = self.qkvb - - # We reset the data for qw (which is the original model parameter) to point - # to the fused weight matrix we have created here - self.qw.data = self.qkvw[:self.qw.shape[0], :] - self.kw.data = self.qkvw[self.qw.shape[0]:2 * self.qw.shape[0], :] - self.vw.data = self.qkvw[self.qw.shape[0] * 2:, :] - - # Assume if one of the biases is not None, then all of them are not None - if self.qb is not None: - self.qb.data = self.qkvb[:self.qw.shape[0]] - self.kb.data = self.qkvb[self.qw.shape[0]:2 * self.qw.shape[0]] - self.vb.data = self.qkvb[self.qw.shape[0] * 2:] - else: - # In ZeRO-3 this will be managed by ZeRO and handled separately in the - # forward of ds_attention - self.module.attention.attn_qw = self.qw - self.module.attention.attn_qb = self.qb - self.module.attention.attn_kw = self.kw - self.module.attention.attn_kb = self.kb - self.module.attention.attn_vw = self.vw - self.module.attention.attn_vb = self.vb - - def get_attn_params(self): - params = super().get_attn_params() - params.extend([self.qw, self.qb, self.kw, self.kb, self.vw, self.vb]) - return params diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gpt2.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gpt2.py deleted file mode 100644 index 7a19aac34b447a7d5116ccceeb7687faa61ea174..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gpt2.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .base import * -from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference -from ..policy import TransformerPolicy - - -class DS_GPT2Container(BaseTransformerContainer): - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - # All model specific things should be defined here instead of the base class. - - def create_module(self, config=None): - _config = config if config is not None else self.ds_model_config - self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) - self.module.config.scale_attention = self.scale_attention - return self.module - - -class HFGPT2LayerPolicy(TransformerPolicy): - _orig_layer_class = None - - def __init__(self, client_module, inference=True): - # HuggingFace GPT2 uses convolutional layer instead of linear layer - super().__init__(inference, linear_layer=False) - self.client_module = client_module - try: - import transformers - HFGPT2LayerPolicy._orig_layer_class = transformers.models.gpt2.modeling_gpt2.GPT2Block - except: - HFGPT2LayerPolicy._orig_layer_class = None - - def get_hidden_heads(self): - return self.client_module.attn.embed_dim, \ - self.client_module.attn.num_heads, \ - self.client_module.ln_1.eps, \ - DEFAULT_INTERMEDIATE_SIZE - - def attention(self, enable_training=False): - return self.client_module.attn.c_attn.weight, \ - self.client_module.attn.c_attn.bias, \ - self.client_module.attn.c_proj.weight, \ - self.client_module.attn.c_proj.bias - - def mlp(self, enable_training=False): - return self.client_module.mlp.c_fc.weight, \ - self.client_module.mlp.c_fc.bias, \ - self.client_module.mlp.c_proj.weight, \ - self.client_module.mlp.c_proj.bias - - def layernorm(self): - return self.client_module.ln_2.weight, \ - self.client_module.ln_2.bias, \ - self.client_module.ln_1.weight, \ - self.client_module.ln_1.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptj.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptj.py deleted file mode 100644 index 17c0a5027a4c94dcc206b69ad6f3fa6d1d746e17..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptj.py +++ /dev/null @@ -1,132 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .base import * -from .features.meta_tensor import MetaTensorContainer -from .features.split_qkv import HybridSplitQKVContainer -from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference -import torch -from torch.nn.parameter import Parameter -from ..policy import TransformerPolicy -from ..policy import transformer_param_names -from ..policy import maybe_copy -from ..policy import maybe_copy_qkv - -from ..policy import maybe_get_lora - - -class DS_GPTJContainer(MetaTensorContainer, HybridSplitQKVContainer, BaseTransformerContainer): - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - # All model specific things should be defined here instead of the base class. - - def create_module(self, config=None): - _config = config if config is not None else self.ds_model_config - self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) - self.module.config.scale_attention = self.scale_attention - return self.module - - def set_lora_params(self): - """ - Necessary to implement for `HybridEngineContainer` - """ - self.lora_params = [ - maybe_get_lora(p) for p in [ - self.policy.client_module.mlp.fc_in, self.policy.client_module.mlp.fc_out, - self.policy.client_module.attn.q_proj, self.policy.client_module.attn.k_proj, - self.policy.client_module.attn.v_proj, self.policy.client_module.attn.out_proj - ] - ] - - def get_lora_matched_pair(self): - fc1_lora, fc2_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() - ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (out_lora, self.dense_w), (q_lora, self.qw), - (k_lora, self.kw), (v_lora, self.vw)] - return ret - - def set_q_k_v(self): - """ - Necessary to implement for `HybridSplitQKVContainer` - """ - self.qw = self.policy.client_module.attn.q_proj.weight - self.qb = None - self.kw = self.policy.client_module.attn.k_proj.weight - self.kb = None - self.vw = self.policy.client_module.attn.v_proj.weight - self.vb = None - - def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): - param_names = ( - 'attn.q_proj.weight', \ - 'attn.k_proj.weight', \ - 'attn.v_proj.weight', \ - 'attn.out_proj.weight', \ - 'mlp.fc_in.weight', \ - 'mlp.fc_in.bias', \ - 'mlp.fc_out.weight', \ - 'mlp.fc_out.bias', \ - 'ln_1.weight', \ - 'ln_1.bias' - ) - maybe_copy_qkv(module.attention, - sd, - weight_quantizer, - mp_replace, - 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], - split_qkv=self.policy.split_qkv) - for i in range(3, 4): - maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], - prefix + param_names[i]) - for i in range(4, 8): - maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i], - prefix + param_names[i]) - for i in range(8, 10): - maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i + 2], - prefix + param_names[i]) - - -class HFGPTJLayerPolicy(TransformerPolicy): - _orig_layer_class = None - - def __init__(self, client_module, inference=True): - super().__init__(inference, scale_attention=True) - self.client_module = client_module - try: - import transformers - HFGPTJLayerPolicy._orig_layer_class = transformers.models.gptj.modeling_gptj.GPTJBlock - except: - HFGPTJLayerPolicy._orig_layer_class = None - - def get_hidden_heads(self): - return self.client_module.attn.embed_dim, \ - self.client_module.attn.num_attention_heads, \ - self.client_module.ln_1.eps, \ - DEFAULT_INTERMEDIATE_SIZE - - def attention(self, enable_training=False): - qw = self.client_module.attn.q_proj.weight - kw = self.client_module.attn.k_proj.weight - vw = self.client_module.attn.v_proj.weight - - qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) - - return qkvw, \ - None, \ - self.client_module.attn.out_proj.weight, \ - None, - - def mlp(self, enable_training=False): - return self.client_module.mlp.fc_in.weight, \ - self.client_module.mlp.fc_in.bias, \ - self.client_module.mlp.fc_out.weight, \ - self.client_module.mlp.fc_out.bias - - def layernorm(self): - return None, \ - None, \ - self.client_module.ln_1.weight, \ - self.client_module.ln_1.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptneo.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptneo.py deleted file mode 100644 index fca673b375e18a9b74741f86a9692bfd386f6a51..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptneo.py +++ /dev/null @@ -1,145 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .base import * -from .features.meta_tensor import MetaTensorContainer -from .features.split_qkv import HybridSplitQKVContainer -from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference -import torch -from torch.nn.parameter import Parameter -from ..policy import TransformerPolicy -from ..policy import transformer_param_names -from ..policy import maybe_copy -from ..policy import maybe_copy_qkv - -from ..policy import maybe_get_lora - - -class DS_GPTNEOContainer(MetaTensorContainer, HybridSplitQKVContainer, BaseTransformerContainer): - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - # All model specific things should be defined here instead of the base class. - - def create_module(self, config=None): - _config = config if config is not None else self.ds_model_config - self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) - self.module.config.scale_attention = self.scale_attention - return self.module - - def set_lora_params(self): - """ - Necessary to implement for `HybridEngineContainer` - """ - self.lora_params = [ - maybe_get_lora(p) for p in [ - self.policy.client_module.mlp.c_fc, self.policy.client_module.mlp.c_proj, - self.policy.client_module.attn.attention.q_proj, self.policy.client_module.attn.attention.k_proj, - self.policy.client_module.attn.attention.v_proj, self.policy.client_module.attn.attention.out_proj - ] - ] - - def set_q_k_v(self): - """ - Necessary to implement for `HybridSplitQKVContainer` - """ - self.qw = self.policy.client_module.attn.attention.q_proj.weight - self.qb = None - self.kw = self.policy.client_module.attn.attention.k_proj.weight - self.kb = None - self.vw = self.policy.client_module.attn.attention.v_proj.weight - self.vb = None - - def get_lora_matched_pair(self): - """ - Necessary to implement for `HybridEngineContainer` - """ - fc1_lora, fc2_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() - ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (out_lora, self.dense_w), (q_lora, self.qw), - (k_lora, self.kw), (v_lora, self.vw)] - return ret - - def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): - param_names = ( - 'attn.attention.q_proj.weight', \ - 'attn.attention.k_proj.weight', \ - 'attn.attention.v_proj.weight', \ - 'attn.attention.out_proj.weight', \ - 'attn.attention.out_proj.bias', \ - 'mlp.c_fc.weight', \ - 'mlp.c_fc.bias', \ - 'mlp.c_proj.weight', \ - 'mlp.c_proj.bias', \ - 'ln_2.weight', \ - 'ln_2.bias', \ - 'ln_1.weight', \ - 'ln_1.bias' - ) - maybe_copy_qkv(module.attention, - sd, - weight_quantizer, - mp_replace, - 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], - split_qkv=self.policy.split_qkv) - for i in range(3, 5): - maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], - prefix + param_names[i]) - for i in range(5, 11): - maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], - prefix + param_names[i]) - for i in range(11, 13): - maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], - prefix + param_names[i]) - - -class HFGPTNEOLayerPolicy(TransformerPolicy): - - def __init__(self, client_module, inference=True): - super().__init__(inference, scale_attention=False) - self.client_module = client_module - try: - import transformers - HFGPTNEOLayerPolicy._orig_layer_class = transformers.models.gpt_neo.modeling_gpt_neo.GPTNeoBlock - except: - HFGPTNEOLayerPolicy._orig_layer_class = None - - def get_hidden_heads(self): - return self.client_module.attn.attention.embed_dim, \ - self.client_module.attn.attention.num_heads, \ - self.client_module.ln_1.eps, \ - DEFAULT_INTERMEDIATE_SIZE - - def get_q_k_v(self): - return self.client_module.attn.attention.q_proj.weight, \ - None, \ - self.client_module.attn.attention.k_proj.weight, \ - None, \ - self.client_module.attn.attention.v_proj.weight, \ - None - - def attention(self, enable_training=False): - qw = self.client_module.attn.attention.q_proj.weight - kw = self.client_module.attn.attention.k_proj.weight - vw = self.client_module.attn.attention.v_proj.weight - - qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) - - return qkvw, \ - None, \ - self.client_module.attn.attention.out_proj.weight, \ - self.client_module.attn.attention.out_proj.bias - - def mlp(self, enable_training=False): - return self.client_module.mlp.c_fc.weight, \ - self.client_module.mlp.c_fc.bias, \ - self.client_module.mlp.c_proj.weight, \ - self.client_module.mlp.c_proj.bias - - def layernorm(self): - return self.client_module.ln_2.weight, \ - self.client_module.ln_2.bias, \ - self.client_module.ln_1.weight, \ - self.client_module.ln_1.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptneox.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptneox.py deleted file mode 100644 index 3978b166b329b341232da61811ecaf5683c5578d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/gptneox.py +++ /dev/null @@ -1,146 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .base import * -from .features.meta_tensor import MetaTensorContainer -from .features.hybrid_megatron import HybridMegatronContainer -from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference -import torch -from ..policy import TransformerPolicy -from ..policy import transformer_param_names -from ..policy import maybe_copy -from packaging import version as pkg_version - -from ..policy import maybe_get_lora - - -class DS_GPTNEOXContainer(MetaTensorContainer, HybridMegatronContainer, BaseTransformerContainer): - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - # All model specific things should be defined here instead of the base class. - - def create_module(self, config=None): - _config = config if config is not None else self.ds_model_config - self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) - self.module.config.scale_attention = self.scale_attention - - if self.megatron_v2: - self.module.config.rotate_half = True - self.module.config.rotate_every_two = False - - return self.module - - def get_lora_matched_pair(self): - """ - Necessary to implement for `HybridEngineContainer` - """ - fc1_lora, fc2_lora, qkv_lora, out_lora = self.get_lora_params() - ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (qkv_lora, self.qkvw), (out_lora, self.dense_w)] - return ret - - def set_lora_params(self): - """ - Necessary to implement for `HybridEngineContainer` - """ - if GPTNEOXLayerPolicy.version == 0: - attention = self.policy.client_module.attention - else: - attention = self.policy.client_module.self_attention - - self.lora_params = [ - maybe_get_lora(p) for p in [ - self.policy.client_module.mlp.dense_h_to_4h, self.policy.client_module.mlp.dense_4h_to_h, - attention.query_key_value, attention.dense - ] - ] - - def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): - param_names = ( - 'attention.query_key_value.weight', \ - 'attention.query_key_value.bias', \ - 'attention.dense.weight', \ - 'attention.dense.bias', \ - 'mlp.dense_h_to_4h.weight', \ - 'mlp.dense_h_to_4h.bias', \ - 'mlp.dense_4h_to_h.weight', \ - 'mlp.dense_4h_to_h.bias', \ - 'post_attention_layernorm.weight', \ - 'post_attention_layernorm.bias', \ - 'input_layernorm.weight', \ - 'input_layernorm.bias' - ) - for i in range(0, 2): - maybe_copy(module.attention, - sd, - weight_quantizer, - mp_replace, - transformer_param_names[i], - prefix + param_names[i], - qkv=True, - megatron_v2=self.policy.is_megatron_v2, - split_qkv=self.policy.split_qkv, - heads=self.policy.client_module.attention.num_attention_heads) - for i in range(2, 4): - maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i], - prefix + param_names[i]) - for i in range(4, 10): - maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i], - prefix + param_names[i]) - for i in range(10, 12): - maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i], prefix + param_names[i]) - - -class GPTNEOXLayerPolicy(TransformerPolicy): - _orig_layer_class = None - version = 0 - - def __init__(self, client_module, inference=True, megatron_v2=True, split_qkv=False): - super().__init__(inference, megatron_v2=megatron_v2, split_qkv=split_qkv) - self.client_module = client_module - if GPTNEOXLayerPolicy._orig_layer_class is None: - if pkg_version.parse(torch.__version__) <= pkg_version.parse("1.2"): - GPTNEOXLayerPolicy._orig_layer_class = None - else: - try: - from transformers import GPTNeoXLayer - GPTNEOXLayerPolicy._orig_layer_class = GPTNeoXLayer - except ImportError: - GPTNEOXLayerPolicy._orig_layer_class = None - - def get_hidden_heads(self): - if GPTNEOXLayerPolicy.version == 0: - attention = self.client_module.attention - else: - attention = self.client_module.self_attention - - return self.client_module.attention.hidden_size, \ - self.client_module.attention.num_attention_heads, \ - self.client_module.input_layernorm.eps, \ - DEFAULT_INTERMEDIATE_SIZE - - def attention(self, enable_training=False): - if GPTNEOXLayerPolicy.version == 0: - attention = self.client_module.attention - else: - attention = self.client_module.self_attention - - return attention.query_key_value.weight, \ - attention.query_key_value.bias, \ - attention.dense.weight, \ - attention.dense.bias - - def mlp(self, enable_training=False): - return self.client_module.mlp.dense_h_to_4h.weight, \ - self.client_module.mlp.dense_h_to_4h.bias, \ - self.client_module.mlp.dense_4h_to_h.weight, \ - self.client_module.mlp.dense_4h_to_h.bias - - def layernorm(self): - return self.client_module.post_attention_layernorm.weight, \ - self.client_module.post_attention_layernorm.bias, \ - self.client_module.input_layernorm.weight, \ - self.client_module.input_layernorm.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/internlm.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/internlm.py deleted file mode 100644 index 31255d4b3ca5a0199ee572219b6d75fd70d478c3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/internlm.py +++ /dev/null @@ -1,181 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import importlib - -import torch -from torch.nn.parameter import Parameter - -from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference -from deepspeed.utils.types import ActivationFuncType, NormType - -from ..policy import (TransformerPolicy, maybe_copy, maybe_copy_geglu, maybe_copy_qkv, maybe_get_lora, - transformer_param_names) -from .base import * -from .features import HybridGatedMLPContainer, HybridSplitQKVContainer - - -class DS_InternLMContainer(HybridGatedMLPContainer, HybridSplitQKVContainer, BaseTransformerContainer): - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - # All model specific things should be defined here instead of the base class. - - def create_module(self, config=None): - _config = config if config is not None else self.ds_model_config - - _config.rotate_half = True - _config.rotate_every_two = False - _config.rotary_dim = self.hidden_size // self.num_attention_heads - self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) - - return self.module - - def set_lora_params(self): - """ - Necessary to implement for `HybridEngineContainer` - """ - self.lora_params = [ - maybe_get_lora(p) for p in [ - self.policy.client_module.mlp.up_proj.weight, self.policy.client_module.mlp.gate_proj.weight, - self.policy.client_module.mlp.down_proj.weight, self.policy.client_module.self_attn.q_proj.weight, - self.policy.client_module.self_attn.k_proj.weight, self.policy.client_module.self_attn.v_proj.weight, - self.policy.client_module.self_attn.o_proj.weight - ] - ] - - def get_lora_matched_pair(self): - up_proj_lora, gate_proj_lora, down_proj_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() - ret = [(up_proj_lora, self.inter_up_w), (gate_proj_lora, self.inter_gate_w), (down_proj_lora, self._4hh_w), - (out_lora, self.dense_w), (q_lora, self.qw), (k_lora, self.kw), (v_lora, self.vw)] - return ret - - def set_q_k_v(self): - """ - Necessary to implement for `HybridSplitQKVContainer` - """ - self.qw = self.policy.client_module.self_attn.q_proj.weight - self.qb = self.policy.client_module.self_attn.q_proj.bias - self.kw = self.policy.client_module.self_attn.k_proj.weight - self.kb = self.policy.client_module.self_attn.k_proj.bias - self.vw = self.policy.client_module.self_attn.v_proj.weight - self.vb = self.policy.client_module.self_attn.v_proj.bias - - def set_mlp_gate(self): - """ - Necessary to implement for `HybridGatedMLPContainer` - """ - self.inter_up_w = self.policy.client_module.mlp.up_proj.weight - self.inter_up_b = None - self.inter_gate_w = self.policy.client_module.mlp.gate_proj.weight - self.inter_gate_b = None - - def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): - param_names = ( - 'self_attn.q_proj.weight', \ - 'self_attn.k_proj.weight', \ - 'self_attn.v_proj.weight', \ - 'self_attn.o_proj.weight', \ - 'mlp.up_proj.weight', \ - 'mlp.gate_proj.weight', \ - 'mlp.down_proj.weight', \ - 'input_layernorm.weight', \ - 'post_attention_layernorm.weight' - 'self_attn.q_proj.bias', \ - 'self_attn.k_proj.bias', \ - 'self_attn.v_proj.bias', \ - 'self_attn.o_proj.bias', \ - ) - - maybe_copy_qkv(module.attention, - sd, - weight_quantizer, - mp_replace, - 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], - split_qkv=self.policy.split_qkv) - maybe_copy_qkv(module.attention, - sd, - weight_quantizer, - mp_replace, - 'attn_qkvb', [prefix + param_names[9], prefix + param_names[10], prefix + param_names[11]], - split_qkv=self.policy.split_qkv) - maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[2], - prefix + param_names[3]) - maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[3], - prefix + param_names[12]) - maybe_copy_geglu(module.mlp, sd, weight_quantizer, mp_replace, 'inter_w', - [prefix + param_names[4], prefix + param_names[5]]) - maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, 'output_w', prefix + param_names[6]) - - maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[8], prefix + param_names[7]) - maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[10], prefix + param_names[8]) - - -class InternLMLayerPolicy(TransformerPolicy): - _orig_layer_class = [] - _orig_layer_class_inited = False - - def __init__(self, client_module, inference=True): - super().__init__( - inference, - mlp_act_func_type=ActivationFuncType.GATED_SILU, - norm_type=NormType.RMSNorm, - ) - self.client_module = client_module - - self._init_orig_layer_class_once() - - def _init_orig_layer_class_once(self): - if InternLMLayerPolicy._orig_layer_class_inited: - return - - for sub_pkg in ['', '.internlm-7b', '.internlm-chat-7b']: - try: - from transformers.utils import TRANSFORMERS_DYNAMIC_MODULE_NAME - module = importlib.import_module(f"{TRANSFORMERS_DYNAMIC_MODULE_NAME}{sub_pkg}.modeling_internlm") - if module.InternLMDecoderLayer not in InternLMLayerPolicy._orig_layer_class: - InternLMLayerPolicy._orig_layer_class.append(module.InternLMDecoderLayer) - except ImportError: - continue - - InternLMLayerPolicy._orig_layer_class_inited = True - - def get_hidden_heads(self): - return self.client_module.self_attn.q_proj.weight.shape[1], \ - self.client_module.self_attn.num_heads, \ - self.client_module.input_layernorm.variance_epsilon, \ - self.client_module.mlp.gate_proj.weight.shape[0] - - def attention(self, enable_training=False): - qw = self.client_module.self_attn.q_proj.weight - kw = self.client_module.self_attn.k_proj.weight - vw = self.client_module.self_attn.v_proj.weight - qb = self.client_module.self_attn.q_proj.bias - kb = self.client_module.self_attn.k_proj.bias - vb = self.client_module.self_attn.v_proj.bias - - qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) - qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) - - return qkvw, \ - qkvb, \ - self.client_module.self_attn.o_proj.weight, \ - self.client_module.self_attn.o_proj.bias - - def mlp(self, enable_training=False): - mlp1_up = self.client_module.mlp.up_proj.weight - mlp1_gate = self.client_module.mlp.gate_proj.weight - mlp2 = self.client_module.mlp.down_proj.weight - - mlp1 = Parameter(torch.cat((mlp1_up, mlp1_gate), dim=0), requires_grad=enable_training) - - return mlp1, None, mlp2, None - - def layernorm(self): - return self.client_module.post_attention_layernorm.weight, \ - None, \ - self.client_module.input_layernorm.weight, \ - None diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/llama.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/llama.py deleted file mode 100644 index 7af333dc1ee4771f52bcc837cf01c847486bdfd7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/llama.py +++ /dev/null @@ -1,164 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .base import * -from .features import HybridSplitQKVContainer, HybridGatedMLPContainer, MetaTensorContainer -from deepspeed.utils.types import ActivationFuncType, NormType -from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference -import torch -from torch.nn.parameter import Parameter - -from ..policy import ( - TransformerPolicy, - transformer_param_names, - maybe_copy, - maybe_copy_qkv, - maybe_copy_geglu, - maybe_get_lora, -) - - -class DS_LLAMAContainer(MetaTensorContainer, HybridGatedMLPContainer, HybridSplitQKVContainer, - BaseTransformerContainer): - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - # All model specific things should be defined here instead of the base class. - - def create_module(self, config=None): - _config = config if config is not None else self.ds_model_config - - _config.rotate_half = True - _config.rotate_every_two = False - _config.rotary_dim = self.hidden_size // self.num_attention_heads - _config.rope_theta = self.policy.client_module.self_attn.rope_theta - self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) - - return self.module - - def set_lora_params(self): - """ - Necessary to implement for `HybridEngineContainer` - """ - self.lora_params = [ - maybe_get_lora(p) for p in [ - self.policy.client_module.mlp.up_proj.weight, self.policy.client_module.mlp.gate_proj.weight, - self.policy.client_module.mlp.down_proj.weight, self.policy.client_module.self_attn.q_proj.weight, - self.policy.client_module.self_attn.k_proj.weight, self.policy.client_module.self_attn.v_proj.weight, - self.policy.client_module.self_attn.o_proj.weight - ] - ] - - def get_lora_matched_pair(self): - up_proj_lora, gate_proj_lora, down_proj_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() - ret = [(up_proj_lora, self.inter_up_w), (gate_proj_lora, self.inter_gate_w), (down_proj_lora, self._4hh_w), - (out_lora, self.dense_w), (q_lora, self.qw), (k_lora, self.kw), (v_lora, self.vw)] - return ret - - def set_q_k_v(self): - """ - Necessary to implement for `HybridSplitQKVContainer` - """ - self.qw = self.policy.client_module.self_attn.q_proj.weight - self.qb = None - self.kw = self.policy.client_module.self_attn.k_proj.weight - self.kb = None - self.vw = self.policy.client_module.self_attn.v_proj.weight - self.vb = None - - def set_mlp_gate(self): - """ - Necessary to implement for `HybridGatedMLPContainer` - """ - self.inter_up_w = self.policy.client_module.mlp.up_proj.weight - self.inter_up_b = None - self.inter_gate_w = self.policy.client_module.mlp.gate_proj.weight - self.inter_gate_b = None - - def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): - param_names = ( - 'self_attn.q_proj.weight', \ - 'self_attn.k_proj.weight', \ - 'self_attn.v_proj.weight', \ - 'self_attn.o_proj.weight', \ - 'mlp.up_proj.weight', \ - 'mlp.gate_proj.weight', \ - 'mlp.down_proj.weight', \ - 'post_attention_layernorm.weight', \ - 'input_layernorm.weight', - ) - - maybe_copy_qkv(module.attention, - sd, - weight_quantizer, - mp_replace, - 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], - split_qkv=self.policy.split_qkv) - for i in range(3, 4): - maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], - prefix + param_names[i]) - maybe_copy_geglu(module.mlp, sd, weight_quantizer, mp_replace, 'inter_w', - [prefix + param_names[4], prefix + param_names[5]]) - maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, 'output_w', prefix + param_names[6]) - - maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[8], prefix + param_names[7]) - maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[10], prefix + param_names[8]) - - # This line is necessary for proper output when kernels + meta tensors are used in Llama models - # TODO: Investigate root-cause and fix meta tensor loading - module.mlp.output_b = None - - -class LLAMALayerPolicy(TransformerPolicy): - - def __init__(self, client_module, inference=True): - super().__init__( - inference, - mlp_act_func_type=ActivationFuncType.GATED_SILU, - norm_type=NormType.RMSNorm, - ) - self.client_module = client_module - try: - import transformers - LLAMALayerPolicy._orig_layer_class = transformers.models.llama.modeling_llama.LlamaDecoderLayer # type: ignore - except: - LLAMALayerPolicy._orig_layer_class = None - - def get_hidden_heads(self): - hidden_heads = ( - self.client_module.self_attn.q_proj.in_features, - self.client_module.self_attn.num_heads, - self.client_module.input_layernorm.variance_epsilon, - self.client_module.mlp.gate_proj.out_features, - ) - return hidden_heads - - def attention(self, enable_training=False): - qw = self.client_module.self_attn.q_proj.weight - kw = self.client_module.self_attn.k_proj.weight - vw = self.client_module.self_attn.v_proj.weight - - qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) - - return qkvw, \ - None, \ - self.client_module.self_attn.o_proj.weight, \ - None - - def mlp(self, enable_training=False): - mlp1_up = self.client_module.mlp.up_proj.weight - mlp1_gate = self.client_module.mlp.gate_proj.weight - mlp2 = self.client_module.mlp.down_proj.weight - - mlp1 = Parameter(torch.cat((mlp1_up, mlp1_gate), dim=0), requires_grad=enable_training) - - return mlp1, None, mlp2, None - - def layernorm(self): - return self.client_module.post_attention_layernorm.weight, \ - None, \ - self.client_module.input_layernorm.weight, \ - None diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/llama2.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/llama2.py deleted file mode 100644 index b531890ab859dfbff51a401819846299d8e77c08..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/llama2.py +++ /dev/null @@ -1,158 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .base import * -from .features import HybridSplitQKVContainer, HybridGatedMLPContainer, MetaTensorContainer -from deepspeed.utils.types import ActivationFuncType, NormType -from deepspeed.model_implementations.transformers.ds_llama2 import DeepSpeedLlama2Inference -import torch -from torch.nn.parameter import Parameter - -from ..policy import ( - TransformerPolicy, - transformer_param_names, - maybe_copy, - maybe_copy_qkv, - maybe_copy_geglu, - maybe_get_lora, -) - - -class DS_LLAMA2Container(MetaTensorContainer, HybridGatedMLPContainer, HybridSplitQKVContainer, - BaseTransformerContainer): - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - # All model specific things should be defined here instead of the base class. - - def create_module(self, config=None): - _config = config if config is not None else self.ds_model_config - - _config.rotate_half = False - _config.rotate_every_two = True - _config.rotary_dim = self.hidden_size // self.num_attention_heads - _config.num_kv = self.policy.client_module.attention.n_kv_heads - self.module = DeepSpeedLlama2Inference(_config, mp_group=self.mp_group) - - return self.module - - def set_lora_params(self): - """ - Necessary to implement for `HybridEngineContainer` - """ - self.lora_params = [ - maybe_get_lora(p) for p in [ - self.policy.client_module.feed_forward.w3.weight, self.policy.client_module.feed_forward.w1.weight, - self.policy.client_module.feed_forward.w2.weight, self.policy.client_module.attention.wq.weight, - self.policy.client_module.attention.wk.weight, self.policy.client_module.attention.wv.weight, - self.policy.client_module.attention.wo.weight - ] - ] - - def get_lora_matched_pair(self): - up_proj_lora, gate_proj_lora, down_proj_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() - ret = [(up_proj_lora, self.inter_up_w), (gate_proj_lora, self.inter_gate_w), (down_proj_lora, self._4hh_w), - (out_lora, self.dense_w), (q_lora, self.qw), (k_lora, self.kw), (v_lora, self.vw)] - return ret - - def set_q_k_v(self): - """ - Necessary to implement for `HybridSplitQKVContainer` - """ - self.qw = self.policy.client_module.attention.wq.weight - self.qb = None - self.kw = self.policy.client_module.attention.wk.weight - self.kb = None - self.vw = self.policy.client_module.attention.wv.weight - self.vb = None - - def set_mlp_gate(self): - """ - Necessary to implement for `HybridGatedMLPContainer` - """ - self.inter_up_w = self.policy.client_module.feed_forward.w2.weight - self.inter_up_b = None - self.inter_gate_w = self.policy.client_module.feed_forward.w1.weight - self.inter_gate_b = None - - def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): - param_names = ( - 'attention.wq.weight', \ - 'attention.wk.weight', \ - 'attention.wv.weight', \ - 'attention.wo.weight', \ - 'feed_forward.w3.weight', \ - 'feed_forward.w1.weight', \ - 'feed_forward.w2.weight', \ - 'ffn_norm.weight', \ - 'attention_norm.weight' - ) - - maybe_copy_qkv(module.attention, - sd, - weight_quantizer, - mp_replace, - 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], - split_qkv=self.policy.split_qkv) - for i in range(3, 4): - maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], - prefix + param_names[i]) - maybe_copy_geglu(module.mlp, sd, weight_quantizer, mp_replace, 'inter_w', - [prefix + param_names[4], prefix + param_names[5]]) - maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, 'output_w', prefix + param_names[6]) - - maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[8], prefix + param_names[7]) - maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[10], prefix + param_names[8]) - - -class LLAMA2LayerPolicy(TransformerPolicy): - - def __init__(self, client_module, inference=True): - super().__init__( - inference, - mlp_act_func_type=ActivationFuncType.GATED_SILU, - norm_type=NormType.RMSNorm, - ) - self.client_module = client_module - try: - import llama - LLAMA2LayerPolicy._orig_layer_class = llama.model.TransformerBlock # type: ignore - except: - LLAMA2LayerPolicy._orig_layer_class = None - - def get_hidden_heads(self): - return self.client_module.attention.wq.weight.shape[1], \ - self.client_module.n_heads, \ - self.client_module.ffn_norm.eps, \ - (self.client_module.feed_forward.w1.weight.shape[0] * \ - deepspeed.comm.get_world_size() if deepspeed.comm.is_initialized() else 1) # this is a hack to inject when model is already partitioned! - - def attention(self, enable_training=False): - qw = self.client_module.attention.wq.weight - kw = self.client_module.attention.wk.weight - vw = self.client_module.attention.wv.weight - - qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) - - return qkvw, \ - None, \ - self.client_module.attention.wo.weight, \ - None - - def mlp(self, enable_training=False): - mlp1_up = self.client_module.feed_forward.w3.weight - mlp1_gate = self.client_module.feed_forward.w1.weight - mlp2 = self.client_module.feed_forward.w2.weight - - mlp1 = Parameter(torch.cat((mlp1_up, mlp1_gate), dim=0), requires_grad=enable_training) - - return mlp1, None, mlp2, None - - def layernorm(self): - return self.client_module.ffn_norm.weight, \ - None, \ - self.client_module.attention_norm.weight, \ - None diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/megatron_gpt.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/megatron_gpt.py deleted file mode 100644 index 2851dd246d99c94235cc6c0abce05f0c73375138..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/megatron_gpt.py +++ /dev/null @@ -1,117 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .base import * -from .features.megatron import MegatronContainer -from deepspeed.model_implementations.transformers.ds_megatron_gpt import DeepSpeedMegatronGPTInference -import torch -from ..policy import TransformerPolicy -from packaging import version as pkg_version - - -class DS_MegatronGPTContainer(MegatronContainer, BaseTransformerContainer): - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - # All model specific things should be defined here instead of the base class. - - def create_module(self, config=None): - _config = config if config is not None else self.ds_model_config - self.module = DeepSpeedMegatronGPTInference(_config, mp_group=self.mp_group) - self.module.config.scale_attention = self.scale_attention - - if self.megatron_v2: - self.module.config.rotate_half = True - self.module.config.rotate_every_two = False - - return self.module - - -# TODO: Megatron GPT MoE inherits from Megatron policy and replaces mlp -# TODO: Generalize MoE overall goal, expand beyond Megatron -class MegatronLayerPolicy(TransformerPolicy): - _orig_layer_class = None - version = 0 - moe_type = 'standard' - megatron_v2 = True - use_mup = False - - def __init__(self, client_module, inference=True): - super().__init__(inference, megatron_v2=MegatronLayerPolicy.megatron_v2, use_mup=MegatronLayerPolicy.use_mup) - self.client_module = client_module - # we use megatron version to differentiate between the old and new - # megatron-lm source code - if MegatronLayerPolicy._orig_layer_class is None: - if pkg_version.parse(torch.__version__) <= pkg_version.parse("1.2"): - MegatronLayerPolicy._orig_layer_class = None - else: - try: - from megatron.model.transformer import ParallelTransformerLayer - MegatronLayerPolicy._orig_layer_class = ParallelTransformerLayer - MegatronLayerPolicy.version = 1 - except ImportError: - MegatronLayerPolicy._orig_layer_class = None - - def get_hidden_heads(self): - if MegatronLayerPolicy.version == 0: - return self.client_module.attention.query_key_value.weight.shape[1], \ - self.client_module.attention.num_attention_heads, \ - self.client_module.input_layernorm.eps, \ - DEFAULT_INTERMEDIATE_SIZE - else: - return self.client_module.self_attention.query_key_value.weight.shape[1], \ - self.client_module.self_attention.num_attention_heads, \ - self.client_module.input_layernorm.eps, \ - DEFAULT_INTERMEDIATE_SIZE - - def attention(self, enable_training=False): - if self.inference: - if MegatronLayerPolicy.version == 0: - attention = self.client_module.attention - else: - attention = self.client_module.self_attention - - return attention.query_key_value.weight, \ - attention.query_key_value.bias, \ - attention.dense.weight, \ - attention.dense.bias - - def mlp(self, moe_type='standard', enable_training=False): - from deepspeed.moe.utils import has_moe_layers - moe, _ = has_moe_layers(self.client_module) - - if moe: - moe_experts = self.client_module.mlp.deepspeed_moe.experts.deepspeed_experts if moe_type == 'standard' else \ - self.client_module.mlp.moe.deepspeed_moe.experts.deepspeed_experts - num_experts = len(moe_experts) - if moe_type == 'standard': - return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \ - [moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \ - [moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \ - [moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)] - else: - - return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \ - [moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \ - [moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \ - [moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)], \ - self.client_module.mlp.mlp.dense_h_to_4h.weight, \ - self.client_module.mlp.mlp.dense_h_to_4h.bias, \ - self.client_module.mlp.mlp.dense_4h_to_h.weight, \ - self.client_module.mlp.mlp.dense_4h_to_h.bias, \ - self.client_module.mlp.coefficient.weight - - else: - return self.client_module.mlp.dense_h_to_4h.weight, \ - self.client_module.mlp.dense_h_to_4h.bias, \ - self.client_module.mlp.dense_4h_to_h.weight, \ - self.client_module.mlp.dense_4h_to_h.bias - - def layernorm(self): - return self.client_module.post_attention_layernorm.weight, \ - self.client_module.post_attention_layernorm.bias, \ - self.client_module.input_layernorm.weight, \ - self.client_module.input_layernorm.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/megatron_gpt_moe.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/megatron_gpt_moe.py deleted file mode 100644 index c4063be05b6c5f232470685a292be51667c027a8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/megatron_gpt_moe.py +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .base import * -from .base_moe import * -from .features.megatron import MegatronContainer -from deepspeed.model_implementations.transformers.ds_megatron_gpt import DeepSpeedMegatronGPTInference -import torch -from .megatron_gpt import MegatronLayerPolicy -from packaging import version as pkg_version - - -class DS_MegatronGPTMoEContainer(MegatronContainer, BaseTransformerMoEContainer): - - def __init__(self, policy, config, model_config, layer_id): - super().__init__(policy, config, model_config, layer_id) - - # All model specific things should be defined here instead of the base class. - - def create_module(self, config=None): - _config = config if config is not None else self.ds_model_config - self.module = DeepSpeedMegatronGPTInference(_config, mp_group=self.mp_group) - self.module.config.scale_attention = self.scale_attention - - if self.megatron_v2: - self.module.config.rotate_half = True - self.module.config.rotate_every_two = False - - return self.module - - -# TODO: Megatron GPT MoE inherits from Megatron policy and replaces mlp -# TODO: Generalize MoE overall goal, expand beyond Megatron -class MegatronMoELayerPolicy(MegatronLayerPolicy): - _orig_layer_class = None - version = 0 - moe_type = 'standard' - num_experts = 1 - - def __init__(self, client_module, inference=True): - super().__init__(inference) - self.client_module = client_module - # we use megatron version to differentiate between the old and new - # megatron-lm source code - if MegatronMoELayerPolicy._orig_layer_class is None: - if pkg_version.parse(torch.__version__) <= pkg_version.parse("1.2"): - MegatronMoELayerPolicy._orig_layer_class = None - else: - try: - from megatron.model.transformer import ParallelTransformerLayer - MegatronMoELayerPolicy._orig_layer_class = ParallelTransformerLayer - except ImportError: - MegatronMoELayerPolicy._orig_layer_class = None - - def get_num_experts(self): - return self.num_experts - - def mlp(self, moe_type='standard', enable_training=False): - # for now, all of this is tightly coupled to megatron-deepspeed moe implementation - # todo: think and refactor this to be more general - - #from deepspeed.moe.utils import has_moe_layers - #moe, _ = has_moe_layers(self.client_module) - - moe_experts = self.client_module.mlp.deepspeed_moe.experts.deepspeed_experts if moe_type == 'standard' else \ - self.client_module.mlp.moe.deepspeed_moe.experts.deepspeed_experts - num_experts = len(moe_experts) - self.num_experts = num_experts - - if moe_type == 'standard': - return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \ - [moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \ - [moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \ - [moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)] - else: - return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \ - [moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \ - [moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \ - [moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)], \ - self.client_module.mlp.mlp.dense_h_to_4h.weight, \ - self.client_module.mlp.mlp.dense_h_to_4h.bias, \ - self.client_module.mlp.mlp.dense_4h_to_h.weight, \ - self.client_module.mlp.mlp.dense_4h_to_h.bias, \ - self.client_module.mlp.coefficient.weight diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/opt.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/opt.py deleted file mode 100644 index eba619c79d6ebf5325c12f42c921b77dcb6e5ad7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/opt.py +++ /dev/null @@ -1,160 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .base import * -from .features import MetaTensorContainer, HybridSplitQKVContainer -from deepspeed.model_implementations.transformers.ds_opt import DeepSpeedOPTInference -import torch -from torch.nn.parameter import Parameter -from ..policy import TransformerPolicy -from ..policy import transformer_param_names -from ..policy import maybe_copy -from ..policy import maybe_copy_qkv -from ..policy import maybe_get_lora -from deepspeed.utils.types import ActivationFuncType - - -class DS_OPTContainer(MetaTensorContainer, HybridSplitQKVContainer, BaseTransformerContainer): - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - # All model specific things should be defined here instead of the base class. - - def create_module(self, config=None): - _config = config if config is not None else self.ds_model_config - self.module = DeepSpeedOPTInference(_config, mp_group=self.mp_group) - self.module.config.scale_attention = self.scale_attention - return self.module - - def set_lora_params(self): - """ - Necessary to implement for `HybridEngineContainer` - """ - self.lora_params = [ - maybe_get_lora(p) for p in [ - self.policy.client_module.fc1, - self.policy.client_module.fc2, - self.policy.client_module.self_attn.q_proj, - self.policy.client_module.self_attn.k_proj, - self.policy.client_module.self_attn.v_proj, - self.policy.client_module.self_attn.out_proj, - ] - ] - - def set_q_k_v(self): - """ - Necessary to implement for `HybridSplitQKVContainer` - """ - self.qw = self.policy.client_module.self_attn.q_proj.weight - self.qb = self.policy.client_module.self_attn.q_proj.bias - self.kw = self.policy.client_module.self_attn.k_proj.weight - self.kb = self.policy.client_module.self_attn.k_proj.bias - self.vw = self.policy.client_module.self_attn.v_proj.weight - self.vb = self.policy.client_module.self_attn.v_proj.bias - - def get_lora_matched_pair(self): - fc1_lora, fc2_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params() - ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (out_lora, self.dense_w), (q_lora, self.qw), - (k_lora, self.kw), (v_lora, self.vw)] - return ret - - def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): - param_names = ( - 'self_attn.q_proj.weight', \ - 'self_attn.k_proj.weight', \ - 'self_attn.v_proj.weight', \ - 'self_attn.q_proj.bias', \ - 'self_attn.k_proj.bias', \ - 'self_attn.v_proj.bias', \ - 'self_attn.out_proj.weight', \ - 'self_attn.out_proj.bias', \ - 'fc1.weight', \ - 'fc1.bias', \ - 'fc2.weight', \ - 'fc2.bias', \ - 'final_layer_norm.weight', \ - 'final_layer_norm.bias', \ - 'self_attn_layer_norm.weight', \ - 'self_attn_layer_norm.bias' - ) - - for i in range(0, 6, 3): - maybe_copy_qkv(module.attention, - sd, - weight_quantizer, - mp_replace, - transformer_param_names[i // 3], - [prefix + param_names[i], prefix + param_names[i + 1], prefix + param_names[i + 2]], - split_qkv=self.policy.split_qkv) - for i in range(6, 8): - maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 4], - prefix + param_names[i]) - for i in range(8, 14): - maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i - 4], - prefix + param_names[i]) - for i in range(14, 16): - maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i - 4], - prefix + param_names[i]) - - -class HFOPTLayerPolicy(TransformerPolicy): - _orig_layer_class = None - - def __init__(self, client_module, inference=True, use_load_prefix=True): - super().__init__(inference, linear_layer=True, pre_attn_norm=True, use_load_prefix=use_load_prefix) - self.client_module = client_module - try: - import transformers - HFOPTLayerPolicy._orig_layer_class = transformers.models.opt.modeling_opt.OPTDecoderLayer - except: - HFOPTLayerPolicy._orig_layer_class = None - - if hasattr(TransformerPolicy, "hf_model_config") and hasattr(TransformerPolicy.hf_model_config, - "activation_function"): - if TransformerPolicy.hf_model_config.activation_function == "relu": - self.mlp_act_func_type = ActivationFuncType.ReLU - elif TransformerPolicy.hf_model_config.activation_function in ["gelu", "gelu_new"]: - self.mlp_act_func_type = ActivationFuncType.GELU - else: - raise ValueError("Unsupported activation function: {}".format( - TransformerPolicy.hf_model_config.activation_function)) - else: - self.mlp_act_func_type = ActivationFuncType.ReLU # default - - def get_hidden_heads(self): - return self.client_module.self_attn.embed_dim, \ - self.client_module.self_attn.num_heads, \ - self.client_module.self_attn_layer_norm.eps, \ - DEFAULT_INTERMEDIATE_SIZE - - def attention(self, enable_training=False): - qw = self.client_module.self_attn.q_proj.weight - qb = self.client_module.self_attn.q_proj.bias - - kw = self.client_module.self_attn.k_proj.weight - kb = self.client_module.self_attn.k_proj.bias - - vw = self.client_module.self_attn.v_proj.weight - vb = self.client_module.self_attn.v_proj.bias - - qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) - qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) - return qkvw, \ - qkvb, \ - self.client_module.self_attn.out_proj.weight, \ - self.client_module.self_attn.out_proj.bias - - def mlp(self, enable_training=False): - return self.client_module.fc1.weight, \ - self.client_module.fc1.bias, \ - self.client_module.fc2.weight, \ - self.client_module.fc2.bias - - def layernorm(self): - return self.client_module.final_layer_norm.weight, \ - self.client_module.final_layer_norm.bias, \ - self.client_module.self_attn_layer_norm.weight, \ - self.client_module.self_attn_layer_norm.bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/unet.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/unet.py deleted file mode 100644 index 48179265553150824855140750da3fc6e41503ef..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/unet.py +++ /dev/null @@ -1,56 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from torch.nn.parameter import Parameter - -from ..policy import DSPolicy -from ...model_implementations.diffusers.unet import DSUNet - - -class UNetPolicy(DSPolicy): - - def __init__(self): - super().__init__() - try: - import diffusers - self._orig_layer_class = diffusers.models.unet_2d_condition.UNet2DConditionModel - except AttributeError: - self._orig_layer_class = diffusers.models.unets.unet_2d_condition.UNet2DConditionModel - except ImportError: - self._orig_layer_class = None - - def match(self, module): - return isinstance(module, self._orig_layer_class) - - def match_replaced(self, module): - return isinstance(module, DSUNet) - - def apply(self, module, enable_cuda_graph=True): - # TODO(cmikeh2): Enable cuda graph should be an inference configuration - return DSUNet(module, enable_cuda_graph=enable_cuda_graph) - - def attention(self, client_module): - qw = client_module.to_q.weight - kw = client_module.to_k.weight - vw = client_module.to_v.weight - - if qw.shape[1] == kw.shape[1]: - qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=False) - - return qkvw, \ - client_module.to_out[0].weight, \ - client_module.to_out[0].bias, \ - qw.shape[-1], \ - client_module.heads - else: - #return None - #kvw = Parameter(torch.cat((kw, vw), dim=0), requires_grad=False) - return qw, \ - kw, vw, \ - client_module.to_out[0].weight, \ - client_module.to_out[0].bias, \ - qw.shape[-1], \ - client_module.heads diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/vae.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/vae.py deleted file mode 100644 index d26d0ef77ca9d10e27bad47859dba589190af7b8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/containers/vae.py +++ /dev/null @@ -1,42 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from ..policy import DSPolicy -from ...model_implementations.diffusers.vae import DSVAE - - -class VAEPolicy(DSPolicy): - - def __init__(self): - super().__init__() - try: - import diffusers - if hasattr(diffusers.models, "autoencoders"): - # Diffusers >= 0.25.0 - # Changes location to 'autoencoders' directory - self._orig_layer_class = diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL - elif hasattr(diffusers.models.vae, "AutoencoderKL"): - # Diffusers < 0.12.0 - self._orig_layer_class = diffusers.models.vae.AutoencoderKL - else: - # Diffusers >= 0.12.0 & < 0.25.0 - # Changes location of AutoencoderKL - self._orig_layer_class = diffusers.models.autoencoder_kl.AutoencoderKL - except ImportError: - self._orig_layer_class = None - - def match(self, module): - return isinstance(module, self._orig_layer_class) - - def match_replaced(self, module): - return isinstance(module, DSVAE) - - def apply(self, module, enable_cuda_graph=True): - # TODO(cmikeh2): Enable cuda graph should be an inference configuration - return DSVAE(module, enable_cuda_graph=enable_cuda_graph) - - # NOTE (lekurile): Should we have a diffusers policy class? - def attention(self, client_module): - pass diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/fusedqkv_utils.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/fusedqkv_utils.py deleted file mode 100644 index 0609c6001dd2f716b3c1953500369dc9e67a4987..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/fusedqkv_utils.py +++ /dev/null @@ -1,231 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -import torch -from deepspeed.utils.logging import warning_once -from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list, get_num_kv_heads, get_n_embd, get_num_attention_heads - - -def split_by_qkvlist_and_refuse(qkv_list, split_size, split_dim=0, cat_dim=0): - qkv_split_list = [torch.split(mat, split_size, dim=split_dim) for mat in qkv_list] - tp_fusedqkv_list = [ - torch.cat([qkv_s[i] for qkv_s in qkv_split_list], dim=cat_dim) for i in range(len(qkv_split_list[0])) - ] - return tp_fusedqkv_list - - -def require_tp_fused_qkvw(name, mp_size): - fused_qkvw_name_list = ['qkv_proj', 'query_key_value', 'attn.Wqkv', 'self_attn.W_pack', 'c_attn'] - - if mp_size == 1: - return False - for fused_name in fused_qkvw_name_list: - if fused_name in name: - return True - return False - - -def prepare_tp_fused_qkvw(module, src, mp_size, gpu_index): - - module_str = str(module).strip() - if src is None: - return - fused_type_dict = { - 'CodeGenBlock': 'codegentype', - 'BloomBlock': 'bloomtype', - 'GLMBlock': 'glmtype', - "MPTBlock": 'glmtype', - "MptBlock": 'glmtype', - "BaichuanLayer": 'glmtype', - "QWenBlock": 'qwentype', - "FalconDecoderLayer": 'bloomtype', - "GPTBigCodeBlock": 'bigcodetype', - "DecoderLayer": 'glmtype', - "Phi3DecoderLayer": "phi3type" - } - - def _codegen_type_transpose(input, mp_size, codegen_mp_num=4): - # codegen_mp_num defined in https://github.com/huggingface/transformers/blob/main/src/transformers/models/codegen/modeling_codegen.py - assert get_num_kv_heads() % ( - mp_size * codegen_mp_num) == 0, "codgen autoTP requires num_kv_heads % (mp_size*codegen_mp_num) == 0" - #input : [3*hidden_dim, hidden_dim](weight) or [3*hidden_dim](bias) - - shape = input.shape - dst_shape = get_shard_size(shape[0], mp_size) - num_mp_blocks = input.reshape(codegen_mp_num, shape[0] // codegen_mp_num, shape[1]) - - #num_mp_blocks : [codegen_mp_num, 3*hidden_dim/codegen_mp_num, :] - src_split = list(torch.split(num_mp_blocks, num_mp_blocks.shape[1] // 3, dim=1)) - src_split = [x.reshape(codegen_mp_num * mp_size, -1, shape[1]) for x in src_split] - - split_fusedqkv = split_by_qkvlist_and_refuse(src_split, get_shard_size(shape[0] // 3, mp_size), 0, 1) - tp_fuseqkv_weight = torch.cat(split_fusedqkv, dim=0).reshape(shape[0], -1) - - return tp_fuseqkv_weight[gpu_index * dst_shape:(gpu_index + 1) * dst_shape] - - def _glm_type_transpose(input, mp_size): - #input : [3*hidden_dim, hidden_dim](weight) or [3*hidden_dim](bias) - - # For chatglm2 & chatglm3(kv_heads=2), need to special handle. - if get_num_kv_heads() == 2: - shape = input.shape - hidden_dim = get_n_embd() - kv_dim = (shape[0] - hidden_dim) // get_num_kv_heads() - q = input[:hidden_dim] - k = input[hidden_dim:hidden_dim + kv_dim] - v = input[hidden_dim + kv_dim:] - q_split = q.split(get_shard_size_list(q.shape[0], mp_size), dim=0) - k_split = k.split(get_shard_size_list(k.shape[0], mp_size), dim=0) - v_split = v.split(get_shard_size_list(v.shape[0], mp_size), dim=0) - return torch.cat((q_split[gpu_index], k_split[gpu_index], v_split[gpu_index]), dim=0) - else: - shape = input.shape - src_split = torch.split(input, shape[0] // 3, dim=0) - - split_fusedqkv = split_by_qkvlist_and_refuse(src_split, get_shard_size_list(shape[0] // 3, mp_size)) - return split_fusedqkv[gpu_index] - - def _bloom_type_transpose(input, mp_size): - shape = input.shape - - split_fusedqkv = input.split(get_shard_size_list(shape[0], mp_size), dim=0) - return split_fusedqkv[gpu_index] - - def _qwen_type_transpose(input, mp_size, module): - if not hasattr(module, "_ds_fusedqkv_entered"): - # Adjust splitting absolute value variables - setattr(module, "_ds_fusedqkv_entered", True) - module.attn.split_size = get_shard_size(module.attn.split_size, mp_size) - return _glm_type_transpose(input, mp_size) - - def _bigcode_type_transpose(input, mp_size): - n_embd = get_n_embd() - q = input[:n_embd] - kv = input[n_embd:] - shape = q.shape - split_q = q.split(get_shard_size_list(shape[0], mp_size), dim=0) - return torch.cat((split_q[gpu_index], kv), dim=0) - - def _phi3_type_transpose(input, mp_size): - num_kv_heads = get_num_kv_heads() - num_heads = get_num_attention_heads() - hidden_size = input.shape[1] - head_dim = hidden_size // num_heads - q_pos = input.shape[0] - 2 * num_kv_heads * head_dim - q = input[:q_pos] - k = input[q_pos:q_pos + num_kv_heads * head_dim] - v = input[q_pos + num_kv_heads * head_dim:] - split_q = q.split(get_shard_size_list(q.shape[0], mp_size), dim=0) - split_k = k.split(get_shard_size_list(k.shape[0], mp_size), dim=0) - split_v = v.split(get_shard_size_list(v.shape[0], mp_size), dim=0) - return torch.cat((split_q[gpu_index], split_k[gpu_index], split_v[gpu_index]), dim=0) - - def _transpose_fused_qkvw(src, mp_size, fused_qkv_type=None, module=None): - - # suppose num_heads=n, q(n)_w means the n-th q head linear weight, the weight format are as following - # bloomtype: [q(1)_w,k(1)_w,v(1)_w,q(2)_w,k(2)_w,v(2)_w,...,q(n)_w,k(n)_w,v(n)_w] - # glmtype: [q(1)_w, q(2)_w,...,q(n)_w,k(1)_w,k(2)_w,...,k(n)_w,v(1)_w,v(2)_w,...,v(n)_w] - # codegentype: [q(1)_w,q(2)_w,...,q(n/t)_w,k(1)_w,k(2)_w,...,k(n/t)_w,v(1)_2,v(2)_w,...v(n/t)_w,q(n/t+1)_w,...], where t is a const defined in model file. - - if fused_qkv_type == 'bloomtype': - return _bloom_type_transpose(src, mp_size) - elif fused_qkv_type == 'codegentype': - return _codegen_type_transpose(src, mp_size) - elif fused_qkv_type == 'glmtype': - return _glm_type_transpose(src, mp_size) - elif fused_qkv_type == 'qwentype': - return _qwen_type_transpose(src, mp_size, module) - elif fused_qkv_type == 'bigcodetype': - return _bigcode_type_transpose(src, mp_size) - elif fused_qkv_type == 'phi3type': - return _phi3_type_transpose(src, mp_size) - - raise ValueError("unknown fused_qkv_type") - - module_name_matches = [k for k in fused_type_dict.keys() if k in module_str] - if module_name_matches: - # There can be overlap with matches (e.g., "DecoderLayer" and "FalconDecoderLayer"). - # We take the longest matching module_name - module_name = max(module_name_matches, key=len) - fused_type = fused_type_dict[module_name] - return _transpose_fused_qkvw(src, mp_size, fused_type, module) - warning_once(f"Unrecognized fusedkqv weight type, default to using bloom type," - f"please check in prepare_tp_fused_qkvw() to avoid potential calculation errors") - return _bloom_type_transpose(src, mp_size) - - -# For share qk type: -# q = [q1,...,q_{n/4}, q_{n/2+1},...,q_{3n/4}, k1,...,k_{n/4}, k_{n/2+1},...,k_{3n/4}] -# k = [q_{n/4+1},...,q_{n/2}, q_{3n/4+1},...,qn, k_{n/4+1},...,k_{n/2}, k{3n/4+1},...,kn] -# Avoid modifying the modeling code. We adjust the value and oproj weight to fit this qk type. -def shard_value_with_share_qk( - weight, - bias, - rank, - world_size, - shard_value=True # True -> shard_value; False -> shard_oproj -): - if shard_value: - total_size = weight.shape[0] - weight_cat_dim = 0 - else: - total_size = weight.shape[1] - weight_cat_dim = 1 - num_heads = get_num_kv_heads() - head_dim = total_size // num_heads - assert (num_heads % world_size == 0) - if world_size > num_heads // 2: - RuntimeError(f"world_size {world_size} is larger than half of num_heads {num_heads}") - head_per_rank = num_heads // world_size - q_head_start = rank * head_per_rank - # mapping q_head to v_head - v_head_ids = [] - i = 0 - # mapping neighbor q_head to v_head - while i < head_per_rank: - v_head_ids.append(q_head_start // 2) - q_head_start += 2 - i = i + 2 - - # mapping neighbor k_head to v_head - v_head_ids.extend([i + num_heads // 2 for i in v_head_ids]) - sharded_weight = [] - sharded_bias = [] - for head_id in v_head_ids: - if shard_value: - sharded_weight.append(weight[head_id * head_dim:(head_id + 1) * head_dim]) - if bias is not None: - sharded_bias.append(bias.data[head_id * head_dim:(head_id + 1) * head_dim]) - else: - sharded_weight.append(weight[:, head_id * head_dim:(head_id + 1) * head_dim]) - sharded_weight = torch.cat(sharded_weight, dim=weight_cat_dim) - if bias is not None: - if shard_value: - sharded_bias = torch.cat(sharded_bias, dim=0) - else: - bias = bias / float(world_size) - return torch.nn.Parameter(sharded_weight), torch.nn.Parameter(sharded_bias) - else: - return torch.nn.Parameter(sharded_weight), None - - -# For phi3 with chunk mlp, adjust the weight order. -def shard_chunk_mlp( - weight, - bias, - rank, - world_size, -): - weight_gate, weight_states = weight.chunk(2, dim=0) - total_size = weight_gate.shape[0] - split_weight_gate = weight_gate.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0) - split_weight_states = weight_states.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0) - shard_weight = torch.cat((split_weight_gate[rank], split_weight_states[rank]), dim=0) - if bias is not None: - bias_gate, bias_states = bias.chunk(2, dim=0) - split_bias_gate = bias_gate.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0) - split_bias_states = bias_states.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0) - return shard_weight, torch.cat((split_bias_gate[rank], split_bias_states[rank]), dim=0) - - return shard_weight, None diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/inject.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/inject.py deleted file mode 100644 index 401da1bd6ef76b268f65e48b77c30a8ed65bfc64..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/inject.py +++ /dev/null @@ -1,112 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import copy -import torch -from deepspeed.ops.transformer import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig - - -def module_inject(layer_obj, model, config, micro_batch_size, max_seq_length, seed, preln, fp16=True): - for name, child in model.named_children(): - if isinstance(child, layer_obj): - print('REPLACING BertLayer') - - cuda_config = DeepSpeedTransformerConfig(batch_size=micro_batch_size, - max_seq_length=max_seq_length, - hidden_size=config.hidden_size, - heads=config.num_attention_heads, - attn_dropout_ratio=config.attention_probs_dropout_prob, - hidden_dropout_ratio=config.hidden_dropout_prob, - num_hidden_layers=config.num_hidden_layers, - initializer_range=config.initializer_range, - seed=seed, - fp16=fp16, - pre_layer_norm=preln) - - new_module = DeepSpeedTransformerLayer(cuda_config) - - # copy relevant state from child -> new module - qw = child.attention.self.query.weight - qb = child.attention.self.query.bias - kw = child.attention.self.key.weight - kb = child.attention.self.key.bias - vw = child.attention.self.value.weight - vb = child.attention.self.value.bias - - qkvw = torch.cat((qw, kw, vw), 0) - qkvb = torch.cat((qb, kb, vb), 0) - - new_module.attn_qkvw.data = qkvw - new_module.attn_qkvb.data = qkvb - new_module.attn_ow.data = child.attention.output.dense.weight - new_module.attn_ob.data = child.attention.output.dense.bias - if preln: - attention_layerNorm = child.PostAttentionLayerNorm - else: - attention_layerNorm = child.attention.output.LayerNorm - new_module.attn_nw.data = attention_layerNorm.weight - new_module.attn_nb.data = attention_layerNorm.bias - if preln: - intermediate_FF = child.intermediate.dense_act - else: - intermediate_FF = child.intermediate.dense - new_module.inter_w.data = intermediate_FF.weight - new_module.inter_b.data = intermediate_FF.bias - new_module.output_w.data = child.output.dense.weight - new_module.output_b.data = child.output.dense.bias - if preln: - transformer_LayerNorm = child.PreAttentionLayerNorm - else: - transformer_LayerNorm = child.output.LayerNorm - new_module.norm_w.data = transformer_LayerNorm.weight - new_module.norm_b.data = transformer_LayerNorm.bias - - setattr(model, name, copy.deepcopy(new_module)) - - else: - module_inject(layer_obj, child, config, micro_batch_size, max_seq_length, seed, preln, fp16) - - return model - - -def test_hi(): - from turing.nvidia_modelingpreln import BertConfig as BertConfigPreLN - from turing.nvidia_modelingpreln import BertForQuestionAnswering as BertForQuestionAnsweringPreLN - from turing.nvidia_modelingpreln import BertLayer - bert_model_config = { - "vocab_size_or_config_json_file": 119547, - "hidden_size": 1024, - "num_hidden_layers": 1, - "num_attention_heads": 16, - "intermediate_size": 4096, - "hidden_act": "gelu", - "hidden_dropout_prob": 0.1, - "attention_probs_dropout_prob": 0.1, - "hidden_dropout_prob": 0.1, - "attention_probs_dropout_prob": 0.1, - "max_position_embeddings": 512, - "type_vocab_size": 2, - "initializer_range": 0.02 - } - bert_config = BertConfigPreLN(**bert_model_config) - base_model = BertForQuestionAnsweringPreLN(bert_config, args=None) - - #base_model = LinearStack() - - test_model = copy.deepcopy(base_model) - test_model = module_inject(BertLayer, test_model, bert_config, 4, 384, 1234) - - print('BASE', base_model) - print('TEST', test_model) - - #base_model.eval() - #test_model.eval() - - #test_input = torch.rand(1, base_model.input_dim) - - #base_output = base_model(test_input) - #test_output = test_model(test_input) - # - #assert torch.allclose(base_output, test_output, atol=3e-8) diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/layers.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/layers.py deleted file mode 100644 index a9faac71361a4aec8c76e6976121ed794b554115..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/layers.py +++ /dev/null @@ -1,821 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed import comm as dist -from torch import nn -from torch.nn import functional as F -from torch.nn.parameter import Parameter -from deepspeed.accelerator import get_accelerator -from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list -from deepspeed.runtime.zero.utils import is_zero_param -from abc import ABC, abstractmethod -from typing import Iterable, Any, Optional, List, Tuple -from .fusedqkv_utils import shard_value_with_share_qk, shard_chunk_mlp, prepare_tp_fused_qkvw -from deepspeed.runtime.tensor_parallel import AUTOTP_MODE -from copy import deepcopy -from typing import Union - -__all__ = [ - "TensorParallel_Layer", "LinearAllreduce", "LinearLayer", "LmHeadLinearAllreduce", "Yuan_LinearAllreduce", - "Yuan_LinearLayer", "GateUpPack_LinearLayer", "Conv_LinearALlreduce", "fused_LinearLayer", "conv_LinearLayer" -] - -DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.INFERENCE -DS_IS_REPLACED_MODULE = 'ds_is_replaced_module' -DS_TENSOR_MODEL_PARALLEL = 'tensor_model_parallel' - - -def get_auto_tp_mode(): - global DEEPSPEED_AUTOTP_MODE - return DEEPSPEED_AUTOTP_MODE - - -def is_autotp_training_mode(): - global DEEPSPEED_AUTOTP_MODE - return DEEPSPEED_AUTOTP_MODE == AUTOTP_MODE.TRAINING - - -def set_autotp_mode(training=False): - """ - Set the DEEPSPEED_AUTOTP_MODE based on the training flag - """ - global DEEPSPEED_AUTOTP_MODE - if training: - DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.TRAINING - else: - DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.INFERENCE - - -def add_bias(input, bias): - if bias is None: - return input - if is_autotp_training_mode(): - # Training mode - avoid inplace to ensure correct autograd - input = input + bias - return input - else: - input += bias - return input - - -class RowParallel(torch.autograd.Function): - """ - A custom autograd function for performing row-wise parallelism. - """ - - @staticmethod - def symbolic(graph, input): - """Symbolic function for tracing.""" - return input - - @staticmethod - def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, is_inference_mode: bool) -> torch.Tensor: - """ - Forward pass. - """ - ctx.group = group - if group == None: - return input - if is_inference_mode: - dist.inference_all_reduce(input, group=group) - else: - dist.all_reduce(input.contiguous(), group=group) - return input - - @staticmethod - def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor, None]: - """ - Backward pass. - """ - return None, grad_output, None - - -class AsyncColumnParallel(torch.autograd.Function): - - @staticmethod - def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, weight, bias) -> torch.Tensor: - """ - Forward pass. - """ - ctx.use_bias = bias is not None - ctx.group = group - output = torch.matmul(input, weight.transpose(-1, -2)) - if bias is not None: - output = add_bias(output, bias) - - ctx.save_for_backward(input, weight) - - return output - - @staticmethod - def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]: - - input, weight = ctx.saved_tensors - grad_input = grad_output.matmul(weight) - handle = dist.all_reduce(grad_input.contiguous(), group=ctx.group, async_op=True) - grad_weight = grad_output.view(-1, grad_output.shape[-1]).t().matmul(input.view(-1, input.shape[-1])) - grad_bias = grad_output.sum(0) if ctx.use_bias else None - handle.wait() - return None, grad_input, grad_weight, grad_bias - - -class ColumnParallel(torch.autograd.Function): - """ - Custom autograd function for column-wise parallelism. - """ - - @staticmethod - def symbolic(graph, input): - """Symbolic function for tracing.""" - return dist.all_reduce(input.contiguous(), dist.get_tensor_model_parallel_group()) - - @staticmethod - def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor) -> torch.Tensor: - """ - Forward pass. - """ - ctx.group = group - return input - - @staticmethod - def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]: - """ - Backward pass. - """ - if ctx.group == None: - return None, grad_output - - dist.all_reduce(grad_output.contiguous(), group=ctx.group) - return None, grad_output - - -class TensorParallel_Layer(nn.Module, ABC): - """ - A base class for model layers with tensor parallelism support. - This class is designed to be extended by specific layers that require distributed - operations and parameter gather/partitioning during inference or training. - - Attributes: - mode (str): The mode of operation[INFERENCE or TRAINING], default is "INFERENCE". - mp_group (Optional[dist.ProcessGroup]): The process group used for model parallelism. - tp_world_size (int): The world size of tensor parallelism, i.e., the number of parallel workers. - tp_index (int): The rank (ID) of the current worker in tensor parallelism. - support_training (bool): Flag indicating whether the layer supports training (default: False). - name (Optional[str]): The name of the layer, if provided. - """ - ##### Initialize Parameter List ##### - - # keep_module_on_host determines whether to keep the module on the host. - # Checkpoints are first loaded to the host (sometimes directly from disk to avoid filling host memory), - # so an additional copy is unnecessary. - keep_module_on_host: bool = False - - ##### Runtime Parameter List ##### - tp_overlap_comm: bool = False - """ Whether to overlap communication with computation. Currently, only allreduce supports overlap. """ - - def __init__(self, mp_group: Optional[dist.ProcessGroup], **kwargs: Any): - """ - Initializes the TensorParallel_Layer with optional model parallelism group and layer name. - - Args: - mp_group (Optional[dist.ProcessGroup]): The process group for model parallelism. - If None, no model parallelism is set. - """ - super().__init__() - self.support_training: bool = False - self.mp_group = mp_group - if mp_group is not None: - self.tp_world_size: int = dist.get_world_size(self.mp_group) - self.tp_index: int = dist.get_rank(self.mp_group) - else: - self.tp_world_size: int = 1 - self.tp_index: int = 0 - - # backward compatibility - self.world_size = self.tp_world_size - self.rank = self.tp_index - - self.name = getattr(self, 'name', None) - if kwargs.get('name') is not None: - self.name = kwargs.get('name') # Set the layer name if provided. - - @classmethod - def set_keep_module_on_host(cls, value: bool): - """ - Set the static variable keep_module_on_host. - - Args: - value (bool): The new value for keep_module_on_host. - """ - cls.keep_module_on_host = value - - @abstractmethod - def forward(self, input): - """ - Forward pass method. Must be implemented by subclasses to define layer-specific operations. - """ - pass - - @abstractmethod - def gather_params(self, params_list): - """ - Gathers parameters across devices for distributed training. Must be implemented by subclasses in "TRAINING" mode. - """ - pass - - @abstractmethod - def _tp_partition(self, params_list: List[torch.Tensor]): - """ - Partitions the parameters for tensor parallelism. - It is necessary to ensure that this function only involves the logic of params partitioning. - """ - pass - - def config_requires_grad(self, weight): - if weight is not None: - if self.is_training_mode(): - if weight.requires_grad is None: - weight.requires_grad = True - else: - weight.requires_grad = False - - def config_tp_params(self, weight): - """ - Configures the weight tensor for training with tensor parallelism. This includes enabling gradients - and associating necessary methods for parameter gathering and partitioning. - - Args: - weight (Optional[torch.Tensor]): The weight tensor to configure for tensor parallelism. - If None, no action is taken. - """ - # # The RNG states have already been synchronized in init_inference. - if self.is_training_mode(): - assert self.support_training, "No implementation of backward." - if weight is not None: - self.config_requires_grad(weight) - weight.gather_params = self.gather_params - weight._tp_partition = self._tp_partition - setattr(weight, DS_TENSOR_MODEL_PARALLEL, True) - setattr(weight, DS_IS_REPLACED_MODULE, True) - - def is_training_mode(self): - global DEEPSPEED_AUTOTP_MODE - return DEEPSPEED_AUTOTP_MODE == AUTOTP_MODE.TRAINING - - def __deepcopy__(self, memo): - # This function is designed for - # 'mp_group' (a 'ProcessGroup') cannot be pickled during deepcopy in some usage. - cls = self.__class__ - new_obj = cls.__new__(cls) - - for key, value in vars(self).items(): - if key == 'mp_group': - new_obj.mp_group = self.mp_group - else: - setattr(new_obj, key, deepcopy(value, memo)) - - memo[id(self)] = new_obj - return new_obj - - def extra_repr(self): - out_features, in_features = None, None - if self.weight is not None: - out_features, in_features = self.weight.ds_shape[-2:] if is_zero_param( - self.weight) else self.weight.shape[-2:] - dtype = self.weight.dtype if self.weight is not None else None - return "in_features={}, out_features={}, bias={}, dtype={}".format(in_features, out_features, self.bias - is not None, dtype) - - def move(self, tensor): - # TODO: consider the timing of deletion - # to save host resources when DP > 1。 - - # keep_module_on_host is used to keep the module on the host. Checkpoints are loaded to the host first (in some - # cases it can be done from the disk even to prevent filling host's memory), thus no need to create a new copy. - if tensor.is_meta: - # Keep tensor in meta device if tensor is meta. - return tensor - else: - device = 'cpu' if self.__class__.keep_module_on_host else get_accelerator().current_device_name() - return_new_copy = not self.__class__.keep_module_on_host - - # Using new tensors help in freeing memory (after split for example) was done before by calling clone(). - # Using copy=True instead of clone() will help in case of cpu --> cpu. - # Otherwise to() will not create a new copy for the view of the full tensor, and it will not be de-referenced. - cloned_tensor = tensor.to(device, copy=return_new_copy) - - if return_new_copy: - # free the memory of the original tensor to reduce memory peak - # Equivalent to directly deleting the tensor reference outside the function. - # see https://github.com/microsoft/DeepSpeed/pull/4353 - tensor.data = torch.empty(0, device=tensor.device) - return cloned_tensor - - -def configure_tensor_parallel_runtime(config): - runtime_keys = ['tp_overlap_comm'] - for key in runtime_keys: - if hasattr(config, key): - setattr(TensorParallel_Layer, key, getattr(config, key)) - - -class GatherReplacedLayerParams: - """ - A context manager for gathering parameters of a replaced layer, enabling partitioning and gathering functionality - based on the configuration of the model. - """ - - def __init__(self, - params: Union[Iterable[torch.Tensor], torch.Tensor], - module: torch.nn.Module, - enabled: bool = True): - """ - Initialize the context manager to handle parameter gathering and partitioning for a replaced layer. - - Args: - params (Iterable or torch.Tensor): A collection or single parameter to manage. - module (torch.nn.Module): The module that these parameters belong to. - enabled (bool): Flag indicating whether the parameter management is enabled (default: True). - """ - self.enabled = enabled - self.module = module - if not enabled: - return - - # Ensure params is a list, whether it's a single param or iterable (e.g., model.parameters()) - if isinstance(params, Iterable) and not isinstance(params, torch.Tensor): - self.params: List[torch.Tensor] = list(params) # Convert generators to a list for multiple iterations - else: - self.params: List[torch.Tensor] = [params] # Wrap single parameter in a list for uniform processing - - # Check if the parameters belong to a replaced layer (indicated by a specific attribute) - if not any(self._is_replaced_module_weight(p) for p in params): - self.enabled = False - return - - def _is_replaced_module_weight(self, param: torch.Tensor) -> bool: - """ - Helper function to determine if a parameter belongs to a replaced module. - - Args: - param (torch.Tensor): The parameter to check. - - Returns: - bool: True if the parameter belongs to a replaced module, False otherwise. - """ - return getattr(param, DS_IS_REPLACED_MODULE, False) - - def __enter__(self) -> None: - """ - Enter the context manager. If enabled, gather parameters for the replaced module. - """ - if self.enabled: - self.params[0].gather_params(self.params) - - def __exit__(self, exc_type, exc_value, traceback) -> None: - """ - Exit the context manager. If enabled, partition the parameters for the replaced module. - """ - #TODO : Check whether there are any missing attributes. - if self.enabled: - self.params[0]._tp_partition(self.params) - - -class LinearAllreduce(TensorParallel_Layer): - - def __init__(self, module, mp_group, **kwargs): - super(LinearAllreduce, self).__init__(mp_group, **kwargs) - self.weight = module.weight - self.bias = module.bias - - self._tp_partition([self.weight, self.bias]) - self.support_training = True - self.config_tp_params(self.weight) - if self.bias is not None: - # bias here is not tp params - self.config_requires_grad(self.bias) - - def forward(self, input): - output = torch.matmul(input, self.weight.transpose(-1, -2)) - output = RowParallel.apply(self.mp_group, output, not self.is_training_mode()) - if self.bias is not None: - output = add_bias(output, self.bias) - return output - - @torch.no_grad() - def gather_params(self, params_list): - - for idx, param in enumerate(params_list): - if param is None or idx > 0: - # don't gather bias - return - params_list[idx].data_partition = param.data - param = param.transpose(0, 1).contiguous() - - output_param = torch.empty(self.tp_world_size * param.shape[0], - param.shape[1], - dtype=param.dtype, - device=param.device) - dist.all_gather_into_tensor(output_param, param, group=self.mp_group) - params_list[idx].data = output_param.transpose(0, 1).contiguous() - return - - @torch.no_grad() - def _tp_partition(self, params_list): - - if not self.is_training_mode(): - self.uneven_partition(params_list) - return - - else: - for idx, param in enumerate(params_list): - if param is None: - # don't slipt bias - return - if idx > 0: # move bias to device at initialization - _partition = self.move(param).detach() - params_list[idx].data = _partition - return - - _partition = torch.chunk(param, self.tp_world_size, dim=-1)[self.tp_index] - - _partition = self.move(_partition).detach() - - params_list[idx].data = _partition - - def uneven_partition(self, params_list): - for idx, param in enumerate(params_list): - if param is None or idx > 0: - # don't slipt bias - return - assert self.name is not None, "The module name must be provided in the initialization." - _partition = params_list[idx].split(get_shard_size_list(params_list[idx].shape[1], self.tp_world_size, - self.name), - dim=1)[self.tp_index] - - _partition = self.move(_partition).detach() - params_list[idx].data = _partition - - -#remove kwargs from partition. -class LinearLayer(TensorParallel_Layer): - - def __init__(self, module, mp_group=None, skip_partition=False, **kwargs): - super(LinearLayer, self).__init__(mp_group, **kwargs) - self.weight = module.weight - self.bias = module.bias - if not skip_partition: - self._tp_partition([self.weight, self.bias]) - self.support_training = True - self.config_tp_params(self.weight) - if self.bias is not None: - self.config_tp_params(self.bias) - - def forward(self, input): - if not self.__class__.tp_overlap_comm: - if getattr(self, 'mp_group', None) is not None: - input = ColumnParallel.apply(self.mp_group, input) - output = torch.matmul(input, self.weight.transpose(-1, -2)) - if self.bias is not None: - output = add_bias(output, self.bias) - else: - output = AsyncColumnParallel.apply(self.mp_group, input, self.weight, self.bias) - - return output - - @torch.no_grad() - def gather_params(self, params_list): - # Does not support uneven shard. - for idx, param in enumerate(params_list): - - params_list[idx].data_partition = param.data - output_param = torch.empty((self.tp_world_size * param.shape[0], *param.shape[1:]), - dtype=param.dtype, - device=param.device) - dist.all_gather_into_tensor(output_param, param, group=self.mp_group) - params_list[idx].data = output_param.contiguous() - - @torch.no_grad() - def _tp_partition(self, params_list): - - if not self.is_training_mode(): - self.uneven_partition(params_list) - return - for idx, param in enumerate(params_list): - if param is None: - return - #split bias if provide - _partition = torch.chunk(param, self.tp_world_size, dim=0)[self.tp_index] - - _partition = self.move(_partition).detach() - - params_list[idx].data = _partition - - def uneven_partition(self, params_list): - - for idx, param in enumerate(params_list): - if param is None: - #split bias if provide - return - assert self.name is not None, "The module name must be provided in the initialization." - _partition = params_list[idx].split(get_shard_size_list(params_list[idx].shape[0], self.tp_world_size, - self.name), - dim=0)[self.tp_index] - - _partition = self.move(_partition).detach() - - params_list[idx].data = _partition - - # for bwc - @classmethod - def from_weights(cls, weight_shape=None, dtype=torch.half, weight=None, bias=None): - if weight is not None: - in_features = weight.shape[1] - out_features = weight.shape[0] - linear = nn.Linear(in_features, out_features, bias=(bias is not None)) - linear.weight.data = weight - if bias is not None: - linear.bias.data = bias - else: - in_features = weight_shape[1] - out_features = weight_shape[0] - linear = nn.Linear(in_features, out_features, bias=(bias is not None)) - return cls(linear, skip_partition=True) - - -class FusedModuleWrapper: - - def __init__(self, fused_module: nn.Module): - self.fused_module = fused_module - - def __getattr__(self, module): - return self.fused_module - - -class fused_LinearLayer(LinearLayer): - - def __init__(self, module, mp_group, skip_partition=False, **kwargs): - assert kwargs.get('fused_module') is not None, "'fused_module' is required but not provided" - # Use the warp class to avoid module circular references. - self.fused_module = FusedModuleWrapper(kwargs.get('fused_module')) - super().__init__(module, mp_group, skip_partition, **kwargs) - - @torch.no_grad() - def _tp_partition(self, params_list): - for idx, param in enumerate(params_list): - if param is None: - return - - _partition = prepare_tp_fused_qkvw(self.fused_module.module, param, self.tp_world_size, self.tp_index) - - _partition = self.move(_partition).detach() - - params_list[idx].data = _partition - - -class conv_LinearLayer(LinearLayer): - - @torch.no_grad() - def _tp_partition(self, params_list): - weight = None - bias = None - if len(params_list) == 1: - weight = params_list[0] - elif len(params_list) == 2: - weight, bias = params_list[0], params_list[1] - _partition = weight.data.split(get_shard_size_list(weight.shape[0], self.tp_world_size, self.name), - dim=1)[self.tp_index] - _partition = self.move(_partition).detach() - weight.data = _partition - - if bias is not None: - _partition = bias.data.split(get_shard_size_list(weight.shape[1], self.tp_world_size, self.name), - dim=0)[self.tp_index] - _partition = self.move(_partition).detach() - - bias.data = _partition - - -#override the subclasses related to weight splitting. -class Yuan_LinearAllreduce(LinearAllreduce): - - #Yuan2 - @torch.no_grad() - def _tp_partition(self, params_list): - weight, bias = shard_value_with_share_qk(params_list[0].data, params_list[1], self.tp_index, - self.tp_world_size, False) - params_list[0].data = weight - if bias is not None: - params_list[1].data = bias - - -class Yuan_LinearLayer(LinearLayer): - #Yuan2 - @torch.no_grad() - def _tp_partition(self, params_list): - weight, bias = shard_value_with_share_qk(params_list[0].data, params_list[1], self.tp_index, - self.tp_world_size, True) - params_list[0].data = self.move(weight).detach() - if bias is not None: - params_list[1].data = self.move(bias).detach() - - -class GateUpPack_LinearLayer(LinearLayer): - # chatGLM2, chatGLM2 - @torch.no_grad() - def _tp_partition(self, params_list): - weight, bias = shard_chunk_mlp(params_list[0].data, params_list[1], self.tp_index, self.tp_world_size) - params_list[0].data = self.move(weight).detach() - if bias is not None: - params_list[1].data = self.move(bias).detach() - - -class Conv_LinearALlreduce(LinearAllreduce): - - @torch.no_grad() - def _tp_partition(self, params_list): - for idx, param in enumerate(params_list): - if param is None: - return - param.data = param.data.transpose(-1, -2).contiguous() - - _partition = param.split(get_shard_size_list(param.shape[0], self.tp_world_size, self.name), - dim=1)[self.tp_index] - - _partition = self.move(_partition).detach() - - params_list[idx].data = _partition - - -#override the subclasses related to fwd/bwd. -class LmHeadLinearAllreduce(LinearAllreduce): - - def __init__(self, module, mp_group, **kwargs): - # set the fixed name before partition - self.name = "lm_head" - - # In some tied_embedding cases, only the lm head is sharded, while the word embedding is not. - # Reinitialization is used to decouple them and prevent the word embedding from being sharded. - # This should also be effective for cases where both are sharded in tied_embedding scenarios. - - # TODO: Training scenario-related tests, is it necessary to re-implement the vocab parallel module? - module.weight = nn.Parameter(module.weight.clone().detach()) - if hasattr(module, 'bias') and module.bias is not None: - module.bias = nn.Parameter(module.bias.clone().detach()) - super().__init__(module, mp_group, **kwargs) - - def forward(self, input): - input_shard_size = get_shard_size(input.shape[-1], self.tp_world_size, "lm_head") - input_shard_offset = sum(get_shard_size_list(input.shape[-1], self.tp_world_size, "lm_head")[0:self.tp_index]) - output = torch.matmul(input[:, :, input_shard_offset:input_shard_offset + input_shard_size], - self.weight.transpose(-1, -2)) - if self.mp_group is not None: - dist.inference_all_reduce(output, group=self.mp_group) - if self.bias is not None: - output = add_bias(output, self.bias) - return output - - -class TensorParallelConv2d(nn.Module): - - def __init__(self, conv, rank, world_size, shard_by_oc): - super().__init__() - self.rank = rank - self.world_size = world_size - self.shard_by_oc = shard_by_oc - self.shard_weights(conv) - - # Split along the input/output channel depending on whether it is the last conv layer. - def shard_weights(self, conv): - if self.shard_by_oc: - total_size = conv.weight.shape[0] - else: - total_size = conv.weight.shape[1] - bias_data = None - cols_per_rank = [0] - for i in range(self.world_size - 1, -1, -1): - cols = total_size // self.world_size - if i < total_size % self.world_size: - cols += 1 - cols_per_rank.append(cols_per_rank[-1] + cols) - weight_data = conv.weight.data - if self.shard_by_oc: - # not last conv layer, split output channel - weight_data = weight_data[cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]] - if conv.bias is not None: - bias_data = conv.bias.data[cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]] - else: - # last conv layer, split input channel - weight_data = weight_data[:, cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]] - if conv.bias is not None: - bias_data = conv.bias.data / float(self.world_size) - self.conv = nn.Conv2d(weight_data.shape[1], weight_data.shape[0], conv.kernel_size, conv.stride, conv.padding, - conv.dilation, conv.groups, conv.bias is not None, conv.padding_mode) - self.conv.weight = torch.nn.Parameter(weight_data) - if conv.bias is not None: - self.conv.bias = torch.nn.Parameter(bias_data) - del conv - - def forward(self, input: torch.Tensor) -> torch.Tensor: - return self.conv(input) - - -class TensorParallelOcShardConv2d(TensorParallelConv2d): - - def __init__(self, conv, rank, world_size): - super().__init__(conv, rank, world_size, True) - - -class TensorParallelIcShardConv2d(TensorParallelConv2d): - - def __init__(self, conv, rank, world_size): - super().__init__(conv, rank, world_size, False) - - def forward(self, input: torch.Tensor) -> torch.Tensor: - out = self.conv(input) - if self.world_size > 1: - dist.inference_all_reduce(out) - return out - - -class Normalize(nn.Module): - - def __init__(self, dim=None, dtype=torch.float, eps=1e-5, weight=None, bias=None): - super(Normalize, self).__init__() - if weight is not None: - self.weight = weight - self.bias = bias - else: - self.norm = nn.LayerNorm(dim, eps=eps).to(dtype).to(get_accelerator().current_device_name()) - self.weight = self.norm.weight - self.bias = self.norm.bias - - self.eps = eps - - def forward(self, input): - return nn.functional.layer_norm(input, input.shape[-1:], self.weight, self.bias, eps=self.eps) - - -class EmbeddingLayer(nn.Module): - - def __init__(self, weight_shape=None, dtype=torch.half, weight=None, bias=None): - super(EmbeddingLayer, self).__init__() - if weight is None: - self.weight = Parameter( - torch.empty(weight_shape[0], - weight_shape[1], - dtype=dtype, - device=get_accelerator().current_device_name())) - else: - self.weight = weight - - def forward(self, input): - return F.embedding(input, self.weight) - - -class OPTEmbedding(EmbeddingLayer): - """ - This module learns positional embeddings up to a fixed maximum size. - """ - - def __init__(self, weight_shape=None, weight=None, bias=None): - # OPT is set up so that if padding_idx is specified then offset the embedding ids by 2 - # and adjust num_embeddings appropriately. Other models don't have this hack - self.offset = 2 - super().__init__(weight_shape, weight=weight) - - def forward(self, attention_mask: torch.LongTensor, past_key_values_length: int = 0, position_ids: int = 0): - """`input_ids_shape` is expected to be [bsz x seqlen].""" - attention_mask = attention_mask.long() - - # create positions depending on attention_mask - positions = (torch.cumsum(attention_mask, dim=1).type_as(attention_mask) * attention_mask).long() - 1 - - # cut positions if `past_key_values_length` is > 0 - positions = positions[:, past_key_values_length:] - - return super().forward(positions + self.offset) - - -class RMSNormalize(nn.Module): - - def __init__(self, dim=None, dtype=torch.float, eps=1e-5, weight=None): - super(RMSNormalize, self).__init__() - if weight is not None: - self.weight = weight - else: - self.weight = nn.Parameter(torch.ones(dim, dtype=dtype, device=get_accelerator().current_device_name())) - - self.eps = eps - - def forward(self, hidden_states): - variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.eps) - if self.weight.dtype in [torch.float16, torch.bfloat16]: - hidden_states = hidden_states.to(self.weight.dtype) - - return hidden_states * self.weight diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/load_checkpoint.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/load_checkpoint.py deleted file mode 100644 index 862628fa7b4b46ad2cd8bf2c3f5706233bda0627..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/load_checkpoint.py +++ /dev/null @@ -1,285 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from torch import nn -from deepspeed.model_implementations.transformers.ds_bloom import DeepSpeedBloomInference -from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference -from deepspeed.model_implementations.transformers.ds_bert import DeepSpeedBERTInference -from deepspeed.model_implementations.transformers.ds_megatron_gpt import DeepSpeedMegatronGPTInference -from deepspeed.model_implementations.transformers.ds_opt import DeepSpeedOPTInference -from deepspeed.model_implementations.transformers.ds_llama2 import DeepSpeedLlama2Inference - -import deepspeed.ops.transformer as transformer_inference -from .layers import LinearLayer, Normalize, EmbeddingLayer, OPTEmbedding, RMSNormalize -import torch -import gc -from deepspeed.accelerator import get_accelerator -import re - - -def load_model_with_checkpoint(r_module, - sd, - mp_replace, - ckpt_type, - ckpt_mp_size, - weight_quantizer=None, - rank=0, - container=None): - error_msgs = [] - - def prefix_check(): - # if keys start with 'model.' or 'transformer.', don't skip level 0 prefix - for key in sd[0].keys(): - # OPT models - if re.match("^model[.]", key): - return False - # BLOOM models - if re.match("^transformer[.]", key): - return False - return True - - skip_level_0_prefix = prefix_check() and container.policy.use_load_prefix - - def transpose(data): - with torch.no_grad(): - data = data.contiguous() - data1 = data.transpose(-1, -2).reshape(-1) - data.reshape(-1).copy_(data1) - data1 = None - return data.reshape(data.shape[-1], data.shape[-2]) - - def load(module, prefix): - args = (sd[0], prefix, {}, True, [], [], error_msgs) - - if hasattr(module, 'weight'): - module.weight = mp_replace.copy(module.weight.data, sd[0][prefix + 'weight']) - if prefix + 'bias' in sd[0].keys(): - if module.bias.data.is_meta: - # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here - module.bias = torch.nn.parameter.Parameter(data=torch.empty_like(module.bias.data, device="cpu"), - requires_grad=module.bias.data.requires_grad) - module.bias = mp_replace.copy(module.bias.data, sd[0][prefix + 'bias']) - args = None - gc.collect() - - def load_transformer_layer(module, prefix): - if ckpt_type == "tp": - - def load_parameters(module, prefix): - for n, p in module.named_parameters(): - if prefix + n in sd[0] and len(n.split('.')) == 1: - if type(sd[0][prefix + n]) is list: - tmp_data, scale = sd[0][prefix + n] - tmp_data = tmp_data - scale = scale.to(get_accelerator().current_device_name()) - # set the quantizer number of groups using the checkpoint scale shape - weight_quantizer.num_groups = scale.shape[0] - else: - tmp_data = sd[0][prefix + n].to(get_accelerator().current_device_name()) - scale = None - src_shape = tmp_data.shape - dst_shape = p.shape - inner_dim = 1 if tmp_data.dtype == torch.int8 else 0 - outer_dim = 0 if tmp_data.dtype == torch.int8 else 1 - if (len(src_shape) == 2 and len(dst_shape) == 2): - if (src_shape[inner_dim] == dst_shape[0] and src_shape[outer_dim] == dst_shape[1]): - if tmp_data.dtype != torch.int8: - p = weight_quantizer.quantize( - transpose(tmp_data) if weight_quantizer.q_int8 else tmp_data) - else: - p = torch.nn.parameter.Parameter(tmp_data, requires_grad=False) - p.scale = scale - setattr(module, n, p) - else: - dim = inner_dim if src_shape[inner_dim] != dst_shape[0] else outer_dim - dim1 = 0 if src_shape[inner_dim] != dst_shape[0] else 1 - if src_shape[dim] > dst_shape[dim1]: - weight_partition = torch.split(tmp_data, dst_shape[dim1], dim=dim)[rank].to( - get_accelerator().current_device_name()) - assert tmp_data.dtype != torch.int8 or scale.numel() > weight_quantizer.num_groups * (rank+1), \ - '''ERROR: We require the quantization scales for larger TP-size when loading INT8 checkpoint!\ - Please use the FP16 checkpoint to generate INT8 checkpoint with the sharding parameters!''' - scale = scale.view(-1)[weight_quantizer.num_groups * (rank + 1):].reshape( - weight_quantizer.num_groups, -1).contiguous() - else: - assert tmp_data.dtype != torch.int8, \ - '''Merging of the checkpoints are not supported when using INT8 checkpoint! \ - Please use a as many GPUs as TP-size for the checkpoint''' - all_data = [ - sd[j][prefix + n] if type(sd[j][prefix + n]) is list else sd[j][prefix + n].to( - get_accelerator().current_device_name()) for j in range(len(sd)) - ] - # Check if the weight tensor is for the QKV parameter - if src_shape[1] == (3 * src_shape[0]) // ckpt_mp_size: - qkv_size = src_shape[outer_dim] // 3 - src_split = [ - torch.split(src[0].data, qkv_size, dim=outer_dim) for src in all_data - ] - - weight_partition = torch.cat([ - torch.cat([qkv_s[i] for qkv_s in src_split], axis=outer_dim) - for i in range(len(src_split[0])) - ], - dim=dim) - else: - weight_partition = torch.cat([ - ad[0].to(get_accelerator().current_device_name()) - if type(ad) is list else ad for ad in all_data - ], - dim=dim) - if tmp_data.dtype == torch.int8: - scale = torch.cat( - [ad[1].to(get_accelerator().current_device_name()) for ad in all_data], - dim=dim) - - if tmp_data.dtype != torch.int8: - weight_partition = weight_quantizer.quantize( - transpose(weight_partition), \ - parallel_dim=(0 if dim == 1 else 1)) if weight_quantizer.q_int8 else \ - weight_quantizer.quantize(weight_partition) - else: - weight_partition = torch.nn.parameter.Parameter(weight_partition, - requires_grad=False) - weight_partition.scale = scale - setattr(module, n, weight_partition) - else: - if src_shape[0] == dst_shape[0]: - p.data.copy_(tmp_data) - else: - if src_shape[0] > dst_shape[0]: - bias_split = torch.split(tmp_data, dst_shape[-1])[rank].to( - get_accelerator().current_device_name()).contiguous() - p.data.copy_(bias_split) - else: - # Check if the weight tensor is for the QKV parameter - if src_shape[0] == (3 * r_module.config.hidden_size) // ckpt_mp_size: - qkv_size = src_shape[0] // 3 - src_split = [ - torch.split(sd[j][prefix + n], qkv_size, dim=0) for j in range(len(sd)) - ] - - p.data.copy_( - torch.cat([ - torch.cat([qkv_s[i] for qkv_s in src_split], axis=0) - for i in range(len(src_split[0])) - ], - dim=0).to(get_accelerator().current_device_name()).contiguous()) - else: - p.data.copy_( - torch.cat([sd[j][prefix + n] for j in range(len(sd))], - dim=0).to(get_accelerator().current_device_name()).contiguous()) - - load_parameters(module, prefix) - for n, child in module.named_children(): - load_parameters(child, prefix + n + '.') - else: - container.load_params(module, sd[0], weight_quantizer, mp_replace, prefix) - - try: - import transformers - OPTLearnedPositionalEmbedding = transformers.models.opt.modeling_opt.OPTLearnedPositionalEmbedding - if hasattr(transformers.models, "llama"): - LlamaRMSNorm = transformers.models.llama.modeling_llama.LlamaRMSNorm - else: - LlamaRMSNorm = None - except: - OPTLearnedPositionalEmbedding = None - try: - from fairscale.nn.model_parallel.layers import ( - ColumnParallelLinear, - ParallelEmbedding, - RowParallelLinear, - ) - except: - ColumnParallelLinear = None - ParallelEmbedding = None - RowParallelLinear = None - try: - from llama.model import RMSNorm - except: - RMSNorm = None - layer_policies = { - nn.Linear: load, - nn.Embedding: load, - nn.LayerNorm: load, - EmbeddingLayer: load, - LinearLayer: load, - Normalize: load, - transformer_inference.DeepSpeedTransformerInference: load_transformer_layer, - DeepSpeedBloomInference: load_transformer_layer, - DeepSpeedGPTInference: load_transformer_layer, - DeepSpeedBERTInference: load_transformer_layer, - DeepSpeedMegatronGPTInference: load_transformer_layer, - DeepSpeedOPTInference: load_transformer_layer, - DeepSpeedLlama2Inference: load_transformer_layer, - OPTLearnedPositionalEmbedding: load, - OPTEmbedding: load, - LlamaRMSNorm: load, - RMSNormalize: load, - ColumnParallelLinear: load, - ParallelEmbedding: load, - RowParallelLinear: load, - RMSNorm: load - } - - all_ds_ids = {} - - def load_module_recursive(module, prefix='', level=0): - for name, child in module.named_children(): - if child.__class__ in layer_policies: - checking_key = prefix + name + '.' - if not any(checking_key in item for item in sd[0].keys()): - if hasattr(child, 'weight') and \ - (hasattr(child.weight, 'ds_id') and \ - child.weight.ds_id in all_ds_ids): - prefix1 = all_ds_ids[child.weight.ds_id] - if child.__class__ is nn.Linear: - child = LinearLayer.from_weights(weight=all_ds_ids[child.weight.ds_id]) - setattr(module, name, child) - continue - child_params = list(child.parameters()) - if len(child_params) > 0 and (child_params[0].numel() == 0 or child_params[0].is_meta): - if child.weight.is_meta: - ds_shape = child.weight.shape - else: - ds_shape = child.weight.ds_shape - if child.__class__ is nn.LayerNorm: - child = Normalize(dim=ds_shape[-1], dtype=child.weight.dtype, eps=child.eps) - setattr(module, name, child) - elif child.__class__ in [nn.Linear, ColumnParallelLinear, RowParallelLinear]: - child = LinearLayer.from_weights(weight_shape=child.weight.shape, - dtype=child.weight.dtype, - bias=child.bias) - setattr(module, name, child) - elif child.__class__ is OPTLearnedPositionalEmbedding: - child = OPTEmbedding(weight_shape=ds_shape) - setattr(module, name, child) - elif child.__class__ in [LlamaRMSNorm, RMSNorm]: - child = RMSNormalize(dim=ds_shape[-1], - dtype=child.weight.dtype, - eps=child.eps if hasattr(child, 'eps') else child.variance_epsilon) - setattr(module, name, child) - else: - ds_id = None - if hasattr(child.weight, 'ds_id'): - ds_id = child.weight.ds_id - child = EmbeddingLayer(weight_shape=ds_shape, dtype=child.weight.dtype) - if ds_id is not None: - all_ds_ids[ds_id] = child.weight - setattr(module, name, child) - layer_policies[child.__class__](child, prefix + name + '.') - else: - load_module_recursive( - child, - prefix if (level == 0 and ckpt_type == 'pp') and skip_level_0_prefix else \ - prefix + name + '.', - level + 1) - - load_module_recursive(r_module) - - for sd_ in sd: - del sd_ - sd = None - gc.collect() diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/module_quantize.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/module_quantize.py deleted file mode 100644 index 1f5b2f8a1d2861a50878e262a61b783a9324ad4e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/module_quantize.py +++ /dev/null @@ -1,72 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - - -def quantize_transformer_layer(orig_layer_impl, model, megatron=False, preln=False): - """ Quantize bert-style transformer layers with DeepSpeed's transformer layer - Arguments: - orig_layer_impl (torch.nn.Module): the original transformer layer implementation to look for, - e.g., transformers.models.bert.modeling_bert.BertLayer or transformers.BertLayer - model (torch.nn.Module): user's nn.module representing their model - - megatron (bool): megatron model-parallel implementation (this is supported for inference only) - preln (bool): does the original layer implementation do pre or post layer norm? - - Note: For Bert kind of models, we inject based on the DeepSpeed-Example models, if not setting huggingface flag. - - Returns: - Updated nn.module with quantized transformer layers - """ - - def quantize_weight(weight): - return weight.to(torch.int8) - - def megatron_layer_quantize(layer): - layer.attention.query_key_value.weight.data = quantize_weight(layer.attention.query_key_value.weight.data) - layer.attention.dense.weight.data = quantize_weight(layer.attention.dense.weight.data) - layer.mlp.dense_h_to_4h.weight.data = quantize_weight(layer.mlp.dense_h_to_4h.weight.data) - layer.mlp.dense_4h_to_h.weight.data = quantize_weight(layer.mlp.dense_4h_to_h.weight.data) - - def bert_layer_quantize(layer): - layer.attention.self.query.weight.data = quantize_weight(layer.attention.self.query.weight.data) - layer.attention.self.key.weight.data = quantize_weight(layer.attention.self.key.weight.data) - layer.attention.self.value.weight.data = quantize_weight(layer.attention.self.value.weight.data) - layer.attention.output.dense.weight.data = quantize_weight(layer.attention.output.dense.weight.data) - if preln: - layer.intermediate.dense_act.weight.data = quantize_weight(layer.intermediate.dense_act.weight.data) - else: - layer.intermediate.dense.weight.data = quantize_weight(layer.intermediate.dense.weight.data) - layer.output.dense.weight.data = quantize_weight(layer.output.dense.weight.data) - - def quantize_fn(child): - if megatron: - # Quantize megatron GPT2 / GPT3 trained model - megatron_layer_quantize(child) - else: - # Quantize either DeepSpeed or HuggingFace trained model - bert_layer_quantize(child) - - return child - - return quantize_module(model=model, orig_class=orig_layer_impl, quantize_fn=quantize_fn) - - -def quantize_module(model, orig_class, quantize_fn): - policy = {orig_class: quantize_fn} - return _quantize_module(model, policy) - - -def _quantize_module(model, policies): - for name, child in model.named_children(): - if child.__class__ in policies: - orig = repr(child) - setattr(model, name, policies[child.__class__](child)) - new = getattr(model, name) - else: - _quantize_module(child, policies) - - return model diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/policy.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/policy.py deleted file mode 100644 index 41df2b85dc0c16fb1b1c21942f5c963fcdc02ad1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/policy.py +++ /dev/null @@ -1,224 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from abc import ABC, abstractmethod -from deepspeed.utils.types import ActivationFuncType, NormType -import torch -from deepspeed.accelerator import get_accelerator - -transformer_param_names = ( - 'attn_qkvw', \ - 'attn_qkvb', \ - 'attn_ow' , \ - 'attn_ob', \ - 'inter_w', \ - 'inter_b', \ - 'output_w', \ - 'output_b', \ - 'attn_nw', \ - 'attn_nb', \ - 'norm_w', \ - 'norm_b') - - -class DSPolicy(ABC): - _orig_layer_class = None - - def __init__(self): - self.cuda_graph_supported = False - - @abstractmethod - def attention(self): - """ - Returns attention qkv and dense parameters - weight: (3*hidden, hidden) and (hidden, hidden) - bias: (3*hidden) and (hidden) - """ - raise NotImplementedError - - -class TransformerPolicy(DSPolicy): - # a static class variable containing the HuggingFace model configuration. - # see e.g., transformers.models.opt.configuration_opt.OPTConfig - hf_model_config = None - - def __init__( - self, - inference=True, - linear_layer=True, - scale_attention=True, - megatron_v2=False, - use_mup=False, - # the type of activation function used in MLP - mlp_act_func_type=ActivationFuncType.GELU, - # applies layer norm before attention if `pre_attn_norm` is set to True - pre_attn_norm=True, - # this flag shows whether or not using prefix in loading the checkpoint - use_load_prefix=False, - # whether or not the qkv is stored in the split-format - split_qkv=True, - # Type of normalization to perform - norm_type=NormType.LayerNorm): - super().__init__() - self.cuda_graph_supported = False - self.inference = inference - self.linear_layer = linear_layer - self.scale_attention = scale_attention - self.is_megatron_v2 = megatron_v2 - self.use_mup = use_mup - self.mlp_act_func_type = mlp_act_func_type - self.pre_attn_norm = pre_attn_norm - self.use_load_prefix = use_load_prefix - self.split_qkv = split_qkv - self.norm_type = norm_type - - @abstractmethod - def attention(self): - """ - Returns attention qkv and dense parameters - weight: (3*hidden, hidden) and (hidden, hidden) - bias: (3*hidden) and (hidden) - """ - raise NotImplementedError - - @abstractmethod - def get_hidden_heads(self): - """ - return hidden_size and number of heads - """ - raise NotImplementedError - - @abstractmethod - def mlp(self): - """ - Returns mlp intermediate and output - weight: (intermediate, hidden) and (hidden, intermediate) - bias: (intermediate) and (hidden) - """ - raise NotImplementedError - - @abstractmethod - def layernorm(self): - """ - Returns LayerNorms used in transformer layer - Post-Attention and pre/post layer norm - gamma and beta with shape: (hidden) - """ - raise NotImplementedError - - -# TODO (lekurile): This function exists in base container as well, consolidate as some point -def transpose(data): - with torch.no_grad(): - data = data.contiguous() - data1 = data.transpose(-1, -2).reshape(-1) - data.reshape(-1).copy_(data1) - data1 = None - return data.reshape(data.shape[-1], data.shape[-2]) - - -# TODO (lekurile): This function exists in megatron feature container as well, consolidate as some point -def _transpose(x, heads=1, mp_replace=None): - heads = heads // mp_replace.mp_size # type: ignore - outer_dim = -1 - attention_head_size = x.shape[outer_dim] // heads - new_x_shape = x.size()[:outer_dim] + (heads, attention_head_size) - x_1 = x.view(*new_x_shape) - (q, k, v) = torch.split(x_1, (x_1.shape[-1] // 3), dim=-1) - if len(q.shape) > 2: - new_shape = (q.shape[0], ) + (-1, ) - return torch.cat((q.reshape(new_shape), k.reshape(new_shape), v.reshape(new_shape)), - dim=outer_dim).reshape(x.shape) - else: - return torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape) - - -# This checks if the parameter exits in the checkpoint file and maybe copies it into the corresponding destination tensor. -# Note that not all parameters are saved in one checkpoint, that's why we always need to check if they exist! -def maybe_copy(module, - sd, - weight_quantizer, - mp_replace, - dst_name, - src_name, - qkv=False, - megatron_v2=False, - split_qkv=False, - heads=1): - if src_name in sd: - dst = getattr(module, dst_name) - tmp = sd[src_name] - if len(dst.shape) == 1: - if split_qkv: - dst = mp_replace.strided_copy(dst, tmp, num_splits=3) - else: - dst = mp_replace.copy(dst, tmp) - if qkv and megatron_v2: - dst = torch.nn.parameter.Parameter(_transpose(dst, heads=heads, mp_replace=mp_replace).contiguous()) - else: - if split_qkv: - dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(tmp if weight_quantizer.q_int8 else \ - (transpose(tmp).contiguous())), num_splits=3, int8=weight_quantizer.q_int8) - else: - if qkv and megatron_v2: - tmp = _transpose(transpose(tmp), heads=heads, mp_replace=mp_replace).contiguous() - if weight_quantizer.q_int8: - tmp = transpose(tmp) - dst = mp_replace.copy(dst, weight_quantizer.quantize(tmp if weight_quantizer.q_int8 else \ - transpose(tmp)), int8=weight_quantizer.q_int8) - setattr(module, dst_name, dst) - - -# Extending the maybe_copy function for when the q, k, and v are in separate parameters! -def maybe_copy_qkv(module, sd, weight_quantizer, mp_replace, dst_name, src_names, split_qkv=False): - if src_names[0] in sd: - q = sd[src_names[0]] - k = sd[src_names[1]] - v = sd[src_names[2]] - qkv_data = torch.cat((q, k, v), dim=0) - dst = getattr(module, dst_name) - if len(dst.shape) == 1: - if split_qkv: - dst = mp_replace.strided_copy(dst, qkv_data.contiguous(), num_splits=3) - else: - dst = mp_replace.copy(dst, qkv_data) - else: - if split_qkv: - dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(qkv_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \ - ((transpose(qkv_data)).contiguous())), num_splits=3, int8=weight_quantizer.q_int8) - else: - dst = mp_replace.copy(dst, weight_quantizer.quantize(qkv_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \ - transpose(qkv_data)), int8=weight_quantizer.q_int8) - setattr(module, dst_name, dst) - - -# Extending the `maybe_copy` function for when mlp1 is in separate parameters for GeGLU -def maybe_copy_geglu(module, sd, weight_quantizer, mp_replace, dst_name, src_names): - if src_names[0] in sd: - reg_proj = sd[src_names[0]] - gate_proj = sd[src_names[1]] - - mlp1_data = torch.cat((reg_proj, gate_proj), dim=0) - dst = getattr(module, dst_name) - - dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(mlp1_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \ - transpose(mlp1_data)), num_splits=2, int8=weight_quantizer.q_int8) - setattr(module, dst_name, dst) - - -def pack_lora_weights(p): - return [ - p.lora_right_weight, \ - p.lora_left_weight, \ - p.lora_scaling - ] - - -def maybe_get_lora(p): - if hasattr(p, 'lora_right_weight'): - lora_param = pack_lora_weights(p) - else: - lora_param = [] - return lora_param diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/replace_module.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/replace_module.py deleted file mode 100644 index ed94a5021feecaac6be0383b104892e80f489b31..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/replace_module.py +++ /dev/null @@ -1,723 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import torch -import tqdm -import deepspeed -import deepspeed.ops.transformer as transformer_inference -from deepspeed.ops.transformer.inference.diffusers_attention import DeepSpeedDiffusersAttention -from deepspeed.ops.transformer.inference.diffusers_transformer_block import DeepSpeedDiffusersTransformerBlock -from deepspeed.ops.transformer.inference.diffusers_2d_transformer import Diffusers2DTransformerConfig -from deepspeed.accelerator import get_accelerator -from .replace_policy import replace_policies, generic_policies -from .auto_tp import AutoTP, ReplaceWithTensorSlicing, Loading -from .layers import TensorParallelOcShardConv2d, TensorParallelIcShardConv2d -from deepspeed.module_inject.layers import is_autotp_training_mode -from deepspeed import comm as dist -from deepspeed.module_inject.tp_shard import set_num_kv_heads, set_n_embd, set_num_attention_heads, set_tp_grain_size - -from .load_checkpoint import load_model_with_checkpoint -import time - -from .utils import policy_to_ds_container -import gc - - -def get_transformer_name(replaced_module): - from .containers import supported_models - from torch.nn import ModuleList - transformer_name = '' - for n, c in replaced_module.named_children(): - if c.__class__ in supported_models: - transformer_name += n + '.' - for name, child in c.named_children(): - if child.__class__ is ModuleList: - transformer_name += name - break - break - return transformer_name - - -class GroupQuantizer: - - def __init__(self, q_int8=True, group_size=1, num_bits=8, num_groups=0): - self.group_size = group_size - self.num_bits = num_bits - self.q_int8 = q_int8 - - self.num_groups = num_groups - - def quantize(self, inputs, qkv=True, count=1, parallel_dim=0): - if not self.q_int8 or not qkv: - inputs = torch.nn.Parameter(inputs, requires_grad=False) - inputs.scale = torch.empty(1) - return inputs - q_range = 2**self.num_bits - num_groups = self.num_groups if self.num_groups > 0 else inputs.shape[0] // self.group_size - inputs = inputs.to(get_accelerator().current_device_name()) - input_flat = inputs.reshape(num_groups, -1).contiguous() - input_min = torch.min(input_flat, dim=1, keepdim=True)[0].float() - input_max = torch.max(input_flat, dim=1, keepdim=True)[0].float() - scale = torch.max(input_min.abs(), input_max.abs()) * 2.0 / (q_range) - input_flat = (input_flat / scale).round().clamp(-q_range // 2, q_range // 2 - 1) - inputs_q = input_flat.reshape(inputs.shape).to(torch.int8).contiguous() - out = torch.nn.Parameter(inputs_q, requires_grad=False) - inputs_split = inputs.split(inputs.shape[parallel_dim] // 2, dim=parallel_dim) - input_flat = [inputs_split[i].reshape(num_groups, -1).contiguous() for i in range(2)] - input_min = [torch.min(input_flat[i], dim=1, keepdim=True)[0].float() for i in range(2)] - input_max = [torch.max(input_flat[i], dim=1, keepdim=True)[0].float() for i in range(2)] - scale1 = [(torch.max(input_min[i].abs(), input_max[i].abs()) * 2.0 / (q_range)).squeeze().unsqueeze(0) - for i in range(2)] - - out.scale = torch.cat([scale.squeeze().unsqueeze(0), scale1[0], scale1[1]], dim=0).reshape(num_groups, - -1).contiguous() - return out - - -def _module_match(module): - for policy in generic_policies: - policy = policy() - if policy.match(module): - return policy - return None - - -def generic_injection(module, dtype=None, enable_cuda_graph=True): - - def replace_attn(child, policy): - policy_attn = policy.attention(child) - if policy_attn is None: - return child - if len(policy_attn) == 5: - qkvw, attn_ow, attn_ob, hidden_size, heads = policy_attn - else: - qw, kw, vw, attn_ow, attn_ob, hidden_size, heads = policy_attn - - config = transformer_inference.DeepSpeedInferenceConfig( - hidden_size=hidden_size, - heads=heads, - dtype=dtype, - triangular_masking=False, - max_out_tokens=4096, - ) - attn_module = DeepSpeedDiffusersAttention(config) - - def transpose(data): - data = data.contiguous() - data.reshape(-1).copy_(data.transpose(-1, -2).contiguous().reshape(-1)) - data = data.reshape(data.shape[-1], data.shape[-2]) - data.to(get_accelerator().current_device_name()) - return data - - if len(policy_attn) == 5: - attn_module.attn_qkvw.data = transpose(qkvw.data) - else: - attn_module.attn_qkvw = None - attn_module.attn_qw.data = transpose(qw.data) - attn_module.attn_kw.data = transpose(kw.data) - attn_module.attn_vw.data = transpose(vw.data) - - attn_module.attn_qkvb = None - attn_module.attn_ow.data = transpose(attn_ow.data) - attn_module.attn_ob.data.copy_(attn_ob.data.to(get_accelerator().current_device_name())) - return attn_module - - def replace_attn_block(child, policy): - config = Diffusers2DTransformerConfig() - return DeepSpeedDiffusersTransformerBlock(child, config) - - if isinstance(module, torch.nn.Module): - pass - else: - if dtype not in [torch.float16, torch.half]: - raise ValueError("Generic injection only supported with FP16") - - try: - import diffusers - if hasattr(diffusers.models.attention, 'CrossAttention'): - cross_attention = diffusers.models.attention.CrossAttention - else: - cross_attention = diffusers.models.attention_processor.Attention - attention_block = diffusers.models.attention.BasicTransformerBlock - new_policies = { - cross_attention: replace_attn, - attention_block: replace_attn_block, - } - except ImportError: - new_policies = {} - - #replace_transformer_layer(None, - # module.text_encoder, - # training=False, - # replace_with_kernel_inject=True, - # triangular_masking=True, - # max_out_tokens=8192) - from ..model_implementations.transformers.clip_encoder import DSClipEncoder - cg_encoder = DSClipEncoder(module.text_encoder, enable_cuda_graph=enable_cuda_graph) - setattr(module, 'text_encoder', cg_encoder) - for name in module.__dict__.keys(): - sub_module = getattr(module, name) - policy = _module_match(sub_module) - - if policy is not None: - - def _replace_module(module, policy): - for name, child in module.named_children(): - _replace_module(child, policy) - if child.__class__ in new_policies: - replaced_module = new_policies[child.__class__](child, policy) - setattr(module, name, replaced_module) - - _replace_module(sub_module, policy) - new_module = policy.apply(sub_module, enable_cuda_graph=enable_cuda_graph) - print(f"**** found and replaced {name} w. {type(new_module)}") - setattr(module, name, new_module) - - -container_g = None - - -def replace_transformer_layer(orig_layer_impl, model, checkpoint_dict, config, model_config): - """ Replace bert-style transformer layers with DeepSpeed's transformer layer - Arguments: - orig_layer_impl (torch.nn.Module): the original transformer layer implementation to look for, - e.g., transformers.models.bert.modeling_bert.BertLayer or transformers.BertLayer - model (torch.nn.Module): user's nn.module representing their model - checkpoint_dict: Dictionary for checkpoint passed from the Inference Engine - config: top-level DS Inference config defined in inference/config.py - model_config: HuggingFace model config passed from the inference/engine.py - Returns: - Updated nn.module with replaced transformer layers - """ - # defining globals as internally defined functions inherit these everywhere - quantize = (config.dtype == torch.int8) - # todo: Refactor later. In future, let's minimize the style used above and use config.** instead - - linear_layer_setting = None - ''' - linear_layer_setting (tuple of modules) [Optional]: shows which two classes are used for linear layers and embedding layers - ''' - micro_batch_size = -1 - seed = -1 - local_rank = -1 - - mp_replace = ReplaceWithTensorSlicing(mp_group=config.tensor_parallel.tp_group, - mp_size=config.tensor_parallel.tp_size) #, out_dim=0, in_dim=1) - - def replace_with_policy(child, policy_cls, triangular_masking, inference=False, layer_id=0): - policy = policy_cls(child, inference=inference) - if not policy.cuda_graph_supported: - # policy says cuda graph is not supported raise an error if set - assert not config.enable_cuda_graph, "cuda graph is not supported with this model, please disable" - - from deepspeed.moe.layer import MoE - moe = False - if hasattr(child, 'mlp') and isinstance(child.mlp, MoE): - num_experts = child.mlp.num_experts - moe = True - - # 1. Create a model-specific container object using the policy object. - _container = policy_to_ds_container(policy=policy, - config=config, - model_config=model_config, - layer_id=layer_id, - child=child) - _container.set_moe(moe) - - # 2. Set the tensor parallelism config - _container.set_tensor_parallel_config(config.tensor_parallel.tp_size, config.tensor_parallel.tp_group) - - # 3. Initialize tensors - _container.initialize_tensors() - - # 4. deal with data types -- needs refactor to use dtype instead of fp16 - if config.dtype in [torch.float16, torch.bfloat16, torch.int8]: - _container.convert_to_required_dtype() - - # 5. Set the quantization config - quantizer = GroupQuantizer(q_int8=quantize) - _container.set_quantization_config(quantizer) - - # 6. create a DS Inference config object - _container.create_ds_model_config() - - # 7. use the config and create the module - _container.create_module() - - # 8. transpose the weights and bias if needed - _container.transpose() - - # 9. deal with tensor parallelism. - _container.apply_tensor_parallelism(mp_replace) - - # 10. copy the tensors from the model-specific container to the new module - _container.copy_data_to_new_module() - - # 11. set global for generic checkpoint loading - global container_g - - if container_g is None: - container_g = _container - - return _container.module - - def replace_wo_policy(module, all_reduce_linears, prefix="", state_dict=None): - #mp_replace = ReplaceWithTensorSlicing(mp_group=config.tensor_parallel.tp_group) - - # 1. Create AutoTP object - _autotp = AutoTP(module, all_reduce_linears, prefix, state_dict, linear_layer_setting, orig_layer_impl, - config.keep_module_on_host) - - # 2. Set the tensor parallelism config - _autotp.set_tensor_parallel_config(config.tensor_parallel.tp_size, config.tensor_parallel.tp_group) - - # 3. Try to get num_key_heads from model_config.num_key_value_heads - if hasattr(model_config, "vision_config"): - if "MllamaVisionEncoderLayer" in str(module): - num_kv_heads = _autotp.get_model_num_kv_heads(model_config.vision_config) - elif hasattr(model_config, "text_config"): - num_kv_heads = _autotp.get_model_num_kv_heads(model_config.text_config) - else: - num_kv_heads = _autotp.get_model_num_kv_heads(model_config) - else: - num_kv_heads = _autotp.get_model_num_kv_heads(model_config) - - # 4. When we have num_kv_heads defined, uneven division is possible, otherwise enforce even division - set_num_kv_heads(num_kv_heads) - - # 4.1 Get n_embd - n_embd = None - multi_query_n_embd_names = ['n_embd', 'hidden_size'] - for name in multi_query_n_embd_names: - if hasattr(model_config, name): - n_embd = getattr(model_config, name) - if n_embd != None: - break - - # 4.2 set n_embd - set_n_embd(n_embd) - - # 4.3 set attention_heads - if hasattr(model_config, 'num_attention_heads'): - set_num_attention_heads(getattr(model_config, 'num_attention_heads')) - - # 4.4 set tp_grain_size - set_tp_grain_size(config.tensor_parallel.tp_grain_size) - - # 5. Set linear policies - _autotp.update_linear_policies() - - # 6. Replace modules - if "lm_head" in all_reduce_linears or "embed_out" in all_reduce_linears: - return _autotp._replace_last_linear_module(module) - return _autotp._replace_module(module) - - def replace_fn(child, _policy, layer_id=0, prefix="", state_dict=None): - training = False # todo: refactor this part to go in the config - if training: - # copy relevant state from child -> new module - new_module = replace_with_policy(child, _policy, config.triangular_masking) - - else: - # copy relevant state from child -> new module - if not is_autotp_training_mode() and config.replace_with_kernel_inject: - new_module = replace_with_policy(child, - _policy, - config.triangular_masking, - inference=True, - layer_id=layer_id) - else: - new_module = replace_wo_policy(child, _policy, prefix=prefix, state_dict=state_dict) - - return new_module - - def set_lm_head(module): - if is_autotp_training_mode(): - # we need to handle autoTP training mode separately. - return - - embedding_weight = None - for n, p in module.named_parameters(): - if "word_embeddings." in n or "embed_tokens." in n or "wte." in n: - embedding_weight = p - if embedding_weight is not None and hasattr(module, "lm_head") and hasattr( - module.lm_head, "weight") and module.lm_head.weight.is_meta: - module.lm_head.weight = embedding_weight - # enable tensor parallel for the last linear - if hasattr(module, "lm_head") and hasattr(module.lm_head, "weight") and isinstance( - module.lm_head, torch.nn.Linear): - module = replace_wo_policy(module, ("lm_head", ), 0, "lm_head") - elif hasattr(module, "embed_out") and hasattr(module.embed_out, "weight") and isinstance( - module.embed_out, torch.nn.Linear): - module = replace_wo_policy(module, ("embed_out", ), 0, "embed_out") - elif hasattr(module, "language_model") and hasattr(module.language_model, "lm_head"): - module = replace_wo_policy(module.language_model, ("lm_head", ), 0, "lm_head") - return module - - def conv2d_parallel_shard_weights(model, rank, world_size): - # add conv policy - shard_oc_name = ["conv1"] - shard_ic_name = ["conv2"] - for name, sub_m in model.named_children(): - for l_name, l_sub_m in sub_m.named_children(): - if l_name in shard_oc_name: - TPConv2d = TensorParallelOcShardConv2d( - l_sub_m, - rank, - world_size, - ) - setattr(sub_m, l_name, TPConv2d) - if l_name in shard_ic_name: - TPConv2d = TensorParallelIcShardConv2d( - l_sub_m, - rank, - world_size, - ) - setattr(sub_m, l_name, TPConv2d) - conv2d_parallel_shard_weights(sub_m, rank, world_size) - - if checkpoint_dict is not None and not config.replace_with_kernel_inject: - # AutoTP shard loading - checkpoint = checkpoint_dict["checkpoints"] - pbar = tqdm.tqdm(total=len(checkpoint), desc=f"Loading {len(checkpoint)} checkpoint shards") - for i in range(len(checkpoint)): - checkpoint_file = os.path.join(config.base_dir, checkpoint[i]) - replaced_module = replace_module(model=model, - orig_class=orig_layer_impl, - replace_fn=replace_fn, - _replace_policy=config.injection_policy_tuple, - checkpoint=checkpoint_file) - pbar.update(1) - gc.collect() - # conv2d tp module replace - # Now is for yuan model. Add model list and conv policy to decide whether to replace conv. - if 'Yuan' in str(replaced_module): - conv2d_parallel_shard_weights(replaced_module, dist.get_rank(), dist.get_world_size()) - else: - replaced_module = replace_module(model=model, - orig_class=orig_layer_impl, - replace_fn=replace_fn, - _replace_policy=config.injection_policy_tuple) - # AutoTP default set lm_head tp - if not config.replace_with_kernel_inject: - replaced_module = set_lm_head(replaced_module) - - quantizer = GroupQuantizer(q_int8=quantize) - world_size = dist.get_world_size() if dist.is_initialized() else 1 - rank = dist.get_rank() if dist.is_initialized() else 0 - if checkpoint_dict is not None and config.replace_with_kernel_inject: - assert container_g.ckpt_load_enabled, \ - f"Meta Tensor checkpoint loading not supported in {container_g.__class__.__name__} container" - start_time = time.time() - checkpoint = checkpoint_dict['checkpoints'] - ckpt_list = checkpoint["tp"] if type(checkpoint) is dict else checkpoint - ckpt_type = checkpoint_dict.get('parallelization', 'pp') - ckpt_mp_size = checkpoint_dict.get('tp_size', len(ckpt_list)) - ckpt_mp_size = checkpoint_dict.get('mp_size', ckpt_mp_size) - base_dir1 = checkpoint_dict.get('base_dir', config.base_dir) - - if ckpt_type == 'pp' and type(checkpoint) is list: - pbar = tqdm.tqdm(total=len(checkpoint), desc=f"Loading {len(checkpoint)} checkpoint shards") - - for i in range(len(checkpoint)): - sd = [torch.load(os.path.join(base_dir1, checkpoint[i]), map_location='cpu', weights_only=False)] - load_model_with_checkpoint(replaced_module, - sd, - mp_replace, - ckpt_type, - ckpt_mp_size, - quantizer, - container=container_g) - pbar.update(1) - else: - num_checkpoints = len(ckpt_list) // ckpt_mp_size - tp_split_size = (world_size / ckpt_mp_size) - sd_offset = int(rank / tp_split_size) - sd_count = int((rank + max(1, tp_split_size)) / tp_split_size) - sd_offset - pbar = tqdm.tqdm(total=num_checkpoints, desc=f"Loading {num_checkpoints} checkpoint shards") - for i in range(num_checkpoints): - pbar.update(1) - ckpt_index = i * ckpt_mp_size + sd_offset - ckpt_files = [ - os.path.join(base_dir1, ckpt_list[ckpt_index + j]) if base_dir1 else ckpt_list[ckpt_index + j] - for j in range(sd_count) - ] - sds = [torch.load(ckpt_file, map_location='cpu', weights_only=False) for ckpt_file in ckpt_files] - load_model_with_checkpoint(replaced_module, - sds, - mp_replace, - ckpt_type, - ckpt_mp_size, - quantizer, - int(rank % tp_split_size), - container=container_g) - sds = [None for _ in sds] - gc.collect() - - if "non_tp" in checkpoint: - pbar = tqdm.tqdm(total=len(checkpoint["non_tp"]), - desc=f"Loading {len(checkpoint['non_tp'])} checkpoint shards") - - for i in range(len(checkpoint["non_tp"])): - pbar.update(1) - ckpt_file = os.path.join(base_dir1, - checkpoint["non_tp"][i]) if base_dir1 else checkpoint["non_tp"][i] - sds = [torch.load(ckpt_file, map_location='cpu', weights_only=False)] - load_model_with_checkpoint(replaced_module, - sds, - mp_replace, - ckpt_type, - ckpt_mp_size, - quantizer, - int(rank % tp_split_size), - container=container_g) - sds = [None for _ in sds] - gc.collect() - set_lm_head(replaced_module) - print(f"checkpoint loading time at rank {rank}: {time.time()-start_time} sec") - - if not is_autotp_training_mode() and config.save_mp_checkpoint_path is not None: - from collections import OrderedDict - import json - num_partitions = 8 - - if checkpoint_dict is None: - ckpt_name = "ds_model" - try: - from transformers.models.bloom.modeling_bloom import BloomForCausalLM - if isinstance(model, BloomForCausalLM): - ckpt_name = "bloom" - except ImportError: - ckpt_name = "ds_model" - else: - ckpt_name = checkpoint_dict['type'] - if dist.is_initialized(): - dist.barrier() - transformer_name = get_transformer_name(replaced_module) - non_tp_ckpt_name = f'non-tp.pt' - ckpt_files = [non_tp_ckpt_name] - os.makedirs(config.save_mp_checkpoint_path, exist_ok=True) - - if not dist.is_initialized() or dist.get_rank() == 0: - print("Saving tp-sharded checkpoints") - torch.save( - OrderedDict({ - k: v - for k, v in dict(replaced_module.state_dict()).items() if transformer_name not in k - }), f'{config.save_mp_checkpoint_path}/{non_tp_ckpt_name}') - - dtype_reprs = { - torch.float32: 'float32', - torch.float16: 'float16', - torch.int8: 'int8', - torch.bfloat16: 'bfloat16' - } - - ckpt_config = json.dumps({ - 'type': ckpt_name, - 'base_dir': f'{config.save_mp_checkpoint_path}', - 'checkpoints': { - "non_tp": ckpt_files, - "tp": [f'tp_{r:0>2d}_{m:0>2d}.pt' for m in range(num_partitions) for r in range(world_size)] - }, - 'version': 1.0, - 'parallelization': 'tp', - 'tp_size': world_size, - 'dtype': dtype_reprs[config.dtype] - }) - with open(f"{config.save_mp_checkpoint_path}/ds_inference_config.json", "w") as cfg: - cfg.write(ckpt_config) - - rep_sd = replaced_module.state_dict() - for n, p in replaced_module.named_parameters(): - if hasattr(p, 'scale'): - rep_sd[n] = [p, p.scale] - keys = list(rep_sd.keys()) - partition_size = (len(keys) // num_partitions + 1) - for m in range(num_partitions): - torch.save( - OrderedDict({ - k: [rep_sd[k], rep_sd[k].scale] if hasattr(rep_sd[k], 'scale') else rep_sd[k] - for k in keys[m * partition_size:(m + 1) * partition_size] if transformer_name in k - }), f'{config.save_mp_checkpoint_path}/tp_{rank:0>2d}_{m:0>2d}.pt') - - return replaced_module - - -def revert_transformer_layer(orig_layer_impl, model, config, preln=False): - """ Revert DeepSpeed's transformer layer back to original bert-style transformer layer - Arguments: - orig_layer_impl (torch.nn.Module): the original transformer layer implementation that was replaced, - e.g., transformers.models.bert.modeling_bert.BertLayer or transformers.BertLayer - model (torch.nn.Module): user's nn.module representing their model - config (dict): model config containing hidden size, attention heads, etc. - Returns: - Updated nn.module with original bert-style transformer layers - """ - - def replace_fn(child, _replace_policy, layer_id): - #from turing.nvidia_modelingpreln import BertLayer - orig_module = orig_layer_impl(config) - - # copy relevant state from child -> original module - qkvw = child.attn_qkvw.data - qkvb = child.attn_qkvb.data - - qw, kw, vw = torch.chunk(qkvw, 3, axis=0) - qb, kb, vb = torch.chunk(qkvb, 3, axis=0) - - orig_module.attention.self.query.weight.data = qw - orig_module.attention.self.query.bias.data = qb - orig_module.attention.self.key.weight.data = kw - orig_module.attention.self.key.bias.data = kb - orig_module.attention.self.value.weight.data = vw - orig_module.attention.self.value.bias.data = vb - - orig_module.attention.output.dense.weight.data = child.attn_ow.data - orig_module.attention.output.dense.bias.data = child.attn_ob.data - - attn_ln_w = child.attn_nw.data - attn_ln_b = child.attn_nb.data - if preln: - orig_module.PostAttentionLayerNorm.weight.data = attn_ln_w - orig_module.PostAttentionLayerNorm.bias.data = attn_ln_b - else: - orig_module.attention.output.LayerNorm.weight.data = attn_ln_w - orig_module.attention.output.LayerNorm.bias.data = attn_ln_b - - inter_ff_w = child.inter_w.data - inter_ff_b = child.inter_b.data - if preln: - orig_module.intermediate.dense_act.weight.data = inter_ff_w - orig_module.intermediate.dense_act.bias.data = inter_ff_b - else: - orig_module.intermediate.dense.weight.data = inter_ff_w - orig_module.intermediate.dense.bias.data = inter_ff_b - - orig_module.output.dense.weight.data = child.output_w.data - orig_module.output.dense.bias.data = child.output_b.data - - transformer_ln_w = child.norm_w.data - transformer_ln_b = child.norm_b.data - if preln: - orig_module.PreAttentionLayerNorm.weight.data = transformer_ln_w - orig_module.PreAttentionLayerNorm.bias.data = transformer_ln_b - else: - orig_module.output.LayerNorm.weight.data = transformer_ln_w - orig_module.output.LayerNorm.bias.data = transformer_ln_b - return orig_module - - return replace_module(model=model, - orig_class=deepspeed.DeepSpeedTransformerLayer, - replace_fn=replace_fn, - _replace_policy=None) - - -def replace_module(model, orig_class, replace_fn, _replace_policy, checkpoint=None): - """ Scan the model for instances of ``orig_clas:`` to replace using ``replace_fn``. - Arguments: - model (torch.nn.Module): the model to augment - orig_class (torch.nn.Module): the module to search for - replace_fn (method): a method to convert instances of ``orig_class`` to the - desired type and return a new instance. - Returns: - A modified ``model``. - """ - sd = None - if checkpoint is not None: - if checkpoint.endswith(".safetensors"): - from safetensors.torch import load_file - sd = load_file(checkpoint) - else: - sd = torch.load(checkpoint, map_location='cpu', weights_only=False) - - policy = {} - if orig_class is not None: - policy.update({orig_class: (replace_fn, _replace_policy)}) - else: - for plcy in replace_policies: - # instantiate a throw-away policy in order to populate the _orig_layer_class - _ = plcy(None) - if isinstance(plcy._orig_layer_class, list): - for orig_layer_class in plcy._orig_layer_class: - policy.update({orig_layer_class: (replace_fn, plcy)}) - elif plcy._orig_layer_class is not None: - policy.update({plcy._orig_layer_class: (replace_fn, plcy)}) - assert len(policy.items()) > 0,\ - "No default policy found! Please specify your policy injection_policy (like {BertLayer:HFBEertLayerPolicy})." +\ - "You can find some samples here: https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/module_inject/replace_policy.py" - - replaced_module, _ = _replace_module(model, policy, state_dict=sd) - return replaced_module - - -from ..pipe import PipelineModule - -import re - - -def skip_level_0_prefix(model, state_dict): - model = str(model) - key = re.search(r": (.*?)Model", model) - if key is None: - key = re.search(r": (.*?)Stack", model) - if key is None: - key = re.match(r"(.*?)Model", model) - # if keys start with 'model.', don't skip level 0 prefix - if state_dict is not None: - for item in state_dict.keys(): - if re.match("^model[.]", item): - return False - if key is not None and key.group(1).lower() in ["bloom", "opt"]: - return True - return False - - -def _replace_module(model, policies, prefix='', layer_id=0, level_id=0, state_dict=None): - """ Traverse model's children recursively and apply any transformations in ``policies``. - Arguments: - model (torch.nn.Module): model to augment - policies (dict): Mapping of source class to replacement function. - Returns: - Modified ``model``. - """ - for name, child in model.named_children(): - if child.__class__ in policies: - replaced_module = policies[child.__class__][0](child, - policies[child.__class__][-1], - layer_id, - prefix=prefix + name, - state_dict=state_dict) - setattr(model, name, replaced_module) - if isinstance(model, PipelineModule): - assert hasattr(model, 'forward_funcs'),\ - "we require pipe-module to have the list of fwd_functions" - model.forward_funcs[model.fwd_map[name]] = replaced_module - layer_id += 1 - else: - checking_key = prefix + name + '.' - if Loading.is_load_module(child) and state_dict is not None: - if any(checking_key in item for item in state_dict): - Loading.load( - child, - state_dict, - checking_key, - ) - else: - continue - if len(child._buffers) != 0 and state_dict is not None: - Loading.load_buffer(child, state_dict, checking_key) - _, layer_id = _replace_module(child, - policies, - prefix if level_id == 0 and skip_level_0_prefix(model, state_dict) else \ - prefix + name + '.', - layer_id=layer_id, - level_id=level_id + 1, - state_dict=state_dict) - - # Add the reset_cache func to the model, so that it can be called in the beginning of text-generation. - model.reset_cache = transformer_inference.DeepSpeedTransformerInference.reset_cache - return model, layer_id diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/replace_policy.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/replace_policy.py deleted file mode 100644 index 2c06e31aaa41ae61b5091a7766b709fac2c534f5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/replace_policy.py +++ /dev/null @@ -1,30 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .containers import HFGPT2LayerPolicy -from .containers import HFBertLayerPolicy -from .containers import BLOOMLayerPolicy -from .containers import HFGPTJLayerPolicy -from .containers import HFGPTNEOLayerPolicy -from .containers import GPTNEOXLayerPolicy -from .containers import HFOPTLayerPolicy -from .containers import MegatronLayerPolicy -from .containers import HFDistilBertLayerPolicy -from .containers import HFCLIPLayerPolicy -from .containers import LLAMALayerPolicy -from .containers import UNetPolicy -from .containers import VAEPolicy -from .containers import LLAMA2LayerPolicy -from .containers import InternLMLayerPolicy - -# transformer-based policies -replace_policies = [ - HFBertLayerPolicy, HFGPTNEOLayerPolicy, GPTNEOXLayerPolicy, HFGPTJLayerPolicy, MegatronLayerPolicy, - HFGPT2LayerPolicy, BLOOMLayerPolicy, HFOPTLayerPolicy, HFCLIPLayerPolicy, HFDistilBertLayerPolicy, - LLAMALayerPolicy, LLAMA2LayerPolicy, InternLMLayerPolicy -] - -# non-transformer-based policies -generic_policies = [UNetPolicy, VAEPolicy] diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/tp_shard.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/tp_shard.py deleted file mode 100644 index ded262edcf61bbb39d5d3e0c1575430aa8d9b080..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/tp_shard.py +++ /dev/null @@ -1,74 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed import comm as dist -global num_kv_heads - - -def set_num_kv_heads(num): - global num_kv_heads - num_kv_heads = num - - -def set_num_attention_heads(num): - global num_attention_heads - num_attention_heads = num - - -def set_n_embd(num): - global n_embd - n_embd = num - - -def set_tp_grain_size(num): - global tp_grain_size - tp_grain_size = num - - -def get_num_kv_heads(): - global num_kv_heads - if 'num_kv_heads' in globals(): - return num_kv_heads - return None - - -def get_num_attention_heads(): - global num_attention_heads - return num_attention_heads - - -def get_shard_size(total_size, mp_size, name=None, rank=None): - global num_kv_heads - last_linear = ["lm_head", "embed_out"] - # MoE MLP layer use near even division will get better perf. - moe_mlp_layer = ["gate_proj", "up_proj", "down_proj", "w1", "w2", "w3"] - not_moe_mlp_layer = True - if name != None and any(s in str(name) for s in moe_mlp_layer): - not_moe_mlp_layer = False - # When we have num_kv_heads defined, uneven division is possible, otherwise enforce near even division - if rank == None: - rank = dist.get_rank() - if num_kv_heads != None and total_size % num_kv_heads == 0 and "mlp" not in str(name) and str( - name) not in last_linear and not_moe_mlp_layer: - my_slices = (num_kv_heads // mp_size) + (1 if rank < (num_kv_heads % mp_size) else 0) - return total_size * my_slices // num_kv_heads - else: - if total_size >= tp_grain_size: - grain_size = total_size // tp_grain_size - return (grain_size // mp_size + (1 if rank < (grain_size % mp_size) else 0)) * tp_grain_size - else: - return total_size // mp_size + (1 if rank < (total_size % mp_size) else 0) - - -def get_n_embd(): - global n_embd - return n_embd - - -def get_shard_size_list(total_size, mp_size, name=None): - shard_sizes = [] - for i in range(mp_size): - shard_sizes.append(get_shard_size(total_size, mp_size, name, i)) - return shard_sizes diff --git a/venv/lib/python3.10/site-packages/deepspeed/module_inject/utils.py b/venv/lib/python3.10/site-packages/deepspeed/module_inject/utils.py deleted file mode 100644 index 42822128f9e11c84660dcd57a2473753b3bc6642..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/module_inject/utils.py +++ /dev/null @@ -1,49 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.utils import log_dist - - -# helper function to map between DS policies and DS containers -def policy_to_ds_container(**kwargs): - from .containers import HFGPT2LayerPolicy, DS_GPT2Container - from .containers import HFBertLayerPolicy, DS_BERTContainer - from .containers import BLOOMLayerPolicy, DS_BloomContainer - from .containers import HFGPTJLayerPolicy, DS_GPTJContainer - from .containers import HFGPTNEOLayerPolicy, DS_GPTNEOContainer - from .containers import GPTNEOXLayerPolicy, DS_GPTNEOXContainer - from .containers import HFOPTLayerPolicy, DS_OPTContainer - from .containers import MegatronLayerPolicy, DS_MegatronGPTContainer - from .containers import HFDistilBertLayerPolicy, DS_DistilBERTContainer - from .containers import LLAMALayerPolicy, DS_LLAMAContainer - from .containers import LLAMA2LayerPolicy, DS_LLAMA2Container - from .containers import InternLMLayerPolicy, DS_InternLMContainer - - policy_to_container = { - HFGPT2LayerPolicy: DS_GPT2Container, - HFBertLayerPolicy: DS_BERTContainer, - BLOOMLayerPolicy: DS_BloomContainer, - HFGPTJLayerPolicy: DS_GPTJContainer, - HFGPTNEOLayerPolicy: DS_GPTNEOContainer, - GPTNEOXLayerPolicy: DS_GPTNEOXContainer, - HFOPTLayerPolicy: DS_OPTContainer, - MegatronLayerPolicy: DS_MegatronGPTContainer, - HFDistilBertLayerPolicy: DS_DistilBERTContainer, - LLAMALayerPolicy: DS_LLAMAContainer, - LLAMA2LayerPolicy: DS_LLAMA2Container, - InternLMLayerPolicy: DS_InternLMContainer - } - - container = None - policy = kwargs['policy'] - assert policy is not None, "Policy cannot be None" - policy_type = type(policy) - - if policy_type not in policy_to_container: - log_dist(f"Policy type {policy_type} not supported", [0]) - else: - container = policy_to_container[policy_type](**kwargs) - - return container diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/moe/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/moe/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index d99cae3bfa6284be26c8634dc8be986b6d12c7b4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/experts.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/experts.cpython-310.pyc deleted file mode 100644 index d87ce8b98d1e3481bf68b16f751ab04966f44283..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/experts.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/layer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/layer.cpython-310.pyc deleted file mode 100644 index f019848b9ea7c715d7444807aa41d9060e900cba..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/layer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/mappings.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/mappings.cpython-310.pyc deleted file mode 100644 index e2eab9b3b9c03d4b496b95104e8fced542a42f72..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/mappings.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/sharded_moe.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/sharded_moe.cpython-310.pyc deleted file mode 100644 index 206b49b79e173fb78caf9fff8418f17b329a9353..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/sharded_moe.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index f49a2bb61dc48ea5552b25d04c884a0df96672a5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/moe/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/experts.py b/venv/lib/python3.10/site-packages/deepspeed/moe/experts.py deleted file mode 100644 index 0863221d7edf93770af126e3fa1a78be3c7f4289..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/moe/experts.py +++ /dev/null @@ -1,38 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import copy -from typing import List, Optional - -import torch -from torch import nn - - -class Experts(nn.Module): - - def __init__(self, expert: nn.Module, num_local_experts: int = 1, expert_group_name: Optional[str] = None) -> None: - super(Experts, self).__init__() - - self.deepspeed_experts = nn.ModuleList([copy.deepcopy(expert) for _ in range(num_local_experts)]) - self.num_local_experts = num_local_experts - - # TODO: revisit allreduce for moe.gate... - for expert in self.deepspeed_experts: - # TODO: Create param groups to handle expert + data case (e.g. param.group = moe_group) - for param in expert.parameters(): - param.allreduce = False - param.group_name = expert_group_name - - def forward(self, inputs: torch.Tensor) -> torch.Tensor: - chunks = inputs.chunk(self.num_local_experts, dim=1) - expert_outputs: List[torch.Tensor] = [] - - for chunk, expert in zip(chunks, self.deepspeed_experts): - out = expert(chunk) - if isinstance(out, tuple): - out = out[0] # Ignore the bias term for now - expert_outputs += [out] - - return torch.cat(expert_outputs, dim=1) diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/layer.py b/venv/lib/python3.10/site-packages/deepspeed/moe/layer.py deleted file mode 100644 index 6777788ab885a618544ced9ed9448e08c5159751..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/moe/layer.py +++ /dev/null @@ -1,132 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Optional, Tuple - -import torch -from torch import nn -from torch.nn import functional as F - -from deepspeed.utils import groups, log_dist -from .experts import Experts -from .sharded_moe import MOELayer, TopKGate - - -class MoE(nn.Module): - """Initialize an MoE layer. - - Arguments: - hidden_size (int): the hidden dimension of the model, importantly this is also the input and output dimension. - expert (nn.Module): the torch module that defines the expert (e.g., MLP, torch.linear). - num_experts (int, optional): default=1, the total number of experts per layer. - ep_size (int, optional): default=1, number of ranks in the expert parallel world or group. - k (int, optional): default=1, top-k gating value, only supports k=1 or k=2. - capacity_factor (float, optional): default=1.0, the capacity of the expert at training time. - eval_capacity_factor (float, optional): default=1.0, the capacity of the expert at eval time. - min_capacity (int, optional): default=4, the minimum capacity per expert regardless of the capacity_factor. - use_residual (bool, optional): default=False, make this MoE layer a Residual MoE (https://arxiv.org/abs/2201.05596) layer. - noisy_gate_policy (str, optional): default=None, noisy gate policy, valid options are 'Jitter', 'RSample' or 'None'. - drop_tokens (bool, optional): default=True, whether to drop tokens - (setting to False is equivalent to infinite capacity). - use_rts (bool, optional): default=True, whether to use Random Token Selection. - use_tutel (bool, optional): default=False, whether to use Tutel optimizations (if installed). - enable_expert_tensor_parallelism (bool, optional): default=False, whether to use tensor parallelism for experts - top2_2nd_expert_sampling (bool, optional): default=True, whether to perform sampling for 2nd expert - """ - - def __init__(self, - hidden_size: int, - expert: nn.Module, - num_experts: int = 1, - ep_size: int = 1, - k: int = 1, - capacity_factor: float = 1.0, - eval_capacity_factor: float = 1.0, - min_capacity: int = 4, - use_residual: bool = False, - noisy_gate_policy: Optional[str] = None, - drop_tokens: bool = True, - use_rts: bool = True, - use_tutel: bool = False, - enable_expert_tensor_parallelism: bool = False, - top2_2nd_expert_sampling: bool = True) -> None: - - super(MoE, self).__init__() - - self.use_residual = use_residual - self.enable_expert_tensor_parallelism = enable_expert_tensor_parallelism - assert num_experts % ep_size == 0, f"Number of experts ({num_experts}) should be divisible by expert parallel size ({ep_size})" - self.ep_size = ep_size - self.expert_group_name = f"ep_size_{self.ep_size}" - self.num_experts = num_experts - self.num_local_experts = num_experts // self.ep_size - - log_dist( - f'Creating MoE layer with num_experts: {num_experts} | num_local_experts: {self.num_local_experts} | expert_parallel_size: {self.ep_size}', - [0]) - - assert noisy_gate_policy is None or noisy_gate_policy in ['None', 'Jitter', 'RSample'], \ - 'Unsupported noisy_gate_policy: ' + noisy_gate_policy - - experts = Experts(expert, self.num_local_experts, self.expert_group_name) - self.deepspeed_moe = MOELayer(TopKGate(hidden_size, num_experts, k, capacity_factor, eval_capacity_factor, - min_capacity, noisy_gate_policy, drop_tokens, use_rts, None, - top2_2nd_expert_sampling), - experts, - self.expert_group_name, - self.ep_size, - self.num_local_experts, - use_tutel=use_tutel) - if self.use_residual: - self.mlp = expert - # coefficient is used for weighted sum of the output of expert and mlp - self.coefficient = nn.Linear(hidden_size, 2) - - def set_deepspeed_parallelism(self, use_data_before_expert_parallel_: bool = False) -> None: - self._create_process_groups(use_data_before_expert_parallel_=use_data_before_expert_parallel_) - - def _create_process_groups(self, use_data_before_expert_parallel_: bool = False) -> None: - # Create process group for a layer if needed - if self.expert_group_name not in groups._get_expert_parallel_group_dict(): - print(f"No existing process group found, creating a new group named: {self.expert_group_name}") - if (groups.mpu is None) or (not self.enable_expert_tensor_parallelism): - # Condition 1 - no groups.mpu means no tensor parallelism - # Condition 2 - disabling expert tensor parallelism on purpose - groups._create_expert_and_data_parallel( - self.ep_size, use_data_before_expert_parallel_=use_data_before_expert_parallel_) - else: - # expert tensor parallelism is enabled - groups._create_expert_data_and_model_parallel( - self.ep_size, mpu=groups.mpu, use_data_before_expert_parallel_=use_data_before_expert_parallel_) - # Set the group handle for the MOELayer (deepspeed_moe) object - self.deepspeed_moe._set_ep_group(groups._get_expert_parallel_group(self.expert_group_name)) - - def forward(self, - hidden_states: torch.Tensor, - used_token: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ MoE forward - - Arguments: - hidden_states (Tensor): input to the layer - used_token (Tensor, optional): default: None, mask only used tokens - - Returns: - A tuple including output, gate loss, and expert count. - - * output (Tensor): output of the model - - * l_aux (Tensor): gate loss value - - * exp_counts (Tensor): expert count - """ - output = self.deepspeed_moe(hidden_states, used_token) - if self.use_residual: - # Residual MoE - output_mlp = self.mlp(hidden_states) - if isinstance(output_mlp, tuple): - output_mlp = output_mlp[0] # Ignore the bias term for now - coef = self.coefficient(hidden_states) - coef = F.softmax(coef, dim=-1) - output = output * coef[..., 0:1] + output_mlp * coef[..., 1:] - return output, self.deepspeed_moe.l_aux, self.deepspeed_moe.exp_counts diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/mappings.py b/venv/lib/python3.10/site-packages/deepspeed/moe/mappings.py deleted file mode 100644 index e57f66b85193d86734c186a916baa6da033a90a1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/moe/mappings.py +++ /dev/null @@ -1,118 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# The file has been adapted from the following Megatron-LM file: -# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/mpu/mappings.py -# Git commit hash: 9dc3c42a84aa656f583703cf8b6b4f79f712b796 -# We retain the following copyright from the original files: - -# Copyright (c) 2020, NVIDIA CORPORATION. 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 deepspeed -from deepspeed.utils.bwc import (bwc_tensor_model_parallel_world_size, bwc_tensor_model_parallel_rank, - bwc_tensor_model_parallel_group) - - -def _gather_tokens(input_, dim=0): - """Gather tensors and concatenate them along a dimension""" - mpu = deepspeed.utils.groups.mpu - - input_ = input_.contiguous() - world_size = bwc_tensor_model_parallel_world_size(mpu) - if world_size == 1: - return input_ - - gather_buffer = torch.empty(world_size * input_.numel(), dtype=input_.dtype, device=input_.device) - deepspeed.comm.all_gather_into_tensor(gather_buffer, input_, group=bwc_tensor_model_parallel_group(mpu)) - if dim == 0: - shape = list(input_.size()) - shape[0] = shape[0] * world_size - output = gather_buffer.view(shape) - else: - tensor_list = [ - gather_buffer.narrow(0, - input_.numel() * i, input_.numel()).view_as(input_) for i in range(world_size) - ] - # Note: torch.cat already creates a contiguous tensor. - output = torch.cat(tensor_list, dim=dim).contiguous() - - return output - - -def _drop_tokens(input_, dim=0): - """Divide a tensor among the tensor parallel ranks""" - mpu = deepspeed.utils.groups.mpu - - total_chunks = bwc_tensor_model_parallel_world_size(mpu) - if total_chunks == 1: - return input_ - this_chunk = bwc_tensor_model_parallel_rank(mpu) - assert input_.shape[ - dim] % total_chunks == 0, f"input dimension {dim} ({input_.shape[dim]}) is not divisible by tensor parallel world size ({total_chunks})" - chunk_size = input_.shape[dim] // total_chunks - - return torch.narrow(input_, dim, this_chunk * chunk_size, chunk_size) - - -class _GatherTokens(torch.autograd.Function): - """All gather tokens among the tensor parallel ranks""" - - @staticmethod - def symbolic(graph, input_, dim): - return _gather_tokens(input_, dim) - - @staticmethod - def forward(ctx, input_, dim): - ctx.dim = dim - return _gather_tokens(input_, dim) - - @staticmethod - def backward(ctx, grad_output): - return _drop_tokens(grad_output, ctx.dim), None - - -class _DropTokens(torch.autograd.Function): - "Divide tokens equally among the tensor parallel ranks" - - @staticmethod - def symbolic(graph, input_, dim): - return _drop_tokens(input_, dim) - - @staticmethod - def forward(ctx, input_, dim): - ctx.dim = dim - return _drop_tokens(input_, dim) - - @staticmethod - def backward(ctx, input_): - return _gather_tokens(input_, ctx.dim), None - - -def gather_tokens(input_, dim=0): - mpu = deepspeed.utils.groups.mpu - if mpu is None or bwc_tensor_model_parallel_world_size(mpu) == 1: - # no tensor parallelism for non-experts - return input_ - return _GatherTokens.apply(input_, dim) - - -def drop_tokens(input_, dim=0): - mpu = deepspeed.utils.groups.mpu - if mpu is None or bwc_tensor_model_parallel_world_size(mpu) == 1: - # no tensor parallelism for non-experts - return input_ - return _DropTokens.apply(input_, dim) diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/sharded_moe.py b/venv/lib/python3.10/site-packages/deepspeed/moe/sharded_moe.py deleted file mode 100644 index 5aacd1a11a1ad12c8d125c55851c13f6c19df747..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/moe/sharded_moe.py +++ /dev/null @@ -1,678 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -The file has been adapted from two fairscale files: - (1) https://github.com/facebookresearch/fairscale/blob/master/fairscale/nn/moe/moe_layer.py - (2) https://github.com/facebookresearch/fairscale/blob/master/fairscale/nn/moe/top2gate.py - Git commit hash: 34df606902a240567a0d898037ece55c2f1336cf - We retain the following license from the original files: -""" - -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. -# -# This source code is licensed under the BSD license found in the -# LICENSE file in the root directory of this source tree. - -from deepspeed.utils.timer import SynchronizedWallClockTimer -from deepspeed.utils import logger -from deepspeed.utils.bwc import bwc_tensor_model_parallel_world_size -from typing import Callable, Dict, TYPE_CHECKING, Any, Optional, Tuple, Union - -import torch -from torch import Tensor -from torch.nn import Module -import torch.nn.functional as F -from deepspeed.utils import groups -from .mappings import drop_tokens, gather_tokens - -if TYPE_CHECKING: - Base = Module[Tensor] -else: - Base = Module - -TOPK_GATE_TIMER = 'topk_gate' -MOE_TIMER = 'moe' -FIRST_ALLTOALL_TIMER = '1st_a2a' -SECOND_ALLTOALL_TIMER = '2nd_a2a' - -uniform_map: Dict[torch.device, Callable] = {} -gumbel_map: Dict[torch.device, Callable] = {} -exp_selection_uniform_map: Dict[torch.device, Callable] = {} - -try: - # To enable Tutel MoE optimizations: - # python3 -m pip install --user --upgrade git+https://github.com/deepspeedai/tutel@v0.1.x - from tutel import moe as tutel_moe - TUTEL_INSTALLED = True -except: - # Fail silently so we don't spam logs unnecessarily if user isn't using tutel - TUTEL_INSTALLED = False - pass - - -def multiplicative_jitter(x, device: torch.device, epsilon=1e-2): - """ - Modified from switch transformer paper. mesh transformers - Multiply values by a random number between 1-epsilon and 1+epsilon. - Makes models more resilient to rounding errors introduced by bfloat16. - This seems particularly important for logits. - Args: - x: a torch.tensor - device: torch.device - epsilon: a floating point value - Returns: - a jittered x. - """ - if epsilon == 0: - return x - uniform = uniform_map.get(device) - if uniform is None: - uniform = torch.distributions.uniform.Uniform(low=torch.tensor(1.0 - epsilon, device=device), - high=torch.tensor(1.0 + epsilon, - device=device)).rsample # type: ignore - uniform_map[device] = uniform - return x * uniform(x.shape) - - -def gumbel_rsample(shape: Tuple, device: torch.device) -> Tensor: - gumbel = gumbel_map.get(device) - if gumbel is None: - one = torch.tensor(1.0, device=device) - zero = torch.tensor(0.0, device=device) - gumbel = torch.distributions.gumbel.Gumbel(zero, one).rsample # type: ignore - gumbel_map[device] = gumbel - return gumbel(shape) - - -from deepspeed import comm as dist - -# einsum dimensions: (g)roup, (s)equence, (e)xpert, (m)odel, (c)apacity -# See https://arxiv.org/pdf/2006.16668.pdf for details. - - -# Based on https://github.com/pytorch/pytorch/pull/40762 -class _AllToAll(torch.autograd.Function): - - @staticmethod - def forward(ctx: Any, group: dist.ProcessGroup, input: Tensor) -> Tensor: # type: ignore - ctx.group = group - input = input.contiguous() - output = torch.empty_like(input) - dist.all_to_all_single(output, input, group=group) - return output - - @staticmethod - def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor]: - return (None, _AllToAll.apply(ctx.group, *grad_output)) - - -# einsum rewrites are on par or more performant -# switch can be bubbled up in future -USE_EINSUM = True - - -# einsum dimensions: (g)roup, (s)equence, (e)xpert, (m)odel, (c)apacity -# See https://arxiv.org/pdf/2006.16668.pdf for details. -def einsum(rule, a, b): - if USE_EINSUM: - return torch.einsum(rule, a, b) - elif rule == 's,se->se': - return a.reshape(a.shape[0], -1) * b - elif rule == 'se,sc->sec': - return a.unsqueeze(2) * b.unsqueeze(1) - elif rule == 'se,se->s': - return torch.bmm(a.unsqueeze(1), b.unsqueeze(2)).reshape(-1) - elif rule == 'se,sec->sec': - return a.unsqueeze(2) * b - elif rule == 'sec,sm->ecm': - s = a.shape[0] - e = a.shape[1] - c = a.shape[2] - m = b.shape[1] - return torch.matmul(a.reshape(s, -1).t(), b).reshape(e, c, m) - elif rule == 'sec,ecm->sm': - return torch.matmul(a.reshape(a.shape[0], -1), b.reshape(-1, b.shape[-1])) - elif rule == 'ks,ksm->sm': - k = b.shape[0] - s = b.shape[1] - m = b.shape[2] - # [k, s] -> [s, k] -> [s, 1, k] - a = a.t().unsqueeze(1) - # [k,s,m] -> [k, sm] -> [sm, k] -> [s, m, k] - b = b.reshape(k, -1).t().reshape(s, m, k) - # bmm([s, 1, k], [s, m, k]^t) -> [s, m, 1] - return torch.bmm(a, b.transpose(1, 2)).squeeze(2) - else: - return torch.einsum(rule, a, b) - - -# The following functions are extracted and scripted -# because otherwise during a torch.jit.trace, the non-Tensor -# values used in the calculations get recorded as constants. -# torch.jit.script coerces them into Tensors and preserves -# their dynamic shapes. This enables ONNX export. -# We can't script the entire top1gating function because it -# includes stateful caching logic which is incompatible with ONNX. - - -@torch.jit.script -def _capacity(gates: Tensor, capacity_factor: Tensor, min_capacity: Tensor) -> Tensor: - # gates has shape of SE - num_tokens = gates.shape[0] - num_experts = gates.shape[1] - # to(torch.int64) works around a bug in torch.onnx.export: - # it should cast k to int64 when converting torch.topk but it doesn't. - capacity = torch.ceil((num_tokens / num_experts) * capacity_factor).to(torch.int64) - if capacity < min_capacity: - capacity = min_capacity.to(torch.int64) - return capacity - - -@torch.jit.script -def _top_idx(source, k): - return torch.topk(source, k=k, dim=0)[1] - - -@torch.jit.script -def _one_hot_to_float(x, num_classes): - return F.one_hot(x, num_classes=num_classes).float() - - -def top1gating(logits: Tensor, - capacity_factor: float, - min_capacity: int, - used_token: Tensor = None, - noisy_gate_policy: Optional[str] = None, - drop_tokens: bool = True, - use_rts: bool = True, - ep_group: Union[torch.distributed.ProcessGroup, None] = None, - use_tutel: bool = False) -> Tuple[Tensor, Tensor, Tensor, Tensor]: - """Implements Top1Gating on logits.""" - if noisy_gate_policy == 'RSample': - logits_w_noise = logits + gumbel_rsample(logits.shape, device=logits.device) - # everything is in fp32 in this function - - gates = F.softmax(logits, dim=1) - capacity = _capacity(gates, torch.tensor(capacity_factor), torch.tensor(min_capacity)) - - # Create a mask for 1st's expert per token - # noisy gating - indices1_s = torch.argmax(logits_w_noise if noisy_gate_policy == 'RSample' else gates, dim=1) - num_experts = int(gates.shape[1]) - mask1 = F.one_hot(indices1_s, num_classes=num_experts) - - # mask only used tokens - if used_token is not None: - mask1 = einsum("s,se->se", used_token, mask1) - - # gating decisions - exp_counts = torch.sum(mask1, dim=0).detach().to(logits.device) - - # if we don't want to drop any tokens - if not drop_tokens: - new_capacity = torch.max(exp_counts).to(logits.device) - # Communicate across expert processes to pick the maximum capacity. - if ep_group is not None: - dist.all_reduce(new_capacity, op=dist.ReduceOp.MAX, group=ep_group) - if groups._get_expert_model_parallel_world_size() == 1: - # If the non-expert is tensor-parallel, we need to pad the capacity to 'tp'. - # This is since we are going to activate drop_tokens() to drop duplicate tokens. - tp = 1 if groups.mpu is None else bwc_tensor_model_parallel_world_size(mpu=groups.mpu) - new_capacity = torch.ceil(new_capacity / tp).mul(tp).to(new_capacity.dtype) - # Make sure the capacity value does not exceed the number of tokens. - capacity = min(new_capacity, torch.tensor(mask1.size(0)).to(new_capacity.device)) - - # Compute l_aux - me = torch.mean(gates, dim=0) - ce = torch.mean(mask1.float(), dim=0) - l_aux = torch.sum(me * ce) * num_experts - - # Random Token Selection - if use_rts: - uniform = exp_selection_uniform_map.get(logits.device) - if uniform is None: - uniform = torch.distributions.uniform.Uniform(low=torch.tensor(0.0, device=logits.device), - high=torch.tensor(1.0, device=logits.device)).rsample - exp_selection_uniform_map[logits.device] = uniform - - mask1_rand = mask1 * uniform(mask1.shape) - else: - mask1_rand = mask1 - - assert logits.shape[ - 0] >= min_capacity, "No. of tokens (batch-size) should be greater than min_capacity. Either set min_capacity to 0 or increase your batch size." - - top_idx = _top_idx(mask1_rand, capacity) - - new_mask1 = mask1 * torch.zeros_like(mask1).scatter_(0, top_idx, 1) - mask1 = new_mask1 - - if use_tutel: - # Tutel doesn't support index values masked with zero - # so we need to replace masked indices with -1 - indices_mask = mask1.sum(dim=1) * num_experts - 1 - indices1_s = torch.min(indices1_s, indices_mask) - - # Compute locations in capacity buffer - if use_tutel: - locations1 = tutel_moe.fast_cumsum_sub_one(mask1) - else: - locations1 = torch.cumsum(mask1, dim=0) - 1 - - if use_tutel: - gates1_s = (gates * mask1).sum(dim=1) - locations1_s = torch.sum(locations1 * mask1, dim=1) - return l_aux, capacity, num_experts, [ - indices1_s, - ], [ - locations1_s, - ], [ - gates1_s, - ], exp_counts - - # Store the capacity location for each token - locations1_s = torch.sum(locations1 * mask1, dim=1) - - # Normalize gate probabilities - mask1_float = mask1.float() - gates = gates * mask1_float - - locations1_sc = _one_hot_to_float(locations1_s, capacity) - combine_weights = einsum("se,sc->sec", gates, locations1_sc) - - dispatch_mask = combine_weights.bool() - - return l_aux, combine_weights, dispatch_mask, exp_counts - - -def top2gating(logits: Tensor, - capacity_factor: float, - min_capacity: int, - drop_tokens: bool = True, - ep_group: Union[torch.distributed.ProcessGroup, None] = None, - top2_2nd_expert_sampling: bool = True) -> Tuple[Tensor, Tensor, Tensor, Tensor]: - """Implements Top2Gating on logits.""" - # everything is in fp32 in this function - gates = F.softmax(logits, dim=1) - - # Create a mask for 1st's expert per token - indices1_s = torch.argmax(gates, dim=1) - num_experts = int(gates.shape[1]) - mask1 = F.one_hot(indices1_s, num_classes=num_experts) - - if top2_2nd_expert_sampling: - # Create a mask for 2nd's expert per token using Gumbel-max trick - # https://timvieira.github.io/blog/post/2014/07/31/gumbel-max-trick/ - logits += gumbel_rsample(logits.shape, device=logits.device) - - # Replace top-expert with min value - logits_except1 = logits.masked_fill(mask1.bool(), float("-inf")) - indices2_s = torch.argmax(logits_except1, dim=1) - mask2 = F.one_hot(indices2_s, num_classes=num_experts) - - # Compute locations in capacity buffer - locations1 = torch.cumsum(mask1, dim=0) - 1 - locations2 = torch.cumsum(mask2, dim=0) - 1 - # Update 2nd's location by accounting for locations of 1st - locations2 += torch.sum(mask1, dim=0, keepdim=True) - - # Compute l_aux - me = torch.mean(gates, dim=0) - ce = torch.mean(mask1.float(), dim=0) - l_aux = torch.mean(me * ce) * num_experts * num_experts - - # gating decisions - exp_counts = torch.sum(mask1 + mask2, dim=0).detach().to(logits.device) - - if drop_tokens: - # Calculate configured capacity and remove locations outside capacity from mask - capacity = _capacity(gates, torch.tensor(capacity_factor * 2), torch.tensor(min_capacity)) - mask1 *= torch.lt(locations1, capacity) - mask2 *= torch.lt(locations2, capacity) - else: - # Do not drop tokens - set capacity according to current expert assignments - new_capacity = torch.max(exp_counts) - if ep_group is not None: - dist.all_reduce(new_capacity, op=dist.ReduceOp.MAX, group=ep_group) - if groups._get_expert_model_parallel_world_size() == 1: - # If the non-expert is tensor-parallel, we need to pad the capacity to 'tp'. - # This is since we are going to activate drop_tokens() to drop duplicate tokens. - tp = 1 if groups.mpu is None else bwc_tensor_model_parallel_world_size(mpu=groups.mpu) - new_capacity = torch.ceil(new_capacity / tp).mul(tp).to(new_capacity.dtype) - capacity = new_capacity - - # Store the capacity location for each token - locations1_s = torch.sum(locations1 * mask1, dim=1) - locations2_s = torch.sum(locations2 * mask2, dim=1) - - # Normalize gate probabilities - mask1_float = mask1.float() - mask2_float = mask2.float() - gates1_s = einsum("se,se->s", gates, mask1_float) - gates2_s = einsum("se,se->s", gates, mask2_float) - denom_s = gates1_s + gates2_s - # Avoid divide-by-zero - denom_s = torch.clamp(denom_s, min=torch.finfo(denom_s.dtype).eps) - gates1_s /= denom_s - gates2_s /= denom_s - - # Calculate combine_weights and dispatch_mask - gates1 = einsum("s,se->se", gates1_s, mask1_float) - gates2 = einsum("s,se->se", gates2_s, mask2_float) - locations1_sc = _one_hot_to_float(locations1_s, capacity) - locations2_sc = _one_hot_to_float(locations2_s, capacity) - combine1_sec = einsum("se,sc->sec", gates1, locations1_sc) - combine2_sec = einsum("se,sc->sec", gates2, locations2_sc) - combine_weights = combine1_sec + combine2_sec - dispatch_mask = combine_weights.bool() - - return l_aux, combine_weights, dispatch_mask, exp_counts - - -def topkgating( - logits: Tensor, - k: int, - capacity_factor: float, - min_capacity: int, - drop_tokens: bool = True, - ep_group: Union[torch.distributed.ProcessGroup, None] = None, - drop_policy: str = "probs", -) -> Tuple[Tensor, Tensor, Tensor, Tensor]: - """Implements TopKGating on logits.""" - - # everything is in fp32 in this function - # get topk gates - top_gate, top_idx = torch.topk(logits, k=k, dim=1) - # gating decisions - gates = F.softmax(logits, dim=1) - num_experts = int(gates.shape[1]) - - # get topk mask - topk_masked_gates = torch.zeros_like(logits).scatter(1, top_idx, top_gate) - - mask = torch.zeros_like(gates, dtype=torch.bool).scatter_(1, top_idx, 1) - - exp_counts = torch.sum(mask, dim=0).detach().to(logits.device) - - # Compute l_aux - me = torch.mean(gates, dim=0) - ce = torch.mean(mask.float(), dim=0) - l_aux = torch.mean(me * ce) * num_experts * num_experts / k - - if drop_tokens: - # Calculate configured capacity and remove locations outside capacity from mask - capacity = _capacity(gates, torch.tensor(capacity_factor * k), torch.tensor(min_capacity)) - # update mask and locations by capacity - - if drop_policy == 'probs': - capacity_probs, capacity_indices = torch.topk(topk_masked_gates, k=capacity, dim=0, sorted=False) - capacity_mask = torch.zeros_like(logits).scatter(0, capacity_indices, 1) - mask = torch.logical_and(mask, capacity_mask) - locations = torch.cumsum(mask, dim=0) - 1 - - elif drop_policy == "position": - locations = torch.cumsum(mask, dim=0) - 1 - mask *= torch.lt(locations, capacity) - else: - raise ValueError(f"Invalid drop_policy: {drop_policy}") - - else: - # Do not drop tokens - set capacity according to current expert assignments - new_capacity = torch.max(exp_counts) - if ep_group is not None: - dist.all_reduce(new_capacity, op=dist.ReduceOp.MAX, group=ep_group) - if groups._get_expert_model_parallel_world_size() == 1: - # If the non-expert is tensor-parallel, we need to pad the capacity to 'tp'. - # This is since we are going to activate drop_tokens() to drop duplicate tokens. - tp = 1 if groups.mpu is None else bwc_tensor_model_parallel_world_size(mpu=groups.mpu) - new_capacity = torch.ceil(new_capacity / tp).mul(tp).to(new_capacity.dtype) - capacity = new_capacity - locations = torch.cumsum(mask, dim=0) - 1 - - # normalize gates - gates_masked = gates * mask - gates_s = torch.sum(gates_masked, dim=-1, keepdim=True) - denom_s = torch.clamp(gates_s, min=torch.finfo(gates_masked.dtype).eps) - gates_masked = gates_masked / denom_s - - # dispatch_mask - locations_sc = _one_hot_to_float((locations * mask), capacity) - - combine_weights = torch.einsum("se,sec->sec", gates_masked, locations_sc) - - dispatch_mask = combine_weights.bool() - - return l_aux, combine_weights, dispatch_mask, exp_counts - - -class TopKGate(Module): - """Gate module which implements Top2Gating as described in Gshard_. - :: - - gate = TopKGate(model_dim, num_experts) - l_aux, combine_weights, dispatch_mask = gate(input) - - .. Gshard_: https://arxiv.org/pdf/2006.16668.pdf - - Args: - model_dim (int): - size of model embedding dimension - num_experts (int): - number of experts in model - """ - - wg: torch.nn.Linear - - def __init__(self, - model_dim: int, - num_experts: int, - k: int = 1, - capacity_factor: float = 1.0, - eval_capacity_factor: float = 1.0, - min_capacity: int = 8, - noisy_gate_policy: Optional[str] = None, - drop_tokens: bool = True, - use_rts: bool = True, - ep_group: Union[torch.distributed.ProcessGroup, None] = None, - top2_2nd_expert_sampling: bool = True) -> None: - super().__init__() - - self.wg = torch.nn.Linear(model_dim, num_experts, bias=False) - self.ep_group = ep_group - self.k = k - self.capacity_factor = capacity_factor - self.eval_capacity_factor = eval_capacity_factor - self.min_capacity = min_capacity - self.noisy_gate_policy = noisy_gate_policy - self.timers = SynchronizedWallClockTimer() - self.wall_clock_breakdown = False - self.gate_time = 0.0 - self.drop_tokens = drop_tokens - self.use_rts = use_rts - self.top2_2nd_expert_sampling = top2_2nd_expert_sampling - - def _set_ep_group(self, ep_group): - assert self.ep_group is None, f'Attempting to override an existing ep_group' - self.ep_group = ep_group - - def forward(self, - input: torch.Tensor, - used_token: torch.Tensor = None, - use_tutel: bool = False) -> Tuple[Tensor, Tensor, Tensor]: # type: ignore - - if self.wall_clock_breakdown: - self.timers(TOPK_GATE_TIMER).start() - - input_fp32 = input.float() - # input jittering - if self.noisy_gate_policy == 'Jitter' and self.training: - input_fp32 = multiplicative_jitter(input_fp32, device=input.device) - logits = torch.nn.functional.linear(input_fp32, weight=self.wg.weight.float(), bias=None) - - if self.k == 1: - gate_output = top1gating(logits, self.capacity_factor if self.training else self.eval_capacity_factor, - self.min_capacity, used_token, self.noisy_gate_policy if self.training else None, - self.drop_tokens, self.use_rts, self.ep_group, use_tutel) - - elif self.k == 2: - gate_output = top2gating(logits, self.capacity_factor if self.training else self.eval_capacity_factor, - self.min_capacity, self.drop_tokens, self.ep_group, self.top2_2nd_expert_sampling) - else: - gate_output = topkgating(logits, self.k, - self.capacity_factor if self.training else self.eval_capacity_factor, - self.min_capacity, self.drop_tokens, self.ep_group) - - if self.wall_clock_breakdown: - self.timers(TOPK_GATE_TIMER).stop() - self.gate_time = self.timers(TOPK_GATE_TIMER).elapsed(reset=False) - - return gate_output - - -class MOELayer(Base): - """MOELayer module which implements MixtureOfExperts as described in Gshard_. - :: - - gate = TopKGate(model_dim, num_experts) - moe = MOELayer(gate, expert) - output = moe(input) - l_aux = moe.l_aux - - .. Gshard_: https://arxiv.org/pdf/2006.16668.pdf - - Args: - gate (torch.nn.Module): - gate network - expert (torch.nn.Module): - expert network - """ - - def __init__(self, - gate: Module, - experts: Module, - ep_group_name, - ep_size, - num_local_experts: int, - use_tutel: bool = False) -> None: - super().__init__() - self.gate = gate - self.experts = experts - self.ep_group = None - self.ep_size = ep_size - self.ep_group_name = ep_group_name - self.num_local_experts = num_local_experts - self.time_falltoall = 0.0 - self.time_salltoall = 0.0 - self.time_moe = 0.0 - self.timers = SynchronizedWallClockTimer() - self.wall_clock_breakdown = False - - self.use_tutel = use_tutel and TUTEL_INSTALLED and gate.k == 1 - - if self.use_tutel: - logger.info('Using Tutel optimizations.') - elif use_tutel and not TUTEL_INSTALLED: - logger.warning("Tutel optimization requested but not installed. " - "Proceeding without Tutel.") - elif use_tutel and TUTEL_INSTALLED and gate.k != 1: - logger.warning("To enable Tutel optimization, use top-1 instead of top-2 gate. " - "Proceeding without Tutel.") - - def _set_ep_group(self, ep_group): - self.ep_group = ep_group - self.gate._set_ep_group(ep_group) - - def forward(self, *input: Tensor, **kwargs: Any) -> Tensor: - - if self.wall_clock_breakdown: - self.timers(MOE_TIMER).start() - - # Implement Algorithm 2 from GShard paper. - d_model = input[0].shape[-1] - - # Initial implementation -> Reshape into S tokens by dropping sequence dimension. - # Reshape into G groups so that each group can distribute tokens equally - # group_size = kwargs['group_size'] if 'group_size' in kwargs.keys() else 1 - reshaped_input = input[0].reshape(-1, d_model) - - if self.use_tutel: - self.l_aux, C, E, indices_, locations_, gates_, self.exp_counts = self.gate(reshaped_input, input[1], True) - S, M = reshaped_input.size(0), reshaped_input.size(1) - - if not hasattr(self, '_tutel_dispatcher'): - self._tutel_dispatcher = tutel_moe.fast_dispatcher(E, C, M, dispatch_dtype=reshaped_input.dtype) - self._tutel_dispatcher.update(indices_, locations_, gates_, capacity=C) - dispatched_input = self._tutel_dispatcher.encode(reshaped_input) - else: - self.l_aux, combine_weights, dispatch_mask, self.exp_counts = self.gate(reshaped_input, input[1]) - dispatched_input = einsum("sec,sm->ecm", dispatch_mask.type_as(input[0]), reshaped_input) - - if self.wall_clock_breakdown: - self.timers(FIRST_ALLTOALL_TIMER).start() - - tensor_model_world_size = bwc_tensor_model_parallel_world_size(groups.mpu) - if tensor_model_world_size > 1: - # If the non-expert is tensor-parallel, - # Whether expert is tensor-parallel or not , it will create - # duplicate tokens on the tensor-parallel ranks. - # drop duplicate tokens also doubles up as a communication - # optimization as we are reducing the all-to-all communication volume. - # 1: for not tensor-parallel expert,drop duplicate tokens to ensure - # both correctness and reduce all-to-all communication. - # 2: for tensor-parallel expert,drop duplicate tokens to reduce all-to-all - # communication volume,before expert execution, it is necessary to perform - # an allgather to ensure correctness, - dispatched_input = drop_tokens(dispatched_input, dim=1) - - dispatched_input = _AllToAll.apply(self.ep_group, dispatched_input) - - if self.wall_clock_breakdown: - self.timers(FIRST_ALLTOALL_TIMER).stop() - self.time_falltoall = self.timers(FIRST_ALLTOALL_TIMER).elapsed(reset=False) - - if tensor_model_world_size > 1 and groups._get_expert_model_parallel_world_size() > 1: - # if both expert and non-expert are tensor-parallel - # the dropped duplicate tokens need to be gathered on each - # tensor parallel rank again to ensure correctness - dispatched_input = gather_tokens(dispatched_input, dim=1) - - # Re-shape after all-to-all: ecm -> gecm - dispatched_input = dispatched_input.reshape(self.ep_size, self.num_local_experts, -1, d_model) - expert_output = self.experts(dispatched_input) - # Re-shape before drop_tokens: gecm -> ecm - expert_output = expert_output.reshape(self.ep_size * self.num_local_experts, -1, d_model) - if tensor_model_world_size > 1 and groups._get_expert_model_parallel_world_size() > 1: - # if both expert and non-expert are tensor-parallel - # drop duplicate tokens to ensure both correctness - # and reduce all-to-all communication. - expert_output = drop_tokens(expert_output, dim=1) - - if self.wall_clock_breakdown: - self.timers(SECOND_ALLTOALL_TIMER).start() - - expert_output = _AllToAll.apply(self.ep_group, expert_output) - - if self.wall_clock_breakdown: - self.timers(SECOND_ALLTOALL_TIMER).stop() - self.time_salltoall = self.timers(SECOND_ALLTOALL_TIMER).elapsed(reset=False) - - if tensor_model_world_size > 1: - # the dropped duplicate tokens need to be gathered on each - # tensor parallel rank again for the tensor-parallel - # non-expert of the next layer. - expert_output = gather_tokens(expert_output, dim=1) - - if self.use_tutel: - combined_output = self._tutel_dispatcher.decode(expert_output.view(E * C, M)) - else: - combined_output = einsum("sec,ecm->sm", combine_weights.type_as(input[0]), expert_output) - - a = combined_output.reshape(input[0].shape) - - if self.wall_clock_breakdown: - self.timers(MOE_TIMER).stop() - self.time_moe = self.timers(MOE_TIMER).elapsed(reset=False) - - return a diff --git a/venv/lib/python3.10/site-packages/deepspeed/moe/utils.py b/venv/lib/python3.10/site-packages/deepspeed/moe/utils.py deleted file mode 100644 index 20866378efac72c96a7e2f56bdf97e5b7b4effb9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/moe/utils.py +++ /dev/null @@ -1,182 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from collections import defaultdict -from typing import Any, Dict, List, Set, Tuple, Union, cast - -import torch -from torch import nn - -from .layer import MoE - - -def has_moe_layers(m: nn.Module) -> Tuple[bool, int]: - has_moe = False - num_experts = 0 - - for module in m.modules(): - if isinstance(module, MoE): - has_moe = True - num_experts = module.num_experts - break - return has_moe, num_experts - - -def is_moe_param(param: torch.Tensor) -> bool: - if hasattr(param, "allreduce") and not param.allreduce: - return True - return False - - -def split_params_into_shared_and_expert_params( - params: List[torch.nn.Parameter]) -> Tuple[List[torch.nn.Parameter], List[torch.nn.Parameter]]: - shared_params: List[nn.Parameter] = [] - expert_params: List[nn.Parameter] = [] - - for p in params: - if is_moe_param(p): - expert_params.append(p) - else: - shared_params.append(p) - return shared_params, expert_params - - -def split_params_grads_into_shared_and_expert_params( - group: List[torch.nn.Parameter]) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: - """Split grad of parameters into grads of non-expert params - and grads of expert params. This is useful while computing - grad-norms for clipping and overflow detection - - group (List[torch.nn.Parameter]): - Args: - The group of parameters to split - - Returns: - Tuple[List[torch.Tensor], List[torch.Tensor]]: - list of gradients for non MoE params, list of gradients of MoE params - """ - expert_grads: List[torch.Tensor] = [] - shared_grads: List[torch.Tensor] = [] - - for p in group: - if p.grad is not None: - if is_moe_param(p): - expert_grads.append(p.grad.to(p.dtype)) - else: - shared_grads.append(p.grad.to(p.dtype)) - return shared_grads, expert_grads - - -def split_params_into_different_moe_groups_for_optimizer( - param_groups: Union[Dict[str, Any], Tuple[Dict[str, Any], ...], List[Dict[str, Any]]], - max_group_size: Union[int, float] = 178956971) -> List[Dict[str, Any]]: - """Split parameters into different MoE groups for optimizer - - Args: - param_groups (Union[Dict[str, Any], Tuple[Dict[str, Any], ...], List[Dict[str, Any]]]) - The list of parameter groups to split - - Returns: - List[Dict[str, Any]]: - list of MoE/non-MoE groups for optimizer - """ - if isinstance(param_groups, tuple): - param_groups = list(param_groups) # Tuple cannot be modified - elif isinstance(param_groups, dict): - param_groups = [param_groups] - elif not isinstance(param_groups, list): - raise ValueError(f"Unknown param group type of {type(param_groups)}") - - # gather all data parallel group names - data_parallel_group_names: Set[str] = set() - for param_group in param_groups: - for param in cast(List[nn.Parameter], param_group["params"]): - if is_moe_param(param): - data_parallel_group_names.add(param.group_name) - - # Create the param MoE groups, leave param assign to next step - group_moe: Dict[str, Dict[str, Dict[str, Any]]] = defaultdict(lambda: defaultdict(dict)) - for param_group in param_groups: - for key in data_parallel_group_names: - group_moe[param_group['name']][key] = { - **param_group, - 'name': key, - 'moe': True, - 'params': [], - } - - # Assign param - for param_group in param_groups: - new_params: List[nn.Parameter] = [] - - for param in cast(List[nn.Parameter], param_group['params']): - if is_moe_param(param): - group_moe[param_group['name']][param.group_name]['params'].append(param) - else: - new_params.append(param) - param_group['params'] = new_params - - # Flatten the moe groups - if max_group_size is not None: - for moe_group in group_moe.values(): - for param_group in moe_group.values(): - cur_group: List[nn.Parameter] = [] - all_groups: List[List[nn.Parameter]] = [] - size_of_cur_group = 0 - - for param in cast(List[nn.Parameter], param_group['params']): - if size_of_cur_group + param.numel() <= max_group_size: - cur_group.append(param) - size_of_cur_group += param.numel() - else: - all_groups.append(cur_group) - cur_group = [param] - size_of_cur_group = param.numel() - - if cur_group: - all_groups.append(cur_group) - - for group in all_groups: - param_groups.append({**param_group, 'params': group}) - else: - for moe_group in group_moe.values(): - for param_group in moe_group.values(): - param_groups.append(param_group) - - return param_groups - - -def is_moe_param_group(param_group): - return param_group.get('moe', False) - - -def configure_moe_param_groups(model_parameters: List): - assert isinstance(model_parameters, list), "model_parameters must be a list" - - for p in model_parameters: - # match torch.optim.Optimizer expectations, - # see: https://github.com/pytorch/pytorch/blob/2ffab6e663b9c6951048b8c8ba82d2cc5ca5c2fc/torch/optim/optimizer.py#L270-L272 - if not isinstance(p, (torch.Tensor, dict)): - raise TypeError("param argument that would be given to the optimizer should be " - f"an iterable of Tensors or dicts, but got {type(p)}") - - # peak at the first element to determine how to proceed - first = model_parameters[0] - - # Case 1: model_parameters is a list of torch.nn.Parameter - # -> need to create moe compatible param groups - if isinstance(first, torch.nn.Parameter): - param_group = {'params': model_parameters, 'name': 'dense-params'} - return split_params_into_different_moe_groups_for_optimizer(param_group) - - # Case 2: model_parameters is a list of param groups List[dict] - # -> moe compatible param groups might already exist, if not create them - elif isinstance(first, dict): - #there are no moe groups created - if not any(['moe' in param_group for param_group in model_parameters]): - return split_params_into_different_moe_groups_for_optimizer(model_parameters) - else: - # moe groups exist, nothing to do - return model_parameters diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/monitor/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index ee220003517d39176acde8ad9c1176048a2242c7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/comet.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/comet.cpython-310.pyc deleted file mode 100644 index 213191f0bdff5bdba17dcf634303bcd64f19c1b7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/comet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/config.cpython-310.pyc deleted file mode 100644 index cb6c34ac797ef10c2e9d3b9e433ae506e4fb6c61..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/csv_monitor.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/csv_monitor.cpython-310.pyc deleted file mode 100644 index 50bd44d010e98247fc67c83615aba4399f3c1739..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/csv_monitor.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/monitor.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/monitor.cpython-310.pyc deleted file mode 100644 index 2aafc2ef90c68744692d64c79ae5d4b663a74fa8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/monitor.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/tensorboard.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/tensorboard.cpython-310.pyc deleted file mode 100644 index 726283ba773ea26b38dde1ed8bc903db0633c6c1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/tensorboard.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index d9cbb9e045371efeed8da59eaf1a1443a239305b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/wandb.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/wandb.cpython-310.pyc deleted file mode 100644 index 05cf6516620132fcd98416180ec997b5fe2fe8dd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/monitor/__pycache__/wandb.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/comet.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/comet.py deleted file mode 100644 index d8bc4017800f3914fd1a1881b0510c3cdee75324..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/monitor/comet.py +++ /dev/null @@ -1,92 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import TYPE_CHECKING, Any, Tuple, List, Dict, Optional - -from .utils import check_comet_availability -from .monitor import Monitor - -import deepspeed.comm as dist - -if TYPE_CHECKING: - import comet_ml - from .config import CometConfig - -Name = str -Value = Any -GlobalSamples = int -Event = Tuple[Name, Value, GlobalSamples] - - -class CometMonitor(Monitor): - - def __init__(self, comet_config: "CometConfig"): - super().__init__(comet_config) - check_comet_availability() - import comet_ml - - self.enabled = comet_config.enabled - self._samples_log_interval = comet_config.samples_log_interval - self._experiment: Optional["comet_ml.ExperimentBase"] = None - - if self.enabled and dist.get_rank() == 0: - self._experiment = comet_ml.start( - api_key=comet_config.api_key, - project=comet_config.project, - workspace=comet_config.workspace, - experiment_key=comet_config.experiment_key, - mode=comet_config.mode, - online=comet_config.online, - ) - - if comet_config.experiment_name is not None: - self._experiment.set_name(comet_config.experiment_name) - - self._events_log_scheduler = EventsLogScheduler(comet_config.samples_log_interval) - - @property - def experiment(self) -> Optional["comet_ml.ExperimentBase"]: - return self._experiment - - @property - def samples_log_interval(self) -> int: - return self._samples_log_interval - - def write_events(self, event_list: List[Event]) -> None: - if not self.enabled or dist.get_rank() != 0: - return None - - for event in event_list: - name = event[0] - value = event[1] - engine_global_samples = event[2] - - if self._events_log_scheduler.needs_logging(name, engine_global_samples): - self._experiment.__internal_api__log_metric__( - name=name, - value=value, - step=engine_global_samples, - ) - - -class EventsLogScheduler: - - def __init__(self, samples_log_interval: int): - self._samples_log_interval = samples_log_interval - self._last_logged_events_samples: Dict[str, int] = {} - - def needs_logging(self, name: str, current_sample: int) -> bool: - if name not in self._last_logged_events_samples: - self._last_logged_events_samples[name] = current_sample - return True - - last_logged_sample = self._last_logged_events_samples[name] - samples_delta = current_sample - last_logged_sample - - if samples_delta >= self._samples_log_interval: - self._last_logged_events_samples[name] = current_sample - return True - - return False diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/config.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/config.py deleted file mode 100644 index 960ce1ba997a231edec53e84855942f3d6ece1d8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/monitor/config.py +++ /dev/null @@ -1,144 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Optional - -from pydantic import model_validator -from deepspeed.runtime.config_utils import DeepSpeedConfigModel - - -def get_monitor_config(param_dict): - monitor_dict = {key: param_dict.get(key, {}) for key in ("tensorboard", "wandb", "csv_monitor", "comet")} - return DeepSpeedMonitorConfig(**monitor_dict) - - -class TensorBoardConfig(DeepSpeedConfigModel): - """Sets parameters for TensorBoard monitor.""" - - enabled: bool = False - """ Whether logging to Tensorboard is enabled. Requires `tensorboard` package is installed. """ - - output_path: str = "" - """ - Path to where the Tensorboard logs will be written. If not provided, the - output path is set under the training script’s launching path. - """ - - job_name: str = "DeepSpeedJobName" - """ Name for the current job. This will become a new directory inside `output_path`. """ - - -class WandbConfig(DeepSpeedConfigModel): - """Sets parameters for WandB monitor.""" - - enabled: bool = False - """ Whether logging to WandB is enabled. Requires `wandb` package is installed. """ - - group: Optional[str] = None - """ Name for the WandB group. This can be used to group together runs. """ - - team: Optional[str] = None - """ Name for the WandB team. """ - - project: str = "deepspeed" - """ Name for the WandB project. """ - - -class CSVConfig(DeepSpeedConfigModel): - """Sets parameters for CSV monitor.""" - - enabled: bool = False - """ Whether logging to local CSV files is enabled. """ - - output_path: str = "" - """ - Path to where the csv files will be written. If not provided, the output - path is set under the training script’s launching path. - """ - - job_name: str = "DeepSpeedJobName" - """ Name for the current job. This will become a new directory inside `output_path`. """ - - -class CometConfig(DeepSpeedConfigModel): - """ - Sets parameters for Comet monitor. For logging data Comet uses - experiment object. - https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/Experiment/ - """ - - enabled: bool = False - """ Whether logging to Comet is enabled. Requires `comet_ml` package is installed. """ - - samples_log_interval: int = 100 - """ Metrics will be submitted to Comet after processing every `samples_log_intervas` samples""" - - project: Optional[str] = None - """ - Comet project name. Can be set through .comet.config file or environment variable COMET_PROJECT_NAME - https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options - """ - - workspace: Optional[str] = None - """ - Comet workspace name. Can be set through .comet.config file or environment variable COMET_WORKSPACE - https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options - """ - - api_key: Optional[str] = None - """ - Comet API key. Can be set through .comet.config file or environment variable COMET_API_KEY - https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options - """ - - experiment_name: Optional[str] = None - """ - The name for comet experiment to be used for logging. - Can be set through .comet.config file or environment variable COMET_EXPERIMENT_NAME - https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options - """ - - experiment_key: Optional[str] = None - """ - The key for comet experiment to be used for logging. Must be an alphanumeric string whose length is between 32 and 50 characters. - Can be set through .comet.config or environment variable COMET_EXPERIMENT_KEY - https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options - """ - - online: Optional[bool] = None - """ - If True, the data will be logged to Comet server, otherwise it will be stored locally in offline experiment - Defaults to True. - """ - - mode: Optional[str] = None - """ - Control how the Comet experiment is started, 3 options are possible.: - - "get": Continue logging to an existing experiment identified by the `experiment_key` value. - - "create": Always creates of a new experiment, useful for HPO sweeps. - - "get_or_create" (default): Starts a fresh experiment if required, or persists logging to an existing one. - """ - - -class DeepSpeedMonitorConfig(DeepSpeedConfigModel): - """Sets parameters for various monitoring methods.""" - - tensorboard: TensorBoardConfig = {} - """ TensorBoard monitor, requires `tensorboard` package is installed. """ - - comet: CometConfig = {} - """ Comet monitor, requires `comet_ml` package is installed """ - - wandb: WandbConfig = {} - """ WandB monitor, requires `wandb` package is installed. """ - - csv_monitor: CSVConfig = {} - """ Local CSV output of monitoring data. """ - - @model_validator(mode="after") - def check_enabled(self): - enabled = self.tensorboard.enabled or self.wandb.enabled or self.csv_monitor.enabled or self.comet.enabled - self.__dict__["enabled"] = enabled - return self diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/csv_monitor.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/csv_monitor.py deleted file mode 100644 index c7a19b14ad8227fc8187bbe161667285b7d0c717..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/monitor/csv_monitor.py +++ /dev/null @@ -1,67 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .monitor import Monitor -import os - -import deepspeed.comm as dist - - -class csvMonitor(Monitor): - - def __init__(self, csv_config): - super().__init__(csv_config) - self.filenames = [] - self.enabled = csv_config.enabled - self.output_path = csv_config.output_path - self.job_name = csv_config.job_name - self.log_dir = self.setup_log_dir() - - def setup_log_dir(self, base=os.path.join(os.path.expanduser("~"), "csv_monitor")): - if self.enabled and dist.get_rank() == 0: - if self.output_path is not None: - log_dir = os.path.join(self.output_path, self.job_name) - # NOTE: This code path currently is never used since the default tensorboard_output_path is an empty string and not None. Saving it in case we want this functionality in the future. - else: - if "DLWS_JOB_ID" in os.environ: - infra_job_id = os.environ["DLWS_JOB_ID"] - elif "DLTS_JOB_ID" in os.environ: - infra_job_id = os.environ["DLTS_JOB_ID"] - else: - infra_job_id = "unknown-job-id" - - csv_monitor_dir_name = os.path.join(infra_job_id, "logs") - log_dir = os.path.join(base, csv_monitor_dir_name, self.job_name) - os.makedirs(log_dir, exist_ok=True) - return log_dir - - def write_events(self, event_list): - if self.enabled and dist.get_rank() == 0: - import csv - # We assume each event_list element is a tensorboard-style tuple in the format: (log_name: String, value, step: Int) - for event in event_list: - log_name = event[0] - value = event[1] - step = event[2] - - # Set the header to the log_name - # Need this check because the deepspeed engine currently formats log strings to separate with '/' - if '/' in log_name: - record_splits = log_name.split('/') - header = record_splits[len(record_splits) - 1] - else: - header = log_name - - # sanitize common naming conventions into filename - filename = log_name.replace('/', '_').replace(' ', '_') - fname = self.log_dir + '/' + filename + '.csv' - - # Open file and record event. Insert header if this is the first time writing - with open(fname, 'a+') as csv_monitor_file: - csv_monitor_writer = csv.writer(csv_monitor_file) - if filename not in self.filenames: - self.filenames.append(filename) - csv_monitor_writer.writerow(['step', header]) - csv_monitor_writer.writerow([step, value]) diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/monitor.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/monitor.py deleted file mode 100644 index e7e26dc483d9707bfc04febda348bff30d25e53b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/monitor/monitor.py +++ /dev/null @@ -1,59 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Support different forms of monitoring such as wandb and tensorboard -""" - -from abc import ABC, abstractmethod -import deepspeed.comm as dist - - -class Monitor(ABC): - - @abstractmethod - def __init__(self, monitor_config): - self.monitor_config = monitor_config - - @abstractmethod - def write_events(self, event_list): - pass - - -from .wandb import WandbMonitor -from .tensorboard import TensorBoardMonitor -from .csv_monitor import csvMonitor -from .comet import CometMonitor - - -class MonitorMaster(Monitor): - - def __init__(self, monitor_config): - super().__init__(monitor_config) - self.tb_monitor = None - self.wandb_monitor = None - self.csv_monitor = None - self.comet_monitor = None - self.enabled = monitor_config.enabled - - if dist.get_rank() == 0: - if monitor_config.tensorboard.enabled: - self.tb_monitor = TensorBoardMonitor(monitor_config.tensorboard) - if monitor_config.wandb.enabled: - self.wandb_monitor = WandbMonitor(monitor_config.wandb) - if monitor_config.csv_monitor.enabled: - self.csv_monitor = csvMonitor(monitor_config.csv_monitor) - if monitor_config.comet.enabled: - self.comet_monitor = CometMonitor(monitor_config.comet) - - def write_events(self, event_list): - if dist.get_rank() == 0: - if self.tb_monitor is not None: - self.tb_monitor.write_events(event_list) - if self.wandb_monitor is not None: - self.wandb_monitor.write_events(event_list) - if self.csv_monitor is not None: - self.csv_monitor.write_events(event_list) - if self.comet_monitor is not None: - self.comet_monitor.write_events(event_list) diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/tensorboard.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/tensorboard.py deleted file mode 100644 index 985c9ed44b6f5be28785699bca4a3638ed6063bd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/monitor/tensorboard.py +++ /dev/null @@ -1,56 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .utils import check_tb_availability -from .monitor import Monitor -import os - -import deepspeed.comm as dist - - -class TensorBoardMonitor(Monitor): - - def __init__(self, tensorboard_config): - super().__init__(tensorboard_config) - check_tb_availability() - - self.summary_writer = None - self.enabled = tensorboard_config.enabled - self.output_path = tensorboard_config.output_path - self.job_name = tensorboard_config.job_name - - if self.enabled and dist.get_rank() == 0: - self.get_summary_writer() - - def get_summary_writer(self, base=os.path.join(os.path.expanduser("~"), "tensorboard")): - if self.enabled and dist.get_rank() == 0: - from torch.utils.tensorboard import SummaryWriter - if self.output_path is not None: - log_dir = os.path.join(self.output_path, self.job_name) - # NOTE: This code path currently is never used since the default output_path is an empty string and not None. Saving it in case we want this functionality in the future. - else: - if "DLWS_JOB_ID" in os.environ: - infra_job_id = os.environ["DLWS_JOB_ID"] - elif "DLTS_JOB_ID" in os.environ: - infra_job_id = os.environ["DLTS_JOB_ID"] - else: - infra_job_id = "unknown-job-id" - - summary_writer_dir_name = os.path.join(infra_job_id, "logs") - log_dir = os.path.join(base, summary_writer_dir_name, self.output_path) - os.makedirs(log_dir, exist_ok=True) - self.summary_writer = SummaryWriter(log_dir=log_dir) - return self.summary_writer - - def write_events(self, event_list, flush=True): - if self.enabled and self.summary_writer is not None and dist.get_rank() == 0: - for event in event_list: - self.summary_writer.add_scalar(*event) - if flush: - self.summary_writer.flush() - - def flush(self): - if self.enabled and self.summary_writer is not None and dist.get_rank() == 0: - self.summary_writer.flush() diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/utils.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/utils.py deleted file mode 100644 index f5530e8532e1a6982ce8acdf20b98aa8fdbb5e81..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/monitor/utils.py +++ /dev/null @@ -1,37 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from packaging import version as pkg_version - - -def check_tb_availability(): - try: - # torch.utils.tensorboard will fail if `tensorboard` is not available, - # see their docs for more details: https://pytorch.org/docs/1.8.0/tensorboard.html - import tensorboard # noqa: F401 # type: ignore - except ImportError: - print('If you want to use tensorboard logging, please `pip install tensorboard`') - raise - - -def check_wandb_availability(): - try: - import wandb # noqa: F401 # type: ignore - except ImportError: - print( - 'If you want to use wandb logging, please `pip install wandb` and follow the instructions at https://docs.wandb.ai/quickstart' - ) - raise - - -def check_comet_availability(): - try: - import comet_ml - comet_version = pkg_version.parse(comet_ml.__version__) - if comet_version < pkg_version.Version("3.41.0"): - raise ImportError("`comet_ml` must have at least version 3.41.0") - except ImportError: - print('If you want to use comet logging, please `pip install "comet_ml>=3.41.0"`') - raise diff --git a/venv/lib/python3.10/site-packages/deepspeed/monitor/wandb.py b/venv/lib/python3.10/site-packages/deepspeed/monitor/wandb.py deleted file mode 100644 index 174a2eb2d3b7e71d37ac5b61c7fdb44352455f07..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/monitor/wandb.py +++ /dev/null @@ -1,38 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .utils import check_wandb_availability -from .monitor import Monitor - -import deepspeed.comm as dist - - -class WandbMonitor(Monitor): - - def __init__(self, wandb_config): - super().__init__(wandb_config) - check_wandb_availability() - import wandb - - self.enabled = wandb_config.enabled - self.group = wandb_config.group - self.team = wandb_config.team - self.project = wandb_config.project - - if self.enabled and dist.get_rank() == 0: - wandb.init(project=self.project, group=self.group, entity=self.team) - - def log(self, data, step=None, commit=None): - if self.enabled and dist.get_rank() == 0: - import wandb - return wandb.log(data, step=step, commit=commit) - - def write_events(self, event_list): - if self.enabled and dist.get_rank() == 0: - for event in event_list: - label = event[0] - value = event[1] - step = event[2] - self.log({label: value}, step=step) diff --git a/venv/lib/python3.10/site-packages/deepspeed/nebula/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/nebula/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nebula/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 8045d0a09285f4a9133159be15e5de540455d934..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/config.cpython-310.pyc deleted file mode 100644 index 30483f724d6cfd67506f7ca131abf51e500a76e5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/constants.cpython-310.pyc deleted file mode 100644 index 61b782db27e5b0249dbde3fe939b8d6d83e2e0c8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nebula/__pycache__/constants.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nebula/config.py b/venv/lib/python3.10/site-packages/deepspeed/nebula/config.py deleted file mode 100644 index dc49185738c92a3173f2fd5c68f1d6ab5a32dd92..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nebula/config.py +++ /dev/null @@ -1,43 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigObject -from deepspeed.nebula.constants import * - - -class DeepSpeedNebulaConfig(DeepSpeedConfigObject): - - def __init__(self, param_dict): - super(DeepSpeedNebulaConfig, self).__init__() - - self.enabled = None - self.persistent_storage_path = None - self.persistent_time_interval = None - self.num_of_version_in_retention = None - self.enable_nebula_load = None - - if NEBULA in param_dict.keys(): - nebula_dict = param_dict[NEBULA] - else: - nebula_dict = {} - - self._initialize(nebula_dict) - - def _initialize(self, nebula_dict): - self.enabled = get_scalar_param(nebula_dict, NEBULA_ENABLED, NEBULA_ENABLED_DEFAULT) - - self.load_path = get_scalar_param(nebula_dict, NEBULA_LOAD_PATH, NEBULA_LOAD_PATH_DEFAULT) - - self.enable_nebula_load = get_scalar_param(nebula_dict, NEBULA_ENABLE_NEBULA_LOAD, - NEBULA_ENABLE_NEBULA_LOAD_DEFAULT) - - self.persistent_storage_path = get_scalar_param(nebula_dict, NEBULA_PERSISTENT_STORAGE_PATH, - NEBULA_PERSISTENT_STORAGE_PATH_DEFAULT) - - self.persistent_time_interval = get_scalar_param(nebula_dict, NEBULA_PERSISTENT_TIME_INTERVAL, - NEBULA_PERSISTENT_TIME_INTERVAL_DEFAULT) - - self.num_of_version_in_retention = get_scalar_param(nebula_dict, NEBULA_NUM_OF_VERSION_IN_RETENTION, - NEBULA_NUM_OF_VERSION_IN_RETENTION_DEFAULT) diff --git a/venv/lib/python3.10/site-packages/deepspeed/nebula/constants.py b/venv/lib/python3.10/site-packages/deepspeed/nebula/constants.py deleted file mode 100644 index 9fa5769b55979e4dcd5c80ead06d2117dcc2ec40..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nebula/constants.py +++ /dev/null @@ -1,73 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -######################################### -# nebula -######################################### -# Nebula. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -NEBULA_FORMAT = ''' -nebula should be enabled as: -"session_params": { - "nebula": { - "enabled": true, - "persistent_storage_path": "/foo/bar", - "persistent_time_interval": 100, - "num_of_version_in_retention": 2, - "enable_nebula_load": true - } -} -''' - -NEBULA = "nebula" - -NEBULA_ENABLED = "enabled" -NEBULA_ENABLED_DEFAULT = False - -# There is a case where customer want to load the checkpoint saved -# by raw torch. Because nebula cannot load torch checkpoint directly -# as they have different folder structures to bring the gap for -# loading(the data are totally same in bytes for torch and nebula -# saving). -# In this case, we must disable nebula load to use raw torch load. -# Customer can just set NEBULA_ENABLE_NEBULA_LOAD to False. Then use -# original way of deepspeed to load, i.e. set the value of "--load". -NEBULA_ENABLE_NEBULA_LOAD = "enable_nebula_load" -NEBULA_ENABLE_NEBULA_LOAD_DEFAULT = True - -# When you want to resume the previous checkpoint saved by nebula, -# you can set NEBULA_LOAD_PATH as the parent folder of checkpoint. -# If NEBULA_LOAD_PATH is None, the NEBULA_PERSISTENT_STORAGE_PATH -# will be the default path to load. -NEBULA_LOAD_PATH = "nebula_load_path" -NEBULA_LOAD_PATH_DEFAULT = None - -# Nebula will save the checkpoint under NEBULA_LOAD_PATH in the -# asynchronous way. -NEBULA_PERSISTENT_STORAGE_PATH = "persistent_storage_path" -NEBULA_PERSISTENT_STORAGE_PATH_DEFAULT = None - -# Time interval to trigger the nebula persistence. -NEBULA_PERSISTENT_TIME_INTERVAL = "persistent_time_interval" -NEBULA_PERSISTENT_TIME_INTERVAL_DEFAULT = 100 - -# Checkpoint number which will be kept in memory. Let us say, -# if the value is 2. Then we have checkpoints 1 and 2 are ready -# now. When it comes to checkpoint 3, the 1 will be removed if -# 1 has been persisted to disk. -NEBULA_NUM_OF_VERSION_IN_RETENTION = "num_of_version_in_retention" -NEBULA_NUM_OF_VERSION_IN_RETENTION_DEFAULT = 2 - -# Nebula envs -NEBULA_EXPORT_ENVS = [ - 'DLTS_JOB_ID', 'DLTS_NUM_WORKER', 'NEBULA_PERSISTENT_STORAGE_PATH', 'NEBULA_PERSISTENT_TIME_INTERVAL', - 'AML_RUN_ID', 'AZUREML_RUN_TOKEN', 'AZUREML_WORKSPACE_SCOPE', 'AZUREML_EXPERIMENT_SCOPE', - 'AZUREML_RUN_HISTORY_SERVICE_ENDPOINT', 'AZUREML_RUN_ID', 'NEBULA_MEMORY_BUFFER_SIZE', - 'AZUREML_PARAMETER_ITPJOB_NAME', 'FC_TASKROLE_NAME', 'FC_TASK_INDEX', 'MASTER_HOST', 'LOCAL_HOST', - 'AZUREML_BLOB_ACCOUNT_NAME', 'AZUREML_BLOB_ACCOUNT_KEY' -] - -# ITP env files -DLTS_POD_ENV_PATH = '/dlts-runtime/env/pod.env' diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/__init__.py deleted file mode 100644 index 6d0de857cbd345169651f0133cee87a992ea907e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .perf_run_sweep import sweep_main, parse_sweep_arguments -from .perf_generate_param import generate_main -from .test_ds_aio import ds_io_main diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index a7f4f55efbf2e87bbf95a74a8ac9a9dcb21471ea..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_args.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_args.cpython-310.pyc deleted file mode 100644 index afef6d87676ebb4051cd7167fba6b055200e7896..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_args.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_basic.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_basic.cpython-310.pyc deleted file mode 100644 index 6d59209874a11edd77a86b0cf342c10a1b39a847..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_basic.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_constants.cpython-310.pyc deleted file mode 100644 index 4b6ddf90739ba5ea094f38596f39f4ad060ec819..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_constants.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_handle.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_handle.cpython-310.pyc deleted file mode 100644 index 3546b4f4c26d2c8e2acff714a46d80d5e283af17..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_handle.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_job.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_job.cpython-310.pyc deleted file mode 100644 index ce739961549d224fd649441a464ee695f686acc6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/ds_aio_job.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/io_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/io_engine.cpython-310.pyc deleted file mode 100644 index d56dc8e0bb915333bc56ed35db97d48bf44b6f4e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/io_engine.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/parse_nvme_stats.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/parse_nvme_stats.cpython-310.pyc deleted file mode 100644 index 147ee6251c076435e6a283295db9a5910afe9b40..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/parse_nvme_stats.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_generate_param.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_generate_param.cpython-310.pyc deleted file mode 100644 index 5045ea233fca73122d632acdcee1ede6c95674d3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_generate_param.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_run_sweep.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_run_sweep.cpython-310.pyc deleted file mode 100644 index 05fbca2a4dd3712a78fe8f88d6ce4196efbdff5f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_run_sweep.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_sweep_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_sweep_utils.cpython-310.pyc deleted file mode 100644 index 0a19b99eb2c2192a25ee994fa81015ff96c744bb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/perf_sweep_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/test_ds_aio.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/test_ds_aio.cpython-310.pyc deleted file mode 100644 index 0ab553e181e434a0a8bc62d4edbdb6b512feacdd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/test_ds_aio.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/test_ds_aio_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/test_ds_aio_utils.cpython-310.pyc deleted file mode 100644 index f69fa66d4a3afe4b0f1ea34a28e994b702d7b421..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/test_ds_aio_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/torch_fastio_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/torch_fastio_engine.cpython-310.pyc deleted file mode 100644 index a1df8cfc46405e00edfe017f50bf7a5b58e660e5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/torch_fastio_engine.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/torch_io.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/torch_io.cpython-310.pyc deleted file mode 100644 index ef0f874e0fd88264967896bbe219b1fb05bc4f24..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/torch_io.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/validate_async_io.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/validate_async_io.cpython-310.pyc deleted file mode 100644 index 36ae3bc0330e52cad10958206f22a4bf28f5eef8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/nvme/__pycache__/validate_async_io.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_args.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_args.py deleted file mode 100644 index 95165d98e3ffeb31684592cd6a5a24018ee08c0a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_args.py +++ /dev/null @@ -1,214 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping optimizer tensors to/from (NVMe) storage devices. -""" - -import argparse -import os -from .test_ds_aio_utils import refine_integer_value -from .ds_aio_constants import AIO_HANDLE, AIO_BASIC, TORCH_FAST_IO, TORCH_IO, VALID_ENGINES -from deepspeed.accelerator import get_accelerator - -MAPPING_DELIMITER = ':' - - -def refine_args(args): - if args.io_size and type(args.io_size) == str: - args.io_size = refine_integer_value(args.io_size) - - if args.block_size and type(args.block_size) == str: - args.block_size = refine_integer_value(args.block_size) - - if args.fast_io_size and type(args.fast_io_size) == str: - args.fast_io_size = refine_integer_value(args.fast_io_size) - - return args - - -def _get_mapping_dict(args): - if args.folder is not None: - d = {i: args.folder for i in range(args.multi_process)} - else: - d = {} - for m in args.folder_to_device_mapping: - fields = m.split(MAPPING_DELIMITER) - d[fields[1]] = fields[0] - - return d - - -def _validate_folder_mapping(args): - no_error = True - error_messages = [] - invalid_mappings = [m for m in args.folder_to_device_mapping if MAPPING_DELIMITER not in m] - if len(invalid_mappings) > 0: - error_messages.append( - f'Missing delimiter ({MAPPING_DELIMITER}) in folder_to_device_mapping {invalid_mappings}') - no_error = False - - folder_list = [m.split(MAPPING_DELIMITER)[0] for m in args.folder_to_device_mapping] - invalid_folders = [d for d in folder_list if not os.path.exists(d)] - if len(invalid_folders) > 0: - error_messages.append(f'Invalid folders in folder_to_device_mapping: {invalid_folders}') - no_error = False - - if args.gpu: - device_list = [int(m.split(MAPPING_DELIMITER)[1]) for m in args.folder_to_device_mapping] - invalid_device_list = [dev_id for dev_id in device_list if not dev_id < get_accelerator().device_count()] - if len(invalid_device_list) > 0: - error_messages.append(f'Invalid device ids in folder_to_device_mapping: {invalid_device_list}') - no_error = False - - return no_error, error_messages - - -def validate_args(args): - no_error = True - error_messages = [] - - if args.folder is not None and len(args.folder_to_device_mapping) > 0: - error_messages.append(f'--folder and --folder_to_device_mapping cannot be specified together.') - no_error = False - elif args.folder is None and len(args.folder_to_device_mapping) == 0: - error_messages.append(f'At least one of --folder or --folder_to_device_mapping must be specified.') - no_error = False - - # Validate --folder - if args.folder is not None and not os.path.exists(args.folder): - no_error = False - error_messages.append(f'Invalid folder in --folder: {args.folder} ') - - # Validate --folder_mapping_to_device - if len(args.folder_to_device_mapping) > 0: - no_mapping_error, mapping_error_messages = _validate_folder_mapping(args) - no_error = no_error and no_mapping_error - error_messages += mapping_error_messages - - # Validate --engine - if args.engine not in VALID_ENGINES: - no_error = False - error_messages.append(f'Invalid engine {args.engine}. Valid options = {VALID_ENGINES}') - - # Validate --engine=torch_io - if args.engine == TORCH_IO: - if args.read: - no_error = False - error_messages.append(f'Read not currently supported for --engine={TORCH_IO}') - - if not no_error: - print(f'Found {len(error_messages)} validation error(s)') - # Validate --gpu, --use_gds - if args.use_gds and not args.gpu: - error_messages.append(f'--gpu must be set to transfer with --use_gds') - no_error = False - - if not no_error: - print(f'Found {len(error_messages)} validation errors') - for i, msg in enumerate(error_messages): - print(f'{i+1}: {msg}') - - return no_error - - -def parse_arguments(): - parser = argparse.ArgumentParser() - - parser.add_argument('--folder', default=None, type=str, help='Folder to use for I/O.') - - parser.add_argument('--folder_to_device_mapping', - default=[], - nargs='+', - help='Specification of mapping of folder to (gpu) device id, (ignored for cpu accesses).' - 'Can be specified multiple times for multi-process runs,' - 'e.g. --folder_to_device_mapping /mnt/nvme0:0 --folder_to_device_mapping /mnt/nvme1:15 --gpu' - 'means access /mnt/nvme0 with gpu 0 and /mnt/nvme1 with gpu 15') - - parser.add_argument('--io_size', type=str, default=None, required=True, help='Number of bytes to read or write.') - - parser.add_argument('--fast_io_size', type=str, default='64M', help='Size of fast_io pinned buffer (bytes).') - - parser.add_argument('--read', action='store_true', help='Perform read I/O (default is write)') - - parser.add_argument('--multi_process', - type=int, - default=1, - help='Number of parallel processes doing I/O (default 1).') - - parser.add_argument('--block_size', - type=str, - default='1M', - help='I/O block size. Can use K, M, or G suffix (default 1M for 1 megabytes).') - - parser.add_argument('--queue_depth', type=int, default=32, help='I/O queue depth (default 32).') - - parser.add_argument('--single_submit', - action='store_true', - help='Submit I/O requests in singles (default is submit queue_depth amount at once.).') - - parser.add_argument( - '--sequential_requests', - action='store_true', - help= - 'Delay I/O request submission until completion of prior requests (default is overlap I/O submission and completion requests.).' - ) - - parser.add_argument('--validate', action='store_true', help='Perform validation of I/O transfer in library.') - - parser.add_argument( - '--engine', - type=str, - default=AIO_HANDLE, - help= - f'Engine to perform I/O. Options are [{AIO_HANDLE}, {AIO_BASIC}, {TORCH_IO}, {TORCH_FAST_IO}]. Default is aio_handle' - ) - - parser.add_argument('--loops', type=int, default=3, help='Count of operation repetitions') - - parser.add_argument('--io_parallel', type=int, default=None, help='Per iop parallelism') - - parser.add_argument('--gpu', action='store_true', help='Use GPU memory') - - parser.add_argument('--use_gds', action='store_true', help='Enable GDS AIO') - - parser.add_argument('--slow_bounce_buffer', - action='store_true', - help='For GPU memory transfers, measure impact of bounce buffer pinning on critical path.') - - parser.add_argument('--torch_legacy_save', action='store_true', help='Use torch legacy save approach') - - parser.add_argument('--use_accelerator_pin_memory', - action='store_true', - help='Obtain pinned (CPU page-locked) tensors from accelerator') - - parser.add_argument('--warmup_loops', type=int, default=1, help='Count of operation warmup repetitions') - - parser.add_argument('--include_warmup_time', action='store_true', help='Include warmup latency in results') - - parser.add_argument('--different_file_each_iteration', - action='store_true', - help='Read/write a different file on each iteration.') - - args = parser.parse_args() - print(f'args = {args}') - return args - - -def get_validated_args(): - args = parse_arguments() - args = refine_args(args) - if not validate_args(args): - quit() - print(f'Successful validation of command line arguments') - args.total_loops = args.warmup_loops + args.loops - peer_tag = 'gpu' if args.gpu else 'process' - args.mapping_dict = _get_mapping_dict(args) - args.mapping_list = [(device_id, folder) for device_id, folder in args.mapping_dict.items()] - assert len(args.mapping_dict) == len(args.mapping_list) - print(f'Configuring {len(args.mapping_list)} {peer_tag} to folder mapping') - for i, (device_id, folder) in enumerate(args.mapping_list): - print(f'[{i}]: {peer_tag} {device_id} <----> {folder}') - - return args diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_basic.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_basic.py deleted file mode 100644 index 7be7bcd1cd67c7741d3a8008af900f3349f798a3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_basic.py +++ /dev/null @@ -1,64 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping optimizer tensors to/from (NVMe) storage devices. -""" - -import os -import time -from deepspeed.ops.aio import AsyncIOBuilder -from .test_ds_aio_utils import task_log, create_filename, create_file, create_page_locked_tensor -from .ds_aio_constants import * - - -class AIOBasic_Engine(object): - - def __init__(self, args, tid, read_op): - self.ctxt = self._create_context(args, tid, read_op) - - def fini(self): - self.ctxt[BUFFER].detach() - self.ctxt[BUFFER] = None - - def read(self, args, tid, loop_id): - start_time = time.time() - AsyncIOBuilder().load().aio_read(self.ctxt[BUFFER], self.ctxt[FILE], args.block_size, args.queue_depth, - args.single_submit, not args.sequential_requests, args.validate) - end_time = time.time() - self.ctxt[ELAPSED_SEC] += end_time - start_time - - def write(self, args, tid, loop_id): - # Avoid overwriting existing files as it could be artificially faster - if os.path.isfile(self.ctxt[FILE]): - os.remove(self.ctxt[FILE]) - - start_time = time.time() - AsyncIOBuilder().load().aio_write(self.ctxt[BUFFER], self.ctxt[FILE], args.block_size, args.queue_depth, - args.single_submit, not args.sequential_requests, args.validate) - end_time = time.time() - self.ctxt[ELAPSED_SEC] += end_time - start_time - - def _create_context(self, args, tid, read_op): - io_string = "Read" if read_op else "Write" - device_id, folder = args.mapping_list[tid] - filename = create_filename(folder, args.read, args.io_size, tid) - if args.read and not (os.path.isfile(filename) and os.path.getsize(filename) == args.io_size): - create_file(filename, args.io_size) - - task_log(tid, f'Allocate tensor of size {args.io_size} bytes') - - buffer = create_page_locked_tensor(args.io_size, True) - - task_log(tid, - f'{io_string} file {filename} of size {args.io_size} bytes from buffer on device {buffer.device}') - - task_log(tid, f'created deepspeed aio basic engine') - - ctxt = {} - ctxt[FILE] = filename - ctxt[NUM_BYTES] = args.io_size - ctxt[BUFFER] = buffer - ctxt[ELAPSED_SEC] = 0 - return ctxt diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_constants.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_constants.py deleted file mode 100644 index 1b07ed8672ef4b4453523a2292141f18ba50d51b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_constants.py +++ /dev/null @@ -1,20 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -AIO_HANDLE = 'aio_handle' -AIO_BASIC = 'aio_basic' -TORCH_IO = 'torch_io' -TORCH_FAST_IO = 'torch_fastio' -VALID_ENGINES = [AIO_HANDLE, AIO_BASIC, TORCH_IO, TORCH_FAST_IO] - -BUFFER = 'buffer' -BOUNCE_BUFFER = 'bounce_buffer' -NUM_BYTES = 'num_bytes' -FILE = 'file' -HANDLE = 'handle' -ELAPSED_SEC = 'elapsed_sec' -FAST_IO_BUFFER = 'fast_io_buffer' -USE_CPU_LOCKED_TENSOR = 'cpu_locked_tensor' -USE_GDS = 'gds' diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_handle.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_handle.py deleted file mode 100644 index efc157078dd95b21691bd6e4821ce905ae69fde1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_handle.py +++ /dev/null @@ -1,126 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping optimizer tensors to/from (NVMe) storage devices. -""" - -import torch -import os -import time -from deepspeed.ops.aio import AsyncIOBuilder -from deepspeed.ops.op_builder import GDSBuilder -from deepspeed.accelerator import get_accelerator -from .test_ds_aio_utils import task_log, create_filename, create_file, create_page_locked_tensor -from .ds_aio_constants import * - - -class AIOHandle_Engine(object): - - def __init__(self, args, tid, read_op): - self.ctxt = self._create_context(args, tid, read_op) - - def fini(self): - for buf in [BUFFER, BOUNCE_BUFFER]: - if self.ctxt[buf] is not None: - if self.ctxt[USE_CPU_LOCKED_TENSOR]: - self.ctxt[HANDLE].free_cpu_locked_tensor(self.ctxt[buf]) - - self.ctxt[buf].detach() - self.ctxt[buf] = None - - def read(self, args, tid, loop_id): - handle = self.ctxt[HANDLE] - - start_time = time.time() - dest_buffer = BOUNCE_BUFFER if self.ctxt[BOUNCE_BUFFER] is not None else BUFFER - ret = handle.pread(self.ctxt[dest_buffer], self.ctxt[FILE][loop_id], args.validate, True) - assert ret != -1 - handle.wait() - if dest_buffer == BOUNCE_BUFFER: - self.ctxt[BUFFER].data.copy_(self.ctxt[BOUNCE_BUFFER].data) - end_time = time.time() - self.ctxt[ELAPSED_SEC].append(end_time - start_time) - - def write(self, args, tid, loop_id): - # Avoid overwriting existing files as it could be artificially faster - # if os.path.isfile(self.ctxt[FILE]): - # os.remove(self.ctxt[FILE]) - - handle = self.ctxt[HANDLE] - start_time = time.time() - if self.ctxt[BOUNCE_BUFFER] is not None: - source_buffer = BOUNCE_BUFFER - self.ctxt[BOUNCE_BUFFER].data.copy_(self.ctxt[BUFFER].data) - else: - source_buffer = BUFFER - ret = handle.pwrite(self.ctxt[source_buffer], self.ctxt[FILE][loop_id], args.validate, True) - assert ret != -1 - handle.wait() - end_time = time.time() - self.ctxt[ELAPSED_SEC].append(end_time - start_time) - - def _create_files(self, args, folder, tid): - if args.different_file_each_iteration: - filenames = [ - create_filename(folder, args.read, args.io_size, f'{tid}_{l}') for l in range(args.total_loops) - ] - else: - filenames = [ - create_filename(folder, args.read, args.io_size, f'{tid}_{0}') for _ in range(args.total_loops) - ] - - if args.read: - for f in filenames: - if not (os.path.isfile(f) and os.path.getsize(f) == args.io_size): - create_file(f, args.io_size) - else: - for f in filenames: - if os.path.isfile(f): - os.remove(f) - - return filenames - - def _create_context(self, args, tid, read_op): - io_string = "Read" if read_op else "Write" - device_id, folder = args.mapping_list[tid] - filenames = self._create_files(args, folder, tid) - - gds = True if args.use_gds else False - io_parallel = args.io_parallel if args.io_parallel else 1 - if gds: - handle = GDSBuilder().load().gds_handle(args.block_size, args.queue_depth, args.single_submit, - not args.sequential_requests, io_parallel) - else: - handle = AsyncIOBuilder().load().aio_handle(args.block_size, args.queue_depth, args.single_submit, - not args.sequential_requests, io_parallel) - task_log(tid, f'Created DeepNVMe handle engine') - - bounce_buffer = None - if args.gpu: - device_name = get_accelerator().device_name(device_id) - buffer = torch.randint(high=128, size=(args.io_size, ), dtype=torch.uint8, device=device_name) - if gds: - handle.pin_device_tensor(buffer) - elif not args.slow_bounce_buffer: - bounce_buffer = create_page_locked_tensor(args.io_size, args.use_accelerator_pin_memory, handle) - else: - buffer = create_page_locked_tensor(args.io_size, args.use_accelerator_pin_memory, handle) - task_log(tid, f'Allocate tensor of size {args.io_size} bytes') - - ctxt = {} - ctxt[FILE] = filenames - ctxt[NUM_BYTES] = args.io_size - ctxt[HANDLE] = handle - ctxt[USE_GDS] = gds - ctxt[BUFFER] = buffer - ctxt[BOUNCE_BUFFER] = bounce_buffer - ctxt[ELAPSED_SEC] = [] - ctxt[USE_CPU_LOCKED_TENSOR] = not args.use_accelerator_pin_memory - - task_log(tid, - f'{io_string} file {filenames} of size {args.io_size} bytes from buffer on device {buffer.device}', - force=True) - - return ctxt diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_job.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_job.py deleted file mode 100644 index 0f9c8b5f1bcc5b475ef76740a88a10345f5d22d1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/ds_aio_job.py +++ /dev/null @@ -1,50 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping tensors to/from (NVMe) storage devices. -""" -import subprocess -import shlex - - -class Job(object): - - def __init__(self, cmd_line, output_file=None, work_dir=None): - self.cmd_line = cmd_line - self.output_file = output_file - self.work_dir = work_dir - self.output_fd = None - - def cmd(self): - return self.cmd_line - - def get_stdout(self): - return self.output_fd - - def get_stderr(self): - return self.output_fd - - def get_cwd(self): - return self.work_dir - - def open_output_file(self): - if self.output_file is not None: - self.output_fd = open(self.output_file, 'w') - - def close_output_file(self): - if self.output_fd is not None: - self.output_fd.close() - self.output_fd = None - - -def run_job(job, verbose=False): - args = shlex.split(' '.join(job.cmd())) - if verbose: - print(f'args = {args}') - job.open_output_file() - proc = subprocess.run(args=args, stdout=job.get_stdout(), stderr=job.get_stderr(), cwd=job.get_cwd()) - job.close_output_file() - assert proc.returncode == 0, \ - f"This command failed: {job.cmd()}" diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/io_engine.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/io_engine.py deleted file mode 100644 index aea66ec5587499dbc9ea899c1ed28607388d9ab1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/io_engine.py +++ /dev/null @@ -1,126 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import time -from multiprocessing import Pool, Barrier - -from .ds_aio_constants import AIO_BASIC, TORCH_FAST_IO, TORCH_IO -from .test_ds_aio_utils import report_results, task_log, task_barrier -from .ds_aio_handle import AIOHandle_Engine -from .ds_aio_basic import AIOBasic_Engine -from .torch_io import TorchIO_Engine -from .torch_fastio_engine import Torch_FastIO_Engine - - -def prepare_operation(args, tid, read_op): - if args.engine == TORCH_IO: - io_engine = TorchIO_Engine(args, tid, read_op) - elif args.engine == AIO_BASIC: - io_engine = AIOBasic_Engine(args, tid, read_op) - elif args.engine == TORCH_FAST_IO: - io_engine = Torch_FastIO_Engine(args, tid, read_op) - else: - io_engine = AIOHandle_Engine(args, tid, read_op) - - return io_engine - - -def prepare_read(pool_params): - args, tid = pool_params - return prepare_operation(args, tid, True) - - -def prepare_write(pool_params): - args, tid = pool_params - return prepare_operation(args, tid, False) - - -def post_operation(pool_params): - _, _, io_engine = pool_params - io_engine.fini() - - -def read_operation(pool_params): - args, tid, loop_id, io_engine = pool_params - return io_engine.read(args, tid, loop_id) - - -def write_operation(pool_params): - args, tid, loop_id, io_engine = pool_params - return io_engine.write(args, tid, loop_id) - - -def get_schedule(args, read_op): - schedule = {} - if read_op: - schedule['pre'] = prepare_read - schedule['post'] = post_operation - schedule['main'] = read_operation - else: - schedule['pre'] = prepare_write - schedule['post'] = post_operation - schedule['main'] = write_operation - - return schedule - - -def io_engine_tasklet(pool_params): - args, tid, read_op = pool_params - num_processes = len(args.mapping_dict) - - # Create schedule - schedule = get_schedule(args, read_op) - task_log(tid, f'schedule = {schedule}') - task_barrier(aio_barrier, num_processes) - - # Run pre task - task_log(tid, f'running pre-task') - io_engine = schedule["pre"]((args, tid)) - task_barrier(aio_barrier, num_processes) - - # Run main tasks in a loop - io_engine.ctxt["main_task_sec"] = [] - for i in range(args.total_loops): - task_log(tid, f'running main task {i}') - start_time = time.time() - schedule["main"]((args, tid, i, io_engine)) - task_barrier(aio_barrier, num_processes) - stop_time = time.time() - io_engine.ctxt["main_task_sec"].append(stop_time - start_time) - - # Run post task - task_log(tid, f'running post-task') - schedule["post"]((args, tid, io_engine)) - task_barrier(aio_barrier, num_processes) - - ctxt = io_engine.ctxt - # return ctxt["main_task_sec"], ctxt["elapsed_sec"], ctxt["num_bytes"] * args.loops - if args.include_warmup_time: - e2e_latency_sec = sum(ctxt["main_task_sec"]) - task_latency_sec = sum(ctxt["elapsed_sec"]) - actual_loops = args.total_loops - else: - e2e_latency_sec = sum(ctxt["main_task_sec"][args.warmup_loops:]) - task_latency_sec = sum(ctxt["elapsed_sec"][args.warmup_loops:]) - actual_loops = args.loops - - l = ctxt["elapsed_sec"] - task_log(tid, f'task_latency_sec = {l}') - return e2e_latency_sec, task_latency_sec, ctxt["num_bytes"] * actual_loops - - -def _init_takslet(b): - global aio_barrier - aio_barrier = b - - -def io_engine_multiprocessing(args, read_op): - num_processes = len(args.mapping_dict) - b = Barrier(num_processes) - pool_params = [(args, p, read_op) for p in range(num_processes)] - with Pool(processes=num_processes, initializer=_init_takslet, initargs=(b, )) as p: - pool_results = p.map(io_engine_tasklet, pool_params) - - report_results(args, read_op, pool_results) diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/parse_nvme_stats.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/parse_nvme_stats.py deleted file mode 100644 index 09c79ada5b3699676a5c6415bf30909e8686f548..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/parse_nvme_stats.py +++ /dev/null @@ -1,148 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping optimizer tensors to/from (NVMe) storage devices. -""" - -import os -import argparse - -READ_SPEED = 'read_speed' -WRITE_SPEED = 'write_speed' - -PERF_METRICS = [READ_SPEED, WRITE_SPEED] - -METRIC_SEARCH = {READ_SPEED: 'E2E Read Speed', WRITE_SPEED: 'E2E Write Speed'} - - -def parse_arguments(): - parser = argparse.ArgumentParser() - - parser.add_argument('--log_dir', type=str, required=True, help='Folder of statistics logs') - - parser.add_argument('--metric', - type=str, - required=True, - help='Performance metric to report: [read_speed|write_speed]') - - args = parser.parse_args() - print(f'args = {args}') - - return args - - -def extract_value(key, file): - INVALID_PREFIXES = ["ds"] - for p in INVALID_PREFIXES: - if key.startswith(p): - return key - try: - if key[0] in ['t', 'd', 'p']: - return int(key[1:]) - if key.startswith("bs"): - if key.endswith('K'): - v = key[2:].split('K') - return int(v[0]) * 1024 - elif key.endswith('M'): - v = key[2:].split('M') - return int(v[0]) * 1024 * 1024 - else: - return int(key[2:]) - except: - print(f"{file}: extract_value fails on {key}") - return None - - return key - - -def get_file_key(file): - f, _ = os.path.splitext(os.path.basename(file)) - fields = f.split('_') - values = [extract_value(k, file) for k in fields] - return tuple(values) - - -def get_thread_count(file): - f, _ = os.path.splitext(os.path.basename(file)) - fields = f.split('_') - for key in fields: - if key[0] == 't': - return int(key[1:]) - return 1 - - -""" -Extract performance metric from log file. -Sample file lines are: -Task Read Latency = 0.031647682189941406 sec -Task Read Speed = 12.342926020792527 GB/sec -E2E Read Latency = 0.031697988510131836 sec -E2E Read Speed = 12.323337169333062 GB/sec - -For the above sample, -metric = "read_speed" corresponds to "E2E Read Speed", and 12.32 will be returned -""" - - -def get_metric(file, metric): - thread_count = get_thread_count(file) - with open(file) as f: - for line in f.readlines(): - if line.startswith(METRIC_SEARCH[metric]): - if metric in [READ_SPEED, WRITE_SPEED]: - fields = line.split() - return float(fields[-2]) - else: - fields = line.split('=') - return float(fields[-1]) - - return None - - -def validate_args(args): - if not args.metric in PERF_METRICS: - print(f'{args.metric} is not a valid performance metrics') - return False - - if not os.path.isdir(args.log_dir): - print(f'{args.log_dir} folder is not existent') - return False - - return True - - -def get_results(log_files, metric): - results = {} - for f in log_files: - file_key = get_file_key(f) - value = get_metric(f, metric) - results[file_key] = value - - return results - - -def get_sorted_results(log_dir, metric): - log_files = [f for f in os.listdir(log_dir) if os.path.isfile(os.path.join(log_dir, f))] - - log_files_path = [os.path.join(log_dir, f) for f in log_files] - results = get_results(log_files_path, metric) - result_keys = list(results.keys()) - sorted_keys = sorted(result_keys) - return sorted_keys, results - - -def main(): - print("Parsing aio statistics") - args = parse_arguments() - - if not validate_args(args): - quit() - - sorted_keys, results = get_sorted_results(args.log_dir, args.metric) - for k in sorted_keys: - print(f'{k} = {results[k]}') - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_generate_param.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_generate_param.py deleted file mode 100644 index 99be75c2d9190637ee43d49d48cf7a911eae6fbf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_generate_param.py +++ /dev/null @@ -1,101 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping optimizer tensors to/from (NVMe) storage devices. -""" -import os -import argparse -import json -from .parse_nvme_stats import READ_SPEED, WRITE_SPEED, get_sorted_results -from .perf_sweep_utils import BENCH_LOG_DIR, READ_LOG_DIR, WRITE_LOG_DIR - - -def parse_arguments(): - parser = argparse.ArgumentParser() - - parser.add_argument('--log_dir', - type=str, - default=BENCH_LOG_DIR, - help=f'Folder of performance sweep logs. Default is {os.path.join(".", BENCH_LOG_DIR)}') - parser.add_argument('--verbose', action='store_true', help='Print debugging information.') - - args = parser.parse_args() - if args.verbose: - print(f'args = {args}') - - return args - - -def validate_args(args): - for d in [READ_LOG_DIR, WRITE_LOG_DIR]: - log_dir = os.path.join(args.log_dir, d) - if not os.path.isdir(log_dir): - print(f'{log_dir} folder is not existent') - return False - - return True - - -def convert_to_param(key): - assert len(key) == 6 - return { - "single_submit": "true" if key[0] == "single" else "false", - "overlap_events": "true" if key[1] == "overlap" else "false", - "num_threads": int(key[5]), - "queue_depth": int(key[3]), - "block_size": int(key[4]) - } - - -def generate_aio_param(read_log_dir, write_log_dir): - _, read_results = get_sorted_results(read_log_dir, READ_SPEED) - _, write_results = get_sorted_results(write_log_dir, WRITE_SPEED) - - read_results_count = len(read_results.items()) - write_results_count = len(write_results.items()) - assert read_results_count == write_results_count, f"Mismatch in number of read & write results: {read_results_count=} != {write_results_count=}" - - combined_perf = {key[1:]: value for key, value in read_results.items()} - for key, value in write_results.items(): - new_key = key[1:] - if new_key in combined_perf: - combined_perf[new_key] += value - else: - combined_perf[new_key] = 0 - - optimal_key = None - optimal_perf = 0.0 - for key, value in combined_perf.items(): - if value > optimal_perf: - optimal_perf = value - optimal_key = key - - aio_param = {"aio": convert_to_param(optimal_key)} - - read_perf_keys = {key[1:]: key for key in read_results.keys()} - write_perf_keys = {key[1:]: key for key in write_results.keys()} - optimal_config_read = read_results.get(read_perf_keys[optimal_key], None) - optimal_config_write = write_results.get(write_perf_keys[optimal_key], None) - - print(f'Best performance (GB/sec): read = {optimal_config_read:5.2f}, write = {optimal_config_write:5.2f}') - print(json.dumps(aio_param, indent=3)) - - -def generate_main(log_dir): - read_log_dir = os.path.join(log_dir, READ_LOG_DIR) - write_log_dir = os.path.join(log_dir, WRITE_LOG_DIR) - generate_aio_param(read_log_dir, write_log_dir) - - -def main(): - args = parse_arguments() - if not validate_args(args): - quit() - print(f'Generate DeepNVMe configuration from {args.log_dir} logs') - generate_main(args.log_dir) - - -if __name__ == "__main__": - generate_main() diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_run_sweep.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_run_sweep.py deleted file mode 100644 index 4560911bcb98d6f4a7af6911ea750738e02377cf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_run_sweep.py +++ /dev/null @@ -1,319 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping optimizer tensors to/from (NVMe) storage devices. -""" -import os -import sys -import argparse -import json -import itertools -import shutil - -from deepspeed.ops.op_builder import AsyncIOBuilder, GDSBuilder -from .ds_aio_job import Job, run_job -from .perf_sweep_utils import READ_OP_DESC, WRITE_OP_DESC, BENCH_LOG_DIR, \ - READ_LOG_DIR, WRITE_LOG_DIR - -OTHER_OPTIONS = '--engine aio_handle' -PERF_SCRIPT = 'ds_io' -DEFAULT_SWEEP_CONFIG = { - "block_size": ["1M", "8M"], - "queue_depth": [32, 128], - "sequential_requests": [True, False], - "single_submit": [False, True], - "io_parallel": [1, 2, 4, 8], -} - - -class SweepConfig(object): - - def __init__(self, args): - self.folder_to_device_mapping = get_ftd_map(args.nvme_dir) - self.search_space = get_sweep_config_dict(args.sweep_config) - self.search_space.update(self.folder_to_device_mapping) - self.read = not args.no_read - self.write = not args.no_write - self.flush_cache = args.flush_page_cache - self.log_dir = args.log_dir - self.verbose = args.verbose - self.other_options = f'{OTHER_OPTIONS} --loops {args.loops} --io_size {args.io_size}' - if args.gpu: - self.other_options += ' --gpu' - if args.gds: - self.other_options += ' --use_gds' - - -def validate_arguments(args): - if not async_io_setup(): - error_msg = """ - Failing because environment is not properly configured for deepspeed async i/o module. - Possible fix: apt install libaio-dev. - """ - print(error_msg) - quit() - - if args.gds and not gds_io_setup(): - error_msg = """ - Failing because environment is not properly configured for deepspeed GDS I/O operator. - """ - print(error_msg) - quit() - - -def parse_sweep_arguments(): - parser = argparse.ArgumentParser() - - parser.add_argument('--nvme_dir', - nargs='+', - required=True, - help='Directory in which to perform I/O tests. A writeable directory on a NVMe device.') - - parser.add_argument('--sweep_config', type=str, default=None, help='Performance sweep configuration json file.') - - parser.add_argument('--no_read', action='store_true', help='Disable read performance measurements.') - - parser.add_argument('--no_write', action='store_true', help='Disable write performance measurements.') - - parser.add_argument('--io_size', - type=str, - default="400M", - help='Number of I/O bytes to read/write for performance measurements.') - - parser.add_argument('--gpu', action='store_true', help='Test tensor transfers between GPU device and NVME device.') - - parser.add_argument('--gds', action='store_true', help='Run the sweep over NVIDIA GPUDirectStorage operator') - - parser.add_argument( - '--flush_page_cache', - action='store_true', - help= - 'Page cache will not be flushed and reported read speeds may be higher than actual ***Requires sudo access***.' - ) - - parser.add_argument( - '--log_dir', - type=str, - default=BENCH_LOG_DIR, - help=f'Output directory for performance log files. Default is {os.path.join(".", BENCH_LOG_DIR)}') - - parser.add_argument('--loops', type=int, default=1, help='Count of operation repetitions') - - parser.add_argument('--verbose', action='store_true', help='Print debugging information.') - - args = parser.parse_args() - if args.verbose: - print(f'args = {args}') - validate_arguments(args) - - return args - - -def dump_cmd_lines(cmd_lines): - print(f'cmd line count = {len(cmd_lines)}') - for i, cmd in enumerate(cmd_lines): - print(f'{i}: {cmd}') - - -def get_ftd_map(nvme_dir_list): - ftd_list = [f'{dir}:{dev}' for dev, dir in enumerate(nvme_dir_list)] - ftd_arg = [' '.join(ftd for ftd in ftd_list)] - return {'folder_to_device_mapping': ftd_arg} - - -def get_sweep_config_dict(sweep_config_json): - if sweep_config_json is None: - return DEFAULT_SWEEP_CONFIG - - with open(sweep_config_json) as fp: - sweep_config = json.load(fp) - return sweep_config - - -def get_sweep_cmd_lines(sweep_config_dict): - - def flatten_options(key, value_list): - flat_list = [] - for v in value_list: - if not type(v) is bool: - flat_list.append(f'--{key} {v}') - elif v: - flat_list.append(f'--{key}') - else: - flat_list.append(' ') - - return flat_list - - flat_list = [flatten_options(key, value) for key, value in sweep_config_dict.items()] - cmd_list = list(itertools.product(*flat_list)) - cmd_list = [list(cmd) for cmd in cmd_list] - #dump_cmd_lines(cmd_list) - return cmd_list - - -def launch_sweep(sweep_jobs, sync_job, flush_cache_job, verbose): - for perf_job in sweep_jobs: - if flush_cache_job is not None: - run_job(sync_job, verbose) - run_job(flush_cache_job, verbose) - - run_job(perf_job, verbose) - - run_job(sync_job, verbose) - - -def create_cmd_tags(cmd_line): - tags = {} - for param_value in cmd_line: - fields = param_value.split() - if len(fields) == 1: - tags[fields[0]] = None - elif len(fields) == 2: - if fields[0] == '--folder_to_device_mapping': - tags[fields[0]] = len(fields[1:]) - else: - tags[fields[0]] = fields[1] - elif len(fields) > 2: - tags[fields[0]] = len(fields[1:]) - return tags - - -def get_log_file(io_op_desc, cmd_line): - QUEUE_DEPTH = "--queue_depth" - BLOCK_SIZE = "--block_size" - SINGLE_SUBMIT = "--single_submit" - SEQUENTIAL_REQUESTS = "--sequential_requests" - FTD_MAP = "--folder_to_device_mapping" - IO_PARALLEL = "--io_parallel" - - tag_map = { - QUEUE_DEPTH: "d", - BLOCK_SIZE: "bs", - SINGLE_SUBMIT: "single", - SEQUENTIAL_REQUESTS: "sequential", - FTD_MAP: "ftd", - IO_PARALLEL: "p" - } - - tag_default = { - QUEUE_DEPTH: 1, - BLOCK_SIZE: "1M", - SINGLE_SUBMIT: "block", - SEQUENTIAL_REQUESTS: "overlap", - FTD_MAP: 1, - IO_PARALLEL: 1 - } - - def get_default_value(tag): - value = tag_default[tag] - if tag in [SINGLE_SUBMIT, SEQUENTIAL_REQUESTS]: - return value - return f'{tag_map[tag]}{value}' - - def get_config_value(tag, value): - tag_key = tag_map[tag] - if value is None: - return tag_key - return f'{tag_key}{value}' - - tag_list = [SINGLE_SUBMIT, SEQUENTIAL_REQUESTS, FTD_MAP, QUEUE_DEPTH, BLOCK_SIZE, IO_PARALLEL] - log_tags = [io_op_desc] - cmd_tags = create_cmd_tags(cmd_line) - for tag in tag_list: - if tag in cmd_tags: - log_tags.append(get_config_value(tag, cmd_tags[tag])) - else: - log_tags.append(get_default_value(tag)) - - log_file = '_'.join(log_tags) - log_file += '.txt' - return log_file - - -def create_perf_jobs(io_op_desc, log_dir, cmd_lines): - py_cmd = [os.path.join(script_path(), PERF_SCRIPT)] - - perf_jobs = [] - for cmd in cmd_lines: - log_file = os.path.join(log_dir, get_log_file(io_op_desc, cmd)) - job = Job(cmd_line=py_cmd + cmd, output_file=log_file) - perf_jobs.append(job) - - return perf_jobs - - -def script_path(): - return os.path.dirname(os.path.realpath(sys.argv[0])) - - -def async_io_setup(): - return AsyncIOBuilder().is_compatible() - - -def gds_io_setup(): - return GDSBuilder().is_compatible() - - -def remove_folder(folder): - assert os.path.isdir(folder), f"Error: cannot remove {folder} - folder not found" - shutil.rmtree(folder) - - -def run_read_sweep(sweep_config, flush_cache_job, sync_job, cmd_lines): - read_cmd_lines = [[f'--read {sweep_config.other_options}'] + cmd for cmd in cmd_lines] - # dump_cmd_lines(cmd_lines) - - log_folder = os.path.join(sweep_config.log_dir, f'{READ_LOG_DIR}') - os.makedirs(log_folder, exist_ok=True) - - perf_jobs = create_perf_jobs(io_op_desc=READ_OP_DESC, log_dir=log_folder, cmd_lines=read_cmd_lines) - - launch_sweep(sweep_jobs=perf_jobs, - sync_job=sync_job, - flush_cache_job=flush_cache_job, - verbose=sweep_config.verbose) - - -def run_write_sweep(sweep_config, flush_cache_job, sync_job, cmd_lines): - write_cmd_lines = [[f'{sweep_config.other_options}'] + cmd for cmd in cmd_lines] - # dump_cmd_lines(write_cmd_lines) - - log_folder = os.path.join(sweep_config.log_dir, f'{WRITE_LOG_DIR}') - os.makedirs(log_folder, exist_ok=True) - - perf_jobs = create_perf_jobs(io_op_desc=WRITE_OP_DESC, log_dir=log_folder, cmd_lines=write_cmd_lines) - - launch_sweep(sweep_jobs=perf_jobs, - sync_job=sync_job, - flush_cache_job=flush_cache_job, - verbose=sweep_config.verbose) - - -def sweep_main(args): - sweep_config = SweepConfig(args) - cmd_lines = get_sweep_cmd_lines(sweep_config.search_space) - - if sweep_config.flush_cache: - flush_cache_job = Job(cmd_line=['sudo', 'bash -c', "'echo 1 > /proc/sys/vm/drop_caches'"]) - else: - flush_cache_job = None - - sync_job = Job(cmd_line=['sync']) - - if sweep_config.read: - run_read_sweep(sweep_config, flush_cache_job, sync_job, cmd_lines) - - if sweep_config.write: - run_write_sweep(sweep_config, flush_cache_job, sync_job, cmd_lines) - - -def main(): - args = parse_sweep_arguments() - print(f"Running DeepNVMe performance sweep on {args.nvme_dir}") - sweep_main(args) - - -if __name__ == "__main__": - sweep_main() diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_sweep_utils.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_sweep_utils.py deleted file mode 100644 index e6832c1baa492d8dfc695ac251a606e49db82f35..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/perf_sweep_utils.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -SCRIPT_PREFIX = '_aio_bench' -WRITE_OP_DESC = 'write' -READ_OP_DESC = 'read' -READ_IO_DIR = f'{SCRIPT_PREFIX}_{READ_OP_DESC}_io' -WRITE_IO_DIR = f'{SCRIPT_PREFIX}_{WRITE_OP_DESC}_io' -BENCH_LOG_DIR = f'{SCRIPT_PREFIX}_logs' -READ_LOG_DIR = f'{SCRIPT_PREFIX}_{READ_OP_DESC}_logs' -WRITE_LOG_DIR = f'{SCRIPT_PREFIX}_{WRITE_OP_DESC}_logs' diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/test_ds_aio.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/test_ds_aio.py deleted file mode 100644 index a71ad3ee97d373e6feda2f43299ab31c51aec56c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/test_ds_aio.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping optimizer tensors to/from (NVMe) storage devices. -""" - -import multiprocessing as mp -from .ds_aio_args import get_validated_args -from .io_engine import io_engine_multiprocessing - - -def ds_io_main(): - print(f'Testing DeepNVMe python frontend') - - args = get_validated_args() - mp.set_start_method('spawn', force=True) - multiprocess_function = io_engine_multiprocessing - multiprocess_function(args, args.read) - - -if __name__ == "__main__": - ds_io_main() diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/test_ds_aio_utils.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/test_ds_aio_utils.py deleted file mode 100644 index 90b994b0b5322893420037b442aa4763ea5c4af1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/test_ds_aio_utils.py +++ /dev/null @@ -1,91 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping optimizer tensors to/from (NVMe) storage devices. -""" - -import os -from .ds_aio_job import Job, run_job -import torch -from deepspeed.accelerator import get_accelerator - -BYTES_PER_GB = 1024**3 -BYTES_PER_MB = 1024**2 -BYTES_PER_KB = 1024 -LOG_TIDS = [0] - - -def task_log(tid, msg, force=False): - if force or tid in LOG_TIDS: - print(f'tid {tid}: {msg}') - - -def task_barrier(barrier, num_parties): - assert barrier.parties == num_parties - barrier.wait() - assert barrier.broken == False - - -def report_results(args, read_op, pool_results): - #print(f'pool_results = {pool_results}') - io_string = 'Read' if read_op else 'Write' - if None in pool_results: - print(f'Failure in one of {args.threads} {io_string} processes') - return - - total_bytes = sum([num_bytes for _, _, num_bytes in pool_results]) - - task_latency_sec = max([sec for _, sec, _ in pool_results]) - task_speed_GB = 0 if task_latency_sec == 0 else total_bytes / task_latency_sec / BYTES_PER_GB - print(f'Task {io_string} Latency = {task_latency_sec} sec') - print(f'Task {io_string} Speed = {task_speed_GB} GB/sec') - - e2e_latency_sec = max([sec for sec, _, _ in pool_results]) - e2e_speed_GB = 0 if e2e_latency_sec == 0 else total_bytes / e2e_latency_sec / BYTES_PER_GB - print(f'E2E {io_string} Latency = {e2e_latency_sec} sec') - print(f'E2E {io_string} Speed = {e2e_speed_GB} GB/sec') - - -def get_block_size_and_count(io_bytes): - if io_bytes > BYTES_PER_MB and io_bytes % BYTES_PER_MB == 0: - block_size = BYTES_PER_MB - block_size_string = '1M' - else: - assert io_bytes % BYTES_PER_KB == 0 - block_size = BYTES_PER_KB - block_size_string = '1K' - block_count = io_bytes / block_size - - return block_size_string, int(block_count) - - -def refine_integer_value(value): - unit_dict = {'K': 1024, 'M': 1024**2, 'G': 1024**3} - - if value[-1] in list(unit_dict.keys()): - int_value = int(value[:-1]) * unit_dict[value[-1]] - return int_value - return int(value) - - -def create_filename(folder, read_op, size, tid): - io_string = "read" if read_op else "write" - return os.path.join(folder, f'_aio_{io_string}_{size}.pt.{tid}') - - -def create_file(filename, num_bytes): - block_size, block_count = get_block_size_and_count(num_bytes) - dd_job = Job(cmd_line=[f'dd if=/dev/urandom of={filename} bs={block_size} count={block_count}']) - print(f'[Start] Create {filename} of {num_bytes} bytes by running {dd_job.cmd()} ....') - run_job(dd_job) - print(f'[Done] Create read file of {num_bytes} bytes by running {dd_job.cmd()} ....') - - -def create_page_locked_tensor(num_elem, use_accelerator, aio_handle=None): - if use_accelerator: - return get_accelerator().pin_memory(torch.randint(high=128, size=(num_elem, ), dtype=torch.uint8, - device='cpu')) - else: - return aio_handle.new_cpu_locked_tensor(num_elem, torch.empty(0, dtype=torch.uint8)) diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/torch_fastio_engine.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/torch_fastio_engine.py deleted file mode 100644 index fd4918d84c396bf7d14827c42a6b338bed5e9f63..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/torch_fastio_engine.py +++ /dev/null @@ -1,87 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import os -import time -from deepspeed.ops.aio import AsyncIOBuilder -from .test_ds_aio_utils import task_log, create_filename, create_file, create_page_locked_tensor -from .ds_aio_constants import * -from deepspeed.io import FastFileWriter - - -class Torch_FastIO_Engine(object): - - def __init__(self, args, tid, read_op): - assert read_op is False, f'Read operation is not currently supported' - self.ctxt = self._create_context(args, tid, read_op) - self.zipfile_serialization = not args.torch_legacy_save - - def fini(self): - if self.ctxt[USE_CPU_LOCKED_TENSOR]: - for buf in [BUFFER, FAST_IO_BUFFER]: - self.ctxt[HANDLE].free_cpu_locked_tensor(self.ctxt[buf]) - - self.ctxt[BUFFER].detach() - self.ctxt[BUFFER] = None - - def read(self, args, tid): - start_time = time.time() - torch.load(f=self.ctxt[FILE], map_location=self.ctxt[BUFFER].device) - end_time = time.time() - self.ctxt[ELAPSED_SEC] += end_time - start_time - - def write(self, args, tid): - # Avoid overwriting existing files as it could be artificially faster - if os.path.isfile(self.ctxt[FILE]): - os.remove(self.ctxt[FILE]) - - ds_file_writer = FastFileWriter(file_path=self.ctxt[FILE], - aio_handle=self.ctxt[HANDLE], - pinned_tensor=self.ctxt[FAST_IO_BUFFER]) - - start_time = time.time() - torch.save(obj=self.ctxt[BUFFER], f=ds_file_writer, _use_new_zipfile_serialization=self.zipfile_serialization) - ds_file_writer.close() # Force flush to storage - end_time = time.time() - self.ctxt[ELAPSED_SEC] += end_time - start_time - ds_file_writer._dump_state() - - def _create_context(self, args, tid, read_op): - io_string = "Read" if read_op else "Write" - device_id, folder = args.mapping_list[tid] - filename = create_filename(folder, args.read, args.io_size, tid) - if args.read and not (os.path.isfile(filename) and os.path.getsize(filename) == args.io_size): - create_file(filename, args.io_size) - - io_parallel = args.io_parallel if args.io_parallel else 1 - aio_handle = AsyncIOBuilder().load().aio_handle(args.block_size, args.queue_depth, args.single_submit, - not args.sequential_requests, io_parallel) - - if args.gpu: - buffer = torch.randint(high=128, size=(args.io_size, ), dtype=torch.uint8, device=f'cuda:{device_id}') - else: - buffer = create_page_locked_tensor(args.io_size, args.use_accelerator_pin_memory, aio_handle) - - task_log(tid, f'Allocate tensor of size {args.io_size} bytes') - - fast_io_buffer = create_page_locked_tensor(args.fast_io_size, args.use_accelerator_pin_memory, aio_handle) - - task_log(tid, f'created torch_fastio engine') - - ctxt = {} - ctxt[FILE] = filename - ctxt[NUM_BYTES] = args.io_size - ctxt[BUFFER] = buffer - ctxt[HANDLE] = aio_handle - ctxt[FAST_IO_BUFFER] = fast_io_buffer - ctxt[ELAPSED_SEC] = 0 - ctxt[USE_CPU_LOCKED_TENSOR] = not args.use_accelerator_pin_memory - - task_log(tid, - f'{io_string} file {filename} of size {args.io_size} bytes from buffer on device {buffer.device}', - force=True) - - return ctxt diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/torch_io.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/torch_io.py deleted file mode 100644 index 3371771c02f031d1612e55bfce6acc44003131c8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/torch_io.py +++ /dev/null @@ -1,64 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import os -import time -from .test_ds_aio_utils import task_log, create_filename, create_file, create_page_locked_tensor -from .ds_aio_constants import * - - -class TorchIO_Engine(object): - - def __init__(self, args, tid, read_op): - self.ctxt = self._create_context(args, tid, read_op) - self.zipfile_serialization = not args.torch_legacy_save - - def fini(self): - self.ctxt[BUFFER].detach() - self.ctxt[BUFFER] = None - - def read(self, args, tid): - start_time = time.time() - torch.load(f=self.ctxt[FILE], map_location=self.ctxt[BUFFER].device) - end_time = time.time() - self.ctxt[ELAPSED_SEC] += end_time - start_time - - def write(self, args, tid): - # Avoid overwriting existing files as it could be artificially faster - if os.path.isfile(self.ctxt[FILE]): - os.remove(self.ctxt[FILE]) - - start_time = time.time() - torch.save(obj=self.ctxt[BUFFER], f=self.ctxt[FILE], _use_new_zipfile_serialization=self.zipfile_serialization) - end_time = time.time() - self.ctxt[ELAPSED_SEC] += end_time - start_time - - def _create_context(self, args, tid, read_op): - io_string = "Read" if read_op else "Write" - device_id, folder = args.mapping_list[tid] - filename = create_filename(folder, args.read, args.io_size, tid) - if args.read and not (os.path.isfile(filename) and os.path.getsize(filename) == args.io_size): - create_file(filename, args.io_size) - - task_log(tid, f'Allocate tensor of size {args.io_size} bytes') - - if args.gpu: - buffer = torch.randint(high=128, size=(args.io_size, ), dtype=torch.uint8, device=f'cuda:{device_id}') - else: - buffer = create_page_locked_tensor(args.io_size, True) - - task_log(tid, - f'{io_string} file {filename} of size {args.io_size} bytes from buffer on device {buffer.device}', - force=True) - - task_log(tid, f'created torch_io engine') - - ctxt = {} - ctxt[FILE] = filename - ctxt[NUM_BYTES] = args.io_size - ctxt[BUFFER] = buffer - ctxt[ELAPSED_SEC] = 0 - return ctxt diff --git a/venv/lib/python3.10/site-packages/deepspeed/nvme/validate_async_io.py b/venv/lib/python3.10/site-packages/deepspeed/nvme/validate_async_io.py deleted file mode 100644 index 10fb638347bcd9132abc365ba8f488551ea7cb01..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/nvme/validate_async_io.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping optimizer tensors to/from (NVMe) storage devices. -""" -from deepspeed.ops.op_builder import AsyncIOBuilder -assert AsyncIOBuilder().is_compatible() -assert AsyncIOBuilder().load() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/__init__.py deleted file mode 100644 index 15179984173c1f7f8bcdaaeb98037aa8a05ce1e9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from . import adam -from . import adagrad -from . import lamb -from . import lion -from . import sparse_attention -from . import transformer -from . import fp_quantizer -from .transformer import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig - -from ..git_version_info import compatible_ops as __compatible_ops__ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 493ed737a4d90c705ae7055709e91f69a6bc24a8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__init__.py deleted file mode 100644 index 5bf2e4721cd61ec3f548d652698f62b7a50772d5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .cpu_adagrad import DeepSpeedCPUAdagrad diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index ba249042d02572d83565a2c0d2a7f8a15847f0ab..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__pycache__/cpu_adagrad.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__pycache__/cpu_adagrad.cpython-310.pyc deleted file mode 100644 index 70a2d4aaffa1d314ef91f6f849c2cdee77c1ef8d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/__pycache__/cpu_adagrad.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/cpu_adagrad.py b/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/cpu_adagrad.py deleted file mode 100644 index dbde6d95f6520eca20b91d3b9328b741b99df074..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/adagrad/cpu_adagrad.py +++ /dev/null @@ -1,98 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed.ops.op_builder import CPUAdagradBuilder -from deepspeed.utils.logging import should_log_le - - -class DeepSpeedCPUAdagrad(torch.optim.Optimizer): - optimizer_id = 0 - - def __init__(self, model_params, lr=1e-2, eps=1e-10, weight_decay=0, amsgrad=False, fp32_optimizer_states=True): - - default_args = dict(lr=lr, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad) - super(DeepSpeedCPUAdagrad, self).__init__(model_params, default_args) - - self.opt_id = DeepSpeedCPUAdagrad.optimizer_id - DeepSpeedCPUAdagrad.optimizer_id = DeepSpeedCPUAdagrad.optimizer_id + 1 - self.fp32_optimizer_states = fp32_optimizer_states - self.ds_opt_adagrad = CPUAdagradBuilder().load() - - self.ds_opt_adagrad.create_adagrad(self.opt_id, lr, eps, weight_decay, should_log_le("info")) - - def __del__(self): - # need to destroy the C++ object explicitly to avoid a memory leak when deepspeed.initialize - # is used multiple times in the same process (notebook or pytest worker) - self.ds_opt_adagrad.destroy_adagrad(self.opt_id) - - def __setstate__(self, state): - super(DeepSpeedCPUAdagrad, self).__setstate__(state) - for group in self.param_groups: - group.setdefault('amsgrad', False) - - @torch.no_grad() - def step(self, closure=None): - """Update the model parameters. - - .. note:: - This method will be called internally by ZeRO-Offload. DeepSpeed - users should still use ``engine.step()`` as shown in the - `Getting Started - `_ guide. - - Args: - closure (callable, optional): closure to compute the loss. - Defaults to ``None``. - - Returns: - loss: if ``closure`` is provided. Otherwise ``None``. - """ - - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # intended device for step - device = torch.device('cpu') - - for group_id, group in enumerate(self.param_groups): - for param_id, p in enumerate(group['params']): - - if p.grad is None: - continue - - assert p.device == device, f"CPUAdagrad param is on {p.device} and must be 'cpu', make " \ - "sure you enabled 'offload_optimizer': 'cpu' in your ZeRO config." - - state = self.state[p] - # State initialization - if len(state) == 0: - #print(f'group {group_id} param {param_id} = {p.numel()}') - state['step'] = 0 - - #use full precision by default unless self.fp32_optimizer_states is off - state_dtype = torch.float if self.fp32_optimizer_states else p.dtype - - #memory_format=torch.preserve_format) - # gradient variances - state['exp_avg_sq'] = torch.zeros_like(p.data, dtype=state_dtype, device='cpu') - #memory_format=torch.preserve_format) - - state['step'] += 1 - - if p.grad.is_sparse == True: - sparse_param = p.sparse_mask(p.grad) - sparse_exp_avg_sq = state['exp_avg_sq'].sparse_mask(p.grad) - self.ds_opt_adagrad.adagrad_update(self.opt_id, state['step'], group['lr'], group['eps'], - group['weight_decay'], sparse_param.values(), p.grad.values(), - sparse_exp_avg_sq.values()) - p[sparse_param.indices()] = sparse_param.values() - state['exp_avg_sq'][sparse_exp_avg_sq.indices()] = sparse_exp_avg_sq.values() - else: - self.ds_opt_adagrad.adagrad_update(self.opt_id, state['step'], group['lr'], group['eps'], - group['weight_decay'], p.data, p.grad.data, state['exp_avg_sq']) - return loss diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__init__.py deleted file mode 100644 index a29bb9447d015d0992c18624dba41a15aa838866..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .cpu_adam import DeepSpeedCPUAdam -from .fused_adam import FusedAdam diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 88219ad6802c24c4a0dac7984770a9a7a37afb7e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/cpu_adam.cpython-310.pyc deleted file mode 100644 index 853b9f5d7821004d5872f84305e92f0d1137b174..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/cpu_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/fused_adam.cpython-310.pyc deleted file mode 100644 index 7072b45d80e0aad20ad6dfd3febe3840b1674898..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/fused_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/multi_tensor_apply.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/multi_tensor_apply.cpython-310.pyc deleted file mode 100644 index ebc7312343790b7d86fc210c625334972a213bf8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/__pycache__/multi_tensor_apply.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/cpu_adam.py deleted file mode 100644 index e0a72a4942577b4db2ac822ee0675a084df1b2bc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/cpu_adam.py +++ /dev/null @@ -1,166 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from cpuinfo import get_cpu_info -from deepspeed.utils import logger -from deepspeed.utils.logging import should_log_le -from deepspeed.ops.op_builder import CPUAdamBuilder - - -class DeepSpeedCPUAdam(torch.optim.Optimizer): - optimizer_id = 0 - - def __init__(self, - model_params, - lr=1e-3, - bias_correction=True, - betas=(0.9, 0.999), - eps=1e-8, - weight_decay=0, - amsgrad=False, - adamw_mode=True, - fp32_optimizer_states=True): - """Fast vectorized implementation of two variations of Adam optimizer on CPU: - - * Adam: A Method for Stochastic Optimization: (https://arxiv.org/abs/1412.6980); - * AdamW: Fixing Weight Decay Regularization in Adam (https://arxiv.org/abs/1711.05101) - - DeepSpeed CPU Adam(W) provides between 5x to 7x speedup over torch.optim.adam(W). - In order to apply this optimizer, the model requires to have its master parameter (in FP32) - reside on the CPU memory. - - To train on a heterogeneous system, such as coordinating CPU and GPU, DeepSpeed offers - the ZeRO-Offload technology which efficiently offloads the optimizer states into CPU memory, - with minimal impact on training throughput. DeepSpeedCPUAdam plays an important role to minimize - the overhead of the optimizer's latency on CPU. Please refer to ZeRO-Offload tutorial - (https://www.deepspeed.ai/tutorials/zero-offload/) for more information on how to enable this technology. - - For calling step function, there are two options available: (1) update optimizer's states and (2) update - optimizer's states and copy the parameters back to GPU at the same time. We have seen that the second - option can bring 30% higher throughput than the doing the copy separately using option one. - - - .. note:: - We recommend using our `config - `_ - to allow :meth:`deepspeed.initialize` to build this optimizer - for you. - - - Arguments: - model_params (iterable): iterable of parameters to optimize or dicts defining - parameter groups. - lr (float, optional): learning rate. (default: 1e-3) - betas (Tuple[float, float], optional): coefficients used for computing - running averages of gradient and its square. (default: (0.9, 0.999)) - eps (float, optional): term added to the denominator to improve - numerical stability. (default: 1e-8) - weight_decay (float, optional): weight decay (L2 penalty) (default: 0) - amsgrad (boolean, optional): whether to use the AMSGrad variant of this - algorithm from the paper `On the Convergence of Adam and Beyond`_ - (default: False) NOT SUPPORTED in DeepSpeed CPUAdam! - adamw_mode: select between Adam and AdamW implementations (default: AdamW) - fp32_optimizer_states: creates momentum and variance in full precision regardless of - the precision of the parameters (default: True) - """ - - default_args = dict(lr=lr, - betas=betas, - eps=eps, - weight_decay=weight_decay, - bias_correction=bias_correction, - amsgrad=amsgrad) - super(DeepSpeedCPUAdam, self).__init__(model_params, default_args) - - cpu_info = get_cpu_info() - self.cpu_vendor = cpu_info["vendor_id_raw"].lower() if "vendor_id_raw" in cpu_info else "unknown" - if "amd" in self.cpu_vendor: - for group_id, group in enumerate(self.param_groups): - for param_id, p in enumerate(group['params']): - if p.dtype == torch.half: - logger.warning("FP16 params for CPUAdam may not work on AMD CPUs") - break - else: - continue - break - - self.opt_id = DeepSpeedCPUAdam.optimizer_id - DeepSpeedCPUAdam.optimizer_id = DeepSpeedCPUAdam.optimizer_id + 1 - self.adam_w_mode = adamw_mode - self.fp32_optimizer_states = fp32_optimizer_states - self.ds_opt_adam = CPUAdamBuilder().load() - - self.ds_opt_adam.create_adam(self.opt_id, lr, betas[0], betas[1], eps, weight_decay, adamw_mode, - should_log_le("info")) - - def __del__(self): - # need to destroy the C++ object explicitly to avoid a memory leak when deepspeed.initialize - # is used multiple times in the same process (notebook or pytest worker) - self.ds_opt_adam.destroy_adam(self.opt_id) - - def __setstate__(self, state): - super(DeepSpeedCPUAdam, self).__setstate__(state) - for group in self.param_groups: - group.setdefault('amsgrad', False) - - @torch.no_grad() - def step(self, closure=None): - """Update the model parameters. - - .. note:: - This method will be called internally by ZeRO-Offload. DeepSpeed - users should still use ``engine.step()`` as shown in the - `Getting Started - `_ guide. - - Args: - closure (callable, optional): closure to compute the loss. - Defaults to ``None``. - - Returns: - loss: if ``closure`` is provided. Otherwise ``None``. - """ - - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # intended device for step - device = torch.device('cpu') - - for group_id, group in enumerate(self.param_groups): - for param_id, p in enumerate(group['params']): - - if p.grad is None: - continue - - assert p.device == device, f"CPUAdam param is on {p.device} and must be 'cpu', make " \ - "sure you enabled 'offload_optimizer': 'cpu' in your ZeRO config." - - state = self.state[p] - # State initialization - if len(state) == 0: - #print(f'group {group_id} param {param_id} = {p.numel()}') - state['step'] = 0 - - #use full precision by default unless self.fp32_optimizer_states is off - state_dtype = torch.float if self.fp32_optimizer_states else p.dtype - - # gradient momentums - state['exp_avg'] = torch.zeros_like(p.data, dtype=state_dtype, device=device) - #memory_format=torch.preserve_format) - # gradient variances - state['exp_avg_sq'] = torch.zeros_like(p.data, dtype=state_dtype, device=device) - #memory_format=torch.preserve_format) - - state['step'] += 1 - beta1, beta2 = group['betas'] - - self.ds_opt_adam.adam_update(self.opt_id, state['step'], group['lr'], beta1, beta2, group['eps'], - group['weight_decay'], group['bias_correction'], p.data, p.grad.data, - state['exp_avg'], state['exp_avg_sq']) - return loss diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/fused_adam.py deleted file mode 100644 index 53f859e9cc87bde8f16760a4c23394ea1207af0c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/fused_adam.py +++ /dev/null @@ -1,195 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Copyright NVIDIA/apex -This file is adapted from fused adam in NVIDIA/apex, commit 6bd01c4 -""" - -import torch -from .multi_tensor_apply import MultiTensorApply - -multi_tensor_applier = MultiTensorApply(2048 * 32) -from deepspeed.accelerator import get_accelerator -from deepspeed.ops.op_builder import FusedAdamBuilder - - -class FusedAdam(torch.optim.Optimizer): - """Implements Adam algorithm. - - Currently GPU-only. Requires Apex to be installed via - ``pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./``. - - This version of fused Adam implements 2 fusions. - - * Fusion of the Adam update's elementwise operations - * A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches. - - :class:`apex.optimizers.FusedAdam` may be used as a drop-in replacement for ``torch.optim.AdamW``, - or ``torch.optim.Adam`` with ``adam_w_mode=False``:: - - opt = apex.optimizers.FusedAdam(model.parameters(), lr = ....) - ... - opt.step() - - :class:`apex.optimizers.FusedAdam` may be used with or without Amp. If you wish to use :class:`FusedAdam` with Amp, - you may choose any ``opt_level``:: - - opt = apex.optimizers.FusedAdam(model.parameters(), lr = ....) - model, opt = amp.initialize(model, opt, opt_level="O0" or "O1 or "O2") - ... - opt.step() - - In general, ``opt_level="O1"`` is recommended. - - - .. warning:: - A previous version of :class:`FusedAdam` allowed a number of additional arguments to ``step``. These additional arguments - are now deprecated and unnecessary. - - Adam was been proposed in `Adam: A Method for Stochastic Optimization`_. - - Arguments: - params (iterable): iterable of parameters to optimize or dicts defining - parameter groups. - lr (float, optional): learning rate. (default: 1e-3) - betas (Tuple[float, float], optional): coefficients used for computing - running averages of gradient and its square. (default: (0.9, 0.999)) - eps (float, optional): term added to the denominator to improve - numerical stability. (default: 1e-8) - weight_decay (float, optional): weight decay (L2 penalty) (default: 0) - amsgrad (boolean, optional): whether to use the AMSGrad variant of this - algorithm from the paper `On the Convergence of Adam and Beyond`_ - (default: False) NOT SUPPORTED in FusedAdam! - adam_w_mode (boolean, optional): Apply L2 regularization or weight decay - True for decoupled weight decay(also known as AdamW) (default: True) - set_grad_none (bool, optional): whether set grad to None when zero_grad() - method is called. (default: True) - - .. _Adam - A Method for Stochastic Optimization: - https://arxiv.org/abs/1412.6980 - .. _On the Convergence of Adam and Beyond: - https://openreview.net/forum?id=ryQu7f-RZ - """ - - def __init__(self, - params, - lr=1e-3, - bias_correction=True, - betas=(0.9, 0.999), - eps=1e-8, - adam_w_mode=True, - weight_decay=0., - amsgrad=False, - set_grad_none=True): - - if amsgrad: - raise RuntimeError('FusedAdam does not support the AMSGrad variant.') - defaults = dict(lr=lr, bias_correction=bias_correction, betas=betas, eps=eps, weight_decay=weight_decay) - super(FusedAdam, self).__init__(params, defaults) - self.adam_w_mode = 1 if adam_w_mode else 0 - self.set_grad_none = set_grad_none - - fused_adam_cuda = FusedAdamBuilder().load() - # Skip buffer - self._dummy_overflow_buf = get_accelerator().IntTensor([0]) - self.multi_tensor_adam = fused_adam_cuda.multi_tensor_adam - - def zero_grad(self): - if self.set_grad_none: - for group in self.param_groups: - for p in group['params']: - p.grad = None - else: - super(FusedAdam, self).zero_grad() - - def step(self, closure=None, grads=None, output_params=None, scale=None, grad_norms=None, grad_scaler=None): - """Performs a single optimization step. - - Arguments: - closure (callable, optional): A closure that reevaluates the model - and returns the loss. - - The remaining arguments are deprecated, and are only retained (for the moment) for error-checking purposes. - """ - if any(p is not None for p in [grads, output_params, scale, grad_norms]): - raise RuntimeError( - 'FusedAdam has been updated. Simply initialize it identically to torch.optim.Adam, and call step() with no arguments.' - ) - loss = None - if closure is not None: - loss = closure() - - for group in self.param_groups: - if len(group['params']) == 0: - continue - bias_correction = 1 if group['bias_correction'] else 0 - beta1, beta2 = group['betas'] - - # assume same step across group now to simplify things - # per parameter step can be easily support by making it tensor, or pass list into kernel - if 'step' not in group: - group['step'] = 0 - - # create lists for multi-tensor apply - g_16, p_16, m_16, v_16 = [], [], [], [] - g_bf, p_bf, m_bf, v_bf = [], [], [], [] - g_32, p_32, m_32, v_32 = [], [], [], [] - - for p in group['params']: - if p.grad is None: - continue - if p.grad.data.is_sparse: - raise RuntimeError( - 'FusedAdam does not support sparse gradients, please consider SparseAdam instead') - - state = self.state[p] - # State initialization - if len(state) == 0: - # DeepSpeed ZeRO 3 processes each subgroup a time, so we need to keep tracking step count for each tensor separately. - # While this is not an issue for ZeRO 1 & 2, since they apply a single optimization step to the whole param group at the same time. - # In order to keep backward compatibility for the existing checkpoints, we use group['state'] to initialize state['step'] if it exists. - state['step'] = group.get('step', 0) - # Exponential moving average of gradient values - state['exp_avg'] = torch.zeros_like(p.data) - # Exponential moving average of squared gradient values - state['exp_avg_sq'] = torch.zeros_like(p.data) - - if p.dtype == torch.float16: - g_16.append(p.grad.data) - p_16.append(p.data) - m_16.append(state['exp_avg']) - v_16.append(state['exp_avg_sq']) - elif p.dtype == torch.bfloat16: - g_bf.append(p.grad) - p_bf.append(p) - m_bf.append(state['exp_avg']) - v_bf.append(state['exp_avg_sq']) - elif p.dtype == torch.float32: - g_32.append(p.grad.data) - p_32.append(p.data) - m_32.append(state['exp_avg']) - v_32.append(state['exp_avg_sq']) - else: - raise RuntimeError('FusedAdam only support fp16, bf16 and fp32.') - - if len(g_16) > 0: - state['step'] += 1 - multi_tensor_applier(self.multi_tensor_adam, self._dummy_overflow_buf, [g_16, p_16, m_16, v_16], - group['lr'], beta1, beta2, group['eps'], state['step'], self.adam_w_mode, - bias_correction, group['weight_decay']) - - if len(g_bf) > 0: - state['step'] += 1 - multi_tensor_applier(self.multi_tensor_adam, self._dummy_overflow_buf, [g_bf, p_bf, m_bf, v_bf], - group['lr'], beta1, beta2, group['eps'], state['step'], self.adam_w_mode, - bias_correction, group['weight_decay']) - - if len(g_32) > 0: - state['step'] += 1 - multi_tensor_applier(self.multi_tensor_adam, self._dummy_overflow_buf, [g_32, p_32, m_32, v_32], - group['lr'], beta1, beta2, group['eps'], state['step'], self.adam_w_mode, - bias_correction, group['weight_decay']) - - return loss diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/multi_tensor_apply.py b/venv/lib/python3.10/site-packages/deepspeed/ops/adam/multi_tensor_apply.py deleted file mode 100644 index 0ba228505cef747eea4fec62f3e68707fa4daa0c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/adam/multi_tensor_apply.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Copyright NVIDIA/apex -This file is adapted from NVIDIA/apex, commit a109f85 -""" - - -class MultiTensorApply(object): - - def __init__(self, chunk_size): - self.chunk_size = chunk_size - - def __call__(self, op, noop_flag_buffer, tensor_lists, *args): - return op(self.chunk_size, noop_flag_buffer, tensor_lists, *args) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/aio/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/aio/__init__.py deleted file mode 100644 index 6aea264237e366b9978c8bd096a2fbc28945705e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/aio/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from ..op_builder import AsyncIOBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/aio/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/aio/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 199dbff1fd824c073030a17d24759bd4efb9fc75..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/aio/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/compile/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/compile/__init__.py deleted file mode 100644 index e38d56359fea64b067373b0e56cfed1dca5dab9a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/compile/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from ..op_builder import DeepCompileBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/compile/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/compile/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index e45f60e10aa7984b09a9b2b42e9e0382da30a6db..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/compile/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adagrad/cpu_adagrad.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adagrad/cpu_adagrad.cpp deleted file mode 100644 index e276ad0856dd620ecbe47231c5ebc4a8a753cb1f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adagrad/cpu_adagrad.cpp +++ /dev/null @@ -1,215 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "cpu_adagrad.h" -#include -#include -#include -#include -#include -#include -#include - -using namespace std::string_literals; -static std::unordered_map> s_optimizers; - -// C++ interface - -template -void Adagrad_Optimizer::Step_1(ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg_sq, - size_t _param_size) -{ - size_t rounded_size = 0; -#if defined(__AVX512__) or defined(__AVX256__) - Step_AVX<1>(&rounded_size, _params, grads, _exp_avg_sq, _param_size); -#endif - if (_param_size > rounded_size) { - float step_size = -1 * _alpha; - for (size_t t = rounded_size; t < _param_size; t += TILE) { - size_t copy_size = TILE; - if ((t + TILE) > _param_size) copy_size = _param_size - t; - size_t offset = copy_size + t; -#pragma omp parallel for - for (size_t k = t; k < offset; k++) { - float grad = (float)grads[k]; - float param = (float)_params[k]; - float momentum = grads[k]; - float variance = _exp_avg_sq[k]; - if (_weight_decay > 0) { grad = param * _weight_decay + grad; } - - variance += grad * grad; - - grad = sqrt(variance); - grad += _eps; - grad = momentum / grad; - param = grad * step_size + param; - _params[k] = param; - // STORE UPDATE TERM TO GRAD'S MEMORY - grads[k] = grad * step_size; - _exp_avg_sq[k] = variance; - } - } - } -} - -template -void Adagrad_Optimizer::Step_4(ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg_sq, - size_t _param_size) -{ - size_t rounded_size = 0; -#if defined(__AVX512__) or defined(__AVX256__) - Step_AVX<4>(&rounded_size, _params, grads, _exp_avg_sq, _param_size); -#endif - if (_param_size > rounded_size) - Step_1((_params + rounded_size), - (grads + rounded_size), - (_exp_avg_sq + rounded_size), - (_param_size - rounded_size)); -} - -int create_adagrad_optimizer(int optimizer_id, - float alpha = 1e-2, - float eps = 1e-8, - float weight_decay = 0, - bool should_log = false) -{ - auto opt = std::make_shared(alpha, eps, weight_decay); - - s_optimizers[optimizer_id] = opt; - - if (should_log) { - std::string avx_type = ""; -#if defined(__AVX512__) - avx_type = "AVX512"; -#else -#if defined(__AVX256__) - avx_type = "AVX2"; -#else - avx_type = "scalar"; -#endif -#endif - - printf("Adagrad Optimizer #%d is created with %s arithmetic capability.\n", - optimizer_id, - avx_type.c_str()); - printf("Config: alpha=%f, weight_decay=%f\n", alpha, weight_decay); - } - - return 0; -} - -template -void Adagrad_Optimizer::Step_8(ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg_sq, - size_t _param_size) -{ - size_t rounded_size = 0; -#if defined(__AVX512__) or defined(__AVX256__) - Step_AVX<8>(&rounded_size, _params, grads, _exp_avg_sq, _param_size); -#endif - if (_param_size > rounded_size) - Step_4((_params + rounded_size), - (grads + rounded_size), - (_exp_avg_sq + rounded_size), - (_param_size - rounded_size)); -} - -template -void step_invoker(std::shared_ptr opt, - void* _params, - void* grads, - void* _exp_avg_sq, - size_t _param_size) -{ - opt->Step_8((ds_params_precision_t*)(_params), - (ds_params_precision_t*)(grads), - (ds_state_precision_t*)(_exp_avg_sq), - _param_size); -} - -std::map, - std::function, void*, void*, void*, size_t)>> - invokers; - -// Fill map with template functions for each type -template -void create_invoker() -{ - invokers[std::tuple(c10::CppTypeToScalarType(), - c10::CppTypeToScalarType())] = - step_invoker; -} -struct InvokerInitializer { - InvokerInitializer() - { - create_invoker(); - create_invoker(); - create_invoker(); - create_invoker(); - create_invoker(); - } -} _invoker_initializer; - -void invoke(std::shared_ptr opt, - torch::Tensor& params, - torch::Tensor& grads, - torch::Tensor& exp_avg_sq, - size_t param_size) -{ - c10::ScalarType params_type = at::typeMetaToScalarType(params.options().dtype()); - c10::ScalarType state_type = at::typeMetaToScalarType(exp_avg_sq.options().dtype()); - - auto it = invokers.find(std::tuple(params_type, state_type)); - if (it == invokers.end()) { - throw std::runtime_error("Adagrad optimizer with param type "s + - c10::toString(params_type) + " and state type "s + - c10::toString(state_type) + - " is not supported on current hardware"s); - } - - it->second(opt, params.data_ptr(), grads.data_ptr(), exp_avg_sq.data_ptr(), param_size); -} - -int ds_adagrad_step(int optimizer_id, - size_t step, - float lr, - float epsilon, - float weight_decay, - torch::Tensor& params, - torch::Tensor& grads, - torch::Tensor& exp_avg_sq) -{ - auto params_c = params.contiguous(); - auto grads_c = grads.contiguous(); - auto exp_avg_sq_c = exp_avg_sq.contiguous(); - - std::shared_ptr opt = - std::static_pointer_cast(s_optimizers[optimizer_id]); - opt->IncrementStep(step); - opt->update_state(lr, epsilon, weight_decay); - - invoke(opt, params_c, grads_c, exp_avg_sq_c, params_c.numel()); - - return 0; -} - -int destroy_adagrad_optimizer(int optimizer_id) -{ - s_optimizers.erase(optimizer_id); - - return 0; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("adagrad_update", &ds_adagrad_step, "DeepSpeed CPU Adagrad update (C++)"); - m.def("create_adagrad", &create_adagrad_optimizer, "DeepSpeed CPU Adagrad (C++)"); - m.def("destroy_adagrad", &destroy_adagrad_optimizer, "DeepSpeed CPU Adagrad destroy (C++)"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/cpu_adam.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/cpu_adam.cpp deleted file mode 100644 index 263c443cb4d4e5f8135f208a6fde9ced513ac8ae..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/cpu_adam.cpp +++ /dev/null @@ -1,13 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "cpu_adam.h" - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("adam_update", &ds_adam_step, "DeepSpeed CPU Adam update (C++)"); - m.def("create_adam", &create_adam_optimizer, "DeepSpeed CPU Adam (C++)"); - m.def("destroy_adam", &destroy_adam_optimizer, "DeepSpeed CPU Adam destroy (C++)"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/cpu_adam_impl.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/cpu_adam_impl.cpp deleted file mode 100644 index 465aae7b9a343d5a3315c1531e11e8648c585b58..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/cpu_adam_impl.cpp +++ /dev/null @@ -1,244 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include -#include -#include -#include -#include -#include -#include -#include "cpu_adam.h" - -using namespace std::string_literals; -static std::unordered_map> s_optimizers; - -// C++ interface - -template -void Adam_Optimizer::Step_1(ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg, - ds_state_precision_t* _exp_avg_sq, - size_t _param_size) -{ - size_t rounded_size = 0; -#if defined(__AVX512__) or defined(__AVX256__) - Step_AVX<1>(&rounded_size, _params, grads, _exp_avg, _exp_avg_sq, _param_size); -#endif - if (_param_size > rounded_size) { - float betta1_minus1 = 1 - _betta1; - float betta2_minus1 = 1 - _betta2; - - float step_size = -1 * _alpha / _bias_correction1; - float w_decay = -1 * _alpha * _weight_decay; - - for (size_t t = rounded_size; t < _param_size; t += TILE) { - size_t copy_size = TILE; - if ((t + TILE) > _param_size) copy_size = _param_size - t; - size_t offset = copy_size + t; -#pragma omp parallel for - for (size_t k = t; k < offset; k++) { - float grad = (float)grads[k]; - float param = (float)_params[k]; - float momentum = _exp_avg[k]; - float variance = _exp_avg_sq[k]; - if (_weight_decay > 0 && !_adamw_mode) { grad = param * _weight_decay + grad; } - momentum = momentum * _betta1; - momentum = grad * betta1_minus1 + momentum; - - variance = variance * _betta2; - grad = grad * grad; - variance = grad * betta2_minus1 + variance; - - grad = sqrt(variance); - grad = grad * _bias_correction2 + _eps; - grad = momentum / grad; - if (_weight_decay > 0 && _adamw_mode) { param += w_decay * param; } - param = grad * step_size + param; - _params[k] = param; - _exp_avg[k] = momentum; - _exp_avg_sq[k] = variance; - } - } - } -} - -template -void Adam_Optimizer::Step_4(ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg, - ds_state_precision_t* _exp_avg_sq, - size_t _param_size) -{ - size_t rounded_size = 0; -#if defined(__AVX512__) or defined(__AVX256__) - Step_AVX<4>(&rounded_size, _params, grads, _exp_avg, _exp_avg_sq, _param_size); -#endif - if (_param_size > rounded_size) - Step_1((_params + rounded_size), - (grads + rounded_size), - (_exp_avg + rounded_size), - (_exp_avg_sq + rounded_size), - (_param_size - rounded_size)); -} - -int create_adam_optimizer(int optimizer_id, - float alpha, - float betta1, - float betta2, - float eps, - float weight_decay, - bool adamw_mode, - bool should_log) -{ - auto opt = - std::make_shared(alpha, betta1, betta2, eps, weight_decay, adamw_mode); - - s_optimizers[optimizer_id] = opt; - - if (should_log) { - std::string avx_type = ""; -#if defined(__AVX512__) - avx_type = "AVX512"; -#else -#if defined(__AVX256__) - avx_type = "AVX2"; -#else - avx_type = "scalar"; -#endif -#endif - - printf("Adam Optimizer #%d is created with %s arithmetic capability.\n", - optimizer_id, - avx_type.c_str()); - printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n", - alpha, - betta1, - betta2, - weight_decay, - (int)adamw_mode); - } - - return 0; -} - -template -void Adam_Optimizer::Step_8(ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg, - ds_state_precision_t* _exp_avg_sq, - size_t _param_size) -{ - size_t rounded_size = 0; -#if defined(__AVX512__) or defined(__AVX256__) - Step_AVX<8>(&rounded_size, _params, grads, _exp_avg, _exp_avg_sq, _param_size); -#endif - if (_param_size > rounded_size) - Step_4((_params + rounded_size), - (grads + rounded_size), - (_exp_avg + rounded_size), - (_exp_avg_sq + rounded_size), - (_param_size - rounded_size)); -} - -template -void step_invoker(std::shared_ptr opt, - void* _params, - void* grads, - void* _exp_avg, - void* _exp_avg_sq, - size_t _param_size) -{ - opt->Step_8((ds_params_precision_t*)(_params), - (ds_params_precision_t*)(grads), - (ds_state_precision_t*)(_exp_avg), - (ds_state_precision_t*)(_exp_avg_sq), - _param_size); -} - -std::map, - std::function, void*, void*, void*, void*, size_t)>> - invokers; - -// Fill map with template functions for each type -template -void create_invoker() -{ - invokers[std::tuple(c10::CppTypeToScalarType(), - c10::CppTypeToScalarType())] = - step_invoker; -} -struct InvokerInitializer { - InvokerInitializer() - { - create_invoker(); - create_invoker(); - create_invoker(); - create_invoker(); - create_invoker(); - } -} _invoker_initializer; - -void invoke(std::shared_ptr opt, - torch::Tensor& params, - torch::Tensor& grads, - torch::Tensor& exp_avg, - torch::Tensor& exp_avg_sq, - size_t param_size) -{ - c10::ScalarType params_type = at::typeMetaToScalarType(params.options().dtype()); - c10::ScalarType state_type = at::typeMetaToScalarType(exp_avg.options().dtype()); - - auto it = invokers.find(std::tuple(params_type, state_type)); - if (it == invokers.end()) { - throw std::runtime_error("Adam optimizer with param type "s + c10::toString(params_type) + - " and state type "s + c10::toString(state_type) + - " is not supported on current hardware"s); - } - - it->second(opt, - params.data_ptr(), - grads.data_ptr(), - exp_avg.data_ptr(), - exp_avg_sq.data_ptr(), - param_size); -} - -int ds_adam_step(int optimizer_id, - size_t step, - float lr, - float beta1, - float beta2, - float epsilon, - float weight_decay, - bool bias_correction, - torch::Tensor& params, - torch::Tensor& grads, - torch::Tensor& exp_avg, - torch::Tensor& exp_avg_sq) -{ - auto params_c = params.contiguous(); - auto grads_c = grads.contiguous(); - auto exp_avg_c = exp_avg.contiguous(); - auto exp_avg_sq_c = exp_avg_sq.contiguous(); - - std::shared_ptr opt = - std::static_pointer_cast(s_optimizers[optimizer_id]); - opt->IncrementStep(step, beta1, beta2); - opt->update_state(lr, epsilon, weight_decay, bias_correction); - - invoke(opt, params_c, grads_c, exp_avg_c, exp_avg_sq_c, params_c.numel()); - - return 0; -} - -int destroy_adam_optimizer(int optimizer_id) -{ - s_optimizers.erase(optimizer_id); - - return 0; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp deleted file mode 100644 index 13b390248608b046dab443f85346b5446a47d722..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp +++ /dev/null @@ -1,25 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include - -void multi_tensor_adam_cuda(int chunk_size, - at::Tensor noop_flag, - std::vector> tensor_lists, - const float lr, - const float beta1, - const float beta2, - const float epsilon, - const int step, - const int mode, - const int bias_correction, - const float weight_decay); - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("multi_tensor_adam", - &multi_tensor_adam_cuda, - "Compute and apply gradient update to parameters for Adam optimizer"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu deleted file mode 100644 index a1fc7d15aec9ad16ddb242fd62c99cbfdca55d0c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu +++ /dev/null @@ -1,203 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Copyright NVIDIA/apex -This file is adapted from fused adam in NVIDIA/apex, commit a109f85 -*/ - -#include -#include -#include -#include -// Another possibility: -// #include - -#include - -#include "multi_tensor_apply.cuh" -#include "type_shim.h" - -#define BLOCK_SIZE 512 -#define ILP 4 - -typedef enum : int { - ADAM_MODE_0 = 0, // L2 regularization mode - ADAM_MODE_1 = 1 // Decoupled weight decay mode(AdamW) -} adamMode_t; - -using MATH_T = float; - -template -struct AdamFunctor { - __device__ __forceinline__ void operator()(int chunk_size, - volatile int* noop_gmem, - TensorListMetadata<4>& tl, - const float beta1, - const float beta2, - const float beta1_correction, - const float beta2_correction, - const float epsilon, - const float lr, - adamMode_t mode, - const float decay) - { - // I'd like this kernel to propagate infs/nans. - // if(*noop_gmem == 1) - // return; - - index_t tensor_loc = tl.block_to_tensor[blockIdx.x]; - - // potentially use to pass in list of scalar - // int tensor_num = tl.start_tensor_this_launch + tensor_loc; - - index_t chunk_idx = tl.block_to_chunk[blockIdx.x]; - index_t n = tl.sizes[tensor_loc]; - - T* g = (T*)tl.addresses[0][tensor_loc]; - g += chunk_idx * chunk_size; - - T* p = (T*)tl.addresses[1][tensor_loc]; - p += chunk_idx * chunk_size; - - T* m = (T*)tl.addresses[2][tensor_loc]; - m += chunk_idx * chunk_size; - - T* v = (T*)tl.addresses[3][tensor_loc]; - v += chunk_idx * chunk_size; - - n -= chunk_idx * chunk_size; - - // see note in multi_tensor_scale_kernel.cu - for (index_t i_start = 0; i_start < n && i_start < chunk_size; - i_start += blockDim.x * ILP) { - MATH_T r_g[ILP]; - MATH_T r_p[ILP]; - MATH_T r_m[ILP]; - MATH_T r_v[ILP]; -#pragma unroll - for (int ii = 0; ii < ILP; ii++) { - int i = i_start + threadIdx.x + ii * blockDim.x; - if (i < n && i < chunk_size) { - r_g[ii] = g[i]; - r_p[ii] = p[i]; - r_m[ii] = m[i]; - r_v[ii] = v[i]; - } else { - r_g[ii] = MATH_T(0); - r_p[ii] = MATH_T(0); - r_m[ii] = MATH_T(0); - r_v[ii] = MATH_T(0); - } - } -#pragma unroll - for (int ii = 0; ii < ILP; ii++) { - if (mode == ADAM_MODE_0) { // L2 - r_g[ii] = r_g[ii] + (decay * r_p[ii]); - r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; - r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii]; - MATH_T next_m_unbiased = r_m[ii] / beta1_correction; - MATH_T next_v_unbiased = r_v[ii] / beta2_correction; - MATH_T denom = sqrtf(next_v_unbiased) + epsilon; - MATH_T update = next_m_unbiased / denom; - r_p[ii] = r_p[ii] - (lr * update); - } else { // weight decay - r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; - r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii]; - MATH_T next_m_unbiased = r_m[ii] / beta1_correction; - MATH_T next_v_unbiased = r_v[ii] / beta2_correction; - MATH_T denom = sqrtf(next_v_unbiased) + epsilon; - MATH_T update = (next_m_unbiased / denom) + (decay * r_p[ii]); - r_p[ii] = r_p[ii] - (lr * update); - } - } -#pragma unroll - for (int ii = 0; ii < ILP; ii++) { - int i = i_start + threadIdx.x + ii * blockDim.x; - if (i < n && i < chunk_size) { - p[i] = r_p[ii]; - m[i] = r_m[ii]; - v[i] = r_v[ii]; - } - } - } - } -}; - -void multi_tensor_adam_cuda(int chunk_size, - at::Tensor noop_flag, - std::vector> tensor_lists, - const float lr, - const float beta1, - const float beta2, - const float epsilon, - const int step, - const int mode, - const int bias_correction, - const float weight_decay) -{ - using namespace at; - - // Handle bias correction mode - float bias_correction1 = 1.0f, bias_correction2 = 1.0f; - if (bias_correction == 1) { - bias_correction1 = 1 - std::pow(beta1, step); - bias_correction2 = 1 - std::pow(beta2, step); - } - - size_t max_size = 0; - bool requires_64bit_indexing = false; - for (auto it = tensor_lists.begin(); it != tensor_lists.end(); it++) { - for (auto it2 = it->begin(); it2 != it->end(); it2++) { - if (it2->numel() > max_size) { - max_size = it2->numel(); - if (max_size >= INT_MAX) { - requires_64bit_indexing = true; - break; - } - } - } - if (requires_64bit_indexing) { break; } - } - - // Assume single type across p,g,m1,m2 now - if (requires_64bit_indexing) { - DISPATCH_DOUBLE_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), - 0, - "adam", - multi_tensor_apply<4>((int64_t)BLOCK_SIZE, - (int64_t)chunk_size, - noop_flag, - tensor_lists, - AdamFunctor(), - beta1, - beta2, - bias_correction1, - bias_correction2, - epsilon, - lr, - (adamMode_t)mode, - weight_decay);) - } else { - DISPATCH_DOUBLE_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), - 0, - "adam", - multi_tensor_apply<4>(BLOCK_SIZE, - chunk_size, - noop_flag, - tensor_lists, - AdamFunctor(), - beta1, - beta2, - bias_correction1, - bias_correction2, - epsilon, - lr, - (adamMode_t)mode, - weight_decay);) - } - - AT_CUDA_CHECK(cudaGetLastError()); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_apply.cuh b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_apply.cuh deleted file mode 100644 index 342376c141be7578553910df43b003d29e0e4fc4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_apply.cuh +++ /dev/null @@ -1,132 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Copyright NVIDIA/apex -This file is adapted from fused adam in NVIDIA/apex, commit a109f85 -*/ - -#include -#include -#include -#include -#include -#include "compat.h" - -#include - -// #include - -// This header is the one-stop shop for all your multi-tensor apply needs. - -// TODO: Kernel arg size limit may be <4KB for some other cards (ie Jetson) -constexpr int depth_to_max_tensors[5] = {110, 64, 48, 36, 30}; -constexpr int depth_to_max_blocks[5] = {320, 320, 320, 320, 320}; - -template -struct TensorListMetadata { - void* addresses[n][depth_to_max_tensors[n - 1]]; - int sizes[depth_to_max_tensors[n - 1]]; - unsigned char block_to_tensor[depth_to_max_blocks[n - 1]]; - int block_to_chunk[depth_to_max_blocks[n - 1]]; // I fear this needs to be a full int. - int start_tensor_this_launch; -}; - -template -__global__ void multi_tensor_apply_kernel(int64_t chunk_size, - volatile int* noop_flag, - T tl, - U callable, - ArgTypes... args) -{ - // Hand the chunk information to the user-supplied functor to process however it likes. - callable(chunk_size, noop_flag, tl, args...); -} - -template -void multi_tensor_apply(int64_t block_size, - int64_t chunk_size, - const at::Tensor& noop_flag, - const std::vector>& tensor_lists, - T callable, - ArgTypes... args) -{ - TORCH_CHECK(tensor_lists.size() == depth, "tensor_lists.size() != depth"); - int len0 = tensor_lists[0].size(); - TORCH_CHECK(len0 > 0, "tensor_lists[0].size() is not > 0"); - auto ref_device = tensor_lists[0][0].device(); - TORCH_CHECK(ref_device.type() == at::kCUDA, "expected input to be on cuda"); - for (int l = 0; l < tensor_lists.size(); l++) // No range-based for because I need indices - { - TORCH_CHECK(tensor_lists[l].size() == len0, "Size mismatch among tensor lists"); - for (int t = 0; t < tensor_lists[l].size(); t++) { - // TODO: Print which tensor fails. - bool contiguous_memory = tensor_lists[l][t].is_contiguous(); -#ifdef VERSION_GE_1_5 - contiguous_memory = (contiguous_memory || - tensor_lists[l][t].is_contiguous(at::MemoryFormat::ChannelsLast)); -#endif - TORCH_CHECK(contiguous_memory, "A tensor was not contiguous."); - TORCH_CHECK(tensor_lists[l][t].device() == ref_device, - "A tensor was not on the same device as the first tensor"); - TORCH_CHECK(tensor_lists[l][t].numel() == tensor_lists[0][t].numel(), "Size mismatch"); - } - } - - int ntensors = tensor_lists[0].size(); - - TensorListMetadata tl; - - const at::cuda::OptionalCUDAGuard device_guard(device_of(tensor_lists[0][0])); - auto stream = at::cuda::getCurrentCUDAStream(); - - tl.start_tensor_this_launch = 0; - int loc_block_info = 0; - int loc_tensor_info = 0; - for (int t = 0; t < ntensors; t++) { - tl.sizes[loc_tensor_info] = tensor_lists[0][t].numel(); - for (int d = 0; d < depth; d++) - tl.addresses[d][loc_tensor_info] = tensor_lists[d][t].data_ptr(); - loc_tensor_info++; - - auto chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size; - - for (auto chunk = 0; chunk < chunks_this_tensor; chunk++) { - // std::cout << chunks_this_tensor << std::endl; - tl.block_to_tensor[loc_block_info] = loc_tensor_info - 1; - tl.block_to_chunk[loc_block_info] = chunk; - loc_block_info++; - - bool tensors_full = (loc_tensor_info == depth_to_max_tensors[depth - 1] && - chunk == chunks_this_tensor - 1); - bool blocks_full = (loc_block_info == depth_to_max_blocks[depth - 1]); - bool last_chunk = (t == ntensors - 1 && chunk == chunks_this_tensor - 1); - if (tensors_full || blocks_full || last_chunk) { - // using accscalar_t = acc_type; - multi_tensor_apply_kernel<<>>( - chunk_size, noop_flag.DATA_PTR(), tl, callable, args...); - - AT_CUDA_CHECK(cudaGetLastError()); - - // Reset. The control flow possibilities here make my brain hurt. - loc_block_info = 0; - if (chunk == chunks_this_tensor - 1) { - // std::cout << "Hit case 1 " << cond1 << " " << cond2 << " " << cond3 << - // std::endl; - loc_tensor_info = 0; - tl.start_tensor_this_launch = t + 1; - } else { - // std::cout << "Hit case 2 " << cond1 << " " << cond2 << " " << cond3 << - // std::endl; - tl.sizes[0] = tl.sizes[loc_tensor_info - 1]; - for (int d = 0; d < depth; d++) - tl.addresses[d][0] = tl.addresses[d][loc_tensor_info - 1]; - loc_tensor_info = 1; - tl.start_tensor_this_launch = t; - } - } - } - } -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.cpp deleted file mode 100644 index 9d7ff5093017f60c06b96554f69630ffe9825918..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.cpp +++ /dev/null @@ -1,341 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "deepspeed_aio_common.h" - -using namespace std; -using namespace std::chrono; - -#define DEBUG_DS_AIO_PERF 0 -#define DEBUG_DS_AIO_SUBMIT_PERF 0 - -static const std::string c_library_name = "deepspeed_aio"; - -static void _report_aio_statistics(const char* tag, - const std::vector>& latencies) - __attribute__((unused)); - -static void _report_aio_statistics(const char* tag, - const std::vector>& latencies) -{ - std::vector lat_usec; - for (auto& lat : latencies) { lat_usec.push_back(lat.count() * 1e6); } - const auto min_lat = *(std::min_element(lat_usec.begin(), lat_usec.end())); - const auto max_lat = *(std::max_element(lat_usec.begin(), lat_usec.end())); - const auto avg_lat = std::accumulate(lat_usec.begin(), lat_usec.end(), 0) / lat_usec.size(); - - std::cout << c_library_name << ": latency statistics(usec) " << tag - << " min/max/avg = " << min_lat << " " << max_lat << " " << avg_lat << std::endl; -} - -static void _get_aio_latencies(std::vector>& raw_latencies, - struct deepspeed_aio_latency_t& summary_latencies) -{ - std::vector lat_usec; - for (auto& lat : raw_latencies) { lat_usec.push_back(lat.count() * 1e6); } - summary_latencies._min_usec = *(std::min_element(lat_usec.begin(), lat_usec.end())); - summary_latencies._max_usec = *(std::max_element(lat_usec.begin(), lat_usec.end())); - summary_latencies._avg_usec = - std::accumulate(lat_usec.begin(), lat_usec.end(), 0) / lat_usec.size(); -} - -static void _do_io_submit_singles(const int64_t n_iocbs, - const int64_t iocb_index, - std::unique_ptr& aio_ctxt, - std::vector>& submit_times) -{ - for (auto i = 0; i < n_iocbs; ++i) { - const auto st = std::chrono::high_resolution_clock::now(); - const auto submit_ret = io_submit(aio_ctxt->_io_ctxt, 1, aio_ctxt->_iocbs.data() + i); - submit_times.push_back(std::chrono::high_resolution_clock::now() - st); -#if DEBUG_DS_AIO_SUBMIT_PERF - printf("submit(usec) %f io_index=%lld buf=%p len=%lu off=%llu \n", - submit_times.back().count() * 1e6, - iocb_index, - aio_ctxt->_iocbs[i]->u.c.buf, - aio_ctxt->_iocbs[i]->u.c.nbytes, - aio_ctxt->_iocbs[i]->u.c.offset); -#endif - assert(submit_ret > 0); - } -} - -static void _do_io_submit_block(const int64_t n_iocbs, - const int64_t iocb_index, - std::unique_ptr& aio_ctxt, - std::vector>& submit_times) -{ - const auto st = std::chrono::high_resolution_clock::now(); - const auto submit_ret = io_submit(aio_ctxt->_io_ctxt, n_iocbs, aio_ctxt->_iocbs.data()); - submit_times.push_back(std::chrono::high_resolution_clock::now() - st); -#if DEBUG_DS_AIO_SUBMIT_PERF - printf("submit(usec) %f io_index=%lld nr=%lld buf=%p len=%lu off=%llu \n", - submit_times.back().count() * 1e6, - iocb_index, - n_iocbs, - aio_ctxt->_iocbs[0]->u.c.buf, - aio_ctxt->_iocbs[0]->u.c.nbytes, - aio_ctxt->_iocbs[0]->u.c.offset); -#endif - assert(submit_ret > 0); -} - -static int _do_io_complete(const int64_t min_completes, - const int64_t max_completes, - std::unique_ptr& aio_ctxt, - std::vector>& reap_times) -{ - const auto start_time = std::chrono::high_resolution_clock::now(); - int64_t n_completes = io_pgetevents(aio_ctxt->_io_ctxt, - min_completes, - max_completes, - aio_ctxt->_io_events.data(), - nullptr, - nullptr); - reap_times.push_back(std::chrono::high_resolution_clock::now() - start_time); - assert(n_completes >= min_completes); - return n_completes; -} - -void do_aio_operation_sequential(const bool read_op, - std::unique_ptr& aio_ctxt, - std::unique_ptr& xfer_ctxt, - deepspeed_aio_config_t* config, - deepspeed_aio_perf_t* perf) -{ - struct io_prep_context prep_ctxt(read_op, xfer_ctxt, aio_ctxt->_block_size, &aio_ctxt->_iocbs); - - const auto num_io_blocks = static_cast( - ceil(static_cast(xfer_ctxt->_num_bytes) / aio_ctxt->_block_size)); -#if DEBUG_DS_AIO_PERF - const auto io_op_name = std::string(read_op ? "read" : "write"); - std::cout << c_library_name << ": start " << io_op_name << " " << xfer_ctxt->_num_bytes - << " bytes with " << num_io_blocks << " io blocks" << std::endl; -#endif - - std::vector> submit_times; - std::vector> reap_times; - const auto max_queue_bytes = - static_cast(aio_ctxt->_queue_depth * aio_ctxt->_block_size); - - auto start = std::chrono::high_resolution_clock::now(); - for (int64_t iocb_index = 0; iocb_index < num_io_blocks; iocb_index += aio_ctxt->_queue_depth) { - const auto start_offset = iocb_index * aio_ctxt->_block_size; - const auto start_buffer = (char*)xfer_ctxt->_mem_buffer + start_offset; - const auto n_iocbs = - min(static_cast(aio_ctxt->_queue_depth), (num_io_blocks - iocb_index)); - const auto num_bytes = min(max_queue_bytes, (xfer_ctxt->_num_bytes - start_offset)); - prep_ctxt.prep_iocbs(n_iocbs, num_bytes, start_buffer, start_offset); - - if (config->_single_submit) { - _do_io_submit_singles(n_iocbs, iocb_index, aio_ctxt, submit_times); - } else { - _do_io_submit_block(n_iocbs, iocb_index, aio_ctxt, submit_times); - } - - _do_io_complete(n_iocbs, n_iocbs, aio_ctxt, reap_times); - } - const std::chrono::duration elapsed = std::chrono::high_resolution_clock::now() - start; - - if (perf) { - _get_aio_latencies(submit_times, perf->_submit); - _get_aio_latencies(reap_times, perf->_complete); - perf->_e2e_usec = elapsed.count() * 1e6; - perf->_e2e_rate_GB = (xfer_ctxt->_num_bytes / elapsed.count() / 1e9); - } - -#if DEBUG_DS_AIO_PERF - _report_aio_statistics("submit", submit_times); - _report_aio_statistics("complete", reap_times); -#endif - -#if DEBUG_DS_AIO_PERF - std::cout << c_library_name << ": runtime(usec) " << elapsed.count() * 1e6 - << " rate(GB/sec) = " << (xfer_ctxt->_num_bytes / elapsed.count() / 1e9) << std::endl; -#endif - -#if DEBUG_DS_AIO_PERF - std::cout << c_library_name << ": finish " << io_op_name << " " << xfer_ctxt->_num_bytes - << " bytes " << std::endl; -#endif -} - -void do_aio_operation_overlap(const bool read_op, - std::unique_ptr& aio_ctxt, - std::unique_ptr& xfer_ctxt, - deepspeed_aio_config_t* config, - deepspeed_aio_perf_t* perf) -{ - struct io_prep_generator io_gen(read_op, xfer_ctxt, aio_ctxt->_block_size); - -#if DEBUG_DS_AIO_PERF - const auto io_op_name = std::string(read_op ? "read" : "write"); - std::cout << c_library_name << ": start " << io_op_name << " " << xfer_ctxt->_num_bytes - << " bytes with " << io_gen._num_io_blocks << " io blocks" << std::endl; -#endif - - std::vector> submit_times; - std::vector> reap_times; - - auto request_iocbs = aio_ctxt->_queue_depth; - auto n_pending_iocbs = 0; - const auto min_completes = 1; - auto start = std::chrono::high_resolution_clock::now(); - while (true) { - const auto n_iocbs = io_gen.prep_iocbs(request_iocbs - n_pending_iocbs, &aio_ctxt->_iocbs); - if (n_iocbs > 0) { - if (config->_single_submit) { - _do_io_submit_singles( - n_iocbs, (io_gen._next_iocb_index - n_iocbs), aio_ctxt, submit_times); - } else { - _do_io_submit_block( - n_iocbs, (io_gen._next_iocb_index - n_iocbs), aio_ctxt, submit_times); - } - } - - n_pending_iocbs += n_iocbs; - assert(n_pending_iocbs <= aio_ctxt->_queue_depth); - - if (n_pending_iocbs == 0) { break; } - - const auto n_complete = - _do_io_complete(min_completes, n_pending_iocbs, aio_ctxt, reap_times); - n_pending_iocbs -= n_complete; - } - - const std::chrono::duration elapsed = std::chrono::high_resolution_clock::now() - start; - - if (perf) { - _get_aio_latencies(submit_times, perf->_submit); - _get_aio_latencies(reap_times, perf->_complete); - perf->_e2e_usec = elapsed.count() * 1e6; - perf->_e2e_rate_GB = (xfer_ctxt->_num_bytes / elapsed.count() / 1e9); - } - -#if DEBUG_DS_AIO_PERF - _report_aio_statistics("submit", submit_times); - _report_aio_statistics("complete", reap_times); -#endif - -#if DEBUG_DS_AIO_PERF - std::cout << c_library_name << ": runtime(usec) " << elapsed.count() * 1e6 - << " rate(GB/sec) = " << (xfer_ctxt->_num_bytes / elapsed.count() / 1e9) << std::endl; -#endif - -#if DEBUG_DS_AIO_PERF - std::cout << c_library_name << ": finish " << io_op_name << " " << xfer_ctxt->_num_bytes - << " bytes " << std::endl; -#endif -} - -void report_file_error(const char* filename, const std::string file_op, const int error_code) -{ - std::string err_msg = file_op + std::string(" failed on ") + std::string(filename) + - " error = " + std::to_string(error_code); - std::cerr << c_library_name << ": " << err_msg << std::endl; -} - -int open_file(const char* filename, const bool read_op) -{ - const int flags = read_op ? (O_RDONLY | O_DIRECT) : (O_WRONLY | O_CREAT | O_DIRECT); -#if defined(__ENABLE_CANN__) - int* flags_ptr = (int*)&flags; - *flags_ptr = read_op ? (O_RDONLY) : (O_WRONLY | O_CREAT); -#endif - const int mode = 0600; - const auto fd = open(filename, flags, mode); - if (fd == -1) { - const auto error_code = errno; - const auto error_msg = read_op ? " open for read " : " open for write "; - report_file_error(filename, error_msg, error_code); - return -1; - } - return fd; -} - -int regular_read(const char* filename, std::vector& buffer) -{ - const auto fd = open(filename, O_RDONLY, 0600); - assert(fd != -1); - struct stat fs; - const auto result = fstat(fd, &fs); - assert(result != -1); - int64_t num_bytes = fs.st_size; - buffer.resize(num_bytes); - int64_t read_bytes = 0; - auto r = 0; - do { - const auto buffer_ptr = buffer.data() + read_bytes; - const auto bytes_to_read = num_bytes - read_bytes; - r = read(fd, buffer_ptr, bytes_to_read); - read_bytes += r; - } while (r > 0); - - if (read_bytes != num_bytes) { - std::cerr << "read error " << " read_bytes (read) = " << read_bytes - << " num_bytes (fstat) = " << num_bytes << std::endl; - } - assert(read_bytes == num_bytes); - close(fd); - return 0; -} - -static bool _validate_buffer(const char* filename, void* aio_buffer, const int64_t num_bytes) -{ - std::vector regular_buffer; - const auto reg_ret = regular_read(filename, regular_buffer); - assert(0 == reg_ret); - std::cout << "regular read of " << filename << " returned " << regular_buffer.size() << " bytes" - << std::endl; - - if (static_cast(regular_buffer.size()) != num_bytes) { return false; } - - return (0 == memcmp(aio_buffer, regular_buffer.data(), regular_buffer.size())); -} - -bool validate_aio_operation(const bool read_op, - const char* filename, - void* aio_buffer, - const int64_t num_bytes) -{ - const auto msg_suffix = std::string("deepspeed_aio_") + - std::string(read_op ? "read()" : "write()") + - std::string("using read()"); - - if (false == _validate_buffer(filename, aio_buffer, num_bytes)) { - std::cout << "Fail: correctness of " << msg_suffix << std::endl; - return false; - } - - std::cout << "Pass: correctness of " << msg_suffix << std::endl; - return true; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.h deleted file mode 100644 index aa4e49f4f4edee792e81bc0581e16ef3ab968811..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.h +++ /dev/null @@ -1,38 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include -#include -#include -#include - -using namespace std; - -void do_aio_operation_sequential(const bool read_op, - std::unique_ptr& aio_ctxt, - std::unique_ptr& xfer_ctxt, - deepspeed_aio_config_t* config, - deepspeed_aio_perf_t* perf); - -void do_aio_operation_overlap(const bool read_op, - std::unique_ptr& aio_ctxt, - std::unique_ptr& xfer_ctxt, - deepspeed_aio_config_t* config, - deepspeed_aio_perf_t* perf); - -int open_file(const char* filename, const bool read_op); - -void report_file_error(const char* filename, const std::string file_op, const int error_code); - -int regular_read(const char* filename, std::vector& buffer); - -bool validate_aio_operation(const bool read_op, - const char* filename, - void* aio_buffer, - const int64_t num_bytes); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.cpp deleted file mode 100644 index 5e34a61065d4b0f16ba904d0e397eeccc149e621..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.cpp +++ /dev/null @@ -1,76 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include - -#include "deepspeed_aio_utils.h" - -using namespace std; - -const int c_block_size = 128 * 1024; -const int c_io_queue_depth = 8; - -deepspeed_aio_config_t::deepspeed_aio_config_t() - : _block_size(c_block_size), - _queue_depth(c_io_queue_depth), - _single_submit(false), - _overlap_events(false), - _lock_memory(false) -{ -} - -deepspeed_aio_config_t::deepspeed_aio_config_t(const int block_size, - const int queue_depth, - const bool single_submit, - const bool overlap_events, - const bool lock_memory) - : _block_size(block_size), - _queue_depth(queue_depth), - _single_submit(single_submit), - _overlap_events(overlap_events), - _lock_memory(lock_memory) -{ -} - -void deepspeed_aio_latency_t::dump(const std::string tag) -{ - std::cout << tag << _min_usec << " " << _max_usec << " " << _avg_usec << " " << std::endl; -} - -void deepspeed_aio_latency_t::accumulate(const struct deepspeed_aio_latency_t& other) -{ - _min_usec += other._min_usec; - _max_usec += other._max_usec; - _avg_usec += other._avg_usec; -} - -void deepspeed_aio_latency_t::scale(const float scaler) -{ - _min_usec *= scaler; - _max_usec *= scaler; - _avg_usec *= scaler; -} - -aio_context::aio_context(const int block_size, const int queue_depth) -{ - _block_size = block_size; - _queue_depth = queue_depth; - for (auto i = 0; i < queue_depth; ++i) { - _iocbs.push_back((struct iocb*)calloc(1, sizeof(struct iocb))); - } - _io_events.resize(queue_depth); - io_queue_init(queue_depth, &_io_ctxt); -} - -aio_context::~aio_context() -{ - for (auto& iocb : _iocbs) { free(iocb); } - _io_events.resize(0); - io_queue_release(_io_ctxt); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.h deleted file mode 100644 index ce6a4e5cdfa75120c24dad599a8bc717765c220d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.h +++ /dev/null @@ -1,59 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include -#include - -#include -#include - -using namespace std; - -struct deepspeed_aio_latency_t { - double _min_usec; - double _max_usec; - double _avg_usec; - - void dump(const std::string tag); - void accumulate(const deepspeed_aio_latency_t&); - void scale(const float value); -}; - -struct deepspeed_aio_perf_t { - deepspeed_aio_latency_t _submit; - deepspeed_aio_latency_t _complete; - double _e2e_usec; - double _e2e_rate_GB; -}; - -struct deepspeed_aio_config_t { - const int _block_size; - const int _queue_depth; - const bool _single_submit; - const bool _overlap_events; - const bool _lock_memory; - - deepspeed_aio_config_t(); - deepspeed_aio_config_t(const int block_size, - const int queue_depth, - const bool single_submit, - const bool overlap_events, - const bool lock_memory); -}; - -struct aio_context { - io_context_t _io_ctxt; - std::vector _io_events; - std::vector _iocbs; - int _block_size; - int _queue_depth; - - aio_context(const int block_size, const int queue_depth); - ~aio_context(); -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.cpp deleted file mode 100644 index c8e577f299ae9b8f1aae41d742d5b69395fc026b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.cpp +++ /dev/null @@ -1,140 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include -#include - -#include "deepspeed_aio_utils.h" - -using namespace std; - -const int c_block_size = 128 * 1024; -const int c_io_queue_depth = 8; - -io_xfer_ctxt::io_xfer_ctxt(const int fd, - const int64_t file_offset, - const int64_t buffer_offset, - const int64_t num_bytes, - const void* buffer) - : _fd(fd), - _file_base_offset(file_offset), - _buffer_base_offset(buffer_offset), - _mem_buffer(buffer), - _num_bytes(num_bytes) -{ -} - -io_prep_context::io_prep_context(const bool read_op, - const std::unique_ptr& xfer_ctxt, - const size_t block_size, - const std::vector* iocbs) - : _read_op(read_op), _xfer_ctxt(xfer_ctxt), _block_size(block_size), _iocbs(iocbs) -{ -} - -void io_prep_context::prep_iocbs(const int n_iocbs, - const size_t num_bytes, - const void* start_buffer, - const int64_t start_offset) -{ - assert(static_cast(n_iocbs) <= _iocbs->size()); - for (auto i = 0; i < n_iocbs; ++i) { - const auto shift = i * _block_size; - const auto xfer_buffer = (char*)start_buffer + _xfer_ctxt->_buffer_base_offset + shift; - const auto xfer_offset = _xfer_ctxt->_file_base_offset + start_offset + shift; - auto byte_count = _block_size; - - if ((shift + _block_size) > num_bytes) { byte_count = num_bytes - shift; } - - if (_read_op) { - io_prep_pread(_iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, byte_count, xfer_offset); - } else { - io_prep_pwrite(_iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, byte_count, xfer_offset); - } - } -} - -io_prep_generator::io_prep_generator(const bool read_op, - const std::unique_ptr& xfer_ctxt, - const size_t block_size) - : _read_op(read_op), - _xfer_ctxt(xfer_ctxt), - _block_size(block_size), - _remaining_bytes(xfer_ctxt->_num_bytes), - _next_iocb_index(0) -{ - _num_io_blocks = - static_cast(ceil(static_cast(xfer_ctxt->_num_bytes) / block_size)); - _remaining_io_blocks = _num_io_blocks; -} - -int io_prep_generator::prep_iocbs(const int n_iocbs, std::vector* iocbs) -{ - if ((_remaining_bytes) == 0 || (_remaining_io_blocks == 0)) { - assert(static_cast(_remaining_bytes) == _remaining_io_blocks); - return 0; - } - - assert(static_cast(n_iocbs) <= iocbs->size()); - - auto actual_n_iocbs = min(static_cast(n_iocbs), _remaining_io_blocks); - for (auto i = 0; i < actual_n_iocbs; ++i, ++_next_iocb_index) { - const auto xfer_buffer = (char*)_xfer_ctxt->_mem_buffer + _xfer_ctxt->_buffer_base_offset + - (_next_iocb_index * _block_size); - const auto xfer_offset = _xfer_ctxt->_file_base_offset + (_next_iocb_index * _block_size); - const auto num_bytes = min(static_cast(_block_size), _remaining_bytes); - if (_read_op) { - io_prep_pread(iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, num_bytes, xfer_offset); - } else { - io_prep_pwrite(iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, num_bytes, xfer_offset); - } - _remaining_bytes -= num_bytes; - } - _remaining_io_blocks -= actual_n_iocbs; - - return actual_n_iocbs; -} - -int64_t get_file_size(const char* filename, int64_t& size) -{ - struct stat st; - if (stat(filename, &st) == -1) { return -1; } - size = st.st_size; - return 0; -} - -int64_t get_fd_file_size(const int fd, int64_t& size) -{ - struct stat st; - if (fstat(fd, &st) == -1) { return -1; } - size = st.st_size; - return 0; -} - -void* ds_page_aligned_alloc(const int64_t size, const bool lock) -{ - void* ptr; - int retval; - - retval = posix_memalign(&ptr, (size_t)sysconf(_SC_PAGESIZE), size); - if (retval) { return nullptr; } - - if (lock == false) { return ptr; } - - auto mlock_ret = mlock(ptr, size); - if (mlock_ret != 0) { - auto mlock_error = errno; - std::cerr << "mlock failed to allocate " << size << " bytes with error no " << mlock_error - << " msg " << strerror(mlock_error) << std::endl; - free(ptr); - return nullptr; - } - - return ptr; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.h deleted file mode 100644 index 8742bf5bff543948032ce19de2fc4bea194a9ca9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.h +++ /dev/null @@ -1,82 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#pragma once - -#include -#include -#include - -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -struct io_xfer_ctxt { - const int _fd; - const int64_t _file_base_offset; - const int64_t _buffer_base_offset; - const void* _mem_buffer; - const int64_t _num_bytes; - - io_xfer_ctxt(const int fd, - const int64_t file_offset, - const int64_t buffer_offset, - const int64_t num_bytes, - const void* buffer); -}; - -struct io_prep_context { - const bool _read_op; - const std::unique_ptr& _xfer_ctxt; - const size_t _block_size; - const std::vector* _iocbs; - - io_prep_context(const bool read_op, - const std::unique_ptr& xfer_ctxt, - const size_t block_size, - const std::vector* iocbs); - - void prep_iocbs(const int n_iocbs, - const size_t num_bytes, - const void* start_buffer, - const int64_t start_offset); -}; - -struct io_prep_generator { - const bool _read_op; - const std::unique_ptr& _xfer_ctxt; - const size_t _block_size; - - int64_t _remaining_bytes; - int64_t _num_io_blocks; - int64_t _remaining_io_blocks; - int64_t _next_iocb_index; - - io_prep_generator(const bool read_op, - const std::unique_ptr& xfer_ctxt, - const size_t block_size); - - int prep_iocbs(const int n_iocbs, std::vector* iocbs); -}; - -void* ds_page_aligned_alloc(const int64_t size, const bool lock = false); - -int64_t get_file_size(const char* filename, int64_t& size); -int64_t get_fd_file_size(const int fd, int64_t& size); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp deleted file mode 100644 index 8387e667b33236e9ffcb7d070baec43cf498d20f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp +++ /dev/null @@ -1,39 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "deepspeed_aio_op_desc.h" - -using namespace std; - -io_op_desc_t::io_op_desc_t(const bool read_op, - const torch::Tensor& buffer, - const int fd, - const char* filename, - const int intra_op_parallelism, - const bool validate, - const int64_t file_offset) - : _read_op(read_op), - _buffer(buffer), - _fd(fd), - _filename((filename == nullptr) ? std::string() : filename), - _file_offset(file_offset), - _intra_op_parallelism(intra_op_parallelism), - _num_bytes_per_thread(static_cast(buffer.nbytes()) / intra_op_parallelism), - _validate(validate) -{ - if (validate) { assert(nullptr != filename); } -} - -char* io_op_desc_t::data_ptr() const { return (char*)_contiguous_buffer.data_ptr(); } - -void io_op_desc_t::finish() {} - -void io_op_desc_t::validate() {} - -void io_op_desc_t::run(const int tid, - std::unique_ptr& aio_ctxt, - deepspeed_aio_config_t* aio_config) -{ -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.h deleted file mode 100644 index cc7f15d7465830b32c273232c4e4353391c4a4d4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.h +++ /dev/null @@ -1,41 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#ifndef _IO_OP_DESC_T_ -#define _IO_OP_DESC_T_ -#include -#include -#include "deepspeed_py_aio.h" - -struct io_op_desc_t { - const bool _read_op; - torch::Tensor _buffer; - int _fd; - std::string _filename; - const int _intra_op_parallelism; - const int64_t _num_bytes_per_thread; - torch::Tensor _contiguous_buffer; - const bool _validate; - const int64_t _file_offset; - - io_op_desc_t(const bool read_op, - const torch::Tensor& buffer, - const int fd, - const char* filename, - const int intra_op_parallelism, - const bool validate, - const int64_t file_offset); - - virtual void run(const int tid, - std::unique_ptr& aio_ctxt, - deepspeed_aio_config_t* aio_config); - - virtual char* data_ptr() const; - - virtual void validate(); - - virtual void finish(); -}; -#endif // _IO_OP_DESC_T_ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.cpp deleted file mode 100644 index 30c3b49143979cd4e71f7e0898b0d454a8099a51..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.cpp +++ /dev/null @@ -1,51 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include "deepspeed_aio_thread.h" - -using namespace std; - -deepspeed_aio_thread_t::deepspeed_aio_thread_t(const int tid, deepspeed_aio_config_t& aio_config) - : _tid(tid), - _aio_config(aio_config), - _aio_ctxt(new aio_context(aio_config._block_size, aio_config._queue_depth)), - _time_to_exit(false) -{ -} - -deepspeed_aio_thread_t::~deepspeed_aio_thread_t() {} - -void deepspeed_aio_thread_t::run() -{ - while (true) { - std::shared_ptr next_io_op = nullptr; - - { - std::unique_lock lock(_work_sync._mutex); - _work_sync._cond_var.wait(lock, - [this] { return (!_work_queue.empty() || _time_to_exit); }); - if (!_work_queue.empty()) { - next_io_op = _work_queue.front(); - _work_queue.pop(); - } - } - - if (next_io_op) { - next_io_op->run(_tid, _aio_ctxt, &_aio_config); - - { - std::lock_guard lock(_complete_sync._mutex); - _complete_queue.push(next_io_op); - } - _complete_sync._cond_var.notify_one(); - } - - if (_time_to_exit) { break; } - } -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.h deleted file mode 100644 index a192804db13d8a9a1a160134ce5f90e0ca091cd7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.h +++ /dev/null @@ -1,38 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include -#include -#include -#include "deepspeed_cpu_op.h" - -struct thread_sync_t { - std::mutex _mutex; - std::condition_variable _cond_var; -}; - -struct deepspeed_aio_thread_t { - const int _tid; - deepspeed_aio_config_t& _aio_config; - - std::unique_ptr _aio_ctxt; - std::queue> _work_queue; - std::queue> _complete_queue; - - bool _time_to_exit; - - struct thread_sync_t _work_sync; - struct thread_sync_t _complete_sync; - - deepspeed_aio_thread_t(const int tid, deepspeed_aio_config_t& aio_config); - - ~deepspeed_aio_thread_t(); - - void run(); -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.cpp deleted file mode 100644 index 3ee1ecd3b4a8c9acd49c516fcb3eb69a088916d6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.cpp +++ /dev/null @@ -1,102 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "deepspeed_cpu_op.h" -#include "deepspeed_pin_tensor.h" - -using namespace std; - -cpu_op_desc_t::cpu_op_desc_t( - const std::unique_ptr& pinned_tensor_mgr, - const bool read_op, - const torch::Tensor& buffer, - const int fd, - const char* filename, - const int intra_op_parallelism, - const bool validate, - const int64_t file_offset) - : io_op_desc_t(read_op, buffer, fd, filename, intra_op_parallelism, validate, file_offset), - _cpu_buffer(buffer), - _pinned_tensor_mgr(pinned_tensor_mgr), - _is_managed_bounce_buffer(false) -{ - // Need to use CPU bounce buffer if buffer is not a page-locked DRAM memory. - _use_bounce_buffer = - !(_buffer.is_cpu() && (_buffer.is_pinned() || _pinned_tensor_mgr->is_managed(_buffer))); - if (_use_bounce_buffer) { - _alloc_bounce_buffer(); - if (!_read_op) { _cpu_buffer.copy_(_buffer); } - } - _contiguous_buffer = _cpu_buffer.contiguous(); -} - -char* cpu_op_desc_t::data_ptr() const { return (char*)_contiguous_buffer.data_ptr(); } - -void cpu_op_desc_t::finish() -{ - if (_use_bounce_buffer) { - if (_read_op) { - if (_buffer.is_cuda()) { - _buffer.copy_(_cpu_buffer.to(torch::Device(torch::kCUDA, _buffer.get_device()), - /*non_blocking=*/true)); - } - if (_buffer.is_xpu()) { _buffer.copy_(_cpu_buffer.to(torch::kXPU)); } - if (_buffer.is_cpu()) { _buffer.copy_(_cpu_buffer); } -#if defined(__ENABLE_CANN__) - if (torch_npu::utils::is_npu(_buffer)) { - auto device = at::Device("npu:0"); - _buffer.copy_(_cpu_buffer.to(device)); - } -#endif - } - - _free_bounce_buffer(); - } -} - -void cpu_op_desc_t::validate() -{ - const auto num_io_bytes = static_cast(_contiguous_buffer.nbytes()); - validate_aio_operation(_read_op, _filename.c_str(), data_ptr(), num_io_bytes); -} - -void cpu_op_desc_t::run(const int tid, - std::unique_ptr& aio_ctxt, - deepspeed_aio_config_t* aio_config) -{ - assert(tid < _intra_op_parallelism); - const auto buffer_base_offset = _num_bytes_per_thread * tid; - const auto file_base_offset = _file_offset + (_num_bytes_per_thread * tid); - - std::unique_ptr xfer_ctxt(new io_xfer_ctxt( - _fd, file_base_offset, buffer_base_offset, _num_bytes_per_thread, data_ptr())); - - if (aio_config->_overlap_events) { - do_aio_operation_overlap(_read_op, aio_ctxt, xfer_ctxt, aio_config, nullptr); - } else { - do_aio_operation_sequential(_read_op, aio_ctxt, xfer_ctxt, aio_config, nullptr); - } -} - -void cpu_op_desc_t::_alloc_bounce_buffer() -{ - auto options = torch::TensorOptions() - .dtype(_buffer.dtype()) - .layout(_buffer.layout()) - .device(torch::kCPU) - .requires_grad(false); - -#if defined(__CUDA_ARCH__) - _cpu_buffer = torch::empty(_buffer.numel(), options).pin_memory(); -#else - _is_managed_bounce_buffer = true; - _cpu_buffer = _pinned_tensor_mgr->alloc(_buffer.numel(), options); -#endif -} - -void cpu_op_desc_t::_free_bounce_buffer() -{ - if (_is_managed_bounce_buffer) { _pinned_tensor_mgr->free(_cpu_buffer); } -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.h deleted file mode 100644 index 7cc648bace8ed680e85c53652788c6d7349f293e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.h +++ /dev/null @@ -1,37 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include -#include "deepspeed_aio_op_desc.h" - -struct cpu_op_desc_t : io_op_desc_t { - torch::Tensor _cpu_buffer; - bool _use_bounce_buffer; - bool _is_managed_bounce_buffer; - const std::unique_ptr& _pinned_tensor_mgr; - - cpu_op_desc_t(const std::unique_ptr& pinned_tensor_mgr, - const bool read_op, - const torch::Tensor& buffer, - const int fd, - const char* filename, - const int intra_op_parallelism, - const bool validate, - const int64_t file_offset); - - void run(const int tid, - std::unique_ptr& aio_ctxt, - deepspeed_aio_config_t* aio_config); - - char* data_ptr() const; - - void validate(); - - void finish(); - - void _alloc_bounce_buffer(); - void _free_bounce_buffer(); -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.cpp deleted file mode 100644 index a97a4ac18ba811dfaad713fdee15248caecf035b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.cpp +++ /dev/null @@ -1,61 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for managing CPU tensors occupying page-locked memory. -*/ - -#include "deepspeed_pin_tensor.h" - -using namespace std; - -deepspeed_pin_tensor_t::~deepspeed_pin_tensor_t() -{ - for (auto iter = _locked_tensors.begin(); iter != _locked_tensors.end(); ++iter) { - munlock(iter->first, iter->second); - std::free((void*)iter->first); - } - _locked_tensors.clear(); -} - -torch::Tensor deepspeed_pin_tensor_t::alloc(const int64_t num_elem, - const torch::TensorOptions& options) -{ - const auto scalar_dtype = torch::typeMetaToScalarType(options.dtype()); - const auto num_bytes = num_elem * torch::elementSize(scalar_dtype); - auto pinned_buffer = ds_page_aligned_alloc(num_bytes, true); - assert(nullptr != pinned_buffer); - - _locked_tensors[pinned_buffer] = num_bytes; - - return at::from_blob(pinned_buffer, static_cast(num_elem), options); -} - -torch::Tensor deepspeed_pin_tensor_t::alloc(const int64_t num_elem, const at::ScalarType& elem_type) -{ - auto options = torch::TensorOptions().dtype(elem_type).device(torch::kCPU).requires_grad(false); - return alloc(num_elem, options); -} - -bool deepspeed_pin_tensor_t::free(torch::Tensor& locked_tensor) -{ - auto addr = locked_tensor.data_ptr(); - if (_locked_tensors.find(addr) != _locked_tensors.end()) { - munlock(addr, _locked_tensors[addr]); - std::free(addr); - _locked_tensors.erase(addr); - return true; - } - - return false; -} - -bool deepspeed_pin_tensor_t::is_managed(const torch::Tensor& buffer) -{ - if (!buffer.is_cpu()) { return false; } - auto addr = buffer.data_ptr(); - if (_locked_tensors.find(addr) != _locked_tensors.end()) { return true; } - return false; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.h deleted file mode 100644 index 4b8ad7e76085877818aa26e511968efb0dc31fd2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.h +++ /dev/null @@ -1,30 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for managing CPU tensors occupying page-locked memory. -TODO: Implement a full-featured manager that -1. Avoid page-locked memory leaks -2. Minimize page-locked memory usage by reducing internal fragmentation -Functionality for managing CPU tensors occupying page-locked memory. -*/ - -#include -#include "deepspeed_py_aio.h" - -struct deepspeed_pin_tensor_t { - std::map _locked_tensors; - - deepspeed_pin_tensor_t() = default; - - ~deepspeed_pin_tensor_t(); - - torch::Tensor alloc(const int64_t num_elem, const at::ScalarType& elem_type); - torch::Tensor alloc(const int64_t num_elem, const torch::TensorOptions& options); - - bool free(torch::Tensor& locked_tensor); - - bool is_managed(const torch::Tensor& buffer); -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.cpp deleted file mode 100644 index 1ff0397043fae2d904fb22d5d0b9eea5af6b6cd4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.cpp +++ /dev/null @@ -1,123 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include -#include -#include - -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "deepspeed_py_aio.h" - -using namespace std; -using namespace std::chrono; - -#define DEBUG_DS_AIO_READ 0 -#define DEBUG_DS_AIO_WRITE 0 - -static const std::string c_library_name = "deepspeed_aio"; - -int deepspeed_py_aio_write(const torch::Tensor& buffer, - const char* filename, - const int block_size, - const int queue_depth, - const bool single_submit, - const bool overlap_events, - const bool validate) -{ - const auto start_time = std::chrono::high_resolution_clock::now(); - deepspeed_aio_config_t config(block_size, queue_depth, single_submit, overlap_events, false); - - const auto fd = open_file(filename, false); - if (fd == -1) { return -1; } - - auto write_buffer = (char*)buffer.data_ptr(); - const auto num_write_bytes = static_cast(buffer.nbytes()); - - std::unique_ptr xfer_ctxt( - new io_xfer_ctxt(fd, 0, 0, num_write_bytes, write_buffer)); - std::unique_ptr aio_ctxt(new aio_context(config._block_size, config._queue_depth)); - - if (config._overlap_events) { - do_aio_operation_overlap(false, aio_ctxt, xfer_ctxt, &config, nullptr); - } else { - do_aio_operation_sequential(false, aio_ctxt, xfer_ctxt, &config, nullptr); - } - const std::chrono::duration aio_time = - std::chrono::high_resolution_clock::now() - start_time; - - close(fd); - - if (validate) { validate_aio_operation(false, filename, write_buffer, num_write_bytes); } - - const std::chrono::duration fn_time = - std::chrono::high_resolution_clock::now() - start_time; - std::cout << "Elapsed time(usec): " << "aio = " << aio_time.count() * 1e6 - << " call = " << fn_time.count() * 1e6 << std::endl; - return 0; -} - -int deepspeed_py_aio_read(torch::Tensor& buffer, - const char* filename, - const int block_size, - const int queue_depth, - const bool single_submit, - const bool overlap_events, - const bool validate) -{ - const auto start_time = std::chrono::high_resolution_clock::now(); - int64_t num_file_bytes; - if (-1 == get_file_size(filename, num_file_bytes)) { - const auto error_code = errno; - report_file_error(filename, " fstat for read", error_code); - return -1; - } - - deepspeed_aio_config_t config(block_size, queue_depth, single_submit, overlap_events, false); - const auto fd = open_file(filename, true); - if (fd == -1) { return -1; } - - auto read_buffer = (char*)buffer.data_ptr(); - assert(static_cast(buffer.nbytes()) == num_file_bytes); - - std::unique_ptr xfer_ctxt( - new io_xfer_ctxt(fd, 0, 0, num_file_bytes, read_buffer)); - std::unique_ptr aio_ctxt(new aio_context(config._block_size, config._queue_depth)); - - if (config._overlap_events) { - do_aio_operation_overlap(true, aio_ctxt, xfer_ctxt, &config, nullptr); - } else { - do_aio_operation_sequential(true, aio_ctxt, xfer_ctxt, &config, nullptr); - } - const std::chrono::duration aio_time = - std::chrono::high_resolution_clock::now() - start_time; - - close(fd); - - if (validate) { validate_aio_operation(true, filename, read_buffer, num_file_bytes); } - - const std::chrono::duration fn_time = - std::chrono::high_resolution_clock::now() - start_time; - std::cout << "Elapsed time(usec): " << "aio = " << aio_time.count() * 1e6 - << " call = " << fn_time.count() * 1e6 << std::endl; - return 0; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.h deleted file mode 100644 index ba794db5440d540967053fc2ae4b184cfb8825d8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.h +++ /dev/null @@ -1,28 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping tensors to/from (NVMe) storage devices. -*/ - -#include -#include -#include - -int deepspeed_py_aio_write(const torch::Tensor& buffer, - const char* filename, - const int block_size, - const int queue_depth, - const bool single_submit, - const bool overlap_events, - const bool validate); - -int deepspeed_py_aio_read(torch::Tensor& buffer, - const char* filename, - const int block_size, - const int queue_depth, - const bool single_submit, - const bool overlap_events, - const bool validate); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp deleted file mode 100644 index 2b1093e9928659bdaac1c14323928c644061155e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp +++ /dev/null @@ -1,28 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include "deepspeed_py_aio_handle.h" -#include - -using namespace std; - -deepspeed_aio_handle_t::deepspeed_aio_handle_t(const int block_size, - const int queue_depth, - const bool single_submit, - const bool overlap_events, - const int intra_op_parallelism) - : deepspeed_io_handle_t(block_size, - queue_depth, - single_submit, - overlap_events, - intra_op_parallelism) -{ -} - -deepspeed_aio_handle_t::~deepspeed_aio_handle_t() {} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.h deleted file mode 100644 index c9fcb6d2b46227251ec56a286c9144935daffe0b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.h +++ /dev/null @@ -1,21 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ -#include -#include -#include "deepspeed_py_io_handle.h" - -struct deepspeed_aio_handle_t : deepspeed_io_handle_t { - deepspeed_aio_handle_t(const int block_size, - const int queue_depth, - const bool single_submit, - const bool overlap_events, - const int intra_op_parallelism); - - ~deepspeed_aio_handle_t(); -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.cpp deleted file mode 100644 index f5480e9d9d836db587cc9ded27c8f437bc294452..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.cpp +++ /dev/null @@ -1,135 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping tensors to/from (NVMe) storage devices. -*/ - -#include "deepspeed_py_copy.h" -#include - -#define ROUND_DOWN(size, step) ((size) & ~((step) - 1)) - -#if defined(__AVX512__) or defined(__AVX256__) -union AVX_Data { -#if defined(__AVX512__) - __m512 data; -#else - __m256 data; -#endif -}; -#endif - -static void helper_memcpy_1(float* dest, float* src, size_t param_size) -{ - size_t rounded_size = 0; - -#if defined(__AVX512__) or defined(__AVX256__) - - rounded_size = ROUND_DOWN(param_size, SIMD_WIDTH); - - for (size_t t = 0; t < rounded_size; t += TILE) { - size_t copy_size = TILE; - if ((t + TILE) > rounded_size) copy_size = rounded_size - t; - size_t offset = copy_size + t; -#pragma omp parallel for - for (size_t i = t; i < offset; i += SIMD_WIDTH) { - AVX_Data src_4; - src_4.data = SIMD_LOAD(src + i); - - SIMD_STORE(dest + i, src_4.data); - } - } - -#endif - - if (param_size > rounded_size) { -#pragma omp parallel for - for (size_t k = rounded_size; k < param_size; k++) { dest[k] = src[k]; } - } -} - -static void helper_memcpy_4(float* dest, float* src, size_t param_size) -{ - size_t rounded_size = 0; - -#if defined(__AVX512__) or defined(__AVX256__) - - rounded_size = ROUND_DOWN(param_size, (SIMD_WIDTH << 2)); - - for (size_t t = 0; t < rounded_size; t += TILE) { - size_t copy_size = TILE; - if ((t + TILE) > rounded_size) copy_size = rounded_size - t; - size_t offset = copy_size + t; -#pragma omp parallel for - for (size_t i = t; i < offset; i += (SIMD_WIDTH << 2)) { - AVX_Data src_4[4]; - src_4[0].data = SIMD_LOAD(src + i); - src_4[1].data = SIMD_LOAD(src + i + SIMD_WIDTH); - src_4[2].data = SIMD_LOAD(src + i + (SIMD_WIDTH << 1)); - src_4[3].data = SIMD_LOAD(src + i + SIMD_WIDTH * 3); - - SIMD_STORE(dest + i, src_4[0].data); - SIMD_STORE(dest + i + SIMD_WIDTH, src_4[1].data); - SIMD_STORE(dest + i + (SIMD_WIDTH << 1), src_4[2].data); - SIMD_STORE(dest + i + SIMD_WIDTH * 3, src_4[3].data); - } - } -#endif - if (param_size > rounded_size) - helper_memcpy_1((dest + rounded_size), (src + rounded_size), (param_size - rounded_size)); -} - -static void helper_mempcy_8(float* dest, float* src, size_t param_size) -{ - size_t rounded_size = 0; - -#if defined(__AVX512__) or defined(__AVX256__) - - rounded_size = ROUND_DOWN(param_size, (SIMD_WIDTH << 2)); - - for (size_t t = 0; t < rounded_size; t += TILE) { - size_t copy_size = TILE; - if ((t + TILE) > rounded_size) copy_size = rounded_size - t; - size_t offset = copy_size + t; -#pragma omp parallel for - for (size_t i = t; i < offset; i += (SIMD_WIDTH << 3)) { - AVX_Data src_4[8]; - src_4[0].data = SIMD_LOAD(src + i); - src_4[1].data = SIMD_LOAD(src + i + SIMD_WIDTH); - src_4[2].data = SIMD_LOAD(src + i + (SIMD_WIDTH << 1)); - src_4[3].data = SIMD_LOAD(src + i + SIMD_WIDTH * 3); - src_4[4].data = SIMD_LOAD(src + i + (SIMD_WIDTH << 2)); - src_4[5].data = SIMD_LOAD(src + i + SIMD_WIDTH * 5); - src_4[6].data = SIMD_LOAD(src + i + SIMD_WIDTH * 6); - src_4[7].data = SIMD_LOAD(src + i + SIMD_WIDTH * 7); - - SIMD_STORE(dest + i, src_4[0].data); - SIMD_STORE(dest + i + SIMD_WIDTH, src_4[1].data); - SIMD_STORE(dest + i + (SIMD_WIDTH << 1), src_4[2].data); - SIMD_STORE(dest + i + SIMD_WIDTH * 3, src_4[3].data); - SIMD_STORE(dest + i + (SIMD_WIDTH << 2), src_4[4].data); - SIMD_STORE(dest + i + SIMD_WIDTH * 5, src_4[5].data); - SIMD_STORE(dest + i + SIMD_WIDTH * 6, src_4[6].data); - SIMD_STORE(dest + i + SIMD_WIDTH * 7, src_4[7].data); - } - } -#endif - if (param_size > rounded_size) - helper_memcpy_4((dest + rounded_size), (src + rounded_size), (param_size - rounded_size)); -} - -int deepspeed_py_memcpy(torch::Tensor& dest, const torch::Tensor& src) -{ - auto dest_c = dest.contiguous(); - auto src_c = src.contiguous(); - - float* dest_ptr = (float*)dest_c.data_ptr(); - float* src_ptr = (float*)src_c.data_ptr(); - - helper_mempcy_8(dest_ptr, src_ptr, dest_c.size(0)); - - return 0; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.h deleted file mode 100644 index f443571a3e7b35f3bc70ec835152c34e5d979137..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.h +++ /dev/null @@ -1,43 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#if (__x86_64__ || __i386__) -#include -#include -#endif - -#include -#include -#include - -#define TILE (1024 * 1024 * 1024) - -#if defined(__AVX512__) -#define SIMD_STORE(a, d) _mm512_storeu_ps(a, d) -#define SIMD_LOAD(x) _mm512_loadu_ps(x) -#define SIMD_SET(x) _mm512_set1_ps(x) -#define SIMD_MUL(x, y) _mm512_mul_ps(x, y) -#define SIMD_FMA(x, y, c) _mm512_fmadd_ps(x, y, c) -#define SIMD_SQRT(x) _mm512_sqrt_ps(x) -#define SIMD_DIV(x, y) _mm512_div_ps(x, y) -#define SIMD_WIDTH 16 -#else -#if defined(__AVX256__) -#define SIMD_STORE(a, d) _mm256_storeu_ps(a, d) -#define SIMD_LOAD(x) _mm256_loadu_ps(x) -#define SIMD_SET(x) _mm256_set1_ps(x) -#define SIMD_MUL(x, y) _mm256_mul_ps(x, y) -#define SIMD_FMA(x, y, c) _mm256_fmadd_ps(x, y, c) -#define SIMD_SQRT(x) _mm256_sqrt_ps(x) -#define SIMD_DIV(x, y) _mm256_div_ps(x, y) -#define SIMD_WIDTH 8 -#endif -#endif - -int deepspeed_py_memcpy(torch::Tensor& dest, const torch::Tensor& src); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.cpp deleted file mode 100644 index 48668a842949a718a2e4ab1fdac2b33026568433..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.cpp +++ /dev/null @@ -1,366 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include "deepspeed_py_io_handle.h" -#include - -#define O_DIRECT_ALIGNMENT 512 - -using namespace std; - -static void _start_aio_thread(std::shared_ptr ctxt) { ctxt->run(); } - -static bool is_valid_bytes_to_read(const char* filename, - const int64_t file_offset, - const int64_t num_bytes_to_read) -{ - int64_t num_file_bytes; - if (-1 == get_file_size(filename, num_file_bytes)) { - const auto error_code = errno; - report_file_error(filename, " fstat for read", error_code); - return false; - } - if ((file_offset + num_bytes_to_read) > num_file_bytes) { - std::cout << filename << ": file_offset + buffer nbytes > file bytes " - << (file_offset + num_bytes_to_read) << " > " << num_file_bytes << std::endl; - } - assert((file_offset + num_bytes_to_read) <= num_file_bytes); - return true; -} - -deepspeed_io_handle_t::deepspeed_io_handle_t(const int block_size, - const int queue_depth, - const bool single_submit, - const bool overlap_events, - const int intra_op_parallelism) - : _aio_ctxt(new aio_context(block_size, queue_depth)), - _single_submit(single_submit), - _overlap_events(overlap_events), - _intra_op_parallelism(intra_op_parallelism), - _aio_config(block_size, queue_depth, single_submit, overlap_events, false), - _num_pending_ops(0), - _pinned_tensor_mgr(new deepspeed_pin_tensor_t()) -{ - for (auto i = 0; i < intra_op_parallelism; ++i) { - _thread_contexts.push_back(std::make_shared(i, _aio_config)); - } - - for (auto& ctxt : _thread_contexts) { - _threads.push_back(std::thread(_start_aio_thread, ctxt)); - } -} - -deepspeed_io_handle_t::~deepspeed_io_handle_t() -{ - _stop_threads(); - for (auto& thr : _threads) { thr.join(); } -} - -const int deepspeed_io_handle_t::get_block_size() const -{ - return _aio_ctxt ? _aio_ctxt->_block_size : -1; -} - -const int deepspeed_io_handle_t::get_queue_depth() const -{ - return _aio_ctxt ? _aio_ctxt->_queue_depth : -1; -} - -const bool deepspeed_io_handle_t::get_single_submit() const { return _single_submit; } - -const bool deepspeed_io_handle_t::get_overlap_events() const { return _overlap_events; } - -const int deepspeed_io_handle_t::get_intra_op_parallelism() const { return _intra_op_parallelism; } - -const int deepspeed_io_handle_t::get_alignment() const -{ - return _intra_op_parallelism * O_DIRECT_ALIGNMENT; -} - -int deepspeed_io_handle_t::read(torch::Tensor& buffer, - const char* filename, - const bool validate, - const int64_t file_offset) -{ - const auto start_time = std::chrono::high_resolution_clock::now(); - - assert(_aio_ctxt); - - int64_t num_file_bytes; - if (-1 == get_file_size(filename, num_file_bytes)) { - const auto error_code = errno; - report_file_error(filename, " fstat for read", error_code); - return -1; - } - assert(static_cast(buffer.nbytes()) == num_file_bytes); - - const auto fd = open_file(filename, true); - if (fd == -1) { return -1; } - - auto read_buffer = (char*)buffer.data_ptr(); - std::unique_ptr xfer_ctxt( - new io_xfer_ctxt(fd, file_offset, 0, num_file_bytes, read_buffer)); - - if (_aio_config._overlap_events) { - do_aio_operation_overlap(true, _aio_ctxt, xfer_ctxt, &_aio_config, nullptr); - } else { - do_aio_operation_sequential(true, _aio_ctxt, xfer_ctxt, &_aio_config, nullptr); - } - - close(fd); - const std::chrono::duration aio_time = - std::chrono::high_resolution_clock::now() - start_time; - - if (validate) { validate_aio_operation(true, filename, read_buffer, num_file_bytes); } - const std::chrono::duration fn_time = - std::chrono::high_resolution_clock::now() - start_time; - std::cout << "Elapsed time(usec): " << "aio = " << aio_time.count() * 1e6 - << " call = " << fn_time.count() * 1e6 << std::endl; - return 0; -} - -int deepspeed_io_handle_t::write(const torch::Tensor& buffer, - const char* filename, - const bool validate, - const int64_t file_offset) -{ - assert(_aio_ctxt); - - const auto start_time = std::chrono::high_resolution_clock::now(); - - const auto fd = open_file(filename, false); - if (fd == -1) { return -1; } - - auto write_buffer = (char*)buffer.data_ptr(); - const auto num_write_bytes = static_cast(buffer.nbytes()); - std::unique_ptr xfer_ctxt( - new io_xfer_ctxt(fd, file_offset, 0, num_write_bytes, write_buffer)); - - if (_aio_config._overlap_events) { - do_aio_operation_overlap(false, _aio_ctxt, xfer_ctxt, &_aio_config, nullptr); - } else { - do_aio_operation_sequential(false, _aio_ctxt, xfer_ctxt, &_aio_config, nullptr); - } - const std::chrono::duration aio_time = - std::chrono::high_resolution_clock::now() - start_time; - - close(fd); - - if (validate) { validate_aio_operation(false, filename, write_buffer, num_write_bytes); } - - const std::chrono::duration fn_time = - std::chrono::high_resolution_clock::now() - start_time; - std::cout << "Elapsed time(usec): " << "aio = " << aio_time.count() * 1e6 - << " call = " << fn_time.count() * 1e6 << std::endl; - return 0; -} - -void deepspeed_io_handle_t::_schedule_aio_work(std::shared_ptr scheduled_op) -{ - for (auto& ctxt : _thread_contexts) { - { - std::lock_guard lock(ctxt->_work_sync._mutex); - ctxt->_work_queue.push(scheduled_op); - } - ctxt->_work_sync._cond_var.notify_one(); - } - _num_pending_ops++; -} - -std::shared_ptr deepspeed_io_handle_t::_wait_for_aio_work() -{ - std::shared_ptr completed_op = nullptr; - for (auto& ctxt : _thread_contexts) { - std::unique_lock lock(ctxt->_complete_sync._mutex); - ctxt->_complete_sync._cond_var.wait(lock, - [ctxt] { return !ctxt->_complete_queue.empty(); }); - completed_op = ctxt->_complete_queue.front(); - ctxt->_complete_queue.pop(); - } - return completed_op; -} - -void deepspeed_io_handle_t::_stop_threads() -{ - assert(0 == _num_pending_ops); - for (auto& ctxt : _thread_contexts) { - { - std::lock_guard lock(ctxt->_work_sync._mutex); - ctxt->_time_to_exit = true; - } - ctxt->_work_sync._cond_var.notify_one(); - } -} - -int deepspeed_io_handle_t::wait() -{ - assert(_num_pending_ops > 0); - auto num_completed_ops = 0; - - while (_num_pending_ops > 0) { - auto completed_op = _wait_for_aio_work(); - - if (completed_op->_validate) { completed_op->validate(); } - - completed_op->finish(); - - if (!completed_op->_filename.empty()) { (completed_op->_fd); } - - --_num_pending_ops; - ++num_completed_ops; - } - - return num_completed_ops; -} - -bool deepspeed_io_handle_t::_is_valid_parallel_aio_op(const bool read_op, const int64_t num_bytes) -{ - const auto op_string = read_op ? "Read" : "Write"; - if (num_bytes % get_intra_op_parallelism()) { - std::cout << "deepspeed_aio failure: parallel " << op_string << " num_bytes = " << num_bytes - << " not divisible by intra op parallelism = " << get_intra_op_parallelism() - << std::endl; - return false; - } - - return true; -} - -std::shared_ptr deepspeed_io_handle_t::_create_io_op_desc( - const bool read_op, - const torch::Tensor& buffer, - const int fd, - const char* filename, - const bool validate, - const int64_t file_offset) -{ - return std::make_shared(_pinned_tensor_mgr, - read_op, - buffer, - fd, - filename, - _intra_op_parallelism, - validate, - file_offset); -} - -int deepspeed_io_handle_t::_pread(const torch::Tensor& buffer, - const int fd, - const char* filename, - const bool validate, - const bool async, - const int64_t file_offset) -{ - auto scheduled_op = _create_io_op_desc(true, buffer, fd, filename, validate, file_offset); - - _schedule_aio_work(scheduled_op); - - if (async) { return 0; } - - return wait(); -} - -int deepspeed_io_handle_t::pread(const torch::Tensor& buffer, - const char* filename, - const bool validate, - const bool async, - const int64_t file_offset) -{ - const auto buffer_bytes = static_cast(buffer.nbytes()); - - if (!is_valid_bytes_to_read(filename, file_offset, buffer_bytes)) { return -1; } - - if (!_is_valid_parallel_aio_op(true, buffer_bytes)) { return -1; } - - const auto fd = open_file(filename, true); - if (fd == -1) { return -1; } - - return _pread(buffer, fd, filename, validate, async, file_offset); -} - -int deepspeed_io_handle_t::_pwrite(const torch::Tensor& buffer, - const int fd, - const char* filename, - const bool validate, - const bool async, - const int64_t file_offset) -{ - auto scheduled_op = _create_io_op_desc(false, buffer, fd, filename, validate, file_offset); - - _schedule_aio_work(scheduled_op); - - if (async) { return 0; } - - return wait(); -} - -int deepspeed_io_handle_t::pwrite(const torch::Tensor& buffer, - const char* filename, - const bool validate, - const bool async, - const int64_t file_offset) -{ - const auto num_write_bytes = static_cast(buffer.nbytes()); - - if (!_is_valid_parallel_aio_op(false, num_write_bytes)) { return -1; } - - const auto fd = open_file(filename, false); - if (fd == -1) { return -1; } - - return _pwrite(buffer, fd, filename, validate, async, file_offset); -} - -int deepspeed_io_handle_t::sync_pread(torch::Tensor& buffer, - const char* filename, - const int64_t file_offset) -{ - return pread(buffer, filename, false, false, file_offset); -} - -int deepspeed_io_handle_t::sync_pwrite(const torch::Tensor& buffer, - const char* filename, - const int64_t file_offset) -{ - return pwrite(buffer, filename, false, false, file_offset); -} - -int deepspeed_io_handle_t::async_pread(torch::Tensor& buffer, - const char* filename, - const int64_t file_offset) -{ - return pread(buffer, filename, false, true, file_offset); -} - -int deepspeed_io_handle_t::async_pwrite(const torch::Tensor& buffer, - const char* filename, - const int64_t file_offset) -{ - return pwrite(buffer, filename, false, true, file_offset); -} - -int deepspeed_io_handle_t::async_pwrite(const torch::Tensor& buffer, - const int fd, - const int64_t file_offset = 0) -{ - const auto num_write_bytes = static_cast(buffer.nbytes()); - if (!_is_valid_parallel_aio_op(false, num_write_bytes)) { return -1; } - - return _pwrite(buffer, fd, nullptr, false, true, file_offset); -} - -at::Tensor deepspeed_io_handle_t::new_cpu_locked_tensor(const int64_t num_elem, - const torch::Tensor& example_tensor) -{ - return _pinned_tensor_mgr->alloc(num_elem, example_tensor.scalar_type()); -} - -bool deepspeed_io_handle_t::free_cpu_locked_tensor(torch::Tensor& locked_tensor) -{ - return _pinned_tensor_mgr->free(locked_tensor); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.h deleted file mode 100644 index 8cb43c5b38e5d8f1981176c6459db0135c5f896d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.h +++ /dev/null @@ -1,109 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include -#include -#include "deepspeed_aio_thread.h" -#include "deepspeed_pin_tensor.h" - -struct deepspeed_io_handle_t { - std::unique_ptr _aio_ctxt; - const bool _single_submit; - const bool _overlap_events; - const int _intra_op_parallelism; - deepspeed_aio_config_t _aio_config; - - std::vector> _thread_contexts; - std::vector _threads; - int _num_pending_ops; - std::unique_ptr _pinned_tensor_mgr; - - deepspeed_io_handle_t(const int block_size, - const int queue_depth, - const bool single_submit, - const bool overlap_events, - const int intra_op_parallelism); - - virtual ~deepspeed_io_handle_t() = 0; - - const int get_block_size() const; - const int get_queue_depth() const; - const bool get_single_submit() const; - const bool get_overlap_events() const; - const int get_intra_op_parallelism() const; - const int get_alignment() const; - - int read(torch::Tensor& buffer, - const char* filename, - const bool validate, - const int64_t file_offset); - - int write(const torch::Tensor& buffer, - const char* filename, - const bool validate, - const int64_t file_offset); - - int pread(const torch::Tensor& buffer, - const char* filename, - const bool validate, - const bool async, - const int64_t file_offset); - - int pwrite(const torch::Tensor& buffer, - const char* filename, - const bool validate, - const bool async, - const int64_t file_offset); - - int sync_pread(torch::Tensor& buffer, const char* filename, const int64_t file_offset); - - int sync_pwrite(const torch::Tensor& buffer, const char* filename, const int64_t file_offset); - - int async_pread(torch::Tensor& buffer, const char* filename, const int64_t file_offset); - - int async_pwrite(const torch::Tensor& buffer, const char* filename, const int64_t file_offset); - int async_pwrite(const torch::Tensor& buffer, const int fd, const int64_t file_offset); - - // TODO: Make API's args to be shape and dtype. - torch::Tensor new_cpu_locked_tensor(const int64_t num_elem, - const torch::Tensor& example_tensor); - - bool free_cpu_locked_tensor(torch::Tensor&); - - int wait(); - - void _stop_threads(); - - void _schedule_aio_work(std::shared_ptr scheduled_op); - - std::shared_ptr _wait_for_aio_work(); - - bool _is_valid_parallel_aio_op(const bool read_op, const int64_t num_bytes); - - int _pread(const torch::Tensor& buffer, - const int fd, - const char* filename, - const bool validate, - const bool async, - const int64_t file_offset); - - int _pwrite(const torch::Tensor& buffer, - const int fd, - const char* filename, - const bool validate, - const bool async, - const int64_t file_offset); - - virtual std::shared_ptr _create_io_op_desc(const bool read_op, - const torch::Tensor& buffer, - const int fd, - const char* filename, - const bool validate, - const int64_t file_offset); -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/py_ds_aio.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/py_ds_aio.cpp deleted file mode 100644 index 62500bf4a6e997d4444dc8f6dfdfdc10173199f5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_lib/py_ds_aio.cpp +++ /dev/null @@ -1,128 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ -#include -#include "deepspeed_py_aio_handle.h" -#include "deepspeed_py_copy.h" -using namespace pybind11::literals; - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("aio_read", &deepspeed_py_aio_read, "DeepSpeed Asynchronous I/O Read"); - - m.def("aio_write", &deepspeed_py_aio_write, "DeepSpeed Asynchronous I/O Write"); - - m.def("deepspeed_memcpy", &deepspeed_py_memcpy, "DeepSpeed Memory Copy"); - - py::class_(m, "aio_handle") - .def(py::init(), - "AIO handle constructor", - "block_size"_a = 1024 * 1024, - "queue_depth"_a = 128, - "single_submit"_a = false, - "overlap_events"_a = false, - "intra_op_parallelism"_a = 1) - - .def("get_block_size", &deepspeed_aio_handle_t::get_block_size) - .def("get_queue_depth", &deepspeed_aio_handle_t::get_queue_depth) - .def("get_single_submit", &deepspeed_aio_handle_t::get_single_submit) - .def("get_overlap_events", &deepspeed_aio_handle_t::get_overlap_events) - .def("get_intra_op_parallelism", &deepspeed_aio_handle_t::get_intra_op_parallelism) - .def("get_alignment", &deepspeed_aio_handle_t::get_alignment) - - .def("read", - &deepspeed_aio_handle_t::read, - "Synchronous and non-parallel file read. Returns count of completed read ops", - "buffer"_a, - "filename"_a, - "validate"_a, - "file_offset"_a = 0) - - .def("write", - &deepspeed_aio_handle_t::write, - "Synchronous and non-parallel file write. Returns count of completed write ops", - "buffer"_a, - "filename"_a, - "validate"_a, - "file_offset"_a = 0) - - .def("pread", - &deepspeed_aio_handle_t::pread, - "Parallel file read with option of asynchronous completion. If synchronous, returns " - "count of completed read ops", - "buffer"_a, - "filename"_a, - "validate"_a, - "async"_a, - "file_offset"_a = 0) - - .def("pwrite", - &deepspeed_aio_handle_t::pwrite, - "Parallel file write with option of asynchronous completion. If synchronous, returns " - "count of completed write ops", - "buffer"_a, - "filename"_a, - "validate"_a, - "async"_a, - "file_offset"_a = 0) - - .def("sync_pread", - &deepspeed_aio_handle_t::sync_pread, - "Synchronous parallel file read. Returns count of completed read ops", - "buffer"_a, - "filename"_a, - "file_offset"_a = 0) - - .def("sync_pwrite", - &deepspeed_aio_handle_t::sync_pwrite, - "Synchronous parallel file write. Returns count of completed write ops", - "buffer"_a, - "filename"_a, - "file_offset"_a = 0) - - .def("async_pread", - &deepspeed_aio_handle_t::async_pread, - "Asynchronous parallel file read. Returns 0 on success. Returns 0 on success, and " - "subsequent wait() returns count of completed ops.", - "buffer"_a, - "filename"_a, - "file_offset"_a = 0) - - .def( - "async_pwrite", - py::overload_cast( - &deepspeed_aio_handle_t::async_pwrite), - "Asynchronous parallel file write. Returns 0 on success, and subsequent wait() returns " - "count of completed ops.", - "buffer"_a, - "filename"_a, - "file_offset"_a = 0) - - .def("async_pwrite", - py::overload_cast( - &deepspeed_aio_handle_t::async_pwrite), - "Asynchronous parallel file write using opened python file object.", - "buffer"_a, - "fd"_a, - "file_offset"_a = 0) - - .def("new_cpu_locked_tensor", - &deepspeed_aio_handle_t::new_cpu_locked_tensor, - "Allocate pinned CPU tensor.", - "num_elem"_a, - "example_tenosr"_a) - - .def("free_cpu_locked_tensor", - &deepspeed_aio_handle_t::free_cpu_locked_tensor, - "Free pinned CPU tensor.", - "tensor"_a) - - .def("wait", - &deepspeed_aio_handle_t::wait, - "Wait for (ongoing) asynchronous operations to complete"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_test/single_process_config.json b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_test/single_process_config.json deleted file mode 100644 index 4a224711b5ff4330ca704af0aaa9a4a4b0ea8ca3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/aio/py_test/single_process_config.json +++ /dev/null @@ -1,34 +0,0 @@ -{ - "block_size": [ - "128K", - "256K", - "1M", - "2M", - "4M", - "8M", - "16M" - ], - "queue_depth": [ - 8, - 16, - 32, - 64 - ], - "io_parallel": [ - 1, - 2, - 4, - 8 - ], - "single_submit": [ - true, - false - ], - "sequential_requests": [ - true, - false - ], - "threads": [ - 1 - ] -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/deepcompile.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/deepcompile.cpp deleted file mode 100644 index 2eca0a33262e69cf8901c8293d6c7859b2619149..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/deepcompile.cpp +++ /dev/null @@ -1,188 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "deepcompile.h" - -#define USE_C10D_NCCL - -namespace dc { - -std::shared_ptr param_registry; -std::unordered_map> executors; -std::shared_ptr reduce_buckets = nullptr; - -c10::intrusive_ptr process_group = nullptr; -c10::intrusive_ptr symm_mem = nullptr; -ncclComm_t nccl_comm; -bool use_symm_mem; -bool clone_custom_op_output; -bool profile = false; -bool pre_div_reduce = true; - -bool sync_before_reduce; // for debugging -bool sync_after_reduce; // for debugging -bool sync_before_allgather; // for debugging -bool sync_after_allgather; // for debugging - -std::vector sizes_to_int_vector(at::IntArrayRef sizes) -{ - std::vector result; - for (int i = 0; i < sizes.size(); i++) { result.push_back(sizes[i]); } - return result; -} - -void enable_profiling(bool enable) { profile = enable; } - -bool is_profiling() { return profile; } - -c10::intrusive_ptr getSymmMemWorkspace(int64_t size) -{ - c10::Device device = c10::Device(c10::kCUDA, c10::cuda::current_device()); - std::vector sizes = {size}; - std::vector strides = {1}; - at::Tensor sym_mem_ws = c10d::symmetric_memory::empty_strided_p2p( - {size}, {1}, c10::ScalarType::Byte, device, process_group->getGroupName(), std::nullopt); - return c10d::symmetric_memory::rendezvous(sym_mem_ws); -} - -void lazy_init_symm_memory() -{ - if (use_symm_mem && !symm_mem) { - int64_t max_param_size = 0; - for (const auto& it : param_registry->getParams()) { - int64_t size = it.second.getDSTensor().numel() * it.second.getDSTensor().element_size(); - if (size > max_param_size) { max_param_size = size; } - } - symm_mem = getSymmMemWorkspace(max_param_size); - } -} - -ncclDataType_t get_nccl_data_type(at::ScalarType scalar_type) -{ - switch (scalar_type) { - case at::kFloat: return ncclFloat; - case at::kHalf: return ncclHalf; - case at::kDouble: return ncclDouble; - case at::kBFloat16: return ncclBfloat16; - case at::kLong: return ncclInt64; - case at::kInt: return ncclInt; - case at::kChar: return ncclInt8; - default: throw std::runtime_error("Unsupported scalar type"); - } -} - -void reset() -{ - executors.clear(); - // We keep the buckets for memory estimation - // reduce_buckets->clear(); -} - -void cleanup() -{ - reset(); - - ncclCommDestroy(nccl_comm); - process_group = nullptr; - symm_mem = nullptr; -} - -at::Tensor reduce_grad(at::Tensor grad_tensor, long graph_id, long ds_id) -{ - if (sync_before_reduce) { c10::cuda::device_synchronize(); } - - assert(hasKey(executors, graph_id)); - if (!profile) { executors[graph_id]->reduceGrad(grad_tensor, ds_id); } - - if (sync_after_reduce) { c10::cuda::device_synchronize(); } - - return at::Tensor(); -} - -at::Tensor reduce_grad_meta(at::Tensor grad_tensor, long graph_id, long ds_id) -{ - return at::Tensor(); -} - -void free_tensors(std::vector tensors) -{ - int64_t THRESHOLD = 10 * 1024 * 1024; - - if (!profile) { - for (auto& tensor : tensors) { - if (tensor.is_cuda() && tensor.numel() > THRESHOLD) { - tensor.record_stream(at::cuda::getCurrentCUDAStream()); - tensor.set_data(torch::empty({0}, tensor.options())); - } - } - } -} - -void free_tensors_meta(std::vector tensors) {} - -void init(c10::intrusive_ptr pg, - int64_t initial_reduce_bucket_size, - bool enable_double_buffer, - bool _use_symm_mem, - bool _clone_custom_op_output, - bool _sync_before_reduce, - bool _sync_after_reduce, - bool _sync_before_allgather, - bool _sync_after_allgather) -{ - process_group = pg; - - ncclUniqueId ncclID; - ncclGetUniqueId(&ncclID); - - // ProcessGroup doesn't have an API to get the CUDA stream for comm calls. - // So we create a NCCL communicator and call NCCL APIs directly. - auto vec = std::vector(reinterpret_cast(&ncclID), - reinterpret_cast(&ncclID) + NCCL_UNIQUE_ID_BYTES); - auto device = torch::Device(torch::kCUDA); - at::Tensor tensor = torch::from_blob(vec.data(), {static_cast(vec.size())}, torch::kUInt8) - .to(torch::Device(torch::kCUDA)); - std::vector bcast_input = {tensor}; - - process_group->broadcast(bcast_input, c10d::BroadcastOptions())->wait(); - - // create a new nccl communicator - std::memcpy(&ncclID, tensor.to(torch::Device(torch::kCPU)).data_ptr(), NCCL_UNIQUE_ID_BYTES); - ncclCommInitRank(&nccl_comm, process_group->getSize(), ncclID, process_group->getRank()); - - param_registry = std::make_shared(); - reduce_buckets = std::make_shared(initial_reduce_bucket_size, - enable_double_buffer); - use_symm_mem = _use_symm_mem; - clone_custom_op_output = _clone_custom_op_output; - - sync_before_reduce = _sync_before_reduce; - sync_after_reduce = _sync_after_reduce; - sync_before_allgather = _sync_before_allgather; - sync_after_allgather = _sync_after_allgather; -} - -void start_forward() -{ - lazy_init_symm_memory(); - for (auto& it : executors) { it.second->startForward(); } -} - -void end_forward() -{ - for (auto& it : executors) { it.second->endForward(); } -} - -void start_backward(bool update) -{ - for (auto& it : executors) { it.second->startBackward(update); } -} - -// We don't call this -// void end_backward(bool update) -// { -// } - -} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/init.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/init.cpp deleted file mode 100644 index 9b80fff4e6a7fab16ed8dacaf2af30a55ea71567..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/init.cpp +++ /dev/null @@ -1,103 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "deepcompile.h" -#include "z1.h" -#include "z2.h" -#include "z3.h" - -TORCH_LIBRARY(dc, m) -{ - m.def("allgather_param(Tensor a, int graph_id, int id) -> Tensor"); - m.def("prefetch_params_fused(int graph_id, Tensor[] params, int[] ids) -> ()"); - m.def("wait_allgather(Tensor a, int graph_id, int id) -> Tensor"); - m.def("release_param(Tensor a, int graph_id, int id, int n_users) -> Tensor"); - m.def("reduce_grad(Tensor a, int graph_id, int id) -> Tensor"); - m.def("free_tensors(Tensor[] a) -> ()"); - m.def("offload_tensor(Tensor a, int id, int id) -> Tensor"); - m.def("reload_tensor(Tensor a, int id, int id) -> Tensor"); - m.def("wait_offload(Tensor a, int id, int id) -> Tensor"); - m.def("wait_reload(Tensor a, int id, int id) -> Tensor"); - m.def("offload_parameter(Tensor a, int id, int id) -> ()"); - m.def("reload_parameter(Tensor a, int id, int id) -> ()"); - - m.def("test_call(Tensor a) -> Tensor"); -} - -TORCH_LIBRARY_IMPL(dc, CPU, m) -{ - m.impl("allgather_param", &dc::allgather_param); - m.impl("prefetch_params_fused", &dc::prefetch_params_fused); - m.impl("wait_allgather", &dc::wait_allgather); - m.impl("release_param", &dc::release_param); - m.impl("reduce_grad", &dc::reduce_grad); - m.impl("free_tensors", &dc::free_tensors); - m.impl("offload_tensor", &dc::offload_tensor); - m.impl("reload_tensor", &dc::reload_tensor); - m.impl("wait_offload", &dc::wait_offload); - m.impl("wait_reload", &dc::wait_reload); - m.impl("offload_parameter", &dc::offload_parameter); - m.impl("reload_parameter", &dc::reload_parameter); - - m.impl("test_call", &dc::test_call); -} - -TORCH_LIBRARY_IMPL(dc, CUDA, m) -{ - m.impl("allgather_param", &dc::allgather_param); - m.impl("prefetch_params_fused", &dc::prefetch_params_fused); - m.impl("wait_allgather", &dc::wait_allgather); - m.impl("release_param", &dc::release_param); - m.impl("reduce_grad", &dc::reduce_grad); - m.impl("free_tensors", &dc::free_tensors); - m.impl("offload_tensor", &dc::offload_tensor); - m.impl("reload_tensor", &dc::reload_tensor); - m.impl("wait_offload", &dc::wait_offload); - m.impl("wait_reload", &dc::wait_reload); - m.impl("offload_parameter", &dc::offload_parameter); - m.impl("reload_parameter", &dc::reload_parameter); - - m.impl("test_call", &dc::test_call); -} - -TORCH_LIBRARY_IMPL(dc, Meta, m) -{ - m.impl("allgather_param", &dc::allgather_param_meta); - m.impl("prefetch_params_fused", &dc::prefetch_params_fused_meta); - m.impl("release_param", &dc::release_param_meta); - m.impl("wait_allgather", &dc::wait_allgather_meta); - m.impl("reduce_grad", &dc::reduce_grad_meta); - m.impl("free_tensors", &dc::free_tensors_meta); - m.impl("reload_parameter", &dc::reload_parameter_meta); - m.impl("offload_parameter", &dc::offload_parameter_meta); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("set_persistent", &dc::set_persistent, "Set persistent flag for a parameter"); - m.def("enable_profiling", &dc::enable_profiling, "Enable profiling"); - m.def("is_profiling", &dc::is_profiling, "Check if profiling is enabled"); - m.def("init", &dc::init, "Set the process group"); - m.def("cleanup", &dc::cleanup, "Cleanup the process group"); - m.def("register_param", &dc::register_param, "Register a parameter"); - m.def("register_graph_z1", - &dc::register_graph_z1, - "Register graph with a list of ds parameter ids"); - m.def("register_graph_z2", - &dc::register_graph_z2, - "Register graph with a list of ds parameter ids"); - m.def("register_z3_param", &dc::register_z3_param, "Register a parameter"); - m.def("register_graph_z3", - &dc::register_graph_z3, - "Register graph with a list of ds parameter ids"); - m.def("start_forward", &dc::start_forward, "Start forward pass"); - m.def("end_forward", &dc::end_forward, "End forward pass"); - m.def("start_backward", &dc::start_backward, "Start backward pass"); - // m.def("end_backward", &dc::end_backward, "End backward pass"); - m.def("cleanup", &dc::cleanup, "Clean up DeepCompile"); - m.def("reset", &dc::reset, "Reset the state"); - m.def("invalidate_gathered_param", &dc::invalidate_gathered_param, "Invalidate gathered param"); - m.def("clear_all_gathered_params", &dc::clear_all_gathered_params, "Clear all gathered params"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/util.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/util.cpp deleted file mode 100644 index 948338028059bc1eff42e00d1d7390c09b58fbb0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/util.cpp +++ /dev/null @@ -1,89 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "deepcompile.h" - -#include - -namespace dc { - -std::string tensorToString(const at::Tensor& t, size_t max_elem, size_t max_str_len) -{ - auto t_cpu = t.flatten() - .slice(0, 0, std::min((int64_t)max_elem, t.numel())) - .to(c10::Device(c10::kCPU), false, true); - - size_t size = std::min(max_elem, productDim(t.sizes())); - - if (t.scalar_type() == c10::ScalarType::Half || t.scalar_type() == c10::ScalarType::BFloat16) { - auto float_ten = t_cpu.to(c10::ScalarType::Float, false, true).contiguous(); - return tensorPtrToString((float*)float_ten.data_ptr(), size, max_str_len); - } else if (t.scalar_type() == c10::ScalarType::Float) { - return tensorPtrToString((float*)t_cpu.data_ptr(), size, max_str_len); - } else if (t.scalar_type() == c10::ScalarType::Double) { - return tensorPtrToString((double*)t_cpu.data_ptr(), size, max_str_len); - } else if (t.scalar_type() == c10::ScalarType::Int) { - int* ptr = static_cast(t_cpu.data_ptr()); - return tensorPtrToString(ptr, size, max_str_len); - } else if (t.scalar_type() == c10::ScalarType::Long) { - long* ptr = static_cast(t_cpu.data_ptr()); - return tensorPtrToString(ptr, size, max_str_len); - } else if (t.scalar_type() == c10::ScalarType::Byte) { - unsigned char* ptr = static_cast(t_cpu.data_ptr()); - std::vector vec; - vec.reserve(size); - for (size_t i = 0; i < size; i++) { - vec.push_back(*ptr); - ptr++; - } - return tensorPtrToString(&vec[0], size, max_str_len); - } else if (t.scalar_type() == c10::ScalarType::Bool) { - bool* ptr = static_cast(t_cpu.data_ptr()); - std::vector vec; - vec.reserve(size); - for (size_t i = 0; i < size; i++) { - vec.push_back(*ptr); - ptr++; - } - return tensorPtrToString(&vec[0], size, max_str_len); - } - std::stringstream ss; - ss << "Failed to convert tensor to string. Invalid type of tensor: " - << toString(t.scalar_type()); - throw std::invalid_argument(ss.str()); -} - -std::string tensorPtrToString(void* ptr, - size_t size, - c10::ScalarType datatype, - size_t max_elem, - size_t max_str_len) -{ - int64_t elem_size = std::min((size_t)max_elem, size); - - if (datatype == c10::ScalarType::Long) { - return tensorPtrToString(static_cast(ptr), elem_size, max_str_len); - } else if (datatype == c10::ScalarType::Int) { - return tensorPtrToString(static_cast(ptr), elem_size, max_str_len); - } else if (datatype == c10::ScalarType::Double) { - return tensorPtrToString(static_cast(ptr), elem_size, max_str_len); - } else if (datatype == c10::ScalarType::Float) { - return tensorPtrToString(static_cast(ptr), elem_size, max_str_len); - } else if (datatype == c10::ScalarType::Half || datatype == c10::ScalarType::BFloat16) { - const auto ten = torch::from_blob(ptr, {(int64_t)elem_size}, datatype); - auto float_ten = ten.to(c10::ScalarType::Float, false, true).contiguous(); - return tensorPtrToString((float*)float_ten.data_ptr(), elem_size, max_str_len); - } - std::stringstream ss; - ss << "Failed to convert tensor ptr to string. Invalid type of tensor: " << toString(datatype); - throw std::invalid_argument(ss.str()); -} - -std::string tensorDimToString(const at::Tensor& t) -{ - const auto dim = t.sizes(); - return join_as_str(dim); -} -} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z1.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z1.cpp deleted file mode 100644 index cbec2dec82ab13e055a433a4f2f96c704f3709cc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z1.cpp +++ /dev/null @@ -1,117 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "z1.h" -#include "deepcompile.h" - -namespace dc { - -class Z1CustomOpExecutor : public CustomOpExecutor { -public: - Z1CustomOpExecutor(c10::intrusive_ptr process_group, - std::shared_ptr param_registry, - std::shared_ptr reduce_buckets, - std::vector ds_ids, - ncclComm_t nccl_comm, - at::cuda::CUDAStream rs_stream, - at::cuda::CUDAStream copy_stream, - bool pre_div_reduce) - : CustomOpExecutor(process_group, - param_registry, - reduce_buckets, - ds_ids, - nccl_comm, - rs_stream, - copy_stream, - pre_div_reduce) - { - } - ~Z1CustomOpExecutor() {} - - at::Tensor reduceGrad(at::Tensor grad_tensor, long ds_id) override - { - if (!hasKey(grad_tensors_, ds_id)) { - grad_tensors_[ds_id] = grad_tensor; - } else { - grad_tensors_[ds_id].add_(grad_tensor); - } - - if (param_updated_) { - CustomOpExecutor::reduceGrad(grad_tensors_[ds_id], ds_id); - grad_tensors_.erase(ds_id); - } - - return at::Tensor(); - } - - void flushReduceBucket(at::ScalarType scalar_type) override - { - if (!hasKey(reduce_tasks_, scalar_type)) { return; } - - blockCopyEvents(scalar_type); - applyPreDivision(scalar_type); - - // NCCL AllReduce operation - ncclGroupStart(); - for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { - ncclResult_t result = ncclAllReduce(t.getSendBuf().data_ptr(), - t.getSendBuf().data_ptr(), - t.getSendBuf().numel(), - get_nccl_data_type(scalar_type), - getReductionOp(), - nccl_comm_, - rs_stream_); - if (result != ncclSuccess) { throw std::runtime_error("NCCL AllReduce failed"); } - } - ncclGroupEnd(); - - // Copy results to gradient buffers - { - at::cuda::CUDAStreamGuard guard(rs_stream_); - for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { - auto param = param_registry_->getParam(t.getDSId()); - auto grad_buf = param.getGradBuffer().flatten(); - - if (grad_buf.numel() == 0) { continue; } - - int64_t offset = param.getOffset(); - auto recv_buf = t.getSendBuf().flatten().index( - {torch::indexing::Slice(offset, offset + grad_buf.numel())}); - grad_buf.copy_(recv_buf); - } - } - - performCleanup(scalar_type); - } - -protected: - std::unordered_map grad_tensors_; -}; - -static at::cuda::CUDAStream rs_stream = at::cuda::getStreamFromPool(true); -static at::cuda::CUDAStream copy_stream = at::cuda::getStreamFromPool(true); - -void register_graph_z1(long graph_id, const std::vector& ds_ids) -{ - executors[graph_id] = std::make_shared(process_group, - param_registry, - reduce_buckets, - ds_ids, - nccl_comm, - rs_stream, - copy_stream, - pre_div_reduce); -} - -void register_param(long ds_id, - const std::vector& ds_shape, - at::Tensor ds_tensor, - at::Tensor grad_buffer, - int64_t offset) -{ - param_registry->registerParam(ds_id, ds_shape, ds_tensor, grad_buffer, false, offset, false); -} - -} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z1.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z1.h deleted file mode 100644 index 1d3607a59b066559ab73b11d23d292e70afd432f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z1.h +++ /dev/null @@ -1,18 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "deepcompile.h" - -#pragma once - -namespace dc { - -void register_graph_z1(long graph_id, const std::vector& ds_ids); -void register_param(long ds_id, - const std::vector& ds_shape, - at::Tensor ds_tensor, - at::Tensor grad_buffer, - int64_t offset); -} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z2.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z2.cpp deleted file mode 100644 index 2effa83f02e516b2aeb8f60b6ba1d2f1dd97cb61..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z2.cpp +++ /dev/null @@ -1,102 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "z2.h" -#include "deepcompile.h" - -namespace dc { - -class Z2CustomOpExecutor : public CustomOpExecutor { -public: - Z2CustomOpExecutor(c10::intrusive_ptr process_group, - std::shared_ptr param_registry, - std::shared_ptr reduce_buckets, - std::vector ds_ids, - ncclComm_t nccl_comm, - at::cuda::CUDAStream rs_stream, - at::cuda::CUDAStream copy_stream, - bool pre_div_reduce) - : CustomOpExecutor(process_group, - param_registry, - reduce_buckets, - ds_ids, - nccl_comm, - rs_stream, - copy_stream, - pre_div_reduce) - { - } - ~Z2CustomOpExecutor() {} - - void endBackward() override - { - if (param_updated_) { - for (auto& it : has_acc_grad_) { it.second = false; } - } - } - - void flushReduceBucket(at::ScalarType scalar_type) override - { - if (!hasKey(reduce_tasks_, scalar_type)) { return; } - - blockCopyEvents(scalar_type); - applyPreDivision(scalar_type); - - // NCCL AllReduce operation - ncclGroupStart(); - for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { - ncclResult_t result = ncclAllReduce(t.getSendBuf().data_ptr(), - t.getSendBuf().data_ptr(), - t.getSendBuf().numel(), - get_nccl_data_type(scalar_type), - getReductionOp(), - nccl_comm_, - rs_stream_); - if (result != ncclSuccess) { throw std::runtime_error("NCCL AllReduce failed"); } - } - ncclGroupEnd(); - - // Copy or accumulate results to gradient buffers - { - at::cuda::CUDAStreamGuard guard(rs_stream_); - for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { - bool acc_grad = has_acc_grad_.at(t.getDSId()); - auto param = param_registry_->getParam(t.getDSId()); - auto grad_buf = param.getGradBuffer().flatten(); - - if (grad_buf.numel() == 0) { continue; } - - int64_t offset = param.getOffset(); - auto recv_buf = t.getSendBuf().flatten().index( - {torch::indexing::Slice(offset, offset + grad_buf.numel())}); - if (acc_grad) { - grad_buf.add_(recv_buf); - } else { - grad_buf.copy_(recv_buf); - } - has_acc_grad_[t.getDSId()] = true; - } - } - - performCleanup(scalar_type); - } -}; - -static at::cuda::CUDAStream rs_stream = at::cuda::getStreamFromPool(true); -static at::cuda::CUDAStream copy_stream = at::cuda::getStreamFromPool(true); - -void register_graph_z2(long graph_id, const std::vector& ds_ids) -{ - executors[graph_id] = std::make_shared(process_group, - param_registry, - reduce_buckets, - ds_ids, - nccl_comm, - rs_stream, - copy_stream, - pre_div_reduce); -} - -} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z2.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z2.h deleted file mode 100644 index cc6c3136c20c9c69d253caafde31c743d2243dda..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z2.h +++ /dev/null @@ -1,14 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "deepcompile.h" - -#pragma once - -namespace dc { - -void register_graph_z2(long graph_id, const std::vector& ds_ids); - -} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z3.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z3.cpp deleted file mode 100644 index 75352c2bafdeb548b08d0e7a2e3fbe84beaafa63..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z3.cpp +++ /dev/null @@ -1,524 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "z3.h" -#include "deepcompile.h" - -namespace dc { - -const size_t TIMEOUT_SYMMETRIC_MEMORY_BARRIER = 60000; - -class Z3CustomOpExecutor : public CustomOpExecutor { -public: - Z3CustomOpExecutor(c10::intrusive_ptr process_group, - std::shared_ptr param_registry, - std::shared_ptr reduce_buckets, - std::vector ds_ids, - ncclComm_t nccl_comm, - at::cuda::CUDAStream ag_stream, - at::cuda::CUDAStream rs_stream, - at::cuda::CUDAStream copy_stream, - at::cuda::CUDAStream offload_stream, - at::cuda::CUDAStream reload_stream, - bool pre_div_reduce) - : CustomOpExecutor(process_group, - param_registry, - reduce_buckets, - ds_ids, - nccl_comm, - rs_stream, - copy_stream, - pre_div_reduce), - ag_stream_(ag_stream), - offload_stream_(offload_stream), - reload_stream_(reload_stream) - { - for (long ds_id : ds_ids_) { - ag_comm_done_events_[ds_id] = - std::make_shared(cudaEventDisableTiming); - ag_comp_done_events_[ds_id] = - std::make_shared(cudaEventDisableTiming); - - param_use_count_[ds_id] = 0; - } - } - ~Z3CustomOpExecutor() {} - - void endBackward() override - { - if (param_updated_) { - for (auto& it : has_acc_grad_) { - it.second = false; - param_registry_->setValid(it.first, false); - } - } - - for (auto& it : reload_buffers_) { - it.second.record_stream(at::cuda::getCurrentCUDAStream()); - } - reload_buffers_.clear(); - } - - void launchAllGather(at::Tensor output_buf, - long ds_id, - c10::intrusive_ptr symm_mem) - { - const DSParam& param = param_registry_->getParam(ds_id); - const at::Tensor& ds_tensor = param.getDSTensor(); - - if (symm_mem == nullptr) { - ncclResult_t result = ncclAllGather(ds_tensor.contiguous().data_ptr(), - output_buf.data_ptr(), - ds_tensor.numel(), - get_nccl_data_type(ds_tensor.scalar_type()), - nccl_comm_, - ag_stream_); - - if (result != ncclSuccess) { throw std::runtime_error("NCCL AllGather failed"); } - } else { - at::cuda::CUDAStreamGuard guard(ag_stream_); - int world_size = process_group_->getSize(); - int rank = process_group_->getRank(); - - at::Tensor local_buf = - symm_mem->get_buffer(rank, ds_tensor.sizes(), ds_tensor.scalar_type(), 0); - local_buf.copy_(ds_tensor, true); - - symm_mem->barrier(0, TIMEOUT_SYMMETRIC_MEMORY_BARRIER); - auto chunks = output_buf.flatten().chunk(world_size); - for (int step = 0; step < world_size; step++) { - int remote_rank = (rank - step + world_size) % world_size; - auto src_buf = symm_mem->get_buffer( - remote_rank, ds_tensor.sizes(), ds_tensor.scalar_type(), 0); - chunks[remote_rank].copy_(src_buf.flatten(), true); - } - symm_mem->barrier(0, TIMEOUT_SYMMETRIC_MEMORY_BARRIER); - } - - param_registry_->registerGatheredParam(ds_id, output_buf); - param_registry_->setValid(ds_id, true); - } - - at::Tensor allgatherParam(long ds_id, - c10::intrusive_ptr symm_mem) - { - if (param_registry_->isValid(ds_id)) { return param_registry_->getGatheredParam(ds_id); } - - const DSParam& param = param_registry_->getParam(ds_id); - const at::Tensor& ds_tensor = param.getDSTensor(); - at::Tensor output_buf = param_registry_->hasGatheredParam(ds_id) - ? param_registry_->getGatheredParam(ds_id) - : torch::empty(param.getShape(), ds_tensor.options()); - - assert(hasKey(ag_comp_done_events_, ds_id)); - ag_comp_done_events_[ds_id]->record(); - ag_comp_done_events_[ds_id]->block(ag_stream_); - - launchAllGather(output_buf, ds_id, symm_mem); - - ag_comm_done_events_[ds_id]->record(ag_stream_); - return output_buf; - } - - void prefetchParamsFused(std::vector ds_ids, - c10::intrusive_ptr symm_mem) - { - std::vector invalid_ds_ids; - for (const auto& ds_id : ds_ids) { - if (!param_registry_->isValid(ds_id)) { invalid_ds_ids.push_back(ds_id); } - } - - std::unordered_map output_bufs; - for (long ds_id : invalid_ds_ids) { - const DSParam& param = param_registry_->getParam(ds_id); - if (param_registry_->hasGatheredParam(ds_id)) { - output_bufs[ds_id] = param_registry_->getGatheredParam(ds_id); - } else { - output_bufs[ds_id] = torch::empty(param.getShape(), param.getDSTensor().options()); - } - } - - for (long ds_id : invalid_ds_ids) { - ag_comp_done_events_[ds_id]->record(); - ag_comp_done_events_[ds_id]->block(ag_stream_); - } - - ncclGroupStart(); - for (long ds_id : invalid_ds_ids) { - assert(hasKey(output_bufs, ds_id)); - launchAllGather(output_bufs.at(ds_id), ds_id, symm_mem); - } - ncclGroupEnd(); - - for (long ds_id : invalid_ds_ids) { ag_comm_done_events_[ds_id]->record(ag_stream_); } - } - - void releaseParam(long ds_id, long n_users) - { - const DSParam& param = param_registry_->getParam(ds_id); - - assert(hasKey(param_use_count_, ds_id)); - if (param_use_count_[ds_id] == 0) { param_use_count_[ds_id] = n_users; } - param_use_count_[ds_id]--; - - if (param_use_count_[ds_id] == 0 && !param.isPersistent()) { - at::Tensor gathered_param = param_registry_->getGatheredParam(ds_id); - - if (gathered_param.defined()) { // gathered param is undefined while profiling - const auto options = gathered_param.options(); - at::Tensor empty_buffer = torch::empty({0}, options); - gathered_param.set_data(empty_buffer); - } - - param_registry_->unregisterGatheredParam(ds_id); - } - } - - at::Tensor waitAllgather(at::Tensor v, long ds_id) - { - assert(hasKey(ag_comm_done_events_, ds_id)); - ag_comm_done_events_[ds_id]->block(at::cuda::getCurrentCUDAStream()); - return v; - } - - void flushReduceBucket(at::ScalarType scalar_type) override - { - if (!hasKey(reduce_tasks_, scalar_type)) { return; } - - blockCopyEvents(scalar_type); - - // Calculate temporary buffer size for accumulated gradients - int64_t tmp_recv_numel = 0; - for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { - if (has_acc_grad_.at(t.getDSId())) { - tmp_recv_numel += param_registry_->getParam(t.getDSId()).getGradBuffer().numel(); - } - } - - // Allocate temporary buffer if needed - at::Tensor tmp_recv_buf = at::Tensor(); - if (tmp_recv_numel > 0) { - at::cuda::CUDAStreamGuard guard(rs_stream_); - tmp_recv_buf = torch::empty({tmp_recv_numel}, - at::TensorOptions().dtype(scalar_type).device(at::kCUDA)); - } - - applyPreDivision(scalar_type); - - // NCCL ReduceScatter operation - ncclGroupStart(); - int64_t offset = 0; - for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { - auto recv_buf = param_registry_->getParam(t.getDSId()).getGradBuffer(); - bool acc_grad = has_acc_grad_.at(t.getDSId()); - - if (acc_grad) { - recv_buf = - tmp_recv_buf.index({torch::indexing::Slice(offset, offset + recv_buf.numel())}); - } - - ncclResult_t result = ncclReduceScatter(t.getSendBuf().data_ptr(), - recv_buf.data_ptr(), - recv_buf.numel(), - get_nccl_data_type(scalar_type), - getReductionOp(), - nccl_comm_, - rs_stream_); - if (result != ncclSuccess) { throw std::runtime_error("NCCL ReduceScatter failed"); } - - if (acc_grad) { offset += recv_buf.numel(); } - } - ncclGroupEnd(); - - // Handle gradient accumulation with temporary buffer - { - at::cuda::CUDAStreamGuard guard(rs_stream_); - int64_t offset = 0; - for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { - bool acc_grad = has_acc_grad_.at(t.getDSId()); - - if (acc_grad) { - auto recv_buf = param_registry_->getParam(t.getDSId()).getGradBuffer(); - recv_buf.add_(tmp_recv_buf.index( - {torch::indexing::Slice(offset, offset + recv_buf.numel())})); - offset += recv_buf.numel(); - } - has_acc_grad_[t.getDSId()] = true; - } - } - - performCleanup(scalar_type); - - // Record stream for temporary buffer to prevent early deallocation - if (tmp_recv_numel > 0) { tmp_recv_buf.record_stream(rs_stream_); } - } - - at::Tensor offloadTensor(at::Tensor tensor, long id) - { - if (!hasKey(offload_events_, id)) { - offload_events_[id] = std::make_shared(cudaEventDisableTiming); - offload_comp_done_events_[id] = - std::make_shared(cudaEventDisableTiming); - - const auto options = at::TensorOptions().pinned_memory(true).device(torch::kCPU); - offload_buffers_[id] = at::empty_like(tensor, options); - } - - offload_comp_done_events_[id]->record(); - offload_comp_done_events_[id]->block(offload_stream_); - { - at::cuda::CUDAStreamGuard guard(offload_stream_); - offload_buffers_.at(id).copy_(tensor, true); - } - - tensor.record_stream(offload_stream_); - - offload_events_[id]->record(offload_stream_); - assert(hasKey(offload_buffers_, id)); - return offload_buffers_.at(id); - } - - at::Tensor reloadTensor(at::Tensor tensor, long id) - { - if (!hasKey(reload_events_, id)) { - reload_events_[id] = std::make_shared(cudaEventDisableTiming); - } - - assert(hasKey(offload_buffers_, id)); - offload_events_[id]->block(reload_stream_); - - at::Tensor ten; - { - at::cuda::CUDAStreamGuard guard(reload_stream_); - - assert(hasKey(offload_buffers_, id)); - at::Tensor buf = offload_buffers_.at(id); - const auto options = at::TensorOptions().device(torch::kCUDA); - ten = at::empty_like(buf, options); - ten.copy_(buf, true); - - reload_buffers_[id] = ten; - } - - reload_events_[id]->record(reload_stream_); - return ten; - } - - at::Tensor waitOffload(at::Tensor tensor, long id) - { - assert(hasKey(offload_events_, id)); - offload_events_[id]->block(at::cuda::getCurrentCUDAStream()); - - assert(hasKey(offload_buffers_, id)); - return offload_buffers_.at(id); - } - - at::Tensor waitReload(at::Tensor tensor, long id) - { - assert(hasKey(reload_events_, id)); - reload_events_[id]->block(at::cuda::getCurrentCUDAStream()); - - assert(hasKey(reload_buffers_, id)); - auto ten = reload_buffers_.at(id); - - // We can't release here because the tensor is still being used - // We will need "freeReloadedTensor" after the last user of the tensor to call - // ".record_stream". As it is a bit complicated, we clear the buffer and do at the end of - // the backward pass for now. reload_buffers_.erase(id); - return ten; - } - - void offloadParameter(at::Tensor tensor, long ds_id) { param_registry_->offload(ds_id); } - void reloadParameter(at::Tensor tensor, long ds_id) { param_registry_->reload(ds_id); } - - bool hasReloadBuffer(long id) { return hasKey(reload_buffers_, id); } - - bool hasParam(long ds_id) const { return hasKey(has_acc_grad_, ds_id); } - -private: - at::cuda::CUDAStream ag_stream_; - at::cuda::CUDAStream offload_stream_; - at::cuda::CUDAStream reload_stream_; - - std::unordered_map> ag_comp_done_events_; - std::unordered_map> ag_comm_done_events_; - - std::unordered_map> offload_events_; - std::unordered_map> offload_comp_done_events_; - std::unordered_map> reload_events_; - std::unordered_map offload_buffers_; - std::unordered_map reload_buffers_; - - std::unordered_map param_use_count_; -}; - -static at::cuda::CUDAStream ag_stream = at::cuda::getStreamFromPool(true); -static at::cuda::CUDAStream rs_stream = at::cuda::getStreamFromPool(true); -static at::cuda::CUDAStream copy_stream = at::cuda::getStreamFromPool(true); -static at::cuda::CUDAStream offload_stream = at::cuda::getStreamFromPool(true); -static at::cuda::CUDAStream reload_stream = at::cuda::getStreamFromPool(true); - -void register_graph_z3(long graph_id, const std::vector& ds_ids) -{ - executors[graph_id] = std::make_shared(process_group, - param_registry, - reduce_buckets, - ds_ids, - nccl_comm, - ag_stream, - rs_stream, - copy_stream, - offload_stream, - reload_stream, - pre_div_reduce); -} - -void register_z3_param(long ds_id, - const std::vector& ds_shape, - at::Tensor ds_tensor, - at::Tensor grad_buffer, - bool persistent) -{ - param_registry->registerParam(ds_id, ds_shape, ds_tensor, grad_buffer, true, 0, persistent); - if (persistent) { param_registry->registerGatheredParam(ds_id, ds_tensor); } -} - -at::Tensor allgather_param(at::Tensor param_tensor, long graph_id, long ds_id) -{ - auto executor = getExecutor(graph_id, executors); - - if (sync_before_allgather) { c10::cuda::device_synchronize(); } - auto ret = executor->allgatherParam(ds_id, symm_mem); - if (sync_after_allgather) { c10::cuda::device_synchronize(); } - return ret; -} - -void set_persistent(long ds_id) -{ - param_registry->setPersistent(ds_id, true); - - // Allocate buffer here - // Memory fragmentation will be more severe if we allocate in forward/backward - for (auto& it : executors) { - if (it.second->hasParam(ds_id)) { - auto executor = getExecutor(it.first, executors); - executor->allgatherParam(ds_id, symm_mem); - } - } -} - -void prefetch_params_fused(long graph_id, - const std::vector params, - const std::vector& ds_ids) -{ - auto executor = getExecutor(graph_id, executors); - executor->prefetchParamsFused(ds_ids, symm_mem); -} - -void prefetch_params_fused_meta(long graph_id, - const std::vector params, - const std::vector& ds_ids) -{ -} - -// for profiling -void invalidate_gathered_param(long ds_id) -{ - const DSParam& param = param_registry->getParam(ds_id); - if (param.isPersistent()) { return; } - - param_registry->unregisterGatheredParam(ds_id); - param_registry->registerGatheredParam(ds_id, at::Tensor()); -} - -void clear_all_gathered_params() -{ - for (const auto& it : param_registry->getParams()) { - long ds_id = it.first; - const DSParam& param = param_registry->getParam(ds_id); - if (param.isPersistent()) { continue; } - if (param_registry->hasGatheredParam(ds_id)) { - param_registry->unregisterGatheredParam(ds_id); - } - } -} - -at::Tensor allgather_param_meta(at::Tensor param_tensor, long graph_id, long ds_id) -{ - const DSParam& param = param_registry->getParam(ds_id); - auto options = param.getDSTensor().options().device(c10::kMeta); - at::Tensor output_buf = torch::empty(param.getShape(), options); - return output_buf; -} - -at::Tensor release_param(at::Tensor dummy, long graph_id, long ds_id, long n_users) -{ - auto executor = getExecutor(graph_id, executors); - executor->releaseParam(ds_id, n_users); - - if (clone_custom_op_output) { return dummy.clone(); } - return dummy; -} - -at::Tensor release_param_meta(at::Tensor dummy, long graph_id, long ds_id, long n_users) -{ - return dummy; -} - -at::Tensor wait_allgather(at::Tensor v, long graph_id, long ds_id) -{ - auto executor = getExecutor(graph_id, executors); - executor->waitAllgather(v, ds_id); - return v; -} - -at::Tensor wait_allgather_meta(at::Tensor v, long graph_id, long ds_id) { return v; } - -at::Tensor offload_tensor(at::Tensor tensor, long graph_id, long id) -{ - auto executor = getExecutor(graph_id, executors); - return executor->offloadTensor(tensor, id); -} - -at::Tensor reload_tensor(at::Tensor tensor, long graph_id, long id) -{ - auto executor = getExecutor(graph_id, executors); - return executor->reloadTensor(tensor, id); -} - -at::Tensor wait_offload(at::Tensor tensor, long graph_id, long id) -{ - auto executor = getExecutor(graph_id, executors); - return executor->waitOffload(tensor, id); -} - -at::Tensor wait_reload(at::Tensor tensor, long graph_id, long id) -{ - auto executor = getExecutor(graph_id, executors); - if (profile && !executor->hasReloadBuffer(id)) { return tensor; } - return executor->waitReload(tensor, id); -} - -at::Tensor test_call(at::Tensor a) -{ - std::cout << "test_call" << std::endl; - return a; -} - -void reload_parameter(at::Tensor tensor, long graph_id, long ds_id) -{ - auto executor = getExecutor(graph_id, executors); - executor->reloadParameter(tensor, ds_id); -} - -void offload_parameter(at::Tensor tensor, long graph_id, long ds_id) -{ - auto executor = getExecutor(graph_id, executors); - executor->offloadParameter(tensor, ds_id); -} -void reload_parameter_meta(at::Tensor param_tensor, long graph_id, long ds_id) {} -void offload_parameter_meta(at::Tensor tensor, long graph_id, long ds_id) {} - -} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z3.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z3.h deleted file mode 100644 index 1031f0c84f7c1447d662c1df4629f50c7c47fce1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/compile/z3.h +++ /dev/null @@ -1,48 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "deepcompile.h" - -#pragma once - -namespace dc { - -void register_graph_z3(long graph_id, const std::vector& ds_ids); -void register_graph_ops_z3(long graph_id, - const std::vector& op_names, - const std::vector& n_args); -void register_bwd_graph_ops_z3(long graph_id, - const std::vector& op_names, - const std::vector& n_args); -void register_z3_param(long ds_id, - const std::vector& ds_shape, - at::Tensor ds_tensor, - at::Tensor grad_buffer, - bool persistent); -at::Tensor allgather_param(at::Tensor param_tensor, long graph_id, long ds_id); -void set_persistent(long ds_id); -void prefetch_params_fused(long graph_id, - const std::vector params, - const std::vector& ds_ids); -void prefetch_params_fused_meta(long graph_id, - const std::vector params, - const std::vector& ds_ids); -// for profiling -void invalidate_gathered_param(long ds_id); -void clear_all_gathered_params(); -at::Tensor allgather_param_meta(at::Tensor param_tensor, long graph_id, long ds_id); -at::Tensor release_param(at::Tensor dummy, long graph_id, long ds_id, long n_users); -at::Tensor release_param_meta(at::Tensor dummy, long graph_id, long ds_id, long n_users); -at::Tensor wait_allgather(at::Tensor v, long graph_id, const long ds_id); -at::Tensor wait_allgather_meta(at::Tensor v, long graph_id, long ds_id); -at::Tensor offload_tensor(at::Tensor tensor, long graph_id, long id); -at::Tensor reload_tensor(at::Tensor tensor, long graph_id, long id); -at::Tensor wait_offload(at::Tensor tensor, long graph_id, long id); -at::Tensor wait_reload(at::Tensor tensor, long graph_id, long id); -void reload_parameter(at::Tensor tensor, long graph_id, long id); -void offload_parameter(at::Tensor tensor, long graph_id, long id); -void reload_parameter_meta(at::Tensor tensor, long graph_id, long id); -void offload_parameter_meta(at::Tensor tensor, long graph_id, long id); -} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/adam/fused_adam.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/adam/fused_adam.cpp deleted file mode 100644 index d25578f410da278dace61fb5b488b14ed1257c9f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/adam/fused_adam.cpp +++ /dev/null @@ -1,48 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "cpu_adam.h" - -// C++ interface - -void multi_tensor_adam(int chunk_size, - at::Tensor noop_flag, - std::vector> tensor_lists, /*gpmv*/ - const float lr, - const float beta1, - const float beta2, - const float epsilon, - const int step, - const int mode, - const int bias_correction, - const float weight_decay) -{ - static bool initialized = false; - if (!initialized) { - create_adam_optimizer(0); - initialized = true; - } - for (int i = 0; i < tensor_lists[0].size(); i++) { - ds_adam_step(0, - step, - lr, - beta1, - beta2, - epsilon, - weight_decay, - bias_correction, - tensor_lists[1][i], - tensor_lists[0][i], - tensor_lists[2][i], - tensor_lists[3][i]); - } -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("multi_tensor_adam", - &multi_tensor_adam, - "Compute and apply gradient update to parameters for Adam optimizer"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/ccl.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/ccl.cpp deleted file mode 100644 index d28509e592668aca100555e18ae47e3aadacb43a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/ccl.cpp +++ /dev/null @@ -1,344 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include - -#include -#include "shm.h" - -// #define DO_PROFILE -#ifdef DO_PROFILE -#include -#include -#endif - -// Communication settings -static int world_rank = -1; -static int world_size = -1; - -static std::set _comm_ids; -static std::set _colors; -static std::vector _ccl_comms; -static ccl::shared_ptr_class sub_kvs; -static std::map, int> group_to_comm_id; - -ccl::communicator& _get_comm_from_group() { return _ccl_comms[0]; } -ccl::communicator& _get_comm_from_group(py::object group) { return _ccl_comms[0]; } -ccl::communicator& _get_comm_from_group(std::vector ranks) -{ - if (group_to_comm_id.find(ranks) != group_to_comm_id.end()) { - auto id = group_to_comm_id.find(ranks); - return _ccl_comms[id->second]; - } - return _ccl_comms[0]; -} - -#define CCLCHECK(cmd) \ - do { \ - cmd; \ - } while (0) - -#define KVS_CREATE_SUCCESS 0 -#define KVS_CREATE_FAILURE -1 - -static bool is_initialized = 0; - -static ccl::shared_ptr_class kvs; - -static bool all_ranks_local_p = false; - -void initialize(int size, int rank, torch::Tensor& kvs_data) -{ - if (is_initialized) return; - - // Check whether all ranks is on the same physical machine. - // If true, we will use an SHM based low latency allreduce - - auto ls_string = std::getenv("LOCAL_SIZE"); - int ls = 0; - if (ls_string != NULL) { ls = std::stoi(std::getenv("LOCAL_SIZE")); } - - if (size >= 1 && size == ls) { all_ranks_local_p = true; } - - world_size = size; - world_rank = rank; - is_initialized = 1; - - ccl::kvs::address_type main_addr; - - if (rank != 0) { - memcpy(main_addr.data(), kvs_data.data_ptr(), main_addr.size()); - kvs = ccl::create_kvs(main_addr); - } - - _ccl_comms.emplace_back(ccl::create_communicator(size, rank, kvs)); - - auto addr_string = std::getenv("MASTER_ADDR"); - if (addr_string == NULL) { addr_string = ""; } - auto port_string = std::getenv("MASTER_PORT"); - if (port_string == NULL) { port_string = ""; } - - if (all_ranks_local_p) { shm_initialize(size, rank, addr_string, port_string); } -} - -/* - rank == 0: create main kvs and return its address - rank == else: return an empty address -*/ -std::vector get_kvs_addr(int rank) -{ - if (rank == 0) { - kvs = ccl::create_main_kvs(); - ccl::kvs::address_type main_addr = kvs->get_address(); - auto ccl_kvs_addr = std::vector(main_addr.begin(), main_addr.end()); - return ccl_kvs_addr; - } else { - ccl::kvs::address_type main_addr; - auto ccl_kvs_addr = std::vector(main_addr.begin(), main_addr.end()); - return ccl_kvs_addr; - } -} - -int get_rank(int group = 0) { return world_rank; } - -int get_world_size(int group = 0) { return world_size; } - -// Find the next ordered, unique value to a set. E.g. <0,1,2,7> --> 3 -int next_unique_val(std::set s) -{ - std::set::iterator itr; - // Base case. Add 0 to start of set. - if (s.empty() || *s.begin() != 0) { - return 0; - // second base case where s = {0} (the case of s = {n != 0} is caught above) - } else if (s.size() == 1) { - return 1; - } else { - int prev_val = *s.begin(); - for (itr = std::next(s.begin()); itr != s.end(); itr++) { - if (*itr != prev_val + 1) { return prev_val + 1; } - prev_val = *itr; - } - return *(s.end()) + 1; - } -} - -std::vector get_sub_kvs_addr(bool first) -{ - if (first) { - sub_kvs = ccl::create_main_kvs(); - ccl::kvs::address_type main_addr = sub_kvs->get_address(); - auto ccl_kvs_addr = std::vector(main_addr.begin(), main_addr.end()); - return ccl_kvs_addr; - } else { - ccl::kvs::address_type main_addr; - auto ccl_kvs_addr = std::vector(main_addr.begin(), main_addr.end()); - return ccl_kvs_addr; - } -} - -void initialize_sub_comm(int size, int rank, torch::Tensor& kvs_data, std::vector ranks) -{ - ccl::kvs::address_type main_addr; - if (rank != 0) { - memcpy(main_addr.data(), kvs_data.data_ptr(), main_addr.size()); - sub_kvs = ccl::create_kvs(main_addr); - } - _ccl_comms.push_back(ccl::create_communicator(size, rank, sub_kvs)); - group_to_comm_id[ranks] = _ccl_comms.size() - 1; -} - -ccl::datatype get_ccl_datatype(c10::ScalarType type) -{ - ccl::datatype ccl_type; - switch (type) { - case c10::ScalarType::Int: ccl_type = ccl::datatype::int32; break; - case c10::ScalarType::Long: ccl_type = ccl::datatype::int64; break; - case c10::ScalarType::Float: ccl_type = ccl::datatype::float32; break; - case c10::ScalarType::Double: ccl_type = ccl::datatype::float64; break; - case c10::ScalarType::BFloat16: ccl_type = ccl::datatype::bfloat16; break; - case c10::ScalarType::Half: ccl_type = ccl::datatype::float16; break; - default: ccl_type = ccl::datatype::int8; - } - return ccl_type; -} - -ccl::reduction get_ccl_reduce_op(py::object op, at::Tensor& input) -{ - py::object ReduceOp = py::module_::import("deepspeed.comm").attr("ReduceOp"); - if (!py::isinstance(op, ReduceOp)) { - throw std::runtime_error("Error: Op must be of type ReduceOp"); - } - - int op_val = py::int_(op.attr("value")); - ccl::reduction ccl_op; - - if (input.scalar_type() == at::kBool) { - if (op_val == (int)py::int_(ReduceOp.attr("SUM").attr("value"))) { - // For bool tensors, map sum to max, which both represent a bitwise or. - // This is to prevent overflow issues with sum, since we use uint8 to - // represent a bool (see cclDataType mapping). - ccl_op = ccl::reduction::max; - } else if (op_val == (int)py::int_(ReduceOp.attr("AVG").attr("value"))) { - throw std::runtime_error("Error: For bool tensors, op must be of type ReduceOp"); - } - } - - if (op_val == (int)py::int_(ReduceOp.attr("SUM").attr("value"))) { - ccl_op = ccl::reduction::sum; - } else if (op_val == (int)py::int_(ReduceOp.attr("MIN").attr("value"))) { - ccl_op = ccl::reduction::min; - } else if (op_val == (int)py::int_(ReduceOp.attr("MAX").attr("value"))) { - ccl_op = ccl::reduction::max; - } else if (op_val == (int)py::int_(ReduceOp.attr("PRODUCT").attr("value"))) { - ccl_op = ccl::reduction::prod; - } else { - throw std::runtime_error("Error: Unrecognized ReduceOp type"); - } - return ccl_op; -} - -void broadcast(torch::Tensor& data, int src, std::vector group, bool async_op) -{ - CCLCHECK(ccl::broadcast(data.data_ptr(), - data.numel(), - get_ccl_datatype(data.scalar_type()), - src, - _get_comm_from_group(group)) - .wait()); -} - -// TODO: implement torch's async_op behavior, document it. -void all_reduce(torch::Tensor& data, py::object op, std::vector group, bool async_op) -{ - CCLCHECK(ccl::allreduce(data.data_ptr(), - data.data_ptr(), - data.numel(), - get_ccl_datatype(data.scalar_type()), - get_ccl_reduce_op(op, data), - _get_comm_from_group(group)) - .wait()); -} - -void all_reduce_caching(torch::Tensor& data, - py::object op, - std::string match_id, - std::vector group, - bool async_op) -{ - ccl::allreduce_attr attr = ccl::default_allreduce_attr; - auto match_str = ccl::v1::string(match_id); - attr.template set(true); - attr.template set(match_str); - // To control this, use operation attribute and set true value for to_cache field and unique - // string (for example, tensor name) for match_id field. Note that: - // match_id should be the same for a specific communication operation across all ranks. - // If the same tensor is a part of different communication operations, match_id should have - // different values for each of these operations. - CCLCHECK(ccl::allreduce(data.data_ptr(), - data.data_ptr(), - data.numel(), - get_ccl_datatype(data.scalar_type()), - get_ccl_reduce_op(op, data), - _get_comm_from_group(group), - attr) - .wait()); -} - -void inference_all_reduce(torch::Tensor& data, py::object op) -{ -#ifdef DO_PROFILE - static double total_time = 0.0; - static double total_time_sq = 0.0; - static int count = -16; // warmup - static double max_time = 0.0; - static double min_time = DBL_MAX; - // make sure all rank reach this point before measuring time - // turn on this if you suspect each rank didn't reach here at the same time (stragger) - // if (all_ranks_local_p) { - // barrier_wait(0, world_size); - //} - auto start = std::chrono::system_clock::now(); -#endif - - static py::object ReduceOp = py::module_::import("deepspeed.comm").attr("ReduceOp"); - static auto ReduceOpSum = (int)py::int_(ReduceOp.attr("SUM").attr("value")); - - assert(py::int_(op.attr("value")) == ReduceOpSum); - - auto numel = data.numel(); - - int data_size = 0; - bool data_type_fallback = false; - - switch (data.scalar_type()) { - case c10::ScalarType::BFloat16: data_size = numel * 2; break; - case c10::ScalarType::Float: data_size = numel * 4; break; - default: data_type_fallback = true; - } - - if (data_type_fallback || !all_ranks_local_p) { - // fallback to oneccl allreduce - CCLCHECK(ccl::allreduce(data.data_ptr(), - data.data_ptr(), - data.numel(), - get_ccl_datatype(data.scalar_type()), - get_ccl_reduce_op(op, data), - _get_comm_from_group()) - .wait()); - } else { - all_reduce_outer_loop(data, numel, data_size); - } - -#ifdef DO_PROFILE - auto end = std::chrono::system_clock::now(); - count++; - if (count > 0) { - double elapsed = std::chrono::duration_cast(end - start).count(); - if (elapsed > max_time) { max_time = elapsed; } - if (elapsed < min_time) { min_time = elapsed; } - total_time += elapsed; - total_time_sq += elapsed * elapsed; - if (world_rank == 0 && count == 1000) { - auto avg = total_time / count; - auto sd = - sqrt(total_time_sq / count - total_time * total_time / (count * count)) / avg * 100; - printf(" C++ kernel\t\t %.2f\t %.2f\t%.2f\t %.2f\n", - min_time, - max_time, - total_time / count, - sd); - } - } -#endif -} - -void barrier(std::vector group, bool async_op) -{ - CCLCHECK(ccl::barrier(_get_comm_from_group(group)).wait()); -} - -std::vector get_available_coll() -{ - std::vector colls{ - "broadcast", "all_reduce", "inference_all_reduce", "all_reduce_caching", "barrier"}; - return colls; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("get_kvs_addr", &get_kvs_addr, "create and get main kvs addr"); - m.def("initialize", &initialize, "ccl initialize"); - m.def("get_rank", &get_rank, "get rank"); - m.def("get_world_size", &get_world_size, "get world size"); - m.def("broadcast", &broadcast, "ccl broadcast"); - m.def("all_reduce", &all_reduce, "ccl all_reduce"); - m.def("inference_all_reduce", &inference_all_reduce, "low latency all_reduce implementation"); - m.def("all_reduce_caching", &all_reduce_caching, "ccl all_reduce with caching"); - m.def("barrier", &barrier, "barrier"); - m.def("initialize_sub_comm", &initialize_sub_comm, "initialize_sub_comm"); - m.def("get_sub_kvs_addr", &get_sub_kvs_addr, "get_sub_kvs_addr"); - m.def("get_available_coll", &get_available_coll, "get_available_coll"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm.cpp deleted file mode 100644 index be44681ca0626a9dcef9ccf5e7ea09b38bf27d6c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm.cpp +++ /dev/null @@ -1,692 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include - -#include -#include -#include -#include -#include -#include "shm.h" - -// #define DO_PROFILE -#ifdef DO_PROFILE -#include -#include -#endif - -// states for collectives -enum coll_state { - coll_begin = 0, - coll_allreduce_naive__copy_in_done, - coll_allreduce_naive__reduce_done, - // alternative state when allreduce is working on alternative buffer - // of the double buffer. - coll_alt1_allreduce_naive__copy_in_done, - coll_alt2_allreduce_naive__copy_in_done, - coll_alt1_allreduce_naive__reduce_done, -}; - -// SHM building blocks -struct SharedData { - const char* name; - int descriptor; - void* bytes; - size_t nbytes; -}; - -void shared_open(SharedData* data, const char* name, size_t nbytes) -{ - int d = shm_open(name, O_RDWR, S_IRUSR | S_IWUSR); - if (d != -1) { - void* bytes = mmap(NULL, nbytes, PROT_READ | PROT_WRITE, MAP_SHARED, d, 0); - data->name = name; - data->descriptor = d; - data->bytes = bytes; - data->nbytes = nbytes; - } else { - if (errno != ENOENT) { - // don't print if shm can not be found because we want to loop over from - // caller again until the other ranks created the shm - printf("shared_open %s failed, errno=%d\n", name, errno); - } - data->descriptor = -1; - } -} - -void shared_create(SharedData* data, const char* name, void* bytes, size_t nbytes) -{ - int d = shm_open(name, O_CREAT | O_RDWR, S_IRUSR | S_IWUSR); - if (d != -1) { - if (nbytes = write(d, bytes, nbytes)) { shared_open(data, name, nbytes); } - } else { - printf("shared_create %s failed\n", name); - } -} - -void shared_close(SharedData* data) -{ - if (data->descriptor != -1) { - munmap(data->bytes, data->nbytes); - shm_unlink(data->name); - } -} - -static int world_size; - -// SHM based allreduce helper functions -// buffer that holds shm name -#define NAME_BUF_SIZE 1000 -#define MAX_BUF_SIZE 1048576 * 32 -#define NAIVE_ALLREDUCE_THRESHOLD 1048576 -#define SHM_BUFFER_NAME "deepspeed_allreduce_buffer" -struct allreduce_workspace { - enum coll_state states[2]; // idx=0 -- state for symmetric_naive_all_reduce - // idx=1 -- state for distributed_naive_all_reduce - // double buffer to avoid syncing between rounds - // offset=0 -- 2*NAIVE_ALLREDUCE_THRESHOLD : buffer for symmetric_naive_all_reduce - // after that : buffer for distributed_naive_all_reduce - char buffer[2 * NAIVE_ALLREDUCE_THRESHOLD + 2 * MAX_BUF_SIZE]; -}; - -#define BUFFER0_OFFSET(current_buffer) current_buffer* NAIVE_ALLREDUCE_THRESHOLD -#define BUFFER1_OFFSET(current_buffer) 2 * NAIVE_ALLREDUCE_THRESHOLD + current_buffer* MAX_BUF_SIZE - -struct allreduce_workspace** workspace; - -// buffer for small messages, double buffer -char** symmetric_buffer[2]; -// buffer for large messages, double buffer -char** distributed_buffer[2]; - -void wait_buffer_state_until_2(int index, - enum coll_state state0, - enum coll_state state1, - int state_group) -{ - volatile enum coll_state* state_ptr = &(workspace[index]->states[state_group]); - - while (1) { - volatile enum coll_state cur_state = *state_ptr; - if (cur_state == state0 || cur_state == state1) break; - } -} - -__m512 cvt_bf16_to_fp32(const __m256i src) __attribute__((target("avx512bw"))); -inline __m512 cvt_bf16_to_fp32(const __m256i src) -{ - auto y = _mm512_cvtepu16_epi32(src); - return _mm512_castsi512_ps(_mm512_bslli_epi128(y, 2)); -} - -inline __m256i cvt_fp32_to_bf16(const __m512 src) __attribute__((target("avx512bw"))); -inline __m256i cvt_fp32_to_bf16(const __m512 src) -{ - __m512i value = _mm512_castps_si512(src); - __m512i nan = _mm512_set1_epi32(0xffff); - auto mask_value = _mm512_cmp_ps_mask(src, src, _CMP_ORD_Q); - __m512i ones = _mm512_set1_epi32(0x1); - __m512i vec_bias = _mm512_set1_epi32(0x7fff); - // uint32_t lsb = (input >> 16) & 1; - auto t_value = _mm512_and_si512(_mm512_srli_epi32(value, 16), ones); - // uint32_t rounding_bias = 0x7fff + lsb; - t_value = _mm512_add_epi32(t_value, vec_bias); - // input += rounding_bias; - t_value = _mm512_add_epi32(t_value, value); - // input = input >> 16; - t_value = _mm512_srli_epi32(t_value, 16); - // Check NaN before converting back to bf16 - t_value = _mm512_mask_blend_epi32(mask_value, nan, t_value); - return _mm512_cvtusepi32_epi16(t_value); -} - -__m512 cvt_fp16_to_fp32(const __m256i src) __attribute__((target("avx512bw"))); -inline __m512 cvt_fp16_to_fp32(const __m256i src) { return _mm512_cvtph_ps(src); } - -inline __m256i cvt_fp32_to_fp16(const __m512 src) __attribute__((target("avx512bw"))); -inline __m256i cvt_fp32_to_fp16(const __m512 src) -{ - return _mm512_cvtps_ph(src, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); -} - -void reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) - __attribute__((target("avx512bw"))); - -void reduce_fp16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) - __attribute__((target("avx512bw"))); - -void reduce_fp32_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) - __attribute__((target("avx512bw"))); - -void reduce_all_buffers(int start_elements, - int num_elements, - c10::ScalarType scalar_type, - int to_buffer_idx, - char* to_buffer, - char** buffers) -{ - switch (scalar_type) { - case c10::ScalarType::BFloat16: - reduce_bf16_buffers(start_elements, num_elements, to_buffer, buffers); - break; - case c10::ScalarType::Half: - reduce_fp16_buffers(start_elements, num_elements, to_buffer, buffers); - break; - case c10::ScalarType::Float: - reduce_fp32_buffers(start_elements, num_elements, to_buffer, buffers); - break; - default: assert(!"Should not get here"); - } -} - -#define CVT_ADD_BF16(x) \ - do { \ - auto in##x##_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[x] + i))); \ - inout_val = _mm512_add_ps(inout_val, in##x##_val); \ - } while (0) - -// Reduce functions down below use vectorized algorithm, the number of bytes processed each -// iteration depends on vector length. 256bit vector ==> 32 bytes, 512bit vector ==> 64 bytes -// If you change implementation of reduce_bf16_buffers, etc. , check whether this number needs -// to be changed -#define VECTOR_LENGTH_IN_BYTES 32 - -void reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) -{ - const int element_size = 2; - const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; - int main_elements = num_elements - (num_elements % vector_length); - int remain_elements = num_elements % vector_length; - - // process aligned part -#pragma omp parallel for - for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; - i += VECTOR_LENGTH_IN_BYTES) { - auto inout_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[0] + i))); - switch (world_size) { - case 16: CVT_ADD_BF16(15); - case 15: CVT_ADD_BF16(14); - case 14: CVT_ADD_BF16(13); - case 13: CVT_ADD_BF16(12); - case 12: CVT_ADD_BF16(11); - case 11: CVT_ADD_BF16(10); - case 10: CVT_ADD_BF16(9); - case 9: CVT_ADD_BF16(8); - case 8: CVT_ADD_BF16(7); - case 7: CVT_ADD_BF16(6); - case 6: CVT_ADD_BF16(5); - case 5: CVT_ADD_BF16(4); - case 4: CVT_ADD_BF16(3); - case 3: CVT_ADD_BF16(2); - case 2: CVT_ADD_BF16(1); - case 1: break; - default: - for (int j = 1; j < world_size; j++) { - auto in_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[j] + i))); - inout_val = _mm512_add_ps(inout_val, in_val); - } - } - _mm256_storeu_si256((__m256i*)(to_buffer + i), cvt_fp32_to_bf16(inout_val)); - } - - // process remaining part - int i = (start_elements + main_elements) * element_size; - while (remain_elements > 0) { - float val = 0.0f; - for (int j = 0; j < world_size; j++) { val += *(at::BFloat16*)(buffers[j] + i); } - *(at::BFloat16*)(to_buffer + i) = val; - remain_elements--; - i += element_size; - } -} - -#define CVT_ADD_FP16(x) \ - do { \ - auto in##x##_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[x] + i))); \ - inout_val = _mm512_add_ps(inout_val, in##x##_val); \ - } while (0) - -void reduce_fp16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) -{ - const int element_size = 2; - const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; - int main_elements = num_elements - (num_elements % vector_length); - int remain_elements = num_elements % vector_length; - - // process aligned part -#pragma omp parallel for - for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; - i += VECTOR_LENGTH_IN_BYTES) { - auto inout_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[0] + i))); - switch (world_size) { - case 16: CVT_ADD_FP16(15); - case 15: CVT_ADD_FP16(14); - case 14: CVT_ADD_FP16(13); - case 13: CVT_ADD_FP16(12); - case 12: CVT_ADD_FP16(11); - case 11: CVT_ADD_FP16(10); - case 10: CVT_ADD_FP16(9); - case 9: CVT_ADD_FP16(8); - case 8: CVT_ADD_FP16(7); - case 7: CVT_ADD_FP16(6); - case 6: CVT_ADD_FP16(5); - case 5: CVT_ADD_FP16(4); - case 4: CVT_ADD_FP16(3); - case 3: CVT_ADD_FP16(2); - case 2: CVT_ADD_FP16(1); - case 1: break; - default: - for (int j = 1; j < world_size; j++) { - auto in_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[j] + i))); - inout_val = _mm512_add_ps(inout_val, in_val); - } - } - _mm256_storeu_si256((__m256i*)(to_buffer + i), cvt_fp32_to_fp16(inout_val)); - } - - // process remaining part - int i = (start_elements + main_elements) * element_size; - while (remain_elements > 0) { - float val = 0.0f; - for (int j = 0; j < world_size; j++) { val += *(at::Half*)(buffers[j] + i); } - *(at::Half*)(to_buffer + i) = val; - remain_elements--; - i += element_size; - } -} - -#define CVT_ADD_F32(x) \ - do { \ - auto in##x##_val = _mm256_loadu_ps((float*)(buffers[x] + i)); \ - inout_val = _mm256_add_ps(inout_val, in##x##_val); \ - } while (0) - -void reduce_fp32_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) -{ - const int element_size = 4; - const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; - int main_elements = num_elements - (num_elements % vector_length); - int remain_elements = num_elements % vector_length; - - // process aligned part -#pragma omp parallel for - for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; - i += VECTOR_LENGTH_IN_BYTES) { - auto inout_val = _mm256_loadu_ps((float*)(buffers[0] + i)); - switch (world_size) { - case 16: CVT_ADD_F32(15); - case 15: CVT_ADD_F32(14); - case 14: CVT_ADD_F32(13); - case 13: CVT_ADD_F32(12); - case 12: CVT_ADD_F32(11); - case 11: CVT_ADD_F32(10); - case 10: CVT_ADD_F32(9); - case 9: CVT_ADD_F32(8); - case 8: CVT_ADD_F32(7); - case 7: CVT_ADD_F32(6); - case 6: CVT_ADD_F32(5); - case 5: CVT_ADD_F32(4); - case 4: CVT_ADD_F32(3); - case 3: CVT_ADD_F32(2); - case 2: CVT_ADD_F32(1); - case 1: break; - default: - for (int j = 1; j < world_size; j++) { - auto in_val = _mm256_loadu_ps((float*)(buffers[j] + i)); - inout_val = _mm256_add_ps(inout_val, in_val); - } - } - _mm256_storeu_ps((float*)(to_buffer + i), inout_val); - } - - // process remaining part - int i = (start_elements + main_elements) * element_size; - while (remain_elements > 0) { - float val = 0.0f; - for (int j = 0; j < world_size; j++) { val += *(float*)(buffers[j] + i); } - *(float*)(to_buffer + i) = val; - remain_elements--; - i += element_size; - } -} - -static bool is_initialized = 0; -static int world_rank; - -void shm_initialize(int size, int rank, char* addr_string, char* port_string) -{ - if (is_initialized) return; - is_initialized = 1; - - world_size = size; - world_rank = rank; - - char shm_name_prefix[NAME_BUF_SIZE]; - char shm_name[NAME_BUF_SIZE]; - snprintf(shm_name_prefix, - NAME_BUF_SIZE, - "%s_%d_%s_%s", - SHM_BUFFER_NAME, - getuid(), - addr_string, - port_string); - // create shared workspace for SHM based allreduce - SharedData allreduce_buffer; - // allocate workspace_buf for current rank - struct allreduce_workspace* workspace_buf; - struct allreduce_workspace* workspace_buf_other; - workspace_buf = (struct allreduce_workspace*)malloc(sizeof(struct allreduce_workspace)); - snprintf(shm_name, NAME_BUF_SIZE, "%s_%d", shm_name_prefix, rank); - shared_create(&allreduce_buffer, shm_name, workspace_buf, sizeof(struct allreduce_workspace)); - workspace_buf = (struct allreduce_workspace*)allreduce_buffer.bytes; - workspace_buf->states[0] = coll_alt2_allreduce_naive__copy_in_done; - workspace_buf->states[1] = coll_begin; - - // create the workspace pointer list - workspace = (struct allreduce_workspace**)malloc(size * sizeof(struct allreduce_workspace*)); - symmetric_buffer[0] = (char**)malloc(size * sizeof(char**)); - symmetric_buffer[1] = (char**)malloc(size * sizeof(char**)); - distributed_buffer[0] = (char**)malloc(size * sizeof(char**)); - distributed_buffer[1] = (char**)malloc(size * sizeof(char**)); - - // map shm of all ranks - for (int i = 0; i < size; i++) { - if (i != rank) { - snprintf(shm_name, NAME_BUF_SIZE, "%s_%d", shm_name_prefix, i); - // printf("open %s, %d\n", shm_name, rank); - do { - shared_open(&allreduce_buffer, shm_name, sizeof(struct allreduce_workspace)); - } while (allreduce_buffer.descriptor == -1 && errno == ENOENT); - workspace_buf_other = (struct allreduce_workspace*)allreduce_buffer.bytes; - workspace[i] = workspace_buf_other; - } else { - workspace[i] = workspace_buf; - } - symmetric_buffer[0][i] = workspace[i]->buffer + BUFFER0_OFFSET(0); - symmetric_buffer[1][i] = workspace[i]->buffer + BUFFER0_OFFSET(1); - distributed_buffer[0][i] = workspace[i]->buffer + BUFFER1_OFFSET(0); - distributed_buffer[1][i] = workspace[i]->buffer + BUFFER1_OFFSET(1); - } -} - -static void parallel_memcpy(void* to, void* from, size_t n_bytes) - __attribute__((target("avx512bw"))); -static void parallel_memcpy(void* to, void* from, size_t n_bytes) -{ - auto aligned_bytes = n_bytes - (n_bytes % VECTOR_LENGTH_IN_BYTES); - // process aligned part -#pragma omp parallel for - for (int i = 0; i < aligned_bytes; i += VECTOR_LENGTH_IN_BYTES) { - auto val = _mm256_loadu_si256((__m256i*)((char*)from + i)); - _mm256_storeu_si256((__m256i*)((char*)to + i), val); - } - - // process remaining part - for (int i = aligned_bytes; i < n_bytes; i++) { *((char*)to + i) = *((char*)from + i); } -} - -#define positive_mod(num, mod) ((((num) % (mod)) + (mod)) % (mod)) -#define rank_mod(rank) positive_mod(rank, world_size) -size_t slice_size(size_t chunk_el, int slice_idx) -{ - size_t slice_size = chunk_el / world_size; - return slice_idx == world_size - 1 ? slice_size + (chunk_el % world_size) : slice_size; -} - -char* slice_data(char* data_ptr, size_t chunk_el, int el_size, int slice_idx) -{ - size_t slice_size = chunk_el / world_size; - size_t el_offset = slice_size * slice_idx; - return data_ptr + el_offset * el_size; -} - -size_t slice_el_start(size_t chunk_el, int slice_idx) -{ - size_t slice_size = chunk_el / world_size; - return slice_size * slice_idx; -} - -/* - Symmetrical naive all_reduce - step 0: before enter the function ith times, state is copy(i-1) - step 1: each rank copy data from input (data_ptr) to SHM buffer[i] - step 2: set own state to copy(i) - step 3: wait each other rank's state equal or later than copy(i) - step 4: reduce across SHM buffer(ith) directly into output (data_ptr) -*/ -void symmetric_naive_all_reduce(char* data_ptr, - c10::ScalarType scalar_type, - size_t chunk_size, - size_t chunk_el) -{ -#ifdef DO_PROFILE - static double total_t1_t0 = 0.0; - static double total_t2_t1 = 0.0; - static double total_t3_t2 = 0.0; - static int count = -16; // warmup - auto t0 = std::chrono::system_clock::now(); -#endif - - /* - We can't have infinite number of buffers and states. 2 sets of buffer - and 3 sets of states is just enough. Consider current rank is in step 3, - with it's own state set to copy(i), the other rank will them have the - following situations: - ------------------------------------------------ - my state | can I proceed? | the other rank state - ================================================ - | N | copy(i-1) - |----------------|--------------------- - copy(i) | Y | copy(i) - |----------------|--------------------- - | Y | copy(i+1) - ------------------------------------------------ - * When I have state as copy(i), the other rank cannot have state - copy(i-2) or before. In that case I'll be in state copy(i-1) and cannot - proceed to copy(i). - * The other rank cannot have state copy(i+2) or beyond because my - state is still copy(i), copy(i+1) is as far as the other rank could go. - * From a rank's POV, all the other ranks can be divided into three sets: - - Lagging ranks: ranks that are still working on previous iteration - - Syncing ranks: ranks that are working on current iteration - - Leading ranks: ranks that are working on next iteration - * We can have 3 sets of states, one set for syncing ranks; one set for - lagging ranks; one set of leading ranks. With 3 sets of states, we can - distinguish between lagging and leading ranks. - * Note from any rank's POV, leading ranks and lagging ranks does not - appear at the same time. Either all other ranks are syncing or - lagging, or all other ranks are syncing or leading. Otherwise leading - and lagging ranks will be 2 iterations apart and this should not happen. - * So we have 2 sets of buffers, one buffer is used by current iter; - one buffer used by either lagging ranks or leading ranks. - */ - const int state_group = 0; - static int current_buffer = 0; - static int state_idx = 0; - - enum coll_state copy_current, copy_next; - - switch (state_idx) { - case 0: - copy_current = coll_allreduce_naive__copy_in_done; - copy_next = coll_alt1_allreduce_naive__copy_in_done; - break; - case 1: - copy_current = coll_alt1_allreduce_naive__copy_in_done; - copy_next = coll_alt2_allreduce_naive__copy_in_done; - break; - case 2: - copy_current = coll_alt2_allreduce_naive__copy_in_done; - copy_next = coll_allreduce_naive__copy_in_done; - break; - default: assert(!"Should not get here."); - } - state_idx = (state_idx + 1) % 3; - - parallel_memcpy(symmetric_buffer[current_buffer][world_rank], data_ptr, chunk_size); - std::atomic_thread_fence(std::memory_order_release); - workspace[world_rank]->states[state_group] = copy_current; - -#ifdef DO_PROFILE - auto t1 = std::chrono::system_clock::now(); -#endif - - for (int i = 0; i < world_size; i++) { - // wait until the other rank copy the buffer - if (i != world_rank) { wait_buffer_state_until_2(i, copy_current, copy_next, state_group); } - } -#ifdef DO_PROFILE - auto t2 = std::chrono::system_clock::now(); -#endif - - // each rank reduce the buffer independently so therre is no need for synchronization afterward - reduce_all_buffers( - 0, chunk_el, scalar_type, world_rank, data_ptr, symmetric_buffer[current_buffer]); - - // switch buffer - current_buffer = 1 - current_buffer; - -#ifdef DO_PROFILE - auto t3 = std::chrono::system_clock::now(); - - count++; - if (count > 0) { - total_t1_t0 += std::chrono::duration_cast(t1 - t0).count(); - total_t2_t1 += std::chrono::duration_cast(t2 - t1).count(); - total_t3_t2 += std::chrono::duration_cast(t3 - t2).count(); - if (world_rank == 0 && count == 1000) { - printf("symmetric_naive_all_reduce time breakdown:\n"); - printf("\tcopy input buffer: %.2f\n", total_t1_t0 / count); - printf("\twait for copy: %.2f\n", total_t2_t1 / count); - printf("\treduce: %.2f\n", total_t3_t2 / count); - } - } -#endif -} - -// naive allreduce distributed, each rank do naive reduce on its slice -void distributed_naive_reduce(char* data_ptr, - c10::ScalarType scalar_type, - size_t chunk_size, - size_t chunk_el) -{ -#ifdef DO_PROFILE - static double total_t1_t0 = 0.0; - static double total_t2_t1 = 0.0; - static double total_t3_t2 = 0.0; - static double total_t4_t3 = 0.0; - static double total_t5_t4 = 0.0; - static int count = -16; // warmup - auto t0 = std::chrono::system_clock::now(); -#endif - - const int state_group = 1; - static int current_buffer = 0; - static int state_idx = 0; - - enum coll_state copy_current, copy_next, reduce_current; - - // similar to symmetric_naive_allreduce, but here we only need two sets of - // states, because distributed naive reduce has two barriers in the algorithm - switch (state_idx) { - case 0: - copy_current = coll_allreduce_naive__copy_in_done; - reduce_current = coll_allreduce_naive__reduce_done; - copy_next = coll_alt1_allreduce_naive__copy_in_done; - break; - case 1: - copy_current = coll_alt1_allreduce_naive__copy_in_done; - reduce_current = coll_alt1_allreduce_naive__reduce_done; - copy_next = coll_allreduce_naive__copy_in_done; - break; - default: assert(!"Should not get here."); - } - state_idx = (state_idx + 1) % 2; - - int data_size = chunk_size / chunk_el; - parallel_memcpy(distributed_buffer[current_buffer][world_rank], data_ptr, chunk_size); - std::atomic_thread_fence(std::memory_order_release); - workspace[world_rank]->states[state_group] = copy_current; - -#ifdef DO_PROFILE - auto t1 = std::chrono::system_clock::now(); -#endif - - for (int i = 0; i < world_size; i++) { - // wait until all the other ranks copy the buffer - if (i != world_rank) - wait_buffer_state_until_2(i, copy_current, reduce_current, state_group); - } - -#ifdef DO_PROFILE - auto t2 = std::chrono::system_clock::now(); -#endif - - // reduce scatter - reduce_all_buffers(slice_el_start(chunk_el, world_rank), - slice_size(chunk_el, world_rank), - scalar_type, - world_rank, - distributed_buffer[current_buffer][world_rank], - distributed_buffer[current_buffer]); - std::atomic_thread_fence(std::memory_order_release); - workspace[world_rank]->states[state_group] = reduce_current; - -#ifdef DO_PROFILE - auto t3 = std::chrono::system_clock::now(); -#endif - - for (int i = 0; i < world_size; i++) { - // wait until all the other ranks reduce the buffer - if (i != world_rank) wait_buffer_state_until_2(i, reduce_current, copy_next, state_group); - } - - auto t4 = std::chrono::system_clock::now(); - - for (int i = 0; i < world_size; i++) { - int rank = (i + world_rank) % world_size; - parallel_memcpy( - slice_data(data_ptr, chunk_el, data_size, rank), - slice_data( - distributed_buffer[current_buffer][rank], chunk_el, chunk_size / chunk_el, rank), - slice_size(chunk_el, rank) * data_size); - } - - current_buffer = 1 - current_buffer; - -#ifdef DO_PROFILE - auto t5 = std::chrono::system_clock::now(); - count++; - if (count > 0) { - total_t1_t0 += std::chrono::duration_cast(t1 - t0).count(); - total_t2_t1 += std::chrono::duration_cast(t2 - t1).count(); - total_t3_t2 += std::chrono::duration_cast(t3 - t2).count(); - total_t4_t3 += std::chrono::duration_cast(t4 - t3).count(); - total_t5_t4 += std::chrono::duration_cast(t5 - t4).count(); - if (world_rank == 0 && count == 1000) { - printf("distributed_naive_reduce time breakdown:\n"); - printf("\tcopy input buffer: %.2f\n", total_t1_t0 / count); - printf("\twait for copy: %.2f\n", total_t2_t1 / count); - printf("\treduce: %.2f\n", total_t3_t2 / count); - printf("\twait for reduce finish: %.2f\n", total_t4_t3 / count); - printf("\tcopy out: %.2f\n", total_t5_t4 / count); - } - } -#endif -} - -void all_reduce_outer_loop(torch::Tensor& data, size_t numel, int data_size) -{ - for (int offset = 0; offset < data_size; offset += MAX_BUF_SIZE) { - auto data_ptr = ((char*)(data.data_ptr()) + offset); - size_t chunk_size = data_size - offset > MAX_BUF_SIZE ? MAX_BUF_SIZE : data_size - offset; - size_t chunk_el = chunk_size / (data_size / numel); - if (chunk_size < NAIVE_ALLREDUCE_THRESHOLD) - symmetric_naive_all_reduce(data_ptr, data.scalar_type(), chunk_size, chunk_el); - else - distributed_naive_reduce(data_ptr, data.scalar_type(), chunk_size, chunk_el); - } -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm.h deleted file mode 100644 index 7f73197a8caa5c4487c79c0c30ef926797bae97b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm.h +++ /dev/null @@ -1,12 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#ifndef __SHM_COLLECTIVES__ -#define __SHM_COLLECTIVES__ -#define VECTOR_LENGTH_IN_BYTES 32 -void shm_initialize(int size, int rank, char* addr_string, char* port_string); -void all_reduce_outer_loop(torch::Tensor& data, size_t numel, int data_size); -void barrier_wait(int root_idx, int num_ranks); -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm_interface.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm_interface.cpp deleted file mode 100644 index 5be5cb799a7bb47f624f5b4952f2947738864d62..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm_interface.cpp +++ /dev/null @@ -1,186 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include - -#include "shm.h" - -// #define DO_PROFILE -#ifdef DO_PROFILE -#include -#include -#endif - -// Communication settings -static int world_rank = -1; -static int world_size = -1; - -static bool is_initialized = 0; - -static bool all_ranks_local_p = false; - -void initialize(int size, int rank) -{ - if (is_initialized) return; - - // Check whether all ranks is on the same physical machine. - // If true, we will use an SHM based low latency allreduce - - auto ls_string = std::getenv("LOCAL_SIZE"); - int ls = 0; - if (ls_string != NULL) { ls = std::stoi(std::getenv("LOCAL_SIZE")); } - - if (size >= 1 && size == ls) { all_ranks_local_p = true; } - - world_size = size; - world_rank = rank; - is_initialized = 1; - - auto addr_string = std::getenv("MASTER_ADDR"); - if (addr_string == NULL) { addr_string = ""; } - auto port_string = std::getenv("MASTER_PORT"); - if (port_string == NULL) { port_string = ""; } - - if (all_ranks_local_p) { shm_initialize(size, rank, addr_string, port_string); } -} - -void inference_all_reduce_(torch::Tensor& data, int op); - -// Success - return 0 -// Fail (cannot hornor the request and need to fall back) - return -1 -void inference_all_reduce_(torch::Tensor& data, int op) -{ - assert(op == 0); -#ifdef DO_PROFILE - static double total_time = 0.0; - static double total_time_sq = 0.0; - static int count = -16; // warmup - static double max_time = 0.0; - static double min_time = DBL_MAX; - // make sure all rank reach this point before measuring time - // turn on this if you suspect each rank didn't reach here at the same time (stragger) - // if (all_ranks_local_p) { barrier_wait(0, world_size); } - auto start = std::chrono::system_clock::now(); -#endif - - auto numel = data.numel(); - - int data_size = 0; - bool data_type_fallback = false; - - switch (data.scalar_type()) { - case c10::ScalarType::BFloat16: data_size = numel * 2; break; - case c10::ScalarType::Half: data_size = numel * 2; break; - case c10::ScalarType::Float: data_size = numel * 4; break; - default: data_type_fallback = true; - } - - if (data_type_fallback) return; - - all_reduce_outer_loop(data, numel, data_size); - -#ifdef DO_PROFILE - auto end = std::chrono::system_clock::now(); - count++; - if (count > 0) { - double elapsed = std::chrono::duration_cast(end - start).count(); - if (elapsed > max_time) { max_time = elapsed; } - if (elapsed < min_time) { min_time = elapsed; } - total_time += elapsed; - total_time_sq += elapsed * elapsed; - if (world_rank == 0 && count == 1000) { - auto avg = total_time / count; - auto sd = - sqrt(total_time_sq / count - total_time * total_time / (count * count)) / avg * 100; - printf(" C++ kernel\t\t %.2f\t %.2f\t%.2f\t %.2f\n", - min_time, - max_time, - total_time / count, - sd); - } - } -#endif - return; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("initialize", &initialize, "shm initialize"); } - -TORCH_LIBRARY(deepspeed, m) -{ - m.def("inference_all_reduce(Tensor self) -> Tensor"); - m.def("inference_all_reduce_(Tensor(a!) self) -> Tensor(a!)"); -} - -torch::Tensor inference_all_reduce_meta(const torch::Tensor& self_) -{ - torch::Tensor result_ = torch::empty_like(self_); - return result_; -} - -torch::Tensor& inference_all_reduce__meta(torch::Tensor& self_) { return self_; } - -torch::Tensor& inference_all_reduce__cpu(torch::Tensor& self_) -{ - TORCH_INTERNAL_ASSERT(self_.device().type() == torch::DeviceType::CPU); - torch::Tensor self_tensor = self_.contiguous(); - inference_all_reduce_(self_tensor, 0); - return self_; -} - -torch::Tensor inference_all_reduce_cpu(const torch::Tensor& self_) -{ - torch::Tensor result = self_.clone(); - inference_all_reduce__cpu(result); - return result; -} - -#include -// The boilerplate functionalization logic, that teaches functionalization -// how to map x_() calls into x() calls. -// Long term, we'd like to not require users to write this logic. -// HOWEVER, if you have a custom op that is mutable, -// You will still need to write an out-of-place version of that op! -at::Tensor& inference_all_reduce__functionalization_glue(at::Tensor& x) -{ - // We expect all tensor inputs to our op to be "functional tensors" - TORCH_INTERNAL_ASSERT(at::functionalization::impl::isFunctionalTensor(x)); - // First, sync and unwrap and functional tensors - at::functionalization::impl::sync(x); - auto x_ = at::functionalization::impl::from_functional_tensor(x); - // Grab the dispatcher entry corresponding to the out-of-place op, "x" - static auto op_handle = c10::Dispatcher::singleton() - // specify namespace::op_name, op_overload_name - .findSchemaOrThrow("deepspeed::inference_all_reduce", "") - // Specify the C++ schema of the out-of-place op. - .typed(); - // Next, redispatch to the out-of-place op, x() (user called x_, we call x) - at::Tensor tmp_output; - { - at::AutoDispatchSkipFunctionalize guard; - tmp_output = op_handle.call(x_); - } - // Finally, tell functionalization about this mutation. - at::functionalization::impl::replace_(x, tmp_output); - at::functionalization::impl::commit_update(x); - at::functionalization::impl::sync(x); - return x; -} - -TORCH_LIBRARY_IMPL(deepspeed, CPU, m) -{ - m.impl("inference_all_reduce", inference_all_reduce_cpu); - m.impl("inference_all_reduce_", inference_all_reduce__cpu); -} - -TORCH_LIBRARY_IMPL(deepspeed, Meta, m) -{ - m.impl("inference_all_reduce", inference_all_reduce_meta); - m.impl("inference_all_reduce_", inference_all_reduce__meta); -} - -TORCH_LIBRARY_IMPL(deepspeed, Functionalize, m) -{ - m.impl("inference_all_reduce_", inference_all_reduce__functionalization_glue); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/lion/fused_lion.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/lion/fused_lion.cpp deleted file mode 100644 index 708df7f0146aa996e1652ed938c038331738e149..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/lion/fused_lion.cpp +++ /dev/null @@ -1,43 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "cpu_lion.h" - -// C++ interface - -void multi_tensor_lion(int chunk_size, - at::Tensor noop_flag, - std::vector> tensor_lists, /*gpmv*/ - const float lr, - const float beta1, - const float beta2, - const int step, - const int mode, - const float weight_decay) -{ - static bool initialized = false; - if (!initialized) { - create_lion_optimizer(0); - initialized = true; - } - for (int i = 0; i < tensor_lists[0].size(); i++) { - ds_lion_step(0, - step, - lr, - beta1, - beta2, - weight_decay, - tensor_lists[1][i], - tensor_lists[0][i], - tensor_lists[2][i]); - } -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("multi_tensor_lion", - &multi_tensor_lion, - "Compute and apply gradient update to parameters for Lion optimizer"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention.cpp deleted file mode 100644 index ac3364539ff1e47bceffcf79298f2bbc91845ac8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention.cpp +++ /dev/null @@ -1,62 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include - -void attention_impl(torch::Tensor& q, - torch::Tensor& k, - torch::Tensor& v, - torch::Tensor& bias1, - torch::Tensor& bias2, - torch::Tensor& o, - torch::Tensor& lse); -void attention(torch::Tensor& q, - torch::Tensor& k, - torch::Tensor& v, - torch::Tensor& bias1, - torch::Tensor& bias2, - torch::Tensor& o, - torch::Tensor& lse) -{ - attention_impl(q, k, v, bias1, bias2, o, lse); -} - -void attention_back_impl(torch::Tensor& go, - torch::Tensor& q, - torch::Tensor& k, - torch::Tensor& v, - torch::Tensor& o, - torch::Tensor& lse, - torch::Tensor& delta, - torch::Tensor& bias1, - torch::Tensor& bias2, - torch::Tensor& gq, - torch::Tensor& gk, - torch::Tensor& gv, - torch::Tensor& gb1, - torch::Tensor& gb2); -void attention_bwd(torch::Tensor& go, - torch::Tensor& q, - torch::Tensor& k, - torch::Tensor& v, - torch::Tensor& o, - torch::Tensor& lse, - torch::Tensor& delta, - torch::Tensor& bias1, - torch::Tensor& bias2, - torch::Tensor& gq, - torch::Tensor& gk, - torch::Tensor& gv, - torch::Tensor& gb1, - torch::Tensor& gb2) -{ - attention_back_impl(go, q, k, v, o, lse, delta, bias1, bias2, gq, gk, gv, gb1, gb2); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("attention", &attention, ""); - m.def("attention_bwd", &attention_bwd, ""); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_back.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_back.cu deleted file mode 100644 index a82c4ec68a13fccfe63c684a481afeebfe378031..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_back.cu +++ /dev/null @@ -1,218 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include -#include -#include "gemm_kernel_utils.h" -#include "kernel_backward.h" -#include "transform/bias_broadcast.h" - -constexpr auto kBlockSizeI = 64; -constexpr auto kBlockSizeJ = 64; - -template - class Broadcast1_, - template - class Broadcast2_> -typename std::enable_if::value>::type attention_back_impl_template( - torch::Tensor& go, - torch::Tensor& q, - torch::Tensor& k, - torch::Tensor& v, - torch::Tensor& o, - torch::Tensor& lse, - torch::Tensor& delta, - torch::Tensor& bias1, - torch::Tensor& bias2, - torch::Tensor& gq, - torch::Tensor& gk, - torch::Tensor& gv, - torch::Tensor& gb1, - torch::Tensor& gb2) -{ - EVOFORMER_CHECK(false, "Unsupported GPU and data type combination") -} - -template - class Broadcast1_, - template - class Broadcast2_> -typename std::enable_if::value>::type attention_back_impl_template( - torch::Tensor& go, - torch::Tensor& q, - torch::Tensor& k, - torch::Tensor& v, - torch::Tensor& o, - torch::Tensor& lse, - torch::Tensor& delta, - torch::Tensor& bias1, - torch::Tensor& bias2, - torch::Tensor& gq, - torch::Tensor& gk, - torch::Tensor& gv, - torch::Tensor& gb1, - torch::Tensor& gb2) -{ - constexpr bool kPreload_ = arch::kMinComputeCapability >= 80; - using Kernel = AttentionBackwardKernel; - int head_size = q.size(-1); - int head_number = q.size(-2); - int seq_length = q.size(-3); - auto q_view = q.view({-1, seq_length, head_number, head_size}); - auto k_view = k.view({-1, seq_length, head_number, head_size}); - auto v_view = v.view({-1, seq_length, head_number, head_size}); - auto o_view = o.view({-1, seq_length, head_number, head_size}); - auto do_view = go.view({-1, seq_length, head_number, head_size}); - auto dk_view = gk.view({-1, seq_length, head_number, head_size}); - auto dv_view = gv.view({-1, seq_length, head_number, head_size}); - auto dq_view = gq.view({-1, seq_length, head_number, head_size}); - auto q_ptr = reinterpret_cast(q.data_ptr()); - auto k_ptr = reinterpret_cast(k.data_ptr()); - auto v_ptr = reinterpret_cast(v.data_ptr()); - auto o_ptr = reinterpret_cast(o.data_ptr()); - auto do_ptr = reinterpret_cast(go.data_ptr()); - auto dk_ptr = reinterpret_cast(gk.data_ptr()); - auto dv_ptr = reinterpret_cast(gv.data_ptr()); - auto dq_ptr = reinterpret_cast(gq.data_ptr()); - auto db1_ptr = gb1.size(0) > 0 ? reinterpret_cast(gb1.data_ptr()) : nullptr; - auto db2_ptr = gb2.size(0) > 0 ? reinterpret_cast(gb2.data_ptr()) : nullptr; - auto lse_ptr = reinterpret_cast(lse.data_ptr()); - auto delta_ptr = reinterpret_cast(delta.data_ptr()); - auto bias1_ptr = reinterpret_cast(bias1.data_ptr()); - auto bias2_ptr = reinterpret_cast(bias2.data_ptr()); - static_assert(Kernel::kKernelComputesDelta, "Kernel must compute delta"); - - typename Kernel::Params p; - p.query_ptr = q_ptr; - p.key_ptr = k_ptr; - p.value_ptr = v_ptr; - p.logsumexp_ptr = lse_ptr; - p.output_ptr = o_ptr; - p.grad_output_ptr = do_ptr; - p.delta_ptr = delta_ptr; - p.grad_query_ptr = dq_ptr; - p.grad_key_ptr = dk_ptr; - p.grad_value_ptr = dv_ptr; - - p.grad_bias1_ptr = db1_ptr; - p.grad_bias2_ptr = db2_ptr; - p.B = q.size(0); - p.N = q.size(1); - p.bias1_ptr = bias1.size(0) ? bias1_ptr : nullptr; - p.bias2_ptr = bias2.size(0) ? bias2_ptr : nullptr; - - p.scale = 1.0f / sqrtf(head_size); - - p.head_dim = head_size; - p.head_dim_value = head_size; - p.num_queries = seq_length; - p.num_keys = seq_length; - p.num_heads = head_number; - - p.q_strideM = q_view.stride(-3); - p.k_strideM = k_view.stride(-3); - p.v_strideM = v_view.stride(-3); - p.gO_strideM = do_view.stride(-3); - p.o_strideH = o_view.stride(-2); - p.q_strideH = q_view.stride(-2); - p.k_strideH = k_view.stride(-2); - p.v_strideH = v_view.stride(-2); - p.o_strideB = o_view.stride(-4); - p.q_strideB = q_view.stride(-4); - p.k_strideB = k_view.stride(-4); - p.v_strideB = v_view.stride(-4); - p.lse_strideB = lse.stride(-3); - p.lse_strideH = lse.stride(-2); - p.delta_strideB = delta.stride(-3); - p.delta_strideH = delta.stride(-2); - p.num_batches = q_view.size(-4); - - p.gO_strideB = do_view.stride(-4); - p.gQ_strideB = dq_view.stride(-4); - p.gK_strideB = dk_view.stride(-4); - p.gV_strideB = dv_view.stride(-4); - p.gO_strideH = do_view.stride(-2); - p.gQ_strideH = dq_view.stride(-2); - p.gK_strideH = dk_view.stride(-2); - p.gV_strideH = dv_view.stride(-2); - - torch::Tensor workspace = torch::empty(p.workspace_size() / 4, lse.options()); - p.workspace = workspace.data_ptr(); - - auto kernel_fn = attention_kernel_backward_batched_impl; - size_t smem_bytes = sizeof(typename Kernel::SharedStorage); - cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, int(smem_bytes)); - if (!Kernel::check_supported(p)) { throw std::runtime_error("Unsupported parameters"); } - kernel_fn<<>>(p); -} - -#define CODE(scalar_t, torch_scalar_t) \ - do { \ - if (bias1.size(0) == 0 && bias2.size(0) == 0) { \ - attention_back_impl_template( \ - go, q, k, v, o, lse, delta, bias1, bias2, gq, gk, gv, gb1, gb2); \ - } else if (bias1.size(0) > 0 && bias2.size(0) > 0) { \ - attention_back_impl_template( \ - go, q, k, v, o, lse, delta, bias1, bias2, gq, gk, gv, gb1, gb2); \ - } else if (bias1.size(0) > 0) { \ - attention_back_impl_template( \ - go, q, k, v, o, lse, delta, bias1, bias2, gq, gk, gv, gb1, gb2); \ - } else { \ - attention_back_impl_template( \ - go, q, k, v, o, lse, delta, bias1, bias2, gq, gk, gv, gb1, gb2); \ - } \ - } while (0) - -void attention_back_impl(torch::Tensor& go, - torch::Tensor& q, - torch::Tensor& k, - torch::Tensor& v, - torch::Tensor& o, - torch::Tensor& lse, - torch::Tensor& delta, - torch::Tensor& bias1, - torch::Tensor& bias2, - torch::Tensor& gq, - torch::Tensor& gk, - torch::Tensor& gv, - torch::Tensor& gb1, - torch::Tensor& gb2) -{ - cudaDeviceProp* prop = at::cuda::getCurrentDeviceProperties(); - DISPATCH_ARCHTAG(prop->major * 10 + prop->minor, - DISPATCH_TYPES(q, { CODE(scalar_t, torch_scalar_t); })); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_cu.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_cu.cu deleted file mode 100644 index 37636c4bf98842543fe1bd055a848cb52d0a6449..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/attention_cu.cu +++ /dev/null @@ -1,160 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include -#include "gemm_kernel_utils.h" -#include "kernel_forward.h" -#include "transform/bias_broadcast.h" - -template - class Broadcast1_, - template - class Broadcast2_> -typename std::enable_if::value>::type attention_impl_template( - torch::Tensor& q, - torch::Tensor& k, - torch::Tensor& v, - torch::Tensor& bias1, - torch::Tensor& bias2, - torch::Tensor& o, - float* lse_ptr) -{ - EVOFORMER_CHECK(false, "Unsupported GPU and data type combination") -} - -template - class Broadcast1_, - template - class Broadcast2_> -typename std::enable_if::value>::type attention_impl_template( - torch::Tensor& q, - torch::Tensor& k, - torch::Tensor& v, - torch::Tensor& bias1, - torch::Tensor& bias2, - torch::Tensor& o, - float* lse_ptr) -{ - // Attention definition goes here, replaced with BroadcastType1 and - // BroadcastType2 - using Attention = AttentionKernel; - - static_assert(!Attention::kNeedsOutputAccumulatorBuffer, - "This test does not support output accumulator buffer"); - int head_size = q.size(-1); - int head_number = q.size(-2); - int seq_length = q.size(-3); - auto q_view = q.view({-1, seq_length, head_number, head_size}); - auto k_view = k.view({-1, seq_length, head_number, head_size}); - auto v_view = v.view({-1, seq_length, head_number, head_size}); - auto o_view = o.view({-1, seq_length, head_number, head_size}); - int batch_size = q_view.size(0); - auto q_ptr = reinterpret_cast(q.data_ptr()); - auto k_ptr = reinterpret_cast(k.data_ptr()); - auto v_ptr = reinterpret_cast(v.data_ptr()); - auto o_ptr = reinterpret_cast(o.data_ptr()); - - auto bias1_ptr = reinterpret_cast(bias1.data_ptr()); - auto bias2_ptr = reinterpret_cast(bias2.data_ptr()); - - typename Attention::Params p; - { // set parameters - p.query_ptr = q_ptr; - p.key_ptr = k_ptr; - p.value_ptr = v_ptr; - p.logsumexp_ptr = lse_ptr; // Only needed for bw - p.output_accum_ptr = nullptr; - p.output_ptr = o_ptr; - p.scale = 1.0f / sqrt(float(head_size)); - - p.bias1_ptr = bias1_ptr; - p.bias2_ptr = bias2_ptr; - p.B = q.size(0); - p.N = q.size(1); - - p.num_heads = head_number; - p.num_batches = batch_size; - p.head_dim = head_size; - p.head_dim_value = head_size; - p.num_queries = seq_length; - p.num_keys = seq_length; - - // All tensors are in BMHK shapes - p.q_strideH = q_view.stride(-2); - p.k_strideH = k_view.stride(-2); - p.v_strideH = v_view.stride(-2); - p.q_strideM = q_view.stride(-3); - p.k_strideM = k_view.stride(-3); - p.v_strideM = v_view.stride(-3); - p.o_strideM = o_view.stride(-3); - p.q_strideB = q_view.stride(-4); - p.k_strideB = k_view.stride(-4); - p.v_strideB = v_view.stride(-4); - } - - constexpr auto kernel_fn = attention_kernel_batched_impl; - int smem_bytes = sizeof(typename Attention::SharedStorage); - if (smem_bytes > 0xc000) { - cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); - } - if (!Attention::check_supported(p)) { throw std::runtime_error("Parameters not supported"); } - kernel_fn<<>>(p); -} - -#define CODE(scalar_t, torch_scalar_t) \ - do { \ - if (bias1.size(0) == 0 && bias2.size(0) == 0) { \ - attention_impl_template(q, k, v, bias1, bias2, o, lse_ptr); \ - } else if (bias1.size(0) == 0) { \ - attention_impl_template(q, k, v, bias1, bias2, o, lse_ptr); \ - } else if (bias2.size(0) == 0) { \ - attention_impl_template(q, k, v, bias1, bias2, o, lse_ptr); \ - } else { \ - attention_impl_template( \ - q, k, v, bias1, bias2, o, lse_ptr); \ - } \ - } while (0) - -// Function to select and call the correct template based on biases sizes -void attention_impl(torch::Tensor& q, - torch::Tensor& k, - torch::Tensor& v, - torch::Tensor& bias1, - torch::Tensor& bias2, - torch::Tensor& o, - torch::Tensor& lse) -{ - auto lse_ptr = lse.size(0) == 0 ? nullptr : reinterpret_cast(lse.data_ptr()); - cudaDeviceProp* prop = at::cuda::getCurrentDeviceProperties(); - DISPATCH_ARCHTAG(prop->major * 10 + prop->minor, - DISPATCH_TYPES(q, { CODE(scalar_t, torch_scalar_t); })); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_grad_bias.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_grad_bias.h deleted file mode 100644 index 17b6479ed8c53711a4b0c5fd7e23bee44ff9e142..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_grad_bias.h +++ /dev/null @@ -1,250 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once -#include -#include -#include "../iterators/predicated_tile_iterator_atomic.h" -#include "cutlass/epilogue/threadblock/epilogue.h" - -namespace cutlass { -namespace epilogue { -namespace threadblock { -template -struct EpilogueTensorOpAffineRankN : public DefaultEpilogueTensorOpAffineRankN { - using Base = DefaultEpilogueTensorOpAffineRankN; - using OutputTileIterator = - cutlass::epilogue::threadblock::PredicatedTileIteratorAffineRankNAtomic< - typename Base::OutputTileThreadMap, - typename Base::ElementOutput, - Rank>; - - using Epilogue = - cutlass::epilogue::threadblock::Epilogue; -}; - -template -struct EpilogueVoltaTensorOpAffineRankN - : public DefaultEpilogueVoltaTensorOpAffineRankN { - using Base = DefaultEpilogueVoltaTensorOpAffineRankN; - using OutputTileIterator = - cutlass::epilogue::threadblock::PredicatedTileIteratorAffineRankNAtomic< - typename Base::OutputTileThreadMap, - typename Base::ElementOutput, - Rank>; - - using Epilogue = - cutlass::epilogue::threadblock::Epilogue; -}; - -template -struct EpilogueTensorOp : public DefaultEpilogueTensorOp { - using Base = DefaultEpilogueTensorOp; - using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIteratorAtomic< - typename Base::OutputTileThreadMap, - typename Base::ElementOutput, - ScatterD, - PermuteDLayout>; - using Epilogue = - cutlass::epilogue::threadblock::Epilogue; -}; - -template -struct EpilogueVoltaTensorOp : public DefaultEpilogueVoltaTensorOp { - using Base = DefaultEpilogueVoltaTensorOp; - using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIteratorAtomic< - typename Base::OutputTileThreadMap, - typename Base::ElementOutput, - ScatterD, - PermuteDLayout>; - using Epilogue = - cutlass::epilogue::threadblock::Epilogue; -}; -} // namespace threadblock -} // namespace epilogue -} // namespace cutlass - -template -struct BiasGradEpilogue { - using Epilogue = - typename cutlass::epilogue::threadblock::EpilogueTensorOp::Epilogue; -}; - -template -struct BiasGradEpilogue { - using Epilogue = - typename cutlass::epilogue::threadblock::EpilogueVoltaTensorOp::Epilogue; -}; - -template -struct BiasGradEpilogueAffineRankN { - using Epilogue = typename cutlass::epilogue::threadblock::EpilogueTensorOpAffineRankN< - Rank, - Shape_, - WarpMmaTensorOp_, - PartitionsK, - OutputOp_, - ElementsPerAccess>::Epilogue; -}; - -template -struct BiasGradEpilogueAffineRankN { - using Epilogue = typename cutlass::epilogue::threadblock::EpilogueVoltaTensorOpAffineRankN< - Rank, - Shape_, - WarpMmaTensorOp_, - PartitionsK, - OutputOp_, - ElementsPerAccess>::Epilogue; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_pipelined.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_pipelined.h deleted file mode 100644 index 3b7b32d6145202ea59d707b025669cbe75a8bc51..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_pipelined.h +++ /dev/null @@ -1,592 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holdvr nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/*! \file - \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. - - File copied from "cutlass/epilogue/threadblock/epilogue.h" - then modified to: - (1) load 2 source fragments at the same time (pipelining) - (2) support reading from a different dtype - (3) pass the row id to the OutputOp if it takes it - (see MemoryEfficientAttentionNormalize) - Note that in general the fragment passed to the OutputOp could - span multiple rows but it does not happen with the configurations we have -*/ - -#pragma once - -#if defined(__CUDACC_RTC__) -#include -#else -#include -#endif - -#include "cutlass/aligned_buffer.h" -#include "cutlass/array.h" -#include "cutlass/cutlass.h" -#include "cutlass/functional.h" -#include "cutlass/layout/tensor.h" -#include "cutlass/layout/vector.h" -#include "cutlass/numeric_types.h" -#include "cutlass/tensor_coord.h" - -#include "cutlass/gemm/gemm.h" - -#include "cutlass/transform/pitch_linear_thread_map.h" -#include "cutlass/transform/threadblock/regular_tile_iterator.h" - -#include "cutlass/epilogue/threadblock/epilogue_base.h" -#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h" -#include "cutlass/numeric_types.h" - -//////////////////////////////////////////////////////////////////////////////// - -namespace cutlass { -namespace epilogue { -namespace threadblock { - -template -struct ApplyEpilogueOp { - static CUTLASS_DEVICE typename Op::FragmentOutput apply( - Op const& output_op, - int row_id, - typename Op::FragmentAccumulator const& accum, - typename Op::FragmentOutput const& source) - { - return output_op(accum, source); - } - static CUTLASS_DEVICE typename Op::FragmentOutput - apply(Op const& output_op, int row_id, typename Op::FragmentAccumulator const& accum) - { - return output_op(accum); - } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Epilogue operator -template ::value), - typename OutputTileSourceIterator_ = - OutputTileIterator_ ///< Tile iterator reading tensors - > -class EpiloguePipelined : public EpilogueBase { -public: - using Base = EpilogueBase; - - using Shape = Shape_; - using WarpMmaOperator = WarpMmaOperator_; - static int const kPartitionsK = PartitionsK; - using OutputTileIterator = OutputTileIterator_; - using OutputTileSourceIterator = OutputTileSourceIterator_; - using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; - using WarpTileIterator = WarpTileIterator_; - using SharedLoadIterator = SharedLoadIterator_; - using OutputOp = OutputOp_; - using Padding = Padding_; - - using Layout = layout::RowMajor; - using LongIndex = typename Layout::LongIndex; - - /// The complete warp-level accumulator tile - using AccumulatorTile = typename Base::AccumulatorTile; - - /// Accumulator element - using ElementAccumulator = typename WarpTileIterator::Element; - - /// Output element - using ElementOutput = typename OutputTileIterator::Element; - using ElementSource = typename OutputTileSourceIterator::Element; - - /// Output access size - static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; - - /// Tensor reference to destination tensor - using TensorRef = typename OutputTileIterator::TensorRef; - - /// Tensor reference to sync tensor - using SyncTensorRef = typename cutlass::TensorRef; - - /// Const tensor reference to source tensor - using ConstTensorRef = typename OutputTileIterator::ConstTensorRef; - - /// Array type used to output - using OutputAccessType = - Array; - using SourceAccessType = Array; - - /// Array type used by output functor - using AccumulatorAccessType = - Array; - - /// Number of warps - using WarpCount = typename Base::WarpCount; - - static int constexpr kSmemTiles = - Base::kFragmentsPerIteration > 1 ? Base::kFragmentsPerIteration : kPartitionsK; - static int constexpr kSmemPointerOffset = - Base::SharedStorage::StorageShape::kCount / kSmemTiles; - -public: - static_assert(OutputTileSourceIterator::Fragment::kElements == - OutputTileIterator::Fragment::kElements, - "Mismatch between input tile and output tile iterator (kElements)"); - static_assert(OutputTileSourceIterator::kIterations == OutputTileIterator::kIterations, - "Mismatch between input tile and output tile iterator (kIterations)"); - static_assert(SharedLoadIterator::Fragment::kElements == - OutputTileIterator::Fragment::kElements, - "Mismatch between shared load iterator and output tile iterator."); - - static_assert(OutputTileIterator::kElementsPerAccess, - "OutputTileIterator::kElementsPerAccess must not be zero."); - - static_assert(!(OutputTileIterator::Fragment::kElements % - OutputTileIterator::kElementsPerAccess), - "Divisibility"); - -private: - /// Loads fragment from shared memory aligned with output tensor - SharedLoadIterator shared_load_iterator_; - -public: - /// Constructor - CUTLASS_DEVICE - EpiloguePipelined(typename Base::SharedStorage& shared_storage, ///< Shared storage object - int thread_idx, ///< ID of a thread within the threadblock - int warp_idx, ///< ID of warp within threadblock - int lane_idx ///< Id of thread within warp - ) - : Base(shared_storage, thread_idx, warp_idx, lane_idx), - shared_load_iterator_(shared_storage.reference(), thread_idx) - { - } - - /// Streams the result to global memory - CUTLASS_DEVICE - void operator()(OutputOp const& output_op, ///< Output operator - OutputTileIterator destination_iterator, ///< Tile iterator for destination - AccumulatorTile const& accumulators, ///< Complete warp-level accumulator tile - OutputTileSourceIterator source_iterator) - { ///< Threadblock tile coordinate in GEMM (in units - ///< of threadblock tiles) - - if (!output_op.is_source_needed()) { - compute_source_not_needed_(output_op, destination_iterator, accumulators); - } else { - compute_source_needed_(output_op, destination_iterator, accumulators, source_iterator); - } - } - CUTLASS_DEVICE - void operator()(OutputOp const& output_op, ///< Output operator - OutputTileIterator destination_iterator, ///< Tile iterator for destination - AccumulatorTile const& accumulators) - { ///< Complete warp-level accumulator tile - compute_source_not_needed_(output_op, destination_iterator, accumulators); - } - -private: - template - struct acc2smem_source_not_needed; - - template - struct acc2smem_source_not_needed> { - template - CUTLASS_DEVICE static void helper(AccumulatorFragmentIterator accum_fragment_iterator, - WarpTileIterator& warp_tile_iterator) - { - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < Advance; i++) { ++accum_fragment_iterator; } - - CUTLASS_PRAGMA_UNROLL - for (int p = 0; p < Base::kFragmentsPerIteration; ++p) { - typename AccumulatorFragmentIterator::Fragment accum_fragment; - - accum_fragment_iterator.load(accum_fragment); - ++accum_fragment_iterator; - - warp_tile_iterator.store(accum_fragment); - if (p < Base::kFragmentsPerIteration - 1) { - warp_tile_iterator.add_pointer_offset(kSmemPointerOffset); - } - } - - if (Base::kFragmentsPerIteration > 1) { - warp_tile_iterator.add_pointer_offset(kSmemPointerOffset * - (1 - Base::kFragmentsPerIteration)); - } - } - - CUTLASS_DEVICE - static void push(size_t pos, - AccumulatorFragmentIterator const& iterator_begin, - WarpTileIterator& warp_tile_iterator) - { - int dummy[] = { - (pos == (Seq * Base::kFragmentsPerIteration)) && - (helper(iterator_begin, warp_tile_iterator), - 0)...}; - - CUTLASS_UNUSED(dummy[0]); - } - }; - - static_assert(kPartitionsK == 1 || Base::kFragmentsPerIteration == 1, - "One of these must be exactly 1."); - - /// Streams the result to global memory - CUTLASS_DEVICE - void compute_source_not_needed_( - OutputOp const& output_op, ///< Output operator - OutputTileIterator destination_iterator, ///< Tile iterator for destination - AccumulatorTile const& accumulators ///< Complete warp-level accumulator tile - ) - { - // - // Iterator over warp-level accumulator fragment - // - - AccumulatorFragmentIterator accum_fragment_iterator(accumulators); - - // - // Iterate over accumulator tile - // - -#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations / Base::kFragmentsPerIteration \ - : 1) - for (int iter = 0; iter < OutputTileIterator::kIterations; - iter += Base::kFragmentsPerIteration) { - // - // Convert and store fragment - // - - __syncthreads(); - - acc2smem_source_not_needed>:: - push(iter, accum_fragment_iterator, this->warp_tile_iterator_); - - __syncthreads(); - - // - // Load fragments from shared memory - // - - CUTLASS_PRAGMA_UNROLL - for (int p = 0; p < Base::kFragmentsPerIteration; ++p) { - typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK]; - - shared_load_iterator_.load(aligned_accum_fragment[0]); - - if (p < Base::kFragmentsPerIteration - 1) { - shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); - } else if (kPartitionsK > 1) { - plus add_fragments; - - CUTLASS_PRAGMA_UNROLL - for (int i = 1; i < kPartitionsK; ++i) { - shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); - shared_load_iterator_.load(aligned_accum_fragment[i]); - aligned_accum_fragment[0] = - add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]); - } - - shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) * - kSmemPointerOffset); - } - - // - // Compute the output result - // - - typename OutputTileIterator::Fragment output_fragment; - - apply_output_operator_source_not_needed_(destination_iterator.thread_start_row(), - output_fragment, - output_op, - aligned_accum_fragment[0]); - - // - // Store the final result - // - - destination_iterator.store(output_fragment); - ++destination_iterator; - } - - if (Base::kFragmentsPerIteration > 1) { - shared_load_iterator_.add_pointer_offset(kSmemPointerOffset * - (1 - Base::kFragmentsPerIteration)); - } - } - } - - template - struct acc2smem_source_needed; - - template - struct acc2smem_source_needed> { - template - CUTLASS_DEVICE static void helper(AccumulatorFragmentIterator accum_fragment_iterator, - WarpTileIterator& warp_tile_iterator) - { - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < Advance; i++) { ++accum_fragment_iterator; } - - typename AccumulatorFragmentIterator::Fragment accum_fragment; - accum_fragment_iterator.load(accum_fragment); - warp_tile_iterator.store(accum_fragment); - } - - CUTLASS_DEVICE - static void push(size_t pos, - AccumulatorFragmentIterator const& iterator_begin, - WarpTileIterator& warp_tile_iterator) - { - int dummy[] = {(pos == Seq) && (helper(iterator_begin, warp_tile_iterator), 0)...}; - } - }; - - /// Streams the result to global memory - CUTLASS_DEVICE - void compute_source_needed_( - OutputOp const& output_op, ///< Output operator - OutputTileIterator destination_iterator, ///< Tile iterator for destination - AccumulatorTile const& accumulators, ///< Complete warp-level accumulator tile - OutputTileSourceIterator source_iterator ///< Threadblock tile coordinate in GEMM (in units - ///< of threadblock tiles) - ) - { - typename OutputTileSourceIterator::Fragment source_fragment[2]; - - source_fragment[0].clear(); - source_iterator.load(source_fragment[0]); - ++source_iterator; - source_fragment[1].clear(); - - // - // Iterator over warp-level accumulator fragment - // - - AccumulatorFragmentIterator accum_fragment_iterator(accumulators); - - // - // Iterate over accumulator tile - // - -#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1) - for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) { - if (iter > 0) { __syncthreads(); } - // - // Load the source for next iteration (pipelining) - // - - if (iter + 1 < OutputTileIterator::kIterations) { - source_iterator.load(source_fragment[(iter + 1) % 2]); - } - ++source_iterator; - acc2smem_source_needed>:: - push(iter, accum_fragment_iterator, this->warp_tile_iterator_); - - __syncthreads(); - - // - // Load fragments from shared memory - // - - typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK]; - - shared_load_iterator_.load(aligned_accum_fragment[0]); - - // If the number of k-slices is > 1 - perform a reduction amongst the - // k-slices - if (kPartitionsK > 1) { - plus add_fragments; - - CUTLASS_PRAGMA_UNROLL - for (int i = 1; i < kPartitionsK; ++i) { - shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); - shared_load_iterator_.load(aligned_accum_fragment[i]); - aligned_accum_fragment[0] = - add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]); - } - - shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) * kSmemPointerOffset); - } - - // - // Compute the output result - // - - typename OutputTileIterator::Fragment output_fragment; - - apply_output_operator_(destination_iterator.thread_start_row(), - output_fragment, - output_op, - aligned_accum_fragment[0], - source_fragment[iter % 2]); - - // - // Store the final result - // - - destination_iterator.store(output_fragment); - ++destination_iterator; - } - } - - /// Helper to invoke the output functor over each vector of output - CUTLASS_DEVICE - void apply_output_operator_(int begin_row, - typename OutputTileIterator::Fragment& output_fragment, - OutputOp const& output_op, ///< Output operator - typename SharedLoadIterator::Fragment const& aligned_accum_fragment, - typename OutputTileSourceIterator::Fragment const& source_fragment) - { - OutputAccessType* output_frag_ptr = reinterpret_cast(&output_fragment); - - AccumulatorAccessType const* compute_frag_ptr = - reinterpret_cast(&aligned_accum_fragment); - - SourceAccessType const* source_frag_ptr = - reinterpret_cast(&source_fragment); - - int const kOutputOpIterations = - OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kOutputOpIterations; ++i) { - // Call the output operator - output_frag_ptr[i] = ApplyEpilogueOp::apply( - output_op, - begin_row + getRowOffset(i * OutputTileIterator::kElementsPerAccess), - compute_frag_ptr[i], - source_frag_ptr[i]); - } - } - - /// Helper to invoke the output functor over each vector of output - CUTLASS_DEVICE - void apply_output_operator_source_not_needed_( - int begin_row, - typename OutputTileIterator::Fragment& output_fragment, - OutputOp const& output_op, ///< Output operator - typename SharedLoadIterator::Fragment const& aligned_accum_fragment) - { - OutputAccessType* output_frag_ptr = reinterpret_cast(&output_fragment); - - AccumulatorAccessType const* compute_frag_ptr = - reinterpret_cast(&aligned_accum_fragment); - - int const kOutputOpIterations = - OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kOutputOpIterations; ++i) { - // Call the output operator - output_frag_ptr[i] = ApplyEpilogueOp::apply( - output_op, - begin_row + getRowOffset(i * OutputTileIterator::kElementsPerAccess), - compute_frag_ptr[i]); - } - } - - // This should be constexpr, but it's only supported on c++14 - static int CUTLASS_HOST_DEVICE getRowOffset(int i) - { - using ThreadMap = typename OutputTileIterator::ThreadMap; - - CUTLASS_PRAGMA_UNROLL - for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { - CUTLASS_PRAGMA_UNROLL - for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { - CUTLASS_PRAGMA_UNROLL - for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { - int row_offset = row * ThreadMap::Delta::kRow + - group * ThreadMap::Delta::kGroup + - cluster * ThreadMap::Delta::kCluster; - int frag_row_idx = - (row + ThreadMap::Iterations::kRow * - (group + ThreadMap::Iterations::kGroup * cluster)); - CUTLASS_PRAGMA_UNROLL - for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { - int frag_idx = ThreadMap::kElementsPerAccess * - (frag_row_idx * ThreadMap::Iterations::kColumn + column); - if (i < frag_idx + ThreadMap::kElementsPerAccess) { return row_offset; } - } - } - } - } - return -1; - } -}; - -//////////////////////////////////////////////////////////////////////////////// - -} // namespace threadblock -} // namespace epilogue -} // namespace cutlass - -//////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_rescale_output.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_rescale_output.h deleted file mode 100644 index f81a09f74f1ebed8e3d46fee52ac220efd98214b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_rescale_output.h +++ /dev/null @@ -1,251 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holdvr nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/*! \file - \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. - - The epilogue rearranges the result of a matrix product through shared memory - to match canonical tensor layouts in global memory. Epilogues support - conversion and reduction operations. - - This is a copy of cutlass/epilogue/threadblock/epilogue.h that can - handle "row_id" as a first argument, as uses it to get the corresponding - `m_prime` / `s_prime` to rescale the output. -*/ - -#pragma once - -#if defined(__CUDACC_RTC__) -#include -#else -#include -#endif - -#include "cutlass/aligned_buffer.h" -#include "cutlass/array.h" -#include "cutlass/cutlass.h" -#include "cutlass/functional.h" -#include "cutlass/layout/tensor.h" -#include "cutlass/layout/vector.h" -#include "cutlass/numeric_types.h" -#include "cutlass/tensor_coord.h" - -#include "cutlass/gemm/gemm.h" - -#include "cutlass/transform/pitch_linear_thread_map.h" -#include "cutlass/transform/threadblock/regular_tile_iterator.h" - -#include "cutlass/epilogue/threadblock/epilogue_base.h" -#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h" -#include "cutlass/numeric_types.h" - -#include "cutlass/array.h" -#include "cutlass/cutlass.h" -#include "cutlass/epilogue/thread/scale_type.h" -#include "cutlass/functional.h" -#include "cutlass/numeric_conversion.h" -#include "cutlass/numeric_types.h" -#include "epilogue_pipelined.h" - -///////////////////////////////////////////////////////////////////////////////////////////////// - -namespace cutlass { -namespace epilogue { -namespace thread { - -///////////////////////////////////////////////////////////////////////////////////////////////// - -/// Applies a linear combination operator to an array of elements. -// output <- alpha * accumulator + beta * source -// with: -// alpha = 1 / s_prime (to normalize when isLast=True, 1 otherwise) -// beta = alpha / m_prime (renormalize the output when the max changes) -// source is the current output -template , - ///< but we use 64 or 32 sometimes when there are not enough data - ///< to store - typename ElementAccumulator_, ///< Accumulator data type - typename ElementCompute_, ///< Data type used to compute linear combination - bool isFirst, - bool isLast, - typename FragmentAlphaBeta_, - FloatRoundStyle Round = FloatRoundStyle::round_to_nearest> -class MemoryEfficientAttentionNormalize { -public: - using ElementOutput = ElementOutput_; - using ElementSource = ElementSource_; - using ElementAccumulator = ElementAccumulator_; - using ElementCompute = ElementCompute_; - - static int const kCount = Count; - - using FragmentOutput = Array; - using FragmentSource = Array; - using FragmentAccumulator = Array; - using ComputeFragment = Array; - using FragmentAlphaBeta = FragmentAlphaBeta_; - - static FloatRoundStyle const kRound = Round; - -private: - // - // Data members - // - - FragmentAlphaBeta const& s_prime_; - FragmentAlphaBeta const& m_prime_; - -public: - /// Constructs the function object, possibly loading from pointers in host - /// memory - CUTLASS_HOST_DEVICE - MemoryEfficientAttentionNormalize(FragmentAlphaBeta const& s_prime, - FragmentAlphaBeta const& m_prime) - : s_prime_(s_prime), m_prime_(m_prime) - { - } - - /// Returns true if source is needed - CUTLASS_HOST_DEVICE - bool is_source_needed() const { return !isFirst; } - - /// Functionally required for serial reduction in the epilogue - CUTLASS_HOST_DEVICE - void set_k_partition(int k_partition, int k_partition_count) {} - - /// Computes linear scaling: D = alpha * accumulator + beta * source - CUTLASS_HOST_DEVICE - FragmentOutput operator()(int row, - FragmentAccumulator const& accumulator, - FragmentSource const& source) const - { - assert(!isFirst); - - // Convert source to internal compute numeric type - NumericArrayConverter source_converter; - NumericArrayConverter - accumulator_converter; - - // Convert to destination numeric type - NumericArrayConverter destination_converter; - - ComputeFragment converted_source = source_converter(source); - ComputeFragment converted_accumulator = accumulator_converter(accumulator); - - // Perform binary operations - ComputeFragment intermediate; - - multiplies mul_add_source; - multiply_add mul_add_accumulator; - - ElementCompute alpha = isLast ? (1 / s_prime_[row]) : 1; - ElementCompute beta = alpha * m_prime_[row]; - - intermediate = mul_add_source(beta, converted_source); // X = beta * C - - intermediate = mul_add_accumulator( - alpha, converted_accumulator, intermediate); // D = alpha * Accum + X - - return destination_converter(intermediate); - } - - /// Computes linear scaling: D = alpha * accumulator - CUTLASS_HOST_DEVICE - FragmentOutput operator()(int row, FragmentAccumulator const& accumulator) const - { - assert(isFirst); - - // Convert source to internal compute numeric type - NumericArrayConverter - accumulator_converter; - - // Convert to destination numeric type - NumericArrayConverter destination_converter; - - ComputeFragment converted_accumulator = accumulator_converter(accumulator); - - ComputeFragment intermediate; - multiplies mul_accumulator; - - ElementCompute alpha = isLast ? (1 / s_prime_[row]) : 1; - - intermediate = mul_accumulator(alpha, converted_accumulator); // X = alpha * C + uniform - - return destination_converter(intermediate); - } -}; - -} // namespace thread - -namespace threadblock { -template -struct ApplyEpilogueOp< - thread::MemoryEfficientAttentionNormalize> { - using Op = thread::MemoryEfficientAttentionNormalize; - static CUTLASS_DEVICE typename Op::FragmentOutput apply( - Op const& output_op, - int row_id, - typename Op::FragmentAccumulator const& accum, - typename Op::FragmentSource const& source) - { - return output_op(row_id, accum, source); - } - static CUTLASS_DEVICE typename Op::FragmentOutput - apply(Op const& output_op, int row_id, typename Op::FragmentAccumulator const& accum) - { - return output_op(row_id, accum); - } -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace threadblock -} // namespace epilogue -} // namespace cutlass - -///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_thread_apply_logsumexp.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_thread_apply_logsumexp.h deleted file mode 100644 index 46fb2bf17c1c1c20400244593dd8a98a91973d0d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/epilogue/epilogue_thread_apply_logsumexp.h +++ /dev/null @@ -1,168 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights - *reserved. SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, - *this list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE - *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE - *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR - *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF - *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS - *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN - *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) - *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE - *POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/*! \file - \brief Functor performing linear combination operations used by epilogues. -*/ - -#pragma once - -#include - -#include "cutlass/array.h" -#include "cutlass/cutlass.h" -#include "cutlass/epilogue/thread/activation.h" -#include "cutlass/functional.h" -#include "cutlass/numeric_conversion.h" -#include "cutlass/numeric_types.h" - -///////////////////////////////////////////////////////////////////////////////////////////////// - -namespace cutlass { -namespace epilogue { -namespace thread { - -///////////////////////////////////////////////////////////////////////////////////////////////// - -namespace detail { - -template -struct ArrayExponential { - CUTLASS_HOST_DEVICE - Array operator()( - Array const& input) const - { - Array result; - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < ElementsPerAccess; ++i) { result[i] = expf(input[i]); } - - return result; - } -}; - -template -struct ArrayExponential { - CUTLASS_DEVICE - Array operator()(Array const& input) const - { - Array result; - - int const kVectorCount = ElementsPerAccess / 2; - - __half2 const* input_ptr = reinterpret_cast<__half2 const*>(input.raw_data()); - __half2* res_ptr = reinterpret_cast<__half2*>(result.raw_data()); - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kVectorCount; ++i) { res_ptr[i] = h2exp(input_ptr[i]); } - - return result; - } -}; -} // namespace detail - -///////////////////////////////////////////////////////////////////////////////////////////////// - -/// Applies: -/// output <- (input - lse).exp() -template -class ApplyLogSumExp { -public: - using ElementOutput = ElementOutput_; - using ElementAccumulator = ElementAccumulator_; - using ElementCompute = ElementCompute_; - using ElementLSE = ElementLSE_; - - static int const kElementsPerAccess = ElementsPerAccess; - static int const kCount = kElementsPerAccess; - static const ScaleType::Kind kScale = cutlass::epilogue::thread::ScaleType::NoBetaScaling; - - using FragmentOutput = Array; - using FragmentAccumulator = Array; - using FragmentCompute = Array; - using FragmentLSE = Array; - using FragmentScaleBias = FragmentLSE; // Used by epilogue_smem_accumulator.h - -public: - // - // Methods - // - - CUTLASS_HOST_DEVICE - ApplyLogSumExp() {} - - /// Returns true if source is needed - CUTLASS_HOST_DEVICE - bool is_source_needed() const { return true; } - - /// Functionally required for serial reduction in the epilogue - CUTLASS_HOST_DEVICE - void set_k_partition(int k_partition, int k_partition_count) {} - - CUTLASS_HOST_DEVICE - FragmentOutput operator()(FragmentAccumulator const& AB, - FragmentLSE const& scale_unused, - // bias used as LSE - FragmentLSE const& bias) const - { - FragmentCompute frag_AB = - NumericArrayConverter()(AB); - FragmentCompute frag_lse_compute = - NumericArrayConverter()(bias); - FragmentCompute frag_compute; - - minus minus_lse; - detail::ArrayExponential apply_exp; - frag_compute = minus_lse(frag_AB, frag_lse_compute); - frag_compute = apply_exp(frag_compute); - - return NumericArrayConverter()( - frag_compute); - } -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace thread -} // namespace epilogue -} // namespace cutlass - -///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma.h deleted file mode 100644 index 75833bbfe7d2f3c4a91db1ddc7f4453b270db968..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma.h +++ /dev/null @@ -1,119 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holdvr nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include "custom_mma_multistage.h" -#include "custom_mma_pipelined.h" -#include "cutlass/gemm/threadblock/mma_multistage.h" -#include "cutlass/gemm/threadblock/mma_pipelined.h" - -template -struct MakeCustomMma; - -template -struct MakeCustomMma, - kMaxK> { - // Reduce the number of stages if we don't need that many - static int constexpr kStages = - kMaxK == cutlass::platform::numeric_limits::max() - ? Stages - : cutlass::const_min(Stages, (kMaxK + int(Shape::kK) - 1) / int(Shape::kK)); - using Mma = cutlass::gemm::threadblock::CustomMmaMultistage; -}; - -template -struct MakeCustomMma, - kMaxK> { - using Mma = cutlass::gemm::threadblock::CustomMmaPipelined; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_base.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_base.h deleted file mode 100644 index bbf91240b90075da6f711dba55cf27e19506f898..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_base.h +++ /dev/null @@ -1,181 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights - *reserved. SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, - *this list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE - *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE - *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR - *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF - *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS - *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN - *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) - *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE - *POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/*! \file - \brief Template for a double-buffered threadblock-scoped GEMM kernel. -*/ - -#pragma once - -#include "cutlass/aligned_buffer.h" -#include "cutlass/arch/memory.h" -#include "cutlass/array.h" -#include "cutlass/cutlass.h" -#include "cutlass/gemm/gemm.h" -#include "cutlass/gemm/threadblock/mma_base.h" -#include "cutlass/matrix_shape.h" -#include "cutlass/numeric_types.h" - -//////////////////////////////////////////////////////////////////////////////// - -namespace cutlass { -namespace gemm { -namespace threadblock { - -//////////////////////////////////////////////////////////////////////////////// - -/// Structure to compute the matrix product targeting CUDA cores and SIMT math -/// instructions. -template < - /// Size of the Gemm problem - concept: gemm::GemmShape<> - typename Shape_, - /// Policy describing tuning details (concept: MmaPolicy) - typename Policy_, - /// Number of stages, - int Stages, - /// Used for partial specialization - typename Enable = bool> -class CustomMmaBase { -public: - ///< Size of the Gemm problem - concept: gemm::GemmShape<> - using Shape = Shape_; - - ///< Policy describing tuning details - using Policy = Policy_; - - // - // Dependent types - // - - /// Warp-level Mma - using Operator = typename Policy::Operator; - - /// Shape describing the overall GEMM computed from shared memory - /// by each warp. - using WarpGemm = typename Policy::Operator::Shape; - - /// Shape describing the number of warps filling the CTA - using WarpCount = - GemmShape; - - /// Number of warp-level GEMM oeprations - static int const kWarpGemmIterations = (WarpGemm::kK / Operator::Policy::MmaShape::kK); - - /// Number of stages - static int const kStages = Stages; - - // - // Nested structs - // - - /// Shared storage object needed by threadblock-scoped GEMM - template - struct OperandSharedStorage { - AlignedBuffer buffer; - using TensorRef = TensorRef; - - CUTLASS_DEVICE - static OperandLayout Layout() - { - return OperandLayout::packed({OperandShape::kRow, OperandShape::kColumn}); - } - - /// Returns a TensorRef to the operand - CUTLASS_HOST_DEVICE - TensorRef ref() { return TensorRef{buffer.data(), Layout()}; } - }; - - /// Shape of the A matrix operand in shared memory - using ShapeA = MatrixShape; - - /// Shape of the B matrix operand in shared memory - using ShapeB = MatrixShape; - - using SharedStorageA = - OperandSharedStorage; - using SharedStorageB = - OperandSharedStorage; - using TensorRefA = typename SharedStorageA::TensorRef; - using TensorRefB = typename SharedStorageB::TensorRef; - - struct SharedStorage { - /// Buffer for A operand - SharedStorageA operand_A; - - /// Buffer for B operand - SharedStorageB operand_B; - }; - -protected: - // - // Data members - // - - /// Iterator to load a warp-scoped tile of A operand from shared memory - typename Operator::IteratorA warp_tile_iterator_A_; - - /// Iterator to load a warp-scoped tile of B operand from shared memory - typename Operator::IteratorB warp_tile_iterator_B_; - -public: - /// Construct from tensor references - CUTLASS_DEVICE - CustomMmaBase( - ///< Shared storage needed for internal use by threadblock-scoped GEMM - SharedStorageA& shared_storageA, - SharedStorageB& shared_storageB, - ///< ID within the threadblock - int thread_idx, - ///< ID of warp - int warp_idx, - ///< ID of each thread within a warp - int lane_idx) - : warp_tile_iterator_A_(shared_storageA.ref(), lane_idx), - warp_tile_iterator_B_(shared_storageB.ref(), lane_idx) - { - } -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace threadblock -} // namespace gemm -} // namespace cutlass - -///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_multistage.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_multistage.h deleted file mode 100644 index 3760ccab852a5ceec65fa6693070a84fec284d8a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_multistage.h +++ /dev/null @@ -1,714 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights - *reserved. SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, - *this list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE - *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE - *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR - *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF - *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS - *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN - *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) - *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE - *POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/*! \file - \brief Template for a double-buffered threadblock-scoped GEMM kernel. -*/ - -#pragma once - -#include "cutlass/aligned_buffer.h" -#include "cutlass/arch/cache_operation.h" -#include "cutlass/arch/memory.h" -#include "cutlass/array.h" -#include "cutlass/cutlass.h" -#include "cutlass/gemm/gemm.h" -#include "cutlass/matrix_shape.h" -#include "cutlass/numeric_types.h" - -#include "custom_mma_base.h" - -///////////////////////////////////////////////////////////////////////////////////////////////// - -namespace cutlass { -namespace gemm { -namespace threadblock { - -///////////////////////////////////////////////////////////////////////////////////////////////// - -/// Structure to compute the matrix product targeting CUDA cores and SIMT math -/// instructions. -template < - /// Size of the Gemm problem - concept: gemm::GemmShape<> - typename Shape_, - /// Iterates over tiles of A operand in global memory - // (concept: ReadableTileIterator | ForwardTileIterator | - // MaskedTileIterator) - typename IteratorA_, - /// Iterates over tiles of A operand in shared memory - /// (concept: WriteableTileIterator | RandomAccessTileIterator) - typename SmemIteratorA_, - /// Cache operation for operand A - cutlass::arch::CacheOperation::Kind CacheOpA, - /// Iterates over tiles of B operand in global memory - // (concept: ReadableTileIterator | ForwardTileIterator | - // MaskedTileIterator) - typename IteratorB_, - /// Iterates over tiles of B operand in shared memory - /// (concept: WriteableTileIterator | RandomAccessTileIterator) - typename SmemIteratorB_, - /// Cache operation for operand B - cutlass::arch::CacheOperation::Kind CacheOpB, - /// Data type of accumulator matrix - typename ElementC_, - /// Data type of accumulator matrix - typename LayoutC_, - /// Policy describing tuning details (concept: MmaPolicy) - typename Policy_, - /// Number of stages, - int Stages, - /// Use zfill or predicate for out-of-bound cp.async - SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, - /// Upper boundon the K dimension - int kMaxK = cutlass::platform::numeric_limits::max(), - /// Used for partial specialization - typename Enable = bool> -class CustomMmaMultistage : public CustomMmaBase { -public: - ///< Base class - using Base = CustomMmaBase; - ///< Size of the Gemm problem - concept: gemm::GemmShape<> - using Shape = Shape_; - ///< Iterates over tiles of A operand in global memory - using IteratorA = IteratorA_; - ///< Iterates over tiles of B operand in global memory - using IteratorB = IteratorB_; - ///< Data type of accumulator matrix - using ElementC = ElementC_; - ///< Layout of accumulator matrix - using LayoutC = LayoutC_; - ///< Policy describing tuning details - using Policy = Policy_; - - using SmemIteratorA = SmemIteratorA_; - using SmemIteratorB = SmemIteratorB_; - - static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; - static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; - - // - // Dependent types - // - - /// Fragment of accumulator tile - using FragmentC = typename Policy::Operator::FragmentC; - - /// Warp-level Mma - using Operator = typename Policy::Operator; - - /// Minimum architecture is Sm80 to support cp.async - using ArchTag = arch::Sm80; - - /// Complex transform on A operand - static ComplexTransform const kTransformA = Operator::kTransformA; - - /// Complex transform on B operand - static ComplexTransform const kTransformB = Operator::kTransformB; - - /// Internal structure exposed for introspection. - struct Detail { - static_assert(Base::kWarpGemmIterations > 1, - "The pipelined structure requires at least two warp-level " - "GEMM operations."); - - /// Number of cp.async instructions to load one stage of operand A - static int const AsyncCopyIterationsPerStageA = IteratorA::ThreadMap::Iterations::kCount; - - /// Number of cp.async instructions to load one stage of operand B - static int const AsyncCopyIterationsPerStageB = IteratorB::ThreadMap::Iterations::kCount; - - /// Number of stages - static int const kStages = Stages; - - /// Number of cp.async instructions to load on group of operand A - static int const kAccessesPerGroupA = - (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / - Base::kWarpGemmIterations; - - /// Number of cp.async instructions to load on group of operand B - static int const kAccessesPerGroupB = - (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / - Base::kWarpGemmIterations; - }; - - static bool const kSmemContainsEntireMat = kMaxK <= Shape::kK * Stages; - static constexpr int kNumStagesConcurrentLoad = kSmemContainsEntireMat ? Stages : Stages - 1; - -private: - using WarpLoadedFragmentA = typename Operator::FragmentA; - using WarpLoadedFragmentB = typename Operator::FragmentB; - using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; - using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; - -private: - // - // Data members - // - - /// Iterator to write threadblock-scoped tile of A operand to shared memory - SmemIteratorA smem_iterator_A_; - - /// Iterator to write threadblock-scoped tile of B operand to shared memory - SmemIteratorB smem_iterator_B_; - - bool prologue_done_; - - // Set to `True` to ensure the accumulator will be zero outside the GEMM - // footprint - bool zero_outside_bounds_; - -public: - /// Construct from tensor references - CUTLASS_DEVICE - CustomMmaMultistage( - ///< Shared storage needed for internal use by threadblock-scoped GEMM - typename Base::SharedStorageA& shared_storageA, - typename Base::SharedStorageB& shared_storageB, - ///< ID within the threadblock - int thread_idx, - ///< ID of warp - int warp_idx, - ///< ID of each thread within a warp - int lane_idx) - : Base(shared_storageA, shared_storageB, thread_idx, warp_idx, lane_idx), - smem_iterator_A_(shared_storageA.ref(), thread_idx), - smem_iterator_B_(shared_storageB.ref(), thread_idx), - prologue_done_(false), - zero_outside_bounds_(false) - { - // Compute warp location within threadblock tile by mapping the warp_id to - // three coordinates: - // _m: the warp's position within the threadblock along the M dimension - // _n: the warp's position within the threadblock along the N dimension - // _k: the warp's position within the threadblock along the K dimension - - int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); - int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); - - int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; - int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; - - // Add per-warp offsets in units of warp-level tiles - this->warp_tile_iterator_A_.add_tile_offset( - {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); - this->warp_tile_iterator_B_.add_tile_offset( - {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); - } - CUTLASS_DEVICE - CustomMmaMultistage( - ///< Shared storage needed for internal use by threadblock-scoped GEMM - typename Base::SharedStorage& st, - ///< ID within the threadblock - int thread_idx, - ///< ID of warp - int warp_idx, - ///< ID of each thread within a warp - int lane_idx) - : CustomMmaMultistage(st.operand_A, st.operand_B, thread_idx, warp_idx, lane_idx) - { - } - - CUTLASS_DEVICE - bool set_prologue_done(bool value) - { - prologue_done_ = value; - return true; - } - - CUTLASS_DEVICE - bool set_zero_outside_bounds(bool value) - { - zero_outside_bounds_ = value; - return true; - } - - template - CUTLASS_DEVICE static void prologue(typename Base::SharedStorage& shared_storage, - ///< iterator over A operand in global memory - IteratorA iterator_A, - ///< iterator over B operand in global memory - IteratorB iterator_B, - int thread_idx, - int problem_size_k) - { - prologue(shared_storage.operand_A, - shared_storage.operand_B, - iterator_A, - iterator_B, - thread_idx, - problem_size_k); - } - - template - CUTLASS_DEVICE static void prologue(typename Base::SharedStorageA& shared_storageA, - typename Base::SharedStorageB& shared_storageB, - ///< iterator over A operand in global memory - IteratorA iterator_A, - ///< iterator over B operand in global memory - IteratorB iterator_B, - int thread_idx, - int problem_size_k) - { - SmemIteratorA smem_iterator_A(shared_storageA.ref(), thread_idx); - SmemIteratorB smem_iterator_B(shared_storageB.ref(), thread_idx); - int32_t iter = (problem_size_k + Base::Shape::kK - 1) / Base::Shape::kK; - _prologue(iterator_A, iterator_B, iter, smem_iterator_A, smem_iterator_B); - } - - CUTLASS_DEVICE - void copy_tiles_and_advance(IteratorA& iterator_A, - IteratorB& iterator_B, - int group_start_A = 0, - int group_start_B = 0) - { - iterator_A.set_iteration_index(group_start_A * IteratorA::kAccessesPerVector); - this->smem_iterator_A_.set_iteration_index(group_start_A); - - // Async Copy for operand A - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { - if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { - typename IteratorA::AccessType* dst_ptr = - reinterpret_cast(this->smem_iterator_A_.get()); - - int const kSrcBytes = sizeof_bits::value * - IteratorA::ThreadMap::kElementsPerAccess / - IteratorA::kAccessesPerVector / 8; - - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { - auto gmem_ptr = iterator_A.get(); - - if (zero_outside_bounds_ || - SharedMemoryClear == SharedMemoryClearOption::kZfill) { - cutlass::arch::cp_async_zfill( - dst_ptr + v, gmem_ptr, iterator_A.valid()); - } else { - cutlass::arch::cp_async( - dst_ptr + v, gmem_ptr, iterator_A.valid()); - } - - ++iterator_A; - } - - ++this->smem_iterator_A_; - } - } - - iterator_B.set_iteration_index(group_start_B * IteratorB::kAccessesPerVector); - this->smem_iterator_B_.set_iteration_index(group_start_B); - - // Async Copy for operand B - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { - if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { - typename IteratorB::AccessType* dst_ptr = - reinterpret_cast(this->smem_iterator_B_.get()); - - int const kSrcBytes = sizeof_bits::value * - IteratorB::ThreadMap::kElementsPerAccess / - IteratorB::kAccessesPerVector / 8; - - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { - auto gmem_ptr = iterator_B.get(); - - if (zero_outside_bounds_ || - SharedMemoryClear == SharedMemoryClearOption::kZfill) { - cutlass::arch::cp_async_zfill( - dst_ptr + v, gmem_ptr, iterator_B.valid()); - } else { - cutlass::arch::cp_async( - dst_ptr + v, gmem_ptr, iterator_B.valid()); - } - - ++iterator_B; - } - ++this->smem_iterator_B_; - } - } - } - - template - CUTLASS_DEVICE static void _prologue(IteratorA& iterator_A, - IteratorB& iterator_B, - int32_t& gemm_k_iterations, - SmemIteratorA& smem_iterator_A_, - SmemIteratorB& smem_iterator_B_) - { - // Issue several complete stages - CUTLASS_PRAGMA_UNROLL - for (int stage = 0; stage < kNumStagesConcurrentLoad; ++stage, --gemm_k_iterations) { - iterator_A.clear_mask(gemm_k_iterations == 0); - iterator_B.clear_mask(gemm_k_iterations == 0); - - iterator_A.set_iteration_index(0); - smem_iterator_A_.set_iteration_index(0); - - // Async Copy for operand A - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { - typename IteratorA::AccessType* dst_ptr = - reinterpret_cast(smem_iterator_A_.get()); - - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { - int const kSrcBytes = sizeof_bits::value * - IteratorA::ThreadMap::kElementsPerAccess / - IteratorA::kAccessesPerVector / 8; - - int src_bytes = (iterator_A.valid() ? kSrcBytes : 0); - - if (kLoadA) { - cutlass::arch::cp_async_zfill( - dst_ptr + v, iterator_A.get(), iterator_A.valid()); - } - - ++iterator_A; - } - - ++smem_iterator_A_; - } - - iterator_B.set_iteration_index(0); - smem_iterator_B_.set_iteration_index(0); - - // Async Copy for operand B - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { - typename IteratorB::AccessType* dst_ptr = - reinterpret_cast(smem_iterator_B_.get()); - - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { - int const kSrcBytes = sizeof_bits::value * - IteratorB::ThreadMap::kElementsPerAccess / - IteratorB::kAccessesPerVector / 8; - - if (kLoadB) { - cutlass::arch::cp_async_zfill( - dst_ptr + v, iterator_B.get(), iterator_B.valid()); - } - - ++iterator_B; - } - - ++smem_iterator_B_; - } - - // Move to the next stage - iterator_A.add_tile_offset({0, 1}); - iterator_B.add_tile_offset({1, 0}); - - smem_iterator_A_.add_tile_offset({0, 1}); - smem_iterator_B_.add_tile_offset({1, 0}); - - // Defines the boundary of a stage of cp.async. - cutlass::arch::cp_async_fence(); - } - } - - /// Perform a threadblock-scoped matrix multiply-accumulate - CUTLASS_DEVICE - void operator()( - ///< problem size of GEMM - int gemm_k_iterations, - ///< destination accumulator tile - FragmentC& accum, - ///< iterator over A operand in global memory - IteratorA iterator_A, - ///< iterator over B operand in global memory - IteratorB iterator_B, - ///< initial value of accumulator - FragmentC const& src_accum) - { - // - // Prologue - // - - if (!prologue_done_) { - _prologue( - iterator_A, iterator_B, gemm_k_iterations, smem_iterator_A_, smem_iterator_B_); - } else if (!kSmemContainsEntireMat) { - _prologue( - iterator_A, iterator_B, gemm_k_iterations, smem_iterator_A_, smem_iterator_B_); - } else { - gemm_k_iterations -= kNumStagesConcurrentLoad; - } - - // Perform accumulation in the 'd' output operand - accum = src_accum; - - // - // Clear the remaining tiles of SMEM. This is a functional requirement for - // some kernels so that all accumulator elements outside the GEMM footprint - // are zero. - // - - if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage) { - /// Iterator to write threadblock-scoped tile of A operand to shared - /// memory - SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_); - - typename IteratorA::AccessType zero_A; - zero_A.clear(); - - last_smem_iterator_A.set_iteration_index(0); - - // Async Copy for operand A - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { - typename IteratorA::AccessType* dst_ptr = - reinterpret_cast(last_smem_iterator_A.get()); - - *dst_ptr = zero_A; - - ++last_smem_iterator_A; - } - - /// Iterator to write threadblock-scoped tile of B operand to shared - /// memory - SmemIteratorB last_smem_iterator_B(this->smem_iterator_B_); - typename IteratorB::AccessType zero_B; - - zero_B.clear(); - last_smem_iterator_B.set_iteration_index(0); - - // Async Copy for operand B - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { - typename IteratorB::AccessType* dst_ptr = - reinterpret_cast(last_smem_iterator_B.get()); - - *dst_ptr = zero_B; - - ++last_smem_iterator_B; - } - } - - // Waits until kStages-2 stages have committed. - cutlass::arch::cp_async_wait(); - __syncthreads(); - - // Pair of fragments used to overlap shared memory loads and math - // instructions - WarpLoadedFragmentA warp_loaded_frag_A[2]; - WarpLoadedFragmentB warp_loaded_frag_B[2]; - WarpTransformedFragmentA warp_transformed_frag_A[2]; - WarpTransformedFragmentB warp_transformed_frag_B[2]; - - Operator warp_mma; - - this->warp_tile_iterator_A_.set_kgroup_index(0); - this->warp_tile_iterator_B_.set_kgroup_index(0); - - this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); - this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); - - ++this->warp_tile_iterator_A_; - ++this->warp_tile_iterator_B_; - - iterator_A.clear_mask(gemm_k_iterations == 0); - iterator_B.clear_mask(gemm_k_iterations == 0); - - int smem_write_stage_idx = Base::kStages - 1; - int smem_read_stage_idx = 0; - - warp_mma.transform(warp_transformed_frag_A[0], - warp_transformed_frag_B[0], - warp_loaded_frag_A[0], - warp_loaded_frag_B[0]); - - // tf32x3 kernels use staging accumulation. warp_mma uses a temporary - // accumulator and this temporary accumulator is added to the final - // accumulator once in every mainloop iteration. - plus plus_accum; - - FragmentC tmp_accum; - - if (platform::is_same::value || - platform::is_same::value) { - tmp_accum.clear(); - } - - // - // Mainloop - // - - CUTLASS_GEMM_LOOP - for (; gemm_k_iterations > (-kNumStagesConcurrentLoad);) { - // - // Loop over GEMM K dimension - // - - // Computes a warp-level GEMM on data held in shared memory - // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate - CUTLASS_PRAGMA_UNROLL - for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) { - // Load warp-level tiles from shared memory, wrapping to k offset if - // this is the last group as the case may be. - - this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % - Base::kWarpGemmIterations); - this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % - Base::kWarpGemmIterations); - - // In case of a non-circular buffer ("kSmemContainsEntireMat") - // make sure we don't load out of bounds data. - if (!kSmemContainsEntireMat || gemm_k_iterations > (-kNumStagesConcurrentLoad) || - warp_mma_k < Base::kWarpGemmIterations - 1) { - this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]); - this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]); - } - - ++this->warp_tile_iterator_A_; - ++this->warp_tile_iterator_B_; - - if (warp_mma_k > 0) - warp_mma.transform(warp_transformed_frag_A[warp_mma_k % 2], - warp_transformed_frag_B[warp_mma_k % 2], - warp_loaded_frag_A[warp_mma_k % 2], - warp_loaded_frag_B[warp_mma_k % 2]); - - if (platform::is_same::value || - platform::is_same::value) { - warp_mma(tmp_accum, - warp_transformed_frag_A[warp_mma_k % 2], - warp_transformed_frag_B[warp_mma_k % 2], - tmp_accum); - - if (warp_mma_k == 0) { - accum = plus_accum(accum, tmp_accum); - tmp_accum.clear(); - } - } else { - warp_mma(accum, - warp_transformed_frag_A[warp_mma_k % 2], - warp_transformed_frag_B[warp_mma_k % 2], - accum); - } - - // Issue global->shared copies for the this stage - if (!kSmemContainsEntireMat && warp_mma_k < Base::kWarpGemmIterations - 1) { - int group_start_iteration_A, group_start_iteration_B; - - group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA; - group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB; - - copy_tiles_and_advance( - iterator_A, iterator_B, group_start_iteration_A, group_start_iteration_B); - } - - if (warp_mma_k + 2 == Base::kWarpGemmIterations) { - if (!kSmemContainsEntireMat) { - int group_start_iteration_A, group_start_iteration_B; - group_start_iteration_A = (warp_mma_k + 1) * Detail::kAccessesPerGroupA; - group_start_iteration_B = (warp_mma_k + 1) * Detail::kAccessesPerGroupB; - - copy_tiles_and_advance(iterator_A, - iterator_B, - group_start_iteration_A, - group_start_iteration_B); - } - - // Inserts a memory fence between stages of cp.async instructions. - cutlass::arch::cp_async_fence(); - - // Waits until kStages-2 stages have committed. - cutlass::arch::cp_async_wait(); - __syncthreads(); - - // Move to the next stage - iterator_A.add_tile_offset({0, 1}); - iterator_B.add_tile_offset({1, 0}); - - this->smem_iterator_A_.add_tile_offset({0, 1}); - this->smem_iterator_B_.add_tile_offset({1, 0}); - - // Add negative offsets to return iterators to the 'start' of the - // circular buffer in shared memory - if (smem_write_stage_idx == (Base::kStages - 1)) { - this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); - this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); - smem_write_stage_idx = 0; - } else { - ++smem_write_stage_idx; - } - - if (!kSmemContainsEntireMat && smem_read_stage_idx == (Base::kStages - 1)) { - this->warp_tile_iterator_A_.add_tile_offset( - {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations}); - this->warp_tile_iterator_B_.add_tile_offset( - {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations, 0}); - smem_read_stage_idx = 0; - } else { - ++smem_read_stage_idx; - } - - --gemm_k_iterations; - iterator_A.clear_mask(gemm_k_iterations == 0); - iterator_B.clear_mask(gemm_k_iterations == 0); - } - - // Do any conversions feeding the first stage at the end of the loop so - // we can start right away on mma instructions - if (warp_mma_k + 1 == Base::kWarpGemmIterations) - warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2], - warp_transformed_frag_B[(warp_mma_k + 1) % 2], - warp_loaded_frag_A[(warp_mma_k + 1) % 2], - warp_loaded_frag_B[(warp_mma_k + 1) % 2]); - } - } - - if (platform::is_same::value || - platform::is_same::value) { - accum = plus_accum(accum, tmp_accum); - } - } -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace threadblock -} // namespace gemm -} // namespace cutlass - -///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_pipelined.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_pipelined.h deleted file mode 100644 index 07b26ca3129968f8d78810df5187dd92c3c166c9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/custom_mma_pipelined.h +++ /dev/null @@ -1,388 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights - *reserved. SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, - *this list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE - *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE - *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR - *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF - *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS - *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN - *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) - *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE - *POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/*! \file - \brief Template for a double-buffered threadblock-scoped GEMM kernel. -*/ - -#pragma once - -#include "cutlass/aligned_buffer.h" -#include "cutlass/array.h" -#include "cutlass/cutlass.h" -#include "cutlass/numeric_conversion.h" - -#include "cutlass/matrix_shape.h" -#include "cutlass/numeric_types.h" - -#include "custom_mma_base.h" -#include "cutlass/gemm/gemm.h" - -///////////////////////////////////////////////////////////////////////////////////////////////// - -namespace cutlass { -namespace gemm { -namespace threadblock { - -///////////////////////////////////////////////////////////////////////////////////////////////// - -/// Structure to compute the matrix product targeting CUDA cores and SIMT math -/// instructions. -template < - /// Size of the Gemm problem - concept: gemm::GemmShape<> - typename Shape_, - /// Iterates over tiles of A operand in global memory - // (concept: ReadableTileIterator | ForwardTileIterator | - // MaskedTileIterator) - typename IteratorA_, - /// Iterates over tiles of A operand in shared memory - /// (concept: WriteableTileIterator | RandomAccessTileIterator) - typename SmemIteratorA_, - /// Iterates over tiles of B operand in global memory - // (concept: ReadableTileIterator | ForwardTileIterator | - // MaskedTileIterator) - typename IteratorB_, - /// Iterates over tiles of B operand in shared memory - /// (concept: WriteableTileIterator | RandomAccessTileIterator) - typename SmemIteratorB_, - /// Data type of accumulator matrix - typename ElementC_, - /// Data type of accumulator matrix - typename LayoutC_, - /// Policy describing tuning details (concept: MmaPolicy) - typename Policy_, - /// Transformation applied to A operand - typename TransformA_ = NumericArrayConverter, - /// - /// Transformation applied to B operand - typename TransformB_ = NumericArrayConverter, - /// Used for partial specialization - typename Enable = bool> -class CustomMmaPipelined : public CustomMmaBase { -public: - ///< Base class - using Base = CustomMmaBase; - - using Shape = Shape_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> - using IteratorA = IteratorA_; ///< Iterates over tiles of A operand in global memory - using IteratorB = IteratorB_; ///< Iterates over tiles of B operand in global memory - using ElementC = ElementC_; ///< Data type of accumulator matrix - using LayoutC = LayoutC_; ///< Layout of accumulator matrix - using Policy = Policy_; ///< Policy describing tuning details - - using SmemIteratorA = SmemIteratorA_; - using SmemIteratorB = SmemIteratorB_; - - using TransformA = TransformA_; - using TransformB = TransformB_; - - // - // Dependent types - // - - /// Fragment of operand A loaded from global memory - using FragmentA = typename IteratorA::Fragment; - - /// Fragment of operand B loaded from global memory - using FragmentB = typename IteratorB::Fragment; - - /// Fragment of accumulator tile - using FragmentC = typename Policy::Operator::FragmentC; - - /// Warp-level Mma - using Operator = typename Policy::Operator; - - /// Obtain the arch tag from the warp-level operator - using ArchTag = typename Policy::Operator::ArchTag; - - /// Complex transform on A operand - static ComplexTransform const kTransformA = Operator::kTransformA; - - /// Complex transform on B operand - static ComplexTransform const kTransformB = Operator::kTransformB; - - // staticaly assert kStages for MmaPipelined is two (Double-buffered pipeline) - static_assert((Base::kStages == 2), "MmaPipelined requires kStages set to value 2"); - - static bool const kSmemContainsEntireMat = false; - -private: - using WarpFragmentA = typename Operator::FragmentA; - using WarpFragmentB = typename Operator::FragmentB; - -protected: - /// Iterator to write threadblock-scoped tile of A operand to shared memory - SmemIteratorA smem_iterator_A_; - - /// Iterator to write threadblock-scoped tile of B operand to shared memory - SmemIteratorB smem_iterator_B_; - -public: - /// Construct from tensor references - CUTLASS_DEVICE - CustomMmaPipelined(typename Base::SharedStorageA& shared_storageA, - typename Base::SharedStorageB& shared_storageB, - int thread_idx, ///< ID within the threadblock - int warp_idx, ///< ID of warp - int lane_idx ///< ID of each thread within a warp - ) - : Base(shared_storageA, shared_storageB, thread_idx, warp_idx, lane_idx), - smem_iterator_A_(shared_storageA.ref(), thread_idx), - smem_iterator_B_(shared_storageB.ref(), thread_idx) - { - // Compute warp location within threadblock tile by mapping the warp_id to - // three coordinates: - // _m: the warp's position within the threadblock along the M dimension - // _n: the warp's position within the threadblock along the N dimension - // _k: the warp's position within the threadblock along the K dimension - - int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); - int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); - - int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; - int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; - - // Add per-warp offsets in units of warp-level tiles - this->warp_tile_iterator_A_.add_tile_offset( - {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); - this->warp_tile_iterator_B_.add_tile_offset( - {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); - } - CUTLASS_DEVICE - CustomMmaPipelined( - ///< Shared storage needed for internal use by threadblock-scoped GEMM - typename Base::SharedStorage& st, - ///< ID within the threadblock - int thread_idx, - ///< ID of warp - int warp_idx, - ///< ID of each thread within a warp - int lane_idx) - : CustomMmaPipelined(st.operand_A, st.operand_B, thread_idx, warp_idx, lane_idx) - { - } - - CUTLASS_DEVICE - bool set_prologue_done(bool value) - { - // NOT IMPLEMENTED FOR PIPELINED - } - - CUTLASS_DEVICE - bool set_zero_outside_bounds(bool value) - { - // NOT NEEDED FOR PIPELINED - // shared memory will always be zero-filled - } - - template - CUTLASS_DEVICE static void prologue(typename Base::SharedStorage& shared_storage, - ///< iterator over A operand in global memory - IteratorA iterator_A, - ///< iterator over B operand in global memory - IteratorB iterator_B, - int thread_idx, - int problem_size_k) - { - prologue(shared_storage.operand_A, - shared_storage.operand_B, - iterator_A, - iterator_B, - thread_idx, - problem_size_k); - } - - template - CUTLASS_DEVICE static void prologue(typename Base::SharedStorageA& shared_storageA, - typename Base::SharedStorageB& shared_storageB, - ///< iterator over A operand in global memory - IteratorA iterator_A, - ///< iterator over B operand in global memory - IteratorB iterator_B, - int thread_idx, - int problem_size_k) - { - // NOT IMPLEMENTED FOR PIPELINED - } - - /// Perform a threadblock-scoped matrix multiply-accumulate - CUTLASS_DEVICE - void operator()( - int gemm_k_iterations, ///< number of iterations of the mainloop - FragmentC& accum, ///< destination accumulator tile - IteratorA iterator_A, ///< iterator over A operand in global memory - IteratorB iterator_B, ///< iterator over B operand in global memory - FragmentC const& src_accum, ///< source accumulator tile - TransformA transform_A = TransformA(), ///< transformation applied to A fragment - TransformB transform_B = TransformB()) - { ///< transformation applied to B fragment - - // - // Prologue - // - - // Perform accumulation in the 'd' output operand - accum = src_accum; - - FragmentA tb_frag_A; - FragmentB tb_frag_B; - - tb_frag_A.clear(); - tb_frag_B.clear(); - - // The last kblock is loaded in the prolog - iterator_A.load(tb_frag_A); - iterator_B.load(tb_frag_B); - - ++iterator_A; - ++iterator_B; - - this->smem_iterator_A_.store(transform_A(tb_frag_A)); - this->smem_iterator_B_.store(transform_B(tb_frag_B)); - - ++this->smem_iterator_A_; - ++this->smem_iterator_B_; - - __syncthreads(); - - // Pair of fragments used to overlap shared memory loads and math - // instructions - WarpFragmentA warp_frag_A[2]; - WarpFragmentB warp_frag_B[2]; - - this->warp_tile_iterator_A_.set_kgroup_index(0); - this->warp_tile_iterator_B_.set_kgroup_index(0); - - this->warp_tile_iterator_A_.load(warp_frag_A[0]); - this->warp_tile_iterator_B_.load(warp_frag_B[0]); - - ++this->warp_tile_iterator_A_; - ++this->warp_tile_iterator_B_; - - Operator warp_mma; - - int smem_write_stage_idx = 1; - - // Avoid reading out of bounds - iterator_A.clear_mask(gemm_k_iterations <= 1); - iterator_B.clear_mask(gemm_k_iterations <= 1); - - // Issue loads during the first warp-level matrix multiply-add *AFTER* - // issuing shared memory loads (which have the tightest latency requirement). - - // - // Mainloop - // - - // Note: The main loop does not support Base::kWarpGemmIterations == 2. - CUTLASS_GEMM_LOOP - for (; gemm_k_iterations > 0; --gemm_k_iterations) { - // - // Loop over GEMM K dimension - // - - CUTLASS_PRAGMA_UNROLL - for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) { - // Load warp-level tiles from shared memory, wrapping to k offset if - // this is the last group as the case may be. - - if (warp_mma_k == Base::kWarpGemmIterations - 1) { - // Write fragments to shared memory - this->smem_iterator_A_.store(transform_A(tb_frag_A)); - - this->smem_iterator_B_.store(transform_B(tb_frag_B)); - - __syncthreads(); - - ++this->smem_iterator_A_; - ++this->smem_iterator_B_; - - // Add negative offsets to return iterators to the 'start' of the - // circular buffer in shared memory - if (smem_write_stage_idx == 1) { - this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); - this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); - } else { - this->warp_tile_iterator_A_.add_tile_offset( - {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations}); - this->warp_tile_iterator_B_.add_tile_offset( - {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations, 0}); - } - - smem_write_stage_idx ^= 1; - } - - this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % - Base::kWarpGemmIterations); - this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % - Base::kWarpGemmIterations); - - this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]); - this->warp_tile_iterator_B_.load(warp_frag_B[(warp_mma_k + 1) % 2]); - - ++this->warp_tile_iterator_A_; - ++this->warp_tile_iterator_B_; - - if (warp_mma_k == 0) { - iterator_A.load(tb_frag_A); - iterator_B.load(tb_frag_B); - - ++iterator_A; - ++iterator_B; - - // Avoid reading out of bounds if this was the last loop iteration - iterator_A.clear_mask(gemm_k_iterations <= 2); - iterator_B.clear_mask(gemm_k_iterations <= 2); - } - - warp_mma(accum, warp_frag_A[warp_mma_k % 2], warp_frag_B[warp_mma_k % 2], accum); - } - } - } -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace threadblock -} // namespace gemm -} // namespace cutlass - -///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/find_default_mma.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/find_default_mma.h deleted file mode 100644 index 163dcbf852595c2f97fe950b5be6576d938d4efb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/find_default_mma.h +++ /dev/null @@ -1,191 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holdvr nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/*! \file - \brief Cutlass provides helper template functions to figure out the right - datastructures to instantiate to run a GEMM with various parameters (see - `cutlass/gemm/threadblock/default_mma.h`). However, due to template - instantiation priority rules, it will only create an MmaMultiStage with - kStages=3 (otherwise creates an MmePipelined - which is not compatible with - FastF32). kStages=3 uses too much shared memory and we want to use kStages=2, - so we just copy-pasted some code from `default_mma.h` and - `default_mma_core.h` files and wrapped this template to allow our usecase. - - This is really only for the FastF32 case - aka using TensorCores with fp32. -*/ - -#pragma once - -#include "cutlass/gemm/threadblock/default_mma.h" -#include "cutlass/gemm/threadblock/default_mma_core_simt.h" -#include "cutlass/gemm/threadblock/default_mma_core_sm70.h" -#include "cutlass/gemm/threadblock/default_mma_core_sm75.h" -#include "cutlass/gemm/threadblock/default_mma_core_sm80.h" - -namespace cutlass { -namespace gemm { -namespace threadblock { - -template < - /// Element type for A matrix operand - typename ElementA, - /// Layout type for A matrix operand - typename LayoutA, - /// Access granularity of A matrix in units of elements - int kAlignmentA, - /// Element type for B matrix operand - typename ElementB, - /// Layout type for B matrix operand - typename LayoutB, - /// Access granularity of B matrix in units of elements - int kAlignmentB, - /// Element type for internal accumulation - typename ElementAccumulator, - /// Layout type for C and D matrix operand - typename LayoutC, - /// Operator class tag - typename OperatorClass, - /// Tag indicating architecture to tune for - typename ArchTag, - /// Threadblock-level tile size (concept: GemmShape) - typename ThreadblockShape, - /// Warp-level tile size (concept: GemmShape) - typename WarpShape, - /// Instruction-level tile size (concept: GemmShape) - typename InstructionShape, - /// Number of stages used in the pipelined mainloop - int Stages, - /// Operation performed by GEMM - typename Operator, - typename Enable_ = void> -struct FindDefaultMma { - static constexpr bool AccumulatorsInRowMajor = false; - static constexpr SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone; - using DefaultMma = cutlass::gemm::threadblock::DefaultMma; -}; - -/// Specialization for sm80 / FastF32 / multistage with kStages=2 -template -struct FindDefaultMma 1)>::type> { - using LayoutC = layout::RowMajor; - using OperatorClass = arch::OpClassTensorOp; - using ArchTag = arch::Sm80; - - using DefaultMma_ = cutlass::gemm::threadblock::DefaultMma; - struct DefaultMma : DefaultMma_ { - using MmaCore_ = typename DefaultMma_::MmaCore; - // Define the threadblock-scoped multistage matrix multiply - using ThreadblockMma = - cutlass::gemm::threadblock::MmaMultistage; - }; -}; - -} // namespace threadblock -} // namespace gemm -} // namespace cutlass diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_accum_lambda_iterator.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_accum_lambda_iterator.h deleted file mode 100644 index 5e2f0cf681bfe10fdc30c2dfcbe390c5f711dae0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_accum_lambda_iterator.h +++ /dev/null @@ -1,347 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holdvr nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include "cutlass/functional.h" -#include "cutlass/gemm/warp/mma_simt_tile_iterator.h" -#include "cutlass/gemm/warp/mma_tensor_op_tile_iterator_sm70.h" -#include "cutlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" -#include "cutlass/matrix_shape.h" - -/* -TensorCores have different accumulator layouts. -This file provides a class to easily map the accumulator -i-th element with the corresponding matrix row/col. -*/ - -template -struct AccumLambdaIteratorSm80 { - static_assert(cutlass::platform::is_same::value, - "only RowMajor is supported"); - - using Policy = typename T::Policy; - using InstructionShape = typename T::InstructionShape; - using OpDelta = typename T::OpDelta; - using Shape = typename T::Shape; - static int const kElementsPerAccess = InstructionShape::kN / 4; - static int const kRowsPerTile = 8; - static int const kAccumulatorRows = InstructionShape::kM / kRowsPerTile; - - static cutlass::MatrixCoord CUTLASS_DEVICE - get_lane_offset(int8_t lane_id, int8_t warp_id, typename T::TensorCoord const& tile_offset) - { - int quad = (lane_id >> 2); - int lane_in_quad = (lane_id & 3); - return cutlass::MatrixCoord( - quad + tile_offset.row() * Shape::kRow, - lane_in_quad * kElementsPerAccess + tile_offset.column() * Shape::kColumn); - } - - template - CUTLASS_DEVICE static void iterateRows(cutlass::MatrixCoord& lane_offset, - FA beginRow, - FB op, - FC endRow) - { - // See cutlass/gemm/warp/mma_tensor_op_tile_iterator.h - CUTLASS_PRAGMA_UNROLL - for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { - CUTLASS_PRAGMA_UNROLL - for (int row = 0; row < kAccumulatorRows; ++row) { - int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + row * kRowsPerTile + - lane_offset.row(); - beginRow(accum_m); - - CUTLASS_PRAGMA_UNROLL - for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { - int mma_accum_start = kAccumulatorRows * kElementsPerAccess * - (mma_n * Policy::MmaIterations::kRow + mma_m); - CUTLASS_PRAGMA_UNROLL - for (int col = 0; col < kElementsPerAccess; ++col) { - int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col + - lane_offset.column(); - int idx = mma_accum_start + row * kElementsPerAccess + col; - op(accum_m, accum_n, idx); - } - } - - endRow(accum_m); - } - } - } - - template - CUTLASS_DEVICE static bool reduceSameRow(int lane_id, DT& myValue, F fn) - { - // In each warp, 4 threads will work on the same row - // - the ones with the same `quad` - auto otherV = __shfl_xor_sync(0xffffffff, myValue, 1); - myValue = fn(myValue, otherV); - otherV = __shfl_xor_sync(0xffffffff, myValue, 2); - myValue = fn(myValue, otherV); - int lane_in_quad = (lane_id & 3); - return lane_in_quad == 0; - } -}; - -template -struct AccumLambdaIteratorSm70 { - static_assert(cutlass::platform::is_same::value, - "only RowMajor is supported"); - - using Policy = typename T::Policy; - using InstructionShape = typename T::InstructionShape; - using OpDelta = typename T::OpDelta; - using Shape = typename T::Shape; - using Element = accum_t; - - static int const kElementsPerPartial = 4; - using EleShapePerPatial = - typename cutlass::platform::conditional::value, - cutlass::MatrixShape<2, 2>, - cutlass::MatrixShape<1, 4>>::type; - static int const kElementsPerMma = 8; - static int const kAccumulatorPatials = 2; - using QuadShapePerPatialMma = cutlass::MatrixShape<4, 4>; - - static cutlass::MatrixCoord CUTLASS_DEVICE - get_lane_offset(int8_t lane_id, int8_t warp_id, typename T::TensorCoord const& tile_offset) - { - int quad = (lane_id >> 2); - int lane_in_quad = (lane_id & 3); - int accum_m, accum_n; - - if (cutlass::platform::is_same::value) { - // (quad[2],quad[0])+lane_in_quad[0] - accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + (lane_in_quad & 1); - // (quad[1])+lane_in_quad[1] - accum_n = ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials + - (lane_in_quad & 2); - } else { - accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + lane_in_quad; // (quad[2],quad[0]) - accum_n = ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials; - } - return cutlass::MatrixCoord(accum_m + tile_offset.row() * Shape::kRow, - accum_n + tile_offset.column() * Shape::kColumn); - } - - template - CUTLASS_DEVICE static bool reduceSameRow(int lane_id, DT& myValue, F fn) - { - static_assert(cutlass::platform::is_same::value, - "update to support non-float accum"); - // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-fragment-mma-884-f16 - // T0 & T2 share same line within a quad - auto otherV = __shfl_xor_sync(0xffffffff, myValue, 1 << 1); - myValue = fn(myValue, otherV); - // quad 0 and quad 2 are on the same lines - otherV = __shfl_xor_sync(0xffffffff, myValue, 1 << 3); - myValue = fn(myValue, otherV); - return (lane_id & ((1 << 1) | (1 << 3))) == 0; - } - - template - CUTLASS_DEVICE static void iterateRows(cutlass::MatrixCoord& lane_offset, - FA beginRow, - FB op, - FC endRow) - { - CUTLASS_PRAGMA_UNROLL - for (int tile_m = 0; tile_m < Policy::TileIterations::kRow; ++tile_m) { - CUTLASS_PRAGMA_UNROLL - for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { - CUTLASS_PRAGMA_UNROLL - for (int m = 0; m < EleShapePerPatial::kRow; ++m) { - int accum_m = tile_m * Policy::InterleavedTile::kRow + - mma_m * QuadShapePerPatialMma::kRow + m * 2 + lane_offset.row(); - beginRow(accum_m); - - CUTLASS_PRAGMA_UNROLL - for (int tile_n = 0; tile_n < Policy::TileIterations::kColumn; ++tile_n) { - CUTLASS_PRAGMA_UNROLL - for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { - CUTLASS_PRAGMA_UNROLL - for (int p = 0; p < kAccumulatorPatials; ++p) { - CUTLASS_PRAGMA_UNROLL - for (int n = 0; n < EleShapePerPatial::kColumn; ++n) { - int mma_accum_start = - (((tile_n * Policy::TileIterations::kRow + tile_m) * - Policy::MmaIterations::kColumn + - mma_n) * - Policy::MmaIterations::kRow + - mma_m) * - kElementsPerMma; - int accum_n = tile_n * Policy::InterleavedTile::kColumn + - mma_n * QuadShapePerPatialMma::kColumn + - p * Policy::InterleavedTile::kColumn / 2 + n + - lane_offset.column(); - int idx = mma_accum_start + p * kElementsPerPartial + - m * EleShapePerPatial::kColumn + n; - op(accum_m, accum_n, idx); - } - } - } - } - endRow(accum_m); - } - } - } - } -}; - -template -struct AccumLambdaIteratorSimt { - using Policy = typename T::Policy; - using Iterations = typename T::Iterations; - using Element = typename T::Element; - using Delta = typename T::Delta; - using Shape = typename T::Shape; - static_assert(cutlass::platform::is_same::value, - "only RowMajor is supported"); - - template - CUTLASS_DEVICE static bool reduceSameRow(int lane_id, DT& myValue, F fn) - { - CUTLASS_PRAGMA_UNROLL - for (int bit = 1; bit < Policy::WarpShape::kColumn; bit *= 2) { - auto otherV = __shfl_xor_sync(0xffffffff, myValue, bit); - myValue = fn(myValue, otherV); - } - return (lane_id & (Policy::WarpShape::kColumn - 1)) == 0; - } - - template - CUTLASS_DEVICE static void iterateRows(cutlass::MatrixCoord& lane_offset, - FA beginRow, - FB op, - FC endRow) - { - CUTLASS_PRAGMA_UNROLL - for (int mma_m = 0; mma_m < Iterations::kRow; ++mma_m) { - CUTLASS_PRAGMA_UNROLL - for (int m = 0; m < Policy::LaneMmaShape::kM; ++m) { - int accum_m = mma_m * Delta::kRow + m + lane_offset.row(); - beginRow(accum_m); - - CUTLASS_PRAGMA_UNROLL - for (int mma_n = 0; mma_n < Iterations::kColumn; ++mma_n) { - int accum_n = mma_n * Policy::WarpShape::kColumn * Policy::LaneMmaShape::kN + - lane_offset.column(); - CUTLASS_PRAGMA_UNROLL - for (int n = 0; n < Policy::LaneMmaShape::kN; ++n) { - int idx = n + Policy::LaneMmaShape::kN * - (mma_n + Iterations::kColumn * - (m + mma_m * Policy::LaneMmaShape::kM)); - op(accum_m, accum_n + n, idx); - } - } - endRow(accum_m); - } - } - } - - static cutlass::MatrixCoord CUTLASS_DEVICE - get_lane_offset(int8_t lane_id, int8_t warp_id, typename T::TensorCoord const& tile_offset) - { - static_assert(cutlass::platform::is_same>::value, - ""); - typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); - - cutlass::MatrixCoord lane_offset = - lane_layout.inverse(lane_id) * - cutlass::MatrixCoord(Policy::LaneMmaShape::kM, Policy::LaneMmaShape::kN); - return lane_offset + tile_offset * cutlass::MatrixCoord(Shape::kRow, Shape::kColumn); - } -}; - -template -struct DefaultMmaAccumLambdaIterator; - -// Simt -template -struct DefaultMmaAccumLambdaIterator< - cutlass::gemm::warp::MmaSimtTileIterator, - accum_t, - kWarpSize> { - using WarpIterator = - typename cutlass::gemm::warp::MmaSimtTileIterator; - using Iterator = AccumLambdaIteratorSimt; -}; - -// TensorOp - Volta -template -struct DefaultMmaAccumLambdaIterator< - cutlass::gemm::warp::MmaVoltaTensorOpAccumulatorTileIterator>, - accum_t, - kWarpSize> { - using WarpIterator = typename cutlass::gemm::warp::MmaVoltaTensorOpAccumulatorTileIterator< - S1, - accum_t, - cutlass::layout::RowMajor, - S2, - cutlass::MatrixShape<1, 1>>; - using Iterator = AccumLambdaIteratorSm70; -}; - -// TensorOp - Sm75+ -template -struct DefaultMmaAccumLambdaIterator< - cutlass::gemm::warp:: - MmaTensorOpAccumulatorTileIterator, - accum_t, - kWarpSize> { - using WarpIterator = typename cutlass::gemm::warp:: - MmaTensorOpAccumulatorTileIterator; - using Iterator = AccumLambdaIteratorSm80; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_from_smem.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_from_smem.h deleted file mode 100644 index 40d3265c7a638396ca5c12282f7da56fa63c925b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm/mma_from_smem.h +++ /dev/null @@ -1,1939 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights - *reserved. SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, - *this list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE - *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE - *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR - *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF - *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS - *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN - *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) - *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE - *POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/*! \file - \brief Template for a double-buffered threadblock-scoped GEMM kernel. -*/ - -#pragma once - -#include "cutlass/aligned_buffer.h" -#include "cutlass/arch/memory.h" -#include "cutlass/array.h" -#include "cutlass/cutlass.h" -#include "cutlass/epilogue/thread/linear_combination.h" -#include "cutlass/epilogue/threadblock/default_epilogue_simt.h" -#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h" -#include "cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" -#include "cutlass/functional.h" -#include "cutlass/gemm/gemm.h" -#include "cutlass/gemm/warp/mma_tensor_op_fragment_iterator.h" -#include "cutlass/matrix_shape.h" -#include "cutlass/numeric_conversion.h" -#include "cutlass/numeric_types.h" -#include "cutlass/platform/platform.h" -#include "cutlass/transform/threadblock/vector_iterator.h" - -#include "../epilogue/epilogue_thread_apply_logsumexp.h" -#include "../gemm/mma_accum_lambda_iterator.h" -#include "../gemm_kernel_utils.h" -#include "../iterators/make_residual_last.h" -#include "../iterators/transpose_warp_iterator.h" -#include "../iterators/warp_iterator_from_smem.h" -#include "cutlass/epilogue/threadblock/epilogue_smem_accumulator.h" -#include "cutlass/gemm/threadblock/mma_base.h" -#include "cutlass/gemm/threadblock/mma_multistage.h" -#include "cutlass/gemm/threadblock/mma_pipelined.h" -#include "cutlass/gemm/warp/mma_tensor_op_tile_access_iterator.h" - -namespace cutlass { -namespace gemm { -namespace threadblock { - -/// Shared storage object needed by accumulator -/// From 13_two_tensor_op_fusion/threadblock/b2b_mma_base_smem_accumulator.h -template -class AccumulatorSharedStorage { -public: - // - // Type definitions - // - using Shape = Shape_; - using Element = Element_; - using Layout = Layout_; - using Padding = Padding_; - - /// Tensor reference to the accumulator - using TensorRefAccum = cutlass::TensorRef; - - /// Shape of the accumulator matrix in shared memory - using ShapeAccum = - cutlass::MatrixShape; - -public: - // - // Data members - // - - /// Buffer for accumulator - cutlass::AlignedBuffer accum; - -public: - // - // Methods - // - - /// Returns a layout object for the Accum matrix - CUTLASS_DEVICE - static Layout LayoutAccum() { return Layout::packed({ShapeAccum::kRow, ShapeAccum::kColumn}); } - - /// Returns a TensorRef to the Accumulator - CUTLASS_HOST_DEVICE - TensorRefAccum accum_ref() { return TensorRefAccum{accum.data(), LayoutAccum()}; } -}; - -//////////////////////////////////////////////////////////////////////////////// -// Taken from -// https://github.com/NVIDIA/cutlass/blob/master/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_base_smem_accumulator.h -//////////////////////////////////////////////////////////////////////////////// - -/// Structure to compute the matrix product targeting CUDA cores and SIMT math -/// instructions. -template < - /// Size of the Gemm problem - concept: gemm::GemmShape<> - typename Shape_, - // Maximum value for K - int kMaxK, - /// Policy describing tuning details (concept: MmaPolicy) - typename Policy_, - /// Number of stages, - int Stages, - /// Used for partial specialization - typename Enable = bool> -class MmaBaseFromSharedMemory { -public: - ///< Size of the Gemm problem - concept: gemm::GemmShape<> - using Shape = Shape_; - - ///< Policy describing tuning details - using Policy = Policy_; - - // - // Dependent types - // - - /// Warp-level Mma - using Operator = typename Policy::Operator; - - /// Shape describing the overall GEMM computed from shared memory - /// by each warp. - using WarpGemm = typename Policy::Operator::Shape; - - /// Shape describing the number of warps filling the CTA - using WarpCount = - GemmShape; - using WarpCount1 = WarpCount; - - /// Number of warp-level GEMM oeprations - static int const kWarpGemmIterations = (WarpGemm::kK / Operator::Policy::MmaShape::kK); - static int const kWarpGemmIterations1 = kWarpGemmIterations; - - /// Number of stages - static int const kStages = Stages; - - /// If this is true, we fill the entire shmem buffer at start - /// and don't need to iterate through it in a circular fashion - static bool const kSmemContainsEntireB = kMaxK <= Shape::kK * kStages; - - /// Tensor reference to the A operand - using TensorRefA = TensorRef; - - /// Tensor reference to the B operand - using TensorRefB = TensorRef; - - // - // Nested structs - // - - /// Shared storage object needed by threadblock-scoped GEMM - class SharedStorage { - public: - // - // Type definitions - // - - /// Shape of the B matrix operand in shared memory - using ShapeB = MatrixShape; - - public: - // - // Data members - // - - /// Buffer for B operand - AlignedBuffer operand_B; - - public: - // - // Methods - // - - /// Returns a layout object for the B matrix - CUTLASS_HOST_DEVICE - static typename Operator::LayoutB LayoutB() - { - return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn}); - } - - /// Returns a TensorRef to the B operand - CUTLASS_HOST_DEVICE - TensorRefB operand_B_ref() { return TensorRefB{operand_B.data(), LayoutB()}; } - }; - -protected: - // - // Data members - // - - // /// Iterator to load a warp-scoped tile of A operand from shared memory - // typename Operator::IteratorA warp_tile_iterator_A_; - - /// Iterator to load a warp-scoped tile of B operand from shared memory - typename Operator::IteratorB warp_tile_iterator_B_; - -public: - /// Construct from tensor references - CUTLASS_DEVICE - MmaBaseFromSharedMemory( - ///< Shared storage needed for internal use by threadblock-scoped GEMM - SharedStorage& shared_storage, - ///< ID within the threadblock - int thread_idx, - ///< ID of warp - int warp_idx, - ///< ID of each thread within a warp - int lane_idx) - : warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx) - { - } -}; - -namespace { - -// has necessary trait compliance with WarpIteratorFromSmem but doesn't do -// anything, can be default initialized, and uses fragment that takes up -// (almost) no space. this warp iterator is selected at compile time when -// elementwise on-the-fly scaling for operand A is disabled, in which case -// operations related to loading scale factors for operand A get wiped out by -// the compiler. -template -class NoOpWarpIteratorScale { -public: - // in pipelined+multistage MMA implementations we keep an array of fragments. - // if we aren't using scaling we don't want to waste registers on fragments - // of scale elements, so ideally this would be sized 0. - // Since arrays of zero-sized objects are not allowed, using size as 1. - // The compiler will most likely wipe it out anyways. - using Fragment = cutlass::Array; - - CUTLASS_HOST_DEVICE - NoOpWarpIteratorScale() {} - - CUTLASS_HOST_DEVICE - NoOpWarpIteratorScale(TensorRef const&, int) {} - - CUTLASS_HOST_DEVICE - NoOpWarpIteratorScale& add_tile_offset(typename TensorRef::TensorCoord const&) { return *this; } - - CUTLASS_HOST_DEVICE - NoOpWarpIteratorScale& operator++() { return *this; } - - CUTLASS_DEVICE - void load(Fragment&) const {} -}; - -// if scaling is enabled, performs fragment elementwise multiplication between -// fragment and its scaling factor. -template -class FragmentElementwiseScaler; - -// specialization for scaling being enabled. -template -class FragmentElementwiseScaler { -public: - // cast scale_frag to correct type then apply elementwise to fragment - CUTLASS_DEVICE - static Fragment apply(Fragment frag, FragmentScale const& scale_frag) - { - Fragment converted_scale_frag = - cutlass::NumericArrayConverter()(scale_frag); - return cutlass::multiplies()(frag, converted_scale_frag); - } -}; - -// specialization for scaling being disabled. doesn't do anything and should -// just get wiped out by the compiler. -template -class FragmentElementwiseScaler { -public: - CUTLASS_DEVICE - static Fragment apply(Fragment frag, FragmentScale const&) { return frag; } -}; -} // namespace - -//////////////////////////////////////////////////////////////////////////////// -// Taken from -// https://github.com/NVIDIA/cutlass/blob/master/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_pipelined_smem_accumulator.h -//////////////////////////////////////////////////////////////////////////////// - -/// Structure to compute the matrix product targeting CUDA cores and SIMT math -/// instructions. -template < - /// Size of the Gemm problem - concept: gemm::GemmShape<> - typename Shape_, - // BEGIN smem - /// Iterates over the intermediate accumulator tile in shared memory - typename WarpIteratorA, - /// whether or not to perform elementwise multiplication of A - // by another matrix (A_scale) that is also kept in shared memory prior - // to matmul A @ B - bool ScaleOperandA_, - // Accumulator type - typename AccumulatorSharedStorage, - // END smem - /// Iterates over tiles of B operand in global memory - // (concept: ReadableTileIterator | ForwardTileIterator | - // MaskedTileIterator) - typename IteratorB_, - /// Iterates over tiles of B operand in shared memory - /// (concept: WriteableTileIterator | RandomAccessTileIterator) - typename SmemIteratorB_, - /// Data type of accumulator matrix - typename ElementC_, - /// Data type of accumulator matrix - typename LayoutC_, - /// Policy describing tuning details (concept: MmaPolicy) - typename Policy_, - /// Transformation applied to B operand - typename TransformB_ = NumericArrayConverter, - /// Used for partial specialization - typename Enable = bool> -class MmaPipelinedFromSharedMemory - : public MmaBaseFromSharedMemory { -public: - ///< Base class - using Base = MmaBaseFromSharedMemory; - - using Shape = Shape_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> - static constexpr bool ScaleOperandA = ScaleOperandA_; - - ///< loads fragments of A_scale from shared memory if operand A scaling is - ///< enabled. otherwise no-op. - using WarpIteratorAScale = typename cutlass::platform::conditional< - ScaleOperandA, - WarpIteratorA, - NoOpWarpIteratorScale>::type; - - using IteratorB = IteratorB_; ///< Iterates over tiles of B operand in global memory - using ElementC = ElementC_; ///< Data type of accumulator matrix - using LayoutC = LayoutC_; ///< Layout of accumulator matrix - using Policy = Policy_; ///< Policy describing tuning details - - using SmemIteratorB = SmemIteratorB_; - - using TransformB = TransformB_; - - // - // Dependent types - // - - /// Fragment of operand B loaded from global memory - using FragmentB = typename IteratorB::Fragment; - - /// Fragment of accumulator tile - using FragmentC = typename Policy::Operator::FragmentC; - - /// Warp-level Mma - using Operator = typename Policy::Operator; - - /// Obtain the arch tag from the warp-level operator - using ArchTag = typename Policy::Operator::ArchTag; - - /// Complex transform on B operand - static ComplexTransform const kTransformB = Operator::kTransformB; - - // staticaly assert kStages for MmaPipelined is two (Double-buffered pipeline) - static_assert((Base::kStages == 2), "MmaPipelined requires kStages set to value 2"); - -private: - using WarpFragmentA = typename Operator::FragmentA; - - /// fragment type of OperandA elementwise scaling matrix. (almost) empty - /// if operand A scaling is disabled. - using WarpFragmentAScale = typename WarpIteratorAScale::Fragment; - - using WarpFragmentB = typename Operator::FragmentB; - - /// applies scaling factor to operand A fragment if operand A scaling is - /// enabled. otherwise no-op. - using FragmentAScaler = - FragmentElementwiseScaler; - -protected: - // /// Iterator to write threadblock-scoped tile of A operand to shared memory - // SmemIteratorA smem_iterator_A_; - - /// Iterator to write threadblock-scoped tile of B operand to shared memory - SmemIteratorB smem_iterator_B_; - - /// Iterator to load a warp-scoped tile of A operand from intermediate - /// accumulator tile - WarpIteratorA warp_tile_iterator_A_; - - /// Iterator to load a warp-scoped tile of A_scale from intermediate - /// accumulator tile (only used if ScaleOperandA_ is true) - WarpIteratorAScale warp_tile_iterator_A_scale_; - -public: - /// constructor for MMA with operand A scaling enabled. - CUTLASS_DEVICE - MmaPipelinedFromSharedMemory( - // shared storage needed for internal use by threadblock-scoped GEMM - typename Base::SharedStorage& shared_storage, - // warp iterator over A tile held in shared memory - WarpIteratorA warp_iter_a, - // warp iterator over A_scale tile held in shared memory - WarpIteratorAScale warp_iter_a_scale, - int thread_idx, - int warp_idx, - int lane_idx) - : Base(shared_storage, thread_idx, warp_idx, lane_idx), - warp_tile_iterator_A_(warp_iter_a), - warp_tile_iterator_A_scale_(warp_iter_a_scale), - smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) - { - // Compute warp location within threadblock tile by mapping the warp_id to - // three coordinates: - // _m: the warp's position within the threadblock along the M dimension - // _n: the warp's position within the threadblock along the N dimension - // _k: the warp's position within the threadblock along the K dimension - int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); - int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); - int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; - int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; - - // Add per-warp offsets in units of warp-level tiles - this->warp_tile_iterator_A_.add_tile_offset( - {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); - this->warp_tile_iterator_A_scale_.add_tile_offset( - {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); - this->warp_tile_iterator_B_.add_tile_offset( - {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); - } - - /// Construct from tensor references - CUTLASS_DEVICE - MmaPipelinedFromSharedMemory( - typename Base::SharedStorage& shared_storage, ///< Shared storage needed for internal use - ///< by threadblock-scoped GEMM - AccumulatorSharedStorage& accumulator_shared_storage, - int thread_idx, ///< ID within the threadblock - int warp_idx, ///< ID of warp - int lane_idx, ///< ID of each thread within a warp - int problem_size_0_n) - : Base(shared_storage, thread_idx, warp_idx, lane_idx), - warp_tile_iterator_A_(accumulator_shared_storage.accum_ref(), lane_idx), - smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) - { - // Compute warp location within threadblock tile by mapping the warp_id to - // three coordinates: - // _m: the warp's position within the threadblock along the M dimension - // _n: the warp's position within the threadblock along the N dimension - // _k: the warp's position within the threadblock along the K dimension - - int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); - int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); - - int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; - int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; - - // Add per-warp offsets in units of warp-level tiles - this->warp_tile_iterator_A_.add_tile_offset( - {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); - this->warp_tile_iterator_B_.add_tile_offset( - {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); - } - - // For API compatibility with MmaMultistageFromSharedMemory - // but not supported as it worsens perf: older gpus < sm80 don't - // support async transfers and have to waste registers - CUTLASS_DEVICE - void set_prologue_done(bool value) {} - CUTLASS_DEVICE - static void prologue(typename Base::SharedStorage& shared_storage, - IteratorB iterator_B1, - int thread_idx, - int problem_size_0_n) - { - } - - CUTLASS_DEVICE - static void drain_cp_asyncs() {} - - /// Perform a threadblock-scoped matrix multiply-accumulate - CUTLASS_DEVICE - void operator()(int gemm_k_iterations, ///< number of iterations of the mainloop - FragmentC& accum, ///< destination accumulator tile - // IteratorA iterator_A, ///< iterator over A - // operand in global memory - IteratorB iterator_B, ///< iterator over B operand in global memory - FragmentC const& src_accum, ///< source accumulator tile - // TransformA transform_A = TransformA(), ///< transformation - // applied to A fragment - TransformB transform_B = TransformB()) - { ///< transformation applied to B fragment - - // - // Prologue - // - - // Perform accumulation in the 'd' output operand - accum = src_accum; - - FragmentB tb_frag_B; - - tb_frag_B.clear(); - - // The last kblock is loaded in the prolog - iterator_B.set_residual_tile(gemm_k_iterations == 1); - iterator_B.load(tb_frag_B); - - ++iterator_B; - - this->smem_iterator_B_.store(transform_B(tb_frag_B)); - - ++this->smem_iterator_B_; - - __syncthreads(); - - // remember that WarpFragmentAScale and WarpIteratorAScale are empty/no-op - // if scaling is disabled. - - // Pair of fragments used to overlap shared memory loads and math - // instructions - WarpFragmentA warp_frag_A[2]; - WarpFragmentAScale warp_frag_A_scale[2]; - WarpFragmentB warp_frag_B[2]; - warp_frag_A[0].clear(); - warp_frag_A_scale[0].clear(); - warp_frag_B[0].clear(); - - this->warp_tile_iterator_B_.set_kgroup_index(0); - - this->warp_tile_iterator_A_.load(warp_frag_A[0]); - this->warp_tile_iterator_A_scale_.load(warp_frag_A_scale[0]); - this->warp_tile_iterator_B_.load(warp_frag_B[0]); - - ++this->warp_tile_iterator_A_; - ++this->warp_tile_iterator_A_scale_; - ++this->warp_tile_iterator_B_; - - Operator warp_mma; - - int smem_write_stage_idx = 1; - - // Avoid reading out of bounds - iterator_B.set_residual_tile(gemm_k_iterations == 2); - iterator_B.clear_mask(gemm_k_iterations <= 1); - - // Issue loads during the first warp-level matrix multiply-add *AFTER* - // issuing shared memory loads (which have the tightest latency requirement). - - // - // Mainloop - // - - // Note: The main loop does not support Base::kWarpGemmIterations == 2. - CUTLASS_GEMM_LOOP - for (; gemm_k_iterations > 0; --gemm_k_iterations) { - // - // Loop over GEMM K dimension - // - - CUTLASS_PRAGMA_UNROLL - for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) { - // Load warp-level tiles from shared memory, wrapping to k offset if - // this is the last group as the case may be. - bool hasNext = true; - - if (warp_mma_k == Base::kWarpGemmIterations - 1) { - // Write fragments to shared memory - this->smem_iterator_B_.store(transform_B(tb_frag_B)); - - __syncthreads(); - - ++this->smem_iterator_B_; - - // Add negative offsets to return iterators to the 'start' of the - // circular buffer in shared memory SMEM: Don't reset iterator A, as - // we are continuing our iteration at this point - if (smem_write_stage_idx == 1) { - this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); - } else { - this->warp_tile_iterator_B_.add_tile_offset( - {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations, 0}); - } - - smem_write_stage_idx ^= 1; - hasNext = gemm_k_iterations > 1; - } - - // Only read the next if we need to - if (hasNext) { - this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % - Base::kWarpGemmIterations); - - this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]); - this->warp_tile_iterator_A_scale_.load(warp_frag_A_scale[(warp_mma_k + 1) % 2]); - this->warp_tile_iterator_B_.load(warp_frag_B[(warp_mma_k + 1) % 2]); - - ++this->warp_tile_iterator_A_; - ++this->warp_tile_iterator_A_scale_; - ++this->warp_tile_iterator_B_; - - if (warp_mma_k == 0) { - iterator_B.load(tb_frag_B); - - ++iterator_B; - - // Avoid reading out of bounds if this was the last loop iteration - iterator_B.set_residual_tile(gemm_k_iterations == 3); - iterator_B.clear_mask(gemm_k_iterations <= 2); - } - } - - warp_mma(accum, - FragmentAScaler::apply(warp_frag_A[warp_mma_k % 2], - warp_frag_A_scale[warp_mma_k % 2]), - warp_frag_B[warp_mma_k % 2], - accum); - } - } - } -}; - -//////////////////////////////////////////////////////////////////////////////// -// Taken from -// https://github.com/NVIDIA/cutlass/blob/master/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_multistage_smem_accumulator.h -//////////////////////////////////////////////////////////////////////////////// - -/// Structure to compute the matrix product targeting CUDA cores and SIMT math -/// instructions. -template < - /// Size of the Gemm problem - concept: gemm::GemmShape<> - typename Shape1_, - /// Iterates over the intermediate accumulator tile in shared memory - typename WarpIteratorA1_, - /// whether or not to perform elementwise multiplication of A - // by another matrix (A_scale) that is also kept in shared memory prior - // to matmul A @ B - bool ScaleOperandA_, - // Accumulator type - typename AccumulatorSharedStorage, - /// Iterates over tiles of B operand in global memory - // (concept: ReadableTileIterator | ForwardTileIterator | - // MaskedTileIterator) - typename IteratorB1_, - /// Iterates over tiles of B operand in shared memory - /// (concept: WriteableTileIterator | RandomAccessTileIterator) - typename SmemIteratorB1_, - /// Cache operation for operand B - cutlass::arch::CacheOperation::Kind CacheOpB1, - /// Data type of accumulator matrix - typename ElementC_, - /// Data type of accumulator matrix - typename LayoutC_, - /// Policy describing tuning details (concept: MmaPolicy) - typename Policy1_, - /// Number of stages, - int Stages_, - int kMaxK_, - /// Used for partial specialization - typename Enable = bool> -class MmaMultistageFromSharedMemory - : public MmaBaseFromSharedMemory { -public: - ///< Base class - using Base = MmaBaseFromSharedMemory; - - ///< Size of the Gemm problem - concept: gemm::GemmShape<> - using Shape1 = Shape1_; - ///< Iterates over tiles of B operand in global memory - using IteratorB1 = IteratorB1_; - using IteratorB = IteratorB1; - ///< Policy describing tuning details - using Policy1 = Policy1_; - - using SmemIteratorB1 = SmemIteratorB1_; - using WarpIteratorA1 = WarpIteratorA1_; ///< Iterates over the intermediate - ///< accumulator tile in shared memory - static constexpr bool ScaleOperandA = ScaleOperandA_; - - ///< warp level iterator over A_scale matrix tile kept in shared memory. - ///< if elementwise A scaling is disabled then everything this does is no-op. - using WarpIteratorAScale = typename cutlass::platform::conditional< - ScaleOperandA, - WarpIteratorA1, - NoOpWarpIteratorScale>::type; - ///< Data type of accumulator matrix - using ElementC = ElementC_; - ///< Layout of accumulator matrix - using LayoutC = LayoutC_; - - static cutlass::arch::CacheOperation::Kind const kCacheOpB1 = CacheOpB1; - static constexpr bool kSmemContainsEntireB = Base::kSmemContainsEntireB; - - // - // Dependent types - // - - /// Fragment of accumulator tile - using FragmentC1 = typename Policy1::Operator::FragmentC; - using FragmentC = FragmentC1; - - /// Warp-level Mma - using Operator1 = typename Policy1::Operator; - - /// Minimum architecture is Sm80 to support cp.async - using ArchTag = arch::Sm80; - - /// Complex transform on B operand - static ComplexTransform const kTransformB1 = Operator1::kTransformB; - - /// Internal structure exposed for introspection. - struct Detail { - static_assert(Base::kWarpGemmIterations1 > 1, - "The pipelined structure requires at least two warp-level " - "GEMM operations."); - - /// Number of cp.async instructions to load one stage of operand B - static int const TBLoadIterationsB1 = IteratorB1::ThreadMap::Iterations::kCount; - - /// Number of cp.async instructions to load on group of operand B - static int const kAccessesPerGroupB1 = - (TBLoadIterationsB1 + Base::kWarpGemmIterations1 - 1) / Base::kWarpGemmIterations1; - }; - - static constexpr int kNumStagesConcurrentLoad = kSmemContainsEntireB ? Base::kStages - : Base::kStages - 1; - -private: - using WarpLoadedFragmentA1 = typename Operator1::FragmentA; - /// fragment of OperandA scale matrix. if operand A scaling is disabled this - /// is (almost) empty. - using WarpLoadedFragmentA1Scale = typename WarpIteratorAScale::Fragment; - using WarpLoadedFragmentB1 = typename Operator1::FragmentB; - using WarpTransformedFragmentA1 = typename Operator1::TransformedFragmentA; - using WarpTransformedFragmentB1 = typename Operator1::TransformedFragmentB; - - /// applies elementwise scaling to fragment of A. if operand A scaling is - /// disabled this is a no-op. - using FragmentAScaler = - FragmentElementwiseScaler; - -private: - // - // Data members - // - - /// Iterator to load a warp-scoped tile of A1 operand from intermediate - /// accumulator tile - WarpIteratorA1 warp_tile_iterator_A1_; - - /// Iterator to load a warp-scoped tile of A1_scale operand from shared memory - /// if operand A scaling is disabled everything this does is a no-op. - WarpIteratorAScale warp_tile_iterator_A1_scale_; - - /// Iterator to write threadblock-scoped tile of B operand to shared memory - SmemIteratorB1 smem_iterator_B1_; - - bool prologue_done_; - -public: - /// constructor for MMA with operand A scaling enabled. - CUTLASS_DEVICE - MmaMultistageFromSharedMemory( - // shared storage needed for internal use by threadblock-scoped GEMM - typename Base::SharedStorage& shared_storage, - // warp level iterator over operand A tile kept in shared memory - WarpIteratorA1 warp_tile_iterator_A1, - // warp level iterator over operand A elementwise scale tile kept in - // shared memory. - WarpIteratorAScale warp_tile_iterator_A1_scale, - int thread_idx, - int warp_idx, - int lane_idx) - : Base(shared_storage, thread_idx, warp_idx, lane_idx), - warp_tile_iterator_A1_(warp_tile_iterator_A1), - warp_tile_iterator_A1_scale_(warp_tile_iterator_A1_scale), - smem_iterator_B1_(shared_storage.operand_B_ref(), thread_idx), - prologue_done_(false) - { - // Compute warp location within threadblock tile by mapping the warp_id to - // three coordinates: - // _m: the warp's position within the threadblock along the M dimension - // _n: the warp's position within the threadblock along the N dimension - // _k: the warp's position within the threadblock along the K dimension - int warp_idx_mn_1 = warp_idx % (Base::WarpCount1::kM * Base::WarpCount1::kN); - int warp_idx_k_1 = warp_idx / (Base::WarpCount1::kM * Base::WarpCount1::kN); - int warp_idx_m_1 = warp_idx_mn_1 % Base::WarpCount1::kM; - int warp_idx_n_1 = warp_idx_mn_1 / Base::WarpCount1::kM; - - // Add per-warp offsets in units of warp-level tiles - warp_tile_iterator_A1_.add_tile_offset( - {warp_idx_m_1, Base::kWarpGemmIterations1 * warp_idx_k_1}); - warp_tile_iterator_A1_scale_.add_tile_offset( - {warp_idx_m_1, Base::kWarpGemmIterations1 * warp_idx_k_1}); - this->warp_tile_iterator_B_.add_tile_offset( - {Base::kWarpGemmIterations1 * warp_idx_k_1, warp_idx_n_1}); - } - - /// Construct from tensor references - CUTLASS_DEVICE - MmaMultistageFromSharedMemory( - typename Base::SharedStorage& shared_storage, ///< Shared storage needed for internal use - ///< by threadblock-scoped GEMM - AccumulatorSharedStorage& accumulator_shared_storage, - ///< ID within the threadblock - int thread_idx, - ///< ID of warp - int warp_idx, - ///< ID of each thread within a warp - int lane_idx, - ///< GEMM0 N is used for accumulator extent - int problem_size_0_n) - : Base(shared_storage, thread_idx, warp_idx, lane_idx), - warp_tile_iterator_A1_(accumulator_shared_storage.accum_ref(), lane_idx), - smem_iterator_B1_(shared_storage.operand_B_ref(), thread_idx), - prologue_done_(false) - { - // Compute warp location within threadblock tile by mapping the warp_id to - // three coordinates: - // _m: the warp's position within the threadblock along the M dimension - // _n: the warp's position within the threadblock along the N dimension - // _k: the warp's position within the threadblock along the K dimension - - int warp_idx_mn_1 = warp_idx % (Base::WarpCount1::kM * Base::WarpCount1::kN); - int warp_idx_k_1 = warp_idx / (Base::WarpCount1::kM * Base::WarpCount1::kN); - - int warp_idx_m_1 = warp_idx_mn_1 % Base::WarpCount1::kM; - int warp_idx_n_1 = warp_idx_mn_1 / Base::WarpCount1::kM; - - // Add per-warp offsets in units of warp-level tiles - warp_tile_iterator_A1_.add_tile_offset( - {warp_idx_m_1, Base::kWarpGemmIterations1 * warp_idx_k_1}); - this->warp_tile_iterator_B_.add_tile_offset( - {Base::kWarpGemmIterations1 * warp_idx_k_1, warp_idx_n_1}); - } - - CUTLASS_DEVICE - void set_prologue_done(bool value) { prologue_done_ = value; } - - CUTLASS_DEVICE - static void prologue(typename Base::SharedStorage& shared_storage, - IteratorB iterator_B1, - int thread_idx, - int problem_size_0_n) - { - SmemIteratorB1 smem_iterator_B1(shared_storage.operand_B_ref(), thread_idx); - _prologue(iterator_B1, - (problem_size_0_n + Base::Shape::kK - 1) / Base::Shape::kK, - smem_iterator_B1); - } - - CUTLASS_DEVICE - static void drain_cp_asyncs() - { - // commit and drain all pending and predicated cp.async pnz from the GEMM - // mainloop - cutlass::arch::cp_async_fence(); - cutlass::arch::cp_async_wait<0>(); - __syncthreads(); - } - - CUTLASS_DEVICE - void copy_tiles_and_advance_1(IteratorB1& iterator_B1, int group_start_B1 = 0) - { - iterator_B1.set_iteration_index(group_start_B1 * IteratorB1::kAccessesPerVector); - this->smem_iterator_B1_.set_iteration_index(group_start_B1); - - // Load for operand B - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < Detail::kAccessesPerGroupB1; ++j) { - if (group_start_B1 + j < Detail::TBLoadIterationsB1) { - typename IteratorB1::AccessType* dst_ptr = - reinterpret_cast( - this->smem_iterator_B1_.get()); - - int const kSrcBytes = sizeof_bits::value * - IteratorB1::ThreadMap::kElementsPerAccess / - IteratorB1::kAccessesPerVector / 8; - - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < IteratorB1::kAccessesPerVector; ++v) { - auto gmem_ptr = iterator_B1.get(); - - cutlass::arch::cp_async_zfill( - dst_ptr + v, gmem_ptr, iterator_B1.valid()); - - ++iterator_B1; - } - ++this->smem_iterator_B1_; - } - } - } - - CUTLASS_DEVICE - static void _prologue(IteratorB& iterator_B1, - int32_t gemm_k_iterations_1, - SmemIteratorB1& smem_iterator_B1_) - { - // Issue several complete stages - CUTLASS_PRAGMA_UNROLL - for (int stage = 0; stage < kNumStagesConcurrentLoad; ++stage, --gemm_k_iterations_1) { - iterator_B1.set_residual_tile(gemm_k_iterations_1 == 1); - iterator_B1.clear_mask(gemm_k_iterations_1 == 0); - - iterator_B1.set_iteration_index(0); - smem_iterator_B1_.set_iteration_index(0); - - // Load for operand B - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < Detail::TBLoadIterationsB1; ++j) { - typename IteratorB1::AccessType* dst_ptr = - reinterpret_cast(smem_iterator_B1_.get()); - - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < IteratorB1::kAccessesPerVector; ++v) { - int const kSrcBytes = sizeof_bits::value * - IteratorB1::ThreadMap::kElementsPerAccess / - IteratorB1::kAccessesPerVector / 8; - - cutlass::arch::cp_async_zfill( - dst_ptr + v, iterator_B1.get(), iterator_B1.valid()); - - ++iterator_B1; - } - - ++smem_iterator_B1_; - } - - // Move to the next stage - iterator_B1.add_tile_offset({1, 0}); - - smem_iterator_B1_.add_tile_offset({1, 0}); - - // Defines the boundary of a stage of cp.async. - cutlass::arch::cp_async_fence(); - } - iterator_B1.set_residual_tile(gemm_k_iterations_1 == 1); - iterator_B1.clear_mask(gemm_k_iterations_1 == 0); - } - - /// Perform a threadblock-scoped matrix multiply-accumulate - CUTLASS_DEVICE - void operator()( - ///< problem size of GEMM - int gemm_k_iterations_1_, - ///< destination accumulator tile - FragmentC1& accum, - ///< iterator over B1 operand in global memory - IteratorB1 iterator_B1, - ///< initial value of accumulator - FragmentC1 const& src_accum) - { - // 2nd Gemm - - // - // Prologue - // - // Perform accumulation in the 'd' output operand - accum = src_accum; - - if (!prologue_done_) { - _prologue(iterator_B1, gemm_k_iterations_1_, smem_iterator_B1_); - } else if (!kSmemContainsEntireB) { - // Restore the iterators increments - - int gemm_k_iterations_1 = gemm_k_iterations_1_; - // Issue several complete stages - CUTLASS_PRAGMA_UNROLL - for (int stage = 0; stage < kNumStagesConcurrentLoad; ++stage, --gemm_k_iterations_1) { - iterator_B1.set_iteration_index(0); - this->smem_iterator_B1_.set_iteration_index(0); - - // Load for operand B - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < Detail::TBLoadIterationsB1; ++j) { - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < IteratorB1::kAccessesPerVector; ++v) { ++iterator_B1; } - ++this->smem_iterator_B1_; - } - iterator_B1.add_tile_offset({1, 0}); - this->smem_iterator_B1_.add_tile_offset({1, 0}); - } - iterator_B1.set_residual_tile(gemm_k_iterations_1 <= 1); - iterator_B1.clear_mask(gemm_k_iterations_1 <= 0); - } - - // DEPBAR+SYNC - cutlass::arch::cp_async_wait(); - __syncthreads(); - - // remember that WarpFragmentAScale and WarpIteratorAScale are no-op/empty - // if scaling is disabled. - - // Pair of fragments used to overlap shared memory loads and math - // instructions - WarpLoadedFragmentA1 warp_loaded_frag_A1[2]; - WarpLoadedFragmentA1Scale warp_loaded_frag_A1_scale[2]; - WarpLoadedFragmentB1 warp_loaded_frag_B1[2]; - WarpTransformedFragmentA1 warp_transformed_frag_A1[2]; - WarpTransformedFragmentB1 warp_transformed_frag_B1[2]; - - Operator1 warp_mma1; - - warp_tile_iterator_A1_.load(warp_loaded_frag_A1[0]); - ++warp_tile_iterator_A1_; - - warp_tile_iterator_A1_scale_.load(warp_loaded_frag_A1_scale[0]); - ++warp_tile_iterator_A1_scale_; - - this->warp_tile_iterator_B_.set_kgroup_index(0); - this->warp_tile_iterator_B_.load(warp_loaded_frag_B1[0]); - ++this->warp_tile_iterator_B_; - - int smem_write_stage_idx = Base::kStages - 1; - int smem_read_stage_idx = 0; - - warp_mma1.transform( - warp_transformed_frag_A1[0], - warp_transformed_frag_B1[0], - FragmentAScaler::apply(warp_loaded_frag_A1[0], warp_loaded_frag_A1_scale[0]), - warp_loaded_frag_B1[0]); - - // tf32x3 kernels use staging accumulation. warp_mma uses a temporary - // accumulator and this temporary accumulator is added to the final - // accumulator once in every mainloop iteration. - plus plus_accum; - - FragmentC1 tmp_accum; - - if (platform::is_same::value || - platform::is_same::value) { - tmp_accum.clear(); - } - - // - // Mainloop - // - - CUTLASS_PRAGMA_UNROLL - for (int gemm_k_iterations_1 = gemm_k_iterations_1_ - (Base::kStages - 1); - gemm_k_iterations_1 > (-Base::kStages + 1); - gemm_k_iterations_1--) { - // - // Loop over GEMM K dimension - // - - // Computes a warp-level GEMM on data held in shared memory - // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate - CUTLASS_PRAGMA_UNROLL - for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations1; ++warp_mma_k) { - // Load warp-level tile from accumulator fragment (A) - // or shared memory (operand B) - this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % - Base::kWarpGemmIterations1); - // skip warp tile loading for the last kgroup (we are out of the buf) - if (gemm_k_iterations_1 > (-Base::kStages + 2) || - warp_mma_k < Base::kWarpGemmIterations1 - 1) { - warp_tile_iterator_A1_.load(warp_loaded_frag_A1[(warp_mma_k + 1) % 2]); - warp_tile_iterator_A1_scale_.load( - warp_loaded_frag_A1_scale[(warp_mma_k + 1) % 2]); - this->warp_tile_iterator_B_.load(warp_loaded_frag_B1[(warp_mma_k + 1) % 2]); - } - ++warp_tile_iterator_A1_; - ++warp_tile_iterator_A1_scale_; - ++this->warp_tile_iterator_B_; - - if (warp_mma_k > 0) - warp_mma1.transform( - warp_transformed_frag_A1[warp_mma_k % 2], - warp_transformed_frag_B1[warp_mma_k % 2], - FragmentAScaler::apply(warp_loaded_frag_A1[warp_mma_k % 2], - warp_loaded_frag_A1_scale[warp_mma_k % 2]), - warp_loaded_frag_B1[warp_mma_k % 2]); - - if (platform::is_same::value || - platform::is_same::value) { - warp_mma1(tmp_accum, - warp_transformed_frag_A1[warp_mma_k % 2], - warp_transformed_frag_B1[warp_mma_k % 2], - tmp_accum); - - if (warp_mma_k == 0) { - accum = plus_accum(accum, tmp_accum); - tmp_accum.clear(); - } - } else { - warp_mma1(accum, - warp_transformed_frag_A1[warp_mma_k % 2], - warp_transformed_frag_B1[warp_mma_k % 2], - accum); - } - - // Issue global->shared copies for the this stage - if (warp_mma_k < Base::kWarpGemmIterations1 - 1) { - int group_start_iteration_B1; - - group_start_iteration_B1 = warp_mma_k * Detail::kAccessesPerGroupB1; - - if (!kSmemContainsEntireB) { - copy_tiles_and_advance_1(iterator_B1, group_start_iteration_B1); - } - } - - if (warp_mma_k + 2 == Base::kWarpGemmIterations1) { - int group_start_iteration_B1; - group_start_iteration_B1 = (warp_mma_k + 1) * Detail::kAccessesPerGroupB1; - - if (!kSmemContainsEntireB) { - copy_tiles_and_advance_1(iterator_B1, group_start_iteration_B1); - } - - // Inserts a memory fence between stages of cp.async instructions. - cutlass::arch::cp_async_fence(); - - // Waits until kStages-2 stages have committed. - arch::cp_async_wait(); - __syncthreads(); - - // Move to the next stage - iterator_B1.add_tile_offset({1, 0}); - - this->smem_iterator_B1_.add_tile_offset({1, 0}); - - // Add negative offsets to return iterators to the 'start' of the - // circular buffer in shared memory - if (!kSmemContainsEntireB) { - if (smem_write_stage_idx == (Base::kStages - 1)) { - this->smem_iterator_B1_.add_tile_offset({-Base::kStages, 0}); - smem_write_stage_idx = 0; - } else { - ++smem_write_stage_idx; - } - - if (smem_read_stage_idx == (Base::kStages - 1)) { - this->warp_tile_iterator_B_.add_tile_offset( - {-Base::kStages * Policy1::kPartitionsK * - Base::kWarpGemmIterations1, - 0}); - smem_read_stage_idx = 0; - } else { - ++smem_read_stage_idx; - } - } - - iterator_B1.set_residual_tile(gemm_k_iterations_1 == 2); - iterator_B1.clear_mask(gemm_k_iterations_1 == 1); - } - - // Do any conversions feeding the first stage at the end of the loop so - // we can start right away on mma instructions - if (warp_mma_k + 1 == Base::kWarpGemmIterations1) - warp_mma1.transform( - warp_transformed_frag_A1[(warp_mma_k + 1) % 2], - warp_transformed_frag_B1[(warp_mma_k + 1) % 2], - FragmentAScaler::apply(warp_loaded_frag_A1[(warp_mma_k + 1) % 2], - warp_loaded_frag_A1_scale[(warp_mma_k + 1) % 2]), - warp_loaded_frag_B1[(warp_mma_k + 1) % 2]); - } - } - - if (platform::is_same::value || - platform::is_same::value) { - accum = plus_accum(accum, tmp_accum); - } - } -}; - -template -struct DefaultWarpIteratorAFromSharedMemory {}; - -// TensorOp - Ampere half -template -struct DefaultWarpIteratorAFromSharedMemory< - cutlass::gemm::GemmShape<32, 32, 32>, - cutlass::gemm::GemmShape<16, 8, 8>, - RegularWarpIterator, - Policy, - typename platform::enable_if<(sizeof_bits::value == 16 && - Policy::Operator::Policy::OpDelta::kRow == 1)>::type> { - static constexpr auto kWarpSize = 32; - using OpDelta = typename Policy::Operator::Policy::OpDelta; - using WarpShape = cutlass::MatrixShape<32, 32>; - - using WarpIterator = - cutlass::gemm::warp::WarpIteratorFromSmem; -}; - -// TensorOp - Ampere f32 -template -struct DefaultWarpIteratorAFromSharedMemory< - WarpShape, - cutlass::gemm::GemmShape<16, 8, 8>, - RegularWarpIterator, - Policy, - typename platform::enable_if<(sizeof_bits::value != 16 || - Policy::Operator::Policy::OpDelta::kRow != 1)>::type> { - using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; - static constexpr auto kWarpSize = 32; - using OpDelta = typename Policy::Operator::Policy::OpDelta; - - using WarpIterator = cutlass::gemm::warp::MmaTensorOpMultiplicandTileAccessIterator< - cutlass::MatrixShape, - cutlass::gemm::Operand::kA, - typename RegularWarpIterator::Element, - cutlass::layout::RowMajor, - cutlass::MatrixShape, - OpDelta::kRow, - kWarpSize>; -}; - -// TensorOp - Volta -template -struct DefaultWarpIteratorAFromSharedMemory, - RegularWarpIterator, - Policy> { - using InstructionShape = cutlass::gemm::GemmShape<16, 16, 4>; - static constexpr auto kWarpSize = 32; - using OpDelta = typename Policy::Operator::Policy::OpDelta; - - using WarpIterator = cutlass::gemm::warp::MmaVoltaTensorOpMultiplicandTileIterator< - cutlass::MatrixShape<32, 32>, // MatrixShape, - cutlass::gemm::Operand::kA, - typename RegularWarpIterator::Element, - cutlass::layout::RowMajorVoltaTensorOpMultiplicandCrosswise<16, 32>, - cutlass::MatrixShape<16, 4>, - OpDelta::kRow, - kWarpSize>; -}; - -// Simt -template -struct DefaultWarpIteratorAFromSharedMemory, - RegularWarpIterator, - Policy> { - using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>; - static constexpr auto kWarpSize = 32; - - // We just use the same iterator, as we reproduced the same shared-memory - // schema. Just modify it to handle non-complete tiles. - using WarpIterator = RegularWarpIterator; -}; - -// Converts a "regular" Mma into their counterpart from shared memory -template -struct DefaultMmaFromSharedMemory; - -// Mma pipelined -template < - /// Size of the Gemm problem - concept: gemm::GemmShape<> - typename Shape_, - /// Iterates over tiles of A operand in global memory - // (concept: ReadableTileIterator | ForwardTileIterator | - // MaskedTileIterator) - typename IteratorA_, - /// Iterates over tiles of A operand in shared memory - /// (concept: WriteableTileIterator | RandomAccessTileIterator) - typename SmemIteratorA_, - /// Iterates over tiles of B operand in global memory - // (concept: ReadableTileIterator | ForwardTileIterator | - // MaskedTileIterator) - typename IteratorB_, - /// Iterates over tiles of B operand in shared memory - /// (concept: WriteableTileIterator | RandomAccessTileIterator) - typename SmemIteratorB_, - /// Data type of accumulator matrix - typename ElementC_, - /// Data type of accumulator matrix - typename LayoutC_, - /// Policy describing tuning details (concept: MmaPolicy) - typename Policy_, - /// Transformation applied to A operand - typename TransformA_, - /// Transformation applied to B operand - typename TransformB_, - typename AccumulatorSharedStorage_, - /// whether or not to apply elementwise multiplication of operand A by - /// another matrix in shared memory before usage in A @ B - bool kScaleOperandA, - bool kTransposeA> -struct DefaultMmaFromSharedMemory, - AccumulatorSharedStorage_, - kScaleOperandA, - kTransposeA> { - static constexpr int kWarpSize = 32; - using SmemAccumulatorLayout = cutlass::layout::RowMajor; - - using RegularMma = MmaPipelined; - - using WarpShape = typename Policy_::Operator::Shape; - using InstructionShape = typename Policy_::Operator::InstructionShape; - using ArchMmaOperator = typename Policy_::Operator; - - static constexpr bool kIsTransposedA = false; - using WarpIteratorA = - typename DefaultWarpIteratorAFromSharedMemory::WarpIterator; - using IteratorB = - typename cutlass::transform::threadblock::MakeIteratorResidualLast::Iterator; - - using Mma = - typename cutlass::gemm::threadblock::MmaPipelinedFromSharedMemory; -}; - -template < - /// Size of the Gemm problem - concept: gemm::GemmShape<> - typename Shape_, - /// Iterates over tiles of A operand in global memory - // (concept: ReadableTileIterator | ForwardTileIterator | - // MaskedTileIterator) - typename IteratorA_, - /// Iterates over tiles of A operand in shared memory - /// (concept: WriteableTileIterator | RandomAccessTileIterator) - typename SmemIteratorA_, - /// Cache operation for operand A - cutlass::arch::CacheOperation::Kind CacheOpA, - /// Iterates over tiles of B operand in global memory - // (concept: ReadableTileIterator | ForwardTileIterator | - // MaskedTileIterator) - typename IteratorB_, - /// Iterates over tiles of B operand in shared memory - /// (concept: WriteableTileIterator | RandomAccessTileIterator) - typename SmemIteratorB_, - /// Cache operation for operand B - cutlass::arch::CacheOperation::Kind CacheOpB, - /// Data type of accumulator matrix - typename ElementC_, - /// Data type of accumulator matrix - typename LayoutC_, - /// Policy describing tuning details (concept: MmaPolicy) - typename Policy_, - /// Number of stages, - int Stages, - /// Use zfill or predicate for out-of-bound cp.async - SharedMemoryClearOption SharedMemoryClear, - typename AccumulatorSharedStorage_, - /// whether or not to apply elementwise multiplication of operand A by - /// another matrix in shared memory before usage in A @ B - bool kScaleOperandA, - bool kTransposeA> -struct DefaultMmaFromSharedMemory, - AccumulatorSharedStorage_, - kScaleOperandA, - kTransposeA> { - static constexpr int kWarpSize = 32; - - using RegularMma = MmaMultistage; - - using WarpShape = typename Policy_::Operator::Shape; - using InstructionShape = typename Policy_::Operator::InstructionShape; - using WarpIteratorA_ = - typename DefaultWarpIteratorAFromSharedMemory::WarpIterator; - using WarpIteratorTranspose = TransposeWarpIterator; - static constexpr bool kIsTransposedA = WarpIteratorTranspose::kSupportsTranspose && kTransposeA; - using WarpIteratorA = typename platform:: - conditional::type; - - static int constexpr kMaxK = kIsTransposedA ? AccumulatorSharedStorage_::Shape::kM - : AccumulatorSharedStorage_::Shape::kN; - // Reduce the number of stages if we don't need that many - static int constexpr kStagesMax = (kMaxK + int(Shape_::kK) - 1) / int(Shape_::kK); - static int constexpr kStages = cutlass::const_min(Stages, kStagesMax); - - using IteratorB = - typename cutlass::transform::threadblock::MakeIteratorResidualLast::Iterator; - using Mma = typename cutlass::gemm::threadblock::MmaMultistageFromSharedMemory< - Shape_, - WarpIteratorA, - kScaleOperandA, - AccumulatorSharedStorage_, - IteratorB, - SmemIteratorB_, - RegularMma::kCacheOpB, - ElementC_, - LayoutC_, - Policy_, - kStages, - kMaxK>; -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -template -struct B2bGemm; - -// Tensor Cores >= Sm75 specialization (Ampere ...) -template < /// Size of the matrix to load (concept: MatrixShape) - typename Shape_, - /// Element type - typename Element_, - /// Layout of operand in memory - typename Layout_, - /// Shape of one matrix product operation (concept: MatrixShape) - typename InstructionShape_, - /// Interval between adjacent *MMA instructions (in units of MMA - /// instructions, concept: MatrixShape) - typename OpDelta_, - typename Operator, - typename scalar_t, - typename WarpShape_, - typename ThreadblockShape_> -struct B2bGemm< - cutlass::gemm::warp:: - MmaTensorOpAccumulatorTileIterator, - Operator, - scalar_t, - WarpShape_, - ThreadblockShape_> { - using IteratorC = typename cutlass::gemm::warp:: - MmaTensorOpAccumulatorTileIterator; - using FragmentC = typename IteratorC::Fragment; - using InstructionShape = InstructionShape_; - using WarpShape = WarpShape_; - using ThreadblockShape = ThreadblockShape_; - using accum_t = Element_; - using lse_scalar_t = float; - - using SmemAccumulatorLayout = cutlass::layout::RowMajor; - - // Iterator to load accumulators (results of matmul in registers) - using FragmentIteratorAccumulator = cutlass::epilogue::warp::FragmentIteratorTensorOp< - WarpShape, - InstructionShape, - accum_t, - typename Operator::Policy::Operator::FragmentC, - cutlass::layout::RowMajor>; - - // Iterator to store to shared-memory - using SmemIteratorD0 = - typename cutlass::epilogue::warp::TileIteratorTensorOp; - using AccumulatorSharedStorage = - cutlass::gemm::threadblock::AccumulatorSharedStorage; - // We need to provide an operation for the epilogue. Let's create an - // operation that does nothing (ScaleType::Nothing), just converts - // from accum_t (float) -> scalar_t (can be half) - using OutputOpNoOp = cutlass::epilogue::thread::LinearCombination< - typename SmemIteratorD0::Element, // ElementOutput - FragmentIteratorAccumulator::Fragment::kElements, - accum_t, // ElementAccumulator - typename SmemIteratorD0::Element, // ElementCompute - cutlass::epilogue::thread::ScaleType::Nothing>; - using Epilogue = cutlass::epilogue::threadblock::EpilogueSmemAccumulator< - SmemIteratorD0, - FragmentIteratorAccumulator, - SmemIteratorD0, // ScaleBiasIterator - // - not used - OutputOpNoOp>; - - // Epilogue 2: with LSE (for backwards pass) - static int const kElementsPerAccess = 2; // TODO: Why 2? - using IteratorAccumulatorLSE = cutlass::transform::threadblock::VectorIterator< - cutlass::transform::threadblock::PredicatedVectorAccessIterator< - // Shape - cutlass::MatrixShape, - // WarpShape - cutlass::MatrixShape, - lse_scalar_t, - cutlass::layout::RowMajor, - kElementsPerAccess>>; - using EpilogueOpApplyLSE = cutlass::epilogue::thread::ApplyLogSumExp< - scalar_t, // ElementOutput_ - lse_scalar_t, // ElementLSE_ - accum_t, // ElementAccumulator_ - accum_t, // ElementCompute_ - 128 / cutlass::sizeof_bits::value - // FragmentIteratorAccumulator::Fragment::kElements - // InstructionShape::kM * InstructionShape::kN / 32 - >; - using EpilogueWithLSE = - cutlass::epilogue::threadblock::EpilogueSmemAccumulator; - - static void CUTLASS_DEVICE accumToSmem(AccumulatorSharedStorage& shared_storage, - FragmentC const& accum, - int lane_id, - cutlass::MatrixCoord const& tile_coords) - { - SmemIteratorD0 smem_iterator_attn(shared_storage.accum_ref(), lane_id); - smem_iterator_attn.add_tile_offset( - tile_coords * cutlass::MatrixCoord{SmemIteratorD0::TileIterations::kRow, - SmemIteratorD0::TileIterations::kColumn}); - Epilogue epilogue; - epilogue(OutputOpNoOp({}), smem_iterator_attn, accum); - } - - static void CUTLASS_DEVICE accumApplyLSEToSmem(AccumulatorSharedStorage& shared_storage, - FragmentC& accum, - lse_scalar_t const* lse, - int32_t lse_extents, - int thread_id, - int warp_id, - int lane_id, - cutlass::MatrixCoord const& tile_coords) - { - constexpr int32_t kAlignLSE = 32; - IteratorAccumulatorLSE iterator_lse( - lse, - {(int32_t)0, (int32_t)ceil_div(lse_extents, kAlignLSE) * kAlignLSE}, - thread_id, - warp_id, - cutlass::MatrixCoord{0, 0} // offset - ); - - SmemIteratorD0 smem_iterator_attn(shared_storage.accum_ref(), lane_id); - smem_iterator_attn.add_tile_offset( - tile_coords * cutlass::MatrixCoord{SmemIteratorD0::TileIterations::kRow, - SmemIteratorD0::TileIterations::kColumn}); - EpilogueWithLSE epilogue; - EpilogueOpApplyLSE minus_lse_exp({}); - epilogue(minus_lse_exp, - smem_iterator_attn, - accum, - // scale - unused - iterator_lse, - // bias - iterator_lse); - } -}; - -// Volta Specialization -// only supported for f16 -template -struct B2bGemm, - float, - cutlass::layout::RowMajor, - cutlass::gemm::GemmShape<16, 16, 4>, - cutlass::MatrixShape<1, 1>>, - Operator, - cutlass::half_t, - WarpShape_, - ThreadblockShape_> { - using IteratorC = cutlass::gemm::warp::MmaVoltaTensorOpAccumulatorTileIterator< - cutlass::MatrixShape<32, 32>, - float, - cutlass::layout::RowMajor, - cutlass::gemm::GemmShape<16, 16, 4>, - cutlass::MatrixShape<1, 1>>; - using scalar_t = cutlass::half_t; - using accum_t = IteratorC::Element; - using WarpShape = WarpShape_; - using ThreadblockShape = ThreadblockShape_; - using FragmentC = IteratorC::Fragment; - using lse_scalar_t = float; - - using SmemAccumulatorLayout = cutlass::layout::RowMajor; - using SmemIteratorD0 = - cutlass::epilogue::warp::TileIteratorVoltaTensorOp, - scalar_t, - SmemAccumulatorLayout>; - - // // Storage in shared-memory for Q.Kt - using AccumulatorSharedStorage = cutlass::gemm::threadblock::AccumulatorSharedStorage< - ThreadblockShape, - scalar_t, - cutlass::layout::RowMajorVoltaTensorOpMultiplicandCrosswise< - 16, - 32>, // typename SmemIteratorD0::TensorLayout, - cutlass::MatrixShape<0, 0> // Padding - >; - - using OutputLayout = cutlass::layout::RowMajorVoltaTensorOpMultiplicandCrosswise<16, 32>; - using TensorRef = cutlass::TensorRef; - using Policy = typename IteratorC::Policy; - using Element = accum_t; - // Those are MmaVoltaTensorOpAccumulatorTileIterator private fields - // Let's copy their values - static int const kElementsPerPartial = 4; - using EleShapePerPatial = - typename cutlass::platform::conditional::value, - cutlass::MatrixShape<2, 2>, - cutlass::MatrixShape<1, 4>>::type; - static int const kElementsPerMma = 8; - static int const kAccumulatorPatials = 2; - using QuadShapePerPatialMma = cutlass::MatrixShape<4, 4>; - - static void CUTLASS_DEVICE accumToSmem(AccumulatorSharedStorage& shared_storage, - FragmentC const& accum, - int lane_id, - cutlass::MatrixCoord const& tile_coords) - { - // ctor - from MmaVoltaTensorOpAccumulatorTileIterator - TensorRef ref_(shared_storage.accum_ref()); - int quad = (lane_id >> 2); - int lane_in_quad = (lane_id & 3); - int accum_m, accum_n; - - if (cutlass::platform::is_same::value) { - // (quad[2],quad[0])+lane_in_quad[0] - accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + (lane_in_quad & 1); - // (quad[1])+lane_in_quad[1] - accum_n = ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials + - (lane_in_quad & 2); - } else { - accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + lane_in_quad; // (quad[2],quad[0]) - accum_n = ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials; - } - cutlass::MatrixCoord lane_offset(accum_m, accum_n); - - // Tile offset - ref_.add_coord_offset(tile_coords * cutlass::MatrixCoord({IteratorC::Shape::kRow, - IteratorC::Shape::kColumn})); - - using AccessType = cutlass::Array; - - // store - from MmaVoltaTensorOpAccumulatorTileIterator - CUTLASS_PRAGMA_UNROLL - for (int tile_n = 0; tile_n < Policy::TileIterations::kColumn; ++tile_n) { - CUTLASS_PRAGMA_UNROLL - for (int tile_m = 0; tile_m < Policy::TileIterations::kRow; ++tile_m) { - CUTLASS_PRAGMA_UNROLL - for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { - CUTLASS_PRAGMA_UNROLL - for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { - int mma_accum_start = (((tile_n * Policy::TileIterations::kRow + tile_m) * - Policy::MmaIterations::kColumn + - mma_n) * - Policy::MmaIterations::kRow + - mma_m) * - kElementsPerMma; - - CUTLASS_PRAGMA_UNROLL - for (int p = 0; p < kAccumulatorPatials; ++p) { - CUTLASS_PRAGMA_UNROLL - for (int m = 0; m < EleShapePerPatial::kRow; ++m) { - int accum_m = tile_m * Policy::InterleavedTile::kRow + - mma_m * QuadShapePerPatialMma::kRow + m * 2; - int accum_n = tile_n * Policy::InterleavedTile::kColumn + - mma_n * QuadShapePerPatialMma::kColumn + - p * Policy::InterleavedTile::kColumn / 2; - int r = (accum_m + lane_offset.row()); - AccessType to_store; - CUTLASS_PRAGMA_UNROLL - for (int n = 0; n < EleShapePerPatial::kColumn; ++n) { - int idx = mma_accum_start + p * kElementsPerPartial + - m * EleShapePerPatial::kColumn + n; - int c = (accum_n + n + lane_offset.column()); - to_store[n] = scalar_t(accum[idx]); - } - int c = (accum_n + lane_offset.column()); - assert(r < 32); - assert(c < 32); - *reinterpret_cast(ref_.data() + ref_.offset({r, c})) = - to_store; - } - } - } - } - } - } - } - - static void CUTLASS_DEVICE accumApplyLSEToSmem(AccumulatorSharedStorage& shared_storage, - typename IteratorC::Fragment& accum, - lse_scalar_t const* lse, - int lse_extent, - int thread_id, - int warp_id, - int lane_id, - cutlass::MatrixCoord const& tile_coords) - { - // Non-optimized way to apply LSE to registers - // NOTE: accum is attn.T - // TODO: Optimize for each architecture - static constexpr int WarpSize = 32; - using AccumLambdaIterator = - typename DefaultMmaAccumLambdaIterator::Iterator; - auto lane_offset = AccumLambdaIterator::get_lane_offset(lane_id, warp_id, tile_coords); - - cutlass::Array lse_prefetched; - lse_prefetched.clear(); - int rowIdx = 0; - int colIdx = 0; - AccumLambdaIterator::iterateRows( - lane_offset, - [&](int accum_m) { - ++rowIdx; - colIdx = 0; - }, - [&](int accum_m, int accum_n, int idx) { - if (rowIdx == 1) { - lse_prefetched[colIdx] = accum_n < lse_extent - ? lse[accum_n] - : platform::numeric_limits::infinity(); - } - accum[idx] = expf(accum[idx] - lse_prefetched[colIdx]); - ++colIdx; - }, - [&](int accum_m) {}); - accumToSmem(shared_storage, accum, lane_id, tile_coords); - } -}; - -// Simt Specialization -// for f32 on Sm70-Sm75 and f16/f32 below - -template -struct B2bGemm, - cutlass::gemm::Operand::kC, - float, - cutlass::layout::RowMajor, - OperatorPolicy, - 1, - 1>, - Operator, - scalar_t, - WarpShape_, - ThreadblockShape_> { - using IteratorC = cutlass::gemm::warp::MmaSimtTileIterator, - cutlass::gemm::Operand::kC, - float, - cutlass::layout::RowMajor, - OperatorPolicy, - 1, - 1>; - using accum_t = typename IteratorC::Element; - using WarpShape = WarpShape_; - using ThreadblockShape = ThreadblockShape_; - using FragmentC = typename IteratorC::Fragment; - using lse_scalar_t = float; - - // Storage in shared-memory for Q.Kt - using AccumulatorSharedStorage = - cutlass::gemm::threadblock::AccumulatorSharedStorage // Padding - >; - - static void CUTLASS_DEVICE accumToSmem(AccumulatorSharedStorage& shared_storage, - FragmentC const& accum, - int lane_id, - cutlass::MatrixCoord const& tile_coords) - { - using Policy = typename IteratorC::Policy; - using Element = typename IteratorC::Element; - using Iterations = typename IteratorC::Iterations; - using Delta = typename IteratorC::Delta; - - auto ref_ = shared_storage.accum_ref(); - // ctor - MmaSimtTileIterator - // compute offset based on thread ID and lane layout - typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); - - MatrixCoord lane_offset = lane_layout.inverse(lane_id) * - MatrixCoord(Policy::LaneMmaShape::kM, Policy::LaneMmaShape::kN); - - ref_.add_coord_offset(lane_offset); - - // Tile offset - ref_.add_coord_offset(tile_coords * cutlass::MatrixCoord({IteratorC::Shape::kRow, - IteratorC::Shape::kColumn})); - - // store - MmaSimtTileIterator - CUTLASS_PRAGMA_UNROLL - for (int mma_n = 0; mma_n < Iterations::kColumn; ++mma_n) { - CUTLASS_PRAGMA_UNROLL - for (int n = 0; n < Policy::LaneMmaShape::kN; ++n) { - CUTLASS_PRAGMA_UNROLL - for (int mma_m = 0; mma_m < Iterations::kRow; ++mma_m) { - CUTLASS_PRAGMA_UNROLL - for (int m = 0; m < Policy::LaneMmaShape::kM; ++m) { - int r = Policy::LaneMmaShape::kM * (mma_m * Policy::WarpShape::kRow) + m; - int c = mma_n * Delta::kColumn + n; - int idx = n + Policy::LaneMmaShape::kN * - (mma_n + Iterations::kColumn * - (m + mma_m * Policy::LaneMmaShape::kM)); - ref_.at({r, c}) = scalar_t(accum[idx]); - } - } - } - } - } - - static void CUTLASS_DEVICE accumApplyLSEToSmem(AccumulatorSharedStorage& shared_storage, - typename IteratorC::Fragment& accum, - lse_scalar_t const* lse, - int lse_extent, - int thread_id, - int warp_id, - int lane_id, - cutlass::MatrixCoord const& tile_coords) - { - // Non-optimized way to apply LSE to registers - // NOTE: accum is attn.T - // TODO: Optimize for each architecture - static constexpr int WarpSize = 32; - using AccumLambdaIterator = - typename DefaultMmaAccumLambdaIterator::Iterator; - auto lane_offset = AccumLambdaIterator::get_lane_offset(lane_id, warp_id, tile_coords); - - cutlass::Array lse_prefetched; - lse_prefetched.clear(); - int rowIdx = 0; - int colIdx = 0; - AccumLambdaIterator::iterateRows( - lane_offset, - [&](int accum_m) { - ++rowIdx; - colIdx = 0; - }, - [&](int accum_m, int accum_n, int idx) { - if (rowIdx == 1) { - lse_prefetched[colIdx] = accum_n < lse_extent - ? lse[accum_n] - : platform::numeric_limits::infinity(); - } - accum[idx] = expf(accum[idx] - lse_prefetched[colIdx]); - ++colIdx; - }, - [&](int accum_m) {}); - accumToSmem(shared_storage, accum, lane_id, tile_coords); - } -}; - -} // namespace threadblock -} // namespace gemm -} // namespace cutlass - -///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm_kernel_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm_kernel_utils.h deleted file mode 100644 index c102234a4dfb040863bc5babad46e5bcfa1fd9b0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/gemm_kernel_utils.h +++ /dev/null @@ -1,253 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holdvr nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include "cutlass/arch/mma.h" - -template -struct CheckArch { - static constexpr bool isPreVolta = arch::kMinComputeCapability < 70; - static constexpr bool isPreAmpere = - arch::kMinComputeCapability < 80 && arch::kMinComputeCapability >= 70; - static constexpr bool isAmpere = arch::kMinComputeCapability >= 80; -#if defined(__CUDA_ARCH__) - static constexpr bool compiler_cc = arch::kMinComputeCapability * 10 <= __CUDA_ARCH__; -#else - static constexpr bool compiler_cc = true; -#endif - static constexpr bool value = (isPreVolta && std::is_same_v) || - (isPreAmpere && !std::is_same_v) || - isAmpere && compiler_cc; -}; - -#define DISPATCH_ARCHTAG(CC, func) \ - { \ - if constexpr (GPU_ARCH >= 80) { \ - if (CC >= 80) { \ - using ArchTag = cutlass::arch::Sm80; \ - func; \ - } else { \ - EVOFORMER_CHECK(false, "Compile flag error. Unexpected GPU"); \ - } \ - } else if constexpr (GPU_ARCH >= 75) { \ - if (CC >= 75) { \ - using ArchTag = cutlass::arch::Sm75; \ - func; \ - } else { \ - EVOFORMER_CHECK(false, "Compile flag error. Unexpected GPU"); \ - } \ - } else if constexpr (GPU_ARCH >= 70) { \ - if (CC >= 70) { \ - using ArchTag = cutlass::arch::Sm70; \ - func; \ - } else { \ - EVOFORMER_CHECK(false, "Compile flag error. Unexpected GPU"); \ - } \ - } else { \ - EVOFORMER_CHECK(false, "Only GPUs with Tensor Core are supported for now"); \ - } \ - } - -#define DISPATCH_TYPES(tensor, func) \ - { \ - if (tensor.scalar_type() == at::ScalarType::Half) { \ - using scalar_t = cutlass::half_t; \ - using torch_scalar_t = at::Half; \ - func; \ - } else if (tensor.scalar_type() == at::ScalarType::BFloat16) { \ - using scalar_t = cutlass::bfloat16_t; \ - using torch_scalar_t = at::BFloat16; \ - func; \ - } else { \ - EVOFORMER_CHECK(false, "Only fp16 and bf16 supported at the moment"); \ - } \ - } - -#define DISPATCH_BOOL(BOOL_V, BOOL_NAME, F) \ - { \ - if (BOOL_V) { \ - constexpr bool BOOL_NAME = true; \ - F(); \ - } else { \ - constexpr bool BOOL_NAME = false; \ - F(); \ - } \ - } - -#ifdef TORCH_CHECK -#define CHECK_ALIGNED_PTR(PTR, ALIGNMENT) \ - EVOFORMER_CHECK(uint64_t(PTR) % ALIGNMENT == 0, #PTR " is not correctly aligned") -#define EVOFORMER_CHECK TORCH_CHECK -#elif defined(__CUDACC_RTC__) -#define CHECK_ALIGNED_PTR(PTR, ALIGNMENT) \ - if (!(uint64_t(PTR) % ALIGNMENT == 0)) { return false; } -#define EVOFORMER_CHECK(COND, ERR) \ - if (!(COND)) { return false; } -#else -#include -#define CHECK_ALIGNED_PTR(PTR, ALIGNMENT) \ - if (!(uint64_t(PTR) % ALIGNMENT == 0)) { \ - std::cerr << #PTR " is not correctly aligned\n"; \ - return false; \ - } -#define EVOFORMER_CHECK(COND, ERR) \ - if (!(COND)) { \ - std::cerr << "[Evoformer Attention]" << "'" #COND "' failed: " << ERR << "\n"; \ - return false; \ - } -#endif - -namespace gemm_kernel_utils { - -template -constexpr CUTLASS_HOST_DEVICE integer ceil_div(integer n, integer m) -{ - return (n + m - 1) / m; -} - -template -constexpr CUTLASS_HOST_DEVICE integer align_up(integer n, integer m) -{ - return ((n + m - 1) / m) * m; -} - -//////////////////////////////////////////////////////////////////////////////// -// Determine the type of GEMM we do (TensorCores or not, Shapes ...) -// TODO: Maybe we could rely on Cutlass's DefaultGemm templates -//////////////////////////////////////////////////////////////////////////////// - -// Fallback to Simt (FMA on cuda cores) if not in a special case below -template -struct DefaultGemmType { - static constexpr int ThreadK = 8; - static constexpr int WarpK = 8; - static constexpr int kMinimumAlignment = 1; - using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>; - using OpClass = cutlass::arch::OpClassSimt; - using Operator = cutlass::arch::OpMultiplyAdd; -}; - -// Specialization for tensorcores with f32 -template -struct DefaultGemmType< - ArchTag, - float, - typename cutlass::platform::enable_if= 80>::type> { - static constexpr int ThreadK = 32; - static constexpr int WarpK = 32; - static constexpr int kMinimumAlignment = 4; - using OpClass = cutlass::arch::OpClassTensorOp; - using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; - using Operator = cutlass::arch::OpMultiplyAddFastF32; -}; - -// Specialization for tensorcores with f16/bf16 - Sm75+ -template -struct DefaultGemmType< - ArchTag, - scalar_t, - typename cutlass::platform::enable_if= 75 && - cutlass::sizeof_bits::value == 16>::type> { - static constexpr int ThreadK = 32; - static constexpr int WarpK = 32; - static constexpr int kMinimumAlignment = 4; - using OpClass = cutlass::arch::OpClassTensorOp; - using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; - using Operator = cutlass::arch::OpMultiplyAdd; -}; - -// Specialization for tensorcores with f16 - Volta -template <> -struct DefaultGemmType { - static constexpr int ThreadK = 32; - static constexpr int WarpK = 32; - static constexpr int kMinimumAlignment = 2; - using OpClass = cutlass::arch::OpClassTensorOp; - using InstructionShape = cutlass::gemm::GemmShape<8, 8, 4>; - using Operator = cutlass::arch::OpMultiplyAdd; -}; - -// Enables to do -// `auto x = kCondition ? fa(arg) : fb(arg)` -// when `fa` and `fb` have different types -template -struct call_conditional; - -template -struct call_conditional { - template - static CUTLASS_HOST_DEVICE auto apply(TA ta, TB tb, Arg arg) -> decltype(ta(arg)) - { - return ta(arg); - } -}; - -template -struct call_conditional { - template - static CUTLASS_HOST_DEVICE auto apply(TA ta, TB tb, Arg arg) -> decltype(tb(arg)) - { - return tb(arg); - } -}; - -//////////////////////////////////////////////////////////////////////////////// -// Mark a variable as warp-uniform - enables some compiler optimizations -// The cheapest way to do it is just to broadcast it from lane 0 -//////////////////////////////////////////////////////////////////////////////// - -CUTLASS_DEVICE int32_t warp_uniform(int32_t value) -{ - return (int32_t)__shfl_sync(0xffffffff, (unsigned)value, 0); -} - -template -CUTLASS_DEVICE T* warp_uniform(T* ptr) -{ - struct { - union { - T* ptr; - uint32_t asInt[2]; - }; - } p; - p.ptr = ptr; - p.asInt[0] = warp_uniform(p.asInt[0]); - p.asInt[1] = warp_uniform(p.asInt[1]); - return p.ptr; -} -} // namespace gemm_kernel_utils diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/epilogue_predicated_tile_iterator.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/epilogue_predicated_tile_iterator.h deleted file mode 100644 index 667f1982d30d295498b952f0a681af29796a1521..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/epilogue_predicated_tile_iterator.h +++ /dev/null @@ -1,691 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights - *reserved. SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, - *this list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE - *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE - *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR - *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF - *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS - *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN - *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) - *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE - *POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/*! \file - \brief Epilogue iterator that supports prefetching - - Mostly copied from "cutlass/epilogue/threadblock/predicated_tile_iterator.h" -*/ - -#pragma once - -#include "cutlass/arch/arch.h" -#include "cutlass/arch/memory.h" -#include "cutlass/array.h" -#include "cutlass/cutlass.h" -#include "cutlass/epilogue/threadblock/output_tile_thread_map.h" -#include "cutlass/epilogue/threadblock/predicated_tile_iterator_params.h" -#include "cutlass/layout/matrix.h" -#include "cutlass/layout/tensor.h" -#include "cutlass/matrix_shape.h" -#include "cutlass/numeric_types.h" -#include "cutlass/tensor_ref.h" -#include "cutlass/transform/pitch_linear_thread_map.h" - -//////////////////////////////////////////////////////////////////////////////// - -namespace cutlass { - -//////////////////////////////////////////////////////////////////////////////// - -namespace epilogue { -namespace threadblock { - -//////////////////////////////////////////////////////////////////////////////// - -/// Tile iterator used to load and store output tile from global memory in -/// epilogue. -/// -/// Satisfies: ReadableTileIterator | PredicatedTileIterator | -/// ForwardTileIterator -/// -template -class PredicatedTileIteratorPrefetch { -public: - using ThreadMap = ThreadMap_; - using Shape = typename ThreadMap::Shape; - - using Element = Element_; - - using Layout = layout::RowMajor; - using TensorRef = TensorRef; - using ConstTensorRef = typename TensorRef::ConstTensorRef; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - using TensorCoord = MatrixCoord; - - static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; - static int const kThreads = ThreadMap::kThreads; - static int const kIterations = ThreadMap::Count::kTile; - - static_assert(ThreadMap::Iterations::kRow > 0, "ThreadMap::Iterations::kRow must be > 0"); - static_assert(ThreadMap::Iterations::kGroup > 0, "ThreadMap::Iterations::kGroup must be > 0"); - static_assert(ThreadMap::Iterations::kCluster > 0, - "ThreadMap::Iterations::kCluster must be > 0"); - static_assert(ThreadMap::Iterations::kColumn > 0, "ThreadMap::Iterations::kColumn must be > 0"); - - /// Fragment object - using Fragment = Array; - - /// Memory access size - using AccessType = AlignedArray; - - // - // Parameters struct - // - - /// Uses a non-template class - struct Params : PredicatedTileIteratorParams { - using Base = PredicatedTileIteratorParams; - - CUTLASS_HOST_DEVICE - Params() {} - - CUTLASS_HOST_DEVICE - Params(Layout const& layout) - : PredicatedTileIteratorParams( - layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess, - make_OutputTileThreadMapDesc()) - { - } - - CUTLASS_HOST_DEVICE - Params(Base const& base) : Base(base) {} - }; - - /// Mask object - struct Mask { - static int const kCount = ThreadMap::Iterations::kColumn; - - /// Predicate state - bool predicates[kCount]; - - // - // Mask - // - CUTLASS_HOST_DEVICE - Mask() { enable(); } - - ///< Efficiently disables all accesses guarded by mask - CUTLASS_HOST_DEVICE void clear() - { - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kCount; ++i) { predicates[i] = false; } - } - - ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask - CUTLASS_DEVICE void enable() - { - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kCount; ++i) { predicates[i] = true; } - } - }; - -private: - // - // Data members - // - - /// Parameters structure containing reference and precomputed state. - PredicatedTileIteratorParams params_; - - /// Byte-level pointer - uint8_t* byte_pointer_; - - /// Array of boolean values to contain steady-state predicates - Mask mask_; - - /// Extent of the matrix tile in rows - Index extent_row_; - - /// Extent of the matrix tile in rows - Index extent_column_; - - /// A thread's starting row position (assuming steady-state predicates have - /// been computed) - Index thread_start_row_; - - /// A thread's starting column - Index thread_start_column_; - - /// Internal state counter - int state_[3]; - - /// Scatter indices - int const* indices_; - - // - // Static asserts about internal strides - // - - static_assert(sizeof(extent_row_) == 4, "Expected 32b extents"); - static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); - static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides"); - -private: - // - // Methods - // - -public: - // - // Methods - // - - /// Constructor - CUTLASS_DEVICE - PredicatedTileIteratorPrefetch(PredicatedTileIteratorParams const& params, - Element* pointer, - TensorCoord extent, - int thread_idx, - TensorCoord threadblock_offset = TensorCoord(), - int const* indices = nullptr) - : params_(params), indices_(indices) - { - TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; - - extent_row_ = extent.row(); - extent_column_ = extent.column(); - - thread_start_row_ = thread_offset.row(); - thread_start_column_ = thread_offset.column(); - - // Initialize predicates - CUTLASS_PRAGMA_UNROLL - for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { - mask_.predicates[c] = - ((thread_offset.column() + ThreadMap::Delta::kColumn * c) < extent.column()); - } - - // Null pointer performs no accesses - if (!pointer) { mask_.clear(); } - - if (ScatterD && !indices) { mask_.clear(); } - - // Initialize pointer - byte_pointer_ = reinterpret_cast(pointer) + - LongIndex(thread_offset.row()) * LongIndex(params_.stride) + - LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess; - - if (ScatterD) { - byte_pointer_ = - reinterpret_cast(pointer) + - LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess; - } - - // Initialize internal state counter - state_[0] = state_[1] = state_[2] = 0; - } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - byte_pointer_ += pointer_offset * sizeof_bits::value / 8; - } - - CUTLASS_DEVICE - void prefetch_all() - { - CUTLASS_PRAGMA_UNROLL - for (int iter = 0; iter < kIterations; ++iter) { - prefetch(); - ++(*this); - } - } - - CUTLASS_DEVICE - void prefetch() - { - uint8_t* byte_pointer = byte_pointer_; - - CUTLASS_PRAGMA_UNROLL - for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { - CUTLASS_PRAGMA_UNROLL - for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { - CUTLASS_PRAGMA_UNROLL - for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { - int row_offset = row * ThreadMap::Delta::kRow + - group * ThreadMap::Delta::kGroup + - cluster * ThreadMap::Delta::kCluster; - - AccessType* memory_pointer = reinterpret_cast(byte_pointer); - - CUTLASS_PRAGMA_UNROLL - for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { - // on windows using unsigned long here gives the error - // error: asm operand type size(4) does not match - // type/size implied by constraint 'l' - uint64_t addr = - (uint64_t)((void*)&memory_pointer[column * ThreadMap::Delta::kColumn / - kElementsPerAccess]); - asm volatile("prefetch.global.L1 [ %1 ];" : "=l"(addr) : "l"(addr)); - } - - if (row + 1 < ThreadMap::Iterations::kRow) { - if (!ScatterD) { byte_pointer += params_.increment_row; } - } - } - - if (group + 1 < ThreadMap::Iterations::kGroup) { - byte_pointer += params_.increment_group; - } - } - - if (cluster + 1 < ThreadMap::Iterations::kCluster) { - byte_pointer += params_.increment_cluster; - } - } - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load_with_byte_offset(Fragment& frag, int64_t byte_offset) const - { - uint8_t* byte_pointer = byte_pointer_; - AccessType* frag_ptr = reinterpret_cast(&frag); - - CUTLASS_PRAGMA_UNROLL - for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { - CUTLASS_PRAGMA_UNROLL - for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { - CUTLASS_PRAGMA_UNROLL - for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { - int frag_row_idx = - (row + ThreadMap::Iterations::kRow * - (group + ThreadMap::Iterations::kGroup * cluster)); - - int row_offset = row * ThreadMap::Delta::kRow + - group * ThreadMap::Delta::kGroup + - cluster * ThreadMap::Delta::kCluster; - - bool row_guard = ((row_offset + thread_start_row_) < extent_row_); - - AccessType* memory_pointer = - reinterpret_cast(byte_pointer + byte_offset); - - if (ScatterD && row_guard) { - assert(indices_); - - memory_pointer = reinterpret_cast( - byte_pointer + byte_offset + - LongIndex(indices_[row_offset + thread_start_row_]) * - LongIndex(params_.stride)); - } - - CUTLASS_PRAGMA_UNROLL - for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { - bool guard = row_guard && mask_.predicates[column]; - - cutlass::arch::global_load( - frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], - (void*)&memory_pointer[column * ThreadMap::Delta::kColumn / - kElementsPerAccess], - guard); - } - - if (row + 1 < ThreadMap::Iterations::kRow) { - if (!ScatterD) { byte_pointer += params_.increment_row; } - } - } - - if (group + 1 < ThreadMap::Iterations::kGroup) { - byte_pointer += params_.increment_group; - } - } - - if (cluster + 1 < ThreadMap::Iterations::kCluster) { - byte_pointer += params_.increment_cluster; - } - } - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load(Fragment& frag) const { load_with_byte_offset(frag, 0); } - - /// Stores a fragment to memory - CUTLASS_DEVICE - void store_with_byte_offset(Fragment const& frag, int64_t byte_offset) const - { - uint8_t* byte_pointer = byte_pointer_; - AccessType const* frag_ptr = reinterpret_cast(&frag); - - CUTLASS_PRAGMA_UNROLL - for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { - CUTLASS_PRAGMA_UNROLL - for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { - CUTLASS_PRAGMA_UNROLL - for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { - int frag_row_idx = - (row + ThreadMap::Iterations::kRow * - (group + ThreadMap::Iterations::kGroup * cluster)); - - int row_offset = row * ThreadMap::Delta::kRow + - group * ThreadMap::Delta::kGroup + - cluster * ThreadMap::Delta::kCluster; - - bool row_guard = ((row_offset + thread_start_row_) < extent_row_); - - AccessType* memory_pointer = - reinterpret_cast(byte_pointer + byte_offset); - - if (ScatterD && row_guard) { - assert(indices_); - - memory_pointer = reinterpret_cast( - byte_pointer + byte_offset + - LongIndex(indices_[row_offset + thread_start_row_]) * - LongIndex(params_.stride)); - } - - CUTLASS_PRAGMA_UNROLL - for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { - bool guard = row_guard && mask_.predicates[column]; - - if (UseCUDAStore) { - if (guard) { - memory_pointer[column * ThreadMap::Delta::kColumn / - kElementsPerAccess] = - frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + - column]; - } - } else { - cutlass::arch::global_store( - frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], - (void*)&memory_pointer[column * ThreadMap::Delta::kColumn / - kElementsPerAccess], - guard); - } - } - - if (row + 1 < ThreadMap::Iterations::kRow) { - if (!ScatterD) { byte_pointer += params_.increment_row; } - } - } - - if (group + 1 < ThreadMap::Iterations::kGroup) { - byte_pointer += params_.increment_group; - } - } - - if (cluster + 1 < ThreadMap::Iterations::kCluster) { - byte_pointer += params_.increment_cluster; - } - } - } - - /// Stores a fragment to memory - CUTLASS_DEVICE - void store(Fragment const& frag) const { store_with_byte_offset(frag, 0); } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void downsample_load_with_byte_offset(Fragment& frag, - int64_t byte_offset, - int convolution_P, - int convolution_Q, - int add_P, - int add_Q, - int problem_N) const - { - uint8_t* byte_pointer = byte_pointer_; - AccessType* frag_ptr = reinterpret_cast(&frag); - - CUTLASS_PRAGMA_UNROLL - for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { - CUTLASS_PRAGMA_UNROLL - for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { - CUTLASS_PRAGMA_UNROLL - for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { - int frag_row_idx = - (row + ThreadMap::Iterations::kRow * - (group + ThreadMap::Iterations::kGroup * cluster)); - - int row_offset = row * ThreadMap::Delta::kRow + - group * ThreadMap::Delta::kGroup + - cluster * ThreadMap::Delta::kCluster; - - bool row_guard = ((row_offset + thread_start_row_) < extent_row_); - - int output_row = row_offset + thread_start_row_; - int output_N = output_row / (convolution_P * convolution_Q); - int output_PQ = output_row % (convolution_P * convolution_Q); - int output_P = output_PQ / convolution_Q; - int output_Q = output_PQ % convolution_Q; - - int input_row = output_N * 2 * convolution_P * 2 * convolution_Q + - (2 * output_P + add_P) * 2 * convolution_Q + 2 * output_Q + - add_Q; - - int64_t byte_offset = (input_row - output_row) * problem_N * sizeof(float); - - AccessType* memory_pointer = - reinterpret_cast(byte_pointer + byte_offset); - - CUTLASS_PRAGMA_UNROLL - for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { - bool guard = row_guard && mask_.predicates[column]; - - cutlass::arch::global_load( - frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], - (void*)&memory_pointer[column * ThreadMap::Delta::kColumn / - kElementsPerAccess], - guard); - } - - if (row + 1 < ThreadMap::Iterations::kRow) { - byte_pointer += params_.increment_row; - } - } - - if (group + 1 < ThreadMap::Iterations::kGroup) { - byte_pointer += params_.increment_group; - } - } - - if (cluster + 1 < ThreadMap::Iterations::kCluster) { - byte_pointer += params_.increment_cluster; - } - } - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void upsample_load_with_byte_offset(Fragment& frag, - int64_t byte_offset, - int convolution_P, - int convolution_Q, - int add_P, - int add_Q, - int problem_N) const - { - uint8_t* byte_pointer = byte_pointer_; - AccessType* frag_ptr = reinterpret_cast(&frag); - - CUTLASS_PRAGMA_UNROLL - for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { - CUTLASS_PRAGMA_UNROLL - for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { - CUTLASS_PRAGMA_UNROLL - for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { - int frag_row_idx = - (row + ThreadMap::Iterations::kRow * - (group + ThreadMap::Iterations::kGroup * cluster)); - - int row_offset = row * ThreadMap::Delta::kRow + - group * ThreadMap::Delta::kGroup + - cluster * ThreadMap::Delta::kCluster; - - bool row_guard = ((row_offset + thread_start_row_) < extent_row_); - - int output_row = row_offset + thread_start_row_; - int output_N = output_row / (convolution_P * convolution_Q); - int output_PQ = output_row % (convolution_P * convolution_Q); - int output_P = output_PQ / convolution_Q; - int output_Q = output_PQ % convolution_Q; - int row_add_P = add_P; - int row_add_Q = add_Q; - if (output_P > convolution_P - 2) row_add_P = 0; - if (output_Q > convolution_Q - 2) row_add_Q = 0; - - int input_row = output_N * (convolution_P / 2) * (convolution_Q / 2) + - ((output_P + row_add_P) / 2) * (convolution_Q / 2) + - (output_Q + row_add_Q) / 2; - - int64_t byte_offset = (input_row - output_row) * problem_N * sizeof(float); - - AccessType* memory_pointer = - reinterpret_cast(byte_pointer + byte_offset); - - CUTLASS_PRAGMA_UNROLL - for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { - bool guard = row_guard && mask_.predicates[column]; - - cutlass::arch::global_load( - frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], - (void*)&memory_pointer[column * ThreadMap::Delta::kColumn / - kElementsPerAccess], - guard); - } - - if (row + 1 < ThreadMap::Iterations::kRow) { - byte_pointer += params_.increment_row; - } - } - - if (group + 1 < ThreadMap::Iterations::kGroup) { - byte_pointer += params_.increment_group; - } - } - - if (cluster + 1 < ThreadMap::Iterations::kCluster) { - byte_pointer += params_.increment_cluster; - } - } - } - - CUTLASS_DEVICE - MatrixCoord thread_start() const - { - return MatrixCoord(thread_start_row_, thread_start_column_); - } - - /// Need to get the thread start row from the tile iterator - CUTLASS_DEVICE - int32_t thread_start_row() const { return thread_start_row_; } - - /// Need to get the thread start row from the tile iterator - CUTLASS_DEVICE - int32_t thread_start_column() const { return thread_start_column_; } - - /// Extent of the matrix in rows - CUTLASS_DEVICE - Index extent_row() const { return extent_row_; } - - /// Extent of the matrix in columns - CUTLASS_DEVICE - Index extent_column() const { return extent_column_; } - - /// Advances to the next position to load or store - CUTLASS_HOST_DEVICE - PredicatedTileIteratorPrefetch& operator++() - { - ++state_[0]; - - if (!ScatterD) { byte_pointer_ += params_.advance_row; } - - thread_start_row_ += ThreadMap::Shape::kRow; - - if (state_[0] == ThreadMap::Count::kRow) { - state_[0] = 0; - ++state_[1]; - byte_pointer_ += params_.advance_group; - - thread_start_row_ += - (ThreadMap::Shape::kGroup - 1) * ThreadMap::Shape::kRow * ThreadMap::Count::kRow; - - if (state_[1] == ThreadMap::Count::kGroup) { - state_[1] = 0; - ++state_[2]; - byte_pointer_ += params_.advance_cluster; - - thread_start_row_ += ThreadMap::Count::kGroup * ThreadMap::Shape::kGroup * - ThreadMap::Count::kRow * ThreadMap::Shape::kRow; - - if (state_[2] == ThreadMap::Count::kCluster) { - state_[2] = 0; - byte_pointer_ += params_.advance_tile; - } - } - } - - return *this; - } - - ///< Efficiently disables all accesses guarded by mask - CUTLASS_DEVICE void clear_mask() { mask_.clear(); } - - ///< Efficiently enables all accesses guarded by mask - CUTLASS_DEVICE void enable_mask() { mask_.enable(); } - - ///< Sets the mask - CUTLASS_DEVICE void get_mask(Mask& mask) const { mask = mask_; } - - ///< Sets the mask - CUTLASS_DEVICE void set_mask(Mask const& mask) { mask_ = mask; } -}; - -template -struct MakePrefetchableIterator { - using Iterator = PredicatedTileIteratorPrefetch; -}; - -/////////////////////////////////////////////////////////////////////////////// - -} // namespace threadblock -} // namespace epilogue -} // namespace cutlass - -//////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/make_residual_last.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/make_residual_last.h deleted file mode 100644 index ff0e324c3a6c0899d8439d92658515fb255f2dd4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/make_residual_last.h +++ /dev/null @@ -1,91 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holdvr nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include "predicated_tile_access_iterator_residual_last.h" -#include "predicated_tile_iterator_residual_last.h" - -namespace cutlass { -namespace transform { -namespace threadblock { - -template -struct MakeIteratorResidualLast; - -template -struct MakeIteratorResidualLast< - PredicatedTileIterator> { - using Iterator = PredicatedTileIteratorResidualLast; -}; - -template -struct MakeIteratorResidualLast> { - using Iterator = PredicatedTileAccessIteratorResidualLast; -}; -} // namespace threadblock -} // namespace transform -} // namespace cutlass diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_access_iterator_residual_last.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_access_iterator_residual_last.h deleted file mode 100644 index 7f6a2430845a16aac5891257c5b0600c5226e9a9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_access_iterator_residual_last.h +++ /dev/null @@ -1,1964 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights - *reserved. SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, - *this list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE - *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE - *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR - *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF - *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS - *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN - *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) - *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE - *POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/*! \file - \brief Templates calculating the address and predicates to the load of tiles - from pitch-linear rank=2 tensors. - - This iterator uses masks to guard out-of-bounds accesses. The first tile - this iterator visits maybe partial, then the remaining tiles are complete. - So, we only need to compute the predicates twice, once before the first tile - and once for the remaining full tiles which can share the same predicates. - - A precomputed "Params" object minimizes the amount of state that must be - stored in registers, and integer addition is used to advance the pointer - through memory. -*/ - -#pragma once - -#include "cutlass/array.h" -#include "cutlass/coord.h" -#include "cutlass/cutlass.h" -#include "cutlass/layout/matrix.h" -#include "cutlass/layout/pitch_linear.h" -#include "cutlass/matrix_shape.h" -#include "cutlass/predicate_vector.h" -#include "cutlass/tensor_ref.h" -#include "cutlass/tensor_view.h" -#include "cutlass/transform/threadblock/predicated_tile_access_iterator_params.h" - -//////////////////////////////////////////////////////////////////////////////// - -//////////////////////////////////////////////////////////////////////////////// - -namespace cutlass { -namespace transform { -namespace threadblock { - -//////////////////////////////////////////////////////////////////////////////// - -/// PredicatedTileAccessIteratorResidualLast -/// -template -class PredicatedTileAccessIteratorResidualLast; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileAccessIteratorResidualLast for pitch-linear -/// data. -/// -template -class PredicatedTileAccessIteratorResidualLast { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may along advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - using Layout = layout::PitchLinear; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - using AccessType = AccessType_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - using UnderlyingPredicates = PredicatedTileAccessIteratorPredicates; - - static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; - - static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), - "Vectors implied by the thread map must be divisible by the access type."); - - using Mask = typename UnderlyingPredicates::Mask; - - /// Uses a non-template class - struct Params : PredicatedTileAccessIteratorParams { - using Base = PredicatedTileAccessIteratorParams; - - // Default ctor - CUTLASS_HOST_DEVICE - Params() {} - - /// Construct the Params object given a pitch-linear tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) - : Base(layout.stride(0), - MakePredicatedTileAccessIteratorDesc()()) - { - } - - CUTLASS_HOST_DEVICE - Params(Base const& base) : Base(base) {} - }; - -private: - /// Internal pointer type permits fast address arithmetic - using BytePointer = char*; - -private: - // - // Data members - // - - UnderlyingPredicates the_predicates; - Mask residual_tile_mask; - - /// Parameters object with precomputed internal state - Params params_; - - /// Internal pointer to first access of tile - BytePointer pointer_; - - /// Below is used when Gather is turned on. We need to record strided_offset - /// and contiguous_offset separated to compute the offset by using - /// - /// offset = contiguous_offset + indices[strided_offset] - /// - - /// Gather indices - int const* indices_; - - Index gather_offset_strided; - -private: - /// Computes predicates based on internally tracked per-thread offset. - CUTLASS_DEVICE - void compute_predicates_( - /// Extent of the matrix window - TensorCoord extent, - /// optionally, simplify predicate calculation during 'steady state' phase - bool is_steady_state = false) - { - the_predicates.compute_predicates_(extent, is_steady_state); - } - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - /// Precomputed parameters object - Params const& params, - /// Pointer to start of tensor - Pointer pointer, - /// Extent of tensor - TensorCoord extent, - /// ID of each participating thread - int thread_id, - /// Initial offset of threadblock - TensorCoord const& threadblock_offset, - /// Gather indices - int const* indices = nullptr) - : params_(params), - pointer_(reinterpret_cast(const_cast(pointer))), - the_predicates(extent), - indices_(indices) - { - the_predicates.set_predicates(thread_id, threadblock_offset); - the_predicates.get_mask(residual_tile_mask); - - // Working around a weird compiler bug happening on P100 for the backward. - // I've seen together: the_predicates.predicates_[0] = 14 (instead of 15) - // residual_tile_mask[0] = 15 (correct) - // - // Adding prints when the value is calculated (in `compute_predicates_`) - // sometimes removes the bug. The consequence is that we skip some - // element of a tensor, leading to wrong results - // Setting `compute_predicates_`'s second argument (`is_steady_state`) to - // true also seems to get rid of the bug - at the cost of twice as many - // comparisons. -#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 700) - constexpr bool kWorkAroundCompilerBug = false; -#else - constexpr bool kWorkAroundCompilerBug = true; -#endif - the_predicates.compute_predicates_(extent, true && !kWorkAroundCompilerBug); - - // update internal pointers - Layout layout(params_.stride_); - - if (!Gather) { - add_pointer_offset(layout(the_predicates.thread_offset_)); - } else { - gather_offset_strided = the_predicates.thread_offset_.strided(); - add_pointer_offset(layout(make_Coord(the_predicates.thread_offset_.contiguous(), 0))); - } - } - - /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - /// Precomputed parameters object - Params const& params, - /// Pointer to start of tensor - Pointer pointer, - /// Extent of tensor - TensorCoord extent, - ///< ID of each participating thread - int thread_id) - : PredicatedTileAccessIteratorResidualLast(params, - pointer, - extent, - thread_id, - make_Coord(0, 0)) - { - } - - /// Overrides the internal iteration index - CUTLASS_HOST_DEVICE - void set_iteration_index(int index) { the_predicates.set_iteration_index(index); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool is_residual_tile) - { - if (is_residual_tile) { the_predicates.set_mask(residual_tile_mask); } - } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - pointer_ += sizeof_bits::value * pointer_offset / 8; - } - - /// Advances an iterator along logical dimensions of matrix in units of whole - /// tiles - CUTLASS_DEVICE - void add_tile_offset(TensorCoord const& tile_offset) - { - if (!Gather) { - if (kAdvanceRank) { - pointer_ += params_.inc_advance_ * LongIndex(tile_offset.strided()); - pointer_ += Shape::kContiguous * tile_offset.contiguous(); - } else { - pointer_ += params_.inc_advance_ * LongIndex(tile_offset.contiguous()); - pointer_ += Shape::kStrided * tile_offset.strided(); - } - } else { - add_pointer_offset(Shape::kContiguous * tile_offset.contiguous()); - gather_offset_strided += Shape::kStrided * tile_offset.strided(); - } - } - - /// Returns a pointer - CUTLASS_HOST_DEVICE - AccessType* get() const - { - if (Gather) { - assert(indices_); - - if (!valid()) { return nullptr; } - - LongIndex contiguous_offset = - the_predicates.iteration_contiguous_ * - (ThreadMap::Delta::kContiguous * sizeof_bits::value / 8) + - the_predicates.iteration_vector_; - int strided_index = gather_offset_strided + - the_predicates.iteration_strided_ * ThreadMap::Delta::kStrided; - - LongIndex strided_offset = indices_[strided_index] * LongIndex(params_.stride_) * - sizeof_bits::value / 8; - - return reinterpret_cast(pointer_ + contiguous_offset + strided_offset); - } - - return reinterpret_cast( - pointer_ + the_predicates.iteration_contiguous_ * - (ThreadMap::Delta::kContiguous * sizeof_bits::value) / - 8) + - the_predicates.iteration_vector_; - } - - /// Increment and return an instance to self. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast& operator++() - { - the_predicates.operator++(); - - ++the_predicates.iteration_vector_; - if (the_predicates.iteration_vector_ < kAccessesPerVector) { return *this; } - - the_predicates.iteration_vector_ = 0; - ++the_predicates.iteration_contiguous_; - - if (the_predicates.iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { - return *this; - } - - // Enter here only if (iteration_contiguous_ == - // ThreadMap::Iteration::kContiguous) - the_predicates.iteration_contiguous_ = 0; - ++the_predicates.iteration_strided_; - - if (the_predicates.iteration_strided_ < ThreadMap::Iterations::kStrided) { - if (!Gather) { pointer_ += params_.inc_strided_; } - - return *this; - } - - // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) - // which means we enter the next tile. - the_predicates.iteration_strided_ = 0; - - if (!Gather) { - // advance to next tile - pointer_ += params_.inc_next_; - - // now return to start tile - if the iterator is subsequently advanced, - // this subtraction as well as the subsequent integer addition are both - // elided by the compiler. - pointer_ -= params_.inc_advance_; - } - - return *this; - } - - /// Increment and return an instance to self. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast operator++(int) - { - PredicatedTileAccessIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { the_predicates.clear_mask(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { the_predicates.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { the_predicates.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { the_predicates.get_mask(mask); } - - /// Returns whether access is valid or not - CUTLASS_HOST_DEVICE - bool valid() const { return the_predicates.valid(); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileAccessIteratorResidualLast for column-major -/// data. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// -template -class PredicatedTileAccessIteratorResidualLast { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may along advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - using Layout = layout::ColumnMajor; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - using AccessType = AccessType_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< - layout::PitchLinearShape, - Element, - layout::PitchLinear, - (kAdvanceRank == 0 ? 0 : 1), - ThreadMap, - AccessType, - Gather>; - - /// Predicate vector stores mask to guard accesses - using Mask = typename UnderlyingIterator::Mask; - - static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; - - /// Parameters object is precomputed state and is host-constructible - class Params { - private: - friend PredicatedTileAccessIteratorResidualLast; - - /// Parameters object - typename UnderlyingIterator::Params params_; - - public: - /// Default ctor - CUTLASS_HOST_DEVICE - Params() {} - - /// Construct the Params object given a pitch-linear tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))){}; - - /// Construct the Params object given a pitch-linear tensor's layout - CUTLASS_HOST_DEVICE - Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} - }; - -private: - // - // Data members - // - - /// Underlying pitch-linear tile iterator - UnderlyingIterator iterator_; - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - ///< Precomputed parameters object - Params const& params, - ///< Pointer to start of tensor - Pointer pointer, - ///< Extent of tensor - TensorCoord extent, - ///< ID of each participating thread - int thread_id, - ///< Initial offset of threadblock - TensorCoord const& threadblock_offset, - int const* indices = nullptr ///< gather/scatter indices, note no support for - ///< gather/scatter at this specialization - ) - : iterator_(params.params_, - pointer, - layout::PitchLinearCoord(extent.row(), extent.column()), - thread_id, - layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column()), - indices) - { - } - - /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileAccessIteratorResidualLast(params, - pointer, - extent, - thread_id, - make_Coord(0, 0)) - { - } - - /// Overrides the internal iteration index - CUTLASS_HOST_DEVICE - void set_iteration_index(int index) { iterator_.set_iteration_index(index); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - iterator_.add_pointer_offset(pointer_offset); - } - - /// Advances an iterator along logical dimensions of matrix in units of whole - /// tiles - CUTLASS_HOST_DEVICE - void add_tile_offset(TensorCoord const& tile_offset) - { - iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); - } - - /// Returns a pointer - CUTLASS_HOST_DEVICE - AccessType* get() const { return reinterpret_cast(iterator_.get()); } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast& operator++() - { - ++iterator_; - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast operator++(int) - { - PredicatedTileAccessIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { iterator_.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { iterator_.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { iterator_.get_mask(mask); } - - /// Returns whether access is valid or not - CUTLASS_HOST_DEVICE - bool valid() { return iterator_.valid(); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileAccessIteratorResidualLast for row-major -/// data. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// -template -class PredicatedTileAccessIteratorResidualLast { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may along advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - using Layout = layout::RowMajor; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - using AccessType = AccessType_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< - layout::PitchLinearShape, - Element, - layout::PitchLinear, - (kAdvanceRank == 0 ? 1 : 0), - ThreadMap, - AccessType, - Gather>; - - static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; - - /// Predicate vector stores mask to guard accesses - using Mask = typename UnderlyingIterator::Mask; - - /// Parameters object is precomputed state and is host-constructible - class Params { - private: - friend PredicatedTileAccessIteratorResidualLast; - - /// Parameters object - typename UnderlyingIterator::Params params_; - - public: - /// Default ctor - CUTLASS_HOST_DEVICE - Params() {} - - /// Construct the Params object given a pitch-linear tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))){}; - - /// Construct the Params object given a pitch-linear tensor's layout - CUTLASS_HOST_DEVICE - Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} - }; - -private: - // - // Data members - // - - /// Underlying pitch-linear tile iterator - UnderlyingIterator iterator_; - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - ///< Precomputed parameters object - Params const& params, - ///< Pointer to start of tensor - Pointer pointer, - ///< Extent of tensor - TensorCoord extent, - ///< ID of each participating thread - int thread_id, - ///< Initial offset of threadblock - TensorCoord const& threadblock_offset, - /// Gather indices - int const* indices = nullptr) - : iterator_(params.params_, - pointer, - layout::PitchLinearCoord(extent.column(), extent.row()), - thread_id, - layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row()), - indices) - { - } - - /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileAccessIteratorResidualLast(params, - pointer, - extent, - thread_id, - make_Coord(0, 0)) - { - } - - /// Overrides the internal iteration index - CUTLASS_HOST_DEVICE - void set_iteration_index(int index) { iterator_.set_iteration_index(index); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - iterator_.add_pointer_offset(pointer_offset); - } - - /// Advances an iterator along logical dimensions of matrix in units of whole - /// tiles - CUTLASS_HOST_DEVICE - void add_tile_offset(TensorCoord const& tile_offset) - { - iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); - } - - /// Returns a pointer - CUTLASS_HOST_DEVICE - AccessType* get() const { return reinterpret_cast(iterator_.get()); } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast& operator++() - { - ++iterator_; - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast operator++(int) - { - PredicatedTileAccessIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { iterator_.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { iterator_.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { iterator_.get_mask(mask); } - - /// Returns whether access is valid or not - CUTLASS_HOST_DEVICE - bool valid() { return iterator_.valid(); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileAccessIteratorResidualLast for affine rank 2 -/// data. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// -template -class PredicatedTileAccessIteratorResidualLast, - AdvanceRank, - ThreadMap_, - AccessType_, - false> { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may along advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - using Layout = layout::AffineRankN<2>; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - using AccessType = AccessType_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - using UnderlyingPredicates = PredicatedTileAccessIteratorPredicates; - - static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; - - static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), - "Vectors implied by the thread map must be divisible by the access type."); - - /// Predicate vector stores mask to guard accesses - using Mask = typename UnderlyingPredicates::Mask; - - /// Parameters object is precomputed state and is host-constructible - class Params { - public: - friend PredicatedTileAccessIteratorResidualLast; - - private: - /// stride of pitch-linear layout (units of Element) - Coord stride_; - /// amount (in byte) to increment pointer to move to next access along - /// contiguous dimension - LongIndex inc_contiguous_; - /// amount (in byte) to increment pointer from first access of current - /// contiguous dimension to first access of next one. - LongIndex inc_strided_; - /// amount (in byte) to increment pointer from last access of current - /// contiguous dimension to first access of next one. - LongIndex inc_next_strided_; - /// amount (in byte) to increment pointer from last access to first access - /// of next tile - LongIndex inc_next_; - /// amount (in byte) to increment pointer from first access of current tile - /// to first access of next tile - LongIndex inc_advance_; - - public: - // Default ctor - CUTLASS_HOST_DEVICE - Params() : stride_(0), inc_contiguous_(0), inc_strided_(0), inc_next_(0), inc_advance_(0) {} - - /// Construct the Params object given a pitch-linear tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) : stride_({layout.stride(0), layout.stride(1)}) - { - inc_contiguous_ = (LongIndex(stride_[0]) * ThreadMap::Delta::kContiguous) * - sizeof_bits::value / 8; - - inc_strided_ = (LongIndex(stride_[1]) * ThreadMap::Delta::kStrided) * - sizeof_bits::value / 8; - - inc_next_strided_ = - inc_strided_ - LongIndex(ThreadMap::Iterations::kContiguous - 1) * inc_contiguous_; - - if (kAdvanceRank) { - // advance along strided dimension - inc_advance_ = - Shape::kStrided * LongIndex(stride_[1]) * sizeof_bits::value / 8; - } else { - // advance along contiguous dimension - inc_advance_ = Shape::kContiguous * stride_[0] * sizeof_bits::value / 8; - } - - inc_next_ = inc_advance_ - - LongIndex(ThreadMap::Iterations::kContiguous - 1) * inc_contiguous_ - - LongIndex(ThreadMap::Iterations::kStrided - 1) * inc_strided_; - }; - }; - -private: - /// Internal pointer type permits fast address arithmetic - using BytePointer = char*; - - // - // Data members - // - - /// Parameters object with precomputed internal state - Params params_; - - /// Internal pointer to first access of tile - BytePointer pointer_; - - UnderlyingPredicates the_predicates; - Mask residual_tile_mask; - -private: - /// Computes predicates based on internally tracked per-thread offset. - CUTLASS_DEVICE - void compute_predicates_( - /// Extent of the matrix window - TensorCoord extent, - /// optionally, simplify predicate calculation during 'steady state' phase - bool is_steady_state = false) - { - the_predicates.compute_predicates_(extent, is_steady_state); - } - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - ///< Precomputed parameters object - Params const& params, - ///< Pointer to start of tensor - Pointer pointer, - ///< Extent of tensor - TensorCoord extent, - ///< ID of each participating thread - int thread_id, - ///< Initial offset of threadblock - TensorCoord const& threadblock_offset, - int const* indices = nullptr ///< gather/scatter indices, note no support for - ///< gather/scatter at this specialization - ) - : params_(params), - pointer_(reinterpret_cast(const_cast(pointer))), - the_predicates(extent) - { - the_predicates.set_predicates(thread_id, threadblock_offset); - - // update internal pointers - Layout layout(params_.stride_); - add_pointer_offset(layout(the_predicates.thread_offset_)); - } - - /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileAccessIteratorResidualLast(params, - pointer, - extent, - thread_id, - make_Coord(0, 0)) - { - } - - /// Overrides the internal iteration index - CUTLASS_HOST_DEVICE - void set_iteration_index(int index) { the_predicates.set_iteration_index(index); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool is_residual_tile) - { - if (is_residual_tile) { the_predicates.set_mask(residual_tile_mask); } - } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - pointer_ += sizeof_bits::value * pointer_offset / 8; - } - - /// Advances an iterator along logical dimensions of matrix in units of whole - /// tiles - CUTLASS_HOST_DEVICE - void add_tile_offset(TensorCoord const& tile_offset) - { - if (kAdvanceRank) { - pointer_ += params_.inc_advance_ * LongIndex(tile_offset[1]); - pointer_ += Shape::kContiguous * tile_offset[0]; - } else { - pointer_ += params_.inc_advance_ * LongIndex(tile_offset[0]); - pointer_ += Shape::kStrided * tile_offset[1]; - } - } - - /// Returns a pointer - CUTLASS_HOST_DEVICE - AccessType* get() const - { - return reinterpret_cast(pointer_) + the_predicates.iteration_vector_; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast& operator++() - { - the_predicates.operator++(); - ++the_predicates.iteration_vector_; - if (the_predicates.iteration_vector_ < kAccessesPerVector) { return *this; } - - the_predicates.iteration_vector_ = 0; - ++the_predicates.iteration_contiguous_; - - if (the_predicates.iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { - pointer_ += params_.inc_contiguous_; - return *this; - } - - // Enter here only if (iteration_contiguous_ == - // ThreadMap::Iteration::kContiguous) - the_predicates.iteration_contiguous_ = 0; - ++the_predicates.iteration_strided_; - - if (the_predicates.iteration_strided_ < ThreadMap::Iterations::kStrided) { - pointer_ += params_.inc_next_strided_; - return *this; - } - - // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) - // which means we enter the next tile. - the_predicates.iteration_strided_ = 0; - - // advance to next tile - pointer_ += params_.inc_next_; - - // now return to start tile - if the iterator is subsequently advanced, this - // subtraction as well as the subsequent integer addition are both elided by - // the compiler. - pointer_ -= params_.inc_advance_; - - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast operator++(int) - { - PredicatedTileAccessIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { the_predicates.clear_mask(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { the_predicates.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { the_predicates.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { the_predicates.get_mask(mask); } - - /// Returns whether access is valid or not - CUTLASS_HOST_DEVICE - bool valid() { return the_predicates.valid(); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileAccessIteratorResidualLast for affine rank 2 -/// column-major data. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// -template -class PredicatedTileAccessIteratorResidualLast { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may along advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - using Layout = layout::AffineRank2ColumnMajor; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - using AccessType = AccessType_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - // Map to the underlying AffineRankN<2> layout - using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< - layout::PitchLinearShape, - Element, - layout::AffineRankN<2>, - (kAdvanceRank == 0 ? 0 : 1), - ThreadMap, - AccessType>; - - static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; - - /// Predicate vector stores mask to guard accesses - using Mask = typename UnderlyingIterator::Mask; - - /// Parameters object is precomputed state and is host-constructible - class Params { - private: - friend PredicatedTileAccessIteratorResidualLast; - - /// Parameters object - typename UnderlyingIterator::Params params_; - - public: - /// Default ctor - CUTLASS_HOST_DEVICE - Params() {} - - /// Construct the Params object given an AffineRankN<2> tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) - : params_(layout::AffineRankN<2>(layout.stride(0), layout.stride(1))){}; - }; - -private: - // - // Data members - // - - /// Underlying AffineRankN<2> tile iterator - UnderlyingIterator iterator_; - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - ///< Precomputed parameters object - Params const& params, - ///< Pointer to start of tensor - Pointer pointer, - ///< Extent of tensor - TensorCoord extent, - ///< ID of each participating thread - int thread_id, - ///< Initial offset of threadblock - TensorCoord const& threadblock_offset, - int const* indices = nullptr ///< gather/scatter indices, note no support for - ///< gather/scatter at this specialization - ) - : iterator_(params.params_, - pointer, - layout::PitchLinearCoord(extent.row(), extent.column()), - thread_id, - layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column())) - { - } - - /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileAccessIteratorResidualLast(params, - pointer, - extent, - thread_id, - make_Coord(0, 0)) - { - } - - /// Overrides the internal iteration index - CUTLASS_HOST_DEVICE - void set_iteration_index(int index) { iterator_.set_iteration_index(index); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - iterator_.add_pointer_offset(pointer_offset); - } - - /// Advances an iterator along logical dimensions of matrix in units of whole - /// tiles - CUTLASS_HOST_DEVICE - void add_tile_offset(TensorCoord const& tile_offset) - { - iterator_.add_tile_offset(make_Coord(tile_offset.row(), tile_offset.column())); - } - - /// Returns a pointer - CUTLASS_HOST_DEVICE - AccessType* get() const { return reinterpret_cast(iterator_.get()); } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast& operator++() - { - ++iterator_; - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast operator++(int) - { - PredicatedTileAccessIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { iterator_.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { iterator_.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { iterator_.get_mask(mask); } - - /// Returns whether access is valid or not - CUTLASS_HOST_DEVICE - bool valid() { return iterator_.valid(); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileAccessIteratorResidualLast for affine rank-2 -/// row-major data. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// -template -class PredicatedTileAccessIteratorResidualLast { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may along advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - using Layout = layout::AffineRank2RowMajor; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - using AccessType = AccessType_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - // Map to the underlying AffineRankN<2> layout - using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< - layout::PitchLinearShape, - Element, - layout::AffineRankN<2>, - (kAdvanceRank == 0 ? 1 : 0), - ThreadMap, - AccessType>; - - static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; - - /// Predicate vector stores mask to guard accesses - using Mask = typename UnderlyingIterator::Mask; - - /// Parameters object is precomputed state and is host-constructible - class Params { - private: - friend PredicatedTileAccessIteratorResidualLast; - - /// Parameters object - typename UnderlyingIterator::Params params_; - - public: - /// Default ctor - CUTLASS_HOST_DEVICE - Params() {} - - /// Construct the Params object given an AffineRankN<2> tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) - : params_(layout::AffineRankN<2>(layout.stride(1), layout.stride(0))){}; - }; - -private: - // - // Data members - // - - /// Underlying AffineRankN<2> tile iterator - UnderlyingIterator iterator_; - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - ///< Precomputed parameters object - Params const& params, - ///< Pointer to start of tensor - Pointer pointer, - ///< Extent of tensor - TensorCoord extent, - ///< ID of each participating thread - int thread_id, - ///< Initial offset of threadblock - TensorCoord const& threadblock_offset, - int const* indices = nullptr ///< gather/scatter indices, note no support for - ///< gather/scatter at this specialization - ) - : iterator_(params.params_, - pointer, - layout::PitchLinearCoord(extent.column(), extent.row()), - thread_id, - layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row())) - { - } - - /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileAccessIteratorResidualLast(params, - pointer, - extent, - thread_id, - make_Coord(0, 0)) - { - } - - /// Overrides the internal iteration index - CUTLASS_HOST_DEVICE - void set_iteration_index(int index) { iterator_.set_iteration_index(index); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - iterator_.add_pointer_offset(pointer_offset); - } - - /// Advances an iterator along logical dimensions of matrix in units of whole - /// tiles - CUTLASS_HOST_DEVICE - void add_tile_offset(TensorCoord const& tile_offset) - { - iterator_.add_tile_offset(make_Coord(tile_offset.column(), tile_offset.row())); - } - - /// Returns a pointer - CUTLASS_HOST_DEVICE - AccessType* get() const { return reinterpret_cast(iterator_.get()); } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast& operator++() - { - ++iterator_; - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast operator++(int) - { - PredicatedTileAccessIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { iterator_.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { iterator_.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { iterator_.get_mask(mask); } - - /// Returns whether access is valid or not - CUTLASS_HOST_DEVICE - bool valid() { return iterator_.valid(); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileAccessIteratorResidualLast for column-major -/// interleaved data. It is mapped to the congruous layout. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// - -template -class PredicatedTileAccessIteratorResidualLast, - AdvanceRank, - ThreadMap_, - AccessType_, - false> { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may along advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - static int const kInterleavedK = InterleavedK; - using Layout = layout::ColumnMajorInterleaved; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - using AccessType = AccessType_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< - layout::PitchLinearShape, - Element, - layout::PitchLinear, - (kAdvanceRank == 0 ? 0 : 1), - ThreadMap, - AccessType>; - - static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; - - /// Predicate vector stores mask to guard accesses - using Mask = typename UnderlyingIterator::Mask; - - /// Parameters object is precomputed state and is host-constructible - class Params { - private: - friend PredicatedTileAccessIteratorResidualLast; - - /// Parameters object - typename UnderlyingIterator::Params params_; - - public: - CUTLASS_HOST_DEVICE - Params() {} - - /// Construct the Params object given a pitch-linear tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} - - CUTLASS_HOST_DEVICE - Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} - }; - -private: - // - // Data members - // - - /// Underlying pitch-linear tile iterator - UnderlyingIterator iterator_; - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - /// Precomputed parameters object - Params const& params, - /// Pointer to start of tensor - Pointer pointer, - /// Extent of tensor - TensorCoord extent, - /// ID of each participating thread - int thread_id, - /// Initial offset of threadblock - TensorCoord const& threadblock_offset, - int const* indices = nullptr ///< gather/scatter indices, note no support for - ///< gather/scatter at this specialization - ) - : iterator_(params.params_, - pointer, - layout::PitchLinearCoord(extent.row() * kInterleavedK, - extent.column() / kInterleavedK), - thread_id, - layout::PitchLinearCoord(threadblock_offset.row() * kInterleavedK, - threadblock_offset.column() / kInterleavedK)) - { - } - - /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileAccessIteratorResidualLast(params, - pointer, - extent, - thread_id, - make_Coord(0, 0)) - { - } - - /// Overrides the internal iteration index - CUTLASS_HOST_DEVICE - void set_iteration_index(int index) { iterator_.set_iteration_index(index); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - iterator_.add_pointer_offset(pointer_offset); - } - - /// Advances an iterator along logical dimensions of matrix in units of whole - /// tiles - CUTLASS_HOST_DEVICE - void add_tile_offset(TensorCoord const& tile_offset) - { - iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); - } - - /// Returns a pointer - CUTLASS_HOST_DEVICE - AccessType* get() const { return reinterpret_cast(iterator_.get()); } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast& operator++() - { - ++iterator_; - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast operator++(int) - { - PredicatedTileAccessIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { iterator_.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { iterator_.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { iterator_.get_mask(mask); } - - /// Returns whether access is valid or not - CUTLASS_HOST_DEVICE - bool valid() { return iterator_.valid(); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileAccessIteratorResidualLast for row-major -/// interleaved data. -// It is mapped to the congruous layout. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// -template -class PredicatedTileAccessIteratorResidualLast, - AdvanceRank, - ThreadMap_, - AccessType_, - false> { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may along advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - static int const kInterleavedK = InterleavedK; - using Layout = layout::RowMajorInterleaved; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - using AccessType = AccessType_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< - layout::PitchLinearShape, - Element, - layout::PitchLinear, - (kAdvanceRank == 0 ? 1 : 0), - ThreadMap, - AccessType>; - - static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; - - /// Predicate vector stores mask to guard accesses - using Mask = typename UnderlyingIterator::Mask; - - /// Parameters object is precomputed state and is host-constructible - class Params { - private: - friend PredicatedTileAccessIteratorResidualLast; - - /// Parameters object - typename UnderlyingIterator::Params params_; - - public: - CUTLASS_HOST_DEVICE - Params() {} - - /// Construct the Params object given a pitch-linear tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} - - CUTLASS_HOST_DEVICE - Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} - }; - -private: - // - // Data members - // - - /// Underlying pitch-linear tile iterator - UnderlyingIterator iterator_; - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - /// Precomputed parameters object - Params const& params, - /// Pointer to start of tensor - Pointer pointer, - /// Extent of tensor - TensorCoord extent, - /// ID of each participating thread - int thread_id, - /// Initial offset of threadblock - TensorCoord const& threadblock_offset, - int const* indices = nullptr ///< gather/scatter indices, note no support for - ///< gather/scatter at this specialization - ) - : iterator_(params.params_, - pointer, - layout::PitchLinearCoord(extent.column() * kInterleavedK, - extent.row() / kInterleavedK), - thread_id, - layout::PitchLinearCoord(threadblock_offset.column() * kInterleavedK, - threadblock_offset.row() / kInterleavedK)) - { - } - - /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast( - Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileAccessIteratorResidualLast(params, - pointer, - extent, - thread_id, - make_Coord(0, 0)) - { - } - - /// Overrides the internal iteration index - CUTLASS_HOST_DEVICE - void set_iteration_index(int index) { iterator_.set_iteration_index(index); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - iterator_.add_pointer_offset(pointer_offset); - } - - /// Advances an iterator along logical dimensions of matrix in units of whole - /// tiles - CUTLASS_HOST_DEVICE - void add_tile_offset(TensorCoord const& tile_offset) - { - iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); - } - - /// Returns a pointer - CUTLASS_HOST_DEVICE - AccessType* get() const { return reinterpret_cast(iterator_.get()); } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast& operator++() - { - ++iterator_; - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileAccessIteratorResidualLast operator++(int) - { - PredicatedTileAccessIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { iterator_.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { iterator_.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { iterator_.get_mask(mask); } - - /// Returns whether access is valid or not - CUTLASS_HOST_DEVICE - bool valid() { return iterator_.valid(); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -} // namespace threadblock -} // namespace transform -} // namespace cutlass - -//////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_atomic.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_atomic.h deleted file mode 100644 index 8d4173f1a6a2e5a073d7da4c129f471d4f394632..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_atomic.h +++ /dev/null @@ -1,886 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once -#include -#include -#include -namespace cutlass { -namespace epilogue { -namespace threadblock { - -template -struct atomic_store {}; - -template -struct atomic_store::value>::type> { - using Element = typename AccessType::Element; - static const int kCount = AccessType::kElements; - - CUTLASS_DEVICE - atomic_store(AccessType const& D, void* ptr, bool pred_guard) - { - static_assert(!(kCount % 2), "kCount must be even"); - half2* p = reinterpret_cast(ptr); - uint const* data = reinterpret_cast(&D); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - : - : "r"((int)pred_guard)); - for (int i = 0; i < kCount / 2; i++) { - asm volatile(" @p red.relaxed.global.add.noftz.f16x2 [%0], %1;\n" - : - : "l"(p + i), "r"(data[i])); - } - asm volatile("}\n" ::); - } -}; - -template -struct atomic_store::value>::type> { - using Element = typename AccessType::Element; - static const int kCount = AccessType::kElements; - - CUTLASS_DEVICE - atomic_store(AccessType const& D, void* ptr, bool pred_guard) - { - Element* p = reinterpret_cast(ptr); - uint const* data = reinterpret_cast(&D); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - : - : "r"((int)pred_guard)); - for (int i = 0; i < kCount; i++) { - asm volatile(" @p red.relaxed.global.add.f32 [%0], %1;\n" - : - : "l"(p + i), "r"(data[i])); - } - asm volatile("}\n" ::); - } -}; - -template -class PredicatedTileIteratorAffineRankNAtomic { -public: - using ThreadMap = ThreadMap_; - using Shape = typename ThreadMap::Shape; - - using Element = Element_; - - using Layout = layout::AffineRankN; - using TensorRef = TensorRef; - using TensorView = TensorView; - using ConstTensorRef = typename TensorRef::ConstTensorRef; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - using TensorCoord = typename Layout::TensorCoord; - - static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; - static int const kThreads = ThreadMap::kThreads; - static int const kIterations = ThreadMap::Count::kTile; - - static_assert(ThreadMap::Iterations::kRow > 0, "ThreadMap::Iterations::kRow must be > 0"); - static_assert(ThreadMap::Iterations::kGroup > 0, "ThreadMap::Iterations::kGroup must be > 0"); - static_assert(ThreadMap::Iterations::kCluster > 0, - "ThreadMap::Iterations::kCluster must be > 0"); - static_assert(ThreadMap::Iterations::kColumn > 0, "ThreadMap::Iterations::kColumn must be > 0"); - static_assert(!(Layout::kRank % 2), - "Layout rank must be even. This assumes the first half of the " - "modes correspond to the 'row' " - "and the second half of the modes correspond to the 'column'"); - - static bool const kBigEndian = false; - - /// Fragment object - using Fragment = Array; - - /// Memory access size - using AccessType = AlignedArray; - - // - // Parameters struct - // - - /// Parameters structure - struct Params { - // - // Data members - // - - Layout layout; - - /// Stride in units of bytes along M modes - Coord stride_m; - - /// Stride in units of bytes along N modes - Coord stride_n; - - /// Fast divmod objects divided by tensor extents - FastDivmod divmod_m[(Layout::kRank == 2) ? 1 : (Layout::kRank / 2 - 1)]; - - /// Fast divmod objects divided by tensor extents - FastDivmod divmod_n[(Layout::kRank == 2) ? 1 : (Layout::kRank / 2 - 1)]; - - int64_t rank2_inc_col; - int64_t rank2_inc_row; - - // - // Methods - // - CUTLASS_HOST_DEVICE - Params() {} - - CUTLASS_HOST_DEVICE - Params(TensorCoord const& extent, Layout const& layout_) : layout(layout_) - { - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < Layout::kRank / 2; ++i) { - stride_m[i] = OffsetBytes(layout_.stride()[i]); - stride_n[i] = OffsetBytes(layout_.stride()[i + Layout::kRank / 2]); - } - - if (kBigEndian) { - // "Big Endian" scheme - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < Layout::kRank / 2 - 1; ++i) { - divmod_m[i] = FastDivmod(extent[i + 1]); - divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2 + 1]); - } - } else { - // "Little Endian" scheme - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < Layout::kRank / 2 - 1; ++i) { - divmod_m[i] = FastDivmod(extent[i]); - divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2]); - } - } - } - - CUTLASS_HOST_DEVICE - Params(Layout const& layout_) : layout(layout_) - { - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < Layout::kRank / 2; ++i) { - stride_m[i] = OffsetBytes(layout_.stride()[i]); - stride_n[i] = OffsetBytes(layout_.stride()[i + Layout::kRank / 2]); - } - - rank2_inc_col = ThreadMap::Delta::kColumn * stride_n[0]; - rank2_inc_row = ThreadMap::Delta::kRow * stride_m[0]; - } - }; - - /// Mask object - struct Mask { - static int const kCount = ThreadMap::Iterations::kColumn; - - /// Predicate state - bool predicates[kCount]; - - // - // Mask - // - CUTLASS_HOST_DEVICE - Mask() { enable(); } - - ///< Efficiently disables all accesses guarded by mask - CUTLASS_HOST_DEVICE void clear() - { - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kCount; ++i) { predicates[i] = false; } - } - - ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask - CUTLASS_DEVICE void enable() - { - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kCount; ++i) { predicates[i] = true; } - } - }; - -private: - // - // Data members - // - - /// Parameters structure containing reference and precomputed state. - Params params_; - - /// Byte-level pointer - uint8_t* byte_pointer_; - - /// Array of boolean values to contain steady-state predicates - Mask mask_; - - /// Extent of the matrix tile in rows - Index extent_row_; - - /// Extent of the matrix tile in columns - Index extent_col_; - - /// A thread's starting row position (assuming steady-state predicates have - /// been computed) - Index thread_start_row_; - - /// A thread's starting column position (assuming steady-state predicates have - /// been computed) - Index thread_start_column_; - - /// Internal state counter - int state_[3]; - - /// Offsets in columns, cached for performance - int64_t offset_modes_n_[ThreadMap::Iterations::kColumn]; - - // - // Static asserts about internal strides - // - - static_assert(sizeof(extent_row_) == 4, "Expected 32b extents"); - static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); - -private: - // - // Methods - // - -public: - // - // Methods - // - - /// Constructor - CUTLASS_DEVICE - PredicatedTileIteratorAffineRankNAtomic( - Params const& params, - Element* pointer, - MatrixCoord extent, - int thread_idx, - MatrixCoord threadblock_offset = MatrixCoord(), - int const* indices = nullptr ///< gather/scatter indices, note no support for - ///< gather/scatter at this specialization - ) - : params_(params) - { - MatrixCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; - - extent_row_ = extent.row(); - extent_col_ = extent.column(); - - thread_start_row_ = thread_offset.row(); - thread_start_column_ = thread_offset.column(); - - if (Layout::kRank > 2) { - // Initialize predicates - CUTLASS_PRAGMA_UNROLL - for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { - // - // Compute coordinate and decompose into N modes - // - - int coord_n = thread_start_column_ + c * ThreadMap::Delta::kColumn; - - mask_.predicates[c] = coord_n < extent.column(); - - Coord modes_n; - - int64_t offset_modes_n = 0; - - if (kBigEndian) { - modes_n = CoordinateDecomposition(coord_n, params_.divmod_n); - - offset_modes_n = dot(modes_n, params_.stride_n); - } else { - modes_n = CoordinateDecompositionLittleEndian( - coord_n, params_.divmod_n); - - offset_modes_n = dot(modes_n, params_.stride_n); - } - - offset_modes_n_[c] = offset_modes_n; - } - - if (!pointer) { mask_.clear(); } - } - - // Initialize pointer - byte_pointer_ = reinterpret_cast(pointer); - - // Initialize internal state counter - state_[0] = state_[1] = state_[2] = 0; - } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - byte_pointer_ += pointer_offset * sizeof_bits::value / 8; - } - - /// Stores a fragment to memory - CUTLASS_DEVICE - void store_with_byte_offset(Fragment const& frag, int64_t byte_offset) - { - uint8_t* byte_pointer = byte_pointer_; - AccessType const* frag_ptr = reinterpret_cast(&frag); - - CUTLASS_PRAGMA_UNROLL - for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { - CUTLASS_PRAGMA_UNROLL - for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { - int row_begin = thread_start_row_ + group * ThreadMap::Delta::kGroup + - cluster * ThreadMap::Delta::kCluster; - int64_t offset_modes_m = row_begin * params_.stride_m[0]; - - CUTLASS_PRAGMA_UNROLL - for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { - int frag_row_idx = - (row + ThreadMap::Iterations::kRow * - (group + ThreadMap::Iterations::kGroup * cluster)); - - // - // Compute coordinate and decompose into M modes - // - - int coord_m = row * ThreadMap::Delta::kRow + row_begin; - - Coord modes_m; - - if (Layout::kRank > 2) { - if (kBigEndian) { - modes_m = CoordinateDecomposition(coord_m, - params_.divmod_m); - } else { - modes_m = CoordinateDecompositionLittleEndian( - coord_m, params_.divmod_m); - } - - offset_modes_m = dot(modes_m, params_.stride_m); - } - - // - // Compute the offset due to modes M - // - - bool row_guard = (coord_m < extent_row_); - int64_t offset_modes_n = thread_start_column_ * params_.stride_n[0]; - - CUTLASS_PRAGMA_UNROLL - for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { - // - // Compute coordinate and decompose into N modes - // - - if (Layout::kRank > 2) { offset_modes_n = offset_modes_n_[column]; } - - // - // Compute the pointer and access - // - bool guard; - if (Layout::kRank > 2) { - guard = row_guard && mask_.predicates[column]; - } else { - guard = (coord_m < extent_row_) && - ((thread_start_column_ + ThreadMap::Delta::kColumn * column) < - extent_col_); - } - - atomic_store( - frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], - (void*)(byte_pointer + offset_modes_m + offset_modes_n + byte_offset), - guard); - - if (Layout::kRank == 2) { offset_modes_n += params_.rank2_inc_col; } - } - - if (Layout::kRank == 2) { offset_modes_m += params_.rank2_inc_row; } - } - } - } - } - - /// Stores a fragment to memory - CUTLASS_DEVICE - void store(Fragment const& frag) { store_with_byte_offset(frag, 0); } - - CUTLASS_DEVICE - void load(Fragment& frag) {} - - /// Advances to the next position to load or store - CUTLASS_HOST_DEVICE - PredicatedTileIteratorAffineRankNAtomic& operator++() - { - ++state_[0]; - thread_start_row_ += ThreadMap::Shape::kRow; - - if (state_[0] == ThreadMap::Count::kRow) { - state_[0] = 0; - ++state_[1]; - - thread_start_row_ += - (ThreadMap::Shape::kGroup - 1) * ThreadMap::Shape::kRow * ThreadMap::Count::kRow; - - if (state_[1] == ThreadMap::Count::kGroup) { - state_[1] = 0; - ++state_[2]; - - thread_start_row_ += ThreadMap::Count::kGroup * ThreadMap::Shape::kGroup * - ThreadMap::Count::kRow * ThreadMap::Shape::kRow; - - if (state_[2] == ThreadMap::Count::kCluster) { state_[2] = 0; } - } - } - - return *this; - } - - ///< Efficiently disables all accesses guarded by mask - CUTLASS_DEVICE void clear_mask() { mask_.clear(); } - - ///< Efficiently enables all accesses guarded by mask - CUTLASS_DEVICE void enable_mask() { mask_.enable(); } - - ///< Sets the mask - CUTLASS_DEVICE void get_mask(Mask& mask) { mask = mask_; } - - ///< Sets the mask - CUTLASS_DEVICE void set_mask(Mask const& mask) { mask_ = mask; } -}; - -template -class PredicatedTileIteratorAtomic { -public: - using ThreadMap = ThreadMap_; - using Shape = typename ThreadMap::Shape; - - using Element = Element_; - - using Layout = layout::RowMajor; - using TensorRef = TensorRef; - using ConstTensorRef = typename TensorRef::ConstTensorRef; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - using TensorCoord = MatrixCoord; - - static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; - static int const kThreads = ThreadMap::kThreads; - static int const kIterations = ThreadMap::Count::kTile; - - static bool constexpr PermuteD = !layout::is_trivial_permute; - - static_assert(ThreadMap::Iterations::kRow > 0, "ThreadMap::Iterations::kRow must be > 0"); - static_assert(ThreadMap::Iterations::kGroup > 0, "ThreadMap::Iterations::kGroup must be > 0"); - static_assert(ThreadMap::Iterations::kCluster > 0, - "ThreadMap::Iterations::kCluster must be > 0"); - static_assert(ThreadMap::Iterations::kColumn > 0, "ThreadMap::Iterations::kColumn must be > 0"); - - /// Fragment object - using Fragment = Array; - - /// Memory access size - using AccessType = AlignedArray; - - // - // Parameters struct - // - - /// Uses a non-template class - struct Params : PredicatedTileIteratorParams { - using Base = PredicatedTileIteratorParams; - - CUTLASS_HOST_DEVICE - Params() {} - - CUTLASS_HOST_DEVICE - Params(Layout const& layout) - : PredicatedTileIteratorParams( - layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess, - make_OutputTileThreadMapDesc()) - { - } - - CUTLASS_HOST_DEVICE - Params(Base const& base) : Base(base) {} - }; - - /// Mask object - struct Mask { - static int const kCount = ThreadMap::Iterations::kColumn; - - /// Predicate state - bool predicates[kCount]; - - // - // Mask - // - CUTLASS_HOST_DEVICE - Mask() { enable(); } - - ///< Efficiently disables all accesses guarded by mask - CUTLASS_HOST_DEVICE void clear() - { - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kCount; ++i) { predicates[i] = false; } - } - - ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask - CUTLASS_DEVICE void enable() - { - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kCount; ++i) { predicates[i] = true; } - } - }; - -private: - // - // Data members - // - - /// Parameters structure containing reference and precomputed state. - PredicatedTileIteratorParams params_; - - /// Byte-level pointer. This pointer is usually for both load() and store(), - /// unless PermuteD is performed. When having PermuteD, byte_pointer_ is only - /// for load(). - uint8_t* byte_pointer_; - - /// Byte-level pointer for store(). Due to PermuteD Op, store_byte_pointer_ - /// may be with different address computation compared to byte_pointer_. - uint8_t* store_byte_pointer_; - - /// Array of boolean values to contain steady-state predicates - Mask mask_; - - /// Extent of the matrix tile in rows - Index extent_row_; - - /// Extent of the matrix tile in rows - Index extent_column_; - - /// A thread's starting row position (assuming steady-state predicates have - /// been computed) - Index thread_start_row_; - - /// A thread's starting column - Index thread_start_column_; - - /// Internal state counter - int state_[3]; - - /// Scatter indices - int const* indices_; - - /// PermuteDLayout - PermuteDLayout permute_layout_; - - // - // Static asserts about internal strides - // - - static_assert(sizeof(extent_row_) == 4, "Expected 32b extents"); - static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); - static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides"); - -private: - // - // Methods - // - -public: - // - // Methods - // - - /// Constructor - CUTLASS_DEVICE - PredicatedTileIteratorAtomic(PredicatedTileIteratorParams const& params, - Element* pointer, - TensorCoord extent, - int thread_idx, - TensorCoord threadblock_offset = TensorCoord(), - int const* indices = nullptr) - : params_(params), - indices_(indices), - permute_layout_(PitchLinearCoord(extent.column(), extent.row()), - params_.stride * kElementsPerAccess / sizeof(AccessType)) - { - TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; - - extent_row_ = extent.row(); - extent_column_ = extent.column(); - - thread_start_row_ = thread_offset.row(); - thread_start_column_ = thread_offset.column(); - - // Initialize predicates - CUTLASS_PRAGMA_UNROLL - for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { - mask_.predicates[c] = - ((thread_offset.column() + ThreadMap::Delta::kColumn * c) < extent.column()); - } - - // Null pointer performs no accesses - if (!pointer) { mask_.clear(); } - - if (ScatterD && !indices) { mask_.clear(); } - - // Initialize byte_pointer_ - byte_pointer_ = reinterpret_cast(pointer) + - LongIndex(thread_offset.row()) * LongIndex(params_.stride) + - LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess; - - if (ScatterD) { - byte_pointer_ = - reinterpret_cast(pointer) + - LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess; - } - - // store_byte_pointer_ is set to be the same with byte_pointer_ unless - // PermuteD is used. - store_byte_pointer_ = PermuteD ? reinterpret_cast(pointer) : byte_pointer_; - - // Initialize internal state counter - state_[0] = state_[1] = state_[2] = 0; - } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - store_byte_pointer_ += pointer_offset * sizeof_bits::value / 8; - byte_pointer_ += pointer_offset * sizeof_bits::value / 8; - } - - /// Stores a fragment to memory - CUTLASS_DEVICE - void store_with_byte_offset(Fragment const& frag, int64_t byte_offset) const - { - uint8_t* byte_pointer = store_byte_pointer_; - AccessType const* frag_ptr = reinterpret_cast(&frag); - - CUTLASS_PRAGMA_UNROLL - for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { - CUTLASS_PRAGMA_UNROLL - for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { - CUTLASS_PRAGMA_UNROLL - for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { - int frag_row_idx = - (row + ThreadMap::Iterations::kRow * - (group + ThreadMap::Iterations::kGroup * cluster)); - - int row_offset = row * ThreadMap::Delta::kRow + - group * ThreadMap::Delta::kGroup + - cluster * ThreadMap::Delta::kCluster; - - bool row_guard = ((row_offset + thread_start_row_) < extent_row_); - - AccessType* memory_pointer = - reinterpret_cast(byte_pointer + byte_offset); - - if (ScatterD && row_guard) { - assert(indices_); - - memory_pointer = reinterpret_cast( - byte_pointer + byte_offset + - LongIndex(indices_[row_offset + thread_start_row_]) * - LongIndex(params_.stride)); - } - - CUTLASS_PRAGMA_UNROLL - for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { - bool guard = row_guard && mask_.predicates[column]; - - if (PermuteD) { - int col_offset = column * ThreadMap::Delta::kColumn; - - int col = col_offset + thread_start_column_; - int row = row_offset + thread_start_row_; - - // Locate memory_pointer - memory_pointer = reinterpret_cast( - byte_pointer + byte_offset + - permute_layout_(PitchLinearCoord(col, row)) * sizeof(AccessType) / - kElementsPerAccess); - } - atomic_store( - frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], - (void*)&memory_pointer[0], - guard); - - if (!PermuteD) { - memory_pointer += (ThreadMap::Delta::kColumn / kElementsPerAccess); - } - } - - if (row + 1 < ThreadMap::Iterations::kRow) { - if (!ScatterD && !PermuteD) { byte_pointer += params_.increment_row; } - } - } - - if (group + 1 < ThreadMap::Iterations::kGroup) { - byte_pointer += params_.increment_group; - } - } - - if (cluster + 1 < ThreadMap::Iterations::kCluster) { - byte_pointer += params_.increment_cluster; - } - } - } - - /// Stores a fragment to memory - CUTLASS_DEVICE - void store(Fragment const& frag) const { store_with_byte_offset(frag, 0); } - - CUTLASS_DEVICE - void load(Fragment& frag) {} - - CUTLASS_DEVICE - MatrixCoord thread_start() const - { - return MatrixCoord(thread_start_row_, thread_start_column_); - } - - /// Need to get the thread start row from the tile iterator - CUTLASS_DEVICE - int32_t thread_start_row() const { return thread_start_row_; } - - /// Need to get the thread start row from the tile iterator - CUTLASS_DEVICE - int32_t thread_start_column() const { return thread_start_column_; } - - /// Extent of the matrix in rows - CUTLASS_DEVICE - Index extent_row() const { return extent_row_; } - - /// Extent of the matrix in columns - CUTLASS_DEVICE - Index extent_column() const { return extent_column_; } - - /// Advances to the next position to load or store - CUTLASS_HOST_DEVICE - PredicatedTileIteratorAtomic& operator++() - { - ++state_[0]; - - if (!ScatterD && !PermuteD) { store_byte_pointer_ += params_.advance_row; } - - if (!ScatterD) { byte_pointer_ += params_.advance_row; } - - thread_start_row_ += ThreadMap::Shape::kRow; - - if (state_[0] == ThreadMap::Count::kRow) { - state_[0] = 0; - ++state_[1]; - byte_pointer_ += params_.advance_group; - store_byte_pointer_ += params_.advance_group; - - thread_start_row_ += - (ThreadMap::Shape::kGroup - 1) * ThreadMap::Shape::kRow * ThreadMap::Count::kRow; - - if (state_[1] == ThreadMap::Count::kGroup) { - state_[1] = 0; - ++state_[2]; - byte_pointer_ += params_.advance_cluster; - store_byte_pointer_ += params_.advance_cluster; - - thread_start_row_ += ThreadMap::Count::kGroup * ThreadMap::Shape::kGroup * - ThreadMap::Count::kRow * ThreadMap::Shape::kRow; - - if (state_[2] == ThreadMap::Count::kCluster) { - state_[2] = 0; - byte_pointer_ += params_.advance_tile; - store_byte_pointer_ += params_.advance_tile; - - thread_start_row_ += ThreadMap::Shape::kGroup * ThreadMap::Shape::kRow * - ThreadMap::Shape::kCluster * ThreadMap::Shape::kTile; - } - } - } - - return *this; - } - - /// Advances a number of positions to load or store - CUTLASS_HOST_DEVICE - PredicatedTileIteratorAtomic& operator+=(int increment) - { - // Row - state_[0] += increment; - int increment_row = state_[0] / ThreadMap::Count::kRow; - state_[0] = state_[0] % ThreadMap::Count::kRow; - - byte_pointer_ += (params_.advance_row * increment); - store_byte_pointer_ += (params_.advance_row * increment); - thread_start_row_ += (ThreadMap::Shape::kRow * increment); - - // Group - state_[1] += increment_row; - int increment_group = state_[1] / ThreadMap::Count::kGroup; - state_[1] = state_[1] % ThreadMap::Count::kGroup; - - byte_pointer_ += (params_.advance_group * increment_row); - store_byte_pointer_ += (params_.advance_group * increment_row); - thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * ThreadMap::Shape::kRow * - ThreadMap::Count::kRow * increment_row; - - // Cluster - state_[2] += increment_group; - int increment_cluster = state_[2] / ThreadMap::Count::kCluster; - state_[2] = state_[2] % ThreadMap::Count::kCluster; - - byte_pointer_ += (params_.advance_cluster * increment_group); - store_byte_pointer_ += (params_.advance_cluster * increment_group); - thread_start_row_ += ThreadMap::Count::kGroup * ThreadMap::Shape::kGroup * - ThreadMap::Count::kRow * ThreadMap::Shape::kRow * increment_group; - - // Tile - byte_pointer_ += (params_.advance_tile * increment_cluster); - store_byte_pointer_ += (params_.advance_tile * increment_cluster); - thread_start_row_ += ThreadMap::Shape::kGroup * ThreadMap::Shape::kRow * - ThreadMap::Shape::kCluster * ThreadMap::Shape::kTile * - increment_cluster; - - return *this; - } - - ///< Efficiently disables all accesses guarded by mask - CUTLASS_DEVICE void clear_mask() { mask_.clear(); } - - ///< Efficiently enables all accesses guarded by mask - CUTLASS_DEVICE void enable_mask() { mask_.enable(); } - - ///< Sets the mask - CUTLASS_DEVICE void get_mask(Mask& mask) const { mask = mask_; } - - ///< Sets the mask - CUTLASS_DEVICE void set_mask(Mask const& mask) { mask_ = mask; } -}; - -} // namespace threadblock -} // namespace epilogue -} // namespace cutlass diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_residual_last.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_residual_last.h deleted file mode 100644 index 629047dbb057384ac8f0d7fa5d557e0c070cd830..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_residual_last.h +++ /dev/null @@ -1,1938 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights - *reserved. SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, - *this list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE - *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE - *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR - *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF - *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS - *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN - *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) - *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE - *POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/*! \file - \brief Templates implementing loading of tiles from pitch-linear rank=2 - tensors. - - This iterator uses masks to guard out-of-bounds accesses. The first tile - this iterator visits maybe partial, then the remaining tiles are complete. - So, we only need to compute the predicates twice, once before the first tile - and once for the remaining full tiles which can share the same predicates. - - A precomputed "Params" object minimizes the amount of state that must be - stored in registers, and integer addition is used to advance the pointer - through memory. -*/ - -#pragma once - -#include "cutlass/arch/memory.h" -#include "cutlass/transform/threadblock/predicated_tile_access_iterator.h" - -//////////////////////////////////////////////////////////////////////////////// - -namespace cutlass { -namespace transform { -namespace threadblock { - -//////////////////////////////////////////////////////////////////////////////// - -/// PredicatedTileIteratorResidualLast -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// -/// Regular tile iterator using a precomputed control structure to minimize -/// register liveness and integer arithmetic. -/// -/// Layout is assumed to be invariant at the time the precomputed "Params" -/// object is constructed. -/// -/// Base pointer and tensor extents may be specified at the time the iterator is -/// constructed. Subsequently, they are assumed to be immutable. -/// -/// Adding a logical coordinate offset may be performed at the time the iterator -/// is constructed. Subsequent additions to logical coordinate offset may be -/// performed but are relatively expensive. -/// -/// Visitation order is intended to first visit a "residual" tile that may be -/// partially full in both the advance dimension and the steady-state dimension. -/// This is assumed to be the last tile in the iteration sequence. Advancing an -/// iterator that has just been constructed moves to the first tile that is full -/// in the advance dimension and recomputes predicates. Subsequent accesses may -/// be performed without updating internal predicates and are efficient in terms -/// of live register state and pointer arithmetic instructions. -/// -/// To be efficient, this assumes the iterator will be dereferenced and advanced -/// at least once outside any looping structure to minimize integer arithmetic. -/// -/// Accesses out of bounds are safe so long as `clear_mask()` is called prior to -/// dereferencing the iterator. -/// -/// -/// Example: -/// -/// An efficient pipeline structure may be constructed as follows: -/// -// template -// __global__ void kernel( -// typename Iterator::Params params, -// typename Iterator::Element *ptr, -// TensorCoord extent) { -// -// typename Iterator::Fragment fragment; -// -// TensorCoord threadblock_offset(0, 0); -// -// Iterator iter(params, ptr, extent, threadIdx.x, threadblock_offsets); -// -// -// fragment = *iter; // load "residue" tile first -// ++iter; // advance to first "steady state" tile and update -// internal masks -// -// -// #pragma unroll -// for (int i = Remaining - 1; i >= 0; --i) { -// -// f(fragment); -// -// if (!i) { -// iter.clear_mask(); // light-weight operation to clear masks - -// subsequent loads become NO-OPs. -// } -// -// fragment = *iter; // load tile during "steady state" phase -// ++iter; // advance to next tile - lightweight due to -// steady-state masks -// } -// } -// -// void host(TensorView view) { -// -// using Iterator = -// transform::threadblock::PredicatedTileIteratorResidualLast; -// -// typename Iterator::Params params(view.layout()); -// -// kernel(params, view.data()); -// } -/// -/// -template -class PredicatedTileIteratorResidualLast; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileIteratorResidualLast for pitch-linear data. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// -template -class PredicatedTileIteratorResidualLast { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - using Layout = layout::PitchLinear; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - /// Type used for internal memory accesses - using AccessType = - AlignedArray::value / 8)>; - - /// Underlying iterator to compute the addresses - using TileAccessIterator = PredicatedTileAccessIteratorResidualLast; - - static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; - - /// Fragment object to be loaded or stored - using Fragment = - cutlass::Array; - - /// Predicate vector stores mask to guard accesses - using Mask = typename TileAccessIterator::Mask; - - /// Parameters object is precomputed state and is host-constructible - class Params { - public: - using Base = typename TileAccessIterator::Params::Base; - - friend PredicatedTileIteratorResidualLast; - - private: - /// Parameters object - typename TileAccessIterator::Params params_; - - public: - /// Construct the Params object given a pitch-linear tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) : params_(layout) {} - - CUTLASS_HOST_DEVICE - Params() {} - - CUTLASS_HOST_DEVICE - Params(Base const& base) : params_(base) {} - }; - -private: - /// Internal pointer type permits fast address arithmetic - using BytePointer = char*; - -private: - // - // Data members - // - - /// Data member to the tile access iterator - TileAccessIterator address_iterator_; - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast( - /// Precomputed parameters object - Params const& params, - /// Pointer to start of tensor - Pointer pointer, - /// Extent of tensor - TensorCoord extent, - /// ID of each participating thread - int thread_id, - /// Initial offset of threadblock - TensorCoord const& threadblock_offset, - /// Gather indices - int const* indices = nullptr) - : address_iterator_(params.params_, pointer, extent, thread_id, threadblock_offset, indices) - { - } - - /// Construct a PredicatedTileIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) - { - } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - address_iterator_.add_pointer_offset(pointer_offset); - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast& operator++() - { - if (kAdvanceRank) - address_iterator_.add_tile_offset({0, 1}); - else - address_iterator_.add_tile_offset({1, 0}); - - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast operator++(int) - { - PredicatedTileIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { address_iterator_.clear_mask(enable); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool enable) { address_iterator_.set_residual_tile(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { address_iterator_.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { address_iterator_.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { address_iterator_.get_mask(mask); } - - CUTLASS_DEVICE - void load_with_pointer_offset(Fragment& frag, Index pointer_offset) - { - load_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); - } - - CUTLASS_DEVICE - void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) - { - AccessType* frag_ptr = reinterpret_cast(&frag); - - CUTLASS_PRAGMA_UNROLL - for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - CUTLASS_PRAGMA_UNROLL - for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < kAccessesPerVector; ++v) { - int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); - - address_iterator_.set_iteration_index(idx); - char const* byte_ptr = - reinterpret_cast(address_iterator_.get()) + byte_offset; - - AccessType const* access_ptr = reinterpret_cast(byte_ptr); - - cutlass::arch::global_load( - frag_ptr[idx], access_ptr, address_iterator_.valid()); - - ++address_iterator_; - } - } - } - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load(Fragment& frag) { load_with_byte_offset(frag, 0); } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) - { - store_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); - } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) - { - address_iterator_.set_iteration_index(0); - AccessType const* frag_ptr = reinterpret_cast(&frag); - - CUTLASS_PRAGMA_UNROLL - for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - CUTLASS_PRAGMA_UNROLL - for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < kAccessesPerVector; ++v) { - int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); - - char* byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; - AccessType* access_ptr = reinterpret_cast(byte_ptr); - - if (address_iterator_.valid()) { *access_ptr = frag_ptr[idx]; } - ++address_iterator_; - } - } - } - } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store(Fragment const& frag) { store_with_byte_offset(frag, 0); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileIteratorResidualLast for pitch-linear data. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// -template -class PredicatedTileIteratorResidualLast { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may along advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - using Layout = layout::ColumnMajor; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - using UnderlyingIterator = - PredicatedTileIteratorResidualLast, - Element, - layout::PitchLinear, - (kAdvanceRank == 0 ? 0 : 1), - ThreadMap, - AccessSize, - Gather>; - - using AccessType = typename UnderlyingIterator::AccessType; - - /// Fragment object to be loaded or stored - using Fragment = - cutlass::Array; - - /// Predicate vector stores mask to guard accesses - using Mask = typename UnderlyingIterator::Mask; - - /// Parameters object is precomputed state and is host-constructible - class Params { - private: - friend PredicatedTileIteratorResidualLast; - - /// Parameters object - typename UnderlyingIterator::Params params_; - - public: - CUTLASS_HOST_DEVICE - Params() {} - - /// Construct the Params object given a pitch-linear tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} - - CUTLASS_HOST_DEVICE - Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} - }; - -private: - // - // Data members - // - - /// Underlying pitch-linear tile iterator - UnderlyingIterator iterator_; - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast( - Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id, ///< ID of each participating thread - TensorCoord const& threadblock_offset, ///< Initial offset of threadblock - int const* indices = nullptr ///< gather/scatter indices, note no support for - ///< gather/scatter at this specialization - ) - : iterator_(params.params_, - pointer, - layout::PitchLinearCoord(extent.row(), extent.column()), - thread_id, - layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column()), - indices) - { - } - - /// Construct a PredicatedTileIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) - { - } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - iterator_.add_pointer_offset(pointer_offset); - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast& operator++() - { - ++iterator_; - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast operator++(int) - { - PredicatedTileIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { iterator_.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { iterator_.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { iterator_.get_mask(mask); } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load_with_pointer_offset(Fragment& frag, Index pointer_offset) - { - iterator_.load_with_pointer_offset(frag, pointer_offset); - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) - { - iterator_.load_with_byte_offset(frag, byte_offset); - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) - { - iterator_.store_with_pointer_offset(frag, pointer_offset); - } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) - { - iterator_.store_with_byte_offset(frag, byte_offset); - } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileIteratorResidualLast for pitch-linear data. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// -template -class PredicatedTileIteratorResidualLast { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may along advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - using Layout = layout::RowMajor; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - using UnderlyingIterator = - PredicatedTileIteratorResidualLast, - Element, - layout::PitchLinear, - (kAdvanceRank == 0 ? 1 : 0), - ThreadMap, - AccessSize, - Gather>; - - using AccessType = typename UnderlyingIterator::AccessType; - - /// Fragment object to be loaded or stored - using Fragment = - cutlass::Array; - - /// Predicate vector stores mask to guard accesses - using Mask = typename UnderlyingIterator::Mask; - - /// Parameters object is precomputed state and is host-constructible - class Params { - private: - friend PredicatedTileIteratorResidualLast; - - /// Parameters object - typename UnderlyingIterator::Params params_; - - public: - CUTLASS_HOST_DEVICE - Params() {} - - /// Construct the Params object given a pitch-linear tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} - - CUTLASS_HOST_DEVICE - Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} - }; - -private: - // - // Data members - // - - /// Underlying pitch-linear tile iterator - UnderlyingIterator iterator_; - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast( - Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id, ///< ID of each participating thread - TensorCoord const& threadblock_offset, ///< Initial offset of threadblock - int const* indices = nullptr ///< Gather indices - ) - : iterator_(params.params_, - pointer, - layout::PitchLinearCoord(extent.column(), extent.row()), - thread_id, - layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row()), - indices) - { - } - - /// Construct a PredicatedTileIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) - { - } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - iterator_.add_pointer_offset(pointer_offset); - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast& operator++() - { - ++iterator_; - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast operator++(int) - { - PredicatedTileIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { iterator_.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { iterator_.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { iterator_.get_mask(mask); } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load_with_pointer_offset(Fragment& frag, Index pointer_offset) - { - iterator_.load_with_pointer_offset(frag, pointer_offset); - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) - { - iterator_.load_with_byte_offset(frag, byte_offset); - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) - { - iterator_.store_with_pointer_offset(frag, pointer_offset); - } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) - { - iterator_.store_with_byte_offset(frag, byte_offset); - } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileIteratorResidualLast for affine rank-2 data. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// -template -class PredicatedTileIteratorResidualLast, - AdvanceRank, - ThreadMap_, - AccessSize, - false> { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - using Layout = layout::AffineRankN<2>; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - /// Type used for internal memory accesses - using AccessType = - AlignedArray::value / 8)>; - - /// Underlying iterator to compute the addresses - using TileAccessIterator = PredicatedTileAccessIteratorResidualLast; - - static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; - - /// Fragment object to be loaded or stored - using Fragment = - cutlass::Array; - - /// Predicate vector stores mask to guard accesses - using Mask = typename TileAccessIterator::Mask; - - /// Parameters object is precomputed state and is host-constructible - class Params { - public: - friend PredicatedTileIteratorResidualLast; - - private: - /// Parameters object - typename TileAccessIterator::Params params_; - - public: - /// Construct the Params object given a pitch-linear tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) : params_(layout) {} - - CUTLASS_HOST_DEVICE - Params() {} - }; - -private: - /// Internal pointer type permits fast address arithmetic - using BytePointer = char*; - -private: - // - // Data members - // - - /// Data member to the tile access iterator - TileAccessIterator address_iterator_; - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast( - /// Precomputed parameters object - Params const& params, - /// Pointer to start of tensor - Pointer pointer, - /// Extent of tensor - TensorCoord extent, - /// ID of each participating thread - int thread_id, - /// Initial offset of threadblock - TensorCoord const& threadblock_offset, - int const* indices = nullptr ///< gather/scatter indices, note no support for - ///< gather/scatter at this specialization - ) - : address_iterator_(params.params_, pointer, extent, thread_id, threadblock_offset) - { - } - - /// Construct a PredicatedTileIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) - { - } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - address_iterator_.add_pointer_offset(pointer_offset); - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast& operator++() - { - if (kAdvanceRank) - address_iterator_.add_tile_offset(make_Coord(0, 1)); - else - address_iterator_.add_tile_offset(make_Coord(1, 0)); - - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast operator++(int) - { - PredicatedTileIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { address_iterator_.clear_mask(enable); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool enable) { address_iterator_.set_residual_tile(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { address_iterator_.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { address_iterator_.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { address_iterator_.get_mask(mask); } - - CUTLASS_DEVICE - void load_with_pointer_offset(Fragment& frag, Index pointer_offset) - { - load_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); - } - - CUTLASS_DEVICE - void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) - { - AccessType* frag_ptr = reinterpret_cast(&frag); - - CUTLASS_PRAGMA_UNROLL - for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - CUTLASS_PRAGMA_UNROLL - for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < kAccessesPerVector; ++v) { - int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); - - address_iterator_.set_iteration_index(idx); - char const* byte_ptr = - reinterpret_cast(address_iterator_.get()) + byte_offset; - - AccessType const* access_ptr = reinterpret_cast(byte_ptr); - - cutlass::arch::global_load( - frag_ptr[idx], access_ptr, address_iterator_.valid()); - - ++address_iterator_; - } - } - } - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load(Fragment& frag) { load_with_byte_offset(frag, 0); } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) - { - store_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); - } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) - { - address_iterator_.set_iteration_index(0); - AccessType const* frag_ptr = reinterpret_cast(&frag); - - CUTLASS_PRAGMA_UNROLL - for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - CUTLASS_PRAGMA_UNROLL - for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < kAccessesPerVector; ++v) { - int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); - - char* byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; - AccessType* access_ptr = reinterpret_cast(byte_ptr); - - if (address_iterator_.valid()) { *access_ptr = frag_ptr[idx]; } - ++address_iterator_; - } - } - } - } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store(Fragment const& frag) { store_with_byte_offset(frag, 0); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileIteratorResidualLast for affine rank 2 -/// column-major data. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// -template -class PredicatedTileIteratorResidualLast { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may along advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - using Layout = layout::AffineRank2ColumnMajor; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - // Map to the underlying AffineRankN<2> layout - using UnderlyingIterator = - PredicatedTileIteratorResidualLast, - Element, - layout::AffineRankN<2>, - (kAdvanceRank == 0 ? 0 : 1), - ThreadMap, - AccessSize>; - - using AccessType = typename UnderlyingIterator::AccessType; - - /// Fragment object to be loaded or stored - using Fragment = - cutlass::Array; - - /// Predicate vector stores mask to guard accesses - using Mask = typename UnderlyingIterator::Mask; - - /// Parameters object is precomputed state and is host-constructible - class Params { - private: - friend PredicatedTileIteratorResidualLast; - - /// Parameters object - typename UnderlyingIterator::Params params_; - - public: - CUTLASS_HOST_DEVICE - Params() {} - - /// Construct the Params object given an AffineRankN<2> tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) - : params_(layout::AffineRankN<2>(layout.stride(0), layout.stride(1))) - { - } - }; - -private: - // - // Data members - // - - /// Underlying AffineRankN<2> tile iterator - UnderlyingIterator iterator_; - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast( - Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id, ///< ID of each participating thread - TensorCoord const& threadblock_offset, ///< Initial offset of threadblock - int const* indices = nullptr ///< gather/scatter indices, note no support for - ///< gather/scatter at this specialization - ) - : iterator_(params.params_, - pointer, - layout::PitchLinearCoord(extent.row(), extent.column()), - thread_id, - layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column())) - { - } - - /// Construct a PredicatedTileIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) - { - } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - iterator_.add_pointer_offset(pointer_offset); - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast& operator++() - { - ++iterator_; - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast operator++(int) - { - PredicatedTileIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { iterator_.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { iterator_.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { iterator_.get_mask(mask); } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load_with_pointer_offset(Fragment& frag, Index pointer_offset) - { - iterator_.load_with_pointer_offset(frag, pointer_offset); - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) - { - iterator_.load_with_byte_offset(frag, byte_offset); - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) - { - iterator_.store_with_pointer_offset(frag, pointer_offset); - } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) - { - iterator_.store_with_byte_offset(frag, byte_offset); - } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileIteratorResidualLast for affine rank 2 -/// row-major data. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// -template -class PredicatedTileIteratorResidualLast { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may along advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - using Layout = layout::AffineRank2RowMajor; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - // Map to the underlying AffineRankN<2> layout - using UnderlyingIterator = - PredicatedTileIteratorResidualLast, - Element, - layout::AffineRankN<2>, - (kAdvanceRank == 0 ? 1 : 0), - ThreadMap, - AccessSize>; - - using AccessType = typename UnderlyingIterator::AccessType; - - /// Fragment object to be loaded or stored - using Fragment = - cutlass::Array; - - /// Predicate vector stores mask to guard accesses - using Mask = typename UnderlyingIterator::Mask; - - /// Parameters object is precomputed state and is host-constructible - class Params { - private: - friend PredicatedTileIteratorResidualLast; - - /// Parameters object - typename UnderlyingIterator::Params params_; - - public: - CUTLASS_HOST_DEVICE - Params() {} - - /// Construct the Params object given an AffineRankN<2> tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) - : params_(layout::AffineRankN<2>(layout.stride(1), layout.stride(0))) - { - } - }; - -private: - // - // Data members - // - - /// Underlying AffineRankN<2> tile iterator - UnderlyingIterator iterator_; - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast( - Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id, ///< ID of each participating thread - TensorCoord const& threadblock_offset, ///< Initial offset of threadblock - int const* indices = nullptr ///< gather/scatter indices, note no support for - ///< gather/scatter at this specialization - ) - : iterator_(params.params_, - pointer, - layout::PitchLinearCoord(extent.column(), extent.row()), - thread_id, - layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row())) - { - } - - /// Construct a PredicatedTileIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) - { - } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - iterator_.add_pointer_offset(pointer_offset); - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast& operator++() - { - ++iterator_; - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast operator++(int) - { - PredicatedTileIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { iterator_.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { iterator_.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { iterator_.get_mask(mask); } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load_with_pointer_offset(Fragment& frag, Index pointer_offset) - { - iterator_.load_with_pointer_offset(frag, pointer_offset); - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) - { - iterator_.load_with_byte_offset(frag, byte_offset); - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) - { - iterator_.store_with_pointer_offset(frag, pointer_offset); - } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) - { - iterator_.store_with_byte_offset(frag, byte_offset); - } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileIteratorResidualLast for interleaved data. -/// It is mapped to the congruous layout. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// - -template -class PredicatedTileIteratorResidualLast, - AdvanceRank, - ThreadMap_, - AccessSize, - false> { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may along advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - static int const kInterleavedK = InterleavedK; - using Layout = layout::ColumnMajorInterleaved; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - using UnderlyingIterator = PredicatedTileIteratorResidualLast< - layout::PitchLinearShape, - Element, - layout::PitchLinear, - (kAdvanceRank == 0 ? 0 : 1), - ThreadMap, - AccessSize>; - - using AccessType = typename UnderlyingIterator::AccessType; - - /// Fragment object to be loaded or stored - using Fragment = - cutlass::Array; - - /// Predicate vector stores mask to guard accesses - using Mask = typename UnderlyingIterator::Mask; - - /// Parameters object is precomputed state and is host-constructible - class Params { - private: - friend PredicatedTileIteratorResidualLast; - - /// Parameters object - typename UnderlyingIterator::Params params_; - - public: - CUTLASS_HOST_DEVICE - Params() {} - - /// Construct the Params object given a pitch-linear tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} - - CUTLASS_HOST_DEVICE - Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} - }; - -private: - // - // Data members - // - - /// Underlying pitch-linear tile iterator - UnderlyingIterator iterator_; - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast( - /// Precomputed parameters object - Params const& params, - /// Pointer to start of tensor - Pointer pointer, - /// Extent of tensor - TensorCoord extent, - /// ID of each participating thread - int thread_id, - /// Initial offset of threadblock - TensorCoord const& threadblock_offset, - int const* indices = nullptr ///< gather/scatter indices, note no support for - ///< gather/scatter at this specialization - ) - : iterator_(params.params_, - pointer, - layout::PitchLinearCoord(extent.row() * kInterleavedK, - extent.column() / kInterleavedK), - thread_id, - layout::PitchLinearCoord(threadblock_offset.row() * kInterleavedK, - threadblock_offset.column() / kInterleavedK)) - { - } - - /// Construct a PredicatedTileIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) - { - } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - iterator_.add_pointer_offset(pointer_offset); - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast& operator++() - { - ++iterator_; - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast operator++(int) - { - PredicatedTileIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { iterator_.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { iterator_.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { iterator_.get_mask(mask); } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load_with_pointer_offset(Fragment& frag, Index pointer_offset) - { - iterator_.load_with_pointer_offset(frag, pointer_offset); - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) - { - iterator_.store_with_pointer_offset(frag, pointer_offset); - } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -/// Specialization of PredicatedTileIteratorResidualLast for interleaved-32 -/// data. It is mapped to the congruous layout. -/// -/// Satisfies: ForwardTileIteratorConcept | -/// ReadableContiguousTileIteratorConcept | -/// WriteableContiguousTileIteratorConcept | -/// MaskedTileIteratorConcept -/// -template -class PredicatedTileIteratorResidualLast, - AdvanceRank, - ThreadMap_, - AccessSize, - false> { -public: - static_assert(AdvanceRank == 0 || AdvanceRank == 1, - "Specialization for pitch-linear iterator may along advance along the " - "contiguous(rank=0) or strided(rank=1) dimension."); - - using Shape = Shape_; - using Element = Element_; - static int const kInterleavedK = InterleavedK; - using Layout = layout::RowMajorInterleaved; - static int const kAdvanceRank = AdvanceRank; - using ThreadMap = ThreadMap_; - - using Index = typename Layout::Index; - using LongIndex = typename Layout::LongIndex; - - using TensorRef = TensorRef; - using TensorView = TensorView; - using TensorCoord = typename Layout::TensorCoord; - - using Pointer = Element*; - using NonConstPointer = typename platform::remove_const::type*; - - using UnderlyingIterator = PredicatedTileIteratorResidualLast< - layout::PitchLinearShape, - Element, - layout::PitchLinear, - (kAdvanceRank == 0 ? 1 : 0), - ThreadMap, - AccessSize>; - - using AccessType = typename UnderlyingIterator::AccessType; - - /// Fragment object to be loaded or stored - using Fragment = - cutlass::Array; - - /// Predicate vector stores mask to guard accesses - using Mask = typename UnderlyingIterator::Mask; - - /// Parameters object is precomputed state and is host-constructible - class Params { - private: - friend PredicatedTileIteratorResidualLast; - - /// Parameters object - typename UnderlyingIterator::Params params_; - - public: - CUTLASS_HOST_DEVICE - Params() {} - - /// Construct the Params object given a pitch-linear tensor's layout - CUTLASS_HOST_DEVICE - Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} - - CUTLASS_HOST_DEVICE - Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} - }; - -private: - // - // Data members - // - - /// Underlying pitch-linear tile iterator - UnderlyingIterator iterator_; - -public: - /// Constructs a TileIterator from its precomputed state, threadblock offset, - /// and thread ID - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast( - /// Precomputed parameters object - Params const& params, - /// Pointer to start of tensor - Pointer pointer, - /// Extent of tensor - TensorCoord extent, - /// ID of each participating thread - int thread_id, - /// Initial offset of threadblock - TensorCoord const& threadblock_offset, - int const* indices = nullptr ///< gather/scatter indices, note no support for - ///< gather/scatter at this specialization - ) - : iterator_(params.params_, - pointer, - layout::PitchLinearCoord(extent.column() * kInterleavedK, - extent.row() / kInterleavedK), - thread_id, - layout::PitchLinearCoord(threadblock_offset.column() * kInterleavedK, - threadblock_offset.row() / kInterleavedK)) - { - } - - /// Construct a PredicatedTileIteratorResidualLast with zero threadblock - /// offset - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object - Pointer pointer, ///< Pointer to start of tensor - TensorCoord extent, ///< Extent of tensor - int thread_id ///< ID of each participating thread - ) - : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) - { - } - - /// Adds a pointer offset in units of Element - CUTLASS_HOST_DEVICE - void add_pointer_offset(LongIndex pointer_offset) - { - iterator_.add_pointer_offset(pointer_offset); - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast& operator++() - { - ++iterator_; - return *this; - } - - /// Advances to the next tile in memory. - /// - /// The first time this method is called, predicates are updated, and the - /// iterator's internal pointer is reverted to the first "steady state" tile. - /// Subsequent calls are lightweight and must only update the internal - /// pointer. - CUTLASS_HOST_DEVICE - PredicatedTileIteratorResidualLast operator++(int) - { - PredicatedTileIteratorResidualLast self(*this); - operator++(); - return self; - } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } - - CUTLASS_HOST_DEVICE - void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } - - /// Clears the predicate set efficiently - CUTLASS_HOST_DEVICE - void enable_mask() { iterator_.enable_mask(); } - - /// Sets the predicate mask, overriding value stored in predicate iterator - CUTLASS_HOST_DEVICE - void set_mask(Mask const& mask) { iterator_.set_mask(mask); } - - /// Gets the mask - CUTLASS_HOST_DEVICE - void get_mask(Mask& mask) { iterator_.get_mask(mask); } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load_with_pointer_offset(Fragment& frag, Index pointer_offset) - { - iterator_.load_with_pointer_offset(frag, pointer_offset); - } - - /// Loads a fragment from memory - CUTLASS_DEVICE - void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) - { - iterator_.store_with_pointer_offset(frag, pointer_offset); - } - - /// Store a fragment to memory - CUTLASS_DEVICE - void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } -}; - -//////////////////////////////////////////////////////////////////////////////// - -} // namespace threadblock -} // namespace transform -} // namespace cutlass - -//////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/transpose_warp_iterator.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/transpose_warp_iterator.h deleted file mode 100644 index 2435c07f8989dcd6ad9c99a2fd530fe6092b1e08..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/transpose_warp_iterator.h +++ /dev/null @@ -1,57 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holdvr nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include "warp_iterator_from_smem.h" - -template -struct TransposeWarpIterator { - using Iterator = char; - static bool constexpr kSupportsTranspose = false; -}; - -template < - /// Operand identity - cutlass::gemm::Operand Operand, - /// Data type of A elements - typename Element, - bool kTranspose> -struct TransposeWarpIterator< - cutlass::gemm::warp::WarpIteratorFromSmem> { - using Iterator = cutlass::gemm::warp::WarpIteratorFromSmem; - static bool constexpr kSupportsTranspose = true; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/warp_iterator_from_smem.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/warp_iterator_from_smem.h deleted file mode 100644 index 7dd59832b4b0f278408552de38849b676d6b668a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/warp_iterator_from_smem.h +++ /dev/null @@ -1,269 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights - *reserved. SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, - *this list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE - *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE - *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR - *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF - *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS - *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN - *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) - *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE - *POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/*! \file - \brief Inspired from - "cutlass/gemm/warp/mma_tensor_op_tile_access_iterator.h" Loads tiles of GEMM - operands from a RowMajor shared-memory layout into registers to use by A100 - TensorCores. - - The difference with "mma_tensor_op_tile_access_iterator.h" is that: - (1) We use "ldmatrix" to load tiles, rather than manual loads (slightly - faster) (2) We support to transpose the operand (eg read `A.transpose()` when - the shared memory holds `A`) - - This is only implemented for the specific shapes. -*/ -#pragma once - -#include - -//////////////////////////////////////////////////////////////////////////////// -namespace cutlass { -namespace gemm { -namespace warp { - -template < - /// Operand identity - Operand Operand_, - /// Data type of A elements - typename Element_, - bool kTranspose = false> -class WarpIteratorFromSmem { -public: - /// Shape of tile to load (concept: MatrixShape) - using Shape = cutlass::MatrixShape<32, 32>; - - /// Operand tag - static Operand const kOperand = Operand_; - - /// Basic check - static_assert( - kOperand == Operand::kA || kOperand == Operand::kB, - "WarpIteratorFromSmem may only be instantiated for A or B operands to warp-level Mma."); - - /// Element type - using Element = Element_; - static_assert(sizeof_bits::value == 16, "Only supported for half"); - - /// Layout of source tile - using Layout = cutlass::layout::RowMajor; - - /// Shape of one matrix product operation (concept: MatrixShape) - using InstructionShape = cutlass::MatrixShape<16, 8>; - - /// Delta between *MMA operations (in units of *MMA operations, concept: - /// MatrixShape) - static int const kOpDelta = 1; - - /// Number of participating threads - static int const kThreads = 32; - - /// TensorRef type for loading element from a tensor - using TensorRef = TensorRef; - - /// Index type - using Index = typename TensorRef::Index; - - /// Long Index type - using LongIndex = typename TensorRef::LongIndex; - - /// Coordinate for an element in the tensor - using TensorCoord = typename TensorRef::TensorCoord; - - /// Number of elements accessed per Shared Memory load - static int const kElementsPerAccess = - (sizeof_bits::value >= 32 ? 1 : 32 / sizeof_bits::value); - - using InstructionCount = MatrixShape; - - static int const kIterations = (kOperand == Operand::kA) ? InstructionCount::kColumn - : InstructionCount::kRow; - -public: - // - // Derived quantities - // - - /// Fragment object holding a thread's part of a tile - using Fragment = - Array; - - /// Memory access type - // using AccessType = AlignedArray; - using AccessType = Array; - - static int constexpr kWarpShapeDivisibleInner = - (kOperand == Operand::kA ? InstructionShape::kColumn : InstructionShape::kRow); - static int constexpr kAccessesInner = (kWarpShapeDivisibleInner / kElementsPerAccess) / 4; - static int const kTilesPerInstruction = InstructionShape::kRow / 8; - -private: - /// Underlying tensor reference - TensorRef ref_; - - /// Origin - MatrixCoord origin_; - - /// Iterations in a tile - int iterations_; - -public: - /// Constructor from TensorRef - CUTLASS_HOST_DEVICE - WarpIteratorFromSmem(TensorRef const& ref, int lane_id) - : WarpIteratorFromSmem(ref, {Shape::kRow, Shape::kColumn}, lane_id) - { - } - CUTLASS_HOST_DEVICE - WarpIteratorFromSmem(TensorRef const& ref, TensorCoord extent, int lane_id) - : ref_(ref), iterations_(0) - { - int ldsm_vec_num = (lane_id >> 3); - if (kOperand == Operand::kA) { - origin_ = MatrixCoord(lane_id % 8, 0); - static_assert(InstructionCount::kRow * kAccessesInner * kTilesPerInstruction == 4, ""); - CUTLASS_PRAGMA_UNROLL - for (int inst_m_idx = 0; inst_m_idx < InstructionCount::kRow; ++inst_m_idx) { - CUTLASS_PRAGMA_UNROLL - for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { - CUTLASS_PRAGMA_UNROLL - for (int access_m_idx = 0; access_m_idx < kTilesPerInstruction; - ++access_m_idx) { - int access_idx = - access_m_idx + - kTilesPerInstruction * (inner_idx + kAccessesInner * inst_m_idx); - - MatrixCoord offset(access_m_idx * 8 + inst_m_idx * InstructionShape::kRow, - inner_idx * 4 * kElementsPerAccess); - - if (access_idx == ldsm_vec_num) { - if (kTranspose) { offset = MatrixCoord(offset.column(), offset.row()); } - origin_ += offset; - } - } - } - } - } else { - origin_ = MatrixCoord(0, lane_id % 8); - static_assert(InstructionCount::kColumn * kAccessesInner == 4, ""); - CUTLASS_PRAGMA_UNROLL - for (int inst_n_idx = 0; inst_n_idx < InstructionCount::kColumn; ++inst_n_idx) { - CUTLASS_PRAGMA_UNROLL - for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { - int access_idx = inner_idx + kAccessesInner * inst_n_idx; - - MatrixCoord offset(inner_idx * 4 * kElementsPerAccess, inst_n_idx * 8); - - if (access_idx == ldsm_vec_num) { - if (kTranspose) { offset = MatrixCoord(offset.column(), offset.row()); } - origin_ += offset; - } - } - } - } - - ref_.add_coord_offset(origin_); - } - - /// Advances an iterator along logical dimensions of matrix in units of whole - /// tiles - CUTLASS_HOST_DEVICE - WarpIteratorFromSmem& add_tile_offset(TensorCoord const& tile_offset) - { - TensorCoord coord_offset(tile_offset.row() * Shape::kRow, - tile_offset.column() * Shape::kColumn); - if (kTranspose) { coord_offset = TensorCoord{coord_offset.column(), coord_offset.row()}; } - origin_ += coord_offset; - - ref_.add_coord_offset(coord_offset); - - return *this; - } - - /// Advances the iterator along the advance dimension - CUTLASS_DEVICE - void advance() - { - if (kOperand == Operand::kA) { - add_tile_offset({0, 1}); - } else { - add_tile_offset({1, 0}); - } - - iterations_ = 0; - } - - /// increase iterations in a tile - CUTLASS_HOST_DEVICE - WarpIteratorFromSmem& operator++() - { - iterations_++; - - if (iterations_ >= kIterations) advance(); - - return *this; - } - - /// Loads a fragment from memory at the location pointed to by the iterator. - CUTLASS_DEVICE - void load(Fragment& frag) const - { - AccessType* access_ptr = reinterpret_cast(&frag); - using LoadLayout = - typename platform::conditional::type; - - MatrixCoord offset; - if (kOperand == Operand::kA) { - offset = MatrixCoord(0, iterations_ * InstructionShape::kColumn); - } else { - offset = MatrixCoord(iterations_ * InstructionShape::kRow, 0); - } - if (kTranspose) { offset = MatrixCoord(offset.column(), offset.row()); } - cutlass::arch::ldsm(access_ptr[0], ref_.data() + ref_.offset(offset)); - } -}; - -//////////////////////////////////////////////////////////////////////////////// - -} // namespace warp -} // namespace gemm -} // namespace cutlass -//////////////////////////////////////////////////////////////////////////////// diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_backward.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_backward.h deleted file mode 100644 index 87e6df18bb04ad52278062f6bb03f40c36edc601..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_backward.h +++ /dev/null @@ -1,1965 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holdvr nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include -#include - -#include - -#include "cutlass/cutlass.h" -#include "cutlass/epilogue/thread/linear_combination.h" -#include "cutlass/epilogue/thread/scale_type.h" -#include "cutlass/fast_math.h" -#include "cutlass/functional.h" -#include "cutlass/gemm/gemm.h" -#include "cutlass/layout/matrix.h" -#include "cutlass/layout/vector.h" -#include "cutlass/numeric_conversion.h" -#include "cutlass/numeric_types.h" -#include "cutlass/tensor_ref.h" - -#include "gemm_kernel_utils.h" - -#include "cutlass/epilogue/thread/linear_combination_relu.h" -#include "cutlass/epilogue/threadblock/epilogue_smem_accumulator.h" -#include "cutlass/epilogue/warp/fragment_iterator_tensor_op.h" -#include "cutlass/epilogue/warp/tile_iterator_tensor_op.h" -#include "cutlass/gemm/device/default_gemm_configuration.h" -#include "cutlass/gemm/kernel/default_gemm.h" -#include "cutlass/gemm/threadblock/default_mma.h" -#include "cutlass/gemm/threadblock/default_mma_core_simt.h" -#include "cutlass/gemm/threadblock/default_mma_core_sm70.h" -#include "cutlass/gemm/threadblock/default_mma_core_sm75.h" -#include "cutlass/gemm/threadblock/default_mma_core_sm80.h" -#include "cutlass/matrix_shape.h" -#include "cutlass/platform/platform.h" -#include "cutlass/transform/threadblock/predicated_tile_iterator.h" -#include "cutlass/transform/threadblock/vector_iterator.h" -#include "epilogue/epilogue_pipelined.h" -#include "iterators/epilogue_predicated_tile_iterator.h" - -#include "epilogue/epilogue_grad_bias.h" -#include "gemm/custom_mma.h" -#include "gemm/find_default_mma.h" -#include "gemm/mma_accum_lambda_iterator.h" -#include "gemm/mma_from_smem.h" -#include "transform/bias_broadcast.h" -#include "transform/tile_smem_loader.h" - -#include - -using namespace gemm_kernel_utils; - -namespace { - -template -struct GmemTile { - /* - Helper functions to efficient store/load RF to gmem - - GEMM accumulators have a particular format on A100, and - it takes some compute/shared-memory to rearrange them to - a RowMajor or ColumnMajor format in global memory through - an Epilogue. The same complexity goes for loading into RF. - - This class loads/stores RF as they are, and can be used for - efficient accumulation across gemms for instance: - - ``` - GmemTile tile; - for (int i = 0; i < N; ++i) { - // ... - - Fragment accum; - if (i == 0) { - accum.clear(); - } else { - tile.load(accum); - } - mma(accum, ...); - if (i < N-1) { - // Store for next GEMM - tile.store(accum); - } else { - // Store in tensor (eg RowMajor) - epilogue(accum); - } - - // ... - } - ``` - */ - - // 128bits per thread - using AccessType = cutlass::Array; - static constexpr int32_t kBytes = sizeof(AccessType); - static constexpr int32_t kStride = kNumThreads * AccessType::kElements; - static constexpr int32_t kNumIters = FragmentType::kElements / AccessType::kElements; - static constexpr int32_t kElementsStored = kNumThreads * FragmentType::kElements; - static_assert(FragmentType::kElements % AccessType::kElements == 0, - "fragment not aligned on 128 bits"); - - float* ptr; - - CUTLASS_DEVICE void load(FragmentType& fragment, int thread_id) - { - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kNumIters; ++i) { - AccessType* __restrict__ gmem_ptr = reinterpret_cast( - ptr + thread_id * AccessType::kElements + i * kStride); - AccessType sub_fragment; - cutlass::arch::global_load(sub_fragment, gmem_ptr, true); - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < AccessType::kElements; ++j) { - fragment[i * AccessType::kElements + j] = sub_fragment[j]; - } - } - } - - CUTLASS_DEVICE void store(FragmentType const& fragment, int thread_id) - { - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kNumIters; ++i) { - AccessType* __restrict__ gmem_ptr = reinterpret_cast( - ptr + thread_id * AccessType::kElements + i * kStride); - AccessType sub_fragment; - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < AccessType::kElements; ++j) { - sub_fragment[j] = fragment[i * AccessType::kElements + j]; - } - cutlass::arch::global_store(sub_fragment, gmem_ptr, true); - } - } -}; - -template -constexpr int getWarpsPerSm() -{ - constexpr bool is_half = !cutlass::platform::is_same::value; - if (Arch::kMinComputeCapability >= 80) { return is_half ? 12 : 8; } - return 8; -} -} // namespace - -template < - // which arch we target (eg `cutlass::arch::Sm80`) - typename ArchTag_, - // input/output type - typename scalar_t_, - // run optimized kernel because memory accesses will be aligned - bool kIsAligned_, - // use dropout if enabled - bool kApplyDropout_, - // when doing a GEMM, preload the next one (uses more shmem) - bool kPreload_, - // block dimensions - int kBlockSizeI_, - int kBlockSizeJ_, - // upperbound on `max(value.shape[-1], query.shape[-1])` - int kMaxK_ = (int)cutlass::platform::numeric_limits::max(), - template class Broadcast1_ = BroadcastNoLoad, - template class Broadcast2_ = BroadcastNoLoad> -struct AttentionBackwardKernel { - using scalar_t = scalar_t_; - using output_t = scalar_t; - using output_accum_t = float; - using lse_scalar_t = float; - using accum_t = float; - using ArchTag = ArchTag_; - static constexpr bool kIsAligned = kIsAligned_; - static constexpr bool kApplyDropout = kApplyDropout_; - static constexpr bool kPreload = kPreload_; - static constexpr int kBlockSizeI = kBlockSizeI_; - static constexpr int kBlockSizeJ = kBlockSizeJ_; - static constexpr int kMaxK = kMaxK_; - - struct Params { - // Input tensors - scalar_t* query_ptr; // [Mq, nH, K] - scalar_t* key_ptr; // [Mk, nH, K] - scalar_t* value_ptr; // [Mk, nH, Kv] - lse_scalar_t* logsumexp_ptr; // [nH, Mq] - scalar_t* output_ptr; // [Mq, nH, Kv] - scalar_t* grad_output_ptr; // [Mq, nH, Kv] - accum_t* delta_ptr; // [nH, Mq] - int32_t* cu_seqlens_q_ptr = nullptr; - int32_t* cu_seqlens_k_ptr = nullptr; - - // Output tensors - output_t* grad_query_ptr; // [Mq, nH, K] - output_t* grad_key_ptr; // [Mk, nH, K] - output_t* grad_value_ptr; // [Mk, nH, Kv] - - accum_t* grad_bias1_ptr = nullptr; - accum_t* grad_bias2_ptr = nullptr; - int32_t B = 0; - int32_t N = 0; - scalar_t* bias1_ptr = nullptr; - scalar_t* bias2_ptr = nullptr; - - // Accumulators - union { - output_accum_t* workspace = nullptr; // [Mq, Kq] + [Mkv, Kq] + [Mkv, Kv] - output_accum_t* workspace_gk; - }; - output_accum_t* workspace_gv; // (will be calculated by the kernel) - output_accum_t* workspace_gq; // (will be calculated by the kernel) - - // Scale - accum_t scale; - - // Dimensions/strides - int32_t head_dim = -1; - int32_t head_dim_value = -1; - int32_t num_queries = -1; - int32_t num_keys = -1; - int32_t num_heads = -1; - - int32_t q_strideM; - int32_t k_strideM; - int32_t v_strideM; - int32_t gO_strideM; - int32_t gB_strideM; - int8_t gQKV_strideM_multiplier = 1; // 3 for packed, 1 otherwise - - // RNG sequence offset based on batch_id and head_id - unsigned long long dropout_batch_head_rng_offset; - float dropout_prob = 0.0f; - - CUTLASS_HOST_DEVICE int32_t o_strideM() const { return head_dim_value * num_heads; } - CUTLASS_HOST_DEVICE int32_t gQ_strideM() const - { - return gQKV_strideM_multiplier * num_heads * head_dim; - } - CUTLASS_HOST_DEVICE int32_t gK_strideM() const - { - return gQKV_strideM_multiplier * num_heads * head_dim; - } - CUTLASS_HOST_DEVICE int32_t gV_strideM() const - { - return gQKV_strideM_multiplier * num_heads * head_dim_value; - } - - // Everything below is only used in `advance_to_block` - // and shouldn't use registers - int64_t o_strideH; - int32_t q_strideH; - int32_t k_strideH; - int32_t v_strideH; - int64_t o_strideB; - int64_t q_strideB; - int64_t k_strideB; - int64_t v_strideB; - int64_t lse_strideB; - int64_t lse_strideH; - int64_t delta_strideB; - int64_t delta_strideH; - int32_t num_batches; - - int64_t gO_strideB = 0; - int64_t gQ_strideB = 0; - int64_t gK_strideB = 0; - int64_t gV_strideB = 0; - int64_t gB_strideB = 0; - int64_t gO_strideH = 0; - int64_t gQ_strideH = 0; - int64_t gK_strideH = 0; - int64_t gV_strideH = 0; - int64_t gB_strideH = 0; - - CUTLASS_DEVICE bool advance_to_block() - { - int64_t batch_id = blockIdx.z; - int32_t head_id = blockIdx.y; - - if (kNeedsAccumGradQ || kNeedsAccumGradK || kNeedsAccumGradV) { - assert(workspace_size() == 0 || workspace != nullptr); - - workspace += (batch_id * num_heads + head_id) * workspace_strideBH(); - workspace = warp_uniform(workspace); - workspace_gv = workspace + workspace_elements_gk(); - workspace_gq = workspace_gv + workspace_elements_gv(); - } else { - workspace = nullptr; - } - - // Advance pointers that depend on the total concatenated - // number of queries, as `num_queries` is modified in the block - // below - dropout_batch_head_rng_offset = batch_id * (num_heads * num_queries * num_keys) + - head_id * (num_queries * num_keys); - logsumexp_ptr += batch_id * lse_strideB + head_id * lse_strideH; - - query_ptr += batch_id * q_strideB + head_id * q_strideH; - key_ptr += batch_id * k_strideB + head_id * k_strideH; - value_ptr += batch_id * v_strideB + head_id * v_strideH; - output_ptr += batch_id * o_strideB + head_id * o_strideH; - grad_output_ptr += batch_id * gO_strideB + head_id * gO_strideH; - delta_ptr += batch_id * delta_strideB + head_id * delta_strideH; - - grad_query_ptr += batch_id * gQ_strideB + head_id * gQ_strideH; - grad_key_ptr += batch_id * gK_strideB + head_id * gK_strideH; - grad_value_ptr += batch_id * gV_strideB + head_id * gV_strideH; - using broadcast_1 = Broadcast1_; - using broadcast_2 = Broadcast2_; - - if (broadcast_1::kEnable && grad_bias1_ptr) { - grad_bias1_ptr += batch_id * num_queries; - } - if (broadcast_2::kEnable && grad_bias2_ptr) { - auto strideB = num_heads * num_queries * num_keys; - auto strideH = num_queries * num_keys; - grad_bias2_ptr += (batch_id / N) * strideB + head_id * strideH; - } - if (broadcast_1::kEnable && bias1_ptr) { - bias1_ptr = broadcast_1::advance(bias1_ptr, - batch_id / N, - batch_id % N, - head_id, - num_queries * N, - num_queries, - 0); - } - if (broadcast_2::kEnable && bias2_ptr) { - auto strideB = num_heads * num_queries * num_keys; - auto strideH = num_queries * num_keys; - bias2_ptr = broadcast_2::advance( - bias2_ptr, batch_id / N, batch_id % N, head_id, strideB, 0, strideH); - } - - num_queries = warp_uniform(num_queries); - num_keys = warp_uniform(num_keys); - - query_ptr = warp_uniform(query_ptr); - key_ptr = warp_uniform(key_ptr); - value_ptr = warp_uniform(value_ptr); - logsumexp_ptr = warp_uniform(logsumexp_ptr); - output_ptr = warp_uniform(output_ptr); - grad_output_ptr = warp_uniform(grad_output_ptr); - delta_ptr = warp_uniform(delta_ptr); - - grad_query_ptr = warp_uniform(grad_query_ptr); - grad_key_ptr = warp_uniform(grad_key_ptr); - grad_value_ptr = warp_uniform(grad_value_ptr); - if (broadcast_1::kEnable) { - grad_bias1_ptr = warp_uniform(grad_bias1_ptr); - bias1_ptr = warp_uniform(bias1_ptr); - } - if (broadcast_2::kEnable) { - grad_bias2_ptr = warp_uniform(grad_bias2_ptr); - bias2_ptr = warp_uniform(bias2_ptr); - } - - return true; - } - - __host__ dim3 getBlocksGrid() const { return dim3(1, num_heads, num_batches); } - __host__ dim3 getThreadsGrid() const { return dim3(kWarpSize * kNumWarpsPerBlock, 1, 1); } - CUTLASS_HOST_DEVICE int64_t workspace_elements_gk() const - { - if (!kNeedsAccumGradK) { return 0; } - return align_up(num_keys, (int32_t)kBlockSizeJ) * - align_up(head_dim, (int32_t)kBlockSizeI); - } - CUTLASS_HOST_DEVICE int64_t workspace_elements_gv() const - { - if (!kNeedsAccumGradV) { return 0; } - return align_up(num_keys, (int32_t)kBlockSizeJ) * - align_up(head_dim_value, (int32_t)kBlockSizeI); - } - CUTLASS_HOST_DEVICE int64_t workspace_elements_gq() const - { - if (!kNeedsAccumGradQ) { return 0; } - if (num_keys <= kBlockSizeJ) { return 0; } - return align_up(num_queries, (int32_t)kBlockSizeI) * - align_up(head_dim, (int32_t)kBlockSizeJ); - } - CUTLASS_HOST_DEVICE int64_t workspace_strideBH() const - { - // Aligned on 128bits - return align_up( - workspace_elements_gk() + workspace_elements_gv() + workspace_elements_gq(), - int64_t(4)); - } - CUTLASS_HOST_DEVICE int64_t workspace_size() const - { - // Returns size of buffer we need to run this kernel - return num_batches * num_heads * workspace_strideBH() * sizeof(float); - } - }; - - static constexpr int64_t kWarpSize = 32; - - // If this is true, we store and accumulate dK/dV in RF - // rather than going back to gmem every time - static constexpr bool kIsHalf = cutlass::sizeof_bits::value <= 16; - static constexpr bool kOutputInRF = kIsHalf && kMaxK <= kBlockSizeI; - static_assert(!kPreload || (kIsHalf && ArchTag::kMinComputeCapability >= 80 && kOutputInRF), - "preload MMA not supported"); - static constexpr bool kPrologueQK = kPreload; - static constexpr bool kPrologueGV = kPreload; - static constexpr bool kPrologueDOV = kPreload; - static constexpr bool kPrologueGQ = kPreload; - static constexpr bool kPrologueGK = kPreload; - - static constexpr int64_t kNumWarpsPerBlock = (kBlockSizeI * kBlockSizeJ) / (32 * 32); - - // Compute delta for the f16 kernels - // TODO: Figure out why it's slower on the f32 kernels - // (something due to RF pressure?) - // TODO: Remove condition on `kOutputInRF` - this is needed to work - // around a compiler bug on V100, not exactly sure why but I spent - // too much time on this already. Reproducible with - // (B, Mq, Mkv, K) = (1, 1, 1, 136) for instance - static constexpr bool kKernelComputesDelta = - kIsHalf && (kOutputInRF || ArchTag::kMinComputeCapability != 70); - - static constexpr bool kNeedsAccumGradQ = - !cutlass::platform::is_same::value; - static constexpr bool kNeedsAccumGradK = - !kOutputInRF && !cutlass::platform::is_same::value; - static constexpr bool kNeedsAccumGradV = - !kOutputInRF && !cutlass::platform::is_same::value; - - // Launch bounds - static constexpr int64_t kNumThreads = kWarpSize * kNumWarpsPerBlock; - static constexpr int64_t kMinBlocksPerSm = - getWarpsPerSm() / kNumWarpsPerBlock; - - using GemmType = DefaultGemmType; - using DefaultConfig = - typename cutlass::gemm::device::DefaultGemmConfiguration; - static constexpr auto kOptimalAlignement = - cutlass::platform::max(DefaultConfig::kAlignmentA, DefaultConfig::kAlignmentB); - static constexpr auto kMinimumAlignment = GemmType::kMinimumAlignment; - - struct MatmulQK { - /* - attn_T = k_j @ q_i.transpose(-2, -1) # matmul - attn_T = (attn_T - logsumexp[i_start:i_end].unsqueeze(1).transpose(-2, - -1)).exp() # epilogue - - with attn_T.shape = (kBlockSizeJ, kBlockSizeI) - */ - using ThreadblockShape = - cutlass::gemm::GemmShape; - using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; - using DefaultMma = typename cutlass::gemm::threadblock::DefaultMma< - scalar_t, // ElementA - cutlass::layout::RowMajor, // LayoutA - kIsAligned ? DefaultConfig::kAlignmentA : GemmType::kMinimumAlignment, - scalar_t, // ElementB - cutlass::layout::ColumnMajor, // LayoutB - kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment, - accum_t, // ElementC - cutlass::layout::RowMajor, // LayoutC - typename GemmType::OpClass, - ArchTag, - ThreadblockShape, - WarpShape, - typename GemmType::InstructionShape, - DefaultConfig::kStages, - typename GemmType::Operator, - false, // AccumulatorsInRowMajor = false, - cutlass::gemm::SharedMemoryClearOption::kNone>; - using MmaCore = typename DefaultMma::MmaCore; - using Mma = typename MakeCustomMma::Mma; - - // used for efficient load of bias tile (Bij) from global memory to shared - // memory - using BiasLoader = - TileSmemLoader, - MmaCore::kThreads, - // input restriction: kv_len has to be a multiple of this value - 128 / cutlass::sizeof_bits::value>; - - // Epilogue to store to shared-memory in a format that we can use later for - // the second matmul - using B2bGemm = - typename cutlass::gemm::threadblock::B2bGemm; - using AccumLambdaIterator = - typename DefaultMmaAccumLambdaIterator::Iterator; - using AccumulatorSharedStorage = typename B2bGemm::AccumulatorSharedStorage; - }; - - struct MatmulGradV { - /* - grad_v[j_start:j_end] += attn_T @ do_i # matmul - - Dimensions: (kBlockSizeJ * kNumWarpsPerBlock, kBlockSizeI, K) - (we might need to iterate multiple times on K) - */ - using ThreadblockShape = - cutlass::gemm::GemmShape; - using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; - using InstructionShape = typename GemmType::InstructionShape; - - using DefaultGemm = - cutlass::gemm::kernel::DefaultGemm; - - // if dropout: - // for computing dVj += (Pij.T * Zij) @ dOi - // Pij_dropped.T = Pij.T * Zij is computed on the fly as fragments of - // Pij.T are loaded in. The reason we do it this way is because Pij.T and - // Zij are reused in later steps, while Pij_dropped.T is only needed in - // this step. computing Pij_dropped.T on the fly allows us to avoid - // keeping all 3 of Pij_dropped.T, Pij.T, and Zij in shared memory at the - // same time. - // if no dropout: - // for computing dVj += Pij.T @ dOi - using DefaultMmaFromSmem = typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< - typename DefaultGemm::Mma, - typename MatmulQK::AccumulatorSharedStorage, - kApplyDropout>; // kScaleOperandA - - using Mma = typename DefaultMmaFromSmem::Mma; - using WarpIteratorA = typename DefaultMmaFromSmem::WarpIteratorA; - using IteratorB = typename Mma::IteratorB; - using WarpCount = typename Mma::WarpCount; - - // Epilogue - using DefaultOutputOp = typename DefaultConfig::EpilogueOutputOp; - using DefaultEpilogue = typename DefaultGemm::Epilogue; - using OutputTileIterator = - typename cutlass::epilogue::threadblock::MakePrefetchableIterator< - typename DefaultEpilogue::OutputTileIterator>::Iterator; - using AccumTileGmem = GmemTile; - }; - - struct MatmulDOIVJ { - /* - doi_t_vj = do_i @ v_j.transpose(-2, -1) # matmul - tmp = (doi_t_vj - Di.unsqueeze(1)) * attn # inplace / epilogue? - */ - using ThreadblockShape = - cutlass::gemm::GemmShape; - using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; - - using ElementC = accum_t; // CSY: Change it for better accuracy - using ElementAccum = accum_t; - - // no-op output op - epilogue just stores result to global memory - using BiasGradEpilogueOutputOp = typename cutlass::epilogue::thread::LinearCombination< - ElementC, - DefaultConfig::EpilogueOutputOp::kCount, - typename DefaultConfig::EpilogueOutputOp::ElementAccumulator, - typename DefaultConfig::EpilogueOutputOp::ElementCompute, - cutlass::epilogue::thread::ScaleType::Nothing>; - - using DefaultGemm = typename cutlass::gemm::kernel::DefaultGemm< - scalar_t, // ElementA - cutlass::layout::RowMajor, // LayoutA - kIsAligned ? DefaultConfig::kAlignmentA : GemmType::kMinimumAlignment, - scalar_t, // ElementB - cutlass::layout::ColumnMajor, // LayoutB - kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment, - ElementC, // ElementC - cutlass::layout::RowMajor, // LayoutC - ElementAccum, // ElementAccumulator - typename GemmType::OpClass, - ArchTag, - ThreadblockShape, - WarpShape, - typename GemmType::InstructionShape, - BiasGradEpilogueOutputOp, // EpilogueOutputOp - void, // ThreadblockSwizzle (not used) - // multiple preloads, dropout Zij tile, and 3 stages push us over shared - // memory capacity on A100. set a ceiling on number of stages to save - // shared memory if dropout is in use. - kPreload && kApplyDropout && (kBlockSizeI * kBlockSizeJ > 64 * 64) - ? cutlass::const_min(2, DefaultConfig::kStages) - : DefaultConfig::kStages, // Stages - false, // SplitKSerial - typename GemmType::Operator, - cutlass::gemm::SharedMemoryClearOption::kNone>; - using Mma = typename MakeCustomMma::Mma; - - // epilogue used to write bias gradient, which is just the output of this - // matmul with some operations applied to the fragment - using BiasGradEpilogue = typename DefaultGemm::Epilogue; - - // Epilogue to store to shared-memory in a format that we can use later for - // the second matmul - using B2bGemm = - typename cutlass::gemm::threadblock::B2bGemm; - using AccumulatorSharedStorage = typename B2bGemm::AccumulatorSharedStorage; - }; - - struct MatmulGradQ { - // grad_q <- tmp @ k_j - using ThreadblockShape = - cutlass::gemm::GemmShape; - using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; - using InstructionShape = typename GemmType::InstructionShape; - - using DefaultGemm = - cutlass::gemm::kernel::DefaultGemm; - - using DefaultMmaFromSmem = typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< - typename DefaultGemm::Mma, - typename MatmulDOIVJ::AccumulatorSharedStorage, - false>; // kScaleOperandA - using Mma = typename DefaultMmaFromSmem::Mma; - using IteratorB = typename Mma::IteratorB; - using WarpCount = typename Mma::WarpCount; - - // Epilogue - using DefaultOutputOp = typename DefaultConfig::EpilogueOutputOp; - using DefaultEpilogue = typename DefaultGemm::Epilogue; - using OutputTileIterator = - typename cutlass::epilogue::threadblock::MakePrefetchableIterator< - typename DefaultEpilogue::OutputTileIterator>::Iterator; - using AccumTileGmem = GmemTile; - }; - struct MatmulGradK { - // grad_k <- tmp.transpose(-2, -1) @ q_i - using ThreadblockShape = - cutlass::gemm::GemmShape; - using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; - using InstructionShape = typename GemmType::InstructionShape; - - using DefaultGemm = - cutlass::gemm::kernel::DefaultGemm; - - using DefaultMmaFromSmemN = typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< - typename DefaultGemm::Mma, - typename MatmulQK::AccumulatorSharedStorage, - false>; // kScaleOperandA - using DefaultMmaFromSmemT = typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< - typename DefaultGemm::Mma, - typename MatmulDOIVJ::AccumulatorSharedStorage, - false, // kScaleOperandA - kPreload>; // kTransposeA - using DefaultMmaFromSmem = - typename cutlass::platform::conditional::type; - using Mma = typename DefaultMmaFromSmem::Mma; - using IteratorB = typename Mma::IteratorB; - using WarpCount = typename Mma::WarpCount; - - // Epilogue - using DefaultOutputOp = typename DefaultConfig::EpilogueOutputOp; - using DefaultEpilogue = typename DefaultGemm::Epilogue; - using OutputTileIterator = - typename cutlass::epilogue::threadblock::MakePrefetchableIterator< - typename DefaultEpilogue::OutputTileIterator>::Iterator; - using AccumTileGmem = GmemTile; - }; - - using broadcast_1 = Broadcast1_; - using broadcast_2 = Broadcast2_; - - // shared storage for keeping Zij matrix. not needed if we aren't using - // dropout, in which case we use an empty array to save shared memory - using ZijSharedStorage = typename cutlass::platform::conditional< - kApplyDropout, - typename MatmulQK::AccumulatorSharedStorage, - // dummy shared storage object that takes up no space. - typename cutlass::gemm::threadblock::AccumulatorSharedStorage< -#ifdef _WIN32 - // windows builds throw the error: - // "type containing an unknown-size array is not allowed" - // if we try to make Zij shared storage zero-sized. - // To get around this just make it sized 1 on windows. - typename cutlass::gemm::GemmShape<1, 1, 0>, -#else - typename cutlass::gemm::GemmShape<0, 0, 0>, -#endif - typename MatmulQK::AccumulatorSharedStorage::Element, - typename MatmulQK::AccumulatorSharedStorage::Layout, - typename cutlass::MatrixShape<0, 0>>>::type; - - struct SharedStoragePrologue { - struct { - cutlass::Array di; // (do_i * o_i).sum(-1) - typename MatmulQK::Mma::SharedStorageA mm_qk_k; - } persistent; - union { - struct { - // part1 - after Q.K / dV / dO.V - union { - // 1. efficient load of bias tile Bij, which is then applied to Pij - // typename MatmulQK::BiasLoader::SmemTile bias; - cutlass::AlignedBuffer bias; - // 4. store Pij. it is needed: - // - in dVj += (Pij.T * Zij) @ dOi - // - in dSij = Pij * (dPij - Di) - // 6. dVj += (Pij.T * Zij) @ dOi - // 10. write to fragment - typename MatmulQK::AccumulatorSharedStorage attn_shared_storage; - }; - // 5. store Zij. it is needed: - // - to compute Pij_dropped = Pij * Zij on the fly as fragments of Pij - // are loaded for the computation of dVj. - // - to compute dPij = (dOi @ Vj.T) * Zij - // 6. used in dVj += (Pij.T * Zij) @ dOi - // 9. used in dPij = dPij_dropped * Zij - ZijSharedStorage zij; - - union { - // 2. prologue for dVj - // 6. workspace for dVj += (Pij.T * Zij) @ dOi - typename MatmulGradV::Mma::SharedStorage mm_gradV; - // 7. dVj epilogue - typename MatmulGradV::DefaultEpilogue::SharedStorage gradV_epilogue; - }; - - // 3. prologue for dPij_dropped - // 8. used in dPij_dropped = dOi @ Vj.T - typename MatmulDOIVJ::Mma::SharedStorage mm_doivj; - } part1; - - struct { - // part2 - dQ - union { - typename MatmulQK::AccumulatorSharedStorage - tmpT_shared_storage; // (from part1) - typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; - }; - typename MatmulGradK::Mma::SharedStorage mm_gradK; // (preload) - typename MatmulGradQ::Mma::SharedStorage mm_gradQ; // (preload) - union { - // store dB = dSij to global memory - typename MatmulDOIVJ::BiasGradEpilogue::SharedStorage gradB_epilogue; - typename MatmulGradQ::DefaultEpilogue::SharedStorage gradQ_epilogue; - }; - - } part2; - - struct { - // part3 - after last iteration on dQ's epilogue / dK - union { - typename MatmulQK::AccumulatorSharedStorage - tmpT_shared_storage; // (from part1) - typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; - }; - typename MatmulGradK::Mma::SharedStorage mm_gradK; // (preload) - typename MatmulGradQ::DefaultEpilogue::SharedStorage gradQ_epilogue_lastIter; - - typename MatmulGradK::DefaultEpilogue::SharedStorage gradK_epilogue; - } part3; - - struct { - // part4 - after last iteration on dK's epilogue / preload next K.Q_t - typename MatmulQK::Mma::SharedStorageB mm_qk_q; - - // If we reach end of current key, dump RF->gmem with "final" epilogues - typename MatmulGradK::DefaultEpilogue::SharedStorage gradK_epilogue_final; - typename MatmulGradV::DefaultEpilogue::SharedStorage gradV_epilogue_final; - } part4; - }; -// =========================================== -#define FIELD(INSIDE_STRUCT, FIELDNAME) \ - CUTLASS_DEVICE auto& FIELDNAME() { return INSIDE_STRUCT.FIELDNAME; } - - FIELD(persistent, di) - FIELD(persistent, mm_qk_k) - FIELD(part1, bias) - FIELD(part1, attn_shared_storage) - FIELD(part1, zij) - FIELD(part1, mm_gradV) - FIELD(part1, gradV_epilogue) - FIELD(part1, mm_doivj) - FIELD(part2, mm_gradK) - FIELD(part2, mm_gradQ) - FIELD(part2, gradB_epilogue) - FIELD(part2, gradQ_epilogue) - FIELD(part2, tmp_shared_storage) - FIELD(part3, tmpT_shared_storage) - FIELD(part3, gradQ_epilogue_lastIter) - FIELD(part3, gradK_epilogue) - FIELD(part4, mm_qk_q) - FIELD(part4, gradK_epilogue_final) - FIELD(part4, gradV_epilogue_final) - }; - - struct SharedStorageNoPrologue { - struct { - cutlass::Array di; // (do_i * o_i).sum(-1) - } persistent; - union { - struct { - // part1 - Q.K matmul - typename MatmulQK::Mma::SharedStorageA mm_qk_k; - typename MatmulQK::Mma::SharedStorageB mm_qk_q; - } part1; - - struct { - // part2 - compute gradV - union { - // 1. efficient load of bias tile Bij, which is then applied to Pij - cutlass::AlignedBuffer bias; - // 2. store Pij to shared memory. it is needed: - // - in this step, where it is used in dVj += (Pij.T * Zij) @ dOi - // - in next step where it is used in dSij = Pij * (dPij - Di) - typename MatmulQK::AccumulatorSharedStorage attn_shared_storage; - }; - // 3. store Zij. it is needed: - // - in this step, where it is used to compute Pij_dropped = Pij * Zij - // on the - // fly as fragments of Pij are loaded for the computation of dVj. - // - later to compute dPij = (dOi @ Vj.T) * Zij - ZijSharedStorage zij; - - union { - typename MatmulGradV::Mma::SharedStorage mm_gradV; - typename MatmulGradV::DefaultEpilogue::SharedStorage gradV_epilogue; - }; - } part2; - - struct { - // part3 - DO.V matmul - union { - // first compute dPij = (dOi @ Vj.T) * Zij - // and dSij = Pij * (dPij - Di) - struct { - // (from part2) - Pij for computing dSij = Pij * (dPij - Di) - typename MatmulQK::AccumulatorSharedStorage attn_shared_storage; - // (from part2) - Zij for computing dPij = dPij_dropped * Zij - ZijSharedStorage zij; - // matmul to compute dOiVj - typename MatmulDOIVJ::Mma::SharedStorage mm_doivj; - }; - // then store dB = dSij to global memory - typename MatmulDOIVJ::BiasGradEpilogue::SharedStorage gradB_epilogue; - }; - } part3; - - struct { - // part4 - compute gradQ - typename MatmulQK::AccumulatorSharedStorage tmpT_shared_storage; // (from part2) - typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; - union { - typename MatmulGradQ::Mma::SharedStorage mm_gradQ; - typename MatmulGradQ::DefaultEpilogue::SharedStorage gradQ_epilogue; - typename MatmulGradQ::DefaultEpilogue::SharedStorage gradQ_epilogue_lastIter; - }; - } part4; - - struct { - // part5 - compute gradK - typename MatmulQK::AccumulatorSharedStorage tmpT_shared_storage; // (from part2) - typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; - union { - typename MatmulGradK::Mma::SharedStorage mm_gradK; - typename MatmulGradK::DefaultEpilogue::SharedStorage gradK_epilogue; - }; - } part5; - - struct { - // part6 - store RF accumulated into gmem - typename MatmulGradK::DefaultEpilogue::SharedStorage gradK_epilogue_final; - typename MatmulGradV::DefaultEpilogue::SharedStorage gradV_epilogue_final; - } part6; - }; -// =========================================== -#define FIELD(INSIDE_STRUCT, FIELDNAME) \ - CUTLASS_DEVICE auto& FIELDNAME() { return INSIDE_STRUCT.FIELDNAME; } - - FIELD(persistent, di) - FIELD(part1, mm_qk_k) - FIELD(part1, mm_qk_q) - FIELD(part2, bias) - FIELD(part2, attn_shared_storage) - FIELD(part2, zij) - FIELD(part2, mm_gradV) - FIELD(part2, gradV_epilogue) - FIELD(part3, mm_doivj) - FIELD(part3, gradB_epilogue) - FIELD(part4, tmpT_shared_storage) - FIELD(part4, tmp_shared_storage) - FIELD(part4, mm_gradQ) - FIELD(part4, gradQ_epilogue) - FIELD(part4, gradQ_epilogue_lastIter) - FIELD(part5, mm_gradK) - FIELD(part5, gradK_epilogue) - FIELD(part6, gradK_epilogue_final) - FIELD(part6, gradV_epilogue_final) - }; - - using SharedStorage = typename cutlass::platform:: - conditional::type; - - struct OutputFragments { - typename MatmulGradV::Mma::FragmentC gradV; - typename MatmulGradK::Mma::FragmentC gradK; - - CUTLASS_DEVICE void clear() - { - gradV.clear(); - gradK.clear(); - } - }; - - static bool __host__ check_supported(Params const& p) - { - CHECK_ALIGNED_PTR(p.query_ptr, kMinimumAlignment); - CHECK_ALIGNED_PTR(p.key_ptr, kMinimumAlignment); - CHECK_ALIGNED_PTR(p.value_ptr, kMinimumAlignment); - CHECK_ALIGNED_PTR(p.output_ptr, kMinimumAlignment); - CHECK_ALIGNED_PTR(p.grad_output_ptr, kMinimumAlignment); - EVOFORMER_CHECK(p.lse_strideH % 8 == 0, "LSE is not correctly aligned"); - EVOFORMER_CHECK(p.lse_strideB % 8 == 0, "LSE is not correctly aligned"); - EVOFORMER_CHECK(p.num_heads <= 1 || p.q_strideH % kMinimumAlignment == 0, - "query is not correctly aligned (strideH)"); - EVOFORMER_CHECK(p.num_heads <= 1 || p.k_strideH % kMinimumAlignment == 0, - "key is not correctly aligned (strideH)"); - EVOFORMER_CHECK(p.num_heads <= 1 || p.v_strideH % kMinimumAlignment == 0, - "value is not correctly aligned (strideH)"); - EVOFORMER_CHECK(p.num_batches <= 1 || p.q_strideB % kMinimumAlignment == 0, - "query is not correctly aligned (strideB)"); - EVOFORMER_CHECK(p.num_batches <= 1 || p.k_strideB % kMinimumAlignment == 0, - "key is not correctly aligned (strideB)"); - EVOFORMER_CHECK(p.num_batches <= 1 || p.v_strideB % kMinimumAlignment == 0, - "value is not correctly aligned (strideB)"); - EVOFORMER_CHECK(p.q_strideM % kMinimumAlignment == 0, - "query is not correctly aligned (strideM)"); - EVOFORMER_CHECK(p.k_strideM % kMinimumAlignment == 0, - "key is not correctly aligned (strideM)"); - EVOFORMER_CHECK(p.v_strideM % kMinimumAlignment == 0, - "value is not correctly aligned (strideM)"); - EVOFORMER_CHECK(p.dropout_prob <= 1.0f && p.dropout_prob >= 0.0f, - "Invalid value for `dropout_prob`"); - EVOFORMER_CHECK(kApplyDropout || p.dropout_prob == 0.0f, - "Set `kApplyDropout`=True to support `dropout_prob > 0`"); - EVOFORMER_CHECK(p.head_dim > 0, "Invalid value for `head_dim`"); - EVOFORMER_CHECK(p.head_dim_value > 0, "Invalid value for `head_dim_value`"); - EVOFORMER_CHECK(p.num_queries > 0, "Invalid value for `num_queries`"); - EVOFORMER_CHECK(p.num_keys > 0, "Invalid value for `num_keys`"); - EVOFORMER_CHECK(p.num_heads > 0, "Invalid value for `num_heads`"); - EVOFORMER_CHECK(p.num_batches > 0, "Invalid value for `num_batches`"); - EVOFORMER_CHECK(p.head_dim <= kMaxK, "kMaxK: Expected `head_dim < kMaxK`"); - EVOFORMER_CHECK(p.head_dim_value <= kMaxK, "kMaxK: Expected `head_dim_value < kMaxK`"); - return true; - } - - static CUTLASS_DEVICE void attention_kernel(Params p) - { - extern __shared__ char smem_buffer[]; - SharedStorage& shared_storage = *((SharedStorage*)smem_buffer); - - uint16_t thread_id = threadIdx.x; - uint8_t warp_id = warp_uniform(thread_id / 32); - uint8_t lane_id = thread_id % 32; - - if (kPrologueQK) { - prologueQkNextIteration(shared_storage, p, 0, 0, warp_id, lane_id); - } - - // Computes (dO*out).sum(-1) and writes it to `p.delta_ptr` - if (kKernelComputesDelta) { - constexpr int kOptimalElements = 128 / cutlass::sizeof_bits::value; - if (p.head_dim_value % kOptimalElements == 0) { - for (int query_start = 0; query_start < p.num_queries; query_start += kBlockSizeI) { - computeDelta(p, query_start, warp_id, lane_id); - } - } else { - for (int query_start = 0; query_start < p.num_queries; query_start += kBlockSizeI) { - computeDelta<1>(p, query_start, warp_id, lane_id); - } - } - __syncthreads(); - } - - OutputFragments output_frags; - - int32_t key_start = 0; - int32_t key_end = p.num_keys / kBlockSizeJ * kBlockSizeJ; - for (; key_start < key_end; key_start += kBlockSizeJ) { - output_frags.clear(); - int32_t query_start = getQueryStart(p, key_start); - int32_t query_end = - query_start + (p.num_queries - query_start) / kBlockSizeI * kBlockSizeI; - for (; query_start < query_end; query_start += kBlockSizeI) { - processBlockIJ( - shared_storage, output_frags, p, query_start, key_start, warp_id, lane_id); - } - // last (partial) query - if (query_start < p.num_queries) { - processBlockIJ( - shared_storage, output_frags, p, query_start, key_start, warp_id, lane_id); - } - if (kOutputInRF) { - writeFragsToGmem( - shared_storage, output_frags, p, key_start, warp_id, lane_id); - } else if (getQueryStart(p, key_start) >= p.num_queries) { - zfillGradKV(p, key_start, warp_id, lane_id); - } - __syncthreads(); - } - // Last (partial) key - if (key_start != p.num_keys) { - output_frags.clear(); - int32_t query_start = getQueryStart(p, key_start); - for (; query_start < p.num_queries; query_start += kBlockSizeI) { - warp_id = warp_uniform(warp_id); - processBlockIJ( - shared_storage, output_frags, p, query_start, key_start, warp_id, lane_id); - } - if (kOutputInRF) { - writeFragsToGmem( - shared_storage, output_frags, p, key_start, warp_id, lane_id); - } else if (getQueryStart(p, key_start) >= p.num_queries) { - zfillGradKV(p, key_start, warp_id, lane_id); - } - } - } - - static CUTLASS_DEVICE void loadDi(cutlass::Array& di, - Params const& p, - int32_t query_start) - { - int32_t thread_id = threadIdx.x + threadIdx.y * blockDim.x; - if (thread_id < kBlockSizeI) { - accum_t di_rf = accum_t(0); - if (query_start + thread_id < p.num_queries) { - di_rf = p.delta_ptr[query_start + thread_id]; - } - di[thread_id] = di_rf; - } - } - - template - static CUTLASS_DEVICE void zfillGradKV(Params const& p, - int32_t key_start, - uint8_t warp_id, - uint8_t lane_id) - { - constexpr int kThreadsPerKey = 8; - constexpr int kParallelKeys = kNumThreads / kThreadsPerKey; - static_assert(kBlockSizeJ % kParallelKeys == 0, ""); - // This function is not really optimized, but should rarely be used - // It's only used when some keys are "useless" and don't attend to - // any query, due to causal masking - int thread_id = 32 * warp_id + lane_id; - int k_shift = lane_id % kThreadsPerKey; - - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < kBlockSizeJ; j += kParallelKeys) { - int key = key_start + j + (thread_id / kThreadsPerKey); - if (!skipBoundsChecks && key >= p.num_keys) { continue; } - auto gv_ptr = p.grad_value_ptr + key * p.gV_strideM(); - auto gk_ptr = p.grad_key_ptr + key * p.gK_strideM(); - - for (int k = k_shift; k < p.head_dim_value; k += kThreadsPerKey) { - gv_ptr[k] = scalar_t(0); - } - for (int k = k_shift; k < p.head_dim; k += kThreadsPerKey) { gk_ptr[k] = scalar_t(0); } - } - } - - template - static CUTLASS_DEVICE void processBlockIJ(SharedStorage& shared_storage, - OutputFragments& output_frags, - Params& p, - int32_t query_start, - int32_t key_start, - uint8_t warp_id, - uint8_t lane_id) - { - cutlass::MatrixCoord no_offset{0, 0}; - accum_t scale = p.scale; - int16_t thread_id = 32 * warp_id + lane_id; - auto rematerializeThreadIds = [&]() { - // Prevents `nvcc` from keeping values deduced from - // `thread_id`, `warp_id`, ... in RF - to reduce register pressure - warp_id = warp_uniform(thread_id / 32); - lane_id = thread_id % 32; - thread_id = 32 * warp_id + lane_id; - }; - - bool isFirstQuery = (query_start == getQueryStart(p, key_start)); - int32_t next_query, next_key; - incrIteration(p, query_start, key_start, next_query, next_key); - bool isLastQuery = next_key != key_start; - __syncthreads(); - loadDi(shared_storage.di(), p, query_start); - - int32_t num_queries_in_block = - skipBoundsChecks ? MatmulQK::Mma::Shape::kN - : warp_uniform(cutlass::fast_min((int32_t)MatmulQK::Mma::Shape::kN, - p.num_queries - query_start)); - int32_t num_keys_in_block = - skipBoundsChecks ? MatmulQK::Mma::Shape::kM - : warp_uniform(cutlass::fast_min((int32_t)MatmulQK::Mma::Shape::kM, - p.num_keys - key_start)); - - auto prologueGradV = [&](int col) { - typename MatmulGradV::Mma::IteratorB iterator_dO( - {int32_t(p.gO_strideM)}, - p.grad_output_ptr + query_start * p.gO_strideM + col, - {num_queries_in_block, p.head_dim_value - col}, - thread_id, - no_offset); - MatmulGradV::Mma::prologue( - shared_storage.mm_gradV(), iterator_dO, thread_id, num_queries_in_block); - }; - auto prologueGradQ = [&](int col) { - typename MatmulGradQ::Mma::IteratorB iterator_K( - {int32_t(p.k_strideM)}, - p.key_ptr + key_start * p.k_strideM + col, - {num_keys_in_block, p.head_dim - col}, - thread_id, - no_offset); - MatmulGradQ::Mma::prologue( - shared_storage.mm_gradQ(), iterator_K, thread_id, num_keys_in_block); - }; - auto prologueGradK = [&](int col) { - typename MatmulGradK::Mma::IteratorB iterator_Q( - {int32_t(p.q_strideM)}, - p.query_ptr + query_start * p.q_strideM + col, - {num_queries_in_block, p.head_dim - col}, - thread_id, - no_offset); - MatmulGradK::Mma::prologue( - shared_storage.mm_gradK(), iterator_Q, thread_id, num_queries_in_block); - }; - auto prologueDOV = [&]() { - typename MatmulDOIVJ::Mma::IteratorA iterator_A( - {int32_t(p.gO_strideM)}, - p.grad_output_ptr + query_start * p.gO_strideM, - {num_queries_in_block, p.head_dim_value}, - thread_id, - no_offset); - typename MatmulDOIVJ::Mma::IteratorB iterator_B({int32_t(p.v_strideM)}, - p.value_ptr + key_start * p.v_strideM, - {p.head_dim_value, num_keys_in_block}, - thread_id, - no_offset); - MatmulDOIVJ::Mma::prologue( - shared_storage.mm_doivj(), iterator_A, iterator_B, thread_id, p.head_dim_value); - }; - - ///////////////////////////////////////////////////////////////////////////////////////////////// - // MatmulQK - ///////////////////////////////////////////////////////////////////////////////////////////////// - { - using Mma = typename MatmulQK::Mma; - - cutlass::gemm::GemmCoord problem_size(num_keys_in_block, - num_queries_in_block, - p.head_dim // k - ); - - // k_j - typename Mma::IteratorA iterator_A({int32_t(p.k_strideM)}, - p.key_ptr + key_start * p.k_strideM, - {problem_size.m(), problem_size.k()}, - thread_id, - no_offset); - - // q_i.transpose(-2, -1) - typename Mma::IteratorB iterator_B({int32_t(p.q_strideM)}, - p.query_ptr + query_start * p.q_strideM, - {problem_size.k(), problem_size.n()}, - thread_id, - no_offset); - - Mma mma( - shared_storage.mm_qk_k(), shared_storage.mm_qk_q(), thread_id, warp_id, lane_id); - - typename Mma::FragmentC accum; - - accum.clear(); - - auto gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; - - // Compute threadblock-scoped matrix multiply-add - mma.set_prologue_done(kPrologueQK); - mma.set_zero_outside_bounds(!skipBoundsChecks); - mma(gemm_k_iterations, accum, iterator_A, iterator_B, accum); - - // Epilogue: add LSE + exp and store that to our shared memory buffer - // shmem <- (matmul_result - - // logsumexp[i_start:i_end].unsqueeze(1)).exp() - int warp_idx_mn_0 = warp_id % (Mma::Base::WarpCount::kM * Mma::Base::WarpCount::kN); - auto output_tile_coords = cutlass::MatrixCoord{ - warp_idx_mn_0 % Mma::Base::WarpCount::kM, warp_idx_mn_0 / Mma::Base::WarpCount::kM}; - - if (broadcast_1::kEnable || broadcast_2::kEnable) { - cutlass::TensorRef bias_tensor_ref( - shared_storage.bias().data(), - cutlass::layout::RowMajor(MatmulQK::ThreadblockShape::kM)); - using Shape = cutlass::MatrixShape; - AttentionBiasEpilogue - bias_epilogue; - bias_epilogue(bias_tensor_ref, - p.bias1_ptr + key_start, - p.bias2_ptr + query_start * p.num_keys + key_start, - thread_id, - {num_queries_in_block, num_keys_in_block}, - p.num_keys); - // Pij += Bij, Pij is in register fragment and Bij is in shared memory - auto lane_offset = MatmulQK::AccumLambdaIterator::get_lane_offset( - lane_id, warp_id, output_tile_coords); - MatmulQK::AccumLambdaIterator::iterateRows( - lane_offset, - [&](int accum_n) {}, - [&](int accum_m, int accum_n, int idx) { - // remember we are transposed - accum[idx] = accum[idx] * scale + bias_tensor_ref.at({accum_n, accum_m}); - }, - [&](int accum_n) {}); - } else { - accum = cutlass::multiplies()(scale, accum); - } - - __syncthreads(); - if (kPrologueGV) { prologueGradV(0); } - if (kPrologueDOV) { prologueDOV(); } - - MatmulQK::B2bGemm::accumApplyLSEToSmem(shared_storage.attn_shared_storage(), - accum, - p.logsumexp_ptr + query_start, - problem_size.n(), - thread_id, - warp_id, - lane_id, - output_tile_coords); - - __syncthreads(); - } - rematerializeThreadIds(); - - ///////////////////////////////////////////////////////////////////////////////////////////////// - // GradV matmul - // - // grad_v[j_start:j_end] += attn_T @ do_i - ///////////////////////////////////////////////////////////////////////////////////////////////// - constexpr bool kSingleIterationGradV = kMaxK <= MatmulGradV::ThreadblockShape::kN; - for (int col = 0; col < (kSingleIterationGradV ? 1 : p.head_dim_value); - col += MatmulGradV::ThreadblockShape::kN) { - using Mma = typename MatmulGradV::Mma; - using AccumTileGmem = typename MatmulGradQ::AccumTileGmem; - - cutlass::gemm::GemmCoord problem_size( - num_keys_in_block, p.head_dim_value - col, num_queries_in_block); - auto createEpilogueIter = [&]() { - return typename MatmulGradV::OutputTileIterator( - typename MatmulGradV::OutputTileIterator::Params{p.gV_strideM()}, - p.grad_value_ptr + key_start * p.gV_strideM() + col, - {num_keys_in_block, p.head_dim_value - col}, - thread_id); - }; - typename Mma::IteratorB iterator_B({int32_t(p.gO_strideM)}, - p.grad_output_ptr + query_start * p.gO_strideM + col, - {num_queries_in_block, p.head_dim_value - col}, - thread_id, - no_offset); - - // if dropout: dVj += (Pij.T * Zij) @ dOi - // otherwise: dVj += Pij.T @ dOi - Mma mma(shared_storage.mm_gradV(), - // operand A: Pij - typename MatmulGradV::WarpIteratorA( - shared_storage.attn_shared_storage().accum_ref(), lane_id), - // if we're using dropout, operand A is Pij_dropped = Pij * Zij - // which is computed on the fly as fragments of Pij are loaded in - typename Mma::WarpIteratorAScale(shared_storage.zij().accum_ref(), lane_id), - thread_id, - warp_id, - lane_id); - - int storage_id = col / MatmulGradV::ThreadblockShape::kN; - AccumTileGmem gmem_tile{p.workspace_gv + storage_id * AccumTileGmem::kElementsStored}; - if (!kOutputInRF) { - if (isFirstQuery || !kNeedsAccumGradV) { - output_frags.gradV.clear(); - } else { - gmem_tile.load(output_frags.gradV, thread_id); - } - } - mma.set_prologue_done(kPrologueGV); - - auto gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; - - // Compute threadblock-scoped matrix multiply-add - __syncthreads(); - - mma(gemm_k_iterations, output_frags.gradV, iterator_B, output_frags.gradV); - __syncthreads(); - if (kPrologueGV && !kSingleIterationGradV && - col + MatmulGradV::ThreadblockShape::kN < p.head_dim_value) { - prologueGradV(col + MatmulGradV::ThreadblockShape::kN); - } - - if (!kOutputInRF) { - if (kNeedsAccumGradV && !isLastQuery) { - gmem_tile.store(output_frags.gradV, thread_id); - } else { - accumulateInGmem(shared_storage.gradV_epilogue(), - output_frags.gradV, - createEpilogueIter(), - isFirstQuery || kNeedsAccumGradV, - warp_id, - lane_id); - } - } - } - __syncthreads(); - ///////////////////////////////////////////////////////////////////////////////////////////////// - // MatmulDOIVJ - ///////////////////////////////////////////////////////////////////////////////////////////////// - { - using Mma = typename MatmulDOIVJ::Mma; - // do_i - typename Mma::IteratorA iterator_A({int32_t(p.gO_strideM)}, - p.grad_output_ptr + query_start * p.gO_strideM, - {num_queries_in_block, p.head_dim_value}, - thread_id, - no_offset); - - // v_j.transpose(-2, -1) - typename Mma::IteratorB iterator_B({int32_t(p.v_strideM)}, - p.value_ptr + key_start * p.v_strideM, - {p.head_dim_value, num_keys_in_block}, - thread_id, - no_offset); - - Mma mma(shared_storage.mm_doivj(), thread_id, warp_id, lane_id); - mma.set_prologue_done(kPrologueDOV); - mma.set_zero_outside_bounds(!skipBoundsChecks); - - typename Mma::FragmentC accum; - - accum.clear(); - - auto gemm_k_iterations = (p.head_dim_value + Mma::Shape::kK - 1) / Mma::Shape::kK; - - // Compute threadblock-scoped matrix multiply-add - mma(gemm_k_iterations, accum, iterator_A, iterator_B, accum); - __syncthreads(); - if (kPrologueGQ) { prologueGradQ(0); } - if (kPrologueGK) { prologueGradK(0); } - - int warp_idx_mn_0 = warp_id % (Mma::Base::WarpCount::kM * Mma::Base::WarpCount::kN); - auto output_tile_coords = cutlass::MatrixCoord{ - warp_idx_mn_0 % Mma::Base::WarpCount::kM, warp_idx_mn_0 / Mma::Base::WarpCount::kM}; - // TODO: This must be terribly inefficient. There must be a better way - // tmp [RF] <- (accum [RF] - Di [smem] ) * attn_T.T [smem] - // attn_shared_storage [smem] <- tmp.T - // tmp_shared_storage [smem] <- tmp - { - using LambdaIterator = - typename DefaultMmaAccumLambdaIterator::Iterator; - auto lane_offset = - LambdaIterator::get_lane_offset(lane_id, warp_id, output_tile_coords); - - auto attn_T = shared_storage.attn_shared_storage().accum_ref(); - accum_t current_di; - // dSij = (dPij - Di) * Pij - LambdaIterator::iterateRows( - lane_offset, - [&](int accum_m) { current_di = shared_storage.di()[accum_m]; }, - [&](int accum_m, int accum_n, int idx) { - if (skipBoundsChecks || - (accum_m < num_queries_in_block && accum_n < num_keys_in_block)) { - accum_t attn = attn_T.at({accum_n, accum_m}); - accum[idx] = (accum[idx] - current_di) * attn; - } else { - accum[idx] = 0; - } - }, - [&](int accum_m) { - - }); - - using DefaultGemm = typename MatmulDOIVJ::DefaultGemm; - using OutputOp = typename MatmulDOIVJ::BiasGradEpilogueOutputOp; - if (broadcast_1::kEnable && p.grad_bias1_ptr) { - using Epilogue = - typename BiasGradEpilogueAffineRankN::Epilogue; - cutlass::layout::AffineRankN<2> layout({0, 1}); - auto dst_ptr = p.grad_bias1_ptr + key_start; - typename Epilogue::OutputTileIterator output_iter( - {layout}, - dst_ptr, - {num_queries_in_block, num_keys_in_block}, - (int)thread_id); - Epilogue epilogue(shared_storage.gradB_epilogue(), - (int)thread_id, - (int)warp_id, - (int)lane_id); - epilogue(OutputOp(1), output_iter, accum); - } - - if (broadcast_2::kEnable && p.grad_bias2_ptr) { - if (broadcast_1::kEnable) { __syncthreads(); } - using Epilogue = - typename BiasGradEpilogue::Epilogue; - typename Epilogue::OutputTileIterator::Params params{p.num_keys}; - auto dst_ptr = p.grad_bias2_ptr + query_start * p.num_keys + key_start; - typename Epilogue::OutputTileIterator output_iter( - params, dst_ptr, {num_queries_in_block, num_keys_in_block}, (int)thread_id); - Epilogue epilogue(shared_storage.gradB_epilogue(), - (int)thread_id, - (int)warp_id, - (int)lane_id); - epilogue(OutputOp(1), output_iter, accum); - } - - accum = accum * scale; - - __syncthreads(); - if (!MatmulGradK::DefaultMmaFromSmem::kIsTransposedA) { - auto tmpT = shared_storage.tmpT_shared_storage().accum_ref(); - // attn <- attn_T.T - LambdaIterator::iterateRows( - lane_offset, - [&](int accum_m) {}, - [&](int accum_m, int accum_n, int idx) { - tmpT.at({accum_n, accum_m}) = scalar_t(accum[idx]); - }, - [&](int accum_m) {}); - } - } - - MatmulDOIVJ::B2bGemm::accumToSmem( - shared_storage.tmp_shared_storage(), accum, lane_id, output_tile_coords); - __syncthreads(); - } - p.head_dim = warp_uniform(p.head_dim); - p.k_strideM = warp_uniform(p.k_strideM); - rematerializeThreadIds(); - ///////////////////////////////////////////////////////////////////////////////////////////////// - // GradQ matmul - // - // grad_q[i_start:i_end] += tmp @ k_j - ///////////////////////////////////////////////////////////////////////////////////////////////// - // Skip the loop & associated branches if we know at compile time the number - // of iterations - constexpr bool kSingleIterationGradQ = kMaxK <= MatmulGradQ::ThreadblockShape::kN; - for (int col = 0; col < (kSingleIterationGradQ ? 1 : p.head_dim); - col += MatmulGradQ::ThreadblockShape::kN) { - using Mma = typename MatmulGradQ::Mma; - using AccumTileGmem = typename MatmulGradQ::AccumTileGmem; - - cutlass::gemm::GemmCoord problem_size( - num_queries_in_block, - false ? MatmulGradQ::ThreadblockShape::kN : p.head_dim - col, - num_keys_in_block); - - // k_j - typename Mma::IteratorB iterator_B({int32_t(p.k_strideM)}, - p.key_ptr + key_start * p.k_strideM + col, - {problem_size.k(), problem_size.n()}, - thread_id, - no_offset); - - auto a = shared_storage.tmp_shared_storage().accum_ref(); - Mma mma(shared_storage.mm_gradQ(), - shared_storage.tmp_shared_storage(), - thread_id, - warp_id, - lane_id, - problem_size.k()); - - typename Mma::FragmentC accum; - - bool isFirst = key_start == 0; - int col_id = col / MatmulGradQ::ThreadblockShape::kN; - int num_cols = - kSingleIterationGradQ ? 1 : ceil_div(p.head_dim, MatmulGradQ::ThreadblockShape::kN); - int storage_id = (col_id + query_start / kBlockSizeI * num_cols); - AccumTileGmem gmem_tile{p.workspace_gq + storage_id * AccumTileGmem::kElementsStored}; - if (isFirst || !kNeedsAccumGradQ) { - accum.clear(); - } else { - gmem_tile.load(accum, thread_id); - } - - auto gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; - - // Compute threadblock-scoped matrix multiply-add - __syncthreads(); - mma.set_prologue_done(kPrologueGQ); - mma(gemm_k_iterations, accum, iterator_B, accum); - __syncthreads(); - bool isLastColumn = kSingleIterationGradQ || - (col + MatmulGradQ::ThreadblockShape::kN >= p.head_dim); - if (kPrologueGQ && !isLastColumn) { - prologueGradQ(col + MatmulGradQ::ThreadblockShape::kN); - } - - // Output results - int32_t next_query, next_key; - incrIteration(p, p.num_queries, key_start, next_query, next_key); - bool isLast = next_query > query_start || next_key >= p.num_keys; - if (kNeedsAccumGradQ && !isLast) { - gmem_tile.store(accum, thread_id); - } else { - typename MatmulGradQ::OutputTileIterator output_it( - typename MatmulGradQ::OutputTileIterator::Params{p.gQ_strideM()}, - p.grad_query_ptr + query_start * p.gQ_strideM() + col, - {problem_size.m(), problem_size.n()}, - thread_id); - accumulateInGmem(isLastColumn - ? shared_storage.gradQ_epilogue_lastIter() - : shared_storage.gradQ_epilogue(), - accum, - output_it, - isFirst || kNeedsAccumGradQ, - warp_id, - lane_id); - } - } - ///////////////////////////////////////////////////////////////////////////////////////////////// - // GradK matmul - // - // grad_k[i_start:i_end] += tmp.transpose(-2, -1) @ q_i - ///////////////////////////////////////////////////////////////////////////////////////////////// - rematerializeThreadIds(); - - constexpr bool kSingleIterationGradK = kMaxK <= MatmulGradK::ThreadblockShape::kN; - for (int col = 0; col < (kSingleIterationGradK ? 1 : p.head_dim); - col += MatmulGradK::ThreadblockShape::kN) { - using Mma = typename MatmulGradK::Mma; - using AccumTileGmem = typename MatmulGradQ::AccumTileGmem; - - cutlass::gemm::GemmCoord problem_size( - num_keys_in_block, - false ? MatmulGradK::ThreadblockShape::kN : p.head_dim - col, - num_queries_in_block); - auto createEpilogueIter = [&]() { - return typename MatmulGradK::OutputTileIterator( - typename MatmulGradK::OutputTileIterator::Params{p.gK_strideM()}, - p.grad_key_ptr + key_start * p.gK_strideM() + col, - {num_keys_in_block, - false ? MatmulGradK::ThreadblockShape::kN : p.head_dim - col}, - thread_id); - }; - - // q_i - typename Mma::IteratorB iterator_B({int32_t(p.q_strideM)}, - p.query_ptr + query_start * p.q_strideM + col, - {problem_size.k(), problem_size.n()}, - thread_id, - no_offset); - - auto getTmp = [&](int) { return &shared_storage.tmp_shared_storage(); }; - auto getTmpT = [&](int) { return &shared_storage.tmpT_shared_storage(); }; - // this is basically: - // opA = kIsTransposedA ? getTmp() : getTmpT(); - bool constexpr kIsTransposedA = MatmulGradK::DefaultMmaFromSmem::kIsTransposedA; - auto& opA = - *call_conditional::apply( - getTmp, getTmpT, 0); - Mma mma(shared_storage.mm_gradK(), opA, thread_id, warp_id, lane_id, problem_size.k()); - - int storage_id = col / MatmulGradK::ThreadblockShape::kN; - AccumTileGmem gmem_tile{p.workspace_gk + storage_id * AccumTileGmem::kElementsStored}; - if (!kOutputInRF) { - if (isFirstQuery || !kNeedsAccumGradK) { - output_frags.gradK.clear(); - } else { - gmem_tile.load(output_frags.gradK, thread_id); - } - } - mma.set_prologue_done(kPrologueGK); - - auto gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; - - // Compute threadblock-scoped matrix multiply-add - __syncthreads(); - - mma(gemm_k_iterations, output_frags.gradK, iterator_B, output_frags.gradK); - __syncthreads(); - bool isLastColumn = kSingleIterationGradK || - col + MatmulGradK::ThreadblockShape::kN >= p.head_dim; - if (kPrologueGK && !isLastColumn) { - prologueGradK(col + MatmulGradK::ThreadblockShape::kN); - } - - if (kPrologueQK && isLastColumn) { - int32_t next_query, next_key; - incrIteration(p, query_start, key_start, next_query, next_key); - DISPATCH_BOOL(next_key != key_start, kForceReloadK, ([&]() { - prologueQkNextIteration( - shared_storage, p, next_query, next_key, warp_id, lane_id); - })); - } - - // Output results - if (!kOutputInRF) { - if (kNeedsAccumGradK && !isLastQuery) { - gmem_tile.store(output_frags.gradK, thread_id); - } else { - accumulateInGmem(isLastColumn - ? shared_storage.gradK_epilogue_final() - : shared_storage.gradK_epilogue(), - output_frags.gradK, - createEpilogueIter(), - isFirstQuery || kNeedsAccumGradK, - warp_id, - lane_id); - __syncthreads(); - } - } - } - } - - static CUTLASS_DEVICE int32_t getQueryStart(Params const& p, int32_t key_start) { return 0; }; - - static CUTLASS_DEVICE void incrIteration(Params const& p, - int32_t query_start, - int32_t key_start, - int32_t& next_query, - int32_t& next_key) - { - next_query = query_start + kBlockSizeI; - next_key = key_start; - if (next_query >= p.num_queries) { - next_key = key_start + kBlockSizeJ; - next_query = getQueryStart(p, next_key); - } - } - - template - static CUTLASS_DEVICE void prologueQkNextIteration(SharedStorage& shared_storage, - Params const& p, - int32_t query_start, - int32_t key_start, - uint8_t warp_id, - uint8_t lane_id) - { - if (query_start >= p.num_queries || key_start >= p.num_keys) { return; } - - static constexpr bool kReloadK = kForceReloadK || !MatmulQK::Mma::kSmemContainsEntireMat; - int thread_id = 32 * warp_id + lane_id; - typename MatmulQK::Mma::IteratorA iterator_A({int32_t(p.k_strideM)}, - p.key_ptr + key_start * p.k_strideM, - {p.num_keys - key_start, p.head_dim}, - thread_id, - cutlass::MatrixCoord{0, 0}); - - typename MatmulQK::Mma::IteratorB iterator_B({int32_t(p.q_strideM)}, - p.query_ptr + query_start * p.q_strideM, - {p.head_dim, p.num_queries - query_start}, - thread_id, - cutlass::MatrixCoord{0, 0}); - - MatmulQK::Mma::prologue(shared_storage.mm_qk_k(), - shared_storage.mm_qk_q(), - iterator_A, - iterator_B, - thread_id, - p.head_dim); - } - - template - static CUTLASS_DEVICE void writeFragsToGmem(SharedStorage& shared_storage, - OutputFragments& output_frags, - Params const& p, - int32_t key_start, - uint8_t warp_id, - uint8_t lane_id) - { - uint16_t thread_id = 32 * warp_id + lane_id; - int32_t num_keys_in_block = - skipBoundsChecks - ? MatmulQK::Mma::Shape::kM - : cutlass::fast_min((int32_t)MatmulQK::Mma::Shape::kM, p.num_keys - key_start); - typename MatmulGradV::OutputTileIterator outputV_it( - typename MatmulGradV::OutputTileIterator::Params{p.gV_strideM()}, - p.grad_value_ptr + key_start * p.gV_strideM(), - {num_keys_in_block, p.head_dim_value}, - thread_id); - accumulateInGmem(shared_storage.gradV_epilogue_final(), - output_frags.gradV, - outputV_it, - true, - warp_id, - lane_id); - - typename MatmulGradK::OutputTileIterator outputK_it( - typename MatmulGradK::OutputTileIterator::Params{p.gK_strideM()}, - p.grad_key_ptr + key_start * p.gK_strideM(), - {num_keys_in_block, false ? MatmulGradK::ThreadblockShape::kN : p.head_dim}, - thread_id); - accumulateInGmem(shared_storage.gradK_epilogue_final(), - output_frags.gradK, - outputK_it, - true, - warp_id, - lane_id); - } - - template - static CUTLASS_DEVICE void accumulateInGmem( - typename MatmulT::DefaultEpilogue::SharedStorage& epilogue_smem, - typename MatmulT::Mma::FragmentC const& accum, - typename MatmulT::OutputTileIterator output_it, - bool first, - uint8_t warp_id, - uint8_t lane_id) - { - using DefaultEpilogue = typename MatmulT::DefaultEpilogue; - using DefaultOutputOp = typename MatmulT::DefaultOutputOp; - using Mma = typename MatmulT::Mma; - int thread_id = 32 * warp_id + lane_id; - DISPATCH_BOOL( - first, kIsFirst, ([&]() { - static constexpr auto ScaleType = - kIsFirst ? cutlass::epilogue::thread::ScaleType::Nothing - : cutlass::epilogue::thread::ScaleType::NoBetaScaling; - using EpilogueOutputOp = typename cutlass::epilogue::thread::LinearCombination< - typename DefaultOutputOp::ElementOutput, - DefaultOutputOp::kCount, - typename DefaultOutputOp::ElementAccumulator, - typename DefaultOutputOp::ElementCompute, - ScaleType>; - using Epilogue = typename cutlass::epilogue::threadblock::EpiloguePipelined< - typename DefaultEpilogue::Shape, - typename Mma::Operator, - DefaultEpilogue::kPartitionsK, - typename MatmulT::OutputTileIterator, - typename DefaultEpilogue::AccumulatorFragmentIterator, - typename DefaultEpilogue::WarpTileIterator, - typename DefaultEpilogue::SharedLoadIterator, - EpilogueOutputOp, - typename DefaultEpilogue::Padding, - DefaultEpilogue::kFragmentsPerIteration, - true // IterationsUnroll - >; - EpilogueOutputOp rescale({1, 1}); - Epilogue epilogue(epilogue_smem, thread_id, warp_id, lane_id); - epilogue(rescale, output_it, accum, output_it); - })); - } - - template - static CUTLASS_DEVICE void computeDelta(Params const& p, - int32_t query_start, - uint8_t warp_id, - uint8_t lane_id) - { - // Each thread computes one value for Delta - // Depending on warp configuration, we might have multiple - // threads of the same warp working on the same row - using AccessType = cutlass::Array; - static_assert(kNumThreads >= kBlockSizeI, ""); - static constexpr int kNumThreadsPerLine = kNumThreads / kBlockSizeI; - int16_t thread_id = 32 * warp_id + lane_id; - - int16_t laneFirstCol = kElementsPerAccess * (lane_id % kNumThreadsPerLine); - int16_t laneRow = thread_id / kNumThreadsPerLine; - bool rowPred = (query_start + laneRow) < p.num_queries; - bool pred = rowPred; - - // on windows, previous syntax __restrict__ AccessType* - // resulted in error: "restrict" is not allowed - const AccessType* __restrict__ grad_output_ptr = reinterpret_cast( - p.grad_output_ptr + (query_start + laneRow) * p.gO_strideM + laneFirstCol); - const AccessType* __restrict__ output_ptr = reinterpret_cast( - p.output_ptr + (query_start + laneRow) * p.o_strideM() + laneFirstCol); - - static constexpr int64_t kMaxIters = kMaxK / (kElementsPerAccess * kNumThreadsPerLine); - constexpr int kPipelineStages = 2; - accum_t delta_value = accum_t(0); - using GlobalLoad = cutlass::arch::global_load; - AccessType frag_grad_output[kPipelineStages]; - AccessType frag_output[kPipelineStages]; - - auto loadAndIncrement = [&](int ld_pos, bool is_valid) { - frag_grad_output[ld_pos].clear(); - frag_output[ld_pos].clear(); - GlobalLoad(frag_grad_output[ld_pos], grad_output_ptr, is_valid); - GlobalLoad(frag_output[ld_pos], output_ptr, is_valid); - grad_output_ptr += kNumThreadsPerLine; - output_ptr += kNumThreadsPerLine; - }; - - CUTLASS_PRAGMA_UNROLL - for (int iter = 0; iter < kPipelineStages - 1; ++iter) { - int ld_pos = iter % kPipelineStages; - pred = pred && (laneFirstCol + iter * kElementsPerAccess * kNumThreadsPerLine) < - p.head_dim_value; - loadAndIncrement(ld_pos, pred); - } - auto columnIteration = [&](int iter) { - // Load for next iter - int ld_pos = (iter + kPipelineStages - 1) % kPipelineStages; - pred = pred && (laneFirstCol + (iter + kPipelineStages - 1) * kElementsPerAccess * - kNumThreadsPerLine) < p.head_dim_value; - loadAndIncrement(ld_pos, pred); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < AccessType::kElements; ++i) { - delta_value += accum_t(frag_output[iter % kPipelineStages][i]) * - accum_t(frag_grad_output[iter % kPipelineStages][i]); - } - }; - - // If we have a small lower-bound for K, we can unroll the loop - if (kMaxK <= 256) { - CUTLASS_PRAGMA_UNROLL - for (int iter = 0; iter < kMaxIters; ++iter) { columnIteration(iter); } - } else { - int num_iters = ceil_div(p.head_dim_value, kElementsPerAccess * kNumThreadsPerLine) * - (kElementsPerAccess * kNumThreadsPerLine); - for (int iter = 0; iter < num_iters; ++iter) { columnIteration(iter); } - } - - // Reduce between workers - static_assert(kNumThreadsPerLine == 1 || kNumThreadsPerLine == 2 || kNumThreadsPerLine == 4, - ""); - CUTLASS_PRAGMA_UNROLL - for (int i = 1; i < kNumThreadsPerLine; i *= 2) { - delta_value = delta_value + __shfl_xor_sync(0xffffffff, delta_value, i); - } - - // Store in gmem - if (rowPred) { p.delta_ptr[query_start + laneRow] = delta_value; } - } -}; - -template -__global__ void __launch_bounds__(AK::kNumThreads, AK::kMinBlocksPerSm) - attention_kernel_backward_batched_impl(typename AK::Params p) -{ - if (!p.advance_to_block()) { return; } - AK::attention_kernel(p); -} - -template -__global__ void __launch_bounds__(AK::kNumThreads, AK::kMinBlocksPerSm) - attention_kernel_backward_batched(typename AK::Params params); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_forward.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_forward.h deleted file mode 100644 index e3b11ebcc6615156179c5f9047acb0d65688be7b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/kernel_forward.h +++ /dev/null @@ -1,986 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holdvr nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include -#include - -#include "cutlass/bfloat16.h" -#include "cutlass/fast_math.h" -#include "cutlass/gemm/gemm.h" -#include "cutlass/layout/matrix.h" -#include "cutlass/layout/vector.h" -#include "cutlass/matrix.h" -#include "cutlass/numeric_types.h" -#include "cutlass/tensor_ref.h" - -#include "cutlass/epilogue/threadblock/default_epilogue_simt.h" -#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h" -#include "cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" -#include "cutlass/gemm/device/default_gemm_configuration.h" -#include "cutlass/gemm/kernel/default_gemm.h" -#include "cutlass/gemm/threadblock/default_mma.h" -#include "cutlass/gemm/threadblock/default_mma_core_simt.h" -#include "cutlass/gemm/threadblock/default_mma_core_sm70.h" -#include "cutlass/gemm/threadblock/default_mma_core_sm75.h" -#include "cutlass/gemm/threadblock/default_mma_core_sm80.h" -#include "cutlass/gemm/threadblock/threadblock_swizzle.h" -#include "cutlass/matrix_shape.h" -#include "cutlass/platform/platform.h" -#include "cutlass/transform/threadblock/predicated_tile_iterator.h" -#include "epilogue/epilogue_pipelined.h" -#include "epilogue/epilogue_rescale_output.h" -#include "gemm/find_default_mma.h" -#include "gemm/mma_from_smem.h" -#include "gemm_kernel_utils.h" -#include "transform/bias_broadcast.h" -#include "transform/tile_smem_loader.h" - -#include - -using namespace gemm_kernel_utils; - -namespace { -template -constexpr int getWarpsPerSm() -{ - return (Arch::kMinComputeCapability >= 80 && !cutlass::platform::is_same::value - ? 16 - : 12); -} -static CUTLASS_DEVICE float atomicMaxFloat(float* addr, float value) -{ - // source: https://stackoverflow.com/a/51549250 - return (value >= 0) ? __int_as_float(atomicMax((int*)addr, __float_as_int(value))) - : __uint_as_float(atomicMin((unsigned int*)addr, __float_as_uint(value))); -} -} // namespace - -template < - // The datatype of Q/K/V - typename scalar_t_, - // Architecture we are targeting (eg `cutlass::arch::Sm80`) - typename ArchTag, - // If Q/K/V are correctly aligned in memory and we can run a fast kernel - bool isAligned_, - int kQueriesPerBlock, - int kKeysPerBlock_, - bool kSingleValueIteration_, // = `value.shape[-1] <= kKeysPerBlock` - // This is quite slower on V100 for some reason - // Set to false if you know at compile-time you will never need dropout - bool kSupportsBias_ = false, - template class Broadcast1_ = BroadcastNoLoad, - template class Broadcast2_ = BroadcastNoLoad> -struct AttentionKernel { - using scalar_t = scalar_t_; - using accum_t = float; - using lse_scalar_t = float; - using output_t = scalar_t; - // Accumulator between 2 iterations - // Using `accum_t` improves perf on f16 at the cost of - // numerical errors - using output_accum_t = accum_t; - static constexpr bool kSupportsBias = kSupportsBias_; - static constexpr int kKeysPerBlock = kKeysPerBlock_; - static constexpr bool kIsAligned = isAligned_; - static constexpr bool kSingleValueIteration = kSingleValueIteration_; - static constexpr int32_t kAlignLSE = 32; // block size of backward - static constexpr bool kPreloadV = - ArchTag::kMinComputeCapability >= 80 && cutlass::sizeof_bits::value == 16; - static constexpr bool kKeepOutputInRF = kSingleValueIteration; - static constexpr bool kNeedsOutputAccumulatorBuffer = - !kKeepOutputInRF && !cutlass::platform::is_same::value; - - static_assert(kQueriesPerBlock % 32 == 0, ""); - static_assert(kKeysPerBlock % 32 == 0, ""); - static constexpr int kNumWarpsPerBlock = kQueriesPerBlock * kKeysPerBlock / (32 * 32); - static constexpr int kWarpSize = 32; - - // Launch bounds - static constexpr int kNumThreads = kWarpSize * kNumWarpsPerBlock; - static constexpr int kMinBlocksPerSm = getWarpsPerSm() / kNumWarpsPerBlock; - - struct Params { - // Input tensors - scalar_t* query_ptr; // [num_queries, num_heads, head_dim] - scalar_t* key_ptr; // [num_keys, num_heads, head_dim] - scalar_t* value_ptr; // [num_keys, num_heads, head_dim_value] - - // Output tensors - output_t* output_ptr; // [num_queries, num_heads, head_dim_value] - output_accum_t* output_accum_ptr; // [num_queries, num_heads, head_dim_value] - lse_scalar_t* logsumexp_ptr; // [num_heads, num_queries] - can be null - - // Scale - accum_t scale; - - // Dimensions/strides - int32_t head_dim; - int32_t head_dim_value; - int32_t num_queries; - int32_t num_keys; - - int32_t q_strideM; - int32_t k_strideM; - int32_t v_strideM; - // int32_t bias_strideM = 0; - - int32_t o_strideM = 0; - - // Everything below is only used in `advance_to_block` - // and shouldn't use registers - int32_t q_strideH; - int32_t k_strideH; - int32_t v_strideH; - // int32_t bias_strideH = 0; - - int64_t q_strideB; - int64_t k_strideB; - int64_t v_strideB; - // int32_t bias_strideB = 0; - - int32_t num_batches; - int32_t num_heads; - - // Parameters for biases - scalar_t* bias1_ptr = nullptr; - scalar_t* bias2_ptr = nullptr; - int32_t B = 0; - int32_t N = 0; - - // Moves pointers to what we should process - // Returns "false" if there is no work to do - CUTLASS_DEVICE bool advance_to_block() - { - auto batch_id = blockIdx.z; - auto head_id = blockIdx.y; - auto query_start = blockIdx.x * kQueriesPerBlock; - - auto lse_dim = ceil_div((int32_t)num_queries, kAlignLSE) * kAlignLSE; - - query_ptr += batch_id * q_strideB; - key_ptr += batch_id * k_strideB; - value_ptr += batch_id * v_strideB; - output_ptr += int64_t(batch_id * num_queries) * o_strideM; - if (output_accum_ptr != nullptr) { - output_accum_ptr += int64_t(batch_id * num_queries) * (head_dim_value * num_heads); - } - - int64_t q_start = 0, k_start = 0; - // Advance to the current batch / head / query_start - query_ptr += (q_start + query_start) * q_strideM + head_id * q_strideH; - key_ptr += k_start * k_strideM + head_id * k_strideH; - - value_ptr += k_start * v_strideM + head_id * v_strideH; - output_ptr += int64_t(q_start + query_start) * o_strideM + head_id * head_dim_value; - - if (output_accum_ptr != nullptr) { - output_accum_ptr += int64_t(q_start + query_start) * (head_dim_value * num_heads) + - head_id * head_dim_value; - } else { - // Accumulate directly in the destination buffer (eg for f32) - output_accum_ptr = (accum_t*)output_ptr; - } - - if (logsumexp_ptr != nullptr) { - // lse[batch_id, head_id, query_start] - logsumexp_ptr += batch_id * lse_dim * num_heads + head_id * lse_dim + query_start; - } - - using broadcast_1 = Broadcast1_; - if (kSupportsBias && broadcast_1::kEnable && bias1_ptr) { - bias1_ptr = broadcast_1::advance(bias1_ptr, - batch_id / N, - batch_id % N, - head_id, - num_queries * N, - num_queries, - 0); - } - using broadcast_2 = Broadcast2_; - if (kSupportsBias && broadcast_2::kEnable && bias2_ptr) { - auto strideB = num_heads * num_queries * num_keys; - auto strideH = num_queries * num_keys; - bias2_ptr = broadcast_2::advance( - bias2_ptr, batch_id / N, batch_id % N, head_id, strideB, 0, strideH); - } - - num_queries -= query_start; - num_batches = 0; // no longer used after - - // If num_queries == 1, and there is only one key head we're wasting - // 15/16th of tensor core compute In that case : - // - we only launch kernels for head_id % kQueriesPerBlock == 0 - // - we iterate over heads instead of queries (strideM = strideH) - if (num_queries == 1 && k_strideH == 0 && v_strideH == 0) { - if (head_id % kQueriesPerBlock != 0) return false; - q_strideM = q_strideH; - num_queries = num_heads; - num_heads = 1; // unused but here for intent - o_strideM = head_dim_value; - } - - // Make sure the compiler knows these variables are the same on all - // the threads of the warp. - query_ptr = warp_uniform(query_ptr); - key_ptr = warp_uniform(key_ptr); - value_ptr = warp_uniform(value_ptr); - output_ptr = warp_uniform(output_ptr); - output_accum_ptr = warp_uniform(output_accum_ptr); - logsumexp_ptr = warp_uniform(logsumexp_ptr); - num_queries = warp_uniform(num_queries); - num_keys = warp_uniform(num_keys); - num_heads = warp_uniform(num_heads); - head_dim = warp_uniform(head_dim); - head_dim_value = warp_uniform(head_dim_value); - o_strideM = warp_uniform(o_strideM); - if (kSupportsBias && broadcast_1::kEnable) { bias1_ptr = warp_uniform(bias1_ptr); } - if (kSupportsBias && broadcast_2::kEnable) { bias2_ptr = warp_uniform(bias2_ptr); } - return true; - } - - __host__ dim3 getBlocksGrid() const - { - return dim3(ceil_div(num_queries, (int32_t)kQueriesPerBlock), num_heads, num_batches); - } - - __host__ dim3 getThreadsGrid() const { return dim3(kWarpSize, kNumWarpsPerBlock, 1); } - }; - - struct MM0 { - /* - In this first matmul, we compute a block of `Q @ K.T`. - While the calculation result is still hot in registers, we update - `mi`, `m_prime`, `s_prime` in shared-memory, and then store this value - into a shared-memory ("AccumulatorSharedStorage") that is used later as - operand A for the second matmul (see MM1) - */ - using GemmType = DefaultGemmType; - - using OpClass = typename GemmType::OpClass; - using DefaultConfig = - typename cutlass::gemm::device::DefaultGemmConfiguration; - static constexpr int kAlignmentA = kIsAligned ? DefaultConfig::kAlignmentA - : GemmType::kMinimumAlignment; - static constexpr int kAlignmentB = kIsAligned ? DefaultConfig::kAlignmentB - : GemmType::kMinimumAlignment; - using ThreadblockShape = - cutlass::gemm::GemmShape; - using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; - using DefaultMma = typename cutlass::gemm::threadblock::FindDefaultMma< - scalar_t, // ElementA, - cutlass::layout::RowMajor, // LayoutA, - kAlignmentA, - scalar_t, // ElementB, - cutlass::layout::ColumnMajor, // LayoutB, - kAlignmentB, - accum_t, - cutlass::layout::RowMajor, // LayoutC, - OpClass, - ArchTag, // ArchTag - ThreadblockShape, // ThreadblockShape - WarpShape, // WarpShape - typename GemmType::InstructionShape, // InstructionShape - DefaultConfig::kStages, // Should use `DefaultConfig::kStages`, but that - // uses too much smem - typename GemmType::Operator // Operator - >::DefaultMma; - using MmaCore = typename DefaultMma::MmaCore; - using IteratorA = typename DefaultMma::IteratorA; - using IteratorB = typename DefaultMma::IteratorB; - using Mma = typename DefaultMma::ThreadblockMma; - using AccumLambdaIterator = - typename DefaultMmaAccumLambdaIterator::Iterator; - static_assert(MmaCore::WarpCount::kM * MmaCore::WarpCount::kN * MmaCore::WarpCount::kK == - kNumWarpsPerBlock, - ""); - - // used for efficient load of bias tile Bij from global to shared memory - using BiasLoader = - TileSmemLoader, - MmaCore::kThreads, - // input restriction: kv_len has to be a multiple of this value - 128 / cutlass::sizeof_bits::value>; - - // Epilogue to store to shared-memory in a format that we can use later for - // the second matmul - using B2bGemm = - typename cutlass::gemm::threadblock::B2bGemm; - using AccumulatorSharedStorage = typename B2bGemm::AccumulatorSharedStorage; - }; - - struct MM1 { - /** - Second matmul: perform `attn @ V` where `attn` is the attention (not - normalized) and stored in shared memory - */ - using GemmType = DefaultGemmType; - - using OpClass = typename GemmType::OpClass; - using DefaultConfig = - typename cutlass::gemm::device::DefaultGemmConfiguration; - static constexpr int kAlignmentA = DefaultConfig::kAlignmentA; // from smem - static constexpr int kAlignmentB = kIsAligned ? DefaultConfig::kAlignmentB - : GemmType::kMinimumAlignment; - using ThreadblockShape = - cutlass::gemm::GemmShape; - using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; - using InstructionShape = typename GemmType::InstructionShape; - - using LayoutB = cutlass::layout::RowMajor; - using DefaultGemm = - cutlass::gemm::kernel::DefaultGemm; - - using DefaultMmaFromSmem = typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< - typename DefaultGemm::Mma, - typename MM0::AccumulatorSharedStorage, - false>; // kScaleOperandA - using Mma = typename DefaultMmaFromSmem::Mma; - using IteratorB = typename Mma::IteratorB; - using WarpCount = typename Mma::WarpCount; - static_assert(WarpCount::kM * WarpCount::kN * WarpCount::kK == kNumWarpsPerBlock, ""); - - using DefaultEpilogue = typename DefaultGemm::Epilogue; - using OutputTileIterator = typename cutlass::epilogue::threadblock::PredicatedTileIterator< - typename DefaultEpilogue::OutputTileIterator::ThreadMap, - output_t>; - using OutputTileIteratorAccum = - typename cutlass::epilogue::threadblock::PredicatedTileIterator< - typename DefaultEpilogue::OutputTileIterator::ThreadMap, - output_accum_t>; - - struct SharedStorageMM1 { - typename Mma::SharedStorage mm; - }; - }; - - static constexpr int64_t kAlignmentQ = MM0::kAlignmentA; - static constexpr int64_t kAlignmentK = MM0::kAlignmentB; - static constexpr int64_t kAlignmentV = 1; - - // Shared storage - depends on kernel params - struct ScalingCoefs { - cutlass::Array m_prime; - cutlass::Array s_prime; - cutlass::Array mi; - }; - - struct SharedStorageEpilogueAtEnd : ScalingCoefs { - struct SharedStorageAfterMM0 { - // Everything here might be overwritten during MM0 - union { - // typename MM0::BiasLoader::SmemTile bias; - cutlass::AlignedBuffer bias; - typename MM0::AccumulatorSharedStorage si; - }; - typename MM1::SharedStorageMM1 mm1; - }; - - union { - typename MM0::Mma::SharedStorage mm0; - SharedStorageAfterMM0 after_mm0; - typename MM1::DefaultEpilogue::SharedStorage epilogue; - }; - - CUTLASS_DEVICE typename MM1::DefaultEpilogue::SharedStorage& epilogue_shared_storage() - { - return epilogue; - } - }; - - struct SharedStorageEpilogueInLoop : ScalingCoefs { - struct SharedStorageAfterMM0 { - // Everything here might be overwritten during MM0 - union { - // typename MM0::BiasLoader::SmemTile bias; - cutlass::AlignedBuffer bias; - typename MM0::AccumulatorSharedStorage si; - }; - typename MM1::SharedStorageMM1 mm1; - typename MM1::DefaultEpilogue::SharedStorage epilogue; - }; - - union { - typename MM0::Mma::SharedStorage mm0; - SharedStorageAfterMM0 after_mm0; - }; - - CUTLASS_DEVICE typename MM1::DefaultEpilogue::SharedStorage& epilogue_shared_storage() - { - return after_mm0.epilogue; - } - }; - - using SharedStorage = - typename cutlass::platform::conditional::type; - - static bool __host__ check_supported(Params const& p) - { - CHECK_ALIGNED_PTR(p.query_ptr, kAlignmentQ); - CHECK_ALIGNED_PTR(p.key_ptr, kAlignmentK); - CHECK_ALIGNED_PTR(p.value_ptr, kAlignmentV); - EVOFORMER_CHECK(p.q_strideM % kAlignmentQ == 0, "query is not correctly aligned (strideM)"); - EVOFORMER_CHECK(p.k_strideM % kAlignmentK == 0, "key is not correctly aligned (strideM)"); - EVOFORMER_CHECK(p.v_strideM % kAlignmentV == 0, "value is not correctly aligned (strideM)"); - EVOFORMER_CHECK(p.num_heads <= 1 || p.q_strideH % kAlignmentQ == 0, - "query is not correctly aligned (strideH)"); - EVOFORMER_CHECK(p.num_heads <= 1 || p.k_strideH % kAlignmentK == 0, - "key is not correctly aligned (strideH)"); - EVOFORMER_CHECK(p.num_heads <= 1 || p.v_strideH % kAlignmentV == 0, - "value is not correctly aligned (strideH)"); - return true; - } - - static void CUTLASS_DEVICE attention_kernel(Params& p) - { - // In this block, we will only ever: - // - read query[query_start:query_end, :] - // - write to output[query_start:query_end, :] - - extern __shared__ char smem_buffer[]; - SharedStorage& shared_storage = *((SharedStorage*)smem_buffer); - auto& m_prime = shared_storage.m_prime; - auto& s_prime = shared_storage.s_prime; - auto& mi = shared_storage.mi; - const uint32_t query_start = blockIdx.x * kQueriesPerBlock; - - static_assert(kQueriesPerBlock < kNumWarpsPerBlock * kWarpSize, ""); - if (thread_id() < kQueriesPerBlock) { - s_prime[thread_id()] = accum_t(0); - m_prime[thread_id()] = -cutlass::platform::numeric_limits::infinity(); - mi[thread_id()] = -cutlass::platform::numeric_limits::infinity(); - } - typename MM1::Mma::FragmentC accum_o; - accum_o.clear(); - - auto createOutputIter = [&](int col) -> typename MM1::OutputTileIterator { - using OutputTileIterator = typename MM1::OutputTileIterator; - return OutputTileIterator( - typename OutputTileIterator::Params{(int32_t)p.o_strideM}, - p.output_ptr, - typename OutputTileIterator::TensorCoord{p.num_queries, p.head_dim_value}, - thread_id(), - {0, col}); - }; - - auto createOutputAccumIter = [&](int col) -> typename MM1::OutputTileIteratorAccum { - using OutputTileIteratorAccum = typename MM1::OutputTileIteratorAccum; - return OutputTileIteratorAccum( - typename OutputTileIteratorAccum::Params{(int32_t)(p.head_dim_value * p.num_heads)}, - p.output_accum_ptr, - typename OutputTileIteratorAccum::TensorCoord{p.num_queries, p.head_dim_value}, - thread_id(), - {0, col}); - }; - - // Iterate through keys - for (int32_t iter_key_start = 0; iter_key_start < p.num_keys; - iter_key_start += kKeysPerBlock) { - int32_t problem_size_0_m = cutlass::fast_min((int32_t)kQueriesPerBlock, p.num_queries); - int32_t problem_size_0_n = - cutlass::fast_min(int32_t(kKeysPerBlock), p.num_keys - iter_key_start); - int32_t const& problem_size_0_k = p.head_dim; - int32_t const& problem_size_1_n = p.head_dim_value; - int32_t const& problem_size_1_k = problem_size_0_n; - - auto prologueV = [&](int blockN) { - typename MM1::Mma::IteratorB iterator_V( - typename MM1::IteratorB::Params{MM1::LayoutB(p.v_strideM)}, - p.value_ptr + iter_key_start * p.v_strideM, - {problem_size_1_k, problem_size_1_n}, - thread_id(), - cutlass::MatrixCoord{0, blockN * MM1::Mma::Shape::kN}); - MM1::Mma::prologue( - shared_storage.after_mm0.mm1.mm, iterator_V, thread_id(), problem_size_1_k); - }; - - __syncthreads(); // Need to have shared memory initialized, and `m_prime` - // updated from end of prev iter - // - // MATMUL: Q.K_t - // - // Computes the block-matrix product of: - // (a) query[query_start:query_end, :] - // with - // (b) key[iter_key_start:iter_key_start + kKeysPerBlock] - // and stores that into `shared_storage.si` - // - - // Compute threadblock location - cutlass::gemm::GemmCoord tb_tile_offset = {0, 0, 0}; - - cutlass::MatrixCoord tb_offset_A{tb_tile_offset.m() * MM0::Mma::Shape::kM, - tb_tile_offset.k()}; - - cutlass::MatrixCoord tb_offset_B{tb_tile_offset.k(), - tb_tile_offset.n() * MM0::Mma::Shape::kN}; - - // Construct iterators to A and B operands - typename MM0::IteratorA iterator_A( - typename MM0::IteratorA::Params(typename MM0::MmaCore::LayoutA(p.q_strideM)), - p.query_ptr, - {problem_size_0_m, problem_size_0_k}, - thread_id(), - tb_offset_A); - - typename MM0::IteratorB iterator_B( - typename MM0::IteratorB::Params(typename MM0::MmaCore::LayoutB(p.k_strideM)), - p.key_ptr + iter_key_start * p.k_strideM, - {problem_size_0_k, problem_size_0_n}, - thread_id(), - tb_offset_B); - - auto my_warp_id = warp_id(); - auto my_lane_id = lane_id(); - - // Construct thread-scoped matrix multiply - typename MM0::Mma mma(shared_storage.mm0, thread_id(), my_warp_id, my_lane_id); - - typename MM0::Mma::FragmentC accum; - - accum.clear(); - - auto gemm_k_iterations = - (problem_size_0_k + MM0::Mma::Shape::kK - 1) / MM0::Mma::Shape::kK; - - // Compute threadblock-scoped matrix multiply-add - mma(gemm_k_iterations, accum, iterator_A, iterator_B, accum); - __syncthreads(); - - if (kPreloadV) { - prologueV(0); - } else { - MM1::Mma::drain_cp_asyncs(); - } - - typename MM0::Mma::Operator::IteratorC::TensorCoord iteratorC_tile_offset = { - (tb_tile_offset.m() * MM0::Mma::WarpCount::kM) + - (my_warp_id % MM0::Mma::WarpCount::kM), - (tb_tile_offset.n() * MM0::Mma::WarpCount::kN) + - (my_warp_id / MM0::Mma::WarpCount::kM)}; - - // multiply by scaling factor - // if (kSupportsBias) { - // accum = - // cutlass::multiplies()(p.scale, - // accum); - // } - - if (kSupportsBias) { - cutlass::TensorRef bias_tensor_ref( - shared_storage.after_mm0.bias.data(), - cutlass::layout::RowMajor(MM0::ThreadblockShape::kN)); - using Shape = - cutlass::MatrixShape; - AttentionBiasEpilogue - bias_epilogue; - bias_epilogue(bias_tensor_ref, - p.bias1_ptr + iter_key_start, - p.bias2_ptr + query_start * p.num_keys + iter_key_start, - thread_id(), - {problem_size_0_m, problem_size_0_n}, - p.num_keys); - // Pij += Bij, Pij is in register fragment and Bij is in shared memory - auto lane_offset = MM0::AccumLambdaIterator::get_lane_offset( - lane_id(), warp_id(), iteratorC_tile_offset); - MM0::AccumLambdaIterator::iterateRows( - lane_offset, - [&](int accum_m) {}, - [&](int accum_m, int accum_n, int idx) { - if (accum_m < problem_size_0_m && accum_n < problem_size_0_n) { - accum[idx] = - accum[idx] * p.scale + bias_tensor_ref.at({accum_m, accum_n}); - } - }, - [&](int accum_m) {}); - } - - DISPATCH_BOOL(iter_key_start == 0, kIsFirst, ([&] { - DISPATCH_BOOL( - p.num_keys - iter_key_start >= kKeysPerBlock, kFullColumns, ([&] { - // Update `mi` from accum stored in registers - // Also does accum[i] <- exp(accum[i] - mi) - iterative_softmax(accum_o, - accum, - mi, - m_prime, - s_prime, - lane_id(), - thread_id(), - warp_id(), - p.num_keys - iter_key_start, - iteratorC_tile_offset, - kSupportsBias ? 1.0f : p.scale); - })); - })); - - // Output results to shared-memory - int warp_idx_mn_0 = - my_warp_id % (MM0::Mma::Base::WarpCount::kM * MM0::Mma::Base::WarpCount::kN); - auto output_tile_coords = - cutlass::MatrixCoord{warp_idx_mn_0 % MM0::Mma::Base::WarpCount::kM, - warp_idx_mn_0 / MM0::Mma::Base::WarpCount::kM}; - - MM0::B2bGemm::accumToSmem( - shared_storage.after_mm0.si, accum, my_lane_id, output_tile_coords); - - __syncthreads(); - - // - // MATMUL: Attn . V - // Run the matmul `attn @ V` for a block of attn and V. - // `attn` is read from shared memory (in `shared_storage_si`) - // `V` is read from global memory (with iterator_B) - // - - const int64_t nBlockN = - kSingleValueIteration - ? 1 - : ceil_div((int64_t)problem_size_1_n, int64_t(MM1::ThreadblockShape::kN)); - for (int blockN = 0; blockN < nBlockN; ++blockN) { - int gemm_k_iterations = - (problem_size_1_k + MM1::Mma::Shape::kK - 1) / MM1::Mma::Shape::kK; - - // Compute threadblock-scoped matrix multiply-add and store it in accum - // (in registers) - if (!kPreloadV) { - __syncthreads(); // we share shmem between mma and epilogue - } - - typename MM1::Mma::IteratorB iterator_V( - typename MM1::IteratorB::Params{MM1::LayoutB(p.v_strideM)}, - p.value_ptr + iter_key_start * p.v_strideM, - {problem_size_1_k, problem_size_1_n}, - thread_id(), - cutlass::MatrixCoord{0, blockN * MM1::Mma::Shape::kN}); - typename MM1::Mma mma_pv(shared_storage.after_mm0.mm1.mm, - shared_storage.after_mm0.si, - (int)thread_id(), - (int)warp_id(), - (int)lane_id(), - (int)problem_size_1_k); - mma_pv.set_prologue_done(kPreloadV); - if (!kKeepOutputInRF) { accum_o.clear(); } - mma_pv(gemm_k_iterations, accum_o, iterator_V, accum_o); - __syncthreads(); - - if (kPreloadV && !kSingleValueIteration && blockN + 1 < nBlockN) { - prologueV(blockN + 1); - } - - if (!kKeepOutputInRF) { - MM1::Mma::drain_cp_asyncs(); - DISPATCH_BOOL( - iter_key_start == 0, kIsFirst, ([&] { - DISPATCH_BOOL( - (iter_key_start + kKeysPerBlock) >= p.num_keys, kIsLast, ([&] { - using DefaultEpilogue = typename MM1::DefaultEpilogue; - using DefaultOp = typename MM1::DefaultConfig::EpilogueOutputOp; - using ElementCompute = typename DefaultOp::ElementCompute; - using EpilogueOutputOp = typename cutlass::epilogue::thread:: - MemoryEfficientAttentionNormalize< - typename cutlass::platform:: - conditional:: - type, - output_accum_t, - DefaultOp::kCount, - typename DefaultOp::ElementAccumulator, - ElementCompute, - kIsFirst, - kIsLast, - cutlass::Array>; - using Epilogue = - typename cutlass::epilogue::threadblock::EpiloguePipelined< - typename DefaultEpilogue::Shape, - typename MM1::Mma::Operator, - DefaultEpilogue::kPartitionsK, - typename cutlass::platform::conditional< - kIsLast, - typename MM1::OutputTileIterator, - typename MM1::OutputTileIteratorAccum>::type, - typename DefaultEpilogue::AccumulatorFragmentIterator, - typename DefaultEpilogue::WarpTileIterator, - typename DefaultEpilogue::SharedLoadIterator, - EpilogueOutputOp, - typename DefaultEpilogue::Padding, - DefaultEpilogue::kFragmentsPerIteration, - true, // IterationsUnroll - typename MM1::OutputTileIteratorAccum // Read - // iterator - >; - - int col = blockN * MM1::Mma::Shape::kN; - auto source_iter = createOutputAccumIter(col); - auto dest_iter = - call_conditional:: - apply(createOutputIter, createOutputAccumIter, col); - EpilogueOutputOp rescale(s_prime, m_prime); - Epilogue epilogue(shared_storage.epilogue_shared_storage(), - thread_id(), - warp_id(), - lane_id()); - epilogue(rescale, dest_iter, accum_o, source_iter); - })); - })); - if (!kSingleValueIteration) { __syncthreads(); } - } - } - __syncthreads(); // we modify `m_prime` after - } - - if (kKeepOutputInRF) { - constexpr bool kIsFirst = true; - constexpr bool kIsLast = true; - using DefaultEpilogue = typename MM1::DefaultEpilogue; - using DefaultOp = typename MM1::DefaultConfig::EpilogueOutputOp; - using ElementCompute = typename DefaultOp::ElementCompute; - using EpilogueOutputOp = - typename cutlass::epilogue::thread::MemoryEfficientAttentionNormalize< - output_t, // output - output_accum_t, // source - DefaultOp::kCount, - typename DefaultOp::ElementAccumulator, // accum - output_accum_t, // compute - kIsFirst, - kIsLast, - cutlass::Array>; - using Epilogue = typename cutlass::epilogue::threadblock::EpiloguePipelined< - typename DefaultEpilogue::Shape, - typename MM1::Mma::Operator, - DefaultEpilogue::kPartitionsK, - typename MM1::OutputTileIterator, // destination - typename DefaultEpilogue::AccumulatorFragmentIterator, - typename DefaultEpilogue::WarpTileIterator, - typename DefaultEpilogue::SharedLoadIterator, - EpilogueOutputOp, - typename DefaultEpilogue::Padding, - DefaultEpilogue::kFragmentsPerIteration, - true, // IterationsUnroll - typename MM1::OutputTileIteratorAccum // source tile - >; - auto dest_iter = createOutputIter(0); - EpilogueOutputOp rescale(s_prime, m_prime); - Epilogue epilogue( - shared_storage.epilogue_shared_storage(), thread_id(), warp_id(), lane_id()); - MM1::Mma::drain_cp_asyncs(); - epilogue(rescale, dest_iter, accum_o); - } - - // 7. Calculate logsumexp - // To make the backward easier, we pad logsumexp with `inf` - // this avoids a few bound checks, and is not more expensive during fwd - static_assert(kQueriesPerBlock < kNumWarpsPerBlock * kWarpSize, ""); - if (p.logsumexp_ptr && thread_id() < kQueriesPerBlock) { - auto lse_dim = ceil_div((int32_t)p.num_queries, kAlignLSE) * kAlignLSE; - if (thread_id() < p.num_queries) { - p.logsumexp_ptr[thread_id()] = - accum_t(mi[thread_id()]) + cutlass::fast_log(accum_t(s_prime[thread_id()])); - } else if (thread_id() < lse_dim) { - p.logsumexp_ptr[thread_id()] = - cutlass::platform::numeric_limits::infinity(); - } - } - } - - template - CUTLASS_DEVICE static void iterative_softmax( - typename WarpIteratorC::Fragment& frag_o, // output so far - typename WarpIteratorC::Fragment& frag, - cutlass::Array& mi, - cutlass::Array& m_prime, - cutlass::Array& s_prime, - int8_t lane_id, - int8_t thread_id, - int8_t warp_id, - int16_t max_col, - typename WarpIteratorC::TensorCoord const& tile_offset, - float scaling) - { - /* Iterates on the accumulator and corresponding position on result matrix - - (1) Update `mi[r]` to the max value of the row `r` - (2) In a second iteration do the following: - (a) accum <- exp(accum - mi) - (b) m_prime <- exp(m_prime - mi) - (c) s_prime <- s_prime * m_prime + sum(accum) - - All of this is done on registers, before we store all of this - on shared memory for the next matmul with Value. - */ - using Fragment = typename WarpIteratorC::Fragment; - using LambdaIterator = - typename DefaultMmaAccumLambdaIterator::Iterator; - // Convert to `accum_t` (rather than double) - constexpr float kLog2e = 1.4426950408889634074; // log_2(e) = M_LOG2E - if (!kIsFirst) { - if (thread_id < kQueriesPerBlock) { m_prime[thread_id] = mi[thread_id]; } - __syncthreads(); - } - - auto lane_offset = LambdaIterator::get_lane_offset(lane_id, warp_id, tile_offset); - - // First update `mi` to the max per-row - { - accum_t max; - LambdaIterator::iterateRows( - lane_offset, - [&](int accum_m) { max = -cutlass::platform::numeric_limits::infinity(); }, - [&](int accum_m, int accum_n, int idx) { - if (kFullColumns || accum_n < max_col) { - max = cutlass::fast_max(max, frag[idx]); - } - }, - [&](int accum_m) { - // Having 4x atomicMax seems faster than reduce within warp - // first... - atomicMaxFloat(&mi[accum_m], max * scaling); - }); - } - frag = cutlass::multiplies()(scaling * kLog2e, frag); - - // Make sure we all share the update values for `mi` - __syncthreads(); - - if (thread_id < kQueriesPerBlock) { - auto m_prime_exp = exp2f(kLog2e * (m_prime[thread_id] - mi[thread_id])); - m_prime[thread_id] = m_prime_exp; - s_prime[thread_id] *= m_prime_exp; - } - __syncthreads(); // Update output fragments - if (kKeepOutputInRF && !kIsFirst) { - accum_t mp; - LambdaIterator::iterateRows( - lane_offset, - [&](int accum_m) { mp = m_prime[accum_m]; }, - [&](int accum_m, int accum_n, int idx) { frag_o[idx] *= mp; }, - [&](int accum_m) {}); - __syncthreads(); - } - // Update accum_m, accum_n, ... - { - accum_t mi_row, total_row; - LambdaIterator::iterateRows( - lane_offset, - [&](int accum_m) { mi_row = kLog2e * mi[accum_m]; }, - [&](int accum_m, int accum_n, int idx) { - frag[idx] = (kFullColumns || accum_n < max_col) ? exp2f(frag[idx] - mi_row) - : accum_t(0.0); - }, - [&](int accum_m) {}); - LambdaIterator::iterateRows( - lane_offset, - [&](int accum_m) { total_row = 0.0; }, - [&](int accum_m, int accum_n, int idx) { total_row += frag[idx]; }, - [&](int accum_m) { - if (LambdaIterator::reduceSameRow( - lane_id, total_row, [](accum_t a, accum_t b) { return a + b; })) { - atomicAdd(&s_prime[accum_m], total_row); - } - }); - } - } - - static CUTLASS_DEVICE int8_t lane_id() { return threadIdx.x; } - static CUTLASS_DEVICE int8_t warp_id() { return threadIdx.y; } - static CUTLASS_DEVICE int16_t thread_id() { return threadIdx.x + threadIdx.y * blockDim.x; } -}; - -template -__global__ void __launch_bounds__(AK::kNumThreads, AK::kMinBlocksPerSm) - attention_kernel_batched_impl(typename AK::Params p) -{ - if (!p.advance_to_block()) { return; } - AK::attention_kernel(p); -} - -template -__global__ void __launch_bounds__(AK::kNumThreads, AK::kMinBlocksPerSm) - attention_kernel_batched(typename AK::Params params); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/bias_broadcast.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/bias_broadcast.h deleted file mode 100644 index 0f15a43574cfe3eba123e03b7835227275e762f4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/bias_broadcast.h +++ /dev/null @@ -1,148 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -// This does nothing. -template -struct BroadcastNoLoad { - using Fragment = - cutlass::Array; - static const bool kEnable = false; - CUTLASS_DEVICE static void load(Fragment& frag, - scalar_t* ptr, - int thread_id, - const cutlass::MatrixCoord& extent, - int stride) - { - } - CUTLASS_DEVICE static scalar_t* - advance(scalar_t* ptr, int B_id, int N_id, int H_id, int strideB, int strideN, int strideH) - { - return ptr; - } -}; - -// This is to load the bias matrix from the global memory with on-the-fly -// broadcast. The shape in global memory is [B, N, 1, 1, L]. Each time we load -// the last dimension as a L row vector, and we further broadcast the L vector -// to a tile of size [L, L] by repeating the L vector L times -template -struct BroadcastA : public BroadcastNoLoad { - using Base = BroadcastNoLoad; - static const bool kEnable = true; - using layout = cutlass::layout::AffineRank2RowMajor; - - using GmemTileIterator = cutlass::transform::threadblock:: - PredicatedTileIterator; - using Fragment = typename GmemTileIterator::Fragment; - - CUTLASS_DEVICE static void load(Fragment& frag, - scalar_t* ptr, - int thread_id, - const cutlass::MatrixCoord& extent, - int stride) - { - GmemTileIterator iter({layout(0, 1)}, ptr, extent, thread_id); - iter.load(frag); - } - - CUTLASS_DEVICE static scalar_t* - advance(scalar_t* ptr, int B_id, int N_id, int H_id, int strideB, int strideN, int strideH) - { - return ptr + B_id * strideB + N_id * strideN; - } -}; - -// This is to load the bias matrix from the global memory with on-the-fly -// broadcast. The shape in global memory is [B, 1, H, L, L]. Each time we load -// a [L, L] matrix. Different N use the same bias matrix when B and H are the -// same. -template -struct BroadcastB : public BroadcastNoLoad { - using Base = BroadcastNoLoad; - static const bool kEnable = true; - using layout = cutlass::layout::RowMajor; - - using GmemTileIterator = cutlass::transform::threadblock:: - PredicatedTileIterator; - using Fragment = typename GmemTileIterator::Fragment; - - CUTLASS_DEVICE static void load(Fragment& frag, - scalar_t* ptr, - int thread_id, - const cutlass::MatrixCoord& extent, - int stride) - { - GmemTileIterator iter({layout(stride)}, ptr, extent, thread_id); - iter.load(frag); - } - - CUTLASS_DEVICE static scalar_t* - advance(scalar_t* ptr, int B_id, int N_id, int H_id, int strideB, int strideN, int strideH) - { - return ptr + B_id * strideB + H_id * strideH; - } -}; - -template - class Broadcast1_, - template - class Broadcast2_> -struct AttentionBiasEpilogue { - using ThreadMap = cutlass::transform::PitchLinearStripminedThreadMap< - cutlass::layout::PitchLinearShape, - kThreads, - 1>; - - using Broadcast1 = Broadcast1_; - using Broadcast2 = Broadcast2_; - - Broadcast1 broadcast1; - Broadcast2 broadcast2; - - using Ref = cutlass::TensorRef; - using SmemTileIterator = cutlass::transform::threadblock:: - RegularTileIterator; - - CUTLASS_DEVICE void operator()(const Ref& ref, - scalar_t* ptr1, - scalar_t* ptr2, - int thread_id, - const cutlass::MatrixCoord& extent, - int stride) - { - static_assert(Broadcast1::Fragment::kElements == Broadcast2::Fragment::kElements, - "The two broadcast fragments must have the same number of " - "elements"); - typename SmemTileIterator::Fragment frag; - frag.clear(); - float* frag_ptr = reinterpret_cast(&frag); - if (Broadcast1::kEnable) { - typename Broadcast1::Fragment frag1; - frag1.clear(); - broadcast1.load(frag1, ptr1, thread_id, extent, stride); - scalar_t* frag1_ptr = reinterpret_cast(&frag1); - for (int i = 0; i < Broadcast1::Fragment::kElements; ++i) { - frag_ptr[i] += static_cast(frag1_ptr[i]); - } - } - if (Broadcast2::kEnable) { - typename Broadcast2::Fragment frag2; - frag2.clear(); - broadcast2.load(frag2, ptr2, thread_id, extent, stride); - scalar_t* frag2_ptr = reinterpret_cast(&frag2); - for (int i = 0; i < Broadcast2::Fragment::kElements; ++i) { - frag_ptr[i] += static_cast(frag2_ptr[i]); - } - } - SmemTileIterator iter(ref, thread_id); - iter.store(frag); - __syncthreads(); - } -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/tile_smem_loader.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/tile_smem_loader.h deleted file mode 100644 index 5f3f0ee507e53bb0ad7a0f8fc3fcd25982c177f7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/transform/tile_smem_loader.h +++ /dev/null @@ -1,93 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights - *reserved. SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, - *this list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holdvr nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE - *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE - *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR - *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF - *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS - *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN - *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) - *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE - *POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once -#include -#include "cutlass/aligned_buffer.h" -#include "cutlass/array.h" -#include "cutlass/coord.h" -#include "cutlass/layout/matrix.h" -#include "cutlass/layout/pitch_linear.h" -#include "cutlass/numeric_types.h" -#include "cutlass/platform/platform.h" -#include "cutlass/transform/pitch_linear_thread_map.h" -#include "cutlass/transform/threadblock/predicated_tile_iterator.h" -#include "cutlass/transform/threadblock/regular_tile_iterator.h" - -template // thread access width in elements -class TileSmemLoader { -public: - using Shape = ThreadblockTileShape; - using SmemTile = cutlass::AlignedBuffer; - - using ThreadMap = cutlass::transform::PitchLinearStripminedThreadMap< - cutlass::layout::PitchLinearShape, // strided - Threads, // Threads - ElementsPerAccess>; // ElementsPerAccess - - using GmemTileIterator = cutlass::transform::threadblock::PredicatedTileIterator< - ThreadblockTileShape, // Shape - scalar_t, // Element - cutlass::layout::RowMajor, // Layout - 0, // AdvanceRank - ThreadMap>; // ThreadMap - - using SmemTileIterator = - cutlass::transform::threadblock::RegularTileIterator; // ThreadMap - - using Fragment = typename GmemTileIterator::Fragment; - - /// load a tile from global memory into shared memory - CUTLASS_DEVICE - static void load(GmemTileIterator tile_load_iter, SmemTileIterator tile_store_iter) - { - Fragment tb_frag; - tb_frag.clear(); - tile_load_iter.load(tb_frag); - tile_store_iter.store(tb_frag); - - __syncthreads(); - } -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cpp deleted file mode 100644 index 1a887b50e1a3d06c7e6f554b6df4399de071c712..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cpp +++ /dev/null @@ -1,124 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "fp_quantize.h" - -#include -#include -#include - -#define DISPATCH_QUANTIZE(T_TYPE, C_TYPE, mantisa, exponent) \ - if (val.options().dtype() == torch::T_TYPE) { \ - launch_quantization((C_TYPE*)val.data_ptr(), \ - (uint8_t*)out.data_ptr(), \ - num_groups, \ - group_size, \ - at::cuda::getCurrentCUDAStream(), \ - q_range, \ - q_bits, \ - q_mantisa_bits, \ - stochastic_rounding); \ - } - -at::Tensor quantize(torch::Tensor& out, - torch::Tensor& val, - int group_size, - int stochastic_rounding, - int q_bits, - int q_mantisa_bits) -{ - int total_elems = at::numel(val); - float q_range = q_bits == 8 ? (q_mantisa_bits == 3 ? 480.0 : 114688.0) : // fp8 ranges - (q_bits == 12 ? 510.0 : // fp12 range - (q_bits == 6 ? 28.0 : // fp6 range - 6.0)); // fp4 range (using power 2); TODO (Reza): add the power-4 - // in case accuracy is not matching! - int num_groups = total_elems / group_size; - - DISPATCH_QUANTIZE(kHalf, __half, 23, 8); -#ifdef BF16_AVAILABLE - DISPATCH_QUANTIZE(kBFloat16, __nv_bfloat16, 23, 8); -#endif - - return out; -} - -#define DISPATCH_DEQUANTIZE(T_TYPE, C_TYPE, mantisa) \ - if (val.options().dtype() == torch::T_TYPE) { \ - launch_dequantization((uint8_t*)val_q.data_ptr(), \ - (C_TYPE*)val.data_ptr(), \ - num_groups, \ - group_size, \ - q_mantisa_bits, \ - q_exponent_bits, \ - at::cuda::getCurrentCUDAStream()); \ - return; \ - } - -void dequantize(torch::Tensor& val, - torch::Tensor& val_q, - int group_size, - int q_mantisa_bits, - int q_exponent_bits) -{ - int total_elems = at::numel(val); - - int num_groups = total_elems / group_size; - - DISPATCH_DEQUANTIZE(kHalf, __half, 10); -#ifdef BF16_AVAILABLE - DISPATCH_DEQUANTIZE(kBFloat16, __nv_bfloat16, 7); -#endif -} - -#define DISPATCH_DEQUANTIZE_INDEX(T_TYPE, C_TYPE, mantisa) \ - if (val.options().dtype() == torch::T_TYPE) { \ - launch_selective_dequantization((uint8_t*)val_q.data_ptr(), \ - (C_TYPE*)val.data_ptr(), \ - (int32_t*)indexes.data_ptr(), \ - num_groups, \ - group_size, \ - num_indexes, \ - q_mantisa_bits, \ - q_exponent_bits, \ - at::cuda::getCurrentCUDAStream()); \ - return; \ - } -void selective_dequantize(torch::Tensor& val, - torch::Tensor& val_q, - torch::Tensor& indexes, - int group_size, - int q_mantisa_bits, - int q_exponent_bits) -{ - int total_elems = at::numel(val); - int num_indexes = indexes.size(0); - int num_groups = total_elems / group_size; - - DISPATCH_DEQUANTIZE_INDEX(kHalf, __half, 10); -#ifdef BF16_AVAILABLE - DISPATCH_DEQUANTIZE_INDEX(kBFloat16, __nv_bfloat16, 7); -#endif -} - -at::Tensor get_scales(torch::Tensor& out, int num_groups) -{ - auto options = at::TensorOptions() - .dtype(torch::kFloat) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - auto scales = - torch::from_blob(out.data_ptr(), {num_groups, 1}, {out.stride(0) / 4, 1}, options); - return scales; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("quantize", &quantize, "quantize function"); - m.def("dequantize", &dequantize, "dequantize function"); - m.def("get_scales", &get_scales, "get scales function"); - m.def("selective_dequantize", &selective_dequantize, "selective dequantize function"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cu deleted file mode 100644 index 66ea7392e0115ab9ebbeb66baad01f9d6fa80750..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cu +++ /dev/null @@ -1,532 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include "context.h" -#include "fp_quantize.h" -#include "memory_access_utils.h" -#include "reduction_utils.h" - -#include -#include - -#include -#include - -#ifdef BF16_AVAILABLE -#include -#endif -#include - -using ROp = reduce::ROpType; - -namespace quantization { - -constexpr int access_granularity = 16; -constexpr int quanitzed_access_granularity = 4; -constexpr int quanitzed_access_granularity_6bits = 2; -constexpr int threads = 256; -constexpr int warps = threads / 32; - -} // namespace quantization - -template -__device__ void round(uint32_t& mantisa, uint32_t& dst_exponent, curandStatePhilox4_32_10_t* state) -{ - constexpr uint32_t mantisa_mask = (1 << (_mantisa_bits - q_mantisa_bits)) - 1; - uint32_t offset = stochastic_rounding ? (curand_poisson(state, 10) & mantisa_mask) - : 1 << (_mantisa_bits - q_mantisa_bits - 1); - mantisa += offset; - dst_exponent += (((mantisa & ~mantisa_mask) == (1 << _mantisa_bits)) ? 1 : 0); -} - -template -__device__ void clip(uint32_t& exponent, uint32_t& mantisa) -{ - constexpr uint32_t max_exponent = (1 << (q_exponent_bits - 1)) + (1 << (_exponent_bits - 1)); - constexpr uint32_t min_exponent = - (1 << (_exponent_bits - 1)) - ((1 << (q_exponent_bits - 1)) - 1); - if (exponent > max_exponent) { - exponent = max_exponent; - mantisa = (((uint32_t)-1) >> (32 - q_mantisa_bits)) << 1; //.11 .. 10 - } - if (exponent < min_exponent) { - exponent = min_exponent; - mantisa = 0; - } -} - -template -__global__ void apply_quantization(T* val, - uint8_t* q_val, - int group_size, - std::pair seed, - float q_range) -{ - int tidx = threadIdx.x; - int wid = tidx >> 5; - int lane = tidx & 0x1f; - int gid = blockIdx.x * quantization::warps + wid; - - constexpr int q_exponent_bits = total_q_bits - q_mantisa_bits - 1; - constexpr uint32_t _mantisa_mask = (1 << _mantisa_bits) - 1; - constexpr uint32_t _exponent_mask = ((1 << _exponent_bits) - 1) << _mantisa_bits; - constexpr uint32_t _sign_mask = 1 << (_mantisa_bits + _exponent_bits); - // CG helpers - cg::thread_block tb = cg::this_thread_block(); - cg::thread_block_tile warp = cg::tiled_partition(tb); - - constexpr uint32_t vector_size = quantization::access_granularity / sizeof(T); - constexpr uint32_t load_stride = vector_size * hw_warp_size; - constexpr uint32_t store_stride = (total_q_bits * vector_size / 8) * hw_warp_size; - const uint32_t thread_offset = lane * vector_size; - const uint32_t store_thread_offset = lane * (total_q_bits * vector_size / 8); - const uint32_t base_load_offset = gid * group_size + thread_offset; - const uint32_t base_store_offset = - gid * ((group_size * total_q_bits / 8) + 4) + - store_thread_offset; // 4-byte for saving the scale per group - const T* load_base_ptr = val + base_load_offset; - T tmp_buf[unroll * vector_size]; - T cur_max; - reduce::init(&cur_max); - - int idx = blockIdx.x * blockDim.x + threadIdx.x; - curandStatePhilox4_32_10_t state; - curand_init(seed.first, idx, seed.second, &state); - -#pragma unroll - for (int i = 0; i < unroll; i++) { - if (i * load_stride + thread_offset < group_size) { - mem_access::load_global( - &tmp_buf[vector_size * i], load_base_ptr + i * load_stride); - for (int j = 0; j < vector_size; j++) - cur_max = reduce::element(cur_max, __habs(tmp_buf[i * vector_size + j])); - } - } - reduce::_block(tb, warp, &cur_max); - - int mantisa_mask = ((1 << q_mantisa_bits) - 1); - mantisa_mask <<= (_mantisa_bits - q_mantisa_bits); - - uint8_t* store_base_ptr = q_val + base_store_offset; - float scale = (float)q_range / conversion::to(cur_max); -#pragma unroll - for (int i = 0; i < unroll; i++) { - if (i * load_stride + thread_offset < group_size) { - uint64_t q_buf = 0; - uint64_t q_buf1 = 0; -#pragma unroll - for (int j = 0; j < vector_size; j++) { - float val_f = conversion::to(tmp_buf[i * vector_size + j]) * scale; - uint32_t* data = reinterpret_cast(&val_f); - uint32_t sign = (data[0] & _sign_mask) >> (_mantisa_bits + _exponent_bits); - uint32_t cur_exponent = (data[0] & _exponent_mask) >> _mantisa_bits; - uint32_t dst_mantisa = (data[0] & _mantisa_mask); - - uint32_t dst_exponent = cur_exponent; - - round<_mantisa_bits, q_mantisa_bits, stochastic_rounding>( - dst_mantisa, dst_exponent, &state); - if (cur_exponent != 0) - clip<_mantisa_bits, _exponent_bits, q_mantisa_bits, q_exponent_bits>( - dst_exponent, dst_mantisa); - - dst_mantisa = (dst_mantisa & mantisa_mask) >> (_mantisa_bits - q_mantisa_bits); - - if (dst_exponent != (1 << q_exponent_bits) - 1) - dst_exponent = (dst_exponent - ((1 << (_exponent_bits - 1)) - 1)) + - (1 << (q_exponent_bits - 1)) - 1; - if (total_q_bits == 8 || total_q_bits == 4 || total_q_bits == 6) - q_buf = q_buf | - ((uint64_t)((uint8_t)(sign << (q_exponent_bits + q_mantisa_bits) | - (dst_exponent << q_mantisa_bits) | dst_mantisa)) - << j * total_q_bits); - else if (total_q_bits == 12) { - if (j < 5) - q_buf = - q_buf | - ((uint64_t)((uint16_t)(sign << (q_exponent_bits + q_mantisa_bits) | - (dst_exponent << q_mantisa_bits) | dst_mantisa)) - << j * total_q_bits); - else - q_buf1 = - q_buf1 | - ((uint64_t)((uint16_t)(sign << (q_exponent_bits + q_mantisa_bits) | - (dst_exponent << q_mantisa_bits) | dst_mantisa)) - << (j - 5) * total_q_bits); - } - } - if (total_q_bits == 12) { - uint64_t last_nibble_mask = 0xf; - last_nibble_mask = q_buf1 & last_nibble_mask; - q_buf = (last_nibble_mask << 60) | q_buf; - q_buf1 >>= 4; - } - uint8_t* int8_data = reinterpret_cast(&q_buf); - uint8_t* int8_data1 = reinterpret_cast(&q_buf1); - if (total_q_bits == 6) { - mem_access::store_global( - store_base_ptr + i * store_stride, int8_data); - mem_access::store_global( - store_base_ptr + i * store_stride + - quantization::quanitzed_access_granularity_6bits, - int8_data + quantization::quanitzed_access_granularity_6bits); - mem_access::store_global( - store_base_ptr + i * store_stride + - quantization::quanitzed_access_granularity_6bits * 2, - int8_data + 2 * quantization::quanitzed_access_granularity_6bits); - } else { - mem_access::store_global( - store_base_ptr + i * store_stride, int8_data); - - if (total_q_bits > 4) { - mem_access::store_global( - store_base_ptr + i * store_stride + - quantization::quanitzed_access_granularity, - int8_data + quantization::quanitzed_access_granularity); - if (total_q_bits == 12) { - mem_access::store_global( - store_base_ptr + i * store_stride + - quantization::quanitzed_access_granularity * 2, - int8_data1); - } - } - } - } - } - if (lane == 0) { - float q_scale = conversion::to(cur_max) / (float)q_range; - uint8_t* scale_as_int8 = reinterpret_cast(&q_scale); - uint32_t scale_offset = - gid * ((group_size * total_q_bits / 8) + 4) + (group_size * total_q_bits / 8); - if (total_q_bits != 6) - mem_access::store_global( - q_val + scale_offset, scale_as_int8); - else { - mem_access::store_global( - q_val + scale_offset, scale_as_int8); - mem_access::store_global( - q_val + scale_offset + quantization::quanitzed_access_granularity_6bits, - scale_as_int8 + quantization::quanitzed_access_granularity_6bits); - } - } -} - -template -__global__ void apply_dequantization(uint8_t* val, T* q_val, int group_size, int total_num_elements) -{ - constexpr uint32_t vector_size = quantization::access_granularity / sizeof(T); - int tidx = (blockIdx.x * blockDim.x + threadIdx.x) * vector_size; - - constexpr int quantized_bits = _mantisa_bits + _exponent_bits + 1; - constexpr int q_exponent_bits = total_q_bits - q_mantisa_bits - 1; - constexpr uint16_t _mantisa_mask = (1 << _mantisa_bits) - 1; - constexpr uint16_t _exponent_mask = ((1 << _exponent_bits) - 1) << _mantisa_bits; - constexpr uint16_t _sign_mask = 1 << (_mantisa_bits + _exponent_bits); - const uint32_t g_index = (tidx / group_size); - const uint32_t group_size_bytes = (group_size * quantized_bits / 8); - const uint8_t* load_base_ptr = - val + g_index * (group_size_bytes + 4) + (tidx % group_size) * quantized_bits / 8; - - int mantisa_mask = ((1 << q_mantisa_bits) - 1); - mantisa_mask <<= (_mantisa_bits - q_mantisa_bits); - - T* store_base_ptr = q_val + tidx; - float scale; - - uint8_t* scale_as_int8 = reinterpret_cast(&scale); - if (quantized_bits == 6) { - mem_access::load_global( - scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes); - mem_access::load_global( - scale_as_int8 + quantization::quanitzed_access_granularity_6bits, - val + g_index * (group_size_bytes + 4) + group_size_bytes + - quantization::quanitzed_access_granularity_6bits); - } else - mem_access::load_global( - scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes); - - if (tidx < total_num_elements) { - uint64_t q_buf_in; - uint64_t q_buf_in1; - uint8_t* int8_data = reinterpret_cast(&q_buf_in); - uint8_t* int8_data1 = reinterpret_cast(&q_buf_in1); - if (quantized_bits == 6) { - mem_access::load_global( - int8_data, load_base_ptr); - mem_access::load_global( - int8_data + quantization::quanitzed_access_granularity_6bits, - load_base_ptr + quantization::quanitzed_access_granularity_6bits); - mem_access::load_global( - int8_data + quantization::quanitzed_access_granularity_6bits * 2, - load_base_ptr + quantization::quanitzed_access_granularity_6bits * 2); - - } else { - mem_access::load_global(int8_data, - load_base_ptr); - if (quantized_bits > 4) { - mem_access::load_global( - int8_data + quantization::quanitzed_access_granularity, - load_base_ptr + quantization::quanitzed_access_granularity); - if (quantized_bits == 12) { - mem_access::load_global( - int8_data1, load_base_ptr + quantization::quanitzed_access_granularity * 2); - } - } - } - T store_buf[vector_size]; - uint16_t* q_buf = reinterpret_cast(store_buf); -#pragma unroll - for (int j = 0; j < vector_size; j++) { - uint16_t new_data; - if (j < 5 || quantized_bits != 12) { - new_data = (uint16_t)(q_buf_in >> (j * quantized_bits)); - } else { - if (j == 5) { - new_data = (uint16_t)(q_buf_in1); - new_data = (uint16_t)((new_data << 4) | (q_buf_in >> 60)); - } else - new_data = (uint16_t)(q_buf_in1 >> ((j - 6) * quantized_bits + 8)); - } - - uint16_t sign = (new_data & _sign_mask) >> (_mantisa_bits + _exponent_bits); - uint16_t dst_exponent = (new_data & _exponent_mask) >> _mantisa_bits; - uint16_t dst_mantisa = (new_data & _mantisa_mask); - - if (dst_exponent != (1 << q_exponent_bits) - 1) - dst_exponent = (dst_exponent - ((1 << (_exponent_bits - 1)) - 1)) + - (1 << (q_exponent_bits - 1)) - 1; - - q_buf[j] = - ((sign << (q_exponent_bits + q_mantisa_bits)) | (dst_exponent << q_mantisa_bits) | - (dst_mantisa << (q_mantisa_bits - _mantisa_bits))); - float up_cast = conversion::to(store_buf[j]); - store_buf[j] = conversion::to(up_cast * scale); - } - mem_access::store_global(store_base_ptr, store_buf); - } -} - -#define LAUNCH_FOR_QUANTIZATION_UNROLL(COUNT) \ - case COUNT: \ - apply_quantization \ - <<>>(val, q_val, group_size, seed, q_range); \ - break; - -template -void launch_quantization(T* val, - uint8_t* q_val, - int num_groups, - int group_size, - cudaStream_t stream, - float q_range, - int q_bits, - int q_mantisa_bits, - int stochastic_rounding) -{ - const dim3 grid((num_groups + quantization::warps - 1) / quantization::warps); - const dim3 block(quantization::threads); - - std::pair seed = FPContext::Instance().IncrementOffset(16); - - constexpr int vals_per_unroll = hw_warp_size * quantization::access_granularity / sizeof(T); - - const int copy_unroll = (group_size + vals_per_unroll - 1) / vals_per_unroll; - QUANT_SWITCH((q_bits - q_mantisa_bits - 1) * q_mantisa_bits + stochastic_rounding, [&] { - switch (copy_unroll) { - LAUNCH_FOR_QUANTIZATION_UNROLL(1) - LAUNCH_FOR_QUANTIZATION_UNROLL(2) - LAUNCH_FOR_QUANTIZATION_UNROLL(3) - LAUNCH_FOR_QUANTIZATION_UNROLL(4) - LAUNCH_FOR_QUANTIZATION_UNROLL(5) - LAUNCH_FOR_QUANTIZATION_UNROLL(6) - } - }); -} -#define INSTANTIATE_LAUNCH_QUANTIZATION(T, mantisa, exponent) \ - template void launch_quantization( \ - T*, uint8_t*, int, int, cudaStream_t, float q_range, int, int, int); -// fp8(E4M3), nearest-rounding -#ifdef BF16_AVAILABLE -INSTANTIATE_LAUNCH_QUANTIZATION(__nv_bfloat16, 23, 8); -#endif -INSTANTIATE_LAUNCH_QUANTIZATION(__half, 23, 8); - -template -void launch_dequantization(uint8_t* val, - T* q_val, - int num_groups, - int group_size, - int q_mantisa_bits, - int q_exponent_bits, - cudaStream_t stream) -{ - int blocks = ((num_groups * group_size) - 1) / - (quantization::threads * (quantization::access_granularity / sizeof(T))) + - 1; - const dim3 grid(blocks); - const dim3 block(quantization::threads); - DEQUANT_SWITCH(q_mantisa_bits * q_exponent_bits, [&] { - apply_dequantization - <<>>(val, q_val, group_size, (num_groups * group_size)); - }); -} -#define INSTANTIATE_LAUNCH_DEQUANTIZATION(T, mantisa) \ - template void launch_dequantization(uint8_t*, T*, int, int, int, int, cudaStream_t); -// fp8(E4M3) -#ifdef BF16_AVAILABLE -INSTANTIATE_LAUNCH_DEQUANTIZATION(__nv_bfloat16, 7); -#endif -INSTANTIATE_LAUNCH_DEQUANTIZATION(__half, 10); - -template -__global__ void apply_selective_dequantization(uint8_t* val, - T* q_val, - int32_t* indexes, - int group_size, - int total_num_elements) -{ - int index = indexes[blockIdx.x]; - constexpr uint32_t vector_size = quantization::access_granularity / sizeof(T); - int tidx = (blockIdx.y * blockDim.x + threadIdx.x) * vector_size; - int input_index = index * total_num_elements + tidx; - constexpr int quantized_bits = _mantisa_bits + _exponent_bits + 1; - constexpr int q_exponent_bits = total_q_bits - q_mantisa_bits - 1; - constexpr uint16_t _mantisa_mask = (1 << _mantisa_bits) - 1; - constexpr uint16_t _exponent_mask = ((1 << _exponent_bits) - 1) << _mantisa_bits; - constexpr uint16_t _sign_mask = 1 << (_mantisa_bits + _exponent_bits); - const uint32_t g_index = (input_index / group_size); - const uint32_t group_size_bytes = (group_size * quantized_bits / 8); - const uint8_t* load_base_ptr = - val + g_index * (group_size_bytes + 4) + (input_index % group_size) * quantized_bits / 8; - - int mantisa_mask = ((1 << q_mantisa_bits) - 1); - mantisa_mask <<= (_mantisa_bits - q_mantisa_bits); - - T* store_base_ptr = q_val + tidx + blockIdx.x * total_num_elements; - float scale; - - uint8_t* scale_as_int8 = reinterpret_cast(&scale); - if (quantized_bits == 6) { - mem_access::load_global( - scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes); - mem_access::load_global( - scale_as_int8 + quantization::quanitzed_access_granularity_6bits, - val + g_index * (group_size_bytes + 4) + group_size_bytes + - quantization::quanitzed_access_granularity_6bits); - } else - mem_access::load_global( - scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes); - - if (tidx < total_num_elements) { - uint64_t q_buf_in; - uint64_t q_buf_in1; - uint8_t* int8_data = reinterpret_cast(&q_buf_in); - uint8_t* int8_data1 = reinterpret_cast(&q_buf_in1); - if (quantized_bits == 6) { - mem_access::load_global( - int8_data, load_base_ptr); - mem_access::load_global( - int8_data + quantization::quanitzed_access_granularity_6bits, - load_base_ptr + quantization::quanitzed_access_granularity_6bits); - mem_access::load_global( - int8_data + quantization::quanitzed_access_granularity_6bits * 2, - load_base_ptr + quantization::quanitzed_access_granularity_6bits * 2); - } else { - mem_access::load_global(int8_data, - load_base_ptr); - if (quantized_bits > 4) { - mem_access::load_global( - int8_data + quantization::quanitzed_access_granularity, - load_base_ptr + quantization::quanitzed_access_granularity); - if (quantized_bits == 12) { - mem_access::load_global( - int8_data1, load_base_ptr + quantization::quanitzed_access_granularity * 2); - } - } - } - T store_buf[vector_size]; - uint16_t* q_buf = reinterpret_cast(store_buf); -#pragma unroll - for (int j = 0; j < vector_size; j++) { - uint16_t new_data; - if (j < 5 || quantized_bits != 12) { - new_data = (uint16_t)(q_buf_in >> (j * quantized_bits)); - } else { - if (j == 5) { - new_data = (uint16_t)(q_buf_in1); - new_data = (uint16_t)((new_data << 4) | (q_buf_in >> 60)); - } else - new_data = (uint16_t)(q_buf_in1 >> ((j - 6) * quantized_bits + 8)); - } - - uint16_t sign = (new_data & _sign_mask) >> (_mantisa_bits + _exponent_bits); - uint16_t dst_exponent = (new_data & _exponent_mask) >> _mantisa_bits; - uint16_t dst_mantisa = (new_data & _mantisa_mask); - - if (dst_exponent != (1 << q_exponent_bits) - 1) - dst_exponent = (dst_exponent - ((1 << (_exponent_bits - 1)) - 1)) + - (1 << (q_exponent_bits - 1)) - 1; - - q_buf[j] = - ((sign << (q_exponent_bits + q_mantisa_bits)) | (dst_exponent << q_mantisa_bits) | - (dst_mantisa << (q_mantisa_bits - _mantisa_bits))); - float up_cast = conversion::to(store_buf[j]); - store_buf[j] = conversion::to(up_cast * scale); - } - mem_access::store_global(store_base_ptr, store_buf); - } -} - -template -void launch_selective_dequantization(uint8_t* val, - T* q_val, - int32_t* indexes, - int num_groups, - int group_size, - int num_indexes, - int q_mantisa_bits, - int q_exponent_bits, - cudaStream_t stream) -{ - int total_elements_per_index = (num_groups / num_indexes) * group_size; - int blocks = (total_elements_per_index - 1) / - (quantization::threads * (quantization::access_granularity / sizeof(T))) + - 1; - const dim3 grid(num_indexes, blocks); - const dim3 block(quantization::threads); - DEQUANT_SWITCH(q_mantisa_bits * q_exponent_bits, [&] { - apply_selective_dequantization - <<>>(val, q_val, indexes, group_size, total_elements_per_index); - }); -} -#define INSTANTIATE_LAUNCH_SELECTIVE_DEQUANTIZATION(T, mantisa) \ - template void launch_selective_dequantization( \ - uint8_t*, T*, int32_t*, int, int, int, int, int, cudaStream_t); -// fp8(E4M3) -#ifdef BF16_AVAILABLE -INSTANTIATE_LAUNCH_SELECTIVE_DEQUANTIZATION(__nv_bfloat16, 7); -#endif -INSTANTIATE_LAUNCH_SELECTIVE_DEQUANTIZATION(__half, 10); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/context.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/context.h deleted file mode 100644 index 5bd9badbcb4fa06091a8c01b0b0d386153a24589..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/context.h +++ /dev/null @@ -1,66 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include -#include -#include -#include -#include "cublas_v2.h" -#include "cuda.h" -#include "curand.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#define WARP_SIZE 32 - -class FPContext { -public: - FPContext() : _seed(42) - { - curandCreateGenerator(&_gen, CURAND_RNG_PSEUDO_DEFAULT); - curandSetPseudoRandomGeneratorSeed(_gen, 123); - } - - virtual ~FPContext() {} - - static FPContext& Instance() - { - static FPContext _ctx; - return _ctx; - } - - curandGenerator_t& GetRandGenerator() { return _gen; } - - cudaStream_t GetCurrentStream() - { - // get current pytorch stream. - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return stream; - } - - std::pair IncrementOffset(uint64_t offset_inc) - { - uint64_t offset = _curr_offset; - _curr_offset += offset_inc; - return std::pair(_seed, offset); - } - - void SetSeed(uint64_t new_seed) { _seed = new_seed; } - -private: - curandGenerator_t _gen; - cublasHandle_t _cublasHandle; - uint64_t _seed; - uint64_t _curr_offset; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/fp_quantize.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/fp_quantize.h deleted file mode 100644 index 60c75541f603f8b58558ecdf1d5658407f4d2857..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/fp_quantize.h +++ /dev/null @@ -1,128 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include - -#include - -#ifdef BF16_AVAILABLE -#include -#endif -#include -#include - -#define QUANT_SWITCH(Q_BITS, ...) \ - [&] { \ - if (12 == Q_BITS) { \ - constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ - constexpr int CONST_Q_BITS = 8; \ - constexpr int CONST_Q_MANTISA_BITS = 3; \ - __VA_ARGS__(); \ - } else if (13 == Q_BITS) { \ - constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ - constexpr int CONST_Q_BITS = 8; \ - constexpr int CONST_Q_MANTISA_BITS = 3; \ - __VA_ARGS__(); \ - } else if (10 == Q_BITS) { \ - constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ - constexpr int CONST_Q_BITS = 8; \ - constexpr int CONST_Q_MANTISA_BITS = 2; \ - __VA_ARGS__(); \ - } else if (11 == Q_BITS) { \ - constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ - constexpr int CONST_Q_BITS = 8; \ - constexpr int CONST_Q_MANTISA_BITS = 2; \ - __VA_ARGS__(); \ - } else if (28 == Q_BITS) { \ - constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ - constexpr int CONST_Q_BITS = 12; \ - constexpr int CONST_Q_MANTISA_BITS = 7; \ - __VA_ARGS__(); \ - } else if (29 == Q_BITS) { \ - constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ - constexpr int CONST_Q_BITS = 12; \ - constexpr int CONST_Q_MANTISA_BITS = 7; \ - __VA_ARGS__(); \ - } else if (6 == Q_BITS) { \ - constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ - constexpr int CONST_Q_BITS = 6; \ - constexpr int CONST_Q_MANTISA_BITS = 2; \ - __VA_ARGS__(); \ - } else if (7 == Q_BITS) { \ - constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ - constexpr int CONST_Q_BITS = 6; \ - constexpr int CONST_Q_MANTISA_BITS = 2; \ - __VA_ARGS__(); \ - } else if (2 == Q_BITS) { \ - constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ - constexpr int CONST_Q_BITS = 4; \ - constexpr int CONST_Q_MANTISA_BITS = 1; \ - __VA_ARGS__(); \ - } else { \ - constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ - constexpr int CONST_Q_BITS = 4; \ - constexpr int CONST_Q_MANTISA_BITS = 1; \ - __VA_ARGS__(); \ - } \ - }() - -#define DEQUANT_SWITCH(Q_MANTISA_EXPONENT_BITS, ...) \ - [&] { \ - if (12 == Q_MANTISA_EXPONENT_BITS) { \ - constexpr int CONST_Q_MANTISA_BITS = 3; \ - constexpr int CONST_Q_EXPONENT_BITS = 4; \ - __VA_ARGS__(); \ - } else if (10 == Q_MANTISA_EXPONENT_BITS) { \ - constexpr int CONST_Q_MANTISA_BITS = 2; \ - constexpr int CONST_Q_EXPONENT_BITS = 5; \ - __VA_ARGS__(); \ - } else if (28 == Q_MANTISA_EXPONENT_BITS) { \ - constexpr int CONST_Q_MANTISA_BITS = 7; \ - constexpr int CONST_Q_EXPONENT_BITS = 4; \ - __VA_ARGS__(); \ - } else if (6 == Q_MANTISA_EXPONENT_BITS) { \ - constexpr int CONST_Q_MANTISA_BITS = 2; \ - constexpr int CONST_Q_EXPONENT_BITS = 3; \ - __VA_ARGS__(); \ - } else { \ - constexpr int CONST_Q_MANTISA_BITS = 1; \ - constexpr int CONST_Q_EXPONENT_BITS = 2; \ - __VA_ARGS__(); \ - } \ - }() - -template -void launch_quantization(T* val, - uint8_t* q_val, - int num_groups, - int group_size, - cudaStream_t stream, - float q_range, - int q_bits, - int q_mantisa_bits, - int stochastic_rounding); - -template -void launch_dequantization(uint8_t* val, - T* q_val, - int num_groups, - int group_size, - int q_mantisa_bits, - int q_exponent_bits, - cudaStream_t stream); - -template -void launch_selective_dequantization(uint8_t* val, - T* q_val, - int32_t* indexes, - int num_groups, - int group_size, - int num_indexes, - int q_mantisa_bits, - int q_exponent_bits, - cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.cpp deleted file mode 100644 index 460330b93d31f1a992efe02a9f3bddcff297f4ee..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.cpp +++ /dev/null @@ -1,154 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include "deepspeed_gds_op.h" - -using namespace std; - -// For when there is more than 1 device -static std::map> base_ptr_registry; - -static void _safe_handle_register(const int fd, CUfileDescr_t& cf_descr, CUfileHandle_t& cf_handle) -{ - memset((void*)&cf_descr, 0, sizeof(CUfileDescr_t)); - cf_descr.handle.fd = fd; - cf_descr.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD; - CUfileError_t status = cuFileHandleRegister(&cf_handle, &cf_descr); - if (status.err != CU_FILE_SUCCESS) { - std::cerr << "file register error:" << cuFileGetErrorString(status) << std::endl; - close(fd); - exit(EXIT_FAILURE); - } -} - -static void* _find_base_ptr(const int64_t device, char* buf_ptr) -{ - void* base_ptr = nullptr; - int64_t last = -1; - int64_t ptr_diff; - for (const auto& value : base_ptr_registry[device]) { - ptr_diff = buf_ptr - (char*)value; - if (last == -1 && ptr_diff >= 0) { - last = ptr_diff; - base_ptr = value; - } else if (ptr_diff < last && ptr_diff >= 0) { - last = ptr_diff; - base_ptr = value; - } - } - if (!base_ptr || buf_ptr < base_ptr) { - std::cerr << "BASE PTR ERROR :" << base_ptr << " BUF PTR " << (void*)buf_ptr << std::endl; - for (const auto& value : base_ptr_registry[device]) { - std::cerr << "BASE PTR AVAIL :" << value << std::endl; - } - exit(EXIT_FAILURE); - } - - return base_ptr; -} - -void gds_op_desc_t::add_buffer_to_registry(const torch::Tensor& buffer) -{ - const int64_t device = buffer.get_device(); - void* reg_ptr = buffer.data_ptr(); - - // TODO: add checking to make sure pointer isn't already in set - const auto it = base_ptr_registry.find(device); - if (it == base_ptr_registry.end()) { - std::set new_ptr_set; - new_ptr_set.insert(reg_ptr); - base_ptr_registry.insert(std::pair>(device, new_ptr_set)); - } else { - base_ptr_registry[device].insert(reg_ptr); - } - - check_cudaruntimecall(cudaSetDevice(device)); - CUfileError_t status = cuFileBufRegister(reg_ptr, buffer.nbytes(), 0); - if (status.err != CU_FILE_SUCCESS) { - std::cerr << "buffer register failed:" << cuFileGetErrorString(status) << std::endl; - exit(EXIT_FAILURE); - } -} - -void gds_op_desc_t::remove_buffer_from_registry(const torch::Tensor& buffer) -{ - const int64_t device = buffer.get_device(); - void* reg_ptr = buffer.data_ptr(); - - // std::cout << "DEREG PTR " << reg_ptr << std::endl; - check_cudaruntimecall(cudaSetDevice(device)); - cuFileBufDeregister(reg_ptr); - - // Remove from tracked registry - base_ptr_registry[device].erase(reg_ptr); -} - -gds_op_desc_t::gds_op_desc_t(const bool read_op, - const torch::Tensor& buffer, - const int fd, - const char* filename, - const int intra_op_parallelism, - const bool validate, - const int64_t file_offset) - : io_op_desc_t(read_op, buffer, fd, filename, intra_op_parallelism, validate, file_offset) -{ - _contiguous_buffer = _buffer.contiguous(); - const int64_t device = _buffer.get_device(); - check_cudaruntimecall(cudaSetDevice(device)); - _base_ptr = _find_base_ptr(device, (char*)_contiguous_buffer.data_ptr()); - - _safe_handle_register(fd, _cf_descr, _cf_handle); -} - -char* gds_op_desc_t::data_ptr() const { return (char*)_contiguous_buffer.data_ptr(); } - -void gds_op_desc_t::finish() { cuFileHandleDeregister(_cf_handle); } - -void gds_op_desc_t::validate() -{ - check_cudaruntimecall(cudaSetDevice(_buffer.get_device())); - const auto cpu_buffer = _buffer.to(torch::kCPU); - const auto num_io_bytes = static_cast(_contiguous_buffer.nbytes()); - validate_aio_operation( - _read_op, _filename.c_str(), (char*)(cpu_buffer.data_ptr()), num_io_bytes); -} - -void gds_op_desc_t::run(const int tid, - std::unique_ptr& aio_ctxt, - deepspeed_aio_config_t* aio_config) -{ - assert(tid < _intra_op_parallelism); - check_cudaruntimecall(cudaSetDevice(_buffer.get_device())); - const auto buf_offset = data_ptr() + (_num_bytes_per_thread * tid) - (char*)_base_ptr; - const auto tid_file_offset = _file_offset + (_num_bytes_per_thread * tid); - - if (_read_op) { - auto ret = - cuFileRead(_cf_handle, _base_ptr, _num_bytes_per_thread, tid_file_offset, buf_offset); - if (ret < 0) { _report_error(ret, errno, tid_file_offset); } - } else { - auto ret = - cuFileWrite(_cf_handle, _base_ptr, _num_bytes_per_thread, tid_file_offset, buf_offset); - if (ret < 0) { _report_error(ret, errno, tid_file_offset); } - } -} - -void gds_op_desc_t::_report_error(const ssize_t return_code, - const int error_num, - const off_t offset) -{ - const auto op_string = _read_op ? "read failed with " : "write failed with "; - const auto error_string = IS_CUFILE_ERR(return_code) ? "cuFile error: " : "posix error: "; - const auto error_code = IS_CUFILE_ERR(return_code) ? cuFileGetErrorString(return_code) - : cuFileGetErrorString(error_num); - std::cerr << op_string << error_string << error_code << " return code = " << return_code - << " filename = " << _filename << " num bytes = " << _num_bytes_per_thread - << " offset = " << offset << std::endl; - exit(EXIT_FAILURE); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.h deleted file mode 100644 index fe2d3cafb8ef7e706330663bc49a559b7b7cc343..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.h +++ /dev/null @@ -1,44 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include -#include -#include -#include -#include - -#include "deepspeed_aio_op_desc.h" -#include "deepspeed_gds_utils.h" - -struct gds_op_desc_t : io_op_desc_t { - CUfileDescr_t _cf_descr; - CUfileHandle_t _cf_handle; - void* _base_ptr; - - gds_op_desc_t(const bool read_op, - const torch::Tensor& buffer, - const int fd, - const char* filename, - const int intra_op_parallelism, - const bool validate, - const int64_t file_offset); - - void run(const int tid, - std::unique_ptr& aio_ctxt, - deepspeed_aio_config_t* aio_config); - - char* data_ptr() const; - - void validate(); - - void finish(); - - void _report_error(const ssize_t return_code, const int error_num, const off_t offset); - - static void add_buffer_to_registry(const torch::Tensor& buffer); - - static void remove_buffer_from_registry(const torch::Tensor& buffer); -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_utils.h deleted file mode 100644 index 12b014d909880da286ea76b264a72ee15d05d2a7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_utils.h +++ /dev/null @@ -1,91 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include - -// CUDA/cuFile includes -#include -#include -#include "cufile.h" - -// Macro for checking cuda errors following a cuda launch or api call -#define cudaCheckError() \ - { \ - cudaError_t e = cudaGetLastError(); \ - if (e != cudaSuccess) { \ - printf("Cuda failure %s:%d: '%s'\n", __FILE__, __LINE__, cudaGetErrorString(e)); \ - exit(EXIT_FAILURE); \ - } \ - } - -#define check_cudadrivercall(fn) \ - do { \ - CUresult res = fn; \ - if (res != CUDA_SUCCESS) { \ - const char* str = nullptr; \ - cuGetErrorName(res, &str); \ - std::cerr << "cuda driver api call failed " << #fn << " res : " << res << ", " \ - << __LINE__ << ":" << str << std::endl; \ - std::cerr << "EXITING program!!!" << std::endl; \ - exit(1); \ - } \ - } while (0) - -#define check_cudaruntimecall(fn) \ - do { \ - cudaError_t res = fn; \ - if (res != cudaSuccess) { \ - const char* str = cudaGetErrorName(res); \ - std::cerr << "cuda runtime api call failed " << #fn << __LINE__ << ":" << str \ - << std::endl; \ - std::cerr << "EXITING program!!!" << std::endl; \ - exit(1); \ - } \ - } while (0) - -#define check_cuFileCall(fn, api_msg) \ - do { \ - CUfileError_t status = fn; \ - if (status.err != CU_FILE_SUCCESS) { \ - std::cout << api_msg << " failed with error " << CUFILE_ERRSTR(status.err) \ - << std::endl; \ - exit(EXIT_FAILURE); \ - } \ - } while (0) - -// -// cuda driver error description -// -static inline const char* GetCuErrorString(CUresult curesult) -{ - const char* descp; - if (cuGetErrorName(curesult, &descp) != CUDA_SUCCESS) descp = "unknown cuda error"; - return descp; -} - -// -// cuFile APIs return both cuFile specific error codes as well as POSIX error codes -// for ease, the below template can be used for getting the error description depending -// on its type. - -// POSIX -template ::value, std::nullptr_t>::type = nullptr> -std::string cuFileGetErrorString(T status) -{ - status = std::abs(status); - return IS_CUFILE_ERR(status) ? std::string(CUFILE_ERRSTR(status)) - : std::string(std::strerror(status)); -} - -// CUfileError_t -template ::value, std::nullptr_t>::type = nullptr> -std::string cuFileGetErrorString(T status) -{ - std::string errStr = cuFileGetErrorString(static_cast(status.err)); - if (IS_CUDA_ERR(status)) errStr.append(".").append(GetCuErrorString(status.cu_err)); - return errStr; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.cpp deleted file mode 100644 index 4d9dc5445043229ede5ef7c9c2daf7f7521d7bd0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.cpp +++ /dev/null @@ -1,118 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* - GPUDirect Storage functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include "deepspeed_py_gds_handle.h" -#include -#include "deepspeed_gds_op.h" - -using namespace std; - -int deepspeed_gds_handle_t::s_cuFile_init = 0; - -deepspeed_gds_handle_t::deepspeed_gds_handle_t(const int block_size, - const int queue_depth, - const bool single_submit, - const bool overlap_events, - const int intra_op_parallelism) - : deepspeed_io_handle_t(block_size, queue_depth, single_submit, overlap_events, 1), - _intra_gds_op_parallelism(intra_op_parallelism) -{ - _init_cuFile(block_size, queue_depth); -} - -deepspeed_gds_handle_t::~deepspeed_gds_handle_t() { _close_cuFile(); } - -const int deepspeed_gds_handle_t::get_intra_op_parallelism() const -{ - return _intra_gds_op_parallelism; -} - -void deepspeed_gds_handle_t::_init_cuFile(const int block_size, const int queue_depth) -{ - if (deepspeed_gds_handle_t::s_cuFile_init == 0) { - std::string depthStr = std::to_string(queue_depth); - std::string threadsStr = std::to_string(_intra_gds_op_parallelism); - std::string json1 = R"({"execution": {"max_io_queue_depth": )" + depthStr + ", "; - std::string json2 = R"("max_request_parallelism": )" + threadsStr + ", "; - std::string json3 = R"("max_io_threads": )" + threadsStr + ", "; - std::string json4 = R"("parallel_io": true, "min_io_threshold_size_kb": 8192}})"; - std::ofstream outFile("local_cufile.json"); - if (outFile.is_open()) { - outFile << json1 + json2 + json3 + json4; - outFile.close(); - } else { - std::cerr << "Can't open local cufile" << std::endl; - exit(EXIT_FAILURE); - } - // TODO: Address the following issues with this code - // (1) Fix C++14 warning - // (2) Create file in a different location than PWD - // (3) Handle multi-GPU/multi-rank scenarios: should cufile be shared, is per-rank cufile - // safe? - putenv("CUFILE_ENV_PATH_JSON=$PWD/local_cufile.json"); - cuFileDriverOpen(); - cudaCheckError(); - size_t direct_io_size = (size_t)block_size / 1024; - CUfileError_t status = cuFileDriverSetMaxDirectIOSize(direct_io_size); - if (status.err != CU_FILE_SUCCESS) { - std::cerr << "file register error:" << cuFileGetErrorString(status) << std::endl; - exit(EXIT_FAILURE); - } - } - deepspeed_gds_handle_t::s_cuFile_init++; -} - -void deepspeed_gds_handle_t::_close_cuFile() -{ - deepspeed_gds_handle_t::s_cuFile_init--; - if (deepspeed_gds_handle_t::s_cuFile_init == 0) { cuFileDriverClose(); } -} - -torch::Tensor deepspeed_gds_handle_t::new_pinned_device_tensor(const size_t num_elem, - const torch::Tensor& example_tensor) -{ - auto options = torch::TensorOptions().dtype(example_tensor.scalar_type()).device(torch::kCUDA); - auto dev_tensor = torch::empty(num_elem, options); - pin_device_tensor(dev_tensor); - return dev_tensor; -} - -bool deepspeed_gds_handle_t::free_pinned_device_tensor(torch::Tensor& buffer) -{ - unpin_device_tensor(buffer); - return true; -} - -bool deepspeed_gds_handle_t::pin_device_tensor(const torch::Tensor& buffer) -{ - gds_op_desc_t::add_buffer_to_registry(buffer); - return true; -} - -bool deepspeed_gds_handle_t::unpin_device_tensor(const torch::Tensor& buffer) -{ - gds_op_desc_t::remove_buffer_from_registry(buffer); - return true; -} - -std::shared_ptr deepspeed_gds_handle_t::_create_io_op_desc( - const bool read_op, - const torch::Tensor& buffer, - const int fd, - const char* filename, - const bool validate, - const int64_t file_offset) -{ - if (buffer.is_cuda()) { - return std::make_shared( - read_op, buffer, fd, filename, _intra_op_parallelism, validate, file_offset); - } - return deepspeed_io_handle_t::_create_io_op_desc( - read_op, buffer, fd, filename, validate, file_offset); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.h deleted file mode 100644 index f8090fc6fde44a14944a392bd7055e99c87e2bac..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.h +++ /dev/null @@ -1,48 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include -#include -#include "deepspeed_py_io_handle.h" - -struct deepspeed_gds_handle_t : deepspeed_io_handle_t { - const int _intra_gds_op_parallelism; - - deepspeed_gds_handle_t(const int block_size, - const int queue_depth, - const bool single_submit, - const bool overlap_events, - const int intra_op_parallelism); - - ~deepspeed_gds_handle_t(); - - torch::Tensor new_pinned_device_tensor(const size_t num_elem, - const torch::Tensor& example_tensor); - - bool free_pinned_device_tensor(torch::Tensor&); - - bool pin_device_tensor(const torch::Tensor& buffer); - - bool unpin_device_tensor(const torch::Tensor& buffer); - - void _init_cuFile(const int block_size, const int queue_depth); - - void _close_cuFile(); - - const int get_intra_op_parallelism() const; - - std::shared_ptr _create_io_op_desc(const bool read_op, - const torch::Tensor& buffer, - const int fd, - const char* filename, - const bool validate, - const int64_t file_offset); - - static int s_cuFile_init; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/py_ds_gds.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/py_ds_gds.cpp deleted file mode 100644 index bc6327bb38068c9c9abbca9dd1830cbae7da6fcc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/gds/py_lib/py_ds_gds.cpp +++ /dev/null @@ -1,141 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Functionality for swapping optimizer tensors to/from (NVMe) storage devices. -*/ - -#include -#include "deepspeed_py_gds_handle.h" -using namespace pybind11::literals; - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - py::class_(m, "gds_handle") - .def(py::init(), - "GDS handle constructor", - "block_size"_a = 1024 * 1024, - "queue_depth"_a = 128, - "single_submit"_a = false, - "overlap_events"_a = false, - "intra_op_parallelism"_a = 1) - - .def("get_block_size", &deepspeed_gds_handle_t::get_block_size) - .def("get_queue_depth", &deepspeed_gds_handle_t::get_queue_depth) - .def("get_single_submit", &deepspeed_gds_handle_t::get_single_submit) - .def("get_overlap_events", &deepspeed_gds_handle_t::get_overlap_events) - .def("get_intra_op_parallelism", &deepspeed_gds_handle_t::get_intra_op_parallelism) - .def("get_alignment", &deepspeed_gds_handle_t::get_alignment) - - .def("read", - &deepspeed_gds_handle_t::read, - "Synchronous and non-parallel file read. Returns count of completed read ops", - "buffer"_a, - "filename"_a, - "validate"_a, - "file_offset"_a = 0) - - .def("write", - &deepspeed_gds_handle_t::write, - "Synchronous and non-parallel file write. Returns count of completed write ops", - "buffer"_a, - "filename"_a, - "validate"_a, - "file_offset"_a = 0) - - .def("pread", - &deepspeed_gds_handle_t::pread, - "Parallel file read with option of parallelism. Returns count of completed read ops", - "buffer"_a, - "filename"_a, - "validate"_a, - "async"_a, - "file_offset"_a = 0) - - .def("pwrite", - &deepspeed_gds_handle_t::pwrite, - "Parallel file write with option of parallelism. Returns count of completed write ops", - "buffer"_a, - "filename"_a, - "validate"_a, - "async"_a, - "file_offset"_a = 0) - - .def("sync_pread", - &deepspeed_gds_handle_t::sync_pread, - "Synchrononous parallel file read. Returns count of completed read ops", - "buffer"_a, - "filename"_a, - "file_offset"_a = 0) - - .def("sync_pwrite", - &deepspeed_gds_handle_t::sync_pwrite, - "Synchronous parallel file write. Returns count of completed write ops", - "buffer"_a, - "filename"_a, - "file_offset"_a = 0) - - .def("async_pread", - &deepspeed_gds_handle_t::async_pread, - "Asynchronous parallel file read. Returns 0 on success. Returns 0 on success, and " - "following wait() returns count of completed ops.", - "buffer"_a, - "filename"_a, - "file_offset"_a = 0) - - .def( - "async_pwrite", - py::overload_cast( - &deepspeed_gds_handle_t::async_pwrite), - "Asynchronous parallel file write. Returns 0 on success, and subsequent wait() returns " - "count of completed ops.", - "buffer"_a, - "filename"_a, - "file_offset"_a = 0) - - .def("async_pwrite", - py::overload_cast( - &deepspeed_gds_handle_t::async_pwrite), - "Asynchronous parallel file write using opened python file object.", - "buffer"_a, - "fd"_a, - "file_offset"_a = 0) - - .def("new_cpu_locked_tensor", - &deepspeed_gds_handle_t::new_cpu_locked_tensor, - "Allocate pinned CPU tensor.", - "num_elem"_a, - "example_tenosr"_a) - - .def("free_cpu_locked_tensor", - &deepspeed_gds_handle_t::free_cpu_locked_tensor, - "Free pinned CPU tensor.", - "tensor"_a) - - .def("new_pinned_device_tensor", - &deepspeed_gds_handle_t::new_pinned_device_tensor, - "Allocate pinned device tensor.", - "num_elem"_a, - "example_tenosr"_a) - - .def("free_pinned_device_tensor", - &deepspeed_gds_handle_t::free_pinned_device_tensor, - "Free pinned device tensor.", - "tensor"_a) - - .def("pin_device_tensor", - &deepspeed_gds_handle_t::pin_device_tensor, - "Pin device tensor.", - "tensor"_a) - - .def("unpin_device_tensor", - &deepspeed_gds_handle_t::unpin_device_tensor, - "Unpin device tensor.", - "tensor"_a) - - .def("wait", - &deepspeed_gds_handle_t::wait, - "Wait for (ongoing) asynchronous operations to complete"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/StopWatch.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/StopWatch.h deleted file mode 100644 index 42e4eb886eaa086e9eff786601ea8cb784a5b4ed..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/StopWatch.h +++ /dev/null @@ -1,103 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once -#ifdef _WIN32 -#include -#else -#include -#endif - -#ifdef _WIN32 - -class Stopwatch { -private: - double m_total_time; - LARGE_INTEGER m_start_time; - -public: - Stopwatch() { m_total_time = 0.0; } - - ~Stopwatch() {} - - void Reset() { m_total_time = 0.0; } - - void Start() { QueryPerformanceCounter(&m_start_time); } - - void Restart() - { - m_total_time = 0.0; - QueryPerformanceCounter(&m_start_time); - } - - void Stop() - { - LARGE_INTEGER frequency; - LARGE_INTEGER stop_time; - QueryPerformanceFrequency(&frequency); - QueryPerformanceCounter(&stop_time); - m_total_time += - ((double)(stop_time.QuadPart - m_start_time.QuadPart) / (double)frequency.QuadPart); - } - - double GetTimeInSeconds() { return m_total_time; } -}; - -#else - -class Stopwatch { -private: - double m_total_time; - struct timespec m_start_time; - bool m_is_started; - -public: - Stopwatch() - { - m_total_time = 0.0; - m_is_started = false; - } - - ~Stopwatch() {} - - void Reset() { m_total_time = 0.0; } - - void Start() - { - clock_gettime(CLOCK_MONOTONIC, &m_start_time); - m_is_started = true; - } - - void Restart() - { - m_total_time = 0.0; - clock_gettime(CLOCK_MONOTONIC, &m_start_time); - m_is_started = true; - } - - void Stop() - { - if (m_is_started) { - m_is_started = false; - - struct timespec end_time; - clock_gettime(CLOCK_MONOTONIC, &end_time); - - m_total_time += (double)(end_time.tv_sec - m_start_time.tv_sec) + - (double)(end_time.tv_nsec - m_start_time.tv_nsec) / 1e9; - } - } - - double GetTimeInSeconds() - { - if (m_is_started) { - Stop(); - Start(); - } - return m_total_time; - } -}; - -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/Timer.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/Timer.h deleted file mode 100644 index b67e04171dafb38151f2cfadaa1a800a84a21667..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/Timer.h +++ /dev/null @@ -1,51 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#ifndef __TIMER_H__ -#define __TIMER_H__ - -#include -#include -#include "cuda.h" - -class GPUTimer { - cudaEvent_t start, stop; - -public: - GPUTimer() - { - cudaEventCreate(&start); - cudaEventCreate(&stop); - } - ~GPUTimer() - { - cudaEventDestroy(start); - cudaEventDestroy(stop); - } - inline void Record() { cudaEventRecord(start); } - inline void Elapsed(float& time_elapsed) - { - cudaEventRecord(stop); - cudaEventSynchronize(stop); - cudaEventElapsedTime(&time_elapsed, start, stop); - } -}; - -class CPUTimer { - std::chrono::high_resolution_clock::time_point start; - -public: - CPUTimer() : start(std::chrono::high_resolution_clock::now()) {} - inline void Reset() { start = std::chrono::high_resolution_clock::now(); } - inline float Elapsed() - { - auto temp = start; - start = std::chrono::high_resolution_clock::now(); - return (float)(std::chrono::duration_cast(start - temp).count() / - 1e3); - } -}; - -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/activation_type.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/activation_type.h deleted file mode 100644 index a44921d5d650d3fb4d68522e95b82cda08a1ad60..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/activation_type.h +++ /dev/null @@ -1,17 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -enum ActivationType { - GELU = 0, - RELU = 1, - SILU = 2, - GEGLU = 3, - ReGLU = 4, - SiGLU = 5, - IDENTITY = 6, - InvalidType = -1 -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/compat.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/compat.h deleted file mode 100644 index 6d54446d472e9dfa6141436f55e0e623be48acb4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/compat.h +++ /dev/null @@ -1,19 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Copyright NVIDIA/apex -This file is adapted from fused adam in NVIDIA/apex, commit a109f85 -*/ - -#ifndef TORCH_CHECK -#define TORCH_CHECK AT_CHECK -#endif - -#ifdef VERSION_GE_1_3 -#define DATA_PTR data_ptr -#else -#define DATA_PTR data -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/context.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/context.h deleted file mode 100644 index cd80f8fbeebe1450a52894b431e3229200b7ff52..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/context.h +++ /dev/null @@ -1,180 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include -#include -#include -#include -#include "cublas_v2.h" -#include "cuda.h" -#include "curand.h" -#include "gemm_test.h" - -#define WARP_SIZE 32 - -#define CUDA_CHECK(callstr) \ - { \ - cudaError_t error_code = callstr; \ - if (error_code != cudaSuccess) { \ - std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \ - assert(0); \ - } \ - } - -#define CUDA_1D_KERNEL_LOOP(i, n) \ - for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); i += blockDim.x * gridDim.x) - -#define CUDA_2D_KERNEL_LOOP(i, n, j, m) \ - for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); i += blockDim.x * gridDim.x) \ - for (size_t j = blockIdx.y * blockDim.y + threadIdx.y; j < (m); j += blockDim.y * gridDim.y) - -#define DS_CUDA_NUM_THREADS 512 -#define DS_MAXIMUM_NUM_BLOCKS 262144 - -inline int DS_GET_BLOCKS(const int N) -{ - return (std::max)( - (std::min)((N + DS_CUDA_NUM_THREADS - 1) / DS_CUDA_NUM_THREADS, DS_MAXIMUM_NUM_BLOCKS), - // Use at least 1 block, since CUDA does not allow empty block - 1); -} - -class TrainingContext { -public: - TrainingContext() : _workspace(nullptr), _seed(42), _curr_offset(0) - { - curandCreateGenerator(&_gen, CURAND_RNG_PSEUDO_DEFAULT); - curandSetPseudoRandomGeneratorSeed(_gen, 123); - cublasStatus_t stat = cublasCreate(&_cublasHandle); - if (stat != CUBLAS_STATUS_SUCCESS) { - // It would be nice to use cublasGetStatusName and - // cublasGetStatusString, but they were only added in CUDA 11.4.2. - auto message = std::string("Failed to create cublas handle: cublasStatus_t was ") + - std::to_string(stat); - std::cerr << message << std::endl; - throw std::runtime_error(message); - } - } - - virtual ~TrainingContext() - { - cublasDestroy(_cublasHandle); - cudaFree(_workspace); - } - - static TrainingContext& Instance() - { - static TrainingContext _ctx; - return _ctx; - } - - void SetWorkSpace(void* workspace) - { - if (!workspace) { throw std::runtime_error("Workspace is null."); } - _workspace = workspace; - } - - void* GetWorkSpace() { return _workspace; } - - curandGenerator_t& GetRandGenerator() { return _gen; } - - cudaStream_t GetCurrentStream() - { - // get current pytorch stream. - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return stream; - } - - cudaStream_t GetNewStream() { return at::cuda::getStreamFromPool(); } - - cublasHandle_t GetCublasHandle() { return _cublasHandle; } - - std::pair IncrementOffset(uint64_t offset_inc) - { - uint64_t offset = _curr_offset; - _curr_offset += offset_inc; - return std::pair(_seed, offset); - } - - void SetSeed(uint64_t new_seed) { _seed = new_seed; } - - void TestGemmFP16(bool test_gemm, int batch_size, int seq_len, int head_num, int size_per_head) - { - // avoid rerun. - if (_gemm_algos.size() > 0) return; - - if (test_gemm) { - cublasHandle_t handle = GetCublasHandle(); - - std::unique_ptr> test_qkv_fw( - new GemmTest<__half>(batch_size * seq_len, // M - head_num * size_per_head, // N - head_num * size_per_head, // K - CUBLAS_OP_T, - CUBLAS_OP_N, - handle)); - - std::unique_ptr> test_inter( - new GemmTest<__half>(batch_size * seq_len, // M - 4 * head_num * size_per_head, // N - head_num * size_per_head, // K - CUBLAS_OP_T, - CUBLAS_OP_N, - handle)); - - std::unique_ptr> test_output( - new GemmTest<__half>(batch_size * seq_len, // M - head_num * size_per_head, // N - 4 * head_num * size_per_head, // K - CUBLAS_OP_T, - CUBLAS_OP_N, - handle)); - - std::unique_ptr> test_attn_scores( - new StridedGemmTest<__half>(batch_size * head_num, // batch - seq_len, // M - seq_len, // N - size_per_head, // K - CUBLAS_OP_T, - CUBLAS_OP_N, - handle)); - - std::unique_ptr> test_attn_context( - new StridedGemmTest<__half>(batch_size * head_num, // batch - size_per_head, // M - seq_len, // N - seq_len, // K - CUBLAS_OP_N, - CUBLAS_OP_N, - handle)); - - _gemm_algos.push_back(test_qkv_fw->TestAlgo(100)); - _gemm_algos.push_back(test_inter->TestAlgo(100)); - _gemm_algos.push_back(test_output->TestAlgo(100)); - _gemm_algos.push_back(test_attn_scores->TestAlgo(100)); - _gemm_algos.push_back(test_attn_context->TestAlgo(100)); - } else { - // Use default algo. - _gemm_algos.push_back(std::array({99, 99, 99})); - _gemm_algos.push_back(std::array({99, 99, 99})); - _gemm_algos.push_back(std::array({99, 99, 99})); - _gemm_algos.push_back(std::array({99, 99, 99})); - _gemm_algos.push_back(std::array({99, 99, 99})); - } - } - - const std::vector>& GetGemmAlgos() const { return _gemm_algos; } - -private: - curandGenerator_t _gen; - cublasHandle_t _cublasHandle; - void* _workspace; - uint64_t _seed; - uint64_t _curr_offset; - std::vector> _gemm_algos; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/conversion_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/conversion_utils.h deleted file mode 100644 index 3a90a3e91ddf7d38d943e2b95a6f731796369e97..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/conversion_utils.h +++ /dev/null @@ -1,640 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include "ds_kernel_utils.h" - -#include - -#ifdef BF16_AVAILABLE -#include -#endif - -namespace conversion { - -// Basic primitive for constructing conversions -template -DS_D_INLINE TO to(FROM val) -{ - return to(val); -} - -// Specializations - -/********************* Identity Conversions *********************/ -/* -Identity conversions are useful in templated functions where we might have -a fixed destination type. For example, I might have a kernel that accepts -__half, __nv_bfloat16, and float but always want to do the core computation -at floating point: - -T mem_value = input[idx]; -float compute_value = conversion::to(mem_value); - -In practice, we should be able to elide the second template parameter: -float compute_val = conversion::to(mem_value); - -In this case, we need an implementation to handle the T = float case - -NOTE: The type inferencing system appears to be unable to handle inferring the first -template parameter, even in the trivial case. -*/ - -// Floating point types -template <> -DS_D_INLINE double to(double val) -{ - return val; -} -template <> -DS_D_INLINE float to(float val) -{ - return val; -} -template <> -DS_D_INLINE __half to(__half val) -{ - return val; -} -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE __nv_bfloat16 to(__nv_bfloat16 val) -{ - return val; -} -#endif - -// Integer types -template <> -DS_D_INLINE int8_t to(int8_t val) -{ - return val; -} -template <> -DS_D_INLINE uint8_t to(uint8_t val) -{ - return val; -} -template <> -DS_D_INLINE int16_t to(int16_t val) -{ - return val; -} -template <> -DS_D_INLINE uint16_t to(uint16_t val) -{ - return val; -} -template <> -DS_D_INLINE int32_t to(int32_t val) -{ - return val; -} -template <> -DS_D_INLINE uint32_t to(uint32_t val) -{ - return val; -} -template <> -DS_D_INLINE int64_t to(int64_t val) -{ - return val; -} -template <> -DS_D_INLINE uint64_t to(uint64_t val) -{ - return val; -} - -// TODO: evaluate if we want bools - -/********************* To Double Conversions *********************/ - -// * to double variants - -// Would normally like to not use C cast, but this is an important enough conversion -// to keep -template <> -DS_D_INLINE double to(float val) -{ -#ifdef PTX_AVAILABLE - double ret_val; - asm("ctv.rn.f64.f32 %0, %1;\n" : "=d"(ret_val) : "f"(val)); - return ret_val; -#else - return double(val); -#endif -} -// Note: there is a CVT instruction for __half -> double, but there's no inline interface -// for passing a single half value -template <> -DS_D_INLINE double to(__half val) -{ - return to(__half2float(val)); -} -template <> -DS_D_INLINE double to(int64_t val) -{ - return __ll2double_rn(val); -} -template <> -DS_D_INLINE double to(int32_t val) -{ - return __int2double_rn(val); -} -template <> -DS_D_INLINE double to(int16_t val) -{ - return __int2double_rn(val); -} -template <> -DS_D_INLINE double to(int8_t val) -{ - return __int2double_rn(val); -} -template <> -DS_D_INLINE double to(uint64_t val) -{ - return __ull2double_rn(val); -} -template <> -DS_D_INLINE double to(uint32_t val) -{ - return __uint2double_rn(val); -} -template <> -DS_D_INLINE double to(uint16_t val) -{ - return __uint2double_rn(val); -} -template <> -DS_D_INLINE double to(uint8_t val) -{ - return __uint2double_rn(val); -} - -// Same applies here -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE double to(__nv_bfloat16 val) -{ - return to(__bfloat162float(val)); -} -#endif - -/********************* To Float Conversions *********************/ - -template <> -DS_D_INLINE float to(double val) -{ - return __double2float_rn(val); -} -template <> -DS_D_INLINE float to(__half val) -{ - return __half2float(val); -} -template <> -DS_D_INLINE float to(int64_t val) -{ - return __ll2float_rn(val); -} -template <> -DS_D_INLINE float to(int32_t val) -{ - return __int2float_rn(val); -} -template <> -DS_D_INLINE float to(int16_t val) -{ - return __int2float_rn(val); -} -template <> -DS_D_INLINE float to(int8_t val) -{ - return __int2float_rn(val); -} -template <> -DS_D_INLINE float to(uint64_t val) -{ - return __ull2float_rn(val); -} -template <> -DS_D_INLINE float to(uint32_t val) -{ - return __uint2float_rn(val); -} -template <> -DS_D_INLINE float to(uint16_t val) -{ - return __uint2float_rn(val); -} -template <> -DS_D_INLINE float to(uint8_t val) -{ - return __uint2float_rn(val); -} - -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE float to(__nv_bfloat16 val) -{ - return __bfloat162float(val); -} -#endif - -/********************* To Float2 Conversions *********************/ -template <> -DS_D_INLINE float2 to(__half2 val) -{ - return __half22float2(val); -} - -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE float2 to(__nv_bfloat162 val) -{ - return __bfloat1622float2(val); -} -#endif - -/********************* To Half Conversions *********************/ -template <> -DS_D_INLINE __half to(double val) -{ -#ifdef __HIP_PLATFORM_AMD__ - float val_f = __double2float_rn(val); - return __float2half(val_f); -#else - return __double2half(val); -#endif -} -template <> -DS_D_INLINE __half to(float val) -{ - return __float2half(val); -} -template <> -DS_D_INLINE __half to(int64_t val) -{ - return __ll2half_rn(val); -} -template <> -DS_D_INLINE __half to(int32_t val) -{ - return __int2half_rn(val); -} -template <> -DS_D_INLINE __half to(int16_t val) -{ - return __short2half_rn(val); -} -template <> -DS_D_INLINE __half to(int8_t val) -{ - return __int2half_rn(val); -} -template <> -DS_D_INLINE __half to(uint64_t val) -{ - return __ull2half_rn(val); -} -template <> -DS_D_INLINE __half to(uint32_t val) -{ - return __uint2half_rn(val); -} -template <> -DS_D_INLINE __half to(uint16_t val) -{ - return __ushort2half_rn(val); -} -template <> -DS_D_INLINE __half to(uint8_t val) -{ - return __uint2half_rn(val); -} - -#ifdef BF16_AVAILABLE -// No direct conversion -template <> -DS_D_INLINE __half to(__nv_bfloat16 val) -{ - return to<__half>(to(val)); -} -#endif - -/********************* To Half2 Conversions *********************/ -template <> -DS_D_INLINE __half2 to(float2 val) -{ - return __float22half2_rn(val); -} -template <> -DS_D_INLINE __half2 to(float val) -{ - return __float2half2_rn(val); -} - -#ifdef BF16_AVAILABLE -// No direct conversion -template <> -DS_D_INLINE __half2 to(__nv_bfloat162 val) -{ - return to<__half2>(to(val)); -} -#endif - -/********************* To BF16 Conversions *********************/ -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE __nv_bfloat16 to(double val) -{ - return __double2bfloat16(val); -} -template <> -DS_D_INLINE __nv_bfloat16 to(float val) -{ - return __float2bfloat16(val); -} -template <> -DS_D_INLINE __nv_bfloat16 to(int64_t val) -{ - return __ll2bfloat16_rn(val); -} -template <> -DS_D_INLINE __nv_bfloat16 to(int32_t val) -{ - return __int2bfloat16_rn(val); -} -template <> -DS_D_INLINE __nv_bfloat16 to(int16_t val) -{ - return __short2bfloat16_rn(val); -} -template <> -DS_D_INLINE __nv_bfloat16 to(int8_t val) -{ - return __int2bfloat16_rn(val); -} -template <> -DS_D_INLINE __nv_bfloat16 to(uint64_t val) -{ - return __ull2bfloat16_rn(val); -} -template <> -DS_D_INLINE __nv_bfloat16 to(uint32_t val) -{ - return __uint2bfloat16_rn(val); -} -template <> -DS_D_INLINE __nv_bfloat16 to(uint16_t val) -{ - return __ushort2bfloat16_rn(val); -} -template <> -DS_D_INLINE __nv_bfloat16 to(uint8_t val) -{ - return __uint2bfloat16_rn(val); -} -#endif - -/********************* To BF162 Conversions *********************/ -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE __nv_bfloat162 to(float2 val) -{ - return __float22bfloat162_rn(val); -} -template <> -DS_D_INLINE __nv_bfloat162 to(float val) -{ - return __float2bfloat162_rn(val); -} -template <> -DS_D_INLINE __nv_bfloat162 to(__half2 val) -{ - return to<__nv_bfloat162>(to(val)); -} -#endif - -/********************* To INT64_T Conversions *********************/ -template <> -DS_D_INLINE int64_t to(double val) -{ - return __double2ll_rn(val); -} -template <> -DS_D_INLINE int64_t to(float val) -{ - return __float2ll_rn(val); -} -template <> -DS_D_INLINE int64_t to(__half val) -{ - return __half2ll_rn(val); -} -// No direct support for integer casts at the C++ level and I don't feel they're so important -// to demand an PTX at this time - -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE int64_t to(__nv_bfloat16 val) -{ - return __bfloat162ll_rn(val); -} -#endif - -/********************* To INT32_T Conversions *********************/ -template <> -DS_D_INLINE int32_t to(double val) -{ - return __double2int_rn(val); -} -template <> -DS_D_INLINE int32_t to(float val) -{ - return __float2int_rn(val); -} -template <> -DS_D_INLINE int32_t to(__half val) -{ - return __half2int_rn(val); -} -// No direct support for integer casts at the C++ level and I don't feel they're so important -// to demand an PTX at this time - -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE int32_t to(__nv_bfloat16 val) -{ - return __bfloat162int_rn(val); -} -#endif - -/********************* To INT16_T Conversions *********************/ -template <> -DS_D_INLINE int16_t to(double val) -{ - return __double2int_rn(val); -} -template <> -DS_D_INLINE int16_t to(float val) -{ - return __float2int_rn(val); -} -template <> -DS_D_INLINE int16_t to(__half val) -{ - return __half2int_rn(val); -} -// No direct support for integer casts at the C++ level and I don't feel they're so important -// to demand an PTX at this time - -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE int16_t to(__nv_bfloat16 val) -{ - return __bfloat162int_rn(val); -} -#endif - -/********************* To INT8_T Conversions *********************/ -template <> -DS_D_INLINE int8_t to(double val) -{ - return __double2int_rn(val); -} -template <> -DS_D_INLINE int8_t to(float val) -{ - return __float2int_rn(val); -} -template <> -DS_D_INLINE int8_t to(__half val) -{ - return __half2int_rn(val); -} -// No direct support for integer casts at the C++ level and I don't feel they're so important -// to demand an PTX at this time - -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE int8_t to(__nv_bfloat16 val) -{ - return __bfloat162int_rn(val); -} -#endif - -/********************* To UINT64_T Conversions *********************/ -template <> -DS_D_INLINE uint64_t to(double val) -{ - return __double2ull_rn(val); -} -template <> -DS_D_INLINE uint64_t to(float val) -{ - return __float2ull_rn(val); -} -template <> -DS_D_INLINE uint64_t to(__half val) -{ - return __half2ull_rn(val); -} -// No direct support for integer casts at the C++ level and I don't feel they're so important -// to demand an PTX at this time - -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE uint64_t to(__nv_bfloat16 val) -{ - return __bfloat162ull_rn(val); -} -#endif - -/********************* To UINT32_T Conversions *********************/ -template <> -DS_D_INLINE uint32_t to(double val) -{ - return __double2uint_rn(val); -} -template <> -DS_D_INLINE uint32_t to(float val) -{ - return __float2uint_rn(val); -} -template <> -DS_D_INLINE uint32_t to(__half val) -{ - return __half2uint_rn(val); -} -// No direct support for integer casts at the C++ level and I don't feel they're so important -// to demand an PTX at this time - -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE uint32_t to(__nv_bfloat16 val) -{ - return __bfloat162uint_rn(val); -} -#endif - -/********************* To UINT16_T Conversions *********************/ -template <> -DS_D_INLINE uint16_t to(double val) -{ - return __double2uint_rn(val); -} -template <> -DS_D_INLINE uint16_t to(float val) -{ - return __float2uint_rn(val); -} -template <> -DS_D_INLINE uint16_t to(__half val) -{ - return __half2uint_rn(val); -} -// No direct support for integer casts at the C++ level and I don't feel they're so important -// to demand an PTX at this time - -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE uint16_t to(__nv_bfloat16 val) -{ - return __bfloat162uint_rn(val); -} -#endif - -/********************* To UINT8_T Conversions *********************/ -template <> -DS_D_INLINE uint8_t to(double val) -{ - return __double2uint_rn(val); -} -template <> -DS_D_INLINE uint8_t to(float val) -{ - return __float2uint_rn(val); -} -template <> -DS_D_INLINE uint8_t to(__half val) -{ - return __half2uint_rn(val); -} -// No direct support for integer casts at the C++ level and I don't feel they're so important -// to demand an PTX at this time - -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE uint8_t to(__nv_bfloat16 val) -{ - return __bfloat162uint_rn(val); -} -#endif - -} // namespace conversion diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_adagrad.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_adagrad.h deleted file mode 100644 index 6f500250f033a1fa222bba0aaafef2f2852f1156..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_adagrad.h +++ /dev/null @@ -1,120 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#define NOMINMAX // Windows idiosyncrasy - // https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c - -#include -#include -#include -#include "simd.h" - -#define STEP(SPAN) \ - template \ - void Step_##SPAN(ds_params_precision_t* _params, \ - ds_params_precision_t* grads, \ - ds_state_precision_t* _exp_avg_sq, \ - size_t _param_size); - -class Adagrad_Optimizer { -public: - Adagrad_Optimizer(float alpha = 1e-2, float eps = 1e-8, float weight_decay = 0) - : _alpha(alpha), _eps(eps), _weight_decay(weight_decay) - { - } - ~Adagrad_Optimizer() {} -#if defined(__AVX512__) or defined(__AVX256__) - template - void Step_AVX(size_t* rounded_size, - ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg_sq, - size_t param_size); -#endif - STEP(1) - STEP(4) - STEP(8) - inline void IncrementStep(size_t step) - { - _step++; - if (_step != step) { _step = step; } - } - inline void update_state(float lr, float epsilon, float weight_decay) - { - _alpha = lr; - _eps = epsilon; - _weight_decay = weight_decay; - } - -private: - float _alpha; - float _eps; - float _weight_decay; - - float _betta1_t; - float _betta2_t; - size_t _step; -}; - -#if defined(__AVX512__) or defined(__AVX256__) -template -void Adagrad_Optimizer::Step_AVX(size_t* rounded_size, - ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg_sq, - size_t _param_size) -{ -#if !defined(__AVX512__) - if (std::is_same_v || - std::is_same_v) { - return; - } -#endif - size_t new_rounded_size = 0; - AVX_Data eps_4; - eps_4.data = SIMD_SET(_eps); - - float step_size = -1 * _alpha; - AVX_Data step_size_4; - step_size_4.data = SIMD_SET(step_size); - - AVX_Data weight_decay4; - if (_weight_decay > 0) weight_decay4.data = SIMD_SET(_weight_decay); - new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span); - for (size_t t = 0; t < new_rounded_size; t += TILE) { - size_t copy_size = TILE; - if ((t + TILE) > new_rounded_size) copy_size = new_rounded_size - t; - size_t offset = copy_size + t; -#pragma omp parallel for - for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { - AVX_Data grad_4[span]; - simd_load(grad_4, grads + i); - - AVX_Data momentum_4[span]; - simd_load(momentum_4, grads + i); - - AVX_Data variance_4[span]; - simd_load(variance_4, _exp_avg_sq + i); - - AVX_Data param_4[span]; - simd_load(param_4, _params + i); - - if (_weight_decay > 0) { simd_fma(grad_4, param_4, weight_decay4, grad_4); } - - simd_fma(variance_4, grad_4, grad_4, variance_4); - simd_sqrt(grad_4, variance_4); - simd_add(grad_4, grad_4, eps_4); - simd_div(grad_4, momentum_4, grad_4); - simd_fma(param_4, grad_4, step_size_4, param_4); - - simd_store(_params + i, param_4); - simd_store(_exp_avg_sq + i, variance_4); - } - } - *rounded_size = new_rounded_size; -} -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_adam.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_adam.h deleted file mode 100644 index a7db6fda370543823fe9ac90b90d66f6e7a002ae..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_adam.h +++ /dev/null @@ -1,220 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#define NOMINMAX // Windows idiosyncrasy - // https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c - -#include -#include -#include -#include "simd.h" - -#define STEP(SPAN) \ - template \ - void Step_##SPAN(ds_params_precision_t* _params, \ - ds_params_precision_t* grads, \ - ds_state_precision_t* _exp_avg, \ - ds_state_precision_t* _exp_avg_sq, \ - size_t _param_size); - -class Adam_Optimizer { -public: - Adam_Optimizer(float alpha = 1e-3, - float betta1 = 0.9, - float betta2 = 0.999, - float eps = 1e-8, - float weight_decay = 0, - bool adamw_mode = true) - : _alpha(alpha), - _betta1(betta1), - _betta2(betta2), - _eps(eps), - _weight_decay(weight_decay), - _betta1_t(1.0), - _betta2_t(1.0), - _step(0), - _adamw_mode(adamw_mode) - { - } - ~Adam_Optimizer() {} - -#if defined(__AVX512__) or defined(__AVX256__) - template - void Step_AVX(size_t* rounded_size, - ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg, - ds_state_precision_t* _exp_avg_sq, - size_t param_size); -#endif - STEP(1) - STEP(4) - STEP(8) - inline void IncrementStep(size_t step, float beta1, float beta2) - { - if (beta1 != _betta1 || beta2 != _betta2) { - _step = step; - _betta1 = beta1; - _betta2 = beta2; - _betta1_t = std::pow(_betta1, step); - _betta2_t = std::pow(_betta2, step); - } else { - _step++; - if (_step != step) { - _betta1_t = std::pow(_betta1, step); - _betta2_t = std::pow(_betta2, step); - _step = step; - } else { - _betta1_t *= _betta1; - _betta2_t *= _betta2; - } - } - } - inline void update_state(float lr, float epsilon, float weight_decay, bool bias_correction) - { - _alpha = lr; - _eps = epsilon; - _weight_decay = weight_decay; - - _bias_correction1 = 1.0f; - _bias_correction2 = 1.0f; - if (bias_correction == 1) { - _bias_correction1 = 1 - _betta1_t; - _bias_correction2 = 1 / sqrt(1 - _betta2_t); - } - } - -private: - float _alpha; - float _betta1; - float _betta2; - float _eps; - float _weight_decay; - - float _betta1_t; - float _betta2_t; - size_t _step; - - float _bias_correction1; - float _bias_correction2; - - bool _adamw_mode; -}; - -#if defined(__AVX512__) or defined(__AVX256__) -template -void Adam_Optimizer::Step_AVX(size_t* rounded_size, - ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg, - ds_state_precision_t* _exp_avg_sq, - size_t _param_size) -{ -#if !defined(__AVX512__) - if (std::is_same_v || - std::is_same_v) { - return; - } -#endif - size_t new_rounded_size = 0; - - AVX_Data betta1_4; - betta1_4.data = SIMD_SET(_betta1); - AVX_Data betta2_4; - betta2_4.data = SIMD_SET(_betta2); - - float betta1_minus1 = 1 - _betta1; - float betta2_minus1 = 1 - _betta2; - AVX_Data betta1_minus1_4; - betta1_minus1_4.data = SIMD_SET(betta1_minus1); - AVX_Data betta2_minus1_4; - betta2_minus1_4.data = SIMD_SET(betta2_minus1); - - AVX_Data bias2_sqrt; - bias2_sqrt.data = SIMD_SET(_bias_correction2); - - AVX_Data eps_4; - eps_4.data = SIMD_SET(_eps); - - float step_size = -1 * _alpha / _bias_correction1; - AVX_Data step_size_4; - step_size_4.data = SIMD_SET(step_size); - - float w_decay = -1 * _alpha * _weight_decay; - AVX_Data weight_decay4; - if (_weight_decay > 0) - weight_decay4.data = (_adamw_mode ? SIMD_SET(w_decay) : SIMD_SET(_weight_decay)); - new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span); - for (size_t t = 0; t < new_rounded_size; t += TILE) { - size_t copy_size = TILE; - if ((t + TILE) > new_rounded_size) copy_size = new_rounded_size - t; - size_t offset = copy_size + t; -#pragma omp parallel for - for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { - AVX_Data grad_4[span]; - simd_load(grad_4, grads + i); - - AVX_Data momentum_4[span]; - simd_load(momentum_4, _exp_avg + i); - - AVX_Data variance_4[span]; - simd_load(variance_4, _exp_avg_sq + i); - - AVX_Data param_4[span]; - simd_load(param_4, _params + i); - - if (_weight_decay > 0 && !_adamw_mode) { - simd_fma(grad_4, param_4, weight_decay4, grad_4); - } - - simd_mul(momentum_4, momentum_4, betta1_4); - simd_fma(momentum_4, grad_4, betta1_minus1_4, momentum_4); - simd_mul(variance_4, variance_4, betta2_4); - simd_mul(grad_4, grad_4, grad_4); - simd_fma(variance_4, grad_4, betta2_minus1_4, variance_4); - simd_sqrt(grad_4, variance_4); - simd_fma(grad_4, grad_4, bias2_sqrt, eps_4); - simd_div(grad_4, momentum_4, grad_4); - - if (_weight_decay > 0 && _adamw_mode) { - simd_fma(param_4, param_4, weight_decay4, param_4); - } - - simd_fma(param_4, grad_4, step_size_4, param_4); - - simd_store(_params + i, param_4); - simd_store(_exp_avg + i, momentum_4); - simd_store(_exp_avg_sq + i, variance_4); - } - } - *rounded_size = new_rounded_size; -} -#endif - -int create_adam_optimizer(int optimizer_id, - float alpha = 1e-3, - float betta1 = 0.9, - float betta2 = 0.999, - float eps = 1e-8, - float weight_decay = 0, - bool adamw_mode = true, - bool should_log = false); - -int ds_adam_step(int optimizer_id, - size_t step, - float lr, - float beta1, - float beta2, - float epsilon, - float weight_decay, - bool bias_correction, - torch::Tensor& params, - torch::Tensor& grads, - torch::Tensor& exp_avg, - torch::Tensor& exp_avg_sq); - -int destroy_adam_optimizer(int optimizer_id); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_lion.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_lion.h deleted file mode 100644 index beaf357a3211773ecb9abe023a48ed39b1276362..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cpu_lion.h +++ /dev/null @@ -1,167 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#define NOMINMAX // Windows idiosyncrasy - // https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c - -#include -#include -#include -#include "simd.h" - -#define STEP(SPAN) \ - template \ - void Step_##SPAN(ds_params_precision_t* _params, \ - ds_params_precision_t* grads, \ - ds_state_precision_t* _exp_avg, \ - size_t _param_size); - -class Lion_Optimizer { -public: - Lion_Optimizer(float alpha = 1e-3, - float betta1 = 0.9, - float betta2 = 0.999, - float weight_decay = 0) - : _alpha(alpha), _betta1(betta1), _betta2(betta2), _weight_decay(weight_decay), _step(0) - { - } - ~Lion_Optimizer() {} - -#if defined(__AVX512__) or defined(__AVX256__) - template - void Step_AVX(size_t* rounded_size, - ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg, - size_t param_size); -#endif - STEP(1) - STEP(4) - STEP(8) - - inline void IncrementStep(size_t step, float beta1, float beta2) - { - _step++; - if (_step != step || beta1 != _betta1 || beta2 != _betta2) { - _step = step; - _betta1 = beta1; - _betta2 = beta2; - } - } - inline void update_state(float lr, float weight_decay) - { - _alpha = lr; - _weight_decay = weight_decay; - } - -private: - float _alpha; - float _betta1; - float _betta2; - float _weight_decay; - size_t _step; -}; - -#if defined(__AVX512__) or defined(__AVX256__) -template -void Lion_Optimizer::Step_AVX(size_t* rounded_size, - ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg, - size_t _param_size) -{ -#if !defined(__AVX512__) - if (std::is_same_v || - std::is_same_v) { - return; - } -#endif - size_t new_rounded_size = 0; - - constexpr float neg1 = -1.0f; - AVX_Data neg1_4; - neg1_4.data = SIMD_SET(neg1); - - AVX_Data betta1_4; - betta1_4.data = SIMD_SET(_betta1); - AVX_Data betta2_4; - betta2_4.data = SIMD_SET(_betta2); - - float betta1_minus1 = 1 - _betta1; - float betta2_minus1 = 1 - _betta2; - AVX_Data betta1_minus1_4; - betta1_minus1_4.data = SIMD_SET(betta1_minus1); - AVX_Data betta2_minus1_4; - betta2_minus1_4.data = SIMD_SET(betta2_minus1); - - float step_size = -_alpha; - AVX_Data step_size_4; - step_size_4.data = SIMD_SET(step_size); - - float after_decay = 1.0f - _alpha * _weight_decay; - AVX_Data after_decay_4; - if (_weight_decay > 0) after_decay_4.data = SIMD_SET(after_decay); - - new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span); - for (size_t t = 0; t < new_rounded_size; t += TILE) { - size_t copy_size = TILE; - if ((t + TILE) > new_rounded_size) copy_size = new_rounded_size - t; - size_t offset = copy_size + t; - -#pragma omp parallel for - for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { - AVX_Data grad_4[span]; - simd_load(grad_4, grads + i); - - AVX_Data momentum_4[span]; - simd_load(momentum_4, _exp_avg + i); - - AVX_Data param_4[span]; - simd_load(param_4, _params + i); - - AVX_Data tmp_4[span]; - - simd_mul(tmp_4, momentum_4, betta1_4); - simd_fma(tmp_4, grad_4, betta1_minus1_4, tmp_4); - // We already used intrinsics, so consider the machine representation fixed. - simd_and(tmp_4, tmp_4, neg1_4); - simd_xor(tmp_4, tmp_4, step_size_4); - if (_weight_decay > 0) { - simd_fma(param_4, param_4, after_decay_4, tmp_4); - } else { - simd_add(param_4, param_4, tmp_4); - } - - simd_mul(momentum_4, momentum_4, betta2_4); - simd_fma(momentum_4, grad_4, betta2_minus1_4, momentum_4); - - simd_store(_params + i, param_4); - simd_store(_exp_avg + i, momentum_4); - } - } - *rounded_size = new_rounded_size; -} -#endif - -int create_lion_optimizer(int optimizer_id, - float alpha = 1e-3, - float betta1 = 0.9, - float betta2 = 0.999, - float weight_decay = 0, - bool should_log = false); - -int ds_lion_step(int optimizer_id, - size_t step, - float lr, - float beta1, - float beta2, - float weight_decay, - torch::Tensor& params, - torch::Tensor& grads, - torch::Tensor& exp_avg); - -int destroy_lion_optimizer(int optimizer_id); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cublas_wrappers.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cublas_wrappers.h deleted file mode 100644 index 2721fb990c7e270e89dbba2dd636587612cb6465..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/cublas_wrappers.h +++ /dev/null @@ -1,101 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include -#include -#include -#include -#ifndef __HIP_PLATFORM_AMD__ -#include -#endif -#ifdef __HIP_PLATFORM_AMD__ -#include -#endif -#include -#include - -int cublas_gemm_ex(cublasHandle_t handle, - cublasOperation_t transa, - cublasOperation_t transb, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const float* A, - const float* B, - float* C, -// TODO HIP: Remove backward compatibility for torch<=2.0 in future -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo algo = rocblas_gemm_algo_standard); -#else - cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT); -#endif - -int cublas_gemm_ex(cublasHandle_t handle, - cublasOperation_t transa, - cublasOperation_t transb, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const __half* A, - const __half* B, - __half* C, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo algo = rocblas_gemm_algo_standard); -#else - cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif - -int cublas_strided_batched_gemm(cublasHandle_t handle, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const float* A, - const float* B, - float* C, - cublasOperation_t op_A, - cublasOperation_t op_B, - int stride_A, - int stride_B, - int stride_C, - int batch, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo algo = rocblas_gemm_algo_standard); -#else - cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT); -#endif - -int cublas_strided_batched_gemm(cublasHandle_t handle, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const __half* A, - const __half* B, - __half* C, - cublasOperation_t op_A, - cublasOperation_t op_B, - int stride_A, - int stride_B, - int stride_C, - int batch, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo algo = rocblas_gemm_algo_standard); -#else - cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/custom_cuda_layers.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/custom_cuda_layers.h deleted file mode 100644 index 21f19749d4cf4741d2264e504d01e74379396e14..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/custom_cuda_layers.h +++ /dev/null @@ -1,324 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include "ds_kernel_utils.h" - -#include -#include -#include -#include -#include - -#include "context.h" -#include "cublas_wrappers.h" - -#define CUDA_CHECK(callstr) \ - { \ - cudaError_t error_code = callstr; \ - if (error_code != cudaSuccess) { \ - std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \ - assert(0); \ - } \ - } - -#define MAX_THREADS 1024 -#define THREADS 256 - -#define MAX_THREAD_STRIDE 32 -#define TILE_DIM 32 - -// Maximum sequence-length support based on the number of threads (2048) allowed in each block and -// this MAX is 8K For higher sequence length we need to use higher Max, like for 64K : 32 -#define MAX_THREAD_ITERATIONS 8 // Maximum 8K -#define MAX_WARP_NUM 32 - -#define MAX_REGISTERS 256 - -#define MAX_REG 256 - -#define WARP_SIZE_BITS 5 - -// Fused bias add with gelu activation -template -void launch_bias_gelu(const T* input, - const T* bias, - T* output, - int intermediate_size, - int batch_size, - cudaStream_t stream); - -template -void launch_gelu(const T* input, - T* output, - int intermediate_size, - int batch_size, - cudaStream_t stream); - -template -void launch_d_gelu(T* d_output, - const T* input, - const T* bias, - int intermediate_size, - int batch_size, - cudaStream_t stream); - -// Custom fused bias add with layer normalization -template -void launch_bias_residual_layer_norm(T* vals, - const T* residual, - const T* gamma, - const T* beta, - float epsilon, - int batch_size, - int hidden_dim, - cudaStream_t stream, - bool preLayerNorm, - bool training, - T* vars, - T* means); - -template -void launch_bias_residual_layer_norm(T* vals, - const T* residual, - const T* gamma, - const T* beta, - float epsilon, - int batch_size, - int hidden_dim, - cudaStream_t stream, - bool preLayerNorm, - bool training, - T* vars); - -template -void launch_layerNorm_backward_fused_add(const T* out_grad1, - const T* out_grad2, - const T* X_data, - const T* vars, - const T* means, - const T* gamma, - T* gamma_grad, - T* betta_grad, - T* inp_grad, - int batch_size, - int hidden_dim, - cudaStream_t stream[2]); -template -void launch_layerNorm_backward_fused_add(const T* out_grad1, - const T* out_grad2, - const T* vals_hat, - const T* vars, - const T* gamma, - T* gamma_grad, - T* betta_grad, - T* inp_grad, - int batch_size, - int hidden_dim, - cudaStream_t stream[2], - bool invertible = false, - const T* betta = nullptr); - -template -void launch_layerNorm_backward(const T* out_grad, - const T* X_data, - const T* vars, - const T* means, - const T* gamma, - T* gamma_grad, - T* betta_grad, - T* inp_grad, - int batch_size, - int hidden_dim, - cudaStream_t stream[2]); - -template -void launch_layerNorm_backward(const T* out_grad, - const T* vals_hat, - const T* vars, - const T* gamma, - T* gamma_grad, - T* betta_grad, - T* inp_grad, - int batch_size, - int hidden_dim, - cudaStream_t stream[2], - bool invertible = false, - const T* betta = nullptr); - -template -void launch_layerNorm_backward_nreversible(const T* out_grad, - const T* vals, - const T* out_grad_trans, - const T* vals_trans, - const T* means, - const T* vars, - const T* gamma, - T* gamma_grad, - T* betta_grad, - T* inp_grad, - int batch_size, - int hidden_dim, - cudaStream_t stream[2]); - -template -void Transpose(const T* inp_mat, T* out_mat, int rows, int cols, cudaStream_t stream); - -template -void launch_attn_softmax_backward(T* out_grad, - const T* soft_inp, - int batch_size, - int heads, - int seq_length, - cudaStream_t stream); - -template -void launch_attn_softmax_backward_v2(T* out_grad, - const T* soft_inp, - int batch_size, - int heads, - int seq_length, - cudaStream_t stream); - -// Custom softmax with scaling and attention mask addition -template -void launch_attn_softmax(T* vals, - const T* attn_mask, - int batch_size, - int heads, - int sequence_length, - cudaStream_t stream); - -template -void launch_transform_0213(T* output, - const T* vals, - int batch_size, - int seq_length, - int hidden_dim, - int heads, - cudaStream_t stream); - -// Custom bias add -template -void launch_bias_add_transform_0213(T* outputs, - const T* vals, - const T* bias, - int batch_size, - int seq_length, - int hidden_dim, - int heads, - cudaStream_t stream, - int trans_count); - -// 4D transform [0, 1, 2, 3] -> [0, 2, 1, 3] -template -void launch_transform4d_0213(T* out, - const T* in, - int batch_size, - int heads, - int seq_length, - int hidden_dim, - cudaStream_t stream, - int trans_count); - -template -void launch_dropout(T* vals, - const T* bias, - uint8_t* mask, - int batch, - int dim, - float ratio, - cudaStream_t stream); - -template -void launch_dropout(T* vals_out, - const T* vals, - uint8_t* mask, - int total_count, - int dim, - float ratio, - cudaStream_t stream, - bool bwd = false); - -template -void launch_dropout(T* out, - const T* vals, - const T* residual, - const T* bias, - uint8_t* mask, - int batch, - int dim, - float ratio, - cudaStream_t stream); - -template -void launch_dropout_grad(T* vals, uint8_t* mask, int total_count, float ratio, cudaStream_t stream); - -template -void launch_dropout_grad(T* vals_out, - const T* vals, - uint8_t* mask, - int total_count, - float ratio, - cudaStream_t stream); - -template -void launch_fuse_transpose_bias_kernel(const T* inp, - T* out, - int rows, - int cols, - cudaStream_t stream); - -void launch_token_sort(int32_t* indices, - int layers, - int batch_size, - int reserved_size, - int original_tokens, - cudaStream_t stream); - -template -void launch_gather_tokens(T* retained_tokens, - T* activations, - int32_t* gather_indices, - int32_t batch_size, - int32_t sampled_tokens, - int32_t channels, - int32_t read_batch_stride, - int32_t read_seq_stride, - int32_t write_batch_stride, - int32_t write_seq_stride, - cudaStream_t stream); - -template -void launch_scatter_tokens(T* all_activations, - T* layer_activations, - int32_t* gather_indices, - int32_t batch_size, - int32_t sampled_tokens, - int32_t channels, - int32_t read_batch_stride, - int32_t read_seq_stride, - int32_t write_batch_stride, - int32_t write_seq_stride, - cudaStream_t stream); - -template -void launch_slice_gpt_mask(T* output_mask, - const T* input_mask, - int batch_size, - int truncated_seq_len, - int orig_seq_len, - cudaStream_t stream); - -template -void launch_slice_bert_mask(T* output_mask, - const T* input_mask, - const int32_t* retained_indices, - int32_t layers, - int32_t batch_size, - int32_t truncated_seq_len, - int32_t orig_seq_len, - cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/deepcompile.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/deepcompile.h deleted file mode 100644 index 810664267d3127321be309b7a6a3ed575a1c7e19..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/deepcompile.h +++ /dev/null @@ -1,615 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#define NOMINMAX // Windows idiosyncrasy - // https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c - -#define USE_C10D_NCCL - -#include -#include - -#include -#include -#include -#include -#include -#include - -#if __has_include() -#include -#else -#include -#endif - -namespace dc { - -template -static bool hasKey(const std::unordered_map& map, const K& key) -{ - return map.find(key) != map.end(); -} - -template -inline std::string to_string(const T& v) -{ - std::stringstream ss; - ss << v; - return ss.str(); -} - -template -size_t productDim(const L& dim) -{ - size_t prod = 1; - for (auto d : dim) { prod *= d; } - return prod; -} - -template -std::string join_as_str(const T& v, const char* delim = ",", const size_t maxlen = 0) -{ - std::stringstream ss; - - if (!v.empty()) { - auto it = v.begin(); - ss << to_string(*it); - it++; - for (; it != v.end(); ++it) { - if (delim) ss << delim; - ss << to_string(*it); - } - } - - std::string s = ss.str(); - if (maxlen > 0 && s.length() > maxlen) { s = s.substr(0, maxlen) + " ..."; } - - return "[" + s + "]"; -} - -template -std::string tensorPtrToString(T* ptr, size_t size, size_t str_len = 100) -{ - std::vector vals; - for (size_t i = 0; i < size; i++) { - vals.push_back(*ptr); - ptr++; - } - return join_as_str(vals, ",", str_len); -} - -std::string tensorPtrToString(void* ptr, - size_t size, - c10::ScalarType datatype, - size_t max_elem = 20, - size_t max_str_len = 100); - -std::string tensorToString(const at::Tensor& t, size_t max_elem = 20, size_t max_str_len = 100); - -std::string tensorDimToString(const at::Tensor& t); - -at::Tensor test_call(at::Tensor param); - -extern c10::intrusive_ptr process_group; -extern c10::intrusive_ptr symm_mem; -extern ncclComm_t nccl_comm; -extern bool use_symm_mem; -extern bool clone_custom_op_output; -extern bool profile; -extern bool pre_div_reduce; - -extern bool sync_before_reduce; // for debugging -extern bool sync_after_reduce; // for debugging -extern bool sync_before_allgather; // for debugging -extern bool sync_after_allgather; // for debugging - -std::vector sizes_to_int_vector(at::IntArrayRef sizes); -void enable_profiling(bool enable); -bool is_profiling(); - -c10::intrusive_ptr getSymmMemWorkspace(int64_t size); -void lazy_init_symm_memory(); -ncclDataType_t get_nccl_data_type(at::ScalarType scalar_type); -void cleanup(); - -class ReduceTask { -public: - ReduceTask(long ds_id, at::Tensor grad, at::Tensor send_buf) - : ds_id_(ds_id), grad_(std::move(grad)), send_buf_(std::move(send_buf)) - { - } - - long getDSId() const { return ds_id_; } - at::Tensor getSendBuf() const { return send_buf_; } - -private: - long ds_id_; - at::Tensor grad_; - at::Tensor send_buf_; -}; - -class ReduceBucket { -public: - ReduceBucket(int64_t size, at::ScalarType scalar_type) : size_(size), scalar_type_(scalar_type) - { - buffer_ = torch::empty({size}, at::TensorOptions().dtype(scalar_type).device(at::kCUDA)); - offset_ = 0; - } - - int64_t getSize() const { return size_; } - int64_t getOffset() const { return offset_; } - at::Tensor getBuffer() const { return buffer_; } - at::ScalarType getScalarType() const { return scalar_type_; } - - void reserve(int64_t size) - { - if (size > size_) { - buffer_ = - torch::empty({size}, at::TensorOptions().dtype(scalar_type_).device(at::kCUDA)); - size_ = size; - } - } - - at::Tensor allocate(int64_t numel) - { - if (offset_ + numel > size_) { - throw std::runtime_error("Buffer size exceeds the reduce bucket size"); - } - - at::Tensor result = buffer_.index({torch::indexing::Slice(offset_, offset_ + numel)}); - offset_ += numel; - return result; - } - - bool shouldFlush(int64_t numel) { return offset_ > 0 && offset_ + numel > size_; } - - void reset() { offset_ = 0; } - -private: - int64_t size_; - int64_t offset_; - at::Tensor buffer_; - at::ScalarType scalar_type_; -}; - -class DoubleBufferedReduceBucket { -public: - DoubleBufferedReduceBucket(int64_t initial_bucket_size, bool enable_double_buffer) - : initial_bucket_size_(initial_bucket_size), enable_double_buffer_(enable_double_buffer) - { - } - - void swap(at::ScalarType scalar_type, - at::cuda::CUDAStream rs_stream, - at::cuda::CUDAStream copy_stream) - { - assert(hasKey(current_buffer_, scalar_type)); - assert(hasKey(current_buffer_events_, scalar_type)); - - current_buffer_.at(scalar_type)->reset(); - current_buffer_events_.at(scalar_type)->record(rs_stream); - - if (enable_double_buffer_) { - assert(hasKey(shadow_buffer_, scalar_type)); - assert(hasKey(shadow_buffer_events_, scalar_type)); - - auto tmp = current_buffer_.at(scalar_type); - current_buffer_[scalar_type] = shadow_buffer_.at(scalar_type); - shadow_buffer_[scalar_type] = tmp; - - auto tmp_event = current_buffer_events_.at(scalar_type); - current_buffer_events_[scalar_type] = shadow_buffer_events_.at(scalar_type); - shadow_buffer_events_[scalar_type] = tmp_event; - } - } - - std::shared_ptr getBuffer(at::ScalarType scalar_type) - { - if (!hasKey(current_buffer_, scalar_type)) { - current_buffer_[scalar_type] = - std::make_shared(initial_bucket_size_, scalar_type); - current_buffer_events_[scalar_type] = - std::make_shared(cudaEventDisableTiming); - - if (enable_double_buffer_) { - shadow_buffer_[scalar_type] = - std::make_shared(initial_bucket_size_, scalar_type); - shadow_buffer_events_[scalar_type] = - std::make_shared(cudaEventDisableTiming); - } - } - - return current_buffer_.at(scalar_type); - } - - std::shared_ptr getEvent(at::ScalarType scalar_type) - { - assert(hasKey(current_buffer_events_, scalar_type)); - return current_buffer_events_.at(scalar_type); - } - - void clear() - { - current_buffer_.clear(); - shadow_buffer_.clear(); - current_buffer_events_.clear(); - shadow_buffer_events_.clear(); - } - -private: - int64_t initial_bucket_size_; - bool enable_double_buffer_; - std::unordered_map> current_buffer_; - std::unordered_map> shadow_buffer_; - std::unordered_map> current_buffer_events_; - std::unordered_map> shadow_buffer_events_; -}; - -class DSParam { -public: - DSParam(long id, - std::vector ds_shape, - at::Tensor ds_tensor, - at::Tensor grad_buffer, - bool partitioned, - int64_t offset, // for Z1 - bool persistent // for Z3 - ) - : id_(id), - shape_(std::move(ds_shape)), - ds_tensor_(ds_tensor), - grad_buffer_(grad_buffer), - partitioned_(partitioned), - offset_(offset), - persistent_(persistent), - offload_stream_(at::cuda::getStreamFromPool()), - reload_stream_(at::cuda::getStreamFromPool()) - { - } - - long getId() const { return id_; } - std::vector getShape() const { return shape_; } - at::Tensor getDSTensor() const - { - // If the reload event exists and is complete, return the reloaded tensor (if defined) - if (reload_done_event_) { - if (!reload_done_event_->query()) { - reload_done_event_->block(at::cuda::getCurrentCUDAStream()); - } - if (ds_reload_tensor_.defined()) { return ds_reload_tensor_; } - } - // Otherwise, if an offload event exists, wait for it to complete - if (offload_done_event_) { - if (!offload_done_event_->query()) { - offload_done_event_->block(at::cuda::getCurrentCUDAStream()); - } - } - return ds_tensor_; - } - at::Tensor getGradBuffer() const { return grad_buffer_; } - bool isPartitioned() const { return partitioned_; } - int64_t getOffset() const { return offset_; } - void setPersistent(bool persistent) { persistent_ = persistent; } - bool isPersistent() const { return persistent_; } - - void offload() - { - // If a reloaded tensor exists, offload its data back to ds_tensor_ - if (ds_reload_tensor_.defined()) { - auto comp_stream = at::cuda::getCurrentCUDAStream(); - comp_done_event_ = std::make_shared(cudaEventDisableTiming); - // Record completion and wait on the offload stream - comp_done_event_->record(comp_stream); - comp_done_event_->block(offload_stream_); - offload_done_event_ = std::make_shared(cudaEventDisableTiming); - - { - at::cuda::CUDAStreamGuard guard(offload_stream_); - ds_tensor_.copy_(ds_reload_tensor_, /*non_blocking=*/true); - ds_reload_tensor_.reset(); // Clear the reloaded tensor - offload_done_event_->record(offload_stream_); - } - // Reset the reload event to indicate that no valid reload is present. - if (reload_done_event_) { reload_done_event_.reset(); } - } - } - - void reload() - { - // Reload only if the current ds_tensor_ is on CPU - if (ds_tensor_.device().is_cpu()) { - auto comp_stream = at::cuda::getCurrentCUDAStream(); - comp_done_event_ = std::make_shared(cudaEventDisableTiming); - // Record and wait on the reload stream - comp_done_event_->record(comp_stream); - comp_done_event_->block(reload_stream_); - reload_done_event_ = std::make_shared(cudaEventDisableTiming); - - { - at::cuda::CUDAStreamGuard guard(reload_stream_); - ds_reload_tensor_ = - at::empty_like(ds_tensor_, ds_tensor_.options().device(torch::kCUDA)); - ds_reload_tensor_.copy_(ds_tensor_, /*non_blocking=*/true); - reload_done_event_->record(reload_stream_); - } - // Reset offload_done_event if it exists to clear any stale offload state. - if (offload_done_event_) { offload_done_event_.reset(); } - } - } - -private: - long id_; - std::vector shape_; - at::Tensor ds_tensor_; - at::Tensor ds_reload_tensor_; - at::Tensor grad_buffer_; - bool partitioned_; - int64_t offset_; // for Z1 - bool persistent_; // for Z3 - mutable bool is_reloaded = false; - - at::cuda::CUDAStream offload_stream_; - at::cuda::CUDAStream reload_stream_; - std::shared_ptr comp_done_event_; - std::shared_ptr offload_done_event_; - std::shared_ptr reload_done_event_; -}; - -class DSParamRegistry { -public: - DSParamRegistry() {} - ~DSParamRegistry() {} - - void registerParam(long ds_id, - const std::vector& ds_shape, - at::Tensor ds_tensor, - at::Tensor grad_buffer, - bool partitioned, - int64_t offset, // for Z1 - bool persistent // for Z3 - ) - { - grad_buffer.zero_(); - params_.emplace( - ds_id, - DSParam(ds_id, ds_shape, ds_tensor, grad_buffer, partitioned, offset, persistent)); - valid_[ds_id] = false; - } - - void registerGatheredParam(long ds_id, at::Tensor ds_tensor) - { - gathered_params_.emplace(ds_id, ds_tensor); - } - - void unregisterGatheredParam(long ds_id) - { - assert(hasKey(gathered_params_, ds_id)); - gathered_params_.erase(ds_id); - valid_[ds_id] = false; - } - - const std::unordered_map& getParams() const { return params_; } - - const DSParam& getParam(long ds_id) const { return params_.at(ds_id); } - const size_t getNumParams() const { return params_.size(); } - const at::Tensor& getGatheredParam(long ds_id) const - { - assert(hasKey(gathered_params_, ds_id)); - return gathered_params_.at(ds_id); - } - bool hasGatheredParam(long ds_id) const { return hasKey(gathered_params_, ds_id); } - void setPersistent(long ds_id, bool persistent) { params_.at(ds_id).setPersistent(persistent); } - void offload(long ds_id) { params_.at(ds_id).offload(); } - void reload(long ds_id) { params_.at(ds_id).reload(); } - - void setValid(long ds_id, bool valid) { valid_[ds_id] = valid; } - bool isValid(long ds_id) const - { - assert(hasKey(valid_, ds_id)); - return valid_.at(ds_id); - } - -private: - std::unordered_map params_; - std::unordered_map gathered_params_; - std::unordered_map valid_; -}; - -class CustomOpExecutor { -public: - CustomOpExecutor(c10::intrusive_ptr process_group, - std::shared_ptr param_registry, - std::shared_ptr reduce_buckets, - std::vector ds_ids, - ncclComm_t nccl_comm, - at::cuda::CUDAStream rs_stream, - at::cuda::CUDAStream copy_stream, - bool pre_div_reduce) - : process_group_(process_group), - param_registry_(std::move(param_registry)), - reduce_buckets_(std::move(reduce_buckets)), - ds_ids_(std::move(ds_ids)), - nccl_comm_(nccl_comm), - rs_stream_(rs_stream), - copy_stream_(copy_stream), - pre_div_reduce_(pre_div_reduce) - { - for (long ds_id : ds_ids_) { - has_acc_grad_[ds_id] = false; - - rs_comp_done_events_[ds_id] = - std::make_shared(cudaEventDisableTiming); - rs_copy_done_events_[ds_id] = - std::make_shared(cudaEventDisableTiming); - } - reduce_counter_ = ds_ids_.size(); - } - ~CustomOpExecutor() {} - - virtual void startForward() {} - - virtual void endForward() {} - - virtual void startBackward(bool update) { param_updated_ = update; } - - virtual void endBackward() {} - - virtual at::Tensor reduceGrad(at::Tensor grad_tensor, long ds_id) - { - int world_size = process_group_->getSize(); - const DSParam& param = param_registry_->getParam(ds_id); - const auto scalar_type = grad_tensor.scalar_type(); - std::shared_ptr reduce_bucket = reduce_buckets_->getBuffer(scalar_type); - - auto comp_stream = at::cuda::getCurrentCUDAStream(); - - if (reduce_bucket->shouldFlush(grad_tensor.numel())) { - int rank = process_group_->getRank(); - - flushReduceBucket(scalar_type); - - // reduce_bucket is swapped in flushReduceBucket if double buffering is enabled - reduce_bucket = reduce_buckets_->getBuffer(scalar_type); - } - - if (grad_tensor.numel() > reduce_bucket->getSize()) { - // extend buckets - at::cuda::stream_synchronize(rs_stream_); - reduce_bucket->reserve(grad_tensor.numel()); - } - - at::Tensor reduce_in_buffer = reduce_bucket->allocate(grad_tensor.numel()); - - // This ensures the order of reduce_scatter -> copy - // Without this block, copy may start while reduce_scatter is still running - reduce_buckets_->getEvent(scalar_type)->block(comp_stream); - auto copy_src = grad_tensor.contiguous().view({-1}).detach(); - // keep references to copy src - reduce_tasks_[scalar_type].emplace_back(ds_id, copy_src, reduce_in_buffer); - - // computation must be done before copy - rs_comp_done_events_[ds_id]->record(comp_stream); - rs_comp_done_events_[ds_id]->block(copy_stream_); - { - at::cuda::CUDAStreamGuard guard(copy_stream_); - reduce_in_buffer.copy_(copy_src, true); - rs_copy_done_events_[ds_id]->record(copy_stream_); - } - - reduce_counter_--; - - if (reduce_counter_ == 0) { - flushAllReduceBuckets(); - - reduce_counter_ = ds_ids_.size(); - - // This synchronization ensures all of reduce calls are done before optimizer's step. - at::cuda::stream_synchronize(rs_stream_); - - endBackward(); - } - - return at::Tensor(); - } - - bool hasParam(long ds_id) const { return hasKey(has_acc_grad_, ds_id); } - -protected: - c10::intrusive_ptr process_group_; - std::shared_ptr param_registry_; - std::shared_ptr reduce_buckets_; - std::vector ds_ids_; - ncclComm_t nccl_comm_; - at::cuda::CUDAStream rs_stream_; - at::cuda::CUDAStream copy_stream_; - - std::unordered_map> rs_comp_done_events_; - std::unordered_map> rs_copy_done_events_; - - size_t reduce_counter_ = 0; - bool param_updated_ = false; - std::unordered_map> reduce_tasks_; - std::unordered_map has_acc_grad_; - bool pre_div_reduce_; - - virtual void flushReduceBucket(at::ScalarType scalar_type) = 0; - - void flushAllReduceBuckets() - { - for (const auto& it : reduce_tasks_) { flushReduceBucket(it.first); } - } - - // Common helper methods for flushReduceBucket implementations - void blockCopyEvents(at::ScalarType scalar_type) - { - for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { - auto copy_done_event = rs_copy_done_events_.at(t.getDSId()); - copy_done_event->block(rs_stream_); - } - } - - void applyPreDivision(at::ScalarType scalar_type) - { - if (pre_div_reduce_) { - at::cuda::CUDAStreamGuard guard(rs_stream_); - for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { - t.getSendBuf().div_(process_group_->getSize()); - } - } - } - - ncclRedOp_t getReductionOp() const { return pre_div_reduce_ ? ncclSum : ncclAvg; } - - void performCleanup(at::ScalarType scalar_type) - { - reduce_buckets_->swap(scalar_type, rs_stream_, copy_stream_); - - // Prevent grad tensor from being released before the copy is done - auto comp_stream = at::cuda::getCurrentCUDAStream(); - for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { - auto copy_done_event = rs_copy_done_events_.at(t.getDSId()); - copy_done_event->block(comp_stream); - } - reduce_tasks_[scalar_type].clear(); - } -}; - -template -std::shared_ptr getExecutor(long graph_id, - const std::unordered_map>& executors) -{ - assert(hasKey(executors, graph_id)); - if (auto executor = std::dynamic_pointer_cast(executors.at(graph_id))) { return executor; } - throw std::runtime_error("Invalid executor type"); -} - -extern std::shared_ptr param_registry; -extern std::unordered_map> executors; -extern std::shared_ptr reduce_buckets; - -at::Tensor reduce_grad(at::Tensor grad_tensor, long graph_id, long ds_id); -at::Tensor reduce_grad_meta(at::Tensor grad_tensor, long graph_id, long ds_id); -void free_tensors(std::vector tensors); -void free_tensors_meta(std::vector tensors); - -void init(c10::intrusive_ptr pg, - int64_t initial_reduce_bucket_size, - bool enable_double_buffer, - bool _use_symm_mem, - bool _clone_custom_op_output, - bool _sync_before_reduce, - bool _sync_after_reduce, - bool _sync_before_allgather, - bool _sync_after_allgather); -void reset(); -void cleanup(); - -void start_forward(); -void end_forward(); -void start_backward(bool update); - -} // namespace dc diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/dequantization_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/dequantization_utils.h deleted file mode 100644 index 37a7cdf27a014a4eae42b9400ae0c2a12132c449..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/dequantization_utils.h +++ /dev/null @@ -1,177 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "conversion_utils.h" -#include "ds_kernel_utils.h" -#include "quantization.h" -#include "quantization_utils.h" - -namespace cg = cooperative_groups; - -#pragma once - -namespace dequantize { -using Type = quantize::Type; - -template -using Params = quantize::Params; - -constexpr int granularity = quantize::granularity; -using PackedInt4 = quantize::PackedInt4; - -constexpr int h_per_chunk = granularity / sizeof(__half); -constexpr int h2_per_chunk = granularity / sizeof(__half2); - -/* -Device function that reads quantized data from global memory, dequantizes -it, and stores it to global memory. -Template Arguments : - numBits - Number of bits in quantized element. int: 4, 8 - qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric - unroll - Number of load steps to internally unroll int - threads - Number of threads to perform dequant int -Function arguments: - global_output - __half pointer in global memory - data - Quantized data in global memory - global_params - Quantization parameters in global memory - elems_per_group - Number of elements in each quantization group - total_elems - Tensor size (note, does not need to be multiple of elems_per_group) -*/ -template -DS_D_INLINE void to_global(__half* global_output, - const int8_t* data, - const float* global_params, - const int elems_per_group, - const int total_elems); - -/* -Device function that quantizes 16 bytes of __half type input data. -Template Arguments : - numBits - Number of bits in quantized element. int : 8 or 4 - qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric -Function Arguments : - local_output - Local array to store dequantized data __half* or __half2* - data - Pointer to quantized input data. int8_t* - Params - Parameters for quantization. Params -*/ -template -DS_D_INLINE void chunk(__half2* local_output, const int8_t* data, Params q_params); - -template -DS_D_INLINE void chunk(T* local_output, const int8_t* data, Params q_params); - -/**************** Implementations ******************/ - -template -DS_D_INLINE void chunk(T* local_output, const int8_t* data, Params q_params) -{ - constexpr int32_t num_elems_packed = 8 / numBits; - constexpr int32_t iters = h_per_chunk / num_elems_packed; - -#pragma unroll - for (int i = 0; i < iters; i++) { - if constexpr (num_elems_packed == 1) { - local_output[i] = q_params.template dequantize(data[i]); - } else { - auto accessible_data = *(PackedInt4*)(&data[i]); - local_output[2 * i] = q_params.template dequantize(accessible_data.low); - local_output[2 * i + 1] = q_params.template dequantize(accessible_data.high); - } - } -} - -template -DS_D_INLINE void chunk(__half2* local_output, const int8_t* data, Params q_params) -{ - __half* local_output_cast = reinterpret_cast<__half*>(local_output); - chunk<__half, numBits>(local_output_cast, data, q_params); -} - -template -DS_D_INLINE void _to_global(T* global_output, - const int8_t* data, - const float* global_params, - const int elems_per_group, - const int total_elems) -{ - cg::thread_block tb = cg::this_thread_block(); - cg::thread_block_tile warp = cg::tiled_partition(tb); - - // Load constants - // TODO(cmikeh2): Refactor into functions? - constexpr int load_granularity = (granularity / (sizeof(T))) / (numBits == 8 ? 1 : 2); - constexpr int load_step_stride = load_granularity * threads; - constexpr int load_block_stride = load_step_stride * unroll; - - // Store constants - constexpr int T_per_chunk = granularity / sizeof(T); - constexpr int store_step_stride = T_per_chunk * threads; - constexpr int store_block_stride = store_step_stride * unroll; - - // Load offsets - const int load_block_offset = tb.group_index().x * load_block_stride; - // Note: we can use `load_granularity` since the dtype is `int8_t`. - const int load_thread_offset = tb.thread_index().x * load_granularity; - const int8_t* load_base = data + load_block_offset + load_thread_offset; - - // Store offsets - const int store_block_offset = tb.group_index().x * store_block_stride; - const int store_thread_offset = tb.thread_index().x * T_per_chunk; - const int elem_id_base = store_block_offset + store_thread_offset; - - int8_t local_load_buffer[load_granularity * unroll]; - T local_dequant_buffer[T_per_chunk * unroll]; - - /* - Note: Splitting this loop in half gave about 3-5% performance increase for reasons that aren't - totally clear to me, so this is a deliberately weird code structure. - */ -#pragma unroll - for (int i = 0; i < unroll; i++) { - const int elem_id_iter = elem_id_base + i * store_step_stride; - - if (elem_id_iter < total_elems) { - mem_access::load_global(local_load_buffer + i * load_granularity, - load_base + i * load_step_stride); - } - } - -#pragma unroll - for (int i = 0; i < unroll; i++) { - const int elem_id_iter = elem_id_base + i * store_step_stride; - if (elem_id_iter < total_elems) { - // TODO(cmikeh2): Can we amortize this division? Perform once on the first iteration and - // use indexing math to do division free interpolation of the successive groups? - const int group_index = elem_id_iter / elems_per_group; - Params q_params(global_params, group_index); - - chunk(local_dequant_buffer + i * T_per_chunk, - local_load_buffer + i * load_granularity, - q_params); - mem_access::store_global(global_output + elem_id_iter, - local_dequant_buffer + i * T_per_chunk); - } - } -} - -template -DS_D_INLINE void to_global(T* global_output, - const int8_t* data, - const float* global_params, - const int elems_per_group, - const int total_elems) -{ - if constexpr (numBits == 4 || numBits == 8) { - _to_global( - global_output, data, global_params, elems_per_group, total_elems); - } else if constexpr (numBits == 3) { - // TODO(cmikeh2): Need this implementation - assert(false); - } else { - assert(false); - } -} - -} // namespace dequantize diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/dropout.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/dropout.h deleted file mode 100644 index c00943a97687c3514ff96f63b4b0540805e9b545..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/dropout.h +++ /dev/null @@ -1,81 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include -#include - -template -class Dropout { -public: - struct Config { - float ratio; - uint32_t dim; - bool training; - - Config(float r, uint32_t d) : ratio(r), dim(d), training(true) {} - - float RATIO() const { return training ? ratio : 0.0; } - inline void SetDim(uint32_t d) { dim = d; } - }; - - Dropout(const Config& config) : _config(config), _mask(nullptr) {} - - virtual ~Dropout() {} - - void Forward(int bsz, T* out, const T* vals, cudaStream_t stream, bool bwd = false) - { - launch_dropout( - out, vals, _mask, bsz * _config.dim, _config.dim, _config.RATIO(), stream, bwd); - } - - void ForwardWithBias(int bsz, T* vals, const T* bias, cudaStream_t stream) - { - launch_dropout(vals, bias, _mask, bsz, _config.dim, _config.RATIO(), stream); - } - - void ForwardWithBias(int bsz, - T* out, - const T* vals, - const T* residual, - const T* bias, - cudaStream_t stream) - { - launch_dropout( - out, vals, residual, bias, _mask, bsz, _config.dim, _config.RATIO(), stream); - } - - void Backward(int bsz, T* d_vals, cudaStream_t stream) - { - launch_dropout_grad(d_vals, _mask, bsz * _config.dim, _config.RATIO(), stream); - } - - void Backward(int bsz, T* d_vals_out, const T* d_vals, cudaStream_t stream) - { - launch_dropout_grad( - d_vals_out, d_vals, _mask, bsz * _config.dim, _config.RATIO(), stream); - } - - bool HasDropout() const { return _config.RATIO() > 0.0; } - - void SetTrainingMode(bool training) { _config.training = training; } - - void SetMask(uint8_t* mask) - { - if (!mask) { throw std::runtime_error("Dropout mask is null."); } - - _mask = mask; - } - - Config GetConfig() const { return _config; } - - inline void SetDimension(uint32_t dim) { _config.SetDim(dim); } - -private: - uint8_t* _mask; - Config _config; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/ds_kernel_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/ds_kernel_utils.h deleted file mode 100644 index f8b16ee6a315f4c2ff5ca72fab077c0534c070e3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/ds_kernel_utils.h +++ /dev/null @@ -1,58 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Centralized header file for preprocessor macros and constants -used throughout the codebase. -*/ - -#pragma once - -#include -#include - -#ifdef BF16_AVAILABLE -#include -#endif - -#define DS_HD_INLINE __host__ __device__ __forceinline__ -#define DS_D_INLINE __device__ __forceinline__ - -#ifdef __HIP_PLATFORM_AMD__ - -// constexpr variant of warpSize for templating -constexpr int hw_warp_size = ROCM_WAVEFRONT_SIZE; -#define HALF_PRECISION_AVAILABLE = 1 -#include -#include - -#else // !__HIP_PLATFORM_AMD__ - -// constexpr variant of warpSize for templating -constexpr int hw_warp_size = 32; - -#if __CUDA_ARCH__ >= 530 -#define HALF_PRECISION_AVAILABLE = 1 -#define PTX_AVAILABLE -#endif // __CUDA_ARCH__ >= 530 - -#if __CUDA_ARCH__ >= 800 -#define ASYNC_COPY_AVAILABLE -#endif // __CUDA_ARCH__ >= 800 - -#include -#include - -#endif //__HIP_PLATFORM_AMD__ - -inline int next_pow2(const int val) -{ - int rounded_val = val - 1; - rounded_val |= rounded_val >> 1; - rounded_val |= rounded_val >> 2; - rounded_val |= rounded_val >> 4; - rounded_val |= rounded_val >> 8; - return rounded_val + 1; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/ds_transformer_cuda.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/ds_transformer_cuda.h deleted file mode 100644 index 4dbd1d4773385db8e40a9e2f0ac8d6ebf2433c7a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/ds_transformer_cuda.h +++ /dev/null @@ -1,189 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include -#include -#include -#include "cublas_v2.h" -#include "cuda.h" -#include "dropout.h" -#include "feed_forward.h" -#include "gelu.h" -#include "general_kernels.h" -#include "normalize_layer.h" -#include "softmax.h" -#include "strided_batch_gemm.h" - -struct BertGemmAlgos { - int m_gemm_qkv_algo; - int m_gemm_inter_algo; - int m_gemm_output_algo; - int m_gemm_batch1_algo; - int m_gemm_batch2_algo; - - BertGemmAlgos() - : m_gemm_qkv_algo(-1), - m_gemm_inter_algo(-1), - m_gemm_output_algo(-1), - m_gemm_batch1_algo(-1), - m_gemm_batch2_algo(-1) - { - } -}; - -template -class BertTransformerLayer { -public: - BertTransformerLayer(unsigned layer_id, - unsigned batch_size, - unsigned hidden_size, - unsigned num_heads, - unsigned intermediate_size, - unsigned seq_length, - float attn_dropout_ratio, - float hidden_output_dropout_ratio, - float layer_norm_eps, - bool pre_or_postLayerNorm, - const std::vector>& gemm_algos, - bool attn_dropout_checkpoint, - bool normalize_invertible, - bool gelu_checkpoint, - bool stochastic_mode); - - virtual ~BertTransformerLayer(); - - void Forward(unsigned bsz, - const T* input_ptr, - const T* input_mask_ptr, - const T* attn_qkvw_ptr, - const T* attn_qkvb_ptr, - const T* attn_ow_ptr, - const T* attn_ob_ptr, - const T* attn_nw_ptr, - const T* attn_nb_ptr, - const T* inter_w_ptr, - const T* inter_b_ptr, - const T* output_w_ptr, - const T* output_b_ptr, - const T* norm_w_ptr, - const T* norm_b_ptr, - T* out_ptr, - T* inp_norm_ptr, - T* q_tf_ptr, - T* k_tf_ptr, - T* v_tf_ptr, - T* softmax_output_ptr, - T* ctx_bufB_ptr, - T* attn_o_inp_ptr, - T* add_res_ptr, - T* ff1_inp_ptr, - T* gelu_inp_ptr, - T* ff2_inp_ptr); - - void Backward(unsigned bsz, - const T* grad_output_ptr, - const T* input_ptr, - const T* output_ptr, - const T* inp_norm_ptr, - const T* q_tf_ptr, - const T* k_tf_ptr, - const T* v_tf_ptr, - const T* softmax_output_ptr, - const T* ctx_bufB_ptr, - const T* attn_o_inp_ptr, - const T* add_res_ptr, - const T* ff1_inp_ptr, - const T* gelu_inp_ptr, - const T* ff2_inp_ptr, - const T* input_mask_ptr, - const T* attn_qkvw_ptr, - const T* attn_ow_ptr, - const T* attn_nw_ptr, - const T* attn_nb_ptr, - const T* inter_w_ptr, - const T* inter_b_ptr, - const T* output_w_ptr, - const T* norm_w_ptr, - const T* norm_b_ptr, - - T* grad_input_ptr, - T* grad_attn_qkvw_ptr, - T* grad_attn_qkvb_ptr, - T* grad_attn_ow_ptr, - T* grad_attn_ob_ptr, - T* grad_attn_nw_ptr, - T* grad_attn_nb_ptr, - T* grad_inter_w_ptr, - T* grad_inter_b_ptr, - T* grad_output_w_ptr, - T* grad_output_b_ptr, - T* grad_norm_w_ptr, - T* grad_norm_b_ptr); - - void SetIntermediateBuffers(uint8_t* attn_prob_dropout_mask_ptr, - uint8_t* attn_output_dropout_mask_ptr, - uint8_t* layer_output_dropout_mask_ptr, - T* layer_norm_var, - T* layer_norm_mean, - T* attn_layer_norm_var, - T* attn_layer_norm_mean); - - inline unsigned GetBatchSize() const { return _batch_size; } - inline unsigned GetNumHeads() const { return _heads; } - inline unsigned GetSeqLength() const { return _seq_length; } - inline unsigned GetIntermediateSize() const { return _intermediate_size; } - - void SetSeqLength(unsigned seq_len); - inline unsigned GetHiddenSize() const { return _hidden_size; } - void SetTrainingMode(bool training); - inline bool IsTrainingMode() const { return _training; } - inline bool GeluCheckpoint() const { return _gelu_checkpoint; } - -private: - void Initialize(); - size_t getWorkspaceSize(int maxBatchSize) const; - - // Params - unsigned _layer_id; - unsigned _batch_size; - unsigned _hidden_size; - unsigned _heads; - unsigned _size_per_head; - unsigned _intermediate_size; - unsigned _seq_length; - - bool _pre_or_postLayerNorm; - - cublasHandle_t _cublasHandle; - cudaStream_t _stream; - - // layers - FeedForward _qkv_linear; - FeedForward _attn_out_linear; - Normalize_Layer _attn_layer_norm; - Normalize_Layer _layer_norm; - Normalize_Layer* _last_normalize; - FeedForward _ff1, _ff2; - Softmax _softmax; - Gelu _gelu; - Dropout _attn_prob_dropout; - Dropout _attn_output_dropout; - Dropout _layer_output_dropout; - StridedBatchGemm _attn_scores; - StridedBatchGemm _attn_context; - - bool _training; - - // Memory saving flags - bool _attn_dropout_checkpoint; - bool _normalize_invertible; - bool _gelu_checkpoint; - - // High Performance flags - bool _stochastic_mode; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/feed_forward.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/feed_forward.h deleted file mode 100644 index d2056403d265e4083efd9da8f7ee9c887e0d9afc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/feed_forward.h +++ /dev/null @@ -1,114 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#ifndef __FEEDFORWARD_H__ -#define __FEEDFORWARD_H__ - -#include -#include -#include -#include "custom_cuda_layers.h" - -template -class FeedForward { -public: - struct Config { - int batchSize, outputSize; - int inputSize; - std::array gemm_algos; - Config(int batch, int outputs, int inputs, const std::array& algos) - : batchSize(batch), outputSize(outputs), inputSize(inputs), gemm_algos(algos) - { - } - }; - - FeedForward(Config config) : config_(config) {} - - ~FeedForward() {} - - void Forward(int bsz, - const T* input_ptr, - const T* weights, - T* out, - cublasHandle_t& _cublasHandle) - { - float alpha = T(1.); - float beta = T(0.); - - cublas_gemm_ex(_cublasHandle, - CUBLAS_OP_T, - CUBLAS_OP_N, - config_.outputSize, - bsz, - config_.inputSize, - &alpha, - &beta, - weights, - input_ptr, - out, -// TODO HIP: Remove backward compatibility for torch<=2.0 in future -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo(config_.gemm_algos[0])); -#else - cublasGemmAlgo_t(config_.gemm_algos[0])); -#endif - } - void Backward(int bsz, - const T* out_grad, - const T* input_ptr, - const T* weights, - T* weights_grad, - T* bias_grad, - cublasHandle_t& _cublasHandle, - cudaStream_t& stream, - T* inp_grad_out = nullptr, - T* out_grad_trans_out = nullptr) - { - float alpha = (T)1.0, beta = (T)0.0; - cublas_gemm_ex(_cublasHandle, - CUBLAS_OP_N, - CUBLAS_OP_T, - config_.inputSize, - config_.outputSize, - bsz, - &alpha, - &beta, - input_ptr, - out_grad, - weights_grad, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo(config_.gemm_algos[1])); -#else - cublasGemmAlgo_t(config_.gemm_algos[1])); -#endif - - cublas_gemm_ex(_cublasHandle, - CUBLAS_OP_N, - CUBLAS_OP_N, - config_.inputSize, - bsz, - config_.outputSize, - &alpha, - &beta, - weights, - out_grad, - inp_grad_out, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo(config_.gemm_algos[2])); -#else - cublasGemmAlgo_t(config_.gemm_algos[2])); -#endif - - launch_fuse_transpose_bias_kernel(out_grad, bias_grad, bsz, config_.outputSize, stream); - } - -private: - Config config_; -}; - -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/gelu.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/gelu.h deleted file mode 100644 index b75c1510bad65101604dff2aaebfb79d5d435531..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/gelu.h +++ /dev/null @@ -1,41 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include -#include -#include "custom_cuda_layers.h" - -template -class Gelu { -public: - struct Config { - uint32_t intermediate_size; - Config(uint32_t inter_size) : intermediate_size(inter_size) {} - }; - - Gelu(const Config& config) : _config(config) {} - - virtual ~Gelu() {} - - void ForwardWithBiasAdd(int bsz, - const T* input_buf, - const T* bias, - T* output, - cudaStream_t stream) - { - launch_bias_gelu(input_buf, bias, output, _config.intermediate_size, bsz, stream); - } - - void Backward(int bsz, T* d_output, const T* input_buf, const T* bias, cudaStream_t stream) - { - launch_d_gelu(d_output, input_buf, bias, _config.intermediate_size, bsz, stream); - } - -private: - Config _config; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/gemm_test.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/gemm_test.h deleted file mode 100644 index de5b55cd3df16f211af75fde895a49b3871e0331..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/gemm_test.h +++ /dev/null @@ -1,350 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#ifndef __HIP_PLATFORM_AMD__ -#include -#endif -#ifdef __HIP_PLATFORM_AMD__ -#include -#endif -#include -#include -#include -#include -#include -#include -#include "StopWatch.h" -#include "cublas_wrappers.h" - -template -void check(T result, char const* const func, const char* const file, int const line) -{ - if (result) { - std::cout << (std::string("CUDA runtime error: ") + +file + ":" + std::to_string(line) + - " \n"); - } -} - -#define check_cuda_error(val) check((val), #val, __FILE__, __LINE__) - -template -class GemmTest { -public: - GemmTest(int m, int n, int k, cublasOperation_t ta, cublasOperation_t tb, cublasHandle_t h) - : M(m), N(n), K(k), transa(ta), transb(tb), handle(h) - { - check_cuda_error(cudaMalloc((void**)&A, sizeof(T) * M * K)); - check_cuda_error(cudaMalloc((void**)&B, sizeof(T) * K * N)); - check_cuda_error(cudaMalloc((void**)&C, sizeof(T) * M * N)); - } - - ~GemmTest() - { - check_cuda_error(cudaFree(A)); - check_cuda_error(cudaFree(B)); - check_cuda_error(cudaFree(C)); - } - - std::array TestAlgo(int loops) - { - float alpha = (T)1.0f; - float beta = (T)0.0f; - - int algo_fw = Run(loops, [=](int algo) { - cublas_gemm_ex(handle, - CUBLAS_OP_T, - CUBLAS_OP_N, - N, - M, - K, - &alpha, - &beta, - B, - A, - C, -// TODO HIP: Remove backward compatibility for torch<=2.0 in future -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - static_cast(algo)); -#else - static_cast(algo)); -#endif - }); - - int algo_bw1 = Run(loops, [=](int algo) { - cublas_gemm_ex(handle, - CUBLAS_OP_N, - CUBLAS_OP_T, - K, - N, - M, - &alpha, - &beta, - A, - C, - B, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - static_cast(algo)); -#else - static_cast(algo)); -#endif - }); - - int algo_bw2 = Run(loops, [=](int algo) { - cublas_gemm_ex(handle, - CUBLAS_OP_N, - CUBLAS_OP_N, - K, - M, - N, - &alpha, - &beta, - B, - C, - A, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - static_cast(algo)); -#else - static_cast(algo)); -#endif - }); - - return std::array({algo_fw, algo_bw1, algo_bw2}); - } - - template - int Run(int loops, Func f) - { - float fast_latency = (std::numeric_limits::max)(); - int fast_algo = 0; - -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - for (int algo = (int)rocblas_gemm_algo_standard; algo <= (int)rocblas_gemm_algo_standard; -#elif defined(__HIP_PLATFORM_AMD__) - for (int algo = (int)HIPBLAS_GEMM_DEFAULT; algo <= (int)HIPBLAS_GEMM_DEFAULT; -#else - for (int algo = (int)CUBLAS_GEMM_DEFAULT_TENSOR_OP; - algo <= (int)CUBLAS_GEMM_ALGO15_TENSOR_OP; -#endif - algo++) { - int warm_up = 5; - for (int i = 0; i < warm_up; ++i) f(algo); - - cudaDeviceSynchronize(); - Stopwatch timer; - timer.Restart(); - - for (int i = 0; i < loops; ++i) f(algo); - - cudaDeviceSynchronize(); - timer.Stop(); - - float avg_latency = (float)timer.GetTimeInSeconds() * 1000 / loops; - - printf("algo-%d: %.3fms\n", algo, avg_latency); - - if (avg_latency < fast_latency) { - fast_latency = avg_latency; - fast_algo = algo; - } - } - - printf("fast_algo %d: %.3f ms\n", fast_algo, fast_latency); - - return fast_algo; - } - -private: - int M, N, K; - cublasHandle_t handle; - cublasOperation_t transa, transb; - T *A, *B, *C; -}; - -template -class StridedGemmTest { -public: - StridedGemmTest(int b, - int m, - int n, - int k, - cublasOperation_t ta, - cublasOperation_t tb, - cublasHandle_t h) - : bsz(b), M(m), N(n), K(k), transa(ta), transb(tb), handle(h) - { - check_cuda_error(cudaMalloc((void**)&A, sizeof(T) * M * K * bsz)); - check_cuda_error(cudaMalloc((void**)&B, sizeof(T) * K * N * bsz)); - check_cuda_error(cudaMalloc((void**)&C, sizeof(T) * M * N * bsz)); - } - - ~StridedGemmTest() - { - check_cuda_error(cudaFree(A)); - check_cuda_error(cudaFree(B)); - check_cuda_error(cudaFree(C)); - } - - std::array TestAlgo(int loops) - { - float alpha = (T)1.0f; - float beta = (T)0.0f; - - int algo_fw = Run(loops, [=](int algo) { - int stride_a = M * K; - int stride_b = N * K; - int stride_c = M * N; - - cublas_strided_batched_gemm(handle, - M, - N, - K, - &alpha, - &beta, - A, - B, - C, - transa, - transb, - stride_a, - stride_b, - stride_c, - bsz, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - static_cast(algo)); -#else - static_cast(algo)); -#endif - }); - - int algo_bw1 = Run(loops, [=](int algo) { - int mb = (transa == CUBLAS_OP_T ? K : M); - int kb = (transa == CUBLAS_OP_T ? M : K); - - int stride_a = mb * N; - int stride_b = N * kb; - int stride_c = M * K; - - // B need to transpose. - cublasOperation_t op_b = (transb == CUBLAS_OP_T ? CUBLAS_OP_N : CUBLAS_OP_T); - - // Calculate d_A. - cublas_strided_batched_gemm(handle, - mb, - kb, - N, - &alpha, - &beta, - (transa == CUBLAS_OP_T ? B : C), - (transa == CUBLAS_OP_T ? C : B), - A, - CUBLAS_OP_N, - op_b, - stride_a, - stride_b, - stride_c, - bsz, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - static_cast(algo)); -#else - static_cast(algo)); -#endif - }); - - int algo_bw2 = Run(loops, [=](int algo) { - // A need to transpose. - cublasOperation_t op_a = (transa == CUBLAS_OP_T ? CUBLAS_OP_N : CUBLAS_OP_T); - - int stride_a = M * K; - int stride_b = M * N; - int stride_c = N * K; - - // Calculate d_B. - cublas_strided_batched_gemm(handle, - K, - N, - M, - &alpha, - &beta, - A, - C, - B, - op_a, - CUBLAS_OP_N, - stride_a, - stride_b, - stride_c, - bsz, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - static_cast(algo)); -#else - static_cast(algo)); -#endif - }); - - return std::array({algo_fw, algo_bw1, algo_bw2}); - } - - template - int Run(int loops, Func f) - { - float fast_latency = (std::numeric_limits::max)(); - int fast_algo = 0; - -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - for (int algo = (int)rocblas_gemm_algo_standard; algo <= (int)rocblas_gemm_algo_standard; -#else -#ifdef __HIP_PLATFORM_AMD__ - for (int algo = (int)CUBLAS_GEMM_DEFAULT_TENSOR_OP; - algo <= (int)CUBLAS_GEMM_DEFAULT_TENSOR_OP; -#else - for (int algo = (int)CUBLAS_GEMM_DEFAULT_TENSOR_OP; - algo <= (int)CUBLAS_GEMM_ALGO15_TENSOR_OP; -#endif -#endif - algo++) { - int warm_up = 5; - for (int i = 0; i < warm_up; ++i) f(algo); - - cudaDeviceSynchronize(); - Stopwatch timer; - timer.Restart(); - - for (int i = 0; i < loops; ++i) f(algo); - - cudaDeviceSynchronize(); - timer.Stop(); - - float avg_latency = (float)timer.GetTimeInSeconds() * 1000 / loops; - - printf("algo-%d: %.3fms\n", algo, avg_latency); - - if (avg_latency < fast_latency) { - fast_latency = avg_latency; - fast_algo = algo; - } - } - - printf("fast_algo %d: %.3f ms\n", fast_algo, fast_latency); - - return fast_algo; - } - -private: - int bsz, M, N, K; - cublasHandle_t handle; - cublasOperation_t transa, transb; - T *A, *B, *C; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/general_kernels.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/general_kernels.h deleted file mode 100644 index bd621d3c4329c1e12ba5aac25404faaa9d1bca3e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/general_kernels.h +++ /dev/null @@ -1,56 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include -#include -#include - -#ifdef __HIP_PLATFORM_AMD__ -#include -#else -#include -#endif -#include - -#include "context.h" -#include "cublas_wrappers.h" - -#define THREADS 256 -#define TILE_DIM 32 - -#define minus_infinity -1 * std::numeric_limits::infinity() - -#define FINAL_MASK 0xffffffff - -template -void launch_fused_add2(T* out, - const T* inp1, - const T* inp2, - int batch_size, - int seq_length, - int hidden_size, - cudaStream_t& stream); - -template -void launch_fused_add4(T* out, - const T* inp1, - const T* inp2, - const T* inp3, - const T* inp4, - int batch_size, - int seq_length, - int hidden_size, - cudaStream_t& stream); - -template -void launch_fused_add3(T* out, - const T* inp1, - const T* inp2, - const T* inp3, - int batch_size, - int seq_length, - int hidden_size, - cudaStream_t& stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/memory_access_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/memory_access_utils.h deleted file mode 100644 index cb990b75bbe8460176a431c4ea0c2573b742dffe..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/memory_access_utils.h +++ /dev/null @@ -1,1144 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include "ds_kernel_utils.h" - -/////////////////////////////// Memory Access Utils /////////////////////////////// -namespace mem_access { - -enum class LoadPolicy { - CacheAll, // Cache at all levels - CacheGlobal, // Cache at L2 only - CacheStreaming // Cache with evict first policy -}; - -enum class StorePolicy { - Writeback, // Cache in L1, write-back on eviction - CacheGlobal, // Bypass L1, write-back on eviction - CacheStreaming // Allocate cache line with evict first policy -}; - -template -__device__ __forceinline__ void load_global(void* dst, const void* src); - -template -__device__ __forceinline__ void load_global(void* dst, const void* src, bool do_access); - -// Shared accesses have no cache policy -template -__device__ __forceinline__ void load_shared(void* dst, const void* src); - -template -__device__ __forceinline__ void load_shared(void* dst, const void* src, bool do_access); - -template -__device__ __forceinline__ void store_global(void* dst, const void* src); - -// Shared accesses have no cache policy -template -__device__ __forceinline__ void store_shared(void* dst, const void* src); - -#ifdef ASYNC_COPY_AVAILABLE -template -__device__ __forceinline__ void memcpy_async(void* shr, const void* gbl); - -template -__device__ __forceinline__ void memcpy_async_nop(void* shr, const void* gbl, bool predicate); - -template -__device__ __forceinline__ void memcpy_async_zero(void* shr, const void* gbl, bool predicate); - -__device__ __forceinline__ void memcpy_async_fence(); - -template -__device__ __forceinline__ void memcpy_async_wait(); - -template -__device__ __forceinline__ void tail_complete_wait(int remaining_stages); -#endif - -// Util for tracking pipeline buffers -// TODO: Evaluate whether this should also be guarded by ASYNC_COPY_AVAILABLE -template -class BufferTracker { -public: - int current_state; - - __device__ __forceinline__ BufferTracker() : current_state(0) {} - - __device__ __forceinline__ int get() - { - int return_val = current_state++; - current_state = (current_state == max ? 0 : current_state); - return return_val; - } -}; - -__device__ __forceinline__ uint32_t lane_id() -{ -#ifdef PTX_AVAILABLE - unsigned int lane_id; - asm volatile("mov.u32 %0, %%laneid;" : "=r"(lane_id)); - return lane_id; -#else - return threadIdx.x & (warpSize - 1); // Portable -#endif -} - -/////////// Load Global /////////// -template <> -__device__ __forceinline__ void load_global<16>(void* dst, const void* src) -{ - uint4* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile("ld.global.ca.v4.u32 {%0, %1, %2, %3}, [%4];\n" - : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) - : "l"(src)); -#else - const uint4* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_global<16>(void* dst, const void* src, bool do_access) -{ - uint4* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %5, 0;\n" - "\tmov.b32 %0, 0;\n" - "\tmov.b32 %1, 0;\n" - "\tmov.b32 %2, 0;\n" - "\tmov.b32 %3, 0;\n" - "\t@p ld.global.v4.u32 {%0, %1, %2, %3}, [%4];\n" - "}\n" - : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) - : "l"(src), "r"((int)do_access)); -#else - const uint4* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0].x = 0; - data[0].y = 0; - data[0].z = 0; - data[0].w = 0; - } -#endif -} - -template <> -__device__ __forceinline__ void load_global<16, LoadPolicy::CacheGlobal>(void* dst, const void* src) -{ - uint4* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile("ld.global.cg.v4.u32 {%0, %1, %2, %3}, [%4];\n" - : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) - : "l"(src)); -#else - const uint4* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_global<16, LoadPolicy::CacheGlobal>(void* dst, - const void* src, - bool do_access) -{ - uint4* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %5, 0;\n" - "\tmov.b32 %0, 0;\n" - "\tmov.b32 %1, 0;\n" - "\tmov.b32 %2, 0;\n" - "\tmov.b32 %3, 0;\n" - "\t@p ld.global.cg.v4.u32 {%0, %1, %2, %3}, [%4];\n" - "}\n" - : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) - : "l"(src), "r"((int)do_access)); -#else - const uint4* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0].x = 0; - data[0].y = 0; - data[0].z = 0; - data[0].w = 0; - } -#endif -} - -template <> -__device__ __forceinline__ void load_global<16, LoadPolicy::CacheStreaming>(void* dst, - const void* src) -{ - uint4* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile("ld.global.cs.v4.u32 {%0, %1, %2, %3}, [%4];\n" - : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) - : "l"(src)); -#else - const uint4* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_global<16, LoadPolicy::CacheStreaming>(void* dst, - const void* src, - bool do_access) -{ - uint4* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %5, 0;\n" - "\tmov.b32 %0, 0;\n" - "\tmov.b32 %1, 0;\n" - "\tmov.b32 %2, 0;\n" - "\tmov.b32 %3, 0;\n" - "\t@p ld.global.cg.v4.u32 {%0, %1, %2, %3}, [%4];\n" - "}\n" - : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) - : "l"(src), "r"((int)do_access)); -#else - const uint4* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0].x = 0; - data[0].y = 0; - data[0].z = 0; - data[0].w = 0; - } -#endif -} - -template <> -__device__ __forceinline__ void load_global<8>(void* dst, const void* src) -{ - uint2* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile("ld.global.ca.v2.u32 {%0, %1}, [%2];\n" - : "=r"(data[0].x), "=r"(data[0].y) - : "l"(src)); -#else - const uint2* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_global<8>(void* dst, const void* src, bool do_access) -{ - uint2* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %3, 0;\n" - "\tmov.b32 %0, 0;\n" - "\tmov.b32 %1, 0;\n" - "\t@p ld.global.v2.u32 {%0, %1}, [%2];\n" - "}\n" - : "=r"(data[0].x), "=r"(data[0].y) - : "l"(src), "r"((int)do_access)); -#else - const uint2* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0].x = 0; - data[0].y = 0; - } -#endif -} - -template <> -__device__ __forceinline__ void load_global<8, LoadPolicy::CacheGlobal>(void* dst, const void* src) -{ - uint2* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile("ld.global.cg.v2.u32 {%0, %1}, [%2];\n" - : "=r"(data[0].x), "=r"(data[0].y) - : "l"(src)); -#else - const uint2* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_global<8, LoadPolicy::CacheGlobal>(void* dst, - const void* src, - bool do_access) -{ - uint2* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %3, 0;\n" - "\tmov.b32 %0, 0;\n" - "\tmov.b32 %1, 0;\n" - "\t@p ld.global.cg.v2.u32 {%0, %1}, [%2];\n" - "}\n" - : "=r"(data[0].x), "=r"(data[0].y) - : "l"(src), "r"((int)do_access)); -#else - const uint2* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0].x = 0; - data[0].y = 0; - } -#endif -} - -template <> -__device__ __forceinline__ void load_global<8, LoadPolicy::CacheStreaming>(void* dst, - const void* src) -{ - uint2* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile("ld.global.cs.v2.u32 {%0, %1}, [%2];\n" - : "=r"(data[0].x), "=r"(data[0].y) - : "l"(src)); -#else - const uint2* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_global<8, LoadPolicy::CacheStreaming>(void* dst, - const void* src, - bool do_access) -{ - uint2* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %3, 0;\n" - "\tmov.b32 %0, 0;\n" - "\tmov.b32 %1, 0;\n" - "\t@p ld.global.cs.v2.u32 {%0, %1}, [%2];\n" - "}\n" - : "=r"(data[0].x), "=r"(data[0].y) - : "l"(src), "r"((int)do_access)); -#else - const uint2* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0].x = 0; - data[0].y = 0; - } -#endif -} - -template <> -__device__ __forceinline__ void load_global<4>(void* dst, const void* src) -{ - int32_t* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile("ld.global.ca.u32 {%0}, [%1];\n" : "=r"(*data) : "l"(src)); -#else - const int32_t* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_global<4>(void* dst, const void* src, bool do_access) -{ - int32_t* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %2, 0;\n" - "\tmov.b32 %0, 0;\n" - "\t@p ld.global.u32 {%0}, [%1];\n" - "}\n" - : "=r"(data[0]) - : "l"(src), "r"((int)do_access)); -#else - const int32_t* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0] = 0; - } -#endif -} - -template <> -__device__ __forceinline__ void load_global<4, LoadPolicy::CacheGlobal>(void* dst, const void* src) -{ - int32_t* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile("ld.global.cg.u32 {%0}, [%1];\n" : "=r"(*data) : "l"(src)); -#else - const int32_t* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_global<4, LoadPolicy::CacheGlobal>(void* dst, - const void* src, - bool do_access) -{ - int32_t* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %2, 0;\n" - "\tmov.b32 %0, 0;\n" - "\t@p ld.global.cg.u32 {%0}, [%1];\n" - "}\n" - : "=r"(data[0]) - : "l"(src), "r"((int)do_access)); -#else - const int32_t* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0] = 0; - } -#endif -} - -template <> -__device__ __forceinline__ void load_global<4, LoadPolicy::CacheStreaming>(void* dst, - const void* src) -{ - int32_t* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile("ld.global.cs.u32 {%0}, [%1];\n" : "=r"(*data) : "l"(src)); -#else - const int32_t* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_global<4, LoadPolicy::CacheStreaming>(void* dst, - const void* src, - bool do_access) -{ - int32_t* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %2, 0;\n" - "\tmov.b32 %0, 0;\n" - "\t@p ld.global.cs.u32 {%0}, [%1];\n" - "}\n" - : "=r"(data[0]) - : "l"(src), "r"((int)do_access)); -#else - const int32_t* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0] = 0; - } -#endif -} - -template <> -__device__ __forceinline__ void load_global<2>(void* dst, const void* src) -{ - int16_t* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile("ld.global.ca.u16 {%0}, [%1];\n" : "=h"(*data) : "l"(src)); -#else - const int16_t* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_global<2>(void* dst, const void* src, bool do_access) -{ - int16_t* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %2, 0;\n" - "\tmov.u16 %0, 0;\n" - "\t@p ld.global.u16 {%0}, [%1];\n" - "}\n" - : "=h"(*data) - : "l"(src), "r"((int)do_access)); -#else - const int16_t* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0] = 0; - } -#endif -} - -template <> -__device__ __forceinline__ void load_global<2, LoadPolicy::CacheGlobal>(void* dst, const void* src) -{ - int16_t* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile("ld.global.cg.u16 {%0}, [%1];\n" : "=h"(*data) : "l"(src)); -#else - const int16_t* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_global<2, LoadPolicy::CacheGlobal>(void* dst, - const void* src, - bool do_access) -{ - int16_t* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %2, 0;\n" - "\tmov.u16 %0, 0;\n" - "\t@p ld.global.cg.u16 {%0}, [%1];\n" - "}\n" - : "=h"(*data) - : "l"(src), "r"((int)do_access)); -#else - const int16_t* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0] = 0; - } -#endif -} - -template <> -__device__ __forceinline__ void load_global<2, LoadPolicy::CacheStreaming>(void* dst, - const void* src) -{ - int16_t* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile("ld.global.cs.u16 {%0}, [%1];\n" : "=h"(*data) : "l"(src)); -#else - const int16_t* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_global<2, LoadPolicy::CacheStreaming>(void* dst, - const void* src, - bool do_access) -{ - int16_t* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %2, 0;\n" - "\tmov.u16 %0, 0;\n" - "\t@p ld.global.cs.u16 {%0}, [%1];\n" - "}\n" - : "=h"(*data) - : "l"(src), "r"((int)do_access)); -#else - const int16_t* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0] = 0; - } -#endif -} - -/////////// Load Shared /////////// -namespace internal { - -#ifdef PTX_AVAILABLE -__device__ __forceinline__ unsigned convert_to_shared(const void* ptr) -{ -#if __CUDACC_VER_MAJOR__ >= 11 - // In CUDA 11 we have a builtin intrinsic - return __cvta_generic_to_shared(ptr); -#else - unsigned ret_val; - asm volatile( - "{\n" - "\t.reg .u64 p1;\n" - "\tcvta.to.shared.u64 p1, %1\n" - "\tcvt.u32.u64 %0, p1;\n" - "}\n" - : "=r"(ret_val) - : "l"(ptr)); - return ret_val; -#endif -} -#endif - -} // namespace internal - -template <> -__device__ __forceinline__ void load_shared<16>(void* dst, const void* src) -{ - uint4* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - unsigned src_shr = internal::convert_to_shared(src); - - asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];\n" - : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) - : "r"(src_shr)); -#else - const uint4* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_shared<16>(void* dst, const void* src, bool do_access) -{ - uint4* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - unsigned src_shr = internal::convert_to_shared(src); - - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %5, 0;\n" - "\tmov.b32 %0, 0;\n" - "\tmov.b32 %1, 0;\n" - "\tmov.b32 %2, 0;\n" - "\tmov.b32 %3, 0;\n" - "\t@p ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];\n" - "}\n" - : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w) - : "r"(src_shr), "r"((int)do_access)); -#else - const uint4* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0].x = 0; - data[0].y = 0; - data[0].z = 0; - data[0].w = 0; - } -#endif -} - -template <> -__device__ __forceinline__ void load_shared<8>(void* dst, const void* src) -{ - uint2* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - unsigned src_shr = internal::convert_to_shared(src); - - asm volatile("ld.shared.v2.u32 {%0, %1}, [%2];\n" - : "=r"(data[0].x), "=r"(data[0].y) - : "r"(src_shr)); -#else - const uint2* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_shared<8>(void* dst, const void* src, bool do_access) -{ - uint2* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - unsigned src_shr = internal::convert_to_shared(src); - - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %3, 0;\n" - "\tmov.b32 %0, 0;\n" - "\tmov.b32 %1, 0;\n" - "\t@p ld.shared.v2.u32 {%0, %1}, [%2];\n" - "}\n" - : "=r"(data[0].x), "=r"(data[0].y) - : "r"(src_shr), "r"((int)do_access)); -#else - const uint2* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0].x = 0; - data[0].y = 0; - } -#endif -} - -template <> -__device__ __forceinline__ void load_shared<4>(void* dst, const void* src) -{ - int32_t* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - unsigned src_shr = internal::convert_to_shared(src); - - asm volatile("ld.shared.u32 {%0}, [%1];\n" : "=r"(*data) : "r"(src_shr)); -#else - const int32_t* src_cast = reinterpret_cast(src); - data[0] = src_cast[0]; -#endif -} - -template <> -__device__ __forceinline__ void load_shared<4>(void* dst, const void* src, bool do_access) -{ - int32_t* data = reinterpret_cast(dst); -#ifdef PTX_AVAILABLE - unsigned src_shr = internal::convert_to_shared(src); - - asm volatile( - "{\n" - "\t.reg .pred p;\n" - "\tsetp.ne.b32 p, %2, 0;\n" - "\tmov.b32 %0, 0;\n" - "\t@p ld.shared.u32 %0, [%1];\n" - "}\n" - : "=r"(data[0]) - : "r"(src_shr), "r"((int)do_access)); -#else - const int32_t* src_cast = reinterpret_cast(src); - if (do_access) { - data[0] = src_cast[0]; - } else { - data[0] = 0; - } -#endif -} - -/////////// Store Global /////////// - -template <> -__device__ __forceinline__ void store_global<16>(void* dst, const void* src) -{ - const uint4* data = reinterpret_cast(src); -#ifdef PTX_AVAILABLE - asm volatile("st.global.wb.v4.u32 [%0], {%1, %2, %3, %4};\n" - : - : "l"(dst), "r"(data[0].x), "r"(data[0].y), "r"(data[0].z), "r"(data[0].w) - : "memory"); -#else - uint4* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -#endif -} - -template <> -__device__ __forceinline__ void store_global<16, StorePolicy::CacheGlobal>(void* dst, - const void* src) -{ - const uint4* data = reinterpret_cast(src); -#ifdef PTX_AVAILABLE - asm volatile("st.global.cg.v4.u32 [%0], {%1, %2, %3, %4};\n" - : - : "l"(dst), "r"(data[0].x), "r"(data[0].y), "r"(data[0].z), "r"(data[0].w) - : "memory"); -#else - uint4* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -#endif -} - -template <> -__device__ __forceinline__ void store_global<16, StorePolicy::CacheStreaming>(void* dst, - const void* src) -{ - const uint4* data = reinterpret_cast(src); -#ifdef PTX_AVAILABLE - asm volatile("st.global.cs.v4.u32 [%0], {%1, %2, %3, %4};\n" - : - : "l"(dst), "r"(data[0].x), "r"(data[0].y), "r"(data[0].z), "r"(data[0].w) - : "memory"); -#else - uint4* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -#endif -} - -template <> -__device__ __forceinline__ void store_global<8>(void* dst, const void* src) -{ - const uint2* data = reinterpret_cast(src); -#ifdef PTX_AVAILABLE - asm volatile("st.global.wb.v2.u32 [%0], {%1, %2};\n" - : - : "l"(dst), "r"(data[0].x), "r"(data[0].y)); -#else - uint2* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -#endif -} - -template <> -__device__ __forceinline__ void store_global<8, StorePolicy::CacheGlobal>(void* dst, - const void* src) -{ - const uint2* data = reinterpret_cast(src); -#ifdef PTX_AVAILABLE - asm volatile("st.global.cg.v2.u32 [%0], {%1, %2};\n" - : - : "l"(dst), "r"(data[0].x), "r"(data[0].y)); -#else - uint2* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -#endif -} - -template <> -__device__ __forceinline__ void store_global<8, StorePolicy::CacheStreaming>(void* dst, - const void* src) -{ - const uint2* data = reinterpret_cast(src); -#ifdef PTX_AVAILABLE - asm volatile("st.global.cs.v2.u32 [%0], {%1, %2};\n" - : - : "l"(dst), "r"(data[0].x), "r"(data[0].y)); -#else - uint2* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -#endif -} - -template <> -__device__ __forceinline__ void store_global<4>(void* dst, const void* src) -{ - const int32_t* data = reinterpret_cast(src); -#ifdef PTX_AVAILABLE - asm volatile("st.global.wb.u32 [%0], %1;\n" : : "l"(dst), "r"(*data)); -#else - int32_t* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -#endif -} - -template <> -__device__ __forceinline__ void store_global<4, StorePolicy::CacheGlobal>(void* dst, - const void* src) -{ - const int32_t* data = reinterpret_cast(src); -#ifdef PTX_AVAILABLE - asm volatile("st.global.cg.u32 [%0], %1;\n" : : "l"(dst), "r"(*data)); -#else - int32_t* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -#endif -} - -template <> -__device__ __forceinline__ void store_global<4, StorePolicy::CacheStreaming>(void* dst, - const void* src) -{ - const int32_t* data = reinterpret_cast(src); -#ifdef PTX_AVAILABLE - asm volatile("st.global.cs.u32 [%0], %1;\n" : : "l"(dst), "r"(*data)); -#else - int32_t* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -#endif -} - -template <> -__device__ __forceinline__ void store_global<2>(void* dst, const void* src) -{ - const int16_t* data = reinterpret_cast(src); - - int16_t* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -} - -template <> -__device__ __forceinline__ void store_global<2, StorePolicy::CacheGlobal>(void* dst, - const void* src) -{ - const int16_t* data = reinterpret_cast(src); - - int16_t* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -} - -template <> -__device__ __forceinline__ void store_global<2, StorePolicy::CacheStreaming>(void* dst, - const void* src) -{ - const int16_t* data = reinterpret_cast(src); - - int16_t* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -} - -/////////// Store Shared /////////// - -template <> -__device__ __forceinline__ void store_shared<16>(void* dst, const void* src) -{ - const uint4* data = reinterpret_cast(src); -#ifdef PTX_AVAILABLE - unsigned dst_int = internal::convert_to_shared(dst); - - asm volatile("st.shared.v4.u32 [%0], {%1, %2, %3, %4};\n" - : - : "r"(dst_int), "r"(data[0].x), "r"(data[0].y), "r"(data[0].z), "r"(data[0].w)); -#else - uint4* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -#endif -} - -template <> -__device__ __forceinline__ void store_shared<8>(void* dst, const void* src) -{ - const uint2* data = reinterpret_cast(src); -#ifdef PTX_AVAILABLE - unsigned dst_int = internal::convert_to_shared(dst); - - asm volatile("st.shared.v2.u32 [%0], {%1, %2};\n" - : - : "r"(dst_int), "r"(data[0].x), "r"(data[0].y)); -#else - uint2* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -#endif -} - -template <> -__device__ __forceinline__ void store_shared<4>(void* dst, const void* src) -{ - const int32_t* data = reinterpret_cast(src); -#ifdef PTX_AVAILABLE - unsigned dst_int = internal::convert_to_shared(dst); - - asm volatile("st.shared.u32 [%0], %1;\n" : : "r"(dst_int), "r"(*data)); -#else - int32_t* dst_cast = reinterpret_cast(dst); - dst_cast[0] = data[0]; -#endif -} - -/////////// Asynchronous Memory Copy /////////// - -#ifdef ASYNC_COPY_AVAILABLE -template -__device__ __forceinline__ void memcpy_async(void* shr, const void* gbl) -{ - static_assert((AccessSize == 4 || AccessSize == 8 || AccessSize == 16)); - unsigned shr_int = internal::convert_to_shared(shr); - - asm volatile("cp.async.ca.shared.global [%0], [%1], %2;\n" - : - : "r"(shr_int), "l"(gbl), "n"(AccessSize)); -} - -template -__device__ __forceinline__ void memcpy_async_nop(void* shr, const void* gbl, bool predicate) -{ - static_assert((AccessSize == 4 || AccessSize == 8 || AccessSize == 16)); - unsigned shr_int = internal::convert_to_shared(shr); - - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" - : - : "r"((int)predicate), "r"(shr_int), "l"(gbl), "n"(AccessSize)); -} - -template -__device__ __forceinline__ void memcpy_async_zero(void* shr, const void* gbl, bool predicate) -{ - static_assert((AccessSize == 4 || AccessSize == 8 || AccessSize == 16)); - unsigned shr_int = internal::convert_to_shared(shr); - int bytes_to_copy = (predicate ? AccessSize : 0); - - asm volatile("cp.async.ca.shared.global [%0], [%1], %2, %3;\n" - : - : "r"(shr_int), "l"(gbl), "n"(AccessSize), "r"(bytes_to_copy)); -} - -template -__device__ __forceinline__ void memcpy_async_zero_nop(void* shr, - const void* gbl, - bool zero_predicate, - bool nop_predicate) -{ - static_assert((AccessSize == 4 || AccessSize == 8 || AccessSize == 16)); - unsigned shr_int = internal::convert_to_shared(shr); - int bytes_to_copy = (zero_predicate ? AccessSize : 0); - - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3, %4;\n" - "}\n" - : - : "r"((int)nop_predicate), "r"(shr_int), "l"(gbl), "n"(AccessSize), "r"(bytes_to_copy)); -} - -// Cache global variants. Separate interface to require deliberate use of them. -__device__ __forceinline__ void memcpy_async_cg(void* shr, const void* gbl) -{ - unsigned shr_int = internal::convert_to_shared(shr); - - asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" : : "r"(shr_int), "l"(gbl)); -} - -__device__ __forceinline__ void memcpy_async_nop_cg(void* shr, const void* gbl, bool predicate) -{ - unsigned shr_int = internal::convert_to_shared(shr); - - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.cg.shared.global [%1], [%2], 16;\n" - "}\n" - : - : "r"((int)predicate), "r"(shr_int), "l"(gbl)); -} - -__device__ __forceinline__ void memcpy_async_zero_cg(void* shr, const void* gbl, bool predicate) -{ - unsigned shr_int = internal::convert_to_shared(shr); - int bytes_to_copy = (predicate ? 16 : 0); - - asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" - : - : "r"(shr_int), "l"(gbl), "r"(bytes_to_copy)); -} - -__device__ __forceinline__ void memcpy_async_zero_nop_cg(void* shr, - const void* gbl, - bool zero_predicate, - bool nop_predicate) -{ - unsigned shr_int = internal::convert_to_shared(shr); - int bytes_to_copy = (zero_predicate ? 16 : 0); - - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.cg.shared.global [%1], [%2], 16, %3;\n" - "}\n" - : - : "r"((int)nop_predicate), "r"(shr_int), "l"(gbl), "r"(bytes_to_copy)); -} - -__device__ __forceinline__ void memcpy_async_fence() { asm volatile("cp.async.commit_group;\n"); } - -template -__device__ __forceinline__ void memcpy_async_wait() -{ - static_assert(stages <= 8); - - asm volatile("cp.async.wait_group %0;\n" : : "n"(stages)); -} - -// TODO: The tail complete should be a known compile time artifact, should try and induce this -// without all of the branches from the call-site. This is a hacky solution. -template <> -__device__ __forceinline__ void tail_complete_wait<1>(int remaining_stages) -{ - if (remaining_stages == 0) memcpy_async_wait<0>(); -} - -template <> -__device__ __forceinline__ void tail_complete_wait<2>(int remaining_stages) -{ - if (remaining_stages == 1) - memcpy_async_wait<1>(); - else if (remaining_stages == 0) - memcpy_async_wait<0>(); -} - -template <> -__device__ __forceinline__ void tail_complete_wait<3>(int remaining_stages) -{ - if (remaining_stages == 2) - memcpy_async_wait<2>(); - else if (remaining_stages == 1) - memcpy_async_wait<1>(); - else if (remaining_stages == 0) - memcpy_async_wait<0>(); -} - -template <> -__device__ __forceinline__ void tail_complete_wait<4>(int remaining_stages) -{ - if (remaining_stages == 3) - memcpy_async_wait<3>(); - else if (remaining_stages == 2) - memcpy_async_wait<2>(); - else if (remaining_stages == 1) - memcpy_async_wait<1>(); - else if (remaining_stages == 0) - memcpy_async_wait<0>(); -} - -template <> -__device__ __forceinline__ void tail_complete_wait<5>(int remaining_stages) -{ - if (remaining_stages == 4) - memcpy_async_wait<4>(); - else if (remaining_stages == 3) - memcpy_async_wait<3>(); - else if (remaining_stages == 2) - memcpy_async_wait<2>(); - else if (remaining_stages == 1) - memcpy_async_wait<1>(); - else if (remaining_stages == 0) - memcpy_async_wait<0>(); -} - -template <> -__device__ __forceinline__ void tail_complete_wait<6>(int remaining_stages) -{ - if (remaining_stages == 5) - memcpy_async_wait<5>(); - else if (remaining_stages == 4) - memcpy_async_wait<4>(); - else if (remaining_stages == 3) - memcpy_async_wait<3>(); - else if (remaining_stages == 2) - memcpy_async_wait<2>(); - else if (remaining_stages == 1) - memcpy_async_wait<1>(); - else if (remaining_stages == 0) - memcpy_async_wait<0>(); -} -#endif - -} // namespace mem_access diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/normalize_layer.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/normalize_layer.h deleted file mode 100644 index b9c719087a6ba0af693589f3c36da452b82e21d7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/normalize_layer.h +++ /dev/null @@ -1,207 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include -#include -#include -#include "custom_cuda_layers.h" - -using namespace std; - -template -class Normalize_Layer { -public: - struct Config { - uint32_t batchSize; - uint32_t seqLength; - uint32_t hiddenDim; - float epsilon; - bool training; - bool useMean; - Config(uint32_t batch, - uint32_t seq, - uint32_t h, - float epsilon = 1e-12, - bool training = true, - bool useMean = true) - : batchSize(batch), - seqLength(seq), - hiddenDim(h), - epsilon(epsilon), - training(training), - useMean(useMean) - { - } - }; - - Normalize_Layer(Config config) - : config_(config), vars(nullptr), means(nullptr), vals_hat(nullptr) - { - } - - ~Normalize_Layer() {} - - void ForwardCheckpoint(int bsz, // batch * seq - T* vals, - const T* residual, - const T* gamma, - const T* betta, - cudaStream_t& stream, - bool preLayerNorm = false) - { - launch_bias_residual_layer_norm(vals, - residual, - gamma, - betta, - config_.epsilon, - bsz, - config_.hiddenDim, - stream, - preLayerNorm, - config_.training, - vars, - means); - } - - void Forward(int bsz, - T* vals, - const T* residual, - const T* gamma, - const T* betta, - cudaStream_t& stream, - bool preLayerNorm = false) - { - launch_bias_residual_layer_norm(vals, - residual, - gamma, - betta, - config_.epsilon, - bsz, - config_.hiddenDim, - stream, - preLayerNorm, - config_.training, - vars); - } - - void Backward(int bsz, - const T* out_grad, - const T* gamma, - T* gamma_grad, - T* betta_grad, - cudaStream_t stream[2], - T* inp_grad_out, - const T* norm_in = nullptr) - { - launch_layerNorm_backward(out_grad, - norm_in, - vars, - means, - gamma, - gamma_grad, - betta_grad, - inp_grad_out, - bsz, - config_.hiddenDim, - stream); - } - - void Backward(int bsz, - const T* out_grad, - const T* gamma, - const T* betta, - T* gamma_grad, - T* betta_grad, - cudaStream_t stream[2], - T* inp_grad_out, - const T* norm_out) - { - launch_layerNorm_backward(out_grad, - norm_out, - vars, - gamma, - gamma_grad, - betta_grad, - inp_grad_out, - bsz, - config_.hiddenDim, - stream, - !config_.useMean, - betta); - } - - void BackwardFusedAdd(int bsz, - const T* out_grad1, - const T* out_grad2, - const T* gamma, - T* gamma_grad, - T* betta_grad, - cudaStream_t stream[2], - T* inp_grad_out, - const T* norm_in = nullptr) - { - launch_layerNorm_backward_fused_add(out_grad1, - out_grad2, - norm_in, - vars, - means, - gamma, - gamma_grad, - betta_grad, - inp_grad_out, - bsz, - config_.hiddenDim, - stream); - } - - void BackwardFusedAdd(int bsz, - const T* out_grad1, - const T* out_grad2, - const T* gamma, - const T* betta, - T* gamma_grad, - T* betta_grad, - cudaStream_t stream[2], - T* inp_grad_out, - const T* norm_out) - { - launch_layerNorm_backward_fused_add(out_grad1, - out_grad2, - norm_out, - vars, - gamma, - gamma_grad, - betta_grad, - inp_grad_out, - bsz, - config_.hiddenDim, - stream, - !config_.useMean, - betta); - } - - inline bool UseMean() const { return config_.useMean; } - - inline void SetVar(T* variance) - { - if (!variance) { throw std::runtime_error("Normalize variance is null."); } - vars = variance; - } - - inline void SetMean(T* mean) - { - if (!mean) { throw std::runtime_error("Normalize mean is null."); } - means = mean; - } - -private: - Config config_; - T* vars; - T* means; - T* vals_hat; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantization.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantization.h deleted file mode 100644 index 5bdc96061a311d78b94ec9051e7f13919454934b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantization.h +++ /dev/null @@ -1,138 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include "ds_kernel_utils.h" - -namespace quantize { - -enum class Type { Symmetric, Asymmetric }; - -struct PackedInt4 { - int8_t high : 4; - int8_t low : 4; -}; - -DS_HD_INLINE bool requires_offset(Type qType) { return qType == Type::Asymmetric; } - -} // namespace quantize - -void launch_quant(int8_t* output_data, - float* params, - const __half* input_data, - const int groups, - const int elems_per_group, - const int num_bits, - const quantize::Type quant_type, - cudaStream_t stream); - -template -void launch_dequantize_kernel(T* dequant_data, - const int8_t* q_data, - const float* q_params, - quantize::Type q_type, - int num_bits, - int elems_per_group, - int total_elems, - cudaStream_t stream); - -void launch_swizzled_quant(int8_t* q_data, - float* q_scales, - const __half* input_data, - int num_bits, - quantize::Type q_type, - int groups, - int elems_per_group, - int pipelining, - int nodes, - int devices_per_node, - cudaStream_t stream); - -void launch_loco_swizzled_quant(int8_t* quantized_data, - float* quantized_scales, - const __half* uncompressed_data, - __half* error_feedback, - const float err_beta, - int num_bits, - quantize::Type quant_type, - int groups, - int elems_per_group, - int pipelining, - int nodes, - int devices_per_node, - cudaStream_t stream); - -void launch_loco_dequant_reduce(int8_t* reduced_data, - float* reduced_scales, - const int8_t* input_data, - const float* input_scales, - int num_gpus, - int num_bits, - quantize::Type quant_type, - int out_groups, - int elems_per_out_group, - int elems_per_in_tensor, - int groups_per_in_tensor, - int elems_per_in_group, - __half2* error_feedback, - const float err_beta, - cudaStream_t stream); - -void launch_dequant_reduce(int8_t* reduced_data, - float* reduced_scales, - const int8_t* input_data, - const float* input_scales, - int num_gpus, - int num_bits, - quantize::Type quant_type, - int out_groups, - int elems_per_out_group, - int elems_per_in_tensor, - int groups_per_in_tensor, - int elems_per_in_group, - cudaStream_t stream); - -template -void launch_fake_quantize_kernel(T* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream); -template -void launch_sr_fake_quantize_kernel(T* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream); -template -void launch_fake_quantize_kernel_asym(T* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream); -template -void launch_sr_fake_quantize_kernel_asym(T* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream); - -void launch_dequantize_int4_to_half_experimental(uint8_t* data_in, - half* data_out, - half* scale_buffer, - half* min_val_buffer, - int num_group, - int group_size, - cudaStream_t stream); - -void launch_dequantize_int8_to_half_experimental(uint8_t* data_in, - half* data_out, - half* scale_buffer, - half* min_val_buffer, - int num_group, - int group_size, - cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantization_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantization_utils.h deleted file mode 100644 index 94958fb455c6f632713aea4f27b7bc1c901bfa45..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantization_utils.h +++ /dev/null @@ -1,469 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include "conversion_utils.h" -#include "ds_kernel_utils.h" -#include "memory_access_utils.h" -#include "quantization.h" -#include "reduction_utils.h" - -#pragma once - -using rop = reduce::ROpType; - -namespace quantize { -constexpr int granularity = 16; -constexpr int h_per_load = granularity / sizeof(__half); -constexpr int h2_per_load = granularity / sizeof(__half2); -constexpr int max_threads = 1024; - -/* -Class to hold the quantization parameters for a given tensor. -Holds the implementation of the quantization operation. -*/ - -template -class Params { -public: - /* - Quantization implementation, supports - 1) 4 Bit - 2) 8 Bit - 3) Symmetric - 4) Asymmetric - Function Arguments : - val : The __half value to quantize. - */ - DS_D_INLINE int8_t quantize(__half val); - - template - DS_D_INLINE T dequantize(int8_t val); - - DS_D_INLINE void store(float* params, int group_index); - - // Initialize from memory - DS_D_INLINE Params(const float* params, int group_index); -}; - -template -class Params { -public: - float scale; - - DS_D_INLINE Params(float max) - { - if (max == 0) { - scale = 1.0; - } else { - scale = (1 << numBits) / (2 * max); - } - } - - DS_D_INLINE int8_t quantize(__half val) - { - constexpr int32_t q_min = -(1 << (numBits - 1)); - constexpr int32_t q_max = (1 << (numBits - 1)) - 1; - - float val_f = conversion::to(val) * scale; - int32_t data_i32 = conversion::to(val_f); - data_i32 = min(max(data_i32, q_min), q_max); - return (int8_t)data_i32; - } - - template - DS_D_INLINE T dequantize(int8_t val) - { - const float val_deq_f = conversion::to(val) * scale; - return conversion::to(val_deq_f); - } - - DS_D_INLINE void store(float* params, int group_index) - { - const float store_scale = 1 / scale; - mem_access::store_global(params + group_index, &store_scale); - } - - DS_D_INLINE Params(const float* params, int group_index) - { - mem_access::load_global(&scale, params + group_index); - } -}; - -template -class Params { -public: - float scale; - float offset; - - DS_D_INLINE Params(float max, float min) - { - if (max == min) { - scale = 1.0; - } else { - scale = ((1 << numBits)) / (max - min); - } - offset = (max + min) / 2; - } - - DS_D_INLINE int8_t quantize(__half val) - { - constexpr int32_t q_min = -(1 << (numBits - 1)); - constexpr int32_t q_max = (1 << (numBits - 1)) - 1; - - float val_f = (conversion::to(val) - offset) * scale; - int32_t data_i32 = conversion::to(val_f); - data_i32 = min(max(data_i32, q_min), q_max); - return (int8_t)data_i32; - } - - template - DS_D_INLINE T dequantize(int8_t val) - { - const float val_deq_f = ((conversion::to(val)) * scale) + offset; - return conversion::to(val_deq_f); - } - - DS_D_INLINE void store(float* params, int group_index) - { - // Codegen should turn this into stg.64 - const float store_scale = 1 / scale; - mem_access::store_global(params + 2 * group_index, &store_scale); - mem_access::store_global(params + 2 * group_index + 1, &offset); - } - - DS_D_INLINE Params(const float* params, int group_index) - { - // Codegen should turn this into ldg.64 - mem_access::load_global(&scale, params + 2 * group_index); - mem_access::load_global(&offset, params + 2 * group_index + 1); - } -}; - -/* -Group stats tracks the necessary statistics about the quantized group -to abstract the particulars for the main loop. -*/ -template -class GroupStats { -public: - DS_D_INLINE void update(__half2 val); - - DS_D_INLINE void reduce(cg::thread_block& tb, cg::thread_block_tile& warp); -}; - -template <> -class GroupStats { -public: - // Symmetric quantization only tracks the maximum absolute value - __half2 cur_max; - float max; - - /* - Technically, this would give bad results if there - are 0 values to process since the reduction would - give -inf instead of 0. We do not consider this - to be a reasonable edge case. - */ - DS_D_INLINE GroupStats() { cur_max = reduce::init(); } - - /* - Updated the running absmax used to calculate params. - Function Arguments : - val : The __half2 value to update the running min and max with. - */ - DS_D_INLINE void update(__half2 val) - { - cur_max = reduce::element(cur_max, __habs2(val)); - } - - /* - Function to return calculated quantization params. - Template Arguments : - numBits - Number of bits in quantized element. int : 8 or 4 - Function Arguments : - tb - Threadblock object. cg::thread_block - warp - Warp object. cg::thread_block_tile - */ - template - DS_D_INLINE Params get_params( - cg::thread_block& tb, - cg::thread_block_tile& warp) - { - const float2 partial_max = conversion::to(cur_max); - float max = reduce::element(partial_max.x, partial_max.y); - - reduce::partitioned_block(tb, warp, max); - Params params(max); - - return params; - } -}; - -template <> -class GroupStats { -public: - __half2 cur_max; - __half2 cur_min; - - /* - Initialize cur_max to -inf, cur_min to inf since - we are doing a true range analysis. - */ - DS_D_INLINE GroupStats() - { - cur_max = reduce::init(); - cur_min = reduce::init(); - } - - /* - Updated the running min and max used to calculate params. - Function Arguments : - val : The __half2 value to update the running min and max with. - */ - DS_D_INLINE void update(__half2 val) - { - cur_max = reduce::element(cur_max, val); - cur_min = reduce::element(cur_min, val); - } - - /* - Function to return calculated quantization params. - Template Arguments : - numBits - Number of bits in quantized element. int : 8 or 4 - Function Arguments : - tb - Threadblock object. cg::thread_block - warp - Warp object. cg::thread_block_tile - */ - template - DS_D_INLINE Params get_params( - cg::thread_block& tb, - cg::thread_block_tile& warp) - { - const float2 partial_max = conversion::to(cur_max); - float max = reduce::element(partial_max.x, partial_max.y); - - const float2 partial_min = conversion::to(cur_min); - float min = reduce::element(partial_min.x, partial_min.y); - - reduce::partitioned_block(tb, warp, max, min); - - Params params(max, min); - - return params; - } -}; - -/* -Device function that quantizes 16 bytes of __half type input data. -Template Arguments : - numBits - Number of bits in quantized element. int : 8 or 4 - qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric -Function Arguments : - local_output - Pointer to local memory to store quantized data. int8_t* - data - Pointer to input data. __half* - Params - Parameters for quantization. Params -*/ -template -DS_D_INLINE void _chunk(int8_t* local_output, const __half* data, Params q_params); - -/* -Device function that quantizes 16 bytes of __half2 type input data. -Template Arguments : - numBits - Number of bits in quantized element. int : 8 or 4 - qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric -Function Arguments : - local_output - Pointer to local memory to store quantized data. int8_t* - data - Pointer to input data. __half2* - Params - Parameters for quantization. Params -*/ -template -DS_D_INLINE void _chunk(int8_t* local_output, const __half2* data, Params q_params); - -/* -Helper function to do serial reduction on register-file arrays. -Template Arguments : - qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric - numChunks - Number of bits in quantized element. int : 8 or 4 -Function Arguments : - local_buffer - Pointer memory with input half2 data to be quantized. -*/ -template -DS_D_INLINE GroupStats _local_serial_reduce(__half2* local_buffer); - -/* -The main loop of the kernel that quantizes array in local memory of __half2 type input data, when -Quantization parameters are pre-computed. -Template Arguments : - qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric - numBits - Number of bits in quantized element. int : 8 or 4 - numChunks - Number of chunks(16 bytes of Input data). int : 8 or 4 -Function Arguments : - local_buffer - Pointer memory with input half2 data to be quantized. - scales - Pointer to output scales. - offsets - Pointer to output offsets. - output_data - Pointer to output data. - elems_per_group - Number of elements to quantize in a group. - q_params - Quantization parameters. -*/ -template -DS_D_INLINE void local_array(cg::thread_block& tb, - cg::thread_block_tile& warp, - __half2* local_buffer, - float* __restrict__ scales, - float* __restrict__ offsets, - int8_t* __restrict__ output_data, - const int& elems_per_group, - const int& groups, - Params q_params); - -/* -The main loop of the kernel that quantizes array in local memory of __half2 type input data. -This function computes quantization parameters for each group. -Template Arguments : - qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric - numBits - Number of bits in quantized element. int : 8 or 4 - numChunks - Number of chunks(16 bytes of Input data). int : 8 or 4 -Function Arguments : - local_buffer - Pointer memory with input half2 data to be quantized. - scales - Pointer to output scales. - offsets - Pointer to output offsets. - output_data - Pointer to output data. - elems_per_group - Number of elements to quantize in a group. -*/ -template -__device__ void local_array(__half2* local_buffer, - float* __restrict__ scales, - float* __restrict__ offsets, - int8_t* __restrict__ output_data, - const int& elems_per_group, - const int& groups); - -template -DS_D_INLINE void _chunk(int8_t* local_output, const __half* data, Params q_params) -{ - constexpr int32_t elems = 16 / sizeof(__half); - constexpr int32_t num_elems_packed = 8 / numBits; - -#pragma unroll - for (int i = 0, oi = 0; i < elems; i += num_elems_packed, oi++) { - if (num_elems_packed == 1) { - // TODO(cmikeh2): refactor to use conversion utils - local_output[i] = q_params.quantize(data[i]); - } else if (num_elems_packed == 2) { - int8_t data_i8_1 = q_params.quantize(data[i]); - int8_t data_i8_2 = q_params.quantize(data[i + 1]); - auto data_i8 = PackedInt4{data_i8_2, data_i8_1}; - local_output[oi] = *((int8_t*)(&data_i8)); - } - } -} - -template -DS_D_INLINE void _chunk(int8_t* local_output, const __half2* data, Params q_params) -{ - const __half* data_cast = reinterpret_cast(data); - _chunk(local_output, data_cast, q_params); -} - -template -DS_D_INLINE GroupStats _local_serial_reduce(__half2* local_buffer) -{ - GroupStats stats; -#pragma unroll - for (int i = 0; i < numChunks * h2_per_load; i++) { stats.update(local_buffer[i]); } - - return stats; -} - -template -DS_D_INLINE void local_array(cg::thread_block& tb, - cg::thread_block_tile& warp, - __half2* local_buffer, - float* __restrict__ global_params, - int8_t* __restrict__ output_data, - const int& elems_per_group, - const int& groups, - Params q_params) -{ - constexpr int num_ele_int8 = 8 / numBits; - constexpr int num_int8_out = quantize::h_per_load / num_ele_int8; - - // Indexing offsets - const int block_num = - (tb.group_index().x * max_threads / threads_per_group) + tb.thread_index().y; - const int block_offset = block_num * elems_per_group; - const int elem_offset = tb.thread_index().x * quantize::h_per_load; - const int base_offset = (block_offset + elem_offset) / num_ele_int8; - const int stride = tb.size() * quantize::h_per_load / num_ele_int8; - - int8_t local_output[num_int8_out]; - - if (tb.thread_index().x == 0 && block_num < groups) { - q_params.store( - global_params, - (tb.group_index().x * max_threads / threads_per_group) + tb.thread_index().y); - } -#pragma unroll - for (int i = 0; i < numChunks; i++) { - if (elem_offset + i * stride * num_ele_int8 < elems_per_group && block_num < groups) { - quantize::_chunk( - local_output, local_buffer + i * quantize::h2_per_load, q_params); - mem_access::store_global(output_data + (base_offset + i * stride), - local_output); - } - } -} - -template -DS_D_INLINE void local_array(cg::thread_block& tb, - cg::thread_block_tile& warp, - __half* local_buffer, - float* __restrict__ global_params, - int8_t* __restrict__ output_data, - const int& elems_per_group, - const int& groups, - Params q_params) -{ - __half2* local_buffer_h2 = reinterpret_cast<__half2*>(local_buffer); - - quantize::local_array( - tb, warp, local_buffer, global_params, output_data, elems_per_group, groups, q_params); -} - -template -__device__ void local_array(__half2* local_buffer, - float* __restrict__ global_params, - int8_t* __restrict__ output_data, - const int& elems_per_group, - const int& groups) -{ - cg::thread_block tb = cg::this_thread_block(); - cg::thread_block_tile warp = cg::tiled_partition(tb); - - auto group_stats = _local_serial_reduce(local_buffer); - auto params = group_stats.template get_params(tb, warp); - - quantize::local_array( - tb, warp, local_buffer, global_params, output_data, elems_per_group, groups, params); -} - -template -__device__ void local_array(__half* local_buffer, - float* __restrict__ global_params, - int8_t* __restrict__ output_data, - const int& elems_per_group, - const int& groups) -{ - __half2* local_buffer_h2 = reinterpret_cast<__half2*>(local_buffer); - quantize::local_array( - local_buffer_h2, global_params, output_data, elems_per_group, groups); -} - -} // namespace quantize diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantizer.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantizer.h deleted file mode 100644 index f4f63160d79b23aaa94ae183a176d6c05024496a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/quantizer.h +++ /dev/null @@ -1,19 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#ifdef __HIP_PLATFORM_AMD__ -#include -#else -#include -#endif - -#include -#include -#include -#include -#include -#include diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/reduction_utils.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/reduction_utils.h deleted file mode 100644 index eb9afb66a894a63889c3b07c33a634ff60f54f36..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/reduction_utils.h +++ /dev/null @@ -1,826 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include "conversion_utils.h" -#include "ds_kernel_utils.h" -#include "memory_access_utils.h" - -namespace cg = cooperative_groups; - -namespace reduce { - -enum class ROpType { - // Addition - Add, - - // Maximum reduction - Max, - - // Minimum reduction - Min, -}; - -constexpr int max_threads = 1024; -constexpr int max_warps = max_threads / hw_warp_size; - -/* -High level API. The API takes in a set of operations and variables -and performs that reduction operation on that variable. The reductions -of each of the arguments are completely independent of each other ( -i.e., the val1-op1 combination has no impact on val2-op2). - -Example usage: -``` cpp -float max_val; -float min_val; -reduce::block(tb, warp, max_val, min_val); -``` - -TODO(cmikeh2): In theory, we might be able to do this sequentially with -device functions and rely on the assembler correctly behaving. My initial -instinct is this won't work, but if it does it would reduce implementation -cost significantly. - -TODO(cmikeh2): We need to support sub-block reductions. The warp intrinsic -currently supports this (more incidentally than anything else). It is not -uncommon in something like softmax or a fused attention kernel to map multiple -reductions to a thread block, but each reduction itself is only scoped -to part of the threads (i.e block size = 512, 128 threads per reduction). -*/ -template -DS_D_INLINE void block(cg::thread_block& tb, cg::thread_block_tile& warp, float& val); - -template -DS_D_INLINE void block(cg::thread_block& tb, - cg::thread_block_tile& warp, - float& val1, - float& val2); - -template -DS_D_INLINE void block(cg::thread_block& tb, - cg::thread_block_tile& warp, - float& val1, - float& val2, - float& val3); - -template -DS_D_INLINE void block(cg::thread_block& tb, - cg::thread_block_tile& warp, - float& val1, - float& val2, - float& val3, - float& val4); - -/* -The partitioned block is a special case of the above where in the warps of a threadblock are -partitioned into separate independent reductions. For example, I might have an 8 warp thread block -in which each pair of warps is processing an independent piece of data. I would then reduce that -data with the something like the following: -``` cpp -float max_val; -reduce::partitioned_block(tb, warp, max_val); -``` -After which, each pair of warps would have coherent data with each other. Note, this API will not -provide correct results if the number of warps per partition is not a power of 2. -*/ -template -DS_D_INLINE void partitioned_block(cg::thread_block& tb, - cg::thread_block_tile& warp, - float& val); - -template -DS_D_INLINE void partitioned_block(cg::thread_block& tb, - cg::thread_block_tile& warp, - float& val1, - float& val2); - -template -DS_D_INLINE void partitioned_block(cg::thread_block& tb, - cg::thread_block_tile& warp, - float& val1, - float& val2, - float& val3); - -template -DS_D_INLINE void partitioned_block(cg::thread_block& tb, - cg::thread_block_tile& warp, - float& val1, - float& val2, - float& val3, - float& val4); - -/* -Single element reduction primitives. Used inside serial collection -loops. - -Example usage: -using rop = reduce::OpType; -float min = init(); -for (int i = 0; i < 4; i++) { - min = reduce::element(min, data[i]); -} -*/ - -template -DS_D_INLINE T element(const T lhs, const T rhs); - -template -DS_D_INLINE T init(); - -/********************** Internal reduction APIs **********************/ - -/* -Single element "reductions". TODO(cmikeh2): this sort of "op" concept -should be refactored into its own implementation at some point. This interface -may be easily expanded for new types/operations, but the typical reductions -we need are covered with min/max/add on float. - -NOTE: there is no mean reduction because that relies on knowledge of how -many values were already reduced into each scalar. Implementing this on top -of reduce should be straightforward (can just wrap the sum reduction) and -would be a good extension of the header. -*/ - -DS_D_INLINE int _warp_rank() -{ - const int thread_rank = - threadIdx.x + threadIdx.y * blockDim.x + threadIdx.z * blockDim.x * blockDim.y; - return thread_rank / hw_warp_size; -} - -/* Float element reduce implementations */ -template <> -DS_D_INLINE float element(const float lhs, const float rhs) -{ - return lhs + rhs; -} - -template <> -DS_D_INLINE double element(const double lhs, const double rhs) -{ - return lhs + rhs; -} - -template <> -DS_D_INLINE float element(const float lhs, const float rhs) -{ - return fmaxf(lhs, rhs); -} - -template <> -DS_D_INLINE float element(const float lhs, const float rhs) -{ - return fminf(lhs, rhs); -} - -/* __half element reduce implementation */ -template <> -DS_D_INLINE __half element(const __half lhs, const __half rhs) -{ - return lhs + rhs; -} - -template <> -DS_D_INLINE __half element(const __half lhs, const __half rhs) -{ -#if __CUDA_ARCH__ >= 800 - // Intrinsic limited to Ampere + newer - return __hmax(lhs, rhs); -#else - return (lhs > rhs) ? lhs : rhs; -#endif -} - -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE __nv_bfloat16 element(const __nv_bfloat16 lhs, const __nv_bfloat16 rhs) -{ -#if __CUDA_ARCH__ >= 800 - // Intrinsic limited to Ampere + newer - return __hmax(lhs, rhs); -#else - return (lhs > rhs) ? lhs : rhs; -#endif -} -#endif - -template <> -DS_D_INLINE __half element(const __half lhs, const __half rhs) -{ -#if __CUDA_ARCH__ >= 800 - // Intrinsic limited to Ampere + newer - return __hmin(lhs, rhs); -#else - return (lhs < rhs) ? lhs : rhs; -#endif -} - -/* __half2 element reduce implementation */ -template <> -DS_D_INLINE __half2 element(const __half2 lhs, const __half2 rhs) -{ - return lhs + rhs; -} - -template <> -DS_D_INLINE __half2 element(const __half2 lhs, const __half2 rhs) -{ -#if __CUDA_ARCH__ >= 800 - return __hmax2(lhs, rhs); -#else - __half2 ret_val; - ret_val.x = (lhs.x > rhs.x) ? lhs.x : rhs.x; - ret_val.y = (lhs.y > rhs.y) ? lhs.y : rhs.y; - return ret_val; -#endif -} - -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE __nv_bfloat162 element(const __nv_bfloat162 lhs, const __nv_bfloat162 rhs) -{ -#if __CUDA_ARCH__ >= 800 - return __hmax2(lhs, rhs); -#else - __nv_bfloat162 ret_val; - ret_val.x = (lhs.x > rhs.x) ? lhs.x : rhs.x; - ret_val.y = (lhs.y > rhs.y) ? lhs.y : rhs.y; - return ret_val; -#endif -} -#endif - -template <> -DS_D_INLINE __half2 element(const __half2 lhs, const __half2 rhs) -{ -#if __CUDA_ARCH__ >= 800 - return __hmin2(lhs, rhs); -#else - __half2 ret_val; - ret_val.x = (lhs.x < rhs.x) ? lhs.x : rhs.x; - ret_val.y = (lhs.y < rhs.y) ? lhs.y : rhs.y; - return ret_val; -#endif -} - -template <> -DS_D_INLINE int32_t element(const int32_t lhs, const int32_t rhs) -{ - return lhs + rhs; -} - -template <> -DS_D_INLINE int32_t element(const int32_t lhs, const int32_t rhs) -{ - return (lhs > rhs) ? lhs : rhs; -} - -template <> -DS_D_INLINE int32_t element(const int32_t lhs, const int32_t rhs) -{ - return (lhs < rhs) ? lhs : rhs; -} - -template <> -DS_D_INLINE uint32_t element(const uint32_t lhs, const uint32_t rhs) -{ - return lhs + rhs; -} - -template <> -DS_D_INLINE uint32_t element(const uint32_t lhs, const uint32_t rhs) -{ - return (lhs > rhs) ? lhs : rhs; -} - -template <> -DS_D_INLINE uint32_t element(const uint32_t lhs, const uint32_t rhs) -{ - return (lhs < rhs) ? lhs : rhs; -} - -template <> -DS_D_INLINE int64_t element(const int64_t lhs, const int64_t rhs) -{ - return lhs + rhs; -} - -template <> -DS_D_INLINE int64_t element(const int64_t lhs, const int64_t rhs) -{ - return (lhs > rhs) ? lhs : rhs; -} - -template <> -DS_D_INLINE int64_t element(const int64_t lhs, const int64_t rhs) -{ - return (lhs < rhs) ? lhs : rhs; -} - -/* -Reduction initialization primitives -*/ -template <> -DS_D_INLINE float init() -{ - return 0.0f; -} -template <> -DS_D_INLINE double init() -{ - return (double)0.0f; -} - -template <> -DS_D_INLINE float init() -{ - // Positive infinity - return INFINITY; -} - -template <> -DS_D_INLINE float init() -{ - // Negative infinity - return -INFINITY; -} - -template <> -DS_D_INLINE __half init() -{ - constexpr __half_raw zero = {0x0000}; - return __half(zero); -} - -template <> -DS_D_INLINE __half init() -{ - constexpr __half_raw inf = {0x7C00}; - return __half(inf); -} - -template <> -DS_D_INLINE __half init() -{ - constexpr __half_raw neg_inf = {0xFC00}; - return __half(neg_inf); -} - -#ifdef BF16_AVAILABLE -template <> -DS_D_INLINE __nv_bfloat16 init() -{ - constexpr __nv_bfloat16_raw neg_inf = {0xFF80}; - return __nv_bfloat16(neg_inf); -} -#endif - -template <> -DS_D_INLINE __half2 init() -{ -#ifdef __HIP_PLATFORM_AMD__ - return __half2{_Float16_2{0x0000, 0x0000}}; -#else - constexpr __half2_raw zero = {0x0000, 0x0000}; - return __half2(zero); -#endif -} - -template <> -DS_D_INLINE __half2 init() -{ -#ifdef __HIP_PLATFORM_AMD__ - return __half2{_Float16_2{0x7C00, 0x7C00}}; -#else - constexpr __half2_raw inf = {0x7C00, 0x7C00}; - return __half2(inf); -#endif -} - -template <> -DS_D_INLINE __half2 init() -{ -#ifdef __HIP_PLATFORM_AMD__ - return __half2{_Float16_2{0xFC00, 0xFC00}}; -#else - constexpr __half2_raw neg_inf = {0xFC00, 0xFC00}; - return __half2(neg_inf); -#endif -} - -template <> -DS_D_INLINE int32_t init() -{ - return 0; -} - -template <> -DS_D_INLINE int32_t init() -{ - return 0x7FFFFFFF; -} - -template <> -DS_D_INLINE int32_t init() -{ - return 0x80000000; -} - -template <> -DS_D_INLINE uint32_t init() -{ - return 0; -} - -template <> -DS_D_INLINE uint32_t init() -{ - return 0xFFFFFFFF; -} - -template <> -DS_D_INLINE uint32_t init() -{ - return 0; -} - -template <> -DS_D_INLINE int64_t init() -{ - return 0; -} - -template <> -DS_D_INLINE int64_t init() -{ - return 0x7FFFFFFFFFFFFFFF; -} - -template <> -DS_D_INLINE int64_t init() -{ - return 0x8000000000000000; -} - -template <> -DS_D_INLINE uint64_t init() -{ - return 0; -} - -template <> -DS_D_INLINE uint64_t init() -{ - return 0xFFFFFFFFFFFFFFFF; -} - -template <> -DS_D_INLINE uint64_t init() -{ - return 0; -} - -template -DS_D_INLINE void init(T* data) -{ - data[0] = init(); -} - -template -DS_D_INLINE void init(T* data) -{ - data[0] = init(); - data[1] = init(); -} - -template -DS_D_INLINE void init(T* data) -{ - data[0] = init(); - data[1] = init(); - data[2] = init(); -} - -template -DS_D_INLINE void init(T* data) -{ - data[0] = init(); - data[1] = init(); - data[2] = init(); - data[3] = init(); -} - -/* -Warp reduction primitives - -`reduction_width` is an unsafe template parameter, that is that -when using `reduction_width` < hw_warp_size the warp is partitioned -into `hw_warp_size` / `reduction_width` groups of partial sums. - -If someone can figure out how to use variadic templates in a reasonable way -here (fold is C++17 only and I don't think helps and recursion feels like -huge overkill that harms readability) that would be wonderful. -*/ - -template -DS_D_INLINE void _warp(cg::thread_block_tile& warp, T* data) -{ -#pragma unroll - for (int i = 1; i < reduce_width; i *= 2) { - data[0] = element(data[0], warp.shfl_xor(data[0], i)); - } -} - -template -DS_D_INLINE void _warp(cg::thread_block_tile& warp, T* data) -{ -#pragma unroll - for (int i = 1; i < reduce_width; i *= 2) { - data[0] = element(data[0], warp.shfl_xor(data[0], i)); - data[1] = element(data[1], warp.shfl_xor(data[1], i)); - } -} - -template -DS_D_INLINE void _warp(cg::thread_block_tile& warp, T* data) -{ -#pragma unroll - for (int i = 1; i < reduce_width; i *= 2) { - data[0] = element(data[0], warp.shfl_xor(data[0], i)); - data[1] = element(data[1], warp.shfl_xor(data[1], i)); - data[2] = element(data[2], warp.shfl_xor(data[2], i)); - } -} - -template -DS_D_INLINE void _warp(cg::thread_block_tile& warp, T* data) -{ -#pragma unroll - for (int i = 1; i < reduce_width; i *= 2) { - data[0] = element(data[0], warp.shfl_xor(data[0], i)); - data[1] = element(data[1], warp.shfl_xor(data[1], i)); - data[2] = element(data[2], warp.shfl_xor(data[2], i)); - data[3] = element(data[3], warp.shfl_xor(data[3], i)); - } -} - -/* -Implementation for primary block reduction that serves both `block` and -`partitioned_block`. - -Total warps refers to the reduction width of the reduction, not -the number of warps in the block (which may exceed that -if the block is partitioned or if we do a conservative bound at -compile time). -*/ -template -DS_D_INLINE void _block(cg::thread_block& tb, - cg::thread_block_tile& warp_arg, - T* data) -{ - constexpr int elems = sizeof...(Ops); - constexpr int bytes = sizeof(T); - // Unused when `partition_size == 1` or total_warps == 1 - __shared__ T reduce_buffer[max_warps * elems]; - -#ifdef __HIP_PLATFORM_AMD__ - const int total_threads = blockDim.x * blockDim.y * blockDim.z; - const int running_warps = total_threads / hw_warp_size; -#else - const int running_warps = warp_arg.meta_group_size(); -#endif - - // Always perform warp-scope reduction - _warp(warp_arg, data); - - // If max_warps == 1 let's skip the runtime check - if (total_warps != 1) { - if (warp_arg.thread_rank() == 0) { -#pragma unroll - for (int i = 0; i < elems; i++) { - mem_access::store_shared(reduce_buffer + elems * _warp_rank() + i, data + i); - } - } - - // Synchronization inside block-uniform conditional is safe - tb.sync(); - - if (_warp_rank() == 0) { - if (warp_arg.thread_rank() < running_warps) { -#pragma unroll - for (int i = 0; i < elems; i++) { - mem_access::load_shared( - data + i, reduce_buffer + elems * warp_arg.thread_rank() + i); - } - } else { - init(data); - } - - _warp(warp_arg, data); - -#pragma unroll - for (int i = 0; i < elems; i++) { - mem_access::store_shared(reduce_buffer + elems * warp_arg.thread_rank() + i, - data + i); - } - } - - // Synchronization inside block-uniform conditional is safe - tb.sync(); - -#pragma unroll - for (int i = 0; i < elems; i++) { - mem_access::load_shared(data + i, reduce_buffer + _warp_rank() * elems + i); - } - } -} - -/* -Main API implementations. For the most part, they just convert the individual -variables into arrays, which makes working with them easier with a single -implementation. In theory, we could use the `_block` implementation as another -option, but the nature of using a pointer is a little less safe and this allows -us to obfuscate the details of the partitioned implementation. -*/ -template -DS_D_INLINE void block(cg::thread_block& tb, cg::thread_block_tile& warp, float& val) -{ - _block(tb, warp, &val); -} - -template -DS_D_INLINE void block(cg::thread_block& tb, - cg::thread_block_tile& warp, - float& val1, - float& val2) -{ - float data[2] = {val1, val2}; - _block(tb, warp, data); - val1 = data[0]; - val2 = data[1]; -} - -template -DS_D_INLINE void block(cg::thread_block& tb, - cg::thread_block_tile& warp, - float& val1, - float& val2, - float& val3) -{ - float data[3] = {val1, val2, val3}; - _block(tb, warp, data); - val1 = data[0]; - val2 = data[1]; - val3 = data[2]; -} - -template -DS_D_INLINE void block(cg::thread_block& tb, - cg::thread_block_tile& warp, - float& val1, - float& val2, - float& val3, - float& val4) -{ - float data[4] = {val1, val2, val3, val4}; - _block(tb, warp, data); - val1 = data[0]; - val2 = data[1]; - val3 = data[2]; - val4 = data[3]; -} - -/* -Note: for the partitioned blocks, the implementation does not support non-power of 2 blocks in order -to shorten block scale reduction length. -*/ -template -DS_D_INLINE void partitioned_block(cg::thread_block& tb, - cg::thread_block_tile& warp, - float& val) -{ - if (num_threads <= hw_warp_size) { - _warp(warp, &val); - } else { - constexpr int num_warps = num_threads / hw_warp_size; - _block(tb, warp, &val); - } -} - -template -DS_D_INLINE void partitioned_block(cg::thread_block& tb, - cg::thread_block_tile& warp, - float& val1, - float& val2) -{ - float data[2] = {val1, val2}; - - if (num_threads <= hw_warp_size) { - _warp(warp, data); - } else { - constexpr int num_warps = num_threads / hw_warp_size; - _block(tb, warp, data); - } - - val1 = data[0]; - val2 = data[1]; -} - -template -DS_D_INLINE void partitioned_block(cg::thread_block& tb, - cg::thread_block_tile& warp, - float& val1, - float& val2, - float& val3) -{ - float data[3] = {val1, val2, val3}; - - if (num_threads <= hw_warp_size) { - _warp(warp, data); - } else { - constexpr int num_warps = num_threads / hw_warp_size; - _block(tb, warp, data); - } - - val1 = data[0]; - val2 = data[1]; - val3 = data[2]; -} - -template -DS_D_INLINE void partitioned_block(cg::thread_block& tb, - cg::thread_block_tile& warp, - float& val1, - float& val2, - float& val3, - float& val4) -{ - float data[4] = {val1, val2, val3, val4}; - - if (num_threads <= hw_warp_size) { - _warp(warp, data); - } else { - constexpr int num_warps = num_threads / hw_warp_size; - _block(tb, warp, data); - } - - val1 = data[0]; - val2 = data[1]; - val3 = data[2]; - val4 = data[3]; -} - -/* -Arg-reduce is a specialization of the above. We only support this with a single reduction -parameter. This only works for max/min reductions. -*/ - -__align__(8) struct IdxReduceResult { - /* - NOTE: ORDERING MATTERS HERE! The idx is the least significant set of bits - and the val is the most significant. Changing the order of this declaration - will break the code. - */ - int idx; - float val; -}; - -template -DS_D_INLINE IdxReduceResult -idx_reduce(cg::thread_block& tb, cg::thread_block_tile& warp, float val, int idx) -{ - IdxReduceResult res = {idx, val}; - - // Clear out the nan. This shouldn't be an issue for our initial applications - if (isnan(val)) res.val = init(); - - // Can do float compares as integers. By packing the index into the lower bits - // we can just do a single int64 rather than a branch, compare, and select. - // One side benefit of this is that it is by nature a stable algorithm and - // will always bias ties to the higher index. - int64_t* res_as_int = reinterpret_cast(&res); - - // The way floating point compare works is normally to perform a sign comparison - // and if they match, then do a comparison of the rest of the bits as unsigned - // integers. Since we are bundling these, that means for negative values we need - // to reverse the sort order, which we can do with an XOR. - if (val < 0) { *res_as_int ^= 0x7fffffff00000000; } - - _block(tb, warp, res_as_int); - - // Sign bit is preserved, so we can check if we need to invert the mantissa back - if (res.val < 0) { *res_as_int ^= 0x7fffffff00000000; } - - return res; -} - -} // namespace reduce diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/simd.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/simd.h deleted file mode 100644 index a205026ec7c1155faf565739e7a407fcfab08c41..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/simd.h +++ /dev/null @@ -1,299 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#if (__x86_64__ || __i386__) -#include -#include -#endif - -#define TILE (128 * 1024 * 1024) -#if defined(__AVX512__) or defined(__AVX256__) -#include - -template -inline T readAs(const void* src) -{ - T res; - std::memcpy(&res, src, sizeof(T)); - return res; -} -template -inline void writeAs(void* dst, const T& val) -{ - std::memcpy(dst, &val, sizeof(T)); -} - -#define ROUND_DOWN(size, step) ((size) & ~((step) - 1)) - -#if defined(__AVX512__) -#define SIMD_STORE(a, d) _mm512_storeu_ps(a, d) -#define SIMD_LOAD(x) _mm512_loadu_ps(x) -#define SIMD_SET(x) _mm512_set1_ps(x) -#define SIMD_ADD(x, y) _mm512_add_ps(x, y) -#define SIMD_MUL(x, y) _mm512_mul_ps(x, y) -#define SIMD_FMA(x, y, c) _mm512_fmadd_ps(x, y, c) -#define SIMD_SQRT(x) _mm512_sqrt_ps(x) -#define SIMD_DIV(x, y) _mm512_div_ps(x, y) -#define SIMD_AND(x, y) _mm512_and_ps(x, y) -#define SIMD_ANDNOT(x, y) _mm512_andnot_ps(x, y) -#define SIMD_OR(x, y) _mm512_or_ps(x, y) -#define SIMD_XOR(x, y) _mm512_xor_ps(x, y) -#define SIMD_WIDTH 16 - -static __m512 load_16_bf16_as_f32(const void* data) -{ - __m256i a = readAs<__m256i>(data); // use memcpy to avoid aliasing - __m512i b = _mm512_cvtepu16_epi32(a); // convert 8 u16 to 8 u32 - __m512i c = _mm512_slli_epi32(b, 16); // logical shift left of all u32 by - // 16 bits (representing bf16->f32) - return readAs<__m512>(&c); // use memcpy to avoid aliasing -} - -static void store_16_f32_as_bf16_nearest(__m512 v, void* data) -{ - __m512i u32 = readAs<__m512i>(&v); - - // flow assuming non-nan: - - // uint32_t rounding_bias = ((U32 >> 16) & 1) + UINT32_C(0x7FFF); - __m512i b = _mm512_srli_epi32(u32, 16); - __m512i lsb_mask = _mm512_set1_epi32(0x00000001); - __m512i c = _mm512_and_si512(b, lsb_mask); - __m512i bias_constant = _mm512_set1_epi32(0x00007fff); - __m512i rounding_bias = _mm512_add_epi32(c, bias_constant); - - // uint16_t res = static_cast((U32 + rounding_bias) >> 16); - __m512i d = _mm512_add_epi32(u32, rounding_bias); - __m512i e = _mm512_srli_epi32(d, 16); - __m256i non_nan_res = _mm512_cvtusepi32_epi16(e); - - // handle nan (exp is all 1s and mantissa != 0) - // if ((x & 0x7fffffffU) > 0x7f800000U) - __m512i mask_out_sign = _mm512_set1_epi32(0x7fffffff); - __m512i non_sign_bits = _mm512_and_si512(u32, mask_out_sign); - __m512i nan_threshold = _mm512_set1_epi32(0x7f800000); - __mmask16 nan_mask = _mm512_cmp_epi32_mask(non_sign_bits, nan_threshold, _MM_CMPINT_GT); - - // mix in results with nans as needed - __m256i nans = _mm256_set1_epi16(0x7fc0); - __m256i res = _mm256_mask_mov_epi16(non_nan_res, nan_mask, nans); - - writeAs(data, res); -} -#define SIMD_LOAD_BF16(x) load_16_bf16_as_f32(x) -#define SIMD_STORE_BF16(x, d) store_16_f32_as_bf16_nearest(d, x) - -#define SIMD_LOAD_FP16(x) _mm512_cvtph_ps(_mm256_castps_si256(_mm256_loadu_ps(x))) -#define SIMD_STORE_FP16(x, d) \ - _mm256_store_ps(x, _mm256_castsi256_ps(_mm512_cvtps_ph(d, _MM_FROUND_TO_NEAREST_INT))) - -#define INTV __m256i -#elif defined(__AVX256__) -#define SIMD_STORE(a, d) _mm256_storeu_ps(a, d) -#define SIMD_LOAD(x) _mm256_loadu_ps(x) -#define SIMD_SET(x) _mm256_set1_ps(x) -#define SIMD_ADD(x, y) _mm256_add_ps(x, y) -#define SIMD_MUL(x, y) _mm256_mul_ps(x, y) -#define SIMD_FMA(x, y, c) _mm256_fmadd_ps(x, y, c) -#define SIMD_SQRT(x) _mm256_sqrt_ps(x) -#define SIMD_DIV(x, y) _mm256_div_ps(x, y) -#define SIMD_AND(x, y) _mm256_and_ps(x, y) -#define SIMD_ANDNOT(x, y) _mm256_andnot_ps(x, y) -#define SIMD_OR(x, y) _mm256_or_ps(x, y) -#define SIMD_XOR(x, y) _mm256_xor_ps(x, y) -#define SIMD_WIDTH 8 - -#define SIMD_LOAD_BF16(x) static_assert(false && "AVX256 does not support BFloat16") -#define SIMD_STORE_BF16(x, d) static_assert(false && "AVX256 does not support BFloat16") -#define SIMD_LOAD_FP16(x) _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*)x)) -#define SIMD_STORE_FP16(x, d) \ - _mm_store_ps(x, _mm_castsi128_ps(_mm256_cvtps_ph(d, _MM_FROUND_TO_NEAREST_INT))) - -#define INTV __m128i -#endif - -union AVX_Data { -#if defined(__AVX512__) - __m512 data; -#elif defined(__AVX256__) - __m256 data; -#endif - // float data_f[16]; -}; - -template -inline typename std::enable_if_t, void> simd_store(T* dst, - AVX_Data* src) -{ - size_t width = SIMD_WIDTH; -#pragma unroll - for (size_t i = 0; i < span; ++i) { SIMD_STORE_FP16((float*)(dst + width * i), src[i].data); } -} - -template -inline typename std::enable_if_t, void> simd_store(T* dst, - AVX_Data* src) -{ -#ifdef __AVX512__ - size_t width = SIMD_WIDTH; -#pragma unroll - for (size_t i = 0; i < span; ++i) { SIMD_STORE_BF16((float*)(dst + width * i), src[i].data); } -#else - throw std::runtime_error("AVX512 required for BFloat16"); -#endif -} - -template -inline typename std::enable_if_t, void> simd_store(T* dst, AVX_Data* src) -{ - size_t width = SIMD_WIDTH; -#pragma unroll - for (size_t i = 0; i < span; ++i) { SIMD_STORE(dst + width * i, src[i].data); } -} - -template -inline typename std::enable_if_t, void> simd_load(AVX_Data* dst, - T* src) -{ - size_t width = SIMD_WIDTH; -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_LOAD_FP16((float*)(src + width * i)); } -} - -template -inline typename std::enable_if_t, void> simd_load(AVX_Data* dst, - T* src) -{ -#ifdef __AVX512__ - size_t width = SIMD_WIDTH; -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_LOAD_BF16((float*)(src + width * i)); } -#else - throw std::runtime_error("AVX512 required for BFloat16"); -#endif -} - -template -inline typename std::enable_if_t, void> simd_load(AVX_Data* dst, T* src) -{ - size_t width = SIMD_WIDTH; -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_LOAD(src + width * i); } -} - -template -inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data src_m_r, AVX_Data* src_a) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { - dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r.data, src_a[i].data); - } -} -template -inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data src_m_r, AVX_Data src_a) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { - dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r.data, src_a.data); - } -} -template -inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data* src_m_r, AVX_Data* src_a) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { - dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r[i].data, src_a[i].data); - } -} -template -inline void simd_sqrt(AVX_Data* dst, AVX_Data* src) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_SQRT(src[i].data); } -} -template -inline void simd_add(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ADD(src_a_l[i].data, src_a_r.data); } -} -template -inline void simd_add(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ADD(src_a_l[i].data, src_a_r[i].data); } -} -template -inline void simd_mul(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_MUL(src_a_l[i].data, src_a_r.data); } -} -template -inline void simd_mul(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_MUL(src_a_l[i].data, src_a_r[i].data); } -} -template -inline void simd_div(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_DIV(src_a_l[i].data, src_a_r[i].data); } -} -template -inline void simd_and(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_AND(src_a_l[i].data, src_a_r.data); } -} -template -inline void simd_and(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_AND(src_a_l[i].data, src_a_r[i].data); } -} -template -inline void simd_andnot(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ANDNOT(src_a_l[i].data, src_a_r.data); } -} -template -inline void simd_andnot(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { - dst[i].data = SIMD_ANDNOT(src_a_l[i].data, src_a_r[i].data); - } -} -template -inline void simd_or(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_OR(src_a_l[i].data, src_a_r.data); } -} -template -inline void simd_or(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_OR(src_a_l[i].data, src_a_r[i].data); } -} -template -inline void simd_xor(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_XOR(src_a_l[i].data, src_a_r.data); } -} -template -inline void simd_xor(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_XOR(src_a_l[i].data, src_a_r[i].data); } -} - -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/softmax.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/softmax.h deleted file mode 100644 index b6320a3f91bec40b63d7f96390a6031a61a9d2cc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/softmax.h +++ /dev/null @@ -1,65 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include -#include -#include "custom_cuda_layers.h" - -#include - -using namespace std; - -template -class Softmax { -public: - struct Config { - size_t batchSize; - size_t heads; - size_t seq_length; - size_t prob_depth; - float temperature; - bool mem_alloc; - Config(size_t batch, size_t h, size_t seq, int prob_size = 0, bool mem_alloc = false) - : batchSize(batch), - heads(h), - seq_length(seq), - prob_depth(prob_size), - temperature(1.0), - mem_alloc(mem_alloc) - { - } - }; - - Softmax(Config config) : config_(config) {} - - ~Softmax() {} - - void Forward(int bsz, T* vals, const T* attn_mask, cudaStream_t& stream) - { - launch_attn_softmax(vals, attn_mask, bsz, config_.heads, config_.seq_length, stream); - } - - void Backward(int bsz, T* out_grad, const T* soft_out, cudaStream_t stream) - { - launch_attn_softmax_backward_v2( - out_grad, soft_out, bsz, config_.heads, config_.seq_length, stream); - } - - inline size_t GetProbDepth() const { return config_.prob_depth; } - - inline size_t GetBatchSize() const { return config_.batchSize; } - - inline size_t GetNumHeads() const { return config_.heads; } - - inline size_t GetSeqLength() const { return config_.seq_length; } - - inline void SetSeqLength(size_t seq_len) { config_.seq_length = seq_len; } - -private: - Config config_; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/strided_batch_gemm.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/strided_batch_gemm.h deleted file mode 100644 index 9767fcf589b82356230016abc2d183b74ab621e2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/strided_batch_gemm.h +++ /dev/null @@ -1,205 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include -#include -#include "context.h" - -template -class StridedBatchGemm { -public: - struct Config { - int batch_size; - int m; - int n; - int k; - float alpha; - float beta; - cublasOperation_t op_A; - cublasOperation_t op_B; - std::array gemm_algos; - - Config(int batch, - int mm, - int nn, - int kk, - float param_alpha, - float param_beta, - cublasOperation_t opA, - cublasOperation_t opB, - const std::array& algos) - : batch_size(batch), - m(mm), - n(nn), - k(kk), - alpha(param_alpha), - beta(param_beta), - op_A(opA), - op_B(opB), - gemm_algos(algos) - { - } - void SetConfig(int mm, int nn, int kk) - { - m = mm; - n = nn; - k = kk; - } - }; - - StridedBatchGemm(const Config& config) : _config(config) {} - - virtual ~StridedBatchGemm() {} - - void Forward(int bsz, T* output, const T* _buffer_a, const T* _buffer_b, cublasHandle_t handle) - { - int stride_a = _config.m * _config.k; - int stride_b = _config.n * _config.k; - int stride_c = _config.m * _config.n; - - cublas_strided_batched_gemm(handle, - _config.m, - _config.n, - _config.k, - &_config.alpha, - &_config.beta, - _buffer_a, - _buffer_b, - output, - _config.op_A, - _config.op_B, - stride_a, - stride_b, - stride_c, - bsz, -// TODO HIP: Remove backward compatibility for torch<=2.0 in future -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo(_config.gemm_algos[0])); -#else - cublasGemmAlgo_t(_config.gemm_algos[0])); -#endif - } - - void ForwardPlusSave(T* output, const T* _buffer_a, const T* _buffer_b, cublasHandle_t handle) - { - int stride_a = _config.m * _config.k; - int stride_b = _config.n * _config.k; - int stride_c = _config.m * _config.n; - - cublas_strided_batched_gemm(handle, - _config.m, - _config.n, - _config.k, - &_config.alpha, - &_config.beta, - _buffer_a, - _buffer_b, - output, - _config.op_A, - _config.op_B, - stride_a, - stride_b, - stride_c, - _config.batch_size, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo(_config.gemm_algos[0])); -#else - cublasGemmAlgo_t(_config.gemm_algos[0])); -#endif - - k_buf = _buffer_a; - q_buf = _buffer_b; - } - - void Backward(int bsz, - const T* d_output, - const T* _buffer_a, - const T* _buffer_b, - cublasHandle_t handle, - T* inpGradA = nullptr, - T* inpGradB = nullptr) - { - int mb = (_config.op_A == CUBLAS_OP_T ? _config.k : _config.m); - int kb = (_config.op_A == CUBLAS_OP_T ? _config.m : _config.k); - - int stride_a = mb * _config.n; - int stride_b = _config.n * kb; - int stride_c = _config.m * _config.k; - - // B need to transpose. - cublasOperation_t op_b = (_config.op_B == CUBLAS_OP_T ? CUBLAS_OP_N : CUBLAS_OP_T); - - // Calculate d_A. - cublas_strided_batched_gemm(handle, - mb, - kb, - _config.n, - &_config.alpha, - &_config.beta, - (_config.op_A == CUBLAS_OP_T ? _buffer_b : d_output), - (_config.op_A == CUBLAS_OP_T ? d_output : _buffer_b), - inpGradA, - CUBLAS_OP_N, - op_b, - stride_a, - stride_b, - stride_c, - bsz, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo(_config.gemm_algos[1])); -#else - cublasGemmAlgo_t(_config.gemm_algos[1])); -#endif - - // A need to transpose. - cublasOperation_t op_a = (_config.op_A == CUBLAS_OP_T ? CUBLAS_OP_N : CUBLAS_OP_T); - - stride_a = _config.m * _config.k; - stride_b = _config.m * _config.n; - stride_c = _config.n * _config.k; - - // Calculate d_B. - cublas_strided_batched_gemm(handle, - _config.k, - _config.n, - _config.m, - &_config.alpha, - &_config.beta, - _buffer_a, - d_output, - inpGradB, - op_a, - CUBLAS_OP_N, - stride_a, - stride_b, - stride_c, - bsz, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo(_config.gemm_algos[2])); -#else - cublasGemmAlgo_t(_config.gemm_algos[2])); -#endif - } - - inline int GetN() const { return _config.k; } - - inline const T* GetBufferA() const { return k_buf; } - - inline const T* GetBufferB() const { return q_buf; } - - inline void SetConfig(int m, int n, int k) { _config.SetConfig(m, n, k); } - -private: - Config _config; - const T* q_buf; - const T* k_buf; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/type_shim.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/type_shim.h deleted file mode 100644 index 78ade5a06daa12cf628e9966ce2ba8424a6e96ba..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/includes/type_shim.h +++ /dev/null @@ -1,124 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* Taken from NVIDIA/apex commit 855808f3fc268e9715d613f3c2e56469d8c986d8 */ -#include - -// Forward/backward compatibility hack around -// https://github.com/pytorch/pytorch/commit/3aeb78079bcd68282fe9117088e138b77318e288 -// pending more future-proof guidance from upstream. -// struct TypeShim -// { -// const at::Type& payload; -// TypeShim(const at::Type& type) : payload(type) {} -// // Enable trivial conversion to a const at::Type& for pre-3aeb78 -// operator const at::Type&(){ return payload; }; -// // Enable dispatch switch statements to take *this directly for post-3aeb78 -// //operator at::ScalarType(){ return payload.; }; -// }; - -#define DISPATCH_FLOAT_AND_HALF(TYPE, LEVEL, NAME, ...) \ - switch (TYPE) { \ - case at::ScalarType::Float: { \ - using scalar_t_##LEVEL = float; \ - __VA_ARGS__; \ - break; \ - } \ - case at::ScalarType::Half: { \ - using scalar_t_##LEVEL = at::Half; \ - __VA_ARGS__; \ - break; \ - } \ - case at::ScalarType::BFloat16: { \ - using scalar_t_##LEVEL = at::BFloat16; \ - __VA_ARGS__; \ - break; \ - } \ - default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ - } - -#define DISPATCH_DOUBLE_FLOAT_AND_HALF(TYPE, LEVEL, NAME, ...) \ - switch (TYPE) { \ - case at::ScalarType::Double: { \ - using scalar_t_##LEVEL = double; \ - __VA_ARGS__; \ - break; \ - } \ - case at::ScalarType::Float: { \ - using scalar_t_##LEVEL = float; \ - __VA_ARGS__; \ - break; \ - } \ - case at::ScalarType::Half: { \ - using scalar_t_##LEVEL = at::Half; \ - __VA_ARGS__; \ - break; \ - } \ - case at::ScalarType::BFloat16: { \ - using scalar_t_##LEVEL = at::BFloat16; \ - __VA_ARGS__; \ - break; \ - } \ - default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ - } - -#define DISPATCH_DOUBLE_AND_FLOAT(TYPE, LEVEL, NAME, ...) \ - switch (TYPE) { \ - case at::ScalarType::Double: { \ - using scalar_t_##LEVEL = double; \ - __VA_ARGS__; \ - break; \ - } \ - case at::ScalarType::Float: { \ - using scalar_t_##LEVEL = float; \ - __VA_ARGS__; \ - break; \ - } \ - default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ - } - -template -__device__ __forceinline__ T -reduce_block_into_lanes(T* x, - T val, - int lanes = 1, - bool share_result = false) // lanes is intended to be <= 32. -{ - int tid = threadIdx.x + threadIdx.y * blockDim.x; - int blockSize = blockDim.x * blockDim.y; // blockSize is intended to be a multiple of 32. - - if (blockSize >= 64) { - x[tid] = val; - __syncthreads(); - } - -#pragma unroll - for (int i = (blockSize >> 1); i >= 64; i >>= 1) { - if (tid < i) x[tid] = x[tid] + x[tid + i]; - __syncthreads(); - } - - T final; - - if (tid < 32) { - if (blockSize >= 64) - final = x[tid] + x[tid + 32]; - else - final = val; - // __SYNCWARP(); - -#pragma unroll - for (int i = 16; i >= lanes; i >>= 1) - final = final + __shfl_down_sync(0xffffffff, final, i); - } - - if (share_result) { - if (tid < lanes) x[tid] = final; // EpilogueOp - // Make sure the smem result is visible to all warps. - __syncthreads(); - } - - return final; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda.cpp deleted file mode 100644 index c629b93517d278f586613dd1d32cd4ac4dc2867a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda.cpp +++ /dev/null @@ -1,113 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include - -// CUDA forward declaration -void fused_lamb_cuda(at::Tensor& p, - at::Tensor& p_copy, - at::Tensor& m, - at::Tensor& v, - at::Tensor& g, - float lr, - float beta1, - float beta2, - float max_coeff, - float min_coeff, - float eps, - float grad_scale, - int step, - int mode, - int bias_correction, - float decay, - at::Tensor& w_l2_i, - at::Tensor& u_l2_i, - at::Tensor& lamb_coeff_val); - -#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") -#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") -#define CHECK_INPUT(x) \ - CHECK_CUDA(x); \ - CHECK_CONTIGUOUS(x) - -// C++ interface -at::Tensor lamb(at::Tensor& p, - at::Tensor& p_copy, - at::Tensor& m, - at::Tensor& v, - at::Tensor& g, - float lr, - float beta1, - float beta2, - float max_coeff, - float min_coeff, - float eps, - float grad_scale, - int step, - int mode, - int bias_correction, - float decay) -{ - CHECK_INPUT(p); - if (p_copy.numel() > 0) CHECK_INPUT(p_copy); - CHECK_INPUT(m); - CHECK_INPUT(v); - CHECK_INPUT(g); - int64_t num_elem = p.numel(); - AT_ASSERTM(m.numel() == num_elem, "number of elements in m and p tensors should be equal"); - AT_ASSERTM(v.numel() == num_elem, "number of elements in v and p tensors should be equal"); - AT_ASSERTM(g.numel() == num_elem, "number of elements in g and p tensors should be equal"); - AT_ASSERTM( - p_copy.numel() == num_elem || p_copy.numel() == 0, - "number of elements in p_copy and p tensors should be equal, or p_copy should be empty"); - - // intermediate for weight L2 reduction - // make sure that the threads per block is at least 512 during the kernel launch otherwise the - // behaviour is unexpected - at::Tensor w_l2_i = at::empty( - {512}, - p.options().dtype(p.type().scalarType() == at::ScalarType::Half ? at::ScalarType::Float - : p.type().scalarType())); - - // intermediate for update L2 reduction - // make sure that the threads per block is at least 512 during the kernel launch otherwise the - // behaviour is unexpected - at::Tensor u_l2_i = at::empty( - {512}, - p.options().dtype(p.type().scalarType() == at::ScalarType::Half ? at::ScalarType::Float - : p.type().scalarType())); - - at::Tensor lamb_coeff_val = at::empty( - {1}, - p.options().dtype(p.type().scalarType() == at::ScalarType::Half ? at::ScalarType::Float - : p.type().scalarType())); - - fused_lamb_cuda(p, - p_copy, - m, - v, - g, - lr, - beta1, - beta2, - max_coeff, - min_coeff, - eps, - grad_scale, - step, - mode, - bias_correction, - decay, - w_l2_i, - u_l2_i, - lamb_coeff_val); - - return lamb_coeff_val; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("lamb", &lamb, "Adam optimized CUDA implementation with LAMB."); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda_kernel.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda_kernel.cu deleted file mode 100644 index d9bacae73457fe1e757f1f6410d64d18ebcf31bb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda_kernel.cu +++ /dev/null @@ -1,478 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include -#include -#include -#include "ATen/ATen.h" -#include "ATen/TensorUtils.h" -#include "ATen/cuda/CUDAContext.h" -#include "ATen/cuda/detail/IndexUtils.cuh" -// #include "ATen/Type.h" -#include "ATen/AccumulateType.h" - -#include - -// #include -#if defined(__HIP_PLATFORM_AMD__) && HIP_VERSION > 305 -#include -#else -#include -#endif -#include -#include - -namespace cg = cooperative_groups; - -// Utility class used to avoid linker errors with extern -// unsized shared memory arrays with templated type -namespace { -// This is the un-specialized struct. Note that we prevent instantiation of this -// struct by putting an undefined symbol in the function body so it won't compile. -template -struct SharedMemory { - // Ensure that we won't compile any un-specialized types - __device__ inline operator T*() - { -#ifndef _WIN32 - extern __device__ void error(void); - error(); -#endif - return NULL; - } -}; - -template <> -struct SharedMemory { - __device__ inline operator float*() - { - extern __shared__ float s_float[]; - return s_float; - } -}; - -template <> -struct SharedMemory { - __device__ inline operator double*() - { - extern __shared__ double s_double[]; - return s_double; - } -}; -} // namespace - -#include "type_shim.h" - -typedef enum { - ADAM_MODE_0 = 0, // eps under square root - ADAM_MODE_1 = 1 // eps outside square root -} adamMode_t; - -// s_a and s_b are in shared memory -// g_a and g_b are in shared memory -template -__device__ void reduce_block_in_shared_memory(T* s_a, T* s_b, T* g_a, T* g_b) -{ - // Handle to thread block group - cg::thread_block cta = cg::this_thread_block(); - - // perform block reduction in shared memory, - unsigned int tid = cta.thread_rank(); - - T a_sum = s_a[tid]; - T b_sum = s_b[tid]; - - cg::sync(cta); - - // do reduction in shared mem - if ((blockSize >= 512) && (tid < 256)) { - s_a[tid] = a_sum = a_sum + s_a[tid + 256]; - s_b[tid] = b_sum = b_sum + s_b[tid + 256]; - } - - cg::sync(cta); - - if ((blockSize >= 256) && (tid < 128)) { - s_a[tid] = a_sum = a_sum + s_a[tid + 128]; - s_b[tid] = b_sum = b_sum + s_b[tid + 128]; - } - - cg::sync(cta); - - if ((blockSize >= 128) && (tid < 64)) { - s_a[tid] = a_sum = a_sum + s_a[tid + 64]; - s_b[tid] = b_sum = b_sum + s_b[tid + 64]; - } - - cg::sync(cta); - -#if (__CUDA_ARCH__ >= 300) || (defined(__HIP_PLATFORM_AMD__) && HIP_VERSION >= 502) - if (tid < 32) { - cg::coalesced_group active = cg::coalesced_threads(); - - // Fetch final intermediate sum from 2nd warp - if (blockSize >= 64) { - a_sum = a_sum + s_a[tid + 32]; - b_sum = b_sum + s_b[tid + 32]; - } - - // Reduce final warp using shuffle - for (int offset = warpSize / 2; offset > 0; offset /= 2) { - a_sum += active.shfl_down(a_sum, offset); - b_sum += active.shfl_down(b_sum, offset); - } - } -#else - if ((blockSize >= 64) && (tid < 32)) { - s_a[tid] = a_sum = a_sum + s_a[tid + 32]; - s_b[tid] = b_sum = b_sum + s_b[tid + 32]; - } - - cg::sync(cta); - - if ((blockSize >= 32) && (tid < 16)) { - s_a[tid] = a_sum = a_sum + s_a[tid + 16]; - s_b[tid] = b_sum = b_sum + s_b[tid + 16]; - } - - cg::sync(cta); - - if ((blockSize >= 16) && (tid < 8)) { - s_a[tid] = a_sum = a_sum + s_a[tid + 8]; - s_b[tid] = b_sum = b_sum + s_b[tid + 8]; - } - - cg::sync(cta); - - if ((blockSize >= 8) && (tid < 4)) { - s_a[tid] = a_sum = a_sum + s_a[tid + 4]; - s_b[tid] = b_sum = b_sum + s_b[tid + 4]; - } - - cg::sync(cta); - - if ((blockSize >= 4) && (tid < 2)) { - s_a[tid] = a_sum = a_sum + s_a[tid + 2]; - s_b[tid] = b_sum = b_sum + s_b[tid + 2]; - } - - cg::sync(cta); - - if ((blockSize >= 2) && (tid < 1)) { - s_a[tid] = a_sum = a_sum + s_a[tid + 1]; - s_b[tid] = b_sum = b_sum + s_b[tid + 1]; - } - - cg::sync(cta); - -#endif - - // write result for this block to global mem - if (tid == 0) { - g_a[blockIdx.x] = (T)a_sum; - g_b[blockIdx.x] = (T)b_sum; - } -} - -template -__device__ void reduce_two_vectors_in_register(T a, T b, T* g_a, T* g_b) -{ - const int threadIdInBlock = cg::this_thread_block().thread_rank(); - - T* s_a = SharedMemory(); - T* s_b = SharedMemory() + cg::this_thread_block().size(); - - s_a[threadIdInBlock] = a; - s_b[threadIdInBlock] = b; - - reduce_block_in_shared_memory(s_a, s_b, g_a, g_b); -} - -template -__global__ void lamb_cuda_kernel_part1( - T* __restrict__ p, - GRAD_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed - T* __restrict__ m, - T* __restrict__ v, - const GRAD_T* __restrict__ g, - const float b1, - const float b2, - const float eps, - const float grad_scale, - const float step_size, - const size_t tsize, - adamMode_t mode, - const float decay, - T* __restrict__ w_l2_i, - T* __restrict__ u_l2_i) -{ - // Assuming 2D grids and 2D blocks - const int blockId = gridDim.x * blockIdx.y + blockIdx.x; - const int threadsPerBlock = blockDim.x * blockDim.y; - const int threadIdInBlock = cg::this_thread_block().thread_rank(); - const int i = (blockId * threadsPerBlock + threadIdInBlock); - const int totThreads = gridDim.x * gridDim.y * threadsPerBlock; - - T reg_w = 0; - T reg_u = 0; - - for (int j = i; j < tsize; j += totThreads) { - T scaled_grad = g[j] / grad_scale; - T pj = p[j]; - m[j] = b1 * m[j] + (1 - b1) * scaled_grad; - v[j] = b2 * v[j] + (1 - b2) * scaled_grad * scaled_grad; - float denom; - if (mode == ADAM_MODE_0) - denom = sqrtf(v[j] + eps); - else // Mode 1 - denom = sqrtf(v[j]) + eps; - T update = (m[j] / denom) + (decay * p[j]); - - reg_u += update * update; - reg_w += pj * pj; - } - - reduce_two_vectors_in_register(reg_w, reg_u, w_l2_i, u_l2_i); -} - -template -__global__ void lamb_cuda_kernel_part2(const size_t tsize, T* __restrict__ g_a, T* __restrict__ g_b) -{ - T* s_a = SharedMemory(); - T* s_b = SharedMemory() + cg::this_thread_block().size(); - - const int threadIdInBlock = cg::this_thread_block().thread_rank(); - - s_a[threadIdInBlock] = g_a[threadIdInBlock]; - s_b[threadIdInBlock] = g_b[threadIdInBlock]; - - if (threadIdInBlock >= tsize) { - s_a[threadIdInBlock] = 0.0; - s_b[threadIdInBlock] = 0.0; - } - - reduce_block_in_shared_memory(s_a, s_b, g_a, g_b); -} - -template -__global__ void lamb_cuda_kernel_part3( - T* __restrict__ p, - GRAD_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed - T* __restrict__ m, - T* __restrict__ v, - const GRAD_T* __restrict__ g, - const float b1, - const float b2, - const float max_coeff, - const float min_coeff, - const float eps, - const float grad_scale, - const float step_size, - const size_t tsize, - adamMode_t mode, - const float decay, - T* __restrict__ w_l2_i, - T* __restrict__ u_l2_i, - T* __restrict__ lamb_coeff_val) -{ - // Assuming 2D grids and 2D blocks - const int blockId = gridDim.x * blockIdx.y + blockIdx.x; - const int threadsPerBlock = blockDim.x * blockDim.y; - const int threadIdInBlock = cg::this_thread_block().thread_rank(); - const int i = (blockId * threadsPerBlock + threadIdInBlock); - const int totThreads = gridDim.x * gridDim.y * threadsPerBlock; - - T reg_w = sqrtf(w_l2_i[0]); - T reg_u = sqrtf(u_l2_i[0]); - - float lamb_coeff = 1.0; - - if (reg_w != 0 && reg_u != 0) { - lamb_coeff = reg_w / reg_u; - if (lamb_coeff > max_coeff) { lamb_coeff = max_coeff; } - if (lamb_coeff < min_coeff) { lamb_coeff = min_coeff; } - } - - if (blockId == 0 && threadIdInBlock == 0) { - lamb_coeff_val[0] = lamb_coeff; - // printf("Cuda Lamb Coeff is %.6f \n",lamb_coeff); - } - - for (int j = i; j < tsize; j += totThreads) { - T pj = (float)p[j]; - T mj = m[j]; - T vj = v[j]; - float denom; - if (mode == ADAM_MODE_0) - denom = sqrtf(vj + eps); - else // Mode 1 - denom = sqrtf(vj) + eps; - T update = (mj / denom) + (decay * pj); - - pj = pj - (step_size * lamb_coeff * update); - p[j] = pj; - if (p_copy != NULL) p_copy[j] = (GRAD_T)pj; - } -} - -void fused_lamb_cuda(at::Tensor& p, - at::Tensor& p_copy, - at::Tensor& m, - at::Tensor& v, - at::Tensor& g, - float lr, - float beta1, - float beta2, - float max_coeff, - float min_coeff, - float eps, - float grad_scale, - int step, - int mode, - int bias_correction, - float decay, - at::Tensor& w_l2_i, - at::Tensor& u_l2_i, - at::Tensor& lamb_coeff) -{ - // using namespace at; - - // Get tensor size - int tsize = p.numel(); - // Determine #threads and #blocks - const int threadsPerBlock = 512; - int num_blocks = (tsize + threadsPerBlock - 1) / threadsPerBlock; - if (num_blocks > 512) num_blocks = 512; - - int smemsize = 0; - - if (p.type().scalarType() == at::ScalarType::Double) - smemsize = 2 * threadsPerBlock * sizeof(double); - else - smemsize = 2 * threadsPerBlock * sizeof(float); - - const dim3 blocks(num_blocks); - const dim3 threads(threadsPerBlock); - - AT_ASSERTM(at::cuda::detail::canUse32BitIndexMath(p), - "parameter tensor is too large to be indexed with int32"); - // Constants - float step_size = 0; - if (bias_correction == 1) { - const float bias_correction1 = 1 - std::pow(beta1, step); - const float bias_correction2 = 1 - std::pow(beta2, step); - step_size = lr * std::sqrt(bias_correction2) / bias_correction1; - } else { - step_size = lr; - } - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - - if (g.type().scalarType() == at::ScalarType::Half) { - // all other values should be fp32 for half gradients - AT_ASSERTM(p.type().scalarType() == at::ScalarType::Float, - "expected parameter to be of float type"); - // dispatch is done on the gradient type - using namespace at; // prevents "toString is undefined" errors - AT_DISPATCH_FLOATING_TYPES_AND_HALF( - g.scalar_type(), "lamb_cuda_kernel", ([&] { - using accscalar_t = at::acc_type; - - lamb_cuda_kernel_part1 - <<>>( - p.data(), - p_copy.numel() ? p_copy.data() : NULL, - m.data(), - v.data(), - g.data(), - beta1, - beta2, - eps, - grad_scale, - step_size, - tsize, - (adamMode_t)mode, - decay, - w_l2_i.data(), - u_l2_i.data()); - - lamb_cuda_kernel_part2 - <<<1, threadsPerBlock, smemsize, stream>>>( - num_blocks, w_l2_i.data(), u_l2_i.data()); - - lamb_cuda_kernel_part3 - <<>>( - p.data(), - p_copy.numel() ? p_copy.data() : NULL, - m.data(), - v.data(), - g.data(), - beta1, - beta2, - max_coeff, - min_coeff, - eps, - grad_scale, - step_size, - tsize, - (adamMode_t)mode, - decay, - w_l2_i.data(), - u_l2_i.data(), - lamb_coeff.data()); - })); - } else { - using namespace at; - AT_DISPATCH_FLOATING_TYPES( - g.scalar_type(), "lamb_cuda_kernel", ([&] { - lamb_cuda_kernel_part1 - <<>>( - p.data(), - NULL, // don't output p_copy for fp32, it's wasted write - m.data(), - v.data(), - g.data(), - beta1, - beta2, - eps, - grad_scale, - step_size, - tsize, - (adamMode_t)mode, - decay, - w_l2_i.data(), - u_l2_i.data()); - - lamb_cuda_kernel_part2 - <<<1, threadsPerBlock, smemsize, stream>>>( - num_blocks, w_l2_i.data(), u_l2_i.data()); - - lamb_cuda_kernel_part3 - <<>>( - p.data(), - NULL, // don't output p_copy for fp32, it's wasted write - m.data(), - v.data(), - g.data(), - beta1, - beta2, - max_coeff, - min_coeff, - eps, - grad_scale, - step_size, - tsize, - (adamMode_t)mode, - decay, - w_l2_i.data(), - u_l2_i.data(), - lamb_coeff.data()); - })); - } - C10_CUDA_CHECK(cudaGetLastError()); -} - -// template __device__ void reduce_two_vectors_in_register(float a, float b, float* g_a, -// float* g_b, cg::grid_group &cgg); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/cpu_lion.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/cpu_lion.cpp deleted file mode 100644 index c5cf3e9e9235662a84a4be2c634e3e6d5ebb735b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/cpu_lion.cpp +++ /dev/null @@ -1,13 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "cpu_lion.h" - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("lion_update", &ds_lion_step, "DeepSpeed CPU Lion update (C++)"); - m.def("create_lion", &create_lion_optimizer, "DeepSpeed CPU Lion (C++)"); - m.def("destroy_lion", &destroy_lion_optimizer, "DeepSpeed CPU Lion destroy (C++)"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/cpu_lion_impl.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/cpu_lion_impl.cpp deleted file mode 100644 index 6a98162314f9f396f2f58d872405856a18011fc8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/cpu_lion_impl.cpp +++ /dev/null @@ -1,221 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include "cpu_lion.h" - -using namespace std::string_literals; -static std::unordered_map> s_optimizers; - -// C++ interface - -template -void Lion_Optimizer::Step_1(ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg, - size_t _param_size) -{ - size_t rounded_size = 0; -#if defined(__AVX512__) or defined(__AVX256__) - Step_AVX<1>(&rounded_size, _params, grads, _exp_avg, _param_size); -#endif - if (_param_size > rounded_size) { - float betta1_minus1 = 1 - _betta1; - float betta2_minus1 = 1 - _betta2; - - float alpha = _alpha; - float after_decay = 1 - alpha * _weight_decay; - - for (size_t t = rounded_size; t < _param_size; t += TILE) { - size_t copy_size = TILE; - if ((t + TILE) > _param_size) copy_size = _param_size - t; - size_t offset = copy_size + t; -#pragma omp parallel for - for (size_t k = t; k < offset; k++) { - float grad = (float)grads[k]; - float param = (float)_params[k]; - float momentum = _exp_avg[k]; - float tmp = momentum * _betta1; - tmp = grad * betta1_minus1 + tmp; - // Rely on portable C++ methods to manipulate the sign bit of a floating-point - // number. - tmp = -std::copysignf(alpha, tmp); - if (_weight_decay > 0) { - param = param * after_decay + tmp; - } else { - param = param + tmp; - } - momentum = momentum * _betta2; - momentum = grad * betta2_minus1 + momentum; - _params[k] = param; - _exp_avg[k] = momentum; - } - } - } -} - -template -void Lion_Optimizer::Step_4(ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg, - size_t _param_size) -{ - size_t rounded_size = 0; -#if defined(__AVX512__) or defined(__AVX256__) - Step_AVX<4>(&rounded_size, _params, grads, _exp_avg, _param_size); -#endif - if (_param_size > rounded_size) - Step_1((_params + rounded_size), - (grads + rounded_size), - (_exp_avg + rounded_size), - (_param_size - rounded_size)); -} - -int create_lion_optimizer(int optimizer_id, - float alpha, - float betta1, - float betta2, - float weight_decay, - bool should_log) -{ - auto opt = std::make_shared(alpha, betta1, betta2, weight_decay); - - s_optimizers[optimizer_id] = opt; - - if (should_log) { - std::string avx_type = ""; -#if defined(__AVX512__) - avx_type = "AVX512"; -#else -#if defined(__AVX256__) - avx_type = "AVX2"; -#else - avx_type = "scalar"; -#endif -#endif - - printf("Lion Optimizer #%d is created with %s arithmetic capability.\n", - optimizer_id, - avx_type.c_str()); - printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f\n", - alpha, - betta1, - betta2, - weight_decay); - } - - return 0; -} - -template -void Lion_Optimizer::Step_8(ds_params_precision_t* _params, - ds_params_precision_t* grads, - ds_state_precision_t* _exp_avg, - size_t _param_size) -{ - size_t rounded_size = 0; -#if defined(__AVX512__) or defined(__AVX256__) - Step_AVX<8>(&rounded_size, _params, grads, _exp_avg, _param_size); -#endif - if (_param_size > rounded_size) - Step_4((_params + rounded_size), - (grads + rounded_size), - (_exp_avg + rounded_size), - (_param_size - rounded_size)); -} - -template -void step_invoker(std::shared_ptr opt, - void* _params, - void* grads, - void* _exp_avg, - size_t _param_size) -{ - opt->Step_8((ds_params_precision_t*)(_params), - (ds_params_precision_t*)(grads), - (ds_state_precision_t*)(_exp_avg), - _param_size); -} - -std::map, - std::function, void*, void*, void*, size_t)>> - invokers; - -// Fill map with template functions for each type -template -void create_invoker() -{ - invokers[std::tuple(c10::CppTypeToScalarType(), - c10::CppTypeToScalarType())] = - step_invoker; -} -struct InvokerInitializer { - InvokerInitializer() - { - create_invoker(); - create_invoker(); - create_invoker(); - create_invoker(); - create_invoker(); - } -} _invoker_initializer; - -void invoke(std::shared_ptr opt, - torch::Tensor& params, - torch::Tensor& grads, - torch::Tensor& exp_avg, - size_t param_size) -{ - c10::ScalarType params_type = at::typeMetaToScalarType(params.options().dtype()); - c10::ScalarType state_type = at::typeMetaToScalarType(exp_avg.options().dtype()); - - auto it = invokers.find(std::tuple(params_type, state_type)); - if (it == invokers.end()) { - throw std::runtime_error("Lion optimizer with param type "s + c10::toString(params_type) + - " and state type "s + c10::toString(state_type) + - " is not supported on current hardware"s); - } - - it->second(opt, params.data_ptr(), grads.data_ptr(), exp_avg.data_ptr(), param_size); -} - -int ds_lion_step(int optimizer_id, - size_t step, - float lr, - float beta1, - float beta2, - float weight_decay, - torch::Tensor& params, - torch::Tensor& grads, - torch::Tensor& exp_avg) -{ - auto params_c = params.contiguous(); - auto grads_c = grads.contiguous(); - auto exp_avg_c = exp_avg.contiguous(); - - std::shared_ptr opt = - std::static_pointer_cast(s_optimizers[optimizer_id]); - opt->IncrementStep(step, beta1, beta2); - opt->update_state(lr, weight_decay); - - invoke(opt, params_c, grads_c, exp_avg_c, params_c.numel()); - - return 0; -} - -int destroy_lion_optimizer(int optimizer_id) -{ - s_optimizers.erase(optimizer_id); - - return 0; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/fused_lion_frontend.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/fused_lion_frontend.cpp deleted file mode 100644 index e523f97ca3098444a07b21394cf6a773ddcc700d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/fused_lion_frontend.cpp +++ /dev/null @@ -1,22 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include - -void multi_tensor_lion_cuda(int chunk_size, - at::Tensor noop_flag, - std::vector> tensor_lists, - const float lr, - const float beta1, - const float beta2, - const int step, - const float weight_decay); - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("multi_tensor_lion", - &multi_tensor_lion_cuda, - "Compute and apply gradient update to parameters for Lion optimizer"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/multi_tensor_apply.cuh b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/multi_tensor_apply.cuh deleted file mode 100644 index 12f41cb49c6bf505db48f1d21e312578f19da836..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/multi_tensor_apply.cuh +++ /dev/null @@ -1,132 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Copyright NVIDIA/apex -This file is adapted from fused adam in NVIDIA/apex, commit a109f85 -*/ - -#include -#include -#include -#include -#include -#include "compat.h" - -#include - -// #include - -// This header is the one-stop shop for all your multi-tensor apply needs. - -// TODO: Kernel arg size limit may be <4KB for some other cards (ie Jetson) -constexpr int depth_to_max_tensors[5] = {110, 64, 48, 36, 30}; -constexpr int depth_to_max_blocks[5] = {320, 320, 320, 320, 320}; - -template -struct TensorListMetadata { - void* addresses[n][depth_to_max_tensors[n - 1]]; - int sizes[depth_to_max_tensors[n - 1]]; - unsigned char block_to_tensor[depth_to_max_blocks[n - 1]]; - int block_to_chunk[depth_to_max_blocks[n - 1]]; // I fear this needs to be a full int. - int start_tensor_this_launch; -}; - -template -__global__ void multi_tensor_apply_kernel(int chunk_size, - volatile int* noop_flag, - T tl, - U callable, - ArgTypes... args) -{ - // Hand the chunk information to the user-supplied functor to process however it likes. - callable(chunk_size, noop_flag, tl, args...); -} - -template -void multi_tensor_apply(int block_size, - int chunk_size, - const at::Tensor& noop_flag, - const std::vector>& tensor_lists, - T callable, - ArgTypes... args) -{ - TORCH_CHECK(tensor_lists.size() == depth, "tensor_lists.size() != depth"); - int len0 = tensor_lists[0].size(); - TORCH_CHECK(len0 > 0, "tensor_lists[0].size() is not > 0"); - auto ref_device = tensor_lists[0][0].device(); - TORCH_CHECK(ref_device.type() == at::kCUDA, "expected input to be on cuda"); - for (int l = 0; l < tensor_lists.size(); l++) // No range-based for because I need indices - { - TORCH_CHECK(tensor_lists[l].size() == len0, "Size mismatch among tensor lists"); - for (int t = 0; t < tensor_lists[l].size(); t++) { - // TODO: Print which tensor fails. - bool contiguous_memory = tensor_lists[l][t].is_contiguous(); -#ifdef VERSION_GE_1_5 - contiguous_memory = (contiguous_memory || - tensor_lists[l][t].is_contiguous(at::MemoryFormat::ChannelsLast)); -#endif - TORCH_CHECK(contiguous_memory, "A tensor was not contiguous."); - TORCH_CHECK(tensor_lists[l][t].device() == ref_device, - "A tensor was not on the same device as the first tensor"); - TORCH_CHECK(tensor_lists[l][t].numel() == tensor_lists[0][t].numel(), "Size mismatch"); - } - } - - int ntensors = tensor_lists[0].size(); - - TensorListMetadata tl; - - const at::cuda::OptionalCUDAGuard device_guard(device_of(tensor_lists[0][0])); - auto stream = at::cuda::getCurrentCUDAStream(); - - tl.start_tensor_this_launch = 0; - int loc_block_info = 0; - int loc_tensor_info = 0; - for (int t = 0; t < ntensors; t++) { - tl.sizes[loc_tensor_info] = tensor_lists[0][t].numel(); - for (int d = 0; d < depth; d++) - tl.addresses[d][loc_tensor_info] = tensor_lists[d][t].data_ptr(); - loc_tensor_info++; - - int chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size; - - for (int chunk = 0; chunk < chunks_this_tensor; chunk++) { - // std::cout << chunks_this_tensor << std::endl; - tl.block_to_tensor[loc_block_info] = loc_tensor_info - 1; - tl.block_to_chunk[loc_block_info] = chunk; - loc_block_info++; - - bool tensors_full = (loc_tensor_info == depth_to_max_tensors[depth - 1] && - chunk == chunks_this_tensor - 1); - bool blocks_full = (loc_block_info == depth_to_max_blocks[depth - 1]); - bool last_chunk = (t == ntensors - 1 && chunk == chunks_this_tensor - 1); - if (tensors_full || blocks_full || last_chunk) { - // using accscalar_t = acc_type; - multi_tensor_apply_kernel<<>>( - chunk_size, noop_flag.DATA_PTR(), tl, callable, args...); - - AT_CUDA_CHECK(cudaGetLastError()); - - // Reset. The control flow possibilities here make my brain hurt. - loc_block_info = 0; - if (chunk == chunks_this_tensor - 1) { - // std::cout << "Hit case 1 " << cond1 << " " << cond2 << " " << cond3 << - // std::endl; - loc_tensor_info = 0; - tl.start_tensor_this_launch = t + 1; - } else { - // std::cout << "Hit case 2 " << cond1 << " " << cond2 << " " << cond3 << - // std::endl; - tl.sizes[0] = tl.sizes[loc_tensor_info - 1]; - for (int d = 0; d < depth; d++) - tl.addresses[d][0] = tl.addresses[d][loc_tensor_info - 1]; - loc_tensor_info = 1; - tl.start_tensor_this_launch = t; - } - } - } - } -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/multi_tensor_lion.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/multi_tensor_lion.cu deleted file mode 100644 index f5fe6dfdd7d0c3fa9200da8592d0a3f56b138428..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/lion/multi_tensor_lion.cu +++ /dev/null @@ -1,126 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Copyright NVIDIA/apex -This file is adapted from fused adam in NVIDIA/apex, commit a109f85 -*/ - -#include -#include -#include -#include -// Another possibility: -// #include - -#include - -#include "multi_tensor_apply.cuh" -#include "type_shim.h" - -#define BLOCK_SIZE 512 -#define ILP 4 - -using MATH_T = float; - -template -struct LionFunctor { - __device__ __forceinline__ void operator()(int chunk_size, - volatile int* noop_gmem, - TensorListMetadata<3>& tl, - const float beta1, - const float beta2, - const float lr, - const float decay) - { - // I'd like this kernel to propagate infs/nans. - // if(*noop_gmem == 1) - // return; - - int tensor_loc = tl.block_to_tensor[blockIdx.x]; - - // potentially use to pass in list of scalar - // int tensor_num = tl.start_tensor_this_launch + tensor_loc; - - int chunk_idx = tl.block_to_chunk[blockIdx.x]; - int n = tl.sizes[tensor_loc]; - - T* g = (T*)tl.addresses[0][tensor_loc]; - g += chunk_idx * chunk_size; - - T* p = (T*)tl.addresses[1][tensor_loc]; - p += chunk_idx * chunk_size; - - T* m = (T*)tl.addresses[2][tensor_loc]; - m += chunk_idx * chunk_size; - - n -= chunk_idx * chunk_size; - - MATH_T after_decay = 1.0f - lr * decay; - - // see note in multi_tensor_scale_kernel.cu - for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) { - MATH_T r_g[ILP]; - MATH_T r_p[ILP]; - MATH_T r_m[ILP]; -#pragma unroll - for (int ii = 0; ii < ILP; ii++) { - int i = i_start + threadIdx.x + ii * blockDim.x; - if (i < n && i < chunk_size) { - r_g[ii] = g[i]; - r_p[ii] = p[i]; - r_m[ii] = m[i]; - } else { - r_g[ii] = MATH_T(0); - r_p[ii] = MATH_T(0); - r_m[ii] = MATH_T(0); - } - } -#pragma unroll - for (int ii = 0; ii < ILP; ii++) { - MATH_T c = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; - MATH_T update = c > 0 ? (-lr) : lr; - r_p[ii] = r_p[ii] * after_decay + update; - r_m[ii] = beta2 * r_m[ii] + (1 - beta2) * r_g[ii]; - } -#pragma unroll - for (int ii = 0; ii < ILP; ii++) { - int i = i_start + threadIdx.x + ii * blockDim.x; - if (i < n && i < chunk_size) { - p[i] = r_p[ii]; - m[i] = r_m[ii]; - } - } - } - } -}; - -void multi_tensor_lion_cuda(int chunk_size, - at::Tensor noop_flag, - std::vector> tensor_lists, - const float lr, - const float beta1, - const float beta2, - const int step, - const float weight_decay) -{ - using namespace at; - - // Assume single type across p,g,m1,m2 now - DISPATCH_DOUBLE_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), - 0, - "lion", - multi_tensor_apply<3>(BLOCK_SIZE, - chunk_size, - noop_flag, - tensor_lists, - LionFunctor(), - beta1, - beta2, - lr, - weight_decay);) - - AT_CUDA_CHECK(cudaGetLastError()); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/dequantize.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/dequantize.cu deleted file mode 100644 index 783ffb7a41e396d73e9c0254826012a1ddc6655c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/dequantize.cu +++ /dev/null @@ -1,72 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "dequantization_utils.h" -#include "memory_access_utils.h" - -namespace cg = cooperative_groups; - -template -__global__ void dequantize_kernel(T* __restrict__ dequant_data, - const int8_t* __restrict__ q_data, - const float* __restrict__ q_params, - int elems_per_group, - int total_elems) -{ - dequantize::to_global( - dequant_data, q_data, q_params, elems_per_group, total_elems); -} - -#define LAUNCH_DEQUANT_KERNEL(num_bits, q_type) \ - dequantize_kernel<<>>( \ - dequant_data, q_data, q_params, elems_per_group, total_elems); - -template -void launch_dequantize_kernel(T* dequant_data, - const int8_t* q_data, - const float* q_params, - quantize::Type q_type, - int num_bits, - int elems_per_group, - int total_elems, - cudaStream_t stream) -{ - constexpr int unroll = 8; - constexpr int threads = 512; - constexpr int elems_per_block = unroll * threads * dequantize::granularity / (sizeof(T)); - - const dim3 block(threads); - const dim3 grid((total_elems + elems_per_block - 1) / elems_per_block); - - // TODO(cmikeh2): It may make sense to tune unroll, there is perf benefit for large - // problem sizes with this large unroll value. - if (num_bits == 8 && q_type == quantize::Type::Symmetric) { - LAUNCH_DEQUANT_KERNEL(8, quantize::Type::Symmetric); - } else if (num_bits == 8 && q_type == quantize::Type::Asymmetric) { - LAUNCH_DEQUANT_KERNEL(8, quantize::Type::Asymmetric); - } else if (num_bits == 4 && q_type == quantize::Type::Symmetric) { - LAUNCH_DEQUANT_KERNEL(4, quantize::Type::Symmetric); - } else if (num_bits == 4 && q_type == quantize::Type::Asymmetric) { - LAUNCH_DEQUANT_KERNEL(4, quantize::Type::Asymmetric); - } -} - -template void launch_dequantize_kernel(__half* dequant_data, - const int8_t* q_data, - const float* q_params, - quantize::Type q_type, - int num_bits, - int elems_per_group, - int total_elems, - cudaStream_t stream); - -template void launch_dequantize_kernel(float* dequant_data, - const int8_t* q_data, - const float* q_params, - quantize::Type q_type, - int num_bits, - int elems_per_group, - int total_elems, - cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/fake_quantizer.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/fake_quantizer.cu deleted file mode 100644 index 4c08cd4cc3d28d9095813e9190a5034a59c5660e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/fake_quantizer.cu +++ /dev/null @@ -1,1028 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include "custom_cuda_layers.h" -#include "memory_access_utils.h" - -namespace cg = cooperative_groups; - -__global__ void fake_quantize_kernel(__half* vals, int group_size, int num_bits) -{ -#if __CUDA_ARCH__ >= 700 || defined(__HIP_PLATFORM_AMD__) - - cg::thread_block b = cg::this_thread_block(); // tb - cg::thread_block_tile<32> g = - cg::tiled_partition<32>(b); // warp, 32 not optimal for AMD which should be 64. - - int gid = threadIdx.x >> 5; - int lane = threadIdx.x & 0x1f; - int warp_num = blockDim.x >> 5; - int id = threadIdx.x; - - constexpr int granularity = 16; - constexpr int vals_per_access = granularity / sizeof(__half); - - __half data[vals_per_access]; - - int group_id = blockIdx.x; - - int thread_index = id * vals_per_access; - int reg_count = 0; - int offset = group_id * group_size; - float max = -10000.0; - for (int thread_index = id * vals_per_access; thread_index < group_size; - thread_index += blockDim.x * vals_per_access) { - mem_access::load_global(data, vals + offset + thread_index); - -#pragma unroll - for (int i = 0; i < vals_per_access; i++) { - if (abs((float)data[i]) > max) max = abs((float)data[i]); - } - } - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) { - auto temp = g.shfl_xor(max, i); - if (max < temp) max = temp; - } - __shared__ float partialMax[WARP_SIZE]; - - if (lane == 0) partialMax[gid] = max; - - b.sync(); - - if (lane < warp_num) max = partialMax[lane]; - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) { - auto temp = g.shfl_down(max, i); - if (max < temp) max = temp; - } - - max = g.shfl(max, 0); - - float q_scale = (float)(1 << num_bits) / (2 * max + 1e-5); - float q_scale_inv = 1 / q_scale; - int q_range_max = (1 << (num_bits - 1)) - 1; - int q_range_min = -(1 << (num_bits - 1)); - - for (int thread_index = id * vals_per_access; thread_index < group_size; - thread_index += blockDim.x * vals_per_access) { - mem_access::load_global(data, vals + offset + thread_index); -#pragma unroll - for (int j = 0; j < vals_per_access; j++) { - float q_data; - q_data = __half2float(data[j]); - q_data = __float2int_rn(q_data * q_scale); - q_data = q_data > (q_range_max) ? (q_range_max) - : (q_data < (q_range_min) ? (q_range_min) : q_data); - data[j] = __float2half_rn(q_data * q_scale_inv); - } - mem_access::store_global(vals + offset + thread_index, data); - } - -#endif -} - -__global__ void fake_quantize_kernel(float* vals, int group_size, int num_bits) -{ - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); - - int gid = threadIdx.x >> 5; - int lane = threadIdx.x & 0x1f; - int warp_num = blockDim.x >> 5; - int id = threadIdx.x; - - constexpr int granularity = 16; - constexpr int vals_per_access = granularity / sizeof(float); - - float data[vals_per_access]; - - int bid = blockIdx.x; - - int thread_index = id * vals_per_access; - - int reg_count = 0; - - int offset = bid * group_size; - - float max = -10000.0; - - for (int thread_index = id * vals_per_access; thread_index < group_size; - thread_index += blockDim.x * vals_per_access) { - mem_access::load_global(data, vals + offset + thread_index); - -#pragma unroll - for (int i = 0; i < vals_per_access; i++) { - if (abs(data[i]) > max) max = abs(data[i]); - } - } - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) { - auto temp = g.shfl_xor(max, i); - if (max < temp) max = temp; - } - __shared__ float partialMax[WARP_SIZE]; - - if (lane == 0) partialMax[gid] = max; - - b.sync(); - - if (lane < warp_num) max = partialMax[lane]; - - b.sync(); - -#pragma unroll - for (int i = 1; i < warp_num; i <<= 1) { - auto temp = g.shfl_down(max, i); - if (max < temp) max = temp; - } - - max = g.shfl(max, 0); - - float q_scale = (1 << num_bits) / (2 * max + 1e-5); - float q_scale_inv = 1 / q_scale; - - int q_range_max = (1 << (num_bits - 1)) - 1; - int q_range_min = -(1 << (num_bits - 1)); - - for (int thread_index = id * vals_per_access; thread_index < group_size; - thread_index += blockDim.x * vals_per_access) { - mem_access::load_global(data, vals + offset + thread_index); -#pragma unroll - for (int j = 0; j < vals_per_access; j++) { - float q_data; - q_data = __float2int_rn(data[j] * q_scale); - q_data = q_data > (q_range_max) ? (q_range_max) - : (q_data < (q_range_min) ? (q_range_min) : q_data); - data[j] = roundf(q_data * q_scale_inv); - } - mem_access::store_global(vals + offset + thread_index, data); - } -} - -template -void launch_fake_quantize_kernel(T* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream) -{ - dim3 grid_dim(group_num); - dim3 block_dim(1024); - - fake_quantize_kernel<<>>( - vals, total_count / group_num, num_bits); -} - -template void launch_fake_quantize_kernel(float* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream); -template void launch_fake_quantize_kernel(__half* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream); - -__global__ void sr_fake_quantize_kernel(__half* vals, - int token_size, - int token_num, - int num_bits, - std::pair seed) -{ -#if __CUDA_ARCH__ >= 700 || defined(__HIP_PLATFORM_AMD__) - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); - - int gid = threadIdx.x >> 5; - int lane = threadIdx.x & 0x1f; - int warp_num = blockDim.x >> 5; - - int idx = blockIdx.x * blockDim.x + threadIdx.x; - - float2* vals_cast = reinterpret_cast(vals); - - __half2 data_low[128]; - __half2 data_high[128]; - - int bid = blockIdx.x; - - curandStatePhilox4_32_10_t state; - curand_init(seed.first, idx, seed.second, &state); - unsigned int tid = threadIdx.x; - int reg_count = 0; - int offset = bid * token_size; - int group_index = bid * token_size + tid; - - int total_count = token_size * token_num; - if (group_index < total_count) { - // float min = 10000.0; - float max = -10000.0; - while (tid < token_size) { - float2 data = vals_cast[offset + tid]; - __half2* data_h = reinterpret_cast<__half2*>(&data); - data_low[reg_count] = data_h[0]; - data_high[reg_count] = data_h[1]; - - float2 data_f[2]; - data_f[0] = __half22float2(data_h[0]); - data_f[1] = __half22float2(data_h[1]); - - if (abs((float)data_f[0].x) > max) max = abs((float)data_f[0].x); - if (abs((float)data_f[0].y) > max) max = abs((float)data_f[0].y); - if (abs((float)data_f[1].x) > max) max = abs((float)data_f[1].x); - if (abs((float)data_f[1].y) > max) max = abs((float)data_f[1].y); - - tid += blockDim.x; - reg_count++; - } - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) { - auto temp = g.shfl_xor(max, i); - if (max < temp) max = temp; - } - - __shared__ float partialMax[WARP_SIZE]; - - if (lane == 0) partialMax[gid] = max; - - b.sync(); - - if (lane < warp_num) max = partialMax[lane]; - -#pragma unroll - for (int i = 1; i < warp_num; i <<= 1) { - auto temp = g.shfl_down(max, i); - if (max < temp) max = temp; - } - - max = g.shfl(max, 0); - - float q_scale_val = (float)(1 << num_bits) / (max * 2 + 1e-5); - float high_q = (float)((1 << (num_bits - 1)) - 1); - float low_q = (float)(-((1 << (num_bits - 1)))); - - for (int i = 0; i < reg_count; i++) { - int token_index = i * blockDim.x + threadIdx.x; - if (token_index < token_size) { - float2 data_f[2]; - data_f[0] = __half22float2(data_low[i]); - data_f[1] = __half22float2(data_high[i]); - - float2 q_data_int[2]; - q_data_int[0].x = (float)((int)(data_f[0].x * q_scale_val)); - q_data_int[0].y = (float)((int)(data_f[0].y * q_scale_val)); - q_data_int[1].x = (float)((int)(data_f[1].x * q_scale_val)); - q_data_int[1].y = (float)((int)(data_f[1].y * q_scale_val)); - - // Stochastic rounding - float4 rand = curand_uniform4(&state); - - float q_error[4]; - q_error[0] = abs(data_f[0].x - (q_data_int[0].x / q_scale_val)) * q_scale_val; - q_error[1] = abs(data_f[0].y - (q_data_int[0].y / q_scale_val)) * q_scale_val; - q_error[2] = abs(data_f[1].x - (q_data_int[1].x / q_scale_val)) * q_scale_val; - q_error[3] = abs(data_f[1].y - (q_data_int[1].y / q_scale_val)) * q_scale_val; - - q_data_int[0].x = - (rand.x < q_error[0] && q_data_int[0].x > low_q && q_data_int[0].x < high_q) - ? (q_data_int[0].x + (data_f[0].x > 0 ? 1 : -1)) - : q_data_int[0].x; - q_data_int[0].y = - (rand.y < q_error[1] && q_data_int[0].y > low_q && q_data_int[0].y < high_q) - ? (q_data_int[0].y + (data_f[0].y > 0 ? 1 : -1)) - : q_data_int[0].y; - q_data_int[1].x = - (rand.w < q_error[2] && q_data_int[1].x > low_q && q_data_int[1].x < high_q) - ? (q_data_int[1].x + (data_f[1].x > 0 ? 1 : -1)) - : q_data_int[1].x; - q_data_int[1].y = - (rand.z < q_error[3] && q_data_int[1].y > low_q && q_data_int[1].y < high_q) - ? (q_data_int[1].y + (data_f[1].y > 0 ? 1 : -1)) - : q_data_int[1].y; - - data_f[0].x = q_data_int[0].x / q_scale_val; - data_f[0].y = q_data_int[0].y / q_scale_val; - data_f[1].x = q_data_int[1].x / q_scale_val; - data_f[1].y = q_data_int[1].y / q_scale_val; - - float2 result; - __half2* result_h = reinterpret_cast<__half2*>(&result); - result_h[0] = __float22half2_rn(data_f[0]); - result_h[1] = __float22half2_rn(data_f[1]); - - vals_cast[offset + token_index] = result; - } - } - } -#endif -} - -__global__ void sr_fake_quantize_kernel(float* vals, - int token_size, - int token_num, - int num_bits, - std::pair seed) -{ - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); - - int gid = threadIdx.x >> 5; - int lane = threadIdx.x & 0x1f; - int warp_num = blockDim.x >> 5; - int id = threadIdx.x; - - int idx = blockIdx.x * blockDim.x + id; - - float4* vals_cast = reinterpret_cast(vals); - - float4 data[128]; - - int bid = blockIdx.x; - int tid = threadIdx.x; - curandStatePhilox4_32_10_t state; - curand_init(seed.first, idx, seed.second, &state); - - int group_index = bid * token_size + threadIdx.x; - int reg_count = 0; - int total_count = token_size * token_num; - if (group_index < total_count) { - // float min = 10000.0; - float max = -10000.0; - - while (tid < token_size) { - data[reg_count] = vals_cast[group_index]; - - if (abs(data[reg_count].x) > max) max = abs(data[reg_count].x); - if (abs(data[reg_count].y) > max) max = abs(data[reg_count].y); - if (abs(data[reg_count].z) > max) max = abs(data[reg_count].z); - if (abs(data[reg_count].w) > max) max = abs(data[reg_count].w); - - group_index += blockDim.x; - tid += blockDim.x; - reg_count++; - } - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) { - auto temp = g.shfl_xor(max, i); - if (max < temp) max = temp; - } - __shared__ float partialMax[WARP_SIZE]; - - if (lane == 0) partialMax[gid] = max; - - b.sync(); - - if (lane < warp_num) max = partialMax[lane]; - -#pragma unroll - for (int i = 1; i < warp_num; i <<= 1) { - auto temp = g.shfl_down(max, i); - if (max < temp) max = temp; - } - - max = g.shfl(max, 0); - - float q_scale_val = (float)(1 << num_bits) / (max * 2 + 1e-5); - float high_q = (float)((1 << (num_bits - 1)) - 1); - float low_q = (float)(-((1 << (num_bits - 1)))); - - int offset = (bid)*token_size; - for (int i = 0; i < reg_count; i++) { - group_index = i * blockDim.x + threadIdx.x; - if (group_index < token_size) { - float4 q_data = data[i]; - - float4 q_data_int; - q_data_int.x = (float)((int)(q_data.x * q_scale_val)); - q_data_int.y = (float)((int)(q_data.y * q_scale_val)); - q_data_int.w = (float)((int)(q_data.w * q_scale_val)); - q_data_int.z = (float)((int)(q_data.z * q_scale_val)); - - // Stochastic rounding - float4 rand = curand_uniform4(&state); - - float q_error[4]; - q_error[0] = abs(q_data.x - (q_data_int.x / q_scale_val)) * q_scale_val; - q_error[1] = abs(q_data.y - (q_data_int.y / q_scale_val)) * q_scale_val; - q_error[2] = abs(q_data.w - (q_data_int.w / q_scale_val)) * q_scale_val; - q_error[3] = abs(q_data.z - (q_data_int.z / q_scale_val)) * q_scale_val; - - q_data_int.x = - (rand.x < q_error[0] && q_data_int.x > low_q && q_data_int.x < high_q) - ? (q_data_int.x + (q_data.x > 0 ? 1 : -1)) - : q_data_int.x; - q_data_int.y = - (rand.y < q_error[1] && q_data_int.y > low_q && q_data_int.y < high_q) - ? (q_data_int.y + (q_data.y > 0 ? 1 : -1)) - : q_data_int.y; - q_data_int.w = - (rand.w < q_error[2] && q_data_int.w > low_q && q_data_int.w < high_q) - ? (q_data_int.w + (q_data.w > 0 ? 1 : -1)) - : q_data_int.w; - q_data_int.z = - (rand.z < q_error[3] && q_data_int.z > low_q && q_data_int.z < high_q) - ? (q_data_int.z + (q_data.z > 0 ? 1 : -1)) - : q_data_int.z; - - q_data_int.x /= q_scale_val; - q_data_int.y /= q_scale_val; - q_data_int.w /= q_scale_val; - q_data_int.z /= q_scale_val; - - vals_cast[group_index + offset] = q_data_int; - } - } - } -} - -template -void launch_sr_fake_quantize_kernel(T* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream) -{ - dim3 block_dim(1024); - dim3 grid_dim(group_num); - - uint64_t inc = total_count / grid_dim.x / block_dim.x; - std::pair seed = TrainingContext::Instance().IncrementOffset(inc); - - sr_fake_quantize_kernel<<>>( - vals, (total_count / group_num) / 4, group_num, num_bits, seed); -} -template void launch_sr_fake_quantize_kernel(float* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream); -template void launch_sr_fake_quantize_kernel(__half* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream); - -__global__ void fake_quantize_kernel_asym(__half* vals, int group_size, int num_bits) -{ -#if __CUDA_ARCH__ >= 700 || defined(__HIP_PLATFORM_AMD__) - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); - - int gid = threadIdx.x >> 5; - int lane = threadIdx.x & 0x1f; - int warp_num = blockDim.x >> 5; - int id = threadIdx.x; - - float2* vals_cast = reinterpret_cast(vals); - - float2 data[MAX_REG]; - - int group_id = blockIdx.x; - - { - int group_index = id; - int reg_count = 0; - int offset = group_id * group_size; - float max = -10000.0; - float min = 10000.0; - - while (group_index < group_size && reg_count < MAX_REG) { - data[reg_count] = vals_cast[offset + group_index]; - __half* data_h = reinterpret_cast<__half*>(&data[reg_count]); - - if (((float)data_h[0]) > max) max = (float)data_h[0]; - if (((float)data_h[1]) > max) max = (float)data_h[1]; - if (((float)data_h[2]) > max) max = (float)data_h[2]; - if (((float)data_h[3]) > max) max = (float)data_h[3]; - - if (((float)data_h[0]) < min) min = (float)data_h[0]; - if (((float)data_h[1]) < min) min = (float)data_h[1]; - if (((float)data_h[2]) < min) min = (float)data_h[2]; - if (((float)data_h[3]) < min) min = (float)data_h[3]; - - group_index += blockDim.x; - reg_count++; - } - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) { - auto temp = g.shfl_xor(max, i); - if (max < temp) max = temp; - } - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) { - auto temp = g.shfl_xor(min, i); - if (min > temp) min = temp; - } - - __shared__ float partialMax[WARP_SIZE]; - __shared__ float partialMin[WARP_SIZE]; - - if (lane == 0) partialMax[gid] = max; - if (lane == 0) partialMin[gid] = min; - - b.sync(); - - if (lane < warp_num) max = partialMax[lane]; - if (lane < warp_num) min = partialMin[lane]; - -#pragma unroll - for (int i = 1; i < warp_num; i <<= 1) { - auto temp = g.shfl_down(max, i); - if (max < temp) max = temp; - } -#pragma unroll - for (int i = 1; i < warp_num; i <<= 1) { - auto temp = g.shfl_down(min, i); - if (min > temp) min = temp; - } - - max = g.shfl(max, 0); - min = g.shfl(min, 0); - - float q_scale = ((max - min) + 1e-5) / (float)(1 << num_bits); - float q_scale_inv = 1 / q_scale; - - for (int i = 0; i < reg_count; i++) { - group_index = i * blockDim.x + id; - if (group_index < group_size) { - __half2* data_h = reinterpret_cast<__half2*>(&data[i]); - float2 q_data[2]; - q_data[0] = __half22float2(data_h[0]); - q_data[1] = __half22float2(data_h[1]); - - float2 q_data_int[2]; - - q_data_int[0].x = roundf((q_data[0].x - min) * q_scale_inv); - q_data_int[0].y = roundf((q_data[0].y - min) * q_scale_inv); - q_data_int[1].x = roundf((q_data[1].x - min) * q_scale_inv); - q_data_int[1].y = roundf((q_data[1].y - min) * q_scale_inv); - - q_data_int[0].x = q_data_int[0].x * q_scale + min; - q_data_int[0].y = q_data_int[0].y * q_scale + min; - q_data_int[1].x = q_data_int[1].x * q_scale + min; - q_data_int[1].y = q_data_int[1].y * q_scale + min; - - data_h[0] = __float22half2_rn(q_data_int[0]); - data_h[1] = __float22half2_rn(q_data_int[1]); - - vals_cast[offset + group_index] = data[i]; - } - } - } -#endif -} - -__global__ void fake_quantize_kernel_asym(float* vals, int group_size, int num_bits) -{ - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); - - int gid = threadIdx.x >> 5; - int lane = threadIdx.x & 0x1f; - int warp_num = blockDim.x >> 5; - int id = threadIdx.x; - - float4* vals_cast = reinterpret_cast(vals); - - float4 data[MAX_REG]; - - int bid = blockIdx.x; - - int group_index = bid * group_size + id; - int reg_count = 0; - - float max = -10000.0; - float min = 10000.0; - - while (id < group_size && reg_count < MAX_REG) { - float4 data_reg = vals_cast[group_index]; - data[reg_count] = data_reg; - - if (data_reg.x > max) max = data_reg.x; - if (data_reg.y > max) max = data_reg.y; - if (data_reg.w > max) max = data_reg.w; - if (data_reg.z > max) max = data_reg.z; - - if (data_reg.x < min) min = data_reg.x; - if (data_reg.y < min) min = data_reg.y; - if (data_reg.w < min) min = data_reg.w; - if (data_reg.z < min) min = data_reg.z; - - group_index += blockDim.x; - id += blockDim.x; - reg_count++; - } - id = threadIdx.x; - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) { - auto temp = g.shfl_xor(max, i); - if (max < temp) max = temp; - } - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) { - auto temp = g.shfl_xor(min, i); - if (min > temp) min = temp; - } - - __shared__ float partialMax[WARP_SIZE]; - __shared__ float partialMin[WARP_SIZE]; - - if (lane == 0) partialMax[gid] = max; - if (lane == 0) partialMin[gid] = min; - - b.sync(); - - if (lane < warp_num) max = partialMax[lane]; - if (lane < warp_num) min = partialMin[lane]; - -#pragma unroll - for (int i = 1; i < warp_num; i <<= 1) { - auto temp = g.shfl_down(max, i); - if (max < temp) max = temp; - } -#pragma unroll - for (int i = 1; i < warp_num; i <<= 1) { - auto temp = g.shfl_down(min, i); - if (min > temp) min = temp; - } - - max = g.shfl(max, 0); - min = g.shfl(min, 0); - - float q_scale = ((max - min) + 1e-5) / (float)(1 << num_bits); - float q_scale_inv = 1 / q_scale; - for (int i = 0; i < reg_count; i++) { - group_index = i * blockDim.x + id; - if (group_index < group_size) { - float4 q_data; - q_data = data[i]; - - float4 q_data_int; - q_data_int.x = roundf((q_data.x - min) * q_scale_inv); - q_data_int.y = roundf((q_data.y - min) * q_scale_inv); - q_data_int.w = roundf((q_data.w - min) * q_scale_inv); - q_data_int.z = roundf((q_data.z - min) * q_scale_inv); - - q_data.x = q_data_int.x * q_scale + min; - q_data.y = q_data_int.y * q_scale + min; - q_data.w = q_data_int.w * q_scale + min; - q_data.z = q_data_int.z * q_scale + min; - - vals_cast[group_index + bid * group_size] = q_data; - } - } -} - -template -void launch_fake_quantize_kernel_asym(T* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream) -{ - dim3 grid_dim(group_num); - dim3 block_dim(1024); - - fake_quantize_kernel_asym<<>>( - vals, (total_count / group_num) / 4, num_bits); -} - -template void launch_fake_quantize_kernel_asym(float* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream); -template void launch_fake_quantize_kernel_asym(__half* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream); - -__global__ void sr_fake_quantize_kernel_asym(__half* vals, - int token_size, - int token_num, - int num_bits, - std::pair seed) -{ -#if __CUDA_ARCH__ >= 700 || defined(__HIP_PLATFORM_AMD__) - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); - - int gid = threadIdx.x >> 5; - int lane = threadIdx.x & 0x1f; - int warp_num = blockDim.x >> 5; - - int idx = blockIdx.x * blockDim.x + threadIdx.x; - - float2* vals_cast = reinterpret_cast(vals); - - __half2 data_low[128]; - __half2 data_high[128]; - - int bid = blockIdx.x; - - curandStatePhilox4_32_10_t state; - curand_init(seed.first, idx, seed.second, &state); - unsigned int tid = threadIdx.x; - int reg_count = 0; - int offset = bid * token_size; - int group_index = bid * token_size + tid; - - int total_count = token_size * token_num; - if (group_index < total_count) { - float min = 10000.0; - float max = -10000.0; - while (tid < token_size) { - float2 data = vals_cast[offset + tid]; - __half2* data_h = reinterpret_cast<__half2*>(&data); - data_low[reg_count] = data_h[0]; - data_high[reg_count] = data_h[1]; - - float2 data_f[2]; - data_f[0] = __half22float2(data_h[0]); - data_f[1] = __half22float2(data_h[1]); - - if (((float)data_f[0].x) > max) max = (float)data_f[0].x; - if (((float)data_f[0].y) > max) max = (float)data_f[0].y; - if (((float)data_f[1].x) > max) max = (float)data_f[1].x; - if (((float)data_f[1].y) > max) max = (float)data_f[1].y; - - if (((float)data_f[0].x) < min) min = (float)data_f[0].x; - if (((float)data_f[0].y) < min) min = (float)data_f[0].y; - if (((float)data_f[1].x) < min) min = (float)data_f[1].x; - if (((float)data_f[1].y) < min) min = (float)data_f[1].y; - - tid += blockDim.x; - reg_count++; - } - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) { - auto temp = g.shfl_xor(max, i); - if (max < temp) max = temp; - } - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) { - auto temp = g.shfl_xor(min, i); - if (min > temp) min = temp; - } - - __shared__ float partialMax[WARP_SIZE]; - __shared__ float partialMin[WARP_SIZE]; - - if (lane == 0) partialMax[gid] = max; - if (lane == 0) partialMin[gid] = min; - - b.sync(); - - if (lane < warp_num) max = partialMax[lane]; - if (lane < warp_num) min = partialMin[lane]; - -#pragma unroll - for (int i = 1; i < warp_num; i <<= 1) { - auto temp = g.shfl_down(max, i); - if (max < temp) max = temp; - } -#pragma unroll - for (int i = 1; i < warp_num; i <<= 1) { - auto temp = g.shfl_down(min, i); - if (min > temp) min = temp; - } - - max = g.shfl(max, 0); - min = g.shfl(min, 0); - - float q_scale_val = ((max - min) + 1e-5) / (float)(1 << num_bits); - float q_scale_val_inv = 1 / q_scale_val; - float high_q = (float)((1 << num_bits) - 1); - - for (int i = 0; i < reg_count; i++) { - int token_index = i * blockDim.x + threadIdx.x; - if (token_index < token_size) { - float2 data_f[2]; - data_f[0] = __half22float2(data_low[i]); - data_f[1] = __half22float2(data_high[i]); - - float2 q_data_int[2]; - q_data_int[0].x = (float)((unsigned int)((data_f[0].x - min) * q_scale_val_inv)); - q_data_int[0].y = (float)((unsigned int)((data_f[0].y - min) * q_scale_val_inv)); - q_data_int[1].x = (float)((unsigned int)((data_f[1].x - min) * q_scale_val_inv)); - q_data_int[1].y = (float)((unsigned int)((data_f[1].y - min) * q_scale_val_inv)); - - // Stochastic rounding - float4 rand = curand_uniform4(&state); - - float q_error[4]; - q_error[0] = - abs(data_f[0].x - ((q_data_int[0].x * q_scale_val) + min)) * q_scale_val_inv; - q_error[1] = - abs(data_f[0].y - ((q_data_int[0].y * q_scale_val) + min)) * q_scale_val_inv; - q_error[2] = - abs(data_f[1].x - ((q_data_int[1].x * q_scale_val) + min)) * q_scale_val_inv; - q_error[3] = - abs(data_f[1].y - ((q_data_int[1].y * q_scale_val) + min)) * q_scale_val_inv; - - q_data_int[0].x = (rand.x < q_error[0] && q_data_int[0].x < high_q) - ? (q_data_int[0].x + 1) - : q_data_int[0].x; - q_data_int[0].y = (rand.y < q_error[1] && q_data_int[0].y < high_q) - ? (q_data_int[0].y + 1) - : q_data_int[0].y; - q_data_int[1].x = (rand.w < q_error[2] && q_data_int[1].x < high_q) - ? (q_data_int[1].x + 1) - : q_data_int[1].x; - q_data_int[1].y = (rand.z < q_error[3] && q_data_int[1].y < high_q) - ? (q_data_int[1].y + 1) - : q_data_int[1].y; - - data_f[0].x = q_data_int[0].x * q_scale_val + min; - data_f[0].y = q_data_int[0].y * q_scale_val + min; - data_f[1].x = q_data_int[1].x * q_scale_val + min; - data_f[1].y = q_data_int[1].y * q_scale_val + min; - - float2 result; - __half2* result_h = reinterpret_cast<__half2*>(&result); - result_h[0] = __float22half2_rn(data_f[0]); - result_h[1] = __float22half2_rn(data_f[1]); - - vals_cast[offset + token_index] = result; - } - } - } -#endif -} - -__global__ void sr_fake_quantize_kernel_asym(float* vals, - int token_size, - int token_num, - int num_bits, - std::pair seed) -{ - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); - - int gid = threadIdx.x >> 5; - int lane = threadIdx.x & 0x1f; - int warp_num = blockDim.x >> 5; - int id = threadIdx.x; - - int idx = blockIdx.x * blockDim.x + id; - - float4* vals_cast = reinterpret_cast(vals); - - float4 data[128]; - - int bid = blockIdx.x; - int tid = threadIdx.x; - curandStatePhilox4_32_10_t state; - curand_init(seed.first, idx, seed.second, &state); - - int group_index = bid * token_size + threadIdx.x; - int reg_count = 0; - int total_count = token_size * token_num; - if (group_index < total_count) { - float min = 10000.0; - float max = -10000.0; - - while (tid < token_size) { - float4 data_reg = vals_cast[group_index]; - data[reg_count] = data_reg; - if (data_reg.x > max) max = data_reg.x; - if (data_reg.y > max) max = data_reg.y; - if (data_reg.w > max) max = data_reg.w; - if (data_reg.z > max) max = data_reg.z; - - if (data_reg.x < min) min = data_reg.x; - if (data_reg.y < min) min = data_reg.y; - if (data_reg.w < min) min = data_reg.w; - if (data_reg.z < min) min = data_reg.z; - - group_index += blockDim.x; - tid += blockDim.x; - reg_count++; - } - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) { - auto temp = g.shfl_xor(max, i); - if (max < temp) max = temp; - } - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) { - auto temp = g.shfl_xor(min, i); - if (min > temp) min = temp; - } - - __shared__ float partialMax[WARP_SIZE]; - __shared__ float partialMin[WARP_SIZE]; - - if (lane == 0) partialMax[gid] = max; - if (lane == 0) partialMin[gid] = min; - - b.sync(); - - if (lane < warp_num) max = partialMax[lane]; - if (lane < warp_num) min = partialMin[lane]; - -#pragma unroll - for (int i = 1; i < warp_num; i <<= 1) { - auto temp = g.shfl_down(max, i); - if (max < temp) max = temp; - } -#pragma unroll - for (int i = 1; i < warp_num; i <<= 1) { - auto temp = g.shfl_down(min, i); - if (min > temp) min = temp; - } - - max = g.shfl(max, 0); - min = g.shfl(min, 0); - - float q_scale_val = ((max - min) + 1e-5) / (float)(1 << num_bits); - float high_q = (float)((1 << num_bits) - 1); - - int offset = (bid)*token_size; - for (int i = 0; i < reg_count; i++) { - group_index = i * blockDim.x + threadIdx.x; - if (group_index < token_size) { - float4 q_data = data[i]; - - float4 q_data_int; - q_data_int.x = (float)((int)((q_data.x - min) / q_scale_val)); - q_data_int.y = (float)((int)((q_data.y - min) / q_scale_val)); - q_data_int.w = (float)((int)((q_data.w - min) / q_scale_val)); - q_data_int.z = (float)((int)((q_data.z - min) / q_scale_val)); - - // Stochastic rounding - float4 rand = curand_uniform4(&state); - - float q_error[4]; - q_error[0] = abs(q_data.x - ((q_data_int.x * q_scale_val) + min)) / q_scale_val; - q_error[1] = abs(q_data.y - ((q_data_int.y * q_scale_val) + min)) / q_scale_val; - q_error[2] = abs(q_data.w - ((q_data_int.w * q_scale_val) + min)) / q_scale_val; - q_error[3] = abs(q_data.z - ((q_data_int.z * q_scale_val) + min)) / q_scale_val; - - q_data_int.x = (rand.x < q_error[0] && q_data_int.x < high_q) ? (q_data_int.x + 1) - : q_data_int.x; - q_data_int.y = (rand.y < q_error[1] && q_data_int.y < high_q) ? (q_data_int.y + 1) - : q_data_int.y; - q_data_int.w = (rand.w < q_error[2] && q_data_int.w < high_q) ? (q_data_int.w + 1) - : q_data_int.w; - q_data_int.z = (rand.z < q_error[3] && q_data_int.z < high_q) ? (q_data_int.z + 1) - : q_data_int.z; - - q_data_int.x = q_data_int.x * q_scale_val + min; - q_data_int.y = q_data_int.y * q_scale_val + min; - q_data_int.w = q_data_int.w * q_scale_val + min; - q_data_int.z = q_data_int.z * q_scale_val + min; - - vals_cast[group_index + offset] = q_data_int; - } - } - } -} -template -void launch_sr_fake_quantize_kernel_asym(T* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream) -{ - dim3 block_dim(1024); - dim3 grid_dim(group_num); - - uint64_t inc = total_count / grid_dim.x / block_dim.x; - std::pair seed = TrainingContext::Instance().IncrementOffset(inc); - - sr_fake_quantize_kernel<<>>( - vals, (total_count / group_num) / 4, group_num, num_bits, seed); -} -template void launch_sr_fake_quantize_kernel_asym(float* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream); -template void launch_sr_fake_quantize_kernel_asym(__half* vals, - int total_count, - int group_num, - int num_bits, - cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/pt_binding.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/pt_binding.cpp deleted file mode 100644 index b48eaacd0881c68ce06052e894f9f32797233cc8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/pt_binding.cpp +++ /dev/null @@ -1,404 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include -#include -#include -#include "quantization.h" - -template -at::Tensor ds_quantize(at::Tensor& vals, int groups, int bits) -{ - auto t_size = vals.sizes(); - int size = 1; - for (auto dim : t_size) size *= dim; - - if ((((size / groups) - 1) / 4096 + 1) <= 256) { - launch_fake_quantize_kernel( - (T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream()); - } - return vals; -} - -template -at::Tensor ds_sr_quantize(at::Tensor& vals, int groups, int bits) -{ - auto t_size = vals.sizes(); - int size = 1; - for (auto dim : t_size) size *= dim; - - if (((size / groups) / 4 / 1024) <= 256) { - launch_sr_fake_quantize_kernel( - (T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream()); - } - return vals; -} - -template -at::Tensor ds_quantize_asym(at::Tensor& vals, int groups, int bits) -{ - auto t_size = vals.sizes(); - int size = 1; - for (auto dim : t_size) size *= dim; - - if ((((size / groups) - 1) / 4096 + 1) <= 256) { - launch_fake_quantize_kernel_asym( - (T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream()); - } - return vals; -} - -template -at::Tensor ds_sr_quantize_asym(at::Tensor& vals, int groups, int bits) -{ - auto t_size = vals.sizes(); - int size = 1; - for (auto dim : t_size) size *= dim; - - if (((size / groups) / 4 / 1024) <= 256) { - launch_sr_fake_quantize_kernel_asym( - (T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream()); - } - return vals; -} - -std::vector quantize_kernel(at::Tensor& input_vals, - int groups, - int numBits, - quantize::Type quantType) -{ - auto dtype = at::kFloat; - auto params_options = at::TensorOptions() - .dtype(dtype) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - const int param_elems = (quantize::requires_offset(quantType)) ? 2 : 1; - auto params = torch::empty({groups, param_elems}, params_options); - - auto output_options = at::TensorOptions() - .dtype(at::kChar) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - auto output_sizes = input_vals.sizes().vec(); - output_sizes[output_sizes.size() - 1] /= numBits == 8 ? 1 : 2; - auto output = torch::empty(output_sizes, output_options); - - const int elems_per_group = at::numel(input_vals) / groups; - - launch_quant((int8_t*)output.data_ptr(), - (float*)params.data_ptr(), - (__half*)input_vals.data_ptr(), - groups, - elems_per_group, - numBits, - quantType, - at::cuda::getCurrentCUDAStream()); - - return {output, params}; -} - -template -at::Tensor dequantize(at::Tensor& quantized_data, - at::Tensor& params, - int groups, - int num_bits, - quantize::Type quant_type) -{ - auto dtype = (std::is_same::value) ? torch::kFloat32 : torch::kFloat16; - auto output_options = at::TensorOptions() - .dtype(dtype) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - auto output_sizes = quantized_data.sizes().vec(); - output_sizes[output_sizes.size() - 1] *= num_bits == 8 ? 1 : 2; - auto output = torch::empty(output_sizes, output_options); - - const int total_elems = at::numel(output); - const int elems_per_group = total_elems / groups; - - launch_dequantize_kernel((T*)output.data_ptr(), - (const int8_t*)quantized_data.data_ptr(), - (const float*)params.data_ptr(), - quant_type, - num_bits, - elems_per_group, - total_elems, - at::cuda::getCurrentCUDAStream()); - - return output; -} - -at::Tensor dequantize_int4_to_half_experimental(at::Tensor& data_in, - at::Tensor& scale_buffer, - at::Tensor& min_val_buffer, - int num_group, - int group_size) -{ - auto output_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA); - auto output = torch::empty({num_group, group_size}, output_options); - - launch_dequantize_int4_to_half_experimental((uint8_t*)data_in.data_ptr(), - (half*)output.data_ptr(), - (half*)scale_buffer.data_ptr(), - (half*)min_val_buffer.data_ptr(), - num_group, - group_size, - at::cuda::getCurrentCUDAStream()); - - return output; -} - -at::Tensor dequantize_int8_to_half_experimental(at::Tensor& data_in, - at::Tensor& scale_buffer, - at::Tensor& min_val_buffer, - int num_group, - int group_size) -{ - auto output_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA); - auto output = torch::empty({num_group, group_size}, output_options); - - launch_dequantize_int8_to_half_experimental((uint8_t*)data_in.data_ptr(), - (half*)output.data_ptr(), - (half*)scale_buffer.data_ptr(), - (half*)min_val_buffer.data_ptr(), - num_group, - group_size, - at::cuda::getCurrentCUDAStream()); - - return output; -} - -std::vector ds_loco_swizzle_quant(at::Tensor& input_vals, - at::Tensor& error_feedback, - float err_beta, - int groups, - int num_bits, - quantize::Type quant_type, - int pipeline_size, - int nodes, - int devices_per_node) -{ - auto scales_options = at::TensorOptions() - .dtype(at::kFloat) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1; - auto scales = torch::empty({groups, scales_elems}, scales_options); - - auto output_options = at::TensorOptions() - .dtype(at::kChar) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - const int quantization_scalar = 8 / num_bits; - const int compressed_vals = at::numel(input_vals) / quantization_scalar; - - auto output = torch::empty({compressed_vals}, output_options); - const int elems_per_group = at::numel(input_vals) / groups; - - launch_loco_swizzled_quant(reinterpret_cast(output.data_ptr()), - reinterpret_cast(scales.data_ptr()), - reinterpret_cast(input_vals.data_ptr()), - reinterpret_cast<__half*>(error_feedback.data_ptr()), - err_beta, - num_bits, - quant_type, - groups, - elems_per_group, - pipeline_size, - nodes, - devices_per_node, - at::cuda::getCurrentCUDAStream()); - - return {output, scales}; -} - -std::vector ds_swizzle_quant(at::Tensor& input_vals, - int groups, - int num_bits, - quantize::Type quant_type, - int pipeline_size, - int nodes, - int devices_per_node) -{ - auto scales_options = at::TensorOptions() - .dtype(at::kFloat) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1; - auto scales = torch::empty({groups, scales_elems}, scales_options); - - auto output_options = at::TensorOptions() - .dtype(at::kChar) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - const int quantization_scalar = 8 / num_bits; - const int compressed_vals = at::numel(input_vals) / quantization_scalar; - - auto output = torch::empty({compressed_vals}, output_options); - const int elems_per_group = at::numel(input_vals) / groups; - - launch_swizzled_quant((int8_t*)output.data_ptr(), - (float*)scales.data_ptr(), - (__half*)input_vals.data_ptr(), - num_bits, - quant_type, - groups, - elems_per_group, - pipeline_size, - nodes, - devices_per_node, - at::cuda::getCurrentCUDAStream()); - - return {output, scales}; -} - -std::vector quantized_reduction(at::Tensor& input_vals, - at::Tensor& input_scales, - int in_groups, - int out_groups, - int num_bits, - quantize::Type quant_type, - int devices_per_node) -{ - auto scales_options = at::TensorOptions() - .dtype(at::kFloat) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1; - auto scales = torch::empty({out_groups, scales_elems}, scales_options); - - auto output_options = at::TensorOptions() - .dtype(at::kChar) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - std::vector sz(input_vals.sizes().begin(), input_vals.sizes().end()); - sz[sz.size() - 1] = sz.back() / devices_per_node; // num of GPU per nodes - const int elems_per_in_tensor = at::numel(input_vals) / devices_per_node; - auto output = torch::empty(sz, output_options); - - const int elems_per_in_group = elems_per_in_tensor / (in_groups / devices_per_node); - const int elems_per_out_group = elems_per_in_tensor / out_groups; - - launch_dequant_reduce((int8_t*)output.data_ptr(), - (float*)scales.data_ptr(), - (const int8_t*)input_vals.data_ptr(), - (const float*)input_scales.data_ptr(), - devices_per_node, - num_bits, - quant_type, - out_groups, - elems_per_out_group, - elems_per_in_tensor, - in_groups / devices_per_node, - elems_per_in_group, - at::cuda::getCurrentCUDAStream()); - return {output, scales}; -} - -std::vector loco_quantized_reduction(at::Tensor& input_vals, - at::Tensor& input_scales, - at::Tensor& error_feedback, - float err_beta, - int in_groups, - int out_groups, - int num_bits, - quantize::Type quant_type, - int devices_per_node) -{ - auto scales_options = at::TensorOptions() - .dtype(at::kFloat) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1; - - auto scales = torch::empty({out_groups, scales_elems}, scales_options); - - auto output_options = at::TensorOptions() - .dtype(at::kChar) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - std::vector sz(input_vals.sizes().begin(), input_vals.sizes().end()); - sz[sz.size() - 1] = sz.back() / devices_per_node; - - const int elems_per_in_tensor = at::numel(input_vals) / devices_per_node; - - auto output = torch::empty(sz, output_options); - - const int elems_per_in_group = elems_per_in_tensor / (in_groups / devices_per_node); - const int elems_per_out_group = elems_per_in_tensor / out_groups; - - launch_loco_dequant_reduce((int8_t*)output.data_ptr(), - (float*)scales.data_ptr(), - (const int8_t*)input_vals.data_ptr(), - (const float*)input_scales.data_ptr(), - devices_per_node, - num_bits, - quant_type, - out_groups, - elems_per_out_group, - elems_per_in_tensor, - in_groups / devices_per_node, - elems_per_in_group, - (__half2*)error_feedback.data_ptr(), - err_beta, - at::cuda::getCurrentCUDAStream()); - - return {output, scales}; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("ds_quantize_fp32", &ds_quantize, "DeepSpeed Quantize with fp32 (CUDA)"); - m.def("ds_quantize_fp16", &ds_quantize<__half>, "DeepSpeed Quantize with fp16 (CUDA)"); - m.def("ds_sr_quantize_fp32", &ds_sr_quantize, "DeepSpeed Quantize with fp32 (CUDA)"); - m.def("ds_sr_quantize_fp16", &ds_sr_quantize<__half>, "DeepSpeed Quantize with fp16 (CUDA)"); - m.def("ds_quantize_asym_fp32", &ds_quantize_asym, "DeepSpeed Quantize with fp32 (CUDA)"); - m.def( - "ds_quantize_asym_fp16", &ds_quantize_asym<__half>, "DeepSpeed Quantize with fp16 (CUDA)"); - m.def("ds_sr_quantize_asym_fp32", - &ds_sr_quantize_asym, - "DeepSpeed Quantize with fp32 (CUDA)"); - m.def("ds_sr_quantize_asym_fp16", - &ds_sr_quantize_asym<__half>, - "DeepSpeed Quantize with fp16 (CUDA)"); - pybind11::enum_(m, "QuantizationType") - .value("Symmetric", quantize::Type::Symmetric) - .value("Asymmetric", quantize::Type::Asymmetric) - .export_values(); - m.def("quantize", &quantize_kernel); - m.def("dequantize", &dequantize<__half>); - m.def("dequantize_fp32", &dequantize); - m.def("dequantize_int4_to_half_experimental", - &dequantize_int4_to_half_experimental, - "Dequantize int4 to half (experimental)"); - m.def("dequantize_int8_to_half_experimental", - &dequantize_int8_to_half_experimental, - "Dequantize int8 to half (experimental)"); - m.def("swizzle_quant", &ds_swizzle_quant); - m.def("quantized_reduction", &quantized_reduction); - m.def("loco_swizzle_quant", &ds_loco_swizzle_quant, "LoCo Swizzled Quantization Kernel"); - m.def("loco_quantized_reduction", - &loco_quantized_reduction, - "LoCo Quantization and Reduction Kernel"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quant_reduce.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quant_reduce.cu deleted file mode 100644 index 4100c5174b809e0db4dcf138dfcfde4f42e75843..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quant_reduce.cu +++ /dev/null @@ -1,557 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include "dequantization_utils.h" -#include "ds_kernel_utils.h" -#include "memory_access_utils.h" -#include "quantization_utils.h" -#include "reduction_utils.h" - -using rop = reduce::ROpType; - -/* -TODO(cmikeh2): Add implementation that better handles larger nodes. It would like make sense -to leverage some parallel reductions here to improve performance. -*/ - -template -__global__ void __launch_bounds__(1024) dequant_reduce(int8_t* reduced_data, - float* reduced_scales, - const int8_t* input_data, - const float* input_scales, - int elems_per_out_group, - int elems_per_in_tensor, - int groups_per_in_tensor, - int elems_per_in_group, - int num_tensors) -{ - cg::thread_block tb = cg::this_thread_block(); - cg::thread_block_tile warp = cg::tiled_partition(tb); - - // NOTE(cmikeh2): This probably could be hardcoded to a larger number, - // but that means even stronger restrictions on the number of elements per group - // A performance analysis here might be beneficial - constexpr int mem_granularity = (numBits == 8) ? 8 : 4; - constexpr int elems_per_load = mem_granularity / sizeof(int8_t); // div by 1 - constexpr int storage_values = 16 / sizeof(__half2); - - const int block_offset = tb.group_index().x * elems_per_out_group; - const int elem_offset = tb.thread_index().x * elems_per_load; - const int base_offset = block_offset + elem_offset; - const int stride = tb.group_dim().x * elems_per_load; - - __half2 local_buffer[totalChunks * storage_values]; - - quantize::GroupStats stats; - -#pragma unroll - for (int i = 0; i < totalChunks; i++) { - __half2* iteration_buffer = local_buffer + i * storage_values; - -#pragma unroll - for (int j = 0; j < storage_values; j++) { - iteration_buffer[j] = reduce::init(); - } - - const int iter_offset = i * stride + base_offset; - const int iter_scale_idx = iter_offset / elems_per_in_group; - bool do_loads = i * stride + elem_offset < elems_per_out_group; - - if (numTensors > 0) { -#pragma unroll - for (int j = 0; j < numTensors; j++) { - if (do_loads) { - int8_t load_buffer[elems_per_load]; - - mem_access::load_global( - load_buffer, input_data + j * elems_per_in_tensor + iter_offset); - - quantize::Params params( - input_scales + j * groups_per_in_tensor, iter_scale_idx); - - __half2 dequant_buffer[storage_values]; - dequantize::chunk(dequant_buffer, load_buffer, params); - -#pragma unroll - for (int k = 0; k < storage_values; k++) { - iteration_buffer[k] = - reduce::element(iteration_buffer[k], dequant_buffer[k]); - } - } - } - } else { -#pragma unroll 4 - for (int j = 0; j < num_tensors; j++) { - if (do_loads) { - int8_t load_buffer[elems_per_load]; - - mem_access::load_global( - load_buffer, input_data + j * elems_per_in_tensor + iter_offset); - - quantize::Params params( - input_scales + j * groups_per_in_tensor, iter_scale_idx); - - __half2 dequant_buffer[storage_values]; - dequantize::chunk(dequant_buffer, load_buffer, params); - -#pragma unroll - for (int k = 0; k < storage_values; k++) { - iteration_buffer[k] = - reduce::element(iteration_buffer[k], dequant_buffer[k]); - } - } - } - } - -#pragma unroll - for (int j = 0; j < storage_values; j++) { stats.update(iteration_buffer[j]); } - } - - auto params = stats.template get_params(tb, warp); - - if (tb.thread_index().x == 0) { params.store(reduced_scales, tb.group_index().x); } - -#pragma unroll - for (int i = 0; i < totalChunks; i++) { - const int iter_offset = i * stride + base_offset; - if (i * stride + elem_offset < elems_per_out_group) { - int8_t local_output[elems_per_load]; - quantize::_chunk( - local_output, local_buffer + i * storage_values, params); - mem_access::store_global(reduced_data + iter_offset, local_output); - } - } -} - -template -int32_t pow2_round(int32_t raw_value) -{ - return (((raw_value - 1) >> Power) + 1) << Power; -} - -#define LAUNCH_DEQUANT_REDUCE(num_chunks) \ - dequant_reduce \ - <<>>(reduced_data, \ - reduced_scales, \ - input_data, \ - input_scales, \ - elems_per_out_group, \ - elems_per_in_tensor, \ - groups_per_in_tensor, \ - elems_per_in_group, \ - num_tensors); - -template -void launch_dequant_reduce_impl(int8_t* reduced_data, - float* reduced_scales, - const int8_t* input_data, - const float* input_scales, - int out_groups, - int elems_per_out_group, - int elems_per_in_tensor, - int groups_per_in_tensor, - int elems_per_in_group, - int num_tensors, - cudaStream_t stream) -{ - // This is a coincidence. This is derived by 8 halves per 16 bytes with 2-way packing for int4 - constexpr int elems_per_thread = numBits; - const int one_step_threads = - next_pow2((elems_per_out_group + elems_per_thread - 1) / (elems_per_thread)); - // TODO(cmikeh2): Tune this - const int threads = (one_step_threads < 1024) ? one_step_threads : 1024; - - dim3 block(threads); - dim3 grid(out_groups); - - const int elems_per_step = threads * elems_per_thread; - const int unroll_raw = (elems_per_out_group + elems_per_step - 1) / elems_per_step; - - const int unroll = (unroll_raw >= 4) ? pow2_round<1>(unroll_raw) : unroll_raw; - - if (unroll == 1) { - // 0-4096 elems - LAUNCH_DEQUANT_REDUCE(1); - } else if (unroll == 2) { - // 4097-8192 etc... - LAUNCH_DEQUANT_REDUCE(2); - } else if (unroll == 3) { - LAUNCH_DEQUANT_REDUCE(3); - } else if (unroll == 4) { - LAUNCH_DEQUANT_REDUCE(4); - } else if (unroll == 6) { - LAUNCH_DEQUANT_REDUCE(6); - } else if (unroll == 8) { - LAUNCH_DEQUANT_REDUCE(8); - } else if (unroll == 10) { - LAUNCH_DEQUANT_REDUCE(10); - } else if (unroll == 12) { - // 48k limit - LAUNCH_DEQUANT_REDUCE(12); - } else { - assert(false); - } -} - -#define LAUNCH_DEQUANT_REDUCE_IMPL(NUM_BITS, NUM_GPUS, QUANT_TYPE) \ - launch_dequant_reduce_impl(reduced_data, \ - reduced_scales, \ - input_data, \ - input_scales, \ - out_groups, \ - elems_per_out_group, \ - elems_per_in_tensor, \ - groups_per_in_tensor, \ - elems_per_in_group, \ - num_gpus, \ - stream); - -void launch_dequant_reduce(int8_t* reduced_data, - float* reduced_scales, - const int8_t* input_data, - const float* input_scales, - int num_gpus, - int num_bits, - quantize::Type quant_type, - int out_groups, - int elems_per_out_group, - int elems_per_in_tensor, - int groups_per_in_tensor, - int elems_per_in_group, - cudaStream_t stream) -{ - if (quant_type == quantize::Type::Symmetric) { - if (num_bits == 4) { - if (num_gpus == 8) { - LAUNCH_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Symmetric); - } else if (num_gpus == 16) { - LAUNCH_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Symmetric); - } else { - LAUNCH_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Symmetric); - } - } else if (num_bits == 8) { - if (num_gpus == 8) { - LAUNCH_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Symmetric); - } else if (num_gpus == 16) { - LAUNCH_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Symmetric); - } else { - LAUNCH_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Symmetric); - } - } - } else if (quant_type == quantize::Type::Asymmetric) { - if (num_bits == 4) { - if (num_gpus == 8) { - LAUNCH_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Asymmetric); - } else if (num_gpus == 16) { - LAUNCH_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Asymmetric); - } else { - LAUNCH_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Asymmetric); - } - } else if (num_bits == 8) { - if (num_gpus == 8) { - LAUNCH_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Asymmetric); - } else if (num_gpus == 16) { - LAUNCH_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Asymmetric); - } else { - LAUNCH_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Asymmetric); - } - } - } -} - -/* -Modified loco_dequant_reduce function that performs dequantization and reduction, -and incorporates error-feedback by updating the error_feedback tensor in-place. -*/ - -template -__global__ void __launch_bounds__(1024) loco_dequant_reduce(int8_t* reduced_data, - float* reduced_scales, - const int8_t* input_data, - const float* input_scales, - int elems_per_out_group, - int elems_per_in_tensor, - int groups_per_in_tensor, - int elems_per_in_group, - int num_tensors, - __half2* error_feedback, - const float err_beta) -{ - cg::thread_block tb = cg::this_thread_block(); - cg::thread_block_tile warp = cg::tiled_partition(tb); - - constexpr int mem_granularity = (numBits == 8) ? 8 : 4; - constexpr int elems_per_load = mem_granularity / sizeof(int8_t); - constexpr int storage_values = 16 / sizeof(__half2); - - const int block_offset = tb.group_index().x * elems_per_out_group; - const int elem_offset = tb.thread_index().x * elems_per_load; - const int base_offset = block_offset + elem_offset; - const int stride = tb.group_dim().x * elems_per_load; - - constexpr int scaling_factor = elems_per_load / storage_values; - const int block_offset_err = block_offset / scaling_factor; - const int elem_offset_err = tb.thread_index().x * storage_values; - const int base_offset_err = block_offset_err + elem_offset_err; - const int stride_err = tb.group_dim().x * storage_values; - - __half2 local_buffer[totalChunks * storage_values]; - __half2 err_buffer[totalChunks * storage_values]; - - quantize::GroupStats stats; - -#pragma unroll - for (int i = 0; i < totalChunks; i++) { - __half2* iteration_buffer = local_buffer + i * storage_values; - __half2* iter_err_buffer = err_buffer + i * storage_values; - -#pragma unroll - for (int j = 0; j < storage_values; j++) { - iteration_buffer[j] = reduce::init(); - } - - const int iter_offset = i * stride + base_offset; - const int iter_offset_err = i * stride_err + base_offset_err; - const int iter_scale_idx = iter_offset / elems_per_in_group; - bool do_loads = i * stride + elem_offset < elems_per_out_group; - - if (numTensors > 0) { -#pragma unroll - for (int j = 0; j < numTensors; j++) { - if (do_loads) { - int8_t load_buffer[elems_per_load]; - - mem_access::load_global( - load_buffer, input_data + j * elems_per_in_tensor + iter_offset); - - quantize::Params params( - input_scales + j * groups_per_in_tensor, iter_scale_idx); - - __half2 dequant_buffer[storage_values]; - dequantize::chunk(dequant_buffer, load_buffer, params); - -#pragma unroll - for (int k = 0; k < storage_values; k++) { - iteration_buffer[k] = - reduce::element(iteration_buffer[k], dequant_buffer[k]); - } - } - } - } else { -#pragma unroll 4 - for (int j = 0; j < num_tensors; j++) { - if (do_loads) { - int8_t load_buffer[elems_per_load]; - - mem_access::load_global( - load_buffer, input_data + j * elems_per_in_tensor + iter_offset); - - quantize::Params params( - input_scales + j * groups_per_in_tensor, iter_scale_idx); - - __half2 dequant_buffer[storage_values]; - dequantize::chunk(dequant_buffer, load_buffer, params); - -#pragma unroll - for (int k = 0; k < storage_values; k++) { - iteration_buffer[k] = - reduce::element(iteration_buffer[k], dequant_buffer[k]); - } - } - } - } - mem_access::load_global( - iter_err_buffer, error_feedback + iter_offset_err, do_loads); -#pragma unroll - for (int k = 0; k < storage_values; k++) { - iteration_buffer[k] = __hadd2(iteration_buffer[k], iter_err_buffer[k]); - stats.update(iteration_buffer[k]); - } - } - - auto params = stats.template get_params(tb, warp); - - // Initialize dequantization parameters based on params - auto de_params = params; - de_params.scale = 1.0f / params.scale; - if constexpr (quantType == quantize::Type::Asymmetric) { de_params.offset = params.offset; } - - if (tb.thread_index().x == 0) { params.store(reduced_scales, tb.group_index().x); } - -#pragma unroll - for (int i = 0; i < totalChunks; i++) { - const int iter_offset = i * stride + base_offset; - const int iter_offset_err = i * stride_err + base_offset_err; - __half2* iteration_buffer = local_buffer + i * storage_values; - __half2* iter_err_buffer = err_buffer + i * storage_values; - - if (i * stride + elem_offset < elems_per_out_group) { - // ----------- Begin Error-Feedback Modification ----------- - int8_t local_output[elems_per_load]; - quantize::_chunk(local_output, iteration_buffer, params); - mem_access::store_global(reduced_data + iter_offset, local_output); - - // Dequantize the quantized output to compute the dequantized value - __half2 dequant_buffer[storage_values]; - dequantize::chunk(dequant_buffer, local_output, de_params); - -#pragma unroll - for (int k = 0; k < storage_values; k++) { - // __half2 to float2 - float2 iter_buf_f = __half22float2(iteration_buffer[k]); - float2 dequant_buf_f = __half22float2(dequant_buffer[k]); - - // Update within float precision - float2 new_error_f; - new_error_f.x = iter_buf_f.x - dequant_buf_f.x; - new_error_f.y = iter_buf_f.y - dequant_buf_f.y; - - float2 iter_err_buf_f = __half22float2(iter_err_buffer[k]); - - iter_err_buf_f.x = err_beta * iter_err_buf_f.x + (1.0f - err_beta) * new_error_f.x; - iter_err_buf_f.y = err_beta * iter_err_buf_f.y + (1.0f - err_beta) * new_error_f.y; - - // float2 back to __half2 - iter_err_buffer[k] = __float22half2_rn(iter_err_buf_f); - } - mem_access::store_global(error_feedback + iter_offset_err, - iter_err_buffer); - } - } -} - -#define LAUNCH_LOCO_DEQUANT_REDUCE(num_chunks) \ - loco_dequant_reduce \ - <<>>(reduced_data, \ - reduced_scales, \ - input_data, \ - input_scales, \ - elems_per_out_group, \ - elems_per_in_tensor, \ - groups_per_in_tensor, \ - elems_per_in_group, \ - num_tensors, \ - error_feedback, \ - err_beta); - -template -void launch_loco_dequant_reduce_impl(int8_t* reduced_data, - float* reduced_scales, - const int8_t* input_data, - const float* input_scales, - int out_groups, - int elems_per_out_group, - int elems_per_in_tensor, - int groups_per_in_tensor, - int elems_per_in_group, - int num_tensors, - __half2* error_feedback, - const float err_beta, - cudaStream_t stream) -{ - constexpr int elems_per_thread = numBits; - const int one_step_threads = - next_pow2((elems_per_out_group + elems_per_thread - 1) / (elems_per_thread)); - const int threads = (one_step_threads < 1024) ? one_step_threads : 1024; - - dim3 block(threads); - dim3 grid(out_groups); - - const int elems_per_step = threads * elems_per_thread; - const int unroll_raw = (elems_per_out_group + elems_per_step - 1) / elems_per_step; - - const int unroll = (unroll_raw >= 4) ? pow2_round<1>(unroll_raw) : unroll_raw; - - if (unroll == 1) { - LAUNCH_LOCO_DEQUANT_REDUCE(1); - } else if (unroll == 2) { - LAUNCH_LOCO_DEQUANT_REDUCE(2); - } else if (unroll == 3) { - LAUNCH_LOCO_DEQUANT_REDUCE(3); - } else if (unroll == 4) { - LAUNCH_LOCO_DEQUANT_REDUCE(4); - } else if (unroll == 6) { - LAUNCH_LOCO_DEQUANT_REDUCE(6); - } else if (unroll == 8) { - LAUNCH_LOCO_DEQUANT_REDUCE(8); - } else if (unroll == 10) { - LAUNCH_LOCO_DEQUANT_REDUCE(10); - } else if (unroll == 12) { - LAUNCH_LOCO_DEQUANT_REDUCE(12); - } else { - assert(false); - } -} - -#define LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(NUM_BITS, NUM_GPUS, QUANT_TYPE) \ - launch_loco_dequant_reduce_impl(reduced_data, \ - reduced_scales, \ - input_data, \ - input_scales, \ - out_groups, \ - elems_per_out_group, \ - elems_per_in_tensor, \ - groups_per_in_tensor, \ - elems_per_in_group, \ - num_gpus, \ - error_feedback, \ - err_beta, \ - stream); - -void launch_loco_dequant_reduce(int8_t* reduced_data, - float* reduced_scales, - const int8_t* input_data, - const float* input_scales, - int num_gpus, - int num_bits, - quantize::Type quant_type, - int out_groups, - int elems_per_out_group, - int elems_per_in_tensor, - int groups_per_in_tensor, - int elems_per_in_group, - __half2* error_feedback, - const float err_beta, - cudaStream_t stream) -{ - if (quant_type == quantize::Type::Symmetric) { - if (num_bits == 4) { - if (num_gpus == 8) { - LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Symmetric); - } else if (num_gpus == 16) { - LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Symmetric); - } else { - LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Symmetric); - } - } else if (num_bits == 8) { - if (num_gpus == 8) { - LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Symmetric); - } else if (num_gpus == 16) { - LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Symmetric); - } else { - LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Symmetric); - } - } - } else if (quant_type == quantize::Type::Asymmetric) { - if (num_bits == 4) { - if (num_gpus == 8) { - LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Asymmetric); - } else if (num_gpus == 16) { - LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Asymmetric); - } else { - LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Asymmetric); - } - } else if (num_bits == 8) { - if (num_gpus == 8) { - LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Asymmetric); - } else if (num_gpus == 16) { - LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Asymmetric); - } else { - LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Asymmetric); - } - } - } -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quantize.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quantize.cu deleted file mode 100644 index b04d0e968ba58a203f2b5e24790053cbd43bc74a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quantize.cu +++ /dev/null @@ -1,151 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "ds_kernel_utils.h" -#include "memory_access_utils.h" -#include "quantization.h" -#include "quantization_utils.h" -#include "reduction_utils.h" - -namespace cg = cooperative_groups; - -/* -Pure quantization kernel with no fusion. -*/ -template -__global__ void cached_quantization(int8_t* __restrict__ output_data, - float* __restrict__ params, - const __half* __restrict__ input_data, - int groups, - int elems_per_group) -{ - cg::thread_block tb = cg::this_thread_block(); - cg::thread_block_tile warp = cg::tiled_partition(tb); - - // Indexing offsets - const int block_offset = - (tb.group_index().x * (max_threads / threads_per_group) * elems_per_group) + - (tb.thread_index().y * elems_per_group); - const int elem_offset = tb.thread_index().x * quantize::h_per_load; - const int base_offset = block_offset + elem_offset; - const int stride = tb.size() * quantize::h_per_load; - - const __half* input_base = input_data + base_offset; //.. - - __half2 local_buffer[UNROLL * internal_unroll * quantize::h2_per_load]; - -#pragma unroll - for (int i = 0; i < UNROLL; i++) { - // Convenience helper, should resolve to register indices and not realize. - __half2* iteration_buffer = local_buffer + i * internal_unroll * quantize::h2_per_load; -#pragma unroll - for (int j = 0; j < internal_unroll; j++) { - const int iteration = i * internal_unroll + j; - mem_access::load_global( - iteration_buffer + j * quantize::h2_per_load, - input_base + iteration * stride, - elem_offset + iteration * stride < elems_per_group); - } - } - - quantize:: - local_array( - local_buffer, params, output_data, elems_per_group, groups); -} - -/********* Launcher methods ***********/ -#define LAUNCH_CACHED_QUANT_CALL(q_bits, quant_type) \ - cached_quantization \ - <<>>(output_data, params, input_data, groups, elems_per_group); - -#define LAUNCH_CACHED_QUANT( \ - q_bits, quant_type, unroll_factor_in, internal_unroll_in, threads_per_group_in) \ - const int unroll_factor = unroll_factor_in; \ - const int internal_unroll_l = internal_unroll_in; \ - const int threads_per_group = threads_per_group_in; \ - if (q_bits == 4) { \ - if (quant_type == quantize::Type::Asymmetric) { \ - LAUNCH_CACHED_QUANT_CALL(4, quantize::Type::Asymmetric) \ - } else { \ - LAUNCH_CACHED_QUANT_CALL(4, quantize::Type::Symmetric) \ - } \ - } else { \ - if (quant_type == quantize::Type::Asymmetric) { \ - LAUNCH_CACHED_QUANT_CALL(8, quantize::Type::Asymmetric) \ - } else { \ - LAUNCH_CACHED_QUANT_CALL(8, quantize::Type::Symmetric) \ - } \ - } - -void launch_quant(int8_t* output_data, - float* params, - const __half* input_data, - const int groups, - const int elems_per_group, - const int num_bits, - const quantize::Type quant_type, - cudaStream_t stream) -{ - constexpr int max_threads = 256; - - constexpr int internal_unroll = 2; - - const bool is_subblock_schedule = (elems_per_group <= 128) ? true : false; - const int h_per_step = is_subblock_schedule ? quantize::h_per_load - : quantize::h_per_load * internal_unroll; - - // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of - // warp-sized blocks rather than stepping up to 64/96 threads - const int one_step_threads = next_pow2((elems_per_group + h_per_step - 1) / h_per_step); - const int threads_per_group = (one_step_threads < max_threads) ? one_step_threads : max_threads; - - const int groups_per_block = - is_subblock_schedule ? (max_threads + threads_per_group - 1) / threads_per_group : 1; - const int groups_launch = (groups_per_block + groups - 1) / groups_per_block; - - dim3 block(threads_per_group, groups_per_block); - dim3 grid(groups_launch); - - const int elems_per_step = threads_per_group * h_per_step; - const int external_unroll = (elems_per_group + elems_per_step - 1) / elems_per_step; - - if (is_subblock_schedule) { - // <=128 - if (threads_per_group == 1) { - LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 1); - } else if (threads_per_group == 2) { - LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 2); - } else if (threads_per_group == 4) { - LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 4); - } else if (threads_per_group == 8) { - LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 8); - } else if (threads_per_group == 16) { - LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 16); - } - } else if (external_unroll == 1) { - // 129 - 4096 elems - // (this can launch with 1-7 warps as well) - LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, internal_unroll, max_threads); - } else if (external_unroll == 2) { - // 4097 - 8192 elems - LAUNCH_CACHED_QUANT(num_bits, quant_type, 2, internal_unroll, max_threads); - } else if (external_unroll == 3) { - // 8193 - 12288 elems - LAUNCH_CACHED_QUANT(num_bits, quant_type, 3, internal_unroll, max_threads); - } else if (external_unroll == 4) { - // 12289 - 16384 elems - LAUNCH_CACHED_QUANT(num_bits, quant_type, 4, internal_unroll, max_threads); - } -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quantize_intX.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quantize_intX.cu deleted file mode 100644 index b26151ab5c8c327d7b5b54239b73ee3ebd6bea8e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/quantize_intX.cu +++ /dev/null @@ -1,281 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include -#include -#include "memory_access_utils.h" - -template -struct alignas(sizeof(T) * N) AlignedArray { - using Element = T; - static const int kElements = N; - - __device__ __host__ AlignedArray() {} - - __device__ __host__ AlignedArray(const T& rhs) - { -#pragma unroll - for (int idx = 0; idx < kElements; ++idx) { this->at(idx) = rhs; } - } - - __device__ __host__ T& operator[](int offset) - { - return reinterpret_cast(this->buffer[offset]); - } - - __device__ __host__ const T& operator[](int offset) const - { - return reinterpret_cast(this->buffer[offset]); - } - - __device__ __host__ T& at(int offset) { return reinterpret_cast(this->buffer[offset]); } - - __device__ __host__ const T& at(int offset) const - { - return reinterpret_cast(this->buffer[offset]); - } - - __device__ __host__ AlignedArray operator+(const AlignedArray& rhs) const - { - AlignedArray ret; - -#pragma unroll - for (int idx = 0; idx < kElements; ++idx) { ret[idx] = this->at(idx) + rhs.at(idx); } - - return ret; - } - - __device__ __forceinline__ void clear() - { -#pragma unroll - for (int idx = 0; idx < kElements; ++idx) { this->at(idx) = Element(0); } - } - - Element buffer[N]; -}; - -template -struct reduce_max { - __device__ __forceinline__ T operator()(const T& lhs, const T& rhs) - { - return lhs > rhs ? lhs : rhs; - } -}; - -template -struct reduce_min { - __device__ __forceinline__ T operator()(const T& lhs, const T& rhs) - { - return lhs < rhs ? lhs : rhs; - } -}; - -template -struct subtract { - __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, - const T& rhs) - { - AlignedArray ret; - -#pragma unroll - for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] - rhs; } - - return ret; - } -}; - -template -struct plus { - __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, - const T& rhs) - { - AlignedArray ret; - -#pragma unroll - for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] + rhs; } - - return ret; - } -}; - -template -struct multiply { - __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, - const T& rhs) - { - AlignedArray ret; - -#pragma unroll - for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] * rhs; } - - return ret; - } -}; - -template -struct clamp { - __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, - const T& min_val, - const T& max_val) - { - AlignedArray ret; - -#pragma unroll - for (int idx = 0; idx < N; ++idx) { - ret[idx] = reduce_max()(reduce_min()(lhs[idx], max_val), min_val); - } - - return ret; - } -}; - -template -struct round_int; - -template -struct round_int { - __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs) - { - AlignedArray ret; - -#pragma unroll - for (int idx = 0; idx < N; ++idx) { ret[idx] = hrint(lhs[idx]); } - - return ret; - } -}; - -template -struct divide { - __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, - const T& rhs) - { - AlignedArray ret; - -#pragma unroll - for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] / rhs; } - - return ret; - } -}; - -template -__device__ __forceinline__ T to_scalar(const AlignedArray& data) -{ - Reducer re; - T res = data[0]; - -#pragma unroll - for (int idx = 1; idx < N; ++idx) { res = re(res, data[idx]); } - - return res; -} - -template -__device__ __forceinline__ AlignedArray int4_to_half( - const AlignedArray& data) -{ - AlignedArray ret; - -#pragma unroll - for (int idx = 0; idx < N * 2; idx += 2) { - ret[idx] = half(int(data[idx / 2] >> 4)); - ret[idx + 1] = half(int(data[idx / 2] & 0xf)); - } - - return ret; -} - -__global__ void dequantize_int4_to_half(uint8_t* data_in, - half* data_out, - half* scale_buffer, - half* min_val_buffer, - int num_group, - int group_size) -{ - using AccessType = AlignedArray; - using AccessTypeOut = AlignedArray; - - for (int idx = threadIdx.x + blockIdx.x * blockDim.x; idx < num_group * group_size / 8; - idx += blockDim.x * gridDim.x) { - int id_group = idx / (group_size / 8); - AccessType value = reinterpret_cast(data_in)[idx]; - half scale = scale_buffer[id_group]; - half min_value = min_val_buffer[id_group]; - - AccessTypeOut output = int4_to_half(value); - output = divide()(output, scale); - output = plus()(output, min_value); - - reinterpret_cast(data_out)[idx] = output; - } -} - -void launch_dequantize_int4_to_half_experimental(uint8_t* data_in, - half* data_out, - half* scale_buffer, - half* min_val_buffer, - int num_group, - int group_size, - cudaStream_t stream) -{ - int num_warp = num_group / 4; - int num_block = num_warp / 8; // 256 trd / block - - dequantize_int4_to_half<<>>( - data_in, data_out, scale_buffer, min_val_buffer, num_group, group_size); -} - -template -__device__ __forceinline__ AlignedArray int8_to_half(const AlignedArray& data) -{ - AlignedArray ret; - -#pragma unroll - for (int idx = 0; idx < N; idx += 1) { ret[idx] = half(int(data[idx])); } - - return ret; -} - -__global__ void dequantize_int8_to_half(uint8_t* data_in, - half* data_out, - half* scale_buffer, - half* min_val_buffer, - int num_group, - int group_size) -{ - using AccessType = AlignedArray; - using AccessTypeOut = AlignedArray; - - for (int idx = threadIdx.x + blockIdx.x * blockDim.x; idx < num_group * group_size / 8; - idx += blockDim.x * gridDim.x) { - int id_group = idx / (group_size / 8); - AccessType value = reinterpret_cast(data_in)[idx]; - half scale = scale_buffer[id_group]; - half min_value = min_val_buffer[id_group]; - - AccessTypeOut output = int8_to_half(value); - output = divide()(output, scale); - output = plus()(output, min_value); - - reinterpret_cast(data_out)[idx] = output; - } -} - -void launch_dequantize_int8_to_half_experimental(uint8_t* data_in, - half* data_out, - half* scale_buffer, - half* min_val_buffer, - int num_group, - int group_size, - cudaStream_t stream) -{ - int num_warp = num_group / 4; - int num_block = num_warp / 8; // 256 trd / block - - dequantize_int8_to_half<<>>( - data_in, data_out, scale_buffer, min_val_buffer, num_group, group_size); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/swizzled_quantize.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/swizzled_quantize.cu deleted file mode 100644 index a4b6096c81af102cab9f71a946024f7117902e2e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/quantization/swizzled_quantize.cu +++ /dev/null @@ -1,427 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "dequantization_utils.h" -#include "memory_access_utils.h" -#include "quantization_utils.h" -#include "reduction_utils.h" - -using rop = reduce::ROpType; - -namespace swiz_quant { -constexpr int max_threads = 512; -constexpr int min_threads = 32; - -constexpr int step_granularity = 2; -constexpr int h_per_step = step_granularity * quantize::h_per_load; -} // namespace swiz_quant - -template -__global__ void swizzled_quant_kernel(int8_t* quantized_data, - float* quantized_scales, - const __half* uncompressed_data, - int elems_per_group, - int nodes, - int devices_per_node) -{ - cg::thread_block tb = cg::this_thread_block(); - cg::thread_block_tile warp = cg::tiled_partition(tb); - - // Indexing offsets, same as normal quantization for in-case - const int block_rank = blockIdx.x + blockIdx.y * gridDim.x + blockIdx.z * gridDim.x * gridDim.y; - const int block_offset = block_rank * elems_per_group; - const int elem_offset = tb.thread_index().x * quantize::h_per_load; - const int base_offset = block_offset + elem_offset; - const int stride = tb.size() * quantize::h_per_load; - const __half* input_base = uncompressed_data + base_offset; - - // Local buffer - __half2 local_buffer[totalChunks * quantize::h2_per_load]; - - quantize::GroupStats stats; -#pragma unroll - for (int i = 0; i < totalChunks; i++) { - __half2* iteration_buffer = local_buffer + i * quantize::h2_per_load; - - mem_access::load_global( - iteration_buffer, input_base + i * stride, elem_offset + i * stride < elems_per_group); - -#pragma unroll - for (int j = 0; j < quantize::h2_per_load; j++) { stats.update(iteration_buffer[j]); } - } - - auto params = stats.template get_params(tb, warp); - - const int partition_id = blockIdx.z; - const int partition_offset = partition_id / devices_per_node; - const int partition_base = (partition_id % devices_per_node) * nodes; - const int pipelining_offset = blockIdx.y * (devices_per_node * nodes); - const int output_partition = (pipelining_offset + partition_base + partition_offset); - - constexpr int out_scalar_effect = 8 / numBits; - const int out_block_rank = output_partition * gridDim.x + blockIdx.x; - const int out_block_offset = out_block_rank * elems_per_group / out_scalar_effect; - const int out_base_offset = out_block_offset + elem_offset / out_scalar_effect; - int8_t* out_base = quantized_data + out_base_offset; - - const int out_stride = stride / out_scalar_effect; - constexpr int num_int8_out = quantize::h_per_load / out_scalar_effect; - - if (tb.thread_index().x == 0) { params.store(quantized_scales, out_block_rank); } - -#pragma unroll - for (int i = 0; i < totalChunks; i++) { - if (i * stride + elem_offset < elems_per_group) { - int8_t local_output[quantize::h_per_load / out_scalar_effect]; - quantize::_chunk( - local_output, local_buffer + i * quantize::h2_per_load, params); - mem_access::store_global(out_base + i * out_stride, local_output); - } - } -} - -#define LAUNCH_SWIZZLE_QUANT(total_chunks, threads) \ - swizzled_quant_kernel<<>>( \ - q_data, q_scales, input_data, elems_per_group, nodes, devices_per_node); - -/* -Swizzled quantization reorganizes the quantized groups in order to better facilitate -communication. As an example of the partitioning scheme we have the following example -of 2 node, 4 device swizzling: - - --- --- --- --- --- --- --- --- -| 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | - --- --- --- --- --- --- --- --- -becomes - --- --- --- --- --- --- --- --- -| 0 | 4 | 1 | 5 | 2 | 6 | 3 | 7 | - --- --- --- --- --- --- --- --- - -Multiple quantization groups may be mapped into a single partition. In order to better support -later pipelining, we may also perform an additional slicing. In two-way slicing, for instance, -the first halves of each partition are concatenated. -*/ - -template -void launch_swizzled_quant_impl(int8_t* q_data, - float* q_scales, - const __half* input_data, - int groups, - int elems_per_group, - int pipelining, - int nodes, - int devices_per_node, - cudaStream_t stream) -{ - const int one_step_threads = - next_pow2((elems_per_group + swiz_quant::h_per_step - 1) / swiz_quant::h_per_step); - const int max_threads = (one_step_threads < swiz_quant::max_threads) ? one_step_threads - : swiz_quant::max_threads; - const int threads = (max_threads < swiz_quant::min_threads) ? swiz_quant::min_threads - : max_threads; - - dim3 block(threads); - const int groups_per_partition = groups / (nodes * devices_per_node); - assert(groups_per_partition % pipelining == 0); - const int contiguous_groups = groups_per_partition / pipelining; - const int partitions = nodes * devices_per_node; - dim3 grid(contiguous_groups, pipelining, partitions); - - const int elems_per_step = threads * swiz_quant::h_per_step; - const int external_unroll = ((elems_per_group + elems_per_step - 1) / elems_per_step); - const int total_unroll = external_unroll * swiz_quant::step_granularity; - - assert(total_unroll % 2 == 0); - - if (threads == 32) { - LAUNCH_SWIZZLE_QUANT(2, 32); - } else if (threads == 64) { - LAUNCH_SWIZZLE_QUANT(2, 64); - } else if (threads == 128) { - LAUNCH_SWIZZLE_QUANT(2, 128); - } else if (threads == 256) { - LAUNCH_SWIZZLE_QUANT(2, 256); - } else if (threads == 512) { - if (total_unroll == 2) { - LAUNCH_SWIZZLE_QUANT(2, 512); - } else if (total_unroll == 4) { - LAUNCH_SWIZZLE_QUANT(4, 512); - } else if (total_unroll == 6) { - LAUNCH_SWIZZLE_QUANT(6, 512); - } else if (total_unroll == 8) { - LAUNCH_SWIZZLE_QUANT(8, 512); - } else if (total_unroll == 10) { - LAUNCH_SWIZZLE_QUANT(10, 512); - } - } -} - -#define DISPATCH_SWIZZLE_QUANT(num_bits, qtype) \ - launch_swizzled_quant_impl(q_data, \ - q_scales, \ - input_data, \ - groups, \ - elems_per_group, \ - pipelining, \ - nodes, \ - devices_per_node, \ - stream); - -void launch_swizzled_quant(int8_t* q_data, - float* q_scales, - const __half* input_data, - int num_bits, - quantize::Type q_type, - int groups, - int elems_per_group, - int pipelining, - int nodes, - int devices_per_node, - cudaStream_t stream) -{ - if (num_bits == 4) { - if (q_type == quantize::Type::Asymmetric) { - DISPATCH_SWIZZLE_QUANT(4, quantize::Type::Asymmetric); - } else if (q_type == quantize::Type::Symmetric) { - DISPATCH_SWIZZLE_QUANT(4, quantize::Type::Symmetric); - } - } else if (num_bits == 8) { - if (q_type == quantize::Type::Asymmetric) { - DISPATCH_SWIZZLE_QUANT(8, quantize::Type::Asymmetric); - } else if (q_type == quantize::Type::Symmetric) { - DISPATCH_SWIZZLE_QUANT(8, quantize::Type::Symmetric); - } - } -} - -template -__global__ void loco_swizzled_quant_kernel(int8_t* quantized_data, - float* quantized_scales, - const __half* uncompressed_data, - __half* error_feedback, - const float err_beta, - int groups, - int elems_per_group, - int pipelining, - int nodes, - int devices_per_node) -{ - cg::thread_block tb = cg::this_thread_block(); - cg::thread_block_tile warp = cg::tiled_partition(tb); - - // Indexing offsets, same as normal quantization for in-case - const int block_rank_data = - blockIdx.x + blockIdx.y * gridDim.x + blockIdx.z * gridDim.x * gridDim.y; - const int block_offset_data = block_rank_data * elems_per_group; - const int elem_offset = tb.thread_index().x * quantize::h_per_load; - const int base_offset_data = block_offset_data + elem_offset; - const int stride = tb.size() * quantize::h_per_load; - const __half* uncompressed_data_base = uncompressed_data + base_offset_data; - - const int partition_id = blockIdx.z; - const int partition_offset = partition_id / devices_per_node; - const int partition_base = (partition_id % devices_per_node) * nodes; - const int pipelining_offset = blockIdx.y * (devices_per_node * nodes); - const int output_partition = (pipelining_offset + partition_base + partition_offset); - const int block_rank_err = output_partition * gridDim.x + blockIdx.x; - - const int block_offset_err = block_rank_err * elems_per_group; - const int base_offset_err = block_offset_err + elem_offset; - __half* error_feedback_base = error_feedback + base_offset_err; - - __half2 local_buffer[totalChunks * quantize::h2_per_load]; - __half2 err_buffer[totalChunks * quantize::h2_per_load]; - - quantize::GroupStats stats; - -#pragma unroll - for (int i = 0; i < totalChunks; i++) { - __half2* iteration_buffer = local_buffer + i * quantize::h2_per_load; - __half2* iter_err_buffer = err_buffer + i * quantize::h2_per_load; - const int i_stride = i * stride; - bool do_loads = (elem_offset + i_stride) < elems_per_group; - - mem_access::load_global( - iteration_buffer, uncompressed_data_base + i_stride, do_loads); - - mem_access::load_global( - iter_err_buffer, error_feedback_base + i_stride, do_loads); - -#pragma unroll - for (int j = 0; j < quantize::h2_per_load; j++) { - iteration_buffer[j] = __hadd2(iteration_buffer[j], iter_err_buffer[j]); - stats.update(iteration_buffer[j]); - } - } - - auto params = stats.template get_params(tb, warp); - - // Initialize dequantization parameters based on params - auto de_params = params; - de_params.scale = 1.0f / params.scale; - if constexpr (quantType == quantize::Type::Asymmetric) { de_params.offset = params.offset; } - - if (threadIdx.x == 0) { params.store(quantized_scales, block_rank_err); } - - constexpr int out_scalar_effect = 8 / numBits; - const int out_block_offset = block_rank_err * elems_per_group / out_scalar_effect; - const int out_base_offset = out_block_offset + elem_offset / out_scalar_effect; - int8_t* out_base = quantized_data + out_base_offset; - - const int out_stride = stride / out_scalar_effect; - constexpr int num_int8_out = quantize::h_per_load / out_scalar_effect; - -#pragma unroll - for (int i = 0; i < totalChunks; i++) { - const int i_stride = i * stride; - __half2* iteration_buffer = local_buffer + i * quantize::h2_per_load; - __half2* iter_err_buffer = err_buffer + i * quantize::h2_per_load; - - if (i_stride + elem_offset < elems_per_group) { - int8_t local_output[quantize::h_per_load / out_scalar_effect]; - quantize::_chunk(local_output, iteration_buffer, params); - mem_access::store_global(out_base + i * out_stride, local_output); - - // Dequantize the quantized output to compute the dequantized value - __half2 dequant_buffer[quantize::h2_per_load]; - dequantize::chunk(dequant_buffer, local_output, de_params); - -// Compute new error: sum - dequant_buffer -#pragma unroll - for (int k = 0; k < quantize::h2_per_load; k++) { - // __half2 to float2 - float2 iter_buf_f = __half22float2(iteration_buffer[k]); - float2 dequant_buf_f = __half22float2(dequant_buffer[k]); - - // Update within float precision - float2 new_error_f; - new_error_f.x = iter_buf_f.x - dequant_buf_f.x; - new_error_f.y = iter_buf_f.y - dequant_buf_f.y; - - float2 iter_err_buf_f = __half22float2(iter_err_buffer[k]); - - iter_err_buf_f.x = err_beta * iter_err_buf_f.x + (1.0f - err_beta) * new_error_f.x; - iter_err_buf_f.y = err_beta * iter_err_buf_f.y + (1.0f - err_beta) * new_error_f.y; - - // float2 back to __half2 - iter_err_buffer[k] = __float22half2_rn(iter_err_buf_f); - } - __half2* error_feedback_base_h2 = reinterpret_cast<__half2*>(error_feedback_base); - mem_access::store_global(error_feedback_base_h2 + i_stride / 2, - iter_err_buffer); - } - } -} - -#define LAUNCH_LOCO_SWIZZLE_QUANT(total_chunks, threads) \ - loco_swizzled_quant_kernel \ - <<>>(output_data, \ - params, \ - input_data, \ - error_feedback, \ - err_beta, \ - groups, \ - elems_per_group, \ - pipelining, \ - nodes, \ - devices_per_node); - -template -void launch_loco_swizzled_quant_impl(int8_t* output_data, - float* params, - const __half* input_data, - __half* error_feedback, - const float err_beta, - int groups, - int elems_per_group, - int pipelining, - int nodes, - int devices_per_node, - cudaStream_t stream) -{ - const int one_step_threads = - next_pow2((elems_per_group + swiz_quant::h_per_step - 1) / swiz_quant::h_per_step); - const int max_threads = (one_step_threads < swiz_quant::max_threads) ? one_step_threads - : swiz_quant::max_threads; - const int threads = (max_threads < swiz_quant::min_threads) ? swiz_quant::min_threads - : max_threads; - - dim3 block(threads); - const int groups_per_partition = groups / (nodes * devices_per_node); - assert(groups_per_partition % pipelining == 0); - const int contiguous_groups = groups_per_partition / pipelining; - const int partitions = nodes * devices_per_node; - dim3 grid(contiguous_groups, pipelining, partitions); - - const int elems_per_step = threads * swiz_quant::h_per_step; - const int external_unroll = ((elems_per_group + elems_per_step - 1) / elems_per_step); - const int total_unroll = external_unroll * swiz_quant::step_granularity; - - assert(total_unroll % 2 == 0); - - if (threads == 32) { - LAUNCH_LOCO_SWIZZLE_QUANT(2, 32); - } else if (threads == 64) { - LAUNCH_LOCO_SWIZZLE_QUANT(2, 64); - } else if (threads == 128) { - LAUNCH_LOCO_SWIZZLE_QUANT(2, 128); - } else if (threads == 256) { - LAUNCH_LOCO_SWIZZLE_QUANT(2, 256); - } else if (threads == 512) { - if (total_unroll == 2) { - LAUNCH_LOCO_SWIZZLE_QUANT(2, 512); - } else if (total_unroll == 4) { - LAUNCH_LOCO_SWIZZLE_QUANT(4, 512); - } else if (total_unroll == 6) { - LAUNCH_LOCO_SWIZZLE_QUANT(6, 512); - } else if (total_unroll == 8) { - LAUNCH_LOCO_SWIZZLE_QUANT(8, 512); - } else if (total_unroll == 10) { - LAUNCH_LOCO_SWIZZLE_QUANT(10, 512); - } - } -} - -#define DISPATCH_LOCO_SWIZZLE_QUANT(num_bits, qtype) \ - launch_loco_swizzled_quant_impl(output_data, \ - params, \ - input_data, \ - error_feedback, \ - err_beta, \ - groups, \ - elems_per_group, \ - pipelining, \ - nodes, \ - devices_per_node, \ - stream); - -void launch_loco_swizzled_quant(int8_t* output_data, - float* params, - const __half* input_data, - __half* error_feedback, - const float err_beta, - int num_bits, - quantize::Type q_type, - int groups, - int elems_per_group, - int pipelining, - int nodes, - int devices_per_node, - cudaStream_t stream) -{ - if (num_bits == 4) { - if (q_type == quantize::Type::Asymmetric) { - DISPATCH_LOCO_SWIZZLE_QUANT(4, quantize::Type::Asymmetric); - } else if (q_type == quantize::Type::Symmetric) { - DISPATCH_LOCO_SWIZZLE_QUANT(4, quantize::Type::Symmetric); - } - } else if (num_bits == 8) { - if (q_type == quantize::Type::Asymmetric) { - DISPATCH_LOCO_SWIZZLE_QUANT(8, quantize::Type::Asymmetric); - } else if (q_type == quantize::Type::Symmetric) { - DISPATCH_LOCO_SWIZZLE_QUANT(8, quantize::Type::Symmetric); - } - } -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/gather_scatter.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/gather_scatter.cu deleted file mode 100644 index 291c8eb063b0e0a56843fb43d83caa5fdb00f6cb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/gather_scatter.cu +++ /dev/null @@ -1,186 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "custom_cuda_layers.h" -#include "memory_access_utils.h" - -namespace cg = cooperative_groups; - -namespace td_data { -constexpr int granularity = 16; -} - -template -__global__ void gather_tokens_impl(T* retained_tokens, - const T* activations, - int32_t* gather_indices, - int32_t sampled_tokens, - int32_t channels, - int32_t read_batch_stride, - int32_t read_seq_stride, - int32_t write_batch_stride, - int32_t write_seq_stride) -{ - constexpr int mem_vals_t = td_data::granularity / sizeof(T); - - cg::thread_block tb = cg::this_thread_block(); - - const int gather_idx = gather_indices[tb.group_index().x * sampled_tokens + tb.group_index().y]; - - const int read_offset = read_batch_stride * tb.group_index().x + read_seq_stride * gather_idx; - const int write_offset = - write_batch_stride * tb.group_index().x + write_seq_stride * tb.group_index().y; - - for (int i = tb.thread_index().x * mem_vals_t; i < channels; i += blockDim.x * mem_vals_t) { - T local_data[mem_vals_t]; - mem_access::load_global(local_data, activations + read_offset + i); - mem_access::store_global(retained_tokens + write_offset + i, - local_data); - } -} - -template -void launch_gather_tokens(T* retained_tokens, - T* activations, - int32_t* gather_indices, - int32_t batch_size, - int32_t sampled_tokens, - int32_t channels, - int32_t read_batch_stride, - int32_t read_seq_stride, - int32_t write_batch_stride, - int32_t write_seq_stride, - cudaStream_t stream) -{ - constexpr int mem_vals_t = td_data::granularity / sizeof(T); - - const int load_steps = (channels + mem_vals_t - 1) / mem_vals_t; - const int threads = (load_steps >= 1024) ? 1024 : load_steps; - - dim3 block(threads); - dim3 grid(batch_size, sampled_tokens); - - gather_tokens_impl<<>>(retained_tokens, - activations, - gather_indices, - sampled_tokens, - channels, - read_batch_stride, - read_seq_stride, - write_batch_stride, - write_seq_stride); -} - -template void launch_gather_tokens(float*, - float*, - int32_t*, - int32_t, - int32_t, - int32_t, - int32_t, - int32_t, - int32_t, - int32_t, - cudaStream_t); - -template void launch_gather_tokens<__half>(__half*, - __half*, - int32_t*, - int32_t, - int32_t, - int32_t, - int32_t, - int32_t, - int32_t, - int32_t, - cudaStream_t); - -template -__global__ void scatter_tokens_impl(T* all_activations, - const T* layer_activations, - int32_t* gather_indices, - int32_t retained_tokens, - int32_t channels, - int32_t read_batch_stride, - int32_t read_seq_stride, - int32_t write_batch_stride, - int32_t write_seq_stride) -{ - constexpr int mem_vals_t = td_data::granularity / sizeof(T); - - cg::thread_block tb = cg::this_thread_block(); - - const int gather_idx = - gather_indices[tb.group_index().x * retained_tokens + tb.group_index().y]; - - const int read_offset = - read_batch_stride * tb.group_index().x + read_seq_stride * tb.group_index().y; - const int write_offset = - write_batch_stride * tb.group_index().x + write_seq_stride * gather_idx; - - for (int i = tb.thread_index().x * mem_vals_t; i < channels; i += mem_vals_t * blockDim.x) { - T local_data[mem_vals_t]; - mem_access::load_global(local_data, - layer_activations + read_offset + i); - mem_access::store_global(all_activations + write_offset + i, - local_data); - } -} - -template -void launch_scatter_tokens(T* all_activations, - T* layer_activations, - int32_t* gather_indices, - int32_t batch_size, - int32_t sampled_tokens, - int32_t channels, - int32_t read_batch_stride, - int32_t read_seq_stride, - int32_t write_batch_stride, - int32_t write_seq_stride, - cudaStream_t stream) -{ - constexpr int mem_vals_t = td_data::granularity / sizeof(T); - - const int load_steps = (channels + mem_vals_t - 1) / mem_vals_t; - const int threads = (load_steps >= 1024) ? 1024 : load_steps; - - dim3 block(threads); - dim3 grid(batch_size, sampled_tokens); - - scatter_tokens_impl<<>>(all_activations, - layer_activations, - gather_indices, - sampled_tokens, - channels, - read_batch_stride, - read_seq_stride, - write_batch_stride, - write_seq_stride); -} - -template void launch_scatter_tokens(float*, - float*, - int32_t*, - int32_t, - int32_t, - int32_t, - int32_t, - int32_t, - int32_t, - int32_t, - cudaStream_t); - -template void launch_scatter_tokens<__half>(__half*, - __half*, - int32_t*, - int32_t, - int32_t, - int32_t, - int32_t, - int32_t, - int32_t, - int32_t, - cudaStream_t); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/pt_binding.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/pt_binding.cpp deleted file mode 100644 index b50476719872e80fe785cac2c17501b24e1a01e4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/pt_binding.cpp +++ /dev/null @@ -1,216 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include -#include "custom_cuda_layers.h" - -torch::Tensor token_sort_(torch::Tensor& unsorted_token_ids, int64_t original_tokens) -{ - const int layers = unsorted_token_ids.size(0); - const int batch_size = unsorted_token_ids.size(1); - const int reserved_tokens = unsorted_token_ids.size(2); - - launch_token_sort(unsorted_token_ids.data_ptr(), - layers, - batch_size, - reserved_tokens, - original_tokens, - c10::cuda::getCurrentCUDAStream()); - - return unsorted_token_ids; -} - -torch::Tensor token_gather(torch::Tensor& activations, - torch::Tensor& sorted_indices, - bool batch_first) -{ - // Activations may be in either [N, S, C] or [S, N, C] while sorted_indices is - // always in [N, retained] - /* - TORCH_CHECK(sorted_indices.size(0) == activations.size(0) || - sorted_indices.size(0) == activations.size(1), - "Unable to match the batch size of the sorted indices to the activation - shape."); TORCH_CHECK(activations.size(2) % 8 == 0, "Channels must be divisible by 8 to align - with vectorized loads."); - */ - // bool batch_first = sorted_indices.size(0) == activations.size(0); - - const int64_t dim_0 = (batch_first) ? sorted_indices.size(0) : sorted_indices.size(1); - const int64_t dim_1 = (batch_first) ? sorted_indices.size(1) : sorted_indices.size(0); - const int64_t dim_2 = activations.size(2); - - auto output = torch::empty({dim_0, dim_1, dim_2}, activations.options()); - - const int batch_size = sorted_indices.size(0); - const int channels = dim_2; - const int retained_tokens = sorted_indices.size(1); - const int read_batch_stride = (batch_first) ? activations.stride(0) : activations.stride(1); - const int read_seq_stride = (batch_first) ? activations.stride(1) : activations.stride(0); - const int write_batch_stride = (batch_first) ? output.stride(0) : output.stride(1); - const int write_seq_stride = (batch_first) ? output.stride(1) : output.stride(0); - - if (activations.options().dtype() == torch::kFloat) { - launch_gather_tokens((float*)output.data_ptr(), - (float*)activations.data_ptr(), - (int32_t*)sorted_indices.data_ptr(), - batch_size, - retained_tokens, - channels, - read_batch_stride, - read_seq_stride, - write_batch_stride, - write_seq_stride, - c10::cuda::getCurrentCUDAStream()); - } else { - launch_gather_tokens((__half*)output.data_ptr(), - (__half*)activations.data_ptr(), - (int32_t*)sorted_indices.data_ptr(), - batch_size, - retained_tokens, - channels, - read_batch_stride, - read_seq_stride, - write_batch_stride, - write_seq_stride, - c10::cuda::getCurrentCUDAStream()); - } - - return output; -} - -torch::Tensor token_scatter_(torch::Tensor& all_activations, - torch::Tensor& layer_activations, - torch::Tensor& sorted_indices, - bool batch_first) -{ - // Activations may be in either [N, S, C] or [S, N, C] while sorted_indices is - // always in [N, retained] - /* - TORCH_CHECK(sorted_indices.size(0) == all_activations.size(0) || - sorted_indices.size(0) == all_activations.size(1), - "Unable to match the batch size of the sorted indices to the activation - shape."); TORCH_CHECK(all_activations.size(2) % 8 != 0, "Channels must be divisible by 8 to - align with vectorized loads."); - */ - // bool batch_first = sorted_indices.size(0) == all_activations.size(0); - - const int batch_size = sorted_indices.size(0); - const int channels = all_activations.size(2); - const int retained_tokens = sorted_indices.size(1); - const int read_batch_stride = (batch_first) ? layer_activations.stride(0) - : layer_activations.stride(1); - const int read_seq_stride = (batch_first) ? layer_activations.stride(1) - : layer_activations.stride(0); - const int write_batch_stride = (batch_first) ? all_activations.stride(0) - : all_activations.stride(1); - const int write_seq_stride = (batch_first) ? all_activations.stride(1) - : all_activations.stride(0); - - if (all_activations.options().dtype() == torch::kFloat) { - launch_scatter_tokens((float*)all_activations.data_ptr(), - (float*)layer_activations.data_ptr(), - (int32_t*)sorted_indices.data_ptr(), - batch_size, - retained_tokens, - channels, - read_batch_stride, - read_seq_stride, - write_batch_stride, - write_seq_stride, - c10::cuda::getCurrentCUDAStream()); - } else { - launch_scatter_tokens((__half*)all_activations.data_ptr(), - (__half*)layer_activations.data_ptr(), - (int32_t*)sorted_indices.data_ptr(), - batch_size, - retained_tokens, - channels, - read_batch_stride, - read_seq_stride, - write_batch_stride, - write_seq_stride, - c10::cuda::getCurrentCUDAStream()); - } - - return all_activations; -} - -torch::Tensor mask_gather_bert(torch::Tensor& dense_mask, torch::Tensor& sorted_indices) -{ - // TORCH_CHECK(dense_mask.dim() == 4) - - const int batch_size = dense_mask.size(0); - const int layers = sorted_indices.size(0); - /* - TORCH_CHECK(layers * batch_size == sorted_indices.size(0), - "Mismatch between the indices and the mask"); - */ - const int orig_seq_len = dense_mask.size(3); - const int truncated_seq_len = sorted_indices.size(2); - - auto output = torch::empty({layers, batch_size, 1, truncated_seq_len, truncated_seq_len}, - dense_mask.options()); - - if (dense_mask.options().dtype() == torch::kFloat) { - launch_slice_bert_mask((float*)output.data_ptr(), - (const float*)dense_mask.data_ptr(), - (const int32_t*)sorted_indices.data_ptr(), - layers, - batch_size, - truncated_seq_len, - orig_seq_len, - c10::cuda::getCurrentCUDAStream()); - } else { - launch_slice_bert_mask((__half*)output.data_ptr(), - (const __half*)dense_mask.data_ptr(), - (const int32_t*)sorted_indices.data_ptr(), - layers, - batch_size, - truncated_seq_len, - orig_seq_len, - c10::cuda::getCurrentCUDAStream()); - } - - return output; -} - -torch::Tensor mask_gather_gpt(torch::Tensor dense_mask, int truncated_seq_len) -{ - // TORCH_CHECK(dense_mask.dim() == 4) - - const int batch_size = dense_mask.size(0); - const int orig_seq_len = dense_mask.size(3); - - auto output = - torch::empty({batch_size, 1, truncated_seq_len, truncated_seq_len}, dense_mask.options()); - - if (dense_mask.options().dtype() == torch::kFloat) { - launch_slice_gpt_mask((float*)output.data_ptr(), - (const float*)dense_mask.data_ptr(), - batch_size, - truncated_seq_len, - orig_seq_len, - c10::cuda::getCurrentCUDAStream()); - } else { - launch_slice_gpt_mask((__half*)output.data_ptr(), - (const __half*)dense_mask.data_ptr(), - batch_size, - truncated_seq_len, - orig_seq_len, - c10::cuda::getCurrentCUDAStream()); - } - - return output; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("token_sort_", &token_sort_, "Comparison free sorting algorithm (CUDA)"); - m.def("token_gather", &token_gather, "Parallel gather of tokens (CUDA)"); - m.def("token_scatter_", &token_scatter_, "Parallel scatter of tokens (CUDA)"); - m.def("mask_gather_bert", &mask_gather_bert, "Token-based mask gather for BERT masking (CUDA)"); - m.def("mask_gather_gpt", &mask_gather_gpt, "Token-based mask gather for GPT masking (CUDA)"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/slice_attn_masks.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/slice_attn_masks.cu deleted file mode 100644 index bc3823b846550cef1049e461ea35c35b53522146..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/slice_attn_masks.cu +++ /dev/null @@ -1,128 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "custom_cuda_layers.h" -#include "memory_access_utils.h" - -namespace cg = cooperative_groups; - -template -__global__ void slice_gpt_mask_impl(T* output_mask, - const T* input_mask, - int truncated_seq_len, - int orig_seq_len) -{ - const int in_batch_stride = orig_seq_len * orig_seq_len; - const int out_batch_stride = truncated_seq_len * truncated_seq_len; - - cg::thread_block tb = cg::this_thread_block(); - - const T* input_mask_block = - input_mask + blockIdx.x * in_batch_stride + blockIdx.y * orig_seq_len; - T* output_mask_block = - output_mask + blockIdx.x * out_batch_stride + blockIdx.y * truncated_seq_len; - - for (int i = tb.thread_index().x; i < truncated_seq_len; i += blockDim.x) { - output_mask_block[i] = input_mask_block[i]; - } -} - -template -void launch_slice_gpt_mask(T* output_mask, - const T* input_mask, - int batch_size, - int truncated_seq_len, - int orig_seq_len, - cudaStream_t stream) -{ - const int threads = (truncated_seq_len >= 1024) ? 1024 : truncated_seq_len; - - dim3 block(threads); - dim3 grid(batch_size, truncated_seq_len); - - slice_gpt_mask_impl - <<>>(output_mask, input_mask, truncated_seq_len, orig_seq_len); -} - -template void launch_slice_gpt_mask(float*, const float*, int, int, int, cudaStream_t); - -template void launch_slice_gpt_mask<__half>(__half*, const __half*, int, int, int, cudaStream_t); - -template -__global__ void slice_bert_mask_impl(T* output_mask, - const T* input_mask, - const int32_t* retained_indices, - int32_t truncated_seq_len, - int32_t orig_seq_len) -{ - const int in_batch_stride = orig_seq_len * orig_seq_len; - const int out_batch_stride = truncated_seq_len * truncated_seq_len; - const int out_layer_stride = out_batch_stride * gridDim.y; - - cg::thread_block tb = cg::this_thread_block(); - - const int out_layer_offset = tb.group_index().x * out_layer_stride; - - const int in_batch_offset = tb.group_index().y * in_batch_stride; - const int out_batch_offset = tb.group_index().y * out_batch_stride; - - const int32_t gather_row = - retained_indices[tb.group_index().y * truncated_seq_len + tb.group_index().z]; - const int in_seq_offset = gather_row * orig_seq_len; - const int out_seq_offset = tb.group_index().z * truncated_seq_len; - - const T* in_sequence = input_mask + in_batch_offset + in_seq_offset; - T* out_sequence = output_mask + out_layer_offset + out_batch_offset + out_seq_offset; - const int32_t* gather_data = retained_indices + tb.group_index().y * truncated_seq_len; - - for (int i = tb.thread_index().x; i < truncated_seq_len; i += blockDim.x) { - out_sequence[i] = in_sequence[gather_data[i]]; - } -} - -/* -Since the Bert mask is not causal like GPT, we can't just generate a set of -masks for the entire model based off a single layer sample. - -We map the kernel as follows: -z-dimension: layer -y-dimension: batch -x-dimension: sequence_offset -*/ -template -void launch_slice_bert_mask(T* output_mask, - const T* input_mask, - const int32_t* retained_indices, - int32_t layers, - int32_t batch_size, - int32_t truncated_seq_len, - int32_t orig_seq_len, - cudaStream_t stream) -{ - const int threads = (truncated_seq_len >= 1024) ? 1024 : truncated_seq_len; - dim3 block(threads); - dim3 grid(layers, batch_size, truncated_seq_len); - - slice_bert_mask_impl<<>>( - output_mask, input_mask, retained_indices, truncated_seq_len, orig_seq_len); -} - -template void launch_slice_bert_mask(float*, - const float*, - const int32_t*, - int32_t, - int32_t, - int32_t, - int32_t, - cudaStream_t); - -template void launch_slice_bert_mask<__half>(__half*, - const __half*, - const int32_t*, - int32_t, - int32_t, - int32_t, - int32_t, - cudaStream_t); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/token_sort.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/token_sort.cu deleted file mode 100644 index 3c1dff49429fc2e0f436ddc9d89a48af490da9c8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/random_ltd/token_sort.cu +++ /dev/null @@ -1,194 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include "custom_cuda_layers.h" -#include "memory_access_utils.h" - -namespace cg = cooperative_groups; - -namespace td_sort { -constexpr int threads = 512; -constexpr int granularity = 16; -constexpr int mem_vals = granularity / sizeof(int32_t); -constexpr int max_buffer_size = (threads + 1) * mem_vals; - -#ifdef __HIP_PLATFORM_AMD__ -constexpr int warp_size = ROCM_WAVEFRONT_SIZE; -#else -constexpr int warp_size = 32; -#endif - -constexpr int max_warps = threads / warp_size; -} // namespace td_sort - -template -__global__ void scan_sort(int32_t* data, int reserved_tokens, int original_tokens) -{ - cg::thread_block tb = cg::this_thread_block(); - cg::thread_block_tile warp = cg::tiled_partition(tb); - - __shared__ int32_t indices_buffer[td_sort::max_buffer_size]; - __shared__ int32_t intermediate_buffer[td_sort::max_warps]; - __shared__ int32_t sorted_indices_buffer[td_sort::max_buffer_size]; - - for (int i = tb.thread_index().x * td_sort::mem_vals; i < original_tokens + 1; - i += tb.group_dim().x * td_sort::mem_vals) { - uint32_t zeros[td_sort::mem_vals] = {0, 0, 0, 0}; - mem_access::store_shared(indices_buffer + i, zeros); - } - - int32_t local_vals[VALS_PER_THREAD]; - - // We flatten layers/batch into a single indexing dimension - int32_t* data_block = data + tb.group_index().x * reserved_tokens; - - // The next two loops really could be fused for a more logical code layout, but don't want to - // move the barrier forward -#pragma unroll - for (int i = 0; i < VALS_PER_THREAD; i++) { - const int iter_idx = i * td_sort::threads + tb.thread_index().x; - if (iter_idx < reserved_tokens) { - mem_access::load_global(local_vals + i, data_block + iter_idx); - } else { - local_vals[i] = 0; - } - } - - tb.sync(); - -#pragma unroll - for (int i = 0; i < VALS_PER_THREAD; i++) { - const int iter_idx = i * td_sort::threads + tb.thread_index().x; - if (iter_idx < reserved_tokens) { - const int32_t one = 1; - mem_access::store_shared(indices_buffer + local_vals[i], &one); - } - } - - tb.sync(); - - int32_t local_input[td_sort::mem_vals]; - mem_access::load_shared( - local_input, indices_buffer + tb.thread_index().x * td_sort::mem_vals); - - int32_t reduce_vals[td_sort::mem_vals]; - reduce_vals[0] = local_input[0]; - -#pragma unroll - for (int i = 1; i < td_sort::mem_vals; i++) { - reduce_vals[i] = local_input[i] + reduce_vals[i - 1]; - } - - int32_t step_1_val = reduce_vals[td_sort::mem_vals - 1]; - // Short span exclusive scan algorithm (less work efficient) -#pragma unroll - for (int i = 1; i < td_sort::warp_size; i *= 2) { - int32_t step_val = warp.shfl_up(step_1_val, i); - step_1_val = (warp.thread_rank() < i) ? step_1_val : step_1_val + step_val; - } - - if (warp.thread_rank() == td_sort::warp_size - 1) { - mem_access::store_shared(intermediate_buffer + warp.meta_group_rank(), - &step_1_val); - } - - tb.sync(); - - if (warp.meta_group_rank() == 0) { - int32_t step_2_val = 0; - if (warp.thread_rank() < td_sort::max_warps) { - mem_access::load_shared(&step_2_val, - intermediate_buffer + warp.thread_rank()); - } - -#pragma unroll - for (int i = 1; i < td_sort::warp_size; i *= 2) { - int32_t step_val = warp.shfl_up(step_2_val, i); - step_2_val = (warp.thread_rank() < i) ? step_2_val : step_2_val + step_val; - } - - if (warp.thread_rank() < td_sort::max_warps) { - mem_access::store_shared(intermediate_buffer + warp.thread_rank(), - &step_2_val); - } - } - - tb.sync(); - - int step_2_val = 0; - if (warp.meta_group_rank() > 0) { - mem_access::load_shared(&step_2_val, - intermediate_buffer + warp.meta_group_rank() - 1); - } - - const int thread_offset = reduce_vals[td_sort::mem_vals - 1]; - -#pragma unroll - for (int i = 0; i < td_sort::mem_vals; i++) { - reduce_vals[i] += step_1_val + step_2_val - thread_offset; - } - mem_access::store_shared( - indices_buffer + tb.thread_index().x * td_sort::mem_vals, reduce_vals); - - if (tb.thread_index().x == 0) { - indices_buffer[original_tokens] = original_tokens - indices_buffer[original_tokens]; - } - tb.sync(); - - for (int i = 0; i < VALS_PER_THREAD; i++) { - const int iter_idx = i * td_sort::threads + tb.thread_index().x; - if (iter_idx < reserved_tokens) { - if (local_vals[i] == 0) { - int zero = 0; - mem_access::store_shared(sorted_indices_buffer, &zero); - } else { - int sorted_idx; - mem_access::load_shared(&sorted_idx, - indices_buffer + local_vals[i] - 1); - mem_access::store_shared(sorted_indices_buffer + sorted_idx, - local_vals + i); - } - } - } - - tb.sync(); - -#pragma unroll - for (int i = 0; i < VALS_PER_THREAD; i++) { - const int iter_idx = i * td_sort::threads + tb.thread_index().x; - if (iter_idx < reserved_tokens) { - int32_t store_val; - mem_access::load_shared(&store_val, sorted_indices_buffer + iter_idx); - mem_access::store_global(data_block + iter_idx, &store_val); - } - } -} - -void launch_token_sort(int32_t* indices, - int layers, - int batch_size, - int reserved_size, - int original_tokens, - cudaStream_t stream) -{ - // Each sort is completely independent, can flatten this dimension - dim3 grid(layers * batch_size); - dim3 block(td_sort::threads); - - const int vals_per_thread = (reserved_size + td_sort::threads - 1) / td_sort::threads; - - if (vals_per_thread == 1) { - scan_sort<1><<>>(indices, reserved_size, original_tokens); - } else if (vals_per_thread == 2) { - scan_sort<2><<>>(indices, reserved_size, original_tokens); - } else if (vals_per_thread == 3) { - scan_sort<3><<>>(indices, reserved_size, original_tokens); - } else if (vals_per_thread == 4) { - scan_sort<4><<>>(indices, reserved_size, original_tokens); - } else { - assert(false); - } -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/sparse_attention/utils.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/sparse_attention/utils.cpp deleted file mode 100644 index 352306ba26128b96236f43764c26d2f4191eb391..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/sparse_attention/utils.cpp +++ /dev/null @@ -1,127 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a - https:github.com/ptillet/torch-blocksparse/blob/master/csrc/utils.cpp -*/ - -#include -#include -#include -#include -#ifdef _OPENMP -#include -#endif - -typedef std::vector> ret_t; - -void segment_blocks(torch::Tensor layout, - torch::Tensor idx, - torch::Tensor scratch, - int max_width, - ret_t& ret) -{ - size_t H = layout.size(0); - size_t M = layout.size(1); - size_t N = layout.size(2); - torch::Tensor tmp = torch::zeros_like(layout); - - auto _tmp = tmp.accessor(); - auto _layout = layout.accessor(); - auto _idx = idx.accessor(); - auto _scratch = scratch.accessor(); - std::vector current(H, 0); - -#ifdef _OPENMP -#pragma omp parallel for -#endif - for (size_t h = 0; h < H; h++) { - // surrounding indices - std::vector ii_left(max_width, -1); - std::vector> ii_top(max_width, std::vector(N, -1)); - - for (size_t m = 0; m < M; m++) { - for (size_t n = 0; n < N; n++) { - int v = _layout[h][m][n]; - if (v == 0) continue; - int n_left = ii_left[max_width - 1]; - int m_top = ii_top[max_width - 1][n]; - int top = (m_top >= 0) ? _tmp[h][m_top][n] : 0; - int left = (n_left >= 0) ? _tmp[h][m][n_left] : 0; - int topleft = (m_top >= 0 && n_left >= 0) ? _tmp[h][m_top][n_left] : 0; - int width = std::min(left, std::min(top, topleft)) + 1; - - // reset width if blocks cannot be - // packed together (i.e., there's a 1 "in the middle") - for (int nn = n_left + 1; nn < n; nn++) - if (ii_top[max_width - 1][nn] > ii_top[max_width - 1][n]) width = 1; - _tmp[h][m][n] = width; - - // update n_left ring buffer - for (int k = 0; k < max_width - 1; k++) ii_left[k] = ii_left[k + 1]; - ii_left[max_width - 1] = n; - - // update ii_top ring buffer - for (int k = 0; k < max_width - 1; k++) ii_top[k][n] = ii_top[k + 1][n]; - ii_top[max_width - 1][n] = m; - - // block is too small -- skip - if (width != max_width) continue; - - // retained blocks are set to zeros - for (size_t km = 0; km < max_width; km++) - for (size_t kn = 0; kn < max_width; kn++) { - int mm = ii_top[km][n]; - int nn = ii_left[kn]; - if (mm < 0 || nn < 0) continue; - _layout[h][mm][nn] = 0; - _tmp[h][mm][nn] = 0; - _scratch[h][current[h]][0] = (int)h; - _scratch[h][current[h]][1] = (int)mm; - _scratch[h][current[h]][2] = (int)nn; - _scratch[h][current[h]][3] = _idx[h][mm][nn]; - current[h]++; - } - } - } - } - std::vector to_cat; - for (size_t h = 0; h < H; h++) - if (current[h] > 0) to_cat.push_back(scratch[h].slice(0, 0, current[h])); - if (!to_cat.empty()) ret.push_back({max_width, torch::cat(to_cat)}); -} - -ret_t sdd_segment(torch::Tensor layout, int start_width) -{ - ret_t ret; - - // block index - torch::Tensor idx = torch::zeros_like(layout); - int current = 0; - int64_t H = layout.size(0); - int64_t M = layout.size(1); - int64_t N = layout.size(2); - auto _layout = layout.accessor(); - auto _idx = idx.accessor(); - for (int64_t h = 0; h < H; h++) - for (int64_t m = 0; m < M; m++) - for (int64_t n = 0; n < N; n++) { - if (_layout[h][m][n] == 0) continue; - _idx[h][m][n] = current++; - } - - // scratch memory - torch::Tensor scratch = torch::empty({H, layout.sum().item(), 4}, layout.dtype()); - - for (int max_width = start_width; max_width > 0; max_width /= 2) - segment_blocks(layout, idx, scratch, max_width, ret); - return ret; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("sdd_segment", &sdd_segment, "SDD segmentation handler"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/csrc/opt_bias_add.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/csrc/opt_bias_add.cu deleted file mode 100644 index d831b372b65f398b43c5cd343b2bf2db67562f56..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/csrc/opt_bias_add.cu +++ /dev/null @@ -1,149 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include "memory_access_utils.h" -#include "spatial_cuda_layers.h" - -/* -Fused bias add variants -*/ - -namespace badd_opt { -constexpr int threads = 256; -constexpr int steps = 2; -constexpr int granularity = 16; -constexpr int vals_per_h = granularity / sizeof(__half); -constexpr int vals_per_h2 = granularity / sizeof(__half2); -constexpr int vals_per_block = threads * steps * vals_per_h; -constexpr int stride = vals_per_h * threads; -} // namespace badd_opt - -__global__ void opt_bias_add(__half* result, - const __half* activation, - const __half* bias, - int seq_len, - int channels) -{ - const int id = blockIdx.x * badd_opt::vals_per_block + threadIdx.x * badd_opt::vals_per_h; - const int stride = badd_opt::vals_per_h * badd_opt::threads; - - for (int i = 0; i < badd_opt::steps; i++) { - if (id + i * badd_opt::stride < seq_len * channels) { - __half2 act_buffer[badd_opt::vals_per_h2]; - __half2 bias_buffer[badd_opt::vals_per_h2]; - - mem_access::load_global(act_buffer, - activation + id + i * stride); - mem_access::load_global(bias_buffer, - bias + ((id + i * stride) % channels)); - - for (int j = 0; j < badd_opt::vals_per_h2; j++) { act_buffer[j] += bias_buffer[j]; } - - mem_access::store_global(result + id + i * stride, act_buffer); - } - } -} - -__global__ void opt_bias_add_add(__half* result, - const __half* activation, - const __half* bias, - const __half* other, - int seq_len, - int channels) -{ - const int id = blockIdx.x * badd_opt::vals_per_block + threadIdx.x * badd_opt::vals_per_h; - const int stride = badd_opt::vals_per_h * badd_opt::threads; - - for (int i = 0; i < badd_opt::steps; i++) { - if (id + i * badd_opt::stride < seq_len * channels) { - __half2 act_buffer[badd_opt::vals_per_h2]; - __half2 bias_buffer[badd_opt::vals_per_h2]; - __half2 other_buffer[badd_opt::vals_per_h2]; - - mem_access::load_global(act_buffer, - activation + id + i * stride); - mem_access::load_global(bias_buffer, - bias + ((id + i * stride) % channels)); - mem_access::load_global(other_buffer, other + id + i * stride); - - for (int j = 0; j < badd_opt::vals_per_h2; j++) { - act_buffer[j] += bias_buffer[j] + other_buffer[j]; - } - - mem_access::store_global(result + id + i * stride, act_buffer); - } - } -} - -__global__ void opt_bias_add_bias_add(__half* result, - const __half* activation, - const __half* bias, - const __half* other, - const __half* other_bias, - int seq_len, - int channels) -{ - const int id = blockIdx.x * badd_opt::vals_per_block + threadIdx.x * badd_opt::vals_per_h; - const int stride = badd_opt::vals_per_h * badd_opt::threads; - - for (int i = 0; i < badd_opt::steps; i++) { - if (id + i * badd_opt::stride < seq_len * channels) { - __half2 act_buffer[badd_opt::vals_per_h2]; - __half2 bias_buffer[badd_opt::vals_per_h2]; - __half2 other_buffer[badd_opt::vals_per_h2]; - __half2 other_bias_buffer[badd_opt::vals_per_h2]; - - mem_access::load_global(act_buffer, - activation + id + i * stride); - mem_access::load_global(bias_buffer, - bias + ((id + i * stride) % channels)); - mem_access::load_global(other_buffer, other + id + i * stride); - mem_access::load_global( - other_bias_buffer, other_bias + ((id + i * stride) % channels)); - - for (int j = 0; j < badd_opt::vals_per_h2; j++) { - act_buffer[j] = - (act_buffer[j] + bias_buffer[j]) + (other_buffer[j] + other_bias_buffer[j]); - } - - mem_access::store_global(result + id + i * stride, act_buffer); - } - } -} - -void launch_opt_bias_add(__half* result, - const __half* activation, - const __half* bias, - const __half* other, - const __half* other_bias, - int batch_size, - int seq_len, - int channels, - cudaStream_t stream) -{ - // Should evaluate `true` for reasonable hidden sizes - assert(channels % badd_opt::vals_per_h == 0); - - const int effective_seq_len = batch_size * seq_len; - const int vals = effective_seq_len * channels; - - dim3 block(badd_opt::threads); - dim3 grid((vals + badd_opt::vals_per_block - 1) / badd_opt::vals_per_block); - - if (!other) { - // We shouldn't have a bias if there's no activation - assert(!other_bias); - - opt_bias_add<<>>( - result, activation, bias, effective_seq_len, channels); - } else if (!other_bias) { - opt_bias_add_add<<>>( - result, activation, bias, other, effective_seq_len, channels); - } else { - opt_bias_add_bias_add<<>>( - result, activation, bias, other, other_bias, effective_seq_len, channels); - } -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/csrc/pt_binding.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/csrc/pt_binding.cpp deleted file mode 100644 index cbf6636a6ee4ac5769a88b7cab7d2a222f00d5cc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/csrc/pt_binding.cpp +++ /dev/null @@ -1,112 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include -#include -#include -#include "spatial_cuda_layers.h" - -ChannelsLastProblem dimension_problem(at::Tensor& input) -{ - ChannelsLastProblem dims; - - if (input.dim() == 4) { - // In some sense this is unsafe (and a reflection of the assumptions made inside - // the C10 options checker). Basically, there's no great way to be sure that - // a tensor is in channels last because a 1x1 image will appear to be in channels - // last even when it isn't. - assert(input.is_contiguous(at::MemoryFormat::ChannelsLast)); - dims.batch_size = input.size(0); - dims.seq_len = input.size(2) * input.size(3); - dims.channels = input.size(1); - } else { - assert(input.is_contiguous()); - dims.batch_size = input.size(0); - dims.seq_len = input.size(1); - dims.channels = input.size(2); - } - - return dims; -} - -at::Tensor seq_unroll_bias_add(at::Tensor& input, at::Tensor& bias) -{ - assert(input.dtype() == at::kHalf); - - // TODO(cmikeh2): Should probably refactor this into a more portable - // description, since it does generalize for channels-last - ChannelsLastProblem problem = dimension_problem(input); - - auto output = at::empty_like(input); - - launch_opt_bias_add((__half*)output.data_ptr(), - (const __half*)input.data_ptr(), - (const __half*)bias.data_ptr(), - nullptr, - nullptr, - problem.batch_size, - problem.seq_len, - problem.channels, - at::cuda::getCurrentCUDAStream()); - - return output; -} - -at::Tensor seq_bias_add_add(at::Tensor& input, at::Tensor& bias, at::Tensor& other) -{ - assert(input.dtype() == at::kHalf); - - // TODO(cmikeh2): Should probably refactor this into a more portable - // description, since it does generalize for channels-last - ChannelsLastProblem problem = dimension_problem(input); - - auto output = at::empty_like(input); - - launch_opt_bias_add((__half*)output.data_ptr(), - (const __half*)input.data_ptr(), - (const __half*)bias.data_ptr(), - (const __half*)other.data_ptr(), - nullptr, - problem.batch_size, - problem.seq_len, - problem.channels, - at::cuda::getCurrentCUDAStream()); - - return output; -} - -at::Tensor seq_bias_add_bias_add(at::Tensor& input, - at::Tensor& bias, - at::Tensor& other, - at::Tensor& other_bias) -{ - assert(input.dtype() == at::kHalf); - - // TODO(cmikeh2): Should probably refactor this into a more portable - // description, since it does generalize for channels-last - ChannelsLastProblem problem = dimension_problem(input); - - auto output = at::empty_like(input); - - launch_opt_bias_add((__half*)output.data_ptr(), - (const __half*)input.data_ptr(), - (const __half*)bias.data_ptr(), - (const __half*)other.data_ptr(), - (const __half*)other_bias.data_ptr(), - problem.batch_size, - problem.seq_len, - problem.channels, - at::cuda::getCurrentCUDAStream()); - - return output; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("nhwc_bias_add", &seq_unroll_bias_add); - m.def("nhwc_bias_add_add", &seq_bias_add_add); - m.def("nhwc_bias_add_bias_add", &seq_bias_add_bias_add); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/includes/spatial_cuda_layers.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/includes/spatial_cuda_layers.h deleted file mode 100644 index 4f56f89f16cd1a200f057ebd651575224be23b7b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/spatial/includes/spatial_cuda_layers.h +++ /dev/null @@ -1,37 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#if __CUDA_ARCH__ >= 530 -#define HALF_PRECISION_AVAILABLE = 1 -#endif - -#ifdef __HIP_PLATFORM_AMD__ -#include -#else -#include -#endif - -#include -#include - -/*********** Group Norm Kernels, Structs, and Helpers ************/ - -struct { - int64_t batch_size; - int64_t seq_len; - int64_t channels; -} typedef ChannelsLastProblem; - -void launch_opt_bias_add(__half* result, - const __half* activation, - const __half* bias, - const __half* other, - const __half* other_bias, - int batch_size, - int seq_len, - int channels, - cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/cublas_wrappers.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/cublas_wrappers.cu deleted file mode 100644 index d982e65b8a81f786b4df550da310725d32ccc826..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/cublas_wrappers.cu +++ /dev/null @@ -1,493 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "cublas_wrappers.h" - -// TODO HIP: Remove backward compatibility for torch<=2.0 in future -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) -int cublas_gemm_ex(rocblas_handle handle, - rocblas_operation transa, - rocblas_operation transb, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const float* A, - const float* B, - float* C, - rocblas_gemm_algo algo) -#else -int cublas_gemm_ex(cublasHandle_t handle, - cublasOperation_t transa, - cublasOperation_t transb, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const float* A, - const float* B, - float* C, - cublasGemmAlgo_t algo) -#endif -{ -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_status status = rocblas_gemm_ex(handle, - transa, - transb, - m, - n, - k, - (const void*)alpha, - (const void*)A, - rocblas_datatype_f32_r, - (transa == rocblas_operation_none) ? m : k, - (const void*)B, - rocblas_datatype_f32_r, - (transb == rocblas_operation_none) ? k : n, - (const void*)beta, - C, - rocblas_datatype_f32_r, - m, - C, - rocblas_datatype_f32_r, - m, - rocblas_datatype_f32_r, - algo, - 0, - 0); -#else - cublasStatus_t status = cublasGemmEx(handle, - transa, - transb, - m, - n, - k, - (const void*)alpha, - (const void*)A, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - (transa == CUBLAS_OP_N) ? m : k, - (const void*)B, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - (transb == CUBLAS_OP_N) ? k : n, - (const void*)beta, - C, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - m, -#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) - HIPBLAS_COMPUTE_32F, -#elif defined(__HIP_PLATFORM_AMD__) - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - algo); -#endif - -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - if (status != rocblas_status_success) { -#else - if (status != CUBLAS_STATUS_SUCCESS) { -#endif - fprintf(stderr, - "!!!! kernel execution error. (m: %d, n: %d, k: %d, error: %d) \n", - m, - n, - k, - (int)status); - return EXIT_FAILURE; - } - return 0; -} - -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) -int cublas_gemm_ex(rocblas_handle handle, - rocblas_operation transa, - rocblas_operation transb, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const __half* A, - const __half* B, - __half* C, - rocblas_gemm_algo algo) -#else -int cublas_gemm_ex(cublasHandle_t handle, - cublasOperation_t transa, - cublasOperation_t transb, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const __half* A, - const __half* B, - __half* C, - cublasGemmAlgo_t algo) -#endif -{ -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_status status = rocblas_gemm_ex(handle, - transa, - transb, - m, - n, - k, - (const void*)alpha, - (const void*)A, - rocblas_datatype_f16_r, - (transa == rocblas_operation_none) ? m : k, - (const void*)B, - rocblas_datatype_f16_r, - (transb == rocblas_operation_none) ? k : n, - (const void*)beta, - (void*)C, - rocblas_datatype_f16_r, - m, - (void*)C, - rocblas_datatype_f16_r, - m, - rocblas_datatype_f32_r, - algo, - 0, - 0); -#else - cublasStatus_t status = cublasGemmEx(handle, - transa, - transb, - m, - n, - k, - (const void*)alpha, - (const void*)A, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_16F, -#else - CUDA_R_16F, -#endif - (transa == CUBLAS_OP_N) ? m : k, - (const void*)B, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_16F, -#else - CUDA_R_16F, -#endif - (transb == CUBLAS_OP_N) ? k : n, - (const void*)beta, - (void*)C, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_16F, -#else - CUDA_R_16F, -#endif - m, -#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) - HIPBLAS_COMPUTE_32F, -#elif defined(__HIP_PLATFORM_AMD__) - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - algo); -#endif - -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - if (status != rocblas_status_success) { -#else - if (status != CUBLAS_STATUS_SUCCESS) { -#endif - fprintf(stderr, - "!!!! kernel execution error. (m: %d, n: %d, k: %d, error: %d) \n", - m, - n, - k, - (int)status); - return EXIT_FAILURE; - } - return 0; -} - -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) -int cublas_strided_batched_gemm(rocblas_handle handle, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const float* A, - const float* B, - float* C, - rocblas_operation op_A, - rocblas_operation op_B, - int stride_A, - int stride_B, - int stride_C, - int batch, - rocblas_gemm_algo algo) -#else -int cublas_strided_batched_gemm(cublasHandle_t handle, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const float* A, - const float* B, - float* C, - cublasOperation_t op_A, - cublasOperation_t op_B, - int stride_A, - int stride_B, - int stride_C, - int batch, - cublasGemmAlgo_t algo) -#endif -{ -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_status status = - rocblas_gemm_strided_batched_ex(handle, - op_A, - op_B, - m, - n, - k, - alpha, - A, - rocblas_datatype_f32_r, - (op_A == rocblas_operation_none) ? m : k, - stride_A, - B, - rocblas_datatype_f32_r, - (op_B == rocblas_operation_none) ? k : n, - stride_B, - beta, - C, - rocblas_datatype_f32_r, - m, - stride_C, - C, - rocblas_datatype_f32_r, - m, - stride_C, - batch, - rocblas_datatype_f32_r, - algo, - 0, - 0); -#else - cublasStatus_t status = cublasGemmStridedBatchedEx(handle, - op_A, - op_B, - m, - n, - k, - alpha, - A, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - (op_A == CUBLAS_OP_N) ? m : k, - stride_A, - B, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - (op_B == CUBLAS_OP_N) ? k : n, - stride_B, - beta, - C, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - m, - stride_C, - batch, -#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) - HIPBLAS_COMPUTE_32F, -#elif defined(__HIP_PLATFORM_AMD__) - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - algo); -#endif - -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - if (status != rocblas_status_success) { -#else - if (status != CUBLAS_STATUS_SUCCESS) { -#endif - fprintf(stderr, - "!!!! kernel execution error. (batch: %d, m: %d, n: %d, k: %d, error: %d) \n", - batch, - m, - n, - k, - (int)status); - return EXIT_FAILURE; - } - return 0; -} - -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) -int cublas_strided_batched_gemm(rocblas_handle handle, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const __half* A, - const __half* B, - __half* C, - rocblas_operation op_A, - rocblas_operation op_B, - int stride_A, - int stride_B, - int stride_C, - int batch, - rocblas_gemm_algo algo) -#else -int cublas_strided_batched_gemm(cublasHandle_t handle, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const __half* A, - const __half* B, - __half* C, - cublasOperation_t op_A, - cublasOperation_t op_B, - int stride_A, - int stride_B, - int stride_C, - int batch, - cublasGemmAlgo_t algo) -#endif -{ -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_status status = - rocblas_gemm_strided_batched_ex(handle, - op_A, - op_B, - m, - n, - k, - alpha, - A, - rocblas_datatype_f16_r, - (op_A == rocblas_operation_none) ? m : k, - stride_A, - B, - rocblas_datatype_f16_r, - (op_B == rocblas_operation_none) ? k : n, - stride_B, - beta, - C, - rocblas_datatype_f16_r, - m, - stride_C, - C, - rocblas_datatype_f16_r, - m, - stride_C, - batch, - rocblas_datatype_f32_r, - algo, - 0, - 0); -#else - cublasStatus_t status = cublasGemmStridedBatchedEx(handle, - op_A, - op_B, - m, - n, - k, - alpha, - A, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_16F, -#else - CUDA_R_16F, -#endif - (op_A == CUBLAS_OP_N) ? m : k, - stride_A, - B, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_16F, -#else - CUDA_R_16F, -#endif - (op_B == CUBLAS_OP_N) ? k : n, - stride_B, - beta, - C, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_16F, -#else - CUDA_R_16F, -#endif - m, - stride_C, - batch, -#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) - HIPBLAS_COMPUTE_32F, -#elif defined(__HIP_PLATFORM_AMD__) - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - algo); -#endif - -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - if (status != rocblas_status_success) { -#else - if (status != CUBLAS_STATUS_SUCCESS) { -#endif - fprintf(stderr, - "!!!! kernel execution error. (m: %d, n: %d, k: %d, error: %d) \n", - m, - n, - k, - (int)status); - return EXIT_FAILURE; - } - - return 0; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/dropout_kernels.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/dropout_kernels.cu deleted file mode 100644 index 97cb3592dcdc4cf0a2dc37e34fc9b13117d9bcc5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/dropout_kernels.cu +++ /dev/null @@ -1,873 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "custom_cuda_layers.h" - -const int unroll_factor = 4; - -__global__ void dropout_kernel(const int N, - const float ratio, - float* out, - const float* Xdata, - uint8_t* mask, - std::pair seed) -{ - const float scale = 1. / (1. - ratio); - int idx = blockIdx.x * blockDim.x + threadIdx.x; - - curandStatePhilox4_32_10_t state; - curand_init(seed.first, idx, seed.second, &state); - - CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) - { - float4 rand = curand_uniform4(&state); - uint8_t m[unroll_factor]; - - m[0] = (uint8_t)(rand.x > ratio); - m[1] = (uint8_t)(rand.y > ratio); - m[2] = (uint8_t)(rand.z > ratio); - m[3] = (uint8_t)(rand.w > ratio); - - int i = j * unroll_factor; - - mask[i] = (uint8_t)m[0]; - mask[i + 1] = (uint8_t)m[1]; - mask[i + 2] = (uint8_t)m[2]; - mask[i + 3] = (uint8_t)m[3]; - - out[i] = Xdata[i] * scale * m[0]; - out[i + 1] = Xdata[i + 1] * scale * m[1]; - out[i + 2] = Xdata[i + 2] * scale * m[2]; - out[i + 3] = Xdata[i + 3] * scale * m[3]; - } - int high_index = - ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; - if (N > high_index) { - float4 rand = curand_uniform4(&state); - float* rand_data = &(rand.x); - int k = 0; - for (int i = high_index; i < N; i++) { - uint8_t m = (uint8_t)(rand_data[k++] > ratio); - out[i] = Xdata[i] * scale * m; - mask[i] = m; - } - } -} - -__global__ void dropout_kernel(const int N, - const float ratio, - __half* out, - const __half* Xdata, - uint8_t* mask, - std::pair seed) -{ - const float scale = 1. / (1. - ratio); - - int idx = blockIdx.x * blockDim.x + threadIdx.x; - - curandStatePhilox4_32_10_t state; - curand_init(seed.first, idx, seed.second, &state); - -#ifdef __STOCHASTIC_MODE__ - - const __half2 h_scale = __float2half2_rn(scale); - const float2* x_cast = reinterpret_cast(Xdata); - float2* out_cast = reinterpret_cast(out); - uint32_t* mask_cast = reinterpret_cast(mask); - - uint32_t m_32; - uint8_t* m = reinterpret_cast(&m_32); - - float2 result_f; - __half2* result_h = reinterpret_cast<__half2*>(&result_f); - __half2 mask_h[2]; - float2 mask_f[2]; - - CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) - { - float2 x_f = x_cast[j]; - __half2* x_h = reinterpret_cast<__half2*>(&x_f); - - float4 rand = curand_uniform4(&state); - - m[0] = (uint8_t)(rand.x > ratio); - m[1] = (uint8_t)(rand.y > ratio); - m[2] = (uint8_t)(rand.z > ratio); - m[3] = (uint8_t)(rand.w > ratio); - - float* mask_f_data = &mask_f[0].x; -#pragma unroll - for (int i = 0; i < unroll_factor; i++) mask_f_data[i] = (float)(m[i]); - - mask_h[0] = __float22half2_rn(mask_f[0]); - mask_h[1] = __float22half2_rn(mask_f[1]); - - result_h[0] = x_h[0] * h_scale * mask_h[0]; - result_h[1] = x_h[1] * h_scale * mask_h[1]; - - out_cast[j] = result_f; - - mask_cast[j] = m_32; - } - -#else - - CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) - { - int i = j * unroll_factor; - - const __half2* vals_half = reinterpret_cast(Xdata + i); - float2 vals_half_f[2]; - vals_half_f[0] = __half22float2(vals_half[0]); - vals_half_f[1] = __half22float2(vals_half[1]); - - uint8_t m[unroll_factor]; - float4 rand = curand_uniform4(&state); - m[0] = (uint8_t)(rand.x > ratio); - m[1] = (uint8_t)(rand.y > ratio); - m[2] = (uint8_t)(rand.z > ratio); - m[3] = (uint8_t)(rand.w > ratio); - - out[i] = __float2half(vals_half_f[0].x * scale * m[0]); - out[i + 1] = __float2half(vals_half_f[0].y * scale * m[1]); - out[i + 2] = __float2half(vals_half_f[1].x * scale * m[2]); - out[i + 3] = __float2half(vals_half_f[1].y * scale * m[3]); - - mask[i] = m[0]; - mask[i + 1] = m[1]; - mask[i + 2] = m[2]; - mask[i + 3] = m[3]; - } - -#endif - int high_index = - ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; - if (N > high_index) { - float4 rand = curand_uniform4(&state); - float* rand_data = &(rand.x); - int k = 0; - for (int i = high_index; i < N; i++) { - uint8_t m = (uint8_t)(rand_data[k++] > ratio); - out[i] = __float2half((float)Xdata[i] * scale * m); - mask[i] = m; - } - } -} - -__global__ void dropout_kernel_bwd(const int N, - const float ratio, - const float* Xdata, - float* out, - uint8_t* mask, - std::pair seed) -{ - const float scale = 1. / (1. - ratio); - CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) - { - int i = j * unroll_factor; - - out[i] = mask[i] ? Xdata[i] * scale : 0.0; - out[i + 1] = mask[i + 1] ? Xdata[i + 1] * scale : 0.0; - out[i + 2] = mask[i + 2] ? Xdata[i + 2] * scale : 0.0; - out[i + 3] = mask[i + 3] ? Xdata[i + 3] * scale : 0.0; - } - int high_index = - ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; - if (N > high_index) { - for (int i = high_index; i < N; i++) { out[i] = mask[i] ? Xdata[i] * scale : 0.0; } - } -} - -__global__ void dropout_kernel_bwd(const int N, - const float ratio, - const __half* Xdata, - __half* out, - uint8_t* mask, - std::pair seed) -{ - const float scale = 1. / (1. - ratio); - -#ifdef __STOCHASTIC_MODE__ - - const __half2 h_scale = __float2half2_rn(scale); - - const float2* x_cast = reinterpret_cast(Xdata); - float2* out_cast = reinterpret_cast(out); - uint32_t* mask_cast = reinterpret_cast(mask); - - CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) - { - float2 x_f = x_cast[j]; - __half2* x_h = reinterpret_cast<__half2*>(&x_f); - - uint32_t m_32 = mask_cast[j]; - uint8_t* m = (uint8_t*)&m_32; - - __half2 mask_h[2]; - float2 mask_f[2]; - - float* mask_f_data = &mask_f[0].x; -#pragma unroll - for (int i = 0; i < unroll_factor; i++) mask_f_data[i] = (float)(m[i]); - -#pragma unroll - for (int i = 0; i < 2; i++) mask_h[i] = __float22half2_rn(mask_f[i]); - - float2 result_f; - __half2* result_h = reinterpret_cast<__half2*>(&result_f); - - result_h[0] = x_h[0] * h_scale * mask_h[0]; - result_h[1] = x_h[1] * h_scale * mask_h[1]; - - out_cast[j] = result_f; - } - -#else - - const __half h_scale = __float2half(scale); - const __half h_zero = __float2half(0.0); - - CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) - { - int i = j * unroll_factor; - - const __half2* vals_half = reinterpret_cast(Xdata + i); - - uint8_t* m = mask + i; - - float2 vals_half_f[2]; - - vals_half_f[0] = __half22float2(vals_half[0]); - vals_half_f[1] = __half22float2(vals_half[1]); - - out[i] = __float2half(vals_half_f[0].x * scale * m[0]); - out[i + 1] = __float2half(vals_half_f[0].y * scale * m[1]); - out[i + 2] = __float2half(vals_half_f[1].x * scale * m[2]); - out[i + 3] = __float2half(vals_half_f[1].y * scale * m[3]); - } - -#endif - int high_index = - ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; - if (N > high_index) { - for (int i = high_index; i < N; i++) { - out[i] = __float2half((float)Xdata[i] * scale * mask[i]); - } - } -} - -template -void launch_dropout(T* out, - const T* vals, - uint8_t* mask, - int total_count, - int dim, - float ratio, - cudaStream_t stream, - bool bwd) -{ - assert(unroll_factor == 4); - - dim3 grid_dim = DS_GET_BLOCKS(total_count / unroll_factor); - dim3 block_dim = DS_CUDA_NUM_THREADS; - - if (dim > 512) { - block_dim.x >>= 1; - grid_dim.x <<= 1; - } - uint64_t inc = total_count / grid_dim.x / block_dim.x; - std::pair seed = TrainingContext::Instance().IncrementOffset(inc); - if (bwd) - dropout_kernel_bwd<<>>( - total_count, ratio, vals, out, mask, seed); - else - dropout_kernel<<>>( - total_count, ratio, out, vals, mask, seed); -} - -template void launch_dropout(float* out, - const float* vals, - uint8_t* mask, - int total_count, - int dim, - float ratio, - cudaStream_t stream, - bool); -template void launch_dropout(__half* out, - const __half* vals, - uint8_t* mask, - int total_count, - int dim, - float ratio, - cudaStream_t stream, - bool); - -__global__ void dropout_grad_kernel(const int N, const float scale, float* Xdata, uint8_t* mask) -{ - CUDA_1D_KERNEL_LOOP(i, N) { Xdata[i] *= scale * mask[i]; } -} - -__global__ void dropout_grad_kernel(const int N, const float scale, __half* Xdata, uint8_t* mask) -{ - const __half2 h_scale = __float2half2_rn(scale); - float2* x_cast = reinterpret_cast(Xdata); - uint32_t* mask_cast = reinterpret_cast(mask); - - CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) - { - float2 x_data = x_cast[j]; - uint32_t m_32 = mask_cast[j]; - uint8_t* m = (uint8_t*)&m_32; - - float2 result_f; - __half2* result_h = reinterpret_cast<__half2*>(&result_f); - -#ifdef __STOCHASTIC_MODE__ - - __half2* x_data_h = reinterpret_cast<__half2*>(&x_data); - __half2 mask_h[2]; - float2 mask_f[2]; - - float* mask_f_data = &mask_f[0].x; -#pragma unroll - for (int i = 0; i < unroll_factor; i++) *(mask_f_data++) = (float)(m[i]); - - mask_h[0] = __float22half2_rn(mask_f[0]); - mask_h[1] = __float22half2_rn(mask_f[1]); - - result_h[0] = x_data_h[0] * h_scale * mask_h[0]; - result_h[1] = x_data_h[1] * h_scale * mask_h[1]; - -#else - - __half* x_data_h = reinterpret_cast<__half*>(&x_data); - float2 result[2]; - - result[0].x = (float)x_data_h[0] * scale * m[0]; - result[0].y = (float)x_data_h[1] * scale * m[1]; - result[1].x = (float)x_data_h[2] * scale * m[2]; - result[1].y = (float)x_data_h[3] * scale * m[3]; - - result_h[0] = __float22half2_rn(result[0]); - result_h[1] = __float22half2_rn(result[1]); - -#endif - x_cast[j] = result_f; - } - int high_index = - ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; - if (N > high_index) { - for (int i = high_index; i < N; i++) { - Xdata[i] = __float2half((float)Xdata[i] * scale * mask[i]); - } - } -} - -template -void launch_dropout_grad(T* vals, uint8_t* mask, int total_count, float ratio, cudaStream_t stream) -{ - assert(unroll_factor == 4); - - const float scale = 1. / (1. - ratio); - dropout_grad_kernel<<>>(total_count, scale, vals, mask); -} - -template void launch_dropout_grad(float* vals, - uint8_t* mask, - int total_count, - float ratio, - cudaStream_t stream); -template void launch_dropout_grad(__half* vals, - uint8_t* mask, - int total_count, - float ratio, - cudaStream_t stream); - -__global__ void dropout_grad_kernel(const int N, - const float scale, - const float* Xdata, - float* out, - uint8_t* mask) -{ - CUDA_1D_KERNEL_LOOP(i, N) { out[i] = Xdata[i] * scale * mask[i]; } -} - -__global__ void dropout_grad_kernel(const int N, - const float scale, - const __half* Xdata, - __half* out, - uint8_t* mask) -{ - const float2* x_cast = reinterpret_cast(Xdata); - float2* out_cast = reinterpret_cast(out); - const uint32_t* mask_cast = reinterpret_cast(mask); - - float2 result_f; - __half2* result_h = reinterpret_cast<__half2*>(&result_f); - - CUDA_1D_KERNEL_LOOP(j, N / unroll_factor) - { - float2 x_data = x_cast[j]; - uint32_t m_32 = mask_cast[j]; - uint8_t* m = (uint8_t*)&m_32; - - __half* x_data_h = reinterpret_cast<__half*>(&x_data); - float2 result[2]; - - result[0].x = (float)x_data_h[0] * scale * m[0]; - result[0].y = (float)x_data_h[1] * scale * m[1]; - result[1].x = (float)x_data_h[2] * scale * m[2]; - result[1].y = (float)x_data_h[3] * scale * m[3]; - - result_h[0] = __float22half2_rn(result[0]); - result_h[1] = __float22half2_rn(result[1]); - - out_cast[j] = result_f; - } - int high_index = - ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; - if (N > high_index) { - for (int i = high_index; i < N; i++) { - out[i] = __float2half((float)Xdata[i] * scale * mask[i]); - } - } -} - -template -void launch_dropout_grad(T* vals_out, - const T* vals, - uint8_t* mask, - int total_count, - float ratio, - cudaStream_t stream) -{ - assert(unroll_factor == 4); - - const float scale = 1. / (1. - ratio); - dropout_grad_kernel<<>>(total_count, scale, vals, vals_out, mask); -} -template void launch_dropout_grad(float*, - const float* vals, - uint8_t* mask, - int total_count, - float ratio, - cudaStream_t stream); -template void launch_dropout_grad(__half*, - const __half* vals, - uint8_t* mask, - int total_count, - float ratio, - cudaStream_t stream); - -__global__ void dropout_kernel(const int N, - const int dim, - const float ratio, - const float* bias, - float* Xdata, - uint8_t* mask, - std::pair seed) -{ - const float scale = 1. / (1. - ratio); - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int tid = threadIdx.x % (dim / unroll_factor); - - curandStatePhilox4_32_10_t state; - curand_init(seed.first, idx, seed.second, &state); - - float4* Xdata_cast = reinterpret_cast(Xdata); - uint32_t* mask_32 = reinterpret_cast(mask); - const float4* bias_cast = reinterpret_cast(bias); - - CUDA_1D_KERNEL_LOOP(j, N) - { - float4 rand = curand_uniform4(&state); - uint32_t m_32; - uint8_t* m = (uint8_t*)&m_32; - - m[0] = (uint8_t)(rand.x > ratio); - m[1] = (uint8_t)(rand.y > ratio); - m[2] = (uint8_t)(rand.z > ratio); - m[3] = (uint8_t)(rand.w > ratio); - - float4 x_data = Xdata_cast[j]; - float4 b_data = bias_cast[j % (dim / unroll_factor)]; - - x_data.x += b_data.x; - x_data.y += b_data.y; - x_data.z += b_data.z; - x_data.w += b_data.w; - - x_data.x = x_data.x * scale * m[0]; - x_data.y = x_data.y * scale * m[1]; - x_data.z = x_data.z * scale * m[2]; - x_data.w = x_data.w * scale * m[3]; - - mask_32[j] = m_32; - Xdata_cast[j] = x_data; - } - int high_index = - ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; - if (N > high_index) { - float4 rand = curand_uniform4(&state); - float* rand_data = &(rand.x); - int k = 0; - for (int i = high_index; i < N; i++) { - float x_data = Xdata[i] + bias[i % dim]; - uint8_t m = (uint8_t)(rand_data[k++] > ratio); - Xdata[i] = x_data * scale * m; - mask[i] = m; - } - } -} - -__global__ void dropout_kernel(const int N, - const int dim, - const float ratio, - const __half* bias, - __half* Xdata, - uint8_t* mask, - std::pair seed) -{ - const float scale = 1. / (1. - ratio); - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int tid = threadIdx.x % (dim / unroll_factor); - - curandStatePhilox4_32_10_t state; - curand_init(seed.first, idx, seed.second, &state); - - float2* Xdata_cast = reinterpret_cast(Xdata); - uint32_t* mask_32 = reinterpret_cast(mask); - const float2* bias_cast = reinterpret_cast(bias); - - CUDA_1D_KERNEL_LOOP(j, N) - { - float4 rand = curand_uniform4(&state); - - float2 data_f; - __half2* data_h = reinterpret_cast<__half2*>(&data_f); - - float2 bias_f; - __half2* bias_h = reinterpret_cast<__half2*>(&bias_f); - - data_f = Xdata_cast[j]; - bias_f = bias_cast[j % (dim / unroll_factor)]; - - float2 data_h_0 = __half22float2(data_h[0]); - float2 data_h_1 = __half22float2(data_h[1]); - - float2 bias_h_0 = __half22float2(bias_h[0]); - float2 bias_h_1 = __half22float2(bias_h[1]); - - data_h_0.x += bias_h_0.x; - data_h_0.y += bias_h_0.y; - data_h_1.x += bias_h_1.x; - data_h_1.y += bias_h_1.y; - - uint32_t m_32; - uint8_t* m = (uint8_t*)&m_32; - - m[0] = (uint8_t)(rand.x > ratio); - m[1] = (uint8_t)(rand.y > ratio); - m[2] = (uint8_t)(rand.z > ratio); - m[3] = (uint8_t)(rand.w > ratio); - - data_h_0.x = __float2half(data_h_0.x * scale * m[0]); - data_h_0.y = __float2half(data_h_0.y * scale * m[1]); - data_h_1.x = __float2half(data_h_1.x * scale * m[2]); - data_h_1.y = __float2half(data_h_1.y * scale * m[3]); - - float2 result_f; - __half2* result_h = reinterpret_cast<__half2*>(&result_f); - - result_h[0] = __float22half2_rn(data_h_0); - result_h[1] = __float22half2_rn(data_h_1); - - Xdata_cast[j] = result_f; - mask_32[j] = m_32; - } - int high_index = - ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; - if (N > high_index) { - float4 rand = curand_uniform4(&state); - float* rand_data = &(rand.x); - int k = 0; - for (int i = high_index; i < N; i++) { - float x_data = (float)Xdata[i] + (float)bias[i % dim]; - uint8_t m = (uint8_t)(rand_data[k++] > ratio); - Xdata[i] = __float2half(x_data * scale * m); - mask[i] = m; - } - } -} - -template -void launch_dropout(T* out, - const T* bias, - uint8_t* mask, - int batch, - int dim, - float ratio, - cudaStream_t stream) -{ - assert(unroll_factor == 4); - - int total_count = batch * dim / unroll_factor; - - dim3 grid_dim = DS_GET_BLOCKS(total_count); - dim3 block_dim = DS_CUDA_NUM_THREADS; - - uint64_t inc = (batch * dim) / grid_dim.x / block_dim.x; - std::pair seed = TrainingContext::Instance().IncrementOffset(inc); - - dropout_kernel<<>>( - total_count, dim, ratio, bias, out, mask, seed); -} - -template void launch_dropout(float*, - const float* bias, - uint8_t* mask, - int batch, - int dim, - float ratio, - cudaStream_t stream); -template void launch_dropout(__half*, - const __half* bias, - uint8_t* mask, - int batch, - int dim, - float ratio, - cudaStream_t stream); - -__global__ void dropout_kernel(const int N, - const int dim, - const float ratio, - const float* input, - const float* residual, - const float* bias, - float* out, - uint8_t* mask, - std::pair seed) -{ - const float scale = 1. / (1. - ratio); - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int tid = threadIdx.x % (dim / unroll_factor); - - curandStatePhilox4_32_10_t state; - curand_init(seed.first, idx, seed.second, &state); - - float4* out_cast = reinterpret_cast(out); - uint32_t* mask_32 = reinterpret_cast(mask); - - const float4* bias_cast = reinterpret_cast(bias); - const float4* residual_cast = reinterpret_cast(residual); - const float4* input_cast = reinterpret_cast(input); - - CUDA_1D_KERNEL_LOOP(j, N) - { - float4 rand = curand_uniform4(&state); - - uint32_t m_32; - uint8_t* m = (uint8_t*)&m_32; - - m[0] = (uint8_t)(rand.x > ratio); - m[1] = (uint8_t)(rand.y > ratio); - m[2] = (uint8_t)(rand.z > ratio); - m[3] = (uint8_t)(rand.w > ratio); - - float4 out_data; - float4 b_data = bias_cast[j % (dim / unroll_factor)]; - float4 res_data = residual_cast[j]; - float4 inp_data = input_cast[j]; - - out_data.x = (b_data.x + inp_data.x); - out_data.y = (b_data.y + inp_data.y); - out_data.z = (b_data.z + inp_data.z); - out_data.w = (b_data.w + inp_data.w); - - out_data.x = out_data.x * scale * m[0]; - out_data.y = out_data.y * scale * m[1]; - out_data.z = out_data.z * scale * m[2]; - out_data.w = out_data.w * scale * m[3]; - - out_data.x += res_data.x; - out_data.y += res_data.y; - out_data.z += res_data.z; - out_data.w += res_data.w; - - mask_32[j] = m_32; - out_cast[j] = out_data; - } - int high_index = - ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; - if (N > high_index) { - float4 rand = curand_uniform4(&state); - float* rand_data = &(rand.x); - int k = 0; - for (int i = high_index; i < N; i++) { - float x_data = input[i] + bias[i % dim]; - uint8_t m = (uint8_t)(rand_data[k++] > ratio); - x_data = x_data * scale * m; - x_data += residual[i]; - - out[i] = x_data; - mask[i] = m; - } - } -} - -__global__ void dropout_kernel(const int N, - const int dim, - const float ratio, - const __half* input, - const __half* residual, - const __half* bias, - __half* out, - uint8_t* mask, - std::pair seed) -{ - const float scale = 1. / (1. - ratio); - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int tid = threadIdx.x % (dim / unroll_factor); - - curandStatePhilox4_32_10_t state; - curand_init(seed.first, idx, seed.second, &state); - - float2* out_cast = reinterpret_cast(out); - uint32_t* mask_32 = reinterpret_cast(mask); - - const float2* bias_cast = reinterpret_cast(bias); - const float2* residual_cast = reinterpret_cast(residual); - const float2* input_cast = reinterpret_cast(input); - - CUDA_1D_KERNEL_LOOP(j, N) - { - float4 rand = curand_uniform4(&state); - - float2 data_f; - __half2* data_h = reinterpret_cast<__half2*>(&data_f); - - float2 bias_f; - __half2* bias_h = reinterpret_cast<__half2*>(&bias_f); - - float2 residual_f; - __half2* residual_h = reinterpret_cast<__half2*>(&residual_f); - - float2 input_f; - __half2* input_h = reinterpret_cast<__half2*>(&input_f); - - bias_f = bias_cast[j % (dim / unroll_factor)]; - residual_f = residual_cast[j]; - input_f = input_cast[j]; - - float2 data_h_0 = __half22float2(data_h[0]); - float2 data_h_1 = __half22float2(data_h[1]); - - float2 bias_h_0 = __half22float2(bias_h[0]); - float2 bias_h_1 = __half22float2(bias_h[1]); - - float2 residual_h_0 = __half22float2(residual_h[0]); - float2 residual_h_1 = __half22float2(residual_h[1]); - - float2 input_h_0 = __half22float2(input_h[0]); - float2 input_h_1 = __half22float2(input_h[1]); - - data_h_0.x = (bias_h_0.x + input_h_0.x); - data_h_0.y = (bias_h_0.y + input_h_0.y); - data_h_1.x = (bias_h_1.x + input_h_1.x); - data_h_1.y = (bias_h_1.y + input_h_1.y); - - uint32_t m_32; - uint8_t* m = (uint8_t*)&m_32; - - m[0] = (uint8_t)(rand.x > ratio); - m[1] = (uint8_t)(rand.y > ratio); - m[2] = (uint8_t)(rand.z > ratio); - m[3] = (uint8_t)(rand.w > ratio); - - data_h_0.x = __float2half(data_h_0.x * scale * m[0]); - data_h_0.y = __float2half(data_h_0.y * scale * m[1]); - data_h_1.x = __float2half(data_h_1.x * scale * m[2]); - data_h_1.y = __float2half(data_h_1.y * scale * m[3]); - - data_h_0.x += residual_h_0.x; - data_h_0.y += residual_h_0.y; - data_h_1.x += residual_h_1.x; - data_h_1.y += residual_h_1.y; - - float2 result_f; - __half2* result_h = reinterpret_cast<__half2*>(&result_f); - - result_h[0] = __float22half2_rn(data_h_0); - result_h[1] = __float22half2_rn(data_h_1); - - out_cast[j] = result_f; - mask_32[j] = m_32; - } - int high_index = - ((((N / unroll_factor) - 1) / blockDim.x + 1) * (unroll_factor * blockDim.x)) + threadIdx.x; - if (N > high_index) { - float4 rand = curand_uniform4(&state); - float* rand_data = &(rand.x); - int k = 0; - for (int i = high_index; i < N; i++) { - float x_data = (float)input[i] + (float)bias[i % dim]; - uint8_t m = (uint8_t)(rand_data[k++] > ratio); - x_data = x_data * scale * m; - x_data += (float)residual[i]; - - out[i] = __float2half(x_data); - mask[i] = m; - } - } -} - -template -void launch_dropout(T* out, - const T* input, - const T* residual, - const T* bias, - uint8_t* mask, - int batch, - int dim, - float ratio, - cudaStream_t stream) -{ - assert(unroll_factor == 4); - - int total_count = batch * dim / unroll_factor; - dim3 grid_dim = DS_GET_BLOCKS(total_count); - dim3 block_dim = DS_CUDA_NUM_THREADS; - - uint64_t inc = (batch * dim) / grid_dim.x / block_dim.x; - std::pair seed = TrainingContext::Instance().IncrementOffset(inc); - - dropout_kernel<<>>( - total_count, dim, ratio, input, residual, bias, out, mask, seed); -} - -template void launch_dropout(float*, - const float*, - const float* residual, - const float* bias, - uint8_t* mask, - int batch, - int dim, - float ratio, - cudaStream_t stream); -template void launch_dropout(__half*, - const __half*, - const __half* residual, - const __half* bias, - uint8_t* mask, - int batch, - int dim, - float ratio, - cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/ds_transformer_cuda.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/ds_transformer_cuda.cpp deleted file mode 100644 index b637bb710c67a4f980571cd876c4caa350f4a71b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/ds_transformer_cuda.cpp +++ /dev/null @@ -1,1055 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include - -#include -#include -#include -#include -#include -#include -#include "Timer.h" -#include "context.h" -#include "cublas_wrappers.h" -#include "custom_cuda_layers.h" -#include "ds_transformer_cuda.h" - -static std::unordered_map> s_transformer_layers; - -const int init_seq_length = 128; - -// C++ interface - -template -unsigned get_workspace_size(unsigned maxBatchSize, - unsigned seq_len, - unsigned hidden_size, - unsigned intermediate_size, - unsigned heads, - bool training, - bool gelu_checkpoint) -{ - unsigned workSpacesize = 4 * (size_t(maxBatchSize) * seq_len * hidden_size); - if (training) { - workSpacesize += 2 * (size_t(maxBatchSize) * seq_len * hidden_size); - workSpacesize += ((std::max)((size_t(maxBatchSize) * seq_len * intermediate_size), - 2 * (size_t(maxBatchSize) * heads * seq_len * seq_len))); - if (gelu_checkpoint) - workSpacesize += 2 * (size_t(maxBatchSize) * seq_len * intermediate_size); - } - return workSpacesize; // * sizeof(T); -} - -// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. -#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") -#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") -#define CHECK_INPUT(x) \ - CHECK_CUDA(x); \ - CHECK_CONTIGUOUS(x) - -template -BertTransformerLayer::BertTransformerLayer(unsigned layer_id, - unsigned batch_size, - unsigned hidden_size, - unsigned num_heads, - unsigned intermediate_size, - unsigned seq_length, - float attn_prob_dropout_ratio, - float hidden_output_dropout_ratio, - float layer_norm_eps, - bool pre_or_postLayerNorm, - const std::vector>& gemm_algos, - bool attn_dropout_checkpoint, - bool normalize_invertible, - bool gelu_checkpoint, - bool stochastic_mode) - : _layer_id(layer_id), - _batch_size(batch_size), - _hidden_size(hidden_size), - _heads(num_heads), - _intermediate_size(intermediate_size), - _seq_length(seq_length), - _training(true), - _pre_or_postLayerNorm(pre_or_postLayerNorm), - _attn_dropout_checkpoint(attn_dropout_checkpoint), - _normalize_invertible(normalize_invertible), - _gelu_checkpoint(gelu_checkpoint), - _stochastic_mode(stochastic_mode), - _stream(TrainingContext::Instance().GetCurrentStream()), - _cublasHandle(TrainingContext::Instance().GetCublasHandle()), - _qkv_linear(typename FeedForward::Config(batch_size * seq_length, - 3 * hidden_size, - hidden_size, - gemm_algos[0])), - _attn_out_linear(typename FeedForward::Config(batch_size * seq_length, - hidden_size, - hidden_size, - gemm_algos[0])), - _attn_layer_norm(typename Normalize_Layer::Config(batch_size, - seq_length, - hidden_size, - layer_norm_eps, - true, - !normalize_invertible)), - _layer_norm(typename Normalize_Layer::Config(batch_size, - seq_length, - hidden_size, - layer_norm_eps, - true, - !normalize_invertible)), - _ff1(typename FeedForward::Config(batch_size * seq_length, - _intermediate_size, - hidden_size, - gemm_algos[1])), - _ff2(typename FeedForward::Config(batch_size * seq_length, - hidden_size, - _intermediate_size, - gemm_algos[2])), - _softmax(typename Softmax::Config(batch_size, num_heads, seq_length)), - _gelu(typename Gelu::Config(_intermediate_size)), - _attn_prob_dropout(typename Dropout::Config(attn_prob_dropout_ratio, _seq_length)), - _attn_output_dropout(typename Dropout::Config(hidden_output_dropout_ratio, _hidden_size)), - _layer_output_dropout(typename Dropout::Config(hidden_output_dropout_ratio, _hidden_size)), - _attn_scores(typename StridedBatchGemm::Config(_batch_size * _heads, - _seq_length, - _seq_length, - _hidden_size / _heads, - (T(1.0) / T(sqrt(_hidden_size / _heads))), - T(0.0), - CUBLAS_OP_T, - CUBLAS_OP_N, - gemm_algos[3])), - _attn_context(typename StridedBatchGemm::Config(_batch_size * _heads, - _hidden_size / _heads, - _seq_length, - _seq_length, - T(1.0), - T(0.0), - CUBLAS_OP_N, - CUBLAS_OP_N, - gemm_algos[4])) -{ - assert(_hidden_size % _heads == 0); - - Initialize(); -} - -template -BertTransformerLayer::~BertTransformerLayer() -{ -} - -template -void BertTransformerLayer::Initialize() -{ -#ifndef __HIP_PLATFORM_AMD__ - if (std::is_same::value) cublasSetMathMode(_cublasHandle, CUBLAS_TENSOR_OP_MATH); -#endif -} - -template -void BertTransformerLayer::Forward(unsigned bsz, - const T* input_ptr, - const T* input_mask_ptr, - const T* attn_qkvw_ptr, - const T* attn_qkvb_ptr, - const T* attn_ow_ptr, - const T* attn_ob_ptr, - const T* attn_nw_ptr, - const T* attn_nb_ptr, - const T* inter_w_ptr, - const T* inter_b_ptr, - const T* output_w_ptr, - const T* output_b_ptr, - const T* norm_w_ptr, - const T* norm_b_ptr, - T* out_ptr, - T* inp_norm_ptr, - T* q_tf_ptr, - T* k_tf_ptr, - T* v_tf_ptr, - T* soft_out_ptr, - T* ctx_bufB_ptr, - T* attn_o_inp_ptr, - T* add_res_ptr, - T* ff1_inp_ptr, - T* gelu_inp_ptr, - T* ff2_inp_ptr) -{ - cublasSetStream(_cublasHandle, _stream); - - if (!_stochastic_mode) cudaStreamSynchronize(_stream); - - T* workspace = static_cast(TrainingContext::Instance().GetWorkSpace()); - size_t small_buf_size = bsz * _seq_length * _hidden_size; - T* buf_0 = workspace; - T* buf_1 = buf_0 + small_buf_size; - T* buf_2 = buf_1; - - if (_normalize_invertible) { - add_res_ptr = buf_1 + 3 * small_buf_size; - buf_2 = add_res_ptr; - } - if (_gelu_checkpoint) buf_2 += small_buf_size; - if (_attn_dropout_checkpoint) - ctx_bufB_ptr = - (_gelu_checkpoint ? (buf_2 + (_intermediate_size / _hidden_size) * small_buf_size) - : (buf_1 + 4 * small_buf_size)); - - int bsz_seq = bsz * _seq_length; - - if (_pre_or_postLayerNorm) { - if (_layer_norm.UseMean()) - _layer_norm.ForwardCheckpoint( - bsz_seq, inp_norm_ptr, input_ptr, norm_w_ptr, norm_b_ptr, _stream, true); - - else - _layer_norm.Forward( - bsz_seq, inp_norm_ptr, input_ptr, norm_w_ptr, norm_b_ptr, _stream, true); - } - - if (_pre_or_postLayerNorm) - _qkv_linear.Forward(bsz_seq, inp_norm_ptr, attn_qkvw_ptr, buf_0, _cublasHandle); - else - _qkv_linear.Forward(bsz_seq, input_ptr, attn_qkvw_ptr, buf_0, _cublasHandle); - - launch_bias_add_transform_0213( - q_tf_ptr, buf_0, attn_qkvb_ptr, bsz, _seq_length, _hidden_size, _heads, _stream, 3); - - int bsz_heads = bsz * _heads; - - // attention scores - _attn_scores.Forward(bsz_heads, soft_out_ptr, k_tf_ptr, q_tf_ptr, _cublasHandle); - - // Softmax + Mask - _softmax.Forward(bsz, soft_out_ptr, input_mask_ptr, _stream); - - // attn prob dropout. - _attn_prob_dropout.Forward(bsz_heads * _seq_length, ctx_bufB_ptr, soft_out_ptr, _stream); - - // attention context - _attn_context.Forward(bsz_heads, buf_1, v_tf_ptr, ctx_bufB_ptr, _cublasHandle); - - launch_transform4d_0213( - attn_o_inp_ptr, buf_1, bsz, _heads, _seq_length, _hidden_size, _stream, 1); - - if (_pre_or_postLayerNorm) - _attn_out_linear.Forward(bsz_seq, attn_o_inp_ptr, attn_ow_ptr, buf_1, _cublasHandle); - else - _attn_out_linear.Forward(bsz_seq, attn_o_inp_ptr, attn_ow_ptr, ff1_inp_ptr, _cublasHandle); - - // attn output dropout. - if (_pre_or_postLayerNorm) - _attn_output_dropout.ForwardWithBias( - bsz_seq, add_res_ptr, buf_1, input_ptr, attn_ob_ptr, _stream); - else - _attn_output_dropout.ForwardWithBias( - bsz_seq, add_res_ptr, ff1_inp_ptr, input_ptr, attn_ob_ptr, _stream); - - if (_pre_or_postLayerNorm) { - if (_attn_layer_norm.UseMean()) - _attn_layer_norm.ForwardCheckpoint( - bsz_seq, ff1_inp_ptr, add_res_ptr, attn_nw_ptr, attn_nb_ptr, _stream, true); - else - _attn_layer_norm.Forward( - bsz_seq, ff1_inp_ptr, add_res_ptr, attn_nw_ptr, attn_nb_ptr, _stream, true); - } else { - if (_attn_layer_norm.UseMean()) - _attn_layer_norm.ForwardCheckpoint( - bsz_seq, ff1_inp_ptr, add_res_ptr, attn_nw_ptr, attn_nb_ptr, _stream, true); - else - _attn_layer_norm.Forward( - bsz_seq, ff1_inp_ptr, add_res_ptr, attn_nw_ptr, attn_nb_ptr, _stream, true); - } - - _ff1.Forward(bsz_seq, - ff1_inp_ptr, - inter_w_ptr, - (_gelu_checkpoint ? ff2_inp_ptr : gelu_inp_ptr), - _cublasHandle); - - _gelu.ForwardWithBiasAdd(bsz_seq, - (_gelu_checkpoint ? ff2_inp_ptr : gelu_inp_ptr), - inter_b_ptr, - (_gelu_checkpoint ? buf_2 : ff2_inp_ptr), - _stream); - - _ff2.Forward( - bsz_seq, (_gelu_checkpoint ? buf_2 : ff2_inp_ptr), output_w_ptr, out_ptr, _cublasHandle); - - // layer output dropout. - if (_pre_or_postLayerNorm) - _layer_output_dropout.ForwardWithBias( - bsz_seq, out_ptr, out_ptr, add_res_ptr, output_b_ptr, _stream); - else - _layer_output_dropout.ForwardWithBias( - bsz_seq, inp_norm_ptr, out_ptr, ff1_inp_ptr, output_b_ptr, _stream); - - if (!_pre_or_postLayerNorm) { - if (_layer_norm.UseMean()) - _layer_norm.ForwardCheckpoint( - bsz_seq, out_ptr, inp_norm_ptr, norm_w_ptr, norm_b_ptr, _stream, true); - else - _layer_norm.Forward( - bsz_seq, out_ptr, inp_norm_ptr, norm_w_ptr, norm_b_ptr, _stream, true); - } -} - -template -void BertTransformerLayer::Backward(unsigned bsz, - const T* grad_output_ptr, - const T* input_ptr, - const T* output_ptr, - const T* inp_norm_ptr, - const T* q_tf_ptr, - const T* k_tf_ptr, - const T* v_tf_ptr, - const T* soft_out_ptr, - const T* ctx_bufB_ptr, - const T* attn_o_inp_ptr, - const T* add_res_ptr, - const T* ff1_inp_ptr, - const T* gelu_inp_ptr, - const T* ff2_inp_ptr, - const T* input_mask_ptr, - const T* attn_qkvw_ptr, - const T* attn_ow_ptr, - const T* attn_nw_ptr, - const T* attn_nb_ptr, - const T* inter_w_ptr, - const T* inter_b_ptr, - const T* output_w_ptr, - const T* norm_w_ptr, - const T* norm_b_ptr, - - T* grad_input_ptr, - T* grad_attn_qkvw_ptr, - T* grad_attn_qkvb_ptr, - T* grad_attn_ow_ptr, - T* grad_attn_ob_ptr, - T* grad_attn_nw_ptr, - T* grad_attn_nb_ptr, - T* grad_inter_w_ptr, - T* grad_inter_b_ptr, - T* grad_output_w_ptr, - T* grad_output_b_ptr, - T* grad_norm_w_ptr, - T* grad_norm_b_ptr) -{ - cublasSetStream(_cublasHandle, _stream); - - if (!_stochastic_mode) cudaStreamSynchronize(_stream); - - T* workspace = static_cast(TrainingContext::Instance().GetWorkSpace()); - size_t small_buf_size = bsz * _seq_length * _hidden_size; - T* buf_0 = workspace; - T* buf_1 = buf_0 + small_buf_size; - T* buf_2 = buf_1 + small_buf_size; - T* buf_3 = buf_2 + small_buf_size; - - T* ff2_buf = (_gelu_checkpoint ? buf_3 + (bsz * _seq_length * _intermediate_size) - : buf_3 + small_buf_size); - T* ctx_bufB_ptr_recomp = ff2_buf + (_seq_length * _seq_length * bsz * _heads); - - cudaStream_t streams[2] = {_stream, _stream}; - - int bsz_seq = bsz * _seq_length; - int bsz_heads = bsz * _heads; - - if (!_pre_or_postLayerNorm) { - if (_layer_norm.UseMean()) - _layer_norm.Backward(bsz_seq, - grad_output_ptr, - norm_w_ptr, - grad_norm_w_ptr, - grad_norm_b_ptr, - streams, - buf_1, - inp_norm_ptr); - - else - _layer_norm.Backward(bsz_seq, - grad_output_ptr, - norm_w_ptr, - norm_b_ptr, - grad_norm_w_ptr, - grad_norm_b_ptr, - streams, - buf_1, - output_ptr); - } - - if (_pre_or_postLayerNorm) - _layer_output_dropout.Backward(bsz_seq, buf_0, grad_output_ptr, _stream); - else - _layer_output_dropout.Backward(bsz_seq, buf_0, buf_1, _stream); - - const T* layer_dropout_buf = _layer_output_dropout.HasDropout() - ? buf_0 - : (_pre_or_postLayerNorm ? grad_output_ptr : buf_1); - - if (_gelu_checkpoint) - _gelu.ForwardWithBiasAdd(bsz_seq, ff2_inp_ptr, inter_b_ptr, buf_2, _stream); - _ff2.Backward(bsz_seq, - layer_dropout_buf, - (_gelu_checkpoint ? buf_2 : ff2_inp_ptr), - output_w_ptr, - grad_output_w_ptr, - grad_output_b_ptr, - _cublasHandle, - _stream, - ff2_buf); - - _gelu.Backward( - bsz_seq, ff2_buf, (_gelu_checkpoint ? ff2_inp_ptr : gelu_inp_ptr), inter_b_ptr, _stream); - - _ff1.Backward(bsz_seq, - ff2_buf, - ff1_inp_ptr, - inter_w_ptr, - grad_inter_w_ptr, - grad_inter_b_ptr, - _cublasHandle, - _stream, - buf_3); - - if (!_pre_or_postLayerNorm) - launch_fused_add2(buf_2, buf_3, buf_1, bsz, _seq_length, _hidden_size, _stream); - - if (_pre_or_postLayerNorm) { - if (_attn_layer_norm.UseMean()) - _attn_layer_norm.BackwardFusedAdd(bsz_seq, - buf_3, - grad_output_ptr, - attn_nw_ptr, - grad_attn_nw_ptr, - grad_attn_nb_ptr, - streams, - buf_0, - add_res_ptr); - - else - _attn_layer_norm.BackwardFusedAdd(bsz_seq, - buf_3, - grad_output_ptr, - attn_nw_ptr, - attn_nb_ptr, - grad_attn_nw_ptr, - grad_attn_nb_ptr, - streams, - buf_0, - ff1_inp_ptr); - } else { - if (_attn_layer_norm.UseMean()) - _attn_layer_norm.Backward(bsz_seq, - buf_2, - attn_nw_ptr, - grad_attn_nw_ptr, - grad_attn_nb_ptr, - streams, - buf_0, - add_res_ptr); - - else - _attn_layer_norm.Backward(bsz_seq, - buf_2, - attn_nw_ptr, - attn_nb_ptr, - grad_attn_nw_ptr, - grad_attn_nb_ptr, - streams, - buf_0, - ff1_inp_ptr); - } - - _attn_output_dropout.Backward(bsz_seq, buf_2, buf_0, _stream); - - T* attn_output_dropout_buf = _attn_output_dropout.HasDropout() ? buf_2 : buf_0; - - _attn_out_linear.Backward(bsz_seq, - attn_output_dropout_buf, - attn_o_inp_ptr, - attn_ow_ptr, - grad_attn_ow_ptr, - grad_attn_ob_ptr, - _cublasHandle, - _stream, - buf_1); - - launch_transform_0213(buf_2, buf_1, bsz, _seq_length, _hidden_size, _heads, _stream); - - if (_attn_prob_dropout.HasDropout()) { - if (_attn_dropout_checkpoint) - _attn_prob_dropout.Forward( - bsz_heads * _seq_length, ctx_bufB_ptr_recomp, soft_out_ptr, _stream, true); - - _attn_context.Backward(bsz_heads, - buf_2, - v_tf_ptr, - (_attn_dropout_checkpoint ? ctx_bufB_ptr_recomp : ctx_bufB_ptr), - _cublasHandle, - buf_3, - ff2_buf); - } else - _attn_context.Backward( - bsz_heads, buf_2, v_tf_ptr, soft_out_ptr, _cublasHandle, buf_3, ff2_buf); - - _attn_prob_dropout.Backward(bsz_heads * _seq_length, ff2_buf, _stream); - - _softmax.Backward(bsz, ff2_buf, soft_out_ptr, _stream); - - _attn_scores.Backward(bsz_heads, ff2_buf, k_tf_ptr, q_tf_ptr, _cublasHandle, buf_2, buf_1); - - launch_transform4d_0213(ff2_buf, buf_1, bsz, _heads, _seq_length, _hidden_size, _stream, 3); - - if (_pre_or_postLayerNorm) - _qkv_linear.Backward(bsz_seq, - ff2_buf, - inp_norm_ptr, - attn_qkvw_ptr, - grad_attn_qkvw_ptr, - grad_attn_qkvb_ptr, - _cublasHandle, - _stream, - buf_2); - else - _qkv_linear.Backward(bsz_seq, - ff2_buf, - input_ptr, - attn_qkvw_ptr, - grad_attn_qkvw_ptr, - grad_attn_qkvb_ptr, - _cublasHandle, - _stream, - buf_2); - - if (_pre_or_postLayerNorm) { - if (_layer_norm.UseMean()) - _layer_norm.BackwardFusedAdd(bsz_seq, - buf_2, - buf_0, - norm_w_ptr, - grad_norm_w_ptr, - grad_norm_b_ptr, - streams, - grad_input_ptr, - input_ptr); - - else - _layer_norm.BackwardFusedAdd(bsz_seq, - buf_2, - buf_0, - norm_w_ptr, - norm_b_ptr, - grad_norm_w_ptr, - grad_norm_b_ptr, - streams, - grad_input_ptr, - inp_norm_ptr); - } else - launch_fused_add2(grad_input_ptr, buf_2, buf_0, bsz, _seq_length, _hidden_size, _stream); -} - -template -void BertTransformerLayer::SetTrainingMode(bool training) -{ - // Dropout will be skipped when not in training model. - _attn_prob_dropout.SetTrainingMode(training); - _attn_output_dropout.SetTrainingMode(training); - _layer_output_dropout.SetTrainingMode(training); -} - -template -void BertTransformerLayer::SetIntermediateBuffers(uint8_t* attn_prob_dropout_mask_ptr, - uint8_t* attn_output_dropout_mask_ptr, - uint8_t* layer_output_dropout_mask_ptr, - T* attn_layer_norm_var, - T* attn_layer_norm_mean, - T* layer_norm_var, - T* layer_norm_mean) -{ - _attn_prob_dropout.SetMask(attn_prob_dropout_mask_ptr); - _attn_output_dropout.SetMask(attn_output_dropout_mask_ptr); - _layer_output_dropout.SetMask(layer_output_dropout_mask_ptr); - - _attn_layer_norm.SetVar(attn_layer_norm_var); - _attn_layer_norm.SetMean(attn_layer_norm_mean); - _layer_norm.SetVar(layer_norm_var); - _layer_norm.SetMean(layer_norm_mean); -} - -template -void BertTransformerLayer::SetSeqLength(unsigned seq_len) -{ - _seq_length = seq_len; - - _softmax.SetSeqLength(_seq_length); - _attn_prob_dropout.SetDimension(_seq_length); - _attn_scores.SetConfig(_seq_length, _seq_length, _hidden_size / _heads); - _attn_context.SetConfig(_hidden_size / _heads, _seq_length, _seq_length); -} - -template -int create_transformer_layer(unsigned layer_id, - unsigned batch_size, - unsigned hidden_dim, - unsigned num_heads, - unsigned intermediate_size, - float attn_dropout_ratio, - float hidden_dropout_ratio, - float layer_norm_eps, - int seed, - bool pre_or_postLayerNorm, - bool test_gemm, - bool attn_dropout_checkpoint, - bool normalize_invertible, - bool gelu_checkpoint, - bool stochastic_mode) -{ - TrainingContext::Instance().SetSeed(seed); - TrainingContext::Instance().TestGemmFP16( - test_gemm, batch_size, init_seq_length, num_heads, hidden_dim / num_heads); - - auto layer = - std::make_shared>(layer_id, - batch_size, - hidden_dim, - num_heads, - intermediate_size, - init_seq_length, - attn_dropout_ratio, - hidden_dropout_ratio, - layer_norm_eps, - pre_or_postLayerNorm, - TrainingContext::Instance().GetGemmAlgos(), - attn_dropout_checkpoint, - normalize_invertible, - gelu_checkpoint, - stochastic_mode); - - s_transformer_layers[layer_id] = layer; - - std::string dtype = (std::is_same::value) ? "half" : "float"; - - std::cout << "layer #" << layer_id << " is created with date type [" << dtype << "]." - << std::endl; - - return 0; -} - -template -std::vector ds_transformer_forward(unsigned layer_id, - const torch::Tensor& input, - const torch::Tensor& input_mask, - const torch::Tensor& attn_qkvw, - const torch::Tensor& attn_qkvb, - const torch::Tensor& attn_ow, - const torch::Tensor& attn_ob, - const torch::Tensor& attn_nw, - const torch::Tensor& attn_nb, - const torch::Tensor& inter_w, - const torch::Tensor& inter_b, - const torch::Tensor& output_w, - const torch::Tensor& output_b, - const torch::Tensor& norm_w, - const torch::Tensor& norm_b, - bool training_mode, - bool prelayernorm, - bool attn_dropout_checkpoint, - bool normalize_invertible, - bool gelu_checkpoint) -{ - CHECK_INPUT(input); - CHECK_INPUT(input_mask); - CHECK_INPUT(attn_qkvw); - CHECK_INPUT(attn_qkvb); - CHECK_INPUT(attn_ow); - CHECK_INPUT(attn_ob); - CHECK_INPUT(attn_nw); - CHECK_INPUT(attn_nb); - CHECK_INPUT(inter_w); - CHECK_INPUT(inter_b); - CHECK_INPUT(output_w); - CHECK_INPUT(output_b); - CHECK_INPUT(norm_w); - CHECK_INPUT(norm_b); - - unsigned bsz = input.size(0); - - const T* input_ptr = (const T*)input.data_ptr(); - const T* input_mask_ptr = (const T*)input_mask.data_ptr(); - const T* attn_qkvw_ptr = (const T*)attn_qkvw.data_ptr(); - const T* attn_qkvb_ptr = (const T*)attn_qkvb.data_ptr(); - const T* attn_ow_ptr = (const T*)attn_ow.data_ptr(); - const T* attn_ob_ptr = (const T*)attn_ob.data_ptr(); - const T* attn_nw_ptr = (const T*)attn_nw.data_ptr(); - const T* attn_nb_ptr = (const T*)attn_nb.data_ptr(); - const T* inter_w_ptr = (const T*)inter_w.data_ptr(); - const T* inter_b_ptr = (const T*)inter_b.data_ptr(); - const T* output_w_ptr = (const T*)output_w.data_ptr(); - const T* output_b_ptr = (const T*)output_b.data_ptr(); - const T* norm_w_ptr = (const T*)norm_w.data_ptr(); - const T* norm_b_ptr = (const T*)norm_b.data_ptr(); - - auto output = torch::empty_like(input); - T* out_ptr = (T*)output.data_ptr(); - - auto options = torch::TensorOptions() - .dtype(input.options().dtype()) - .layout(torch::kStrided) - .device(torch::kCUDA) - .requires_grad(true); - - auto uint8_options = torch::TensorOptions() - .dtype(torch::kInt8) - .layout(torch::kStrided) - .device(torch::kCUDA) - .requires_grad(false); - - std::shared_ptr> layer = - std::static_pointer_cast>(s_transformer_layers[layer_id]); - - unsigned seq_len = layer->GetSeqLength(); - if (input.size(1) != seq_len) { - seq_len = input.size(1); - layer->SetSeqLength(seq_len); - } - - auto workspace = torch::empty({get_workspace_size(bsz, - seq_len, - layer->GetHiddenSize(), - layer->GetIntermediateSize(), - layer->GetNumHeads(), - layer->IsTrainingMode(), - layer->GeluCheckpoint())}, - options); - TrainingContext::Instance().SetWorkSpace((T*)workspace.data_ptr()); - - auto inp_norm = ((prelayernorm || !normalize_invertible) ? torch::empty_like(input) : output); - auto add_res = (normalize_invertible ? inp_norm : torch::empty_like(input)); - auto attn_o_inp = torch::empty_like(input); - auto qkv_tf = torch::empty({(bsz * seq_len), output_w.size(0) * 3}, options); - - auto attn_prob_dropout_mask = - torch::empty({(bsz * layer->GetNumHeads() * seq_len), seq_len}, uint8_options); - auto attn_output_dropout_mask = - torch::empty({(bsz * seq_len), layer->GetHiddenSize()}, uint8_options); - auto layer_output_dropout_mask = - torch::empty({(bsz * seq_len), layer->GetHiddenSize()}, uint8_options); - - auto attn_layer_norm_var = torch::empty({(bsz * seq_len)}, options); - auto attn_layer_norm_mean = torch::empty({(bsz * seq_len)}, options); - auto layer_norm_var = torch::empty({(bsz * seq_len)}, options); - auto layer_norm_mean = torch::empty({(bsz * seq_len)}, options); - - T* inp_norm_ptr = (T*)inp_norm.data_ptr(); - T* add_res_ptr = (T*)add_res.data_ptr(); - T* q_tf_ptr = (T*)qkv_tf.data_ptr(); - T* k_tf_ptr = q_tf_ptr + (bsz * seq_len * output_w.size(0)); //(T*)k_tf.data_ptr(); - T* v_tf_ptr = k_tf_ptr + (bsz * seq_len * output_w.size(0)); //(T*)v_tf.data_ptr(); - T* attn_o_inp_ptr = (T*)attn_o_inp.data_ptr(); - - torch::Tensor ff2_inp = torch::empty({(bsz * seq_len), output_w.size(1)}, options); - torch::Tensor gelu_inp = - (gelu_checkpoint ? ff2_inp : torch::empty({(bsz * seq_len), output_w.size(1)}, options)); - auto ff1_inp = torch::empty_like(input); - T* ff2_inp_ptr = (T*)ff2_inp.data_ptr(); - T* gelu_inp_ptr = (T*)gelu_inp.data_ptr(); - T* ff1_inp_ptr = (T*)ff1_inp.data_ptr(); - - torch::Tensor soft_out = - torch::empty({(bsz * layer->GetNumHeads() * seq_len), seq_len}, options); - torch::Tensor ctx_bufB = - (attn_dropout_checkpoint - ? soft_out - : torch::empty({(bsz * layer->GetNumHeads() * seq_len), seq_len}, options)); - T* soft_out_ptr = (T*)soft_out.data_ptr(); - T* ctx_bufB_ptr = (T*)ctx_bufB.data_ptr(); - - layer->SetTrainingMode(training_mode); - layer->SetIntermediateBuffers((uint8_t*)attn_prob_dropout_mask.data_ptr(), - (uint8_t*)attn_output_dropout_mask.data_ptr(), - (uint8_t*)layer_output_dropout_mask.data_ptr(), - (T*)attn_layer_norm_var.data_ptr(), - (T*)attn_layer_norm_mean.data_ptr(), - (T*)layer_norm_var.data_ptr(), - (T*)layer_norm_mean.data_ptr()); - - layer->Forward(bsz, - input_ptr, - input_mask_ptr, - attn_qkvw_ptr, - attn_qkvb_ptr, - attn_ow_ptr, - attn_ob_ptr, - attn_nw_ptr, - attn_nb_ptr, - inter_w_ptr, - inter_b_ptr, - output_w_ptr, - output_b_ptr, - norm_w_ptr, - norm_b_ptr, - out_ptr, - inp_norm_ptr, - q_tf_ptr, - k_tf_ptr, - v_tf_ptr, - soft_out_ptr, - ctx_bufB_ptr, - attn_o_inp_ptr, - add_res_ptr, - ff1_inp_ptr, - gelu_inp_ptr, - ff2_inp_ptr); - - return {output, - inp_norm, - qkv_tf, - soft_out, - ctx_bufB, - attn_o_inp, - add_res, - ff1_inp, - gelu_inp, - ff2_inp, - attn_prob_dropout_mask, - attn_output_dropout_mask, - layer_output_dropout_mask, - attn_layer_norm_var, - attn_layer_norm_mean, - layer_norm_var, - layer_norm_mean}; -} - -template -std::vector ds_transformer_backward(unsigned layer_id, - const torch::Tensor& grad_output, - const torch::Tensor& output, - const torch::Tensor& inp_norm, - const torch::Tensor& qkv_tf, - const torch::Tensor& soft_out, - const torch::Tensor& ctx_bufB, - const torch::Tensor& attn_o_inp, - const torch::Tensor& add_res, - const torch::Tensor& ff1_inp, - const torch::Tensor& gelu_inp, - const torch::Tensor& ff2_inp, - const torch::Tensor& attn_prob_dropout_mask, - const torch::Tensor& attn_output_dropout_mask, - const torch::Tensor& layer_output_dropout_mask, - const torch::Tensor& attn_layer_norm_var, - const torch::Tensor& attn_layer_norm_mean, - const torch::Tensor& layer_norm_var, - const torch::Tensor& layer_norm_mean, - const torch::Tensor& input, - const torch::Tensor& input_mask, - const torch::Tensor& attn_qkvw, - const torch::Tensor& attn_qkvb, - const torch::Tensor& attn_ow, - const torch::Tensor& attn_ob, - const torch::Tensor& attn_nw, - const torch::Tensor& attn_nb, - const torch::Tensor& inter_w, - const torch::Tensor& inter_b, - const torch::Tensor& output_w, - const torch::Tensor& output_b, - const torch::Tensor& norm_w, - const torch::Tensor& norm_b) -{ - auto g_output = grad_output.contiguous(); - CHECK_INPUT(g_output); - CHECK_INPUT(output); - CHECK_INPUT(inp_norm); - CHECK_INPUT(qkv_tf); - CHECK_INPUT(add_res); - CHECK_INPUT(soft_out); - CHECK_INPUT(ctx_bufB); - CHECK_INPUT(attn_o_inp); - CHECK_INPUT(ff1_inp); - CHECK_INPUT(gelu_inp); - CHECK_INPUT(ff2_inp); - CHECK_INPUT(input); - CHECK_INPUT(input_mask); - CHECK_INPUT(attn_qkvw); - CHECK_INPUT(attn_qkvb); - CHECK_INPUT(attn_ow); - CHECK_INPUT(attn_ob); - CHECK_INPUT(attn_nw); - CHECK_INPUT(attn_nb); - CHECK_INPUT(inter_w); - CHECK_INPUT(inter_b); - CHECK_INPUT(output_w); - CHECK_INPUT(output_b); - CHECK_INPUT(norm_w); - CHECK_INPUT(norm_b); - - unsigned bsz = g_output.size(0); - - std::shared_ptr> layer = - std::static_pointer_cast>(s_transformer_layers[layer_id]); - - unsigned seq_len = layer->GetSeqLength(); - if (g_output.size(1) != seq_len) { - seq_len = g_output.size(1); - layer->SetSeqLength(seq_len); - } - auto options = torch::TensorOptions() - .dtype(g_output.options().dtype()) - .layout(torch::kStrided) - .device(torch::kCUDA) - .requires_grad(true); - auto workspace = torch::empty({get_workspace_size(bsz, - seq_len, - layer->GetHiddenSize(), - layer->GetIntermediateSize(), - layer->GetNumHeads(), - layer->IsTrainingMode(), - layer->GeluCheckpoint())}, - options); - TrainingContext::Instance().SetWorkSpace((T*)workspace.data_ptr()); - - auto grad_input = torch::empty_like(input); - auto grad_attn_qkvw = torch::empty_like(attn_qkvw); - auto grad_attn_qkvb = torch::empty_like(attn_qkvb); - auto grad_attn_ow = torch::empty_like(attn_ow); - auto grad_attn_ob = torch::empty_like(attn_ob); - auto grad_attn_nw = torch::empty_like(attn_nw); - auto grad_attn_nb = torch::empty_like(attn_nb); - auto grad_inter_w = torch::empty_like(inter_w); - auto grad_inter_b = torch::empty_like(inter_b); - auto grad_output_w = torch::empty_like(output_w); - auto grad_output_b = torch::empty_like(output_b); - auto grad_norm_w = torch::empty_like(norm_w); - auto grad_norm_b = torch::empty_like(norm_b); - - // inputs. - const T* grad_output_ptr = (const T*)g_output.data_ptr(); - const T* input_ptr = (const T*)input.data_ptr(); - const T* output_ptr = (const T*)output.data_ptr(); - const T* inp_norm_ptr = (const T*)inp_norm.data_ptr(); - const T* q_tf_ptr = (const T*)qkv_tf.data_ptr(); - const T* add_res_ptr = (const T*)add_res.data_ptr(); - const T* k_tf_ptr = - q_tf_ptr + (bsz * layer->GetSeqLength() * output_w.size(0)); //(const T*)k_tf.data_ptr(); - const T* v_tf_ptr = - k_tf_ptr + (bsz * layer->GetSeqLength() * output_w.size(0)); //(const T*)v_tf.data_ptr(); - const T* ff1_inp_ptr = (const T*)ff1_inp.data_ptr(); - const T* gelu_inp_ptr = (const T*)gelu_inp.data_ptr(); - const T* ff2_inp_ptr = (const T*)ff2_inp.data_ptr(); - const T* ctx_bufB_ptr = (const T*)ctx_bufB.data_ptr(); - const T* soft_out_ptr = (const T*)soft_out.data_ptr(); - const T* attn_o_inp_ptr = (const T*)attn_o_inp.data_ptr(); - const T* input_mask_ptr = (const T*)input_mask.data_ptr(); - const T* attn_qkvw_ptr = (const T*)attn_qkvw.data_ptr(); - const T* attn_ow_ptr = (const T*)attn_ow.data_ptr(); - const T* attn_nw_ptr = (const T*)attn_nw.data_ptr(); - const T* attn_nb_ptr = (const T*)attn_nb.data_ptr(); - const T* inter_w_ptr = (const T*)inter_w.data_ptr(); - const T* inter_b_ptr = (const T*)inter_b.data_ptr(); - const T* output_w_ptr = (const T*)output_w.data_ptr(); - const T* norm_w_ptr = (const T*)norm_w.data_ptr(); - const T* norm_b_ptr = (const T*)norm_b.data_ptr(); - - // outputs. - T* grad_input_ptr = (T*)grad_input.data_ptr(); - T* grad_attn_qkvw_ptr = (T*)grad_attn_qkvw.data_ptr(); - T* grad_attn_qkvb_ptr = (T*)grad_attn_qkvb.data_ptr(); - T* grad_attn_ow_ptr = (T*)grad_attn_ow.data_ptr(); - T* grad_attn_ob_ptr = (T*)grad_attn_ob.data_ptr(); - T* grad_attn_nw_ptr = (T*)grad_attn_nw.data_ptr(); - T* grad_attn_nb_ptr = (T*)grad_attn_nb.data_ptr(); - T* grad_inter_w_ptr = (T*)grad_inter_w.data_ptr(); - T* grad_inter_b_ptr = (T*)grad_inter_b.data_ptr(); - T* grad_output_w_ptr = (T*)grad_output_w.data_ptr(); - T* grad_output_b_ptr = (T*)grad_output_b.data_ptr(); - T* grad_norm_w_ptr = (T*)grad_norm_w.data_ptr(); - T* grad_norm_b_ptr = (T*)grad_norm_b.data_ptr(); - - layer->SetIntermediateBuffers((uint8_t*)attn_prob_dropout_mask.data_ptr(), - (uint8_t*)attn_output_dropout_mask.data_ptr(), - (uint8_t*)layer_output_dropout_mask.data_ptr(), - (T*)attn_layer_norm_var.data_ptr(), - (T*)attn_layer_norm_mean.data_ptr(), - (T*)layer_norm_var.data_ptr(), - (T*)layer_norm_mean.data_ptr()); - - layer->Backward(bsz, - grad_output_ptr, - input_ptr, - output_ptr, - inp_norm_ptr, - q_tf_ptr, - k_tf_ptr, - v_tf_ptr, - soft_out_ptr, - ctx_bufB_ptr, - attn_o_inp_ptr, - add_res_ptr, - ff1_inp_ptr, - gelu_inp_ptr, - ff2_inp_ptr, - input_mask_ptr, - attn_qkvw_ptr, - attn_ow_ptr, - attn_nw_ptr, - attn_nb_ptr, - inter_w_ptr, - inter_b_ptr, - output_w_ptr, - norm_w_ptr, - norm_b_ptr, - - grad_input_ptr, - grad_attn_qkvw_ptr, - grad_attn_qkvb_ptr, - grad_attn_ow_ptr, - grad_attn_ob_ptr, - grad_attn_nw_ptr, - grad_attn_nb_ptr, - grad_inter_w_ptr, - grad_inter_b_ptr, - grad_output_w_ptr, - grad_output_b_ptr, - grad_norm_w_ptr, - grad_norm_b_ptr); - - return {grad_input, - grad_attn_qkvw, - grad_attn_qkvb, - grad_attn_ow, - grad_attn_ob, - grad_attn_nw, - grad_attn_nb, - grad_inter_w, - grad_inter_b, - grad_output_w, - grad_output_b, - grad_norm_w, - grad_norm_b}; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("forward_fp32", - &ds_transformer_forward, - "DeepSpeed Transformer forward with fp32 (CUDA)"); - m.def("forward_fp16", - &ds_transformer_forward<__half>, - "DeepSpeed Transformer forward with fp16 (CUDA)"); - m.def("backward_fp32", - &ds_transformer_backward, - "DeepSpeed Transformer backward with fp32 (CUDA)"); - m.def("backward_fp16", - &ds_transformer_backward<__half>, - "DeepSpeed Transformer backward with fp16 (CUDA)"); - m.def("create_transformer_layer_fp32", - &create_transformer_layer, - "Create DeepSpeed Transformer Transformer Layer with fp32 (CUDA)"); - m.def("create_transformer_layer_fp16", - &create_transformer_layer<__half>, - "Create DeepSpeed Transformer Transformer Layer with fp16 (CUDA)"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/gelu_kernels.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/gelu_kernels.cu deleted file mode 100644 index 273891b91923056fe3468600827f538c9cee5009..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/gelu_kernels.cu +++ /dev/null @@ -1,335 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "custom_cuda_layers.h" - -inline __device__ float gelu(const float x) -{ - const float sqrt_param = 0.79788456080286535587989211986876f; - const float mul_param = 0.044715; - return x * 0.5f * (1.0f + tanhf(sqrt_param * (x + mul_param * x * x * x))); -} - -inline __device__ float d_gelu(const float x) -{ - const float sqrt_param = 0.79788456080286535587989211986876f; - const float mul_param = 0.044715; - - float x2mul = x * x * mul_param; - float tan_h = tanhf(sqrt_param * (x + x * x2mul)); - float dg1 = 0.5f * (1.0f + tan_h); - float dg2 = x * 0.5f * sqrt_param * (1 - tan_h * tan_h); - float dg3 = dg2 * 3 * x2mul; - return (dg1 + dg2 + dg3); -} - -/* -Fused bias add with GELU - -Loads a vector of 4 elements each iteration, for stride -iterations. It was written with the intention to launch 256 thread -threadblocks, so to launch for bert-large, we would set ITERATIONS -to 4. This is currently done automatically as a heuristic, setting -the number of iterations as blocks of 1024. - -For FP16, the values are loaded from memory as __half, but converted -to FP32 for the arithmetic itself, to prevent numerous overflow on -the intermediate hyperbolic tangent, since there's no intrinsic -that computes it directly. -*/ - -__global__ void gelu_kernel(const float* input, float* vals, int row_stride, int iterations) -{ - int row = blockIdx.x; - int id = threadIdx.x; - int loop_stride = blockDim.x; - - const float4* input_cast = reinterpret_cast(input); - float4* vals_cast = reinterpret_cast(vals); - - for (int i = 0; i < iterations; i++) { - if (i * loop_stride + id < row_stride) { - float4 data = input_cast[row * row_stride + i * loop_stride + id]; - - data.x = gelu(data.x); - data.y = gelu(data.y); - data.z = gelu(data.z); - data.w = gelu(data.w); - - vals_cast[row * row_stride + i * loop_stride + id] = data; - } - } -} - -__global__ void gelu_kernel(const __half* input, __half* vals, int row_stride, int iterations) -{ -#ifdef HALF_PRECISION_AVAILABLE - int row = blockIdx.x; - int id = threadIdx.x; - int loop_stride = blockDim.x; - - const float2* input_cast = reinterpret_cast(input); - float2* vals_cast = reinterpret_cast(vals); - - for (int i = 0; i < iterations; i++) { - if (i * loop_stride + id < row_stride) { - float2 vals_vec = input_cast[row * row_stride + i * loop_stride + id]; - - __half2* vals_half = reinterpret_cast<__half2*>(&vals_vec); - - float2 low_data = __half22float2(vals_half[0]); - float2 high_data = __half22float2(vals_half[1]); - - low_data.x = gelu(low_data.x); - low_data.y = gelu(low_data.y); - high_data.x = gelu(high_data.x); - high_data.y = gelu(high_data.y); - - vals_half[0] = __float22half2_rn(low_data); - vals_half[1] = __float22half2_rn(high_data); - - vals_cast[row * row_stride + i * loop_stride + id] = vals_vec; - } - } -#endif -} - -__global__ void fused_bias_gelu(const float* input, - const float* bias, - float* vals, - int row_stride, - int iterations) -{ - int row = blockIdx.x; - int id = threadIdx.x; - int loop_stride = blockDim.x; - - const float4* input_cast = reinterpret_cast(input); - float4* vals_cast = reinterpret_cast(vals); - const float4* bias_cast = reinterpret_cast(bias); - - for (int i = 0; i < iterations; i++) { - if (i * loop_stride + id < row_stride) { - float4 data = input_cast[row * row_stride + i * loop_stride + id]; - float4 bias_data = bias_cast[i * loop_stride + id]; - - data.x += bias_data.x; - data.y += bias_data.y; - data.z += bias_data.z; - data.w += bias_data.w; - - data.x = gelu(data.x); - data.y = gelu(data.y); - data.z = gelu(data.z); - data.w = gelu(data.w); - - vals_cast[row * row_stride + i * loop_stride + id] = data; - } - } -} - -__global__ void fused_bias_gelu(const __half* input, - const __half* bias, - __half* vals, - int row_stride, - int iterations) -{ -#ifdef HALF_PRECISION_AVAILABLE - int row = blockIdx.x; - int id = threadIdx.x; - int loop_stride = blockDim.x; - - const float2* input_cast = reinterpret_cast(input); - float2* vals_cast = reinterpret_cast(vals); - const float2* bias_cast = reinterpret_cast(bias); - - for (int i = 0; i < iterations; i++) { - if (i * loop_stride + id < row_stride) { - float2 vals_vec = input_cast[row * row_stride + i * loop_stride + id]; - float2 bias_vec = bias_cast[i * loop_stride + id]; - - __half2* vals_half = reinterpret_cast<__half2*>(&vals_vec); - __half2* bias_half = reinterpret_cast<__half2*>(&bias_vec); - - float2 low_data = __half22float2(vals_half[0]); - float2 high_data = __half22float2(vals_half[1]); - - float2 low_bias = __half22float2(bias_half[0]); - float2 high_bias = __half22float2(bias_half[1]); - - low_data.x += low_bias.x; - low_data.y += low_bias.y; - high_data.x += high_bias.x; - high_data.y += high_bias.y; - - low_data.x = gelu(low_data.x); - low_data.y = gelu(low_data.y); - high_data.x = gelu(high_data.x); - high_data.y = gelu(high_data.y); - - vals_half[0] = __float22half2_rn(low_data); - vals_half[1] = __float22half2_rn(high_data); - - vals_cast[row * row_stride + i * loop_stride + id] = vals_vec; - } - } -#endif -} - -__global__ void d_gelu_func(float* d_output, - const float* gelu_input, - const float* bias, - int row_stride, - int iterations) -{ - int row = blockIdx.x; - int id = threadIdx.x; - int loop_stride = blockDim.x; - - float4* d_output_cast = reinterpret_cast(d_output); - const float4* gelu_input_cast = reinterpret_cast(gelu_input); - const float4* bias_cast = reinterpret_cast(bias); - - for (int i = 0; i < iterations; i++) { - if (i * loop_stride + id < row_stride) { - float4 output_data = d_output_cast[row * row_stride + i * loop_stride + id]; - float4 gelu_input_data = gelu_input_cast[row * row_stride + i * loop_stride + id]; - float4 bias_data = bias_cast[i * loop_stride + id]; - - gelu_input_data.x += bias_data.x; - gelu_input_data.y += bias_data.y; - gelu_input_data.z += bias_data.z; - gelu_input_data.w += bias_data.w; - - output_data.x *= d_gelu(gelu_input_data.x); - output_data.y *= d_gelu(gelu_input_data.y); - output_data.z *= d_gelu(gelu_input_data.z); - output_data.w *= d_gelu(gelu_input_data.w); - - d_output_cast[row * row_stride + i * loop_stride + id] = output_data; - } - } -} - -__global__ void d_gelu_func(__half* d_output, - const __half* gelu_input, - const __half* bias, - int row_stride, - int iterations) -{ -#ifdef HALF_PRECISION_AVAILABLE - int row = blockIdx.x; - int id = threadIdx.x; - int loop_stride = blockDim.x; - - float2* d_output_cast = reinterpret_cast(d_output); - const float2* gelu_input_cast = reinterpret_cast(gelu_input); - const float2* bias_cast = reinterpret_cast(bias); - -#pragma unroll - for (int i = 0; i < iterations; i++) { - if (i * loop_stride + id < row_stride) { - float2 output_data = d_output_cast[row * row_stride + i * loop_stride + id]; - float2 gelu_input_data = gelu_input_cast[row * row_stride + i * loop_stride + id]; - float2 bias_vec = bias_cast[i * loop_stride + id]; - - __half2* output_data_half = reinterpret_cast<__half2*>(&output_data); - __half2* gelu_input_data_half = reinterpret_cast<__half2*>(&gelu_input_data); - __half2* bias_half = reinterpret_cast<__half2*>(&bias_vec); - - float2 output_half_0 = __half22float2(output_data_half[0]); - float2 output_half_1 = __half22float2(output_data_half[1]); - - float2 gelu_input_half_0 = __half22float2(gelu_input_data_half[0]); - float2 gelu_input_half_1 = __half22float2(gelu_input_data_half[1]); - - float2 bias_half_0 = __half22float2(bias_half[0]); - float2 bias_half_1 = __half22float2(bias_half[1]); - - gelu_input_half_0.x += bias_half_0.x; - gelu_input_half_0.y += bias_half_0.y; - gelu_input_half_1.x += bias_half_1.x; - gelu_input_half_1.y += bias_half_1.y; - - output_half_0.x *= d_gelu(gelu_input_half_0.x); - output_half_0.y *= d_gelu(gelu_input_half_0.y); - output_half_1.x *= d_gelu(gelu_input_half_1.x); - output_half_1.y *= d_gelu(gelu_input_half_1.y); - - float2 result; - __half2* result_half2 = reinterpret_cast<__half2*>(&result); - - result_half2[0] = __float22half2_rn(output_half_0); - result_half2[1] = __float22half2_rn(output_half_1); - - d_output_cast[row * row_stride + i * loop_stride + id] = result; - } - } -#endif -} - -template -void launch_bias_gelu(const T* input, - const T* bias, - T* output, - int intermediate_size, - int batch_size, - cudaStream_t stream) -{ - int iterations = (intermediate_size + 1023) / 1024; - int threads = (intermediate_size - 1) / (iterations * 4) + 1; - dim3 block_dims(threads); - dim3 grid_dims(batch_size); - - fused_bias_gelu<<>>( - input, bias, output, intermediate_size / 4, iterations); -} - -template -void launch_gelu(const T* input, - T* output, - int intermediate_size, - int batch_size, - cudaStream_t stream) -{ - int iterations = (intermediate_size + 1023) / 1024; - int threads = (intermediate_size - 1) / (iterations * 4) + 1; - dim3 block_dims(threads); - dim3 grid_dims(batch_size); - - gelu_kernel<<>>( - input, output, intermediate_size / 4, iterations); -} - -template void launch_bias_gelu(const float*, const float*, float*, int, int, cudaStream_t); -template void launch_bias_gelu<__half>(const __half*, - const __half*, - __half*, - int, - int, - cudaStream_t); - -template void launch_gelu(const float*, float*, int, int, cudaStream_t); -template void launch_gelu<__half>(const __half*, __half*, int, int, cudaStream_t); - -template -void launch_d_gelu(T* d_output, - const T* input, - const T* bias, - int intermediate_size, - int batch_size, - cudaStream_t stream) -{ - int iterations = (intermediate_size + 1023) / 1024; - int threads = (intermediate_size - 1) / (iterations * 4) + 1; - dim3 block_dims(threads); - dim3 grid_dims(batch_size); - - d_gelu_func<<>>( - d_output, input, bias, intermediate_size / 4, iterations); -} - -template void launch_d_gelu(float*, const float*, const float*, int, int, cudaStream_t); -template void launch_d_gelu<__half>(__half*, const __half*, const __half*, int, int, cudaStream_t); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/general_kernels.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/general_kernels.cu deleted file mode 100644 index a987eec5ef0bdc87153422eae635507f44188802..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/general_kernels.cu +++ /dev/null @@ -1,416 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "general_kernels.h" - -namespace cg = cooperative_groups; - -template -__global__ void column_sum_reduce(const T* __restrict__ inp, - T* __restrict__ out, - int rows, - int width) -{ - __shared__ float tile[TILE_DIM][TILE_DIM + 1]; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int idx = blockDim.x * blockIdx.x + threadIdx.x; - - int y_stride = width * TILE_DIM; - - float localSum = 0; - - // Loop across matrix height - if (idx < width) { - int offset = threadIdx.y * width + idx; - for (int r = threadIdx.y; r < rows; r += TILE_DIM) { - localSum += (float)inp[offset]; - offset += y_stride; - } - } - - tile[threadIdx.x][threadIdx.y] = localSum; - - __syncthreads(); - - // Sum the shared buffer. - float sum = tile[threadIdx.y][threadIdx.x]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < TILE_DIM; i <<= 1) sum += g.shfl_down(sum, i); - - if (threadIdx.x == 0) { - int pos = blockIdx.x * TILE_DIM + threadIdx.y; - if (pos < width) out[pos] = sum; - } -} - -template -void launch_fuse_transpose_bias_kernel(const T* inp, - T* out, - int rows, - int cols, - cudaStream_t stream); - -template <> -void launch_fuse_transpose_bias_kernel(const float* inp, - float* out, - int rows, - int cols, - cudaStream_t stream) -{ - // assert(rows % TILE_DIM == 0); - // assert(cols % TILE_DIM == 0); - - dim3 grid_dim((cols - 1) / TILE_DIM + 1); - dim3 block_dim(TILE_DIM, TILE_DIM); - - column_sum_reduce<<>>(inp, out, rows, cols); -} - -template <> -void launch_fuse_transpose_bias_kernel<__half>(const __half* inp, - __half* out, - int rows, - int cols, - cudaStream_t stream) -{ - // assert(rows % TILE_DIM == 0); - // assert(cols % TILE_DIM == 0); - - dim3 grid_dim((cols - 1) / TILE_DIM + 1); - dim3 block_dim(TILE_DIM, TILE_DIM); - - column_sum_reduce<__half><<>>(inp, out, rows, cols); -} - -__global__ void fused_add2_kernel(const int N, float* out, const float* inp1, const float* inp2) -{ - const float4* inp1_4 = reinterpret_cast(inp1); - const float4* inp2_4 = reinterpret_cast(inp2); - float4* out_4 = reinterpret_cast(out); - - CUDA_1D_KERNEL_LOOP(j, N) - { - float4 val; - float4 inp1_reg = inp1_4[j]; - float4 inp2_reg = inp2_4[j]; - - val.x = inp1_reg.x + inp2_reg.x; - val.y = inp1_reg.y + inp2_reg.y; - val.z = inp1_reg.z + inp2_reg.z; - val.w = inp1_reg.w + inp2_reg.w; - - out_4[j] = val; - } -} - -__global__ void fused_add2_kernel(const int N, __half* out, const __half* inp1, const __half* inp2) -{ - float2 inp1_4; - float2 inp2_4; - - __half2* inp1_h = reinterpret_cast<__half2*>(&inp1_4); - __half2* inp2_h = reinterpret_cast<__half2*>(&inp2_4); - - const float2* inp1_arr = reinterpret_cast(inp1); - const float2* inp2_arr = reinterpret_cast(inp2); - - CUDA_1D_KERNEL_LOOP(j, N) - { - inp1_4 = inp1_arr[j]; - inp2_4 = inp2_arr[j]; - - float2 inp1_h_f_0 = __half22float2(inp1_h[0]); - float2 inp1_h_f_1 = __half22float2(inp1_h[1]); - - float2 inp2_h_f_0 = __half22float2(inp2_h[0]); - float2 inp2_h_f_1 = __half22float2(inp2_h[1]); - - inp1_h_f_0.x += inp2_h_f_0.x; - inp1_h_f_0.y += inp2_h_f_0.y; - inp1_h_f_1.x += inp2_h_f_1.x; - inp1_h_f_1.y += inp2_h_f_1.y; - - float2 val_f; - __half2* val_h = reinterpret_cast<__half2*>(&val_f); - - val_h[0] = __float22half2_rn(inp1_h_f_0); - val_h[1] = __float22half2_rn(inp1_h_f_1); - - float2* out_4 = reinterpret_cast(out); - out_4[j] = val_f; - } -} - -template <> -void launch_fused_add2(float* out, - const float* inp1, - const float* inp2, - int batch_size, - int seq_length, - int hidden_dim, - cudaStream_t& stream) -{ - int total_count = batch_size * seq_length * hidden_dim / 4; - dim3 grid_dim = DS_GET_BLOCKS(total_count); //(batch_size * seq_length); - - dim3 block_dim = DS_CUDA_NUM_THREADS; //(hidden_dim / 4); - - fused_add2_kernel<<>>(total_count, out, inp1, inp2); -} - -template <> -void launch_fused_add2<__half>(__half* out, - const __half* inp1, - const __half* inp2, - int batch_size, - int seq_length, - int hidden_dim, - cudaStream_t& stream) -{ - int total_count = batch_size * seq_length * hidden_dim / 4; - dim3 grid_dim = DS_GET_BLOCKS(total_count); //(batch_size * seq_length); - - dim3 block_dim = DS_CUDA_NUM_THREADS; //(hidden_dim / 4); - - fused_add2_kernel<<>>(total_count, out, inp1, inp2); -} - -__global__ void fused_add3_kernel(float* out, - const float* inp1, - const float* inp2, - const float* inp3, - int size, - int row_stride) -{ - int row = blockIdx.x; - int id = threadIdx.x; - - const float4* inp1_4 = reinterpret_cast(inp1); - const float4* inp2_4 = reinterpret_cast(inp2); - const float4* inp3_4 = reinterpret_cast(inp3); - - float4* out_4 = reinterpret_cast(out); - - float4 val; - float4 inp1_reg = inp1_4[row * row_stride + id]; - float4 inp2_reg = inp2_4[row * row_stride + id]; - float4 inp3_reg = inp3_4[row * row_stride + id]; - - val.x = inp1_reg.x + inp2_reg.x + inp3_reg.x; - val.y = inp1_reg.y + inp2_reg.y + inp3_reg.y; - val.z = inp1_reg.z + inp2_reg.z + inp3_reg.z; - val.w = inp1_reg.w + inp2_reg.w + inp3_reg.w; - - out_4[row * row_stride + id] = val; -} - -__global__ void fused_add3_kernel(__half* out, - const __half* inp1, - const __half* inp2, - const __half* inp3, - int size, - int row_stride) -{ - int row = blockIdx.x; - int id = threadIdx.x; - const float2* inp1_arr = reinterpret_cast(inp1); - const float2* inp2_arr = reinterpret_cast(inp2); - const float2* inp3_arr = reinterpret_cast(inp3); - - float2 inp1_4 = inp1_arr[row * row_stride + id]; - float2 inp2_4 = inp2_arr[row * row_stride + id]; - float2 inp3_4 = inp3_arr[row * row_stride + id]; - - __half2* inp1_h = reinterpret_cast<__half2*>(&inp1_4); - __half2* inp2_h = reinterpret_cast<__half2*>(&inp2_4); - __half2* inp3_h = reinterpret_cast<__half2*>(&inp3_4); - - float2 inp1_h_f_0 = __half22float2(inp1_h[0]); - float2 inp1_h_f_1 = __half22float2(inp1_h[1]); - - float2 inp2_h_f_0 = __half22float2(inp2_h[0]); - float2 inp2_h_f_1 = __half22float2(inp2_h[1]); - - float2 inp3_h_f_0 = __half22float2(inp3_h[0]); - float2 inp3_h_f_1 = __half22float2(inp3_h[1]); - - inp1_h_f_0.x += (inp2_h_f_0.x + inp3_h_f_0.x); - inp1_h_f_0.y += (inp2_h_f_0.y + inp3_h_f_0.y); - inp1_h_f_1.x += (inp2_h_f_1.x + inp3_h_f_1.x); - inp1_h_f_1.y += (inp2_h_f_1.y + inp3_h_f_1.y); - - float2 val_f; - __half2* val_h = reinterpret_cast<__half2*>(&val_f); - - val_h[0] = __float22half2_rn(inp1_h_f_0); - val_h[1] = __float22half2_rn(inp1_h_f_1); - - float2* out_4 = reinterpret_cast(out); - out_4[row * row_stride + id] = val_f; -} - -template <> -void launch_fused_add3(float* out, - const float* inp1, - const float* inp2, - const float* inp3, - int batch_size, - int seq_length, - int hidden_size, - cudaStream_t& stream) -{ - dim3 grid_dim(batch_size * seq_length); - - dim3 block_dim(hidden_size / 4); - - fused_add3_kernel<<>>( - out, inp1, inp2, inp3, (batch_size * seq_length * hidden_size), hidden_size / 4); -} - -template <> -void launch_fused_add3<__half>(__half* out, - const __half* inp1, - const __half* inp2, - const __half* inp3, - int batch_size, - int seq_length, - int hidden_size, - cudaStream_t& stream) -{ - dim3 grid_dim(batch_size * seq_length); - - dim3 block_dim(hidden_size / 4); - - fused_add3_kernel<<>>( - out, inp1, inp2, inp3, (batch_size * seq_length * hidden_size), hidden_size / 4); -} - -__global__ void fused_add4_kernel(float* out, - const float* inp1, - const float* inp2, - const float* inp3, - const float* inp4, - int size, - int row_stride) -{ - int row = blockIdx.x; - int id = threadIdx.x; - - const float4* inp1_4 = reinterpret_cast(inp1); - const float4* inp2_4 = reinterpret_cast(inp2); - const float4* inp3_4 = reinterpret_cast(inp3); - const float4* inp4_4 = reinterpret_cast(inp4); - float4* out_4 = reinterpret_cast(out); - - float4 val; - float4 inp1_reg = inp1_4[row * row_stride + id]; - float4 inp2_reg = inp2_4[row * row_stride + id]; - float4 inp3_reg = inp3_4[row * row_stride + id]; - float4 inp4_reg = inp4_4[row * row_stride + id]; - - val.x = inp1_reg.x + inp2_reg.x + inp3_reg.x + inp4_reg.x; - val.y = inp1_reg.y + inp2_reg.y + inp3_reg.y + inp4_reg.y; - val.z = inp1_reg.z + inp2_reg.z + inp3_reg.z + inp4_reg.z; - val.w = inp1_reg.w + inp2_reg.w + inp3_reg.w + inp4_reg.w; - - out_4[row * row_stride + id] = val; -} - -__global__ void fused_add4_kernel(__half* out, - const __half* inp1, - const __half* inp2, - const __half* inp3, - const __half* inp4, - int size, - int row_stride) -{ - int row = blockIdx.x; - int id = threadIdx.x; - const float2* inp1_arr = reinterpret_cast(inp1); - const float2* inp2_arr = reinterpret_cast(inp2); - const float2* inp3_arr = reinterpret_cast(inp3); - const float2* inp4_arr = reinterpret_cast(inp4); - - float2 inp1_4 = inp1_arr[row * row_stride + id]; - float2 inp2_4 = inp2_arr[row * row_stride + id]; - float2 inp3_4 = inp3_arr[row * row_stride + id]; - float2 inp4_4 = inp4_arr[row * row_stride + id]; - - __half2* inp1_h = reinterpret_cast<__half2*>(&inp1_4); - __half2* inp2_h = reinterpret_cast<__half2*>(&inp2_4); - __half2* inp3_h = reinterpret_cast<__half2*>(&inp3_4); - __half2* inp4_h = reinterpret_cast<__half2*>(&inp4_4); - - float2 inp1_h_f_0 = __half22float2(inp1_h[0]); - float2 inp1_h_f_1 = __half22float2(inp1_h[1]); - - float2 inp2_h_f_0 = __half22float2(inp2_h[0]); - float2 inp2_h_f_1 = __half22float2(inp2_h[1]); - - float2 inp3_h_f_0 = __half22float2(inp3_h[0]); - float2 inp3_h_f_1 = __half22float2(inp3_h[1]); - - float2 inp4_h_f_0 = __half22float2(inp4_h[0]); - float2 inp4_h_f_1 = __half22float2(inp4_h[1]); - - inp1_h_f_0.x += (inp2_h_f_0.x + inp3_h_f_0.x + inp4_h_f_0.x); - inp1_h_f_0.y += (inp2_h_f_0.y + inp3_h_f_0.y + inp4_h_f_0.y); - inp1_h_f_1.x += (inp2_h_f_1.x + inp3_h_f_1.x + inp4_h_f_1.x); - inp1_h_f_1.y += (inp2_h_f_1.y + inp3_h_f_1.y + inp4_h_f_1.y); - - float2 val_f; - __half2* val_h = reinterpret_cast<__half2*>(&val_f); - - val_h[0] = __float22half2_rn(inp1_h_f_0); - val_h[1] = __float22half2_rn(inp1_h_f_1); - - float2* out_4 = reinterpret_cast(out); - out_4[row * row_stride + id] = val_f; -} - -template <> -void launch_fused_add4(float* out, - const float* inp1, - const float* inp2, - const float* inp3, - const float* inp4, - int batch_size, - int seq_length, - int hidden_size, - cudaStream_t& stream) -{ - dim3 grid_dim(batch_size * seq_length); - - dim3 block_dim(hidden_size / 4); - - fused_add4_kernel<<>>( - out, inp1, inp2, inp3, inp4, (batch_size * seq_length * hidden_size), hidden_size / 4); -} - -template <> -void launch_fused_add4<__half>(__half* out, - const __half* inp1, - const __half* inp2, - const __half* inp3, - const __half* inp4, - int batch_size, - int seq_length, - int hidden_size, - cudaStream_t& stream) -{ - dim3 grid_dim(batch_size * seq_length); - - dim3 block_dim(hidden_size / 4); - - fused_add4_kernel<<>>( - out, inp1, inp2, inp3, inp4, (batch_size * seq_length * hidden_size), hidden_size / 4); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu deleted file mode 100644 index bbb8a7f00b1f5a0be87a6bce7c4c65ff3ee84145..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu +++ /dev/null @@ -1,199 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "conversion_utils.h" -#ifdef __HIP_PLATFORM_AMD__ -#include "hip/hip_cooperative_groups.h" -#else -#include "cooperative_groups.h" -#endif -#include "ds_kernel_utils.h" -#include "inference_cuda_layers.h" -#include "memory_access_utils.h" - -#ifndef __HIP_PLATFORM_AMD__ -#include -#endif - -namespace cg = cooperative_groups; - -namespace rot_half { -constexpr int threads = 256; -} // namespace rot_half - -template -__global__ void apply_rotary_pos_half(T* mixed_query, - T* key_layer, - unsigned rotary_dim, - unsigned seq_len, - unsigned seq_offset, - unsigned num_heads, - unsigned head_size, - unsigned total_count, - float rope_theta, - int max_out_tokens) -{ - constexpr int T_per_thread = granularity / sizeof(T); - constexpr int heads_per_block = rot_half::threads / threadsPerHead; - - cg::thread_block tb = cg::this_thread_block(); - cg::thread_block_tile head_group = cg::tiled_partition(tb); - - const int head_idx = blockIdx.x * heads_per_block + threadIdx.x / threadsPerHead; - const int cur_seq_idx = head_idx % seq_len; - const int offset = head_idx * head_size; - const int k_offset = (cur_seq_idx + (head_idx / seq_len) * max_out_tokens) * head_size; - - const int seq_idx = cur_seq_idx + seq_offset; - const int half_dim = rotary_dim >> 1; - const int half_dim_threads = half_dim / T_per_thread; - - if (head_idx < total_count) { - const int base_neuron_idx = head_group.thread_rank() * T_per_thread; - - T q[T_per_thread], k[T_per_thread]; - mem_access::load_global(q, mixed_query + offset + base_neuron_idx); - mem_access::load_global(k, key_layer + k_offset + base_neuron_idx); - -#pragma unroll - for (int i = 0; i < T_per_thread; i++) { - const int neuron_idx = base_neuron_idx + i; - if (neuron_idx < rotary_dim) { - float inv_freq = (float)((neuron_idx % half_dim) * 2) / (float)rotary_dim; - inv_freq = 1.0 / powf(rope_theta, inv_freq) * (float)seq_idx; - - float rotary_sign = (neuron_idx > (half_dim - 1) ? -1.0 : 1.0); - float q_rot = conversion::to(q[i]) * rotary_sign; - float k_rot = conversion::to(k[i]) * rotary_sign; - - const int target_lane = (neuron_idx < half_dim) - ? head_group.thread_rank() + half_dim_threads - : head_group.thread_rank() - half_dim_threads; - - const float q_rot_temp = head_group.shfl(q_rot, target_lane); - const float k_rot_temp = head_group.shfl(k_rot, target_lane); - - q[i] = conversion::to(conversion::to(q[i]) * cosf(inv_freq) + - q_rot_temp * sinf(inv_freq)); - k[i] = conversion::to(conversion::to(k[i]) * cosf(inv_freq) + - k_rot_temp * sinf(inv_freq)); - } - } - - mem_access::store_global(mixed_query + offset + base_neuron_idx, q); - mem_access::store_global(key_layer + k_offset + base_neuron_idx, k); - } -} - -#define LAUNCH_ROT_POS_EMB_HALF(HEAD_THREADS, ALIGNMENT) \ - apply_rotary_pos_half<<>>(mixed_query, \ - key_layer, \ - rotary_dim, \ - seq_len, \ - offset, \ - num_heads, \ - head_size, \ - total_count, \ - rope_theta, \ - max_out_tokens); - -#if defined(__HIP_PLATFORM_AMD__) and ROCM_WAVEFRONT_SIZE == 64 -#define LAUNCH_FOR_ALIGNMENT(ALIGNMENT) \ - if (threads_per_head == 4) { \ - LAUNCH_ROT_POS_EMB_HALF(4, ALIGNMENT); \ - } else if (threads_per_head == 8) { \ - LAUNCH_ROT_POS_EMB_HALF(8, ALIGNMENT); \ - } else if (threads_per_head == 16) { \ - LAUNCH_ROT_POS_EMB_HALF(16, ALIGNMENT); \ - } else if (threads_per_head == 32) { \ - LAUNCH_ROT_POS_EMB_HALF(32, ALIGNMENT); \ - } else if (threads_per_head == 64) { \ - LAUNCH_ROT_POS_EMB_HALF(64, ALIGNMENT); \ - } else { \ - assert(false); \ - } -#else -#define LAUNCH_FOR_ALIGNMENT(ALIGNMENT) \ - if (threads_per_head == 4) { \ - LAUNCH_ROT_POS_EMB_HALF(4, ALIGNMENT); \ - } else if (threads_per_head == 8) { \ - LAUNCH_ROT_POS_EMB_HALF(8, ALIGNMENT); \ - } else if (threads_per_head == 16) { \ - LAUNCH_ROT_POS_EMB_HALF(16, ALIGNMENT); \ - } else if (threads_per_head == 32) { \ - LAUNCH_ROT_POS_EMB_HALF(32, ALIGNMENT); \ - } else { \ - assert(false); \ - } -#endif - -template -void launch_apply_rotary_pos_emb(T* mixed_query, - T* key_layer, - unsigned head_size, - unsigned seq_len, - unsigned rotary_dim, - unsigned offset, - unsigned num_heads, - unsigned batch, - float rope_theta, - cudaStream_t stream, - int max_out_tokens) -{ - const int half_dim = rotary_dim >> 1; - - int alignment = sizeof(T); - if (half_dim % (16 / sizeof(T)) == 0) { - alignment = 16; - } else if (half_dim % (8 / sizeof(T)) == 0) { - alignment = 8; - } else if (half_dim % (4 / sizeof(T)) == 0) { - alignment = 4; - } else { - assert(false); - } - const int T_per_elem = alignment / sizeof(T); - - int total_count = batch * num_heads * seq_len; - - const int padded_head_size = next_pow2(head_size); - - assert(padded_head_size <= hw_warp_size * T_per_elem); - - const int threads_per_head = padded_head_size / T_per_elem; - const int heads_per_block = rot_half::threads / threads_per_head; - - dim3 block(rot_half::threads); - dim3 grid((total_count + heads_per_block - 1) / heads_per_block); - - if (alignment == 4) { - LAUNCH_FOR_ALIGNMENT(4); - } else if (alignment == 8) { - LAUNCH_FOR_ALIGNMENT(8); - } else if (alignment == 16) { - LAUNCH_FOR_ALIGNMENT(16); - } else { - assert(false); - } -} - -#define INSTANTIATE_LAUNCH_ROTARY_POS_EMB(T) \ - template void launch_apply_rotary_pos_emb(T*, \ - T*, \ - unsigned, \ - unsigned, \ - unsigned, \ - unsigned, \ - unsigned, \ - unsigned, \ - float, \ - cudaStream_t, \ - int); - -INSTANTIATE_LAUNCH_ROTARY_POS_EMB(float); -#ifdef BF16_AVAILABLE -INSTANTIATE_LAUNCH_ROTARY_POS_EMB(__nv_bfloat16); -#endif -INSTANTIATE_LAUNCH_ROTARY_POS_EMB(__half); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/dequantize.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/dequantize.cu deleted file mode 100644 index 7a8e7ca446b0fbc4154fd0198c8e301f77cc7815..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/dequantize.cu +++ /dev/null @@ -1,153 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "conversion_utils.h" -#include "inference_cuda_layers.h" - -#define MAX_QUANTIZE_GROUPING 1024 - -#define loop_unroll 1 -#define loop_unroll_bits 1 - -template -__global__ void dequantize_kernel(T* output, - const int8_t* input, - const float* qscale, - int output_size, - int hidden_dim, - int groups, - int merge_count) -{ - unsigned merge_hidden = hidden_dim >> merge_count; - unsigned quantization_stride = (merge_hidden * output_size) / groups; - - unsigned bid = blockIdx.x; - unsigned tid = threadIdx.x; - - while (tid < output_size) { - unsigned w_index = bid / merge_hidden; - unsigned q_index = tid + bid * output_size; - - auto q = input[q_index]; - - unsigned merge_hidden_total = w_index * merge_hidden; - unsigned scale_index = - ((((bid - merge_hidden_total) + tid * merge_hidden) / quantization_stride) - << merge_count) + - w_index; - - float scale_data = qscale[scale_index]; - - output[q_index] = conversion::to(scale_data * (float)q); - tid += blockDim.x; - } -} - -template -void launch_dequantize(T* output, - const int8_t* input, - const float* qscale, - unsigned output_size, - unsigned hidden_dim, - unsigned groups, - unsigned merge_count, - cudaStream_t stream) -{ - unsigned threads = 1024; - dim3 block_dims(threads); - dim3 grid_dims(hidden_dim); - - dequantize_kernel<<>>( - output, input, qscale, output_size, hidden_dim, groups, merge_count); -} - -#define INSTANTIATE_DEQUANTIZE_MERGE(T) \ - template void launch_dequantize( \ - T*, const int8_t*, const float*, unsigned, unsigned, unsigned, unsigned, cudaStream_t); - -INSTANTIATE_DEQUANTIZE_MERGE(float); -#ifdef BF16_AVAILABLE -INSTANTIATE_DEQUANTIZE_MERGE(__nv_bfloat16); -#endif -INSTANTIATE_DEQUANTIZE_MERGE(__half); - -__global__ void dequantize_kernel(float* output, - const int8_t* input, - const float* qscale, - int hidden_dim, - unsigned merge_hidden, - int cnt) -{ -} - -template -__global__ void dequantize_kernel(T* output, - const int8_t* input, - const float* qscale, - unsigned hidden_dim, - unsigned merge_hidden, - int cnt) -{ - unsigned bid = blockIdx.x * gridDim.y + blockIdx.y; - unsigned tid = threadIdx.x; - - float local_scale = qscale[blockIdx.x]; - - const float* input_cast = reinterpret_cast(input); - float2* output_cast = reinterpret_cast(output); - - input_cast += bid * merge_hidden; - output_cast += bid * merge_hidden; - - for (int c = 0; c < cnt; c++) { - if (tid < merge_hidden) { - float q = input_cast[tid]; - int8_t* q_int8 = (int8_t*)&q; - - float2 q_f; - T* q_h = (T*)&q_f; - - q_h[0] = conversion::to(local_scale * (float)q_int8[0]); - q_h[1] = conversion::to(local_scale * (float)q_int8[1]); - q_h[2] = conversion::to(local_scale * (float)q_int8[2]); - q_h[3] = conversion::to(local_scale * (float)q_int8[3]); - output_cast[tid] = q_f; - tid += blockDim.x; - } - } -} - -template -void launch_dequantize(T* output, - const int8_t* input, - const float* qscale, - unsigned output_size, - unsigned hidden_dim, - unsigned groups, - cudaStream_t stream) -{ - unsigned threads = 1024; - hidden_dim /= 4; - unsigned thd_cnt = (hidden_dim - 1) / threads + 1; - - assert(output_size % groups == 0); - unsigned blocks = output_size / groups; - - dim3 block_dims(threads); - dim3 grid_dims(groups, blocks); - - dequantize_kernel<<>>( - output, input, qscale, hidden_dim, hidden_dim, thd_cnt); -} - -#define INSTANTIATE_DEQUANTIZE_NO_MERGE(T) \ - template void launch_dequantize( \ - T*, const int8_t*, const float*, unsigned, unsigned, unsigned, cudaStream_t); - -INSTANTIATE_DEQUANTIZE_NO_MERGE(float); -#ifdef BF16_AVAILABLE -INSTANTIATE_DEQUANTIZE_NO_MERGE(__nv_bfloat16); -#endif -INSTANTIATE_DEQUANTIZE_NO_MERGE(__half); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/gelu.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/gelu.cu deleted file mode 100644 index 97857bc3f70bfa044cbf26e4a82bf6dd781f75b9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/gelu.cu +++ /dev/null @@ -1,716 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "conversion_utils.h" -#include "inference_cuda_layers.h" -#include "memory_access_utils.h" - -namespace cg = cooperative_groups; -#define MAX_CAP 4 -#define MAX_SEQ 2048 - -// only used to avoid compilation error due to lack of definition. -#ifndef BF16_AVAILABLE -#if defined(__CUDA_BF16_H__) -static_assert(sizeof(__nv_bfloat162) == sizeof(__half2), - "CUDA's __nv_bfloat162 doesn't match __half2 size"); -#else -// Fallback to simple typedef only if CUDA doesn't provide it -using __nv_bfloat162 = __half2; -#endif -#endif - -inline __device__ float gelu(const float x) -{ - constexpr float sqrt_param = 0.79788456080286535587989211986876f; - constexpr float mul_param = 0.044715; - return x * 0.5f * (1.0f + tanhf(sqrt_param * (x + mul_param * x * x * x))); -} - -/* -In-place gelu(biasAdd(x)) for channels last -*/ -template -__global__ void fused_bias_gelu(T* input, const T* bias, int total_count, int intermediate_size) -{ - // Input restriction: intermediate_size % vals_per_access == 0 - constexpr int granularity = 16; - constexpr int values_per_access = granularity / sizeof(T); - const int offset = (blockIdx.x * blockDim.x + threadIdx.x) * values_per_access; - - if (offset < total_count) { - T data[values_per_access]; - T data_bias[values_per_access]; - mem_access::load_global(data, input + offset); - mem_access::load_global( - data_bias, bias + (offset % intermediate_size), bias != nullptr); - -#pragma unroll - for (int i = 0; i < values_per_access; i++) { - float data_f = conversion::to(data[i]); - float bias_f = conversion::to(data_bias[i]); - data[i] = conversion::to(gelu(data_f + bias_f)); - } - - mem_access::store_global(input + offset, data); - } -} - -template -void launch_bias_gelu(T* input, - const T* bias, - int intermediate_size, - int batch_size, - cudaStream_t stream) -{ - constexpr int threads = 1024; - constexpr int granularity = 16; - - const int total_count = batch_size * intermediate_size; - const int elems_per_block = threads * (granularity / sizeof(T)); - dim3 block_dims(threads); - dim3 grid_dims((total_count + elems_per_block - 1) / elems_per_block); - - fused_bias_gelu<<>>( - input, bias, total_count, intermediate_size); -} - -#define INSTANTIATE_LAUNCH_BIAS_GELU(T) \ - template void launch_bias_gelu(T*, const T*, int, int, cudaStream_t); - -INSTANTIATE_LAUNCH_BIAS_GELU(float) -#ifdef BF16_AVAILABLE -INSTANTIATE_LAUNCH_BIAS_GELU(__nv_bfloat16) -#endif -INSTANTIATE_LAUNCH_BIAS_GELU(__half) - -/* -In-place channels-last bias add -*/ -template -__global__ void fused_bias_add(T* input, const T* bias, int total_count, int intermediate_size) -{ - // Input restriction: intermediate_size % vals_per_access == 0 - constexpr int granularity = 16; - constexpr int values_per_access = granularity / sizeof(T); - const int offset = (blockIdx.x * blockDim.x + threadIdx.x) * values_per_access; - - if (offset < total_count) { - T data[values_per_access]; - T data_bias[values_per_access]; - mem_access::load_global(data, input + offset); - mem_access::load_global( - data_bias, bias + (offset % intermediate_size), bias != nullptr); - -#pragma unroll - for (int i = 0; i < values_per_access; i++) { - float data_f = conversion::to(data[i]); - float bias_f = conversion::to(data_bias[i]); - data[i] = conversion::to(data_f + bias_f); - } - - mem_access::store_global(input + offset, data); - } -} - -template -void launch_bias_add(T* input, - const T* bias, - int intermediate_size, - int batch_size, - cudaStream_t stream) -{ - constexpr int threads = 1024; - constexpr int granularity = 16; - - const int total_count = batch_size * intermediate_size; - const int elems_per_block = threads * (granularity / sizeof(T)); - dim3 block_dims(threads); - dim3 grid_dims((total_count + elems_per_block - 1) / elems_per_block); - - fused_bias_add<<>>( - input, bias, total_count, intermediate_size); -} - -#define INSTANTIATE_LAUNCH_BIAS_ADD(T) \ - template void launch_bias_add(T*, const T*, int, int, cudaStream_t); - -INSTANTIATE_LAUNCH_BIAS_ADD(float) -#ifdef BF16_AVAILABLE -INSTANTIATE_LAUNCH_BIAS_ADD(__nv_bfloat16) -#endif -INSTANTIATE_LAUNCH_BIAS_ADD(__half) - -__global__ void fused_bias_residual(float* residual, - const float* hidden_state, - const float* attn, - const float* bias, - const float* attn_bias, - const int total_count, - const int intermediate_size, - const float mp_scale, - const bool preln) -{ - float4* res_fl4_ptr = reinterpret_cast(residual); - const float4* hs_fl4_ptr = reinterpret_cast(hidden_state); - const float4* attn_fl4_ptr = reinterpret_cast(attn); - const float4* bias_fl4_ptr = reinterpret_cast(bias); - const float4* attn_bias_fl4_ptr = reinterpret_cast(attn_bias); - const int offset = blockIdx.x * blockDim.x + threadIdx.x; - - if (offset < total_count) { - float4 res_fl4 = res_fl4_ptr[offset]; - const float4 hs_fl4 = hs_fl4_ptr[offset]; - const float4 attn_fl4 = attn_fl4_ptr[offset]; - const float4 bias_fl4 = bias_fl4_ptr[offset % intermediate_size]; - const float4 attn_bias_fl4 = attn_bias_fl4_ptr[offset % intermediate_size]; - if (preln) { - // residual = (residual + attention + bias + attention_bias) * - // mp_scale + hidden_state - res_fl4.x = - (res_fl4.x + attn_fl4.x + bias_fl4.x + attn_bias_fl4.x) * mp_scale + (hs_fl4.x); - res_fl4.y = - (res_fl4.y + attn_fl4.y + bias_fl4.y + attn_bias_fl4.y) * mp_scale + (hs_fl4.y); - res_fl4.z = - (res_fl4.z + attn_fl4.z + bias_fl4.z + attn_bias_fl4.z) * mp_scale + (hs_fl4.z); - res_fl4.w = - (res_fl4.w + attn_fl4.w + bias_fl4.w + attn_bias_fl4.w) * mp_scale + (hs_fl4.w); - } else { - // residual += hidden_state + bias - res_fl4.x = res_fl4.x + hs_fl4.x + bias_fl4.x; - res_fl4.y = res_fl4.y + hs_fl4.y + bias_fl4.y; - res_fl4.z = res_fl4.z + hs_fl4.z + bias_fl4.z; - res_fl4.w = res_fl4.w + hs_fl4.w + bias_fl4.w; - } - res_fl4_ptr[offset] = res_fl4; - } -} - -template -__global__ void fused_bias_residual(T* residual, - const T* hidden_state, - const T* attn, - const T* bias, - const T* attn_bias, - const int total_count, - const int intermediate_size, - const float mp_scale, - const bool preln) -{ - using T2 = - typename std::conditional::value, __half2, __nv_bfloat162>::type; - float2* res_fl2_ptr = reinterpret_cast(residual); - const float2* hs_fl2_ptr = reinterpret_cast(hidden_state); - const float2* attn_fl2_ptr = reinterpret_cast(attn); - const float2* bias_fl2_ptr = reinterpret_cast(bias); - const float2* attn_bias_fl2_ptr = reinterpret_cast(attn_bias); - const int offset = blockIdx.x * blockDim.x + threadIdx.x; - - if (offset < total_count) { - float2 res_fl2 = res_fl2_ptr[offset]; - const float2 hs_fl2 = hs_fl2_ptr[offset]; - const float2 attn_fl2 = attn_fl2_ptr[offset]; - const float2 bias_fl2 = bias_fl2_ptr[offset % intermediate_size]; - const float2 attn_bias_fl2 = attn_bias_fl2_ptr[offset % intermediate_size]; - - T2* res_half2 = reinterpret_cast(&res_fl2); - const T2* hs_half2 = reinterpret_cast(&hs_fl2); - const T2* attn_half2 = reinterpret_cast(&attn_fl2); - const T2* bias_half2 = reinterpret_cast(&bias_fl2); - const T2* attn_bias_half2 = reinterpret_cast(&attn_bias_fl2); - - float2 res_low = conversion::to(res_half2[0]); - float2 res_high = conversion::to(res_half2[1]); - - const float2 hs_low = conversion::to(hs_half2[0]); - const float2 hs_high = conversion::to(hs_half2[1]); - - const float2 attn_low = conversion::to(attn_half2[0]); - const float2 attn_high = conversion::to(attn_half2[1]); - - const float2 bias_low = conversion::to(bias_half2[0]); - const float2 bias_high = conversion::to(bias_half2[1]); - - const float2 attn_bias_low = conversion::to(attn_bias_half2[0]); - const float2 attn_bias_high = conversion::to(attn_bias_half2[1]); - - if (preln) { - // residual = (residual + attention + bias + attention_bias) * - // mp_scale + hidden_state - res_low.x = - (res_low.x + attn_low.x + bias_low.x + attn_bias_low.x) * mp_scale + hs_low.x; - res_low.y = - (res_low.y + attn_low.y + bias_low.y + attn_bias_low.y) * mp_scale + hs_low.y; - res_high.x = - (res_high.x + attn_high.x + bias_high.x + attn_bias_high.x) * mp_scale + hs_high.x; - res_high.y = - (res_high.y + attn_high.y + bias_high.y + attn_bias_high.y) * mp_scale + hs_high.y; - } else { - // residual += hidden_state + bias - res_low.x = (res_low.x + hs_low.x + bias_low.x); - res_low.y = (res_low.y + hs_low.y + bias_low.y); - res_high.x = (res_high.x + hs_high.x + bias_high.x); - res_high.y = (res_high.y + hs_high.y + bias_high.y); - } - res_half2[0] = conversion::to(res_low); - res_half2[1] = conversion::to(res_high); - - res_fl2_ptr[offset] = res_fl2; - } -} - -template -void launch_bias_residual(T* residual, - T* hidden_state, - T* attn, - T* bias, - T* attn_bias, - int batch, - int hidden_dim, - int mp_size, - bool preln, - cudaStream_t stream) -{ - int total_count = batch * hidden_dim / 4; - dim3 block_dims(1024); - dim3 grid_dims((total_count - 1) / 1024 + 1); // (batch_size); - - fused_bias_residual<<>>(residual, - hidden_state, - attn, - bias, - attn_bias, - total_count, - hidden_dim / 4, - 1.0 / mp_size, - preln); -} - -#define INSTANTIATE_LAUNCH_BIAS_RESIDUAL(T) \ - template void launch_bias_residual(T*, T*, T*, T*, T*, int, int, int, bool, cudaStream_t); - -INSTANTIATE_LAUNCH_BIAS_RESIDUAL(float); -#ifdef BF16_AVAILABLE -INSTANTIATE_LAUNCH_BIAS_RESIDUAL(__nv_bfloat16); -#endif -INSTANTIATE_LAUNCH_BIAS_RESIDUAL(__half); - -__global__ void gptj_residual_add(float* residual, - const float* hidden_state, - const float* attn, - const float* bias, - const float* attn_bias, - const int total_count, - const int intermediate_size, - const float mp_scale) -{ - float4* res_fl4_ptr = reinterpret_cast(residual); - const float4* hs_fl4_ptr = reinterpret_cast(hidden_state); - const float4* attn_fl4_ptr = reinterpret_cast(attn); - const float4* bias_fl4_ptr = reinterpret_cast(bias); - const float4* attn_bias_fl4_ptr = reinterpret_cast(attn_bias); - const int offset = blockIdx.x * blockDim.x + threadIdx.x; - - if (offset < total_count) { - float4 res_fl4 = res_fl4_ptr[offset]; - const float4 hs_fl4 = hs_fl4_ptr[offset]; - const float4 attn_fl4 = attn_fl4_ptr[offset]; - const float4 bias_fl4 = bias_fl4_ptr[offset % intermediate_size]; - - if (attn_bias) { - float4 attn_bias_fl4 = attn_bias_fl4_ptr[offset % intermediate_size]; - // residual += attention_bias - res_fl4.x += attn_bias_fl4.x; - res_fl4.y += attn_bias_fl4.y; - res_fl4.z += attn_bias_fl4.z; - res_fl4.w += attn_bias_fl4.w; - } - // residual = hidden_state + attention + (residual + bias) * mp_scale - res_fl4.x = hs_fl4.x + attn_fl4.x + (res_fl4.x + bias_fl4.x) * mp_scale; - res_fl4.y = hs_fl4.y + attn_fl4.y + (res_fl4.y + bias_fl4.y) * mp_scale; - res_fl4.z = hs_fl4.z + attn_fl4.z + (res_fl4.z + bias_fl4.z) * mp_scale; - res_fl4.w = hs_fl4.w + attn_fl4.w + (res_fl4.w + bias_fl4.w) * mp_scale; - - res_fl4_ptr[offset] = res_fl4; - } -} - -template -__global__ void gptj_residual_add(T* residual, - const T* hidden_state, - const T* attn, - const T* bias, - const T* attn_bias, - const int total_count, - const int intermediate_size, - const float mp_scale) -{ - using T2 = - typename std::conditional::value, __half2, __nv_bfloat162>::type; - float2* res_fl2_ptr = reinterpret_cast(residual); - const float2* hs_fl2_ptr = reinterpret_cast(hidden_state); - const float2* attn_fl2_ptr = reinterpret_cast(attn); - const float2* bias_fl2_ptr = reinterpret_cast(bias); - const float2* attn_bias_fl2_ptr = reinterpret_cast(attn_bias); - const int offset = blockIdx.x * blockDim.x + threadIdx.x; - - if (offset < total_count) { - float2 res_fl2 = res_fl2_ptr[offset]; - const float2 hs_fl2 = hs_fl2_ptr[offset]; - const float2 attn_fl2 = attn_fl2_ptr[offset]; - const float2 bias_fl2 = bias_fl2_ptr[offset % intermediate_size]; - - T2* res_half2 = reinterpret_cast(&res_fl2); - const T2* hs_half2 = reinterpret_cast(&hs_fl2); - const T2* attn_half2 = reinterpret_cast(&attn_fl2); - const T2* bias_half2 = reinterpret_cast(&bias_fl2); - - float2 res_low = conversion::to(res_half2[0]); - float2 res_high = conversion::to(res_half2[1]); - - const float2 hs_low = conversion::to(hs_half2[0]); - const float2 hs_high = conversion::to(hs_half2[1]); - - const float2 attn_low = conversion::to(attn_half2[0]); - const float2 attn_high = conversion::to(attn_half2[1]); - - const float2 bias_low = conversion::to(bias_half2[0]); - const float2 bias_high = conversion::to(bias_half2[1]); - - if (attn_bias) { - const float2 attn_bias_fl2 = attn_bias_fl2_ptr[offset % intermediate_size]; - const T2* attn_bias_half2 = reinterpret_cast(&attn_bias_fl2); - const float2 attn_bias_low = conversion::to(attn_bias_half2[0]); - const float2 attn_bias_high = conversion::to(attn_bias_half2[1]); - // residual += attention_bias - res_low.x += attn_bias_low.x; - res_low.y += attn_bias_low.y; - res_high.x += attn_bias_high.x; - res_high.y += attn_bias_high.y; - } - // residual = hidden_state + attention + (residual + bias) * mp_scale - res_low.x = attn_low.x + hs_low.x + (res_low.x + bias_low.x) * mp_scale; - res_low.y = attn_low.y + hs_low.y + (res_low.y + bias_low.y) * mp_scale; - res_high.x = attn_high.x + hs_high.x + (res_high.x + bias_high.x) * mp_scale; - res_high.y = attn_high.y + hs_high.y + (res_high.y + bias_high.y) * mp_scale; - - res_half2[0] = conversion::to(res_low); - res_half2[1] = conversion::to(res_high); - - res_fl2_ptr[offset] = res_fl2; - } -} - -template -void launch_gptj_residual_add(T* residual, - T* hidden_state, - T* attn, - T* bias, - T* attn_bias, - int hidden_dim, - int batch, - int mp_size, - cudaStream_t stream) -{ - int total_count = batch * hidden_dim / 4; - dim3 block_dims(1024); - dim3 grid_dims((total_count - 1) / 1024 + 1); // (batch_size); - - gptj_residual_add<<>>( - residual, hidden_state, attn, bias, attn_bias, total_count, hidden_dim / 4, 1.0 / mp_size); -} - -#define INSTANTIATE_GPT_RES_ADD(T) \ - template void launch_gptj_residual_add(T*, T*, T*, T*, T*, int, int, int, cudaStream_t); - -INSTANTIATE_GPT_RES_ADD(float); -INSTANTIATE_GPT_RES_ADD(__half); -#ifdef BF16_AVAILABLE -INSTANTIATE_GPT_RES_ADD(__nv_bfloat16); -#endif - -template -__global__ void moe_res_matmul(T* residual, T* coef, T* mlp_out, int seq_len, int hidden_dim) -{ - constexpr int granularity = 16; - constexpr int vals_per_access = granularity / sizeof(T); - - T* residual_seq = residual + blockIdx.x * hidden_dim; - T* mlp_out_seq = mlp_out + blockIdx.x * hidden_dim; - - for (unsigned tid = threadIdx.x * vals_per_access; tid < hidden_dim; - tid += blockDim.x * vals_per_access) { - T mlp[vals_per_access]; - T res[vals_per_access]; - T coef1[vals_per_access]; - T coef2[vals_per_access]; - - mem_access::load_global(mlp, mlp_out_seq + tid); - mem_access::load_global(res, residual_seq + tid); - mem_access::load_global(coef1, coef + tid); - mem_access::load_global(coef2, coef + tid + hidden_dim); - -#pragma unroll - for (int idx = 0; idx < vals_per_access; idx++) { - mlp[idx] = mlp[idx] * coef2[idx] + res[idx] * coef1[idx]; - } - - mem_access::store_global(mlp_out_seq + tid, mlp); - } -} - -template -void launch_moe_res_matmul(T* residual, - T* coef, - T* mlp_out, - int seq_len, - int hidden_dim, - cudaStream_t stream) -{ - dim3 grid_dim(seq_len); - dim3 block_dim(1024); - moe_res_matmul<<>>( - residual, coef, mlp_out, seq_len, hidden_dim); -} - -#define INSTANTIATE_LAUNCH_MOE_RES_MATMUL(T) \ - template void launch_moe_res_matmul(T*, T*, T*, int, int, cudaStream_t); - -INSTANTIATE_LAUNCH_MOE_RES_MATMUL(float); -#ifdef BF16_AVAILABLE -INSTANTIATE_LAUNCH_MOE_RES_MATMUL(__nv_bfloat16); -#endif -INSTANTIATE_LAUNCH_MOE_RES_MATMUL(__half); - -template -__global__ void pad_data_kernel(T* padded_output, T* output, int head_size, int padded_head_size) -{ - using T2 = - typename std::conditional::value, __half2, __nv_bfloat162>::type; - float4* padded_output_cast = reinterpret_cast(padded_output); - float4* output_cast = reinterpret_cast(output); - int bid = blockIdx.x * (blockDim.y) + threadIdx.y; - int idx = threadIdx.x; - padded_output_cast += (bid * padded_head_size); - output_cast += (bid * head_size); - float4 ZERO; - const T2 zero_h = conversion::to(0.f); - T2* ZERO_h = reinterpret_cast(&ZERO); -#pragma unroll - for (int i = 0; i < 4; i++) ZERO_h[i] = zero_h; - if (idx < head_size) - padded_output_cast[idx] = output_cast[idx]; - else - padded_output_cast[idx] = ZERO; -} - -__global__ void pad_data_kernel(float* padded_output, - float* output, - int head_size, - int padded_head_size) -{ -} - -template -void pad_data(T* padded_output, - T* output, - int bsz, - int head_size, - int padded_head_size, - cudaStream_t stream) -{ - dim3 grid_dim((bsz - 1) / 16 + 1); - dim3 block_dim(padded_head_size / 8, 16); - pad_data_kernel<<>>( - padded_output, output, head_size / 8, padded_head_size / 8); -} - -#define INSTANTIATE_PAD_DATA(T) template void pad_data(T*, T*, int, int, int, cudaStream_t stream); - -INSTANTIATE_PAD_DATA(float); -INSTANTIATE_PAD_DATA(__half); -#ifdef BF16_AVAILABLE -INSTANTIATE_PAD_DATA(__nv_bfloat16); -#endif - -template -__global__ void pad_head_seq_kernel(T* padded_output, - T* output, - int seq_len, - int padded_seq_len, - int head_size, - int padded_head_size) -{ - using T2 = - typename std::conditional::value, __half2, __nv_bfloat162>::type; - float4* padded_output_cast = reinterpret_cast(padded_output); - float4* output_cast = reinterpret_cast(output); - int bsz = blockIdx.x; - int bid = blockIdx.y * (blockDim.y) + threadIdx.y; - int idx = threadIdx.x; - padded_output_cast += (bsz * padded_seq_len + bid) * padded_head_size; - output_cast += (bsz * seq_len + bid) * head_size; - float4 ZERO; - const T2 zero_h = conversion::to(0.f); - T2* ZERO_h = reinterpret_cast(&ZERO); -#pragma unroll - for (int i = 0; i < 4; i++) ZERO_h[i] = zero_h; - - if (idx < head_size && bid < seq_len) - padded_output_cast[idx] = output_cast[idx]; - else - padded_output_cast[idx] = ZERO; -} - -__global__ void pad_head_seq_kernel(float* padded_output, - float* output, - int seq_len, - int padded_seq_len, - int head_size, - int padded_head_size) -{ -} - -template -void pad_head_seq(T* padded_output, - T* output, - int bsz, - int seq_len, - int padded_seq_len, - int head_size, - int padded_head_size, - cudaStream_t stream) -{ - dim3 grid_dim(bsz, padded_seq_len / 16); - dim3 block_dim(padded_head_size / 8, 16); - pad_head_seq_kernel<<>>( - padded_output, output, seq_len, padded_seq_len, head_size / 8, padded_head_size / 8); -} - -#define INSTANTIATE_PAD_HEAD_SEQ(T) \ - template void pad_head_seq(T*, T*, int, int, int, int, int, cudaStream_t); - -INSTANTIATE_PAD_HEAD_SEQ(__half); -#ifdef BF16_AVAILABLE -INSTANTIATE_PAD_HEAD_SEQ(__nv_bfloat16); -#endif -INSTANTIATE_PAD_HEAD_SEQ(float); - -// TODO(cmikeh2): evaluate different GeLU performance -__device__ __forceinline__ float old_gelu(float val) -{ - // 1 / sqrt(2) - constexpr float rsqrt_2 = 0.707106769084930419922; - return val * 0.5f * (1.0f + erff(val * rsqrt_2)); -} - -namespace fused_geglu { -constexpr int threads = 256; -constexpr int steps = 2; -constexpr int granularity = 16; -} // namespace fused_geglu - -__device__ __forceinline__ float silu(float val) { return val / (1.0f + expf(-val)); } - -template -__global__ void fused_gate_activation(T* output, - const T* activation, - const T* bias, - int base_channels, - int output_stride, - int total_elems) -{ - constexpr int T_per_access = fused_geglu::granularity / sizeof(T); - constexpr int T_per_step = T_per_access * fused_geglu::threads; - constexpr int T_per_block = T_per_step * fused_geglu::steps; - - const int id = blockIdx.x * T_per_block + threadIdx.x * T_per_access; - -#pragma unroll - for (int i = 0; i < fused_geglu::steps; i++) { - T activation_buffer_1[T_per_access]; - T activation_buffer_2[T_per_access]; - T bias_buffer_1[T_per_access]; - T bias_buffer_2[T_per_access]; - - const int iter_id = id + T_per_step * i; - if (iter_id < total_elems) { - const int channel_id = iter_id % base_channels; - const int seq_id = iter_id / base_channels; - const int seq_offset = seq_id * base_channels * 2; - - mem_access::load_global(activation_buffer_1, - activation + seq_offset + channel_id); - mem_access::load_global( - activation_buffer_2, activation + seq_offset + channel_id + base_channels); - mem_access::load_global( - bias_buffer_1, bias + channel_id, bias != nullptr); - mem_access::load_global( - bias_buffer_2, bias + channel_id + base_channels, bias != nullptr); - - // Since the GeLU is going to happen at float, might as well - // convert -#pragma unroll - for (int v = 0; v < T_per_access; v++) { - T hidden_state = activation_buffer_1[v] + bias_buffer_1[v]; - T pre_gate = activation_buffer_2[v] + bias_buffer_2[v]; - float pre_gate_f = conversion::to(pre_gate); - float gate_f = (useGelu) ? old_gelu(pre_gate_f) : silu(pre_gate_f); - T gate = conversion::to(gate_f); - activation_buffer_1[v] = hidden_state * gate; - } - - mem_access::store_global( - output + seq_id * output_stride + channel_id, activation_buffer_1); - } - } -} - -template -void launch_gated_activation(T* output, - const T* activation, - const T* bias, - int rows, - int output_stride, - int elems_per_row, - bool use_gelu, - cudaStream_t stream) -{ - /* - Fused bias GEGLU is a variant of the gated activation functions. - The input here is a matrix of [batch, seq_len, 2 * intermediate_dim] - where the second half of the channels act as GeLU gates for the first - half. - */ - - // Re-derive the above figures - constexpr int T_per_access = fused_geglu::granularity / sizeof(T); - constexpr int T_per_step = T_per_access * fused_geglu::threads; - constexpr int T_per_block = T_per_step * fused_geglu::steps; - - const int base_channels = elems_per_row / 2; - const int total_elems = base_channels * rows; - - dim3 block(fused_geglu::threads); - dim3 grid((total_elems + T_per_block - 1) / T_per_block); - - if (use_gelu) { - fused_gate_activation<<>>( - output, activation, bias, base_channels, output_stride, total_elems); - } else { - fused_gate_activation<<>>( - output, activation, bias, base_channels, output_stride, total_elems); - } -} - -#define INSTANTIATE_LAUNCH_GATED_ACTIVATION(T) \ - template void launch_gated_activation( \ - T*, const T*, const T*, int, int, int, bool, cudaStream_t); - -INSTANTIATE_LAUNCH_GATED_ACTIVATION(__half); -#ifdef BF16_AVAILABLE -INSTANTIATE_LAUNCH_GATED_ACTIVATION(__nv_bfloat16); -#endif -INSTANTIATE_LAUNCH_GATED_ACTIVATION(float); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/layer_norm.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/layer_norm.cu deleted file mode 100644 index e5e7e89c9d2069678bc445004053f8ad5f9809b1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/layer_norm.cu +++ /dev/null @@ -1,503 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "conversion_utils.h" -#include "ds_kernel_utils.h" -#include "inference_cuda_layers.h" -#include "memory_access_utils.h" -#include "reduction_utils.h" - -namespace cg = cooperative_groups; -using rop = reduce::ROpType; - -namespace ln { -constexpr int granularity = 16; -} // namespace ln - -/* -Primary layer norm implementation. Assumes elems_per_row % 8 -is equal to 0. - -Args: - output: buffer for output data - vals: buffer for input data - gamma: gain for normalization - beta: bias for normalization - epsilon: numeric stability - elems_per_row: number of elements each block will normalize -*/ -template -__global__ void fused_ln(T* output, - const T* vals, - const T* gamma, - const T* beta, - float epsilon, - int elems_per_row) -{ - constexpr int T_per_load = ln::granularity / sizeof(T); - - cg::thread_block tb = cg::this_thread_block(); - cg::thread_block_tile warp = cg::tiled_partition(tb); - - // X-dimension of the block - const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) + - (tb.thread_index().y * elems_per_row); - const int thread_offset = tb.thread_index().x * T_per_load; - const int base_offset = block_offset + thread_offset; - const int stride = blockDim.x * T_per_load; - - float sum = reduce::init(); - - const T* input_base = vals + base_offset; - - T local_buffer[unRoll * T_per_load]; - -#pragma unRoll - for (int i = 0; i < unRoll; i++) { - T* iteration_buffer = local_buffer + i * T_per_load; - - mem_access::load_global( - iteration_buffer, input_base + i * stride, thread_offset + i * stride < elems_per_row); - -#pragma unRoll - for (int j = 0; j < T_per_load; j++) { - float vals_up_cast = conversion::to(iteration_buffer[j]); - sum = reduce::element(sum, vals_up_cast); - } - } - - reduce::partitioned_block(tb, warp, sum); - const float mean = sum / elems_per_row; - - float mean_diff = reduce::init(); - -#pragma unRoll - for (int i = 0; i < unRoll; i++) { -#pragma unRoll - for (int j = 0; j < T_per_load; j++) { - // Using a 0 value here skews the variance, have to if-guard - if (thread_offset + i * stride < elems_per_row) { - float diff = (conversion::to(local_buffer[i * T_per_load + j]) - mean); - mean_diff = reduce::element(mean_diff, diff * diff); - } - } - } - - reduce::partitioned_block(tb, warp, mean_diff); - const float variance = mean_diff / elems_per_row; - const float denom = __frsqrt_rn(variance + epsilon); - - // const T mean_compute = conversion::to(mean); - // const T denom_compute = conversion::to(denom); - - T* block_output = output + block_offset; - -#pragma unRoll - for (int i = 0; i < unRoll; i++) { - T* iteration_buffer = local_buffer + i * T_per_load; - const int iter_idx = i * stride + thread_offset; - const bool do_loads = iter_idx < elems_per_row; - - T gamma_local[T_per_load], beta_local[T_per_load]; - - mem_access::load_global(gamma_local, gamma + iter_idx, do_loads); - mem_access::load_global(beta_local, beta + iter_idx, do_loads); - -#pragma unRoll - for (int j = 0; j < T_per_load; j++) { - float val = conversion::to(iteration_buffer[j]); - val = (val - mean) * denom; - val = - val * conversion::to(gamma_local[j]) + conversion::to(beta_local[j]); - iteration_buffer[j] = conversion::to(val); - } - - if (do_loads) { - mem_access::store_global(block_output + iter_idx, iteration_buffer); - } - } -} - -#define LAUNCH_FUSED_LN(unRollFactor, threadsPerGroup, maxThreads) \ - fused_ln \ - <<>>(output, vals, gamma, beta, epsilon, elems_per_row); - -template -void launch_fused_ln(T* output, - const T* vals, - const T* gamma, - const T* beta, - float epsilon, - int rows, - int elems_per_row, - cudaStream_t stream) -{ - // 8 for __half, 4 for float - constexpr int T_per_load = ln::granularity / sizeof(T); - - constexpr int maxThreads = 256; - - // For Flaoat, unRoll 4, for __half, unRoll 2 - constexpr int internal_unRoll = sizeof(T) == 4 ? 4 : 2; - - const bool is_subblock_schedule = (elems_per_row <= 128) ? true : false; - const int h_per_step = is_subblock_schedule ? T_per_load : T_per_load * internal_unRoll; - - // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of - // warp-sized blocks rather than stepping up to 64/96 threads - const int one_step_threads = next_pow2((elems_per_row + h_per_step - 1) / h_per_step); - const int threadsPerGroup = (one_step_threads < maxThreads) ? one_step_threads : maxThreads; - - const int groups_per_block_max = - is_subblock_schedule ? (maxThreads + threadsPerGroup - 1) / threadsPerGroup : 1; - const int groups_per_block = (rows < groups_per_block_max) ? rows : groups_per_block_max; - const int groups_launch = (groups_per_block + rows - 1) / groups_per_block; - - dim3 block(threadsPerGroup, groups_per_block); - dim3 grid(groups_launch); - - const int elems_per_step = threadsPerGroup * h_per_step; - const int external_unRoll = (elems_per_row + elems_per_step - 1) / elems_per_step; - - if (is_subblock_schedule) { - // <=128 - if (threadsPerGroup == 1) { - LAUNCH_FUSED_LN(1, 1, maxThreads); - } else if (threadsPerGroup == 2) { - LAUNCH_FUSED_LN(1, 2, maxThreads); - } else if (threadsPerGroup == 4) { - LAUNCH_FUSED_LN(1, 4, maxThreads); - } else if (threadsPerGroup == 8) { - LAUNCH_FUSED_LN(1, 8, maxThreads); - } else if (threadsPerGroup == 16) { - LAUNCH_FUSED_LN(1, 16, maxThreads); - } - } else if (external_unRoll == 1) { - // 129 - 4096 elems - // (this can launch with 1-7 warps as well) - LAUNCH_FUSED_LN(1 * internal_unRoll, maxThreads, maxThreads); - } else if (external_unRoll == 2) { - // 4097 - 8192 elems - LAUNCH_FUSED_LN(2 * internal_unRoll, maxThreads, maxThreads); - } else if (external_unRoll == 3) { - // 8193 - 12288 elems - LAUNCH_FUSED_LN(3 * internal_unRoll, maxThreads, maxThreads); - } else if (external_unRoll == 4) { - // 12289 - 16384 elems - LAUNCH_FUSED_LN(4 * internal_unRoll, maxThreads, maxThreads); - } -} - -#define INSTANTIATE_FUSED_LN(T) \ - template void launch_fused_ln(T*, const T*, const T*, const T*, float, int, int, cudaStream_t); - -INSTANTIATE_FUSED_LN(__half); -#ifdef BF16_AVAILABLE -INSTANTIATE_FUSED_LN(__nv_bfloat16); -#endif -INSTANTIATE_FUSED_LN(float); - -/* -Fused resiual + bias + layer norm implementation. Assumes elems_per_row % 8 -is equal to 0. - -TODO(cmikeh2): Goal is to deprecate this implementation. The bias + residual -need to be fused into compute-bound producer operations. - -Args: - output: buffer for output data - res_output: output of residual addition - vals: buffer for input data - residual: residual data - bias: bias of of input data - gamma: gain for normalization - beta: bias for normalization - epsilon: numeric stability - elems_per_row: number of elements each block will normalize -Template arg: - StoreResidual: controls whether the residual calculation is stored - or not. When set to false, the input `res_output` is unused. -*/ -template -__global__ void fused_residual_ln(T* output, - T* res_output, - const T* vals, - const T* residual, - const T* bias, - const T* gamma, - const T* beta, - float epsilon, - int elems_per_row) -{ - constexpr int T_per_load = ln::granularity / sizeof(T); - - cg::thread_block tb = cg::this_thread_block(); - cg::thread_block_tile warp = cg::tiled_partition(tb); - - // X-dimension of the block - const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) + - (tb.thread_index().y * elems_per_row); - const int thread_offset = tb.thread_index().x * T_per_load; - const int base_offset = block_offset + thread_offset; - const int stride = tb.size() * T_per_load; - - float sum = reduce::init(); - - const T* input_base = vals + base_offset; - const T* residual_base = residual + base_offset; - const T* bias_base = bias + thread_offset; - - T local_buffer[unRoll * T_per_load]; - - // Unlike a vanilla layernorm, since we're fusing the two adds as well - // an inner unRoll seems to be less valuable. If anything, a double unRoll - // makes the most sense if we find we are having performance issues. -#pragma unRoll - for (int i = 0; i < unRoll; i++) { - T* iteration_buffer = local_buffer + i * T_per_load; - T residual_buffer[T_per_load]; - T bias_buffer[T_per_load]; - - mem_access::load_global( - iteration_buffer, input_base + i * stride, thread_offset + i * stride < elems_per_row); - mem_access::load_global(residual_buffer, - residual_base + i * stride, - thread_offset + i * stride < elems_per_row); - mem_access::load_global( - bias_buffer, bias_base + i * stride, thread_offset + i * stride < elems_per_row); - -#pragma unRoll - for (int j = 0; j < T_per_load; j++) { - float vals_up_cast = conversion::to(iteration_buffer[j]); - float res_up_cast = conversion::to(residual_buffer[j]); - float bias_up_cast = conversion::to(bias_buffer[j]); - vals_up_cast = vals_up_cast + bias_up_cast + res_up_cast; - sum = reduce::element(sum, vals_up_cast); - iteration_buffer[j] = conversion::to(vals_up_cast); - } - - if (preLnResidual && (thread_offset + i * stride < elems_per_row)) { - mem_access::store_global(res_output + base_offset + i * stride, - iteration_buffer); - } - } - - reduce::partitioned_block(tb, warp, sum); - const float mean = sum / elems_per_row; - - float mean_diff = reduce::init(); -#pragma unRoll - for (int i = 0; i < unRoll; i++) { -#pragma unRoll - for (int j = 0; j < T_per_load; j++) { - // Using a 0 value here skews the variance, have to if-guard - if (thread_offset + i * stride < elems_per_row) { - float diff = (conversion::to(local_buffer[i * T_per_load + j]) - mean); - mean_diff = reduce::element(mean_diff, diff * diff); - } - } - } - - reduce::partitioned_block(tb, warp, mean_diff); - const float variance = mean_diff / elems_per_row; - const float denom = __frsqrt_rn(variance + epsilon); - - T* block_output = output + block_offset; - -#pragma unRoll - for (int i = 0; i < unRoll; i++) { - T* iteration_buffer = local_buffer + i * T_per_load; - const int iter_idx = i * stride + thread_offset; - const bool do_loads = iter_idx < elems_per_row; - - T gamma_local[T_per_load], beta_local[T_per_load]; - - mem_access::load_global(gamma_local, gamma + iter_idx, do_loads); - mem_access::load_global(beta_local, beta + iter_idx, do_loads); - -#pragma unRoll - for (int j = 0; j < T_per_load; j++) { - // iteration_buffer[j] = (iteration_buffer[j] - mean_compute) * denom_compute; - // iteration_buffer[j] = iteration_buffer[j] * gamma_local[j] + beta_local[j]; - float val = conversion::to(iteration_buffer[j]); - val = (val - mean) * denom; - val = - val * conversion::to(gamma_local[j]) + conversion::to(beta_local[j]); - iteration_buffer[j] = conversion::to(val); - } - - if (do_loads) { - mem_access::store_global(block_output + iter_idx, iteration_buffer); - } - } -} - -// TODO(cmikeh2): There's a bunch of redundancy here that needs to be removed/simplified. -#define LAUNCH_FUSED_RES_LN(unRollFactor, threadsPerGroup, maxThreads) \ - fused_residual_ln \ - <<>>( \ - output, nullptr, vals, residual, bias, gamma, beta, epsilon, elems_per_row); - -template -void launch_fused_residual_ln(T* output, - const T* vals, - const T* residual, - const T* bias, - const T* gamma, - const T* beta, - float epsilon, - int rows, - int elems_per_row, - cudaStream_t stream) -{ - // 8 for __half, 4 for float - constexpr int T_per_load = ln::granularity / sizeof(T); - - constexpr int maxThreads = 256; - - // For Flaoat, unRoll 4, for __half, unRoll 2 - constexpr int internal_unRoll = sizeof(T) == 4 ? 4 : 2; - - const bool is_subblock_schedule = (elems_per_row <= 128) ? true : false; - const int h_per_step = is_subblock_schedule ? T_per_load : T_per_load * internal_unRoll; - - // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of - // warp-sized blocks rather than stepping up to 64/96 threads - const int one_step_threads = next_pow2((elems_per_row + h_per_step - 1) / h_per_step); - const int threadsPerGroup = (one_step_threads < maxThreads) ? one_step_threads : maxThreads; - - const int groups_per_block_max = - is_subblock_schedule ? (maxThreads + threadsPerGroup - 1) / threadsPerGroup : 1; - const int groups_per_block = (rows < groups_per_block_max) ? rows : groups_per_block_max; - const int groups_launch = (groups_per_block + rows - 1) / groups_per_block; - - dim3 block(threadsPerGroup, groups_per_block); - dim3 grid(groups_launch); - - const int elems_per_step = threadsPerGroup * h_per_step; - const int external_unRoll = (elems_per_row + elems_per_step - 1) / elems_per_step; - - if (is_subblock_schedule) { - // <=128 - if (threadsPerGroup == 1) { - LAUNCH_FUSED_RES_LN(1, 1, maxThreads); - } else if (threadsPerGroup == 2) { - LAUNCH_FUSED_RES_LN(1, 2, maxThreads); - } else if (threadsPerGroup == 4) { - LAUNCH_FUSED_RES_LN(1, 4, maxThreads); - } else if (threadsPerGroup == 8) { - LAUNCH_FUSED_RES_LN(1, 8, maxThreads); - } else if (threadsPerGroup == 16) { - LAUNCH_FUSED_RES_LN(1, 16, maxThreads); - } - } else if (external_unRoll == 1) { - // 129 - 4096 elems - // (this can launch with 1-7 warps as well) - LAUNCH_FUSED_RES_LN(1 * internal_unRoll, maxThreads, maxThreads); - } else if (external_unRoll == 2) { - // 4097 - 8192 elems - LAUNCH_FUSED_RES_LN(2 * internal_unRoll, maxThreads, maxThreads); - } else if (external_unRoll == 3) { - // 8193 - 12288 elems - LAUNCH_FUSED_RES_LN(3 * internal_unRoll, maxThreads, maxThreads); - } else if (external_unRoll == 4) { - // 12289 - 16384 elems - LAUNCH_FUSED_RES_LN(4 * internal_unRoll, maxThreads, maxThreads); - } -} - -#define LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(unRollFactor, threadsPerGroup, maxThreads) \ - fused_residual_ln \ - <<>>( \ - norm_output, res_output, vals, residual, bias, gamma, beta, epsilon, elems_per_row); - -template -void launch_fused_residual_ln_store_pre_ln_res(T* norm_output, - T* res_output, - const T* vals, - const T* residual, - const T* bias, - const T* gamma, - const T* beta, - float epsilon, - int rows, - int elems_per_row, - cudaStream_t stream) -{ - // 8 for __half, 4 for float - constexpr int T_per_load = ln::granularity / sizeof(T); - - constexpr int maxThreads = 256; - - // For Flaoat, unRoll 4, for __half, unRoll 2 - constexpr int internal_unRoll = sizeof(T) == 4 ? 4 : 2; - - const bool is_subblock_schedule = (elems_per_row <= 128) ? true : false; - const int h_per_step = is_subblock_schedule ? T_per_load : T_per_load * internal_unRoll; - - // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of - // warp-sized blocks rather than stepping up to 64/96 threads - const int one_step_threads = next_pow2((elems_per_row + h_per_step - 1) / h_per_step); - const int threadsPerGroup = (one_step_threads < maxThreads) ? one_step_threads : maxThreads; - - const int groups_per_block_max = - is_subblock_schedule ? (maxThreads + threadsPerGroup - 1) / threadsPerGroup : 1; - const int groups_per_block = (rows < groups_per_block_max) ? rows : groups_per_block_max; - const int groups_launch = (groups_per_block + rows - 1) / groups_per_block; - - dim3 block(threadsPerGroup, groups_per_block); - dim3 grid(groups_launch); - - const int elems_per_step = threadsPerGroup * h_per_step; - const int external_unRoll = (elems_per_row + elems_per_step - 1) / elems_per_step; - - if (is_subblock_schedule) { - // <=128 - if (threadsPerGroup == 1) { - LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 1, maxThreads); - } else if (threadsPerGroup == 2) { - LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 2, maxThreads); - } else if (threadsPerGroup == 4) { - LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 4, maxThreads); - } else if (threadsPerGroup == 8) { - LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 8, maxThreads); - } else if (threadsPerGroup == 16) { - LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 16, maxThreads); - } - } else if (external_unRoll == 1) { - // 129 - 4096 elems - // (this can launch with 1-7 warps as well) - LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1 * internal_unRoll, maxThreads, maxThreads); - } else if (external_unRoll == 2) { - // 4097 - 8192 elems - LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(2 * internal_unRoll, maxThreads, maxThreads); - } else if (external_unRoll == 3) { - // 8193 - 12288 elems - LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(3 * internal_unRoll, maxThreads, maxThreads); - } else if (external_unRoll == 4) { - // 12289 - 16384 elems - LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(4 * internal_unRoll, maxThreads, maxThreads); - } -} - -#define INSTANTIATE_RES_LN(T) \ - template void launch_fused_residual_ln( \ - T*, const T*, const T*, const T*, const T*, const T*, float, int, int, cudaStream_t); - -#define INSTANTIATE_PRE_LN_RES(T) \ - template void launch_fused_residual_ln_store_pre_ln_res( \ - T*, T*, const T*, const T*, const T*, const T*, const T*, float, int, int, cudaStream_t); - -INSTANTIATE_RES_LN(__half); -INSTANTIATE_RES_LN(float); -#ifdef BF16_AVAILABLE -INSTANTIATE_RES_LN(__nv_bfloat16); -#endif - -INSTANTIATE_PRE_LN_RES(__half); -INSTANTIATE_PRE_LN_RES(float); -#ifdef BF16_AVAILABLE -INSTANTIATE_PRE_LN_RES(__nv_bfloat16); -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/pointwise_ops.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/pointwise_ops.cu deleted file mode 100644 index 0301ff7770426994fc76cf3e6087586287c96c0f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/pointwise_ops.cu +++ /dev/null @@ -1,74 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include "conversion_utils.h" -#include "ds_kernel_utils.h" -#include "memory_access_utils.h" - -namespace pwise { -constexpr int granularity = 16; -constexpr int unroll = 4; -constexpr int threads = 256; -} // namespace pwise - -template -__global__ void vector_add_kernel(T* out, const T* a, const T* b, float gamma, int num_elems) -{ - constexpr int T_per_access = pwise::granularity / sizeof(T); - - const int block_offset = blockIdx.x * pwise::threads * pwise::unroll * T_per_access; - const int thread_offset = threadIdx.x * T_per_access; - const int total_offset = block_offset + thread_offset; - constexpr int stride = pwise::threads * T_per_access; - -#pragma unroll - for (int i = 0; i < pwise::unroll; i++) { - T temp_buf_a[T_per_access], temp_buf_b[T_per_access]; - - const int iter_idx = total_offset + i * stride; - - mem_access::load_global(temp_buf_a, a + iter_idx, iter_idx < num_elems); - mem_access::load_global(temp_buf_b, b + iter_idx, iter_idx < num_elems); - -#pragma unroll - for (int j = 0; j < T_per_access; j++) { - float up_cast_a = conversion::to(temp_buf_a[j]); - float up_cast_b = conversion::to(temp_buf_b[j]); - temp_buf_a[j] = conversion::to((gamma * up_cast_a) + up_cast_b); - } - - if (iter_idx < num_elems) { - mem_access::store_global(out + iter_idx, temp_buf_a); - } - } -} - -template -void launch_vector_add(T* out, - const T* a, - const T* b, - float gamma, - int num_elems, - cudaStream_t stream) -{ - constexpr int T_per_access = pwise::granularity / sizeof(T); - constexpr int T_per_block = pwise::threads * T_per_access * pwise::unroll; - - dim3 block(pwise::threads); - dim3 grid((num_elems + T_per_block - 1) / T_per_block); - - vector_add_kernel<<>>(out, a, b, gamma, num_elems); -} - -#define INSTANTIATE_VECTOR_ADD(T) \ - template void launch_vector_add( \ - T * out, const T* a, const T* b, float gamma, int num_elems, cudaStream_t stream); - -INSTANTIATE_VECTOR_ADD(float) -INSTANTIATE_VECTOR_ADD(__half) -#ifdef BF16_AVAILABLE -INSTANTIATE_VECTOR_ADD(__nv_bfloat16) -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/pt_binding.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/pt_binding.cpp deleted file mode 100644 index 19dbe73726f74bbcdcc62c6db1537a009c872b71..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/pt_binding.cpp +++ /dev/null @@ -1,2044 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include -#include -#include -#include "inference_context.h" -#include "inference_cublas_wrappers.h" -#include "inference_cuda_layers.h" - -std::array gemm_algos = std::array({99, 99, 99}); - -// NOTE: This activation function type enum should be always in sync -// with the python counterpart, otherwise the casting from python binding -// will be incorrect. -enum class ActivationFuncType { UNKNOWN = 0, GELU = 1, ReLU = 2, GATED_GELU = 3, GATED_SILU = 4 }; - -enum class NormType { UNKNOWN = 0, LayerNorm = 1, GroupNorm = 2, RMSNorm = 3 }; - -enum class TransformerType : uint8_t { UNKNOWN = 0, GPTType = 1, BERTType = 2 }; - -// NOTE: this is a temporary and dodgy solution to distinguish GPT and BERT style models -// based on the dimensions of the corresponding attention mask. -inline auto infer_transformer_type(at::Tensor& attn_mask) -> TransformerType -{ - auto attn_mask_num_dims = attn_mask.sizes().size(); - - if (attn_mask_num_dims > 2) { - return TransformerType::GPTType; - } else if (attn_mask_num_dims == 2) { - return TransformerType::BERTType; - } else { - return TransformerType::UNKNOWN; - } -} - -// infer stride of attention mask memory layout based on the model type. -inline auto get_attn_mask_stride(at::Tensor& attn_mask) -> int -{ - auto trnsfrmr_type = infer_transformer_type(attn_mask); - - if (trnsfrmr_type == TransformerType::GPTType) { - return attn_mask.size(2); - } else if (trnsfrmr_type == TransformerType::BERTType) { - // Bert style models have always a mask stride of 1. - return 1; - } else if (trnsfrmr_type == TransformerType::UNKNOWN) { - return 0; - } - - // this is just to make the compiler happy. - return 0; -} - -template -at::Tensor ds_softmax(at::Tensor& attn_scores, - at::Tensor& attn_mask, - at::Tensor& alibi, - bool triangular, - bool recompute, - bool local_attention, - int window_size, - bool async_op, - float layer_scale, - int head_offset, - int mp_size) -{ - auto attn_scores_c = attn_scores.contiguous(); - int bsz = attn_scores_c.size(0); - - int seq_len = attn_scores_c.size(1); - int len = attn_scores_c.sizes().size(); - if (len > 2) seq_len = attn_scores_c.size(2); - - int soft_len = attn_scores_c.size(2); - if (len > 3) soft_len = attn_scores_c.size(3); - - int heads = 1; - if (len > 1) heads = attn_scores_c.size(1); - - auto mask_stride = get_attn_mask_stride(attn_mask); - - launch_attn_softmax_v2((T*)attn_scores_c.data_ptr(), - (attn_mask.sizes().size() > 1 ? (T*)attn_mask.data_ptr() : nullptr), - (alibi.sizes().size() > 1 ? (T*)alibi.data_ptr() : nullptr), - layer_scale, - triangular, - recompute, - local_attention, - window_size, - bsz, - heads, - seq_len, - soft_len, - head_offset, - mask_stride, - mp_size, - InferenceContext::Instance().GetCurrentStream(async_op)); - - return attn_scores_c; -} - -template -void allocate_workspace(unsigned hidden_dim, - unsigned num_heads, - unsigned prompt_length, - unsigned batch_size, - unsigned num_layers, - unsigned mp_size = 1, - bool external_cache = false, - unsigned rank = 0, - unsigned max_out_tokens = 1024, - unsigned min_out_tokens = 1) -{ - InferenceContext::Instance().GenWorkSpace(num_layers, - num_heads, - batch_size, - prompt_length, - hidden_dim, - mp_size, - external_cache, - sizeof(T), - rank, - max_out_tokens, - min_out_tokens); -} - -template -at::Tensor einsum_sec_sm_ecm(at::Tensor& Q, at::Tensor& W) -{ - auto options = at::TensorOptions() - .dtype(Q.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); - float alpha = 1; - float gemm_beta = 0.0; - - /* - // Reallocate memory if we received a new prompt - if (!workspace || input.size(1) != 1) { - allocate_workspace(W.size(1), InferenceContext::Instance().GetMaxTokenLength(), - Q.size(0), 1, head_size); workspace = (T*)InferenceContext::Instance().GetWorkSpace(); - } - */ - - auto O = at::from_blob(workspace, {Q.size(1), Q.size(2), W.size(1)}, options); - unsigned m = W.size(1); - unsigned n = Q.size(1) * Q.size(2); - unsigned k = Q.size(0); - cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), - CUBLAS_OP_N, - CUBLAS_OP_T, - m, - n, - k, - &alpha, - &gemm_beta, - (T*)W.data_ptr(), - (T*)Q.data_ptr(), - (T*)O.data_ptr(), -// TODO HIP: Remove backward compatibility for torch<=2.0 in future -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif - return O; -} - -template -void attention_unfused(at::Tensor& prev_key_cont, - at::Tensor& query_cont, - at::Tensor& attn_mask, - at::Tensor& prev_value_cont, - at::Tensor& output, - int& bsz, - int& seq_len, - int& soft_len, - int& heads, - float& norm_factor, - bool triangular, - bool recompute, - bool local_attention, - int window_size) -{ - auto options = at::TensorOptions() - .dtype(query_cont.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - float alpha = norm_factor; - float gemm_beta = 0.0; - auto attn_score = at::empty({bsz, heads, seq_len, soft_len}, options); - int k = prev_value_cont.size(2) / heads; - - auto mask_stride = get_attn_mask_stride(attn_mask); - - cublasSetStream(InferenceContext::Instance().GetCublasHandle(), - InferenceContext::Instance().GetCurrentStream()); - cublas_strided_batched_gemm(InferenceContext::Instance().GetCublasHandle(), - soft_len, - seq_len, - k, - &alpha, - &gemm_beta, - (T*)prev_key_cont.data_ptr(), - (T*)query_cont.data_ptr(), - (T*)attn_score.data_ptr(), - CUBLAS_OP_N, - CUBLAS_OP_N, - soft_len * k, - seq_len * k, - seq_len * soft_len, - bsz * heads, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif - launch_attn_softmax_v2((T*)attn_score.data_ptr(), - (T*)(attn_mask.sizes().size() > 1 ? attn_mask.data_ptr() : nullptr), - (T*)nullptr, - 1.0, - triangular, - recompute, - local_attention, - window_size, - bsz, - heads, - seq_len, - soft_len, - 0, - mask_stride, - 1, - InferenceContext::Instance().GetCurrentStream(false)); - alpha = 1.0; - cublas_strided_batched_gemm(InferenceContext::Instance().GetCublasHandle(), - k, - seq_len, - soft_len, - &alpha, - &gemm_beta, - (T*)prev_value_cont.data_ptr(), - (T*)attn_score.data_ptr(), - (T*)output.data_ptr(), - CUBLAS_OP_N, - CUBLAS_OP_N, - soft_len * k, - seq_len * soft_len, - seq_len * k, - bsz * heads, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif -} - -template -std::vector ds_softmax_context1(at::Tensor& query, - at::Tensor& prev_key, - at::Tensor& new_key, - at::Tensor& attn_mask, - at::Tensor& prev_value, - at::Tensor& new_value, - int heads, - float norm_factor, - bool merging, - bool triangular, - bool local_attention, - int window_size, - bool no_masking) -{ - auto query_cont = query.contiguous(); - auto prev_key_cont = prev_key.contiguous(); - auto prev_value_cont = prev_value.contiguous(); - - int new_size = (new_value.sizes().size() > 1 ? new_value.size(1) : 0); - - // Attn_Score [ batch Head Sequence-length Softmax-length] - - int bsz = query_cont.size(0); - int seq_len = query_cont.size(1); - int soft_len = prev_value.size(1); - - auto options = at::TensorOptions() - .dtype(query_cont.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - auto output = - at::empty({prev_value.size(0), heads, seq_len, prev_value.size(2) / heads}, options); - attention_unfused(prev_key_cont, - query_cont, - attn_mask, //(no_masking ? nullptr : (T*)attn_mask.data_ptr()), - prev_value_cont, - output, - bsz, - seq_len, - soft_len, - heads, - norm_factor, - (triangular && (new_size == 0)), - (new_size == 0), - local_attention, - window_size); - - return {output, prev_key, prev_value}; -} - -template -void ds_softmax_internal(T* attn_scores, - at::Tensor& attn_mask, - at::Tensor& alibi, - float& layer_scale, - bool triangular, - bool recompute, - bool local_attention, - int window_size, - int bsz, - int seq_len, - int soft_len, - int heads) -{ - auto mask_stride = get_attn_mask_stride(attn_mask); - - launch_attn_softmax_v2((T*)attn_scores, - (attn_mask.sizes().size() > 1 ? (T*)attn_mask.data_ptr() : nullptr), - (alibi.sizes().size() > 1 ? (T*)alibi.data_ptr() : nullptr), - layer_scale, - triangular, - recompute, - local_attention, - window_size, - bsz, - heads, - seq_len, - soft_len, - 0, - mask_stride, - 1, - at::cuda::getCurrentCUDAStream()); -} - -template -void attention_unfused(T* prev_key_cont, - T* query_cont, - at::Tensor& attn_mask, - T* prev_value_cont, - T* output, - unsigned& bsz, - int& k, - unsigned& seq_len, - unsigned& soft_len, - int& heads, - float& norm_factor, - bool triangular, - bool recompute, - bool local_attention, - int window_size, - at::Tensor& alibi, - int layer_id) -{ - float layer_scale = alibi.sizes().size() > 1 ? std::max(1, layer_id) : 1.0; - float alpha = norm_factor * norm_factor / layer_scale; - float gemm_beta = 0.0; - T* workspace = (T*)InferenceContext::Instance().GetAttentionUnfusedWorkspace(); - - cublasSetStream(InferenceContext::Instance().GetCublasHandle(), - InferenceContext::Instance().GetCurrentStream()); - cublas_strided_batched_gemm(InferenceContext::Instance().GetCublasHandle(), - soft_len, - seq_len, - k, - &alpha, - &gemm_beta, - (T*)prev_key_cont, - (T*)query_cont, - workspace, - CUBLAS_OP_T, - CUBLAS_OP_N, - InferenceContext::Instance().GetMaxTokenLength() * k, - seq_len * k, - seq_len * soft_len, - bsz * heads, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif - ds_softmax_internal(workspace, - attn_mask, - alibi, - layer_scale, - triangular, - recompute, - local_attention, - window_size, - bsz, - seq_len, - soft_len, - heads); - alpha = 1.0; - cublas_strided_batched_gemm(InferenceContext::Instance().GetCublasHandle(), - k, - seq_len, - soft_len, - &alpha, - &gemm_beta, - (T*)prev_value_cont, - workspace, - (T*)output, - CUBLAS_OP_N, - CUBLAS_OP_N, - InferenceContext::Instance().GetMaxTokenLength() * k, - seq_len * soft_len, - seq_len * k, - bsz * heads, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif -} - -void reset_cache() { InferenceContext::Instance().reset_tokens(); } - -template -std::vector ds_softmax_context(at::Tensor& query_key_value, - at::Tensor& attn_mask, - int rotary_dim, - bool rotate_half, - bool rotate_every_two, - int heads, - int num_kv, - float norm_factor, - bool triangular, - bool local_attention, - int window_size, - bool no_masking, - unsigned layer_id, - unsigned num_layers, - at::Tensor& alibi, - float rope_theta, - bool is_prompt, - std::optional token_idx, - std::optional position_ids) -{ - unsigned bsz = query_key_value.size(0); - unsigned seq_len = query_key_value.size(1); - int k = query_key_value.size(2) / (heads + 2 * (num_kv > 0 ? num_kv : heads)); - unsigned hidden_dim = heads * k; - - is_prompt = (seq_len > 1); - - if (is_prompt) InferenceContext::Instance().reset_tokens(seq_len); - unsigned soft_len = InferenceContext::Instance().current_tokens(); - - auto options = at::TensorOptions() - .dtype(query_key_value.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); - size_t buf_size = bsz * seq_len * hidden_dim; - auto output = torch::from_blob(workspace + 4 * buf_size, {bsz, seq_len, hidden_dim}, options); - - auto query_cont = workspace + 5 * buf_size; - size_t offset = - 10 * (hidden_dim * bsz * InferenceContext::Instance().GetMaxTokenLength()) + - layer_id * 2 * bsz * InferenceContext::Instance().GetMaxTokenLength() * hidden_dim; - unsigned all_tokens = soft_len; - auto kv_cache = workspace + offset + (hidden_dim / heads) * (is_prompt ? 0 : soft_len - 1); - size_t value_offset = bsz * InferenceContext::Instance().GetMaxTokenLength() * hidden_dim; - - T* temp_buf = (T*)output.data_ptr() + at::numel(output); - launch_bias_add_transform_0213((T*)query_cont, - kv_cache, - kv_cache + value_offset, - (T*)query_key_value.data_ptr(), - nullptr, - bsz, - seq_len, - (is_prompt ? 0 : soft_len - 1), - soft_len, - hidden_dim, - heads, - (num_kv > 0 ? num_kv : heads), - rotary_dim, - rotate_half, - rotate_every_two, - InferenceContext::Instance().GetCurrentStream(), - 3, - InferenceContext::Instance().GetMaxTokenLength(), - rope_theta); - if (rotary_dim > 0 && rotate_half) - launch_apply_rotary_pos_emb(query_cont, - kv_cache, - k, - seq_len, - rotary_dim, - (is_prompt ? 0 : soft_len - 1), - heads, - bsz, - rope_theta, - InferenceContext::Instance().GetCurrentStream(), - InferenceContext::Instance().GetMaxTokenLength()); - - attention_unfused(workspace + offset, - (T*)query_cont, - attn_mask, - workspace + offset + value_offset, - temp_buf, - bsz, - k, - seq_len, - all_tokens, - heads, - norm_factor, - (triangular && is_prompt), - is_prompt, - local_attention, - window_size, - alibi, - layer_id); - launch_transform4d_0213((T*)output.data_ptr(), - temp_buf, - bsz, - heads, - seq_len, - output.size(2), - InferenceContext::Instance().GetCurrentStream(false), - 1); - - if (layer_id == num_layers - 1) InferenceContext::Instance().advance_tokens(); - auto prev_key = torch::from_blob( - workspace + offset, - {bsz, heads, all_tokens, k}, - {hidden_dim * static_cast(InferenceContext::Instance().GetMaxTokenLength()), - k * static_cast(InferenceContext::Instance().GetMaxTokenLength()), - k, - 1}, - options); - - auto prev_value = torch::from_blob( - workspace + offset + value_offset, - {bsz, heads, all_tokens, k}, - {hidden_dim * static_cast(InferenceContext::Instance().GetMaxTokenLength()), - k * static_cast(InferenceContext::Instance().GetMaxTokenLength()), - k, - 1}, - options); - - return {output, prev_key, prev_value}; -} - -template -at::Tensor ds_bias_gelu(at::Tensor& input, at::Tensor& bias) -{ - auto input_cont = input.contiguous(); - - int bsz = input_cont.size(0) * input_cont.size(1); - int intermediate_size = input_cont.size(2); - - launch_bias_gelu((T*)input_cont.data_ptr(), - (T*)bias.data_ptr(), - intermediate_size, - bsz, - InferenceContext::Instance().GetCurrentStream()); - return input_cont; -} - -#define DISPATCH_GATED_ACT(T_TYPE, C_TYPE) \ - if (activation.options().dtype() == torch::T_TYPE) { \ - launch_gated_activation((C_TYPE*)output.data_ptr(), \ - (const C_TYPE*)activation.data_ptr(), \ - (const C_TYPE*)bias.data_ptr(), \ - rows, \ - out_channels, \ - channels, \ - activation_type == ActivationFuncType::GATED_GELU, \ - InferenceContext::Instance().GetCurrentStream()); \ - } - -at::Tensor ds_gated_activation(at::Tensor& activation, at::Tensor& bias, int actFun) -{ - /* - Used in FF of Stable diffusion - */ - - const ActivationFuncType activation_type = static_cast(actFun); - - assert(activation_type == ActivationFuncType::GATED_GELU || - activation_type == ActivationFuncType::GATED_SILU); - - const int batch_size = activation.size(0); - const int seq_len = activation.size(1); - const int channels = activation.size(2); - - const int rows = batch_size * seq_len; - // Dimensionality is cut in half - const int out_channels = channels / 2; - - auto output = at::empty({batch_size, seq_len, out_channels}, activation.options()); - - DISPATCH_GATED_ACT(kFloat, float); - DISPATCH_GATED_ACT(kHalf, __half); -#ifdef BF16_AVAILABLE - DISPATCH_GATED_ACT(kBFloat16, __nv_bfloat16); -#endif - - return output; -} - -template -at::Tensor ds_bias_relu(at::Tensor& input, at::Tensor& bias) -{ - auto input_cont = input.contiguous(); - - int bsz = input_cont.size(0) * input_cont.size(1); - int intermediate_size = input_cont.size(2); - - launch_bias_relu((T*)input_cont.data_ptr(), - (T*)bias.data_ptr(), - intermediate_size, - bsz, - InferenceContext::Instance().GetCurrentStream()); - return input_cont; -} - -template -at::Tensor ds_bias_add(at::Tensor& input, at::Tensor& bias) -{ - auto input_cont = input.contiguous(); - - int bsz = input_cont.size(0) * input_cont.size(1); - int hidden_size = input_cont.size(2); - - launch_bias_add((T*)input_cont.data_ptr(), - (T*)bias.data_ptr(), - hidden_size, - bsz, - InferenceContext::Instance().GetCurrentStream()); - return input_cont; -} - -template -at::Tensor ds_bias_residual(at::Tensor& input, at::Tensor& residual, at::Tensor& bias) -{ - auto input_cont = input.contiguous(); - auto residual_cont = residual.contiguous(); - - int bsz = input_cont.size(0) * input_cont.size(1); - // launch_bias_residual((T*)input_cont.data_ptr(), - // (T*)residual_cont.data_ptr(), - // (T*)bias.data_ptr(), - // bsz, - // input_cont.size(2), - // (bias.size(0) > 1), - // InferenceContext::Instance().GetCurrentStream()); - return input_cont; -} - -#define DISPATCH_LAYER_NORM(T_TYPE, C_TYPE) \ - if (input.options().dtype() == torch::T_TYPE) { \ - launch_fused_ln((C_TYPE*)output.data_ptr(), \ - (const C_TYPE*)input.data_ptr(), \ - (const C_TYPE*)gamma.data_ptr(), \ - (const C_TYPE*)beta.data_ptr(), \ - epsilon, \ - rows, \ - elems_per_row, \ - InferenceContext::Instance().GetCurrentStream()); \ - } - -at::Tensor ds_layer_norm(at::Tensor& input, at::Tensor& gamma, at::Tensor& beta, float epsilon) -{ - const int rows = input.size(0) * input.size(1); - const int elems_per_row = input.size(2); - auto output = at::empty_like(input); - - DISPATCH_LAYER_NORM(kFloat, float); - DISPATCH_LAYER_NORM(kHalf, __half); -#ifdef BF16_AVAILABLE - DISPATCH_LAYER_NORM(kBFloat16, __nv_bfloat16); -#endif - - return output; -} - -#define DISPATCH_RMS_NORM(T_TYPE, C_TYPE) \ - if (input.options().dtype() == torch::T_TYPE) { \ - launch_rms_norm((C_TYPE*)output.data_ptr(), \ - (C_TYPE*)nullptr, \ - (const C_TYPE*)input.data_ptr(), \ - (const C_TYPE*)nullptr, \ - (const C_TYPE*)gamma.data_ptr(), \ - epsilon, \ - rows, \ - elems_per_row, \ - InferenceContext::Instance().GetCurrentStream()); \ - } - -at::Tensor ds_rms_norm(at::Tensor& input, at::Tensor& gamma, float epsilon) -{ - // Get number of dims of tensor - int num_dims = input.dim(); - const int rows = (num_dims == 2) ? input.size(0) : input.size(0) * input.size(1); - const int elems_per_row = (num_dims == 2) ? input.size(1) : input.size(2); - - auto output = at::empty_like(input); - - DISPATCH_RMS_NORM(kFloat, float); - DISPATCH_RMS_NORM(kHalf, __half); -#ifdef BF16_AVAILABLE - DISPATCH_RMS_NORM(kBFloat16, __nv_bfloat16); -#endif - - return output; -} - -#define DISPATCH_PRE_RMS_NORM(T_TYPE, C_TYPE) \ - if (input.options().dtype() == torch::T_TYPE) { \ - launch_rms_norm((C_TYPE*)output.data_ptr(), \ - (C_TYPE*)res_out.data_ptr(), \ - (const C_TYPE*)input.data_ptr(), \ - (const C_TYPE*)residual.data_ptr(), \ - (const C_TYPE*)gamma.data_ptr(), \ - epsilon, \ - rows, \ - elems_per_row, \ - InferenceContext::Instance().GetCurrentStream()); \ - } - -std::vector ds_pre_rms_norm(at::Tensor& input, - at::Tensor& residual, - at::Tensor& gamma, - float epsilon) -{ - // Get number of dims of tensor - int num_dims = input.dim(); - const int rows = (num_dims == 2) ? input.size(0) : input.size(0) * input.size(1); - const int elems_per_row = (num_dims == 2) ? input.size(1) : input.size(2); - - auto output = at::empty_like(input); - auto res_out = at::empty_like(residual); - - DISPATCH_PRE_RMS_NORM(kFloat, float); - DISPATCH_PRE_RMS_NORM(kHalf, __half); -#ifdef BF16_AVAILABLE - DISPATCH_PRE_RMS_NORM(kBFloat16, __nv_bfloat16); -#endif - - return {output, res_out}; -} - -template -void ds_layer_norm_internal(T* workspace, - at::Tensor& input, - at::Tensor& gamma, - at::Tensor& beta, - float epsilon) -{ - int bsz = input.size(0) * input.size(1); - launch_fused_ln(workspace, - (const T*)input.data_ptr(), - (const T*)gamma.data_ptr(), - (const T*)beta.data_ptr(), - epsilon, - bsz, - input.size(2), - InferenceContext::Instance().GetCurrentStream()); -} - -#define DISPATCH_LAYER_NORM_RESIDUAL(T_TYPE, C_TYPE) \ - if (input.options().dtype() == torch::T_TYPE) { \ - launch_fused_residual_ln((C_TYPE*)output.data_ptr(), \ - (const C_TYPE*)input.data_ptr(), \ - (const C_TYPE*)residual.data_ptr(), \ - (const C_TYPE*)bias.data_ptr(), \ - (const C_TYPE*)gamma.data_ptr(), \ - (const C_TYPE*)beta.data_ptr(), \ - epsilon, \ - rows, \ - elems_per_row, \ - InferenceContext::Instance().GetCurrentStream()); \ - } - -/* Currently only used in unit testing */ -at::Tensor ds_layer_norm_residual(at::Tensor& input, - at::Tensor& bias, - at::Tensor& residual, - at::Tensor& gamma, - at::Tensor& beta, - float epsilon) -{ - const int rows = input.size(0) * input.size(1); - const int elems_per_row = input.size(2); - auto output = at::empty_like(input); - - DISPATCH_LAYER_NORM_RESIDUAL(kFloat, float); - DISPATCH_LAYER_NORM_RESIDUAL(kHalf, __half); -#ifdef BF16_AVAILABLE - DISPATCH_LAYER_NORM_RESIDUAL(kBFloat16, __nv_bfloat16); -#endif - - return output; -} - -#define DISPATCH_PRE_LAYER_NORM_RESIDUAL(T_TYPE, C_TYPE) \ - if (input.options().dtype() == torch::T_TYPE) { \ - launch_fused_residual_ln_store_pre_ln_res( \ - (C_TYPE*)norm_output.data_ptr(), \ - (C_TYPE*)res_output.data_ptr(), \ - (const C_TYPE*)input.data_ptr(), \ - (const C_TYPE*)residual.data_ptr(), \ - (const C_TYPE*)bias.data_ptr(), \ - (const C_TYPE*)gamma.data_ptr(), \ - (const C_TYPE*)beta.data_ptr(), \ - epsilon, \ - rows, \ - elems_per_row, \ - InferenceContext::Instance().GetCurrentStream()); \ - } - -/* Currently only used in unit testing */ -std::vector ds_layer_norm_residual_store_pre_ln_res(at::Tensor& input, - at::Tensor& bias, - at::Tensor& residual, - at::Tensor& gamma, - at::Tensor& beta, - float epsilon) -{ - const int rows = input.size(0) * input.size(1); - const int elems_per_row = input.size(2); - auto norm_output = at::empty_like(input); - auto res_output = at::empty_like(input); - - DISPATCH_PRE_LAYER_NORM_RESIDUAL(kFloat, float); - DISPATCH_PRE_LAYER_NORM_RESIDUAL(kHalf, __half); -#ifdef BF16_AVAILABLE - DISPATCH_PRE_LAYER_NORM_RESIDUAL(kBFloat16, __nv_bfloat16); -#endif - - return {norm_output, res_output}; -} - -template -void quantized_gemm(void* output, - T* input, - at::Tensor& weight, - at::Tensor& qscale, - int groups, - int bsz, - int hidden_size) -{ - // T* weight16 = (T*)InferenceContext::Instance().GetWorkSpace() + 12 * hidden_size * bsz; - - auto options = at::TensorOptions() - .dtype(at::kHalf) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - auto tmp = torch::empty(weight.sizes(), options); - T* weight16 = (T*)tmp.data_ptr(); - launch_dequantize(weight16, - (int8_t*)weight.data_ptr(), - (float*)qscale.data_ptr(), - weight.size(0), - weight.size(1), - groups, - InferenceContext::Instance().GetCurrentStream()); - - float alpha = (T)1.0; - float gemm_beta = (T)0.0; - cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), - CUBLAS_OP_T, - CUBLAS_OP_N, - weight.size(0), - bsz, - weight.size(1), - &alpha, - &gemm_beta, - weight16, - (T*)input, - (T*)output, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif -} - -template -at::Tensor qkv_unfused_cublas(at::Tensor& output, - at::Tensor& input, - at::Tensor& weight, - at::Tensor& q_scale, - at::Tensor& bias, - at::Tensor& gamma, - at::Tensor& beta, - const float epsilon, - bool add_bias, - bool q_int8, - bool transposed_mode) -{ - int bsz = input.size(0) * input.size(1); - T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); - workspace += (3 * bsz * input.size(2)); - ds_layer_norm_internal(workspace, input, gamma, beta, epsilon); - - if (q_int8) { - quantized_gemm( - output.data_ptr(), workspace, weight, q_scale, q_scale.size(0), bsz, input.size(2)); - } else { - float alpha = (T)1.0; - float gemm_beta = (T)0.0; - - cublasSetStream(InferenceContext::Instance().GetCublasHandle(), - InferenceContext::Instance().GetCurrentStream()); - cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), - (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), - CUBLAS_OP_N, - weight.size(transposed_mode ? 0 : 1), - bsz, - input.size(2), - &alpha, - &gemm_beta, - (T*)weight.data_ptr(), - workspace, - (T*)output.data_ptr(), -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif - } - if (add_bias) - launch_bias_add((T*)output.data_ptr(), - (T*)bias.data_ptr(), - (transposed_mode || q_int8) ? weight.size(0) : weight.size(1), - bsz, - InferenceContext::Instance().GetCurrentStream()); - return torch::from_blob(workspace, input.sizes(), input.options()); -} - -template -std::vector ds_rms_qkv(at::Tensor& input, - at::Tensor& weight, - at::Tensor& q_scale, - at::Tensor& gamma, - const float epsilon, - bool q_int8, - bool transposed_mode) -{ - const int bsz = input.size(0) * input.size(1); - T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); - T* rms_norm_ptr = workspace + (3 * bsz * input.size(2)); - int out_size = (transposed_mode || q_int8) ? weight.size(0) : weight.size(1); - - auto options = at::TensorOptions() - .dtype(input.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - auto rms_norm = at::from_blob(rms_norm_ptr, input.sizes(), options); - auto output = at::from_blob(workspace, {input.size(0), input.size(1), out_size}, options); - - launch_rms_norm((T*)rms_norm.data_ptr(), - (T*)nullptr, - (const T*)input.data_ptr(), - (const T*)nullptr, - (const T*)gamma.data_ptr(), - epsilon, - bsz, - input.size(2), - InferenceContext::Instance().GetCurrentStream()); - - if (q_int8) { - quantized_gemm((T*)output.data_ptr(), - (T*)rms_norm.data_ptr(), - weight, - q_scale, - q_scale.size(0), - bsz, - input.size(2)); - } else { - float alpha = (T)1.0; - float gemm_beta = (T)0.0; - - cublasSetStream(InferenceContext::Instance().GetCublasHandle(), - InferenceContext::Instance().GetCurrentStream()); - cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), - (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), - CUBLAS_OP_N, - weight.size(transposed_mode ? 0 : 1), - bsz, - input.size(2), - &alpha, - &gemm_beta, - (T*)weight.data_ptr(), - (T*)rms_norm.data_ptr(), - (T*)output.data_ptr(), -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif - } - - return {output, rms_norm}; -} - -template -std::vector ds_qkv_gemm(at::Tensor& input, - at::Tensor& weight, - at::Tensor& q_scale, - at::Tensor& bias, - at::Tensor& gamma, - at::Tensor& beta, - const float epsilon, - bool add_bias, - bool q_int8, - bool transposed_mode) -{ - int bsz = input.size(0) * input.size(1); - T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); - int out_size = (transposed_mode || q_int8) ? weight.size(0) : weight.size(1); - - auto options = at::TensorOptions() - .dtype(input.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - auto output = at::from_blob(workspace, {input.size(0), input.size(1), out_size}, options); - auto inp_norm = qkv_unfused_cublas(output, - input, - weight, - q_scale, - bias, - gamma, - beta, - epsilon, - add_bias, - q_int8, - transposed_mode); - - return {output, inp_norm}; -} - -template -void quantized_gemm(at::Tensor& output, - at::Tensor& input, - at::Tensor& weight, - at::Tensor& qscale, - int groups, - int merge_count) -{ - int bsz = input.size(0) * input.size(1); - auto options = at::TensorOptions() - .dtype(input.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - auto weight16 = at::empty({weight.size(0), weight.size(1)}, options); - - launch_dequantize((T*)weight16.data_ptr(), - (int8_t*)weight.data_ptr(), - (float*)qscale.data_ptr(), - weight.size(0), - weight.size(1), - groups, - merge_count, - InferenceContext::Instance().GetCurrentStream()); - - float alpha = (T)1.0; - float gemm_beta = (T)0.0; - cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), - CUBLAS_OP_T, - CUBLAS_OP_N, - weight.size(0), - bsz, - input.size(2), - &alpha, - &gemm_beta, - (T*)weight16.data_ptr(), - (T*)input.data_ptr(), - (T*)output.data_ptr(), -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif -} - -template -at::Tensor ds_linear_layer(at::Tensor& input, - at::Tensor& weight, - at::Tensor& bias, - bool add_bias, - bool do_flash_attn, - int num_heads, - bool transposed_mode, - float rope_theta) -{ - auto input_cont = input.contiguous(); - auto options = at::TensorOptions() - .dtype(input_cont.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - int head_size = input_cont.size(2) / num_heads; - int bsz = input.size(0) * input.size(1); - int out_size = transposed_mode ? weight.size(0) : weight.size(1); - T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); - auto output = at::from_blob(workspace, {input.size(0), input.size(1), out_size}, options); - - float alpha = (T)1.0; - float gemm_beta = (T)0.0; - cublasSetStream(InferenceContext::Instance().GetCublasHandle(), - InferenceContext::Instance().GetCurrentStream()); - - cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), - (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), - CUBLAS_OP_N, - weight.size(transposed_mode ? 0 : 1), - bsz, - input_cont.size(2), - &alpha, - &gemm_beta, - (T*)weight.data_ptr(), - (T*)input_cont.data_ptr(), - (T*)output.data_ptr(), -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif - if (add_bias) - launch_bias_add((T*)output.data_ptr(), - (T*)bias.data_ptr(), - weight.size(transposed_mode ? 0 : 1), - bsz, - InferenceContext::Instance().GetCurrentStream()); - bool add_padding = (head_size % 32 != 0 && head_size < 64) || (head_size % 64 != 0); - if (do_flash_attn) { - if (add_padding) { - int padded_head_size = head_size < 32 ? 32 : (head_size < 64 ? 64 : 128); - auto padded_output = workspace + output.numel(); - auto final_output = - padded_output + (input.size(0) * input.size(1) * 3 * num_heads * padded_head_size); - pad_data(padded_output, - workspace, - 3 * bsz * num_heads, - head_size, - padded_head_size, - InferenceContext::Instance().GetCurrentStream()); - - launch_bias_add_transform_0213( - final_output, - final_output + (input.size(0) * input.size(1) * num_heads * padded_head_size), - final_output + (input.size(0) * input.size(1) * 2 * num_heads * padded_head_size), - padded_output, - nullptr, - input.size(0), - input.size(1), - 0, - input.size(1), - (num_heads * padded_head_size), - num_heads, - -1, - -1, - false, - false, - InferenceContext::Instance().GetCurrentStream(), - 3, - input.size(1), - rope_theta); - return at::from_blob(final_output, - {3, input.size(0), num_heads, input.size(1), padded_head_size}, - options); - // return at::from_blob(padded_output, {input.size(0) * input.size(1), 3, num_heads, - // padded_head_size}, options); - } else { - auto final_output = workspace + output.numel(); - launch_bias_add_transform_0213( - final_output, - final_output + (input.size(0) * input.size(1) * input_cont.size(2)), - final_output + (input.size(0) * input.size(1) * 2 * input_cont.size(2)), - workspace, - nullptr, - input.size(0), - input.size(1), - 0, - input.size(1), - input_cont.size(2), - num_heads, - -1, - -1, - false, - false, - InferenceContext::Instance().GetCurrentStream(), - 3, - input.size(1), - rope_theta); - return at::from_blob( - final_output, {3, input.size(0), num_heads, input.size(1), head_size}, options); - // return at::from_blob(workspace, {input.size(0) * input.size(1), 3, num_heads, - // head_size}, options); - } - - } else - return output; -} - -template -std::vector add_padding(at::Tensor& query, at::Tensor& key, at::Tensor& value) -{ - int head_size = query.size(3); - int padded_head_size = head_size < 32 ? 32 : (head_size < 64 ? 64 : 128); - T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); - T* key_pad_ptr = workspace + padded_head_size * query.size(0) * query.size(1) * query.size(2); - T* value_pad_ptr = key_pad_ptr + padded_head_size * query.size(0) * query.size(1) * 128; - pad_head_seq(workspace, - (T*)query.data_ptr(), - query.size(0) * query.size(1), - query.size(2), - query.size(2), - head_size, - padded_head_size, - InferenceContext::Instance().GetCurrentStream()); - pad_head_seq(key_pad_ptr, - (T*)key.data_ptr(), - query.size(0) * query.size(1), - key.size(2), - 128, - head_size, - padded_head_size, - InferenceContext::Instance().GetCurrentStream()); - pad_head_seq(value_pad_ptr, - (T*)value.data_ptr(), - query.size(0) * query.size(1), - key.size(2), - 128, - head_size, - padded_head_size, - InferenceContext::Instance().GetCurrentStream()); - return { - at::from_blob(workspace, - {query.size(0), query.size(1), query.size(2), padded_head_size}, - query.options()), - at::from_blob( - key_pad_ptr, {query.size(0), query.size(1), 128, padded_head_size}, query.options()), - at::from_blob( - value_pad_ptr, {query.size(0), query.size(1), 128, padded_head_size}, query.options())}; -} - -template -std::vector padd_add_transform(at::Tensor& query, - at::Tensor& key, - at::Tensor& value, - int heads, - bool add_padding) -{ - int head_size = query.size(2) / heads; - int key_value_length = add_padding ? 128 : key.size(1); - int padded_head_size = add_padding ? (head_size < 32 ? 32 : (head_size < 64 ? 64 : 128)) - : head_size; - T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); - T* key_pad_ptr = workspace + padded_head_size * query.size(0) * heads * query.size(1); - T* value_pad_ptr = key_pad_ptr + padded_head_size * query.size(0) * heads * key_value_length; - launch_pad_add_transform_0213(workspace, - (T*)query.data_ptr(), - query.size(0), - query.size(2), - query.size(1), - query.size(1), - heads, - padded_head_size, - InferenceContext::Instance().GetCurrentStream()); - launch_pad_add_transform_0213(key_pad_ptr, - (T*)key.data_ptr(), - key.size(0), - key.size(2), - key.size(1), - key_value_length, - heads, - padded_head_size, - InferenceContext::Instance().GetCurrentStream()); - launch_pad_add_transform_0213(value_pad_ptr, - (T*)value.data_ptr(), - value.size(0), - value.size(2), - value.size(1), - key_value_length, - heads, - padded_head_size, - InferenceContext::Instance().GetCurrentStream()); - return { - at::from_blob( - workspace, {query.size(0), heads, query.size(1), padded_head_size}, query.options()), - at::from_blob(key_pad_ptr, - {query.size(0), heads, key_value_length, padded_head_size}, - query.options()), - at::from_blob(value_pad_ptr, - {query.size(0), heads, key_value_length, padded_head_size}, - query.options())}; -} - -template -at::Tensor ds_vector_matmul(at::Tensor& input, - at::Tensor& weight, - bool async_op, - at::Tensor& q_scale, - bool q_int8, - bool transposed_mode) -{ - auto options = at::TensorOptions() - .dtype(input.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - int out_size = (q_int8 || transposed_mode) ? weight.size(0) : weight.size(1); - int bsz = input.size(0) * input.size(1); - - T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); - auto output = at::from_blob(workspace, {input.size(0), input.size(1), out_size}, options); - if (q_int8) { - quantized_gemm(output.data_ptr(), - (T*)input.data_ptr(), - weight, - q_scale, - q_scale.size(0), - bsz, - input.size(2)); - } else { - float alpha = (T)1.0; - float gemm_beta = (T)0.0; - cublasSetStream(InferenceContext::Instance().GetCublasHandle(), - InferenceContext::Instance().GetCurrentStream(async_op)); - cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), - (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), - CUBLAS_OP_N, - weight.size(transposed_mode ? 0 : 1), - bsz, - input.size(2), - &alpha, - &gemm_beta, - (T*)weight.data_ptr(), - (T*)input.data_ptr(), - (T*)output.data_ptr(), -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif - } - return output; -} - -template -at::Tensor ds_vector_matmul_int8(at::Tensor& input, - at::Tensor& weight, - at::Tensor& q_scale, - int groups, - int merge_count) -{ - auto input_cont = input.contiguous(); - auto options = at::TensorOptions() - .dtype(input_cont.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - auto output = at::empty({input_cont.size(0), input_cont.size(1), weight.size(1)}, options); - - quantized_gemm(output, input_cont, weight, q_scale, groups, merge_count); - return output; -} - -template -at::Tensor mlp_unfused_cublas(at::Tensor& output, - at::Tensor& input, - at::Tensor& residual, - at::Tensor& input_bias, - at::Tensor& weight, - at::Tensor& weight1, - at::Tensor& bias, - at::Tensor& gamma, - at::Tensor& beta, - const float epsilon, - bool preLayerNorm, - bool mlp_after_attn, - at::Tensor& q_scale, - at::Tensor& q_scale1, - bool q_int8, - ActivationFuncType act_func_type, - bool transposed_mode) -{ - int bsz = input.size(0) * input.size(1); - T* inp_norm = (T*)InferenceContext::Instance().GetWorkSpace() + torch::numel(input) + - torch::numel(output); - T* intermediate = inp_norm + torch::numel(input); - - if (mlp_after_attn) { - launch_fused_residual_ln((T*)inp_norm, - (const T*)input.data_ptr(), - (const T*)residual.data_ptr(), - (const T*)input_bias.data_ptr(), - (const T*)gamma.data_ptr(), - (const T*)beta.data_ptr(), - epsilon, - bsz, - input.size(2), - InferenceContext::Instance().GetCurrentStream()); - } else { - ds_layer_norm_internal(inp_norm, input, gamma, beta, epsilon); - } - if (q_int8) { - quantized_gemm( - intermediate, inp_norm, weight, q_scale, q_scale.size(0), bsz, input.size(2)); - } else { - float alpha = (T)1.0; - float gemm_beta = (T)0.0; - cublasSetStream(InferenceContext::Instance().GetCublasHandle(), - InferenceContext::Instance().GetCurrentStream()); - cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), - (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), - CUBLAS_OP_N, - weight.size(transposed_mode ? 0 : 1), - bsz, - input.size(2), - &alpha, - &gemm_beta, - (T*)weight.data_ptr(), - inp_norm, - intermediate, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif - } - if (act_func_type == ActivationFuncType::GELU) { - launch_bias_gelu(intermediate, - (T*)bias.data_ptr(), - (transposed_mode || q_int8) ? weight.size(0) : weight.size(1), - bsz, - InferenceContext::Instance().GetCurrentStream()); - } else if (act_func_type == ActivationFuncType::ReLU) { - launch_bias_relu(intermediate, - (T*)bias.data_ptr(), - (transposed_mode || q_int8) ? weight.size(0) : weight.size(1), - bsz, - InferenceContext::Instance().GetCurrentStream()); - } - - if (q_int8) { - quantized_gemm(output.data_ptr(), - intermediate, - weight1, - q_scale1, - q_scale1.size(0), - bsz, - input.size(2)); - } else { - float alpha = (T)1.0; - float gemm_beta = (T)0.0; - cublasSetStream(InferenceContext::Instance().GetCublasHandle(), - InferenceContext::Instance().GetCurrentStream()); - cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), - (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), - CUBLAS_OP_N, - weight1.size(transposed_mode ? 0 : 1), - bsz, - weight1.size(transposed_mode ? 1 : 0), - &alpha, - &gemm_beta, - (T*)weight1.data_ptr(), - intermediate, - (T*)output.data_ptr(), -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif - } - - return torch::from_blob(inp_norm, input.sizes(), input.options()); -} - -template -std::vector ds_mlp_gemm(at::Tensor& input, - at::Tensor& residual, - at::Tensor& input_bias, - at::Tensor& weight_interm, - at::Tensor& weight_out, - at::Tensor& bias, - at::Tensor& gamma, - at::Tensor& beta, - const float epsilon, - bool preLayerNorm, - bool mlp_after_attn, - at::Tensor& q_scale, - at::Tensor& q_scale1, - bool q_int8, - int activation_type, - bool transposed_mode) -{ - auto options = at::TensorOptions() - .dtype(input.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - int out_size = (q_int8 || transposed_mode) ? weight_out.size(0) : weight_out.size(1); - auto output = - at::from_blob((T*)InferenceContext::Instance().GetWorkSpace() + torch::numel(input), - {input.size(0), input.size(1), out_size}, - options); - int bsz = input.size(0) * input.size(1); - - auto act_func_type = static_cast(activation_type); - auto res_add = mlp_unfused_cublas(output, - mlp_after_attn ? input : residual, - residual, - input_bias, - weight_interm, - weight_out, - bias, - gamma, - beta, - epsilon, - preLayerNorm, - mlp_after_attn, - q_scale, - q_scale1, - q_int8, - act_func_type, - transposed_mode); - - return {output, res_add}; -} - -template -std::vector ds_rms_mlp_gemm(at::Tensor& input, - at::Tensor& residual, - at::Tensor& weight_interm, - at::Tensor& weight_out, - at::Tensor& gamma, - const float epsilon, - at::Tensor& q_scale, - at::Tensor& q_scale1, - bool q_int8, - int activation_type, - bool transposed_mode) -{ - const int bsz = input.size(0) * input.size(1); - const size_t input_neurons = input.size(2); - const size_t mlp_1_out_neurons = transposed_mode ? weight_interm.size(0) - : weight_interm.size(1); - const size_t mlp_2_in_neurons = transposed_mode ? weight_out.size(1) : weight_out.size(0); - - auto options = at::TensorOptions() - .dtype(input.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - T* output_ptr = (T*)InferenceContext::Instance().GetWorkSpace() + torch::numel(input); - T* inp_norm_ptr = output_ptr + torch::numel(input); - T* intermediate_ptr = inp_norm_ptr + torch::numel(input); - - auto output = at::from_blob(output_ptr, input.sizes(), options); - auto inp_norm = at::from_blob(inp_norm_ptr, input.sizes(), options); - auto intermediate_gemm = - at::from_blob(intermediate_ptr, - {input.size(0), input.size(1), static_cast(mlp_1_out_neurons)}, - options); - - auto act_func_type = static_cast(activation_type); - - // RMS Norm, we'll update the residual in-place - launch_rms_norm((T*)inp_norm.data_ptr(), - (T*)residual.data_ptr(), - (const T*)input.data_ptr(), - (const T*)residual.data_ptr(), - (const T*)gamma.data_ptr(), - epsilon, - bsz, - input_neurons, - InferenceContext::Instance().GetCurrentStream()); - - if (q_int8) { - quantized_gemm(intermediate_ptr, - (T*)inp_norm.data_ptr(), - weight_interm, - q_scale, - q_scale.size(0), - bsz, - input_neurons); - } else { - float alpha = (T)1.0; - float gemm_beta = (T)0.0; - cublasSetStream(InferenceContext::Instance().GetCublasHandle(), - InferenceContext::Instance().GetCurrentStream()); - cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), - (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), - CUBLAS_OP_N, - mlp_1_out_neurons, - bsz, - input_neurons, - &alpha, - &gemm_beta, - (T*)weight_interm.data_ptr(), - (T*)inp_norm.data_ptr(), - intermediate_ptr, -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif - } - - if (act_func_type == ActivationFuncType::GELU) { - launch_bias_gelu(intermediate_ptr, - (T*)nullptr, - mlp_1_out_neurons, - bsz, - InferenceContext::Instance().GetCurrentStream()); - } else if (act_func_type == ActivationFuncType::ReLU) { - launch_bias_relu(intermediate_ptr, - (T*)nullptr, - mlp_1_out_neurons, - bsz, - InferenceContext::Instance().GetCurrentStream()); - } else if (act_func_type == ActivationFuncType::GATED_GELU) { - launch_gated_activation(intermediate_ptr, - (const T*)intermediate_ptr, - (const T*)nullptr, - bsz, - mlp_1_out_neurons, - mlp_1_out_neurons, - true, - InferenceContext::Instance().GetCurrentStream()); - } else if (act_func_type == ActivationFuncType::GATED_SILU) { - launch_gated_activation(intermediate_ptr, - (const T*)intermediate_ptr, - (const T*)nullptr, - bsz, - mlp_1_out_neurons, - mlp_1_out_neurons, - false, - InferenceContext::Instance().GetCurrentStream()); - } - - if (q_int8) { - quantized_gemm(output.data_ptr(), - intermediate_ptr, - weight_out, - q_scale1, - q_scale1.size(0), - bsz, - input.size(2)); - } else { - float alpha = (T)1.0; - float gemm_beta = (T)0.0; - cublasSetStream(InferenceContext::Instance().GetCublasHandle(), - InferenceContext::Instance().GetCurrentStream()); - cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), - (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), - CUBLAS_OP_N, - input_neurons, - bsz, - mlp_2_in_neurons, - &alpha, - &gemm_beta, - (T*)weight_out.data_ptr(), - intermediate_ptr, - (T*)output.data_ptr(), -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard, -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP, -#endif - mlp_1_out_neurons); - } - - return {output, residual}; -} - -template -at::Tensor fused_gemm_gelu(at::Tensor& input, - at::Tensor& weight, - at::Tensor& weight_scale, - at::Tensor& bias, - at::Tensor& weight_out, - at::Tensor& weight_out_scale, - bool q_int8, - bool transposed_mode) -{ - auto options = at::TensorOptions() - .dtype(input.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - int intm_dim = (transposed_mode || q_int8) ? weight.size(0) : weight.size(1); - - // auto output = at::from_blob((T*)InferenceContext::Instance().GetWorkSpace() + - // torch::numel(input), - // {input.size(0), input.size(1), out_size}, - // options); - // T* intermediate = (T*)input.data_ptr() + torch::numel(input); - auto intermediate = at::empty({input.size(0), input.size(1), intm_dim}, options); - - int bsz = input.size(0) * input.size(1); - - float alpha = (T)1.0; - float gemm_beta = (T)0.0; - if (q_int8) { - quantized_gemm(intermediate.data_ptr(), - (T*)input.data_ptr(), - weight, - weight_scale, - weight_scale.size(0), - bsz, - input.size(2)); - } else { - cublasSetStream(InferenceContext::Instance().GetCublasHandle(), - InferenceContext::Instance().GetCurrentStream()); - cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), - (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), - CUBLAS_OP_N, - intm_dim, - bsz, - input.size(2), - &alpha, - &gemm_beta, - (T*)weight.data_ptr(), - (T*)input.data_ptr(), - (T*)intermediate.data_ptr(), -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif - } - launch_bias_gelu((T*)intermediate.data_ptr(), - (T*)bias.data_ptr(), - intm_dim, - bsz, - InferenceContext::Instance().GetCurrentStream()); - - int out_size = (transposed_mode || q_int8) ? weight_out.size(0) : weight_out.size(1); - auto output = at::empty({input.size(0), input.size(1), out_size}, options); - if (q_int8) { - quantized_gemm(output.data_ptr(), - (T*)intermediate.data_ptr(), - weight_out, - weight_out_scale, - weight_out_scale.size(0), - bsz, - input.size(2)); - } else { - cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), - (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), - CUBLAS_OP_N, - out_size, - bsz, - intm_dim, - &alpha, - &gemm_beta, - (T*)weight_out.data_ptr(), - (T*)intermediate.data_ptr(), - (T*)output.data_ptr(), -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_gemm_algo_standard); -#else - CUBLAS_GEMM_DEFAULT_TENSOR_OP); -#endif - } - // cudaEventRecord(InferenceContext::Instance().GetCompEvent(2), - // InferenceContext::Instance().GetCurrentStream(true)); - return output; -} - -template -at::Tensor& residual_add_bias(at::Tensor& hidden_state, - at::Tensor& residual, - const at::Tensor& attention_output, - const at::Tensor& attention_bias, - const at::Tensor& final_bias, - const int mp_size, - const bool mlp_after_attn, - const bool add_bias, - const bool preln) -{ - int bsz = residual.size(0) * residual.size(1); - int hidden_size = residual.size(2); - if (mlp_after_attn) - launch_bias_residual(static_cast(residual.data_ptr()), - static_cast(hidden_state.data_ptr()), - static_cast(attention_output.data_ptr()), - static_cast(final_bias.data_ptr()), - static_cast(attention_bias.data_ptr()), - bsz, - hidden_size, - mp_size, - preln, - InferenceContext::Instance().GetCurrentStream()); - else - launch_gptj_residual_add( - static_cast(residual.data_ptr()), - static_cast(hidden_state.data_ptr()), - static_cast(attention_output.data_ptr()), - static_cast(final_bias.data_ptr()), - static_cast((add_bias ? attention_bias.data_ptr() : nullptr)), - hidden_size, - bsz, - mp_size, - InferenceContext::Instance().GetCurrentStream()); - return residual; -} - -#define DISPATCH_VECTOR_ADD(T_TYPE, C_TYPE) \ - if (a.scalar_type() == at::k##T_TYPE) { \ - launch_vector_add((C_TYPE*)(a.data_ptr()), \ - (const C_TYPE*)(a.data_ptr()), \ - (const C_TYPE*)(b.data_ptr()), \ - gamma, \ - total_elems, \ - InferenceContext::Instance().GetCurrentStream()); \ - } - -at::Tensor& _vector_add(at::Tensor& a, at::Tensor& b, float gamma) -{ - const int total_elems = a.numel(); - - DISPATCH_VECTOR_ADD(Float, float) - DISPATCH_VECTOR_ADD(Half, __half) -#ifdef BF16_AVAILABLE - DISPATCH_VECTOR_ADD(BFloat16, __nv_bfloat16) -#endif - - return a; -} - -std::vector apply_rotary_pos_emb(at::Tensor& mixed_query, - at::Tensor& key_layer, - unsigned rotary_dim, - unsigned offset, - unsigned num_heads, - bool rotate_half, - float rope_theta) -{ - auto query_cont = mixed_query.contiguous(); - auto key_cont = key_layer.contiguous(); - - unsigned bsz = mixed_query.size(0); - unsigned head_size = mixed_query.size(2) / num_heads; - unsigned seq_len = mixed_query.size(1); - - if (mixed_query.scalar_type() == at::kFloat) - launch_apply_rotary_pos_emb((float*)query_cont.data_ptr(), - (float*)key_cont.data_ptr(), - head_size, - seq_len, - rotary_dim, - offset, - num_heads, - bsz, - rope_theta, - InferenceContext::Instance().GetCurrentStream(), - InferenceContext::Instance().GetMaxTokenLength()); - else - launch_apply_rotary_pos_emb<__half>((__half*)query_cont.data_ptr(), - (__half*)key_cont.data_ptr(), - head_size, - seq_len, - rotary_dim, - offset, - num_heads, - bsz, - rope_theta, - InferenceContext::Instance().GetCurrentStream(), - InferenceContext::Instance().GetMaxTokenLength()); - return {query_cont, key_cont}; -} - -#define DISPATCH_MOE_RESIDUAL(T_TYPE, C_TYPE) \ - if (moe_res.scalar_type() == torch::T_TYPE) { \ - launch_moe_res_matmul((C_TYPE*)moe_res.data_ptr(), \ - (C_TYPE*)coef.data_ptr(), \ - (C_TYPE*)output.data_ptr(), \ - M, \ - N, \ - InferenceContext::Instance().GetCurrentStream()); \ - } - -at::Tensor moe_res_matmul(at::Tensor& moe_res, at::Tensor& coef, at::Tensor& output) -{ - int M = moe_res.size(0) * moe_res.size(1); - int N = moe_res.size(2); - InferenceContext::Instance().SynchComm(); - - DISPATCH_MOE_RESIDUAL(kFloat, float) - DISPATCH_MOE_RESIDUAL(kHalf, __half) -#ifdef BF16_AVAILABLE - DISPATCH_MOE_RESIDUAL(kBFloat16, __nv_bfloat16) -#endif - - return output; -} - -void ds_release_workspace() { InferenceContext::Instance().release_workspace(); } - -bool ds_retake_workspace() { return InferenceContext::Instance().retake_workspace(); } - -template -at::Tensor ds_dequantize(at::Tensor& weight, at::Tensor& qscale, int groups) -{ - auto options = at::TensorOptions() - .dtype(torch::kFloat16) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - auto weight16 = at::empty({weight.size(0), weight.size(1)}, options); - - launch_dequantize((T*)weight16.data_ptr(), - (int8_t*)weight.data_ptr(), - (float*)qscale.data_ptr(), - weight.size(0), - weight.size(1), - groups, - InferenceContext::Instance().GetCurrentStream()); - - return weight16; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("softmax_context_int8", - &ds_softmax_context1<__half>, - "DeepSpeed attention with int8 (CUDA)"); - - // The following functions handle type dispatching internally - m.def("gated_activation", &ds_gated_activation, "DeepSpeed Bias GEGLU (CUDA)"); - m.def("layer_norm", &ds_layer_norm, "DeepSpeed layer norm (CUDA)"); - m.def( - "_layer_norm_residual", &ds_layer_norm_residual, "DeepSpeed layer norm + residual (CUDA)"); - m.def("layer_norm_residual_store_pre_ln_res", - &ds_layer_norm_residual_store_pre_ln_res, - "DeepSpeed layer norm + store pre Layernorm residual (CUDA)"); - m.def("rms_norm", &ds_rms_norm, "DeepSpeed rms norm (CUDA)"); - m.def("pre_rms_norm", &ds_pre_rms_norm, "DeepSpeed pre rms norm (CUDA)"); - m.def("_vector_add", &_vector_add, "DeepSpeed vector add (CUDA)"); - m.def("apply_rotary_pos_emb", &apply_rotary_pos_emb, "DeepSpeed mlp with fp16 (CUDA)"); - m.def("moe_res_matmul", &moe_res_matmul, "DeepSpeed moe residual matmul (CUDA)"); - m.def("reset_cache", &reset_cache, "Reset Cache for generation tasks"); - m.def("release_workspace", &ds_release_workspace, "DeepSpeed Release Workspace"); - m.def("retake_workspace", &ds_retake_workspace, "DeepSpeed Retake Workspace"); - - // The following functions are templated and need to be explicitly instantiated and bound - // to different python methods -#define DEF_OPS(_name, _dtype) \ - m.def("softmax_" #_name, &ds_softmax<_dtype>, "DeepSpeed SoftMax with " #_name " (CUDA)"); \ - m.def("softmax_context_" #_name, \ - &ds_softmax_context<_dtype>, \ - "DeepSpeed attention with " #_name " (CUDA)"); \ - m.def("bias_gelu_" #_name, &ds_bias_gelu<_dtype>, "DeepSpeed Gelu with " #_name " (CUDA)"); \ - m.def("bias_add_" #_name, &ds_bias_add<_dtype>, "DeepSpeed Bias Add with " #_name " (CUDA)"); \ - m.def("bias_relu_" #_name, &ds_bias_relu<_dtype>, "DeepSpeed ReLU with " #_name " (CUDA)"); \ - m.def("bias_residual_" #_name, \ - &ds_bias_residual<_dtype>, \ - "DeepSpeed residual-bias add with " #_name " (CUDA)"); \ - m.def("qkv_gemm_" #_name, &ds_qkv_gemm<_dtype>, "DeepSpeed qkv gemm with " #_name " (CUDA)"); \ - m.def("rms_qkv_gemm_" #_name, \ - &ds_rms_qkv<_dtype>, \ - "DeepSpeed rms qkv gemm with " #_name " (CUDA)"); \ - m.def("mlp_gemm_" #_name, &ds_mlp_gemm<_dtype>, "DeepSpeed mlp with " #_name " (CUDA)"); \ - m.def("rms_mlp_gemm_" #_name, \ - &ds_rms_mlp_gemm<_dtype>, \ - "DeepSpeed rms mlp gemm with " #_name " (CUDA)"); \ - m.def("vector_matmul_" #_name, \ - &ds_vector_matmul<_dtype>, \ - "DeepSpeed vector-MM with " #_name " (CUDA)"); \ - m.def("linear_layer_" #_name, \ - &ds_linear_layer<_dtype>, \ - "DeepSpeed linear_layer with " #_name " (CUDA)"); \ - m.def("fused_gemm_gelu_" #_name, \ - &fused_gemm_gelu<_dtype>, \ - "DeepSpeed mlp with " #_name " (CUDA)"); \ - m.def("residual_add_bias_" #_name, \ - &residual_add_bias<_dtype>, \ - "DeepSpeed residual add with " #_name " (CUDA)"); \ - m.def("einsum_sec_sm_ecm_" #_name, \ - &einsum_sec_sm_ecm<_dtype>, \ - "DeepSpeed vector-MM with " #_name " (CUDA)"); \ - m.def("add_padding_" #_name, \ - &add_padding<_dtype>, \ - "DeepSpeed residual add with " #_name " (CUDA)"); \ - m.def("pad_transform_" #_name, \ - &padd_add_transform<_dtype>, \ - "DeepSpeed residual add with " #_name " (CUDA)"); \ - m.def("allocate_workspace_" #_name, \ - &allocate_workspace<_dtype>, \ - "DeepSpeed memory allocation for GPT inference with " #_name " (CUDA)"); \ - m.def("dequantize_" #_name, \ - &ds_dequantize<_dtype>, \ - "DeepSpeed dequantize with " #_name " (CUDA)"); - - DEF_OPS(fp32, float); - DEF_OPS(fp16, __half); -#ifdef BF16_AVAILABLE - DEF_OPS(bf16, __nv_bfloat16); -#endif -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/relu.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/relu.cu deleted file mode 100644 index 40926b776cf2f0bd3c46f44c6eda2ed3ec6bfce5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/relu.cu +++ /dev/null @@ -1,71 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "conversion_utils.h" -#include "inference_cuda_layers.h" -#include "memory_access_utils.h" - -namespace cg = cooperative_groups; -#define MAX_CAP 4 -#define MAX_SEQ 2048 - -inline __device__ float relu(const float x) { return x < 0 ? 0 : x; } - -/* -In-place relu(biasAdd(x)) for channels last -*/ -template -__global__ void fused_bias_relu(T* input, const T* bias, int total_count, int intermediate_size) -{ - // Input restriction: intermediate_size % vals_per_access == 0 - constexpr int granularity = 16; - constexpr int values_per_access = granularity / sizeof(T); - const int offset = (blockIdx.x * blockDim.x + threadIdx.x) * values_per_access; - - if (offset < total_count) { - T data[values_per_access]; - T data_bias[values_per_access]; - mem_access::load_global(data, input + offset); - mem_access::load_global( - data_bias, bias + (offset % intermediate_size), bias != nullptr); - -#pragma unroll - for (int i = 0; i < values_per_access; i++) { - float data_f = conversion::to(data[i]); - float bias_f = conversion::to(data_bias[i]); - data[i] = conversion::to(relu(data_f + bias_f)); - } - - mem_access::store_global(input + offset, data); - } -} - -template -void launch_bias_relu(T* input, - const T* bias, - int intermediate_size, - int batch_size, - cudaStream_t stream) -{ - constexpr int threads = 1024; - constexpr int granularity = 16; - - const int total_count = batch_size * intermediate_size; - const int elems_per_block = threads * (granularity / sizeof(T)); - dim3 block_dims(threads); - dim3 grid_dims((total_count + elems_per_block - 1) / elems_per_block); - - fused_bias_relu<<>>( - input, bias, total_count, intermediate_size); -} - -#define INSTANTIATE_LAUNCH_BIAS_RELU(T) \ - template void launch_bias_relu(T*, const T*, int, int, cudaStream_t); - -INSTANTIATE_LAUNCH_BIAS_RELU(float) -#ifdef BF16_AVAILABLE -INSTANTIATE_LAUNCH_BIAS_RELU(__nv_bfloat16) -#endif -INSTANTIATE_LAUNCH_BIAS_RELU(__half) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/rms_norm.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/rms_norm.cu deleted file mode 100644 index 5f72a4193752cc54b8ee4e1b375e19ef5c3d1214..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/rms_norm.cu +++ /dev/null @@ -1,263 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "conversion_utils.h" -#include "ds_kernel_utils.h" -#include "inference_cuda_layers.h" -#include "memory_access_utils.h" -#include "reduction_utils.h" - -namespace cg = cooperative_groups; -using rop = reduce::ROpType; - -namespace rms { -constexpr int granularity = 16; -} // namespace rms - -template -__global__ void rms_norm(T* output, const T* vals, const T* gamma, float epsilon, int elems_per_row) -{ - constexpr int T_per_load = rms::granularity / sizeof(T); - - cg::thread_block tb = cg::this_thread_block(); - cg::thread_block_tile warp = cg::tiled_partition(tb); - - // X-dimension of the block - const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) + - (tb.thread_index().y * elems_per_row); - const int thread_offset = tb.thread_index().x * T_per_load; - const int base_offset = block_offset + thread_offset; - const int stride = blockDim.x * T_per_load; - - float var_sum = reduce::init(); - - const T* input_base = vals + base_offset; - - T local_buffer[UNROLL * T_per_load]; - -#pragma unroll - for (int i = 0; i < UNROLL; i++) { - T* iteration_buffer = local_buffer + (i * T_per_load); - - mem_access::load_global(iteration_buffer, - input_base + (i * stride), - thread_offset + (i * stride) < elems_per_row); - -#pragma unroll - for (int j = 0; j < T_per_load; j++) { - float up_cast = conversion::to(iteration_buffer[j]); - float sq_val = up_cast * up_cast; - var_sum = reduce::element(var_sum, sq_val); - } - } - - reduce::partitioned_block(tb, warp, var_sum); - const float var = var_sum / elems_per_row; - const T denom = conversion::to(__frsqrt_rn(var + epsilon)); - - T* block_output = output + block_offset; - -#pragma unroll - for (int i = 0; i < UNROLL; i++) { - T* iteration_buffer = local_buffer + (i * T_per_load); - const int iter_idx = i * stride + thread_offset; - const bool do_loads = (iter_idx < elems_per_row); - - T gamma_local[T_per_load]; - - mem_access::load_global(gamma_local, gamma + iter_idx, do_loads); - -#pragma unroll - for (int j = 0; j < T_per_load; j++) { - iteration_buffer[j] *= denom; - iteration_buffer[j] *= gamma_local[j]; - } - - if (do_loads) { - mem_access::store_global(block_output + iter_idx, iteration_buffer); - } - } -} - -template -__global__ void pre_rms_norm(T* output, - T* res_out, - const T* vals, - const T* residual, - const T* gamma, - float epsilon, - int elems_per_row) -{ - constexpr int T_per_load = rms::granularity / sizeof(T); - - cg::thread_block tb = cg::this_thread_block(); - cg::thread_block_tile warp = cg::tiled_partition(tb); - - // X-dimension of the block - const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) + - (tb.thread_index().y * elems_per_row); - const int thread_offset = tb.thread_index().x * T_per_load; - const int base_offset = block_offset + thread_offset; - const int stride = blockDim.x * T_per_load; - - float var_sum = reduce::init(); - - const T* input_base = vals + base_offset; - const T* residual_base = residual + base_offset; - T* res_output = res_out + base_offset; - - T local_buffer[UNROLL * T_per_load]; - -#pragma unroll - for (int i = 0; i < UNROLL; i++) { - T* iteration_buffer = local_buffer + (i * T_per_load); - T residual_buffer[T_per_load]; - - const int iter_offset = i * stride + thread_offset; - const bool do_loads = (iter_offset < elems_per_row); - - mem_access::load_global( - iteration_buffer, input_base + (i * stride), do_loads); - mem_access::load_global( - residual_buffer, residual_base + (i * stride), do_loads); - -#pragma unroll - for (int j = 0; j < T_per_load; j++) { - iteration_buffer[j] += residual_buffer[j]; - float vals_up_cast = conversion::to(iteration_buffer[j]); - - var_sum = reduce::element(var_sum, vals_up_cast * vals_up_cast); - } - - if (do_loads) { - mem_access::store_global(res_output + i * stride, iteration_buffer); - } - } - - reduce::partitioned_block(tb, warp, var_sum); - const float var = var_sum / elems_per_row; - const T denom = conversion::to(__frsqrt_rn(var + epsilon)); - - T* block_output = output + block_offset; - -#pragma unroll - for (int i = 0; i < UNROLL; i++) { - T* iteration_buffer = local_buffer + (i * T_per_load); - const int iter_idx = i * stride + thread_offset; - const bool do_loads = (iter_idx < elems_per_row); - - T gamma_local[T_per_load]; - - mem_access::load_global(gamma_local, gamma + iter_idx, do_loads); - -#pragma unroll - for (int j = 0; j < T_per_load; j++) { - iteration_buffer[j] *= denom; - iteration_buffer[j] *= gamma_local[j]; - } - - if (do_loads) { - mem_access::store_global(block_output + iter_idx, iteration_buffer); - } - } -} - -#define LAUNCH_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ - rms_norm \ - <<>>(norm_output, vals, gamma, epsilon, elems_per_row); - -#define LAUNCH_PRE_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ - pre_rms_norm<<>>( \ - norm_output, res_output, vals, residual, gamma, epsilon, elems_per_row); - -#define LAUNCH_ALL_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ - if (pre_norm) { \ - LAUNCH_PRE_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ - } else { \ - LAUNCH_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ - } - -template -void launch_rms_norm(T* norm_output, - T* res_output, - const T* vals, - const T* residual, - const T* gamma, - float epsilon, - int rows, - int elems_per_row, - cudaStream_t stream) -{ - // 8 for __half, 4 for float - constexpr int T_per_load = rms::granularity / sizeof(T); - constexpr int maxThreads = 256; - constexpr int internalUnroll = sizeof(T) == 4 ? 4 : 2; - - const bool is_subblock_schedule = (elems_per_row <= 128) ? true : false; - const int h_per_step = is_subblock_schedule ? T_per_load : T_per_load * internalUnroll; - - // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of - // warp-sized blocks rather than stepping up to 64/96 threads - const int one_step_threads = next_pow2((elems_per_row + h_per_step - 1) / h_per_step); - const int threads_per_group = (one_step_threads < maxThreads) ? one_step_threads : maxThreads; - - const int groups_per_block_max = - is_subblock_schedule ? (maxThreads + threads_per_group - 1) / threads_per_group : 1; - const int groups_per_block = (rows < groups_per_block_max) ? rows : groups_per_block_max; - const int groups_launch = (groups_per_block + rows - 1) / groups_per_block; - - dim3 block(threads_per_group, groups_per_block); - dim3 grid(groups_launch); - - const int elems_per_step = threads_per_group * h_per_step; - const int external_unRoll = (elems_per_row + elems_per_step - 1) / elems_per_step; - - bool pre_norm = (residual == nullptr) ? false : true; - - if (is_subblock_schedule) { - // <=128 - if (threads_per_group == 1) { - LAUNCH_ALL_RMS_NORM(1, 1, maxThreads); - } else if (threads_per_group == 2) { - LAUNCH_ALL_RMS_NORM(1, 2, maxThreads); - } else if (threads_per_group == 4) { - LAUNCH_ALL_RMS_NORM(1, 4, maxThreads); - } else if (threads_per_group == 8) { - LAUNCH_ALL_RMS_NORM(1, 8, maxThreads); - } else if (threads_per_group == 16) { - LAUNCH_ALL_RMS_NORM(1, 16, maxThreads); - } - } else if (external_unRoll == 1) { - // 129 - 4096 elems - // (this can launch with 1-7 warps as well) - LAUNCH_ALL_RMS_NORM(1 * internalUnroll, maxThreads, maxThreads); - } else if (external_unRoll == 2) { - // 4097 - 8192 elems - LAUNCH_ALL_RMS_NORM(2 * internalUnroll, maxThreads, maxThreads); - } else if (external_unRoll == 3) { - // 8193 - 12288 elems - LAUNCH_ALL_RMS_NORM(3 * internalUnroll, maxThreads, maxThreads); - } else if (external_unRoll == 4) { - // 12289 - 16384 elems - LAUNCH_ALL_RMS_NORM(4 * internalUnroll, maxThreads, maxThreads); - } -} - -#define INSTANTIATE_LAUNCH_RMS_NORM(T) \ - template void launch_rms_norm(T * norm_output, \ - T * res_output, \ - const T* vals, \ - const T* residual, \ - const T* gamma, \ - float epsilon, \ - int rows, \ - int elems_per_row, \ - cudaStream_t stream); - -INSTANTIATE_LAUNCH_RMS_NORM(float) -INSTANTIATE_LAUNCH_RMS_NORM(__half) -#ifdef BF16_AVAILABLE -INSTANTIATE_LAUNCH_RMS_NORM(__nv_bfloat16) -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/softmax.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/softmax.cu deleted file mode 100644 index bb06cc149ef4fb29dae84a7b16173d1babfa5e50..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/softmax.cu +++ /dev/null @@ -1,562 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include "conversion_utils.h" -#include "inference_cuda_layers.h" - -#ifndef __HIP_PLATFORM_AMD__ -#include -#endif -#include -#include -#include - -#define MAX_REG_SIZE 8 - -#define minus_infinity -10000.0 - -void CheckCudaErrorAux(const char* file, unsigned line) -{ - cudaError_t err = cudaGetLastError(); - if (err == cudaSuccess) return; - std::cerr << cudaGetErrorString(err) << "(" << err << ") at " << file << ":" << line - << std::endl; - throw std::runtime_error("CUDA ERROR!!!\n"); -} - -#define CUDA_CHECK_ERROR() CheckCudaErrorAux(__FILE__, __LINE__) - -namespace cg = cooperative_groups; - -template -__global__ void attn_softmax_v2(T* vals, - T* mask, - T* alibi, - float layer_scale, - bool triangular, - bool recompute, - bool local_attention, - int window_size, - int total_count, - int heads, - int sequence_length, - int num_seq, - int head_offset, - int mask_stride, - int mp_size, - int reduceWidth) -{ - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - float2 low_data[MAX_REG_SIZE]; - float2 high_data[MAX_REG_SIZE]; - const T zero_h = conversion::to(0.f); - - int wid = threadIdx.x >> 5; - int lane = threadIdx.x & 0x1f; - int warp_num = blockDim.x >> 5; - - int reduce_blocks = reduceWidth >> 5; - int seq_lane = threadIdx.x % reduceWidth; - - __shared__ float partialSum[MAX_WARP_NUM]; - - int iter_offset = blockIdx.x * (warp_num / reduce_blocks) + (wid / reduce_blocks); - int batch_idx = iter_offset / (num_seq * heads); - int alibi_offset = batch_idx * heads * mp_size + head_offset; - int mask_offset = batch_idx * mask_stride + (iter_offset % mask_stride); - - if (iter_offset < total_count) { - vals += (iter_offset * sequence_length); - - alibi_offset = (alibi_offset + ((iter_offset / num_seq) % heads)) * sequence_length; - mask_offset = mask_offset * sequence_length; - int seq_id = iter_offset % num_seq; - - int real_seq_id = seq_id + (num_seq == sequence_length ? 0 : sequence_length); - int window_stride4 = (local_attention && (real_seq_id >> 2) > (window_size >> 2)) - ? (real_seq_id >> 2) - (window_size >> 2) - : 0; - int window_stride = - (local_attention && real_seq_id >= window_size) ? real_seq_id - window_size : -1; - - float max_val = minus_infinity; - // if (lane == 0) printf("%d, %d: %d \n", wid, blockIdx.x, mask_offset); - for (int i = 0; i < iterations; i++) { - int data_id = i * (reduceWidth << 2) + (seq_lane); - bool check = (data_id >> 2) >= window_stride4; - bool low_x_check = check && (data_id < sequence_length) && - (!triangular || (data_id <= seq_id)) && (data_id > window_stride); - bool low_y_check = check && ((data_id + reduceWidth) < sequence_length) && - (!triangular || ((data_id + reduceWidth) <= seq_id)) && - ((data_id + reduceWidth) > window_stride); - bool high_x_check = check && ((data_id + reduceWidth * 2) < sequence_length) && - (!triangular || ((data_id + reduceWidth * 2) <= seq_id)) && - ((data_id + reduceWidth * 2) > window_stride); - bool high_y_check = check && ((data_id + reduceWidth * 3) < sequence_length) && - (!triangular || ((data_id + reduceWidth * 3) <= seq_id)) && - ((data_id + reduceWidth * 3) > window_stride); - - if (mask && alibi) { - low_data[i].x = low_x_check - ? conversion::to(vals[data_id]) * layer_scale + - (conversion::to(alibi[data_id + alibi_offset])) + - (conversion::to(mask[data_id + mask_offset])) - : minus_infinity; - low_data[i].y = - low_y_check - ? conversion::to(vals[data_id + reduceWidth]) * layer_scale + - (conversion::to(alibi[data_id + alibi_offset + reduceWidth])) + - (conversion::to(mask[data_id + mask_offset + reduceWidth])) - : minus_infinity; - high_data[i].x = - high_x_check - ? conversion::to(vals[data_id + reduceWidth * 2]) * layer_scale + - (conversion::to( - alibi[data_id + alibi_offset + reduceWidth * 2])) + - (conversion::to(mask[data_id + mask_offset + reduceWidth * 2])) - : minus_infinity; - high_data[i].y = - high_y_check - ? conversion::to(vals[data_id + reduceWidth * 3]) * layer_scale + - (conversion::to( - alibi[data_id + alibi_offset + reduceWidth * 3])) + - (conversion::to(mask[data_id + mask_offset + reduceWidth * 3])) - : minus_infinity; - } else if (mask) { - low_data[i].x = low_x_check - ? conversion::to(vals[data_id]) * layer_scale + - (conversion::to(mask[data_id + mask_offset])) - : minus_infinity; - low_data[i].y = - low_y_check - ? conversion::to(vals[data_id + reduceWidth]) * layer_scale + - (conversion::to(mask[data_id + mask_offset + reduceWidth])) - : minus_infinity; - high_data[i].x = - high_x_check - ? conversion::to(vals[data_id + reduceWidth * 2]) * layer_scale + - (conversion::to(mask[data_id + mask_offset + reduceWidth * 2])) - : minus_infinity; - high_data[i].y = - high_y_check - ? conversion::to(vals[data_id + reduceWidth * 3]) * layer_scale + - (conversion::to(mask[data_id + mask_offset + reduceWidth * 3])) - : minus_infinity; - } else if (alibi) { - low_data[i].x = low_x_check - ? conversion::to(vals[data_id]) * layer_scale + - (conversion::to(alibi[data_id + alibi_offset])) - : minus_infinity; - low_data[i].y = - low_y_check - ? conversion::to(vals[data_id + reduceWidth]) * layer_scale + - (conversion::to(alibi[data_id + alibi_offset + reduceWidth])) - : minus_infinity; - high_data[i].x = - high_x_check - ? conversion::to(vals[data_id + reduceWidth * 2]) * layer_scale + - (conversion::to( - alibi[data_id + alibi_offset + reduceWidth * 2])) - : minus_infinity; - high_data[i].y = - high_y_check - ? conversion::to(vals[data_id + reduceWidth * 3]) * layer_scale + - (conversion::to( - alibi[data_id + alibi_offset + reduceWidth * 3])) - : minus_infinity; - } else { - low_data[i].x = low_x_check ? conversion::to(vals[data_id]) * layer_scale - : minus_infinity; - low_data[i].y = - low_y_check ? conversion::to(vals[data_id + reduceWidth]) * layer_scale - : minus_infinity; - high_data[i].x = - high_x_check - ? conversion::to(vals[data_id + reduceWidth * 2]) * layer_scale - : minus_infinity; - high_data[i].y = - high_y_check - ? conversion::to(vals[data_id + reduceWidth * 3]) * layer_scale - : minus_infinity; - } - - // if(lane == 0) printf("%f , %d, %d \n", low_data[i].x, data_id, seq_id); - max_val = (low_data[i].x > max_val ? low_data[i].x : max_val); - max_val = (low_data[i].y > max_val ? low_data[i].y : max_val); - max_val = (high_data[i].x > max_val ? high_data[i].x : max_val); - max_val = (high_data[i].y > max_val ? high_data[i].y : max_val); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) { - auto temp = g.shfl_xor(max_val, i); - max_val = (temp > max_val ? temp : max_val); - } - - if (reduceWidth > WARP_SIZE) { - if (lane == 0) partialSum[wid] = max_val; - b.sync(); - - if (lane < warp_num) max_val = partialSum[lane]; - - b.sync(); - - for (int i = 1; i < reduce_blocks; i *= 2) { - auto temp = g.shfl_xor(max_val, i); - max_val = (temp > max_val ? temp : max_val); - } - - max_val = g.shfl(max_val, threadIdx.x / WARP_SIZE); - } - float sum = 0; - for (int i = 0; i < iterations; i++) { - low_data[i].x = __expf(low_data[i].x - max_val); - low_data[i].y = __expf(low_data[i].y - max_val); - high_data[i].x = __expf(high_data[i].x - max_val); - high_data[i].y = __expf(high_data[i].y - max_val); - - sum += (low_data[i].x + low_data[i].y + high_data[i].x + high_data[i].y); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) sum += g.shfl_xor(sum, i); - - if (reduceWidth > WARP_SIZE) { - if (lane == 0) partialSum[wid] = sum; - b.sync(); - - if (lane < warp_num) sum = partialSum[lane]; - - b.sync(); - - for (int i = 1; i < reduce_blocks; i *= 2) { sum += g.shfl_xor(sum, i); } - - sum = g.shfl(sum, threadIdx.x / WARP_SIZE); - } - sum += 1e-6; - for (int i = 0; i < iterations; i++) { - int data_id = i * (reduceWidth << 2) + (seq_lane); - if (data_id < sequence_length) { - vals[data_id] = conversion::to(low_data[i].x / sum); - if ((data_id + reduceWidth) < sequence_length) - vals[data_id + reduceWidth] = conversion::to(low_data[i].y / sum); - if ((data_id + reduceWidth * 2) < sequence_length) - vals[data_id + reduceWidth * 2] = conversion::to(high_data[i].x / sum); - if ((data_id + reduceWidth * 3) < sequence_length) - vals[data_id + reduceWidth * 3] = conversion::to(high_data[i].y / sum); - } - } - } -} - -template -__global__ void attn_softmax_v2(float* vals, - float* attn_mask, - float* alibi, - float layer_scale, - bool triangular, - bool recompute, - bool local_attention, - int window_size, - int total_count, - int heads, - int sequence_length, - int num_seq, - int head_offset, - int mask_stride, - int mp_size, - int reduceWidth) -{ - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - float4 data[MAX_REG_SIZE]; - - int wid = threadIdx.x >> 5; - int lane = threadIdx.x & 0x1f; - int warp_num = blockDim.x >> 5; - - int reduce_blocks = reduceWidth >> 5; - int seq_lane = threadIdx.x % reduceWidth; - - __shared__ float partialSum[MAX_WARP_NUM]; - - int iter_offset = blockIdx.x * (warp_num / reduce_blocks) + (wid / reduce_blocks); - if (iter_offset < total_count) { - vals += (iter_offset * sequence_length); - - int batch_idx = iter_offset / (num_seq * heads); - int mask_offset = batch_idx * mask_stride + (iter_offset % mask_stride); - mask_offset = mask_offset * sequence_length; - int seq_id = iter_offset % num_seq; - - int real_seq_id = seq_id + (num_seq == sequence_length ? 0 : sequence_length); - int window_stride4 = (local_attention && (real_seq_id >> 2) > (window_size >> 2)) - ? (real_seq_id >> 2) - (window_size >> 2) - : 0; - int window_stride = - (local_attention && real_seq_id >= window_size) ? real_seq_id - window_size : -1; - - float max_val = minus_infinity; - - for (int i = 0; i < iterations; i++) { - int data_id = i * (reduceWidth << 2) + (seq_lane); - bool check = (data_id >> 2) >= window_stride4; - bool x_check = check && (data_id < sequence_length) && - (!triangular || (data_id <= seq_id)) && (data_id > window_stride); - bool y_check = check && ((data_id + reduceWidth) < sequence_length) && - (!triangular || ((data_id + reduceWidth) <= seq_id)) && - ((data_id + reduceWidth) > window_stride); - bool z_check = check && ((data_id + reduceWidth * 2) < sequence_length) && - (!triangular || ((data_id + reduceWidth * 2) <= seq_id)) && - ((data_id + reduceWidth * 2) > window_stride); - bool w_check = check && ((data_id + reduceWidth * 3) < sequence_length) && - (!triangular || ((data_id + reduceWidth * 3) <= seq_id)) && - ((data_id + reduceWidth * 3) > window_stride); - - if (attn_mask) { - data[i].x = x_check ? vals[data_id] + attn_mask[data_id + mask_offset] - : minus_infinity; - data[i].y = y_check ? vals[data_id + reduceWidth] + - attn_mask[data_id + mask_offset + reduceWidth] - : minus_infinity; - data[i].z = z_check ? vals[data_id + reduceWidth * 2] + - attn_mask[data_id + mask_offset + reduceWidth * 2] - : minus_infinity; - data[i].w = w_check ? vals[data_id + reduceWidth * 3] + - attn_mask[data_id + mask_offset + reduceWidth * 3] - : minus_infinity; - } else { - data[i].x = x_check ? vals[data_id] : minus_infinity; - data[i].y = y_check ? vals[data_id + reduceWidth] : minus_infinity; - data[i].z = z_check ? vals[data_id + reduceWidth * 2] : minus_infinity; - data[i].w = w_check ? vals[data_id + reduceWidth * 3] : minus_infinity; - } - - max_val = (data[i].x > max_val ? data[i].x : max_val); - max_val = (data[i].y > max_val ? data[i].y : max_val); - max_val = (data[i].z > max_val ? data[i].z : max_val); - max_val = (data[i].w > max_val ? data[i].w : max_val); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) { - auto temp = g.shfl_xor(max_val, i); - max_val = (temp > max_val ? temp : max_val); - } - - if (reduceWidth > WARP_SIZE) { - if (lane == 0) partialSum[wid] = max_val; - b.sync(); - - if (lane < warp_num) max_val = partialSum[lane]; - - b.sync(); - - for (int i = 1; i < reduce_blocks; i *= 2) { - auto temp = g.shfl_xor(max_val, i); - max_val = (temp > max_val ? temp : max_val); - } - - max_val = g.shfl(max_val, threadIdx.x / WARP_SIZE); - } - - float sum = 0; - for (int i = 0; i < iterations; i++) { - data[i].x = __expf(data[i].x - max_val); - data[i].y = __expf(data[i].y - max_val); - data[i].z = __expf(data[i].z - max_val); - data[i].w = __expf(data[i].w - max_val); - - sum += (data[i].x + data[i].y + data[i].z + data[i].w); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) sum += g.shfl_xor(sum, i); - - if (reduceWidth > WARP_SIZE) { - if (lane == 0) partialSum[wid] = sum; - b.sync(); - - if (lane < warp_num) sum = partialSum[lane]; - - b.sync(); - - for (int i = 1; i < reduce_blocks; i *= 2) { sum += g.shfl_xor(sum, i); } - - sum = g.shfl(sum, threadIdx.x / WARP_SIZE); - } - sum += 1e-6; - - for (int i = 0; i < iterations; i++) { - int data_id = i * (reduceWidth << 2) + (seq_lane); - if (data_id < sequence_length) { - vals[data_id] = data[i].x / sum; - if ((data_id + reduceWidth) < sequence_length) - vals[data_id + reduceWidth] = data[i].y / sum; - if ((data_id + reduceWidth * 2) < sequence_length) - vals[data_id + reduceWidth * 2] = data[i].z / sum; - if ((data_id + reduceWidth * 3) < sequence_length) - vals[data_id + reduceWidth * 3] = data[i].w / sum; - } - } - } -} - -#define LAUNCH_ATTN_SOFTMAX_V2(iterations) \ - attn_softmax_v2<<>>(vals, \ - mask, \ - alibi, \ - layer_scale, \ - triangular, \ - recompute, \ - local_attention, \ - window_size, \ - total_count, \ - heads, \ - sequence_length, \ - num_seq, \ - head_offset, \ - mask_stride, \ - mp_size, \ - reduce_width); - -template -void launch_attn_softmax_v2(T* vals, - T* mask, - T* alibi, - float layer_scale, - bool triangular, - bool recompute, - bool local_attention, - int window_size, - int batch_size, - int heads, - int num_seq, - int sequence_length, - int head_offset, - int mask_stride, - int mp_size, - cudaStream_t stream) -{ - const int total_count = batch_size * heads * num_seq; - - // Scheduling Overview - // 4 element unroll with power of 2 `reduce_width` threads to a ceiling of `attn_threads` - // Each block should be partitioned into as many `reduce_width` blocks - // as can be fit. - constexpr int attn_threads = 256; - constexpr int min_reduce_width = hw_warp_size; - constexpr int internal_unroll = 4; - - // Handle internal unroll then round to next power of 2. Bump up to minimum granularity. - const int thread_steps_rounded = - next_pow2((sequence_length + internal_unroll - 1) / internal_unroll); - const int thread_steps_schedule = - (thread_steps_rounded < min_reduce_width) ? min_reduce_width : thread_steps_rounded; - // Bound reduce width to the number of threads - const int reduce_width = (thread_steps_schedule < attn_threads) ? thread_steps_schedule - : attn_threads; - // Scale for the excess - const int iterations = thread_steps_schedule / reduce_width; - // Should be safe since reduce_width is capped to attn_threads - const int partitions = attn_threads / reduce_width; - - // Launch params - dim3 grid((total_count + partitions - 1) / partitions); - dim3 block(attn_threads); - - if (sequence_length <= 32768) { - if (iterations == 1) { - LAUNCH_ATTN_SOFTMAX_V2(1); - } else if (iterations == 2) { - LAUNCH_ATTN_SOFTMAX_V2(2); - } else if (iterations == 4) { - LAUNCH_ATTN_SOFTMAX_V2(4); - } else if (iterations == 8) { - LAUNCH_ATTN_SOFTMAX_V2(8); - } else if (iterations == 16) { - LAUNCH_ATTN_SOFTMAX_V2(16); - } else if (iterations == 32) { - LAUNCH_ATTN_SOFTMAX_V2(32); - } else if (iterations == 64) { - LAUNCH_ATTN_SOFTMAX_V2(64); - } - } else - throw std::runtime_error("Unsupport Seq_Length!"); -} - -#define INSTANTIATE_LAUNCH_ATTN_SOFTMAX_V2(T) \ - template void launch_attn_softmax_v2(T* vals, \ - T* mask, \ - T* alibi, \ - float layer_scale, \ - bool triangular, \ - bool recompute, \ - bool local_attention, \ - int window_size, \ - int batch_size, \ - int heads, \ - int num_seq, \ - int sequence_length, \ - int head_offset, \ - int mask_stride, \ - int mp_size, \ - cudaStream_t stream); - -INSTANTIATE_LAUNCH_ATTN_SOFTMAX_V2(float); -#ifdef BF16_AVAILABLE -INSTANTIATE_LAUNCH_ATTN_SOFTMAX_V2(__nv_bfloat16); -#endif -INSTANTIATE_LAUNCH_ATTN_SOFTMAX_V2(__half); - -#define DEF_ATTN_SOFTMAX_V2_HALF(_iter) \ - template __global__ void attn_softmax_v2<__half, _iter>(__half * vals, \ - __half * mask, \ - __half * alibi, \ - float layer_scale, \ - bool triangular, \ - bool recompute, \ - bool local_attention, \ - int window_size, \ - int total_count, \ - int heads, \ - int sequence_length, \ - int num_seq, \ - int head_offset, \ - int mask_stride, \ - int mp_size, \ - int reduceWidth) - -#define DEF_ATTN_SOFTMAX_V2_BF16(_iter) \ - template __global__ void attn_softmax_v2<__nv_bfloat16, _iter>(__nv_bfloat16 * vals, \ - __nv_bfloat16 * mask, \ - __nv_bfloat16 * alibi, \ - float layer_scale, \ - bool triangular, \ - bool recompute, \ - bool local_attention, \ - int window_size, \ - int total_count, \ - int heads, \ - int sequence_length, \ - int num_seq, \ - int head_offset, \ - int mask_stride, \ - int mp_size, \ - int reduceWidth) - -#define FOREACH_ITERATIONS(cb) \ - cb(1); \ - cb(2); \ - cb(4); \ - cb(8); \ - cb(16); \ - cb(32); \ - cb(64) - -FOREACH_ITERATIONS(DEF_ATTN_SOFTMAX_V2_HALF); -#ifdef BF16_AVAILABLE -FOREACH_ITERATIONS(DEF_ATTN_SOFTMAX_V2_BF16); -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/transform.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/transform.cu deleted file mode 100644 index e7624363021eef46738cc3a144af987c16aa51d5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/transform.cu +++ /dev/null @@ -1,733 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#ifndef __HIP_PLATFORM_AMD__ -#include -#endif -#include "conversion_utils.h" -#include "inference_cuda_layers.h" -namespace cg = cooperative_groups; - -// only used to avoid compilation error due to lack of definition. -#ifndef BF16_AVAILABLE -#if defined(__CUDA_BF16_H__) -static_assert(sizeof(__nv_bfloat162) == sizeof(__half2), - "CUDA's __nv_bfloat162 doesn't match __half2 size"); -#else -// Fallback to simple typedef only if CUDA doesn't provide it -using __nv_bfloat162 = __half2; -#endif -#endif - -// Bias add - -__global__ void bias_add_transform_0213(float* output, - float* k_cache, - float* v_cache, - const float* vals, - const float* bias, - int hidden_dim, - int seq_length, - unsigned seq_offset, - int heads, - int head_stride, - int num_kv, - int rotary_dim, - bool rotate_half, - bool rotate_every_two, - int head_ext, - int max_out_tokens, - float rope_theta) -{ - int d0_stride = hidden_dim * seq_length; - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - - int d0 = blockIdx.x; // Batch - int d1 = blockIdx.y; // Sequence ID (0-127) - int cnt = blockIdx.z / head_ext; // Hidden count - int d2 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head (0-11) - int d3 = threadIdx.x; // Values (groups of 4) - - int d2_out_stride = d2_stride * (cnt == 0 ? seq_length : max_out_tokens); - int d0_out_stride = hidden_dim * (cnt == 0 ? seq_length : max_out_tokens); - - const float4* vals_vec = reinterpret_cast(vals); - float4* output_vec = - reinterpret_cast(cnt == 0 ? output : (cnt == 1 ? k_cache : v_cache)); - - vals_vec += (d0 * (d1_stride + num_kv * 2 * d2_stride) * seq_length); - vals_vec += d1 * (d1_stride + num_kv * 2 * d2_stride); - vals_vec += (cnt == 0 ? 0 : d1_stride) + (cnt == 0 ? 0 : (cnt - 1) * num_kv * d2_stride); - vals_vec += ((cnt == 0 ? d2 : (d2 / head_stride)) * d2_stride); - - output_vec += (d1 * d2_stride); - output_vec += (d0 * d0_out_stride); - output_vec += (d2 * d2_out_stride); - - unsigned seq_id = d1 + seq_offset; - float4 inputs = vals_vec[d3]; - int lane = d3 & 0x1f; - if (cnt < 2 && rotary_dim > 0 && d3 < rotary_dim) { - float4 q = vals_vec[d3]; - float2* q_f = reinterpret_cast(&q); - if (rotate_every_two) { -#pragma unroll - for (int o = 0; o < 2; o++) { - float inv_freq = (float)(((d3 << 1) + o) * 2) / (float)(rotary_dim << 2); - inv_freq = 1.0 / powf(rope_theta, inv_freq) * (float)seq_id; - q_f[o].x = (-1.0 * q_f[o].y * sinf(inv_freq) + q_f[o].x * cosf(inv_freq)); - q_f[o].y = (q_f[o].x * sinf(inv_freq) + q_f[o].y * cosf(inv_freq)); - } - } - output_vec[d3] = q; - } else - output_vec[d3] = inputs; -} - -#define ATTN_H 3 -#define MAX_SEQ_LINE 10 - -template -__global__ void bias_add_transform_0213(T* output, // q - T* k_cache, - T* v_cache, - const T* vals, // qkv - const T* bias, - int hidden_dim, - int seq_length, - unsigned seq_offset, - int all_tokens, - int heads, - int head_stride, - int num_kv, - int rotary_dim, - bool rotate_half, - bool rotate_every_two, - int head_ext, - int max_out_tokens, - float rope_theta) -{ - using T2 = - typename std::conditional::value, __half2, __nv_bfloat162>::type; - unsigned half_dim = (rotary_dim << 3) >> 1; - int d0_stride = hidden_dim * seq_length; - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - - int d0 = blockIdx.x; // Batch - int d1 = blockIdx.y; // Sequence ID (0-127) - int cnt = blockIdx.z / head_ext; // Hidden count - int d2 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head (0-11) - int d3 = threadIdx.x; // Values (groups of 4) - - int d2_out_stride = d2_stride * (cnt == 0 ? seq_length : max_out_tokens); - int d0_out_stride = hidden_dim * (cnt == 0 ? seq_length : max_out_tokens); - - float4 vals_arr; - float4 output_arr; - - T2* vals_half = reinterpret_cast(&vals_arr); - T2* output_half = reinterpret_cast(&output_arr); - - const float4* vals_vec = reinterpret_cast(vals); - float4* output_vec = - reinterpret_cast(cnt == 0 ? output : (cnt == 1 ? k_cache : v_cache)); - - vals_vec += (d0 * (d1_stride + num_kv * 2 * d2_stride) * seq_length); - vals_vec += (d1 * (d1_stride + num_kv * 2 * d2_stride)); - vals_vec += (cnt == 0 ? 0 : d1_stride) + (cnt == 0 ? 0 : (cnt - 1) * num_kv * d2_stride); - vals_vec += ((cnt == 0 ? d2 : (d2 / head_stride)) * d2_stride); - - output_vec += (d1 * d2_stride); - output_vec += (d0 * d0_out_stride); - output_vec += (d2 * d2_out_stride); - - unsigned seq_id = d1 + seq_offset; - - int lane = d3 & 0x1f; - if (cnt < 2 && rotary_dim > 0 && d3 < rotary_dim) { - float4 q = vals_vec[d3]; - T2* q_h = reinterpret_cast(&q); - if (rotate_every_two) { -#pragma unroll - for (int o = 0; o < 4; o++) { - float inv_freq = (float)(((d3 << 2) + o) * 2) / (float)(rotary_dim << 3); - inv_freq = 1.0 / powf(rope_theta, inv_freq) * (float)seq_id; - float q_data[2]; - q_data[0] = conversion::to(q_h[o].x); - q_data[1] = conversion::to(q_h[o].y); - q_h[o].x = conversion::to(-1.0 * q_data[1] * sinf(inv_freq) + - q_data[0] * cosf(inv_freq)); - q_h[o].y = - conversion::to(q_data[0] * sinf(inv_freq) + q_data[1] * cosf(inv_freq)); - } - } - output_vec[d3] = q; - } else - output_vec[d3] = vals_vec[d3]; -} - -// [B S C*H] - > C * [B A S N] -template <> -void launch_bias_add_transform_0213(float* output, - float* k_cache, - float* v_cache, - const float* vals, - const float* bias, - int batch_size, - int seq_length, - unsigned seq_offset, - int all_tokens, - int hidden_dim, - int heads, - int num_kv, - int rotary_dim, - bool rotate_half, - bool rotate_every_two, - cudaStream_t stream, - int trans_count, - int max_out_tokens, - float rope_theta) -{ - hidden_dim >>= 2; - int head_ext = (hidden_dim - 1) / MAX_THREADS + 1; - - dim3 block_dim(hidden_dim / heads, (heads / head_ext)); - dim3 grid_dim(batch_size, seq_length, (trans_count * head_ext)); - - bias_add_transform_0213<<>>(output, - k_cache, - v_cache, - vals, - bias, - hidden_dim, - seq_length, - seq_offset, - heads, - num_kv > 0 ? (heads / num_kv) : 1, - num_kv > 0 ? num_kv : heads, - rotary_dim >> 2, - rotate_half, - rotate_every_two, - head_ext, - max_out_tokens, - rope_theta); -} - -template -void launch_bias_add_transform_0213(T* output, - T* k_cache, - T* v_cache, - const T* vals, - const T* bias, - int batch_size, - int seq_length, - unsigned seq_offset, - int all_tokens, - int hidden_dim, - int heads, - int num_kv, - int rotary_dim, - bool rotate_half, - bool rotate_every_two, - cudaStream_t stream, - int trans_count, - int max_out_tokens, - float rope_theta) -{ - hidden_dim >>= 3; - int head_ext = 1; // (hidden_dim - 1) / MAX_THREADS + 1; - dim3 block_dim(hidden_dim / heads, (heads / head_ext)); - dim3 grid_dim(batch_size, seq_length, (trans_count * head_ext)); - bias_add_transform_0213<<>>(output, - k_cache, - v_cache, - vals, - bias, - hidden_dim, - seq_length, - seq_offset, - all_tokens, - heads, - num_kv > 0 ? (heads / num_kv) : 1, - num_kv > 0 ? num_kv : heads, - rotary_dim >> 3, - rotate_half, - rotate_every_two, - head_ext, - max_out_tokens, - rope_theta); -} - -#define INSTANTIATE_LAUNCH_BIAS_ADD_TRANSFORM_0213(T) \ - template void launch_bias_add_transform_0213(T*, \ - T*, \ - T*, \ - const T*, \ - const T*, \ - int, \ - int, \ - unsigned, \ - int, \ - int, \ - int, \ - int, \ - int, \ - bool, \ - bool, \ - cudaStream_t, \ - int, \ - int, \ - float) - -#ifdef BF16_AVAILABLE -INSTANTIATE_LAUNCH_BIAS_ADD_TRANSFORM_0213(__nv_bfloat16); -#endif -INSTANTIATE_LAUNCH_BIAS_ADD_TRANSFORM_0213(__half); - -// Bias add - -__global__ void pad_add_transform_0213(float* output, - const float* vals, - int hidden_dim, - int seq_length, - int padded_seq_len, - int heads, - int padded_head_size) -{ -} - -template -__global__ void pad_add_transform_0213(T* output, - const T* vals, - int hidden_dim, - int seq_length, - int padded_seq_len, - int heads, - int padded_head_size) -{ - using T2 = - typename std::conditional::value, __half2, __nv_bfloat162>::type; - float4 ZERO; - const T2 zero_h = conversion::to(0.f); - T2* ZERO_h = reinterpret_cast(&ZERO); -#pragma unroll - for (int i = 0; i < 4; i++) ZERO_h[i] = zero_h; - - int d0_stride = hidden_dim * seq_length; - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - - int d0 = blockIdx.x; // Batch - int d1 = blockIdx.y * blockDim.z + threadIdx.z; // Sequence ID (0-127) - int d2 = threadIdx.y; // Head (0-11) - int d3 = threadIdx.x; // Values (groups of 4) - - int d2_out_stride = padded_head_size * padded_seq_len; - int d0_out_stride = heads * d2_out_stride; - - const float4* vals_vec = reinterpret_cast(vals); - float4* output_vec = reinterpret_cast(output); - - vals_vec += (d0 * d0_stride); - vals_vec += (d1 * d1_stride); - vals_vec += (d2 * d2_stride); - - output_vec += (d1 * padded_head_size); - output_vec += (d0 * d0_out_stride); - output_vec += (d2 * d2_out_stride); - - if (d3 < d2_stride && d1 < seq_length) - output_vec[d3] = vals_vec[d3]; - else - output_vec[d3] = ZERO; -} - -// [B S C*H] - > C * [B A S N] -template <> -void launch_pad_add_transform_0213(float* output, - const float* vals, - int batch_size, - int hidden_dim, - int seq_length, - int padded_seq_len, - int heads, - int padded_head_size, - cudaStream_t stream) -{ -} - -template -void launch_pad_add_transform_0213(T* output, - const T* vals, - int batch_size, - int hidden_dim, - int seq_length, - int padded_seq_len, - int heads, - int padded_head_size, - cudaStream_t stream) -{ - hidden_dim >>= 3; - dim3 block_dim((padded_head_size >> 3), heads, 2); - dim3 grid_dim(batch_size, padded_seq_len / 2); - pad_add_transform_0213<<>>( - output, vals, hidden_dim, seq_length, padded_seq_len, heads, padded_head_size >> 3); -} - -#define INSTANTIATE_LAUNCH_PAD_ADD_TRANSFORM_0213_SIMPLE(T) \ - template void launch_pad_add_transform_0213( \ - T*, const T*, int, int, int, int, int, int, cudaStream_t); - -INSTANTIATE_LAUNCH_PAD_ADD_TRANSFORM_0213_SIMPLE(__half); -#ifdef BF16_AVAILABLE -INSTANTIATE_LAUNCH_PAD_ADD_TRANSFORM_0213_SIMPLE(__nv_bfloat16); -#endif - -// Bias add -template -__global__ void bias_add_transform_0213(T* output, - const T* vals, - const T* bias, - int hidden_dim, - int seq_length, - int heads, - int head_ext); - -template <> -__global__ void bias_add_transform_0213(float* output, - const float* vals, - const float* bias, - int hidden_dim, - int seq_length, - int heads, - int head_ext) -{ - int d0_stride = hidden_dim * seq_length; - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - - int d0_out_stride = d0_stride; - int d1_out_stride = d2_stride; - int d2_out_stride = d2_stride * seq_length; - - int d0 = blockIdx.x; // Batch - int d1 = blockIdx.y; // Sequence ID (0-127) - int cnt = blockIdx.z / head_ext; // Hidden count - int d2 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head (0-11) - int d3 = threadIdx.x; // Values (groups of 4) - - const float4* vals_vec = reinterpret_cast(vals); - const float4* bias_vec = reinterpret_cast(bias); - float4* output_vec = reinterpret_cast(output); - - float4 inputs = vals_vec[d0 * d0_stride * (gridDim.z / head_ext) + cnt * d1_stride + - d1 * d1_stride * (gridDim.z / head_ext) + d2 * d2_stride + d3]; - float4 biases = bias_vec[cnt * d1_stride + d2 * d2_stride + d3]; - - float4 outputs; - outputs.x = inputs.x + biases.x; - outputs.y = inputs.y + biases.y; - outputs.z = inputs.z + biases.z; - outputs.w = inputs.w + biases.w; - - output_vec[cnt * d0_out_stride * gridDim.x + d0 * d0_out_stride + d1 * d1_out_stride + - d2 * d2_out_stride + d3] = outputs; -} - -template -__global__ void bias_add_transform_0213(T* output, - const T* vals, - const T* bias, - int hidden_dim, - int seq_length, - int heads, - int head_ext) -{ - using T2 = - typename std::conditional::value, __half2, __nv_bfloat162>::type; - int d0_stride = hidden_dim * seq_length; - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - - int d2_out_stride = d2_stride * seq_length; - - int d0 = blockIdx.x; // Batch - int d1 = blockIdx.y; // Sequence ID (0-127) - int cnt = blockIdx.z / head_ext; // Hidden count - int d2 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head (0-11) - int d3 = threadIdx.x; // Values (groups of 4) - - float4 vals_arr; - float4 bias_arr; - float4 output_arr; - T2* vals_half = reinterpret_cast(&vals_arr); - T2* bias_half = reinterpret_cast(&bias_arr); - T2* output_half = reinterpret_cast(&output_arr); - - const float4* vals_vec = reinterpret_cast(vals); - const float4* bias_vec = reinterpret_cast(bias); - float4* output_vec = reinterpret_cast(output); - - vals_vec += (d0 * d0_stride * (gridDim.z / head_ext)); - vals_vec += (d1 * d1_stride * (gridDim.z / head_ext)); - vals_vec += (cnt * d1_stride); - vals_vec += (d2 * d2_stride); - - bias_vec += (cnt * d1_stride); - bias_vec += (d2 * d2_stride); - - output_vec += (cnt * d0_stride * gridDim.x); - output_vec += (d1 * d2_stride); - output_vec += (d0 * d0_stride); - output_vec += (d2 * d2_out_stride); - - bias_arr = bias_vec[d3]; - vals_arr = vals_vec[d3]; - - output_half[0] = vals_half[0] + bias_half[0]; - output_half[1] = vals_half[1] + bias_half[1]; - output_half[2] = vals_half[2] + bias_half[2]; - output_half[3] = vals_half[3] + bias_half[3]; - output_vec[d3] = output_arr; -} - -template -__global__ void bias_add_transform_0213_v2(T* output, - const T* vals, - const T* bias, - int hidden_dim, - int seq_length, - int heads) -{ - using T2 = - typename std::conditional::value, __half2, __nv_bfloat162>::type; - __shared__ float4 in_data[3072]; - - int d0_stride = hidden_dim * seq_length; - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - int iteration_stride = d1_stride * blockDim.z; // Hidden * 3 / 8 - int batch_stride = d0_stride * blockDim.z; // Hidden * S * 3 / 8 - - int d0_out_stride = d0_stride; - int d1_out_stride = d2_stride; - int d2_out_stride = d2_stride * seq_length; - - int d0 = blockIdx.x; // Batch - int d1 = blockIdx.y; // Sequence ID (0-127) - int cnt = threadIdx.z; // blockIdx.z; // Hidden count - int d2 = threadIdx.y; // Head (0-11) - int d3 = threadIdx.x; // Values (groups of 4) - - float4 vals_arr[1]; - float4 bias_arr[1]; - float4 output_arr[1]; - T2* vals_half = reinterpret_cast(vals_arr); - T2* bias_half = reinterpret_cast(bias_arr); - T2* output_half = reinterpret_cast(output_arr); - - const float4* vals_vec = reinterpret_cast(vals); - const float4* bias_vec = reinterpret_cast(bias); - float4* output_vec = reinterpret_cast(output); - - int iter_index = cnt * d1_stride + d2 * d2_stride + d3; - int input_offset = d0 * batch_stride + d1 * (iteration_stride << 1); - bias_arr[0] = bias_vec[iter_index]; - -#pragma unroll - for (int iter = 0; iter < 2; iter++) { - int iter_id = iter * iteration_stride + iter_index; - vals_arr[0] = vals_vec[input_offset + iter_id]; - - output_half[0] = vals_half[0] + bias_half[0]; - output_half[1] = vals_half[1] + bias_half[1]; - output_half[2] = vals_half[2] + bias_half[2]; - output_half[3] = vals_half[3] + bias_half[3]; - - in_data[iter_id] = output_arr[0]; - } - __syncthreads(); - - iteration_stride = blockDim.z * (blockDim.y >> 1); - int matrix_stride = (d0_out_stride * gridDim.x); - int head_count = (d2 >> 1) + cnt * (blockDim.y >> 1); - - int out_index = d0 * d0_out_stride + d1 * (d1_out_stride << 1) + d3 + (d2 % 2) * d2_stride; - -#pragma unroll - for (int iter = 0; iter < 2; iter++) { - int iter_row = (iter * iteration_stride) + head_count; - int iter_offset = - (iter_row % blockDim.y) * d2_out_stride + (iter_row / blockDim.y) * matrix_stride; - output_vec[out_index + iter_offset] = - in_data[iter_row * d2_stride + d3 + (d2 % 2) * (d1_stride * blockDim.z)]; - } -} - -template -__global__ void transform4d_0213(T* out, - const T* in, - int heads, - int seq_length, - int hidden_dim, - int head_ext); - -template <> -__global__ void transform4d_0213(float* out, - const float* in, - int heads, - int seq_length, - int hidden_dim, - int head_ext) -{ - int d0_stride = hidden_dim * seq_length; - int d1_stride = d0_stride / heads; - int d2_stride = hidden_dim / heads; - - int d0_out_stride = d0_stride; - int d1_out_stride = d2_stride; - int d2_out_stride = hidden_dim; - - int d0 = blockIdx.x; // Batch - int d1 = blockIdx.y / ((seq_length - 1) / blockDim.y + 1); // Head - int d2 = (threadIdx.y + blockDim.y * blockIdx.y) % seq_length; - int cnt = blockIdx.z; - int d3 = threadIdx.x; // Values (groups of 8) - - if (d2 < seq_length) { - const float4* in_vec = reinterpret_cast(in); - float4* out_vec = reinterpret_cast(out); - - float4 vals_vec = in_vec[cnt * d0_stride * gridDim.x + d0 * d0_stride + d1 * d1_stride + - d2 * d2_stride + d3]; - out_vec[d0 * d0_out_stride * gridDim.z + cnt * d2_out_stride + d1 * d1_out_stride + - d2 * d2_out_stride * gridDim.z + d3] = vals_vec; - } -} - -template -__global__ void transform4d_0213(T* out, - const T* in, - int heads, - int seq_length, - int hidden_dim, - int head_ext) -{ - int d0_stride = hidden_dim * (seq_length / head_ext); - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - - int d0 = blockIdx.x; // Batch - int d1 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head - int d2 = blockIdx.z / head_ext; // Sequence - int cnt = blockIdx.y; // Hidden count - int d3 = threadIdx.x; // Values (groups of 8) - - const float4* in_vec = reinterpret_cast(in); - float4* out_vec = reinterpret_cast(out); - - in_vec += (cnt * d0_stride * gridDim.x); - in_vec += (d0 * d0_stride); - in_vec += (d2 * d2_stride); - in_vec += (d1 * d2_stride * seq_length); - - out_vec += (cnt * d1_stride); - out_vec += (d1 * d2_stride); - out_vec += (d0 * d0_stride * gridDim.y); - out_vec += (d2 * d1_stride * gridDim.y); - - out_vec[d3] = in_vec[d3]; -} - -template -__global__ void transform4d_0213_v2(T* out, const T* in, int heads, int seq_length, int hidden_dim) -{ - __shared__ float4 in_data[3072]; - - int d0_stride = hidden_dim * seq_length; - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - - int d0 = blockIdx.x; // Batch - int d1 = threadIdx.y; // Head - int d2 = blockIdx.y; // Sequence - int cnt = threadIdx.z; // Hidden count - int d3 = threadIdx.x; // Values (groups of 8) - - const float4* in_vec = reinterpret_cast(in); - float4* out_vec = reinterpret_cast(out); - - int input_offset = d0 * d0_stride + d2 * (d2_stride << 1) + d3 + (d1 % 2) * d2_stride; - int head_count = (d1 >> 1) + cnt * (blockDim.y >> 1); - int iteration_stride = blockDim.z * (blockDim.y >> 1); - int matrix_stride = (d0_stride * gridDim.x); - -#pragma unroll - for (int iter = 0; iter < 2; iter++) { - int iter_row = iter * iteration_stride + head_count; - int iter_offset = (iter_row % blockDim.y) * d2_stride; - - in_data[d3 + iter_offset + (iter_row / blockDim.y + (d1 % 2) * blockDim.z) * d1_stride] = - in_vec[input_offset + iter_offset * seq_length + - (iter_row / blockDim.y) * matrix_stride]; - } - __syncthreads(); - - iteration_stride = d1_stride * blockDim.z; - int iter_index = cnt * d1_stride + d1 * d2_stride + d3; - int output_offset = d0 * d0_stride * blockDim.z + d2 * (iteration_stride << 1); - -#pragma unroll - for (int iter = 0; iter < 2; iter++) { - int iter_id = iter * iteration_stride + iter_index; - out_vec[output_offset + iter_id] = in_data[iter_id]; - } -} - -// 3 * [B A S N] - > [B S C*H] -template <> -void launch_transform4d_0213(float* out, - const float* in, - int batch_size, - int heads, - int seq_length, - int hidden_dim, - cudaStream_t stream, - int trans_count) -{ - hidden_dim >>= 2; - dim3 grid_dims(batch_size, heads * ((seq_length - 1) / 8 + 1), trans_count); - dim3 block_dims(hidden_dim / heads, 8); - transform4d_0213 - <<>>(out, in, heads, seq_length, hidden_dim, 1); -} - -template -void launch_transform4d_0213(T* out, - const T* in, - int batch_size, - int heads, - int seq_length, - int hidden_dim, - cudaStream_t stream, - int trans_count) -{ - hidden_dim >>= 3; - int head_ext = (hidden_dim - 1) / MAX_THREADS + 1; - dim3 grid_dims(batch_size, trans_count, (seq_length * head_ext)); - dim3 block_dims(hidden_dim / heads, (heads / head_ext)); - transform4d_0213<<>>( - out, in, heads, seq_length, hidden_dim, head_ext); -} - -#define INSTANTIATE_2B_LAUNCH_TRANSFORM4D(T) \ - template void launch_transform4d_0213(T*, const T*, int, int, int, int, cudaStream_t, int); - -INSTANTIATE_2B_LAUNCH_TRANSFORM4D(__half) -#ifdef BF16_AVAILABLE -INSTANTIATE_2B_LAUNCH_TRANSFORM4D(__nv_bfloat16) -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_context.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_context.h deleted file mode 100644 index 378fd4e5e990d5be9b642e2092ab30370506b191..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_context.h +++ /dev/null @@ -1,292 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include -#include -#include -#include -#include "cublas_v2.h" -#include "cuda.h" - -#define MEGABYTE (1024 * 1024) -#define GIGABYTE (1024 * 1024 * 1024) - -// TODO: refactor out -#define WARP_SIZE 32 - -#define CUDA_CHECK(callstr) \ - { \ - cudaError_t error_code = callstr; \ - if (error_code != cudaSuccess) { \ - std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \ - assert(0); \ - } \ - } - -#define CUDA_1D_KERNEL_LOOP(i, n) \ - for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); i += blockDim.x * gridDim.x) - -#define CUDA_2D_KERNEL_LOOP(i, n, j, m) \ - for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); i += blockDim.x * gridDim.x) \ - for (size_t j = blockIdx.y * blockDim.y + threadIdx.y; j < (m); j += blockDim.y * gridDim.y) - -#define DS_CUDA_NUM_THREADS 512 -#define DS_MAXIMUM_NUM_BLOCKS 262144 - -inline int DS_GET_BLOCKS(const int N) -{ - return std::max( - std::min((N + DS_CUDA_NUM_THREADS - 1) / DS_CUDA_NUM_THREADS, DS_MAXIMUM_NUM_BLOCKS), - // Use at least 1 block, since CUDA does not allow empty block - 1); -} - -class InferenceContext { -public: - InferenceContext() - : _workspace(nullptr), - _seed(42), - _curr_offset(0), - _stream(0), - _free_memory_size(0), - _num_tokens(1), - _attention_unfused_workspace_offset(0), - _workSpaceSize(0) - { - _workSpaceSize = 0; - _workspace = 0; - - cublasStatus_t stat = cublasCreate(&_cublasHandle); - if (stat != CUBLAS_STATUS_SUCCESS) { - // It would be nice to use cublasGetStatusName and - // cublasGetStatusString, but they were only added in CUDA 11.4.2. - auto message = std::string("Failed to create cublas handle: cublasStatus_t was ") + - std::to_string(stat); - std::cerr << message << std::endl; - throw std::runtime_error(message); - } -#ifndef __HIP_PLATFORM_AMD__ - cublasSetMathMode(_cublasHandle, CUBLAS_TENSOR_OP_MATH); -#endif - cudaEventCreate(&_comp1_event); - cudaEventCreate(&_comp2_event); - cudaEventCreate(&_comp_event); - cudaEventCreate(&_comm_event); - } - - virtual ~InferenceContext() - { - cublasDestroy(_cublasHandle); - cudaFree(_workspace); - cudaEventDestroy(_comp1_event); - cudaEventDestroy(_comp2_event); - cudaEventDestroy(_comp_event); - cudaEventDestroy(_comm_event); - } - - static InferenceContext& Instance() - { - static InferenceContext _ctx; - return _ctx; - } - - void GenWorkSpace(const unsigned& num_layers, - const unsigned& num_heads, - const size_t& batch_size, - const size_t& prompt_len, - const size_t& hidden_dim, - const unsigned& mp_size, - const bool& external_cache, - const size_t& elem_size, - const unsigned& rank, - unsigned max_out_tokens, - unsigned min_out_tokens) - { - size_t total_size; - if (!_free_memory_size) { cudaMemGetInfo(&_free_memory_size, &total_size); } - - // Flash attention requires padded heads and we'll conservatively allocate - // for that here. Flash attention is only enabled for head size <= 128 right now - const int head_size = hidden_dim / num_heads; - const int padded_head_size = head_size <= 32 ? 32 : (head_size <= 64 ? 64 : 128); - const int effective_head_size = (head_size > 128) ? head_size : padded_head_size; - - size_t activation_size = 10 * (num_heads * effective_head_size) * batch_size; - // Other sequence length dimension is added when the final workSpaceSize is calculated - size_t temp_size = batch_size * (num_heads / mp_size) * max_out_tokens; - size_t cache_size = - num_layers * batch_size * ((num_heads * effective_head_size) / mp_size) * 2; - size_t minimal_requirements = - temp_size + (_free_memory_size > GIGABYTE ? 500 : 100) * MEGABYTE; - if (_free_memory_size < minimal_requirements) { - printf("Requested:\t%lu\nFree:\t%lu\nTotal:\t%lu\n", - minimal_requirements, - _free_memory_size, - total_size); - throw std::runtime_error("Workspace can't be allocated, no enough memory."); - } - - _max_seq_len = ((_free_memory_size - minimal_requirements) / elem_size) / - (activation_size + temp_size + cache_size); - _max_seq_len = std::min((size_t)max_out_tokens, _max_seq_len); - size_t workSpaceSize = ((external_cache ? (activation_size + temp_size) - : (activation_size + temp_size + cache_size))) * - _max_seq_len * elem_size; - temp_size *= _max_seq_len * elem_size; - - if (_max_seq_len < min_out_tokens) { - printf( - "Allocatable workspace available (%ld tokens) is less than minimum requested " - "workspace (%d tokens)\n", - _max_seq_len, - min_out_tokens); - throw std::runtime_error("Workspace can't be allocated, not enough memory"); - } - - if (!_workspace) { - assert(_workspace == nullptr); - cudaMalloc(&_workspace, workSpaceSize); - } else if (_workSpaceSize < workSpaceSize) { - cudaFree(_workspace); - cudaMalloc(&_workspace, workSpaceSize); - } - if (rank == 0 && (!_workspace || _workSpaceSize < workSpaceSize)) - printf( - "------------------------------------------------------\n" - "Free memory : %f (GigaBytes) \n" - "Total memory: %f (GigaBytes) \n" - "Requested memory: %f (GigaBytes) \n" - "Setting maximum total tokens (input + output) to %lu \n" - "WorkSpace: %p \n" - "------------------------------------------------------\n", - (float)_free_memory_size / GIGABYTE, - (float)total_size / GIGABYTE, - (float)workSpaceSize / GIGABYTE, - _max_seq_len, - _workspace); - - if (!_workspace) { - printf("Requested:\t%lu\nFree:\t%lu\nTotal:\t%lu\n", - workSpaceSize, - _free_memory_size, - total_size); - throw std::runtime_error("Workspace is null."); - } - _workSpaceSize = workSpaceSize; - _attention_unfused_workspace_offset = workSpaceSize - temp_size; - } - inline size_t GetMaxTokenLength() const { return _max_seq_len; } - - cudaEvent_t GetCompEvent(int id) { return id == 1 ? _comp1_event : _comp2_event; } - - size_t get_workspace_size() const { return _workSpaceSize; } - void* GetWorkSpace() { return _workspace; } - void* GetAttentionUnfusedWorkspace() - { - return (char*)_workspace + _attention_unfused_workspace_offset; - } - - inline unsigned new_token(unsigned layer_id) - { - if (layer_id == 0) _token_length++; - return _token_length; - } - - inline void reset_tokens(unsigned initial_tokens = 1) - { - _num_tokens = initial_tokens; - } //_token_length = 0; } - - inline unsigned current_tokens() const { return _num_tokens; } - - inline void advance_tokens() { _num_tokens++; } - - cudaStream_t GetCommStream(bool async_op = false) - { - if (!_comm_stream) - _comm_stream = async_op ? at::cuda::getStreamFromPool(true) - : at::cuda::getCurrentCUDAStream(); - return _comm_stream; - } - cudaStream_t GetCurrentStream(bool other_stream = false) - { - // get current pytorch stream. - if (other_stream) { - if (!_stream) _stream = at::cuda::getStreamFromPool(true); - return _stream; - } - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - return stream; - } - - void release_workspace() - { - cudaFree(_workspace); - _workspace = nullptr; - } - bool retake_workspace() - { - if (_workspace != nullptr || _workSpaceSize == 0) return true; - cudaMalloc(&_workspace, _workSpaceSize); - return _workspace != nullptr; - } - cublasHandle_t GetCublasHandle() { return _cublasHandle; } - - std::pair IncrementOffset(uint64_t offset_inc) - { - uint64_t offset = _curr_offset; - _curr_offset += offset_inc; - return std::pair(_seed, offset); - } - - void SetSeed(uint64_t new_seed) { _seed = new_seed; } - - const std::vector>& GetGemmAlgos() const { return _gemm_algos; } - - inline void SynchComp() - { - cudaEventRecord(_comp_event, _comp_stream); - cudaStreamWaitEvent(_comm_stream, _comp_event, 0); - } - inline void SynchComm() - { - cudaEventRecord(_comm_event, _comm_stream); - cudaStreamWaitEvent(_comp_stream, _comm_event, 0); - } - -private: - cublasHandle_t _cublasHandle; - - cudaEvent_t _comp_event; - cudaEvent_t _comm_event; - - void* _workspace; - // offset from _workspace for attention unfused memory - size_t _attention_unfused_workspace_offset; - uint64_t _seed; - uint64_t _curr_offset; - - size_t _workSpaceSize; - size_t _free_memory_size; - - size_t _max_seq_len; - - cudaEvent_t _comp1_event; - cudaEvent_t _comp2_event; - - cudaStream_t _stream; - - unsigned _token_length; - unsigned _num_tokens; - std::vector> _gemm_algos; - - cudaStream_t _comp_stream; - cudaStream_t _comm_stream; - - std::unordered_map _world_sizes; -}; diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_cublas_wrappers.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_cublas_wrappers.h deleted file mode 100644 index 40c3e443941d24174ee27c2728cd95907e654074..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_cublas_wrappers.h +++ /dev/null @@ -1,504 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include -#include -#include -#ifdef BF16_AVAILABLE -#include -#endif -#include -#include -#ifndef __HIP_PLATFORM_AMD__ -#include -#endif -#include - -// TODO HIP: Remove backward compatibility for torch<=2.0 in future -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) -int cublas_gemm_ex(rocblas_handle handle, - rocblas_operation transa, - rocblas_operation transb, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const float* A, - const float* B, - float* C, - rocblas_gemm_algo algo, - int b_stride = -1) -#else -int cublas_gemm_ex(cublasHandle_t handle, - cublasOperation_t transa, - cublasOperation_t transb, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const float* A, - const float* B, - float* C, - cublasGemmAlgo_t algo, - int b_stride = -1) -#endif -{ - const int ldb = (b_stride == -1) ? ((transb == CUBLAS_OP_N) ? k : n) : b_stride; -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_status status = rocblas_gemm_ex(handle, - transa, - transb, - m, - n, - k, - (const void*)alpha, - (const void*)A, - rocblas_datatype_f32_r, - (transa == rocblas_operation_none) ? m : k, - (const void*)B, - rocblas_datatype_f32_r, - ldb, - (const void*)beta, - C, - rocblas_datatype_f32_r, - m, - C, - rocblas_datatype_f32_r, - m, - rocblas_datatype_f32_r, - algo, - 0, - 0); -#else - cublasStatus_t status = cublasGemmEx(handle, - transa, - transb, - m, - n, - k, - (const void*)alpha, - (const void*)A, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - (transa == CUBLAS_OP_N) ? m : k, - (const void*)B, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - ldb, - (const void*)beta, - C, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - m, -#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) - HIPBLAS_COMPUTE_32F, -#elif defined(__HIP_PLATFORM_AMD__) - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - algo); -#endif - -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - if (status != rocblas_status_success) { -#else - if (status != CUBLAS_STATUS_SUCCESS) { -#endif - fprintf(stderr, - "!!!! kernel execution error. (m: %d, n: %d, k: %d, error: %d) \n", - m, - n, - k, - (int)status); - return EXIT_FAILURE; - } - return 0; -} - -template -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) -int cublas_gemm_ex(rocblas_handle handle, - rocblas_operation transa, - rocblas_operation transb, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const T* A, - const T* B, - T* C, - rocblas_gemm_algo algo, - int b_stride = -1) -#else -int cublas_gemm_ex(cublasHandle_t handle, - cublasOperation_t transa, - cublasOperation_t transb, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const T* A, - const T* B, - T* C, - cublasGemmAlgo_t algo, - int b_stride = -1) -#endif -{ - const int ldb = (b_stride == -1) ? ((transb == CUBLAS_OP_N) ? k : n) : b_stride; -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - constexpr auto rocblas_dtype_16 = std::is_same::value ? rocblas_datatype_f16_r - : rocblas_datatype_bf16_r; - rocblas_status status = rocblas_gemm_ex(handle, - transa, - transb, - m, - n, - k, - (const void*)alpha, - (const void*)A, - rocblas_dtype_16, - (transa == rocblas_operation_none) ? m : k, - (const void*)B, - rocblas_dtype_16, - ldb, - (const void*)beta, - (void*)C, - rocblas_dtype_16, - m, - (void*)C, - rocblas_dtype_16, - m, - rocblas_datatype_f32_r, - algo, - 0, - 0); -#else -#ifdef __HIP_PLATFORM_AMD__ - constexpr auto cublas_dtype_16 = std::is_same::value ? HIPBLAS_R_16F : HIPBLAS_R_16B; -#else - constexpr auto cublas_dtype_16 = std::is_same::value ? CUDA_R_16F : CUDA_R_16BF; -#endif - cublasStatus_t status = cublasGemmEx(handle, - transa, - transb, - m, - n, - k, - (const void*)alpha, - (const void*)A, - cublas_dtype_16, - (transa == CUBLAS_OP_N) ? m : k, - (const void*)B, - cublas_dtype_16, - ldb, - (const void*)beta, - (void*)C, - cublas_dtype_16, - m, -#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) - HIPBLAS_COMPUTE_32F, -#elif defined(__HIP_PLATFORM_AMD__) - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - algo); -#endif - -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - if (status != rocblas_status_success) { -#else - if (status != CUBLAS_STATUS_SUCCESS) { -#endif - fprintf(stderr, - "!!!! kernel execution error. (m: %d, n: %d, k: %d, error: %d) \n", - m, - n, - k, - (int)status); - return EXIT_FAILURE; - } - return 0; -} - -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) -int cublas_strided_batched_gemm(rocblas_handle handle, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const float* A, - const float* B, - float* C, - rocblas_operation op_A, - rocblas_operation op_B, - int stride_A, - int stride_B, - int stride_C, - int batch, - rocblas_gemm_algo algo) -#else -int cublas_strided_batched_gemm(cublasHandle_t handle, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const float* A, - const float* B, - float* C, - cublasOperation_t op_A, - cublasOperation_t op_B, - int stride_A, - int stride_B, - int stride_C, - int batch, - cublasGemmAlgo_t algo) -#endif -{ -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - rocblas_status status = - rocblas_gemm_strided_batched_ex(handle, - op_A, - op_B, - m, - n, - k, - alpha, - A, - rocblas_datatype_f32_r, - (op_A == rocblas_operation_none) ? m : k, - stride_A, - B, - rocblas_datatype_f32_r, - (op_B == rocblas_operation_none) ? k : n, - stride_B, - beta, - C, - rocblas_datatype_f32_r, - m, - stride_C, - C, - rocblas_datatype_f32_r, - m, - stride_C, - batch, - rocblas_datatype_f32_r, - algo, - 0, - 0); -#else - cublasStatus_t status = cublasGemmStridedBatchedEx(handle, - op_A, - op_B, - m, - n, - k, - alpha, - A, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - (op_A == CUBLAS_OP_N) ? m : k, - stride_A, - B, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - (op_B == CUBLAS_OP_N) ? k : n, - stride_B, - beta, - C, -#ifdef __HIP_PLATFORM_AMD__ - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - m, - stride_C, - batch, -#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) - HIPBLAS_COMPUTE_32F, -#elif defined(__HIP_PLATFORM_AMD__) - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - algo); -#endif - -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - if (status != rocblas_status_success) { -#else - if (status != CUBLAS_STATUS_SUCCESS) { -#endif - fprintf(stderr, - "!!!! kernel execution error. (batch: %d, m: %d, n: %d, k: %d, error: %d) \n", - batch, - m, - n, - k, - (int)status); - return EXIT_FAILURE; - } - return 0; -} - -template -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) -int cublas_strided_batched_gemm(rocblas_handle handle, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const T* A, - const T* B, - T* C, - rocblas_operation op_A, - rocblas_operation op_B, - int stride_A, - int stride_B, - int stride_C, - int batch, - rocblas_gemm_algo algo) -#else -int cublas_strided_batched_gemm(cublasHandle_t handle, - int m, - int n, - int k, - const float* alpha, - const float* beta, - const T* A, - const T* B, - T* C, - cublasOperation_t op_A, - cublasOperation_t op_B, - int stride_A, - int stride_B, - int stride_C, - int batch, - cublasGemmAlgo_t algo) -#endif -{ -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - constexpr auto rocblas_dtype_16 = std::is_same::value ? rocblas_datatype_f16_r - : rocblas_datatype_bf16_r; - rocblas_status status = - rocblas_gemm_strided_batched_ex(handle, - op_A, - op_B, - m, - n, - k, - alpha, - A, - rocblas_dtype_16, - (op_A == rocblas_operation_none) ? m : k, - stride_A, - B, - rocblas_dtype_16, - (op_B == rocblas_operation_none) ? k : n, - stride_B, - beta, - C, - rocblas_dtype_16, - m, - stride_C, - C, - rocblas_dtype_16, - m, - stride_C, - batch, - rocblas_datatype_f32_r, - algo, - 0, - 0); -#else -#ifdef __HIP_PLATFORM_AMD__ - constexpr auto cublas_dtype_16 = std::is_same::value ? HIPBLAS_R_16F : HIPBLAS_R_16B; -#else - constexpr auto cublas_dtype_16 = std::is_same::value ? CUDA_R_16F : CUDA_R_16BF; -#endif - cublasStatus_t status = cublasGemmStridedBatchedEx(handle, - op_A, - op_B, - m, - n, - k, - alpha, - A, - cublas_dtype_16, - (op_A == CUBLAS_OP_N) ? m : k, - stride_A, - B, - cublas_dtype_16, - (op_B == CUBLAS_OP_N) ? k : n, - stride_B, - beta, - C, - cublas_dtype_16, - m, - stride_C, - batch, -#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2) - HIPBLAS_COMPUTE_32F, -#elif defined(__HIP_PLATFORM_AMD__) - HIPBLAS_R_32F, -#else - CUDA_R_32F, -#endif - algo); -#endif - -#if defined(__HIP_PLATFORM_AMD__) && \ - ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0)) - if (status != rocblas_status_success) { -#else - if (status != CUBLAS_STATUS_SUCCESS) { -#endif - fprintf(stderr, - "!!!! kernel execution error. (m: %d, n: %d, k: %d, error: %d) \n", - m, - n, - k, - (int)status); - return EXIT_FAILURE; - } - - return 0; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_cuda_layers.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_cuda_layers.h deleted file mode 100644 index dcc020483687f751f5997863307504449743fa35..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/inference/includes/inference_cuda_layers.h +++ /dev/null @@ -1,248 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#include "ds_kernel_utils.h" - -#include -#ifdef BF16_AVAILABLE -#include -#endif -#include -#include -#include -#include -#include - -#define MAX_WARP_NUM 32 -#define WARP_SIZE 32 - -#define MAX_THREADS 1024 -#define SMs 80 - -#define MAX_REGISTERS 256 - -template -void launch_attn_softmax_v2(T* vals, - T* mask, - T* alibi, - float layer_scale, - bool triangular, - bool recompute, - bool local_attention, - int window_size, - int batch_size, - int heads, - int num_seq, - int sequence_length, - int offset, - int mask_stride, - int mp_size, - cudaStream_t stream); - -// Fused bias add with gelu activation -template -void launch_bias_gelu(T* input, - const T* bias, - int intermediate_size, - int batch_size, - cudaStream_t stream); - -template -void launch_gated_activation(T* output, - const T* activation, - const T* bias, - int rows, - int output_stride, - int elems_per_row, - bool use_gelu, - cudaStream_t stream); - -// Fused bias add with relu activation -template -void launch_bias_relu(T* input, - const T* bias, - int intermediate_size, - int batch_size, - cudaStream_t stream); - -template -void launch_bias_add(T* input, const T* bias, int hidden_size, int batch_size, cudaStream_t stream); - -template -void launch_bias_residual(T* input, - T* output, - T* attn, - T* bias, - T* attn_bias, - int batch, - int hidden_dim, - int mp_size, - bool preln, - cudaStream_t stream); - -template -void launch_fused_ln(T* output, - const T* vals, - const T* gamma, - const T* beta, - float epsilon, - int rows, - int elems_per_row, - cudaStream_t stream); - -template -void launch_fused_residual_ln(T* output, - const T* vals, - const T* residual, - const T* bias, - const T* gamma, - const T* beta, - float epsilon, - int rows, - int elems_per_row, - cudaStream_t stream); - -template -void launch_fused_residual_ln_store_pre_ln_res(T* norm_output, - T* res_output, - const T* vals, - const T* residual, - const T* bias, - const T* gamma, - const T* beta, - float epsilon, - int rows, - int elems_per_row, - cudaStream_t stream); - -template -void launch_rms_norm(T* norm_output, - T* res_output, - const T* vals, - const T* residual, - const T* gamma, - float epsilon, - int rows, - int elems_per_row, - cudaStream_t stream); - -template -void launch_dequantize(T* output, - const int8_t* input, - const float* qscale, - unsigned output_size, - unsigned hidden_dim, - unsigned groups, - unsigned merge_count, - cudaStream_t stream); - -template -void launch_dequantize(T* output, - const int8_t* input, - const float* qscale, - unsigned output_size, - unsigned hidden_dim, - unsigned groups, - cudaStream_t stream); -template -void launch_gptj_residual_add(T* input, - T* output, - T* attn, - T* bias, - T* attn_bias, - int batch, - int head_size, - int mp_size, - cudaStream_t stream); - -template -void launch_apply_rotary_pos_emb(T* mixed_query, - T* key_layer, - unsigned head_size, - unsigned seq_len, - unsigned rotary_dim, - unsigned offset, - unsigned num_heads, - unsigned batch, - float rope_theta, - cudaStream_t stream, - int max_out_tokens); - -template -void launch_moe_res_matmul(T* residual, - T* coef, - T* mlp_out, - int seq_len, - int hidden_dim, - cudaStream_t stream); - -// 4D transform [0, 1, 2, 3] -> [0, 2, 1, 3] -template -void launch_transform4d_0213(T* out, - const T* in, - int batch_size, - int heads, - int seq_length, - int hidden_dim, - cudaStream_t stream, - int trans_count); -template -void launch_bias_add_transform_0213(T* outputs, - T* vals, - T* vals1, - const T* vals2, - const T* bias, - int batch_size, - int seq_length, - unsigned seq_offset, - int seq_length1, - int hidden_dim, - int heads, - int num_kv, - int rotary_dim, - bool rotate_half, - bool rotate_every_two, - cudaStream_t stream, - int trans_count, - int max_out_tokens, - float rope_theta); -template -void pad_data(T* padded_output, - T* output, - int bsz, - int head_size, - int padded_head_size, - cudaStream_t stream); - -template -void pad_head_seq(T* padded_output, - T* output, - int bsz, - int seq_len, - int padded_seq_len, - int head_size, - int padded_head_size, - cudaStream_t stream); - -template -void launch_pad_add_transform_0213(T* output, - const T* vals, - int batch_size, - int hidden_dim, - int seq_length, - int padded_seq_len, - int heads, - int padded_head_size, - cudaStream_t stream); - -template -void launch_vector_add(T* out, - const T* a, - const T* b, - float gamma, - int num_elems, - cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/normalize_kernels.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/normalize_kernels.cu deleted file mode 100644 index 8b9dbaea6c86d98ed7e431073684b89b8d1499ee..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/normalize_kernels.cu +++ /dev/null @@ -1,2134 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "custom_cuda_layers.h" - -namespace cg = cooperative_groups; - -/* -Fused bias add, residual (elementwise) add, and normalization layer. - -For FP16, this kernel does not promote to FP32 in order to utilize the 2x throughput for -__half2 instructions, and avoid the conversion overhead (1/8 of __hal2 arithmetic). - -For specific launch constraints, see the launch functions. -*/ - -#define NORM_REG (MAX_REGISTERS / 4) - -__global__ void fused_bias_residual_layer_norm(float* vals, - const float* residual, - const float* gamma, - const float* beta, - float epsilon, - bool preLayerNorm, - bool training, - float* vars, - float* means, - int row_stride) -{ - int iteration_stride = blockDim.x; - int iterations = row_stride / iteration_stride; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int row = blockIdx.x; - int id = threadIdx.x; - int gid = id / WARP_SIZE; - - float vals_arr[NORM_REG]; - __shared__ float shr[MAX_WARP_NUM]; - - residual += (row * row_stride); - vals += (row * row_stride); - - float sum = 0.f; - int high_index = iterations * iteration_stride + id; -#pragma unroll - for (int i = 0; i < iterations; i++) { - vals_arr[i] = residual[i * iteration_stride + id]; - sum += vals_arr[i]; - } - if (high_index < row_stride) { - vals_arr[iterations] = residual[high_index]; - sum += vals_arr[iterations]; - iterations++; - } - - for (int i = 1; i < 32; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) shr[gid] = sum; - - b.sync(); - - if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) sum = shr[g.thread_rank()]; - -#if !defined(__STOCHASTIC_MODE__) || __CUDA_ARCH__ < 700 - b.sync(); -#endif - - for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { - sum += g.shfl_down(sum, i); - } - - sum = g.shfl(sum, 0); - float mean = sum / row_stride; - if (training) - if (threadIdx.x == 0) means[row] = mean; - float variance = 0.f; - for (int i = 0; i < iterations; i++) { - vals_arr[i] -= mean; - variance += vals_arr[i] * vals_arr[i]; - } - - for (int i = 1; i < 32; i *= 2) { variance += g.shfl_down(variance, i); } - - if (g.thread_rank() == 0) shr[gid] = variance; - - b.sync(); - - if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) variance = shr[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - b.sync(); -#endif - - for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { - variance += g.shfl_down(variance, i); - } - variance = g.shfl(variance, 0); - variance /= row_stride; - variance += epsilon; - if (training) - if (threadIdx.x == 0) vars[row] = variance; - - iterations = row_stride / iteration_stride; - for (int i = 0; i < iterations; i++) { - vals_arr[i] = vals_arr[i] * rsqrtf(variance); - vals_arr[i] = - vals_arr[i] * gamma[i * iteration_stride + id] + beta[i * iteration_stride + id]; - vals[i * iteration_stride + id] = vals_arr[i]; - } - if ((high_index) < row_stride) { - vals_arr[iterations] = vals_arr[iterations] * rsqrtf(variance); - vals_arr[iterations] = vals_arr[iterations] * gamma[high_index] + beta[high_index]; - vals[high_index] = vals_arr[iterations]; - } -} - -__global__ void fused_bias_residual_layer_norm(__half* vals, - const __half* residual, - const __half* gamma, - const __half* beta, - float epsilon, - bool preLayerNorm, - bool training, - __half* vars, - __half* means, - int row_stride) -{ -#ifdef HALF_PRECISION_AVAILABLE - int iteration_stride = blockDim.x; - int iterations = row_stride / iteration_stride; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); - - int row = blockIdx.x; - int id = threadIdx.x; - int gid = id >> WARP_SIZE_BITS; - - float2 vals_f[NORM_REG]; - __shared__ float shr[MAX_WARP_NUM]; - - __half2* vals_cast = reinterpret_cast<__half2*>(vals); - const __half2* residual_cast = reinterpret_cast(residual); - - residual_cast += (row * row_stride); - vals_cast += (row * row_stride); - - float sum = 0.f; - int high_index = iterations * iteration_stride + id; -#pragma unroll - for (int i = 0; i < iterations; i++) { - vals_f[i] = __half22float2(residual_cast[i * iteration_stride + id]); - sum += vals_f[i].x; - sum += vals_f[i].y; - } - if ((high_index) < row_stride) { - vals_f[iterations] = __half22float2(residual_cast[high_index]); - sum += vals_f[iterations].x; - sum += vals_f[iterations].y; - iterations++; - } - - for (int i = 1; i < 32; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) shr[gid] = sum; - - b.sync(); - - if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) sum = shr[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - b.sync(); -#endif - - for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { - sum += g.shfl_down(sum, i); - } - sum = g.shfl(sum, 0); - float mean = sum / (row_stride * 2); - - float variance = 0.f; - for (int i = 0; i < iterations; i++) { - vals_f[i].x -= mean; - vals_f[i].y -= mean; - variance += vals_f[i].x * vals_f[i].x; - variance += vals_f[i].y * vals_f[i].y; - } - - for (int i = 1; i < 32; i *= 2) { variance += g.shfl_down(variance, i); } - - if (g.thread_rank() == 0) shr[gid] = variance; - - b.sync(); - - if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) variance = shr[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - b.sync(); -#endif - - for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { - variance += g.shfl_down(variance, i); - } - variance = g.shfl(variance, 0); - variance /= (row_stride * 2); - variance += epsilon; - - __half2 variance_h = __float2half2_rn(variance); - const __half2* gamma_cast = reinterpret_cast(gamma); - const __half2* beta_cast = reinterpret_cast(beta); - - if (training && threadIdx.x == 0) { - vars[row] = __float2half(variance); - means[row] = __float2half(mean); - } - iterations = row_stride / iteration_stride; - for (int i = 0; i < iterations; i++) { - __half2 vals_arr = __float22half2_rn(vals_f[i]); - vals_arr = vals_arr * h2rsqrt(variance_h); - vals_arr = - vals_arr * gamma_cast[i * iteration_stride + id] + beta_cast[i * iteration_stride + id]; - vals_cast[i * iteration_stride + id] = vals_arr; - } - if ((high_index) < row_stride) { - __half2 vals_arr = __float22half2_rn(vals_f[iterations]); - vals_arr = vals_arr * h2rsqrt(variance_h); - vals_arr = vals_arr * gamma_cast[high_index] + beta_cast[high_index]; - vals_cast[high_index] = vals_arr; - } -#endif -} - -template -void launch_bias_residual_layer_norm(T* vals, - const T* residual, - const T* gamma, - const T* beta, - float epsilon, - int batch_size, - int hidden_dim, - cudaStream_t stream, - bool preLayerNorm, - bool training, - T* vars, - T* means); - -template <> -void launch_bias_residual_layer_norm(float* vals, - const float* residual, - const float* gamma, - const float* beta, - float epsilon, - int batch_size, - int hidden_dim, - cudaStream_t stream, - bool preLayerNorm, - bool training, - float* vars, - float* means) -{ - int threads = THREADS; - - dim3 grid_dim(batch_size); - - if (hidden_dim > 16384 && hidden_dim <= 32768) - threads <<= 1; - else if (hidden_dim > 32768 && hidden_dim <= 65536) - threads <<= 2; - else if (hidden_dim > 65536) - throw std::runtime_error("Unsupport hidden_dim."); - - dim3 block_dim(threads); - - fused_bias_residual_layer_norm<<>>( - vals, residual, gamma, beta, epsilon, preLayerNorm, training, vars, means, hidden_dim); -} - -template <> -void launch_bias_residual_layer_norm<__half>(__half* vals, - const __half* residual, - const __half* gamma, - const __half* beta, - float epsilon, - int batch_size, - int hidden_dim, - cudaStream_t stream, - bool preLayerNorm, - bool training, - __half* vars, - __half* means) -{ - int threads = 128; - - dim3 grid_dim(batch_size); - - if (hidden_dim > 8192 && hidden_dim <= 16384) - threads <<= 1; - else if (hidden_dim > 16384 && hidden_dim <= 32768) - threads <<= 2; - else if (hidden_dim > 32768 && hidden_dim <= 65536) - threads <<= 3; - else if (hidden_dim > 65536) - throw std::runtime_error("Unsupport hidden_dim."); - - dim3 block_dim(threads); - - fused_bias_residual_layer_norm<<>>( - vals, residual, gamma, beta, epsilon, preLayerNorm, training, vars, means, hidden_dim / 2); -} - -__global__ void fused_bias_residual_layer_norm(float* vals, - const float* residual, - const float* gamma, - const float* beta, - float epsilon, - bool preLayerNorm, - bool training, - float* vars, - int row_stride) -{ - int iteration_stride = blockDim.x; - int iterations = row_stride / iteration_stride; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); - - int row = blockIdx.x; - int id = threadIdx.x; - int gid = id / 32; - - float vals_arr[NORM_REG]; - __shared__ float shr[MAX_WARP_NUM]; - - residual += (row * row_stride); - vals += (row * row_stride); - - float sum = 0.f; - int high_index = iterations * iteration_stride + id; -#pragma unroll - for (int i = 0; i < iterations; i++) { - vals_arr[i] = residual[i * iteration_stride + id]; - sum += vals_arr[i]; - } - if ((high_index) < row_stride) { - vals_arr[iterations] = residual[high_index]; - sum += vals_arr[iterations]; - iterations++; - } - - for (int i = 1; i < 32; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) shr[gid] = sum; - - b.sync(); - - if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) sum = shr[g.thread_rank()]; - -#if !defined(__STOCHASTIC_MODE__) || __CUDA_ARCH__ < 700 - b.sync(); -#endif - - for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { - sum += g.shfl_down(sum, i); - } - - sum = g.shfl(sum, 0); - float mean = sum / row_stride; - float variance = 0.f; - for (int i = 0; i < iterations; i++) { - vals_arr[i] -= mean; - variance += vals_arr[i] * vals_arr[i]; - } - - for (int i = 1; i < 32; i *= 2) { variance += g.shfl_down(variance, i); } - - if (g.thread_rank() == 0) shr[gid] = variance; - - b.sync(); - - if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) variance = shr[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - b.sync(); -#endif - - for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { - variance += g.shfl_down(variance, i); - } - variance = g.shfl(variance, 0); - variance /= row_stride; - variance += epsilon; - if (training) - if (threadIdx.x == 0) vars[row] = variance; - - iterations = row_stride / iteration_stride; - for (int i = 0; i < iterations; i++) { - vals_arr[i] = vals_arr[i] * rsqrtf(variance); - vals_arr[i] = - vals_arr[i] * gamma[i * iteration_stride + id] + beta[i * iteration_stride + id]; - vals[i * iteration_stride + id] = vals_arr[i]; - } - if ((high_index) < row_stride) { - vals_arr[iterations] = vals_arr[iterations] * rsqrtf(variance); - vals_arr[iterations] = vals_arr[iterations] * gamma[high_index] + beta[high_index]; - vals[high_index] = vals_arr[iterations]; - } -} - -__global__ void fused_bias_residual_layer_norm(__half* vals, - const __half* residual, - const __half* gamma, - const __half* beta, - float epsilon, - bool preLayerNorm, - bool training, - __half* vars, - int row_stride) -{ -#ifdef HALF_PRECISION_AVAILABLE - - int iteration_stride = blockDim.x; - int iterations = row_stride / iteration_stride; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); - - int row = blockIdx.x; - int id = threadIdx.x; - int gid = id >> WARP_SIZE_BITS; - - float2 vals_f[NORM_REG]; - __shared__ float shr[MAX_WARP_NUM]; - - __half2* vals_cast = reinterpret_cast<__half2*>(vals); - const __half2* residual_cast = reinterpret_cast(residual); - - residual_cast += (row * row_stride); - vals_cast += (row * row_stride); - - float sum = 0.f; - int high_index = iterations * iteration_stride + id; -#pragma unroll - for (int i = 0; i < iterations; i++) { - vals_f[i] = __half22float2(residual_cast[i * iteration_stride + id]); - sum += vals_f[i].x; - sum += vals_f[i].y; - } - if ((high_index) < row_stride) { - vals_f[iterations] = __half22float2(residual_cast[high_index]); - sum += vals_f[iterations].x; - sum += vals_f[iterations].y; - iterations++; - } - - for (int i = 1; i < 32; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) shr[gid] = sum; - - b.sync(); - - if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) sum = shr[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - b.sync(); -#endif - - for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { - sum += g.shfl_down(sum, i); - } - sum = g.shfl(sum, 0); - float mean = sum / (row_stride * 2); - - float variance = 0.f; - for (int i = 0; i < iterations; i++) { - vals_f[i].x -= mean; - vals_f[i].y -= mean; - variance += vals_f[i].x * vals_f[i].x; - variance += vals_f[i].y * vals_f[i].y; - } - - for (int i = 1; i < 32; i *= 2) { variance += g.shfl_down(variance, i); } - - if (g.thread_rank() == 0) shr[gid] = variance; - - b.sync(); - - if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) variance = shr[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - b.sync(); -#endif - - for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) { - variance += g.shfl_down(variance, i); - } - variance = g.shfl(variance, 0); - variance /= (row_stride * 2); - variance += epsilon; - - __half2 variance_h = __float2half2_rn(variance); - const __half2* gamma_cast = reinterpret_cast(gamma); - const __half2* beta_cast = reinterpret_cast(beta); - - if (training && threadIdx.x == 0) vars[row] = __float2half(variance); - - iterations = row_stride / iteration_stride; - for (int i = 0; i < iterations; i++) { - __half2 vals_arr = __float22half2_rn(vals_f[i]); - vals_arr = vals_arr * h2rsqrt(variance_h); - vals_arr = - vals_arr * gamma_cast[i * iteration_stride + id] + beta_cast[i * iteration_stride + id]; - vals_cast[i * iteration_stride + id] = vals_arr; - } - if ((high_index) < row_stride) { - __half2 vals_arr = __float22half2_rn(vals_f[iterations]); - vals_arr = vals_arr * h2rsqrt(variance_h); - vals_arr = vals_arr * gamma_cast[high_index] + beta_cast[high_index]; - vals_cast[high_index] = vals_arr; - } -#endif -} - -template -void launch_bias_residual_layer_norm(T* vals, - const T* residual, - const T* gamma, - const T* beta, - float epsilon, - int batch_size, - int hidden_dim, - cudaStream_t stream, - bool preLayerNorm, - bool training, - T* vars); - -/* -To tune this launch the following restrictions must be met: - -For float: -row_stride == hidden_size -threads * iterations == row_stride -threads is in [32, 64, 128, 256, 512, 1024] - -For half: -row_stride == hidden_size / 2 -threads * iterations == row_stride -threads is in [32, 64, 128, 256, 512, 1024] - -*/ - -template <> -void launch_bias_residual_layer_norm(float* vals, - const float* residual, - const float* gamma, - const float* beta, - float epsilon, - int batch_size, - int hidden_dim, - cudaStream_t stream, - bool preLayerNorm, - bool training, - float* vars) -{ - int threads = THREADS; - - dim3 grid_dim(batch_size); - - // There are some limitations to call below functions, now just enumerate the situations. - - if (hidden_dim > 16384 && hidden_dim <= 32768) - threads <<= 1; - else if (hidden_dim > 32768 && hidden_dim <= 65536) - threads <<= 2; - else if (hidden_dim > 65536) - throw std::runtime_error("Unsupport hidden_dim."); - - dim3 block_dim(threads); - - fused_bias_residual_layer_norm<<>>( - vals, residual, gamma, beta, epsilon, preLayerNorm, training, vars, hidden_dim); -} - -template <> -void launch_bias_residual_layer_norm<__half>(__half* vals, - const __half* residual, - const __half* gamma, - const __half* beta, - float epsilon, - int batch_size, - int hidden_dim, - cudaStream_t stream, - bool preLayerNorm, - bool training, - __half* vars) -{ - int threads = 128; - - dim3 grid_dim(batch_size); - - // There are some limitations to call below functions, now just enumerate the situations. - - if (hidden_dim > 8192 && hidden_dim <= 16384) - threads <<= 1; - else if (hidden_dim > 16384 && hidden_dim <= 32768) - threads <<= 2; - else if (hidden_dim > 32768 && hidden_dim <= 65536) - threads <<= 3; - else if (hidden_dim > 65536) - throw std::runtime_error("Unsupport hidden_dim."); - - dim3 block_dim(threads); - fused_bias_residual_layer_norm<<>>( - vals, residual, gamma, beta, epsilon, preLayerNorm, training, vars, hidden_dim / 2); -} - -/* Normalize Gamma & Betta gradients - * Compute gradients using either X_hat or - * normalize input (invertible). - * Combine transpose with gradients computation. - */ - -template -__global__ void LayerNormBackward1(const T* __restrict__ out_grad, - const T* __restrict__ vals_hat, - const T* __restrict__ gamma, - const T* __restrict__ betta, - T* __restrict__ gamma_grad, - T* __restrict__ betta_grad, - int rows, - int width, - bool invertible) -{ - __shared__ float betta_buffer[TILE_DIM][TILE_DIM + 1]; - __shared__ float gamma_buffer[TILE_DIM][TILE_DIM + 1]; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int idx = blockDim.x * blockIdx.x + threadIdx.x; - int offset = threadIdx.y * width + idx; - int y_stride = width * TILE_DIM; - - float betta_reg = (invertible ? (float)betta[idx] : 0.0f); - float gamma_reg = (float)gamma[idx]; - - // Loop across matrix height - float betta_tmp = 0; - float gamma_tmp = 0; - for (int r = threadIdx.y; r < rows; r += TILE_DIM) { - float grad = (float)out_grad[offset]; - float val = (invertible ? ((float)vals_hat[offset] - betta_reg) / gamma_reg - : (float)vals_hat[offset]); - betta_tmp += grad; - gamma_tmp += (val * grad); - - offset += y_stride; - } - - betta_buffer[threadIdx.x][threadIdx.y] = betta_tmp; - gamma_buffer[threadIdx.x][threadIdx.y] = gamma_tmp; - - __syncthreads(); - - // Sum the shared buffer. - float s1 = betta_buffer[threadIdx.y][threadIdx.x]; - float s2 = gamma_buffer[threadIdx.y][threadIdx.x]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < TILE_DIM; i <<= 1) { - s1 += g.shfl_down(s1, i); - s2 += g.shfl_down(s2, i); - } - - if (threadIdx.x == 0) { - int pos = blockIdx.x * TILE_DIM + threadIdx.y; - betta_grad[pos] = s1; - gamma_grad[pos] = s2; - } -} - -/* Normalize Gamma & Betta gradients - * Compute gradients using the input to - * the normalize. - * Combine transpose with gradients computation. - */ - -template -__global__ void LayerNormBackward1(const T* __restrict__ out_grad, - const T* __restrict__ X_data, - const T* __restrict__ vars, - const T* __restrict__ means, - T* __restrict__ gamma_grad, - T* __restrict__ betta_grad, - int rows, - int width) -{ - __shared__ float betta_buffer[TILE_DIM][TILE_DIM + 1]; - __shared__ float gamma_buffer[TILE_DIM][TILE_DIM + 1]; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int idx = blockDim.x * blockIdx.x + threadIdx.x; - int offset = threadIdx.y * width + idx; - int y_stride = width * TILE_DIM; - - int pos = blockIdx.x * TILE_DIM + threadIdx.y; - // Loop across matrix height - - float betta_tmp = 0; - float gamma_tmp = 0; - for (int r = threadIdx.y; r < rows; r += TILE_DIM) { - float grad = (float)out_grad[offset]; - float val = (float)X_data[offset]; - val = (val - (float)means[r]) * rsqrtf((float)vars[r]); - betta_tmp += grad; - gamma_tmp += (val * grad); - - offset += y_stride; - } - - betta_buffer[threadIdx.x][threadIdx.y] = betta_tmp; - gamma_buffer[threadIdx.x][threadIdx.y] = gamma_tmp; - - __syncthreads(); - - // Sum the shared buffer. - float s1 = betta_buffer[threadIdx.y][threadIdx.x]; - float s2 = gamma_buffer[threadIdx.y][threadIdx.x]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < TILE_DIM; i <<= 1) { - s1 += g.shfl_down(s1, i); - s2 += g.shfl_down(s2, i); - } - - if (threadIdx.x == 0) { - betta_grad[pos] = s1; - gamma_grad[pos] = s2; - } -} -/* - -/* Backward Normalize (Input-Gradient) - * Using the means and variances from the input - * This type of backward is invertible! - * We do the backward using the X_hat (X - u) / sqrt(variance) or the output of Normalization. - */ - -__global__ void LayerNormBackward2(const float* out_grad, - const float* vals_hat, - const float* gamma, - const float* betta, - const float* vars, - float* inp_grad, - bool invertible, - int row_stride) -{ - int iteration_stride = blockDim.x; - int iterations = row_stride / iteration_stride; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int row = blockIdx.x; - int id = threadIdx.x; - int wid = id / WARP_SIZE; - int warp_num = iteration_stride >> WARP_SIZE_BITS; - __shared__ float partialSum[MAX_WARP_NUM]; - - out_grad += (row * row_stride); - vals_hat += (row * row_stride); - inp_grad += (row * row_stride); - - float vals_arr[NORM_REG]; - float vals_hat_arr[NORM_REG]; - int high_index = iterations * iteration_stride + id; -#pragma unroll - for (int i = 0; i < iterations; i++) { - float gamma_reg = gamma[i * iteration_stride + id]; - vals_arr[i] = out_grad[i * iteration_stride + id]; - vals_arr[i] *= gamma_reg; - vals_hat_arr[i] = - (invertible ? (vals_hat[i * iteration_stride + id] - betta[i * iteration_stride + id]) / - gamma_reg - : vals_hat[i * iteration_stride + id]); - } - if ((high_index) < row_stride) { - float gamma_reg = gamma[high_index]; - vals_arr[iterations] = out_grad[high_index]; - vals_arr[iterations] *= gamma_reg; - vals_hat_arr[iterations] = - (invertible ? (vals_hat[high_index] - betta[high_index]) / gamma_reg - : vals_hat[high_index]); - iterations++; - } - - float var_reg = vars[row]; - - float sum = 0; - for (int i = 0; i < iterations; i++) { - sum += vals_hat_arr[i] * vals_arr[i] * - sqrtf(var_reg); // dval_hat = gamma * (x - u) * out_grad - vals_arr[i] *= rsqrtf(var_reg); // dvar_inv = gamma * out_grad / sqrt(var) - } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - - sum = g.shfl(sum, 0); - sum /= row_stride; - - for (int i = 0; i < iterations; i++) { vals_arr[i] += ((-sum * vals_hat_arr[i]) / var_reg); } - - sum = 0; - for (int i = 0; i < iterations; i++) { sum += vals_arr[i]; } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - sum = g.shfl(sum, 0); - sum /= row_stride; - - iterations = row_stride / iteration_stride; - for (int i = 0; i < iterations; i++) inp_grad[i * iteration_stride + id] = (vals_arr[i] - sum); - if ((high_index) < row_stride) inp_grad[high_index] = (vals_arr[iterations] - sum); -} - -__global__ void LayerNormBackward2(const __half* out_grad, - const __half* vals_hat, - const __half* gamma, - const __half* betta, - const __half* vars, - __half* inp_grad, - bool invertible, - int row_stride) -{ -#ifdef HALF_PRECISION_AVAILABLE - int iteration_stride = blockDim.x; - int iterations = row_stride / iteration_stride; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int row = blockIdx.x; - int id = threadIdx.x; - int wid = id / WARP_SIZE; - int warp_num = iteration_stride >> WARP_SIZE_BITS; - __shared__ float partialSum[MAX_WARP_NUM]; - - __half2 vals_arr[NORM_REG]; - float2 vals_arr_f[NORM_REG]; - __half2 vals_hat_arr[NORM_REG]; - - __half2* inp_grad_h = reinterpret_cast<__half2*>(inp_grad); - const __half2* out_grad_h = reinterpret_cast(out_grad); - const __half2* vals_hat_h = reinterpret_cast(vals_hat); - - inp_grad_h += (row * row_stride); - out_grad_h += (row * row_stride); - vals_hat_h += (row * row_stride); - - const __half2* gamma_h = reinterpret_cast(gamma); - const __half2* betta_h = (invertible ? reinterpret_cast(betta) : nullptr); - int high_index = iterations * iteration_stride + id; -#pragma unroll - for (int i = 0; i < iterations; i++) { - __half2 gamma_reg = gamma_h[i * iteration_stride + id]; - vals_arr[i] = out_grad_h[i * iteration_stride + id]; - vals_arr[i] *= gamma_reg; - vals_hat_arr[i] = - (invertible - ? (vals_hat_h[i * iteration_stride + id] - betta_h[i * iteration_stride + id]) / - gamma_reg - : vals_hat_h[i * iteration_stride + id]); - } - if ((high_index) < row_stride) { - __half2 gamma_reg = gamma_h[high_index]; - vals_arr[iterations] = out_grad_h[high_index]; - vals_arr[iterations] *= gamma_reg; - vals_hat_arr[iterations] = - (invertible ? (vals_hat_h[high_index] - betta_h[high_index]) / gamma_reg - : vals_hat_h[high_index]); - iterations++; - } - __half var_h = vars[row]; - __half2 var_reg = __halves2half2(var_h, var_h); - - float sum = 0.f; - for (int i = 0; i < iterations; i++) { - __half2 result_h = (vals_hat_arr[i] * vals_arr[i] * h2sqrt(var_reg)); - float2 result_f = __half22float2(result_h); - sum += result_f.x; - sum += result_f.y; - vals_arr[i] *= h2rsqrt(var_reg); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - - sum = g.shfl(sum, 0); - sum /= (2 * row_stride); - __half2 sum_h = __float2half2_rn(sum); - - for (int i = 0; i < iterations; i++) { - __half2 temp = ((-sum_h * vals_hat_arr[i]) / (var_reg)); - vals_arr_f[i] = __half22float2(vals_arr[i]); - float2 temp_f = __half22float2(temp); - vals_arr_f[i].x += temp_f.x; - vals_arr_f[i].y += temp_f.y; - } - sum = 0.f; - - for (int i = 0; i < iterations; i++) { - sum += (vals_arr_f[i].x); - sum += (vals_arr_f[i].y); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - - sum = g.shfl(sum, 0); - sum /= (2 * row_stride); - - iterations = row_stride / iteration_stride; - for (int i = 0; i < iterations; i++) { - vals_arr_f[i].x -= sum; - vals_arr_f[i].y -= sum; - __half2 temp = __float22half2_rn(vals_arr_f[i]); - - inp_grad_h[i * iteration_stride + id] = temp; - } - if ((high_index) < row_stride) { - vals_arr_f[iterations].x -= sum; - vals_arr_f[iterations].y -= sum; - __half2 temp = __float22half2_rn(vals_arr_f[iterations]); - - inp_grad_h[high_index] = temp; - } -#endif -} - -template <> -void launch_layerNorm_backward(const float* out_grad, - const float* vals_hat, - const float* vars, - const float* gamma, - float* gamma_grad, - float* betta_grad, - float* inp_grad, - int batch, - int hidden_dim, - cudaStream_t stream[2], - bool invertible, - const float* betta) -{ - int threads = THREADS; - - dim3 grid_dim(hidden_dim / TILE_DIM); - dim3 block_dim(TILE_DIM, TILE_DIM); - - LayerNormBackward1<<>>( - out_grad, vals_hat, gamma, betta, gamma_grad, betta_grad, batch, hidden_dim, invertible); - - dim3 grid_dim2(batch); - - if (hidden_dim > 16384 && hidden_dim <= 32768) - threads <<= 1; - else if (hidden_dim > 32768 && hidden_dim <= 65536) - threads <<= 2; - else if (hidden_dim > 65536) - throw std::runtime_error("Unsupport hidden_dim."); - - dim3 block_dim2(threads); - - LayerNormBackward2<<>>( - out_grad, vals_hat, gamma, betta, vars, inp_grad, invertible, hidden_dim); -} - -template <> -void launch_layerNorm_backward<__half>(const __half* out_grad, - const __half* vals_hat, - const __half* vars, - const __half* gamma, - __half* gamma_grad, - __half* betta_grad, - __half* inp_grad, - int batch, - int hidden_dim, - cudaStream_t stream[2], - bool invertible, - const __half* betta) -{ - int threads = THREADS; - - dim3 grid_dim(hidden_dim / TILE_DIM); - dim3 block_dim(TILE_DIM, TILE_DIM); - - // LayerNormBackward1<__half><<>>( - // out_grad, vals_hat, gamma, betta, gamma_grad, betta_grad, batch, hidden_dim, invertible); - - dim3 grid_dim2(batch); - - if (hidden_dim > 8192 && hidden_dim <= 16384) - threads <<= 1; - else if (hidden_dim > 16384 && hidden_dim <= 32768) - threads <<= 2; - else if (hidden_dim > 32768 && hidden_dim <= 65536) - threads <<= 3; - else if (hidden_dim > 65536) - throw std::runtime_error("Unsupport hidden_dim."); - - dim3 block_dim2(threads / 2); - - LayerNormBackward2<<>>( - out_grad, vals_hat, gamma, betta, vars, inp_grad, invertible, hidden_dim / 2); -} - -/* Backward Normalize (Input-Gradient) - * Using the means and variances from the input - * This type of backward is not invertible! - * We do the backward using the input (X) - */ - -__global__ void LayerNormBackward2(const float* out_grad, - const float* X_vals, - const float* gamma, - const float* vars, - const float* means, - float* inp_grad, - int row_stride) -{ - int iteration_stride = blockDim.x; - int iterations = row_stride / iteration_stride; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int row = blockIdx.x; - int id = threadIdx.x; - int wid = id >> WARP_SIZE_BITS; - int warp_num = iteration_stride >> WARP_SIZE_BITS; - __shared__ float partialSum[MAX_WARP_NUM]; - - out_grad += (row * row_stride); - X_vals += (row * row_stride); - inp_grad += (row * row_stride); - - float vals_arr[NORM_REG]; - int high_index = iterations * iteration_stride + id; -#pragma unroll - for (int i = 0; i < iterations; i++) { - float gamma_reg = gamma[i * iteration_stride + id]; - vals_arr[i] = out_grad[i * iteration_stride + id]; - vals_arr[i] *= gamma_reg; - } - if ((high_index) < row_stride) { - float gamma_reg = gamma[high_index]; - vals_arr[iterations] = out_grad[high_index]; - vals_arr[iterations] *= gamma_reg; - iterations++; - } - - float var_reg = vars[row]; - float mean_reg = means[row]; - - float sum = 0; - float xu[NORM_REG]; - for (int i = 0; i < iterations; i++) { - xu[i] = (X_vals[i * iteration_stride + id] - mean_reg); - sum += vals_arr[i] * xu[i]; - vals_arr[i] *= rsqrtf(var_reg); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - - sum = g.shfl(sum, 0); - sum /= row_stride; - - for (int i = 0; i < iterations; i++) { - vals_arr[i] += (-sum * xu[i] * rsqrtf(var_reg) / (var_reg)); - } - - sum = 0; - for (int i = 0; i < iterations; i++) { sum += vals_arr[i]; } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - sum = g.shfl(sum, 0); - sum /= row_stride; - - iterations = row_stride / iteration_stride; - for (int i = 0; i < iterations; i++) inp_grad[i * iteration_stride + id] = (vals_arr[i] - sum); - if ((high_index) < row_stride) inp_grad[high_index] = (vals_arr[iterations] - sum); -} - -__global__ void LayerNormBackward2(const __half* out_grad, - const __half* X_vals, - const __half* gamma, - const __half* vars, - const __half* means, - __half* inp_grad, - int row_stride) -{ -#ifdef HALF_PRECISION_AVAILABLE - int iteration_stride = blockDim.x; - int iterations = row_stride / iteration_stride; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int row = blockIdx.x; - int id = threadIdx.x; - int wid = id >> WARP_SIZE_BITS; - int warp_num = iteration_stride >> WARP_SIZE_BITS; - - __shared__ float partialSum[MAX_WARP_NUM]; - - __half2 vals_arr[NORM_REG]; - float2 vals_arr_f[NORM_REG]; - __half2 xu[NORM_REG]; - - __half2* inp_grad_h = reinterpret_cast<__half2*>(inp_grad); - const __half2* out_grad_h = reinterpret_cast(out_grad); - const __half2* vals_hat_h = reinterpret_cast(X_vals); - - inp_grad_h += (row * row_stride); - out_grad_h += (row * row_stride); - vals_hat_h += (row * row_stride); - - const __half2* gamma_h = reinterpret_cast(gamma); - int high_index = iterations * iteration_stride + id; - - __half mean_h = means[row]; - __half2 mean_reg = __halves2half2(mean_h, mean_h); -#pragma unroll - for (int i = 0; i < iterations; i++) { - __half2 gamma_reg = gamma_h[i * iteration_stride + id]; - vals_arr[i] = out_grad_h[i * iteration_stride + id]; - vals_arr[i] *= gamma_reg; // out_grad * gamma - xu[i] = (vals_hat_h[i * iteration_stride + id] - mean_reg); - } - if ((high_index) < row_stride) { - __half2 gamma_reg = gamma_h[high_index]; - vals_arr[iterations] = out_grad_h[high_index]; - vals_arr[iterations] *= gamma_reg; // out_grad * gamma - xu[iterations] = (vals_hat_h[high_index] - mean_reg); - iterations++; - } - __half var_h = vars[row]; - __half2 var_reg = __halves2half2(var_h, var_h); - - float sum = 0.f; - for (int i = 0; i < iterations; i++) { - __half2 result_h = (xu[i] * vals_arr[i]); - float2 result_f = __half22float2(result_h); - sum += result_f.x; - sum += result_f.y; - vals_arr[i] *= h2rsqrt(var_reg); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - - sum = g.shfl(sum, 0); - sum /= (2 * row_stride); - __half2 sum_h = __float2half2_rn(sum); - - for (int i = 0; i < iterations; i++) { - __half2 xu_grad = ((-sum_h * xu[i] * h2rsqrt(var_reg)) / (var_reg)); - vals_arr_f[i] = __half22float2(vals_arr[i]); - float2 xu_grad_f = __half22float2(xu_grad); - vals_arr_f[i].x += xu_grad_f.x; - vals_arr_f[i].y += xu_grad_f.y; - } - - sum = 0.f; - for (int i = 0; i < iterations; i++) { - sum += (vals_arr_f[i].x); - sum += (vals_arr_f[i].y); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - - sum = g.shfl(sum, 0); - sum /= (2 * row_stride); - - iterations = row_stride / iteration_stride; - for (int i = 0; i < iterations; i++) { - vals_arr_f[i].x -= sum; - vals_arr_f[i].y -= sum; - __half2 temp = __float22half2_rn(vals_arr_f[i]); - inp_grad_h[i * iteration_stride + id] = temp; - } - if ((high_index) < row_stride) { - vals_arr_f[iterations].x -= sum; - vals_arr_f[iterations].y -= sum; - __half2 temp = __float22half2_rn(vals_arr_f[iterations]); - inp_grad_h[high_index] = temp; - } -#endif -} - -template <> -void launch_layerNorm_backward(const float* out_grad, - const float* X_data, - const float* vars, - const float* means, - const float* gamma, - float* gamma_grad, - float* betta_grad, - float* inp_grad, - int batch, - int hidden_dim, - cudaStream_t stream[2]) -{ - int threads = THREADS; - - dim3 grid_dim(hidden_dim / TILE_DIM); - dim3 block_dim(TILE_DIM, TILE_DIM); - - LayerNormBackward1<<>>( - out_grad, X_data, vars, means, gamma_grad, betta_grad, batch, hidden_dim); - - dim3 grid_dim2(batch); - - if (hidden_dim > 16384 && hidden_dim <= 32768) - threads <<= 1; - else if (hidden_dim > 32768 && hidden_dim <= 65536) - threads <<= 2; - else if (hidden_dim > 65536) - throw std::runtime_error("Unsupport hidden_dim."); - - dim3 block_dim2(threads); - LayerNormBackward2<<>>( - out_grad, X_data, gamma, vars, means, inp_grad, hidden_dim); -} - -template <> -void launch_layerNorm_backward<__half>(const __half* out_grad, - const __half* X_data, - const __half* vars, - const __half* means, - const __half* gamma, - __half* gamma_grad, - __half* betta_grad, - __half* inp_grad, - int batch, - int hidden_dim, - cudaStream_t stream[2]) -{ - int threads = THREADS; - - dim3 grid_dim(hidden_dim / TILE_DIM); - dim3 block_dim(TILE_DIM, TILE_DIM); - - LayerNormBackward1<__half><<>>( - out_grad, X_data, vars, means, gamma_grad, betta_grad, batch, hidden_dim); - - dim3 grid_dim2(batch); - - if (hidden_dim > 8192 && hidden_dim <= 16384) - threads <<= 1; - else if (hidden_dim > 16384 && hidden_dim <= 32768) - threads <<= 2; - else if (hidden_dim > 32768 && hidden_dim <= 65536) - threads <<= 3; - else if (hidden_dim > 65536) - throw std::runtime_error("Unsupport hidden_dim."); - - dim3 block_dim2(threads / 2); - LayerNormBackward2<<>>( - out_grad, X_data, gamma, vars, means, inp_grad, hidden_dim / 2); -} - -template -__global__ void LayerNormBackward1_fused_add(const T* __restrict__ out_grad1, - const T* __restrict__ out_grad2, - const T* __restrict__ vals_hat, - const T* __restrict__ gamma, - const T* __restrict__ betta, - T* __restrict__ gamma_grad, - T* __restrict__ betta_grad, - int rows, - int width, - bool invertible) -{ - __shared__ float betta_buffer[TILE_DIM][TILE_DIM + 1]; - __shared__ float gamma_buffer[TILE_DIM][TILE_DIM + 1]; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int idx = blockDim.x * blockIdx.x + threadIdx.x; - int offset = threadIdx.y * width + idx; - int y_stride = width * TILE_DIM; - - float betta_reg = (invertible ? (float)betta[idx] : 0.0f); - float gamma_reg = (float)gamma[idx]; - - // Loop across matrix height - float betta_tmp = 0; - float gamma_tmp = 0; - for (int r = threadIdx.y; r < rows; r += TILE_DIM) { - float grad = (float)out_grad1[offset] + (float)out_grad2[offset]; - float val = (invertible ? ((float)vals_hat[offset] - betta_reg) / gamma_reg - : (float)vals_hat[offset]); - betta_tmp += grad; - gamma_tmp += (val * grad); - - offset += y_stride; - } - - betta_buffer[threadIdx.x][threadIdx.y] = betta_tmp; - gamma_buffer[threadIdx.x][threadIdx.y] = gamma_tmp; - - __syncthreads(); - - // Sum the shared buffer. - float s1 = betta_buffer[threadIdx.y][threadIdx.x]; - float s2 = gamma_buffer[threadIdx.y][threadIdx.x]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < TILE_DIM; i <<= 1) { - s1 += g.shfl_down(s1, i); - s2 += g.shfl_down(s2, i); - } - - if (threadIdx.x == 0) { - int pos = blockIdx.x * TILE_DIM + threadIdx.y; - betta_grad[pos] = s1; - gamma_grad[pos] = s2; - } -} - -template -__global__ void LayerNormBackward1_fused_add(const T* __restrict__ out_grad1, - const T* __restrict__ out_grad2, - const T* __restrict__ X_data, - const T* __restrict__ vars, - const T* __restrict__ means, - T* __restrict__ gamma_grad, - T* __restrict__ betta_grad, - int rows, - int width) -{ - __shared__ float betta_buffer[TILE_DIM][TILE_DIM + 1]; - __shared__ float gamma_buffer[TILE_DIM][TILE_DIM + 1]; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int idx = blockDim.x * blockIdx.x + threadIdx.x; - int offset = threadIdx.y * width + idx; - int y_stride = width * TILE_DIM; - - int pos = blockIdx.x * TILE_DIM + threadIdx.y; - // Loop across matrix height - - float betta_tmp = 0; - float gamma_tmp = 0; - for (int r = threadIdx.y; r < rows; r += TILE_DIM) { - float grad = (float)out_grad1[offset] + (float)out_grad2[offset]; - float val = (float)X_data[offset]; - val = (val - (float)means[r]) * rsqrtf((float)vars[r]); - betta_tmp += grad; - gamma_tmp += (val * grad); - - offset += y_stride; - } - - betta_buffer[threadIdx.x][threadIdx.y] = betta_tmp; - gamma_buffer[threadIdx.x][threadIdx.y] = gamma_tmp; - - __syncthreads(); - - // Sum the shared buffer. - float s1 = betta_buffer[threadIdx.y][threadIdx.x]; - float s2 = gamma_buffer[threadIdx.y][threadIdx.x]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < TILE_DIM; i <<= 1) { - s1 += g.shfl_down(s1, i); - s2 += g.shfl_down(s2, i); - } - - if (threadIdx.x == 0) { - betta_grad[pos] = s1; - gamma_grad[pos] = s2; - } -} - -__global__ void LayerNormBackward2_fused_add(const float* out_grad1, - const float* out_grad2, - const float* vals_hat, - const float* gamma, - const float* betta, - const float* vars, - float* inp_grad, - bool invertible, - int row_stride) -{ - int iteration_stride = blockDim.x; - int iterations = row_stride / iteration_stride; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int row = blockIdx.x; - int id = threadIdx.x; - int wid = id / WARP_SIZE; - int warp_num = iteration_stride >> WARP_SIZE_BITS; - __shared__ float partialSum[MAX_WARP_NUM]; - - out_grad1 += (row * row_stride); - out_grad2 += (row * row_stride); - vals_hat += (row * row_stride); - inp_grad += (row * row_stride); - - float vals_arr[NORM_REG]; - float vals_hat_arr[NORM_REG]; - int high_index = iterations * iteration_stride + id; -#pragma unroll - for (int i = 0; i < iterations; i++) { - float gamma_reg = gamma[i * iteration_stride + id]; - vals_arr[i] = out_grad1[i * iteration_stride + id]; - vals_arr[i] *= gamma_reg; - vals_hat_arr[i] = - (invertible ? (vals_hat[i * iteration_stride + id] - betta[i * iteration_stride + id]) / - gamma_reg - : vals_hat[i * iteration_stride + id]); - } - if ((high_index) < row_stride) { - float gamma_reg = gamma[high_index]; - vals_arr[iterations] = out_grad1[high_index]; - vals_arr[iterations] *= gamma_reg; - vals_hat_arr[iterations] = - (invertible ? (vals_hat[high_index] - betta[high_index]) / gamma_reg - : vals_hat[high_index]); - iterations++; - } - - float var_reg = vars[row]; - - float sum = 0; - for (int i = 0; i < iterations; i++) { - sum += vals_hat_arr[i] * vals_arr[i] * sqrtf(var_reg); - vals_arr[i] *= rsqrtf(var_reg); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - - sum = g.shfl(sum, 0); - sum /= row_stride; - - for (int i = 0; i < iterations; i++) { vals_arr[i] += ((-sum * vals_hat_arr[i]) / var_reg); } - - sum = 0; - for (int i = 0; i < iterations; i++) { sum += vals_arr[i]; } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - sum = g.shfl(sum, 0); - sum /= row_stride; - - iterations = row_stride / iteration_stride; - for (int i = 0; i < iterations; i++) - inp_grad[i * iteration_stride + id] = - (vals_arr[i] - sum) + out_grad2[i * iteration_stride + id]; - if ((high_index) < row_stride) - inp_grad[high_index] = (vals_arr[iterations] - sum) + out_grad2[high_index]; -} - -__global__ void LayerNormBackward2_fused_add(const __half* out_grad1, - const __half* out_grad2, - const __half* vals_hat, - const __half* gamma, - const __half* betta, - const __half* vars, - __half* inp_grad, - bool invertible, - int row_stride) -{ -#ifdef HALF_PRECISION_AVAILABLE - int iteration_stride = blockDim.x; - int iterations = row_stride / iteration_stride; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int row = blockIdx.x; - int id = threadIdx.x; - int wid = id / WARP_SIZE; - int warp_num = iteration_stride >> WARP_SIZE_BITS; - __shared__ float partialSum[MAX_WARP_NUM]; - - __half2 vals_arr[NORM_REG]; - float2 vals_arr_f[NORM_REG]; - __half2 vals_hat_arr[NORM_REG]; - - // float2 result[iterations]; - - __half2* inp_grad_h = reinterpret_cast<__half2*>(inp_grad); - const __half2* out_grad_h1 = reinterpret_cast(out_grad1); - const __half2* out_grad_h2 = reinterpret_cast(out_grad2); - const __half2* vals_hat_h = reinterpret_cast(vals_hat); - - inp_grad_h += (row * row_stride); - out_grad_h1 += (row * row_stride); - out_grad_h2 += (row * row_stride); - vals_hat_h += (row * row_stride); - - const __half2* gamma_h = reinterpret_cast(gamma); - const __half2* betta_h = (invertible ? reinterpret_cast(betta) : nullptr); - int high_index = iterations * iteration_stride + id; -#pragma unroll - for (int i = 0; i < iterations; i++) { - __half2 gamma_reg = gamma_h[i * iteration_stride + id]; - vals_arr[i] = out_grad_h1[i * iteration_stride + id]; - vals_arr[i] *= gamma_reg; // out_grad * gamma - vals_hat_arr[i] = - (invertible - ? (vals_hat_h[i * iteration_stride + id] - betta_h[i * iteration_stride + id]) / - gamma_reg - : vals_hat_h[i * iteration_stride + id]); - } - if ((high_index) < row_stride) { - __half2 gamma_reg = gamma_h[high_index]; - vals_arr[iterations] = out_grad_h1[high_index]; - vals_arr[iterations] *= gamma_reg; // out_grad * gamma - vals_hat_arr[iterations] = - (invertible ? (vals_hat_h[high_index] - betta_h[high_index]) / gamma_reg - : vals_hat_h[high_index]); - iterations++; - } - __half var_h = vars[row]; - __half2 var_reg = __halves2half2(var_h, var_h); - - float sum = 0.f; - for (int i = 0; i < iterations; i++) { - __half2 result_h = (vals_hat_arr[i] * vals_arr[i] * h2sqrt(var_reg)); - float2 result_f = __half22float2(result_h); - sum += result_f.x; - sum += result_f.y; - vals_arr[i] *= h2rsqrt(var_reg); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - - sum = g.shfl(sum, 0); - sum /= (2 * row_stride); - __half2 sum_h = __float2half2_rn(sum); - - for (int i = 0; i < iterations; i++) { - __half2 temp = ((-sum_h * vals_hat_arr[i]) / (var_reg)); - vals_arr_f[i] = __half22float2(vals_arr[i]); - float2 temp_f = __half22float2(temp); - vals_arr_f[i].x += temp_f.x; - vals_arr_f[i].y += temp_f.y; - } - sum = 0.f; - for (int i = 0; i < iterations; i++) { - sum += (vals_arr_f[i].x); - sum += (vals_arr_f[i].y); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - - sum = g.shfl(sum, 0); - sum /= (2 * row_stride); - - iterations = row_stride / iteration_stride; - for (int i = 0; i < iterations; i++) { - vals_arr_f[i].x -= sum; - vals_arr_f[i].y -= sum; - __half2 temp = __float22half2_rn(vals_arr_f[i]); - - inp_grad_h[i * iteration_stride + id] = temp + out_grad_h2[i * iteration_stride + id]; - } - if ((high_index) < row_stride) { - vals_arr_f[iterations].x -= sum; - vals_arr_f[iterations].y -= sum; - __half2 temp = __float22half2_rn(vals_arr_f[iterations]); - - inp_grad_h[high_index] = temp + out_grad_h2[high_index]; - } -#endif -} - -template <> -void launch_layerNorm_backward_fused_add(const float* out_grad1, - const float* out_grad2, - const float* vals_hat, - const float* vars, - const float* gamma, - float* gamma_grad, - float* betta_grad, - float* inp_grad, - int batch, - int hidden_dim, - cudaStream_t stream[2], - bool invertible, - const float* betta) -{ - int threads = THREADS; - - dim3 grid_dim(hidden_dim / TILE_DIM); - dim3 block_dim(TILE_DIM, TILE_DIM); - LayerNormBackward1<<>>( - out_grad1, vals_hat, gamma, betta, gamma_grad, betta_grad, batch, hidden_dim, invertible); - - dim3 grid_dim2(batch); - - if (hidden_dim > 16384 && hidden_dim <= 32768) - threads <<= 1; - else if (hidden_dim > 32768 && hidden_dim <= 65536) - threads <<= 2; - else if (hidden_dim > 65536) - throw std::runtime_error("Unsupport hidden_dim."); - - dim3 block_dim2(threads); - LayerNormBackward2_fused_add<<>>( - out_grad1, out_grad2, vals_hat, gamma, betta, vars, inp_grad, invertible, hidden_dim); -} - -template <> -void launch_layerNorm_backward_fused_add<__half>(const __half* out_grad1, - const __half* out_grad2, - const __half* vals_hat, - const __half* vars, - const __half* gamma, - __half* gamma_grad, - __half* betta_grad, - __half* inp_grad, - int batch, - int hidden_dim, - cudaStream_t stream[2], - bool invertible, - const __half* betta) -{ - int threads = THREADS; - - dim3 grid_dim(hidden_dim / TILE_DIM); - dim3 block_dim(TILE_DIM, TILE_DIM); - - LayerNormBackward1<__half><<>>( - out_grad1, vals_hat, gamma, betta, gamma_grad, betta_grad, batch, hidden_dim, invertible); - - dim3 grid_dim2(batch); - - if (hidden_dim > 8192 && hidden_dim <= 16384) - threads <<= 1; - else if (hidden_dim > 16384 && hidden_dim <= 32768) - threads <<= 2; - else if (hidden_dim > 32768 && hidden_dim <= 65536) - threads <<= 3; - else if (hidden_dim > 65536) - throw std::runtime_error("Unsupport hidden_dim."); - - dim3 block_dim2(threads / 2); - LayerNormBackward2_fused_add<<>>( - out_grad1, out_grad2, vals_hat, gamma, betta, vars, inp_grad, invertible, hidden_dim / 2); -} - -/* Backward Normalize (Input-Gradient) - * Using the means and variances from the input - * This type of backward is not invertible! - * We do the backward using the input (X) - */ - -__global__ void LayerNormBackward2_fused_add(const float* out_grad1, - const float* out_grad2, - const float* X_vals, - const float* gamma, - const float* vars, - const float* means, - float* inp_grad, - int row_stride) -{ - int iteration_stride = blockDim.x; - int iterations = row_stride / iteration_stride; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int row = blockIdx.x; - int id = threadIdx.x; - int wid = id / WARP_SIZE; - int warp_num = iteration_stride >> WARP_SIZE_BITS; - __shared__ float partialSum[MAX_WARP_NUM]; - - float vals_arr[NORM_REG]; - float vals_hat_arr[NORM_REG]; - - out_grad1 += (row * row_stride); - out_grad2 += (row * row_stride); - X_vals += (row * row_stride); - inp_grad += (row * row_stride); - int high_index = iterations * iteration_stride + id; -#pragma unroll - for (int i = 0; i < iterations; i++) { - float gamma_reg = gamma[i * iteration_stride + id]; - vals_arr[i] = out_grad1[i * iteration_stride + id]; - vals_arr[i] *= gamma_reg; - vals_hat_arr[i] = X_vals[i * iteration_stride + id]; - } - if ((high_index) < row_stride) { - float gamma_reg = gamma[high_index]; - vals_arr[iterations] = out_grad1[high_index]; - vals_arr[iterations] *= gamma_reg; - vals_hat_arr[iterations] = X_vals[high_index]; - iterations++; - } - - float var_reg = vars[row]; - float mean_reg = means[row]; - - float sum = 0; - float xu[NORM_REG]; - for (int i = 0; i < iterations; i++) { - xu[i] = (vals_hat_arr[i] - mean_reg); - sum += vals_arr[i] * xu[i]; - vals_arr[i] *= rsqrtf(var_reg); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - - sum = g.shfl(sum, 0); - sum /= row_stride; - - for (int i = 0; i < iterations; i++) { - vals_arr[i] += (-sum * xu[i] * rsqrtf(var_reg) / (var_reg)); - } - - sum = 0; - for (int i = 0; i < iterations; i++) { sum += vals_arr[i]; } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - sum = g.shfl(sum, 0); - sum /= row_stride; - - iterations = row_stride / iteration_stride; - for (int i = 0; i < iterations; i++) - inp_grad[i * iteration_stride + id] = - (vals_arr[i] - sum) + out_grad2[i * iteration_stride + id]; - if ((high_index) < row_stride) - inp_grad[high_index] = (vals_arr[iterations] - sum) + out_grad2[high_index]; -} - -__global__ void LayerNormBackward2_fused_add(const __half* out_grad1, - const __half* out_grad2, - const __half* X_vals, - const __half* gamma, - const __half* vars, - const __half* means, - __half* inp_grad, - int row_stride) -{ -#ifdef HALF_PRECISION_AVAILABLE - int iteration_stride = blockDim.x; - int iterations = row_stride / iteration_stride; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int row = blockIdx.x; - int id = threadIdx.x; - int wid = id / WARP_SIZE; - int warp_num = iteration_stride >> WARP_SIZE_BITS; - - __shared__ float partialSum[MAX_WARP_NUM]; - - __half2 vals_arr[NORM_REG]; - float2 vals_arr_f[NORM_REG]; - __half2 vals_hat_arr[NORM_REG]; - - __half2* inp_grad_h = reinterpret_cast<__half2*>(inp_grad); - const __half2* out_grad_h1 = reinterpret_cast(out_grad1); - const __half2* out_grad_h2 = reinterpret_cast(out_grad2); - const __half2* vals_hat_h = reinterpret_cast(X_vals); - - out_grad_h1 += (row * row_stride); - out_grad_h2 += (row * row_stride); - inp_grad_h += (row * row_stride); - vals_hat_h += (row * row_stride); - - const __half2* gamma_h = reinterpret_cast(gamma); - int high_index = iterations * iteration_stride + id; -#pragma unroll - for (int i = 0; i < iterations; i++) { - __half2 gamma_reg = gamma_h[i * iteration_stride + id]; - vals_arr[i] = out_grad_h1[i * iteration_stride + id]; - vals_arr[i] *= gamma_reg; // out_grad * gamma - vals_hat_arr[i] = vals_hat_h[i * iteration_stride + id]; - } - if ((high_index) < row_stride) { - __half2 gamma_reg = gamma_h[high_index]; - vals_arr[iterations] = out_grad_h1[high_index]; - vals_arr[iterations] *= gamma_reg; // out_grad * gamma - vals_hat_arr[iterations] = vals_hat_h[high_index]; - iterations++; - } - - __half mean_h = means[row]; - __half var_h = vars[row]; - __half2 var_reg = __halves2half2(var_h, var_h); - __half2 mean_reg = __halves2half2(mean_h, mean_h); - __half2 xu[NORM_REG]; - - float sum = 0.f; - for (int i = 0; i < iterations; i++) { - xu[i] = (vals_hat_arr[i] - mean_reg); - __half2 result_h = (xu[i] * vals_arr[i]); - float2 result_f = __half22float2(result_h); - sum += result_f.x; - sum += result_f.y; - vals_arr[i] *= h2rsqrt(var_reg); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - - sum = g.shfl(sum, 0); - sum /= (2 * row_stride); - __half2 sum_h = __float2half2_rn(sum); - - for (int i = 0; i < iterations; i++) { - __half2 xu_grad = ((-sum_h * xu[i] * h2rsqrt(var_reg)) / (var_reg)); - vals_arr_f[i] = __half22float2(vals_arr[i]); - float2 xu_grad_f = __half22float2(xu_grad); - vals_arr_f[i].x += xu_grad_f.x; - vals_arr_f[i].y += xu_grad_f.y; - } - - sum = 0.f; - for (int i = 0; i < iterations; i++) { - sum += (vals_arr_f[i].x); - sum += (vals_arr_f[i].y); - } - - for (int i = 1; i < WARP_SIZE; i *= 2) { sum += g.shfl_down(sum, i); } - - if (g.thread_rank() == 0) partialSum[wid] = sum; - - __syncthreads(); - - if (g.thread_rank() < warp_num) sum = partialSum[g.thread_rank()]; - -#ifndef __STOCHASTIC_MODE__ - __syncthreads(); -#endif - - for (int i = 1; i < warp_num; i *= 2) sum += g.shfl_down(sum, i); - - sum = g.shfl(sum, 0); - sum /= (2 * row_stride); - - iterations = row_stride / iteration_stride; - for (int i = 0; i < iterations; i++) { - vals_arr_f[i].x -= sum; - vals_arr_f[i].y -= sum; - __half2 temp = __float22half2_rn(vals_arr_f[i]); - inp_grad_h[i * iteration_stride + id] = temp + out_grad_h2[i * iteration_stride + id]; - } - if ((high_index) < row_stride) { - vals_arr_f[iterations].x -= sum; - vals_arr_f[iterations].y -= sum; - __half2 temp = __float22half2_rn(vals_arr_f[iterations]); - inp_grad_h[high_index] = temp + out_grad_h2[high_index]; - } -#endif -} - -template <> -void launch_layerNorm_backward_fused_add(const float* out_grad1, - const float* out_grad2, - const float* X_data, - const float* vars, - const float* means, - const float* gamma, - float* gamma_grad, - float* betta_grad, - float* inp_grad, - int batch, - int hidden_dim, - cudaStream_t stream[2]) -{ - int threads = THREADS; - - dim3 grid_dim(hidden_dim / TILE_DIM); - dim3 block_dim(TILE_DIM, TILE_DIM); - - LayerNormBackward1<<>>( - out_grad1, X_data, vars, means, gamma_grad, betta_grad, batch, hidden_dim); - - dim3 grid_dim2(batch); - - if (hidden_dim > 16384 && hidden_dim <= 32768) - threads <<= 1; - else if (hidden_dim > 32768 && hidden_dim <= 65536) - threads <<= 2; - else if (hidden_dim > 65536) - throw std::runtime_error("Unsupport hidden_dim."); - - dim3 block_dim2(threads); - LayerNormBackward2_fused_add<<>>( - out_grad1, out_grad2, X_data, gamma, vars, means, inp_grad, hidden_dim); -} - -template <> -void launch_layerNorm_backward_fused_add<__half>(const __half* out_grad1, - const __half* out_grad2, - const __half* X_data, - const __half* vars, - const __half* means, - const __half* gamma, - __half* gamma_grad, - __half* betta_grad, - __half* inp_grad, - int batch, - int hidden_dim, - cudaStream_t stream[2]) -{ - int threads = THREADS; - - dim3 grid_dim(hidden_dim / TILE_DIM); - dim3 block_dim(TILE_DIM, TILE_DIM); - - LayerNormBackward1<__half><<>>( - out_grad1, X_data, vars, means, gamma_grad, betta_grad, batch, hidden_dim); - - dim3 grid_dim2(batch); - - if (hidden_dim > 8192 && hidden_dim <= 16384) - threads <<= 1; - else if (hidden_dim > 16384 && hidden_dim <= 32768) - threads <<= 2; - else if (hidden_dim > 32768 && hidden_dim <= 65536) - threads <<= 3; - else if (hidden_dim > 65536) - throw std::runtime_error("Unsupport hidden_dim."); - - dim3 block_dim2(threads / 2); - LayerNormBackward2_fused_add<<>>( - out_grad1, out_grad2, X_data, gamma, vars, means, inp_grad, hidden_dim / 2); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/softmax_kernels.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/softmax_kernels.cu deleted file mode 100644 index d9efed286462f1be971922c2345efb8d2049e227..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/softmax_kernels.cu +++ /dev/null @@ -1,701 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include "custom_cuda_layers.h" -#include "general_kernels.h" - -namespace cg = cooperative_groups; - -dim3 get_attn_softmax_grid(int batch_size, int heads, int sequence_length, int threads) -{ - int seq_length4 = sequence_length / 4; - int block_compute_size = - (seq_length4 < threads ? (int)pow(2.0, floor(log2((float)(threads / seq_length4)))) : 1); - // Note that the Y and Z dimensions are limited to 65535, while X is basically unlimited: - // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#features-and-technical-specifications - // The batch size is typically relatively small, while the sequence length could potentially be - // arbitrarily large. We therefore place the batch size second to avoid hitting the Y limit. - unsigned x = heads * sequence_length / block_compute_size; - unsigned y = batch_size; - return {x, y}; -} - -// Fused attention + softmax -template -__global__ void attn_softmax(float* vals, - const float* attn_mask, - int heads, - int seq_length, - int iterations) -{ - __shared__ float partialSum[MAX_WARP_NUM]; - - int warp_num = blockDim.x >> WARP_SIZE_BITS; - - int iteration_stride = blockDim.x; - int block_width = blockStride * seq_length; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int batch = blockIdx.y; - int row = blockIdx.x; - int max_threads_in_sequence = std::max(seq_length, tbSeq); - int seq_lane = threadIdx.x % max_threads_in_sequence; - - int data_offset = batch * (gridDim.x * block_width) + row * block_width + - (threadIdx.x / max_threads_in_sequence) * seq_length; - int mask_offset = batch * seq_length; - - int wid = threadIdx.x >> WARP_SIZE_BITS; - int lane = threadIdx.x & 0x1f; - - float4* val_cast = reinterpret_cast(vals); - const float4* attn_mask_cast = reinterpret_cast(attn_mask); - - float4 data[MAX_THREAD_ITERATIONS]; - - float max_val = minus_infinity; - - for (int i = 0; i < iterations; i++) { - int data_id = i * iteration_stride + seq_lane; - if (data_id < seq_length) { - float4 mask = attn_mask_cast[mask_offset + data_id]; - data[i] = val_cast[data_offset + data_id]; - - data[i].x += mask.x; - data[i].y += mask.y; - data[i].z += mask.z; - data[i].w += mask.w; - - max_val = (data[i].x > max_val ? data[i].x : max_val); - max_val = (data[i].y > max_val ? data[i].y : max_val); - max_val = (data[i].z > max_val ? data[i].z : max_val); - max_val = (data[i].w > max_val ? data[i].w : max_val); - } else { - data[i].x = minus_infinity; - data[i].y = minus_infinity; - data[i].z = minus_infinity; - data[i].w = minus_infinity; - } - } - - for (int i = 1; i < tbSize; i *= 2) { - auto temp = g.shfl_xor(max_val, i); - max_val = (temp > max_val ? temp : max_val); - } - - if (seq_length > tbSize) { - if (lane == 0) partialSum[wid] = max_val; - b.sync(); - - if (lane < warp_num) max_val = partialSum[lane]; - -#ifndef __STOCHASTIC_MODE__ - b.sync(); -#endif - - int iters = warp_num; - if (seq_length < iteration_stride) - iters = warp_num / (iteration_stride / max_threads_in_sequence); - - for (int i = 1; i < iters; i *= 2) { - auto temp = g.shfl_xor(max_val, i); - max_val = (temp > max_val ? temp : max_val); - } - - max_val = g.shfl(max_val, threadIdx.x / tbSize); - } - - float sum = 0; - for (int i = 0; i < iterations; i++) { - data[i].x = __expf(data[i].x - max_val); - data[i].y = __expf(data[i].y - max_val); - data[i].z = __expf(data[i].z - max_val); - data[i].w = __expf(data[i].w - max_val); - - sum += (data[i].x + data[i].y + data[i].z + data[i].w); - } - - for (int i = 1; i < tbSize; i *= 2) { sum += g.shfl_xor(sum, i); } - - if (seq_length > tbSize) { - if (lane == 0) partialSum[wid] = sum; - b.sync(); - - if (lane < warp_num) sum = partialSum[lane]; - -#ifndef __STOCHASTIC_MODE__ - b.sync(); -#endif - - int iters = warp_num; - if (seq_length < iteration_stride) - iters = warp_num / (iteration_stride / max_threads_in_sequence); - - for (int i = 1; i < iters; i *= 2) { sum += g.shfl_xor(sum, i); } - - sum = g.shfl(sum, threadIdx.x / tbSize); - } - - sum += 1e-6; - - for (int i = 0; i < iterations; i++) { - data[i].x /= sum; - data[i].y /= sum; - data[i].z /= sum; - data[i].w /= sum; - - int data_id = i * iteration_stride + seq_lane; - if (data_id < seq_length) val_cast[data_offset + data_id] = data[i]; - } -} - -template -__global__ void attn_softmax(__half* vals, - const __half* attn_mask, - int heads, - int seq_length, - int iterations) -{ -#ifdef HALF_PRECISION_AVAILABLE - __shared__ float partialSum[MAX_WARP_NUM]; - - int warp_num = blockDim.x >> WARP_SIZE_BITS; - - int iteration_stride = blockDim.x; - int block_width = blockStride * seq_length; - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int batch = blockIdx.y; - int row = blockIdx.x; - int max_threads_in_sequence = std::max(seq_length, tbSeq); - int seq_lane = threadIdx.x % max_threads_in_sequence; - - int data_offset = batch * (gridDim.x * block_width) + row * block_width + - (threadIdx.x / max_threads_in_sequence) * seq_length; - int mask_offset = batch * seq_length; - - int wid = threadIdx.x >> WARP_SIZE_BITS; - int lane = threadIdx.x & 0x1f; - - float2* val_cast = reinterpret_cast(vals); - const float2* attn_mask_cast = reinterpret_cast(attn_mask); - - val_cast += data_offset; - attn_mask_cast += mask_offset; - - float2 low_data[MAX_THREAD_ITERATIONS]; - float2 high_data[MAX_THREAD_ITERATIONS]; - - float max_val = minus_infinity; - - for (int i = 0; i < iterations; i++) { - int data_id = i * iteration_stride + seq_lane; - if (data_id < seq_length) { - float2 data = val_cast[data_id]; - float2 mask = attn_mask_cast[data_id]; - - __half2* data_arr = reinterpret_cast<__half2*>(&data); - __half2* mask_arr = reinterpret_cast<__half2*>(&mask); - - low_data[i] = __half22float2(data_arr[0]); - high_data[i] = __half22float2(data_arr[1]); - float2 low_mask = __half22float2(mask_arr[0]); - float2 high_mask = __half22float2(mask_arr[1]); - - low_data[i].x += low_mask.x; - low_data[i].y += low_mask.y; - high_data[i].x += high_mask.x; - high_data[i].y += high_mask.y; - - max_val = (low_data[i].x > max_val ? low_data[i].x : max_val); - max_val = (low_data[i].y > max_val ? low_data[i].y : max_val); - max_val = (high_data[i].x > max_val ? high_data[i].x : max_val); - max_val = (high_data[i].y > max_val ? high_data[i].y : max_val); - } - } - - for (int i = 1; i < tbSize; i *= 2) { - auto temp = g.shfl_xor(max_val, i); - max_val = (temp > max_val ? temp : max_val); - } - - if (seq_length > tbSize) { - if (lane == 0) partialSum[wid] = max_val; - b.sync(); - - if (lane < warp_num) max_val = partialSum[lane]; - -#ifndef __STOCHASTIC_MODE__ - b.sync(); -#endif - - int iters = warp_num; - if (seq_length < iteration_stride) - iters = warp_num / (iteration_stride / max_threads_in_sequence); - - for (int i = 1; i < iters; i *= 2) { - auto temp = g.shfl_xor(max_val, i); - max_val = (temp > max_val ? temp : max_val); - } - - max_val = g.shfl(max_val, threadIdx.x / tbSize); - } - - float sum = 0; - for (int i = 0; i < iterations; i++) { - int data_id = i * iteration_stride + seq_lane; - if (data_id < seq_length) { - low_data[i].x = __expf(low_data[i].x - max_val); - low_data[i].y = __expf(low_data[i].y - max_val); - high_data[i].x = __expf(high_data[i].x - max_val); - high_data[i].y = __expf(high_data[i].y - max_val); - - sum += (low_data[i].x + low_data[i].y + high_data[i].x + high_data[i].y); - } - } - - for (int i = 1; i < tbSize; i *= 2) { sum += g.shfl_xor(sum, i); } - - if (seq_length > tbSize) { - if (lane == 0) partialSum[wid] = sum; - b.sync(); - - if (lane < warp_num) sum = partialSum[lane]; - -#ifndef __STOCHASTIC_MODE__ - b.sync(); -#endif - - int iters = warp_num; - if (seq_length < iteration_stride) - iters = warp_num / (iteration_stride / max_threads_in_sequence); - - for (int i = 1; i < iters; i *= 2) { sum += g.shfl_xor(sum, i); } - - sum = g.shfl(sum, threadIdx.x / tbSize); - } - - sum += 1e-6; - - for (int i = 0; i < iterations; i++) { - int data_id = i * iteration_stride + seq_lane; - if (data_id < seq_length) { - float2 result_f; - __half2* result_h = reinterpret_cast<__half2*>(&result_f); - - low_data[i].x /= sum; - low_data[i].y /= sum; - high_data[i].x /= sum; - high_data[i].y /= sum; - - result_h[0] = __float22half2_rn(low_data[i]); - result_h[1] = __float22half2_rn(high_data[i]); - - val_cast[data_id] = result_f; - } - } - -#endif -} - -template -void launch_attn_softmax(T*, const T*, int, int, int, cudaStream_t); - -template <> -void launch_attn_softmax(float* vals, - const float* attn_mask, - int batch_size, - int heads, - int sequence_length, - cudaStream_t stream) -{ - const int threads = 128; - int seq_length4 = sequence_length / 4; - - dim3 grid_dim = get_attn_softmax_grid(batch_size, heads, sequence_length, threads); - - int subblock_max_workload = MAX_THREAD_ITERATIONS * 4 * threads; - - dim3 block_dim(seq_length4 > threads ? ((sequence_length + subblock_max_workload - 1) / - subblock_max_workload * threads) - : threads); - int iterations = - (sequence_length < subblock_max_workload ? (seq_length4 + threads - 1) / threads - : MAX_THREAD_ITERATIONS); - - if (sequence_length <= 8) - attn_softmax<2, (threads / 2), 2> - <<>>(vals, attn_mask, heads, seq_length4, iterations); - else if (sequence_length <= 16) - attn_softmax<4, (threads / 4), 4> - <<>>(vals, attn_mask, heads, seq_length4, iterations); - else if (sequence_length <= 32) - attn_softmax<8, (threads / 8), 8> - <<>>(vals, attn_mask, heads, seq_length4, iterations); - else if (sequence_length <= 64) - attn_softmax<16, (threads / 16), 16> - <<>>(vals, attn_mask, heads, seq_length4, iterations); - else if (sequence_length <= 128) - attn_softmax<32, (threads / 32), 32> - <<>>(vals, attn_mask, heads, seq_length4, iterations); - else if (sequence_length <= 256) - attn_softmax<32, (threads / 64), 64> - <<>>(vals, attn_mask, heads, seq_length4, iterations); - else { - const int threads = 256; - dim3 grid_dim = get_attn_softmax_grid(batch_size, heads, sequence_length, threads); - - int subblock_max_workload = MAX_THREAD_ITERATIONS * 4 * threads; - - dim3 block_dim(seq_length4 > threads ? ((sequence_length + subblock_max_workload - 1) / - subblock_max_workload * threads) - : threads); - iterations = - (sequence_length < subblock_max_workload ? (seq_length4 + threads - 1) / threads - : MAX_THREAD_ITERATIONS); - if (sequence_length <= 512) - attn_softmax<32, (threads / 128), 128><<>>( - vals, attn_mask, heads, seq_length4, iterations); - else if (sequence_length < (MAX_THREADS * MAX_THREAD_ITERATIONS * 4)) - attn_softmax<32, 1, 128><<>>( - vals, attn_mask, heads, seq_length4, iterations); - else - throw std::runtime_error( - "Unsupport Seq_Length! Check the restriction of the max_threads and " - "max_thread_iterations!"); - } -} - -template <> -void launch_attn_softmax<__half>(__half* vals, - const __half* attn_mask, - int batch_size, - int heads, - int sequence_length, - cudaStream_t stream) -{ - const int threads = 128; - int seq_length4 = sequence_length / 4; - - dim3 grid_dim = get_attn_softmax_grid(batch_size, heads, sequence_length, threads); - - int subblock_max_workload = MAX_THREAD_ITERATIONS * 4 * threads; - - dim3 block_dim(seq_length4 > threads ? ((sequence_length + subblock_max_workload - 1) / - subblock_max_workload * threads) - : threads); - - int iterations = - (sequence_length < subblock_max_workload ? (seq_length4 + threads - 1) / threads - : MAX_THREAD_ITERATIONS); - - if (sequence_length <= 8) - attn_softmax<2, (threads / 2), 2> - <<>>(vals, attn_mask, heads, seq_length4, iterations); - else if (sequence_length <= 16) - attn_softmax<4, (threads / 4), 4> - <<>>(vals, attn_mask, heads, seq_length4, iterations); - else if (sequence_length <= 32) - attn_softmax<8, (threads / 8), 8> - <<>>(vals, attn_mask, heads, seq_length4, iterations); - else if (sequence_length <= 64) - attn_softmax<16, (threads / 16), 16> - <<>>(vals, attn_mask, heads, seq_length4, iterations); - else if (sequence_length <= 128) - attn_softmax<32, (threads / 32), 32> - <<>>(vals, attn_mask, heads, seq_length4, iterations); - else if (sequence_length <= 256) - attn_softmax<32, (threads / 64), 64> - <<>>(vals, attn_mask, heads, seq_length4, iterations); - else { - const int threads = 256; - dim3 grid_dim = get_attn_softmax_grid(batch_size, heads, sequence_length, threads); - - int subblock_max_workload = MAX_THREAD_ITERATIONS * 4 * threads; - - dim3 block_dim(seq_length4 > threads ? ((sequence_length + subblock_max_workload - 1) / - subblock_max_workload * threads) - : threads); - iterations = - (sequence_length < subblock_max_workload ? (seq_length4 + threads - 1) / threads - : MAX_THREAD_ITERATIONS); - if (sequence_length <= 512) - attn_softmax<32, (threads / 128), 128><<>>( - vals, attn_mask, heads, seq_length4, iterations); - else if (sequence_length < (MAX_THREADS * MAX_THREAD_ITERATIONS * 4)) - attn_softmax<32, 1, 128><<>>( - vals, attn_mask, heads, seq_length4, iterations); - else - throw std::runtime_error( - "Unsupport Seq_Length! Check the restriction of the max_threads and " - "max_thread_iterations!"); - } -} - -template -__global__ void softmax_backward_kernel(T* out_grad, const T* soft_inp, int seq_length) -{ - __shared__ float partialSum[MAX_WARP_NUM]; - - int warp_num = blockDim.x >> WARP_SIZE_BITS; // warp-count = num_threads / WARP_SIZE (32) - - int iteration_stride = blockDim.x; - int block_width = blockStride * seq_length; - - int iterations = (seq_length < (MAX_THREAD_ITERATIONS * iteration_stride) - ? (seq_length + iteration_stride - 1) / iteration_stride - : MAX_THREAD_ITERATIONS); - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int row = blockIdx.x; - int id = threadIdx.x; - - int wid = id >> WARP_SIZE_BITS; - int lane = id & 0x1f; - - T val_reg[MAX_THREAD_ITERATIONS]; - T soft_reg[MAX_THREAD_ITERATIONS]; - float grad_reg = 0.0f; - -#pragma unroll - for (int i = 0; i < iterations; i++) { - int data_id = i * iteration_stride + id; - if (data_id < block_width) { - val_reg[i] = out_grad[row * block_width + data_id]; - soft_reg[i] = soft_inp[row * block_width + data_id]; - - grad_reg += ((float)val_reg[i] * - (float)soft_reg[i]); // if done in half, the multiplication, we may lose - // 2% of accuracy in computation!! - } - } - for (int i = 1; i < tbSize; i *= 2) grad_reg += g.shfl_xor(grad_reg, i); - - if (seq_length > tbSize) { - if (lane == 0) partialSum[wid] = grad_reg; - b.sync(); - - if (lane < warp_num) grad_reg = partialSum[lane]; - - int iters = warp_num; - if (seq_length < iteration_stride) iters = warp_num / (iteration_stride / seq_length); - - for (int i = 1; i < iters; i *= 2) grad_reg += g.shfl_xor(grad_reg, i); - - grad_reg = g.shfl(grad_reg, id / tbSize); - } - - for (int i = 0; i < iterations; i++) { - int data_id = i * iteration_stride + id; - if (data_id < block_width) { - float temp = (float)soft_reg[i] * ((float)val_reg[i] - grad_reg); - out_grad[row * block_width + data_id] = (T)temp; - } - } -} - -template -__global__ void softmax_backward_kernel_v2(T* grad /* input & output*/, - const T* output, - int softmax_length) -{ - int batch_idx = blockIdx.x * blockDim.y + threadIdx.y; - int offset = batch_idx * softmax_length + threadIdx.x; - - grad += offset; - output += offset; - - T grad_reg[ITERATIONS]; - T output_reg[ITERATIONS]; - float sum = 0.0; - -#pragma unroll - for (int i = 0; i < ITERATIONS; ++i) { - int curr_idx = threadIdx.x + i * WARP_SIZE; - if (curr_idx < softmax_length) { - grad_reg[i] = grad[i * WARP_SIZE]; - output_reg[i] = output[i * WARP_SIZE]; - sum += (float)grad_reg[i] * (float)output_reg[i]; - } - } - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - for (int i = 1; i < WARP_SIZE; i <<= 1) sum += g.shfl_xor(sum, i); - -#pragma unroll - for (int i = 0; i < ITERATIONS; ++i) { - int curr_idx = threadIdx.x + i * WARP_SIZE; - if (curr_idx < softmax_length) - grad[i * WARP_SIZE] = (float)output_reg[i] * ((float)grad_reg[i] - sum); - } -} - -__global__ void softmax_backward_kernel_arbitrary_length(__half* grad /* input & output*/, - const __half* output, - int softmax_length) -{ - int batch_idx = blockIdx.x * blockDim.y + threadIdx.y; - int offset = batch_idx * softmax_length + threadIdx.x; - - const float4* output_cast = reinterpret_cast(output); - float4* grad_cast = reinterpret_cast(grad); - - grad_cast += offset; - output_cast += offset; - - float sum = 0.0; - int curr_idx = threadIdx.x; - while (curr_idx < softmax_length) { - float4 out_reg = output_cast[curr_idx]; - float4 grad_reg = grad_cast[curr_idx]; - __half2* out_h = reinterpret_cast<__half2*>(&out_reg); - __half2* grad_h = reinterpret_cast<__half2*>(&grad_reg); -#pragma unroll - for (int m = 0; m < 4; m++) grad_h[m] *= out_h[m]; - sum += ((float)grad_h[0].x + (float)grad_h[0].y + (float)grad_h[1].x + (float)grad_h[1].y) + - ((float)grad_h[2].x + (float)grad_h[2].y + (float)grad_h[3].x + (float)grad_h[3].y); - curr_idx += WARP_SIZE; - } - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) sum += g.shfl_xor(sum, i); - - curr_idx = threadIdx.x; - while (curr_idx < softmax_length) { - float4 out_reg = output_cast[curr_idx]; - float4 grad_reg = grad_cast[curr_idx]; - __half* grad_h = reinterpret_cast<__half*>(&grad_reg); - __half* out_h = reinterpret_cast<__half*>(&out_reg); - -#pragma unroll - for (int m = 0; m < 8; m++) grad_h[m] = (float)out_h[m] * ((float)grad_h[m] - sum); - - grad_cast[curr_idx] = grad_reg; - curr_idx += WARP_SIZE; - } -} - -__global__ void softmax_backward_kernel_arbitrary_length(float* grad /* input & output*/, - const float* output, - int softmax_length) -{ - int batch_idx = blockIdx.x * blockDim.y + threadIdx.y; - int offset = batch_idx * softmax_length + threadIdx.x; - - const float4* output_cast = reinterpret_cast(output); - float4* grad_cast = reinterpret_cast(grad); - - grad_cast += offset; - output_cast += offset; - - float sum = 0.0; - int curr_idx = threadIdx.x; - while (curr_idx < softmax_length) { - float4 out_reg = output_cast[curr_idx]; - float4 grad_reg = grad_cast[curr_idx]; - - grad_reg.x *= out_reg.x; - grad_reg.y *= out_reg.y; - grad_reg.z *= out_reg.z; - grad_reg.w *= out_reg.w; - sum += (grad_reg.x + grad_reg.y + grad_reg.z + grad_reg.w); - - curr_idx += WARP_SIZE; - } - - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - -#pragma unroll - for (int i = 1; i < WARP_SIZE; i <<= 1) sum += g.shfl_xor(sum, i); - - curr_idx = threadIdx.x; - while (curr_idx < softmax_length) { - float4 out_reg = output_cast[curr_idx]; - float4 grad_reg = grad_cast[curr_idx]; - grad_reg.x = out_reg.x * (grad_reg.x - sum); - grad_reg.y = out_reg.y * (grad_reg.y - sum); - grad_reg.z = out_reg.z * (grad_reg.z - sum); - grad_reg.w = out_reg.w * (grad_reg.w - sum); - - grad_cast[curr_idx] = grad_reg; - curr_idx += WARP_SIZE; - } -} - -template -void launch_attn_softmax_backward_v2(T* out_grad, - const T* soft_inp, - int batch_size, - int heads, - int seq_length, - cudaStream_t stream) -{ - const int warps_per_block = 4; - dim3 grid_dim(batch_size * heads * seq_length / warps_per_block); - dim3 block_dim(WARP_SIZE, warps_per_block); - - if (seq_length <= 32) - softmax_backward_kernel_v2 - <<>>(out_grad, soft_inp, seq_length); - else if (seq_length <= 64) - softmax_backward_kernel_v2 - <<>>(out_grad, soft_inp, seq_length); - else if (seq_length <= 128) - softmax_backward_kernel_v2 - <<>>(out_grad, soft_inp, seq_length); - else if (seq_length <= 256) - softmax_backward_kernel_v2 - <<>>(out_grad, soft_inp, seq_length); - else if (seq_length <= 384) - softmax_backward_kernel_v2 - <<>>(out_grad, soft_inp, seq_length); - else if (seq_length <= 512) - softmax_backward_kernel_v2 - <<>>(out_grad, soft_inp, seq_length); - else if (seq_length <= 768) - softmax_backward_kernel_v2 - <<>>(out_grad, soft_inp, seq_length); - else if (seq_length <= 1024) - softmax_backward_kernel_v2 - <<>>(out_grad, soft_inp, seq_length); - else if (seq_length <= 2048) - softmax_backward_kernel_v2 - <<>>(out_grad, soft_inp, seq_length); - else if (seq_length <= 4096) - softmax_backward_kernel_v2 - <<>>(out_grad, soft_inp, seq_length); - else if (seq_length <= 8192) - softmax_backward_kernel_v2 - <<>>(out_grad, soft_inp, seq_length); - else - softmax_backward_kernel_arbitrary_length<<>>( - out_grad, soft_inp, seq_length / (4 << ((sizeof(T) & 2) >> 1))); -} - -template void launch_attn_softmax_backward_v2<__half>(__half* out_grad, - const __half* soft_inp, - int batch_size, - int heads, - int seq_length, - cudaStream_t stream); -template void launch_attn_softmax_backward_v2(float* out_grad, - const float* soft_inp, - int batch_size, - int heads, - int seq_length, - cudaStream_t stream); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/transform_kernels.cu b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/transform_kernels.cu deleted file mode 100644 index dbecc31536c6c9355c4a3e243f2a0219bc476e6e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/transformer/transform_kernels.cu +++ /dev/null @@ -1,580 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "custom_cuda_layers.h" - -#define rows_trans 16 -#define cols_trans 16 - -template -__global__ void Transpose_Kernel(const T* inp, T* out, int row_width, int col_width) -{ - __shared__ T data_block[rows_trans * (cols_trans + 1)]; - - int r = threadIdx.x / cols_trans; - int c = threadIdx.x % cols_trans; - - int m = row_width / cols_trans; - - int i = blockIdx.x / m * rows_trans + r; - int j = blockIdx.x % m * cols_trans + c; - - int row_stride = rows_trans / ((rows_trans * cols_trans + THREADS - 1) / THREADS); - - for (int k = 0; k < rows_trans; k += row_stride) - data_block[(k + r) * cols_trans + c] = inp[(i + k) * row_width + j]; - - __syncthreads(); - - i = blockIdx.x % m * rows_trans + r; - j = blockIdx.x / m * cols_trans + c; - - for (int k = 0; k < rows_trans; k += row_stride) - out[(i + k) * col_width + j] = data_block[c * cols_trans + r + k]; -} - -template <> -void Transpose<__half>(const __half* inp_mat, - __half* out_mat, - int rows, - int cols, - cudaStream_t stream) -{ - int threads = THREADS; - - Transpose_Kernel<__half><<<(rows * cols + threads - 1) / threads, threads, 0, stream>>>( - inp_mat, out_mat, cols, rows); -} - -template <> -void Transpose(const float* inp_mat, float* out_mat, int rows, int cols, cudaStream_t stream) -{ - int threads = THREADS; - - Transpose_Kernel<<<(rows * cols + threads - 1) / threads, threads, 0, stream>>>( - inp_mat, out_mat, cols, rows); -} - -template -__global__ void transform_0213(T* output, - const T* vals, - int hidden_dim, - int seq_length, - int heads, - int head_ext); - -template <> -__global__ void transform_0213(float* output, - const float* vals, - int hidden_dim, - int seq_length, - int heads, - int head_ext) -{ - int d0_stride = hidden_dim * seq_length; - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - - int d0_out_stride = d0_stride; - int d1_out_stride = d2_stride; - int d2_out_stride = d2_stride * seq_length; - - int d0 = blockIdx.x; // Batch - int d1 = blockIdx.y / head_ext; // Sequence ID (0-127) - int d2 = threadIdx.y + (blockIdx.y % head_ext) * (heads / head_ext); // Head (0-11) - int d3 = threadIdx.x; // Values (groups of 4) - - const float4* vals_vec = reinterpret_cast(vals); - float4* output_vec = reinterpret_cast(output); - - float4 inputs = vals_vec[d0 * d0_stride + d1 * d1_stride + d2 * d2_stride + d3]; - output_vec[d0 * d0_out_stride + d1 * d1_out_stride + d2 * d2_out_stride + d3] = inputs; -} - -template <> -__global__ void transform_0213<__half>(__half* output, - const __half* vals, - int hidden_dim, - int seq_length, - int heads, - int head_ext) -{ -#ifdef HALF_PRECISION_AVAILABLE - - int d0_stride = hidden_dim * seq_length; - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - - int d0_out_stride = d0_stride; - int d1_out_stride = d2_stride; - int d2_out_stride = d2_stride * seq_length; - - int d0 = blockIdx.x; // Batch - int d1 = blockIdx.y / head_ext; // Sequence ID (0-127) - int d2 = threadIdx.y + (blockIdx.y % head_ext) * (heads / head_ext); // Head (0-11) - int d3 = threadIdx.x; // Values (groups of 4) - - float4 vals_arr[1]; - - const float4* vals_vec = reinterpret_cast(vals); - float4* output_vec = reinterpret_cast(output); - - vals_arr[0] = vals_vec[d0 * d0_stride + d1 * d1_stride + d2 * d2_stride + d3]; - output_vec[d0 * d0_out_stride + d1 * d1_out_stride + d2 * d2_out_stride + d3] = vals_arr[0]; -#endif -} - -template <> -void launch_transform_0213(float* output, - const float* vals, - int batch_size, - int seq_length, - int hidden_dim, - int heads, - cudaStream_t stream) -{ - hidden_dim >>= 2; - int head_ext = (hidden_dim - 1) / MAX_THREADS + 1; - dim3 block_dim(hidden_dim / heads, (heads / head_ext)); - dim3 grid_dim(batch_size, (seq_length * head_ext)); - - transform_0213 - <<>>(output, vals, hidden_dim, seq_length, heads, head_ext); -} - -template <> -void launch_transform_0213<__half>(__half* output, - const __half* vals, - int batch_size, - int seq_length, - int hidden_dim, - int heads, - cudaStream_t stream) -{ - hidden_dim >>= 3; - int head_ext = (hidden_dim - 1) / MAX_THREADS + 1; - dim3 block_dim(hidden_dim / heads, (heads / head_ext)); - dim3 grid_dim(batch_size, (seq_length * head_ext)); - transform_0213<__half> - <<>>(output, vals, hidden_dim, seq_length, heads, head_ext); -} - -// Bias add -template -__global__ void bias_add_transform_0213(T* output, - const T* vals, - const T* bias, - int hidden_dim, - int seq_length, - int heads, - int head_ext); - -template <> -__global__ void bias_add_transform_0213(float* output, - const float* vals, - const float* bias, - int hidden_dim, - int seq_length, - int heads, - int head_ext) -{ - int d0_stride = hidden_dim * seq_length; - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - - int d0_out_stride = d0_stride; - int d1_out_stride = d2_stride; - int d2_out_stride = d2_stride * seq_length; - - int d0 = blockIdx.x; // Batch - int d1 = blockIdx.y; // Sequence ID (0-127) - int cnt = blockIdx.z / head_ext; // Hidden count - int d2 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head (0-11) - int d3 = threadIdx.x; // Values (groups of 4) - - const float4* vals_vec = reinterpret_cast(vals); - const float4* bias_vec = reinterpret_cast(bias); - float4* output_vec = reinterpret_cast(output); - - float4 inputs = vals_vec[d0 * d0_stride * (gridDim.z / head_ext) + cnt * d1_stride + - d1 * d1_stride * (gridDim.z / head_ext) + d2 * d2_stride + d3]; - float4 biases = bias_vec[cnt * d1_stride + d2 * d2_stride + d3]; - - float4 outputs; - outputs.x = inputs.x + biases.x; - outputs.y = inputs.y + biases.y; - outputs.z = inputs.z + biases.z; - outputs.w = inputs.w + biases.w; - - output_vec[cnt * d0_out_stride * gridDim.x + d0 * d0_out_stride + d1 * d1_out_stride + - d2 * d2_out_stride + d3] = outputs; -} - -#define ATTN_H 3 -#define MAX_SEQ_LINE 10 - -template <> -__global__ void bias_add_transform_0213<__half>(__half* output, - const __half* vals, - const __half* bias, - int hidden_dim, - int seq_length, - int heads, - int head_ext) -{ -#ifdef HALF_PRECISION_AVAILABLE - - int d0_stride = hidden_dim * seq_length; - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - - int d2_out_stride = d2_stride * seq_length; - - int d0 = blockIdx.x; // Batch - int d1 = blockIdx.y; // Sequence ID (0-127) - int cnt = blockIdx.z / head_ext; // Hidden count - int d2 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head (0-11) - int d3 = threadIdx.x; // Values (groups of 4) - - float4 vals_arr; - float4 bias_arr; - float4 output_arr; - __half2* vals_half = reinterpret_cast<__half2*>(&vals_arr); - __half2* bias_half = reinterpret_cast<__half2*>(&bias_arr); - __half2* output_half = reinterpret_cast<__half2*>(&output_arr); - - const float4* vals_vec = reinterpret_cast(vals); - const float4* bias_vec = reinterpret_cast(bias); - float4* output_vec = reinterpret_cast(output); - - vals_vec += (d0 * d0_stride * (gridDim.z / head_ext)); - vals_vec += (d1 * d1_stride * (gridDim.z / head_ext)); - vals_vec += (cnt * d1_stride); - vals_vec += (d2 * d2_stride); - - bias_vec += (cnt * d1_stride); - bias_vec += (d2 * d2_stride); - - output_vec += (cnt * d0_stride * gridDim.x); - output_vec += (d1 * d2_stride); - output_vec += (d0 * d0_stride); - output_vec += (d2 * d2_out_stride); - - bias_arr = bias_vec[d3]; - vals_arr = vals_vec[d3]; - -#if defined(__ACC_HALF__) - output_half[0] = vals_half[0] + bias_half[0]; - output_half[1] = vals_half[1] + bias_half[1]; - output_half[2] = vals_half[2] + bias_half[2]; - output_half[3] = vals_half[3] + bias_half[3]; -#else - float2 bias_arr_f[4]; - float2 vals_arr_f[4]; -#pragma unroll - for (int l = 0; l < 4; l++) { - bias_arr_f[l] = __half22float2(bias_half[l]); - vals_arr_f[l] = __half22float2(vals_half[l]); - vals_arr_f[l].x += bias_arr_f[l].x; - vals_arr_f[l].y += bias_arr_f[l].y; - output_half[l] = __float22half2_rn(vals_arr_f[l]); - } -#endif - output_vec[d3] = output_arr; - -#endif -} - -__global__ void bias_add_transform_0213_v2(__half* output, - const __half* vals, - const __half* bias, - int hidden_dim, - int seq_length, - int heads) -{ -#ifdef HALF_PRECISION_AVAILABLE - __shared__ float4 in_data[3072]; - - int d0_stride = hidden_dim * seq_length; - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - int iteration_stride = d1_stride * blockDim.z; // Hidden * 3 / 8 - int batch_stride = d0_stride * blockDim.z; // Hidden * S * 3 / 8 - - int d0_out_stride = d0_stride; - int d1_out_stride = d2_stride; - int d2_out_stride = d2_stride * seq_length; - - int d0 = blockIdx.x; // Batch - int d1 = blockIdx.y; // Sequence ID (0-127) - int cnt = threadIdx.z; // blockIdx.z; // Hidden count - int d2 = threadIdx.y; // Head (0-11) - int d3 = threadIdx.x; // Values (groups of 4) - - float4 vals_arr[1]; - float4 bias_arr[1]; - float4 output_arr[1]; - __half2* vals_half = reinterpret_cast<__half2*>(vals_arr); - __half2* bias_half = reinterpret_cast<__half2*>(bias_arr); - __half2* output_half = reinterpret_cast<__half2*>(output_arr); - - const float4* vals_vec = reinterpret_cast(vals); - const float4* bias_vec = reinterpret_cast(bias); - float4* output_vec = reinterpret_cast(output); - - int iter_index = cnt * d1_stride + d2 * d2_stride + d3; - int input_offset = d0 * batch_stride + d1 * (iteration_stride << 1); - bias_arr[0] = bias_vec[iter_index]; - -#pragma unroll - for (int iter = 0; iter < 2; iter++) { - int iter_id = iter * iteration_stride + iter_index; - vals_arr[0] = vals_vec[input_offset + iter_id]; - - output_half[0] = vals_half[0] + bias_half[0]; - output_half[1] = vals_half[1] + bias_half[1]; - output_half[2] = vals_half[2] + bias_half[2]; - output_half[3] = vals_half[3] + bias_half[3]; - - in_data[iter_id] = output_arr[0]; - } - __syncthreads(); - - iteration_stride = blockDim.z * (blockDim.y >> 1); - int matrix_stride = (d0_out_stride * gridDim.x); - int head_count = (d2 >> 1) + cnt * (blockDim.y >> 1); - - int out_index = d0 * d0_out_stride + d1 * (d1_out_stride << 1) + d3 + (d2 % 2) * d2_stride; - -#pragma unroll - for (int iter = 0; iter < 2; iter++) { - int iter_row = (iter * iteration_stride) + head_count; - int iter_offset = - (iter_row % blockDim.y) * d2_out_stride + (iter_row / blockDim.y) * matrix_stride; - output_vec[out_index + iter_offset] = - in_data[iter_row * d2_stride + d3 + (d2 % 2) * (d1_stride * blockDim.z)]; - } -#endif -} - -// [B S C*H] - > C * [B A S N] -template <> -void launch_bias_add_transform_0213(float* output, - const float* vals, - const float* bias, - int batch_size, - int seq_length, - int hidden_dim, - int heads, - cudaStream_t stream, - int trans_count) -{ - hidden_dim >>= 2; - int head_ext = (hidden_dim - 1) / MAX_THREADS + 1; - - dim3 block_dim(hidden_dim / heads, (heads / head_ext)); - dim3 grid_dim(batch_size, seq_length, (trans_count * head_ext)); - - bias_add_transform_0213<<>>( - output, vals, bias, hidden_dim, seq_length, heads, head_ext); -} - -template <> -void launch_bias_add_transform_0213<__half>(__half* output, - const __half* vals, - const __half* bias, - int batch_size, - int seq_length, - int hidden_dim, - int heads, - cudaStream_t stream, - int trans_count) -{ - hidden_dim >>= 3; - if (hidden_dim > 128 || hidden_dim < 16) { - int head_ext = (hidden_dim - 1) / MAX_THREADS + 1; - dim3 block_dim(hidden_dim / heads, (heads / head_ext)); - dim3 grid_dim(batch_size, seq_length, (trans_count * head_ext)); - bias_add_transform_0213<__half><<>>( - output, vals, bias, hidden_dim, seq_length, heads, head_ext); - } else { - dim3 block_dim(hidden_dim / heads, heads, trans_count); - dim3 grid_dim(batch_size, seq_length / 2); - bias_add_transform_0213_v2<<>>( - output, vals, bias, hidden_dim, seq_length, heads); - } -} - -template -__global__ void transform4d_0213(T* out, - const T* in, - int heads, - int seq_length, - int hidden_dim, - int head_ext); - -template <> -__global__ void transform4d_0213(float* out, - const float* in, - int heads, - int seq_length, - int hidden_dim, - int head_ext) -{ - int d0_stride = hidden_dim * seq_length; - int d1_stride = d0_stride / heads; - int d2_stride = hidden_dim / heads; - - int d0_out_stride = d0_stride; - int d1_out_stride = d2_stride; - int d2_out_stride = hidden_dim; - - int d0 = blockIdx.x; // Batch - int d1 = blockIdx.y / ((seq_length - 1) / blockDim.y + 1); // Head - int d2 = (threadIdx.y + blockDim.y * blockIdx.y) % seq_length; - int cnt = blockIdx.z; - int d3 = threadIdx.x; // Values (groups of 8) - - if (d2 < seq_length) { - const float4* in_vec = reinterpret_cast(in); - float4* out_vec = reinterpret_cast(out); - - float4 vals_vec = in_vec[cnt * d0_stride * gridDim.x + d0 * d0_stride + d1 * d1_stride + - d2 * d2_stride + d3]; - out_vec[d0 * d0_out_stride * gridDim.z + cnt * d2_out_stride + d1 * d1_out_stride + - d2 * d2_out_stride * gridDim.z + d3] = vals_vec; - } -} - -template <> -__global__ void transform4d_0213<__half>(__half* out, - const __half* in, - int heads, - int seq_length, - int hidden_dim, - int head_ext) -{ -#ifdef HALF_PRECISION_AVAILABLE - - int d0_stride = hidden_dim * (seq_length / head_ext); - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - - int d0 = blockIdx.x; // Batch - int d1 = threadIdx.y + (blockIdx.z % head_ext) * (heads / head_ext); // Head - int d2 = blockIdx.z / head_ext; // Sequence - int cnt = blockIdx.y; // Hidden count - int d3 = threadIdx.x; // Values (groups of 8) - - const float4* in_vec = reinterpret_cast(in); - float4* out_vec = reinterpret_cast(out); - - in_vec += (cnt * d0_stride * gridDim.x); - in_vec += (d0 * d0_stride); - in_vec += (d2 * d2_stride); - in_vec += (d1 * d2_stride * seq_length); - - out_vec += (cnt * d1_stride); - out_vec += (d1 * d2_stride); - out_vec += (d0 * d0_stride * gridDim.y); - out_vec += (d2 * d1_stride * gridDim.y); - - out_vec[d3] = in_vec[d3]; - -#endif -} - -__global__ void transform4d_0213_v2(__half* out, - const __half* in, - int heads, - int seq_length, - int hidden_dim) -{ -#ifdef HALF_PRECISION_AVAILABLE - __shared__ float4 in_data[3072]; - - int d0_stride = hidden_dim * seq_length; - int d1_stride = hidden_dim; - int d2_stride = hidden_dim / heads; - - int d0 = blockIdx.x; // Batch - int d1 = threadIdx.y; // Head - int d2 = blockIdx.y; // Sequence - int cnt = threadIdx.z; // Hidden count - int d3 = threadIdx.x; // Values (groups of 8) - - const float4* in_vec = reinterpret_cast(in); - float4* out_vec = reinterpret_cast(out); - - int input_offset = d0 * d0_stride + d2 * (d2_stride << 1) + d3 + (d1 % 2) * d2_stride; - int head_count = (d1 >> 1) + cnt * (blockDim.y >> 1); - int iteration_stride = blockDim.z * (blockDim.y >> 1); - int matrix_stride = (d0_stride * gridDim.x); - -#pragma unroll - for (int iter = 0; iter < 2; iter++) { - int iter_row = iter * iteration_stride + head_count; - int iter_offset = (iter_row % blockDim.y) * d2_stride; - - in_data[d3 + iter_offset + (iter_row / blockDim.y + (d1 % 2) * blockDim.z) * d1_stride] = - in_vec[input_offset + iter_offset * seq_length + - (iter_row / blockDim.y) * matrix_stride]; - } - __syncthreads(); - - iteration_stride = d1_stride * blockDim.z; - int iter_index = cnt * d1_stride + d1 * d2_stride + d3; - int output_offset = d0 * d0_stride * blockDim.z + d2 * (iteration_stride << 1); - -#pragma unroll - for (int iter = 0; iter < 2; iter++) { - int iter_id = iter * iteration_stride + iter_index; - out_vec[output_offset + iter_id] = in_data[iter_id]; - } -#endif -} - -// 3 * [B A S N] - > [B S C*H] -template <> -void launch_transform4d_0213(float* out, - const float* in, - int batch_size, - int heads, - int seq_length, - int hidden_dim, - cudaStream_t stream, - int trans_count) -{ - hidden_dim >>= 2; - dim3 grid_dims(batch_size, heads * ((seq_length - 1) / 8 + 1), trans_count); - dim3 block_dims(hidden_dim / heads, 8); - transform4d_0213 - <<>>(out, in, heads, seq_length, hidden_dim, 1); -} - -template <> -void launch_transform4d_0213<__half>(__half* out, - const __half* in, - int batch_size, - int heads, - int seq_length, - int hidden_dim, - cudaStream_t stream, - int trans_count) -{ - hidden_dim >>= 3; - if (hidden_dim > 128 || hidden_dim < 16) { - int head_ext = (hidden_dim - 1) / MAX_THREADS + 1; - dim3 grid_dims(batch_size, trans_count, (seq_length * head_ext)); - dim3 block_dims(hidden_dim / heads, (heads / head_ext)); - transform4d_0213<__half><<>>( - out, in, heads, seq_length, hidden_dim, head_ext); - } else { - dim3 grid_dims(batch_size, seq_length / 2); - dim3 block_dims(hidden_dim / heads, heads, trans_count); - transform4d_0213_v2<<>>( - out, in, heads, seq_length, hidden_dim); - } -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/py_ds_utils.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/py_ds_utils.cpp deleted file mode 100644 index df5c9f361c616f07fd13640804200dca25d71191..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/py_ds_utils.cpp +++ /dev/null @@ -1,25 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* - Collection of system utilities. -*/ - -#include -#include "tensor_cast.h" -using namespace pybind11::literals; - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("cast_to_byte_tensor", - py::overload_cast(&cast_to_byte_tensor), - "Cast a 1-dimensional tensor of any type to byte tensor.", - "src_tensor"_a); - - m.def("cast_to_byte_tensor", - py::overload_cast&>(&cast_to_byte_tensor), - "Cast a multi-dimensional tensor of any type to byte tensor.", - "src_tensor"_a); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/tensor_cast.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/tensor_cast.cpp deleted file mode 100644 index 8352bc72bb0286e30268e77f243787c2923a0d40..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/tensor_cast.cpp +++ /dev/null @@ -1,26 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "tensor_cast.h" - -at::Tensor cast_to_byte_tensor(at::Tensor& src_tensor) -{ - if (src_tensor.nbytes() <= 1) return src_tensor; - - auto options = torch::TensorOptions() - .dtype(torch::kUInt8) - .layout(src_tensor.layout()) - .device(src_tensor.device()); - return at::from_blob( - src_tensor.data_ptr(), static_cast(src_tensor.nbytes()), options); -} - -std::vector cast_to_byte_tensor(std::vector& tensor_list) -{ - std::vector byte_tensors; - for (auto src_tensor : tensor_list) { byte_tensors.push_back(cast_to_byte_tensor(src_tensor)); } - - return byte_tensors; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/tensor_cast.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/tensor_cast.h deleted file mode 100644 index 86155567122d592c58979606a21aa9c3f00d5511..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/utils/tensor_cast.h +++ /dev/null @@ -1,16 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* - Utilities for type casting torch tensors without data movement. -*/ - -#include -#include - -using namespace std; -at::Tensor cast_to_byte_tensor(at::Tensor& src_tensor); - -std::vector cast_to_byte_tensor(std::vector& tensor_list); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adagrad/cpu_adagrad.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adagrad/cpu_adagrad.cpp deleted file mode 100644 index dc727f8fa2168ef7897f685de6ac5d80aaad8309..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adagrad/cpu_adagrad.cpp +++ /dev/null @@ -1,196 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include "cpu_adagrad.h" -#include -#include -#include -#include -#include -#include - -static std::unordered_map> s_optimizers; - -// C++ interface - -void Adagrad_Optimizer::Step_1(float* _params, - float* grads, - float* _exp_avg_sq, - size_t _param_size, - ds_half_precision_t* dev_params, - bool half_precision) -{ - size_t rounded_size = 0; -#if defined(__AVX512__) or defined(__AVX256__) - Step_AVX<1>( - &rounded_size, _params, grads, _exp_avg_sq, _param_size, dev_params, half_precision); -#endif - if (_param_size > rounded_size) { - float step_size = -1 * _alpha; - ds_half_precision_t* grads_cast_h; - ds_half_precision_t* params_cast_h; - if (half_precision) { - grads_cast_h = reinterpret_cast(grads); - params_cast_h = reinterpret_cast(_params); - } - for (size_t t = rounded_size; t < _param_size; t += TILE) { - size_t copy_size = TILE; - if ((t + TILE) > _param_size) copy_size = _param_size - t; - size_t offset = copy_size + t; -#pragma omp parallel for - for (size_t k = t; k < offset; k++) { - float grad = half_precision ? (float)grads_cast_h[k] : grads[k]; - float param = half_precision ? (float)params_cast_h[k] : _params[k]; - float momentum = grads[k]; - float variance = _exp_avg_sq[k]; - if (_weight_decay > 0) { grad = param * _weight_decay + grad; } - - variance += grad * grad; - - grad = sqrt(variance); - grad += _eps; - grad = momentum / grad; - param = grad * step_size + param; - if (half_precision) - params_cast_h[k] = (ds_half_precision_t)param; - else - _params[k] = param; - // STORE UPDATE TERM TO GRAD'S MEMORY - grads[k] = grad * step_size; - _exp_avg_sq[k] = variance; - } - } - } -} - -void Adagrad_Optimizer::Step_4(float* _params, - float* grads, - float* _exp_avg_sq, - size_t _param_size, - ds_half_precision_t* dev_params, - bool half_precision) -{ - size_t rounded_size = 0; -#if defined(__AVX512__) or defined(__AVX256__) - Step_AVX<4>( - &rounded_size, _params, grads, _exp_avg_sq, _param_size, dev_params, half_precision); -#endif - if (_param_size > rounded_size) - Step_1((_params + rounded_size), - (grads + rounded_size), - (_exp_avg_sq + rounded_size), - (_param_size - rounded_size), - (dev_params != nullptr ? (dev_params + rounded_size) : dev_params), - half_precision); -} - -int create_adagrad_optimizer(int optimizer_id, - float alpha = 1e-2, - float eps = 1e-8, - float weight_decay = 0, - bool should_log = false) -{ - auto opt = std::make_shared(alpha, eps, weight_decay); - - s_optimizers[optimizer_id] = opt; - - if (should_log) { - std::string avx_type = ""; -#if defined(__AVX512__) - avx_type = "AVX512"; -#else -#if defined(__AVX256__) - avx_type = "AVX2"; -#else - avx_type = "scalar"; -#endif -#endif - - printf("Adagrad Optimizer #%d is created with %s arithmetic capability.\n", - optimizer_id, - avx_type.c_str()); - printf("Config: alpha=%f, weight_decay=%f\n", alpha, weight_decay); - } - - return 0; -} - -void Adagrad_Optimizer::Step_8(float* _params, - float* grads, - float* _exp_avg_sq, - size_t _param_size, - ds_half_precision_t* dev_params, - bool half_precision) -{ - size_t rounded_size = 0; -#if defined(__AVX512__) or defined(__AVX256__) - Step_AVX<8>( - &rounded_size, _params, grads, _exp_avg_sq, _param_size, dev_params, half_precision); -#endif - if (_param_size > rounded_size) - Step_4((_params + rounded_size), - (grads + rounded_size), - (_exp_avg_sq + rounded_size), - (_param_size - rounded_size), - (dev_params != nullptr ? (dev_params + rounded_size) : dev_params), - half_precision); -} - -int ds_adagrad_step(int optimizer_id, - size_t step, - float lr, - float epsilon, - float weight_decay, - torch::Tensor& params, - torch::Tensor& grads, - torch::Tensor& exp_avg_sq) -{ - auto params_c = params.contiguous(); - auto grads_c = grads.contiguous(); - auto exp_avg_sq_c = exp_avg_sq.contiguous(); - - float* params_ptr = (float*)params_c.data_ptr(); - float* grads_ptr = (float*)grads_c.data_ptr(); - float* exp_avg_sq_ptr = (float*)exp_avg_sq_c.data_ptr(); - - std::shared_ptr opt = - std::static_pointer_cast(s_optimizers[optimizer_id]); - opt->IncrementStep(step); - opt->update_state(lr, epsilon, weight_decay); - opt->Step_8(params_ptr, grads_ptr, exp_avg_sq_ptr, params_c.numel()); - - return 0; -} - -int ds_adagrad_step_plus_copy(int optimizer_id, - size_t step, - float lr, - float epsilon, - float weight_decay, - torch::Tensor& params, - torch::Tensor& grads, - torch::Tensor& exp_avg_sq, - torch::Tensor& gpu_params) -{ - assert(false); - return 0; -} - -int destroy_adagrad_optimizer(int optimizer_id) -{ - s_optimizers.erase(optimizer_id); - - return 0; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("adagrad_update", &ds_adagrad_step, "DeepSpeed CPU Adagrad update (C++)"); - m.def("adagrad_update_copy", - &ds_adagrad_step_plus_copy, - "DeepSpeed CPU Adagrad update and param copy (C++)"); - m.def("create_adagrad", &create_adagrad_optimizer, "DeepSpeed CPU Adagrad (C++)"); - m.def("destroy_adagrad", &destroy_adagrad_optimizer, "DeepSpeed CPU Adagrad destroy (C++)"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/fused_adam_frontend.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/fused_adam_frontend.cpp deleted file mode 100644 index 13b390248608b046dab443f85346b5446a47d722..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/fused_adam_frontend.cpp +++ /dev/null @@ -1,25 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include - -void multi_tensor_adam_cuda(int chunk_size, - at::Tensor noop_flag, - std::vector> tensor_lists, - const float lr, - const float beta1, - const float beta2, - const float epsilon, - const int step, - const int mode, - const int bias_correction, - const float weight_decay); - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("multi_tensor_adam", - &multi_tensor_adam_cuda, - "Compute and apply gradient update to parameters for Adam optimizer"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_adam.dp.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_adam.dp.cpp deleted file mode 100644 index 0720a020247a3ebf80fc231e11333ea56cb09924..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_adam.dp.cpp +++ /dev/null @@ -1,159 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Copyright NVIDIA/apex -This file is adapted from fused adam in NVIDIA/apex, commit a109f85 -*/ - -#include -#include -#include - -#include - -#include -#include "multi_tensor_apply.dp.hpp" -#include "type_shim.h" - -#define BLOCK_SIZE 512 -#define ILP 4 - -typedef enum : int { - ADAM_MODE_0 = 0, // L2 regularization mode - ADAM_MODE_1 = 1 // Decoupled weight decay mode(AdamW) -} adamMode_t; - -using MATH_T = float; - -template -struct AdamFunctor { - __inline__ __attribute__((always_inline)) void operator()(int chunk_size, - volatile int* noop_gmem, - TensorListMetadata<4>& tl, - const float beta1, - const float beta2, - const float beta1_correction, - const float beta2_correction, - const float epsilon, - const float lr, - adamMode_t mode, - const float decay) - { - auto item_ct1 = sycl::ext::oneapi::experimental::this_nd_item<3>(); - int tensor_loc = tl.block_to_tensor[item_ct1.get_group(2)]; - - int chunk_idx = tl.block_to_chunk[item_ct1.get_group(2)]; - int n = tl.sizes[tensor_loc]; - - T* g = (T*)tl.addresses[0][tensor_loc]; - g += chunk_idx * chunk_size; - - T* p = (T*)tl.addresses[1][tensor_loc]; - p += chunk_idx * chunk_size; - - T* m = (T*)tl.addresses[2][tensor_loc]; - m += chunk_idx * chunk_size; - - T* v = (T*)tl.addresses[3][tensor_loc]; - v += chunk_idx * chunk_size; - - n -= chunk_idx * chunk_size; - - // see note in multi_tensor_scale_kernel.cu - for (int i_start = 0; i_start < n && i_start < chunk_size; - i_start += item_ct1.get_local_range(2) * ILP) { - MATH_T r_g[ILP]; - MATH_T r_p[ILP]; - MATH_T r_m[ILP]; - MATH_T r_v[ILP]; -#pragma unroll - for (int ii = 0; ii < ILP; ii++) { - int i = i_start + item_ct1.get_local_id(2) + ii * item_ct1.get_local_range(2); - if (i < n && i < chunk_size) { - r_g[ii] = g[i]; - r_p[ii] = p[i]; - r_m[ii] = m[i]; - r_v[ii] = v[i]; - } else { - r_g[ii] = MATH_T(0); - r_p[ii] = MATH_T(0); - r_m[ii] = MATH_T(0); - r_v[ii] = MATH_T(0); - } - } -#pragma unroll - for (int ii = 0; ii < ILP; ii++) { - if (mode == ADAM_MODE_0) { // L2 - r_g[ii] = r_g[ii] + (decay * r_p[ii]); - r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; - r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii]; - MATH_T next_m_unbiased = r_m[ii] / beta1_correction; - MATH_T next_v_unbiased = r_v[ii] / beta2_correction; - MATH_T denom = sycl::sqrt(next_v_unbiased) + epsilon; - MATH_T update = next_m_unbiased / denom; - r_p[ii] = r_p[ii] - (lr * update); - } else { // weight decay - r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; - r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii]; - MATH_T next_m_unbiased = r_m[ii] / beta1_correction; - MATH_T next_v_unbiased = r_v[ii] / beta2_correction; - MATH_T denom = sycl::sqrt(next_v_unbiased) + epsilon; - MATH_T update = (next_m_unbiased / denom) + (decay * r_p[ii]); - r_p[ii] = r_p[ii] - (lr * update); - } - } -#pragma unroll - for (int ii = 0; ii < ILP; ii++) { - int i = i_start + item_ct1.get_local_id(2) + ii * item_ct1.get_local_range(2); - if (i < n && i < chunk_size) { - p[i] = r_p[ii]; - m[i] = r_m[ii]; - v[i] = r_v[ii]; - } - } - } - } -}; - -void multi_tensor_adam_cuda(int chunk_size, - at::Tensor noop_flag, - std::vector> tensor_lists, - const float lr, - const float beta1, - const float beta2, - const float epsilon, - const int step, - const int mode, - const int bias_correction, - const float weight_decay) -{ - using namespace at; - - // Handle bias correction mode - float bias_correction1 = 1.0f, bias_correction2 = 1.0f; - if (bias_correction == 1) { - bias_correction1 = 1 - std::pow(beta1, step); - bias_correction2 = 1 - std::pow(beta2, step); - } - - // Assume single type across p,g,m1,m2 now - DISPATCH_DOUBLE_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), - 0, - "adam", - multi_tensor_apply<4>(BLOCK_SIZE, - chunk_size, - noop_flag, - tensor_lists, - AdamFunctor(), - beta1, - beta2, - bias_correction1, - bias_correction2, - epsilon, - lr, - (adamMode_t)mode, - weight_decay);) -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_apply.dp.hpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_apply.dp.hpp deleted file mode 100644 index 14a130e2a23e906ba9b305ee3968aa4e27c17cf8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_apply.dp.hpp +++ /dev/null @@ -1,226 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Copyright NVIDIA/apex -This file is adapted from fused adam in NVIDIA/apex, commit a109f85 -*/ - -#include -#include -#include -#include -#include -#include "compat.h" - -#include -#include -#include - -namespace at { -namespace cuda { -sycl::queue* getCurrentCUDAStream() -{ - c10::xpu::XPUStream stream = c10::xpu::getCurrentXPUStream(); - auto& queue = stream.queue(); - return &queue; -} - -sycl::queue* getStreamFromPool(bool) -{ - // not implemented - return nullptr; -} -} // namespace cuda -} // namespace at -// #include - -// This header is the one-stop shop for all your multi-tensor apply needs. - -// TODO: Kernel arg size limit may be <4KB for some other cards (ie Jetson) -constexpr int depth_to_max_tensors[5] = {110, 64, 48, 36, 30}; -constexpr int depth_to_max_blocks[5] = {320, 320, 320, 320, 320}; - -template -struct TensorListMetadata { - void* addresses[n][depth_to_max_tensors[n - 1]]; - int sizes[depth_to_max_tensors[n - 1]]; - unsigned char block_to_tensor[depth_to_max_blocks[n - 1]]; - int block_to_chunk[depth_to_max_blocks[n - 1]]; // I fear this needs to be a full int. - int start_tensor_this_launch; -}; - -template -class multi_tensor_apply_kernel { -public: - multi_tensor_apply_kernel(int chunk_size, - volatile int* noop_flag, - T tl, - U callable, - ArgTypes... args) - : chunk_size(chunk_size), noop_flag(noop_flag), tl(tl), callable(callable), args(args...) - { - } - - // This should be identical to original __global__ function - static void inline __global__function(int chunk_size, - volatile int* noop_flag, - T tl, - U callable, - ArgTypes... args) - { - callable(chunk_size, noop_flag, tl, args...); - } - - // If global function template contains parameter pack, - // we only deal with parameter pack at the end of template parameter list - template - static void inline __tuple_expand_driver(int chunk_size, - volatile int* noop_flag, - T tl, - U callable, - Tuple args, - std::index_sequence) - { - __global__function(chunk_size, noop_flag, tl, callable, std::get(args)...); - } - - // - // Because __global__ function can't really use any reference types, we can sure that args - // are all good behaviors - // - void operator()(sycl::nd_item<3>) const - { - __tuple_expand_driver(chunk_size, - noop_flag, - tl, - callable, - args, - std::make_index_sequence()); - } - -private: - int chunk_size; - volatile int* noop_flag; - T tl; - U callable; - std::tuple args; -}; - -// to make sure multi_tensor_apply_kernel can be used in sycl::buffer -namespace sycl { -template -struct is_device_copyable> : std::true_type {}; -} // namespace sycl - -template -void multi_tensor_apply(int block_size, - int chunk_size, - const at::Tensor& noop_flag, - const std::vector>& tensor_lists, - T callable, - ArgTypes... args) -{ - TORCH_CHECK(tensor_lists.size() == depth, "tensor_lists.size() != depth"); - int len0 = tensor_lists[0].size(); - TORCH_CHECK(len0 > 0, "tensor_lists[0].size() is not > 0"); - auto ref_device = tensor_lists[0][0].device(); - TORCH_CHECK(ref_device.type() == at::kXPU, "expected input to be on cuda"); - for (int l = 0; l < tensor_lists.size(); l++) // No range-based for because I need indices - { - TORCH_CHECK(tensor_lists[l].size() == len0, "Size mismatch among tensor lists"); - for (int t = 0; t < tensor_lists[l].size(); t++) { - // TODO: Print which tensor fails. - bool contiguous_memory = tensor_lists[l][t].is_contiguous(); -#ifdef VERSION_GE_1_5 - contiguous_memory = (contiguous_memory || - tensor_lists[l][t].is_contiguous(at::MemoryFormat::ChannelsLast)); -#endif - TORCH_CHECK(contiguous_memory, "A tensor was not contiguous."); - TORCH_CHECK(tensor_lists[l][t].device() == ref_device, - "A tensor was not on the same device as the first tensor"); - TORCH_CHECK(tensor_lists[l][t].numel() == tensor_lists[0][t].numel(), "Size mismatch"); - } - } - - int ntensors = tensor_lists[0].size(); - - TensorListMetadata tl; - - /* const at::cuda::OptionalCUDAGuard device_guard(device_of(tensor_lists[0][0])); */ - auto stream = at::cuda::getCurrentCUDAStream(); - - tl.start_tensor_this_launch = 0; - int loc_block_info = 0; - int loc_tensor_info = 0; - for (int t = 0; t < ntensors; t++) { - tl.sizes[loc_tensor_info] = tensor_lists[0][t].numel(); - for (int d = 0; d < depth; d++) - tl.addresses[d][loc_tensor_info] = tensor_lists[d][t].data_ptr(); - loc_tensor_info++; - - int chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size; - - for (int chunk = 0; chunk < chunks_this_tensor; chunk++) { - // std::cout << chunks_this_tensor << std::endl; - tl.block_to_tensor[loc_block_info] = loc_tensor_info - 1; - tl.block_to_chunk[loc_block_info] = chunk; - loc_block_info++; - - bool tensors_full = (loc_tensor_info == depth_to_max_tensors[depth - 1] && - chunk == chunks_this_tensor - 1); - bool blocks_full = (loc_block_info == depth_to_max_blocks[depth - 1]); - bool last_chunk = (t == ntensors - 1 && chunk == chunks_this_tensor - 1); - if (tensors_full || blocks_full || last_chunk) { - // using accscalar_t = acc_type; - /* multi_tensor_apply_kernel, T, ArgTypes...> - * fn(chunk_size, noop_flag.DATA_PTR(), tl, callable, args...); */ - if constexpr (sizeof(multi_tensor_apply_kernel( - chunk_size, noop_flag.DATA_PTR(), tl, callable, args...)) < - 2048) { - ((sycl::queue*)(stream)) - ->parallel_for( - sycl::nd_range<3>(sycl::range<3>(1, 1, loc_block_info) * - sycl::range<3>(1, 1, block_size), - sycl::range<3>(1, 1, block_size)), - multi_tensor_apply_kernel( - chunk_size, noop_flag.DATA_PTR(), tl, callable, args...)); - } else { - auto capture = multi_tensor_apply_kernel( - chunk_size, noop_flag.DATA_PTR(), tl, callable, args...); - sycl::buffer params(const_cast(&capture), - sycl::range<1>(1)); - stream->submit([&](sycl::handler& cgh) { - auto device_params = - params.template get_access(cgh); - cgh.parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, loc_block_info) * - sycl::range<3>(1, 1, block_size), - sycl::range<3>(1, 1, block_size)), - [=](sycl::nd_item<3> item) { device_params[0](item); }); - }); - } - 0; - - // Reset. The control flow possibilities here make my brain hurt. - loc_block_info = 0; - if (chunk == chunks_this_tensor - 1) { - // std::cout << "Hit case 1 " << cond1 << " " << cond2 << " " << cond3 << - // std::endl; - loc_tensor_info = 0; - tl.start_tensor_this_launch = t + 1; - } else { - // std::cout << "Hit case 2 " << cond1 << " " << cond2 << " " << cond3 << - // std::endl; - tl.sizes[0] = tl.sizes[loc_tensor_info - 1]; - for (int d = 0; d < depth; d++) - tl.addresses[d][0] = tl.addresses[d][loc_tensor_info - 1]; - loc_tensor_info = 1; - tl.start_tensor_this_launch = t; - } - } - } - } -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/common/custom_cuda_kernel.dp.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/common/custom_cuda_kernel.dp.cpp deleted file mode 100644 index cfd004ef13577c02f958fa18d116131398443562..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/common/custom_cuda_kernel.dp.cpp +++ /dev/null @@ -1,92 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include - -inline void has_capability_or_fail(const sycl::device& dev, - const std::initializer_list& props) -{ - for (const auto& it : props) { - if (dev.has(it)) continue; - switch (it) { - case sycl::aspect::fp64: - throw std::runtime_error("'double' is not supported in '" + - dev.get_info() + "' device"); - break; - case sycl::aspect::fp16: - throw std::runtime_error("'half' is not supported in '" + - dev.get_info() + "' device"); - break; - default: -#define __SYCL_ASPECT(ASPECT, ID) \ - case sycl::aspect::ASPECT: return #ASPECT; -#define __SYCL_ASPECT_DEPRECATED(ASPECT, ID, MESSAGE) __SYCL_ASPECT(ASPECT, ID) -#define __SYCL_ASPECT_DEPRECATED_ALIAS(ASPECT, ID, MESSAGE) - auto getAspectNameStr = [](sycl::aspect AspectNum) -> std::string { - switch (AspectNum) { -#include -#include - default: return "unknown aspect"; - } - }; -#undef __SYCL_ASPECT_DEPRECATED_ALIAS -#undef __SYCL_ASPECT_DEPRECATED -#undef __SYCL_ASPECT - throw std::runtime_error("'" + getAspectNameStr(it) + "' is not supported in '" + - dev.get_info() + "' device"); - } - break; - } -} - -void param_update_kernel(const float* input, sycl::half* output, int size) -{ - auto item_ct1 = sycl::ext::oneapi::experimental::this_nd_item<3>(); - int id = item_ct1.get_group(2) * item_ct1.get_local_range(2) + item_ct1.get_local_id(2); - - if (id < size) { output[id] = (sycl::half)input[id]; } -} - -void launch_param_update(const float* input, sycl::half* output, int size, sycl::queue* stream) -{ - int threads = 1024; - - sycl::range<3> grid_dim(1, 1, (size - 1) / threads + 1); - sycl::range<3> block_dim(1, 1, threads); - - { - has_capability_or_fail(stream->get_device(), {sycl::aspect::fp16}); - stream->parallel_for( - sycl::nd_range<3>(grid_dim * block_dim, block_dim), - [=](sycl::nd_item<3> item_ct1) { param_update_kernel(input, output, size); }); - } -} - -void param_update_kernel_half(const float* input, sycl::half* output, int size) -{ - auto item_ct1 = sycl::ext::oneapi::experimental::this_nd_item<3>(); - int id = item_ct1.get_group(2) * item_ct1.get_local_range(2) + item_ct1.get_local_id(2); - sycl::half2* output_cast = reinterpret_cast(output); - if (id < size) { - float input_f = input[id]; - sycl::half2* input_h = reinterpret_cast(&input_f); - output_cast[id] = *input_h; - } -} - -void launch_param_update_half(const float* input, sycl::half* output, int size, sycl::queue* stream) -{ - int threads = 1024; - size /= 2; - sycl::range<3> grid_dim(1, 1, (size - 1) / threads + 1); - sycl::range<3> block_dim(1, 1, threads); - - { - has_capability_or_fail(stream->get_device(), {sycl::aspect::fp16}); - stream->parallel_for( - sycl::nd_range<3>(grid_dim * block_dim, block_dim), - [=](sycl::nd_item<3> item_ct1) { param_update_kernel_half(input, output, size); }); - } -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/compat.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/compat.h deleted file mode 100644 index 6d54446d472e9dfa6141436f55e0e623be48acb4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/compat.h +++ /dev/null @@ -1,19 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -Copyright NVIDIA/apex -This file is adapted from fused adam in NVIDIA/apex, commit a109f85 -*/ - -#ifndef TORCH_CHECK -#define TORCH_CHECK AT_CHECK -#endif - -#ifdef VERSION_GE_1_3 -#define DATA_PTR data_ptr -#else -#define DATA_PTR data -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adagrad.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adagrad.h deleted file mode 100644 index 660f860917f6c8b345825e6110a7508ba19a584f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adagrad.h +++ /dev/null @@ -1,120 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#define NOMINMAX // Windows idiosyncrasy - // https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c - -#include -#include -#include "simd.h" - -typedef unsigned short ds_half_precision_t; - -#define STEP(SPAN) \ - void Step_##SPAN(float* _params, \ - float* grads, \ - float* _exp_avg_sq, \ - size_t _param_size, \ - ds_half_precision_t* dev_param = nullptr, \ - bool half_precision = false); - -class Adagrad_Optimizer { -public: - Adagrad_Optimizer(float alpha = 1e-2, float eps = 1e-8, float weight_decay = 0) - : _alpha(alpha), _eps(eps), _weight_decay(weight_decay) - { - } - ~Adagrad_Optimizer() {} -#if defined(__AVX512__) or defined(__AVX256__) - template - void Step_AVX(size_t* rounded_size, - float* _params, - float* grads, - float* _exp_avg_sq, - size_t param_size, - ds_half_precision_t* dev_param = nullptr, - bool half_precision = false); -#endif - STEP(1) - STEP(4) - STEP(8) - inline void IncrementStep(size_t step) - { - _step++; - if (_step != step) { _step = step; } - } - inline void update_state(float lr, float epsilon, float weight_decay) - { - _alpha = lr; - _eps = epsilon; - _weight_decay = weight_decay; - } - -private: - float _alpha; - float _eps; - float _weight_decay; - - float _betta1_t; - float _betta2_t; - size_t _step; -}; - -#if defined(__AVX512__) or defined(__AVX256__) -template -void Adagrad_Optimizer::Step_AVX(size_t* rounded_size, - float* _params, - float* grads, - float* _exp_avg_sq, - size_t _param_size, - ds_half_precision_t* dev_params, - bool half_precision) -{ - size_t new_rounded_size = 0; - AVX_Data eps_4; - eps_4.data = SIMD_SET(_eps); - - float step_size = -1 * _alpha; - AVX_Data step_size_4; - step_size_4.data = SIMD_SET(step_size); - - AVX_Data weight_decay4; - if (_weight_decay > 0) weight_decay4.data = SIMD_SET(_weight_decay); - new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span); - for (size_t t = 0; t < new_rounded_size; t += TILE) { - size_t copy_size = TILE; - if ((t + TILE) > new_rounded_size) copy_size = new_rounded_size - t; - size_t offset = copy_size + t; -#pragma omp parallel for - for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { - AVX_Data grad_4[span]; - simd_load(grad_4, grads + i, half_precision); - - AVX_Data momentum_4[span]; - simd_load(momentum_4, grads + i, false); - - AVX_Data variance_4[span]; - simd_load(variance_4, _exp_avg_sq + i, false); - - AVX_Data param_4[span]; - simd_load(param_4, _params + i, half_precision); - - if (_weight_decay > 0) { simd_fma(grad_4, param_4, weight_decay4, grad_4); } - - simd_fma(variance_4, grad_4, grad_4, variance_4); - simd_sqrt(grad_4, variance_4); - simd_add(grad_4, grad_4, eps_4); - simd_div(grad_4, momentum_4, grad_4); - simd_fma(param_4, grad_4, step_size_4, param_4); - - simd_store(_params + i, param_4, half_precision); - simd_store(_exp_avg_sq + i, variance_4, false); - } - } - *rounded_size = new_rounded_size; -} -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adam.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adam.h deleted file mode 100644 index 7bc0364c569d00d01b9f0b526d21d3c57fb478dc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adam.h +++ /dev/null @@ -1,237 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#define NOMINMAX // Windows idiosyncrasy - // https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c - -#include -#include -#include -#include "simd.h" - -#include -typedef unsigned short ds_half_precision_t; - -#define STEP(SPAN) \ - void Step_##SPAN(float* _params, \ - float* grads, \ - float* _exp_avg, \ - float* _exp_avg_sq, \ - size_t _param_size, \ - ds_half_precision_t* dev_param = nullptr, \ - bool half_precision = false); - -class Adam_Optimizer { -public: - Adam_Optimizer(float alpha = 1e-3, - float betta1 = 0.9, - float betta2 = 0.999, - float eps = 1e-8, - float weight_decay = 0, - bool adamw_mode = true) - : _alpha(alpha), - _betta1(betta1), - _betta2(betta2), - _eps(eps), - _weight_decay(weight_decay), - _betta1_t(1.0), - _betta2_t(1.0), - _step(0), - _adamw_mode(adamw_mode) - { - } - ~Adam_Optimizer() {} - -#if defined(__AVX512__) or defined(__AVX256__) - template - void Step_AVX(size_t* rounded_size, - float* _params, - float* grads, - float* _exp_avg, - float* _exp_avg_sq, - size_t param_size, - ds_half_precision_t* dev_param = nullptr, - bool half_precision = false); -#endif - STEP(1) - STEP(4) - STEP(8) - inline void IncrementStep(size_t step, float beta1, float beta2) - { - if (beta1 != _betta1 || beta2 != _betta2) { - _step = step; - _betta1 = beta1; - _betta2 = beta2; - _betta1_t = std::pow(_betta1, step); - _betta2_t = std::pow(_betta2, step); - } else { - _step++; - if (_step != step) { - _betta1_t = std::pow(_betta1, step); - _betta2_t = std::pow(_betta2, step); - _step = step; - } else { - _betta1_t *= _betta1; - _betta2_t *= _betta2; - } - } - } - inline void update_state(float lr, float epsilon, float weight_decay, bool bias_correction) - { - _alpha = lr; - _eps = epsilon; - _weight_decay = weight_decay; - - _bias_correction1 = 1.0f; - _bias_correction2 = 1.0f; - if (bias_correction == 1) { - _bias_correction1 = 1 - _betta1_t; - _bias_correction2 = 1 / sqrt(1 - _betta2_t); - } - } - -private: - float _alpha; - float _betta1; - float _betta2; - float _eps; - float _weight_decay; - - float _betta1_t; - float _betta2_t; - size_t _step; - - float _bias_correction1; - float _bias_correction2; - - bool _adamw_mode; -}; - -#if defined(__AVX512__) or defined(__AVX256__) -template -void Adam_Optimizer::Step_AVX(size_t* rounded_size, - float* _params, - float* grads, - float* _exp_avg, - float* _exp_avg_sq, - size_t _param_size, - ds_half_precision_t* dev_params, - bool half_precision) -{ - size_t new_rounded_size = 0; - int rshft = half_precision ? 1 : 0; - - AVX_Data betta1_4; - betta1_4.data = SIMD_SET(_betta1); - AVX_Data betta2_4; - betta2_4.data = SIMD_SET(_betta2); - - float betta1_minus1 = 1 - _betta1; - float betta2_minus1 = 1 - _betta2; - AVX_Data betta1_minus1_4; - betta1_minus1_4.data = SIMD_SET(betta1_minus1); - AVX_Data betta2_minus1_4; - betta2_minus1_4.data = SIMD_SET(betta2_minus1); - - AVX_Data bias2_sqrt; - bias2_sqrt.data = SIMD_SET(_bias_correction2); - - AVX_Data eps_4; - eps_4.data = SIMD_SET(_eps); - - float step_size = -1 * _alpha / _bias_correction1; - AVX_Data step_size_4; - step_size_4.data = SIMD_SET(step_size); - - float w_decay = -1 * _alpha * _weight_decay; - AVX_Data weight_decay4; - if (_weight_decay > 0) - weight_decay4.data = (_adamw_mode ? SIMD_SET(w_decay) : SIMD_SET(_weight_decay)); - new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span); - for (size_t t = 0; t < new_rounded_size; t += TILE) { - size_t copy_size = TILE; - if ((t + TILE) > new_rounded_size) copy_size = new_rounded_size - t; - size_t offset = copy_size + t; -#pragma omp parallel for - for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { - AVX_Data grad_4[span]; - simd_load(grad_4, grads + (i >> rshft), half_precision); - - AVX_Data momentum_4[span]; - simd_load(momentum_4, _exp_avg + i, false); - - AVX_Data variance_4[span]; - simd_load(variance_4, _exp_avg_sq + i, false); - - AVX_Data param_4[span]; - simd_load(param_4, _params + (i >> rshft), half_precision); - - if (_weight_decay > 0 && !_adamw_mode) { - simd_fma(grad_4, param_4, weight_decay4, grad_4); - } - - simd_mul(momentum_4, momentum_4, betta1_4); - simd_fma(momentum_4, grad_4, betta1_minus1_4, momentum_4); - simd_mul(variance_4, variance_4, betta2_4); - simd_mul(grad_4, grad_4, grad_4); - simd_fma(variance_4, grad_4, betta2_minus1_4, variance_4); - simd_sqrt(grad_4, variance_4); - simd_fma(grad_4, grad_4, bias2_sqrt, eps_4); - simd_div(grad_4, momentum_4, grad_4); - - if (_weight_decay > 0 && _adamw_mode) { - simd_fma(param_4, param_4, weight_decay4, param_4); - } - - simd_fma(param_4, grad_4, step_size_4, param_4); - - simd_store(_params + (i >> rshft), param_4, half_precision); - simd_store(_exp_avg + i, momentum_4, false); - simd_store(_exp_avg_sq + i, variance_4, false); - } - } - *rounded_size = new_rounded_size; -} -#endif - -int create_adam_optimizer(int optimizer_id, - float alpha = 1e-3, - float betta1 = 0.9, - float betta2 = 0.999, - float eps = 1e-8, - float weight_decay = 0, - bool adamw_mode = true, - bool should_log = false); - -int ds_adam_step(int optimizer_id, - size_t step, - float lr, - float beta1, - float beta2, - float epsilon, - float weight_decay, - bool bias_correction, - torch::Tensor& params, - torch::Tensor& grads, - torch::Tensor& exp_avg, - torch::Tensor& exp_avg_sq); - -int ds_adam_step_plus_copy(int optimizer_id, - size_t step, - float lr, - float beta1, - float beta2, - float epsilon, - float weight_decay, - bool bias_correction, - torch::Tensor& params, - torch::Tensor& grads, - torch::Tensor& exp_avg, - torch::Tensor& exp_avg_sq, - torch::Tensor& gpu_params); - -int destroy_adam_optimizer(int optimizer_id); diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/simd.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/simd.h deleted file mode 100644 index 097e2d8585ccbdb83f1c249c1b6d89272a134999..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/simd.h +++ /dev/null @@ -1,198 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#pragma once - -#if (__x86_64__ || __i386__) -#include -#include -#endif - -#define TILE (128 * 1024 * 1024) -#if defined(__AVX512__) or defined(__AVX256__) - -#define ROUND_DOWN(size, step) ((size) & ~((step) - 1)) - -#if defined(__AVX512__) -#define SIMD_STORE(a, d) _mm512_storeu_ps(a, d) -#define SIMD_LOAD(x) _mm512_loadu_ps(x) -#define SIMD_SET(x) _mm512_set1_ps(x) -#define SIMD_ADD(x, y) _mm512_add_ps(x, y) -#define SIMD_MUL(x, y) _mm512_mul_ps(x, y) -#define SIMD_FMA(x, y, c) _mm512_fmadd_ps(x, y, c) -#define SIMD_SQRT(x) _mm512_sqrt_ps(x) -#define SIMD_DIV(x, y) _mm512_div_ps(x, y) -#define SIMD_AND(x, y) _mm512_and_ps(x, y) -#define SIMD_ANDNOT(x, y) _mm512_andnot_ps(x, y) -#define SIMD_OR(x, y) _mm512_or_ps(x, y) -#define SIMD_XOR(x, y) _mm512_xor_ps(x, y) -#define SIMD_WIDTH 16 - -#define SIMD_LOAD2(x, h) \ - ((h) ? _mm512_cvtph_ps(_mm256_castps_si256(_mm256_loadu_ps(x))) : _mm512_loadu_ps(x)) -#define SIMD_STORE2(x, d, h) \ - ((h) ? _mm256_store_ps(x, _mm256_castsi256_ps(_mm512_cvtps_ph(d, _MM_FROUND_TO_NEAREST_INT))) \ - : _mm512_storeu_ps(x, d)) - -#define INTV __m256i -#elif defined(__AVX256__) -#define SIMD_STORE(a, d) _mm256_storeu_ps(a, d) -#define SIMD_LOAD(x) _mm256_loadu_ps(x) -#define SIMD_SET(x) _mm256_set1_ps(x) -#define SIMD_ADD(x, y) _mm256_add_ps(x, y) -#define SIMD_MUL(x, y) _mm256_mul_ps(x, y) -#define SIMD_FMA(x, y, c) _mm256_fmadd_ps(x, y, c) -#define SIMD_SQRT(x) _mm256_sqrt_ps(x) -#define SIMD_DIV(x, y) _mm256_div_ps(x, y) -#define SIMD_AND(x, y) _mm256_and_ps(x, y) -#define SIMD_ANDNOT(x, y) _mm256_andnot_ps(x, y) -#define SIMD_OR(x, y) _mm256_or_ps(x, y) -#define SIMD_XOR(x, y) _mm256_xor_ps(x, y) -#define SIMD_WIDTH 8 - -#define SIMD_LOAD2(x, h) \ - ((h) ? _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*)x)) : _mm256_loadu_ps(x)) -#define SIMD_STORE2(x, d, h) \ - ((h) ? _mm_store_ps(x, _mm_castsi128_ps(_mm256_cvtps_ph(d, _MM_FROUND_TO_NEAREST_INT))) \ - : _mm256_storeu_ps(x, d)) - -#define INTV __m128i -#endif - -union AVX_Data { -#if defined(__AVX512__) - __m512 data; -#elif defined(__AVX256__) - __m256 data; -#endif - // float data_f[16]; -}; - -template -inline void simd_store(float* dst, AVX_Data* src, bool half_precision) -{ - size_t width = (half_precision ? SIMD_WIDTH / 2 : SIMD_WIDTH); -#pragma unroll - for (size_t i = 0; i < span; ++i) { SIMD_STORE2(dst + width * i, src[i].data, half_precision); } -} -template -inline void simd_load(AVX_Data* dst, float* src, bool half_precision) -{ - size_t width = (half_precision ? 1 : SIMD_WIDTH); -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_LOAD2(src + width * i, half_precision); } -} -template -inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data src_m_r, AVX_Data* src_a) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { - dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r.data, src_a[i].data); - } -} -template -inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data src_m_r, AVX_Data src_a) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { - dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r.data, src_a.data); - } -} -template -inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data* src_m_r, AVX_Data* src_a) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { - dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r[i].data, src_a[i].data); - } -} -template -inline void simd_sqrt(AVX_Data* dst, AVX_Data* src) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_SQRT(src[i].data); } -} -template -inline void simd_add(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ADD(src_a_l[i].data, src_a_r.data); } -} -template -inline void simd_add(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ADD(src_a_l[i].data, src_a_r[i].data); } -} -template -inline void simd_mul(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_MUL(src_a_l[i].data, src_a_r.data); } -} -template -inline void simd_mul(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_MUL(src_a_l[i].data, src_a_r[i].data); } -} -template -inline void simd_div(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_DIV(src_a_l[i].data, src_a_r[i].data); } -} -template -inline void simd_and(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_AND(src_a_l[i].data, src_a_r.data); } -} -template -inline void simd_and(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_AND(src_a_l[i].data, src_a_r[i].data); } -} -template -inline void simd_andnot(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ANDNOT(src_a_l[i].data, src_a_r.data); } -} -template -inline void simd_andnot(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { - dst[i].data = SIMD_ANDNOT(src_a_l[i].data, src_a_r[i].data); - } -} -template -inline void simd_or(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_OR(src_a_l[i].data, src_a_r.data); } -} -template -inline void simd_or(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_OR(src_a_l[i].data, src_a_r[i].data); } -} -template -inline void simd_xor(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_XOR(src_a_l[i].data, src_a_r.data); } -} -template -inline void simd_xor(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) -{ -#pragma unroll - for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_XOR(src_a_l[i].data, src_a_r[i].data); } -} - -#endif diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/type_shim.h b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/type_shim.h deleted file mode 100644 index 1897afd1fea248f497acf3fb6ba049baa1ca56c3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/includes/type_shim.h +++ /dev/null @@ -1,155 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* Taken from NVIDIA/apex commit 855808f3fc268e9715d613f3c2e56469d8c986d8 */ -#include -/* #include */ -#include - -// Forward/backward compatibility hack around -// https://github.com/pytorch/pytorch/commit/3aeb78079bcd68282fe9117088e138b77318e288 -// pending more future-proof guidance from upstream. -// struct TypeShim -// { -// const at::Type& payload; -// TypeShim(const at::Type& type) : payload(type) {} -// // Enable trivial conversion to a const at::Type& for pre-3aeb78 -// operator const at::Type&(){ return payload; }; -// // Enable dispatch switch statements to take *this directly for post-3aeb78 -// //operator at::ScalarType(){ return payload.; }; -// }; - -#define DISPATCH_FLOAT_AND_HALF(TYPE, LEVEL, NAME, ...) \ - switch (TYPE) { \ - case at::ScalarType::Float: { \ - using scalar_t_##LEVEL = float; \ - __VA_ARGS__; \ - break; \ - } \ - case at::ScalarType::Half: { \ - using scalar_t_##LEVEL = at::Half; \ - __VA_ARGS__; \ - break; \ - } \ - case at::ScalarType::BFloat16: { \ - using scalar_t_##LEVEL = at::BFloat16; \ - __VA_ARGS__; \ - break; \ - } \ - default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ - } - -#define DISPATCH_DOUBLE_FLOAT_AND_HALF(TYPE, LEVEL, NAME, ...) \ - switch (TYPE) { \ - case at::ScalarType::Double: { \ - using scalar_t_##LEVEL = double; \ - __VA_ARGS__; \ - break; \ - } \ - case at::ScalarType::Float: { \ - using scalar_t_##LEVEL = float; \ - __VA_ARGS__; \ - break; \ - } \ - case at::ScalarType::Half: { \ - using scalar_t_##LEVEL = at::Half; \ - __VA_ARGS__; \ - break; \ - } \ - case at::ScalarType::BFloat16: { \ - using scalar_t_##LEVEL = at::BFloat16; \ - __VA_ARGS__; \ - break; \ - } \ - default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ - } - -#define DISPATCH_DOUBLE_AND_FLOAT(TYPE, LEVEL, NAME, ...) \ - switch (TYPE) { \ - case at::ScalarType::Double: { \ - using scalar_t_##LEVEL = double; \ - __VA_ARGS__; \ - break; \ - } \ - case at::ScalarType::Float: { \ - using scalar_t_##LEVEL = float; \ - __VA_ARGS__; \ - break; \ - } \ - default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ - } - -template -__inline__ __attribute__((always_inline)) T -reduce_block_into_lanes(T* x, - T val, - int lanes = 1, - bool share_result = false) // lanes is intended to be <= 32. -{ - auto item_ct1 = sycl::ext::oneapi::experimental::this_nd_item<3>(); - int tid = item_ct1.get_local_id(2) + item_ct1.get_local_id(1) * item_ct1.get_local_range(2); - int blockSize = item_ct1.get_local_range(2) * - item_ct1.get_local_range(1); // blockSize is intended to be a multiple of 32. - - if (blockSize >= 64) { - x[tid] = val; - /* - DPCT1118:1: SYCL group functions and algorithms must be encountered in converged control - flow. You may need to adjust the code. - */ - /* - DPCT1065:6: Consider replacing sycl::nd_item::barrier() with - sycl::nd_item::barrier(sycl::access::fence_space::local_space) for better performance if - there is no access to global memory. - */ - item_ct1.barrier(); - } - -#pragma unroll - for (int i = (blockSize >> 1); i >= 64; i >>= 1) { - if (tid < i) x[tid] = x[tid] + x[tid + i]; - /* - DPCT1118:2: SYCL group functions and algorithms must be encountered in converged control - flow. You may need to adjust the code. - */ - /* - DPCT1065:7: Consider replacing sycl::nd_item::barrier() with - sycl::nd_item::barrier(sycl::access::fence_space::local_space) for better performance if - there is no access to global memory. - */ - item_ct1.barrier(); - } - - T final; - - if (tid < 32) { - if (blockSize >= 64) - final = x[tid] + x[tid + 32]; - else - final = val; - // __SYNCWARP(); - -#pragma unroll - for (int i = 16; i >= lanes; i >>= 1) - final = final + __shfl_down_sync(0xffffffff, final, i); - } - - if (share_result) { - if (tid < lanes) x[tid] = final; // EpilogueOp - // Make sure the smem result is visible to all warps. - /* - DPCT1118:3: SYCL group functions and algorithms must be encountered in converged control - flow. You may need to adjust the code. - */ - /* - DPCT1065:8: Consider replacing sycl::nd_item::barrier() with - sycl::nd_item::barrier(sycl::access::fence_space::local_space) for better performance if - there is no access to global memory. - */ - item_ct1.barrier(); - } - - return final; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/packbits/packing.cpp b/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/packbits/packing.cpp deleted file mode 100644 index fb09bb16efdda5b578c239c35be7ee493dc87c75..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/csrc/xpu/packbits/packing.cpp +++ /dev/null @@ -1,100 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -#include -#include -#include -#include - -using namespace sycl; -using namespace xpu; - -void packbitskernel(const float* input, uint8_t* output, const int input_size, id<1> item_ct1) -{ - // get the sign bit of each float and pack them into byte - int i = item_ct1; - for (int j = 0; j < 8; ++j) { - int k = i * 8 + j; - int bit = k < input_size && (!sycl::signbit(input[k])); - output[i] |= bit << (7 - j); - } -} - -void unpackbitskernel(const uint8_t* input, float* output, id<1> item_ct1) -{ - // use the bit value to set float, bit 0 -> float -1, bit 1 -> float 1 - int i = item_ct1; - output[i] = (float((input[i / 8] >> (7 - i % 8)) & 1) - 0.5) * 2; -} - -sycl::queue get_current_queue(at::Device device) -{ - c10::impl::VirtualGuardImpl impl(device.type()); - c10::Stream _stream = impl.getStreamFromGlobalPool(device, /*isHighPriority=*/false); - sycl::queue queue = xpu::get_queue_from_stream(_stream); - return queue; -} - -/* -pack float tensor into uint8 tensor. Every eight float elements get packed into one uint8 -if float x >= 0, will be packed as a '1' bit, or will be packed as '0' -Arguments: - tensor: A bool tensor that get packed. - input_size: numel of input tensor - rank: device id in order to get corresponding stream -*/ -at::Tensor packbits(at::Tensor tensor, int input_size, int rank) -{ - at::Device device = "xpu:" + std::to_string(rank); - sycl::queue q = get_current_queue(device); - - int packed_size = (input_size + 7) / 8; - auto unit8_options = at::TensorOptions().dtype(at::kByte).device(at::kXPU); - at::Tensor packed = torch::zeros({packed_size}, unit8_options); - - float* input = (float*)tensor.data_ptr(); - uint8_t* output = (uint8_t*)packed.data_ptr(); - - auto event = q.submit([&](sycl::handler& cgh) { - cgh.parallel_for<>(range(packed_size), [=](id<1> item_ct1) { - packbitskernel(input, output, input_size, item_ct1); - }); - }); - - return packed; -} - -/* -unpack uint8 tensor into float tensor. Every uint8 element get unpacked into eight float -a '1' bit will be converted to a float(1), a '0' bit will be converted to a float(-1). -Arguments: - tensor: A uint8 tensor that get unpacked. - input_size: numel of input tensor - rank: device id in order to get corresponding stream -*/ -at::Tensor unpackbits(at::Tensor tensor, int input_size, int rank) -{ - at::Device device = "xpu:" + std::to_string(rank); - sycl::queue q = get_current_queue(device); - - auto float_options = at::TensorOptions().dtype(at::kFloat).device(at::kXPU); - at::Tensor unpacked = torch::empty({input_size * 8}, float_options); - - uint8_t* input = (uint8_t*)tensor.data_ptr(); - float* output = (float*)unpacked.data_ptr(); - - auto event = q.submit([&](sycl::handler& cgh) { - cgh.parallel_for<>(range(input_size * 8), - [=](id<1> item_ct1) { unpackbitskernel(input, output, item_ct1); }); - }); - - return unpacked; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("packbits", &packbits, "DeepSpeed XPU packbits (C++)"); - m.def("unpackbits", &unpackbits, "DeepSpeed XPU unpackbits (C++)"); -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__init__.py deleted file mode 100644 index 1c5fd280fc32661bae081ad8824c23efa36c9c97..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .evoformer_attn import DS4Sci_EvoformerAttention, EvoformerFusedAttention diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index a12e2f2a54200db66cc3c0ff9fe637c23c37b6ab..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__pycache__/evoformer_attn.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__pycache__/evoformer_attn.cpython-310.pyc deleted file mode 100644 index 572e6cc40a018cedc713cd4f67326c33eab8e2c3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/__pycache__/evoformer_attn.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/evoformer_attn.py b/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/evoformer_attn.py deleted file mode 100644 index da5843d6de31ed94c27c466bbcf5ae5966bb620f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/deepspeed4science/evoformer_attn.py +++ /dev/null @@ -1,106 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import numpy as np -from deepspeed.ops.op_builder import EvoformerAttnBuilder -from deepspeed.accelerator import get_accelerator - -kernel_ = None - - -def _attention(Q, K, V, bias1, bias2): - assert Q.shape[-3] > 16, "seq_len must be greater than 16" - O = torch.empty_like(Q, dtype=Q.dtype) - assert get_accelerator().on_accelerator(Q), "Q must be on cuda" - assert get_accelerator().on_accelerator(K), "K must be on cuda" - assert get_accelerator().on_accelerator(V), "V must be on cuda" - assert get_accelerator().on_accelerator(bias1), "bias1 must be on cuda" - assert get_accelerator().on_accelerator(bias2), "bias2 must be on cuda" - global kernel_ - if kernel_ is None: - kernel_ = EvoformerAttnBuilder().load() - nheads = Q.shape[-2] - nq = (Q.shape[-3] + 31) // 32 * 32 - nb = np.prod(Q.shape[:-3]) - lse = torch.empty((nb, nheads, nq), dtype=torch.float32, device=Q.device) - kernel_.attention(Q, K, V, bias1, bias2, O, lse) - return O, lse - - -def attention_bwd(dO, Q, K, V, O, lse, bias1, bias2, bias1_grad, bias2_grad): - assert max(Q.shape[-1], V.shape[-1]) <= 64, "Hidden size is too large. Need to change kMax to a larger value" - dQ = torch.empty_like(Q, dtype=Q.dtype) - dK = torch.empty_like(K, dtype=K.dtype) - dV = torch.empty_like(V, dtype=V.dtype) - assert get_accelerator().on_accelerator(dO), "dO must be on cuda" - assert get_accelerator().on_accelerator(Q), "Q must be on cuda" - assert get_accelerator().on_accelerator(K), "K must be on cuda" - assert get_accelerator().on_accelerator(V), "V must be on cuda" - assert get_accelerator().on_accelerator(O), "O must be on cuda" - global kernel_ - if kernel_ is None: - kernel_ = EvoformerAttnBuilder().load() - delta = torch.empty_like(lse) - if bias1_grad: - dB1 = torch.zeros_like(bias1, dtype=torch.float32) - else: - dB1 = torch.tensor([], dtype=torch.float32, device=bias1.device) - if bias2_grad: - dB2 = torch.zeros_like(bias2, dtype=torch.float32) - else: - dB2 = torch.tensor([], dtype=torch.float32, device=bias2.device) - kernel_.attention_bwd(dO, Q, K, V, O, lse, delta, bias1, bias2, dQ, dK, dV, dB1, dB2) - return dQ, dK, dV, dB1.to(dO.dtype), dB2.to(dO.dtype) - - -class EvoformerFusedAttention(torch.autograd.Function): - - @staticmethod - def forward(ctx, q, k, v, bias1=None, bias2=None): - """ - q, k, v: are in shape [*, L, H, D] - """ - bias1_ = bias1.contiguous() if bias1 is not None else torch.tensor([], dtype=q.dtype, device=q.device) - bias2_ = bias2.contiguous() if bias2 is not None else torch.tensor([], dtype=q.dtype, device=q.device) - q = q.contiguous() - k = k.contiguous() - v = v.contiguous() - o, lse = _attention(q, k, v, bias1_, bias2_) - ctx.save_for_backward(q, k, v, o, lse, bias1_, bias2_) - return o - - @staticmethod - def backward(ctx, grad_output): - q, k, v, o, lse, bias1, bias2 = ctx.saved_tensors - is_b1_grad = bias1.numel() != 0 and ctx.needs_input_grad[3] - is_b2_grad = bias2.numel() != 0 and ctx.needs_input_grad[4] - dQ, dK, dV, dB1, dB2 = attention_bwd(grad_output, q, k, v, o, lse, bias1, bias2, is_b1_grad, is_b2_grad) - if not is_b1_grad: - dB1 = None - if not is_b2_grad: - dB2 = None - return dQ, dK, dV, dB1, dB2 - - -def DS4Sci_EvoformerAttention(Q, K, V, biases): - assert len(biases) <= 2 - - if (len(biases) == 0): - biases.append(None) - - if (len(biases) == 1): - biases.append(None) - - bias_1_shape = lambda x: (x.shape[0], x.shape[1], 1, 1, x.shape[2]) - bias_2_shape = lambda x: (x.shape[0], 1, x.shape[3], x.shape[2], x.shape[2]) - - if biases[0] is not None: - assert biases[0].shape == bias_1_shape(Q), "bias1 shape is incorrect" - - if biases[1] is not None: - assert biases[1].shape == bias_2_shape(Q), "bias2 shape is incorrect" - - return EvoformerFusedAttention.apply(Q, K, V, biases[0], biases[1]) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__init__.py deleted file mode 100644 index f9cf23373c265a1c0521a9eaa3b713cdd7aca57c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .quantize import FP_Quantize, Quantizer -from .fp8_gemm import matmul_fp8 diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 3bdb0d9c972e4d6353b211a31e59fc9105786d6d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm.cpython-310.pyc deleted file mode 100644 index 24194fccc374fa7f9cd20898e6df9753d0ebcc2e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm_triton.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm_triton.cpython-310.pyc deleted file mode 100644 index b973add6fe3e02b5e5138e3588ffef65323ee64a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm_triton.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/quantize.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/quantize.cpython-310.pyc deleted file mode 100644 index c2761ae6b6d35c85c520ad268a07204d1bb8a372..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/__pycache__/quantize.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm.py deleted file mode 100644 index db4fa5ae2c9271b480b6d209da28f0f4fdbbe85d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm.py +++ /dev/null @@ -1,28 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -######## Fused MoE kernel ######### -# These kernels are implemented for -# fusing GeMM with dequantization of -# fp8 weight data when using bit-16 -# activation. -################################### - -import torch - - -def matmul_fp8(inp, weight, scale, quantization_group_size, quantizer): - from deepspeed import get_accelerator - - if not get_accelerator().is_triton_supported(): - return matmul_fp8_fallback(inp, weight, scale, quantization_group_size, quantizer) - else: - # Import dynamically to prevent failures on systems without triton. - from .fp8_gemm_triton import matmul_fp8_triton - return matmul_fp8_triton(inp, weight, scale, quantization_group_size) - - -def matmul_fp8_fallback(inp, weight, scale, quantization_group_size, quantizer): - return torch.matmul(inp, quantizer.dequantize(weight, scale=scale)) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm_triton.py b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm_triton.py deleted file mode 100644 index 086525cc64425558399886cc534c57a039ee0a41..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm_triton.py +++ /dev/null @@ -1,165 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -######## Fused MoE kernel ######### -# These kernels are implemented for -# fusing GeMM with dequantization of -# fp8 weight data when using bit-16 -# activation. -################################### - -import torch -import triton -import triton.language as tl - - -@triton.jit -def matmul_kernel_fp8_bf16(inp_ptr, weight_ptr, out_ptr, scale_ptr, M, N, K, stride_am, stride_ak, stride_bk, - stride_bn, stride_cm, stride_cn, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, - quantization_group_size: tl.constexpr): - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - - offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N - offs_k = tl.arange(0, BLOCK_SIZE_K) - - inp_data = inp_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) - weight_data = weight_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) - weight_ptrs_offset = offs_k[:, None] * (stride_bk // quantization_group_size) + ( - (pid_n * BLOCK_SIZE_N) // quantization_group_size) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - inp = tl.load(inp_data, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) - weight = tl.load(weight_data, mask=offs_k[:, None] < K, other=0.0) - scale = tl.load(scale_ptr + weight_ptrs_offset + ((k * BLOCK_SIZE_K * stride_bk) // quantization_group_size)) - # Dequantize weight (fp8 -> bf16) - w = (weight & 0x80).to(tl.uint16) << 8 - w = w | ((weight & 0x7f).to(tl.uint16) << 4) - w = (w + 0x3C00).to(tl.uint16) - w = (w.to(tl.bfloat16, bitcast=True).to(tl.float32) * scale).to(tl.bfloat16) - - inp_data += BLOCK_SIZE_K * stride_ak - weight_data += BLOCK_SIZE_K * stride_bk - - accumulator += tl.dot(inp, w) - - out = accumulator.to(tl.bfloat16) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - out_data = out_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] - tl.store(out_data, out, mask=(offs_cm[:, None] < M) & (offs_cn[None, :] < N)) - - -@triton.jit -def matmul_kernel_fp8_fp16(inp_ptr, weight_ptr, out_ptr, scale_ptr, M, N, K, stride_am, stride_ak, stride_bk, - stride_bn, stride_cm, stride_cn, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, - quantization_group_size: tl.constexpr): - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - - offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N - offs_k = tl.arange(0, BLOCK_SIZE_K) - - inp_data = inp_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) - weight_data = weight_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) - weight_ptrs_offset = offs_k[:, None] * (stride_bk // quantization_group_size) + ( - (pid_n * BLOCK_SIZE_N) // quantization_group_size) - - weight = tl.load(weight_data, mask=offs_k[:, None] < K, other=0.0) - scale = tl.load(scale_ptr + weight_ptrs_offset) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - inp = tl.load(inp_data, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) - # Dequantize weight (fp8 -> fp16) - w = (((weight & 0x80) << 8) | ((weight & 0x7f) << 7)).to(tl.uint16) - w = (w + 0x2000).to(tl.uint16) - w = (w.to(tl.float16, bitcast=True) * scale).to(tl.float16) - - inp_data += BLOCK_SIZE_K * stride_ak - weight_data += BLOCK_SIZE_K * stride_bk - - weight = tl.load(weight_data, mask=offs_k[:, None] < K - (k + 1) * BLOCK_SIZE_K, other=0.0) - scale = tl.load(scale_ptr + (weight_ptrs_offset + - (((k + 1) * BLOCK_SIZE_K * stride_bk) // quantization_group_size))) - - accumulator += tl.dot(inp, w) - - out = accumulator.to(tl.float16) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - out_data = out_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] - tl.store(out_data, out, mask=(offs_cm[:, None] < M) & (offs_cn[None, :] < N)) - - -def matmul_fp8_triton(inp, weight, scale, quantization_group_size): - - assert inp.shape[1] == weight.shape[0], \ - f"Incompatible dimensions (input: {inp.shape}, weight: {weight.shape})" - - M, K = inp.shape - K, N = weight.shape - - out = torch.empty((M, N), device=inp.device, dtype=inp.dtype) - - # GEMM tuning parameters! - # TODO: Add a more configurable tuning for selecting the best GeMM - BLOCK_SIZE_M = 16 if M <= 16 else 32 if M <= 32 else 64 if M <= 64 else 128 - BLOCK_SIZE_N = 64 - BLOCK_SIZE_K = max(64, quantization_group_size) - GROUP_SIZE_M = 8 - num_stages = 4 - num_warps = 4 - if M >= 256: - BLOCK_SIZE_M = 256 - BLOCK_SIZE_N = 128 - BLOCK_SIZE_K = max(128, quantization_group_size) - num_stages = 3 - num_warps = 8 - - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), ) - kernel = matmul_kernel_fp8_bf16 if inp.dtype == torch.bfloat16 else matmul_kernel_fp8_fp16 - kernel[grid](inp, - weight, - out, - scale, - M, - N, - K, - inp.stride(0), - inp.stride(1), - weight.stride(0), - weight.stride(1), - out.stride(0), - out.stride(1), - quantization_group_size=quantization_group_size, - BLOCK_SIZE_M=BLOCK_SIZE_M, - BLOCK_SIZE_N=BLOCK_SIZE_N, - BLOCK_SIZE_K=BLOCK_SIZE_K, - GROUP_SIZE_M=GROUP_SIZE_M, - num_stages=num_stages, - num_warps=num_warps) - return out diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/quantize.py b/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/quantize.py deleted file mode 100644 index 47b3b08c7e03d94a98cc6c79af071beb82e40860..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/fp_quantizer/quantize.py +++ /dev/null @@ -1,167 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import abc -from abc import ABC - -import gc -from deepspeed.ops.op_builder import FPQuantizerBuilder -from deepspeed.accelerator import get_accelerator - -fp_quant_module = None - - -class Quantizer(ABC): - """ - Abstract Quantizer class that implements quantize/dequantize methods. - - Arguments: - group_size (int, optional): number of values or elements that are grouped - together for the quantization process. - """ - - def __init__(self, group_size=512) -> None: - self.group_size = group_size - - @abc.abstractmethod - def quantize(self, - input, - q_bits=8, - q_mantisa_bits=3, - stochastic_mode=False, - return_meta_tensor=False) -> torch.Tensor: - ... - - @abc.abstractmethod - def dequantize(self, input_q, fp_out=None, q_bits=8, q_mantisa_bits=3, scale=None) -> torch.Tensor: - ... - - -class FP_Quantize(Quantizer): - - def __init__(self, quantization_config) -> None: - global fp_quant_module - super().__init__(group_size=quantization_config.group_size) - if fp_quant_module is None: - fp_quant_module = FPQuantizerBuilder().load() - self.cuda_impl = getattr(fp_quant_module, "CUDA_IMPL", True) - self.q_config = quantization_config - - self.orig_dtype = None - self.num_groups = None - self.input_q = None - self.scale = None - - def quantize(self, - input, - q_bits=8, - q_mantisa_bits=3, - stochastic_mode=False, - return_meta_tensor=False) -> torch.Tensor: - assert input.dtype == torch.bfloat16, "only support bf16 for now" - if return_meta_tensor: - assert q_bits == 8, "meta tensor is only supported with q_bit=8" - - self.orig_dtype = input.dtype - self.orig_shape = input.shape - - if q_bits == 8: - pass - elif q_bits == 12: - q_mantisa_bits = 4 - elif q_bits == 6: - q_mantisa_bits = 2 - elif q_bits == 4: - q_mantisa_bits = 1 - else: - assert (0), \ - f"Missing {q_bits}-quantization, please add the template arguments for the kernel to support this precision!" - self.num_groups = input.numel() // self.group_size - self.input_q = torch.ones(self.num_groups, - int(self.group_size * q_bits) // 8 + 4, - dtype=torch.uint8, - device=input.device) - out = fp_quant_module.quantize(self.input_q, input, self.group_size, stochastic_mode, q_bits, q_mantisa_bits) - if return_meta_tensor: - data, self.scale = out.split(self.group_size, dim=-1) - data = data.contiguous().reshape(input.shape) - self.scale = self.scale.contiguous() - del self.input_q - del out - gc.collect() - get_accelerator().empty_cache() - return data, self.scale - - return out - - def to(self, *args, **kwargs): - # Intermediate tensors may need to be moved to different devices - if hasattr(self, 'input_q'): - self.input_q = self.input_q.to(*args, **kwargs) - if hasattr(self, 'scale'): - self.scale = self.scale.to(*args, **kwargs) - - def get_scales(self): - return fp_quant_module.get_scales(self.scale, self.num_groups) - - def dequantize(self, input_q, fp_out=None, q_bits=8, q_mantisa_bits=3, scale=None) -> torch.Tensor: - assert (self.orig_dtype is not None), \ - "[De-quantization Error]: you need to call quantize before dequantizing!" - fp_out = torch.empty(self.orig_shape, dtype=self.orig_dtype, - device=input_q.device) if fp_out is None else fp_out - if q_bits == 8: - pass - elif q_bits == 12: - q_mantisa_bits = 4 - elif q_bits == 6: - q_mantisa_bits = 2 - elif q_bits == 4: - q_mantisa_bits = 1 - else: - assert (0), \ - f"Missing {q_bits}-dequantization, please add the template arguments for the kernel to support this precision!" - - if scale is not None: - assert input_q.numel() == fp_out.numel(), \ - f'[De-quantization Error]: quantized data should have the same size as original tensor when scale is not None!' - input_q = torch.cat([input_q.reshape(-1, self.group_size), scale], dim=-1).contiguous() - fp_quant_module.dequantize(fp_out, input_q, self.group_size, q_mantisa_bits, q_bits - q_mantisa_bits - 1) - return fp_out - - def selective_dequantize(self, - input_q, - indexes, - fp_out=None, - q_bits=8, - q_mantisa_bits=3, - scale=None) -> torch.Tensor: - assert (not hasattr(self, 'orig_shape') or len(self.orig_shape) == 3), \ - "Selective-Dequantization works on 3d tensor only! Please reshape the tensor before calling dequantize function." - assert (self.orig_dtype is not None), \ - "[De-quantization Error]: you need to call quantize before dequantizing!" - fp_out = torch.empty( - (indexes.shape[0], - *self.orig_shape[1:]), dtype=self.orig_dtype, device=input_q.device) if fp_out is None else fp_out - if q_bits == 8: - pass - elif q_bits == 12: - q_mantisa_bits = 4 - elif q_bits == 6: - q_mantisa_bits = 2 - elif q_bits == 4: - q_mantisa_bits = 1 - else: - assert (0), \ - f"Missing {q_bits}-dequantization, please add the template arguments for the kernel to support this precision!" - - if scale is not None: - assert input_q.numel() == fp_out.numel(), \ - f'[De-quantization Error]: quantized data should have the same size as original tensor when scale is not None!' - input_q = torch.cat([input_q.reshape(-1, self.group_size), scale], dim=-1).contiguous() - - fp_quant_module.selective_dequantize(fp_out, input_q, indexes, self.group_size, q_mantisa_bits, - q_bits - q_mantisa_bits - 1) - return fp_out diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/gds/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/gds/__init__.py deleted file mode 100644 index 3c0762c810766ebc7bc832e737d7e773216475b9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/gds/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from ..op_builder import GDSBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/gds/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/gds/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index e3498d2be482d8e2b822b5ca1445fcfd9e2b4273..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/gds/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__init__.py deleted file mode 100644 index 81bc07e827c88e3f7c134eab78c761813fdf7826..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .fused_lamb import FusedLamb diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 1bef2dbda816956cc8737c73b9ddf832c5391992..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__pycache__/fused_lamb.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__pycache__/fused_lamb.cpython-310.pyc deleted file mode 100644 index 819f035891f9b80e747479172f986bb159593ad1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/__pycache__/fused_lamb.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/fused_lamb.py b/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/fused_lamb.py deleted file mode 100644 index 6ccd9d4c6b066601bb2f6e5d8d8def6be09fd22b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/lamb/fused_lamb.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Copyright NVIDIA/apex -This file is adapted from NVIDIA/apex/optimizer/fused_adam and implements the LAMB optimizer -""" -import types -import torch -from deepspeed.ops.op_builder import FusedLambBuilder - - -class FusedLamb(torch.optim.Optimizer): - """Implements the LAMB algorithm. Currently GPU-only. - - LAMB was proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes. - https://arxiv.org/abs/1904.00962 - - Arguments: - params (iterable): iterable of parameters to optimize or dicts defining - parameter groups. - lr (float, optional): learning rate. (default: 1e-3) - bias_correction (bool, optional): bias correction (default: True) - betas (Tuple[float, float], optional): coefficients used for computing - running averages of gradient and its square. (default: (0.9, 0.999)) - eps (float, optional): term added to the denominator to improve - numerical stability. (default: 1e-8) - eps_inside_sqrt (boolean, optional): in the 'update parameters' step, - adds eps to the bias-corrected second moment estimate before - evaluating square root instead of adding it to the square root of - second moment estimate as in the original paper. (default: False) - weight_decay (float, optional): weight decay (L2 penalty) (default: 0) - max_grad_norm (float, optional): value used to clip global grad norm - (default: 0.0) - max_coeff(float, optional): maximum value of the lamb coefficient (default: 10.0) - min_coeff(float, optional): minimum value of the lamb coefficient (default: 0.01) - amsgrad (boolean, optional): NOT SUPPORTED in FusedLamb! - """ - - def __init__(self, - params, - lr=1e-3, - bias_correction=True, - betas=(0.9, 0.999), - eps=1e-8, - eps_inside_sqrt=False, - weight_decay=0., - max_grad_norm=0., - max_coeff=10.0, - min_coeff=0.01, - amsgrad=False): - self.fused_lamb_cuda = FusedLambBuilder().load() - - if amsgrad: - raise RuntimeError('FusedLamb does not support the AMSGrad variant.') - defaults = dict(lr=lr, - bias_correction=bias_correction, - betas=betas, - eps=eps, - weight_decay=weight_decay, - max_grad_norm=max_grad_norm, - max_coeff=max_coeff, - min_coeff=min_coeff) - super(FusedLamb, self).__init__(params, defaults) - self.eps_mode = 0 if eps_inside_sqrt else 1 - self.lamb_coeffs = [] - - def step(self, closure=None, grads=None, output_params=None, scale=1., grad_norms=None): - """Performs a single optimization step. - - Arguments: - closure (callable, optional): A closure that reevaluates the model - and returns the loss. - grads (list of tensors, optional): weight gradient to use for the - optimizer update. If gradients have type torch.half, parameters - are expected to be in type torch.float. (default: None) - output params (list of tensors, optional): A reduced precision copy - of the updated weights written out in addition to the regular - updated weights. Have to be of same type as gradients. (default: None) - scale (float, optional): factor to divide gradient tensor values - by before applying to weights. (default: 1) - """ - loss = None - if closure is not None: - loss = closure() - - if grads is None: - grads_group = [None] * len(self.param_groups) - # backward compatibility - # assuming a list/generator of parameter means single group - elif isinstance(grads, types.GeneratorType): - grads_group = [grads] - elif type(grads[0]) != list: - grads_group = [grads] - else: - grads_group = grads - - if output_params is None: - output_params_group = [None] * len(self.param_groups) - elif isinstance(output_params, types.GeneratorType): - output_params_group = [output_params] - elif type(output_params[0]) != list: - output_params_group = [output_params] - else: - output_params_group = output_params - - if grad_norms is None: - grad_norms = [None] * len(self.param_groups) - - #remove the previous coeffs - del self.lamb_coeffs[:] - - for group, grads_this_group, output_params_this_group, grad_norm_group in zip( - self.param_groups, grads_group, output_params_group, grad_norms): - if grads_this_group is None: - grads_this_group = [None] * len(group['params']) - if output_params_this_group is None: - output_params_this_group = [None] * len(group['params']) - - if grad_norm_group is None: - grad_norm_group = [None] * len(group['params']) - elif not isinstance(grad_norm_group, list): - grad_norm_group = [grad_norm_group] - - bias_correction = 1 if group['bias_correction'] else 0 - - for p, grad, output_param, grad_norm in zip(group['params'], grads_this_group, output_params_this_group, - grad_norm_group): - - # compute combined scale factor for this group - combined_scale = scale - if group['max_grad_norm'] > 0: - # norm is in fact norm*scale - clip = ((grad_norm / scale) + 1e-6) / group['max_grad_norm'] - if clip > 1: - combined_scale = clip * scale - - #note: p.grad should not ever be set for correct operation of mixed precision optimizer that sometimes sends None gradients - if p.grad is None and grad is None: - continue - if grad is None: - grad = p.grad.data - if grad.is_sparse: - raise RuntimeError('FusedLamb does not support sparse gradients') - - state = self.state[p] - - # State initialization - if len(state) == 0: - state['step'] = 0 - # Exponential moving average of gradient values - state['exp_avg'] = torch.zeros_like(p.data) - # Exponential moving average of squared gradient values - state['exp_avg_sq'] = torch.zeros_like(p.data) - - exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] - beta1, beta2 = group['betas'] - max_coeff = group['max_coeff'] - min_coeff = group['min_coeff'] - - state['step'] += 1 - - out_p = torch.tensor([], dtype=torch.float) if output_param is None else output_param - lamb_coeff = self.fused_lamb_cuda.lamb(p.data, out_p, exp_avg, exp_avg_sq, grad, group['lr'], beta1, - beta2, max_coeff, min_coeff, group['eps'], combined_scale, - state['step'], self.eps_mode, bias_correction, - group['weight_decay']) - self.lamb_coeffs.append(lamb_coeff) - return loss - - def get_lamb_coeffs(self): - lamb_coeffs = [lamb_coeff.item() for lamb_coeff in self.lamb_coeffs] - return lamb_coeffs diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__init__.py deleted file mode 100644 index 2f90e5ec2e808057a26cc36f2d7caedecb44c86e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .cpu_lion import DeepSpeedCPULion -from .fused_lion import FusedLion diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 915ddd5b12c8e9624e97bae763c4d3c9277d7cd4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/cpu_lion.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/cpu_lion.cpython-310.pyc deleted file mode 100644 index 57689253a6ccf0e599cacf22a1ee7603efb48fd0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/cpu_lion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/fused_lion.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/fused_lion.cpython-310.pyc deleted file mode 100644 index 3d484992084b17f9527f6d12f870b9b00b8aedcf..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/fused_lion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/multi_tensor_apply.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/multi_tensor_apply.cpython-310.pyc deleted file mode 100644 index e4f0420c5f3b895d5cd24f591fa52ce6fc56aadb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/__pycache__/multi_tensor_apply.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/cpu_lion.py b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/cpu_lion.py deleted file mode 100644 index 03342a3fcd34fdb852434a97e03fa7f99426c7d0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/cpu_lion.py +++ /dev/null @@ -1,127 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from cpuinfo import get_cpu_info -from deepspeed.utils import logger -from deepspeed.utils.logging import should_log_le -from deepspeed.ops.op_builder import CPULionBuilder - - -class DeepSpeedCPULion(torch.optim.Optimizer): - optimizer_id = 0 - - def __init__(self, model_params, lr=1e-3, betas=(0.9, 0.999), weight_decay=0, fp32_optimizer_states=True): - """Fast vectorized implementation of Lion optimizer on CPU: - - See Symbolic Discovery of Optimization Algorithms (https://doi.org/10.48550/arXiv.2302.06675). - - .. note:: - We recommend using our `config - `_ - to allow :meth:`deepspeed.initialize` to build this optimizer - for you. - - - Arguments: - model_params (iterable): iterable of parameters to optimize or dicts defining - parameter groups. - lr (float, optional): learning rate. (default: 1e-3) - betas (Tuple[float, float], optional): coefficients used for computing - running averages of gradient and its square. (default: (0.9, 0.999)) - weight_decay (float, optional): weight decay (L2 penalty) (default: 0) - full_precision_optimizer_states: creates momentum and variance in full precision regardless of - the precision of the parameters (default: True) - """ - - default_args = dict(lr=lr, betas=betas, weight_decay=weight_decay) - super(DeepSpeedCPULion, self).__init__(model_params, default_args) - - cpu_info = get_cpu_info() - self.cpu_vendor = cpu_info["vendor_id_raw"].lower() if "vendor_id_raw" in cpu_info else "unknown" - if "amd" in self.cpu_vendor: - for group_id, group in enumerate(self.param_groups): - for param_id, p in enumerate(group['params']): - if p.dtype == torch.half: - logger.warning("FP16 params for CPULion may not work on AMD CPUs") - break - else: - continue - break - - self.opt_id = DeepSpeedCPULion.optimizer_id - DeepSpeedCPULion.optimizer_id = DeepSpeedCPULion.optimizer_id + 1 - self.fp32_optimizer_states = fp32_optimizer_states - self.ds_opt_lion = CPULionBuilder().load() - - self.ds_opt_lion.create_lion(self.opt_id, lr, betas[0], betas[1], weight_decay, should_log_le("info")) - - def __del__(self): - # need to destroy the C++ object explicitly to avoid a memory leak when deepspeed.initialize - # is used multiple times in the same process (notebook or pytest worker) - self.ds_opt_lion.destroy_lion(self.opt_id) - - def __setstate__(self, state): - super(DeepSpeedCPULion, self).__setstate__(state) - for group in self.param_groups: - group.setdefault('amsgrad', False) - - @torch.no_grad() - def step(self, closure=None): - """Update the model parameters. - - .. note:: - This method will be called internally by ZeRO-Offload. DeepSpeed - users should still use ``engine.step()`` as shown in the - `Getting Started - `_ guide. - - Args: - closure (callable, optional): closure to compute the loss. - Defaults to ``None``. - - Returns: - loss: if ``closure`` is provided. Otherwise ``None``. - """ - - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - # intended device for step - device = torch.device('cpu') - - for group_id, group in enumerate(self.param_groups): - for param_id, p in enumerate(group['params']): - - if p.grad is None: - continue - - assert p.device == device, f"CPULion param is on {p.device} and must be 'cpu', make " \ - "sure you enabled 'offload_optimizer': 'cpu' in your ZeRO config." - - state = self.state[p] - # State initialization - if len(state) == 0: - #print(f'group {group_id} param {param_id} = {p.numel()}') - state['step'] = 0 - - #use full precision by default unless self.fp32_optimizer_states is off - state_dtype = torch.float if self.fp32_optimizer_states else p.dtype - - # gradient momentums - state['exp_avg'] = torch.zeros_like(p.data, dtype=state_dtype, device=device) - #memory_format=torch.preserve_format) - # gradient variances - state['exp_avg_sq'] = torch.zeros_like(p.data, dtype=state_dtype, device=device) - #memory_format=torch.preserve_format) - - state['step'] += 1 - beta1, beta2 = group['betas'] - - self.ds_opt_lion.lion_update(self.opt_id, state['step'], group['lr'], beta1, beta2, - group['weight_decay'], p.data, p.grad.data, state['exp_avg']) - return loss diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/fused_lion.py b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/fused_lion.py deleted file mode 100644 index 7332a7f96361a1a05d770d945a90efe6b24ef217..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/fused_lion.py +++ /dev/null @@ -1,131 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -This file is modified from fused_adam.py -""" - -import torch -from .multi_tensor_apply import MultiTensorApply - -multi_tensor_applier = MultiTensorApply(2048 * 32) -from deepspeed.accelerator import get_accelerator -from deepspeed.ops.op_builder import FusedLionBuilder - - -class FusedLion(torch.optim.Optimizer): - """Implements Lion algorithm. - - Currently GPU-only. - - Arguments: - params (iterable): iterable of parameters to optimize or dicts defining - parameter groups. - lr (float, optional): learning rate. (default: 1e-3) - betas (Tuple[float, float], optional): coefficients used for computing - running averages of gradient and its square. (default: (0.9, 0.999)) - weight_decay (float, optional): weight decay (L2 penalty) (default: 0) - set_grad_none (bool, optional): whether set grad to None when zero_grad() - method is called. (default: True) - - .. _Symbolic Discovery of Optimization Algorithms: - https://doi.org/10.48550/arXiv.2302.06675 - """ - - def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), weight_decay=0., set_grad_none=True): - - defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay) - super(FusedLion, self).__init__(params, defaults) - self.set_grad_none = set_grad_none - - fused_lion_cuda = FusedLionBuilder().load() - # Skip buffer - self._dummy_overflow_buf = get_accelerator().IntTensor([0]) - self.multi_tensor_lion = fused_lion_cuda.multi_tensor_lion - - def zero_grad(self): - if self.set_grad_none: - for group in self.param_groups: - for p in group['params']: - p.grad = None - else: - super(FusedLion, self).zero_grad() - - def step(self, closure=None, grads=None, output_params=None, scale=None, grad_norms=None, grad_scaler=None): - """Performs a single optimization step. - - Arguments: - closure (callable, optional): A closure that reevaluates the model - and returns the loss. - - The remaining arguments are deprecated, and are only retained (for the moment) for error-checking purposes. - """ - if any(p is not None for p in [grads, output_params, scale, grad_norms]): - raise RuntimeError('FusedLion has been updated.') - loss = None - if closure is not None: - loss = closure() - - for group in self.param_groups: - if len(group['params']) == 0: - continue - beta1, beta2 = group['betas'] - - # assume same step across group now to simplify things - # per parameter step can be easily support by making it tensor, or pass list into kernel - if 'step' not in group: - group['step'] = 0 - - # create lists for multi-tensor apply - g_16, p_16, m_16 = [], [], [] - g_bf, p_bf, m_bf = [], [], [] - g_32, p_32, m_32 = [], [], [] - - for p in group['params']: - if p.grad is None: - continue - if p.grad.data.is_sparse: - raise NotImplementedError('FusedLion does not support sparse gradients') - - state = self.state[p] - # State initialization - if len(state) == 0: - # DeepSpeed ZeRO 3 processes each subgroup a time, so we need to keep tracking step count for each tensor separately. - # While this is not an issue for ZeRO 1 & 2, since they apply a single optimization step to the whole param group at the same time. - # In order to keep backward compatibility for the existing checkpoints, we use group['state'] to initialize state['step'] if it exists. - state['step'] = group.get('step', 0) - # Exponential moving average of gradient values - state['exp_avg'] = torch.zeros_like(p.data) - - if p.dtype == torch.float16: - g_16.append(p.grad.data) - p_16.append(p.data) - m_16.append(state['exp_avg']) - elif p.dtype == torch.bfloat16: - g_bf.append(p.grad) - p_bf.append(p) - m_bf.append(state['exp_avg']) - elif p.dtype == torch.float32: - g_32.append(p.grad.data) - p_32.append(p.data) - m_32.append(state['exp_avg']) - else: - raise RuntimeError('FusedLion only support fp16, bf16 and fp32.') - - if len(g_16) > 0: - state['step'] += 1 - multi_tensor_applier(self.multi_tensor_lion, self._dummy_overflow_buf, [g_16, p_16, m_16], group['lr'], - beta1, beta2, state['step'], group['weight_decay']) - - if len(g_bf) > 0: - state['step'] += 1 - multi_tensor_applier(self.multi_tensor_lion, self._dummy_overflow_buf, [g_bf, p_bf, m_bf], group['lr'], - beta1, beta2, state['step'], group['weight_decay']) - - if len(g_32) > 0: - state['step'] += 1 - multi_tensor_applier(self.multi_tensor_lion, self._dummy_overflow_buf, [g_32, p_32, m_32], group['lr'], - beta1, beta2, state['step'], group['weight_decay']) - - return loss diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/multi_tensor_apply.py b/venv/lib/python3.10/site-packages/deepspeed/ops/lion/multi_tensor_apply.py deleted file mode 100644 index 0ba228505cef747eea4fec62f3e68707fa4daa0c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/lion/multi_tensor_apply.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Copyright NVIDIA/apex -This file is adapted from NVIDIA/apex, commit a109f85 -""" - - -class MultiTensorApply(object): - - def __init__(self, chunk_size): - self.chunk_size = chunk_size - - def __call__(self, op, noop_flag_buffer, tensor_lists, *args): - return op(self.chunk_size, noop_flag_buffer, tensor_lists, *args) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__init__.py deleted file mode 100644 index afe48159933c75045e9445a577d545b7ae260073..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__init__.py +++ /dev/null @@ -1,53 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import sys -import os -import pkgutil -import importlib - -from .builder import get_default_compute_capabilities, OpBuilder - -# Do not remove, required for abstract accelerator to detect if we have a deepspeed or 3p op_builder -__deepspeed__ = True - -# List of all available op builders from deepspeed op_builder -try: - import deepspeed.ops.op_builder # noqa: F401 # type: ignore - op_builder_dir = "deepspeed.ops.op_builder" -except ImportError: - op_builder_dir = "op_builder" - -__op_builders__ = [] - -this_module = sys.modules[__name__] - - -def builder_closure(member_name): - if op_builder_dir == "op_builder": - # during installation time cannot get builder due to torch not installed, - # return closure instead - def _builder(): - from deepspeed.accelerator import get_accelerator - builder = get_accelerator().create_op_builder(member_name) - return builder - - return _builder - else: - # during runtime, return op builder class directly - from deepspeed.accelerator import get_accelerator - builder = get_accelerator().get_op_builder(member_name) - return builder - - -# reflect builder names and add builder closure, such as 'TransformerBuilder()' creates op builder wrt current accelerator -for _, module_name, _ in pkgutil.iter_modules([os.path.dirname(this_module.__file__)]): - if module_name != 'all_ops' and module_name != 'builder': - module = importlib.import_module(f".{module_name}", package=op_builder_dir) - for member_name in module.__dir__(): - if member_name.endswith('Builder') and member_name != "OpBuilder" and member_name != "CUDAOpBuilder": - # assign builder name to variable with same name - # the following is equivalent to i.e. TransformerBuilder = "TransformerBuilder" - this_module.__dict__[member_name] = builder_closure(member_name) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 31b16fcdde12b70ca256c9e1f56b0eed8de2c20e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/all_ops.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/all_ops.cpython-310.pyc deleted file mode 100644 index 8e2ea343b24aaf34644b72419379c9b8414a3a18..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/all_ops.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/async_io.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/async_io.cpython-310.pyc deleted file mode 100644 index a4b378bdce86424372bfdca95a044880f47fbe38..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/async_io.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/builder.cpython-310.pyc deleted file mode 100644 index c42dd6b1a31c5935651f779f3bf1793480084a2b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/builder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adagrad.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adagrad.cpython-310.pyc deleted file mode 100644 index 179c01c29b836af966e6660a6d204e1460d324e7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adagrad.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adam.cpython-310.pyc deleted file mode 100644 index 41ba016eccad83a5f43c5ab0799aead24863562a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_lion.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_lion.cpython-310.pyc deleted file mode 100644 index ed07f53a5f3ac67bfb052a7bb098075a02f75a82..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_lion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/dc.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/dc.cpython-310.pyc deleted file mode 100644 index 02e6d63f4c80c8b90e3977391e693a2dfb3047df..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/dc.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/evoformer_attn.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/evoformer_attn.cpython-310.pyc deleted file mode 100644 index 9c937d3760c41a85a6ac6623241a16214e743ba7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/evoformer_attn.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fp_quantizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fp_quantizer.cpython-310.pyc deleted file mode 100644 index 6a1e7cbd9b2b9ae57698bb559f3a20e350ff15ed..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fp_quantizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_adam.cpython-310.pyc deleted file mode 100644 index d38644f7b9abb5a2c4aa0140d6d552ab06d4c4e0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lamb.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lamb.cpython-310.pyc deleted file mode 100644 index 4060487e61f2364d39d7d1c55a43fa6b4e55b7e1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lamb.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lion.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lion.cpython-310.pyc deleted file mode 100644 index 163133ab38feca34de1f7f2bc9c401dc44e743c9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/gds.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/gds.cpython-310.pyc deleted file mode 100644 index 14b3c1a68427f1ff28ad6164beec96e1ecffb178..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/gds.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/inference_core_ops.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/inference_core_ops.cpython-310.pyc deleted file mode 100644 index e154874ef49c7189e04b8b29725276042ec8400e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/inference_core_ops.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/inference_cutlass_builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/inference_cutlass_builder.cpython-310.pyc deleted file mode 100644 index 79056246654e7e796876e47883c15ba1319052fe..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/inference_cutlass_builder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/quantizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/quantizer.cpython-310.pyc deleted file mode 100644 index b76fe296e68a7e7725be39e9f3dac9f44f0e9301..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/quantizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_ops.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_ops.cpython-310.pyc deleted file mode 100644 index 37e310c96336647289aa3f26a01cc628814ebdb4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_ops.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_utils.cpython-310.pyc deleted file mode 100644 index bb642b4d541fabef240bb807a82d67ad0ff3299e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/random_ltd.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/random_ltd.cpython-310.pyc deleted file mode 100644 index d193d3ffda639f93a5fc2e9564c52733d66c42f4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/random_ltd.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/sparse_attn.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/sparse_attn.cpython-310.pyc deleted file mode 100644 index 9fb0a063861b02d89496018168d1c38f3492158a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/sparse_attn.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/spatial_inference.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/spatial_inference.cpython-310.pyc deleted file mode 100644 index fb695b224f3113b9a6ffc1745084e1cb095b98bf..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/spatial_inference.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/stochastic_transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/stochastic_transformer.cpython-310.pyc deleted file mode 100644 index 59b8e29b17d49867fe2a0db0269e36779f029ec4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/stochastic_transformer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/transformer.cpython-310.pyc deleted file mode 100644 index 7b3a30fd82c041fa79668543032be1256f31ecfa..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/transformer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/transformer_inference.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/transformer_inference.cpython-310.pyc deleted file mode 100644 index 819193b74ad8b078f2ac03abe77bddcfe9e802ac..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/transformer_inference.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index f667907ddc531df5e04ebeb874f09211d5952c99..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/all_ops.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/all_ops.py deleted file mode 100644 index ff11ca180072f7dc72918b417b5e6e44eb53e3c4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/all_ops.py +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import pkgutil -import importlib -try: - # during installation time accelerator is visible, otherwise return deepspeed.accelerator - from accelerator import get_accelerator -except ImportError: - from deepspeed.accelerator import get_accelerator - -# List of all available ops - -# reflect all builder names into __op_builders__ -op_builder_dir = get_accelerator().op_builder_dir() -op_builder_module = importlib.import_module(op_builder_dir) -__op_builders__ = [] - -for _, module_name, _ in pkgutil.iter_modules([os.path.dirname(op_builder_module.__file__)]): - # avoid self references - if module_name != 'all_ops' and module_name != 'builder': - module = importlib.import_module("{}.{}".format(op_builder_dir, module_name)) - for member_name in module.__dir__(): - if member_name.endswith('Builder'): - # append builder to __op_builders__ list - builder = get_accelerator().create_op_builder(member_name) - __op_builders__.append(builder) - -ALL_OPS = {op.name: op for op in __op_builders__ if op is not None} -accelerator_name = get_accelerator()._name diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/async_io.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/async_io.py deleted file mode 100644 index f59cc6810c6f04a9c9b323093b9bfa202294f9ee..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/async_io.py +++ /dev/null @@ -1,113 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import shutil -import subprocess - -from .builder import TorchCPUOpBuilder - - -class AsyncIOBuilder(TorchCPUOpBuilder): - BUILD_VAR = "DS_BUILD_AIO" - NAME = "async_io" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.aio.{self.NAME}_op' - - def lib_sources(self): - src_list = [ - 'csrc/aio/py_lib/deepspeed_py_io_handle.cpp', 'csrc/aio/py_lib/deepspeed_py_aio.cpp', - 'csrc/aio/py_lib/deepspeed_py_aio_handle.cpp', 'csrc/aio/py_lib/deepspeed_aio_thread.cpp', - 'csrc/aio/common/deepspeed_aio_utils.cpp', 'csrc/aio/common/deepspeed_aio_common.cpp', - 'csrc/aio/common/deepspeed_aio_types.cpp', 'csrc/aio/py_lib/deepspeed_cpu_op.cpp', - 'csrc/aio/py_lib/deepspeed_aio_op_desc.cpp', 'csrc/aio/py_lib/deepspeed_py_copy.cpp', - 'csrc/aio/py_lib/deepspeed_pin_tensor.cpp' - ] - return src_list - - def sources(self): - return self.lib_sources() + ['csrc/aio/py_lib/py_ds_aio.cpp'] - - def include_paths(self): - import torch - if self.build_for_cpu: - CUDA_INCLUDE = [] - elif not self.is_rocm_pytorch(): - CUDA_INCLUDE = [os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")] - else: - CUDA_INCLUDE = [ - os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include"), - os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include", "rocrand"), - os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include", "hiprand"), - ] - return ['csrc/aio/py_lib', 'csrc/aio/common'] + CUDA_INCLUDE - - def cxx_args(self): - # -O0 for improved debugging, since performance is bound by I/O - args = super().cxx_args() - import torch - TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2]) - if not (TORCH_MAJOR >= 2 and TORCH_MINOR >= 1): - args.remove('-std=c++17') - args.append('-std=c++14') - args += ['-Wall', '-O0', '-shared', '-fPIC', '-Wno-reorder'] - return args - - def extra_ldflags(self): - if self.build_for_cpu: - return ['-fopenmp'] - - import torch.utils.cpp_extension - CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME - if CUDA_HOME is None: - ldflags = ['-laio'] # the ROCM case - else: - CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64") - ldflags = [f'-L{CUDA_HOME}', f'-L{CUDA_LIB64}', '-laio', '-lcuda', '-lcudart'] - return ldflags - - def check_for_libaio_pkg(self): - libs = dict( - dpkg=["-l", "libaio-dev", "apt"], - pacman=["-Q", "libaio", "pacman"], - rpm=["-q", "libaio-devel", "yum"], - ) - - found = False - for pkgmgr, data in libs.items(): - flag, lib, tool = data - path = shutil.which(pkgmgr) - if path is not None: - cmd = [pkgmgr, flag, lib] - result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - if result.wait() == 0: - found = True - else: - self.warning(f"{self.NAME}: please install the {lib} package with {tool}") - break - return found - - def is_compatible(self, verbose=False): - # Check for the existence of libaio by using distutils - # to compile and link a test program that calls io_submit, - # which is a function provided by libaio that is used in the async_io op. - # If needed, one can define -I and -L entries in CFLAGS and LDFLAGS - # respectively to specify the directories for libaio.h and libaio.so. - aio_compatible = self.has_function('io_submit', ('aio', )) - if verbose and not aio_compatible: - self.warning(f"{self.NAME} requires the dev libaio .so object and headers but these were not found.") - - # Check for the libaio package via known package managers - # to print suggestions on which package to install. - self.check_for_libaio_pkg() - - self.warning( - "If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found." - ) - return super().is_compatible(verbose) and aio_compatible diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/builder.py deleted file mode 100644 index aef7d50358ddfed809c1a7744db4a37f16037e8c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/builder.py +++ /dev/null @@ -1,861 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import re -import sys -import time -import importlib -from pathlib import Path -import subprocess -import shlex -import shutil -import tempfile -import distutils.ccompiler -import distutils.log -import distutils.sysconfig -from distutils.errors import CompileError, LinkError -from abc import ABC, abstractmethod -from typing import List - -YELLOW = '\033[93m' -END = '\033[0m' -WARNING = f"{YELLOW} [WARNING] {END}" - -DEFAULT_TORCH_EXTENSION_PATH = "/tmp/torch_extensions" -DEFAULT_COMPUTE_CAPABILITIES = "6.0;6.1;7.0" - -try: - import torch -except ImportError: - print(f"{WARNING} unable to import torch, please install it if you want to pre-compile any deepspeed ops.") -else: - TORCH_MAJOR = int(torch.__version__.split('.')[0]) - TORCH_MINOR = int(torch.__version__.split('.')[1]) - - -class MissingCUDAException(Exception): - pass - - -class CUDAMismatchException(Exception): - pass - - -def installed_cuda_version(name=""): - import torch.utils.cpp_extension - cuda_home = torch.utils.cpp_extension.CUDA_HOME - if cuda_home is None: - raise MissingCUDAException("CUDA_HOME does not exist, unable to compile CUDA op(s)") - # Ensure there is not a cuda version mismatch between torch and nvcc compiler - output = subprocess.check_output([cuda_home + "/bin/nvcc", "-V"], universal_newlines=True) - output_split = output.split() - release_idx = output_split.index("release") - release = output_split[release_idx + 1].replace(',', '').split(".") - # Ignore patch versions, only look at major + minor - cuda_major, cuda_minor = release[:2] - return int(cuda_major), int(cuda_minor) - - -def get_default_compute_capabilities(): - compute_caps = DEFAULT_COMPUTE_CAPABILITIES - # Update compute capability according to: https://en.wikipedia.org/wiki/CUDA#GPUs_supported - import torch.utils.cpp_extension - if torch.utils.cpp_extension.CUDA_HOME is not None: - if installed_cuda_version()[0] == 11: - if installed_cuda_version()[1] >= 0: - compute_caps += ";8.0" - if installed_cuda_version()[1] >= 1: - compute_caps += ";8.6" - if installed_cuda_version()[1] >= 8: - compute_caps += ";9.0" - elif installed_cuda_version()[0] == 12: - compute_caps += ";8.0;8.6;9.0" - if installed_cuda_version()[1] >= 8: - compute_caps += ";10.0;12.0" - return compute_caps - - -# list compatible minor CUDA versions - so that for example pytorch built with cuda-11.0 can be used -# to build deepspeed and system-wide installed cuda 11.2 -cuda_minor_mismatch_ok = { - 10: ["10.0", "10.1", "10.2"], - 11: ["11.0", "11.1", "11.2", "11.3", "11.4", "11.5", "11.6", "11.7", "11.8"], - 12: ["12.0", "12.1", "12.2", "12.3", "12.4", "12.5", "12.6", "12.8", "12.9"], # There is no CUDATk 12.7 -} - - -def assert_no_cuda_mismatch(name=""): - cuda_major, cuda_minor = installed_cuda_version(name) - sys_cuda_version = f'{cuda_major}.{cuda_minor}' - torch_cuda_version = ".".join(torch.version.cuda.split('.')[:2]) - # This is a show-stopping error, should probably not proceed past this - if sys_cuda_version != torch_cuda_version: - if (cuda_major in cuda_minor_mismatch_ok and sys_cuda_version in cuda_minor_mismatch_ok[cuda_major] - and torch_cuda_version in cuda_minor_mismatch_ok[cuda_major]): - print(f"Installed CUDA version {sys_cuda_version} does not match the " - f"version torch was compiled with {torch.version.cuda} " - "but since the APIs are compatible, accepting this combination") - return True - elif os.getenv("DS_SKIP_CUDA_CHECK", "0") == "1": - print( - f"{WARNING} DeepSpeed Op Builder: Installed CUDA version {sys_cuda_version} does not match the " - f"version torch was compiled with {torch.version.cuda}." - "Detected `DS_SKIP_CUDA_CHECK=1`: Allowing this combination of CUDA, but it may result in unexpected behavior." - ) - return True - raise CUDAMismatchException( - f">- DeepSpeed Op Builder: Installed CUDA version {sys_cuda_version} does not match the " - f"version torch was compiled with {torch.version.cuda}, unable to compile " - "cuda/cpp extensions without a matching cuda version.") - return True - - -class OpBuilder(ABC): - _rocm_version = None - _rocm_gpu_arch = None - _rocm_wavefront_size = None - _is_rocm_pytorch = None - _is_sycl_enabled = None - _loaded_ops = {} - - def __init__(self, name): - self.name = name - self.jit_mode = False - self.build_for_cpu = False - self.enable_bf16 = False - self.error_log = None - - @abstractmethod - def absolute_name(self): - ''' - Returns absolute build path for cases where the op is pre-installed, e.g., deepspeed.ops.adam.cpu_adam - will be installed as something like: deepspeed/ops/adam/cpu_adam.so - ''' - pass - - @abstractmethod - def sources(self): - ''' - Returns list of source files for your op, relative to root of deepspeed package (i.e., DeepSpeed/deepspeed) - ''' - pass - - def hipify_extension(self): - pass - - def sycl_extension(self): - pass - - @staticmethod - def validate_torch_version(torch_info): - install_torch_version = torch_info['version'] - current_torch_version = ".".join(torch.__version__.split('.')[:2]) - if install_torch_version != current_torch_version: - raise RuntimeError("PyTorch version mismatch! DeepSpeed ops were compiled and installed " - "with a different version than what is being used at runtime. " - f"Please re-install DeepSpeed or switch torch versions. " - f"Install torch version={install_torch_version}, " - f"Runtime torch version={current_torch_version}") - - @staticmethod - def validate_torch_op_version(torch_info): - if not OpBuilder.is_rocm_pytorch(): - current_cuda_version = ".".join(torch.version.cuda.split('.')[:2]) - install_cuda_version = torch_info['cuda_version'] - if install_cuda_version != current_cuda_version: - raise RuntimeError("CUDA version mismatch! DeepSpeed ops were compiled and installed " - "with a different version than what is being used at runtime. " - f"Please re-install DeepSpeed or switch torch versions. " - f"Install CUDA version={install_cuda_version}, " - f"Runtime CUDA version={current_cuda_version}") - else: - current_hip_version = ".".join(torch.version.hip.split('.')[:2]) - install_hip_version = torch_info['hip_version'] - if install_hip_version != current_hip_version: - raise RuntimeError("HIP version mismatch! DeepSpeed ops were compiled and installed " - "with a different version than what is being used at runtime. " - f"Please re-install DeepSpeed or switch torch versions. " - f"Install HIP version={install_hip_version}, " - f"Runtime HIP version={current_hip_version}") - - @staticmethod - def is_rocm_pytorch(): - if OpBuilder._is_rocm_pytorch is not None: - return OpBuilder._is_rocm_pytorch - - _is_rocm_pytorch = False - try: - import torch - except ImportError: - pass - else: - if TORCH_MAJOR > 1 or (TORCH_MAJOR == 1 and TORCH_MINOR >= 5): - _is_rocm_pytorch = hasattr(torch.version, 'hip') and torch.version.hip is not None - if _is_rocm_pytorch: - from torch.utils.cpp_extension import ROCM_HOME - _is_rocm_pytorch = ROCM_HOME is not None - OpBuilder._is_rocm_pytorch = _is_rocm_pytorch - return OpBuilder._is_rocm_pytorch - - @staticmethod - def is_sycl_enabled(): - if OpBuilder._is_sycl_enabled is not None: - return OpBuilder._is_sycl_enabled - - _is_sycl_enabled = False - try: - result = subprocess.run(["c2s", "--version"], capture_output=True) - except: - pass - else: - _is_sycl_enabled = True - - OpBuilder._is_sycl_enabled = _is_sycl_enabled - return OpBuilder._is_sycl_enabled - - @staticmethod - def installed_rocm_version(): - if OpBuilder._rocm_version: - return OpBuilder._rocm_version - - ROCM_MAJOR = '0' - ROCM_MINOR = '0' - ROCM_VERSION_DEV_RAW = "" - if OpBuilder.is_rocm_pytorch(): - from torch.utils.cpp_extension import ROCM_HOME - rocm_ver_file = Path(ROCM_HOME).joinpath(".info/version") - if rocm_ver_file.is_file(): - with open(rocm_ver_file, 'r') as file: - ROCM_VERSION_DEV_RAW = file.read() - elif "rocm" in torch.__version__: - ROCM_VERSION_DEV_RAW = torch.__version__.split("rocm")[1] - if ROCM_VERSION_DEV_RAW != "": - ROCM_MAJOR = ROCM_VERSION_DEV_RAW.split('.')[0] - ROCM_MINOR = ROCM_VERSION_DEV_RAW.split('.')[1] - else: - # Look in /usr/include/rocm-version.h - rocm_ver_file = Path("/usr/include/rocm_version.h") - if rocm_ver_file.is_file(): - with open(rocm_ver_file, 'r') as file: - for ln in file.readlines(): - if "#define ROCM_VERSION_MAJOR" in ln: - ROCM_MAJOR = re.findall(r'\S+', ln)[2] - elif "#define ROCM_VERSION_MINOR" in ln: - ROCM_MINOR = re.findall(r'\S+', ln)[2] - if ROCM_MAJOR == '0': - assert False, "Could not detect ROCm version" - - OpBuilder._rocm_version = (int(ROCM_MAJOR), int(ROCM_MINOR)) - return OpBuilder._rocm_version - - @staticmethod - def get_rocm_gpu_arch(): - if OpBuilder._rocm_gpu_arch: - return OpBuilder._rocm_gpu_arch - rocm_info = Path("/opt/rocm/bin/rocminfo") - if (not rocm_info.is_file()): - rocm_info = Path("rocminfo") - rocm_gpu_arch_cmd = str(rocm_info) + " | grep -o -m 1 'gfx.*'" - try: - result = subprocess.check_output(rocm_gpu_arch_cmd, shell=True) - rocm_gpu_arch = result.decode('utf-8').strip() - except subprocess.CalledProcessError: - rocm_gpu_arch = "" - OpBuilder._rocm_gpu_arch = rocm_gpu_arch - return OpBuilder._rocm_gpu_arch - - @staticmethod - def get_rocm_wavefront_size(): - if OpBuilder._rocm_wavefront_size: - return OpBuilder._rocm_wavefront_size - - rocm_info = Path("/opt/rocm/bin/rocminfo") - if (not rocm_info.is_file()): - rocm_info = Path("rocminfo") - rocm_wavefront_size_cmd = str( - rocm_info) + " | grep -Eo -m1 'Wavefront Size:[[:space:]]+[0-9]+' | grep -Eo '[0-9]+'" - try: - result = subprocess.check_output(rocm_wavefront_size_cmd, shell=True) - rocm_wavefront_size = result.decode('utf-8').strip() - except subprocess.CalledProcessError: - rocm_wavefront_size = "32" - OpBuilder._rocm_wavefront_size = rocm_wavefront_size - return OpBuilder._rocm_wavefront_size - - def include_paths(self): - ''' - Returns list of include paths, relative to root of deepspeed package (i.e., DeepSpeed/deepspeed) - ''' - return [] - - def nvcc_args(self): - ''' - Returns optional list of compiler flags to forward to nvcc when building CUDA sources - ''' - return [] - - def cxx_args(self): - ''' - Returns optional list of compiler flags to forward to the build - ''' - return [] - - def is_compatible(self, verbose=False): - ''' - Check if all non-python dependencies are satisfied to build this op - ''' - return True - - def extra_ldflags(self): - return [] - - def has_function(self, funcname, libraries, library_dirs=None, verbose=False): - ''' - Test for existence of a function within a tuple of libraries. - - This is used as a smoke test to check whether a certain library is available. - As a test, this creates a simple C program that calls the specified function, - and then distutils is used to compile that program and link it with the specified libraries. - Returns True if both the compile and link are successful, False otherwise. - ''' - tempdir = None # we create a temporary directory to hold various files - filestderr = None # handle to open file to which we redirect stderr - oldstderr = None # file descriptor for stderr - try: - # Echo compile and link commands that are used. - if verbose: - distutils.log.set_verbosity(1) - - # Create a compiler object. - compiler = distutils.ccompiler.new_compiler(verbose=verbose) - - # Configure compiler and linker to build according to Python install. - distutils.sysconfig.customize_compiler(compiler) - - # Create a temporary directory to hold test files. - tempdir = tempfile.mkdtemp() - - # Define a simple C program that calls the function in question - prog = "void %s(void); int main(int argc, char** argv) { %s(); return 0; }" % (funcname, funcname) - - # Write the test program to a file. - filename = os.path.join(tempdir, 'test.c') - with open(filename, 'w') as f: - f.write(prog) - - # Redirect stderr file descriptor to a file to silence compile/link warnings. - if not verbose: - filestderr = open(os.path.join(tempdir, 'stderr.txt'), 'w') - oldstderr = os.dup(sys.stderr.fileno()) - os.dup2(filestderr.fileno(), sys.stderr.fileno()) - - # Workaround for behavior in distutils.ccompiler.CCompiler.object_filenames() - # Otherwise, a local directory will be used instead of tempdir - drive, driveless_filename = os.path.splitdrive(filename) - root_dir = driveless_filename[0] if os.path.isabs(driveless_filename) else '' - output_dir = os.path.join(drive, root_dir) - - # Attempt to compile the C program into an object file. - cflags = shlex.split(os.environ.get('CFLAGS', "")) - objs = compiler.compile([filename], output_dir=output_dir, extra_preargs=self.strip_empty_entries(cflags)) - - # Attempt to link the object file into an executable. - # Be sure to tack on any libraries that have been specified. - ldflags = shlex.split(os.environ.get('LDFLAGS', "")) - compiler.link_executable(objs, - os.path.join(tempdir, 'a.out'), - extra_preargs=self.strip_empty_entries(ldflags), - libraries=libraries, - library_dirs=library_dirs) - - # Compile and link succeeded - return True - - except CompileError: - return False - - except LinkError: - return False - - except: - return False - - finally: - # Restore stderr file descriptor and close the stderr redirect file. - if oldstderr is not None: - os.dup2(oldstderr, sys.stderr.fileno()) - if filestderr is not None: - filestderr.close() - - # Delete the temporary directory holding the test program and stderr files. - if tempdir is not None: - shutil.rmtree(tempdir) - - def strip_empty_entries(self, args): - ''' - Drop any empty strings from the list of compile and link flags - ''' - return [x for x in args if len(x) > 0] - - def cpu_arch(self): - try: - from cpuinfo import get_cpu_info - except ImportError as e: - cpu_info = self._backup_cpuinfo() - if cpu_info is None: - return "-march=native" - - try: - cpu_info = get_cpu_info() - except Exception as e: - self.warning(f"{self.name} attempted to use py-cpuinfo but failed (exception type: {type(e)}, {e}), " - "falling back to lscpu to get this information.") - cpu_info = self._backup_cpuinfo() - if cpu_info is None: - return "-march=native" - - if cpu_info['arch'].startswith('PPC_'): - # gcc does not provide -march on PowerPC, use -mcpu instead - return '-mcpu=native' - return '-march=native' - - def get_cuda_compile_flag(self): - try: - if not self.is_rocm_pytorch(): - assert_no_cuda_mismatch(self.name) - return "-D__ENABLE_CUDA__" - except MissingCUDAException: - print(f"{WARNING} {self.name} cuda is missing or is incompatible with installed torch, " - "only cpu ops can be compiled!") - return '-D__DISABLE_CUDA__' - return '-D__DISABLE_CUDA__' - - def _backup_cpuinfo(self): - # Construct cpu_info dict from lscpu that is similar to what py-cpuinfo provides - if not self.command_exists('lscpu'): - self.warning(f"{self.name} attempted to query 'lscpu' after failing to use py-cpuinfo " - "to detect the CPU architecture. 'lscpu' does not appear to exist on " - "your system, will fall back to use -march=native and non-vectorized execution.") - return None - result = subprocess.check_output(['lscpu']) - result = result.decode('utf-8').strip().lower() - - cpu_info = {} - cpu_info['arch'] = None - cpu_info['flags'] = "" - if 'genuineintel' in result or 'authenticamd' in result: - cpu_info['arch'] = 'X86_64' - if 'avx512' in result: - cpu_info['flags'] += 'avx512,' - elif 'avx512f' in result: - cpu_info['flags'] += 'avx512f,' - if 'avx2' in result: - cpu_info['flags'] += 'avx2' - elif 'ppc64le' in result: - cpu_info['arch'] = "PPC_" - - return cpu_info - - def simd_width(self): - try: - from cpuinfo import get_cpu_info - except ImportError as e: - cpu_info = self._backup_cpuinfo() - if cpu_info is None: - return '-D__SCALAR__' - - try: - cpu_info = get_cpu_info() - except Exception as e: - self.warning(f"{self.name} attempted to use py-cpuinfo but failed (exception type: {type(e)}, {e}), " - "falling back to lscpu to get this information.") - cpu_info = self._backup_cpuinfo() - if cpu_info is None: - return '-D__SCALAR__' - - if cpu_info['arch'] == 'X86_64': - if 'avx512' in cpu_info['flags'] or 'avx512f' in cpu_info['flags']: - return '-D__AVX512__' - elif 'avx2' in cpu_info['flags']: - return '-D__AVX256__' - return '-D__SCALAR__' - - def command_exists(self, cmd): - if '|' in cmd: - cmds = cmd.split("|") - else: - cmds = [cmd] - valid = False - for cmd in cmds: - safe_cmd = ["bash", "-c", f"type {cmd}"] - result = subprocess.Popen(safe_cmd, stdout=subprocess.PIPE) - valid = valid or result.wait() == 0 - - if not valid and len(cmds) > 1: - print(f"{WARNING} {self.name} requires one of the following commands '{cmds}', but it does not exist!") - elif not valid and len(cmds) == 1: - print(f"{WARNING} {self.name} requires the '{cmd}' command, but it does not exist!") - return valid - - def warning(self, msg): - self.error_log = f"{msg}" - print(f"{WARNING} {msg}") - - def deepspeed_src_path(self, code_path): - if os.path.isabs(code_path): - return code_path - else: - return os.path.join(Path(__file__).parent.parent.absolute(), code_path) - - def builder(self): - from torch.utils.cpp_extension import CppExtension - include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] - return CppExtension(name=self.absolute_name(), - sources=self.strip_empty_entries(self.sources()), - include_dirs=include_dirs, - extra_compile_args={'cxx': self.strip_empty_entries(self.cxx_args())}, - extra_link_args=self.strip_empty_entries(self.extra_ldflags())) - - def load(self, verbose=True): - if self.name in __class__._loaded_ops: - return __class__._loaded_ops[self.name] - - from deepspeed.git_version_info import installed_ops, torch_info, accelerator_name - from deepspeed.accelerator import get_accelerator - if installed_ops.get(self.name, False) and accelerator_name == get_accelerator()._name: - # Ensure the op we're about to load was compiled with the same - # torch/cuda versions we are currently using at runtime. - self.validate_torch_version(torch_info) - if torch.cuda.is_available() and isinstance(self, CUDAOpBuilder): - self.validate_torch_op_version(torch_info) - - op_module = importlib.import_module(self.absolute_name()) - __class__._loaded_ops[self.name] = op_module - return op_module - else: - return self.jit_load(verbose) - - def jit_load(self, verbose=True): - if not self.is_compatible(verbose): - raise RuntimeError( - f"Unable to JIT load the {self.name} op due to it not being compatible due to hardware/software issue. {self.error_log}" - ) - try: - import ninja # noqa: F401 # type: ignore - except ImportError: - raise RuntimeError(f"Unable to JIT load the {self.name} op due to ninja not being installed.") - - if isinstance(self, CUDAOpBuilder) and not self.is_rocm_pytorch(): - self.build_for_cpu = not torch.cuda.is_available() - - self.jit_mode = True - from torch.utils.cpp_extension import load - - start_build = time.time() - sources = [os.path.abspath(self.deepspeed_src_path(path)) for path in self.sources()] - extra_include_paths = [os.path.abspath(self.deepspeed_src_path(path)) for path in self.include_paths()] - - # Torch will try and apply whatever CCs are in the arch list at compile time, - # we have already set the intended targets ourselves we know that will be - # needed at runtime. This prevents CC collisions such as multiple __half - # implementations. Stash arch list to reset after build. - torch_arch_list = None - if "TORCH_CUDA_ARCH_LIST" in os.environ: - torch_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST") - os.environ["TORCH_CUDA_ARCH_LIST"] = "" - - nvcc_args = self.strip_empty_entries(self.nvcc_args()) - cxx_args = self.strip_empty_entries(self.cxx_args()) - - cxx_args.append("-UC10_USE_GLOG") - nvcc_args.append("-UC10_USE_GLOG") - if isinstance(self, CUDAOpBuilder): - if not self.build_for_cpu and self.enable_bf16: - cxx_args.append("-DBF16_AVAILABLE") - nvcc_args.append("-DBF16_AVAILABLE") - nvcc_args.append("-U__CUDA_NO_BFLOAT16_OPERATORS__") - nvcc_args.append("-U__CUDA_NO_BFLOAT162_OPERATORS__") - nvcc_args.append("-U__CUDA_NO_BFLOAT16_CONVERSIONS__") - - if self.is_rocm_pytorch(): - cxx_args.append("-D__HIP_PLATFORM_AMD__=1") - os.environ["PYTORCH_ROCM_ARCH"] = self.get_rocm_gpu_arch() - cxx_args.append('-DROCM_WAVEFRONT_SIZE=%s' % self.get_rocm_wavefront_size()) - - op_module = load(name=self.name, - sources=self.strip_empty_entries(sources), - extra_include_paths=self.strip_empty_entries(extra_include_paths), - extra_cflags=cxx_args, - extra_cuda_cflags=nvcc_args, - extra_ldflags=self.strip_empty_entries(self.extra_ldflags()), - with_cuda=True if (isinstance(self, CUDAOpBuilder) and not self.build_for_cpu) else None, - verbose=verbose) - - build_duration = time.time() - start_build - if verbose: - print(f"Time to load {self.name} op: {build_duration} seconds") - - # Reset arch list so we are not silently removing it for other possible use cases - if torch_arch_list: - os.environ["TORCH_CUDA_ARCH_LIST"] = torch_arch_list - - __class__._loaded_ops[self.name] = op_module - - return op_module - - -class CUDAOpBuilder(OpBuilder): - - def compute_capability_args(self, cross_compile_archs=None): - """ - Returns nvcc compute capability compile flags. - - 1. `TORCH_CUDA_ARCH_LIST` takes priority over `cross_compile_archs`. - 2. If neither is set default compute capabilities will be used - 3. Under `jit_mode` compute capabilities of all visible cards will be used plus PTX - - Format: - - - `TORCH_CUDA_ARCH_LIST` may use ; or whitespace separators. Examples: - - TORCH_CUDA_ARCH_LIST="6.1;7.5;8.6;9.0;10.0" pip install ... - TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6 9.0 10.0+PTX" pip install ... - - - `cross_compile_archs` uses ; separator. - - """ - ccs = [] - if self.jit_mode: - # Compile for underlying architectures since we know those at runtime - for i in range(torch.cuda.device_count()): - CC_MAJOR, CC_MINOR = torch.cuda.get_device_capability(i) - cc = f"{CC_MAJOR}.{CC_MINOR}" - if cc not in ccs: - ccs.append(cc) - ccs = sorted(ccs) - ccs[-1] += '+PTX' - else: - # Cross-compile mode, compile for various architectures - # env override takes priority - cross_compile_archs_env = os.environ.get('TORCH_CUDA_ARCH_LIST', None) - if cross_compile_archs_env is not None: - if cross_compile_archs is not None: - print( - f"{WARNING} env var TORCH_CUDA_ARCH_LIST={cross_compile_archs_env} overrides cross_compile_archs={cross_compile_archs}" - ) - cross_compile_archs = cross_compile_archs_env.replace(' ', ';') - else: - if cross_compile_archs is None: - cross_compile_archs = get_default_compute_capabilities() - ccs = cross_compile_archs.split(';') - - ccs = self.filter_ccs(ccs) - if len(ccs) == 0: - raise RuntimeError( - f"Unable to load {self.name} op due to no compute capabilities remaining after filtering") - - args = [] - self.enable_bf16 = True - for cc in ccs: - num = cc[0] + cc[1].split('+')[0] - args.append(f'-gencode=arch=compute_{num},code=sm_{num}') - if cc[1].endswith('+PTX'): - args.append(f'-gencode=arch=compute_{num},code=compute_{num}') - - if int(cc[0]) <= 7: - self.enable_bf16 = False - - return args - - def filter_ccs(self, ccs: List[str]): - """ - Prune any compute capabilities that are not compatible with the builder. Should log - which CCs have been pruned. - """ - return [cc.split('.') for cc in ccs] - - def version_dependent_macros(self): - # Fix from apex that might be relevant for us as well, related to https://github.com/NVIDIA/apex/issues/456 - version_ge_1_1 = [] - if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 0): - version_ge_1_1 = ['-DVERSION_GE_1_1'] - version_ge_1_3 = [] - if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2): - version_ge_1_3 = ['-DVERSION_GE_1_3'] - version_ge_1_5 = [] - if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 4): - version_ge_1_5 = ['-DVERSION_GE_1_5'] - return version_ge_1_1 + version_ge_1_3 + version_ge_1_5 - - def is_compatible(self, verbose=False): - return super().is_compatible(verbose) - - def builder(self): - try: - if not self.is_rocm_pytorch(): - assert_no_cuda_mismatch(self.name) - self.build_for_cpu = False - except MissingCUDAException: - self.build_for_cpu = True - - if self.build_for_cpu: - from torch.utils.cpp_extension import CppExtension as ExtensionBuilder - else: - from torch.utils.cpp_extension import CUDAExtension as ExtensionBuilder - include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] - compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} if self.build_for_cpu else \ - {'cxx': self.strip_empty_entries(self.cxx_args()), \ - 'nvcc': self.strip_empty_entries(self.nvcc_args())} - - if not self.build_for_cpu and self.enable_bf16: - compile_args['cxx'].append("-DBF16_AVAILABLE") - compile_args['nvcc'].append("-DBF16_AVAILABLE") - - if self.is_rocm_pytorch(): - compile_args['cxx'].append("-D__HIP_PLATFORM_AMD__=1") - #cxx compiler args are required to compile cpp files - compile_args['cxx'].append('-DROCM_WAVEFRONT_SIZE=%s' % self.get_rocm_wavefront_size()) - #nvcc compiler args are required to compile hip files - compile_args['nvcc'].append('-DROCM_WAVEFRONT_SIZE=%s' % self.get_rocm_wavefront_size()) - if self.get_rocm_gpu_arch(): - os.environ["PYTORCH_ROCM_ARCH"] = self.get_rocm_gpu_arch() - - cuda_ext = ExtensionBuilder(name=self.absolute_name(), - sources=self.strip_empty_entries(self.sources()), - include_dirs=include_dirs, - libraries=self.strip_empty_entries(self.libraries_args()), - extra_compile_args=compile_args, - extra_link_args=self.strip_empty_entries(self.extra_ldflags())) - - if self.is_rocm_pytorch(): - # hip converts paths to absolute, this converts back to relative - sources = cuda_ext.sources - curr_file = Path(__file__).parent.parent # ds root - for i in range(len(sources)): - src = Path(sources[i]) - if src.is_absolute(): - sources[i] = str(src.relative_to(curr_file)) - else: - sources[i] = str(src) - cuda_ext.sources = sources - return cuda_ext - - def hipify_extension(self): - if self.is_rocm_pytorch(): - from torch.utils.hipify import hipify_python - hipify_python.hipify( - project_directory=os.getcwd(), - output_directory=os.getcwd(), - header_include_dirs=self.include_paths(), - includes=[os.path.join(os.getcwd(), '*')], - extra_files=[os.path.abspath(s) for s in self.sources()], - show_detailed=True, - is_pytorch_extension=True, - hipify_extra_files_only=True, - ) - - def cxx_args(self): - if sys.platform == "win32": - return ['-O2'] - else: - return ['-O3', '-std=c++17', '-g', '-Wno-reorder'] - - def nvcc_args(self): - if self.build_for_cpu: - return [] - args = ['-O3'] - if self.is_rocm_pytorch(): - ROCM_MAJOR, ROCM_MINOR = self.installed_rocm_version() - args += [ - '-std=c++17', '-U__HIP_NO_HALF_OPERATORS__', '-U__HIP_NO_HALF_CONVERSIONS__', - '-U__HIP_NO_HALF2_OPERATORS__', - '-DROCM_VERSION_MAJOR=%s' % ROCM_MAJOR, - '-DROCM_VERSION_MINOR=%s' % ROCM_MINOR - ] - else: - try: - nvcc_threads = int(os.getenv("DS_NVCC_THREADS", "")) - if nvcc_threads <= 0: - raise ValueError("") - except ValueError: - nvcc_threads = min(os.cpu_count(), 8) - - cuda_major, cuda_minor = installed_cuda_version() - if cuda_major > 10: - if cuda_major == 12 and cuda_minor >= 5: - std_lib = '-std=c++20' - else: - std_lib = '-std=c++17' - else: - std_lib = '-std=c++14' - args += [ - '-allow-unsupported-compiler' if sys.platform == "win32" else '', '--use_fast_math', std_lib, - '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__', - f'--threads={nvcc_threads}' - ] - if os.environ.get('DS_DEBUG_CUDA_BUILD', '0') == '1': - args.append('--ptxas-options=-v') - args += self.compute_capability_args() - return args - - def libraries_args(self): - if self.build_for_cpu: - return [] - - if sys.platform == "win32": - return ['cublas', 'curand'] - else: - return [] - - -class TorchCPUOpBuilder(CUDAOpBuilder): - - def get_cuda_lib64_path(self): - import torch - if not self.is_rocm_pytorch(): - CUDA_LIB64 = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "lib64") - if not os.path.exists(CUDA_LIB64): - CUDA_LIB64 = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "lib") - else: - CUDA_LIB64 = os.path.join(torch.utils.cpp_extension.ROCM_HOME, "lib") - return CUDA_LIB64 - - def extra_ldflags(self): - if self.build_for_cpu: - return ['-fopenmp'] - - if not self.is_rocm_pytorch(): - ld_flags = ['-lcurand'] - if not self.build_for_cpu: - ld_flags.append(f'-L{self.get_cuda_lib64_path()}') - return ld_flags - - return [] - - def cxx_args(self): - args = [] - if not self.build_for_cpu: - CUDA_LIB64 = self.get_cuda_lib64_path() - - args += super().cxx_args() - args += [ - f'-L{CUDA_LIB64}', - '-lcudart', - '-lcublas', - '-g', - ] - - CPU_ARCH = self.cpu_arch() - SIMD_WIDTH = self.simd_width() - CUDA_ENABLE = self.get_cuda_compile_flag() - args += [ - CPU_ARCH, - '-fopenmp', - SIMD_WIDTH, - CUDA_ENABLE, - ] - - return args diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__init__.py deleted file mode 100644 index 7084db8469f1dae201796c44f911f539459e4757..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' - -from .comm import CCLCommBuilder, ShareMemCommBuilder -from .fused_adam import FusedAdamBuilder -from .cpu_adam import CPUAdamBuilder -from .no_impl import NotImplementedBuilder -from .async_io import AsyncIOBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index f3351925e39d041e72faa3dd8168f1d2ec174b5e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/async_io.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/async_io.cpython-310.pyc deleted file mode 100644 index 6489a0777eda56a12ae400415df66988830150a2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/async_io.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/builder.cpython-310.pyc deleted file mode 100644 index 96b8fdf538eacd2f4d0264bb9204e5704dc0e087..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/builder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/comm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/comm.cpython-310.pyc deleted file mode 100644 index 10c284beb50a5390d1b665812525d90e684a1a64..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/comm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/cpu_adam.cpython-310.pyc deleted file mode 100644 index 3f4b8f66a497f2469e8a10cfe293a2751e474048..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/cpu_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/fused_adam.cpython-310.pyc deleted file mode 100644 index 910a2f6beef0aafb3d9ee8e22e2dfee799d9ee45..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/fused_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/no_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/no_impl.cpython-310.pyc deleted file mode 100644 index a8d0955c878e510412665104f5a19f2fd34d04a5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/no_impl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/async_io.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/async_io.py deleted file mode 100644 index dcb9feabcfc3c7c3081d812fafb7d8816a031eb3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/async_io.py +++ /dev/null @@ -1,91 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import shutil -import subprocess - -from .builder import CPUOpBuilder - - -class AsyncIOBuilder(CPUOpBuilder): - BUILD_VAR = "DS_BUILD_AIO" - NAME = "async_io" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.aio.{self.NAME}_op' - - def lib_sources(self): - src_list = [ - 'csrc/aio/py_lib/deepspeed_py_io_handle.cpp', 'csrc/aio/py_lib/deepspeed_py_aio.cpp', - 'csrc/aio/py_lib/deepspeed_py_aio_handle.cpp', 'csrc/aio/py_lib/deepspeed_aio_thread.cpp', - 'csrc/aio/common/deepspeed_aio_utils.cpp', 'csrc/aio/common/deepspeed_aio_common.cpp', - 'csrc/aio/common/deepspeed_aio_types.cpp', 'csrc/aio/py_lib/deepspeed_cpu_op.cpp', - 'csrc/aio/py_lib/deepspeed_aio_op_desc.cpp', 'csrc/aio/py_lib/deepspeed_py_copy.cpp', - 'csrc/aio/py_lib/deepspeed_pin_tensor.cpp' - ] - return src_list - - def sources(self): - return self.lib_sources() + ['csrc/aio/py_lib/py_ds_aio.cpp'] - - def include_paths(self): - return ['csrc/aio/py_lib', 'csrc/aio/common'] - - def cxx_args(self): - # -O0 for improved debugging, since performance is bound by I/O - args = super().cxx_args() - import torch - TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2]) - if not (TORCH_MAJOR >= 2 and TORCH_MINOR >= 1): - args.remove('-std=c++17') - args.append('-std=c++14') - args += ['-Wall', '-O0', '-shared', '-fPIC', '-Wno-reorder'] - return args - - def extra_ldflags(self): - return ['-laio', '-fopenmp'] - - def check_for_libaio_pkg(self): - libs = dict( - dpkg=["-l", "libaio-dev", "apt"], - pacman=["-Q", "libaio", "pacman"], - rpm=["-q", "libaio-devel", "yum"], - ) - - found = False - for pkgmgr, data in libs.items(): - flag, lib, tool = data - path = shutil.which(pkgmgr) - if path is not None: - cmd = [pkgmgr, flag, lib] - result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - if result.wait() == 0: - found = True - else: - self.warning(f"{self.NAME}: please install the {lib} package with {tool}") - break - return found - - def is_compatible(self, verbose=False): - # Check for the existence of libaio by using distutils - # to compile and link a test program that calls io_submit, - # which is a function provided by libaio that is used in the async_io op. - # If needed, one can define -I and -L entries in CFLAGS and LDFLAGS - # respectively to specify the directories for libaio.h and libaio.so. - aio_compatible = self.has_function('io_submit', ('aio', )) - if verbose and not aio_compatible: - self.warning(f"{self.NAME} requires the dev libaio .so object and headers but these were not found.") - - # Check for the libaio package via known package managers - # to print suggestions on which package to install. - self.check_for_libaio_pkg() - - self.warning( - "If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found." - ) - return super().is_compatible(verbose) and aio_compatible diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/builder.py deleted file mode 100644 index d881842ad0b18654bc26cb3b54d345a4cd160da4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/builder.py +++ /dev/null @@ -1,40 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os - -try: - # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed - # if successful this also means we're doing a local install and not JIT compile path - from op_builder import __deepspeed__ # noqa: F401 # type: ignore - from op_builder.builder import OpBuilder -except ImportError: - from deepspeed.ops.op_builder.builder import OpBuilder - - -class CPUOpBuilder(OpBuilder): - - def builder(self): - from torch.utils.cpp_extension import CppExtension as ExtensionBuilder - include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] - compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} - - cpp_ext = ExtensionBuilder(name=self.absolute_name(), - sources=self.strip_empty_entries(self.sources()), - include_dirs=include_dirs, - libraries=self.strip_empty_entries(self.libraries_args()), - extra_compile_args=compile_args) - - return cpp_ext - - def cxx_args(self): - args = ['-O3', '-g', '-Wno-reorder'] - CPU_ARCH = self.cpu_arch() - SIMD_WIDTH = self.simd_width() - args += [CPU_ARCH, '-fopenmp', SIMD_WIDTH] - return args - - def libraries_args(self): - return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/comm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/comm.py deleted file mode 100644 index fec960b63b2ec16cb7ab6a92b509a5d97b901cef..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/comm.py +++ /dev/null @@ -1,71 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -from .builder import CPUOpBuilder - - -class CCLCommBuilder(CPUOpBuilder): - BUILD_VAR = "DS_BUILD_CCL_COMM" - NAME = "deepspeed_ccl_comm" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.comm.{self.NAME}_op' - - def sources(self): - return ['csrc/cpu/comm/ccl.cpp', 'csrc/cpu/comm/shm.cpp'] - - def include_paths(self): - includes = ['csrc/cpu/includes'] - return includes - - def cxx_args(self): - return ['-O2', '-fopenmp'] - - def is_compatible(self, verbose=False): - # TODO: add soft compatibility check for private binary release. - # a soft check, as in we know it can be trivially changed. - return super().is_compatible(verbose) - - def extra_ldflags(self): - ccl_root_path = os.environ.get("CCL_ROOT") - if ccl_root_path is None: - raise ValueError( - "Didn't find CCL_ROOT, install oneCCL from https://github.com/oneapi-src/oneCCL and source its environment variable" - ) - return [] - else: - return ['-lccl', f'-L{ccl_root_path}/lib'] - - -class ShareMemCommBuilder(CPUOpBuilder): - BUILD_VAR = "DS_BUILD_SHM_COMM" - NAME = "deepspeed_shm_comm" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.comm.{self.NAME}_op' - - def sources(self): - return ['csrc/cpu/comm/shm_interface.cpp', 'csrc/cpu/comm/shm.cpp'] - - def include_paths(self): - includes = ['csrc/cpu/includes'] - return includes - - def cxx_args(self): - return ['-O2', '-fopenmp'] - - def is_compatible(self, verbose=False): - # TODO: add soft compatibility check for private binary release. - # a soft check, as in we know it can be trivially changed. - return super().is_compatible(verbose) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/cpu_adam.py deleted file mode 100644 index 0c8438aea40d4ffd588a496844d692c9a17faaa6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/cpu_adam.py +++ /dev/null @@ -1,27 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import CPUOpBuilder - - -class CPUAdamBuilder(CPUOpBuilder): - BUILD_VAR = "DS_BUILD_CPU_ADAM" - NAME = "cpu_adam" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adam.{self.NAME}_op' - - def sources(self): - return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] - - def libraries_args(self): - args = super().libraries_args() - return args - - def include_paths(self): - return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/fused_adam.py deleted file mode 100644 index 34b43825b09024136afdc44e916349ea3f5ce5ad..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/fused_adam.py +++ /dev/null @@ -1,23 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import CPUOpBuilder - - -class FusedAdamBuilder(CPUOpBuilder): - BUILD_VAR = "DS_BUILD_FUSED_ADAM" - NAME = "fused_adam" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adam.{self.NAME}_op' - - def sources(self): - return ['csrc/cpu/adam/fused_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] - - def include_paths(self): - return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/no_impl.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/no_impl.py deleted file mode 100644 index 69d114a9f1c0b0defc482c1fb143c261fc466125..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu/no_impl.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import CPUOpBuilder - - -class NotImplementedBuilder(CPUOpBuilder): - BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" - NAME = "deepspeed_not_implemented" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.comm.{self.NAME}_op' - - def load(self, verbose=True): - raise ValueError("This op had not been implemented on CPU backend.") - - def sources(self): - return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_adagrad.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_adagrad.py deleted file mode 100644 index c05f7148895000c10cdfebef7461858a1f987fbb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_adagrad.py +++ /dev/null @@ -1,27 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import TorchCPUOpBuilder - - -class CPUAdagradBuilder(TorchCPUOpBuilder): - BUILD_VAR = "DS_BUILD_CPU_ADAGRAD" - NAME = "cpu_adagrad" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adagrad.{self.NAME}_op' - - def sources(self): - return ['csrc/adagrad/cpu_adagrad.cpp'] - - def libraries_args(self): - args = super().libraries_args() - return args - - def include_paths(self): - return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_adam.py deleted file mode 100644 index 7f4c0847a8c4f32b5f3cacd40f88eed8eea904f9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_adam.py +++ /dev/null @@ -1,27 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import TorchCPUOpBuilder - - -class CPUAdamBuilder(TorchCPUOpBuilder): - BUILD_VAR = "DS_BUILD_CPU_ADAM" - NAME = "cpu_adam" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adam.{self.NAME}_op' - - def sources(self): - return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] - - def libraries_args(self): - args = super().libraries_args() - return args - - def include_paths(self): - return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_lion.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_lion.py deleted file mode 100644 index 9a60d99773b31252c9d002a0cf3282185d91a07c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/cpu_lion.py +++ /dev/null @@ -1,27 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import TorchCPUOpBuilder - - -class CPULionBuilder(TorchCPUOpBuilder): - BUILD_VAR = "DS_BUILD_CPU_LION" - NAME = "cpu_lion" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.lion.{self.NAME}_op' - - def sources(self): - return ['csrc/lion/cpu_lion.cpp', 'csrc/lion/cpu_lion_impl.cpp'] - - def libraries_args(self): - args = super().libraries_args() - return args - - def include_paths(self): - return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/dc.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/dc.py deleted file mode 100644 index 15b25bf3393e45d3ea2936cf7ca57fee7c49adeb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/dc.py +++ /dev/null @@ -1,40 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import TorchCPUOpBuilder - - -class DeepCompileBuilder(TorchCPUOpBuilder): - BUILD_VAR = "DS_BUILD_DEEP_COMPILE" - NAME = "dc" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.{self.NAME}_op' - - def sources(self): - return [ - 'csrc/compile/deepcompile.cpp', 'csrc/compile/init.cpp', 'csrc/compile/z1.cpp', 'csrc/compile/z2.cpp', - 'csrc/compile/z3.cpp', 'csrc/compile/util.cpp' - ] - - def libraries_args(self): - args = super().libraries_args() - return args - - def include_paths(self): - import os - import torch - if self.build_for_cpu: - CUDA_INCLUDE = [] - elif not self.is_rocm_pytorch(): - CUDA_INCLUDE = [os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")] - else: - CUDA_INCLUDE = [ - os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include"), - ] - return ['csrc/includes', 'csrc/compile'] + CUDA_INCLUDE diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/evoformer_attn.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/evoformer_attn.py deleted file mode 100644 index 7f68ccf8729085028d445fb3d7f76d6f51556bfe..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/evoformer_attn.py +++ /dev/null @@ -1,93 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import CUDAOpBuilder, installed_cuda_version -import os - - -class EvoformerAttnBuilder(CUDAOpBuilder): - BUILD_VAR = "DS_BUILD_EVOFORMER_ATTN" - NAME = "evoformer_attn" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - self.cutlass_path = os.environ.get('CUTLASS_PATH') - - def absolute_name(self): - return f'deepspeed.ops.{self.NAME}_op' - - def extra_ldflags(self): - if not self.is_rocm_pytorch(): - return ['-lcurand'] - else: - return [] - - def sources(self): - src_dir = 'csrc/deepspeed4science/evoformer_attn' - return [f'{src_dir}/attention.cpp', f'{src_dir}/attention_back.cu', f'{src_dir}/attention_cu.cu'] - - def nvcc_args(self): - args = super().nvcc_args() - try: - import torch - except ImportError: - self.warning("Please install torch if trying to pre-compile kernels") - return args - major = torch.cuda.get_device_properties(0).major #ignore-cuda - minor = torch.cuda.get_device_properties(0).minor #ignore-cuda - args.append(f"-DGPU_ARCH={major}{minor}") - return args - - def is_compatible(self, verbose=False): - try: - import torch - except ImportError: - if verbose: - self.warning("Please install torch if trying to pre-compile kernels") - return False - if self.cutlass_path is None: - if verbose: - self.warning("Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH") - return False - if os.path.exists(f'{self.cutlass_path}/CHANGELOG.md'): - with open(f'{self.cutlass_path}/CHANGELOG.md', 'r') as f: - if '3.1.0' not in f.read(): - if verbose: - self.warning("Please use CUTLASS version >= 3.1.0") - return False - else: - # pip install nvidia-cutlass package - try: - import cutlass - except ImportError: - if verbose: - self.warning("Please pip install nvidia-cutlass if trying to pre-compile kernels") - return False - cutlass_major, cutlass_minor = cutlass.__version__.split('.')[:2] - cutlass_compatible = (int(cutlass_major) >= 3 and int(cutlass_minor) >= 1) - if not cutlass_compatible: - if verbose: - self.warning("Please use CUTLASS version >= 3.1.0") - return False - - cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda - if cuda_capability < 7: - if verbose: - self.warning("Please use a GPU with compute capability >= 7.0") - cuda_okay = False - if torch_cuda_major < 11 or sys_cuda_major < 11: - if verbose: - self.warning("Please use CUDA 11+") - cuda_okay = False - return super().is_compatible(verbose) and cuda_okay - - def include_paths(self): - includes = [f'{self.cutlass_path}/include', f'{self.cutlass_path}/tools/util/include'] - return includes diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fp_quantizer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fp_quantizer.py deleted file mode 100644 index 2b962ac2c1fea18d1024359623f3faf82228f253..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fp_quantizer.py +++ /dev/null @@ -1,120 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -try: - from packaging import version as pkg_version -except ImportError: - pkg_version = None - -from .builder import CUDAOpBuilder, installed_cuda_version - - -class FPQuantizerBuilder(CUDAOpBuilder): - BUILD_VAR = "DS_BUILD_FP_QUANTIZER" - NAME = "fp_quantizer" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.fp_quantizer.{self.NAME}_op' - - def is_compatible(self, verbose=False): - try: - import torch - except ImportError: - if verbose: - self.warning("Please install torch if trying to pre-compile inference kernels") - return False - - cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda - if cuda_capability < 8: - if verbose: - self.warning("NVIDIA Inference is only supported on Ampere and newer architectures") - cuda_okay = False - if cuda_capability >= 8: - if torch_cuda_major < 11 or sys_cuda_major < 11: - if verbose: - self.warning("On Ampere and higher architectures please use CUDA 11+") - cuda_okay = False - - try: - import triton - except ImportError: - if verbose: - self.warning( - f"please install triton==2.3.0, 2.3.1 or 3.0.0 if you want to use the FP Quantizer Kernels") - return False - - # triton 2.3.{0,1} and 3.0.0 are ok. - allowed_versions = ("2.3", "3.0", "3.1", "3.2") - if pkg_version: - allowed = (pkg_version.parse(v) for v in allowed_versions) - installed_triton = pkg_version.parse(triton.__version__) - triton_mismatch = all(installed_triton.major != a.major or installed_triton.minor != a.minor - for a in allowed) - else: - installed_triton = triton.__version__ - major, minor, _ = installed_triton.split(".") - allowed = (v.split(".") for v in allowed_versions) - triton_mismatch = all(major != v[0] or minor != v[1] for v in allowed) - - if triton_mismatch: - if verbose: - self.warning( - f"FP Quantizer is using an untested triton version ({installed_triton}), only 2.3.{0,1} and 3.0.0 are known to be compatible with these kernels" - ) - return False - - return super().is_compatible(verbose) and cuda_okay - - def filter_ccs(self, ccs): - ccs_retained = [] - ccs_pruned = [] - for cc in [cc.split('.') for cc in ccs]: - if int(cc[0]) >= 8: - ccs_retained.append(cc) - else: - ccs_pruned.append(cc) - if len(ccs_pruned) > 0: - self.warning(f"Filtered compute capabilities {ccs_pruned}") - return ccs_retained - - def sources(self): - return [ - "csrc/fp_quantizer/fp_quantize.cu", - "csrc/fp_quantizer/fp_quantize.cpp", - ] - - def extra_ldflags(self): - if not self.is_rocm_pytorch(): - return ['-lcurand'] - else: - return [] - - def include_paths(self): - return ['csrc/fp_quantizer/includes', 'csrc/includes'] - - @staticmethod - def get_default_quant_dtype(): - import torch - return torch.uint8 - - @staticmethod - def get_quant_range(q_bits=None): - if q_bits == 8: - return 480 - elif q_bits == 6: - return 28. - elif q_bits == 12: - return 510. - else: - assert (0), \ - "Please specify the right quantization range for the selected precision!" diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_adam.py deleted file mode 100644 index ac6e4eeaaea5d9b2c2ee70de3d4261c6348abe94..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_adam.py +++ /dev/null @@ -1,37 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import CUDAOpBuilder - -import sys - - -class FusedAdamBuilder(CUDAOpBuilder): - BUILD_VAR = "DS_BUILD_FUSED_ADAM" - NAME = "fused_adam" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adam.{self.NAME}_op' - - def sources(self): - return ['csrc/adam/fused_adam_frontend.cpp', 'csrc/adam/multi_tensor_adam.cu'] - - def include_paths(self): - return ['csrc/includes', 'csrc/adam'] - - def cxx_args(self): - args = super().cxx_args() - return args + self.version_dependent_macros() - - def nvcc_args(self): - nvcc_flags = ['-O3'] + self.version_dependent_macros() - if not self.is_rocm_pytorch(): - nvcc_flags.extend( - ['-allow-unsupported-compiler' if sys.platform == "win32" else '', '-lineinfo', '--use_fast_math'] + - self.compute_capability_args()) - return nvcc_flags diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_lamb.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_lamb.py deleted file mode 100644 index f0cb557706b3f7e6d4ab73d7a999a70aebecbeb3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_lamb.py +++ /dev/null @@ -1,40 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import CUDAOpBuilder - -import sys - - -class FusedLambBuilder(CUDAOpBuilder): - BUILD_VAR = 'DS_BUILD_FUSED_LAMB' - NAME = "fused_lamb" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.lamb.{self.NAME}_op' - - def sources(self): - return ['csrc/lamb/fused_lamb_cuda.cpp', 'csrc/lamb/fused_lamb_cuda_kernel.cu'] - - def include_paths(self): - return ['csrc/includes'] - - def cxx_args(self): - args = super().cxx_args() - return args + self.version_dependent_macros() - - def nvcc_args(self): - nvcc_flags = ['-O3'] + self.version_dependent_macros() - if self.is_rocm_pytorch(): - ROCM_MAJOR, ROCM_MINOR = self.installed_rocm_version() - nvcc_flags += ['-DROCM_VERSION_MAJOR=%s' % ROCM_MAJOR, '-DROCM_VERSION_MINOR=%s' % ROCM_MINOR] - else: - nvcc_flags.extend( - ['-allow-unsupported-compiler' if sys.platform == "win32" else '', '-lineinfo', '--use_fast_math'] + - self.compute_capability_args()) - return nvcc_flags diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_lion.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_lion.py deleted file mode 100644 index b900a8f2369df14d65a73235d360db21583ee4ab..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/fused_lion.py +++ /dev/null @@ -1,37 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import CUDAOpBuilder - -import sys - - -class FusedLionBuilder(CUDAOpBuilder): - BUILD_VAR = "DS_BUILD_FUSED_LION" - NAME = "fused_lion" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.lion.{self.NAME}_op' - - def sources(self): - return ['csrc/lion/fused_lion_frontend.cpp', 'csrc/lion/multi_tensor_lion.cu'] - - def include_paths(self): - return ['csrc/includes', 'csrc/lion'] - - def cxx_args(self): - args = super().cxx_args() - return args + self.version_dependent_macros() - - def nvcc_args(self): - nvcc_flags = ['-O3'] + self.version_dependent_macros() - if not self.is_rocm_pytorch(): - nvcc_flags.extend( - ['-allow-unsupported-compiler' if sys.platform == "win32" else '', '-lineinfo', '--use_fast_math'] + - self.compute_capability_args()) - return nvcc_flags diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/gds.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/gds.py deleted file mode 100644 index d815ccea387a78a0a6256b6f981b81346a3c230b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/gds.py +++ /dev/null @@ -1,66 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -from .async_io import AsyncIOBuilder - - -class GDSBuilder(AsyncIOBuilder): - BUILD_VAR = "DS_BUILD_GDS" - NAME = "gds" - - def __init__(self): - super().__init__() - - def absolute_name(self): - return f'deepspeed.ops.gds.{self.NAME}_op' - - def lib_sources(self): - src_list = ['csrc/gds/py_lib/deepspeed_py_gds_handle.cpp', 'csrc/gds/py_lib/deepspeed_gds_op.cpp'] - return super().lib_sources() + src_list - - def sources(self): - return self.lib_sources() + ['csrc/gds/py_lib/py_ds_gds.cpp'] - - def cxx_args(self): - return super().cxx_args() + ['-lcufile'] - - def include_paths(self): - import torch - CUDA_INCLUDE = [os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")] - return ['csrc/aio/py_lib', 'csrc/aio/common'] + CUDA_INCLUDE - - def extra_ldflags(self): - return super().extra_ldflags() + ['-lcufile'] - - def is_compatible(self, verbose=False): - if self.is_rocm_pytorch(): - if verbose: - self.warning(f'{self.NAME} is not compatible with ROCM') - return False - - try: - import torch.utils.cpp_extension - except ImportError: - if verbose: - self.warning("Please install torch if trying to pre-compile GDS") - return False - - CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME - if CUDA_HOME is None: - if verbose: - self.warning("Please install torch CUDA if trying to pre-compile GDS with CUDA") - return False - - CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64") - gds_compatible = self.has_function(funcname="cuFileDriverOpen", - libraries=("cufile", ), - library_dirs=( - CUDA_HOME, - CUDA_LIB64, - ), - verbose=verbose) - - return gds_compatible and super().is_compatible(verbose) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__init__.py deleted file mode 100644 index 5ad1b9a7f891eae0ceec8c26109ba2a235ef6349..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' - -from .cpu_adam import CPUAdamBuilder -from .fused_adam import FusedAdamBuilder -from .transformer_inference import InferenceBuilder -from .no_impl import NotImplementedBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 5398ff40de82847f3354ade3e38478276e47fdeb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/builder.cpython-310.pyc deleted file mode 100644 index 690a1ae30990ba5595a96aba83f3ecbfb79fc955..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/builder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/cpu_adam.cpython-310.pyc deleted file mode 100644 index efdd52aadd0a9fd4de4c6ed3a70e8871e274dbdc..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/cpu_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fp_quantizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fp_quantizer.cpython-310.pyc deleted file mode 100644 index 49284bf77d00a456973111371bf943c1290374d4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fp_quantizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fused_adam.cpython-310.pyc deleted file mode 100644 index 40ed757197be177e85b2d09582ba16ad321a5a85..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fused_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/no_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/no_impl.cpython-310.pyc deleted file mode 100644 index 11843bd9d64b245718c12f44ab31095737d50cd7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/no_impl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/transformer_inference.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/transformer_inference.cpython-310.pyc deleted file mode 100644 index 023127cfea9bd52a20960b766640f6b23541022b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/transformer_inference.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/builder.py deleted file mode 100644 index 11e710a8ee4833bc66e82c4ed3241f7fe8a2e977..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/builder.py +++ /dev/null @@ -1,38 +0,0 @@ -# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os - -try: - # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed - # if successful this also means we're doing a local install and not JIT compile path - from op_builder import __deepspeed__ # noqa: F401 # type: ignore - from op_builder.builder import OpBuilder -except ImportError: - from deepspeed.ops.op_builder.builder import OpBuilder - - -class CPUOpBuilder(OpBuilder): - - def builder(self): - from torch.utils.cpp_extension import CppExtension as ExtensionBuilder - include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] - compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} - - cpp_ext = ExtensionBuilder(name=self.absolute_name(), - sources=self.strip_empty_entries(self.sources()), - include_dirs=include_dirs, - libraries=self.strip_empty_entries(self.libraries_args()), - extra_compile_args=compile_args) - - return cpp_ext - - def cxx_args(self): - args = ['-O3', '-g', '-Wno-reorder'] - return args - - def libraries_args(self): - return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/cpu_adam.py deleted file mode 100644 index 58eea2698ebb9eb6a8269a527a902792f1c7a6bc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/cpu_adam.py +++ /dev/null @@ -1,28 +0,0 @@ -# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import CPUOpBuilder - - -class CPUAdamBuilder(CPUOpBuilder): - BUILD_VAR = "DS_BUILD_CPU_ADAM" - NAME = "cpu_adam" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adam.{self.NAME}_op' - - def sources(self): - return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] - - def libraries_args(self): - args = super().libraries_args() - return args - - def include_paths(self): - return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/fp_quantizer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/fp_quantizer.py deleted file mode 100644 index c74affb55045a167f8dd3bb51b54f1487a815587..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/fp_quantizer.py +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright (c) 2024 Habana Labs, Ltd. an Intel Company -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -try: - # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed - # if successful this also means we're doing a local install and not JIT compile path - from op_builder import __deepspeed__ # noqa: F401 # type: ignore - from op_builder.builder import OpBuilder -except ImportError: - from deepspeed.ops.op_builder.builder import OpBuilder - - -class FPQuantizerBuilder(OpBuilder): - BUILD_VAR = "DS_BUILD_FP_QUANTIZER" - NAME = "fp_quantizer" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.fp_quantizer.{self.NAME}_op' - - def sources(self): - return [] - - def load(self, verbose=True): - return FPQuantizer - - @staticmethod - def get_default_quant_dtype(): - return torch.float8_e4m3fn - - @staticmethod - def get_quant_range(q_bits=None): - import habana_frameworks.torch.utils.experimental as htexp - if htexp._get_device_type() == htexp.synDeviceType.synDeviceGaudi2: - dtype = torch.float8_e4m3fnuz - else: - dtype = torch.float8_e4m3fn - return torch.finfo(dtype).max - - -class FPQuantizer: - CUDA_IMPL = False - - @classmethod - def selective_dequantize(cls, val_q, scales, indexes, group_size, q_mantisa_bits, q_exponent_bits): - assert False, "Selective dequantize isn't implemented for HPU!" - - @classmethod - def dequantize(cls, fp_out, input_q, scale, group_size, q_mantisa_bits, q_exponent_bits): - orig_shape = fp_out.shape - orig_dtype = fp_out.dtype - dequant_out = torch.ops.hpu.cast_from_fp8(input_q, (1.0 / scale), orig_dtype).view(orig_shape) - fp_out.copy_(dequant_out) - return fp_out - - @classmethod - def quantize(cls, out, val, scale, group_size, stochastic_rounding, q_bits, q_mantisa_bits): - assert q_bits == 8, "Quantize on HPU only supports quantization to FP8" - assert q_mantisa_bits == 3, "Quantize on HPU only supports q_mantissa_bits = 3" - assert out.dtype.is_floating_point, "Quantization on HPU is only to float dtypes" - - num_groups, group_size = out.shape - - # Reshape the tensor - val_reshaped = val.view(num_groups, group_size).float() - # Calculate the scale - max_vals = val_reshaped.abs().max(dim=1, keepdim=True)[0] - q_range = torch.finfo(out.dtype).max - tmp_scale = q_range / max_vals - scale.copy_(tmp_scale) - # Copy quantized - quant, _ = torch.ops.hpu.cast_to_fp8_v2(val_reshaped, scale, stochastic_rounding, dtype=out.dtype) - out.copy_(quant) - - return out - - @classmethod - def get_scales(cls, out, num_groups): - return out diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/fused_adam.py deleted file mode 100644 index 5acb121668e37d1f7b8b660b831d35ae863d2a1a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/fused_adam.py +++ /dev/null @@ -1,105 +0,0 @@ -# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -try: - # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed - # if successful this also means we're doing a local install and not JIT compile path - from op_builder import __deepspeed__ # noqa: F401 # type: ignore - from op_builder.builder import OpBuilder -except ImportError: - from deepspeed.ops.op_builder.builder import OpBuilder - -try: - import torch - import math -except ImportError as e: - pass - - -class HPUFusedAdam: - htcore = None - is_lazy_mode = None - - @staticmethod - def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, - bias_correction, weight_decay, *args): - - if HPUFusedAdam.htcore is None: - from habana_frameworks.torch import core as htcore - from habana_frameworks.torch.utils.internal import is_lazy - HPUFusedAdam.htcore = htcore - HPUFusedAdam.is_lazy_mode = is_lazy() - - htcore = HPUFusedAdam.htcore - - htcore.step_closure._mark_step_if_lazy() - step_size = lr - if bias_correction: - bias_correction1 = 1.0 - pow(beta1, step) - bias_correction2 = 1.0 - pow(beta2, step) - step_size = step_size * math.sqrt(bias_correction2) / bias_correction1 - - neg_step = -step_size - neg_step_t = (torch.tensor([neg_step], dtype=torch.float, - requires_grad=False).to(tensor_lists[1][0].dtype).to(tensor_lists[1][0].device, - non_blocking=True)) - - weight_decay = weight_decay if adam_w_mode else 0 - - # since lr is fed into the kernel as tensor, perform the scalar multiplication of wd here - # NOTE: TODO if lr is updated every step, then we need to convert it as tensor and - # perform weight decay unconditonally. - modified_wd = 1.0 - weight_decay * lr - - if HPUFusedAdam.is_lazy_mode: - torch.ops.hpu.optimizer_adamw( - tensor_lists[0], - tensor_lists[1], - tensor_lists[2], - tensor_lists[3], - neg_step_t, - beta1, - beta2, - epsilon, - modified_wd, - ) - else: - modified_wd_t = (torch.tensor([modified_wd], dtype=torch.float, requires_grad=False).to( - tensor_lists[1][0].dtype).to(tensor_lists[1][0].device, non_blocking=True)) - torch.ops.hpu.optimizer_adamw( - tensor_lists[0], - tensor_lists[1], - tensor_lists[2], - tensor_lists[3], - neg_step_t, - beta1, - beta2, - epsilon, - modified_wd_t, - modified_wd != 1.0, - ) - - htcore.step_closure._mark_step_if_lazy() - - -class FusedAdamBuilder(OpBuilder): - BUILD_VAR = "DS_BUILD_FUSED_ADAM" - NAME = "fused_adam" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adam.{self.NAME}_op' - - def sources(self): - return [] - - def include_paths(self): - return [] - - def load(self, verbose=True): - return HPUFusedAdam diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/no_impl.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/no_impl.py deleted file mode 100644 index 140d65b48defbea439483b63fd545e52de92209e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/no_impl.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import CPUOpBuilder - - -class NotImplementedBuilder(CPUOpBuilder): - BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" - NAME = "deepspeed_not_implemented" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.comm.{self.NAME}_op' - - def load(self, verbose=True): - raise ValueError("This op had not been implemented on HPU backend.") - - def sources(self): - return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/transformer_inference.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/transformer_inference.py deleted file mode 100644 index e397c99200ec9155c933bc0691050b1f1e85cf7b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/hpu/transformer_inference.py +++ /dev/null @@ -1,39 +0,0 @@ -# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 -import importlib - -# DeepSpeed Team - -try: - # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed - # if successful this also means we're doing a local install and not JIT compile path - from op_builder import __deepspeed__ # noqa: F401 - from op_builder.builder import OpBuilder -except ImportError: - from deepspeed.ops.op_builder.builder import OpBuilder - - -class InferenceBuilder(OpBuilder): - BUILD_VAR = "DS_BUILD_TRANSFORMER_INFERENCE" - NAME = "transformer_inference" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=self.NAME) - - def absolute_name(self): - return f"deepspeed.ops.transformer.inference.{self.NAME}_op" - - def sources(self): - return [] - - def load(self, verbose=True): - if self.name in __class__._loaded_ops: - return __class__._loaded_ops[self.name] - - from deepspeed.git_version_info import installed_ops # noqa: F401 - if installed_ops.get(self.name, False): - op_module = importlib.import_module(self.absolute_name()) - __class__._loaded_ops[self.name] = op_module - return op_module diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/inference_core_ops.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/inference_core_ops.py deleted file mode 100644 index b6665ebb76180c533f9700db222fca45314ecd19..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/inference_core_ops.py +++ /dev/null @@ -1,98 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os - -from .builder import CUDAOpBuilder, installed_cuda_version - - -class InferenceCoreBuilder(CUDAOpBuilder): - BUILD_VAR = "DS_BUILD_INFERENCE_CORE_OPS" - NAME = "inference_core_ops" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.inference.v2.kernels{self.NAME}' - - def is_compatible(self, verbose=False): - try: - import torch - except ImportError: - if verbose: - self.warning("Please install torch if trying to pre-compile inference kernels") - return False - - cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda - if cuda_capability < 6: - if verbose: - self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") - cuda_okay = False - if cuda_capability >= 8: - if torch_cuda_major < 11 or sys_cuda_major < 11: - if verbose: - self.warning("On Ampere and higher architectures please use CUDA 11+") - cuda_okay = False - return super().is_compatible(verbose) and cuda_okay - - def filter_ccs(self, ccs): - ccs_retained = [] - ccs_pruned = [] - for cc in [cc.split('.') for cc in ccs]: - if int(cc[0]) >= 6: - ccs_retained.append(cc) - else: - ccs_pruned.append(cc) - if len(ccs_pruned) > 0: - self.warning(f"Filtered compute capabilities {ccs_pruned}") - return ccs_retained - - def get_prefix(self): - ds_path = self.deepspeed_src_path("deepspeed") - return "deepspeed" if os.path.isdir(ds_path) else ".." - - def sources(self): - sources = [ - "inference/v2/kernels/core_ops/core_ops.cpp", - "inference/v2/kernels/core_ops/bias_activations/bias_activation.cpp", - "inference/v2/kernels/core_ops/bias_activations/bias_activation_cuda.cu", - "inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm.cpp", - "inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm_cuda.cu", - "inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm.cpp", - "inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm_cuda.cu", - "inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels.cpp", - "inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels_cuda.cu", - "inference/v2/kernels/core_ops/cuda_linear/linear_kernels.cpp", - "inference/v2/kernels/core_ops/cuda_linear/linear_kernels_cuda.cu", - ] - - prefix = self.get_prefix() - sources = [os.path.join(prefix, src) for src in sources] - return sources - - def extra_ldflags(self): - return [] - - def include_paths(self): - sources = [ - 'inference/v2/kernels/core_ops/bias_activations', - 'inference/v2/kernels/core_ops/blas_kernels', - 'inference/v2/kernels/core_ops/cuda_layer_norm', - 'inference/v2/kernels/core_ops/cuda_rms_norm', - 'inference/v2/kernels/core_ops/gated_activations', - 'inference/v2/kernels/core_ops/cuda_linear', - 'inference/v2/kernels/includes', - ] - - prefix = self.get_prefix() - sources = [os.path.join(prefix, src) for src in sources] - - return sources diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/inference_cutlass_builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/inference_cutlass_builder.py deleted file mode 100644 index a4a607288ca857f6812fdd925413493822bfdcb9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/inference_cutlass_builder.py +++ /dev/null @@ -1,95 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -import os - -from .builder import CUDAOpBuilder, installed_cuda_version - - -class InferenceCutlassBuilder(CUDAOpBuilder): - BUILD_VAR = "DS_BUILD_CUTLASS_OPS" - NAME = "cutlass_ops" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.inference.v2.kernels.cutlass_ops.{self.NAME}' - - def is_compatible(self, verbose=False): - try: - import torch - except ImportError: - if verbose: - self.warning("Please install torch if trying to pre-compile inference kernels") - return False - - cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda - if cuda_capability < 6: - if verbose: - self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") - cuda_okay = False - if cuda_capability >= 8: - if torch_cuda_major < 11 or sys_cuda_major < 11: - if verbose: - self.warning("On Ampere and higher architectures please use CUDA 11+") - cuda_okay = False - return super().is_compatible(verbose) and cuda_okay - - def filter_ccs(self, ccs): - ccs_retained = [] - ccs_pruned = [] - for cc in [cc.split('.') for cc in ccs]: - if int(cc[0]) >= 8: - # Only support Ampere and newer - ccs_retained.append(cc) - else: - ccs_pruned.append(cc) - if len(ccs_pruned) > 0: - self.warning(f"Filtered compute capabilities {ccs_pruned}") - return ccs_retained - - def get_prefix(self): - ds_path = self.deepspeed_src_path("deepspeed") - return "deepspeed" if os.path.isdir(ds_path) else ".." - - def sources(self): - sources = [ - "inference/v2/kernels/cutlass_ops/cutlass_ops.cpp", - "inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm.cu", - "inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm.cu", - ] - - prefix = self.get_prefix() - sources = [os.path.join(prefix, src) for src in sources] - return sources - - def extra_ldflags(self): - import dskernels - lib_path = dskernels.library_path() - prefix = self.get_prefix() - lib_path = os.path.join(prefix, lib_path) - lib_path = self.deepspeed_src_path(lib_path) - - args = [f'-L{lib_path}', '-ldeepspeedft'] - if self.jit_load: - args.append(f'-Wl,-rpath,{lib_path}') - return args - - def include_paths(self): - sources = [ - 'inference/v2/kernels/includes', - 'inference/v2/kernels/cutlass_ops/mixed_gemm', - 'inference/v2/kernels/cutlass_ops/moe_gemm', - 'inference/v2/kernels/cutlass_ops/shared_resources/', - ] - - prefix = self.get_prefix() - sources = [os.path.join(prefix, src) for src in sources] - return sources diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__init__.py deleted file mode 100644 index db12afbbf20ec7a5329b95110dbceaf5e8607d10..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Copyright (c) 2024 Cambricon Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' - -# MLU related operators will be added in the future. -from .no_impl import NotImplementedBuilder -from .cpu_adagrad import CPUAdagradBuilder -from .cpu_adam import CPUAdamBuilder -from .fused_adam import FusedAdamBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 1ef60702ac2ebb5661c99104be0932a984bbb036..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/builder.cpython-310.pyc deleted file mode 100644 index 39e4b026a6da1f48c4596fea31aa1795dc2331a7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/builder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adagrad.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adagrad.cpython-310.pyc deleted file mode 100644 index e430f3f9f67bd361a8c4f9d05cf543a11b90d144..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adagrad.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adam.cpython-310.pyc deleted file mode 100644 index bb948533db205fb5915702a716946b58feddcd5f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/fused_adam.cpython-310.pyc deleted file mode 100644 index 7a29a1de8bf10090a68678dac91b262522fa7c8a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/fused_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/no_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/no_impl.cpython-310.pyc deleted file mode 100644 index 123ffdb541aa216f2bc18664cf92c080b08ea94a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/no_impl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/builder.py deleted file mode 100644 index 17b9723ffcc1e3c81ea2e594d20b352b0a2a5bdf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/builder.py +++ /dev/null @@ -1,35 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Copyright (c) 2024 Cambricon Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -try: - # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed - # if successful this also means we're doing a local install and not JIT compile path - from op_builder import __deepspeed__ # noqa: F401 # type: ignore - from op_builder.builder import OpBuilder -except ImportError: - from deepspeed.ops.op_builder.builder import OpBuilder - - -class MLUOpBuilder(OpBuilder): - - def builder(self): - from torch.utils.cpp_extension import CppExtension as ExtensionBuilder - - compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} - - cpp_ext = ExtensionBuilder(name=self.absolute_name(), - sources=self.strip_empty_entries(self.sources()), - include_dirs=self.strip_empty_entries(self.include_paths()), - libraries=self.strip_empty_entries(self.libraries_args()), - extra_compile_args=compile_args) - - return cpp_ext - - def cxx_args(self): - return ['-O3', '-g', '-Wno-reorder'] - - def libraries_args(self): - return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/cpu_adagrad.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/cpu_adagrad.py deleted file mode 100644 index 68b7bbe514eea6d7714a6f8b6135ebd622a1b7f9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/cpu_adagrad.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Copyright (c) 2024 Cambricon Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import MLUOpBuilder - - -class CPUAdagradBuilder(MLUOpBuilder): - BUILD_VAR = "DS_BUILD_CPU_ADAGRAD" - NAME = "cpu_adagrad" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adagrad.{self.NAME}_op' - - def sources(self): - return ['csrc/adagrad/cpu_adagrad.cpp'] - - def include_paths(self): - return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/cpu_adam.py deleted file mode 100644 index b3c8e476bf3960c3ff467f362a72947cf73babf8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/cpu_adam.py +++ /dev/null @@ -1,28 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Copyright (c) 2024 Cambricon Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import MLUOpBuilder - - -class CPUAdamBuilder(MLUOpBuilder): - BUILD_VAR = "DS_BUILD_CPU_ADAM" - NAME = "cpu_adam" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adam.{self.NAME}_op' - - def sources(self): - return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] - - def libraries_args(self): - args = super().libraries_args() - return args - - def include_paths(self): - return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/fused_adam.py deleted file mode 100644 index 2343e55568c489e1ff5d2354e88b676889c7211c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/fused_adam.py +++ /dev/null @@ -1,43 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Copyright (c) 2024 Cambricon Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import MLUOpBuilder - -try: - import torch -except ImportError as e: - pass - - -class MLUFusedAdam: - - @staticmethod - def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, - bias_correction, weight_decay, *args): - - torch.ops.torch_mlu.fused_adam(noop_flag_buffer, tensor_lists[0], tensor_lists[1], tensor_lists[2], - tensor_lists[3], lr, beta1, beta2, epsilon, step, adam_w_mode, bias_correction, - weight_decay) - - -class FusedAdamBuilder(MLUOpBuilder): - BUILD_VAR = "DS_BUILD_FUSED_ADAM" - NAME = "fused_adam" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adam.{self.NAME}_op' - - def sources(self): - return [] - - def include_paths(self): - return [] - - def load(self, verbose=True): - return MLUFusedAdam diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/no_impl.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/no_impl.py deleted file mode 100644 index 375c148b4a5e06e53fb87cddbdf988708d425995..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/mlu/no_impl.py +++ /dev/null @@ -1,34 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Copyright (c) 2024 Cambricon Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import MLUOpBuilder - - -class NotImplementedBuilder(MLUOpBuilder): - BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" - NAME = "deepspeed_not_implemented" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.comm.{self.NAME}_op' - - def load(self, verbose=True): - raise ValueError("This op had not been implemented on MLU backend.") - - def sources(self): - return [] - - def cxx_args(self): - return [] - - def extra_ldflags(self): - return [] - - def include_paths(self): - return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__init__.py deleted file mode 100644 index e512e25a2a29e1b463bdc373722a64c755b89098..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' - -from .fused_adam import FusedAdamBuilder -from .async_io import AsyncIOBuilder -from .no_impl import NotImplementedBuilder -from .cpu_adam import CPUAdamBuilder -from .cpu_adagrad import CPUAdagradBuilder -from .cpu_lion import CPULionBuilder -from .inference import InferenceBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 163b776004947a0fcd43aaaece695ab6663a1c88..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/async_io.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/async_io.cpython-310.pyc deleted file mode 100644 index 22b411e53b4e03b9733a7900b4971907c7188d9f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/async_io.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/builder.cpython-310.pyc deleted file mode 100644 index 2b63f9280cbe007d5420f751b607d975d2114c0d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/builder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adagrad.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adagrad.cpython-310.pyc deleted file mode 100644 index 682d69e5c95ffb79af54327a274424ceac2362e6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adagrad.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adam.cpython-310.pyc deleted file mode 100644 index 5ddc97d0c66f39282c6d390ee5c8ce11cd2e8444..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_lion.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_lion.cpython-310.pyc deleted file mode 100644 index 482f1b0df3694d15814abb805d3077ed2d4e7f91..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_lion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/fused_adam.cpython-310.pyc deleted file mode 100644 index 75b92c9246cd28a18a8aeda5e1fa456c422c9bd4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/fused_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/inference.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/inference.cpython-310.pyc deleted file mode 100644 index 6ec55a3d8a4658f45b24bcbbf787e07f1d82bdea..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/inference.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/no_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/no_impl.cpython-310.pyc deleted file mode 100644 index 64a1bc262eab3ef11b59ae46b902533387945a91..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/__pycache__/no_impl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/async_io.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/async_io.py deleted file mode 100644 index 9a574ffe2dfbd2b52f42296460e348c6f619207c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/async_io.py +++ /dev/null @@ -1,103 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import shutil -import subprocess - -from .builder import NPUOpBuilder - - -class AsyncIOBuilder(NPUOpBuilder): - BUILD_VAR = "DS_BUILD_AIO" - NAME = "async_io" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.aio.{self.NAME}_op' - - def sources(self): - return [ - 'csrc/aio/py_lib/deepspeed_py_copy.cpp', 'csrc/aio/py_lib/py_ds_aio.cpp', - 'csrc/aio/py_lib/deepspeed_py_aio.cpp', 'csrc/aio/py_lib/deepspeed_py_aio_handle.cpp', - 'csrc/aio/py_lib/deepspeed_aio_thread.cpp', 'csrc/aio/common/deepspeed_aio_utils.cpp', - 'csrc/aio/common/deepspeed_aio_common.cpp', 'csrc/aio/common/deepspeed_aio_types.cpp', - 'csrc/aio/py_lib/deepspeed_pin_tensor.cpp' - ] - - def include_paths(self): - args = super().include_paths() - args += ['csrc/aio/py_lib', 'csrc/aio/common'] - return args - - def cxx_args(self): - args = super().cxx_args() - # -O0 for improved debugging, since performance is bound by I/O - CPU_ARCH = self.cpu_arch() - SIMD_WIDTH = self.simd_width() - import torch # Keep this import here to avoid errors when building DeepSpeed wheel without torch installed - TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2]) - if TORCH_MAJOR >= 2 and TORCH_MINOR >= 1: - CPP_STD = '-std=c++17' - else: - CPP_STD = '-std=c++14' - return args + [ - '-g', - '-Wall', - '-O0', - CPP_STD, - '-shared', - '-fPIC', - '-Wno-reorder', - CPU_ARCH, - '-fopenmp', - SIMD_WIDTH, - '-laio', - ] - - def extra_ldflags(self): - args = super().extra_ldflags() - return args + ['-laio'] - - def check_for_libaio_pkg(self): - libs = dict( - dpkg=["-l", "libaio-dev", "apt"], - pacman=["-Q", "libaio", "pacman"], - rpm=["-q", "libaio-devel", "yum"], - ) - - found = False - for pkgmgr, data in libs.items(): - flag, lib, tool = data - path = shutil.which(pkgmgr) - if path is not None: - cmd = [pkgmgr, flag, lib] - result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - if result.wait() == 0: - found = True - else: - self.warning(f"{self.NAME}: please install the {lib} package with {tool}") - break - return found - - def is_compatible(self, verbose=False): - # Check for the existence of libaio by using distutils - # to compile and link a test program that calls io_submit, - # which is a function provided by libaio that is used in the async_io op. - # If needed, one can define -I and -L entries in CFLAGS and LDFLAGS - # respectively to specify the directories for libaio.h and libaio.so. - aio_compatible = self.has_function('io_pgetevents', ('aio', )) - if verbose and not aio_compatible: - self.warning(f"{self.NAME} requires the dev libaio .so object and headers but these were not found.") - - # Check for the libaio package via known package managers - # to print suggestions on which package to install. - self.check_for_libaio_pkg() - - self.warning( - "If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found." - ) - return super().is_compatible(verbose) and aio_compatible diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/builder.py deleted file mode 100644 index 0dea2e78915e2a939d5c1839c2e2d88e948e10f1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/builder.py +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import re -import os -try: - import torch_npu -except ImportError as e: - pass - -try: - # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed - # if successful this also means we're doing a local install and not JIT compile path - from op_builder import __deepspeed__ # noqa: F401 # type: ignore - from op_builder.builder import OpBuilder -except ImportError: - from deepspeed.ops.op_builder.builder import OpBuilder - - -class NPUOpBuilder(OpBuilder): - _ascend_path = None - _torch_npu_path = None - _cann_version = None - - def __init__(self, name): - super().__init__(name) - self._ascend_path = self.installed_cann_path() - self._torch_npu_path = os.path.join(os.path.dirname(os.path.abspath(torch_npu.__file__))) - try: - self._cann_version = self.installed_cann_version(self.name) - except BaseException: - print(f"{self.name} ascend_cann is missing, npu ops cannot be compiled!") - - def cann_defs(self): - if self._cann_version: - return '-D__ENABLE_CANN__' - return '-D__DISABLE_CANN__' - - def installed_cann_path(self): - if "ASCEND_HOME_PATH" in os.environ or os.path.exists(os.environ["ASCEND_HOME_PATH"]): - return os.environ["ASCEND_HOME_PATH"] - return None - - def installed_cann_version(self, name=""): - ascend_path = self.installed_cann_path() - assert ascend_path is not None, "CANN_HOME does not exist, unable to compile NPU op(s)" - cann_version = "" - for dirpath, _, filenames in os.walk(os.path.realpath(ascend_path)): - if cann_version: - break - install_files = [file for file in filenames if re.match(r"ascend_.*_install\.info", file)] - if install_files: - filepath = os.path.join(dirpath, install_files[0]) - with open(filepath, "r") as f: - for line in f: - if line.find("version") != -1: - cann_version = line.strip().split("=")[-1] - break - return cann_version - - def include_paths(self): - paths = super().include_paths() - paths += [os.path.join(self._ascend_path, 'include'), os.path.join(self._torch_npu_path, 'include')] - return paths - - def cxx_args(self): - args = super().cxx_args() - args += ['-O3', '-std=c++17', '-g', '-Wno-reorder', '-fopenmp'] - args += ['-fstack-protector-all', '-Wl,-z,relro,-z,now,-z,noexecstack', '-Wl,--disable-new-dtags,--rpath'] - args += [ - self.cann_defs(), - self.cpu_arch(), - self.simd_width(), '-L' + os.path.join(self._ascend_path, 'lib64'), - '-L' + os.path.join(self._torch_npu_path, 'lib') - ] - return args - - def extra_ldflags(self): - flags = super().extra_ldflags() - flags += [ - '-L' + os.path.join(self._ascend_path, 'lib64'), '-lascendcl', - '-L' + os.path.join(self._torch_npu_path, 'lib'), '-ltorch_npu' - ] - return flags diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_adagrad.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_adagrad.py deleted file mode 100644 index 161bc82efe1ca01660fdeedd30079a8f10f1d269..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_adagrad.py +++ /dev/null @@ -1,25 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import NPUOpBuilder - - -class CPUAdagradBuilder(NPUOpBuilder): - BUILD_VAR = "DS_BUILD_CPU_ADAGRAD" - NAME = "cpu_adagrad" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adagrad.{self.NAME}_op' - - def sources(self): - return ['csrc/adagrad/cpu_adagrad.cpp'] - - def include_paths(self): - args = super().include_paths() - args += ['csrc/includes'] - return args diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_adam.py deleted file mode 100644 index a4e9569c0f336122cd003a2df5e196527d84666c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_adam.py +++ /dev/null @@ -1,25 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import NPUOpBuilder - - -class CPUAdamBuilder(NPUOpBuilder): - BUILD_VAR = "DS_BUILD_CPU_ADAM" - NAME = "cpu_adam" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adam.{self.NAME}_op' - - def sources(self): - return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] - - def include_paths(self): - args = super().include_paths() - args += ['csrc/includes'] - return args diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_lion.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_lion.py deleted file mode 100644 index 6917e0fd03d08dec42e71479110224d577b55b5b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/cpu_lion.py +++ /dev/null @@ -1,25 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import NPUOpBuilder - - -class CPULionBuilder(NPUOpBuilder): - BUILD_VAR = "DS_BUILD_CPU_LION" - NAME = "cpu_lion" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.lion.{self.NAME}_op' - - def sources(self): - return ['csrc/lion/cpu_lion.cpp', 'csrc/lion/cpu_lion_impl.cpp'] - - def include_paths(self): - args = super().include_paths() - args += ['csrc/includes'] - return args diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/fused_adam.py deleted file mode 100644 index d32103db7055669f8ecce578bb4ef7703e3c07ef..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/fused_adam.py +++ /dev/null @@ -1,74 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import NPUOpBuilder - -try: - import torch_npu -except ImportError as e: - pass - - -class NPUFusedAdam: - - @staticmethod - def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, - bias_correction, weight_decay, *args): - bias_correction1 = beta1**(step - 1) - bias_correction2 = beta2**(step - 1) - - # iteration group['params'] - for i in range(len(tensor_lists[0])): - grad_flat = tensor_lists[0][i] - param_flat = tensor_lists[1][i] - m_flat = tensor_lists[2][i] - v_flat = tensor_lists[3][i] - - if adam_w_mode: - param_flat.data, m_flat, v_flat = torch_npu.npu_apply_adam_w( - bias_correction1, - bias_correction2, - lr, - weight_decay, - beta1, - beta2, - epsilon, - grad_flat, - None, # max_grad_norm - False, # amsgrad - False, # maximize - out=(param_flat.data, m_flat, v_flat)) - else: - param_flat.data, m_flat, v_flat = torch_npu.npu_apply_adam( - bias_correction1, - bias_correction2, - lr, - beta1, - beta2, - epsilon, - grad_flat, - False, # use_locking - False, # use_nesterov - out=(param_flat.data, m_flat, v_flat)) - - -class FusedAdamBuilder(NPUOpBuilder): - BUILD_VAR = "DS_BUILD_FUSED_ADAM" - NAME = "fused_adam" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adam.{self.NAME}_op' - - def sources(self): - return [] - - def include_paths(self): - return [] - - def load(self, verbose=True): - return NPUFusedAdam diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/inference.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/inference.py deleted file mode 100644 index 46f28c0d401161f70431776a5a53387235ebb5ce..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/inference.py +++ /dev/null @@ -1,307 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from enum import IntEnum -from .builder import NPUOpBuilder - -try: - import torch - import torch_npu -except ImportError as e: - pass - - -class ActivationFuncType(IntEnum): - UNKNOWN = 0 - GELU = 1 - ReLU = 2 - GATED_GELU = 3 - GATED_SILU = 4 - - -class InferenceContext: - _workspace = None - - _seed = 42 - _curr_offset = 0 - _stream = 0 - _free_memory_size = 0 - _num_tokens = 1 - _attention_unfused_workspace_offset = 0 - _workSpaceSize = 0 - - workSpaceSize = 0 - kv_caches = None - - @staticmethod - def reset_tokens(initial_tokens=1): - InferenceContext._num_tokens = initial_tokens - - @staticmethod - def current_tokens(): - return InferenceContext._num_tokens - - @staticmethod - def GetWorkSpace(): - return InferenceContext._workspace - - -class NPUInference: - - @staticmethod - def layer_norm(inputs, gamma, beta, epsilon): - return torch.nn.functional.layer_norm(inputs, [inputs.shape[-1]], gamma, beta, eps=epsilon) - - @staticmethod - def _qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose): - inp_norm = torch.nn.functional.layer_norm(inputs, (inputs.shape[2], ), gamma, beta, eps) - weight = weight.t() if transpose else weight - tmp = torch.matmul(inp_norm, weight) - if add_bias: - tmp += bias - output = [tmp, inp_norm] - return output - - @staticmethod - def qkv_gemm_fp16(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose): - return NPUInference._qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose) - - @staticmethod - def qkv_gemm_bf16(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose): - return NPUInference._qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose) - - @staticmethod - def qkv_gemm_fp32(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose): - return NPUInference._qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose) - - @staticmethod - def _bias_add_transform_0213(vals, bias, hidden_dim, seq_length, seq_offset, heads, num_kv, rotary_dim, - rotate_half, rotate_every_two, rope_theta): - bsz, _, _ = vals.shape - q = vals[..., :hidden_dim].reshape(bsz, seq_length, heads, -1) - k = vals[..., hidden_dim:hidden_dim + num_kv * (hidden_dim // heads)].reshape(bsz, seq_length, num_kv, -1) - v = vals[..., hidden_dim + num_kv * (hidden_dim // heads):] - - if rotary_dim > 0 and rotate_every_two: - # sin, cos may use cache - seq_id = torch.arange(0, seq_length).to("npu") - inv_freq = torch.arange(0, rotary_dim, 2) / rotary_dim - inv_freq = inv_freq.to("npu") - inv_freq = 1.0 / torch.pow(rope_theta, inv_freq) - inv_freq = torch.outer(seq_id, inv_freq) - sin = inv_freq.sin() - cos = inv_freq.cos() - # shape: [bsz=1, seq_len, heads=1, rotary_dim] - sin = sin.view(-1, seq_length, 1, rotary_dim // 2).repeat_interleave(2, dim=-1) - cos = cos.view(-1, seq_length, 1, rotary_dim // 2).repeat_interleave(2, dim=-1) - - q_pos, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] - k_pos, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] - - q_pos = torch_npu.npu_rotary_mul(q_pos, cos, sin) - q = torch.cat([q_pos, q_pass], dim=-1) - k_pos = torch_npu.npu_rotary_mul(k_pos, cos, sin) - k = torch.cat([k_pos, k_pass], dim=-1) - - output = q.reshape(bsz, seq_length, -1).contiguous() # [b, s, H] - k_cache = k.reshape(bsz, seq_length, heads, -1).transpose(1, 2).contiguous() # [b, n, s, d] - v_cache = v.reshape(bsz, seq_length, heads, -1).transpose(1, 2).contiguous() # [b, n, s, d] - return output, k_cache, v_cache - - @staticmethod - def _softmax_context(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads, num_kv, - norm_factor, triangular_masking, local_attention, window_size, no_masking, layer_id, - num_layers, alibi, rope_theta): - bsz, seq_len, k = query_key_value.size() - k = k // (heads + 2 * (num_kv if num_kv > 0 else heads)) - hidden_dim = heads * k - - is_promt = seq_len > 1 - if not InferenceContext.kv_caches: - InferenceContext.kv_caches = [[None, None] for _ in range(num_layers)] - if is_promt: - InferenceContext.reset_tokens(seq_len) - InferenceContext.kv_caches[layer_id] = [None, None] - - soft_len = InferenceContext.current_tokens() - workspace = InferenceContext.GetWorkSpace() - seq_offset = 0 if is_promt else soft_len - 1 - - q, k, v = NPUInference._bias_add_transform_0213(vals=query_key_value, - bias=None, - hidden_dim=hidden_dim, - seq_length=seq_len, - seq_offset=seq_offset, - heads=heads, - num_kv=num_kv if num_kv > 0 else heads, - rotary_dim=rotary_dim, - rotate_half=rotate_half, - rotate_every_two=rotate_every_two, - rope_theta=rope_theta) - - if not is_promt: - k_cache, v_cache = InferenceContext.kv_caches[layer_id] - if k_cache is not None: - k = torch.cat([k_cache, k], dim=2) - v = torch.cat([v_cache, v], dim=2) - InferenceContext.kv_caches[layer_id] = [k, v] - seq_len = k.shape[2] - - layer_scale = max(1, layer_id) if len(alibi.size()) > 1 else 1.0 - alpha = norm_factor * norm_factor / layer_scale - - output = torch_npu.npu_fusion_attention(q, - k.transpose(1, 2).reshape(bsz, seq_len, -1).contiguous(), - v.transpose(1, 2).reshape(bsz, seq_len, -1).contiguous(), - heads, - "BSH", - pse=None, - padding_mask=None, - atten_mask=attn_mask.bool(), - scale=alpha, - pre_tockens=65536, - next_tockens=65536, - keep_prob=1, - inner_precise=0)[0] - - return output, k, v - - @staticmethod - def softmax_context_fp16(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads, num_kv, - norm_factor, triangular_masking, local_attention, window_size, no_masking, layer_id, - num_layers, alibi, rope_theta): - return NPUInference._softmax_context(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, - heads, num_kv, norm_factor, triangular_masking, local_attention, - window_size, no_masking, layer_id, num_layers, alibi, rope_theta) - - @staticmethod - def softmax_context_bf16(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads, num_kv, - norm_factor, triangular_masking, local_attention, window_size, no_masking, layer_id, - num_layers, alibi, rope_theta): - return NPUInference._softmax_context(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, - heads, num_kv, norm_factor, triangular_masking, local_attention, - window_size, no_masking, layer_id, num_layers, alibi, rope_theta) - - @staticmethod - def softmax_context_fp32(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads, num_kv, - norm_factor, triangular_masking, local_attention, window_size, no_masking, layer_id, - num_layers, alibi, rope_theta): - return NPUInference._softmax_context(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, - heads, num_kv, norm_factor, triangular_masking, local_attention, - window_size, no_masking, layer_id, num_layers, alibi, rope_theta) - - @staticmethod - def _vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode): - if transposed_mode: - return torch.matmul(input, weight.t()) - return torch.matmul(input, weight) - - @staticmethod - def vector_matmul_fp16(input, weight, async_op, q_scale, q_int8, transposed_mode): - return NPUInference._vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode) - - @staticmethod - def vector_matmul_bf16(input, weight, async_op, q_scale, q_int8, transposed_mode): - return NPUInference._vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode) - - @staticmethod - def vector_matmul_fp32(input, weight, async_op, q_scale, q_int8, transposed_mode): - return NPUInference._vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode) - - @staticmethod - def _mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm, - mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose): - if mlp_after_attn: - residual_add = torch.nn.functional.layer_norm(input + residual + input_bias, (input.shape[-1], ), gamma, - beta, eps) - else: - residual_add = torch.nn.functional.layer_norm(input, (input.shape[-1], ), gamma, beta, eps) - - weight_interm = weight_interm.t() if transpose else weight_interm - tmp = torch.matmul(residual_add, weight_interm) - if mlp_act_func_type == ActivationFuncType.GELU: - tmp = torch.nn.functional.gelu(tmp + bias) - elif mlp_act_func_type == ActivationFuncType.ReLU: - tmp = torch.nn.functional.relu(tmp + bias) - else: - raise Exception('Unsupported ActivationFuncType {}'.format(mlp_act_func_type)) - output = torch.matmul(tmp, weight_out.t()) - return output, residual_add - - @staticmethod - def mlp_gemm_fp16(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm, - mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose): - return NPUInference._mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, - pre_layer_norm, mlp_after_attn, interm_scale, out_scale, dtype, - mlp_act_func_type, transpose) - - @staticmethod - def mlp_gemm_bf16(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm, - mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose): - return NPUInference._mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, - pre_layer_norm, mlp_after_attn, interm_scale, out_scale, dtype, - mlp_act_func_type, transpose) - - @staticmethod - def mlp_gemm_fp32(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm, - mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose): - return NPUInference._mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, - pre_layer_norm, mlp_after_attn, interm_scale, out_scale, dtype, - mlp_act_func_type, transpose) - - @staticmethod - def _residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size, - mlp_after_attn, add_bias, pre_layer_norm): - if mlp_after_attn: - if pre_layer_norm: - tmp = (residual.float() + attention_output.float() + attention_bias.float() + - final_bias.float()) / mp_size + hidden_state.float() - else: - tmp = residual.float() + hidden_state.float() + final_bias.float() - else: - if add_bias: - residual += attention_bias.float() - tmp = hidden_state.float() + attention_output.float() + (residual.float() + final_bias.float()) / mp_size - - input_dtype = hidden_state.dtype - residual.set_(tmp.to(input_dtype)) - - @staticmethod - def residual_add_bias_fp16(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size, - mlp_after_attn, add_bias, pre_layer_norm): - return NPUInference._residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, - mp_size, mlp_after_attn, add_bias, pre_layer_norm) - - @staticmethod - def residual_add_bias_bf16(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size, - mlp_after_attn, add_bias, pre_layer_norm): - return NPUInference._residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, - mp_size, mlp_after_attn, add_bias, pre_layer_norm) - - @staticmethod - def residual_add_bias_fp32(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size, - mlp_after_attn, add_bias, pre_layer_norm): - return NPUInference._residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, - mp_size, mlp_after_attn, add_bias, pre_layer_norm) - - -class InferenceBuilder(NPUOpBuilder): - BUILD_VAR = "DS_BUILD_TRANSFORMER_INFERENCE" - NAME = "transformer_inference" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.transformer.inference.{self.NAME}_op' - - def sources(self): - return [] - - def include_paths(self): - return [] - - def load(self, verbose=True): - return NPUInference diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/no_impl.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/no_impl.py deleted file mode 100644 index 5b1771fabc22f6ad13161231900a737bba733e68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/npu/no_impl.py +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import NPUOpBuilder - - -class NotImplementedBuilder(NPUOpBuilder): - BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" - NAME = "deepspeed_not_implemented" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.comm.{self.NAME}_op' - - def load(self, verbose=True): - raise ValueError("This op had not been implemented on NPU backend.") - - def sources(self): - return [] - - def cxx_args(self): - return [] - - def extra_ldflags(self): - return [] - - def include_paths(self): - return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/quantizer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/quantizer.py deleted file mode 100644 index 0b5348e5af96e28e256524b38efaff5d50863d17..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/quantizer.py +++ /dev/null @@ -1,38 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import CUDAOpBuilder - - -class QuantizerBuilder(CUDAOpBuilder): - BUILD_VAR = "DS_BUILD_QUANTIZER" - NAME = "quantizer" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.quantizer.{self.NAME}_op' - - def sources(self): - return [ - 'csrc/quantization/pt_binding.cpp', - 'csrc/quantization/fake_quantizer.cu', - 'csrc/quantization/quantize.cu', - 'csrc/quantization/quantize_intX.cu', - 'csrc/quantization/dequantize.cu', - 'csrc/quantization/swizzled_quantize.cu', - 'csrc/quantization/quant_reduce.cu', - ] - - def include_paths(self): - return ['csrc/includes'] - - def extra_ldflags(self): - if not self.is_rocm_pytorch(): - return ['-lcurand'] - else: - return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/ragged_ops.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/ragged_ops.py deleted file mode 100644 index 0df28cc2282a0beee59e4e5ae4d84c3a94ef9aa3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/ragged_ops.py +++ /dev/null @@ -1,118 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os - -from .builder import CUDAOpBuilder, installed_cuda_version - - -class RaggedOpsBuilder(CUDAOpBuilder): - BUILD_VAR = "DS_BUILD_RAGGED_DEVICE_OPS" - NAME = "ragged_device_ops" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.inference.v2.kernels.ragged_ops.{self.NAME}' - - def is_compatible(self, verbose=False): - try: - import torch - except ImportError: - if verbose: - self.warning("Please install torch if trying to pre-compile inference kernels") - return False - - cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda - if cuda_capability < 6: - if verbose: - self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") - cuda_okay = False - if cuda_capability >= 8: - if torch_cuda_major < 11 or sys_cuda_major < 11: - if verbose: - self.warning("On Ampere and higher architectures please use CUDA 11+") - cuda_okay = False - return super().is_compatible(verbose) and cuda_okay - - def filter_ccs(self, ccs): - ccs_retained = [] - ccs_pruned = [] - for cc in [cc.split('.') for cc in ccs]: - if int(cc[0]) >= 8: - # Blocked flash has a dependency on Ampere + newer - ccs_retained.append(cc) - else: - ccs_pruned.append(cc) - if len(ccs_pruned) > 0: - self.warning(f"Filtered compute capabilities {ccs_pruned}") - return ccs_retained - - def get_prefix(self): - ds_path = self.deepspeed_src_path("deepspeed") - return "deepspeed" if os.path.isdir(ds_path) else ".." - - def sources(self): - sources = [ - "inference/v2/kernels/ragged_ops/ragged_ops.cpp", - "inference/v2/kernels/ragged_ops/atom_builder/atom_builder.cpp", - "inference/v2/kernels/ragged_ops/blocked_flash/blocked_flash.cpp", - "inference/v2/kernels/ragged_ops/embed/embed.cpp", - "inference/v2/kernels/ragged_ops/embed/embed_cuda.cu", - "inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.cpp", - "inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary_cuda.cu", - "inference/v2/kernels/ragged_ops/logits_gather/logits_gather.cpp", - "inference/v2/kernels/ragged_ops/logits_gather/logits_gather_cuda.cu", - "inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.cpp", - "inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter_cuda.cu", - "inference/v2/kernels/ragged_ops/moe_gather/moe_gather.cpp", - "inference/v2/kernels/ragged_ops/moe_gather/moe_gather_cuda.cu", - "inference/v2/kernels/ragged_ops/ragged_helpers/ragged_kernel_helpers.cpp", - "inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.cpp", - "inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating_cuda.cu", - ] - - prefix = self.get_prefix() - sources = [os.path.join(prefix, src) for src in sources] - return sources - - def extra_ldflags(self): - import dskernels - lib_path = dskernels.library_path() - - prefix = self.get_prefix() - lib_path = os.path.join(prefix, lib_path) - lib_path = self.deepspeed_src_path(lib_path) - - args = [f'-L{lib_path}', '-lblockedflash'] - if self.jit_load: - args.append(f'-Wl,-rpath,{lib_path}') - return args - - def include_paths(self): - sources = [ - 'inference/v2/kernels/includes', - 'inference/v2/kernels/ragged_ops', - 'inference/v2/kernels/ragged_ops/atom_builder', - 'inference/v2/kernels/ragged_ops/blocked_flash', - 'inference/v2/kernels/ragged_ops/embed', - 'inference/v2/kernels/ragged_ops/includes', - 'inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary', - 'inference/v2/kernels/ragged_ops/logits_gather', - 'inference/v2/kernels/ragged_ops/moe_gather', - 'inference/v2/kernels/ragged_ops/moe_scatter', - 'inference/v2/kernels/ragged_ops/ragged_helpers', - 'inference/v2/kernels/ragged_ops/top_k_gating', - ] - - prefix = self.get_prefix() - sources = [os.path.join(prefix, src) for src in sources] - return sources diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/ragged_utils.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/ragged_utils.py deleted file mode 100644 index 208c9f833ebe97c39593e66c46a53b71190fedfd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/ragged_utils.py +++ /dev/null @@ -1,80 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os - -from .builder import CUDAOpBuilder, installed_cuda_version - - -class RaggedUtilsBuilder(CUDAOpBuilder): - BUILD_VAR = "DS_BUILD_RAGGED_OPS" - NAME = "ragged_ops" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.inference.v2.{self.NAME}' - - def is_compatible(self, verbose=False): - try: - import torch - except ImportError: - if verbose: - self.warning("Please install torch if trying to pre-compile inference kernels") - return False - - cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda - if cuda_capability < 6: - if verbose: - self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") - cuda_okay = False - if cuda_capability >= 8: - if torch_cuda_major < 11 or sys_cuda_major < 11: - if verbose: - self.warning("On Ampere and higher architectures please use CUDA 11+") - cuda_okay = False - return super().is_compatible(verbose) and cuda_okay - - def filter_ccs(self, ccs): - ccs_retained = [] - ccs_pruned = [] - for cc in [cc.split('.') for cc in ccs]: - if int(cc[0]) >= 6: - ccs_retained.append(cc) - else: - ccs_pruned.append(cc) - if len(ccs_pruned) > 0: - self.warning(f"Filtered compute capabilities {ccs_pruned}") - return ccs_retained - - def get_prefix(self): - ds_path = self.deepspeed_src_path("deepspeed") - return "deepspeed" if os.path.isdir(ds_path) else ".." - - def sources(self): - sources = [ - "inference/v2/ragged/csrc/fast_host_buffer.cu", - "inference/v2/ragged/csrc/ragged_ops.cpp", - ] - - prefix = self.get_prefix() - sources = [os.path.join(prefix, src) for src in sources] - return sources - - def extra_ldflags(self): - return [] - - def include_paths(self): - include_dirs = ['inference/v2/ragged/includes', 'inference/v2/kernels/includes'] - prefix = self.get_prefix() - includes = [os.path.join(prefix, include_dir) for include_dir in include_dirs] - - return includes diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/random_ltd.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/random_ltd.py deleted file mode 100644 index 54af7150fb36f9eb8bd6a295648d8a11da8d8373..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/random_ltd.py +++ /dev/null @@ -1,34 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import CUDAOpBuilder - - -class RandomLTDBuilder(CUDAOpBuilder): - BUILD_VAR = "DS_BUILD_RANDOM_LTD" - NAME = "random_ltd" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.{self.NAME}_op' - - def extra_ldflags(self): - if not self.is_rocm_pytorch(): - return ['-lcurand'] - else: - return [] - - def sources(self): - return [ - 'csrc/random_ltd/pt_binding.cpp', 'csrc/random_ltd/gather_scatter.cu', - 'csrc/random_ltd/slice_attn_masks.cu', 'csrc/random_ltd/token_sort.cu' - ] - - def include_paths(self): - includes = ['csrc/includes'] - return includes diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__init__.py deleted file mode 100644 index 2a6eb8bbfa2c52c3e61ddbca734011c2ec5d630b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__init__.py +++ /dev/null @@ -1,36 +0,0 @@ -# Copyright (c) Microsoft Corporation. - -# SPDX-License-Identifier: Apache-2.0 - -# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights -# reserved. -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# Redistributions of source code must retain the above copyright notice, -# this list of conditions and the following disclaimer. -# Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# Neither the name of the copyright holder nor the names of its contributors -# may be used to endorse or promote products derived from this software -# without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE -# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -# INTERRUPTION) -# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, -# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY -# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY -# OF SUCH DAMAGE. - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' - -from .no_impl import NotImplementedBuilder -from .cpu_adam import CPUAdamBuilder -from .fused_adam import FusedAdamBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 0a72c6e77032dc7c1510f512000ece69b9f15304..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/builder.cpython-310.pyc deleted file mode 100644 index ad16d07da9757ba5c138fd4922d60b97fd7f70d2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/builder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/cpu_adam.cpython-310.pyc deleted file mode 100644 index 3cf4c2a96759c73f826fabf776775c391c185679..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/cpu_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/fused_adam.cpython-310.pyc deleted file mode 100644 index ea9ef88279f91862f7b4ae0d5808f533656e1fc5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/fused_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/no_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/no_impl.cpython-310.pyc deleted file mode 100644 index 62ec296ca461c110a9ed79b7ca22ab0138d59783..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/no_impl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/builder.py deleted file mode 100644 index 81f0e98c076837650bd5a3cd9a23493361e6b02e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/builder.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights -# reserved. -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# Redistributions of source code must retain the above copyright notice, -# this list of conditions and the following disclaimer. -# Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# Neither the name of the copyright holder nor the names of its contributors -# may be used to endorse or promote products derived from this software -# without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE -# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -# INTERRUPTION) -# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, -# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY -# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY -# OF SUCH DAMAGE. - -# DeepSpeed Team - -try: - # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed - # if successful this also means we're doing a local install and not JIT compile path - from op_builder import __deepspeed__ # noqa: F401 # type: ignore - from op_builder.builder import OpBuilder -except ImportError: - from deepspeed.ops.op_builder.builder import OpBuilder - - -class SDAAOpBuilder(OpBuilder): - - def builder(self): - from torch.utils.cpp_extension import CppExtension as ExtensionBuilder - - compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} - - cpp_ext = ExtensionBuilder(name=self.absolute_name(), - sources=self.strip_empty_entries(self.sources()), - include_dirs=self.strip_empty_entries(self.include_paths()), - libraries=self.strip_empty_entries(self.libraries_args()), - extra_compile_args=compile_args) - - return cpp_ext - - def cxx_args(self): - return ['-O3', '-g', '-Wno-reorder'] - - def libraries_args(self): - return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/cpu_adam.py deleted file mode 100644 index b38a71a6275d3b962ed2bd7d8f625ba3298418dd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/cpu_adam.py +++ /dev/null @@ -1,53 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights -# reserved. -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# Redistributions of source code must retain the above copyright notice, -# this list of conditions and the following disclaimer. -# Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# Neither the name of the copyright holder nor the names of its contributors -# may be used to endorse or promote products derived from this software -# without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE -# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -# INTERRUPTION) -# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, -# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY -# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY -# OF SUCH DAMAGE. - -# DeepSpeed Team - -from .builder import SDAAOpBuilder - - -class CPUAdamBuilder(SDAAOpBuilder): - BUILD_VAR = "DS_BUILD_CPU_ADAM" - NAME = "cpu_adam" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adam.{self.NAME}_op' - - def sources(self): - return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] - - def libraries_args(self): - args = super().libraries_args() - return args - - def include_paths(self): - return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/fused_adam.py deleted file mode 100644 index 73a2dff4145961a909258278e64a6fd094ca0206..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/fused_adam.py +++ /dev/null @@ -1,67 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights -# reserved. -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# Redistributions of source code must retain the above copyright notice, -# this list of conditions and the following disclaimer. -# Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# Neither the name of the copyright holder nor the names of its contributors -# may be used to endorse or promote products derived from this software -# without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE -# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -# INTERRUPTION) -# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, -# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY -# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY -# OF SUCH DAMAGE. - -# DeepSpeed Team - -from .builder import SDAAOpBuilder - -try: - import torch -except ImportError as e: - pass - - -class SDAAFusedAdam: - - @staticmethod - def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, - bias_correction, weight_decay, *args): - g_tensor_lis, p_tensor_lis, m_tensor_lis, v_tensor_lis = tensor_lists - torch.ops.sdaa.fused_adam(g_tensor_lis, p_tensor_lis, m_tensor_lis, v_tensor_lis, [], beta1, beta2, epsilon, - lr, weight_decay, adam_w_mode, step, bias_correction) - - -class FusedAdamBuilder(SDAAOpBuilder): - BUILD_VAR = "DS_BUILD_FUSED_ADAM" - NAME = "fused_adam" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adam.{self.NAME}_op' - - def sources(self): - return [] - - def include_paths(self): - return [] - - def load(self, verbose=True): - return SDAAFusedAdam diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/no_impl.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/no_impl.py deleted file mode 100644 index 10a8b8f48652a8abbd9978fb444332036571acb3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sdaa/no_impl.py +++ /dev/null @@ -1,59 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights -# reserved. -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# Redistributions of source code must retain the above copyright notice, -# this list of conditions and the following disclaimer. -# Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# Neither the name of the copyright holder nor the names of its contributors -# may be used to endorse or promote products derived from this software -# without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE -# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -# INTERRUPTION) -# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, -# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY -# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY -# OF SUCH DAMAGE. - -# DeepSpeed Team - -from .builder import SDAAOpBuilder - - -class NotImplementedBuilder(SDAAOpBuilder): - BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" - NAME = "deepspeed_not_implemented" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.comm.{self.NAME}_op' - - def load(self, verbose=True): - raise ValueError("This op had not been implemented on SDAA backend.") - - def sources(self): - return [] - - def cxx_args(self): - return [] - - def extra_ldflags(self): - return [] - - def include_paths(self): - return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sparse_attn.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sparse_attn.py deleted file mode 100644 index b8e5e8a6fd9aaa62372f72c0b78d250994b3e98e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/sparse_attn.py +++ /dev/null @@ -1,89 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import OpBuilder - -try: - from packaging import version as pkg_version -except ImportError: - pkg_version = None - - -class SparseAttnBuilder(OpBuilder): - BUILD_VAR = "DS_BUILD_SPARSE_ATTN" - NAME = "sparse_attn" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.sparse_attention.{self.NAME}_op' - - def sources(self): - return ['csrc/sparse_attention/utils.cpp'] - - def cxx_args(self): - return ['-O2', '-fopenmp'] - - def is_compatible(self, verbose=False): - # Check to see if llvm and cmake are installed since they are dependencies - #required_commands = ['llvm-config|llvm-config-9', 'cmake'] - #command_status = list(map(self.command_exists, required_commands)) - #deps_compatible = all(command_status) - - if self.is_rocm_pytorch(): - if verbose: - self.warning(f'{self.NAME} is not compatible with ROCM') - return False - - try: - import torch - except ImportError: - if verbose: - self.warning(f"unable to import torch, please install it first") - return False - - # torch-cpu will not have a cuda version - if torch.version.cuda is None: - cuda_compatible = False - if verbose: - self.warning(f"{self.NAME} cuda is not available from torch") - else: - major, minor = torch.version.cuda.split('.')[:2] - cuda_compatible = (int(major) == 10 and int(minor) >= 1) or (int(major) >= 11) - if not cuda_compatible: - if verbose: - self.warning(f"{self.NAME} requires CUDA version 10.1+") - - TORCH_MAJOR = int(torch.__version__.split('.')[0]) - TORCH_MINOR = int(torch.__version__.split('.')[1]) - torch_compatible = (TORCH_MAJOR == 1 and TORCH_MINOR >= 5) - if not torch_compatible: - if verbose: - self.warning( - f'{self.NAME} requires a torch version >= 1.5 and < 2.0 but detected {TORCH_MAJOR}.{TORCH_MINOR}') - try: - import triton - except ImportError: - # auto-install of triton is broken on some systems, reverting to manual install for now - # see this issue: https://github.com/deepspeedai/DeepSpeed/issues/1710 - if verbose: - self.warning(f"please install triton==1.0.0 if you want to use sparse attention") - return False - - if pkg_version: - installed_triton = pkg_version.parse(triton.__version__) - triton_mismatch = installed_triton != pkg_version.parse("1.0.0") - else: - installed_triton = triton.__version__ - triton_mismatch = installed_triton != "1.0.0" - - if triton_mismatch: - if verbose: - self.warning( - f"using untested triton version ({installed_triton}), only 1.0.0 is known to be compatible") - return False - - return super().is_compatible(verbose) and torch_compatible and cuda_compatible diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/spatial_inference.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/spatial_inference.py deleted file mode 100644 index d6c5fa6611563bc0a45fe712e27f0f06f56bf2d0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/spatial_inference.py +++ /dev/null @@ -1,47 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import CUDAOpBuilder, installed_cuda_version - - -class SpatialInferenceBuilder(CUDAOpBuilder): - BUILD_VAR = "DS_BUILD_SPATIAL_INFERENCE" - NAME = "spatial_inference" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.spatial.{self.NAME}_op' - - def is_compatible(self, verbose=False): - try: - import torch - except ImportError: - if verbose: - self.warning("Please install torch if trying to pre-compile inference kernels") - return False - - cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major - if cuda_capability >= 8: - if torch_cuda_major < 11 or sys_cuda_major < 11: - if verbose: - self.warning("On Ampere and higher architectures please use CUDA 11+") - cuda_okay = False - return super().is_compatible(verbose) and cuda_okay - - def sources(self): - return [ - 'csrc/spatial/csrc/opt_bias_add.cu', - 'csrc/spatial/csrc/pt_binding.cpp', - ] - - def include_paths(self): - return ['csrc/spatial/includes', 'csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/stochastic_transformer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/stochastic_transformer.py deleted file mode 100644 index 52b02a3c629e9d537b8e22139f48f5335396e4ff..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/stochastic_transformer.py +++ /dev/null @@ -1,22 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .transformer import TransformerBuilder - - -class StochasticTransformerBuilder(TransformerBuilder): - BUILD_VAR = "DS_BUILD_STOCHASTIC_TRANSFORMER" - NAME = "stochastic_transformer" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.transformer.{self.NAME}_op' - - def nvcc_args(self): - args = super().nvcc_args() - args.append('-D__STOCHASTIC_MODE__') - return args diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/transformer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/transformer.py deleted file mode 100644 index 8db30fdc67919474679715014cbf8d3f470bedf4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/transformer.py +++ /dev/null @@ -1,36 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import CUDAOpBuilder - - -class TransformerBuilder(CUDAOpBuilder): - BUILD_VAR = "DS_BUILD_TRANSFORMER" - NAME = "transformer" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.transformer.{self.NAME}_op' - - def extra_ldflags(self): - if not self.is_rocm_pytorch(): - return ['-lcurand'] - else: - return [] - - def sources(self): - return [ - 'csrc/transformer/ds_transformer_cuda.cpp', 'csrc/transformer/cublas_wrappers.cu', - 'csrc/transformer/transform_kernels.cu', 'csrc/transformer/gelu_kernels.cu', - 'csrc/transformer/dropout_kernels.cu', 'csrc/transformer/normalize_kernels.cu', - 'csrc/transformer/softmax_kernels.cu', 'csrc/transformer/general_kernels.cu' - ] - - def include_paths(self): - includes = ['csrc/includes'] - return includes diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/transformer_inference.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/transformer_inference.py deleted file mode 100644 index 642aed56a192e4b5170a62d8cc457bb451ae8646..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/transformer_inference.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import CUDAOpBuilder, installed_cuda_version - - -class InferenceBuilder(CUDAOpBuilder): - BUILD_VAR = "DS_BUILD_TRANSFORMER_INFERENCE" - NAME = "transformer_inference" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.transformer.inference.{self.NAME}_op' - - def is_compatible(self, verbose=False): - try: - import torch - except ImportError: - if verbose: - self.warning("Please install torch if trying to pre-compile inference kernels") - return False - - cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major - if cuda_capability < 6: - if verbose: - self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") - cuda_okay = False - if cuda_capability >= 8: - if torch_cuda_major < 11 or sys_cuda_major < 11: - if verbose: - self.warning("On Ampere and higher architectures please use CUDA 11+") - cuda_okay = False - return super().is_compatible(verbose) and cuda_okay - - def filter_ccs(self, ccs): - ccs_retained = [] - ccs_pruned = [] - for cc in [cc.split('.') for cc in ccs]: - if int(cc[0]) >= 6: - ccs_retained.append(cc) - else: - ccs_pruned.append(cc) - if len(ccs_pruned) > 0: - self.warning(f"Filtered compute capabilities {ccs_pruned}") - return ccs_retained - - def sources(self): - return [ - 'csrc/transformer/inference/csrc/pt_binding.cpp', - 'csrc/transformer/inference/csrc/gelu.cu', - 'csrc/transformer/inference/csrc/relu.cu', - 'csrc/transformer/inference/csrc/layer_norm.cu', - 'csrc/transformer/inference/csrc/rms_norm.cu', - 'csrc/transformer/inference/csrc/softmax.cu', - 'csrc/transformer/inference/csrc/dequantize.cu', - 'csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu', - 'csrc/transformer/inference/csrc/transform.cu', - 'csrc/transformer/inference/csrc/pointwise_ops.cu', - ] - - def extra_ldflags(self): - if not self.is_rocm_pytorch(): - return ['-lcurand'] - else: - return [] - - def include_paths(self): - return ['csrc/transformer/inference/includes', 'csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/utils.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/utils.py deleted file mode 100644 index 927ff1b361bd0876c065b6c5f8904c6668fd2ace..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/utils.py +++ /dev/null @@ -1,20 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import OpBuilder - - -class UtilsBuilder(OpBuilder): - BUILD_VAR = "DS_BUILD_UTILS" - NAME = "utils" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.{self.NAME}_op' - - def sources(self): - return ['csrc/utils/tensor_cast.cpp', 'csrc/utils/py_ds_utils.cpp'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__init__.py deleted file mode 100644 index 093f90b30234359d37913bb36919f3259d4cab1c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .cpu_adam import CPUAdamBuilder -from .cpu_adagrad import CPUAdagradBuilder -from .fused_adam import FusedAdamBuilder -from .async_io import AsyncIOBuilder -from .inference import InferenceBuilder -from .flash_attn import FlashAttentionBuilder -from .no_impl import NotImplementedBuilder -from .packbits import PackbitsBuilder diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 48c6a39f32acaeaadaa8b7b18e00eb17bbe4ec20..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/async_io.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/async_io.cpython-310.pyc deleted file mode 100644 index 956fa57de5d741b9fb34112f83260e3b8100abc5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/async_io.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/builder.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/builder.cpython-310.pyc deleted file mode 100644 index 794726222ea34077fc50f11f2df36f2dfbcca70d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/builder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adagrad.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adagrad.cpython-310.pyc deleted file mode 100644 index 09d5c83cc0ddaf3021da2e2962105a9489e241dc..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adagrad.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adam.cpython-310.pyc deleted file mode 100644 index 581a5afa69b7593d420fe43d33f0e81cd606b161..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/flash_attn.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/flash_attn.cpython-310.pyc deleted file mode 100644 index cd7a8adaba9df0125c674adca0a08a76d0b1fdf0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/flash_attn.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/fused_adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/fused_adam.cpython-310.pyc deleted file mode 100644 index eade3a536ecb05b0ad584abd0e7f8489bf1bb573..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/fused_adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/inference.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/inference.cpython-310.pyc deleted file mode 100644 index 30840b286e6ac27c0de9feb2fa9e7d449f0a0d73..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/inference.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/no_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/no_impl.cpython-310.pyc deleted file mode 100644 index 05726b588cc3494a3f68654eb6f108d629f87d51..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/no_impl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/packbits.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/packbits.cpython-310.pyc deleted file mode 100644 index 6cb3d61bbbf1711c1c848c4484bdbb40a935e0a7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/packbits.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/async_io.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/async_io.py deleted file mode 100644 index 8ec030880368dff58337b9e78bfd80bad9f31f97..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/async_io.py +++ /dev/null @@ -1,106 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import shutil -import subprocess - -from .builder import OpBuilder - - -class AsyncIOBuilder(OpBuilder): - BUILD_VAR = "DS_BUILD_AIO" - NAME = "async_io" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.aio.{self.NAME}_op' - - def sources(self): - return [ - 'csrc/aio/py_lib/deepspeed_py_copy.cpp', - 'csrc/aio/py_lib/py_ds_aio.cpp', - 'csrc/aio/py_lib/deepspeed_py_aio.cpp', - 'csrc/aio/py_lib/deepspeed_py_aio_handle.cpp', - 'csrc/aio/py_lib/deepspeed_aio_thread.cpp', - 'csrc/aio/common/deepspeed_aio_utils.cpp', - 'csrc/aio/common/deepspeed_aio_common.cpp', - 'csrc/aio/common/deepspeed_aio_types.cpp', - 'csrc/aio/py_lib/deepspeed_pin_tensor.cpp', - 'csrc/aio/py_lib/deepspeed_py_io_handle.cpp', - 'csrc/aio/py_lib/deepspeed_cpu_op.cpp', - 'csrc/aio/py_lib/deepspeed_aio_op_desc.cpp', - ] - - def include_paths(self): - return ['csrc/aio/py_lib', 'csrc/aio/common'] - - def cxx_args(self): - import torch - # -O0 for improved debugging, since performance is bound by I/O - CPU_ARCH = self.cpu_arch() - SIMD_WIDTH = self.simd_width() - TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2]) - if TORCH_MAJOR >= 2 and TORCH_MINOR >= 1: - CPP_STD = '-std=c++17' - else: - CPP_STD = '-std=c++14' - return [ - '-g', - '-Wall', - '-O0', - CPP_STD, - '-shared', - '-fPIC', - '-Wno-reorder', - CPU_ARCH, - '-fopenmp', - SIMD_WIDTH, - '-laio', - ] - - def extra_ldflags(self): - return ['-laio'] - - def check_for_libaio_pkg(self): - libs = dict( - dpkg=["-l", "libaio-dev", "apt"], - pacman=["-Q", "libaio", "pacman"], - rpm=["-q", "libaio-devel", "yum"], - ) - - found = False - for pkgmgr, data in libs.items(): - flag, lib, tool = data - path = shutil.which(pkgmgr) - if path is not None: - cmd = [pkgmgr, flag, lib] - result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - if result.wait() == 0: - found = True - else: - self.warning(f"{self.NAME}: please install the {lib} package with {tool}") - break - return found - - def is_compatible(self, verbose=False): - # Check for the existence of libaio by using distutils - # to compile and link a test program that calls io_submit, - # which is a function provided by libaio that is used in the async_io op. - # If needed, one can define -I and -L entries in CFLAGS and LDFLAGS - # respectively to specify the directories for libaio.h and libaio.so. - aio_compatible = self.has_function('io_pgetevents', ('aio', )) - if verbose and not aio_compatible: - self.warning(f"{self.NAME} requires the dev libaio .so object and headers but these were not found.") - - # Check for the libaio package via known package managers - # to print suggestions on which package to install. - self.check_for_libaio_pkg() - - self.warning( - "If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found." - ) - return super().is_compatible(verbose) and aio_compatible diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/builder.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/builder.py deleted file mode 100644 index 81b15f197f4375b6138f50cf865b2dd64194c629..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/builder.py +++ /dev/null @@ -1,131 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import time -import importlib - -try: - # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed - # if successful this also means we're doing a local install and not JIT compile path - from op_builder import __deepspeed__ # noqa: F401 # type: ignore - from op_builder.builder import OpBuilder -except ImportError: - from deepspeed.ops.op_builder.builder import OpBuilder - - -class SYCLOpBuilder(OpBuilder): - - def builder(self): - try: - from intel_extension_for_pytorch.xpu.cpp_extension import DPCPPExtension - except ImportError: - from intel_extension_for_pytorch.xpu.utils import DPCPPExtension - include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] - print("dpcpp sources = {}".format(self.sources())) - dpcpp_ext = DPCPPExtension(name=self.absolute_name(), - sources=self.strip_empty_entries(self.sources()), - include_dirs=include_dirs, - extra_compile_args={ - 'cxx': self.strip_empty_entries(self.cxx_args()), - }, - extra_link_args=self.strip_empty_entries(self.fixed_aotflags())) - return dpcpp_ext - - def version_dependent_macros(self): - try: - from op_builder.builder import TORCH_MAJOR, TORCH_MINOR - except ImportError: - from deepspeed.ops.op_builder.builder import TORCH_MAJOR, TORCH_MINOR - # Fix from apex that might be relevant for us as well, related to https://github.com/NVIDIA/apex/issues/456 - version_ge_1_1 = [] - if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 0): - version_ge_1_1 = ['-DVERSION_GE_1_1'] - version_ge_1_3 = [] - if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2): - version_ge_1_3 = ['-DVERSION_GE_1_3'] - version_ge_1_5 = [] - if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 4): - version_ge_1_5 = ['-DVERSION_GE_1_5'] - return version_ge_1_1 + version_ge_1_3 + version_ge_1_5 - - def cxx_args(self): - cxx_flags = [ - '-fsycl', '-fsycl-targets=spir64_gen', '-g', '-gdwarf-4', '-O3', '-std=c++17', '-fPIC', '-DMKL_ILP64', - '-fno-strict-aliasing' - ] - if os.environ.get('USE_MKL_GEMM'): - cxx_flags.append('-DUSE_MKL_GEMM') - return cxx_flags - - def extra_ldflags(self): - return [ - '-fPIC', '-fsycl', '-fsycl-targets=spir64_gen', '-fsycl-max-parallel-link-jobs=8', - '-Xs "-options -cl-poison-unsupported-fp64-kernels,cl-intel-enable-auto-large-GRF-mode"', - '-Xs "-device pvc"', '-Wl,-export-dynamic' - ] - - def fixed_aotflags(self): - return [ - '-fsycl', '-fsycl-targets=spir64_gen', '-fsycl-max-parallel-link-jobs=8', '-Xs', - "-options -cl-poison-unsupported-fp64-kernels,cl-intel-enable-auto-large-GRF-mode", '-Xs', "-device pvc" - ] - - def load(self, verbose=True): - from deepspeed.git_version_info import installed_ops, torch_info, accelerator_name # noqa: F401 - from deepspeed.accelerator import get_accelerator - if installed_ops.get(self.name, False) and accelerator_name == get_accelerator()._name: - return importlib.import_module(self.absolute_name()) - else: - return self.jit_load(verbose) - - def jit_load(self, verbose=True): - if not self.is_compatible(verbose): - raise RuntimeError( - f"Unable to JIT load the {self.name} op due to it not being compatible due to hardware/software issue. {self.error_log}" - ) - try: - import ninja # noqa: F401 - except ImportError: - raise RuntimeError(f"Unable to JIT load the {self.name} op due to ninja not being installed.") - - self.jit_mode = True - from intel_extension_for_pytorch.xpu.cpp_extension import load - - start_build = time.time() - # Recognize relative paths as absolute paths for jit load - - sources = [self.deepspeed_src_path(path) for path in self.sources()] - extra_include_paths = [self.deepspeed_src_path(path) for path in self.include_paths()] - - # Torch will try and apply whatever CCs are in the arch list at compile time, - # we have already set the intended targets ourselves we know that will be - # needed at runtime. This prevents CC collisions such as multiple __half - # implementations. Stash arch list to reset after build. - ''' - torch_arch_list = None - if "TORCH_CUDA_ARCH_LIST" in os.environ: - torch_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST") - os.environ["TORCH_CUDA_ARCH_LIST"] = "" - ''' - - op_module = load( - name=self.name, - sources=self.strip_empty_entries(sources), - extra_include_paths=self.strip_empty_entries(extra_include_paths), - extra_cflags=self.strip_empty_entries(self.cxx_args()), - # extra_cuda_cflags=self.strip_empty_entries(self.nvcc_args()), - extra_ldflags=self.strip_empty_entries(self.extra_ldflags()), - verbose=verbose) - - build_duration = time.time() - start_build - if verbose: - print(f"Time to load {self.name} op: {build_duration} seconds") - ''' - # Reset arch list so we are not silently removing it for other possible use cases - if torch_arch_list: - os.environ["TORCH_CUDA_ARCH_LIST"] = torch_arch_list - ''' - return op_module diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/cpu_adagrad.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/cpu_adagrad.py deleted file mode 100644 index 18f80848e1b800de7f8826405bb9ea9f500db263..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/cpu_adagrad.py +++ /dev/null @@ -1,23 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import SYCLOpBuilder - - -class CPUAdagradBuilder(SYCLOpBuilder): - BUILD_VAR = "DS_BUILD_CPU_ADAGRAD" - NAME = "cpu_adagrad" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adagrad.{self.NAME}_op' - - def sources(self): - return ['csrc/xpu/adagrad/cpu_adagrad.cpp', 'csrc/xpu/common/custom_cuda_kernel.dp.cpp'] - - def include_paths(self): - return ['csrc/xpu/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/cpu_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/cpu_adam.py deleted file mode 100644 index 9d5fdcd3e8ab51b84bb2fec137b95ee96fee8570..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/cpu_adam.py +++ /dev/null @@ -1,27 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import SYCLOpBuilder - - -class CPUAdamBuilder(SYCLOpBuilder): - BUILD_VAR = "DS_BUILD_CPU_ADAM" - NAME = "cpu_adam" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adam.{self.NAME}_op' - - def sources(self): - return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] - - def libraries_args(self): - args = super().libraries_args() - return args - - def include_paths(self): - return ['csrc/includes'] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/flash_attn.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/flash_attn.py deleted file mode 100644 index c8c2674d5d27c87377d3a64bfbba401a44ae9d52..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/flash_attn.py +++ /dev/null @@ -1,53 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -from .builder import SYCLOpBuilder - - -class FlashAttentionBuilderObject(): - - def __init__(self): - pass - - # general functions - def flash_attn_func_v2(self, q, k, v, dropout_p, softmax_scale, is_causal): - try: - import torch - import intel_extension_for_pytorch # noqa - return torch.nn.functional.scaled_dot_product_attention(q, - k, - v, - dropout_p=dropout_p, - is_causal=is_causal, - scale=softmax_scale) - except ImportError: - raise ImportError( - "Please install pytorch and intel_extension_for_pytorch to include scaled dot product attention.") - - -class FlashAttentionBuilder(SYCLOpBuilder): - BUILD_VAR = "DS_BUILD_FlashAttention" - NAME = "flash_attn" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.{self.NAME}_op' - - def sources(self): - return - - def include_paths(self): - return [] - - def extra_ldflags(self): - return [] - - def cxx_args(self): - return [] - - def load(self): - return FlashAttentionBuilderObject() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/fused_adam.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/fused_adam.py deleted file mode 100644 index 0e0f1a66f8e648305d2ba6fd0ad85a384e6f8bd3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/fused_adam.py +++ /dev/null @@ -1,26 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -from .builder import SYCLOpBuilder - - -class FusedAdamBuilder(SYCLOpBuilder): - BUILD_VAR = "DS_BUILD_FUSED_ADAM" - NAME = "fused_adam" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.adam.{self.NAME}_op' - - def sources(self): - return ['csrc/xpu/adam/fused_adam_frontend.cpp', 'csrc/xpu/adam/multi_tensor_adam.dp.cpp'] - - def include_paths(self): - return ['csrc/xpu/includes', 'csrc/xpu/adam'] - - def cxx_args(self): - args = super().cxx_args() - return args + self.version_dependent_macros() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/inference.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/inference.py deleted file mode 100644 index a9ac4f84c2ca9ba885e586934efd406b6469bd9e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/inference.py +++ /dev/null @@ -1,39 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -from .builder import SYCLOpBuilder - - -class InferenceBuilder(SYCLOpBuilder): - BUILD_VAR = "DS_BUILD_TRANSFORMER_INFERENCE" - NAME = "transformer_inference" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.transformer.inference.{self.NAME}_op' - - def sources(self): - return - - def include_paths(self): - return [] - - def extra_ldflags(self): - return [] - - def cxx_args(self): - return [] - - def load(self): - try: - import intel_extension_for_pytorch - if hasattr(intel_extension_for_pytorch, "deepspeed"): - return intel_extension_for_pytorch.deepspeed.transformer_inference.transformer_inference - else: - return intel_extension_for_pytorch.xpu.deepspeed - except ImportError: - raise ImportError("Please install intel-extension-for-pytorch >= 2.1.30 to include DeepSpeed kernels.") diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/no_impl.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/no_impl.py deleted file mode 100644 index 8b294f70c2791db9686e8d463ad736ff2d7c90c5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/no_impl.py +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .builder import SYCLOpBuilder - - -class NotImplementedBuilder(SYCLOpBuilder): - BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" - NAME = "deepspeed_not_implemented" - - def __init__(self, name=None): - name = self.NAME if name is None else name - super().__init__(name=name) - - def absolute_name(self): - return f'deepspeed.ops.{self.NAME}_op' - - def load(self, verbose=True): - raise ValueError("This op had not been implemented on XPU backend.") - - def sources(self): - return [] - - def cxx_args(self): - return [] - - def extra_ldflags(self): - return [] - - def include_paths(self): - return [] diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/packbits.py b/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/packbits.py deleted file mode 100644 index cf5b5ebc59e407ae20f478151846e282680fe4ae..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/op_builder/xpu/packbits.py +++ /dev/null @@ -1,26 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -from .builder import SYCLOpBuilder - - -class PackbitsBuilder(SYCLOpBuilder): - BUILD_VAR = "DS_BUILD_PACK_BITS" - NAME = "pack_bits" - - def __init__(self): - super().__init__(name=self.NAME) - - def absolute_name(self): - return f'deepspeed.ops.{self.NAME}_op' - - def sources(self): - return ['csrc/xpu/packbits/packing.cpp'] - - def include_paths(self): - return ['csrc/xpu/includes'] - - def cxx_args(self): - args = super().cxx_args() - return args + self.version_dependent_macros() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__init__.py deleted file mode 100644 index 0a5d1da5e3ae0fa097e7313ddb1328c4f910801d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .quantizer import ds_quantizer diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 6306206049784dca7dacb649eb9d2ced11d7bb5d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__pycache__/quantizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__pycache__/quantizer.cpython-310.pyc deleted file mode 100644 index 26af8c3833a98b313b5a2610089d29d7772369f5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/__pycache__/quantizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/quantizer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/quantizer.py deleted file mode 100644 index eb4bfd35700075f3b32db329c5f7026b39bef520..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/quantizer/quantizer.py +++ /dev/null @@ -1,29 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from deepspeed.ops.op_builder import QuantizerBuilder - -# Cuda modules will be imported if needed -quantizer_cuda_module = None - - -def ds_quantizer(input, groups=1, bit_num=8, sr=False, asym=False): - # Load cuda modules if needed - global quantizer_cuda_module - if quantizer_cuda_module is None: - quantizer_cuda_module = QuantizerBuilder().load() - if sr: - if asym: - quantize_func = quantizer_cuda_module.ds_sr_quantize_asym_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_sr_quantize_asym_fp32 - else: - quantize_func = quantizer_cuda_module.ds_sr_quantize_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_sr_quantize_fp32 - else: - if asym: - quantize_func = quantizer_cuda_module.ds_quantize_asym_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_quantize_asym_fp32 - else: - quantize_func = quantizer_cuda_module.ds_quantize_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_quantize_fp32 - return quantize_func(input, groups, bit_num) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__init__.py deleted file mode 100644 index 1ade5811ce76d5b7bf51881db5ff0824c1fd379d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .dropping_utils import gpt_sample_tokens, bert_sample_tokens, GatherTokens, ScatterTokens diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index ef9ef6640ab54329bf9783fa46d9e67dc179513e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__pycache__/dropping_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__pycache__/dropping_utils.cpython-310.pyc deleted file mode 100644 index 5b78227479c82388b4267866246d669e0da5a8c7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/__pycache__/dropping_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/dropping_utils.py b/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/dropping_utils.py deleted file mode 100644 index dd36c94537f896831291458514b513ea840791e2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/random_ltd/dropping_utils.py +++ /dev/null @@ -1,132 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from deepspeed.ops.op_builder import RandomLTDBuilder -""" -Returns: - sampled_indices: [layers, batch_size, reserved_length] - new_mask: [batch_size, 1, reserved_length, reserved_length] -""" - -random_ltd_module = None - - -def gpt_sample_tokens(reserved_length: int, - seq_length: int, - batch_size: int, - layers: int = 1, - device: str = 'cpu', - attn_mask: torch.Tensor = None): - - prob_dist = torch.ones((layers * batch_size, seq_length), device=device) - sampled_indices = torch.multinomial(prob_dist, reserved_length) - - sampled_indices = sampled_indices.reshape(layers, batch_size, reserved_length).to(torch.int32) - global random_ltd_module - if random_ltd_module is None: - random_ltd_module = RandomLTDBuilder().load() - sampled_indices = random_ltd_module.token_sort_(sampled_indices, seq_length) - - # Not certain the optimized kernel is actually better here, cause it kind of screws - # with alignment right if the sequence length is not divisible by like 16 - # new_mask = random_ltd_module.mask_gather_gpt(attn_mask, reserved_length) - if attn_mask is not None: - new_mask = attn_mask[:, :, :reserved_length, :reserved_length] - else: - new_mask = None - - return sampled_indices, new_mask - - -""" -Returns: - sampled_indices: [layers, batch_size, reserved_length] - new_mask: [layers, batch_size, 1, reserved_length, reserved_length] -""" - - -def bert_sample_tokens(reserved_length: int, - seq_length: int, - batch_size: int, - layers: int = 1, - device: str = 'cpu', - attn_mask: torch.Tensor = None): - assert attn_mask is not None - prob_dist = torch.ones((layers * batch_size, seq_length), device=device) - sampled_indices = torch.multinomial(prob_dist, reserved_length) - - sampled_indices = sampled_indices.reshape(layers, batch_size, reserved_length).to(torch.int32) - global random_ltd_module - if random_ltd_module is None: - random_ltd_module = RandomLTDBuilder().load() - - sampled_indices = random_ltd_module.token_sort_(sampled_indices, seq_length) - dtype = sampled_indices.dtype - - sampled_indices = sampled_indices.to(torch.long) - new_mask = [] - for l in range(layers): - tmp_mask_list = [] - for i in range(batch_size): - mask_tmp = attn_mask[i:i + 1, :, sampled_indices[l][i], :] - tmp_mask_list.append(mask_tmp[:, :, :, sampled_indices[l][i]]) - new_mask.append(torch.cat(tmp_mask_list, dim=0)) - - return sampled_indices.to(dtype), new_mask - - -class GatherTokens(torch.autograd.Function): - - @staticmethod - def forward(ctx, activations: torch.Tensor, sorted_indices: torch.Tensor, batch_first: bool): - global random_ltd_module - if random_ltd_module is None: - random_ltd_module = RandomLTDBuilder().load() - ctx.save_for_backward(activations, sorted_indices) - ctx.batch_first = batch_first - return activations, random_ltd_module.token_gather(activations, sorted_indices, batch_first) - - @staticmethod - def backward(ctx, a_gradients: torch.Tensor, g_gradients: torch.Tensor): - - g_gradients = g_gradients.contiguous() - global random_ltd_module - if random_ltd_module is None: - random_ltd_module = RandomLTDBuilder().load() - activations, sorted_indices = ctx.saved_tensors - batch_first = ctx.batch_first - - return random_ltd_module.token_scatter_(a_gradients, g_gradients, sorted_indices, batch_first), None, None - - -class ScatterTokens(torch.autograd.Function): - - @staticmethod - def forward(ctx, all_activations: torch.Tensor, layer_activations: torch.Tensor, sorted_indices: torch.Tensor, - batch_first: bool): - global random_ltd_module - if random_ltd_module is None: - random_ltd_module = RandomLTDBuilder().load() - scatter_results = random_ltd_module.token_scatter_(all_activations.clone(), layer_activations, sorted_indices, - batch_first) - - ctx.save_for_backward(sorted_indices) - ctx.batch_first = batch_first - return scatter_results - - @staticmethod - def backward(ctx, out_gradients: torch.Tensor): - - out_gradients = out_gradients.contiguous() - global random_ltd_module - if random_ltd_module is None: - random_ltd_module = RandomLTDBuilder().load() - sorted_indices, = ctx.saved_tensors - batch_first = ctx.batch_first - - ret_val = random_ltd_module.token_gather(out_gradients, sorted_indices, batch_first) - return out_gradients, ret_val, None, None diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__init__.py deleted file mode 100644 index e357257869f794a06d575bfa378769f8e6d3d43c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .sparsity_config import SparsityConfig, DenseSparsityConfig, FixedSparsityConfig, VariableSparsityConfig, BigBirdSparsityConfig, BSLongformerSparsityConfig, LocalSlidingWindowSparsityConfig -from .sparse_self_attention import SparseSelfAttention -from .bert_sparse_self_attention import BertSparseSelfAttention -from .sparse_attention_utils import SparseAttentionUtils diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 24c275de33737fe3c0e358918721a5dfaeae2a59..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/bert_sparse_self_attention.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/bert_sparse_self_attention.cpython-310.pyc deleted file mode 100644 index 42d6388ddbc7cb88ce610a27435d1151125c6069..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/bert_sparse_self_attention.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/matmul.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/matmul.cpython-310.pyc deleted file mode 100644 index bcbd172a00bd7a657c5fbe7f047ea45af54d39b4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/matmul.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/softmax.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/softmax.cpython-310.pyc deleted file mode 100644 index 9d3815a0fc1970f72dea70e85801f8aa5bb9a67a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/softmax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_attention_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_attention_utils.cpython-310.pyc deleted file mode 100644 index 500a614b1c1f7bdf19c082044b9dcca25c279cff..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_attention_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_self_attention.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_self_attention.cpython-310.pyc deleted file mode 100644 index 39a403c907f50f3a741bc8acbefd49ecc62d5b41..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_self_attention.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparsity_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparsity_config.cpython-310.pyc deleted file mode 100644 index 9276165446a1f96afd9643de68a9abf474c7b004..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparsity_config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/bert_sparse_self_attention.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/bert_sparse_self_attention.py deleted file mode 100644 index 37f065e4863111275aae82bcd0bc1ae8513a0896..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/bert_sparse_self_attention.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from torch import nn -from deepspeed.ops.sparse_attention import SparseSelfAttention, FixedSparsityConfig - - -class BertSparseSelfAttention(nn.Module): - """Implements Sparse Self Attention layer of Bert model based on https://github.com/deepspeedai/DeepSpeedExamples/blob/master/bing_bert/nvidia/modelingpreln.py#L373 - - For more information please see, TODO DeepSpeed Sparse Transformer. - - For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial. - """ - - def __init__( - self, - config, - # SparsityConfig parameters needs to be set accordingly - sparsity_config=FixedSparsityConfig(num_heads=4)): - """Initialize the bert sparse self attention layer. - - Note) you can use any of the provided sparsity configs or simply add yours! - - Arguments: - config: required: Bert model config - sparsity_config: optional: this parameter determines sparsity pattern configuration; it is based on FixedSparsityConfig class. - """ - - super(BertSparseSelfAttention, self).__init__() - if config.hidden_size % config.num_attention_heads != 0: - raise ValueError("The hidden size (%d) is not a multiple of the number of attention " - "heads (%d)" % (config.hidden_size, config.num_attention_heads)) - self.num_attention_heads = config.num_attention_heads - self.attention_head_size = int(config.hidden_size / config.num_attention_heads) - self.all_head_size = self.num_attention_heads * self.attention_head_size - - self.query = nn.Linear(config.hidden_size, self.all_head_size) - self.key = nn.Linear(config.hidden_size, self.all_head_size) - self.value = nn.Linear(config.hidden_size, self.all_head_size) - - self.sparse_self_attention = SparseSelfAttention(sparsity_config) - - def transpose_for_scores(self, x): - new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size) - x = x.view(*new_x_shape) - return x.permute(0, 2, 1, 3) - - def forward(self, hidden_states, attention_mask): - """Applies forward phase of bert sparse self attention - - Arguments: - hidden_states: required: hidden_states tensor of the bert model - attn_mask: required: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported - - Return: - context_layer: a dense tensor containing attention context - """ - mixed_query_layer = self.query(hidden_states) - mixed_key_layer = self.key(hidden_states) - mixed_value_layer = self.value(hidden_states) - - query_layer = self.transpose_for_scores(mixed_query_layer) - key_layer = self.transpose_for_scores(mixed_key_layer) - value_layer = self.transpose_for_scores(mixed_value_layer) - - context_layer = self.sparse_self_attention(query_layer, - key_layer, - value_layer, - key_padding_mask=attention_mask) - - context_layer = context_layer.permute(0, 2, 1, 3).contiguous() - new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size, ) - context_layer = context_layer.view(*new_context_layer_shape) - return context_layer diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/matmul.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/matmul.py deleted file mode 100644 index b30028fffbaafaf6869156406e0f8d70c8e97538..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/matmul.py +++ /dev/null @@ -1,819 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a -# https://github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/matmul.py -import importlib -import torch - -import triton -import triton.language as tl -import triton._C.libtriton as libtriton -from deepspeed.accelerator import get_accelerator - - -@triton.jit -def _kernel(A, B, C, stride_za, stride_ha, stride_ma, stride_ka, stride_zb, stride_hb, stride_kb, stride_nb, stride_zc, - stride_hc, stride_mc, stride_nc, DS0, DS1, SDD_K, SDD_off_width, lut, locks, nlocks, **meta): - TM = meta['TM'] - TN = meta['TN'] - TK = meta['TK'] - TZ = meta['TZ'] - BLOCK = meta['BLOCK'] - #------------# - #- Prologue -# - #------------# - pid0 = tl.program_id(0) - pid1 = tl.program_id(1) - pidz = tl.program_id(2) - if meta['SDD']: - pid1 = pid1 + SDD_off_width - blockidm = tl.arange(0, TM) // BLOCK - blockidn = tl.arange(0, TN) // BLOCK - offlutm = blockidm * (TN // BLOCK) * 4 - offlutn = blockidn * 4 - header = lut + pid1 * (TM // BLOCK) * (TN // BLOCK) * 4 - z = tl.load(header + 0) - i = tl.load(header + 1 + offlutm) - j = tl.load(header + 2 + offlutn) - AS1 = SDD_K // TZ - lockid = tl.where(TZ > 1, 1, 0) - offka = pid0 * AS1 - offkb = pid0 * AS1 - offmc = 0 - offnc = 0 - offpa = 0 - offpb = 0 - maxid = TZ - offhc = 0 - offha = z - offhb = z - ram = i * BLOCK + (tl.arange(0, TM) % BLOCK) - rbn = j * BLOCK + (tl.arange(0, TN) % BLOCK) - else: - header = lut + pid0 * 6 - offset = tl.load(header + 0) - AS1 = tl.load(header + 1) - column = tl.load(header + 2) - depth = tl.load(header + 3) - lockid = tl.load(header + 4) - maxid = tl.load(header + 5) - pinc = lut + offset - offhc = depth - if meta['DSD']: - # output offset - offnc = pid1 * TN - offmc = column * TM - offpc = 0 - # dense input offset - offnb = pid1 * TN - offkb = tl.load(pinc) - offkb = tl.multiple_of(offkb, 8) # compiler hint - offpb = 0 - # sparse input offset - offma = 0 - offka = 0 - offpa = tl.load(pinc + 1) - offpa = tl.multiple_of(offpa, 8) # compiler hint - offpa = offpa * BLOCK * BLOCK - offha = 0 - offhb = depth - else: - # output offset - offmc = pid1 * TM - offnc = column * TN - offpc = 0 - # dense input offset - offma = pid1 * TM - offka = tl.load(pinc) - offka = tl.multiple_of(offka, 8) # compiler hint - offpa = 0 - # sparse input offset - offnb = 0 - offkb = 0 - offpb = tl.load(pinc + 1) - offpb = tl.multiple_of(offpb, 8) # compiler hint - offpb = offpb * BLOCK * BLOCK - offha = depth - offhb = 0 - ram = offma + tl.arange(0, TM) - rbn = offnb + tl.arange(0, TN) - - # initialize a, b pointers - rka = offka + tl.arange(0, TK) - rkb = offkb + tl.arange(0, TK) - pa = A + pidz * stride_za + offha * stride_ha + offpa + ram[:, None] * stride_ma + rka[None, :] * stride_ka - pb = B + pidz * stride_zb + offhb * stride_hb + offpb + rbn[None, :] * stride_nb + rkb[:, None] * stride_kb - if meta['DDS']: - checkam = ram[:, None] < DS0 - else: - checkam = AS1 > 0 - if meta['DSD']: - checkbn = rbn[None, :] < DS0 - else: - checkbn = AS1 > 0 - a = tl.load(pa, mask=checkam, other=0.) - b = tl.load(pb, mask=checkbn, other=0.) - - ## ---------------- ## - ## Inner Loop ## - ## ---------------- ## - acc = tl.zeros((TM, TN), dtype=tl.float32) - for k in range(AS1, 0, -TK): - acc += tl.dot(a, b) - if meta['SDD']: - inc_a = TK * stride_ka - inc_b = TK * stride_kb - else: - pinc += 2 - if meta['DSD']: - inc_b = tl.load(pinc) - inc_a = tl.load(pinc + 1) - inc_b = tl.multiple_of(inc_b, 8) - inc_a = tl.multiple_of(inc_a, 8) - inc_b = inc_b * stride_kb - if meta['DDS']: - inc_a = tl.load(pinc) - inc_b = tl.load(pinc + 1) - inc_a = tl.multiple_of(inc_a, 8) - inc_b = tl.multiple_of(inc_b, 8) - inc_a = inc_a * stride_ka - pa += inc_a - pb += inc_b - # pre-fetch - checkak = k > TK - checkbk = k > TK - checka = checkam & checkak - checkb = checkbn & checkbk - a = tl.load(pa, mask=checka) - b = tl.load(pb, mask=checkb) - c = acc.to(C.dtype.element_ty) - - if meta['SDD']: - checkc = True - rr_blockidm = tl.arange(0, TM) // BLOCK - rr_blockidn = tl.arange(0, TN) // BLOCK - rr_offlutm = rr_blockidm * (TN // BLOCK) * 4 - rr_offlutn = rr_blockidn * 4 - off_bkid = 3 + rr_offlutm[:, None] + rr_offlutn[None, :] - bkid = tl.load(header + off_bkid) - offpc = bkid * BLOCK * BLOCK - rcm = tl.arange(0, TM) % BLOCK - rcn = tl.arange(0, TN) % BLOCK - else: - rcm = offmc + tl.arange(0, TM) - rcn = offnc + tl.arange(0, TN) - if meta['DSD']: - checkc = rcn[None, :] < DS0 - if meta['DDS']: - checkc = rcm[:, None] < DS0 - - pc = C + offpc + offhc * stride_hc + pidz * stride_zc + rcm[:, None] * stride_mc + rcn[None, :] * stride_nc - # write-back directly - if lockid == 0: - tl.store(pc, c, mask=checkc) - # accumulate partial results using spin-locks - else: - plock = locks + tl.program_id(2) * nlocks * tl.num_programs(1) + tl.program_id(1) * nlocks + lockid - 1 - pcount = plock + tl.num_programs(2) * tl.num_programs(1) * nlocks - while tl.atomic_cas(plock, 0, 1) == 1: - pass - count = tl.load(pcount) - if count == 0: - tl.store(pc, c, mask=checkc) - else: - d = tl.load(pc, mask=checkc) - tl.store(pc, d + c, mask=checkc) - tl.atomic_xchg(pcount, (count + 1) % maxid) - tl.atomic_xchg(plock, 0) - - -############## -# MAIN API # -############## -class _sparse_matmul(torch.autograd.Function): - - sdd_cache = dict() - dsd_cache = dict() - dds_cache = dict() - locks = dict() - - # Given an array sizes representing reduction size for each - # column of a block-mode matrix multiplication, - # performs load-balancing to achieve more smaller reductions - # between `seg_size` elements - @staticmethod - def load_balance(sizes, block): - #global triton - #if triton is None: - # triton = importlib.import_module('triton') - # segment size - # heuristics taken from OpenAI blocksparse code - # https://github.com/openai/blocksparse/blob/master/blocksparse/matmul.py#L95 - max_size = sizes.max() - min_size = sizes[sizes != 0].min() - #if max_size > min_size * 2.0: - # seg_max = max(triton.cdiv(max_size, 4), min_size*2) - #else: - # seg_max = max_size - seg_max = max_size - seg_min = max(triton.cdiv(seg_max, 4), 4) - # split reduction into segments - div = sizes // seg_max - rem = sizes % seg_max - packs = div + (sizes < seg_min).long() + (rem >= seg_min).long() - width = packs.sum() - segments = torch.empty(width, dtype=sizes.dtype) - column = torch.empty_like(segments) - lockid = torch.zeros_like(segments) - maxid = torch.zeros_like(segments) - nlocks = 0 - current = 0 - col_idx = 0 - for i in range(len(sizes)): - d, r = div[i], rem[i] - isempty = sizes[i] < seg_min - last = current + d + (r >= seg_min) + isempty - # column id - column[current:last] = col_idx - # lock id - if d > 1 or (d == 1 and r >= seg_min): - nlocks += 1 - lockid[current:last] = nlocks - maxid[current:last] = last - current - # segment size - segments[current:current + d] = seg_max - if r < seg_min and not isempty: - segments[current + d - 1] += r - if r >= seg_min or isempty: - segments[current + d] = r - current = last - col_idx += 1 - offsets = torch.zeros_like(segments) - offsets[1:] = torch.cumsum(segments[:-1], dim=0) - return segments, column, lockid, maxid, offsets - - @staticmethod - def get_locks(size, dev): - if dev not in _sparse_matmul.locks or \ - size > _sparse_matmul.locks[dev].size(0): - _sparse_matmul.locks[dev] = torch.zeros(size, dtype=torch.int32, device=dev) - return _sparse_matmul.locks[dev] - - ########################## - # SPARSE = DENSE x DENSE # - ########################## - - @staticmethod - def make_sdd_lut(layout, block, dtype, device): - #_sparse_matmul._load_utils() - #start_width = 64 // block - #segmented = _sparse_matmul.sdd_segment(layout.type(torch.int32), start_width) - start_width = (128 if block > 16 else 32) // block - layout = layout.type(torch.int32) - segmented = libtriton.superblock(layout.data_ptr(), layout.shape[0], layout.shape[1], layout.shape[2], - start_width) - luts, widths, packs = [], [], [] - for size, nnz in segmented: - """ width = nnz.shape[0] // (size * size) - h = nnz[:, 0] - i = nnz[:, 1] - j = nnz[:, 2] - b = nnz[:, 3] - lut = torch.stack((h, i, j, b), dim=1).view(-1).contiguous() - luts.append(lut.type(torch.int32).to(device)) - widths.append(width) - packs.append(size) """ - nnz = nnz.reshape(-1, 4) - width = nnz.shape[0] // (size * size) - luts.append(torch.from_numpy(nnz).type(torch.int32).to(device)) - widths.append(width) - packs.append(size) - # create locks - return luts, None, widths, packs - - @staticmethod - def _sdd_matmul(a, b, trans_a, trans_b, trans_c, spdims, block, luts, num_locks, widths, packs, bench, time): - if trans_c: - a, b = b, a - trans_a, trans_b = not trans_b, not trans_a - AS0 = a.size(0) - # Shape check - a_dim = -2 if trans_a else -1 - b_dim = -1 if trans_b else -2 - a_inner, b_inner = a.shape[a_dim], b.shape[b_dim] - if a_inner != b_inner: - raise ValueError(f"Size of tensor A along the {a_dim} dim ({a_inner}) must match size " - f"of tensor B along the {b_dim} dim ({b_inner})") - if a_inner % 16 != 0: - raise ValueError('Reduction size for SDD must be a multiple of 16') - - batch_size = a.size(0) - a_outer = a.size(3 if trans_a else 2) - dtype = a.dtype - is_16_multiple = a_inner % 16 == 0 - is_32_multiple = a_inner % 32 == 0 - is_64_multiple = a_inner % 64 == 0 - if not is_16_multiple: - raise ValueError('Reduction size for SDD must be a multiple of 16') - device = a.device - # create kernel - total_width = sum([width * pack * pack for width, pack in zip(widths, packs)]) - c = torch.empty((batch_size, total_width, block, block), dtype=dtype, device=a.device) - for lut, width, pack in zip(luts, widths, packs): - F32TK = [8, 16] - F16TK = [16] - F16TK += [32] if is_32_multiple else [] - F16TK += [64] if is_64_multiple else [] - TK = {torch.float32: F32TK, torch.float16: F16TK}[dtype] - num_lock = 1 - meta = { - 'TM': block * pack, - 'TN': block * pack, - 'BLOCK': block, - 'TK': TK[0], - 'TZ': 1, - 'SDD': True, - 'DSD': False, - 'DDS': False - } - # create output - locks = _sparse_matmul.get_locks(2 * width * AS0 * num_lock, a.device) - # maximum grid size is 65535 - # so operation might be decomposed into multiple - # kernel calls - max_width = 49152 - total = 0 if bench else None - for off_width in range(0, width, max_width): - grid = lambda meta: [meta['TZ'], min(max_width, width - off_width), batch_size] - _kernel[grid](a, - b, - c, - a.stride(0), - a.stride(1), - a.stride(3 if trans_a else 2), - a.stride(2 if trans_a else 3), - b.stride(0), - b.stride(1), - b.stride(3 if trans_b else 2), - b.stride(2 if trans_b else 3), - c.stride(0), - c.stride(0), - c.stride(2), - c.stride(3), - a_outer, - a_outer, - a_inner, - off_width, - lut, - locks, - num_lock, - num_warps=4, - **meta) - # save for backward pass - return c - - ########################## - # DENSE = DENSE x SPARSE # - ########################## - - # Given a binary layout of 0s and 1s, - # Construct look-up table for efficient execution on GPUs - @staticmethod - def make_dxx_lut(layout, block, step, trans, device, transform=lambda idx: idx): - # load-balancing - _empty = torch.tensor([], dtype=torch.int64, device=layout.device) - segments = _empty.clone() - column = _empty.clone() - depth = _empty.clone() - lockid = _empty.clone() - maxid = _empty.clone() - offsets = _empty.clone() - current_offset = 0 - current_maxid = 0 - for z in range(layout.size(0)): - if trans: - sizes = torch.sum(layout[z, :, :], 1) - else: - sizes = torch.sum(layout[z, :, :], 0) - z_segments, z_column, z_lockid, z_maxid, z_offsets = _sparse_matmul.load_balance(sizes, block) - z_depth = z * torch.ones_like(z_segments) - z_lockid[z_lockid > 0] += current_maxid - current_maxid = z_lockid.max() - # concatenate depth - segments = torch.cat((segments, z_segments)) - column = torch.cat((column, z_column)) - depth = torch.cat((depth, z_depth)) - maxid = torch.cat((maxid, z_maxid)) - offsets = torch.cat((offsets, current_offset + z_offsets)) - lockid = torch.cat((lockid, z_lockid)) - current_offset += layout[z, :, :].sum() - segments *= step - # pointer increments - if trans: - nnz = layout.nonzero() - else: - nnz = layout.transpose(1, 2).nonzero() - num_blocks = nnz.size(0) - offsets = torch.min(offsets, (num_blocks - 1) * torch.ones_like(offsets)) - idx = transform(nnz[:, 2] * block) - xincs = idx.clone() - xincs[1:] -= idx[:-1] - # divide block into multiple steps - div = block // step - xincs = xincs.view(-1, 1).repeat(1, div) - xincs[:, 1:] = step - xincs[:, 0] -= (div - 1) * step - # first increment for each reduction is actually the offset - xincs[offsets[segments > 0], 0] = idx[offsets[segments > 0]] - xincs = xincs.view(-1) - # block-mode input increments - if trans: - widx = torch.arange(num_blocks) - else: - widx = _empty.clone() - current_offset = 0 - for z in range(layout.size(0)): - layoutw = layout[z, :, :].clone() - msum = layoutw.sum() - layoutw[layoutw > 0] = 1 + torch.arange(msum) - widx = torch.cat((widx, current_offset + layoutw.T[layoutw.T > 0] - 1)) - current_offset += msum - widx = widx - wincs = widx * block * block - wincs[1:] -= widx[:-1] * block * block - wincs = wincs.view(-1, 1).repeat(1, div) - if trans: - wincs[:, 1:] = step - wincs[:, 0] -= (div - 1) * step - else: - wincs[:, 1:] = step * block - wincs[:, 0] -= (div - 1) * step * block - wincs[offsets[segments > 0], 0] = widx[offsets[segments > 0]] - wincs = wincs.view(-1) - # adjust offset and segment size - offsets *= 2 * div - segments *= div - # create header - width = column.size(0) - offsets += 6 * width - header = torch.stack((offsets, segments, column, depth, lockid, maxid), dim=1).view(-1).contiguous() - incs = torch.stack((xincs, wincs), dim=1).view(-1).contiguous() - incs = torch.cat((incs, torch.zeros(2, device=incs.device, dtype=incs.dtype))) - # create lut - lut = torch.cat((header, incs)) - lut = lut.type(torch.int32).to(device) - # create locks - num_locks = max(1, lockid.max()) - return lut, num_locks, width, None - - @staticmethod - def _dds_matmul(a, b, trans_a, trans_b, trans_c, spdims, block, lut, num_locks, width, packs, bench, time): - global triton - if triton is None: - triton = importlib.import_module('triton') - - # shapes / dtypes - AS0 = a.size(0) - AS1 = a.size(1) - AS2 = a.size(3 if trans_a else 2) - AS3 = a.size(2 if trans_a else 3) - BS0 = spdims[0] - BS1 = block * spdims[2 if trans_b else 1] - BS2 = block * spdims[1 if trans_b else 2] - dtype = a.dtype - # kernel - meta = {'TN': block, 'TM': 128, 'TK': 16, 'BLOCK': block, 'TZ': 1, 'SDD': False, 'DSD': False, 'DDS': True} - # output - CS0 = AS0 - CS1 = AS1 - CS2 = BS2 if trans_c else AS2 - CS3 = AS2 if trans_c else BS2 - locks = _sparse_matmul.get_locks(2 * AS0 * AS2 // 32 * num_locks, a.device) - c = torch.empty((CS0, CS1, CS2, CS3), dtype=dtype, device=a.device) - grid = lambda meta: [width, triton.cdiv(AS2, meta['TM']), AS0] - _kernel[grid](a, - b, - c, - a.stride(0), - a.stride(1), - a.stride(3 if trans_a else 2), - a.stride(2 if trans_a else 3), - b.stride(0), - b.stride(1), - b.stride(3 if trans_b else 2), - b.stride(2 if trans_b else 3), - c.stride(0), - c.stride(1), - c.stride(3 if trans_c else 2), - c.stride(2 if trans_c else 3), - AS2, - BS2, - 0, - 0, - lut, - locks, - num_locks, - num_warps=4, - **meta) - return c - - @staticmethod - def _dsd_matmul(a, b, trans_a, trans_b, trans_c, spdims, block, lut, num_locks, width, packs, bench, time): - global triton - if triton is None: - triton = importlib.import_module('triton') - - # shapes / dtypes - AS0 = spdims[0] - AS1 = block * spdims[2 if trans_a else 1] - AS2 = block * spdims[1 if trans_a else 2] - BS0 = b.size(0) - BS1 = b.size(1) - BS2 = b.size(3 if trans_b else 2) - BS3 = b.size(2 if trans_b else 3) - dtype = a.dtype - # kernel - - meta = {'TM': block, 'TN': 128, 'TK': 16, 'BLOCK': block, 'TZ': 1, 'SDD': False, 'DSD': True, 'DDS': False} - # output - CS0 = BS0 - CS1 = BS1 - CS2 = BS3 if trans_c else AS1 - CS3 = AS1 if trans_c else BS3 - locks = _sparse_matmul.get_locks(2 * BS0 * BS3 // 32 * num_locks, a.device) - c = torch.empty((CS0, CS1, CS2, CS3), dtype=dtype, device=a.device) - grid = lambda meta: [width, triton.cdiv(BS3, meta['TN']), BS0] - _kernel[grid](a, - b, - c, - a.stride(0), - a.stride(1), - a.stride(3 if trans_a else 2), - a.stride(2 if trans_a else 3), - b.stride(0), - b.stride(1), - b.stride(3 if trans_b else 2), - b.stride(2 if trans_b else 3), - c.stride(0), - c.stride(1), - c.stride(2), - c.stride(3), - BS3, - AS1, - 0, - 0, - lut, - locks, - num_locks, - num_warps=4, - **meta) - return c - - fn = {'sdd': _sdd_matmul.__get__(object), 'dsd': _dsd_matmul.__get__(object), 'dds': _dds_matmul.__get__(object)} - - @staticmethod - def forward(ctx, a, b, trans_a, trans_b, trans_c, mode, spdims, block, c_lut, c_num_locks, c_width, c_packs, - c_bench, c_time, da_lut, da_num_locks, da_width, da_packs, da_bench, da_time, db_lut, db_num_locks, - db_width, db_packs, db_bench, db_time): - c = _sparse_matmul.fn[mode](a, b, trans_a, trans_b, trans_c, spdims, block, c_lut, c_num_locks, c_width, - c_packs, c_bench, c_time) - # save for backward - ctx.save_for_backward(a, b) - ctx.da_num_locks = da_num_locks - ctx.da_lut = da_lut - ctx.da_width = da_width - ctx.da_packs = da_packs - ctx.da_bench = da_bench - ctx.da_time = da_time - ctx.db_lut = db_lut - ctx.db_num_locks = db_num_locks - ctx.db_width = db_width - ctx.db_bench = db_bench - ctx.db_packs = db_packs - ctx.db_time = db_time - ctx.mode = mode - ctx.spdims = spdims - ctx.block = block - ctx.trans_a = trans_a - ctx.trans_b = trans_b - return c - - @staticmethod - def backward(ctx, dc): - # saved for backward - a, b = ctx.saved_tensors - mode = ctx.mode - # gradients w.r.t. a - if ctx.needs_input_grad[0]: - mode_da = mode[1] + mode[0] + mode[2] - da = _sparse_matmul.fn[mode_da](dc, b, False, not ctx.trans_b, ctx.trans_a, ctx.spdims, ctx.block, - ctx.da_lut, ctx.da_num_locks, ctx.da_width, ctx.da_packs, ctx.da_bench, - ctx.da_time) - # gradients w.r.t. b - if ctx.needs_input_grad[1]: - mode_db = mode[2] + mode[1] + mode[0] - db = _sparse_matmul.fn[mode_db](a, dc, not ctx.trans_a, False, ctx.trans_b, ctx.spdims, ctx.block, - ctx.db_lut, ctx.db_num_locks, ctx.db_width, ctx.db_packs, ctx.db_bench, - ctx.db_time) - return da, db, None, None, None,\ - None, None, None, None,\ - None, None, None, None, None, None,\ - None, None, None, None, None, None,\ - None, None, None, None, None, None - - -class MatMul: - """Block-Sparse MatMul class; this class handles three types of matrix-multiplication: - - sparse = dense X dense - - dense = sparse X dense - - dense = dense X sparse - - For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 - """ - - def make_lut(self, dtype, device): - """Generates the sparsity layout/s used in block-sparse matmul - """ - key = (dtype, device) - if key in self.lut_cache: - return self.lut_cache[key] - # C look-up table - layout, block = self.layout, self.block - step = 16 - if self.mode == 'sdd': - c_lut, c_num_locks, c_width, c_packs = _sparse_matmul.make_sdd_lut(layout, block, dtype, device) - elif self.mode == 'dsd': - c_lut, c_num_locks, c_width, c_packs = _sparse_matmul.make_dxx_lut(layout, block, step, not self.trans_a, - device) - elif self.mode == 'dds': - c_lut, c_num_locks, c_width, c_packs = _sparse_matmul.make_dxx_lut(layout, block, step, self.trans_b, - device) - # DA look-up table - if self.mode == 'sdd': - da_lut, da_num_locks, da_width, da_packs = _sparse_matmul.make_dxx_lut(layout, block, step, True, device) - elif self.mode == 'dsd': - da_lut, da_num_locks, da_width, da_packs = _sparse_matmul.make_sdd_lut(layout, block, dtype, device) - elif self.mode == 'dds': - da_lut, da_num_locks, da_width, da_packs = _sparse_matmul.make_dxx_lut(layout, block, step, - not self.trans_b, device) - # DB look-up table - if self.mode == 'sdd': - db_lut, db_num_locks, db_width, db_packs = _sparse_matmul.make_dxx_lut(layout, block, step, False, device) - elif self.mode == 'dsd': - db_lut, db_num_locks, db_width, db_packs = _sparse_matmul.make_dxx_lut(layout, block, step, self.trans_a, - device) - elif self.mode == 'dds': - db_lut, db_num_locks, db_width, db_packs = _sparse_matmul.make_sdd_lut(layout, block, dtype, device) - self.lut_cache[key] = (c_lut, c_num_locks, c_width, c_packs,\ - da_lut, da_num_locks, da_width, da_packs,\ - db_lut, db_num_locks, db_width, db_packs) - return self.lut_cache[key] - - def __init__(self, layout, block, mode, trans_a=False, trans_b=False, bench=False): - """Initialize the Block-Sparse MatMul class. - - Arguments: - layout: required: sparsity layout tensor - block: required: an integer determining the block size. - mode: required: a string determining type of matmul; ('sdd') sparse = dense X dense, ('dsd') dense = sparse X dense, ('dds') dense = dense X sparse - trans_a: optional: a boolean determining if multiplication needs to be applied on transpose of input a; default is false - trans_b: optional: a boolean determining if multiplication needs to be applied on transpose of input b; default is false - bench: optional: set if you want to do benchmarking - """ - - if mode not in ['sdd', 'dsd', 'dds']: - raise NotImplementedError('Supported modes are: sdd, dsd, dds') - # look-up table cache - self.lut_cache = dict() - # attributes - self.trans_a = trans_a - self.trans_b = trans_b - self.mode = mode - self.block = block - self.layout = layout - layout_dim = layout.ndim - assert layout_dim in (2, 3), "Layout should be a 2 or 3 dimensional tensor of 0s and 1s" - if not mode == 'sdd': - # Dims to be reduced on the 'inside' of the matmul, either -1 or -2 - trans_dense, trans_sparse, sparse_inner = (trans_b, trans_a, -1) if mode == 'dsd' else (trans_a, trans_b, - -2) - self.dense_inner_dim = -((sparse_inner % 2) + 1) if not trans_dense else sparse_inner - sparse_inner = sparse_inner if not trans_sparse else -((sparse_inner % 2) + 1) - - # Inner dim of the dense input should be equal to the inner dim of the sparse input - self.dense_inner_size = layout.shape[sparse_inner] * block - # Expected shape for sparse inputs - self.sparse_shape = (layout.sum().item(), block, block) - - # Support using the same layout across attention heads etc. - if layout_dim == 2: - layout = layout.unsqueeze(0) - - layout = layout.long() # Above code assumes the layout tensor is an integral type - - self.spdims = layout.shape - # timings - self.bench = bench - self.time_c = None - self.time_da = None - self.time_db = None - - # pad shapes of a tensor to make it - # compatible with kernel calls - @staticmethod - def _pad_shape(x, is_sparse): - max_dim = 3 if is_sparse else 4 - for i in range(max_dim - x.dim()): - x = x.unsqueeze(0) - return x - - def __call__(self, a, b): - """Applies Block-Sparse MatMul. - - For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 - - Arguments: - a: required: a dense/block-sparse tensor; first input of mat-mul - b: required: a dense/block-sparse tensor; second input of mat-mul - - Return: - c: a dense/block-sparse tensor result of a X b - """ - - - c_lut, c_num_locks, c_width, c_packs,\ - da_lut, da_num_locks, da_width, da_packs,\ - db_lut, db_num_locks, db_width, db_packs = self.make_lut(a.dtype, a.device) - # timings - time_c = [None] - time_da = [None] - time_db = [None] - - original_dims = max(a.ndim, b.ndim) - a, b = self._validate_inputs(a, b) - - # pad shapes with ones - a = MatMul._pad_shape(a, self.mode == 'dsd') - b = MatMul._pad_shape(b, self.mode == 'dds') - # execute - - c = _sparse_matmul.apply(a, b, self.trans_a, self.trans_b, False, self.mode, self.spdims, self.block, c_lut, - c_num_locks, c_width, c_packs, self.bench, time_c, da_lut, da_num_locks, da_width, - da_packs, self.bench, time_da, db_lut, db_num_locks, db_width, db_packs, self.bench, - time_db) - - # This removes any leading singleton dimensions we may have added to the tensor that weren't in the input - dims_to_trim = c.ndim - original_dims - for _ in range(dims_to_trim): - c = c.squeeze(0) - - self.time_c = time_c[0] - self.time_da = time_da[0] - self.time_db = time_db[0] - return c - - def _validate_inputs(self, a, b): - if a.device != b.device: - raise ValueError(f"Inputs must be on the same device; got {a.device} for tensor A " - f"and {b.device} for tensor B") - if not get_accelerator().on_accelerator(a): - raise ValueError("Only GPU devices are supported for now") - - # When autocast is enabled, torch.matmul autocasts to float16, so we do the same here - if torch.is_autocast_enabled(): - a, b = a.half(), b.half() - elif a.dtype != b.dtype: - raise ValueError(f"Inputs must be the same dtype; got {a.dtype} for A and {b.dtype} for B") - - mode, trans_a, trans_b = self.mode, self.trans_a, self.trans_b - if mode != 'sdd': - # One input is sparse - dense, dense_name, sparse, sparse_name = (a, 'A', b, 'B') if mode == 'dds' else (b, 'B', a, 'A') - dense_inner = dense.shape[self.dense_inner_dim] - if dense_inner != self.dense_inner_size: - raise ValueError(f"Expected tensor {dense_name} to have size {self.dense_inner_size} at dim " - f"{self.dense_inner_dim % dense.ndim}, got {dense_inner}.") - - if sparse.shape[-len(self.sparse_shape):] != self.sparse_shape: - raise ValueError(f"Expected tensor with trailing dimensions of shape {self.sparse_shape} for argument " - f"{sparse_name}, got {sparse.shape}") - - def add_extra_dims(x): - # Add extra leading singleton dimensions if needed - dims_needed = 4 - x.ndim - if dims_needed > 0: - singletons = [1] * dims_needed - x = x.view(*singletons, *x.shape) - elif dims_needed < 0: - raise ValueError("Tensors with more than 4 dimensions are not currently supported") - - return x - - # Pad shapes with leading singleton dimensions - a = add_extra_dims(a) - b = add_extra_dims(b) - - return a, b diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/softmax.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/softmax.py deleted file mode 100644 index debee5688fe3f06699a710dbaa35c09601189a4f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/softmax.py +++ /dev/null @@ -1,296 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a -# https://github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/matmul.py - -import torch - -import triton -import triton.language as tl - - -def next_power_of_2(n): - n -= 1 - n |= n >> 1 - n |= n >> 2 - n |= n >> 4 - n |= n >> 8 - n |= n >> 16 - n += 1 - return n - - -def num_warps(n): - if n < 512: - return 4 - if n < 2048: - return 8 - return 16 - - -@triton.heuristics({'num_warps': lambda *args, **meta: num_warps(args[6] * meta['BLOCK'])}) -@triton.heuristics({'TN': lambda *args, **meta: next_power_of_2(args[6] * meta['BLOCK'])}) -@triton.jit -def _forward(X, scale, LUT, RPE, KP_M, ATTN_M, sizemax, stride_zx, stride_zrpe, stride_hrpe, stride_srpe, stride_zkpm, - stride_zattnm, **meta): - TN = meta['TN'] - BLOCK = meta['BLOCK'] - pidhm = tl.program_id(0) - pidz = tl.program_id(1) - # create index ranges - rxm = pidhm % BLOCK - rbm = pidhm // BLOCK - rxn = tl.arange(0, TN) % BLOCK - rbn = tl.arange(0, TN) // BLOCK - # extract information from LUT - header = LUT + rbm * 2 - size = tl.load(header + 0) - offset = tl.load(header + 1) - check = rbn < size - rbmn = tl.where(check, rbn, size - 1) - # block id and column id - blockid = tl.load(LUT + offset + rbmn * 4 + 0) - columnid = tl.load(LUT + offset + rbmn * 4 + 1) - rowid = tl.load(LUT + offset + rbmn * 4 + 2) - headid = tl.load(LUT + offset + rbmn * 4 + 3) - # pointers to X - px = X + pidz * stride_zx + blockid * BLOCK * BLOCK + rxm * BLOCK + rxn - x = tl.load(px, mask=check, other=-float('inf')) - x = x.to(tl.float32) - # apply scale - if meta['APPLY_SCALE']: - x = x * scale - # apply RPE - if meta['APPLY_RPE']: - prpe = RPE + pidz * stride_zrpe + headid * stride_hrpe + columnid * BLOCK + rowid * BLOCK * stride_srpe + rxm * stride_srpe + rxn - rpe = tl.load(prpe, mask=check, other=0) - x = x + rpe - # apply key-padding mask - if meta['APPLY_KP_MASK']: - pkp_m = KP_M + pidz * stride_zkpm + columnid * BLOCK + rxn - kp_m = tl.load(pkp_m, mask=check, other=-float('inf')) - if meta['KP_MASK_MUL']: - kp_m = tl.where(kp_m == 0, -float('inf'), 0.) - x = x + kp_m - # apply attention mask - if meta['APPLY_ATTN_MASK']: - pattn_m = ATTN_M + columnid * BLOCK + rowid * BLOCK * stride_zattnm + rxm * stride_zattnm + rxn - attn_m = tl.load(pattn_m, mask=check, other=-float('inf')) - if meta['ATTN_MASK_MUL']: - attn_m = tl.where(attn_m == 0, -float('inf'), 0.) - x = x + attn_m - # computation - x = tl.softmax(x) - tl.store(px, x, mask=check) - - -@triton.heuristics({'num_warps': lambda *args, **meta: num_warps(args[4] * meta['BLOCK'])}) -@triton.heuristics({'TN': lambda *args, **meta: next_power_of_2(args[4]) * meta['BLOCK']}) -@triton.jit -def _backward(X, scale, DX, LUT, sizemax, stride_zx, stride_zdx, **meta): - pidhm = tl.program_id(0) - pidz = tl.program_id(1) - TN = meta['TN'] - BLOCK = meta['BLOCK'] - # create index ranges - rxm = pidhm % BLOCK - rbm = pidhm // BLOCK - rxn = tl.arange(0, TN) % BLOCK - rbn = tl.arange(0, TN) // BLOCK - # extract information from look-up table - header = LUT + rbm * 2 - size = tl.load(header + 0) - offset = tl.load(header + 1) - # bounds checking on lut - check = rbn < size - rbmn = tl.where(check, rbn, size - 1) - # initialize pointers to block-sparse input - blockid = tl.load(LUT + offset + rbmn * 4) - X = X + pidz * stride_zx + blockid * BLOCK * BLOCK + rxm * BLOCK + rxn - DX = DX + pidz * stride_zdx + blockid * BLOCK * BLOCK + rxm * BLOCK + rxn - # compute fused softmax backward - x = tl.load(X, mask=check, other=0) - dx = tl.load(DX, mask=check, other=0) - x = x.to(tl.float32) - dx = dx.to(tl.float32) - y = x * (dx - tl.sum(x * dx, 0)) * scale - tl.store(DX, y, mask=check) - - -class _sparse_softmax(torch.autograd.Function): - - bwd_kernels = dict() - - @staticmethod - def make_lut(layout, block, device): - _empty = torch.tensor([], dtype=torch.int64, device=layout.device) - sizes = _empty.clone() - # sizes along rows - for h in range(layout.shape[0]): - sizes = torch.cat((sizes, layout[h, :, :].sum(-1))) - # offsets in block format - offsets = torch.zeros_like(sizes) - offsets[1:] = torch.cumsum(sizes[:-1], dim=0) - # block indices - idx = torch.arange(layout.sum()) - head = layout.nonzero()[:, 0] - rows = layout.nonzero()[:, 1] - columns = layout.nonzero()[:, 2] - core = torch.stack((idx, columns, rows, head), dim=1).view(-1) - # construct look-up table - offsets = offsets * 4 + 2 * sizes.numel() - header = torch.stack((sizes, offsets), dim=1).view(-1) - lut = torch.cat((header, core)).type(torch.int32).to(device) - return lut, int(sizes.max()) - - @staticmethod - def forward(ctx, x, scale, rpe, key_padding_mask, attn_mask, kp_mask_mode, attn_mask_mode, spdims, block, lut, - num_blocks, maxlut, bench, time): - - apply_scale = False if scale == 1.0 else True - - # handle None rpe - if rpe is None: - apply_rpe = False - stride_zrpe, stride_hrpe, stride_srpe = 0, 0, 0 - rpe = torch.empty(0, dtype=x.dtype, device=x.device) - else: - apply_rpe = True - stride_zrpe, stride_hrpe, stride_srpe = rpe.stride(0), rpe.stride(1), rpe.stride(2) - - # handle None key_padding_mask - if key_padding_mask is None: - apply_kp_mask = False - stride_zkpm = 0 - key_padding_mask = torch.empty(0, dtype=x.dtype, device=x.device) - else: - apply_kp_mask = True - stride_zkpm = key_padding_mask.stride(0) - - # handle None attention_mask - if attn_mask is None: - apply_attn_mask = False - stride_zattnm = 0 - attn_mask = torch.empty(0, dtype=x.dtype, device=x.device) - else: - apply_attn_mask = True - stride_zattnm = attn_mask.stride(0) - - # run kernel - M = x.shape[0] - meta = { - 'BLOCK': block, - 'APPLY_SCALE': apply_scale, - 'APPLY_RPE': apply_rpe, - 'APPLY_KP_MASK': apply_kp_mask, - 'APPLY_ATTN_MASK': apply_attn_mask, - 'KP_MASK_MUL': kp_mask_mode == 'mul', - 'ATTN_MASK_MUL': attn_mask_mode == 'mul', - } - grid = lambda opt: [spdims[0] * spdims[1] * block, M] - _forward[grid](x, scale, lut, rpe, key_padding_mask, attn_mask, maxlut, x.stride(0),\ - stride_zrpe, stride_hrpe, stride_srpe, stride_zkpm, stride_zattnm, **meta) - - # save to context - ctx.mark_dirty(x) - ctx.save_for_backward(x, lut) - ctx.spdims = spdims - ctx.block = block - ctx.maxlut = maxlut - ctx.scale = scale - ctx.apply_scale = apply_scale - ctx.apply_rpe = apply_rpe - ctx.apply_kp_mask = apply_kp_mask - ctx.apply_attn_mask = apply_attn_mask - ctx.kp_mask_mode = kp_mask_mode - ctx.attn_mask_mode = attn_mask_mode - return x - - @staticmethod - def backward(ctx, dx): - - # retrieve from context - x, lut = ctx.saved_tensors - # run kernel - M = x.shape[0] - grid = lambda opt: [ctx.spdims[0] * ctx.spdims[1] * ctx.block, M] - _backward[grid](x, ctx.scale, dx, lut, ctx.maxlut, x.stride(0), dx.stride(0), BLOCK=ctx.block) - return dx, None, None, None, None, None, None, None, None, None, None, None, None, None, None - - -class Softmax: - """Block-Sparse Softmax class; this class computes softmax on a block sparse matrix. It is also able to apply either/all of the following masks: - - relative position embedding - - key padding mask - - attention mask - - For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 - """ - - def sparse_softmax(*args, **kwargs): - return _sparse_softmax.apply(*args, **kwargs) - - def make_lut(self, device): - """Generates the sparsity layout used in block-sparse softmax - """ - key = (device, ) - if key not in self.lut_cache: - self.lut_cache[key] = _sparse_softmax.make_lut(self.layout, self.block, device) - return self.lut_cache[key] - - def __init__(self, layout, block, bench=False): - """Initialize the Block-Sparse Softmax class. - - Arguments: - layout: required: sparsity layout tensor - block: required: an integer determining the block size. - bench: optional: set if you want to do benchmarking - """ - - self.num_blocks = layout.sum().item() - self.spdims = layout.shape - self.layout = layout - self.block = block - self.bench = bench - self.lut_cache = dict() - - def __call__(self, - x, - scale=1., - rpe=None, - key_padding_mask=None, - attn_mask=None, - key_padding_mask_mode='add', - attn_mask_mode='add'): - """Applies softmax on a Block-Sparse input tensor. - - For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 - - Arguments: - x: required: a block-sparse tensor that softmax is applied on it; computation will be in place and result will be returned in the same tensor - scale: optional: a float value; x values will be multiplied by this value before normalization. Default value is 1.0. - rpe: optional: a tensor same dimension as x that is used as relative position embedding - key_padding_mask: optional: a mask tensor of size (BatchSize X SequenceLength) - attn_mask: optional: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported - key_padding_mask_mode: optional: a boolean determining if key_padding_mask needs to be added or multiplied - attn_mask_mode: optional: a boolean determining if attn_mask needs to be added or multiplied - - Return: - x: a block-sparse tensor contains normalized input x using softmax; and masks applied if given - """ - - time_y = [None] - if rpe is not None and rpe.dtype != x.dtype: - raise ValueError('relative position embedding must be %s' % x.dtype) - if attn_mask is not None and attn_mask.dtype != x.dtype: - raise ValueError('Attention mask must be %s' % x.dtype) - if key_padding_mask is not None and key_padding_mask.dtype != x.dtype: - raise ValueError('Key padding mask must be %s' % x.dtype) - lut, maxlut = self.make_lut(x.device) - x = Softmax.sparse_softmax(x, scale, rpe, key_padding_mask, attn_mask, key_padding_mask_mode, attn_mask_mode, - self.spdims, self.block, lut, self.num_blocks, maxlut, self.bench, time_y) - self.time_y = time_y[0] - return x diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparse_attention_utils.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparse_attention_utils.py deleted file mode 100644 index ccb0f940dff65839beac579f81c4dfb7e499e6bb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparse_attention_utils.py +++ /dev/null @@ -1,208 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from torch.nn import functional as F -from deepspeed.ops.sparse_attention import BertSparseSelfAttention, SparsityConfig -''' -This file contains few utility functions to handle adapting pretrained model with sparse self-attention module. -''' - - -class SparseAttentionUtils: - """This class provides some utility functions that are use integrating sparse attention into transformer models. - Such utilities include extending position embeddings, replacing current self-attention layer with sparse attention, padding sequences to multiple of block size, etc. - - """ - - @staticmethod - def extend_position_embedding(model, max_position): - """This function extends the position embedding weights of a model loaded from a checkpoint. - It assumes the new max position is bigger than the original max length. - - Arguments: - model: required: a transformer model - max_position: required: an integer determining new position embedding size - Return: - model: updated model; in which position embedding weights have been extended based on new size - """ - - if hasattr(model, 'bert'): - original_max_position = model.bert.embeddings.position_embeddings.weight.size(0) - assert max_position > original_max_position - extend_multiples = max(1, max_position // original_max_position) - model.bert.embeddings.position_embeddings.weight.data = model.bert.embeddings.position_embeddings.weight.repeat( - extend_multiples, 1) - elif hasattr(model, 'roberta'): - # RoBERTa has positions 0 & 1 reserved, so embedding size is max position + 2 - original_max_position, embed_size = model.roberta.embeddings.position_embeddings.weight.shape - original_max_position -= 2 - extend_multiples = max(1, max_position // original_max_position) - assert max_position > original_max_position - max_position += 2 - extended_position_embedding = model.roberta.embeddings.position_embeddings.weight.new_empty( - max_position, embed_size) - k = 2 - for i in range(extend_multiples): - extended_position_embedding[k:( - k + original_max_position)] = model.roberta.embeddings.position_embeddings.weight[2:] - k += original_max_position - model.roberta.embeddings.position_embeddings.weight.data = extended_position_embedding - else: - raise ValueError( - 'Please extend \"extend_position_embedding\" function to support your model type. It currently only supports \"bert\" & \"roberta\"!' - ) - - model.config.max_position_embeddings = max_position - print(f'Extended position embeddings to {original_max_position * extend_multiples}') - - return model - - @staticmethod - def update_tokenizer_model_max_length(tokenizer, max_position): - """This function updates the position embedding length of a tokenizer to a new max position. - - Arguments: - tokenizer: required: a transformer tokenizer - max_position: required: an integer determining new position embedding size - Return: - tokenizer: updated tokenizer; in which model maximum length has been extended based on new size - """ - - tokenizer.model_max_length = max_position - tokenizer.init_kwargs['model_max_length'] = max_position - print(f'updated tokenizer model max imum length to {max_position}') - - return tokenizer - - @staticmethod - def replace_model_self_attention_with_sparse_self_attention( - model, - max_position, - # SparsityConfig parameters needs to be set accordingly - sparsity_config=SparsityConfig(num_heads=4)): - """This function replaces the self attention layers in model encoder with sparse self attention. - It currently supports bert and roberta model and can be easily extended to any other models following similar steps here. - For sparsityConfig, refer to the config class. - - Arguments: - model: required: a transformer model - max_position: required: an integer determining new position embedding size - sparsity_config: optional: this parameter determines sparsity pattern configuration; it is based on SparsityConfig class - - Return: - model: updated model; in which self attention layer has been replaced with DeepSpeed Sparse Self Attention layer. - """ - - if hasattr(model, 'bert'): - model.config.max_position_embeddings = max_position - model.replace_self_attention_layer_with_sparse_self_attention_layer(model.config, model.bert.encoder.layer, - sparsity_config) - elif hasattr(model, 'roberta'): - model.config.max_position_embeddings = max_position + 2 - model.replace_self_attention_layer_with_sparse_self_attention_layer(model.config, - model.roberta.encoder.layer, - sparsity_config) - else: - raise ValueError( - 'Please extend \"update_model_self_attention_to_sparse_self_attention\" function to support \ - your model type. It currently only supports \"bert\" & \"roberta\"!') - return model - - @staticmethod - def replace_self_attention_layer_with_sparse_self_attention_layer( - config, - layers, - # SparsityConfig parameters needs to be set accordingly - sparsity_config=SparsityConfig(num_heads=4)): - """This function replaces the self attention layers in attention layer with sparse self attention. - For sparsityConfig, refer to the config class. - - Arguments: - config: required: transformer model config - layers: required: transformer model attention layers - sparsity_config: optional: this parameter determines sparsity pattern configuration; it is based on SparsityConfig class - - Return: - layers: updated attention layers; in which self attention layers have been replaced with DeepSpeed Sparse Self Attention layer. - """ - - for layer in layers: - deepspeed_sparse_self_attn = BertSparseSelfAttention(config, sparsity_config) - deepspeed_sparse_self_attn.query = layer.attention.self.query - deepspeed_sparse_self_attn.key = layer.attention.self.key - deepspeed_sparse_self_attn.value = layer.attention.self.value - - layer.attention.self = deepspeed_sparse_self_attn - - return layers - - @staticmethod - def pad_to_block_size(block_size, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds, - pad_token_id, model_embeddings): - """This function pads input tokens and attention mask on sequence length dimension to be multiple of block size. - This is a requirement for Sparse Transformer in which the self attention layer works on sequences of length multiple of block size. - It needs to be called in your model, such as BertModel, right before you calculate the embedding outputs. - Note) - 1- instead of passing your embedding layer to this function, you can simply add this function to your model. It can be more simplified if given attention_mask and/or token_type_ids are none. - 2- you need to call unpad function before returning your model output to unpad the encoder sequence output. - - Arguments: - block_size: required: an integer determining the block size of sparsity config. - pad_token_id: required: an integer determining the pad token from the model config; such as bert.config.pad_token_id. - input_ids: a torch.LongTensor of shape [batch_size, sequence_length] with the word token indices in the vocabulary - attention_mask: a torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max input sequence length in the current batch. It's the mask that we typically use for attention when a batch has varying length sentences. - token_type_ids: a torch.LongTensor of shape [batch_size, sequence_length] with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to a `sentence B` token (see BERT paper for more details). - position_ids: a torch.LongTensor of shape [batch_size, sequence_length] with the indices of positions of each input sequence tokens in the position embeddings. - inputs_embeds: an optional torch.FloatTensor of shape [batch_size, sequence_length, hidden_size] that contains embedded representation and can be passed instead of input_ids directly. - model_embeddings: an optional object. If inputs_embeds are not none, this will be your model embeddings such as BertEmbeddings from your model such as BertModel. You can move this function inside your model and use self.embeddings instead of passing this parameter. - - Return: - pad_len: an integer determining how much inputs have been padded to transfer sequence length dimension to multiple of block size. - input_ids: if input_ids are not none padded input_ids otherwise none. - attention_mask: if attention_mask is not none padded attention_mask otherwise none. - token_type_ids: if token_type_ids are not none padded token_type_ids otherwise none. - position_ids: if position_ids are not none padded position_ids otherwise none. - inputs_embeds: if inputs_embeds are not none padded inputs_embeds otherwise none. - """ - - batch_size, seq_len = input_ids.shape if input_ids is not None else inputs_embeds.shape[:-1] - - pad_len = (block_size - seq_len % block_size) % block_size - if pad_len > 0: - if inputs_embeds is not None: - pad_input_ids = inputs_embeds.new_full((batch_size, pad_len), pad_token_id, dtype=torch.long) - pad_inputs_embeds = model_embeddings(pad_input_ids) - inputs_embeds = torch.cat([inputs_embeds, pad_inputs_embeds], dim=-2) - # may not be needed as input_ids are not used if inputs_embeds are given - if input_ids is not None: - input_ids = F.pad(input_ids, (0, pad_len), value=pad_token_id) - if position_ids is not None: - # pad position_id with pad_token_id - position_ids = F.pad(position_ids, (0, pad_len), value=pad_token_id) - # pad attention mask without attention on the padding tokens - attention_mask = F.pad(attention_mask, (0, pad_len), value=False) - # pad token_type_ids with token_type_id = 0 - token_type_ids = F.pad(token_type_ids, (0, pad_len), value=0) - - return pad_len, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds - - @staticmethod - def unpad_sequence_output(pad_len, sequence_output): - """This function unpads sequence output if inputs of the model were padded. - This is a requirement for Sparse Transformer in which the self attention layer works on sequences of length multiple of block size. - It needs to be called in your model, such as BertModel, right before you return the model outputs. - - Arguments: - pad_len: required: an integer determining how much model inputs have been padded to transfer sequence length dimension to multiple of block size. - sequence_output: required: sequence output of the encoder layer. - - Return: - sequence_output: unpaded sequence output of the encoder layer. - """ - - if (pad_len > 0): - sequence_output = sequence_output[:, :-pad_len] - return sequence_output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparse_self_attention.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparse_self_attention.py deleted file mode 100644 index b673c4561902e943981ca3008fae53ec73c0cd73..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparse_self_attention.py +++ /dev/null @@ -1,149 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch.nn as nn -import torch -from torch import distributed as dist -from deepspeed.ops.sparse_attention import SparsityConfig - - -class SparseSelfAttention(nn.Module): - """Implements an efficient Sparse Self Attention of Transformer layer based on `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 - - For more information please see, TODO DeepSpeed Sparse Transformer. - - For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial. - """ - - def __init__( - self, - # SparsityConfig parameters needs to be set accordingly - sparsity_config=SparsityConfig(num_heads=4), - key_padding_mask_mode='add', - attn_mask_mode='mul', - max_seq_length=2048): - """Initialize the sparse self attention layer. - Arguments: - sparsity_config: optional: this parameter determines sparsity pattern configuration; it is based on SparsityConfig class. - key_padding_mask_mode: optional: a string determining if key padding mask needs to be added, `add`, or be multiplied, `mul`. - attn_mask_mode: optional: a string determining if attention mask needs to be added, `add`, or be multiplied, `mul`. - max_seq_length: optional: the maximum sequence length this sparse attention module will be applied to; it controls the size of the master_layout. - """ - super().__init__() - - # sparsity information - self.sparsity_config = sparsity_config - - # initialize sparse layout and register as buffer - master_layout = self.sparsity_config.make_layout(max_seq_length) - self.register_buffer("master_layout", master_layout) - self._need_layout_synchronization = True - - # mask modes - self.key_padding_mask_mode = key_padding_mask_mode - self.attn_mask_mode = attn_mask_mode - - ops = dict() - - def get_layout(self, L): - # if layout is never synchronized across GPUs, broadcast the layout from global rank 0 - if self._need_layout_synchronization and dist.is_initialized(): - dist.broadcast(self.master_layout, src=0) - self._need_layout_synchronization = False - - if (L % self.sparsity_config.block != 0): - raise ValueError( - f'Sequence Length, {L}, needs to be dividable by Block size {self.sparsity_config.block}!') - - num_blocks = L // self.sparsity_config.block - return self.master_layout[..., :num_blocks, :num_blocks].cpu() # layout needs to be a CPU tensor - - # add to cache - def get_ops(self, H, L): - from deepspeed.ops.sparse_attention.matmul import MatMul - from deepspeed.ops.sparse_attention.softmax import Softmax - if L not in SparseSelfAttention.ops: - sparsity_layout = self.get_layout(L) - sparse_dot_sdd_nt = MatMul(sparsity_layout, self.sparsity_config.block, 'sdd', trans_a=False, trans_b=True) - - sparse_dot_dsd_nn = MatMul(sparsity_layout, - self.sparsity_config.block, - 'dsd', - trans_a=False, - trans_b=False) - - sparse_softmax = Softmax(sparsity_layout, self.sparsity_config.block) - - SparseSelfAttention.ops[L] = (sparse_dot_sdd_nt, sparse_dot_dsd_nn, sparse_softmax) - return SparseSelfAttention.ops[L] - - def transpose_key_for_scores(self, x, L): - bsz, num_heads, seq_len, head_dim = x.size() - if seq_len != L: - return x.permute(0, 1, 3, 2) - return x - - def transpose_mask_for_sparse(self, qtype, x, is_key_padding_mask=False): - x = x.type(qtype) - if is_key_padding_mask: - xdim = x.dim() - for d in range(xdim - 1, 0, -1): - x = x.squeeze(dim=d) - return x - return x.squeeze() - - # forward pass - def forward(self, query, key, value, rpe=None, key_padding_mask=None, attn_mask=None): - """Applies forward phase of sparse self attention - - Arguments: - query: required: query tensor - key: required: key tensor - value: required: value tensor - rpe: optional: a tensor same dimension as x that is used as relative position embedding - key_padding_mask: optional: a mask tensor of size (BatchSize X SequenceLength) - attn_mask: optional: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported - key_padding_mask_mode: optional: a boolean determining if key_padding_mask needs to be added or multiplied - attn_mask_mode: optional: a boolean determining if attn_mask needs to be added or multiplied - - Return: - attn_output: a dense tensor containing attention context - """ - assert query.dtype == torch.half, "sparse attention only supports training in fp16 currently, please file a github issue if you need fp32 support" - bsz, num_heads, tgt_len, head_dim = query.size() - - # transpose back key if it is already transposed - key = self.transpose_key_for_scores(key, tgt_len) - - # check that operation is supported - if query.shape != key.shape or key.shape != value.shape: - raise NotImplementedError('only self-attention is supported for now') - - # squeeze key_padding_mask if it is given - if key_padding_mask is not None: - key_padding_mask = self.transpose_mask_for_sparse(query.dtype, key_padding_mask, is_key_padding_mask=True) - - # squeeze attn_mask if it is given - if attn_mask is not None: - attn_mask = self.transpose_mask_for_sparse(query.dtype, attn_mask) - - # cache look-up table computations etc - sparse_dot_sdd_nt, sparse_dot_dsd_nn, sparse_softmax = self.get_ops(num_heads, tgt_len) - - scaling = float(head_dim)**-0.5 - - # attention scores - attn_output_weights = sparse_dot_sdd_nt(query, key) - attn_output_weights = sparse_softmax(attn_output_weights, - scale=scaling, - rpe=rpe, - key_padding_mask=key_padding_mask, - attn_mask=attn_mask, - key_padding_mask_mode=self.key_padding_mask_mode, - attn_mask_mode=self.attn_mask_mode) - - # outputs - attn_output = sparse_dot_dsd_nn(attn_output_weights, value) - return attn_output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparsity_config.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparsity_config.py deleted file mode 100644 index 1f59c4b469282ebfdb26fdb7c40c61d738231030..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/sparsity_config.py +++ /dev/null @@ -1,727 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import random - - -class SparsityConfig: - """Abstract Configuration class to store `sparsity configuration of a self attention layer`. - It contains shared property of different block-sparse sparsity patterns. However, each class needs to extend it based on required property and functionality. - """ - - def __init__(self, num_heads, block=16, different_layout_per_head=False): - """Initialize the Sparsity Pattern Config. - - For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial - - Arguments: - num_heads: required: an integer determining number of attention heads of the layer. - block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. - different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. - """ - - self.num_heads = num_heads - self.block = block - self.different_layout_per_head = different_layout_per_head - self.num_layout_heads = num_heads if different_layout_per_head else 1 - - def setup_layout(self, seq_len): - """Create layout tensor for the given sequence length - - Arguments: - seq_len: required: an integer determining number of attention heads of the layer. - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) for sparsity layout of all head; initialized with zero - """ - - if (seq_len % self.block != 0): - raise ValueError(f'Sequence Length, {seq_len}, needs to be dividable by Block size {self.block}!') - num_blocks = seq_len // self.block - # TODO Currently we allocate layout per head; needs to be updated if heads share a single layout. - layout = torch.zeros((self.num_heads, num_blocks, num_blocks), dtype=torch.int64) - return layout - - def check_and_propagate_first_head_layout(self, layout): - """If all heads require same sparsity layout, it propagate first head layout to all heads - - Arguments: - layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head - """ - - if not self.different_layout_per_head: - layout[1:self.num_heads, :, :] = layout[0, :, :] - return layout - - -class DenseSparsityConfig(SparsityConfig): - """Configuration class to store `Dense` configuration. - In reality, this is not sparse and all blocks are used. We keep it for the sake of comparison and comprehension. - """ - - def __init__(self, num_heads, block=16, different_layout_per_head=False): - """Initialize the Dense Sparsity Pattern Config. - In reality, this is not sparse and all blocks are used. We keep it for the sake of comparison and comprehension. - - Arguments: - num_heads: required: an integer determining number of attention heads of the layer. - seq_len: required: an integer determining number of attention heads of the layer. - different_layout_per_head: optional: this is just for the sake of consistency with other sparsity formats; can ignore it for DenseSparsityConfig - """ - - super().__init__(num_heads, block, different_layout_per_head) - - def make_layout(self, seq_len): - """Set 1 to all blocks of the layout meaning the pattern is dense; not sparse. - - Arguments: - seq_len: required: an integer determining the underling sequence length; must be <= max sequence length - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; for dense everything is 1 - """ - - layout = self.setup_layout(seq_len) - layout[:, :, :] = 1 - return layout - - -class FixedSparsityConfig(SparsityConfig): - """Configuration class to store `Fixed` sparsity configuration. - For more details about this sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509; this has been customized. - This class extends parent class of `SparsityConfig` and customizes it for `Fixed` sparsity. - """ - - def __init__(self, - num_heads, - block=16, - different_layout_per_head=False, - num_local_blocks=4, - num_global_blocks=1, - attention='bidirectional', - horizontal_global_attention=False, - num_different_global_patterns=1): - """Initialize `Fixed` Sparsity Pattern Config. - - For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial - - Arguments: - num_heads: required: an integer determining number of attention heads of the layer. - block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. - different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. - num_local_blocks: optional: an integer determining the number of blocks in local attention window. - num_global_blocks: optional: an integer determining how many consecutive blocks in a local window is used as the representative of the window for global attention. - attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. - horizontal_global_attention: optional: a boolean determining if blocks that are global representative of a local window, also attend to all other blocks. This is valid only if attention type is `bidirectional`. Looking at the attention matrix, that means global attention not only includes the vertical blocks, but also horizontal blocks. - num_different_global_patterns: optional: an integer determining number of different global attentions layouts. While global attention can be fixed by which block/s are representative of any local window, since there are multi-heads, each head can use a different global representative. For example, with 4 blocks local window and global attention size of 1 block, we can have 4 different versions in which the first, Second, third, or forth block of each local window can be global representative of that window. This parameter determines how many of such patterns we want. Of course, there is a limitation based on num_local_blocks and num_global_blocks. - """ - - super().__init__(num_heads, block, different_layout_per_head) - - self.num_local_blocks = num_local_blocks - - if (num_local_blocks % num_global_blocks != 0): - raise ValueError( - f'Number of blocks in a local window, {num_local_blocks}, must be dividable by number of global blocks, {num_global_blocks}!' - ) - self.num_global_blocks = num_global_blocks - - if (attention != 'unidirectional' and attention != 'bidirectional'): - raise NotImplementedError('only \"uni/bi-directional\" attentions are supported for now!') - self.attention = attention - - if (attention != 'bidirectional' and horizontal_global_attention): - raise ValueError('only \"bi-directional\" attentions can support horizontal global attention!') - self.horizontal_global_attention = horizontal_global_attention - - if (num_different_global_patterns > 1 and not different_layout_per_head): - raise ValueError( - f'Number of different layouts cannot be more than one when you have set a single layout for all heads! Set different_layout_per_head to True.' - ) - if (num_different_global_patterns > (num_local_blocks // num_global_blocks)): - raise ValueError( - f'Number of layout versions (num_different_global_patterns), {num_different_global_patterns}, cannot be larger than number of local window blocks divided by number of global blocks, {num_local_blocks} / {num_global_blocks} = {num_local_blocks//num_global_blocks}!' - ) - self.num_different_global_patterns = num_different_global_patterns - - def set_local_layout(self, h, layout): - """Sets local attention layout used by the given head in the sparse attention. - - Arguments: - h: required: an integer determining head index - layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local layout is set - """ - - num_blocks = layout.shape[1] - for i in range(0, num_blocks, self.num_local_blocks): - end = min(i + self.num_local_blocks, num_blocks) - for row in range(i, end): - for col in range(i, (row + 1 if self.attention == 'unidirectional' else end)): - layout[h, row, col] = 1 - return layout - - def set_global_layout(self, h, layout): - """Sets global attention layout used by the given head in the sparse attention. - - Currently we set global blocks starting from the last block of a local window to the first one. That means if a local window consists of 4 blocks and global attention size is one block, we use block #4 in each local window as global. If we have different layout per head, then other heads will get #3, #2, and #1. And if we have more heads (and different layout has set) than num of global attentions, multiple head may have same global attentions. - Note) if horizontal_global_attention is set, global blocks will be set both horizontally and vertically. - - Arguments: - h: required: an integer determining head index - layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which global layout is set - """ - - num_blocks = layout.shape[1] - first_global_block_idx = self.num_local_blocks - ( - 1 + h % self.num_different_global_patterns) * self.num_global_blocks - - # set all global blocks except the last one if (in last local window) - end = num_blocks - (num_blocks % self.num_local_blocks) - for i in range(first_global_block_idx, end, self.num_local_blocks): - - # vertical global attention - first_row = 0 if self.attention == 'bidirectional' else i - #(((i // self.num_local_blocks) + 1) * self.num_local_blocks) - #if (first_row < num_blocks): - layout[h, first_row:, i:i + self.num_global_blocks] = 1 - - # horizontal global attention; only in bidirectional attention - if (self.horizontal_global_attention): - layout[h, i:i + self.num_global_blocks, :] = 1 - - # set last global blocks; handle possible short last local window - if (end < num_blocks): - start = min(end + first_global_block_idx, num_blocks - self.num_global_blocks) - end = start + self.num_global_blocks - - # vertical global attention - first_row = 0 if self.attention == 'bidirectional' else start - #(((start // self.num_local_blocks) + 1) * self.num_local_blocks) - #if (first_row < num_blocks): - layout[h, first_row:, start:end] = 1 - - # horizontal global attention - if (self.horizontal_global_attention): - layout[h, start:end, :] = 1 - return layout - - def make_layout(self, seq_len): - """Generates `Fixed` sparsity layout used by each head in the sparse attention. - - Arguments: - seq_len: required: an integer determining number of attention heads of the layer. - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `Fixed` sparsity layout of all head - """ - - layout = self.setup_layout(seq_len) - for h in range(0, self.num_layout_heads): - layout = self.set_local_layout(h, layout) - layout = self.set_global_layout(h, layout) - - layout = self.check_and_propagate_first_head_layout(layout) - return layout - - -class VariableSparsityConfig(SparsityConfig): - """Configuration class to store `Variable` sparsity configuration. - This layout is an extension of FixedSparsityConfig in which: - - user can set random layout; default value is zero means no random block - - user can provide a list of local block sizes - - user can provide a list of global block indices. - - For more details about `Fixed` sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509; this has been customized. - This class extends parent class of `SparsityConfig` and customizes it for `Fixed` sparsity. - """ - - def __init__(self, - num_heads, - block=16, - different_layout_per_head=False, - num_random_blocks=0, - local_window_blocks=[4], - global_block_indices=[0], - global_block_end_indices=None, - attention='bidirectional', - horizontal_global_attention=False): - """Initialize `Variable` Sparsity Pattern Config. - - For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial - - Arguments: - num_heads: required: an integer determining number of attention heads of the layer. - block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. - different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. Currently this sparsity config can only assign single layout to all heads; needs to be extended for different layout per head. - num_random_blocks: optional: an integer determining the number of random blocks in each block row. - local_window_blocks: optional: a list of integers determining the number of blocks in each local attention window. It assumes first number determines # of blocks in the first local window, second the second window, ..., and the last number determines the number of blocks in the remaining local windows. - global_block_indices: optional: a list of integers determining which blocks are considered as global attention. Given indices, determine the blocks that all other token blocks attend to and they attend to all other token blocks. Default value is only index 0. Notice that if global_block_end_indices parameter is set, this parameter is used as starting index of each global window. - global_block_end_indices: optional: a list of integers determining end indices of global window blocks. By default this is not used. But if it is set, it must have the same size of global_block_indices parameter, and combining this two parameters, for each index i, blocks from global_block_indices[i] to global_block_end_indices[i] (exclusive) are considered as global attention. - num_global_blocks: optional: an integer determining how many consecutive blocks in a local window is used as the representative of the window for global attention. - attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. - horizontal_global_attention: optional: a boolean determining if blocks that are global representative of a local window, also attend to all other blocks. This is valid only if attention type is `bidirectional`. Looking at the attention matrix, that means global attention not only includes the vertical blocks, but also horizontal blocks. - """ - - super().__init__(num_heads, block, different_layout_per_head) - - self.num_random_blocks = num_random_blocks - self.local_window_blocks = local_window_blocks - self.global_block_indices = global_block_indices - - if (global_block_end_indices is not None): - if (len(global_block_indices) != len(global_block_end_indices)): - raise ValueError( - f'Global block start indices length, {len(global_block_indices)}, must be same as global block end indices length, {len(global_block_end_indices)}!' - ) - for _, (start_idx, end_idx) in enumerate(zip(global_block_indices, global_block_end_indices)): - if start_idx >= end_idx: - raise ValueError( - f'Global block start index, {start_idx}, must be smaller than global block end index, {end_idx}!' - ) - self.global_block_end_indices = global_block_end_indices - - if (attention != 'unidirectional' and attention != 'bidirectional'): - raise NotImplementedError('only \"uni/bi-directional\" attentions are supported for now!') - self.attention = attention - - if (attention != 'bidirectional' and horizontal_global_attention): - raise ValueError('only \"bi-directional\" attentions can support horizontal global attention!') - self.horizontal_global_attention = horizontal_global_attention - - def set_random_layout(self, h, layout): - """Sets random attention layout used by the given head in the sparse attention. - Note) By default, it assumes there will be a unique random block layout for all heads; unless `different_layout_per_head` parameter is set in which each head can have a different random layout. - - Arguments: - h: required: an integer determining head index - layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which random layout is set - """ - - num_blocks = layout.shape[1] - if (num_blocks < self.num_random_blocks): - raise ValueError( - f'Number of random blocks, {self.num_random_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' - ) - for row in range(0, num_blocks): - rnd_cols = random.sample(range(0, num_blocks), self.num_random_blocks) - layout[h, row, rnd_cols] = 1 - return layout - - def set_local_layout(self, h, layout): - """Sets local attention layout used by the given head in the sparse attention. - Arguments: - h: required: an integer determining head index - layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local layout is set - """ - - num_blocks = layout.shape[1] - start_block_idx = 0 - end_block_idx = 0 - for block_size in self.local_window_blocks: - end_block_idx += block_size - end_block_idx = min(end_block_idx, num_blocks) - for row in range(start_block_idx, end_block_idx): - for col in range(start_block_idx, (row + 1 if self.attention == 'unidirectional' else end_block_idx)): - layout[h, row, col] = 1 - start_block_idx += block_size - - # if there is any remaining not attended part, use the lats local window block size as local window for the remaining applicable local windows - for i in range(start_block_idx, num_blocks, block_size): - end_block_idx = min(i + block_size, num_blocks) - for row in range(i, end_block_idx): - for col in range(i, (row + 1 if self.attention == 'unidirectional' else end_block_idx)): - layout[h, row, col] = 1 - return layout - - def set_global_layout(self, h, layout): - """Sets global attention layout used by the given head in the sparse attention. - - Arguments: - h: required: an integer determining head index - layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which global layout is set - """ - - num_blocks = layout.shape[1] - if (self.global_block_end_indices is None): - for idx in self.global_block_indices: - # if global block idx is in the range of the sequence blocks - if (idx < num_blocks): - #global rows - if (self.horizontal_global_attention): - layout[h, idx, :] = 1 - - #global columns - first_row = 0 if self.attention == 'bidirectional' else idx - layout[h, first_row:, idx] = 1 - else: - for _, (start_idx, end_idx) in enumerate(zip(self.global_block_indices, self.global_block_end_indices)): - # if global block idx is in the range of the sequence blocks - if (start_idx < num_blocks): - end_idx = min(end_idx, num_blocks) - #global rows - if (self.horizontal_global_attention): - layout[h, start_idx:end_idx, :] = 1 - - #global columns - first_row = 0 if self.attention == 'bidirectional' else start_idx - layout[h, first_row:, start_idx:end_idx] = 1 - return layout - - def make_layout(self, seq_len): - """Generates `Variable` sparsity layout used by each head in the sparse attention. - - Arguments: - seq_len: required: an integer determining number of attention heads of the layer. - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `Variable` sparsity layout of all head - """ - - layout = self.setup_layout(seq_len) - for h in range(0, self.num_layout_heads): - layout = self.set_random_layout(h, layout) - layout = self.set_local_layout(h, layout) - layout = self.set_global_layout(h, layout) - - layout = self.check_and_propagate_first_head_layout(layout) - return layout - - -class BigBirdSparsityConfig(SparsityConfig): - """Configuration class to store `BigBird` sparsity configuration. - For more details about this sparsity config, please see `Big Bird: Transformers for Longer Sequences`: https://arxiv.org/pdf/2007.14062.pdf - This class extends parent class of `SparsityConfig` and customizes it for `BigBird` sparsity. - """ - - def __init__(self, - num_heads, - block=16, - different_layout_per_head=False, - num_random_blocks=1, - num_sliding_window_blocks=3, - num_global_blocks=1, - attention='bidirectional'): - """Initialize the BigBird Sparsity Pattern Config. - - For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial - - Arguments: - num_heads: required: an integer determining number of attention heads of the layer. - block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. - different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. - num_random_blocks: optional: an integer determining the number of random blocks in each block row. - num_sliding_window_blocks: optional: an integer determining the number of blocks in sliding local attention window. - num_global_blocks: optional: an integer determining how many consecutive blocks, starting from index 0, are considered as global attention. Global block tokens will be attended by all other block tokens and will attend to all other block tokens as well. - attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. - """ - - super().__init__(num_heads, block, different_layout_per_head) - - self.num_random_blocks = num_random_blocks - self.num_sliding_window_blocks = num_sliding_window_blocks - self.num_global_blocks = num_global_blocks - - if (attention != 'unidirectional' and attention != 'bidirectional'): - raise NotImplementedError('only \"uni/bi-directional\" attentions are supported for now!') - self.attention = attention - - def set_random_layout(self, h, layout): - """Sets random attention layout used by the given head in the sparse attention. - Note) By default, it assumes there will be a unique random block layout for all heads; unless `different_layout_per_head` parameter is set in which each head can have a different random layout. - - Arguments: - h: required: an integer determining head index - layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which random layout is set - """ - - num_blocks = layout.shape[1] - if (num_blocks < self.num_random_blocks): - raise ValueError( - f'Number of random blocks, {self.num_random_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' - ) - - for row in range(0, num_blocks): - sample_range = range(0, num_blocks) if self.attention == 'bidirectional' else range(0, row + 1) - rnd_cols = random.sample(sample_range, self.num_random_blocks) - layout[h, row, rnd_cols] = 1 - return layout - - def set_sliding_window_layout(self, h, layout): - """Sets sliding local attention layout used by the given head in the sparse attention. - - Arguments: - h: required: an integer determining head index - layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local sliding window layout is set - """ - - num_blocks = layout.shape[1] - if (num_blocks < self.num_sliding_window_blocks): - raise ValueError( - f'Number of sliding window blocks, {self.num_sliding_window_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' - ) - - w = self.num_sliding_window_blocks // 2 - for row in range(0, num_blocks): - start = max(0, row - w) - end = min(row + w + 1, num_blocks) - layout[h, row, start:end] = 1 - return layout - - def set_global_layout_itc(self, h, layout): - """Sets global attention layout used by the given head in the sparse attention. - - Arguments: - h: required: an integer determining head index - layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which global layout is set - """ - - num_blocks = layout.shape[1] - if (num_blocks < self.num_global_blocks): - raise ValueError( - f'Number of global blocks, {self.num_global_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' - ) - - #global rows - layout[h, 0:self.num_global_blocks, :] = 1 - - #global columns - layout[h, :, 0:self.num_global_blocks] = 1 - - if self.attention == 'unidirectional': - # zero out anything attending to the future - layout = torch.tril(layout) - - return layout - - def make_layout(self, seq_len): - """Generates `BigBird` sparsity layout used by each head in the sparse attention. - - Arguments: - seq_len: required: an integer determining number of attention heads of the layer. - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `BigBird` sparsity layout of all head - """ - - layout = self.setup_layout(seq_len) - for h in range(0, self.num_layout_heads): - layout = self.set_random_layout(h, layout) - layout = self.set_sliding_window_layout(h, layout) - layout = self.set_global_layout_itc(h, layout) - - layout = self.check_and_propagate_first_head_layout(layout) - return layout - - -class BSLongformerSparsityConfig(SparsityConfig): - """Configuration class to store edited `Longformer` sparsity configuration. - - Note) this is a block-sparse version of the Longformer which is slightly different than original Longformer; which is element-wise sparsity. - - For more details about this sparsity config, please see `Longformer: The Long-Document Transformer`: https://arxiv.org/pdf/2004.05150.pdf - This class extends parent class of `SparsityConfig` and customizes it for `Longformer` sparsity. - """ - - def __init__(self, - num_heads, - block=16, - different_layout_per_head=False, - num_sliding_window_blocks=3, - global_block_indices=[0], - global_block_end_indices=None, - attention='bidirectional'): - """Initialize the edited `Longformer` Sparsity Pattern Config. - - For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial - - Arguments: - num_heads: required: an integer determining number of attention heads of the layer. - block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. - different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. - - num_sliding_window_blocks: optional: an integer determining the number of blocks in sliding local attention window. - global_block_indices: optional: a list of integers determining which blocks are considered as global attention. Given indices, determine the blocks that all other token blocks attend to and they attend to all other token blocks. Default value is only index 0. Notice that if global_block_end_indices parameter is set, this parameter is used as starting index of each global window. - global_block_end_indices: optional: a list of integers determining end indices of global window blocks. By default this is not used. But if it is set, it must have the same size of global_block_indices parameter, and combining this two parameters, for each index i, blocks from global_block_indices[i] to global_block_end_indices[i] (exclusive) are considered as global attention. - attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. - """ - - super().__init__(num_heads, block, different_layout_per_head) - - self.num_sliding_window_blocks = num_sliding_window_blocks - self.global_block_indices = global_block_indices - self.attention = attention - - if (global_block_end_indices is not None): - if (len(global_block_indices) != len(global_block_end_indices)): - raise ValueError( - f'Global block start indices length, {len(global_block_indices)}, must be same as global block end indices length, {len(global_block_end_indices)}!' - ) - for _, (start_idx, end_idx) in enumerate(zip(global_block_indices, global_block_end_indices)): - if start_idx >= end_idx: - raise ValueError( - f'Global block start index, {start_idx}, must be smaller than global block end index, {end_idx}!' - ) - self.global_block_end_indices = global_block_end_indices - - def set_sliding_window_layout(self, h, layout): - """Sets sliding local attention layout used by the given head in the sparse attention. - - Arguments: - h: required: an integer determining head index - layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local sliding window layout is set - """ - - num_blocks = layout.shape[1] - if (num_blocks < self.num_sliding_window_blocks): - raise ValueError( - f'Number of sliding window blocks, {self.num_sliding_window_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' - ) - - w = self.num_sliding_window_blocks // 2 - for row in range(0, num_blocks): - start = max(0, row - w) - end = min(row + w + 1, num_blocks) - layout[h, row, start:end] = 1 - return layout - - def set_global_layout(self, h, layout): - """Sets global attention layout used by the given head in the sparse attention. - - Arguments: - h: required: an integer determining head index - layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which global layout is set - """ - - num_blocks = layout.shape[1] - if (self.global_block_end_indices is None): - for idx in self.global_block_indices: - # if global block idx is in the range of the sequence blocks - if (idx < num_blocks): - #global rows - layout[h, idx, :] = 1 - - #global columns - layout[h, :, idx] = 1 - else: - for _, (start_idx, end_idx) in enumerate(zip(self.global_block_indices, self.global_block_end_indices)): - # if global block idx is in the range of the sequence blocks - if (start_idx < num_blocks): - end_idx = min(end_idx, num_blocks) - #global rows - layout[h, start_idx:end_idx, :] = 1 - - #global columns - layout[h, :, start_idx:end_idx] = 1 - if self.attention == 'unidirectional': - layout = torch.tril(layout) - return layout - - def make_layout(self, seq_len): - """Generates edited `Longformer` sparsity layout used by each head in the sparse attention. - - Arguments: - seq_len: required: an integer determining number of attention heads of the layer. - - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `BSLongformer` sparsity layout of all head - """ - - layout = self.setup_layout(seq_len) - for h in range(0, self.num_layout_heads): - layout = self.set_sliding_window_layout(h, layout) - layout = self.set_global_layout(h, layout) - - layout = self.check_and_propagate_first_head_layout(layout) - return layout - - -class LocalSlidingWindowSparsityConfig(SparsityConfig): - """Configuration class to store `Local Sliding Window` sparsity configuration - a purely-local sliding window attention. - This class extends parent class of `SparsityConfig` and customizes it for `Local` sparsity. - """ - - def __init__(self, num_heads, block=16, num_sliding_window_blocks=3, attention='unidirectional'): - """Initialize the Local Sliding Window Sparsity Pattern Config. - For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial - Arguments: - num_heads: required: an integer determining number of attention heads of the layer. - block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. - num_sliding_window_blocks: optional: an integer determining the number of blocks in sliding local attention window. - attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. - """ - - super().__init__(num_heads, block) - self.num_sliding_window_blocks = num_sliding_window_blocks - self.attention = attention - - def set_sliding_window_layout(self, h, layout): - """Sets sliding local attention layout used by the given head in the sparse attention. - Arguments: - h: required: an integer determining head index - layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local sliding window layout is set - """ - - num_blocks = layout.shape[1] - if (num_blocks < self.num_sliding_window_blocks): - raise ValueError( - f'Number of sliding window blocks, {self.num_sliding_window_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' - ) - - w = self.num_sliding_window_blocks // 2 - for row in range(0, num_blocks): - start = max(0, row - w) - end = min(row + w + 1, num_blocks) if self.attention == "bidirectional" else row + 1 - layout[h, row, start:end] = 1 - return layout - - def make_layout(self, seq_len): - """Generates `Local Sliding Window` sparsity layout used by each head in the sparse attention. - Arguments: - seq_len: required: an integer determining number of attention heads of the layer. - Return: - layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `BigBird` sparsity layout of all head - """ - - layout = self.setup_layout(seq_len) - for h in range(0, self.num_layout_heads): - layout = self.set_sliding_window_layout(h, layout) - layout = self.check_and_propagate_first_head_layout(layout) - return layout diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/__init__.py deleted file mode 100644 index 18d58b5af4bf387ffc8b4e90645e3bd578ee10fe..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/__init__.py +++ /dev/null @@ -1,37 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import sys -import os - - -def _build_file_index(directory, suffix='.tr'): - """Build an index of source files and their basenames in a given directory. - - Args: - directory (string): the directory to index - suffix (string): index files with this suffix - - Returns: - list: A list of tuples of the form [(basename, absolute path), ...] - """ - - index = [] - - for fname in os.listdir(directory): - if fname.endswith(suffix): - basename = fname[:fname.rfind(suffix)] # strip the suffix - path = os.path.join(directory, fname) - index.append((basename, path)) - - return index - - -# Go over all local source files and parse them as strings -_module = sys.modules[_build_file_index.__module__] -_directory = os.path.dirname(os.path.realpath(__file__)) -for name, fname in _build_file_index(_directory): - with open(fname, 'r') as fin: - setattr(_module, name, fin.read()) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index c248736d3fddb49f3f34113bf5d9e71f5ad8e61e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/matmul.tr b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/matmul.tr deleted file mode 100644 index 9be6aa2aed7ccfb0d438d2e4d766d30b21d8f744..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/matmul.tr +++ /dev/null @@ -1,208 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a - https:github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/matmul.py -*/ - -__global__ void NAME (TYPE* A __readonly __noalias __aligned(16), - TYPE* B __readonly __noalias __aligned(16), - TYPE* C __noalias __aligned(16), - int lda __multipleof(8), - int ldb __multipleof(8), - int ldc __multipleof(8), - long stride_za __multipleof(8), - long stride_zb __multipleof(8), - long stride_zc __multipleof(8), - long stride_ha __multipleof(8), - long stride_hb __multipleof(8), - long stride_hc __multipleof(8), - int DS0, int DS1, - int SDD_K __multipleof(16), - int SDD_off_width, - int* lut, int* locks, int nlocks) { - /* ---------------- */ - /* Prologue */ - /* ---------------- */ - // program ids - int pid0 = get_program_id(0); - int pid1 = get_program_id(1); - int pidz = get_program_id(2); -#ifdef SDD - // load LUT header - pid1 = pid1 + SDD_off_width; - int blockidm[TM] = (0 ... TM) / BLOCK; - int blockidn[TN] = (0 ... TN) / BLOCK; - int offlutm[TM] = blockidm*(TN/BLOCK)*4; - int offlutn[TN] = blockidn*4; - int *header = lut + pid1 * (TM/BLOCK) * (TN/BLOCK) * 4; - int z = *(header + 0); - int i[TM] = *(header + 1 + offlutm); - int j[TN] = *(header + 2 + offlutn); - int AS1 = SDD_K / TZ; - int lockid = select(TZ > 1, 1, 0); - int offka = pid0 * AS1; - int offkb = pid0 * AS1; - int offmc = 0; - int offnc = 0; - int offpa = 0; - int offpb = 0; - int maxid = TZ; - int offhc = 0; - int offha = z; - int offhb = z; - int ram[TM] = i*BLOCK + ((0 ... TM) % BLOCK); - int rbn[TN] = j*BLOCK + ((0 ... TN) % BLOCK); -#else - // load LUT header - int *header = lut + pid0 * 6; - int offset = *(header + 0); - int AS1 = *(header + 1); - int column = *(header + 2); - int depth = *(header + 3); - int lockid = *(header + 4); - int maxid = *(header + 5); - int *pinc = lut + offset; - int offhc = depth; -#ifdef DSD - // output offset - int offnc = pid1 * TN; - int offmc = column * TM; - int offpc = 0; - // dense input offset - int offnb = pid1 * TN; - int offkb __multipleof(8) = *pinc; - int offpb = 0; - // sparse input offset - int offma = 0; - int offka = 0; - long offpa __multipleof(8) = *(pinc + 1); - offpa = offpa * BLOCK * BLOCK; - int offha = 0; - int offhb = depth; -#endif -#ifdef DDS - // output offset - int offmc = pid1 * TM; - int offnc = column * TN; - int offpc = 0; - // dense input offset - int offma = pid1 * TM; - int offka __multipleof(8) = *pinc; - int offpa = 0; - // sparse input offset - int offnb = 0; - int offkb = 0; - long offpb __multipleof(8) = *(pinc + 1); - offpb = offpb * BLOCK * BLOCK; - int offha = depth; - int offhb = 0; -#endif - int ram[TM] = offma + 0 ... TM; - int rbn[TN] = offnb + 0 ... TN; -#endif - // initialize a, b pointers - int rka[TK] = offka + 0 ... TK; - int rkb[TK] = offkb + 0 ... TK; - TYPE* pa[TM, TK] = A + pidz * stride_za + offha * stride_ha + offpa + ram[:, newaxis] * STRIDE_AM + rka[newaxis, :] * STRIDE_AK; - TYPE* pb[TK, TN] = B + pidz * stride_zb + offhb * stride_hb + offpb + rbn[newaxis, :] * STRIDE_BN + rkb[:, newaxis] * STRIDE_BK; - // pre-fetch -#ifdef DDS - bool checkam[TM, TK] = ram[:, newaxis] < DS0; -#else - bool checkam[TM, TK] = AS1 > 0; -#endif -#ifdef DSD - bool checkbn[TK, TN] = rbn[newaxis, :] < DS0; -#else - bool checkbn[TK, TN] = AS1 > 0; -#endif - TYPE a[TM, TK] = checkam ? *pa : 0; - TYPE b[TK, TN] = checkbn ? *pb : 0; - - /* ---------------- */ - /* Inner Loop */ - /* ---------------- */ - // create result tile - float acc[TM, TN] = 0; - int step = TK; - for(int k = AS1; k > 0; k -= step) { - acc += a @ b; - // update pointers -#ifdef SDD - int inc_a = TK * STRIDE_AK; - int inc_b = TK * STRIDE_BK; -#else - pinc += 2; -#ifdef DSD - int inc_b __multipleof(8) = *pinc; - int inc_a __multipleof(8) = *(pinc + 1); - inc_b = inc_b * STRIDE_BK; -#endif -#ifdef DDS - int inc_a __multipleof(8) = *pinc; - int inc_b __multipleof(8) = *(pinc + 1); - inc_a = inc_a * STRIDE_AK; -#endif -#endif - pa += inc_a; - pb += inc_b; - // pre-fetch - bool checkak[TM, TK] = k > TK; - bool checkbk[TK, TN] = k > TK; - bool checka[TM, TK] = checkam && checkak; - bool checkb[TK, TN] = checkbk && checkbn; - a = *?(checka)pa; - b = *?(checkb)pb; - } - TYPE c[TM, TN] = acc; - - /* ---------------- */ - /* Epilogue */ - /* ---------------- */ - // initialize c pointers -#ifdef SDD - bool checkc[TM, TN] = 1; - // rematerialize - int rr_blockidm[TM] = (0 ... TM) / BLOCK; - int rr_blockidn[TN] = (0 ... TN) / BLOCK; - int rr_offlutm[TM] = rr_blockidm*(TN/BLOCK)*4; - int rr_offlutn[TN] = rr_blockidn*4; - int off_bkid[TM, TN] = 3 + rr_offlutm[:, newaxis] + rr_offlutn[newaxis, :]; - int bkid[TM, TN] = *(header + off_bkid); - long offpc[TM, TN] = bkid * BLOCK * BLOCK; - // range within blocks - int rcm[TM] = (0 ... TM) % BLOCK; - int rcn[TN] = (0 ... TN) % BLOCK; -#else - int rcm[TM] = offmc + 0 ... TM; - int rcn[TN] = offnc + 0 ... TN; -#ifdef DSD - bool checkc[TM, TN] = rcn[newaxis, :] < DS0; -#endif -#ifdef DDS - bool checkc[TM, TN] = rcm[:, newaxis] < DS0; -#endif -#endif - TYPE* pc[TM, TN] = C + offpc + offhc*stride_hc + pidz*stride_zc + rcm[:, newaxis]*STRIDE_CM + rcn[newaxis, :]*STRIDE_CN; - // write-back directly - if(lockid == 0) { - *?(checkc) pc = c; - } - // accumulate partial result using spin-locks - else { - int *plock = locks + get_program_id(2)*nlocks*get_num_programs(1) + get_program_id(1)*nlocks + lockid - 1; - int *pcount = plock + get_num_programs(2)*get_num_programs(1)*nlocks; - for(int repeat = 1; repeat == 1; repeat = atomic_cas(plock, 0, 1)); - int count = *pcount; - if(count == 0) - *?(checkc) pc = c; - else - *?(checkc) pc = c + *?(checkc)pc; - atomic_xchg(pcount, (count + 1) % maxid); - atomic_xchg(plock, 0); - } - } diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_bwd.tr b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_bwd.tr deleted file mode 100644 index 1a90f41d94945e1d6d6f52e6beaea94fa52cdda8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_bwd.tr +++ /dev/null @@ -1,61 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a - https:github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/softmax.py -*/ - -__global__ void softmax_bwd(TYPE * X __readonly __noalias __aligned(16), - float scale, - TYPE* DX __readonly __noalias __aligned(16), - int* LUT, - int sizemax, - long stride_zx __multipleof(BLOCK), - long stride_zdx __multipleof(BLOCK)) { - int pidhm = get_program_id(0); - int pidz = get_program_id(1); - - // create index ranges - int rxm = pidhm % BLOCK; - int rbm = pidhm / BLOCK; - int rxn[TN] = (0 ... TN) % BLOCK; - int rbn[TN] = (0 ... TN) / BLOCK; - - // extract information from look-up table - int* header = LUT + rbm * 2; - int size = *(header + 0); - int offset = *(header + 1); - - // bounds checking on lut - bool check[TN] = rbn < size; - int rbmn[TN] = check ? rbn : size - 1; - - // initialize pointers to block-sparse input - long blockid[TN] = *(LUT + offset + rbmn*4); - - TYPE* px[TN] = X + pidz * stride_zx - + blockid * BLOCK * BLOCK - + rxm * BLOCK - + rxn; - - TYPE* pdx[TN] = DX + pidz * stride_zdx - + blockid * BLOCK * BLOCK - + rxm * BLOCK - + rxn; - - // compute fused softmax backward - TYPE x[TN] = check ? *px : 0; - TYPE dx[TN] = check ? *pdx : 0; - float Fdx[TN] = dx; - float Fx[TN] = x; - float Fxdx[TN] = Fdx*Fx; - float Fxdxsum = Fxdx[+]; - float Fy[TN] = Fx * (Fdx - Fxdxsum) * scale; - TYPE y[TN] = Fy; - - // write-back - *? (check)pdx = y; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_fwd.tr b/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_fwd.tr deleted file mode 100644 index ebd317d9469b47f7e2ee3032d3aabf57b5620a73..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_fwd.tr +++ /dev/null @@ -1,143 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// SPDX-License-Identifier: Apache-2.0 - -// DeepSpeed Team - -/* -DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a - https:github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/softmax.py -*/ - -__global__ void softmax_fwd(TYPE *X __readonly __noalias __aligned(16), - float scale, - int *LUT __readonly __noalias __aligned(16), - TYPE *RPE __readonly __noalias __aligned(16), - TYPE *KP_M __readonly __noalias __aligned(16), - TYPE *ATTN_M __readonly __noalias __aligned(16), - int num_blocks, - int sizemax, - long stride_zx __multipleof(BLOCK), - long stride_zrpe __multipleof(BLOCK), - int stride_hrpe __multipleof(BLOCK), - int stride_srpe __multipleof(BLOCK), - int stride_zkpm __multipleof(BLOCK), - int stride_zattnm __multipleof(BLOCK)){ - int pidhm = get_program_id(0); - int pidz = get_program_id(1); - - // create index ranges - int rxm = pidhm % BLOCK; - int rbm = pidhm / BLOCK; - int rxn[TN] = (0 ... TN) % BLOCK; - int rbn[TN] = (0 ... TN) / BLOCK; - - // extract information from look-up table - int* header = LUT + rbm * 2; - int size = *(header + 0); - int offset = *(header + 1); - - bool check[TN] = rbn < size; - int rbmn[TN] = check ? rbn : size - 1; - - // block id and column id - long blockid [TN] = *(LUT + offset + rbmn*4 + 0); - long columnid[TN] = *(LUT + offset + rbmn*4 + 1); - long rowid [TN] = *(LUT + offset + rbmn*4 + 2); - long headid [TN] = *(LUT + offset + rbmn*4 + 3); - - // pointers to X - TYPE* px[TN] = X + pidz * stride_zx - + blockid * BLOCK * BLOCK - + rxm * BLOCK - + rxn; -#ifdef APPLY_RPE - // pointers to relative position embedding - TYPE* prpe[TN] = RPE + pidz * stride_zrpe - + headid * stride_hrpe - + columnid * BLOCK - + rowid * BLOCK * stride_srpe - + rxm * stride_srpe - + rxn; -#endif - -#ifdef APPLY_KP_MASK - // pointers to key padding mask - TYPE* pkp_m[TN] = KP_M + pidz * stride_zkpm - + columnid * BLOCK - + rxn; -#endif - -#ifdef APPLY_ATTN_MASK - // pointers to attention mask - TYPE* pattn_m[TN] = ATTN_M + columnid * BLOCK - + rowid * BLOCK * stride_zattnm - + rxm * stride_zattnm - + rxn; -#endif - - // load input - TYPE x[TN] = check ? *px : -INFINITY; - -#ifdef APPLY_RPE - // load relative position embedding - TYPE rpe[TN] = check ? *prpe : 0; -#endif - -#ifdef APPLY_KP_MASK - // load key-padding mask - TYPE kp_m[TN] = check ? *pkp_m : -INFINITY; -#endif - -#ifdef APPLY_ATTN_MASK - // load attention mask - TYPE attn_m[TN] = check ? *pattn_m : -INFINITY; -#endif - - // compute softmax in float -#ifdef APPLY_RPE - float Frpe[TN] = rpe; -#endif - -#ifdef APPLY_KP_MASK - float Fkp_m[TN] = kp_m; -#endif - -#ifdef APPLY_ATTN_MASK - float Fattn_m[TN] = attn_m; -#endif - -#ifdef KP_MASK_MUL - Fkp_m = (Fkp_m == 0) ? (float[TN])-INFINITY : 0; -#endif - -#ifdef ATTN_MASK_MUL - Fattn_m = (Fattn_m == 0) ? (float[TN])-INFINITY : 0; -#endif - - float Fx[TN] = x; - -#ifdef APPLY_SCALE - Fx = Fx * scale; // apply scale -#endif - -#ifdef APPLY_RPE - Fx = Fx + Frpe; // apply relative position embedding -#endif - -#ifdef APPLY_KP_MASK - Fx = Fx + Fkp_m; // apply key padding mask -#endif - -#ifdef APPLY_ATTN_MASK - Fx = Fx + Fattn_m; // apply attention mask -#endif - - float Fxmax = Fx[max]; - float Fy[TN] = exp(Fx - Fxmax); - float Fysum = (check ? Fy : 0)[+]; - - // write-back in half/float - TYPE y[TN] = Fy; - TYPE ysum = Fysum; - *?(check)px = y / ysum; -} diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__init__.py deleted file mode 100644 index b80fe2b4ba714611b6bcb652d9e559d87c7ed6fd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .transformer import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig -from .inference.config import DeepSpeedInferenceConfig -from ...model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference -from .inference.moe_inference import DeepSpeedMoEInferenceConfig, DeepSpeedMoEInference diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index d40e38a15fccf16639bf041eda206def43db479c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__pycache__/transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__pycache__/transformer.cpython-310.pyc deleted file mode 100644 index 25953a0a79bca2f45e8e8e33112ad68855089524..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/__pycache__/transformer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__init__.py deleted file mode 100644 index c8b31a90eac2850fa7554214760c00df9815db25..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .config import DeepSpeedInferenceConfig -from ....model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference -from .moe_inference import DeepSpeedMoEInferenceConfig, DeepSpeedMoEInference diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 1fccdc794e975c1715426602a0131da16329bd6a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/bias_add.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/bias_add.cpython-310.pyc deleted file mode 100644 index ac10b402be0d1e377c8e1d04558b3f2c7c1a2638..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/bias_add.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/config.cpython-310.pyc deleted file mode 100644 index 46c012eda8fd71bb8faa1c7720934e8d0b9acffa..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_2d_transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_2d_transformer.cpython-310.pyc deleted file mode 100644 index 6cca2957c9857a3f845085e0b6331c051e295afd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_2d_transformer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_attention.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_attention.cpython-310.pyc deleted file mode 100644 index 8c65d09e74f5dc34c7c5f6805e4d02feebd18d63..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_attention.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_transformer_block.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_transformer_block.cpython-310.pyc deleted file mode 100644 index aa024f67be7410ad10443fa4ca13efbfb52da562..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/diffusers_transformer_block.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/ds_attention.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/ds_attention.cpython-310.pyc deleted file mode 100644 index 9a3c9c97ef21b6fc25dc923ba32760493b2e2529..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/ds_attention.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/ds_mlp.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/ds_mlp.cpython-310.pyc deleted file mode 100644 index c59b1dbc94c8fc8e3cc70248c90628563fd717c0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/ds_mlp.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/moe_inference.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/moe_inference.cpython-310.pyc deleted file mode 100644 index 227a21d78223723e650d4719beb6c00e169b5737..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/moe_inference.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/triton_ops.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/triton_ops.cpython-310.pyc deleted file mode 100644 index 073603c404114efb97c3a8232e511a7946eee205..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/__pycache__/triton_ops.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/bias_add.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/bias_add.py deleted file mode 100644 index 253784f001aeb3431a0d60812e4d2068c0bd5455..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/bias_add.py +++ /dev/null @@ -1,26 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Optional -import torch -from deepspeed.ops.op_builder import SpatialInferenceBuilder - -spatial_cuda_module = None - - -def nhwc_bias_add(activation: torch.Tensor, - bias: torch.Tensor, - other: Optional[torch.Tensor] = None, - other_bias: Optional[torch.Tensor] = None) -> torch.Tensor: - global spatial_cuda_module - if spatial_cuda_module is None: - spatial_cuda_module = SpatialInferenceBuilder().load() - - if other is None: - return spatial_cuda_module.nhwc_bias_add(activation, bias) - elif other_bias is None: - return spatial_cuda_module.nhwc_bias_add_add(activation, bias, other) - else: - return spatial_cuda_module.nhwc_bias_add_bias_add(activation, bias, other, other_bias) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/config.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/config.py deleted file mode 100644 index c0dd29f4f9627e79a294089323d8bb83728e13db..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/config.py +++ /dev/null @@ -1,133 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import json -import torch -from deepspeed.utils.types import ActivationFuncType, NormType - - -class TransformerConfig(): - - def __init__(self, hidden_size, intermediate_size, heads, num_hidden_layers): - self.layer_id = -1 - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.heads = heads - self.num_hidden_layers = num_hidden_layers - - -class DeepSpeedInferenceConfig(TransformerConfig): - """Initialize the DeepSpeed Transformer Config. - Arguments: - hidden_size: The hidden size of the transformer layer - intermediate_size: The intermediate size of the feed-forward part of transformer layer - heads: The number of heads in the self-attention of the transformer layer - num_hidden_layers: The number of transformer layers - layer_norm_eps: The epsilon value for the layer norm - local_rank: Optional: The rank of GPU running the transformer kernel, it is not required - to use if the model already set the current device, otherwise need to set it - so that the transformer kernel can work on the right device - mp_size (optional): This argument is mainly used to create the parameters on the kernel side - using model-parallel architecture. If the client model already takes care of this, there is no - need to pass this argument. - pre_layer_norm: Select between Pre-LN or Post-LN transformer architecture - stochastic_mode: Enable for high performance, please note that this flag has some level of - non-determinism and can produce different results on different runs. However, we have seen - that by enabling it, the pretraining tasks such as BERT are not affected and can obtain - a high accuracy level. On the other hand, for the downstream tasks, such as fine-tuning, we recommend - to turn it off in order to be able to reproduce the same result through the regular kernel execution. - - scale_attention: If true, both q and k are scaled by 1/sqrt(attention_heads) before attention computation. - return_tuple: if True, returns the transformer output as a tuple, otherwise returns as a tensor - bigscience_bloom: This flag is added temporarily for supporting the BLOOM-176B model architecture. - use_triton: This flag is to enable triton kernels in inference or not. - invert_mask: If True, the attention mask is inverted when passed to attention block. - """ - - def __init__(self, - hidden_size=-1, - intermediate_size=-1, - heads=-1, - num_hidden_layers=-1, - layer_norm_eps=1e-12, - local_rank=-1, - mp_size=1, - dtype=torch.float16, - pre_layer_norm=True, - norm_type=NormType.LayerNorm, - stochastic_mode=False, - scale_attention=True, - triangular_masking=True, - local_attention=False, - window_size=256, - rotary_dim=-1, - rotate_half=False, - rotate_every_two=True, - return_tuple=True, - mlp_after_attn=True, - mlp_act_func_type=ActivationFuncType.GELU, - training_mp_size=1, - bigscience_bloom=False, - max_out_tokens=1024, - min_out_tokens=1, - enable_qkv_quantization=False, - use_mup=False, - scale_attn_by_inverse_layer_idx=False, - return_single_tuple=False, - set_empty_params=False, - transposed_mode=False, - use_triton=False, - triton_autotune=False, - num_kv=-1, - rope_theta=10000, - invert_mask=True): - super(DeepSpeedInferenceConfig, - self).__init__(hidden_size, (intermediate_size if intermediate_size > 0 else 4 * hidden_size), heads, - num_hidden_layers) - self.dtype = dtype - self.pre_layer_norm = pre_layer_norm - self.norm_type = norm_type - self.local_rank = local_rank - self.stochastic_mode = stochastic_mode - self.epsilon = layer_norm_eps - self.mp_size = mp_size - self.scale_attention = scale_attention - self.triangular_masking = triangular_masking - self.local_attention = local_attention - self.window_size = window_size - self.rotary_dim = rotary_dim - self.rotate_half = rotate_half - self.rotate_every_two = rotate_every_two - self.return_tuple = return_tuple - self.mlp_after_attn = mlp_after_attn - self.mlp_act_func_type = mlp_act_func_type - self.training_mp_size = training_mp_size - self.bigscience_bloom = bigscience_bloom - self.max_out_tokens = max_out_tokens - self.min_out_tokens = min_out_tokens - self.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx - self.enable_qkv_quantization = enable_qkv_quantization - self.use_mup = use_mup - self.return_single_tuple = return_single_tuple - self.set_empty_params = set_empty_params - self.transposed_mode = transposed_mode - self.use_triton = use_triton - self.triton_autotune = triton_autotune - self.num_kv = num_kv - self.rope_theta = rope_theta - self.invert_mask = invert_mask - - @classmethod - def from_dict(cls, json_object): - config = DeepSpeedInferenceConfig() - for key, value in json_object.items(): - config.__dict__[key] = value - return config - - @classmethod - def from_json_file(cls, json_file): - with open(json_file, "r", encoding='utf-8') as reader: - text = reader.read() - return cls.from_dict(json.loads(text)) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_2d_transformer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_2d_transformer.py deleted file mode 100644 index fa4c6d53f871d326d0a51f2bc9ae7c71743befda..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_2d_transformer.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - - -class Diffusers2DTransformerConfig(): - - def __init__(self, int8_quantization=False): - self.int8_quantization = int8_quantization diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_attention.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_attention.py deleted file mode 100644 index 3c2340ccfc6f9f2bc66fd1f1fa10d4a7c2b2aaf9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_attention.py +++ /dev/null @@ -1,192 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import math -import torch -from torch.autograd import Function -import torch.nn as nn -from packaging import version as pkg_version -from deepspeed.utils.logging import log_dist -from deepspeed.accelerator import get_accelerator -from deepspeed.ops.transformer.inference.op_binding.workspace import WorkspaceOp -from deepspeed.ops.transformer.inference.op_binding.softmax_context import SoftmaxContextOp -from deepspeed.ops.transformer.inference.op_binding import LinearOp -from deepspeed.ops.transformer.inference.op_binding.pad_transform import PadTransformOp - -minus_inf = -10000.0 -triton_flash_attn = None - - -def load_triton_flash_attn(): - global triton_flash_attn - try: - import triton - except ImportError: - raise ImportError("Please install triton 2.0+ or `pip install deepspeed[sd]`") - - if pkg_version.parse(triton.__version__) < pkg_version.parse("2.0"): - raise ImportError("Please install triton 2.0+ or `pip install deepspeed[sd]`") - - from .triton_ops import triton_flash_attn - - -class DeepSpeedDiffusersAttentionFunction(Function): - - @staticmethod - def forward(ctx, input, context, input_mask, config, attn_qkvw, attn_qw, attn_kw, attn_vw, attn_qkvb, - num_attention_heads_per_partition, norm_factor, hidden_size_per_partition, attn_ow, attn_ob, - do_out_bias, score_context_func, linear_func, pad_transform_func, triton_flash_attn_kernel, - rope_theta): - - def _transpose_for_context(x): - x = x.permute(0, 2, 1, 3) - new_x_layer_shape = x.size()[:-2] + \ - (hidden_size_per_partition,) - return x.reshape(*new_x_layer_shape) - - def _transpose_for_scores(x): - attention_head_size = x.shape[-1] // num_attention_heads_per_partition - new_x_shape = x.size()[:-1] + (num_attention_heads_per_partition, attention_head_size) - x = x.reshape(*new_x_shape) - x = x.permute(0, 2, 1, 3) - return x.contiguous() - - def selfAttention_fp(input, context, input_mask): - if config.dtype in [torch.half, torch.float16] and input.dtype == torch.float32: - input = input.half() - head_size = input.shape[-1] // config.heads - do_flash_attn = (head_size <= 128) - scale = (1 / norm_factor) * (1 / norm_factor) - if do_flash_attn and context is None: - qkv_out = linear_func(input, attn_qkvw, attn_qkvb if attn_qkvb is not None else attn_qkvw, attn_qkvb - is not None, do_flash_attn, config.heads, False, rope_theta) - - context_layer = triton_flash_attn_kernel(qkv_out[0], qkv_out[1], qkv_out[2], scale, - input.shape[-2] % 128 == 0) - context_layer = _transpose_for_context(context_layer[:, :, :, :head_size]) - - else: - do_flash_attn = False - if context is not None: - query = torch.matmul(input, attn_qw) - key = torch.matmul(context, attn_kw) - value = torch.matmul(context, attn_vw) - else: - qkv = torch.matmul(input, attn_qkvw) - query, key, value = qkv.chunk(3, dim=-1) - query = query.contiguous() - key = key.contiguous() - value = value.contiguous() - query, key, value = pad_transform_func(query, key, value, config.heads, do_flash_attn) - attention_scores = (torch.matmul(query, key.transpose(-1, -2)) * scale).softmax(dim=-1) - context_layer = _transpose_for_context(torch.matmul(attention_scores, value)) - - output = linear_func(context_layer, attn_ow, attn_ob, do_out_bias, False, config.heads, False, rope_theta) - return output - - output = selfAttention_fp(input, context, input_mask) - - return output - - @staticmethod - def backward(ctx, grad_output, grad_output1, grad_output2, grad_output3): - raise RuntimeError('You are running with DeepSpeed Inference mode. \ - Please switch to Training mode for running backward!') - - -class DeepSpeedDiffusersAttention(nn.Module): - """Initialize the DeepSpeed Transformer Layer. - Arguments: - layer_id: The layer index starting from 0, e.g. if model has 24 transformer layers, - layer_id will be 0,1,2...23 when each layer object is instantiated - config: An object of DeepSpeedInferenceConfig - """ - layer_id = 0 - - def __init__( - self, - config, - ): - super(DeepSpeedDiffusersAttention, self).__init__() - - self.config = config - self.config.layer_id = DeepSpeedDiffusersAttention.layer_id - DeepSpeedDiffusersAttention.layer_id += 1 - device = get_accelerator().current_device_name() if config.bigscience_bloom else 'cpu' - qkv_size_per_partition = (self.config.hidden_size // self.config.mp_size) * 3 - - data_type = self.config.dtype - data_type_fp = torch.half if self.config.dtype == torch.int8 else self.config.dtype - - if DeepSpeedDiffusersAttention.layer_id == 1: - log_dist(f"DeepSpeed-Attention config: {self.config.__dict__}", [0]) - - self.attn_qkvw = nn.Parameter(torch.empty(self.config.hidden_size, - qkv_size_per_partition, - dtype=data_type, - device=device), - requires_grad=False) - self.attn_kw = nn.Parameter(torch.empty(self.config.hidden_size, - self.config.hidden_size, - dtype=data_type, - device=device), - requires_grad=False) - self.attn_vw = nn.Parameter(torch.empty(self.config.hidden_size, - self.config.hidden_size, - dtype=data_type, - device=device), - requires_grad=False) - self.attn_qw = nn.Parameter(torch.empty(self.config.hidden_size, - self.config.hidden_size, - dtype=data_type, - device=device), - requires_grad=False) - self.attn_qkvb = nn.Parameter(torch.empty(qkv_size_per_partition, dtype=data_type_fp, device=device), - requires_grad=False) - out_size_per_partition = self.config.hidden_size // self.config.mp_size - self.attn_ow = nn.Parameter(torch.empty(out_size_per_partition, - self.config.hidden_size, - dtype=data_type, - device=device), - requires_grad=False) - - self.attn_ob = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), - requires_grad=False) - self.do_out_bias = True - - if triton_flash_attn is None: - load_triton_flash_attn() - self.triton_flash_attn_kernel = triton_flash_attn() - self.num_attention_heads_per_partition = self.config.heads // self.config.mp_size - self.hidden_size_per_partition = self.config.hidden_size // self.config.mp_size - self.hidden_size_per_attention_head = self.config.hidden_size // self.config.heads - - self.norm_factor = math.sqrt(math.sqrt(self.config.hidden_size // self.config.heads)) - - if self.config.scale_attn_by_inverse_layer_idx is True: - self.norm_factor *= math.sqrt(self.config.layer_id + 1) - # https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/gpt2/modeling_gpt2.py#L191 - - self.workspace = WorkspaceOp(self.config) - self.score_context_func = SoftmaxContextOp(self.config) - self.linear_func = LinearOp(self.config) - self.pad_transform_func = PadTransformOp(self.config) - - def allocate_workspace(self, size): - # Allocate memory only on first layer forward - if self.config.layer_id == 0: - self.workspace.allocate_workspace(self.config.hidden_size, self.config.heads, size[1], size[0], - DeepSpeedDiffusersAttention.layer_id, self.config.mp_size, False, 0, - self.config.max_out_tokens, self.config.min_out_tokens) - - def forward(self, input, context=None, input_mask=None): - self.allocate_workspace(input.size()) - output = DeepSpeedDiffusersAttentionFunction.apply( - input, context, input_mask, self.config, self.attn_qkvw, self.attn_qw, self.attn_kw, self.attn_vw, - self.attn_qkvb, self.num_attention_heads_per_partition, self.norm_factor, self.hidden_size_per_partition, - self.attn_ow, self.attn_ob, self.do_out_bias, self.score_context_func, self.linear_func, - self.pad_transform_func, self.triton_flash_attn_kernel, self.config.rope_theta) - - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_transformer_block.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_transformer_block.py deleted file mode 100644 index d01638f36e40141dc91276dcb484aeb96865048a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/diffusers_transformer_block.py +++ /dev/null @@ -1,87 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import torch.nn as nn - -from deepspeed import module_inject -from .diffusers_attention import DeepSpeedDiffusersAttention -from .bias_add import nhwc_bias_add -from .diffusers_2d_transformer import Diffusers2DTransformerConfig -from deepspeed.utils.types import ActivationFuncType -from .op_binding.gated_activation import GatedActivationOp -from .op_binding.layer_norm import LayerNormOp - - -class DeepSpeedDiffusersTransformerBlock(nn.Module): - - def __init__(self, equivalent_module: nn.Module, config: Diffusers2DTransformerConfig): - super(DeepSpeedDiffusersTransformerBlock, self).__init__() - self.quantizer = module_inject.GroupQuantizer(q_int8=config.int8_quantization) - # Ensure ops are built by the time we start running - self.config = config - - self.ff1_w = self.quantizer.quantize( - nn.Parameter(equivalent_module.ff.net[0].proj.weight.data, requires_grad=False)) - self.ff1_b = nn.Parameter(equivalent_module.ff.net[0].proj.bias.data, requires_grad=False) - self.ff2_w = self.quantizer.quantize(nn.Parameter(equivalent_module.ff.net[2].weight.data, - requires_grad=False)) - self.ff2_b = nn.Parameter(equivalent_module.ff.net[2].bias.data, requires_grad=False) - - self.norm1_g = nn.Parameter(equivalent_module.norm1.weight.data, requires_grad=False) - self.norm1_b = nn.Parameter(equivalent_module.norm1.bias.data, requires_grad=False) - self.norm1_eps = equivalent_module.norm1.eps - - self.norm2_g = nn.Parameter(equivalent_module.norm2.weight.data, requires_grad=False) - self.norm2_b = nn.Parameter(equivalent_module.norm2.bias.data, requires_grad=False) - self.norm2_eps = equivalent_module.norm2.eps - - self.norm3_g = nn.Parameter(equivalent_module.norm3.weight.data, requires_grad=False) - self.norm3_b = nn.Parameter(equivalent_module.norm3.bias.data, requires_grad=False) - self.norm3_eps = equivalent_module.norm3.eps - - self.attn_1 = equivalent_module.attn1 - self.attn_2 = equivalent_module.attn2 - - # Pull the bias in if we can - if isinstance(self.attn_1, DeepSpeedDiffusersAttention): - self.attn_1.do_out_bias = False - self.attn_1_bias = self.attn_1.attn_ob - else: - self.attn_1_bias = nn.Parameter(torch.zeros_like(self.norm2_g), requires_grad=False) - - # Pull the bias in if we can - if isinstance(self.attn_2, DeepSpeedDiffusersAttention): - self.attn_2.do_out_bias = False - self.attn_2_bias = self.attn_2.attn_ob - else: - self.attn_2_bias = nn.Paramaeter(torch.zeros_like(self.norm3_g), requires_grad=False) - - self.gated_activation = GatedActivationOp() - self.layer_norm = LayerNormOp() - - def forward(self, hidden_states, context=None, timestep=None, **kwargs): - # In v0.12.0 of diffuser, several new kwargs were added. Capturing - # those with kwargs to maintain backward compatibility - - # In v0.11.0 of diffusers, the kwarg was changed from 'context' to 'encoder_hidden_states' - # This is so we can support older and newer versions of diffusers - if "encoder_hidden_states" in kwargs and kwargs["encoder_hidden_states"] is not None: - context = kwargs["encoder_hidden_states"] - - out_norm_1 = self.layer_norm(hidden_states, self.norm1_g, self.norm1_b, self.norm1_eps) - out_attn_1 = self.attn_1(out_norm_1) - - out_norm_2, out_attn_1 = self.layer_norm.layer_norm_residual_store_pre_ln_res( - out_attn_1, self.attn_1_bias, hidden_states, self.norm2_g, self.norm2_b, self.norm2_eps) - out_attn_2 = self.attn_2(out_norm_2, context=context) - out_norm_3, out_attn_2 = self.layer_norm.layer_norm_residual_store_pre_ln_res( - out_attn_2, self.attn_2_bias, out_attn_1, self.norm3_g, self.norm3_b, self.norm3_eps) - - out_ff1 = nn.functional.linear(out_norm_3, self.ff1_w) - out_geglu = self.gated_activation(out_ff1, self.ff1_b, ActivationFuncType.GATED_GELU) - - out_ff2 = nn.functional.linear(out_geglu, self.ff2_w) - return nhwc_bias_add(out_ff2, self.ff2_b, other=out_attn_2) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/ds_attention.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/ds_attention.py deleted file mode 100644 index 24f710d224949a5b7684742b7ff7cf96eb007122..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/ds_attention.py +++ /dev/null @@ -1,301 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import math -import torch -import torch.nn as nn -from deepspeed import comm as dist -from deepspeed.accelerator import get_accelerator -from .op_binding import LinearOp, VectorMatMulOp, SoftmaxContextOp, QKVGemmOp, SoftmaxOp - -minus_inf = -10000.0 - - -class DeepSpeedSelfAttention(nn.Module): - num_layers = 0 - _qkv_buffers = [] - - def __init__(self, config, mp_group=None, q_scales=None, q_groups=1, merge_count=1): - super(DeepSpeedSelfAttention, self).__init__() - self.config = config - data_type = self.config.dtype - data_type_fp = torch.half if self.config.dtype == torch.int8 else self.config.dtype - self.config.layer_id = DeepSpeedSelfAttention.num_layers - DeepSpeedSelfAttention.num_layers = DeepSpeedSelfAttention.num_layers + 1 - device = get_accelerator().current_device_name() #if config.bigscience_bloom else 'cpu' - if self.config.set_empty_params: - self.attn_qw = None - self.attn_qb = None - self.attn_kw = None - self.attn_kb = None - self.attn_vw = None - self.attn_vb = None - self.attn_qkvw = None - self.attn_qkvb = None - self.attn_ow = None - self.attn_ob = None - else: - qkv_size_per_partition = (self.config.hidden_size // self.config.mp_size) * 3 if config.num_kv < 0 else \ - ((self.config.heads + self.config.num_kv * 2) // self.config.mp_size) * (self.config.hidden_size // self.config.heads) - self.attn_qkvw = nn.Parameter(torch.empty(self.config.hidden_size, - qkv_size_per_partition, - dtype=data_type, - device=device), - requires_grad=False) - self.attn_qkvb = nn.Parameter(torch.empty(qkv_size_per_partition, dtype=data_type_fp, device=device), - requires_grad=False) - out_size_per_partition = self.config.hidden_size // self.config.mp_size - self.attn_ow = nn.Parameter(torch.empty(out_size_per_partition, - self.config.hidden_size, - dtype=data_type, - device=device), - requires_grad=False) - - self.attn_ob = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), - requires_grad=False) - - self.num_attention_heads_per_partition = self.config.heads // self.config.mp_size - self.num_kv_partition = self.config.num_kv // self.config.mp_size - self.hidden_size_per_partition = self.config.hidden_size // self.config.mp_size - self.hidden_size_per_attention_head = self.config.hidden_size // self.config.heads - - self.mp_group = mp_group - - # used for quantization - self.q_scales = q_scales - self.q_groups = q_groups - self.merge_count = int(math.log2(merge_count)) - - self.norm_factor = math.sqrt(self.config.hidden_size // self.config.heads) - if not config.use_mup: - self.norm_factor = math.sqrt(self.norm_factor) - - if self.config.scale_attn_by_inverse_layer_idx is True: - self.norm_factor *= math.sqrt(self.config.layer_id + 1) - # https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/gpt2/modeling_gpt2.py#L191 - - self.qkv_func = QKVGemmOp(config) - self.score_context_func = SoftmaxContextOp(config) - self.linear_func = LinearOp(config) - self.vector_matmul_func = VectorMatMulOp(config) - if len(DeepSpeedSelfAttention._qkv_buffers) == 0: - DeepSpeedSelfAttention._qkv_buffers = [ - torch.empty(self.hidden_size_per_partition * 3, - self.config.hidden_size, - dtype=data_type_fp, - device=device), - torch.empty(self.hidden_size_per_partition * 3, dtype=data_type_fp, device=device) - ] - - def compute_attention(self, qkv_out, input_mask, layer_past, alibi, is_prompt, token_idx, position_ids): - if isinstance(qkv_out, list) or isinstance(qkv_out, tuple): - qkv_out = qkv_out[0] - - no_masking = input_mask is None or input_mask is False - - if no_masking: - input_mask = torch.empty(1) - - attn_key_value = self.score_context_func( - query_key_value=qkv_out, - attn_mask=((1 - input_mask).to(qkv_out.dtype) * - minus_inf) if input_mask.dtype == torch.int64 else input_mask, - heads=self.num_attention_heads_per_partition, - num_kv=self.num_kv_partition, - norm_factor=(1 / self.norm_factor if self.config.scale_attention else 1.0), - no_masking=no_masking, - layer_id=self.config.layer_id, - num_layers=DeepSpeedSelfAttention.num_layers, - alibi=alibi, - is_prompt=is_prompt, - token_idx=token_idx, - position_ids=position_ids) - - context_layer, key_layer, value_layer = attn_key_value - return context_layer, key_layer, value_layer - - def _merge_qkv(self): - qvkw = DeepSpeedSelfAttention._qkv_buffers[0] - qvkw[:self.hidden_size_per_partition, :] = self.attn_qw # type: ignore - qvkw[self.hidden_size_per_partition:2 * self.hidden_size_per_partition, :] = self.attn_kw # type: ignore - qvkw[2 * self.hidden_size_per_partition:, :] = self.attn_vw # type: ignore - if self.attn_qb is not None: - qvkb = DeepSpeedSelfAttention._qkv_buffers[1] - qvkb[:self.hidden_size_per_partition] = self.attn_qb - qvkb[self.hidden_size_per_partition:2 * self.hidden_size_per_partition] = self.attn_kb # type: ignore - qvkb[2 * self.hidden_size_per_partition:] = self.attn_vb # type: ignore - return DeepSpeedSelfAttention._qkv_buffers - - def forward(self, - input, - input_mask, - head_mask=None, - layer_past=None, - get_present=False, - encoder_hidden_states=None, - encoder_attention_mask=None, - output_attentions=False, - norm_w=None, - norm_b=None, - alibi=None, - **kwargs): - if self.attn_qkvw is None: - self._attn_qkvw, self._attn_qkvb = self._merge_qkv() - else: - self._attn_qkvw = self.attn_qkvw - self._attn_qkvb = self.attn_qkvb - if not self.config.pre_layer_norm: - qkv_out = self.linear_func(input=input, - weight=self._attn_qkvw, - bias=self._attn_qkvb, - add_bias=self.attn_qkvb is not None, - do_flash_attn=False, - num_heads=self.num_attention_heads_per_partition, - num_layers=DeepSpeedSelfAttention.num_layers) - else: - qkv_out = self.qkv_func(input=input, - weight=self._attn_qkvw, - bias=self._attn_qkvb, - gamma=norm_w, - beta=norm_b) - - is_prompt = kwargs.get("first_token", qkv_out[0].shape[1] > 1) - token_idx = kwargs.get("token_idx", None) - position_ids = kwargs.get("position_ids", None) - - context_layer, key_layer, value_layer = self.compute_attention(qkv_out=qkv_out, - input_mask=input_mask, - layer_past=layer_past, - alibi=alibi, - is_prompt=is_prompt, - token_idx=token_idx, - position_ids=position_ids) - - output = self.vector_matmul_func(input=context_layer, weight=self.attn_ow) - inp_norm = qkv_out[-1] - - if self.config.mlp_after_attn and self.mp_group is not None and dist.get_world_size(group=self.mp_group) > 1: - dist.all_reduce(output, group=self.mp_group) - return (output, key_layer, value_layer, context_layer, inp_norm) - - -class BloomSelfAttention(DeepSpeedSelfAttention): - - def __init__(self, *args, **kwargs): - super(BloomSelfAttention, self).__init__(*args, **kwargs) - self.softmax_func = SoftmaxOp(self.config) - - ########### This part is taken/modified form the HF modeling_bloom.py ################ - # Reference: https://github.com/huggingface/transformers/blob/main/src/transformers/models/bloom/modeling_bloom.py - - def _transpose_for_context(self, x): - x = x.permute(0, 2, 1, 3).contiguous() - new_x_layer_shape = x.size()[:-2] + \ - (self.hidden_size_per_partition,) - return x.view(*new_x_layer_shape).contiguous() - - def _split_tensor_along_last_dim(self, tensor, num_partitions, contiguous_split_chunks=True): - """Split a tensor along its last dimension. - - Args: - tensor: ([`torch.tensor`], *required*): - input tensor to split - num_partitions ([`int`], *required*): - number of partitions to split the tensor - contiguous_split_chunks ([`bool`], *optional*, default=`False`):: - If True, make each chunk contiguous in memory. - """ - # Get the size and dimension. - last_dim = tensor.dim() - 1 - numerator, denominator = tensor.size()[last_dim], num_partitions - if not (numerator % denominator == 0): - raise ValueError(f"{numerator} is not divisible by {denominator}") - last_dim_size = numerator // denominator - # Split. - tensor_list = torch.split(tensor, last_dim_size, dim=last_dim) - # Note: torch.split does not create contiguous tensors by default. - if contiguous_split_chunks: - return tuple(chunk.contiguous() for chunk in tensor_list) - - return tensor_list - - def compute_attention(self, qkv_out, input_mask, layer_past, alibi, is_prompt, token_idx, position_ids): - if isinstance(qkv_out, list) or isinstance(qkv_out, tuple): - qkv_out = qkv_out[0] - - no_masking = input_mask is None - - if no_masking: - input_mask = torch.empty(1) - - mixed_x_layer = qkv_out - alibi = alibi.to(get_accelerator().current_device_name()) - head_dim = self.hidden_size_per_partition // self.num_attention_heads_per_partition - new_tensor_shape = mixed_x_layer.size()[:-1] + (self.num_attention_heads_per_partition, 3 * head_dim) - mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) - - query_layer, key_layer, value_layer = self._split_tensor_along_last_dim(mixed_x_layer, 3) - - # [batch_size, head_dim, q_length, k_length] - output_size = (query_layer.size(0), query_layer.size(2), query_layer.size(1), key_layer.size(1)) - # [batch_size, q_length, num_heads, head_dim] -> [q_length, batch_size * num_heads, head_dim] - query_layer = query_layer.transpose(1, 2).reshape(output_size[0] * output_size[1], output_size[2], -1) - # [batch_size, k_length, num_heads, head_dim] -> [k_length, batch_size * num_heads, head_dim] - key_layer = key_layer.transpose(1, 2).reshape(output_size[0] * output_size[1], output_size[3], - -1).transpose(-1, -2) - value_layer = value_layer.transpose(1, 2).reshape(output_size[0] * output_size[1], output_size[3], -1) - if layer_past is not None: - past_key, past_value = layer_past - # concatenate along seq_length dimension -> [batch_size, qk_length, num_heads, head_dim] - key_layer = torch.cat((past_key.type_as(key_layer), key_layer), dim=-1) - value_layer = torch.cat((past_value.type_as(value_layer), value_layer), dim=-2) - - presents = (key_layer, value_layer) - # Raw attention scores. [batch_size * num_heads, q_length, k_length] - matmul_result = torch.matmul(query_layer, key_layer) - # change view to [batch_size, num_heads, q_length, k_length] - attention_scores = matmul_result.view(output_size[0], output_size[1], output_size[2], -1) - - offset = dist.get_rank() * self.num_attention_heads_per_partition if dist.is_initialized() else 0 - target_dtype = torch.float16 if self.config.dtype == torch.int8 else self.config.dtype - - # When using the hybrid engine with BLOOM, input_mask needs to be converted from torch.bool -> torch.int64 - if input_mask.dtype == torch.bool: - input_mask = input_mask.long() - - # Invert input_mask per transformer implementation (eg, in BLOOM, it's already inverted) - if self.config.invert_mask: - input_mask = 1 - input_mask - - attention_probs = self.softmax_func(attn_scores=attention_scores, - attn_mask=input_mask.to(target_dtype) * minus_inf, - alibi=alibi, - triangular=(self.config.triangular_masking - and (attention_scores.shape[-2] > 1)), - recompute=False, - local_attention=False, - window_size=1, - async_op=False, - layer_scale=1 / (self.norm_factor * self.norm_factor), - head_offset=offset) - - # change view [batch_size x num_heads, q_length, k_length] - attention_probs_reshaped = attention_probs.view(*matmul_result.shape) - - # matmul: [batch_size * num_heads, q_length, head_dim] - context_layer = torch.bmm(attention_probs_reshaped, value_layer) - - # change view [batch_size, num_heads, q_length, head_dim] - context_layer = context_layer.view( - context_layer.size(0) // self.num_attention_heads_per_partition, self.num_attention_heads_per_partition, - context_layer.size(1), context_layer.shape[-1]) - - context_layer = self._transpose_for_context(context_layer) - key_layer = presents[0] - value_layer = presents[1] - - return context_layer, key_layer, value_layer - - ###################### End of HF modeling_bloom addition ######################## diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/ds_mlp.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/ds_mlp.py deleted file mode 100644 index 36de06db920fb51adf31809e02d74a797494938c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/ds_mlp.py +++ /dev/null @@ -1,124 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import math -import torch -import torch.nn as nn -from deepspeed import comm as dist -from deepspeed.utils.types import GATED_ACTIVATION_TYPES -from deepspeed.accelerator import get_accelerator -from .op_binding import MLPGemmOp, VectorMatMulOp, GELUGemmOp, ResidualAddOp - - -class DeepSpeedMLP(nn.Module): - _inter_w_buffers = [] - - def __init__(self, config, mp_group=None, q_scales=None, q_groups=1, merge_count=1, mlp_extra_grouping=False): - super(DeepSpeedMLP, self).__init__() - - self.config = config - - data_type = torch.int8 if self.config.dtype == torch.int8 else self.config.dtype - data_type_fp = torch.half if self.config.dtype == torch.int8 else self.config.dtype - device = get_accelerator().current_device_name() - - proj_factor = 2 if self.config.mlp_act_func_type in GATED_ACTIVATION_TYPES else 1 - self.config.intermediate_size = self.config.intermediate_size if self.config.intermediate_size > 0 else 4 * self.config.hidden_size - self.intm_w_sz_per_partition = self.config.intermediate_size * proj_factor // self.config.mp_size - self.intm_o_sz_per_partition = self.config.intermediate_size // self.config.mp_size - - if self.config.set_empty_params: - self.attn_nw = None - self.attn_nb = None - self.inter_w = None - self.inter_b = None - self.inter_up_w = None - self.inter_up_b = None - self.inter_gate_w = None - self.inter_gate_b = None - self.output_w = None - self.output_b = None - else: - self.attn_nw = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), - requires_grad=False) - self.attn_nb = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), - requires_grad=False) - - self.inter_w = nn.Parameter(torch.empty(self.config.hidden_size, - self.intm_w_sz_per_partition, - dtype=data_type, - device=device), - requires_grad=False) - self.inter_b = nn.Parameter(torch.empty(self.intm_w_sz_per_partition, dtype=data_type_fp, device=device), - requires_grad=False) - self.output_w = nn.Parameter(torch.empty(self.intm_o_sz_per_partition, - self.config.hidden_size, - dtype=data_type, - device=device), - requires_grad=False) - self.output_b = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), - requires_grad=False) - - # used for quantization - self.q_scales = q_scales - self.q_groups = q_groups * 2 if mlp_extra_grouping else q_groups - self.merge_count = int(math.log2(merge_count)) - self.mp_group = mp_group - - self.mlp_gemm_func = MLPGemmOp(config) - self.vector_matmul_func = VectorMatMulOp(config) - self.fused_gemm_gelu = GELUGemmOp(config) - self.residual_add_func = ResidualAddOp(config) - - if len(DeepSpeedMLP._inter_w_buffers) == 0: - DeepSpeedMLP._inter_w_buffers = [ - torch.empty(self.intm_w_sz_per_partition, self.config.hidden_size, dtype=data_type, device=device), - torch.empty(self.intm_w_sz_per_partition, dtype=data_type_fp, device=device) - ] - - def _merge_inter_w(self): - inter_w = DeepSpeedMLP._inter_w_buffers[0] - inter_w[:self.intm_w_sz_per_partition // 2, :] = self.inter_up_w # type: ignore - inter_w[self.intm_w_sz_per_partition // 2:, :] = self.inter_gate_w # type: ignore - if self.inter_up_b is not None: - inter_b = DeepSpeedMLP._inter_w_buffers[1] - inter_b[:self.intm_w_sz_per_partition // 2] = self.inter_up_b # type: ignore - inter_b[self.intm_w_sz_per_partition // 2:] = self.inter_gate_b # type: ignore - return DeepSpeedMLP._inter_w_buffers - - def forward(self, input, residual, residual_norm, bias): - if self.inter_w is None: - self._inter_w, self._inter_b = self._merge_inter_w() - else: - self._inter_w = self.inter_w - self._inter_b = self.inter_b - - residual_add = None - if self.attn_nw is None: - output = self.fused_gemm_gelu(input=residual_norm, - weight=self._inter_w, - bias=self._inter_b, - weight_out=self.output_w) - else: - output, residual_add = self.mlp_gemm_func(input=input, - residual=residual, - weight_interm=self._inter_w, - weight_out=self.output_w, - input_bias=bias, - bias=self._inter_b, - gamma=self.attn_nw, - beta=self.attn_nb) - - residual = self.residual_add_func(hidden_state=output, - residual=residual, - add_bias=bias is not None, - attention_output=input, - attention_bias=bias if bias is not None else self.output_b, - final_bias=self.output_b, - residual_add=residual_add) - if self.mp_group is not None and dist.get_world_size(group=self.mp_group) > 1: - dist.all_reduce(residual, group=self.mp_group) - - return residual diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/moe_inference.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/moe_inference.py deleted file mode 100644 index 3a9785985d196704abe162cd6cef6f75438c7e7a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/moe_inference.py +++ /dev/null @@ -1,350 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import json -import math -import torch -from torch.autograd import Function -import torch.nn as nn -from .ds_attention import DeepSpeedSelfAttention -from .config import DeepSpeedInferenceConfig -from .op_binding import SoftmaxOp, VectorMatMulOp, GELUGemmOp -from .op_binding.bias_residual import BiasResidualOp -from .op_binding.einsum_sec_sm_ecm import EinsumSecSmEcmOp -from .op_binding.layer_norm import LayerNormOp -from ....moe.sharded_moe import TopKGate -from deepspeed import comm as dist -from .op_binding.moe_res_matmul import MoEResMatmulOp - - -class DeepSpeedMoEInferenceConfig(DeepSpeedInferenceConfig): - """Initialize the DeepSpeed Transformer Config. - Arguments: - hidden_size: The hidden size of the transformer layer - intermediate_size: The intermediate size of the feed-forward part of transformer layer - heads: The number of heads in the self-attention of the transformer layer - num_hidden_layers: The number of transformer layers - layer_norm_eps: The epsilon value for the layer norm - local_rank: Optional: The rank of GPU running the transformer kernel, it is not required - to use if the model already set the current device, otherwise need to set it - so that the transformer kernel can work on the right device - mp_size (optional): This argument is mainly used to create the parameters on the kernel side - using model-parallel architecture. If the client model already takes care of this, there is no - need to pass this argument. - fp16: Enable half-precision computation - bf16: Enable bf16 floating point computation - pre_layer_norm: Select between Pre-LN or Post-LN transformer architecture - stochastic_mode: Enable for high performance, please note that this flag has some level of - non-determinism and can produce different results on different runs. However, we have seen - that by enabling it, the pretraining tasks such as BERT are not affected and can obtain - a high accuracy level. On the other hand, for the downstream tasks, such as fine-tuning, we recommend - to turn it off in order to be able to reproduce the same result through the regular kernel execution. - - scale_attention: If true, both q and k are scaled by 1/sqrt(attention_heads) before attention computation. - return_tuple: if True, returns the transformer output as a tuple, otherwise returns as a tensor - """ - - def __init__(self, - hidden_size=-1, - intermediate_size=-1, - heads=-1, - num_hidden_layers=-1, - layer_norm_eps=1e-12, - local_rank=-1, - mp_size=1, - fp16=False, - bf16=False, - q_int8=False, - pre_layer_norm=True, - stochastic_mode=False, - scale_attention=True, - triangular_masking=True, - local_attention=False, - window_size=256, - return_tuple=True, - moe_experts=1, - global_experts=1, - k=1, - capacity_factor=1., - eval_capacity_factor=1., - min_capacity=1, - noisy_gate_policy=None, - drop_tokens=True, - use_rts=False, - mlp_type='standard', - scale_attn_by_inverse_layer_idx=False): - super(DeepSpeedMoEInferenceConfig, - self).__init__(hidden_size, (intermediate_size if intermediate_size > 0 else 4 * hidden_size), heads, - num_hidden_layers, layer_norm_eps, local_rank, mp_size, fp16, bf16, q_int8, - pre_layer_norm, stochastic_mode, scale_attention, triangular_masking, local_attention, - window_size, return_tuple) - self.moe_experts = moe_experts - self.k = k - self.capacity_factor = capacity_factor - self.eval_capacity_factor = eval_capacity_factor - self.min_capacity = min_capacity - self.noisy_gate_policy = noisy_gate_policy - self.drop_tokens = drop_tokens - self.use_rts = use_rts - self.global_experts = global_experts - self.mlp_type = mlp_type - self.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx - - @classmethod - def from_dict(cls, json_object): - config = DeepSpeedInferenceConfig() - for key, value in json_object.items(): - config.__dict__[key] = value - return config - - @classmethod - def from_json_file(cls, json_file): - with open(json_file, "r", encoding='utf-8') as reader: - text = reader.read() - return cls.from_dict(json.loads(text)) - - -class DeepSpeedMLPFunction(Function): - - @staticmethod - def forward(ctx, input, inter_w, inter_b, config, output_b, output_w, q_scales, q_groups, merge_count, mp_group, - async_op, gelu_gemm_func, vector_matmul_func): - if config.q_int8: - intermediate = gelu_gemm_func(input, inter_w, inter_b, config.epsilon, q_scales[2], - (q_groups * (2**merge_count)), config.pre_layer_norm) - output = vector_matmul_func(intermediate, output_w, q_scales[3], q_groups, (merge_count)) - else: - output = gelu_gemm_func(input, inter_w, inter_b, output_w, config.epsilon, config.pre_layer_norm, async_op) - if mp_group is not None and dist.get_world_size(group=mp_group) > 1: - dist.all_reduce(output, group=mp_group, async_op=async_op) - - return output + output_b - - @staticmethod - def backward(ctx, grad_output): - raise RuntimeError('You are running with DeepSpeed Inference mode. \ - Please switch to Training mode for running backward!') - - -class DeepSpeedMoEMLP(nn.Module): - - def __init__(self, config, q_scales=None, q_groups=1, merge_count=1, mlp_extra_grouping=False, mp_group=None): - super(DeepSpeedMoEMLP, self).__init__() - - self.config = config - self.attn_nw = nn.Parameter(torch.Tensor(self.config.hidden_size)) - self.attn_nb = nn.Parameter(torch.Tensor(self.config.hidden_size)) - interm_size = self.config.intermediate_size // (1 if mp_group is None else dist.get_world_size(group=mp_group)) - self.inter_w = nn.Parameter(torch.Tensor(self.config.hidden_size, interm_size)) - self.inter_b = nn.Parameter(torch.Tensor(interm_size)) - self.output_w = nn.Parameter(torch.Tensor((interm_size), self.config.hidden_size)) - self.output_b = nn.Parameter(torch.Tensor(self.config.hidden_size)) - - # used for quantization - self.q_scales = q_scales - self.q_groups = q_groups * 2 if mlp_extra_grouping else q_groups - self.merge_count = int(math.log2(merge_count)) - self.mp_group = mp_group - self.gelu_gemm_func = GELUGemmOp(self.config) - self.vector_matmul_func = VectorMatMulOp(self.config) - - def forward(self, input, async_op=False): - return DeepSpeedMLPFunction.apply(input, self.inter_w, self.inter_b, self.config, self.output_b, self.output_w, - self.q_scales, self.q_groups, self.merge_count, self.mp_group, async_op, - self.gelu_gemm_func, self.vector_matmul_func) - - -class DeepSpeedMoEInference(nn.Module): - """Initialize the DeepSpeed MoE Transformer Layer. - Arguments: - layer_id: The layer index starting from 0, e.g. if model has 24 transformer layers, - layer_id will be 0,1,2...23 when each layer object is instantiated - config: An object of DeepSpeedInferenceConfig - mp_group: Model parallelism group initialized on the modeling side. - quantize_scales: This argument groups all the layers' scales used for quantization - quantize_groups: Number of groups used for quantizing the model - merge_count: Shows the number of model-parallel checkpoints merged before running inference. - We use this argument to control the quantization scale for the model parameters if a bigger - quantize-grouping than 1 is used. - mlp_extra_grouping: This flag is used to show a 2x higher number of groups used for the MLP part - of a Transformer layer. We use this feature for quantization to reduce the convergence impact - for specific downstream tasks. - """ - layer_id = 0 - - def __init__(self, - config, - mp_group=None, - ep_group=None, - expert_mp_group=None, - quantize_scales=None, - quantize_groups=1, - merge_count=1, - mlp_extra_grouping=False): - super(DeepSpeedMoEInference, self).__init__() - - self.config = config - self.config.layer_id = DeepSpeedMoEInference.layer_id - - assert self.config.dtype != torch.bfloat16, "DeepSpeed MoE Transformer Inference not yet tested for bfloat support" - - DeepSpeedMoEInference.layer_id += 1 - self.attention = DeepSpeedSelfAttention(self.config, mp_group, quantize_scales, quantize_groups, merge_count) - self.attn_nw = nn.Parameter(torch.Tensor(self.config.hidden_size)) - self.attn_nb = nn.Parameter(torch.Tensor(self.config.hidden_size)) - - self.norm_w = nn.Parameter(torch.Tensor(self.config.hidden_size)) - self.norm_b = nn.Parameter(torch.Tensor(self.config.hidden_size)) - - if config.mlp_type == 'residual': - self.res_mlp = DeepSpeedMoEMLP(config, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping, - mp_group) - self.res_coef = nn.Parameter(torch.Tensor(self.config.hidden_size, 2)) - self.coef_func = SoftmaxOp(self.config) - self.vector_matmul_func = VectorMatMulOp(self.config) - - config.mp_size = 1 - self.mlp = nn.ModuleList( - DeepSpeedMoEMLP(config, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping, expert_mp_group) - for i in range(self.config.moe_experts)) - - self.moe_gate = TopKGate(self.config.hidden_size, self.config.global_experts, self.config.k, - self.config.capacity_factor, self.config.eval_capacity_factor, - self.config.min_capacity, self.config.noisy_gate_policy, self.config.drop_tokens, - self.config.use_rts, self.ep_group) - - self.ep_group = ep_group - self.mp_group = mp_group - self.expert_mp_group = expert_mp_group - - print("DeepSpeed MoE Transformer Inference config is ", self.config.__dict__) - - self.bias_residual_func = BiasResidualOp(self.config) - self.ds_layernorm = LayerNormOp(self.config) - self.einsum_sec_sm_ecm = EinsumSecSmEcmOp(self.config) - self.moe_res_matmul = MoEResMatmulOp(self.config) - - def res_coef_func(self, inp, async_op): - inp = self.vector_matmul_func(inp, self.res_coef, async_op) - return self.coef_func(inp, torch.empty(1), False, False, False, 256, async_op) - - def moe_gate_einsum(self, attention_output): - _, combined_weights, dispatch_mask, _ = self.moe_gate( - attention_output.view(-1, self.config.hidden_size), - None, - ) - dispatched_attention = self.einsum_sec_sm_ecm(dispatch_mask.type_as(attention_output), - attention_output.view(-1, self.config.hidden_size)) - return dispatched_attention, combined_weights - - def expert_exec(self, dispatched_input): - dispatched_input = dispatched_input.reshape(self.config.global_experts // self.config.moe_experts, - self.config.moe_experts, -1, self.config.hidden_size) - - chunks = dispatched_input.chunk(self.config.moe_experts, dim=1) - expert_outputs = torch.empty(( - self.config.moe_experts, - chunks[0].shape[0], - ) + chunks[0].shape[2:], - dtype=dispatched_input.dtype, - device=dispatched_input.device) - for chunk, expert in zip(chunks, range(len(self.mlp))): - expert_outputs[expert] = self.mlp[expert](chunk.view(-1, dispatched_input.shape[-2], - dispatched_input.shape[-1])) - return expert_outputs - - def _alltoall(self, dispatched_attention): - if dist.get_world_size(group=self.ep_group) > 1: - dispatched_input = torch.empty_like(dispatched_attention) - dist.all_to_all_single(dispatched_input, dispatched_attention, group=self.ep_group) - return dispatched_input - else: - return dispatched_attention - - def scale_expert_output(self, attention_output, expert_output, combined_weights): - combined_output = torch.matmul( - combined_weights.type_as(attention_output).reshape(combined_weights.shape[0], -1), - expert_output.reshape(-1, expert_output.shape[-1])) - return combined_output.reshape(attention_output.shape) - - def forward(self, - input, - input_mask=None, - attention_mask=None, - head_mask=None, - layer_past=None, - get_key_value=False, - get_present=False, - encoder_output=None, - enc_dec_attn_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - use_cache=False, - output_attentions=False): - get_present = (get_present or get_key_value or use_cache) - input_mask = input_mask if attention_mask is None else attention_mask - input_type = input.dtype - - if (self.config.dtype in [torch.float16, torch.int8]) and input_type == torch.float: - input = input.half() - - with torch.no_grad(): - attention_output = self.attention(input, input_mask, head_mask, layer_past, get_present, - encoder_hidden_states, encoder_attention_mask, output_attentions, - self.norm_w, self.norm_b) - - if get_present: - attention_output, p_key, p_value = attention_output[0:3] - presents = (p_key, p_value) - elif output_attentions: - attention_output, _, _, context_output = attention_output[0:4] - else: - attention_output = attention_output[0] - - residual_add = attention_output + self.attention.attn_ob - attention_output = self.ds_layernorm(residual_add, self.attn_nw, self.attn_nb, self.config.epsilon) - - if self.config.mlp_type == 'residual': - res_mlp_out = self.res_mlp(attention_output, async_op=True) - res_coef_out = self.res_coef_func(attention_output, async_op=True) - - if self.expert_mp_group is not None: - world_size = dist.get_world_size(group=self.expert_mp_group) - gather_buffer = torch.empty(world_size * attention_output.numel(), - dtype=attention_output.dtype, - device=attention_output.device) - dist.all_gather_into_tensor(gather_buffer, attention_output, group=self.expert_mp_group) - attention_output = gather_buffer.view(-1, *attention_output.size()[1:]) - - ############## MoE Gating + Experts ############### - dispatched_attention, combined_weights = self.moe_gate_einsum(attention_output) - dispatched_input = self._alltoall(dispatched_attention) - expert_outputs = self.expert_exec(dispatched_input) - expert_output = self._alltoall(expert_outputs) - output = self.scale_expert_output(attention_output, expert_output, combined_weights) - ################################################ - - if self.expert_mp_group is not None: - output = output.split(output.shape[0] // dist.get_world_size(group=self.expert_mp_group), - dim=0)[dist.get_rank(group=self.expert_mp_group)] - - if self.config.mlp_type == 'residual': - self.moe_res_matmul(res_mlp_out, res_coef_out, output) - - output = self.bias_residual_func(output, residual_add, torch.empty(1)) - - if not self.config.pre_layer_norm: - output = self.ds_layernorm(output, self.norm_w, self.norm_b, self.config.epsilon) - - if input_type != output.dtype: - output = output.to(input_type) - - if get_present: - output = (output, presents) - - if self.config.return_tuple: - return output if type(output) is tuple else (output, ) - else: - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__init__.py deleted file mode 100644 index 20b7bf12a917865a753d9db41851562edc8cb337..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .linear import LinearOp -from .vector_matmul import VectorMatMulOp -from .softmax_context import SoftmaxContextOp -from .qkv_gemm import QKVGemmOp -from .softmax import SoftmaxOp -from .mlp_gemm import MLPGemmOp -from .gelu_gemm import GELUGemmOp -from .residual_add import ResidualAddOp diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 012102c60704f4047de723a1ff4367950069d4ed..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/base.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/base.cpython-310.pyc deleted file mode 100644 index f04241896075d4bcf6162c8a52e2c08bfc61b03a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_add.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_add.cpython-310.pyc deleted file mode 100644 index c2e81d0dfd5b72c207377fce956919de089fb482..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_add.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_gelu.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_gelu.cpython-310.pyc deleted file mode 100644 index 7e6c8db957654add351ab1a0390932618365fea6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_gelu.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_relu.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_relu.cpython-310.pyc deleted file mode 100644 index b7d8f33064ee43aa838b594ec35db9df8a430f60..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_relu.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_residual.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_residual.cpython-310.pyc deleted file mode 100644 index 6b81cbd18ecc0f49e4cbfba4d85692d3688710ba..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/bias_residual.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/einsum_sec_sm_ecm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/einsum_sec_sm_ecm.cpython-310.pyc deleted file mode 100644 index bb411252369571c7b94392eec6edc73d92e31417..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/einsum_sec_sm_ecm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/gated_activation.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/gated_activation.cpython-310.pyc deleted file mode 100644 index 2c07b4e5659091b4d14bc9f85f3492ca6d2875e6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/gated_activation.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/gelu_gemm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/gelu_gemm.cpython-310.pyc deleted file mode 100644 index 58ca4bc87ea6eb511de28ef9dd5f01cabd833e58..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/gelu_gemm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/layer_norm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/layer_norm.cpython-310.pyc deleted file mode 100644 index e1864d0c499379053f6989b479282bbd6ded512f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/layer_norm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/linear.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/linear.cpython-310.pyc deleted file mode 100644 index 624c384dd3f8b08673a3af0c09d34d0e989afb1c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/linear.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/mlp_gemm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/mlp_gemm.cpython-310.pyc deleted file mode 100644 index df781f3207bbf022318d81b351dde3843fcc9114..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/mlp_gemm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/moe_res_matmul.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/moe_res_matmul.cpython-310.pyc deleted file mode 100644 index 6d57f9436b5b96b3534b27a226b2622192916c37..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/moe_res_matmul.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/pad_transform.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/pad_transform.cpython-310.pyc deleted file mode 100644 index 33d94db7b710e52e4de0027303868526ff3372f9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/pad_transform.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/pre_rms_norm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/pre_rms_norm.cpython-310.pyc deleted file mode 100644 index ac6fd8ce74310f49cbb928c24c23071a288dd1e8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/pre_rms_norm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/qkv_gemm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/qkv_gemm.cpython-310.pyc deleted file mode 100644 index 7e680b4c826475e552278871f7bae0b8059e1ba9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/qkv_gemm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/residual_add.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/residual_add.cpython-310.pyc deleted file mode 100644 index 1bf4da17454efbbc98ec993561112b4d797f75d5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/residual_add.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/rms_norm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/rms_norm.cpython-310.pyc deleted file mode 100644 index f99b0c38a99b83657f88899c9648c15f04a3dd59..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/rms_norm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax.cpython-310.pyc deleted file mode 100644 index 9debe159dc3f10ac29b6b28b7534f5c17d952d6f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax_context.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax_context.cpython-310.pyc deleted file mode 100644 index 76f4d12318e319bd0845b838333bf0f753dcb22a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/softmax_context.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_add.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_add.cpython-310.pyc deleted file mode 100644 index 5e7e5d7860166e9eaa2ca1639044c95be466234e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_add.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_matmul.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_matmul.cpython-310.pyc deleted file mode 100644 index 9b84447b84ffb34d479c4561d005a54327d3a7f2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/vector_matmul.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/workspace.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/workspace.cpython-310.pyc deleted file mode 100644 index b901a6197f3083afe5d889a2c764cc9b7b36db18..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/__pycache__/workspace.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/base.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/base.py deleted file mode 100644 index 5a997f95d5cc53e3323180ad7cd70ed9e8210a1b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/base.py +++ /dev/null @@ -1,20 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from ..config import DeepSpeedInferenceConfig - -from deepspeed.ops.op_builder import InferenceBuilder - - -class BaseOp(torch.nn.Module): - inference_module = None - - def __init__(self, config: DeepSpeedInferenceConfig): - super(BaseOp, self).__init__() - self.config = config - if BaseOp.inference_module is None: - builder = InferenceBuilder() - BaseOp.inference_module = builder.load() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_add.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_add.py deleted file mode 100644 index d2ae38f546eb431ccd368b5da730cfe6a901d431..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_add.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp - - -class BiasAddOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig): - super(BiasAddOp, self).__init__(config) - - try: - if self.config.dtype == torch.float16: - self.bias_add_func = self.inference_module.bias_add_fp16 - elif self.config.dtype == torch.bfloat16: - self.bias_add_func = self.inference_module.bias_add_bf16 - else: - self.bias_add_func = self.inference_module.bias_add_fp32 - except AttributeError: - self.bias_add_func = self.bias_add_fallback - - @classmethod - def bias_add_fallback(cls, input, bias): - return torch.add(input, bias) - - def forward(self, activation: torch.Tensor, bias: torch.Tensor): - return self.bias_add_func(activation, bias) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_gelu.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_gelu.py deleted file mode 100644 index f0fee0b0d06ea906d3c3229d5b3b27cd15c4f2c3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_gelu.py +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import torch.nn.functional as F -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp - - -class BiasGeluOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig): - super(BiasGeluOp, self).__init__(config) - - try: - if self.config.dtype == torch.float16: - self.bias_gelu_func = self.inference_module.bias_gelu_fp16 - elif self.config.dtype == torch.bfloat16: - self.bias_gelu_func = self.inference_module.bias_gelu_bf16 - else: - self.bias_gelu_func = self.inference_module.bias_gelu_fp32 - except AttributeError: - self.bias_gelu_func = self.bias_gelu_fallback - - @classmethod - def bias_gelu_fallback(cls, activations, bias): - # Expected behavior is that of casting to float32 internally and using the tanh approximation - return F.gelu(activations.to(torch.float32) + bias.to(torch.float32), approximate='tanh').to(activations.dtype) - - def forward(self, activation: torch.Tensor, bias: torch.Tensor): - return self.bias_gelu_func(activation, bias) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_relu.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_relu.py deleted file mode 100644 index ccfade1d9524562fb30f9d758cadd33191f53a45..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_relu.py +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import torch.nn.functional as F -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp - - -class BiasReluOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig): - super(BiasReluOp, self).__init__(config) - - try: - if self.config.dtype == torch.float16: - self.bias_relu_func = self.inference_module.bias_relu_fp16 - elif self.config.dtype == torch.bfloat16: - self.bias_relu_func = self.inference_module.bias_relu_bf16 - else: - self.bias_relu_func = self.inference_module.bias_relu_fp32 - except AttributeError: - self.bias_relu_func = self.bias_relu_fallback - - @classmethod - def bias_relu_fallback(cls, activations, bias): - # Expected behavior is that of casting to float32 internally - return F.relu(activations.to(torch.float32) + bias.to(torch.float32)).to(activations.dtype) - - def forward(self, activation: torch.Tensor, bias: torch.Tensor): - return self.bias_relu_func(activation, bias) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_residual.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_residual.py deleted file mode 100644 index ecad50e10ffef489eea41547c036124df057acb2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/bias_residual.py +++ /dev/null @@ -1,29 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp - - -class BiasResidualOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig): - super(BiasResidualOp, self).__init__(config) - - try: - if self.config.dtype in [torch.float16, torch.int8]: - self.bias_residual_func = self.inference_module.bias_residual_fp16 - else: - self.bias_residual_func = self.inference_module.bias_residual_fp32 - except AttributeError: - self.bias_residual_func = self.bias_residual_fallback - - @classmethod - def bias_residual_fallback(cls, output, residual, bias): - raise NotImplementedError("bias residual fallback isn't implemented") - - def forward(self, output, residual, bias): - return self.bias_residual_func(output, residual, bias) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/einsum_sec_sm_ecm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/einsum_sec_sm_ecm.py deleted file mode 100644 index f34b10f786d11b87c0c8943ecba375c079451484..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/einsum_sec_sm_ecm.py +++ /dev/null @@ -1,29 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp - - -class EinsumSecSmEcmOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig): - super(EinsumSecSmEcmOp, self).__init__(config) - - try: - if self.config.dtype in [torch.float16, torch.int8]: - self.einsum_sec_sm_ecm_func = self.inference_module.einsum_sec_sm_ecm_fp16 - else: - self.einsum_sec_sm_ecm_func = self.inference_module.einsum_sec_sm_ecm_fp32 - except AttributeError: - self.einsum_sec_sm_ecm_func = self.einsum_sec_sm_ecm_fallback - - @classmethod - def einsum_sec_sm_ecm_fallback(cls, Q, W): - raise NotImplementedError("einsum sec sm ecm fallback isn't implemented") - - def forward(self, Q, W): - return self.einsum_sec_sm_ecm_func(Q, W) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/gated_activation.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/gated_activation.py deleted file mode 100644 index d28d818ce4b33fbe5429ed9dd2c3e458baa52316..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/gated_activation.py +++ /dev/null @@ -1,40 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import torch.nn.functional as F -from deepspeed.utils.types import ActivationFuncType -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp - - -class GatedActivationOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig = None): - if config is None: - config = DeepSpeedInferenceConfig() - super(GatedActivationOp, self).__init__(config) - try: - self.gated_activation_func = self.inference_module.gated_activation - except AttributeError: - self.gated_activation_func = self.gated_activation_fallback - - @classmethod - def gated_activation_fallback(cls, activation, bias, activation_func_type): - # Expected behavior is that of casting to float32 internally - # Explicitly using the default GeLU - activation_func = None - activations = activation + bias.reshape(1, 1, -1) - hidden_states, gate = activations.chunk(2, dim=-1) - - if activation_func_type == ActivationFuncType.GATED_SILU: - activation_func = F.silu - elif activation_func_type == ActivationFuncType.GATED_GELU: - activation_func = F.gelu - - return hidden_states * activation_func(gate.to(torch.float32)).to(activations.dtype) - - def forward(self, activation: torch.Tensor, bias: torch.Tensor, activation_func_type: ActivationFuncType): - return self.gated_activation_func(activation, bias, activation_func_type) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/gelu_gemm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/gelu_gemm.py deleted file mode 100644 index 60bbb4b48bdb3977f037ac56f5a3008e3bea2af0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/gelu_gemm.py +++ /dev/null @@ -1,52 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import torch.nn.functional as F -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp -import deepspeed - - -class GELUGemmOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig): - super(GELUGemmOp, self).__init__(config) - try: - if self.config.dtype == torch.int8: - self.fused_gemm_gelu = self.inference_module.fused_gemm_gelu_int8 - elif self.config.dtype == torch.float16: - if deepspeed.HAS_TRITON and self.config.use_triton and self.config.dtype == torch.float16: - from deepspeed.ops.transformer.inference.triton.ops import fused_gemm_gelu as _triton_fused_gemm_gelu - self.fused_gemm_gelu = _triton_fused_gemm_gelu # type: ignore - else: - self.fused_gemm_gelu = self.inference_module.fused_gemm_gelu_fp16 # type: ignore - elif self.config.dtype == torch.bfloat16: - self.fused_gemm_gelu = self.inference_module.fused_gemm_gelu_bf16 # type: ignore - else: - self.fused_gemm_gelu = self.inference_module.fused_gemm_gelu_fp32 # type: ignore - except AttributeError: - self.fused_gemm_gelu = self.gelu_gemm_fallback - - def gelu_gemm_fallback(self, input, weight, scale, bias, out, out_scale, dtype, transpose): - tmp = torch.matmul(input, weight) - tmp = F.gelu(tmp.to(torch.float32) + bias.to(torch.float32), approximate="tanh").to(tmp.dtype) - output = torch.matmul(tmp, out) - - return output - - def forward(self, input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, weight_out: torch.Tensor): - - output = self.fused_gemm_gelu( - input, - weight, - weight.scale if hasattr(weight, 'scale') else torch.empty(1), # type: ignore - bias, - weight_out, - weight_out.scale if hasattr(weight_out, 'scale') else torch.empty(1), # type: ignore - self.config.dtype == torch.int8, - self.config.transposed_mode) - - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/layer_norm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/layer_norm.py deleted file mode 100644 index 31219a58ac3c820ddb2d9a301860274ad84e5755..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/layer_norm.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import torch.nn.functional as F -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp - - -class LayerNormOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig = None): - super(LayerNormOp, self).__init__(config) - try: - if config is None: - self.layer_norm_func = self.inference_module.layer_norm - elif self.config.dtype in [torch.float16, torch.int8]: - self.layer_norm_func = self.inference_module.layer_norm_fp16 - else: - self.layer_norm_func = self.inference_module.layer_norm_fp32 - except AttributeError: - self.layer_norm_func = self.layer_norm_fallback - - @classmethod - def layer_norm_residual(cls, vals, bias, res, gamma, beta, epsilon): - channels = gamma.shape[0] - dtype = gamma.dtype - vals_f = vals.to(torch.float32) - bias_f = bias.to(torch.float32).reshape(1, 1, -1) - res_f = res.to(torch.float32) - gamma_f = gamma.to(torch.float32) - beta_f = beta.to(torch.float32) - return F.layer_norm(vals_f + bias_f + res_f, (channels, ), weight=gamma_f, bias=beta_f, eps=epsilon).to(dtype) - - @classmethod - def layer_norm_residual_store_pre_ln_res(cls, vals, bias, res, gamma, beta, epsilon): - channels = gamma.shape[0] - dtype = gamma.dtype - vals_f = vals.to(torch.float32) - bias_f = bias.to(torch.float32).reshape(1, 1, -1) - res_f = res.to(torch.float32) - gamma_f = gamma.to(torch.float32) - beta_f = beta.to(torch.float32) - res_output = vals_f + bias_f + res_f - norm_output = F.layer_norm(res_output, (channels, ), weight=gamma_f, bias=beta_f, eps=epsilon).to(dtype) - return norm_output, res_output.to(dtype) - - @classmethod - def layer_norm_fallback(cls, vals, gamma, beta, epsilon): - channels = gamma.shape[0] - dtype = gamma.dtype - vals_f = vals.to(torch.float32) - gamma_f = gamma.to(torch.float32) - beta_f = beta.to(torch.float32) - return F.layer_norm(vals_f, (channels, ), weight=gamma_f, bias=beta_f, eps=epsilon).to(dtype) - - def forward(self, vals, gamma, beta, epsilon): - return self.layer_norm_func(vals, gamma, beta, epsilon) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/linear.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/linear.py deleted file mode 100644 index b8decb6dc5ea4333dd2ae117db352936de9b6c63..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/linear.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp -import deepspeed - - -class LinearOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig): - super(LinearOp, self).__init__(config) - try: - if self.config.dtype in [torch.float16, torch.int8]: - if deepspeed.HAS_TRITON and self.config.use_triton and self.config.dtype == torch.float16: - from deepspeed.ops.transformer.inference.triton.ops import linear_func as _triton_linear_func - self.linear_func = _triton_linear_func - triton_autotune = config.triton_autotune and config.layer_id == 0 - if triton_autotune: - __class__._triton_autotune(2, self.config.max_out_tokens, self.config.hidden_size) - else: - self.linear_func = self.inference_module.linear_layer_fp16 - self.linear_func = self.inference_module.linear_layer_fp16 - elif self.config.dtype == torch.bfloat16: - self.linear_func = self.inference_module.linear_layer_bf16 - else: - self.linear_func = self.inference_module.linear_layer_fp32 - except AttributeError: - self.linear_func = self.linear_fallback - - def linear_fallback(self, input, weight, bias, add_bias, do_flash_attn, num_heads, transpose, rope_theta): - raise NotImplementedError - - def forward(self, - input: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - add_bias: bool, - do_flash_attn: bool, - num_heads: int, - external_cache: bool = None, - num_layers: int = None): - qkv_out = self.linear_func(input, weight, bias, add_bias, do_flash_attn, num_heads, - self.config.transposed_mode, self.config.rope_theta) - return qkv_out - - @staticmethod - def _triton_autotune(min_seqlen, max_seqlen, hidden_size, dtype=torch.float16): - from deepspeed.ops.transformer.inference.triton.matmul_ext import Fp16Matmul, matmul - seqlen = [(min_seqlen + i) - for i in range(0, max_seqlen - min_seqlen + Fp16Matmul._cache_stride + 1, Fp16Matmul._cache_stride)] - Fp16Matmul._read_autotune_table() - for N in seqlen: - A = torch.randn((N, hidden_size), dtype=dtype, device='cuda') - B = torch.randn((hidden_size, 3 * hidden_size), dtype=dtype, device='cuda') - matmul(A, B) - Fp16Matmul._update_autotune_table() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/mlp_gemm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/mlp_gemm.py deleted file mode 100644 index 5f1f915ec02165b74ba5a8b7ec6dc238f9534d72..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/mlp_gemm.py +++ /dev/null @@ -1,124 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Optional - -import torch -import torch.nn.functional as F -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp -from deepspeed.utils.types import NormType -from .pre_rms_norm import PreRMSNormOp - - -class MLPGemmOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig): - super(MLPGemmOp, self).__init__(config) - try: - if self.config.norm_type == NormType.LayerNorm: - if self.config.dtype in [ - torch.float16, torch.int8 - ]: # non-triton cuda kernel has a higher performance in MLP than mlp_gemm_func in triton.ops - self.mlp_gemm_func = self.inference_module.mlp_gemm_fp16 # type: ignore - elif self.config.dtype == torch.bfloat16: - self.mlp_gemm_func = self.inference_module.mlp_gemm_bf16 - else: - self.mlp_gemm_func = self.inference_module.mlp_gemm_fp32 # type: ignore - elif self.config.norm_type == NormType.RMSNorm: - if self.config.dtype in [torch.float16, torch.int8]: - self.mlp_gemm_func = self.inference_module.rms_mlp_gemm_fp16 # type: ignore - elif self.config.dtype == torch.bfloat16: - self.mlp_gemm_func = self.inference_module.rms_mlp_gemm_bf16 - else: - self.mlp_gemm_func = self.inference_module.rms_mlp_gemm_fp32 # type: ignore - except AttributeError: - if self.config.norm_type == NormType.LayerNorm: - self.mlp_gemm_func = self.mlp_gemm_fallback - elif self.config.norm_type == NormType.RMSNorm: - self.mlp_gemm_func = self.rms_mlp_gemm_fallback - self.pre_rms_norm = PreRMSNormOp() - - def mlp_gemm_fallback(self, input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, - pre_layer_norm, mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, - transpose): - if mlp_after_attn: - residual_add = F.layer_norm(input + residual + input_bias, (input.shape[2], ), gamma, beta, eps) - tmp = torch.matmul(residual_add, weight_interm.t() if transpose else weight_interm) - tmp = F.gelu(tmp + bias) - output = torch.matmul(tmp, weight_out.t() if transpose else weight_out) - - return output, residual_add - else: - raise NotImplementedError - - def rms_mlp_gemm_fallback(self, input, residual, weight_interm, weight_out, gamma, eps, interm_scale, out_scale, - dtype, mlp_act_func_type, transpose): - inp_norm, residual = self.pre_rms_norm(input, residual, gamma, eps) - tmp = torch.matmul(inp_norm.view([-1, inp_norm.size(2)]), weight_interm.t() if transpose else weight_interm) - up_proj, gate_proj = tmp.chunk(2, dim=1) - - from deepspeed.utils.types import ActivationFuncType - if mlp_act_func_type == ActivationFuncType.GELU: - intermediate = F.gelu(gate_proj) - elif mlp_act_func_type == ActivationFuncType.ReLU: - intermediate = F.relu(gate_proj) - elif mlp_act_func_type == ActivationFuncType.GATED_GELU: - intermediate = F.gelu(gate_proj) - elif mlp_act_func_type == ActivationFuncType.GATED_SILU: - intermediate = F.silu(gate_proj) - else: - raise f"rms_mlp_gemm_fallback not implemented for activation type {mlp_act_func_type}" - - intermediate = intermediate * up_proj - - output = torch.matmul(intermediate, weight_out.t() if transpose else weight_out) - output = output.view([input.size(0), input.size(1), -1]) - - return [output, residual] - - def forward(self, - input: torch.Tensor, - residual: torch.Tensor, - weight_interm: torch.Tensor, - weight_out: torch.Tensor, - input_bias: Optional[torch.Tensor] = None, - bias: Optional[torch.Tensor] = None, - gamma: Optional[torch.Tensor] = None, - beta: Optional[torch.Tensor] = None): - if self.config.norm_type == NormType.LayerNorm: - output, residual_add = self.mlp_gemm_func( - input, - residual, - input_bias, - weight_interm, - weight_out, - bias, - gamma, - beta, - self.config.epsilon, - self.config.pre_layer_norm, - self.config.mlp_after_attn, - weight_interm.scale if hasattr(weight_interm, 'scale') else torch.empty(1), # type: ignore - weight_out.scale if hasattr(weight_out, 'scale') else torch.empty(1), # type: ignore - self.config.dtype == torch.int8, - self.config.mlp_act_func_type, - self.config.transposed_mode) - else: - if input_bias is not None: - input += input_bias - output, residual_add = self.mlp_gemm_func( - input, - residual, - weight_interm, - weight_out, - gamma, - self.config.epsilon, - weight_interm.scale if hasattr(weight_interm, 'scale') else torch.empty(1), # type: ignore - weight_out.scale if hasattr(weight_out, 'scale') else torch.empty(1), # type: ignore - self.config.dtype == torch.int8, - self.config.mlp_act_func_type, - self.config.transposed_mode) - return output, residual_add diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/moe_res_matmul.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/moe_res_matmul.py deleted file mode 100644 index ef3558c8bc889442159aa73eec22315d42481fb4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/moe_res_matmul.py +++ /dev/null @@ -1,29 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp - - -class MoEResMatmulOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig = None): - if config is None: - config = DeepSpeedInferenceConfig() - super(MoEResMatmulOp, self).__init__(config) - try: - self.moe_res_matmul_func = self.inference_module.moe_res_matmul - except AttributeError: - self.moe_res_matmul_func = self.moe_res_matmul_fallback - - @classmethod - def moe_res_matmul_fallback(cls, residual, coef, output): - coef_t = coef.transpose(1, 2).contiguous() - coef1, coef2 = torch.split(coef_t, split_size_or_sections=coef_t.shape[len(coef_t.shape) - 1] // 2, dim=-1) - return residual * coef1 + output * coef2 - - def forward(self, residual, coef, output): - return self.moe_res_matmul_func(residual, coef, output) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/pad_transform.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/pad_transform.py deleted file mode 100644 index 876fefc3bcfbd34e2f64a30703132649a96224ec..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/pad_transform.py +++ /dev/null @@ -1,26 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp - - -class PadTransformOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig = None): - if config is None: - config = DeepSpeedInferenceConfig() - super(PadTransformOp, self).__init__(config) - try: - self.pad_transform_func = self.inference_module.pad_transform_fp16 - except AttributeError: - self.pad_transform_func = self.pad_transform_fallback - - @staticmethod - def pad_transform_fallback(query, key, value, heads, do_flash_attn): - raise NotImplementedError("pad_transform fallback is not implemented.") - - def forward(self, query, key, value, heads, do_flash_attn): - return self.pad_transform_func(query, key, value, heads, do_flash_attn) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/pre_rms_norm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/pre_rms_norm.py deleted file mode 100644 index 7969d20f0527bd6fb1fce2f666bd5e5683296cb9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/pre_rms_norm.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp -from .rms_norm import RMSNormOp - - -class PreRMSNormOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig = None): - if config is None: - config = DeepSpeedInferenceConfig() - super(PreRMSNormOp, self).__init__(config) - try: - self.pre_rms_norm_func = self.inference_module.pre_rms_norm - except AttributeError: - self.pre_rms_norm_func = self.pre_rms_norm_fallback - - @staticmethod - def pre_rms_norm_fallback(vals, residual, gamma, epsilon): - residual = vals.to(torch.float32) + residual.to(torch.float32) - vals = residual - - return RMSNormOp.rms_norm_fallback(vals, gamma, epsilon), residual.to(gamma.dtype) - - def forward(self, vals, residual, gamma, epsilon): - return self.pre_rms_norm_func(vals, residual, gamma, epsilon) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/qkv_gemm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/qkv_gemm.py deleted file mode 100644 index 9ff5366fae5d6a267a44852f1aace7a98d3ad3dc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/qkv_gemm.py +++ /dev/null @@ -1,94 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import torch.nn.functional as F -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp -from .rms_norm import RMSNormOp -import deepspeed -from deepspeed.utils.types import NormType - - -class QKVGemmOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig): - super(QKVGemmOp, self).__init__(config) - try: - if self.config.norm_type == NormType.LayerNorm: - if self.config.dtype in [torch.float16, torch.int8]: - if deepspeed.HAS_TRITON and self.config.use_triton and self.config.dtype == torch.float16: - from deepspeed.ops.transformer.inference.triton.ops import qkv_gemm_func as _triton_qkv_gemm_func - self.qkv_gemm_func = _triton_qkv_gemm_func - triton_autotune = config.triton_autotune and config.layer_id == 0 - if triton_autotune: - __class__._triton_autotune(2, self.config.max_out_tokens, self.config.hidden_size) - else: - self.qkv_gemm_func = self.inference_module.qkv_gemm_fp16 # type: ignore - elif self.config.dtype == torch.bfloat16: - self.qkv_gemm_func = self.inference_module.qkv_gemm_bf16 - else: - self.qkv_gemm_func = self.inference_module.qkv_gemm_fp32 # type: ignore - elif self.config.norm_type == NormType.RMSNorm: - if self.config.dtype in [torch.float16, torch.int8]: - self.qkv_gemm_func = self.inference_module.rms_qkv_gemm_fp16 # type: ignore - elif self.config.dtype == torch.bfloat16: - self.qkv_gemm_func = self.inference_module.rms_qkv_gemm_bf16 - else: - self.qkv_gemm_func = self.inference_module.rms_qkv_gemm_fp32 # type: ignore - except AttributeError: - if self.config.norm_type == NormType.LayerNorm: - self.qkv_gemm_func = self.qkv_gemm_fallback - elif self.config.norm_type == NormType.RMSNorm: - self.qkv_gemm_func = self.rms_qkv_gemm_fallback - - @staticmethod - def _triton_autotune(min_seqlen, max_seqlen, hidden_size, dtype=torch.float16): - from deepspeed.ops.transformer.inference.triton.matmul_ext import Fp16Matmul, matmul - seqlen = [(min_seqlen + i) - for i in range(0, max_seqlen - min_seqlen + Fp16Matmul._cache_stride + 1, Fp16Matmul._cache_stride)] - Fp16Matmul._read_autotune_table() - for N in seqlen: - A = torch.randn((N, hidden_size), dtype=dtype, device='cuda') - B = torch.randn((hidden_size, 3 * hidden_size), dtype=dtype, device='cuda') - matmul(A, B) - Fp16Matmul._update_autotune_table() - - @staticmethod - def qkv_gemm_fallback(input, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose): - inp_norm = F.layer_norm(input, (input.shape[2], ), gamma, beta, eps) - tmp = torch.matmul(inp_norm, weight.t() if transpose else weight) - if add_bias: - tmp += bias - output = [tmp, inp_norm] - - return output - - @staticmethod - def rms_qkv_gemm_fallback(input, weight, q_scale, gamma, eps, q_int8, transpose): - inp_norm = RMSNormOp.rms_norm_fallback(input, gamma, eps) - tmp = torch.matmul(inp_norm, weight.t() if transpose else weight) - output = [tmp, inp_norm] - - return output - - def forward(self, input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, gamma: torch.Tensor, - beta: torch.Tensor): - - add_bias = bias is not None - bias = bias if add_bias else torch.empty(1) # type: ignore - q_scale = weight.scale if hasattr(weight, 'scale') else torch.empty(1) # type: ignore - q_int8 = self.config.dtype == torch.int8 - - if self.config.norm_type == NormType.LayerNorm: - output, norm = self.qkv_gemm_func(input, weight, q_scale, bias, gamma, beta, self.config.epsilon, add_bias, - q_int8, self.config.transposed_mode) - else: - output, norm = self.qkv_gemm_func(input, weight, q_scale, gamma, self.config.epsilon, q_int8, - self.config.transposed_mode) - if add_bias: - output += bias - - return output, norm diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/residual_add.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/residual_add.py deleted file mode 100644 index 93b229c5d1ac0165f82ccc322769f11dd61d890e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/residual_add.py +++ /dev/null @@ -1,73 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from typing import Optional - -from .vector_add import VectorAddOp -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp - - -class ResidualAddOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig): - super(ResidualAddOp, self).__init__(config) - try: - if self.config.dtype in [torch.float16, torch.int8]: - self.residual_add_func = self.inference_module.residual_add_bias_fp16 - elif self.config.dtype == torch.bfloat16: - self.residual_add_func = self.inference_module.residual_add_bias_bf16 - else: - self.residual_add_func = self.inference_module.residual_add_bias_fp32 - except AttributeError: - self.residual_add_func = self.residual_add_fallback - self.vector_add = VectorAddOp() - - @staticmethod - def res_add_bias(hidden_state, residual, attn_output, attn_bias, final_bias, add_attn_bias, mp_size): - hidden_state += attn_output + (residual + final_bias) / mp_size - if add_attn_bias: - hidden_state += attn_bias / mp_size - - return hidden_state - - @staticmethod - def residual_add_fallback(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size, - mlp_after_attn, add_bias, pre_layer_norm): - if mlp_after_attn: - if pre_layer_norm: - tmp = (residual.float() + attention_output.float() + attention_bias.float() + - final_bias.float()) / mp_size + hidden_state.float() - else: - tmp = residual.float() + hidden_state.float() + final_bias.float() - else: - tmp = ResidualAddOp.res_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, - add_bias, mp_size) - residual.copy_(tmp.to(hidden_state.dtype)) - - return residual - - def forward(self, - hidden_state: torch.Tensor, - residual: torch.Tensor, - add_bias: bool, - attention_output: Optional[torch.Tensor] = None, - residual_add: Optional[torch.Tensor] = None, - attention_bias: Optional[torch.Tensor] = None, - final_bias: Optional[torch.Tensor] = None): - - if final_bias is None and attention_bias is None: - residual = self.vector_add(residual + attention_output, hidden_state, 1.0 / self.config.mp_size) - else: - if not self.config.pre_layer_norm and residual_add is not None: - # only use residual add if its set and we are not pre layer norm - residual = residual_add - - self.residual_add_func(hidden_state, residual, attention_output, attention_bias, final_bias, - self.config.mp_size, self.config.mlp_after_attn, add_bias, - self.config.pre_layer_norm) - - return residual diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/rms_norm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/rms_norm.py deleted file mode 100644 index 128883ce5d438b03718496abe1e26134d9e7f440..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/rms_norm.py +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp - - -class RMSNormOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig = None): - if config is None: - config = DeepSpeedInferenceConfig() - super(RMSNormOp, self).__init__(config) - try: - self.rms_norm_func = self.inference_module.rms_norm - except AttributeError: - self.rms_norm_func = self.rms_norm_fallback - - @staticmethod - def rms_norm_fallback(vals, gamma, epsilon): - variance = vals.to(torch.float32).pow(2).mean(-1, keepdim=True) - vals = vals * torch.rsqrt(variance + epsilon) - - if gamma.dtype in [torch.float16, torch.bfloat16]: - vals = vals.to(gamma.dtype) - - return gamma * vals - - def forward(self, vals, gamma, epsilon): - return self.rms_norm_func(vals, gamma, epsilon) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/softmax.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/softmax.py deleted file mode 100644 index 2e08541596fa8dc89f8ac893254dc773ed5ee284..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/softmax.py +++ /dev/null @@ -1,71 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import torch.nn.functional as F -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp -from deepspeed.ops.transformer.inference.op_binding.workspace import InferenceContext - - -class SoftmaxOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig): - super(SoftmaxOp, self).__init__(config) - self.num_attention_heads_per_partition = config.heads // config.mp_size - try: - if self.config.dtype in [torch.float16, torch.int8]: - self.softmax_func = self.inference_module.softmax_fp16 - elif self.config.dtype == torch.bfloat16: - self.softmax_func = self.inference_module.softmax_bf16 - else: - self.softmax_func = self.inference_module.softmax_fp32 - except AttributeError: - self.softmax_func = self.softmax_fallback - - @staticmethod - def softmax_fallback(attn_scores, attn_mask, alibi, triangular, recompute, local_attention, window_size, async_op, - layer_scale, head_offset, mp_size): - scores_len = len(attn_scores.size()) - heads = 1 - if scores_len > 1: - heads = attn_scores.size()[1] - num_attention_heads_per_partition = heads // mp_size - - if alibi is not None: - if len(alibi.shape) == 1: - alibi = None - else: - alibi = alibi[head_offset:head_offset + num_attention_heads_per_partition] - if attn_mask is not None and len(attn_mask.shape) == 1: - attn_mask = None - input_dtype = attn_scores.dtype - attn_scores *= layer_scale - - if alibi is not None: - attn_scores += alibi - if attn_mask is not None: - # expand atten_mask from two dim into 4 dim, insert two dims in the middle - if len(attn_mask.shape) == 2: - attn_mask = attn_mask[:, None, None, :] - attn_scores += attn_mask - if triangular: - if attn_scores.shape[2] == 1: # query using kv cache - token_idx = InferenceContext.Instance().current_tokens() - tri = torch.arange(attn_scores.shape[2], device=attn_scores.device).ge(token_idx) - else: - tri = ~torch.tril(torch.ones(attn_scores.size(), device=attn_scores.device)).to(bool) - attn_scores = torch.masked_fill(attn_scores, tri, float('-inf')) - output = F.softmax(attn_scores, dim=-1, dtype=torch.float32).to(input_dtype) - - return output - - def forward(self, attn_scores: torch.Tensor, attn_mask: torch.Tensor, alibi: torch.Tensor, triangular: bool, - recompute: bool, local_attention: bool, window_size: int, async_op: bool, layer_scale: float, - head_offset: int): - output = self.softmax_func(attn_scores, attn_mask, alibi, triangular, recompute, local_attention, window_size, - async_op, layer_scale, head_offset, self.config.mp_size) - - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/softmax_context.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/softmax_context.py deleted file mode 100644 index d745df678e93bf3e359288610156c7cb8392e2b0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/softmax_context.py +++ /dev/null @@ -1,144 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed import comm as dist -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp -from .softmax import SoftmaxOp -from deepspeed.ops.transformer.inference.op_binding.workspace import InferenceContext - - -class SoftmaxContextOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig): - super(SoftmaxContextOp, self).__init__(config) - try: - if self.config.dtype in [torch.float16, torch.int8]: - self.softmax_context_func = self.inference_module.softmax_context_fp16 - elif self.config.dtype == torch.bfloat16: - self.softmax_context_func = self.inference_module.softmax_context_bf16 - else: - self.softmax_context_func = self.inference_module.softmax_context_fp32 - except AttributeError: - self.softmax_context_func = self.softmax_context_fallback - - @staticmethod - def transform4d_0213(x, seq_length): - assert x.dim() == 3, F"Dim {x.dim()} is not supported" - batch_size, num_heads, seq_length_head_dim = x.shape - head_dim = seq_length_head_dim // seq_length - x = x.view(batch_size, num_heads, seq_length, head_dim) - x = x.permute(0, 2, 1, 3) - - return x - - @staticmethod - def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep <= 1 or num_key_value_heads == 1: - return hidden_states - - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - @staticmethod - def bias_add_transform_0213(input, bias, num_heads, trans_count, perform_bias=False): - assert trans_count == 1 or trans_count == 3, F"Trans count {trans_count} is not supported" - assert input.dim() == 3, F"Dim {input.dim()} is not supported" - input_biased = torch.add(input, bias) if perform_bias else input - batch_size, seq_length, value_size = input_biased.shape - hid_dim = value_size // trans_count - head_dim = hid_dim // num_heads - - if trans_count == 1: - query_layer = input.view(batch_size, seq_length, num_heads, head_dim) - query_layer = query_layer.permute(0, 2, 1, 3) - key_layer = torch.zeros_like(query_layer) - value_layer = torch.zeros_like(query_layer) - return query_layer, key_layer, value_layer - - qkv_layers = input.view(batch_size, seq_length, 3, num_heads, head_dim) - query_layer, key_layer, value_layer = qkv_layers[..., 0, :, :], qkv_layers[..., 1, :, :], qkv_layers[..., - 2, :, :] - query_layer = query_layer.transpose(1, 2) - key_layer = key_layer.transpose(1, 2) - value_layer = value_layer.transpose(1, 2) - - return query_layer, key_layer, value_layer - - def softmax_context_fallback(self, query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads, - num_kv, norm_factor, triangular_masking, local_attention, window_size, no_masking, - layer_id, num_layers, alibi, rope_theta, is_prompt, token_idx, position_ids): - bat_0213_query, bat_0213_key, bat_0213_value = self.bias_add_transform_0213( - query_key_value, None, heads, 3, False) - - if rotary_dim > 0 and rotate_half: - from transformers.models.llama.modeling_llama import apply_rotary_pos_emb - - rotary = InferenceContext.Instance().get_rotary(rotary_dim, rope_theta, bat_0213_value.device) - cos, sin = rotary(bat_0213_value, InferenceContext.Instance().get_max_tokens_num()) - bat_0213_query, bat_0213_key = apply_rotary_pos_emb(bat_0213_query, bat_0213_key, cos, sin, position_ids) - - bat_0213_key, bat_0213_value = InferenceContext.Instance().update_cache(layer_id, token_idx, is_prompt, - bat_0213_key, bat_0213_value) - - bat_0213_key = self.repeat_kv(bat_0213_key, num_kv) - bat_0213_value = self.repeat_kv(bat_0213_value, num_kv) - - bsz = query_key_value.shape[0] - head_dim = query_key_value.shape[2] // (heads * 3) - - bmm_output = torch.bmm(bat_0213_query.reshape(bsz * heads, bat_0213_query.shape[2], head_dim), - bat_0213_key.reshape(bsz * heads, bat_0213_key.shape[2], head_dim).transpose(1, 2)) - - layer_scale = 1.0 - if alibi is not None and len(alibi.shape) > 1: - layer_scale = max(1, layer_id).to(float) - - alpha = norm_factor * norm_factor / layer_scale - bmm_output *= alpha - bmm_output_reshape = bmm_output.reshape(bsz, heads, bmm_output.shape[1], bmm_output.shape[2]) - - recompute = is_prompt - if attn_mask is not None and len(attn_mask.shape) > 1 and attn_mask.shape[-1] < bmm_output_reshape.shape[3]: - attn_mask = torch.nn.functional.pad(attn_mask, (0, bmm_output_reshape.shape[3] - attn_mask.shape[-1]), - value=torch.finfo(attn_mask.dtype).min) - softmax_output = SoftmaxOp.softmax_fallback(bmm_output_reshape, attn_mask, alibi, triangular_masking, - recompute, local_attention, window_size, None, layer_scale, 0, 1) - - output = torch.bmm(softmax_output.reshape(bsz * heads, softmax_output.shape[2], softmax_output.shape[3]), - bat_0213_value.reshape(bsz * heads, bat_0213_value.shape[2], head_dim)) - - output = output.reshape(bsz, heads, output.shape[1], head_dim) - output = output.reshape(bsz, heads, output.shape[2] * head_dim) - input_seq_len = query_key_value.shape[1] - t4d_0123_output = self.transform4d_0213(output, input_seq_len) - t4d_0123_output = t4d_0123_output.reshape(bsz, t4d_0123_output.shape[1], heads * head_dim) - - if layer_id == num_layers - 1: - InferenceContext.Instance().advance_tokens() - - return t4d_0123_output, bat_0213_key, bat_0213_value - - def forward(self, query_key_value: torch.Tensor, attn_mask: torch.Tensor, heads: int, num_kv: int, - norm_factor: float, no_masking: bool, layer_id: int, num_layers: int, alibi: torch.Tensor, - is_prompt: bool, token_idx: torch.Tensor, position_ids: torch.Tensor): - - if alibi is not None: - batch_heads = query_key_value.shape[0] * heads - offset = dist.get_rank() * batch_heads if dist.is_initialized() else 0 - alibi = alibi[offset:batch_heads + offset, :, :] - else: - alibi = torch.empty(1) - - output = self.softmax_context_func(query_key_value, attn_mask, self.config.rotary_dim, self.config.rotate_half, - self.config.rotate_every_two, heads, num_kv, norm_factor, - self.config.triangular_masking, self.config.local_attention, - self.config.window_size, no_masking, layer_id, num_layers, alibi, - self.config.rope_theta, is_prompt, token_idx, position_ids) - - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/vector_add.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/vector_add.py deleted file mode 100644 index 015340a1084b5a64557737928462c5a1a79ae2ba..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/vector_add.py +++ /dev/null @@ -1,28 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp - - -class VectorAddOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig = None): - if config is None: - config = DeepSpeedInferenceConfig() - super(VectorAddOp, self).__init__(config) - try: - self.vector_add_func = self.inference_module._vector_add - except AttributeError: - self.vector_add_func = self.vector_add_fallback - - @classmethod - def vector_add_fallback(cls, a, b, gamma): - """Based on csrc/transformer/inference/csrc/pt_binding.cpp code of _vector_add""" - dtype = a.dtype - return (gamma * a.float() + b.float()).to(dtype) - - def forward(self, a, b, gamma): - return self.vector_add_func(a, b, gamma) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/vector_matmul.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/vector_matmul.py deleted file mode 100644 index cabab8d8c4ab146bde0f7d0d69a9056a299694c1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/vector_matmul.py +++ /dev/null @@ -1,54 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp -import deepspeed - - -class VectorMatMulOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig): - super(VectorMatMulOp, self).__init__(config) - try: - if self.config.dtype == torch.float16: - if deepspeed.HAS_TRITON and config.use_triton: - from deepspeed.ops.transformer.inference.triton.ops import vector_matmul_func as _triton_vector_matmul_func - self.vector_matmul_func = _triton_vector_matmul_func - triton_autotune = config.triton_autotune and config.layer_id == 0 - if triton_autotune: - __class__._triton_autotune(2, self.config.max_out_tokens, self.config.hidden_size) - else: - self.vector_matmul_func = self.inference_module.vector_matmul_fp16 - elif self.config.dtype == torch.int8: - self.vector_matmul_func = self.inference_module.vector_matmul_int8 - elif self.config.dtype == torch.bfloat16: - self.vector_matmul_func = self.inference_module.vector_matmul_bf16 - else: - self.vector_matmul_func = self.inference_module.vector_matmul_fp32 - except AttributeError: - self.vector_matmul_func = self.vector_matmul_fallback - - def vector_matmul_fallback(self, input, weight, async_op, q_scale, q_int8, transpose): - return torch.matmul(input, weight.t() if transpose else weight) - - def forward(self, input: torch.Tensor, weight: torch.Tensor, async_op: bool = False): - q_scale = weight.scale if hasattr(weight, 'scale') else torch.empty(1) - q_int8 = self.config.dtype == torch.int8 - output = self.vector_matmul_func(input, weight, async_op, q_scale, q_int8, self.config.transposed_mode) - return output - - @staticmethod - def _triton_autotune(min_seqlen, max_seqlen, hidden_size, dtype=torch.float16): - from deepspeed.ops.transformer.inference.triton.matmul_ext import Fp16Matmul, matmul - seqlen = [(min_seqlen + i) - for i in range(0, max_seqlen - min_seqlen + Fp16Matmul._cache_stride + 1, Fp16Matmul._cache_stride)] - Fp16Matmul._read_autotune_table() - for N in seqlen: - A = torch.randn((N, hidden_size), dtype=dtype, device='cuda') - B = torch.randn((hidden_size, hidden_size), dtype=dtype, device='cuda') - matmul(A, B) - Fp16Matmul._update_autotune_table() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/workspace.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/workspace.py deleted file mode 100644 index 19de7d9576afe44bfbf2e9b4c1f24a842a05ebbf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/workspace.py +++ /dev/null @@ -1,222 +0,0 @@ -# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from ..config import DeepSpeedInferenceConfig -from .base import BaseOp - -minus_inf = -10000.0 -key_idx = 0 -value_idx = 1 - - -class InferenceContext: - - __instance = None - - def __init__(self): - self.kv_cache = None - self.kv_cache_elem_dtype = None - self.num_tokens = 1 - self.kv_cache_num_layers = None - self.kv_cache_size = None - self.max_out_tokens = None - self.rotary = None - self.allocate_called = False - self.static_shapes = True - - @classmethod - def Instance(cls): - if InferenceContext.__instance is None: - InferenceContext.__instance = InferenceContext() - return InferenceContext.__instance - - def gen_workspace(self, num_layers, num_heads, batch_size, prompt_len, hidden_dim, mp_size, external_cache, - elem_dtype, rank, max_out_tokens, min_out_tokens): - self.allocate_called = True - self.kv_cache = None - if not external_cache: - self.kv_cache_num_layers = num_layers - self.max_out_tokens = max_out_tokens - head_size = hidden_dim // num_heads - self.kv_cache_size = torch.Size([batch_size, (num_heads // mp_size), max_out_tokens, head_size]) - self.kv_cache_elem_dtype = elem_dtype - self.num_tokens = 0 - self.static_shapes = True - return True - - def retake_workspace(self): - return True - - def _retake_workspace(self): - assert self.allocate_called, "retake workspace called before allocate workspace" - - import deepspeed.accelerator as accelerator - if self.kv_cache is None: - self.kv_cache = [] - for layer in range(self.kv_cache_num_layers): - self.kv_cache.append((torch.zeros(self.kv_cache_size, - dtype=self.kv_cache_elem_dtype, - device=accelerator.get_accelerator().device_name()), - torch.zeros(self.kv_cache_size, - dtype=self.kv_cache_elem_dtype, - device=accelerator.get_accelerator().device_name()))) - - return True - - def update_cache(self, layer_id, token_idx, is_prompt, bat_0213_key, bat_0213_value): - has_workspace = self._retake_workspace() - assert has_workspace, "Could not allocate workspace" - - # Update current token - if is_prompt: - self.static_shapes = True - if token_idx is None: - self.static_shapes = False - InferenceContext.Instance().reset_tokens(bat_0213_key.shape[2]) - else: - InferenceContext.Instance().reset_tokens(token_idx) - - if token_idx is None: - token_idx = InferenceContext.Instance().current_tokens() - - bsz = bat_0213_key.shape[0] - - # Update cache content - if is_prompt: - cache_max_seq = self.kv_cache_size[2] - cache_max_head_dim = self.kv_cache_size[3] - seq = bat_0213_key.shape[2] - - mask = torch.arange(cache_max_seq, device=bat_0213_key.device) - mask = mask.ge(token_idx) - mask = mask.unsqueeze(-1) - mask = mask.expand([cache_max_seq, cache_max_head_dim]) - - self.kv_cache[layer_id][key_idx][:bsz, :, :seq, :].copy_(bat_0213_key) - self.kv_cache[layer_id][key_idx][:bsz, :].masked_fill_(mask, 0) - self.kv_cache[layer_id][value_idx][:bsz, :, :seq, :].copy_(bat_0213_value) - self.kv_cache[layer_id][value_idx][:bsz, :].masked_fill_(mask, 0) - else: - if self.static_shapes: - assert type(token_idx) == torch.Tensor, "token_idx is expected to be torch.Tensor" - self.kv_cache[layer_id][key_idx][:bsz].index_copy_(2, token_idx - 1, bat_0213_key) - self.kv_cache[layer_id][value_idx][:bsz].index_copy_(2, token_idx - 1, bat_0213_value) - else: - assert type(token_idx) == int, "token_idx is expected to be int" - self.kv_cache[layer_id][key_idx][:bsz, :, token_idx - 1:token_idx, :] = bat_0213_key - self.kv_cache[layer_id][value_idx][:bsz, :, token_idx - 1:token_idx, :] = bat_0213_value - - bat_0213_key = self.kv_cache[layer_id][key_idx][:bsz] - bat_0213_value = self.kv_cache[layer_id][value_idx][:bsz] - - if not self.static_shapes: - bat_0213_key = bat_0213_key[:, :, :token_idx, :] - bat_0213_value = bat_0213_value[:, :, :token_idx, :] - - return bat_0213_key, bat_0213_value - - def release_workspace(self): - self.kv_cache = None - self.rotary = None - - def reset_tokens(self, initial_tokens=1): - self.num_tokens = initial_tokens - - def current_tokens(self): - return self.num_tokens - - def advance_tokens(self): - self.num_tokens = self.num_tokens + 1 - - def get_kv_cache(self): - return self.kv_cache - - def get_rotary(self, rotary_dim, rope_theta, device=None): - if self.rotary is None: - from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding - - self.rotary = LlamaRotaryEmbedding(rotary_dim, base=rope_theta, device=device) - - return self.rotary - - def get_max_tokens_num(self): - return self.max_out_tokens - - -class WorkspaceOp(BaseOp): - - def __init__(self, config: DeepSpeedInferenceConfig = None): - if config is None: - config = DeepSpeedInferenceConfig() - self.inference_context = InferenceContext.Instance() - self._is_allocated = False - try: - super(WorkspaceOp, self).__init__(config) - if config.dtype == torch.float32: - self.allocate_workspace_func = self.inference_module.allocate_workspace_fp32 - elif config.dtype == torch.bfloat16: - self.allocate_workspace_func = self.inference_module.allocate_workspace_bf16 - else: - self.allocate_workspace_func = self.inference_module.allocate_workspace_fp16 - self.release_workspace_func = self.inference_module.release_workspace - self.retake_workspace_func = self.inference_module.retake_workspace - self.reset_cache_func = self.inference_module.reset_cache - except (ValueError, AttributeError) as e: - print(f"Using fallback functions in workspace because of {e}") - if config.dtype == torch.float32: - self.allocate_workspace_func = self.allocate_workspace_fp32_fallback - elif config.dtype == torch.bfloat16: - self.allocate_workspace_func = self.allocate_workspace_bf16_fallback - else: - self.allocate_workspace_func = self.allocate_workspace_fp16_fallback - self.release_workspace_func = self.release_workspace_fallback - self.retake_workspace_func = self.retake_workspace_fallback - self.reset_cache_func = self.reset_cache_fallback - - def allocate_workspace(self, *args, **kwargs): - self._is_allocated = True - return self.allocate_workspace_func(*args, **kwargs) - - def release_workspace(self): - self._is_allocated = False - return self.release_workspace_func() - - def reset_cache(self): - return self.reset_cache_func() if self.reset_cache_func else None - - def retake_workspace(self): - return self.retake_workspace_func() if self.retake_workspace_func else None - - def allocate_workspace_fp32_fallback(self, hidden_dim, num_heads, prompt_length, batch_size, num_layers, mp_size, - external_cache, rank, max_out_tokens, min_out_tokens): - return self.inference_context.gen_workspace(num_layers, num_heads, batch_size, prompt_length, hidden_dim, - mp_size, external_cache, torch.float, rank, max_out_tokens, - min_out_tokens) - - def allocate_workspace_bf16_fallback(self, hidden_dim, num_heads, prompt_length, batch_size, num_layers, mp_size, - external_cache, rank, max_out_tokens, min_out_tokens): - return self.inference_context.gen_workspace(num_layers, num_heads, batch_size, prompt_length, hidden_dim, - mp_size, external_cache, torch.bfloat16, rank, max_out_tokens, - min_out_tokens) - - def allocate_workspace_fp16_fallback(self, hidden_dim, num_heads, prompt_length, batch_size, num_layers, mp_size, - external_cache, rank, max_out_tokens, min_out_tokens): - return self.inference_context.gen_workspace(num_layers, num_heads, batch_size, prompt_length, hidden_dim, - mp_size, external_cache, torch.half, rank, max_out_tokens, - min_out_tokens) - - def reset_cache_fallback(self): - return self.inference_context.reset_tokens() - - def release_workspace_fallback(self): - return self.inference_context.release_workspace() - - def retake_workspace_fallback(self): - return self.inference_context.retake_workspace() - - def is_allocated(self): - return self._is_allocated diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__init__.py deleted file mode 100644 index b7d1968df62a99849992c8b1e93698d9e51cec30..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .residual_add import residual_add_bias -from .layer_norm import layer_norm, layer_norm_residual -from .gelu import gelu -from .softmax import softmax -from .ops import * -from .matmul_ext import fp16_matmul, matmul_4d, score_4d_matmul, context_4d_matmul diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 1e6a0176bfe1be058db64f53eae02a9382752b2d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/attention.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/attention.cpython-310.pyc deleted file mode 100644 index e52b580672056b6eee52fcfc2f4a2faa0c6e3328..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/attention.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/gelu.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/gelu.cpython-310.pyc deleted file mode 100644 index 57812e209d38a9a88276c453bdc6fc0b37cfc82d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/gelu.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/layer_norm.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/layer_norm.cpython-310.pyc deleted file mode 100644 index f72d3f22e65db520fa6b9e99a51e0f13b347d739..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/layer_norm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/matmul_ext.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/matmul_ext.cpython-310.pyc deleted file mode 100644 index 882737cb3d1e647047fb379983a0f1869ea277ad..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/matmul_ext.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/mlp.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/mlp.cpython-310.pyc deleted file mode 100644 index f826ca6c54b90878f90cbb4dab4b2459092d04ba..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/mlp.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/ops.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/ops.cpython-310.pyc deleted file mode 100644 index 0153896a84a9fdea6257af8fa0dc8dc26d2df89d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/ops.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/residual_add.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/residual_add.cpython-310.pyc deleted file mode 100644 index bfacdec3c8db9f24f8f16441c79ae40491c32a7e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/residual_add.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/softmax.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/softmax.cpython-310.pyc deleted file mode 100644 index 10c6794cd9703fc2caf9a29ac11ec0b935bbd207..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/softmax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/triton_matmul_kernel.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/triton_matmul_kernel.cpython-310.pyc deleted file mode 100644 index 07988c76d8954d9ef3c2dc716f20d05554e5cd39..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/__pycache__/triton_matmul_kernel.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/attention.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/attention.py deleted file mode 100644 index 6845d91b06be91ffe3e0eedb7f59fd79609c9f92..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/attention.py +++ /dev/null @@ -1,397 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import math -import torch -import torch.nn as nn -import triton -import triton.language as tl -from deepspeed.accelerator import get_accelerator -from deepspeed import comm as dist -from deepspeed.ops.transformer.inference.op_binding import LinearOp, VectorMatMulOp, SoftmaxContextOp, QKVGemmOp -from deepspeed.ops.transformer.inference.triton import ( - softmax, - score_4d_matmul, - context_4d_matmul, -) - -minus_inf = -10000.0 - - -class TritonSelfAttention(nn.Module): - num_layers = 0 - - def __init__(self, config, mp_group=None, q_scales=None, q_groups=1, merge_count=1, qkv_merging=False): - super(TritonSelfAttention, self).__init__() - self.config = config - data_type = self.config.dtype - data_type_fp = torch.half if self.config.dtype == torch.int8 else self.config.dtype - assert data_type_fp == torch.half, "triton supports fp16 data_type_fp" - - self.config.layer_id = TritonSelfAttention.num_layers - TritonSelfAttention.num_layers = TritonSelfAttention.num_layers + 1 - device = get_accelerator().current_device_name() #if config.bigscience_bloom else 'cpu' - - assert config.mp_size == 1, "mp_size has to be 1 with triton attention yet" - if self.config.set_empty_params: - self.attn_qw = None - self.attn_qb = None - self.attn_kw = None - self.attn_kb = None - self.attn_vw = None - self.attn_vb = None - self.attn_qkvw = None - self.attn_qkvb = None - self.attn_ow = None - self.attn_ob = None - else: - qkv_size_per_partition = (self.config.hidden_size // self.config.mp_size) * 3 - self.attn_qkvw = nn.Parameter(torch.empty(self.config.hidden_size, - qkv_size_per_partition, - dtype=data_type, - device=device), - requires_grad=False) - self.attn_qkvb = nn.Parameter(torch.empty(qkv_size_per_partition, dtype=data_type_fp, device=device), - requires_grad=False) - # self-ouput weights - out_size_per_partition = self.config.hidden_size // self.config.mp_size - self.attn_ow = nn.Parameter(torch.empty(out_size_per_partition, - self.config.hidden_size, - dtype=data_type, - device=device), - requires_grad=False) - - self.attn_ob = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), - requires_grad=False) - - self.num_attention_heads_per_partition = self.config.heads // self.config.mp_size - self.hidden_size_per_partition = self.config.hidden_size // self.config.mp_size - self.hidden_size_per_attention_head = self.config.hidden_size // self.config.heads - - self.mp_group = mp_group - self.use_flash = False - # triton flash attention is enabled when the compute capability >= 8.0 - if get_accelerator().is_triton_supported(): - self.use_flash = True - - # used for quantization - self.q_scales = q_scales - self.q_groups = q_groups - self.merge_count = int(math.log2(merge_count)) - - self.norm_factor = math.sqrt(self.config.hidden_size // self.config.heads) - if not config.use_mup: - self.norm_factor = math.sqrt(self.norm_factor) - - if self.config.scale_attn_by_inverse_layer_idx is True: - self.norm_factor *= math.sqrt(self.config.layer_id + 1) - # https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/gpt2/modeling_gpt2.py#L191 - - triton_autotune = self.config.triton_autotune and self.config.layer_id == 0 - self.qkv_func = QKVGemmOp(config) - self.score_context_func = SoftmaxContextOp(config) - self.linear_func = LinearOp(config) - self.vector_matmul_func = VectorMatMulOp(config) - - self.hidden_size = config.hidden_size - self.head_size = config.hidden_size // config.heads - self.scale = (1 / self.norm_factor / self.norm_factor if self.config.scale_attention else 1.0 - ) # making it back to 1/sqrt(head_size) - self.triangular_masking = self.config.triangular_masking - - # triton autotune table update for score/context matmul - if triton_autotune: - print(f"running triton autotune for regular attention kernel") - __class__._triton_autotune(2, self.config.max_out_tokens, self.head_size, self.config.hidden_size, - self.triangular_masking, self.scale) - - @staticmethod - def _triton_autotune(min_seqlen, - max_seqlen, - head_size, - hidden_size, - triangular_masking, - scale, - dtype=torch.float16): - from deepspeed.ops.transformer.inference.triton.matmul_ext import Fp16Matmul, score_4d_matmul, context_4d_matmul - seqlen = [(min_seqlen + i) - for i in range(0, max_seqlen - min_seqlen + Fp16Matmul._cache_stride + 1, Fp16Matmul._cache_stride)] - Fp16Matmul._read_autotune_table() - for N in seqlen: - qkv = torch.randn((1, N, 3 * hidden_size), dtype=dtype, device='cuda') - output = score_4d_matmul(qkv, head_size, triangular_masking, scale) - context_4d_matmul(output, qkv, head_size) - Fp16Matmul._update_autotune_table() - - def ds_compute_attention(self, qkv_out, input_mask, layer_past, alibi, is_prompt, token_idx, position_ids): - if isinstance(qkv_out, list): - qkv_out = qkv_out[0] - - no_masking = input_mask is None - - if no_masking: - input_mask = torch.empty(1) - - attn_key_value = self.score_context_func( - query_key_value=qkv_out, - attn_mask=((1 - input_mask).to(qkv_out.dtype) * - minus_inf) if input_mask.dtype == torch.int64 else input_mask, - heads=self.num_attention_heads_per_partition, - norm_factor=(1 / self.norm_factor if self.config.scale_attention else 1.0), - no_masking=no_masking, - layer_id=self.config.layer_id, - num_layers=TritonSelfAttention.num_layers, - alibi=alibi, - is_prompt=is_prompt, - token_idx=token_idx, - position_ids=position_ids) - - context_layer, key_layer, value_layer = attn_key_value - return context_layer, key_layer, value_layer - - def forward( - self, - input, - input_mask, - head_mask=None, - layer_past=None, - get_present=False, # not used - encoder_hidden_states=None, # not used - encoder_attention_mask=None, # not used - triangularutput_attentions=False, # not used - norm_w=None, - norm_b=None, - alibi=None, - use_triton_attention=True, - **kwargs): - - if not self.config.pre_layer_norm: - qkv_out = self.linear_func(input=input, - weight=self.attn_qkvw, - bias=self.attn_qkvb, - add_bias=self.attn_qkvb is not None, - do_flash_attn=False, - num_heads=self.num_attention_heads_per_partition, - num_layers=TritonSelfAttention.num_layers) - qkv = qkv_out - else: - qkv_out = self.qkv_func(input=input, - weight=self.attn_qkvw, - bias=(self.attn_qkvb if self.attn_qkvb is not None else norm_b), - gamma=norm_w, - beta=norm_b) - qkv = qkv_out[0] - - if use_triton_attention and (alibi is None): - context_layer = _triton_attention(qkv=qkv, - input_mask=input_mask, - scale=self.scale, - layer_past=layer_past, - alibi=alibi, - head_size=self.head_size, - use_triton_flash=self.use_flash, - use_cuda_flash=False, - triangular=self.triangular_masking) - key_layer, value_layer = qkv[:, :, self.hidden_size:2 * self.hidden_size], qkv[:, :, 2 * self.hidden_size:] - else: - is_prompt = kwargs.get("first_token", qkv_out[0].shape[1] > 1) - token_idx = kwargs.get("token_idx", None) - position_ids = kwargs.get("position_ids", None) - context_layer, key_layer, value_layer = self.ds_compute_attention(qkv_out=qkv_out, - input_mask=input_mask, - layer_past=layer_past, - alibi=alibi, - is_prompt=is_prompt, - toke_idx=token_idx, - position_ids=position_ids) - output = self.vector_matmul_func(input=context_layer, weight=self.attn_ow) - - inp_norm = qkv_out[-1] - - if self.config.mlp_after_attn and self.mp_group is not None and dist.get_world_size(group=self.mp_group) > 1: - dist.all_reduce(output, group=self.mp_group) - - return (output, key_layer, value_layer, context_layer, inp_norm) - - -global inference_module - - -def _triton_attention(qkv, - input_mask, - layer_past, - alibi, - scale, - head_size, - triangular=False, - use_cuda_flash=False, - use_triton_flash=False, - use_ds_attention=False): - if isinstance(qkv, list): - qkv = qkv[0] - - assert alibi is None, "layer_past not supported in alibi yet" - - if use_triton_flash: - output = _triton_packed_flash(qkv, - head_size, - input_mask, - scale, - causal=triangular, - add_mask=(not triangular and input_mask is not None)) - else: - output = score_4d_matmul(qkv, head_size, triangular, scale) - if triangular: - output = softmax(output) - else: - output = softmax(output, input_mask) - output = context_4d_matmul(output, qkv, head_size) - - return output - - -''' -flash attention 2 -modified the triton kernel in -https://github.com/openai/triton/blob/08c16589573621fcb8cd5a9c3b8a0537077f876d/python/tutorials/06-fused-attention.py -''' - - -@triton.jit -def _flash_packed_kernel( - QKV, - mask, - ADD_MASK: tl.constexpr, - IS_CAUSAL: tl.constexpr, - sm_scale, - Out, - stride_qz, - stride_qn, - stride_qm, - stride_mz, - stride_oz, - stride_on, - Z, - H, - N_CTX, - P_SEQ, - hidden_size, - BLOCK_M: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - BLOCK_N: tl.constexpr, -): - start_m = tl.program_id(0) - off_hz = tl.program_id(1) - batch = off_hz // H - head = off_hz % H - - q_offset = batch * stride_qz + head * BLOCK_DMODEL - k_offset = q_offset + hidden_size - v_offset = k_offset + hidden_size - - # initialize offsets - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - offs_n = tl.arange(0, BLOCK_N) - offs_d = tl.arange(0, BLOCK_DMODEL) - - q_ptrs = QKV + q_offset + offs_m[:, None] * stride_qn + offs_d[None, :] - k_ptrs = QKV + hidden_size + q_offset + offs_n[:, None] * stride_qn + offs_d[None, :] - v_ptrs = QKV + 2 * hidden_size + q_offset + offs_n[:, None] * stride_qn + offs_d[None, :] - - # mask - off_mask = batch * stride_mz + offs_n[None, :] - mask_ptrs = mask + off_mask - - # initialize pointer to m and l - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - l_i = tl.zeros([BLOCK_M], dtype=tl.float32) - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) - # scale sm_scale by log_2(e) and use - # 2^x instead of exp in the loop because CSE and LICM - # don't work as expected with `exp` in the loop - qk_scale = sm_scale * 1.44269504 - # load q: it will stay in SRAM throughout - q = tl.load(q_ptrs, mask=offs_m[:, None] < N_CTX, other=0.0) - q = (q * qk_scale).to(tl.float16) - # loop over k, v and update accumulator - lo = 0 - hi = P_SEQ + (start_m + 1) * BLOCK_M if IS_CAUSAL else N_CTX + P_SEQ - for start_n in range(lo, hi, BLOCK_N): - # -- load k, v -- - k = tl.load(k_ptrs + start_n * stride_qn, mask=(start_n + offs_n)[:, None] < N_CTX, other=0.0) - v = tl.load(v_ptrs + start_n * stride_qn, mask=(start_n + offs_n)[:, None] < N_CTX, other=0.0) - # -- compute qk --- - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float16) - - if ADD_MASK: - mask_val = tl.load(mask_ptrs) - mask_ptrs += BLOCK_N - qk = qk + mask_val.to(tl.float32) - - if IS_CAUSAL: - qk = tl.where(P_SEQ + offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf")) - - qk += tl.dot(q, tl.trans(k), out_dtype=tl.float16) - qk += tl.where((start_n + offs_n)[None, :] < N_CTX, 0, minus_inf) - # -- compute scaling constant --- - m_i_new = tl.maximum(m_i, tl.max(qk, 1)) - alpha = tl.math.exp2(m_i - m_i_new) - p = tl.math.exp2(qk - m_i_new[:, None]) - # -- scale and update acc -- - acc_scale = l_i * 0 + alpha # workaround some compiler bug - acc *= acc_scale[:, None] - acc += tl.dot(p.to(tl.float16), v.to(tl.float16)) - # -- update m_i and l_i -- - l_i = l_i * alpha + tl.sum(p, 1) - m_i = m_i_new - - # write back l and m - acc = acc / l_i[:, None] - o_offset = batch * stride_oz + head * BLOCK_DMODEL - out_ptrs = Out + o_offset + (offs_m[:, None] * stride_on + offs_d[None, :]) - tl.store(out_ptrs, acc.to(tl.float16), mask=offs_m[:, None] < N_CTX) - - -def _triton_packed_flash(qkv, head_size, mask, sm_scale, causal=False, add_mask=True): - heads = qkv.shape[-1] // 3 // head_size - hidden_size = qkv.shape[-1] // 3 - - BLOCK_M = 128 - BLOCK_N = 64 if head_size <= 64 else 32 - - o = torch.empty((qkv.shape[0], qkv.shape[1], hidden_size), device=qkv.device, dtype=torch.half) - if mask is None: - mask = torch.empty(0) - add_mask = False - - grid = (triton.cdiv(qkv.shape[1], BLOCK_M), qkv.shape[0] * heads, 1) - num_stages = 4 if head_size <= 64 else 3 - num_warps = 4 - P_SEQ = 0 - - _flash_packed_kernel[grid](qkv, - mask, - add_mask, - causal, - sm_scale, - o, - qkv.stride(0), - qkv.stride(1), - qkv.stride(2), - mask.stride(1) if add_mask else 0, - o.stride(0), - o.stride(1), - qkv.shape[0], - heads, - qkv.shape[1], - P_SEQ, - hidden_size, - BLOCK_M=BLOCK_M, - BLOCK_N=BLOCK_N, - BLOCK_DMODEL=head_size, - num_warps=num_warps, - num_stages=num_stages) - - return o diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/gelu.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/gelu.py deleted file mode 100644 index 738d7d96a1c9d57c3de3558452922737b9b4f7b4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/gelu.py +++ /dev/null @@ -1,38 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import triton -import triton.language as tl -from deepspeed.accelerator import get_accelerator - - -@triton.jit -def gelu_functor(x): - # Using approximation introduces greater parity errors. - # return tl.sigmoid(1.702 * x) * x - return x * 0.5 * (1.0 + tl.math.erf(x / 1.41421356237)) - - -@triton.jit -def gelu_kernel(x_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): - pid = tl.program_id(axis=0) - block_start = pid * BLOCK_SIZE - offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - x = tl.load(x_ptr + offsets, mask=mask) - output = gelu_functor(x) - tl.store(output_ptr + offsets, output, mask=mask) - - -def gelu(activations: torch.Tensor) -> torch.Tensor: - assert activations.is_contiguous() - assert get_accelerator().on_accelerator(activations) - - output = torch.empty_like(activations) - n_elements = output.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) - gelu_kernel[grid](activations, output, n_elements, BLOCK_SIZE=1024) - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/layer_norm.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/layer_norm.py deleted file mode 100644 index d3f313d2ac3d8205702dc2ceb82856154e9ddb2c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/layer_norm.py +++ /dev/null @@ -1,249 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import triton -import triton.language as tl -''' -layer-normalization -modified the triton kernel in -https://github.com/openai/triton/blob/34817ecc954a6f4ca7b4dfb352fdde1f8bd49ca5/python/tutorials/05-layer-norm.py -''' - - -@triton.jit -def layer_norm_kernel( - Out, - A, - Weight, - Bias, - stride, - N, - eps, - BLOCK_SIZE: tl.constexpr, -): - # position of elements processed by this program - row = tl.program_id(0) - Out += row * stride - A += row * stride - # compute mean - mean = 0 - _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - a = tl.load(A + cols, mask=cols < N, other=0.0).to(tl.float32) - _mean += a - mean = tl.sum(_mean, axis=0) / N - # compute variance - _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - a = tl.load(A + cols, mask=cols < N, other=0.0).to(tl.float32) - a = tl.where(cols < N, a - mean, 0.0) - _var += a * a - var = tl.sum(_var, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - # multiply by weight and add bias - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - mask = cols < N - weight = tl.load(Weight + cols, mask=mask) - bias = tl.load(Bias + cols, mask=mask) - a = tl.load(A + cols, mask=mask, other=0.0).to(tl.float32) - a_hat = (a - mean) * rstd - out = a_hat * weight + bias - # # write-back - tl.store(Out + cols, out, mask=mask) - - -@triton.jit -def layer_norm_residual_kernel( - Out, - A, - Residual, - ln_input, - Weight, - Bias, - stride, - N, - eps, - BLOCK_SIZE: tl.constexpr, -): - # position of elements processed by this program - row = tl.program_id(0) - Out += row * stride - A += row * stride - Residual += row * stride - ln_input += row * stride - # compute mean - mean = 0 - _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - a = tl.load(A + cols, mask=cols < N, other=0.0).to(tl.float32) - res = tl.load(Residual + cols, mask=cols < N, other=0.0).to(tl.float32) - a = a + res - tl.store(ln_input + cols, a, mask=cols < N) - _mean += a - mean = tl.sum(_mean, axis=0) / N - # compute variance - _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - a = tl.load(ln_input + cols, mask=cols < N, other=0.0).to(tl.float32) - a = tl.where(cols < N, a - mean, 0.0) - _var += a * a - var = tl.sum(_var, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - # multiply by weight and add bias - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - mask = cols < N - weight = tl.load(Weight + cols, mask=mask) - bias = tl.load(Bias + cols, mask=mask) - a = tl.load(ln_input + cols, mask=mask, other=0.0).to(tl.float32) - a_hat = (a - mean) * rstd - out = a_hat * weight + bias - # write-back - tl.store(Out + cols, out, mask=mask) - - -@triton.jit -def layer_norm_residual_bias_kernel( - Out, - A, - Residual, - InputBias, - ln_input, - Weight, - Bias, - stride, - N, - eps, - BLOCK_SIZE: tl.constexpr, -): - # position of elements processed by this program - row = tl.program_id(0) - Out += row * stride - A += row * stride - Residual += row * stride - ln_input += row * stride - # compute mean - mean = 0 - _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - a = tl.load(A + cols, mask=cols < N, other=0.0).to(tl.float32) - res = tl.load(Residual + cols, mask=cols < N, other=0.0).to(tl.float32) - b = tl.load(InputBias + cols, mask=cols < N, other=0.0).to(tl.float32) - a = a + b + res - tl.store(ln_input + cols, a, mask=cols < N) - _mean += a - mean = tl.sum(_mean, axis=0) / N - # compute variance - _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - a = tl.load(ln_input + cols, mask=cols < N, other=0.0).to(tl.float32) - a = tl.where(cols < N, a - mean, 0.0) - _var += a * a - var = tl.sum(_var, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - # multiply by weight and add bias - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - mask = cols < N - weight = tl.load(Weight + cols, mask=mask) - bias = tl.load(Bias + cols, mask=mask) - a = tl.load(ln_input + cols, mask=mask, other=0.0).to(tl.float32) - a_hat = (a - mean) * rstd - out = a_hat * weight + bias - # write-back - tl.store(Out + cols, out, mask=mask) - - -def layer_norm(a, weight, bias, eps): - assert a.is_contiguous() - assert weight.is_contiguous() - assert bias.is_contiguous() - - # allocate output - out = torch.empty_like(a) - # reshape input data into 2D tensor - a_arg = a.view(-1, a.shape[-1]) - M, N = a_arg.shape - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // a.element_size() - BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - BLOCK_SIZE = max(BLOCK_SIZE, 128) - BLOCK_SIZE = min(BLOCK_SIZE, 4096) - BLOCK_SIZE = BLOCK_SIZE if N <= 4096 else 8192 - # heuristics for number of warps - num_warps = min(max(BLOCK_SIZE // 256, 1), 8) - layer_norm_kernel[(M, )]( - out, - a_arg, - weight, - bias, - a_arg.stride(0), - N, - eps, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - ) - return out - - -def layer_norm_residual(a, input_bias, residual, weight, bias, eps): - assert a.is_contiguous() - assert weight.is_contiguous() - assert bias.is_contiguous() - assert residual.is_contiguous() - - # allocate output and scratch-pad for residual addition - out = torch.empty_like(a) - ln_input = torch.empty_like(a) - # reshape input data into 2D tensor - a_arg = a.view(-1, a.shape[-1]) - residual = residual.view(-1, residual.shape[-1]) - M, N = a_arg.shape - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // a.element_size() - BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) - BLOCK_SIZE = max(BLOCK_SIZE, 128) - BLOCK_SIZE = min(BLOCK_SIZE, 4096) - BLOCK_SIZE = BLOCK_SIZE if N <= 4096 else 8192 - # heuristics for number of warps - num_warps = min(max(BLOCK_SIZE // 256, 1), 8) - if input_bias is None: - layer_norm_residual_kernel[(M, )]( - out, - a_arg, - residual, - ln_input, - weight, - bias, - a_arg.stride(0), - N, - eps, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - ) - else: - layer_norm_residual_bias_kernel[(M, )]( - out, - a_arg, - residual, - input_bias, - ln_input, - weight, - bias, - a_arg.stride(0), - N, - eps, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - ) - return out diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/matmul_ext.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/matmul_ext.py deleted file mode 100644 index 9be4b0098c371d689581e3317bba701472a33ddc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/matmul_ext.py +++ /dev/null @@ -1,480 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import triton -import os -from filelock import FileLock -import deepspeed.ops.transformer.inference.triton.triton_matmul_kernel as triton_matmul_kernel -import pickle -from io import open -import deepspeed -from pathlib import Path -import atexit -import subprocess - - -# ----------------------------------------------------------------------------- -# util class/functions for triton -def is_nfs_path(path): - if os.name == 'nt': - return False - - # Normalize the path to get the absolute path - path = os.path.abspath(path) - - # Use the 'df' command to find the file system type for the given path - try: - output = subprocess.check_output(['df', '-T', path], encoding='utf-8') - except subprocess.CalledProcessError: - return False # Command failed - - # Process the output of 'df -T' to check for 'nfs' in the filesystem type column - lines = output.strip().split('\n') - if len(lines) > 1: # The first line is headers - fs_type = lines[1].split()[1].lower() # File system type is the second column - return 'nfs' in fs_type - return False - - -class TritonCacheDir: - _warning_printed = False - - @staticmethod - def warn_if_nfs(cache_dir): - if is_nfs_path(cache_dir) and not TritonCacheDir._warning_printed: - print( - f"Warning: The cache directory for DeepSpeed Triton autotune, {cache_dir}, appears to be on an NFS system. While this is generally acceptable, if you experience slowdowns or hanging when DeepSpeed exits, it is recommended to set the TRITON_CACHE_DIR environment variable to a non-NFS path." - ) - TritonCacheDir._warning_printed = True - return - - @staticmethod - def default_cache_dir(): - tmp_path = os.path.join(Path.home(), ".triton", "autotune") - return tmp_path - - -def bias_add_activation(C, bias=None, activation=""): - if bias is not None: - C += bias - # activation - if activation == "relu": - relu = torch.nn.Relu() - C = relu(C) - elif activation == "leaky_relu": - leaky_relu = torch.nn.LeakyReLU(0.01) - C = leaky_relu(C) - elif activation == "gelu": - sigmoid = torch.nn.Sigmoid() - C = sigmoid(1.702 * C) * C - elif activation == "sigmoid": - sigmoid = torch.nn.Sigmoid() - C = sigmoid(C) - return C - - -class AutotuneCacheManager: - """ - Cache manager for autotune - """ - - def __init__(self, key): - self.key = key - self.file_path = None - self.lock_path = None - # if caching is enabled, get the lock and bin path - self.cache_dir = os.environ.get('TRITON_CACHE_DIR', TritonCacheDir.default_cache_dir()) - TritonCacheDir.warn_if_nfs(self.cache_dir) - if self.cache_dir: - os.makedirs(self.cache_dir, exist_ok=True) - self.file_path = os.path.join(self.cache_dir, self.key + ".pickle") - self.lock_path = self.file_path + ".lock" - - def has_file(self): - return self.file_path and os.path.exists(self.file_path) - - def put(self, table): - if self.file_path: - assert self.lock_path is not None - with FileLock(self.lock_path): - with open(self.file_path + ".tmp", 'wb') as handle: - pickle.dump(table, handle) - os.replace(self.file_path + ".tmp", self.file_path) - - def load(self): - if os.path.exists(self.file_path): - with open(self.file_path, 'rb') as handle: - loaded_dict = pickle.load(handle) - return loaded_dict - else: - return None - - -# ----------------------------------------------------------------------------- -# triton matmul class - - -class MatmulExt(torch.autograd.Function): - """ - a wrapper class that can call different triton matmul kernels depending on the input parameters - """ - - @staticmethod - def forward(A, B, bias=None, activation="", use_triton=True, update_autotune_table=False): - """ - A: input, activation matrix A - B: input, weight matrix B - """ - matmul = None - quantize_activation = False - Batch = 0 - - if len(A.shape) == 3: # if A is 3d-tensor where batch index is given as 0-axis - assert A.is_contiguous(), "matrix A must be contiguous" - Batch, M, K = A.shape - A = A.view(-1, K) - - # fp16 activation and fp16 weight matmul into fp16 output - matmul = fp16_matmul - C = matmul.forward(A, B, use_triton=use_triton, bias=bias, activation=activation) - - if matmul and update_autotune_table: - matmul._update_autotune_table() - - if Batch > 0: - C = C.view(Batch, M, -1) - - return C - - -class TritonMatmul(torch.autograd.Function): - """ - triton matmul kernel superclass - """ - - def __init__(self): - pass - - @staticmethod - def _ref_forward(A, B, ref_dtype=torch.float32): - C = torch.matmul(A.type(ref_dtype), B.type(ref_dtype)) - return C - - @staticmethod - def _read_autotune_table(cache_key, triton_kernel): - cache_manager = AutotuneCacheManager(cache_key) - table = cache_manager.load() - if table: - triton_kernel.cache = table - - @staticmethod - def _write_autotune_table(cache_key, triton_kernel): - cache_manager = AutotuneCacheManager(cache_key) - cache_manager.put(triton_kernel.cache) - - @staticmethod - def _update_autotune_table(cache_key, triton_kernel): - cache_manager = AutotuneCacheManager(cache_key) - autotune_table = cache_manager.load() - if autotune_table is None: - autotune_table = dict() - autotune_table.update(triton_kernel.cache) # always overwrite with the new autotune results - cache_manager = AutotuneCacheManager(cache_key) - cache_manager.put(autotune_table) - - @staticmethod - def forward( - A, - B, - ref_dtype=torch.float32, # fp32 only - bias=None, - activation=""): - C = torch.matmul(A.type(ref_dtype), B.type(ref_dtype)) - C = bias_add_activation(C, bias, activation) - return C - - -class Fp16Matmul(TritonMatmul): - """ - fp16 matrix multiplication kernel - dtypes: fp16 x fp16 = fp16 - """ - - _2d_kernel = triton_matmul_kernel._fp_matmul - _4d_kernel = triton_matmul_kernel.matmul_4d_kernel - _cache_stride = 32 - - def __init__(self, read_cache=True): - super().__init__() - if read_cache: - __class__._read_autotune_table() - - def skip_autotune(self): - __class__._2d_kernel.configs = [__class__._2d_kernel.configs[0]] - __class__._4d_kernel.configs = [__class__._4d_kernel.configs[0]] - - @staticmethod - def forward(A, B, use_triton=True, bias=None, activation=""): - if use_triton: - device = A.device - # handle non-contiguous inputs if necessary - if A.stride(0) > 1 and A.stride(1) > 1: - A = A.contiguous() - if B.stride(0) > 1 and B.stride(1) > 1: - B = B.contiguous() - # checks constraints - assert A.shape[1] == B.shape[0], "incompatible dimensions" - M, K = A.shape - _, N = B.shape - # allocates output - C = torch.empty((M, N), device=device, dtype=A.dtype) - # accumulator types - ACC_TYPE = triton.language.float32 if A.dtype in [torch.float16, torch.bfloat16, torch.float32 - ] else triton.language.int32 - # launch kernel - grid = lambda META: (triton.cdiv(M, META['BLOCK_M']) * triton.cdiv(N, META['BLOCK_N']), META['SPLIT_K']) - __class__._2d_kernel[grid](A, - B, - C, - M, - N, - K, - bias, - A.stride(0), - A.stride(1), - B.stride(0), - B.stride(1), - C.stride(0), - C.stride(1), - M // __class__._cache_stride, - N // __class__._cache_stride, - K // __class__._cache_stride, - GROUP_M=8, - ACC_TYPE=ACC_TYPE, - BIAS_ADD=(0 if bias is None else 1), - ACTIVATION=activation) - else: - C = torch.matmul(A, B) - return C - - @staticmethod - def _matmul_4d(a, b): - assert a.shape[-1] == b.shape[-2], "incompatible dimensions" - assert a.is_contiguous(), "matrix A must be contiguous" - assert b.is_contiguous(), "matrix B must be contiguous" - - B, H, M, K = a.shape - B, H, K, N = b.shape - - assert K > 1, "inner-product dimension K should be larger than 1" - - c = torch.empty((B, H, M, N), device=a.device, dtype=a.dtype) - - grid = lambda META: ( - triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), - H, - B, - ) - - __class__._4d_kernel[grid]( - a, - b, - c, - M, - N, - K, - M // __class__._cache_stride, - N // __class__._cache_stride, - K // __class__._cache_stride, - a.stride(0), - a.stride(1), - a.stride(2), - a.stride(3), - b.stride(0), - b.stride(1), - b.stride(2), - b.stride(3), - c.stride(0), - c.stride(1), - c.stride(2), - c.stride(3), - scale=-1.0, - MASK=False, - ) - return c - - @staticmethod - def _score_4d_matmul(input, head_size, input_mask, scale=-1.0): - assert input.is_contiguous(), "matrix input must be contiguous" - - batches = input.shape[0] - d_model = input.shape[-1] // 3 - num_of_heads = d_model // head_size - - q = input[:, :, :d_model] - k = input[:, :, d_model:d_model * 2] - - q = q.view(batches, -1, num_of_heads, head_size) - k = k.view(batches, -1, num_of_heads, head_size) - - # checks constraints - assert q.shape == k.shape, "incompatible dimensions" - B, M, H, K = q.shape - B, N, H, K = k.shape - - assert K > 1, "inner-product dimension K should be larger than 1" - - # allocates output - output = torch.empty((B, H, M, N), device=q.device, dtype=q.dtype) - grid = lambda META: ( - triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), - H, - B, - ) - __class__._4d_kernel[grid]( - q, - k, - output, - M, - N, - K, - M // __class__._cache_stride, - N // __class__._cache_stride, - K // __class__._cache_stride, - q.stride(0), - q.stride(2), - q.stride(1), - q.stride(3), - k.stride(0), - k.stride(2), - k.stride(3), - k.stride(1), - output.stride(0), - output.stride(1), - output.stride(2), - output.stride(3), - scale=scale, - MASK=False, - ) - return output - - @staticmethod - def _context_4d_matmul(prob, input, head_size): - assert prob.is_contiguous(), "matrix prob must be contiguous" - assert input.is_contiguous(), "matrix input must be contiguous" - - batches = input.shape[0] - d_model = input.shape[-1] // 3 - num_of_heads = d_model // head_size - - v = input[:, :, d_model * 2:] - - v = v.view(batches, -1, num_of_heads, head_size) - - # checks constraints - assert (prob.shape[0] == v.shape[0] and prob.shape[1] == v.shape[2] and prob.shape[2] == v.shape[1] - and prob.shape[3] == v.shape[1]), "incompatible dimensions" - B, H, M, K = prob.shape - B, K, H, N = v.shape - - assert K > 1, "inner-product dimension K should be larger than 1" - - # allocates output - output = torch.empty((B, M, H, N), device=v.device, dtype=v.dtype) - grid = lambda META: ( - triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), - H, - B, - ) - - __class__._4d_kernel[grid]( - prob, - v, - output, - M, - N, - K, - M // __class__._cache_stride, - N // __class__._cache_stride, - K // __class__._cache_stride, - prob.stride(0), - prob.stride(1), - prob.stride(2), - prob.stride(3), - v.stride(0), - v.stride(2), - v.stride(1), - v.stride(3), - # Here we also transpose the output when writing to memory. - output.stride(0), - output.stride(2), - output.stride(1), - output.stride(3), - scale=-1, - MASK=False, - ) - return output.view(batches, -1, d_model) - - @staticmethod - def _ref_forward(A, B, ref_dtype=torch.float32, bias=None, activation=""): - C = torch.matmul(A.type(ref_dtype), B.type(ref_dtype)) - C = bias_add_activation(C, bias, activation) - return C - - @staticmethod - def _check_parity(A, - B, - output_dtype, - SA=None, - SB=None, - qblock_size=None, - ref_dtype=torch.float32, - tol=0.01, - use_triton=True, - bias=None, - activation=""): - torch_output = __class__._ref_forward(A, B, ref_dtype=ref_dtype, bias=bias, activation=activation) - triton_output = __class__.forward(A, B, use_triton=use_triton, bias=bias, activation=activation) - assert torch.allclose(triton_output.cpu().type(torch_output.dtype), torch_output.cpu(), rtol=tol) - print(f"{__class__.__name__}: PASSed the parity check") - return triton_output, torch_output - - @staticmethod - def _read_autotune_table(): - TritonMatmul._read_autotune_table(__class__.__name__ + "_2d_kernel", __class__._2d_kernel) - TritonMatmul._read_autotune_table(__class__.__name__ + "_4d_kernel", __class__._4d_kernel) - - @staticmethod - def _write_autotune_table(): - TritonMatmul._write_autotune_table(__class__.__name__ + "_2d_kernel", __class__._2d_kernel) - TritonMatmul._write_autotune_table(__class__.__name__ + "_4d_kernel", __class__._4d_kernel) - - @staticmethod - def _update_autotune_table(): - TritonMatmul._update_autotune_table(__class__.__name__ + "_2d_kernel", __class__._2d_kernel) - TritonMatmul._update_autotune_table(__class__.__name__ + "_4d_kernel", __class__._4d_kernel) - - -# ----------------------------------------------------------------------------- -# mapping -if deepspeed.HAS_TRITON: - fp16_matmul = Fp16Matmul() - matmul = MatmulExt.forward - matmul_4d = fp16_matmul._matmul_4d - score_4d_matmul = fp16_matmul._score_4d_matmul - context_4d_matmul = fp16_matmul._context_4d_matmul -else: - fp16_matmul = None - matmul = None - matmul_4d = None - score_4d_matmul = None - context_4d_matmul = None - - -@atexit.register -def matmul_ext_update_autotune_table(): - if deepspeed.HAS_TRITON: - fp16_matmul._update_autotune_table() diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/mlp.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/mlp.py deleted file mode 100644 index 1708080b27efb0671d361b7cdeaff4b262cf0ce8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/mlp.py +++ /dev/null @@ -1,81 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import math -import torch.nn as nn -from deepspeed.accelerator import get_accelerator -from deepspeed import comm as dist -from ..op_binding import MLPGemmOp, VectorMatMulOp, GELUGemmOp, ResidualAddOp - - -class TritonMLP(nn.Module): - - def __init__(self, config, mp_group=None, q_scales=None, q_groups=1, merge_count=1, mlp_extra_grouping=False): - super(TritonMLP, self).__init__() - - self.config = config - data_type = self.config.dtype - data_type_fp = torch.half if self.config.dtype == torch.int8 else self.config.dtype - device = get_accelerator().current_device_name() - self.attn_nw = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), - requires_grad=False) - self.attn_nb = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), - requires_grad=False) - intm_size_per_partition = self.config.intermediate_size // self.config.mp_size - self.inter_w = nn.Parameter(torch.empty(self.config.hidden_size, - intm_size_per_partition, - dtype=data_type, - device=device), - requires_grad=False) - self.inter_b = nn.Parameter(torch.empty(intm_size_per_partition, dtype=data_type_fp, device=device), - requires_grad=False) - self.output_w = nn.Parameter(torch.empty(intm_size_per_partition, - self.config.hidden_size, - dtype=data_type, - device=device), - requires_grad=False) - self.output_b = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), - requires_grad=False) - - # used for quantization - self.q_scales = q_scales - self.q_groups = q_groups * 2 if mlp_extra_grouping else q_groups - self.merge_count = int(math.log2(merge_count)) - self.mp_group = mp_group - - self.mlp_gemm_func = MLPGemmOp(config) - self.vector_matmul_func = VectorMatMulOp(config) - self.fused_gemm_gelu = GELUGemmOp(config) - self.residual_add_func = ResidualAddOp(config) - - def forward(self, input, residual, residual_norm, bias): - residual_add = None - if self.attn_nw is None: - output = self.fused_gemm_gelu(input=residual_norm, - weight=self.inter_w, - bias=self.inter_b, - weight_out=self.output_w) - else: - output, residual_add = self.mlp_gemm_func(input=input, - residual=residual, - input_bias=bias, - weight_interm=self.inter_w, - weight_out=self.output_w, - bias=self.inter_b, - gamma=self.attn_nw, - beta=self.attn_nb) - residual = self.residual_add_func(hidden_state=output, - residual=residual, - attention_output=input, - attention_bias=bias if bias is not None else self.output_b, - final_bias=self.output_b, - add_bias=bias is not None, - residual_add=residual_add) - - if self.mp_group is not None and dist.get_world_size(group=self.mp_group) > 1: - dist.all_reduce(residual, group=self.mp_group) - - return residual diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/ops.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/ops.py deleted file mode 100644 index dbed453137806a66d13fbef62292297e1c30c3f4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/ops.py +++ /dev/null @@ -1,123 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import deepspeed.ops.transformer.inference.triton.matmul_ext as matmul_ext -from deepspeed.ops.transformer.inference.op_binding.layer_norm import LayerNormOp -from deepspeed.ops.transformer.inference.triton.layer_norm import layer_norm, layer_norm_residual -from deepspeed.utils.types import ActivationFuncType - - -def vector_matmul_func(input, weight, async_op, q_scale, q_int8, transposed_mode): - assert not transposed_mode and not async_op and not q_int8 - return matmul_ext.matmul(input, weight, bias=None, activation="", use_triton=True) - - -def fused_gemm_gelu(input, - weight, - weight_scale, - bias, - weight_out, - weight_out_scale, - epsilon, - pre_layer_norm, - q_int8, - async_op, - transposed_mode, - use_triton_ln=True): - assert not transposed_mode - - # activation - activation = "gelu" - - # intermediate fc in FF - intm_out = matmul_ext.matmul(input, weight, bias=bias, activation=activation, use_triton=True) - - # output fc in FF - ff_out = matmul_ext.matmul( - intm_out, - weight_out, - bias=None, - activation="", # bias added layer with residual_add + bias + layerNorm layer - use_triton=True) - return ff_out - - -def linear_func(input, weight, bias, add_bias, do_flash_attn, num_heads, transposed_mode=False): - assert not transposed_mode and not do_flash_attn - qkv_out = matmul_ext.matmul(input, weight, bias=(bias if add_bias else None), activation="", use_triton=True) - - return qkv_out - - -def mlp_gemm_func(input, - residual, - input_bias, - weight_interm, - weight_out, - bias, - gamma, - beta, - epsilon, - pre_layer_norm, - mlp_after_attn, - weight_interm_scale, - weight_out_scale, - q_int8, - mlp_act_func_type, - transposed_mode, - use_triton_ln=True): - assert not transposed_mode - - # residual add and layerNorm after attention - if use_triton_ln: - mlp_input = layer_norm_residual(input, input_bias, residual, gamma, beta, epsilon) - else: - mlp_input = LayerNormOp.layer_norm_residual(input, input_bias, residual, gamma, beta, epsilon) - - # activation - if ActivationFuncType(mlp_act_func_type) == ActivationFuncType.GELU: - activation = "gelu" - elif ActivationFuncType(mlp_act_func_type) == ActivationFuncType.ReLU: - activation = "relu" - else: - activation = "" - - # intermediate fc in FF - intm_out = matmul_ext.matmul(mlp_input, weight_interm, bias=bias, activation=activation, use_triton=True) - # output fc in FF - ff_out = matmul_ext.matmul( - intm_out, - weight_out, - bias=None, - activation="", # bias added layer with residual_add + bias + layerNorm layer - use_triton=True) - - return ff_out, mlp_input - - -def qkv_gemm_func( - input, - weight, - q_scale, - bias, - gamma, - beta, - epsilon, - add_bias, - q_int8, - transposed_mode=False, - use_triton_ln=True, -): - - assert not transposed_mode - # residual add and layerNorm after attention - if use_triton_ln: - qkv_input = layer_norm(input, gamma, beta, epsilon) - else: - qkv_input = LayerNormOp()(input, gamma, beta, epsilon) - - qkv_out = matmul_ext.matmul(qkv_input, weight, bias=(bias if add_bias else None), activation="", use_triton=True) - - return qkv_out, qkv_input diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/residual_add.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/residual_add.py deleted file mode 100644 index 063e7a7e4a2d9740ad641b3116ea2cec37b472f4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/residual_add.py +++ /dev/null @@ -1,88 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import triton -import triton.language as tl -from deepspeed.accelerator import get_accelerator - - -@triton.jit -def residual_add_bias_kernel( - hidden_state_ptr, - residual_ptr, - attn_output_ptr, - hidden_state_size, - attn_bias_ptr, - final_bias_ptr, - bias_size, - output_ptr, - mp_size: tl.constexpr, - mlp_after_attn: tl.constexpr, - pre_attn_norm: tl.constexpr, - add_attn_bias: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(axis=0) - - block_start = pid * BLOCK_SIZE - - offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = offsets < hidden_state_size - - bias_offsets = offsets % bias_size - bias_mask = bias_offsets < bias_size - - tl_hidden_state = tl.load(hidden_state_ptr + offsets, mask=mask) - tl_residual = tl.load(residual_ptr + offsets, mask=mask) - tl_attn_output = tl.load(attn_output_ptr + offsets, mask=mask) - tl_attn_bias = tl.load(attn_bias_ptr + bias_offsets, mask=bias_mask) - tl_final_bias = tl.load(final_bias_ptr + bias_offsets, mask=bias_mask) - - if mlp_after_attn: - if pre_attn_norm: - output = tl_hidden_state + (tl_residual + tl_final_bias + tl_attn_output + tl_attn_bias) / mp_size - else: - output = tl_hidden_state + tl_residual + tl_final_bias - else: - output = tl_hidden_state + tl_attn_output + (tl_residual + tl_final_bias) / mp_size - if add_attn_bias: - output += tl_attn_bias / mp_size - - tl.store(output_ptr + offsets, output, mask=mask) - - -def residual_add_bias(hidden_state: torch.Tensor, residual: torch.Tensor, attn_output: torch.Tensor, - attn_bias: torch.Tensor, final_bias: torch.Tensor, mp_size: int, mlp_after_attn: bool, - add_attn_bias: bool, pre_attn_norm: bool): - # check that all tensors are on the same device - assert get_accelerator().on_accelerator(hidden_state) \ - and get_accelerator().on_accelerator(residual) \ - and get_accelerator().on_accelerator(attn_output) \ - and get_accelerator().on_accelerator(attn_bias) \ - and get_accelerator().on_accelerator(final_bias) - - # check that all tensors have the same dtype - assert hidden_state.dtype == residual.dtype == attn_output.dtype \ - == attn_bias.dtype == final_bias.dtype - - # check that all tensors have the right shape - assert hidden_state.shape == residual.shape == attn_output.shape - assert attn_bias.shape == final_bias.shape - assert attn_bias.shape[0] == hidden_state.shape[2] - - output = torch.empty_like(hidden_state) - - hidden_state_size = output.numel() - bias_size = attn_bias.numel() - - grid = lambda meta: (triton.cdiv(hidden_state_size, meta['BLOCK_SIZE']), ) - - residual_add_bias_kernel[grid](hidden_state, residual, attn_output, hidden_state_size,\ - attn_bias, final_bias, bias_size, output, mp_size, mlp_after_attn, pre_attn_norm, \ - add_attn_bias, \ - BLOCK_SIZE=1024) - - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/softmax.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/softmax.py deleted file mode 100644 index 1ee10d63e6cf8bfa6723856b53b7ca9ec30d3fdd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/softmax.py +++ /dev/null @@ -1,89 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import triton -import triton.language as tl -''' -softmax -modified the triton kernel in -https://github.com/openai/triton/blob/34817ecc954a6f4ca7b4dfb352fdde1f8bd49ca5/python/tutorials/02-fused-softmax.py -''' - - -@triton.jit -def softmax_kernel(output_ptr, input_ptr, stride, n_cols, BLOCK_SIZE: tl.constexpr): - row_idx = tl.program_id(0) - row_start_ptr = input_ptr + row_idx * stride - col_offsets = tl.arange(0, BLOCK_SIZE) - input_ptrs = row_start_ptr + col_offsets - row = tl.load(input_ptrs, mask=col_offsets < n_cols, other=-float('inf')).to(tl.float32) - row_minus_max = row - tl.max(row, axis=0) - numerator = tl.exp(row_minus_max) - denominator = tl.sum(numerator, axis=0) - softmax_output = numerator / denominator - output_row_start_ptr = output_ptr + row_idx * stride - output_ptrs = output_row_start_ptr + col_offsets - tl.store(output_ptrs, softmax_output, mask=col_offsets < n_cols) - - -@triton.jit -def masked_softmax_kernel(output_ptr, input_ptr, stride, mask_ptr, mask_stride, n_cols, BLOCK_SIZE: tl.constexpr): - row_idx = tl.program_id(0) - row_start_ptr = input_ptr + row_idx * stride - col_offsets = tl.arange(0, BLOCK_SIZE) - input_ptrs = row_start_ptr + col_offsets - mask_ptrs = mask_ptr + col_offsets + row_idx * mask_stride # mask_stride is 0 for 1d mask - row = tl.load(input_ptrs, mask=col_offsets < n_cols, other=-float('inf')).to(tl.float32) - mask = tl.load(mask_ptrs, mask=col_offsets < n_cols, other=0).to(tl.float32) - row_minus_max = row - tl.max(row, axis=0) - row_minus_max = row_minus_max + mask - numerator = tl.exp(row_minus_max) - denominator = tl.sum(numerator, axis=0) - softmax_output = numerator / denominator - output_row_start_ptr = output_ptr + row_idx * stride - output_ptrs = output_row_start_ptr + col_offsets - tl.store(output_ptrs, softmax_output, mask=col_offsets < n_cols) - - -def softmax(input: torch.Tensor, mask: torch.Tensor = None, dim=-1) -> torch.Tensor: - assert input.is_contiguous() - assert (dim == -1) or (dim == len(input.shape) - 1), "Only dim=-1 is supported" - - use_mask = False if mask is None else True - input_arg = input.view(-1, input.shape[-1]) - n_rows, n_cols = input_arg.shape - BLOCK_SIZE = max(triton.next_power_of_2(n_cols), 2) - num_warps = 4 - if BLOCK_SIZE >= 2048: - num_warps = 8 - if BLOCK_SIZE >= 4096: - num_warps = 16 - # Allocate output - output = torch.empty_like(input) - if use_mask: - assert mask.is_contiguous() - mask = mask.view(-1, mask.shape[-1]) - mask_stride = mask.shape[-1] if mask.shape[-2] > 1 else 0 - masked_softmax_kernel[(n_rows, )]( - output, - input, - input_arg.stride(0), - mask, - mask_stride, - n_cols, - num_warps=num_warps, - BLOCK_SIZE=BLOCK_SIZE, - ) - else: - softmax_kernel[(n_rows, )]( - output, - input, - input_arg.stride(0), - n_cols, - num_warps=num_warps, - BLOCK_SIZE=BLOCK_SIZE, - ) - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/triton_matmul_kernel.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/triton_matmul_kernel.py deleted file mode 100644 index e2128e046df049ddbd846131b3dc6001083e991b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/triton_matmul_kernel.py +++ /dev/null @@ -1,398 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import triton -import triton.language as tl -from .gelu import gelu_functor -import torch - -AUTOTUNE_TOP_K = 10 -SKIP_AUTOTUNE = False - - -def _triton_ops_matmul_early_config_prune(configs, named_args): - device = torch.cuda.current_device() #ignore-cuda - capability = torch.cuda.get_device_capability() #ignore-cuda - # BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K, num_warps, num_stages - dtsize = named_args['A'].element_size() - dtype = named_args['A'].dtype - - # 1. make sure we have enough smem - pruned_configs = [] - for config in configs: - kw = config.kwargs - BLOCK_M, BLOCK_N, BLOCK_K, num_stages = \ - kw['BLOCK_M'], kw['BLOCK_N'], kw['BLOCK_K'], config.num_stages - - max_shared_memory = triton.runtime.driver.utils.get_device_properties(device)["max_shared_mem"] - required_shared_memory = (BLOCK_M + BLOCK_N) * BLOCK_K * num_stages * dtsize - if required_shared_memory <= max_shared_memory: - pruned_configs.append(config) - - return pruned_configs - - -def _fp16_matmul_prune_config(configs, named_args, skip_autotune=SKIP_AUTOTUNE): - if skip_autotune: - configs = [configs[0]] - else: - configs = _triton_ops_matmul_early_config_prune(configs, named_args) - return configs - - -""" -fp16 matmul implementation is adapted from triton matmul: -https://github.com/openai/triton/blob/34817ecc954a6f4ca7b4dfb352fdde1f8bd49ca5/python/triton/ops/matmul.py -""" - - -@triton.autotune( - configs=[ - # basic configs for compute-bound matmuls - triton.Config({ - 'BLOCK_M': 128, - 'BLOCK_N': 256, - 'BLOCK_K': 32, - 'SPLIT_K': 1 - }, num_stages=3, num_warps=8), - triton.Config({ - 'BLOCK_M': 256, - 'BLOCK_N': 128, - 'BLOCK_K': 32, - 'SPLIT_K': 1 - }, num_stages=3, num_warps=8), - triton.Config({ - 'BLOCK_M': 256, - 'BLOCK_N': 64, - 'BLOCK_K': 32, - 'SPLIT_K': 1 - }, num_stages=4, num_warps=4), - triton.Config({ - 'BLOCK_M': 64, - 'BLOCK_N': 256, - 'BLOCK_K': 32, - 'SPLIT_K': 1 - }, num_stages=4, num_warps=4), - triton.Config({ - 'BLOCK_M': 128, - 'BLOCK_N': 128, - 'BLOCK_K': 32, - 'SPLIT_K': 1 - }, num_stages=4, num_warps=4), - triton.Config({ - 'BLOCK_M': 128, - 'BLOCK_N': 64, - 'BLOCK_K': 32, - 'SPLIT_K': 1 - }, num_stages=4, num_warps=4), - triton.Config({ - 'BLOCK_M': 64, - 'BLOCK_N': 128, - 'BLOCK_K': 32, - 'SPLIT_K': 1 - }, num_stages=4, num_warps=4), - triton.Config({ - 'BLOCK_M': 128, - 'BLOCK_N': 32, - 'BLOCK_K': 32, - 'SPLIT_K': 1 - }, num_stages=4, num_warps=4), - triton.Config({ - 'BLOCK_M': 64, - 'BLOCK_N': 32, - 'BLOCK_K': 32, - 'SPLIT_K': 1 - }, num_stages=5, num_warps=2), - ], - key=['CACHE_M', 'CACHE_N', 'CACHE_K'], - prune_configs_by={ - 'early_config_prune': _fp16_matmul_prune_config, - 'perf_model': None, - 'top_k': AUTOTUNE_TOP_K - }, -) -@triton.heuristics({ - 'EVEN_K': lambda args: args['K'] % (args['BLOCK_K'] * args['SPLIT_K']) == 0, -}) -@triton.jit -def _fp_matmul( - A, - B, - C, - M, - N, - K, - bias, - stride_am, - stride_ak, - stride_bk, - stride_bn, - stride_cm, - stride_cn, - CACHE_M, - CACHE_N, - CACHE_K, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, - GROUP_M: tl.constexpr, - SPLIT_K: tl.constexpr, - EVEN_K: tl.constexpr, - ACC_TYPE: tl.constexpr, - BIAS_ADD: tl.constexpr, - ACTIVATION: tl.constexpr, -): - # matrix multiplication - pid = tl.program_id(0) - pid_z = tl.program_id(1) - grid_m = (M + BLOCK_M - 1) // BLOCK_M - grid_n = (N + BLOCK_N - 1) // BLOCK_N - # re-order program ID for better L2 performance - width = GROUP_M * grid_n - group_id = pid // width - group_size = min(grid_m - group_id * GROUP_M, GROUP_M) - pid_m = group_id * GROUP_M + (pid % group_size) - pid_n = (pid % width) // (group_size) - # do matrix multiplication - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) - rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) - rk = pid_z * BLOCK_K + tl.arange(0, BLOCK_K) - # pointers - A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) - B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE) - for k in range(K, 0, -BLOCK_K * SPLIT_K): - if EVEN_K: - a = tl.load(A) - b = tl.load(B) - else: - a = tl.load(A, mask=rk[None, :] < k, other=0.) - b = tl.load(B, mask=rk[:, None] < k, other=0.) - acc += tl.dot(a, b) - A += BLOCK_K * SPLIT_K * stride_ak - B += BLOCK_K * SPLIT_K * stride_bk - # bias addition - if BIAS_ADD: - bias_offset = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - bias_ptr = bias + bias_offset - b = tl.load(bias_ptr, mask=bias_offset < N) - acc = acc + b[None, :] - # activation - if ACTIVATION == "relu": - acc = tl.where(acc >= 0, acc, 0) - elif ACTIVATION == "leaky_relu": - acc = tl.where(acc >= 0, acc, 0.01 * acc) - elif ACTIVATION == "gelu": - #acc = tl.sigmoid(1.702 * acc) * acc - acc = gelu_functor(acc) - elif ACTIVATION == "sigmoid": - acc = tl.sigmoid(acc) # sigmoid - acc = acc.to(C.dtype.element_ty) - # rematerialize rm and rn to save registers - rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn) - mask = (rm < M)[:, None] & (rn < N)[None, :] - # handles write-back with reduction-splitting - if SPLIT_K == 1: - tl.store(C, acc, mask=mask) - else: - tl.atomic_add(C, acc, mask=mask) - - -def matmul_4d_prune_config(configs, named_args, skip_autotune=SKIP_AUTOTUNE): - if skip_autotune: - configs = [configs[0]] - else: - device = torch.cuda.current_device() #ignore-cuda - capability = torch.cuda.get_device_capability() #ignore-cuda - # BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K, num_warps, num_stages - dtsize = named_args['a_ptr'].element_size() - dtype = named_args['a_ptr'].dtype - - # make sure we have enough smem - pruned_configs = [] - for config in configs: - kw = config.kwargs - BLOCK_M, BLOCK_N, BLOCK_K, num_stages = \ - kw['BLOCK_SIZE_M'], kw['BLOCK_SIZE_N'], kw['BLOCK_SIZE_K'], config.num_stages - - max_shared_memory = triton.runtime.driver.utils.get_device_properties(device)["max_shared_mem"] - required_shared_memory = (BLOCK_M + BLOCK_N) * BLOCK_K * num_stages * dtsize - if required_shared_memory <= max_shared_memory: - pruned_configs.append(config) - configs = pruned_configs - return configs - - -@triton.autotune( - configs=[ - triton.Config( - { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 64, - "GROUP_SIZE_M": 8 - }, - num_stages=1, # this is mainly for unit test, to minimize the share memory usage - num_warps=8), - triton.Config( - { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 32, - "GROUP_SIZE_M": 8, - }, - num_stages=4, - num_warps=4, - ), - triton.Config( - { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 32, - "GROUP_SIZE_M": 8, - }, - num_stages=4, - num_warps=4, - ), - triton.Config( - { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 32, - "GROUP_SIZE_M": 8, - }, - num_stages=4, - num_warps=4, - ), - triton.Config( - { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 32, - "GROUP_SIZE_M": 8, - }, - num_stages=4, - num_warps=4, - ), - triton.Config( - { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 32, - "GROUP_SIZE_M": 8, - }, - num_stages=5, - num_warps=2, - ), - triton.Config( - { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 32, - "GROUP_SIZE_M": 8, - }, - num_stages=5, - num_warps=2, - ), - ], - key=['CACHE_M', 'CACHE_N', 'CACHE_K'], - prune_configs_by={ - 'early_config_prune': matmul_4d_prune_config, - 'perf_model': None, - 'top_k': AUTOTUNE_TOP_K - }, -) -@triton.jit -def matmul_4d_kernel( - # Pointers to matrices - a_ptr, - b_ptr, - c_ptr, - # Matrix dimensions - M, - N, - K, - CACHE_M, - CACHE_N, - CACHE_K, - stride_ab, - stride_ah, - stride_am, - stride_ak, - stride_bb, - stride_bh, - stride_bk, - stride_bn, - stride_cb, - stride_ch, - stride_cm, - stride_cn, - scale, - # Meta-parameters - BLOCK_SIZE_M: tl.constexpr, - BLOCK_SIZE_N: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr, - MASK: tl.constexpr, -): - """Kernel for computing the matmul C = A x B. - A has shape (M, K), B has shape (K, N) and C has shape (M, N) - """ - pid = tl.program_id(axis=0) - head = tl.program_id(axis=1) - batch = tl.program_id(axis=2) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + (pid % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - - if MASK: - if (pid_m + 1) * BLOCK_SIZE_M - 1 < pid_n * BLOCK_SIZE_N: - c = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=c_ptr.dtype.element_ty) - float("inf") - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - c_ptrs = (c_ptr + batch * stride_cb + head * stride_ch + stride_cm * offs_cm[:, None] + - stride_cn * offs_cn[None, :]) - tl.store(c_ptrs, c) - return - - offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - offs_k = tl.arange(0, BLOCK_SIZE_K) - a_ptrs = (a_ptr + batch * stride_ab + head * stride_ah + - (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)) - b_ptrs = (b_ptr + batch * stride_bb + head * stride_bh + - (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - for k in range(0, K, BLOCK_SIZE_K): - a_mask = (offs_am[:, None] < M) & (offs_k[None, :] + k < K) - b_mask = (offs_k[:, None] + k < K) & (offs_bn[None, :] < N) - a = tl.load(a_ptrs, mask=a_mask, other=0.) - b = tl.load(b_ptrs, mask=b_mask, other=0.) - accumulator += tl.dot(a, b) - a_ptrs += BLOCK_SIZE_K * stride_ak - b_ptrs += BLOCK_SIZE_K * stride_bk - - c = accumulator.to(c_ptr.dtype.element_ty) - if scale > 0: - c = c * scale.to(c_ptr.dtype.element_ty) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - if MASK: - c += tl.where(offs_cm[:, None] >= offs_cn[None, :], 0, float("-inf")) - c_ptrs = (c_ptr + batch * stride_cb + head * stride_ch + stride_cm * offs_cm[:, None] + - stride_cn * offs_cn[None, :]) - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) - tl.store(c_ptrs, c, mask=c_mask) diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton_ops.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton_ops.py deleted file mode 100644 index f98f45ef638eb875454c9f9a40e9ad1bbd347320..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton_ops.py +++ /dev/null @@ -1,166 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Inspired by original Triton implementation: -https://github.com/openai/triton/blob/release/2.1.x/python/tutorials/06-fused-attention.py -""" - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _fwd_kernel( - Q, - K, - V, - sm_scale, - Out, - stride_qz, - stride_qh, - stride_qm, - stride_qk, - stride_kz, - stride_kh, - stride_kn, - stride_kk, - stride_vz, - stride_vh, - stride_vk, - stride_vn, - stride_oz, - stride_oh, - stride_om, - stride_on, - Z, - H, - N_CTX, - BLOCK_M: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - BLOCK_N: tl.constexpr, -): - start_m = tl.program_id(0) - off_hz = tl.program_id(1) - qvk_offset = off_hz * stride_qh - Q_block_ptr = tl.make_block_ptr(base=Q + qvk_offset, - shape=(N_CTX, BLOCK_DMODEL), - strides=(stride_qm, stride_qk), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, BLOCK_DMODEL), - order=(1, 0)) - K_block_ptr = tl.make_block_ptr(base=K + qvk_offset, - shape=(BLOCK_DMODEL, N_CTX), - strides=(stride_kk, stride_kn), - offsets=(0, 0), - block_shape=(BLOCK_DMODEL, BLOCK_N), - order=(0, 1)) - V_block_ptr = tl.make_block_ptr(base=V + qvk_offset, - shape=(N_CTX, BLOCK_DMODEL), - strides=(stride_vk, stride_vn), - offsets=(0, 0), - block_shape=(BLOCK_N, BLOCK_DMODEL), - order=(1, 0)) - # initialize offsets - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - offs_n = tl.arange(0, BLOCK_N) - # initialize pointer to m and l - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - l_i = tl.zeros([BLOCK_M], dtype=tl.float32) - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) - # scale sm_scale by log_2(e) and use - # 2^x instead of exp in the loop because CSE and LICM - # don't work as expected with `exp` in the loop - qk_scale = sm_scale * 1.44269504 - # load q: it will stay in SRAM throughout - q = tl.load(Q_block_ptr) - q = (q * qk_scale).to(tl.float16) - # loop over k, v and update accumulator - lo = 0 - #hi = (start_m + 1) * BLOCK_M if IS_CAUSAL else N_CTX - hi = N_CTX - #hi = (start_m + 1) * BLOCK_M - for start_n in range(lo, hi, BLOCK_N): - # -- load k, v -- - k = tl.load(K_block_ptr) - v = tl.load(V_block_ptr) - # -- compute qk --- - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - #if IS_CAUSAL: - #qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf")) - qk += tl.dot(q, k) - # -- compute scaling constant --- - m_i_new = tl.maximum(m_i, tl.max(qk, 1)) - alpha = tl.math.exp2(m_i - m_i_new) - p = tl.math.exp2(qk - m_i_new[:, None]) - # -- scale and update acc -- - acc_scale = l_i * 0 + alpha # workaround some compiler bug - acc *= acc_scale[:, None] - acc += tl.dot(p.to(tl.float16), v) - # -- update m_i and l_i -- - l_i = l_i * alpha + tl.sum(p, 1) - m_i = m_i_new - # update pointers - K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N)) - V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0)) - # write back l and m - acc = acc / l_i[:, None] - #l_ptrs = L + off_hz * N_CTX + offs_m - #tl.store(l_ptrs, m_i + tl.math.log2(l_i)) - # write back O - O_block_ptr = tl.make_block_ptr(base=Out + qvk_offset, - shape=(N_CTX, BLOCK_DMODEL), - strides=(stride_om, stride_on), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, BLOCK_DMODEL), - order=(1, 0)) - tl.store(O_block_ptr, acc.to(tl.float16)) - - -class triton_flash_attn(torch.nn.Module): - - def __init__(self, ): - super(triton_flash_attn, self).__init__() - - def forward(self, q, k, v, sm_scale, block_128=True): - BLOCK = 128 if block_128 else 64 - # shape constraints - Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] - o = torch.empty_like(q) - grid = (triton.cdiv(q.shape[2], BLOCK), q.shape[0] * q.shape[1]) - num_warps = 4 if Lk <= 64 else 8 - - _fwd_kernel[grid]( - q, - k, - v, - sm_scale, - o, - q.stride(0), - q.stride(1), - q.stride(2), - q.stride(3), - k.stride(0), - k.stride(1), - k.stride(2), - k.stride(3), - v.stride(0), - v.stride(1), - v.stride(2), - v.stride(3), - o.stride(0), - o.stride(1), - o.stride(2), - o.stride(3), - k.shape[0], - k.shape[1], - k.shape[2], - BLOCK_M=BLOCK, - BLOCK_N=BLOCK, - BLOCK_DMODEL=Lk, - num_warps=num_warps, - num_stages=1, - ) - return o diff --git a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/transformer.py b/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/transformer.py deleted file mode 100644 index bfd4d60dcb1ceafff2d87b62a78a8cf480ab9448..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/ops/transformer/transformer.py +++ /dev/null @@ -1,412 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import json -import math -import torch -from torch import nn -from torch.autograd import Function -from deepspeed.accelerator import get_accelerator -from deepspeed.ops.op_builder import TransformerBuilder, StochasticTransformerBuilder - -# Cuda modules will be imported if needed -transformer_cuda_module = None -stochastic_transformer_cuda_module = None - - -class TransformerConfig(): - - def __init__(self, batch_size, hidden_size, intermediate_size, heads, attn_dropout_ratio, hidden_dropout_ratio, - num_hidden_layers, initializer_range): - self.layer_id = -1 - self.batch_size = batch_size - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.heads = heads - self.attn_dropout_ratio = attn_dropout_ratio - self.hidden_dropout_ratio = hidden_dropout_ratio - self.num_hidden_layers = num_hidden_layers - self.initializer_range = initializer_range - - -class DeepSpeedTransformerConfig(TransformerConfig): - """Initialize the DeepSpeed Transformer Config. - - Arguments: - batch_size: The maximum batch size used for running the kernel on each GPU - - hidden_size: The hidden size of the transformer layer - - intermediate_size: The intermediate size of the feed-forward part of transformer layer - - heads: The number of heads in the self-attention of the transformer layer - - attn_dropout_ratio: The ratio of dropout for the attention's output - - hidden_dropout_ratio: The ratio of dropout for the transformer's output - - num_hidden_layers: The number of transformer layers - - initializer_range: BERT model's initializer range for initializing parameter data - - local_rank: Optional: The rank of GPU running the transformer kernel, it is not required - to use if the model already set the current device, otherwise need to set it - so that the transformer kernel can work on the right device - - seed: The random seed for the dropout layers - - fp16: Enable half-precision computation - - pre_layer_norm: Select between Pre-LN or Post-LN transformer architecture - - normalize_invertible: Optional: Enable invertible LayerNorm execution (dropping the input activation), - default is False - - gelu_checkpoint: Optional: Enable checkpointing of Gelu activation output to save memory, - default is False - - adjust_init_range: Optional: Set as True (default) if the model adjusts the weight initial values of - its self-attention output and layer output, False keeps the initializer_range no change. - See the adjustment below: - output_std = self.config.initializer_range / math.sqrt(2.0 * num_layers) - - attn_dropout_checkpoint: Optional: Enable checkpointing of attention dropout to save memory, - default is False - - stochastic_mode: Enable for high performance, please note that this flag has some level of - non-determinism and can produce different results on different runs. However, we have seen - that by enabling it, the pretraining tasks such as BERT are not affected and can obtain - a high accuracy level. On the other hand, for the downstream tasks, such as fine-tuning, we recommend - to turn it off in order to be able to reproduce the same result through the regular kernel execution. - - return_tuple: Enable if using the return_tuple interface style for sending out the forward results. - - training: Enable for training rather than inference. - """ - - def __init__(self, - batch_size=-1, - hidden_size=-1, - intermediate_size=-1, - heads=-1, - attn_dropout_ratio=-1, - hidden_dropout_ratio=-1, - num_hidden_layers=-1, - initializer_range=-1, - layer_norm_eps=1e-12, - local_rank=-1, - seed=-1, - fp16=False, - pre_layer_norm=True, - normalize_invertible=False, - gelu_checkpoint=False, - adjust_init_range=True, - attn_dropout_checkpoint=False, - stochastic_mode=False, - return_tuple=False, - training=True): - super(DeepSpeedTransformerConfig, - self).__init__(batch_size, hidden_size, - (intermediate_size if intermediate_size > 0 else 4 * hidden_size), heads, - attn_dropout_ratio, hidden_dropout_ratio, num_hidden_layers, initializer_range) - self.fp16 = fp16 - self.pre_layer_norm = pre_layer_norm - self.local_rank = local_rank - self.seed = seed - self.normalize_invertible = normalize_invertible - self.gelu_checkpoint = gelu_checkpoint # True: if higher batch size is required - self.adjust_init_range = adjust_init_range - self.test_gemm = False - self.layer_norm_eps = layer_norm_eps - self.training = training - self.is_grad_enabled = True - self.attn_dropout_checkpoint = attn_dropout_checkpoint - self.stochastic_mode = stochastic_mode - self.return_tuple = return_tuple - - @classmethod - def from_dict(cls, json_object): - config = DeepSpeedTransformerConfig() - for key, value in json_object.items(): - config.__dict__[key] = value - return config - - @classmethod - def from_json_file(cls, json_file): - with open(json_file, "r", encoding='utf-16') as reader: - text = reader.read() - return cls.from_dict(json.loads(text)) - - -class DeepSpeedTransformerFunction(Function): - - @staticmethod - def forward(ctx, input, input_mask, self, grads, layer_id, attn_qkvw, attn_qkvb, attn_ow, attn_ob, attn_nw, - attn_nb, inter_w, inter_b, output_w, output_b, norm_w, norm_b, config): - - cuda_module = stochastic_transformer_cuda_module if config.stochastic_mode else transformer_cuda_module - forward_func = cuda_module.forward_fp16 if config.fp16 else cuda_module.forward_fp32 - - inp_size = input.size() - if inp_size[1] % 16 != 0: - input = torch.cat( - (input, - torch.randn( - (inp_size[0], (16 - (inp_size[1] % 16)), inp_size[2]), device=input.device, dtype=input.dtype)), - 1) - input_mask = torch.cat((input_mask, torch.ones((inp_size[0], input_mask.shape[1], input_mask.shape[2], \ - (16 - (inp_size[1] % 16))), device=input_mask.device, dtype=input_mask.dtype) * -10000), 3) - - (output, inp_norm, qkv_tf, soft_inp, ctx_bufB, attn_o_inp, add_res, ff1_inp, gelu_inp, ff2_inp, - attn_prob_dropout_mask, attn_output_dropout_mask, layer_output_dropout_mask, attn_layer_norm_var, - attn_layer_norm_mean, layer_norm_var, layer_norm_mean) = forward_func( - config.layer_id, input, input_mask, attn_qkvw, attn_qkvb, attn_ow, attn_ob, attn_nw, attn_nb, inter_w, - inter_b, output_w, output_b, norm_w, norm_b, config.training and config.is_grad_enabled, - config.pre_layer_norm, config.attn_dropout_checkpoint, config.normalize_invertible, - config.gelu_checkpoint) - - # For testing only. - if grads is not None: - for i in [2]: - attn_qkvw.register_hook(lambda x, i=i, self=self: grads.append([ - x[i * attn_ow.size(0):(i + 1) * attn_ow.size(0)], ("Q_W" if i == 0 else "K_W" if i == 1 else "V_W") - ])) - for i in [2]: - attn_qkvb.register_hook(lambda x, i=i, self=self: grads.append([ - x[i * attn_ow.size(0):(i + 1) * attn_ow.size(0)], ("Q_B" if i == 0 else "K_B" if i == 1 else "V_B") - ])) - - attn_ow.register_hook(lambda x, self=self: grads.append([x, "O_W"])) - attn_ob.register_hook(lambda x, self=self: grads.append([x, "O_B"])) - attn_nw.register_hook(lambda x, self=self: grads.append([x, "N2_W"])) - attn_nb.register_hook(lambda x, self=self: grads.append([x, "N2_B"])) - inter_w.register_hook(lambda x, self=self: grads.append([x, "int_W"])) - inter_b.register_hook(lambda x, self=self: grads.append([x, "int_B"])) - output_w.register_hook(lambda x, self=self: grads.append([x, "out_W"])) - output_b.register_hook(lambda x, self=self: grads.append([x, "out_B"])) - norm_w.register_hook(lambda x, self=self: grads.append([x, "norm_W"])) - norm_b.register_hook(lambda x, self=self: grads.append([x, "norm_B"])) - - if config.is_grad_enabled and config.training: - if (config.pre_layer_norm and config.normalize_invertible): - ctx.save_for_backward(input_mask, attn_qkvw, attn_qkvb, attn_ow, attn_ob, attn_nw, attn_nb, inter_w, - inter_b, output_w, output_b, norm_w, norm_b) - else: - ctx.save_for_backward(output, input, input_mask, attn_qkvw, attn_qkvb, attn_ow, attn_ob, attn_nw, - attn_nb, inter_w, inter_b, output_w, output_b, norm_w, norm_b) - - ctx.config = config - if (config.pre_layer_norm or not config.normalize_invertible): - ctx.inp_norm = inp_norm - - ctx.qkv_tf = qkv_tf - ctx.soft_inp = soft_inp - if not config.attn_dropout_checkpoint: - ctx.ctx_bufB = ctx_bufB - - ctx.attn_o_inp = attn_o_inp - if not config.normalize_invertible: - ctx.add_res = add_res - - ctx.attn_layer_norm_mean = attn_layer_norm_mean - ctx.layer_norm_mean = layer_norm_mean - - ctx.ff1_inp = ff1_inp - if not config.gelu_checkpoint: - ctx.gelu_inp = gelu_inp - - ctx.ff2_inp = ff2_inp - ctx.attn_prob_dropout_mask = attn_prob_dropout_mask - ctx.attn_output_dropout_mask = attn_output_dropout_mask - ctx.layer_output_dropout_mask = layer_output_dropout_mask - ctx.attn_layer_norm_var = attn_layer_norm_var - ctx.layer_norm_var = layer_norm_var - - if inp_size[1] % 16 != 0: - output = torch.narrow(output, 1, 0, inp_size[1]) - - if config.return_tuple: - return (output, ) # outputs -> (output) : outputs[0] = output - else: - return output - - @staticmethod - def backward(ctx, grad_output): - bsz = grad_output.shape[0] - grad_output_shape = grad_output.size() - if grad_output_shape[1] % 16 != 0: - grad_output = torch.cat((grad_output, torch.zeros((bsz, (16 - (grad_output_shape[1] % 16)), \ - grad_output_shape[2]), device=grad_output.device, dtype=grad_output.dtype)), 1) - - assert ctx.config.training - - if (ctx.config.pre_layer_norm and ctx.config.normalize_invertible): - (input_mask, attn_qkvw, attn_qkvb, attn_ow, attn_ob, attn_nw, attn_nb, inter_w, inter_b, output_w, - output_b, norm_w, norm_b) = ctx.saved_tensors - else: - (output, input, input_mask, attn_qkvw, attn_qkvb, attn_ow, attn_ob, attn_nw, attn_nb, inter_w, inter_b, - output_w, output_b, norm_w, norm_b) = ctx.saved_tensors - - cuda_module = stochastic_transformer_cuda_module if ctx.config.stochastic_mode else transformer_cuda_module - backward_func = cuda_module.backward_fp16 if ctx.config.fp16 else cuda_module.backward_fp32 - - (grad_input, grad_attn_qkvw, grad_attn_qkvb, grad_attn_ow, grad_attn_ob, grad_attn_nw, grad_attn_nb, - grad_inter_w, grad_inter_b, grad_output_w, grad_output_b, grad_norm_w, grad_norm_b) = backward_func( - ctx.config.layer_id, grad_output, - (ctx.inp_norm if (ctx.config.pre_layer_norm and ctx.config.normalize_invertible) else output), - (ctx.inp_norm if (ctx.config.pre_layer_norm or not ctx.config.normalize_invertible) else input), - ctx.qkv_tf, ctx.soft_inp, (ctx.soft_inp if ctx.config.attn_dropout_checkpoint else ctx.ctx_bufB), - ctx.attn_o_inp, (ctx.ff1_inp if ctx.config.normalize_invertible else ctx.add_res), ctx.ff1_inp, - (ctx.ff2_inp if ctx.config.gelu_checkpoint else ctx.gelu_inp), ctx.ff2_inp, ctx.attn_prob_dropout_mask, - ctx.attn_output_dropout_mask, ctx.layer_output_dropout_mask, ctx.attn_layer_norm_var, - ctx.attn_layer_norm_mean, ctx.layer_norm_var, ctx.layer_norm_mean, - (ctx.inp_norm if - (ctx.config.pre_layer_norm and ctx.config.normalize_invertible) else input), input_mask, attn_qkvw, - attn_qkvb, attn_ow, attn_ob, attn_nw, attn_nb, inter_w, inter_b, output_w, output_b, norm_w, norm_b) - - # This appears to be an effective way to release context memory - ctx.qkv_tf = None - ctx.soft_inp = None - ctx.ctx_bufB = None - ctx.gelu_inp = None - ctx.ff2_inp = None - ctx.attn_o_inp = None - ctx.ff1_inp = None - ctx.add_res = None - ctx.inp_norm = None - ctx.config = None - ctx.attn_layer_norm_mean = None - ctx.layer_norm_mean = None - ctx.attn_prob_dropout_mask = None - ctx.attn_output_dropout_mask = None - ctx.layer_output_dropout_mask = None - ctx.attn_layer_norm_var = None - ctx.layer_norm_var = None - - if grad_output_shape[1] % 16 != 0: - grad_input = torch.narrow(grad_input, 1, 0, grad_output_shape[1]) - - return (grad_input, None, None, None, None, grad_attn_qkvw, grad_attn_qkvb, grad_attn_ow, grad_attn_ob, - grad_attn_nw, grad_attn_nb, grad_inter_w, grad_inter_b, grad_output_w, grad_output_b, grad_norm_w, - grad_norm_b, None) - - -class DeepSpeedTransformerLayer(nn.Module): - """Initialize the DeepSpeed Transformer Layer. - - Static variable: - layer_id: The layer-index counter starting from 0 and incrementing by 1 every time a layer object is instantiated, - e.g. if a model has 24 transformer layers, layer_id goes from 0 to 23. - Arguments: - config: An object of DeepSpeedTransformerConfig - - initial_weights: Optional: Only used for unit test - - initial_biases: Optional: Only used for unit test - """ - layer_id = 0 - - def __init__(self, config, initial_weights=None, initial_biases=None): - super(DeepSpeedTransformerLayer, self).__init__() - - self.config = config - self.config.layer_id = DeepSpeedTransformerLayer.layer_id - DeepSpeedTransformerLayer.layer_id = DeepSpeedTransformerLayer.layer_id + 1 - - print("DeepSpeed Transformer config is ", self.config.__dict__) - - if self.config.local_rank >= 0: - get_accelerator().set_device(self.config.local_rank) - - if initial_weights is None and initial_biases is None: - self.attn_qkvw = nn.Parameter(torch.Tensor(self.config.hidden_size * 3, self.config.hidden_size)) - self.attn_qkvb = nn.Parameter(torch.Tensor(self.config.hidden_size * 3)) - self.attn_ow = nn.Parameter(torch.Tensor(self.config.hidden_size, self.config.hidden_size)) - self.attn_ob = nn.Parameter(torch.Tensor(self.config.hidden_size)) - self.attn_nw = nn.Parameter(torch.Tensor(self.config.hidden_size)) - self.attn_nb = nn.Parameter(torch.Tensor(self.config.hidden_size)) - self.inter_w = nn.Parameter(torch.Tensor(self.config.intermediate_size, self.config.hidden_size)) - self.inter_b = nn.Parameter(torch.Tensor(self.config.intermediate_size)) - self.output_w = nn.Parameter(torch.Tensor(self.config.hidden_size, self.config.intermediate_size)) - self.output_b = nn.Parameter(torch.Tensor(self.config.hidden_size)) - self.norm_w = nn.Parameter(torch.Tensor(self.config.hidden_size)) - self.norm_b = nn.Parameter(torch.Tensor(self.config.hidden_size)) - self.init_transformer_weights(self.config.adjust_init_range) - else: - # For testing only. - q = initial_weights[0].data - k = initial_weights[1].data - v = initial_weights[2].data - - self.attn_qkvw = nn.Parameter(torch.cat((q, k, v))) - #self.attn_qkvw[i * self.config.hidden_size:(i + 1) * self.config.hidden_size] = \ - # initial_weights[i].clone() - #torch.empty_like(initial_weights[i]).data.copy_(initial_weights[i].data) - self.attn_qkvb = nn.Parameter(torch.Tensor(self.config.hidden_size * 3)) - self.attn_qkvb.data.zero_() - self.attn_ow = initial_weights[3] - self.attn_ob = initial_biases[3] - self.attn_nw = initial_weights[4] - self.attn_nb = initial_biases[4] - self.inter_w = initial_weights[5] - self.inter_b = initial_biases[5] - self.output_w = initial_weights[6] - self.output_b = initial_biases[6] - self.norm_w = initial_weights[7] - self.norm_b = initial_biases[7] - - # Load cuda modules if needed - global transformer_cuda_module, stochastic_transformer_cuda_module - if transformer_cuda_module is None and not self.config.stochastic_mode: - transformer_cuda_module = TransformerBuilder().load() - if stochastic_transformer_cuda_module is None and self.config.stochastic_mode: - stochastic_transformer_cuda_module = StochasticTransformerBuilder().load() - - # create the layer in cuda kernels. - cuda_module = stochastic_transformer_cuda_module if self.config.stochastic_mode else transformer_cuda_module - create_layer_func = cuda_module.create_transformer_layer_fp16 if self.config.fp16 else cuda_module.create_transformer_layer_fp32 - - create_layer_func(self.config.layer_id, self.config.batch_size, self.config.hidden_size, self.config.heads, - self.config.intermediate_size, self.config.attn_dropout_ratio, - self.config.hidden_dropout_ratio, self.config.layer_norm_eps, self.config.seed, - self.config.pre_layer_norm, self.config.test_gemm, self.config.attn_dropout_checkpoint, - self.config.normalize_invertible, self.config.gelu_checkpoint, self.config.stochastic_mode) - - def init_transformer_weights(self, adjust_init_range=False): - num_layers = self.config.num_hidden_layers - output_std = self.config.initializer_range - if adjust_init_range and self.config.local_rank == 0: - print("Accounting for accumulation on the residual path") - output_std = self.config.initializer_range / math.sqrt(2.0 * num_layers) - - self.attn_qkvw.data.normal_(mean=0.0, std=self.config.initializer_range) - self.attn_qkvb.data.zero_() - self.attn_ow.data.normal_(mean=0.0, std=output_std) - self.attn_ob.data.zero_() - self.attn_nw.data.fill_(1.0) - self.attn_nb.data.zero_() - self.inter_w.data.normal_(mean=0.0, std=self.config.initializer_range) - self.inter_b.data.zero_() - self.output_w.data.normal_(mean=0.0, std=output_std) - self.output_b.data.zero_() - self.norm_w.data.fill_(1.0) - self.norm_b.data.zero_() - - def forward(self, - hidden_states, - attention_mask=None, - head_mask=None, - layer_head_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - past_key_value=None, - output_attentions=False, - grads=None): - self.config.is_grad_enabled = torch.is_grad_enabled() - self.config.training = self.training - return DeepSpeedTransformerFunction.apply(hidden_states, attention_mask, self, grads, self.config.layer_id, - self.attn_qkvw, self.attn_qkvb, self.attn_ow, self.attn_ob, - self.attn_nw, self.attn_nb, self.inter_w, self.inter_b, - self.output_w, self.output_b, self.norm_w, self.norm_b, self.config) diff --git a/venv/lib/python3.10/site-packages/deepspeed/pipe/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/pipe/__init__.py deleted file mode 100644 index f25471713b4568b43216fe16786a46904575ae5e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/pipe/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from ..runtime.pipe import PipelineModule, LayerSpec, TiedLayerSpec diff --git a/venv/lib/python3.10/site-packages/deepspeed/pipe/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/pipe/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 587abed6296d262c3b55789e267ffd47109f8762..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/pipe/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/profiling/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/profiling/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 06b12d70150cbf207db9598fad677733f1ab4ae9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/config.cpython-310.pyc deleted file mode 100644 index c933b8b9b89694fa6d41125bd3d24327d1562b00..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/constants.cpython-310.pyc deleted file mode 100644 index 3c0459e6637a1ea570b7fb87c6d5a21fedd37fce..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/profiling/__pycache__/constants.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/config.py b/venv/lib/python3.10/site-packages/deepspeed/profiling/config.py deleted file mode 100644 index e4f06630ea6f6801883046efb5f61df3e7132f63..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/profiling/config.py +++ /dev/null @@ -1,46 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigObject -from deepspeed.profiling.constants import * - - -class DeepSpeedFlopsProfilerConfig(DeepSpeedConfigObject): - - def __init__(self, param_dict): - super(DeepSpeedFlopsProfilerConfig, self).__init__() - - self.enabled = None - self.recompute_fwd_factor = None - self.profile_step = None - self.module_depth = None - self.top_modules = None - - if FLOPS_PROFILER in param_dict.keys(): - flops_profiler_dict = param_dict[FLOPS_PROFILER] - else: - flops_profiler_dict = {} - - self._initialize(flops_profiler_dict) - - def _initialize(self, flops_profiler_dict): - self.enabled = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_ENABLED, FLOPS_PROFILER_ENABLED_DEFAULT) - - self.recompute_fwd_factor = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_RECOMPUTE_FWD_FACTOR, - FLOPS_PROFILER_RECOMPUTE_FWD_FACTOR_DEFAULT) - - self.profile_step = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_PROFILE_STEP, - FLOPS_PROFILER_PROFILE_STEP_DEFAULT) - - self.module_depth = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_MODULE_DEPTH, - FLOPS_PROFILER_MODULE_DEPTH_DEFAULT) - - self.top_modules = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_TOP_MODULES, - FLOPS_PROFILER_TOP_MODULES_DEFAULT) - - self.detailed = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_DETAILED, FLOPS_PROFILER_DETAILED_DEFAULT) - - self.output_file = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_OUTPUT_FILE, - FLOPS_PROFILER_OUTPUT_FILE_DEFAULT) diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/constants.py b/venv/lib/python3.10/site-packages/deepspeed/profiling/constants.py deleted file mode 100644 index 0374303d7d969c76cd6771b2b18bf46f94df8d04..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/profiling/constants.py +++ /dev/null @@ -1,47 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -######################################### -# flops profiler -######################################### -# Flops profiler. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -FLOPS_PROFILER_FORMAT = ''' -flops profiler should be enabled as: -"session_params": { - "flops_profiler": { - "enabled": true, - "recompute_fwd_factor": 0.0, - "profile_step": 1, - "module_depth": -1, - "top_modules": 3, - "detailed": true, - "output_file": null - } -} -''' - -FLOPS_PROFILER = "flops_profiler" - -FLOPS_PROFILER_ENABLED = "enabled" -FLOPS_PROFILER_ENABLED_DEFAULT = False - -FLOPS_PROFILER_RECOMPUTE_FWD_FACTOR = "recompute_fwd_factor" -FLOPS_PROFILER_RECOMPUTE_FWD_FACTOR_DEFAULT = 0.0 - -FLOPS_PROFILER_PROFILE_STEP = "profile_step" -FLOPS_PROFILER_PROFILE_STEP_DEFAULT = 1 - -FLOPS_PROFILER_MODULE_DEPTH = "module_depth" -FLOPS_PROFILER_MODULE_DEPTH_DEFAULT = -1 - -FLOPS_PROFILER_TOP_MODULES = "top_modules" -FLOPS_PROFILER_TOP_MODULES_DEFAULT = 1 - -FLOPS_PROFILER_DETAILED = "detailed" -FLOPS_PROFILER_DETAILED_DEFAULT = True - -FLOPS_PROFILER_OUTPUT_FILE = "output_file" -FLOPS_PROFILER_OUTPUT_FILE_DEFAULT = None diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__init__.py deleted file mode 100644 index 980d8a5e92fc4c3491f822479d05c5e2776233df..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .profiler import * diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 1bdc72be53afd5521f231e53f84455751f809815..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__pycache__/profiler.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__pycache__/profiler.cpython-310.pyc deleted file mode 100644 index 7143ac310b1a352fe164b889b0f13bad0ae3a52c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/__pycache__/profiler.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/profiler.py b/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/profiler.py deleted file mode 100644 index 1f051077c36c24b272f031d629ac5ee84f88aa3f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/profiling/flops_profiler/profiler.py +++ /dev/null @@ -1,1297 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import time -import torch -import torch.nn as nn -import torch.nn.functional as F -from functools import partial -from typing import List, Optional -from collections import OrderedDict -import numpy as np -from deepspeed.accelerator import get_accelerator -from deepspeed.utils import logger -from deepspeed.moe.layer import MoE -from deepspeed.utils.timer import FORWARD_GLOBAL_TIMER, BACKWARD_GLOBAL_TIMER, STEP_GLOBAL_TIMER -from deepspeed.utils.torch import required_torch_version -import einops - -Tensor = torch.Tensor - -module_flop_count = [] -module_mac_count = [] -old_functions = {} - -DEFAULT_PRECISION = 2 - - -class FlopsProfiler(object): - """Measures the latency, number of estimated floating-point operations and parameters of each module in a PyTorch model. - - The flops-profiler profiles the forward pass of a PyTorch model and prints the model graph with the measured profile attached to each module. It shows how latency, flops and parameters are spent in the model and which modules or layers could be the bottleneck. It also outputs the names of the top k modules in terms of aggregated latency, flops, and parameters at depth l with k and l specified by the user. The output profile is computed for each batch of input. - The DeepSpeed flops profiler can be used with the DeepSpeed runtime or as a standalone package. - When using DeepSpeed for model training, the flops profiler can be configured in the deepspeed_config file and no user code change is required. - - If using the profiler as a standalone package, one imports the flops_profiler package and use the APIs. - - Here is an example for usage in a typical training workflow: - - .. code-block:: python - - model = Model() - prof = FlopsProfiler(model) - - for step, batch in enumerate(data_loader): - if step == profile_step: - prof.start_profile() - - loss = model(batch) - - if step == profile_step: - flops = prof.get_total_flops(as_string=True) - params = prof.get_total_params(as_string=True) - prof.print_model_profile(profile_step=profile_step) - prof.end_profile() - - loss.backward() - optimizer.step() - - To profile a trained model in inference, use the `get_model_profile` API. - - Args: - object (torch.nn.Module): The PyTorch model to profile. - """ - - def __init__(self, model, ds_engine=None, recompute_fwd_factor=0.0): - self.model = model - self.ds_engine = ds_engine - self.recompute_fwd_factor = recompute_fwd_factor - self.started = False - self.func_patched = False - - def start_profile(self, ignore_list=None): - """Starts profiling. - - Extra attributes are added recursively to all the modules and the profiled torch.nn.functionals are monkey patched. - - Args: - ignore_list (list, optional): the list of modules to ignore while profiling. Defaults to None. - """ - logger.info("Flops profiler started") - self.reset_profile() - _patch_functionals() - _patch_tensor_methods() - _patch_miscellaneous_operations() - - def register_module_hooks(module, ignore_list): - if ignore_list and type(module) in ignore_list: - return - - # if computing the flops of a module directly - if type(module) in MODULE_HOOK_MAPPING: - if not hasattr(module, "__flops_handle__"): - module.__flops_handle__ = module.register_forward_hook(MODULE_HOOK_MAPPING[type(module)]) - return - - # if computing the flops of the functionals in a module - def pre_hook(module, input): - module_flop_count.append([]) - module_mac_count.append([]) - - if not hasattr(module, "__pre_hook_handle__"): - module.__pre_hook_handle__ = module.register_forward_pre_hook(pre_hook) - - def post_hook(module, input, output): - if module_flop_count: - module.__flops__ += sum([elem[1] for elem in module_flop_count[-1]]) - module_flop_count.pop() - module.__macs__ += sum([elem[1] for elem in module_mac_count[-1]]) - module_mac_count.pop() - - if not hasattr(module, "__post_hook_handle__"): - module.__post_hook_handle__ = module.register_forward_hook(post_hook) - - def start_time_hook(module, input): - get_accelerator().synchronize() - module.__start_time__ = time.time() - - if not hasattr(module, "__start_time_hook_handle__"): - module.__start_time_hook_handle__ = module.register_forward_pre_hook(start_time_hook) - - def end_time_hook(module, input, output): - get_accelerator().synchronize() - module.__duration__ += time.time() - module.__start_time__ - - if not hasattr(module, "__end_time_hook_handle__"): - module.__end_time_hook_handle__ = module.register_forward_hook(end_time_hook) - - self.model.apply(partial(register_module_hooks, ignore_list=ignore_list)) - self.started = True - self.func_patched = True - - def stop_profile(self): - """Stop profiling. - - All torch.nn.functionals are restored to their originals. - """ - if self.started and self.func_patched: - _reload_functionals() - _reload_tensor_methods() - _reload_miscellaneous_operations() - self.func_patched = False - - def remove_profile_attrs(module): - if hasattr(module, "__pre_hook_handle__"): - module.__pre_hook_handle__.remove() - del module.__pre_hook_handle__ - if hasattr(module, "__post_hook_handle__"): - module.__post_hook_handle__.remove() - del module.__post_hook_handle__ - if hasattr(module, "__flops_handle__"): - module.__flops_handle__.remove() - del module.__flops_handle__ - if hasattr(module, "__start_time_hook_handle__"): - module.__start_time_hook_handle__.remove() - del module.__start_time_hook_handle__ - if hasattr(module, "__end_time_hook_handle__"): - module.__end_time_hook_handle__.remove() - del module.__end_time_hook_handle__ - - self.model.apply(remove_profile_attrs) - - def reset_profile(self): - """Resets the profiling. - - Adds or resets the extra attributes. - """ - - def get_param_count_and_ep(param): - """ - Return the number of parameters in the layer, whether the layer is an MoE layer, - and its expert parallelism size if so - """ - prefix = 'ep_size_' - offset = len(prefix) - expert_parallelism = 0 - if getattr(param, "group_name", "").startswith(prefix): - try: - expert_parallelism = int(param.group_name[offset:]) - except ValueError: - pass - return param.numel(), expert_parallelism, param.element_size() - - def add_or_reset_attrs(module): - module.__flops__ = 0 - module.__macs__ = 0 - module.__params__ = module.__expert_params__ = module.__model_expert_params__ = 0 - parameters = (get_param_count_and_ep(p) for p in module.parameters()) - for num_params, expert_parallelism, per_param_size in parameters: - params = num_params if not expert_parallelism else 0 - expert_params = num_params if expert_parallelism else 0 - # number of expert parameters taking into account other expert parallel groups - model_expert_params = num_params * expert_parallelism - module.__params__ += params - module.__expert_params__ += expert_params - module.__model_expert_params__ += model_expert_params - module.__start_time__ = 0 - module.__duration__ = 0 - - self.model.apply(add_or_reset_attrs) - - def end_profile(self): - """Ends profiling. - - The added attributes and handles are removed recursively on all the modules. - """ - if not self.started: - return - self.stop_profile() - self.started = False - - def remove_profile_attrs(module): - if hasattr(module, "__flops__"): - del module.__flops__ - if hasattr(module, "__macs__"): - del module.__macs__ - if hasattr(module, "__params__"): - del module.__params__ - if hasattr(module, "__expert_params__"): - del module.__expert_params__ - if hasattr(module, "__model_expert_params__"): - del module.__model_expert_params__ - if hasattr(module, "__start_time__"): - del module.__start_time__ - if hasattr(module, "__duration__"): - del module.__duration__ - - self.model.apply(remove_profile_attrs) - logger.info("Flops profiler finished") - - def get_total_flops(self, as_string=False): - """Returns the total flops of the model. - - Args: - as_string (bool, optional): whether to output the flops as string. Defaults to False. - - Returns: - The number of multiply-accumulate operations of the model forward pass. - """ - total_flops = get_module_flops(self.model) - return number_to_string(total_flops) if as_string else total_flops - - def get_total_macs(self, as_string=False): - """Returns the total MACs of the model. - - Args: - as_string (bool, optional): whether to output the flops as string. Defaults to False. - - Returns: - The number of multiply-accumulate operations of the model forward pass. - """ - total_macs = get_module_macs(self.model) - return macs_to_string(total_macs) if as_string else total_macs - - def get_total_duration(self, as_string=False): - """Returns the total duration of the model forward pass. - - Args: - as_string (bool, optional): whether to output the duration as string. Defaults to False. - - Returns: - The latency of the model forward pass. - """ - total_duration = get_module_duration(self.model) - return duration_to_string(total_duration) if as_string else total_duration - - def get_total_params(self, as_string=False): - """Returns the total number of parameters stored per rank. - - Args: - as_string (bool, optional): whether to output the parameters as string. Defaults to False. - - Returns: - The total number of parameters stored per rank. - """ - total_params = self.model.__expert_params__ + self.model.__params__ - return params_to_string(total_params) if as_string else total_params - - def is_expert_tensor_parallelism_enabled(self): - for _, module in self.model.named_modules(): - if isinstance(module, MoE) and hasattr(module, 'enable_expert_tensor_parallelism'): - return module.enable_expert_tensor_parallelism - return False - - def print_model_profile(self, profile_step=1, module_depth=-1, top_modules=1, detailed=True, output_file=None): - """Prints the model graph with the measured profile attached to each module. - - Args: - profile_step (int, optional): The global training step at which to profile. Note that warm up steps are needed for accurate time measurement. - module_depth (int, optional): The depth of the model to which to print the aggregated module information. When set to -1, it prints information from the top to the innermost modules (the maximum depth). - top_modules (int, optional): Limits the aggregated profile output to the number of top modules specified. - detailed (bool, optional): Whether to print the detailed model profile. - output_file (str, optional): Path to the output file. If None, the profiler prints to stdout. - """ - if not self.started: - return - import sys - import os.path - original_stdout = None - f = None - if output_file and output_file != "": - dir_path = os.path.dirname(os.path.abspath(output_file)) - if not os.path.exists(dir_path): - os.makedirs(dir_path) - original_stdout = sys.stdout - f = open(output_file, "w") - sys.stdout = f - - total_flops = self.get_total_flops() - total_macs = self.get_total_macs() - total_duration = self.get_total_duration() - total_params = self.get_total_params() - expert_tensor_parallelism = None # silence the linters - total_model_expert_params = total_model_nonexpert_params = 0 - if self.ds_engine: - total_model_nonexpert_params = self.model.__params__ * self.ds_engine.mp_world_size - if self.ds_engine.has_moe_layers: - expert_tensor_parallelism = self.ds_engine.mp_world_size if self.is_expert_tensor_parallelism_enabled( - ) else 1 - total_model_expert_params = self.model.__model_expert_params__ * expert_tensor_parallelism - - self.flops = total_flops - self.macs = total_macs - self.params = total_params - - print("\n-------------------------- DeepSpeed Flops Profiler --------------------------") - print(f'Profile Summary at step {profile_step}:') - print("Notations:\n" - "data parallel size (dp_size), model parallel size(mp_size),\n" - "number of parameters (params), number of multiply-accumulate operations(MACs),\n" - "number of floating-point operations (flops), floating-point operations per second (FLOPS),\n" - "fwd latency (forward propagation latency), bwd latency (backward propagation latency),\n" - "step (weights update latency), iter latency (sum of fwd, bwd and step latency)\n") - line_fmt = '{:<70} {:<8}' - if self.ds_engine: - print(line_fmt.format('world size: ', self.ds_engine.world_size)) - print(line_fmt.format('data parallel size: ', self.ds_engine.dp_world_size)) - print(line_fmt.format('model parallel size: ', self.ds_engine.mp_world_size)) - print(line_fmt.format('batch size per GPU: ', self.ds_engine.train_micro_batch_size_per_gpu())) - if self.ds_engine.has_moe_layers: - print(line_fmt.format('expert tensor parallelism enabled: ', expert_tensor_parallelism > 1)) - - print(line_fmt.format('params per GPU: ', params_to_string(total_params))) - if total_model_expert_params > 0: - print( - line_fmt.format('params of model: ', - params_to_string(total_model_nonexpert_params + total_model_expert_params))) - print(line_fmt.format(' non-expert params of model: ', params_to_string(total_model_nonexpert_params))) - print(line_fmt.format(' expert params of model: ', params_to_string(total_model_expert_params))) - else: - print( - line_fmt.format('params of model = params per GPU * mp_size: ', - params_to_string(total_model_nonexpert_params))) - - print(line_fmt.format('fwd MACs per GPU: ', macs_to_string(total_macs))) - - print(line_fmt.format('fwd flops per GPU: ', number_to_string(total_flops))) - - print( - line_fmt.format('fwd flops of model = fwd flops per GPU * mp_size: ', - number_to_string(total_flops * (self.ds_engine.mp_world_size if self.ds_engine else 1)))) - - fwd_latency = self.get_total_duration() - if self.ds_engine and self.ds_engine.wall_clock_breakdown(): - fwd_latency = self.ds_engine.timers(FORWARD_GLOBAL_TIMER).elapsed(False) / 1000.0 - print(line_fmt.format('fwd latency: ', duration_to_string(fwd_latency))) - print( - line_fmt.format('fwd FLOPS per GPU = fwd flops per GPU / fwd latency: ', - flops_to_string(total_flops / fwd_latency))) - - if self.ds_engine and self.ds_engine.wall_clock_breakdown(): - bwd_factor = 2 + self.recompute_fwd_factor - bwd_latency = self.ds_engine.timers(BACKWARD_GLOBAL_TIMER).elapsed(False) / 1000.0 - step_latency = self.ds_engine.timers(STEP_GLOBAL_TIMER).elapsed(False) / 1000.0 - print(line_fmt.format('bwd latency: ', duration_to_string(bwd_latency))) - print( - line_fmt.format(f'bwd FLOPS per GPU = {bwd_factor:g} * fwd flops per GPU / bwd latency: ', - flops_to_string(bwd_factor * total_flops / bwd_latency))) - print( - line_fmt.format( - f'fwd+bwd FLOPS per GPU = {bwd_factor + 1:g} * fwd flops per GPU / (fwd+bwd latency): ', - flops_to_string((bwd_factor + 1) * total_flops / (fwd_latency + bwd_latency)))) - - print(line_fmt.format('step latency: ', duration_to_string(step_latency))) - - iter_latency = fwd_latency + bwd_latency + step_latency - print(line_fmt.format('iter latency: ', duration_to_string(iter_latency))) - print( - line_fmt.format(f'FLOPS per GPU = {bwd_factor + 1:g} * fwd flops per GPU / iter latency: ', - flops_to_string((bwd_factor + 1) * total_flops / iter_latency))) - - samples_per_iter = self.ds_engine.train_micro_batch_size_per_gpu() * self.ds_engine.world_size - print(line_fmt.format('samples/second: ', round(samples_per_iter / iter_latency, DEFAULT_PRECISION))) - - def flops_repr(module): - params = module.__params__ + module.__expert_params__ - flops = get_module_flops(module) - macs = get_module_macs(module) - duration = get_module_duration(module) - items = [ - "{} = {:g}% Params".format( - params_to_string(params), - round(100 * params / total_params, DEFAULT_PRECISION) if total_params else 0), - "{} = {:g}% MACs".format(macs_to_string(macs), - round(100 * macs / total_macs, DEFAULT_PRECISION) if total_macs else 0), - "{} = {:g}% latency".format( - duration_to_string(duration), - round(100 * duration / total_duration, DEFAULT_PRECISION) if total_duration else 0), - flops_to_string(round(flops / duration, DEFAULT_PRECISION) if duration else 0), - ] - original_extra_repr = module.original_extra_repr() - if original_extra_repr: - items.append(original_extra_repr) - return ", ".join(items) - - def add_extra_repr(module): - flops_extra_repr = flops_repr.__get__(module) - if module.extra_repr != flops_extra_repr: - module.original_extra_repr = module.extra_repr - module.extra_repr = flops_extra_repr - assert module.extra_repr != module.original_extra_repr - - def del_extra_repr(module): - if hasattr(module, "original_extra_repr"): - module.extra_repr = module.original_extra_repr - del module.original_extra_repr - - self.model.apply(add_extra_repr) - - print("\n----------------------------- Aggregated Profile per GPU -----------------------------") - self.print_model_aggregated_profile(module_depth=module_depth, top_modules=top_modules) - - if detailed: - print("\n------------------------------ Detailed Profile per GPU ------------------------------") - print( - "Each module profile is listed after its name in the following order: \nparams, percentage of total params, MACs, percentage of total MACs, fwd latency, percentage of total fwd latency, fwd FLOPS" - ) - print( - "\nNote: 1. A module can have torch.nn.module or torch.nn.functional to compute logits (e.g. CrossEntropyLoss). They are not counted as submodules, thus not to be printed out. However they make up the difference between a parent's MACs (or latency) and the sum of its submodules'.\n2. Number of floating-point operations is a theoretical estimation, thus FLOPS computed using that could be larger than the maximum system throughput.\n3. The fwd latency listed in the top module's profile is directly captured at the module forward function in PyTorch, thus it's less than the fwd latency shown above which is captured in DeepSpeed.\n" - ) - print(self.model) - - self.model.apply(del_extra_repr) - - print("------------------------------------------------------------------------------") - - if output_file: - sys.stdout = original_stdout - f.close() - - def print_model_aggregated_profile(self, module_depth=-1, top_modules=1): - """Prints the names of the top top_modules modules in terms of aggregated time, flops, and parameters at depth module_depth. - - Args: - module_depth (int, optional): the depth of the modules to show. Defaults to -1 (the innermost modules). - top_modules (int, optional): the number of top modules to show. Defaults to 1. - """ - info = {} - if not hasattr(self.model, "__flops__"): - print("no __flops__ attribute in the model, call this function after start_profile and before end_profile") - return - - def walk_module(module, curr_depth, info): - if curr_depth not in info: - info[curr_depth] = {} - if module.__class__.__name__ not in info[curr_depth]: - info[curr_depth][module.__class__.__name__] = [ - 0, - 0, - 0, - ] # macs, params, time - info[curr_depth][module.__class__.__name__][0] += get_module_macs(module) - info[curr_depth][module.__class__.__name__][1] += module.__params__ + module.__expert_params__ - info[curr_depth][module.__class__.__name__][2] += get_module_duration(module) - has_children = len(module._modules.items()) != 0 - if has_children: - for child in module.children(): - walk_module(child, curr_depth + 1, info) - - walk_module(self.model, 0, info) - - depth = module_depth - if module_depth == -1: - depth = len(info) - 1 - - print(f'Top {top_modules} modules in terms of params, MACs or fwd latency at different model depths:') - - for d in range(depth): - num_items = min(top_modules, len(info[d])) - - sort_macs = { - k: macs_to_string(v[0]) - for k, v in sorted(info[d].items(), key=lambda item: item[1][0], reverse=True)[:num_items] - } - sort_params = { - k: params_to_string(v[1]) - for k, v in sorted(info[d].items(), key=lambda item: item[1][1], reverse=True)[:num_items] - } - sort_time = { - k: duration_to_string(v[2]) - for k, v in sorted(info[d].items(), key=lambda item: item[1][2], reverse=True)[:num_items] - } - - print(f"depth {d}:") - print(f" params - {sort_params}") - print(f" MACs - {sort_macs}") - print(f" fwd latency - {sort_time}") - - -def _prod(dims): - p = 1 - for v in dims: - p *= v - return p - - -def _linear_flops_compute(input, weight, bias=None): - out_features = weight.shape[0] - macs = input.numel() * out_features - return 2 * macs, macs - - -def _relu_flops_compute(input, inplace=False): - return input.numel(), 0 - - -def _prelu_flops_compute(input: Tensor, weight: Tensor): - return input.numel(), 0 - - -def _elu_flops_compute(input: Tensor, alpha: float = 1.0, inplace: bool = False): - return input.numel(), 0 - - -def _leaky_relu_flops_compute(input: Tensor, negative_slope: float = 0.01, inplace: bool = False): - return input.numel(), 0 - - -def _relu6_flops_compute(input: Tensor, inplace: bool = False): - return input.numel(), 0 - - -def _silu_flops_compute(input: Tensor, inplace: bool = False): - return input.numel(), 0 - - -def _gelu_flops_compute(input, **kwargs): - return input.numel(), 0 - - -def _pool_flops_compute(input, - kernel_size, - stride=None, - padding=0, - dilation=None, - ceil_mode=False, - count_include_pad=True, - divisor_override=None, - return_indices=None): - return input.numel(), 0 - - -def _conv_flops_compute(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): - assert weight.shape[1] * groups == input.shape[1] - - batch_size = input.shape[0] - in_channels = input.shape[1] - out_channels = weight.shape[0] - kernel_dims = list(weight.shape[2:]) - input_dims = list(input.shape[2:]) - - length = len(input_dims) - - strides = stride if type(stride) is tuple else (stride, ) * length - dilations = dilation if type(dilation) is tuple else (dilation, ) * length - if isinstance(padding, str): - if padding == 'valid': - paddings = (0, ) * length - elif padding == 'same': - paddings = () - for d, k in zip(dilations, kernel_dims): - total_padding = d * (k - 1) - paddings += (total_padding // 2, ) - elif isinstance(padding, tuple): - paddings = padding - else: - paddings = (padding, ) * length - - output_dims = [] - for idx, input_dim in enumerate(input_dims): - output_dim = (input_dim + 2 * paddings[idx] - (dilations[idx] * - (kernel_dims[idx] - 1) + 1)) // strides[idx] + 1 - output_dims.append(output_dim) - - filters_per_channel = out_channels // groups - conv_per_position_macs = int(_prod(kernel_dims)) * in_channels * filters_per_channel - active_elements_count = batch_size * int(_prod(output_dims)) - overall_conv_macs = conv_per_position_macs * active_elements_count - overall_conv_flops = 2 * overall_conv_macs - - bias_flops = 0 - if bias is not None: - bias_flops = out_channels * active_elements_count - - return int(overall_conv_flops + bias_flops), int(overall_conv_macs) - - -def _conv_trans_flops_compute( - input, - weight, - bias=None, - stride=1, - padding=0, - output_padding=0, - groups=1, - dilation=1, -): - batch_size = input.shape[0] - in_channels = input.shape[1] - out_channels = weight.shape[1] - kernel_dims = list(weight.shape[2:]) - input_dims = list(input.shape[2:]) - - length = len(input_dims) - - paddings = padding if type(padding) is tuple else (padding, ) * length - strides = stride if type(stride) is tuple else (stride, ) * length - dilations = dilation if type(dilation) is tuple else (dilation, ) * length - - output_dims = [] - for idx, input_dim in enumerate(input_dims): - - output_dim = (input_dim + 2 * paddings[idx] - (dilations[idx] * - (kernel_dims[idx] - 1) + 1)) // strides[idx] + 1 - output_dims.append(output_dim) - - paddings = padding if type(padding) is tuple else (padding, padding) - strides = stride if type(stride) is tuple else (stride, stride) - dilations = dilation if type(dilation) is tuple else (dilation, dilation) - - filters_per_channel = out_channels // groups - conv_per_position_macs = int(_prod(kernel_dims)) * in_channels * filters_per_channel - active_elements_count = batch_size * int(_prod(input_dims)) - overall_conv_macs = conv_per_position_macs * active_elements_count - overall_conv_flops = 2 * overall_conv_macs - - bias_flops = 0 - if bias is not None: - bias_flops = out_channels * batch_size * int(_prod(output_dims)) - - return int(overall_conv_flops + bias_flops), int(overall_conv_macs) - - -def _batch_norm_flops_compute( - input, - running_mean, - running_var, - weight=None, - bias=None, - training=False, - momentum=0.1, - eps=1e-05, -): - has_affine = weight is not None - if training: - # estimation - return input.numel() * (5 if has_affine else 4), 0 - flops = input.numel() * (2 if has_affine else 1) - return flops, 0 - - -def _layer_norm_flops_compute( - input: Tensor, - normalized_shape: List[int], - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - eps: float = 1e-5, -): - has_affine = weight is not None - # estimation - return input.numel() * (5 if has_affine else 4), 0 - - -def _group_norm_flops_compute(input: Tensor, - num_groups: int, - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - eps: float = 1e-5): - has_affine = weight is not None - # estimation - return input.numel() * (5 if has_affine else 4), 0 - - -def _instance_norm_flops_compute( - input: Tensor, - running_mean: Optional[Tensor] = None, - running_var: Optional[Tensor] = None, - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - use_input_stats: bool = True, - momentum: float = 0.1, - eps: float = 1e-5, -): - has_affine = weight is not None - # estimation - return input.numel() * (5 if has_affine else 4), 0 - - -def _upsample_flops_compute(*args, **kwargs): - input = args[0] - size = kwargs.get('size', None) - if size is None and len(args) > 1: - size = args[1] - - if size is not None: - if isinstance(size, tuple) or isinstance(size, list): - return int(_prod(size)), 0 - else: - return int(size), 0 - - scale_factor = kwargs.get('scale_factor', None) - if scale_factor is None and len(args) > 2: - scale_factor = args[2] - assert scale_factor is not None, "either size or scale_factor should be defined" - - flops = input.numel() - if isinstance(scale_factor, (list, tuple)): - # see documentation of `F.interpolate` - # the spatial dims are defined as the last `n-2` dims of the tensor - assert len(scale_factor) == input.ndim - 2 - flops *= _prod(scale_factor) - else: - flops *= scale_factor**(input.ndim - 2) - - return int(flops), 0 - - -def _softmax_flops_compute(input, dim=None, _stacklevel=3, dtype=None): - return input.numel(), 0 - - -def _embedding_flops_compute( - input, - weight, - padding_idx=None, - max_norm=None, - norm_type=2.0, - scale_grad_by_freq=False, - sparse=False, -): - return 0, 0 - - -def _dropout_flops_compute(input, p=0.5, training=True, inplace=False): - return 0, 0 - - -def _matmul_flops_compute(input, other, *, out=None): - """ - Count flops for the matmul operation. - """ - macs = _prod(input.shape) * other.shape[-1] - return 2 * macs, macs - - -def _addmm_flops_compute(input, mat1, mat2, *, beta=1, alpha=1, out=None): - """ - Count flops for the addmm operation. - """ - macs = _prod(mat1.shape) * mat2.shape[-1] - return 2 * macs + _prod(input.shape), macs - - -def _einsum_flops_compute(equation, *operands): - """ - Count flops for the einsum operation. - """ - equation = equation.replace(" ", "") - input_shapes = [o.shape for o in operands] - - # Re-map equation so that same equation with different alphabet - # representations will look the same. - letter_order = OrderedDict((k, 0) for k in equation if k.isalpha()).keys() - mapping = {ord(x): 97 + i for i, x in enumerate(letter_order)} - equation = equation.translate(mapping) - - np_arrs = [np.zeros(s) for s in input_shapes] - optim = np.einsum_path(equation, *np_arrs, optimize="optimal")[1] - for line in optim.split("\n"): - if "optimized flop" in line.lower(): - flop = int(float(line.split(":")[-1])) - return flop, 0 - raise NotImplementedError("Unsupported einsum operation.") - - -def _einops_einsum_flops_compute(*args): - """ - Count flops for the einops.einsum operation. - """ - *operands, equation = args - input_shapes = [o.shape for o in operands] - - # Re-map equation so that same equation with different alphabet - # representations will look the same. - letter_order = OrderedDict((k, 0) for k in equation if k.isalpha()).keys() - mapping = {ord(x): 97 + i for i, x in enumerate(letter_order)} - equation = equation.translate(mapping) - - np_arrs = [np.zeros(s) for s in input_shapes] - optim = np.einsum_path(equation, *np_arrs, optimize="optimal")[1] - for line in optim.split("\n"): - if "optimized flop" in line.lower(): - flop = int(float(line.split(":")[-1])) - return flop, 0 - - raise NotImplementedError("Unsupported einops.einsum operation.") - - -def _tensor_addmm_flops_compute(self, mat1, mat2, *, beta=1, alpha=1, out=None): - """ - Count flops for the tensor addmm operation. - """ - macs = _prod(mat1.shape) * mat2.shape[-1] - return 2 * macs + _prod(self.shape), macs - - -def _mul_flops_compute(input, other, *, out=None): - return _elementwise_flops_compute(input, other) - - -def _add_flops_compute(input, other, *, alpha=1, out=None): - return _elementwise_flops_compute(input, other) - - -def _elementwise_flops_compute(input, other): - if not torch.is_tensor(input): - if torch.is_tensor(other): - return _prod(other.shape), 0 - else: - return 1, 0 - elif not torch.is_tensor(other): - return _prod(input.shape), 0 - else: - dim_input = len(input.shape) - dim_other = len(other.shape) - max_dim = max(dim_input, dim_other) - - final_shape = [] - for i in range(max_dim): - in_i = input.shape[i] if i < dim_input else 1 - ot_i = other.shape[i] if i < dim_other else 1 - if in_i > ot_i: - final_shape.append(in_i) - else: - final_shape.append(ot_i) - flops = _prod(final_shape) - return flops, 0 - - -def _attn_flops_compute(q, k, v, *args, **kwargs): - """ - Count flops for the scaled_dot_product_attention operation. - """ - macs = _prod(q.shape) * k.shape[-2] - macs += _prod(q.shape[:-1]) * k.shape[-2] * v.shape[-1] - return 2 * macs, macs - - -def wrapFunc(func, funcFlopCompute): - oldFunc = func - name = func.__str__ - old_functions[name] = oldFunc - - def newFunc(*args, **kwds): - flops, macs = funcFlopCompute(*args, **kwds) - if module_flop_count: - module_flop_count[-1].append((name, flops)) - if module_mac_count and macs: - module_mac_count[-1].append((name, macs)) - return oldFunc(*args, **kwds) - - newFunc.__str__ = func.__str__ - - return newFunc - - -def _patch_functionals(): - # FC - F.linear = wrapFunc(F.linear, _linear_flops_compute) - - # convolutions - F.conv1d = wrapFunc(F.conv1d, _conv_flops_compute) - F.conv2d = wrapFunc(F.conv2d, _conv_flops_compute) - F.conv3d = wrapFunc(F.conv3d, _conv_flops_compute) - - # conv transposed - F.conv_transpose1d = wrapFunc(F.conv_transpose1d, _conv_trans_flops_compute) - F.conv_transpose2d = wrapFunc(F.conv_transpose2d, _conv_trans_flops_compute) - F.conv_transpose3d = wrapFunc(F.conv_transpose3d, _conv_trans_flops_compute) - - # activations - F.relu = wrapFunc(F.relu, _relu_flops_compute) - F.prelu = wrapFunc(F.prelu, _prelu_flops_compute) - F.elu = wrapFunc(F.elu, _elu_flops_compute) - F.leaky_relu = wrapFunc(F.leaky_relu, _leaky_relu_flops_compute) - F.relu6 = wrapFunc(F.relu6, _relu6_flops_compute) - if hasattr(F, "silu"): - F.silu = wrapFunc(F.silu, _silu_flops_compute) - F.gelu = wrapFunc(F.gelu, _gelu_flops_compute) - - # Normalizations - F.batch_norm = wrapFunc(F.batch_norm, _batch_norm_flops_compute) - F.layer_norm = wrapFunc(F.layer_norm, _layer_norm_flops_compute) - F.instance_norm = wrapFunc(F.instance_norm, _instance_norm_flops_compute) - F.group_norm = wrapFunc(F.group_norm, _group_norm_flops_compute) - - # poolings - F.avg_pool1d = wrapFunc(F.avg_pool1d, _pool_flops_compute) - F.avg_pool2d = wrapFunc(F.avg_pool2d, _pool_flops_compute) - F.avg_pool3d = wrapFunc(F.avg_pool3d, _pool_flops_compute) - F.max_pool1d = wrapFunc(F.max_pool1d, _pool_flops_compute) - F.max_pool2d = wrapFunc(F.max_pool2d, _pool_flops_compute) - F.max_pool3d = wrapFunc(F.max_pool3d, _pool_flops_compute) - F.adaptive_avg_pool1d = wrapFunc(F.adaptive_avg_pool1d, _pool_flops_compute) - F.adaptive_avg_pool2d = wrapFunc(F.adaptive_avg_pool2d, _pool_flops_compute) - F.adaptive_avg_pool3d = wrapFunc(F.adaptive_avg_pool3d, _pool_flops_compute) - F.adaptive_max_pool1d = wrapFunc(F.adaptive_max_pool1d, _pool_flops_compute) - F.adaptive_max_pool2d = wrapFunc(F.adaptive_max_pool2d, _pool_flops_compute) - F.adaptive_max_pool3d = wrapFunc(F.adaptive_max_pool3d, _pool_flops_compute) - - # upsample - F.upsample = wrapFunc(F.upsample, _upsample_flops_compute) - F.interpolate = wrapFunc(F.interpolate, _upsample_flops_compute) - - # softmax - F.softmax = wrapFunc(F.softmax, _softmax_flops_compute) - - # embedding - F.embedding = wrapFunc(F.embedding, _embedding_flops_compute) - - # attn - scaled_dot_product_attention added in torch 2.0+ - if required_torch_version(min_version=2.0): - F.scaled_dot_product_attention = wrapFunc(F.scaled_dot_product_attention, _attn_flops_compute) - - -def _patch_tensor_methods(): - torch.matmul = wrapFunc(torch.matmul, _matmul_flops_compute) - torch.Tensor.matmul = wrapFunc(torch.Tensor.matmul, _matmul_flops_compute) - torch.Tensor.__matmul__ = wrapFunc(torch.Tensor.__matmul__, _matmul_flops_compute) - torch.mm = wrapFunc(torch.mm, _matmul_flops_compute) - torch.Tensor.mm = wrapFunc(torch.Tensor.mm, _matmul_flops_compute) - torch.bmm = wrapFunc(torch.bmm, _matmul_flops_compute) - torch.Tensor.bmm = wrapFunc(torch.Tensor.bmm, _matmul_flops_compute) - - torch.addmm = wrapFunc(torch.addmm, _addmm_flops_compute) - torch.Tensor.addmm = wrapFunc(torch.Tensor.addmm, _tensor_addmm_flops_compute) - - torch.mul = wrapFunc(torch.mul, _mul_flops_compute) - torch.Tensor.mul = wrapFunc(torch.Tensor.mul, _mul_flops_compute) - - torch.add = wrapFunc(torch.add, _add_flops_compute) - torch.Tensor.add = wrapFunc(torch.Tensor.add, _add_flops_compute) - - torch.einsum = wrapFunc(torch.einsum, _einsum_flops_compute) - - torch.baddbmm = wrapFunc(torch.baddbmm, _tensor_addmm_flops_compute) - - -def _patch_miscellaneous_operations(): - einops.einsum = wrapFunc(einops.einsum, _einops_einsum_flops_compute) - - -def _reload_functionals(): - # torch.nn.functional does not support importlib.reload() - F.linear = old_functions[F.linear.__str__] - F.conv1d = old_functions[F.conv1d.__str__] - F.conv2d = old_functions[F.conv2d.__str__] - F.conv3d = old_functions[F.conv3d.__str__] - F.conv_transpose1d = old_functions[F.conv_transpose1d.__str__] - F.conv_transpose2d = old_functions[F.conv_transpose2d.__str__] - F.conv_transpose3d = old_functions[F.conv_transpose3d.__str__] - F.relu = old_functions[F.relu.__str__] - F.prelu = old_functions[F.prelu.__str__] - F.elu = old_functions[F.elu.__str__] - F.leaky_relu = old_functions[F.leaky_relu.__str__] - F.relu6 = old_functions[F.relu6.__str__] - if hasattr(F, "silu"): - F.silu = old_functions[F.silu.__str__] - F.gelu = old_functions[F.gelu.__str__] - F.batch_norm = old_functions[F.batch_norm.__str__] - F.layer_norm = old_functions[F.layer_norm.__str__] - F.instance_norm = old_functions[F.instance_norm.__str__] - F.group_norm = old_functions[F.group_norm.__str__] - F.avg_pool1d = old_functions[F.avg_pool1d.__str__] - F.avg_pool2d = old_functions[F.avg_pool2d.__str__] - F.avg_pool3d = old_functions[F.avg_pool3d.__str__] - F.max_pool1d = old_functions[F.max_pool1d.__str__] - F.max_pool2d = old_functions[F.max_pool2d.__str__] - F.max_pool3d = old_functions[F.max_pool3d.__str__] - F.adaptive_avg_pool1d = old_functions[F.adaptive_avg_pool1d.__str__] - F.adaptive_avg_pool2d = old_functions[F.adaptive_avg_pool2d.__str__] - F.adaptive_avg_pool3d = old_functions[F.adaptive_avg_pool3d.__str__] - F.adaptive_max_pool1d = old_functions[F.adaptive_max_pool1d.__str__] - F.adaptive_max_pool2d = old_functions[F.adaptive_max_pool2d.__str__] - F.adaptive_max_pool3d = old_functions[F.adaptive_max_pool3d.__str__] - F.upsample = old_functions[F.upsample.__str__] - F.interpolate = old_functions[F.interpolate.__str__] - F.softmax = old_functions[F.softmax.__str__] - F.embedding = old_functions[F.embedding.__str__] - - -def _reload_tensor_methods(): - torch.matmul = old_functions[torch.matmul.__str__] - torch.Tensor.matmul = old_functions[torch.Tensor.matmul.__str__] - torch.mm = old_functions[torch.mm.__str__] - torch.Tensor.mm = old_functions[torch.Tensor.mm.__str__] - torch.bmm = old_functions[torch.matmul.__str__] - torch.Tensor.bmm = old_functions[torch.Tensor.bmm.__str__] - torch.addmm = old_functions[torch.addmm.__str__] - torch.Tensor.addmm = old_functions[torch.Tensor.addmm.__str__] - torch.mul = old_functions[torch.mul.__str__] - torch.Tensor.mul = old_functions[torch.Tensor.mul.__str__] - torch.add = old_functions[torch.add.__str__] - torch.Tensor.add = old_functions[torch.Tensor.add.__str__] - - torch.einsum = old_functions[torch.einsum.__str__] - - torch.baddbmm = old_functions[torch.baddbmm.__str__] - - -def _reload_miscellaneous_operations(): - einops.einsum = old_functions[einops.einsum.__str__] - - -def _rnn_flops(flops, rnn_module, w_ih, w_hh, input_size): - gates_size = w_ih.shape[0] - # matrix matrix mult ih state and internal state - flops += 2 * w_ih.shape[0] * w_ih.shape[1] - gates_size - # matrix matrix mult hh state and internal state - flops += 2 * w_hh.shape[0] * w_hh.shape[1] - gates_size - if isinstance(rnn_module, (nn.RNN, nn.RNNCell)): - # add both operations - flops += rnn_module.hidden_size - elif isinstance(rnn_module, (nn.GRU, nn.GRUCell)): - # hadamard of r - flops += rnn_module.hidden_size - # adding operations from both states - flops += rnn_module.hidden_size * 3 - # last two hadamard _product and add - flops += rnn_module.hidden_size * 3 - elif isinstance(rnn_module, (nn.LSTM, nn.LSTMCell)): - # adding operations from both states - flops += rnn_module.hidden_size * 4 - # two hadamard _product and add for C state - flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size - # final hadamard - flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size - return flops - - -def _rnn_forward_hook(rnn_module, input, output): - flops = 0 - # input is a tuple containing a sequence to process and (optionally) hidden state - inp = input[0] - batch_size = inp.shape[0] - seq_length = inp.shape[1] - num_layers = rnn_module.num_layers - - for i in range(num_layers): - w_ih = rnn_module.__getattr__("weight_ih_l" + str(i)) - w_hh = rnn_module.__getattr__("weight_hh_l" + str(i)) - if i == 0: - input_size = rnn_module.input_size - else: - input_size = rnn_module.hidden_size - flops = _rnn_flops(flops, rnn_module, w_ih, w_hh, input_size) - if rnn_module.bias: - b_ih = rnn_module.__getattr__("bias_ih_l" + str(i)) - b_hh = rnn_module.__getattr__("bias_hh_l" + str(i)) - flops += b_ih.shape[0] + b_hh.shape[0] - - flops *= batch_size - flops *= seq_length - if rnn_module.bidirectional: - flops *= 2 - rnn_module.__flops__ += int(flops) - - -def _rnn_cell_forward_hook(rnn_cell_module, input, output): - flops = 0 - inp = input[0] - batch_size = inp.shape[0] - w_ih = rnn_cell_module.__getattr__("weight_ih") - w_hh = rnn_cell_module.__getattr__("weight_hh") - input_size = inp.shape[1] - flops = _rnn_flops(flops, rnn_cell_module, w_ih, w_hh, input_size) - if rnn_cell_module.bias: - b_ih = rnn_cell_module.__getattr__("bias_ih") - b_hh = rnn_cell_module.__getattr__("bias_hh") - flops += b_ih.shape[0] + b_hh.shape[0] - - flops *= batch_size - rnn_cell_module.__flops__ += int(flops) - - -MODULE_HOOK_MAPPING = { - # RNN - nn.RNN: _rnn_forward_hook, - nn.GRU: _rnn_forward_hook, - nn.LSTM: _rnn_forward_hook, - nn.RNNCell: _rnn_cell_forward_hook, - nn.LSTMCell: _rnn_cell_forward_hook, - nn.GRUCell: _rnn_cell_forward_hook, -} - - -def macs_to_string(macs, units=None, precision=DEFAULT_PRECISION): - return f"{number_to_string(macs, units=units, precision=precision)}MACs" - - -def number_to_string(num, units=None, precision=DEFAULT_PRECISION): - if units is None: - if num >= 1e12: - magnitude, units = 1e12, "T" - elif num >= 1e9: - magnitude, units = 1e9, "G" - elif num >= 1e6: - magnitude, units = 1e6, "M" - elif num >= 1e3: - magnitude, units = 1e3, "K" - elif num >= 1 or num == 0: - magnitude, units = 1, "" - elif num >= 1e-3: - magnitude, units = 1e-3, "m" - else: - magnitude, units = 1e-6, "u" - else: - if units == "T": - magnitude = 1e12 - elif units == "G": - magnitude = 1e9 - elif units == "M": - magnitude = 1e6 - elif units == "K": - magnitude = 1e3 - elif units == "m": - magnitude = 1e-3 - elif units == "u": - magnitude = 1e-6 - else: - magnitude = 1 - return f"{round(num / magnitude, precision):g} {units}" - - -def flops_to_string(flops, units=None, precision=DEFAULT_PRECISION): - return f"{number_to_string(flops, units=units, precision=precision)}FLOPS" - - -def bytes_to_string(b, units=None, precision=DEFAULT_PRECISION): - return f"{number_to_string(b, units=units, precision=precision)}B" - - -def params_to_string(params_num, units=None, precision=DEFAULT_PRECISION): - units = units.replace("B", "G") if units else units - return number_to_string(params_num, units=units, precision=precision).replace("G", "B").strip() - - -def duration_to_string(duration, units=None, precision=DEFAULT_PRECISION): - return f"{number_to_string(duration, units=units, precision=precision)}s" - - - # can not iterate over all submodules using self.model.modules() - # since modules() returns duplicate modules only once -def get_module_flops(module): - sum = module.__flops__ - # iterate over immediate children modules - for child in module.children(): - sum += get_module_flops(child) - return sum - - -def get_module_macs(module): - sum = module.__macs__ - # iterate over immediate children modules - for child in module.children(): - sum += get_module_macs(child) - return sum - - -def get_module_duration(module): - duration = module.__duration__ - if duration == 0: # e.g. ModuleList - for m in module.children(): - duration += get_module_duration(m) - return duration - - -def get_model_profile(model, - input_shape=None, - args=[], - kwargs={}, - print_profile=True, - detailed=True, - module_depth=-1, - top_modules=1, - warm_up=1, - as_string=True, - output_file=None, - ignore_modules=None, - mode='forward'): - """Returns the total floating-point operations, MACs, and parameters of a model. - - Example: - - .. code-block:: python - - model = torchvision.models.alexnet() - batch_size = 256 - flops, macs, params = get_model_profile(model=model, input_shape=(batch_size, 3, 224, 224))) - - Args: - model ([torch.nn.Module]): the PyTorch model to be profiled. - input_shape (tuple): input shape to the model. If specified, the model takes a tensor with this shape as the only positional argument. - args (list): list of positional arguments to the model. - kwargs (dict): dictionary of keyword arguments to the model. - print_profile (bool, optional): whether to print the model profile. Defaults to True. - detailed (bool, optional): whether to print the detailed model profile. Defaults to True. - module_depth (int, optional): the depth into the nested modules. Defaults to -1 (the inner most modules). - top_modules (int, optional): the number of top modules to print in the aggregated profile. Defaults to 3. - warm_up (int, optional): the number of warm-up steps before measuring the latency of each module. Defaults to 1. - as_string (bool, optional): whether to print the output as string. Defaults to True. - output_file (str, optional): path to the output file. If None, the profiler prints to stdout. - ignore_modules ([type], optional): the list of modules to ignore during profiling. Defaults to None. - - Returns: - The number of floating-point operations, multiply-accumulate operations (MACs), and parameters in the model. - """ - assert isinstance(model, nn.Module), "model must be a PyTorch module" - prof = FlopsProfiler(model) - model.eval() - - if input_shape is not None: - assert type(input_shape) is tuple, "input_shape must be a tuple" - assert len(input_shape) >= 1, "input_shape must have at least one element" - try: - input = torch.ones(()).new_empty( - (*input_shape, ), - dtype=next(model.parameters()).dtype, - device=next(model.parameters()).device, - ) - except StopIteration: - input = torch.ones(()).new_empty((*input_shape, )) - - args = [input] - assert (len(args) > 0) or (len(kwargs) > 0), "args and/or kwargs must be specified if input_shape is None" - - logger.info("Flops profiler warming-up...") - for _ in range(warm_up): - if kwargs: - if mode == 'forward': - _ = model(*args, **kwargs) - if mode == 'generate': - _ = model.generate(*args, **kwargs) - else: - if mode == 'forward': - _ = model(*args) - if mode == 'generate': - _ = model.generate(*args) - prof.start_profile(ignore_list=ignore_modules) - - if kwargs: - if mode == 'forward': - _ = model(*args, **kwargs) - if mode == 'generate': - _ = model.generate(*args, **kwargs) - else: - if mode == 'forward': - _ = model(*args) - if mode == 'generate': - _ = model.generate(*args) - - flops = prof.get_total_flops() - macs = prof.get_total_macs() - params = prof.get_total_params() - if print_profile: - prof.print_model_profile(profile_step=warm_up, - module_depth=module_depth, - top_modules=top_modules, - detailed=detailed, - output_file=output_file) - - prof.end_profile() - if as_string: - return number_to_string(flops), macs_to_string(macs), params_to_string(params) - - return flops, macs, params diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/__init__.py deleted file mode 100644 index 208299fb8c50f73468d293b6fa5dca71649d62e7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 244cce2d968ae43ae2af3d3e3091e59629dac594..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/base_optimizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/base_optimizer.cpython-310.pyc deleted file mode 100644 index 935099cf58376956c4aa48e43b33c3d4f109c8cb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/base_optimizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/bf16_optimizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/bf16_optimizer.cpython-310.pyc deleted file mode 100644 index 2f1f4fb58297c10dcb4ac05a20a64295acac0ef8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/bf16_optimizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/compiler.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/compiler.cpython-310.pyc deleted file mode 100644 index 74e10a2913bb155d71495a6c4e2fa292bbf7a27d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/compiler.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/config.cpython-310.pyc deleted file mode 100644 index 3633baf1b5995e9e6f144c2758d0b83e9f02203e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/config_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/config_utils.cpython-310.pyc deleted file mode 100644 index a15ee659fb9175c6a2b66ea12e8698ce7706d0d0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/config_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/constants.cpython-310.pyc deleted file mode 100644 index acd297ba1d81ed63d6bd7f60ebfaf911aae96626..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/constants.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/dataloader.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/dataloader.cpython-310.pyc deleted file mode 100644 index f095f2cfb62d26de447fb110ea276e72b53ec6fd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/dataloader.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/eigenvalue.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/eigenvalue.cpython-310.pyc deleted file mode 100644 index fa87a925d865a118c2ec2ae6300850c8a2830340..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/eigenvalue.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/engine.cpython-310.pyc deleted file mode 100644 index 6ba1be5fc39297b0c2a298bb19c346494aa3cc85..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/engine.cpython-310.pyc +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a60d37ed4603ba2d09236d2a5c99fcaf4862224b2404215258104c6ddd878d60 -size 130592 diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/hybrid_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/hybrid_engine.cpython-310.pyc deleted file mode 100644 index 752e43150e33a01fd5b233a65c648d353dcb6c58..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/hybrid_engine.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/lr_schedules.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/lr_schedules.cpython-310.pyc deleted file mode 100644 index 4c95d76f6e8d82ef1b1f985db8a44345f9ca580a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/lr_schedules.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/precision_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/precision_config.cpython-310.pyc deleted file mode 100644 index c2cd14efc968707f89e0ce054edbbdf6d5a5a504..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/precision_config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/progressive_layer_drop.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/progressive_layer_drop.cpython-310.pyc deleted file mode 100644 index 506a6d77e0e439d0edaf5a16b49a5b52b8279fc8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/progressive_layer_drop.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/quantize.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/quantize.cpython-310.pyc deleted file mode 100644 index 9dd0305714d7a9dfce33f01105993f5cde1d4bd0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/quantize.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/sparse_tensor.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/sparse_tensor.cpython-310.pyc deleted file mode 100644 index da793330a8a8fa49753fb6ba7dd68f26285610b0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/sparse_tensor.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/state_dict_factory.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/state_dict_factory.cpython-310.pyc deleted file mode 100644 index 8bea0b44c741a7e77a8cfc4c4083ec6bf590428b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/state_dict_factory.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/torch_autocast.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/torch_autocast.cpython-310.pyc deleted file mode 100644 index 4c9bdfa072a9d533652f584925dd48ab1a06c40d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/torch_autocast.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index 14b5b7effa380bd687a9df0679284657fbd661b8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/weight_quantizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/weight_quantizer.cpython-310.pyc deleted file mode 100644 index 1d46e8d5a8323e12a90609aa6d6082024387a143..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/__pycache__/weight_quantizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 16728f1309a65ffd1b0550081df0e1c9f3da3847..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/checkpointing.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/checkpointing.cpython-310.pyc deleted file mode 100644 index c186426c2d6eb7d2a48d385b2b2addd1bcc0e82f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/checkpointing.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/config.cpython-310.pyc deleted file mode 100644 index 58d4cccec19e7c640516e23a01894c316e8519df..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py deleted file mode 100644 index 08c4b81937f918425d6da976b9c8391b30eee741..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py +++ /dev/null @@ -1,1142 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Use to partition the activations stored for backward propagation -Therefore reduces the memory consumption -Also implements CPU checkpointing and contiguous memory checkpointing -Reduces memory consumption and memory fragmentation - -Code for rng checkpointing taken from NVIDIA Megatron-LM mpu/random.py -b886b7bb972afe72bac0f5de4f42a4a7bae8ebef -""" - -# Parts of the code here are adapted from PyTorch -# repo: https://github.com/pytorch/pytorch -import copy -import torch -import contextlib -from deepspeed import comm as dist -import weakref - -import mmap -from torch import _C - -from deepspeed.runtime.config import DeepSpeedConfig -from deepspeed.utils import logger -from deepspeed.runtime.utils import copy_to_device, move_to_device, see_memory_usage -from deepspeed.utils.timer import SynchronizedWallClockTimer as Timers, FORWARD_GLOBAL_TIMER -from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank -from deepspeed.accelerator import get_accelerator -from deepspeed.runtime import compiler - -# DeepSpeed Checkpointing Enabled or Disabled -deepspeed_checkpointing_enabled = False - -# MP parameters -mpu = None - -#set default values -mp_rank = 0 -mp_size = 1 -mp_group = None - -# Model Parameters -num_layers = None - -# Checkpointing buffers -contiguous_data_buffers = [] -data_offsets = [] - -contiguous_size_buffers = [] -size_offsets = [] - -timers = None - -# optimization flags -PARTITION_ACTIVATIONS = False -CPU_CHECKPOINT = False -CONTIGUOUS_CHECKPOINTING = False -SYNCHRONIZE = False -PROFILE_TIME = False - -# Default name for the model parallel rng tracker. -_MODEL_PARALLEL_RNG_TRACKER_NAME = 'model-parallel-rng' - - -def detach_variable(inputs, device=None): - if isinstance(inputs, tuple): - out = [] - for inp in inputs: - if not isinstance(inp, torch.Tensor): - out.append(inp) - continue - - requires_grad = inp.requires_grad - - if device is not None: - x = inp.to(device=device) - else: - x = inp - - x = x.detach() - x.requires_grad = requires_grad - out.append(x) - return tuple(out) - else: - raise RuntimeError("Only tuple of tensors is supported. Got Unsupported input type: ", type(inputs).__name__) - - -def _set_cuda_rng_state(new_state, device=-1): - """Sets the random number generator state of the current GPU. - - Arguments: - new_state (torch.ByteTensor): The desired state - This function is adapted from PyTorch repo (torch.cuda.set_rng_state) #ignore-cuda - with a single change: the input state is not cloned. Cloning caused - major performance issues for +4 GPU cases. - """ - if hasattr(_C, '_cuda_setRNGState') and callable(_C._cuda_setRNGState): - # older PyTorch - def cb(): - with get_accelerator().device(device): - _C._cuda_setRNGState(new_state) - else: - # newer PyTorch - if device == -1: - device = torch.device(get_accelerator().device_name()) - elif isinstance(device, str): - device = torch.device(device) - elif isinstance(device, int): - device = torch.device(get_accelerator().device_name(), device) - - def cb(): - idx = device.index - if idx is None: - idx = get_accelerator().current_device() - default_generator = get_accelerator().default_generator(idx) - default_generator.set_state(new_state) - - get_accelerator().lazy_call(cb) - - -class CudaRNGStatesTracker: - """Tracker for the cuda RNG states. - - Using the `add` method, a cuda rng state is initialized based on - the input `seed` and is assigned to `name`. Later, by forking the - rng state, we can perform operations and return to our starting - cuda state. - """ - - def __init__(self): - # Map from a string name to the cuda rng state. - self.states_ = {} - # Seeds are just for book keeping and ensure no seed is set twice. - self.seeds_ = set() - - def reset(self): - """Set to the initial state (no tracker).""" - self.states_ = {} - self.seeds_ = set() - - def get_states(self): - """Get rng states. Copy the dictionary so we have direct - pointers to the states, not just a pointer to the dictionary.""" - return copy.copy(self.states_) - - def set_states(self, states): - """Set the rng states. For efficiency purposes, we do not check - the size of seed for compatibility.""" - self.states_ = states - - def add(self, name, seed): - """Track the rng state.""" - # Check seed is not already used. - if seed in self.seeds_: - raise Exception('seed {} already exists'.format(seed)) - self.seeds_.add(seed) - # Check that state is not already defined. - if name in self.states_: - raise Exception('cuda rng state {} already exists'.format(name)) - # Get the current rng state. - orig_rng_state = get_accelerator().get_rng_state() - # Set the new state and store it. - get_accelerator().manual_seed(seed) - self.states_[name] = get_accelerator().get_rng_state() - # Reset rng state to what it was. - _set_cuda_rng_state(orig_rng_state) - - @contextlib.contextmanager - def fork(self, name=_MODEL_PARALLEL_RNG_TRACKER_NAME): - """Fork the cuda rng state, perform operations, and exit with - the original state.""" - # Check if we have added the state - if name not in self.states_: - raise Exception('cuda rng state {} is not added'.format(name)) - # Store current rng state. - orig_cuda_rng_state = get_accelerator().get_rng_state() - # Set rng state to the desired one - _set_cuda_rng_state(self.states_[name]) - # Do the stuff we wanted to do. - try: - yield - finally: - # Update the current rng state for later use. - self.states_[name] = get_accelerator().get_rng_state() - # And set the state to the original state we started with. - _set_cuda_rng_state(orig_cuda_rng_state) - - -# RNG tracker object. -_CUDA_RNG_STATE_TRACKER = CudaRNGStatesTracker() - - -def get_cuda_rng_tracker(): - """Get cuda rng tracker.""" - return _CUDA_RNG_STATE_TRACKER - - -def model_parallel_cuda_manual_seed(seed): - """Initialize model parallel cuda seed. - - This function should be called after the model parallel is - initialized. Also, no get_accelerator().manual_seed should be called - after this function. Basically, this is replacement for that - function. - Two set of RNG states are tracked: - default state: This is for data parallelism and is the same among a - set of model parallel GPUs but different across - different model parallel groups. This is used for - example for dropout in the non-model-parallel regions. - model-parallel state: This state is different among a set of model - parallel GPUs, but the same across data parallel - groups. This is used for example for dropout in - model parallel regions. - """ - global mpu - - tp_rank = bwc_tensor_model_parallel_rank(mpu) - - # 2718 is just for fun and any POSITIVE value will work. - offset = seed + 2718 - model_parallel_seed = offset + tp_rank - # Data parallel gets the original seed. - data_parallel_seed = seed - - if dist.get_rank() == 0: - logger.info( - '> initializing model parallel cuda seeds on global rank {}, ' - 'model parallel rank {}, and data parallel rank {} with ' - 'model parallel seed: {} and data parallel seed: {}'.format(dist.get_rank(), tp_rank, - mpu.get_data_parallel_rank(), - model_parallel_seed, data_parallel_seed), ) - _CUDA_RNG_STATE_TRACKER.reset() - # Set the default state. - get_accelerator().manual_seed(data_parallel_seed) - # and model parallel state. - _CUDA_RNG_STATE_TRACKER.add(_MODEL_PARALLEL_RNG_TRACKER_NAME, model_parallel_seed) - - -def model_parallel_reconfigure_tp_seed(seed): - global mpu - tp_rank = bwc_tensor_model_parallel_rank(mpu) - model_parallel_seed = seed + 2718 + tp_rank - with _CUDA_RNG_STATE_TRACKER.fork(): - get_accelerator().manual_seed(model_parallel_seed) - - -def get_partition_start(item): - global mp_rank, mp_size, mp_group - size = item.numel() - partition_size = size / mp_size - start = partition_size * mp_rank - return int(start) - - -def get_partition_size(item): - global mp_rank, mp_size, mp_group - size = item.numel() - assert size % mp_size == 0, "Doesn't handle if partition activation if item is not divisible by mp size" - partition_size = size / mp_size - return int(partition_size) - - -def gather_partitioned_activations(tensors, device=None): - global mp_rank, mp_size, mp_group - assert len(tensors) % 2 == 0, f'Expected even count of tensors, instead got {len(tensors)}' - inputs = [] - num_args = int(len(tensors) / 2) - for i in range(num_args): - - item = tensors[2 * i] - size = tensors[2 * i + 1] - - if not is_activation_to_checkpoint(item): - inputs.append(item) - continue - - # don't need to do all_gather if model parallel is not enabled - if mp_group is None or mp_size == 1: - item = item.view(list(size.numpy())) - if device is not None: - item = item.to(device) - inputs.append(item) - continue - - partition_size = item.numel() - tensor_size = partition_size * mp_size - if device is not None: - flat_tensor = torch.zeros([tensor_size], dtype=item.dtype, device=device) - else: - flat_tensor = torch.zeros([tensor_size], dtype=item.dtype, device=item.device) - part = flat_tensor.narrow(0, partition_size * mp_rank, partition_size) - part.copy_(item) - dist.all_gather_into_tensor(flat_tensor, part, group=mp_group) - input_tensor = flat_tensor.view(list(size.numpy())) - item.data = input_tensor.data - - inputs.append(item) - - return tuple(inputs) - - -def extract_tensors(all_objects): - """ - Separate objects in list/tuple into tensors and non-tensors and create a mapping to enable re-aggregation. - The order of tensors and non-tensors is preserved in their respective output groups. - - Parameters: - all_objects (list/tuple): Objects containing tensors and non-tensors to be split. - - Returns: - tuple: Containing tensors, non-tensors, and bools of whether each position in original list/tuple was a tensor. - - """ - tensor_objects = [v for v in all_objects if torch.is_tensor(v)] - non_tensor_objects = [v for v in all_objects if not torch.is_tensor(v)] - tensor_flags = [torch.is_tensor(v) for v in all_objects] - if type(all_objects) is tuple: - return tuple(tensor_objects), tuple(non_tensor_objects), tuple(tensor_flags) - return tensor_objects, non_tensor_objects, tensor_flags - - -def merge_tensors(tensor_objects, non_tensor_objects, tensor_flags): - """ - Merge two lists (or tuples) of tensors and non-tensors using a mapping of positions in merged list (or tuple). - - Parameters: - tensor_objects (list/tuple): Tensors to merge. - non_tensor_objects (list/tuple): Non-tensors to merge. - tensor_flags (list/tuple): Indicates whether each position in output is a tensor. - - Returns: - tuple: Merge of tensors and non-tensors - """ - merged_objects = [] - tensor_idx = 0 - non_tensor_idx = 0 - - real_tensor_flags = None - - # remove the flags that are assigned to the size of the flattened tensors - if PARTITION_ACTIVATIONS: - real_tensor_flags = [] - previous_flag = False - for flag in tensor_flags: - if previous_flag: - previous_flag = False - continue - previous_flag = flag - real_tensor_flags.append(flag) - else: - real_tensor_flags = tensor_flags - - for is_tensor in real_tensor_flags: - if is_tensor: - merged_objects.append(tensor_objects[tensor_idx]) - tensor_idx += 1 - else: - merged_objects.append(non_tensor_objects[non_tensor_idx]) - non_tensor_idx += 1 - - return tuple(merged_objects) - - -def is_activation_to_checkpoint(item): - """ - Is an activation to be checkpointed - """ - global mp_size - extra_flag = (not hasattr(item, 'no_checkpointing')) or (hasattr(item, 'no_checkpointing') - and item.no_checkpointing == False) - return torch.is_tensor(item) and item.is_floating_point() and item.numel() >= mp_size and extra_flag - - -def partition_activations(args, cpu_checkpoint, contiguous_checkpoint): - global contiguous_data_buffers, data_offsets - - inputs = [] - num_non_fp_tensors = 0 - - for arg_index, item in enumerate(args): - if not is_activation_to_checkpoint(item): - inputs.append(item) - num_non_fp_tensors += 1 - continue - - i = arg_index - num_non_fp_tensors - partition_size = get_partition_size(item) - partition = item.detach().contiguous().view(-1).narrow(0, get_partition_start(item), partition_size).clone() - - buffer_device = torch.device('cpu') if cpu_checkpoint else partition.device - - if contiguous_checkpoint: - if i >= len(contiguous_data_buffers): - tensor_list = [ - torch.tensor(()).new_empty([partition_size], dtype=partition.dtype, device=buffer_device) - for _ in range(num_layers) - ] - contiguous_data_buffers.append(tensor_list) - data_offsets.append(0) - elif contiguous_data_buffers[i] is None: - tensor_list = [ - torch.tensor(()).new_empty([partition_size], dtype=partition.dtype, device=buffer_device) - for _ in range(num_layers) - ] - contiguous_data_buffers[i] = tensor_list - data_offsets[i] = 0 - - # Because the 'new_empty' returns uninitialized pages, - # the pages need to be populated during the cudaMemcpy time - # which increases the data copy time. To avoid this, we - # pre-populate these pages by simply writing 0 ahead of - # the actual cudaMemcpy operation time. Due to the - # previously launched GPU kernels, there is a small - # window of time here for CPUs to populate pages asynchronously. - contiguous_data_buffers[i][data_offsets[i]].data[range( - 0, contiguous_data_buffers[i][data_offsets[i]].data.shape[0], - int(mmap.PAGESIZE / contiguous_data_buffers[i][data_offsets[i]].data.element_size()))] = 0 - - contiguous_partition = contiguous_data_buffers[i][data_offsets[i]].data.copy_(partition.data) - data_offsets[i] = data_offsets[i] + 1 - inputs.append(contiguous_partition) - else: - partition = partition.cpu() if CPU_CHECKPOINT else partition - inputs.append(partition) - - return inputs - - -def get_partitioned_activations_for_backward(args, inputs, contiguous_checkpoint): - global contiguous_size_buffers, size_offsets - - new_args = [] - num_non_fp_tensors = 0 - - for arg_index, (arg, inp) in enumerate(zip(args, inputs)): - size = torch.tensor(arg.size()) if torch.is_tensor(arg) else None - if not is_activation_to_checkpoint(arg): - new_args.append(arg) - new_args.append(size) - num_non_fp_tensors += 1 - continue - - arg.data = torch.empty([], device=arg.device).data - arg.saved_data = inp.data - - new_args.append(arg) - i = arg_index - num_non_fp_tensors - - if contiguous_checkpoint: - numel = size.numel() - if i >= len(contiguous_size_buffers): - tmp = torch.tensor(()) - contiguous_size_buffers.append( - tmp.new_empty([numel * num_layers], dtype=size.dtype, device=size.device)) - size_offsets.append(0) - elif contiguous_size_buffers[i] is None: - tmp = torch.tensor(()) - contiguous_size_buffers[i] = tmp.new_empty([numel * num_layers], dtype=size.dtype, device=size.device) - size_offsets[i] = 0 - - contiguous_size = contiguous_size_buffers[i].narrow(0, size_offsets[i], numel).data.copy_(size.data) - contiguous_size = contiguous_size.view_as(size) - size_offsets[i] = size_offsets[i] + numel - new_args.append(contiguous_size) - else: - new_args.append(size) - - return new_args - - -def get_cpu_activations_for_backward(args, inputs): - new_args = [] - for i, (arg, inp) in enumerate(zip(args, inputs)): - if not is_activation_to_checkpoint(arg): - new_args.append(arg) - continue - - arg.data = torch.empty([], device=arg.device).data - arg.saved_data = inp.data - new_args.append(arg) - - return new_args - - -class CheckpointFunction(torch.autograd.Function): - """This function is adapted from torch.utils.checkpoint with - two main changes: - 1) torch.cuda.set_rng_state is replaced with `_set_cuda_rng_state` #ignore-cuda - 2) the states in the model parallel tracker are also properly - tracked/set/reset. - 3) Performance activation partitioning, contiguous memory optimization - 4) CPU Checkpointing - 5) Profile forward and backward functions - """ - - @staticmethod - def forward(ctx, run_function, all_outputs, *args): - global mpu, timers, SYNCHRONIZE, PROFILE_TIME - - def save_args_for_backward(*all_args): - tensor_args, non_tensor_args, tensor_flags = extract_tensors(all_objects=all_args) - ctx.deepspeed_saved_tensors = tensor_args - ctx.non_tensor_args = non_tensor_args - ctx.tensor_flags = tensor_flags - - if SYNCHRONIZE: - get_accelerator().synchronize() - - if timers is None and PROFILE_TIME: - timers = Timers() - - if PROFILE_TIME: - timers(FORWARD_GLOBAL_TIMER).start() - - ctx.run_function = run_function - global num_layers - global mp_rank, mp_size, mp_group - global contiguous_data_buffers, contiguous_size_buffers - global data_offsets, size_offsets - global PARTITION_ACTIVATIONS, buffer_0, buffer_1, buffer_0_offset, buffer_1_offset - - cuda_device = get_accelerator().current_device_name() - transport_stream = get_accelerator().Stream(device=cuda_device) - - if PARTITION_ACTIVATIONS: - inputs = partition_activations(args, CPU_CHECKPOINT, CONTIGUOUS_CHECKPOINTING) - elif CPU_CHECKPOINT: - inputs = copy_to_device(args, device=torch.device('cpu'), criterion_func=is_activation_to_checkpoint) - - # just in case something funky is happening such as reuse of inputs - inputs_cuda = copy_to_device(args, device=cuda_device, criterion_func=is_activation_to_checkpoint) - - # Copy the rng states. - ctx.fwd_cpu_rng_state = torch.get_rng_state() - ctx.fwd_cuda_rng_state = get_accelerator().get_rng_state() - ctx.fwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states() - - see_memory_usage("Before running forward on the layer", force=False) - # ctx.save_for_backward(*args) - with torch.no_grad(): - outputs = run_function(*inputs_cuda) - - see_memory_usage("After running forward on the layer", force=False) - del inputs_cuda - - if PARTITION_ACTIVATIONS: - new_args = get_partitioned_activations_for_backward(args, inputs, CONTIGUOUS_CHECKPOINTING) - assert len(new_args) % 2 == 0, f'save_for_backward called with odd number of args, {len(new_args)}' - save_args_for_backward(*new_args) - elif CPU_CHECKPOINT: - new_args = get_cpu_activations_for_backward(args, inputs) - save_args_for_backward(*new_args) - else: - save_args_for_backward(*args) - - if PROFILE_TIME: - timers(FORWARD_GLOBAL_TIMER).stop() - timers.log([FORWARD_GLOBAL_TIMER]) - if SYNCHRONIZE: - get_accelerator().synchronize() - - # Tensors returned from forward() may not be differentiable. - if torch.is_tensor(outputs): - non_grad_outputs = [outputs] if not outputs.is_floating_point() else [] - else: - non_grad_outputs = [o for o in outputs if torch.is_tensor(o) and not o.is_floating_point()] - ctx.mark_non_differentiable(*non_grad_outputs) - - if torch.is_tensor(outputs): - all_outputs += [outputs] - return outputs - else: - all_outputs += outputs - outputs, _, _ = extract_tensors(all_objects=outputs) - return tuple(outputs) - - @staticmethod - def backward(ctx, *grads): - global timers - see_memory_usage("In backward", force=False) - # removing pointers to the contiguous buffer memory - # so that they can be garbage collected once the checkpoints - # have been used - if SYNCHRONIZE: - get_accelerator().synchronize() - if PROFILE_TIME: - timers('backward').start() - - if CONTIGUOUS_CHECKPOINTING: - global data_offsets, size_offsets - global contiguous_data_buffers, contiguous_size_buffers - - for buffers in contiguous_data_buffers: - buffers = [] - - # frees up all the pointers to the checkpoints except for the ones - # stored by save for backward - contiguous_data_buffers = [] - contiguous_size_buffers = [] - data_offsets = [] - size_offsets = [] - - see_memory_usage("In backward checkpointing code", force=False) - if not torch.autograd._is_checkpoint_valid(): - raise RuntimeError("Checkpointing is not compatible with .grad(), " - "please use .backward() if possible") - - global PARTITION_ACTIVATIONS - cuda_device = get_accelerator().current_device_name() - transport_stream = get_accelerator().Stream(device=cuda_device) - # Rebuild deepspeed_saved_tensors - for t in ctx.deepspeed_saved_tensors: - if t is not None and hasattr(t, 'saved_data') and t.saved_data is not None: - t.data = t.saved_data.to(t.device) - t.saved_data = None - - if PARTITION_ACTIVATIONS: - # with get_accelerator().stream(transport_stream): - inputs = gather_partitioned_activations(ctx.deepspeed_saved_tensors, - device=cuda_device if CPU_CHECKPOINT else None) - detached_inputs = detach_variable(inputs) - elif CPU_CHECKPOINT: - inputs = move_to_device(ctx.deepspeed_saved_tensors, cuda_device, is_activation_to_checkpoint) - detached_inputs = detach_variable(inputs) - else: - inputs = ctx.deepspeed_saved_tensors - detached_inputs = detach_variable(inputs) - - # Add non tensor input args - detached_inputs = merge_tensors(tensor_objects=detached_inputs, - non_tensor_objects=ctx.non_tensor_args, - tensor_flags=ctx.tensor_flags) - - # Store the current states. - bwd_cpu_rng_state = torch.get_rng_state() - bwd_cuda_rng_state = get_accelerator().get_rng_state() - bwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states() - - # Set the states to what it used to be before the forward pass. - torch.set_rng_state(ctx.fwd_cpu_rng_state) - _set_cuda_rng_state(ctx.fwd_cuda_rng_state) - get_cuda_rng_tracker().set_states(ctx.fwd_cuda_rng_state_tracker) - - # if PARTITION_ACTIVATIONS: - # current_stream=get_accelerator().current_stream() - # current_stream.wait_stream(transport_stream) - - see_memory_usage("In backward checkpointing code before forward", force=False) - - with torch.enable_grad(): - outputs = ctx.run_function(*detached_inputs) - - see_memory_usage("In backward checkpointing code after forward", force=False) - # Set the states back to what it was at the start of this function. - torch.set_rng_state(bwd_cpu_rng_state) - _set_cuda_rng_state(bwd_cuda_rng_state) - get_cuda_rng_tracker().set_states(bwd_cuda_rng_state_tracker) - - if isinstance(outputs, torch.Tensor): - outputs = (outputs, ) - - # Filter out non tensor outputs - outputs, _, _ = extract_tensors(all_objects=outputs) - - # Construct arguments to autograd.backward(). - # This is usually just outputs and grads, but forward() can return tensors that - # are not differentiable. - output_tensors = [] - grad_tensors = [] - for out, grad in zip(outputs, grads): - if out.requires_grad: - output_tensors.append(out) - grad_tensors.append(grad) - - see_memory_usage("In backward checkpointing code before backward", force=False) - - torch.autograd.backward(output_tensors, grad_tensors) - - # Force clear our stashed tensors to prevent a memory leak in certain scenarios - ctx.deepspeed_saved_tensors = None - ctx.non_tensor_args = None - ctx.tensor_flags = None - - see_memory_usage("After backward checkpointing code after backward", force=False) - - if PROFILE_TIME: - timers('backward').stop() - timers.log(['backward']) - if SYNCHRONIZE: - get_accelerator().synchronize() - ret_list = [None, None] # first None for ctx - for inp in detached_inputs: - if torch.is_tensor(inp): - ret_list.append(inp.grad) - else: - ret_list.append(None) - - return tuple(ret_list) - - -def non_reentrant_checkpoint(function, *args): - """This function is union of `torch.utils.checkpoint._checkpoint_without_reentrant` and `CheckpointFunction` in this module - - This function is aim to solve the back probagation error raised from all input requires no grad. - * has already been implemented in pytorch for a while, the solution is stable at most time except for jit module mode. - * can help to solve the issue which is hacked by `deepspeed.runtime.pipe.module.PipelineModule._is_checkpointable` - - Main modifications compared to the implementation of torch: - 1. adapt to the signature of `checkpoint` function in this module - 2. solve the non-deterministic by random state management consistent with deepspeed `CheckpointFunction` - 3. when there is partition or cpu checkpointing, gather them in the unpack_hook during back probagation - 4. make all after backward blocks in the hook which will executed after all leaf nodes backward execution. - 5. above 4. is inspired by `torch.autograd.graph.register_multi_grad_hook`, which is only implemented after 2.0.0 - """ - global mpu, timers, SYNCHRONIZE, PROFILE_TIME - - deepspeed_saved_tensors = None - non_tensor_args = None - tensor_flags = None - - def save_args_for_backward(*all_args): - """keep this function to reduce the modification from original implementation""" - nonlocal deepspeed_saved_tensors, non_tensor_args, tensor_flags - tensor_args, non_tensor_args, tensor_flags = extract_tensors(all_objects=all_args) - deepspeed_saved_tensors = tensor_args - non_tensor_args = non_tensor_args - tensor_flags = tensor_flags - - if SYNCHRONIZE: - get_accelerator().synchronize() - - if timers is None and PROFILE_TIME: - timers = Timers() - - if PROFILE_TIME: - timers(FORWARD_GLOBAL_TIMER).start() - - global num_layers - global mp_rank, mp_size, mp_group - global contiguous_data_buffers, contiguous_size_buffers - global data_offsets, size_offsets - global PARTITION_ACTIVATIONS, buffer_0, buffer_1, buffer_0_offset, buffer_1_offset - - cuda_device = get_accelerator().current_device_name() - transport_stream = get_accelerator().Stream(device=cuda_device) - - if PARTITION_ACTIVATIONS: - inputs = partition_activations(args, CPU_CHECKPOINT, CONTIGUOUS_CHECKPOINTING) - elif CPU_CHECKPOINT: - inputs = copy_to_device(args, device=torch.device('cpu'), criterion_func=is_activation_to_checkpoint) - - # just in case something funky is happening such as reuse of inputs - inputs_cuda = copy_to_device(args, device=cuda_device, criterion_func=is_activation_to_checkpoint) - - # Copy the rng states. - fwd_cpu_rng_state = torch.get_rng_state() - fwd_cuda_rng_state = get_accelerator().get_rng_state() - fwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states() - - if PARTITION_ACTIVATIONS: - new_args = get_partitioned_activations_for_backward(args, inputs, CONTIGUOUS_CHECKPOINTING) - assert len(new_args) % 2 == 0, f'save_for_backward called with odd number of args, {len(new_args)}' - save_args_for_backward(*new_args) - elif CPU_CHECKPOINT: - new_args = get_cpu_activations_for_backward(args, inputs) - save_args_for_backward(*new_args) - else: - save_args_for_backward(*args) - - class Holder(): - """the place holder object used as activations to save memory""" - pass - - # weakref seems utilized to discover the tensor deletion before a whole - # forward backward pair loop finished - storage: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() - weak_holder_list = [] - leaf_tensors = [] - backward_visited_leaf_nodes = 0 - - def checkpoint_pack(tensor_from_forward): - """used to record the activation order in the `weak_holder_list` - - the activation order in holder list is consistent between the first forward and recomputing forward. - * the jit compiled forward will break the order consistency * - """ - res = Holder() - weak_holder_list.append(weakref.ref(res)) - - # if this is a leaf tensor, save it for backward progression trace - # leaf tensor used to be input or parameters, which is not activations and - # has no memory overhead - if tensor_from_forward.requires_grad and tensor_from_forward.is_leaf: - leaf_tensors.append(tensor_from_forward) - return res - - def checkpoint_unpack(holder_from_backward): - """retrieve the activations from recompute""" - nonlocal deepspeed_saved_tensors, non_tensor_args, tensor_flags - - # if this is the first step of backward probagation, recompute the graph and save - # all the activations with the same order as `checkpoint_pack` does - if len(storage) == 0: - unpack_counter = 0 - - def replay_pack(tensor_from_replay): - """save recompute activations""" - nonlocal unpack_counter - unpack_counter += 1 - - if weak_holder_list[unpack_counter - 1]() is None: - return - - detached_activations = tensor_from_replay.detach() - storage[weak_holder_list[unpack_counter - 1]()] = detached_activations - - return - - def replay_unpack(none_value): - """recompute graph need not to backward""" - raise RuntimeError("You are calling backwards on a tensor that is never exposed.") - - global timers - see_memory_usage("In backward", force=False) - # removing pointers to the contiguous buffer memory - # so that they can be garbage collected once the checkpoints - # have been used - if SYNCHRONIZE: - get_accelerator().synchronize() - if PROFILE_TIME: - timers('backward').start() - - if CONTIGUOUS_CHECKPOINTING: - global data_offsets, size_offsets - global contiguous_data_buffers, contiguous_size_buffers - - for buffers in contiguous_data_buffers: - buffers = [] - - # frees up all the pointers to the checkpoints except for the ones - # stored by save for backward - contiguous_data_buffers = [] - contiguous_size_buffers = [] - data_offsets = [] - size_offsets = [] - - see_memory_usage("In backward checkpointing code", force=False) - if not torch.autograd._is_checkpoint_valid(): - raise RuntimeError("Checkpointing is not compatible with .grad(), " - "please use .backward() if possible") - - global PARTITION_ACTIVATIONS - cuda_device = get_accelerator().current_device_name() - transport_stream = get_accelerator().Stream(device=cuda_device) - - # gather inputs which is partitioned or checkpointed before first forward - if PARTITION_ACTIVATIONS: - # with get_accelerator().stream(transport_stream): - inputs = gather_partitioned_activations(deepspeed_saved_tensors, - device=cuda_device if CPU_CHECKPOINT else None) - detached_inputs = detach_variable(inputs) - elif CPU_CHECKPOINT: - inputs = move_to_device(deepspeed_saved_tensors, cuda_device, is_activation_to_checkpoint) - detached_inputs = detach_variable(inputs) - else: - inputs = deepspeed_saved_tensors - detached_inputs = detach_variable(inputs) - - # Add non tensor input args - detached_inputs = merge_tensors(tensor_objects=detached_inputs, - non_tensor_objects=non_tensor_args, - tensor_flags=tensor_flags) - - # Store the current states. - bwd_cpu_rng_state = torch.get_rng_state() - bwd_cuda_rng_state = get_accelerator().get_rng_state() - bwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states() - - # Set the states to what it used to be before the forward pass. - torch.set_rng_state(fwd_cpu_rng_state) - _set_cuda_rng_state(fwd_cuda_rng_state) - get_cuda_rng_tracker().set_states(fwd_cuda_rng_state_tracker) - - see_memory_usage("In backward checkpointing code before forward", force=False) - with torch.enable_grad(), torch.autograd.graph.saved_tensors_hooks(replay_pack, replay_unpack): - _unused = function(*detached_inputs) - - see_memory_usage("In backward checkpointing code after forward", force=False) - # Set the states back to what it was at the start of this function. - torch.set_rng_state(bwd_cpu_rng_state) - _set_cuda_rng_state(bwd_cuda_rng_state) - get_cuda_rng_tracker().set_states(bwd_cuda_rng_state_tracker) - - deepspeed_saved_tensors = None - non_tensor_args = None - tensor_flags = None - - if holder_from_backward not in storage: - raise RuntimeError("Attempt to retrieve a tensor saved by autograd multiple times without checkpoint" - " recomputation being triggered in between, this is not currently supported.") - - return storage[holder_from_backward] - - def after_backward_hook(_nonuse_grads): - """the hook registered to all leaf tensors""" - nonlocal leaf_tensors, backward_visited_leaf_nodes - backward_visited_leaf_nodes += 1 - - if backward_visited_leaf_nodes == len(leaf_tensors): - see_memory_usage("After backward checkpointing code after backward", force=False) - - if PROFILE_TIME: - timers('backward').stop() - timers.log(['backward']) - if SYNCHRONIZE: - get_accelerator().synchronize() - - with torch.autograd.graph.saved_tensors_hooks(checkpoint_pack, checkpoint_unpack): - outputs = function(*inputs_cuda) - if PROFILE_TIME or SYNCHRONIZE: - for leaf_tensor in leaf_tensors: - leaf_tensor.register_hook(after_backward_hook) - - see_memory_usage("After running forward on the layer", force=False) - - if PROFILE_TIME: - timers(FORWARD_GLOBAL_TIMER).stop() - timers.log([FORWARD_GLOBAL_TIMER]) - if SYNCHRONIZE: - get_accelerator().synchronize() - - all_outputs = [] - if torch.is_tensor(outputs): - all_outputs += [outputs] - else: - all_outputs += outputs - - if len(all_outputs) == 1: - return all_outputs[0] - else: - return tuple(all_outputs) - - -@compiler.disable # WA from Pytorch repo for compile + zero 3 accuracy issue -def checkpoint(function, *args): - """Checkpoint a model or part of the model. - This has been directly copied from torch.utils.checkpoint. """ - - all_outputs = [] - CheckpointFunction.apply(function, all_outputs, *args) - if len(all_outputs) == 1: - return all_outputs[0] - else: - return tuple(all_outputs) - - -def partition_activations_in_checkpoint(partition_activation): - global PARTITION_ACTIVATIONS - PARTITION_ACTIVATIONS = partition_activation - if dist.get_rank() == 0: - logger.info(f"**************Partition Activations {PARTITION_ACTIVATIONS}************") - - -def set_num_layers(nlayers): - global num_layers - num_layers = nlayers - - -def reset(): - """Resets memory buffers related to contiguous memory optimizations. - Should be called during eval when multiple forward propagations are - computed without any backward propagation that usually clears these - buffers. - Arguments: - None - - Return: - None - """ - if CONTIGUOUS_CHECKPOINTING: - global data_offsets, size_offsets - global contiguous_data_buffers, contiguous_size_buffers - - for buffers in contiguous_data_buffers: - buffers = [] - - # frees up all the pointers to the checkpoints except for the ones - # stored by save for backward - contiguous_data_buffers = [] - contiguous_size_buffers = [] - data_offsets = [] - size_offsets = [] - - -def _configure_using_config_file(config, mpu=None): - global num_layers, PARTITION_ACTIVATIONS, CONTIGUOUS_CHECKPOINTING, \ - CPU_CHECKPOINT, SYNCHRONIZE, PROFILE_TIME - - config = DeepSpeedConfig(config, mpu=mpu).activation_checkpointing_config - if dist.get_rank() == 0: - logger.info(config.repr()) - PARTITION_ACTIVATIONS = config.partition_activations - CONTIGUOUS_CHECKPOINTING = config.contiguous_memory_optimization - num_layers = config.number_checkpoints - CPU_CHECKPOINT = config.cpu_checkpointing - SYNCHRONIZE = config.synchronize_checkpoint_boundary - PROFILE_TIME = config.profile - - -def _configure_defaults(): - - global mpu, num_layers, deepspeed_checkpointing_enabled - - global PARTITION_ACTIVATIONS, CONTIGUOUS_CHECKPOINTING, \ - CPU_CHECKPOINT, SYNCHRONIZE, PROFILE_TIME - - PARTITION_ACTIVATIONS = False - CONTIGUOUS_CHECKPOINTING = False - num_layers = False - CPU_CHECKPOINT = False - SYNCHRONIZE = False - PROFILE_TIME = False - deepspeed_checkpointing_enabled = True - - -def configure( - mpu_, - deepspeed_config=None, - partition_activations=None, - contiguous_checkpointing=None, - num_checkpoints=None, - checkpoint_in_cpu=None, - synchronize=None, - profile=None, -): - """Configure DeepSpeed Activation Checkpointing. - - Arguments: - mpu_: Optional: An object that implements the following methods - get_model_parallel_rank/group/world_size, and get_data_parallel_rank/group/world_size - - deepspeed_config: Optional: DeepSpeed Config json file when provided will be used to - configure DeepSpeed Activation Checkpointing - - partition_activations: Optional: Partitions activation checkpoint across model parallel - GPUs when enabled. By default False. Will overwrite deepspeed_config if provided - - contiguous_checkpointing: Optional: Copies activation checkpoints to a contiguous memory - buffer. Works only with homogeneous checkpoints when partition_activations is enabled. - Must provide num_checkpoints. By default False. Will overwrite deepspeed_config if - provided - - num_checkpoints: Optional: Number of activation checkpoints stored during the forward - propagation of the model. Used to calculate the buffer size for contiguous_checkpointing - Will overwrite deepspeed_config if provided - - checkpoint_in_cpu: Optional: Moves the activation checkpoint to CPU. Only works with - partition_activation. Default is false. Will overwrite deepspeed_config if provided - - synchronize: Optional: Performs get_accelerator().synchronize() at the beginning and end of - each call to deepspeed.checkpointing.checkpoint for both forward and backward pass. - By default false. Will overwrite deepspeed_config if provided - - profile: Optional: Logs the forward and backward time for each - deepspeed.checkpointing.checkpoint invocation. Will overwrite deepspeed_config - if provided - - Returns: - None - """ - global mpu, num_layers, deepspeed_checkpointing_enabled - - global PARTITION_ACTIVATIONS, CONTIGUOUS_CHECKPOINTING, \ - CPU_CHECKPOINT, SYNCHRONIZE, PROFILE_TIME - - _configure_defaults() - - if mpu_ is not None: - mpu = mpu_ - - if deepspeed_config is not None: - _configure_using_config_file(deepspeed_config, mpu=mpu) - - if partition_activations is not None: - PARTITION_ACTIVATIONS = partition_activations - - if contiguous_checkpointing is not None: - CONTIGUOUS_CHECKPOINTING = contiguous_checkpointing - - if num_checkpoints is not None: - num_layers = num_checkpoints - - if checkpoint_in_cpu is not None: - CPU_CHECKPOINT = checkpoint_in_cpu - - if synchronize is not None: - SYNCHRONIZE = synchronize - - if profile is not None: - PROFILE_TIME = profile - - if CONTIGUOUS_CHECKPOINTING: - assert PARTITION_ACTIVATIONS, "Contiguous Checkpointing is only available with partitioned activations. Set partitioned activations to true in deepspeed config" - if CONTIGUOUS_CHECKPOINTING: - assert num_layers is not None, "Must specify the number of layers with contiguous memory checkpointing" - - global mp_rank, mp_size, mp_group - - if mpu is not None: - if hasattr(mpu, 'get_tensor_model_parallel_rank'): - mp_rank = mpu.get_tensor_model_parallel_rank() - mp_size = mpu.get_tensor_model_parallel_world_size() - mp_group = mpu.get_tensor_model_parallel_group() - else: - mp_rank = mpu.get_model_parallel_rank() - mp_size = mpu.get_model_parallel_world_size() - mp_group = mpu.get_model_parallel_group() - - #print configuration only once - see_memory_usage("After configuration", force=False) - if dist.get_rank() == 0: - logger.info(f"Activation Checkpointing Information") - logger.info(f"----Partition Activations {PARTITION_ACTIVATIONS}, CPU CHECKPOINTING {CPU_CHECKPOINT}") - logger.info(f"----contiguous Memory Checkpointing {CONTIGUOUS_CHECKPOINTING} with {num_layers} total layers") - logger.info(f"----Synchronization {SYNCHRONIZE}") - logger.info(f"----Profiling time in checkpointing {PROFILE_TIME}") - - -def is_configured(): - """True if deepspeed activation checkpointing has been configured - by calling deepspeed.checkpointing.configure, else returns false - - Arguments: - None - - Return: - True of configured, else False - """ - return deepspeed_checkpointing_enabled diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/config.py deleted file mode 100644 index dc07388a95da039b50bad87b4aa57b12f4e41f6f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/activation_checkpointing/config.py +++ /dev/null @@ -1,94 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigObject - -######################################### -# DeepSpeed Activation Checkpointing -######################################### -# Activation Checkpointing Allows to save memory by only keeping a select few -#activations for the backpropagation. -ACTIVATION_CHKPT_FORMAT = ''' -Activation Checkpointing should be configured as: -"session_params": { - "activation_checkpointing": { - "partitioned_activations": [true|false], - "number_checkpoints": 100, - "contiguous_memory_optimization": [true|false], - "cpu_checkpointing": [true|false], - "profile": [true|false], - "synchronize_checkpoint_boundary": [true|false], - } -} -''' - -ACT_CHKPT_PARTITION_ACTIVATIONS = 'partition_activations' -ACT_CHKPT_PARTITION_ACTIVATIONS_DEFAULT = False - -ACT_CHKPT_NUMBER_CHECKPOINTS = 'number_checkpoints' -ACT_CHKPT_NUMBER_CHECKPOINTS_DEFAULT = None - -ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION = 'contiguous_memory_optimization' -ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION_DEFAULT = False - -ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY = 'synchronize_checkpoint_boundary' -ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY_DEFAULT = False - -ACT_CHKPT_PROFILE = 'profile' -ACT_CHKPT_PROFILE_DEFAULT = False - -ACT_CHKPT_CPU_CHECKPOINTING = 'cpu_checkpointing' -ACT_CHKPT_CPU_CHECKPOINTING_DEFAULT = False - -ACT_CHKPT = 'activation_checkpointing' - -ACT_CHKPT_DEFAULT = { - ACT_CHKPT_PARTITION_ACTIVATIONS: ACT_CHKPT_PARTITION_ACTIVATIONS_DEFAULT, - ACT_CHKPT_NUMBER_CHECKPOINTS: ACT_CHKPT_NUMBER_CHECKPOINTS_DEFAULT, - ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION: ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION_DEFAULT, - ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY: ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY_DEFAULT, - ACT_CHKPT_PROFILE: ACT_CHKPT_PROFILE_DEFAULT, - ACT_CHKPT_CPU_CHECKPOINTING: ACT_CHKPT_CPU_CHECKPOINTING_DEFAULT -} - - -class DeepSpeedActivationCheckpointingConfig(DeepSpeedConfigObject): - - def __init__(self, param_dict): - super(DeepSpeedActivationCheckpointingConfig, self).__init__() - - self.partition_activations = None - self.contiguous_memory_optimization = None - self.cpu_checkpointing = None - self.number_checkpoints = None - self.synchronize_checkpoint_boundary = None - self.profile = None - - if ACT_CHKPT in param_dict.keys(): - act_chkpt_config_dict = param_dict[ACT_CHKPT] - else: - act_chkpt_config_dict = ACT_CHKPT_DEFAULT - - self._initialize(act_chkpt_config_dict) - - def _initialize(self, act_chkpt_config_dict): - self.partition_activations = get_scalar_param(act_chkpt_config_dict, ACT_CHKPT_PARTITION_ACTIVATIONS, - ACT_CHKPT_PARTITION_ACTIVATIONS_DEFAULT) - - self.contiguous_memory_optimization = get_scalar_param(act_chkpt_config_dict, - ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION, - ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION_DEFAULT) - - self.cpu_checkpointing = get_scalar_param(act_chkpt_config_dict, ACT_CHKPT_CPU_CHECKPOINTING, - ACT_CHKPT_CPU_CHECKPOINTING_DEFAULT) - - self.number_checkpoints = get_scalar_param(act_chkpt_config_dict, ACT_CHKPT_NUMBER_CHECKPOINTS, - ACT_CHKPT_NUMBER_CHECKPOINTS_DEFAULT) - - self.profile = get_scalar_param(act_chkpt_config_dict, ACT_CHKPT_PROFILE, ACT_CHKPT_PROFILE_DEFAULT) - - self.synchronize_checkpoint_boundary = get_scalar_param(act_chkpt_config_dict, - ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY, - ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY_DEFAULT) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/base_optimizer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/base_optimizer.py deleted file mode 100644 index 726124027131c4edf110ef0c080e07abcf5f19ff..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/base_optimizer.py +++ /dev/null @@ -1,81 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import torch - -from deepspeed.utils import logger -from deepspeed.utils.tensor_fragment import map_to_flat_opt_states -from deepspeed.runtime.utils import bwc_tensor_model_parallel_rank, see_memory_usage -from deepspeed.runtime.torch_autocast import get_autocast_dtype, is_autocast_initialized - - -class DeepSpeedOptimizer(object): - pass - - -class ZeROOptimizer(DeepSpeedOptimizer): - - def load_hp_checkpoint_state_from_checkpoint_dir(self, lp_groups_name: str, checkpoint_dir: str) -> None: - checkpoint_dir = os.path.join(checkpoint_dir, "zero") - optim_state_path = os.path.join(checkpoint_dir, "optimizer_state.pt") - assert os.path.isfile( - optim_state_path), f'{optim_state_path} containing optimizer global state is missing! Cannot proceed.' - optim_sd = torch.load(optim_state_path, weights_only=False) - - self._load_global_state(optim_sd) - - tp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu) - if self.mpu is None: - logger.warning("MPU is not provided, setting tp size to 1 in checkpoint loading.") - tp_world_size = 1 - else: - tp_world_size = self.mpu.get_slice_parallel_world_size() if hasattr(self.mpu, "get_slice_parallel_world_size") \ - else self.mpu.get_tensor_model_parallel_world_size() - - for i, (param_group, - loaded_param_group) in enumerate(zip(self.optimizer.param_groups, optim_sd['param_groups'])): - # We have an assumption that all params in the same param_group have the same keys - opt_keys = set() - steps = [] - - lp_groups = getattr(self, lp_groups_name) - for lp in lp_groups[i]: - if lp._hp_mapping is not None: - #print(f"Loading {self.param_names[lp]} {tp_rank=} {tp_world_size=}") - step = lp.load_hp_checkpoint_state(os.path.join(checkpoint_dir, self.param_names[lp]), tp_rank, - tp_world_size) - for key in lp._hp_mapping.get_optim_state_keys(): - opt_keys.add(key) - steps.append(step) - - hp_param = param_group['params'][0] - assert all(step == steps[0] for step in steps), f"Steps {steps} are not equal" - if steps[0] is not None: - self.optimizer.state[hp_param]['step'] = steps[0] - - map_to_flat_opt_states(hp_param, lp_groups[i], self.optimizer.state, opt_keys) - - for key, value in loaded_param_group.items(): - if key == 'params': - continue - param_group[key] = value - - def report_ipg_memory_usage(self, tag, param_elems, dtype=None): - dtypes = self.ipg_buckets.keys() if dtype is None else [dtype] - - for dt in dtypes: - bucket = self.ipg_buckets[dt] - elem_count = bucket.elements + param_elems - percent_of_bucket_size = (100.0 * elem_count) // self.reduce_bucket_size - see_memory_usage( - f"{tag}: elems in_bucket {dt} {bucket.elements} param {param_elems} max_percent {percent_of_bucket_size}" - ) - - def get_param_comm_dtype(self, param): - if is_autocast_initialized(): - return get_autocast_dtype(param) - else: - return self.communication_data_type diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/bf16_optimizer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/bf16_optimizer.py deleted file mode 100644 index cd00e9e43b00d03c7b66c445afa68de6700579a1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/bf16_optimizer.py +++ /dev/null @@ -1,566 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from collections import OrderedDict -import torch -import sys -from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors -from deepspeed import comm as dist -from deepspeed.runtime.constants import PIPE_REPLICATED -from deepspeed.runtime.base_optimizer import ZeROOptimizer -from packaging import version as pkg_version -from deepspeed.git_version_info import version -from deepspeed.runtime.utils import (get_global_norm_of_tensors, clip_tensors_by_global_norm, DummyOptim, - align_dense_tensors, all_gather_dp_groups, is_model_parallel_parameter, - see_memory_usage, graph_process, get_norm_with_moe_layers) -from deepspeed.utils import link_hp_params, lazy_init_hp_params_optimizer_state, fragment_address, groups -from deepspeed.moe.utils import is_moe_param, is_moe_param_group -from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank -from deepspeed.utils.torch import register_grad_hook -from deepspeed.checkpoint import enable_universal_checkpoint -from deepspeed.checkpoint.constants import (DS_VERSION, PARTITION_COUNT, BASE_OPTIMIZER_STATE, - SINGLE_PARTITION_OF_FP32_GROUPS, CLIP_GRAD, GROUP_PADDINGS, - PARAM_SLICE_MAPPINGS) - -setattr(sys.modules[__name__], 'fragment_address', fragment_address) - - -def print_rank_0(message, debug=False, force=False): - if dist.get_rank() == 0 and (debug or force): - print(message) - - -class BF16_Optimizer(ZeROOptimizer): - - def __init__(self, - init_optimizer, - param_names, - bfloat16_config, - mpu=None, - clip_grad=0.0, - norm_type=2, - allgather_bucket_size=5000000000, - dp_process_group=None, - timers=None, - grad_acc_dtype=None, - graph_harvesting=False, - has_moe_layers=False): - super().__init__() - see_memory_usage('begin bf16_optimizer', force=True) - self.timers = timers - self.optimizer = init_optimizer - self.param_names = param_names - self.using_real_optimizer = not isinstance(self.optimizer, DummyOptim) - - assert bfloat16_config.enabled, f"BF16Optimizer: requires bfloat16 to be enabled" - assert grad_acc_dtype in [torch.float32, torch.bfloat16 - ], f"BF16Optimizer: Unsupported gradient accumulation data type: {grad_acc_dtype}" - self.grad_acc_dtype = grad_acc_dtype - - self.immediate_grad_update = bfloat16_config.immediate_grad_update - - self.clip_grad = clip_grad - self.norm_type = norm_type - self.mpu = mpu - self.allgather_bucket_size = int(allgather_bucket_size) - self.dp_process_group = dp_process_group - self.dp_rank = dist.get_rank(group=self.dp_process_group) - self.has_moe_layers = has_moe_layers - self.non_expert_gradients = [] - self.real_dp_process_group = [dp_process_group for i in range(len(self.optimizer.param_groups))] - if self.has_moe_layers: - self._configure_moe_settings() - - # Use torch (un)flatten ops - self.flatten = _flatten_dense_tensors - self.unflatten = _unflatten_dense_tensors - - #align nccl all-gather send buffers to 4-bye boundary - self.nccl_start_alignment_factor = 2 # 4-byte alignment/sizeof(fp16) = 2 - - # Build BF16/FP32 groups - self.bf16_groups = [] - self.bf16_groups_flat = [] - self.bf16_partitioned_groups = [] - - self.fp32_groups_flat_partition = [] - - # Maintain different fp32 gradients views for convenience - self.fp32_groups_gradients = [] - self.fp32_groups_gradient_dict = {} - self.fp32_groups_gradients_flat = [] - self.fp32_groups_actual_gradients_flat = [] - self.fp32_groups_gradient_flat_partition = [] - self.fp32_groups_has_gradients = [] - - self.group_paddings = [] - self.graph_harvesting = graph_harvesting - if self.using_real_optimizer: - self._setup_for_real_optimizer() - - see_memory_usage('end bf16_ optimizer', force=True) - - def destroy(self): - for i, _ in enumerate(self.optimizer.param_groups): - for p in self.bf16_groups[i]: - if getattr(p, '_hp_mapping', None): - p._hp_mapping = None - for hook in self._grad_acc_hooks: - hook.remove() - print_rank_0("Removed grad acc hooks") - - def _configure_moe_settings(self): - assert any( - [is_moe_param_group(group) for group in self.optimizer.param_groups] - ), "The model has moe layers, but None of the param groups are marked as MoE. Create a param group with 'moe' key set to True before creating optimizer" - - for i, group in enumerate(self.optimizer.param_groups): - if is_moe_param_group(group): - assert all([is_moe_param(param) - for param in group['params']]), "All params in MoE group must be MoE params" - self.real_dp_process_group[i] = groups._get_expert_data_parallel_group(group['name']) - self.expert_gradients = {} - if self.has_moe_layers: - for key in groups._get_expert_data_parallel_group_dict().keys(): - self.expert_gradients[key] = [] - - def _setup_for_real_optimizer(self): - self.partition_count = [dist.get_world_size(group=pg) for pg in self.real_dp_process_group] - - for i, param_group in enumerate(self.optimizer.param_groups): - real_dp_world_size = dist.get_world_size(group=self.real_dp_process_group[i]) - see_memory_usage(f'before initializing group {i}', force=True) - - partition_id = dist.get_rank(group=self.real_dp_process_group[i]) - - # grab the original list - trainable_parameters = [param for param in param_group['params'] if param.requires_grad] - self.bf16_groups.append(trainable_parameters) - - # create flat bf16 params - self.bf16_groups_flat.append( - self._flatten_dense_tensors_aligned(self.bf16_groups[i], - self.nccl_start_alignment_factor * real_dp_world_size)) - # Make bf16 params point to flat tensor storage - self._update_storage_to_flattened_tensor(tensor_list=self.bf16_groups[i], - flat_tensor=self.bf16_groups_flat[i]) - - # divide flat weights into equal sized partitions - partition_size = self.bf16_groups_flat[i].numel() // real_dp_world_size - bf16_dp_partitions = [ - self.bf16_groups_flat[i].narrow(0, dp_index * partition_size, partition_size) - for dp_index in range(real_dp_world_size) - ] - self.bf16_partitioned_groups.append(bf16_dp_partitions) - - # create fp32 params partition - self.fp32_groups_flat_partition.append(bf16_dp_partitions[partition_id].clone().float().detach()) - self.fp32_groups_flat_partition[i].requires_grad = True - - num_elem_list = [t.numel() for t in self.bf16_groups[i]] - - # create fp32 gradients - fp32_flat_buffer = torch.zeros_like(self.bf16_groups_flat[i], dtype=self.grad_acc_dtype) - self.fp32_groups_gradients_flat.append(fp32_flat_buffer) - if self.has_moe_layers and is_moe_param_group(param_group): - self.expert_gradients[param_group['name']].append(fp32_flat_buffer) - else: - self.non_expert_gradients.append(fp32_flat_buffer) - - # track individual fp32 gradients for entire model - fp32_gradients = self._split_flat_tensor(flat_tensor=self.fp32_groups_gradients_flat[i], - num_elem_list=num_elem_list) - self.fp32_groups_gradients.append(fp32_gradients) - self.fp32_groups_gradient_dict[i] = fp32_gradients - - # flat tensor corresponding to actual fp32 gradients (i.e., minus alignment padding) - length_without_padding = sum(num_elem_list) - self.fp32_groups_actual_gradients_flat.append( - torch.narrow(self.fp32_groups_gradients_flat[i], 0, 0, length_without_padding)) - - # flat tensor corresponding to gradient partition - self.fp32_groups_gradient_flat_partition.append( - torch.narrow(self.fp32_groups_gradients_flat[i], 0, partition_id * partition_size, partition_size)) - - # track fp32 gradient updates - self.fp32_groups_has_gradients.append([False] * len(self.bf16_groups[i])) - - # Record padding required for alignment - if partition_id == dist.get_world_size(group=self.real_dp_process_group[i]) - 1: - padding = self.bf16_groups_flat[i].numel() - length_without_padding - else: - padding = 0 - - self.group_paddings.append(padding) - - # update optimizer param groups to reference fp32 params partition - param_group['params'] = [self.fp32_groups_flat_partition[i]] - - see_memory_usage(f'after initializing group {i}', force=True) - - self._grad_acc_hooks = [] - if self.immediate_grad_update: - self.create_grad_acc_hooks() - - # Need optimizer states initialized before linking lp to optimizer state - self._link_all_hp_params() - self._hp_optimizer_states_linked = False - self._enable_universal_checkpoint() - self._param_slice_mappings = self._create_param_mapping() - - def _enable_universal_checkpoint(self): - for lp_param_group in self.bf16_groups: - enable_universal_checkpoint(param_list=lp_param_group) - - def _create_param_mapping(self): - param_mapping = [] - for i, _ in enumerate(self.optimizer.param_groups): - param_mapping_per_group = OrderedDict() - for lp in self.bf16_groups[i]: - if lp._hp_mapping is not None: - lp_name = self.param_names[lp] - param_mapping_per_group[lp_name] = lp._hp_mapping.get_hp_fragment_address() - param_mapping.append(param_mapping_per_group) - - return param_mapping - - def _link_all_hp_params(self): - for i, _ in enumerate(self.optimizer.param_groups): - real_dp_world_size = dist.get_world_size(group=self.real_dp_process_group[i]) - - # Link bf16 and fp32 params in partition - partition_id = dist.get_rank(group=self.real_dp_process_group[i]) - partition_size = self.bf16_groups_flat[i].numel() // real_dp_world_size - flat_hp_partition = self.fp32_groups_flat_partition[i] - link_hp_params(lp_param_list=self.bf16_groups[i], - flat_hp_partition=flat_hp_partition, - gradient_dict=self.fp32_groups_gradient_dict, - offload_gradient_dict=None, - use_offload=False, - param_group_index=i, - partition_start=partition_id * partition_size, - partition_size=partition_size, - dp_group=self.real_dp_process_group[i]) - - def _lazy_init_hp_params_optimizer_state(self): - if not self._hp_optimizer_states_linked: - for i, _ in enumerate(self.optimizer.param_groups): - lazy_init_hp_params_optimizer_state(self.bf16_groups[i], self.fp32_groups_flat_partition[i], - self.optimizer.state) - self._hp_optimizer_states_linked = True - - def _split_flat_tensor(self, flat_tensor, num_elem_list): - assert sum(num_elem_list) <= flat_tensor.numel() - tensor_list = [] - offset = 0 - for num_elem in num_elem_list: - dense_tensor = torch.narrow(flat_tensor, 0, offset, num_elem) - tensor_list.append(dense_tensor) - offset += num_elem - - return tensor_list - - def _update_storage_to_flattened_tensor(self, tensor_list, flat_tensor): - updated_params = self.unflatten(flat_tensor, tensor_list) - for p, q in zip(tensor_list, updated_params): - p.data = q.data - - def _flatten_dense_tensors_aligned(self, tensor_list, alignment): - return self.flatten(align_dense_tensors(tensor_list, alignment)) - - @torch.no_grad() - def step(self, closure=None): - if closure is not None: - raise NotImplementedError(f'{self.__class__} does not support closure.') - - non_expert_grads_for_norm, expert_grads_for_norm = self.get_grads_for_norm() - non_expert_groups_norm = get_global_norm_of_tensors(input_tensors=non_expert_grads_for_norm, - mpu=self.mpu, - norm_type=self.norm_type, - use_graph=self.graph_harvesting) - all_groups_norm = non_expert_groups_norm - if self.has_moe_layers: - all_groups_norm = get_norm_with_moe_layers(non_expert_groups_norm, - mpu=self.mpu, - expert_tensors=expert_grads_for_norm, - norm_type=self.norm_type) - - self._global_grad_norm = all_groups_norm - - assert all_groups_norm > 0. - if self.clip_grad > 0.: - clip_tensors_by_global_norm(input_tensors=self.get_grads_for_norm(for_clipping=True), - max_norm=self.clip_grad, - global_norm=all_groups_norm, - mpu=self.mpu, - use_graph=self.graph_harvesting) - - for param_partition, grad_partition in zip(self.fp32_groups_flat_partition, - self.fp32_groups_gradient_flat_partition): - # In case of grad acc dtype different than FP32, need to cast to high precision. - param_partition.grad = grad_partition.to( - param_partition.dtype) if grad_partition.dtype != param_partition.dtype else grad_partition - - self.optimizer.step() - - if self.grad_acc_dtype is not torch.float32: - for param_partition in self.fp32_groups_flat_partition: - param_partition.grad = None - - # We need to link optimizer state after the first step() call - self._lazy_init_hp_params_optimizer_state() - - self.update_lp_params() - - self.clear_hp_grads() - - def backward(self, loss, retain_graph=False, update_hp_grads=True, clear_lp_grads=False, **bwd_kwargs): - """Perform a backward pass and copy the low-precision gradients to the - high-precision copy. - - We copy/accumulate to the high-precision grads now to prevent accumulating in the - bf16 grads after successive backward() calls (i.e., grad accumulation steps > 1) - - The low-precision grads are deallocated during this procedure. - """ - self.clear_lp_grads() - loss.backward(retain_graph=retain_graph, **bwd_kwargs) - - if update_hp_grads: - self.update_hp_grads(clear_lp_grads=clear_lp_grads) - - @torch.no_grad() - def _update_hp_grad(self, lp, group_idx, param_idx, clear_lp_grads): - if lp.grad is None: - return - - hp_grad = self.fp32_groups_gradients[group_idx][param_idx] - assert hp_grad is not None, \ - f'high precision param has no gradient, lp param_id = {id(lp)} group_info = [{group_idx}][{param_idx}]' - - hp_grad.data.add_(lp.grad.data.to(hp_grad.dtype).view(hp_grad.shape)) - lp._hp_grad = hp_grad - self.fp32_groups_has_gradients[group_idx][param_idx] = True - - # clear gradients - if clear_lp_grads: - lp.grad.zero_() - - @torch.no_grad() - def _update_hp_grads_func(self, clear_lp_grads=False): - for i, group in enumerate(self.bf16_groups): - for j, lp in enumerate(group): - self._update_hp_grad(lp, i, j, clear_lp_grads) - - @torch.no_grad() - def update_hp_grads(self, clear_lp_grads=False): - if self.immediate_grad_update: - return - - if self.graph_harvesting: - graph_process(False, self._update_hp_grads_func, clear_lp_grads) - else: - self._update_hp_grads_func(clear_lp_grads) - #cpu op - for i, group in enumerate(self.bf16_groups): - for j, lp in enumerate(group): - if lp.grad is None: - continue - self.fp32_groups_has_gradients[i][j] = True - - @torch.no_grad() - def get_grads_for_reduction(self): - if self.has_moe_layers: - return self.non_expert_gradients, self.expert_gradients - return self.non_expert_gradients, {} - - @torch.no_grad() - def get_grads_for_norm(self, for_clipping=False): - """ - Returns: - tuple[list[Tensor], dict[ep_name, List[Tensor]] | list: - If for_clipping, return all gradients. - Otherwise, separate and return dict of expert_grad and list of non_expert_grad - """ - # (grads, expert_group_name) - expert_grads_for_norm = {} - - # grads - non_expert_grads_for_norm = [] - all_grads_for_clip = [] - - tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu) - assert len(self.bf16_groups) == len(self.optimizer.param_groups) - for i, group in enumerate(self.bf16_groups): - for j, lp in enumerate(group): - if not for_clipping: - if hasattr(lp, PIPE_REPLICATED) and lp.ds_pipe_replicated: - continue - - # skip duplicated parameters. perform norm only on cards with tp_rank=0. - # non-duplicated parameters include: - # - Parameters with tp: Use allreducesum of mp_group. - # - Moe Parameters with ep: Use allreducesum of ep_group. - if not (tensor_mp_rank == 0 or is_model_parallel_parameter(lp) or is_moe_param(lp)): - continue - - if not self.fp32_groups_has_gradients[i][j]: - continue - if not for_clipping: - param_group = self.optimizer.param_groups[i] - if self.has_moe_layers and is_moe_param_group(param_group): - if param_group['name'] not in expert_grads_for_norm: - expert_grads_for_norm[param_group['name']] = [] - expert_grads_for_norm[param_group['name']].append(self.fp32_groups_gradients[i][j]) - else: - non_expert_grads_for_norm.append(self.fp32_groups_gradients[i][j]) - else: - all_grads_for_clip.append(self.fp32_groups_gradients[i][j]) - if not for_clipping: - return non_expert_grads_for_norm, expert_grads_for_norm - return all_grads_for_clip - - @torch.no_grad() - def update_lp_params(self): - for i, (bf16_partitions, - fp32_partition) in enumerate(zip(self.bf16_partitioned_groups, self.fp32_groups_flat_partition)): - partition_id = dist.get_rank(group=self.real_dp_process_group[i]) - bf16_partitions[partition_id].data.copy_(fp32_partition.data) - - all_gather_dp_groups(groups_flat=self.bf16_groups_flat, - partitioned_param_groups=self.bf16_partitioned_groups, - dp_process_group=self.real_dp_process_group, - start_alignment_factor=self.nccl_start_alignment_factor, - allgather_bucket_size=self.allgather_bucket_size) - - def clear_hp_grads(self): - for flat_gradients in self.fp32_groups_gradients_flat: - flat_gradients.zero_() - - for i, group in enumerate(self.fp32_groups_gradients): - self.fp32_groups_has_gradients[i] = [False] * len(group) - - def clear_lp_grads(self, set_to_none=False): - - # using zero_() fixed memory address for graph replay - if self.graph_harvesting: - assert not set_to_none, "graph harvesting is incompatible with setting lp grads to None" - - zero_grads_list = [] - for group in self.bf16_groups: - for param in group: - if set_to_none: - param.grad = None - elif param.grad is not None: - if param.grad.grad_fn is not None: - param.grad.detach_() - zero_grads_list.append(param.grad) - if not set_to_none and len(zero_grads_list) > 0: - torch._foreach_zero_(zero_grads_list) - - def zero_grad(self, set_to_none=True): - self.clear_lp_grads(set_to_none) - self.clear_hp_grads() - - def state_dict(self): - state_dict = {} - state_dict[CLIP_GRAD] = self.clip_grad - state_dict[BASE_OPTIMIZER_STATE] = self.optimizer.state_dict() - state_dict[SINGLE_PARTITION_OF_FP32_GROUPS] = self.fp32_groups_flat_partition - state_dict[GROUP_PADDINGS] = self.group_paddings - state_dict[PARTITION_COUNT] = self.partition_count - state_dict[DS_VERSION] = version - state_dict[PARAM_SLICE_MAPPINGS] = self._param_slice_mappings - - return state_dict - - # Restore base optimizer fp32 weights bfloat16 weights - def _restore_from_bit16_weights(self): - for i, (bf16_partitions, - fp32_partition) in enumerate(zip(self.bf16_partitioned_groups, self.fp32_groups_flat_partition)): - partition_id = dist.get_rank(group=self.real_dp_process_group[i]) - fp32_partition.data.copy_(bf16_partitions[partition_id].data) - - def refresh_fp32_params(self): - self._restore_from_bit16_weights() - - def load_state_dict(self, - state_dict_list, - checkpoint_folder=None, - load_optimizer_states=True, - load_from_fp32_weights=False, - load_serial=None, - param_shapes=None): - if checkpoint_folder: - self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights) - else: - self._load_legacy_checkpoint(state_dict_list, load_optimizer_states, load_from_fp32_weights) - - def _load_legacy_checkpoint(self, state_dict_list, load_optimizer_states=True, load_from_fp32_weights=False): - - dp_rank = dist.get_rank(group=self.dp_process_group) - current_rank_sd = state_dict_list[dp_rank] - - ckpt_version = current_rank_sd.get(DS_VERSION, False) - assert ckpt_version, f"Empty ds_version in checkpoint, not clear how to proceed" - ckpt_version = pkg_version.parse(ckpt_version) - - self.clip_grad = current_rank_sd.get(CLIP_GRAD, self.clip_grad) - - if load_optimizer_states: - print(f"_load_legacy_checkpoint current_rank_sd[BASE_OPTIMIZER_STATE]") - self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE]) - - if load_from_fp32_weights: - for current, saved in zip(self.fp32_groups_flat_partition, - current_rank_sd[SINGLE_PARTITION_OF_FP32_GROUPS]): - src_tensor = _get_padded_tensor(saved, current.numel()) - current.data.copy_(src_tensor.data) - - if load_optimizer_states: - self._link_all_hp_params() - - def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights): - self.load_hp_checkpoint_state_from_checkpoint_dir("bf16_groups", checkpoint_folder) - - def _load_global_state(self, sd): - pass - - @property - def param_groups(self): - """Forward the wrapped optimizer's parameters.""" - return self.optimizer.param_groups - - @property - def state(self): - """Forward the wrapped optimizer's states.""" - return self.optimizer.state - - def accumulate_hp_grads_and_remove_lp(self, lp_param, group_idx, param_idx): - assert self.immediate_grad_update - self._update_hp_grad(lp_param, group_idx, param_idx, clear_lp_grads=False) - - def create_grad_acc_hooks(self): - for i, param_group in enumerate(self.bf16_groups): - for j, param in enumerate(param_group): - if param.requires_grad: - - def wrapper(param, i, j): - - def accumulate_hp_grads_and_remove_lp(*notneeded): - self.accumulate_hp_grads_and_remove_lp(param, i, j) - - self._grad_acc_hooks.append(register_grad_hook(param, accumulate_hp_grads_and_remove_lp)) - - wrapper(param, i, j) - - -def _get_padded_tensor(src_tensor, size): - if src_tensor.numel() >= size: - return src_tensor - padded_tensor = torch.zeros(size, dtype=src_tensor.dtype, device=src_tensor.device) - slice_tensor = torch.narrow(padded_tensor, 0, 0, src_tensor.numel()) - slice_tensor.data.copy_(src_tensor.data) - return padded_tensor diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__init__.py deleted file mode 100644 index a6484db58e0608eafd7ec33ba871c5982c939aed..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' - -from .fast_checkpoint_engine import FastCheckpointEngine -from .torch_checkpoint_engine import TorchCheckpointEngine -from .decoupled_checkpoint_engine import DecoupledCheckpointEngine -from .checkpoint_engine import CheckpointCommitInfo -from .utils import create_checkpoint_engine diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 6f7092ee2fc910edeb981d2effee7953a47f46ce..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/checkpoint_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/checkpoint_engine.cpython-310.pyc deleted file mode 100644 index dbf77ba59e77a173d0c048ebf1aaca8122f75969..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/checkpoint_engine.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/decoupled_checkpoint_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/decoupled_checkpoint_engine.cpython-310.pyc deleted file mode 100644 index 41cc9c69b73767b55b29f4601d09a487a8ba5df6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/decoupled_checkpoint_engine.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/fast_checkpoint_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/fast_checkpoint_engine.cpython-310.pyc deleted file mode 100644 index 231f221a9ab25f40780795c4735fcd4e3f5e9ed7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/fast_checkpoint_engine.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/nebula_checkpoint_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/nebula_checkpoint_engine.cpython-310.pyc deleted file mode 100644 index 42adbf4b7cea1423d2485ed9b047d73845037ee6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/nebula_checkpoint_engine.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/torch_checkpoint_engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/torch_checkpoint_engine.cpython-310.pyc deleted file mode 100644 index 786b24e04ee7abda9d7147d5b6ce8f8878779389..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/torch_checkpoint_engine.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index ff792fd1f460138608352c27dbacb8456de25b0b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/checkpoint_engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/checkpoint_engine.py deleted file mode 100644 index 03254b935a2ac688ab0790e2ac0bc5ee88398855..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/checkpoint_engine.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os - -import abc -from abc import ABC - -from dataclasses import dataclass - - -@dataclass -class CheckpointCommitInfo(object): - tag: str - save_dir: str - save_latest: bool - - -class CheckpointEngine(ABC): - # init checkpoint engine for save/load - def __init__(self, config_params=None): - self.name = None - - @abc.abstractmethod - def create(self, info: CheckpointCommitInfo): - # create checkpoint on give tag for save/load. - ... - - @abc.abstractmethod - def save(self, state_dict, path: str): - ... - - def makedirs(self, path, exist_ok=False): - os.makedirs(path, exist_ok=exist_ok) - - @abc.abstractmethod - def load(self, path: str, map_location=None): - ... - - @abc.abstractmethod - def commit(self, info: CheckpointCommitInfo): - # to tell checkpoint services if all files are ready. - ... - - def is_data_parallel_writer(self, dp_rank): - return dp_rank == 0 - - def is_decoupled(self): - return False - - def set_commit_info(self, info: CheckpointCommitInfo): - pass - - def get_commit_info(self): - return None - - def cleanup(self): - pass diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/decoupled_checkpoint_engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/decoupled_checkpoint_engine.py deleted file mode 100644 index d6c0ea88836e376deb54a72f2157c88f907fc52f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/decoupled_checkpoint_engine.py +++ /dev/null @@ -1,160 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import torch.multiprocessing as mp -from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \ - CheckpointEngine, CheckpointCommitInfo -from deepspeed.runtime.checkpoint_engine.fast_checkpoint_engine import FastCheckpointEngine -from deepspeed import comm as dist -from deepspeed.runtime.utils import get_checkpoint_folder_size - -from enum import Enum - - -class DecoupledEvent(Enum): - SAVE_EVENT = 1 - COMMIT_EVENT = 2 - EXIT_EVENT = 3 - - -class CheckpointSize(object): - - def __init__(self): - self._pre = None - self._post = None - self._gigabytes = None - - def gb_size(self): - return self._gigabytes - - def set_pre_size(self, size): - self._pre = size - - def set_post_size(self, size): - self._post = size - self._gigabytes = (self._post - self._pre) / (1024**3) - - -def init_decoupled_checkpoint(config_params, dp_writer_config, save_event, save_queue, optimize_dp_state): - checkpoint_engine = FastCheckpointEngine(config_params, dp_writer_config, optimize_dp_state) - print(f'Created FastCheckpointEngine for Decoupled Checkpointing') - save_path_list = [] - while True: - (save_info, event_type) = save_queue.get() - if event_type == DecoupledEvent.SAVE_EVENT and save_info is not None: - state_dict, save_path = save_info - # print(f'Received decoupled checkpoint request for {save_path=}') - save_path_list.append(save_path) - checkpoint_engine.save(state_dict, save_path) - del state_dict - # print(f'Completed decoupled checkpoint request for {save_path=}') - - if event_type == DecoupledEvent.COMMIT_EVENT: - # print(f'Recieved commit request for {save_path_list=}') - save_path_list = [] - save_event.set() - - if event_type == DecoupledEvent.EXIT_EVENT: - # print(f'Received decoupled exit request') - break - - -ENGINE_NAME = "DecoupledCheckpointEngine" - - -class DecoupledCheckpointEngine(CheckpointEngine): - - def __init__(self, config_params, dp_writer_config, optimize_dp_state): - if mp.get_start_methods(allow_None=False) is None: - mp.set_start_method('spawn') - super().__init__(config_params) - self.name = ENGINE_NAME - self.dp_writer_config = dp_writer_config - self.commit_info = None - self.checkpoint_size = CheckpointSize() - self.global_rank = dist.get_rank() - self.optimize_dp_state = optimize_dp_state - if dp_writer_config is None: - self.save_event = None - self.save_queue = None - self.ckpt_process = None - self.local_rank = None - print( - f'[{ENGINE_NAME}]: No checkpoint process self.global_rank={self.global_rank} self.dp_writer_config={self.dp_writer_config}' - ) - else: - self.save_event = mp.Event() - self.save_queue = mp.SimpleQueue() - engine_args = (config_params, dp_writer_config, self.save_event, self.save_queue, self.optimize_dp_state) - self.ckpt_process = mp.Process(target=init_decoupled_checkpoint, args=engine_args) - self.ckpt_process.start() - self.local_rank = dp_writer_config.local_rank - print( - f'[{ENGINE_NAME}]: Create checkpoint process self.global_rank={self.global_rank} self.ckpt_process.pid={self.ckpt_process.pid} self.dp_writer_config={self.dp_writer_config}' - ) - - def __del__(self): - self.cleanup() - - def create(self, info: CheckpointCommitInfo): - self.commit_info = info - if self.checkpoint_size.gb_size() is None: - pre_size = get_checkpoint_folder_size(info.save_dir, info.tag, self.local_rank) - self.checkpoint_size.set_pre_size(pre_size) - - def load(self, path: str, map_location=None): - sd = torch.load(path, map_location=map_location) - return sd - - def save(self, state_dict, path: str): - if self.ckpt_process is None: - return - save_info = (state_dict, path) - self.save_queue.put((save_info, DecoupledEvent.SAVE_EVENT)) - - def commit(self, info: CheckpointCommitInfo): - assert info == self.commit_info - if self.ckpt_process is not None: - self.save_queue.put((None, DecoupledEvent.COMMIT_EVENT)) - # print(f'[begin] wait for decoupled complete for {info.tag}') - self.save_event.wait() - # print(f'[end] wait for decoupled complete for {info.tag}') - self.save_event.clear() - self.commit_info = None - - if self.checkpoint_size.gb_size() is None: - dist.barrier() - post_size = get_checkpoint_folder_size(info.save_dir, info.tag, self.local_rank) - self.checkpoint_size.set_post_size(post_size) - - if self.global_rank == 0: - print( - f'{self.name} self.global_rank={self.global_rank} created checkpoint of {round(self.checkpoint_size.gb_size(), 2)} GB' - ) - - return True - - def get_commit_info(self): - # print(f'getting commit info {self.commit_info=}') - return self.commit_info - - def is_decoupled(self): - return True - - def cleanup(self): - # print(f'Inside {self.name} cleanup') - - if self.get_commit_info() is not None: - self.commit(self.commit_info) - - if self.ckpt_process is not None: - self.save_queue.put((None, DecoupledEvent.EXIT_EVENT)) - self.ckpt_process.join() - self.ckpt_process = None - self.save_queue = None - - def is_data_parallel_writer(self, dp_rank): - return self.ckpt_process is not None diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/fast_checkpoint_engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/fast_checkpoint_engine.py deleted file mode 100644 index 4bfecf810dc7984fe1aa412763df915cd56fb8ca..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/fast_checkpoint_engine.py +++ /dev/null @@ -1,50 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \ - CheckpointEngine, CheckpointCommitInfo -from deepspeed.runtime.model_checkpointing import ( - CHECKPOINT_WRITER, - CHECKPOINT_SERIALIZATION, - CheckpointWriterFactory, -) - - -class FastCheckpointEngine(CheckpointEngine): - - def __init__(self, config_params, dp_writer_config, optimize_dp_state): - super().__init__(config_params) - self.name = 'FastCheckpointEngine' - self.serialization_enabled = config_params.checkpoint_config[CHECKPOINT_SERIALIZATION] - self.optimize_dp_state = optimize_dp_state - if dp_writer_config is None: - self._writer = None - else: - self._writer = CheckpointWriterFactory(writer_config=config_params.checkpoint_config[CHECKPOINT_WRITER], - aio_config=config_params.aio_config, - dp_writer_config=dp_writer_config) - - def create(self, info: CheckpointCommitInfo): - pass - - def save(self, state_dict, path: str): - if self._writer is None: - return - - torch.save(obj=state_dict, - f=self._writer.create_writer(path, self.optimize_dp_state), - _use_new_zipfile_serialization=self.serialization_enabled) - self._writer.release_writer() - - def load(self, path: str, map_location=None): - sd = torch.load(path, map_location=map_location) - return sd - - def commit(self, info: CheckpointCommitInfo): - return True - - def is_data_parallel_writer(self, dp_rank): - return self._writer is not None diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/nebula_checkpoint_engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/nebula_checkpoint_engine.py deleted file mode 100644 index 485e21411fb8332a167458770de31d71680e678d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/nebula_checkpoint_engine.py +++ /dev/null @@ -1,108 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import torch -import torch_nebula - -from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \ - CheckpointEngine, CheckpointCommitInfo -from deepspeed.utils import logger, log_dist -from deepspeed.nebula.constants import * - - -def _get_tag_from_path(path): - return os.path.basename(os.path.dirname(path)) - - -class NebulaCheckpointEngine(CheckpointEngine): - - def __init__(self, config_params=None): - super().__init__(config_params) - self.name = "NebulaCheckpointEngine" - self.checkpoint = None - self.tag_flag = None - self.enable_nebula_load = config_params.enable_nebula_load - self.nebula_load_path = config_params.load_path - if self.nebula_load_path is None: - self.nebula_load_path = config_params.persistent_storage_path - - nebula_config_params = { - NEBULA_PERSISTENT_STORAGE_PATH: config_params.persistent_storage_path, - NEBULA_PERSISTENT_TIME_INTERVAL: config_params.persistent_time_interval, - NEBULA_NUM_OF_VERSION_IN_RETENTION: config_params.num_of_version_in_retention, - } - torch_nebula.init(**nebula_config_params) - - def create(self, info: CheckpointCommitInfo): - log_dist(f"[Nebula] Start Checkpoint for tag:{info.tag}", ranks=[0]) - # -2 means: customer needs to explicitly tell nebula - # current checkpoint is complete by commit methond. - self.checkpoint = torch_nebula.Checkpoint(info.tag, -2) - - def save(self, state_dict, path: str): - log_dist(f"[Nebula] Create dummy files for loading.") - torch.save("", path) - - tag = _get_tag_from_path(path) - partititon_name = os.path.basename(path) - logger.info(f"[Nebula] Saving {partititon_name} under tag {tag}...") - self.checkpoint.save(partititon_name, state_dict) - logger.info(f"[Nebula] Saved {partititon_name} under tag {tag}.") - - def load(self, path: str, map_location=None): - tag = _get_tag_from_path(path) - first_load_flag = self.tag_flag is None or self.tag_flag == tag - if not self.enable_nebula_load and first_load_flag: - self.tag_flag = tag - logger.info(f"[Nebula] Disable nebula load. Loading checkpoint from {path} ...") - partition = torch.load(path, map_location=map_location, weights_only=False) - logger.info(f"[Nebula] Disable nebula load. Loaded checkpoint from {path} .") - return partition - - partition_name = os.path.basename(path) - logger.info(f"[Nebula] Loading {path} under tag {tag} from nebula path {self.nebula_load_path}...") - - checkpoint = None - if tag in (None, 'latest', 'latest_universal'): - # In some cases, there is the inconsistent tag between deepspeed metadata (latest file) - # and nebula metadata, will lead to the failure on loading with deepspeed tag. Then we - # will try to load the valid latest checkpoint from nebula(tier3 > tier1). So, in summary - # when met failure loading for given tag, the loading priority would be like: - # nebula tier3 latest > nebula tier1 latest. - checkpoint = torch_nebula.get_latest_checkpoint(persist_path=self.nebula_load_path) - else: - checkpoint = torch_nebula.get_checkpoint(tag=tag, persist_path=self.nebula_load_path) - - if checkpoint is None or (checkpoint is not None and checkpoint.tag == ''): - logger.info( - f"Unable to find valid checkpoint tag:{tag} from Nebula, try to get latest checkpoint again from nebula {self.nebula_load_path} path!" - ) - # nebula tier3 latest - checkpoint = torch_nebula.get_latest_checkpoint(persist_path=self.nebula_load_path) - if checkpoint is None or (checkpoint is not None and checkpoint.tag == ''): - logger.info( - f"Unable to find latest checkpoint from Nebula tier3, try to get latest checkpoint again from nebula tier1 path!" - ) - # nebula tier1 latest - checkpoint = torch_nebula.get_latest_checkpoint() - logger.warning(f"Unable to find valid checkpoint from Nebula under tag:{tag}.") - return None - - tag = checkpoint.tag - self.tag_flag = -1 - partition = checkpoint.load(partition_name, map_location=map_location) - logger.info(f"[Nebula] Loaded {path} under tag {tag} from {self.nebula_load_path}.") - return partition - - def commit(self, info: CheckpointCommitInfo): - tag = info.tag - # nebula commit will be call when all files under give tag are ready to be persisted in the async way. - logger.info(f"[Nebula] all files for {tag} are saved in tier1. It is ready to start persisting") - commit_rls = self.checkpoint.commit() - if not commit_rls: - logger.error(f"[Nebula] failed to commit the checkpoint, please check the log.") - return False - return commit_rls diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/torch_checkpoint_engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/torch_checkpoint_engine.py deleted file mode 100644 index 9014677f1d5c5a50f4df5fc011a11489924b1e04..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/torch_checkpoint_engine.py +++ /dev/null @@ -1,43 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed.utils import log_dist -from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \ - CheckpointEngine, CheckpointCommitInfo -from deepspeed.runtime.model_checkpointing import CHECKPOINT_SERIALIZATION - -ENGINE_NAME = "TorchCheckpointEngine" - - -class TorchCheckpointEngine(CheckpointEngine): - - def __init__(self, config_params=None): - super().__init__(config_params) - self.name = ENGINE_NAME - if config_params is None: - self.zipfile_serialization = False - else: - self.zipfile_serialization = config_params.checkpoint_config[CHECKPOINT_SERIALIZATION] - log_dist(f'[{ENGINE_NAME}] Initialized with serialization = {self.zipfile_serialization}', ranks=[0]) - - def create(self, info: CheckpointCommitInfo): - log_dist(f"[Torch] Checkpoint {info.tag} is begin to save!", ranks=[0]) - pass - - def save(self, state_dict, path: str): - # log_dist(f"[Torch] Saving [begin] {path}... {self.zipfile_serialization=}", ranks=[0]) - torch.save(state_dict, path, _use_new_zipfile_serialization=self.zipfile_serialization) - # log_dist(f"[Torch] Saving [end] {path}... {self.zipfile_serialization=}", ranks=[0]) - - def load(self, path: str, map_location=None): - log_dist(f"[Torch] Begin Load checkpoint from {path}...", ranks=[0]) - partition = torch.load(path, map_location=map_location, weights_only=False) - log_dist(f"[Torch] End Load checkpoint from {path}...", ranks=[0]) - return partition - - def commit(self, info: CheckpointCommitInfo): - #logger.info(f"[Torch] Checkpoint {tag} is ready now!") - return True diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/utils.py deleted file mode 100644 index 8fa5e9725092338aee41b6951688c1cdad11b2e2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/utils.py +++ /dev/null @@ -1,38 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.runtime.model_checkpointing.constants import * -from deepspeed.runtime.model_checkpointing.utils import create_data_parallel_writer_config -from deepspeed.utils import logger - -from .decoupled_checkpoint_engine import DecoupledCheckpointEngine -from .fast_checkpoint_engine import FastCheckpointEngine -from .torch_checkpoint_engine import TorchCheckpointEngine - - -def create_checkpoint_engine(config_params, groups, zero_stage, has_moe_layers, optimize_dp_state): - if config_params is not None: - if config_params.checkpoint_config[CHECKPOINT_WRITER] is not None: - writer_config = config_params.checkpoint_config[CHECKPOINT_WRITER] - dp_writer_config = create_data_parallel_writer_config( - groups=groups, - parallel_unit=writer_config[CHECKPOINT_DATA_PARALLEL], - zero_stage=zero_stage, - has_moe_layers=has_moe_layers) - if writer_config[CHECKPOINT_WRITER_DECOUPLED]: - return DecoupledCheckpointEngine(config_params, dp_writer_config, optimize_dp_state) - else: - return FastCheckpointEngine(config_params, dp_writer_config, optimize_dp_state) - - if config_params is not None and config_params.nebula_config.enabled: - try: - from .nebula_checkpoint_engine import NebulaCheckpointEngine - except ImportError as err: - logger.error(f"No torch_nebula was found! Will fall back to torch.save. Details: {err}") - return TorchCheckpointEngine(config_params) - else: - return NebulaCheckpointEngine(config_params=config_params.nebula_config) - - return TorchCheckpointEngine(config_params) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 64e70db95100a0e60ca0eccdde878fa49a673ff5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/coalesced_collectives.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/coalesced_collectives.cpython-310.pyc deleted file mode 100644 index 56b53e7618b3cfd994037be7ef05147124312ed5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/coalesced_collectives.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/compressed.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/compressed.cpython-310.pyc deleted file mode 100644 index 54bd84ed04398619556000e2f9d47e742c10f085..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/compressed.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/hccl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/hccl.cpython-310.pyc deleted file mode 100644 index 427c567885a1b6b85127d06140a44209d7b02d9f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/hccl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/mpi.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/mpi.cpython-310.pyc deleted file mode 100644 index 0136194c61e3fb4a841750ab42c41bae2c92d803..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/mpi.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/nccl.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/nccl.cpython-310.pyc deleted file mode 100644 index d52e786fb516fd33de388a22abb1ff5fe7196f4a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/__pycache__/nccl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/coalesced_collectives.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/coalesced_collectives.py deleted file mode 100644 index 2fadce52222cda680600253ae840f84b89bda7ed..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/coalesced_collectives.py +++ /dev/null @@ -1,218 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -batched collective operations for overhead amortization and better -bandwidth utilization -""" - -import math -from typing import List, Any -import torch -from torch import Tensor -from deepspeed import comm as dist -from deepspeed.comm import ProcessGroup, all_to_all_single -from deepspeed.accelerator import get_accelerator -from deepspeed.utils import instrument_w_nvtx -from deepspeed.ops import op_builder -from deepspeed.utils import logger - - -def _torch_reduce_scatter_fn(input_tensor: Tensor, output_tensor: Tensor, group=None, async_op=False, prof=False): - return instrument_w_nvtx(dist.reduce_scatter_fn)(output_tensor, input_tensor, group=group, async_op=False) - - -quantizer_module = None - - -@instrument_w_nvtx -@torch.no_grad() -def all_to_all_quant_reduce(tensors: List[Tensor], groups: {}) -> List[Tensor]: - global quantizer_module - if quantizer_module is None: - quantizer_module = op_builder.QuantizerBuilder().load() - local_world_size = get_accelerator().device_count() - global_world_size = dist.get_world_size() - num_nodes = global_world_size // local_world_size - this_rank = dist.get_rank() - intra_idx = int(this_rank / local_world_size) - inter_idx = this_rank % local_world_size - output_lst: List[Tensor] = [None] * len(tensors) - for idx, tensor in enumerate(tensors): - if tensor.dim() == 1: - output_lst[idx] = reduce_scatter_coalesced([tensor])[0] - elif tensor.numel() % (2 * global_world_size) != 0: - # Due to the constraint of 2-stage all-to-all, the input tensor must be divisible by 2 * global_world_size - # Otherwise, all-to-all cannot be performed because of shape mismatch. - # See more at https://github.com/deepspeedai/DeepSpeed/pull/5056 - logger.warning( - f"qgZ falls back to reduce_scatter because tensor size = {tensor.numel()} is not divisible by (2 * global_world_size) = {2 * global_world_size}. Please consider allocating a new world to enable qgZ" - ) - output_lst[idx] = reduce_scatter_coalesced([tensor])[0] - else: - intra_quant_group = max(tensor.shape[0], tensor.shape[1], global_world_size) - - inter_quant_group = intra_quant_group // local_world_size - intra_quant_int4, intra_q_scales = quantizer_module.swizzle_quant(tensor, intra_quant_group, 4, - quantizer_module.Symmetric, 1, num_nodes, - local_world_size) - local_output = torch.empty_like(intra_quant_int4) - scale_output = torch.empty_like(intra_q_scales) - all_to_all_single(local_output, intra_quant_int4, group=groups[f'local_{intra_idx}']) - all_to_all_single(scale_output, intra_q_scales, group=groups[f'local_{intra_idx}']) - global_input_tensor, global_scales = quantizer_module.quantized_reduction( - local_output, scale_output, intra_quant_group, inter_quant_group, 4, quantizer_module.Symmetric, - local_world_size) - global_output = torch.empty_like(global_input_tensor) - global_scale_output = torch.empty_like(global_scales) - all_to_all_single(global_output, global_input_tensor, group=groups[f'global_{inter_idx}']) - all_to_all_single(global_scale_output, global_scales, group=groups[f'global_{inter_idx}']) - final_output = quantizer_module.dequantize(global_output, global_scale_output, global_scale_output.numel(), - 4, quantizer_module.Symmetric) - assert final_output.numel( - ) % num_nodes == 0, f"final_output.numel()={final_output.numel()} is not divisible by num_nodes={num_nodes}" - output_lst[idx] = (sum(list(final_output.chunk(num_nodes))) / num_nodes).view(-1) - return output_lst - - -@instrument_w_nvtx -@torch.no_grad() -def all_to_all_loco_quant_reduce( - params: List[Tensor], - groups: {}, - loco_param: Any = None, -) -> List[Tensor]: - global quantizer_module - global loco_idx - if quantizer_module is None: - quantizer_module = op_builder.QuantizerBuilder().load() - local_world_size = get_accelerator().device_count() - global_world_size = dist.get_world_size() - num_nodes = global_world_size // local_world_size - this_rank = dist.get_rank() - intra_idx = int(this_rank / local_world_size) - inter_idx = this_rank % local_world_size - output_lst: List[Tensor] = [None] * len(params) - for idx, p in enumerate(params): - tensor = p.grad - if tensor.dim() == 1: - output_lst[idx] = reduce_scatter_coalesced([tensor])[0] - elif tensor.numel() % (2 * global_world_size) != 0: - # Due to the constraint of 2-stage all-to-all, the input tensor must be divisible by 2 * global_world_size - # Otherwise, all-to-all cannot be performed because of shape mismatch. - # See more at https://github.com/deepspeedai/DeepSpeed/pull/5056 - logger.warning( - f"qgZ falls back to reduce_scatter because tensor size = {tensor.numel()} is not divisible by (2 * global_world_size) = {2 * global_world_size}. Please consider allocating a new world to enable qgZ" - ) - output_lst[idx] = reduce_scatter_coalesced([tensor])[0] - else: - err_beta = loco_param['err_beta'] - reset_T = loco_param['reset_T'] - if not hasattr(p, 'intra_ef_buf') or loco_idx > reset_T: - loco_idx = 0 - intra_err = torch.zeros_like(p.grad) - inter_err = torch.zeros(tensor.numel() // local_world_size, device=tensor.device, dtype=tensor.dtype) - else: - intra_err = quantizer_module.dequantize(p.intra_ef_buf[0], p.intra_ef_buf[1], - p.intra_ef_buf[1].numel(), 8, quantizer_module.Symmetric) - inter_err = quantizer_module.dequantize(p.inter_ef_buf[0], p.inter_ef_buf[1], - p.inter_ef_buf[1].numel(), 8, quantizer_module.Symmetric) - - intra_quant_group = max(tensor.shape[0], tensor.shape[1], global_world_size) - inter_quant_group = intra_quant_group // local_world_size - intra_quant_int4, intra_q_scales = quantizer_module.loco_swizzle_quant(tensor, intra_err, err_beta, - intra_quant_group, 4, - quantizer_module.Symmetric, 1, - num_nodes, local_world_size) - local_output = torch.empty_like(intra_quant_int4) - scale_output = torch.empty_like(intra_q_scales) - all_to_all_single(local_output, intra_quant_int4, group=groups[f'local_{intra_idx}']) - all_to_all_single(scale_output, intra_q_scales, group=groups[f'local_{intra_idx}']) - - p.intra_ef_buf = quantizer_module.quantize(intra_err, intra_quant_group, 8, quantizer_module.Symmetric) - - global_input_tensor, global_scales = quantizer_module.loco_quantized_reduction( - local_output, scale_output, inter_err, err_beta, intra_quant_group, inter_quant_group, 4, - quantizer_module.Symmetric, local_world_size) - - global_output = torch.empty_like(global_input_tensor) - global_scale_output = torch.empty_like(global_scales) - all_to_all_single(global_output, global_input_tensor, group=groups[f'global_{inter_idx}']) - all_to_all_single(global_scale_output, global_scales, group=groups[f'global_{inter_idx}']) - - p.inter_ef_buf = quantizer_module.quantize(inter_err, inter_quant_group, 8, quantizer_module.Symmetric) - - final_output = quantizer_module.dequantize(global_output, global_scale_output, global_scale_output.numel(), - 4, quantizer_module.Symmetric) - assert final_output.numel( - ) % num_nodes == 0, f"final_output.numel()={final_output.numel()} is not divisible by num_nodes={num_nodes}" - output_lst[idx] = (sum(list(final_output.chunk(num_nodes))) / num_nodes).view(-1) - loco_idx += 1 - - return output_lst - - -@instrument_w_nvtx -@torch.no_grad() -def reduce_scatter_coalesced( - tensors: List[Tensor], - group: ProcessGroup = None, -) -> List[Tensor]: - """simultaneously reduce-scatter a list of tensors - this can be done more - efficiently than individual reduce scatter calls - TODO. see if PyTorch team wants a c++ version of this for ProcessGroupNCCL - """ - this_rank = dist.get_rank(group) - world_sz = dist.get_world_size(group) - - partition_lst_for_each_tensor = [None] * len(tensors) - for tensor_idx, tensor in enumerate(tensors): - flattened_tensor = tensor.view(-1) - chunk_sz = math.ceil(tensor.numel() / world_sz) - partition_lst_for_each_tensor[tensor_idx] = [ - flattened_tensor[rank * chunk_sz:rank * chunk_sz + chunk_sz] for rank in range(0, world_sz) - ] - - padded_partition_sz_for_each_tensor = tuple(math.ceil(t.numel() / world_sz) for t in tensors) - - if len(tensors) == 1 and tensors[0].numel() % world_sz == 0: - # if there's only one tensor being reduced and we don't need to pad - # we have an opportunity to avoid a memory allocation - tensor_partition_flat_buffer = tensors[0].view(-1) - else: - # interleave tensor partitions such that the correct reduced partitions of each tensor - # end up at each rank - tensor_partitions_lst_with_padding = [] - for rank in range(world_sz): - for tensor_idx in range(len(tensors)): - # add tensor content - tensor_chunk = partition_lst_for_each_tensor[tensor_idx][rank] - tensor_partitions_lst_with_padding.append(tensor_chunk) - - # add padding if necessary - padding_sz = padded_partition_sz_for_each_tensor[tensor_idx] - tensor_chunk.numel() - if padding_sz > 0: - tensor_partitions_lst_with_padding.append( - torch.empty(padding_sz, dtype=tensor_chunk.dtype, device=tensor_chunk.device)) - - tensor_partition_flat_buffer = instrument_w_nvtx(torch.cat)(tensor_partitions_lst_with_padding) - - tensor_partition_flat_buffer.div_(world_sz) # pre-divide - tensor_partition_buffer_for_each_rank: List[Tensor] = torch.chunk(tensor_partition_flat_buffer, world_sz) - - # batched reduce-scatter call - _torch_reduce_scatter_fn(tensor_partition_flat_buffer, - tensor_partition_buffer_for_each_rank[this_rank], - group=group) - - # reverse procedure of the interleaving done previously, done on the - # result of the batched reduce-scatter - output_lst: List[Tensor] = [None] * len(tensors) - offset = 0 - for tensor_idx in range(len(tensors)): - output_lst[tensor_idx] = tensor_partition_buffer_for_each_rank[this_rank].narrow( - 0, offset, partition_lst_for_each_tensor[tensor_idx][this_rank].numel()) - - offset += padded_partition_sz_for_each_tensor[tensor_idx] - return output_lst diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/compressed.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/compressed.py deleted file mode 100644 index 2c5482eb1ad769bb3a0ecd87b4a3489d31ce28db..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/compressed.py +++ /dev/null @@ -1,137 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import numpy as np -import torch -import deepspeed.comm as dist -from deepspeed.accelerator import get_accelerator -from deepspeed.ops.op_builder import PackbitsBuilder - - -class CompressedBackend(object): - - def __init__(self, mpu=None): - if mpu is None: - self.world_group = dist.new_group(ranks=range(dist.get_world_size())) - else: - self.mpu = mpu - self.world_group = self.mpu.get_data_parallel_group() - self.size = dist.get_world_size(group=self.world_group) - self.rank = dist.get_rank(group=self.world_group) - self.packer = PackbitsBuilder().load() - - def my_igather(self, rank, size, group, sendbuf, recvbuf, root): - req = [] - if rank == root: - for idx in range(size): - if idx != rank: - req.append(dist.irecv(recvbuf[idx], src=idx, group=group)) - else: - recvbuf[rank] = sendbuf - else: - req.append(dist.isend(sendbuf, group=group, dst=root)) - return req - - def my_gather(self, rank, size, group, sendbuf, recvbuf, root): - if rank == root: - for idx in range(size): - if idx != rank: - dist.recv(recvbuf[idx], src=idx, group=group) - else: - recvbuf[rank] = sendbuf - else: - dist.send(sendbuf, group=group, dst=root) - - def pack(self, buffer, size): - # pack float tensor into uint8 tensor - packed = self.packer.packbits(buffer.float(), buffer.numel(), self.rank) - return packed.reshape(size, -1) - - def unpack(self, buffer, size, dtype): - # unpack uint8 to float tensor - unpacked = self.packer.unpackbits(buffer, buffer.numel(), self.rank) - return unpacked.reshape(size, -1).to(dtype) - - def compressed_allreduce(self, buffer_m: torch.tensor, worker_error, server_error, local_rank): - original_shape = buffer_m.size() - if len(original_shape) > 1: - buffer_m = torch.flatten(buffer_m) - - # align size of original_buffer and error - original_size = buffer_m.numel() - worker_error_size = worker_error.numel() - if original_size != worker_error_size: - empty_tensor = torch.zeros(worker_error_size - original_size, device=buffer_m.device) - buffer_m = torch.cat([buffer_m, empty_tensor]) - - buffer_m.add_(worker_error) - worker_scale = torch.linalg.norm(buffer_m) / np.sqrt(torch.numel(buffer_m)) - - worker_error.set_(buffer_m - worker_scale * buffer_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) - - sign_list_packed_tmp = self.pack(buffer_m, self.size).type(torch.int8) - - recvbuf_sign = torch.zeros([self.size, len(sign_list_packed_tmp[self.rank])], - dtype=sign_list_packed_tmp[0].dtype, - device=sign_list_packed_tmp.device) - - sign_list_packed = [sign_list_packed_tmp[idx] for idx in range(self.size)] - - recvbuf_scale = [ - torch.zeros(1, dtype=worker_scale.dtype, device=get_accelerator().current_device_name()) - for _ in range(self.size) - ] - - # communication phase 1 - # all to all for sign - dist.all_to_all_single(recvbuf_sign, torch.stack(sign_list_packed), group=self.world_group) - # all gather for scale - dist.all_gather(recvbuf_scale, worker_scale, group=self.world_group) - - flattened_recvbuf_sign = recvbuf_sign.type(torch.uint8).flatten() - compensated_server_m = self.unpack(flattened_recvbuf_sign, self.size, torch.float32) \ - .mul_(torch.stack(recvbuf_scale).mul_(1 / self.size)).sum(0) - - compensated_server_m.add_(server_error) - - server_scale = torch.linalg.norm(compensated_server_m) / np.sqrt(compensated_server_m.numel()) - - server_error.set_(compensated_server_m - - server_scale * compensated_server_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) - - server_sign_packed = self.pack(compensated_server_m, 1).type(torch.int8) - - # recvbuf_sign_server - recvbuf_sign_server_tmp = torch.zeros([self.size, len(server_sign_packed[0])], - dtype=recvbuf_sign.dtype, - device=server_sign_packed.device) - - recvbuf_sign_server = [recvbuf_sign_server_tmp[idx] for idx in range(self.size)] - - # recvbuf_scale_server - recvbuf_scale_server_tmp = torch.zeros([self.size, 1], - dtype=worker_scale.dtype, - device=server_sign_packed.device) - - recvbuf_scale_server = [recvbuf_scale_server_tmp[idx] for idx in range(self.size)] - - # communication Phase 2 - dist.all_gather(recvbuf_sign_server, server_sign_packed[0], group=self.world_group) - dist.all_gather(recvbuf_scale_server, server_scale, group=self.world_group) - - recvbuf_sign_server = torch.stack(recvbuf_sign_server) - - flattened_recvbuf_sign_server = recvbuf_sign_server.type(torch.uint8).flatten() - - buffer_m.data.copy_( - self.unpack(flattened_recvbuf_sign_server, self.size, - torch.float32).mul_(recvbuf_scale_server_tmp).flatten().data) - - if original_size != worker_error_size: - buffer_m = buffer_m[0:original_size] - if len(original_shape) > 1: - buffer_m = buffer_m.reshape(original_shape) - - return buffer_m diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/hccl.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/hccl.py deleted file mode 100644 index b8639c7da4c99327e4a6afe4e226d6c4a224886b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/hccl.py +++ /dev/null @@ -1,124 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import numpy as np -import torch -import torch_npu -import deepspeed.comm as dist - - -class HcclBackend(object): - - def __init__(self, mpu=None): - if mpu is None: - self.world_group = dist.new_group(ranks=range(dist.get_world_size())) - else: - self.mpu = mpu - self.world_group = self.mpu.get_data_parallel_group() - self.size = dist.get_world_size(group=self.world_group) - self.rank = dist.get_rank(group=self.world_group) - - def my_igather(self, rank, size, group, sendbuf, recvbuf, root): - req = [] - if rank == root: - for idx in range(size): - if idx != rank: - req.append(dist.irecv(recvbuf[idx], src=idx, group=group)) - else: - recvbuf[rank] = sendbuf - else: - req.append(dist.isend(sendbuf, group=group, dst=root)) - return req - - def my_gather(self, rank, size, group, sendbuf, recvbuf, root): - if rank == root: - for idx in range(size): - if idx != rank: - dist.recv(recvbuf[idx], src=idx, group=group) - else: - recvbuf[rank] = sendbuf - else: - dist.send(sendbuf, group=group, dst=root) - - def compressed_allreduce(self, buffer_m: torch.tensor, worker_error, server_error, local_rank): - original_shape = buffer_m.size() - if len(original_shape) > 1: - buffer_m = torch.flatten(buffer_m) - - # align size of original_buffer and error - original_size = buffer_m.numel() - worker_error_size = worker_error.numel() - if original_size != worker_error_size: - empty_tensor = torch.zeros(worker_error_size - original_size, device=buffer_m.device) - buffer_m = torch.cat([buffer_m, empty_tensor]) - - buffer_m.add_(worker_error) - worker_scale = torch.linalg.norm(buffer_m) / np.sqrt(torch.numel(buffer_m)) - - worker_error.set_(buffer_m - worker_scale * buffer_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) - - sign_list_packed_tmp = torch_npu.npu_sign_bits_pack(buffer_m, self.size).type(torch.int8) - - recvbuf_sign = torch.zeros([self.size, len(sign_list_packed_tmp[self.rank])], - dtype=sign_list_packed_tmp[0].dtype, - device=sign_list_packed_tmp.device) - - sign_list_packed = [sign_list_packed_tmp[idx] for idx in range(self.size)] - - recvbuf_scale = [ - torch.zeros(1, dtype=worker_scale.dtype, device=torch.device(local_rank)) for _ in range(self.size) - ] - - # communication phase 1 - # all to all for sign - dist.all_to_all_single(recvbuf_sign, torch.stack(sign_list_packed), group=self.world_group) - # all gather for scale - dist.all_gather(recvbuf_scale, worker_scale, group=self.world_group) - - flattened_recvbuf_sign = recvbuf_sign.type(torch.uint8).flatten() - compensated_server_m = torch_npu.npu_sign_bits_unpack(flattened_recvbuf_sign, self.size, torch.float32) \ - .mul_(torch.stack(recvbuf_scale).mul_(1 / self.size)).sum(0) - - compensated_server_m.add_(server_error) - - server_scale = torch.linalg.norm(compensated_server_m) / np.sqrt(compensated_server_m.numel()) - - server_error.set_(compensated_server_m - - server_scale * compensated_server_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) - - server_sign_packed = torch_npu.npu_sign_bits_pack(compensated_server_m, 1).type(torch.int8) - - # recvbuf_sign_server - recvbuf_sign_server_tmp = torch.zeros([self.size, len(server_sign_packed[0])], - dtype=recvbuf_sign.dtype, - device=server_sign_packed.device) - - recvbuf_sign_server = [recvbuf_sign_server_tmp[idx] for idx in range(self.size)] - - # recvbuf_scale_server - recvbuf_scale_server_tmp = torch.zeros([self.size, 1], - dtype=worker_scale.dtype, - device=server_sign_packed.device) - - recvbuf_scale_server = [recvbuf_scale_server_tmp[idx] for idx in range(self.size)] - - # communication Phase 2 - dist.all_gather(recvbuf_sign_server, server_sign_packed[0], group=self.world_group) - dist.all_gather(recvbuf_scale_server, server_scale, group=self.world_group) - - recvbuf_sign_server = torch.stack(recvbuf_sign_server) - - flattened_recvbuf_sign_server = recvbuf_sign_server.type(torch.uint8).flatten() - - buffer_m.data.copy_( - torch_npu.npu_sign_bits_unpack(flattened_recvbuf_sign_server, self.size, - torch.float32).mul_(recvbuf_scale_server_tmp).flatten().data) - - if original_size != worker_error_size: - buffer_m = buffer_m[0:original_size] - if len(original_shape) > 1: - buffer_m = buffer_m.reshape(original_shape) - - return buffer_m diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/mpi.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/mpi.py deleted file mode 100644 index bc544787aa7a7de5181d83897e725991f572e5ce..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/mpi.py +++ /dev/null @@ -1,215 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import cupy -import time -import numpy as np -from mpi4py import MPI - -from deepspeed.runtime.compression.cupy import CupyBackend - - -class MpiBackend(object): - - def __init__(self, cuda_aware): - self.comm = MPI.COMM_WORLD - self.rank = self.comm.Get_rank() - self.size = self.comm.Get_size() - self.cuda_aware = cuda_aware - self.compression_backend = CupyBackend() - - def my_igather(self, rank, size, comm, sendbuf, recbuf, root): - req = [] - if rank == root: - for idx in range(size): - if idx != rank: - req.append(comm.Irecv(recbuf[idx], source=idx)) - else: - recbuf[rank] = sendbuf - else: - req.append(comm.Isend(sendbuf, dest=root)) - return req - - def gather_cuda(self, rank, world_size, comm, cupy_sign_list_packed, cupy_recvbuf_sign, cupy_worker_scale, - cupy_recvbuf_scale): - # We do in-place operations on cupy buffers so we do not return any buffers - requests = [] - for idx in range(world_size): - req_sign = self.my_igather(rank, world_size, comm, cupy_sign_list_packed[idx], cupy_recvbuf_sign, root=idx) - requests += req_sign - - for idx in range(world_size): - req_scale = self.my_igather(rank, world_size, comm, cupy_worker_scale, cupy_recvbuf_scale, root=idx) - requests += req_scale - - MPI.Request.Waitall(requests) - - def gather_host(self, rank, world_size, comm, cupy_sign_list_packed, cupy_recvbuf_sign, cupy_worker_scale, - cupy_recvbuf_scale): - - # In-place operations are not possible for newly created cupy arrays - # so we need to return the new buffers - numpy_recvbuf_sign = np.zeros([world_size, cupy_sign_list_packed[rank].size], - dtype=cupy_sign_list_packed[0].dtype) - numpy_recvbuf_scale = np.zeros([world_size, 1], dtype=cupy_worker_scale.dtype) - - # 1. convert from cupy to numpy - numpy_sign_list_packed = cupy_sign_list_packed - - for idx in range(world_size): - numpy_sign_list_packed[idx] = cupy.asnumpy(cupy_sign_list_packed[idx]) - - numpy_worker_scale = cupy.asnumpy(cupy_worker_scale) - numpy_recvbuf_scale = cupy.asnumpy(cupy_recvbuf_scale) - - cupy.cuda.get_current_stream().synchronize() - - # 2. use numpy buffers for communication - requests = [] - - for idx in range(world_size): - req_sign = self.my_igather(rank, - world_size, - comm, - numpy_sign_list_packed[idx], - numpy_recvbuf_sign, - root=idx) - requests += req_sign - - for idx in range(world_size): - req_scale = self.my_igather(rank, world_size, comm, numpy_worker_scale, numpy_recvbuf_scale, root=idx) - requests += req_scale - - MPI.Request.Waitall(requests) - - # 3. Convert back from numpy to cupy - cupy_recvbuf_sign = cupy.asarray(numpy_recvbuf_sign) - for idx in range(world_size): - cupy_sign_list_packed[idx] = cupy.asarray(numpy_sign_list_packed[idx]) - - cupy_worker_scale = cupy.asarray(numpy_worker_scale) - cupy_recvbuf_scale = cupy.asarray(numpy_recvbuf_scale) - cupy.cuda.get_current_stream().synchronize() - - return cupy_sign_list_packed, cupy_recvbuf_sign, cupy_worker_scale, cupy_recvbuf_scale - - def allgather_cuda(self, comm, cupy_server_sign_packed, cupy_recvbuf_sign_server, cupy_server_scale, - cupy_recvbuf_scale_server): - comm.Allgather(cupy_server_sign_packed, cupy_recvbuf_sign_server) - comm.Allgather(cupy_server_scale, cupy_recvbuf_scale_server) - - def allgather_host(self, comm, cupy_server_sign_packed, cupy_recvbuf_sign_server, cupy_server_scale, - cupy_recvbuf_scale_server): - - # 1. Convert cupy to numpy - numpy_recvbuf_sign_server = np.zeros([comm.Get_size(), cupy_server_sign_packed.size], - dtype=cupy_server_sign_packed.dtype) - numpy_recvbuf_scale_server = np.zeros([comm.Get_size(), 1], dtype=cupy_server_scale.dtype) - - numpy_server_sign_packed = cupy.asnumpy(cupy_server_sign_packed) - numpy_recvbuf_sign_server = cupy.asnumpy(cupy_recvbuf_sign_server) - numpy_server_scale = cupy.asnumpy(cupy_server_scale) - numpy_recvbuf_scale_server = cupy.asnumpy(cupy_recvbuf_scale_server) - cupy.cuda.get_current_stream().synchronize() - - # 2. Communicate numpy buffers - comm.Allgather(numpy_server_sign_packed, numpy_recvbuf_sign_server) - comm.Allgather(numpy_server_scale, numpy_recvbuf_scale_server) - comm.Barrier() - - # 3. Convert numpy back to cupy - cupy_server_sign_packed = cupy.asarray(numpy_server_sign_packed) - cupy_recvbuf_sign_server = cupy.asarray(numpy_recvbuf_sign_server) - cupy_server_scale = cupy.asarray(numpy_server_scale) - cupy_recvbuf_scale_server = cupy.asarray(numpy_recvbuf_scale_server) - cupy.cuda.get_current_stream().synchronize() - - return cupy_server_sign_packed, cupy_recvbuf_sign_server, cupy_server_scale, cupy_recvbuf_scale_server - - def compressed_allreduce(self, buffer_m: torch.tensor, worker_error, server_error, local_rank): - - all_start_time = time.time() - original_shape = buffer_m.size() - if len(original_shape) > 1: - buffer_m = torch.flatten(buffer_m) - original_size = buffer_m.numel() - worker_error_size = worker_error.numel() - cupy.cuda.Device(local_rank).use() - - if original_size != worker_error_size: - empty_tensor = torch.zeros(worker_error_size - original_size, device=buffer_m.device) - buffer_m = torch.cat([buffer_m, empty_tensor]) - - buffer_m.add_(worker_error) - worker_scale = torch.linalg.norm(buffer_m) / np.sqrt(torch.numel(buffer_m)) - worker_error.set_(buffer_m - worker_scale * buffer_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) - - cupy_sign_list_packed = self.compression_backend.compress_by_chunk( - self.compression_backend.torch2cupy(buffer_m.sign_().add_(1).bool()), self.size) - cupy_worker_scale = self.compression_backend.torch2cupy(worker_scale) - - cupy_recvbuf_sign = cupy.zeros([self.size, cupy_sign_list_packed[self.rank].size], - dtype=cupy_sign_list_packed[0].dtype) - cupy_recvbuf_scale = cupy.zeros([self.size, 1], dtype=cupy_worker_scale.dtype) - - # Communication Phase 1 - gather_start = time.time() - if self.cuda_aware: - self.gather_cuda(self.rank, self.size, self.comm, cupy_sign_list_packed, cupy_recvbuf_sign, - cupy_worker_scale, cupy_recvbuf_scale) - else: - _, cupy_recvbuf_sign, _, cupy_recvbuf_scale = self.gather_host(self.rank, self.size, self.comm, - cupy_sign_list_packed, cupy_recvbuf_sign, - cupy_worker_scale, cupy_recvbuf_scale) - gather_end = time.time() - - # cupy_sign_list_packed, cupy_worker_scale, worker_scale = None, None, None - cupy_sign_list_packed = None - - compensated_server_m = self.compression_backend.cupy2torch( - (cupy.unpackbits(cupy_recvbuf_sign.flatten())).reshape(self.size, -1)).float().add_(-0.5).mul_(2.0).mul_( - self.compression_backend.cupy2torch(cupy_recvbuf_scale).mul_(1 / self.size)).sum(0) - compensated_server_m.add_(server_error) - server_scale = torch.linalg.norm(compensated_server_m) / np.sqrt(compensated_server_m.numel()) - server_error.set_(compensated_server_m - - server_scale * compensated_server_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) - - cupy_server_scale = self.compression_backend.torch2cupy(server_scale) - - cupy_server_sign_packed = self.compression_backend.compress_by_chunk( - self.compression_backend.torch2cupy(compensated_server_m.sign_().add_(1).bool()), 1) - compensated_server_m = None - - cupy_recvbuf_sign_server = cupy.zeros([self.size, cupy_server_sign_packed[0].size], - dtype=cupy_recvbuf_sign.dtype) - cupy_recvbuf_scale_server = cupy.zeros([self.size, 1], dtype=cupy_recvbuf_scale.dtype) - # cupy_recvbuf_sign, cupy_recvbuf_scale = None, None - cupy_recvbuf_sign = None - - # Communication Phase 2 - if self.cuda_aware: - self.allgather_cuda(self.comm, cupy_server_sign_packed[0], cupy_recvbuf_sign_server, cupy_server_scale, - cupy_recvbuf_scale_server) - else: - _, cupy_recvbuf_sign_server, _, cupy_recvbuf_scale_server = self.allgather_host( - self.comm, cupy_server_sign_packed[0], cupy_recvbuf_sign_server, cupy_server_scale, - cupy_recvbuf_scale_server) - - # cupy_server_sign_packed, cupy_server_scale, server_scale = None, None, None - cupy_server_sign_packed = None - - buffer_m.data.copy_( - self.compression_backend.cupy2torch((cupy.unpackbits(cupy_recvbuf_sign_server.flatten())).reshape( - self.size, -1)).float().add_(-0.5).mul_(2.0).mul_( - self.compression_backend.cupy2torch(cupy_recvbuf_scale_server)).flatten().data) - if original_size != worker_error_size: - buffer_m = buffer_m[0:original_size] - if len(original_shape) > 1: - buffer_m = buffer_m.reshape(original_shape) - - # cupy_recvbuf_sign_server, cupy_recvbuf_scale_server = None, None - - return buffer_m diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/nccl.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/nccl.py deleted file mode 100644 index a57b7519a295e584e77561843c52fa025bfaa66d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/comm/nccl.py +++ /dev/null @@ -1,166 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed import comm as dist -import cupy -import numpy as np - -from deepspeed.runtime.compression.cupy import CupyBackend -from deepspeed.utils.torch import required_torch_version -from deepspeed.accelerator import get_accelerator - - -class NcclBackend(object): - - def __init__(self, mpu=None): - if mpu is None: - self.world_group = dist.new_group(ranks=range(dist.get_world_size())) - else: - self.mpu = mpu - self.world_group = self.mpu.get_data_parallel_group() - self.rank = dist.get_rank(group=self.world_group) - self.size = dist.get_world_size(group=self.world_group) - self.compression_backend = CupyBackend() - self.bool_not_supported = required_torch_version(min_version=1.10) - - def my_igather(self, rank, size, group, sendbuf, recvbuf, root): - req = [] - if rank == root: - for idx in range(size): - if idx != rank: - req.append(dist.irecv(recvbuf[idx], src=idx, group=group)) - else: - recvbuf[rank] = sendbuf - else: - req.append(dist.isend(sendbuf, group=group, dst=root)) - return req - - def my_gather(self, rank, size, group, sendbuf, recvbuf, root): - if rank == root: - for idx in range(size): - if idx != rank: - dist.recv(recvbuf[idx], src=idx, group=group) - else: - recvbuf[rank] = sendbuf - else: - dist.send(sendbuf, group=group, dst=root) - - def compressed_allreduce(self, buffer_m: torch.tensor, worker_error, server_error, local_rank): - - # all_start_time = time.time() - original_shape = buffer_m.size() - if len(original_shape) > 1: - buffer_m = torch.flatten(buffer_m) - original_size = buffer_m.numel() - worker_error_size = worker_error.numel() - cupy.cuda.Device(local_rank).use() - - if original_size != worker_error_size: - empty_tensor = torch.zeros(worker_error_size - original_size, device=buffer_m.device) - buffer_m = torch.cat([buffer_m, empty_tensor]) - - buffer_m.add_(worker_error) - worker_scale = torch.linalg.norm(buffer_m) / np.sqrt(buffer_m.numel()) - worker_error.set_(buffer_m - worker_scale * buffer_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) - - if self.bool_not_supported: - cupy_sign_list_packed = self.compression_backend.compress_by_chunk( - self.compression_backend.torch2cupy(buffer_m.sign_().add_(1).bool().to(dtype=torch.uint8)), self.size) - else: - cupy_sign_list_packed = self.compression_backend.compress_by_chunk( - self.compression_backend.torch2cupy(buffer_m.sign_().add_(1).bool()), self.size) - cupy_worker_scale = self.compression_backend.torch2cupy(worker_scale) - - cupy_recvbuf_sign = cupy.zeros([self.size, cupy_sign_list_packed[self.rank].size], - dtype=cupy_sign_list_packed[0].dtype) - # cupy_recvbuf_scale = cupy.zeros([self.size, 1], dtype=cupy_worker_scale.dtype) - - sign_list_packed = [ - self.compression_backend.cupy2torch(cupy_sign_list_packed[idx]) for idx in range(self.size) - ] - - # worker_scale = self.compression_backend.cupy2torch(cupy_worker_scale) - recvbuf_sign = self.compression_backend.cupy2torch(cupy_recvbuf_sign) - #recvbuf_scale = self.compression_backend.cupy2torch(cupy_recvbuf_scale) - recvbuf_scale = [ - torch.zeros(1, dtype=worker_scale.dtype, device=torch.device(get_accelerator().device_name(local_rank))) - for i in range(self.size) - ] - - # communication phase 1 - # gather_start = time.time() - # Alltoall for sign - dist.all_to_all_single(recvbuf_sign, torch.stack(sign_list_packed), group=self.world_group) - # Allgather for scale - dist.all_gather(recvbuf_scale, worker_scale, group=self.world_group) - - # gather_end = time.time() - - # cupy_sign_list_packed, sign_list_packed, cupy_worker_scale, worker_scale = None, None, None, None - cupy_sign_list_packed = None - - cupy_recvbuf_sign = self.compression_backend.torch2cupy(recvbuf_sign) - #cupy_recvbuf_scale = self.compression_backend.torch2cupy(torch.stack(recvbuf_scale)) - - compensated_server_m = self.compression_backend.cupy2torch( - (cupy.unpackbits(cupy_recvbuf_sign.flatten())).reshape(self.size, -1)).float().add_(-0.5).mul_(2.0).mul_( - torch.stack(recvbuf_scale).mul_(1 / self.size)).sum(0) - compensated_server_m.add_(server_error) - server_scale = torch.linalg.norm(compensated_server_m) / np.sqrt(compensated_server_m.numel()) - server_error.set_(compensated_server_m - - server_scale * compensated_server_m.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) - - # cupy_server_scale = self.compression_backend.torch2cupy(server_scale) - - if self.bool_not_supported: - cupy_server_sign_packed = self.compression_backend.compress_by_chunk( - self.compression_backend.torch2cupy(compensated_server_m.sign_().add_(1).bool().to(dtype=torch.uint8)), - 1) - else: - cupy_server_sign_packed = self.compression_backend.compress_by_chunk( - self.compression_backend.torch2cupy(compensated_server_m.sign_().add_(1).bool()), 1) - compensated_server_m = None - - cupy_recvbuf_sign_server = cupy.zeros([self.size, cupy_server_sign_packed[0].size], - dtype=cupy_recvbuf_sign.dtype) - # cupy_recvbuf_sign, recvbuf_sign = None, None - cupy_recvbuf_sign = None - - server_sign_packed = [self.compression_backend.cupy2torch(cupy_server_sign_packed[0])] - recvbuf_sign_server = [ - self.compression_backend.cupy2torch(cupy_recvbuf_sign_server[idx]) for idx in range(self.size) - ] - - # server_scale = self.compression_backend.cupy2torch(cupy_server_scale) - cupy_recvbuf_scale_server = cupy.zeros([self.size, 1], dtype=cupy_worker_scale.dtype) - # cupy_recvbuf_scale, recvbuf_scale = None, None - - recvbuf_scale_server = [ - self.compression_backend.cupy2torch(cupy_recvbuf_scale_server[idx]) for idx in range(self.size) - ] - - # Communication Phase 2 - dist.all_gather(recvbuf_sign_server, server_sign_packed[0], group=self.world_group) - dist.all_gather(recvbuf_scale_server, server_scale, group=self.world_group) - - cupy_server_sign_packed = None - - # need to convert from a tensor list to a single tensor - # dist.all_gather only provides a tensor list as the recv/output buffer - recvbuf_sign_server = torch.stack(recvbuf_sign_server) - - cupy_recvbuf_sign_server = self.compression_backend.torch2cupy(recvbuf_sign_server) - - buffer_m.data.copy_( - self.compression_backend.cupy2torch((cupy.unpackbits(cupy_recvbuf_sign_server.flatten())).reshape( - self.size, -1)).float().add_(-0.5).mul_(2.0).mul_( - self.compression_backend.cupy2torch(cupy_recvbuf_scale_server)).flatten().data) - if original_size != worker_error_size: - buffer_m = buffer_m[0:original_size] - if len(original_shape) > 1: - buffer_m = buffer_m.reshape(original_shape) - - return buffer_m diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/compiler.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/compiler.py deleted file mode 100644 index be778b83f8bb8f3b9416fdf1d5717532367a8e4b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/compiler.py +++ /dev/null @@ -1,30 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed.utils.torch import required_torch_version - -try: - from torch.compiler import is_compiling as torch_is_compiling -except ImportError: - try: - from torch._dynamo.external_utils import is_compiling as torch_is_compiling - except ImportError: - # Torch does not have compiler support - torch_is_compiling = lambda: False - - -def is_compile_supported(): - return required_torch_version(min_version=2.1) - - -def disable(func): - if is_compile_supported(): - return torch.compiler.disable(func) - return func - - -def is_compiling(): - return torch_is_compiling() diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 5b8d439bfbf4ad2903d701886ec163f5b55a05f7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__pycache__/cupy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__pycache__/cupy.cpython-310.pyc deleted file mode 100644 index b743bf3e021e3d27f8a70a8b82bf7d726917e94d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/__pycache__/cupy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/cupy.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/cupy.py deleted file mode 100644 index 7133ac04ed2b65bd496656870aebbf95ac04d3ad..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/compression/cupy.py +++ /dev/null @@ -1,26 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import cupy -from torch.utils.dlpack import to_dlpack -from torch.utils.dlpack import from_dlpack - - -class CupyBackend(object): - - def __init__(self): - pass - - def torch2cupy(self, tensor): - return cupy.from_dlpack(to_dlpack(tensor)) - - def cupy2torch(self, cupy_tensor): - return from_dlpack(cupy_tensor) - - def compress_by_chunk(self, cupy_bool_tensor, num_chunks): - packed_sign = cupy.packbits(cupy_bool_tensor) - sign_list_packed = cupy.split(packed_sign, num_chunks) - cupy.cuda.get_current_stream().synchronize() - return sign_list_packed diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/config.py deleted file mode 100644 index 6a18ff193aec1e4d47c1feb1b015e94693547a3e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/config.py +++ /dev/null @@ -1,997 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -from typing import Union -from enum import Enum - -import torch -import json -import hjson -import copy -import base64 - -from .constants import * -from .config_utils import ( - get_scalar_param, - dict_raise_error_on_duplicate_keys, - ScientificNotationEncoder, -) -from .zero.config import get_zero_config, ZeroStageEnum -from .activation_checkpointing.config import DeepSpeedActivationCheckpointingConfig -from ..comm.config import DeepSpeedCommsConfig -from ..monitor.config import get_monitor_config -from ..inference.config import WeightQuantConfig -from .precision_config import get_bfloat16_config, get_float16_config -from ..compile.config import CompileConfig - -from deepspeed import comm as dist -from deepspeed.runtime.config_utils import DeepSpeedConfigModel - -from ..git_version_info import version as __version__ -from ..utils import logger - -from ..elasticity import ( - elasticity_enabled, - compute_elastic_config, - ensure_immutable_elastic_config, -) -from ..elasticity.config import ElasticityConfigError -from ..elasticity.constants import ( - ELASTICITY, - IGNORE_NON_ELASTIC_BATCH_INFO, - IGNORE_NON_ELASTIC_BATCH_INFO_DEFAULT, - MODEL_PARALLEL_SIZE, - MODEL_PARALLEL_SIZE_DEFAULT, - NUM_GPUS_PER_NODE, - NUM_GPUS_PER_NODE_DEFAULT, -) - -from ..profiling.config import DeepSpeedFlopsProfilerConfig -from ..autotuning.config import DeepSpeedAutotuningConfig -from ..nebula.config import DeepSpeedNebulaConfig - -from ..compression.config import get_compression_config, get_quantize_enabled -from ..compression.constants import * -from .swap_tensor.aio_config import get_aio_config -from .model_checkpointing.config import get_checkpoint_config - -from .tensor_parallel import get_tensor_parallel_config -from .data_pipeline.config import get_data_efficiency_enabled, get_data_efficiency_config, get_curriculum_enabled_legacy, get_curriculum_params_legacy -from .data_pipeline.constants import * - -from ..utils.config import get_timers_config - -TENSOR_CORE_ALIGN_SIZE = 8 - -ADAGRAD_OPTIMIZER = 'adagrad' -ADAM_OPTIMIZER = 'adam' -ADAMW_OPTIMIZER = 'adamw' -LAMB_OPTIMIZER = 'lamb' -ONEBIT_ADAM_OPTIMIZER = 'onebitadam' -ZERO_ONE_ADAM_OPTIMIZER = 'zerooneadam' -ONEBIT_LAMB_OPTIMIZER = 'onebitlamb' -MUADAM_OPTIMIZER = 'muadam' -MUADAMW_OPTIMIZER = 'muadamw' -MUSGD_OPTIMIZER = 'musgd' -LION_OPTIMIZER = 'lion' -DEEPSPEED_OPTIMIZERS = [ - ADAGRAD_OPTIMIZER, ADAM_OPTIMIZER, ADAMW_OPTIMIZER, LAMB_OPTIMIZER, ONEBIT_ADAM_OPTIMIZER, ONEBIT_LAMB_OPTIMIZER, - ZERO_ONE_ADAM_OPTIMIZER, MUADAM_OPTIMIZER, MUADAMW_OPTIMIZER, MUSGD_OPTIMIZER, LION_OPTIMIZER -] - -# extra optimizer parameters for adam/adamw -TORCH_ADAM_PARAM = "torch_adam" - -# default to adamw logic for adam/adamw optimizers unless user explicitly opts out -ADAM_W_MODE = "adam_w_mode" -ADAM_W_MODE_DEFAULT = True - - -class DeepSpeedConfigError(Exception): - pass - - -class DtypeEnum(Enum): - # The torch dtype must always be the first value (so we return torch.dtype) - fp16 = torch.float16, "torch.float16", "fp16", "float16", "half" - fp32 = torch.float32, "torch.float32", "fp32", "float32", "float" - int8 = torch.int8, "torch.int8", "int8" - bf16 = torch.bfloat16, "torch.bfloat16", "bf16", "bfloat16" - - # Copied from https://stackoverflow.com/a/43210118 - # Allows us to use multiple values for each Enum index and returns first - # listed value when Enum is called - def __new__(cls, *values): - obj = object.__new__(cls) - # first value is canonical value - obj._value_ = values[0] - for other_value in values[1:]: - cls._value2member_map_[other_value] = obj - obj._all_values = values - return obj - - def __repr__(self): - return "<%s.%s: %s>" % ( - self.__class__.__name__, - self._name_, - ", ".join([repr(v) for v in self._all_values]), - ) - - -def get_pld_enabled(param_dict): - if PROGRESSIVE_LAYER_DROP in param_dict.keys(): - return get_scalar_param(param_dict[PROGRESSIVE_LAYER_DROP], PLD_ENABLED, PLD_ENABLED_DEFAULT) - else: - return False - - -def get_pld_params(param_dict): - if PROGRESSIVE_LAYER_DROP in param_dict.keys(): - pld_params = copy.copy(param_dict[PROGRESSIVE_LAYER_DROP]) - pld_params.pop(PLD_ENABLED) - return pld_params - else: - return False - - -def get_amp_enabled(param_dict): - if AMP in param_dict.keys(): - return get_scalar_param(param_dict[AMP], AMP_ENABLED, AMP_ENABLED_DEFAULT) - else: - return False - - -def get_amp_params(param_dict): - if AMP in param_dict.keys(): - amp_params = copy.copy(param_dict[AMP]) - amp_params.pop(AMP_ENABLED) - return amp_params - else: - return False - - -def get_torch_autocast_enabled(param_dict): - if TORCH_AUTOCAST in param_dict.keys(): - return get_scalar_param(param_dict[TORCH_AUTOCAST], TORCH_AUTOCAST_ENABLED, TORCH_AUTOCAST_ENABLED_DEFAULT) - else: - return False - - -def get_torch_autocast_dtype(param_dict): - if TORCH_AUTOCAST in param_dict: - if TORCH_AUTOCAST_DTYPE in param_dict[TORCH_AUTOCAST]: - try: - return DtypeEnum(param_dict[TORCH_AUTOCAST][TORCH_AUTOCAST_DTYPE]).value - except KeyError: - raise ValueError( - f"Invalid dtype for torch autocast: {param_dict[TORCH_AUTOCAST][TORCH_AUTOCAST_DTYPE]}") - return None - - -def get_lower_precision_safe_modules(param_dict): - if TORCH_AUTOCAST in param_dict: - if TORCH_AUTOCAST_LOWER_PRECISION_SAFE_MODULES in param_dict[TORCH_AUTOCAST]: - module_names_with_package = param_dict[TORCH_AUTOCAST][TORCH_AUTOCAST_LOWER_PRECISION_SAFE_MODULES] - if not all(isinstance(module_name, str) for module_name in module_names_with_package): - raise ValueError( - f"Invalid module names for torch autocast: {module_names_with_package}. Expected list of strings.") - return module_names_with_package - return None - - -def get_gradient_accumulation_steps(param_dict): - return get_scalar_param(param_dict, GRADIENT_ACCUMULATION_STEPS, GRADIENT_ACCUMULATION_STEPS_DEFAULT) - - -def get_sparse_gradients_enabled(param_dict): - return get_scalar_param(param_dict, SPARSE_GRADIENTS, SPARSE_GRADIENTS_DEFAULT) - - -def get_communication_data_type(param_dict, - comm_type=COMMUNICATION_DATA_TYPE, - comm_data_type_default=COMMUNICATION_DATA_TYPE_DEFAULT): - val = get_scalar_param(param_dict, comm_type, comm_data_type_default) - val = val.lower() if val is not None else val - if val is None: - return val # we must determine it by other parameters - elif val == "fp32": - return torch.float32 - elif val == "fp16": - return torch.float16 - elif val == "bf16": - return torch.bfloat16 - - raise ValueError(f"Invalid communication_data_type. Supported data types: ['fp16', 'bf16', 'fp32']. Got: {val}") - - -def get_prescale_gradients(param_dict): - return get_scalar_param(param_dict, PRESCALE_GRADIENTS, PRESCALE_GRADIENTS_DEFAULT) - - -def get_gradient_predivide_factor(param_dict): - return get_scalar_param(param_dict, GRADIENT_PREDIVIDE_FACTOR, GRADIENT_PREDIVIDE_FACTOR_DEFAULT) - - -def get_steps_per_print(param_dict): - return get_scalar_param(param_dict, STEPS_PER_PRINT, STEPS_PER_PRINT_DEFAULT) - - -def get_disable_allgather(param_dict): - return get_scalar_param(param_dict, DISABLE_ALLGATHER, DISABLE_ALLGATHER_DEFAULT) - - -def get_dump_state(param_dict): - return get_scalar_param(param_dict, DUMP_STATE, DUMP_STATE_DEFAULT) - - -def get_gradient_clipping(param_dict): - return get_scalar_param(param_dict, GRADIENT_CLIPPING, GRADIENT_CLIPPING_DEFAULT) - - -def get_graph_harvesting(param_dict): - return get_scalar_param(param_dict, GRAPH_HARVESTING, GRAPH_HARVESTING_DEFAULT) - - -def get_sparse_attention(param_dict): - if SPARSE_ATTENTION in param_dict.keys(): - sparsity = param_dict[SPARSE_ATTENTION] - mode = get_sparse_attention_mode(sparsity) - - if mode == SPARSE_DENSE_MODE: - return get_sparse_dense_config(sparsity) - elif mode == SPARSE_FIXED_MODE: - return get_sparse_fixed_config(sparsity) - elif mode == SPARSE_VARIABLE_MODE: - return get_sparse_variable_config(sparsity) - elif mode == SPARSE_BIGBIRD_MODE: - return get_sparse_bigbird_config(sparsity) - elif mode == SPARSE_BSLONGFORMER_MODE: - return get_sparse_bslongformer_config(sparsity) - else: - raise NotImplementedError(f"Given sparsity mode, {mode}, has not been implemented yet!") - - else: - return None - - -def get_sparse_dense_config(sparsity): - block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) - return {SPARSE_MODE: SPARSE_DENSE_MODE, SPARSE_BLOCK: block} - - -def get_sparse_fixed_config(sparsity): - block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) - different_layout_per_head = get_scalar_param( - sparsity, - SPARSE_DIFFERENT_LAYOUT_PER_HEAD, - SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT, - ) - num_local_blocks = get_scalar_param(sparsity, SPARSE_NUM_LOCAL_BLOCKS, SPARSE_NUM_LOCAL_BLOCKS_DEFAULT) - num_global_blocks = get_scalar_param(sparsity, SPARSE_NUM_GLOBAL_BLOCKS, SPARSE_NUM_GLOBAL_BLOCKS_DEFAULT) - attention = get_scalar_param(sparsity, SPARSE_ATTENTION_TYPE, SPARSE_ATTENTION_TYPE_DEFAULT) - horizontal_global_attention = get_scalar_param( - sparsity, - SPARSE_HORIZONTAL_GLOBAL_ATTENTION, - SPARSE_HORIZONTAL_GLOBAL_ATTENTION_DEFAULT, - ) - num_different_global_patterns = get_scalar_param( - sparsity, - SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS, - SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS_DEFAULT, - ) - - return { - SPARSE_MODE: SPARSE_FIXED_MODE, - SPARSE_BLOCK: block, - SPARSE_DIFFERENT_LAYOUT_PER_HEAD: different_layout_per_head, - SPARSE_NUM_LOCAL_BLOCKS: num_local_blocks, - SPARSE_NUM_GLOBAL_BLOCKS: num_global_blocks, - SPARSE_ATTENTION_TYPE: attention, - SPARSE_HORIZONTAL_GLOBAL_ATTENTION: horizontal_global_attention, - SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS: num_different_global_patterns, - } - - -def get_sparse_variable_config(sparsity): - block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) - different_layout_per_head = get_scalar_param( - sparsity, - SPARSE_DIFFERENT_LAYOUT_PER_HEAD, - SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT, - ) - num_random_blocks = get_scalar_param(sparsity, SPARSE_NUM_RANDOM_BLOCKS, SPARSE_NUM_RANDOM_BLOCKS_DEFAULT) - local_window_blocks = get_scalar_param(sparsity, SPARSE_LOCAL_WINDOW_BLOCKS, SPARSE_LOCAL_WINDOW_BLOCKS_DEFAULT) - global_block_indices = get_scalar_param(sparsity, SPARSE_GLOBAL_BLOCK_INDICES, SPARSE_GLOBAL_BLOCK_INDICES_DEFAULT) - global_block_end_indices = get_scalar_param( - sparsity, - SPARSE_GLOBAL_BLOCK_END_INDICES, - SPARSE_GLOBAL_BLOCK_END_INDICES_DEFAULT, - ) - attention = get_scalar_param(sparsity, SPARSE_ATTENTION_TYPE, SPARSE_ATTENTION_TYPE_DEFAULT) - horizontal_global_attention = get_scalar_param( - sparsity, - SPARSE_HORIZONTAL_GLOBAL_ATTENTION, - SPARSE_HORIZONTAL_GLOBAL_ATTENTION_DEFAULT, - ) - - return { - SPARSE_MODE: SPARSE_VARIABLE_MODE, - SPARSE_BLOCK: block, - SPARSE_DIFFERENT_LAYOUT_PER_HEAD: different_layout_per_head, - SPARSE_NUM_RANDOM_BLOCKS: num_random_blocks, - SPARSE_LOCAL_WINDOW_BLOCKS: local_window_blocks, - SPARSE_GLOBAL_BLOCK_INDICES: global_block_indices, - SPARSE_GLOBAL_BLOCK_END_INDICES: global_block_end_indices, - SPARSE_ATTENTION_TYPE: attention, - SPARSE_HORIZONTAL_GLOBAL_ATTENTION: horizontal_global_attention, - } - - -def get_sparse_bigbird_config(sparsity): - block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) - different_layout_per_head = get_scalar_param( - sparsity, - SPARSE_DIFFERENT_LAYOUT_PER_HEAD, - SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT, - ) - num_random_blocks = get_scalar_param(sparsity, SPARSE_NUM_RANDOM_BLOCKS, SPARSE_NUM_RANDOM_BLOCKS_DEFAULT) - num_sliding_window_blocks = get_scalar_param( - sparsity, - SPARSE_NUM_SLIDING_WINDOW_BLOCKS, - SPARSE_NUM_SLIDING_WINDOW_BLOCKS_DEFAULT, - ) - num_global_blocks = get_scalar_param(sparsity, SPARSE_NUM_GLOBAL_BLOCKS, SPARSE_NUM_GLOBAL_BLOCKS_DEFAULT) - - return { - SPARSE_MODE: SPARSE_BIGBIRD_MODE, - SPARSE_BLOCK: block, - SPARSE_DIFFERENT_LAYOUT_PER_HEAD: different_layout_per_head, - SPARSE_NUM_RANDOM_BLOCKS: num_random_blocks, - SPARSE_NUM_SLIDING_WINDOW_BLOCKS: num_sliding_window_blocks, - SPARSE_NUM_GLOBAL_BLOCKS: num_global_blocks, - } - - -def get_sparse_bslongformer_config(sparsity): - block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) - different_layout_per_head = get_scalar_param( - sparsity, - SPARSE_DIFFERENT_LAYOUT_PER_HEAD, - SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT, - ) - num_sliding_window_blocks = get_scalar_param( - sparsity, - SPARSE_NUM_SLIDING_WINDOW_BLOCKS, - SPARSE_NUM_SLIDING_WINDOW_BLOCKS_DEFAULT, - ) - global_block_indices = get_scalar_param(sparsity, SPARSE_GLOBAL_BLOCK_INDICES, SPARSE_GLOBAL_BLOCK_INDICES_DEFAULT) - global_block_end_indices = get_scalar_param( - sparsity, - SPARSE_GLOBAL_BLOCK_END_INDICES, - SPARSE_GLOBAL_BLOCK_END_INDICES_DEFAULT, - ) - - return { - SPARSE_MODE: SPARSE_BSLONGFORMER_MODE, - SPARSE_BLOCK: block, - SPARSE_DIFFERENT_LAYOUT_PER_HEAD: different_layout_per_head, - SPARSE_NUM_SLIDING_WINDOW_BLOCKS: num_sliding_window_blocks, - SPARSE_GLOBAL_BLOCK_INDICES: global_block_indices, - SPARSE_GLOBAL_BLOCK_END_INDICES: global_block_end_indices, - } - - -def get_sparse_attention_mode(param_dict): - if SPARSE_MODE in param_dict.keys(): - return param_dict[SPARSE_MODE] - else: - return SPARSE_MODE_DEFAULT - - -def get_sparse_attention_type(param_dict): - if SPARSE_ATTENTION_TYPE in param_dict.keys(): - return param_dict[SPARSE_ATTENTION_TYPE] - else: - return SPARSE_ATTENTION_TYPE_DEFAULT - - -def get_pipeline_config(param_dict): - """Parses pipeline engine configuration. """ - default_pipeline = { - "stages": "auto", - "partition": "best", - "seed_layers": False, - "activation_checkpoint_interval": 0, - "pipe_partitioned": True, - "grad_partitioned": True, - } - config = default_pipeline - for key, val in param_dict.get("pipeline", {}).items(): - config[key] = val - return config - - -def get_optimizer_name(param_dict): - if OPTIMIZER in param_dict.keys() and TYPE in param_dict[OPTIMIZER].keys(): - return param_dict[OPTIMIZER][TYPE] - else: - return OPTIMIZER_TYPE_DEFAULT - - -def get_optimizer_params(param_dict): - if (get_optimizer_name(param_dict) is not None and OPTIMIZER_PARAMS in param_dict[OPTIMIZER].keys()): - return param_dict[OPTIMIZER][OPTIMIZER_PARAMS] - else: - return None - - -def get_optimizer_gradient_clipping(param_dict): - optimizer_params = get_optimizer_params(param_dict) - if optimizer_params is not None and MAX_GRAD_NORM in optimizer_params.keys(): - return optimizer_params[MAX_GRAD_NORM] - else: - return None - - -def get_optimizer_legacy_fusion(param_dict): - if OPTIMIZER in param_dict.keys() and LEGACY_FUSION in param_dict[OPTIMIZER].keys(): - return param_dict[OPTIMIZER][LEGACY_FUSION] - else: - return LEGACY_FUSION_DEFAULT - - -def get_zero_allow_untested_optimizer(param_dict): - return get_scalar_param(param_dict, ZERO_ALLOW_UNTESTED_OPTIMIZER, ZERO_ALLOW_UNTESTED_OPTIMIZER_DEFAULT) - - -def get_zero_force_ds_cpu_optimizer(param_dict): - return get_scalar_param(param_dict, ZERO_FORCE_DS_CPU_OPTIMIZER, ZERO_FORCE_DS_CPU_OPTIMIZER_DEFAULT) - - -def get_scheduler_name(param_dict): - if SCHEDULER in param_dict.keys() and TYPE in param_dict[SCHEDULER].keys(): - return param_dict[SCHEDULER][TYPE] - else: - return SCHEDULER_TYPE_DEFAULT - - -def get_scheduler_params(param_dict): - if (get_scheduler_name(param_dict) is not None and SCHEDULER_PARAMS in param_dict[SCHEDULER].keys()): - return param_dict[SCHEDULER][SCHEDULER_PARAMS] - else: - return None - - -def get_train_batch_size(param_dict): - return get_scalar_param(param_dict, TRAIN_BATCH_SIZE, TRAIN_BATCH_SIZE_DEFAULT) - - -def get_train_micro_batch_size_per_gpu(param_dict): - return get_scalar_param( - param_dict, - TRAIN_MICRO_BATCH_SIZE_PER_GPU, - TRAIN_MICRO_BATCH_SIZE_PER_GPU_DEFAULT, - ) - - -def get_wall_clock_breakdown(param_dict): - return get_scalar_param(param_dict, WALL_CLOCK_BREAKDOWN, WALL_CLOCK_BREAKDOWN_DEFAULT) - - -def get_memory_breakdown(param_dict): - return get_scalar_param(param_dict, MEMORY_BREAKDOWN, MEMORY_BREAKDOWN_DEFAULT) - - -class HybridEngineConfig(DeepSpeedConfigModel): - enabled: bool = False - max_out_tokens: int = 512 - inference_tp_size: int = 1 - release_inference_cache: bool = False - pin_parameters: bool = True - tp_gather_partition_size: int = 8 - - -def get_hybrid_engine_config(param_dict): - hybrid_engine_config_dict = param_dict.get("hybrid_engine", {}) - hybrid_engine_config = HybridEngineConfig(**hybrid_engine_config_dict) - return hybrid_engine_config - - -def get_expert_data_topo_config(param_dict): - return get_scalar_param(param_dict, USE_DATA_BEFORE_EXPERT_PARALLEL, USE_DATA_BEFORE_EXPERT_PARALLEL_DEFAULT) - - -def get_eigenvalue_config(param_dict): - if get_quantize_enabled(param_dict): - param_dict = param_dict[QUANTIZE_TRAINING] - assert not get_eigenvalue_enabled(param_dict), "Eigenvalue based MoQ is temporarily disabled" - return ( - get_eigenvalue_enabled(param_dict), - get_eigenvalue_verbose(param_dict), - get_eigenvalue_max_iter(param_dict), - get_eigenvalue_tol(param_dict), - get_eigenvalue_stability(param_dict), - get_eigenvalue_gas_boundary_resolution(param_dict), - get_eigenvalue_layer_name(param_dict), - get_eigenvalue_layer_num(param_dict), - ) - else: - return ( - EIGENVALUE_ENABLED_DEFAULT, - EIGENVALUE_VERBOSE_DEFAULT, - EIGENVALUE_MAX_ITER_DEFAULT, - EIGENVALUE_TOL_DEFAULT, - EIGENVALUE_STABILITY_DEFAULT, - EIGENVALUE_GAS_BOUNDARY_RESOLUTION_DEFAULT, - EIGENVALUE_LAYER_NAME_DEFAULT, - EIGENVALUE_LAYER_NUM_DEFAULT, - ) - - -def get_eigenvalue_enabled(param_dict): - if EIGENVALUE in param_dict.keys(): - return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_ENABLED, EIGENVALUE_ENABLED_DEFAULT) - else: - return EIGENVALUE_ENABLED_DEFAULT - - -def get_eigenvalue_verbose(param_dict): - if EIGENVALUE in param_dict.keys(): - return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_VERBOSE, EIGENVALUE_VERBOSE_DEFAULT) - else: - return EIGENVALUE_VERBOSE_DEFAULT - - -def get_eigenvalue_max_iter(param_dict): - if EIGENVALUE in param_dict.keys(): - return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_MAX_ITER, EIGENVALUE_MAX_ITER_DEFAULT) - else: - return EIGENVALUE_MAX_ITER_DEFAULT - - -def get_eigenvalue_tol(param_dict): - if EIGENVALUE in param_dict.keys(): - return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_TOL, EIGENVALUE_TOL_DEFAULT) - else: - return EIGENVALUE_TOL_DEFAULT - - -def get_eigenvalue_stability(param_dict): - if EIGENVALUE in param_dict.keys(): - return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_STABILITY, EIGENVALUE_STABILITY_DEFAULT) - else: - return EIGENVALUE_STABILITY_DEFAULT - - -def get_eigenvalue_gas_boundary_resolution(param_dict): - if EIGENVALUE in param_dict.keys(): - return get_scalar_param( - param_dict[EIGENVALUE], - EIGENVALUE_GAS_BOUNDARY_RESOLUTION, - EIGENVALUE_GAS_BOUNDARY_RESOLUTION_DEFAULT, - ) - else: - return EIGENVALUE_GAS_BOUNDARY_RESOLUTION_DEFAULT - - -def get_eigenvalue_layer_name(param_dict): - if EIGENVALUE in param_dict.keys(): - return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_LAYER_NAME, EIGENVALUE_LAYER_NAME_DEFAULT) - else: - return EIGENVALUE_LAYER_NAME_DEFAULT - - -def get_eigenvalue_layer_num(param_dict): - if EIGENVALUE in param_dict.keys(): - return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_LAYER_NUM, EIGENVALUE_LAYER_NUM_DEFAULT) - else: - return EIGENVALUE_LAYER_NUM_DEFAULT - - -def get_checkpoint_params(param_dict): - return param_dict.get(CHECKPOINT, {}) - - -def get_data_types_params(param_dict): - return param_dict.get(DATA_TYPES, {}) - - -def get_checkpoint_tag_validation_mode(checkpoint_params): - tag_validation_mode = checkpoint_params.get(CHECKPOINT_TAG_VALIDATION, CHECKPOINT_TAG_VALIDATION_DEFAULT) - tag_validation_mode = tag_validation_mode.upper() - if tag_validation_mode in CHECKPOINT_TAG_VALIDATION_MODES: - return tag_validation_mode - else: - raise DeepSpeedConfigError( - "Checkpoint config contains invalid tag_validation " - f"value of {tag_validation_mode}, expecting one of {CHECKPOINT_TAG_VALIDATION_MODES}") - - -def get_checkpoint_parallel_write_pipeline(checkpoint_params): - par_write_params = checkpoint_params.get(CHECKPOINT_PARALLEL_WRITE, {}) - par_write_pipeline = par_write_params.get(CHECKPOINT_PARALLEL_WRITE_PIPELINE_STAGE, - CHECKPOINT_PARALLEL_WRITE_PIPELINE_STAGE_DEFAULT) - if par_write_pipeline in [True, False]: - return par_write_pipeline - else: - raise DeepSpeedConfigError("checkpoint::parallel_write::pipeline_stage " - f"value of '{par_write_pipeline}' is invalid, expecting: true or false") - - -def get_dataloader_drop_last(param_dict): - return get_scalar_param(param_dict, DATALOADER_DROP_LAST, DATALOADER_DROP_LAST_DEFAULT) - - -'''Write deepspeed config files by modifying basic templates. -Can be used for quickly changing parameters via command line parameters.''' - - -class DeepSpeedConfigWriter: - - def __init__(self, data=None): - self.data = data if data is not None else {} - - def add_config(self, key, value): - self.data[key] = value - - def load_config(self, filename): - self.data = json.load(open(filename, "r"), object_pairs_hook=dict_raise_error_on_duplicate_keys) - - def write_config(self, filename): - with open(filename, "w") as outfile: - json.dump(self.data, outfile) - - -class DeepSpeedConfig(object): - - def __init__(self, config: Union[str, dict], mpu=None, mesh_device=None): - super(DeepSpeedConfig, self).__init__() - if isinstance(config, dict): - self._param_dict = config - elif os.path.exists(config): - self._param_dict = hjson.load(open(config, "r"), object_pairs_hook=dict_raise_error_on_duplicate_keys) - else: - try: - config_decoded = base64.urlsafe_b64decode(config).decode('utf-8') - self._param_dict = hjson.loads(config_decoded) - except (UnicodeDecodeError, AttributeError): - raise ValueError( - f"Expected a string path to an existing deepspeed config, or a dictionary or a valid base64. Received: {config}" - ) - - try: - self.global_rank = dist.get_rank() - if mpu is not None: - # Ulysses SP - if not hasattr(mpu, "get_data_parallel_world_size"): - self.world_size = dist.get_world_size() / mpu.get_sequence_parallel_world_size() - else: - self.world_size = mpu.get_data_parallel_world_size() - elif mesh_device is not None: - self.world_size = dist.get_world_size(mesh_device.get_group(mesh_dim="data_parallel")) - else: - # HF zero.init case where there is no mpu - if "sequence_parallel_size" in config: - self.world_size = dist.get_world_size() / config["sequence_parallel_size"] - else: - self.world_size = dist.get_world_size() - except: - self.global_rank = 0 - self.world_size = 1 - logger.info(f"Config mesh_device {mesh_device} world_size = {self.world_size}") - # If elastic-mode enabled, update compute + update _param_dict - self.elasticity_enabled = elasticity_enabled(self._param_dict) - if self.elasticity_enabled: - logger.info("DeepSpeed elasticity support enabled") - final_batch_size, valid_gpus, micro_batch_size = compute_elastic_config( - ds_config=self._param_dict, - target_deepspeed_version=__version__, - world_size=self.world_size, - ) - - elastic_dict = self._param_dict[ELASTICITY] - - # Ensure the resource scheduler saw the same elastic config we are using at runtime - ensure_immutable_elastic_config(runtime_elastic_config_dict=elastic_dict) - - self.elastic_model_parallel_size = elastic_dict.get(MODEL_PARALLEL_SIZE, MODEL_PARALLEL_SIZE_DEFAULT) - if self.elastic_model_parallel_size < 1: - raise ElasticityConfigError("Model-Parallel size cannot be less than 1, " - f"given model-parallel size: {self.elastic_model_parallel_size}") - - self.num_gpus_per_node = elastic_dict.get(NUM_GPUS_PER_NODE, NUM_GPUS_PER_NODE_DEFAULT) - if self.num_gpus_per_node < 1: - raise ElasticityConfigError("NUmber of GPUs per node cannot be less than 1, " - f"given number of GPUs per node: {self.num_gpus_per_node}") - - ignore_non_elastic_batch_info = elastic_dict.get(IGNORE_NON_ELASTIC_BATCH_INFO, - IGNORE_NON_ELASTIC_BATCH_INFO_DEFAULT) - - if not ignore_non_elastic_batch_info: - batch_params = [ - TRAIN_BATCH_SIZE, - TRAIN_MICRO_BATCH_SIZE_PER_GPU, - GRADIENT_ACCUMULATION_STEPS, - ] - if any(map(lambda t: t in self._param_dict, batch_params)): - raise ElasticityConfigError("One or more batch related parameters were found in your " \ - f"ds_config ({TRAIN_BATCH_SIZE}, {TRAIN_MICRO_BATCH_SIZE_PER_GPU}, and/or " \ - f"{GRADIENT_ACCUMULATION_STEPS}). These parameters *will not be used* since " \ - "elastic training is enabled, which takes control of these parameters. " \ - "If you want to suppress this error (the parameters will be silently ignored) " \ - f"please set {IGNORE_NON_ELASTIC_BATCH_INFO}':true in your elasticity config.") - - # micro_bsz * world_size * gas = total_batch_size - # gas = total_batch_size // (micro_bsz * world_size) - gradient_accu_steps = final_batch_size // (micro_batch_size * self.world_size) - - if TRAIN_BATCH_SIZE in self._param_dict: - logger.warning("[Elasticity] overriding training_batch_size: " - f"{self._param_dict[TRAIN_BATCH_SIZE]} -> {final_batch_size}") - if TRAIN_MICRO_BATCH_SIZE_PER_GPU in self._param_dict: - logger.warning("[Elasticity] overriding train_micro_batch_size_per_gpu: " - f"{self._param_dict[TRAIN_MICRO_BATCH_SIZE_PER_GPU]} -> {micro_batch_size}") - if GRADIENT_ACCUMULATION_STEPS in self._param_dict: - logger.warning("[Elasticity] overriding gradient_accumulation_steps: " - f"{self._param_dict[GRADIENT_ACCUMULATION_STEPS]} -> {gradient_accu_steps}") - - logger.info(f"[Elasticity] valid GPU counts: {valid_gpus}") - - self._param_dict[TRAIN_BATCH_SIZE] = final_batch_size - self._param_dict[TRAIN_MICRO_BATCH_SIZE_PER_GPU] = micro_batch_size - self._param_dict[GRADIENT_ACCUMULATION_STEPS] = gradient_accu_steps - - # Pass a copy so that user json is unmodified, e.g. for logging - self._initialize_params(copy.copy(self._param_dict)) - self._configure_train_batch_size() - self._do_sanity_check() - - def _initialize_params(self, param_dict): - self.train_batch_size = get_train_batch_size(param_dict) - self.train_micro_batch_size_per_gpu = get_train_micro_batch_size_per_gpu(param_dict) - self.gradient_accumulation_steps = get_gradient_accumulation_steps(param_dict) - self.steps_per_print = get_steps_per_print(param_dict) - self.dump_state = get_dump_state(param_dict) - - self.disable_allgather = get_disable_allgather(param_dict) - self.communication_data_type = get_communication_data_type(param_dict) - self.seq_parallel_communication_data_type = get_communication_data_type( - param_dict, SEQ_PARALLEL_COMMUNICATION_DATA_TYPE, SEQ_PARALLEL_COMMUNICATION_DATA_TYPE_DEFAULT) - self.prescale_gradients = get_prescale_gradients(param_dict) - self.gradient_predivide_factor = get_gradient_predivide_factor(param_dict) - self.sparse_gradients_enabled = get_sparse_gradients_enabled(param_dict) - - self.zero_config = get_zero_config(param_dict) - self.mics_shard_size = self.zero_config.mics_shard_size - self.mics_hierarchial_params_gather = self.zero_config.mics_hierarchical_params_gather - self.zero_optimization_stage = self.zero_config.stage - self.zero_enabled = self.zero_optimization_stage > 0 - - self.activation_checkpointing_config = DeepSpeedActivationCheckpointingConfig(param_dict) - - self.comms_config = DeepSpeedCommsConfig(param_dict) - self.monitor_config = get_monitor_config(param_dict) - - self.gradient_clipping = get_gradient_clipping(param_dict) - self.float16_config = get_float16_config(param_dict) - self.bfloat16_config = get_bfloat16_config(param_dict) - assert not (self.float16_config.enabled - and self.bfloat16_config.enabled), 'bfloat16 and fp16 modes cannot be simultaneously enabled' - - self.amp_enabled = get_amp_enabled(param_dict) - self.amp_params = get_amp_params(param_dict) - - self.torch_autocast_enabled = get_torch_autocast_enabled(param_dict) - self.torch_autocast_dtype = get_torch_autocast_dtype(param_dict) - self.torch_autocast_lower_precision_safe_modules = get_lower_precision_safe_modules(param_dict) - - self.compression_config = get_compression_config(param_dict) - self.graph_harvesting = get_graph_harvesting(param_dict) - - self.optimizer_name = get_optimizer_name(param_dict) - if (self.optimizer_name is not None and self.optimizer_name.lower() in DEEPSPEED_OPTIMIZERS): - self.optimizer_name = self.optimizer_name.lower() - - self.optimizer_params = get_optimizer_params(param_dict) - self.optimizer_legacy_fusion = get_optimizer_legacy_fusion(param_dict) - - self.zero_allow_untested_optimizer = get_zero_allow_untested_optimizer(param_dict) - - self.zero_force_ds_cpu_optimizer = get_zero_force_ds_cpu_optimizer(param_dict) - - self.scheduler_name = get_scheduler_name(param_dict) - self.scheduler_params = get_scheduler_params(param_dict) - - self.flops_profiler_config = DeepSpeedFlopsProfilerConfig(param_dict) - self.wall_clock_breakdown = (get_wall_clock_breakdown(param_dict) | self.flops_profiler_config.enabled) - self.memory_breakdown = get_memory_breakdown(param_dict) - self.autotuning_config = DeepSpeedAutotuningConfig(param_dict) - - ( - self.eigenvalue_enabled, - self.eigenvalue_verbose, - self.eigenvalue_max_iter, - self.eigenvalue_tol, - self.eigenvalue_stability, - self.eigenvalue_gas_boundary_resolution, - self.eigenvalue_layer_name, - self.eigenvalue_layer_num, - ) = get_eigenvalue_config(param_dict) - - self.use_data_before_expert_parallel_ = get_expert_data_topo_config(param_dict) - self.hybrid_engine = get_hybrid_engine_config(param_dict) - - self.sparse_attention = get_sparse_attention(param_dict) - self.pipeline = get_pipeline_config(param_dict) - - self.pld_enabled = get_pld_enabled(param_dict) - self.pld_params = get_pld_params(param_dict) - - self.curriculum_enabled_legacy = get_curriculum_enabled_legacy(param_dict) - self.curriculum_params_legacy = get_curriculum_params_legacy(param_dict) - - self.data_efficiency_enabled = get_data_efficiency_enabled(param_dict) - self.data_efficiency_config = get_data_efficiency_config(param_dict) - - checkpoint_params = get_checkpoint_params(param_dict) - validation_mode = get_checkpoint_tag_validation_mode(checkpoint_params) - self.checkpoint_tag_validation_enabled = (validation_mode != ValidationMode.IGNORE) - self.checkpoint_tag_validation_fail = validation_mode == ValidationMode.FAIL - self.load_universal_checkpoint = checkpoint_params.get(LOAD_UNIVERSAL_CHECKPOINT, - LOAD_UNIVERSAL_CHECKPOINT_DEFAULT) - - self.use_node_local_storage = checkpoint_params.get(USE_NODE_LOCAL_STORAGE_CHECKPOINT, - USE_NODE_LOCAL_STORAGE_CHECKPOINT_DEFAULT) - - data_types_params = get_data_types_params(param_dict) - self.grad_accum_dtype = data_types_params.get(GRAD_ACCUM_DTYPE, GRAD_ACCUM_DTYPE_DEFAULT) - - par_write_pipe = get_checkpoint_parallel_write_pipeline(checkpoint_params) - self.checkpoint_parallel_write_pipeline = par_write_pipe - - self.aio_config = get_aio_config(param_dict) - - self.dataloader_drop_last = get_dataloader_drop_last(param_dict) - - self.nebula_config = DeepSpeedNebulaConfig(param_dict) - self.checkpoint_config = get_checkpoint_config(param_dict) - - self.weight_quantization_config = WeightQuantConfig( - **param_dict['weight_quantization']) if 'weight_quantization' in param_dict else None - - self.compile_config = CompileConfig(**param_dict.get('compile', {})) - - self.timers_config = get_timers_config(param_dict) - self.tensor_parallel_config = get_tensor_parallel_config(param_dict) - - def _batch_assertion(self): - - train_batch = self.train_batch_size - micro_batch = self.train_micro_batch_size_per_gpu - grad_acc = self.gradient_accumulation_steps - - assert (train_batch > 0), f"Train batch size: {train_batch} has to be greater than 0" - - assert (micro_batch > 0), f"Micro batch size per gpu: {micro_batch} has to be greater than 0" - - assert (grad_acc > 0), f"Gradient accumulation steps: {grad_acc} has to be greater than 0" - - assert train_batch == micro_batch * grad_acc * self.world_size, ( - f"Check batch related parameters. train_batch_size is not equal " - "to micro_batch_per_gpu * gradient_acc_step * world_size " - f"{train_batch} != {micro_batch} * {grad_acc} * {self.world_size}") - - def _set_batch_related_parameters(self): - - train_batch = self.train_batch_size - micro_batch = self.train_micro_batch_size_per_gpu - grad_acc = self.gradient_accumulation_steps - - #print(f"in: train_batch = {train_batch}, micro_batch={micro_batch}") - - # all values are provided nothing needs to be set - if train_batch is not None and micro_batch is not None and grad_acc is not None: - return - - # global_accumulation_steps needs to be set - elif train_batch is not None and micro_batch is not None: - grad_acc = train_batch // micro_batch - grad_acc //= self.world_size - self.gradient_accumulation_steps = grad_acc - - # micro_batch_per_gpu needs to be set - elif train_batch is not None and grad_acc is not None: - micro_batch = train_batch // self.world_size - micro_batch //= grad_acc - self.train_micro_batch_size_per_gpu = micro_batch - - # train_batch_size needs to be set - elif micro_batch is not None and grad_acc is not None: - train_batch_size = micro_batch * grad_acc - train_batch_size *= self.world_size - self.train_batch_size = train_batch_size - - # gradient_accumulation_steps and micro_batch_per_gpus is set - elif train_batch is not None: - self.gradient_accumulation_steps = 1 - self.train_micro_batch_size_per_gpu = train_batch // self.world_size - - # train_batch_size and gradient_accumulation_step is set - elif micro_batch is not None: - self.train_batch_size = micro_batch * self.world_size - self.gradient_accumulation_steps = 1 - - # either none of the three parameters are provided or just gradient_accumulation_step is provided - else: - assert False, \ - 'Either train_batch_size or train_micro_batch_size_per_gpu needs to be provided' - - #print(f"final: {self.train_batch_size=} {self.train_micro_batch_size_per_gpu=} {self.gradient_accumulation_steps=}") - - def _configure_train_batch_size(self): - self._set_batch_related_parameters() - self._batch_assertion() - - def _do_sanity_check(self): - self._do_error_check() - - self._do_warning_check() - - def print_user_config(self): - logger.info(" json = {}".format( - json.dumps( - self._param_dict, - sort_keys=True, - indent=4, - cls=ScientificNotationEncoder, - separators=(",", ":"), - ))) - - def print(self, name): - logger.info("{}:".format(name)) - for arg in sorted(vars(self)): - if arg != "_param_dict": - dots = "." * (29 - len(arg)) - logger.info(" {} {} {}".format(arg, dots, getattr(self, arg))) - - self.print_user_config() - - def _do_error_check(self): - assert (self.train_micro_batch_size_per_gpu - ), "DeepSpeedConfig: {} is not defined".format(TRAIN_MICRO_BATCH_SIZE_PER_GPU) - - assert ( - self.gradient_accumulation_steps), "DeepSpeedConfig: {} is not defined".format(GRADIENT_ACCUMULATION_STEPS) - - if self.zero_enabled: - assert (self.zero_optimization_stage - <= ZeroStageEnum.max_stage), "DeepSpeedConfig: Maximum supported ZeRO stage is {}".format( - ZeroStageEnum.max_stage) - - if self.float16_config.fp16_master_weights_and_grads: - assert self.zero_enabled and self.zero_optimization_stage == ZeroStageEnum.gradients, "Fp16_master_weights_and_grads is only supported with ZeRO Stage 2 for now." - - def _do_warning_check(self): - fp16_enabled = self.float16_config.enabled - - vocabulary_size = self._param_dict.get(VOCABULARY_SIZE, VOCABULARY_SIZE_DEFAULT) - if vocabulary_size and vocabulary_size % TENSOR_CORE_ALIGN_SIZE != 0: - logger.warning( - "DeepSpeedConfig: vocabulary size {} is not aligned to {}, may import tensor core utilization.".format( - vocabulary_size, TENSOR_CORE_ALIGN_SIZE)) - - if (self.optimizer_params is not None and MAX_GRAD_NORM in self.optimizer_params.keys() - and self.optimizer_params[MAX_GRAD_NORM] > 0): - if fp16_enabled: - if self.global_rank == 0: - logger.warning("DeepSpeedConfig: In FP16 mode, DeepSpeed will pass {}:{} to FP16 wrapper".format( - MAX_GRAD_NORM, self.optimizer_params[MAX_GRAD_NORM])) - else: - if self.global_rank == 0: - logger.warning( - "DeepSpeedConfig: In FP32 mode, DeepSpeed does not permit MAX_GRAD_NORM ({}) > 0, setting to zero" - .format(self.optimizer_params[MAX_GRAD_NORM])) - self.optimizer_params[MAX_GRAD_NORM] = 0.0 diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/config_utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/config_utils.py deleted file mode 100644 index 2ae6921c363b5201ce191c9630ea9daa3774e19b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/config_utils.py +++ /dev/null @@ -1,212 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Collection of DeepSpeed configuration utilities -""" -import collections -import json -import torch -from functools import reduce -from pydantic import BaseModel, ConfigDict, field_serializer - -from deepspeed.utils import logger - - -class DeepSpeedConfigModel(BaseModel): - """ - This class should be used as a base for all DeepSpeed configs. It extends - pydantic.BaseModel to allow for deprecated fields. To enable this feature, - add deprecated=True to pydantic.Field: - - my_dep_field: int = Field(0, deprecated=True) - - Deprecated Field kwargs: - - deprecated: [True|False], default False - Enables / Disables deprecated fields - - deprecated_msg: str, default "" - Message to include with deprecation warning - - new_param: str, default "" - Name of the field replacing the deprecated field - - set_new_param: [True|False], default True - If new_param is provided, enables setting the value of that param with - deprecated field value - - new_param_fn: callable, default (lambda x: x) - If new_param is provided and set_new_param is True, this function will - modify the value of the deprecated field before placing that value in - the new_param field - - Example: - my_new_field is replacing a deprecated my_old_field. The expected type - for my_new_field is int while the expected type for my_old_field is - str. We want to maintain backward compatibility with our configs, so we - define the fields with: - - class MyExampleConfig(DeepSpeedConfigModel): - my_new_field: int = 0 - my_old_field: str = Field('0', - deprecated=True, - new_param='my_new_field', - new_param_fn=(lambda x: int(x))) - """ - - def __init__(self, strict=False, **data): - if (not strict): # This is temporary until we refactor all DS configs, allows HF to load models - data = {k: v for k, v in data.items() if (v != "auto" or k == "replace_method")} - super().__init__(**data) - self._deprecated_fields_check() - - def _process_deprecated_field(self, dep_field): - # Get information about the deprecated field - pydantic_config = self - fields_set = pydantic_config.model_fields_set - kwargs = type(pydantic_config).model_fields[dep_field].json_schema_extra - new_param_fn = kwargs.get("new_param_fn", lambda x: x) - param_value = new_param_fn(getattr(pydantic_config, dep_field)) - new_field = kwargs.get("new_param", "") - dep_msg = kwargs.get("deprecated_msg", "") - if dep_field in fields_set: - logger.warning(f"Config parameter {dep_field} is deprecated" + - (f" use {new_field} instead" if new_field else "") + (f". {dep_msg}" if dep_msg else "")) - # Check if there is a new param and if it should be set with a value - if new_field and kwargs.get("set_new_param", True): - # Remove the deprecate field if there is a replacing field - try: - delattr(pydantic_config, dep_field) - except Exception as e: - logger.error(f"Tried removing deprecated '{dep_field}' from config") - raise e - - # Set new param value - new_param_nested = new_field.split(".") - if len(new_param_nested) > 1: - # If the new param exists in a subconfig, we need to get - # the fields set for that subconfig - pydantic_config = reduce(getattr, new_param_nested[:-1], pydantic_config) - fields_set = pydantic_config.model_fields_set - new_param_name = new_param_nested[-1] - assert ( - new_param_name not in fields_set - ), f"Cannot provide deprecated parameter '{dep_field}' and replacing parameter '{new_field}' together" - # A custom function for converting the old param value to new param value can be provided - try: - setattr(pydantic_config, new_param_name, param_value) - except Exception as e: - logger.error(f"Tried setting value for '{new_field}' with value from deprecated '{dep_field}'") - raise e - - def _deprecated_fields_check(self): - fields = type(self).model_fields - for field_name, field_info in fields.items(): - if field_info.json_schema_extra and field_info.json_schema_extra.get("deprecated", False): - self._process_deprecated_field(field_name) - - model_config = ConfigDict( - validate_default=True, - validate_assignment=True, - use_enum_values=True, - populate_by_name=True, - extra="forbid", - arbitrary_types_allowed=True, - protected_namespaces=(), - ) - - @field_serializer("dtype", check_fields=False) - def serialize_torch_dtype(dtype: torch.dtype) -> str: - return str(dtype) - - -def get_config_default(config, field_name): - assert field_name in config.model_fields, f"'{field_name}' is not a field in {config}" - assert not config.model_fields.get( - field_name).is_required(), f"'{field_name}' is a required field and does not have a default value" - return config.model_fields.get(field_name).get_default() - - -class pp_int(int): - """ - A wrapper for integers that will return a custom string or comma-formatted - string of the integer. For example, print(pp_int(1e5)) will return - "10,000". This is useful mainly for auto-generated documentation purposes. - """ - - def __new__(cls, val, custom_print_str=None): - inst = super().__new__(cls, val) - inst.custom_print_str = custom_print_str - return inst - - def __repr__(self): - if self.custom_print_str: - return self.custom_print_str - return f"{self.real:,}" - - -# adapted from https://stackoverflow.com/a/50701137/9201239 -class ScientificNotationEncoder(json.JSONEncoder): - """ - This class overrides ``json.dumps`` default formatter. - - This version keeps everything as normal except formats numbers bigger than 1e3 using scientific notation. - - Just pass ``cls=ScientificNotationEncoder`` to ``json.dumps`` to activate it - - """ - - def iterencode(self, o, _one_shot=False, level=0): - indent = self.indent if self.indent is not None else 4 - prefix_close = " " * level * indent - level += 1 - prefix = " " * level * indent - if isinstance(o, bool): - return "true" if o else "false" - elif isinstance(o, float) or isinstance(o, int): - if o > 1e3: - return f"{o:e}" - else: - return f"{o}" - elif isinstance(o, collections.abc.Mapping): - x = [f'\n{prefix}"{k}": {self.iterencode(v, level=level)}' for k, v in o.items()] - return "{" + ", ".join(x) + f"\n{prefix_close}" + "}" - elif isinstance(o, collections.abc.Sequence) and not isinstance(o, str): - return f"[{ f', '.join(map(self.iterencode, o)) }]" - return "\n, ".join(super().iterencode(o, _one_shot)) - - -class DeepSpeedConfigObject(object): - """ - For json serialization - """ - - def repr(self): - return self.__dict__ - - def __repr__(self): - return json.dumps( - self.__dict__, - sort_keys=True, - indent=4, - cls=ScientificNotationEncoder, - ) - - -def get_scalar_param(param_dict, param_name, param_default_value): - return param_dict.get(param_name, param_default_value) - - -def get_list_param(param_dict, param_name, param_default_value): - return param_dict.get(param_name, param_default_value) - - -def get_dict_param(param_dict, param_name, param_default_value): - return param_dict.get(param_name, param_default_value) - - -def dict_raise_error_on_duplicate_keys(ordered_pairs): - """Reject duplicate keys.""" - d = dict((k, v) for k, v in ordered_pairs) - if len(d) != len(ordered_pairs): - counter = collections.Counter([pair[0] for pair in ordered_pairs]) - keys = [key for key, value in counter.items() if value > 1] - raise ValueError("Duplicate keys in DeepSpeed config: {}".format(keys)) - return d diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/constants.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/constants.py deleted file mode 100644 index 7fe5076e5d9043f1553b4f686d4fb2d39afa32d8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/constants.py +++ /dev/null @@ -1,492 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.accelerator import get_accelerator - -############################################# -# Routes -############################################# -ROUTE_TRAIN = "train" -ROUTE_EVAL = "eval" -ROUTE_PREDICT = "predict" -ROUTE_ENCODE = "encode" - -############################################# -# Batch size -############################################# -TRAIN_BATCH_SIZE = "train_batch_size" -TRAIN_BATCH_SIZE_DEFAULT = None - -############################################# -# Sparse attention -############################################# -SPARSE_ATTENTION = "sparse_attention" -SPARSE_DENSE_MODE = "dense" -SPARSE_FIXED_MODE = "fixed" -SPARSE_VARIABLE_MODE = "variable" -SPARSE_BIGBIRD_MODE = "bigbird" -SPARSE_BSLONGFORMER_MODE = "bslongformer" -SPARSE_MODE = "mode" -SPARSE_MODE_DEFAULT = SPARSE_FIXED_MODE -SPARSE_BLOCK = "block" -SPARSE_BLOCK_DEFAULT = 16 -SPARSE_DIFFERENT_LAYOUT_PER_HEAD = "different_layout_per_head" -SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT = False -SPARSE_NUM_LOCAL_BLOCKS = "num_local_blocks" -SPARSE_NUM_LOCAL_BLOCKS_DEFAULT = 4 -SPARSE_NUM_GLOBAL_BLOCKS = "num_global_blocks" -SPARSE_NUM_GLOBAL_BLOCKS_DEFAULT = 1 -SPARSE_ATTENTION_TYPE = "attention" -SPARSE_ATTENTION_TYPE_DEFAULT = "bidirectional" -SPARSE_HORIZONTAL_GLOBAL_ATTENTION = "horizontal_global_attention" -SPARSE_HORIZONTAL_GLOBAL_ATTENTION_DEFAULT = False -SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS = "num_different_global_patterns" -SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS_DEFAULT = 1 -SPARSE_NUM_RANDOM_BLOCKS = "num_random_blocks" -SPARSE_NUM_RANDOM_BLOCKS_DEFAULT = 0 -SPARSE_LOCAL_WINDOW_BLOCKS = "local_window_blocks" -SPARSE_LOCAL_WINDOW_BLOCKS_DEFAULT = [4] -SPARSE_GLOBAL_BLOCK_INDICES = "global_block_indices" -SPARSE_GLOBAL_BLOCK_INDICES_DEFAULT = [0] -SPARSE_GLOBAL_BLOCK_END_INDICES = "global_block_end_indices" -SPARSE_GLOBAL_BLOCK_END_INDICES_DEFAULT = None -SPARSE_NUM_SLIDING_WINDOW_BLOCKS = "num_sliding_window_blocks" -SPARSE_NUM_SLIDING_WINDOW_BLOCKS_DEFAULT = 3 - -############################################# -# Optimizer and lr scheduler -############################################# -OPTIMIZER = "optimizer" -OPTIMIZER_TYPE_DEFAULT = None -OPTIMIZER_PARAMS = "params" -TYPE = "type" -LEGACY_FUSION = "legacy_fusion" -LEGACY_FUSION_DEFAULT = False -SCHEDULER = "scheduler" -SCHEDULER_TYPE_DEFAULT = None -SCHEDULER_PARAMS = "params" -MAX_GRAD_NORM = 'max_grad_norm' - -############################################# -# Optimizer and lr scheduler -############################################# -ZERO_ALLOW_UNTESTED_OPTIMIZER = "zero_allow_untested_optimizer" -ZERO_ALLOW_UNTESTED_OPTIMIZER_DEFAULT = False -ZERO_FORCE_DS_CPU_OPTIMIZER = "zero_force_ds_cpu_optimizer" -ZERO_FORCE_DS_CPU_OPTIMIZER_DEFAULT = True - -# Steps -STEPS_PER_PRINT = "steps_per_print" -STEPS_PER_PRINT_DEFAULT = None - -######################################### -# Training micro batch size per GPU -######################################### -# Batch size for one training step. This is used when the -# TRAIN_BATCH_SIZE cannot fit in GPU memory to determine -# the number of gradient accumulation steps. By default, this -# is set to None. Users can configure in ds_config.json as below example: -TRAIN_MICRO_BATCH_SIZE_PER_GPU = ''' -TRAIN_MICRO_BATCH_SIZE_PER_GPU is defined in this format: -"train_micro_batch_size_per_gpu": 1 -''' -TRAIN_MICRO_BATCH_SIZE_PER_GPU = "train_micro_batch_size_per_gpu" -TRAIN_MICRO_BATCH_SIZE_PER_GPU_DEFAULT = None - -######################################### -# Gradient Accumulation -######################################### -# Gradient accumulation feature. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -GRADIENT_ACCUMULATION_FORMAT = ''' -Gradient Accumulation should be of the format: -"gradient_accumulation_steps": 1 -''' -GRADIENT_ACCUMULATION_STEPS = "gradient_accumulation_steps" -GRADIENT_ACCUMULATION_STEPS_DEFAULT = None - -# DeepSpeed CSR gradient sparsity -SPARSE_GRADIENTS = "sparse_gradients" -SPARSE_GRADIENTS_DEFAULT = False - -######################################### -# BFLOAT16 support -######################################### -# BFLOAT16 feature. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -BFLOAT16_FORMAT = ''' -BFLOAT16 parameters should be of the format: -"bf16": { - "enabled": true, - "immediate_grad_update": false, - "check_overflow": false -} -''' -BFLOAT16 = "bf16" -BFLOAT16_OLD = "bfloat16" # keeping for backwards compatibility - -BFLOAT16_ENABLED = "enabled" -BFLOAT16_ENABLED_DEFAULT = False - -CHECK_OVERFLOW = "check_overflow" -BFLOAT16_CHECK_OVERFLOW_DEFAULT = False - -# BFLOAT16 optimizer immediate gradient update -BFLOAT16_IMMEDIATE_GRAD_UPDATE = "immediate_grad_update" -BFLOAT16_IMMEDIATE_GRAD_UPDATE_DEFAULT = True - -######################################### -# FP16 support -######################################### -# FP16 feature. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -FP16_FORMAT = ''' -FP16 parameters should be of the format: -"fp16": { - "enabled": true, - "auto_cast": false, - "loss_scale": 0, - "initial_scale_power": 16, - "loss_scale_window": 1000, - "hysteresis": 2, - "consecutive_hysteresis": false, - "min_loss_scale": 1 -} -''' -FP16 = "fp16" - -FP16_ENABLED = "enabled" -FP16_ENABLED_DEFAULT = False - -# FP16 loss scale, zero means using dynamic scaling -FP16_LOSS_SCALE = "loss_scale" -FP16_LOSS_SCALE_DEFAULT = 0 - -FP16_AUTO_CAST = "auto_cast" -FP16_AUTO_CAST_DEFAULT = False - -# FP16 initial dynamic scale loss power -FP16_INITIAL_SCALE_POWER = "initial_scale_power" -FP16_INITIAL_SCALE_POWER_DEFAULT = 16 - -# FP16 loss scale window -FP16_LOSS_SCALE_WINDOW = "loss_scale_window" -FP16_LOSS_SCALE_WINDOW_DEFAULT = 1000 - -# FP16 hysteresis -FP16_HYSTERESIS = "hysteresis" -FP16_HYSTERESIS_DEFAULT = 2 - -# FP16 consecutive hysteresis -FP16_CONSECUTIVE_HYSTERESIS = "consecutive_hysteresis" -FP16_CONSECUTIVE_HYSTERESIS_DEFAULT = False - -# FP16 min loss scale -FP16_MIN_LOSS_SCALE = "min_loss_scale" -FP16_MIN_LOSS_SCALE_DEFAULT = 1 - -# FP16 master and grads -FP16_MASTER_WEIGHTS_AND_GRADS = "fp16_master_weights_and_grads" -FP16_MASTER_WEIGHTS_AND_GRADS_DEFAULT = False - -######################################### -# Apex AMP support -######################################### -# Use Apex AMP for mixed precision support, all parameters (other than 'enabled') will be passed to -# amp.initialize(model, optimizer, **amp_params) -# See apex documentation for supported parameters/features: https://nvidia.github.io/apex/amp.html#apex.amp.initialize -AMP_FORMAT = ''' -"amp" { - "enabled: true, - "opt_level": "O1", - ... -} -''' -AMP = "amp" - -AMP_ENABLED = "enabled" -AMP_ENABLED_DEFAULT = False - -######################################### -# Torch AMP support -######################################### -TORCH_AUTOCAST_FORMAT = ''' -PyTorch autocast config should be of the format: -"torch_autocast": { - "enabled": true, - "dtype": "bfloat16", - "lower_precision_safe_modules": [ - "torch.nn.modules.linear.Linear", - "torch.nn.modules.conv.Conv2d" - ] -} -''' -TORCH_AUTOCAST = "torch_autocast" - -TORCH_AUTOCAST_ENABLED = "enabled" -TORCH_AUTOCAST_ENABLED_DEFAULT = False -TORCH_AUTOCAST_DTYPE = "dtype" -TORCH_AUTOCAST_LOWER_PRECISION_SAFE_MODULES = "lower_precision_safe_modules" - -######################################### -# Gradient clipping -######################################### -# Gradient clipping. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -GRADIENT_CLIPPING_FORMAT = ''' -Gradient clipping should be enabled as: -"gradient_clipping": 1.0 -''' -GRADIENT_CLIPPING = 'gradient_clipping' -GRADIENT_CLIPPING_DEFAULT = 0. - -######################################### -# Capture graph for short kernels sequences -######################################### -# Graph harvesting. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -GRAPH_HARVESTING_FORMAT = ''' -Graph harvesting should be enabled as: -"graph_harvesting": true -''' -GRAPH_HARVESTING = 'graph_harvesting' -GRAPH_HARVESTING_DEFAULT = False - -######################################### -# Communication data type -######################################### -# Supported types: ['none', 'fp16', 'fp32'] -# By default, this feature is not enabled ('none' value) -# Users can configure in ds_config.json as below example: -COMMUNICATION_DATA_TYPE_FORMAT = ''' -Communication data type should be set as: -"communication_data_type": "fp32" -''' -COMMUNICATION_DATA_TYPE = "communication_data_type" -COMMUNICATION_DATA_TYPE_DEFAULT = None - -########################################################### -# Gradient communication data type for sequence parallelism -########################################################### -# Supported types: ['fp16', 'bf16','fp32'] -# Default value is fp32 -# Users can configure in ds_config.json as below example: -SEQ_PARALLEL_COMMUNICATION_DATA_TYPE_FORMAT = ''' -Optional comm data type for seq paralleism should be set as: -"seq_parallel_communication_data_type": "fp32" -''' -SEQ_PARALLEL_COMMUNICATION_DATA_TYPE = "seq_parallel_communication_data_type" - -if get_accelerator().device_name == 'cuda' and get_accelerator().communication_backend_version() >= (2, 27, 3): - # nccl>=2.27.3 uses fp32 accumulation for half precision inputs, so there is no need to waste compute and memory to manually upcast to fp32 unless the user wants it and then override - SEQ_PARALLEL_COMMUNICATION_DATA_TYPE_DEFAULT = None -else: - SEQ_PARALLEL_COMMUNICATION_DATA_TYPE_DEFAULT = "fp32" - -SEQ_PARALLEL_COMMUNICATION_DATA_TYPE_DEFAULT = "fp32" - -######################################### -# Scale/predivide gradients before allreduce -######################################### -# Prescale gradients. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -PRESCALE_GRADIENTS_FORMAT = ''' -Gradient prescaling should be enabled as: -"prescale_gradients": true -''' -PRESCALE_GRADIENTS = "prescale_gradients" -PRESCALE_GRADIENTS_DEFAULT = False - -GRADIENT_PREDIVIDE_FACTOR_FORMAT = ''' -Gradient predivide factor should be enabled as: -"gradient_predivide_factor": 1.0 -''' -GRADIENT_PREDIVIDE_FACTOR = "gradient_predivide_factor" -GRADIENT_PREDIVIDE_FACTOR_DEFAULT = 1.0 - -######################################### -# Disable AllGather -######################################### -# Disable AllGather. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -DISABLE_ALLGATHER_FORMAT = ''' -Disable AllGather should be enabled as: -"disable_allgather": true -''' -DISABLE_ALLGATHER = "disable_allgather" -DISABLE_ALLGATHER_DEFAULT = False - -######################################### -# Dump DeepSpeed state -######################################### -# Dump State. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -DUMP_STATE_FORMAT = ''' -Dump state should be enabled as: -"dump_state": true -''' -DUMP_STATE = 'dump_state' -DUMP_STATE_DEFAULT = False - -######################################### -# Vocabulary size -######################################### -# Vocabulary size. -# Users can configure in ds_config.json as below example: -VOCABULARY_SIZE_FORMAT = ''' -Vocabulary size can be specified as: -"vocabulary_size": 1024 -''' -VOCABULARY_SIZE = 'vocabulary_size' -VOCABULARY_SIZE_DEFAULT = None - -######################################### -# Wall block breakdown -######################################### -# Wall clock breakdown. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -WALL_CLOCK_BREAKDOWN_FORMAT = ''' -Wall block breakdown should be enabled as: -"wall_clock_breakdown": true -''' -WALL_CLOCK_BREAKDOWN = 'wall_clock_breakdown' -WALL_CLOCK_BREAKDOWN_DEFAULT = False - -MEMORY_BREAKDOWN = 'memory_breakdown' -MEMORY_BREAKDOWN_DEFAULT = False - -######################################### -# Eigenvalue -######################################### -# Eigenvalue computation. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -EIGENVALUE_FORMAT = ''' -Tensorboard can be specified as: -"eigenvalue": { - "enabled": true, - "verbose": true, - "max_iter": 100, - "tol": 1e-2, - "stability": 1e-6 -} -''' -EIGENVALUE = "eigenvalue" - -# Tensorboard enable signal -EIGENVALUE_ENABLED = "enabled" -EIGENVALUE_ENABLED_DEFAULT = False - -EIGENVALUE_VERBOSE = "verbose" -EIGENVALUE_VERBOSE_DEFAULT = False - -EIGENVALUE_MAX_ITER = "max_iter" -EIGENVALUE_MAX_ITER_DEFAULT = 100 - -EIGENVALUE_TOL = "tol" -EIGENVALUE_TOL_DEFAULT = 1e-2 - -EIGENVALUE_STABILITY = "stability" -EIGENVALUE_STABILITY_DEFAULT = 1e-6 - -EIGENVALUE_GAS_BOUNDARY_RESOLUTION = "gas_boundary_resolution" -EIGENVALUE_GAS_BOUNDARY_RESOLUTION_DEFAULT = 1 - -EIGENVALUE_LAYER_NAME = "layer_name" -EIGENVALUE_LAYER_NAME_DEFAULT = "bert.encoder.layer" - -EIGENVALUE_LAYER_NUM = "layer_num" -EIGENVALUE_LAYER_NUM_DEFAULT = 0 - -######################################### -# Progressive Layer Drop (PLD) -######################################### -PROGRESSIVE_LAYER_DROP = "progressive_layer_drop" - -# PLD enable signal -PLD_ENABLED = "enabled" -PLD_ENABLED_DEFAULT = False - -PLD_THETA = "theta" -PLD_THETA_DEFAULT = 1.0 - -PLD_GAMMA = "gamma" -PLD_GAMMA_DEFAULT = 0.001 - - -######################################### -# Validation modes -######################################### -class ValidationMode: - WARN = "WARN" - IGNORE = "IGNORE" - FAIL = "FAIL" - - -######################################### -# Checkpoint config params -######################################### -# "checkpoint": { -# tag_validation=["Ignore"|"Warn"|"Fail"] -# load_universal=false -# use_node_local_storage=false -# parallel_write: { -# pipeline_stage: [True|False] -# } -# } -CHECKPOINT = "checkpoint" -CHECKPOINT_TAG_VALIDATION = "tag_validation" -CHECKPOINT_TAG_VALIDATION_DEFAULT = ValidationMode.WARN -CHECKPOINT_TAG_VALIDATION_MODES = [ValidationMode.WARN, ValidationMode.IGNORE, ValidationMode.FAIL] - -LOAD_UNIVERSAL_CHECKPOINT = "load_universal" -LOAD_UNIVERSAL_CHECKPOINT_DEFAULT = False - -USE_NODE_LOCAL_STORAGE_CHECKPOINT = "use_node_local_storage" -USE_NODE_LOCAL_STORAGE_CHECKPOINT_DEFAULT = False - -CHECKPOINT_PARALLEL_WRITE = "parallel_write" -CHECKPOINT_PARALLEL_WRITE_PIPELINE_STAGE = "pipeline_stage" -CHECKPOINT_PARALLEL_WRITE_PIPELINE_STAGE_DEFAULT = False - -######################################### -# Data types config params -######################################### -# "data_types": { -# grad_accum_dtype=["bf16"|"fp16"|"fp32"] -# } -# } - -DATA_TYPES = "data_types" -GRAD_ACCUM_DTYPE = "grad_accum_dtype" -GRAD_ACCUM_DTYPE_DEFAULT = None - -######################################### -# Drop the last incomplete Batch -# ######################################### -# dataloader_drop_last. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -DATALOADER_DROP_LAST_FORMAT = ''' -The last incomplete batch can be dropped by setting: -"dataloader_drop_last": True -''' -DATALOADER_DROP_LAST = "dataloader_drop_last" -DATALOADER_DROP_LAST_DEFAULT = False - -######################################### -# PIPELINE PARALLELISM -######################################### -PIPE_REPLICATED = 'ds_pipe_replicated' - -######################################### -# DATA PARALLELISM -######################################### -DATA_PARALLEL_GROUP = "data_parallel_group" -GLOBAL_RANK = "global_rank" - -######################################### -# EXPERT-DATA PARALLELISM TOPO Config -######################################### -USE_DATA_BEFORE_EXPERT_PARALLEL = "use_data_before_expert_parallelism" -USE_DATA_BEFORE_EXPERT_PARALLEL_DEFAULT = False diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 42d17520f9eadaf688bbf529515a763529f34331..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/config.cpython-310.pyc deleted file mode 100644 index 9c38643a6e000c21441cac2db44c1fe842112eec..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/constants.cpython-310.pyc deleted file mode 100644 index 7a430f1619fbc870ec0189e4b0c1cea0dd886338..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/constants.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/curriculum_scheduler.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/curriculum_scheduler.cpython-310.pyc deleted file mode 100644 index c7bfa2bfeff8e1416ac25670620e9bf05a320774..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/__pycache__/curriculum_scheduler.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/config.py deleted file mode 100644 index 690ce97034e452c2dbb451c3f71ebfaa579ad02d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/config.py +++ /dev/null @@ -1,193 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .constants import * -import copy -from ..config_utils import get_scalar_param - - -# TODO: Reducing config verbosity by returning None or {} when disabled. -# One challenge is that we still need to somehow include the default values, -# for example the *_ENABLED has default of false. -def get_data_efficiency_config(param_dict): - output = {} - output[DATA_EFFICIENCY_ENABLED] = get_data_efficiency_enabled(param_dict) - output[DATA_EFFICIENCY_SEED] = get_data_efficiency_seed(param_dict) - if DATA_EFFICIENCY not in param_dict.keys(): - param_dict[DATA_EFFICIENCY] = {} - sub_param_dict = param_dict[DATA_EFFICIENCY] - output[DATA_SAMPLING] = get_data_sampling(sub_param_dict) - output[DATA_ROUTING] = get_data_routing(sub_param_dict) - return output - - -def get_data_efficiency_enabled(param_dict): - if DATA_EFFICIENCY in param_dict.keys(): - return get_scalar_param(param_dict[DATA_EFFICIENCY], DATA_EFFICIENCY_ENABLED, DATA_EFFICIENCY_ENABLED_DEFAULT) - else: - return False - - -def get_data_efficiency_seed(param_dict): - if DATA_EFFICIENCY in param_dict.keys(): - return get_scalar_param(param_dict[DATA_EFFICIENCY], DATA_EFFICIENCY_SEED, DATA_EFFICIENCY_SEED_DEFAULT) - else: - return DATA_EFFICIENCY_SEED_DEFAULT - - -def get_data_sampling(param_dict): - sub_param_dict = param_dict.get(DATA_SAMPLING, {}) - output = copy.copy(sub_param_dict) - output[DATA_SAMPLING_ENABLED] = get_data_sampling_enabled(param_dict) - output[DATA_SAMPLING_NUM_EPOCHS] = get_data_sampling_num_epochs(param_dict) - output[DATA_SAMPLING_NUM_WORKERS] = get_data_sampling_num_workers(param_dict) - output[DATA_SAMPLING_PIN_MEMORY] = get_data_sampling_pin_memory(param_dict) - output[CURRICULUM_LEARNING] = get_curriculum_learning(sub_param_dict) - output[DYNAMIC_BATCHING] = get_dynamic_batching(sub_param_dict) - return output - - -def get_data_sampling_enabled(param_dict): - if DATA_SAMPLING in param_dict.keys(): - return get_scalar_param(param_dict[DATA_SAMPLING], DATA_SAMPLING_ENABLED, DATA_SAMPLING_ENABLED_DEFAULT) - else: - return False - - -def get_data_sampling_num_epochs(param_dict): - if DATA_SAMPLING in param_dict.keys(): - return get_scalar_param(param_dict[DATA_SAMPLING], DATA_SAMPLING_NUM_EPOCHS, DATA_SAMPLING_NUM_EPOCHS_DEFAULT) - else: - return DATA_SAMPLING_NUM_EPOCHS_DEFAULT - - -def get_data_sampling_num_workers(param_dict): - if DATA_SAMPLING in param_dict.keys(): - return get_scalar_param(param_dict[DATA_SAMPLING], DATA_SAMPLING_NUM_WORKERS, - DATA_SAMPLING_NUM_WORKERS_DEFAULT) - else: - return DATA_SAMPLING_NUM_WORKERS_DEFAULT - - -def get_data_sampling_pin_memory(param_dict): - if DATA_SAMPLING in param_dict.keys(): - return get_scalar_param(param_dict[DATA_SAMPLING], DATA_SAMPLING_PIN_MEMORY, DATA_SAMPLING_PIN_MEMORY_DEFAULT) - else: - return DATA_SAMPLING_PIN_MEMORY_DEFAULT - - -def get_curriculum_learning(param_dict): - output = {} - output[CURRICULUM_LEARNING_ENABLED] = get_curriculum_learning_enabled(param_dict) - if CURRICULUM_LEARNING not in param_dict.keys(): - param_dict[CURRICULUM_LEARNING] = {} - sub_param_dict = param_dict[CURRICULUM_LEARNING] - if output[CURRICULUM_LEARNING_ENABLED]: - assert CURRICULUM_LEARNING_METRICS in sub_param_dict.keys( - ), f"Curriculum learning is enabled, {CURRICULUM_LEARNING_METRICS} must be specified" - for key, val in get_curriculum_learning_params(param_dict).items(): - output[key] = val - return output - - -def get_dynamic_batching(param_dict): - output = copy.copy(param_dict.get(DYNAMIC_BATCHING, {})) - output[DYNAMIC_BATCHING_ENABLED] = bool(output.get(DYNAMIC_BATCHING_ENABLED, DYNAMIC_BATCHING_ENABLED_DEFAULT)) - output[DYNAMIC_BATCHING_LR_SCALING_METHOD] = str( - output.get(DYNAMIC_BATCHING_LR_SCALING_METHOD, DYNAMIC_BATCHING_LR_SCALING_METHOD_DEFAULT)) - output[DYNAMIC_BATCHING_MIN_BATCH_SIZE] = int( - output.get(DYNAMIC_BATCHING_MIN_BATCH_SIZE, DYNAMIC_BATCHING_MIN_BATCH_SIZE_DEFAULT)) - output[DYNAMIC_BATCHING_MAX_BATCH_SIZE] = int(output[DYNAMIC_BATCHING_MAX_BATCH_SIZE]) \ - if DYNAMIC_BATCHING_MAX_BATCH_SIZE in output.keys() \ - else DYNAMIC_BATCHING_MAX_BATCH_SIZE_DEFAULT - output[DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER] = str( - output.get(DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER, DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER_DEFAULT)) - if output[DYNAMIC_BATCHING_ENABLED]: - assert DYNAMIC_BATCHING_MAX_TOKENS in output.keys( - ), f"Dynamic batching is enabled, so {DYNAMIC_BATCHING_MAX_TOKENS} must be specified" - output[DYNAMIC_BATCHING_MAX_TOKENS] = int(output[DYNAMIC_BATCHING_MAX_TOKENS]) - output[DYNAMIC_BATCHING_VERBOSE] = bool(output.get(DYNAMIC_BATCHING_VERBOSE, False)) - return output - - -def get_curriculum_learning_enabled(param_dict): - if CURRICULUM_LEARNING in param_dict.keys(): - return get_scalar_param(param_dict[CURRICULUM_LEARNING], CURRICULUM_LEARNING_ENABLED, - CURRICULUM_LEARNING_ENABLED_DEFAULT) - else: - return False - - -def get_curriculum_learning_params(param_dict): - if CURRICULUM_LEARNING in param_dict.keys(): - curriculum_learning_params = copy.copy(param_dict[CURRICULUM_LEARNING]) - curriculum_learning_params.pop(CURRICULUM_LEARNING_ENABLED) - return curriculum_learning_params - else: - return {} - - -def get_curriculum_enabled_legacy(param_dict): - if CURRICULUM_LEARNING_LEGACY in param_dict.keys(): - return get_scalar_param(param_dict[CURRICULUM_LEARNING_LEGACY], CURRICULUM_ENABLED_LEGACY, - CURRICULUM_ENABLED_DEFAULT_LEGACY) - else: - return False - - -def get_curriculum_params_legacy(param_dict): - if CURRICULUM_LEARNING_LEGACY in param_dict.keys(): - curriculum_params = copy.copy(param_dict[CURRICULUM_LEARNING_LEGACY]) - curriculum_params.pop(CURRICULUM_ENABLED_LEGACY) - return curriculum_params - else: - return False - - -def get_data_routing(param_dict): - output = {} - output[DATA_ROUTING_ENABLED] = get_data_routing_enabled(param_dict) - if DATA_ROUTING not in param_dict.keys(): - param_dict[DATA_ROUTING] = {} - sub_param_dict = param_dict[DATA_ROUTING] - output[RANDOM_LTD] = get_random_ltd(sub_param_dict) - - return output - - -def get_data_routing_enabled(param_dict): - if DATA_ROUTING in param_dict.keys(): - return get_scalar_param(param_dict[DATA_ROUTING], DATA_ROUTING_ENABLED, DATA_ROUTING_ENABLED_DEFAULT) - else: - return False - - -def get_random_ltd(param_dict): - output = {} - output[RANDOM_LTD_ENABLED] = RANDOM_LTD_ENABLED_DEFAULT - output[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE] = {} - output[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE][ - RANDOM_LTD_LAYER_TOKEN_LR_ENABLED] = RANDOM_LTD_LAYER_TOKEN_LR_ENABLED_DEFAULT - if get_random_ltd_enabled(param_dict): - output[RANDOM_LTD_ENABLED] = get_random_ltd_enabled(param_dict) - for key, val in get_random_ltd_params(param_dict).items(): - output[key] = val - return output - - -def get_random_ltd_enabled(param_dict): - if RANDOM_LTD in param_dict.keys(): - return get_scalar_param(param_dict[RANDOM_LTD], RANDOM_LTD_ENABLED, RANDOM_LTD_ENABLED_DEFAULT) - else: - return False - - -def get_random_ltd_params(param_dict): - if RANDOM_LTD in param_dict.keys(): - random_ltd_params = copy.copy(param_dict[RANDOM_LTD]) - random_ltd_params.pop(RANDOM_LTD_ENABLED) - return random_ltd_params - else: - return {} diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/constants.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/constants.py deleted file mode 100644 index 73cc69c1f606a3c77e6969cc32c57ef7cd8e5111..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/constants.py +++ /dev/null @@ -1,136 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Data efficiency library - See sample config at https://www.deepspeed.ai/docs/config-json/data-efficiency -""" -DATA_EFFICIENCY = "data_efficiency" -DATA_EFFICIENCY_ENABLED = "enabled" -DATA_EFFICIENCY_ENABLED_DEFAULT = False -DATA_EFFICIENCY_SEED = "seed" -DATA_EFFICIENCY_SEED_DEFAULT = 1234 - -######################################### -# Data efficiency - Data Sampling -######################################### -DATA_SAMPLING = "data_sampling" -DATA_SAMPLING_ENABLED = "enabled" -DATA_SAMPLING_ENABLED_DEFAULT = False -DATA_SAMPLING_NUM_EPOCHS = "num_epochs" -DATA_SAMPLING_NUM_EPOCHS_DEFAULT = 1000 -DATA_SAMPLING_NUM_WORKERS = "num_workers" -DATA_SAMPLING_NUM_WORKERS_DEFAULT = 0 -DATA_SAMPLING_PIN_MEMORY = "pin_memory" -DATA_SAMPLING_PIN_MEMORY_DEFAULT = False - -######################################### -# Data efficiency - Data Sampling - Curriculum Learning -######################################### -CURRICULUM_LEARNING = "curriculum_learning" -CURRICULUM_LEARNING_ENABLED = "enabled" -CURRICULUM_LEARNING_ENABLED_DEFAULT = False -CURRICULUM_LEARNING_CLUSTER_PATH = "data_cluster_path" -CURRICULUM_LEARNING_METRICS = "curriculum_metrics" -CURRICULUM_LEARNING_SAMPLE_PATH = "index_to_sample_path" -CURRICULUM_LEARNING_METRIC_PATH = "index_to_metric_path" -CURRICULUM_LEARNING_CLUSTERING_TYPE = "clustering_type" -CURRICULUM_LEARNING_SINGLE_CLUSTER = "single_cluster" -CURRICULUM_LEARNING_CLUSTER_PREFIX = "cluster" -CURRICULUM_LEARNING_DIFFICULTY_TYPE = "difficulty_type" -CURRICULUM_LEARNING_VALUE_BASED = "value" -CURRICULUM_LEARNING_PERCENTILE_BASED = "percentile" -CURRICULUM_LEARNING_MIN_DIFFICULTY = "min_difficulty" -CURRICULUM_LEARNING_MAX_DIFFICULTY = "max_difficulty" -CURRICULUM_LEARNING_SCHEDULE_TYPE = "schedule_type" -CURRICULUM_LEARNING_SCHEDULE_CONFIG = "schedule_config" -CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY = "difficulty" -CURRICULUM_LEARNING_SCHEDULE_MAX_STEP = "max_step" -CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP = "total_curriculum_step" -CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP = "difficulty_step" -CURRICULUM_LEARNING_SCHEDULE_ROOT_DEGREE = "root_degree" -CURRICULUM_LEARNING_SCHEDULE_FIXED_DISCRETE = "fixed_discrete" -CURRICULUM_LEARNING_SCHEDULE_FIXED_ROOT = "fixed_root" -CURRICULUM_LEARNING_SCHEDULE_FIXED_LINEAR = "fixed_linear" -CURRICULUM_LEARNING_SCHEDULE_CUSTOM = "custom" -CURRICULUM_LEARNING_CURRENT_DIFFICULTY = "current_difficulty" - -CURRICULUM_LEARNING_BATCH = "batch" -CURRICULUM_LEARNING_CONSUMED_SAMPLES = "consumed_samples" -CURRICULUM_LEARNING_STEP = "curriculum_step" -CURRICULUM_LEARNING_CURRENT_DIFFICULTIES = "current_difficulties" -CURRICULUM_LEARNING_DATA_CLUSTER_PATHS = "data_cluster_paths" -CURRICULUM_LEARNING_DATA_CLUSTER_CURRENT_POSITION = "data_cluster_current_position" -CURRICULUM_LEARNING_NP_RNG_STATE = "np_rng_state" - -######################################### -# Data efficiency - Dynamic batching and LR scaling -######################################### -DYNAMIC_BATCHING = "dynamic_batching" -DYNAMIC_BATCHING_ENABLED = "enabled" -DYNAMIC_BATCHING_ENABLED_DEFAULT = False -DYNAMIC_BATCHING_METRICS_PATH = "metrics_path" -DYNAMIC_BATCHING_LR_SCALING_METHOD = "lr_scaling_method" # "linear" / "sqrt" / "none" -DYNAMIC_BATCHING_LR_SCALING_METHOD_DEFAULT = "linear" -DYNAMIC_BATCHING_MIN_BATCH_SIZE = "min_batch_size" -DYNAMIC_BATCHING_MIN_BATCH_SIZE_DEFAULT = 1 -DYNAMIC_BATCHING_MAX_BATCH_SIZE = "max_batch_size" -DYNAMIC_BATCHING_MAX_BATCH_SIZE_DEFAULT = None -DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER = "sequence_picking_order" # "random" / "seqlen" / "dataloader" -DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER_DEFAULT = "dataloader" # "random" / "seqlen" / "dataloader" -DYNAMIC_BATCHING_MAX_TOKENS = "max_tokens" -DYNAMIC_BATCHING_VERBOSE = "verbose" - -######################################### -# Curriculum Learning legacy implementation -######################################### -CURRICULUM_LEARNING_LEGACY = "curriculum_learning" - -CURRICULUM_ENABLED_LEGACY = "enabled" -CURRICULUM_ENABLED_DEFAULT_LEGACY = False - -######################################### -# Data efficiency - Data Routing -######################################### -DATA_ROUTING = "data_routing" -DATA_ROUTING_ENABLED = "enabled" -DATA_ROUTING_ENABLED_DEFAULT = False - -######################################### -# Data efficiency - Data Routing - Random LTD -######################################### -RANDOM_LTD = "random_ltd" -RANDOM_LTD_ENABLED = "enabled" -RANDOM_LTD_ENABLED_DEFAULT = False - -RANDOM_LTD_MODEL_MASK_NAME = "model_mask_name" -RANDOM_LTD_MODEL_TYPE = "model_type" -RANDOM_LTD_MICRO_BATCH_SIZE = "micro_batch_size" -RANDOM_LTD_GLOBAL_BATCH_SIZE = "global_batch_size" -RANDOM_LTD_SAMPLE_INDEX = "sample_idx" -RANDOM_LTD_ATTENTION_MASK = "attention_mask" -RANDOM_LTD_HIDDEN_STATE_ORDER = "hidden_state_order" -RANDOM_LTD_LAYER_NUM = "random_ltd_layer_num" -RANDOM_LTD_LAYER_ID = "random_ltd_layer_id" -RANDOM_LTD_TOTAL_LAYER_NUM = "total_layer_num" -RANDOM_LTD_CONSUMED_LAYER_TOKENS = "consumed_layer_tokens" - -# scheduler -RANDOM_LTD_SCHEDULER = "random_ltd_schedule" -RANDOM_LTD_MAX_VALUE = "max_value" -RANDOM_LTD_MIN_VALUE = "min_value" -RANDOM_LTD_CURRENT_VALUE = "current_value" -RANDOM_LTD_SCHEDULE_CONFIG = "schedule_config" -RANDOM_LTD_INCREASE_STEP = "seq_per_step" -RANDOM_LTD_REQUIRE_STEP = "require_steps" -RANDOM_LTD_SCHEDULER_TYPE = "schedule_type" -RANDOM_LTD_CURR_STEP = "current_steps" - -# learning rate schedulers -RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE = "layer_token_lr_schedule" -RANDOM_LTD_LAYER_TOKEN_LR_ENABLED = "enabled" -RANDOM_LTD_LAYER_TOKEN_LR_ENABLED_DEFAULT = False -RANDOM_LTD_TOTAL_LAYER_TOKENS = "total_layer_tokens" -RANDOM_LTD_WARMUP_TYPE = "warmup_type" -RANDOM_LTD_WARMUP_LAYER_TOKENS = "warmup_layer_tokens" diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/curriculum_scheduler.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/curriculum_scheduler.py deleted file mode 100644 index 23d747957dc4647e06fad0a94e5e4b071b6f6e23..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/curriculum_scheduler.py +++ /dev/null @@ -1,158 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import math -from deepspeed.utils import logger -from .constants import * - - -class CurriculumScheduler(object): - - def __init__(self, config): - super().__init__() - self.state = {} - assert CURRICULUM_LEARNING_MIN_DIFFICULTY in config, \ - f"Curriculum learning requires the config '{CURRICULUM_LEARNING_MIN_DIFFICULTY}'" - assert CURRICULUM_LEARNING_MAX_DIFFICULTY in config, \ - f"Curriculum learning requires the config '{CURRICULUM_LEARNING_MAX_DIFFICULTY}'" - assert CURRICULUM_LEARNING_SCHEDULE_TYPE in config, \ - f"Curriculum learning requires the config '{CURRICULUM_LEARNING_SCHEDULE_TYPE}'" - self.state[CURRICULUM_LEARNING_MIN_DIFFICULTY] = config[CURRICULUM_LEARNING_MIN_DIFFICULTY] - self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY] = config[CURRICULUM_LEARNING_MAX_DIFFICULTY] - self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] = config[CURRICULUM_LEARNING_MIN_DIFFICULTY] - self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] = config[CURRICULUM_LEARNING_SCHEDULE_TYPE] - self.first_step = True - if config[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_DISCRETE: - """ - The schedule_config is a list of difficulty and a list of max - step belonging to each difficulty. Example json config: - "schedule_config": { - "difficulty": [1,2,3], - "max_step": [5,10] - } - The "max_step" has one less element than "difficulty", because - the last difficulty will be used for all following steps. - The self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] is a dictionary of - difficulty : [max step for this difficulty, next difficulty]. - """ - assert CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ - f"Curriculum learning with fixed_discrete schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY}'" - assert CURRICULUM_LEARNING_SCHEDULE_MAX_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ - f"Curriculum learning with fixed_discrete schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_MAX_STEP}'" - assert len(config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_MAX_STEP]) > 0 - assert len(config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY]) > 0 - assert len(config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY]) == len( - config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_MAX_STEP]) + 1 - self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] = config[CURRICULUM_LEARNING_SCHEDULE_CONFIG] - elif config[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_ROOT: - """ - The schedule_config includes: - total_curriculum_step: how many steps the curriculum learning takes to go - from min difficulty to max difficulty. - difficulty_step: the difficulty level determined every time must - be a multiple of this difficulty_step. This is used to determine - the step of difficulty increase, and to ensure the use of NVIDIA - Tensor Core acceleration (requires multiple of 8 (FP16) or - 16 (INT8)). - root_degree: the degree of the root function. Degree of 2 means - square root and degree of 3 means cube root. Degree of 1 is - equivalent to linear. - "schedule_config": { - "total_curriculum_step": 30000, - "difficulty_step": 8, - "root_degree": 2 - } - """ - assert CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ - f"Curriculum learning with fixed_root schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP}'" - assert CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ - f"Curriculum learning with fixed_root schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP}'" - assert CURRICULUM_LEARNING_SCHEDULE_ROOT_DEGREE in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ - f"Curriculum learning with fixed_root schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_ROOT_DEGREE}'" - if config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP] % 8 != 0: - logger.warning( - f'When using seqlen metric, the difficulty_step for curriculum learning has to be multiple of 8 (for FP16 data) or 16 (for INT8 data) to enable NVIDIA Tensor Core acceleration. Disregard this warning if this is unrelated to your metric/hardware.' - ) - self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] = config[CURRICULUM_LEARNING_SCHEDULE_CONFIG] - elif config[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_LINEAR: - """ - The schedule_config is the same as CURRICULUM_LEARNING_SCHEDULE_FIXED_ROOT but without the - root_degree. - "schedule_config": { - "total_curriculum_step": 30000, - "difficulty_step": 8 - } - """ - assert CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ - f"Curriculum learning with fixed_linear schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP}'" - assert CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ - f"Curriculum learning with fixed_linear schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP}'" - if config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP] % 8 != 0: - logger.warning( - f'When using seqlen metric, the difficulty_step for curriculum learning has to be multiple of 8 (for FP16 data) or 16 (for INT8 data) to enable NVIDIA Tensor Core acceleration. Disregard this warning if this is unrelated to your metric/hardware.' - ) - self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] = config[CURRICULUM_LEARNING_SCHEDULE_CONFIG] - elif config[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_CUSTOM: - """ - Fully customized schedule. User need to provide a custom schedule - function by using the set_custom_curriculum_learning_schedule API - in deepspeed/runtime/engine.py - """ - self.custom_get_difficulty = None - else: - raise RuntimeError('Unsupported curriculum schedule type') - - def get_current_difficulty(self): - return self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] - - def set_current_difficulty(self, difficulty): - self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] = difficulty - - def set_custom_get_difficulty(self, schedule_function): - self.custom_get_difficulty = schedule_function - - def get_state(self): - return self.state - - def set_state(self, state): - self.state = state - - def __fixed_discrete_get_difficulty(self, global_steps): - s_state = self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] - if global_steps > s_state[CURRICULUM_LEARNING_SCHEDULE_MAX_STEP][-1]: - return s_state[CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY][-1] - for i in range(len(s_state[CURRICULUM_LEARNING_SCHEDULE_MAX_STEP])): - if global_steps <= s_state[CURRICULUM_LEARNING_SCHEDULE_MAX_STEP][i]: - return s_state[CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY][i] - - def __fixed_root_get_difficulty(self, global_steps, root_degree=None): - s_state = self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] - if root_degree is None: - root_degree = s_state[CURRICULUM_LEARNING_SCHEDULE_ROOT_DEGREE] - next_difficulty = (float(global_steps) / s_state[CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP])**(1.0 / root_degree) - next_difficulty = math.floor( - next_difficulty * - (self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY] - self.state[CURRICULUM_LEARNING_MIN_DIFFICULTY]) + - self.state[CURRICULUM_LEARNING_MIN_DIFFICULTY]) - next_difficulty -= (next_difficulty % s_state[CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP]) - next_difficulty = min(next_difficulty, self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY]) - return next_difficulty - - def get_difficulty(self, global_steps): - if self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_DISCRETE: - return self.__fixed_discrete_get_difficulty(global_steps) - elif self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_LINEAR: - return self.__fixed_root_get_difficulty(global_steps, 1) - elif self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_ROOT: - return self.__fixed_root_get_difficulty(global_steps) - elif self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_CUSTOM: - return self.custom_get_difficulty(global_steps) - else: - raise RuntimeError('Unsupported curriculum schedule type') - - def update_difficulty(self, global_steps): - if self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] < self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY]: - self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] = self.get_difficulty(global_steps) - return self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 7b9c1c5c49c7dc7d59762c4c08f4b5c4171e75d9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/basic_layer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/basic_layer.cpython-310.pyc deleted file mode 100644 index 164053097aff638211a00394303d9a447dc5f616..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/basic_layer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/helper.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/helper.cpython-310.pyc deleted file mode 100644 index 3d728dbbee2d39cd0ccb30aba76a2b4bfe8136cb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/helper.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/scheduler.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/scheduler.cpython-310.pyc deleted file mode 100644 index 47555803744b1c39fa8580eef887f7070fc68374..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/scheduler.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index 0f64c99133a047df7c11c034240b480a250ee7d4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/basic_layer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/basic_layer.py deleted file mode 100644 index 780a88c3d5a2fb016f2c59bbe3a88fbaff6499b0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/basic_layer.py +++ /dev/null @@ -1,113 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.utils import logger -from torch import Tensor -from torch.nn import Module -from ..constants import * -from deepspeed.ops.random_ltd.dropping_utils import gpt_sample_tokens, bert_sample_tokens, GatherTokens, ScatterTokens - - -#####based on the paper random-ltd: https://arxiv.org/abs/2211.11586 -class RandomLayerTokenDrop(Module): - """ - A layer wrapper for random LTD - """ - - def __init__(self, layer: Module): - super(RandomLayerTokenDrop, self).__init__() - self.random_ltd_layer = layer - self.reserved_length = None #config['max_value'] - self.random_ltd_scheduler = None - self.max_length = None - self.reserved_length = -1 - self.curr_seq = -1 - self.batch_first = False - - def init_config(self, config, scheduler, random_ltd_layer_id): - self.random_ltd_scheduler = scheduler - self.random_ltd_layer_id = random_ltd_layer_id - self.max_length = self.random_ltd_scheduler.state[RANDOM_LTD_MAX_VALUE] - - self.mask_name = config[RANDOM_LTD_MODEL_MASK_NAME] - self.micro_bs = config[RANDOM_LTD_MICRO_BATCH_SIZE] - self.random_ltd_num_layer = self.random_ltd_scheduler.random_ltd_layer_num - hs_order = config[RANDOM_LTD_HIDDEN_STATE_ORDER] - self.model_type = config[RANDOM_LTD_MODEL_TYPE] - - if hs_order == 'batch_seq_dim': - self.get_hidden_tensor_shape = self.get_bsh - self.batch_first = True - elif hs_order == 'seq_batch_dim': - self.get_hidden_tensor_shape = self.get_sbh - self.batch_first = False - else: - logger.warning( - "************For now, we only support batch_seq_dim or seq_batch_dim inputs. You can easily \ - your own input dimension orders************") - raise NotImplementedError - - if self.model_type == 'encoder': - self.index_generator = bert_sample_tokens - elif self.model_type == 'decoder': - self.index_generator = gpt_sample_tokens - else: - logger.warning("************For now, we only support encoder-only or decoder-only models************") - raise NotImplementedError - - def get_bsh(self, hidden_stats): - self.curr_seq, self.curr_micro_batch = hidden_stats.size()[1], hidden_stats.size()[0] - - def get_sbh(self, hidden_stats): - self.curr_seq, self.curr_micro_batch = hidden_stats.size()[0], hidden_stats.size()[1] - - def forward(self, hidden_states, **kwargs) -> Tensor: - if self.random_ltd_scheduler is not None: - self.reserved_length = self.random_ltd_scheduler.get_current_seq() - self.get_hidden_tensor_shape(hidden_states) - if self.training and self.random_ltd_scheduler is not None and self.reserved_length < self.curr_seq: - if self.mask_name is not None: - mask = kwargs[self.mask_name] - else: - mask = None - if self.random_ltd_layer_id == 0: - sampled_indices, part_attention_mask = self.index_generator(self.reserved_length,\ - self.curr_seq, \ - self.curr_micro_batch, \ - self.random_ltd_num_layer, \ - hidden_states.device, mask) - self.random_ltd_scheduler.state[RANDOM_LTD_SAMPLE_INDEX] = sampled_indices - self.random_ltd_scheduler.state[RANDOM_LTD_ATTENTION_MASK] = part_attention_mask - else: - sampled_indices = self.random_ltd_scheduler.state[RANDOM_LTD_SAMPLE_INDEX] - part_attention_mask = self.random_ltd_scheduler.state[RANDOM_LTD_ATTENTION_MASK] - - hidden_states, part_hidden_states = GatherTokens.apply(hidden_states, - sampled_indices[self.random_ltd_layer_id, :, :], - self.batch_first) - if self.mask_name is not None: - if self.model_type == 'encoder': - kwargs[self.mask_name] = part_attention_mask[self.random_ltd_layer_id] - else: - kwargs[self.mask_name] = part_attention_mask - - outputs = self.random_ltd_layer(part_hidden_states, **kwargs) - - if isinstance(outputs, tuple): - hidden_states = ScatterTokens.apply(hidden_states, outputs[0], - sampled_indices[self.random_ltd_layer_id, :, :], self.batch_first) - my_list = list(outputs) - my_list[0] = hidden_states - return tuple(my_list) - elif isinstance(outputs, Tensor): - hidden_states = ScatterTokens.apply(hidden_states, outputs, - sampled_indices[self.random_ltd_layer_id, :, :], self.batch_first) - return hidden_states - else: - logger.warning("************For now, we only support tuple and tensor output. \ - You need to adjust the output according to the layer in your model************") - raise NotImplementedError - else: - return self.random_ltd_layer(hidden_states, **kwargs) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/helper.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/helper.py deleted file mode 100644 index 150182d77bcfda20b1aa1aabd4b8785542ca9d1b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/helper.py +++ /dev/null @@ -1,46 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .basic_layer import RandomLayerTokenDrop -from collections import OrderedDict -from deepspeed.compression.helper import recursive_getattr, recursive_setattr - - -def convert_to_random_ltd(model, convert_type): - if hasattr(model, 'module'): - c_model = model.module - else: - c_model = model - - for name, module in c_model.named_modules(): - - if isinstance(module, convert_type): - old_module = recursive_getattr(c_model, name) - new_module = RandomLayerTokenDrop(old_module) - recursive_setattr(c_model, name, new_module) - - model.random_ltd_initialize() - return model - - -def save_without_random_ltd(model): - if hasattr(model, 'module'): - c_model = model.module - else: - c_model = model - - model_dic = c_model.state_dict() - return remove_random_ltd_state_dict(model_dic) - - -def remove_random_ltd_state_dict(state_dict): - new_state_dict = OrderedDict() - for key, value in state_dict.items(): - if '.random_ltd_layer' in key: - new_key = ''.join(key.split('.random_ltd_layer')) - else: - new_key = key - new_state_dict[new_key] = value - return new_state_dict diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/scheduler.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/scheduler.py deleted file mode 100644 index 8a0b031d4f633976f438d5151973e58afa77712e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/scheduler.py +++ /dev/null @@ -1,107 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import math - -from deepspeed.utils import logger -# from deepspeed.runtime.lr_schedules import WarmupLR -from ..constants import * - -#####based on the paper random-ltd: https://arxiv.org/abs/2211.11586 - - -class BaseScheduler(object): - - def __init__(self): - self.state = {} - - def __fixed_root_get_value(self, global_steps, root_degree=None): - s_state = self.state[RANDOM_LTD_SCHEDULE_CONFIG] - if root_degree is None: - root_degree = s_state['root_degree'] - next_seq = (float(global_steps) / s_state[RANDOM_LTD_REQUIRE_STEP])**(1.0 / root_degree) - next_seq = math.floor(next_seq * (self.state[RANDOM_LTD_MAX_VALUE] - self.state[RANDOM_LTD_MIN_VALUE]) + - self.state[RANDOM_LTD_MIN_VALUE]) - next_seq -= (next_seq % s_state[RANDOM_LTD_INCREASE_STEP]) - next_seq = min(next_seq, self.state[RANDOM_LTD_MAX_VALUE]) - return next_seq - - def get_value(self, global_steps): - if self.state[RANDOM_LTD_SCHEDULER_TYPE] == 'fixed_linear': - return self.__fixed_root_get_value(global_steps, 1) - else: - raise RuntimeError('Unsupported random LTD schedule type') - - -class RandomLTDScheduler(BaseScheduler): - - def __init__(self, config): - super().__init__() - self.model_layer_num = config[RANDOM_LTD_TOTAL_LAYER_NUM] - self.random_ltd_layer_num = config[RANDOM_LTD_LAYER_NUM] - self.config_schedule = config[RANDOM_LTD_SCHEDULER] - self.global_batch_size = config[RANDOM_LTD_GLOBAL_BATCH_SIZE] - self.reset_to_init() - - if config[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE][RANDOM_LTD_LAYER_TOKEN_LR_ENABLED]: - logger.warning("**********Work In Progress************") - raise NotImplementedError - - self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] = 0 - - # self.first_step = True - def get_total_layer_tokens(self, train_iters): - for step in range(train_iters): - self.update_seq(step) - return self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] - - def reset_to_init(self): - if self.config_schedule is not None: - self.state[RANDOM_LTD_MIN_VALUE] = self.config_schedule[RANDOM_LTD_MIN_VALUE] - self.state[RANDOM_LTD_MAX_VALUE] = self.config_schedule[RANDOM_LTD_MAX_VALUE] - self.state[RANDOM_LTD_CURRENT_VALUE] = self.config_schedule[RANDOM_LTD_MIN_VALUE] - self.state[RANDOM_LTD_SCHEDULE_CONFIG] = self.config_schedule[RANDOM_LTD_SCHEDULE_CONFIG] - self.state[RANDOM_LTD_SCHEDULER_TYPE] = self.config_schedule[RANDOM_LTD_SCHEDULER_TYPE] - self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] = 0 - self.state[RANDOM_LTD_CURR_STEP] = -1 - - def get_current_seq(self): - return self.state[RANDOM_LTD_CURRENT_VALUE] - - def set_current_seq(self, seq_length): - self.state[RANDOM_LTD_CURRENT_VALUE] = seq_length - - def get_random_ltd_layer_num(self): - return self.random_ltd_layer_num - - def get_state(self): - return self.state - - def set_state(self, state): - self.state = state - - def update_seq(self, global_steps): - if self.state[RANDOM_LTD_CURRENT_VALUE] < self.state[RANDOM_LTD_MAX_VALUE]: - self.state[RANDOM_LTD_CURRENT_VALUE] = self.get_value(global_steps) - if global_steps != self.state[RANDOM_LTD_CURR_STEP]: - self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] += self.global_batch_size*(self.state[RANDOM_LTD_CURRENT_VALUE] * self.random_ltd_layer_num \ - + self.state[RANDOM_LTD_MAX_VALUE] * (self.model_layer_num - self.random_ltd_layer_num)) - self.state[RANDOM_LTD_CURR_STEP] = global_steps - - def state_dict(self): - return { - RANDOM_LTD_CONSUMED_LAYER_TOKENS: self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS], - RANDOM_LTD_CURR_STEP: self.state[RANDOM_LTD_CURR_STEP], - RANDOM_LTD_CURRENT_VALUE: self.state[RANDOM_LTD_CURRENT_VALUE], - RANDOM_LTD_MIN_VALUE: self.state[RANDOM_LTD_MIN_VALUE], - RANDOM_LTD_MAX_VALUE: self.state[RANDOM_LTD_MAX_VALUE], - } - - def load_state_dict(self, state_dict): - self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] = state_dict[RANDOM_LTD_CONSUMED_LAYER_TOKENS] - self.state[RANDOM_LTD_CURR_STEP] = state_dict[RANDOM_LTD_CURR_STEP] - self.state[RANDOM_LTD_CURRENT_VALUE] = state_dict[RANDOM_LTD_CURRENT_VALUE] - self.state[RANDOM_LTD_MIN_VALUE] = state_dict[RANDOM_LTD_MIN_VALUE] - self.state[RANDOM_LTD_MAX_VALUE] = state_dict[RANDOM_LTD_MAX_VALUE] diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/utils.py deleted file mode 100644 index afcfef2ef4dc04279b7fa3ee49916519f33ffbec..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_routing/utils.py +++ /dev/null @@ -1,27 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - - -def bsh_decoder_gather(reserved_length, hidden_states, mask): - # random-layer-token-drop - rand_list = [] - part_hidden_states = [] # batch, seq, hidden ## different from megatron - for k in range(hidden_states.size(0)): - B_tmp = torch.randperm(hidden_states.size(1), device=hidden_states.device)[:reserved_length] - B = B_tmp.sort()[0] - rand_list.append(B) - part_hidden_states.append(hidden_states[k:k + 1, B, :]) - - part_hidden_states = torch.cat(part_hidden_states, dim=0) - part_mask = mask[:, :, :reserved_length, :reserved_length] - return part_hidden_states, rand_list, part_mask - - -def bsh_decoder_scatter(hidden_states, part_hidden_states, rand_list): - for k in range(hidden_states.size(0)): - hidden_states[k, rand_list[k], :] = part_hidden_states[k, :, :] - return hidden_states diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index a4047518331986d5b489904158578900a000246a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_analyzer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_analyzer.cpython-310.pyc deleted file mode 100644 index 2ee54cd1ec6962063e1ec303c0149e1d62c8d4d9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_analyzer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_sampler.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_sampler.cpython-310.pyc deleted file mode 100644 index 50356307a566ce38fa7dec96be530f5af22dbca5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_sampler.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/indexed_dataset.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/indexed_dataset.cpython-310.pyc deleted file mode 100644 index f8e04aa78bd52f3a9e6fc1a78100fe14533b1aa4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/indexed_dataset.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index aa34ee0ebb20b74193ea439f699b773fbd322066..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/variable_batch_size_and_lr.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/variable_batch_size_and_lr.cpython-310.pyc deleted file mode 100644 index bd86079e7a4a286ba58496256000d75d4a1af2e2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/variable_batch_size_and_lr.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py deleted file mode 100644 index f82c684ec6e229597736831d7ea56ca8c3bdbaa0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py +++ /dev/null @@ -1,885 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import sys -from collections import defaultdict -import csv -import time -from multiprocessing import Process, Manager -import numpy as np -import torch -from torch.utils.data import BatchSampler, SequentialSampler, DataLoader, Subset - -import deepspeed.comm as dist -from deepspeed.utils import logger -from deepspeed.runtime.data_pipeline.data_sampling.indexed_dataset import MMapIndexedDataset, valid_dtypes -from deepspeed.runtime.data_pipeline.data_sampling.utils import split_dataset, split_index, create_mmap_dataset_builder, close_mmap_dataset_builder, find_fit_int_dtype - - -class DataAnalyzer(object): - - def __init__(self, - dataset, - num_workers=1, - worker_id=0, - num_threads=1, - num_threads_reduce=1, - specific_threads=[], - batch_size=1, - metric_names=[], - metric_functions=[], - metric_types=[], - metric_dtypes=[], - save_path="./", - collate_fn=None, - custom_map_init=None, - custom_map_update=None, - custom_map_finalize=None, - custom_reduce=None, - sample_indices=None): - super().__init__() - self.dataset = dataset - self.num_workers = num_workers - self.worker_id = worker_id - self.num_threads = num_threads - self.num_threads_reduce = num_threads_reduce - self.specific_threads = specific_threads - self.batch_size = batch_size - self.metric_names = metric_names - self.metric_functions = metric_functions - self.metric_types = metric_types - self.metric_dtypes = metric_dtypes - self.save_path = save_path - self.collate_fn = collate_fn - self.custom_map_init = custom_map_init - self.custom_map_update = custom_map_update - self.custom_map_finalize = custom_map_finalize - self.custom_reduce = custom_reduce - self.sample_indices = sample_indices - - def init_metric_results(self, thread_id, metric_names, metric_types, metric_dtypes, save_path, worker_id): - metric_results = [] - for m_idx in range(len(metric_names)): - metric_name, metric_type, metric_dtype = metric_names[m_idx], \ - metric_types[m_idx], metric_dtypes[m_idx] - assert metric_dtype in valid_dtypes, f"metric_dtype {metric_dtype} not supported. Supported dtypes {valid_dtypes}" - metric_save_path = f"{save_path}/{metric_name}/worker{worker_id}_thread{thread_id}/" - os.makedirs(metric_save_path, exist_ok=True) - if metric_type == 'single_value_per_sample': - sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric" - sample_to_metric_builder = create_mmap_dataset_builder(sample_to_metric_fname, metric_dtype) - metric_to_sample_fname = f"{metric_save_path}/{metric_name}_metric_to_sample" - os.system(f"rm -rf {metric_to_sample_fname}*") - metric_to_sample_dict = defaultdict(list) - metric_results.append({ - "sample_to_metric_fname": sample_to_metric_fname, - "sample_to_metric_builder": sample_to_metric_builder, - "metric_to_sample_fname": metric_to_sample_fname, - "metric_to_sample_dict": metric_to_sample_dict - }) - elif metric_type == 'accumulate_value_over_samples': - metric_value = None - metric_value_fname = f"{metric_save_path}/{metric_name}_metric_value" - metric_results.append({"metric_value": metric_value, "metric_value_fname": metric_value_fname}) - return metric_results - - def update_metric_results(self, - data, - metric_types, - metric_dtypes, - metric_functions, - metric_results, - batch_start_idx=0): - for m_idx in range(len(metric_types)): - metric_type, metric_dtype, metric_function, metric_result = metric_types[m_idx], \ - metric_dtypes[m_idx], metric_functions[m_idx], metric_results[m_idx] - metric_values = metric_function(data) - - assert torch.is_tensor(metric_values) or isinstance(metric_values, np.ndarray), \ - "metric_function must return a tensor or array" - assert metric_values.dtype == metric_dtype, \ - f"metric_function result dtype {metric_values.dtype} does not match metric_dtype {metric_dtype}" - if isinstance(metric_values, np.ndarray): - metric_values = torch.from_numpy(metric_values) - - if metric_type == 'single_value_per_sample': - for row in range(metric_values.size()[0]): - sample_idx = batch_start_idx + row # sample idx following dataset iteration order - if isinstance(data, dict) and 'index' in data: # Megatron use case, idx provided in 'index' field - sample_idx = data['index'][row][0].item() - elif self.sample_indices is not None: # user defined shuffling of indices - sample_idx = self.sample_indices[sample_idx] - metric_result["sample_to_metric_builder"].add_item(metric_values[row].reshape(-1)) - metric_result["metric_to_sample_dict"][metric_values[row].item()].append(sample_idx) - for m_value in metric_result["metric_to_sample_dict"]: - if len(metric_result["metric_to_sample_dict"][m_value]) > 100: - metric_fname = metric_result["metric_to_sample_fname"] - with open(f"{metric_fname}_{m_value}.csv", 'a') as f: - writer = csv.writer(f) - writer.writerows([metric_result["metric_to_sample_dict"][m_value]]) - metric_result["metric_to_sample_dict"][m_value] = [] - elif metric_type == 'accumulate_value_over_samples': - if metric_result["metric_value"] is None: - metric_result["metric_value"] = metric_values - else: - metric_result["metric_value"].add_(metric_values) - - def finalize_metric_results(self, metric_types, metric_dtypes, metric_results): - for m_idx in range(len(metric_types)): - metric_type, metric_dtype, metric_result = metric_types[m_idx], \ - metric_dtypes[m_idx], metric_results[m_idx] - if metric_type == 'single_value_per_sample': - metric_fname = metric_result["sample_to_metric_fname"] - close_mmap_dataset_builder(metric_result["sample_to_metric_builder"], metric_fname) - for m_value in metric_result["metric_to_sample_dict"]: - if len(metric_result["metric_to_sample_dict"][m_value]) > 0: - metric_fname = metric_result["metric_to_sample_fname"] - with open(f"{metric_fname}_{m_value}.csv", 'a') as f: - writer = csv.writer(f) - writer.writerows([metric_result["metric_to_sample_dict"][m_value]]) - metric_result["metric_to_sample_dict"][m_value] = [] - elif metric_type == 'accumulate_value_over_samples': - if metric_result["metric_value"] is not None: - metric_value_builder = create_mmap_dataset_builder(metric_result["metric_value_fname"], - metric_dtype) - metric_value_builder.add_item(metric_result["metric_value"].reshape(-1)) - close_mmap_dataset_builder(metric_value_builder, metric_result["metric_value_fname"]) - - def run_map_helper(self, thread_id): - start_idx, end_idx = self.thread_splits[thread_id][0], \ - self.thread_splits[thread_id][1] - logger.info(f"worker {self.worker_id} thread {thread_id}: start working " \ - f"on data subset {start_idx} to {end_idx}") - thread_dataset = Subset(self.dataset, list(range(start_idx, end_idx))) - sampler = BatchSampler(SequentialSampler(thread_dataset), batch_size=self.batch_size, drop_last=False) - iterator = iter( - DataLoader(thread_dataset, - batch_sampler=sampler, - num_workers=0, - collate_fn=self.collate_fn, - pin_memory=False)) - if self.custom_map_init is None: - metric_results = self.init_metric_results(thread_id, self.metric_names, self.metric_types, - self.metric_dtypes, self.save_path, self.worker_id) - else: - metric_results = self.custom_map_init(thread_id, self.metric_names, self.metric_types, self.metric_dtypes, - self.save_path, self.worker_id) - total_sample = len(thread_dataset) - processed_sample = 0 - start = time.time() - while True: - try: - data = next(iterator) - batch_start_idx = start_idx + processed_sample - if self.custom_map_update is None: - self.update_metric_results(data, self.metric_types, self.metric_dtypes, self.metric_functions, - metric_results, batch_start_idx) - else: - self.custom_map_update(data, self.metric_types, self.metric_dtypes, self.metric_functions, - metric_results, batch_start_idx) - processed_sample += len(data) - duration = (time.time() - start) / 3600.0 - remain_duration = duration * total_sample / processed_sample - duration - logger.info( - f"worker {self.worker_id} thread {thread_id}: {processed_sample} " \ - f"out of {total_sample} processed in {duration:.2f} hr, " \ - f"estimated to finish in {remain_duration:.2f} hr") - except StopIteration: - logger.info(f"worker {self.worker_id} thread {thread_id}: reach end of file") - break - if self.custom_map_finalize is None: - self.finalize_metric_results(self.metric_types, self.metric_dtypes, metric_results) - else: - self.custom_map_finalize(self.metric_types, self.metric_dtypes, metric_results) - logger.info(f"worker {self.worker_id} thread {thread_id}: finished") - - def run_map(self): - self.worker_splits, self.thread_splits = split_dataset(self.dataset, self.num_workers, self.worker_id, - self.num_threads) - if len(self.specific_threads) > 0: - threads_to_run = self.specific_threads - else: - threads_to_run = list(range(self.num_threads)) - if self.num_threads > 1: - p = [] - for thread in threads_to_run: - p.append(Process(target=self.run_map_helper, args=(thread, ))) - p[thread].start() - - for thread in threads_to_run: - p[thread].join() - else: - assert self.num_threads == 1 - self.run_map_helper(0) - - def get_metric_value_percentiles(self, metric_name, num_sample_per_value, total_num_samples): - logger.info(f"Checking the value percentiles of metric {metric_name}...") - processed_samples = 0 - current_percentile = 5 - for key in sorted(num_sample_per_value.keys()): - processed_samples += num_sample_per_value[key] - if processed_samples >= total_num_samples * current_percentile / 100.0: - logger.info(f"Metric {metric_name} {current_percentile}th percentile: {key}") - current_percentile += 5 - - def merge_gather_map_stats(self, num_workers, num_threads, num_threads_reduce, t_idx_reduce, metric_save_path, - metric_name, return_dict): - results = [] - for w_idx in range(num_workers): - for t_idx in range(num_threads): - if (w_idx * num_threads + t_idx) % num_threads_reduce == t_idx_reduce: - w_metric_save_path = f"{metric_save_path}/worker{w_idx}_thread{t_idx}/" - w_sample_to_metric_fname = f"{w_metric_save_path}/{metric_name}_sample_to_metric" - w_sample_to_metric = MMapIndexedDataset(w_sample_to_metric_fname, skip_warmup=True) - unique_v = list(np.unique(w_sample_to_metric)) - sample_to_metric_count = len(w_sample_to_metric) - logger.info(f"Finished gathering map stats from worker {w_idx} thread {t_idx}.") - results.append([unique_v, sample_to_metric_count]) - return_dict[t_idx_reduce] = results - - def merge_sample_to_metric(self, t_idx_reduce, metric_save_path, metric_name, metric_value_dtype, - map_worker_thread): - sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric_thread{t_idx_reduce}" - sample_to_metric_builder = create_mmap_dataset_builder(sample_to_metric_fname, metric_value_dtype) - for w_t in map_worker_thread: - w_metric_save_path = f"{metric_save_path}/worker{w_t[0]}_thread{w_t[1]}/" - w_sample_to_metric_fname = f"{w_metric_save_path}/{metric_name}_sample_to_metric" - w_data = MMapIndexedDataset(w_sample_to_metric_fname, skip_warmup=True) - for row in range(len(w_data)): - sample_to_metric_builder.add_item(torch.tensor(w_data[row].astype(np.int64), dtype=torch.long)) - logger.info(f"Finished merge_sample_to_metric from worker {w_t[0]} thread {w_t[1]}.") - close_mmap_dataset_builder(sample_to_metric_builder, sample_to_metric_fname) - - def merge_metric_to_sample(self, t_idx_reduce, metric_save_path, metric_name, sample_idx_dtype, metric_value_dtype, - unique_metric_values, num_workers, num_threads): - index_to_sample_fname = f"{metric_save_path}/{metric_name}_index_to_sample_thread{t_idx_reduce}" - index_to_sample_builder = create_mmap_dataset_builder(index_to_sample_fname, sample_idx_dtype) - index_to_metric_fname = f"{metric_save_path}/{metric_name}_index_to_metric_thread{t_idx_reduce}" - index_to_metric_builder = create_mmap_dataset_builder(index_to_metric_fname, metric_value_dtype) - for unique_v in unique_metric_values: - samples = [] - for w_idx in range(num_workers): - for t_idx in range(num_threads): - w_metric_save_path = f"{metric_save_path}/worker{w_idx}_thread{t_idx}/" - w_metric_to_sample_fname = f"{w_metric_save_path}/{metric_name}_metric_to_sample_{unique_v}.csv" - if os.path.isfile(w_metric_to_sample_fname): - with open(w_metric_to_sample_fname, 'r') as f: - datareader = csv.reader(f) - for row in datareader: - samples += [int(x) for x in row] - index_to_sample_builder.add_item(torch.tensor(samples, dtype=torch.long)) - index_to_metric_builder.add_item(torch.tensor([unique_v], dtype=torch.long)) - logger.info(f"Finished reducing metric {metric_name} value {unique_v}.") - close_mmap_dataset_builder(index_to_sample_builder, index_to_sample_fname) - close_mmap_dataset_builder(index_to_metric_builder, index_to_metric_fname) - - def merge_map_results(self, dataset, metric_names, metric_types, save_path, num_workers, num_threads, - num_threads_reduce): - total_num_samples = len(dataset) - sample_idx_dtype = find_fit_int_dtype(0, total_num_samples - 1) - logger.info( - f"Total number of data samples: {total_num_samples}. Will use {sample_idx_dtype} to store the sample indexes." - ) - for m_idx in range(len(metric_names)): - metric_name, metric_type = metric_names[m_idx], metric_types[m_idx] - if metric_type == 'single_value_per_sample': - metric_save_path = f"{save_path}/{metric_name}/" - sample_to_metric_count = 0 - unique_metric_values = set([]) - manager = Manager() - return_dict = manager.dict() - p = [] - for t_idx_reduce in range(num_threads_reduce): - p.append( - Process(target=self.merge_gather_map_stats, - args=( - num_workers, - num_threads, - num_threads_reduce, - t_idx_reduce, - metric_save_path, - metric_name, - return_dict, - ))) - p[t_idx_reduce].start() - for t_idx_reduce in range(num_threads_reduce): - p[t_idx_reduce].join() - for t_idx_reduce in range(num_threads_reduce): - results = return_dict[t_idx_reduce] - for res in results: - unique_metric_values = unique_metric_values.union(set(res[0])) - sample_to_metric_count += res[1] - value_max = max(unique_metric_values) - value_min = min(unique_metric_values) - assert sample_to_metric_count == total_num_samples, "The number of samples in map result files are not correct. It's possible that some map worker didn't finish successfully." - metric_value_dtype = find_fit_int_dtype(value_min, value_max) - logger.info( - f"Metric {metric_name} has values between {value_min} and {value_max}. Will use {metric_value_dtype} to store the metric values." - ) - - # sample_to_metric - map_worker_thread = [] - for w_idx in range(num_workers): - for t_idx in range(num_threads): - map_worker_thread.append([w_idx, t_idx]) - thread_splits = split_index(0, len(map_worker_thread), num_threads_reduce) - p = [] - for t_idx_reduce in range(num_threads_reduce): - start_idx, end_idx = thread_splits[t_idx_reduce][0], thread_splits[t_idx_reduce][1] - p.append( - Process(target=self.merge_sample_to_metric, - args=( - t_idx_reduce, - metric_save_path, - metric_name, - metric_value_dtype, - map_worker_thread[start_idx:end_idx], - ))) - p[t_idx_reduce].start() - for t_idx_reduce in range(num_threads_reduce): - p[t_idx_reduce].join() - - sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric" - sample_to_metric_builder = create_mmap_dataset_builder(sample_to_metric_fname, metric_value_dtype) - for t_idx_reduce in range(num_threads_reduce): - chunk_fname = f"{metric_save_path}/{metric_name}_sample_to_metric_thread{t_idx_reduce}" - logger.info(f"Merging file {chunk_fname}") - sample_to_metric_builder.merge_file_(chunk_fname) - close_mmap_dataset_builder(sample_to_metric_builder, sample_to_metric_fname) - sample_to_metric = MMapIndexedDataset(sample_to_metric_fname, skip_warmup=True) - assert len(sample_to_metric) == total_num_samples - - # metric_to_sample - unique_metric_values = list(sorted(unique_metric_values)) - thread_splits = split_index(0, len(unique_metric_values), num_threads_reduce) - p = [] - for t_idx_reduce in range(num_threads_reduce): - start_idx, end_idx = thread_splits[t_idx_reduce][0], thread_splits[t_idx_reduce][1] - p.append( - Process(target=self.merge_metric_to_sample, - args=( - t_idx_reduce, - metric_save_path, - metric_name, - sample_idx_dtype, - metric_value_dtype, - unique_metric_values[start_idx:end_idx], - num_workers, - num_threads, - ))) - p[t_idx_reduce].start() - for t_idx_reduce in range(num_threads_reduce): - p[t_idx_reduce].join() - index_to_sample_fname = f"{metric_save_path}/{metric_name}_index_to_sample" - index_to_sample_builder = create_mmap_dataset_builder(index_to_sample_fname, sample_idx_dtype) - index_to_metric_fname = f"{metric_save_path}/{metric_name}_index_to_metric" - index_to_metric_builder = create_mmap_dataset_builder(index_to_metric_fname, metric_value_dtype) - for t_idx_reduce in range(num_threads_reduce): - chunk_is_fname = f"{metric_save_path}/{metric_name}_index_to_sample_thread{t_idx_reduce}" - logger.info(f"Merging file {chunk_is_fname}") - index_to_sample_builder.merge_file_(chunk_is_fname) - chunk_im_fname = f"{metric_save_path}/{metric_name}_index_to_metric_thread{t_idx_reduce}" - logger.info(f"Merging file {chunk_im_fname}") - index_to_metric_builder.merge_file_(chunk_im_fname) - close_mmap_dataset_builder(index_to_sample_builder, index_to_sample_fname) - close_mmap_dataset_builder(index_to_metric_builder, index_to_metric_fname) - - num_sample_per_value = DataAnalyzer.output_index_to_sample_percentile( - index_to_sample_fname, index_to_metric_fname, metric_name, metric_save_path, total_num_samples, - sample_idx_dtype) - self.get_metric_value_percentiles(metric_name, num_sample_per_value, total_num_samples) - elif metric_type == 'accumulate_value_over_samples': - metric_save_path = f"{save_path}/{metric_name}/" - metric_value = None - for w_idx in range(num_workers): - for t_idx in range(num_threads): - w_metric_save_path = f"{metric_save_path}/worker{w_idx}_thread{t_idx}/" - w_metric_value_fname = f"{w_metric_save_path}/{metric_name}_metric_value" - w_metric_value = MMapIndexedDataset(w_metric_value_fname, skip_warmup=True) - if metric_value is None: - metric_value = np.copy(w_metric_value[0]) - else: - metric_value += np.copy(w_metric_value[0]) - value_max = int(max(metric_value)) - value_min = int(min(metric_value)) - metric_value_dtype = find_fit_int_dtype(value_min, value_max) - metric_value_fname = f"{metric_save_path}/{metric_name}_metric_value" - metric_value_builder = create_mmap_dataset_builder(metric_value_fname, metric_value_dtype) - metric_value_builder.add_item(torch.tensor(metric_value.astype(np.int64), dtype=torch.long)) - close_mmap_dataset_builder(metric_value_builder, metric_value_fname) - - @staticmethod - def output_index_to_sample_percentile(index_to_sample_fname, index_to_metric_fname, metric_name, metric_save_path, - total_num_samples, sample_idx_dtype): - """ read index_to_metric and index_to_sample files and write distribution to index_to_sample_percentage_merged """ - num_sample_per_value = {} - index_to_sample = MMapIndexedDataset(index_to_sample_fname, skip_warmup=True) - index_to_metric = MMapIndexedDataset(index_to_metric_fname, skip_warmup=True) - index_to_sample_merged_fname = f"{metric_save_path}/{metric_name}_index_to_sample_percentile_merged" - index_to_sample_merged_builder = create_mmap_dataset_builder(index_to_sample_merged_fname, sample_idx_dtype) - for v_idx in range(len(index_to_sample)): - if v_idx > 0: - assert index_to_metric[v_idx] > index_to_metric[v_idx - 1] - num_sample_per_value[index_to_metric[v_idx][0]] = len(index_to_sample[v_idx]) - assert sum(list(num_sample_per_value.values())) == total_num_samples - merge_step = max(1, len(index_to_sample) // 100) - for v_idx in range(0, len(index_to_sample), merge_step): - merged_samples = np.copy( - np.concatenate(index_to_sample[v_idx:min(len(index_to_sample), (v_idx + merge_step))], axis=None)) - index_to_sample_merged_builder.add_item(torch.tensor(merged_samples.astype(np.int64), dtype=torch.long)) - logger.info(f"Finished merging index_to_sample {v_idx} to {v_idx+merge_step}.") - close_mmap_dataset_builder(index_to_sample_merged_builder, index_to_sample_merged_fname) - return num_sample_per_value - - def run_reduce(self): - if self.custom_reduce is None: - self.merge_map_results(self.dataset, self.metric_names, self.metric_types, self.save_path, - self.num_workers, self.num_threads, self.num_threads_reduce) - else: - self.custom_reduce(self.dataset, self.metric_names, self.metric_types, self.save_path, self.num_workers, - self.num_threads, self.num_threads_reduce) - - def run_map_reduce(self, comm_group=None): - self.run_map() - # wait for the mapping operation, where all nodes outputs their own (partial) result files - dist.barrier(group=comm_group) - if self.worker_id == 0: - self.run_reduce() - # wait for the reduce, where rank 0 merges all (partial) files. Dataset can then be used by all nodes. - dist.barrier(group=comm_group) - - -class DistributedDataAnalyzer(object): - - def __init__( - self, - dataset, - num_workers=1, - num_threads=1, - worker_id=0, - batch_size=1, - metric_names=[], - metric_functions=[], - metric_types=[], - save_path="./", - collate_fn=None, - device='cuda', - comm_group=None, - sample_indices=None, - ) -> None: - self.dataset = dataset - self.batch_size = batch_size - self.metric_names = metric_names - self.metric_functions = metric_functions - self.metric_types = metric_types - self.save_path = save_path - self.collate_fn = collate_fn - self.device = device - self.sample_indices = sample_indices - self.num_threads = num_threads - self.worker_id = worker_id - - if not dist.is_initialized(): - dist.init_distributed() - - # comm_group and worker_id+num_workers are mutually exclusive - self.comm_group = comm_group - if self.comm_group is None: - # self.comm_group = deepspeed.utils.groups._clone_world_group() - self.num_workers = num_workers - self.worker_id = worker_id - else: - self.num_workers = self.comm_group.size() - self.worker_id = self.comm_group.rank() - - if self.worker_id == 0: - logger.info(f"Distributed data analyzer initialized with {self.num_workers} workers.") - - def run_map_helper(self, thread_id=0, metric_queues=None): - thread_start_idx, thread_end_idx = self.thread_splits[thread_id][0], self.thread_splits[thread_id][1] - worker_dataset = Subset(self.dataset, list(range(thread_start_idx, thread_end_idx))) - sampler = BatchSampler(SequentialSampler(worker_dataset), batch_size=self.batch_size, drop_last=False) - dataloader = DataLoader(dataset=worker_dataset, - batch_sampler=sampler, - num_workers=0, - collate_fn=self.collate_fn, - pin_memory=False) - - # set initial results list - metric_results = [] - for metric_type in self.metric_types: - assert metric_type in ['single_value_per_sample', 'accumulate_value_over_samples'], \ - f"metric_type {metric_type} not implemented." - metric_results.append([] if metric_type == 'single_value_per_sample' else None) - - # iterate dataloader and store metric results - batch_start_idx = thread_start_idx - for data in dataloader: - for m_idx in range(len(self.metric_names)): - metric_type, metric_function = self.metric_types[m_idx], self.metric_functions[m_idx] - metric_values = metric_function(data) - assert torch.is_tensor(metric_values) or isinstance(metric_values, np.ndarray), \ - "metric_function must return a tensor or array" - if isinstance(metric_values, np.ndarray): - metric_values = torch.from_numpy(metric_values) - assert metric_values.dtype in valid_dtypes, \ - f"metric_function result dtype {metric_values.dtype} not supported. Supported dtypes {valid_dtypes}" - - if metric_type == 'single_value_per_sample': - for row in range(metric_values.size()[0]): - value = metric_values[row].item() - sample_idx = batch_start_idx + row # sample idx following dataset iteration order - if isinstance(data, dict) and 'index' in data: # Megatron use case - sample_idx = data['index'][row][0].item() - elif self.sample_indices is not None: # user defined shuffling of indices - sample_idx = self.sample_indices[sample_idx] - metric_results[m_idx].append((value, sample_idx)) - elif metric_type == 'accumulate_value_over_samples': - if metric_results[m_idx] is None: - metric_results[m_idx] = metric_values - else: - metric_results[m_idx].add_(metric_values) - batch_start_idx += len(data) - - if self.num_threads == 1: - return metric_results - - # copy metric_results to the shared queue - assert metric_queues - for m_idx in range(len(self.metric_names)): - results = metric_results[m_idx] - if torch.is_tensor(results): - results = results.item() if results.dim() == 0 else results.tolist() - try: - metric_queues[m_idx].put((thread_id, results)) - except Exception as e: - logger.error(f"Error putting metric results to queue: {e}") - sys.exit(1) - - def run_map_reduce(self): - - # setup individual dataloaders - self.worker_splits, self.thread_splits = split_dataset(self.dataset, - self.num_workers, - self.worker_id, - num_threads=self.num_threads) - node_start_idx, node_end_idx = self.worker_splits[self.worker_id] - logger.info(f"worker {self.worker_id} working on data subset {node_start_idx} to {node_end_idx}.") - - if self.num_threads in [0, 1, None]: - metric_results = self.run_map_helper() - metric_results = [torch.tensor(m).to(self.device) for m in metric_results] - else: - - # create a shared queue of results per metric to be populated by individual threads - with Manager() as manager: - metric_queues = [manager.Queue() for _ in self.metric_names] - threads = [ - Process(target=self.run_map_helper, args=(t, metric_queues)) for t in range(self.num_threads) - ] - for thread in threads: - thread.start() - for thread in threads: - thread.join() - - # gather results from shared queues into metric_results - metric_results = [None for _ in self.metric_names] - for m_idx, (queue, metric_type) in enumerate(zip(metric_queues, self.metric_types)): - while not queue.empty(): - t_idx, t_results = queue.get() - t_start_idx, t_end_idx = self.thread_splits[t_idx] - if t_start_idx >= t_end_idx: # no results from this thread - continue #corner case for small datasets and high thread count - t_results = torch.tensor(t_results) - if metric_type == 'single_value_per_sample': - # add thread results to the metric_results list, ordered by thread idx - if metric_results[m_idx] is None: # initialize if needed - metric_results[m_idx] = torch.zeros(node_end_idx - node_start_idx, - t_results.size(1)).to(self.device) - metric_results[m_idx][t_start_idx - node_start_idx:t_end_idx - node_start_idx] = t_results - else: - if metric_results[m_idx] is None: # initialize if needed - metric_results[m_idx] = torch.zeros(t_results.size()).to(self.device) - metric_results[m_idx].add_(t_results) - - # compute dtype for sample ids - total_num_samples = len(self.dataset) - sample_idx_dtype = find_fit_int_dtype(0, total_num_samples - 1) - logger.info(f"Total number of data samples: {total_num_samples}.") - logger.info(f"Will use {sample_idx_dtype} to store the sample indexes.") - - for m_idx in range(len(self.metric_names)): - metric_values, metric_name, metric_type = \ - metric_results[m_idx], self.metric_names[m_idx], self.metric_types[m_idx] - metric_save_path = f"{self.save_path}/{metric_name}/" - os.makedirs(metric_save_path, exist_ok=True) - - if metric_type == 'single_value_per_sample': - - # Compute sample and metric value dtypes based on range - values, samples = metric_values[:, 0], metric_values[:, 1] - value_min, value_max = Dist.min_max(values, self.comm_group) - sample_min, sample_max = Dist.min_max(samples, self.comm_group) - metric_value_dtype = find_fit_int_dtype(value_min, value_max) - sample_value_dtype = find_fit_int_dtype(sample_min, sample_max) - - # sample_to_metric maps sample ids to metric values, as a list of metric values - sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric" - values = [torch.tensor([x]) for x in metric_values[:, 0]] - self.file_write_ordered(values, sample_to_metric_fname, metric_value_dtype) - - # distributed sorting by values, gives an ordered disjoint subset of keys on nodes - metric_values = Dist.sample_sort(metric_values, self.comm_group, self.num_workers) - metric_to_samples_dict = {} - if len(metric_values) > 0: - for value, sample in metric_values: - if value.item() not in metric_to_samples_dict: - metric_to_samples_dict[value.item()] = [] - metric_to_samples_dict[value.item()].append(sample.item()) - - # index_to_metric and index_to_sample serialize a dicitonary from metric to samples - # index_to_metric stores a key per row, index_to_sample stores the values per row - values = [torch.tensor([x]) for x in metric_to_samples_dict.keys()] - samples = [torch.tensor(metric_to_samples_dict[x]) for x in metric_to_samples_dict.keys()] - index_to_metric_fname = f"{metric_save_path}/{metric_name}_index_to_metric" #dict keys - index_to_sample_fname = f"{metric_save_path}/{metric_name}_index_to_sample" #dict values - self.file_write_ordered(values, index_to_metric_fname, metric_value_dtype) - self.file_write_ordered(samples, index_to_sample_fname, sample_value_dtype) - - if self.worker_id == 0: - DataAnalyzer.output_index_to_sample_percentile(index_to_sample_fname, index_to_metric_fname, - metric_name, metric_save_path, total_num_samples, - sample_idx_dtype) - dist.barrier(self.comm_group) - - elif metric_type == 'accumulate_value_over_samples': - metric_value_fname = f"{metric_save_path}/{metric_name}_metric_value" - dist.reduce(metric_values, dst=0, op=dist.ReduceOp.SUM, group=self.comm_group) - metric_value_dtype = find_fit_int_dtype(metric_values.min(), metric_values.max()) - - if self.worker_id == 0: - builder = create_mmap_dataset_builder(metric_value_fname, metric_value_dtype) - builder.add_item(metric_values.cpu()) - close_mmap_dataset_builder(builder, metric_value_fname) - dist.barrier(self.comm_group) - - def file_write_ordered(self, tensor_list, fname, numpy_dtype): - """ MPI_file_write_ordered extended to write a list of tensors, by one rank, iteratively """ - - # each node has a list of rows (tensors) to be written to the file. - # we will serialize it in order to communicate it in one comm step. - - tkwargs = dict(dtype=torch.int64, device=self.device) - - # 1. gather on rank 0 the number of rows to be sent/recv - row_count = torch.tensor([len(tensor_list)], **tkwargs) - row_counts = torch.zeros(self.num_workers, **tkwargs) - dist.all_gather_into_tensor(row_counts, row_count, group=self.comm_group) - assert row_counts[self.worker_id] == row_count == len(tensor_list), "all_gather failed" - - # 2. gather on rank 0 the sizes of the rows to be sent/recv - row_len = torch.tensor([len(l) for l in tensor_list], **tkwargs) - row_lens = Dist.gather_v(row_len, 0, self.comm_group, self.num_workers, self.worker_id) - - # 4. gather on rank 0 of the total size (sum of all row lengths) to be received - size = torch.tensor([sum(row_len).item()], **tkwargs) - sizes = torch.zeros(self.num_workers, **tkwargs) - dist.all_gather_into_tensor(sizes, size, group=self.comm_group) - assert sizes[self.worker_id] == size.item(), "all_gather did not return the same sizes" #sanity check - - # method to deserializes a buffer into rows of different lengths and write them to file - def write_buffer_to_file(buff, src, builder): - assert self.worker_id == 0, "only rank 0 can write to file" - - # collect all buffers and write them at once - buff = buff.cpu().detach().numpy() - row_offsets = np.cumsum([0] + row_lens[src].tolist()) - arr_list = [] - for i in range(len(row_lens[src])): - arr_list.append(buff[row_offsets[i]:row_offsets[i + 1]]) - builder.add_items(arr_list) - - # 5. rank 0 prepares output folder and file - if self.worker_id == 0: - os.makedirs(os.path.dirname(fname), exist_ok=True) - builder = create_mmap_dataset_builder(fname, numpy_dtype) - - # iterate through ranks that have data to be sent/recv/written - for src in [rank for rank, count in enumerate(row_counts) if count > 0]: - - dist.barrier(group=self.comm_group) - if self.worker_id == 0 and src == 0: # rank 0's write its own data - buffer = torch.cat(tensor_list, dim=0).to(self.device) - write_buffer_to_file(buffer, 0, builder) - elif self.worker_id == 0 and src > 0: # rank 0 receives other rank's data and writes it - buffer = torch.empty(sizes[src].item(), dtype=buffer.dtype, device=buffer.device) - err = dist.recv(buffer, src=src, group=self.comm_group, tag=src) - assert err == src and len(buffer) > 0, "recv failed" - write_buffer_to_file(buffer, src, builder) - elif self.worker_id == src: # current rank sends data to rank 0 - buffer = torch.cat(tensor_list, dim=0).to(self.device) - dist.send(buffer, 0, group=self.comm_group, tag=src) - - # rank 0 closes the file - if self.worker_id == 0: - close_mmap_dataset_builder(builder, fname) # close file - dist.barrier(self.comm_group) - - -class Dist: - """ auxiliary class to perform distributed operations on tensors""" - - @staticmethod - def min_max(tensor, comm_group): - """ given a distributed tensor, return the min/max values across all ranks""" - - value_min, value_max = tensor.min(), tensor.max() - dist.reduce(value_min, 0, op=dist.ReduceOp.MIN, group=comm_group) - dist.reduce(value_max, 0, op=dist.ReduceOp.MAX, group=comm_group) - return value_min.item(), value_max.item() - - @staticmethod - def gather_v(tensor, dst, comm_group, num_workers, worker_id): - """ MPI_Gatherv. gather tensors of variable sizes in a single rank """ - - # gather the number of rows to be sent/recv - size = torch.tensor([len(tensor)], dtype=torch.int64, device=tensor.device) - sizes = torch.zeros(num_workers, dtype=torch.int64, device=tensor.device) - dist.all_gather_into_tensor(sizes, size, group=comm_group) - assert sizes[worker_id] == size, "all_gather failed" - - # all_gather requires all tensors to be of same size so we need to pad them - max_size = max(sizes).item() - buffer = torch.empty(max_size, dtype=tensor.dtype, device=tensor.device) - buffer[0:size] = tensor.data - buffer_list = None - if worker_id == 0: # create padded recv buffers - buffer_list = [torch.empty(max_size, dtype=tensor.dtype, device=tensor.device) for _ in range(num_workers)] - dist.gather(buffer, buffer_list, dst=dst, group=comm_group) - - # revert padding and return value - if worker_id == 0: - buffer_list = [r[:s.item()] for r, s in zip(buffer_list, sizes)] - return buffer_list - - @staticmethod - def sample_sort(tensor, comm_group, num_workers, n_samples=100): - """ perform a distributed random sort of a tensor, and returns the sorted partial tensor""" - device, dims = tensor.device, tensor.size()[1] - - # 1 - sort rows by first column, then second column, then third, etc... - tensor = torch.tensor(sorted(tensor.tolist()), dtype=tensor.dtype, device=tensor.device) - - # 2 - collect few samples per rank - idx = torch.round(torch.linspace(0, len(tensor) - 1, n_samples)).to(int) - samples = tensor[idx][:, 0].contiguous().to(device) #only first column, all but last row - - # 2 - Allgather samples - all_samples = [torch.zeros(n_samples, dtype=samples.dtype, device=device) for _ in range(num_workers)] - dist.all_gather(all_samples, samples, group=comm_group) - all_samples = torch.cat(all_samples, dim=0).to(device) - - # 3 - Sort all samples and collect the ranges of each rank as equidistant - all_samples = all_samples.sort()[0] - idx = torch.round(torch.linspace(0, len(all_samples) - 1, num_workers + 1)).to(int) - ranges = all_samples[idx] # range of each rank r as ranges[r] <= x < ranges[r+1] - ranges[-1] += 1 # increase upper limit of last rank so that x < ranges[r+1]. - - # 4 - collect elements to send to each rank, based on the rank ranges - send = [] - for rank in range(num_workers): - mask = (tensor[:, 0] >= ranges[rank]) & (tensor[:, 0] < ranges[rank + 1]) - send.append(tensor[mask]) - - # 5. all to all to communicate the sizes to be sent/recv - send_count = [torch.tensor([len(s) * dims], dtype=torch.int64, device=device) for s in send] - recv_count = list(torch.empty([num_workers], dtype=torch.int64, device=device).chunk(num_workers)) - dist.all_to_all(recv_count, send_count, group=comm_group) - - # 6. all-to-all-v to communicate the elements to be sent/recv as a single tensor - send = torch.cat(send, dim=0).flatten().to(device) - recv = torch.zeros(sum(recv_count), dtype=send.dtype).to(device) - send_count = [s.item() for s in send_count] # convert to list of ints - recv_count = [r.item() for r in recv_count] - dist.all_to_all_single(recv, send, recv_count, send_count, group=comm_group) - del send - - # 7. the received tensor is the 1D disjoint subset of the distributed tensor. - # We will recover the original dimensionality and sort it by columns again. - recv = recv.view(-1, dims) - recv = torch.tensor(sorted(recv.tolist()), dtype=recv.dtype, device=recv.device) - return recv - - -def test_compare_both_data_analyzers(dataset): - """ given a dataset, compare file and memory based data analyser""" - - id = lambda t: t.to(torch.int64) # identity - batch_sum = lambda t: id(t).sum() #sum batch - num_threads = 4 - kwargs = dict( - dataset=dataset, - batch_size=2**10, - worker_id=int(os.environ['RANK']), - num_workers=int(os.environ['WORLD_SIZE']), - metric_names=["mod", "batch_sum"], - metric_functions=[id, batch_sum], - metric_types=['single_value_per_sample', 'accumulate_value_over_samples'], - num_threads=num_threads, - ) - - dda = DistributedDataAnalyzer( - save_path="./output_dist", - device=f"cuda:{int(os.environ['LOCAL_RANK'])}", - **kwargs, - ) - start_time = time.time() - dda.run_map_reduce() - if dda.worker_id == 0: - print("DistributedDataAnalyzer runtime: %s seconds " % (time.time() - start_time)) - - da = DataAnalyzer(num_threads_reduce=num_threads, - save_path="./output_disk", - metric_dtypes=[torch.int64, torch.int64], - **kwargs) - start_time = time.time() - da.run_map_reduce() - if da.worker_id == 0: - print("DataAnalyzer runtime: %s seconds " % (time.time() - start_time)) - - output_paths = [ - "batch_sum/batch_sum_metric_value.bin", "batch_sum/batch_sum_metric_value.idx", \ - "mod/mod_index_to_metric.bin", "mod/mod_index_to_metric.idx", \ - "mod/mod_index_to_sample.bin", "mod/mod_index_to_sample.idx", \ - "mod/mod_index_to_sample_percentile_merged.bin", "mod/mod_index_to_sample_percentile_merged.idx", \ - "mod/mod_sample_to_metric.bin", "mod/mod_sample_to_metric.idx" - ] - - if dda.worker_id == 0: - for path in output_paths: - with open(os.path.join(da.save_path, path), 'rb') as f1, \ - open(os.path.join(dda.save_path, path), 'rb') as f2: - # if files have suffix .bin, they should be identical - if path.endswith(".bin"): - assert f1.read() == f2.read(), f"files {path} are not identical." - elif f1.read() != f2.read(): - print(f"files {path} are not identical.") - dist.barrier() - dist.destroy_process_group() - - -if __name__ == "__main__": - - class TestDataset(torch.utils.data.Dataset): - - def __init__(self, size=10_000_000): - self.values = [(x + 7) % 10_000 for x in range(size)] - self.size = size - - __len__ = lambda self: self.size - __getitem__ = lambda self, idx: self.values[idx] - - test_compare_both_data_analyzers(TestDataset()) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py deleted file mode 100644 index 100bef3f7946c8d2b7280e234abc5e625b6c68f9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py +++ /dev/null @@ -1,349 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -coding=utf-8 - Copyright (c) 2020, NVIDIA CORPORATION. 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. -Part of this code was adopted from https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/data/data_samplers.py -""" - -import torch -import os -import numpy as np - -import deepspeed.comm as dist -from deepspeed.utils import logger -from deepspeed.accelerator import get_accelerator -from ..constants import * -from ..curriculum_scheduler import CurriculumScheduler -from .indexed_dataset import MMapIndexedDataset -from .utils import create_mmap_dataset_builder, close_mmap_dataset_builder, find_fit_int_dtype - - -class DeepSpeedDataSampler(object): - - def __init__(self, - data_efficiency_config, - one_epoch_total_samples, - micro_batch_size, - data_parallel_rank, - data_parallel_size, - data_parallel_group, - gradient_accumulation_steps, - global_rank, - drop_last=True): - # Keep a copy of input params for later use. - self.data_efficiency_config = data_efficiency_config - self.one_epoch_total_samples = one_epoch_total_samples - self.index_dtype = find_fit_int_dtype(0, one_epoch_total_samples) - self.total_samples = one_epoch_total_samples * self.data_efficiency_config[DATA_SAMPLING][ - DATA_SAMPLING_NUM_EPOCHS] - self.micro_batch_size = micro_batch_size - self.data_parallel_rank = data_parallel_rank - self.data_parallel_group = data_parallel_group - self.micro_batch_times_data_parallel_size = \ - self.micro_batch_size * data_parallel_size - self.gradient_accumulation_steps = gradient_accumulation_steps - self.global_batch_size = self.micro_batch_times_data_parallel_size * \ - self.gradient_accumulation_steps - self.global_rank = global_rank - self.drop_last = drop_last - self.np_rng = np.random.default_rng(self.data_efficiency_config[DATA_EFFICIENCY_SEED]) - self.state = {} - self.batch = [] - self.consumed_samples = 0 - if self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_ENABLED]: - self.curriculum_step = 0 - self.current_difficulties = {} - self.data_cluster_paths = [] - self.data_cluster_current_position = [] - self.curriculum_schedulers = {} - self.curriculum_index_to_sample = {} - self.curriculum_index_to_metric = {} - self.difficulty_type = {} - self.clustering_type = {} - self.data_1epoch_size = None - if self.global_rank == 0: - self.data_clusters = [] - self.data_cluster_sizes = [] - cluster_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ - CURRICULUM_LEARNING_CLUSTER_PATH] - if not os.path.exists(cluster_path): - os.makedirs(cluster_path) - for metric in self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS]: - self.curriculum_schedulers[metric] = CurriculumScheduler( - data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS][metric]) - self.difficulty_type[metric] = data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ - CURRICULUM_LEARNING_METRICS][metric][CURRICULUM_LEARNING_DIFFICULTY_TYPE] - self.clustering_type[metric] = data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ - CURRICULUM_LEARNING_METRICS][metric][CURRICULUM_LEARNING_CLUSTERING_TYPE] - if self.global_rank == 0: - if self.clustering_type[metric] != CURRICULUM_LEARNING_SINGLE_CLUSTER: - self.curriculum_index_to_sample[metric] = MMapIndexedDataset( - data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS] - [metric][CURRICULUM_LEARNING_SAMPLE_PATH], - skip_warmup=True) - if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED: - self.curriculum_index_to_metric[metric] = MMapIndexedDataset( - data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS] - [metric][CURRICULUM_LEARNING_METRIC_PATH], - skip_warmup=True) - - # Sanity checks. - assert self.total_samples > 0, \ - 'no sample to consume: {}'.format(self.total_samples) - assert self.micro_batch_size > 0 - assert data_parallel_size > 0 - assert self.data_parallel_rank < data_parallel_size, \ - 'data_parallel_rank should be smaller than data size: {}, ' \ - '{}'.format(self.data_parallel_rank, data_parallel_size) - - def __len__(self): - return self.total_samples - - def set_custom_curriculum_learning_schedule(self, schedule_func_dict): - for metric in self.curriculum_schedulers: - if metric in schedule_func_dict: - self.curriculum_schedulers[metric].set_custom_get_difficulty(schedule_func_dict[metric]) - - def get_start_end_idx(self, batch_len=None): - """ - given the length of a minibatch (defaults to micro-batch size * data_parallel_size), - return the start and end indices of the current data parallel rank - """ - batch_len = batch_len or self.micro_batch_times_data_parallel_size - start_idx_fn = lambda r: round(r * batch_len / self.data_parallel_group.size()) - start_idx = start_idx_fn(self.data_parallel_rank) - end_idx = start_idx_fn(self.data_parallel_rank + 1) - return start_idx, end_idx - - def get_sample_based_on_metric_value(self, metric, value_start, value_end): - new_samples = None - for row in range(len(self.curriculum_index_to_sample[metric])): - if self.curriculum_index_to_metric[metric][row] <= value_end and self.curriculum_index_to_metric[metric][ - row] > value_start: - row_samples = np.copy(self.curriculum_index_to_sample[metric][row]) - new_samples = row_samples if new_samples is None else np.concatenate( - (new_samples, row_samples), axis=None) - return new_samples - - def get_sample_based_on_metric_percentile(self, metric, percentile_start, percentile_end): - new_samples = None - if self.data_1epoch_size is None: - self.data_1epoch_size = sum(len(x) for x in self.curriculum_index_to_sample[metric]) - max_percentile = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS][ - metric][CURRICULUM_LEARNING_MAX_DIFFICULTY] - sample_per_percentile = self.data_1epoch_size // max_percentile - start_count = sample_per_percentile * percentile_start - end_count = sample_per_percentile * percentile_end - if percentile_end == max_percentile: - end_count = self.data_1epoch_size - current_count = 0 - for row in range(len(self.curriculum_index_to_sample[metric])): - row_size = len(self.curriculum_index_to_sample[metric][row]) - if current_count + row_size > start_count: - row_start = max(0, start_count - current_count) - if current_count + row_size <= end_count: - row_end = row_size - else: - row_end = end_count - current_count - row_samples = np.copy(self.curriculum_index_to_sample[metric][row][row_start:row_end]) - new_samples = row_samples if new_samples is None else np.concatenate( - (new_samples, row_samples), axis=None) - current_count += row_size - if current_count >= end_count: - break - return new_samples - - def get_new_cluster(self, previous_difficulties): - cluster_fname = CURRICULUM_LEARNING_CLUSTER_PREFIX - for metric in self.curriculum_schedulers: - cluster_fname = f"{cluster_fname}_{metric}{self.current_difficulties[metric]}" - cluster_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ - CURRICULUM_LEARNING_CLUSTER_PATH] - cluster_path = f"{cluster_path}/{cluster_fname}" - if self.global_rank == 0: - new_cluster = None - need_clustering = 0 - for metric in self.clustering_type: - if self.clustering_type[metric] != CURRICULUM_LEARNING_SINGLE_CLUSTER: - need_clustering += 1 - if need_clustering > 1: - for metric in self.curriculum_schedulers: - if self.clustering_type[metric] == CURRICULUM_LEARNING_SINGLE_CLUSTER: - metric_cluster = np.arange(start=0, - stop=self.one_epoch_total_samples, - step=1, - dtype=self.index_dtype) - else: - if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED: - metric_cluster = self.get_sample_based_on_metric_value(metric, float('-inf'), - self.current_difficulties[metric]) - elif self.difficulty_type[metric] == CURRICULUM_LEARNING_PERCENTILE_BASED: - metric_cluster = self.get_sample_based_on_metric_percentile( - metric, 0, self.current_difficulties[metric]) - new_cluster = metric_cluster if new_cluster is None else \ - np.intersect1d(new_cluster, metric_cluster, assume_unique=True) - for cluster in self.data_clusters: - new_cluster = np.setdiff1d(new_cluster, cluster[0], assume_unique=True) - else: - if len(self.data_clusters) == 0: - new_cluster = np.arange(start=0, stop=self.one_epoch_total_samples, step=1, dtype=self.index_dtype) - for metric in self.curriculum_schedulers: - if self.clustering_type[metric] != CURRICULUM_LEARNING_SINGLE_CLUSTER: - if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED: - new_cluster = self.get_sample_based_on_metric_value(metric, previous_difficulties[metric], - self.current_difficulties[metric]) - elif self.difficulty_type[metric] == CURRICULUM_LEARNING_PERCENTILE_BASED: - new_cluster = self.get_sample_based_on_metric_percentile( - metric, previous_difficulties[metric], self.current_difficulties[metric]) - if new_cluster is not None and len(new_cluster) > 0: - logger.info( - f"new data cluster (previous_difficulties {previous_difficulties}, current_difficulties {self.current_difficulties}) with size {len(new_cluster)} generated." - ) - self.np_rng.shuffle(new_cluster) - cluster_builder = create_mmap_dataset_builder(cluster_path, self.index_dtype) - cluster_builder.add_item_numpy(new_cluster) - close_mmap_dataset_builder(cluster_builder, cluster_path) - self.data_clusters.append(MMapIndexedDataset(cluster_path, skip_warmup=True)) - self.data_cluster_sizes.append(len(self.data_clusters[-1][0])) - else: - logger.info( - f"new data cluster (previous_difficulties {previous_difficulties}, current_difficulties {self.current_difficulties}) has no matched data thus skipped." - ) - dist.barrier(group=self.data_parallel_group) - if os.path.isfile(f"{cluster_path}.bin"): - self.data_cluster_paths.append(cluster_fname) - self.data_cluster_current_position.append(0) - - def sample_from_clusters(self): - num_clusters = len(self.data_clusters) - weight_sum = sum(self.data_cluster_sizes) - weights = [x / weight_sum for x in self.data_cluster_sizes] - samples = self.np_rng.choice(num_clusters, self.global_batch_size, replace=True, p=weights) - samples = np.bincount(samples, minlength=num_clusters) - return samples - - def reshuffle_clusters(self, cidx): - cluster_fname = self.data_cluster_paths[cidx] - cluster_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ - CURRICULUM_LEARNING_CLUSTER_PATH] - cluster_path = f"{cluster_path}/{cluster_fname}" - cluster = np.copy(self.data_clusters[cidx][0]) - self.np_rng.shuffle(cluster) - cluster_builder = create_mmap_dataset_builder(cluster_path, self.index_dtype) - cluster_builder.add_item_numpy(cluster) - close_mmap_dataset_builder(cluster_builder, cluster_path) - self.data_clusters[cidx] = MMapIndexedDataset(cluster_path, skip_warmup=True) - - def get_sample_from_cluster(self, cidx, num_samples): - start_idx = self.data_cluster_current_position[cidx] - samples = list(np.copy(self.data_clusters[cidx][0][start_idx:(start_idx + num_samples)])) - self.data_cluster_current_position[cidx] += num_samples - if len(samples) < num_samples: - num_samples_remained = num_samples - len(samples) - logger.info(f"reshuffling cluster {cidx}.") - self.reshuffle_clusters(cidx) - samples += list(np.copy(self.data_clusters[cidx][0][:num_samples_remained])) - self.data_cluster_current_position[cidx] = num_samples_remained - return samples - - def get_next_global_batch(self): - if self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_ENABLED]: - self.curriculum_step += 1 - new_cluster = False - previous_difficulties = {} - for metric in self.curriculum_schedulers: - next_difficulty = self.curriculum_schedulers[metric].update_difficulty(self.curriculum_step) - if metric not in self.current_difficulties or \ - next_difficulty != self.current_difficulties[metric]: - new_cluster = True - if metric in self.current_difficulties: - previous_difficulties[metric] = self.current_difficulties[metric] - else: - if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED: - previous_difficulties[metric] = float('-inf') - elif self.difficulty_type[metric] == CURRICULUM_LEARNING_PERCENTILE_BASED: - previous_difficulties[metric] = 0 - self.current_difficulties[metric] = next_difficulty - if new_cluster: - self.get_new_cluster(previous_difficulties) - if self.global_rank == 0: - samples_per_cluster = self.sample_from_clusters() - batch = [] - for cidx in range(len(samples_per_cluster)): - batch += self.get_sample_from_cluster(cidx, samples_per_cluster[cidx]) - self.np_rng.shuffle(batch) - - # broadcast tensor must have same shape across participants. So we fill batch with -1s when not full - assert len(batch) <= self.global_batch_size - batch += [-1] * (self.global_batch_size - len(batch)) - batch = torch.tensor(batch, device=get_accelerator().current_device_name(), dtype=torch.long).view(-1) - else: - batch = torch.empty(self.global_batch_size, - device=get_accelerator().current_device_name(), - dtype=torch.long) - dist.broadcast(batch, 0, group=self.data_parallel_group) - batch = batch[batch != -1] # remove trailing -1s used to fill incomplete batch tensor - self.batch = batch.tolist() - - def __iter__(self): - while self.consumed_samples <= self.total_samples: - if len(self.batch) == 0: - self.get_next_global_batch() - current_batch = self.batch[:self.micro_batch_times_data_parallel_size] - self.batch = self.batch[self.micro_batch_times_data_parallel_size:] - if len(current_batch) == self.micro_batch_times_data_parallel_size or \ - (len(current_batch) > 0 and not self.drop_last): - start_idx, end_idx = self.get_start_end_idx(len(current_batch)) - yield current_batch[start_idx:end_idx] - self.consumed_samples += len(current_batch) - current_batch = [] - - def state_dict(self): - return { - CURRICULUM_LEARNING_BATCH: self.batch, - CURRICULUM_LEARNING_CONSUMED_SAMPLES: self.consumed_samples, - CURRICULUM_LEARNING_STEP: self.curriculum_step, - CURRICULUM_LEARNING_CURRENT_DIFFICULTIES: self.current_difficulties, - CURRICULUM_LEARNING_DATA_CLUSTER_PATHS: self.data_cluster_paths, - CURRICULUM_LEARNING_DATA_CLUSTER_CURRENT_POSITION: self.data_cluster_current_position, - CURRICULUM_LEARNING_NP_RNG_STATE: np.random.get_state() - } - - def load_state_dict(self, state_dict): - self.batch = state_dict[CURRICULUM_LEARNING_BATCH] - self.consumed_samples = state_dict[CURRICULUM_LEARNING_CONSUMED_SAMPLES] - self.curriculum_step = state_dict[CURRICULUM_LEARNING_STEP] - self.current_difficulties = state_dict[CURRICULUM_LEARNING_CURRENT_DIFFICULTIES] - self.data_cluster_paths = state_dict[CURRICULUM_LEARNING_DATA_CLUSTER_PATHS] - self.data_cluster_current_position = state_dict[CURRICULUM_LEARNING_DATA_CLUSTER_CURRENT_POSITION] - np.random.set_state(state_dict[CURRICULUM_LEARNING_NP_RNG_STATE]) - cluster_root_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ - CURRICULUM_LEARNING_CLUSTER_PATH] - # Backward compatibility: previously data_cluster_paths were stored as - # absolute paths. Now we changed it to just the file name so that even - # if user moved the cluster files, the checkpoint loading still works - # as long as user set the correct new CURRICULUM_LEARNING_CLUSTER_PATH - # in deepspeed json config. - for idx in range(len(self.data_cluster_paths)): - if '/' in self.data_cluster_paths[idx]: - self.data_cluster_paths[idx] = self.data_cluster_paths[idx].split('/')[-1] - if self.global_rank == 0: - for cluster_fname in self.data_cluster_paths: - cluster_path = f"{cluster_root_path}/{cluster_fname}" - self.data_clusters.append(MMapIndexedDataset(cluster_path, skip_warmup=True)) - self.data_cluster_sizes.append(len(self.data_clusters[-1][0])) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/indexed_dataset.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/indexed_dataset.py deleted file mode 100644 index 872d05de01453aea76fc850de77da3acbbcea76c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/indexed_dataset.py +++ /dev/null @@ -1,627 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Part of this code was adopted from https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/data/indexed_dataset.py -""" - -# Copyright (c) Facebook, Inc. and its affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -# copied from fairseq/fairseq/data/indexed_dataset.py -# Removed IndexedRawTextDataset since it relied on Fairseq dictionary -# other slight modifications to remove fairseq dependencies -# Added document index to index file and made it accessible. -# An empty sentence no longer separates documents. - -# Some of the fixes/improvements are adopted from -# https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/main/megatron/data/indexed_dataset.py - -from functools import lru_cache -import os -import shutil -import struct -from itertools import accumulate - -import numpy as np -import torch - - -def __best_fitting_dtype(vocab_size=None): - if vocab_size is not None and vocab_size < 65500: - return np.uint16 - else: - return np.int32 - - -def get_available_dataset_impl(): - return ['lazy', 'cached', 'mmap'] - - -def infer_dataset_impl(path): - if IndexedDataset.exists(path): - with open(index_file_path(path), 'rb') as f: - magic = f.read(8) - if magic == IndexedDataset._HDR_MAGIC: - return 'cached' - elif magic == MMapIndexedDataset.Index._HDR_MAGIC[:8]: - return 'mmap' - else: - return None - else: - print(f"Dataset does not exist: {path}") - print("Path should be a basename that both .idx and .bin can be appended to get full filenames.") - return None - - -def make_builder(out_file, impl, vocab_size=None): - if impl == 'mmap': - return MMapIndexedDatasetBuilder(out_file, dtype=__best_fitting_dtype(vocab_size)) - else: - return IndexedDatasetBuilder(out_file) - - -def make_dataset(path, impl, skip_warmup=False): - if not IndexedDataset.exists(path): - print(f"Dataset does not exist: {path}") - print("Path should be a basename that both .idx and .bin can be appended to get full filenames.") - return None - if impl == 'infer': - impl = infer_dataset_impl(path) - if impl == 'lazy' and IndexedDataset.exists(path): - return IndexedDataset(path) - elif impl == 'cached' and IndexedDataset.exists(path): - return IndexedCachedDataset(path) - elif impl == 'mmap' and MMapIndexedDataset.exists(path): - return MMapIndexedDataset(path, skip_warmup) - print(f"Unknown dataset implementation: {impl}") - return None - - -def dataset_exists(path, impl): - if impl == 'mmap': - return MMapIndexedDataset.exists(path) - else: - return IndexedDataset.exists(path) - - -def read_longs(f, n): - a = np.empty(n, dtype=np.int64) - f.readinto(a) - return a - - -def write_longs(f, a): - f.write(np.array(a, dtype=np.int64)) - - -# valid metric_dtypes as numpy and torch types -dtypes = { - 1: (np.uint8, torch.uint8), - 2: (np.int8, torch.int8), - 3: (np.int16, torch.int16), - 4: (np.int32, torch.int32), - 5: (np.int64, torch.int64), - 6: (np.uint16, None), - 7: (np.uint32, None), - 8: (np.uint64, None), -} - -valid_dtypes = set([dt[0] for dt in dtypes.values()] + [dt[1] for dt in dtypes.values() if dt[1] is not None]) - - -def code(dtype): - for c, (np_dt, torch_dt) in dtypes.items(): - if dtype in [np_dt, torch_dt]: - return c - raise ValueError(f"{dtype} not supported. Supported types: {valid_dtypes}") - - -def index_file_path(prefix_path): - return prefix_path + '.idx' - - -def data_file_path(prefix_path): - return prefix_path + '.bin' - - -def create_doc_idx(sizes): - doc_idx = [0] - for i, s in enumerate(sizes): - if s == 0: - doc_idx.append(i + 1) - return doc_idx - - -class IndexedDataset(torch.utils.data.Dataset): - """Loader for IndexedDataset""" - _HDR_MAGIC = b'TNTIDX\x00\x00' - - def __init__(self, path): - super().__init__() - self.path = path - self.data_file = None - self.read_index(path) - - def read_index(self, path): - with open(index_file_path(path), 'rb') as f: - magic = f.read(8) - assert magic == self._HDR_MAGIC, ('Index file doesn\'t match expected format. ' - 'Make sure that --dataset-impl is configured properly.') - version = f.read(8) - assert struct.unpack('= self._len: - raise IndexError('index out of range') - - def __del__(self): - if self.data_file: - self.data_file.close() - - # @lru_cache(maxsize=8) - def __getitem__(self, idx): - if not self.data_file: - self.read_data(self.path) - if isinstance(idx, int): - i = idx - self.check_index(i) - tensor_size = self.sizes[self.dim_offsets[i]:self.dim_offsets[i + 1]] - a = np.empty(tensor_size, dtype=self.dtype) - self.data_file.seek(self.data_offsets[i] * self.element_size) - self.data_file.readinto(a) - return a - elif isinstance(idx, slice): - start, stop, step = idx.indices(len(self)) - if step != 1: - raise ValueError("Slices into indexed_dataset must be contiguous") - sizes = self.sizes[self.dim_offsets[start]:self.dim_offsets[stop]] - size = sum(sizes) - a = np.empty(size, dtype=self.dtype) - self.data_file.seek(self.data_offsets[start] * self.element_size) - self.data_file.readinto(a) - offsets = list(accumulate(sizes)) - sents = np.split(a, offsets[:-1]) - return sents - - def __len__(self): - return self._len - - def num_tokens(self, index): - return self.sizes[index] - - def size(self, index): - return self.sizes[index] - - @staticmethod - def exists(path): - return (os.path.exists(index_file_path(path)) and os.path.exists(data_file_path(path))) - - @property - def supports_prefetch(self): - return False # avoid prefetching to save memory - - -class IndexedCachedDataset(IndexedDataset): - - def __init__(self, path): - super().__init__(path) - self.cache = None - self.cache_index = {} - - @property - def supports_prefetch(self): - return True - - def prefetch(self, indices): - if all(i in self.cache_index for i in indices): - return - if not self.data_file: - self.read_data(self.path) - indices = sorted(set(indices)) - total_size = 0 - for i in indices: - total_size += self.data_offsets[i + 1] - self.data_offsets[i] - self.cache = np.empty(total_size, dtype=self.dtype) - ptx = 0 - self.cache_index.clear() - for i in indices: - self.cache_index[i] = ptx - size = self.data_offsets[i + 1] - self.data_offsets[i] - a = self.cache[ptx:ptx + size] - self.data_file.seek(self.data_offsets[i] * self.element_size) - self.data_file.readinto(a) - ptx += size - if self.data_file: - # close and delete data file after prefetch so we can pickle - self.data_file.close() - self.data_file = None - - # @lru_cache(maxsize=8) - def __getitem__(self, idx): - if isinstance(idx, int): - i = idx - self.check_index(i) - tensor_size = self.sizes[self.dim_offsets[i]:self.dim_offsets[i + 1]] - a = np.empty(tensor_size, dtype=self.dtype) - ptx = self.cache_index[i] - np.copyto(a, self.cache[ptx:ptx + a.size]) - return a - elif isinstance(idx, slice): - # Hack just to make this work, can optimizer later if necessary - sents = [] - for i in range(*idx.indices(len(self))): - sents.append(self[i]) - return sents - - -class IndexedDatasetBuilder(object): - - def __init__(self, out_file, dtype=np.int32): - self.out_file = open(out_file, 'wb') - self.dtype = dtype - self.data_offsets = [0] - self.dim_offsets = [0] - self.sizes = [] - self.element_size = self.dtype().itemsize - self.doc_idx = [0] - - def add_item(self, tensor): - bytes = self.out_file.write(np.array(tensor.numpy(), dtype=self.dtype)) - self.data_offsets.append(self.data_offsets[-1] + bytes / self.element_size) - for s in tensor.size(): - self.sizes.append(s) - self.dim_offsets.append(self.dim_offsets[-1] + len(tensor.size())) - - def end_document(self): - self.doc_idx.append(len(self.sizes)) - - def merge_file_(self, another_file): - index = IndexedDataset(another_file) - assert index.dtype == self.dtype - - doc_offset = len(self.sizes) - - begin = self.data_offsets[-1] - for data_offset in index.data_offsets[1:]: - self.data_offsets.append(begin + data_offset) - self.sizes.extend(index.sizes) - begin = self.dim_offsets[-1] - for dim_offset in index.dim_offsets[1:]: - self.dim_offsets.append(begin + dim_offset) - self.doc_idx.extend((doc_offset + index.doc_idx)[1:]) - - with open(data_file_path(another_file), 'rb') as f: - while True: - data = f.read(1024) - if data: - self.out_file.write(data) - else: - break - - def finalize(self, index_file): - self.out_file.close() - index = open(index_file, 'wb') - index.write(b'TNTIDX\x00\x00') - index.write(struct.pack(' [0, 10, 30, 35] - if arr.size > 1: - arr[1:] = arr[:-1] - if arr.size > 0: - arr[0] = 0 - - -def get_pointers_with_total(sizes, elemsize, dtype): - """Return a numpy array of type np.dtype giving the byte offsets. - - Multiplies values in the sizes array by elemsize (bytes), - and then computes an exclusive scan to get byte offsets. - Returns the total number of bytes as second item in a tuple. - """ - - # scale values in sizes array by elemsize to get sizes in bytes - pointers = np.array(sizes, dtype=dtype) - pointers *= elemsize - np.cumsum(pointers, axis=0, out=pointers) - - # get total number of bytes from all sizes (last element) - bytes_last = pointers[-1] if len(sizes) > 0 else 0 - - # convert to byte offsets - exscan_from_cumsum_(pointers) - - return pointers, bytes_last - - -class MMapIndexedDataset(torch.utils.data.Dataset): - - class Index(object): - _HDR_MAGIC = b'MMIDIDX\x00\x00' - - @classmethod - def writer(cls, path, dtype): - - class _Writer(object): - - def __enter__(self): - self._file = open(path, 'wb') - - self._file.write(cls._HDR_MAGIC) - self._file.write(struct.pack('= 0: - if max_value <= 255: - return np.uint8 - elif max_value <= 65535: - return np.uint16 - elif max_value <= 4294967295: - return np.uint32 - else: - return np.uint64 - else: - if max_value <= 127 and min_value >= -128: - return np.int8 - elif max_value <= 32767 and min_value >= -32768: - return np.int16 - elif max_value <= 2147483647 and min_value >= -2147483648: - return np.int32 - else: - return np.int64 - - -def split_index(start_idx, end_idx, num_partitions): - partition_boundaries = np.linspace(start_idx, end_idx, dtype=int, num=num_partitions + 1) - return [(partition_boundaries[i], partition_boundaries[i + 1]) for i in range(num_partitions)] - - -def split_dataset(dataset, num_workers, worker_id, num_threads): - worker_splits = split_index(0, len(dataset), num_workers) - thread_splits = split_index(worker_splits[worker_id][0], worker_splits[worker_id][1], num_threads) - return worker_splits, thread_splits - - -def create_mmap_dataset_builder(fname, dtype): - logger.info(f"Creating mmap dataset builder at {fname}.") - return MMapIndexedDatasetBuilder(f"{fname}.bin", dtype=dtype) - - -def close_mmap_dataset_builder(builder, fname): - builder.end_document() - builder.finalize(f"{fname}.idx") - logger.info(f"Finalized mmap dataset builder at {fname}.") diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/variable_batch_size_and_lr.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/variable_batch_size_and_lr.py deleted file mode 100644 index c9a39bbc53b509479aa2c18f3ddfc2e0abdb7a3c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/data_pipeline/data_sampling/variable_batch_size_and_lr.py +++ /dev/null @@ -1,492 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# support/questions/maintenance: github user @brunomaga or @deepspeedai/deepspeed - -import random -import torch -import os -import numpy as np -from torch.optim.lr_scheduler import LRScheduler -from torch.optim.optimizer import Optimizer -from torch.utils.data import DataLoader, DistributedSampler -from deepspeed.utils import logger -from deepspeed.runtime.pipe.engine import PipelineEngine -from deepspeed.runtime.data_pipeline.constants import * -from deepspeed.runtime.data_pipeline.data_sampling.indexed_dataset import MMapIndexedDataset -from deepspeed.runtime.data_pipeline.data_sampling.data_analyzer import DistributedDataAnalyzer -import pathlib - - -def batch_by_seqlens( - seqlens, - max_tokens, - sequence_ids_per_mb=None, - min_batch_size=1, - max_batch_size=None, - sequence_picking_order="dataloader", - effective_batch_size=1, - required_microbatches_of_same_size=False, - verbose=False, - seed=None, -): - """ - Yield mini-batches of indices bucketed by size. Batches may contain sequences of different lengths. - Similar to "Attention is all you need", Section 5.1: - "sequence pairs were batched together by approximate sequence length. Each training batch - contained a set of sequence pairs containing approximately X source tokens and X target tokens" - - Arguments: - - `seqlens`: a list of difficulties (metric values) for every sample in the dataset; - - `max_tokens`: maximum cap in total difficulty in a batch; - - `min_batch_size`: smallest allowed size of a batch; - - `min_batch_size`: largest allowed size of a batch; - - `sequence_picking_order`: order in which to process samples: "dataloader" (default), "random" or "seqlen" (ascending) - - `effective_batch_size`: effective batch size; - - `required_microbatches_of_same_size`: enable if each mini-batch (in a total of `batch_size_multiple` - micro-batches per batch), should have all micro-batches with the same batch size ie the same - number of sequences. - - `verbose`: print debug information; - - `seed`: random seed for reproducibility; - - Returns: - - `microbatch_ids`: list of tuple of batch id and samples ids per microbatch - - `batch_sizes`: the effective batch size of each batch, used for to compute the scaled LR - - `batch_max_seqlens`: the max seqlen across all microbatches in a batch - """ - - assert sequence_picking_order in ["random", "seqlen", "dataloader"] - if sequence_ids_per_mb is None: - metrics = list(zip(seqlens, range(len(seqlens)))) # use all samples - else: - metrics = list(zip(np.array(seqlens)[sequence_ids_per_mb], sequence_ids_per_mb)) - - if sequence_picking_order == 'random': - metric_random = random.Random(seed) - metric_random.shuffle(metrics) - if sequence_picking_order == 'seqlen': - metrics = sorted(metrics) - - # go through metrics, warn user, and filter samples that alone exceed the max batch threshold - long_ids = [idx for val, idx in metrics if val > max_tokens] - if len(long_ids) > 0: - logger.warning(f"Data indices {long_ids} ignored as metrics exceed {max_tokens}.") - logger.info(f"Original dataset length: {len(metrics)}. New dataset length: {len(long_ids)}") - metrics = [m for m in metrics if m[1] not in long_ids] - - def is_microbatch_valid(metrics): - if min_batch_size and len(metrics) < min_batch_size: return False # insufficient sample count - if max_batch_size and len(metrics) > max_batch_size: return False # too many samples - if sum([m[0] for m in metrics]) > max_tokens: return False # exceeds max - return True - - # go through all samples and pack then in microbatches of metric sums below the threshold - # `required_microbatches_of_same_size` means all minibatches in a batch must be of equal size - equal_size_multiple = effective_batch_size if required_microbatches_of_same_size else 1 - microbatches = [] - batch_init = 0 - while batch_init < len(metrics): - - # we iterate over possible effective batch sizes (groups of microbatches of same size) - valid_batch_end = batch_init - for batch_end in range(batch_init + equal_size_multiple, len(metrics), equal_size_multiple): - - # attempt effective batch - batch = metrics[batch_init:batch_end] - - # pick interleaved samples for each microbatch to help with load balancing - # (in the ordered use case), and to replicate what the distributed sampler does. - mbs = [batch[b::equal_size_multiple] for b in range(equal_size_multiple)] - - # if they are all valid micro-batches, keep them until you find longer mbatches, if any - is_batch_valid = all([is_microbatch_valid(mb) for mb in mbs]) - if is_batch_valid: - valid_batch_end = batch_end - - if batch_init == valid_batch_end: break # last batch is not valid (size zero), so we are done - batch = metrics[batch_init:valid_batch_end] - mbs = [batch[b::equal_size_multiple] for b in range(equal_size_multiple)] - batch_init += sum([len(l) for l in mbs]) - microbatches += mbs - - # make sure we give the same number of (micro-)batches to each dataloader by trimming the dataset - assert len(microbatches) >= effective_batch_size, "not enough datapoints to create a single sample per dataloader" - microbatches = microbatches[:len(microbatches) - len(microbatches) % effective_batch_size] - - #compute the effective batch size for each microbatch. - batch_sizes, batch_max_seqlens, microbatch_ids = [], [], [] - for rank in range(0, len(microbatches), effective_batch_size): - batch_id = rank // effective_batch_size - mbs = microbatches[rank:rank + effective_batch_size] - # compute the number of samples (not tokens) in this batch (not microbatch) - n_sequences = sum([len(mb) for mb in mbs]) - # compute the longest sequence (as number of tokens) in this batch (not microbatch) - sequence_ids_per_mb = [[m[1] for m in metrics] for metrics in mbs] - sequence_lens_per_mb = [[m[0] for m in metrics] for metrics in mbs] - batch_max_seqlen = max([max(seqlens) for seqlens in sequence_lens_per_mb]) - batch_and_mb_ids = zip([batch_id] * effective_batch_size, sequence_ids_per_mb) - batch_sizes.append(n_sequences) - batch_max_seqlens.append(batch_max_seqlen) - microbatch_ids += batch_and_mb_ids - if verbose: - n_tokens_per_mb = [sum([m[0] for m in mb]) for mb in mbs] - n_sequences_per_mb = [len(mb) for mb in mbs] - assert all([n <= max_tokens for n in n_tokens_per_mb]), "size of microbatch exceeds max tokens" - logger.info( - f"Batch id {batch_id} contains in total {len(mbs)} microbatches or {n_sequences} sequences. "\ - f"n_sequences per microbatch {n_sequences_per_mb}. "\ - f"n_tokens per microbatch {n_tokens_per_mb}. "\ - f"sequence ids per microbatch: {sequence_ids_per_mb}. "\ - f"sequence lengths per microbatch: {sequence_lens_per_mb}.") - - # return the sample ids of each microbatch, and the batch sizes - assert len(batch_sizes) == len(microbatch_ids) // effective_batch_size - return microbatch_ids, batch_sizes, batch_max_seqlens - - -def scale_lr(base_batch_size, batch_size, base_lr=1, method="linear"): - """ given a reference lr and batch_size, compute the new LR for a given batch size """ - if method == "linear": - # Linear Scaling Rule: "When the minibatch size is multiplied by k, multiply the learning - # rate by k" (Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour, Goyal et al) - return base_lr * batch_size / base_batch_size - if method == "sqrt": - # Square Root scaling: "when multiplying the batch size by k, multiply the learning rate - # by √k, to keep the variance in the gradient expectation constant" - # (A. Krizhevsky. One weird trick for parallelizing convolutional neural networks) - return base_lr * torch.sqrt(batch_size / base_batch_size) - elif method == None or method.upper() == "NONE": - return base_lr - raise ValueError("Unknown scaling method: {}".format(method)) - - -def dataloader_for_variable_batch_size( - dataset, - microbatch_ids, - batch_max_seqlens, - dataloader_rank=0, - dataloader_batch_size=1, - dataloader_num_replicas=1, - dataloader_collate_fn=None, - dataloader_num_workers=2, - dataloader_pin_memory=False, - required_microbatches_of_same_seqlen=False, - sample_padding_fn=None, -): - - # equidistantly distribute the microbatches across the replicas in an interleaved fashion. - sampler = DistributedSampler( - dataset=microbatch_ids, - num_replicas=dataloader_num_replicas, - rank=dataloader_rank, - shuffle=False, - drop_last=False, - ) - - # collate function wraps user-defined collate function to the variable batch data - def collate_fn_wrapper(list_microbatch_ids): - # each batch is a list of sample ids that fill up to the max tokens per batch - # we return the collated batch of all dataset samples of all input batches. - batch = [] - for batch_id, microbatch_ids in list_microbatch_ids: - batch_data = [dataset[idx] for idx in microbatch_ids] - if required_microbatches_of_same_seqlen: - assert sample_padding_fn is not None, \ - "padding dataloader_padding_fn must be provided if required_microbatches_of_same_seqlen is True" - max_seqlen = batch_max_seqlens[batch_id] - assert all([len(sample) <= max_seqlen for sample in batch_data]), \ - "some samples are longer than the computed max seqlen for the batch those samples belong to" - batch_data = [sample_padding_fn(sample, max_seqlen) for sample in batch_data] - batch += batch_data - return dataloader_collate_fn(batch) if dataloader_collate_fn else batch - - dataloader = DataLoader( - dataset=microbatch_ids, - batch_size=dataloader_batch_size, - sampler=sampler, - num_workers=dataloader_num_workers, - collate_fn=collate_fn_wrapper, - pin_memory=dataloader_pin_memory, - ) - - deepspeed_io_kwargs = dict( - dataset=microbatch_ids, - batch_size=dataloader_batch_size, - pin_memory=dataloader_pin_memory, - data_sampler=sampler, - collate_fn=collate_fn_wrapper, - num_local_io_workers=dataloader_num_workers, - ) - - return dataloader, deepspeed_io_kwargs - - -class VariableBatchSizeLR(LRScheduler): - """ an LR scheduler that scales the LR of a given scheduler's LR """ - - @property - def optimizer(self): - return self.base_lr_scheduler.optimizer - - def __init__(self, - lr_scheduler, - base_batch_size, - batch_sizes, - dataloader, - lr_scaling_method="linear", - last_epoch=-1, - verbose=False): - self.batch_sizes = batch_sizes - self.base_batch_size = base_batch_size - self.lr_scaling_method = lr_scaling_method - self.dataloader = dataloader - self.base_lr_scheduler = lr_scheduler - # the following exist in LRScheduler but not in DeepSpeed's LRScheduler so we redefine them here - self.base_lrs = self.base_lr_scheduler.get_lr() - self.last_epoch = last_epoch - self.verbose = verbose - self.step(0) # scale LR for first sample in the dataloader - - def state_dict(self): - return { - 'base_lr_scheduler': self.base_lr_scheduler.state_dict() - } | { - 'base_batch_size': self.base_batch_size, - 'lr_scaling_method': self.lr_scaling_method, - 'batch_sizes': self.batch_sizes, - 'base_lrs': self.base_lrs, - 'last_epoch': self.last_epoch, - 'verbose': self.verbose, - } - - def load_state_dict(self, state_dict): - self.base_lr_scheduler.load_state_dict(state_dict['base_lr_scheduler']) - self.base_batch_size = state_dict['base_batch_size'] - self.lr_scaling_method = state_dict['lr_scaling_method'] - self.batch_sizes = state_dict['batch_sizes'] - self.base_lrs = state_dict['base_lrs'] - self.last_epoch = state_dict['last_epoch'] - self.verbose = state_dict['verbose'] - - def get_last_lr(self): - return self.base_lr_scheduler._last_lr - - def get_lr(self): - return [group['lr'] for group in self.base_lr_scheduler.optimizer.param_groups] - - def step(self, epoch=None): - # call the base scheduler's step method to get LR for next epoch - # Note: optimizer.step precedes lr_scheduler.step(), so the stepping workflow is: - # init: lr_scheduler.step(0) --> set LR for epoch 0 - # epoch 0: optimizer.step(); lr_scheduler.step(1) --> set LR for epoch 1 - # epoch 1: optimizer.step(); lr_scheduler.step(2) --> set LR for epoch 2 - - # reset unscaled LRs (to the original scheduler's one) to be able to step the base LR scheduler - # Note: epoch==0: reset LR scheduler; epoch==None: scale LR for next epoch; - unscaled_lrs = self.base_lrs if epoch == 0 else self.get_last_lr() - for group, lr in zip(self.base_lr_scheduler.optimizer.param_groups, unscaled_lrs): - group['lr'] = lr - - self.base_lr_scheduler.step(epoch) # set unscaled lr, _step_count, last_epoch, _last_lr for new epoch - - # scale the learning rate for the the next iteration for each parameter group. - self.last_epoch = self.last_epoch + 1 if epoch is None else epoch - # batch sizes are precomputed and stored in batch_sizes se we loop around to get the next one - batch_size = self.batch_sizes[self.last_epoch % len(self.batch_sizes)] - for group in self.base_lr_scheduler.optimizer.param_groups: - group['lr'] = scale_lr(self.base_batch_size, batch_size, group['lr'], self.lr_scaling_method) - - if self.verbose: - logger.info( - f"Next batch id {self.last_epoch}. "\ - f"Reference batch_size {self.base_batch_size} and lr {unscaled_lrs}. "\ - f"Scaled batch_size {batch_size} and lr {self.get_lr()}.") - - -def lr_scheduler_for_variable_batch_size(base_batch_size, - batch_sizes, - dataloader, - lr_scheduler_or_optimizer, - lr_scaling_method='linear', - verbose=False): - """ - returns a class that provides an LR scheduler that scales the learning rate at every - iteration taking into account the batch size of that iteration. - If learning rate is constant, ie no LR scheduler, then the base LR will be taken from the - constant LR values in the optimizer param groups. Otherwise from the scheduler's LR. - - Arguments: - - `base_batch_size`: the batch size that the base LR in the optimizer or scheduler refers to; - - `lr_scaling_method`: method to use to scale LR - see `scale_lr()`; - - `lr_scheduler_or_optimizer`: one instance of `LRScheduler` or `Optimizer` to be used as base; - - `batch_sizes`: the effective batch size of each batch in the dataloader; - - Returns the new LRScheduler - """ - - class StubLRScheduler(LRScheduler): - """ a stub LR scheduler that does not change the LR, keeps it constant """ - - def get_lr(self) -> float: - return self.base_lrs - - if isinstance(lr_scheduler_or_optimizer, Optimizer): - lr_scheduler = StubLRScheduler(lr_scheduler_or_optimizer) - elif hasattr(lr_scheduler_or_optimizer, 'optimizer'): #LRScheduler or DeepSpeed 'object' schedulers - assert isinstance(lr_scheduler_or_optimizer.optimizer, Optimizer) - lr_scheduler = lr_scheduler_or_optimizer - else: - raise ValueError("Unknown type for lr_scheduler_or_optimizer: {}".format(type(lr_scheduler_or_optimizer))) - - return VariableBatchSizeLR(lr_scheduler=lr_scheduler, - base_batch_size=base_batch_size, - batch_sizes=batch_sizes, - dataloader=dataloader, - lr_scaling_method=lr_scaling_method, - verbose=verbose) - - -def get_dataloader_and_lr_scheduler_for_variable_batch_size_deepspeed(dataset, - engine, - dataset_seqlens=None, - dataset_filter_ids=None, - dataloader_collate_fn=None, - sample_padding_fn=None, - batch_seqlens_fn=None): - """ - a simplified call to get_dataloader_and_lr_scheduler_for_variable_batch_size for the deepspeed runtime. - Needs the seqlens of every sample. It will try three alternatives: - - if `dataset_seqlens` is provided by user, use that. - - otherwise, looks for the seqlen metric path (in the connfig) that contains the output of the Data Analyzer - - otherwise, use the user-provided function `batch_seqlens_fn` and call Data Analyzer to output seqlen metric - See `batch_by_seqlens()` for arguments and more documentation. - """ - data_efficiency_config = engine._config.data_efficiency_config - data_sampling_config = data_efficiency_config[DATA_SAMPLING] - batching_config = data_sampling_config[DYNAMIC_BATCHING] - assert batching_config[DYNAMIC_BATCHING_ENABLED], "Dynamic batching is not enabled in the config" - - if dataset_seqlens is None: - # In seqlen provided by user, look for the seqlen metric that was output by the Data Analyzer - # (see the main in deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py for an example) - metrics_path = batching_config[DYNAMIC_BATCHING_METRICS_PATH] - sample_to_seqlen_path = os.path.join(metrics_path, "seqlen/seqlen_sample_to_metric") - if not (os.path.exists(f"{sample_to_seqlen_path}.bin") and os.path.exists(f"{sample_to_seqlen_path}.idx")): - # if the metric files are not found, we run the DataAnalyzer to write the metric files - msg = f"Cannot find metric files for sequence length in {sample_to_seqlen_path}.idx or *.bin." - msg += " We will run data analyzer to generated them..." - logger.warning(msg) - - if batch_seqlens_fn is None: - raise ValueError("sample_seqlen_fn must be provided if dataset_seqlens is not provided") - - DistributedDataAnalyzer( - dataset=dataset, - metric_functions=[batch_seqlens_fn], - collate_fn=dataloader_collate_fn, - batch_size=2**10, # batch size for map-reduce, not training - num_workers=engine.world_size, - worker_id=engine.global_rank, - save_path=pathlib.Path(metrics_path), - metric_types=['single_value_per_sample'], - metric_names=["seqlen"], - device=engine.device, - ).run_map_reduce() - - dataset_seqlens = MMapIndexedDataset(sample_to_seqlen_path, skip_warmup=True) - assert len(dataset_seqlens) == len(dataset), \ - "Seqlens size does not match the input dataset size. If you changed the dataset, delete the metrics_path folder." - - # TODO we are copying all seqlens into memory, we should adapt the code to use an iterative streamer - # and use the other files output by DataAnalyzer that returns an ordered dictionary of seqlen to sample ids - dataset_seqlens = np.array(list(dataset_seqlens), dtype=np.int64).flatten() # from Nx1 to N - - dataloader, lr_scheduler, deepspeed_io_kwargs = get_dataloader_and_lr_scheduler_for_variable_batch_size( - dataset=dataset, - dataset_filter_ids=dataset_filter_ids, - dataset_seqlens=dataset_seqlens, - effective_batch_size=engine.train_batch_size(), - max_tokens=batching_config[DYNAMIC_BATCHING_MAX_TOKENS], - lr_scaling_method=batching_config[DYNAMIC_BATCHING_LR_SCALING_METHOD], - sequence_picking_order=batching_config[DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER], - min_batch_size=batching_config[DYNAMIC_BATCHING_MIN_BATCH_SIZE], - max_batch_size=batching_config[DYNAMIC_BATCHING_MAX_BATCH_SIZE], - dataloader_batch_size=engine.train_micro_batch_size_per_gpu(), - dataloader_rank=engine.data_parallel_group.rank(), - dataloader_num_replicas=engine.data_parallel_group.size(), - dataloader_num_workers=data_sampling_config[DATA_SAMPLING_NUM_WORKERS], - dataloader_collate_fn=dataloader_collate_fn, - dataloader_pin_memory=data_sampling_config[DATA_SAMPLING_PIN_MEMORY], - sample_padding_fn=sample_padding_fn, - lr_scheduler_or_optimizer=engine.lr_scheduler or engine.optimizer, - required_microbatches_of_same_size=isinstance(engine, PipelineEngine), - required_microbatches_of_same_seqlen=isinstance(engine, PipelineEngine), - verbose=batching_config[DYNAMIC_BATCHING_VERBOSE], - seed=data_efficiency_config[DATA_EFFICIENCY_SEED], - ) - return dataloader, lr_scheduler, deepspeed_io_kwargs - - -def get_dataloader_and_lr_scheduler_for_variable_batch_size( - dataset, - dataset_seqlens, - max_tokens, - effective_batch_size, - dataset_filter_ids=None, - lr_scaling_method="linear", - min_batch_size=1, - max_batch_size=None, - sequence_picking_order="dataloader", - dataloader_batch_size=1, - dataloader_rank=0, - dataloader_num_replicas=1, - dataloader_num_workers=0, - dataloader_collate_fn=None, - dataloader_pin_memory=False, - lr_scheduler_or_optimizer=None, - required_microbatches_of_same_size=False, - required_microbatches_of_same_seqlen=False, - sample_padding_fn=None, - verbose=False, - seed=None, -): - """ returns a dataloader and LR scheduler for the variable batch size. see `batch_by_seqlens()` for details. """ - - # effective_batch_size = train_micro_batch_size_per_gpu * gradient_accumulation_steps * number of dataloaders - microbatch_ids, batch_sizes, batch_max_seqlens = batch_by_seqlens( - seqlens=dataset_seqlens, - max_tokens=max_tokens, - sequence_ids_per_mb=dataset_filter_ids, - min_batch_size=min_batch_size, - max_batch_size=max_batch_size, - sequence_picking_order=sequence_picking_order, - effective_batch_size=effective_batch_size, - required_microbatches_of_same_size=required_microbatches_of_same_size, - verbose=verbose, - seed=seed, - ) - - dataloader, deepspeed_io_kwargs = dataloader_for_variable_batch_size( - dataset=dataset, - microbatch_ids=microbatch_ids, - batch_max_seqlens=batch_max_seqlens, - dataloader_rank=dataloader_rank, - dataloader_num_replicas=dataloader_num_replicas, - dataloader_batch_size=dataloader_batch_size, - dataloader_collate_fn=dataloader_collate_fn, - dataloader_num_workers=dataloader_num_workers, - dataloader_pin_memory=dataloader_pin_memory, - required_microbatches_of_same_seqlen=required_microbatches_of_same_seqlen, - sample_padding_fn=sample_padding_fn, - ) - - lr_scheduler = lr_scheduler_for_variable_batch_size(base_batch_size=effective_batch_size, - batch_sizes=batch_sizes, - lr_scaling_method=lr_scaling_method, - lr_scheduler_or_optimizer=lr_scheduler_or_optimizer, - dataloader=dataloader, - verbose=verbose) - - return dataloader, lr_scheduler, deepspeed_io_kwargs diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/dataloader.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/dataloader.py deleted file mode 100644 index 499473b4ced81ba5ae5a447e32e03a2971b63c8f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/dataloader.py +++ /dev/null @@ -1,162 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from torch.utils.data import DataLoader, RandomSampler -from torch.utils.data.distributed import DistributedSampler -from deepspeed.accelerator import get_accelerator - -from deepspeed.runtime.data_pipeline.data_sampling.data_sampler import DeepSpeedDataSampler -from deepspeed.runtime.data_pipeline.constants import CURRICULUM_LEARNING, \ - DATA_EFFICIENCY, DATA_SAMPLING_NUM_WORKERS -from deepspeed.runtime.constants import GRADIENT_ACCUMULATION_STEPS, \ - DATA_PARALLEL_GROUP, GLOBAL_RANK - - -class RepeatingLoader: - - def __init__(self, loader): - """Wraps an iterator to allow for infinite iteration. This is especially useful - for DataLoader types that we wish to automatically restart upon completion. - - Args: - loader (iterator): The data loader to repeat. - """ - self.loader = loader - self.data_iter = iter(self.loader) - - def __iter__(self): - return self - - def __next__(self): - try: - batch = next(self.data_iter) - except StopIteration: - self.data_iter = iter(self.loader) - batch = next(self.data_iter) - return batch - - -class DeepSpeedDataLoader(object): - - def __init__(self, - dataset, - batch_size, - pin_memory, - local_rank, - tput_timer, - collate_fn=None, - num_local_io_workers=None, - data_sampler=None, - data_parallel_world_size=None, - data_parallel_rank=None, - dataloader_drop_last=False, - deepspeed_dataloader_config={}): - self.deepspeed_dataloader_config = deepspeed_dataloader_config - self.tput_timer = tput_timer - self.batch_size = batch_size - self.curriculum_learning_enabled = False - if CURRICULUM_LEARNING in deepspeed_dataloader_config: - self.curriculum_learning_enabled = deepspeed_dataloader_config[CURRICULUM_LEARNING] - - if self.curriculum_learning_enabled: - data_sampler = DeepSpeedDataSampler(self.deepspeed_dataloader_config[DATA_EFFICIENCY], - len(dataset), - self.batch_size, - data_parallel_rank, - data_parallel_world_size, - self.deepspeed_dataloader_config[DATA_PARALLEL_GROUP], - self.deepspeed_dataloader_config[GRADIENT_ACCUMULATION_STEPS], - self.deepspeed_dataloader_config[GLOBAL_RANK], - drop_last=dataloader_drop_last) - device_count = get_accelerator().device_count() - num_local_io_workers = self.deepspeed_dataloader_config[DATA_SAMPLING_NUM_WORKERS] - else: - if local_rank >= 0: - if data_sampler is None: - data_sampler = DistributedSampler(dataset=dataset, - num_replicas=data_parallel_world_size, - rank=data_parallel_rank) - device_count = 1 - else: - if data_sampler is None: - data_sampler = RandomSampler(dataset) - device_count = get_accelerator().device_count() - batch_size *= device_count - - if num_local_io_workers is None: - num_local_io_workers = 2 * device_count - - self.num_local_io_workers = num_local_io_workers - self.data_sampler = data_sampler - self.dataset = dataset - self.collate_fn = collate_fn - self.device_count = device_count - self.batch_size = batch_size - self.pin_memory = pin_memory - self.data = None - self.dataloader_drop_last = dataloader_drop_last - self.post_process_func = None - - if self.dataloader_drop_last: - self.len = len(self.data_sampler) // self.batch_size - else: - from math import ceil - self.len = ceil(len(self.data_sampler) / self.batch_size) - - def __iter__(self): - self._create_dataloader() - return self - - def __len__(self): - return self.len - - def __next__(self): - if self.tput_timer: - self.tput_timer.start() - if self.curriculum_learning_enabled: - data = next(self.data_iterator) - if self.post_process_func is not None: - data = self.post_process_func(data, self.data_sampler.state_dict()) - return data - else: - return next(self.data) - - def _create_dataloader(self): - if self.curriculum_learning_enabled: - if self.collate_fn is None: - self.dataloader = DataLoader(self.dataset, - pin_memory=self.pin_memory, - batch_sampler=self.data_sampler, - num_workers=self.num_local_io_workers) - else: - self.dataloader = DataLoader(self.dataset, - pin_memory=self.pin_memory, - batch_sampler=self.data_sampler, - collate_fn=self.collate_fn, - num_workers=self.num_local_io_workers) - self.data_iterator = iter(self.dataloader) - return self.dataloader - else: - if self.collate_fn is None: - self.dataloader = DataLoader(self.dataset, - batch_size=self.batch_size, - pin_memory=self.pin_memory, - sampler=self.data_sampler, - num_workers=self.num_local_io_workers, - drop_last=self.dataloader_drop_last) - else: - self.dataloader = DataLoader(self.dataset, - batch_size=self.batch_size, - pin_memory=self.pin_memory, - sampler=self.data_sampler, - collate_fn=self.collate_fn, - num_workers=self.num_local_io_workers, - drop_last=self.dataloader_drop_last) - self.data = (x for x in self.dataloader) - - return self.dataloader - - -# DataLoader([(torch.randn(3, 3), torch.tensor(i % 2)) for i in range(10)], batch_size=2)) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__init__.py deleted file mode 100644 index 208299fb8c50f73468d293b6fa5dca71649d62e7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 84896d3405ff4c4ffa848a370325c25006f66a8a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/async_linear.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/async_linear.cpython-310.pyc deleted file mode 100644 index 6e4848577a040c5205a63d6d9ed99b2c38d8e898..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/async_linear.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/transformer.cpython-310.pyc deleted file mode 100644 index fde5ed460d63794267dd092b8dfeeb92fc4342c0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/__pycache__/transformer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/async_linear.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/async_linear.py deleted file mode 100644 index 8e01da500409cd603c17af646e902cfe2a82525b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/async_linear.py +++ /dev/null @@ -1,140 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Adapted from https://github.com/NVIDIA/Megatron-LM/blob/23.08/megatron/core/tensor_parallel/layers.py - -import torch -from torch.nn.parameter import Parameter -import torch.nn.functional as F -from deepspeed.accelerator import get_accelerator -import deepspeed.comm as dist -from typing import Callable - -TP_group = None - - -class DominoAsyncColumnParallelLinearImpl(torch.autograd.Function): - - @staticmethod - def forward(ctx, inp, weight, bias, handle_dic, h_id): # inp: (b, s, k), weight: (m, k), bias (m) - ctx.save_for_backward(inp, weight, bias) - ctx.handle_dic = handle_dic - ctx.h_id = h_id - output = torch.matmul(inp, weight.t()) # (b, s, k) @ (k, m) -> (b, s, m) - if bias is not None: # bias (m) - output = output + bias - return output - - @staticmethod - def backward(ctx, grad_output): - inp, weight, bias = ctx.saved_tensors - grad_input = grad_weight = grad_bias = None - grad_input = torch.matmul(grad_output, weight) # (b, s, m) @ (m, k) -> (b, s, k) - handle = dist.all_reduce(grad_input, group=TP_group, async_op=True) - ctx.handle_dic[ctx.h_id] = handle - grad_output = grad_output.view(grad_output.shape[0] * grad_output.shape[1], grad_output.shape[2]) # (b*s, m) - - inp = inp.view(inp.shape[0] * inp.shape[1], inp.shape[2]) # (b*s, k) - grad_weight = torch.matmul(grad_output.t(), inp) # (m, b*s) @ (b*s, k) -> (m, k) - - if bias is not None: - grad_bias = grad_output.sum(dim=0) # (b*s, m) -> (m) - return grad_input, grad_weight, grad_bias, None, None - - -class DominoAsyncColumnParallelLinear(torch.nn.Module): - - def __init__(self, - input_size, - output_size, - _tp_group, - config, - init_method: Callable, - bias=True, - skip_bias_add=False): - super(DominoAsyncColumnParallelLinear, self).__init__() - - self.skip_bias_add = skip_bias_add - - global TP_group - if TP_group == None: - TP_group = _tp_group - - self.weight = Parameter( - torch.empty( - output_size, - input_size, - device=get_accelerator().current_device_name(), - dtype=config.params_dtype, - )) - if config.perform_initialization: - init_method(self.weight) - - if bias: - self.bias = Parameter( - torch.empty(output_size, device=get_accelerator().current_device_name(), dtype=config.params_dtype)) - - if config.perform_initialization: - with torch.no_grad(): - self.bias.zero_() - else: - self.register_parameter('bias', None) - - def forward(self, input_: torch.Tensor, handle_dic, h_id): - - bias = self.bias if not self.skip_bias_add else None - - output = DominoAsyncColumnParallelLinearImpl.apply(input_, self.weight, bias, handle_dic, h_id) - - output_bias = self.bias if self.skip_bias_add else None - return output, output_bias - - -class RowParallelLinearNoComm(torch.nn.Module): - - def __init__( - self, - input_size: int, - output_size: int, - config, - init_method: Callable, - bias: bool = True, - stride: int = 1, - skip_bias_add: bool = False, - ): - super(RowParallelLinearNoComm, self).__init__() - - self.skip_bias_add = skip_bias_add - - self.weight = Parameter( - torch.empty( - output_size, - input_size, - device=get_accelerator().current_device_name(), - dtype=config.params_dtype, - )) - if config.perform_initialization: - init_method(self.weight) - if bias: - self.bias = Parameter( - torch.empty( - output_size, - device=get_accelerator().current_device_name(), - dtype=config.params_dtype, - )) - - if config.perform_initialization: - with torch.no_grad(): - self.bias.zero_() - else: - self.register_parameter('bias', None) - - def forward(self, input_): - bias = self.bias if not self.skip_bias_add else None - - output = F.linear(input_, self.weight, bias) - - output_bias = self.bias if self.skip_bias_add else None - return output, output_bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/transformer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/transformer.py deleted file mode 100644 index 3dfb133373b5aa7686f1dd9aac903add10b4e84d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/domino/transformer.py +++ /dev/null @@ -1,605 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import torch.nn.functional as F -import enum -import deepspeed.comm as dist - -from .async_linear import DominoAsyncColumnParallelLinear, RowParallelLinearNoComm - - -class LayerType(enum.Enum): - encoder = 1 - decoder = 2 - - -class AttnType(enum.Enum): - self_attn = 1 - cross_attn = 2 - - -class AttnMaskType(enum.Enum): - padding = 1 - causal = 2 - - -class ModelType(enum.Enum): - encoder_or_decoder = 1 - encoder_and_decoder = 2 - - -class DominoUtil: - - BATCH_0 = "BATCH0" - - BATCH_1 = "BATCH1" - - HANDLE_DIC = {"BATCH0": None, "BATCH1": None} - - -class DominoModule(torch.nn.Module): - """extensions of torch Module.""" - - def __init__(self, ): - super(DominoModule, self).__init__() - - -def _Wait_bwd_comm(input_, dic_, h_id): - return NoOper.apply(input_, dic_, h_id) - - -class NoOper(torch.autograd.Function): - - @staticmethod - def symbolic(graph, input_, handle_dic, h_id): - return input_ - - @staticmethod - def forward(ctx, input_, handle_dic, h_id): - ctx.handle_dic = handle_dic - ctx.h_id = h_id - return input_ - - @staticmethod - def backward(ctx, grad_output): - handle = ctx.handle_dic[ctx.h_id] - handle.wait() - return grad_output, None, None - - -class CoreAttention(DominoModule): - - def __init__(self, config, tp_world_size, attn_mask_type=AttnMaskType.causal): - super(CoreAttention, self).__init__() - - self.attn_mask_type = attn_mask_type - - projection_size = config.kv_channels * config.num_attention_heads - - # Per attention head and per partition values. - assert projection_size % tp_world_size == 0, f"projection size {projection_size} should be multiple of TP world size {tp_world_size}" - self.hidden_size_per_partition = projection_size // tp_world_size - self.attention_dropout_rate = config.attention_dropout - - def forward(self, query_layer, key_layer, value_layer, attention_mask): - - context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, - key_layer, - value_layer, - attn_mask=None, - dropout_p=self.attention_dropout_rate, - is_causal=True, - scale=None) - - # [b, np, sq, hn] --> [sq, b, np, hn] - context_layer = context_layer.permute(2, 0, 1, 3).contiguous() - - # [sq, b, np, hn] --> [sq, b, hp] - new_context_layer_shape = context_layer.size()[:-2] + \ - (self.hidden_size_per_partition,) - context_layer = context_layer.view(*new_context_layer_shape) - - return context_layer - - -class ShardedAttention(DominoModule): - """Sharded self-attention layer class. - Only support self attention and causal attention mask for now. - """ - - def __init__(self, - config, - mpu, - apply_rotary_pos_emb, - layer_number, - attention_type=AttnType.self_attn, - attn_mask_type=AttnMaskType.causal): - super(ShardedAttention, self).__init__() - - assert attention_type == AttnType.self_attn, "Only support self_attn for now!" - - self.layer_number = max(1, layer_number) - self.attention_type = attention_type - self.attn_mask_type = attn_mask_type - self.params_dtype = config.params_dtype - self.apply_rotary_pos_emb = apply_rotary_pos_emb - - query_projection_size = config.kv_channels * config.num_attention_heads - kv_projection_size = config.kv_channels * config.num_attention_heads - - tp_world_size = mpu.get_tensor_model_parallel_world_size() - self.hidden_size_per_attention_head = query_projection_size // config.num_attention_heads - self.num_attention_heads_per_partition = config.num_attention_heads // tp_world_size - - qkv_projection_per_partition = (query_projection_size + 2 * kv_projection_size) // tp_world_size - - self.query_key_value = DominoAsyncColumnParallelLinear(config.hidden_size, - qkv_projection_per_partition, - mpu.get_tensor_model_parallel_group(), - config=config, - init_method=config.init_method, - bias=config.add_bias_linear) - - self.core_attention = CoreAttention(config, tp_world_size, self.attn_mask_type) - - query_projection_size_per_partition = query_projection_size // tp_world_size - - # Output. - self.dense = RowParallelLinearNoComm(query_projection_size_per_partition, - config.hidden_size, - config=config, - init_method=config.output_layer_init_method, - bias=config.add_bias_linear, - skip_bias_add=True) - - def forward(self, hidden_states, attention_mask, micro_batch_num, rotary_pos_emb=None): - # hidden_states: [sq, b, h] - - mixed_x_layer, _ = self.query_key_value(hidden_states, DominoUtil.HANDLE_DIC, micro_batch_num) - - new_tensor_shape = mixed_x_layer.size()[:-1] + ( - self.num_attention_heads_per_partition, - 3 * self.hidden_size_per_attention_head, - ) - - mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) - - mixed_x_layer = mixed_x_layer.permute(1, 2, 0, 3).contiguous() - - (query_layer, key_layer, value_layer) = torch.split(mixed_x_layer, [ - self.hidden_size_per_attention_head, self.hidden_size_per_attention_head, - self.hidden_size_per_attention_head - ], - dim=3) - - query_layer = query_layer.view(query_layer.size(0), query_layer.size(1), -1, - self.hidden_size_per_attention_head) - - if rotary_pos_emb is not None: - if isinstance(rotary_pos_emb, tuple): - rotary_pos_emb = rotary_pos_emb - else: - rotary_pos_emb = ((rotary_pos_emb, ) * 2) - q_pos_emb, k_pos_emb = rotary_pos_emb - query_layer = self.apply_rotary_pos_emb(query_layer, q_pos_emb) - key_layer = self.apply_rotary_pos_emb(key_layer, k_pos_emb) - - context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask) - - output, bias = self.dense(context_layer) - return output, bias - - def domino_core_attention_forward(self, mixed_x_layer, attention_mask, rotary_pos_emb=None): - # hidden_states: [sq, b, h] - - # To illustrate the difference between intra-layer overlap and inter-layer overlap - # mixed_x_layer, _ = self.query_key_value(hidden_states, handle_dic, micro_batch_num) - - new_tensor_shape = mixed_x_layer.size()[:-1] + ( - self.num_attention_heads_per_partition, - 3 * self.hidden_size_per_attention_head, - ) - - mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) - - mixed_x_layer = mixed_x_layer.permute(1, 2, 0, 3).contiguous() - - (query_layer, key_layer, value_layer) = torch.split(mixed_x_layer, [ - self.hidden_size_per_attention_head, self.hidden_size_per_attention_head, - self.hidden_size_per_attention_head - ], - dim=3) - - query_layer = query_layer.view(query_layer.size(0), query_layer.size(1), -1, - self.hidden_size_per_attention_head) - - if rotary_pos_emb is not None: - if isinstance(rotary_pos_emb, tuple): - rotary_pos_emb = rotary_pos_emb - else: - rotary_pos_emb = ((rotary_pos_emb, ) * 2) - q_pos_emb, k_pos_emb = rotary_pos_emb - query_layer = self.apply_rotary_pos_emb(query_layer, q_pos_emb) - key_layer = self.apply_rotary_pos_emb(key_layer, k_pos_emb) - - context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask) - - # output, bias = self.dense(context_layer) - # return output, bias - - return context_layer - - -class bias_dropout_add(torch.nn.Module): - - def __init__(self, prob: float): - super(bias_dropout_add, self).__init__() - self.dropout = torch.nn.Dropout(prob) - - def forward(self, x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: - if bias is not None: - x = x + bias - out = self.dropout(x) - out = out + residual - return out - - -class DominoTransformerLayer(DominoModule): - """A domino single transformer layer. - [s, b, h] -> [s, b, h] - """ - - def __init__(self, - config, - mpu, - apply_rotary_pos_emb, - layer_number, - layer_type=LayerType.encoder, - self_attn_mask_type=AttnMaskType.causal, - drop_path_rate=0.): - - super(DominoTransformerLayer, self).__init__() - self.layer_number = layer_number - self.layer_type = layer_type - - self.apply_residual_connection_post_layernorm \ - = config.apply_residual_connection_post_layernorm - - self.llama_model = False - - self.input_layernorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layernorm_epsilon) - - # Self attention. - self.self_attention = ShardedAttention(config, - mpu, - apply_rotary_pos_emb, - layer_number, - attention_type=AttnType.self_attn, - attn_mask_type=self_attn_mask_type) - - self.hidden_dropout = config.hidden_dropout - - self.post_attention_layernorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layernorm_epsilon) - - # MLP - ffn_hidden_size = config.ffn_hidden_size - if config.gated_linear_unit: - ffn_hidden_size *= 2 - - self.output_size_c = config.ffn_hidden_size - self.input_size_c = config.hidden_size - self.input_size_r = config.ffn_hidden_size - self.output_size_r = self.input_size_c - - tp_world_size = mpu.get_tensor_model_parallel_world_size() - self.TP_group = mpu.get_tensor_model_parallel_group() - self.output_size_per_partition = self.output_size_c // tp_world_size - self.input_size_per_partition = self.input_size_r // tp_world_size - - self.linear_fc1 = DominoAsyncColumnParallelLinear(self.input_size_c, - self.output_size_per_partition, - mpu.get_tensor_model_parallel_group(), - config=config, - init_method=config.init_method, - bias=config.add_bias_linear) - - self.mlp_activation_func = F.gelu - - self.linear_fc2 = RowParallelLinearNoComm(self.input_size_per_partition, - self.output_size_r, - config=config, - init_method=config.output_layer_init_method, - bias=config.add_bias_linear, - skip_bias_add=True) - - self.bias_dropout_add_func = bias_dropout_add(self.hidden_dropout) - - def forward(self, hidden_states, attention_mask, rotary_pos_emb=None): - - hidden_states0, hidden_states1 = hidden_states - - layernorm_output0 = self.input_layernorm(hidden_states0) - layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) - - # Micro batch 0: attention - attention_output0, attention_bias0 = self.self_attention(layernorm_output0, - attention_mask, - DominoUtil.BATCH_0, - rotary_pos_emb=rotary_pos_emb) - - fwd_handle0 = dist.all_reduce(attention_output0, group=self.TP_group, async_op=True) - # End of Micro batch 0: attention - - # Micro batch 1: attention - layernorm_output1 = self.input_layernorm(hidden_states1) - layernorm_output1 = _Wait_bwd_comm(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) - - attention_output1, attention_bias1 = self.self_attention(layernorm_output1, - attention_mask, - DominoUtil.BATCH_1, - rotary_pos_emb=rotary_pos_emb) - fwd_handle1 = dist.all_reduce(attention_output1, group=self.TP_group, async_op=True) - - # Micro batch 0: Residual connection. - fwd_handle0.wait() - if self.apply_residual_connection_post_layernorm: - residual0 = layernorm_output0 - else: - residual0 = hidden_states0 - - layernorm_input0 = self.bias_dropout_add_func(attention_output0, attention_bias0, residual0) - - layernorm_output0 = self.post_attention_layernorm(layernorm_input0) - layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) - - if self.apply_residual_connection_post_layernorm: - residual0 = layernorm_output0 - else: - residual0 = layernorm_input0 - # End of Micro batch 0: Residual connection. - - # ------------ MLP ------------ - # Micro batch 0: MLP - output0, _ = self.linear_fc1(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) - output0 = self.mlp_activation_func(output0) - - # Micro batch 1: Residual connection. - fwd_handle1.wait() - if self.apply_residual_connection_post_layernorm: - residual1 = layernorm_output1 - else: - residual1 = hidden_states1 - - layernorm_input1 = self.bias_dropout_add_func(attention_output1, attention_bias1, residual1) - - layernorm_output1 = self.post_attention_layernorm(layernorm_input1) - layernorm_output1 = _Wait_bwd_comm(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) - - if self.apply_residual_connection_post_layernorm: - residual1 = layernorm_output1 - else: - residual1 = layernorm_input1 - # End of Micro batch 1: Residual connection. - - hidden_states0, last_mlp_bias = self.linear_fc2(output0) - fwd_handle0 = dist.all_reduce(hidden_states0, group=self.TP_group, async_op=True) - # End of Micro batch 0: MLP - - # Micro batch 1: MLP - output1, _ = self.linear_fc1(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) - output1 = self.mlp_activation_func(output1) - - hidden_states1, last_mlp_bias = self.linear_fc2(output1) - - fwd_handle1 = dist.all_reduce(hidden_states1, group=self.TP_group, async_op=True) - # End of Micro batch 1: MLP - - # ------------ End of MLP ------------ - - fwd_handle0.wait() - hidden_states0 = self.bias_dropout_add_func(hidden_states0, last_mlp_bias, residual0) - - fwd_handle1.wait() - hidden_states1 = self.bias_dropout_add_func(hidden_states1, last_mlp_bias, residual1) - - return hidden_states0, hidden_states1 - - -class DominoTransformer(DominoModule): - """Transformer class.""" - - def __init__(self, - config, - mpu, - apply_rotary_pos_emb, - model_type, - layer_type=LayerType.encoder, - self_attn_mask_type=AttnMaskType.causal, - post_layer_norm=True, - pre_process=True, - post_process=True, - drop_path_rate=0.0): - super(DominoTransformer, self).__init__() - - self.layer_type = layer_type - self.model_type = model_type - self.post_layer_norm = post_layer_norm - self.post_process = post_process - self.input_tensor = None - self.drop_path_rate = drop_path_rate - self.TP_group = mpu.get_tensor_model_parallel_group() - - if not dist.is_initialized(): - dist.init_distributed() - assert dist.is_initialized(), "deepspeed.comm failed to initialize!" - - self.num_layers = config.num_layers - - self.drop_path_rates = [rate.item() for rate in torch.linspace(0, self.drop_path_rate, config.num_layers)] - - def build_layer(layer_number): - - current_layer_type = layer_type - return DominoTransformerLayer(config, - mpu, - apply_rotary_pos_emb, - layer_number, - layer_type=current_layer_type, - self_attn_mask_type=self_attn_mask_type, - drop_path_rate=self.drop_path_rates[layer_number - 1]) - - self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)]) - - if self.post_process and self.post_layer_norm: - self.final_layernorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layernorm_epsilon) - - self._forward_impl = self.inter_layer_overlap_forward - if config.domino_intra_layer_overlap: - self._forward_impl = self.intra_layer_overlap_forward - - def forward(self, hidden_states, attention_mask, rotary_pos_emb=None): - - return self._forward_impl(hidden_states, attention_mask, rotary_pos_emb) - - def inter_layer_overlap_forward(self, hidden_states, attention_mask, rotary_pos_emb=None): - # hidden_states: [s, b, h] - - hidden_states0, hidden_states1 = torch.chunk(hidden_states, chunks=2, dim=1) - - last_mlp_bias = None - fwd_handle0, fwd_handle1 = None, None - residual0, residual1 = None, None - - layernorm_output0 = self.layers[0].input_layernorm(hidden_states0) - layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) - - for index in range(self.num_layers): - - # Micro batch 0: attention - attention_output0, _ = self.layers[index].self_attention.query_key_value( - layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) - attention_output0 = self.layers[index].self_attention.domino_core_attention_forward( - attention_output0, attention_mask, rotary_pos_emb=rotary_pos_emb) - - # Micro batch 1: Residual connection - if index > 0: - fwd_handle1.wait() - hidden_states1 = self.layers[index - 1].bias_dropout_add_func(hidden_states1, last_mlp_bias, residual1) - - layernorm_output1 = self.layers[index].input_layernorm(hidden_states1) - layernorm_output1 = _Wait_bwd_comm(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) - # End of Micro batch 1: Residual connection - - attention_output0, attention_bias0 = self.layers[index].self_attention.dense(attention_output0) - - fwd_handle0 = dist.all_reduce(attention_output0, group=self.TP_group, async_op=True) - # End of Micro batch 0: attention - - # Micro batch 1: attention - attention_output1, _ = self.layers[index].self_attention.query_key_value( - layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) - attention_output1 = self.layers[index].self_attention.domino_core_attention_forward( - attention_output1, attention_mask, rotary_pos_emb=rotary_pos_emb) - - # Micro batch 0: Residual connection. - fwd_handle0.wait() - if self.layers[index].apply_residual_connection_post_layernorm: - residual0 = layernorm_output0 - else: - residual0 = hidden_states0 - - layernorm_input0 = self.layers[index].bias_dropout_add_func(attention_output0, attention_bias0, residual0) - - layernorm_output0 = self.layers[index].post_attention_layernorm(layernorm_input0) - layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) - - if self.layers[index].apply_residual_connection_post_layernorm: - residual0 = layernorm_output0 - else: - residual0 = layernorm_input0 - # End of Micro batch 0: Residual connection. - - attention_output1, attention_bias1 = self.layers[index].self_attention.dense(attention_output1) - fwd_handle1 = dist.all_reduce(attention_output1, group=self.TP_group, async_op=True) - # End of Micro batch 1: attention - - # ------------ MLP ------------ - # Micro batch 0: MLP - output0, _ = self.layers[index].linear_fc1(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) - output0 = self.layers[index].mlp_activation_func(output0) - - # Micro batch 1: Residual connection. - fwd_handle1.wait() - if self.layers[index].apply_residual_connection_post_layernorm: - residual1 = layernorm_output1 - else: - residual1 = hidden_states1 - - layernorm_input1 = self.layers[index].bias_dropout_add_func(attention_output1, attention_bias1, residual1) - - layernorm_output1 = self.layers[index].post_attention_layernorm(layernorm_input1) - layernorm_output1 = _Wait_bwd_comm(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) - - if self.layers[index].apply_residual_connection_post_layernorm: - residual1 = layernorm_output1 - else: - residual1 = layernorm_input1 - # End of Micro batch 1: Residual connection. - - hidden_states0, last_mlp_bias = self.layers[index].linear_fc2(output0) - fwd_handle0 = dist.all_reduce(hidden_states0, group=self.TP_group, async_op=True) - # End of Micro batch 0: MLP - - # Micro batch 1: MLP - output1, _ = self.layers[index].linear_fc1(layernorm_output1, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_1) - output1 = self.layers[index].mlp_activation_func(output1) - - # Micro batch 0: Residual connection. - fwd_handle0.wait() - hidden_states0 = self.layers[index].bias_dropout_add_func(hidden_states0, last_mlp_bias, residual0) - - if index < self.num_layers - 1: - layernorm_output0 = self.layers[index + 1].input_layernorm(hidden_states0) - layernorm_output0 = _Wait_bwd_comm(layernorm_output0, DominoUtil.HANDLE_DIC, DominoUtil.BATCH_0) - # End of Micro batch 0: Residual connection. - - hidden_states1, last_mlp_bias = self.layers[index].linear_fc2(output1) - - fwd_handle1 = dist.all_reduce(hidden_states1, group=self.TP_group, async_op=True) - # End of Micro batch 1: MLP - - # ------------ End of MLP ------------ - - if self.post_process and self.post_layer_norm: - hidden_states0 = self.final_layernorm(hidden_states0) - - index = self.num_layers - 1 - - fwd_handle1.wait() - hidden_states1 = self.layers[index].bias_dropout_add_func(hidden_states1, last_mlp_bias, residual1) - - if self.post_process and self.post_layer_norm: - hidden_states1 = self.final_layernorm(hidden_states1) - - hidden_states = torch.cat([hidden_states0, hidden_states1], dim=1) - - return hidden_states - - def intra_layer_overlap_forward(self, hidden_states, attention_mask, rotary_pos_emb=None): - - hidden_states = torch.chunk(hidden_states, chunks=2, dim=1) - - for index in range(self.num_layers): - layer = self.layers[index] - hidden_states = layer(hidden_states, attention_mask, rotary_pos_emb) - - hidden_states0, hidden_states1 = hidden_states - if self.post_process and self.post_layer_norm: - hidden_states0 = self.final_layernorm(hidden_states0) - hidden_states1 = self.final_layernorm(hidden_states1) - - hidden_states = torch.cat([hidden_states0, hidden_states1], dim=1) - return hidden_states diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/eigenvalue.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/eigenvalue.py deleted file mode 100644 index a82d8b1d5c7a04d793940fe8c85d23e15be91c69..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/eigenvalue.py +++ /dev/null @@ -1,153 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed.utils import log_dist -import numpy as np -import logging -from deepspeed.utils.torch import required_torch_version - - -class Eigenvalue(object): - - def __init__(self, - verbose=False, - max_iter=100, - tol=1e-2, - stability=0, - gas_boundary_resolution=1, - layer_name='', - layer_num=0): - super().__init__() - - self.verbose = verbose - self.max_iter = max_iter - self.tol = tol - self.stability = stability - self.gas_boundary_resolution = gas_boundary_resolution - self.layer_name = layer_name - self.layer_num = layer_num - - assert len(self.layer_name) > 0 and layer_num > 0 - - log_dist( - f'enabled eigenvalue with verbose={verbose}, max_iter={max_iter}, tol={tol}, stability={stability}, gas_boundary_resolution={gas_boundary_resolution}, layer_name={layer_name}, layer_num={layer_num}', - ranks=[0]) - - # Replace all nan/pos-inf/neg-inf to zero - def nan_to_num(self, x): - if required_torch_version(min_version=1.8): - return torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0) - else: - # Fallback to numpy based implementation for backwards-compatibility with PyTorch 1.7 or older versions. - device = x.device - x = x.cpu().numpy() - x = np.nan_to_num(x=x, copy=False, nan=0.0, posinf=0.0, neginf=0.0) - return torch.from_numpy(x).to(device) - - def normalize(self, v): - norm_squared = self.inner_product(v, v) - norm = norm_squared**0.5 + self.stability - normalized_vectors = [vector / norm for vector in v] - normalized_vectors = [self.nan_to_num(vector) for vector in normalized_vectors] - return normalized_vectors - - def inner_product(self, xs, ys): - return sum([torch.sum(x * y) for (x, y) in zip(xs, ys)]) - - def get_layers(self, module): - scope_names = self.layer_name.split('.') - assert len(scope_names) > 0 - - m = module - for name in scope_names: - assert hasattr(m, name), "layer_name configuration is invalid." - m = getattr(m, name) - - return m - - def compute_eigenvalue(self, module, device=None, scale=1.0): - block_eigenvalue = [] - param_keys = [] - layers = self.get_layers(module) - - for block in range(self.layer_num): - model_block = layers[block] - - # We found this randn() has obvious accuracy impact in some cases, save/recover random state here. - rng_state = torch.random.get_rng_state() - if device is None: - v = [ - torch.randn(p.size()) for p in model_block.parameters() - if p.grad is not None and p.grad.grad_fn is not None - ] - else: - v = [ - torch.randn(p.size(), device=device) for p in model_block.parameters() - if p.grad is not None and p.grad.grad_fn is not None - ] - torch.random.set_rng_state(rng_state) - - grads = [ - param.grad for param in model_block.parameters() - if param.grad is not None and param.grad.grad_fn is not None - ] - params = [ - param for param in model_block.parameters() - if param.grad is not None and param.grad.grad_fn is not None - ] - - layer_keys = [id(p) for p in model_block.parameters()] - param_keys.append(layer_keys) - - v = self.normalize(v) - - # Disable eigenvalue if the model doesn't support second order gradients computation, - # e.g. when enabling DS transformer kernel. - if len(grads) == 0 or len(params) == 0: - log_dist(f'The model does NOT support eigenvalue computation.', ranks=[0], level=logging.WARNING) - return [] - - i = 0 - eigenvalue_current, eigenvalue_previous = 1., 0. - - while (i < self.max_iter) and abs(eigenvalue_current) > 0 and (abs( - (eigenvalue_current - eigenvalue_previous) / eigenvalue_current) - >= self.tol): # test convergence criteria - eigenvalue_previous = eigenvalue_current - - Hv = torch.autograd.grad(grads, params, grad_outputs=v, only_inputs=True, retain_graph=True) - #Hv = [hv.float() for hv in Hv] - Hv = [self.nan_to_num(hv).float() for hv in Hv] - - eigenvalue_current = self.inner_product(Hv, v).item() - - v = self.normalize(Hv) - v = [x / scale for x in v] - i += 1 - - eigenvalue_current *= scale - block_eigenvalue.append(eigenvalue_current) - - if self.verbose: - log_dist(f'block: {block}, power iteration: {i}, eigenvalue: {eigenvalue_current}', ranks=[0]) - - block_eigenvalue = self.post_process(block_eigenvalue) - - if self.verbose: - log_dist(f'post processed block_eigenvalue: {block_eigenvalue}', ranks=[0]) - - # {param_id: (eigenvalue, layer_id)} - ev_dict = {} - for i, (layer_keys, value) in enumerate(zip(param_keys, block_eigenvalue)): - ev_dict.update(dict.fromkeys(layer_keys, (value, i))) - - return ev_dict - - # 1. Map all eigenvalues to [0, 1.0]. - # 2. Some layers can't generate valid eigenvalues on fp16 precision, use 1.0 instead. - def post_process(self, value_list): - max_value = abs(max(value_list, key=abs)) - return [abs(v) / max_value if v != 0.0 else 1.0 for v in value_list] diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/engine.py deleted file mode 100644 index 76b83a716ffee91d78c7d2dea44233ba8928cfdf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/engine.py +++ /dev/null @@ -1,4074 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import re -import stat -import torch -import hashlib -from collections import defaultdict, OrderedDict, deque -from shutil import copyfile -import gc - -from torch.nn.modules import Module -from torch.nn.parameter import Parameter -from torch.optim import Optimizer -from torch.optim.lr_scheduler import _LRScheduler -from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors -from contextlib import contextmanager - -from typing import Callable, Dict, Union, Iterable, Container, List - -import deepspeed - -from deepspeed import comm as dist -from deepspeed.runtime.utils import see_memory_usage, DummyOptim -from .zero.offload_config import OffloadDeviceEnum, OffloadStateTypeEnum -from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer -from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus -from deepspeed.runtime.zero.utils import is_zero_supported_optimizer, ZeRORuntimeException -from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload -from deepspeed.runtime.zero.config import ZERO_OPTIMIZATION - -from deepspeed.runtime.fp16.fused_optimizer import FP16_Optimizer -from deepspeed.runtime.fp16.unfused_optimizer import FP16_UnfusedOptimizer -from deepspeed.runtime.bf16_optimizer import BF16_Optimizer - -from deepspeed.linear.optimized_linear import LoRAOptimizedLinear -from deepspeed.module_inject.layers import GatherReplacedLayerParams, configure_tensor_parallel_runtime -from deepspeed.runtime.config import DEEPSPEED_OPTIMIZERS, \ - ADAGRAD_OPTIMIZER, ADAM_OPTIMIZER, ADAMW_OPTIMIZER, LAMB_OPTIMIZER, ONEBIT_ADAM_OPTIMIZER, ONEBIT_LAMB_OPTIMIZER, \ - TORCH_ADAM_PARAM, ADAM_W_MODE, ADAM_W_MODE_DEFAULT, ZERO_ONE_ADAM_OPTIMIZER, MUADAM_OPTIMIZER, MUADAMW_OPTIMIZER, \ - MUSGD_OPTIMIZER, LION_OPTIMIZER - -from deepspeed.runtime.model_checkpointing.constants import ValidationMode, \ - CHECKPOINT_TAG_VALIDATION, CHECKPOINT_WRITER, CHECKPOINT_SERIALIZATION - -from deepspeed.runtime.dataloader import DeepSpeedDataLoader -from deepspeed.runtime.constants import \ - ROUTE_TRAIN, ROUTE_PREDICT, ROUTE_EVAL, \ - PLD_THETA, PLD_GAMMA, BFLOAT16, FP16, AMP, GRADIENT_ACCUMULATION_STEPS, \ - DATA_PARALLEL_GROUP, GLOBAL_RANK -from deepspeed.runtime.zero.config import ZeroStageEnum -from deepspeed.compression import compression_scheduler -from deepspeed.compression.constants import \ - WEIGHT_QUANTIZE_IN_FORWARD_ENABLED, \ - WEIGHT_QUANTIZATION, SHARED_PARAMETERS, \ - WEIGHT_QUANTIZE_ENABLED, \ - WEIGHT_QUANTIZE_GROUPS, \ - WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE, \ - WEIGHT_QUANTIZE_CHANGE_RATIO, \ - WEIGHT_QUANTIZE_TYPE, \ - WEIGHT_QUANTIZE_ROUNDING, \ - WEIGHT_QUANTIZE_VERBOSE, \ - WEIGHT_QUANTIZE_KERNEL -from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, FROZEN_PARAM_FRAGMENTS -from deepspeed.checkpoint.utils import clone_tensors_for_torch_save -from deepspeed.runtime.sparse_tensor import SparseTensor - -from deepspeed.runtime import lr_schedules -from deepspeed.utils import groups -from deepspeed.utils import logger, log_dist, instrument_w_nvtx -from deepspeed.utils.timer import NoopTimer, ThroughputTimer, SynchronizedWallClockTimer, \ - FORWARD_MICRO_TIMER, BACKWARD_MICRO_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_REDUCE_MICRO_TIMER, \ - STEP_MICRO_TIMER, \ - FORWARD_GLOBAL_TIMER, BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_GLOBAL_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER, \ - STEP_GLOBAL_TIMER -from deepspeed.utils.debug import debug_extract_module_and_param_names, debug_clear_module_and_param_names -from deepspeed.monitor.monitor import MonitorMaster -from deepspeed.runtime.progressive_layer_drop import ProgressiveLayerDrop -from deepspeed.runtime.utils import clip_grad_norm_, compare_tensors_in_structures -from deepspeed.runtime.eigenvalue import Eigenvalue -from deepspeed.runtime.data_pipeline.constants import DATA_SAMPLING, \ - DATA_ROUTING, DATA_SAMPLING_ENABLED, CURRICULUM_LEARNING, \ - CURRICULUM_LEARNING_ENABLED, DATA_SAMPLING_NUM_WORKERS, RANDOM_LTD, \ - RANDOM_LTD_ENABLED, RANDOM_LTD_LAYER_ID, RANDOM_LTD_LAYER_NUM, \ - RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE, RANDOM_LTD_LAYER_TOKEN_LR_ENABLED, \ - RANDOM_LTD_GLOBAL_BATCH_SIZE, RANDOM_LTD_MICRO_BATCH_SIZE, DATA_EFFICIENCY -from deepspeed.runtime.data_pipeline.curriculum_scheduler import CurriculumScheduler -from deepspeed.runtime.checkpoint_engine import (create_checkpoint_engine, TorchCheckpointEngine, CheckpointCommitInfo) - -from deepspeed.runtime.data_pipeline.data_routing.scheduler import RandomLTDScheduler -from deepspeed.runtime.data_pipeline.data_routing.helper import remove_random_ltd_state_dict -from deepspeed.runtime.data_pipeline.data_routing.basic_layer import RandomLayerTokenDrop - -from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint -from deepspeed.runtime.torch_autocast import init_autocast_params, get_default_autocast_lower_precision_modules, validate_nested_autocast - -from .pipe.module import PipelineModule -from .utils import get_ma_status -from .compiler import is_compile_supported -from ..ops.adam import FusedAdam -from ..moe.sharded_moe import TopKGate, MOELayer -from ..moe.layer import MoE -from ..moe.utils import is_moe_param, configure_moe_param_groups -from ..git_version_info import version - -from deepspeed.profiling.flops_profiler.profiler import FlopsProfiler -from deepspeed.utils.logging import print_json_dist, print_configuration - -from deepspeed.accelerator import get_accelerator - -from deepspeed.runtime.config import DtypeEnum - -from deepspeed.compile.util import is_deepcompile_supported, get_deepcompile_handle, deepcompile_backward_prologue -from deepspeed.compile.backend import register_compile_pass, opt_passes -from deepspeed.compile.passes import zero3_compile, prefetch, selective_gather, offload_adam_states -from deepspeed.compile.init_z1 import init_z1 -from deepspeed.compile.init_z3 import init_z3 - -MEMORY_OPT_ALLREDUCE_SIZE = 500000000 - -DeepSpeedOptimizerCallable = \ - Callable[[Union[Iterable[Parameter], Dict[str, Iterable]]], Optimizer] -DeepSpeedSchedulerCallable = Callable[[Optimizer], _LRScheduler] - -try: - import apex - from apex import amp - APEX_INSTALLED = True -except ImportError: - # Fail silently so we don't spam logs unnecessarily if user isn't using amp - APEX_INSTALLED = False - - -def split_half_float_double_sparse(tensors): - device_type = get_accelerator().device_name() - supported_types = get_accelerator().supported_dtypes() - - for t in tensors: - assert t.dtype in supported_types, f"attempting to reduce an unsupported grad type: {t.dtype}" - - sparse_tensor_buckets, dense_tensor_buckets = [], [] - for i, dtype in enumerate(supported_types): - sparse_bucket, dense_bucket = [], [] - for t in tensors: - if t.dtype == dtype: - if isinstance(t, SparseTensor): - sparse_bucket.append(t) - else: - dense_bucket.append(t) - if sparse_bucket: - sparse_tensor_buckets.append((dtype, sparse_bucket)) - if dense_bucket: - dense_tensor_buckets.append((dtype, dense_bucket)) - return sparse_tensor_buckets, dense_tensor_buckets - - -class EngineTimers(object): - r"""Wallclock timers for DeepSpeedEngine""" - - def __init__(self, enable_micro_timers, enable_global_timers): - self.forward_timers = [] - self.backward_timers = [] - self.backward_inner_timers = [] - self.backward_reduce_timers = [] - self.step_timers = [] - self.global_timers = [] - self.micro_timers = [] - - if enable_micro_timers: - self.forward_timers += [FORWARD_MICRO_TIMER] - self.backward_timers += [BACKWARD_MICRO_TIMER] - self.backward_inner_timers += [BACKWARD_INNER_MICRO_TIMER] - self.backward_reduce_timers += [BACKWARD_REDUCE_MICRO_TIMER] - self.step_timers += [STEP_MICRO_TIMER] - self.micro_timers += [ - FORWARD_MICRO_TIMER, BACKWARD_MICRO_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_REDUCE_MICRO_TIMER, - STEP_MICRO_TIMER - ] - - if enable_global_timers: - self.forward_timers += [FORWARD_GLOBAL_TIMER] - self.backward_timers += [BACKWARD_GLOBAL_TIMER] - self.backward_inner_timers += [BACKWARD_INNER_GLOBAL_TIMER] - self.backward_reduce_timers += [BACKWARD_REDUCE_GLOBAL_TIMER] - self.step_timers += [STEP_GLOBAL_TIMER] - self.global_timers += [ - FORWARD_GLOBAL_TIMER, BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_GLOBAL_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER, - STEP_GLOBAL_TIMER - ] - - -class DeepSpeedEngine(Module): - r"""DeepSpeed engine for training.""" - - def __init__(self, - args, - model, - optimizer=None, - model_parameters=None, - training_data=None, - lr_scheduler=None, - mpu=None, - dist_init_required=None, - collate_fn=None, - config=None, - config_class=None, - mesh_device=None, - dont_change_device=False): - super(DeepSpeedEngine, self).__init__() - self.dont_change_device = dont_change_device - self.client_optimizer = optimizer - self.client_lr_scheduler = lr_scheduler - self.training_data = training_data - self.collate_fn = collate_fn - self.mpu = mpu - self.all_to_all_group = None - self.data_parallel_group = None - self.global_steps = 0 - self.global_samples = 0 - self.micro_steps = 0 - self.skipped_steps = 0 - self.gradient_average = True - self.warn_unscaled_loss = True - self.config = config - self._config = config_class - self.loaded_checkpoint_mp_world_size = None - self.loaded_checkpoint_dp_world_size = None - self.enable_backward_allreduce = True - self.inside_no_sync_ctxt = False - self.progressive_layer_drop = None - self.eigenvalue = None - self.block_eigenvalue = None - self.gas_boundary_ctr = 0 - self.dist_backend = get_accelerator().communication_backend_name() - self.has_moe_layers = False - self.num_experts = [] - self.gate_modules = [] - self.moe_layers = [] - self._step_applied = False - self._global_grad_norm = None - self.use_ds_comm = False # False --> Use torch.dist, True --> Use ds.comm backend. - self.checkpoint_engine = None - self.optimizer = None - self.basic_optimizer = None - self.lr_scheduler = None - - self._is_gradient_accumulation_boundary = None - self.scale_wrt_gas = None - self.losses = None - self.mesh_device = mesh_device - - # for debug purposes - can then debug print: debug_get_module_name(module) - debug_extract_module_and_param_names(model) - - if self.mesh_device: - groups.mesh_device = self.mesh_device - - self._do_args_sanity_check(args) - self._configure_with_arguments(args, mpu) - self._do_sanity_check() - if self.autotp_size() > 1: - self._configure_tensor_parallel(model, self.tensor_parallel_config()) - see_memory_usage(f"DeepSpeed Engine: After args sanity test", force=self.memory_breakdown()) - if mpu is not None: - if self.elasticity_enabled(): - if not self.is_elastic_model_parallel_supported(): - assert not self.elasticity_enabled(), ("Elasticity is not currently supported" - " with model parallelism.") - - self._set_distributed_vars(args) - - dist.configure(self._config) - - self.monitor = MonitorMaster(self._config.monitor_config) - - see_memory_usage( - f"DeepSpeed Engine: Before configure distributed model", - force=self.memory_breakdown(), - ) - - self.pipeline_parallelism = isinstance(model, PipelineModule) - - # Configure distributed model - self._configure_distributed_model(model) - - if not self.is_deepcompile_enabled(): - self.module_forward_pre_hook = self._create_module_forward_pre_hook() - self.module_forward_post_hook = self._create_module_forward_post_hook() - - # needed for zero_to_fp32 weights reconstruction to remap nameless data to state_dict - self.param_names = {param: name for name, param in model.named_parameters()} - - self._get_model_parameters() - - see_memory_usage(f"DeepSpeed Engine: After configure distributed model") - - # Configure wall clock timers - self.timers = SynchronizedWallClockTimer() - # Throughput timer - self.tput_timer = ThroughputTimer(self._config.timers_config, - batch_size=self.train_batch_size(), - steps_per_output=self.steps_per_print(), - monitor_memory=False) - - log_dist(f"DeepSpeed Flops Profiler Enabled: {self.flops_profiler_enabled()}", ranks=[0]) - - if self.flops_profiler_enabled(): - self.flops_profiler = FlopsProfiler(self.module, self, self.flops_profiler_recompute_fwd_factor()) - - if training_data: - self.training_dataloader = self.deepspeed_io(training_data) - else: - self.training_dataloader = None - - # Configure optimizer and scheduler - has_optimizer = False - - if optimizer or self.optimizer_name(): - has_optimizer = True - # If no parameters given by init default to module parameters - if model_parameters is None: - model_parameters = self.module.parameters() - - # Convert model parameters from generator to list - if not isinstance(model_parameters, list): - model_parameters = list(model_parameters) - - if self.torch_autocast_enabled(): - init_autocast_params(self, self.torch_autocast_dtype(), self.torch_autocast_lower_precision_safe_modules()) - - if has_optimizer: - self._configure_optimizer(optimizer, model_parameters) - self._configure_lr_scheduler() - self._report_progress(0) - elif self.zero_optimization(): - # no optim selected but zero is enabled - self.optimizer = self._configure_zero_optimizer(optimizer=None) - elif self.bfloat16_enabled(): - self.optimizer = self._configure_bf16_optimizer(optimizer=None) - - # Hook optimizer for snip_momentum pruning - if hasattr(model, 'pruners'): - from ..compression.helper import rewrite_optimizer_step - self.optimizer.pruners = model.pruners - rewrite_optimizer_step(self.optimizer) - - # Bookkeeping for sparse support - self.sparse_tensor_module_names = set() - # if self.sparse_gradients_enabled(): - for name, module in self.module.named_modules(): - if isinstance(module, (torch.nn.Embedding, torch.nn.EmbeddingBag)) and self.sparse_gradients_enabled(): - self.sparse_tensor_module_names.add(name + ".weight") - logger.info("Will convert {} to sparse tensor during training".format(name)) - - self._optimized_linear_offload_setup() - - self.save_non_zero_checkpoint = False - self.save_zero_checkpoint = False - if not isinstance(self.optimizer, DeepSpeedZeRoOffload): - self._configure_checkpointing() - - if self.eigenvalue_enabled(): - self.eigenvalue = self._configure_eigenvalue() - - if self.pld_enabled(): - self.progressive_layer_drop = self._configure_progressive_layer_drop() - - if self.curriculum_enabled_legacy(): - self.curriculum_scheduler_legacy = self._configure_curriculum_scheduler_legacy() - - if self.random_ltd_enabled(): - random_ltd_config = self.random_ltd_config() - random_ltd_config[RANDOM_LTD_GLOBAL_BATCH_SIZE] = self.train_batch_size() - random_ltd_config[RANDOM_LTD_MICRO_BATCH_SIZE] = self.train_micro_batch_size_per_gpu() - self.random_ltd_scheduler = self._configure_random_ltd_scheduler(random_ltd_config) - - # Engine timers - - self.engine_timers = EngineTimers(enable_micro_timers=self.wall_clock_breakdown(), - enable_global_timers=self.wall_clock_breakdown() - or self.flops_profiler_enabled()) - - if self.global_rank == 0: - self._config.print("DeepSpeedEngine configuration") - if self.dump_state(): - print_configuration(self, "DeepSpeedEngine") - - # Use torch (un)flatten ops - self.flatten = _flatten_dense_tensors - self.unflatten = _unflatten_dense_tensors - - self._is_compiled = False - if is_deepcompile_supported(): - # Predefined compile passes - self.register_compile_pass(zero3_compile.NAME, zero3_compile.add_z3_gather_release) - self.register_compile_pass(prefetch.NAME, prefetch.schedule_prefetch) - self.register_compile_pass(selective_gather.NAME, selective_gather.selective_gather) - self.register_compile_pass(offload_adam_states.NAME, offload_adam_states.move_opt_states) - - def _optimized_linear_offload_setup(self): - self.optimized_linear_base_weight_sharding = False - self.optimized_linear_lora_enabled = False - offload_ratio = None - for _, module in self.module.named_modules(): - if isinstance(module, LoRAOptimizedLinear): - self.optimized_linear_lora_enabled = True - offload_ratio = None - if offload_ratio is not None: - assert offload_ratio == module.lora_config.offload_ratio, \ - "all lora_config offload ratios should be the same across the model" - offload_ratio = module.lora_config.offload_ratio - if module.zero_shards > 1: - # set attr so checkpoint saving can handle BWS properly - self.optimized_linear_base_weight_sharding = True - - if offload_ratio is None: - # Nothing enabled, do nothing - return - - total_params = 0 - for _, p in self.module.named_parameters(): - if hasattr(p, 'ds_optim_param'): - total_params += p.numel() - - offload_limit = total_params * offload_ratio - logger.info(f'offloading {offload_ratio*100}% of eligible params, specifically {offload_limit} params') - total_offloaded = 0 - for _, p in self.module.named_parameters(): - if hasattr(p, 'ds_optim_param'): - if total_offloaded < offload_limit: - total_offloaded += p.numel() - p.ds_offload = True - p.offload() - else: - p.ds_offload = False - - def _configure_tensor_parallel(self, model, tp_config): - self._configure_tensor_parallel_states(model) - configure_tensor_parallel_runtime(tp_config) - - def _configure_tensor_parallel_states(self, model): - """ - Configures the tensor parallel states for the model. - This includes setting up the tensor parallel groups, initializing the TP mesh, - and registering a pre-hook to ensure that the Dataloader inputs are consistent across ranks. - """ - self._set_client_model(model) - # sanity check - # currently, the compatibility between 'autotp' and 'zero > 1' has not been validated - assert self.zero_optimization_stage( - ) <= 2, "Currently, the compatibility between 'autotp' and 'zero_stage = 3' has not been validated" - - self.mpu = groups - self.mpu._init_tp_mesh_device(tensor_model_parallel_size=self.autotp_size()) - - self.first_dataloader_check = None - - def check_dataloader_inputs_same_across_ranks(module, args, kwargs): - - def broadcast_and_check(args, bcast_rank, bcast_group): - if isinstance(args, tuple): - args = list(args) - if len(args) > 0: - if self.mpu.get_tensor_model_parallel_rank() == 0: - _src_args = [args] - dist.broadcast_object_list(object_list=_src_args, - src=bcast_rank, - group=bcast_group, - device=get_accelerator().current_device()) - # Rank 0 does not need to compare with itself - is_equal = True - else: - _src_args = [None] - dist.broadcast_object_list(object_list=_src_args, - src=bcast_rank, - group=bcast_group, - device=get_accelerator().current_device()) - - is_equal = compare_tensors_in_structures(args, _src_args[0]) - - equal_tensor = torch.tensor(is_equal, - dtype=self.communication_data_type, - device=get_accelerator().current_device()) - dist.all_reduce(equal_tensor, group=bcast_group) - assert torch.equal( - equal_tensor, - torch.tensor(groups.get_tensor_model_parallel_world_size(), - dtype=self.communication_data_type, - device=get_accelerator().current_device()) - ), "Data inconsistency within the TP group. Please check the Dataloader implementation to ensure consistency." - - bcast_rank = self.mpu.get_tensor_model_parallel_src_rank() - bcast_group = self.mpu.get_tensor_model_parallel_group() - - broadcast_and_check(args, bcast_rank, bcast_group) - broadcast_and_check(kwargs, bcast_rank, bcast_group) - - logger.info(f":The Dataloader has passed the TP group consistency check.") - self.first_dataloader_check.remove() - - self.first_dataloader_check = self.module.register_forward_pre_hook(check_dataloader_inputs_same_across_ranks, - prepend=True, - with_kwargs=True) - - def __del__(self): - self.destroy() - - def destroy(self): - if self.optimizer is not None and hasattr(self.optimizer, 'destroy'): - self.optimizer.destroy() - if self.is_deepcompile_enabled(): - get_deepcompile_handle().cleanup() - debug_clear_module_and_param_names() - - if self.checkpoint_engine is not None and self.checkpoint_engine.is_decoupled(): - self.checkpoint_engine.cleanup() - - def _get_model_parameters(self): - if self.autotuning_profile_model_info(): - self.autotuning_model_info = {} - num_params = 0 - trainable_num_params = 0 - - for p in self.module.parameters(): - # since user code might call deepspeed.zero.Init() before deepspeed.initialize(), need to check the attribute to check if the parameter is partitioned in zero 3 already or not - n = 0 - if hasattr(p, "ds_tensor"): # if the parameter is partitioned in zero 3 - n += p.ds_numel - else: # if the parameter is not partitioned in zero 3 yet - n += p.numel() - num_params += n - if p.requires_grad: - trainable_num_params += n - if self.global_rank == 0: - self.autotuning_model_info["num_params"] = num_params * self.mp_world_size - self.autotuning_model_info["trainable_num_params"] = trainable_num_params * self.mp_world_size - - logger.info(f"model parameter = {num_params}") - - def get_batch_info(self): - """Get all training batch related settings. - Returns: - train_batch_size (int): The effective training batch size. This is the amount of data - samples that leads to one step of model update. - train_micro_batch_size_per_gpu (int): Batch size to be processed by one GPU in one - step (without gradient accumulation). - gradient_accumulation_steps (int): Number of training steps to accumulate gradients - before averaging and applying them. - """ - return ( - self.train_batch_size, - self.train_micro_batch_size_per_gpu, - self.gradient_accumulation_steps, - ) - - def set_train_batch_size(self, train_batch_size): - """Adjust the global batch size by increasing or decreasing the number of - micro-batches (i.e., gradient accumulation steps). The size of each micro-batch - (i.e., ``train_micro_batch_size_per_gpu``) is not changed. - Args: - train_batch_size (int): The new global batch size for training. - Raises: - ValueError: if ``train_batch_size`` is not divisible by the - configured micro-batch size and data parallelism. - """ - if train_batch_size % (self.train_micro_batch_size_per_gpu() * self.dp_world_size) != 0: - #print(f'{train_batch_size=} {self.train_micro_batch_size_per_gpu()=} {self.dp_world_size=}') - raise ValueError(f'Train batch size must be divisible by micro-batch data parallelism') - new_gas = train_batch_size // (self.train_micro_batch_size_per_gpu() * self.dp_world_size) - # overwrite config - self._config.train_batch_size = train_batch_size - self._config.gradient_accumulation_steps = new_gas - - def set_train_micro_batch_size(self, micro_batch_size): - """Adjust the micro batch size(i.e., the micro batch size in every data parallel group), - while keep the gradient accumulation steps the same. - Args: - micro_batch_size (int): The new micro batch size for training. - """ - # overwrite config - new_global_batch_size = micro_batch_size * self._config.gradient_accumulation_steps * self.dp_world_size - self._config.train_batch_size = new_global_batch_size - self._config.train_micro_batch_size_per_gpu = micro_batch_size - - def set_data_post_process_func(self, post_process_func): - if self.training_dataloader is not None: - self.training_dataloader.post_process_func = post_process_func - - def set_custom_curriculum_learning_schedule(self, schedule_func_dict): - if self.training_dataloader is not None and self.curriculum_learning_enabled(): - self.training_dataloader.data_sampler.set_custom_curriculum_learning_schedule(schedule_func_dict) - - def get_global_grad_norm(self) -> float: - """Return the 2-norm of all gradients. If there is model parallelism, - the norm will be global. - The computed norm will be cached and reused until the next step() pass. - .. note:: - In the presence of model parallelism, this is a collective call - and acts as a barrier among ``mpu.get_model_parallel_group()``. - Returns: - float: norm - """ - return self._global_grad_norm - - def __getattr__(self, name): - """ - Pass through attributes defined in the model if they are not overridden by ds-engine. - """ - - _module = {} - if "module" in self.__dict__: - _module = self.__dict__['module'] - if name in dir(self): - return getattr(self, name) - elif name in dir(_module): - return getattr(_module, name) - else: - raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") - - def checkpoint_serialization_enabled(self): - return self._config.checkpoint_config[CHECKPOINT_SERIALIZATION] - - def checkpoint_writer_enabled(self): - return self._config.checkpoint_config[CHECKPOINT_WRITER] is not None - - def checkpoint_tag_validation_enabled(self): - return self._config.checkpoint_config[CHECKPOINT_TAG_VALIDATION] != ValidationMode.IGNORE - - def checkpoint_tag_validation_fail(self): - return self._config.checkpoint_config[CHECKPOINT_TAG_VALIDATION] == ValidationMode.FAIL - - def elasticity_enabled(self): - return self._config.elasticity_enabled - - def is_elastic_model_parallel_supported(self): - if self.elasticity_enabled(): - # Add code for finding number of GPUs per node automatically - if self._config.num_gpus_per_node % self._config.elastic_model_parallel_size == 0: - return True - else: - return False - - def pld_enabled(self): - return self._config.pld_enabled - - def pld_params(self): - return self._config.pld_params - - def pld_theta(self): - return self.pld_params()[PLD_THETA] - - def pld_gamma(self): - return self.pld_params()[PLD_GAMMA] - - def eigenvalue_enabled(self): - return self._config.eigenvalue_enabled - - def eigenvalue_verbose(self): - return self._config.eigenvalue_verbose - - def eigenvalue_max_iter(self): - return self._config.eigenvalue_max_iter - - def eigenvalue_tol(self): - return self._config.eigenvalue_tol - - def eigenvalue_stability(self): - return self._config.eigenvalue_stability - - def eigenvalue_gas_boundary_resolution(self): - return self._config.eigenvalue_gas_boundary_resolution - - def eigenvalue_layer_name(self): - return self._config.eigenvalue_layer_name - - def eigenvalue_layer_num(self): - return self._config.eigenvalue_layer_num - - def curriculum_enabled_legacy(self): - return self._config.curriculum_enabled_legacy - - def curriculum_params_legacy(self): - return self._config.curriculum_params_legacy - - def data_efficiency_enabled(self): - return self._config.data_efficiency_enabled - - def data_efficiency_config(self): - return self._config.data_efficiency_config - - def data_sampling_enabled(self): - return self._config.data_efficiency_config[DATA_SAMPLING][DATA_SAMPLING_ENABLED] - - def data_sampling_config(self): - return self._config.data_efficiency_config[DATA_SAMPLING] - - def curriculum_learning_enabled(self): - return self._config.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_ENABLED] - - def curriculum_learning_config(self): - return self._config.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING] - - def random_ltd_enabled(self): - return self._config.data_efficiency_config[DATA_ROUTING][RANDOM_LTD][RANDOM_LTD_ENABLED] - - def random_ltd_config(self): - return self._config.data_efficiency_config[DATA_ROUTING][RANDOM_LTD] - - def random_ltd_initialize(self): - assert self.random_ltd_enabled() - random_ltd_config = self.random_ltd_config() - random_ltd_queue = deque([x for x in sorted(random_ltd_config[RANDOM_LTD_LAYER_ID])]) - count = 0 - for name, layer in self.module.named_modules(): - if isinstance(layer, RandomLayerTokenDrop): - if len(random_ltd_queue) != 0 and str(random_ltd_queue[0]) in name: ###[1,2,3] - layer.init_config(random_ltd_config, self.random_ltd_scheduler, count) - random_ltd_queue.popleft() - count += 1 - - if random_ltd_config[RANDOM_LTD_LAYER_NUM] != count: - raise ValueError(f'random_ltd_layer_num {random_ltd_config[RANDOM_LTD_LAYER_NUM]} must be \ - equivalent to the len of random_ltd_layer_id {count}') - - if random_ltd_config[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE][RANDOM_LTD_LAYER_TOKEN_LR_ENABLED]: - assert self.client_lr_scheduler is None - raise ValueError(f'not yet support') - #self.lr_scheduler = lr_schedules.WarmupLayerTokenDecayLR(self.optimizer, self.random_ltd_scheduler) - - def get_sequence_parallel_group(self): - return self.seq_parallel_group - - def wall_clock_breakdown(self): - return self._config.wall_clock_breakdown - - def flops_profiler_enabled(self): - return self._config.flops_profiler_config.enabled or self.autotuning_enabled() - - def flops_profiler_recompute_fwd_factor(self): - return self._config.flops_profiler_config.recompute_fwd_factor - - def flops_profiler_profile_step(self): - step = self._config.flops_profiler_config.profile_step - if self._config.autotuning_config.enabled: - step = self.autotuning_start_profile_step() - return step - - def flops_profiler_module_depth(self): - return self._config.flops_profiler_config.module_depth - - def flops_profiler_top_modules(self): - return self._config.flops_profiler_config.top_modules - - def flops_profiler_detailed(self): - if self._config.autotuning_config.enabled: - return False - return self._config.flops_profiler_config.detailed - - def flops_profiler_output_file(self): - return self._config.flops_profiler_config.output_file - - def memory_breakdown(self): - return self._config.memory_breakdown - - def autotuning_enabled(self): - return self._config.autotuning_config.enabled - - def autotuning_start_profile_step(self): - return self._config.autotuning_config.start_profile_step - - def autotuning_end_profile_step(self): - return self._config.autotuning_config.end_profile_step - - def autotuning_metric_path(self): - path = self._config.autotuning_config.metric_path - if not path: - path = os.path.join(os.getcwd(), "autotuning_metric.json") - return path - - def autotuning_model_info_path(self): - path = self._config.autotuning_config.model_info_path - if not path: - path = os.path.join(os.getcwd(), "autotuning_model_info.json") - return path - - def autotuning_metric(self): - return self._config.autotuning_config.metric - - def autotuning_profile_model_info(self): - return self.autotuning_enabled( - ) and self._config.autotuning_config.model_info and self._config.autotuning_config.model_info.get( - "profile", False) - - def sparse_gradients_enabled(self): - return self._config.sparse_gradients_enabled - - def train_batch_size(self): - return self._config.train_batch_size - - def train_micro_batch_size_per_gpu(self): - return self._config.train_micro_batch_size_per_gpu - - def optimizer_name(self): - return (self.client_optimizer.__class__.__name__ if self.client_optimizer else self._config.optimizer_name) - - def optimizer_params(self): - return self._config.optimizer_params - - def optimizer_legacy_fusion(self): - return self._config.optimizer_legacy_fusion - - def scheduler_name(self): - return self._config.scheduler_name - - def scheduler_params(self): - return self._config.scheduler_params - - def quantize_training(self): - return ( - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS] - [WEIGHT_QUANTIZE_IN_FORWARD_ENABLED], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_ENABLED], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_GROUPS], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS] - [WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_CHANGE_RATIO], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_TYPE], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_ROUNDING], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_VERBOSE], - self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_KERNEL], - ) - - def zero_optimization(self): - return self._config.zero_enabled - - def zero_allow_untested_optimizer(self): - return self._config.zero_allow_untested_optimizer - - def zero_force_ds_cpu_optimizer(self): - return self._config.zero_force_ds_cpu_optimizer - - def zero_reduce_scatter(self): - return self._config.zero_config.reduce_scatter - - def zero_overlap_comm(self): - return self._config.zero_config.overlap_comm - - def zero_offload_optimizer(self): - return self._config.zero_config.offload_optimizer - - def zero_offload_param(self): - return self._config.zero_config.offload_param - - def zero_use_cpu_optimizer(self): - if self._config.zero_config.offload_optimizer is not None: - return self._config.zero_config.offload_optimizer.device in [OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme] - return False - - def zero_cpu_offload(self): - if self._config.zero_config.offload_optimizer is not None: - return self._config.zero_config.offload_optimizer.device == OffloadDeviceEnum.cpu - return False - - def zero_partial_offload(self): - return getattr(self._config.zero_config.offload_optimizer, "ratio", 1.0) - - def zero_sub_group_size(self): - return self._config.zero_config.sub_group_size - - def zero_optimization_stage(self): - return self._config.zero_optimization_stage - - def mics_shard_size(self): - return self._config.mics_shard_size - - def zero_reduce_bucket_size(self): - return self._config.zero_config.reduce_bucket_size - - def zero_multi_rank_bucket_allreduce(self): - return self._config.zero_config.use_multi_rank_bucket_allreduce - - def zero_allgather_bucket_size(self): - return self._config.zero_config.allgather_bucket_size - - def zero_optimization_partition_gradients(self): - return self.zero_optimization_stage() >= ZeroStageEnum.gradients - - def zero_optimization_partition_weights(self): - return self.zero_optimization_stage() >= ZeroStageEnum.weights - - def is_first_weights_partition_group(self): - ret = True if self.mics_shard_size() < 0 \ - and self.zero_optimization_partition_weights() else False - if self.mics_shard_size() > 0 and self.global_rank < self.mics_shard_size(): - ret = True - return ret - - def zero_contiguous_gradients(self): - return self._config.zero_config.contiguous_gradients - - def zero_load_from_fp32_weights(self): - return self._config.zero_config.load_from_fp32_weights - - def zero_elastic_checkpoint(self): - return self._config.zero_config.elastic_checkpoint - - def zero_nvme_offload_optimizer(self): - return getattr(self.optimizer, "swap_optimizer", False) - - def zero_max_live_parameters(self): - return self._config.zero_config.max_live_parameters - - def zero_max_reuse_distance(self): - return self._config.zero_config.max_reuse_distance - - def zero_prefetch_bucket_size(self): - return self._config.zero_config.prefetch_bucket_size - - def zero_module_granularity_threshold(self): - return self._config.zero_config.module_granularity_threshold - - def zero_param_persistence_threshold(self): - return self._config.zero_config.param_persistence_threshold - - def zero_model_persistence_threshold(self): - return self._config.zero_config.model_persistence_threshold - - def zero_gather_16bit_weights_on_model_save(self): - return self._config.zero_config.gather_16bit_weights_on_model_save - - def zero_grad_hooks(self): - return self._config.zero_config.grad_hooks - - def zero_legacy_stage1(self): - return self._config.zero_config.legacy_stage1 - - def zero_ignore_unused_parameters(self): - return self._config.zero_config.ignore_unused_parameters - - def tensor_parallel_config(self): - return self._config.tensor_parallel_config - - def autotp_size(self): - return self._config.tensor_parallel_config.autotp_size - - def graph_harvesting(self): - return self._config.graph_harvesting - - def fp16_enabled(self): - return self._config.float16_config.enabled - - def bfloat16_enabled(self): - return self._config.bfloat16_config.enabled - - def fp16_master_weights_and_gradients(self): - return self._config.float16_config.fp16_master_weights_and_grads - - def amp_enabled(self): - return self._config.amp_enabled - - def amp_params(self): - return self._config.amp_params - - def torch_autocast_enabled(self) -> bool: - return self._config.torch_autocast_enabled - - def torch_autocast_dtype(self) -> torch.dtype: - return self._config.torch_autocast_dtype - - def torch_autocast_lower_precision_safe_modules(self) -> List[str]: - module_names = self._config.torch_autocast_lower_precision_safe_modules - return get_default_autocast_lower_precision_modules() if module_names is None else module_names - - def fp16_auto_cast(self): - return self._config.float16_config.auto_cast - - def loss_scale(self): - return self._config.float16_config.loss_scale - - def gradient_accumulation_steps(self): - return self._config.gradient_accumulation_steps - - def use_node_local_storage(self): - return self._config.use_node_local_storage - - def load_universal_checkpoint(self): - return self._config.load_universal_checkpoint - - @property - def communication_data_type(self): - res = self._config.communication_data_type - if res is not None: - return res - - if self.fp16_enabled(): - return torch.float16 - - if self.bfloat16_enabled(): - return torch.bfloat16 - - return torch.float32 - - @communication_data_type.setter - def communication_data_type(self, value): - self._config.communication_data_type = value - - def postscale_gradients(self): - return not self._config.prescale_gradients - - def gradient_predivide_factor(self): - return self._config.gradient_predivide_factor - - def steps_per_print(self): - return self._config.steps_per_print - - def zero_allgather_partitions(self): - return self._config.zero_config.allgather_partitions - - def zero_round_robin_gradients(self): - return self._config.zero_config.round_robin_gradients - - def zero_hpz_partition_size(self): - return self._config.zero_config.zero_hpz_partition_size - - def zero_quantized_weights(self): - return self._config.zero_config.zero_quantized_weights - - def zero_quantized_nontrainable_weights(self): - return self._config.zero_config.zero_quantized_nontrainable_weights - - def zero_quantized_gradients(self): - return self._config.zero_config.zero_quantized_gradients - - def zeropp_loco_param(self): - return self._config.zero_config.zeropp_loco_param - - def zero_log_trace_cache_warnings(self): - return self._config.zero_config.log_trace_cache_warnings - - def dump_state(self): - return self._config.dump_state - - def gradient_clipping(self): - return self._config.gradient_clipping - - def dynamic_loss_scale(self): - return self._config.float16_config.loss_scale == 0 - - def initial_dynamic_scale(self): - return self._config.float16_config.initial_dynamic_scale() - - def dynamic_loss_scale_args(self): - return self._config.float16_config.dynamic_loss_scale_args() - - def swap_tensor_config(self): - return self._config.swap_tensor_config - - def aio_config(self): - return self._config.aio_config - - def get_data_types(self): - model_dtype = torch.float32 - if self.fp16_enabled(): - model_dtype = torch.float16 - elif self.bfloat16_enabled(): - model_dtype = torch.bfloat16 - - if self._config.grad_accum_dtype is None: - if model_dtype == torch.bfloat16 and not self.zero_optimization(): - grad_accum_dtype = torch.float32 - else: - grad_accum_dtype = model_dtype - else: - grad_accum_dtype = DtypeEnum(self._config.grad_accum_dtype).value - - return (model_dtype, grad_accum_dtype) - - def _optimizer_has_ckpt_event_prologue(self): - return self.optimizer is not None and hasattr(self.optimizer, 'checkpoint_event_prologue') - - def _optimizer_has_ckpt_event_epilogue(self): - return self.optimizer is not None and hasattr(self.optimizer, 'checkpoint_event_epilogue') - - def _configure_lr_scheduler(self): - if self.client_lr_scheduler: - if isinstance(self.client_lr_scheduler, Callable): - log_dist('DeepSpeed using client callable to create LR scheduler', ranks=[0]) - self.lr_scheduler = self.client_lr_scheduler(self.basic_optimizer) - else: - log_dist('DeepSpeed using client LR scheduler', ranks=[0]) - self.lr_scheduler = self.client_lr_scheduler - else: - # load lr scheduler from json configuration if lr scheduler is not defined and passed in - lr_scheduler = self._scheduler_from_config(self.optimizer) - log_dist(f"DeepSpeed using configured LR scheduler = {self.scheduler_name()}", ranks=[0]) - self.lr_scheduler = lr_scheduler - - log_dist(f'DeepSpeed LR Scheduler = {self.lr_scheduler}', ranks=[0]) - - def _configure_checkpointing(self): - # Enable optimization to parallelize checkpointing of DP state - optimize_dp_state = not self.zero_optimization_partition_weights() - self.checkpoint_engine = create_checkpoint_engine(config_params=self._config, - groups=groups, - zero_stage=self.zero_optimization_stage(), - has_moe_layers=self.has_moe_layers, - optimize_dp_state=optimize_dp_state) - - dp_rank = groups._get_sequence_data_parallel_rank() - rank = self.local_rank if self.use_node_local_storage() else dp_rank - - # Determine if this data parallel process needs to store the model checkpoint - if self.checkpoint_engine.is_data_parallel_writer(rank) \ - or (self.zero_optimization_partition_weights() and self.is_first_weights_partition_group()): - self.save_non_zero_checkpoint = True - - if self.zero_optimization() or self.bfloat16_enabled(): - param_rank = dist.get_rank(group=self.optimizer.dp_process_group) - - # Only the first parameter parallel process needs to store the - # optimizer state checkpoints for zero - self.save_zero_checkpoint = param_rank == dp_rank - - def _scheduler_from_config(self, optimizer): - scheduler_name = self.scheduler_name() - if scheduler_name is not None: - if hasattr(lr_schedules, scheduler_name): - scheduler = getattr(lr_schedules, scheduler_name) - else: - assert hasattr(torch.optim.lr_scheduler, - scheduler_name), f"DeepSpeed does not recognize LR scheduler {scheduler_name}" - - scheduler = getattr(torch.optim.lr_scheduler, scheduler_name) - - scheduler_params = self.scheduler_params() - instantiated_scheduler = scheduler(optimizer, **scheduler_params) - return instantiated_scheduler - else: - return None - - def _set_distributed_vars(self, args): - device_rank = args.device_rank if args is not None and hasattr(args, 'device_rank') else self.local_rank - if device_rank >= 0: - get_accelerator().set_device(device_rank) - self.device = torch.device(get_accelerator().device_name(device_rank)) - self.world_size = dist.get_world_size() - self.global_rank = dist.get_rank() - else: - self.world_size = 1 - self.global_rank = 0 - self.device = get_accelerator().device() - - # Configure based on command line arguments - def _configure_with_arguments(self, args, mpu): - # After the distributed backend is initialized we are guaranteed the LOCAL_RANK - # environment variable is set. We must align args.local_rank to this value for - # backwards compatibility with scripts relying on [args|self].local_rank containing - # the correct local rank info. _do_args_sanity_check will ensure this is the case. - - if "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ: - ompi_local_rank = os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK") - local_rank = os.environ.get('LOCAL_RANK', ompi_local_rank) - assert ompi_local_rank == local_rank, f"LOCAL_RANK ({local_rank}) != OMPI_COMM_WORLD_LOCAL_RANK ({ompi_local_rank}), " \ - "not sure how to proceed as we're seeing conflicting local rank info." - os.environ['LOCAL_RANK'] = local_rank - - self.local_rank = int(os.environ['LOCAL_RANK']) - if hasattr(args, 'local_rank'): - args.local_rank = self.local_rank - - # Validate command line arguments - def _do_args_sanity_check(self, args): - assert "LOCAL_RANK" in os.environ or "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ, "DeepSpeed requires the LOCAL_RANK environment " \ - "variable, it is set by the deepspeed launcher, deepspeed.init_distributed, or the torch's launcher. If using a " \ - "different launcher please ensure LOCAL_RANK is set prior to initializing deepspeed." - - if hasattr(args, 'local_rank') and args.local_rank is not None: - assert isinstance(args.local_rank, - int), f"args.local_rank of {args.local_rank} is an unknown type {type(args.local_rank)}" - if args.local_rank >= 0: - env_local_rank = int(os.environ.get("LOCAL_RANK")) - assert ( - env_local_rank == args.local_rank - ), f"Mismatch in local rank setting, args.local_rank={args.local_rank} but env['LOCAL_RANK']={env_local_rank}." - - def _is_supported_optimizer(self, optimizer_name): - return (optimizer_name in DEEPSPEED_OPTIMIZERS or getattr(torch.optim, optimizer_name, None) is not None) - - def _supported_optims(self): - FairseqOptimizer = None - try: - from fairseq.optim.fairseq_optimizer import FairseqOptimizer - except ImportError: - pass - - expected_optim_types = [Optimizer] - if FairseqOptimizer: - # fairseq optims are not torch.optim objects - expected_optim_types.append(FairseqOptimizer) - return expected_optim_types - - # Validate configuration based on command line arguments - def _do_sanity_check(self): - if self.fp16_enabled() and not get_accelerator().is_fp16_supported(): - raise ValueError("Type fp16 is not supported on your device.") - - if self.bfloat16_enabled() and not get_accelerator().is_bf16_supported(): - raise ValueError("Type bf16 is not supported on your device.") - - expected_optim_types = self._supported_optims() - expected_optim_types += [type(None), Callable] - assert isinstance(self.client_optimizer, tuple(expected_optim_types)), \ - f'Client Optimizer is of unexpected type {type(self.client_optimizer)}' - - if not self.client_optimizer: - if self.optimizer_name() is not None: - assert self._is_supported_optimizer( - self.optimizer_name()), "{} is not a supported DeepSpeed Optimizer".format(self.optimizer_name()) - - if (self.optimizer_name() == LAMB_OPTIMIZER or self.optimizer_name() == ONEBIT_LAMB_OPTIMIZER): - assert (self.dynamic_loss_scale()), "DeepSpeed {} optimizer requires dynamic loss scaling".format( - self.optimizer_name()) - - # Detect invalid combinations of client optimizer and client scheduler - if isinstance(self.client_lr_scheduler, _LRScheduler): - assert isinstance(self.client_optimizer, Optimizer), \ - f'Client Optimizer (type = {type(self.client_optimizer)} is not instantiated but Client LR Scheduler is instantiated' - - def _broadcast_model(self): - - def is_replicated(p): - if hasattr(p, "ds_status") and p.ds_status is not ZeroParamStatus.AVAILABLE: - return False - elif hasattr(p, 'ds_optim_param'): - # do not broadcast OptimizedLinear parameters, they are unique per base weight shard - return False - return True - - for n, p in self.module.named_parameters(): - # Broadcast the model for different parameters - if is_moe_param(p): - if torch.is_tensor(p) and is_replicated(p): - dist.broadcast(p.data, - groups._get_expert_broadcast_src_rank(p.group_name), - group=self.expert_data_parallel_group[p.group_name]) - else: - if torch.is_tensor(p) and is_replicated(p): - dist.broadcast(p.data, groups._get_broadcast_src_rank(), group=self.seq_data_parallel_group) - - @staticmethod - def __check_params(model: Module, dtype: torch.dtype) -> None: - return - if not all(param.dtype == dtype for param in model.parameters()) and dist.get_rank() == 0: - raise ValueError(f"{dtype} is enabled but the following parameters have dtype that is " - f"not {dtype}: " - f"{[(n, p.dtype) for n, p in model.named_parameters() if p.dtype != dtype]}") - - def _set_client_model(self, model): - # register client model in _modules so that nn.module methods work correctly - modules = self.__dict__.get('_modules') - modules['module'] = model - # register module attribute in engine but avoid getattr - self.__dict__['module'] = model - - def _configure_distributed_model(self, model): - self._set_client_model(model) - is_zero_init_model = self.zero_optimization_partition_weights() and any( - [hasattr(param, "ds_id") for param in self.module.parameters()]) - - if self.fp16_enabled(): - if is_zero_init_model: - self.__check_params(self.module, torch.half) - self.module.half() - elif self.bfloat16_enabled(): - if is_zero_init_model: - self.__check_params(self.module, torch.bfloat16) - self.module.bfloat16() - else: - self.__check_params(self.module, torch.float) - - # zero.Init() handles device placement of model - if not (self.dont_change_device or is_zero_init_model): - self.module.to(self.device) - - # MoE related initialization - for _, module in self.module.named_modules(): - if isinstance(module, MoE): - self.has_moe_layers = True - self.num_experts.append(module.num_experts) - - if self.has_moe_layers: - for _, module in self.module.named_modules(): - if isinstance(module, TopKGate): - self.gate_modules.append(module) - if self.wall_clock_breakdown(): - module.wall_clock_breakdown = True - if isinstance(module, MOELayer): - self.moe_layers.append(module) - if self.wall_clock_breakdown(): - module.wall_clock_breakdown = True - - # Pass the mpu from here to groups. For subsequent use, just query groups - if self.mpu is not None: - groups.mpu = self.mpu - - # Set deepspeed parallelism spec. for the model including expert parallelism - for _, module in self.module.named_modules(): - if hasattr(module, 'set_deepspeed_parallelism'): - module.set_deepspeed_parallelism(self._config.use_data_before_expert_parallel_) - - # Query the groups module to get information about various parallel groups - self.local_all_to_all_group = None - if self.zero_quantized_gradients(): - message = "Using LoCo quantized gradients" if self.zeropp_loco_param() else "Using quantized gradients" - log_dist(message, ranks=[0]) - self.local_all_to_all_group = groups._get_local_all_to_all_group() - self.data_parallel_group = groups._get_data_parallel_group() - self.dp_world_size = groups._get_data_parallel_world_size() - self.seq_data_parallel_group = groups._get_sequence_data_parallel_group() - self.seq_dp_world_size = groups._get_sequence_data_parallel_world_size() - self.mp_world_size = groups._get_model_parallel_world_size() - self.expert_parallel_group = groups._get_expert_parallel_group_dict() - self.expert_data_parallel_group = groups._get_expert_data_parallel_group_dict() - self.sequence_parallel_size = groups._get_sequence_parallel_world_size() - if self.sequence_parallel_size > 1: - self.communication_data_type = self._config.seq_parallel_communication_data_type - self.seq_parallel_group = groups._get_sequence_parallel_group() - - if dist.get_rank() == 0: - summary = "********** distributed groups summary **********\n" - summary += f"\t {self.dp_world_size=}\n" - summary += f"\t {self.mp_world_size=}\n" - summary += f"\t {self.seq_dp_world_size=}\n" - summary += f"\t {self.sequence_parallel_size=}\n" - summary += "***********************************************" - logger.info(summary) - - if not (self.amp_enabled() or is_zero_init_model): - self._broadcast_model() - - # check if parameters are duplicated in optimizer param_groups - def _check_for_duplicates(self, optimizer): - for name, param in self.module.named_parameters(): - param_id = id(param) - - def ids_list(group): - return [id(param) for param in group] - - occurrence = sum([ - ids_list(group['params']).count(param_id) if param_id in ids_list(group['params']) else 0 - for group in optimizer.param_groups - ]) - assert occurrence <= 1, f"Parameter with name: {name} occurs multiple times in optimizer.param_groups. Make sure it only appears once to prevent undefined behavior." - - def _do_optimizer_sanity_check(self, basic_optimizer): - model_dtype, grad_accum_dtype = self.get_data_types() - zero_enabled = self.zero_optimization() - amp_enabled = self.amp_enabled() - # config based assertions - assert ( - not (amp_enabled and zero_enabled) - ), "Amp and ZeRO are not currently compatible, please use (legacy) fp16 mode which performs similar to amp opt_mode=O2" - if zero_enabled: - if not is_zero_supported_optimizer(basic_optimizer): - assert ( - self.zero_allow_untested_optimizer() - ), 'You are using an untested ZeRO Optimizer. Please add <"zero_allow_untested_optimizer": true> in the configuration file to use it.' - - if self.global_rank == 0: - logger.warning("**** You are using ZeRO with an untested optimizer, proceed with caution *****") - if model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32 and self.zero_optimization_stage( - ) == 1 and not self.zero_cpu_offload(): - return BFLOAT16 - return ZERO_OPTIMIZATION - elif amp_enabled: - if model_dtype != grad_accum_dtype: - raise NotImplementedError( - "Model data type and gradient accumulation data type must be equal to use Amp") - if model_dtype == torch.bfloat16 or model_dtype == torch.float16: - raise NotImplementedError("Cannot enable both amp with (legacy) fp16 or bfloat16 mode") - try: - logger.info("Initializing Apex amp from: {}".format(amp.__path__)) - except NameError: - # If apex/amp is available it will be imported above - raise RuntimeError("Unable to import apex/amp, please make sure it is installed") - return AMP - # data type checks - elif model_dtype == grad_accum_dtype: - if model_dtype == torch.bfloat16: - if self.pipeline_parallelism: - logger.warning( - "**** BF16 gradient accumulation is not safe numerically with large number of accumulation steps, proceed with caution *****" - ) - return BFLOAT16 - else: - raise NotImplementedError( - "Bfloat16 wrapper must use a gradient accumulation type of fp32, enable ZeRO to use Bfloat16 gradient accumulation" - ) - if model_dtype == torch.float16: - return FP16 - # else optimizer_wrapper = None - elif model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32: - return BFLOAT16 - else: - raise NotImplementedError("unsupported mix of model dtype and gradient accumulation type") - - return None - - # Configure optimizer - def _configure_optimizer(self, client_optimizer, model_parameters): - if client_optimizer is None: - if self.has_moe_layers: - model_parameters = configure_moe_param_groups(model_parameters) - basic_optimizer = self._configure_basic_optimizer(model_parameters) - log_dist(f"Using DeepSpeed Optimizer param name {self.optimizer_name()} as basic optimizer", ranks=[0]) - else: - if isinstance(client_optimizer, tuple(self._supported_optims())): - basic_optimizer = client_optimizer - log_dist('Using client Optimizer as basic optimizer', ranks=[0]) - else: - basic_optimizer = client_optimizer(model_parameters) - log_dist('Using client callable to create basic optimizer', ranks=[0]) - - if self.zero_use_cpu_optimizer() and not isinstance(basic_optimizer, deepspeed.ops.adam.DeepSpeedCPUAdam): - if self.zero_force_ds_cpu_optimizer(): - msg = f'You are using ZeRO-Offload with a client provided optimizer ({type(basic_optimizer)}) which in most cases will yield poor performance. Please either use deepspeed.ops.adam.DeepSpeedCPUAdam or set an optimizer in your ds-config (https://www.deepspeed.ai/docs/config-json/#optimizer-parameters). If you really want to use a custom optimizer w. ZeRO-Offload and understand the performance impacts you can also set <"zero_force_ds_cpu_optimizer": false> in your configuration file.' - raise ZeRORuntimeException(msg) - - basic_optimizer.param_groups[:] = [pg for pg in basic_optimizer.param_groups if len(pg["params"]) != 0] - log_dist("Removing param_group that has no 'params' in the basic Optimizer", ranks=[0]) - - self._check_for_duplicates(basic_optimizer) - - self.basic_optimizer = basic_optimizer - log_dist(f"DeepSpeed Basic Optimizer = {basic_optimizer.__class__.__name__}", ranks=[0]) - - optimizer_wrapper = self._do_optimizer_sanity_check(basic_optimizer) - - if optimizer_wrapper == ZERO_OPTIMIZATION: - self.optimizer = self._configure_zero_optimizer(basic_optimizer) - elif optimizer_wrapper == AMP: - amp_params = self.amp_params() - log_dist(f"Initializing AMP with these params: {amp_params}", ranks=[0]) - model, self.optimizer = amp.initialize(self.module, basic_optimizer, **amp_params) - self._set_client_model(model) - self._broadcast_model() - # TODO: maybe need to broadcast experts differently? - elif optimizer_wrapper == FP16: - self.optimizer = self._configure_fp16_optimizer(basic_optimizer) - elif optimizer_wrapper == BFLOAT16: - self.optimizer = self._configure_bf16_optimizer(basic_optimizer) - else: - self.optimizer = basic_optimizer - - log_dist("DeepSpeed Final Optimizer = {}".format(self.optimizer.__class__.__name__), ranks=[0]) - - self.compression_scheduler = self._configure_compression_scheduler() - self.quantizer = self._configure_quantization() - - def _configure_basic_optimizer(self, model_parameters): - optimizer_parameters = self.optimizer_params() - if optimizer_parameters is None: - optimizer_parameters = {} - # print(optimizer_parameters.keys()) - if "max_grad_norm" in optimizer_parameters.keys(): - raise ValueError( - "'max_grad_norm' is not supported as an optimizer parameter, please switch to using the deepspeed parameter 'gradient_clipping' see: https://www.deepspeed.ai/docs/config-json/#gradient-clipping for more details" - ) - - if self.optimizer_name() in [ADAM_OPTIMIZER, ADAMW_OPTIMIZER]: - torch_adam = optimizer_parameters.pop(TORCH_ADAM_PARAM, False) - adam_w_mode = optimizer_parameters.pop(ADAM_W_MODE, ADAM_W_MODE_DEFAULT) - - # Optimizer name of Adam forces AdamW logic unless adam_w_mode is explicitly set - effective_adam_w_mode = self.optimizer_name() == ADAMW_OPTIMIZER or adam_w_mode - - if torch_adam: - if not effective_adam_w_mode: - optimizer = torch.optim.Adam(model_parameters, **optimizer_parameters) - else: - optimizer = torch.optim.AdamW(model_parameters, **optimizer_parameters) - else: - if self.zero_use_cpu_optimizer(): - from deepspeed.ops.adam import DeepSpeedCPUAdam - optimizer = DeepSpeedCPUAdam(model_parameters, - **optimizer_parameters, - adamw_mode=effective_adam_w_mode) - else: - from deepspeed.ops.adam import FusedAdam - - optimizer = FusedAdam( - model_parameters, - **optimizer_parameters, - adam_w_mode=effective_adam_w_mode, - ) - - elif self.optimizer_name() == ADAGRAD_OPTIMIZER: - if self.zero_use_cpu_optimizer(): - from deepspeed.ops.adagrad import DeepSpeedCPUAdagrad - optimizer = DeepSpeedCPUAdagrad(model_parameters, **optimizer_parameters) - else: - optimizer = torch.optim.Adagrad(model_parameters, **optimizer_parameters) - elif self.optimizer_name() == LAMB_OPTIMIZER: - from deepspeed.ops.lamb import FusedLamb - - optimizer = FusedLamb(model_parameters, **optimizer_parameters) - elif self.optimizer_name() == ONEBIT_ADAM_OPTIMIZER: - assert not self.zero_optimization(), "1bit-Adam is not compatible with ZeRO" - from deepspeed.runtime.fp16.onebit.adam import OnebitAdam - - optimizer = OnebitAdam(model_parameters, self, **optimizer_parameters) - if not self.fp16_enabled(): - logger.warning(f"Currently the convergence of 1-bit Adam is only verified under FP16") - elif self.optimizer_name() == ZERO_ONE_ADAM_OPTIMIZER: - assert not self.zero_optimization(), "0/1 Adam is not compatible with ZeRO" - from deepspeed.runtime.fp16.onebit.zoadam import ZeroOneAdam - - optimizer = ZeroOneAdam(model_parameters, self, **optimizer_parameters) - if not self.fp16_enabled(): - logger.warning(f'Currently the convergence of 0/1 Adam is only verified under FP16') - elif self.optimizer_name() == ONEBIT_LAMB_OPTIMIZER: - assert not self.zero_optimization(), "1bit-Lamb is not compatible with ZeRO" - from deepspeed.runtime.fp16.onebit.lamb import OnebitLamb - - optimizer = OnebitLamb(model_parameters, self, **optimizer_parameters) - if not self.fp16_enabled(): - logger.warning(f"Currently the convergence of 1-bit Lamb is only verified under FP16") - elif self.optimizer_name() == LION_OPTIMIZER: - if self.zero_use_cpu_optimizer(): - from deepspeed.ops.lion import DeepSpeedCPULion - optimizer = DeepSpeedCPULion(model_parameters, **optimizer_parameters) - else: - from deepspeed.ops.lion import FusedLion - optimizer = FusedLion(model_parameters, **optimizer_parameters) - elif self.optimizer_name() == MUADAM_OPTIMIZER: - try: - from mup import MuAdam - except ImportError: - logger.error(f"Install mup to use MuAdam optimizer") - optimizer = MuAdam(model_parameters, **optimizer_parameters) - elif self.optimizer_name() == MUADAMW_OPTIMIZER: - try: - from mup import MuAdamW - except ImportError: - logger.error(f"Install mup to use MuAdamW optimizer") - optimizer = MuAdamW(model_parameters, **optimizer_parameters) - elif self.optimizer_name() == MUSGD_OPTIMIZER: - try: - from mup import MuSGD - except ImportError: - logger.error(f"Install mup to use MuSGD optimizer") - optimizer = MuSGD(model_parameters, **optimizer_parameters) - else: - torch_optimizer = getattr(torch.optim, self.optimizer_name()) - optimizer = torch_optimizer(model_parameters, **optimizer_parameters) - return optimizer - - def _configure_compression_scheduler(self): - return compression_scheduler(self.module, self._config.compression_config) - - def _configure_random_ltd_scheduler(self, configs): - return RandomLTDScheduler(configs) - - def _configure_quantization(self): - ( - quantize_weight_in_forward, - quantize_enabled, - q_groups, - q_mixed_fp16, - q_change_ratio, - q_type, - q_rounding, - q_verbose, - use_quantizer_kernel, - ) = self.quantize_training() - if quantize_enabled and not quantize_weight_in_forward: - assert self.fp16_enabled( - ), "MoQ (quantize in optimization step) weight quantization is only supported for FP16" - quantizer = None - if quantize_enabled and not quantize_weight_in_forward: - from deepspeed.runtime.quantize import Quantizer - - quantizer = Quantizer( - q_groups, - q_mixed_fp16, - q_change_ratio, - q_type, - q_rounding, - q_verbose, - self.eigenvalue_enabled(), - use_quantizer_kernel, - self.eigenvalue_layer_num() if self.eigenvalue_enabled() else 0, - ) - return quantizer - - def _configure_fp16_optimizer(self, optimizer): - initial_dynamic_scale = self.initial_dynamic_scale() - dynamic_loss_args = self.dynamic_loss_scale_args() - clip_grad = self.gradient_clipping() - - if APEX_INSTALLED: - fused_opts = (apex.optimizers.FusedAdam, FusedAdam) - else: - fused_opts = FusedAdam - - if isinstance(optimizer, fused_opts) \ - or self.optimizer_name() in [ONEBIT_ADAM_OPTIMIZER, ZERO_ONE_ADAM_OPTIMIZER]: - if self.dynamic_loss_scale(): - log_dist(f'Creating fp16 optimizer with dynamic loss scale', ranks=[0]) - timers = self.timers if self.wall_clock_breakdown() else NoopTimer() - optimizer = FP16_Optimizer( - optimizer, - deepspeed=self, - dynamic_loss_scale=True, - initial_dynamic_scale=initial_dynamic_scale, - dynamic_loss_args=dynamic_loss_args, - mpu=self.mpu, - clip_grad=clip_grad, - fused_adam_legacy=self.optimizer_legacy_fusion(), - timers=timers, - has_moe_layers=self.has_moe_layers, - ) - else: - log_dist(f'Creating fp16 optimizer with static loss scale: {self.loss_scale()}', ranks=[0]) - timers = self.timers if self.wall_clock_breakdown() else NoopTimer() - optimizer = FP16_Optimizer( - optimizer, - deepspeed=self, - static_loss_scale=self.loss_scale(), - mpu=self.mpu, - clip_grad=clip_grad, - fused_adam_legacy=self.optimizer_legacy_fusion(), - timers=timers, - has_moe_layers=self.has_moe_layers, - ) - else: - log_dist(f'Creating fp16 unfused optimizer with dynamic loss scale', ranks=[0]) - optimizer = FP16_UnfusedOptimizer( - optimizer, - deepspeed=self, - static_loss_scale=self.loss_scale(), - dynamic_loss_scale=self.dynamic_loss_scale(), - dynamic_loss_args=dynamic_loss_args, - mpu=self.mpu, - clip_grad=clip_grad, - fused_lamb_legacy=self.optimizer_name() == LAMB_OPTIMIZER, - ) - - return optimizer - - def _configure_bf16_optimizer(self, optimizer): - clip_grad = self.gradient_clipping() - - if optimizer is None: - optimizer = DummyOptim(list(self.module.parameters())) - - log_dist('Creating BF16 optimizer', ranks=[0]) - - timers = self.timers if self.wall_clock_breakdown() else NoopTimer() - optimizer = BF16_Optimizer(optimizer, - self.param_names, - bfloat16_config=self._config.bfloat16_config, - mpu=self.mpu, - clip_grad=clip_grad, - allgather_bucket_size=self.zero_allgather_bucket_size(), - dp_process_group=self.seq_data_parallel_group, - timers=timers, - grad_acc_dtype=self.get_data_types()[1], - graph_harvesting=self.graph_harvesting(), - has_moe_layers=self.has_moe_layers) - - return optimizer - - def _configure_zero_optimizer(self, optimizer): - zero_stage = self.zero_optimization_stage() - - mics_shard_size = self.mics_shard_size() - model_dtype, gradient_accumulation_dtype = self.get_data_types() - - if self.bfloat16_enabled(): - check_grad_overflow = self._config.bfloat16_config.check_grad_overflow - elif self.fp16_enabled(): - check_grad_overflow = True - else: - check_grad_overflow = False - - timers = self.timers if self.wall_clock_breakdown() else NoopTimer() - - if optimizer is None: - optimizer = DummyOptim(list(self.module.parameters())) - - if self.zero_legacy_stage1(): - raise Exception( - "The deprecated version of ZeRO Stage 1 is not supported in deepspeed >= 0.5.9. Please downgrade to a version less than 0.5.9 if you need to use this deprecated version of ZeRO." - ) - - if zero_stage <= ZeroStageEnum.gradients: - overlap_comm = self.zero_overlap_comm() - contiguous_gradients = self.zero_contiguous_gradients() - round_robin_gradients = self.zero_round_robin_gradients() - assert not isinstance(optimizer, DummyOptim), "zero stage {} requires an optimizer".format(zero_stage) - - log_dist(f'Creating {model_dtype} ZeRO stage {zero_stage} optimizer', ranks=[0]) - - if isinstance(self.module, PipelineModule): - if overlap_comm: - logger.warning("Pipeline parallelism does not support overlapped communication, will be disabled.") - overlap_comm = False - optimizer = DeepSpeedZeroOptimizer( - optimizer, - self.param_names, - timers=timers, - static_loss_scale=self.loss_scale(), - dynamic_loss_scale=self.dynamic_loss_scale(), - dynamic_loss_args=self.dynamic_loss_scale_args(), - clip_grad=self.gradient_clipping(), - contiguous_gradients=contiguous_gradients, - reduce_bucket_size=self.zero_reduce_bucket_size(), - use_multi_rank_bucket_allreduce=self.zero_multi_rank_bucket_allreduce(), - allgather_bucket_size=self.zero_allgather_bucket_size(), - dp_process_group=self.seq_data_parallel_group, - expert_parallel_group=self.expert_parallel_group if self.has_moe_layers else None, - expert_data_parallel_group=self.expert_data_parallel_group if self.has_moe_layers else None, - reduce_scatter=self.zero_reduce_scatter(), - overlap_comm=overlap_comm, - offload_optimizer_config=self.zero_offload_optimizer(), - mpu=self.mpu, - postscale_gradients=self.postscale_gradients(), - gradient_predivide_factor=self.gradient_predivide_factor(), - gradient_accumulation_steps=self.gradient_accumulation_steps(), - ignore_unused_parameters=self.zero_ignore_unused_parameters(), - partition_grads=zero_stage == ZeroStageEnum.gradients, - round_robin_gradients=round_robin_gradients, - has_moe_layers=self.has_moe_layers, - fp16_master_weights_and_gradients=self.fp16_master_weights_and_gradients(), - gradient_accumulation_dtype=gradient_accumulation_dtype, - communication_data_type=self.communication_data_type, - elastic_checkpoint=self.zero_elastic_checkpoint(), - check_grad_overflow=check_grad_overflow) - - elif zero_stage == ZeroStageEnum.weights: - assert not self.has_moe_layers, "MoE not supported with Stage 3" - if isinstance(optimizer, DummyOptim): - log_dist("Creating ZeRO Offload", ranks=[0]) - zero_param_parallel_group = groups._get_zero_param_intra_parallel_group() - if self.zero_hpz_partition_size() > 1 and zero_param_parallel_group is None: - self._set_zero_group_parallelism() - zero_param_parallel_group = groups._get_zero_param_intra_parallel_group() - optimizer = DeepSpeedZeRoOffload( - self.module, - timers=timers, - ds_config=self.config, - overlap_comm=self.zero_overlap_comm(), - prefetch_bucket_size=self.zero_prefetch_bucket_size(), - max_reuse_distance=self.zero_max_reuse_distance(), - max_live_parameters=self.zero_max_live_parameters(), - param_persistence_threshold=self.zero_param_persistence_threshold(), - model_persistence_threshold=self.zero_model_persistence_threshold(), - offload_param_config=self.zero_offload_param(), - mpu=self.mpu, - zero_param_parallel_group=zero_param_parallel_group, - zero_quantized_weights=self.zero_quantized_weights(), - zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights(), - zero_module_granularity_threshold=self.zero_module_granularity_threshold(), - log_trace_cache_warnings=self.zero_log_trace_cache_warnings(), - ) - else: - log_dist( - f'Creating fp16 ZeRO stage {zero_stage} optimizer,' - f' MiCS is enabled {mics_shard_size>0},' - f' Hierarchical params gather {self._config.mics_hierarchial_params_gather}', - ranks=[0]) - if mics_shard_size > 0: - return self._return_mics_optimizer(optimizer, timers) - - log_dist(f'Creating {model_dtype} ZeRO stage {zero_stage} optimizer', ranks=[0]) - from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 - optimizer = DeepSpeedZeroOptimizer_Stage3( - self.module, - optimizer, - timers=timers, - ds_config=self.config, - static_loss_scale=self.loss_scale(), - dynamic_loss_scale=self.dynamic_loss_scale(), - dynamic_loss_args=self.dynamic_loss_scale_args(), - clip_grad=self.gradient_clipping(), - contiguous_gradients=self.zero_contiguous_gradients(), - reduce_bucket_size=self.zero_reduce_bucket_size(), - prefetch_bucket_size=self.zero_prefetch_bucket_size(), - max_reuse_distance=self.zero_max_reuse_distance(), - max_live_parameters=self.zero_max_live_parameters(), - param_persistence_threshold=self.zero_param_persistence_threshold(), - model_persistence_threshold=self.zero_model_persistence_threshold(), - dp_process_group=self.seq_data_parallel_group, - all2all_process_group=self.local_all_to_all_group, - reduce_scatter=self.zero_reduce_scatter(), - overlap_comm=self.zero_overlap_comm(), - offload_optimizer_config=self.zero_offload_optimizer(), - offload_param_config=self.zero_offload_param(), - sub_group_size=self.zero_sub_group_size(), - offload_ratio=self.zero_partial_offload(), - mpu=self.mpu, - postscale_gradients=self.postscale_gradients(), - gradient_predivide_factor=self.gradient_predivide_factor(), - gradient_accumulation_steps=self.gradient_accumulation_steps(), - aio_config=self.aio_config(), - gradient_accumulation_dtype=gradient_accumulation_dtype, - communication_data_type=self.communication_data_type, - zero_hpz_partition_size=self.zero_hpz_partition_size(), - zero_quantized_weights=self.zero_quantized_weights(), - zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights(), - zero_module_granularity_threshold=self.zero_module_granularity_threshold(), - zeropp_loco_param=self.zeropp_loco_param(), - log_trace_cache_warnings=self.zero_log_trace_cache_warnings(), - ) - - else: - raise NotImplementedError("ZeRO stage {} not implemented".format(zero_stage)) - - return optimizer - - def _return_mics_optimizer(self, basic_optimizer, timers): - from deepspeed.runtime.zero.mics import MiCS_Optimizer - model_dtype, gradient_accumulation_dtype = self.get_data_types() - optimizer = MiCS_Optimizer(self.module, - basic_optimizer, - timers=timers, - ds_config=self.config, - static_loss_scale=self.loss_scale(), - dynamic_loss_scale=self.dynamic_loss_scale(), - dynamic_loss_args=self.dynamic_loss_scale_args(), - clip_grad=self.gradient_clipping(), - contiguous_gradients=self.zero_contiguous_gradients(), - reduce_bucket_size=self.zero_reduce_bucket_size(), - prefetch_bucket_size=self.zero_prefetch_bucket_size(), - max_reuse_distance=self.zero_max_reuse_distance(), - max_live_parameters=self.zero_max_live_parameters(), - param_persistence_threshold=self.zero_param_persistence_threshold(), - model_persistence_threshold=self.zero_model_persistence_threshold(), - dp_process_group=self.seq_data_parallel_group, - reduce_scatter=self.zero_reduce_scatter(), - overlap_comm=self.zero_overlap_comm(), - offload_optimizer_config=self.zero_offload_optimizer(), - offload_param_config=self.zero_offload_param(), - sub_group_size=self.zero_sub_group_size(), - mpu=self.mpu, - postscale_gradients=self.postscale_gradients(), - gradient_predivide_factor=self.gradient_predivide_factor(), - gradient_accumulation_steps=self.gradient_accumulation_steps(), - aio_config=self.aio_config(), - gradient_accumulation_dtype=gradient_accumulation_dtype, - communication_data_type=self.communication_data_type) - return optimizer - - def _configure_eigenvalue(self): - eigenvalue = Eigenvalue( - verbose=self.eigenvalue_verbose(), - max_iter=self.eigenvalue_max_iter(), - tol=self.eigenvalue_tol(), - stability=self.eigenvalue_stability(), - gas_boundary_resolution=self.eigenvalue_gas_boundary_resolution(), - layer_name=self.eigenvalue_layer_name(), - layer_num=self.eigenvalue_layer_num(), - ) - - return eigenvalue - - def _configure_progressive_layer_drop(self): - pld = ProgressiveLayerDrop(theta=self.pld_theta(), gamma=self.pld_gamma()) - - return pld - - def _configure_curriculum_scheduler_legacy(self): - scheduler = CurriculumScheduler(self.curriculum_params_legacy()) - return scheduler - - @staticmethod - def is_map_style_dataset(obj): - return hasattr(obj, "__getitem__") and hasattr(obj, "__len__") - - @staticmethod - def is_iterable_style_dataset(obj): - return isinstance(obj, torch.utils.data.IterableDataset) # hasattr(obj, "__iter__") should work as well - - def dataloader_drop_last(self): - return self._config.dataloader_drop_last - - def was_step_applied(self) -> bool: - """Returns True if the latest ``step()`` produced in parameter updates. - Note that a ``False`` return is not an error condition. Steps are frequently - no-ops, such as between gradient accumulation boundaries or when overflows - occur. - Returns: - bool: Whether the latest ``step()`` modified model parameters. - """ - return self._step_applied - - def deepspeed_io(self, - dataset, - batch_size=None, - route=ROUTE_TRAIN, - pin_memory=True, - data_sampler=None, - collate_fn=None, - num_local_io_workers=None): - if not (self.is_map_style_dataset(dataset) or self.is_iterable_style_dataset(dataset)): - raise ValueError("Training data must be a torch Dataset") - - if batch_size is None: - batch_size = self.train_micro_batch_size_per_gpu() - - if collate_fn is None: - collate_fn = self.collate_fn - - # Currently we only use timer in train route - deepspeed_io_timer = None - if route == ROUTE_TRAIN: - deepspeed_io_timer = self.tput_timer - - # If mpu is provided, forward world size and parallel rank to sampler. - data_parallel_world_size = self.dp_world_size - data_parallel_rank = self.global_rank - if self.mpu is not None: - data_parallel_world_size = self.mpu.get_data_parallel_world_size() - data_parallel_rank = self.mpu.get_data_parallel_rank() - - if data_sampler is None and (route == ROUTE_PREDICT or route == ROUTE_EVAL): - data_sampler = torch.utils.data.DistributedSampler( - dataset, - num_replicas=data_parallel_world_size, - rank=data_parallel_rank, - shuffle=False, - ) - - deepspeed_dataloader_config = {} - if self.curriculum_learning_enabled(): - deepspeed_dataloader_config = { - CURRICULUM_LEARNING: self.curriculum_learning_enabled(), - DATA_EFFICIENCY: self.data_efficiency_config(), - DATA_PARALLEL_GROUP: self.data_parallel_group, - GRADIENT_ACCUMULATION_STEPS: self.gradient_accumulation_steps(), - GLOBAL_RANK: self.global_rank, - DATA_SAMPLING_NUM_WORKERS: self.data_sampling_config()[DATA_SAMPLING_NUM_WORKERS] - } - return DeepSpeedDataLoader(dataset=dataset, - batch_size=batch_size, - pin_memory=pin_memory, - collate_fn=collate_fn, - local_rank=self.local_rank, - tput_timer=deepspeed_io_timer, - num_local_io_workers=num_local_io_workers, - data_sampler=data_sampler, - data_parallel_world_size=data_parallel_world_size, - data_parallel_rank=data_parallel_rank, - dataloader_drop_last=self.dataloader_drop_last(), - deepspeed_dataloader_config=deepspeed_dataloader_config) - - def train(self, mode=True): - r"""""" - - self.warn_unscaled_loss = True - self.module.train(mode) - - def eval(self): - r"""""" - - self.warn_unscaled_loss = True - self.module.train(False) - - def _scale_loss_by_gas(self, prescaled_loss, eval_micro_batches=None): - # In pipeline evaluation, there is an option to use different micro-bs, which creates different number of - # micro batches, thus the training gas, is not valid in this case. need to use the number of eval_micro_batches - scaling_factor = self.gradient_accumulation_steps() if eval_micro_batches is None else eval_micro_batches - if isinstance(prescaled_loss, torch.Tensor): - scaled_loss = prescaled_loss / scaling_factor - elif isinstance(prescaled_loss, tuple) or isinstance(prescaled_loss, list): - scaled_loss = [] - for l in prescaled_loss: - if isinstance(l, torch.Tensor): - scaled_loss.append(l / scaling_factor) - else: - scaled_loss.append(l) - else: - scaled_loss = prescaled_loss - if self.warn_unscaled_loss: - logger.warning(f"DeepSpeed unable to scale loss because of type: {type(prescaled_loss)}") - self.warn_unscaled_loss = False - - return scaled_loss - - def _create_module_forward_pre_hook(self): - - def _module_forward_pre_hook(module, inputs, kwargs): - return self._forward_prologue(inputs, kwargs) - - return self.module.register_forward_pre_hook(_module_forward_pre_hook, prepend=False, with_kwargs=True) - - def _create_module_forward_post_hook(self): - - def _module_forward_post_hook(module, input, output): - self._forward_epilogue() - - return self.module.register_forward_hook(_module_forward_post_hook) - - def _forward_prologue(self, inputs, kwargs): - return_modified = False - - if not self.autotuning_profile_model_info(): - see_memory_usage("Engine before forward", force=self.memory_breakdown()) - - flops_profiler_active = (self.flops_profiler_enabled() - and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0) - - # used to check quantization happens at step 0! - if self.global_steps == 0 and hasattr(self, "compression_scheduler"): - self.compression_scheduler.step(step_zero_check=True) - if self.quantizer: - tensor_to_quantize = self.optimizer.bit16_groups if self.zero_optimization_stage( - ) == 2 else self.optimizer.fp16_groups - if self.compression_scheduler.weight_quantization_enabled: - self.quantizer.quantize( - tensor_to_quantize, - (self.optimizer.overflow if self.fp16_enabled() else False), - self.eigenvalue_enabled(), - None, - ) - return_modified = True - - if flops_profiler_active: - self.flops_profiler.start_profile(ignore_list=None) - - if kwargs is not None: - if self.module.training: - if self.progressive_layer_drop: - kwargs.update(self.progressive_layer_drop.get_state()) - - if self.__class__.__name__ != "PipelineEngine": - # TODO: The above if condition is a HACK since for PipelineEngine - # it's difficult to inject argument in forward pass. - if self.module.training and self.curriculum_enabled_legacy(): - self.curriculum_scheduler_legacy.update_difficulty(self.global_steps + 1) - if self.curriculum_params_legacy()["curriculum_type"] == "seqlen": - kwargs.update({"curriculum_seqlen": self.curriculum_scheduler_legacy.get_current_difficulty()}) - return_modified = True - - if self.module.training and self.random_ltd_enabled(): - self.random_ltd_scheduler.update_seq(self.global_steps) - - if self.training_dataloader is None: - self.tput_timer.start() - - self._start_timers(self.engine_timers.forward_timers) - - if self.zero_optimization_partition_weights(): - # Enable automated discovery of external parameters by indicating that - # we are in a forward pass. - for module in self.module.modules(): - module._parameters._in_forward = True - - if self.fp16_auto_cast(): - inputs = self._cast_inputs_half(inputs) - return_modified = True - - if return_modified: - return inputs, kwargs - - def _forward_epilogue(self): - if self.zero_optimization_partition_weights(): - # Disable automated discovery of external parameters - for module in self.module.modules(): - module._parameters._in_forward = False - - self._stop_timers(self.engine_timers.forward_timers) - - flops_profiler_active = (self.flops_profiler_enabled() - and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0) - - if flops_profiler_active: - self.flops_profiler.stop_profile() - - if not self.autotuning_profile_model_info(): - see_memory_usage("Engine after forward", force=self.memory_breakdown()) - - @instrument_w_nvtx - def forward(self, *inputs, **kwargs): - r"""Execute forward propagation - Arguments: - *inputs: Variable length input list - **kwargs: variable length keyword arguments - """ - if self.autotuning_profile_model_info(): - ma = get_ma_status() - - if self.is_deepcompile_enabled() and hasattr(self, "launch_compile_passes"): - # We can't have this in forward prologue as the compiler compiles hooks including the forward prologue. - self.launch_compile_passes(self.global_steps) - - validate_nested_autocast(self) - with torch.autocast(device_type=get_accelerator().device_name(), - dtype=self.torch_autocast_dtype(), - enabled=self.torch_autocast_enabled()): - loss = self.module(*inputs, **kwargs) - - if self.autotuning_profile_model_info(): - activation_mem = get_ma_status() - ma - self.autotuning_model_info["activation_mem_per_gpu"] = activation_mem - print_json_dist(self.autotuning_model_info, [0], path=self.autotuning_model_info_path()) - exit() - - return loss - - def _cast_inputs_half(self, inputs): - if isinstance(inputs, (list, tuple)): - new_inputs = [] - for v in inputs: - new_inputs.append(self._cast_inputs_half(v)) - return inputs.__class__(new_inputs) - elif isinstance(inputs, dict): - new_inputs = {} - for k, v in inputs.items(): - new_inputs[k] = self._cast_inputs_half(v) - return new_inputs - elif hasattr(inputs, 'half') and inputs.is_floating_point(): - return inputs.half() - else: - return inputs - - def print_forward_breakdown(self, fwd_time): - gate_time = 0.0 - moe_time = 0.0 - falltoall = 0.0 - salltoall = 0.0 - - for gate in self.gate_modules: - #logger.info(f"Individual TopK gate time: {gate.gate_time:.2f} ms") - gate_time += gate.gate_time - - for l in self.moe_layers: - #logger.info(f"MoE layer; total: {l.time_moe:.2f} ms, first alltoall: {l.time_falltoall:.2f}, second alltoall: {l.time_salltoall:.2f}") - moe_time += l.time_moe - falltoall += l.time_falltoall - salltoall += l.time_salltoall - - # TODO: Allreduce/average them across ranks for more accurate timing. - - # if deepspeed.comm.get_rank() == 0: - log_dist( - f"time (ms) | fwd: {fwd_time:.2f} (fwd_moe: {moe_time:.2f}, 1st_a2a: {falltoall:.2f}, 2nd_a2a: {salltoall:.2f}, top_k: {gate_time:.2f})", - ranks=[0]) - - @instrument_w_nvtx - def allreduce_gradients(self, bucket_size=MEMORY_OPT_ALLREDUCE_SIZE): - # Pass (PP) gas boundary flag to optimizer (required for zero) - self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary() - # ZeRO stage >= 2 communicates during non gradient accumulation boundaries as well - if self.zero_optimization_partition_gradients(): - self.optimizer.overlapping_partition_gradients_reduce_epilogue() - - # Communicate only at gradient accumulation boundaries - elif self.is_gradient_accumulation_boundary(): - if self.zero_optimization_stage() == ZeroStageEnum.optimizer_states and hasattr( - self.optimizer, 'reduce_gradients'): - self.optimizer.reduce_gradients(pipeline_parallel=self.pipeline_parallelism) - else: - grads = None - self.buffered_allreduce_fallback(grads=grads, elements_per_buffer=bucket_size) - - def _backward_prologue(self, loss, scale_wrt_gas=True): - see_memory_usage("Engine before backward", force=self.memory_breakdown()) - if self.scale_wrt_gas is not None: - scale_wrt_gas = self.scale_wrt_gas - - # scale loss w.r.t. gradient accumulation if reduction is not disabled - do_gradient_reduction = self.enable_backward_allreduce and not self.inside_no_sync_ctxt and not self.is_deepcompile_enabled( - ) - if do_gradient_reduction and self.gradient_accumulation_steps() > 1 and scale_wrt_gas: - loss = self._scale_loss_by_gas(loss.float()) - - # Log training loss - mean_loss = loss.mean().detach() - self.losses = mean_loss if self.losses is None else self.losses + mean_loss - if self.monitor.enabled: - if self.is_gradient_accumulation_boundary(): - if self.global_rank == 0: - self.summary_events = [( - f"Train/Samples/train_loss", - self.losses.item(), - self.global_samples, - )] - self.monitor.write_events(self.summary_events) - - if self.is_deepcompile_enabled(): - deepcompile_backward_prologue(self.is_gradient_accumulation_boundary()) - - return loss - - def _backward_epilogue(self): - self._start_timers(self.engine_timers.backward_reduce_timers) - if self.enable_backward_allreduce and not self.inside_no_sync_ctxt: - # Traditional code path that allreduces the module parameter grads - self.allreduce_gradients() - - self._stop_timers(self.engine_timers.backward_reduce_timers) - see_memory_usage("Engine after backward", force=self.memory_breakdown()) - - def _do_optimizer_backward(self, loss, retain_graph): - self._start_timers(self.engine_timers.backward_inner_timers) - if self.zero_optimization(): - self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary() - self.optimizer.backward(loss, retain_graph=retain_graph) - elif self.amp_enabled(): - # AMP requires delaying unscale when inside gradient accumulation boundaries - # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations - delay_unscale = not self.is_gradient_accumulation_boundary() - with amp.scale_loss(loss, self.optimizer, delay_unscale=delay_unscale) as scaled_loss: - scaled_loss.backward(retain_graph=retain_graph) - elif self.fp16_enabled(): - if self.eigenvalue_enabled(): - self.optimizer.backward(loss, create_graph=True, retain_graph=True) - else: - self.optimizer.backward(loss, retain_graph=retain_graph) - elif self.bfloat16_enabled(): - self.optimizer.backward(loss, retain_graph=retain_graph) - else: - if self.eigenvalue_enabled(): - loss.backward(create_graph=True, retain_graph=True) - else: - loss.backward(retain_graph=retain_graph) - self._stop_timers(self.engine_timers.backward_inner_timers) - - @contextmanager - def no_sync(self): - r""" - Context manager to disable gradient reduction during backward pass. - This context manager has the following effects on other DeepSpeed features: - 1. Incompatible with ZeRO stage 2/3 which rely on reduction for gradient partitioning. - 2. It is illegal to call engine.step() within the context manager. - 3. Tracking of gradient accumulation steps is disabled. - """ - assert not self.zero_optimization_partition_gradients(), \ - f"no_sync context manager is incompatible with gradient partitioning logic of ZeRO stage {self.zero_optimization_stage()}" - - assert not self.inside_no_sync_ctxt, f"no_sync context manager reentry is unsupported" - - self.inside_no_sync_ctxt = True - try: - yield - finally: - self.inside_no_sync_ctxt = False - - @instrument_w_nvtx - def backward(self, loss, retain_graph=False, scale_wrt_gas=True): - r"""Execute backward pass on the loss - Arguments: - loss: Torch tensor on which to execute backward propagation - retain_graph: bool, default: false - forward on user defined choice of retain_graph - """ - assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ - "must provide optimizer during init in order to use backward" - - self._start_timers(self.engine_timers.backward_timers) - loss = self._backward_prologue(loss, scale_wrt_gas) - self._do_optimizer_backward(loss, retain_graph) - self._backward_epilogue() - self._stop_timers(self.engine_timers.backward_timers) - - return loss - - def is_gradient_accumulation_boundary(self): - """ - Query whether the current micro-batch is at the boundary of - gradient accumulation, and thus will trigger gradient reductions and - an optimizer step. - - Returns: - bool: if the current step is a gradient accumulation boundary. - - """ - if self._is_gradient_accumulation_boundary is None: - return (self.micro_steps + 1) % \ - self.gradient_accumulation_steps() == 0 - else: - return self._is_gradient_accumulation_boundary - - def set_gradient_accumulation_boundary(self, is_boundary): - """ - Manually overrides the DeepSpeed engine's gradient accumulation boundary state, this is an optional - feature and should be used with care. The state should be set before to the intended - value before each forward/backward. The final forward/backward should have the - boundary state set to True. This style allows client code to only call engine.step() once after all - the gradient accumulation passes are complete. See example below: - .. code-block:: python - engine.set_gradient_accumulation_boundary(False) - for _ in range(gradient_accumulation_steps - 1): - micro_batch = next(data_loader) - loss = engine(micro_batch) - engine.backward(loss) - engine.set_gradient_accumulation_boundary(True) - micro_batch = next(data_loader) - loss = engine(micro_batch) - engine.backward(loss) - engine.step() - Arguments: - is_boundary (bool): are we at a gradient accumulation boundary or not? - """ - self._is_gradient_accumulation_boundary = is_boundary - self.optimizer.is_gradient_accumulation_boundary = is_boundary - - def zero_grad(self): - """ - Zero parameter grads. - """ - for param_name, param in self.module.named_parameters(): - param.grad = None - - def clip_fp32_gradients(self): - clip_grad_norm_(parameters=self.module.parameters(), max_norm=self.gradient_clipping(), mpu=self.mpu) - - def _take_model_step(self, lr_kwargs, block_eigenvalue={}): - if self.gradient_clipping() > 0.0: - if not (self.fp16_enabled() or self.bfloat16_enabled() or self.amp_enabled() or self.zero_optimization()): - self.clip_fp32_gradients() - elif self.amp_enabled(): - # AMP's recommended way of doing clipping - # https://nvidia.github.io/apex/advanced.html#gradient-clipping - master_params = amp.master_params(self.optimizer) - clip_grad_norm_(parameters=master_params, max_norm=self.gradient_clipping(), mpu=self.mpu) - self.optimizer.step() - - if hasattr(self.optimizer, '_global_grad_norm'): - self._global_grad_norm = self.optimizer._global_grad_norm - - # Quantize the updated parameter if there is no overflow - if self.quantizer: - tensor_to_quantize = self.optimizer.bit16_groups if self.zero_optimization_stage( - ) == 2 else self.optimizer.fp16_groups - if self.compression_scheduler.weight_quantization_enabled: - self.quantizer.quantize( - tensor_to_quantize, - (self.optimizer.overflow if self.fp16_enabled() else False), - self.eigenvalue_enabled(), - block_eigenvalue, - ) - # zero grad in basic optimizer could be unreliable and may not exhibit - # the behavior that we want - if self.bfloat16_enabled(): - # TODO: Temporary until bf16_optimizer and zero_optimizer are integrated - if self.zero_optimization() and hasattr(self.optimizer, "zero_grad"): - self.optimizer.zero_grad() - else: - pass - elif self.zero_optimization() or self.fp16_enabled() or self.amp_enabled(): - self.optimizer.zero_grad() - else: - self.zero_grad() - - # Check overflow here since in DS fp16 optimizer, the overflow is updated in above step() function. - overflow = False - if hasattr(self.optimizer, "overflow"): - overflow = self.optimizer.overflow - self._step_applied = not overflow - - if overflow: - self.skipped_steps += 1 - else: - self.compression_scheduler.step() - if self.lr_scheduler is not None: - try: - self.lr_scheduler.step(**(lr_kwargs or {})) - except TypeError: - # XXX Hack to work with Megatron 2.0 and DeepSpeed pipelines. - # We don't currently have a way to specify lr_kwargs from - # pipe_engine.train_batch() - self.lr_scheduler.step(self.train_batch_size()) - - if self.steps_per_print() is not None: - report_progress = self.global_rank == 0 if self.global_rank else True - if report_progress and (self.global_steps + 1) % self.steps_per_print() == 0: - self._report_progress(self.global_steps + 1) - - self.losses = None - self.global_steps += 1 - self.global_samples += self.train_batch_size() - - def step(self, lr_kwargs=None): - r"""Execute the weight update step after forward and backward propagation - on effective_train_batch. - """ - assert not self.inside_no_sync_ctxt, \ - "It is illegal to call Engine.step() inside no_sync context manager" - - see_memory_usage("Engine before step", force=self.memory_breakdown()) - - # Check early because self.global_steps is incremented at some point here. - # TODO: Delay self.global_steps increment until very end of this function. - flops_profiler_active = self.flops_profiler_enabled( - ) and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0 - - self._start_timers(self.engine_timers.step_timers) - - assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ - "must provide optimizer during init in order to use step" - - report_progress = False - - self._step_applied = False # assume False, will flip to True - - # Update the model when we reach gradient accumulation boundaries - if self.is_gradient_accumulation_boundary(): - self.gas_boundary_ctr += 1 - - if self.checkpoint_engine.is_decoupled(): - self._commit_decoupled_checkpoint() - - if (self.eigenvalue_enabled() and (self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() == 0) - and self.quantizer.any_precision_switch()): - log_dist(f"computing eigenvalue...", ranks=[0]) - self.block_eigenvalue = self.eigenvalue.compute_eigenvalue(self.module, self.device, - self.optimizer.cur_scale) - - if self.progressive_layer_drop: - self.progressive_layer_drop.update_state(self.global_steps) - - if (self.eigenvalue_enabled() and not self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() - and self.quantizer.any_precision_switch()): - self._take_model_step(lr_kwargs, self.block_eigenvalue) - else: - self._take_model_step(lr_kwargs) - - report_progress = self.global_rank == 0 if self.global_rank else True - - self.tput_timer.stop(global_step=self.is_gradient_accumulation_boundary(), report_speed=report_progress) - - self._stop_timers(self.engine_timers.step_timers) - - # Log learning rate - if self.monitor.enabled: - if self.is_gradient_accumulation_boundary(): - if self.global_rank == 0: - self.summary_events = [(f"Train/Samples/lr", self.get_lr()[0], self.global_samples)] - - if self.fp16_enabled() and hasattr(self.optimizer, "cur_scale"): - self.summary_events.append(( - f"Train/Samples/loss_scale", - self.optimizer.cur_scale, - self.global_samples, - )) - - if (self.eigenvalue_enabled() - and not self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution()): - ev_values = self.block_eigenvalue.values() - for i in range(len(ev_values)): - self.summary_events.append(( - f"Train/Eigenvalues/ModelBlockParam_{i}", - self.ev_values[i][0], - self.global_samples, - )) - self.monitor.write_events(self.summary_events) - - # Check flops profiling - if flops_profiler_active: - if self.autotuning_enabled(): - self.flops = self.flops_profiler.get_total_flops() * 3 - self.fwd_duration = self.flops_profiler.get_total_duration() - else: - self.flops_profiler.print_model_profile( - profile_step=self.global_steps, - module_depth=self.flops_profiler_module_depth(), - top_modules=self.flops_profiler_top_modules(), - detailed=self.flops_profiler_detailed(), - output_file=self.flops_profiler_output_file(), - ) - self.flops_profiler.end_profile() - - if self.autotuning_enabled() and self.global_steps == (self.autotuning_end_profile_step() + 1): - self._autotuning_exit() - - if self.wall_clock_breakdown(): - # Log micro timing and reset - self.timers.log(names=self.engine_timers.micro_timers, memory_breakdown=self.memory_breakdown()) - - if self.wall_clock_breakdown() or self.flops_profiler_enabled(): - # Log global timing and reset - if self.is_gradient_accumulation_boundary(): - if self.monitor.enabled: - self._write_monitor() - - if self.has_moe_layers: - fwd_time = self.timers(FORWARD_GLOBAL_TIMER).elapsed(reset=False) - self.print_forward_breakdown(fwd_time=fwd_time) - - self.timers.log(self.engine_timers.global_timers) - - self.micro_steps += 1 - see_memory_usage("Engine after step", force=self.memory_breakdown()) - - def _start_timers(self, timer_names): - for name in timer_names: - self.timers(name).start() - - def _stop_timers(self, timer_names): - record = self.is_gradient_accumulation_boundary() and \ - self.flops_profiler_enabled() and \ - (self.global_steps >= self.flops_profiler_profile_step()) - for name in timer_names: - self.timers(name).stop(record=record) - - def _autotuning_exit(self): - if self.global_rank == 0: - msg = self.timers.get_mean([ - FORWARD_GLOBAL_TIMER, - BACKWARD_GLOBAL_TIMER, - STEP_GLOBAL_TIMER, - ], reset=False) - titer = 0.0 - titer += msg[FORWARD_GLOBAL_TIMER] if FORWARD_GLOBAL_TIMER in msg else 0 - titer += msg[BACKWARD_GLOBAL_TIMER] if BACKWARD_GLOBAL_TIMER in msg else 0 - titer += msg[STEP_GLOBAL_TIMER] if STEP_GLOBAL_TIMER in msg else 0 - titer *= self.gradient_accumulation_steps() - msg["latency"] = titer - msg["FLOPS_per_gpu"] = self.flops * 1_000_000 * self.gradient_accumulation_steps() / titer - msg["throughput"] = self.train_batch_size() * 1_000_000 / \ - msg["latency"] - print_json_dist(msg, [0], path=self.autotuning_metric_path()) - log_dist( - f"Wrote metrics to {self.autotuning_metric_path()}, {os.path.abspath(self.autotuning_metric_path())}", - ranks=[0]) - import atexit - atexit.register(print, "Autotuning: done with running current ds config.") - exit() - - def _write_monitor(self): - if self.global_rank == 0: - self.summary_events = [ - ( - f"Train/Samples/elapsed_time_ms_forward", - self.timers(FORWARD_GLOBAL_TIMER).elapsed(reset=False), - self.global_samples, - ), - ( - f"Train/Samples/elapsed_time_ms_backward", - self.timers(BACKWARD_GLOBAL_TIMER).elapsed(reset=False), - self.global_samples, - ), - ( - f"Train/Samples/elapsed_time_ms_backward_inner", - self.timers(BACKWARD_INNER_GLOBAL_TIMER).elapsed(reset=False), - self.global_samples, - ), - ( - f"Train/Samples/elapsed_time_ms_backward_allreduce", - self.timers(BACKWARD_REDUCE_GLOBAL_TIMER).elapsed(reset=False), - self.global_samples, - ), - ( - f"Train/Samples/elapsed_time_ms_step", - self.timers(STEP_GLOBAL_TIMER).elapsed(reset=False), - self.global_samples, - ), - ] - self.monitor.write_events(self.summary_events) - - def _get_optimizer_param(self, param_name): - result = [] - if not self.optimizer: - return result - for group in self.optimizer.param_groups: - if param_name in group: - result.append(group[param_name]) - else: - result.append(0.0) - return result - - def get_lr(self): - return self._get_optimizer_param("lr") - - def get_type(self): - return self._get_optimizer_param("type") - - def get_mom(self): - if self.optimizer_name() in ["SGD", "RMSprop"]: - return self._get_optimizer_param("momentum") - else: - return self._get_optimizer_param("betas") - - def get_pld_theta(self): - if self.progressive_layer_drop: - return self.progressive_layer_drop.get_theta() - else: - return None - - def _report_progress(self, step): - lr = self.get_lr() - mom = self.get_mom() - log_dist(f"step={step}, skipped={self.skipped_steps}, lr={lr}, mom={mom}", ranks=[0]) - - def allreduce_bucket(self, bucket, dp_group, dp_world_size=None): - tensor = self.flatten(bucket) - - tensor_to_allreduce = tensor - - if self.communication_data_type != tensor.dtype: - tensor_to_allreduce = tensor.to(self.communication_data_type) - - if dp_world_size is None: - dp_world_size = dist.get_world_size(group=dp_group) - if self.postscale_gradients(): - if self.gradient_predivide_factor() != 1.0: - tensor_to_allreduce.mul_(1.0 / self.gradient_predivide_factor()) - - dist.all_reduce(tensor_to_allreduce, group=dp_group) - if self.gradient_average: - if self.gradient_predivide_factor() != dp_world_size: - tensor_to_allreduce.mul_(self.gradient_predivide_factor() / dp_world_size) - else: - tensor_to_allreduce.mul_(1. / dp_world_size) - dist.all_reduce(tensor_to_allreduce, group=dp_group) - - if self.communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: - tensor.copy_(tensor_to_allreduce) - - return tensor - - def allreduce_and_copy(self, small_bucket, dp_group, dp_world_size=None): - allreduced = self.allreduce_bucket(small_bucket, dp_group, dp_world_size) - for buf, synced in zip(small_bucket, self.unflatten(allreduced, small_bucket)): - buf.copy_(synced) - - def allreduce_no_retain(self, bucket, dp_group, numel_per_bucket=500000000, dp_world_size=None): - small_bucket = [] - numel = 0 - for tensor in bucket: - small_bucket.append(tensor) - numel = numel + tensor.numel() - if numel > numel_per_bucket: - self.allreduce_and_copy(small_bucket, dp_group, dp_world_size) - small_bucket = [] - numel = 0 - if len(small_bucket) > 0: - self.allreduce_and_copy(small_bucket, dp_group, dp_world_size) - - def _get_gradients_for_reduction(self): - non_expert_grads = [] - expert_grads = {} - if self.has_moe_layers: - for key in self.expert_data_parallel_group.keys(): - expert_grads[key] = [] - - for param_name, param in self.module.named_parameters(): - if not param.requires_grad: - continue - - if param.grad is None: - # In cases where there is an imbalance of empty grads across - # ranks we must create empty grads, this will ensure that every - # rank is reducing the same size. In some cases it may make - # sense in the future to support the ability to average not - # w.r.t. world size but with a different value. - param.grad = torch.zeros(param.size(), dtype=param.dtype, device=param.device) - - grad_data = param.grad.data - if param_name in self.sparse_tensor_module_names or grad_data.is_sparse: - # Call param.grad without data to avoid problem with setting of updated grads - grad_data = SparseTensor(param.grad) - - if is_moe_param(param): - expert_grads[param.group_name].append(grad_data) - else: - non_expert_grads.append(grad_data) - - return non_expert_grads, expert_grads - - def _reduce_non_expert_gradients(self, grads, elements_per_buffer): - split_sparse_tensor_buckets, split_dense_tensor_buckets = split_half_float_double_sparse(grads) - if self.pipeline_parallelism: - dp_group = self.mpu.get_data_parallel_group() - dp_world_size = dist.get_world_size(dp_group) - else: - dp_group = groups._get_sequence_data_parallel_group() - dp_world_size = dist.get_world_size(dp_group) / float(self.sequence_parallel_size) - for _, sparse_bucket_tuple in enumerate(split_sparse_tensor_buckets): - if sparse_bucket_tuple: - bucket_type, sparse_bucket = sparse_bucket_tuple - self.sparse_allreduce_no_retain(sparse_bucket, dp_group=dp_group, dp_world_size=dp_world_size) - - for _, dense_bucket_tuple in enumerate(split_dense_tensor_buckets): - if dense_bucket_tuple: - bucket_type, dense_bucket = dense_bucket_tuple - self.allreduce_no_retain(dense_bucket, - dp_group=dp_group, - numel_per_bucket=elements_per_buffer, - dp_world_size=dp_world_size) - - def _reduce_expert_gradients(self, expert_grads, elements_per_buffer): - # to maintain the gradients value unaffected by ep_size setting, - # utilize dp_world_size for allreduce average - dp_world_size = dist.get_world_size(groups._get_data_parallel_group()) - for ep_name, expert_grads_group in expert_grads.items(): - ep_dp_group = groups._get_expert_data_parallel_group(ep_name) - split_sparse_tensor_buckets, split_dense_tensor_buckets = split_half_float_double_sparse( - expert_grads_group) - - for _, sparse_bucket_tuple in enumerate(split_sparse_tensor_buckets): - if sparse_bucket_tuple: - bucket_type, sparse_bucket = sparse_bucket_tuple - self.sparse_allreduce_no_retain(sparse_bucket, dp_group=ep_dp_group, dp_world_size=dp_world_size) - - for _, dense_bucket_tuple in enumerate(split_dense_tensor_buckets): - if dense_bucket_tuple: - bucket_type, dense_bucket = dense_bucket_tuple - # Separate between diff groups - self.allreduce_no_retain(dense_bucket, - dp_group=ep_dp_group, - numel_per_bucket=elements_per_buffer, - dp_world_size=dp_world_size) - - def buffered_allreduce_fallback(self, grads=None, elements_per_buffer=500000000): - if grads is None: - if hasattr(self.optimizer, "get_grads_for_reduction"): - # This is currently for BF16 optimizer - non_expert_grads, expert_grads = self.optimizer.get_grads_for_reduction() - else: - non_expert_grads, expert_grads = self._get_gradients_for_reduction() - else: - assert not self.has_moe_layers, "attempting to reduce grads in unsupported way w.r.t. MoE" - non_expert_grads = grads - - self._reduce_non_expert_gradients(non_expert_grads, elements_per_buffer) - - if self.has_moe_layers: - self._reduce_expert_gradients(expert_grads, elements_per_buffer) - - def sparse_allreduce_no_retain(self, bucket, dp_group, dp_world_size=None): - allreduced_sparses = self.sparse_allreduce_bucket(bucket, dp_group, dp_world_size) - # Densify sparse tensor and copy back to original location - for tensor in allreduced_sparses: - if tensor.is_sparse: - tensor.orig_dense_tensor.data = tensor.to_coo_tensor() - else: - tensor.orig_dense_tensor.copy_(tensor.to_dense()) - - def sparse_allreduce_bucket(self, bucket, dp_group, dp_world_size=None): - sparse_list = [] - for sparse in bucket: - sparse_list.append(self.sparse_allreduce(sparse, dp_group, dp_world_size)) - return sparse_list - - def sparse_allreduce(self, sparse, dp_group, dp_world_size=None): - original_data_type = sparse.values.dtype - if self.communication_data_type != sparse.values.dtype: - if self.communication_data_type in (torch.float16, torch.bfloat16): - indices = sparse.indices.to(torch.int32) - else: - indices = sparse.indices - values = sparse.values.to(self.communication_data_type) - else: - indices = sparse.indices - values = sparse.values - - if dp_world_size is None: - dp_world_size = dist.get_world_size(group=dp_group) - if self.postscale_gradients(): - if self.gradient_average: - values.mul_(self.gradient_predivide_factor() / (dp_world_size)) - else: - values.mul_(1. / (dp_world_size)) - - indices_device_list = self.sparse_all_gather(indices, dp_group) - values_device_list = self.sparse_all_gather(values, dp_group) - - sparse.indices = torch.cat(indices_device_list).to(torch.long) - sparse.values = torch.cat(values_device_list).to(original_data_type) - return sparse - - def sparse_all_gather(self, value, dp_group): - my_size = torch.LongTensor([value.size()[0]]).to(self.device) - all_sizes = self.all_gather_scalar(my_size, dp_group) - max_size = torch.cat(all_sizes).max() - fill_size = max_size - my_size - - assert value.dim() in [1, 2] - if value.dim() == 1: - if fill_size > 0: - value = torch.cat([value, value.new_empty(fill_size)]) - tensor_list = [value.new_empty(max_size) for _ in range(dist.get_world_size(group=dp_group))] - else: - if fill_size > 0: - value = torch.cat([value, value.new_empty(fill_size, value.size()[1])]) - tensor_list = [ - value.new_empty(max_size, - value.size()[1]) for _ in range(dist.get_world_size(group=dp_group)) - ] - - dist.all_gather(tensor_list, value, group=dp_group) - tensors = [] - for dev_idx, t in enumerate(tensor_list): - size = all_sizes[dev_idx][0] - tensors.append(t.index_select(0, torch.arange(size, dtype=torch.long, device=self.device))) - - return tensors - - def all_gather_scalar(self, value, dp_group): - tensor_list = [value.new_zeros(value.size()) for _ in range(dist.get_world_size(group=dp_group))] - dist.all_gather(tensor_list, value, group=dp_group) - return tensor_list - - def module_state_dict(self, destination=None, prefix="", keep_vars=False, exclude_frozen_parameters=False): - sd = self.module.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars) - - # Remove frozen parameter weights from state_dict if specified - if exclude_frozen_parameters: - for n, p in self.module.named_parameters(): - if not p.requires_grad and n in sd: - del sd[n] - - if self.random_ltd_enabled(): - sd = remove_random_ltd_state_dict(sd) - return sd - - @staticmethod - def load_moe_state_dict(checkpoint_path, - tag, - state_dict, - old_moe_load, - model=None, - mpu=None, - num_experts=1, - checkpoint_engine=TorchCheckpointEngine()): - if old_moe_load: - expp_rank = groups._get_expert_data_parallel_rank(groups._get_max_expert_size_name()) - - num_local_experts = max(num_experts) // groups._get_expert_parallel_world_size( - groups._get_max_expert_size_name()) - for local_expert_id in range(num_local_experts): - global_expert_id = expp_rank * num_local_experts + local_expert_id - expert_state_dict = checkpoint_engine.load( - DeepSpeedEngine._get_expert_ckpt_name( - checkpoint_path, - -1, # -1 means ignore layer_id - global_expert_id, - tag, - mpu), - map_location=torch.device('cpu')) - - # Updating global -> local expert ids - moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.' - for key in list(expert_state_dict.keys()): - local_key = key.replace(f'{moe_str_prefix}{global_expert_id}', - f'{moe_str_prefix}{local_expert_id}') - expert_state_dict[local_key] = expert_state_dict.pop(key) - state_dict.update(expert_state_dict) - - else: - moe_layer_id = 0 - for n_module, module in model.named_modules(): - if isinstance(module, MoE): # and deepspeed.comm.get_rank() == 0: - group_name = module.expert_group_name - num_local_experts = module.num_local_experts - expp_rank = groups._get_expert_parallel_rank(group_name) - # loop all local_experts - for local_expert_id in range(num_local_experts): - global_expert_id = expp_rank * num_local_experts + local_expert_id - expert_state_dict = checkpoint_engine.load(DeepSpeedEngine._get_expert_ckpt_name( - checkpoint_path, moe_layer_id, global_expert_id, tag, mpu), - map_location=torch.device('cpu')) - # print(expert_state_dict.keys()) - # Updating global -> local expert ids - moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.' - for key in list(expert_state_dict.keys()): - local_key = key.replace(f'{moe_str_prefix}{global_expert_id}', - f'{moe_str_prefix}{local_expert_id}') - expert_state_dict[local_key] = expert_state_dict.pop(key) - state_dict.update(expert_state_dict) - moe_layer_id += 1 - - def load_module_state_dict(self, checkpoint, strict=True, custom_load_fn=None, fetch_z3_params=False): - if fetch_z3_params: - params_to_fetch = [ - p for p in self.module.parameters() - if hasattr(p, 'ds_id') and p.ds_status == ZeroParamStatus.NOT_AVAILABLE - ] - else: - params_to_fetch = [] - - with deepspeed.zero.GatheredParameters(params_to_fetch, modifier_rank=0): - module_state_dict = checkpoint['module'] - if custom_load_fn: - custom_load_fn(src=module_state_dict, dst=self.module) - else: - self.module.load_state_dict( - module_state_dict, # TODO - strict=strict) - - if checkpoint.get(FROZEN_PARAM_FRAGMENTS, None) is not None: - saved_frozen_params = checkpoint[FROZEN_PARAM_FRAGMENTS] - for param in self.module.parameters(): - if param.requires_grad: - continue - if param not in self.param_names: - raise ValueError(f"failed to find frozen {param} in named params") - name = self.param_names[param] - if hasattr(param, 'ds_id'): - param.ds_tensor.data.copy_(saved_frozen_params[name].data) - else: - param.data.copy_(saved_frozen_params[name].data) - - def _get_zero_ckpt_prefix(self, dp_rank, bf16_mode): - return f'{"bf16_" if bf16_mode else ""}zero_pp_rank_{dp_rank}' - - def _get_rank_zero_ckpt_name(self, checkpoints_path, tag, mp_rank, dp_rank, bf16_mode): - file_prefix = self._get_zero_ckpt_prefix(dp_rank, bf16_mode=bf16_mode) - zero_ckpt_name = os.path.join( - checkpoints_path, - str(tag), - f"{file_prefix}_mp_rank_{mp_rank:02d}_optim_states.pt", - ) - return zero_ckpt_name - - def _get_zero_ckpt_name(self, checkpoints_path, tag): - mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() - pp_rank = dist.get_rank(group=self.optimizer.dp_process_group) - bf16_mode = self.bfloat16_enabled() - return self._get_rank_zero_ckpt_name(checkpoints_path, tag, mp_rank, pp_rank, bf16_mode) - - def _get_ckpt_name(self, checkpoints_path, tag, mp_placeholder=None): - if mp_placeholder is not None: - mp_rank_str = mp_placeholder - else: - mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() - mp_rank_str = f"{mp_rank:02d}" - - if self.zero_optimization_partition_weights(): - if self.load_universal_checkpoint(): - filename = "zero_pp_rank_0" - else: - filename = "zero_pp_rank_{}".format(dist.get_rank(group=self.optimizer.dp_process_group)) - ckpt_name = os.path.join( - checkpoints_path, - str(tag), - f"{filename}_mp_rank_{mp_rank_str}_model_states.pt", - ) - else: - ckpt_name = os.path.join( - checkpoints_path, - str(tag), - "mp_rank_" + mp_rank_str + "_model_states.pt", - ) - return ckpt_name - - def _get_optimizer_ckpt_name(self, checkpoints_path, tag, expp_rank): - mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() - ckpt_name = os.path.join(checkpoints_path, str(tag), - f'expp_rank_{expp_rank}_mp_rank_{mp_rank:02d}_optim_states.pt') - return ckpt_name - - @staticmethod - def _get_expert_ckpt_name(checkpoints_path, layer_id, expert_id, tag, mpu=None): - mp_rank = 0 if mpu is None else mpu.get_model_parallel_rank() - if layer_id <= -1: - # Used to support old checkpoint loading - ckpt_name = os.path.join(checkpoints_path, '' if tag is None else str(tag), - f'expert_{expert_id}_mp_rank_{mp_rank:02d}_model_states.pt') - else: - # Used to support new checkpoint loading - ckpt_name = os.path.join(checkpoints_path, '' if tag is None else str(tag), - f'layer_{layer_id}_expert_{expert_id}_mp_rank_{mp_rank:02d}_model_states.pt') - return ckpt_name - - def _get_all_ckpt_names(self, checkpoints_path, tag): - # It is required that (checkpoints_path, tag) are consistent among all ranks. - ckpt_file_pattern = self._get_ckpt_name(checkpoints_path, tag, mp_placeholder="*") - import glob - - ckpt_files = glob.glob(ckpt_file_pattern) - ckpt_files.sort() - return ckpt_files - - def load_checkpoint(self, - load_dir, - tag=None, - load_module_strict=True, - load_optimizer_states=True, - load_lr_scheduler_states=True, - load_module_only=False, - custom_load_fn=None): - """ - Load training checkpoint - - Arguments: - load_dir: Required. Directory to load the checkpoint from - tag: Checkpoint tag used as a unique identifier for checkpoint, if not provided will attempt to load tag in 'latest' file - load_module_strict: Optional. Boolean to strictly enforce that the keys in state_dict of module and checkpoint match. - load_optimizer_states: Optional. Boolean to load the training optimizer states from Checkpoint. Ex. ADAM's momentum and variance - load_lr_scheduler_states: Optional. Boolean to add the learning rate scheduler states from Checkpoint. - load_module_only: Optional. Boolean to load only the model weights from the checkpoint. Ex. warmstarting. - custom_load_fn: Optional. Custom model load function. - - Returns: - A tuple of ``load_path`` and ``client_state``. - *``load_path``: Path of the loaded checkpoint. ``None`` if loading the checkpoint failed. - *``client_state``: State dictionary used for loading required training states in the client code. - - Important: under ZeRO3, one cannot load checkpoint with ``engine.load_checkpoint()`` right - after ``engine.save_checkpoint()``. It is because ``engine.module`` is partitioned, and - ``load_checkpoint()`` wants a pristine model. If insisting to do so, please reinitialize engine - before ``load_checkpoint()``. - - """ - - if tag is None: - latest_tag = "latest_universal" if self.load_universal_checkpoint() else "latest" - latest_path = os.path.join(load_dir, latest_tag) - if os.path.isfile(latest_path): - with open(latest_path, "r") as fd: - tag = fd.read().strip() - else: - if self.load_universal_checkpoint(): - raise ValueError(f'Invalid for universal checkpoint: {latest_path} does not exist') - else: - logger.warning( - f"Unable to find latest file at {latest_path}, if trying to load latest " - "checkpoint please ensure this file exists or pass an explicit checkpoint tag when loading a checkpoint." - ) - return None, None - - if self._optimizer_has_ckpt_event_prologue(): - # Prepare for checkpoint load by ensuring all parameters are partitioned - self.optimizer.checkpoint_event_prologue() - - load_path, client_states = self._load_checkpoint(load_dir, - tag, - load_module_strict=load_module_strict, - load_optimizer_states=load_optimizer_states, - load_lr_scheduler_states=load_lr_scheduler_states, - load_module_only=load_module_only, - custom_load_fn=custom_load_fn) - - load_zero_checkpoint = load_path is not None and (self.zero_optimization() or self.bfloat16_enabled()) - if load_zero_checkpoint: - if (load_optimizer_states and not load_module_only) or self.load_universal_checkpoint(): - success = self._load_zero_checkpoint(load_dir, tag, load_optimizer_states=load_optimizer_states) - else: - success = False - if not success: - self.optimizer._restore_from_bit16_weights() - - if self.zero_nvme_offload_optimizer(): - from shutil import copytree, disk_usage - offload_dir = self.optimizer.optimizer_swapper.swap_folder - offload_ckpt_dir = os.path.join(load_dir, tag, "offloaded_tensors") - _, _, free = disk_usage(offload_dir) - logger.info( - f"Copying NVMe offload checkpoint from {offload_ckpt_dir} to {offload_dir}, {free / 1e9:,.2f} GB free on target filesystem..." - ) - copytree(offload_ckpt_dir, offload_dir, dirs_exist_ok=True) - _, _, free = disk_usage(offload_dir) - logger.info(f"Copying complete! {free / 1e9:,.2f} GB free on target filesystem") - self.optimizer.reset_swap_buffers() - - if self._optimizer_has_ckpt_event_epilogue(): - self.optimizer.checkpoint_event_epilogue() - - if self.load_universal_checkpoint() and not self.zero_optimization_partition_weights(): - self.optimizer.update_lp_params() - - return load_path, client_states - - def _load_checkpoint(self, - load_dir, - tag, - load_module_strict=True, - load_optimizer_states=True, - load_lr_scheduler_states=True, - load_module_only=False, - custom_load_fn=None): - - from deepspeed.runtime.state_dict_factory import SDLoaderFactory - - ckpt_list = self._get_all_ckpt_names(load_dir, tag) - sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine=self.checkpoint_engine) - - is_pipe_parallel = isinstance(self.module, PipelineModule) - - mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() - load_path, checkpoint, _ = sd_loader.load(self.mp_world_size, mp_rank, is_pipe_parallel=is_pipe_parallel) - - if checkpoint is None: - return None, None - - fetch_z3_params = False - if self.zero_optimization_partition_weights() and not load_optimizer_states: - checkpoint['module'] = get_fp32_state_dict_from_zero_checkpoint(load_dir) - fetch_z3_params = True - - if is_pipe_parallel: - # Pipeline parallelism uses this to load its own checkpoint files. - self._curr_ckpt_path = os.path.join(load_dir, tag) - - if self.has_moe_layers: - # print(checkpoint.keys()) - old_moe_load = False - if not isinstance(checkpoint['num_experts'], list): - old_moe_load = True - DeepSpeedEngine.load_moe_state_dict(load_dir, - tag, - state_dict=checkpoint['module'], - old_moe_load=old_moe_load, - model=self.module, - mpu=self.mpu, - num_experts=self.num_experts, - checkpoint_engine=self.checkpoint_engine) - if not self.load_universal_checkpoint(): - self.load_module_state_dict(checkpoint=checkpoint, - strict=load_module_strict, - custom_load_fn=custom_load_fn, - fetch_z3_params=fetch_z3_params) - - self.loaded_checkpoint_dp_world_size = checkpoint['dp_world_size'] - - optim_checkpoint = None - if load_module_only: - deepspeed_states = ['module'] - if self.optimizer is not None and hasattr(self.optimizer, 'refresh_fp32_params'): - self.optimizer.refresh_fp32_params() - else: - has_zero_optimizer_state = self.zero_optimization() or self.bfloat16_enabled() - if load_optimizer_states and self.optimizer is not None and not has_zero_optimizer_state: - if self.has_moe_layers: - largest_group_name = groups._get_max_expert_size_name() - expp_rank = groups._get_expert_parallel_rank(largest_group_name) - optim_load_path = self._get_optimizer_ckpt_name(load_dir, tag, expp_rank) - optim_checkpoint = self.checkpoint_engine.load(optim_load_path, map_location=torch.device('cpu')) - else: - optim_checkpoint = checkpoint - - if self.fp16_enabled() or self.bfloat16_enabled(): - self.optimizer.load_state_dict(optim_checkpoint['optimizer'], - load_optimizer_states=load_optimizer_states) - else: - optim_checkpoint = checkpoint - - self.optimizer.load_state_dict(optim_checkpoint['optimizer']) - - if load_lr_scheduler_states and self.lr_scheduler is not None: - self.lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) - - if self.random_ltd_enabled() and self.random_ltd_scheduler is not None and 'random_ltd' in checkpoint: - self.random_ltd_scheduler.load_state_dict(checkpoint['random_ltd']) - - if self.training_dataloader is not None and self.curriculum_learning_enabled( - ) and 'data_sampler' in checkpoint: - self.training_dataloader.data_sampler.load_state_dict(checkpoint['data_sampler']) - - def get_sparse_tensor_module_names(original_set, loaded_set, original_parameters, loaded_parameters): - result = set() - - for name in original_set: - if name in loaded_parameters and name not in loaded_set: - continue # parameter existed in previous model and was not sparse - result.add(name) - - for name in loaded_set: - if name in original_parameters: - result.add(name) # parameter exists in both configs and it was sparse - - return result - - if 'sparse_tensor_module_names' in checkpoint: - sparse_tensor_module_names = checkpoint['sparse_tensor_module_names'] - elif 'csr_tensor_module_names' in checkpoint: - sparse_tensor_module_names = checkpoint['csr_tensor_module_names'] - else: - sparse_tensor_module_names = None - if sparse_tensor_module_names is not None: - if load_module_strict: - self.sparse_tensor_module_names = sparse_tensor_module_names - else: - self.sparse_tensor_module_names = get_sparse_tensor_module_names( - self.sparse_tensor_module_names, sparse_tensor_module_names, - dict(self.module.named_parameters()), checkpoint["module"]) - - self.global_steps = checkpoint['global_steps'] - self.global_samples = checkpoint.get('global_samples', self.global_steps * self.train_batch_size()) - self.skipped_steps = checkpoint['skipped_steps'] - self.loaded_checkpoint_mp_world_size = checkpoint['mp_world_size'] - deepspeed_states = [ - 'module', 'sparse_tensor_module_names', 'skipped_steps', 'global_steps', 'dp_world_size', - 'mp_world_size', 'data_sampler', 'random_ltd' - ] - client_state = {} - - if load_lr_scheduler_states: - deepspeed_states.append('lr_scheduler') - if load_optimizer_states: - deepspeed_states.append('optimizer') - - client_state = {key: value for key, value in checkpoint.items() if not key in deepspeed_states} - - if optim_checkpoint is not None: - client_state['optimizer'] = optim_checkpoint['optimizer'] - - return load_path, client_state - - def _load_zero_checkpoint(self, load_dir, tag, load_optimizer_states=True): - - load_serial = None - # When use loading checkpoint serial, checkpoint loading start from local rank 0, - # all other local rank would be paused, waiting for its rank-1 peer ready and its notification. - if self._config.zero_config.pipeline_loading_checkpoint: - assert self.zero_optimization_stage( - ) == ZeroStageEnum.weights, "Only stage3 support for pipeline checkpoint loading" - load_serial = torch.zeros(1).to(self.device) - if dist.get_local_rank() != 0: - dist.recv(tensor=load_serial, src=dist.get_rank() - 1) - if self.load_universal_checkpoint(): - zero_sd_list = None - checkpoint_folder = f'{os.path.join(load_dir, tag)}' - else: - if load_optimizer_states and self.seq_dp_world_size != self.loaded_checkpoint_dp_world_size: - raise ZeRORuntimeException("The checkpoint being loaded used a DP " \ - f"world size of {self.loaded_checkpoint_dp_world_size} but the " \ - f"current world size is {self.seq_dp_world_size}. Automatic adjustment " \ - "of ZeRO's optimizer state partitioning with a new world size is not " \ - "currently supported.") - checkpoint_folder = None - zero_sd_list = self._get_all_zero_checkpoints(load_dir, tag) - if zero_sd_list is None: - return False - - param_shapes = self._get_zero_param_shapes() - self.optimizer.load_state_dict(state_dict_list=zero_sd_list, - load_optimizer_states=load_optimizer_states, - load_from_fp32_weights=self.zero_load_from_fp32_weights(), - checkpoint_folder=checkpoint_folder, - load_serial=load_serial, - param_shapes=param_shapes) - - if self.load_universal_checkpoint(): - logger.info(f'loaded universal zero checkpoints from {checkpoint_folder} for rank {self.global_rank}') - else: - logger.info(f"loading {len(zero_sd_list)} zero partition checkpoints for rank {self.global_rank}") - return True - - def _get_mp_rank_zero_checkpoint_names(self, load_dir, tag, mp_rank, dp_world_size, bf16_mode): - zero_ckpt_names = [] - for dp_rank in range(dp_world_size): - ckpt_name = self._get_rank_zero_ckpt_name(checkpoints_path=load_dir, - tag=tag, - mp_rank=mp_rank, - dp_rank=dp_rank, - bf16_mode=bf16_mode) - zero_ckpt_names.append(ckpt_name) - - return zero_ckpt_names - - def _get_all_zero_checkpoint_names(self, load_dir, tag, bf16_mode): - mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() - zero_ckpt_names = self._get_mp_rank_zero_checkpoint_names(load_dir=load_dir, - tag=tag, - mp_rank=mp_rank, - dp_world_size=self.loaded_checkpoint_dp_world_size, - bf16_mode=bf16_mode) - for i, ckpt_name in enumerate(zero_ckpt_names): - if not os.path.exists(ckpt_name): - # transparently handle the old file pattern for optim_states - if "optim_states.pt" in ckpt_name: - ckpt_name_try = ckpt_name.replace("_optim_states.pt", "optim_states.pt") - if os.path.exists(ckpt_name_try): - zero_ckpt_names[i] = ckpt_name_try - continue - - return zero_ckpt_names - - def _get_all_zero_checkpoint_state_dicts(self, zero_ckpt_names): - zero_sd_list = [] - for i, ckpt_name in enumerate(zero_ckpt_names): - _state = None - if ckpt_name is None: - _state = {OPTIMIZER_STATE_DICT: None} - # Fully load state for current rank - elif self.zero_elastic_checkpoint() or dist.get_rank(group=self.optimizer.dp_process_group) == i: - _state = self.checkpoint_engine.load( - ckpt_name, - map_location='cpu', - ) - else: - _state = {OPTIMIZER_STATE_DICT: None} - zero_sd_list.append(_state) - - zero_optimizer_sd = [sd[OPTIMIZER_STATE_DICT] for sd in zero_sd_list] - logger.info(f"successfully read {len(zero_optimizer_sd)} ZeRO state_dicts for rank {self.global_rank}") - return zero_optimizer_sd - - def _get_all_zero_checkpoints(self, load_dir, tag): - for bf16_mode in [self.bfloat16_enabled(), not self.bfloat16_enabled()]: - zero_ckpt_names = self._get_all_zero_checkpoint_names(load_dir, tag, bf16_mode) - if zero_ckpt_names is not None: - # Warn if loading checkpoint of different bit16 type - if bf16_mode is not self.bfloat16_enabled(): - checkpoint_bit16 = BFLOAT16 if bf16_mode else FP16 - engine_bit16 = BFLOAT16 if self.bfloat16_enabled() else FP16 - logger.warning(f'Loading {checkpoint_bit16} zero checkpoints into {engine_bit16} training engine') - return self._get_all_zero_checkpoint_state_dicts(zero_ckpt_names) - - return None - - def _checkpoint_tag_validation(self, tag): - if self.checkpoint_tag_validation_enabled(): - s_hash = hashlib.sha1(tag.encode()) - bhash = torch.ByteTensor([s_hash.digest()]).flatten().to(self.device) - max_bhash = bhash.clone() - min_bhash = bhash.clone() - dist.all_reduce(max_bhash, op=dist.ReduceOp.MAX) - dist.all_reduce(min_bhash, op=dist.ReduceOp.MIN) - valid = all(min_bhash == bhash) and all(max_bhash == bhash) - msg = (f"[rank={dist.get_rank()}] The checkpoint tag name '{tag}' is not consistent across " - "all ranks. Including rank unique information in checkpoint tag could cause issues when " - "restoring with different world sizes.") - if self.checkpoint_tag_validation_fail(): - assert valid, msg - elif not valid: - logger.warning(msg) - - def save_checkpoint(self, save_dir, tag=None, client_state={}, save_latest=True, exclude_frozen_parameters=False): - """Save training checkpoint - - Arguments: - save_dir: Required. Directory for saving the checkpoint - tag: Optional. Checkpoint tag used as a unique identifier for the checkpoint, global step is - used if not provided. Tag name must be the same across all ranks. - client_state: Optional. State dictionary used for saving required training states in the client code. - save_latest: Optional. Save a file 'latest' pointing to the latest saved checkpoint. - exclude_frozen_parameters: Optional. Exclude frozen parameters from checkpointed state. - Important: all processes must call this method and not just the process with rank 0. It is - because each process needs to save its master weights and scheduler+optimizer states. This - method will hang waiting to synchronize with other processes if it's called just for the - process with rank 0. - - """ - if self._optimizer_has_ckpt_event_prologue(): - # Custom preparation for checkpoint save, if applicable - self.optimizer.checkpoint_event_prologue() - - rank = self.local_rank if self.use_node_local_storage() else self.global_rank - - # This is to make sure the checkpoint names are created without collision - # There seems to be issue creating them in parallel - - # Ensure save_dir directory exists - if rank == 0: - self.checkpoint_engine.makedirs(save_dir, exist_ok=True) - dist.barrier() - - if tag is None: - tag = f"global_step{self.global_steps}" - - # Ensure tag is a string - tag = str(tag) - commit_info = CheckpointCommitInfo(tag=tag, save_dir=save_dir, save_latest=save_latest) - - self.checkpoint_engine.create(commit_info) - - # Ensure checkpoint tag is consistent across ranks - self._checkpoint_tag_validation(tag) - - if self.has_moe_layers: - self.save_non_zero_checkpoint = False - self._create_checkpoint_file(save_dir, tag, False) - self._save_moe_checkpoint(save_dir, - tag, - client_state=client_state, - exclude_frozen_parameters=exclude_frozen_parameters) - - # We distribute the task of saving layer checkpoint files among - # data parallel instances, so all procs should call _save_checkpoint. - # All procs then call module_state_dict(), but only procs of data - # parallel rank 0 save the general model params. - if not self.has_moe_layers: - self._create_checkpoint_file(save_dir, tag, False) - self._save_checkpoint(save_dir, - tag, - client_state=client_state, - exclude_frozen_parameters=exclude_frozen_parameters) - - if self.save_zero_checkpoint: - self._create_zero_checkpoint_files(save_dir, tag) - self._save_zero_checkpoint(save_dir, tag) - - if self.zero_nvme_offload_optimizer(): - from shutil import copytree, disk_usage - offload_dir = self.optimizer.optimizer_swapper.swap_folder - offload_ckpt_dir = os.path.join(save_dir, tag, "offloaded_tensors") - _, _, free = disk_usage(save_dir) - logger.info( - f"Copying NVMe offload files from {offload_dir} to {offload_ckpt_dir}, {free / 1e9:,.2f} GB free on target filesystem..." - ) - copytree(offload_dir, - offload_ckpt_dir, - ignore=lambda _, dir_list: list(filter(lambda x: 'gradient' in x, dir_list)), - dirs_exist_ok=False) - _, _, free = disk_usage(save_dir) - logger.info(f"Copying complete! {free / 1e9:,.2f} GB free on target filesystem") - - if self._optimizer_has_ckpt_event_epilogue(): - self.optimizer.checkpoint_event_epilogue() - - # Save latest checkpoint tag - if not self.checkpoint_engine.is_decoupled(): - self.checkpoint_engine.commit(tag) - if save_latest and self.global_rank == 0: - with open(os.path.join(save_dir, 'latest'), 'w') as fd: - fd.write(tag) - - dist.barrier() - - return True - - def _commit_decoupled_checkpoint(self): - assert self.checkpoint_engine.is_decoupled(), \ - f'{self.checkpoint_engine} is not a Decoupled Checkpoint Engine' - - commit_info = self.checkpoint_engine.get_commit_info() - if commit_info is None: - return - - self.checkpoint_engine.commit(commit_info) - - if self.global_rank == 0 and commit_info.save_latest: - with open(os.path.join(commit_info.save_dir, 'latest'), 'w') as fd: - fd.write(commit_info.tag) - - dist.barrier() - - def _get_non_moe_state_dict(self, full_state_dict): - """ - Get the state dict of the non-moe layers - """ - for key in list(full_state_dict.keys()): - if 'expert' in key and 'moe.gate.wg.weight' not in key: - full_state_dict.pop(key) - - return full_state_dict - - def _save_moe_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False): - save_path = self._get_ckpt_name(save_dir, tag) - - # A hack to save the checkpointing directory. Pipeline parallelism overrides - # module_state_dict() and uses this path to save the model. module_state_dict() - # then instead just returns None. - - # Using layer_#_export_# to save the model's expert state_dict - moe_layer_id = 0 - for n_module, module in self.module.named_modules(): - if isinstance(module, MoE): # and deepspeed.comm.get_rank() == 0: - group_name = module.expert_group_name - num_local_experts = module.num_local_experts - expp_rank = groups._get_expert_parallel_rank(group_name) - exp_dp_rank = groups._get_expert_data_parallel_rank(group_name) - # print(expp_rank, exp_dp_rank) - # if exp_dp_rank != 0: - if not self.checkpoint_engine.is_data_parallel_writer(exp_dp_rank): - moe_layer_id += 1 - continue - - # get all moe parameters - moe_state_dict = {} - for n, p in module.state_dict().items(): - if 'expert' in n and 'moe.gate.wg.weight' not in n: - moe_state_dict[n_module + '.' + n] = p - moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.' - # print(moe_state_dict.keys()) # until now, everything is fine. So the bug happens at next few lines - # Reorder the moe name rank, so that each checkpoint only has one expert - experts_state_dict = defaultdict(dict) - for key in list(moe_state_dict.keys()): - m = re.match(f".*{moe_str_prefix}([0-9]+).*", key) - - local_expert_id = None - if not m: - logger.warning(f'No expert found in key {key}.') - else: - local_expert_id = m.group(1) - - global_expert_id = expp_rank * \ - num_local_experts + int(local_expert_id) - expert_key = key.replace(f'{moe_str_prefix}{local_expert_id}', - f'{moe_str_prefix}{global_expert_id}') - # truncating extra tensor (shared) storage - truncated = moe_state_dict.pop(key).clone().detach() - experts_state_dict[str(global_expert_id)][expert_key] = truncated - - # let save the moe parameters - for global_expert_id, expert_state_dict in experts_state_dict.items(): - # save the moe parameters - moe_save_path = self._get_expert_ckpt_name(save_dir, moe_layer_id, global_expert_id, tag, self.mpu) - if self.random_ltd_enabled(): - expert_state_dict = remove_random_ltd_state_dict(expert_state_dict) - saveable_state_dict = clone_tensors_for_torch_save(expert_state_dict) - self.checkpoint_engine.save(saveable_state_dict, moe_save_path) - moe_layer_id += 1 - - self._curr_ckpt_path = os.path.join(save_dir, tag) - - largest_group_name = groups._get_max_expert_size_name() - expp_rank = groups._get_expert_parallel_rank(largest_group_name) - exp_dp_rank = groups._get_expert_data_parallel_rank(largest_group_name) - - # In the case of E + D parallelism, only the - # first expert parallel group should save the expert weights - # since each expert parallel group is a copy of the model's experts - if not self.checkpoint_engine.is_data_parallel_writer(exp_dp_rank): - return - - # Save optimizer states. They are different across each exp parallel rank. - optimizer_state = { - 'optimizer': self.optimizer.state_dict() if self.optimizer and not self.zero_optimization() else None - } - # TODO: why use BufferedWriter not the path - file_path = self._get_optimizer_ckpt_name(save_dir, tag, expp_rank) - saveable_state_dict = clone_tensors_for_torch_save(optimizer_state) - self.checkpoint_engine.save(saveable_state_dict, file_path) - - # Load flow uses below saved file for model parameters, RNG and more - if groups._get_data_parallel_rank() == 0: - # Get non-moe parameters - # Classes DeepSpeedEngine and PipelineEngine have different behavior for method module_state_dict. - # DeepSpeedEngine returns the state dict, where PipelineEngine saves the state dict and returns None. - # We need to get the state dict, therefore, call to DeepSpeedEngine (base class for PipelineEngine) - model_state_dict = self._get_non_moe_state_dict( - DeepSpeedEngine.module_state_dict(self, exclude_frozen_parameters=exclude_frozen_parameters)) - - # TODO: update num experts info,.. in checkpoint - state = { - 'module': - model_state_dict, - 'lr_scheduler': - self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None, - 'data_sampler': - self.training_dataloader.data_sampler.state_dict() if - (self.training_dataloader is not None and self.curriculum_learning_enabled()) else None, - 'random_ltd': - self.random_ltd_scheduler.state_dict() if self.random_ltd_enabled() else None, - 'sparse_tensor_module_names': - self.sparse_tensor_module_names, - 'skipped_steps': - self.skipped_steps, - 'global_steps': - self.global_steps, - 'global_samples': - self.global_samples, - 'dp_world_size': - self.dp_world_size, - 'mp_world_size': - self.mp_world_size, - 'num_experts': - self.num_experts - } - state.update(client_state) - logger.info(f'Saving model checkpoint: {save_path}') - saveable_state_dict = clone_tensors_for_torch_save(state) - self.checkpoint_engine.save(saveable_state_dict, save_path) - - def _create_checkpoint_file(self, save_dir, tag, zero_checkpoint): - name_function = (self._get_zero_ckpt_name if zero_checkpoint else self._get_ckpt_name) - try: - checkpoint_name = name_function(save_dir, tag) - path = os.path.dirname(checkpoint_name) - self.checkpoint_engine.makedirs(path, exist_ok=True) - except: - logger.error(f"Failed saving model checkpoint to {save_dir} with tag {tag}") - return False - - return True - - def _create_zero_checkpoint_files(self, save_dir, tag): - success = True - # zero checkpoint files are created sequentially - for rank in range(dist.get_world_size(self.optimizer.dp_process_group)): - if rank == self.global_rank: - success = self._create_checkpoint_file(save_dir, tag, True) - - return success - - def _save_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False): - - save_path = self._get_ckpt_name(save_dir, tag) - - zero_optimizer_state = self.zero_optimization() or self.bfloat16_enabled() - - save_frozen_param = self.zero_optimization_partition_gradients() and not exclude_frozen_parameters - - # A hack to save the checkpointing directory. Pipeline parallelism overrides - # module_state_dict() and uses this path to save the model. module_state_dict() - # then instead just returns None. The module_state_dict() implementation in - # PipelineEngine expects the save path to be set in self._curr_ckpt_path. - self._curr_ckpt_path = os.path.join(save_dir, tag) - module = self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters) - self._curr_ckpt_path = None - - state = dict(module=module, - buffer_names=self._get_buffer_names(), - optimizer=self.optimizer.state_dict() if self.optimizer and not zero_optimizer_state else None, - param_shapes=self._get_zero_param_shapes() if self.optimizer and zero_optimizer_state else None, - frozen_param_shapes=self._get_zero_frozen_param_attributes(self._get_param_shape_func) - if save_frozen_param else None, - shared_params=self._get_shared_params() if self.optimizer and zero_optimizer_state else None, - frozen_param_fragments=self._get_zero_frozen_param_attributes(self._get_param_fragment_func) - if save_frozen_param else None, - lr_scheduler=self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None, - data_sampler=self.training_dataloader.data_sampler.state_dict() if - (self.training_dataloader is not None and self.curriculum_learning_enabled()) else None, - random_ltd=self.random_ltd_scheduler.state_dict() if self.random_ltd_enabled() else None, - sparse_tensor_module_names=self.sparse_tensor_module_names, - skipped_steps=self.skipped_steps, - global_steps=self.global_steps, - global_samples=self.global_samples, - dp_world_size=self.seq_dp_world_size, - mp_world_size=self.mp_world_size, - ds_config=self.config, - ds_version=version) - state.update(client_state) - log_dist(message=f'Saving model checkpoint: {save_path}', ranks=[0]) - - if self.save_non_zero_checkpoint: - self.checkpoint_engine.save(state_dict=state, path=save_path) - - def _get_buffer_names(self): - buffer_names = [] - - # we save buffer names so that we could extract later the real buffers from the saved - # state_dict["module"] in the non-zero checkpoint - the buffers are already there but they - # are intermixed with param placeholders - - # have to traverse the tree to be able to skip non-persistent buffers - def get_layer_named_buffers(module, prefix=""): - for name, buf in module.named_buffers(recurse=False): - if buf is not None and name not in module._non_persistent_buffers_set: - buffer_names.append(prefix + name) - - for name, child in module.named_children(): - if child is not None: - get_layer_named_buffers(child, prefix + name + ".") - - get_layer_named_buffers(self.module, prefix="") - - return buffer_names - - def _get_param_shape_func(self, param): - return param.ds_shape if hasattr(param, 'ds_id') else param.shape - - def _get_param_fragment_func(self, param): - return param.ds_tensor.detach().cpu() if hasattr(param, 'ds_id') else param.detach().cpu() - - def _get_zero_frozen_param_attributes(self, attr_func): - frozen_param_fragments = OrderedDict() - - for param in self.module.parameters(): - if param.requires_grad: - continue - if param not in self.param_names: - raise ValueError(f"failed to find frozen {param} in named params") - name = self.param_names[param] - frozen_param_fragments[name] = attr_func(param) - - return frozen_param_fragments - - def _get_zero_param_shapes(self): - """Returns a dict of name to shape mapping, only for the flattened fp32 weights saved by the - optimizer. the names are exactly as in state_dict. The order is absolutely important, since - the saved data is just flattened data with no identifiers and requires reconstruction in the - same order it was saved. - We can't rely on self.module.named_parameters() to get the saved tensors, as some params - will be missing and others unsaved and then it'd be impossible to reconstruct state_dict - from the flattened weights. - optimizer.bit16_groups seems to be the easiest to use as it's in all zeroX versions. - """ - param_group_shapes = [] - cnt = 0 - numel = 0 - - # zero2 started using a round_robin_bit16_groups which is a shuffled version of bit16_groups - - # if we don't use it, we get parameters ordered incorrectly - if hasattr(self.optimizer, "round_robin_bit16_groups"): - bit16_groups = self.optimizer.round_robin_bit16_groups - elif self.bfloat16_enabled() and hasattr(self.optimizer, "bf16_groups"): - bit16_groups = self.optimizer.bf16_groups - else: - bit16_groups = self.optimizer.bit16_groups if self.zero_optimization_stage( - ) == 2 else self.optimizer.fp16_groups - - for bit16_group in bit16_groups: - param_shapes = OrderedDict() - for param in bit16_group: - cnt += 1 - numel += param.ds_numel if hasattr(param, "ds_numel") else param.numel() - shape = param.ds_shape if hasattr(param, "ds_shape") else param.shape - if param not in self.param_names: - raise ValueError(f"failed to find optimizer param in named params") - name = self.param_names[param] - param_shapes[name] = shape - - # uncomment to debug zero_to_fp32.py problems - # if self.global_rank == 0: print(f"saving param {name} {shape} (numel={shape.numel()})") - param_group_shapes.append(param_shapes) - # if self.global_rank == 0: print(f"Total saved {numel} numels in {cnt} params") - - return param_group_shapes - - def _get_shared_params(self): - """ - Returns a dict of shared params, which can later be used to reconstruct the original state dict, - e.g. in `zero_to_fp32`. Each dict entry is a pair of param names, where the key is the name - of the variable that isn't stored and the value is the actual param holding data. - """ - shared_index = {} - shared_params_by_full_name = {} - - is_zero3_model = (self.zero_optimization_partition_weights() - and any(hasattr(param, "ds_id") for param in self.module.parameters())) - - def get_layer_state_dict(module, prefix=""): - # handle params - for name, param in module.named_parameters(recurse=False): - if param is None or (is_zero3_model and not hasattr(param, "ds_id")): - continue - key = prefix + name - - # When weights are manged by stage 3, we can't rely on param.data_ptr() as it will be reused - # as weights get gathered and reduced, but param.ds_id is unique across all zero weights - # (and shared params will have the same param.ds_id) - param_id = param.ds_id if is_zero3_model else param.data_ptr() - - if param_id in shared_index: - # shared weights - #print(f"`{key}` is shared with `{shared_index[param_id]}`") - shared_params_by_full_name[key] = shared_index[param_id] - else: - shared_index[param_id] = key - - for name, child in module.named_children(): - if child is not None: - get_layer_state_dict(child, prefix + name + ".") - - if dist.get_rank() == 0: - get_layer_state_dict(self.module, prefix="") - - return shared_params_by_full_name - - def _copy_recovery_script(self, save_path): - base_dir = os.path.dirname(os.path.dirname(__file__)) - script = "zero_to_fp32.py" - src = os.path.join(base_dir, "utils", script) - dst = os.path.join(save_path, script) - #logger.info(f"creating recovery script {dst}") - copyfile(src, dst) - self._change_recovery_script_permissions(dst) - - def _change_recovery_script_permissions(self, dst): - # make executable (safeguard for file shares - Azure as example) - try: - os.chmod(dst, os.stat(dst).st_mode | stat.S_IEXEC) - except (FileNotFoundError, PermissionError) as e: - #this message is used in unit test TestZeRONonDistributed - logger.info( - f'Warning: Could not change permissions for {dst} due to error: {e}. Continuing without changing permissions.' - ) - - def _save_zero_checkpoint(self, save_path, tag): - zero_checkpoint_name = self._get_zero_ckpt_name(save_path, tag) - zero_sd = dict(optimizer_state_dict=self.optimizer.state_dict(), ds_config=self.config, ds_version=version) - self.checkpoint_engine.save(zero_sd, zero_checkpoint_name) - - if self.global_rank == 0: - self._copy_recovery_script(save_path) - ckpt_type = 'zero' if self.zero_optimization() else 'bf16_zero' - #logger.info(f'{ckpt_type} checkpoint saved {zero_checkpoint_name}') - - def _replace_module_consolidated_state_dict(self): - """ - Get a full non-partitioned state_dict with fp16 weights on cpu. - Important: this function must be called on all ranks and not just rank 0. - This is similar to nn.Module.state_dict (modelled after _save_to_state_dict) - This method is used for tensor parallel training. - - Returns: - OrderedDict: The consolidated state dictionary if the current process rank is 0, otherwise None. - """ - #TODO: If we use both Zero3 and tensor parallel simultaneously - # we need to consolidate the gather mechanisms of both. - state_dict = OrderedDict() if dist.get_rank() == 0 else None - - def get_layer_state_dict(module, prefix=""): - with GatherReplacedLayerParams(list(module.parameters(recurse=False)), module, enabled=True): - for name, param in module.named_parameters(recurse=False): - if param is None: - continue - key = prefix + name - if (dist.get_rank() == 0): - state_dict[key] = param.detach().cpu() - # print(key,module, param.detach().cpu().shape) - - for name, child in module.named_children(): - if child is not None: - get_layer_state_dict(child, prefix + name + ".") - - get_layer_state_dict(self.module, prefix="") - - # ensure that all GPU communication tasks are completed before the process exits - get_accelerator().synchronize() - return state_dict - - def _consolidated_16bit_state_dict(self, exclude_frozen_parameters=False): - """ - Consolidate the 16-bit state dictionary. - """ - if self.zero_optimization_stage() == ZeroStageEnum.weights: - return self._zero3_consolidated_16bit_state_dict(exclude_frozen_parameters) - elif self.autotp_size() > 1: - return self._replace_module_consolidated_state_dict() - - raise ValueError("consolidated_16bit_state_dict is only applicable to cases where weights are partitioned, " - "including Zero Stage 3 and tensor parallelism.") - - def _zero3_consolidated_16bit_state_dict(self, exclude_frozen_parameters=False): - """ - Get a full non-partitioned state_dict with fp16 weights on cpu. - Important: this function must be called on all ranks and not just rank 0. - This is similar to nn.Module.state_dict (modelled after _save_to_state_dict), but: - 1. consolidates the weights from different partitions on gpu0 - 2. works on one layer at a time to require as little gpu0 memory as possible, by - moving the already consolidated weights to cpu - 3. takes care to keep the shared params shared when gradually copying the params to cpu - Returns: - a consolidated fp16 ``state_dict`` on cpu on rank 0, ``None`` on other ranks - """ - if not self.zero_optimization_partition_weights(): - raise ValueError("this function requires ZeRO-3 mode") - - state_dict = OrderedDict() if dist.get_rank() == 0 else None - shared_params = {} - - def get_layer_state_dict(module, prefix=""): - # gather one layer at a time to be memory-efficient - # must use modifier_rank=0 to release GPU memory after each layer gathered - #see_memory_usage("before GatheredParameters", force=True) - with deepspeed.zero.GatheredParameters(list(module.parameters(recurse=False)), modifier_rank=0): - if dist.get_rank() == 0: - # handle params - for name, param in module.named_parameters(recurse=False): - if param is None or (exclude_frozen_parameters and not param.requires_grad): - continue - key = prefix + name - # can't rely on param.data_ptr() as it will be reused as weights gets - # gathered and reduced, but param.ds_id is unique across all zero weights - # (and shared params will have the same param.ds_id) - if param.ds_id in shared_params: - # shared weights - #print(f"`{key}` is shared with `{shared_params[param.ds_id]}`") - state_dict[key] = state_dict[shared_params[param.ds_id]] - else: - state_dict[key] = param.detach().cpu() - shared_params[param.ds_id] = key - #print(f"param {param.ds_id} {param.shape} {key} ") - - # now buffers - not sure if need to take care of potentially shared weights here - for name, buf in module.named_buffers(recurse=False): - if (buf is not None and name not in module._non_persistent_buffers_set): - state_dict[prefix + name] = buf.detach().cpu() - #see_memory_usage("after GatheredParameters", force=True) - - for name, child in module.named_children(): - if child is not None: - get_layer_state_dict(child, prefix + name + ".") - - # Prepare for checkpoint save by ensuring all parameters are partitioned - if self._optimizer_has_ckpt_event_prologue(): - self.optimizer.checkpoint_event_prologue() - - see_memory_usage("before get_layer_state_dict", force=False) - get_layer_state_dict(self.module, prefix="") - see_memory_usage("after get_layer_state_dict", force=False) - - if self._optimizer_has_ckpt_event_epilogue(): - self.optimizer.checkpoint_event_epilogue() - - return state_dict - - def save_fp16_model(self, save_dir, save_filename="pytorch_model.bin"): - """has been renamed to save_16bit_model, keeping this around for backwards - compatibility""" - return self.save_16bit_model(save_dir, save_filename) - - def save_16bit_model(self, save_dir, save_filename="pytorch_model.bin", exclude_frozen_parameters=False): - """ - Save 16bit model weights - - This method saves the 16bit model weights at the desired destination. - - Arguments: - save_dir: Required. Directory for saving the model - save_filename: Optional. Filename to save to. Defaults to ``pytorch_model.bin`` - exclude_frozen_parameters: Optional. Exclude frozen parameters from checkpointed state. - - Returns: - ``True`` when a model has been saved, ``False`` otherwise. It will not be saved if - stage3_gather_16bit_weights_on_model_save is ``False``. - - Important: all processes must call this method and not just the process with rank 0. It is - because the processes need to work in sync to gather the weights. This method will hang - waiting to synchronize with other processes if it's called just for the process with rank 0. - - """ - - path = os.path.join(save_dir, save_filename) - - if self.zero_optimization_partition_weights(): - if self.zero_gather_16bit_weights_on_model_save(): - # consolidation is expensive in time and memory and therefore isn't a default - state_dict = self._zero3_consolidated_16bit_state_dict( - exclude_frozen_parameters=exclude_frozen_parameters) - else: - # the model will be bogus if not consolidated so don't confuse the user by saving it - logger.info( - f"Did not save the model {path} because stage3_gather_16bit_weights_on_model_save is False") - return False - else: - state_dict = self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters) - - tag = f"global_step{self.global_steps}" - tag = str(tag) - commit_info = CheckpointCommitInfo(tag=tag, save_dir=save_dir, save_latest=False) - self.checkpoint_engine.create(commit_info) - - if dist.get_rank() == 0: - self.checkpoint_engine.makedirs(save_dir, exist_ok=True) - logger.info(f"Saving model weights to {path}, tag: {tag}") - self.checkpoint_engine.save(state_dict, path) - - self.checkpoint_engine.commit(tag) - - return True - - def empty_partition_cache(self): - """ - Release GPU memory consumed by offloaded model parameters. - """ - if hasattr(self.optimizer, 'empty_partition_cache'): - self.optimizer.empty_partition_cache() - gc.collect() - get_accelerator().empty_cache() - - def compile(self, backend=get_accelerator().get_compile_backend(), compile_kwargs={}, schedule=None) -> None: - """Compile the module using the specified backend and kwargs. - If a compiler_fn is set, it will be used instead of torch.compile(). - """ - # Avoid graph breaks - deepspeed.utils.nvtx.enable_nvtx = False - - if not is_compile_supported(): - raise RuntimeError("compile is not supported in your version of PyTorch.") - - if self.is_compiled: - return - - if 'backend' in compile_kwargs: - logger.warning("The `backend` in `compile_kwargs` will be overridden. Use the `backend` argument instead.") - - logger.info(f"Compiling deepcompile={self.is_deepcompile_enabled()} backend={backend}") - - enable_deepcompile = self.is_deepcompile_enabled() - if enable_deepcompile and self.zero_optimization_stage() != ZeroStageEnum.optimizer_states \ - and self.zero_optimization_stage() != ZeroStageEnum.weights \ - and self.zero_optimization_stage() != ZeroStageEnum.gradients: - logger.info( - f"Currently DeepCompile supports ZeRO stage 1, 2, or 3 only, but ZeRO stage is set to {self.zero_optimization_stage()}. Falling back to the torch compiler." - ) - enable_deepcompile = False - - if enable_deepcompile: - - if schedule is not None: - - def passes_name_to_fn(passes): - for p in passes: - assert callable(p) or p in opt_passes, f"Unknown pass {p}" - return [p if callable(p) else opt_passes[p] for p in passes] - - schedule = [(step, passes_name_to_fn(passes)) for step, passes in schedule] - - assert backend in ['inductor', 'eager'], f"Backend {backend} is not supported for DeepCompile." - - compile_config = self._config.compile_config - if (("zero_optimization" in self.config and "offload_optimizer" in self.config["zero_optimization"] - and "offload_param" in self.config["zero_optimization"]) - and self._config.zero_config.offload_param.device == "cpu" - and self._config.zero_config.offload_optimizer.device == "cpu"): - compile_config.offload_parameters = True - if self.zero_optimization_stage() == ZeroStageEnum.optimizer_states: - backend = init_z1(self, backend, compile_config, compile_kwargs, schedule) - elif self.zero_optimization_stage() == ZeroStageEnum.gradients: - backend = init_z1(self, backend, compile_config, compile_kwargs, schedule, use_z2=True) - elif self.zero_optimization_stage() == ZeroStageEnum.weights: - backend = init_z3(self, backend, compile_config, compile_kwargs, schedule) - - # create new dict to avoid modifying original dict - self.module.compile(**{**compile_kwargs, 'backend': backend}) - - self._is_compiled = True - - def get_compile_time(self): - from deepspeed.compile.backend import opt_pass_times - return opt_pass_times - - def register_compile_pass(self, pass_name: str, pass_fn: Callable) -> None: - register_compile_pass(pass_name, pass_fn) - - def is_deepcompile_enabled(self): - return self._config.compile_config.deepcompile - - @property - def is_compiled(self) -> bool: - return self._is_compiled - - def offload_states(self, - include: Container[OffloadStateTypeEnum] = None, - device: OffloadDeviceEnum = OffloadDeviceEnum.cpu, - pin_memory: bool = True, - non_blocking: bool = False) -> None: - """Offload the engine's states to the specified device. - - Arguments: - include: Optional. The set of states to offload. If not provided, all states are offloaded. - device: Optional. The device to move the ZeRO optimizer buffers to. Currently only `OffloadDeviceEnum.cpu` is supported. - pin_memory: Optional. Whether to pin the memory of the offloaded states. - non_blocking: Optional. Whether to offload the states asynchronously. - """ - assert self.zero_optimization_stage( - ) == ZeroStageEnum.weights, "Moving buffers across devices is supported only for ZeRO stage 3." - - opt_offload_config = self.zero_offload_optimizer() - assert opt_offload_config is None or opt_offload_config.device == OffloadDeviceEnum.none, "Moving states across devices is not supported for offloaded optimizer states." - param_offload_config = self.zero_offload_param() - assert param_offload_config is None or param_offload_config.device == OffloadDeviceEnum.none, "Moving states across devices is not supported for offloaded parameters." - - assert not isinstance( - self.optimizer, - DeepSpeedZeRoOffload), "Moving states across devices is not supported without an optimizer." - - if device == OffloadDeviceEnum.none: - logger.warning("No device specified for offloading states.") - return - - if device == OffloadDeviceEnum.nvme: - raise ValueError("NVMe offload is not supported for offloading states.") - - self.optimizer.offload_states(include=include, device=device, pin_memory=pin_memory, non_blocking=non_blocking) - - def reload_states(self, non_blocking: bool = False) -> None: - """Reload the engine states to the original device. - - Arguments: - non_blocking: Optional. Whether to offload the states asynchronously. - """ - assert self.zero_optimization_stage( - ) == ZeroStageEnum.weights, "Moving buffers back is supported only for ZeRO stage 3." - - assert not isinstance( - self.optimizer, - DeepSpeedZeRoOffload), "Moving states across devices is not supported without an optimizer." - - self.optimizer.reload_states(non_blocking=non_blocking) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__init__.py deleted file mode 100644 index 6c5067f71c8faf166bc78e88f9b62e8627dda7c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -'''Copyright The Microsoft DeepSpeed Team''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 87a13af244e12fd50badbe51120fa582c4cb4438..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/fused_optimizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/fused_optimizer.cpython-310.pyc deleted file mode 100644 index e31230449e156aa69cd4b49d307a392c34f51d16..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/fused_optimizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/loss_scaler.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/loss_scaler.cpython-310.pyc deleted file mode 100644 index 8a3385ed6644c5c3dca42ac1c2218f62f1b36494..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/loss_scaler.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/unfused_optimizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/unfused_optimizer.cpython-310.pyc deleted file mode 100644 index d05f83a5b6907d9c0ef4d8cf2ad928fdd04993ec..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/__pycache__/unfused_optimizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/fused_optimizer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/fused_optimizer.py deleted file mode 100644 index 706fb6ec78e38f9750f62f0addc0b104b1af91be..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/fused_optimizer.py +++ /dev/null @@ -1,526 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Copyright NVIDIA/apex -This file is adapted from FP16_Optimizer in NVIDIA/apex -""" - -import torch -from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors -from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer -from deepspeed.runtime.utils import get_global_norm, get_flattened_grad_norm, CheckOverflow, get_weight_norm, get_norm_with_moe_layers, is_model_parallel_parameter -from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE -from deepspeed.utils import logger, log_dist -from deepspeed.utils.torch import required_torch_version -from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, CLIP_GRAD -from deepspeed.accelerator import get_accelerator -from deepspeed.moe.utils import is_moe_param_group -from deepspeed.runtime.constants import PIPE_REPLICATED -from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank - -OVERFLOW_CHECK_TIMER = 'overflow_check' -COMPUTE_NORM_TIMER = 'compute_norm' -UNSCALE_AND_CLIP_TIMER = 'unscale_and_clip' -BASIC_STEP_TIMER = 'basic_step' -UPDATE_FP16_TIMER = 'update_fp16' - -OVERFLOW_TIMERS = [COMPUTE_NORM_TIMER, OVERFLOW_CHECK_TIMER] -STEP_TIMERS = OVERFLOW_TIMERS + [UNSCALE_AND_CLIP_TIMER, BASIC_STEP_TIMER, UPDATE_FP16_TIMER] - - -class FP16_Optimizer(DeepSpeedOptimizer): - """ - FP16 Optimizer for training fp16 models. Handles loss scaling. - - For usage example please see, TODO: DeepSpeed V2 Tutorial - """ - - def __init__(self, - init_optimizer, - deepspeed=None, - static_loss_scale=1.0, - dynamic_loss_scale=False, - initial_dynamic_scale=2**32, - dynamic_loss_args=None, - verbose=True, - mpu=None, - clip_grad=0.0, - fused_adam_legacy=False, - has_moe_layers=False, - timers=None): - - self.fused_adam_legacy = fused_adam_legacy - self.timers = timers - self.deepspeed = deepspeed - self.has_moe_layers = has_moe_layers - self.using_pipeline = getattr(self.deepspeed, 'pipeline_parallelism', False) - if not get_accelerator().is_available(): - raise SystemError("Cannot use fp16 without accelerator.") - self.optimizer = init_optimizer - - # param flattened by groups - self.fp16_groups = [] - self.fp16_groups_flat = [] - self.fp32_groups_flat = [] - - self.flatten_grad_norm_mask_list = [] - self.has_executed_step = False - self._global_grad_norm = 0. - - # loop to deal with groups - for i, param_group in enumerate(self.optimizer.param_groups): - # push this group to list before modify - self.fp16_groups.append(param_group['params']) - # init fp16 weight buffer, flattened - self.fp16_groups_flat.append(_flatten_dense_tensors([p.clone().detach() for p in self.fp16_groups[i]])) - # set model fp16 weight to slices of flattened buffer - updated_params = _unflatten_dense_tensors(self.fp16_groups_flat[i], self.fp16_groups[i]) - for p, q in zip(self.fp16_groups[i], updated_params): - p.data = q.data - # init master weight, flattened - self.fp32_groups_flat.append(self.fp16_groups_flat[i].clone().float().detach()) - # modify optimizer of have flat master weight - self.fp32_groups_flat[i].requires_grad = True # keep this in case internal optimizer uses it - param_group['params'] = [self.fp32_groups_flat[i]] - - # we may have a way of fusing dynamic scale. Do not support for now - if dynamic_loss_scale: - self.dynamic_loss_scale = True - self.cur_iter = 0 - self.last_overflow_iter = -1 - self.scale_factor = 2 - - if dynamic_loss_args is None: - self.cur_scale = initial_dynamic_scale - self.scale_window = 1000 - self.min_loss_scale = 1 - else: - self.cur_scale = dynamic_loss_args[INITIAL_LOSS_SCALE] - self.scale_window = dynamic_loss_args[SCALE_WINDOW] - self.min_loss_scale = dynamic_loss_args[MIN_LOSS_SCALE] - else: - self.dynamic_loss_scale = False - self.cur_iter = 0 - self.cur_scale = static_loss_scale - self.verbose = verbose - - self.custom_loss_scaler = False - self.external_loss_scale = None - - self.clip_grad = clip_grad - self.norm_type = 2 - - if required_torch_version(max_version=0.4): - self.clip_grad_norm = torch.nn.utils.clip_grad_norm - else: - self.clip_grad_norm = torch.nn.utils.clip_grad_norm_ - - #model parallel object - self.mpu = mpu - - self.overflow = False - self.overflow_checker = CheckOverflow(self.fp16_groups, mpu=self.mpu, deepspeed=deepspeed) - self.initialize_optimizer_states() - - def initialize_optimizer_states(self): - for i, group in enumerate(self.fp16_groups): - self.fp32_groups_flat[i].grad = torch.zeros(self.fp32_groups_flat[i].size(), - device=self.fp32_groups_flat[i].device) - - self.optimizer.step() - - for i, group in enumerate(self.fp16_groups): - self.fp32_groups_flat[i].grad = None - - return - - def zero_grad(self, set_to_none=True): - """ - Zero FP16 parameter grads. - """ - # For speed, set model fp16 grad to None by default - for group in self.fp16_groups: - for p in group: - if set_to_none: - p.grad = None - else: - if p.grad is not None: - p.grad.detach_() - p.grad.zero_() - - def step_fused_adam(self, closure=None): - """ - Not supporting closure. - """ - - # First compute norm for all group so we know if there is overflow - grads_groups_flat = [] - norm_groups = [] - for i, group in enumerate(self.fp16_groups): - grads_groups_flat.append( - _flatten_dense_tensors([ - torch.zeros(p.size(), dtype=p.dtype, device=p.device) if p.grad is None else p.grad for p in group - ])) - norm_groups.append(get_weight_norm(grads_groups_flat[i], mpu=self.mpu)) - - self.overflow = self.overflow_checker.check_using_norm(norm_groups) - prev_scale = self.cur_scale - self._update_scale(self.overflow) - - if self.overflow: - if self.verbose: - logger.info("[deepspeed] fp16 dynamic loss scale overflow! Skipping step. Attempted loss " - "scale: {}, reducing to {}".format(prev_scale, self.cur_scale)) - return self.overflow - - scaled_grad_norm = get_global_norm(norm_list=norm_groups) - - combined_scale = self.unscale_and_clip_grads(grads_groups_flat, scaled_grad_norm, apply_scale=False) - - # Stash unscaled gradient norm - self._global_grad_norm = scaled_grad_norm / self.cur_scale - - # norm is in fact norm*cur_scale - self.optimizer.step(grads=[[g] for g in grads_groups_flat], - output_params=[[p] for p in self.fp16_groups_flat], - scale=combined_scale, - grad_norms=norm_groups) - # TODO: we probably don't need this? just to be safe - for i in range(len(norm_groups)): - updated_params = _unflatten_dense_tensors(self.fp16_groups_flat[i], self.fp16_groups[i]) - for p, q in zip(self.fp16_groups[i], updated_params): - p.data = q.data - return self.overflow - - def set_lr(self, lr): - """Set the learning rate.""" - for param_group in self.optimizer.param_groups: - param_group["lr"] = lr - - def get_lr(self): - """Return the current learning rate.""" - return self.optimizer.param_groups[0]["lr"] - - def override_loss_scale(self, loss_scale): - if loss_scale != self.external_loss_scale: - logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}') - self.custom_loss_scaler = True - self.external_loss_scale = loss_scale - - def _require_avoid_recompute_norm(self, p, tensor_model_parallel_rank): - # for filtering replicated tensors from tensor - if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: - return True - if (tensor_model_parallel_rank > 0) and not is_model_parallel_parameter(p): - return True - - def _get_norm_mask_idx(self, group): - """The function preserves the parallel information for norm - from unflattened gradients. - - Args: - group (Iterable[Tensor] ): params group - - Returns: - torch.Tensor: A 2D tensor containing index ranges for each group, - where each row represents a [start index, end index]. - """ - group_mask_idx_list = [] - grad_flat_st_idx = 0 - grad_flat_en_idx = 0 - - for p in group: - grad_flat_en_idx = grad_flat_st_idx + p.numel() - if p.grad is not None and self._require_avoid_recompute_norm(p, bwc_tensor_model_parallel_rank(self.mpu)): - # merge range - if len(group_mask_idx_list) > 0 and grad_flat_st_idx == group_mask_idx_list[-1][-1]: - group_mask_idx_list[-1][-1] = grad_flat_en_idx - else: - group_mask_idx_list.append([grad_flat_st_idx, grad_flat_en_idx]) - grad_flat_st_idx = grad_flat_en_idx - - return torch.tensor(group_mask_idx_list, device=get_accelerator().current_device_name()) - - def step(self, closure=None): - """ - Not supporting closure. - """ - - if self.fused_adam_legacy: - return self.step_fused_adam() - - # First determine if there is overflow. - if self.timers: - self.timers(OVERFLOW_CHECK_TIMER).start() - fp16_params = [] - for i, group in enumerate(self.fp16_groups): - fp16_params.extend([p for p in group if p.grad is not None]) - self.overflow = self.overflow_checker.has_overflow(fp16_params) - if self.timers: - self.timers(OVERFLOW_CHECK_TIMER).stop() - prev_scale = self.cur_scale - self._update_scale(self.overflow) - if self.overflow: - if self.verbose: - log_dist( - "Overflow detected. Skipping step. Attempted loss " - f"scale: {prev_scale}, reducing to {self.cur_scale}", - ranks=[0]) - # Clear gradients - for i, group in enumerate(self.fp16_groups): - for p in group: - p.grad = None - - if self.timers: - self.timers.log(OVERFLOW_TIMERS) - return self.overflow - - grads_groups_flat = [] - non_experts_grads_for_norm = [] - expert_grads_for_norm = {} - assert len(self.fp16_groups) == len(self.optimizer.param_groups) - - for i, group in enumerate(self.fp16_groups): - data_type = self.fp32_groups_flat[i].dtype - - grads_groups_flat.append( - _flatten_dense_tensors([ - torch.zeros(p.size(), dtype=data_type, device=p.device) if p.grad is None else p.grad.to(data_type) - for p in group - ])) - - self.fp32_groups_flat[i].grad = grads_groups_flat[i] - param_group = self.optimizer.param_groups[i] - - # split expert and non_expert grads for norm - if self.has_moe_layers and is_moe_param_group(param_group): - if param_group['name'] not in expert_grads_for_norm: - expert_grads_for_norm[param_group['name']] = [] - - expert_grads_for_norm[param_group['name']].append(self.fp32_groups_flat[i]) - else: - # retrieves the required mask for calculating the norm of flat_grad - # perform this collect operation only once - if not self.has_executed_step: - cur_flat_grad_norm_mask = self._get_norm_mask_idx(group) - self.flatten_grad_norm_mask_list.append(cur_flat_grad_norm_mask) - - non_experts_grads_for_norm.append(self.fp32_groups_flat[i]) - - for p in group: - p.grad = None - - if self.timers: - self.timers(COMPUTE_NORM_TIMER).start() - - all_groups_norm = get_flattened_grad_norm(non_experts_grads_for_norm, - mpu=self.mpu, - grad_norm_mask=self.flatten_grad_norm_mask_list) - - if self.has_moe_layers: - all_groups_norm = get_norm_with_moe_layers(all_groups_norm, - mpu=self.mpu, - expert_tensors=expert_grads_for_norm, - norm_type=self.norm_type) - - scaled_global_grad_norm = get_global_norm(norm_list=[all_groups_norm]) - if self.timers: - self.timers(COMPUTE_NORM_TIMER).stop() - - # Stash unscaled gradient norm - self._global_grad_norm = scaled_global_grad_norm / self.cur_scale - - if self.timers: - self.timers(UNSCALE_AND_CLIP_TIMER).start() - self.unscale_and_clip_grads(grads_groups_flat, scaled_global_grad_norm) - if self.timers: - self.timers(UNSCALE_AND_CLIP_TIMER).stop() - - if self.timers: - self.timers(BASIC_STEP_TIMER).start() - self.optimizer.step() - if self.timers: - self.timers(BASIC_STEP_TIMER).stop() - - #get rid of the fp32 gradients. Not needed anymore - for group in self.fp32_groups_flat: - group.grad = None - - if self.timers: - self.timers(UPDATE_FP16_TIMER).start() - - for i in range(len(self.fp16_groups)): - updated_params = _unflatten_dense_tensors(self.fp32_groups_flat[i], self.fp16_groups[i]) - for p, q in zip(self.fp16_groups[i], updated_params): - p.data.copy_(q.data) - self.has_executed_step = True - if self.timers: - self.timers(UPDATE_FP16_TIMER).stop() - - if self.timers: - self.timers.log(STEP_TIMERS) - - return self.overflow - - def unscale_and_clip_grads(self, grad_groups_flat, total_norm, apply_scale=True): - # compute combined scale factor for this group - combined_scale = self.cur_scale - if self.clip_grad > 0.: - # norm is in fact norm*scale - clip = ((total_norm / self.cur_scale) + 1e-6) / self.clip_grad - if clip > 1: - combined_scale = clip * self.cur_scale - - if apply_scale: - for grad in grad_groups_flat: - grad.data.mul_(1. / combined_scale) - - return combined_scale - - def backward(self, loss, create_graph=False, retain_graph=False): - """ - :attr:`backward` performs the following steps: - - 1. fp32_loss = loss.float() - 2. scaled_loss = fp32_loss*loss_scale - 3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves - """ - if self.custom_loss_scaler: - scaled_loss = self.external_loss_scale * loss - scaled_loss.backward() - else: - scaled_loss = (loss.float()) * self.cur_scale - scaled_loss.backward(create_graph=create_graph, retain_graph=retain_graph) - - def _update_scale(self, skip): - if self.dynamic_loss_scale: - prev_scale = self.cur_scale - if skip: - self.cur_scale = max(self.cur_scale / self.scale_factor, self.min_loss_scale) - self.last_overflow_iter = self.cur_iter - if self.verbose: - logger.info(f"\nGrad overflow on iteration {self.cur_iter}") - logger.info(f"Reducing dynamic loss scale from {prev_scale} to {self.cur_scale}") - else: - # Ensure self.scale_window updates since last overflow - stable_interval = (self.cur_iter - self.last_overflow_iter) - 1 - if (stable_interval > 0) and (stable_interval % self.scale_window == 0): - self.cur_scale *= self.scale_factor - if self.verbose: - logger.info(f"No Grad overflow for {self.scale_window} iterations") - logger.info(f"Increasing dynamic loss scale from {prev_scale} to {self.cur_scale}") - else: - if skip: - logger.info("Grad overflow on iteration: %s", self.cur_iter) - logger.info("Using static loss scale of: %s", self.cur_scale) - self.cur_iter += 1 - return - - # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state" - def _get_state(self): - return self.optimizer.state - - def _set_state(self, value): - self.optimizer.state = value - - state = property(_get_state, _set_state) - - # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups" - # (for example, to adjust the learning rate) - def _get_param_groups(self): - return self.optimizer.param_groups - - def _set_param_groups(self, value): - self.optimizer.param_groups = value - - param_groups = property(_get_param_groups, _set_param_groups) - - def state_dict(self): - """ - Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. - This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict - of the contained Pytorch optimizer. - Example:: - checkpoint = {} - checkpoint['model'] = model.state_dict() - checkpoint['optimizer'] = optimizer.state_dict() - torch.save(checkpoint, "saved.pth") - """ - state_dict = {} - state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale - state_dict['cur_scale'] = self.cur_scale - state_dict['cur_iter'] = self.cur_iter - if state_dict['dynamic_loss_scale']: - state_dict['last_overflow_iter'] = self.last_overflow_iter - state_dict['scale_factor'] = self.scale_factor - state_dict['scale_window'] = self.scale_window - state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict() - state_dict['fp32_groups_flat'] = self.fp32_groups_flat - state_dict[CLIP_GRAD] = self.clip_grad - return state_dict - - # Refresh fp32 master params from fp16 copies - def refresh_fp32_params(self): - for current, saved in zip(self.fp32_groups_flat, self.fp16_groups_flat): - current.data.copy_(saved.data) - - def load_state_dict(self, state_dict, load_optimizer_states=True): - """ - Loads a state_dict created by an earlier call to state_dict(). - If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``, - whose parameters in turn came from ``model``, it is expected that the user - will call ``model.load_state_dict()`` before - ``fp16_optimizer_instance.load_state_dict()`` is called. - Example:: - model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half() - optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) - optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0) - ... - checkpoint = torch.load("saved.pth") - model.load_state_dict(checkpoint['model']) - optimizer.load_state_dict(checkpoint['optimizer']) - """ - # I think it should actually be ok to reload the optimizer before the model. - self.dynamic_loss_scale = state_dict['dynamic_loss_scale'] - self.cur_scale = state_dict['cur_scale'] - self.cur_iter = state_dict['cur_iter'] - if state_dict['dynamic_loss_scale']: - self.last_overflow_iter = state_dict['last_overflow_iter'] - self.scale_factor = state_dict['scale_factor'] - self.scale_window = state_dict['scale_window'] - if load_optimizer_states: - self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT]) - self.clip_grad = state_dict[CLIP_GRAD] - # At this point, the optimizer's references to the model's fp32 parameters are up to date. - # The optimizer's hyperparameters and internal buffers are also up to date. - # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still - # out of date. There are two options. - # 1: Refresh the master params from the model's fp16 params. - # This requires less storage but incurs precision loss. - # 2: Save and restore the fp32 master copies separately. - # We choose option 2. - # - # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device - # of their associated parameters, because it's possible those buffers might not exist yet in - # the current optimizer instance. In our case, as long as the current FP16_Optimizer has been - # constructed in the same way as the one whose state_dict we are loading, the same master params - # are guaranteed to exist, so we can just copy_() from the saved master params. - for current, saved in zip(self.fp32_groups_flat, state_dict['fp32_groups_flat']): - current.data.copy_(saved.data) - - def __repr__(self): - return repr(self.optimizer) - - # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale" - def _get_loss_scale(self): - if self.custom_loss_scaler: - return self.external_loss_scale - else: - return self.cur_scale - - def _set_loss_scale(self, value): - self.loss_scaler.cur_scale = value - - loss_scale = property(_get_loss_scale, _set_loss_scale) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/loss_scaler.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/loss_scaler.py deleted file mode 100644 index 3f5ed019fc78fdfe0a7a5e610fdce6f3b8498e60..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/loss_scaler.py +++ /dev/null @@ -1,272 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Copyright (c) 2019, NVIDIA CORPORATION. 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. -Taken and modified for DeepSpeed from: - https://github.com/NVIDIA/Megatron-LM/blob/master/fp16/loss_scaler.py -Commit: 93ab4bea59dc5cbf97c079d313741866af4deac9 -""" - -import torch -from deepspeed.runtime.config_utils import DeepSpeedConfigObject -from deepspeed import comm as dist -from deepspeed.utils import logger - -INITIAL_LOSS_SCALE = 'init_scale' -SCALE_WINDOW = 'scale_window' -DELAYED_SHIFT = 'delayed_shift' -CONSECUTIVE_HYSTERESIS = 'consecutive_hysteresis' -MIN_LOSS_SCALE = 'min_scale' - - -# item() is a recent addition, so this helps with backward compatibility. -def to_python_float(t): - if hasattr(t, 'item'): - return t.item() - return t[0] - - -class LossScalerBase(DeepSpeedConfigObject): - """LossScalarBase - Base class for a loss scaler - """ - - def __init__(self, cur_scale): - super(LossScalerBase, self).__init__() - self.cur_scale = cur_scale - self.dynamic = False - - @property - def loss_scale(self): - return self.cur_scale - - def scale_gradient(self, module, grad_in, grad_out): - return tuple(self.loss_scale * g for g in grad_in) - - def update_scale(self, overflow): - pass - - def backward(self, loss, retain_graph=False): - scaled_loss = loss * self.loss_scale - scaled_loss.backward(retain_graph=retain_graph) - # print(f'LossScalerBackward: {scaled_loss=}') - - -class LossScaler(LossScalerBase): - """ - Class that manages a static loss scale. This class is intended to interact with - :class:`FP16_Optimizer`, and should not be directly manipulated by the user. - - Use of :class:`LossScaler` is enabled via the ``static_loss_scale`` argument to - :class:`FP16_Optimizer`'s constructor. - - Args: - scale (float, optional, default=1.0): The loss scale. - """ - - def __init__(self, scale=1): - super(LossScaler, self).__init__(scale) - - # `params` is a list / generator of torch.Variable - def has_overflow(self, params): - return False - - # `x` is a torch.Tensor - def _has_inf_or_nan(x): - return False - - -class DynamicLossScaler(LossScalerBase): - """ - Class that manages dynamic loss scaling. It is recommended to use :class:`DynamicLossScaler` - indirectly, by supplying ``dynamic_loss_scale=True`` to the constructor of - :class:`FP16_Optimizer`. However, it's important to understand how :class:`DynamicLossScaler` - operates, because the default options can be changed using the - the ``dynamic_loss_args`` argument to :class:`FP16_Optimizer`'s constructor. - - Loss scaling is designed to combat the problem of underflowing gradients encountered at long - times when training fp16 networks. Dynamic loss scaling begins by attempting a very high loss - scale. Ironically, this may result in OVERflowing gradients. If overflowing gradients are - encountered, :class:`DynamicLossScaler` informs :class:`FP16_Optimizer` that an overflow has - occurred. - :class:`FP16_Optimizer` then skips the update step for this particular iteration/minibatch, - and :class:`DynamicLossScaler` adjusts the loss scale to a lower value. - If a certain number of iterations occur without overflowing gradients detected, - :class:`DynamicLossScaler` increases the loss scale once more. - In this way :class:`DynamicLossScaler` attempts to "ride the edge" of - always using the highest loss scale possible without incurring overflow. - - Args: - init_scale (float, optional, default=2**32): Initial loss scale attempted by :class:`DynamicLossScaler.` - scale_factor (float, optional, default=2.0): Factor used when adjusting the loss scale. If an overflow is encountered, the loss scale is readjusted to loss scale/``scale_factor``. If ``scale_window`` consecutive iterations take place without an overflow, the loss scale is readjusted to loss_scale*``scale_factor``. - scale_window (int, optional, default=1000): Number of consecutive iterations without an overflow to wait before increasing the loss scale. - consecutive_hysteresis (bool, optional, default=False): Whether to refill hysteresis if we reach an iteration that doesn't overflow - """ - - def __init__(self, - init_scale, - scale_window, - min_scale, - delayed_shift, - consecutive_hysteresis, - raise_error_at_min_scale=True, - dtype=torch.half): - super(DynamicLossScaler, self).__init__(init_scale) - self.cur_iter = 0 - self.last_overflow_iter = -1 - self.scale_factor = 2.0 - self.scale_window = scale_window - self.min_scale = min_scale - self.delayed_shift = delayed_shift - self.cur_hysteresis = delayed_shift - self.consecutive_hysteresis = consecutive_hysteresis - self.raise_error_at_min_scale = raise_error_at_min_scale - self.dynamic = True - self.dtype = dtype - - # `params` is a list / generator of torch.Variable - def has_overflow_serial(self, params): - for p in params: - if p.grad is not None and self._has_inf_or_nan(p.grad.data): - return True - - return False - - # `x` is a torch.Tensor - def _has_inf_or_nan(x): - try: - # if x is half, the .float() incurs an additional deep copy, but it's necessary if - # Pytorch's .sum() creates a one-element tensor of the same type as x - # (which is true for some recent version of pytorch). - cpu_sum = float(x.float().sum()) - # More efficient version that can be used if .sum() returns a Python scalar - # cpu_sum = float(x.sum()) - except RuntimeError as instance: - # We want to check if inst is actually an overflow exception. - # RuntimeError could come from a different error. - # If so, we still want the exception to propagate. - if "value cannot be converted" not in instance.args[0]: - raise - return True - else: - if cpu_sum in [float('inf'), -float('inf')] or cpu_sum != cpu_sum: - return True - return False - - # `overflow` is boolean indicating whether the gradient overflowed - def update_scale(self, overflow): - if overflow: - # self.cur_scale /= self.scale_factor - if self.delayed_shift == 1 or self.cur_hysteresis == 1: - if (self.cur_scale == self.min_scale) and self.raise_error_at_min_scale: - raise Exception( - "Current loss scale already at minimum - cannot decrease scale anymore. Exiting run.") - else: - next_scale = max(self.cur_scale / self.scale_factor, self.min_scale) - if dist.get_rank() == 0: - overflow_msg = f"[deepspeed] OVERFLOW! Rank {dist.get_rank()} Skipping step." - if self.dtype == torch.half: - overflow_msg += f" Attempted loss scale: {int(self.cur_scale)}, reducing to {int(next_scale)}" - logger.info(overflow_msg) - self.cur_scale = next_scale - else: - if dist.get_rank() == 0: - overflow_msg = f"[deepspeed] OVERFLOW! Rank {dist.get_rank()} Skipping step." - if self.dtype == torch.half: - overflow_msg += f" Attempted loss scale: {int(self.cur_scale)}, but hysteresis is {self.cur_hysteresis}. Reducing hysteresis to {self.cur_hysteresis-1}" - logger.info(overflow_msg) - self.cur_hysteresis -= 1 - self.last_overflow_iter = self.cur_iter - else: - if self.consecutive_hysteresis: - if dist.get_rank() == 0: - hysteresis_msg = f"Consecutive hysteresis is enabled. Restoring hysteresis to {self.delayed_shift}" - logger.info(hysteresis_msg) - self.cur_hysteresis = self.delayed_shift - - stable_interval = (self.cur_iter - self.last_overflow_iter) - 1 - if (stable_interval > 0) and (stable_interval % self.scale_window == 0): - if not self.consecutive_hysteresis: - self.cur_hysteresis = self.delayed_shift - self.cur_scale *= self.scale_factor - self.cur_iter += 1 - - -# Although loss scaling is only defined for fp16, yet for backwards compatibility -# we still create a scaler for other dtypes (fp32, bf16) which does not perform any scaling. -def CreateLossScaler(dtype, static_loss_scale, dynamic_scaling, dynamic_loss_args): - if dtype == torch.half and dynamic_scaling: - assert dynamic_loss_args is not None, f"Dynamic loss scaling parameters must be defined." - return DynamicLossScaler(dtype=dtype, **dynamic_loss_args) - - loss_scale_value = static_loss_scale if dtype == torch.half else 1.0 - return LossScaler(scale=loss_scale_value) - - -############################################################## -# Example usage below here -- assuming it's in a separate file -############################################################## -""" -TO-DO separate out into an example. -if __name__ == "__main__": - import torch - from torch.autograd import Variable - from dynamic_loss_scaler import DynamicLossScaler - - # N is batch size; D_in is input dimension; - # H is hidden dimension; D_out is output dimension. - N, D_in, H, D_out = 64, 1000, 100, 10 - - # Create random Tensors to hold inputs and outputs, and wrap them in Variables. - x = Variable(torch.randn(N, D_in), requires_grad=False) - y = Variable(torch.randn(N, D_out), requires_grad=False) - - w1 = Variable(torch.randn(D_in, H), requires_grad=True) - w2 = Variable(torch.randn(H, D_out), requires_grad=True) - parameters = [w1, w2] - - learning_rate = 1e-6 - optimizer = torch.optim.SGD(parameters, lr=learning_rate) - loss_scaler = DynamicLossScaler() - - for t in range(500): - y_pred = x.mm(w1).clamp(min=0).mm(w2) - loss = (y_pred - y).pow(2).sum() * loss_scaler.loss_scale - print('Iter {} loss scale: {}'.format(t, loss_scaler.loss_scale)) - print('Iter {} scaled loss: {}'.format(t, loss.data[0])) - print('Iter {} unscaled loss: {}'.format(t, loss.data[0] / loss_scaler.loss_scale)) - - # Run backprop - optimizer.zero_grad() - loss.backward() - - # Check for overflow - has_overflow = DynamicLossScaler.has_overflow(parameters) - - # If no overflow, unscale grad and update as usual - if not has_overflow: - for param in parameters: - param.grad.data.mul_(1. / loss_scaler.loss_scale) - optimizer.step() - # Otherwise, don't do anything -- ie, skip iteration - else: - print('fp16 dynamic loss scale overflow!') - - # Update loss scale for next iteration - loss_scaler.update_scale(has_overflow) - -""" diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__init__.py deleted file mode 100644 index ac6f7e9784ce60f6c1b4a9134b73b4e415337641..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .adam import OnebitAdam -from .lamb import OnebitLamb -from .zoadam import ZeroOneAdam diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index a1d97208a48f40cb56e8695e7f7abc747a982489..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/adam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/adam.cpython-310.pyc deleted file mode 100644 index d6ffe01b1b105ac2db2116434034fb73131c732e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/adam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/lamb.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/lamb.cpython-310.pyc deleted file mode 100644 index ad36fd135f8a75f0c170cd3e8241797c5994870d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/lamb.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/zoadam.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/zoadam.cpython-310.pyc deleted file mode 100644 index 728edd6a08e7a28f3d959ac8acae041eaf0fa28c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/zoadam.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/adam.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/adam.py deleted file mode 100644 index fa817573f734801070de5391d161695389351ec3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/adam.py +++ /dev/null @@ -1,310 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import types -import torch -import numpy as np -from deepspeed.accelerator import get_accelerator -from deepspeed.utils.torch import required_torch_version -from deepspeed import comm as dist - - -class OnebitAdam(torch.optim.Optimizer): - """Implements the 1-bit Adam algorithm. Currently GPU-only. - For usage example please see https://www.deepspeed.ai/tutorials/onebit-adam/ - For technical details please read https://arxiv.org/abs/2102.02888 - - Arguments: - params (iterable): iterable of parameters to optimize or dicts defining - parameter groups. - lr (float, optional): learning rate. (default: 1e-3) - freeze_step (int, optional): Number of steps for warmup (uncompressed) - stage before we start using compressed communication. (default 100000) - betas (Tuple[float, float], optional): coefficients used for computing - running averages of gradient and its square. (default: (0.9, 0.999)) - eps (float, optional): term added to the denominator to improve - numerical stability. (default: 1e-8) - weight_decay (float, optional): weight decay (L2 penalty) (default: 0) - amsgrad (boolean, optional): whether to use the AMSGrad variant of this - algorithm from the paper `On the Convergence of Adam and Beyond`_ - (default: False) NOT SUPPORTED in 1-bit Adam! - eps_inside_sqrt (boolean, optional): in the 'update parameters' step, - adds eps to the bias-corrected second moment estimate before - evaluating square root instead of adding it to the square root of - second moment estimate as in the original paper. (default: False) - cuda_aware (boolean, required): Set True if the underlying MPI implementation - supports CUDA-Aware communication. (default: False) - comm_backend_name (string, optional): Set to 'mpi' if needed. (default: 'nccl') - .. _Adam\\: A Method for Stochastic Optimization: - https://arxiv.org/abs/1412.6980 - .. _On the Convergence of Adam and Beyond: - https://openreview.net/forum?id=ryQu7f-RZ - """ - - def __init__(self, - params, - deepspeed=None, - lr=1e-3, - freeze_step=100000, - bias_correction=True, - betas=(0.9, 0.999), - eps=1e-8, - eps_inside_sqrt=False, - weight_decay=0., - max_grad_norm=0., - amsgrad=False, - cuda_aware=False, - comm_backend_name='nccl'): - - if amsgrad: - raise RuntimeError('1-bit Adam does not support the AMSGrad variant.') - - defaults = dict(lr=lr, - bias_correction=bias_correction, - betas=betas, - eps=eps, - weight_decay=weight_decay, - max_grad_norm=max_grad_norm) - - super(OnebitAdam, self).__init__(params, defaults) - self.eps_mode = 0 if eps_inside_sqrt else 1 - self.comm_time = 0.0 - self.step_time = 0.0 - self.ave_step = 1 - self.bk_time = 0.0 - - self.deepspeed = deepspeed - self.adam_freeze_key = False - self.initialize = False - self.freeze_step = freeze_step - self.cuda_aware = cuda_aware - self.using_pipeline = False - - self.comm_backend_name = comm_backend_name - - assert dist.is_initialized(), "Please initialize the torch distributed backend." - # Empty initializer. Set handle based on the comm backend as follows. - self.comm_backend_handle = None - if self.comm_backend_name == 'nccl': - assert ( - required_torch_version(min_version=1.8) - ), "Please use torch 1.8 or greater to enable NCCL backend in 1-bit Adam. Alternatively, please specify 'mpi' as the 'comm_backend_name' in config file to proceed with the MPI backend" - from deepspeed.runtime.comm.nccl import NcclBackend - self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') - self.comm_backend_handle = NcclBackend(self.deepspeed.mpu) - elif self.comm_backend_name == 'mpi': - from deepspeed.runtime.comm.mpi import MpiBackend - self.comm_backend_handle = MpiBackend(cuda_aware) - elif self.comm_backend_name == 'hccl': - from deepspeed.runtime.comm.hccl import HcclBackend - self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') - self.comm_backend_handle = HcclBackend(self.deepspeed.mpu) - elif self.comm_backend_name == 'compressed': - from deepspeed.runtime.comm.compressed import CompressedBackend - self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') - self.comm_backend_handle = CompressedBackend(self.deepspeed.mpu) - self.size = self.comm_backend_handle.size - - self.divider = int(self.size * 8 / np.gcd(self.size, 8)) - - def step(self, closure=None, grads=None): - """Performs a single optimization step. - Arguments: - closure (callable, optional): A closure that reevaluates the model - and returns the loss. - grads (list of tensors, optional): weight gradient to use for the - optimizer update. If gradients have type torch.half, parameters - are expected to be in type torch.float. (default: None) - output params (list of tensors, optional): A reduced precision copy - of the updated weights written out in addition to the regular - updated weights. Have to be of same type as gradients. (default: None) - scale (float, optional): factor to divide gradient tensor values - by before applying to weights. (default: 1) - """ - loss = None - if closure is not None: - loss = closure() - - gather_time = 0 - allgather_time = 0 - all_time = 0 - - if self.adam_freeze_key is False: - v_diff_buffer = 0.0 - - if grads is None: - grads_group = [None] * len(self.param_groups) - # backward compatibility - # assuming a list/generator of parameter means single group - elif isinstance(grads, types.GeneratorType): - grads_group = [grads] - elif type(grads[0]) != list: - grads_group = [grads] - else: - grads_group = grads - - for group, grads_this_group in zip(self.param_groups, grads_group): - if grads_this_group is None: - grads_this_group = [None] * len(group['params']) - - bias_correction = 1 if group['bias_correction'] else 0 - - for p, grad in zip(group['params'], grads_this_group): - if p.grad is None and grad is None: - continue - if grad is None: - grad = p.grad.data - if grad.is_sparse: - raise RuntimeError('1-bit Adam does not support sparse gradients') - - state = self.state[p] - - # State initialization - if len(state) == 0: - state['step'] = 0 - # Exponential moving average of gradient values - state['exp_avg'] = torch.zeros_like(p.data) - # Exponential moving average of squared gradient values - state['exp_avg_sq'] = torch.zeros_like(p.data) - - if not self.initialize or (self.adam_freeze_key and 'worker_error' not in state.keys()): - state['tensor_size'] = torch.numel(p.data) - state['corrected_tensor_size'] = state['tensor_size'] - - if state['tensor_size'] % (self.size * self.divider) != 0: - state['corrected_tensor_size'] += ((self.size * self.divider) - (state['tensor_size'] % - (self.size * self.divider))) - state['server_chunk_size'] = state['corrected_tensor_size'] // self.size - get_accelerator().empty_cache() - state['worker_error'] = torch.zeros(state['corrected_tensor_size'], device=p.device) - state['server_error'] = torch.zeros(state['server_chunk_size'], device=p.device) - get_accelerator().empty_cache() - self.adam_freeze_key = True - if not self.initialize and dist.get_rank() == 0: - print("Cupy Buffers Initialized Successfully.") - - exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] - beta1, beta2 = group['betas'] - - state['step'] += 1 - - if self.adam_freeze_key is False: - exp_avg.mul_(beta1).add_(1 - beta1, grad) - exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) - grad = None - if self.initialize: - update = exp_avg / (exp_avg_sq.sqrt() + group['eps']) - - else: - if 'non_freeze' in group.keys() and group['non_freeze'] is True: - dist.all_reduce(grad) - grad.mul_(1 / dist.get_world_size()) - exp_avg.mul_(beta1).add_(1 - beta1, grad) - exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) - grad = None - else: - if self.initialize is True: - exp_avg.mul_(beta1).add_(1 - beta1, grad) - grad = None - - if self.size > 1: - exp_avg.set_( - self.comm_backend_handle.compressed_allreduce(exp_avg, state['worker_error'], - state['server_error'], - self.deepspeed.local_rank)) - # Because 1-bit compression cannot represent exact zero, it is required to - # provide a momentum mask for those params that have constant exact zeros in their - # momentums, otherwise the compression error would keep accumulating. - # For example, for BERT pre-training seq 128, bert.embeddings.position_embeddings.weight - # always have exact zeros in its momentum for row 129 to 512, because it only - # learns up to seq length 128 while the model supports up to 512 seq length. - # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py.) - if 'exp_avg_mask' in group: - if exp_avg.device != group['exp_avg_mask'].device: - group['exp_avg_mask'] = group['exp_avg_mask'].to(device=exp_avg.device) - exp_avg.mul_(group['exp_avg_mask']) - - if self.initialize: - update = exp_avg / (exp_avg_sq.sqrt() + group['eps']) - - if self.initialize: - if group['weight_decay'] > 0.0: - update += group['weight_decay'] * p.data - with torch.no_grad(): - p.add_(-group['lr'] * update) - - if not self.initialize: - print('Pop out errors', flush=True) - state.pop('worker_error') - state.pop('server_error') - - if not self.initialize: - self.adam_freeze_key = False - self.initialize = True - print(f"Finished the initialization step at rank {dist.get_rank()}") - return loss - - if self.adam_freeze_key is False: - if state['step'] >= self.freeze_step: - print('OnebitAdam - starting compressed communication') - self.adam_freeze_key = True - if self.using_pipeline: - self.deepspeed.pipeline_enable_backward_allreduce = False - else: - self.deepspeed.enable_backward_allreduce = False - - return loss - - def load_state_dict(self, state_dict): - """ - Overrides load_state_dict() to add special handling when loading checkpoints - """ - # Because at different stage exp_avg_mask may change (e.g., - # BERT pre-training seqlen 128 and 512 ), we don't use the exp_avg_mask - # in checkpoints but always use the one user provided in training script. - # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py.) - # Thus here we keep the exp_avg_mask unchanged when loading checkpoint - for i, group in enumerate(self.param_groups): - if 'exp_avg_mask' in group: - state_dict['param_groups'][i]['exp_avg_mask'] = group['exp_avg_mask'] - elif 'exp_avg_mask' not in group and 'exp_avg_mask' in state_dict['param_groups'][i]: - state_dict['param_groups'][i].pop('exp_avg_mask') - super().load_state_dict(state_dict) - if self.state[self.param_groups[0]['params'][0]]['step'] < self.freeze_step: - if dist.get_rank() == 0: - print("Checkpoint loaded and OnebitAdam warmup stage starts/continues.") - if self.adam_freeze_key is True: - self.adam_freeze_key = False - if self.using_pipeline: - self.deepspeed.pipeline_enable_backward_allreduce = True - else: - self.deepspeed.enable_backward_allreduce = True - else: - if dist.get_rank() == 0: - print("Checkpoint loaded and OnebitAdam compression stage starts/continues.") - if self.adam_freeze_key is False: - self.adam_freeze_key = True - if self.using_pipeline: - self.deepspeed.pipeline_enable_backward_allreduce = False - else: - self.deepspeed.enable_backward_allreduce = False - # We reset the compression errors when loading checkpoints for 3 reasons: - # 1) The worker and server error at each GPU are distinct, so in current implementation - # only rank 0's errors are saved in the checkpoint. Thus we have to reset the errors. - # If we want to save them correctly we need O(num_gpu*model_size) memory in order to - # gather all the error, which is a very large memory requirement. It's possible to save - # them in a distributed way, but it will make the checkpoint saving/loading much more complicated. - # 2) Even if we are able to save the compression errors correctly, you need to have the - # exact same number of GPUs in order to load them correctly. - # 3) We verified on BERT pre-training that occasionally resetting the compression error - # at checkpoint loading does not affect the convergence. - # However, please avoid frequent checkpoint loading which could break the error - # compensation mechanism thus affect the convergence. - for group in self.param_groups: - for p in group['params']: - if 'worker_error' in self.state[p]: - self.state[p].pop('worker_error') - if 'server_error' in self.state[p]: - self.state[p].pop('server_error') diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/lamb.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/lamb.py deleted file mode 100644 index 9e7bae816ecdee5aa7e5aa2349a62f1b43ea5e10..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/lamb.py +++ /dev/null @@ -1,447 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import types -import torch -import numpy as np -from deepspeed import comm as dist -from deepspeed.utils.torch import required_torch_version -from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors -from deepspeed.accelerator import get_accelerator - - -class OnebitLamb(torch.optim.Optimizer): - """Implements the 1-bit Lamb algorithm. Currently GPU-only. - For usage example please see https://www.deepspeed.ai/tutorials/onebit-lamb/ - For technical details please see our paper https://arxiv.org/abs/2104.06069. - - Arguments: - params (iterable): iterable of parameters to optimize or dicts defining - parameter groups. - lr (float, optional): learning rate. (default: 1e-3) - freeze_step (int, optional): Number of steps for warmup (uncompressed) - stage before we start using compressed communication. (default 100000) - betas (Tuple[float, float], optional): coefficients used for computing - running averages of gradient and its square. (default: (0.9, 0.999)) - eps (float, optional): term added to the denominator to improve - numerical stability. (default: 1e-8) - weight_decay (float, optional): weight decay (L2 penalty) (default: 0) - max_coeff(float, optional): maximum value of the lamb coefficient (default: 10.0) - min_coeff(float, optional): minimum value of the lamb coefficient (default: 0.01) - amsgrad (boolean, optional): whether to use the AMSGrad variant of this - algorithm from the paper `On the Convergence of Adam and Beyond`_ - (default: False) NOT SUPPORTED in 1-bit Lamb! - eps_inside_sqrt (boolean, optional): in the 'update parameters' step, - adds eps to the bias-corrected second moment estimate before - evaluating square root instead of adding it to the square root of - second moment estimate as in the original paper. (default: False) - cuda_aware (boolean, required): Set True if the underlying MPI implementation - supports CUDA-Aware communication. (default: False) - comm_backend_name (string, optional): Set to 'mpi' if needed. (default: 'nccl') - coeff_beta (float, optional): coefficient used for computing - running averages of lamb coefficient (default: 0.9) note that you may want to - increase or decrease this beta depending on the freeze_step you choose, as - 1/(1 - coeff_beta) should be smaller than or equal to freeze_step - factor_max (float, optional): maximum value of scaling factor to the frozen lamb - coefficient during compression stage (default: 4.0) - factor_min (float, optional): minimum value of scaling factor to the frozen lamb - coefficient during compression stage (default: 0.5) - factor_threshold (float, optional): threshold of how much the scaling factor can - fluctuate between steps (default: 0.1) - .. _Large Batch Optimization for Deep Learning\\: Training BERT in 76 minutes: - https://arxiv.org/abs/1904.00962 - .. _Adam\\: A Method for Stochastic Optimization: - https://arxiv.org/abs/1412.6980 - .. _On the Convergence of Adam and Beyond: - https://openreview.net/forum?id=ryQu7f-RZ - """ - - def __init__(self, - params, - deepspeed=None, - lr=1e-3, - freeze_step=100000, - bias_correction=True, - betas=(0.9, 0.999), - eps=1e-8, - eps_inside_sqrt=False, - weight_decay=0., - max_grad_norm=0., - max_coeff=10.0, - min_coeff=0.01, - amsgrad=False, - cuda_aware=False, - comm_backend_name='nccl', - coeff_beta=0.9, - factor_max=4.0, - factor_min=0.5, - factor_threshold=0.1): - - if amsgrad: - raise RuntimeError('1-bit Lamb does not support the AMSGrad variant.') - - defaults = dict(lr=lr, - bias_correction=bias_correction, - betas=betas, - eps=eps, - weight_decay=weight_decay, - max_grad_norm=max_grad_norm, - max_coeff=max_coeff, - min_coeff=min_coeff) - - super(OnebitLamb, self).__init__(params, defaults) - self.eps_mode = 0 if eps_inside_sqrt else 1 - self.deepspeed = deepspeed - self.lamb_freeze_key = False - self.initialize = False - self.freeze_step = freeze_step - self.cuda_aware = cuda_aware - self.coeff_beta = coeff_beta - self.factor_max = factor_max - self.factor_min = factor_min - self.factor_threshold = factor_threshold - self.using_pipeline = False - - self.comm_backend_name = comm_backend_name - - assert dist.is_initialized(), "Please initialize the torch distributed backend." - # Empty initializer. Set handle based on the comm backend as follows. - self.comm_backend_handle = None - if self.comm_backend_name == 'nccl': - assert ( - required_torch_version(min_version=1.8) - ), "Please use torch 1.8 or greater to enable NCCL backend in 1-bit Adam. Alternatively, please specify 'mpi' as the 'comm_backend_name' in config file to proceed with the MPI backend" - from deepspeed.runtime.comm.nccl import NcclBackend - self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') - self.comm_backend_handle = NcclBackend(self.deepspeed.mpu) - elif self.comm_backend_name == 'mpi': - from deepspeed.runtime.comm.mpi import MpiBackend - self.comm_backend_handle = MpiBackend(cuda_aware) - elif self.comm_backend_name == 'hccl': - from deepspeed.runtime.comm.hccl import HcclBackend - self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') - self.comm_backend_handle = HcclBackend(self.deepspeed.mpu) - elif self.comm_backend_name == 'compressed': - from deepspeed.runtime.comm.compressed import CompressedBackend - self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') - self.comm_backend_handle = CompressedBackend(self.deepspeed.mpu) - - self.size = self.comm_backend_handle.size - - self.divider = int(self.size * 8 / np.gcd(self.size, 8)) - - self.exp_avg_flat = [] - self.dummy_exp_avg = {} - self.corrected_tensor_sizes = [] - self.server_chunk_sizes = [] - self.worker_errors = [] - self.server_errors = [] - - self.lamb_coeffs = [] - - def step(self, closure=None, grads=None): - """Performs a single optimization step. - Arguments: - closure (callable, optional): A closure that reevaluates the model - and returns the loss. - grads (list of tensors, optional): weight gradient to use for the - optimizer update. If gradients have type torch.half, parameters - are expected to be in type torch.float. (default: None) - """ - loss = None - if closure is not None: - loss = closure() - - if grads is None: - grads_group = [None] * len(self.param_groups) - # backward compatibility - # assuming a list/generator of parameter means single group - elif isinstance(grads, types.GeneratorType): - grads_group = [grads] - elif type(grads[0]) != list: - grads_group = [grads] - else: - grads_group = grads - - # remove the previous stats - del self.lamb_coeffs[:] - - if self.lamb_freeze_key: - exp_avg_last_step = [] - for group in self.param_groups: - exp_avg_last_step.append([self.state[p]['exp_avg'].detach().clone() for p in group['params']]) - if 'scaling_coeff' not in self.state[self.param_groups[0]['params'][0]]: - # Compute the scaling_coeff for each momentum at the end of warmup stage. - # This is used to reduce compression error during compression stage. - momentum_scales = [] - for group in self.param_groups: - momentum_scales.append([(torch.linalg.vector_norm(self.state[p]['exp_avg']) / - np.sqrt(torch.numel(self.state[p]['exp_avg']))).item() - for p in group['params']]) - united_scale = sum([sum(x) for x in momentum_scales]) / sum([len(x) for x in momentum_scales]) - for i, group in enumerate(self.param_groups): - for j, p in enumerate(group['params']): - self.state[p]['scaling_coeff'] = united_scale / momentum_scales[i][j] - - for group, grads_this_group in zip(self.param_groups, grads_group): - if grads_this_group is None: - grads_this_group = [None] * len(group['params']) - - bias_correction = 1 if group['bias_correction'] else 0 - - for p, grad in zip(group['params'], grads_this_group): - if p.grad is None and grad is None: - continue - if grad is None: - grad = p.grad.data - if grad.is_sparse: - raise RuntimeError('1-bit Lamb does not support sparse gradients') - - state = self.state[p] - - # State initialization - if len(state) == 0 or (len(state) == 1 and 'scaling_coeff' in state.keys()): - state['step'] = 0 - state['lamb_coeff_freeze'] = 0.0 - state['last_factor'] = 1.0 - # Exponential moving average of gradient values - state['exp_avg'] = torch.zeros_like(p.data) - # Exponential moving average of squared gradient values - state['exp_avg_sq'] = torch.zeros_like(p.data) - state['exp_avg_sq_fresh'] = torch.zeros_like(p.data) - - if not self.initialize: - self.lamb_freeze_key = True - - exp_avg, exp_avg_sq, exp_avg_sq_fresh = state['exp_avg'], state['exp_avg_sq'], state[ - 'exp_avg_sq_fresh'] - beta1, beta2 = group['betas'] - max_coeff = group['max_coeff'] - min_coeff = group['min_coeff'] - - state['step'] += 1 - - if self.lamb_freeze_key is False: - # warmup stage, baseline Lamb optimization - exp_avg.mul_(beta1).add_(1 - beta1, grad) - exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) - if state['step'] == self.freeze_step: - exp_avg_sq_fresh.data = exp_avg_sq.detach().clone() - grad = None - if self.initialize: - weight_norm = p.data.pow(2).sum().sqrt() - update = exp_avg / (exp_avg_sq.sqrt() + group['eps']) - if group['weight_decay'] > 0.0: - update += group['weight_decay'] * p.data - update_norm = update.pow(2).sum().sqrt() - lamb_coeff = 1.0 - if weight_norm != 0 and update_norm != 0: - lamb_coeff = (weight_norm / update_norm).item() - if lamb_coeff > max_coeff: - lamb_coeff = max_coeff - if lamb_coeff < min_coeff: - lamb_coeff = min_coeff - if lamb_coeff != 1.0: - state['lamb_coeff_freeze'] = self.coeff_beta * state['lamb_coeff_freeze'] + ( - 1 - self.coeff_beta) * lamb_coeff - self.lamb_coeffs.append(lamb_coeff) - with torch.no_grad(): - p.add_(-group['lr'] * lamb_coeff * update) - else: - # compression stage, update each momentum locally, then - # communicate based on the compressed_allreduce below - if self.initialize: - exp_avg.mul_(beta1).add_(1 - beta1, grad) - exp_avg.mul_(self.state[p]['scaling_coeff']) - grad = None - - # init fused momentum - if len(self.exp_avg_flat) == 0: - momentum_groups = [] - tensor_size = 0 - for group in self.param_groups: - for p in group['params']: - momentum_groups.append(self.state[p]['exp_avg']) - tensor_size += torch.numel(p.data) - corrected_tensor_size = tensor_size - if tensor_size % (self.size * self.divider) != 0: - difference = ((self.size * self.divider) - (tensor_size % (self.size * self.divider))) - corrected_tensor_size += difference - self.dummy_exp_avg[0] = torch.zeros(difference, device=momentum_groups[0].data.device) - momentum_groups.append(self.dummy_exp_avg[0]) - self.corrected_tensor_sizes.append(corrected_tensor_size) - self.server_chunk_sizes.append(corrected_tensor_size // self.size) - - self.exp_avg_flat.append(_flatten_dense_tensors([p.detach().clone() for p in momentum_groups])) - updated_params = _unflatten_dense_tensors(self.exp_avg_flat[0], momentum_groups) - for p, q in zip(momentum_groups, updated_params): - p.data = q.data - - if self.initialize and len(self.worker_errors) == 0: - get_accelerator().empty_cache() - for i in range(len(self.exp_avg_flat)): - self.worker_errors.append( - torch.zeros(self.corrected_tensor_sizes[i], device=self.exp_avg_flat[i].device)) - self.server_errors.append(torch.zeros(self.server_chunk_sizes[i], device=self.exp_avg_flat[i].device)) - get_accelerator().empty_cache() - - if self.lamb_freeze_key: - if self.size > 1: - for i in range(len(self.exp_avg_flat)): - if not self.initialize: - get_accelerator().empty_cache() - self.worker_errors.append( - torch.zeros(self.corrected_tensor_sizes[i], device=self.exp_avg_flat[i].device)) - self.server_errors.append( - torch.zeros(self.server_chunk_sizes[i], device=self.exp_avg_flat[i].device)) - get_accelerator().empty_cache() - if dist.get_rank() == 0: - print("Cupy Buffers Initialized Successfully.") - - self.comm_backend_handle.compressed_allreduce(self.exp_avg_flat[i], self.worker_errors[0], - self.server_errors[0], self.deepspeed.local_rank) - - if dist.get_rank() == 0: - print('Pop out errors', flush=True) - del self.worker_errors[:] - del self.server_errors[:] - else: - self.comm_backend_handle.compressed_allreduce(self.exp_avg_flat[i], self.worker_errors[i], - self.server_errors[i], self.deepspeed.local_rank) - - if self.lamb_freeze_key and self.initialize: - for i, group in enumerate(self.param_groups): - bias_correction = 1 if group['bias_correction'] else 0 - - for j, p in enumerate(group['params']): - state = self.state[p] - exp_avg, exp_avg_sq, exp_avg_sq_fresh = state['exp_avg'], state['exp_avg_sq'], state[ - 'exp_avg_sq_fresh'] - beta1, beta2 = group['betas'] - exp_avg.div_(self.state[p]['scaling_coeff']) - # Because 1-bit compression cannot represent exact zero, it is required to - # provide a momentum mask for those params that have constant exact zeros in their - # momentums, otherwise the compression error would keep accumulating. - # For example, for BERT pre-training seq 128, bert.embeddings.position_embeddings.weight - # always have exact zeros in its momentum for row 129 to 512, because it only - # learns up to seq length 128 while the model supports up to 512 seq length. - # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py about how - # to add this exp_avg_mask for BERT pre-training.) - if 'exp_avg_mask' in group: - if exp_avg.device != group['exp_avg_mask'].device: - group['exp_avg_mask'] = group['exp_avg_mask'].to(device=exp_avg.device) - exp_avg.mul_(group['exp_avg_mask']) - - grad_reconstruct = ((exp_avg - exp_avg_last_step[i][j] * beta1) / (1 - beta1)) - exp_avg_sq_fresh.mul_(beta2).addcmul_(1 - beta2, grad_reconstruct, grad_reconstruct) - denom = exp_avg_sq.sqrt() + group['eps'] - update_prelim = exp_avg / denom - - if group['weight_decay'] > 0.0: - update = update_prelim + group['weight_decay'] * p.data - else: - update = update_prelim - - lamb_coeff = 1.0 - update_norm = update.pow(2).sum().sqrt() - denom_real = exp_avg_sq_fresh.sqrt() + group['eps'] - factor = (denom / denom_real).max().item() - if group['weight_decay'] > 0.0: - update_ratio = min(1.0, (update_prelim.pow(2).sum().sqrt() / update_norm).item()) - factor = factor * update_ratio + (1.0 - update_ratio) - if factor > self.factor_max: - factor = self.factor_max - if factor < self.factor_min: - factor = self.factor_min - if factor > state['last_factor'] * (1.0 + self.factor_threshold): - factor = state['last_factor'] * (1.0 + self.factor_threshold) - if factor < state['last_factor'] * (1.0 - self.factor_threshold): - factor = state['last_factor'] * (1.0 - self.factor_threshold) - state['last_factor'] = factor - lamb_coeff = state['lamb_coeff_freeze'] * factor - self.lamb_coeffs.append(lamb_coeff) - with torch.no_grad(): - p.add_(-group['lr'] * lamb_coeff * update) - del exp_avg_last_step[:] - exp_avg_last_step = None - - if not self.initialize: - self.lamb_freeze_key = False - self.initialize = True - print(f"Finished the initialization step at rank {dist.get_rank()}") - return loss - - if self.lamb_freeze_key is False: - if state['step'] >= self.freeze_step: - print('OnebitLamb - starting compressed communication') - self.lamb_freeze_key = True - if self.using_pipeline: - self.deepspeed.pipeline_enable_backward_allreduce = False - else: - self.deepspeed.enable_backward_allreduce = False - - return loss - - def load_state_dict(self, state_dict): - """ - Overrides load_state_dict() to add special handling when loading checkpoints - """ - # Because at different stage exp_avg_mask may change (e.g., - # BERT pre-training seqlen 128 and 512 ), we don't use the exp_avg_mask - # in checkpoints but always use the one user provided in training script. - # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py.) - # Thus here we keep the exp_avg_mask unchanged when loading checkpoint - for i, group in enumerate(self.param_groups): - if 'exp_avg_mask' in group: - state_dict['param_groups'][i]['exp_avg_mask'] = group['exp_avg_mask'] - elif 'exp_avg_mask' not in group and 'exp_avg_mask' in state_dict['param_groups'][i]: - state_dict['param_groups'][i].pop('exp_avg_mask') - super().load_state_dict(state_dict) - # need to reset the fused momentum since loading states will break the linking - del self.exp_avg_flat[:] - self.dummy_exp_avg.clear() - del self.corrected_tensor_sizes[:] - del self.server_chunk_sizes[:] - if self.state[self.param_groups[0]['params'][0]]['step'] < self.freeze_step: - if dist.get_rank() == 0: - print("Checkpoint loaded and OnebitLamb warmup stage starts/continues.") - if self.lamb_freeze_key is True: - self.lamb_freeze_key = False - if self.using_pipeline: - self.deepspeed.pipeline_enable_backward_allreduce = True - else: - self.deepspeed.enable_backward_allreduce = True - for group in self.param_groups: - for p in group['params']: - self.state[p]['lamb_coeff_freeze'] = 0.0 - self.state[p]['last_factor'] = 1.0 - if 'scaling_coeff' in self.state[p]: - self.state[p].pop('scaling_coeff') - else: - if dist.get_rank() == 0: - print("Checkpoint loaded and OnebitLamb compression stage starts/continues.") - if self.lamb_freeze_key is False: - self.lamb_freeze_key = True - if self.using_pipeline: - self.deepspeed.pipeline_enable_backward_allreduce = False - else: - self.deepspeed.enable_backward_allreduce = False - # We reset the compression errors when loading checkpoints for 3 reasons: - # 1) The worker and server error at each GPU are distinct, so in current implementation - # only rank 0's errors are saved in the checkpoint. Thus we have to reset the errors. - # If we want to save them correctly we need O(num_gpu*model_size) memory in order to - # gather all the error, which is a very large memory requirement. It's possible to save - # them in a distributed way, but it will make the checkpoint saving/loading much more complicated. - # 2) Even if we are able to save the compression errors correctly, you need to have the - # exact same number of GPUs in order to load them correctly. - # 3) We verified on BERT pre-training that occasionally resetting the compression error - # at checkpoint loading does not affect the convergence. - # However, please avoid frequent checkpoint loading which could break the error - # compensation mechanism thus affect the convergence. - del self.worker_errors[:] - del self.server_errors[:] - - def get_lamb_coeffs(self): - return self.lamb_coeffs diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/zoadam.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/zoadam.py deleted file mode 100644 index 70282ec41714b202a7105925d4b051cd7571ecea..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/onebit/zoadam.py +++ /dev/null @@ -1,365 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import types -import torch -import numpy as np -from deepspeed.accelerator import get_accelerator -from deepspeed.utils.torch import required_torch_version -from deepspeed import comm as dist - - -class ZeroOneAdam(torch.optim.Optimizer): - """ - Implements the 0/1 Adam algorithm. Currently GPU-only. - For usage example please see https://www.deepspeed.ai/tutorials/zero-one-adam/ - For technical details please read https://arxiv.org/abs/2202.06009 - - Arguments: - params (iterable): iterable of parameters to optimize or dicts defining - parameter groups. - lr (float, optional): learning rate. (default: 1e-3) - betas (Tuple[float, float], optional): coefficients used for computing - running averages of gradient and its square. (default: (0.9, 0.999)) - eps (float, optional): term added to the denominator to improve - numerical stability. (default: 1e-8) - weight_decay (float, optional): weight decay (L2 penalty) (default: 0) - var_freeze_step (int, optional): The latest step to update the variance, - using the notation from https://arxiv.org/abs/2202.06009, it denotes the - max{i|i in T_v}. Note that this is different from the freeze step from the - 1-bit Adam. The var_freeze_step is usually the end of the learning rate warmup - and thus does not require tuning. (default: 100000) - var_update_scaler (int, optional): The interval to update the variance. Note that - the update policy for variance follows an exponential rule, where var_update_scaler - denotes the kappa in the 0/1 Adam paper. (default: 16) - local_step_scaler (int, optional): The interval to scale the local steps interval - according to the learning rate policy. (default: 32678) - local_step_clipper (int, optional): The largest interval for local steps with - learning rate policy. This corresponds to the variable H in the 0/1 Adam paper. - (default: 16) - amsgrad (boolean, optional): whether to use the AMSGrad variant of this - algorithm from the paper `On the Convergence of Adam and Beyond`_ - (default: False) NOT SUPPORTED in 0/1 Adam! - eps_inside_sqrt (boolean, optional): in the 'update parameters' step, - adds eps to the bias-corrected second moment estimate before - evaluating square root instead of adding it to the square root of - second moment estimate as in the original paper. (default: False) - cuda_aware (boolean, required): Set True if the underlying MPI implementation - supports CUDA-Aware communication. (default: False) - comm_backend_name (string, optional): Set to 'mpi' if needed. (default: 'nccl') - .. _Adam\\: A Method for Stochastic Optimization: - https://arxiv.org/abs/1412.6980 - .. _On the Convergence of Adam and Beyond: - https://openreview.net/forum?id=ryQu7f-RZ - """ - - def __init__(self, - params, - deepspeed=None, - lr=1e-3, - bias_correction=True, - betas=(0.9, 0.999), - eps=1e-8, - eps_inside_sqrt=False, - weight_decay=0., - max_grad_norm=0., - var_freeze_step=100000, - var_update_scaler=16, - local_step_scaler=32678, - local_step_clipper=16, - amsgrad=False, - cuda_aware=False, - comm_backend_name='nccl'): - - if amsgrad: - raise RuntimeError('0/1 Adam does not support the AMSGrad variant.') - - defaults = dict(lr=lr, - bias_correction=bias_correction, - betas=betas, - eps=eps, - weight_decay=weight_decay, - max_grad_norm=max_grad_norm) - - super(ZeroOneAdam, self).__init__(params, defaults) - self.eps_mode = 0 if eps_inside_sqrt else 1 - self.deepspeed = deepspeed - self.initialize = False - self.cuda_aware = cuda_aware - self.using_pipeline = False - - self.var_freeze_step = var_freeze_step - self.var_update_scaler = var_update_scaler - self.local_step_scaler = local_step_scaler - self.local_step_clipper = local_step_clipper - self.freeze_key = False - self.reinitial_error_buffer = False - - self.comm_backend_name = comm_backend_name - - assert dist.is_initialized(), "Please initialize the torch distributed backend." - # Empty initializer. Set handle based on the comm backend as follows. - self.comm_backend_handle = None - if self.comm_backend_name == 'nccl': - assert ( - required_torch_version(min_version=1.8) - ), "Please use torch 1.8 or greater to enable NCCL backend in 0/1 Adam. Alternatively, please specify 'mpi' as the 'comm_backend_name' in config file to proceed with the MPI backend" - from deepspeed.runtime.comm.nccl import NcclBackend - self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') - self.comm_backend_handle = NcclBackend(self.deepspeed.mpu) - elif self.comm_backend_name == 'mpi': - from deepspeed.runtime.comm.mpi import MpiBackend - self.comm_backend_handle = MpiBackend(cuda_aware) - elif self.comm_backend_name == 'hccl': - from deepspeed.runtime.comm.hccl import HcclBackend - self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') - self.comm_backend_handle = HcclBackend(self.deepspeed.mpu) - elif self.comm_backend_name == 'compressed': - from deepspeed.runtime.comm.compressed import CompressedBackend - self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') - self.comm_backend_handle = CompressedBackend(self.deepspeed.mpu) - self.size = self.comm_backend_handle.size - - self.divider = int(self.size * 8 / np.gcd(self.size, 8)) - - def step(self, closure=None, grads=None): - """Performs a single optimization step. - Arguments: - closure (callable, optional): A closure that reevaluates the model - and returns the loss. - grads (list of tensors, optional): weight gradient to use for the - optimizer update. If gradients have type torch.half, parameters - are expected to be in type torch.float. (default: None) - output params (list of tensors, optional): A reduced precision copy - of the updated weights written out in addition to the regular - updated weights. Have to be of same type as gradients. (default: None) - scale (float, optional): factor to divide gradient tensor values - by before applying to weights. (default: 1) - """ - loss = None - if closure is not None: - loss = closure() - - if grads is None: - grads_group = [None] * len(self.param_groups) - # backward compatibility - # assuming a list/generator of parameter means single group - elif isinstance(grads, types.GeneratorType): - grads_group = [grads] - elif type(grads[0]) != list: - grads_group = [grads] - else: - grads_group = grads - - for group, grads_this_group in zip(self.param_groups, grads_group): - if grads_this_group is None: - grads_this_group = [None] * len(group['params']) - - bias_correction = 1 if group['bias_correction'] else 0 - - for p, grad in zip(group['params'], grads_this_group): - if p.grad is None and grad is None: - continue - if grad is None: - grad = p.grad.data - if grad.is_sparse: - raise RuntimeError('0/1 Adam does not support sparse gradients') - - state = self.state[p] - - # State initialization - if len(state) == 0: - state['step'] = 0 - # Exponential moving average of gradient values - state['exp_avg'] = torch.zeros_like(p.data) - # Exponential moving average of squared gradient values - state['exp_avg_sq'] = torch.zeros_like(p.data) - - if not self.initialize or 'worker_error' not in state.keys(): - # Some scalars to help scale the variance update/local step policies - state['var_interval'] = 1 - state['var_counter'] = 0 - state['local_step_interval'] = 1 - state['local_step_counter'] = 0 - state['lrs'] = 0 - state['tensor_size'] = torch.numel(p.data) - state['corrected_tensor_size'] = state['tensor_size'] - - if state['tensor_size'] % (self.size * self.divider) != 0: - state['corrected_tensor_size'] += ((self.size * self.divider) - (state['tensor_size'] % - (self.size * self.divider))) - state['server_chunk_size'] = state['corrected_tensor_size'] // self.size - get_accelerator().empty_cache() - state['worker_error'] = torch.zeros(state['corrected_tensor_size'], device=p.device) - state['server_error'] = torch.zeros(state['server_chunk_size'], device=p.device) - # Accumulation of momentum, i.e., the u variable in the 0/1 Adam paper - state['momentum_accumulator'] = torch.zeros_like(p.data) - get_accelerator().empty_cache() - # self.freeze_key = True - if not self.initialize and dist.get_rank() == 0: - print("Cupy Buffers Initialized Successfully.") - - exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] - comm_buffer = state['momentum_accumulator'] - beta1, beta2 = group['betas'] - - state['step'] += 1 - - if self.initialize: - if self.freeze_key is False: - if state['step'] % state['var_interval'] == 0: - exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) - exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) - else: - if self.size > 1: - with torch.no_grad(): - grad_onebit = self.comm_backend_handle.compressed_allreduce( - grad, state['worker_error'], state['server_error'], self.deepspeed.local_rank) - if 'exp_avg_mask' in group: - if grad_onebit.device != group['exp_avg_mask'].device: - group['exp_avg_mask'] = group['exp_avg_mask'].to(device=grad_onebit.device) - grad_onebit.mul_(group['exp_avg_mask']) - exp_avg.mul_(beta1).add_(1 - beta1, grad_onebit) - else: - exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) - state['lrs'] += group['lr'] - grad = None - - if not self.initialize: - if self.size > 1: - comm_buffer.set_( - self.comm_backend_handle.compressed_allreduce(comm_buffer, state['worker_error'], - state['server_error'], - self.deepspeed.local_rank)) - if 'exp_avg_mask' in group: - if comm_buffer.device != group['exp_avg_mask'].device: - group['exp_avg_mask'] = group['exp_avg_mask'].to(device=comm_buffer.device) - comm_buffer.mul_(group['exp_avg_mask']) - - if self.initialize: - update = exp_avg / (exp_avg_sq.sqrt() + group['eps']) - if group['weight_decay'] > 0.0: - update += group['weight_decay'] * p.data - with torch.no_grad(): - p.data.add_(-group['lr'] * update) - if self.freeze_key is True: - comm_buffer.add_(-group['lr'] * update) - if state['step'] % state['local_step_interval'] == 0 and self.freeze_key: - with torch.no_grad(): - p.data.add_(-1 * comm_buffer) - comm_buffer.mul_(exp_avg_sq.sqrt() + group['eps']) - if self.size > 1: - comm_buffer.copy_( - self.comm_backend_handle.compressed_allreduce(comm_buffer, state['worker_error'], - state['server_error'], - self.deepspeed.local_rank)) - if 'exp_avg_mask' in group: - if comm_buffer.device != group['exp_avg_mask'].device: - group['exp_avg_mask'] = group['exp_avg_mask'].to(device=comm_buffer.device) - comm_buffer.mul_(group['exp_avg_mask']) - exp_avg.zero_().add_(comm_buffer / state['lrs'], alpha=-1) - p.data.add_(comm_buffer / (exp_avg_sq.sqrt() + group['eps'])) - comm_buffer.zero_() - - state['lrs'] = 0 - - # According to 0/1 Adam theory, a fixed variance would allow more accurate estimation of momentum - # However, in practice, we can also disable the manual freezing of variance, since the interval of - # updating variance will increase exponentially, so that it has negligible effect on the estimation. - if self.freeze_key is False: - if state['step'] % state['var_interval'] == 0: - state['var_counter'] += 1 - if state['var_counter'] == self.var_update_scaler: - state['var_counter'] = 0 - state['var_interval'] *= 2 - if (state['step'] + 1) % state['var_interval'] == 0: - if self.using_pipeline: - self.deepspeed.pipeline_enable_backward_allreduce = True - else: - self.deepspeed.enable_backward_allreduce = True - else: - if self.using_pipeline: - self.deepspeed.pipeline_enable_backward_allreduce = False - else: - self.deepspeed.enable_backward_allreduce = False - else: - state['local_step_counter'] += 1 - if state['local_step_counter'] == self.local_step_scaler: - state['local_step_counter'] = 0 - state['local_step_interval'] = min(self.local_step_clipper, - state['local_step_interval'] * 2) - - if not self.initialize: - print('Pop out errors', flush=True) - self.freeze_key = False - state.pop('worker_error') - state.pop('server_error') - - if not self.initialize: - self.initialize = True - print(f"Finished the initialization step at rank {dist.get_rank()}") - return loss - - if self.state[self.param_groups[0]['params'][0]]['step'] > self.var_freeze_step: - self.freeze_key = True - if self.using_pipeline: - self.deepspeed.pipeline_enable_backward_allreduce = False - else: - self.deepspeed.enable_backward_allreduce = False - - if self.freeze_key is True and self.reinitial_error_buffer is False: - # We need to reinitialize the error buffers when local step > 1 since - # the errors will be logged for different metrics (gradient vs. accumulated momentum). - for group in self.param_groups: - for p in group['params']: - self.state[p]['worker_error'].zero_() - self.state[p]['server_error'].zero_() - self.reinitial_error_buffer = True - - return loss - - def load_state_dict(self, state_dict): - """ - Overrides load_state_dict() to add special handling when loading checkpoints - """ - # Because at different stage exp_avg_mask may change (e.g., - # BERT pre-training seqlen 128 and 512 ), we don't use the exp_avg_mask - # in checkpoints but always use the one user provided in training script. - # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py.) - # Thus here we keep the exp_avg_mask unchanged when loading checkpoint - for i, group in enumerate(self.param_groups): - if 'exp_avg_mask' in group: - state_dict['param_groups'][i]['exp_avg_mask'] = group['exp_avg_mask'] - elif 'exp_avg_mask' not in group and 'exp_avg_mask' in state_dict['param_groups'][i]: - state_dict['param_groups'][i].pop('exp_avg_mask') - super().load_state_dict(state_dict) - if self.state[self.param_groups[0]['params'][0]]['step'] < self.var_freeze_step: - self.var_freeze_key = False - if (self.state[self.param_groups[0]['params'][0]]['step'] + - 1) % self.state[self.param_groups[0]['params'][0]]['var_interval'] == 0: - if self.using_pipeline: - self.deepspeed.pipeline_enable_backward_allreduce = True - else: - self.deepspeed.enable_backward_allreduce = True - else: - if self.using_pipeline: - self.deepspeed.pipeline_enable_backward_allreduce = False - else: - self.deepspeed.enable_backward_allreduce = False - else: - self.var_freeze_key = True - if self.using_pipeline: - self.deepspeed.pipeline_enable_backward_allreduce = False - else: - self.deepspeed.enable_backward_allreduce = False - self.reinitial_error_buffer = False - for group in self.param_groups: - for p in group['params']: - if 'worker_error' in self.state[p]: - self.state[p].pop('worker_error') - if 'server_error' in self.state[p]: - self.state[p].pop('server_error') - if 'momentum_accumulator' in self.state[p]: - self.state[p].pop('momentum_accumulator') diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/unfused_optimizer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/unfused_optimizer.py deleted file mode 100644 index 530355f846e2c04dfeaaf1886d79f67aa764908f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/fp16/unfused_optimizer.py +++ /dev/null @@ -1,428 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Copyright NVIDIA/apex -This file is adapted from FP16_Optimizer in NVIDIA/apex -""" - -from deepspeed.moe.utils import split_params_grads_into_shared_and_expert_params -import torch -from torch._utils import _flatten_dense_tensors - -from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer -from deepspeed.runtime.utils import get_global_norm, CheckOverflow, get_weight_norm -from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE -from deepspeed.utils import logger -from deepspeed.utils.torch import required_torch_version -from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT -from deepspeed.accelerator import get_accelerator -from deepspeed import comm as dist - - -class FP16_UnfusedOptimizer(DeepSpeedOptimizer): - """ - FP16 Optimizer without weight fusion to support LAMB optimizer - - For usage example please see, TODO: DeepSpeed V2 Tutorial - """ - - def __init__(self, - init_optimizer, - deepspeed=None, - static_loss_scale=1.0, - dynamic_loss_scale=False, - dynamic_loss_args=None, - verbose=True, - mpu=None, - clip_grad=0.0, - fused_lamb_legacy=False): - - self.fused_lamb_legacy = fused_lamb_legacy - self._global_grad_norm = 0. - - if dist.get_rank() == 0: - logger.info(f'Fused Lamb Legacy : {self.fused_lamb_legacy} ') - - if not get_accelerator().is_available(): - raise SystemError("Cannot use fp16 without accelerator.") - self.optimizer = init_optimizer - - # param groups - self.fp16_groups = [] - self.fp32_groups = [] - - # loop to deal with groups - for i, param_group in enumerate(self.optimizer.param_groups): - #fp16 weights that represents the actual model weights - self.fp16_groups.append(param_group['params']) - - #creating a fp32 copy of the weights that will be updated first then - #copied to fp16 weights - fp32_group = [p.clone().float().detach() for p in param_group['params']] - - #in case the internal optimizer needs it - for p in fp32_group: - p.requires_grad = True - - #setting the param groups in the optimizer to point to fp32 - #note these are not the weights used by the model - #the model uses the fp16 version that we added to fp16_group - self.fp32_groups.append(fp32_group) - param_group['params'] = self.fp32_groups[i] - - # we may have a way of fusing dynamic scale. Do not support for now - if dynamic_loss_scale: - self.dynamic_loss_scale = True - self.cur_iter = 0 - self.last_overflow_iter = -1 - self.scale_factor = 2.0 - if dynamic_loss_args is None: - self.cur_scale = 1.0 * 2**16 - self.scale_window = 1000 - self.min_loss_scale = 0.25 - else: - self.cur_scale = dynamic_loss_args[INITIAL_LOSS_SCALE] - self.scale_window = dynamic_loss_args[SCALE_WINDOW] - self.min_loss_scale = dynamic_loss_args[MIN_LOSS_SCALE] - else: - self.dynamic_loss_scale = False - self.cur_iter = 0 - self.cur_scale = static_loss_scale - - self.custom_loss_scaler = False - self.external_loss_scale = None - - self.verbose = verbose - - self.clip_grad = clip_grad - self.norm_type = 2 - - if required_torch_version(max_version=0.4): - self.clip_grad_norm = torch.nn.utils.clip_grad_norm - else: - self.clip_grad_norm = torch.nn.utils.clip_grad_norm_ - - self.mpu = mpu - - self.overflow = False - self.overflow_checker = CheckOverflow(self.fp16_groups, mpu=self.mpu, deepspeed=deepspeed) - - self.initialize_optimizer_states() - - def zero_grad(self, set_to_none=True): - """ - Zero FP16 parameter grads. - """ - # FP32 grad should never exist outside of the step function - # For speed, set model fp16 grad to None by default - for group in self.fp16_groups: - for p in group: - if set_to_none: - p.grad = None - else: - if p.grad is not None: - p.grad.detach_() - p.grad.zero_() - - def step_fused_lamb(self, closure=None): - """ - Not supporting closure. - """ - # First compute norm for all group so we know if there is overflow - grads_groups_flat = [] - grads_groups = [] - norm_groups = [] - expert_norm_groups = [] - for i, group in enumerate(self.fp16_groups): - grads = [ - torch.zeros(p.size(), dtype=p.dtype, device=p.device) if p.grad is None else p.grad for p in group - ] - grads_groups.append(grads) - grads_groups_flat.append(_flatten_dense_tensors(grads)) - grads_for_norm, expert_grads_for_norm = split_params_grads_into_shared_and_expert_params(group) - norm_group_value = 0.0 - if len(grads_for_norm) > 0: - norm_group_value = get_weight_norm(_flatten_dense_tensors(grads_for_norm), mpu=self.mpu) - norm_groups.append(norm_group_value) - expert_norm_group_value = 0.0 - if len(expert_grads_for_norm) > 0: - expert_norm_group_value = get_weight_norm(_flatten_dense_tensors(expert_grads_for_norm), mpu=self.mpu) - expert_norm_groups.append(expert_norm_group_value) - - self.overflow = self.overflow_checker.check_using_norm(norm_groups + expert_norm_groups) - prev_scale = self.cur_scale - - self._update_scale(self.overflow) - if self.overflow: - if self.verbose: - logger.info("[deepspeed] fp16 dynamic loss scale overflow! Skipping step. Attempted loss " - "scale: {}, reducing to {}".format(prev_scale, self.cur_scale)) - return self.overflow - - self._global_grad_norm = get_global_norm(norm_list=norm_groups) - combined_scale = self.unscale_and_clip_grads(self._global_grad_norm, apply_scale=False) - self.optimizer.step(grads=grads_groups, output_params=self.fp16_groups, scale=combined_scale) - - for fp32_group, fp16_group in zip(self.fp32_groups, self.fp16_groups): - for idx, (fp32_param, fp16_param) in enumerate(zip(fp32_group, fp16_group)): - - #remove the fp32 grad - fp32_param.grad = None - - #copy data from fp32 to fp16 - fp16_param.data.copy_(fp32_param.data) - - return self.overflow - - def set_lr(self, lr): - """Set the learning rate.""" - for param_group in self.optimizer.param_groups: - param_group["lr"] = lr - - def get_lr(self): - """Return the current learning rate.""" - return self.optimizer.param_groups[0]["lr"] - - def override_loss_scale(self, loss_scale): - if loss_scale != self.external_loss_scale: - logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}') - self.custom_loss_scaler = True - self.external_loss_scale = loss_scale - - def step(self, closure=None): - """ - Not supporting closure. - """ - - if self.fused_lamb_legacy: - return self.step_fused_lamb() - - self.overflow = self.overflow_checker.check() - prev_scale = self.cur_scale - - self._update_scale(self.overflow) - if self.overflow: - if self.verbose: - logger.info("[deepspeed] fp16 dynamic loss scale overflow! Skipping step. Attempted loss " - "scale: {}, reducing to {}".format(prev_scale, self.cur_scale)) - return self.overflow - - norm_groups = [] - for i, group in enumerate(self.fp16_groups): - grads_for_norm, _ = split_params_grads_into_shared_and_expert_params(group) - norm_group_value = 0.0 - if len(grads_for_norm) > 0: - norm_group_value = get_weight_norm(grads_for_norm, mpu=self.mpu) - norm_groups.append(norm_group_value) - - # copying gradients to fp32 to work with fp32 parameters - for fp32_param, fp16_param in zip(self.fp32_groups[i], self.fp16_groups[i]): - if fp16_param.grad is None: - fp32_param.grad = torch.zeros(fp16_param.size(), dtype=fp32_param.dtype, device=fp32_param.device) - else: - fp32_param.grad = fp16_param.grad.to(fp32_param.dtype) - - self._global_grad_norm = get_global_norm(norm_list=norm_groups) - self.unscale_and_clip_grads(self._global_grad_norm) - - self.optimizer.step() - - for fp32_group, fp16_group in zip(self.fp32_groups, self.fp16_groups): - for idx, (fp32_param, fp16_param) in enumerate(zip(fp32_group, fp16_group)): - - #remove the fp32 grad - fp32_param.grad = None - - #copy data from fp32 to fp16 - fp16_param.data.copy_(fp32_param.data) - - return self.overflow - - def unscale_and_clip_grads(self, total_norm, apply_scale=True): - # compute combined scale factor for this group - combined_scale = self.cur_scale - if self.clip_grad > 0.: - # norm is in fact norm*scale - clip = ((total_norm / self.cur_scale) + 1e-6) / self.clip_grad - if clip > 1: - combined_scale = clip * self.cur_scale - - if apply_scale: - for group in self.fp32_groups: - for param in group: - if param.grad is not None: - param.grad.data.mul_(1. / combined_scale) - - return combined_scale - - def backward(self, loss, create_graph=False, retain_graph=False): - """ - :attr:`backward` performs the following steps: - - 1. fp32_loss = loss.float() - 2. scaled_loss = fp32_loss*loss_scale - 3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves - """ - if self.custom_loss_scaler: - scaled_loss = self.external_loss_scale * loss - scaled_loss.backward() - else: - scaled_loss = (loss.float()) * self.cur_scale - scaled_loss.backward(create_graph=create_graph, retain_graph=retain_graph) - - def _update_scale(self, skip): - if self.dynamic_loss_scale: - prev_scale = self.cur_scale - if skip: - self.cur_scale = max(self.cur_scale / self.scale_factor, self.min_loss_scale) - self.last_overflow_iter = self.cur_iter - if self.verbose: - logger.info("Grad overflow on iteration: %s", self.cur_iter) - logger.info(f"Reducing dynamic loss scale from {prev_scale} to {self.cur_scale}") - else: - # Ensure self.scale_window updates since last overflow - stable_interval = (self.cur_iter - self.last_overflow_iter) - 1 - if (stable_interval > 0) and (stable_interval % self.scale_window == 0): - self.cur_scale *= self.scale_factor - if self.verbose: - logger.info(f"No Grad overflow for {self.scale_window} iterations") - logger.info(f"Increasing dynamic loss scale from {prev_scale} to {self.cur_scale}") - else: - if skip: - logger.info("Grad overflow on iteration %s", self.cur_iter) - logger.info("Using static loss scale of %s", self.cur_scale) - self.cur_iter += 1 - return - - # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state" - def _get_state(self): - return self.optimizer.state - - def _set_state(self, value): - self.optimizer.state = value - - state = property(_get_state, _set_state) - - # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups" - # (for example, to adjust the learning rate) - def _get_param_groups(self): - return self.optimizer.param_groups - - def _set_param_groups(self, value): - self.optimizer.param_groups = value - - param_groups = property(_get_param_groups, _set_param_groups) - - # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale" - def _get_loss_scale(self): - if self.custom_loss_scaler: - return self.external_loss_scale - else: - return self.cur_scale - - def _set_loss_scale(self, value): - self.loss_scaler.cur_scale = value - - loss_scale = property(_get_loss_scale, _set_loss_scale) - - def state_dict(self): - """ - Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. - This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict - of the contained Pytorch optimizer. - Example:: - checkpoint = {} - checkpoint['model'] = model.state_dict() - checkpoint['optimizer'] = optimizer.state_dict() - torch.save(checkpoint, "saved.pth") - """ - state_dict = {} - state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale - state_dict['cur_scale'] = self.cur_scale - state_dict['cur_iter'] = self.cur_iter - if state_dict['dynamic_loss_scale']: - state_dict['last_overflow_iter'] = self.last_overflow_iter - state_dict['scale_factor'] = self.scale_factor - state_dict['scale_window'] = self.scale_window - state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict() - state_dict['fp32_groups'] = self.fp32_groups - return state_dict - - # Refresh fp32 master params from fp16 copies - def refresh_fp32_params(self): - for current_group, saved_group in zip(self.fp32_groups, self.fp16_groups): - for current, saved in zip(current_group, saved_group): - current.data.copy_(saved.data) - - def load_state_dict(self, state_dict, load_optimizer_states=True): - """ - Loads a state_dict created by an earlier call to state_dict(). - If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``, - whose parameters in turn came from ``model``, it is expected that the user - will call ``model.load_state_dict()`` before - ``fp16_optimizer_instance.load_state_dict()`` is called. - Example:: - model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half() - optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) - optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0) - ... - checkpoint = torch.load("saved.pth") - model.load_state_dict(checkpoint['model']) - optimizer.load_state_dict(checkpoint['optimizer']) - """ - # I think it should actually be ok to reload the optimizer before the model. - self.dynamic_loss_scale = state_dict['dynamic_loss_scale'] - self.cur_scale = state_dict['cur_scale'] - self.cur_iter = state_dict['cur_iter'] - if state_dict['dynamic_loss_scale']: - self.last_overflow_iter = state_dict['last_overflow_iter'] - self.scale_factor = state_dict['scale_factor'] - self.scale_window = state_dict['scale_window'] - - if load_optimizer_states: - self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT]) - # At this point, the optimizer's references to the model's fp32 parameters are up to date. - # The optimizer's hyperparameters and internal buffers are also up to date. - # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still - # out of date. There are two options. - # 1: Refresh the master params from the model's fp16 params. - # This requires less storage but incurs precision loss. - # 2: Save and restore the fp32 master copies separately. - # We choose option 2. - # - # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device - # of their associated parameters, because it's possible those buffers might not exist yet in - # the current optimizer instance. In our case, as long as the current FP16_Optimizer has been - # constructed in the same way as the one whose state_dict we are loading, the same master params - # are guaranteed to exist, so we can just copy_() from the saved master params. - for current_group, saved_group in zip(self.fp32_groups, state_dict['fp32_groups']): - for current, saved in zip(current_group, saved_group): - current.data.copy_(saved.data) - - def __repr__(self): - return repr(self.optimizer) - - def initialize_optimizer_states(self): - for i, group in enumerate(self.fp16_groups): - for param in group: - param.grad = torch.zeros(param.size(), - dtype=param.dtype, - device=get_accelerator().current_device_name()) - - for i, group in enumerate(self.fp32_groups): - for param in group: - param.grad = torch.zeros(param.size(), - dtype=param.dtype, - device=get_accelerator().current_device_name()) - - self.optimizer.step() - - for i, group in enumerate(self.fp16_groups): - for param in group: - param.grad = None - - for i, group in enumerate(self.fp32_groups): - for param in group: - param.grad = None diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/hybrid_engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/hybrid_engine.py deleted file mode 100644 index b96228c90b02315fd63e547dbc28962009affd7c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/hybrid_engine.py +++ /dev/null @@ -1,445 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed.inference.config import DeepSpeedInferenceConfig -from deepspeed.module_inject.replace_policy import replace_policies -from deepspeed.module_inject.utils import policy_to_ds_container -from .engine import DeepSpeedEngine -from .utils import TLinear, get_inactive_params -from deepspeed.runtime.zero import GatheredParameters -import time -import gc -import math -from deepspeed import comm as dist -from deepspeed.accelerator import get_accelerator -from torch import nn -from deepspeed.utils import logger -from deepspeed.module_inject.layers import LinearLayer, Normalize, EmbeddingLayer, OPTEmbedding -from ..ops.transformer.inference.op_binding.workspace import WorkspaceOp - -try: - import transformers - OPTLearnedPositionalEmbedding = transformers.models.opt.modeling_opt.OPTLearnedPositionalEmbedding -except: - OPTLearnedPositionalEmbedding = None - - -class DeepSpeedHybridEngine(DeepSpeedEngine): - r"""DeepSpeed engine for training and inference.""" - inference_mp_group = None - - def __init__(self, args, model, **kwargs): - - super().__init__(args, model, **kwargs) - - # synch seed between all GPUs - _rng_state = get_accelerator().get_rng_state().to(get_accelerator().current_device_name()) - dist.broadcast(_rng_state, 0) - get_accelerator().set_rng_state(_rng_state.cpu()) - - self.Z3_enabled = (self._config.zero_config.stage == 3) - self.gather_all_layers = self._config.hybrid_engine.pin_parameters - - # inference containers / fwds - self._inference_containers = [] - self._orig_modules = [] - self._orig_fwds = [] - self.create_inference_module() - - # Performance stats - self._t_start = None - self._total_latency = 0 - self._iters = 0 - self._training_start_time = None - self._generate_latency = 0 - self._training_latency = 0 - self._total_batch_size = None - self._gather_latency = 0 - - self.is_lora_fused = False - self.workspace = WorkspaceOp() - - def convert_to_linear_transposed(self, model): - - def _replace_linear_layer(r_module, parent_type=None, prev_type=None): - for name, child in r_module.named_children(): - if child.__class__ in [torch.nn.Linear] and \ - (parent_type is torch.nn.ModuleList or prev_type is torch.nn.ModuleList): - setattr(r_module, name, TLinear(child, name)) - else: - _replace_linear_layer(child, type(r_module), prev_type=parent_type) - return r_module - - _replace_linear_layer(model) - - def new_inference_container(self, orig_layer, policy_cls, layer_id): - policy = policy_cls(orig_layer, inference=True) - - if self._config.float16_config.enabled: - inference_dtype = torch.float16 - elif self._config.bfloat16_config.enabled: - inference_dtype = torch.bfloat16 - else: - inference_dtype = torch.float32 - - _container = policy_to_ds_container( - policy=policy, - config=DeepSpeedInferenceConfig( - set_empty_params=True, - dtype=inference_dtype, - max_out_tokens=self._config.hybrid_engine.max_out_tokens, - min_out_tokens=self._config.hybrid_engine.max_out_tokens, - transposed_mode=True, - ), - model_config=self.module.config if hasattr(self.module, 'config') else None, - layer_id=layer_id, - child=orig_layer) - - if self.mpu is not None: - if hasattr(self.mpu, 'get_model_parallel_world_size'): - _container.set_tensor_parallel_config(self.mpu.get_model_parallel_world_size(), - self.mpu.get_model_parallel_group()) - else: - _container.set_tensor_parallel_config(self.mpu.get_tensor_model_parallel_world_size(), - self.mpu.get_tensor_model_parallel_group()) - else: - _container.set_tensor_parallel_config(self._config.hybrid_engine.inference_tp_size, self.mp_group) - _container.initialize_tensors(enable_training=True) - _container.create_ds_model_config() - _container.create_module() - _container.set_params_wo_copy(Z3_enabled=self.Z3_enabled) - return _container - - def populate_all_inference_policies(self): - self.inference_policies = {} - for plcy in replace_policies: - _ = plcy(None) - if isinstance(plcy._orig_layer_class, list): - for orig_layer_class in plcy._orig_layer_class: - self.inference_policies.update({orig_layer_class: (self.new_inference_container, plcy)}) - elif plcy._orig_layer_class is not None: - self.inference_policies.update({plcy._orig_layer_class: (self.new_inference_container, plcy)}) - self.inference_policies.update({ - nn.Linear: (LinearLayer, ), - nn.Embedding: (EmbeddingLayer, ), - nn.LayerNorm: (Normalize, ), - OPTLearnedPositionalEmbedding: (OPTEmbedding, ) - }) - - def _fuse_lora_layer(self, layer_id): - self._inference_containers[layer_id].fuse_lora() - - def fuse_lora_weight(self): - for layer_id in range(len(self.layer_params)): - self._fuse_lora_layer(layer_id) - - def _unfuse_lora_layer(self, layer_id): - self._inference_containers[layer_id].unfuse_lora() - - def unfuse_lora_weight(self): - for layer_id in range(len(self.layer_params)): - self._unfuse_lora_layer(layer_id) - - def unfuse_lora_weight_non_pinned(self): - for layer_id in range(len(self.layer_params)): - non_active_params = get_inactive_params(self.layer_params[layer_id]) - non_active_lora_params = get_inactive_params(self.layer_lora_params[layer_id]) - non_active_params.extend(non_active_lora_params) - - with GatheredParameters(non_active_params): - self._unfuse_lora_layer(layer_id) - - def retake_inference_cache(self): - if self._config.hybrid_engine.release_inference_cache: - retake_success = self.workspace.retake_workspace() - - if not retake_success: - logger.warning("Unable to acquire workspace on first attempt, emptying cache and retrying.") - gc.collect() - get_accelerator().empty_cache() - retake_success = self.workspace.retake_workspace() - - if not retake_success: - raise RuntimeError("Unable to retake inference workspace.") - - def generate(self, *inputs, **kwargs): - if self._total_batch_size is None: - bsz = inputs[0].shape[0] if len(inputs) > 0 else \ - kwargs['input_ids'].shape[0] - self._total_batch_size = bsz * dist.get_world_size() - - self._t0 = time.time() - - if self.Z3_enabled and self.gather_all_layers: - if self._config.hybrid_engine.inference_tp_size > 1: - non_tp_params = [] - for other_layer in self._other_layers: - non_tp_params.extend(list(other_layer.parameters())) - - partition_size = self._config.hybrid_engine.tp_gather_partition_size - - layer_groups = math.ceil(len(self.layer_params) / partition_size) - for lg in range(layer_groups): - non_active_params = [] - non_active_lora_params = [] - for layer_id in range(lg * partition_size, min(len(self.layer_params), (lg + 1) * partition_size), - 1): - non_tp_params.extend(self.layer_params[layer_id][:4]) - non_active_params.extend(get_inactive_params(self.layer_params[layer_id])) - non_active_params.extend(get_inactive_params(self.layer_lora_params[layer_id])) - with GatheredParameters(non_active_params): - for layer_id in range(lg * partition_size, - min(len(self.layer_params), (lg + 1) * partition_size), 1): - if len(self.all_lora_params) > 0: - self._fuse_lora_layer(layer_id) - - if self.mpu is not None: - self._inference_containers[layer_id].apply_tensor_parallelism(self.mp_replace, - reversed_dim=True) - - # TODO(cmikeh2) Evaluate if this can be deferred when release_inference_cache - # is enabled. - gc.collect() - get_accelerator().empty_cache() - - self._gather_latency = time.time() - self._t0 - - input_shape = inputs[0].shape if len(inputs) > 0 else \ - kwargs['input_ids'].shape - output = torch.zeros( - (input_shape[0] * self._config.hybrid_engine.inference_tp_size, ) + input_shape[1:], - dtype=inputs[0].dtype if len(inputs) > 0 else kwargs['input_ids'].dtype, - device=inputs[0].device if len(inputs) > 0 else kwargs['input_ids'].device) - input_cont = inputs[0].contiguous() if len(inputs) > 0 else kwargs['input_ids'].contiguous() - dist.all_gather_into_tensor(output, input_cont, group=self.mp_group) - - if len(inputs) > 0: - inputs = (output, *inputs[1:]) - else: - kwargs['input_ids'] = output - - self.retake_inference_cache() - - non_active_params = get_inactive_params(non_tp_params) - with GatheredParameters(non_active_params): - generate_ret_vals = self._generate(*inputs, **kwargs) - - for layer_id in range(len(self.layer_params)): - self._inference_containers[layer_id].release_memory() - - rank = dist.get_rank(group=self.mp_group) - generate_ret_vals = generate_ret_vals[input_shape[0] * rank:input_shape[0] * (rank + 1)] - - else: - non_active_layers = get_inactive_params(self.all_layers_params) - non_active_lora_params = get_inactive_params(self.all_lora_params) - non_active_layers.extend(non_active_lora_params) - with GatheredParameters(non_active_layers): - self._gather_latency = time.time() - self._t0 - - if len(self.all_lora_params) > 0: - self.fuse_lora_weight() - - self.retake_inference_cache() - generate_ret_vals = self._generate(*inputs, **kwargs) - - if len(self.all_lora_params) > 0: - self.unfuse_lora_weight() - else: - if len(self.all_lora_params) > 0 and (not self.Z3_enabled): - self.fuse_lora_weight() - - self.retake_inference_cache() - generate_ret_vals = self._generate(*inputs, **kwargs) - - if len(self.all_lora_params) > 0: - if (not self.Z3_enabled): - self.unfuse_lora_weight() - else: - self.unfuse_lora_weight_non_pinned() - self.is_lora_fused = False - - if self._config.hybrid_engine.release_inference_cache: - self.workspace.release_workspace() - gc.collect() - get_accelerator().empty_cache() - - self._generate_latency = time.time() - self._t0 - self._gather_latency - - return generate_ret_vals - - def create_inference_containers(self, module, layer_id=0): - for name, child in module.named_children(): - if child.__class__ in self.inference_policies: - if self.inference_policies[child.__class__][0] == self.new_inference_container: - self._inference_containers.append(self.inference_policies[child.__class__][0]( - child, self.inference_policies[child.__class__][-1], layer_id)) - self._orig_modules.append(child) - self._orig_fwds.append(child.forward) - - self.layer_params.append(self._inference_containers[layer_id].get_all_params()) - - self.lora_params.append(self._inference_containers[layer_id].get_lora_params()) - self.layer_lora_params.append([]) - for lora_param in self.lora_params[layer_id]: - self.layer_lora_params[layer_id].extend(lora_param[:-1]) - self.all_lora_params.extend(lora_param[:-1]) - - layer_id += 1 - else: - if self.inference_policies[child.__class__][0] == LinearLayer: - self._other_layers.append(self.inference_policies[child.__class__][0](module=child, - mp_group=None, - skip_partition=True)) - else: - self._other_layers.append(self.inference_policies[child.__class__][0]( - weight=child.weight, bias=child.bias if hasattr(child, 'bias') else None)) - self._orig_modules_others.append(child) - self._orig_fwds_others.append(child.forward) - else: - self.create_inference_containers(child, layer_id=layer_id) - - def create_inference_module(self): - self.layer_params = [] - self.layer_lora_params = [] - self.lora_params = [] - self.all_lora_params = [] - - self._other_layers = [] - self._orig_modules_others = [] - self._orig_fwds_others = [] - - if self._config.hybrid_engine.inference_tp_size > 1: - if self.mpu is None: - global_rank = dist.get_rank() - world_size = dist.get_world_size() - mp_group_id = global_rank // self._config.hybrid_engine.inference_tp_size - num_mp_groups = world_size // self._config.hybrid_engine.inference_tp_size - for mp_group_id in range(num_mp_groups): - ranks = list( - range(mp_group_id * self._config.hybrid_engine.inference_tp_size, \ - (mp_group_id + 1) * self._config.hybrid_engine.inference_tp_size, \ - 1) - ) - mp_group = dist.new_group(ranks) - if global_rank in ranks: - # mp_group is used for broader collective - self.mp_group = mp_group - - # mp_replace is used for container tensor slicing - from deepspeed.module_inject import ReplaceWithTensorSlicing - self.mp_replace = ReplaceWithTensorSlicing( - mp_group=self.mp_group, - mp_size=self._config.hybrid_engine.inference_tp_size, - out_dim=0, - in_dim=1) - - else: - self.mp_group = self.mpu.get_model_parallel_group() if hasattr(self.mpu, 'get_model_parallel_group') else \ - self.mpu.get_tensor_model_parallel_group() - - from deepspeed.module_inject import ReplaceWithTensorSlicing - self.mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group, - mp_size=self._config.hybrid_engine.inference_tp_size, - out_dim=0, - in_dim=1) - else: - self.mp_group = None - self.mp_replace = None - self.populate_all_inference_policies() - self.all_layers_params = list(self.module.parameters()) - self.create_inference_containers(self.module) - - if len(self._inference_containers) > 0: - self._generate = self.module.generate - self.module.generate = self.generate - - self._t0 = time.time() - - def _zero3_forward(self, layer_id): - - def run_forward(*inputs, **kwargs): - non_active_params = get_inactive_params(self.layer_params[layer_id]) - non_active_lora_params = get_inactive_params(self.layer_lora_params[layer_id]) - non_active_params.extend(non_active_lora_params) - - with GatheredParameters(non_active_params): - if len(self.all_lora_params) > 0: - # Use the is_lora_fused flag to prevent multiple fusion in Z3 with non-pinned memory - if not self.is_lora_fused: - self._fuse_lora_layer(layer_id) - # Set the is_lora_fused to true when reaching the last layer - if layer_id == len(self.layer_params) - 1: - self.is_lora_fused = True - return self._inference_containers[layer_id].module.forward(*inputs, **kwargs) - - return run_forward - - def eval(self): - if self._t_start is not None: - latency = time.time() - self._t_start - self._total_latency = self._total_latency + latency - self._iters = self._iters + 1 - if not dist.is_initialized() or dist.get_rank() == 0: - if self._total_batch_size is not None: - cur_samples_p_sec = f'|CurSamplesPerSec={(1 / latency * self._total_batch_size):.2f} ' - avg_samples_p_sec = f'|AvgSamplesPerSec={(1 / (self._total_latency / self._iters) * self._total_batch_size):.2f}' - else: - cur_samples_p_sec = '' - avg_samples_p_sec = '' - others = latency - (self._generate_latency + self._training_latency) - print(f'|E2E latency={(latency):.2f}s ' + \ - f'|Gather latency={self._gather_latency:.2f}s ({(self._gather_latency / latency * 100):.2f}%) ' - f'|Generate time={(self._generate_latency):.2f}s ({(self._generate_latency / latency * 100):.2f}%) ' + \ - f'|Training time={(self._training_latency):.2f}s ({(self._training_latency / latency * 100):.2f}%) ' + \ - f'|Others={others:.2f} ({(others / latency * 100):.2f}%)' + \ - cur_samples_p_sec + \ - avg_samples_p_sec) - self._t_start = time.time() - self._training_latency = 0 - super().eval() - if len(self._inference_containers) > 0: - for i, (orig_module, inference_container) in enumerate(zip(self._orig_modules, - self._inference_containers)): - if self.Z3_enabled and not self.gather_all_layers: - orig_module.forward = self._zero3_forward(i) - else: - orig_module.forward = inference_container.module.forward - - inference_container.transform_for_inference() - - if not self.Z3_enabled or self.gather_all_layers: - for orig_module, inference_layer in zip(self._orig_modules_others, self._other_layers): - orig_module.forward = inference_layer.forward - if self.Z3_enabled: - gc.collect() - get_accelerator().empty_cache() - if self._t_start is None: - self._t_start = time.time() - - def train(self, mode=True): - if mode and len(self._orig_modules) > 0: - for inference_container, orig_module, orig_fwd in zip(self._inference_containers, self._orig_modules, - self._orig_fwds): - inference_container.transform_for_training() - orig_module.forward = orig_fwd - for orig_module, orig_fwd in zip(self._orig_modules_others, self._orig_fwds_others): - orig_module.forward = orig_fwd - super().train(mode) - if mode: - self._training_start_time = time.time() - - def step(self, lr_kwargs=None): - super().step(lr_kwargs=lr_kwargs) - - if len(self._inference_containers) > 0: - if not self.Z3_enabled: - for inference_container in self._inference_containers: - inference_container.reset_params() - - if self._training_start_time is not None: - self._training_latency += (time.time() - self._training_start_time) - self._training_start_time = time.time() diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/lr_schedules.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/lr_schedules.py deleted file mode 100644 index bb0866e0d52d19bc6c5271007946ee6b1a6125f6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/lr_schedules.py +++ /dev/null @@ -1,881 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Implementation of learning rate schedules. - -Taken and modified from PyTorch v1.0.1 source -https://github.com/pytorch/pytorch/blob/v1.1.0/torch/optim/lr_scheduler.py -""" - -import argparse -from torch.optim import Optimizer -import math -from deepspeed.utils import logger - -LR_SCHEDULE = 'lr_schedule' -LR_RANGE_TEST = 'LRRangeTest' -ONE_CYCLE = 'OneCycle' -WARMUP_LR = 'WarmupLR' -WARMUP_DECAY_LR = 'WarmupDecayLR' -WARMUP_COSINE_LR = 'WarmupCosineLR' -VALID_LR_SCHEDULES = [LR_RANGE_TEST, ONE_CYCLE, WARMUP_LR, WARMUP_DECAY_LR, WARMUP_COSINE_LR] - -LR_RANGE_TEST_MIN_LR = 'lr_range_test_min_lr' -LR_RANGE_TEST_STEP_RATE = 'lr_range_test_step_rate' -LR_RANGE_TEST_STEP_SIZE = 'lr_range_test_step_size' -LR_RANGE_TEST_STAIRCASE = 'lr_range_test_staircase' - -EDGE_VALUE = 'edge_value' -MID_VALUE = 'mid_value' - -CYCLE_FIRST_STEP_SIZE = 'cycle_first_step_size' -CYCLE_FIRST_STAIR_COUNT = 'cycle_first_stair_count' -CYCLE_SECOND_STEP_SIZE = 'cycle_second_step_size' -CYCLE_SECOND_STAIR_COUNT = 'cycle_second_stair_count' -DECAY_STEP_SIZE = 'decay_step_size' - -CYCLE_MIN_LR = 'cycle_min_lr' -CYCLE_MAX_LR = 'cycle_max_lr' -DECAY_LR_RATE = 'decay_lr_rate' - -CYCLE_MIN_MOM = 'cycle_min_mom' -CYCLE_MAX_MOM = 'cycle_max_mom' -DECAY_MOM_RATE = 'decay_mom_rate' - -WARMUP_MIN_LR = 'warmup_min_lr' -WARMUP_MAX_LR = 'warmup_max_lr' -WARMUP_NUM_STEPS = 'warmup_num_steps' -WARMUP_TYPE = 'warmup_type' -WARMUP_LOG_RATE = 'log' -WARMUP_LINEAR_RATE = 'linear' - -WARMUP_MIN_RATIO = 'warmup_min_ratio' -COS_MIN_RATIO = 'cos_min_ratio' - -TOTAL_NUM_STEPS = 'total_num_steps' - - -def add_tuning_arguments(parser): - group = parser.add_argument_group('Convergence Tuning', 'Convergence tuning configurations') - - # LR scheduler - group.add_argument('--lr_schedule', type=str, default=None, help='LR schedule for training.') - - # Learning rate range test - group.add_argument("--lr_range_test_min_lr", type=float, default=0.001, help='Starting lr value.') - group.add_argument("--lr_range_test_step_rate", type=float, default=1.0, help='scaling rate for LR range test.') - group.add_argument("--lr_range_test_step_size", type=int, default=1000, help='training steps per LR change.') - group.add_argument("--lr_range_test_staircase", - type=bool, - default=False, - help='use staircase scaling for LR range test.') - - # OneCycle schedule - group.add_argument("--cycle_first_step_size", - type=int, - default=1000, - help='size of first step of 1Cycle schedule (training steps).') - group.add_argument("--cycle_first_stair_count", - type=int, - default=-1, - help='first stair count for 1Cycle schedule.') - group.add_argument("--cycle_second_step_size", - type=int, - default=-1, - help='size of second step of 1Cycle schedule (default first_step_size).') - group.add_argument("--cycle_second_stair_count", - type=int, - default=-1, - help='second stair count for 1Cycle schedule.') - group.add_argument("--decay_step_size", - type=int, - default=1000, - help='size of intervals for applying post cycle decay (training steps).') - - # 1Cycle LR - group.add_argument("--cycle_min_lr", type=float, default=0.01, help='1Cycle LR lower bound.') - group.add_argument("--cycle_max_lr", type=float, default=0.1, help='1Cycle LR upper bound.') - group.add_argument("--decay_lr_rate", type=float, default=0.0, help='post cycle LR decay rate.') - - # 1Cycle Momentum - group.add_argument('--cycle_momentum', default=False, action='store_true', help='Enable 1Cycle momentum schedule.') - group.add_argument("--cycle_min_mom", type=float, default=0.8, help='1Cycle momentum lower bound.') - group.add_argument("--cycle_max_mom", type=float, default=0.9, help='1Cycle momentum upper bound.') - group.add_argument("--decay_mom_rate", type=float, default=0.0, help='post cycle momentum decay rate.') - - # Warmup LR - group.add_argument('--warmup_min_lr', type=float, default=0, help='WarmupLR minimum/initial LR value') - group.add_argument('--warmup_max_lr', type=float, default=0.001, help='WarmupLR maximum LR value.') - group.add_argument('--warmup_num_steps', type=int, default=1000, help='WarmupLR step count for LR warmup.') - group.add_argument('--warmup_type', - type=str, - default=WARMUP_LOG_RATE, - help='WarmupLR increasing function during warmup') - - # WarmUP cos LR - group.add_argument("--warmup_min_ratio", type=float, default=0.01, help='Cosine LR lower bound.') - group.add_argument("--cos_min_ratio", type=float, default=0.01, help='Cosine LR lower bound.') - - return parser - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser = add_tuning_arguments(parser) - - lr_sched_args, unknown_args = parser.parse_known_args() - return lr_sched_args, unknown_args - - -def override_lr_range_test_params(args, params): - if hasattr(args, LR_RANGE_TEST_MIN_LR) and args.lr_range_test_min_lr is not None: - params[LR_RANGE_TEST_MIN_LR] = args.lr_range_test_min_lr - - if hasattr(args, LR_RANGE_TEST_STEP_RATE) and args.lr_range_test_step_rate is not None: - params[LR_RANGE_TEST_STEP_RATE] = args.lr_range_test_step_rate - - if hasattr(args, LR_RANGE_TEST_STEP_SIZE) and args.lr_range_test_step_size is not None: - params[LR_RANGE_TEST_STEP_SIZE] = args.lr_range_test_step_size - - if hasattr(args, LR_RANGE_TEST_STAIRCASE) and args.lr_range_test_staircase is not None: - params[LR_RANGE_TEST_STAIRCASE] = args.lr_range_test_staircase - - -def override_1cycle_params(args, params): - if hasattr(args, CYCLE_FIRST_STEP_SIZE) and args.cycle_first_step_size is not None: - params[CYCLE_FIRST_STEP_SIZE] = args.cycle_first_step_size - - if hasattr(args, CYCLE_FIRST_STAIR_COUNT) and args.cycle_first_stair_count is not None: - params[CYCLE_FIRST_STAIR_COUNT] = args.cycle_first_stair_count - - if hasattr(args, CYCLE_SECOND_STEP_SIZE) and args.cycle_second_step_size is not None: - params[CYCLE_SECOND_STEP_SIZE] = args.cycle_second_step_size - - if hasattr(args, CYCLE_SECOND_STAIR_COUNT) and args.cycle_second_stair_count is not None: - params[CYCLE_SECOND_STAIR_COUNT] = args.cycle_second_stair_count - - if hasattr(args, DECAY_STEP_SIZE) and args.decay_step_size is not None: - params[DECAY_STEP_SIZE] = args.decay_step_size - - # 1Cycle LR params - if hasattr(args, CYCLE_MIN_LR) and args.cycle_min_lr is not None: - params[CYCLE_MIN_LR] = args.cycle_min_lr - - if hasattr(args, CYCLE_MAX_LR) and args.cycle_max_lr is not None: - params[CYCLE_MAX_LR] = args.cycle_max_lr - - if hasattr(args, DECAY_LR_RATE) and args.decay_lr_rate is not None: - params[DECAY_LR_RATE] = args.decay_lr_rate - - # 1Cycle MOM params - if hasattr(args, CYCLE_MIN_MOM) and args.cycle_min_mom is not None: - params[CYCLE_MIN_MOM] = args.cycle_min_mom - - if hasattr(args, CYCLE_MAX_MOM) and args.cycle_max_mom is not None: - params[CYCLE_MAX_MOM] = args.cycle_max_mom - - if hasattr(args, DECAY_MOM_RATE) and args.decay_mom_rate is not None: - params[DECAY_MOM_RATE] = args.decay_mom_rate - - -def override_warmupLR_params(args, params): - if hasattr(args, WARMUP_MIN_LR) and args.warmup_min_lr is not None: - params[WARMUP_MIN_LR] = args.warmup_min_lr - - if hasattr(args, WARMUP_MAX_LR) and args.warmup_max_lr is not None: - params[WARMUP_MAX_LR] = args.warmup_max_lr - - if hasattr(args, WARMUP_NUM_STEPS) and args.warmup_num_steps is not None: - params[WARMUP_NUM_STEPS] = args.warmup_num_steps - - if hasattr(args, WARMUP_TYPE) and args.warmup_type is not None: - params[WARMUP_TYPE] = args.warmup_type - - -def override_params(args, params): - # LR range test params - override_lr_range_test_params(args, params) - - # 1Cycle params - override_1cycle_params(args, params) - - # WarmupLR params - override_warmupLR_params(args, params) - - -def get_config_from_args(args): - if not hasattr(args, LR_SCHEDULE) or args.lr_schedule is None: - return None, '--{} not specified on command line'.format(LR_SCHEDULE) - - if not args.lr_schedule in VALID_LR_SCHEDULES: - return None, '{} is not supported LR schedule'.format(args.lr_schedule) - - config = {} - config['type'] = args.lr_schedule - config['params'] = {} - - if args.lr_schedule == LR_RANGE_TEST: - override_lr_range_test_params(args, config['params']) - elif args.lr_schedule == ONE_CYCLE: - override_1cycle_params(args, config['params']) - else: - override_warmupLR_params(args, config['params']) - - return config, None - - -def get_lr_from_config(config): - if not 'type' in config: - return None, 'LR schedule type not defined in config' - - if not 'params' in config: - return None, 'LR schedule params not defined in config' - - lr_schedule = config['type'] - lr_params = config['params'] - - if not lr_schedule in VALID_LR_SCHEDULES: - return None, '{} is not a valid LR schedule'.format(lr_schedule) - - if lr_schedule == LR_RANGE_TEST: - return lr_params[LR_RANGE_TEST_MIN_LR], '' - if lr_schedule == ONE_CYCLE: - return lr_params[CYCLE_MAX_LR], '' - # Warmup LR - return lr_params[WARMUP_MAX_LR], '' - - -def update_lr(param_groups, lrs): - for param_group, lr in zip(param_groups, lrs): - param_group['lr'] = lr - return [group['lr'] for group in param_groups] - - -""" -Only optimizers that are subclass of torch.optim.Optimizer are supported. So check the passed optimizer and wrapped -optimizer to see if requirement is satisfied. -TODO: Looking under the hood to examine the wrapped optimizer is a hack that requires a better long-term fix. -""" - - -def get_torch_optimizer(optimizer): - if isinstance(optimizer, Optimizer): - return optimizer - - if hasattr(optimizer, 'optimizer') and isinstance(optimizer.optimizer, Optimizer): - return optimizer.optimizer - - raise TypeError('{} is not a subclass of torch.optim.Optimizer'.format(type(optimizer).__name__)) - - -class LRRangeTest(object): - """Sets the learning rate of each parameter group according to - learning rate range test (LRRT) policy. The policy increases learning - rate starting from a base value with a constant frequency, as detailed in - the paper `A disciplined approach to neural network hyper-parameters: Part 1 `_ - - LRRT policy is used for finding maximum LR that trains a model without divergence, and can be used to - configure the LR boundaries for Cyclic LR schedules. - - LRRT changes the learning rate after every batch. - `step` should be called after a batch has been used for training. - - Args: - optimizer (Optimizer): Wrapped optimizer. - lr_range_test_min_lr (float or list): Initial learning rate which is the - lower boundary in the range test for each parameter group. - lr_range_test_step_size (int): Interval of training steps to increase learning rate. Default: 2000 - lr_range_test_step_rate (float): Scaling rate for range test. Default: 1.0 - lr_range_test_staircase (bool): Scale in staircase fashion, rather than continuous. Default: False. - last_batch_iteration (int): The index of the last batch. This parameter is used when - resuming a training job. Since `step()` should be invoked after each - batch instead of after each epoch, this number represents the total - number of *batches* computed, not the total number of epochs computed. - When last_batch_iteration=-1, the schedule is started from the beginning. - Default: -1 - - Example: - >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) - >>> scheduler = LRRangeTest(optimizer) - >>> data_loader = torch.utils.data.DataLoader(...) - >>> for epoch in range(10): - >>> for batch in data_loader: - >>> train_batch(...) - >>> scheduler.step() - - _A disciplined approach to neural network hyper-parameters: Part 1 -- learning rate, batch size, momentum, and weight decay: - https://arxiv.org/abs/1803.09820 -""" - - def __init__(self, - optimizer: Optimizer, - lr_range_test_min_lr: float = 1e-3, - lr_range_test_step_size: int = 2000, - lr_range_test_step_rate: float = 1.0, - lr_range_test_staircase: bool = False, - last_batch_iteration: int = -1): - - self.optimizer = get_torch_optimizer(optimizer) - - if isinstance(lr_range_test_min_lr, list) or isinstance(lr_range_test_min_lr, tuple): - if len(lr_range_test_min_lr) != len(self.optimizer.param_groups): - raise ValueError("expected {} lr_range_test_min_lr, got {}".format(len(self.optimizer.param_groups), - len(lr_range_test_min_lr))) - self.min_lr = list(lr_range_test_min_lr) - else: - self.min_lr = [lr_range_test_min_lr] * len(self.optimizer.param_groups) - - self.step_size = lr_range_test_step_size - self.step_rate = lr_range_test_step_rate - self.last_batch_iteration = last_batch_iteration - self.staircase = lr_range_test_staircase - self.interval_fn = self._staircase_interval if lr_range_test_staircase else self._continuous_interval - - if last_batch_iteration == -1: - self._last_lr = update_lr(self.optimizer.param_groups, self.min_lr) - - def _staircase_interval(self): - return math.floor(float(self.last_batch_iteration + 1) / self.step_size) - - def _continuous_interval(self): - return float(self.last_batch_iteration + 1) / self.step_size - - def _get_increase(self): - return (1 + self.step_rate * self.interval_fn()) - - def get_lr(self): - lr_increase = self._get_increase() - return [lr_range_test_min_lr * lr_increase for lr_range_test_min_lr in self.min_lr] - - def get_last_lr(self): - """ Return last computed learning rate by current scheduler. - """ - assert getattr(self, '_last_lr', None) is not None, "need to call step() first" - return self._last_lr - - def step(self, batch_iteration=None): - if batch_iteration is None: - batch_iteration = self.last_batch_iteration + 1 - self.last_batch_iteration = batch_iteration - self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) - - def state_dict(self): - return {'last_batch_iteration': self.last_batch_iteration} - - def load_state_dict(self, sd): - self.last_batch_iteration = sd['last_batch_iteration'] - - -class OneCycle(object): - """Sets the learning rate of each parameter group according to - 1Cycle learning rate policy (1CLR). 1CLR is a variation of the - Cyclical Learning Rate (CLR) policy that involves one cycle followed by - decay. The policy simultaneously cycles the learning rate (and momentum) - between two boundaries with a constant frequency, as detailed in - the paper `A disciplined approach to neural network hyper-parameters`_. - - 1CLR policy changes the learning rate after every batch. - `step` should be called after a batch has been used for training. - - This implementation was adapted from the github repo: `PyTorch `_. - - Args: - optimizer (Optimizer): Wrapped optimizer. - cycle_min_lr (float or list): Initial learning rate which is the - lower boundary in the cycle for each parameter group. - cycle_max_lr (float or list): Upper learning rate boundaries in the cycle - for each parameter group. Functionally, - it defines the cycle amplitude (cycle_max_lr - cycle_min_lr). - The lr at any cycle is the sum of cycle_min_lr - and some scaling of the amplitude; therefore - cycle_max_lr may not actually be reached depending on - scaling function. - decay_lr_rate(float): Decay rate for learning rate. Default: 0. - cycle_first_step_size (int): Number of training iterations in the - increasing half of a cycle. Default: 2000 - cycle_second_step_size (int): Number of training iterations in the - decreasing half of a cycle. If cycle_second_step_size is None, - it is set to cycle_first_step_size. Default: None - cycle_first_stair_count(int): Number of stairs in first half of cycle phase. This means - lr/mom are changed in staircase fashion. Default 0, means staircase disabled. - cycle_second_stair_count(int): Number of stairs in second half of cycle phase. This means - lr/mom are changed in staircase fashion. Default 0, means staircase disabled. - decay_step_size (int): Intervals for applying decay in decay phase. Default: 0, means no decay. - cycle_momentum (bool): If ``True``, momentum is cycled inversely - to learning rate between 'cycle_min_mom' and 'cycle_max_mom'. - Default: True - cycle_min_mom (float or list): Initial momentum which is the - lower boundary in the cycle for each parameter group. - Default: 0.8 - cycle_max_mom (float or list): Upper momentum boundaries in the cycle - for each parameter group. Functionally, - it defines the cycle amplitude (cycle_max_mom - cycle_min_mom). - The momentum at any cycle is the difference of cycle_max_mom - and some scaling of the amplitude; therefore - cycle_min_mom may not actually be reached depending on - scaling function. Default: 0.9 - decay_mom_rate (float): Decay rate for momentum. Default: 0. - last_batch_iteration (int): The index of the last batch. This parameter is used when - resuming a training job. Since `step()` should be invoked after each - batch instead of after each epoch, this number represents the total - number of *batches* computed, not the total number of epochs computed. - When last_batch_iteration=-1, the schedule is started from the beginning. - Default: -1 - - Example: - >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) - >>> scheduler = OneCycle(optimizer, 0.0001, 0.0010) - >>> data_loader = torch.utils.data.DataLoader(...) - >>> for epoch in range(10): - >>> for batch in data_loader: - >>> train_batch(...) - >>> scheduler.step() - - - .. _A disciplined approach to neural network hyper-parameters: Part 1 -- learning rate, batch size, momentum, and weight decay: https://arxiv.org/abs/1803.09820 - """ - - def __init__(self, - optimizer, - cycle_min_lr, - cycle_max_lr, - decay_lr_rate=0., - cycle_first_step_size=2000, - cycle_second_step_size=None, - cycle_first_stair_count=0, - cycle_second_stair_count=None, - decay_step_size=0, - cycle_momentum=True, - cycle_min_mom=0.8, - cycle_max_mom=0.9, - decay_mom_rate=0., - last_batch_iteration=-1): - - self.optimizer = get_torch_optimizer(optimizer) - - # Initialize cycle shape - self._initialize_cycle(cycle_first_step_size, cycle_second_step_size, cycle_first_stair_count, - cycle_second_stair_count, decay_step_size) - - # Initialize cycle lr - self._initialize_lr(self.optimizer, cycle_min_lr, cycle_max_lr, decay_lr_rate, last_batch_iteration) - - # Initialize cyclic momentum - self.cycle_momentum = cycle_momentum - if cycle_momentum: - self._initialize_momentum(self.optimizer, cycle_min_mom, cycle_max_mom, decay_mom_rate, - last_batch_iteration) - # Initialize batch iteration tracker - self.last_batch_iteration = last_batch_iteration - - # Configure cycle shape - - def _initialize_cycle(self, cycle_first_step_size, cycle_second_step_size, cycle_first_stair_count, - cycle_second_stair_count, decay_step_size): - cycle_first_step_size = float(cycle_first_step_size) - cycle_second_step_size = float( - cycle_second_step_size) if cycle_second_step_size is not None else cycle_first_step_size - - self.total_size = cycle_first_step_size + cycle_second_step_size - self.step_ratio = cycle_first_step_size / self.total_size - self.first_stair_count = cycle_first_stair_count - self.second_stair_count = cycle_first_stair_count if cycle_second_stair_count is None else cycle_second_stair_count - self.decay_step_size = decay_step_size - - if math.isclose(self.decay_step_size, 0): - self.skip_lr_decay = True - self.skip_mom_decay = True - else: - self.skip_lr_decay = False - self.skip_mom_decay = False - - # Configure lr schedule - def _initialize_lr(self, optimizer, cycle_min_lr, cycle_max_lr, decay_lr_rate, last_batch_iteration): - self.min_lrs = [cycle_min_lr] * len(optimizer.param_groups) - if last_batch_iteration == -1: - for lr, group in zip(self.min_lrs, optimizer.param_groups): - group['lr'] = lr - - self.max_lrs = [cycle_max_lr] * len(optimizer.param_groups) - self.decay_lr_rate = decay_lr_rate - - if math.isclose(self.decay_lr_rate, 0): - self.skip_lr_decay = True - - # Configure momentum schedule - def _initialize_momentum(self, optimizer, cycle_min_mom, cycle_max_mom, decay_mom_rate, last_batch_iteration): - if 'betas' not in optimizer.defaults: - optimizer_name = type(optimizer).__name__ - logger.warning( - f"cycle_momentum is disabled because optimizer {optimizer_name} does not support momentum, no betas attribute in defaults" - ) - self.cycle_momentum = False - return - - self.decay_mom_rate = decay_mom_rate - self.min_moms = [(cycle_min_mom, 0.99)] * len(optimizer.param_groups) - self.max_moms = [(cycle_max_mom, 0.99)] * len(optimizer.param_groups) - - if last_batch_iteration == -1: - for momentum, group in zip(self.min_moms, optimizer.param_groups): - group['betas'] = momentum - - if math.isclose(self.decay_mom_rate, 0): - self.skip_mom_decay = True - - def _get_scale_factor(self): - batch_iteration = (self.last_batch_iteration + 1) - cycle = math.floor(1 + batch_iteration / self.total_size) - x = 1. + batch_iteration / self.total_size - cycle - if x <= self.step_ratio: - scale_factor = x / self.step_ratio - else: - scale_factor = (x - 1) / (self.step_ratio - 1) - - return scale_factor - - def _get_cycle_mom(self): - scale_factor = self._get_scale_factor() - momentums = [] - for base_betas, max_betas in zip(self.min_moms, self.max_moms): - cycle_min_mom = base_betas[0] - cycle_max_mom = max_betas[0] - base_height = (cycle_max_mom - cycle_min_mom) * scale_factor - momentum = cycle_max_mom - base_height - momentums.append((momentum, base_betas[1])) - return momentums - - def _get_cycle_lr(self): - scale_factor = self._get_scale_factor() - lrs = [] - for cycle_min_lr, cycle_max_lr in zip(self.min_lrs, self.max_lrs): - base_height = (cycle_max_lr - cycle_min_lr) * scale_factor - lr = cycle_min_lr + base_height - lrs.append(lr) - - return lrs - - def _get_decay_mom(self, decay_batch_iteration): - if self.skip_mom_decay: - return self.max_moms - - decay_interval = decay_batch_iteration / self.decay_step_size - mom_decay_factor = (1 + self.decay_mom_rate * decay_interval) - momentums = [(beta0 * mom_decay_factor, beta1) for beta0, beta1 in self.max_moms] - - return momentums - - def _get_decay_lr(self, decay_batch_iteration): - """Calculates the learning rate at batch index. This function is used - after the cycle completes and post cycle decaying of lr/mom is enabled. - This function treats `self.last_batch_iteration` as the last batch index. - """ - if self.skip_lr_decay: - return self.min_lrs - - decay_interval = decay_batch_iteration / self.decay_step_size - lr_decay_factor = (1 + self.decay_lr_rate * decay_interval) - lrs = [cycle_min_lr / lr_decay_factor for cycle_min_lr in self.min_lrs] - - return lrs - - def get_lr(self): - """Calculates the learning rate at batch index. This function treats - `self.last_batch_iteration` as the last batch index. - """ - if self.last_batch_iteration < self.total_size: - return self._get_cycle_lr() - return self._get_decay_lr(self.last_batch_iteration - self.total_size + 1) - - def get_mom(self): - """Calculates the momentum at batch index. This function treats - `self.last_batch_iteration` as the last batch index. - """ - if not self.cycle_momentum: - return None - - if self.last_batch_iteration < self.total_size: - return self._get_cycle_mom() - return self._get_decay_mom(self.last_batch_iteration - self.total_size + 1) - - def get_last_lr(self): - """ Return last computed learning rate by current scheduler. - """ - assert getattr(self, '_last_lr', None) is not None, "need to call step() first" - return self._last_lr - - def step(self, batch_iteration=None): - """ Updates the optimizer with the learning rate for the last batch index. - `self.last_batch_iteration` is treated as the last batch index. - - If self.cycle_momentum is true, also updates optimizer momentum. - """ - if batch_iteration is None: - batch_iteration = self.last_batch_iteration + 1 - - self.last_batch_iteration = batch_iteration - self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) - - if self.cycle_momentum: - momentums = self.get_mom() - for param_group, momentum in zip(self.optimizer.param_groups, momentums): - param_group['betas'] = momentum - - def state_dict(self): - return {'last_batch_iteration': self.last_batch_iteration} - - def load_state_dict(self, sd): - self.last_batch_iteration = sd['last_batch_iteration'] - - -class WarmupLR(object): - """Increase the learning rate of each parameter group from min lr to max lr - over warmup_num_steps steps, and then fix at max lr. - - Args: - optimizer (Optimizer): Wrapped optimizer. - warmup_min_lr (float or list): minimum learning rate. Default: 0 - warmup_max_lr (float or list): maximum learning rate. Default: 0.001 - warmup_num_steps (int): number of steps to warm up from min_lr to max_lr. Default: 1000 - warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log - last_batch_iteration (int): The index of the last batch. Default: -1. - Example: - >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) - >>> scheduler = WarmupLR(optimizer) - >>> data_loader = torch.utils.data.DataLoader(...) - >>> for epoch in range(10): - >>> for batch in data_loader: - >>> train_batch(...) - >>> scheduler.step() - - """ - - def __init__(self, - optimizer: Optimizer, - warmup_min_lr: float = 0.0, - warmup_max_lr: float = None, - warmup_num_steps: int = 1000, - warmup_type: str = WARMUP_LOG_RATE, - last_batch_iteration: int = -1): - - self.optimizer = get_torch_optimizer(optimizer) - - if warmup_max_lr is None: - warmup_max_lr = [group['lr'] for group in self.optimizer.param_groups][0] - - self.min_lrs = self._format_param(self.optimizer, warmup_min_lr, "min_lr") - self.max_lrs = self._format_param(self.optimizer, warmup_max_lr, "max_lr") - self.delta_lrs = [big - small for big, small in zip(self.max_lrs, self.min_lrs)] - self.warmup_num_steps = max(2, warmup_num_steps) - # Currently only support linear and log function - if warmup_type not in {WARMUP_LOG_RATE, WARMUP_LINEAR_RATE}: - logger.warning(f"Using unknown warmup_type: {warmup_type}. The increasing function " - f"is set to default (log)") - warmup_type = WARMUP_LOG_RATE - self.warmup_type = warmup_type - self.inverse_log_warm_up = 1.0 / math.log(self.warmup_num_steps) - self.last_batch_iteration = last_batch_iteration - # Initialize lr in optimizer - if last_batch_iteration == -1: - self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) - - def get_lr(self): - if self.last_batch_iteration < 0: - logger.warning("Attempting to get learning rate from scheduler before it has started") - return self.min_lrs - gamma = self._get_gamma() - return [min_lr + (delta_lr * gamma) for min_lr, delta_lr in zip(self.min_lrs, self.delta_lrs)] - - def get_last_lr(self): - """ Return last computed learning rate by current scheduler. - """ - assert getattr(self, '_last_lr', None) is not None, "need to call step() first" - return self._last_lr - - def step(self, last_batch_iteration=None): - if last_batch_iteration is None: - last_batch_iteration = self.last_batch_iteration + 1 - self.last_batch_iteration = last_batch_iteration - self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) - - def state_dict(self): - return {'last_batch_iteration': self.last_batch_iteration} - - def load_state_dict(self, sd): - self.last_batch_iteration = sd['last_batch_iteration'] - - def _get_gamma(self): - if self.last_batch_iteration < self.warmup_num_steps: - if self.warmup_type == WARMUP_LOG_RATE: - return self.inverse_log_warm_up * math.log(self.last_batch_iteration + 1) - elif self.warmup_type == WARMUP_LINEAR_RATE: - return self.last_batch_iteration / self.warmup_num_steps - return 1.0 - - def _format_param(self, optimizer, param_value, param_name): - if isinstance(param_value, list) or isinstance(param_value, tuple): - if len(param_value) != len(optimizer.param_groups): - raise ValueError("expected {} value for {}, got {}".format(len(optimizer.param_groups), param_name, - FileNotFoundError(param_value))) - return list(param_value) - return [param_value] * len(optimizer.param_groups) - - -class WarmupDecayLR(WarmupLR): - """Increase the learning rate of each parameter group from min lr to max lr - over warmup_num_steps steps, and then decay at linear rate over the remaining training steps. - - Args: - optimizer (Optimizer): Wrapped optimizer. - total_num_steps (int): total number of training steps - warmup_min_lr (float or list): minimum learning rate. Default: 0 - warmup_max_lr (float or list): maximum learning rate. Default: 0.001 - warmup_num_steps (int): number of steps to warm up from min_lr to max_lr. Default: 1000 - warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log - last_batch_iteration (int): The index of the last batch. Default: -1. - Example: - >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) - >>> scheduler = WarmupDecayLR(optimizer, 1000000) - >>> data_loader = torch.utils.data.DataLoader(...) - >>> for epoch in range(10): - >>> for batch in data_loader: - >>> train_batch(...) - >>> scheduler.step() - - """ - - def __init__(self, - optimizer: Optimizer, - total_num_steps: int, - warmup_min_lr: float = 0.0, - warmup_max_lr: float = 0.001, - warmup_num_steps: int = 1000, - warmup_type: str = WARMUP_LOG_RATE, - last_batch_iteration: int = -1): - - self.total_num_steps = total_num_steps - super(WarmupDecayLR, self).__init__(optimizer, warmup_min_lr, warmup_max_lr, warmup_num_steps, warmup_type, - last_batch_iteration) - if self.total_num_steps < self.warmup_num_steps: - logger.warning('total_num_steps {} is less than warmup_num_steps {}'.format( - total_num_steps, warmup_num_steps)) - - def _get_gamma(self): - if self.last_batch_iteration < self.warmup_num_steps: - if self.warmup_type == WARMUP_LOG_RATE: - return self.inverse_log_warm_up * math.log(self.last_batch_iteration + 1) - elif self.warmup_type == WARMUP_LINEAR_RATE: - return self.last_batch_iteration / self.warmup_num_steps - return max( - 0.0, - float(self.total_num_steps - self.last_batch_iteration) / - float(max(1.0, self.total_num_steps - self.warmup_num_steps))) - - -class WarmupCosineLR(object): - """Increase the learning rate of each parameter group from min lr ratio to max lr ratio - over warmup_num_steps steps, and then decay at cosine rate over the remaining training steps to min cosine ratio. - - Args: - optimizer (Optimizer): Wrapped optimizer. - total_num_steps (int): total number of training steps - warmup_min_ratio (float or list): warmup start learning rate ratio. Default: 0 - warmup_num_steps (int): number of steps to warm up from warmup_min_ratio to 1.0. Default: 1000 - warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log - cos_min_ratio (float): cosine end learning rate ratio. Default: 0.0001 - last_batch_iteration (int): The index of the last batch. Default: -1. - Example: - >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) - >>> scheduler = WarmupCosineLR(optimizer, 1000000) - >>> data_loader = torch.utils.data.DataLoader(...) - >>> for epoch in range(10): - >>> for batch in data_loader: - >>> train_batch(...) - >>> scheduler.step() - - """ - - def __init__(self, - optimizer: Optimizer, - total_num_steps: int, - warmup_min_ratio: float = 0.0, - warmup_num_steps: int = 1000, - cos_min_ratio: float = 0.0001, - warmup_type: str = WARMUP_LOG_RATE, - last_batch_iteration: int = -1): - - self.optimizer = get_torch_optimizer(optimizer) - - self.total_num_steps = total_num_steps - self.last_batch_iteration = last_batch_iteration - self.cos_min_ratio = cos_min_ratio - - self.warmup_type = warmup_type - self.warmup_min_ratio = warmup_min_ratio - self.warmup_num_steps = max(2, warmup_num_steps) - self.inverse_log_warm_up = 1.0 / math.log(self.warmup_num_steps) - - if self.total_num_steps < self.warmup_num_steps: - logger.warning('total_num_steps {} is less than warmup_num_steps {}'.format( - total_num_steps, warmup_num_steps)) - self.org_lrs = [group['lr'] for group in self.optimizer.param_groups] - - # Initialize lrs in optimizer groups - if last_batch_iteration == -1: - self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) - - def get_lr_ratio(self): - if self.last_batch_iteration < 0: - logger.warning("Attempting to get learning rate from scheduler before it has started") - return [0.0] - - if self.last_batch_iteration < self.warmup_num_steps: - if self.warmup_type == WARMUP_LOG_RATE: - ratio = self.inverse_log_warm_up * math.log(self.last_batch_iteration + 1) - elif self.warmup_type == WARMUP_LINEAR_RATE: - ratio = self.last_batch_iteration / self.warmup_num_steps - ratio_delta = 1. - self.warmup_min_ratio - ratio = self.warmup_min_ratio + ratio * ratio_delta - return ratio - - real_last_step = self.last_batch_iteration - self.warmup_num_steps + 1 - real_total_steps = self.total_num_steps - self.warmup_num_steps - ratio_delta = 1. - self.cos_min_ratio - ratio = (1 + math.cos(math.pi * real_last_step / real_total_steps)) / 2 - ratio = max(0.0, self.cos_min_ratio + ratio_delta * ratio) - return ratio - - def step(self, last_batch_iteration=None): - if last_batch_iteration is None: - last_batch_iteration = self.last_batch_iteration + 1 - self.last_batch_iteration = last_batch_iteration - self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) - - def get_lr(self): - if self.last_batch_iteration < 0: - logger.warning("Attempting to get learning rate from scheduler before it has started") - return [0.0] - lr_ratio = self.get_lr_ratio() - return [org_lr * lr_ratio for org_lr in self.org_lrs] - - def get_last_lr(self): - """ Return last computed learning rate by current scheduler. - """ - assert getattr(self, '_last_lr', None) is not None, "need to call step() first" - return self._last_lr - - def state_dict(self): - return {'last_batch_iteration': self.last_batch_iteration} - - def load_state_dict(self, sd): - self.last_batch_iteration = sd['last_batch_iteration'] - - def _format_param(self, optimizer, param_value, param_name): - if isinstance(param_value, list) or isinstance(param_value, tuple): - if len(param_value) != len(optimizer.param_groups): - raise ValueError("expected {} value for {}, got {}".format(len(optimizer.param_groups), param_name, - FileNotFoundError(param_value))) - return list(param_value) - return [param_value] * len(optimizer.param_groups) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__init__.py deleted file mode 100644 index 5e60b03ac67112149c351045c3d9e117633d2b3f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .constants import * -from .writer_factory import CheckpointWriterFactory diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index b07dfb2fcc29406002c779623bb3db131e1b46f2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/config.cpython-310.pyc deleted file mode 100644 index 8a86f5f95921e1b13802640a5e515e5cdd1a6a83..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/constants.cpython-310.pyc deleted file mode 100644 index cbe7b3206b90c8081a08d76dacce8a8c7cc73209..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/constants.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/data_parallel_writer_factory.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/data_parallel_writer_factory.cpython-310.pyc deleted file mode 100644 index 4141c05accd43bc22d1f3e5e4dc5514c4edf2fd5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/data_parallel_writer_factory.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index b62b0a0ccb79a671444198bf15f38aa788d24290..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/writer_factory.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/writer_factory.cpython-310.pyc deleted file mode 100644 index 62aebef8d3d924280c4d1b7347e1f983d1c32d3a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/__pycache__/writer_factory.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/config.py deleted file mode 100644 index d5a579fe31f09668e59647b433e427447df01367..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/config.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.runtime.config_utils import get_scalar_param -from .constants import * - -VALID_VALUES = { - CHECKPOINT_TAG_VALIDATION: CHECKPOINT_TAG_VALIDATION_MODES, - CHECKPOINT_WRITER_TYPE: CHECKPOINT_WRITER_TYPES, - CHECKPOINT_DATA_PARALLEL: CHECKPOINT_DATA_PARALLEL_UNITS -} - -CHECKPOINT_DEFAULT_DICT = { - CHECKPOINT_TAG_VALIDATION: CHECKPOINT_TAG_VALIDATION_DEFAULT, - CHECKPOINT_SERIALIZATION: CHECKPOINT_SERIALIZATION_DEFAULT, - CHECKPOINT_WRITER: CHECKPOINT_WRITER_DEFAULT -} - - -def _validate_config_values(config_name, config_dict, valid_values): - for key, value in config_dict.items(): - if value is None: - continue - if key in valid_values.keys(): - assert value in valid_values[key], \ - f"{config_name} contains invalid value {value} for {key}, expecting one of {valid_values[key]}" - - -def _make_upper_case(value): - return value if value is None else value.upper() - - -def get_checkpoint_writer_config(param_dict): - writer_dict = param_dict.get(CHECKPOINT_WRITER, None) - if writer_dict is None: - return CHECKPOINT_WRITER_DEFAULT - - writer_config = { - CHECKPOINT_WRITER_TYPE: - _make_upper_case(get_scalar_param(writer_dict, CHECKPOINT_WRITER_TYPE, CHECKPOINT_WRITER_TYPE_DEFAULT)), - CHECKPOINT_IO_BUFFER_SIZE: - get_scalar_param(writer_dict, CHECKPOINT_IO_BUFFER_SIZE, CHECKPOINT_IO_BUFFER_SIZE_DEFAULT), - CHECKPOINT_IO_BUFFER_DOUBLE: - get_scalar_param(writer_dict, CHECKPOINT_IO_BUFFER_DOUBLE, CHECKPOINT_IO_BUFFER_DOUBLE_DEFAULT), - CHECKPOINT_IO_STATISTICS: - get_scalar_param(writer_dict, CHECKPOINT_IO_STATISTICS, CHECKPOINT_IO_STATISTICS_DEFAULT), - CHECKPOINT_DATA_PARALLEL: - _make_upper_case(get_scalar_param(writer_dict, CHECKPOINT_DATA_PARALLEL, CHECKPOINT_DATA_PARALLEL_DEFAULT)), - CHECKPOINT_WRITER_DECOUPLED: - get_scalar_param(writer_dict, CHECKPOINT_WRITER_DECOUPLED, CHECKPOINT_WRITER_DECOUPLED_DEFAULT), - CHECKPOINT_IO_MULTIPLIER: - get_scalar_param(writer_dict, CHECKPOINT_IO_MULTIPLIER, CHECKPOINT_IO_MULTIPLIER_DEFAULT), - } - _validate_config_values(CHECKPOINT_WRITER, writer_config, VALID_VALUES) - - return writer_config - - -def get_checkpoint_config(param_dict): - checkpoint_dict = param_dict.get(CHECKPOINT, None) - if checkpoint_dict is None: - return CHECKPOINT_DEFAULT_DICT - - checkpoint_config = { - CHECKPOINT_TAG_VALIDATION: - get_scalar_param(checkpoint_dict, CHECKPOINT_TAG_VALIDATION, CHECKPOINT_TAG_VALIDATION_DEFAULT).upper(), - CHECKPOINT_SERIALIZATION: - get_scalar_param(checkpoint_dict, CHECKPOINT_SERIALIZATION, CHECKPOINT_SERIALIZATION_DEFAULT), - CHECKPOINT_WRITER: - get_checkpoint_writer_config(checkpoint_dict) - } - - _validate_config_values(CHECKPOINT, checkpoint_config, VALID_VALUES) - - return checkpoint_config diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/constants.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/constants.py deleted file mode 100644 index 3b9bd549af92b1458e65d0b0702da659e6d2fd04..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/constants.py +++ /dev/null @@ -1,85 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - - -######################################### -# Validation modes -######################################### -class ValidationMode: - WARN = "WARN" - IGNORE = "IGNORE" - FAIL = "FAIL" - - -######################################### -# Checkpoint config params -######################################### -# "checkpoint": {tag_validation=["Ignore"|"Warn"|"Fail"]} -CHECKPOINT_FORMAT = ''' -"checkpoint": { - "tag_validation": [Ignore|Warn|Fail], - "checkpoint_serialization": False, - "writer": { - "type": [mock|python|fast], - "decoupled": [True|False] - "io_buffer_size": 64e6, - "io_buffer_double": True, - "show_statistics": False, - "data_parallel": [replica|socket|machine], - "io_multiplier": 1, - } -} -''' -CHECKPOINT = "checkpoint" -CHECKPOINT_TAG_VALIDATION = "tag_validation" -CHECKPOINT_TAG_VALIDATION_DEFAULT = ValidationMode.WARN -CHECKPOINT_TAG_VALIDATION_MODES = [ValidationMode.WARN, ValidationMode.IGNORE, ValidationMode.FAIL] - -CHECKPOINT_SERIALIZATION = "checkpoint_serialization" -CHECKPOINT_SERIALIZATION_DEFAULT = True - -CHECKPOINT_WRITER = "writer" -CHECKPOINT_WRITER_DEFAULT = None - -CHECKPOINT_WRITER_TYPE = "type" - - -class CheckpointWriterType: - MOCK = "MOCK" - PYTHON = "PYTHON" - FAST = "FAST" - - -CHECKPOINT_WRITER_TYPE_DEFAULT = CheckpointWriterType.FAST -CHECKPOINT_WRITER_TYPES = [CheckpointWriterType.MOCK, CheckpointWriterType.PYTHON, CheckpointWriterType.FAST] - -CHECKPOINT_IO_BUFFER_SIZE = "io_buffer_size" -CHECKPOINT_IO_BUFFER_SIZE_DEFAULT = 64 * (1024**2) - -CHECKPOINT_IO_BUFFER_DOUBLE = "io_buffer_double" -CHECKPOINT_IO_BUFFER_DOUBLE_DEFAULT = True - -CHECKPOINT_IO_MULTIPLIER = "io_multiplier" -CHECKPOINT_IO_MULTIPLIER_DEFAULT = 1 - -CHECKPOINT_IO_STATISTICS = "show_statistics" -CHECKPOINT_IO_STATISTICS_DEFAULT = False - -CHECKPOINT_DATA_PARALLEL = "data_parallel" -CHECKPOINT_DATA_PARALLEL_DEFAULT = None - - -class CheckpointDataParallel: - REPLICA = "REPLICA" - SOCKET = "SOCKET" - MACHINE = "MACHINE" - - -CHECKPOINT_DATA_PARALLEL_UNITS = [ - CheckpointDataParallel.REPLICA, CheckpointDataParallel.SOCKET, CheckpointDataParallel.MACHINE -] - -CHECKPOINT_WRITER_DECOUPLED = "decoupled" -CHECKPOINT_WRITER_DECOUPLED_DEFAULT = False diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/data_parallel_writer_factory.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/data_parallel_writer_factory.py deleted file mode 100644 index 15ed235e47066c139e0db662c2c718b436103434..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/data_parallel_writer_factory.py +++ /dev/null @@ -1,216 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from dataclasses import dataclass -from deepspeed.checkpoint.reshape_utils import partition_data -from deepspeed.runtime.zero.config import ZeroStageEnum -from .constants import * - - -@dataclass -class DataParallelWriterConfig(object): - world_size: int - rank: int - global_rank: int - local_rank: int - pure_dp: bool - - -class DataParallelWriterFactory(object): - - def __init__(self, uni_parallel_info, parallel_unit): - self._uni_parallel_info = uni_parallel_info - self._parallel_unit = parallel_unit - if parallel_unit == CheckpointDataParallel.SOCKET: - self._num_resources = uni_parallel_info.num_sockets - else: - self._num_resources = uni_parallel_info.num_machines - self._ranks_per_resource = max(1, self._uni_parallel_info.global_world_size // self._num_resources) - - def create_config(self, zero_stage, has_moe_layers): - if zero_stage == ZeroStageEnum.weights: - return self._create_config(1, 0) - - if has_moe_layers: - writer_config = self._get_expert_data_parallel_config() - else: - writer_config = self._get_data_parallel_config() - - if writer_config is None and zero_stage >= ZeroStageEnum.optimizer_states: - return self._create_config(1, 0) - - return writer_config - - def _create_config(self, world_size, rank): - return DataParallelWriterConfig(world_size=world_size, - rank=rank, - global_rank=self._uni_parallel_info.global_rank, - local_rank=self._uni_parallel_info.local_rank, - pure_dp=self._uni_parallel_info.pure_dp) - - def _get_expert_data_parallel_config(self): - ep_info = self._uni_parallel_info.ep_info - if self._parallel_unit is None: - dp_rank = ep_info.dp_rank - return self._create_config(1, 0) if dp_rank == 0 else None - - assert self._uni_parallel_info.pure_dp, \ - f'3D parallelism is not yet supported for data parallel checkpointing.' - - if self._parallel_unit == CheckpointDataParallel.REPLICA or ep_info.ep_world_size == 1: - return self._get_parallel_write_for_ddp(ep_info.dp_world_size, ep_info.dp_rank) - - return self._get_expert_parallel_write_for_2d() - - def _get_expert_parallel_write_for_2d(self): - ep_info = self._uni_parallel_info.ep_info - - def _get_expert_slice_resources(expert_resources, resource_name): - ep_world_size = ep_info.ep_world_size - slices_per_resource = min(self._ranks_per_resource, ep_world_size) - assert slices_per_resource <= len(expert_resources) - - ep_num_resources = len(expert_resources) - assert ep_num_resources % slices_per_resource == 0, f'{resource_name}: Expected ep_num_resources={ep_num_resources} to multiple of slices_per_resource={slices_per_resource} for ep_world_size={ep_world_size}' - - slice_partitions = partition_data(expert_resources, slices_per_resource) - # print( - # f'edp_resource_partition: self._uni_parallel_info.global_rank={self._uni_parallel_info.global_rank} expert_resources={expert_resources} slices_per_resource={slices_per_resource} ep_world_size={ep_world_size} slice_partitions={slice_partitions}' - # ) - resource_index = ep_info.ep_rank % slice_resources - return slice_partitions[resource_index] - - dp_ranks = ep_info.dp_peer_ranks - expert_resources = [r // self._ranks_per_resource for r in dp_ranks] - slice_resources = _get_expert_slice_resources(expert_resources, self._parallel_unit) - assert all([idx < self._num_resources for idx in expert_resources]), \ - f'Detected invalid resource index in expert_resources={expert_resources}, self._num_resources={self._num_resources}' - return self._assign_resources_to_tensor_slice(slice_resources, ep_info.ep_rank, dp_ranks) - - def _get_data_parallel_config(self): - mpu_info = self._uni_parallel_info.mpu_info - if self._parallel_unit is None: - dp_rank = self._uni_parallel_info.dp_rank if mpu_info is None else mpu_info.dp_rank - return self._create_config(1, 0) if dp_rank == 0 else None - - if self._uni_parallel_info.pure_dp: - return self._get_parallel_write_for_ddp(self._uni_parallel_info.global_world_size, - self._uni_parallel_info.global_rank) - - if self._parallel_unit == CheckpointDataParallel.REPLICA: - return self._create_config(mpu_info.dp_world_size, mpu_info.dp_rank) - - return self._get_parallel_write_for_3d() - - def _get_parallel_write_for_3d(self): - mpu_info = self._uni_parallel_info.mpu_info - my_global_rank = self._uni_parallel_info.global_rank - - def _expand_resources(resource_list, new_size): - old_size = len(resource_list) - if old_size >= new_size: - return resource_list - - assert new_size % old_size == 0, f'Expect new_size={new_size} to be multiple of old_size={old_size}' - multiplier = new_size // old_size - new_resource_list = [] - for r in resource_list: - new_resource_list += [r] * multiplier - # print(f'expand_resources: {my_global_rank=} {old_size=} {new_size=} {resource_list=} {new_resource_list=}') - return new_resource_list - - # Getting resource partition for a tensor slice is a 2-step process - # 1. Get resource partitions for all pipeline stages. A pipeline stage is a 2D grid of size TP x DP - def _get_pipeline_stage_resources(resource_indices): - num_resources = len(resource_indices) - pp_world_size = mpu_info.pp_world_size - if num_resources < pp_world_size: - resource_indices = _expand_resources(resource_indices, pp_world_size) - num_resources = pp_world_size - global_resource_partitions = partition_data(resource_indices, pp_world_size) - pp_rank = mpu_info.pp_rank - return global_resource_partitions[pp_rank] - - # 2. Get resource partition for tensor slice. A tensor slice is a 1D vector of size DP - def _get_tensor_slice_resources(resource_indices, resource_name): - pipe_stage_resources = _get_pipeline_stage_resources(resource_indices) - tp_world_size = mpu_info.tp_world_size - if len(pipe_stage_resources) < tp_world_size: - pipe_stage_resources = _expand_resources(pipe_stage_resources, tp_world_size) - tp_num_resources = len(pipe_stage_resources) - assert tp_num_resources % tp_world_size == 0, \ - f'{resource_name}: Expected tp_num_resources={tp_num_resources} to multiple of tp_world_size={tp_world_size}' - - pipe_stage_resource_partitions = partition_data(pipe_stage_resources, tp_world_size) - tp_rank = mpu_info.tp_rank - return pipe_stage_resource_partitions[tp_rank] - - def _get_model_parallel_slice_resources(): - # Get resources of my dp peer ranks - resources = [(r // self._ranks_per_resource) for r in mpu_info.dp_peer_ranks] - if len(resources) < self._ranks_per_resource: - resources = _expand_resources(resources, self._ranks_per_resource) - - resource_partitions = partition_data(resources, self._ranks_per_resource) - mp_rank = (mpu_info.pp_rank * mpu_info.tp_world_size) + mpu_info.tp_rank - slice_rank = mp_rank % self._ranks_per_resource - return resource_partitions[slice_rank] - - num_slices = mpu_info.tp_world_size * mpu_info.pp_world_size - if num_slices > self._ranks_per_resource: - slice_resources = _get_model_parallel_slice_resources() - else: - all_resources = list(range(self._num_resources)) - slice_resources = _get_tensor_slice_resources(all_resources, self._parallel_unit) - - return self._assign_resources_to_tensor_slice(slice_resources, mpu_info.tp_rank, mpu_info.dp_peer_ranks) - - def _get_slice_writers(self, slice_resources, my_dp_ranks): - resource_map = {} - for res in slice_resources: - resource_map[res] = [r for r in my_dp_ranks if (r // self._ranks_per_resource) == res] - - # Only one writer per resource, and we conventionally pick the first rank as writer. - return [ranks[0] for ranks in resource_map.values()] - - def _assign_resources_to_tensor_slice(self, slice_resources, my_slice_index, my_dp_ranks): - my_global_rank = self._uni_parallel_info.global_rank - slice_writer_ranks = self._get_slice_writers(slice_resources, my_dp_ranks) - my_resource_index = my_global_rank // self._ranks_per_resource - print( - f'resource_assign: my_global_rank={my_global_rank} my_slice_index={my_slice_index} my_dp_ranks={my_dp_ranks} slice_resources={slice_resources} slice_writer_ranks={slice_writer_ranks}' - ) - if my_resource_index in slice_resources and my_global_rank in slice_writer_ranks: - my_writer_index = (my_global_rank - slice_writer_ranks[0]) // self._ranks_per_resource - num_slice_writers = len(slice_writer_ranks) - print( - f'slice_writer: my_global_rank={my_global_rank} my_writer_index={my_writer_index} num_slice_writers={num_slice_writers}' - ) - return self._create_config(num_slice_writers, my_writer_index) - - return None - - def _get_parallel_write_for_ddp(self, dp_world_size, dp_rank): - if self._parallel_unit == CheckpointDataParallel.REPLICA: - return self._create_config(dp_world_size, dp_rank) - - num_machines = self._uni_parallel_info.num_machines - if self._parallel_unit == CheckpointDataParallel.SOCKET: - if dp_world_size == num_machines: - # There is one rank per machine - return self._create_config(num_machines, dp_rank) - - num_sockets = self._uni_parallel_info.num_sockets - ranks_per_socket = dp_world_size // num_sockets - if dp_rank % ranks_per_socket == 0: - return self._create_config(num_sockets, dp_rank // ranks_per_socket) - else: - return None - - ranks_per_machine = dp_world_size // num_machines - if dp_rank % ranks_per_machine == 0: - return self._create_config(num_machines, self._uni_parallel_info.machine_rank) - - return None diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/utils.py deleted file mode 100644 index e212008a9277f1804217a8aed7936004f25ee5fb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/utils.py +++ /dev/null @@ -1,84 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -from dataclasses import dataclass -from deepspeed import comm as dist -from deepspeed.constants import CROSS_RANK, CROSS_SIZE, LOCAL_RANK -from .data_parallel_writer_factory import DataParallelWriterFactory - -# TODO: parse socket number from env. -SOCKETS_PER_MACHINE = 2 - - -@dataclass -class MPUInfo(object): - pp_world_size: int - pp_rank: int - tp_world_size: int - tp_rank: int - dp_world_size: int - dp_peer_ranks: list - dp_rank: int - - -def _create_model_parallel_info(mpu): - return MPUInfo(pp_world_size=mpu.get_pipeline_model_parallel_world_size(), - pp_rank=mpu.get_pipeline_model_parallel_rank(), - tp_world_size=mpu.get_tensor_model_parallel_world_size(), - tp_rank=mpu.get_tensor_model_parallel_rank(), - dp_world_size=mpu.get_data_parallel_world_size(), - dp_peer_ranks=mpu.get_data_parallel_group_ranks(), - dp_rank=mpu.get_data_parallel_rank()) - - -@dataclass -class ExpertParallelInfo(object): - ep_world_size: int - ep_rank: int - dp_world_size: int - dp_peer_ranks: list - dp_rank: int - - -def _create_expert_parallel_info(groups): - group_name = groups._get_max_expert_size_name() - return ExpertParallelInfo(ep_world_size=groups._get_expert_parallel_world_size(group_name), - ep_rank=groups._get_expert_parallel_rank(group_name), - dp_world_size=groups._get_expert_data_parallel_world_size(group_name), - dp_peer_ranks=groups._get_expert_data_parallel_group_ranks(group_name), - dp_rank=groups._get_expert_data_parallel_rank(group_name)) - - -@dataclass -class UniversalParallelInfo(object): - global_world_size: int - global_rank: int - local_rank: int - mpu_info: MPUInfo - ep_info: ExpertParallelInfo - pure_dp: bool - num_machines: int - machine_rank: int - num_sockets: int - - -def create_universal_parallel_info(groups, has_moe_layers): - return UniversalParallelInfo(global_world_size=dist.get_world_size(), - global_rank=dist.get_rank(), - local_rank=int(os.environ[LOCAL_RANK]), - mpu_info=None if groups.mpu is None else _create_model_parallel_info(groups.mpu), - ep_info=_create_expert_parallel_info(groups) if has_moe_layers else None, - pure_dp=groups.mpu is None - or groups.mpu.get_data_parallel_world_size() == dist.get_world_size(), - num_machines=int(os.environ[CROSS_SIZE]), - machine_rank=int(os.environ[CROSS_RANK]), - num_sockets=int(os.environ[CROSS_SIZE]) * SOCKETS_PER_MACHINE) - - -def create_data_parallel_writer_config(groups, parallel_unit, zero_stage, has_moe_layers): - uni_parallel_info = create_universal_parallel_info(groups, has_moe_layers) - writer_factory = DataParallelWriterFactory(uni_parallel_info, parallel_unit) - return writer_factory.create_config(zero_stage, has_moe_layers) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/writer_factory.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/writer_factory.py deleted file mode 100644 index a8c324530ae52468219489ff2c9480288f2d547f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/model_checkpointing/writer_factory.py +++ /dev/null @@ -1,95 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed.ops.op_builder import AsyncIOBuilder, GDSBuilder -from deepspeed.io import MockFileWriter, PyFileWriter, FastFileWriter, FastFileWriterConfig -from deepspeed.runtime.swap_tensor.constants import * -from .constants import * -from deepspeed.accelerator import get_accelerator - - -class CheckpointWriterFactory(object): - - def __init__(self, writer_config, aio_config, dp_writer_config): - self._type = writer_config[CHECKPOINT_WRITER_TYPE] - self._io_buffer_size = writer_config[CHECKPOINT_IO_BUFFER_SIZE] - self._io_buffer_double = writer_config[CHECKPOINT_IO_BUFFER_DOUBLE] - self._data_parallel_writer = dp_writer_config - self._io_multiplier = writer_config[CHECKPOINT_IO_MULTIPLIER] - if self._data_parallel_writer.pure_dp: - self._show_statistics = writer_config[CHECKPOINT_IO_STATISTICS] and self._data_parallel_writer is not None - else: - self._show_statistics = writer_config[CHECKPOINT_IO_STATISTICS] and self._data_parallel_writer is not None - self._io_buffer = None - self._dnvme_handle = None - self._writer = None - self._use_gds = False - - if self._type == CheckpointWriterType.FAST: - self._use_gds = aio_config[AIO_USE_GDS] - if self._use_gds: - self._setup_for_gds(aio_config) - else: - self._setup_for_aio(aio_config) - print( - f'WriterFactory: self._data_parallel_writer={self._data_parallel_writer} self._show_statistics={self._show_statistics}' - ) - - def create_writer(self, file_path, optimize_dp_state): - assert self._writer is None, \ - f'Cannot create checkpoint writer for {file_path} because writer is currently used for {self._writer.file_path()}.\ - Must call writer.release() before reusing to avoid this error.' - - if self._type == CheckpointWriterType.MOCK: - self._writer = MockFileWriter(file_path) - elif self._type == CheckpointWriterType.PYTHON: - self._writer = PyFileWriter(file_path) - else: - if optimize_dp_state: - num_parallel_writers = self._data_parallel_writer.world_size * self._io_multiplier - writer_rank = self._data_parallel_writer.rank - file_path = f'{file_path}-{writer_rank}.{num_parallel_writers}' - # print(f'create_dp_writer: {self._data_parallel_writer.global_rank=} {writer_rank=} {num_parallel_writers=} {file_path=}') - else: - num_parallel_writers = 1 - writer_rank = 0 - # print(f'create_rank0_writer: {self._data_parallel_writer.global_rank=} {writer_rank=} {num_parallel_writers=} {file_path=}') - - config = FastFileWriterConfig(dnvme_handle=self._dnvme_handle, - pinned_tensor=self._io_buffer, - double_buffer=self._io_buffer_double, - num_parallel_writers=num_parallel_writers, - writer_rank=writer_rank, - global_rank=self._data_parallel_writer.global_rank) - self._writer = FastFileWriter(file_path=file_path, config=config) - - return self._writer - - def release_writer(self): - self._writer.close() - if self._show_statistics: - self._writer._dump_state() - self._writer = None - - def _setup_for_aio(self, aio_config): - self._io_buffer = torch.zeros(self._io_buffer_size, dtype=torch.uint8, device='cpu').pin_memory() - self._dnvme_handle = AsyncIOBuilder().load().aio_handle( - block_size=aio_config[AIO_BLOCK_SIZE], - queue_depth=aio_config[AIO_QUEUE_DEPTH], - single_submit=aio_config[AIO_SINGLE_SUBMIT], - overlap_events=aio_config[AIO_OVERLAP_EVENTS], - intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM]) - - def _setup_for_gds(self, aio_config): - self._io_buffer = torch.zeros(self._io_buffer_size, - dtype=torch.uint8, - device=get_accelerator().current_device_name()) - self._dnvme_handle = GDSBuilder().load().gds_handle(block_size=aio_config[AIO_BLOCK_SIZE], - queue_depth=aio_config[AIO_QUEUE_DEPTH], - single_submit=aio_config[AIO_SINGLE_SUBMIT], - overlap_events=aio_config[AIO_OVERLAP_EVENTS], - intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM]) - self._dnvme_handle.pin_device_tensor(self._io_buffer) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__init__.py deleted file mode 100644 index cf58eb1801af7cc02f3cc2057a5c843cd078b575..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .module import PipelineModule, LayerSpec, TiedLayerSpec -from .topology import ProcessTopology diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index de46cb67cfba49ca168ccf1763a60ce6f93ada8b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/engine.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/engine.cpython-310.pyc deleted file mode 100644 index 86ab8e38986732b58303dead4124039016424930..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/engine.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/module.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/module.cpython-310.pyc deleted file mode 100644 index 387ffcfdc213ae71153764368f28466d17df861e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/module.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/p2p.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/p2p.cpython-310.pyc deleted file mode 100644 index ff65b9693087591111f1e1c9008522e0388e1247..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/p2p.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/schedule.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/schedule.cpython-310.pyc deleted file mode 100644 index 4d950a390fc006d7f30a41945296cb438ec43b0f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/schedule.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/topology.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/topology.cpython-310.pyc deleted file mode 100644 index 60203c98b07edda9f58f655dc5544aa36bf064f3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/__pycache__/topology.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/engine.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/engine.py deleted file mode 100644 index 6ef00748bf1f3b74991173c390410761c460c73f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/engine.py +++ /dev/null @@ -1,1426 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from types import MethodType -from collections import OrderedDict -from functools import reduce -from operator import mul - -import torch -from deepspeed import comm as dist - -from deepspeed.utils import logger -from deepspeed.utils.timer import ThroughputTimer -from deepspeed.accelerator import get_accelerator -from deepspeed.runtime.bf16_optimizer import BF16_Optimizer - -from ..engine import DeepSpeedEngine, MEMORY_OPT_ALLREDUCE_SIZE -from deepspeed.utils.timer import FORWARD_MICRO_TIMER, FORWARD_GLOBAL_TIMER, BACKWARD_MICRO_TIMER, \ - BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_INNER_GLOBAL_TIMER, \ - BACKWARD_REDUCE_MICRO_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER, \ - STEP_MICRO_TIMER, STEP_GLOBAL_TIMER - -from ..utils import PartitionedTensor -from ..dataloader import RepeatingLoader -from ..zero.config import ZeroStageEnum -from ..activation_checkpointing import checkpointing as ds_checkpointing - -from .module import PipelineModule, PipelineError -from . import p2p -from . import schedule - -TARGET_ID = -2 -LOG_STAGE = -2 -DATA_PARALLEL_ID = -2 - -BATCH_INPUT_TIMER = 'batch_input' -TRAIN_BATCH_TIMER = 'train_batch' -PIPE_SEND_OUTPUT_TIMER = 'pipe_send_output' -PIPE_SEND_GRAD_TIMER = 'pipe_send_grad' -PIPE_RECV_INPUT_TIMER = 'pipe_recv_input' -PIPE_RECV_GRAD_TIMER = 'pipe_recv_grad' - -# The buffer size to store the meta data for each tensor. -TENSOR_META_SIZE = 256 - - -def is_even(number): - return number % 2 == 0 - - -mem_alloced = 0 -mem_cached = 0 - - -def _tensor_bytes(tensor): - return tensor.numel() * tensor.element_size() - - -class PipelineEngine(DeepSpeedEngine): - """ A training engine hybrid pipeline, data, and model parallel training. - - This engine is created by ``deepspeed.initialize()`` when a :class:`PipelineModule` - is provided. - """ - ID_TO_DTYPE = [ - torch.float32, torch.float64, torch.complex64, torch.complex128, torch.float16, torch.bfloat16, torch.uint8, - torch.int8, torch.int16, torch.int32, torch.int64, torch.bool - ] - DTYPE_TO_ID = {dtype: id_ for id_, dtype in enumerate(ID_TO_DTYPE)} - - def __init__(self, has_bool_tensors=False, *super_args, **super_kwargs): - super().__init__(*super_args, **super_kwargs) - assert isinstance(self.module, PipelineModule), "model must base PipelineModule" - - assert self.zero_optimization_stage( - ) < ZeroStageEnum.gradients, "ZeRO-2 and ZeRO-3 are incompatible with pipeline parallelism" - - # We schedule the all-reduces, so disable it in super().backward() - self.enable_backward_allreduce = False - self.has_bool_tensors = has_bool_tensors - self.eval_return_logits = False - self.outputs = None - # BF16 Optimizer is hardcoded for fp32 gradient accumulation - self.using_bf16_optimizer = type(self.optimizer) == BF16_Optimizer - - # used to disable the pipeline all-reduce when used with 1-bit Adam/1-bit LAMB - self.pipeline_enable_backward_allreduce = True - - if self.elasticity_enabled(): - if not self.is_elastic_model_parallel_supported(): - assert not self.elasticity_enabled(), "Elasticity is not currently supported" \ - " with pipeline parallelism." - - # pipeline step for logging - self.log_batch_step_id = -1 - - self.micro_batch_size = self.train_micro_batch_size_per_gpu() - self.micro_batches = self.gradient_accumulation_steps() - - # Set Grid and Communication Groups - self.grid = self.module._grid - if self.grid.get_global_rank() == 0: - logger.info(f'CONFIG: micro_batches={self.micro_batches} ' - f'micro_batch_size={self.micro_batch_size}') - - self.global_rank = self.grid.get_global_rank() - - assert self.dp_world_size == self.grid.data_parallel_size - assert self.train_batch_size() == \ - self.micro_batch_size * self.micro_batches * self.grid.data_parallel_size - - # Set Stage Inf - self.num_stages = self.grid.pipe_parallel_size - self.stage_id = self.grid.get_stage_id() - self.prev_stage = self.stage_id - 1 - self.next_stage = self.stage_id + 1 - - self.data_iterator = None - self.batch_fn = None - - self._force_grad_boundary = False - - self.batch_timer = ThroughputTimer(self._config.timers_config, - batch_size=self.train_batch_size(), - logging_fn=self.tput_log, - monitor_memory=False, - steps_per_output=self.steps_per_print()) - - # PipelineEngine needs to handle data loading specially due to only the first - # and last stages loading inputs/labels. We construct a sampler that uses - if self.training_data: - self._build_data_iter(self.training_data) - - self.is_pipe_parallel = self.grid.pipe_parallel_size > 1 - self.is_data_parallel = self.grid.data_parallel_size > 1 - self.is_model_parallel = self.grid.model_parallel_size > 1 - - # Partition input/output buffers - # XXX temporarily disable while I revert some partition hacks. - assert isinstance(self._config.pipeline['pipe_partitioned'], bool) - assert isinstance(self._config.pipeline['grad_partitioned'], bool) - self.is_pipe_partitioned = self.is_model_parallel and self._config.pipeline['pipe_partitioned'] - self.is_grad_partitioned = self.is_model_parallel and self._config.pipeline['grad_partitioned'] - logger.info(f'is_pipe_partitioned= {self.is_pipe_partitioned} ' - f'is_grad_partitioned= {self.is_grad_partitioned}') - - model_parameters = filter(lambda p: p.requires_grad, self.module.parameters()) - num_params = sum([p.numel() for p in model_parameters]) - unique_params = num_params - # Subtract tied parameters if we don't own them - if self.module.tied_comms: - tied_params = 0 - for key, d in self.module.tied_comms.items(): - if self.global_rank != min(d['ranks']): - tied_params += sum(p.numel() for p in d['module'].parameters()) - unique_params -= tied_params - params_tensor = torch.LongTensor(data=[num_params, unique_params]).to(self.device) - dist.all_reduce(params_tensor, group=self.grid.get_model_parallel_group()) - params_tensor = params_tensor.tolist() - total_params = params_tensor[0] - unique_params = params_tensor[1] - if self.grid.data_parallel_id == 0: - logger.info(f'RANK={self.global_rank} ' - f'STAGE={self.stage_id} ' - f'LAYERS={self.module._local_stop - self.module._local_start} ' - f'[{self.module._local_start}, {self.module._local_stop}) ' - f'STAGE_PARAMS={num_params} ({num_params/1e6:0.3f}M) ' - f'TOTAL_PARAMS={total_params} ({total_params/1e6:0.3f}M) ' - f'UNIQUE_PARAMS={unique_params} ({unique_params/1e6:0.3f}M)') - - #initialize peer-2-peer communication and allreduce groups - if self.is_pipe_parallel: - p2p.init_process_groups(self.grid) - - # Pipeline buffers - self.num_pipe_buffers = 0 - self.pipe_buffers = { - 'inputs': [], # batch input and received activations - 'labels': [], # labels from batch input - 'outputs': [], # activations - 'output_tensors': [], # tensor object to preserve backward graph - } - self.pipe_recv_buf = None - self.grad_layer = None - self._grad_layer_buf = [] - - self.meta_buffer = None - - self.first_output_send = True - self.first_gradient_send = True - self.pipe_partition_input_meta_cache = None - self.pipe_partition_output_meta_cache = None - self.pipe_partition_grad_meta_cache = None - self.grad_partition_grad_layer_meta_cache = None - - #stores the loss for the current micro batch being processed - self.loss = torch.tensor(0.0).to(self.device) - - #stores the loss for the entire batch - self.total_loss = None - self.total_additional_losses = None - self.agg_loss = torch.tensor(0.0, requires_grad=False).to(self.device) - self.dp_group_loss = torch.tensor(0.0, requires_grad=False).to(self.device) - - # stores aggregated-DP train final loss and aggregated-DP additional losses, if any - # additional losses are stored as dict: {loss-name: agg-loss} - self.agg_train_loss = None - self.agg_additional_losses = None - - if self._config.pipeline['activation_checkpoint_interval'] > 0: - self.module.activation_checkpoint_interval = self._config.pipeline['activation_checkpoint_interval'] - # set use_reentrant default to True. - if self._config.pipeline.get('use_reentrant') is None: - self._config.pipeline['use_reentrant'] = True - if self._config.pipeline['use_reentrant'] is False: - # set activation_checkpoint_func to non_reentrant_checkpoint func. - self.module.activation_checkpoint_func = ds_checkpointing.non_reentrant_checkpoint - if self.grid.get_global_rank() == 0: - logger.info(f'CONFIG: activation_checkpoint_func=non_reentrant_checkpoint') - if self.module.activation_checkpoint_interval > 0: - self.module._precompute_checkpointable_values() - - self.module.checkpoint_parallel_write_pipeline = self._config.checkpoint_parallel_write_pipeline - - if self.is_last_stage(): - self.loss_model = self.module.loss_fn - - self.has_attention_mask = self.module.__class__.__name__ == 'GPT2ModelPipe' - # Initialize pipeline communicators. Just send a 0. - if is_even(self.stage_id): - if not self.is_last_stage(): - p2p.send(self.loss, self.next_stage) - if not self.is_first_stage(): - p2p.recv(self.loss, self.prev_stage) - else: - if not self.is_first_stage(): - p2p.recv(self.loss, self.prev_stage) - if not self.is_last_stage(): - p2p.send(self.loss, self.next_stage) - - # XXX look into timer reporting timing - # Initialize some timers because of early weirdness. - if self.wall_clock_breakdown(): - self.timers(FORWARD_MICRO_TIMER).start() - self.timers(FORWARD_MICRO_TIMER).stop() - self.timers(BACKWARD_MICRO_TIMER).start() - self.timers(BACKWARD_MICRO_TIMER).stop() - self.timers(BACKWARD_INNER_MICRO_TIMER).start() - self.timers(BACKWARD_INNER_MICRO_TIMER).stop() - self.timers(BACKWARD_REDUCE_MICRO_TIMER).start() - self.timers(BACKWARD_REDUCE_MICRO_TIMER).stop() - self.timers(BACKWARD_REDUCE_GLOBAL_TIMER).start() - self.timers(BACKWARD_REDUCE_GLOBAL_TIMER).stop() - self.timers(STEP_MICRO_TIMER).start() - self.timers(STEP_MICRO_TIMER).stop() - - self.dynamic_shape = self.module.dynamic_shape - - def set_has_attention_mask(self, value): - assert isinstance(value, bool) - self.has_attention_mask = value - - def _build_data_iter(self, dataset): - sampler = torch.utils.data.distributed.DistributedSampler(dataset, - num_replicas=self.dp_world_size, - rank=self.mpu.get_data_parallel_rank(), - shuffle=False) - # Build a loader and make it repeating. - pipe_dataloader = self.deepspeed_io(dataset, data_sampler=sampler) - pipe_dataloader = RepeatingLoader(pipe_dataloader) - self.set_dataloader(pipe_dataloader) - - def _exec_reduce_tied_grads(self): - # We need to run this first to write to self.averaged_gradients; - # since this class turns `enable_backward_allreduce` off, - # `self.overlapping_partition_gradients_reduce_epilogue()` defined in the DeepSpeedEngine - # never actually runs. I suspect this is because of efficiency problems; get_flat_partition in - # stage2.py might do something expensive; someone will have to look into that later. But - # in the meantime, this fixes ZeRO2 + Pipelining enough to run a demo. Further profiling - # needed to decide if it actually breaks everything. - # (see https://github.com/EleutherAI/gpt-neox/issues/62#issuecomment-761471944) - if self.zero_optimization_partition_gradients(): - self.optimizer.overlapping_partition_gradients_reduce_epilogue() - - weight_group_list = self.module.get_tied_weights_and_groups() - for weight, group in weight_group_list: - grad = weight._hp_grad if self.using_bf16_optimizer else weight.grad - if grad is not None: - dist.all_reduce(grad, group=group) - - def _exec_reduce_grads(self): - self._force_grad_boundary = True - if self.pipeline_enable_backward_allreduce: - if self.using_bf16_optimizer: - # PP+BF16 work for ZeRO Stage 1 - self._bf16_reduce_grads() - else: - self.allreduce_gradients(bucket_size=MEMORY_OPT_ALLREDUCE_SIZE) - self._force_grad_boundary = False - - def _bf16_reduce_grads(self): - self.buffered_allreduce_fallback(grads=None, elements_per_buffer=MEMORY_OPT_ALLREDUCE_SIZE) - - def _reserve_pipe_buffers(self, num_buffers): - """Ensure that each pipeline buffer has at least ``num_buffers`` slots. - - This method only reserves slots and does not allocate tensors. - - Args: - num_buffers (int): The number of buffers to reserve. - """ - if self.num_pipe_buffers >= num_buffers: - return - - num_added = num_buffers - self.num_pipe_buffers - for key in self.pipe_buffers: - self.pipe_buffers[key].extend([None] * num_added) - self.num_pipe_buffers = num_buffers - - def reset_activation_shape(self): - """Reset the buffers when the shape of activation and gradient change. - For example, for curriculum learning that changes the seqlen of each - sample, we need to call this whenever the seqlen is going to change. - """ - self.first_output_send = True - self.pipe_recv_buf = None - self.grad_layer = None - self._grad_layer_buf = [] - self.meta_buffer = None - - self.pipe_partition_input_meta_cache = None - self.pipe_partition_output_meta_cache = None - self.pipe_partition_grad_meta_cache = None - self.grad_partition_grad_layer_meta_cache = None - - def train_batch(self, data_iter=None): - """Progress the pipeline to train the next batch of data. The engine will ingest - ``self.train_batch_size()`` total samples collectively across all workers. - - - An iterator that over training data should be provided as an argument - unless ``deepspeed.initialize()`` was provided a training set. In that event, - the training data will automatically be read. - - - .. warning:: - A total of ``self.gradient_accumulation_steps()`` entries will be pulled - from ``data_iter`` by each pipeline. There must be sufficient - data left in ``data_iter`` or else a ``StopIteration`` will halt training. - - DeepSpeed provides a convenience class :class:`deepspeed.utils.RepeatingLoader` - that wraps data loaders to automatically restart upon a ``StopIteration``. - - Args: - data_iter (Iterator, optional): Iterator of training data. - - Returns: - The arithmetic mean of the losses computed this batch. - """ - if not torch._C.is_grad_enabled(): - raise RuntimeError(f'train_batch() requires gradients enabled. Use eval_batch() instead.') - - # Curriculum learning could change activation shape - if self.curriculum_enabled_legacy(): - new_difficulty = self.curriculum_scheduler_legacy.update_difficulty( \ - self.global_steps + 1) - if self.global_steps == 0 or self.curriculum_scheduler_legacy.first_step: - self.reset_activation_shape() - self.curriculum_scheduler_legacy.first_step = False - elif new_difficulty != self.curriculum_scheduler_legacy.get_difficulty( \ - self.global_steps): - self.reset_activation_shape() - - if data_iter is not None: - self.set_dataiterator(data_iter) - - self.module.train() - self.total_loss = None - self.total_additional_losses = None - self._compute_loss = True - - # Do the work - self.timers(TRAIN_BATCH_TIMER).start() - sched = schedule.TrainSchedule(micro_batches=self.micro_batches, - stages=self.num_stages, - stage_id=self.stage_id) - self._exec_schedule(sched) - - with torch.no_grad(): - self.agg_train_loss = self._aggregate_total_loss() - - self.timers(TRAIN_BATCH_TIMER).stop() - - if self.steps_per_print() is not None and self.global_steps % self.steps_per_print() == 0: - if self.global_rank == 0: - elapsed = self.timers(TRAIN_BATCH_TIMER).elapsed(reset=True) / 1000.0 - iter_time = elapsed / self.steps_per_print() - tput = self.train_batch_size() / iter_time - log_str = f'steps: {self.global_steps} loss: {self.agg_train_loss:0.4f} ' - if self.agg_additional_losses is not None: - for loss_name, loss_value in self.agg_additional_losses.items(): - log_str += f'{loss_name}: {loss_value.item():0.4f} ' - log_str += f'iter time (s): {iter_time:0.3f} samples/sec: {tput:0.3f}' - print(log_str) - else: - self.timers(TRAIN_BATCH_TIMER).elapsed(reset=True) - - # Monitoring - if self.global_rank == 0 and self.monitor.enabled: - self.summary_events = [(f'Train/Samples/train_loss', self.agg_train_loss.mean().item(), - self.global_samples)] - self.monitor.write_events(self.summary_events) - - if self.steps_per_print() is not None and self.wall_clock_breakdown( - ) and self.global_steps % self.steps_per_print() == 0: - self.timers.log([ - PIPE_SEND_OUTPUT_TIMER, - PIPE_SEND_GRAD_TIMER, - PIPE_RECV_INPUT_TIMER, - PIPE_RECV_GRAD_TIMER, - ]) - - # TODO: should return precisely what loss returned and allow others to be queried? - return self.agg_train_loss - - def eval_batch(self, - data_iter, - return_logits=False, - compute_loss=True, - reduce_output='avg', - bcast_loss=True, - num_micro_batches=None): - """Evaluate the pipeline on a batch of data from ``data_iter``. The - engine will evaluate ``self.train_batch_size()`` total samples - collectively across all workers. - - This method is equivalent to: - - .. code-block:: python - - module.eval() - with torch.no_grad(): - output = module(batch) - - .. warning:: - A total of ``self.gradient_accumulation_steps()`` entries will be pulled - from ``data_iter`` by each pipeline. There must be sufficient - data left in ``data_iter`` or else a ``StopIteration`` will halt training. - - DeepSpeed provides a convenience class :class:`deepspeed.utils.RepeatingLoader` - that wraps data loaders to automatically restart upon a ``StopIteration``. - - Args: - data_iter (Iterator): Iterator of data to evaluate. - - Returns: - The arithmetic mean of the losses computed this batch. - """ - self.eval_return_logits = return_logits - self.module.eval() - - # Curriculum learning could change activation shape - if self.curriculum_enabled_legacy(): - new_difficulty = self.curriculum_scheduler_legacy.update_difficulty( \ - self.global_steps + 1) - if self.global_steps == 0 or self.curriculum_scheduler_legacy.first_step: - self.reset_activation_shape() - self.curriculum_scheduler_legacy.first_step = False - elif new_difficulty != self.curriculum_scheduler_legacy.get_difficulty( \ - self.global_steps): - self.reset_activation_shape() - - eval_output = None - - self._compute_loss = compute_loss - - # Use the provided data iterator - train_iterator = self.data_iterator - self.set_dataiterator(data_iter) - - # set the number micro batches in case the user chose value than training - micro_batches = self.micro_batches if num_micro_batches is None else num_micro_batches - - # Do the work - sched = schedule.InferenceSchedule(micro_batches=micro_batches, stages=self.num_stages, stage_id=self.stage_id) - - # prevent dead-lock with multiple evals sequence - dist.barrier() - - with torch.no_grad(): - self._exec_schedule(sched) - - if self.is_last_stage(): - eval_output = self._reduce_outputs(self.fwd_outputs, reduce=reduce_output, micro_batches=micro_batches) - - if compute_loss and (bcast_loss or self.monitor.enabled): - eval_output = self._bcast_pipe_scalar(eval_output) - - if self.global_rank == 0 and self.monitor.enabled: - self.summary_events = [(f'Train/Samples/eval_loss', eval_output.mean().item(), self.global_samples)] - self.monitor.write_events(self.summary_events) - - # Restore the training iterator - self.set_dataiterator(train_iterator) - - # Reset any buffers that may have been populated during the forward passes. - #ds_checkpointing.reset() - self.eval_return_logits = False - if return_logits: - outputs = self.outputs - self.outputs = None - return eval_output, outputs - return eval_output - - def set_train_batch_size(self, train_batch_size): - """Adjust the global batch size by increasing or decreasing the number of - micro-batches (i.e., gradient accumulation steps). The size of each micro-batch - (i.e., ``train_micro_batch_size_per_gpu``) is not changed. - Args: - train_batch_size (int): The new global batch size for training. - Raises: - ValueError: if ``train_batch_size`` is not divisible by the - configured micro-batch size and data parallelism. - """ - super().set_train_batch_size(train_batch_size) - self.micro_batches = self.gradient_accumulation_steps() - - def is_first_stage(self): - """True if this process is in the first stage in the pipeline.""" - return self.stage_id == 0 - - def is_last_stage(self): - """True if this process is in the last stage in the pipeline.""" - return self.stage_id == self.num_stages - 1 - - def _reduce_outputs(self, outputs, reduce='avg', reduce_dp=True, micro_batches=None): - if reduce is None: - return outputs - - if reduce.lower() == 'avg': - # first sum over all microbatches - if torch.is_tensor(outputs[0]): - reduced = sum(outputs) - else: - assert isinstance(outputs, (list, tuple)) - reduced = [torch.zeros_like(o) for o in outputs[0]] - for idx, out in outputs: - reduced[idx] += out - - # Average over the microbatches - reduced = self._scale_loss_by_gas(reduced, eval_micro_batches=micro_batches) - - # Average over DP groups - if reduce_dp and self.is_data_parallel: - if torch.is_tensor(reduced): - dist.all_reduce(reduced, group=self.mpu.get_data_parallel_group()) - reduced /= self.dp_world_size - else: - for idx in range(len(reduced)): - dist.all_reduce(reduced[idx], group=self.mpu.get_data_parallel_group()) - reduced[idx] /= self.dp_world_size - - return reduced - else: - raise NotImplementedError(f'reduction type {reduce} not supported.') - - def _bcast_pipe_scalar(self, data, src_rank=None, dtype=torch.float32): - # Default to last stage (e.g., for broadcasting loss) - if src_rank is None: - src_rank = self.grid.stage_to_global(self.num_stages - 1) - assert src_rank in self.grid.pp_group - - if self.global_rank == src_rank: - result = data.clone().detach().type(dtype).to(self.device) - else: - result = torch.Tensor([0.]).type(dtype).to(self.device) - - dist.broadcast(tensor=result, src=src_rank, group=self.mpu.get_pipe_parallel_group()) - - return result - - def _aggregate_total_loss(self): - # Scale loss, average among DP ranks, and bcast loss to the rest of my DP group - if self.is_last_stage(): - # Scale loss and additional losses, if any - loss = self._scale_loss_by_gas(self.total_loss) - self.agg_additional_losses = self.total_additional_losses - if self.agg_additional_losses is not None: - self.agg_additional_losses = OrderedDict({ - loss_name: self._scale_loss_by_gas(_loss.clone().detach()) - for loss_name, _loss in self.agg_additional_losses.items() - }) - - self.dp_group_loss = loss.clone().detach() - agg_loss = self.dp_group_loss.clone().detach() - #print(f'RANK={self.global_rank} bcast SENDER src={self.global_rank} group={self.grid.pp_group}', flush=True) - - # Average loss across all data-parallel groups - if self.is_data_parallel: - if self.agg_additional_losses is None: - dist.all_reduce(agg_loss, group=self.mpu.get_data_parallel_group()) - agg_loss /= self.dp_world_size - else: - # use a single reduce op for agg_loss and additional losses, if any - assert '__train_loss__' not in self.agg_additional_losses.keys() - tensors = OrderedDict({'__train_loss__': agg_loss}) - tensors.update(self.agg_additional_losses.items()) - flat_tensor = torch.cat([t.clone().reshape(-1).detach() for t in tensors.values()]) - dist.all_reduce(flat_tensor, group=self.mpu.get_data_parallel_group()) - flat_tensor /= self.dp_world_size - offset = 0 - reduced_tensor = {} - for name, t in tensors.items(): - n_elem = t.numel() - reduced_tensor[name] = flat_tensor[offset:offset + n_elem].clone().detach().reshape(t.shape) - offset += n_elem - agg_loss = reduced_tensor['__train_loss__'] - self.agg_additional_losses = OrderedDict( - {name: reduced_tensor[name] - for name in self.agg_additional_losses.keys()}) - - assert self.global_rank in self.grid.pp_group - losses = [self.dp_group_loss, agg_loss] - if self.agg_additional_losses is not None: - losses += list(self.agg_additional_losses.values()) - losses = torch.stack(losses).float() - if self.is_pipe_parallel: - dist.broadcast(tensor=losses, src=self.global_rank, group=self.mpu.get_pipe_parallel_group()) - else: - # Get loss from last stage - src_rank = self.grid.stage_to_global(self.num_stages - 1) - assert src_rank in self.grid.pp_group - # losses to reduce are: dp_group_loss, agg_loss, model additional losses - # therefore: 2 + n_additional_losses - additional_losses = self.module.get_additional_losses() - n_additional_losses = 0 if additional_losses is None else len(additional_losses) - losses = torch.Tensor([0.] * (2 + n_additional_losses)).to(self.device) - dist.broadcast(tensor=losses, src=src_rank, group=self.grid.get_pipe_parallel_group()) - self.dp_group_loss = losses[0].clone().detach() - agg_loss = losses[1].clone().detach() - if additional_losses is not None: - self.agg_additional_losses = OrderedDict({ - name: losses[2 + i].clone().detach() - for i, name in enumerate(additional_losses.keys()) - }) - return agg_loss - - def set_dataloader(self, loader): - """""" - if self.is_first_stage() or self.is_last_stage(): - self.training_dataloader = loader - self.data_iterator = iter(self.training_dataloader) - - def set_dataiterator(self, iterator): - """ Store an iterator to sample for training data. """ - if self.is_first_stage() or self.is_last_stage(): - self.training_dataloader = None - self.data_iterator = iterator - - def set_batch_fn(self, fn): - """Execute a post-processing function on input data. - - Args: - fn (function): The function to run. - """ - self.batch_fn = fn - - def is_gradient_accumulation_boundary(self): - """True if the engine is executing a gradient reduction or optimizer step instruction. - - This is overridden from :class:`DeepSpeedEngine` to force reductions - and steps when the pipeline engine is instructed to do so. - - Returns: - bool: whether reductions and optimizer steps should occur. - """ - return self._force_grad_boundary - - def log_for_device(self, *msg): - if LOG_STAGE == self.stage_id or LOG_STAGE == -1: - if DATA_PARALLEL_ID == self.grid.data_parallel_id or DATA_PARALLEL_ID == -1: - print( - f'RANK={dist.get_rank()} ' - f'PIPE-ID={self.stage_id} ' - f'DATA-ID={self.grid.data_parallel_id} ' - f'MBATCH-ID={self.microbatch_id} ' - f'STEP-ID={self.log_batch_step_id} ' - '::', - *msg, - flush=True) - - def tput_log(self, *msg): - if self.global_rank == 0 and self.global_steps % self.steps_per_print() == 0: - print(*msg) - - def _next_batch(self): - # If using 3D parallelism, only some first-stage ranks may do IO - batch = None - if self.data_iterator is not None: - batch = next(self.data_iterator) - - # Any post-processing, like broadcasting across a slice-parallel group. - if self.batch_fn: - batch = self.batch_fn(batch) - - return batch - - def _exec_forward_pass(self, buffer_id): - self.tput_timer.start() - self.mem_status('BEFORE FWD', reset_max=True) - - if isinstance(self.pipe_buffers['inputs'][buffer_id], tuple): - inputs = tuple(t.clone() for t in self.pipe_buffers['inputs'][buffer_id]) - else: - inputs = self.pipe_buffers['inputs'][buffer_id].clone() - - # collect the partitioned input from the previous stage - if self.is_pipe_partitioned and not self.is_first_stage(): - if self.pipe_partition_input_meta_cache is None: - self.pipe_partition_input_meta_cache = inputs[0].to('cpu') - part_input = PartitionedTensor.from_meta(meta=self.pipe_partition_input_meta_cache, - local_part=inputs[1], - group=self.grid.get_slice_parallel_group()) - - inputs = (part_input.full(), *inputs[2:]) - inputs[0].requires_grad = True - # skip mask - #inputs[1].requires_grad = True - part_input = None - inputs = inputs[0] if len(inputs) == 1 else inputs - self.pipe_buffers['inputs'][buffer_id] = inputs - - # inputs has no gradient because it is from a cloned tensor - outputs = super().forward(inputs) - - # Reset activation checkpointing buffers. - # Need to call this between evaluation iterations - if not self.module.training: - ds_checkpointing.reset() - - # Partition the outputs if we are not the last stage - if self.is_pipe_partitioned and not self.is_last_stage(): - if isinstance(outputs, tuple): - first_output = outputs[0] - # TODO: Improve pipe partitioning to pass multiple tensors that require grads - assert all([torch.is_tensor(elt) and elt.requires_grad is False for elt in outputs[1:]]) - outputs_tail = outputs[1:] - elif torch.is_tensor(outputs): - first_output = outputs - outputs_tail = [] - else: - raise ValueError("expecting a tensor or a tuple of tensors") - part = PartitionedTensor(tensor=first_output, group=self.grid.get_slice_parallel_group()) - # Clear the large output data, but save the computation graph - first_output.data = torch.zeros(1, device=first_output.data.device) - self.pipe_buffers['output_tensors'][buffer_id] = first_output - # Inject the partitioned tensor into the output before sending - outputs = (part.to_meta(), part.data(), *outputs_tail) - part = None - - self.pipe_buffers['outputs'][buffer_id] = outputs - - # Optionally compute loss on the last device - if self.is_last_stage(): - if self._compute_loss and self.module.loss_fn is not None: - labels = self.pipe_buffers['labels'][buffer_id] - self.loss = self.module.loss_fn(outputs, labels) - else: - # Some models just return loss from forward() - self.loss = outputs - if self.eval_return_logits: - self.outputs = outputs - - if isinstance(self.loss, torch.Tensor): - self.fwd_outputs.append(self.loss.detach()) - else: - self.fwd_outputs.append([l.detach() for l in self.loss]) - - def add_to_total_loss(_total_loss, _loss): - if isinstance(_loss, torch.Tensor): - if _total_loss is None: - _total_loss = torch.zeros_like(_loss) - _total_loss += _loss.detach() - else: - if _total_loss is None: - _total_loss = [torch.zeros_like(_l) for _l in _loss] - for _idx, _l in enumerate(_loss): - _total_loss[_idx] += _l.detach() - return _total_loss - - self.total_loss = add_to_total_loss(self.total_loss, self.loss) - - # aggregate additional losses across gradient accumulation steps - additional_losses = self.module.get_additional_losses() - if additional_losses is not None: - if self.total_additional_losses is None: - self.total_additional_losses = OrderedDict() - for name, loss in additional_losses.items(): - total = self.total_additional_losses[name] if name in self.total_additional_losses else None - self.total_additional_losses[name] = add_to_total_loss(total, loss) - - def _exec_backward_pass(self, buffer_id): - assert self.optimizer is not None, "must provide optimizer during " \ - "init in order to use backward" - - self.mem_status('BEFORE BWD', reset_max=True) - - # The last stage just runs backward on the loss using DeepSpeed's typical - # mechanisms. - if self.is_last_stage(): - super().backward(self.loss) - self.mem_status('AFTER BWD') - return - - outputs = self.pipe_buffers['outputs'][buffer_id] - - if self.wall_clock_breakdown(): - self.timers(BACKWARD_MICRO_TIMER).start() - self.timers(BACKWARD_GLOBAL_TIMER).start() - self.timers(BACKWARD_INNER_MICRO_TIMER).start() - self.timers(BACKWARD_INNER_GLOBAL_TIMER).start() - - # Reconstruct if we previously partitioned the output. We must be - # careful to also restore the computational graph of the tensors we partitioned. - if self.is_pipe_partitioned: - if self.is_grad_partitioned: - if self.pipe_partition_output_meta_cache is None: - self.pipe_partition_output_meta_cache = outputs[0].to('cpu') - part_output = PartitionedTensor.from_meta(meta=self.pipe_partition_output_meta_cache, - local_part=outputs[1], - group=self.grid.get_slice_parallel_group()) - self.pipe_buffers['output_tensors'][buffer_id].data = part_output.full() - outputs = (self.pipe_buffers['output_tensors'][buffer_id], *outputs[2:]) - else: - # Already restored from partition - self.pipe_buffers['output_tensors'][buffer_id].data = outputs[0] - outputs = (self.pipe_buffers['output_tensors'][buffer_id], *outputs[1:]) - - grad_tensors = self.grad_layer - if self.is_grad_partitioned: - #print(f'RANK={self.global_rank} BEFORE-BWD restoring grad={self.grad_layer[0].size()} {self.grad_layer[1].size()}') - if self.grad_partition_grad_layer_meta_cache is None: - self.grad_partition_grad_layer_meta_cache = self.grad_layer[0].to('cpu') - part_grad = PartitionedTensor.from_meta(meta=self.grad_partition_grad_layer_meta_cache, - local_part=self.grad_layer[1], - group=self.grid.get_slice_parallel_group()) - grad_tensors = (part_grad.full(), *grad_tensors[2:]) - part_grad = None - #print(f'RANK={self.global_rank} BEFORE-BWD restored grad={self.grad_layer[0].size()} {self.grad_layer[1].size()}') - - if self.using_bf16_optimizer and not self.is_last_stage(): - # manually call because we don't call optimizer.backward() - self.optimizer.clear_lp_grads() - - # This handles either a single tensor or tuple of tensors. - if isinstance(outputs, tuple): - out_tensors = [t for t in outputs if t.is_floating_point()] - assert len(out_tensors) == len(grad_tensors) - torch.autograd.backward(tensors=out_tensors, grad_tensors=grad_tensors) - else: - torch.autograd.backward(tensors=(outputs, ), grad_tensors=(grad_tensors, )) - - if self.using_bf16_optimizer and not self.is_last_stage(): - # manually call because we don't call optimizer.backward() - if not self._config.bfloat16_config.immediate_grad_update: - self.optimizer.update_hp_grads(clear_lp_grads=False) - - # Free up the memory from the output of forward() - self.pipe_buffers['output_tensors'][buffer_id] = None - self.pipe_buffers['outputs'][buffer_id] = None - grad_tensors = None - - if self.wall_clock_breakdown(): - self.timers(BACKWARD_INNER_MICRO_TIMER).stop() - self.timers(BACKWARD_INNER_GLOBAL_TIMER).stop() - self.timers(BACKWARD_MICRO_TIMER).stop() - self.timers(BACKWARD_GLOBAL_TIMER).stop() - - self.mem_status('AFTER BWD') - - def _exec_load_micro_batch(self, buffer_id): - if self.wall_clock_breakdown(): - self.timers(BATCH_INPUT_TIMER).start() - - batch = self._next_batch() - - if self.is_first_stage(): - loaded = None - if torch.is_tensor(batch[0]): - loaded = batch[0].clone().to(self.device).detach() - if self._config.pipeline['activation_checkpoint_interval'] > 0 and self._config.pipeline[ - 'use_reentrant']: - loaded.requires_grad = loaded.is_floating_point() - else: - assert isinstance(batch[0], (tuple, list)) - # Assume list or tuple - loaded = [] - for x in batch[0]: - assert torch.is_tensor(x) - mine = x.clone().detach().to(self.device) - if self._config.pipeline['activation_checkpoint_interval'] > 0 and self._config.pipeline[ - 'use_reentrant']: - mine.requires_grad = mine.is_floating_point() - loaded.append(mine) - loaded = tuple(loaded) - - self.pipe_buffers['inputs'][buffer_id] = loaded - - if self.is_last_stage(): - loaded = batch[1] - if torch.is_tensor(batch[1]): - loaded = batch[1].to(self.device) - # XXX: torch 1.6.0 DataLoader will auto convert tuple to list - elif isinstance(batch[1], (tuple, list)): - loaded = [] - for x in batch[1]: - assert torch.is_tensor(x) - x = x.to(self.device).detach() - loaded.append(x) - loaded = tuple(loaded) - - self.pipe_buffers['labels'][buffer_id] = loaded - - if self.wall_clock_breakdown(): - self.timers(BATCH_INPUT_TIMER).stop() - - def _send_tensor_meta(self, buffer, recv_stage): - """ Communicate metadata about upcoming p2p transfers. - - Metadata is communicated in this order: - * type (0: tensor, 1: list) - * num_tensors if type=list - foreach tensor in buffer: - * ndims - * shape - """ - meta_buffer = torch.empty(TENSOR_META_SIZE, dtype=torch.int32, device=self.device) - if isinstance(buffer, torch.Tensor): - meta_buf_list = [ - 0, # type of data (0: tensor, 1: list (unused), 2: tuple) - self.DTYPE_TO_ID[buffer.dtype], # dtype - len(buffer.size()) # ndims - ] - meta_buf_list.extend(buffer.size()) - assert len( - meta_buf_list - ) <= TENSOR_META_SIZE, f"Buffer for metadata is too small. Current buffer size: {TENSOR_META_SIZE} but required {len(meta_buf_list)}" - meta_buffer[:len(meta_buf_list)].copy_(torch.tensor(meta_buf_list, dtype=torch.int32)) - p2p.send(meta_buffer, recv_stage) - - elif isinstance(buffer, tuple): - meta_buf_list = [ - 2, # type of data (0: tensor, 1: list (unused), 2: tuple) - len(buffer) # num_tensors - ] - - for tensor in buffer: - assert isinstance(tensor, torch.Tensor) - meta_buf_list.append(self.DTYPE_TO_ID[tensor.dtype]) - meta_buf_list.append(len(tensor.size())) - meta_buf_list.extend(tensor.size()) - - assert len( - meta_buf_list - ) <= TENSOR_META_SIZE, f"Buffer for metadata is too small. Current buffer size: {TENSOR_META_SIZE} but required {len(meta_buf_list)}" - meta_buffer[:len(meta_buf_list)].copy_(torch.tensor(meta_buf_list, dtype=torch.int32)) - p2p.send(meta_buffer, recv_stage) - - else: - raise NotImplementedError(f'Could not send meta type {type(buffer)}') - - # Useful for performance debugging. - ''' - if self.grid.data_parallel_id == 0: - print(f'STAGE={self.stage_id} pipe-send-volume: {send_bytes/1024**2:0.2f}MB') - ''' - - def _recv_tensor_meta(self, send_stage): - """Receive metadata about upcoming p2p transfers and return allocated buffers. - - Returns: - Allocated buffer for receiving from send_stage. - """ - buffer = torch.empty(TENSOR_META_SIZE, dtype=torch.int32, device=self.device) - p2p.recv(buffer, send_stage) - - recv_type = buffer[0].item() - - # A single tensor will be sent. - if recv_type == 0: - recv_dtype = self.ID_TO_DTYPE[buffer[1].item()] - recv_ndims = buffer[2].item() - recv_shape = buffer[3:3 + recv_ndims].tolist() - return self._allocate_or_extend_buffers(0, recv_shape, recv_dtype) - - # List or tuple of tensors (recv_type == 1 (list) is currently unused) - elif recv_type == 1 or recv_type == 2: - num_tensors = buffer[1].item() - - buffers = [] - offset = 2 - for idx in range(num_tensors): - recv_dtype = self.ID_TO_DTYPE[buffer[offset].item()] - recv_ndims = buffer[offset + 1].item() - recv_shape = buffer[offset + 2:offset + 2 + recv_ndims].tolist() - offset += 2 + recv_ndims - - buffers.append(self._allocate_or_extend_buffers(idx, recv_shape, recv_dtype)) - - # Convert to tuples if requested. - if recv_type == 2: - buffers = tuple(buffers) - return buffers - - else: - raise NotImplementedError(f'Could not receive type {type(recv_type)}') - - def _exec_send_activations(self, buffer_id): - if self.wall_clock_breakdown(): - self.timers(PIPE_SEND_OUTPUT_TIMER).start() - - outputs = self.pipe_buffers['outputs'][buffer_id] - - # NCCL does not like to send torch.BoolTensor types, so cast the mask to half(). - # We could do char, but with half() we can eventually flatten with other fp16 - # messages (TODO) - if self.has_attention_mask or self.has_bool_tensors: - outputs = list(outputs) - outputs[-1] = outputs[-1].half() - outputs = tuple(outputs) - - if self.dynamic_shape or self.first_output_send: - self.first_output_send = False - self._send_tensor_meta(outputs, self.next_stage) - - if isinstance(outputs, torch.Tensor): - p2p.send(outputs, self.next_stage) - elif isinstance(outputs, tuple): - for idx, buffer in enumerate(outputs): - p2p.send(buffer, self.next_stage) - else: - raise NotImplementedError('Could not send output of type ' - f'{type(outputs)}') - - # Restore the boolean tensor - if self.has_attention_mask or self.has_bool_tensors: - outputs = list(outputs) - outputs[-1] = outputs[-1].bool() - outputs = tuple(outputs) - - if self.wall_clock_breakdown(): - self.timers(PIPE_SEND_OUTPUT_TIMER).stop() - - def _exec_send_grads(self, buffer_id): - if self.wall_clock_breakdown(): - self.timers(PIPE_SEND_GRAD_TIMER).start() - - inputs = self.pipe_buffers['inputs'][buffer_id] - - # Partition the gradient - if self.is_grad_partitioned: - if isinstance(inputs, tuple): - first_input = inputs[0] - assert all([torch.is_tensor(elt) for elt in inputs[1:]]) - inputs_grad_tail = [elt.grad for elt in inputs[1:]] - elif torch.is_tensor(inputs): - first_input = inputs - inputs_grad_tail = [] - else: - raise ValueError("expecting a tensor or a tuple of tensors") - assert torch.is_tensor(first_input) - part = PartitionedTensor(tensor=first_input.grad, group=self.grid.get_slice_parallel_group()) - - inputs = (part.to_meta(), part.data(), *inputs_grad_tail) - - # XXX Terrible hack - # Drop the attention mask from the input buffer here. It does not have - # a grad that needs to be communicated. We free the buffer immediately - # after, so no need to restore it. The receiver also has a hack that skips - # the recv. This is because NCCL does not let us send torch.BoolTensor :-(. - if self.has_attention_mask or self.has_bool_tensors: - inputs = list(inputs) - inputs.pop() - inputs = tuple(inputs) - - if isinstance(inputs, torch.Tensor): - assert inputs.grad is not None - p2p.send(inputs.grad, self.prev_stage) - else: - # XXX terrible hacky branch - if self.is_grad_partitioned: - # First two sends are partitioned gradient - p2p.send(inputs[0], self.prev_stage) - p2p.send(inputs[1], self.prev_stage) - else: - for idx, buffer in enumerate(inputs): - # Skip tensors that will not produce a grad - if not buffer.is_floating_point(): - assert buffer.grad is None - continue - assert buffer.grad is not None - p2p.send(buffer.grad, self.prev_stage) - - # We can free up the input buffer now - self.pipe_buffers['inputs'][buffer_id] = None - - if self.wall_clock_breakdown(): - self.timers(PIPE_SEND_GRAD_TIMER).stop() - - def _exec_recv_activations(self, buffer_id): - if self.wall_clock_breakdown(): - self.timers(PIPE_RECV_INPUT_TIMER).start() - - recvd = None - - # Allocate the buffer if necessary - if self.dynamic_shape or self.pipe_recv_buf is None: - self.pipe_recv_buf = self._recv_tensor_meta(self.prev_stage) - - if isinstance(self.pipe_recv_buf, torch.Tensor): - p2p.recv(self.pipe_recv_buf, self.prev_stage) - recvd = self.pipe_recv_buf.clone().detach() - recvd.requires_grad = recvd.is_floating_point() - else: - assert isinstance(self.pipe_recv_buf, tuple) - recvd = [None] * len(self.pipe_recv_buf) - for idx, buffer in enumerate(self.pipe_recv_buf): - assert torch.is_tensor(buffer) - # XXX hardcode meta type - if self.is_pipe_partitioned and idx == 0 and buffer.dtype != torch.long: - if self.meta_buffer is None: - self.meta_buffer = torch.zeros(buffer.size(), dtype=torch.long, device=self.device) - buffer = self.meta_buffer - - p2p.recv(buffer, self.prev_stage) - recvd[idx] = buffer.clone().detach() - - # NCCL does not like to send torch.BoolTensor types, so un-cast the - # attention mask - if self.has_attention_mask or self.has_bool_tensors: - recvd[-1] = recvd[-1].bool() - - recvd = tuple(recvd) - - for buffer in recvd: - buffer.requires_grad = buffer.is_floating_point() - - self.pipe_buffers['inputs'][buffer_id] = recvd - - if self.wall_clock_breakdown(): - self.timers(PIPE_RECV_INPUT_TIMER).stop() - - def _exec_recv_grads(self, buffer_id): - if self.wall_clock_breakdown(): - self.timers(PIPE_RECV_GRAD_TIMER).start() - - outputs = self.pipe_buffers['outputs'][buffer_id] - # XXX these shapes are hardcoded for Megatron - # Restore partitioned output if it was partitioned and we are sending full gradients - if self.is_pipe_partitioned and not self.is_grad_partitioned: - if self.pipe_partition_grad_meta_cache is None: - self.pipe_partition_grad_meta_cache = outputs[0].to('cpu') - part_output = PartitionedTensor.from_meta(meta=self.pipe_partition_grad_meta_cache, - local_part=outputs[1], - group=self.grid.get_slice_parallel_group()) - outputs[0].data = part_output.full() - outputs = (outputs[0], *outputs[2:]) - # save for backward - self.pipe_buffers['outputs'][buffer_id] = outputs - - # Allocate gradient if necessary - if self.dynamic_shape or self.grad_layer is None: - if isinstance(outputs, torch.Tensor): - self.grad_layer = self._allocate_or_extend_buffers(0, list(outputs.size()), outputs.dtype) - else: - # XXX This is a HACK - # When we exchange activations/gradients, the two pipe stages - # need to issue the send/recv with the same buffer sizes or - # else there is a deadlock. The is_floating_point() filter is - # used to avoid sending gradients for tensors that do not - # produce gradients. When TP>1, we partition the first - # activations/gradients across TP ranks to save communication - # volume and memory. That partitioned tensor is represented as - # two tensors: a 1/TPth chunk of the original data and also a - # small LongTensor storing the metadata used to reconstruct on - # the other side. When combined, the floating point filter also - # filtered out the metadata tensor. This quick (hacky) fix just - # branches on is_grad_partitioned so we don't filter out the - # metadata tensor. - if self.is_grad_partitioned: - sizes_and_dtypes = [(list(t.size()), t.dtype) - for t in outputs[:2]] + [(list(t.size()), t.dtype) - for t in outputs[2:] if t.is_floating_point()] - else: - sizes_and_dtypes = [(list(t.size()), t.dtype) for t in outputs if t.is_floating_point()] - - self.grad_layer = [ - self._allocate_or_extend_buffers(i, size, dtype) - for i, (size, dtype) in enumerate(sizes_and_dtypes) - ] - - if isinstance(self.grad_layer, torch.Tensor): - p2p.recv(self.grad_layer, self.next_stage) - else: - assert isinstance(outputs, tuple) - for idx, buffer in enumerate(self.grad_layer): - # XXX GPT-2 hack - if self.is_grad_partitioned and idx == 0 and buffer.dtype != torch.long: - buffer.data = torch.zeros(buffer.size(), dtype=torch.long, device=self.device) - p2p.recv(buffer, self.next_stage) - - if self.wall_clock_breakdown(): - self.timers(PIPE_RECV_GRAD_TIMER).stop() - - def _exec_optimizer_step(self, lr_kwargs=None): - if self.wall_clock_breakdown(): - self.timers(STEP_MICRO_TIMER).start() - self.timers(STEP_GLOBAL_TIMER).start() - self.mem_status('BEFORE STEP', reset_max=True) - - self._force_grad_boundary = True - self._take_model_step(lr_kwargs) - self._force_grad_boundary = False - - self.mem_status('AFTER STEP') - - if self.global_rank == 0 and self.monitor.enabled: - self.summary_events = [(f'Train/Samples/lr', self.get_lr()[0], self.global_samples)] - if self.fp16_enabled() and hasattr(self.optimizer, 'cur_scale'): - self.summary_events.append( - (f'Train/Samples/loss_scale', self.optimizer.cur_scale, self.global_samples)) - self.monitor.write_events(self.summary_events) - - if self.wall_clock_breakdown(): - self.timers(STEP_MICRO_TIMER).stop() - self.timers(STEP_GLOBAL_TIMER).stop() - if self.global_steps % self.steps_per_print() == 0: - self.timers.log([ - BATCH_INPUT_TIMER, - FORWARD_MICRO_TIMER, - BACKWARD_MICRO_TIMER, - BACKWARD_INNER_MICRO_TIMER, - BACKWARD_REDUCE_MICRO_TIMER, - STEP_MICRO_TIMER, - ]) - if self.global_steps % self.steps_per_print() == 0: - self.timers.log([ - FORWARD_GLOBAL_TIMER, - BACKWARD_GLOBAL_TIMER, - BACKWARD_INNER_GLOBAL_TIMER, - BACKWARD_REDUCE_GLOBAL_TIMER, - STEP_GLOBAL_TIMER, - ]) - - def _allocate_zeros(self, shape, **kwargs): - """ Allocate a tensor of zeros on the engine's device. - - Arguments: - shape: the shape of the tensor to allocate - kwargs: passed to torch.zeros() - - Returns: - A tensor from torch.zeros() allocated on self.device. - """ - if "dtype" not in kwargs: - if self.fp16_enabled(): - kwargs["dtype"] = torch.half - if self.bfloat16_enabled(): - kwargs["dtype"] = torch.bfloat16 - - return torch.zeros(shape, device=self.device, **kwargs) - - def _allocate_buffer(self, shape, num_buffers=-1, **kwargs): - buffers = [] - if num_buffers == -1: - num_buffers = self.num_pipe_buffers - for count in range(num_buffers): - buffers.append(self._allocate_zeros(shape, **kwargs)) - return buffers - - def _allocate_or_extend_buffers(self, idx, shape, dtype): - numel = reduce(mul, shape) if len(shape) > 0 else 1 - if len(self._grad_layer_buf) <= idx or self._grad_layer_buf[idx].numel() < numel: - new_buf = self._allocate_buffer(shape, dtype=dtype, num_buffers=1)[0] - if len(self._grad_layer_buf) <= idx: - self._grad_layer_buf.append(new_buf) - else: - self._grad_layer_buf[idx] = new_buf - return self._grad_layer_buf[idx] - else: - return self._grad_layer_buf[idx].flatten()[:numel].view(shape) - - def forward(self, *args, **kwargs): - """Disabled for pipeline parallel training. See ``train_batch()``. """ - raise PipelineError("Only train_batch() is accessible in pipeline mode.") - - def backward(self, *args, **kwargs): - """Disabled for pipeline parallel training. See ``train_batch()``. """ - raise PipelineError("Only train_batch() is accessible in pipeline mode.") - - def step(self, *args, **kwargs): - """Disabled for pipeline parallel training. See ``train_batch()``. """ - raise PipelineError("Only train_batch() is accessible in pipeline mode.") - - def mem_status(self, msg, print_rank=-1, reset_max=False): - return - global mem_alloced, mem_cached - if not self.global_steps == 0 or not self.global_steps == 9: - #return - pass - if self.mpu.get_data_parallel_rank() != 0: - return - - if self.global_rank != 0: - return - - rank = self.global_rank - if print_rank != -1 and rank != print_rank: - return - - get_accelerator().synchronize() - - if reset_max: - get_accelerator().reset_max_memory_cached() - get_accelerator().reset_max_memory_allocated() - - new_alloced = get_accelerator().memory_allocated() - new_cached = get_accelerator().memory_cached() - - delta_alloced = new_alloced - mem_alloced - delta_cached = new_cached - mem_cached - - mem_cached = new_cached - mem_alloced = new_alloced - - max_alloced = get_accelerator().max_memory_allocated() - max_cached = get_accelerator().max_memory_cached() - - # convert to GB for printing - new_alloced /= 1024**3 - new_cached /= 1024**3 - delta_alloced /= 1024**3 - delta_cached /= 1024**3 - max_alloced /= 1024**3 - max_cached /= 1024**3 - - print( - f'RANK={rank} STAGE={self.stage_id} STEP={self.global_steps} MEMSTATS', msg, - f'current alloc={new_alloced:0.4f}GB (delta={delta_alloced:0.4f}GB max={max_alloced:0.4f}GB) ' - f'current cache={new_cached:0.4f}GB (delta={delta_cached:0.4f}GB max={max_cached:0.4f}GB)') - - def module_state_dict(self, exclude_frozen_parameters=False): - """Override hack to save a pipe model and return the directory path of the save. - - This method should only be called by DeepSpeed's ``save_checkpoint()``. The - recommended way of saving a ``PipelineModule`` outside of ``save_checkpoint()`` - is ``save_state_dict()``. - - Returns: - None - """ - assert isinstance(self.module, PipelineModule) - assert self._curr_ckpt_path is not None, \ - "PipelineEngine expects module_state_dict() to be called from save_checkpoint()" - - self.module.save_state_dict(self._curr_ckpt_path, - checkpoint_engine=self.checkpoint_engine, - exclude_frozen_params=exclude_frozen_parameters) - return None - - def load_module_state_dict(self, checkpoint, strict=True, custom_load_fn=None, fetch_z3_params=False): - """Override hack to instead use a directory path. - - This is important because pipeline models checkpoint by layer instead of rank. - - If ``state_dict`` is not ``None`` or a ``str``, we revert to ``super()`` expecting a ``dict``. - - Args: - state_dict (str, None): unused - strict (bool, optional): Strict state loading. Defaults to True. - """ - assert custom_load_fn is None, "custom_load_fn not supported w. pipeline parallelism" - state_dict = checkpoint if self.has_moe_layers else checkpoint['module'] - if (state_dict is not None) and (not isinstance(state_dict, str)): - super().load_module_state_dict(state_dict, strict) - return - - self.module.load_state_dir(load_dir=self._curr_ckpt_path, - strict=strict, - checkpoint_engine=self.checkpoint_engine) - - # A map of PipeInstruction types to methods. Each method will be executed with the - # kwargs provided to the PipeInstruction from the scheduler. - _INSTRUCTION_MAP = { - schedule.OptimizerStep: _exec_optimizer_step, - schedule.ReduceGrads: _exec_reduce_grads, - schedule.ReduceTiedGrads: _exec_reduce_tied_grads, - schedule.LoadMicroBatch: _exec_load_micro_batch, - schedule.ForwardPass: _exec_forward_pass, - schedule.BackwardPass: _exec_backward_pass, - schedule.SendActivation: _exec_send_activations, - schedule.RecvActivation: _exec_recv_activations, - schedule.SendGrad: _exec_send_grads, - schedule.RecvGrad: _exec_recv_grads, - } - - def _exec_schedule(self, pipe_schedule): - # Reserve and reset buffers. - self._reserve_pipe_buffers(pipe_schedule.num_pipe_buffers()) - self.fwd_outputs = [] - - # For each step in the schedule - for step_cmds in pipe_schedule: - # For each instruction in the step - for cmd in step_cmds: - if type(cmd) not in self._INSTRUCTION_MAP: - raise RuntimeError(f'{self.__class__.__name__} does not understand instruction {repr(cmd)}') - - # Equivalent to: self._exec_forward_pass(buffer_id=0) - self._exec_instr = MethodType(self._INSTRUCTION_MAP[type(cmd)], self) - self._exec_instr(**cmd.kwargs) - - def get_additional_losses(self): - return self.agg_additional_losses diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/module.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/module.py deleted file mode 100644 index ad2463f0769f66fb597d80e82b1a811f80b12439..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/module.py +++ /dev/null @@ -1,695 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -import glob - -import re as regex - -from functools import partial - -import torch -import torch.nn as nn -from deepspeed import comm as dist - -from deepspeed.utils import logger -from .. import utils as ds_utils -from ..activation_checkpointing import checkpointing -from .topology import PipeDataParallelTopology, PipelineParallelGrid -from deepspeed.runtime.state_dict_factory import SDLoaderFactory -from deepspeed.accelerator import get_accelerator -from deepspeed.checkpoint.utils import clone_tensors_for_torch_save - - -class PipelineError(Exception): - """Errors related to the use of deepspeed.PipelineModule """ - - -class LayerSpec: - """Building block for specifying pipeline-parallel modules. - - LayerSpec stores the type information and parameters for each stage in a - PipelineModule. For example: - - .. code-block:: python - - nn.Sequence( - torch.nn.Linear(self.in_dim, self.hidden_dim, bias=False), - torch.nn.Linear(self.hidden_hidden, self.out_dim) - ) - - becomes - - .. code-block:: python - - layer_specs = [ - LayerSpec(torch.nn.Linear, self.in_dim, self.hidden_dim, bias=False), - LayerSpec(torch.nn.Linear, self.hidden_hidden, self.out_dim)] - ] - """ - - def __init__(self, typename, *module_args, **module_kwargs): - self.typename = typename - self.module_args = module_args - self.module_kwargs = module_kwargs - - if not issubclass(typename, nn.Module): - raise RuntimeError('LayerSpec only supports torch.nn.Module types.') - - if dist.is_initialized(): - self.global_rank = dist.get_rank() - else: - self.global_rank = -1 - - def __repr__(self): - return ds_utils.call_to_str(self.typename.__name__, self.module_args, self.module_kwargs) - - def build(self, log=False): - """Build the stored specification.""" - if log: - logger.info(f'RANK={self.global_rank} building {repr(self)}') - - return self.typename(*self.module_args, **self.module_kwargs) - - -class TiedLayerSpec(LayerSpec): - - def __init__(self, key, typename, *module_args, forward_fn=None, tied_weight_attr=['weight'], **module_kwargs): - super().__init__(typename, *module_args, **module_kwargs) - self.key = key - self.forward_fn = forward_fn - self.tied_weight_attr = [tied_weight_attr] if type(tied_weight_attr) == str else tied_weight_attr - - -class PipelineModule(nn.Module): - """Modules to be parallelized with pipeline parallelism. - - The key constraint that enables pipeline parallelism is the - representation of the forward pass as a sequence of layers - and the enforcement of a simple interface between them. The - forward pass is implicitly defined by the module ``layers``. The key - assumption is that the output of each layer can be directly fed as - input to the next, like a ``torch.nn.Sequence``. The forward pass is - implicitly: - - .. code-block:: python - - def forward(self, inputs): - x = inputs - for layer in self.layers: - x = layer(x) - return x - - .. note:: - Pipeline parallelism is not compatible with ZeRO-2 and ZeRO-3. - - Args: - layers (Iterable): A sequence of layers defining pipeline structure. Can be a ``torch.nn.Sequential`` module. - num_stages (int, optional): The degree of pipeline parallelism. If not specified, ``topology`` must be provided. - topology (``deepspeed.runtime.pipe.ProcessTopology``, optional): Defines the axes of parallelism axes for training. Must be provided if ``num_stages`` is ``None``. - loss_fn (callable, optional): Loss is computed ``loss = loss_fn(outputs, label)`` - seed_layers(bool, optional): Use a different seed for each layer. Defaults to False. - seed_fn(type, optional): The custom seed generating function. Defaults to random seed generator. - base_seed (int, optional): The starting seed. Defaults to 1234. - partition_method (str, optional): The method upon which the layers are partitioned. Defaults to 'parameters'. - activation_checkpoint_interval (int, optional): The granularity activation checkpointing in terms of number of layers. 0 disables activation checkpointing. - activation_checkpoint_func (callable, optional): The function to use for activation checkpointing. Defaults to ``deepspeed.checkpointing.checkpoint``. - checkpointable_layers (list[str], optional): List of layer class names that are eligible for checkpointing. For GPT models, - ParallelTransformerLayerPipe is always checkpointed regardless of this list. If None, all layers with parameters are - considered checkpointable. Defaults to None. - dynamic_shape: Allows dynamic shapes of inputs. This might have a performance impact. - """ - - def __init__(self, - layers, - num_stages=None, - topology=None, - loss_fn=None, - seed_layers=False, - seed_fn=None, - base_seed=1234, - partition_method='parameters', - activation_checkpoint_interval=0, - activation_checkpoint_func=checkpointing.checkpoint, - checkpointable_layers=None, - dynamic_shape=False): - - super().__init__() - - if num_stages is None and topology is None: - raise RuntimeError('must provide num_stages or topology') - - self.micro_offset = 0 - - self.loss_fn = loss_fn - - self.checkpointable_layers = checkpointable_layers - if checkpointable_layers is not None: - assert isinstance(checkpointable_layers, list), "param `checkpointable_layers` must be type of list." - - self.seed_layers = seed_layers - self.seed_fn = seed_fn - self.base_seed = base_seed - if dist.get_rank() == 0: - try: - seed_str = self.seed_fn.__name__ - except AttributeError: - seed_str = None - print(f'SEED_LAYERS={self.seed_layers} BASE_SEED={self.base_seed} SEED_FN={seed_str}') - - # Setup world info - self.world_group = dist.new_group(ranks=range(dist.get_world_size())) - self.global_rank = dist.get_rank(group=self.world_group) - self.world_size = dist.get_world_size(group=self.world_group) - self.local_rank = int(os.environ.get("LOCAL_RANK", None)) - assert self.local_rank is not None - - if topology: - self._topo = topology - self.num_stages = self._topo.get_dim('pipe') - else: - self.num_stages = num_stages - if topology is None: - if self.world_size % self.num_stages != 0: - raise RuntimeError( - f'num_stages ({self.num_stages}) must divide distributed world size ({self.world_size})') - dp = self.world_size // num_stages - topology = PipeDataParallelTopology(num_pp=num_stages, num_dp=dp) - self._topo = topology - - # Construct communicators for pipeline topology - self._grid = PipelineParallelGrid(process_group=self.world_group, topology=self._topo) - - self.stage_id = self._topo.get_coord(self.global_rank).pipe - - # Initialize partition information - self._layer_specs = list(layers) - self._num_layers = len(self._layer_specs) - self._local_start = 0 - self._local_stop = None - self._partition_layers(method=partition_method) - - self.forward_funcs = [] - self.fwd_map = {} - self.tied_modules = nn.ModuleDict() - self.tied_weight_attrs = {} - - # Offset the random seed by the stage ID. - #newseed = get_accelerator().initial_seed() + self._grid.get_stage_id() - #ds_utils.set_random_seed(newseed) - - self.activation_checkpoint_interval = activation_checkpoint_interval - - self.activation_checkpoint_func = activation_checkpoint_func - - #storage for precomputed checkpointeble results - self.is_checkpointable_results = [] - self.is_checkpointable_results_interval = None - - # if configuration use_reentrant = False, self.activation_checkpoint_func will be set to ``checkpointing.non_reentrant_checkpoint`` - - #with torch.random.fork_rng(devices=[get_accelerator().current_device_name()]): - self._build() - self.to(get_accelerator().device_name(self.local_rank)) - - self.tied_comms = self._index_tied_modules() - self._synchronize_tied_weights() - - self.dynamic_shape = dynamic_shape - - def _precompute_checkpointable_values(self): - if self.activation_checkpoint_interval > 0 and self.is_checkpointable_results_interval != self.activation_checkpoint_interval: - num_layers = len(self.forward_funcs) - self.interval_was_zero = False - for start_idx in range(0, num_layers, self.activation_checkpoint_interval): - end_idx = min(start_idx + self.activation_checkpoint_interval, num_layers) - funcs = self.forward_funcs[start_idx:end_idx] - self.is_checkpointable_results.append(self._is_checkpointable(funcs)) - self.is_checkpointable_results_interval = self.activation_checkpoint_interval - - def _build(self): - specs = self._layer_specs - - for local_idx, layer in enumerate(specs[self._local_start:self._local_stop]): - layer_idx = local_idx + self._local_start - if self.seed_layers: - if self.seed_fn: - self.seed_fn(self.base_seed + layer_idx) - else: - ds_utils.set_random_seed(self.base_seed + layer_idx) - - # Recursively build PipelineModule objects - if isinstance(layer, PipelineModule): - raise NotImplementedError('RECURSIVE BUILD NOT YET IMPLEMENTED') - - # LayerSpec objects contain an nn.Module that should be allocated now. - elif isinstance(layer, nn.Module): - name = str(layer_idx) - self.forward_funcs.append(layer) - self.fwd_map.update({name: len(self.forward_funcs) - 1}) - self.add_module(name, layer) - - # TiedLayerSpec objects contain an nn.Module that should be allocated now. - elif isinstance(layer, TiedLayerSpec): - # Build and register the module if we haven't seen it before. - if layer.key not in self.tied_modules: - self.tied_modules[layer.key] = layer.build() - self.tied_weight_attrs[layer.key] = layer.tied_weight_attr - - if layer.forward_fn is None: - # Just use forward() - self.forward_funcs.append(self.tied_modules[layer.key]) - else: - # User specified fn with args (module, input) - self.forward_funcs.append(partial(layer.forward_fn, self.tied_modules[layer.key])) - - # LayerSpec objects contain an nn.Module that should be allocated now. - elif isinstance(layer, LayerSpec): - module = layer.build() - name = str(layer_idx) - self.forward_funcs.append(module) - self.fwd_map.update({name: len(self.forward_funcs) - 1}) - self.add_module(name, module) - - # Last option: layer may be a functional (e.g., lambda). We do nothing in - # that case and just use it in forward() - else: - self.forward_funcs.append(layer) - - # All pipeline parameters should be considered as model parallel in the context - # of our FP16 optimizer - for p in self.parameters(): - p.ds_pipe_replicated = False - - def _get_frozen_parameter_names(self, layer): - """ Get names of frozen parameters in the layer. - - Returns: - A list of frozen parameter names - """ - if isinstance(layer, LayerSpec): - l = layer.build() - return [n for n, p in l.named_parameters() if not p.requires_grad] - elif isinstance(layer, nn.Module): - return [n for n, p in layer.named_parameters() if not p.requires_grad] - - return [] - - def _count_layer_params(self): - """Count the trainable parameters in individual layers. - - This routine will only build one layer at a time. - - Returns: - A list of the number of parameters in each layer. - """ - param_counts = [0] * len(self._layer_specs) - for idx, layer in enumerate(self._layer_specs): - if isinstance(layer, LayerSpec): - l = layer.build() - params = filter(lambda p: p.requires_grad, l.parameters()) - param_counts[idx] = sum(p.numel() for p in params) - elif isinstance(layer, nn.Module): - params = filter(lambda p: p.requires_grad, layer.parameters()) - param_counts[idx] = sum(p.numel() for p in params) - return param_counts - - def _find_layer_type(self, layername): - idxs = [] - typeregex = regex.compile(layername, regex.IGNORECASE) - for idx, layer in enumerate(self._layer_specs): - name = None - if isinstance(layer, LayerSpec): - name = layer.typename.__name__ - elif isinstance(layer, nn.Module): - name = layer.__class__.__name__ - else: - try: - name = layer.__name__ - except AttributeError: - continue - if typeregex.search(name): - idxs.append(idx) - - if len(idxs) == 0: - raise RuntimeError(f"Partitioning '{layername}' found no valid layers to partition.") - return idxs - - def forward(self, forward_input): - # We need to offset the seed by the microbatch ID. Save it in a local var to - # ensure it is preserved in the closure. Otherwise checkpointed forward funcs - # will see a different offset. - self.micro_offset += 1 - - def exec_range_func(start, end): - ''' Helper function to be used with checkpoint() - Adapted from torch.utils.checkpoint:checkpoint_sequential() - ''' - local_micro_offset = self.micro_offset + 1 - - def exec_func(*inputs): - # Single tensor inputs need to be unwrapped - if len(inputs) == 1: - inputs = inputs[0] - for idx, layer in enumerate(self.forward_funcs[start:end]): - self.curr_layer = idx + self._local_start - if self.seed_layers: - new_seed = (self.base_seed * local_micro_offset) + self.curr_layer - if self.seed_fn: - self.seed_fn(new_seed) - else: - ds_utils.set_random_seed(new_seed) - - inputs = layer(inputs) - return inputs - - return exec_func - - if self.activation_checkpoint_interval == 0: - func = exec_range_func(0, len(self.forward_funcs)) - x = func(forward_input) - else: - num_layers = len(self.forward_funcs) - x = forward_input - for start_idx, is_checkpointable_result in \ - zip(range(0, num_layers, self.activation_checkpoint_interval), self.is_checkpointable_results): - - end_idx = min(start_idx + self.activation_checkpoint_interval, num_layers) - - funcs = self.forward_funcs[start_idx:end_idx] - # Since we either pass tensors or tuples of tensors without unpacking, we - # need to be careful not to double-wrap tensors with tuple. - if not isinstance(x, tuple): - x = (x, ) - - if is_checkpointable_result: - x = self.activation_checkpoint_func(exec_range_func(start_idx, end_idx), *x) - else: - x = exec_range_func(start_idx, end_idx)(*x) - return x - - def _partition_layers(self, method='uniform'): - num_stages = self._topo.get_dim('pipe') - stage_id = self._topo.get_coord(self.global_rank).pipe - - if self.global_rank == 0: - logger.info(f'Partitioning pipeline stages with method {method}') - - method = method.lower() - - # Each stage gets a simple uniform number of layers. - if method == 'uniform': - num_layers = len(self._layer_specs) - self.parts = ds_utils.partition_uniform(num_items=num_layers, num_parts=num_stages) - elif method == 'parameters': - param_counts = self._count_layer_params() - self.parts = ds_utils.partition_balanced(weights=param_counts, num_parts=num_stages) - elif method.startswith('type:'): - layertype = method.split(':')[1] - binary_weights = [0] * len(self._layer_specs) - for idx in self._find_layer_type(layertype): - binary_weights[idx] = 1 - self.parts = ds_utils.partition_balanced(weights=binary_weights, num_parts=num_stages) - elif method == 'profile': - raise NotImplementedError(f'Partitioning method {method} not implemented.') - else: - raise NotImplementedError(f'Partitioning method {method} not implemented.') - - # Print some information on the partitioning. - if self.global_rank == 0: - for stage in range(num_stages): - start = self.parts[stage] - stop = self.parts[stage + 1] - print(f'stage={stage} layers={stop - start}') - for idx, layer in enumerate(self._layer_specs[start:stop]): - name = str(layer) - if isinstance(layer, LayerSpec): - name = layer.typename.__name__ - if isinstance(layer, nn.Module): - name = layer.__class__.__name__ - else: - try: - name = layer.__name__ - except AttributeError: - pass - print(f' {idx+start:2d}: {name}') - if self.loss_fn: - try: - print(f' loss: {self.loss_fn.__name__}') - except AttributeError: - print(f' loss: {self.loss_fn.__class__.__name__}') - - self._set_bounds(start=self.parts[stage_id], stop=self.parts[stage_id + 1]) - - @staticmethod - def _recursive_getattr(module: torch.nn.Module, attr_name: str) -> torch.Tensor: - '''Allow getting an attribute like "linear.weight"''' - weight = module - for item in attr_name.split("."): - weight = getattr(weight, item) - return weight - - def allreduce_tied_weight_gradients(self): - '''All reduce the gradients of the tied weights between tied stages''' - for key, comm in self.tied_comms.items(): - for attr_name in comm['weight_attr']: - weight = self._recursive_getattr(self.tied_modules[key], attr_name) - dist.all_reduce(weight.grad, group=comm['group']) - - def get_tied_weights_and_groups(self): - weight_group_list = [] - for key, comm in self.tied_comms.items(): - for attr_name in comm['weight_attr']: - weight = self._recursive_getattr(self.tied_modules[key], attr_name) - weight_group_list.append((weight, comm['group'])) - return weight_group_list - - def _synchronize_tied_weights(self): - for key, comm in self.tied_comms.items(): - for attr_name in comm['weight_attr']: - dist.broadcast( - self._recursive_getattr(comm['module'], attr_name), - src=min(comm['ranks']), - group=comm['group'], - ) - - def _index_tied_modules(self): - ''' Build communication structures for tied modules. ''' - tied_comms = {} - if self._topo.get_dim('pipe') == 1: - return tied_comms - - specs = self._layer_specs - tie_keys = set(s.key for s in specs if isinstance(s, TiedLayerSpec)) - # Since Python 3.7, "Dictionary order is guaranteed to be insertion order." - # Sort tie_keys here so that orders of self.tied_comms.items() are consistent - # among ranks. - for key in sorted(tie_keys): - # Find the layers that the tied module appears in - tied_layers = [] - for idx, layer in enumerate(specs): - if isinstance(layer, TiedLayerSpec) and layer.key == key: - tied_layers.append(idx) - # Find all stages with this tied module - # TODO: Would be nice to remove the nested data/model parallelism loops and - # TODO: instead generalize in some way, since we really just care about the - # TODO: stage that owns the tied layer. Then loop over each (dp, mp, ...) - # TODO: fiber to generate process groups. - tied_stages = set(self.stage_owner(idx) for idx in tied_layers) - for dp in range(self._grid.data_parallel_size): - for mp in range(self._grid.get_slice_parallel_world_size()): - tied_ranks = [] - for s in sorted(tied_stages): - if self._grid.get_slice_parallel_world_size() > 1: - tied_ranks.append(self._grid.stage_to_global(stage_id=s, data=dp, model=mp)) - else: - tied_ranks.append(self._grid.stage_to_global(stage_id=s, data=dp)) - group = dist.new_group(ranks=tied_ranks) - - # Record this tied module if we own a local copy of it. - if self.global_rank in tied_ranks: - assert key in self.tied_modules - if key in self.tied_modules: - tied_comms[key] = { - 'ranks': tied_ranks, - 'group': group, - 'weight_attr': self.tied_weight_attrs[key], - 'module': self.tied_modules[key], - } - # Only count the tied module once in the eyes of the FP16 optimizer - if self.global_rank != tied_ranks[0]: - for p in self.tied_modules[key].parameters(): - p.ds_pipe_replicated = True - ''' - if len(tied_comms) > 0: - print(f'RANK={self.global_rank} tied_comms={tied_comms}') - ''' - - return tied_comms - - def partitions(self): - return self.parts - - def stage_owner(self, layer_idx): - assert 0 <= layer_idx < self._num_layers - for stage in range(self._topo.get_dim('pipe')): - if self.parts[stage] <= layer_idx < self.parts[stage + 1]: - return stage - raise RuntimeError(f'Layer {layer_idx} not owned? parts={self.parts}') - - def _set_bounds(self, start=None, stop=None): - """Manually define the range of layers that will be built on this process. - - These boundaries are treated as list slices and so start is inclusive and stop is - exclusive. The default of None for both results in all layers being built - locally. - """ - self._local_start = start - self._local_stop = stop - - def set_checkpoint_interval(self, interval): - assert interval >= 0 - self.checkpoint_interval = interval - - def topology(self): - """ ProcessTopology object to query process mappings. """ - return self._topo - - def mpu(self): - return self._grid - - def num_pipeline_stages(self): - return self._topo.get_dim('pipe') - - def ckpt_prefix(self, checkpoints_path, tag): - """Build a prefix for all checkpoint files written by this module. """ - # All checkpoint files start with this - rank_name = 'module' - - # Data parallelism is omitted from the naming convention because we are agnostic - # to this in the checkpoint. - omit_dims = frozenset(['data']) - axes = [a for a in self._grid._topo.get_axis_names() if a not in omit_dims] - for dim in axes: - rank = getattr(self._grid._topo.get_coord(rank=self.global_rank), dim) - rank_name += f'-{dim}_{rank:02d}' - - ckpt_name = os.path.join(checkpoints_path, str(tag), rank_name) - return ckpt_name - - def ckpt_layer_path(self, ckpt_dir, local_layer_idx): - """Customize a prefix for a specific pipeline module layer. """ - idx = local_layer_idx + self._local_start - layer_ckpt_path = os.path.join(ckpt_dir, f'layer_{idx:02d}') - rank_repr = self._grid._topo.get_rank_repr(rank=self.global_rank) - if rank_repr != '': - layer_ckpt_path += f'-{rank_repr}' - layer_ckpt_path += '-model_states.pt' - return layer_ckpt_path - - def ckpt_layer_path_list(self, ckpt_dir, local_layer_idx): - """Get all ckpt file list for a specific pipeline module layer. """ - idx = local_layer_idx + self._local_start - layer_ckpt_path = os.path.join(ckpt_dir, f'layer_{idx:02d}-') - layer_ckpt_path += "*model_states.pt" - ckpt_files = glob.glob(layer_ckpt_path) - ckpt_files.sort() - return ckpt_files - - def save_state_dict(self, save_dir, checkpoint_engine, exclude_frozen_params=False): - # TODO: Need to validate interaction of checkpoint_parallel_write_pipeline and fastwriter - - # Processes having the same model parallel rank on different data parallel instances - # have identical layer weights. We can distribute the task of saving the layer weights - # among the data parallel ranks. For example, if a pipeline stage has 9 layers and - # if there are 2 data parallel instances, rank 0 will save the first 5 layers and - # rank 1 will save the last 4. - dp_rank = self._grid.data_parallel_id - dp_size = self._grid.data_parallel_size - num_layers = len(self.forward_funcs) - if self.checkpoint_parallel_write_pipeline: - # spread layers evenly across data parallel ranks - offsets = ds_utils.partition_uniform(num_layers, dp_size) - start, end = offsets[dp_rank], offsets[dp_rank + 1] - else: - # data parallel rank 0 writes all layers - if dp_rank != 0: - return - start, end = 0, num_layers - layer_list = self.forward_funcs[start:end] - - checkpoint_engine.makedirs(save_dir, exist_ok=True) - for idx, layer in enumerate(layer_list): - model_ckpt_path = self.ckpt_layer_path(save_dir, start + idx) - if not hasattr(layer, 'state_dict'): - continue - - orig_state_dict = layer.state_dict() - if exclude_frozen_params: - for n in self._get_frozen_parameter_names(layer): - del orig_state_dict[n] - final_state_dict = clone_tensors_for_torch_save(orig_state_dict) - checkpoint_engine.save(state_dict=final_state_dict, path=model_ckpt_path) - - def load_state_dir(self, load_dir, checkpoint_engine, strict=True): - for idx, layer in enumerate(self.forward_funcs): - # Functions, etc. will not have state_dicts - if not hasattr(layer, 'load_state_dict'): - continue - - # get all checkpoint files for the layer. - model_ckpt_list = self.ckpt_layer_path_list(load_dir, idx) - mp_rank = self._grid.get_slice_parallel_rank() - mp_world_size = self._grid.get_slice_parallel_world_size() - - sd_loader = SDLoaderFactory.get_sd_loader(model_ckpt_list, - version=2.0, - checkpoint_engine=checkpoint_engine) - load_path, checkpoint, _ = sd_loader.load(mp_world_size, mp_rank, module_key=None, is_pipe_parallel=True) - - layer.load_state_dict(checkpoint, strict=strict) - - # if self._grid.data_parallel_id == 0: - # logger.info( - # f'RANK={self.global_rank} Loaded layer={idx+self._local_start} file={load_path}' - # ) - - self._synchronize_tied_weights() - - def _is_checkpointable(self, funcs): - - if self.activation_checkpoint_func is not checkpointing.non_reentrant_checkpoint: - # This hook excludes the embedding layer - # because only non_reentrant_checkpoint can accept inputs with requires_grad=False - # otherwise, the backward of the embedding layer won't receive gradients. - if self.__class__.__name__ in ('GPTModelPipe', 'GPT2ModelPipe'): - # For GPT models, checkpoint both transformer layers and any additional - # layers specified in checkpointable_layers (if provided) - return all('ParallelTransformerLayerPipe' in f.__class__.__name__ or ( - self.checkpointable_layers is not None and f.__class__.__name__ in self.checkpointable_layers) - for f in funcs) - - if self.checkpointable_layers is not None: - # For non-GPT models, only checkpoint layers specified in checkpointable_layers - return all(f.__class__.__name__ in self.checkpointable_layers for f in funcs) - - # Default behavior: checkpoint any layer that has parameters - params = [f.parameters() for f in funcs if isinstance(f, torch.nn.Module)] - return any(len(list(p)) > 0 for p in params) - - def get_additional_losses(self): - """ Returns model specific additional losses for reporting - - Return a dictionary of {"loss name": loss_value} or None if no additional losses. - """ - return None - - def compile(self, *args, **kwargs): - for idx, layer in enumerate(self.forward_funcs): - if isinstance(layer, nn.Module): - layer.compile(*args, **kwargs) - else: - new_layer = torch.compile(layer, *args, **kwargs) - self.forward_funcs[idx] = new_layer diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/p2p.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/p2p.py deleted file mode 100644 index ed6d80b8d4fbf5131d773d72bec8574f3105bac9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/p2p.py +++ /dev/null @@ -1,182 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import msgpack -import typing - -import torch -from deepspeed import comm as dist - -from deepspeed.utils.torch import required_torch_version -from deepspeed.accelerator import get_accelerator - -_groups = None -_grid = None - -_async = [] - - -def can_send_recv() -> bool: - return required_torch_version(min_version=1.8) - - -#initializes adjacent process groups -#run this only after deepspeed.init_distributed() has been called -def init_process_groups(grid): - global _groups, _grid - _grid = grid - - assert _grid.pipe_parallel_size > 1, "There is no pipeline parallelism" - - if not can_send_recv(): - _groups = [dist.new_group(ranks=group) for group in _grid.p2p_groups] - - -def _is_valid_send_recv(src_stage, dest_stage): - first_stage = 0 - last_stage = _grid.pipe_parallel_size - 1 - assert abs(src_stage-dest_stage) == 1 or \ - (src_stage == first_stage and dest_stage == last_stage) or \ - (src_stage == last_stage and dest_stage == first_stage), \ - "Functionality currently limited to send and receive between adjacent ranks only" - - -def send(tensor, dest_stage, async_op=False): - global _groups - assert async_op == False, "Doesn't support async_op true" - src_stage = _grid.get_stage_id() - _is_valid_send_recv(src_stage, dest_stage) - - dest_rank = _grid.stage_to_global(stage_id=dest_stage) - if async_op: - global _async - op = dist.isend(tensor, dest_rank) - _async.append(op) - else: - - if can_send_recv(): - return dist.send(tensor, dest_rank) - else: - group = _get_send_recv_group(src_stage, dest_stage) - src_rank = _grid.stage_to_global(stage_id=src_stage) - return dist.broadcast(tensor, src_rank, group=group, async_op=async_op) - - -def recv(tensor, src_stage, async_op=False): - global _groups - assert async_op == False, "Doesn't support async_op true" - dest_stage = _grid.get_stage_id() - _is_valid_send_recv(src_stage, dest_stage) - - src_rank = _grid.stage_to_global(stage_id=src_stage) - - if async_op: - global _async - op = dist.irecv(tensor, src_rank) - _async.append(op) - else: - if can_send_recv(): - return dist.recv(tensor, src_rank) - else: - group = _get_send_recv_group(src_stage, dest_stage) - return dist.broadcast(tensor, src_rank, group=group, async_op=async_op) - - -def wait(): - global _async - for op in _async: - op.wait() - _async = [] - - get_accelerator().synchronize() - - -def send_obj(msg: typing.Any, dest: int): - """Send an arbitrary python object to ``dest``. - - Note: ``msg`` must be serializable by msgpack. - - WARN: This incurs a CPU -> GPU transfer and should be used sparingly - for performance reasons. - - Args: - msg (typing.Any): The object to send. - dest (int): Destination rank. - """ - # serialize the message - msg = msgpack.packb(msg) - # construct a tensor to send - msg = torch.ByteTensor(torch.ByteStorage.from_buffer(msg)).to(get_accelerator().device_name()) - - # Send meta and message - length_tensor = torch.tensor([len(msg)], dtype=torch.long).to(get_accelerator().device_name()) - dist.send(length_tensor, dst=dest) - dist.send(msg, dst=dest) - - -def recv_obj(sender: int) -> typing.Any: - """Receive an arbitrary python object from ``sender``. - - WARN: This incur a CPU <-> GPU transfers and should be used sparingly - for performance reasons. - - Args: - sender (int): The rank sending the message. - """ - # Get message meta - length = torch.tensor([0], dtype=torch.long).to(get_accelerator().device_name()) - dist.recv(length, src=sender) - - # Receive and deserialize - msg = torch.empty(length.item(), dtype=torch.uint8).to(get_accelerator().device_name()) - dist.recv(msg, src=sender) - - msg = msgpack.unpackb(msg.cpu().numpy().tobytes()) - - def _to(x): - """Recursively move to the current device.""" - if torch.is_tensor(x): - return x.to(get_accelerator().device_name()) - if isinstance(x, (tuple, list)): - ret = [_to(x_) for x_ in x] - if isinstance(x, tuple): - ret = tuple(ret) - return ret - # handle kwargs - if isinstance(x, dict): - ret = dict() - for key, val in x.items(): - ret[_to(key)] = _to(val) - return ret - - # Anything else is a no-op - return x - - msg = _to(msg) - return msg - - -def _get_send_recv_group(src_stage, dest_stage): - '''the group id is always the smaller rank unless its a wrap around''' - - stage_id = None - - first_stage = 0 - last_stage = _grid.pipe_parallel_size - 1 - - if (src_stage == first_stage and dest_stage == last_stage - or dest_stage == first_stage and src_stage == last_stage): - stage_id = last_stage - elif src_stage > dest_stage: - stage_id = dest_stage - else: - stage_id = src_stage - '''group_id corresponds to group of [group_id, group_id+1] - unless group_id is the rank of the last stage - in which case group_id corresponds to group[group_id-num_stages+1, group_id] - ''' - group_id = _grid.stage_to_global(stage_id=stage_id) - - return _groups[group_id] diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/schedule.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/schedule.py deleted file mode 100644 index 21bf24dad8edcffd54d05d0c7092ab225e1e39dc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/schedule.py +++ /dev/null @@ -1,494 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from ..utils import call_to_str - -from abc import ABC, abstractmethod - - -class PipeSchedule(ABC): - """Directs the execution of a pipeline engine by generating sequences of - :class:`PipeInstruction`. - - Schedules are generators that yield sequences of - :class:`PipeInstruction` to process the micro-batches in one batch. - Each yielded step is atomic in the sense that a barrier - synchronization can be placed between successive steps without - deadlock. - - Below is an example schedule that implements data parallelism with gradient accumulation: - - .. code-block:: python - - class DataParallelSchedule(PipeSchedule): - def steps(self): - for step_id in range(self.micro_batches): - cmds = [ - LoadMicroBatch(buffer_id=0), - ForwardPass(buffer_id=0), - BackwardPass(buffer_id=0), - ] - if step_id == self.micro_batches - 1: - cmds.extend([ - ReduceGrads(), - OptimizerStep(), - ]) - yield cmds - - def num_pipe_buffers(self): - return 1 - - Args: - micro_batches (int): The number of micro-batches that comprise a batch. - stages (int): The number of pipeline stages. - stage_id (int): The pipe stage that will execute the generated schedule. - """ - - def __init__(self, micro_batches, stages, stage_id): - super().__init__() - self.micro_batches = micro_batches - self.stages = stages - self.stage_id = stage_id - self.prev_stage = self.stage_id - 1 - self.next_stage = self.stage_id + 1 - - @abstractmethod - def steps(self): - """Yield a list of :class:`PipeInstruction` for each step in the schedule. - - .. note:: - Schedules must implement ``steps()`` to define the schedule. - - Returns: - Instructions to be executed as one step of the pipeline - """ - pass - - def num_pipe_buffers(self): - """The number of pipeline buffers that will be used by this stage. - - .. note:: - Schedules should specialize ``num_pipe_buffers()`` for memory savings at scale. - - Returns: - The number of buffers for the engine to allocate. - """ - return self.micro_batches - - def _valid_micro_batch(self, micro_batch_id): - return 0 <= micro_batch_id < self.micro_batches - - def _valid_stage(self, stage_id): - return 0 <= stage_id < self.stages - - @property - def stage(self): - """Stage index used to configure this schedule.""" - return self.stage_id - - @property - def num_stages(self): - """The number of total pipeline stages used to configure this schedule.""" - return self.stages - - @property - def num_micro_batches(self): - """The number of total micro_batches used to configure this schedule.""" - return self.micro_batches - - @property - def is_first_stage(self): - """True if the configured ``stage_id`` is the first stage in the pipeline.""" - return self.stage_id == 0 - - @property - def is_last_stage(self): - """True if the configured ``stage_id`` is the last stage in the pipeline.""" - return self.stage_id == self.stages - 1 - - def _buffer_idx(self, micro_batch_id): - """Map a micro-batch index to a pipeline buffer index. - - This method uses a cyclic allocation strategy. - - Args: - micro_batch_id (int): The micro-batch index relative to the beginning of the schedule. - - Returns: - int: The index of the buffer that should store data. - """ - assert self._valid_micro_batch(micro_batch_id) - return micro_batch_id % self.num_pipe_buffers() - - def __iter__(self): - self.it = None - return self - - def __next__(self): - if self.it is None: - self.it = self.steps() - return next(self.it) - - -class InferenceSchedule(PipeSchedule): - """A schedule for inferencing batches using pipeline parallelism. - """ - - def steps(self): - """""" - prev_micro_batch_id = -1 - total_steps = self.micro_batches + self.stages - 1 - for step_id in range(total_steps): - cmds = [] - micro_batch_id = step_id - self.stage_id - - # Alternate send/recv buffers - if _is_even(self.stage_id): - recv_buf = step_id % 2 - send_buf = (step_id + 1) % 2 - else: - recv_buf = (step_id + 1) % 2 - send_buf = step_id % 2 - - if self.is_first_stage or self.is_last_stage: - if self._valid_micro_batch(micro_batch_id): - cmds.append(LoadMicroBatch(recv_buf)) - - if _is_even(self.stage_id): - if self._valid_stage(self.next_stage): - if self._valid_micro_batch(micro_batch_id - 1): - cmds.append(SendActivation(send_buf)) - if self._valid_stage(self.prev_stage): - if self._valid_micro_batch(micro_batch_id): - cmds.append(RecvActivation(recv_buf)) - else: - if self._valid_stage(self.prev_stage): - if self._valid_micro_batch(micro_batch_id): - cmds.append(RecvActivation(recv_buf)) - - if self._valid_stage(self.next_stage): - if self._valid_micro_batch(micro_batch_id - 1): - cmds.append(SendActivation(send_buf)) - - if self._valid_micro_batch(micro_batch_id): - cmds.append(ForwardPass(recv_buf)) - - yield cmds - - def num_pipe_buffers(self): - """Only two pipeline buffers are required for inferencing. - - Returns: - ``2`` - """ - return 2 - - -class TrainSchedule(PipeSchedule): - """A schedule for training a batch using hybrid parallelism. - - Pipeline parallelism is extracted through gradient accumulation and thus - convergence follows that of a data parallel approach with the same batch - size. - """ - - def steps(self): - """""" - prev_micro_batch_id = -1 - total_steps = 2 * (self.micro_batches + self.stages - 1) - for step_id in range(total_steps): - # Map the step of the pipeline to the micro-batch id and also whether it is a - # forward or backward pass step. - micro_batch_id, is_forward = self._step_to_micro_batch(step_id) - - if self._valid_micro_batch(prev_micro_batch_id): - prev_buffer = self._buffer_idx(prev_micro_batch_id) - if self._valid_micro_batch(micro_batch_id): - curr_buffer = self._buffer_idx(micro_batch_id) - - cmds = [] - - # Exchange activations - if is_forward: - if self._valid_micro_batch(prev_micro_batch_id) and self._valid_stage(self.prev_stage): - cmds.append(SendGrad(prev_buffer)) - if self._valid_micro_batch(micro_batch_id) and self._valid_stage(self.prev_stage): - cmds.append(RecvActivation(curr_buffer)) - else: - if self._valid_micro_batch(micro_batch_id) and self._valid_stage(self.next_stage): - cmds.append(RecvGrad(curr_buffer)) - if self._valid_micro_batch(prev_micro_batch_id) and self._valid_stage(self.next_stage): - cmds.append(SendActivation(prev_buffer)) - - # First/last stage loads - if self.stage_id == 0 or self.stage_id == self.stages - 1: - if is_forward and self._valid_micro_batch(micro_batch_id): - cmds.append(LoadMicroBatch(curr_buffer)) - - # Computation - if self._valid_micro_batch(micro_batch_id): - if is_forward: - cmds.append(ForwardPass(curr_buffer)) - else: - cmds.append(BackwardPass(curr_buffer)) - - # Model step at the end of the batch - if step_id == total_steps - 1: - cmds.append(ReduceTiedGrads()) - cmds.append(ReduceGrads()) - cmds.append(OptimizerStep()) - - # Prepare state for next time - prev_micro_batch_id = micro_batch_id - yield cmds - - def num_pipe_buffers(self): - """Return the number of pipeline buffers required for this stage. - - This is equivalent to the maximum number of in-flight forward passes, - since we need to remember the activations of forward passes in order - to run backpropagation. For synchronous 1F1B, this is equivalent to - the index difference between this stage and the last stage. - """ - buffers = min(self.stages - self.stage_id, self.micro_batches) - return max(2, buffers) - - def _step_to_micro_batch(self, step_id): - if _is_even(step_id) and _is_even(self.stage_id): - micro_batch_id = self._even_step_forward_id(step_id) - is_forward = True - - elif _is_odd(step_id) and _is_odd(self.stage_id): - micro_batch_id = self._odd_step_forward_id(step_id) - is_forward = True - - elif _is_even(step_id) and _is_odd(self.stage_id): - micro_batch_id = self._even_step_backward_id(step_id) - is_forward = False - - elif _is_odd(step_id) and _is_even(self.stage_id): - micro_batch_id = self._odd_step_backward_id(step_id) - is_forward = False - - else: - assert False - - return micro_batch_id, is_forward - - def _even_step_forward_id(self, step_id): - base = step_id // 2 - micro_batch_id = int(base - self.stage_id // 2) - return micro_batch_id - - def _odd_step_forward_id(self, step_id): - base = (step_id - 1) // 2 - micro_batch_id = int(base - self.stage_id // 2) - return micro_batch_id - - def _even_step_backward_id(self, step_id): - base = step_id // 2 - micro_batch_id = int(base - self.stages + (self.stage_id + 1) // 2) - return micro_batch_id - - def _odd_step_backward_id(self, step_id): - base = ((step_id - 1) // 2) - self.stages + 1 - micro_batch_id = int(base + self.stage_id // 2) - return micro_batch_id - - -class DataParallelSchedule(PipeSchedule): - """An example schedule that trains using traditional data parallelism with gradient - accumulation. - """ - - def steps(self): - """""" - for step_id in range(self.micro_batches): - cmds = [ - LoadMicroBatch(buffer_id=0), - ForwardPass(buffer_id=0), - BackwardPass(buffer_id=0), - ] - if step_id == self.micro_batches - 1: - cmds.extend([ - ReduceGrads(), - OptimizerStep(), - ]) - yield cmds - - def num_pipe_buffers(self): - """Only one pipeline buffer needed. - """ - return 1 - - -class PipeInstruction: - """Base class for all instructions to be executed by the pipeline engine. - - All keyword arguments are stored as members similar to a ``namedtuple``. These are - then accessible to the :class:`PipeEngine` during execution. - - Args: - kwargs (optional): keyword arguments to store as members - """ - - def __init__(self, **kwargs): - self.name = self.__class__.__name__ - self.kwargs = kwargs - for key, val in kwargs.items(): - setattr(self, key, val) - - def __repr__(self): - return call_to_str(self.name, **self.kwargs) - - -class OptimizerStep(PipeInstruction): - """Performs one step with the optimizer and zeros gradients. - - .. note:: Should be issued after :class:`ReduceGrads` and :class:`ReduceTiedGrads`. - - .. note:: Can be a synchronization point among data-parallel ranks. - """ - pass - - -class ReduceGrads(PipeInstruction): - """Reduce the computed gradients among data-parallel processes within the stage. - """ - pass - - -class ReduceTiedGrads(PipeInstruction): - """Reduce the computed gradients of tied modules within a pipeline-parallel group. - - .. warning:: - The stages included in this synchronization point are not known until - the model is partitioned among pipeline stages. In the worst case, it - includes all pipeline stages. This instruction should be scheduled - carefully to avoid deadlocks. - """ - pass - - -class BufferOpInstruction(PipeInstruction): - """A pipeline instruction that operates on pipeline buffer(s). - - Args: - buffer_id (int): the index of the pipeline buffer() to modify. - """ - - def __init__(self, buffer_id, **kwargs): - super().__init__(buffer_id=buffer_id, **kwargs) - - -# IO -class LoadMicroBatch(BufferOpInstruction): - """Load a micro-batch into a buffer. - - Roughly: - - .. code-block:: python - - buffers['inputs'][buffer_id] = next(data_iter) - """ - pass - - -# Compute -class ForwardPass(BufferOpInstruction): - """Compute a forward pass. - - Roughly: - - .. code-block:: python - - buffers['outputs'][buffer_id] = forward(buffers['inputs'][buffer_id]) - """ - pass - - -class BackwardPass(BufferOpInstruction): - """Compute a backward pass and accumulate gradients. - - Roughly: - - .. code-block:: python - - outputs = buffers['outputs'][buffer_id] - gradients = buffers['gradients'][buffer_id] - torch.autograd.backward(tensors=outputs, - grad_tensors=gradients) - """ - pass - - -# Communication -class SendActivation(BufferOpInstruction): - """Send activations to the next stage in the pipeline. - - Roughly: - - .. code-block:: python - - send(buffers['outputs'][buffer_id]) - - .. note:: - The communication is blocking and must be paired with a :class:`RecvActivation` - on the next pipeline stage to avoid deadlock. - """ - pass - - -class RecvActivation(BufferOpInstruction): - """Receive activations from the previous stage in the pipeline. - - Roughly: - - .. code-block:: python - - buffers['inputs'][buffer_id] = recv() - - .. note:: - The communication is blocking and must be paired with a :class:`SendActivation` - on the previous pipeline stage to avoid deadlock. - """ - pass - - -class SendGrad(BufferOpInstruction): - """Send computed gradients to the previous pipeline stage. - with respect to the received activations - - .. note:: - Only received tensors with ``requires_grad==True`` will produce gradients. - Missing gradients will be replaced with ``None`` on the receiving stage. - - .. note:: - The communication is blocking and must be paired with a :class:`RecvGrad` - on the previous pipeline stage to avoid deadlock. - """ - pass - - -class RecvGrad(BufferOpInstruction): - """Receive computed gradients the next pipeline stage. - - .. note:: - Only activations with ``requires_grad==True`` will produce gradients. - Missing gradients will be replaced with ``None``. - - .. note:: - The communication is blocking and must be paired with a :class:`SendGrad` - on the next pipeline stage to avoid deadlock. - """ - pass - - -def _is_even(x): - return x % 2 == 0 - - -def _is_odd(x): - return x % 2 != 0 diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/topology.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/topology.py deleted file mode 100644 index 4b107762636412ce2d455fe004fac863f6ef8f91..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/pipe/topology.py +++ /dev/null @@ -1,472 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed import comm as dist - -from collections import namedtuple -from itertools import product as cartesian_product - - -class ProcessTopology: - """ Manages the mapping of n-dimensional Cartesian coordinates to linear - indices. This mapping is used to map the rank of processes to the grid - for various forms of parallelism. - - Each axis of the tensor is accessed by its name. The provided ordering - of the axes defines the layout of the topology. ProcessTopology uses a "row-major" - layout of the tensor axes, and so axes=['x', 'y'] would map coordinates (x,y) and - (x,y+1) to adjacent linear indices. If instead axes=['y', 'x'] was used, coordinates - (x,y) and (x+1,y) would be adjacent. - - Some methods return ProcessCoord namedtuples. - """ - - def __init__(self, axes, dims): - """Create a mapping of n-dimensional tensor coordinates to linear indices. - - Arguments: - axes (list): the names of the tensor axes - dims (list): the dimension (length) of each axis of the topology tensor - """ - - self.axes = axes # names of each topology axis - self.dims = dims # length of each topology axis - - # This is actually a class that lets us hash {'row':3, 'col':2} mappings - self.ProcessCoord = namedtuple('ProcessCoord', axes) - - self.mapping = {} - ranges = [range(d) for d in dims] - # example: 1, (0,0,1) - for global_rank, coord in enumerate(cartesian_product(*ranges)): - key = {axis: coord[self.axes.index(axis)] for axis in self.axes} - key = self.ProcessCoord(**key) - # for example, {ProcessCoord(row=0, col=1) : 1} - self.mapping[key] = global_rank - - def get_rank(self, **coord_kwargs): - """Return the global rank of a process via its coordinates. - - Coordinates are specified as kwargs. For example: - - >>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3]) - >>> X.get_rank(x=0, y=1) - 1 - """ - if len(coord_kwargs) != len(self.axes): - raise ValueError('get_rank() does not support slices. Use filter_match())') - - key = self.ProcessCoord(**coord_kwargs) - assert key in self.mapping, f'key {coord_kwargs} invalid' - return self.mapping[key] - - def get_axis_names(self): - """Return a list of the axis names in the ordering of the topology. """ - return self.axes - - def get_rank_repr(self, rank, omit_axes=['data', 'pipe'], inner_sep='_', outer_sep='-'): - """Return a string representation of a rank. - - This method is primarily used for checkpointing model data. - - For example: - >>> topo = Topo(axes=['a', 'b'], dims=[2, 2]) - >>> topo.get_rank_repr(rank=3) - 'a_01-b_01' - >>> topo.get_rank_repr(rank=3, omit_axes=['a']) - 'b_01' - - Args: - rank (int): A rank in the topology. - omit_axes (list, optional): Axes that should not be in the representation. Defaults to ['data', 'pipe']. - inner_sep (str, optional): [description]. Defaults to '_'. - outer_sep (str, optional): [description]. Defaults to '-'. - - Returns: - str: A string representation of the coordinate owned by ``rank``. - """ - omit_axes = frozenset(omit_axes) - axes = [a for a in self.get_axis_names() if a not in omit_axes] - names = [] - for ax in axes: - ax_rank = getattr(self.get_coord(rank=rank), ax) - names.append(f'{ax}{inner_sep}{ax_rank:02d}') - return outer_sep.join(names) - - def get_dim(self, axis): - """Return the number of processes along the given axis. - - For example: - >>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3]) - >>> X.get_dim('y') - 3 - """ - if axis not in self.axes: - return 0 - return self.dims[self.axes.index(axis)] - - def get_coord(self, rank): - """Return the coordinate owned by a process rank. - - The axes of the returned namedtuple can be directly accessed as members. For - example: - >>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3]) - >>> coord = X.get_coord(rank=1) - >>> coord.x - 0 - >>> coord.y - 1 - """ - for coord, idx in self.mapping.items(): - if idx == rank: - return coord - raise ValueError(f'rank {rank} not found in topology.') - - def get_axis_comm_lists(self, axis): - """ Construct lists suitable for a communicator group along axis ``axis``. - - Example: - >>> topo = Topo(axes=['pipe', 'data', 'model'], dims=[2, 2, 2]) - >>> topo.get_axis_comm_lists('pipe') - [ - [0, 4], # data=0, model=0 - [1, 5], # data=0, model=1 - [2, 6], # data=1, model=0 - [3, 7], # data=1, model=1 - ] - - Returns: - A list of lists whose coordinates match in all axes *except* ``axis``. - """ - - # We don't want to RuntimeError because it allows us to write more generalized - # code for hybrid parallelisms. - if axis not in self.axes: - return [] - - # Grab all axes but `axis` - other_axes = [a for a in self.axes if a != axis] - - lists = [] - - # Construct all combinations of coords with other_axes - ranges = [range(self.get_dim(a)) for a in other_axes] - for coord in cartesian_product(*ranges): - other_keys = {a: coord[other_axes.index(a)] for a in other_axes} - # now go over all ranks in `axis`. - sub_list = [] - for axis_key in range(self.get_dim(axis)): - key = self.ProcessCoord(**other_keys, **{axis: axis_key}) - sub_list.append(self.mapping[key]) - lists.append(sub_list) - - return lists - - def filter_match(self, **filter_kwargs): - """Return the list of ranks whose coordinates match the provided criteria. - - Example: - >>> X = ProcessTopology(axes=['pipe', 'data', 'model'], dims=[2, 2, 2]) - >>> X.filter_match(pipe=0, data=1) - [2, 3] - >>> [X.get_coord(rank) for rank in X.filter_match(pipe=0, data=1)] - [ProcessCoord(pipe=0, data=1, model=0), ProcessCoord(pipe=0, data=1, model=1)] - - Arguments: - **filter_kwargs (dict): criteria used to select coordinates. - - Returns: - The list of ranks whose coordinates match filter_kwargs. - """ - - def _filter_helper(x): - for key, val in filter_kwargs.items(): - if getattr(x, key) != val: - return False - return True - - coords = filter(_filter_helper, self.mapping.keys()) - return [self.mapping[coord] for coord in coords] - - def get_axis_list(self, axis, idx): - """Returns the list of global ranks whose coordinate in an axis is idx. - - For example: - >>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3]) - >>> X.get_axis_list(axis='x', idx=0) - [0, 1, 2] - >>> X.get_axis_list(axis='y', idx=0) - [0, 3] - """ - - # This could be faster by generating the desired keys directly instead of - # filtering. - axis_num = self.axes.index(axis) - ranks = [self.mapping[k] for k in self.mapping.keys() if k[axis_num] == idx] - return ranks - - def world_size(self): - return len(self.mapping) - - def __str__(self): - return str(self.mapping) - - -def _prime_factors(N): - """ Returns the prime factorization of positive integer N. """ - if N <= 0: - raise ValueError("Values must be strictly positive.") - - primes = [] - while N != 1: - for candidate in range(2, N + 1): - if N % candidate == 0: - primes.append(candidate) - N //= candidate - break - return primes - - -class PipeDataParallelTopology(ProcessTopology): - """ A topology specialization for hybrid data and pipeline parallelism. - - Uses data parallelism on the last dimension to encourage gradient - reductions to use high-bandwidth intra-node links and lower-volume - pipeline communications to use low-bandwidth inter-node links. - """ - - def __init__(self, num_pp, num_dp): - super().__init__(axes=['pipe', 'data'], dims=[num_pp, num_dp]) - - -class PipeModelDataParallelTopology(ProcessTopology): - """ A topology for hybrid pipeline, model, and data parallelism. """ - - def __init__(self, num_pp, num_mp, num_dp): - super().__init__(axes=['pipe', 'data', 'model'], dims=[num_pp, num_dp, num_mp]) - - -class PipelineParallelGrid: - """Implements a grid object that stores the data parallel ranks - corresponding to each of the model parallel stages - - The grid object organizes the processes in a distributed pytorch job - into a 2D grid, of stage_id and data_parallel_id. - - self.stage_id and self.data_parallel_id stores the stage id - and the data parallel id of current process. - - self.dp_group groups the processes by stage_id. - self.dp_group[i], is a list containing all process ranks whose - stage_id is i. - - self.p2p_groups stores a list of tuple, where each tuple - stores process ranks of adjacent stages for a given data_parallel_id. - For example if num_stage is 5 then a tuple [7,8] represents stages [3, 4], - with data_parallel id = 1. A stage wrap around will appear as non-adjacent ranks, - for example tuple [4,0] with representing wrap-around stage 4 and 0, for - data_parallel_id = 0, or similarly [9,5] represents wrapped around stages [4,0] - for data_parallel_id = 1. - """ - - def __init__(self, topology=None, process_group=None): - # TODO use process_group if provided - self.global_rank = dist.get_rank() - self.world_size = dist.get_world_size() - if topology is not None: - if self.global_rank == 0: - print('Using topology:', topology) - self._topo = topology - else: - num_pp = 1 - num_dp = 1 - for idx, prime in enumerate(_prime_factors(self.world_size)): - if idx % 2 == 0: - num_pp *= prime - else: - num_dp *= prime - self._topo = PipeDataParallelTopology(num_dp=num_dp, num_pp=num_pp) - self.data_parallel_size = max(self._topo.get_dim('data'), 1) - self.pipe_parallel_size = max(self._topo.get_dim('pipe'), 1) - self.model_parallel_size = max(self._topo.get_dim('model'), 1) - self.slice_parallel_size = self.model_parallel_size - assert self._is_grid_valid(), "Invalid Grid" - - self.stage_id = self.get_stage_id() - self.data_parallel_id = self.get_data_parallel_id() - - # Create new ProcessGroups for all model parallelism. DeepSpeedLight uses these - # to detect overflow, etc. - self.ds_model_proc_group = None - self.ds_model_rank = -1 - for dp in range(self.data_parallel_size): - ranks = sorted(self._topo.get_axis_list(axis='data', idx=dp)) - if self.global_rank == 0: - #print(f'RANK={self.global_rank} building DeepSpeed model group: {ranks}') - pass - proc_group = dist.new_group(ranks=ranks) - if self.global_rank in ranks: - self.ds_model_proc_group = proc_group - self.ds_model_world_size = len(ranks) - self.ds_model_rank = ranks.index(self.global_rank) - assert self.ds_model_rank > -1 - assert self.ds_model_proc_group is not None - - # Create new ProcessGroup for gradient all-reduces - these are the data parallel groups - self.dp_group = [] - self.dp_groups = self._topo.get_axis_comm_lists('data') - for g in self.dp_groups: - proc_group = dist.new_group(ranks=g) - if self.global_rank in g: - self.dp_group = g - self.dp_proc_group = proc_group - - self.is_first_stage = (self.stage_id == 0) - self.is_last_stage = (self.stage_id == (self.pipe_parallel_size - 1)) - - self.p2p_groups = self._build_p2p_groups() - - # Create new ProcessGroup for pipeline collectives - these are pipe parallel groups - self.pp_group = [] - self.pp_proc_group = None - self.pipe_groups = self._topo.get_axis_comm_lists('pipe') - for ranks in self.pipe_groups: - if self.global_rank == 0: - #print(f'RANK={self.global_rank} building pipeline group: {ranks}') - pass - proc_group = dist.new_group(ranks=ranks) - if self.global_rank in ranks: - self.pp_group = ranks - self.pp_proc_group = proc_group - assert self.pp_proc_group is not None - - # Create new ProcessGroup for model (tensor-slicing) collectives - - # Short circuit case without model parallelism. - # TODO: it would be nice if topology had bcast semantics to avoid this branching - # case? - if self.model_parallel_size == 1: - for group_rank in range(self.world_size): - group_rank = [group_rank] - group = dist.new_group(ranks=group_rank) - if group_rank[0] == self.global_rank: - self.slice_group = group_rank - self.slice_proc_group = group - return - else: - self.mp_group = [] - self.model_groups = self._topo.get_axis_comm_lists('model') - for g in self.model_groups: - proc_group = dist.new_group(ranks=g) - if self.global_rank in g: - self.slice_group = g - self.slice_proc_group = proc_group - - def get_stage_id(self): - return self._topo.get_coord(rank=self.global_rank).pipe - - def get_data_parallel_id(self): - return self._topo.get_coord(rank=self.global_rank).data - - def _build_p2p_groups(self): - """Groups for sending and receiving activations and gradients across model - parallel stages. - """ - comm_lists = self._topo.get_axis_comm_lists('pipe') - p2p_lists = [] - for rank in range(self.world_size): - for l in comm_lists: - assert len(l) == self.pipe_parallel_size - if rank in l: - idx = l.index(rank) - buddy_rank = l[(idx + 1) % self.pipe_parallel_size] - p2p_lists.append([rank, buddy_rank]) - break # next global rank - assert len(p2p_lists) == self.world_size - return p2p_lists - - def _is_grid_valid(self): - ranks = 1 - for ax in self._topo.get_axis_names(): - ranks *= self._topo.get_dim(ax) - return ranks == dist.get_world_size() - - #returns the global rank of the process with the provided stage id - #which has the same data_parallel_id as caller process - def stage_to_global(self, stage_id, **kwargs): - me = self._topo.get_coord(self.global_rank) - transform = me._replace(pipe=stage_id, **kwargs)._asdict() - return self._topo.get_rank(**transform) - - def topology(self): - return self._topo - - # MPU functions for DeepSpeed integration - def get_global_rank(self): - return self.global_rank - - def get_pipe_parallel_rank(self): - """ The stage of the pipeline this rank resides in. """ - return self.get_stage_id() - - def get_pipeline_model_parallel_rank(self): - return self.get_pipe_parallel_rank() - - def get_pipe_parallel_world_size(self): - """ The number of stages in the pipeline. """ - return self.pipe_parallel_size - - def get_pipeline_model_parallel_world_size(self): - return self.get_pipe_parallel_world_size() - - def get_pipe_parallel_group(self): - """ The group of ranks within the same pipeline. """ - return self.pp_proc_group - - def get_data_parallel_rank(self): - """ Which pipeline this rank resides in. """ - return self.data_parallel_id - - def get_data_parallel_world_size(self): - """ The number of pipelines. """ - return self.data_parallel_size - - def get_data_parallel_group(self): - """ The group of ranks within the same stage of all pipelines. """ - return self.dp_proc_group - - def get_data_parallel_group_ranks(self): - """ List of ranks in the data parallel group. """ - return self.dp_group - - # These are model parallel groups across all types of model parallelism. - # Deepspeed uses them to detect overflow, etc. - def get_model_parallel_rank(self): - return self.ds_model_rank - - def get_model_parallel_world_size(self): - return self.ds_model_world_size - - def get_model_parallel_group(self): - return self.ds_model_proc_group - - # For Megatron-style tensor slicing - def get_slice_parallel_rank(self): - if 'model' in self._topo.get_axis_names(): - return self._topo.get_coord(rank=self.global_rank).model - else: - return 0 - - def get_tensor_model_parallel_rank(self): - return self.get_slice_parallel_rank() - - def get_slice_parallel_world_size(self): - return self.slice_parallel_size - - def get_tensor_model_parallel_world_size(self): - return self.get_slice_parallel_world_size() - - def get_slice_parallel_group(self): - return self.slice_proc_group diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/precision_config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/precision_config.py deleted file mode 100644 index 1c7fab8bd234ef3ca87a8feccd2fd8aca463e915..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/precision_config.py +++ /dev/null @@ -1,146 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.runtime.config_utils import DeepSpeedConfigModel -from .fp16.loss_scaler import ( - INITIAL_LOSS_SCALE, - SCALE_WINDOW, - DELAYED_SHIFT, - CONSECUTIVE_HYSTERESIS, - MIN_LOSS_SCALE, -) - -######################################### -# BFLOAT16 support -######################################### -# BFLOAT16 feature. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -BFLOAT16_FORMAT = ''' -BFLOAT16 parameters should be of the format: -"bf16": { - "enabled": true, - "immediate_grad_update": false, - "check_grad_overflow": false -} -''' -BFLOAT16 = "bf16" -BFLOAT16_OLD = "bfloat16" # keeping for backwards compatibility - - -def get_bfloat16_config(param_dict): - bf16_config_dict = param_dict.get(BFLOAT16, None) - if bf16_config_dict is None: - bf16_config_dict = param_dict.get(BFLOAT16_OLD, {}) - return DeepSpeedBF16Config(**bf16_config_dict) - - -class DeepSpeedBF16Config(DeepSpeedConfigModel): - """ - For bfloat16 configuration - """ - - enabled: bool = False - """ - Enable bfloat16 mixed-precision training/inference - """ - - immediate_grad_update: bool = False - """ - Apply gradient updates immediately rather than delayed. - """ - - check_grad_overflow: bool = False - """ - Check for gradient overflows and underflows - """ - - -######################################### -# FP16 support -######################################### -# FP16 feature. By default, this feature is not enabled. -# Users can configure in ds_config.json as below example: -FP16_FORMAT = ''' -FP16 parameters should be of the format: -"fp16": { - "enabled": true, - "auto_cast": false, - "loss_scale": 0, - "initial_scale_power": 16, - "loss_scale_window": 1000, - "hysteresis": 2, - "consecutive_hysteresis": false, - "min_loss_scale": 1 -} -''' -FP16 = "fp16" - - -def get_float16_config(param_dict): - fp16_config_dict = param_dict.get(FP16, {}) - return DeepSpeedFP16Config(**fp16_config_dict) - - -class DeepSpeedFP16Config(DeepSpeedConfigModel): - """ - For float16 configuration - """ - - enabled: bool = False - """ - Enable fp16 mixed-precision training/inference - """ - - auto_cast: bool = False - """ - Automatically cast inputs to fp16 - """ - - loss_scale: float = 0 - """ - Loss scaling value. Default value of 0 means dynamic loss scaling instead of static loss scale. - """ - - initial_scale_power: int = 16 - """ - For dynamic loss scaling, set initial loss scale to 2^{initial_scale_power}. - """ - - loss_scale_window: int = 1000 - """ - Iteration intervals for raising/lowering dynamic loss scale value. - """ - - hysteresis: int = 2 - """ - Delay shift in dynamic loss scaling. - """ - - consecutive_hysteresis: bool = False - """ - Refill hysteresis if iteration does not overflow/underflow. - """ - - min_loss_scale: int = 1 - """ - Minimum dynamic loss scale value. - """ - - fp16_master_weights_and_grads: bool = False - """ - Maintain master weights in optimizer state as fp16 instead of fp32 (valid with DeepSpeedCPUAdam only). - """ - - def initial_dynamic_scale(self): - return 2**self.initial_scale_power - - def dynamic_loss_scale_args(self): - return { - INITIAL_LOSS_SCALE: 2**self.initial_scale_power, - SCALE_WINDOW: self.loss_scale_window, - DELAYED_SHIFT: self.hysteresis, - CONSECUTIVE_HYSTERESIS: self.consecutive_hysteresis, - MIN_LOSS_SCALE: self.min_loss_scale, - } diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/progressive_layer_drop.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/progressive_layer_drop.py deleted file mode 100644 index bd3be8f13011a4c9322839ff8a631778e7fb6847..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/progressive_layer_drop.py +++ /dev/null @@ -1,40 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import numpy as np -from deepspeed.utils import log_dist - - -class ProgressiveLayerDrop(object): - r""" Progressive Layer Dropping (PLD) for model training. - This implements the PLD technique for compressed model training - from this paper: https://arxiv.org/pdf/2010.13369.pdf - Args: - theta (float): a hyper-parameter that controls the trade-off between training time and robustness. - The lower the theta value, the faster the training speed. Default value: 0.5. - gamma (float): a hyper-parameter that controls how fast the drop ratio increases. Default value: 0.001. - """ - - def __init__(self, theta=0.5, gamma=0.001): - super().__init__() - - self.theta = theta - self.gamma = gamma - self.current_theta = 1.0 - log_dist(f'Enabled progressive layer dropping (theta = {self.theta})', ranks=[0]) - - def get_state(self): - kwargs = {'progressive_layer_drop': True, 'pld_theta': self.get_theta()} - return kwargs - - def get_theta(self): - return self.current_theta - - def update_state(self, global_step): - - def _prob(x, gamma, p): - return (1. - p) * np.exp(-gamma * x) + p - - self.current_theta = _prob(global_step, self.gamma, self.theta) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/quantize.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/quantize.py deleted file mode 100644 index ad553138b2d7c8095aa1568f56a1f1a0b25e16d3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/quantize.py +++ /dev/null @@ -1,180 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import math -from deepspeed.utils import logger -from deepspeed.ops.quantizer import ds_quantizer - -TWO_D_PARAMS = 6 - - -class Quantizer(object): - - def __init__(self, - q_groups=1, - q_mixed_fp16=False, - q_change_ratio=0.01, - q_type=0, - q_rounding=0, - q_verbose=False, - q_eigenvalue=False, - use_quantizer_kernel=False, - layer_num=0): - - self.q_groups = q_groups - self.q_mixed_fp16 = q_mixed_fp16 - self.q_change_ratio = q_change_ratio - self.q_type = q_type - self.qsteps = 0 - self.quantize_real_ratio = 1.000 - self.q_verbose = q_verbose - self.q_eigenvalue = q_eigenvalue - self.use_quantizer_kernel = use_quantizer_kernel - self.q_rounding = q_rounding - self.layer_num = layer_num - - def any_precision_switch(self): - # Temporary disabled functionality - if self.layer_num == 0: - return True - result = False - for index in range(self.layer_num): - if self.q_start_bits[index] != self.q_target_bits: - next_step = self.qsteps + (TWO_D_PARAMS * (self.layer_num if self.layer_num != 0 else 1)) - if next_step >= self.q_period[index]: - result = True - return result - - def quantize(self, parameter_group, overflow, eigenvalue_enabled, block_eigenvalue={}): - - if overflow and not eigenvalue_enabled: - return - - self.step() - - self.update_fp16_ratio() - - for i in range(len(parameter_group)): - for p in parameter_group[i]: - if len(p.size()) > 1 and hasattr(p, "start_bits") and p.start_bits: - param_id = id(p) - if block_eigenvalue is None: - eigenvalue, layer_id = None, 0 - else: - eigenvalue, layer_id = block_eigenvalue[param_id] if param_id in block_eigenvalue else (None, - 0) - if eigenvalue is not None: - factor = 1 + math.floor(eigenvalue * 4) - p.data = self.compute_quantization(p.data, layer_id, factor) - else: - p.data = self.compute_quantization(p, layer_id) - - def step(self): - self.qsteps += 1 - - def quantize_highbit(self, inputs, num_bits): - - q_range = 2**num_bits - input_flat = inputs.reshape(self.q_groups, -1) - g_min = input_flat.amin(dim=-1, keepdim=True) - g_max = input_flat.amax(dim=-1, keepdim=True) - - # Random number generator (Uniform) - if self.q_rounding == 'nearest': - p = 0. - else: - p = input_flat.new(input_flat.shape).uniform_(-0.5, 0.5) - - if self.q_type == 'symmetric': - scale = 2 * torch.max(torch.abs(g_min), torch.abs(g_max)) / q_range - zero_point = 0. - input_flat = (input_flat / scale + p).round().clamp(-(q_range >> 1), (q_range >> 1) - 1) * scale - elif self.q_type == 'asymmetric': - scale = (g_max - g_min) / q_range - zero_point = (g_min / scale).round() * scale - input_flat = ((input_flat - zero_point) / scale + p).round().clamp(0, (q_range - 1)) * scale + zero_point - output = input_flat.reshape(inputs.shape).contiguous() - return output - - def quantize_tenary(self, inputs): - input_flat = inputs.reshape(self.q_groups, -1) - n = input_flat.shape[1] - m = input_flat.norm(p=1, dim=1).div(n) - thres = (0.7 * m).view(-1, 1) #.expand_as(input_flat) - pos = (input_flat > thres).type(inputs.type()) - neg = (input_flat < -thres).type(inputs.type()) - mask = (input_flat.abs() > thres).type(inputs.type()) - alpha = ((mask * input_flat).abs().sum(dim=1) / mask.sum(dim=1)).view(-1, 1) - output = alpha * pos - alpha * neg - output = output.reshape(inputs.shape).contiguous() - return output - - def quantize_binary(self, inputs): - input_flat = inputs.reshape(self.q_groups, -1) - n = input_flat.shape[1] - m = input_flat.norm(p=1, dim=1, keepdim=True).div(n) - output = input_flat.sign().mul(m) - output = output.reshape(inputs.shape).contiguous() - return output - - def mixed_fp16_quantize(self, input, input_q, index): - if self.q_mixed_fp16 and self.q_start_bits[index] >= (self.q_target_bits - 1): - input_q = input * self.quantize_real_ratio + (1 - self.quantize_real_ratio) * input_q - return input_q - return input_q - - def compute_quantization(self, input, index=0, factor=1): - # fixing the quantization bits based on the training steps - # when reducing 1 bit at each period, we increase the period - # to go slowly toward the target quantization bits - # the period and starting bit can be configured - - if input.start_bits != input.target_bits: - if self.qsteps >= input.q_period: - self.quantize_real_ratio = 1.0 - input.q_period <<= 1 - input.q_period *= factor - input.start_bits -= 1 - if self.q_verbose: - logger.info( - f'Quantization settings: current bit-precision = {input.start_bits}, step = {self.qsteps}, quantization period = {input.q_period}, index = {index}' - ) - assert (input.start_bits >= input.target_bits), \ - 'Quantization bit is lower than target precision bits!' - - if self.use_quantizer_kernel: - if input.start_bits <= 2: - raise ValueError('Quantization bit is too low, please do it without quantization kernel!') - input_q = ds_quantizer(input.data.clone(), - self.q_groups, - input.start_bits, - asym=False if self.q_type == 'symmetric' else True, - sr=False if self.q_rounding == 'nearest_neighbor' else True) - else: - if input.start_bits >= 3: - input_flat = self.quantize_highbit(input.data, input.start_bits) - elif input.start_bits == 2: - assert self.q_type == 'symmetric', 'Quantization type is not symmetric!' - assert self.q_rounding == 'nearest', 'Quantization rounding is not nearest_neighbor!' - input_flat = self.quantize_tenary(input.data) - elif input.start_bits == 1: - assert self.q_type == 'symmetric', 'Quantization type is not symmetric!' - assert self.q_rounding == 'nearest', 'Quantization rounding is not nearest_neighbor!' - input_flat = self.quantize_binary(input.data) - if self.use_quantizer_kernel: - return self.mixed_fp16_quantize(input.data, input_q, index) - else: - if self.q_mixed_fp16 and input.start_bits >= input.target_bits - 1: - input_flat = self.quantize_real_ratio * input.data + \ - (1 - self.quantize_real_ratio) * input_flat - return input_flat - - def update_fp16_ratio(self): - if self.q_mixed_fp16: - if self.quantize_real_ratio > 0: - self.quantize_real_ratio -= self.q_change_ratio - else: - self.quantize_real_ratio = 0.000 diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__init__.py deleted file mode 100644 index d8cb728da375e05de5023b4196d686c25e1c4e63..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# Copyright (c) The DeepSpeed Contributors -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index fbc27b819b0be9c03bf9ed941bbf946490c291ee..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/parallel_state_sp.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/parallel_state_sp.cpython-310.pyc deleted file mode 100644 index cc0cbfff9bef56f88258c00593ded44f65a4ed84..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/parallel_state_sp.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/ulysses_sp.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/ulysses_sp.cpython-310.pyc deleted file mode 100644 index 8c70881c86b51a22080fdb21502cdc7a93560ed7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/ulysses_sp.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/parallel_state_sp.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/parallel_state_sp.py deleted file mode 100644 index c16880e15d42467f29f8103746f175b2698f4e6c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/parallel_state_sp.py +++ /dev/null @@ -1,90 +0,0 @@ -# Copyright (c) The DeepSpeed Contributors -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -This is a slimmed-down version of parallel_state.py (mpu) from Megatron-Deepspeed -""" - -from deepspeed import comm as dist - -# Sequence parallel groups to handle both data and sequence parallelisms. -# These groups are used to reduce gradients and shard parameters and optimizer stages for ZeRO. -_SEQUENCE_PARALLEL_GROUP = None -_SEQUENCE_DATA_PARALLEL_GROUP = None - - -def initialize_sequence_parallel(sequence_parallel_size: int) -> None: - """Initialize sequence parallel groups.""" - - assert dist.is_initialized() - world_size: int = dist.get_world_size() - - if world_size < sequence_parallel_size: - raise RuntimeError(f"world_size ({world_size}) is less than sequence_parallel_size {sequence_parallel_size}") - - if sequence_parallel_size <= 1: - raise ValueError(f"sequence_parallel_size must be greater than 1, got {sequence_parallel_size}") - - if world_size % sequence_parallel_size != 0: - raise RuntimeError( - f"world_size ({world_size}) is not divisible by sequence_parallel_size {sequence_parallel_size})") - - data_parallel_size: int = world_size // sequence_parallel_size - sequence_data_parallel_size: int = sequence_parallel_size * data_parallel_size - num_sequence_parallel_groups: int = world_size // sequence_parallel_size - num_sequence_data_parallel_groups: int = world_size // sequence_parallel_size // data_parallel_size - - rank = dist.get_rank() - - # Build the sequence parallel groups. - global _SEQUENCE_PARALLEL_GROUP - assert _SEQUENCE_PARALLEL_GROUP is None, "sequence parallel group is already initialized" - for i in range(num_sequence_parallel_groups): - ranks = range(i * sequence_parallel_size, (i + 1) * sequence_parallel_size) - group = dist.new_group(ranks) - if rank in ranks: - _SEQUENCE_PARALLEL_GROUP = group - - # Build the sequence data parallel groups. - global _SEQUENCE_DATA_PARALLEL_GROUP - assert _SEQUENCE_DATA_PARALLEL_GROUP is None, "sequence data parallel group is already initialized" - all_data_sequence_parallel_group_ranks = [] - for i in range(num_sequence_data_parallel_groups): - ranks = range(i * sequence_data_parallel_size, (i + 1) * sequence_data_parallel_size) - group = dist.new_group(ranks) - all_data_sequence_parallel_group_ranks.append(list(ranks)) - if rank in ranks: - _SEQUENCE_DATA_PARALLEL_GROUP = group - - -def get_sequence_parallel_group(): - """Get the sequence parallel group the caller rank belongs to.""" - assert _SEQUENCE_PARALLEL_GROUP is not None, "sequence parallel group is not initialized" - return _SEQUENCE_PARALLEL_GROUP - - -def get_sequence_data_parallel_group(): - """Get the sequence parallel group the caller rank belongs to.""" - assert _SEQUENCE_DATA_PARALLEL_GROUP is not None, "sequence data parallel group is not initialized" - return _SEQUENCE_DATA_PARALLEL_GROUP - - -def get_sequence_parallel_world_size(): - """Return world size for the sequence parallel group.""" - return dist.get_world_size(group=get_sequence_parallel_group()) - - -def get_sequence_data_parallel_world_size(): - """Return world size for the sequence parallel group.""" - return dist.get_world_size(group=get_sequence_data_parallel_group()) - - -def get_sequence_parallel_rank(): - """Return my rank for the sequence parallel group.""" - return dist.get_rank(group=get_sequence_parallel_group()) - - -def get_sequence_data_parallel_rank(): - """Return my rank for the sequence data parallel group.""" - return dist.get_rank(group=get_sequence_data_parallel_group()) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/ulysses_sp.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/ulysses_sp.py deleted file mode 100644 index 3b1becc166df386723e06b21d26f067d9545d544..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/sequence_parallel/ulysses_sp.py +++ /dev/null @@ -1,1240 +0,0 @@ -# Copyright (c) The DeepSpeed Contributors -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -*** Arctic Long Sequence Training (ALST) components *** - -1. Ulysses Sequence Parallelism for HF Transformers implements an efficient way of training on long sequences by employing sequence parallelism and attention head parallelism. -2. ALST enables even longer sequence lengths using a bag of tricks: -- Activation checkpoint offload to CPU -- Tiled MLP compute -- Liger-kernel -- PYTORCH_CUDA_ALLOC_CONF - -ALST features found in this module: - -- `UlyssesSPAttentionHF` - port of UlyssesAttention from Megatron-Deepspeed plus modern MHA-variations -- `UlyssesSPDataLoaderAdapter` - DL adapter to shard the normal DL batches to be used by `UlyssesSPAttentionHF` -- `SequenceTiledCompute` - generic autograd function to perform compute after tiling on the sequence dimension -- `TiledMLP` - a specific autograd function to perform tiled MLP (it's much easier to understand before trying to grok `SequenceTiledCompute`) - -This module implements Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences: https://arxiv.org/abs/2506.13996 - -For integration docs see: https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-pallellism/ - -The other ALST features live inside -https://github.com/snowflakedb/ArcticTraining/blob/main/projects/sequence-parallelism/ - -""" - -from collections import defaultdict -from deepspeed.runtime.utils import see_memory_usage -from deepspeed.sequence.layer import _DimZeroAllToAll -from einops import rearrange -from packaging import version -from torch import Tensor -from torch.utils.data import DataLoader -from typing import Any -from typing import Tuple -import deepspeed.comm as dist -import importlib.metadata -import math -import torch -import torch.distributed.nn - - -class UlyssesSPAttentionHF(torch.nn.Module): - """Re-Implementation of deepspeed.sequence.layer.DistributedAttention. This implementation enforces the input shape - to be standard [sl, bs, hc, hs] form. Any deviation from this shape will raise an error. - - The primary reason for the re-implementation is to make this less error prone, and remove what seemed like bugs in scenarios where batch size > 1 and when using different versions of - flash attention each of which takes different input shape. Those should be handled by - the actual attn implementation, and not by this module. - - This class then has been further adapted to work with HF Transformers' supported attention mechanism. - - Dimension annotation: - bs = bs - hc = head count - hc_l = head count local - hs = head_size - sl = seqlen - sl_l = seqlen local - ws = world_size - em = embedding (hidden size) - em_l = embedding (hidden size) local - - Arguments: - attn: normal attention implementation from transformers.modeling_utils.ALL_ATTENTION_FUNCTIONS - local_seq_length (int): local sequence length per GPU - global_seq_length (int): actual sequence length - batch_size (int): batch size - attn_head_size (int): size of each attention head - attn_head_count (int): total number of attention heads - kv_head_count (int): total number of kv heads - num_hidden_layers (int): total number of layers - process_group (dist.ProcessGroup): Ulysses process group - seq_length_is_variable (bool): whether global seqlen may change between batches - - - Extras: - - set self.skip_all_but_last_attention_debug_mode to True to enable fast debug which will skip calling all core attn layers but the last one, it will produce garbage of course quality-wise. - """ - - def __init__( - self, - attn, - local_seq_length: int, - global_seq_length: int, - batch_size: int, - attn_head_count: int, - attn_head_size: int, - kv_head_count: int, - num_hidden_layers: int, - process_group: dist.ProcessGroup, - seq_length_is_variable: bool = False, - ) -> None: - super().__init__() - self.attn = attn - self.process_group = process_group - self.world_size = dist.get_world_size(process_group) - self.sp_rank = dist.get_rank(process_group) - - self.local_seq_length = local_seq_length - self.global_seq_length = global_seq_length - self.batch_size = batch_size - self.seq_length_is_variable = seq_length_is_variable - - self.attn_head_size = attn_head_size - self.attn_head_count = attn_head_count - self.global_kv_head_count = kv_head_count - - self.num_hidden_layers = num_hidden_layers - self.skip_all_but_last_attention_debug_mode = False - self.rotating_layer_counter = 0 # used for dev work - - self.local_q_head_count = attn_head_count // self.world_size - - # if we have 4 kv heads and sp 8, we need to replicate kv heads 2x - self.kv_replication_factor = self.world_size // kv_head_count - if self.kv_replication_factor > 1: - self.local_kv_head_count = 1 - else: - self.local_kv_head_count = kv_head_count // self.world_size - - transformers_version_min = "4.51.3" - transformers_version_have = importlib.metadata.version("transformers") - if version.parse(transformers_version_have) < version.parse(transformers_version_min): - raise ValueError( - f"transformers>={transformers_version_min} is required, but you have transformers=={transformers_version_have}" - ) - - if self.attn_head_count % self.world_size != 0: - raise ValueError(f"Attention head count {attn_head_count} is not divisible by SP size {self.world_size}") - if not (self.global_kv_head_count % self.world_size == 0 or self.world_size % self.global_kv_head_count == 0): - raise ValueError( - f"KV attention head count {self.global_kv_head_count} is not divisible by SP size {self.world_size} or" - " vice versa") - - # [sl_l bs hc hs] - self.required_query_shape = torch.Size([local_seq_length, batch_size, attn_head_count, attn_head_size]) - self.required_key_value_shape = torch.Size([local_seq_length, batch_size, kv_head_count, attn_head_size]) - - # [sl bs em_l] - self.required_context_shape = torch.Size( - [global_seq_length, batch_size, attn_head_size * attn_head_count // self.world_size]) - - def _combine_local_sequences(self, query, key, value) -> Tuple[Tensor, Tensor, Tensor]: - - def combine_sequence(input, head_type): - """ - expects inputs in shape: [sl_l bs hc hs] - returns output in shape: [sl bs hc_l hs] - - local_head_count could be different for k,v vs q if it's not an MHA situation - """ - if head_type == "q": - local_head_count = self.local_q_head_count - else: # kv - local_head_count = self.local_kv_head_count - - # MQA and some GQA cases: - if self.kv_replication_factor > 1: - # local_head_count *= self.kv_replication_factor - # replicate heads to the kv_replication_factor on hc dimension [sl_l bs hc hs] - so dim=2 - input = input.repeat_interleave(self.kv_replication_factor, dim=2) - - # [sl_l bs hc hs] -> [sl_l bs ws hc_l hs] - input = input.reshape( - [self.local_seq_length, self.batch_size, self.world_size, local_head_count, self.attn_head_size]) - - input = rearrange(input, "sl_l bs ws hc_l hs -> ws sl_l bs hc_l hs").contiguous() - - output = _DimZeroAllToAll.apply(self.process_group, input) - - # [ws sl_l bs hc_l hs] -> [sl bs hc_l hs] - output = output.reshape([self.global_seq_length, *output.shape[2:]]).contiguous() - - # [sl bs hc_l hs] - return output - - return ( - combine_sequence(query, head_type="q"), - combine_sequence(key, head_type="kv"), - combine_sequence(value, head_type="kv"), - ) - - def _partition_global_sequence(self, input) -> Tensor: - """ - expects input in shape: [sl bs em_l] - returns output in shape: [sl_l bs em] - """ - - # [sl bs em_l] -> [ws sl_l bs em_l] - input = input.reshape([ - self.world_size, - self.local_seq_length, - self.batch_size, - self.attn_head_size * self.attn_head_count // self.world_size, - ]).contiguous() - - output = _DimZeroAllToAll.apply(self.process_group, input) - output = rearrange(output, "ws sl_l bs em_l -> sl_l bs ws em_l") - - # [sl_l bs ws em_l] -> [sl_l bs em] - output = output.reshape([*output.shape[:2], -1]).contiguous() - - # [sl_l bs em] - return output - - def forward( - self, - module: torch.nn.Module, - query: Tensor, - key: Tensor, - value: Tensor, - attention_mask: Tensor, - *args: Any, - **kwargs: Any, - ) -> Tensor: - """forward - - Arguments: - query (Tensor): query input to the layer - key (Tensor): key input to the layer - value (Tensor): value input to the layer - attention_mask (Tensor): Attention mask - args: other args - - Returns: - * output (Tensor): context output - """ - # HF incoming shapes are: - # [batch_size, num_heads, seqlen, head_size] - # UlyssesSPAttentionHF expects: - # [seqlen, batch_size, num_heads, head_size] - # print_rank0(f"{query.shape=}") - # print_rank0(f"{key.shape=}") - # print_rank0(f"{value.shape=}") - # print_rank0(f"{self.required_input_shape=}") - current_local_seq_length = query.shape[2] - if self.seq_length_is_variable and current_local_seq_length != self.required_query_shape[0]: - self.local_seq_length = current_local_seq_length - self.global_seq_length = current_local_seq_length * self.world_size - # update the required seqlen shapes - self.required_query_shape = torch.Size([self.local_seq_length] + list(self.required_query_shape)[1:]) - self.required_key_value_shape = torch.Size([self.local_seq_length] + - list(self.required_key_value_shape)[1:]) - self.required_context_shape = torch.Size([self.global_seq_length] + list(self.required_context_shape)[1:]) - - # make the blocks contiguous as early as possible to minimize fragmentation - query = rearrange(query, "bs hc sl hs -> sl bs hc hs") # .contiguous() - key = rearrange(key, "bs hc sl hs -> sl bs hc hs") # .contiguous() - value = rearrange(value, "bs hc sl hs -> sl bs hc hs") # .contiguous() - - # core attn like FA2 expects an unsharded `position_ids` - without which packed samples - # will return loss=nan. - # - # XXX: need to figure out if we can do the same for SDPA - as it doesn't require this and - # wants an attention mask, so possibly doing this for FA2 only? - # - # Ideally we would passing the original unsharded position_ids - but we have no way to pass - # it here as HF Transformers drops unexpected keys in `batch` - so either we need to stash - # it somewhere in UlyssesSPDataLoaderAdapter and retrieve it here or we could gather it once - # per batch and stash it inside `module` arg - I already have a machinery to figure out - # which layer number is being called below in the skip_all_but_last_attention_debug_mode - # code where rotating_layer_counter is used - so we could calculate it on the first layer - # and re-use on the remaining layers - if "position_ids" in kwargs: - position_ids_list = [torch.empty_like(kwargs["position_ids"]) for _ in range(self.world_size)] - dist.all_gather(position_ids_list, kwargs["position_ids"], group=self.process_group) - kwargs["position_ids"] = torch.cat(position_ids_list, dim=1) - - # please don't remove the white-space vertical alignment in the error message - assert query.shape == self.required_query_shape, ( - f"[{dist.get_rank()}]: query input tensor does not match the required shape\n " - f" {self.required_query_shape}:\n {query.shape=}\n {key.shape=}\n {value.shape=}") - assert key.shape == value.shape == self.required_key_value_shape, ( - f"[{dist.get_rank()}]: key or value input tensor does not match the required shape\n " - f" {self.required_key_value_shape}:\n {query.shape=}\n {key.shape=}\n {value.shape=}") - - # expects: [sl_l bs hc hs] - query_layer, key_layer, value_layer = self._combine_local_sequences(query, key, value) - # returns: [sl bs hc_l hs] - - query_layer = rearrange(query_layer, "sl bs hc_l hs -> bs hc_l sl hs").contiguous() - key_layer = rearrange(key_layer, "sl bs hc_l hs -> bs hc_l sl hs").contiguous() - value_layer = rearrange(value_layer, "sl bs hc_l hs -> bs hc_l sl hs").contiguous() - - # crucial in the case of MQA and some GQA cases we need to fix `module.num_key_value_groups` - # XXX: could move this somewhere to do it only once per run - if self.kv_replication_factor > 1: - module.num_key_value_groups = query_layer.size(-3) // key_layer.size(-3) - - if not self.skip_all_but_last_attention_debug_mode: - # expects: [bs hc_l sl hs] - context_layer, attn_weights = self.attn(module, query_layer, key_layer, value_layer, attention_mask, *args, - **kwargs) - # returns [bs sl hc_l hs] - else: - # we need this hack during development in order to be able to check memory fitting w/o - # waiting for 3h to compute 1.5M seqlen attention, because it's quadratic in dense - # attention, so we skip all but the last core attention call - we want the last one to - # still get the memory usage approximately close to the real memory usage. of course - # the loss will be wrong when we do that. - self.rotating_layer_counter = (self.rotating_layer_counter + 1) % self.num_hidden_layers - # we detect the last layer by module counting since we know how many layers there are - if self.rotating_layer_counter % self.num_hidden_layers == 0: - # do the real pass - context_layer, attn_weights = self.attn(module, query_layer, key_layer, value_layer, attention_mask, - *args, **kwargs) - else: - # this feeds bogus data of the right shape - good enough for quick debug - context_layer = rearrange(query_layer, "bs hc_l sl ... -> bs sl hc_l ...") - attn_weights = None - - # [bs sl hc_l hs] -> [sl bs hc_l hs]' - context_layer = rearrange(context_layer, "bs sl ... -> sl bs ...") - context_layer = context_layer.reshape([*context_layer.shape[:2], -1]) - - assert ( - context_layer.shape == self.required_context_shape - ), f"The context shape {context_layer.shape} is not of the expected shape {self.required_context_shape}" - - # expects: [sl bs em_l] - output = self._partition_global_sequence(context_layer) - # returns: [sl_l bs em] - - output = rearrange(output, "sl_l bs ... -> bs sl_l ...") - - output = output.reshape([*output.shape[:2], -1]) - - # expects [bs sl em] - return output, attn_weights - - @classmethod - def register_with_transformers( - cls, - model_name_or_path, - core_attn_implementation, - sequence_parallel_size, - max_length, - micro_batch_size, - seq_length_is_variable=True, - ): - """ - Register "ulysses" attn_implementation with HF transformers and return mpu (Megatron-LM-style parallel state object). - If sequence_parallel_size==1 do nothng and return None. - - """ - if sequence_parallel_size == 1: - return None - - from transformers import AutoConfig - from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS - - import deepspeed.runtime.sequence_parallel.parallel_state_sp as mpu - - mpu.initialize_sequence_parallel(sequence_parallel_size=sequence_parallel_size) - - # we don't have the model yet at this stage - hf_model_config = AutoConfig.from_pretrained(model_name_or_path) - if core_attn_implementation not in ["flash_attention_2", "sdpa"]: - # notes on the excluded ones: - # - eager: The problem is that `eager` wants an attention_mask and it creates the wrong attention mask it seems if we don't provide one - it's possible that we could somehow solve this, but it's also unlikely someone will want to use the slow eager attention with sequence parallelism - # - flex_attention: haven't tried - - raise ValueError( - f"{core_attn_implementation} attn_implementation isn't currently supported by Ulysses sequence" - " parallelism. Set core_attn_implementation arg to either 'flash_attention_2' or 'sdpa'.") - - if core_attn_implementation not in ALL_ATTENTION_FUNCTIONS: - raise ValueError( - f"{core_attn_implementation} is not a valid attn_implementation. The choices are {ALL_ATTENTION_FUNCTIONS.valid_keys()}" - ) - core_attn_function = ALL_ATTENTION_FUNCTIONS[core_attn_implementation] - uattn = UlyssesSPAttentionHF( - attn=core_attn_function, - local_seq_length=max_length // mpu.get_sequence_parallel_world_size(), - global_seq_length=max_length, - batch_size=micro_batch_size, - attn_head_count=hf_model_config.num_attention_heads, - attn_head_size=getattr(hf_model_config, "head_dim", - hf_model_config.hidden_size // hf_model_config.num_attention_heads), - kv_head_count=hf_model_config.num_key_value_heads, - num_hidden_layers=hf_model_config.num_hidden_layers, - process_group=mpu.get_sequence_parallel_group(), - seq_length_is_variable=seq_length_is_variable, - ) - - def uattn_wrapper( - module: torch.nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor, - *args, - **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: - - # We are relaying on position_ids for SP to work so attention_mask has to be None - # the problem is that HF currently doesn't know anything about ALL_ATTENTION_FUNCTIONS["ulysses"] so it doesn't make a special case like for "flash_attention_2" and "sdpa" and it creates an attention mask on the fly and it breaks things. - attention_mask = None - - attn_output, attn_weights = uattn( - module, - query, - key, - value, - attention_mask, - # XXX: fixme - *args, - **kwargs, - ) - return attn_output, attn_weights - - # We don't do: ALL_ATTENTION_FUNCTIONS.register("ulysses", uattn_wrapper) - # The problem with this approach is that we are missing on all the special use cases in HF Transformers that do things like: if self.config._attn_implementation == "flash_attention_2": ... - # So instead we hack `ALL_ATTENTION_FUNCTIONS` to override all existing keys with our implementation, since it only gets used at the point of calling the attention and that's what we want, all other code branches relying on the original core `attn_implementation` will still be executed. This is what we called "Being John Malkovich" - for key in ALL_ATTENTION_FUNCTIONS.keys(): - ALL_ATTENTION_FUNCTIONS[key] = uattn_wrapper - - return mpu - - -class UlyssesSPDataLoaderAdapter: - - def __init__( - self, - dl: DataLoader, - sp_rank: int, - sp_group, - sp_world_size, - device, - ): - """ - This a DataLoader adapter which wraps around any existing DataLoader. It is used in conjunction with Ulysses to perform batch sharding on the sequence dimension. - - It gathers 1 sample from each participating rank, using the DL it wraps, then shards each of them and sends back to the ranks. So that when dl->iter->next is called, we end up with: - - rank 0: getting batch 0 shard 0 - - rank 1: getting batch 0 shard 1 - ... - - rank n: getting batch 0 shard n - which is used to compute the batch (from rank0) using all SP ranks. - - When the next iteration starts and dl->iter->next is called, we end up with: - - rank 0: getting batch 1 shard 0 - - rank 1: getting batch 1 shard 1 - ... - - rank n: getting batch 1 shard n - which is used to compute a second batch (from rank1) using all SP ranks. - - This continues until SP iterations are performed. At this point we need to get more data and so the above repeats. - - The key thing to understand is that all SP ranks participate in processing a single DL sample. So instead of normal DataParallel we perform a sort of SP over DP. - - When SP number of iterations is completed it's an equivalent of performing a single iteration with normal DP. - - If more tokens need to be consumed per step use the gradient accumulation feature. - - Arguments: - - `dl`: an existing DataLoader object to wrap - - `sp_rank`: SP rank - - `sp_group`: SP group - - `sp_world_size`: SP world size - - `device`: cuda device - - Returns: - Another DataLoader object - - Here are the current assumptions on the inputs fetched by dl->iter->next - - the batch is a dict with at least the keys: `input_ids`, `labels`, `position_ids` - but can have any additional keys necessary. - - the tensor values get sharded, the non-tensor values are passed along as is - """ - - self.dl = dl - self.sp_rank = sp_rank - self.sp_group = sp_group - self.sp_world_size = sp_world_size - self.device = device - - self.iter = iter(dl) - self.micro_batches: list[Any] = [] - - def __len__(self): - return len(self.dl) * self.sp_world_size - - def __iter__(self): - return self - - def __next__(self): - if len(self.micro_batches) == 0: - self.refill() - - return self.micro_batches.pop(0) - - def refill(self): - # this will raise StopIteration when empty - batch = next(self.iter) - micro_batches = defaultdict(dict) - # XXX: replace with more efficient all-to-all? - - # we have batches of variable seqlen so in order to do all_gather on batches - we need to know the exact length of each tensor on each rank - seqlen = torch.tensor(batch["input_ids"].shape[1], dtype=torch.int64, device=self.device) - seqlens = [torch.zeros(1, dtype=torch.int64, device=self.device) for _ in range(self.sp_world_size)] - dist.all_gather(seqlens, seqlen, group=self.sp_group) - seqlens = [x[0].item() for x in seqlens] - - for k in batch.keys(): - if torch.is_tensor(batch[k]): - batch[k] = batch[k].to(self.device) - with torch.no_grad(): - tensor_list = [ - torch.zeros((batch[k].shape[0], seqlens[i]), dtype=batch[k].dtype, device=batch[k].device) - for i in range(self.sp_world_size) - ] - dist.all_gather(tensor_list, batch[k], group=self.sp_group) - else: - tensor_list = [None for _ in range(self.sp_world_size)] - dist.all_gather_object(tensor_list, batch[k], group=self.sp_group) - - for rank, tensor in enumerate(tensor_list): - micro_batches[rank][k] = tensor - - del tensor_list - del batch - - for batch in micro_batches.values(): - seq_length = len(batch["input_ids"][0]) - - if seq_length % self.sp_world_size != 0: - raise ValueError(f"batch's seqlen={seq_length} isn't divisible by sp-size={self.sp_world_size}") - chunk_len = seq_length // self.sp_world_size - - # because we have to gather logits from all sp ranks we have to do the loss function ourselves - # therefore remove labels to avoid an attempt to calculate loss by transformers - labels = batch.pop("labels") - labels = torch.nn.functional.pad(labels, (0, 1), value=-100) - batch["shift_labels"] = labels[..., 1:].contiguous() - # free up temp memory - del labels - - # batch sharding - for k in batch.keys(): - # leave non-tensors alone - if not torch.is_tensor(batch[k]): - continue - # at seqlen>10M and 32+ gpus this can take GBs of memory so keep the prefill buffer on cpu - batch[k] = batch[k][:, chunk_len * self.sp_rank:chunk_len * (self.sp_rank + 1)].cpu() - - self.micro_batches.append(batch) - - -def sequence_tiled_compute( - fn, - seqlen, - shards, - kwargs_to_shard, - kwargs_to_pass, - grad_requiring_tensor_key, - compute_params=None, - output_unshard_dimension=1, - output_reduction="mean", -): - """ - This is a wrapper for SequenceTiledCompute which we need since torch.autograd.Function can't work with dicts of tensors (in backward it has to return a grad value and not a dict that may have a non-None grad value). It's also useful for setting default values which we can't do either in torch.autograd.Function. - - Args: - - `fn`: the function to call on sharded inputs - - `seqlen`: total seqlen of the seqlen dimension - - `shards`: how many shards to use - - `kwargs_to_shard`: this dict will be passed to `fn` as `**kwargs` after sharding on seqlen dimension - - `kwargs_to_pass`: this dict will be passed to `fn` as is, as `**kwargs` - - `grad_requiring_tensor_key`: which main key requires grads - - `compute_params`: a list of weights engaged in the compute. Default: `None` (only needed when using DeepSpeed ZeRO) - - `output_reduction`: None, "mean" or "sum": Default: "mean" - - `output_unshard_dimension`: the dimension to concat the outputs on: Default: 1 (seqlen dim) - - Returns: - - unsharded output with an optional reduction applied, depending on the `output_reduction` value: - `None` - return the unsharded output tensor - `"mean"` - apply mean - `"sum"` - apply sum - - Please note that this implementation doesn't require DeepSpeed and can work without it. `compute_params` can remain `None` in such a case. - - """ - args_to_shard = kwargs_to_shard.values() - keys_to_shard = list(kwargs_to_shard.keys()) - args_to_pass = kwargs_to_pass.values() - keys_to_pass = list(kwargs_to_pass.keys()) - - return SequenceTiledCompute.apply( - fn, - seqlen, - shards, - keys_to_shard, - keys_to_pass, - grad_requiring_tensor_key, - compute_params, - output_unshard_dimension, - output_reduction, - *args_to_shard, - *args_to_pass, - ) - - -class SequenceTiledCompute(torch.autograd.Function): - """ - A generic autograd function to perform a tiled compute. - - Please note that this implementation doesn't require DeepSpeed and can work without it. `compute_params` can remain `None` in such a case. - - For an easier to understand example see TiledMLP - which is the same as this autograd function but without the generalization code. - """ - - @staticmethod - def forward( - ctx, - fn, - seqlen, - shards, - keys_to_shard, - keys_to_pass, - grad_requiring_tensor_key, - compute_params, - output_unshard_dimension, - output_reduction, - *args, - ) -> torch.Tensor: - """ - for args and return values see `sequence_tiled_compute`'s doc - - Currently we assume that all kwargs_to_shard values have a shape of `[bs, seqlen, ...]` and we shard on seqlen dimension - """ - ctx.fn = fn - ctx.seqlen = seqlen - ctx.shards = shards - ctx.grad_requiring_tensor_key = grad_requiring_tensor_key - ctx.compute_params = [p for p in compute_params if p.requires_grad] - ctx.output_unshard_dimension = output_unshard_dimension - - with torch.no_grad(): - args = list(args) - ctx.total_args = len(args) - ctx.grad_requiring_tensor_key_index = (keys_to_shard + keys_to_pass).index(grad_requiring_tensor_key) - - kwargs_to_shard = {k: args.pop(0) for k in keys_to_shard} - kwargs_to_pass = {k: args.pop(0) for k in keys_to_pass} - ctx.kwargs_to_shard = kwargs_to_shard - ctx.kwargs_to_pass = kwargs_to_pass - - with torch.no_grad(): - shard_step = math.ceil(seqlen / shards) - output_shards = [] - - for i in range(shards): - output = fn( - **{ - k: v[:, i * shard_step:(i + 1) * shard_step] - for k, v in kwargs_to_shard.items() - }, - **kwargs_to_pass, - ) - output_shards.append(output) - - if output_unshard_dimension == 0: - # this is just the shape=[1] loss use-case, not sure if it's generic enough - output_unsharded = torch.cat([l.unsqueeze(0) for l in output_shards], dim=output_unshard_dimension) - else: - output_unsharded = torch.cat(output_shards, dim=output_unshard_dimension) # .clone().detach() - - if output_reduction is None: - return output_unsharded - elif output_reduction == "mean": - return output_unsharded.mean() - elif output_reduction == "sum": - return output_unsharded.sum() - else: - raise ValueError(f"unknown value {output_reduction}: valid values are: none/mean/sum") - - @staticmethod - def backward(ctx, *grads) -> torch.Tensor: - fn = ctx.fn - shards = ctx.shards - kwargs_to_shard = ctx.kwargs_to_shard - kwargs_to_pass = ctx.kwargs_to_pass - - grad_requiring_tensor_key = ctx.grad_requiring_tensor_key - grad_requiring_tensor_key_index = ctx.grad_requiring_tensor_key_index - compute_params = ctx.compute_params - output_unshard_dimension = ctx.output_unshard_dimension - grad_requiring_tensor = kwargs_to_shard[grad_requiring_tensor_key] - - grad_requiring_tensor_requires_grad = grad_requiring_tensor.requires_grad - grad_requiring_tensor = grad_requiring_tensor.detach() - # detach() unsets `grad_requiring_tensor.requires_grad`, so restore it - grad_requiring_tensor.requires_grad_(grad_requiring_tensor_requires_grad) - - incoming_grad = grads[0] - grad_requiring_tensor_grad = torch.zeros_like(grad_requiring_tensor) - - kwargs_to_shard_shards = { - k: list(torch.chunk(kwargs_to_shard[k], chunks=shards, dim=1)) - for k in kwargs_to_shard.keys() - } - - # if seqlen is not exactly divisible by shards the last step will be shorter than shard_step - shard_step = kwargs_to_shard_shards[grad_requiring_tensor_key][0].numel() - for i in range(shards): - - # when fn involves one or more model weights deepspeed will normally push a grad to - # reduce per sub-module call, so since we only want it to add a grad for the last - # shard's call , we signal to zero not to add new gradients to reduce until the last - # shard when all gradients have been accumulated. An example for such a call is - # `model.lm_head(hidden_states)` - if compute_params is not None: - if i + 1 < shards: - for param in compute_params: - param.ds_grad_is_ready = False - else: - # last shard, can add the grad - for param in compute_params: - param.ds_grad_is_ready = True - - kwargs_to_shard_shard = {k: kwargs_to_shard_shards[k].pop(0) for k in kwargs_to_shard_shards.keys()} - grad_requiring_tensor_shard = kwargs_to_shard_shard[grad_requiring_tensor_key] - - grad_requiring_tensor_shard.requires_grad_(grad_requiring_tensor_requires_grad) - - shard_offset = i * shard_step - # this will enable gradual population of the pre-allocated - # `grad_requiring_tensor_shard.grad` during `torch.autograd.backward` calls - grad_requiring_tensor_shard.grad = (grad_requiring_tensor_grad.view(-1).narrow( - 0, shard_offset, grad_requiring_tensor_shard.numel()).view_as(grad_requiring_tensor_shard)) - - with torch.enable_grad(): - output = fn(**kwargs_to_shard_shard, **kwargs_to_pass) - - if output_unshard_dimension == 0: - # loss use-case - torch.autograd.backward(output, incoming_grad) - else: - incoming_grad_shard = (incoming_grad.view(-1).narrow( - 0, shard_offset, grad_requiring_tensor_shard.numel()).view_as(grad_requiring_tensor_shard)) - torch.autograd.backward(output, incoming_grad_shard) - - # positional args - grad_outputs = [None] * 9 - # inject the grad for the position of forward input that is grad-requiring - arg_outputs = [None] * ctx.total_args - arg_outputs[grad_requiring_tensor_key_index] = grad_requiring_tensor_grad - - return tuple(grad_outputs + arg_outputs) - - -class TiledMLP(torch.autograd.Function): - """ - Perform a tiled MLP computation to massively reduce memory usage needed to compute MLP when using very long sequence lengths - - For a general tiled compute implementation that can handle any `forward` see `SequenceTiledCompute` - - Args: - - fn: the function to call on sharded inputs - - `self`: the MLP nn.Module object - - `x`: the input to MLP.forward (`hidden_states`) - - `shards`: how many shards to use - - compute_params: a list of weights engaged in the compute Default: `None` (only needed when using DeepSpeed ZeRO) - - Returns: - - the computed `hidden_states` - - Here is an example that monkey patches HF Transformers' LLamaMLP: - - def tiled_mlp_forward(self, x): - bs, seqlen, hidden = x.shape - num_shards = math.ceil(seqlen / hidden) - # to avoid deadlocks get all ranks to agree on the same num_shards by using the max value - tensor = torch.tensor(num_shards, device=x.device) - dist.all_reduce(tensor, op=dist.ReduceOp.MAX) - num_shards = tensor.item() - compute_params = [self.down_proj.weight, self.gate_proj.weight, self.up_proj.weight] - - def mlp_forward(self, x): - return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - - return TiledMLP.apply( - mlp_forward, - self, - x, - num_shards, - compute_params, - ) - - # this needs to be done before the model is instantiated - from transformers.models.llama import modeling_llama - modeling_llama.LlamaMLP.forward = tiled_mlp_forward - """ - - @staticmethod - def forward( - ctx, - fn, - self, - x, - shards, - compute_params, - ) -> torch.Tensor: - ctx.fn = fn - ctx.self = self - ctx.shards = shards - ctx.compute_params = [p for p in compute_params if p.requires_grad] - ctx.save_for_backward(x) - - x_shards = list(torch.chunk(x, chunks=shards, dim=1)) - with torch.no_grad(): - output_shards = [fn(self, x_shard) for x_shard in x_shards] - output_unsharded = torch.cat(output_shards, dim=1) - - return output_unsharded - - @staticmethod - def backward(ctx, *grads) -> torch.Tensor: - fn = ctx.fn - (x, ) = ctx.saved_tensors - self = ctx.self - shards = ctx.shards - compute_params = ctx.compute_params - - x_requires_grad = x.requires_grad - x = x.detach() - # detach() unsets `x.requires_grad`, so restore it - x.requires_grad_(x_requires_grad) - - incoming_grad = grads[0] - x_grad = torch.zeros_like(x) - x_shards = list(torch.chunk(x, chunks=shards, dim=1)) - - shard_step = x_shards[0].numel() - for i, x_shard in enumerate(x_shards): - - # Tell deepspeed not to add a new grad to its ipg bucket until the last shard is run - if compute_params is not None: - if i + 1 < shards: - for param in compute_params: - param.ds_grad_is_ready = False - else: - # last shard, can add the grad - for param in compute_params: - param.ds_grad_is_ready = True - - x_shard.requires_grad_(x_requires_grad) - - shard_offset = i * shard_step - x_shard.grad = x_grad.view(-1).narrow(0, shard_offset, x_shard.numel()).view_as(x_shard) - incoming_grad_shard = incoming_grad.view(-1).narrow(0, shard_offset, x_shard.numel()).view_as(x_shard) - with torch.enable_grad(): - output = fn(self, x_shard) - torch.autograd.backward(output, incoming_grad_shard) - - return (None, None, x_grad, None, None) - - -class AutogradComputeMLP(torch.autograd.Function): - """ - This is a simplified example to override the normal MLP via an autograd function - then tiling can be added - this simplified version was useful to detect a leak in Deepspeed, so let's keep it. - - Here is an example of performing the monkey patching on LlamaMLP - - def mlp_forward_new(self, x): - - def mlp_forward(self, x): - return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - - return AutogradComputeMLP.apply(mlp_forward, self, x) - - from transformers.models.llama import modeling_llama - modeling_llama.LlamaMLP.forward = mlp_forward_new - """ - - @staticmethod - def forward( - ctx, - fn, - self, - x, - ) -> torch.Tensor: - ctx.fn = fn - ctx.self = self - ctx.save_for_backward(x) - - with torch.no_grad(): - return fn(self, x) - - @staticmethod - def backward(ctx, *grads) -> torch.Tensor: - fn = ctx.fn - (x, ) = ctx.saved_tensors - self = ctx.self - - x1 = x.detach() - x1.requires_grad = x.requires_grad - with torch.enable_grad(): - output = fn(self, x1) - - torch.autograd.backward(output, grads[0]) - return (None, None, x1.grad, None) - - -########################################################### -### below are older versions that some might still want ### -########################################################### - - -class TiledLoss(torch.autograd.Function): - - @staticmethod - def forward(ctx, loss_fn, logits, vocab_size, shift_labels, shards) -> torch.Tensor: - """ - - This is a memory efficient loss autograd function that takes the existing logits and performs loss calculation in shards. - - This one is an SFT-aware version, therefore it takes care of special cases where the whole shard is made of -100 labels and which requires then a special care. - - Note: logits seqlen dimension doesn't have to be divisible by shards, the last shard will be shorter than the rest. The calculating of the number of shards is in the example. - - Here is an example of using it: - - def loss(self, batch) -> torch.Tensor: - batch = to_device(batch, self.device) - shift_labels = batch.pop("shift_labels") - outputs = self.model(**batch, use_cache=False) - logits = outputs.logits - - if all((shift_labels == -100).squeeze()): - # this is the case where all labels in a micro-batch are -100 (very common for SFT if the seqlen is short) - CE returns `nan` in this case, so we don't want to call loss and instead create a differentiable loss `0` which will also set all the grads to `0` in `backward` - the effect of this is akin to a perfect score where the model needs no adjustment since grads will be all zeros. - loss = (logits.sum() * 0.0).float() - - num_shards: Any = "auto" - if num_shards == "auto": - # parameterize to about 1GB fp32 logits shards - slice_size_in_gb = 1 - size_in_gb = logits.numel() * 4 / 2**30 # fp32 - # the sp shard's seqlen sp shard can be easily not divisible by the derived number of chunked loss shards, so we use the uppper ceiling and allow the last chunk to be shorter than the rest - num_shards = math.ceil(size_in_gb / slice_size_in_gb) - # print(f"derived {num_shards} shards for size {size_in_gb}GB") - if num_shards > 1: - # if shards == 1 this will lead to a higher memory usage then calling the normal loss function, so don't do that. - loss = TiledLoss.apply( - self.model_unwrapped.loss_function, - logits, - self.model_unwrapped.config.vocab_size, - shift_labels, - num_shards, - ) - else: - loss = self.model_unwrapped.loss_function( - logits=logits, - labels=None, - vocab_size=self.model_unwrapped.config.vocab_size, - shift_labels=shift_labels, - ) - - return loss - - - """ - ctx.save_for_backward(logits, shift_labels) - ctx.loss_fn = loss_fn - ctx.vocab_size = vocab_size - ctx.shards = shards - - with torch.no_grad(): - seqlen = shift_labels.shape[1] - shard_step = math.ceil(seqlen / shards) - loss_shards = [] - total_good_items = 0 - - # since -100s are ignored we have to perform a weighted average on each loss slice as each slice may contribute a different number of non- -100 labels - # if seqlen / shards != 0 - the last chunk is just shorter than the rest but no data is ignored - for i in range(shards): - # XXX: here and everywhere don't make a copy, pass the slice or perhaps narrow/view? - shift_labels_shard = shift_labels[:, i * shard_step:(i + 1) * shard_step] - if all((shift_labels_shard == -100).squeeze()): - continue # ignore this shard - loss_shard = loss_fn( - logits=logits[:, i * shard_step:(i + 1) * shard_step, :], - labels=None, - vocab_size=vocab_size, - shift_labels=shift_labels_shard, - ) - good_items = sum((shift_labels_shard != -100).squeeze()) - loss_shards.append(loss_shard * good_items) - total_good_items += good_items - total_loss = torch.cat([l.unsqueeze(0) for l in loss_shards], dim=0).sum() - weighted_loss = total_loss / total_good_items - - return weighted_loss - - @staticmethod - def backward(ctx, *grads) -> torch.Tensor: - logits, shift_labels = ctx.saved_tensors - loss_fn = ctx.loss_fn - vocab_size = ctx.vocab_size - shards = ctx.shards - - grad = grads[0] - logits_grad = torch.zeros_like(logits) - logits_shards = list(torch.chunk(logits, chunks=shards, dim=1)) - shift_labels_shards = list(torch.chunk(shift_labels, chunks=shards, dim=1)) - - # if seqlen is not exactly divisible by shards the last step will be shorter than shard_step - shard_step = logits_shards[0].numel() - for i in range(shards): - logits_shard = logits_shards.pop(0) - shift_labels_shard = shift_labels_shards.pop(0) - - shard_offset = i * shard_step - # this will enable gradual population of the pre-allocated `logits_shard.grad` during `torch.autograd.backward` calls - logits_shard.grad = (logits_grad.view(-1).narrow(0, shard_offset, - logits_shard.numel()).view_as(logits_shard)) - - with torch.enable_grad(): - if all((shift_labels_shard == -100).squeeze()): - # fake loss calculation, since CE will return nan, but grads will be set - # a normal loss_fn upcasts logits to float so match it - loss_shard = (logits_shard.sum() * 0.0).float() - else: - loss_shard = loss_fn( - logits=logits_shard.requires_grad_(), - labels=None, - vocab_size=vocab_size, - shift_labels=shift_labels_shard, - ) - - torch.autograd.backward(loss_shard, grad) - - logits_grad /= shards - - # only logits (2nd arg) needs grads - return None, logits_grad, None, None, None - - -# This is the original implementation/integration of UlyssesSP into the training loop, which was superseded by using UlyssesSPDataLoaderAdapter which did all the sharding and pull the shards from the DL -# -# There are 2 issues with this implementation: -# - it's complex and difficult to integrate into various training scenarios -# - it could lead to a huge number of tokens per step - e.g. 32 ranks of 15M seqlen -> 0.5B token step - which is very wasteful -# -# Therefore if you want to use UlyssesSP via UlyssesSPFwdLossBwdWithLogits with its fwd/loss/bwd for those don't want to use UlyssesSPDataLoaderAdapter - here is how it should be installed into the sub-trainer class: -# class SFTTrainer(Trainer): -# def sp_fwd_loss_bwd(self, batch) -> torch.Tensor: -# batch = to_device(batch, self.device) -# -# from arctic_training.trainer.trainer import UlyssesAttentionHFFwdLossBwdWithLogits -# ulysses = UlyssesAttentionHFFwdLossBwdWithLogits( -# model=self.model, -# model_unwrapped=self.model_unwrapped, -# device=self.device, -# num_loss_logit_shards="auto", -# ) -# return ulysses.sp_fwd_loss_bwd(batch) - - -class UlyssesSPFwdLossBwdWithLogits: - - def __init__(self, model, model_unwrapped, device, num_loss_logit_shards="auto", **kwargs): - - self.model = model - self.model_unwrapped = model_unwrapped - self.device = device - self.num_loss_logit_shards = num_loss_logit_shards - self.kwargs = kwargs - - from deepspeed.utils import groups - - self.sp_group = groups._get_sequence_parallel_group() - self.sp_world_size = groups._get_sequence_parallel_world_size() - self.sp_rank = groups._get_sequence_parallel_rank() - - def sp_fwd_loss_bwd(self, batch) -> torch.Tensor: - - see_memory_usage(f"entered sp_fwd_loss_bwd", force=True) - - # ensure shapes are correct - if not (batch["input_ids"].shape == batch["position_ids"].shape == batch["labels"].shape): - raise ValueError( - f'Borked batch {batch["input_ids"].shape=} != {batch["position_ids"].shape=} !=' - f' {batch["labels"].shape=}) in DataLoader->iter->next, cannot continue with Ulysses Sequence' - " parallelism") - - # gather DL batches into super-batches - # Important: DL doesn't always yield max_length batches. Different ranks may have different seqlen and each could be <= max_length (but always divisible by 256) - - micro_batches: list[Any] = defaultdict(dict) - # Efficient gathering of batch inputs across ranks: - # The problem is that our DL doesn't guarantee the same seqlen on all ranks and may give, 3x 1024 and 1x 768 on 4 gpus for max_length 1024. so 3 options we have to be able to gather batches are: - # 1. use all_gather_object - which allows different shapes - but potentially introducing an undesired overhead - 2x pickle calls - # 2. use all_gather and change DL pad to make sure that all ranks always get the same input shape - this creates its own overhead since if we say have ranks with seqlen 512, 768, 1024, 1024 - now we will need to process 4x 1024 seqlens - # 3. use all_gather and post gathering truncate tensors to their intended length - another overhead of allocating and truncating tensors - # using approach (1) for now but might want to benchmark later the other 2 approaches - - # XXX: if using all_gather_object we can gather the whole batch at once and not per-key! so can drop the loop for that approach - - # we have batches of variable seqlen so in order to do all_gather on batches - we need to know the exact length of each tensor on each rank - seqlen = torch.tensor(batch["input_ids"].shape[1], dtype=torch.int64, device=self.device) - # print(seqlen) - seqlens = [torch.zeros(1, dtype=torch.int64, device=self.device) for _ in range(self.sp_world_size)] - dist.all_gather(seqlens, seqlen, group=self.sp_group) - seqlens = [x[0].item() for x in seqlens] - - for k in batch.keys(): - batch[k] = batch[k].to(self.device) - with torch.no_grad(): - tensor_list = [ - torch.zeros((batch[k].shape[0], seqlens[i]), dtype=batch[k].dtype, device=batch[k].device) - for i in range(self.sp_world_size) - ] - dist.all_gather(tensor_list, batch[k], group=self.sp_group) - - # gathering on the data dimension - # will be concatenating and later splitting again for the more general case - # batch[k] = torch.cat(tensor_list, dim=1) - for rank, tensor in enumerate(tensor_list): - micro_batches[rank][k] = tensor - - del tensor_list - del batch - - # we need to chunk twice - each time on SP size level - # - the first time is because we artificially made the seqlen SP-times longer - # - the second time is because of the Ulysses algorithm - - see_memory_usage("after gathering", force=False) - - self.model.set_gradient_accumulation_boundary(False) - - losses = [] - for sub_step_id in range(self.sp_world_size): - batch = micro_batches[sub_step_id] - seq_length = len(batch["input_ids"][0]) - - if seq_length % self.sp_world_size != 0: - raise ValueError( - f"{sub_step_id=}: batch's seqlen={seq_length} isn't divisible by sp-size={self.sp_world_size}") - chunk_len = int(seq_length / self.sp_world_size) - - # to enable the correct mean calculation across shards before sharding the micro batch: - # 1. count the number of non- `-100`` elements per shard - # 2. and subtract one more element because of label shifting - non_skipped_items = {} - for rank in range(self.sp_world_size): - non_skipped = (batch["labels"][:, chunk_len * rank:chunk_len * (rank + 1)] != -100).sum().item() - if non_skipped > 1: - non_skipped -= 1 - non_skipped_items[rank] = non_skipped - - # because we have to gather logits from all sp ranks we have to do the loss function ourselves - # therefore remove labels to avoid an attempt to calculate loss by transformers - labels = batch.pop("labels") - labels = torch.nn.functional.pad(labels, (0, 1), value=-100) - batch["shift_labels"] = labels[..., 1:].contiguous() - # free up temp memory - del labels - - # batch sharding - for k in batch.keys(): - batch[k] = batch[k][:, chunk_len * self.sp_rank:chunk_len * (self.sp_rank + 1)].to(self.device) - - shift_labels = batch.pop("shift_labels") - - outputs = self.forward(batch) - loss = self.compute_loss(labels=None, shift_labels=shift_labels) - - # free up temp mem (e.g. outputs.logits are huge) - del outputs - - # differentiable loss aggregation across ranks - losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=self.sp_group) - - # since each shard may have a variable number of skipped elemented - need to calculate a weighted mean depending on each rank's contribution - this will also take care of loss=0 when all elements are -100 in a shard - # XXX: not expecting a total of 0-non-skipped items for div - loss = sum(losses_per_rank[rank] * non_skipped_items[rank] - for rank in range(self.sp_world_size)) / sum(non_skipped_items.values()) - - self.backward() - - losses.append(loss.detach().item()) - - self.model.set_gradient_accumulation_boundary(True) - - # for per-iteration reporting - if len(losses) == 0: - loss = float("nan") - else: - loss = sum(losses) / len(losses) - - return loss - - def forward(self, batch): - # critical: the labels shouldn't be in batch - outputs = self.model(**batch, use_cache=False) - self.logits = outputs.logits - return outputs - - def compute_loss(self, labels, shift_labels): - if all((shift_labels == -100).squeeze()): - # this is the case where all labels in a micro-batch are -100 (very common for SFT) - CE returns `nan` in this case, so we don't want to call loss and instead create a differentiable loss `0` which will also set all the grads to `0` in `backward` - the effect of this is akin to a perfect score where the model needs no adjustment since grads will be all zeros. - # XXX: should this be float and not the original dtype? - loss = (self.logits.sum() * 0.0).float() - else: - if self.num_loss_logit_shards == "auto": - # parameterize to about 1GB fp32 logits shards - slice_size_in_gb = 1 # XXX: make configurable? - size_in_gb = self.logits.numel() * 4 / 2**30 # fp32 - # the sp shard's seqlen sp shard can be easily not divisible by the derived number of chunked loss shards, so we use the uppper ceiling and allow the last chunk to be shorter than the rest - self.num_loss_logit_shards = math.ceil(size_in_gb / slice_size_in_gb) - # print(f"derived {self.num_loss_logit_shards} shards for size {size_in_gb}GB") - if self.num_loss_logit_shards > 1: - loss = TiledLoss.apply( - self.model_unwrapped.loss_function, - self.logits, - self.model_unwrapped.config.vocab_size, - shift_labels, - self.num_loss_logit_shards, - ) - else: - # XXX: for some reason this fails with zero1 - loss = self.model_unwrapped.loss_function( - logits=self.logits, - labels=None, - vocab_size=self.model_unwrapped.config.vocab_size, - shift_labels=shift_labels, - ) - - self.loss = loss - return loss - - def backward(self): - self.model.backward(self.loss) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/sparse_tensor.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/sparse_tensor.py deleted file mode 100644 index 291ba5f0c78668b793bea12749bd0a188b680ece..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/sparse_tensor.py +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Implementation of a compressed sparse tensor. Similar in -functionality to TensorFlow's IndexedSlices implementation. -""" - -import torch - - -class SparseTensor(object): - """ Compressed Sparse Tensor """ - - def __init__(self, dense_tensor=None): - self.orig_dense_tensor = dense_tensor - self.dtype = self.orig_dense_tensor.dtype - self.is_sparse = dense_tensor.is_sparse - if dense_tensor is not None: - if dense_tensor.is_sparse: - dense_tensor = dense_tensor.coalesce() - self.indices = dense_tensor.indices().flatten() - self.values = dense_tensor.values() - else: - result = torch.sum(dense_tensor, dim=1) - self.indices = result.nonzero().flatten() - self.values = dense_tensor[self.indices] - self.dense_size = list(dense_tensor.size()) - else: - self.indices = None - self.values = None - self.dense_size = None - - def to_coo_tensor(self): - return torch.sparse_coo_tensor(self.indices.unsqueeze(0), self.values, self.dense_size) - - @staticmethod - def type(): - return "deepspeed.SparseTensor" - - def to_dense(self): - it = self.indices.unsqueeze(1) - full_indices = torch.cat([it for _ in range(self.dense_size[1])], dim=1) - return self.values.new_zeros(self.dense_size).scatter_add_(0, full_indices, self.values) - - def sparse_size(self): - index_size = list(self.indices.size()) - index_size = index_size[0] - value_size = list(self.values.size()) - value_size = value_size[0] * value_size[1] - dense_size = self.dense_size[0] * self.dense_size[1] - return index_size + value_size, dense_size - - def add(self, b): - assert self.dense_size == b.dense_size - self.indices = torch.cat([self.indices, b.indices]) - self.values = torch.cat([self.values, b.values]) - - def __str__(self): - sparse_size, dense_size = self.sparse_size() - return "DeepSpeed.SparseTensor(indices_size={}, values_size={}, " \ - "dense_size={}, device={}, reduction_factor={})".format( - self.indices.size(), self.values.size(), self.dense_size, - self.indices.get_device(), dense_size / sparse_size - ) - - def __repr__(self): - return self.__str__() diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/state_dict_factory.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/state_dict_factory.py deleted file mode 100644 index c2db85d1ba58f16dbe59b1dcd3d3066ab11ff6d2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/state_dict_factory.py +++ /dev/null @@ -1,427 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import os -import copy -import collections -import json -from abc import ABC, abstractmethod - -from deepspeed.utils import logger -from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine - -from .weight_quantizer import WeightQuantization - -AUTO_MODULE_KEY = 'auto' - - -class SDLoaderFactory: - - @staticmethod - def get_sd_loader_json(json_file, checkpoint_engine): - if isinstance(json_file, str): - with open(json_file) as f: - data = json.load(f) - else: - assert isinstance(json_file, dict) - data = json_file - sd_type = data['type'] - ckpt_list = data['checkpoints'] - version = data['version'] - ckpt_type = data.get('parallelization', 'pp') - mp_size = data.get('mp_size', 0) - if sd_type.lower() in ['bloom', 'ds_model']: - return data - return SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine, sd_type, version) - - @staticmethod - def get_sd_loader(ckpt_list, checkpoint_engine, sd_type='Megatron', version=None): - if sd_type == 'Megatron': - return MegatronSDLoader(ckpt_list, version, checkpoint_engine) - else: - assert False, '{} checkpoint type is not supported'.format(sd_type) - - -class SDLoaderBase(ABC): - - def __init__(self, ckpt_list, version, checkpoint_engine): - self.module_key = None - self.ckpt_list = ckpt_list - self.version = version - self.checkpoint_engine = TorchCheckpointEngine() if checkpoint_engine is None else checkpoint_engine - self.check_ckpt_list() - - def load(self, - mp_world_size, - mp_rank, - module_key=AUTO_MODULE_KEY, - is_pipe_parallel=False, - quantize=False, - quantize_bits=8, - quantize_groups=64, - mlp_extra_grouping=True): - self.module_key = module_key - num_ckpt = len(self.ckpt_list) - idx = mp_rank * num_ckpt // mp_world_size - """ We have multiple cases to handle here for both training and inference: - 1. PipeModule loading mp_rank_*.pt files, is_pipe_parallel=True, module_key is not None - a. if no mp_size/pp_size resizing occurs, for both training & inference, loading - the mp_rank related checkpoint directly. - b. if has mp_size/pp_size resizing, only Megatron model inference is supported, - in this case each mp_rank_*.pt have same content, we will load the first checkpoint - file (idx=0), to avoid idx exceeding file list boundary. - - 2. PipeModule loading layer_*.pt files, is_pipe_parallel=True, module_key is None - a. if no mp_size resizing occurs, for both training & inference, loading - the mp_rank related checkpoint directly. - b. if has mp_size resizing, only Megatron model inference is supported, - checkpoint file(s) will be merged/split according to mp_rank, mp_world_size and - checkpoint file list. - - 3. Non-PipeModule loading mp_rank_*.pt files, is_pipe_parallel=False - Same with case (2). - """ - if is_pipe_parallel and module_key is not None and mp_world_size != num_ckpt: - mp_world_size = num_ckpt - idx = 0 - - load_path = self.ckpt_list[idx] - - merge_count = 1 - if num_ckpt == mp_world_size: - assert os.path.exists(load_path) - #logger.info(f'rank: {mp_rank} loading checkpoint: {load_path}') - sd = self.checkpoint_engine.load(load_path, map_location=lambda storage, \ - loc: storage) - - if quantize: - quantizer = WeightQuantization(mlp_extra_grouping=mlp_extra_grouping, mp_size=mp_world_size) - sd_module, all_scales = quantizer.sd_quantize_megatron(self.get_module(sd), quantize_bits, - quantize_groups) - self.set_module(sd, sd_module) - else: - all_scales = None - elif num_ckpt > mp_world_size: - sd, all_scales, merge_count = self.merge_state_dict(mp_world_size, mp_rank, quantize, \ - quantize_bits, quantize_groups, mlp_extra_grouping) - else: - sd, all_scales = self.split_state_dict(mp_world_size, mp_rank, quantize, quantize_bits, \ - quantize_groups, mlp_extra_grouping) - return load_path, sd, (all_scales, merge_count) - - def get_merge_state_dicts(self, mp_world_size, mp_rank): - num_ckpt = len(self.ckpt_list) - assert num_ckpt % mp_world_size == 0, 'Invalid checkpoints and world size for sd merge' - - num_to_merge = num_ckpt // mp_world_size - ckpt_list = [self.ckpt_list[i] for i in range(num_to_merge * mp_rank, num_to_merge * (mp_rank + 1))] - - logger.info(f"mp_rank: {mp_rank}, ckpt_list: {ckpt_list}") - sd_list = [self.checkpoint_engine.load(ckpt, map_location=lambda storage, loc: storage) for ckpt in ckpt_list] - return sd_list - - def get_split_state_dict(self, mp_world_size, mp_rank): - num_ckpt = len(self.ckpt_list) - assert mp_world_size % num_ckpt == 0, 'Invalid checkpoints and world size for sd split' - - num_to_split = mp_world_size // num_ckpt - ckpt_index = mp_rank // num_to_split - ckpt_offset = mp_rank % num_to_split - - logger.info(f"mp_rank: {mp_rank}, ckpt_list: {self.ckpt_list[ckpt_index]}, offset: {ckpt_offset}") - - sd = self.checkpoint_engine.load(self.ckpt_list[ckpt_index], map_location=lambda storage, loc: storage) - - return sd, num_to_split, ckpt_offset - - def _choose_module_key(self, sd): - assert not ('module' in sd - and 'model' in sd), "checkpoint has both 'model' and 'module' keys, not sure how to proceed" - assert 'module' in sd or 'model' in sd, "checkpoint contains neither 'model' or 'module' keys, not sure how to proceed" - if 'module' in sd: - return 'module' - elif 'model' in sd: - return 'model' - - def get_module(self, sd): - if self.module_key is None: - return sd - elif self.module_key == AUTO_MODULE_KEY: - return sd[self._choose_module_key(sd)] - else: - return sd[self.module_key] - - def set_module(self, sd, module): - if self.module_key is None: - sd = module - elif self.module_key == AUTO_MODULE_KEY: - sd[self._choose_module_key(sd)] = module - else: - sd[self.module_key] = module - return sd - - def check_ckpt_list(self): - #logger.info(f'checkpoint file list: {self.ckpt_list}') - assert len(self.ckpt_list) > 0 - - sd = self.checkpoint_engine.load(self.ckpt_list[0], map_location=lambda storage, loc: storage) - - # check checkpoint count is same with saved mp_world_size - if 'mp_world_size' in sd.keys(): - assert len(self.ckpt_list) == sd[ - 'mp_world_size'], f"checkpoint count {len(self.ckpt_list)} is different from saved mp_world_size {sd['mp_world_size']}" - - @abstractmethod - def merge_state_dict(self, mp_world_size, mp_rank, quantize, quantize_bits, groups, mlp_extra_grouping): - pass - - @abstractmethod - def split_state_dict(self, mp_world_size, mp_rank, quantize, quantize_bits, groups, mlp_extra_grouping): - pass - - @abstractmethod - def sanity_check(self, ckpt_file_name): - pass - - -class MegatronSDLoader(SDLoaderBase): - - def __init__(self, ckpt_list, version, checkpoint_engine): - super().__init__(ckpt_list, version, checkpoint_engine) - """ - ## Q/K/V data need special processing - key: transformer.layers.0.attention.query_key_value.weight, shape: torch.Size([3192, 4256]) - key: transformer.layers.0.attention.query_key_value.bias, shape: torch.Size([3192]) - - ## merge or split on axis=0 - key: word_embeddings.weight, shape: torch.Size([12672, 4256]) - key: transformer.layers.0.mlp.dense_h_to_4h.bias, shape: torch.Size([4256]) - key: transformer.layers.0.mlp.dense_h_to_4h.weight, shape: torch.Size([4256, 4256]) - - ## merge or split on axis=1 - key: transformer.layers.0.attention.dense.weight, shape: torch.Size([4256, 1064]) - key: transformer.layers.0.mlp.dense_4h_to_h.weight, shape: torch.Size([4256, 4256]) - - ## no change required - key: transformer.layers.0.mlp.dense_4h_to_h.bias, shape: torch.Size([4256]) - key: transformer.final_layernorm.weight, shape: torch.Size([4256]) - key: transformer.final_layernorm.bias, shape: torch.Size([4256]) - key: transformer.layers.0.attention.dense.bias, shape: torch.Size([4256]) - key: transformer.layers.0.post_attention_layernorm.weight, shape: torch.Size([4256]) - key: transformer.layers.0.post_attention_layernorm.bias, shape: torch.Size([4256]) - key: transformer.layers.0.input_layernorm.weight, shape: torch.Size([4256]) - key: transformer.layers.0.input_layernorm.bias, shape: torch.Size([4256]) - key: position_embeddings.weight, shape: torch.Size([1024, 4256]) - """ - - def merge_query_key_value(self, param_list, ckpt_ver): - """ - Up to now we found 3 Q/K/V parameter formats in different Megatron checkpoint versions: - - 1. version 0, there is no version information saved in checkpoint. - format: [(3 * np * hn), h] - 2. version 1.0 - format: [(np * hn * 3), h] - 3. version 2.0 - format: [(np * 3 * hn), h] - - h: hidden size - n: number of attention heads - p: number of model parallel partitions - np: n/p - hn: h/n - """ - - new_qkv = None - if ckpt_ver == 0: - # [(3 * np * hn), h] - assert param_list[0].shape[0] % 3 == 0 - size_qkv = param_list[0].shape[0] // 3 - split_tensors = [torch.split(param, size_qkv, dim=0) for param in param_list] - - tensors = [] - for i in range(3): - tensor_tuple = [t[i] for t in split_tensors] - tensors.append(torch.cat(tensor_tuple, axis=0)) - new_qkv = torch.cat(tensors, axis=0) - elif ckpt_ver == 1.0 or ckpt_ver == 2.0: - # [(np * hn * 3), h] or [(np * 3 * hn), h] - new_qkv = torch.cat(param_list, axis=0) - else: - assert False, f'checkpoint version: {ckpt_ver} is not supported' - - return new_qkv - - def split_query_key_value(self, param, num_to_split, offset, ckpt_ver): - """ - Up to now we found 3 Q/K/V parameter formats in different Megatron checkpoint versions: - - 1. version 0, there is no version information saved in checkpoint. - format: [(3 * np * hn), h] - 2. version 1.0 - format: [(np * hn * 3), h] - 3. version 2.0 - format: [(np * 3 * hn), h] - - h: hidden size - n: number of attention heads - p: number of model parallel partitions - np: n/p - hn: h/n - """ - - new_qkv = None - if ckpt_ver == 0: - # [(3 * np * hn), h] - assert param.shape[0] % 3 == 0 - size_qkv = param.shape[0] // 3 - split_tensors = torch.split(param, size_qkv, dim=0) - - assert split_tensors[0].shape[0] % num_to_split == 0 - split_size = split_tensors[0].shape[0] // num_to_split - - tensors = [] - for i in range(3): - tensors.append(torch.split(split_tensors[i], split_size, dim=0)[offset]) - new_qkv = torch.cat(tensors, axis=0) - elif ckpt_ver == 1.0 or ckpt_ver == 2.0: - # [(np * hn * 3), h] or [(np * 3 * hn), h] - assert param.shape[0] % num_to_split == 0 - size_qkv = param.shape[0] // num_to_split - split_tensors = torch.split(param, size_qkv, dim=0) - new_qkv = split_tensors[offset] - else: - assert False, f'checkpoint version: {ckpt_ver} is not supported' - - return new_qkv - - def merge_state_dict(self, - mp_world_size, - mp_rank, - quantize=False, - quantize_bits=8, - groups=64, - mlp_extra_grouping=True): - self.sanity_check(self.ckpt_list[0]) - - sd_list = self.get_merge_state_dicts(mp_world_size, mp_rank) - ds_sd = copy.deepcopy(sd_list[0]) - new_client_sd = collections.OrderedDict() - - client_sd_list = [self.get_module(sd) for sd in sd_list] - keys = client_sd_list[0].keys() - - ckpt_ver = self.get_checkpoint_version(ds_sd) - logger.info(f"checkpoint version: {ckpt_ver}") - if quantize: - quantizer = WeightQuantization(mlp_extra_grouping=mlp_extra_grouping, mp_size=mp_world_size) - - for key in keys: - value_list = [sd[key] for sd in client_sd_list] - - if "attention.dense.weight" in key or "mlp.dense_4h_to_h.weight" in key: - if quantize: - value_list = quantizer.Quantize(value_list, quantize_bits, groups, key=key, merge_dim=1) - new_client_sd[key] = torch.cat(value_list, axis=1) - elif "attention.query_key_value" in key: - if quantize and "attention.query_key_value.weight" in key: - value_list = quantizer.Quantize(value_list, quantize_bits, groups, key=key) - new_client_sd[key] = torch.cat(value_list, axis=0) - else: - if quantize: - new_client_sd[key] = torch.cat(value_list, axis=0) - else: - new_client_sd[key] = self.merge_query_key_value(value_list, ckpt_ver) - elif "mlp.dense_h_to_4h.weight" in key or "word_embeddings.weight" in key or "mlp.dense_h_to_4h.bias" in key: - if quantize and "mlp.dense_h_to_4h.weight" in key: - value_list = quantizer.Quantize(value_list, quantize_bits, groups, key=key) - new_client_sd[key] = torch.cat(value_list, axis=0) - else: - new_client_sd[key] = value_list[0] - if quantize: - all_scales = quantizer.merge_scales() - ds_sd = self.set_module(ds_sd, new_client_sd) - - return ds_sd, (all_scales if quantize else None), len(client_sd_list) - - def split_state_dict(self, - mp_world_size, - mp_rank, - quantize=False, - quantize_bits=8, - groups=64, - mlp_extra_grouping=True): - #self.sanity_check(self.ckpt_list[0]) - - sd, num_to_split, ckpt_offset = self.get_split_state_dict(mp_world_size, mp_rank) - ds_sd = copy.deepcopy(sd) - new_client_sd = collections.OrderedDict() - - client_sd = self.get_module(sd) - - ckpt_ver = self.get_checkpoint_version(ds_sd) - logger.info(f"checkpoint version: {ckpt_ver}") - - if quantize: - quantizer = WeightQuantization(mlp_extra_grouping=mlp_extra_grouping, mp_size=mp_world_size) - - for key in client_sd.keys(): - value = client_sd[key] - - if "attention.dense.weight" in key or "mlp.dense_4h_to_h.weight" in key: - assert value.shape[1] % num_to_split == 0 - split_size = value.shape[1] // num_to_split - if quantize: - q_vals = quantizer.Quantize([value], quantize_bits, groups, key) - value = q_vals[0] - new_client_sd[key] = torch.split(value, split_size, dim=1)[ckpt_offset] - elif "attention.query_key_value" in key: - if quantize and "attention.query_key_value.weight" in key: - q_vals = quantizer.Quantize([value], quantize_bits, groups, key) - value = q_vals[0] - new_client_sd[key] = self.split_query_key_value(value, num_to_split, ckpt_offset, ckpt_ver) - elif "mlp.dense_h_to_4h.weight" in key or "word_embeddings.weight" in key or "mlp.dense_h_to_4h.bias" in key or "final_linear.weight" in key: - assert value.shape[0] % num_to_split == 0 - split_size = value.shape[0] // num_to_split - if quantize and "mlp.dense_h_to_4h.weight" in key: - q_vals = quantizer.Quantize([value], quantize_bits, groups, key) - value = q_vals[0] - new_client_sd[key] = torch.split(value, split_size, dim=0)[ckpt_offset] - else: - new_client_sd[key] = value - - if quantize: - all_scales = quantizer.merge_scales_split(num_to_split) - - ds_sd = self.set_module(ds_sd, new_client_sd) - - return ds_sd, (all_scales if quantize else None) - - def sanity_check(self, ckpt_file_name): - keys_to_check = [ - "attention.dense.weight", "mlp.dense_4h_to_h.weight", "attention.query_key_value", - "mlp.dense_h_to_4h.weight", "mlp.dense_h_to_4h.bias" - ] - - sd = self.checkpoint_engine.load(ckpt_file_name, map_location=lambda storage, loc: storage) - - # partial_key is a sub-string of one key in the sd - def check_key_exist(partial_key, sd): - keys = sd.keys() - found = False - for k in keys: - if partial_key in k: - found = True - break - return found - - for key in keys_to_check: - assert check_key_exist(key, - self.get_module(sd)), f'key: {key} is not found in the checkpoint {ckpt_file_name}' - - def get_checkpoint_version(self, state_dict): - # Use 0 if version info doesn't exist - return self.version if self.version is not None else state_dict.get('checkpoint_version', 0) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__init__.py deleted file mode 100644 index 006dfd6dcbc6f7675733ba10938552cdc13f7bcb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -from .utils import MIN_SWAPPABLE_BYTES diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 2293668ed88341df3f4604eac7d7cdf0ca1c5056..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/aio_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/aio_config.cpython-310.pyc deleted file mode 100644 index 2263b50560c511b5296f3cdb55c6ded2c3062de9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/aio_config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/async_swapper.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/async_swapper.cpython-310.pyc deleted file mode 100644 index 262f5bd2bae2b7b124a4196b191ba501f58dfcca..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/async_swapper.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/constants.cpython-310.pyc deleted file mode 100644 index 992b88f494df544aca5df256d8044ece7ff7484f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/constants.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/optimizer_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/optimizer_utils.cpython-310.pyc deleted file mode 100644 index 5966fd5b9962585be989925b335d63279e85f64d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/optimizer_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_optimizer_swapper.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_optimizer_swapper.cpython-310.pyc deleted file mode 100644 index 5c4c0fb6bb5dcd81c69334166b0cc8cdb3d5f376..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_optimizer_swapper.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_param_swapper.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_param_swapper.cpython-310.pyc deleted file mode 100644 index 0f32fa9b9daf3f4ca3c6ff531dd7f75a1b08bccd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_param_swapper.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/pipelined_optimizer_swapper.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/pipelined_optimizer_swapper.cpython-310.pyc deleted file mode 100644 index de18b50786787b23c393e2fe9b059378a53219e4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/pipelined_optimizer_swapper.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index b0b8c185cf79821179ca005914c3e87e6e94084b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/aio_config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/aio_config.py deleted file mode 100644 index be6c7d93c86a4f5e268391f717217db1a687d0dc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/aio_config.py +++ /dev/null @@ -1,43 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.runtime.config_utils import get_scalar_param -from deepspeed.runtime.swap_tensor.constants import * -from deepspeed.accelerator import get_accelerator - -AIO_DEFAULT_DICT = { - AIO_BLOCK_SIZE: AIO_BLOCK_SIZE_DEFAULT, - AIO_QUEUE_DEPTH: AIO_QUEUE_DEPTH_DEFAULT, - AIO_INTRA_OP_PARALLELISM: AIO_INTRA_OP_PARALLELISM_DEFAULT, - AIO_SINGLE_SUBMIT: AIO_SINGLE_SUBMIT_DEFAULT, - AIO_OVERLAP_EVENTS: AIO_OVERLAP_EVENTS_DEFAULT, - AIO_USE_GDS: AIO_USE_GDS_DEFAULT -} - - -def get_aio_config(param_dict): - if AIO in param_dict.keys() and param_dict[AIO] is not None: - aio_dict = param_dict[AIO] - aio_config = { - AIO_BLOCK_SIZE: - get_scalar_param(aio_dict, AIO_BLOCK_SIZE, AIO_BLOCK_SIZE_DEFAULT), - AIO_QUEUE_DEPTH: - get_scalar_param(aio_dict, AIO_QUEUE_DEPTH, AIO_QUEUE_DEPTH_DEFAULT), - AIO_INTRA_OP_PARALLELISM: - get_scalar_param(aio_dict, AIO_INTRA_OP_PARALLELISM, AIO_INTRA_OP_PARALLELISM_DEFAULT), - AIO_SINGLE_SUBMIT: - get_scalar_param(aio_dict, AIO_SINGLE_SUBMIT, AIO_SINGLE_SUBMIT_DEFAULT), - AIO_OVERLAP_EVENTS: - get_scalar_param(aio_dict, AIO_OVERLAP_EVENTS, AIO_OVERLAP_EVENTS_DEFAULT), - AIO_USE_GDS: - get_scalar_param(aio_dict, AIO_USE_GDS, AIO_USE_GDS_DEFAULT) - } - - if aio_config[AIO_USE_GDS]: - assert get_accelerator().device_name() == 'cuda', 'GDS currently only supported for CUDA accelerator' - - return aio_config - - return AIO_DEFAULT_DICT diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/async_swapper.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/async_swapper.py deleted file mode 100644 index b808721537fef5e3907bd3b8819b2e2d0b2f6936..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/async_swapper.py +++ /dev/null @@ -1,175 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping tensors to/from (NVMe) storage devices. -""" -import torch - -from deepspeed import comm as dist -from deepspeed.utils.logging import logger -from deepspeed.runtime.swap_tensor.utils import swap_out_tensors, SwapBuffer -from deepspeed.accelerator import get_accelerator - -INVALID_BUFFER_INDEX = -1 -ASYNC_SWAPPER_WAIT_TIMER = 'async_swap_gradient_wait' - - -class AsyncTensorSwapper(object): - - def __init__(self, aio_handle, numel_alignment, timers): - self.free_buffer_index = [] - self.swapping_buffer_index = [] - self.ready_buffer_index = [] - self.current_buffer_index = INVALID_BUFFER_INDEX - self.all_buffers = [] - self.aio_handle = aio_handle - self.numel_alignment = numel_alignment - self.max_numel = 0 - self.num_pending_swaps = 0 - self.timers = timers - self.timer_names = set() - self.num_elements_swapped = 0 - self.dtype = None - - def has_buffers(self): - return len(self.all_buffers) > 0 - - def add_buffers(self, buffer_list): - assert len(self.all_buffers) == 0 - assert all([get_accelerator().is_pinned(buffer) for buffer in buffer_list]) - dtype = buffer_list[0].dtype - assert all([buffer.dtype == dtype for buffer in buffer_list]) - - self.dtype = dtype - self.all_buffers = [SwapBuffer(buffer) for buffer in buffer_list] - self.free_buffer_index += [i for i in range(len(self.all_buffers))] - self.max_numel = max([buffer.numel() for buffer in buffer_list]) - self.timer_names = set() - - def get_timer_names(self): - return list(self.timer_names) - - def release_buffers(self): - self._report_statistics('Swapped out[Before flush]') - self._flush_buffers_until_complete() - self._report_statistics('Swapped out[After flush]') - - pinned_buffers = [buf.buffer for buf in self.all_buffers] - self.all_buffers = [] - self.free_buffer_index = [] - self.current_buffer_index = INVALID_BUFFER_INDEX - self.num_elements_swapped = 0 - self.dtype = None - - return pinned_buffers - - def swap_out_tensors(self, tensor_list, path_list): - for tensor, swap_path in zip(tensor_list, path_list): - self._swap_out_tensor(tensor, swap_path) - - def _report_statistics(self, message): - if dist.get_rank() == 0: - element_size = torch.tensor([], dtype=self.dtype).element_size() - swapped_GB = (self.num_elements_swapped * element_size) / (1024**3) - logger.debug(f'{message} num_elems = {self.num_elements_swapped}, {swapped_GB:5.2f} GB') - - def _swap_out_tensor(self, tensor, swap_path): - assert len(self.all_buffers) > 0 - - aligned_numel = self._io_aligned_numel(tensor.numel()) - assert aligned_numel <= self.max_numel - - self._make_swap_space(aligned_numel) - assert self.current_buffer_index != INVALID_BUFFER_INDEX - - swap_buffer = self._get_current_buffer() - swap_buffer.insert_tensor(tensor, swap_path, aligned_numel) - - def _make_swap_space(self, numel): - if self.current_buffer_index == INVALID_BUFFER_INDEX: - self._allocate_buffer() - return - - if not self._get_current_buffer().has_space(numel): - if len(self.free_buffer_index) > 0: - self._flush_ready_buffers() - else: - self._flush_buffers_until_complete() - self._allocate_buffer() - - def _io_aligned_numel(self, numel): - remainder = numel % self.numel_alignment - return numel if remainder == 0 else (numel + self.numel_alignment - remainder) - - def _allocate_buffer(self): - assert self.current_buffer_index == INVALID_BUFFER_INDEX - assert len(self.all_buffers) > 0 - assert len(self.free_buffer_index) > 0 - self.current_buffer_index = self.free_buffer_index[-1] - self.free_buffer_index = self.free_buffer_index[:-1] - - def _flush_ready_buffers(self): - if self.current_buffer_index != INVALID_BUFFER_INDEX: - self.ready_buffer_index.append(self.current_buffer_index) - self.current_buffer_index = INVALID_BUFFER_INDEX - - self._swap_out_ready_buffers() - - def _flush_buffers_until_complete(self): - self._flush_ready_buffers() - assert len(self.ready_buffer_index) == 0 - - self._wait_for_swap_complete() - assert len(self.swapping_buffer_index) == 0 - assert len(self.free_buffer_index) == len(self.all_buffers) - - def _swap_out_ready_buffers(self): - for buffer_index in self.ready_buffer_index: - buffer = self._get_buffer(buffer_index) - swap_tensors = buffer.get_swap_tensors() - swap_paths = buffer.get_swap_paths() - self.num_pending_swaps += len(swap_tensors) - swap_out_tensors(self.aio_handle, swap_tensors, swap_paths) - - self.swapping_buffer_index += self.ready_buffer_index - self.ready_buffer_index = [] - - def _wait_for_swap_complete(self): - assert len(self.swapping_buffer_index) > 0 - - self._start_timer(ASYNC_SWAPPER_WAIT_TIMER) - assert self.aio_handle.wait() == self.num_pending_swaps - self._stop_timer(ASYNC_SWAPPER_WAIT_TIMER) - self.timer_names.add(ASYNC_SWAPPER_WAIT_TIMER) - - self.num_pending_swaps = 0 - - for buffer_index in self.swapping_buffer_index: - buffer = self._get_buffer(buffer_index) - self.num_elements_swapped += buffer.get_num_elem() - buffer.reset() - - self.free_buffer_index += self.swapping_buffer_index - assert len(self.free_buffer_index) <= len(self.all_buffers) - self.swapping_buffer_index = [] - - def _get_buffer(self, index): - assert index != INVALID_BUFFER_INDEX - return self.all_buffers[index] - - def _get_current_buffer(self): - return self._get_buffer(self.current_buffer_index) - - def _start_timer(self, name): - if self.timers: - self.timers(name).start() - - def _stop_timer(self, name): - if self.timers: - self.timers(name).stop() - - def _log_timers(self, name_list, force=False): - if self.timers and force: - self.timers.log(name_list) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/constants.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/constants.py deleted file mode 100644 index c1207749eac6915b37c9d3d01f93f2e06e99250e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/constants.py +++ /dev/null @@ -1,30 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -AIO -""" -AIO_FORMAT = ''' -"aio": { - "block_size": 1048576, - "queue_depth": 8, - "intra_op_parallelism": 1, - "single_submit": false, - "overlap_events": true, - "use_gds": false -} -''' -AIO = "aio" -AIO_BLOCK_SIZE = "block_size" -AIO_BLOCK_SIZE_DEFAULT = 1048576 -AIO_QUEUE_DEPTH = "queue_depth" -AIO_QUEUE_DEPTH_DEFAULT = 8 -AIO_INTRA_OP_PARALLELISM = "intra_op_parallelism" -AIO_INTRA_OP_PARALLELISM_DEFAULT = 1 -AIO_SINGLE_SUBMIT = "single_submit" -AIO_SINGLE_SUBMIT_DEFAULT = False -AIO_OVERLAP_EVENTS = "overlap_events" -AIO_OVERLAP_EVENTS_DEFAULT = True -AIO_USE_GDS = "use_gds" -AIO_USE_GDS_DEFAULT = False diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/optimizer_utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/optimizer_utils.py deleted file mode 100644 index 6729fd28b8fc8acd33dd552d6e70d86b326d2ac2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/optimizer_utils.py +++ /dev/null @@ -1,527 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping tensors to/from (NVMe) storage devices. -""" - -import os -import torch - -from deepspeed import comm as dist -from deepspeed.utils.logging import logger -from deepspeed.runtime.swap_tensor.constants import * -from deepspeed.runtime.swap_tensor.utils import swap_in_tensors, swap_out_tensors, \ - MIN_AIO_BYTES, AIO_ALIGNED_BYTES, get_sized_buffers -from deepspeed.runtime.swap_tensor.utils import SwapBufferManager, SwapBufferPool -from deepspeed.accelerator import get_accelerator - - -class FlattenedTensorSwapInfo(object): - - def __init__(self, path, length, offset): - self.path = path - self.offset = offset - self.length = length - - -class SwapTensorContext(object): - - def __init__(self, tensor, swap_folder): - self.compute_tensor = tensor - self.swap_tensor = torch.Tensor() - self.swap_path = os.path.join(swap_folder, f'{OptimizerSwapper.parameter_id(tensor)}.tensor.swp') - - def release_memory(self): - self.compute_tensor.data = torch.Tensor() - self.swap_tensor.data = torch.Tensor() - - def set_buffers(self, compute_buffer, swap_buffer): - self.compute_tensor.data = compute_buffer.data - self.swap_tensor.data = swap_buffer.data - - -class OptimizerStateSwapInfo(object): - - def __init__(self, parameter, numel, base_folder): - self.tensors = [] - self.param_id = OptimizerSwapper.parameter_id(parameter) - self.swap_folder = base_folder - self.swapped_gradients = {} - self.unswapped_gradients = {} - self.tensor_numel = numel - self.tensor_dtype = parameter.dtype - self.tensor_device = parameter.device - self.has_state_tensors = False - self.swap_buffers = [] - self._add_tensors([parameter]) - - def numel(self): - return self.tensor_numel - - def has_gradients(self): - return bool(self.swapped_gradients) or bool(self.unswapped_gradients) - - def _add_tensors(self, tensor_list): - for t in tensor_list: - self.tensors.append(SwapTensorContext(t, self.swap_folder)) - - def add_state_tensors(self, tensor_list): - self.has_state_tensors = True - self._add_tensors(tensor_list) - - def num_tensors(self): - return len(self.tensors) - - def device(self): - return self.tensor_device - - def dtype(self): - return self.tensor_dtype - - def release_memory(self): - for t in self.tensors: - t.release_memory() - - def get_compute_tensors(self): - return [t.compute_tensor for t in self.tensors] - - def get_swap_paths(self): - return [t.swap_path for t in self.tensors] - - def get_swap_buffers_and_paths(self, pinned): - swap_buffers = [] - swap_paths = [] - select_tensors = [t for t in self.tensors if get_accelerator().is_pinned(t.compute_tensor) == pinned] - for t in select_tensors: - swap_buffers.append(t.swap_tensor if pinned else t.compute_tensor) - swap_paths.append(t.swap_path) - return swap_buffers, swap_paths - - def get_or_create_gradient_paths(self, offsets, lengths): - gradient_paths = [] - for offset, length in zip(offsets, lengths): - if not offset in self.swapped_gradients.keys(): - path = os.path.join(self.swap_folder, f'{self.param_id}_gradient_{offset}_{length}.tensor.swp') - self.swapped_gradients[offset] = FlattenedTensorSwapInfo(path, length, offset) - - gradient_paths.append(self.swapped_gradients[offset].path) - - return gradient_paths - - def set_swap_buffers(self, buffers, aligned_numel): - num_tensors = len(self.tensors) - compute_lengths = [self.numel()] * num_tensors - compute_buffers = get_sized_buffers(buffers, compute_lengths) - swap_lengths = [aligned_numel] * num_tensors - swap_buffers = get_sized_buffers(buffers, swap_lengths) - - for i, t in enumerate(self.tensors): - t.set_buffers(compute_buffer=compute_buffers[i], swap_buffer=swap_buffers[i]) - - def get_swap_gradient_buffers(self, swap_buffer): - assert self.numel() <= swap_buffer.numel() - return [swap_buffer.narrow(0, grad.offset, grad.length) for grad in self.swapped_gradients.values()] - - def get_swap_gradient_paths(self): - return [grad.path for grad in self.swapped_gradients.values()] - - def get_unpinned_state_tensors(self): - return [t.compute_tensor for t in self.tensors if not get_accelerator().is_pinned(t.compute_tensor)] - - def read_unswapped_gradients(self, dest_buffer): - num_elem_count = 0 - for offset, grad_partition in self.unswapped_gradients.items(): - dst_tensor = dest_buffer.narrow(0, offset, grad_partition.numel()) - dst_tensor.data.copy_(grad_partition.data) - num_elem_count += grad_partition.numel() - - return num_elem_count - - def write_unswapped_gradients(self, src_buffer): - num_elem_count = 0 - for offset, grad_partition in self.unswapped_gradients.items(): - src_tensor = src_buffer.narrow(0, offset, grad_partition.numel()) - grad_partition.data.copy_(src_tensor.data) - num_elem_count += grad_partition.numel() - - return num_elem_count - - def release_unswapped_gradients(self): - self.unswapped_gradients = {} - - -SWAPPER_DEBUG_MODE = False -SWAP_OUT_GRADIENT_TIMER = 'swap_out_gradient' - - -class OptimizerSwapper(object): - - @staticmethod - def parameter_id(param): - return param.ds_id - - def __init__(self, swap_config, aio_config, base_folder, optimizer, largest_numel, device, dtype, timers): - self.swap_config = swap_config - self.aio_config = aio_config - - # NVMe swap management - self.swap_params_info = {} - self.swap_element_size = torch.tensor([], dtype=dtype).element_size() - self.swap_folder = os.path.join(base_folder, 'optimizer', f'rank{dist.get_rank()}') - os.makedirs(self.swap_folder, exist_ok=True) - - self.optimizer = optimizer - - # Read/Write alignment for each thread during Intra-request parallelism - self.min_aio_bytes = max(MIN_AIO_BYTES, aio_config[AIO_BLOCK_SIZE]) - self.aligned_bytes = AIO_ALIGNED_BYTES * aio_config[AIO_INTRA_OP_PARALLELISM] - self.numel_alignment = self.aligned_bytes // self.swap_element_size - - # Swap buffer management - self.largest_numel = self._io_aligned_numel(largest_numel) - self.dtype = dtype - self.swap_buffer_manager = SwapBufferManager(num_elems=self.largest_numel, - count=swap_config.buffer_count, - dtype=dtype) - - # Timers - self.timers = timers - self.timer_names = set() - - # Print exclusion list - self.print_exclude_list = [ - 'optimizer', - 'swap_buffer_manager', - 'swap_params_info', - 'timers', - 'timer_names', - ] - - def purge_state(self): - for swap_info in self.swap_params_info.values(): - swap_info.tensors = [swap_info.tensors[0]] - swap_info.has_state_tensors = False - - def is_swappable_tensor(self, tensor=None, numel=None): - assert tensor is not None or numel is not None, "Either tensor or numel must be provided" - if tensor is not None: - return self.min_aio_bytes <= (tensor.numel() * self.swap_element_size) - return self.min_aio_bytes <= (numel * self.swap_element_size) - - def init_timers(self): - self.timer_names = set() - - def log_timers(self): - if self.timer_names: - self._log_timers(list(self.timer_names), force=True) - - def pre_backward(self): - self.init_timers() - - def post_backward(self): - pass - - def _flush_gradient_swapper(self, gradient_swapper): - if gradient_swapper.has_buffers(): - self._start_timer(SWAP_OUT_GRADIENT_TIMER) - pinned_buffers = gradient_swapper.release_buffers() - self.swap_buffer_manager.free(pinned_buffers) - self._stop_timer(SWAP_OUT_GRADIENT_TIMER) - self.timer_names.add(SWAP_OUT_GRADIENT_TIMER) - self.timer_names.update(gradient_swapper.get_timer_names()) - - def _swap_out_gradients(self, parameter, gradient_offsets, gradient_tensors, gradient_swapper): - if not OptimizerSwapper.parameter_id(parameter) in self.swap_params_info.keys(): - return - - swap_info = self.swap_params_info[OptimizerSwapper.parameter_id(parameter)] - - swappable_tensors = [] - swappable_offsets = [] - swappable_lengths = [] - - aligned_gradients, aligned_offsets = self._adjust_for_misaligned_lengths(tensors=gradient_tensors, - offsets=gradient_offsets) - - self._start_timer(SWAP_OUT_GRADIENT_TIMER) - for tensor, offset in zip(aligned_gradients, aligned_offsets): - if not self.is_swappable_tensor(tensor=tensor): - swap_info.unswapped_gradients[offset] = tensor - continue - - swappable_tensors.append(tensor) - swappable_offsets.append(offset) - swappable_lengths.append(tensor.numel()) - - if len(swappable_tensors) > 0: - if not gradient_swapper.has_buffers(): - pinned_buffers = self.swap_buffer_manager.allocate_all(num_elems=self.largest_numel, dtype=self.dtype) - - gradient_swapper.add_buffers(pinned_buffers) - - swappable_paths = swap_info.get_or_create_gradient_paths(swappable_offsets, swappable_lengths) - - gradient_swapper.swap_out_tensors(tensor_list=swappable_tensors, path_list=swappable_paths) - - self._stop_timer(SWAP_OUT_GRADIENT_TIMER) - self.timer_names.add(SWAP_OUT_GRADIENT_TIMER) - - def _initialize_from_swapped_fp16_params(self, aio_handle, fp16_partitions_info, fp16_num_elems, - fp16_pinned_buffers, fp32_parameters): - assert len(fp32_parameters) == len(fp16_partitions_info) - assert len(fp32_parameters) == len(fp16_num_elems) - assert all([get_accelerator().is_pinned(buffer) for buffer in fp16_pinned_buffers]) - - fp32_swap_paths = self._get_swap_paths(parameters=fp32_parameters, num_elems=fp16_num_elems) - - fp32_pinned_buffers = self.swap_buffer_manager.allocate_all(num_elems=self.largest_numel, dtype=self.dtype) - - fp16_buffer_numel = [buf.numel() for buf in fp16_pinned_buffers] - assert all([numel >= self.largest_numel for numel in fp16_buffer_numel]), \ - f"numel of fp16 buffers {fp16_buffer_numel} is too small for initializing fp32 params {self.largest_numel}" - - fp32_swap_buffers = SwapBufferPool(fp32_pinned_buffers) - fp16_swap_buffers = SwapBufferPool(fp16_pinned_buffers) - - curr_index = 0 - while curr_index < len(fp32_parameters): - fp16_pinned_tensors = self._swap_in_fp16_params(aio_handle=aio_handle, - fp16_num_elems=fp16_num_elems[curr_index:], - fp16_partitions_info=fp16_partitions_info[curr_index:], - fp16_swap_buffers=fp16_swap_buffers) - - if dist.get_rank() == 0 and SWAPPER_DEBUG_MODE: - for i, tensor in enumerate(fp16_pinned_tensors): - true_index = curr_index + i - logger.info( - f'swap_in_fp16_param: fp32_id = {OptimizerSwapper.parameter_id(fp32_parameters[true_index])} index = {true_index} orig_num_elem = {fp16_num_elems[true_index]}, swap_num_elem = {fp16_pinned_tensors[i].numel()}' - ) - - swap_out_count = self._swap_out_fp16_params(aio_handle=aio_handle, - fp32_swap_paths=fp32_swap_paths[curr_index:], - fp32_swap_buffers=fp32_swap_buffers, - fp16_pinned_tensors=fp16_pinned_tensors) - assert swap_out_count == len(fp16_pinned_tensors), \ - f"{swap_out_count} does not match {len(fp16_pinned_tensors)}" - - fp16_swap_buffers.reset() - fp32_swap_buffers.reset() - curr_index += swap_out_count - - self.swap_buffer_manager.free(fp32_pinned_buffers) - - def _swap_in_fp16_params(self, aio_handle, fp16_num_elems, fp16_partitions_info, fp16_swap_buffers): - assert len(fp16_num_elems) > 0 - - swapped_fp16_tensors = [] - swap_tensors = [] - swap_paths = [] - unswapped_srcs = [] - unswapped_dsts = [] - - for i, numel in enumerate(fp16_num_elems): - pinned_tensor, _ = fp16_swap_buffers.allocate_tensor(numel, None, numel) - if pinned_tensor is None: - break - - swapped_fp16_tensors.append(pinned_tensor) - offset = 0 - for tensor, partition_numel, partition_path in fp16_partitions_info[i]: - dst_tensor = pinned_tensor.narrow(0, offset, partition_numel) - if partition_path is None: - unswapped_srcs.append(tensor) - unswapped_dsts.append(dst_tensor) - else: - swap_paths.append(partition_path) - swap_tensors.append(dst_tensor) - offset += partition_numel - - assert len(swapped_fp16_tensors) + len(unswapped_srcs) > 0 - ret = swap_in_tensors(aio_handle, swap_tensors, swap_paths) - for src, dst in zip(unswapped_srcs, unswapped_dsts): - dst.data.copy_(src.data) - - assert len(swap_tensors) == aio_handle.wait() - - return swapped_fp16_tensors - - def _swap_out_fp16_params(self, aio_handle, fp32_swap_paths, fp32_swap_buffers, fp16_pinned_tensors): - - assert len(fp16_pinned_tensors) <= len(fp32_swap_paths) - swap_out_count = 0 - for i, fp16_tensor in enumerate(fp16_pinned_tensors): - if not fp32_swap_buffers.has_space(fp16_tensor.numel()): - fp32_swap_buffers.swap_out(aio_handle) - fp32_swap_buffers.reset() - - pinned_tensor, _ = fp32_swap_buffers.insert_tensor(fp16_tensor, fp32_swap_paths[i], - self._io_aligned_numel(fp16_tensor.numel())) - assert pinned_tensor is not None - swap_out_count += 1 - - if len(fp32_swap_buffers.get_swap_tensors()) > 0: - fp32_swap_buffers.swap_out(aio_handle) - - return swap_out_count - - def _initialize_parameters(self, parameters, src_tensors, aio_handle): - assert len(parameters) == len(src_tensors) - - swap_paths = self._get_swap_paths(parameters=parameters, num_elems=[src.numel() for src in src_tensors]) - - SWAP_INIT_TIMER = "swap_init_write" - self._start_timer(SWAP_INIT_TIMER) - - pinned_buffers = self.swap_buffer_manager.allocate_all(num_elems=self.largest_numel, dtype=self.dtype) - assert pinned_buffers is not None - - self._swap_out_unpinned_tensors(aio_handle=aio_handle, - unpinned_tensors=src_tensors, - dest_paths=swap_paths, - pinned_buffers=pinned_buffers) - - if dist.get_rank() == 0 and SWAPPER_DEBUG_MODE: - for i, tensor in enumerate(src_tensors): - logger.info( - f'copy_in_fp16_param: fp32_id = {OptimizerSwapper.parameter_id(parameters[i])} index = {i}, swap_num_elem = {src_tensors[i].numel()}' - ) - - self.swap_buffer_manager.free(pinned_buffers) - - self._stop_timer(SWAP_INIT_TIMER) - self._log_timers([SWAP_INIT_TIMER]) - - def _get_swap_paths(self, parameters, num_elems): - swap_info_list = [ - self._create_param_swap_info(parameter=p, - numel=numel) \ - for p, numel in zip(parameters, num_elems) - ] - assert len(swap_info_list) == len(num_elems) - - swap_paths = [info.tensors[0].swap_path for info in swap_info_list] - return swap_paths - - def _swap_out_unpinned_tensors(self, aio_handle, unpinned_tensors, dest_paths, pinned_buffers): - - swap_buffer_count = len(pinned_buffers) - unpinned_tensor_count = len(unpinned_tensors) - - for i in range(0, unpinned_tensor_count, swap_buffer_count): - swap_tensor_count = min((unpinned_tensor_count - i), swap_buffer_count) - - src_tensors = unpinned_tensors[i:(i + swap_tensor_count)] - compute_lengths = [t.numel() for t in src_tensors] - compute_buffers = get_sized_buffers(pinned_buffers, compute_lengths) - - for dst, src in zip(compute_buffers, src_tensors): - dst.data.copy_(src.data) - - swap_lengths = [self._io_aligned_numel(t.numel()) for t in src_tensors] - swap_buffers = get_sized_buffers(pinned_buffers, swap_lengths) - - swap_paths = dest_paths[i:(i + swap_tensor_count)] - swap_out_tensors(aio_handle, swap_buffers, swap_paths) - - assert aio_handle.wait() == swap_tensor_count - - def _adjust_for_misaligned_lengths(self, tensors, offsets): - new_tensors = [] - new_offsets = [] - - for orig_tensor, orig_offset in zip(tensors, offsets): - if not self.is_swappable_tensor(tensor=orig_tensor): - new_tensors.append(orig_tensor) - new_offsets.append(orig_offset) - continue - - remainder = orig_tensor.numel() % self.numel_alignment - if remainder == 0: - new_tensors.append(orig_tensor) - new_offsets.append(orig_offset) - continue - - # Split into two by making remainder a tensor - aligned_length = (orig_tensor.numel() // self.numel_alignment) * self.numel_alignment - new_tensors.append(orig_tensor.narrow(0, 0, aligned_length)) - new_offsets.append(orig_offset) - - # remainder tensor - new_tensors.append(orig_tensor.narrow(0, aligned_length, remainder)) - new_offsets.append(orig_offset + aligned_length) - - return new_tensors, new_offsets - - def _retrieve_unswapped_grad_partitions(self, swap_info, dest_buffer): - UNSWAPPED_READ_GRADIENTS = 'unswapped_read_gradients' - self._start_timer(UNSWAPPED_READ_GRADIENTS) - tensor_count = len(swap_info.unswapped_gradients) - num_elem_count = swap_info.read_unswapped_gradients(dest_buffer) - self._stop_timer(UNSWAPPED_READ_GRADIENTS) - self._log_timers([UNSWAPPED_READ_GRADIENTS]) - - # It should be safe to discard unswapped gradient partitions - swap_info.release_unswapped_gradients() - - if SWAPPER_DEBUG_MODE: - logger.info( - f'optimizer_retrieve_unswapped_gradients: param={swap_info.param_id} tensor_count={tensor_count} elem_count={num_elem_count}' - ) - - def _get_state_tensors(self, parameter): - if not parameter in self.optimizer.state: - return [] - - tensor_list = [] - for state_name, value in self.optimizer.state[parameter].items(): - if torch.is_tensor(value) and self.is_swappable_tensor(tensor=value): - value.ds_id = state_name + '-' + parameter.ds_id - tensor_list.append(value) - - return tensor_list - - def _update_param_state_info(self, swap_info, parameter): - if not swap_info.has_state_tensors: - state_tensors = self._get_state_tensors(parameter) - if state_tensors: - swap_info.add_state_tensors(state_tensors) - - def _create_param_swap_info(self, parameter, numel): - param_id = OptimizerSwapper.parameter_id(parameter) - assert not param_id in self.swap_params_info - - self.swap_params_info[param_id] = OptimizerStateSwapInfo(parameter=parameter, - numel=numel, - base_folder=self.swap_folder) - swap_info = self.swap_params_info[param_id] - - self._update_param_state_info(swap_info, parameter) - - return swap_info - - def _get_param_swap_info(self, parameter): - param_id = OptimizerSwapper.parameter_id(parameter) - swap_info = self.swap_params_info.get(param_id, None) - - if swap_info is not None: - self._update_param_state_info(swap_info, parameter) - - return swap_info - - def _start_timer(self, name): - if self.timers: - self.timers(name).start() - - def _stop_timer(self, name): - if self.timers: - self.timers(name).stop() - - def _log_timers(self, name_list, force=False): - if self.timers and (SWAPPER_DEBUG_MODE or force): - self.timers.log(name_list) - - def _io_aligned_numel(self, numel): - remainder = numel % self.numel_alignment - return numel if remainder == 0 else (numel + self.numel_alignment - remainder) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py deleted file mode 100644 index 52b873ba58a1f7816aae4f79527773eb26870a34..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py +++ /dev/null @@ -1,229 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping optimizer tensors to/from (NVMe) storage devices. -""" - -from deepspeed.utils.logging import logger -from deepspeed.ops.op_builder import AsyncIOBuilder -from deepspeed import comm as dist - -from deepspeed.runtime.swap_tensor.constants import * -from deepspeed.runtime.swap_tensor.utils import swap_in_tensors, swap_out_tensors, print_object, \ - get_sized_buffers -from deepspeed.runtime.swap_tensor.async_swapper import AsyncTensorSwapper -from deepspeed.runtime.swap_tensor.optimizer_utils import OptimizerSwapper -from deepspeed.accelerator import get_accelerator - -DEBUG_MODE = False - -SWAP_IN_PARAM_TIMER = 'swap_in_param' -SWAP_OUT_PARAM_TIMER = 'swap_out_param' -SWAP_IN_GRADIENT_TIMER = 'swap_in_gradient' - - -class PartitionedOptimizerSwapper(OptimizerSwapper): - - def __init__(self, swap_config, aio_config, base_folder, optimizer, largest_numel, device, dtype, timers): - super(PartitionedOptimizerSwapper, self).__init__(swap_config, aio_config, base_folder, optimizer, - largest_numel, device, dtype, timers) - - aio_op = AsyncIOBuilder().load() - self.aio_handle = aio_op.aio_handle(block_size=aio_config[AIO_BLOCK_SIZE], - queue_depth=aio_config[AIO_QUEUE_DEPTH], - single_submit=aio_config[AIO_SINGLE_SUBMIT], - overlap_events=aio_config[AIO_OVERLAP_EVENTS], - intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM]) - - # Overlap swapping out - self.gradient_swapper = AsyncTensorSwapper(aio_handle=self.aio_handle, - numel_alignment=self.numel_alignment, - timers=self.timers) - - self.print_exclude_list += ['aio_handle', 'gradient_swapper', 'print_exclude_list'] - - if dist.get_rank() == 0: - print_object(obj=self, name='PartitionedOptimizerSwapper', exclude_list=self.print_exclude_list) - - def initialize_parameters(self, parameters, src_tensors): - self._initialize_parameters(parameters=parameters, src_tensors=src_tensors, aio_handle=self.aio_handle) - - def initialize_from_swapped_fp16_params(self, fp16_partitions_info, fp16_num_elems, fp16_pinned_buffers, - fp32_parameters): - self._initialize_from_swapped_fp16_params(aio_handle=self.aio_handle, - fp16_partitions_info=fp16_partitions_info, - fp16_num_elems=fp16_num_elems, - fp16_pinned_buffers=fp16_pinned_buffers, - fp32_parameters=fp32_parameters) - - def flush_gradients(self): - self._flush_gradient_swapper(self.gradient_swapper) - - def release_swap_buffers(self, parameter): - swap_info = self._get_param_swap_info(parameter) - if swap_info is None: - return - swap_info.release_memory() - - self.swap_buffer_manager.free(swap_info.swap_buffers) - swap_info.swap_buffers = [] - - def swap_in_optimizer_state(self, parameter, async_parameter=None): - swap_info = self._get_param_swap_info(parameter) - if swap_info is None: - return - - self._flush_gradient_swapper(self.gradient_swapper) - - required_buffer_count = swap_info.num_tensors() + (1 if swap_info.has_gradients() else 0) - aligned_numel = self._io_aligned_numel(swap_info.numel()) - pinned_buffers = self.swap_buffer_manager.allocate(num_elems=aligned_numel, - count=required_buffer_count, - dtype=parameter.dtype) - assert pinned_buffers is not None - swap_info.swap_buffers = pinned_buffers.copy() - - self._start_timer(SWAP_IN_PARAM_TIMER) - self._swap_in_parameter(aio_handle=self.aio_handle, - parameter=parameter, - dest_buffers=pinned_buffers[:swap_info.num_tensors()]) - self._stop_timer(SWAP_IN_PARAM_TIMER) - self.timer_names.add(SWAP_IN_PARAM_TIMER) - - if swap_info.has_gradients(): - self._start_timer(SWAP_IN_GRADIENT_TIMER) - self._swap_in_gradients(aio_handle=self.aio_handle, parameter=parameter, dest_buffer=pinned_buffers[-1]) - self._stop_timer(SWAP_IN_GRADIENT_TIMER) - self.timer_names.add(SWAP_IN_GRADIENT_TIMER) - - def _swap_out_optimizer_state(self, swap_info): - pinned_tensors, pinned_paths = swap_info.get_swap_buffers_and_paths(True) - WRITE_TIMER = 'swap_submit_write' - self._start_timer(WRITE_TIMER) - - swap_out_tensors(self.aio_handle, pinned_tensors, pinned_paths) - assert self.aio_handle.wait() == len(pinned_tensors) - - unpinned_tensors, unpinned_paths = swap_info.get_swap_buffers_and_paths(False) - if len(unpinned_tensors) > 0: - pinned_buffers = self.swap_buffer_manager.allocate_all(num_elems=self.largest_numel, dtype=self.dtype) - self._swap_out_unpinned_tensors(aio_handle=self.aio_handle, - unpinned_tensors=unpinned_tensors, - dest_paths=unpinned_paths, - pinned_buffers=pinned_buffers) - swap_info.swap_buffers += pinned_buffers.copy() - - self._stop_timer(WRITE_TIMER) - self._log_timers([WRITE_TIMER]) - - def writeback_optimizer_state_and_gradients(self, parameter, write_opt_state, write_gradients): - swap_info = self._get_param_swap_info(parameter=parameter) - - if swap_info is None: - return - - if write_opt_state: - self._swap_out_optimizer_state(swap_info) - - if write_gradients and swap_info.has_gradients(): - param_gradients = swap_info.swapped_gradients.values() - swap_buffers = [parameter.grad.narrow(0, grad.offset, grad.length) for grad in param_gradients] - swap_paths = [grad.path for grad in param_gradients] - swap_out_tensors(self.aio_handle, swap_buffers, swap_paths) - assert len(swap_buffers) == self.aio_handle.wait() - if swap_info.unswapped_gradients: - swap_info.write_unswapped_gradients(src_buffer=parameter.grad) - - self.release_swap_buffers(parameter) - - def swap_out_optimizer_state(self, parameter, async_swap=False): - swap_info = self._get_param_swap_info(parameter=parameter) - - if swap_info is None: - return - - swap_bytes = sum( - [self._io_aligned_numel(t.numel()) * t.element_size() for t in swap_info.get_compute_tensors()]) - - self._start_timer(SWAP_OUT_PARAM_TIMER) - self._swap_out_optimizer_state(swap_info) - self.release_swap_buffers(parameter) - self._stop_timer(SWAP_OUT_PARAM_TIMER) - self.timer_names.add(SWAP_OUT_PARAM_TIMER) - - if DEBUG_MODE and dist.get_rank() == 0: - logger.info(f'optimizer_param_swap_out: {(swap_bytes/(1024**3)):5.2f} GB') - - def swap_out_gradients(self, parameter, gradient_offsets, gradient_tensors): - self._swap_out_gradients(parameter=parameter, - gradient_offsets=gradient_offsets, - gradient_tensors=gradient_tensors, - gradient_swapper=self.gradient_swapper) - - def _swap_in_parameter(self, aio_handle, parameter, dest_buffers): - swap_info = self._get_param_swap_info(parameter) - if swap_info is None: - return - - num_swap_tensors = swap_info.num_tensors() - assert num_swap_tensors <= len(dest_buffers) - - swap_lengths = [self._io_aligned_numel(swap_info.numel())] * num_swap_tensors - swap_buffers = get_sized_buffers(dest_buffers, swap_lengths) - - compute_lengths = [swap_info.numel()] * num_swap_tensors - compute_buffers = get_sized_buffers(dest_buffers, compute_lengths) - - READ_TIMER = 'swap_submit_read_param' - WAIT_TIMER = 'swap_wait_read_param' - - self._start_timer(READ_TIMER) - swap_in_tensors(aio_handle, swap_buffers, swap_info.get_swap_paths()) - self._stop_timer(READ_TIMER) - - swap_bytes = sum([buffer.numel() * buffer.element_size() for buffer in swap_buffers]) - - self._start_timer(WAIT_TIMER) - aio_handle.wait() - self._stop_timer(WAIT_TIMER) - - swap_info.set_swap_buffers(dest_buffers, self._io_aligned_numel(swap_info.numel())) - - self._log_timers([READ_TIMER, WAIT_TIMER]) - if DEBUG_MODE and dist.get_rank() == 0: - logger.info(f'optimizer_param_swap_in: {(swap_bytes/(1024**3)):5.2f} GB') - - def _swap_in_pinned_gradients(self, aio_handle, parameter, gradient_tensor): - swap_info = self.swap_params_info[OptimizerSwapper.parameter_id(parameter)] - param_gradients = swap_info.swapped_gradients.values() - swap_buffers = [gradient_tensor.narrow(0, grad.offset, grad.length) for grad in param_gradients] - swap_paths = [grad.path for grad in param_gradients] - SWAP_READ_GRADIENTS = 'swap_submit_read_gradient' - SWAP_WAIT_GRADIENTS = 'swap_submit_wait_gradient' - self._start_timer(SWAP_READ_GRADIENTS) - swap_in_tensors(aio_handle, swap_buffers, swap_paths) - self._stop_timer(SWAP_READ_GRADIENTS) - - self._start_timer(SWAP_WAIT_GRADIENTS) - assert len(swap_buffers) == aio_handle.wait() - self._stop_timer(SWAP_WAIT_GRADIENTS) - - self._log_timers([SWAP_READ_GRADIENTS, SWAP_WAIT_GRADIENTS]) - - def _swap_in_gradients(self, aio_handle, parameter, dest_buffer): - swap_info = self.swap_params_info.get(OptimizerSwapper.parameter_id(parameter), None) - if not (swap_info and swap_info.has_gradients()): - return - - assert get_accelerator().is_pinned(dest_buffer) - assert parameter.numel() <= dest_buffer.numel() - - parameter.grad = dest_buffer.narrow(0, 0, parameter.numel()) - - if swap_info.swapped_gradients: - self._swap_in_pinned_gradients(aio_handle, parameter, parameter.grad) - - if swap_info.unswapped_gradients: - self._retrieve_unswapped_grad_partitions(swap_info=swap_info, dest_buffer=parameter.grad) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py deleted file mode 100644 index 3e00a1860b0b20522e90b9c53437207c2052c92a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py +++ /dev/null @@ -1,419 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping tensors to/from (NVMe) storage devices. -""" - -import os -import shutil -from enum import Enum -import torch -from deepspeed import comm as dist -from deepspeed.accelerator import get_accelerator -from deepspeed.ops.op_builder import AsyncIOBuilder -from deepspeed.ops.op_builder import GDSBuilder -from .constants import * -from .utils import swap_in_tensors, swap_out_tensors, MIN_AIO_BYTES, AIO_ALIGNED_BYTES, print_object, SwapBufferPool - - -def print_rank_0(message, debug=False, force=False): - if dist.get_rank() == 0 and (debug or force): - print(message) - - -class PartitionedParamStatus(Enum): - # Partitioned parameters are present and ready for use - AVAILABLE = 1 - - # partitioned params are in some non-memory device - NOT_AVAILABLE = 2 - - # partitioned params are being read from some non-memory device. - INFLIGHT = 3 - - -class AsyncPartitionedParameterSwapper(object): - - def __init__(self, ds_config, model_dtype): - - self.dtype = model_dtype - - #set swap buffers, create aio handles - self._configure_aio(ds_config) - - #mapping from param id to path - self.id_to_path = {} - - #mapping from pram_id to buffer id - self.param_id_to_buffer_id = {} - - # mapping from param_id to swap buffer - self.param_id_to_swap_buffer = {} - - #number of elements in the param - self.param_id_to_numel = {} - - self.pending_writes = 0 - self.pending_reads = 0 - - #keep track of async swap in params and buffers - self.inflight_params = [] - self.inflight_swap_in_buffers = [] - self.inflight_numel = 0 - - #keep track of available params - self.available_params = set() - self.available_numel = 0 - - # for swapping out from partitioned fp32 params - self.partitioned_swap_buffer = None - self.partitioned_swap_pool = None - - self.invalid_buffer = torch.tensor(1).half() - - if dist.get_rank() == 0: - exclude_list = ['aio_read_handle', 'aio_write_handle', 'buffers'] - print_object(obj=self, name='AsyncPartitionedParameterSwapper', exclude_list=exclude_list) - - def available_swap_in_buffers(self): - return len(self.available_buffer_ids) - - def _configure_aio(self, ds_config): - self.swap_config = ds_config.zero_config.offload_param - torch_dtype_string = str(self.dtype).split(".")[1] - self.swap_folder = os.path.join(self.swap_config.nvme_path, 'zero_stage_3', f'{torch_dtype_string}params', - f'rank{dist.get_rank()}') - shutil.rmtree(self.swap_folder, ignore_errors=True) - os.makedirs(self.swap_folder, exist_ok=True) - - self.swap_element_size = torch.tensor([], dtype=self.dtype).element_size() - - self.aio_config = ds_config.aio_config - - self.use_gds = self.aio_config[AIO_USE_GDS] - self.aio_handle = GDSBuilder().load(verbose=False).gds_handle if self.use_gds else AsyncIOBuilder().load( - verbose=False).aio_handle - - # Read/Write alignment for each thread during Intra-request parallelism - self.min_aio_bytes = max(MIN_AIO_BYTES, self.aio_config[AIO_BLOCK_SIZE]) - self.aligned_bytes = AIO_ALIGNED_BYTES * self.aio_config[AIO_INTRA_OP_PARALLELISM] - self.numel_alignment = self.aligned_bytes // self.swap_element_size - - self.elements_per_buffer = self.swap_config.buffer_size - self.aligned_elements_per_buffer = self._io_aligned_numel(self.elements_per_buffer) - self.param_buffer_count = self.swap_config.buffer_count - - self.available_buffer_ids = [i for i in range(self.param_buffer_count)] - self.reserved_buffer_ids = [] - - self.aio_read_handle = self.aio_handle(block_size=self.aio_config[AIO_BLOCK_SIZE], - queue_depth=self.aio_config[AIO_QUEUE_DEPTH], - single_submit=self.aio_config[AIO_SINGLE_SUBMIT], - overlap_events=self.aio_config[AIO_OVERLAP_EVENTS], - intra_op_parallelism=self.aio_config[AIO_INTRA_OP_PARALLELISM]) - - self.aio_write_handle = self.aio_handle(block_size=self.aio_config[AIO_BLOCK_SIZE], - queue_depth=self.aio_config[AIO_QUEUE_DEPTH], - single_submit=self.aio_config[AIO_SINGLE_SUBMIT], - overlap_events=self.aio_config[AIO_OVERLAP_EVENTS], - intra_op_parallelism=self.aio_config[AIO_INTRA_OP_PARALLELISM]) - - buffer_device = get_accelerator().device_name() if self.use_gds else "cpu" - self.buffers = torch.empty(int(self.aligned_elements_per_buffer * self.param_buffer_count), - dtype=self.dtype, - device=buffer_device, - requires_grad=False) - if self.use_gds: - self.aio_read_handle.pin_device_tensor(self.buffers) - else: - self.buffers = get_accelerator().pin_memory(self.buffers, align_bytes=0) - - self.swap_out_params = [] - - #Check if partitioned param or numel in a tensor is swappable or not - def swappable_tensor(self, param=None, numel=None): - if param is not None: - assert numel is None, "Both parma and numel cannot be provided" - numel = param.ds_tensor.ds_numel - if numel is not None: - return self.min_aio_bytes <= numel * self.swap_element_size - assert False, "Either param or numel must be provided" - - def get_path(self, param, must_exist=False): - paths = self._get_swap_paths([param], must_exist=must_exist) - return paths[0] - - def _get_swap_paths(self, params, must_exist=False): - paths = [] - for param in params: - param_id = param.ds_id - if param_id in self.id_to_path.keys(): - param_path = self.id_to_path[param_id] - else: - assert not must_exist, f"Path for param id {param_id} does not exist" - param_path = os.path.join(self.swap_folder, f'{param_id}_param.tensor.swp') - - self.id_to_path[param_id] = param_path - paths.append(param_path) - - return paths - - def _get_swap_buffers(self, params): - buffers = [] - for param in params: - param_id = param.ds_id - assert param_id in self.param_id_to_swap_buffer.keys(), \ - f'param {param_id} has not been assigned a swap buffer' - buffers.append(self.param_id_to_swap_buffer[param_id]) - - return buffers - - def _track_numel(self, params): - for param in params: - assert param.ds_tensor is not None, "Partitioned tensor is None" - self.param_id_to_numel[param.ds_id] = param.ds_tensor.ds_numel - - def _allocate_and_return_buffers_for_swap_in(self, params): - compute_buffers = [] - swap_buffers = [] - - for param in params: - param_id = param.ds_id - assert param_id in self.param_id_to_numel.keys(), f" Number of elements in param {param_id} is unknown" - assert param_id not in self.param_id_to_buffer_id.keys( - ), f"param {param_id} already assigned swap buffer id {self.param_id_to_buffer_id[param_id]}" - assert param_id not in self.param_id_to_swap_buffer.keys( - ), f"param {param_id} has already been assigned a swap buffer" - - buffer_id = self.available_buffer_ids.pop() - print_rank_0(f"param {param.ds_id} is assigned swap in buffer id {buffer_id} ") - self.param_id_to_buffer_id[param_id] = buffer_id - aligned_swap_numel = self._io_aligned_numel(self.param_id_to_numel[param_id]) - swap_buffer = self.buffers.narrow(0, int(buffer_id * self.aligned_elements_per_buffer), aligned_swap_numel) - - self.param_id_to_swap_buffer[param_id] = swap_buffer - compute_buffer = swap_buffer.narrow(0, 0, self.param_id_to_numel[param_id]) - compute_buffers.append(compute_buffer) - swap_buffers.append(swap_buffer) - - return compute_buffers, swap_buffers - - #waits for inflight nvme write to complete - def synchronize_writes(self): - if self.pending_writes == 0: - return - assert self.pending_writes == self.aio_write_handle.wait() - self.pending_writes = 0 - self.remove_partition_and_release_buffers(self.swap_out_params) - self.swap_out_params = [] - - #waits for inflight nvme reads to complete - def synchronize_reads(self): - if self.pending_reads == 0: - return - - assert self.pending_reads == self.aio_read_handle.wait() - - self.pending_reads = 0 - - for param, swap_in_buffer in zip(self.inflight_params, self.inflight_swap_in_buffers): - param_id = param.ds_id - compute_buffer = swap_in_buffer.narrow(0, 0, self.param_id_to_numel[param_id]) - param.ds_tensor.data = compute_buffer.data - param.ds_tensor.status = PartitionedParamStatus.AVAILABLE - - self.available_params.update([param.ds_id for param in self.inflight_params]) - self.available_numel += self.inflight_numel - - self.inflight_params = [] - self.inflight_swap_in_buffers = [] - self.inflight_numel = 0 - - #Removes the memory assignment and releases the buffers - #Should only be executed after swapping out the tensors - def remove_partition_and_release_buffers(self, params): - for param in params: - param_id = param.ds_id - - if param_id in self.param_id_to_buffer_id.keys(): - - buffer_id = self.param_id_to_buffer_id[param_id] - - assert buffer_id is not None, "Missing buffer id for releasing" - - self.available_buffer_ids.append(buffer_id) - del self.param_id_to_buffer_id[param_id] - del self.param_id_to_swap_buffer[param_id] - print_rank_0(f"param {param.ds_id} releases buffer id {buffer_id} ") - - if param_id in self.available_params: - self.available_params.remove(param_id) - self.available_numel -= self.param_id_to_numel[param_id] - - param.ds_tensor.data = self.invalid_buffer.data - param.ds_tensor.status = PartitionedParamStatus.NOT_AVAILABLE - - #writes from in memory to nvme. Does not release the buffers - def _swap_out(self, params, async_op=True): - - swap_out_paths = self._get_swap_paths(params) - swap_out_params = self._get_swap_buffers(params) - self._track_numel(params) - - swap_out_tensors(self.aio_write_handle, swap_out_params, swap_out_paths) - - self.pending_writes += len(swap_out_params) - self.swap_out_params += params - - if not async_op: - self.synchronize_writes() - - #blocking swap out followed by releasing the memory buffers - def swap_out_and_release(self, params, async_op=False, force_buffer_release=False): - if async_op: - assert force_buffer_release, "Should not release preallocated buffers without completing the swap out. Set force_buffer_release to True to do it anyways" - self._swap_out(params, async_op=async_op) - - # book keeping function for inflight swap in - def _update_inflight_swap_in(self, params, swap_in_buffers, inflight_numel): - self.inflight_params.extend(params) - self.inflight_swap_in_buffers.extend(swap_in_buffers) - self.inflight_numel += inflight_numel - - for param in params: - param.ds_tensor.status = PartitionedParamStatus.INFLIGHT - - self.pending_reads += len(params) - - #assigns an in memory buffer and swaps in from nvme - def swap_in(self, params, async_op=True, swap_in_buffers=None): - - assert all([param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE - for param in params]), "Some params are already available or in flight" - swap_in_paths = self._get_swap_paths(params) - - if swap_in_buffers is None: - if len(self.available_buffer_ids) < len(swap_in_paths): - ids = [p.ds_id for p in params] - print_rank_0( - f'Not enough swap in buffers {len(self.available_buffer_ids)} for {len(swap_in_paths)} params, ids = {ids}', - force=True) - print_rank_0( - f'Num inflight: params {len(self.inflight_params)}, buffers {len(self.inflight_swap_in_buffers)}, numel = {self.inflight_numel}', - force=True) - print_rank_0( - f'Num available params: count = {len(self.available_params)}, ids = {self.available_params}, numel = {self.available_numel}', - force=True) - - assert len(swap_in_paths) <= len( - self.available_buffer_ids - ), f"Not enough buffers {len(self.available_buffer_ids)} for swapping {len(swap_in_paths)}" - compute_buffers, swap_in_buffers = self._allocate_and_return_buffers_for_swap_in(params) - inflight_numel = sum([t.numel() for t in compute_buffers]) - else: - inflight_numel = sum([t.numel() for t in swap_in_buffers]) - - swap_in_tensors(self.aio_read_handle, swap_in_buffers, swap_in_paths) - - self._update_inflight_swap_in(params, swap_in_buffers, inflight_numel) - - if not async_op: - self.synchronize_reads() - - # Enables swapping into buffer that is out the control of swapper. This is always synchronous - def swap_into_buffer(self, param, dest_buffer): - assert param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE, f"param {param.ds_id} is already available or inflight" - - require_swap_buffer = not (get_accelerator().is_pinned(dest_buffer) - and self._is_io_aligned(dest_buffer.numel())) - - if require_swap_buffer: - assert len(self.available_buffer_ids) > 0, f"No buffer available to swap param {param.ds_id}." - compute_buffers, swap_in_buffers = self._allocate_and_return_buffers_for_swap_in([param]) - inflight_numel = compute_buffers[0].numel() - else: - swap_in_buffers = [dest_buffer] - inflight_numel = dest_buffer.numel() - - swap_in_paths = self._get_swap_paths([param]) - - swap_in_tensors(self.aio_read_handle, swap_in_buffers, swap_in_paths) - self._update_inflight_swap_in([param], swap_in_buffers, inflight_numel) - self.synchronize_reads() - - if require_swap_buffer: - dest_buffer.data.copy_(param.ds_tensor.data) - # Release swap buffer memory assignment. Note, this will mark the parameter not available. - self.remove_partition_and_release_buffers([param]) - - #assign a buffer to a param and return the buffer - def get_buffer(self, param, numel): - param_id = param.ds_id - - assert self.available_swap_in_buffers( - ) > 0, f"No swap buffers to allocate for fp16 param {param_id} of numel = {numel}" - assert numel <= self.elements_per_buffer, f"More elements {numel} than buffer size {self.elements_per_buffer}" - - self.param_id_to_numel[param_id] = numel - buffer_id = self.available_buffer_ids.pop() - self.param_id_to_buffer_id[param_id] = buffer_id - aligned_swap_numel = self._io_aligned_numel(self.param_id_to_numel[param_id]) - swap_buffer = self.buffers.narrow(0, int(buffer_id * self.aligned_elements_per_buffer), aligned_swap_numel) - - self.param_id_to_swap_buffer[param_id] = swap_buffer - compute_buffer = swap_buffer.narrow(0, 0, self.param_id_to_numel[param_id]) - print_rank_0(f"param {param.ds_id} is assigned swap in buffer id {buffer_id}") - return compute_buffer - - def reserve_available_buffers(self): - buffers = [] - for id in self.available_buffer_ids: - buffers.append( - self.buffers.narrow(0, int(id * self.aligned_elements_per_buffer), - int(self.aligned_elements_per_buffer))) - self.reserved_buffer_ids.append(id) - - self.available_buffer_ids = [] - return buffers - - def release_reserved_buffers(self): - for id in self.reserved_buffer_ids: - self.available_buffer_ids.append(id) - self.reserved_buffer_ids = [] - - def _io_aligned_numel(self, numel): - remainder = numel % self.numel_alignment - return numel if remainder == 0 else (numel + self.numel_alignment - remainder) - - def _is_io_aligned(self, numel): - return (numel % self.numel_alignment) == 0 - - def reserve_partitioned_swap_space(self, partition_num_elems): - aligned_numel = sum([self._io_aligned_numel(numel) for numel in partition_num_elems]) - self.partitioned_swap_buffer = get_accelerator().pin_memory(torch.zeros(aligned_numel, - device='cpu', - dtype=self.dtype), - align_bytes=0) - self.partitioned_swap_pool = SwapBufferPool([self.partitioned_swap_buffer]) - - def swap_out_partitioned_params(self, dst_fp16_params, src_fp32_params): - assert self.partitioned_swap_buffer is not None, f'partitioned swap buffers for fp16 params not initialized' - assert self.partitioned_swap_pool is not None, f'partitioned swap pool for fp16 params not initialized' - assert len(dst_fp16_params) == len(src_fp32_params), \ - f'mismatch in number of fp16 params {len(dst_fp16_params)} and fp32 params {len(src_fp32_params)}' - - fp16_swap_paths = self._get_swap_paths(dst_fp16_params, must_exist=True) - self.synchronize_writes() - self.partitioned_swap_pool.reset() - for i, fp32_tensor in enumerate(src_fp32_params): - swap_tensor, _ = self.partitioned_swap_pool.insert_tensor(fp32_tensor, fp16_swap_paths[i], - self._io_aligned_numel(fp32_tensor.numel())) - assert swap_tensor is not None - dst_fp16_params[i].ds_tensor.status = PartitionedParamStatus.AVAILABLE - - self.partitioned_swap_pool.swap_out(self.aio_write_handle) - - for param in dst_fp16_params: - param.ds_tensor.status = PartitionedParamStatus.NOT_AVAILABLE diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py deleted file mode 100644 index 1ff570ed3bb9dda06c68de5be8a7c105b98694d7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py +++ /dev/null @@ -1,242 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping optimizer tensors to/from (NVMe) storage devices. -""" - -from deepspeed.ops.op_builder import AsyncIOBuilder -from deepspeed import comm as dist -import torch - -from deepspeed.runtime.swap_tensor.constants import * -from deepspeed.runtime.swap_tensor.utils import swap_in_tensors, swap_out_tensors, print_object -from deepspeed.runtime.swap_tensor.async_swapper import AsyncTensorSwapper -from deepspeed.runtime.swap_tensor.utils import get_sized_buffer -from deepspeed.runtime.swap_tensor.optimizer_utils import OptimizerSwapper - - -class OptimizerSwapOp(object): - - def __init__(self, aio_handle, read_op, param_info, allocated_buffers, state_buffers, num_ops): - self.aio_handle = aio_handle - self.read_op = read_op - self.param_info = param_info - self.allocated_buffers = allocated_buffers - self.state_buffers = state_buffers - self.wait_required = True - self.num_ops = num_ops - - def is_parameter(self, parameter): - return OptimizerSwapper.parameter_id(parameter) == self.param_info.param_id - - def wait(self): - assert self.wait_required - assert self.aio_handle.wait() == self.num_ops - self.wait_required = False - - -SYNC_SWAP_IN = 'sync_swap_in' -ASYNC_SWAP_IN = 'async_swap_in' -SYNC_SWAP_OUT = 'sync_swap_out' -ASYNC_SWAP_OUT = 'async_swap_out' - -SWAP_IN_STATE_TIMER = 'swap_in_state' -SWAP_OUT_STATE_TIMER = 'swap_out_state' -SWAP_OUT_GRADIENT_TIMER = 'swap_out_gradient' -ASYNC_SWAP_IN_STATE_TIMER = "async_swap_in_state" -ASYNC_SWAP_OUT_STATE_TIMER = 'async_swap_out_state' - - -class PipelinedOptimizerSwapper(OptimizerSwapper): - - def __init__(self, swap_config, aio_config, base_folder, optimizer, largest_numel, device, dtype, timers): - super(PipelinedOptimizerSwapper, self).__init__(swap_config, aio_config, base_folder, optimizer, largest_numel, - device, dtype, timers) - - aio_op = AsyncIOBuilder().load() - self.write_aio_handle = aio_op.aio_handle(block_size=aio_config[AIO_BLOCK_SIZE], - queue_depth=aio_config[AIO_QUEUE_DEPTH], - single_submit=aio_config[AIO_SINGLE_SUBMIT], - overlap_events=aio_config[AIO_OVERLAP_EVENTS], - intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM]) - - self.read_aio_handle = aio_op.aio_handle(block_size=aio_config[AIO_BLOCK_SIZE], - queue_depth=aio_config[AIO_QUEUE_DEPTH], - single_submit=aio_config[AIO_SINGLE_SUBMIT], - overlap_events=aio_config[AIO_OVERLAP_EVENTS], - intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM]) - - # Overlap gradient swap out - self.gradient_swapper = AsyncTensorSwapper(aio_handle=self.write_aio_handle, - numel_alignment=self.numel_alignment, - timers=self.timers) - - self.async_swap_in = swap_config.pipeline_read - self.async_swap_out = swap_config.pipeline_write - - self.swap_ops = {SYNC_SWAP_IN: None, ASYNC_SWAP_IN: None, SYNC_SWAP_OUT: None, ASYNC_SWAP_OUT: None} - - self.print_exclude_list += [ - 'gradient_swapper', 'read_aio_handle', 'write_aio_handle', 'swap_ops', 'print_exclude_list' - ] - - if dist.get_rank() == 0: - print_object(obj=self, name='PipelinedOptimizerSwapper', exclude_list=self.print_exclude_list) - - def initialize_parameters(self, parameters, src_tensors): - self._initialize_parameters(parameters=parameters, src_tensors=src_tensors, aio_handle=self.write_aio_handle) - - def initialize_from_swapped_fp16_params(self, fp16_partitions_info, fp16_num_elems, fp16_pinned_buffers, - fp32_parameters): - self._initialize_from_swapped_fp16_params(aio_handle=self.write_aio_handle, - fp16_partitions_info=fp16_partitions_info, - fp16_num_elems=fp16_num_elems, - fp16_pinned_buffers=fp16_pinned_buffers, - fp32_parameters=fp32_parameters) - - def flush_gradients(self): - self._flush_gradient_swapper(self.gradient_swapper) - - def swap_in_optimizer_state(self, parameter, async_parameter): - assert parameter is not None - assert self.swap_ops[SYNC_SWAP_IN] is None - - self._flush_gradient_swapper(self.gradient_swapper) - - self._start_timer(SWAP_IN_STATE_TIMER) - - if self.swap_ops[ASYNC_SWAP_IN]: - assert self.swap_ops[ASYNC_SWAP_IN].is_parameter(parameter) - self.swap_ops[SYNC_SWAP_IN] = self.swap_ops[ASYNC_SWAP_IN] - self.swap_ops[ASYNC_SWAP_IN] = None - else: - self.swap_ops[SYNC_SWAP_IN] = self._swap_in_optimizer_state(aio_handle=self.read_aio_handle, - parameter=parameter) - - if self.swap_ops[SYNC_SWAP_IN]: - self.swap_ops[SYNC_SWAP_IN].wait() - - if self.async_swap_in and async_parameter is not None: - assert self.swap_ops[ASYNC_SWAP_IN] is None - self.swap_ops[ASYNC_SWAP_IN] = self._swap_in_optimizer_state(aio_handle=self.read_aio_handle, - parameter=async_parameter) - - self._stop_timer(SWAP_IN_STATE_TIMER) - self.timer_names.add(SWAP_IN_STATE_TIMER) - - def swap_out_optimizer_state(self, parameter, async_swap): - self._start_timer(SWAP_OUT_STATE_TIMER) - - if self.swap_ops[ASYNC_SWAP_OUT]: - self._start_timer(ASYNC_SWAP_OUT_STATE_TIMER) - self._complete_swap_out(ASYNC_SWAP_OUT) - self._stop_timer(ASYNC_SWAP_OUT_STATE_TIMER) - self.timer_names.add(ASYNC_SWAP_OUT_STATE_TIMER) - - assert self.swap_ops[SYNC_SWAP_IN] is not None - assert not self.swap_ops[SYNC_SWAP_IN].wait_required - swap_op = self._swap_out_optimizer_state(aio_handle=self.write_aio_handle, - parameter=parameter, - swap_in_op=self.swap_ops[SYNC_SWAP_IN]) - self.swap_ops[SYNC_SWAP_IN] = None - - if self.async_swap_out and async_swap: - self.swap_ops[ASYNC_SWAP_OUT] = swap_op - else: - self.swap_ops[SYNC_SWAP_OUT] = swap_op - self._complete_swap_out(SYNC_SWAP_OUT) - - self._stop_timer(SWAP_OUT_STATE_TIMER) - self.timer_names.add(SWAP_OUT_STATE_TIMER) - - def swap_out_gradients(self, parameter, gradient_offsets, gradient_tensors): - self._swap_out_gradients(parameter=parameter, - gradient_offsets=gradient_offsets, - gradient_tensors=gradient_tensors, - gradient_swapper=self.gradient_swapper) - - def _complete_swap_out(self, swap_out_type): - self.swap_ops[swap_out_type].wait() - for buffer in self.swap_ops[swap_out_type].state_buffers: - buffer = torch.Tensor() - self.swap_buffer_manager.free(self.swap_ops[swap_out_type].allocated_buffers) - self.swap_ops[swap_out_type] = None - - def _swap_out_optimizer_state(self, aio_handle, parameter, swap_in_op): - assert swap_in_op.is_parameter(parameter) - - allocated_buffers = swap_in_op.allocated_buffers.copy() - swap_buffers = swap_in_op.state_buffers.copy() - - param_info = swap_in_op.param_info - self._update_param_state_info(param_info, parameter) - unpinned_tensors = param_info.get_unpinned_state_tensors() - - if len(unpinned_tensors) > 0: - new_alloc_buffers = self.swap_buffer_manager.allocate(num_elems=self._io_aligned_numel(param_info.numel()), - count=len(unpinned_tensors), - dtype=param_info.dtype()) - assert new_alloc_buffers is not None - - allocated_buffers += new_alloc_buffers - swap_buffers += new_alloc_buffers - - for pinned_dst, unpinned_src in zip(new_alloc_buffers, unpinned_tensors): - dst = get_sized_buffer(pinned_dst, unpinned_src.numel()) - dst.data.copy_(unpinned_src.data) - - swap_paths = param_info.get_swap_paths() - assert len(swap_paths) == len(swap_buffers) - - swap_out_tensors(aio_handle, swap_buffers, swap_paths) - - swap_out_op = OptimizerSwapOp(aio_handle=aio_handle, - param_info=param_info, - read_op=False, - allocated_buffers=allocated_buffers, - state_buffers=swap_buffers, - num_ops=len(swap_buffers)) - - return swap_out_op - - def _swap_in_optimizer_state(self, aio_handle, parameter): - param_info = self._get_param_swap_info(parameter) - if param_info is None: - return None - - num_swap_tensors = param_info.num_tensors() - required_buffer_count = num_swap_tensors + (1 if param_info.has_gradients() else 0) - aligned_numel = self._io_aligned_numel(param_info.numel()) - allocated_buffers = self.swap_buffer_manager.allocate(num_elems=aligned_numel, - count=required_buffer_count, - dtype=parameter.dtype) - assert allocated_buffers is not None, \ - f"PipelinedOptimizerSwapper ran out of swap buffers, try increasing 'buffer_count'" - - state_buffers = allocated_buffers[:num_swap_tensors] - param_info.set_swap_buffers(state_buffers, aligned_numel) - - swap_buffers = state_buffers.copy() - swap_paths = param_info.get_swap_paths() - - if param_info.has_gradients(): - parameter.grad = allocated_buffers[-1].narrow(0, 0, param_info.numel()) - if param_info.swapped_gradients: - swap_buffers += param_info.get_swap_gradient_buffers(parameter.grad) - swap_paths += param_info.get_swap_gradient_paths() - - swap_in_tensors(aio_handle, swap_buffers, swap_paths) - - if param_info.unswapped_gradients: - self._retrieve_unswapped_grad_partitions(swap_info=param_info, dest_buffer=parameter.grad) - - swap_in_op = OptimizerSwapOp(aio_handle=aio_handle, - param_info=param_info, - read_op=True, - allocated_buffers=allocated_buffers, - state_buffers=state_buffers, - num_ops=len(swap_buffers)) - - return swap_in_op diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/utils.py deleted file mode 100644 index 60a4004381d9d9080023bc843c1eb425afc51377..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/swap_tensor/utils.py +++ /dev/null @@ -1,241 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Functionality of swapping tensors to/from (NVMe) storage devices. -""" - -import torch -from deepspeed.utils.logging import logger -from deepspeed.accelerator import get_accelerator - -from deepspeed import comm as dist - -MIN_AIO_BYTES = 1024**2 -AIO_ALIGNED_BYTES = 1024 -MIN_SWAPPABLE_BYTES = MIN_AIO_BYTES - - -def swap_in_tensors(swap_handle, tensor_buffers, swap_paths): - for buffer, path in zip(tensor_buffers, swap_paths): - assert (swap_handle.async_pread(buffer, path, 0) == 0) - - -def swap_out_tensors(swap_handle, tensor_buffers, swap_paths): - for buffer, path in zip(tensor_buffers, swap_paths): - assert (swap_handle.async_pwrite(buffer, path, 0) == 0) - - -def print_object(obj, name, exclude_list=[]): - logger.info('{}:'.format(name)) - for arg in sorted(vars(obj)): - if not arg in exclude_list: - dots = '.' * (29 - len(arg)) - logger.info(' {} {} {}'.format(arg, dots, getattr(obj, arg))) - - -class SwapBuffer(object): - - def __init__(self, buffer): - self.buffer = buffer - self.reset() - - def reset(self): - self.offset = 0 - self.swap_tensors = {} - self.compute_tensors = {} - self.swap_paths = {} - self.num_elem = 0 - - def insert_tensor(self, tensor, swap_path, aligned_numel): - swap_tensor, compute_tensor = self.allocate_tensor(swap_path, tensor.numel(), aligned_numel) - compute_tensor.data.copy_(tensor.data) - return swap_tensor, compute_tensor - - def allocate_tensor(self, swap_path, numel, aligned_numel): - assert self.has_space(aligned_numel) - assert not self.offset in self.swap_tensors - - allocate_offset = self.offset - swap_tensor = self.buffer.narrow(0, allocate_offset, aligned_numel) - dest_tensor = swap_tensor.narrow(0, 0, numel) - - self.swap_tensors[allocate_offset] = swap_tensor - self.compute_tensors[allocate_offset] = dest_tensor - self.swap_paths[allocate_offset] = swap_path - self.offset += aligned_numel - self.num_elem += numel - - return self.swap_tensors[allocate_offset], self.compute_tensors[allocate_offset] - - def has_space(self, numel): - return (self.offset + numel) <= self.buffer.numel() - - def get_swap_tensors(self): - return [tensor for tensor in self.swap_tensors.values()] - - def get_swap_paths(self): - return [path for path in self.swap_paths.values()] - - def get_compute_tensors(self): - return [tensor for tensor in self.compute_tensors.values()] - - def get_num_elem(self): - return self.num_elem - - def get_swap_tensor(self, offset): - return self.swap_tensors.get(offset, None) - - def get_compute_tensor(self, offset): - return self.compute_tensors.get(offset, None) - - def get_swap_path(self, offset): - return self.swap_paths(offset, None) - - -class SwapBufferPool(object): - - def __init__(self, buffers): - assert all([get_accelerator().is_pinned(buf) for buf in buffers]) - self.buffers = [SwapBuffer(buf) for buf in buffers] - self.current_index = 0 - - def reset(self): - self.current_index = 0 - for buffer in self.buffers: - buffer.reset() - - def allocate_tensor(self, numel, swap_path, aligned_numel): - if self.has_space(aligned_numel): - swap_tensor, compute_tensor = self._get_current_buffer().allocate_tensor(swap_path, numel, aligned_numel) - return swap_tensor, compute_tensor - - return None, None - - def insert_tensor(self, tensor, swap_path, aligned_numel): - if self.has_space(aligned_numel): - swap_tensor, compute_tensor = self._get_current_buffer().insert_tensor(tensor, swap_path, aligned_numel) - return swap_tensor, compute_tensor - - return None, None - - def get_swap_tensors(self): - swap_tensors = [] - for buffer in self._get_used_buffers(): - swap_tensors += buffer.get_swap_tensors() - - return swap_tensors - - def get_swap_paths(self): - swap_paths = [] - for buffer in self._get_used_buffers(): - swap_paths += buffer.get_swap_paths() - - return swap_paths - - def get_compute_tensors(self): - compute_tensors = [] - for buffer in self._get_used_buffers(): - compute_tensors += buffer.get_compute_tensors() - - return compute_tensors - - def has_space(self, numel): - if self._get_current_buffer().has_space(numel): - return True - - if self.current_index == len(self.buffers) - 1: - return False - - self.current_index += 1 - return self._get_current_buffer().has_space(numel) - - def swap_out(self, aio_handle, async_op=False): - swap_tensors = self.get_swap_tensors() - swap_paths = self.get_swap_paths() - assert all([p is not None for p in swap_paths]) - - swap_out_tensors(aio_handle, swap_tensors, swap_paths) - - if not async_op: - assert len(swap_tensors) == aio_handle.wait() - - def swap_in(self, aio_handle, async_op=False): - swap_tensors = self.get_swap_tensors() - swap_paths = self.get_swap_paths() - assert all([p is not None for p in swap_paths]) - - swap_in_tensors(aio_handle, swap_tensors, swap_paths) - - if not async_op: - assert len(swap_tensors) == aio_handle.wait() - - def _get_current_buffer(self): - return self.buffers[self.current_index] - - def _get_used_buffers(self): - return self.buffers[:self.current_index + 1] - - -class SwapBufferManager(object): - - def __init__(self, num_elems, count, dtype): - self.num_elems = num_elems - self.count = count - self.dtype = dtype - self.all_buffers = [ - get_accelerator().pin_memory(torch.zeros(num_elems, device='cpu', dtype=dtype), align_bytes=0) - for _ in range(count) - ] - self.free_buffer_index = [i for i in range(count)] - self.used_buffer_index = {} - self.gigabytes = (self.all_buffers[0].element_size() * num_elems * count) / (1024**3) - - if dist.get_rank() == 0: - exclude_list = ['all_buffers'] - print_object(obj=self, name='SwapBufferManager', exclude_list=exclude_list) - - def allocate(self, num_elems, count, dtype): - assert dtype == self.dtype - assert num_elems <= self.num_elems - if count > len(self.free_buffer_index): - return None - - used_indices = self.free_buffer_index[-count:] - self.free_buffer_index = self.free_buffer_index[:-count] - - buffers = [] - for i in used_indices: - tmp_buffer = self.all_buffers[i].narrow(0, 0, num_elems) - buffers.append(tmp_buffer) - self.used_buffer_index[id(tmp_buffer)] = i - return buffers - - def allocate_all(self, num_elems, dtype): - return self.allocate(num_elems=num_elems, count=len(self.free_buffer_index), dtype=dtype) - - def free(self, buffers): - buffer_ids = [] - for buf in buffers: - buffer_ids.append(id(buf)) - - assert all([b_id in self.used_buffer_index for b_id in buffer_ids]) - - for b_id in buffer_ids: - self.free_buffer_index.append(self.used_buffer_index[b_id]) - del (self.used_buffer_index[b_id]) - - -def get_sized_buffer(buffer, num_elems): - assert num_elems <= buffer.numel(), \ - f'num_elems {num_elems} > buffer {buffer.numel()}' - return buffer.narrow(0, 0, num_elems) if num_elems < buffer.numel() else buffer - - -def get_sized_buffers(buffer_list, num_elems_list): - swap_buffers = [ - get_sized_buffer(buffer, num_elems) \ - for buffer, num_elems in zip(buffer_list, num_elems_list) - ] - return swap_buffers diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__init__.py deleted file mode 100644 index 3882393453512ac915a9786bbfcad349c249fb4a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .config import AUTOTP_MODE, get_tensor_parallel_config -from .tp_manager import TpTrainingManager diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 6e8235a66fd97e18bfef37dde0e26d1a8f65eea9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/config.cpython-310.pyc deleted file mode 100644 index ecbd9932baa7120398c4114eb11c69c1178b88bd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/tp_manager.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/tp_manager.cpython-310.pyc deleted file mode 100644 index 1317aeeb609f8359f3a906d85cb9700ca9846760..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/__pycache__/tp_manager.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/config.py deleted file mode 100644 index 957984e9f8b339c195e94c685a9b1aa10e26cfbf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/config.py +++ /dev/null @@ -1,84 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from enum import Enum -from deepspeed.runtime.config_utils import DeepSpeedConfigModel -import torch -from pydantic import Field -from typing import Optional - - -class AUTOTP_MODE(Enum): - TRAINING = "TRAINING" - INFERENCE = "INFERENCE" - - -class TPConfig(DeepSpeedConfigModel): - """ Configure tensor parallelism settings """ - - tp_size: int = 1 - """ Number of devices to split the model across using tensor parallelism. """ - - tp_grain_size: int = 1 - "The variable required by the autoTP parser has not been activated in training yet" - "as it depends on the gather logic that supports uneven partitioning. " - "Desired MLP/lm_head tp size granularity. DNN library favors tensor size in granularity of power of 2, we pick 64 as a default size." - - mpu: object = None - """ - A model parallelism unit object that implements - ``get_{model,data}_parallel_{rank,group,world_size}()``. - """ - - tp_group: object = None - - -class TPTrainingConfig(DeepSpeedConfigModel): - - dtype: torch.dtype = torch.float16 - """ - Desired model data type, will convert model to this type. - """ - - autotp_size: int = 0 - """ - In automatic tensor-parallelism training, 'tensor_parallel_size' - When set to 0, indicates that it is disabled. - """ - tp_overlap_comm: bool = False - """ Whether to overlap communication with computation. Currently, only allreduce supports overlap. """ - - tensor_parallel: TPConfig = Field({}, alias="tp") - """ - Configuration for tensor parallelism used to split the model across several - GPUs. Expects a dictionary containing values for :any:`DeepSpeedTPConfig`. - """ - - injection_policy_tuple: Optional[tuple] = None - #The following parameters are required by autoTP parser. - ######################################## - keep_module_on_host: bool = False - """ - When loading checkpoints to model parameters, they are moved to the device. In very large models - this might fill the device and cause OOM. Setting this flag to true, will keep checkpoints on - host and not move them directly to the device (giving an option to quantize checkpoint data before - moving it to the device for example). - """ - - replace_with_kernel_inject: bool = Field(False, alias="kernel_inject") - """ - Set to true to inject inference kernels for models such as, Bert, GPT2, - GPT-Neo and GPT-J. Otherwise, the injection_dict provides the names of two - linear layers as a tuple: - `(attention_output projection, transformer output projection)` - """ - ######################################## - - -def get_tensor_parallel_config(ds_config): - - if 'tensor_parallel' in ds_config: - return TPTrainingConfig(**ds_config['tensor_parallel']) - return TPTrainingConfig() diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/tp_manager.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/tp_manager.py deleted file mode 100644 index cf0b5a75c92aa3d426747917735835d1d11afe5f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/tensor_parallel/tp_manager.py +++ /dev/null @@ -1,66 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from .config import TPTrainingConfig, TPConfig -from deepspeed.utils import groups -import deepspeed.comm as dist - - -class TpTrainingManager(): - - def __init__(self, model, tp_size, dtype): - self.module = model - self.config = self._initialize_config(dtype) - - from deepspeed.module_inject.auto_tp import AutoTP - from deepspeed import get_accelerator - - # Parse model configuration - parser_dict = AutoTP.tp_parser(model) - print("AutoTP: ", parser_dict) - - # Initialize TP configuration and model - self._initialize_tp_config(tp_size) - self._get_model_config_generate() - - # Synchronize random number generator state across devices - _rng_state = get_accelerator().get_rng_state().to(get_accelerator().current_device_name()) - dist.broadcast(_rng_state, groups.get_tensor_model_parallel_src_rank(), self.tp_config.tp_group) - get_accelerator().set_rng_state(_rng_state.cpu()) - - # Apply injection policies - self._apply_policies(parser_dict) - - def _initialize_config(self, dtype): - """Initialize and return the DeepSpeed TP training configuration.""" - config = TPTrainingConfig() - config.dtype = dtype - return config - - def _apply_policies(self, parser_dict): - """Apply injection policies to the parsed modules.""" - for client_module, injection_policy in parser_dict: - self.config.injection_policy_tuple = injection_policy - self._apply_injection_policy(self.config, client_module) - - def _apply_injection_policy(self, config, client_module=None): - from deepspeed.module_inject import replace_transformer_layer - """Apply the given injection policy to a client module.""" - if isinstance(self.module, torch.nn.Module): - replace_transformer_layer(client_module, self.module, None, self.config, self.model_config) - - def _initialize_tp_config(self, tp_size): - """Perform TP configuration initialization.""" - self.tp_config = TPConfig() - self.tp_config.tp_size = tp_size - - groups._init_tp_mesh_device(tp_size) - self.tp_config.tp_group = groups.get_tensor_model_parallel_group() - self.config.tensor_parallel = self.tp_config - - def _get_model_config_generate(self): - """Generate and apply HF model configuration.""" - self.model_config = getattr(self.module, 'config', None) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/torch_autocast.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/torch_autocast.py deleted file mode 100644 index faaa457e33bbd4d33a1816ec3e2ecbad1499f05b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/torch_autocast.py +++ /dev/null @@ -1,100 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Iterable, Set, List, Union -import importlib - -import torch -from deepspeed.utils import logger - -LOWER_PRECISION_SAFE_MODULES = [ - torch.nn.Linear, - torch.nn.Conv1d, - torch.nn.Conv2d, - torch.nn.Conv3d, -] - -TORCH_AUTOCAST_INITIALIZED = False -_WARNED_NESTED_AUTOCAST = False - - -def _validate_auto_cast_settings(engine): - - assert not engine.fp16_enabled(), "Cannot enable both torch autocast and fp16" - assert not engine.bfloat16_enabled(), "Cannot enable both torch autocast and bfloat16" - assert not engine.zero_quantized_weights(), "Cannot enable both torch autocast and zero quantized weights" - - assert all(p.dtype == torch.float32 - for p in engine.parameters()), "All parameters must be float32 for torch autocast" - assert engine.communication_data_type == torch.float32, "Communication data type must be float32 for torch autocast" - - -def init_autocast_params(engine, dtype: torch.dtype, - torch_autocast_lower_precision_safe_modules: Union[None, List[str]]) -> None: - - _validate_auto_cast_settings(engine) - model = engine.module - - if torch_autocast_lower_precision_safe_modules is None: - lower_precision_safe_module_classes = LOWER_PRECISION_SAFE_MODULES - else: - lower_precision_safe_module_classes = [] - for module_name in torch_autocast_lower_precision_safe_modules: - try: - package_name, class_name = module_name.rsplit('.', 1) - module = importlib.import_module(package_name) - class_ = getattr(module, class_name) - lower_precision_safe_module_classes.append(class_) - except Exception as e: - raise ValueError(f"Failed to import lower precision safe module {module_name}: {e}") - - for module in model.modules(): - if module.__class__ in lower_precision_safe_module_classes: - for p in module.parameters(recurse=False): - p.autocast_dtype = dtype - - global TORCH_AUTOCAST_INITIALIZED - TORCH_AUTOCAST_INITIALIZED = True - - -def is_autocast_initialized() -> bool: - return TORCH_AUTOCAST_INITIALIZED - - -def get_default_autocast_lower_precision_modules() -> List[str]: - return [f"{cls.__module__}.{cls.__name__}" for cls in LOWER_PRECISION_SAFE_MODULES] - - -def get_autocast_dtype(param: torch.nn.Parameter) -> torch.dtype: - return param.autocast_dtype if hasattr(param, "autocast_dtype") else param.dtype - - -def has_autocast_dtype(param: torch.nn.Parameter) -> bool: - return hasattr(param, "autocast_dtype") - - -def get_all_autocast_dtypes(params: Iterable) -> Set[torch.dtype]: - return {get_autocast_dtype(p) for p in params} - - -def sort_dtypes(dtypes: List[torch.dtype]) -> List[torch.dtype]: - return sorted(dtypes, key=str) - - -def validate_nested_autocast(engine): - global _WARNED_NESTED_AUTOCAST - - if torch.is_autocast_enabled(): - if engine.torch_autocast_enabled(): - if not _WARNED_NESTED_AUTOCAST: - logger.warning( - "DeepSpeed detected torch.autocast context outside the engine. " - "This is unnecessary when torch.autocast is already enabled through the DeepSpeed config.") - _WARNED_NESTED_AUTOCAST = True - else: - raise AssertionError( - "torch.autocast is enabled outside DeepSpeed, but not in the DeepSpeed config. " - "Please enable torch.autocast through the DeepSpeed config to ensure the correct communication dtype is used." - ) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/utils.py deleted file mode 100644 index fa65b9041faf95447385637325dbcc03c1cf58bd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/utils.py +++ /dev/null @@ -1,1185 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -""" -Copyright NVIDIA/Megatron - -Helper functions and classes from multiple sources. -""" - -from collections.abc import Iterable -import os -import psutil -import gc -from math import sqrt - -from numpy import prod - -import torch -from torch.nn import functional as F -try: - from torch._six import inf -except ModuleNotFoundError: - from torch import inf -from typing import Union, List, Dict -from deepspeed import comm as dist -from deepspeed.moe.utils import is_moe_param -from deepspeed.utils import groups, logger -from deepspeed.utils.bwc import (bwc_tensor_model_parallel_rank, bwc_pipeline_parallel_world_size, - bwc_pipeline_parallel_group) -from deepspeed.runtime.constants import PIPE_REPLICATED -from deepspeed.accelerator import get_accelerator -from deepspeed.module_inject.policy import transpose - -torch_memory_reserved = get_accelerator().memory_reserved -torch_max_memory_reserved = get_accelerator().max_memory_reserved - - -class DummyOptim(): - """ - Dummy optimizer presents model parameters as a param group, this is - primarily used to allow ZeRO-3 without an optimizer - """ - - def __init__(self, params): - self.param_groups = [] - self.param_groups.append({'params': params}) - - -graph_cache = {} - - -def graph_process(replay_first_step, func, *args, **kwargs): - # `func` should only contain operations on the GPU - # Please ensure that the memory address of the data required by 'func' remains constant - if func.__name__ not in graph_cache: - cuda_stream = get_accelerator().Stream() - cuda_stream.wait_stream(get_accelerator().current_stream()) - with get_accelerator().stream(cuda_stream): - func(*args, **kwargs) - get_accelerator().current_stream().wait_stream(cuda_stream) - graph_cache[func.__name__] = get_accelerator().create_graph() - with get_accelerator().capture_to_graph(graph_cache[func.__name__]): - func(*args, **kwargs) - if replay_first_step: - get_accelerator().replay_graph(graph_cache[func.__name__]) - else: - get_accelerator().replay_graph(graph_cache[func.__name__]) - - -def noop_decorator(func): - return func - - -class noop_context(object): - - def __init__(self): - pass - - def __enter__(self): - pass - - def __exit__(self, exc_type, exc_val, exc_tb): - pass - - -def ensure_directory_exists(filename): - """Create the directory path to ``filename`` if it does not already exist. - - Args: - filename (str): A file path. - """ - dirname = os.path.dirname(filename) - os.makedirs(dirname, exist_ok=True) - - -def set_random_seed(seed): - """Set the random seed for common PRNGs used during training: random, numpy, and torch. - - Args: - seed (int): the seed to use - """ - import numpy - import random - random.seed(seed) - numpy.random.seed(seed) - torch.manual_seed(seed) - - -def is_model_parallel_parameter(p) -> bool: - if hasattr(p, 'model_parallel') and p.model_parallel: - return True - - if hasattr(p, 'tensor_model_parallel') and p.tensor_model_parallel: - return True - - return False - - -def copy_to_device(item, device, criterion_func): - """ - Return a copy of tensor on specified device. - Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts. - Parameters: - item: tensor to copy or (possibly nested) container of tensors to copy. - device: target device - criterion_func: Function to restrict copy operation to items meet criterion - - Returns: - None - """ - if criterion_func(item): - return item.to(device) - elif isinstance(item, list): - return [copy_to_device(v, device, criterion_func) for v in item] - elif isinstance(item, tuple): - return tuple([copy_to_device(v, device, criterion_func) for v in item]) - elif isinstance(item, dict): - return {k: copy_to_device(v, device, criterion_func) for k, v in item.items()} - else: - return item - - -def move_to_device(item, device, criterion_func=None): - """ - Move tensor on to specified device by changing the storage. - Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts. - Parameters: - item: tensor to move or (possibly nested) container of tensors to move. - device: target device - criterion_func: Function to restrict move operation to items meet criterion, defaults to `None` which is an equivalent to always move - - Returns: - None - """ - if (criterion_func is not None and criterion_func(item)): - device_copy = item.to(device) - item.data = device_copy.data - return item - elif isinstance(item, list): - return [move_to_device(v, device, criterion_func) for v in item] - elif isinstance(item, tuple): - return tuple([move_to_device(v, device, criterion_func) for v in item]) - elif isinstance(item, dict): - return {k: move_to_device(v, device, criterion_func) for k, v in item.items()} - else: - return item.to(device) - - -def get_norm_with_moe_layers_fast(all_groups_norm, group): - # This implementation standardizes the grad_norm across ranks. A more precise implementation can be found in 'get_norm_with_moe_layers'. - # Need to allreduce (avg) the norms across different ranks because moe params will not be synced during allreduce - scaled_norm = all_groups_norm * 1.0 / float(dist.get_world_size(group=group)) - scaled_norm_tensor = torch.tensor(scaled_norm, device=get_accelerator().current_device_name(), dtype=torch.float) - dist.all_reduce(scaled_norm_tensor, group=group) - all_groups_norm = scaled_norm_tensor.item() - #print(f"old = {all_groups_norm_old} and new = {all_groups_norm} at rank: {deepspeed.comm.get_rank()}") - return all_groups_norm - - -class CheckOverflow(object): - '''Checks for overflow in gradient across parallel process''' - - def __init__(self, param_groups=None, mpu=None, zero_reduce_scatter=False, deepspeed=None): - self.mpu = mpu - self.params = [] if param_groups else None - self.zero_reduce_scatter = zero_reduce_scatter - self.deepspeed = deepspeed - self.has_moe_params = False - if param_groups: - for group in param_groups: - for param in group: - self.params.append(param) - if is_moe_param(param): - self.has_moe_params = True - - def check_using_norm(self, norm_group, reduce_overflow=True): - # TODO: I don't think reduce_overflow is needed if mpu is None - overflow = -1 in norm_group - overflow_gpu = get_accelerator().FloatTensor([overflow]) - if self.has_moe_params: - # In this case, we need to do an all_reduce across - # the expert_parallel_group, so that if there was - # an overflow due to expert weights, we detect it - - # Only need to check groups.get_largest_expert_parallel_group() - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=groups._get_max_expert_parallel_group()) - if self.mpu is not None: - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_model_parallel_group()) - elif reduce_overflow: - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX) - dist.barrier() - overflow = overflow_gpu[0].item() - return bool(overflow) - - def check(self, param_groups=None): - params = [] - has_moe_params = False - if param_groups is None: - params = self.params - has_moe_params = self.has_moe_params - else: - assert param_groups is not None, \ - "self.params and param_groups both cannot be none" - - for group in param_groups: - for param in group: - params.append(param) - if is_moe_param(param): - has_moe_params = True - - return self.has_overflow(params, has_moe_params=has_moe_params) - - # `params` is a list / generator of torch.Variable - def has_overflow_serial(self, params): - for i, p in enumerate(params): - if p.grad is not None and self._has_inf_or_nan(p.grad.data, i): - return True - return False - - def has_overflow(self, params, has_moe_params=None): - if has_moe_params is None: - has_moe_params = self.has_moe_params - overflow = self.has_overflow_serial(params) - # Since each model parallel GPU carries only part of the model, - # make sure overflow flag is synced across all the model parallel GPUs - overflow_gpu = get_accelerator().ByteTensor([overflow]) - # deepspeed.comm.all_reduce(overflow_gpu, - # op=deepspeed.comm.ReduceOp.MAX, - # group=mpu.get_model_parallel_group()) - if has_moe_params: - # All reduce this across expert_parallel_group, so that if an expert - # overflows, we detect it here - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=groups._get_max_expert_parallel_group()) - if self.zero_reduce_scatter: - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=dist.get_world_group()) - elif self.mpu is not None: - if self.deepspeed is not None: - using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') - if (using_pipeline and self.deepspeed.pipeline_enable_backward_allreduce - is False) or (not using_pipeline and self.deepspeed.enable_backward_allreduce is False): - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_data_parallel_group()) - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_model_parallel_group()) - elif self.deepspeed is not None and self.deepspeed.enable_backward_allreduce is False: - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=dist.get_world_group()) - - overflow = overflow_gpu[0].item() - return bool(overflow) - - # `x` is a torch.Tensor - @staticmethod - def _has_inf_or_nan(x, i): - try: - # if x is half, the .float() incurs an additional deep copy, but it's necessary if - # Pytorch's .sum() creates a one-element tensor of the same type as x - # (which is true for some recent version of pytorch). - cpu_sum = float(x.float().sum()) - # More efficient version that can be used if .sum() returns a Python scalar - # cpu_sum = float(x.sum()) - except RuntimeError as instance: - # We want to check if inst is actually an overflow exception. - # RuntimeError could come from a different error. - # If so, we still want the exception to propagate. - if "value cannot be converted" not in instance.args[0]: - raise - return True - else: - if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum: - return True - return False - - -def _handle_overflow(cpu_sum, x, i): - import math - rank = dist.get_rank() - if rank == 0: - t_i = -1 - for v_i, v in enumerate(x.data.contiguous().view(-1)): - if not math.isfinite(float(v)): - t_i = v_i - break - logger.info(f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}") - - -def get_global_norm(norm_list): - """ Compute total from a list of norms - """ - total_norm = 0.0 - for norm in norm_list: - total_norm += norm**2.0 - # logger.info(f'norm_list = {norm_list} global = {sqrt(total_norm)}') - return sqrt(total_norm) - - -def clip_grad_norm_(parameters, max_norm, norm_type=2, mpu=None): - """Clips gradient norm of an iterable of parameters. - - This has been adapted from Nvidia megatron. We add norm averaging - to consider MoE params when calculating norm as they will result - in different norms across different ranks. - - This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and - added functionality to handle model parallel parameters. Note that - the gradients are modified in place. - - Arguments: - parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a - single Tensor that will have gradients normalized - max_norm (float or int): max norm of the gradients - norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for - infinity norm. - - Returns: - Total norm of the parameters (viewed as a single vector). - """ - if isinstance(parameters, torch.Tensor): - parameters = [parameters] - parameters = list(filter(lambda p: p.grad is not None, parameters)) - norm_type = float(norm_type) - all_norms = [] - if norm_type == inf: - for p in parameters: - all_norms.append(p.grad.data.abs().max().float()) - total_norm = torch.stack(all_norms).max() - total_norm = total_norm.to(get_accelerator().current_device_name()) - # Take max across all GPUs. - if mpu is not None: - dist.all_reduce(total_norm, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) - else: - total_norm = 0 - for p in parameters: - if mpu is not None: - if (mpu.get_model_parallel_rank() == 0) or is_model_parallel_parameter(p): - param_norm = p.grad.data.detach().float().norm(norm_type) - all_norms.append(param_norm) - else: - param_norm = p.grad.data.detach().float().norm(norm_type) - all_norms.append(param_norm) - if len(all_norms) > 0: - total_norm = torch.stack(all_norms).square().sum().float() - else: - total_norm = get_accelerator().FloatTensor([0.0]) - total_norm = total_norm.to(get_accelerator().current_device_name()) - # Sum across all model parallel GPUs. - if mpu is not None: - dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) - total_norm = total_norm.pow(1. / norm_type) - - # Need to average total_norm across different GPUs due to the presence of moe params - pg = groups._get_data_parallel_group() - scaled_norm = total_norm * 1.0 / float(dist.get_world_size(group=pg)) - scaled_norm_tensor = scaled_norm - - dist.all_reduce(scaled_norm_tensor, group=pg) - total_norm = scaled_norm_tensor - total_norm = total_norm.to(parameters[0].device) - - max_norm = torch.tensor([float(max_norm)], device=total_norm.device) - clip_coef = max_norm / (total_norm + 1e-6) - tmp_tensor = torch.tensor([1.0], device=clip_coef.device) - clip_coef = torch.min(tmp_tensor, clip_coef) - for p in parameters: - p.grad.data.mul_(clip_coef) - return total_norm - - -def get_flattened_grad_norm(parameters, norm_type=2, mpu=None, grad_norm_mask=None): - """Get grad norm of an iterable of parameters. - - This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and - added functionality to handle model parallel parameters. Note that - the gradients are modified in place. Taken from Nvidia Megatron. - - Arguments: - parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a - single Tensor that will have gradients normalized - norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for - infinity norm. - grad_norm_mask (List[Tensor]): A list of Tensor, where - each Tensor is a 2D Tensor containing ranges of [start_index, end_index]. - Returns: - Total norm of the parameters (viewed as a single vector). - """ - if isinstance(parameters, torch.Tensor): - parameters = [parameters] - parameters = list(filter(lambda p: p.grad is not None, parameters)) - - norm_type = float(norm_type) - if norm_type == inf: - total_norm = max(p.grad.data.abs().max() for p in parameters) - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - # Take max across all GPUs. - if mpu is not None: - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) - total_norm = total_norm_cuda[0].item() - else: - total_norm = 0. - for idx, p in enumerate(parameters): - # Use grad_norm_mask to avoid redundant computation of flattened gradient norm - if grad_norm_mask is not None and len(grad_norm_mask[idx]) > 0: - - # A loop-free implementation to create a mask tensor based on a range list - # which is logically equivalent to the following implementation. - # # mask_tensor_ = torch.zeros_like(p, device=p.device, dtype=bool) - # # for mask_idx in grad_norm_mask[idx]: - # # mask_tensor_[mask_idx[0]:mask_idx[1]] = True - cum_sum_pairs = torch.tensor([1, -1], device=get_accelerator().current_device_name(), - dtype=p.dtype).repeat(grad_norm_mask[idx].shape[0], 1) - mask_tensor = torch.zeros(p.shape[0] + 1, - device=get_accelerator().current_device_name(), - dtype=p.dtype) - mask_tensor = mask_tensor.scatter_(0, grad_norm_mask[idx].view(-1), - cum_sum_pairs.view(-1)).cumsum(0).bool()[:-1] - - param_norm = torch.masked_fill(p.grad.data, mask_tensor, 0).float().norm(norm_type) - - else: - param_norm = p.grad.data.float().norm(norm_type) - total_norm += param_norm.item()**norm_type - - # Sum across all model parallel GPUs. - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - if mpu is not None: - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) - total_norm = total_norm_cuda[0].item()**(1. / norm_type) - - if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm: - total_norm = -1 - - return total_norm - - -def get_grad_zeros(parameters, mpu=None): - """Compute the number of grads with zero values. - - This is adapted from get_grad_norm - - Arguments: - parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a - single Tensor that will have gradients normalized - - Returns: - Total number of params with zero values (viewed as a single vector). - """ - if isinstance(parameters, torch.Tensor): - parameters = [parameters] - parameters = list(filter(lambda p: p.grad is not None, parameters)) - - total_zeros = 0. - tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu) - for p in parameters: - # Pipeline parallelism may replicate parameters. Avoid multi-counting. - if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: - continue - - # Filter to avoid over-counting replicated tensors from tensor - # model parallelism - if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p): - continue - - count_zeros = p.grad.numel() - torch.count_nonzero(p.grad) - total_zeros += count_zeros.item() - - # Sum across all model parallel GPUs. - total_zeros_cuda = get_accelerator().FloatTensor([float(total_zeros)]) - if mpu is not None: - dist.all_reduce(total_zeros_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) - total_zeros = total_zeros_cuda[0].item() - - return total_zeros - - -def get_weight_norm(parameters, norm_type=2, mpu=None): - """Get norm of an iterable of parameters. - - This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and - added functionality to handle model parallel parameters. Note that - the gradients are modified in place. Taken from Nvidia Megatron. - - Arguments: - parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a - single Tensor that will have gradients normalized - norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for - infinity norm. - - Returns: - Total norm of the parameters (viewed as a single vector). - -1 if the norm value is NaN or Inf. - """ - if isinstance(parameters, torch.Tensor): - parameters = [parameters] - - norm_type = float(norm_type) - if norm_type == inf: - total_norm = max(p.data.abs().max() for p in parameters) - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - # Take max across all GPUs. - if mpu is not None: - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) - total_norm = total_norm_cuda[0].item() - else: - total_norm = 0. - tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu) - for p in parameters: - # Pipeline parallelism may replicate parameters. Avoid multi-counting. - if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: - continue - - # Filter to avoid over-counting replicated tensors from tensor - # model parallelism - if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p): - continue - - param_norm = p.data.float().norm(norm_type) - total_norm += param_norm**norm_type - - # Sum across all model parallel GPUs. - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - if mpu is not None: - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) - total_norm = total_norm_cuda[0].item()**(1. / norm_type) - - if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm: - total_norm = -1 - - return total_norm - - -def prefix_sum_inc(weights): - """ Compute an inclusive prefix sum. - - Example: - >>> prefix_sum_inc([3,4,5]) - [3, 7, 12] - """ - weights_ = [w for w in weights] - for x in range(1, len(weights_)): - weights_[x] += weights_[x - 1] - return weights_ - - -def partition_uniform(num_items, num_parts): - import numpy - parts = [0] * (num_parts + 1) - # First check for the trivial edge case - if num_items <= num_parts: - for p in range(num_parts + 1): - parts[p] = min(p, num_items) - return parts - - chunksize = num_items // num_parts - residual = num_items - (chunksize * num_parts) - - parts = numpy.arange(0, (num_parts + 1) * chunksize, chunksize) - - for i in range(residual): - parts[i + 1:] += 1 - parts = parts.tolist() - - return parts - - -def partition_balanced(weights, num_parts): - """ - use dynamic programming solve `The Linear Partition Problem`. - see https://www8.cs.umu.se/kurser/TDBAfl/VT06/algorithms/BOOK/BOOK2/NODE45.HTM - """ - import numpy as np - n = len(weights) - m = num_parts - - if n <= m: - return partition_uniform(n, m) - - dp_max = np.full((n + 1, m + 1), np.inf) - dp_min = np.full((n + 1, m + 1), np.inf) - dp_cost = np.full((n + 1, m + 1), np.inf) - position = np.zeros((n + 1, m + 1), dtype=int) - prefix_sum = np.zeros((n + 1)) - prefix_sum[1:] = np.cumsum(weights) - - dp_max[0, 0] = 0 - dp_cost[0, 0] = 0 - for i in range(1, n + 1): - for j in range(1, min(i, m) + 1): - for k in range(i): - max_sum = max(dp_max[k, j - 1], prefix_sum[i] - prefix_sum[k]) - min_sum = min(dp_min[k, j - 1], prefix_sum[i] - prefix_sum[k]) - cost = max_sum - min_sum - if dp_cost[i, j] >= cost: - dp_cost[i, j] = cost - dp_max[i, j] = max_sum - dp_min[i, j] = min_sum - position[i, j] = k - - parts = [n] - for i in reversed(range(1, m + 1)): - parts.append(position[parts[-1], i]) - parts.reverse() - - return parts - - -class PartitionedTensor: - - def __init__(self, tensor, group, partition_meta=None): - super().__init__() - - self.group = group - self.num_parts = dist.get_world_size(group=self.group) - self.rank = dist.get_rank(group=self.group) - self.orig_size = list(tensor.size()) - self.orig_device = tensor.device - self.local_data, self.partition = self._partition_tensor(tensor) - self.even_split = tensor.numel() % self.num_parts == 0 - - @classmethod - def from_meta(cls, meta, local_part, group, device=get_accelerator().device_name()): - assert meta.dtype == torch.long - dummy = torch.ones(dist.get_world_size(group=group)) - part_obj = cls(tensor=dummy, group=group) - - meta = meta.tolist() - - # [N, list0, ..., listN-1] - part_obj.orig_size = meta[1:(1 + meta[0])] - meta = meta[1 + meta[0]:] - - part_obj.orig_device = device - part_obj.local_data = local_part.detach() - - part_obj.group = group - - # Partition is encoded like the rowptr of a CSR matrix: - # [num_parts, rank, 0, part_1, ..., part_num_parts] - # TODO: support shuffle between different partition granularities - assert part_obj.num_parts == meta[0] - assert part_obj.rank == meta[1] - part_obj.partition = meta[2:] # length num_parts+1 - - return part_obj - - def _partition_tensor(self, tensor): - partition = partition_uniform(num_items=tensor.numel(), num_parts=self.num_parts) - start = partition[self.rank] - length = partition[self.rank + 1] - start - tensor_part = tensor.detach().contiguous().view(-1).narrow(0, start=start, length=length).clone() - - return tensor_part, partition - - def full(self, device=None): - if device is None: - device = self.orig_device - - # Allocate the full tensor as a flat buffer. - full_numel = prod(self.full_size()) - flat_tensor = torch.zeros([full_numel], dtype=self.local_data.dtype, device=device) - if self.even_split: - # Collect the full tensor - dist.all_gather_into_tensor(flat_tensor, self.local_data, group=self.group) - else: - for part_id in range(self.num_parts): - part_size = self.partition[part_id + 1] - self.partition[part_id] - buf = flat_tensor.narrow(0, start=self.partition[part_id], length=part_size) - if part_id == self.rank: - buf.copy_(self.local_data) - dist.broadcast(buf, part_id, self.group) - return flat_tensor.view(self.full_size()).clone().detach() - - def to_meta(self): - """Returns a torch.LongTensor that encodes partitioning information. - - Can be used along with ``data()`` to serialize a ``PartitionedTensor`` for - communication. - - Returns: - torch.LongTensor: a tensor encoding the meta-information for the partitioning - """ - meta = [] - meta.append(len(self.orig_size)) - meta += list(self.orig_size) - meta.append(self.num_parts) - meta.append(self.rank) - meta += self.partition - return torch.LongTensor(data=meta).to(self.orig_device) - - def data(self): - return self.local_data - - def local_size(self): - return self.local_data.size() - - def full_size(self): - return self.orig_size - - -mem_alloced = 0 -mem_cached = 0 - - -def memory_status(msg, print_rank=-1, reset_max=False): - global mem_alloced, mem_cached - - rank = dist.get_rank() - if print_rank != -1 and rank != print_rank: - return - - get_accelerator().synchronize() - - if reset_max: - get_accelerator().reset_max_memory_cached() - get_accelerator().reset_max_memory_allocated() - - new_alloced = get_accelerator().memory_allocated() - new_cached = get_accelerator().memory_cached() - - delta_alloced = new_alloced - mem_alloced - delta_cached = new_cached - mem_cached - - mem_cached = new_cached - mem_alloced = new_alloced - - max_alloced = get_accelerator().max_memory_allocated() - max_cached = get_accelerator().max_memory_cached() - - # convert to GB for printing - new_alloced /= 1024**3 - new_cached /= 1024**3 - delta_alloced /= 1024**3 - delta_cached /= 1024**3 - max_alloced /= 1024**3 - max_cached /= 1024**3 - - print( - f'RANK={rank} MEMSTATS', msg, f'device={get_accelerator().current_device_name()} ' - f'current alloc={new_alloced:0.4f}GB (delta={delta_alloced:0.4f}GB max={max_alloced:0.4f}GB) ' - f'current cache={new_cached:0.4f}GB (delta={delta_cached:0.4f}GB max={max_cached:0.4f}GB)') - - -def get_ma_status(): - if dist.is_initialized() and not dist.get_rank() == 0: - return 0 - return get_accelerator().memory_allocated() - - -def empty_cache(): - get_accelerator().empty_cache() - get_accelerator().reset_peak_memory_stats() - - -def see_memory_usage(message, force=False): - if not force: - return - if dist.is_initialized() and not dist.get_rank() == 0: - return - - # python doesn't do real-time garbage collection so do it explicitly to get the correct RAM reports - gc.collect() - - # Print message except when distributed but not rank 0 - logger.info(message) - logger.info(f"MA {round(get_accelerator().memory_allocated() / (1024 * 1024 * 1024),2 )} GB \ - Max_MA {round(get_accelerator().max_memory_allocated() / (1024 * 1024 * 1024),2)} GB \ - CA {round(torch_memory_reserved() / (1024 * 1024 * 1024),2)} GB \ - Max_CA {round(torch_max_memory_reserved() / (1024 * 1024 * 1024))} GB ") - - vm_stats = psutil.virtual_memory() - used_GB = round(((vm_stats.total - vm_stats.available) / (1024**3)), 2) - logger.info(f'CPU Virtual Memory: used = {used_GB} GB, percent = {vm_stats.percent}%') - - # get the peak memory to report correct data, so reset the counter for the next call - get_accelerator().reset_peak_memory_stats() - - -def call_to_str(base, *args, **kwargs): - """Construct a string representation of a call. - - Args: - base (str): name of the call - args (tuple, optional): args to ``base`` - kwargs (dict, optional): kwargs supplied to ``base`` - - Returns: - str: A string representation of base(*args, **kwargs) - """ - name = f'{base}(' - if args: - name += ', '.join(repr(arg) for arg in args) - if kwargs: - name += ', ' - if kwargs: - name += ', '.join(f'{key}={repr(arg)}' for key, arg in kwargs.items()) - name += ')' - return name - - -def get_only_unique_item(items): - item_set = set(items) - if len(item_set) != 1: - raise RuntimeError(f"expected there to be only one unique element in {items}") - unique_item, = item_set - - return unique_item - - -def mask_nan_or_inf_with_val_inplace(input, device=None, val=-1.): - norm_is_inf = input.isinf() - norm_is_nan = input.isnan() - inf_or_nan = norm_is_nan.logical_or(norm_is_inf) - err = torch.tensor(-1.0, device=device, dtype=torch.float) - input.masked_fill_(inf_or_nan, err) - - -def get_global_norm_of_tensors(input_tensors, norm_type=2, mpu=None, use_graph=False, moe_ep_group=None): - """Get norm of an iterable of tensors. - - This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and - added functionality to handle model parallel parameters. Taken from Nvidia Megatron. - - Arguments: - input_tensors (Iterable[Tensor]): an iterable of Tensors will have norm computed - norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for - infinity norm. - - Returns: - Total norm of the tensors (viewed as a single vector). - """ - assert isinstance(input_tensors, Iterable), f'expected Iterable type not {type(input_tensors)}' - assert all([torch.is_tensor(t) for t in input_tensors]), f'expected list of only tensors' - - norm_type = float(norm_type) - all_norms = [] - if norm_type == inf: - for t in input_tensors: - all_norms.append(t.data.abs().max().float()) - total_norm = torch.stack(all_norms).max() - device_total_norm = total_norm.to(get_accelerator().current_device_name()) - # Max across model parallel - if mpu is not None: - # For MoE grads, max over model parallel only if MoE-TP is enabled - if moe_ep_group is None or groups._get_expert_model_parallel_world_size() > 1: - dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) - # If MoE grads and MoE-TP disabled, max over pipeline parallel - elif bwc_pipeline_parallel_world_size(mpu) > 1: - dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=bwc_pipeline_parallel_group(mpu)) - - # MoE grads: max across expert parallel group - if moe_ep_group is not None: - dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=moe_ep_group) - total_norm = device_total_norm.to(input_tensors[0].device) - else: - - if 'norm_tensors_compute_buffer' not in graph_cache or len( - graph_cache['norm_tensors_compute_buffer']) != len(input_tensors): - graph_cache['norm_tensors_compute_buffer'] = [ - torch.empty([], dtype=torch.float, device=get_accelerator().current_device_name()) - for t in input_tensors - ] - compute_buffer = graph_cache['norm_tensors_compute_buffer'] - - def _norm_tensors(tensor_list, _compute_buffer, _norm_type): - for i, t in enumerate(tensor_list): - _compute_buffer[i].data.copy_(t.data.float().norm(_norm_type)**_norm_type) - if i != 0: - _compute_buffer[0].data.add_(_compute_buffer[i].data) - - if use_graph: - graph_process(False, _norm_tensors, input_tensors, compute_buffer, norm_type) - else: - _norm_tensors(input_tensors, compute_buffer, norm_type) - - device_total_norm = compute_buffer[0].float().detach() - - # Sum across model parallel - if mpu is not None: - # For MoE grads, sum over model parallel only if MoE-TP is enabled - if moe_ep_group is None or groups._get_expert_model_parallel_world_size() > 1: - dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) - # If MoE grads and MoE-TP disabled, sum over pipeline parallel - elif bwc_pipeline_parallel_world_size(mpu) > 1: - dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=bwc_pipeline_parallel_group(mpu)) - - # MoE grads: sum across expert parallel group - if moe_ep_group is not None: - dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=moe_ep_group) - total_norm = device_total_norm.to(input_tensors[0].device).pow(1. / norm_type) - - mask_nan_or_inf_with_val_inplace(total_norm, device=total_norm.device) - - return total_norm - - -def clip_tensors_by_global_norm(input_tensors, max_norm=1.0, global_norm=None, mpu=None, eps=1e-6, use_graph=False): - """Clip list of tensors by global norm. - Args: - input_tensors: List of tensors to be clipped - global_norm (float, optional): Precomputed norm. Defaults to None. - mpu (optional): model parallelism unit. Defaults to None. - eps (float, optional): epsilon value added to grad norm. Defaults to 1e-6 - Returns: - float: the global norm - """ - if global_norm is None: - global_norm = get_global_norm_of_tensors(input_tensors, mpu=mpu, use_graph=use_graph) - clip_coef = max_norm / (global_norm + eps) - if clip_coef < 1: - if use_graph: - - def clip_tensors(_tensor_list, _clip_coef_tensor): - for t in _tensor_list: - t.detach().mul_(_clip_coef_tensor) - - if 'clip_coef_tensor' not in graph_cache: - # Alloc memory - graph_cache['clip_coef_tensor'] = torch.tensor(clip_coef, - dtype=torch.float32).to(get_accelerator().device_name()) - clip_coef_tensor = graph_cache['clip_coef_tensor'] - clip_coef_tensor.copy_(torch.tensor(clip_coef, dtype=torch.float32)) - graph_process(False, clip_tensors, input_tensors, clip_coef_tensor) - - else: - for t in input_tensors: - t.detach().mul_(clip_coef) - return global_norm - - -def align_dense_tensors(tensor_list, alignment): - num_elements = sum(t.numel() for t in tensor_list) - remaining = num_elements % alignment - - if remaining: - elements_to_add = alignment - remaining - pad_tensor = torch.zeros(elements_to_add, device=tensor_list[0].device, dtype=tensor_list[0].dtype) - padded_tensor_list = tensor_list + [pad_tensor] - else: - padded_tensor_list = tensor_list - - return padded_tensor_list - - -def all_gather_into_tensor_dp_groups(groups_flat, partitioned_param_groups, dp_process_group): - for group_id, (group_flat, partitioned_params) in enumerate(zip(groups_flat, partitioned_param_groups)): - partition_id = dist.get_rank(group=dp_process_group[group_id]) - dp_world_size = dist.get_world_size(group=dp_process_group[group_id]) - if dp_world_size == 1: - # no groups share optimizer states - # pipeline parallel with bf16 will default call this even if dp size = 1. - continue - dist.all_gather_into_tensor(group_flat, partitioned_params[partition_id], dp_process_group[group_id]) - - -def all_gather_dp_groups(groups_flat, partitioned_param_groups, dp_process_group, start_alignment_factor, - allgather_bucket_size): - if dist.has_all_gather_into_tensor(): - return all_gather_into_tensor_dp_groups(groups_flat, partitioned_param_groups, dp_process_group) - - for group_id, partitioned_params in enumerate(partitioned_param_groups): - # Sequential AllGather Best of both worlds - partition_id = dist.get_rank(group=dp_process_group[group_id]) - dp_world_size = dist.get_world_size(group=dp_process_group[group_id]) - - if dp_world_size == 1: - # no groups share optimizer states - # pipeline parallel with bf16 will default call this even if dp size = 1. - continue - num_shards = max(1, partitioned_params[partition_id].numel() * dp_world_size // allgather_bucket_size) - - shard_size = partitioned_params[partition_id].numel() // num_shards - - # Enforce nccl/rccl alignment of start location of each shard - shard_size = shard_size - (shard_size % start_alignment_factor) - - num_elements = shard_size - - assert shard_size * num_shards <= partitioned_params[partition_id].numel() - - for shard_id in range(num_shards): - - if shard_id == (num_shards - 1): - num_elements = partitioned_params[partition_id].numel() - shard_id * shard_size - - shard_list = [] - for dp_id in range(dp_world_size): - curr_shard = partitioned_params[dp_id].narrow(0, shard_id * shard_size, num_elements).detach() - shard_list.append(curr_shard) - - dist.all_gather(shard_list, shard_list[partition_id], dp_process_group[group_id]) - - -def get_tensor_bytes(item): - if torch.is_tensor(item): - return item.numel() * item.element_size() - elif isinstance(item, list): - return sum([get_tensor_bytes(v) for v in item]) - elif isinstance(item, tuple): - return sum([get_tensor_bytes(v) for v in item]) - elif isinstance(item, dict): - return sum([get_tensor_bytes(v) for v in item.values()]) - else: - return 0 - - -def _get_folder_size(folder): - size = 0 - for path, _, files in os.walk(folder): - size += sum([os.path.getsize(os.path.join(path, f)) for f in files]) - return size - - -def get_checkpoint_folder_size(save_dir, tag, local_rank=None): - if local_rank == 0: - folder = os.path.join(save_dir, tag) - size_tensor = torch.tensor(_get_folder_size(folder)).to(get_accelerator().device_name()) - else: - size_tensor = torch.tensor(0).to(get_accelerator().device_name()) - - dist.reduce(tensor=size_tensor, dst=0) - return int(size_tensor) - - -class TLinear(torch.nn.Linear): - - def __init__(self, orig_layer, name=""): - self.name = name - super().__init__(orig_layer.weight.shape[1], orig_layer.weight.shape[0], bias=(orig_layer.bias is not None)) - self.weight.data = transpose(orig_layer.weight.data) - self.bias = orig_layer.bias - self._fwd_func = self._fwd_bias_add if self.bias is not None else self._fwd - - def _fwd(self, input): - return F.linear(input, self.weight) - - def _fwd_bias_add(self, input): - return F.linear(input, self.weight, bias=self.bias) - - def forward(self, input): - return self._fwd_func(input) - - -def get_inactive_params(param_list): - from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus - return [param for param in param_list if (hasattr(param, 'ds_id') and \ - param.ds_status == ZeroParamStatus.NOT_AVAILABLE)] - - -def get_norm_with_moe_layers(non_expert_norm, mpu, expert_tensors, norm_type=2): - """ Compute the global norm with MoE experts - - Inputs: - non_expert_norm (float) : the calculated norm of the non-expert params - expert_tensors (Dict[ep_name, List[Tensor]): Dictionary of expert group name to list of grad tensors - norm_type (int): the norm to use - - Returns: - if norm is (-/+) inf, returns -1 - otherwise the global norm (float) - """ - - def to_tensor(v): - return get_accelerator().FloatTensor(float(v)).detach() - - group_norms = [non_expert_norm] - for exp_name, tensors in expert_tensors.items(): - group_norm = get_global_norm_of_tensors(input_tensors=tensors, - mpu=mpu, - norm_type=norm_type, - use_graph=False, - moe_ep_group=groups._get_expert_parallel_group(exp_name)) - group_norms.append(group_norm) - - # check if all norms are valid - group_norms = torch.stack([to_tensor(norm) for norm in group_norms]) - if group_norms.eq(-1).any(): - return -1 - - # combine norms - if norm_type == inf: - total_norm = group_norms.max().item() - else: - total_norm = group_norms.pow(norm_type).sum() - total_norm = total_norm.item()**(1. / norm_type) - if total_norm == float('inf') or total_norm == -float('inf'): - total_norm = -1 - - return total_norm - - -def _make_offload_state_key(key): - return f"{key}_offload_buffer" - - -def offload_adam_states(optimizer, device, pin_memory: bool = False, non_blocking: bool = False): - """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" - - def move_key(state, key): - offload_buf_key = _make_offload_state_key(key) - if offload_buf_key not in state: - state[offload_buf_key] = torch.empty_like(state[key], device=device) - if pin_memory: - state[offload_buf_key] = get_accelerator().pin_memory(state[offload_buf_key]) - state[offload_buf_key].copy_(state[key], non_blocking=non_blocking) - state[key].data = state[offload_buf_key] - - for _, state in optimizer.state.items(): - if "exp_avg" in state: - move_key(state, "exp_avg") - if "exp_avg_sq" in state: - move_key(state, "exp_avg_sq") - - -def reload_adam_states(optimizer, device, non_blocking: bool = False): - """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" - - def move_back_key(state, key): - state[key].data = state[_make_offload_state_key(key)].to(device, non_blocking=non_blocking) - - for _, state in optimizer.state.items(): - if "exp_avg" in state: - move_back_key(state, "exp_avg") - if "exp_avg_sq" in state: - move_back_key(state, "exp_avg_sq") - - -def compare_tensors_in_structures(inputs1: Union[List, Dict], inputs2: Union[List, Dict]) -> bool: - """ - Compare two lists or dictionaries for equality, including any tensors they may contain. - - Args: - inputs1: First input, either a list or a dictionary. - inputs2: Second input, either a list or a dictionary. - - Returns: - True if inputs1 and inputs2 are equal; False otherwise. - """ - if type(inputs1) != type(inputs2): # Ensure types match - return False - - if isinstance(inputs1, list) and isinstance(inputs2, list): - if len(inputs1) != len(inputs2): - return False - for val1, val2 in zip(inputs1, inputs2): - if isinstance(val1, torch.Tensor) and isinstance(val2, torch.Tensor): - val1 = val1.to(get_accelerator().current_device()) - val2 = val2.to(get_accelerator().current_device()) - if not torch.equal(val1, val2): - return False - elif val1 != val2: - return False - return True - - elif isinstance(inputs1, dict) and isinstance(inputs2, dict): - if inputs1.keys() != inputs2.keys(): - return False - for key in inputs1: - val1, val2 = inputs1[key], inputs2[key] - if isinstance(val1, torch.Tensor) and isinstance(val2, torch.Tensor): - val1 = val1.to(get_accelerator().current_device()) - val2 = val2.to(get_accelerator().current_device()) - if not torch.equal(val1, val2): - return False - elif val1 != val2: - return False - return True - - return False diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/weight_quantizer.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/weight_quantizer.py deleted file mode 100644 index 95d737614e594ee200d13617654714d6ae9d26f6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/weight_quantizer.py +++ /dev/null @@ -1,153 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from ..module_inject.replace_policy import HFBertLayerPolicy, replace_policies -from deepspeed.accelerator import get_accelerator - - -class WeightQuantization(object): - - def __init__(self, mlp_extra_grouping=True, mp_size=1): - self.dense_scales = [] - self.qkv_scales = [] - self.mlp4hh_scales = [] - self.mlph4h_scales = [] - self.mlp_extra_grouping = mlp_extra_grouping - self.mp_size = mp_size - - def quantize_data(self, data, quantize_bits, groups, key=None): - data_groups = torch.split(data.float().view(-1), data.numel() // groups) - max_d = [max(g.max(), g.min().abs()) for g in data_groups] - data_scale = [float(1 << quantize_bits) / (2 * mx + 1e-5) for mx in max_d] - data_int = [(g * s) for g, s in zip(data_groups, data_scale)] - data_int = [ - di.round().clamp(-(1 << (quantize_bits - 1)), (((1 << (quantize_bits - 1)) - 1))) for di in data_int - ] - data_int = torch.cat(data_int).reshape(data.shape) - data_int = data_int.to(torch.int8) - data_scale = torch.cat([s.unsqueeze(0).unsqueeze(0) for s in data_scale]) - return data_int, data_scale - - def is_mlp(self, data, merge_count=1): - return ((self.mp_size *data.shape[0] * merge_count) / data.shape[1] == 4 or \ - (self.mp_size *data.shape[1] * merge_count) / data.shape[0] == 4) - - def is_qkv(self, data): - return ((self.mp_size * data.shape[0]) / data.shape[1] == 3 or \ - (self.mp_size * data.shape[1]) / data.shape[0] == 3) - - def Quantize(self, value_list, quantize_bits, groups, key, merge_dim=0): - if self.mlp_extra_grouping and self.is_mlp(value_list[0], merge_count=len(value_list)): - groups *= 2 - q_scale = [] - index = 0 - for data in value_list: - data_int, data_scale = self.quantize_data(data, quantize_bits, groups, key) - q_scale.append(data_scale) - value_list[index] = data_int - index += 1 - q_scale = (1 / - torch.cat(q_scale, dim=merge_dim).to(get_accelerator().current_device_name()).view(-1).unsqueeze(0)) - if "mlp.dense_4h_to_h.weight" in key: - self.mlp4hh_scales.append(q_scale) - elif "mlp.dense_h_to_4h.weight" in key: - self.mlph4h_scales.append(q_scale) - elif "attention.query_key_value.weight" in key: - self.qkv_scales.append(q_scale) - else: - self.dense_scales.append(q_scale) - return value_list - - def merge_layer_scales(self, layer_scales): - max_dim = max([s.shape[-1] for s in layer_scales]) - layer_scales = [ - torch.cat((s, torch.zeros((1, max_dim - s.shape[-1]), device=get_accelerator().current_device_name())), - dim=-1) if s.shape[-1] < max_dim else s for s in layer_scales - ] - return torch.cat(layer_scales).unsqueeze(0) - - def merge_scales(self): - all_scales = [] - for dense_scale, qkv_scale, m4hh_scale, mh4h_scale in \ - zip(self.dense_scales, self.qkv_scales, self.mlp4hh_scales, self.mlph4h_scales): - all_scales.append(self.merge_layer_scales([qkv_scale, dense_scale, mh4h_scale, m4hh_scale])) - return torch.cat(all_scales) - - def merge_scales_split(self, split_count): - all_scales = [[] for _ in range(split_count)] - for dense_scale, qkv_scale, m4hh_scale, mh4h_scale in \ - zip(self.dense_scales, self.qkv_scales, self.mlp4hh_scales, self.mlph4h_scales): - dense_scale = torch.split(dense_scale, dense_scale.numel() // split_count) - qkv_scale = torch.split(qkv_scale, qkv_scale.numel() // split_count) - m4hh_scale = torch.split(m4hh_scale, m4hh_scale.numel() // split_count) - mh4h_scale = torch.split(mh4h_scale, mh4h_scale.numel() // split_count) - for s in range(split_count): - all_scales[s].append( - torch.cat([ - torch.cat((qkv_scale[s], torch.zeros_like(qkv_scale[s])), dim=1), - torch.cat((dense_scale[s], torch.zeros_like(dense_scale[s])), dim=1), mh4h_scale[s], - m4hh_scale[s] - ]).unsqueeze(0)) - for scales_a in all_scales: - torch.cat(scales_a) - return all_scales - - def sd_quantize_megatron(self, sd, quantize_bits, groups): - keys = sd.keys() - for key in keys: - value_list = [sd[key]] - if "attention.dense.weight" in key or "mlp.dense_4h_to_h.weight" in key or \ - "mlp.dense_h_to_4h.weight" in key or "attention.query_key_value.weight" in key: - value_list = self.Quantize(value_list, quantize_bits, groups, key=key) - sd[key] = value_list[0] - - all_scales = self.merge_scales() - return sd, all_scales - - def model_quantize(self, model, quantize_policy, quantize_bits, groups): - all_scales = [] - - def quantize_fn(layer, policy_cls): - policy = policy_cls(layer) - - _, qkvw, _, dense_w, _, _ = policy.attention() - _, _h4h_w, _, _4hh_w, _ = policy.mlp() - keys = [qkvw, dense_w, _h4h_w, _4hh_w] - layer_scales = [] - - for key in range(len(keys)): - if self.mlp_extra_grouping and self.is_mlp(keys[key]): - data_quantized, data_scale = self.quantize_data(keys[key], quantize_bits, groups * 2) - elif policy_cls is HFBertLayerPolicy and self.is_qkv(keys[key]): - data_quantized, data_scale = self.quantize_data(keys[key], quantize_bits, groups * 3) - else: - data_quantized, data_scale = self.quantize_data(keys[key], quantize_bits, groups) - keys[key].copy_(data_quantized) - layer_scales.append((1 / data_scale.to(get_accelerator().current_device_name()).view(-1).unsqueeze(0))) - all_scales.append(self.merge_layer_scales(layer_scales)) - return layer - - def _quantize_module(model, policies): - for name, child in model.named_children(): - if child.__class__ in policies: - quantize_fn, replace_policy = policies[child.__class__] - setattr(model, name, quantize_fn(child, replace_policy)) - else: - _quantize_module(child, policies) - - return model - - policy = {} - if quantize_policy is not None: - for layer_name, replace_policy in quantize_policy.items(): - policy.update({layer_name: (quantize_fn, replace_policy)}) - else: - for plcy in replace_policies: - policy.update({plcy._orig_layer_class: (quantize_fn, plcy)}) - - quantized_module = _quantize_module(model, policy) - - return quantized_module, torch.cat(all_scales) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__init__.py deleted file mode 100644 index 8b045dbc8f974c8f6403e2edb006b7c14c3d9cc1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .partition_parameters import ZeroParamType -from .partition_parameters import ZeroParamStatus -from .partition_parameters import Init -from .partition_parameters import GatheredParameters -from .partition_parameters import register_external_parameter -from .parameter_offload import DeepSpeedZeRoOffload -from .partition_parameters import DeepSpeedTensorOverride - -from .tiling import TiledLinear -from .tiling import TiledLinearReturnBias - -from .mics import MiCS_Init - -from .stage3 import unwrap_model_for_generation diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 651f08dd511a96e4a21939cde2e2819d0c74c743..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/config.cpython-310.pyc deleted file mode 100644 index 4765943f3629e264afcb9c3f3927ee9de1adc7db..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/contiguous_memory_allocator.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/contiguous_memory_allocator.cpython-310.pyc deleted file mode 100644 index 40967e9016e649a721d53cfdb8c90b26ecfe375f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/contiguous_memory_allocator.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/linear.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/linear.cpython-310.pyc deleted file mode 100644 index bac3dece5ff224eddb484be529608474d4ed8209..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/linear.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/mics.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/mics.cpython-310.pyc deleted file mode 100644 index 695917fa3480b80a2db8a838ad855c3f8767ea29..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/mics.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/mics_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/mics_utils.cpython-310.pyc deleted file mode 100644 index ac86028dad41449f134347f2b7db95afa6be6a05..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/mics_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/offload_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/offload_config.cpython-310.pyc deleted file mode 100644 index b11737c1691b4e101a442be47e1846db48155fbf..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/offload_config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/offload_states.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/offload_states.cpython-310.pyc deleted file mode 100644 index 3983e201168ef565e94bbfb72e9df96faf8f6e4e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/offload_states.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/parameter_offload.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/parameter_offload.cpython-310.pyc deleted file mode 100644 index f1c0491dfe2c429e74621e90ccfbc8b23d439200..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/parameter_offload.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partition_parameters.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partition_parameters.cpython-310.pyc deleted file mode 100644 index 81d620b1a031b9ab5f98580c9a2f9135844dad2a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partition_parameters.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_coordinator.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_coordinator.cpython-310.pyc deleted file mode 100644 index cbaf4a12736aa4ed1db876e1975e49afba95b2dd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_coordinator.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_profiler.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_profiler.cpython-310.pyc deleted file mode 100644 index 78d67e0ff01caed58f8ce07de90d03041b7bb242..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/partitioned_param_profiler.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/stage3.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/stage3.cpython-310.pyc deleted file mode 100644 index 6f299ff6c7d120b2d80ddaa3a4ff5c17f840db3a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/stage3.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/stage_1_and_2.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/stage_1_and_2.cpython-310.pyc deleted file mode 100644 index 894765eeada933e775a88e612ee202faf43d216d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/stage_1_and_2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/test.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/test.cpython-310.pyc deleted file mode 100644 index 0fcdf6fd73ab7b7c8b5232a16c7424c5cc00ce6c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/test.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/tiling.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/tiling.cpython-310.pyc deleted file mode 100644 index 0cea5e0070e1901cdca9f59b441459facb747fec..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/tiling.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index 786bdc412518f1f51790b12ca9aadfda38de7d64..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/config.py deleted file mode 100644 index 19ee9b51702e786773c3fd1781e3342031e805cc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/config.py +++ /dev/null @@ -1,361 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import sys -from typing import Optional, Dict, Any -from enum import Enum -from pydantic import Field, model_validator -from deepspeed.runtime.config_utils import get_scalar_param, pp_int, DeepSpeedConfigModel -from deepspeed.utils import logger -from .offload_config import DeepSpeedZeroOffloadParamConfig, DeepSpeedZeroOffloadOptimizerConfig, OffloadDeviceEnum - -# ZeRO optimization. By default, this optimization is not enabled. -# Users have to configure the desired optimization (0 means disabled) in params.json as below example: -ZERO_FORMAT = """ -ZeRO optimization should be enabled as: -"session_params": { - "zero_optimization": { - "stage": [0|1|2], - "stage3_max_live_parameters" : 1000000000, - "stage3_max_reuse_distance" : 1000000000, - "stage3_use_all_reduce_for_fetch_params": [true|false], - "stage3_module_granularity_threshold": 0, - "allgather_partitions": [true|false], - "use_multi_rank_bucket_allreduce": [true|false], - "allgather_bucket_size": 500000000, - "reduce_scatter": [true|false], - "contiguous_gradients" : [true|false] - "overlap_comm": [true|false], - "reduce_bucket_size": 500000000, - "load_from_fp32_weights": [true|false], - "cpu_offload": [true|false] (deprecated), - "cpu_offload_param" : [true|false] (deprecated), - "cpu_offload_use_pin_memory": [true|false] (deprecated), - "sub_group_size" : 1000000000000, - "offload_param": {...}, - "offload_optimizer": {...}, - "ignore_unused_parameters": [true|false], - "round_robin_gradients": [true|false], - "zero_hpz_partition_size": 1, - "zero_quantized_weights": [true|false], - "zero_quantized_nontrainable_weights": [true|false], - "zero_quantized_gradients": [true|false], - "memory_efficient_linear": [true|false], - "override_module_apply": [true|false], - "zeropp_loco_param": {...}, - "log_trace_cache_warnings" : [true|false], - } -} -""" - -ZERO_OPTIMIZATION = "zero_optimization" - - -def read_zero_config_deprecated(param_dict): - zero_config_dict = {} - zero_config_dict["stage"] = 1 if param_dict[ZERO_OPTIMIZATION] else 0 - if zero_config_dict["stage"] > 0: - zero_config_dict["allgather_bucket_size"] = get_scalar_param(param_dict, "allgather_size", 5e8) - logger.warning( - "DeepSpeedConfig: this format of ZeRO optimization setup is deprecated. Please use the following format: {}". - format(ZERO_FORMAT)) - return zero_config_dict - - -def get_zero_config(param_dict): - if ZERO_OPTIMIZATION in param_dict: - zero_config_dict = param_dict[ZERO_OPTIMIZATION] - if isinstance(zero_config_dict, bool): - zero_config_dict = read_zero_config_deprecated(param_dict) - else: - zero_config_dict = {} - return DeepSpeedZeroConfig(**zero_config_dict) - - -class ZeroStageEnum(int, Enum): - """ Enum class for possible zero stages """ - disabled = 0 - optimizer_states = 1 - gradients = 2 - weights = 3 - max_stage = 3 - - -class DeepSpeedZeroConfig(DeepSpeedConfigModel): - """ - Sets parameters for ZeRO optimizations. - """ - - stage: ZeroStageEnum = 0 - """ - Chooses different stages of ZeRO Optimizer. Stage 0, 1, 2, and 3 refer - to disabled, optimizer state partitioning, and optimizer+gradient state - partitioning, and optimizer+gradient+parameter partitioning, respectively. - """ - - contiguous_gradients: bool = True - """ - Copies the gradients to a contiguous buffer as they are produced. Avoids - memory fragmentation during backward pass. - """ - - reduce_scatter: bool = True - """ - Uses reduce or reduce scatter instead of allreduce to average gradients - """ - - reduce_bucket_size: int = Field(pp_int(5e8), ge=0) - """ - Number of elements reduced/allreduced at a time. Limits the memory required - for the allgather for large model sizes - """ - - use_multi_rank_bucket_allreduce: bool = True - """ - Combine the reduce buckets of the different ranks and do an All-Reduce instead of multiple Reduce ops. - This feature is useful when the model is small and we want to scale it on too many GPUs which therefore - reduces the message sizes of each packet. - """ - - allgather_partitions: bool = True - """ - Chooses between allgather collective or a series of broadcast collectives - to gather updated parameters from all the GPUs at the end of each step - """ - - allgather_bucket_size: int = Field(pp_int(5e8), ge=0) - """ - Number of elements allgathered at a time. Limits the memory required for - the allgather for large model sizes - """ - - overlap_comm: Optional[bool] = None # None for dynamic default value (see validator `overlap_comm_valid` below) - """ - Attempts to overlap the reduction of the gradients with backward computation - """ - - load_from_fp32_weights: bool = True - """ - Boolean indicating whether to initialize fp32 master weights from fp32 - copies in checkpoint (no precision loss) or from model's fp16 copies (with - precision loss). This can be used to initialize optimizer state even when - checkpoint is missing optimizer state. - """ - - elastic_checkpoint: bool = False - """ - Enable loading checkpoint that was saved by job with different GPU count. - No longer supported. - """ - - offload_param: Optional[DeepSpeedZeroOffloadParamConfig] = None - """ - Enable offloading of model parameters to CPU or NVMe. This frees up GPU - memory for larger models or batch sizes. Valid only with stage 3. Expects a - dictionary containing values for :any:`DeepSpeedZeroOffloadParamConfig`. - """ - - offload_optimizer: Optional[DeepSpeedZeroOffloadOptimizerConfig] = None - """ - Enable offloading of optimizer state to CPU or NVMe, and optimizer - computation to CPU. This frees up GPU memory for larger models or batch - sizes. Valid for ZeRO stage 1, 2, 3. Expects a dictionary containing values - for :any:`DeepSpeedZeroOffloadOptimizerConfig`. - """ - - sub_group_size: int = Field(pp_int(1e9), ge=0) - """ - Tile size for parameter processing to fit massive models (with trillions of - parameters). Used by ZeRO3-Offload and ZeRO-Infinity - """ - - cpu_offload_param: Optional[bool] = Field( - None, - json_schema_extra={ - "deprecated": True, - "new_param": "offload_param", - "new_param_fn": (lambda val: DeepSpeedZeroOffloadParamConfig(device=OffloadDeviceEnum.cpu) - if val else None) - }, - ) - """ Deprecated, please use ``offload_param`` """ - - cpu_offload_use_pin_memory: Optional[bool] = Field( - None, - json_schema_extra={ - "deprecated": True, - "new_param": "offload_param or offload_optimizer", - "set_new_param": False - }, - ) - """ Deprecated, please use ``offload_param`` or ``offload_optimizer`` """ - - cpu_offload: Optional[bool] = Field( - None, - json_schema_extra={ - "deprecated": - True, - "new_param": - "offload_optimizer", - "new_param_fn": (lambda val: DeepSpeedZeroOffloadOptimizerConfig(device=OffloadDeviceEnum.cpu) - if val else None) - }, - ) - """ Deprecated, please use ``offload_optimizer`` """ - - prefetch_bucket_size: int = Field(pp_int(5e7), ge=0, alias="stage3_prefetch_bucket_size") - """ - Maximum number of parameter elements to fetch ahead of use. Used by ZeRO3, - ZeRO3-Offload, ZeRO-Infinity, and ZeRO-Inference. - """ - - param_persistence_threshold: int = Field(pp_int(1e5), ge=0, alias="stage3_param_persistence_threshold") - """ - Do not partition parameters smaller than this threshold. Smaller values use - less memory, but can greatly increase communication (especially - latency-bound messages). - """ - - model_persistence_threshold: int = Field(pp_int(sys.maxsize, "sys.maxsize"), - ge=0, - alias="stage3_model_persistence_threshold") - """ - Maximum number of parameter elements that can be persisted in GPU and not - partitioned. This imposes an upper bound on the number of unpartitioned - parameters resulting from param_persistence_threshold setting. Used by - ZeRO3-Offload, ZeRO-Infinity and ZeRO-Inference. - """ - - max_live_parameters: int = Field(pp_int(1e9), ge=0, alias="stage3_max_live_parameters") - """ - The maximum number of parameters resident per GPU before releasing. Smaller - values use less memory, but perform more communication. - """ - - max_reuse_distance: int = Field(pp_int(1e9), ge=0, alias="stage3_max_reuse_distance") - """ - Do not release a parameter if it will be reused within this threshold of - parameters. Smaller values use less memory, but perform more communication. - """ - - gather_16bit_weights_on_model_save: bool = Field(False, alias="stage3_gather_16bit_weights_on_model_save") - """ - Consolidate the weights before saving the model by ``save_16bit_model()``. - Since the weights are partitioned across GPUs, they aren’t part of - ``state_dict``, so this function automatically gathers the weights when - this option is enabled and then saves the fp16 model weights. - """ - - module_granularity_threshold: int = Field(pp_int(0), alias="stage3_module_granularity_threshold") - """ - The granularity of a module is determined by the ratio of "parameter_count / (1 + descendant count)". - ZeRO3 classifies modules with a granularity below the threshold as fine-grained, - which are treated as integral units during parameter fetching. This reduces host overhead - and the separate allgather overhead introduced by hooks for fine-grained layers when fetching parameters. - """ - - use_all_reduce_for_fetch_params: bool = Field(False, alias="stage3_use_all_reduce_for_fetch_params") - """ - Use all_reduce op when fetching module parameters at stage3. This improves performance by reducing - the overhead of concatenation and slicing on the host. - """ - - stage3_gather_fp16_weights_on_model_save: bool = Field(False, - json_schema_extra={ - "deprecated": True, - "new_param": "gather_16bit_weights_on_model_save" - }) - """ Deprecated, please use ``gather_16bit_weights_on_model_save`` """ - - ignore_unused_parameters: bool = True - """ - Unused parameters in modules may be unexpected in static networks, but - could be normal in dynamic networks. This controls whether or not training - should terminate with an error message when unused parameters are detected. - This is set to ``True`` by default, which means unused parameters are - ignored and training continues. Now is just used in stage 2. - """ - - legacy_stage1: bool = False - """ - For backward-compatibility enable old ZeRO stage 1 implementation. Use at - your own risk, will be deprecated soon. - """ - - round_robin_gradients: bool = False - """ - Stage 1 and 2 optimization for CPU offloading that parallelizes gradient - copying to CPU memory among ranks by fine-grained gradient partitioning. - Performance benefit grows with gradient accumulation steps (more copying - between optimizer steps) or GPU count (increased parallelism). - """ - zero_hpz_partition_size: int = Field(1, ge=0) - """ - Number of ranks in zero parameters partitioning secondary group - """ - zero_quantized_weights: bool = False - """ - Boolean indicating whether to quantize zero parameters (weights) - for efficient all_gather comm - """ - zero_quantized_nontrainable_weights: bool = False - """ - Boolean indicating whether to quantize non-trainable zero parameters (weights) - for efficient memory usage and communication. Different from zero_quantized_weights - that stores the weights in original precision and only perform quantization during communication, - this flag will store the weights in quantized precision. This is useful for LoRA training. - """ - zero_quantized_gradients: bool = False - """ - Boolean indicating whether to use quantized zero gradients - for efficient all_2_all_reduce comm - """ - zeropp_loco_param: Optional[Dict[str, Any]] = None - """ - This dictionary contains parameters for using LoCo-Zero++, with two key parameters: - - `err_beta`: A coefficient for the moving average of quantization errors before and after gradient computation. - It ranges between 0 and 1, with a default value of 0.8. - - `reset_T`: The number of steps after which the moving-average error buffer is cleared. The default value is 1024. - These parameters can be adjusted based on performance needs. Example configuration in ds config: - "zeropp_loco_param": { "err_beta": 0.8, "reset_T": 1024 }. - See LoCo paper for more details: (https://arxiv.org/abs/2407.04480). - """ - - mics_shard_size: int = Field(-1, json_schema_extra={"new_param": "mics_shard_size"}) - - mics_hierarchical_params_gather: bool = False - - memory_efficient_linear: bool = True - """ - Use memory efficient linear implementation, for Stage 3. - """ - """ - Whether force load checkpoint in pipeline mode, current only for Stage 3. - """ - pipeline_loading_checkpoint: bool = False - - override_module_apply: bool = True - """ - Override nn.Module apply function, for Stage 3. - """ - - log_trace_cache_warnings: bool = False - """ - Whether to log warnings from trace cache, such as invalidation events. - """ - - # Validators - @model_validator(mode="after") - def overlap_comm_valid(self): - if self.overlap_comm is None: - self.overlap_comm = self.stage == ZeroStageEnum.weights - return self - - @model_validator(mode="after") - def offload_ratio_check(self): - offload_config = self.offload_optimizer - if offload_config and offload_config.ratio < 1.0: - assert self.stage == ZeroStageEnum.weights, "Partial offloading only supported for ZeRO Stage 3." - return self diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/contiguous_memory_allocator.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/contiguous_memory_allocator.py deleted file mode 100644 index 35b3d5c7dd5d333e621cc554a73bb3c90e2a450e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/contiguous_memory_allocator.py +++ /dev/null @@ -1,287 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from deepspeed import comm as dist - - -def print_rank_0(message): - if dist.get_rank() == 0: - print(message) - - -class ContiguousMemoryAllocator(object): - - def __init__(self, size, dtype, device): - self.buffer = torch.zeros(size, dtype=dtype, device=device) - - #address to contiguous size available - self.contiguous_sizes = {} - - self.contiguous_sizes[0] = size - - #tensor id to its address - self.tensor_addresses = {} - - #tensor address to its size - self.tensor_sizes = {} - - #tensor address to ids - self.tensor_ids = {} - - #id to tensors - self.tensor_map = {} - - #id to params. Maps each tensor buffer to list of parameters that uses it - self.id_to_params = {} - - self.total_size = size - self.total_free = size - self.largest_contiguous = size - self.max_allocated = 0 - - self.count = 0 - - #create a tensor of size from the pre-allocated buffer - #if not enough free space will fail - #if not enough contiguous space, will defragment and allocate - def allocate_tensor(self, size): - free_before = self.total_free - - assert size <= self.total_free, "Not enough memory in buffer. Allocation failed" - if self.largest_contiguous < size: - print_rank_0("Needs defragmentation to allocate. Before Defragmentation:") - self.print_allocation(resolution=100) - self._defragment_memory() - #set the param data to the new tensor buffer locations - self._reset_param_data() - print_rank_0("After defragmentation:") - self.print_allocation(resolution=100) - - self.total_free = self.total_free - size - - allocated = self.total_size - self.total_free - if allocated > self.max_allocated: - self.max_allocated = allocated - - tensor_address = self._get_new_tensor_address(size) - - ret_tensor = self._get_new_tensor(tensor_address, size) - print_rank_0( - f"Free before allocation {free_before}. Allocating {size}. Free after allocation {self.total_free}. Max allocated {self.max_allocated}" - ) - assert self.total_free + size == free_before, "Allocation bookkeeping error" - - return ret_tensor - - #assigns the tensor data to the param data and keeps track of the assignment - #any change the underlying buffer from defragmentation will cause a - #reassignment of the param data - def assign_to_param(self, tensor, param, numel, shape): - tensor_id = id(tensor) - - assert tensor_id in self.tensor_map.keys(), "No such tensor allocated by the allocator." - assert tensor.numel() >= numel, "Assert tensor buffer does is not large enough" - assert not tensor_id in self.id_to_params.keys(), "This tensor has already been assigned to a param" - - self.id_to_params[tensor_id] = [param] - - replicated_tensor = tensor.narrow(0, 0, numel).view(shape) - param.data = replicated_tensor.data - param.contiguous_tensor_id = tensor_id - - #deletes the tensor and frees up the underlying buffer - def release_tensor(self, tensor): - free_before = self.total_free - tensor_id = id(tensor) - tensor_size = tensor.numel() - self._release_tensor(tensor_id) - self._unassign_params(tensor_id) - self.total_free += tensor_size - print_rank_0( - f"Free before release {free_before}. Released {tensor.numel()}. Total free after {self.total_free}.") - assert self.total_free - tensor_size == free_before, "Release bookkeeping error" - - def release_tensor_with_id(self, tensor_id): - free_before = self.total_free - assert tensor_id in self.tensor_map.keys(), "Invalid tensor id" - tensor = self.tensor_map[tensor_id] - tensor_size = tensor.numel() - self._release_tensor(tensor_id) - self._unassign_params(tensor_id) - self.total_free += tensor_size - print_rank_0( - f"Free before release {free_before}. Released {tensor.numel()}. Total free after {self.total_free}.") - assert self.total_free - tensor_size == free_before, "Release bookkeeping error" - - #shows the current memory allocation at specified resolution - def print_allocation(self, resolution=200): - total_size = self.buffer.numel() * 1.0 - empty = [] - for addr, size in self.contiguous_sizes.items(): - start = int(addr * resolution / total_size) - end = int((addr + size) * resolution / total_size) - empty.extend(range(start, end)) - s = '' - for i in range(resolution): - s += '.' if i in empty else '|' - print_rank_0(s) - - def max_allocated(self): - return self.max_allocated - - #to be called after defragmentation that moves the tensor buffers - #this call reassigns the data of all the parameters using the tensor buffers - def _reset_param_data(self): - for id, tensor in self.tensor_map.items(): - for param in self.id_to_params[id]: - param.data = tensor.narrow(0, 0, param.numel()).view(param.data.shape).data - - def _unassign_params(self, tensor_id): - if tensor_id in self.id_to_params.keys(): - del self.id_to_params[tensor_id] - - def _release_tensor(self, tensor_id): - assert tensor_id in self.tensor_addresses, f"Tensor id {tensor_id} not found" - - address = self.tensor_addresses[tensor_id] - contiguous_size = self.tensor_map[tensor_id].numel() - - del self.tensor_addresses[tensor_id] - del self.tensor_ids[address] - del self.tensor_map[tensor_id] - del self.tensor_sizes[address] - - self._consolidate_address(address, contiguous_size) - self.largest_contiguous = self._largest_contiguous() - - def _consolidate_address(self, address, contiguous_size): - - #consolidate next buffer - end_address = address + contiguous_size - if end_address in self.contiguous_sizes: - contiguous_size += self.contiguous_sizes[end_address] - del self.contiguous_sizes[end_address] - - #consolidate previous buffer - for addr, size in self.contiguous_sizes.items(): - if addr + size == address: - del self.contiguous_sizes[addr] - contiguous_size += size - address = addr - break - - self.contiguous_sizes[address] = contiguous_size - - def _defragment_memory(self): - empty_addresses = sorted(self.contiguous_sizes.keys()) - tensor_addresses = sorted(self.tensor_addresses.values()) - - tensor_index = 0 - - while tensor_index < len(tensor_addresses): - - empty_addr = empty_addresses[0] - empty_size = self.contiguous_sizes[empty_addr] - - tensor_addr = tensor_addresses[tensor_index] - tensor_size = self.tensor_sizes[tensor_addr] - tensor_id = self.tensor_ids[tensor_addr] - tensor = self.tensor_map[self.tensor_ids[tensor_addr]] - - assert tensor_size == tensor.numel(), \ - f"Size mismatch. {tensor_size} is allocated at addr {tensor_addr} but tensor size is {tensor.numel()} " - - assert empty_addr != tensor_addr, \ - f"Cannot have same empty address {empty_addr} and tensor address {tensor_addr}" - - if empty_addr < tensor_addr: - - if empty_size >= tensor_size: - dest_buffer = self.buffer.narrow(0, empty_addr, tensor_size) - src_buffer = self.buffer.narrow(0, tensor_addr, tensor_size) - dest_buffer.data.copy_(src_buffer.data) - else: - - #print_rank_0(f'empty addr : {empty_addr}, empty size {empty_size} tensor addr {tensor_addr} tensor size {tensor_size}') - src_addr = tensor_addr - dest_addr = empty_addr - while src_addr < (tensor_addr + tensor_size): - copy_size = min(empty_size, tensor_addr + tensor_size - src_addr) - - dest_buffer = self.buffer.narrow(0, dest_addr, copy_size) - src_buffer = self.buffer.narrow(0, src_addr, copy_size) - - dest_buffer.data.copy_(src_buffer.data) - - src_addr += copy_size - dest_addr += copy_size - - self._replace_old_address_with_new(tensor_id, empty_addr) - - tensor_index += 1 - - else: - tensor_index += 1 - - empty_addresses = sorted(self.contiguous_sizes.keys()) - - def _replace_old_address_with_new(self, tensor_id, new_address): - - tensor = self.tensor_map[tensor_id] - tensor_size = tensor.numel() - tensor.data = self.buffer.narrow(0, new_address, tensor_size).data - - self._release_tensor(tensor_id) - self._mark_as_occupied(new_address, tensor_size) - - self.tensor_ids[new_address] = tensor_id - self.tensor_map[tensor_id] = tensor - self.tensor_addresses[tensor_id] = new_address - self.tensor_sizes[new_address] = tensor_size - - def _get_new_tensor_address(self, size): - tensor_address = None - for address, contiguous_size in self.contiguous_sizes.items(): - if contiguous_size >= size and \ - (tensor_address is None or \ - contiguous_size < self.contiguous_sizes[tensor_address]): - tensor_address = address - assert tensor_address is not None, "address cannot be None" - return tensor_address - - def _get_new_tensor(self, address, size): - available_contiguous_size = self.contiguous_sizes[address] - - assert size <= available_contiguous_size, \ - f"Tensor numel {size} is large than available contiguous size {available_contiguous_size}" - self.count += 1 - new_tensor = self.buffer.narrow(0, address, size) - tensor_id = id(new_tensor) - self.tensor_addresses[tensor_id] = address - self.tensor_sizes[address] = size - - self.tensor_ids[address] = tensor_id - self.tensor_map[tensor_id] = new_tensor - - self._mark_as_occupied(address, size) - - return new_tensor - - def _largest_contiguous(self): - if len(self.contiguous_sizes) > 0: - return max([size for _, size in self.contiguous_sizes.items()]) - else: - return 0 - - def _mark_as_occupied(self, address, size): - available_contiguous_size = self.contiguous_sizes[address] - del self.contiguous_sizes[address] - - if available_contiguous_size != size: - self.contiguous_sizes[address + size] = available_contiguous_size - size - - self.largest_contiguous = self._largest_contiguous() diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/linear.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/linear.py deleted file mode 100644 index 8c8db60768ebca9f7f972c13e0d3d20176158c96..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/linear.py +++ /dev/null @@ -1,185 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -#Linear Module to use with ZeRO Stage 3 to allow for parameter memory release -#after the module execution during forward -#Instead of saving variables using save_for_backward, we save variable ids -#Allowing us to retrieve the variable without creating pointer to it -#Which allows for underlying tensor to be garbage collected -#When partitioned as needed by the Zero Stage 3 optimizer -#TODO instead of patching Linear module, we could patch the ctx.save_for_backward -#ctx.saved_tensors so that this approach works for all nn modules that are built upon -#torch.nn.function. However the issue is that many modules uses C++ implementations -#which does not have pytorch implementation. Eg torch.addmm which acts as a functional -#when implemented outside of torch.autograd.Function - -import math -import functools - -import torch -from torch import Tensor -from torch.nn.parameter import Parameter -from torch.nn import init -from torch.nn.modules.module import Module -from deepspeed.runtime.utils import noop_decorator -from deepspeed import comm as dist -from deepspeed.accelerator import get_accelerator - - -def print_rank_0(message, debug=False, force=False): - if dist.get_rank() == 0 and (debug or force): - print(message) - - -try: - # Fix `torch.[device].amp.custom_fwd/bwd` FutureWarning in torch 2.4 - if hasattr(torch, 'amp') and hasattr(torch.amp, 'custom_fwd') and hasattr(torch.amp, 'custom_bwd'): - autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=get_accelerator().device_name()) - autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=get_accelerator().device_name()) - else: - # original implementation - autocast_custom_fwd = get_accelerator().amp().custom_fwd - autocast_custom_bwd = get_accelerator().amp().custom_bwd -except (ImportError, AttributeError) as exp: - autocast_custom_fwd = noop_decorator - autocast_custom_bwd = noop_decorator - - -class LinearFunctionForZeroStage3(torch.autograd.Function): - - # Note that both forward and backward are @staticmethods - @staticmethod - @autocast_custom_fwd - # bias is an optional argument - def forward(ctx, input, weight, bias=None): - - ctx.save_for_backward(input, weight, bias) - - if input.dim() == 2 and bias is not None: - # fused op is marginally faster - ret = torch.addmm(bias, input, weight.t()) - else: - output = input.matmul(weight.t()) - if bias is not None: - output += bias - ret = output - - return ret - - # This function has only a single output, so it gets only one gradient - @staticmethod - @autocast_custom_bwd - def backward(ctx, grad_output): - # This is a pattern that is very convenient - at the top of backward - # unpack saved_tensors and initialize all gradients w.r.t. inputs to - # None. Thanks to the fact that additional trailing Nones are - # ignored, the return statement is simple even when the function has - # optional inputs. - input, weight, bias = ctx.saved_tensors - - grad_input = grad_weight = grad_bias = None - - #print(f"backward shaped grad_output {grad_output.shape}, input {input.shape}, weight {weight.shape} and bias {bias.shape if bias is not None else None}") - # These needs_input_grad checks are optional and there only to - # improve efficiency. If you want to make your code simpler, you can - # skip them. Returning gradients for inputs that don't require it is - # not an error. - if ctx.needs_input_grad[0]: - #print(f"Computing grad input weight {weight.shape} grad_output {grad_output.shape}") - grad_input = grad_output.matmul(weight) - #print(f"Computed grad input {grad_input.shape}") - if ctx.needs_input_grad[1]: - #print("Computing grad weight") - dim = grad_output.dim() - if dim > 2: - grad_weight = grad_output.reshape(-1, - grad_output.shape[-1]).t().matmul(input.reshape(-1, input.shape[-1])) - else: - grad_weight = grad_output.t().matmul(input) - #print(f"Computed grad weight grad_weight {grad_weight.shape}") - if bias is not None and ctx.needs_input_grad[2]: - #print("Computing grad bias") - if dim > 2: - grad_bias = grad_output.sum([i for i in range(dim - 1)]) - else: - grad_bias = grad_output.sum(0) - #print("Done computing grad bias") - #print("needs bias") - #print(f"backward shaped grad_input {grad_input.shape}, grad_weight {grad_weight.shape}, grad_bias {grad_bias.shape if grad_bias is not None else None}") - return grad_input, grad_weight, grad_bias - - -def zero3_linear_wrap(input, weight, bias=None): - if bias is None: - return LinearFunctionForZeroStage3.apply(input, weight) - else: - return LinearFunctionForZeroStage3.apply(input, weight, bias) - - -class LinearModuleForZeroStage3(Module): - r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b`. - The weights are pre-transposed and stored as A^T instead of transposing during each - forward. Memory savings proportional to the parameter size. - - Args: - in_features: size of each input sample - out_features: size of each output sample - bias: If set to ``False``, the layer will not learn an additive bias. - Default: ``True`` - - Shape: - - Input: :math:`(N, *, H_{in})` where :math:`*` means any number of - additional dimensions and :math:`H_{in} = \text{in\_features}` - - Output: :math:`(N, *, H_{out})` where all but the last dimension - are the same shape as the input and :math:`H_{out} = \text{out\_features}`. - - Attributes: - weight: the learnable weights of the module of shape - :math:`(\text{out\_features}, \text{in\_features})`. The values are - initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where - :math:`k = \frac{1}{\text{in\_features}}` - bias: the learnable bias of the module of shape :math:`(\text{out\_features})`. - If :attr:`bias` is ``True``, the values are initialized from - :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where - :math:`k = \frac{1}{\text{in\_features}}` - - Examples:: - - >>> m = nn.Linear(20, 30) - >>> input = torch.randn(128, 20) - >>> output = m(input) - >>> print(output.size()) - torch.Size([128, 30]) - """ - __constants__ = ['in_features', 'out_features'] - in_features: int - out_features: int - weight: Tensor - - def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None: - super(LinearModuleForZeroStage3, self).__init__() - print("Building ZeRO module") - self.in_features = in_features - self.out_features = out_features - self.weight = Parameter(torch.Tensor(out_features, in_features)) - if bias: - self.bias = Parameter(torch.Tensor(out_features)) - else: - self.register_parameter('bias', None) - self.reset_parameters() - - def reset_parameters(self) -> None: - init.kaiming_uniform_(self.weight, a=math.sqrt(5)) - if self.bias is not None: - fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight) - bound = 1 / math.sqrt(fan_in) - init.uniform_(self.bias, -bound, bound) - - def forward(self, input: Tensor) -> Tensor: - return LinearFunctionForZeroStage3.apply(input, self.weight, self.bias) - - def extra_repr(self) -> str: - return 'in_features={}, out_features={}, bias={}'.format(self.in_features, self.out_features, self.bias - is not None) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/mics.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/mics.py deleted file mode 100644 index 628bf86a61da23182b2f33cc7b6edab2c1aee41f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/mics.py +++ /dev/null @@ -1,472 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -# SPDX-License-Identifier: Apache-2.0 - -import sys -from typing import List - -import deepspeed -import torch -from deepspeed import comm as dist -from deepspeed.runtime.zero.utils import is_zero_param -from deepspeed.runtime.zero.mics_utils import (MiCS_CommGroups, create_mics_comm_groups, scale_tensors) -from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload -from deepspeed.runtime.zero.partition_parameters import Init, AllGatherCoalescedHandle, ZeroParamStatus -from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 -from deepspeed.utils import instrument_w_nvtx, log_dist, logger -from deepspeed.accelerator import get_accelerator -from torch import Tensor -from torch.nn import Parameter - - -def has_hierarchical_all_gather_groups(comm_groups: MiCS_CommGroups): - result = False - if comm_groups.param_intra_node_group is not None and comm_groups.param_inter_node_shard_group is not None: - result = True - return result - - -class MiCS_AllGatherCoalescedHandle(AllGatherCoalescedHandle): - """ This handle assumes that no need to - copy data out from a contiguous tensor - """ - - def __init__(self, allgather_handle, params: List[Parameter], partitions: List[Tensor], world_size: int) -> None: - super().__init__(allgather_handle, params, partitions, world_size) - - def wait(self, **kwargs) -> None: - """ - """ - # let the current stream to op - try: - # print("HANDLE", self.allgather_handle) - instrument_w_nvtx(self.allgather_handle.wait)() - except (ValueError, RuntimeError) as e: - log_dist( - f"WARNING: Runtime Error while waiting the collective all-gather, possibly due to the _IllegalWork", - ranks=[0]) - log_dist(f"Error message: {e}", ranks=[0]) - - if self.complete: - return - - for _, param in enumerate(self.params): - assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight" - param.ds_status = ZeroParamStatus.AVAILABLE - - self.complete = True - - -class MiCS_Init(Init): - - def __init__(self, - module=None, - data_parallel_group=None, - sequence_data_parallel_group=None, - mem_efficient_linear=True, - remote_device=None, - pin_memory=False, - config_dict_or_path=None, - config=None, - enabled=True, - dtype=None, - mpu=None): - """A context manager to partition the model parameters during the model - construction with MiCS partition strategy. Model states are partitioned - to the number of devices specified via ``mics_shard_size`` field in the - deepspeed config json file. The context manager also introduces - hierarchical communication method to reduce the cost of inter-node - communications, which can be enabled with - ``mics_hierarchical_params_gather`` field in deepspeed config. - - Args: - module (``torch.nn.Module``, optional): If provided, partition the model as - if it was constructed in the context. - data_parallel_group (``deepspeed.comm`` process group, optional): - The group of processes to partition among. Defaults to all processes. - Synonymous with sequence data parallel group for param partitioning - across both sequence and data parallel groups. - mem_efficient_linear (bool, optional): Replace - torch.nn.functional.linear with an implementation that allows - DeepSpeed to partition parameters. Defaults to ``True``. - remote_device (string, optional): The initial device to store model - weights e.g., ``cpu``, ``nvme``. Passing ``"cpu"`` will create the model in CPU - memory. The model may still be moved to GPU based on the - offload settings for training. Defaults to param offload device if a config is - defined, otherwise GPU. - pin_memory (bool, optional): Potentially increase performance by - using pinned memory for model weights. ``remote_device`` must be - ``"cpu"``. Defaults to pin_memory value in config, otherwise ``False``. - config_dict_or_path (dict or ``json file``, optional): If provided, provides configuration - for swapping fp16 params to NVMe. - config (dict or ``json file``, optional): Deprecated, use config_dict_or_path instead. - enabled (bool, optional): If ``False``, this context has no - effect. Defaults to ``True``. - dtype (``dtype``, optional): Can be used to change the data type of the parameters. - Supported options are ``torch.half`` and ``torch.float``. Defaults to ``None`` - mpu (``object``, optional): A model parallelism unit object that implements get_{model,data}_parallel_{rank,group,world_size}. - - This context follows the same logic as ``deepspeed.zero.Init()``, but - with the modification for partition size of each parameter. - - Examples - -------- - - #. Allocate a model and partition it among all processes: - - .. code-block:: python - # the config_dict_or_path is required to let the context manager know - # how partition the parameters. - # The configuration has to include the field ``mics_shard_size`` - with deepspeed.zero.MiCS_Init(config_dict_or_path=ds_config): - model = MyLargeModel() - - - #. Allocate a model in pinned CPU memory and partition it among a subgroup of processes: - - .. code-block:: python - - with deepspeed.zero.MiCS_Init(data_parallel_group=mpu.get_data_parallel_group(), - remote_device="cpu", - pin_memory=True - config_dict_or_path=ds_config): - model = MyLargeModel() - - - #. Partition an already-allocated model in CPU memory: - - .. code-block:: python - - model = deepspeed.zero.MiCS_Init(module=model, - config_dict_or_path=ds_config) - """ - - assert config_dict_or_path is not None, "Must provide configuration for MiCS Initialization" - _ds_config = deepspeed.runtime.config.DeepSpeedConfig(config_dict_or_path, mpu) - if not dist.is_initialized(): - dist.init_distributed() - assert dist.is_initialized(), "Parameters cannot be scattered without initializing deepspeed.comm" - - if data_parallel_group is None: - ds_process_group = dist.get_world_group() - else: - ds_process_group = data_parallel_group - - if sequence_data_parallel_group is not None: - logger.warning( - f"sequence_data_parallel_group' is deprecated and will be removed. Use 'data_parallel_group' instead.") - if data_parallel_group is not None: - raise ValueError( - "Both 'data_parallel_group' and 'sequence_data_parallel_group' were specified. Please provide only one of these arguments." - ) - self.ds_process_group = sequence_data_parallel_group - - self.mics_comm_groups = create_mics_comm_groups( - _ds_config.mics_shard_size, - ds_process_group, - hierarchical_allgather=_ds_config.mics_hierarchial_params_gather, - mpu=mpu) - - super().__init__(module, data_parallel_group, mem_efficient_linear, remote_device, pin_memory, - config_dict_or_path, config, enabled, dtype, mpu) - - def _convert_to_deepspeed_param(self, param): - super()._convert_to_deepspeed_param(param) - # attach communication groups to every param - param.comm = self.mics_comm_groups - - # record existing all_gather_coalesced implementation - # so that we can fallback later - old_all_gather_coalesced = param.all_gather_coalesced - - def _param_all_gather_coalesced(params, param_buffers=None, **kwargs): - """""" - mics_comm_groups: MiCS_CommGroups = params[0].comm - hierarchical_all_gather = has_hierarchical_all_gather_groups(mics_comm_groups) - if dist.has_coalescing_manager() and hierarchical_all_gather: - return self._hierarchical_all_gather_params(params, param_buffers) - elif dist.has_coalescing_manager(): - return self._flat_all_gather_with_coalescing_manager(params, param_buffers) - else: - return old_all_gather_coalesced(params, **kwargs) - - # change the all_gather_coalesced method - param.all_gather_coalesced = _param_all_gather_coalesced - - def _pre_all_gather(self, params, params_buffers=None): - # fetches from nvme if the partition is not available and in nvme - self._ensure_availability_of_partitioned_params(params) - - for param in params: - if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: - raise RuntimeError(param.ds_summary()) - param.ds_status = ZeroParamStatus.INFLIGHT - - # ensure that each rank has params in same order. the allgather - # is done by flattening the parameter list into a single tensor that - # can be allgathered in a single call - this means that if each rank - # gives a list of the same parameters in a different order we will - # silently get incorrect parameter values, and have very difficult - # to debug correctness issues. - params = sorted(params, key=lambda p: p.ds_id) - return params, params_buffers - - def _flat_all_gather_with_coalescing_manager(self, params, params_buffers=None): - """""" - # must have to change the status of the param - # and ensure they are on the device - params, params_buffers = self._pre_all_gather(params, params_buffers) - - mics_comm_groups: MiCS_CommGroups = params[0].comm - param_shard_size = mics_comm_groups.param_shard_size - - output_tensors = [] - input_tensors = [] - for i, p in enumerate(params): - t_size = p.ds_tensor.ds_numel * param_shard_size - if params_buffers is not None and params_buffers[i] is not None: - assert params_buffers[i].numel( - ) == t_size, f'params_to_gather_buffers[{i}] size {params_buffers[i].numel()} does not match with t_size {t_size}' - flat_out = params_buffers[i] - else: - flat_out = torch.empty(t_size, dtype=p.dtype, device=self.local_device, requires_grad=False).view(-1) - output_tensors.append(flat_out) - _flat_input = p.ds_tensor.data.view(-1) - input_tensors.append(_flat_input) - - all_gather_handle = dist.all_gather_coalesced(output_tensors, - input_tensors, - group=mics_comm_groups.param_shard_group, - async_op=True) - - for idx, param in enumerate(params): - param.data = output_tensors[idx].narrow(0, 0, param.ds_numel).view(param.ds_shape).data - - return MiCS_AllGatherCoalescedHandle(allgather_handle=all_gather_handle, - params=params, - partitions=[], - world_size=param_shard_size) - - def _hierarchical_all_gather_params(self, params, params_buffers=None): - """""" - params, params_buffers = self._pre_all_gather(params, params_buffers) - - mics_comm_groups: MiCS_CommGroups = params[0].comm - local_rank = dist.get_rank(group=mics_comm_groups.param_intra_node_group) - inter_node_comm_group = mics_comm_groups.param_inter_node_shard_group - intra_node_comm_group = mics_comm_groups.param_intra_node_group - param_shard_size = mics_comm_groups.param_shard_size - - inter_node_size = dist.get_world_size(group=inter_node_comm_group) - intra_node_size = dist.get_world_size(group=intra_node_comm_group) - param_tensors = [] - for i, p in enumerate(params): - param_size = p.ds_tensor.ds_numel * param_shard_size - if params_buffers is not None and params_buffers[i] is not None: - assert params_buffers[i].numel( - ) == param_size, f'param_buffers[{i}] size {params_buffers[i].numel()} does not match with param_size {param_size}' - param_tensor = params_buffers[i] - else: - param_tensor = torch.empty(param_size, dtype=p.dtype, device=self.local_device, - requires_grad=False).view(-1) - param_tensors.append(param_tensor) - - # inter node all-gather - inter_outputs = [] - inter_inputs = [] - for i, p in enumerate(params): - inter_size = p.ds_tensor.ds_numel * inter_node_size - _out = param_tensors[i].narrow(0, local_rank * inter_size, inter_size) - inter_outputs.append(_out) - inter_inputs.append(p.ds_tensor.data.view(-1).to(self.local_device)) - # sync enqueue - dist.all_gather_coalesced(inter_outputs, inter_inputs, group=inter_node_comm_group, async_op=False) - - # intra node all-gather - intra_outputs = [] - intra_inputs = [] - for i, p in enumerate(params): - # partition param into multiple chunks for allgather - # because inter-node all-gather outputs are in a continues memory - # while in param memory, those inter-node data are placed in different - # location. - # each chunk is an intra-node output - param_chunk = param_tensors[i].view( - (inter_node_size, intra_node_size, p.ds_tensor.ds_numel)).narrow(1, local_rank, 1) - param_chunk.copy_(inter_outputs[i].detach().clone().view(param_chunk.size())) - output_chunks = torch.chunk(param_tensors[i], inter_node_size) - for j, _out in enumerate(output_chunks): - intra_chunk_size = intra_node_size * p.ds_tensor.ds_numel - local_offset = local_rank * p.ds_tensor.ds_numel - _in = param_tensors[i].narrow(0, j * intra_chunk_size + local_offset, p.ds_tensor.ds_numel) - intra_outputs.append(_out) - intra_inputs.append(_in) - - all_gather_handle = dist.all_gather_coalesced(intra_outputs, - intra_inputs, - group=intra_node_comm_group, - async_op=True) - for i, param in enumerate(params): - param.data = param_tensors[i].narrow(0, 0, param.ds_numel).view(param.ds_shape).data - - return MiCS_AllGatherCoalescedHandle( - allgather_handle=all_gather_handle, - params=params, - partitions=[], - world_size=param_shard_size, - ) - - def get_partition_dp_group(self, param): - return param.comm.param_shard_group - - def get_partition_rank(self): - return self.mics_comm_groups.param_shard_rank - - @property - def num_partitions(self): - return self.mics_comm_groups.param_shard_size - - -class MiCS_Offload(DeepSpeedZeRoOffload): - """ Wrapper to change the behavior for parameter sharding - """ - - def _convert_to_zero_parameters(self, ds_config, module, mpu): - """ overload the parent class function for convert the parameters - - """ - log_dist(f'Convert to zero parameters from MiCS Offload manager', ranks=[0]) - non_zero_params = [p for p in module.parameters() if not is_zero_param(p)] - if non_zero_params: - zero_params = [p for p in module.parameters() if is_zero_param(p)] - if zero_params: - zero_params[0].convert_to_zero_parameters(param_list=non_zero_params) - else: - group = None - if mpu: - group = mpu.get_data_parallel_group() - - MiCS_Init(module=module, - data_parallel_group=group, - dtype=self.dtype, - config_dict_or_path=ds_config, - remote_device=self.offload_device, - pin_memory=self.offload_param_pin_memory, - mpu=mpu) - - -class MiCS_Optimizer(DeepSpeedZeroOptimizer_Stage3): - """ - MiCS Optimizer - """ - - def __init__(self, - module, - init_optimizer, - timers, - ds_config, - static_loss_scale=1, - dynamic_loss_scale=False, - dynamic_loss_args=None, - verbose=True, - contiguous_gradients=True, - reduce_bucket_size=500000000, - prefetch_bucket_size=50000000, - max_reuse_distance=1000000000, - max_live_parameters=1000000000, - param_persistence_threshold=100000, - model_persistence_threshold=sys.maxsize, - dp_process_group=None, - reduce_scatter=True, - overlap_comm=False, - offload_optimizer_config=None, - offload_param_config=None, - sub_group_size=1000000000000, - offload_ratio=0.0, - mpu=None, - clip_grad=0, - gradient_accumulation_dtype=torch.float16, - communication_data_type=torch.float16, - postscale_gradients=True, - gradient_predivide_factor=1, - gradient_accumulation_steps=1, - elastic_checkpoint=False, - aio_config=None): - - log_dist("Init MiCS optimizer", ranks=[0]) - super().__init__(module, init_optimizer, timers, ds_config, static_loss_scale, dynamic_loss_scale, - dynamic_loss_args, verbose, contiguous_gradients, reduce_bucket_size, prefetch_bucket_size, - max_reuse_distance, max_live_parameters, param_persistence_threshold, - model_persistence_threshold, dp_process_group, reduce_scatter, overlap_comm, - offload_optimizer_config, offload_param_config, sub_group_size, offload_ratio, mpu, clip_grad, - gradient_accumulation_dtype, communication_data_type, postscale_gradients, - gradient_predivide_factor, gradient_accumulation_steps, elastic_checkpoint, aio_config) - first_param = next(module.parameters()) - # overload the dp_process_group and partition_count - assert hasattr(first_param, "comm"), " ".join([ - "Sharded parameters don't have the MiCS_CommGroups attached.", - "Might due to the use of deepspeed.zero.Init context for initializing the weights.", - "To use MiCS sharding, please use deepspeed.zero.MiCS_Init instead for initializing parameter." - ]) - self.dp_process_group = first_param.comm.param_shard_group - self.partition_count = first_param.comm.param_shard_size - - def initialize_ds_offload( - self, - *args, - **kwargs, - ): - return MiCS_Offload(*args, **kwargs) - - def partition_grads(self, params_to_release: List[Parameter], grad_partitions: List[Tensor]) -> None: - grad_buffers = super().partition_grads(params_to_release, grad_partitions) - # perform all-reduce among replication groups - # the function will perform accumulation boundary check - self.allreduce_mics_shard_grads(params_to_release, grad_buffers) - - @instrument_w_nvtx - def allreduce_mics_shard_grads(self, params, partitioned_grads_buffers: List[Tensor]): - """ - """ - # TODO: improve the condition check - if not self.is_gradient_accumulation_boundary or \ - len(partitioned_grads_buffers) == 0: - return - - mics_comm_groups: MiCS_CommGroups = params[0].comm - param_repli_group = mics_comm_groups.param_repli_group - param_repli_size = mics_comm_groups.param_repli_size - - if param_repli_size is None or param_repli_size <= 1: - return - if not get_accelerator().on_accelerator(partitioned_grads_buffers[0]): - raise RuntimeError("Local sharding has no support for CPU offloading") - - if dist.has_all_reduce_coalesced(): - scale_tensors(partitioned_grads_buffers, param_repli_size) - dist.all_reduce_coalesced(tensors=partitioned_grads_buffers, group=param_repli_group) - else: - # manually coalescing all-reduce - aggregated_buffer: Tensor = torch.cat(partitioned_grads_buffers) - aggregated_buffer.div_(param_repli_size) - dist.all_reduce(aggregated_buffer, group=param_repli_group) - offset = 0 - for grad_buff in partitioned_grads_buffers: - grad_buff.view(-1).copy_(aggregated_buffer.narrow(0, offset, grad_buff.numel())) - offset += grad_buff.numel() - - def load_state_dict(self, - state_dict_list, - load_optimizer_states=True, - load_from_fp32_weights=False, - checkpoint_folder=None, - load_serial=None): - r""" Loading the ZeRO-3/MiCS partitioned checkpoints - Because the self.dp_process_group is replaced with the communicator for - partition group we can call the load_state_dict logic from ZeRO-3. - """ - super().load_state_dict(state_dict_list, load_optimizer_states, load_from_fp32_weights, checkpoint_folder) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/mics_utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/mics_utils.py deleted file mode 100644 index 06b83160bd6c4cd93f5250a12c67f5abda01d920..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/mics_utils.py +++ /dev/null @@ -1,203 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -# SPDX-License-Identifier: Apache-2.0 - -import os -from dataclasses import dataclass -from typing import List - -import numpy as np -import torch -from torch import Tensor - -from deepspeed import comm as dist -from deepspeed.accelerator import get_accelerator -from deepspeed.utils import logger - - -def _log_rank0(msg): - if dist.get_rank() == 0: - logger.info(msg) - - -@torch.jit.script -def scale_tensors(tensors: List[Tensor], scale: int): - for t in tensors: - t.div_(scale) - - -@dataclass -class MiCS_CommGroups: - """""" - param_shard_group = None - param_shard_size = -1 - param_shard_rank = -1 - - param_repli_group = None - param_repli_size = -1 - param_repli_rank = -1 - - param_intra_node_group = None - param_inter_node_shard_group = None - - -def create_mics_comm_groups( - shard_size, - dp_group, - hierarchical_allgather=False, - mpu=None, -): - """ - create shard-group, replicate-group from config_file - TODO: consider broadcast the config from rank0 - - Returns: - MiCS_CommGroups - """ - # env var for debugging purpose - ndevices_per_node = int(os.environ.get("NDEV_PER_NODE", get_accelerator().device_count())) - _log_rank0(f'creating MiCS communication groups with per node device size {ndevices_per_node}') - groups = MiCS_CommGroups() - - if mpu is not None: - assert dp_group == mpu.get_data_parallel_group() - - # full size of the world - world_size = dist.get_world_size() - # global rank - global_rank = dist.get_rank() - - config = _generate_mics_config(world_size, ndevices_per_node, shard_size, 1) - ranks_of_shard_group = config['shard_groups'] - ranks_of_repli_group = config['replicate_groups'] - if len(ranks_of_repli_group) == 0: - assert len(ranks_of_shard_group) == 1, "replicate groups are empty only for single shard group" - for r in ranks_of_shard_group[0]: - ranks_of_repli_group.append([r]) - - # for simplicity - assert _sizes_all_same(ranks_of_repli_group), "replicate groups must have the same size" - assert _sizes_all_same(ranks_of_shard_group), "shard groups must have the same size" - - assert sum([len(g) for g in ranks_of_shard_group]) == dist.get_world_size(), "all sharded ranks " - if len(ranks_of_shard_group) > 1: # if only shard on one group then no need for replicate groups - assert len(ranks_of_shard_group) == len( - ranks_of_repli_group[0]), "number of shard groups must equal to the size of each replicate group" - - global_rank = dist.get_rank() - # create shard groups - for shard_ranks in ranks_of_shard_group: - _group = dist.new_group(shard_ranks) - if global_rank in shard_ranks: - groups.param_shard_group = _group - groups.param_shard_size = len(shard_ranks) - groups.param_shard_rank = dist.get_rank(_group) - logger.info(f'rank {global_rank}, shard group' - f' {groups.param_shard_rank}/{dist.get_world_size(group=_group)}') - - # create replicate groups - for repli_ranks in ranks_of_repli_group: - if len(repli_ranks) > 1: - _group = dist.new_group(repli_ranks) - if global_rank in repli_ranks: - groups.param_repli_group = _group - groups.param_repli_size = len(repli_ranks) - groups.param_repli_rank = dist.get_rank(group=_group) - logger.info(f'rank {global_rank} ' - f'replicate group {groups.param_repli_rank}/{dist.get_world_size(group=_group)}') - else: - groups.param_repli_group = None - groups.param_repli_size = 1 - groups.param_repli_rank = 0 - logger.info(f'rank {global_rank} replicate group 0/1') - - # assign shard group size as world size - assert groups.param_shard_size == len(ranks_of_shard_group[0]) - - if hierarchical_allgather: - # create hierarchy inter-node, intra-node groups - # n_span_nodes = config['shard_span'] - n_span_nodes = config['span_nodes'] - assert n_span_nodes > 1, "sharding spans on single node, no need for hierarchy allgather" - assert len(ranks_of_shard_group[0]) % n_span_nodes == 0 - - n_gpu_per_node = len(ranks_of_shard_group[0]) // n_span_nodes - intra_node_ranks_group = [] - inter_node_ranks_group = [] - for shard_group in ranks_of_shard_group: - _intra_node_ranks = [] - for i in range(0, len(shard_group), n_gpu_per_node): - _intra_node_ranks.append(shard_group[i:i + n_gpu_per_node]) - _inter_node_ranks = [] - for i in range(n_gpu_per_node): - _ranks = [_g[i] for _g in _intra_node_ranks] - _inter_node_ranks.append(_ranks) - - intra_node_ranks_group.append(_intra_node_ranks) - inter_node_ranks_group.append(_inter_node_ranks) - - _log_rank0(f"create for hierarchy all-gather groups: intra nodes {intra_node_ranks_group}") - _log_rank0(f"create for hierarchy all-gather groups: inter nodes {inter_node_ranks_group}") - - # create communicators - for shard_group in intra_node_ranks_group: - for intra_node_ranks in shard_group: - _group = dist.new_group(intra_node_ranks) - if global_rank in intra_node_ranks: - groups.param_intra_node_group = _group - _log_rank0(f'create group for intra node ranks {intra_node_ranks}') - - for shard_group in inter_node_ranks_group: - for inter_node_ranks in shard_group: - _group = dist.new_group(inter_node_ranks) - if global_rank in inter_node_ranks: - groups.param_inter_node_shard_group = _group - _log_rank0(f'create group for inter node ranks {inter_node_ranks}') - return groups - - -def _generate_mics_config(world_size, ndev_per_node, shard_size, pp_size=1): - """Generating the configuration for sharding This shard config generation assume - that the pipeline stages are partitioned in order, i.e., first ranks - hold the stage0, etc. - - Args: - - shard_size (int): zero3 data-parallel shard size, FIXME: - change the name later - - pp_size (int): pipeline parallel size, currently, only work with - pipeline parallelism + zero - - """ - assert world_size % pp_size == 0 - assert (world_size // pp_size) % shard_size == 0, \ - f"dp group size is not dividable by dp_shard_size, "\ - f" (world_size {world_size}, pp_size {pp_size}, dp_shard_size {shard_size})" - - config = {} - shard_groups = np.arange(world_size).reshape(-1, shard_size) - replicate_groups = [] - for i in range(shard_size): - same_shard_ranks = shard_groups[:, i].tolist() - n_ranks = len(same_shard_ranks) - replicate_size = n_ranks // pp_size - replicate_groups.extend([same_shard_ranks[j:j + replicate_size] for j in range(0, n_ranks, replicate_size)]) - - config['replicate_groups'] = replicate_groups - config['shard_groups'] = shard_groups.tolist() - config["span_nodes"] = len(shard_groups[0]) // ndev_per_node - return config - - -def _sizes_all_same(groups): - """all groups have same length""" - all_same = True - for g in groups: - if len(g) != len(groups[0]): - return False - return all_same diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/offload_config.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/offload_config.py deleted file mode 100644 index ca35d7a7d1691c8ec4e6a7016cdb22b31ef459ac..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/offload_config.py +++ /dev/null @@ -1,109 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from enum import Enum -from pathlib import Path -from pydantic import Field, model_validator -from typing import Optional - -from deepspeed.runtime.config_utils import DeepSpeedConfigModel, pp_int - - -class OffloadDeviceEnum(str, Enum): - """ Enum for valid offload devices """ - none = "none" - cpu = "cpu" - nvme = "nvme" - - -class DeepSpeedZeroOffloadParamConfig(DeepSpeedConfigModel): - """ Set options for parameter offload. Valid only with stage 3. """ - - device: OffloadDeviceEnum = "none" - """ - Device memory to offload model parameters. Supported options are `cpu` and - `nvme`. - """ - - nvme_path: Optional[Path] = None - """ Filesystem path for NVMe device for parameter offloading. """ - - buffer_count: int = Field(5, ge=0) - """ Number of buffers in buffer pool for parameter offloading to NVMe. """ - - buffer_size: int = Field(pp_int(1e8), ge=0) - """ Size of buffers in buffer pool for parameter offloading to NVMe. """ - - max_in_cpu: int = Field(pp_int(1e9), ge=0) - """ - Number of parameter elements to maintain in CPU memory when offloading to - NVMe is enabled. - """ - - pin_memory: bool = False - """ - Offload to page-locked CPU memory. This could boost throughput at the cost - of extra memory overhead. - """ - - -class DeepSpeedZeroOffloadOptimizerConfig(DeepSpeedConfigModel): - """ Set options for optimizer offload. Valid with stage 1, 2, and 3. """ - - device: OffloadDeviceEnum = "none" - """ - Device memory to offload optimizer state. Supported options are `cpu` and - `nvme`. Optimizer computation is offload to CPU regardless of device option. - """ - - nvme_path: Optional[Path] = None - """ Filesystem path for NVMe device for optimizer state offloading. """ - - buffer_count: int = Field(4, ge=0) - """ - Number of buffers in buffer pool for optimizer state offloading to NVMe. - This should be at least the number of states maintained per parameter by - the optimizer. For example, Adam optimizer has 4 states (parameter, - gradient, momentum, and variance). - """ - - pin_memory: bool = False - """ - Offload to page-locked CPU memory. This could boost throughput at the cost - of extra memory overhead. - """ - - pipeline_read: bool = False - """ - For tile-based optimizer step processing, overlap read of next tile with - computation of current tile. Used in ZeRO-Infinity. - """ - - pipeline_write: bool = False - """ - For tile-based optimizer step processing, overlap write of previous tile - with computation of current tile. - """ - - fast_init: bool = False - """ Enable fast optimizer initialization when offloading to NVMe. """ - - ratio: float = Field(1.0, ge=0.0, le=1.0) - """ Percentage of offloaded optimizer states to CPU Adam. Only valid with ZeRO Stage 3.""" - - @model_validator(mode="after") - def set_pipeline(self): - pipeline = self.pipeline_read or self.pipeline_write - self.__dict__["pipeline"] = pipeline - return self - - -class OffloadStateTypeEnum(str, Enum): - """ Enum for internal buffer types """ - optim_states = "optim_states" - hp_params = "hp_params" - lp_params = "lp_params" - lp_grads = "lp_grads" - contiguous_grad_buffer = "contiguous_grad_buffer" diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/offload_states.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/offload_states.py deleted file mode 100644 index d2c9b87a2bcaf302e09055277a5e3e311ca6ae60..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/offload_states.py +++ /dev/null @@ -1,71 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Set -import torch - -from deepspeed.accelerator import get_accelerator -from deepspeed.runtime.zero.offload_config import OffloadStateTypeEnum - - -def _make_offload_state_key(key): - return f"{key}_offload_buffer" - - -def offload_adam_states(optimizer, device, pin_memory: bool = False, non_blocking: bool = False): - """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" - - def move_key(state, key): - offload_buf_key = _make_offload_state_key(key) - if offload_buf_key not in state: - state[offload_buf_key] = torch.empty_like(state[key], device=device) - if pin_memory: - state[offload_buf_key] = get_accelerator().pin_memory(state[offload_buf_key]) - state[offload_buf_key].copy_(state[key], non_blocking=non_blocking) - state[key].data = state[offload_buf_key] - - for _, state in optimizer.state.items(): - if "exp_avg" in state: - move_key(state, "exp_avg") - if "exp_avg_sq" in state: - move_key(state, "exp_avg_sq") - - -def reload_adam_states(optimizer, device, non_blocking: bool = False): - """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" - - def move_back_key(state, key): - state[key].data = state[_make_offload_state_key(key)].to(device, non_blocking=non_blocking) - - for _, state in optimizer.state.items(): - if "exp_avg" in state: - move_back_key(state, "exp_avg") - if "exp_avg_sq" in state: - move_back_key(state, "exp_avg_sq") - - -def get_state_devices(model, state: OffloadStateTypeEnum) -> Set[torch.device]: - """Retrieve the devices of the specified state of the model. - - Args: - model (DeepSpeedEngine): The model whose device allocations are to be checked. - state (OffloadStateTypeEnum): The specific state for which the devices should be retrieved. - - Returns: - Set[torch.device]: A set of devices of the specified state. - - """ - if state == OffloadStateTypeEnum.hp_params: - return set(model.optimizer.get_hp_param_device(p) for p in model.parameters()) - elif state == OffloadStateTypeEnum.lp_params: - return set(p.ds_tensor.device for p in model.parameters()) - elif state == OffloadStateTypeEnum.lp_grads: - return {model.optimizer.grad_partitions_flat_buffer.device} - elif state == OffloadStateTypeEnum.optim_states: - return set(model.optimizer.get_hp_param_device(p, "exp_avg") for p in model.parameters()) | \ - set(model.optimizer.get_hp_param_device(p, "exp_avg_sq") for p in model.parameters()) - elif state == OffloadStateTypeEnum.contiguous_grad_buffer: - return set(bucket.buffer.device for bucket in model.optimizer.ipg_buckets.values() - if bucket.buffer is not None) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py deleted file mode 100644 index 000d0ebde1c7fbfe9e225ebd999e8c1d1a5e185c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py +++ /dev/null @@ -1,589 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import sys -import torch -from collections import OrderedDict -from deepspeed.utils import z3_leaf_module, set_z3_leaf_module -from deepspeed.runtime.utils import see_memory_usage -from deepspeed.runtime.zero.utils import apply_to_tensors_only, is_zero_param -from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum -from deepspeed.runtime.zero.partition_parameters import _init_external_params -from deepspeed.runtime.zero.partition_parameters import * -from deepspeed.runtime.zero.partitioned_param_coordinator import PartitionedParameterCoordinator, InflightParamRegistry, iter_params -from deepspeed.accelerator import get_accelerator -from deepspeed import utils - -FWD_MODULE_STACK = list() - - -#for each tensor in outputs run the forward_function and register backward_function as hook -def _apply_forward_and_backward_to_tensors_only(module, forward_function, backward_function, outputs): - if type(outputs) is tuple: - touched_outputs = [] - for output in outputs: - touched_output = _apply_forward_and_backward_to_tensors_only(module, forward_function, backward_function, - output) - touched_outputs.append(touched_output) - return tuple(touched_outputs) - elif type(outputs) is torch.Tensor: - forward_function(outputs) - if outputs.requires_grad: - outputs.register_hook(backward_function) - return outputs - else: - return outputs - - -class ZeROOrderedDict(OrderedDict): - - def __init__(self, parent_module, *args, **kwargs): - """A replacement for ``collections.OrderedDict`` to detect external ZeRO params. - - Args: - parent_module (``collections.OrderedDict``): the collection to replace - """ - - super().__init__(*args, **kwargs) - self._parent_module = parent_module - self._in_forward = False - - def __reduce__(self): - r0, _, *r2 = super().__reduce__() - return (r0, (self._parent_module, )) + tuple(r2) - - def __getitem__(self, key): - param = super().__getitem__(key) - - # Params can be registered as None (e.g., bias) - if param is None: - return param - - # TODO: only weaken this check during compilation - if hasattr(param, "ds_status") and param.ds_status == ZeroParamStatus.NOT_AVAILABLE: - if self._parent_module._parameters._in_forward: - register_external_parameter(FWD_MODULE_STACK[-1], param) - param.all_gather() - print_rank_0(f'Registering external parameter from getter {key} ds_id = {param.ds_id}', force=False) - - return param - - -def _inject_parameters(module, cls): - for module in module.modules(): - module._original_parameters = module._parameters - - if cls == ZeROOrderedDict: - new_param = cls(parent_module=module) - else: - new_param = cls() - - for key, param in module._parameters.items(): - new_param[key] = param - - module._parameters = new_param - - -class DeepSpeedZeRoOffload(object): - - def __init__( - self, - module, - timers, - ds_config, - overlap_comm=True, - prefetch_bucket_size=50000000, - max_reuse_distance=1000000000, - max_live_parameters=1000000000, - param_persistence_threshold=100000, - model_persistence_threshold=sys.maxsize, - dp_process_group=None, - offload_param_config=None, - mpu=None, - zero_param_parallel_group=None, - zero_quantized_weights=False, - zero_quantized_nontrainable_weights=False, - zero_module_granularity_threshold=0, - log_trace_cache_warnings=False, - ): - - see_memory_usage("DeepSpeedZeRoOffload initialize [begin]", force=True) - - print_rank_0(f"initialized {__class__.__name__} with args: {locals()}", force=False) - - self.module = module - self.timers = timers - self.dtype = list(module.parameters())[0].dtype - self.dp_process_group = dp_process_group - self.offload_device = None - self.offload_param_pin_memory = False - self.zero_param_parallel_group = zero_param_parallel_group - self.zero_quantized_weights = zero_quantized_weights - self.zero_quantized_nontrainable_weights = zero_quantized_nontrainable_weights - self.log_trace_cache_warnings = log_trace_cache_warnings - - if offload_param_config is not None and offload_param_config.device != OffloadDeviceEnum.none: - self.offload_device = offload_param_config.device - self.offload_param_pin_memory = offload_param_config.pin_memory - - self._convert_to_zero_parameters(ds_config, module, mpu) - - for m in module.modules(): - _init_external_params(m) - - _inject_parameters(module, ZeROOrderedDict) - - self.param_numel_persistence_threshold = int(param_persistence_threshold) - self.model_persistence_threshold = int(model_persistence_threshold) - self.persistent_parameters = self.mark_persistent_parameters(self.param_numel_persistence_threshold, - self.model_persistence_threshold) - - self._prefetch_bucket_sz = int(prefetch_bucket_size) - self._max_reuse_distance_in_numel = int(max_reuse_distance) - self._max_available_parameters_in_numel = int(max_live_parameters) - self.__allgather_stream = None if get_accelerator().is_synchronized_device() else get_accelerator().Stream( - ) if overlap_comm else get_accelerator().default_stream() - - if not hasattr(module, "ds_inflight_param_registry"): - module.ds_inflight_param_registry = InflightParamRegistry() - self.__inflight_param_registry = module.ds_inflight_param_registry - - self.fast_sharding_for_leaf_module = False - - if zero_module_granularity_threshold > 0: - self.min_granularity_value = sys.maxsize - self.min_granularity_layer = None - self.granularity_info = set() - self.z3_leaf_layers = [] - self._set_z3_leaf_modules_by_threshold(module, zero_module_granularity_threshold) - self.fast_sharding_for_leaf_module = True - - self.param_coordinator = PartitionedParameterCoordinator( - prefetch_bucket_sz=self._prefetch_bucket_sz, - max_reuse_distance_in_numel=self._max_reuse_distance_in_numel, - max_available_parameters_in_numel=self._max_available_parameters_in_numel, - allgather_stream=self.__allgather_stream, - inflight_param_registry=self.__inflight_param_registry, - prefetch_nvme=self.offload_device == OffloadDeviceEnum.nvme, - timers=self.timers, - zero_quantized_weights=self.zero_quantized_weights, - zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights, - fast_sharding_for_leaf_module=self.fast_sharding_for_leaf_module, - log_trace_cache_warnings=self.log_trace_cache_warnings, - ) - - self.forward_hooks = [] - self.backward_hooks = [] - - self.setup_zero_stage3_hooks() - print_rank_0( - f'Created module hooks: forward = {len(self.forward_hooks)}, backward = {len(self.backward_hooks)}', - force=False) - - see_memory_usage("DeepSpeedZeRoOffload initialize [end]", force=True) - - @instrument_w_nvtx - def partition_all_parameters(self): - """Partitioning Parameters that were not partitioned usually if parameters - of modules whose input parameters do not require grad computation do not - trigger post call and will therefore will remain unpartitioned""" - self.get_param_coordinator().release_and_reset_all(self.module) - for param in iter_params(self.module, recurse=True): - if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: - raise RuntimeError(f"{param.ds_summary()} expected to be released") - - def get_param_coordinator(self): - return self.param_coordinator - - def empty_partition_cache(self): - self.partition_all_parameters() - - def _convert_to_zero_parameters(self, ds_config, module, mpu): - non_zero_params = [p for p in module.parameters() if not is_zero_param(p)] - if non_zero_params: - zero_params = [p for p in module.parameters() if is_zero_param(p)] - if zero_params: - zero_params[0].convert_to_zero_parameters(param_list=non_zero_params) - else: - group = None - # parallel_state_sp doesn't have get_data_parallel_group - if mpu and hasattr(mpu, "get_data_parallel_group"): - group = mpu.get_data_parallel_group() - - Init(module=module, - data_parallel_group=group, - dtype=self.dtype, - config_dict_or_path=ds_config, - remote_device=self.offload_device, - pin_memory=self.offload_param_pin_memory, - mpu=mpu, - zero_param_parallel_group=self.zero_param_parallel_group, - zero_quantized_weights=self.zero_quantized_weights, - zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights) - - def destroy(self): - self._remove_module_hooks() - - def _remove_module_hooks(self): - num_forward_hooks = len(self.forward_hooks) - num_backward_hooks = len(self.backward_hooks) - - for hook in self.forward_hooks: - hook.remove() - - for hook in self.backward_hooks: - hook.remove() - - self.fwd_pre_hook.remove() - - print_rank_0(f'Deleted module hooks: forward = {num_forward_hooks}, backward = {num_backward_hooks}', - force=False) - - def setup_zero_stage3_hooks(self): - self.hierarchy = 0 - - #reset step if in inference mode - @instrument_w_nvtx - def _start_of_forward_hook(module, *args): - - self.get_param_coordinator().reset_step() - - self.fwd_pre_hook = self.module.register_forward_pre_hook(_start_of_forward_hook) - - #likely one of them should be enough but just to be safe - self._register_deepspeed_module(self.module) - - # Add top module to stack trace - global FWD_MODULE_STACK - FWD_MODULE_STACK.append(self.module) - - def mark_persistent_parameters(self, param_threshold, model_threshold): - persistent_params = [] - total_persistent_parameters = 0 - params_count = 0 - for name, param in self.module.named_parameters(recurse=True): - if param.ds_numel + total_persistent_parameters > model_threshold: - continue - - if param.ds_numel <= param_threshold: - params_count += 1 - param.ds_persist = True - persistent_params.append(param) - total_persistent_parameters += param.ds_numel - - print_rank_0( - f"Parameter Offload - Persistent parameters statistics: param_count = {params_count}, numel = {total_persistent_parameters}", - force=True) - - return persistent_params - - def _register_deepspeed_module(self, module, count=[0]): - my_count = count[0] - module.ds_id = my_count - - #print(f"{module.__class__} : {module.ds_id}") - - if z3_leaf_module(module): - for param in module.parameters(): - param.ds_z3_leaf_module = module - else: - for child in module.children(): - count[0] = count[0] + 1 - self._register_deepspeed_module(child, count=count) - - @torch.compiler.disable - def _pre_forward_module_hook(module, *args): - self.pre_sub_module_forward_function(module) - - @instrument_w_nvtx - def _post_forward_module_hook(module, input, output): - - global FWD_MODULE_STACK - FWD_MODULE_STACK.pop() - if output is None: - output = [] - elif not isinstance(output, (list, tuple)): - if torch.is_tensor(output): - output = [output] - else: - #print(f'got UNKNOWN type {type(output)}') - outputs = [] - output = output if isinstance(output, dict) else vars(output) - for name, val in output.items(): - if not name.startswith('__') and torch.is_tensor(val): - outputs.append(val) - output = outputs - - for item in filter(lambda item: is_zero_param(item) or hasattr(item, 'ds_param_alias'), output): - key = id(item) if hasattr(item, 'ds_id') else id(item.ds_param_alias) - actual_external_param = item if hasattr(item, 'ds_id') else item.ds_param_alias - - if not any(key in m._external_params for m in FWD_MODULE_STACK): - actual_external_param.is_external_param = True - module_to_register = FWD_MODULE_STACK[-1] - register_external_parameter(module_to_register, actual_external_param) - print_rank_0( - f'Registering dangling parameter for module {module_to_register.__class__.__name__}, ds_id = {actual_external_param.ds_id}.', - force=False) - - # It's possible that the parameter was already external to the completed module. If so, remove it the - # registration as it will be covered by the outer module instead. - if key in module._external_params: - print_rank_0( - f' Unregistering nested dangling parameter from module {module.__class__.__name__}, ds_id = {actual_external_param.ds_id}', - force=False) - unregister_external_parameter(module, actual_external_param) - - actual_external_param.all_gather() - - self.post_sub_module_forward_function(module) - - def _bwd_hook_unexpected_inputs_msg(value): - return f"A module has unknown inputs or outputs type ({type(value)}) and the tensors embedded in it cannot be detected. " \ - "The ZeRO-3 hooks designed to trigger before or after backward pass of the module relies on knowing the input and " \ - "output tensors and therefore may not get triggered properly." - - def _pre_backward_module_hook(module, inputs, output): - - return apply_to_tensors_only(module.pre_bwd_fn.apply, - output, - warning_msg_fn=_bwd_hook_unexpected_inputs_msg) - - #This is an alternate to doing _post_backward_module_hook - #it uses tensor.register_hook instead of using torch.autograd.Function - def _alternate_post_backward_module_hook(module, inputs): - module.ds_grads_remaining = 0 - - #print(f"Before Forward {module.__class__.__name__}") - - def _run_after_backward_hook(*unused): - module.ds_grads_remaining = module.ds_grads_remaining - 1 - if module.ds_grads_remaining == 0: - #print(f"After backward {module.__class__.__name__}") - self.post_sub_module_backward_function(module) - - def _run_before_forward_function(input): - if input.requires_grad: - module.ds_grads_remaining += 1 - - return _apply_forward_and_backward_to_tensors_only(module, _run_before_forward_function, - _run_after_backward_hook, inputs) - - @torch.compiler.disable - def _post_backward_module_hook(module, inputs): - module.ds_grads_remaining = 0 - - return apply_to_tensors_only(module.post_bwd_fn.apply, - inputs, - warning_msg_fn=_bwd_hook_unexpected_inputs_msg) - - # Pre forward hook - self.forward_hooks.append(module.register_forward_pre_hook(_pre_forward_module_hook)) - - # Post forward hook - self.forward_hooks.append(module.register_forward_hook(_post_forward_module_hook)) - - # Pre backward hook - if not hasattr(module, "pre_bwd_fn"): - - @instrument_w_nvtx - def _run_before_backward_function(sub_module): - # some models (e.g. Albert) may run multiple forwards on the same layer in a loop - # before doing backwards, so each backward will need a pre-fetch - using reference - # counting to support this scenario - #print(f"COUNTER before: {sub_module.applied_pre_backward_ref_cnt}") - if sub_module.applied_pre_backward_ref_cnt > 0: - self.pre_sub_module_backward_function(sub_module) - sub_module.applied_pre_backward_ref_cnt -= 1 - #print(f"COUNTER after: {sub_module.applied_pre_backward_ref_cnt}") - - class PreBackwardFunctionForModule(torch.autograd.Function): - - @staticmethod - def forward(ctx, outputs): - # Capture `module` and _run_before_backward_function - ctx.module = module - ctx.pre_backward_function = _run_before_backward_function - if not hasattr(ctx.module, "applied_pre_backward_ref_cnt"): - ctx.module.applied_pre_backward_ref_cnt = 0 - ctx.module.applied_pre_backward_ref_cnt += 1 - outputs = outputs.detach() - return outputs - - @staticmethod - def backward(ctx, *args): - ctx.pre_backward_function(ctx.module) - return args - - module.pre_bwd_fn = PreBackwardFunctionForModule - - self.backward_hooks.append(module.register_forward_hook(_pre_backward_module_hook)) - - # post backward hook - if not hasattr(module, "post_bwd_fn"): - - @instrument_w_nvtx - def _run_after_backward_function(sub_module): - if sub_module.ds_grads_remaining == 0: - self.post_sub_module_backward_function(sub_module) - - class PostBackwardFunctionModule(torch.autograd.Function): - - @staticmethod - def forward(ctx, output): - ctx.module = module - if output.requires_grad: - #TODO SOME TIMES post backward does not seem to be triggered debug in detail - #Should only cause increase in memory not correctness issue - #if output.grad_fn.__class__.__name__ == 'ViewBackward': - # ctx.view=True - # print(f"Warning view tensor for input to module : {module.__class__.__name__}. Backward hooks may not trigger properly") - #assert len(module.parameters(recurse=False)), "The input tensor to the module is a view, and autograd Function or register_hook is not triggered with view tensors." - #if module.ds_grads_remaining == 0: - # print(f"Before Forward: {ctx.module.__class__.__name__}") - module.ds_grads_remaining += 1 - ctx.post_backward_function = _run_after_backward_function - output = output.detach() - return output - - @staticmethod - def backward(ctx, *args): - ctx.module.ds_grads_remaining = ctx.module.ds_grads_remaining - 1 - if ctx.module.ds_grads_remaining == 0: - ctx.post_backward_function(ctx.module) - return args - - module.post_bwd_fn = PostBackwardFunctionModule - - self.backward_hooks.append(module.register_forward_pre_hook(_post_backward_module_hook)) - - @torch.no_grad() - def pre_sub_module_forward_function(self, sub_module): - see_memory_usage(f"Before sub module function {sub_module.__class__.__name__}", force=False) - - global FWD_MODULE_STACK - FWD_MODULE_STACK.append(sub_module) - - param_coordinator = self.get_param_coordinator() - param_coordinator.trace_prologue(sub_module) - if param_coordinator.is_record_trace(): - param_coordinator.record_module(sub_module) - param_coordinator.fetch_sub_module(sub_module, forward=True) - - see_memory_usage(f"Before sub module function {sub_module.__class__.__name__} after fetch", force=False) - - @torch.no_grad() - def post_sub_module_forward_function(self, sub_module): - see_memory_usage( - f"After sub module function {sub_module.__class__.__name__} {sub_module.ds_id} before release", - force=False) - - param_coordinator = self.get_param_coordinator() - param_coordinator.release_sub_module(sub_module, forward=True) - - see_memory_usage( - f"After sub module function {sub_module.__class__.__name__} {sub_module.ds_id} after release", - force=False) - - @torch.no_grad() - def pre_sub_module_backward_function(self, sub_module): - # assert sub_module.training, "backward pass is invalid for module in evaluation mode" - param_coordinator = self.get_param_coordinator() - param_coordinator.trace_prologue(sub_module) - if param_coordinator.is_record_trace(): - param_coordinator.record_module(sub_module) - param_coordinator.fetch_sub_module(sub_module, forward=False) - - @torch.no_grad() - def post_sub_module_backward_function(self, sub_module): - # assert sub_module.training, "backward pass is invalid for module in evaluation mode" - see_memory_usage( - f"After sub module backward function {sub_module.__class__.__name__} {sub_module.ds_id} before release", - force=False) - - self.get_param_coordinator().release_sub_module(sub_module, forward=False) - - see_memory_usage( - f"After sub module backward function {sub_module.__class__.__name__} {sub_module.ds_id} after release", - force=False) - - def _set_z3_leaf_modules_by_threshold(self, module, zero_module_granularity_threshold): - - self._get_granularity_recursively(module) - print_rank_0(f"{'MODULE NAME'.ljust(30)}|{'GRANULARITY VALUE'.rjust(20)}", force=True) - for granularity in self.granularity_info: - print_rank_0(granularity, force=True) - - if self.min_granularity_value <= zero_module_granularity_threshold: - self._set_leaf_by_threshold_preorder(module, zero_module_granularity_threshold) - utils.logger.info( - f"z3_leaf_module was set by stage3_module_granularity_threshold:{zero_module_granularity_threshold}") - for layer in self.z3_leaf_layers: - print_rank_0(f"{layer.__class__.__name__}:{layer.ds_model_granularity}", force=True) - else: - utils.logger.warning( - f"The smallest module granularity is [{self.min_granularity_layer}:{self.min_granularity_value}]. "\ - f"To make stage3_module_granularity_threshold effective, you need to set stage3_module_granularity_threshold >= {self.min_granularity_value}. "\ - f"Current Value:{zero_module_granularity_threshold}" - ) - - def _get_granularity_recursively(self, module): - """This function is used to recursively obtain the granularity of each module.""" - - # avoid setting as leaf for particularly large models, even if the granularity is very small - # an oversized leaf module increases the number of live parameters, introducing memory overhead - Z3_MAX_LEAF_SIZE = 1e9 - - if not list(module.parameters()): - # skip Modules without parameters, such as GELU, etc. - module.ds_model_granularity = sys.maxsize - return 0, 0 - - num_layers = 0 - num_params = 0 - num_params += sum(p.ds_numel for p in module.parameters(recurse=False)) - if not any(module.children()): - # torch leaf module - module.ds_model_granularity = sys.maxsize - return 1, num_params - - for child in module.children(): - layers_in_child, params_in_child = self._get_granularity_recursively(child) - num_layers += layers_in_child - num_params += params_in_child - - if module.__class__.__name__ in torch.nn.modules.container.__all__: - # Do not set container modules like ModuleList as leaf modules - # as this will prevent hooks from being set on their children - # and they may do not invoke the forward method - module.ds_model_granularity = sys.maxsize - return num_layers, num_params - - num_layers += 1 - ds_model_granularity = (num_params // num_layers) if num_params <= Z3_MAX_LEAF_SIZE else sys.maxsize - module.ds_model_granularity = ds_model_granularity - # module.ds_model_num_layers = num_layers - # module.ds_model_num_params = num_params - if self.min_granularity_value > ds_model_granularity: - self.min_granularity_value = ds_model_granularity - self.min_granularity_layer = module.__class__.__name__ - self.granularity_info.add(f"{module.__class__.__name__.ljust(30)}|{str(ds_model_granularity).rjust(20)}") - - return num_layers, num_params - - def _set_leaf_by_threshold_preorder(self, module, granularity_treshhold): - '''Set modules as leaf modules based on the threshold, prioritizing parent nodes.''' - - num_params = sum(p.ds_numel for p in module.parameters()) - if num_params == 0: - # skip Modules without parameters, such as GELU, etc. - return - if module.ds_model_granularity <= granularity_treshhold: - set_z3_leaf_module(module, True) - self.z3_leaf_layers.append(module) - return - - for sub_module in module.children(): - self._set_leaf_by_threshold_preorder(sub_module, granularity_treshhold) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py deleted file mode 100644 index 3482d30e35b9aa54c89102b5ab1a6b2439ba464b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py +++ /dev/null @@ -1,2329 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import math -import os -import types -from typing import Callable, Iterable -from enum import Enum -import functools -import itertools -from typing import List -from collections import defaultdict -import logging -import torch -from torch import Tensor -from deepspeed import comm as dist -from torch.nn import Module -from torch.nn import Parameter - -from .linear import zero3_linear_wrap - -from deepspeed.utils import groups -import deepspeed -from ..utils import see_memory_usage, get_only_unique_item -from deepspeed.runtime.zero.config import DeepSpeedZeroConfig -from deepspeed.runtime.zero.utils import assert_ints_same_as_other_ranks, is_zero_param -from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum -from deepspeed.runtime.config_utils import get_config_default -from deepspeed.utils import instrument_w_nvtx, logger -from deepspeed.comm.comm import init_distributed -from deepspeed.utils.debug import (debug_param2name_id_shape, debug_param2name_id_shape_device, debug_module2name, - debug_param2name_id, debug_param2name_id_shape_status) -from deepspeed.accelerator import get_accelerator -from ..swap_tensor.partitioned_param_swapper import AsyncPartitionedParameterSwapper, PartitionedParamStatus -from deepspeed.inference.quantization.utils import _quantize_param, WEIGHT_QUANTIZATION_LAYERS, wrap_quantized_functional, wrap_load_from_state_dict -from deepspeed.runtime.torch_autocast import sort_dtypes, get_autocast_dtype, has_autocast_dtype - -partitioned_param_data_shape = [0] -zero_init_context = 0 -top_level_context = None - - -class DeepSpeedTensorOverride(Enum): - dtype = 1 - device = 2 - - -DEFAULT_TENSOR_OVERRIDES = [DeepSpeedTensorOverride.dtype, DeepSpeedTensorOverride.device] - - -def get_allgather_dtype(param, param_ds_tensor): - autocast = has_autocast_dtype(param) - return get_autocast_dtype(param) if autocast else param_ds_tensor.dtype - - -class NoGatherHandle: - - def __init__(self, param: Parameter) -> None: - if param.ds_status != ZeroParamStatus.INFLIGHT: - raise RuntimeError(f"expected param {param.ds_summary()} to be available") - - if hasattr(param.ds_tensor, "ds_quant_scale"): - param.data = Init.quantizer_module.dequantize(param.ds_tensor.data, param.ds_tensor.ds_quant_scale).to( - device=get_accelerator().current_device_name(), non_blocking=True).view(param.ds_shape) - else: - param.data = param.ds_tensor.data.to(device=get_accelerator().current_device_name(), - non_blocking=True).view(param.ds_shape) - self.__param = param - - def wait(self, **kwargs) -> None: - if not get_accelerator().resolves_data_dependency(): - get_accelerator().current_stream().synchronize() - self.__param.ds_status = ZeroParamStatus.AVAILABLE - - -class NoGatherCoalescedHandle: - - def __init__(self, params: List[Parameter]) -> None: - self.__params = params - self.__complete = False - - for param in self.__params: - if param.ds_status != ZeroParamStatus.INFLIGHT: - raise RuntimeError(f"expected param {param.ds_summary()} to not be available") - if hasattr(param.ds_tensor, "ds_quant_scale"): - param.data = Init.quantizer_module.dequantize(param.ds_tensor.data, param.ds_tensor.ds_quant_scale).to( - device=get_accelerator().current_device_name(), non_blocking=True).view(param.ds_shape) - else: - param.data = param.ds_tensor.data.to(device=get_accelerator().current_device_name(), - non_blocking=True).view(param.ds_shape) - - @instrument_w_nvtx - def wait(self, **kwargs) -> None: - if self.__complete: - return - - if not get_accelerator().resolves_data_dependency(): - get_accelerator().current_stream().synchronize() - for param in self.__params: - assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight" - param.ds_status = ZeroParamStatus.AVAILABLE - - self.__complete = True - - -def _dist_allgather_fn(input_tensor: Tensor, output_tensor: Tensor, group=None): - return instrument_w_nvtx(dist.allgather_fn)(output_tensor, input_tensor, group=group, async_op=True) - - -def print_rank_0(message, debug=False, force=False): - rank = dist.get_rank() - if rank == 0 and (debug or force): - print(message) - # other variations - # - print for all ranks w/o interleaving - # printflock(f"[{rank}] {message}") - # - print to log file per rank - # log_rank_file(rank, message) - - -def debug_rank0(msg: str) -> None: - if dist.get_rank() == 0: - logger.debug(msg) - - -def _init_external_params(module): - if not hasattr(module, '_external_params'): - module._external_params = {} - - def external_parameters(self): - return self._external_params.items() - - def all_parameters(self): - return itertools.chain(self.named_parameters(self, recurse=False), external_parameters(self)) - - module.ds_external_parameters = types.MethodType(external_parameters, module) - module.all_parameters = types.MethodType(all_parameters, module) - - -def register_external_parameter(module, parameter): - """Instruct DeepSpeed to coordinate ``parameter``'s collection and partitioning in - the forward and backward passes of ``module``. - - This is used when a parameter is accessed outside of its owning module's - ``forward()``. DeepSpeed must know to collect it from its partitioned - state and when to release the memory. - - .. note:: - This is only applicable to training with ZeRO stage 3. - - Args: - module (``torch.nn.Module``): The module that requires ``parameter`` in its forward pass. - parameter (``torch.nn.Parameter``): The parameter to register. - - Raises: - RuntimeError: If ``parameter`` is not of type ``torch.nn.Parameter``. - - - Examples - ======== - - #. Register a weight that is used in another module's forward pass (line 6). - Parameter ``layer1.weight`` is used by ``layer2`` (line 11). - - .. code-block:: python - :linenos: - :emphasize-lines: 6,11 - - class ModuleZ3(torch.nn.Module): - def __init__(self, *args): - super().__init__(self, *args) - self.layer1 = SomeLayer() - self.layer2 = OtherLayer() - deepspeed.zero.register_external_parameter(self, self.layer1.weight) - - def forward(self, input): - x = self.layer1(input) - # self.layer1.weight is required by self.layer2.forward - y = self.layer2(x, self.layer1.weight) - return y - """ - if not isinstance(parameter, torch.nn.Parameter): - raise RuntimeError('Parameter is not a torch.nn.Parameter') - - if not hasattr(module, '_external_params'): - _init_external_params(module) - - key = id(parameter) - module._external_params[key] = parameter - - -def unregister_external_parameter(module, parameter): - """Reverses the effects of :meth:`register_external_parameter`. - - Args: - module (``torch.nn.Module``): The module to affect. - parameter (``torch.nn.Parameter``): The parameter to unregister. - - Raises: - RuntimeError: If ``parameter`` is not of type ``torch.nn.Parameter``. - RuntimeError: If ``parameter`` is not a registered external parameter of ``module``. - """ - if not isinstance(parameter, torch.nn.Parameter): - raise RuntimeError('Parameter is not a torch.nn.Parameter') - - if not hasattr(module, '_external_params') or id(parameter) not in module._external_params: - raise RuntimeError('Parameter is not a registered external parameter of module.') - - key = id(parameter) - del module._external_params[key] - - -class ZeroParamType(Enum): - - # same as regular pytorch parameters - NORMAL = 1 - - # parameters are partitioned across data parallel process - PARTITIONED = 2 - - # the parameter is held with a unique process rank - # and is not available on all other process - REMOTE = 3 - - -class ZeroParamStatus(Enum): - # parameters are fully present and ready for use on all processes - AVAILABLE = 1 - - # parameters are either partitioned or remote in some or all process - NOT_AVAILABLE = 2 - - # parameters are being gathered. - INFLIGHT = 3 - - -_orig_torch_tensor = torch.tensor -_orig_torch_empty = torch.empty -_orig_torch_zeros = torch.zeros -_orig_torch_ones = torch.ones -_orig_torch_full = torch.full -_orig_torch_arange = torch.arange -_orig_torch_eye = torch.eye -_orig_torch_randn = torch.randn - - -def zero_wrapper_for_fp_tensor_constructor(fn: Callable, target_fp_dtype: torch.dtype, - target_device: torch.device) -> Callable: - - def wrapped_fn(*args, **kwargs) -> Tensor: - if kwargs.get("device", None) is None and target_device is not None: - kwargs['device'] = target_device - tensor: Tensor = fn(*args, **kwargs) - if target_fp_dtype is not None and tensor.is_floating_point(): - tensor.data = tensor.data.to(target_fp_dtype) - - return tensor - - return wrapped_fn - - -def get_new_tensor_fn_for_dtype(target_fp_dtype: torch.dtype, target_device: torch.device) -> Callable: - - def new_tensor(cls, *args, **kwargs) -> Tensor: - if not args: - args = (0, ) - if target_device is None: - tensor = _orig_torch_empty(0).new_empty(*args, **kwargs) - else: - tensor = _orig_torch_empty(0, device=target_device).new_empty(*args, **kwargs) - - if tensor.is_floating_point() and target_fp_dtype is not None: - tensor = tensor.to(target_fp_dtype) - - return tensor - - return new_tensor - - -# https://stackoverflow.com/a/63851681/9201239 -def get_all_subclasses(cls, include_root=True): - subclass_list = [] - - def recurse(cl): - for subclass in cl.__subclasses__(): - subclass_list.append(subclass) - recurse(subclass) - - recurse(cls) - - ret = set(subclass_list) - if include_root: - ret.add(cls) - return ret - - -@instrument_w_nvtx -def free_param(param: Parameter) -> None: - """Free underlying storage of a parameter.""" - assert not param.ds_active_sub_modules, param.ds_summary() - if get_accelerator().on_accelerator(param.data): - # need to make sure that we don't free the parameter while it is still - # being used for computation - if not get_accelerator().is_synchronized_device(): - param.data.record_stream(get_accelerator().current_stream()) - # param.data doesn't store anything meaningful in partitioned state - param.data = torch.empty(0, dtype=param.dtype, device=param.device) - param.ds_status = ZeroParamStatus.NOT_AVAILABLE - - -reuse_buffers = False -temp_contiguous_tensor = None -empty_buffers = {} - - -# Inserts _post_init_method at the end of init method -# for all sub classes of torch.nn.Module -class InsertPostInitMethodToModuleSubClasses(object): - num_module_parameters = 0 - num_module_elements = 0 - - def __init__(self, enabled=True, mem_efficient_linear=True, ds_config=None, dtype=None): - self.mem_efficient_linear = mem_efficient_linear - self.enabled = enabled - self._set_dtype(ds_config, dtype) - assert self.dtype in [ - torch.half, torch.bfloat16, torch.float - ], f"Invalid data type {self.dtype}, allowed values are [torch.half, torch.bfloat16, torch.float]" - self.wrapped_cls = set() - self.skip_init_depth = 0 - - self.quantized_initialization = None - if ds_config is not None and ds_config.weight_quantization_config and ds_config.weight_quantization_config.quantized_initialization: - self.quantized_initialization = ds_config.weight_quantization_config.quantized_initialization - - def __enter__(self): - if not self.enabled: - return - - global zero_init_context - if zero_init_context == 0: - self.patch_init_and_builtins() - global top_level_context - top_level_context = self - - zero_init_context += 1 - - def __exit__(self, exc_type, exc_value, traceback): - if not self.enabled: - return - - global zero_init_context - zero_init_context -= 1 - - # Exiting the top level context - if zero_init_context == 0: - self.unpatch_init_and_builtins() - global top_level_context - top_level_context = None - - if dist.get_rank() == 0: - billion_elems = InsertPostInitMethodToModuleSubClasses.num_module_elements / 1e9 - num_params = InsertPostInitMethodToModuleSubClasses.num_module_parameters - logger.info( - f"finished initializing model - num_params = {num_params}, num_elems = {billion_elems:.2f}B") - - # Now that we cleaned up the metaclass injection, raise the exception. - if exc_type is not None: - return False - - # To be implemented by inheriting classes - def _post_init_method(self, module): - pass - - def _set_dtype(self, ds_config, dtype): - if ds_config is not None and dtype is None: - if ds_config.bfloat16_config.enabled and ds_config.float16_config.enabled: - raise RuntimeError("bfloat16 and fp16 cannot be enabled at once") - - if ds_config.bfloat16_config.enabled: - self.dtype = torch.bfloat16 - elif ds_config.float16_config.enabled: - self.dtype = torch.half - else: - self.dtype = torch.float - else: - self.dtype = dtype or torch.float16 if get_accelerator().is_fp16_supported( - ) else torch.bfloat16 if get_accelerator().is_bf16_supported else torch.float32 - - def patch_init_and_builtins(self): - - def apply_with_gather(orig_module_apply_fn: Callable) -> Callable: - """many models make use of child modules like Linear or Embedding which - perform their own weight initialization in their __init__ methods, - but will then have more weight initialization in a parent module's __init__ - method that modifies weights of child modules, which is typically done - using the Module.apply method. - - since the Init context manager partitions child modules immediately after - they are initialized, without modifying apply we would entirely skip - any initialization done by parent modules. - - to get around this issue, we wrap the function passed to Module.apply - so that the applied function is applied to child modules correctly. - """ - - def get_wrapped_fn_to_apply(fn_to_apply: Callable) -> Callable: - if hasattr(fn_to_apply, "wrapped"): - return fn_to_apply - - @functools.wraps(fn_to_apply) - def wrapped_fn_to_apply(module_to_apply_fn_to: Module) -> None: - """gathers parameters before calling apply function. afterwards - parameters are broadcasted to ensure consistency across all ranks - then re-partitioned. - - takes the following steps: - 1. allgathers parameters for the current module being worked on - 2. calls the original function - 3. broadcasts root rank's parameters to the other ranks - 4. re-partitions the parameters - """ - - # TODO Delay error checking for dangling partitioned parameters to post module init - # raise RuntimeError(f"not all parameters for {module_to_apply_fn_to.__class__.__name__}, " - # f"were zero params, is it possible that the parameters were " - # f"overwritten after they were initialized? " - # f"params: {[p for p in module_to_apply_fn_to.parameters(recurse=False)]} ") - - params_to_apply_fn_to: Iterable[Parameter] = list( - sorted([p for p in module_to_apply_fn_to.parameters(recurse=False) if is_zero_param(p)], - key=lambda p: p.ds_id)) - - for param in params_to_apply_fn_to: - param.all_gather() - - fn_to_apply(module_to_apply_fn_to) - - for param in params_to_apply_fn_to: - dist.broadcast(param.data, 0, group=param.ds_process_group) - - for param in params_to_apply_fn_to: - param.partition(has_been_updated=True) - - wrapped_fn_to_apply.wrapped = True - - return wrapped_fn_to_apply - - @functools.wraps(orig_module_apply_fn) - def wrapped_apply(module: Module, fn_to_apply: Callable) -> None: - orig_module_apply_fn(module, get_wrapped_fn_to_apply(fn_to_apply)) - - return wrapped_apply - - def hook_for_skip_init(module): - # this function is intended for handling the logic of torch.nn.utils.skip_init - # skip_init:module_cls(*args, **kwargs).to_empty(device=final_device), where kwargs['device']='meta' - # the function call occurs between module_cls(*args, **kwargs) and to_empty(device=final_device). - def partition_after_empty_init(f): - - @functools.wraps(f) - def wrapper(module, *args, **kwargs): - _module = f(module, *args, **kwargs) - # here is the post-hook for module.apply(empty_like...) - # after module.apply(empty_like...), the module has completed its empty init on real device - # since skip_init won't involve any computations or weight adjustments, we can directly utilize post_init - self._post_init_method(_module) - return _module - - return wrapper - - def post_wrapper_to_empty(f): - # append some wrapper restoration after to_empty() call - @functools.wraps(f) - def wrapper(*args, **kwargs): - res = f(*args, **kwargs) - # restore _apply hook - for subclass in get_all_subclasses(torch.nn.modules.module.Module): - _disable_class_apply(subclass) - # self restore - module.to_empty = f - return res - - return wrapper - - def _enable_class_apply(cls): - if '_apply' in cls.__dict__: - cls._old_apply_of_skip_init_hook = cls._apply - cls._apply = partition_after_empty_init(cls._apply) - - def _disable_class_apply(cls): - if hasattr(cls, '_old_apply_of_skip_init_hook'): - cls._apply = cls._old_apply_of_skip_init_hook - - # add hooks for to_empty: apply_(empty_like) - for subclass in get_all_subclasses(torch.nn.modules.module.Module): - _enable_class_apply(subclass) - - # add a restore hook when exiting skip_init - module.to_empty = post_wrapper_to_empty(module.to_empty) - - def partition_after(f): - - @functools.wraps(f) - def wrapper(module, *args, **kwargs): - - # important logic: We want to run post_init only after child's __init__ is - # completed, and do nothing after __init__ of any of its parents and grandparents in - # the inheritance ancestry. This way the partitioning will need to happen only once - # when the whole object is ready to be partitioned and not before. This is because - # often the child module will need to tweak the weights - for example running a - # custom weights init function. So if a parent created the weights param, the child - # won't need to gather it in order to tweak it - - print_rank_0(f'Before initializing {module.__class__.__name__}', force=False) - - is_child_module = False - if not hasattr(module, "_ds_child_entered"): - # child's __init__ was called, since parents all see the same object they can now skip post_init - is_child_module = True - setattr(module, "_ds_child_entered", True) - - init_on_meta = 'device' in kwargs and kwargs['device'] == 'meta' - if init_on_meta: - self.skip_init_depth += 1 - - f(module, *args, **kwargs) - if init_on_meta and self.skip_init_depth == 1: - # check and handle the logic of empty_init - hook_for_skip_init(module) - if is_child_module: - # child's __init__ is done, now we can run a single post_init on the child object - delattr(module, "_ds_child_entered") - - print_rank_0(f'Running post_init for {module.__class__.__name__}', force=False) - if self.skip_init_depth == 0: - self._post_init_method(module) - - print_rank_0(f'After initializing followed by post init for {module.__class__.__name__}', force=False) - if init_on_meta: - self.skip_init_depth -= 1 - - return wrapper - - def _enable_class(cls): - if '__init__' in cls.__dict__: - cls._old_init = cls.__init__ - cls.__init__ = partition_after(cls.__init__) - - def _init_subclass(cls, **kwargs): - if '__init__' in cls.__dict__: - cls._old_init = cls.__init__ - cls.__init__ = partition_after(cls.__init__) - - # Replace .__init__() for all existing subclasses of torch.nn.Module recursively - for subclass in get_all_subclasses(torch.nn.modules.module.Module): - _enable_class(subclass) - - # holding onto some methods so we can put them back the way they were in __exit__ - torch.nn.modules.module.Module._old_init_subclass = torch.nn.modules.module.Module.__init_subclass__ - torch.nn.modules.module.Module._old_apply = torch.nn.modules.module.Module.apply - torch.Tensor.__old_new__ = torch.Tensor.__new__ - - # Replace .__init__() for future subclasses of torch.nn.Module - torch.nn.modules.module.Module.__init_subclass__ = classmethod(_init_subclass) - if Init.override_module_apply: - torch.nn.modules.module.Module.apply = apply_with_gather(torch.nn.modules.module.Module._old_apply) - - if self.tensor_overrides: - self._add_tensor_creation_wrappers() - - if self.mem_efficient_linear: - print_rank_0( - "nn.functional.linear has been overridden with a more memory efficient version. This will persist unless manually reset.", - force=False) - if not hasattr(InsertPostInitMethodToModuleSubClasses, "linear_bk"): - InsertPostInitMethodToModuleSubClasses.linear_bk = torch.nn.functional.linear - torch.nn.functional.linear = zero3_linear_wrap - - if self.quantized_initialization: - print_rank_0("nn.functional.linear has been overridden with quantized linear version.", force=False) - torch.nn.functional.linear = wrap_quantized_functional(torch.nn.functional.linear) - torch.nn.functional.embedding = wrap_quantized_functional(torch.nn.functional.embedding) - for cls in WEIGHT_QUANTIZATION_LAYERS: - cls._load_from_state_dict = wrap_load_from_state_dict(cls._load_from_state_dict) - - logger.info("Enable Zero3 engine with INT4 quantization.") - - self.patched = True - - def unpatch_init_and_builtins(self): - if self.patched: - - def _disable_class(cls): - if hasattr(cls, '_old_init'): - cls.__init__ = cls._old_init - - for subclass in get_all_subclasses(torch.nn.modules.module.Module): - _disable_class(subclass) - - # putting methods back the way we found them - torch.nn.modules.module.Module.__init_subclass__ = torch.nn.modules.module.Module._old_init_subclass - if Init.override_module_apply: - torch.nn.modules.module.Module.apply = torch.nn.modules.module.Module._old_apply - - if self.tensor_overrides: - self._remove_tensor_creation_wrappers() - - self.patched = False - - def _add_tensor_creation_wrappers(self): - if DeepSpeedTensorOverride.dtype in self.tensor_overrides: - target_fp_dtype = self.dtype - else: - target_fp_dtype = None - if DeepSpeedTensorOverride.device in self.tensor_overrides: - target_device = self.local_device - else: - target_device = None - - torch.Tensor.__new__ = get_new_tensor_fn_for_dtype(target_fp_dtype=target_fp_dtype, - target_device=target_device) - torch.tensor = zero_wrapper_for_fp_tensor_constructor(_orig_torch_tensor, - target_fp_dtype=target_fp_dtype, - target_device=target_device) - torch.empty = zero_wrapper_for_fp_tensor_constructor(_orig_torch_empty, - target_fp_dtype=target_fp_dtype, - target_device=target_device) - torch.zeros = zero_wrapper_for_fp_tensor_constructor(_orig_torch_zeros, - target_fp_dtype=target_fp_dtype, - target_device=target_device) - torch.ones = zero_wrapper_for_fp_tensor_constructor(_orig_torch_ones, - target_fp_dtype=target_fp_dtype, - target_device=target_device) - torch.full = zero_wrapper_for_fp_tensor_constructor(_orig_torch_full, - target_fp_dtype=target_fp_dtype, - target_device=target_device) - torch.arange = zero_wrapper_for_fp_tensor_constructor(_orig_torch_arange, - target_fp_dtype=target_fp_dtype, - target_device=target_device) - torch.eye = zero_wrapper_for_fp_tensor_constructor(_orig_torch_eye, - target_fp_dtype=target_fp_dtype, - target_device=target_device) - torch.randn = zero_wrapper_for_fp_tensor_constructor(_orig_torch_randn, - target_fp_dtype=target_fp_dtype, - target_device=target_device) - - def _remove_tensor_creation_wrappers(self): - torch.Tensor.__new__ = torch.Tensor.__old_new__ - torch.tensor = _orig_torch_tensor - torch.empty = _orig_torch_empty - torch.zeros = _orig_torch_zeros - torch.ones = _orig_torch_ones - torch.full = _orig_torch_full - torch.arange = _orig_torch_arange - torch.eye = _orig_torch_eye - torch.randn = _orig_torch_randn - - -def shutdown_init_context(): - """ - This function is used to initialize deepspeed engine inside the context of Init. - We need to remove the wrappers but keep the context. - """ - if top_level_context: - top_level_context.unpatch_init_and_builtins() - - -def restore_init_context(): - """ - This function is used to restore the wrappers after deepspeed engine is initialized. - """ - if top_level_context: - top_level_context.patch_init_and_builtins() - - -class AllGatherHandle: - - def __init__(self, handle, param: Parameter, quantization=None, original_dtype=None) -> None: - if param.ds_status != ZeroParamStatus.INFLIGHT: - raise RuntimeError(f"expected param {param.ds_summary()} to be available") - - # Only one of original_dtype or quantization is provided - assert (original_dtype is None) != (quantization is None) - - self.__handle = handle - self.__param = param - self.__quantization = quantization - self.__original_dtype = original_dtype - - def wait(self, handle_dependency=True) -> None: - instrument_w_nvtx(self.__handle.wait)() - - if self.__original_dtype: - self.__param.data = self.__param.data.to(self.__original_dtype) - elif self.__quantization: - instrument_w_nvtx(self.__quantization.quant_handle.wait)() - self.__param.data = self.__quantization.backend.dequantize( - self.__quantization.quantized_param, self.__quantization.scale_buffer).to(self.__param.device) - self.__param.ds_status = ZeroParamStatus.AVAILABLE - - -class AllGatherCoalescedHandle: - - data_buffer = [] - - def __init__( - self, - allgather_handle, - params: List[Parameter], - partitions: List[Tensor], - world_size: int, - use_secondary_tensor=False, - quantization=None, - ) -> None: - self.allgather_handle = allgather_handle - self.params = params - self.partitions = partitions - self.world_size = world_size - self.use_secondary_tensor = use_secondary_tensor - self.complete = False - self.quantization = quantization - - for param in self.params: - if param.ds_status != ZeroParamStatus.INFLIGHT: - raise RuntimeError(f"expected param {param.ds_summary()} to not be available") - - @instrument_w_nvtx - def wait(self, handle_dependency=True) -> None: - if self.complete: - return - - instrument_w_nvtx(self.allgather_handle.wait)() - - if self.quantization: - instrument_w_nvtx(self.quantization.quant_handle.wait)() - flat_tensor = self.quantization.backend.dequantize( - self.quantization.quantized_param, self.quantization.scale_buffer).to(self.params[0].device) - - self.partitions: List[Parameter] = [] - for i in range(self.world_size): - self.partitions.append( - flat_tensor.narrow(0, self.quantization.partition_sz * i, self.quantization.partition_sz)) - - # split the single tensor out into individual tensors - param_offset = 0 - for param in self.params: - assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight" - partitions: List[Tensor] = [] - ds_tensor_numel = param.ds_tensor.ds_numel - if self.use_secondary_tensor: - ds_tensor_numel *= param.ds_secondary_tensor_num_of_groups - for rank in range(self.world_size): - param_start = rank * ds_tensor_numel - if param_start < param.ds_numel: - part_to_copy = self.partitions[rank].narrow(0, param_offset, - min(param.ds_numel - param_start, ds_tensor_numel)) - partitions.append(part_to_copy) - # Note that dtypes of param and partitions can be different (currently for torch.autocast support) - param.data = instrument_w_nvtx(torch.cat)(partitions).view(param.ds_shape).to(param.ds_tensor.dtype) - param.ds_status = ZeroParamStatus.AVAILABLE - if not get_accelerator().is_synchronized_device() and handle_dependency: - for part_to_copy in partitions: - part_to_copy.record_stream(get_accelerator().current_stream()) - - param_offset += ds_tensor_numel - - self.complete = True - if not get_accelerator().is_synchronized_device() and not handle_dependency: - # if the device needs to handle dependencies and opts for explicit processing outside the function. - AllGatherCoalescedHandle.data_buffer.append(partitions) - - @staticmethod - def free_buffer(): - AllGatherCoalescedHandle.data_buffer = [] - - -class MultipleAllGatherHandles: - - def __init__(self, handles: List[AllGatherCoalescedHandle]): - self.handles = handles - - def wait(self, handle_dependency=True) -> None: - for handle in self.handles: - handle.wait(handle_dependency) - - -class AllReduceCoalescedHandle: - - def __init__(self, handle, params: List[Parameter]) -> None: - self.handle = handle - self.params = params - self.complete = False - - for param in self.params: - if param.ds_status != ZeroParamStatus.INFLIGHT: - raise RuntimeError(f"expected param {param.ds_summary()} to not be available") - - @instrument_w_nvtx - def wait(self) -> None: - if self.complete: - return - - instrument_w_nvtx(self.handle.wait)() - - for param in self.params: - assert param.ds_status == ZeroParamStatus.INFLIGHT, f"expected param {param.ds_summary()} to be inflight" - param.ds_status = ZeroParamStatus.AVAILABLE - - self.complete = True - - -class QuantizationInfo: - # a placeholder object to store all quant related vars used in handles - def __init__(self) -> None: - self.quantized_param = None - self.backend = None - self.quant_handle = None - self.scale_buffer = None - - -class CUDAQuantizer: - async_flag = True - target_group_size = 8000 # the optimal size is 4k, so we set the target to be below 8k - group_size_cache = dict() - quantizer_cuda_module = None - - def __init__(self) -> None: - if CUDAQuantizer.quantizer_cuda_module is None: - CUDAQuantizer.quantizer_cuda_module = deepspeed.ops.op_builder.QuantizerBuilder().load() - - def quantize(self, param, groups=None): - if groups is None: - try: - groups = self.group_size_cache[param.numel()] - except KeyError: - groups = math.ceil(param.numel() / self.target_group_size) - while groups < param.numel(): - if param.numel() % (8 * groups) == 0: - break - groups += 1 - while True: - if param.numel() % (8 * groups * 2) == 0 and param.numel( - ) / groups > self.target_group_size: #hard limit of 16k group_size - groups *= 2 - else: - break - assert ( - param.numel() % (8 * groups) == 0 - ), f"Qantized weight requires the number of weights be a multiple of 8. Yet {param.numel()} cannot be divided by 8*{groups}" - assert (param.numel() / groups < 16000), f"{param.numel()} / {groups} is larger than 16k" - assert param.numel( - ) > groups, f"Adaptive grouping algorithm cannot find a group size for input tensor of size {param.numel()}" - self.group_size_cache[param.numel()] = groups - return self.quantizer_cuda_module.quantize(param.to(get_accelerator().device_name()), groups, 8, - self.quantizer_cuda_module.Symmetric) - - def dequantize(self, quantized_param, scale): - return self.quantizer_cuda_module.dequantize(quantized_param, scale, scale.numel(), 8, - self.quantizer_cuda_module.Symmetric) - - -def _no_gather_coalesced(params: Iterable[Parameter]) -> AllGatherCoalescedHandle: - for param in params: - if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: - raise RuntimeError(f"expect param.ds_status == ZeroParamStatus.NOT_AVAILABLE, got{param.ds_summary()}") - param.ds_status = ZeroParamStatus.INFLIGHT - - params = sorted(params, key=lambda p: p.ds_id) - if len(params) == 1: - param, = params - return NoGatherHandle(param) - return NoGatherCoalescedHandle(params) - - -# Replaces all parameters in module with Scattered Parameters -class Init(InsertPostInitMethodToModuleSubClasses): - param_id = 0 - param_persistence_threshold = get_config_default(DeepSpeedZeroConfig, "param_persistence_threshold") - model_persistence_threshold = get_config_default(DeepSpeedZeroConfig, "model_persistence_threshold") - num_persisted_parameters = 0 - num_persisted_elements = 0 - apply_param_persistence = False - override_module_apply = get_config_default(DeepSpeedZeroConfig, "override_module_apply") - - def __init__(self, - module=None, - data_parallel_group=None, - mem_efficient_linear=True, - remote_device=None, - pin_memory=False, - config_dict_or_path=None, - config=None, - enabled=True, - dtype=None, - mpu=None, - zero_param_parallel_group=None, - zero_quantized_weights=False, - zero_quantized_nontrainable_weights=False, - sequence_data_parallel_group=None, - param_swapper=None, - tensor_overrides=DEFAULT_TENSOR_OVERRIDES): - """A context to enable massive model construction for training with - ZeRO-3. Models are automatically partitioned (or, sharded) across the - system and converted to half precision. - - Args: - module (``torch.nn.Module``, optional): If provided, partition the model as - if it was constructed in the context. - data_parallel_group (``deepspeed.comm`` process group, optional): - The group of processes to partition among. Defaults to all processes. - Synonymous with sequence data parallel group for param partitioning - across both sequence and data parallel groups. - mem_efficient_linear (bool, optional): Replace - torch.nn.functional.linear with an implementation that allows - DeepSpeed to partition parameters. Defaults to ``True``. - remote_device (string, optional): The initial device to store model - weights e.g., ``cpu``, ``nvme``. Passing ``"cpu"`` will create the model in CPU - memory. The model may still be moved to GPU based on the - offload settings for training. Defaults to param offload device if a config is - defined, otherwise GPU. - pin_memory (bool, optional): Potentially increase performance by - using pinned memory for model weights. ``remote_device`` must be - ``"cpu"``. Defaults to pin_memory value in config, otherwise ``False``. - config_dict_or_path (dict or ``json file``, optional): If provided, provides configuration - for swapping fp16 params to NVMe and other things like ``dtype``. - config (dict or ``json file``, optional): Deprecated, use config_dict_or_path instead. - enabled (bool, optional): If ``False``, this context has no - effect. Defaults to ``True``. - dtype (``dtype``, optional): Can be used to change the data type of the parameters. - Supported options are ``torch.half`` and ``torch.float``. Defaults to ``None`` - mpu (``object``, optional): A model parallelism unit object that implements get_{model,data}_parallel_{rank,group,world_size}. - zero_param_parallel_group(``object``, optional): Parallel (comm) group for dual partitioning of ZeRO params. - zero_quantized_weights (bool, optional): If ``True``, turn on quantized weights in all gather weights. Default is ``False`` - zero_quantized_nontrainable_weights (bool, optional): If ``True``, nontrainable weights will be stored in quantized format. Default is ``False`` - param_swapper (``deepspeed.runtime.swap_tensor.partitioned_param_swapper.AsyncPartitionedParameterSwapper``, optional): [Experimental] Use existing parameter swapper. Defaults to ``None``. - This argument will be removed in the near future. - tensor_overrides ([`deepspeed.runtime.zero.DeepSpeedTensorOverride`], optional): Tensor attributes to override. Defaults to overriding dtype and device. - - This context accelerates model initialization and enables models that - are too large to allocate in their entirety in CPU memory. It has the - following effects: - - #. allocates tensors to either GPU or CPU memory or NVMe - #. converts floating point tensors to half precision - #. immediately partitions tensors among the group of data-parallel devices - #. (*optional*) replaces ``torch.nn.functional.linear`` with a more - memory-efficient implementation - - These modifications allow for models that exceed the size of local CPU/GPU - memory/NVMe, but fit within the total NVMe capacity (*i.e.*, aggregate CPU - or GPU memory or NVMe) across all nodes. Consider initializing a model with one - trillion parameters, whose weights occupy two terabytes (TB) in half - precision. The initial CPU allocation in full precision requires 4TB of - memory *per process*, and so a system with 8 GPUs per node would need 32TB of - CPU memory due to data-parallel redundancies. Instead, by immediately - partitioning tensors we remove the redundancies. The result is that - regardless of the number of GPUs, we still only require the original 4TB. This - allows for a linear increase in model size with the aggregate system memory. - For example, if a node has 1TB of memory and 8 GPUs, we could fit a trillion - parameter model with 4 nodes and 32 GPUs. - - Important: If the fp16 weights of the model can't fit onto a single GPU memory - this feature must be used. - - .. note:: - Initializes ``deepspeed.comm`` if it has not already been done so. - See :meth:`deepspeed.init_distributed` for more information. - - .. note:: - Only applicable to training with ZeRO-3. - - Examples - -------- - - #. Allocate a model and partition it among all processes: - - .. code-block:: python - - with deepspeed.zero.Init(): - model = MyLargeModel() - - - #. Allocate a model in pinned CPU memory and partition it among a subgroup of processes: - - .. code-block:: python - - with deepspeed.zero.Init(data_parallel_group=mpu.get_data_parallel_group(), - remote_device="cpu", - pin_memory=True): - model = MyLargeModel() - - - #. Partition an already-allocated model in CPU memory: - - .. code-block:: python - - model = deepspeed.zero.Init(module=model) - """ - if config is not None: - config_dict_or_path = config - logger.warning('zero.Init: the `config` argument is deprecated. Please use `config_dict_or_path` instead.') - _ds_config = deepspeed.runtime.config.DeepSpeedConfig(config_dict_or_path, - mpu) if config_dict_or_path is not None else None - if _ds_config is not None: - mem_efficient_linear = _ds_config.zero_config.memory_efficient_linear - - self.tensor_overrides = tensor_overrides - super().__init__(enabled=enabled, mem_efficient_linear=mem_efficient_linear, ds_config=_ds_config, dtype=dtype) - if not dist.is_initialized(): - init_distributed() - assert dist.is_initialized(), "Parameters cannot be scattered without initializing deepspeed.comm" - - if data_parallel_group is None: - self.ds_process_group = dist.get_world_group() - else: - self.ds_process_group = data_parallel_group - - if sequence_data_parallel_group is not None: - logger.warning( - f"sequence_data_parallel_group' is deprecated and will be removed. Use 'data_parallel_group' instead.") - if data_parallel_group is not None: - raise ValueError( - "Both 'data_parallel_group' and 'sequence_data_parallel_group' were specified. Please provide only one of these arguments." - ) - self.ds_process_group = sequence_data_parallel_group - - self.rank = dist.get_rank(group=self.ds_process_group) - self.dp_world_size = dist.get_world_size(group=self.ds_process_group) - - self.zero_param_process_group = zero_param_parallel_group - if _ds_config is not None and _ds_config.zero_config.zero_hpz_partition_size > 1 and self.zero_param_process_group is None: - groups._create_zero_param_parallel_group(_ds_config.zero_config.zero_hpz_partition_size) - self.zero_param_process_group = groups._get_zero_param_intra_parallel_group() - - self.num_ranks_in_param_group = self.dp_world_size - self.rank_in_group = self.rank - self.num_param_groups = 1 - - if self.zero_param_process_group is not None: - self.num_ranks_in_param_group = groups._get_zero_param_intra_parallel_group_world_size() - self.num_param_groups = int(self.dp_world_size / self.num_ranks_in_param_group) - self.rank_in_group = groups._get_zero_param_intra_parallel_rank_in_mygroup() - print_rank_0(f"hpZeRO group size: {self.num_ranks_in_param_group}", force=True) - - logger.debug( - "hpZeRO partition parameter my rank in world {} my rank in group {} ranks in my param partition group: {} " - .format(self.rank, self.rank_in_group, groups._get_zero_param_intra_parallel_group_ranks())) - - # Local device is the device where the parameters are consumed, must be default device. - # It is the device where parameters are fully instantiated using allgather - self.local_device = torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])) - get_accelerator().set_device(self.local_device) - - self.quantized_weights = zero_quantized_weights - if _ds_config is not None and _ds_config.zero_config.zero_quantized_weights and not self.quantized_weights: - self.quantized_weights = _ds_config.zero_config.zero_quantized_weights - self.quantized_nontrainable_weights = zero_quantized_nontrainable_weights - if _ds_config is not None and _ds_config.zero_config.zero_quantized_nontrainable_weights and not self.quantized_nontrainable_weights: - self.quantized_nontrainable_weights = _ds_config.zero_config.zero_quantized_nontrainable_weights - - self.module = module - if (self.quantized_weights or self.quantized_nontrainable_weights): - self.quantizer_module = CUDAQuantizer() - print_rank_0(f'Using quantizer for weights: {self.quantizer_module.__class__.__name__}', force=True) - - if _ds_config is not None: - Init.override_module_apply = _ds_config.zero_config.override_module_apply - - if _ds_config.zero_config.offload_param is not None: - remote_device = _ds_config.zero_config.offload_param.device - pin_memory = _ds_config.zero_config.offload_param.pin_memory - - self._validate_remote_device(remote_device, _ds_config) - - # Remote device is the device where parameter partitions are stored - # It can be same as local_device or it could be CPU or NVMe. - self.remote_device = self.local_device if remote_device in [None, OffloadDeviceEnum.none] else remote_device - self.pin_memory = pin_memory if (self.remote_device in [OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme - ]) else False - - # Enable fp16 param swapping to NVMe - if self.remote_device == OffloadDeviceEnum.nvme: - self.param_swapper = param_swapper or AsyncPartitionedParameterSwapper(_ds_config, self.dtype) - else: - self.param_swapper = None - - # If we are provided an already-allocated module to prepare. - if module is not None: - assert isinstance(module, torch.nn.Module) - self._convert_to_zero_parameters(module.parameters(recurse=True)) - - self.use_all_gather_into_tensor = dist.has_all_gather_into_tensor() - if not self.use_all_gather_into_tensor: - logger.info(f"all_gather_into_tensor API is not available in torch {torch.__version__}") - - self.use_all_reduce_for_fetch_params = get_config_default(DeepSpeedZeroConfig, - "use_all_reduce_for_fetch_params") - if _ds_config is not None: - self.use_all_reduce_for_fetch_params = _ds_config.zero_config.use_all_reduce_for_fetch_params - - def _update_persist_config(self, ds_config): - Init.apply_param_persistence = True - Init.param_persistence_threshold = ds_config.zero_config.param_persistence_threshold - Init.model_persistence_threshold = ds_config.zero_config.model_persistence_threshold // self.num_partitions - - def _zero_init_param(self, param): - self._convert_to_deepspeed_param(param) - if dist.get_world_group() == self.get_dp_process_group(): - dist.broadcast(param.data, 0, self.get_dp_process_group()) - else: - dist.broadcast(param.data, dist.get_global_rank(self.get_dp_process_group(), 0), - self.get_dp_process_group()) - param.partition() - - def _convert_to_zero_parameters(self, param_list): - for param in param_list: - if is_zero_param(param): - continue - - param.data = param.data.to(self.local_device) - self._zero_init_param(param) - - def _validate_remote_device(self, remote_device, ds_config): - if ds_config is not None: - if remote_device in [None, OffloadDeviceEnum.cpu]: - if ds_config.zero_config.offload_param is not None: - offload_param_device = ds_config.zero_config.offload_param.device - assert offload_param_device != OffloadDeviceEnum.nvme, \ - f"'device' in DeepSpeed Config cannot be {offload_param_device} if remote device is {remote_device}." - - if remote_device == OffloadDeviceEnum.nvme: - assert ds_config.zero_config.offload_param is not None, \ - f'"offload_param" must be defined in DeepSpeed Config if remote device is {OffloadDeviceEnum.nvme}.' - - assert ds_config.zero_config.offload_param.nvme_path is not None, \ - f'"nvme_path" in DeepSpeed Config cannot be None if remote device is {OffloadDeviceEnum.nvme}' - - def _post_init_method(self, module): - #see_memory_usage(f"Before converting params in {module.__class__.__name__}", force=False) - print_rank_0(f'Converting Params in {module.__class__.__name__}', force=False) - see_memory_usage(f"Before converting and partitioning params in {module.__class__.__name__}", force=False) - - for name, param in module.named_parameters(recurse=False): - print_rank_0(f'Analyzing param {name} in {module.__class__.__name__}', force=False) - InsertPostInitMethodToModuleSubClasses.num_module_parameters += 1 - InsertPostInitMethodToModuleSubClasses.num_module_elements += param.numel() - if not is_zero_param(param): - if not get_accelerator().on_accelerator(param): - param.data = param.data.to(self.local_device) - - if name == 'weight' and self.quantized_initialization and type(module) in WEIGHT_QUANTIZATION_LAYERS: - _quantize_param(param, self.quantized_initialization) - - self._zero_init_param(param) - print_rank_0( - f"Partitioning param {debug_param2name_id_shape(param)} module={debug_module2name(module)}") - - see_memory_usage( - f"Param count {InsertPostInitMethodToModuleSubClasses.num_module_elements}. After converting and partitioning params in {module.__class__.__name__}", - force=False) - - def _convert_to_deepspeed_param(self, param): - - # Partitioned, Normal, Remote - param.ds_param_type = ZeroParamType.PARTITIONED - - # Replicated vs Partitioned vs Inflight - param.ds_status = ZeroParamStatus.AVAILABLE - - # Stores the shape of the original tensor - param.ds_shape = param.shape - - # Stores the number of elements in the original parameter without padding - param.ds_numel = param.numel() - - # Stores the partitioned copy of the tensor - param.ds_tensor = None - - # Keeps track of how many active sub-modules need this param at any given point in time - param.ds_active_sub_modules = set() - - # If this flag is true, then the parameters are replicated throughput training - # And only partitioned before the step - if Init.apply_param_persistence and param.ds_numel <= Init.param_persistence_threshold and Init.num_persisted_elements + param.ds_numel <= Init.model_persistence_threshold: - param.ds_persist = True - Init.num_persisted_parameters += 1 - Init.num_persisted_elements += param.ds_numel - else: - param.ds_persist = False - - param.is_external_param = False - - # The group that the parameter is scattered across. - param.ds_process_group = self.ds_process_group - - # Stores the secondary partitioned copy of the tensor - param.ds_secondary_tensor = None - - #Process group for secondary partition all (group) gather - param.ds_zero_param_process_group = self.zero_param_process_group - param.ds_secondary_tensor_group_size = self.num_ranks_in_param_group - param.ds_secondary_tensor_num_of_groups = self.num_param_groups - - # This is set to the Async Param swapper if remote device is nvme - # else this is set to None - param.nvme_swapper = self.param_swapper - - # DeepSpeed Param ID - param.ds_id = Init.param_id - Init.param_id += 1 - - def all_gather(param_list=None, async_op=False, hierarchy=0): - cls = param - if param_list is None: - param_list = [cls] - return self._all_gather(param_list, async_op=async_op, hierarchy=hierarchy) - - def _all_gather_dtype(params, world_size, rank_in_group, ds_process_group, allgather_dtype): - # make sure all params have the same dtype - dtype = params[0].dtype # we assume len(params) > 0 - assert all(p.dtype == dtype for p in params), "all params must have the same dtype" - - partition_sz = sum(p.ds_tensor.ds_numel for p in params) - - use_secondary_tensor = params[0].ds_secondary_tensor is not None - - if use_secondary_tensor: - partition_sz = sum(p.ds_tensor.ds_numel * p.ds_secondary_tensor_num_of_groups for p in params) - - flat_tensor = torch.empty(partition_sz * world_size, - dtype=allgather_dtype, - device=get_accelerator().current_device_name(), - requires_grad=False) - - partitions: List[Parameter] = [] - for i in range(world_size): - partitions.append(flat_tensor.narrow(0, partition_sz * i, partition_sz)) - - if use_secondary_tensor: - instrument_w_nvtx(torch.cat)([ - p.ds_secondary_tensor.to(get_accelerator().current_device_name()).to(allgather_dtype) - for p in params - ], - out=partitions[rank_in_group]) - else: - instrument_w_nvtx(torch.cat)( - [p.ds_tensor.to(get_accelerator().current_device_name()).to(allgather_dtype) for p in params], - out=partitions[rank_in_group]) - handle = _dist_allgather_fn(partitions[rank_in_group], flat_tensor, ds_process_group) - #Fix get_partition_dp_group(params[0])) - - return AllGatherCoalescedHandle( - allgather_handle=handle, - params=params, - partitions=partitions, - world_size=world_size, - use_secondary_tensor=use_secondary_tensor, - ) - - @instrument_w_nvtx - def all_gather_coalesced(params: Iterable[Parameter], - safe_mode: bool = False, - quantize: bool = False) -> AllGatherCoalescedHandle: - - # fetches from nvme if the partition is not available and in nvme - self._ensure_availability_of_partitioned_params(params) - - if self.num_partitions == 1: - return _no_gather_coalesced(params) - - for param in params: - if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: - raise RuntimeError(param.ds_summary()) - param.ds_status = ZeroParamStatus.INFLIGHT - - #use appropriate all gather process group - ds_process_group = self.ds_process_group - rank_in_group = self.rank - world_size = self.dp_world_size - use_secondary_tensor = params[0].ds_secondary_tensor is not None - if self.zero_param_process_group and use_secondary_tensor: - ds_process_group = self.zero_param_process_group #intragroup - rank_in_group = self.rank_in_group - world_size = self.num_ranks_in_param_group - - #pprint(dir(ds_process_group)) - # ensure that each rank has params in same order. the allgather - # is done by flattening the parameter list into a single tensor that - # can be allgathered in a single call - this means that if each rank - # gives a list of the same parameters in a different order we will - # silently get incorrect parameter values, and have very difficult - # to debug correctness issues. - params = sorted(params, key=lambda p: p.ds_id) - - if logger.isEnabledFor(logging.DEBUG): - debug_rank0(f"-allgather_coalesced: {[p.ds_id for p in params]}") - - if safe_mode: - # ensure that same list (with same ordering) of parameters are - # being allgathered across all ranks, otherwise could mix - # data between tensors. - assert_ints_same_as_other_ranks([p.ds_id for p in params]) - # ensure that tensors from each rank agree on the same ds_numel - # otherwise could mix data between tensors. - assert_ints_same_as_other_ranks([p.ds_tensor.ds_numel for p in params]) - - if len(params) == 1: - # have an opportunity to avoid some intermediate memory allocations - param = params[0] - buffer_size = math.ceil(param.ds_numel / world_size) * world_size - if use_secondary_tensor: - buffer_size = param.ds_secondary_tensor.shape[0] * world_size #make sure out is appropriately sized - - param_ds_tensor = param.ds_secondary_tensor if use_secondary_tensor else param.ds_tensor - - original_dtype = param_ds_tensor.dtype - if quantize: - allgather_dtype = torch.int8 - else: - allgather_dtype = get_allgather_dtype(param, param_ds_tensor) - - param_buffer = torch.empty( - buffer_size, - dtype=allgather_dtype, - device=get_accelerator().current_device_name(), - requires_grad=False, - ) - if not quantize: - handles = _dist_allgather_fn( - param_ds_tensor.to(get_accelerator().current_device_name()).to(allgather_dtype), - param_buffer, - ds_process_group, - ) - param.data = param_buffer.narrow(0, 0, param.ds_numel).view(param.ds_shape).to(param.device) - return AllGatherHandle(handles, param, original_dtype=original_dtype) - else: - if hasattr(param_ds_tensor, "ds_quant_scale"): - scales = param_ds_tensor.ds_quant_scale - quantized_param = param_ds_tensor.data - else: - quantized_param, scales = self.quantizer_module.quantize(param_ds_tensor) - handle = _dist_allgather_fn(quantized_param.to(get_accelerator().current_device_name()), - param_buffer, ds_process_group) - - quant_scale_buffer = torch.empty( - scales.numel() * world_size, - dtype=scales.dtype, - device=get_accelerator().current_device_name(), - requires_grad=False, - ) - quant_handle = _dist_allgather_fn(scales.to(get_accelerator().current_device_name()), - quant_scale_buffer, ds_process_group) - quant_info = QuantizationInfo() - quant_info.quantized_param = param_buffer.narrow(0, 0, param.ds_numel).view(param.ds_shape).to( - param.device) - quant_info.backend = self.quantizer_module - quant_info.quant_handle = quant_handle - quant_info.scale_buffer = quant_scale_buffer - return AllGatherHandle(handle, param, quantization=quant_info) - - else: - if self.use_all_reduce_for_fetch_params and not quantize and not use_secondary_tensor: - - # Use all_reduce instead of all_gather to fetch the module params - flat_buffer_size = sum(p.ds_numel_aligned for p in params) - flat_tensor = torch.zeros(flat_buffer_size, - dtype=get_only_unique_item(p.ds_tensor.dtype for p in params), - device=get_accelerator().current_device_name(), - requires_grad=False) - start_param = 0 - for param in params: - param.data = flat_tensor.narrow(0, start_param, param.ds_numel).view(param.ds_shape) - start = start_param + param.ds_tensor.ds_numel * self.get_partition_rank() - flat_tensor.narrow(0, start, param.ds_tensor.ds_numel).copy_(param.ds_tensor) - - start_param += param.ds_numel - - handle = dist.all_reduce(flat_tensor, group=ds_process_group, async_op=True) - - return AllReduceCoalescedHandle(handle=handle, params=params) - else: - if not quantize: - dtype_params = defaultdict(list) - for p in params: - allgather_dtype = get_allgather_dtype(p, p.ds_tensor) - dtype_params[allgather_dtype].append(p) - handles = [] - for dtype in sort_dtypes(dtype_params.keys()): - handles.append( - _all_gather_dtype(params, world_size, rank_in_group, ds_process_group, - allgather_dtype)) - - return MultipleAllGatherHandles(handles) - - else: - partition_sz = sum(p.ds_tensor.ds_numel for p in params) - - if use_secondary_tensor: - partition_sz = sum(p.ds_tensor.ds_numel * p.ds_secondary_tensor_num_of_groups - for p in params) - - flat_tensor = torch.empty(partition_sz * world_size, - dtype=torch.int8, - device=get_accelerator().current_device_name(), - requires_grad=False) - - if use_secondary_tensor: - if hasattr(params[0].ds_secondary_tensor, "ds_quant_scale"): - quantized_param = instrument_w_nvtx(torch.cat)([ - p.ds_secondary_tensor.data.to(get_accelerator().current_device_name()) - for p in params - ]) - scales = instrument_w_nvtx(torch.cat)([ - p.ds_secondary_tensor.ds_quant_scale.to(get_accelerator().current_device_name()) - for p in params - ]) - else: - quantized_param, scales = self.quantizer_module.quantize( - instrument_w_nvtx(torch.cat)([ - p.ds_secondary_tensor.to(get_accelerator().current_device_name()) - for p in params - ])) - else: - if hasattr(params[0].ds_tensor, "ds_quant_scale"): - quantized_param = instrument_w_nvtx(torch.cat)( - [p.ds_tensor.data.to(get_accelerator().current_device_name()) for p in params]) - scales = instrument_w_nvtx(torch.cat)([ - p.ds_tensor.ds_quant_scale.to(get_accelerator().current_device_name()) - for p in params - ]) - else: - quantized_param, scales = self.quantizer_module.quantize( - instrument_w_nvtx(torch.cat)( - [p.ds_tensor.to(get_accelerator().current_device_name()) for p in params])) - quant_scale_buffer = torch.empty( - scales.numel() * world_size, - dtype=torch.float32, - device=get_accelerator().current_device_name(), - requires_grad=False, - ) - handle = _dist_allgather_fn(quantized_param, flat_tensor, ds_process_group) - quant_handle = _dist_allgather_fn(scales, quant_scale_buffer, ds_process_group) - quant_info = QuantizationInfo() - quant_info.quantized_param = flat_tensor - quant_info.backend = self.quantizer_module - quant_info.quant_handle = quant_handle - quant_info.scale_buffer = quant_scale_buffer - quant_info.partition_sz = partition_sz - quant_info.world_size = world_size - return AllGatherCoalescedHandle( - allgather_handle=handle, - params=params, - partitions=None, - world_size=world_size, - use_secondary_tensor=use_secondary_tensor, - quantization=quant_info, - ) - - def partition(param_list=None, hierarchy=0, has_been_updated=False, free_data=True): - cls = param - print_rank_0(f"{'--'*hierarchy}----Partitioning param {debug_param2name_id_shape_device(cls)}", - force=False) - if param_list is None: - param_list = [cls] - self._partition(param_list, has_been_updated=has_been_updated, free_data=True) - - def reduce_gradients_at_owner(param_list=None, hierarchy=0): - cls = param - if param_list is None: - param_list = [cls] - print_rank_0( - f"{'--'*hierarchy}----Reducing Gradients for param with ids {[param.ds_id for param in param_list]} to owner" - ) - self._reduce_scatter_gradients(param_list) - - def partition_gradients(param_list=None, partition_buffers=None, hierarchy=0, accumulate=False): - cls = param - print_rank_0( - f"{'--'*hierarchy}----Partitioning param gradient with id {debug_param2name_id_shape_device(cls)}") - if param_list is None: - param_list = [cls] - if isinstance(partition_buffers, torch.Tensor): - partition_buffers = [partition_buffers] - - self._partition_gradients(param_list, partition_buffers=partition_buffers, accumulate=accumulate) - - def aligned_size(): - return self._aligned_size(param) - - def padding_size(): - return self._padding_size(param) - - def partition_numel(): - return self._partition_numel(param) - - def item_override(): - param.all_gather() - return param._orig_item() - - def ds_summary(slf: torch.Tensor, use_debug_name: bool = False) -> dict: - return { - "id": debug_param2name_id(slf) if use_debug_name else slf.ds_id, - "status": slf.ds_status.name, - "numel": slf.numel(), - "ds_numel": slf.ds_numel, - "shape": tuple(slf.shape), - "ds_shape": tuple(slf.ds_shape), - "requires_grad": slf.requires_grad, - "grad_shape": tuple(slf.grad.shape) if slf.grad is not None else None, - "persist": slf.ds_persist, - "active_sub_modules": slf.ds_active_sub_modules, - "ds_tensor.shape": slf.ds_tensor.shape if slf.ds_tensor is not None else None - } - - def convert_to_zero_parameters(param_list): - self._convert_to_zero_parameters(param_list) - - def allgather_before(func: Callable) -> Callable: - - def wrapped(*args, **kwargs): - param.all_gather() - return func(*args, **kwargs) - - return wrapped - - # Collectives for gathering and partitioning parameters - param.all_gather = all_gather - param.all_gather_coalesced = all_gather_coalesced - param.partition = partition - - # Collective for averaging gradients - param.reduce_gradients_at_owner = reduce_gradients_at_owner - param.partition_gradients = partition_gradients - - # Partitioning size utilities - param.aligned_size = aligned_size - param.padding_size = padding_size - param.partition_numel = partition_numel - param.ds_summary = types.MethodType(ds_summary, param) - - param.item = allgather_before(param.item) - - param.convert_to_zero_parameters = convert_to_zero_parameters - - def _aligned_size(self, param): - return param.ds_numel + self._padding_size(param) - - def _padding_size(self, param): - remainder = param.ds_numel % self.num_partitions - return (self.num_partitions - remainder) if remainder else 0 - - def _partition_numel(self, param): - return param.ds_tensor.ds_numel - - def _ensure_availability_of_partitioned_params(self, params): - swap_in_list = [] - swap_in_flight = [] - for param in params: - if param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE: - assert param.ds_tensor.final_location == OffloadDeviceEnum.nvme and param.ds_status == ZeroParamStatus.NOT_AVAILABLE - swap_in_list.append(param) - if param.ds_tensor.status == PartitionedParamStatus.INFLIGHT: - assert param.ds_tensor.final_location == OffloadDeviceEnum.nvme and param.ds_status == ZeroParamStatus.NOT_AVAILABLE - swap_in_flight.append(param) - if len(swap_in_list) > 0: - swap_in_list[0].nvme_swapper.swap_in(swap_in_list, async_op=False) - elif len(swap_in_flight) > 0: - swap_in_flight[0].nvme_swapper.synchronize_reads() - - @instrument_w_nvtx - def _all_gather(self, param_list, async_op=False, hierarchy=None): - - # fetches from nvme if the partition is not available and in nvme - self._ensure_availability_of_partitioned_params(param_list) - - handles = [] - all_gather_list = [] - for param in param_list: - if param.ds_status == ZeroParamStatus.NOT_AVAILABLE: - if async_op: - handle = self._allgather_param(param, async_op=async_op, hierarchy=hierarchy) - param.ds_status = ZeroParamStatus.INFLIGHT # if async_op else ZeroParamStatus.AVAILABLE - handles.append(handle) - else: - all_gather_list.append(param) - # note: param_list may contain params that are already in flight / aviailable. So we need to use all_gather_list - if not async_op: - if len(all_gather_list) == 1: - ret_value = self._allgather_params(all_gather_list, hierarchy=hierarchy) - else: - all_gather_quantize_list = [] - all_gather_nonquantize_list = [] - for param in all_gather_list: - if hasattr(param.ds_tensor, - "ds_quant_scale") or (hasattr(param, "ds_secondary_tensor") - and hasattr(param.ds_secondary_tensor, "ds_quant_scale")): - all_gather_quantize_list.append(param) - else: - all_gather_nonquantize_list.append(param) - # _allgather_params_coalesced always return None - self._allgather_params_coalesced(all_gather_nonquantize_list, hierarchy, quantize=False) - self._allgather_params_coalesced(all_gather_quantize_list, hierarchy, quantize=True) - for param in all_gather_list: - param.ds_status = ZeroParamStatus.AVAILABLE - return None - - return handles - - def _partition(self, param_list, force=False, has_been_updated=False, free_data=True): - for param in param_list: - print_rank_0(f"Before Partitioning Param {param.ds_id}", force=False) - if self.zero_param_process_group is not None: - self._partition_param_sec(param) - self._partition_param(param, has_been_updated=has_been_updated, free_data=True) - - param.ds_status = ZeroParamStatus.NOT_AVAILABLE - # if param.ds_tensor is not None: - # assert id(param.data) == id(param.ds_tensor.data), \ - # "After the parameters are initially partitioned, make sure we are not recreating the partition." - #print_rank_0(f"After Partitioning Param {param.ds_id} {param.ds_tensor.size()} {param.ds_tensor}",force=False) - @instrument_w_nvtx - def _partition_param(self, param, buffer=None, has_been_updated=False, free_data=True): - assert param.ds_status is not ZeroParamStatus.INFLIGHT, f" {param} Cannot partition a param in flight" - global reuse_buffers - print_rank_0(f"Param id {param.ds_id} status is {param.ds_status}", force=False) - if param.ds_status is ZeroParamStatus.AVAILABLE: - print_rank_0(f"Partitioning param id {param.ds_id} reuse buffers {reuse_buffers}", force=False) - # if reuse_buffers and False: - # numel = buffer.numel() - # buffer = param.data.view(-1) - # print_rank_0( - # "Returning buffer for param {param.ds_id} with numel {param.ds_numel} to empty buffers", - # force=False) - # if numel in empty_buffers: - # empty_buffers[numel].append(buffer) - - # if deepspeed.comm.get_rank(): - # print(f"Releasing {param.data.numel()}") - - if param.ds_tensor is not None and not has_been_updated: ##param already partitioned - - #print_rank_0(f"Param {param.ds_id} pri {param.ds_tensor.size()} loc? {param.ds_tensor.final_location}", force=True) - #param.data = param.ds_tensor.data - - see_memory_usage(f'Before partitioning param {param.ds_id} {param.shape}', force=False) - # param.data does not store anything meaningful in partitioned state - if free_data: - free_param(param) - see_memory_usage(f'After partitioning param {param.ds_id} {param.shape}', force=False) - - if param.ds_tensor.final_location == OffloadDeviceEnum.nvme: - print_rank_0(f"Param {param.ds_id} partition released since it exists in nvme", force=False) - param.nvme_swapper.remove_partition_and_release_buffers([param]) - print_rank_0( - f"after swap Param {param.ds_id} {param.ds_tensor.shape} partition released since it exists in nvme", - force=False) - - return - - tensor_size = self._aligned_size(param) - partition_size = tensor_size // self.num_partitions - if param.ds_tensor is None: - final_location = None - if self.remote_device == OffloadDeviceEnum.nvme and self.param_swapper.swappable_tensor( - numel=partition_size): - final_location = OffloadDeviceEnum.nvme - buffer = self.param_swapper.get_buffer(param, partition_size) - partitioned_tensor = torch.empty(0, dtype=param.dtype, device=buffer.device) - partitioned_tensor.data = buffer.data - print_rank_0(f"ID {param.ds_id} Initializing partition for the first time for nvme offload.") - - else: - if param.ds_persist: - device = self.local_device - elif self.remote_device == OffloadDeviceEnum.nvme: - device = OffloadDeviceEnum.cpu - else: - device = self.remote_device - - partitioned_tensor = torch.empty(partition_size, dtype=param.dtype, device=device) - # quantize the tensor if it's not trainable - if not param.requires_grad and self.quantized_nontrainable_weights: - partitioned_tensor, partitioned_tensor.ds_quant_scale = self.quantizer_module.quantize( - partitioned_tensor) - - if device == OffloadDeviceEnum.cpu and self.pin_memory: - partitioned_tensor = get_accelerator().pin_memory(partitioned_tensor) - - partitioned_tensor.requires_grad = False - param.ds_tensor = partitioned_tensor - param.ds_tensor.ds_numel = partition_size - param.ds_tensor.status = PartitionedParamStatus.AVAILABLE - param.ds_tensor.final_location = final_location - param.ds_numel_aligned = tensor_size - - start = partition_size * self.get_partition_rank() - end = start + partition_size - - one_dim_param = param.contiguous().view(-1) - - if start < param.ds_numel and end <= param.ds_numel: - src_tensor = one_dim_param.narrow(0, start, partition_size) - - with torch.no_grad(): - # make sure param.ds_tensor requires_grad always be false, - # otherwise, torch tracer will complain. - param.ds_tensor.copy_(src_tensor) - - #partitioned_tensor = src_tensor.clone().detach().to(self.remote_device) - - else: - # partitioned_tensor = torch.zeros(partition_size, - # dtype=param.dtype, - # device=self.remote_device ) - - if start < param.ds_numel: - elems_to_copy = param.ds_numel - start - with torch.no_grad(): - # make sure param.ds_tensor requires_grad always be false, - # otherwise, torch tracer will complain. - param.ds_tensor.narrow(0, 0, - elems_to_copy).copy_(one_dim_param.narrow(0, start, elems_to_copy)) - - #print(f"Remote device {self.remote_device}") - - #param.ds_tensor = partitioned_tensor - - #param.data = param.ds_tensor.data - - # param.data does not store anything meaningful in partitioned state - - see_memory_usage(f'Before partitioning param {param.ds_id} {param.shape}', force=False) - free_param(param) - see_memory_usage(f'After partitioning param {param.ds_id} {param.shape}', force=False) - - if param.ds_tensor.final_location == OffloadDeviceEnum.nvme: - self.param_swapper.swap_out_and_release([param]) - print_rank_0(f"ID {param.ds_id} Offloaded to nvme offload and buffers released.") - see_memory_usage(f"ID {param.ds_id} Offloaded to nvme offload and buffers released.", force=False) - - print_rank_0(f"ID {param.ds_id} partitioned type {param.dtype} dev {param.device} shape {param.shape}") - - @instrument_w_nvtx - def _partition_param_sec(self, param, buffer=None, has_been_updated=False): - assert param.ds_status is not ZeroParamStatus.INFLIGHT, f" {param} Cannot partition a param in flight" - global reuse_buffers - ##support for NVME secondary param offload - #print_rank_0(f"SEC Param id {param.ds_id} status is {param.ds_status}", force=True) - if param.ds_status is ZeroParamStatus.AVAILABLE: - if param.ds_secondary_tensor is not None and not has_been_updated: ##param already partitioned - return - #check padding - tensor_size = self._aligned_size(param) - partition_size = tensor_size // self.dp_world_size - - secondary_partition_size = int(tensor_size // self.num_ranks_in_param_group) - if param.ds_secondary_tensor is None: - final_location = None - secondary_partitioned_tensor = torch.empty(secondary_partition_size, - dtype=param.dtype, - device=self.remote_device) - - if self.pin_memory: - secondary_partitioned_tensor = secondary_partitioned_tensor.pin_memory() - # quantize the tensor if it's not trainable - if not param.requires_grad and self.quantized_nontrainable_weights: - secondary_partitioned_tensor, secondary_partitioned_tensor.ds_quant_scale = self.quantizer_module.quantize( - secondary_partitioned_tensor) - secondary_partitioned_tensor.requires_grad = False - param.ds_secondary_tensor = secondary_partitioned_tensor - param.ds_secondary_tensor.ds_numel = secondary_partition_size - param.ds_secondary_tensor.status = PartitionedParamStatus.AVAILABLE - param.ds_secondary_tensor.final_location = final_location - - #use rank in group for secondary tensor - secondary_start = secondary_partition_size * self.rank_in_group - - secondary_end = secondary_start + secondary_partition_size - - one_dim_param = param.contiguous().view(-1) - - # ds_numel is unpadded, so the last chunk of the secondary tensor might not be secondary_partition_size - sec_numel = max(0, min(param.ds_numel - secondary_start, secondary_partition_size)) - - # copy from full tensor to secondary tensor - param.ds_secondary_tensor.narrow(0, 0, - sec_numel).copy_(one_dim_param.narrow(0, secondary_start, sec_numel)) - - # TODO: This is a temporary fix to avoid the issue that 2nd tensor all-gather happens before 2nd tensor partition is done - if not get_accelerator().resolves_data_dependency(): - get_accelerator().current_stream().synchronize() - - print_rank_0(f"{param.ds_id} partitioned type {param.dtype} dev {param.device} shape {param.shape}", - force=False) - - def _param_status(self, param): - if param.ds_tensor is not None: - print_rank_0( - f"Param id {param.ds_id}, param status: {param.ds_status}, param numel {param.ds_numel}, partitioned numel {param.ds_tensor.numel()}, data numel {param.data.numel()}" - ) - else: - print_rank_0( - f"Param id {param.ds_id}, param status: {param.ds_status}, param numel {param.ds_numel}, partitioned ds_tensor {param.ds_tensor}, data numel {param.data.numel()}" - ) - - def _allgather_param(self, param, async_op=False, hierarchy=0): - - partition_size = param.ds_tensor.ds_numel - - tensor_size = partition_size * self.num_partitions - aligned_param_size = self._aligned_size(param) - assert tensor_size == aligned_param_size, f'param id {param.ds_id} aligned size {aligned_param_size} does not match tensor size {tensor_size}' - - print_rank_0( - f"{'--'* hierarchy}---- Before allocating allgather param {debug_param2name_id_shape_status(param)} partition size={partition_size}" - ) - - see_memory_usage( - f'Before allocate allgather param {debug_param2name_id_shape_status(param)} partition_size={partition_size} ', - force=False) - flat_tensor = torch.zeros(aligned_param_size, dtype=param.dtype, device=param.device).view(-1) - see_memory_usage( - f'After allocate allgather param {debug_param2name_id_shape_status(param)} {aligned_param_size} {partition_size} ', - force=False) - - if not get_accelerator().resolves_data_dependency(): - get_accelerator().synchronize() - - print_rank_0( - f"{'--'* hierarchy}----allgather param with {debug_param2name_id_shape_status(param)} partition size={partition_size}" - ) - # if not flat_tensor.numel() > 100000: - # replicated_tensor = flat_tensor.narrow(0, - # 0, - # param.ds_numel).view(param.ds_shape) - # param.data = replicated_tensor.data - # return None - if self.use_all_gather_into_tensor: - handle = dist.all_gather_into_tensor(flat_tensor, - param.ds_tensor.to(get_accelerator().device_name()), - group=self.get_partition_dp_group(param), - async_op=async_op) - else: - partitions = [] - for i in range(self.num_partitions): - partitions.append(flat_tensor.narrow(0, partition_size * i, partition_size)) - - if i == dist.get_rank(group=self.get_partition_dp_group(param)): - partitions[i].data.copy_(param.ds_tensor.data, non_blocking=True) - - handle = dist.all_gather(partitions, - partitions[self.get_partition_rank()], - group=self.get_partition_dp_group(param), - async_op=async_op) - - replicated_tensor = flat_tensor.narrow(0, 0, param.ds_numel).view(param.ds_shape) - param.data = replicated_tensor.data - return handle - - def _allgather_params_coalesced(self, param_list, hierarchy=0, quantize=False): - """ blocking call - avoid explicit memory copy in _allgather_params - """ - if len(param_list) == 0: - return - - if self.num_partitions == 1: - handle = _no_gather_coalesced(param_list) - handle.wait() - return None - - # collect local tensors and partition sizes - partition_sizes = [] - local_tensors = [] - if quantize: - quantize_scale_sizes = [] - quantize_scale_tensors = [] - for param in param_list: - partition_sizes.append(param.ds_tensor.ds_numel) - local_tensors.append(param.ds_tensor.to(get_accelerator().device_name())) - if quantize: - quantize_scale_sizes.append(param.ds_tensor.ds_quant_scale.numel()) - quantize_scale_tensors.append(param.ds_tensor.ds_quant_scale.to(get_accelerator().device_name())) - # allocate memory for allgather params - allgather_params = [] - if quantize: - allgather_quantize_scale = [] - for psize in partition_sizes: - tensor_size = psize * self.num_partitions - flat_tensor = torch.empty(tensor_size, dtype=param_list[0].ds_tensor.dtype, - device=self.local_device).view(-1) - flat_tensor.requires_grad = False - allgather_params.append(flat_tensor) - if quantize: - for psize in quantize_scale_sizes: - tensor_size = psize * self.num_partitions - flat_tensor = torch.empty(tensor_size, - dtype=param_list[0].ds_tensor.ds_quant_scale.dtype, - device=self.local_device).view(-1) - flat_tensor.requires_grad = False - allgather_quantize_scale.append(flat_tensor) - - # launch - launch_handles = [] - launch_quantize_handles = [] - for param_idx, param in enumerate(param_list): - input_tensor = local_tensors[param_idx].view(-1) - - if self.use_all_gather_into_tensor: - # try the _all_gather_base from Pytorch master - h = dist.all_gather_into_tensor(allgather_params[param_idx], - input_tensor, - group=self.get_partition_dp_group(param), - async_op=True) - if quantize: - quantize_handle = dist.all_gather_into_tensor(allgather_quantize_scale[param_idx], - quantize_scale_tensors[param_idx], - group=self.get_partition_dp_group(param), - async_op=True) - launch_quantize_handles.append(quantize_handle) - else: - output_list = [] - for i in range(self.num_partitions): - psize = partition_sizes[param_idx] - partition = allgather_params[param_idx].narrow(0, i * psize, psize) - output_list.append(partition) - if not get_accelerator().on_accelerator(partition): - logger.warning( - f'param {param_idx}, partition {i} is not on CUDA, partition shape {partition.size()}') - - # back to old all_gather function - h = dist.all_gather(output_list, input_tensor, group=self.get_partition_dp_group(param), async_op=True) - if quantize: - output_scale_list = [] - for i in range(self.num_partitions): - psize = quantize_scale_sizes[param_idx] - partition = allgather_quantize_scale[param_idx].narrow(0, i * psize, psize) - output_scale_list.append(partition) - quant_handle = dist.all_gather(output_scale_list, - quantize_scale_tensors[param_idx], - group=self.get_partition_dp_group(param), - async_op=True) - launch_quantize_handles.append(quant_handle) - launch_handles.append(h) - - # Wait ensures the operation is enqueued, but not necessarily complete. - launch_handles[-1].wait() - if quantize: - for quant_handle in launch_quantize_handles: - quant_handle.wait() - - # assign to param.data (not copy) - for i, param in enumerate(param_list): - gathered_tensor = allgather_params[i] - if quantize: - gathered_tensor = self.quantizer_module.dequantize(gathered_tensor, allgather_quantize_scale[i]) - param.data = gathered_tensor.narrow(0, 0, param.ds_numel).view(param.ds_shape).data - - # guarantee the communication to be completed - if not get_accelerator().resolves_data_dependency(): - get_accelerator().synchronize() - - return None - - @torch.no_grad() - def _allgather_params(self, param_list, hierarchy=0): - if len(param_list) == 0: - return - - partition_size = sum([param.ds_tensor.ds_numel for param in param_list]) - - tensor_size = partition_size * self.num_partitions - flat_tensor = torch.empty(tensor_size, dtype=param_list[0].ds_tensor.dtype, device=self.local_device) - partitions = [] - for i in range(self.num_partitions): - start = partition_size * i - - partitions.append(flat_tensor.narrow(0, start, partition_size)) - - if i == self.get_partition_rank(): - offset = 0 - for param in param_list: - param_numel = param.ds_tensor.ds_numel - - partitions[i].narrow(0, offset, param_numel).copy_(param.ds_tensor.data) - - offset += param_numel - - if hasattr(param_list[0], 'ds_quant_scale'): - scale_size = sum([param.ds_tensor.ds_quant_scale.numel() for param in param_list]) - scale_tensor_size = scale_size * self.world_size - flat_scale_tensor = torch.empty(scale_tensor_size, - dtype=param_list[0].ds_tensor.ds_quant_scale.dtype, - device=self.local_device) - scale_partitions = [] - for i in range(self.world_size): - start = scale_tensor_size * i - scale_partitions.append(flat_scale_tensor.narrow(0, start, scale_tensor_size)) - if i == self.rank: - offset = 0 - for param in param_list: - param_scale_numel = param.ds_tensor.ds_quant_scale.ds_numel - - scale_partitions[i].narrow(0, offset, - param_scale_numel).copy_(param.ds_tensor.ds_quant_scale.data) - - offset += param_scale_numel - - dist.all_gather_into_tensor(flat_tensor, - partitions[self.get_partition_rank()], - group=self.get_partition_dp_group(param), - async_op=False) - if hasattr(param_list[0], 'ds_quant_scale'): - dist.all_gather(flat_scale_tensor, - param_list[0].ds_quant_scale, - group=self.get_partition_dp_group(param), - async_op=False) - param_offset = 0 - - for param in param_list: - param_partition_size = param.ds_tensor.ds_numel - param_size = param.ds_numel - replicated_tensor = torch.empty(param.ds_shape, dtype=param.ds_tensor.dtype, device=self.local_device) - - for i in range(self.num_partitions): - - start = i * partition_size - - param_start = i * param_partition_size - - if param_start < param_size: - numel_to_copy = min(param_size - param_start, param_partition_size) - - part_to_copy = partitions[i].narrow(0, param_offset, numel_to_copy) - - replicated_tensor.view(-1).narrow(0, param_start, numel_to_copy).copy_(part_to_copy) - #param_offset += param.data.numel() - param_offset += param.ds_tensor.ds_numel - if hasattr(param_list[0], 'ds_quant_scale'): - replicated_tensor = self.quantizer_module.dequantize(replicated_tensor, flat_scale_tensor) - param.data = replicated_tensor.data - - return None - - def _reduce_scatter_gradients(self, param_list): - #print_rank_0([param.grad for param in param_list]) - #assert any([param.grad is None for param in param_list]), "None gradients cannot be reduce scattered" - - handles_and_reduced_partitions = [] - for param in param_list: - assert param.grad.numel( - ) == param.ds_numel, f"{param.grad.numel()} != {param.ds_numel} Cannot reduce scatter gradients whose size is not same as the params" - - handles_and_reduced_partitions.append(self._reduce_scatter_gradient(param)) - - for param, (handle, reduced_partition) in zip(param_list, handles_and_reduced_partitions): - if handle is not None: - handle.wait() - - # some ranks may have partitions that are padded to go beyond the grad size. - # For these ranks the output of reduce scatter is a separate buffer and needs - # to be copied in - partition_size = param.ds_tensor.ds_numel - start = self.get_partition_rank() * partition_size - end = start + partition_size - #print_rank_0("REduce scatter was executed for param {param.ds_id}") - if start < param.ds_numel < end: - elements = param.ds_numel - start - param.grad.view(-1).narrow(0, start, elements).copy_(reduced_partition.narrow(0, 0, elements)) - - def _reduce_scatter_gradient(self, param): - - partition_size = param.ds_tensor.ds_numel - #output = torch.empty(partition_size, dtype=param.dtype, device=param.device) - - total_size = partition_size * self.num_partitions - input_list = [] - - for i in range(self.num_partitions): - - start = i * partition_size - end = start + partition_size - - #print("before reduce scatter gradients") - if start < param.ds_numel and end <= param.ds_numel: - input = param.grad.view(-1).narrow(0, start, partition_size) - else: - input = torch.zeros(partition_size, dtype=param.dtype, device=param.device) - - if start < param.ds_numel: - elements = param.ds_numel - start - input.narrow(0, 0, elements).copy_(param.grad.view(-1).narrow(0, start, elements)) - #print("after reduce scatter gradients") - input_list.append(input) - - rank = dist.get_rank(group=self.get_partition_dp_group(param)) - handle = dist.reduce_scatter(input_list[rank], - input_list, - group=self.get_partition_dp_group(param), - async_op=True) - - return handle, input_list[rank] - - def _partition_gradients(self, param_list, partition_buffers=None, accumulate=False): - if partition_buffers is None: - partition_buffers = [None] * len(param_list) - - for param, partition_buffer in zip(param_list, partition_buffers): - self._partition_gradient(param, partition_buffer=partition_buffer, accumulate=accumulate) - - def _partition_gradient(self, param, partition_buffer=None, accumulate=False): - - #import pdb;pdb.set_trace() - # param.grad=None - # param.grad.test() - print_rank_0( - f"Partitioning param {param.ds_id} gradient of size {param.grad.numel()} type {param.grad.dtype} part_size {param.ds_tensor.ds_numel}" - ) - see_memory_usage("Before partitioning gradients", force=False) - partition_size = param.ds_tensor.ds_numel - - if partition_buffer is None: - assert not accumulate, "No buffer to accumulate to" - partition_buffer = torch.zeros(partition_size, dtype=param.dtype, device=param.device) - else: - assert partition_buffer.numel( - ) >= partition_size, f"The partition buffer size {partition_buffer.numel()} should match the size of param.ds_tensor {partition_size}" - - rank = dist.get_rank(group=self.get_partition_dp_group(param)) - start = partition_size * rank - end = start + partition_size - - dest_tensor_full_buffer = partition_buffer.view(-1).narrow(0, 0, partition_size) - - #print("before partition gradients") - if start < param.ds_numel: - elements = min(param.ds_numel - start, partition_size) - - dest_tensor = dest_tensor_full_buffer.narrow(0, 0, elements) - src_tensor = param.grad.view(-1).narrow(0, start, elements) - - # just copy the grad partition to the buffer - if not accumulate: - dest_tensor.copy_(src_tensor) - - # if source and destination are on same device, - # add to the provided buffer - elif src_tensor.device == dest_tensor.device: - dest_tensor.add_(src_tensor) - - # if source and destination are on different device, copy first to src - # then add and move back to the destination. This seems to run faster - # when src is gpu and dest is cpu - # adding directly to cpu is very slow - else: - acc_tensor = torch.empty(src_tensor.numel(), dtype=param.dtype, device=param.device) - - acc_tensor.copy_(dest_tensor) - acc_tensor.add_(src_tensor) - dest_tensor.copy_(acc_tensor) - - # partition_buffer.view(-1).narrow( - # 0, - # 0, - # elements).copy_(param.grad.view(-1).narrow(0, - # start, - # elements)) - - #print("after partition gradients") - param.grad.data = dest_tensor_full_buffer.data - see_memory_usage("After partitioning gradients", force=False) - - def get_partition_dp_group(self, param): - return param.ds_process_group - - def get_partition_rank(self): - """subclass can overload to specify different relative rank in - parameter partition group""" - return self.rank - - @property - def num_partitions(self): - return self.dp_world_size - - def get_dp_process_group(self): - """ Return the communication group with all data-parallel ranks """ - return self.ds_process_group - - -class GatheredParameters: - - def __init__(self, params, modifier_rank=None, fwd_module=None, enabled=True): - """A context that collects parameters that were partitioned via a - :class:`deepspeed.zero.Init` context. The parameters are partitioned - again upon exit. - - Args: - params (``torch.nn.Parameter``): A single parameter, or an iterable of parameters (list, tuple, generator) of parameters to collect. - It's assumed that all parameters are zero params. - modifier_rank (int, optional): If specified, this rank's parameter will be - broadcasted on exit from the context. This argument is required if ``params`` are - modified, so that all processes have a consistent view of the data. Defaults - to ``None``. - fwd_module (``torch.nn.Module``, optional): If specified, ``params`` will be - registered as external parameters of ``fwd_module``. See :meth:`deepspeed.zero.register_external_parameter`. - enabled (bool, optional): If ``False``, this context is a no-op. Defaults to ``True``. - - Important: Make sure to use ``modifier_rank`` that is not ``None`` (e.g., ``modifier_rank=0``) - if you need the GPU memory allocated by gather to be released upon exit from the context manager. - - Important: if ``params`` isn't an iterable of parameters or a single parameter it'll be silently ignored! - - Examples - ======== - - #. Allocate a partitioned module, initialize its weight on rank 0, and update all - processes. - - .. code-block:: python - - with deepspeed.zero.Init(): - linear = torch.nn.Linear(1000,1000) - - with deepspeed.zero.GatheredParameters(linear.weight, - modifier_rank=0): - if deepspeed.comm.get_rank() == 0: - linear.weight.zero_() - - with deepspeed.zero.GatheredParameters(linear.weight, - modifier_rank=0): - if deepspeed.comm.get_rank() == 0: - linear.weight.zero_() - - #. Collect a partitioned weight to pass to another module during - training. The parameter will be registered as an external parameter - and made available during the backward pass. - - .. code-block:: python - :emphasize-lines: 6 - - def forward(self, input): - x = self.layer1(input) - - # self.layer1.weight is required by self.layer2.forward - with deepspeed.zero.GatheredParameters(self.layer1.weight, - fwd_module=self): - y = self.layer2(x, self.layer1.weight) - return y - - - #. Pretrained model loading - - .. code-block:: python - - with deepspeed.zero.Init(): - model = MyModel() - - state_dict = torch.load(model_path, map_location="cpu") - - def load(module: nn.Module, prefix=""): - # because zero3 puts placeholders in model params, this context - # manager gathers (unpartitions) the params of the current layer, then loads from - # the state dict and then re-partitions them again - with deepspeed.zero.GatheredParameters(list(module.parameters(recurse=False)), modifier_rank=0): - if deepspeed.comm.get_rank() == 0: - module._load_from_state_dict(state_dict, prefix) - - for name, child in module._modules.items(): - if child is not None: - load(child, prefix + name + ".") - - load(model, prefix="") - - If this approach is not used, then the full model will first be copied to each GPU. For models - bigger than the memory of a single GPU, this method is required. - """ - - self.enabled = enabled - if not enabled: - return - - if isinstance(params, Iterable) and not isinstance(params, torch.Tensor): - # deal with generators like model.parameters() - # must convert to list to be able to iterate more than once if we get a generator - params = list(params) - else: - # single param - params = [params] - # enable if at least one is zero-param, otherwise a noop - if not any(is_zero_param(p) for p in params): - self.enabled = False - return - - self.params = [p for p in params if hasattr(p, "ds_id")] - self.params = sorted( - set(self.params), key=lambda x: x.ds_id - ) # remove the duplicates to prevent racing condition, we must also make sure the order is the same on all ranks otherwise we'll get deadlocks - self.src_rank = None - if modifier_rank is not None: - if self.params[0].ds_process_group == dist.get_world_group(): - self.src_rank = modifier_rank - else: - # A group was specified; convert DP rank to global rank - self.src_rank = dist.get_global_rank(self.params[0].ds_process_group, modifier_rank) - self.fwd_module = fwd_module - if self.fwd_module is not None: - # is a no-op if already registered - for p in self.params: - register_external_parameter(self.fwd_module, p) - - def __enter__(self): - if not self.enabled: - return - self.params[0].all_gather(param_list=self.params) - - def __exit__(self, *exc): - if not self.enabled: - return - if self.src_rank is None: - self.params[0].partition(param_list=self.params, has_been_updated=False) - return - - handles = [dist.broadcast(p.data, self.src_rank, group=p.ds_process_group, async_op=True) for p in self.params] - for h in handles: - h.wait() - self.params[0].partition(param_list=self.params, has_been_updated=True) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py deleted file mode 100644 index 80b20cb1d3e883d91a5dce95a16a496c42b9faef..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py +++ /dev/null @@ -1,585 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from dataclasses import dataclass -import collections -from collections import UserDict -from typing import Deque, Set - -from deepspeed import comm as dist -from deepspeed.utils import z3_leaf_module -from deepspeed.utils.logging import logger -from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum -from deepspeed.runtime.zero.partition_parameters import * -from deepspeed.runtime.zero.partitioned_param_profiler import PartitionedParameterProfiler -from deepspeed.runtime.swap_tensor.partitioned_param_swapper import PartitionedParamStatus -from deepspeed.utils.debug import debug_param2name_id_shape -from deepspeed.accelerator import get_accelerator -import deepspeed.runtime.compiler as compiler -from deepspeed.runtime.compiler import is_compiling - -import logging - -ENABLE_PROFILER = False - - -def debug_rank0(message: str) -> None: - if dist.get_rank() == 0: - logger.debug(message) - - -@instrument_w_nvtx -def get_all_parameters(sub_module, recurse=False): - return itertools.chain(sub_module.named_parameters(recurse=recurse), sub_module.ds_external_parameters()) - - -@compiler.disable -def iter_params(module: Module, recurse=False) -> Iterable[Parameter]: - return map(lambda pair: pair[1], get_all_parameters(module, recurse)) - - -class ZeRoTraceMode(Enum): - # Record trace of the network during a single forward+backward (for training) or forward (for inference) - RECORD = 1 - # Use recorded network trace to optimize current forward+backward or forward - COMPLETE = 2 - # Recorded trace does not match current forward+backward or forward pass. - INVALID = 3 - - -class InflightParamRegistry(UserDict): - """registry for parameters in flight""" - - def __setitem__(self, param: Parameter, handle: AllGatherCoalescedHandle) -> None: - if param in self.data: - raise RuntimeError(f"{param.ds_summary()} already in registry") - if param.ds_status != ZeroParamStatus.INFLIGHT: - raise RuntimeError(f"attempted to add non-inflight parameter to registry {param.ds_summary()}") - self.data[param] = handle - - -class PartitionedParameterCoordinator: - FORWARD_FETCH_SUBMIT = 'forward_fetch_submit' - FORWARD_FETCH_WAIT = 'forward_fetch_wait' - FORWARD_PREFETCH_SUBMIT = 'forward_prefetch_submit' - BACKWARD_FETCH_SUBMIT = 'backward_fetch_submit' - BACKWARD_FETCH_WAIT = 'backward_fetch_wait' - BACKWARD_PREFETCH_SUBMIT = 'backward_prefetch_submit' - FORWARD_ALL_GATHER = 'forward_all_gather' - BACKWARD_ALL_GATHER = 'backward_all_gather' - """Handles partitioning and gathering of parameters.""" - - @dataclass - class __ParamInTrace: - param: Parameter - step_id_last_used_at: int - - def __init__( - self, - prefetch_bucket_sz: int, - max_reuse_distance_in_numel: int, - max_available_parameters_in_numel: int, - allgather_stream: get_accelerator().Stream, - inflight_param_registry: InflightParamRegistry, - prefetch_nvme: bool = False, - timers=None, - zero_quantized_weights=False, - zero_quantized_nontrainable_weights=False, - fast_sharding_for_leaf_module=False, - log_trace_cache_warnings=False, - ) -> None: - # mapping of param -> handle for each param that is currently in flight - self.__inflight_param_registry = inflight_param_registry - # keeps track of the number of submodules invoked so far. - self.__step_id: int = 0 - # network tracing mode - self.__trace_mode: ZeRoTraceMode = ZeRoTraceMode.INVALID - # sequence of submodules/parameters in forward pass + backward pass - self.__submodule_order: Iterable[Module] = [] - self.__param_order: Iterable[__class__.__ParamInTrace] = [] - self.__most_recent_step_id_param_fetched_for = collections.defaultdict(lambda: int(-1e10)) - self.__step_id_module_fetched_for = collections.defaultdict(lambda: collections.deque()) - # number of available params, and max number of available params - self.__n_available_params: int = 0 - self.__max_n_available_params: int = max_available_parameters_in_numel - # max distance between two use of the module beyond which module is released - self.__max_reuse_dist_in_numel: int = max_reuse_distance_in_numel - # queue for parameters to fetch. parameters will be popped off the left - # side of the dequeue as they are fetched - self.__param_queue: Deque[__class__.__ParamInTrace] = None - self.__prefetch_bucket_sz: int = prefetch_bucket_sz - self.__prefetch_nvme: bool = prefetch_nvme - self.hierarchy: int = 0 - self.zero_quantized_weights = zero_quantized_weights - self.zero_quantized_nontrainable_weights = zero_quantized_nontrainable_weights - - # stream that will be used for allgather operations - self.__allgather_stream: get_accelerator().Stream = allgather_stream - - # limit the number of fetch events that can be queued at once - # otherwise, what happens is memory is allocated by the host thread at the - # time of the call, but not used until later by the asynchronous cuda stream. - # allowing an infinite number of these to queue up causes a lot of memory - # pressure that then becomes detrimental to performance. - # this is a much less elegant way of fixing this vs something like using - # cudaMallocAsync/cudaFreeAsync. Choosing to not expose this to the user now - # because ideally in the future its replaced by an async allocation - # mechanism which doesn't require any configuration by the user. - self.__ongoing_fetch_events: Deque[get_accelerator().Event] = collections.deque() - # TODO. make this configurable via JSON - self.__max_ongoing_fetch_events: int = 2 - self.__profiler = PartitionedParameterProfiler(timers if ENABLE_PROFILER else None) - - # Whether to log trace cache warnings, e.g. invalidation events - self.__log_trace_cache_warnings = log_trace_cache_warnings - - # whether to enable fast fetch for the z3 leaf module. - # this will improve fetch speed but will not break down leaf module parameters to alleviate memory pressure. - self.fast_sharding_for_leaf_module = fast_sharding_for_leaf_module - - """Tracing and Tracking - TODO. consider performing trace before initializing PartitionedParameterCoordinator - and passing trace results into constructor. This way all the code in here can - just assume that the trace is complete and the results can be entirely - immutable. - - Bookkeeping operations used to track where we are in the forward/backward pass - """ - - def _clear_trace_structures(self) -> None: - self.__submodule_order = [] - self.__param_order = [] - self.__most_recent_step_id_param_fetched_for = collections.defaultdict(lambda: int(-1e10)) - self.__param_queue = None - - def is_complete_trace(self) -> bool: - return self.__trace_mode == ZeRoTraceMode.COMPLETE - - def is_invalid_trace(self) -> bool: - return self.__trace_mode == ZeRoTraceMode.INVALID - - def is_record_trace(self) -> bool: - return self.__trace_mode == ZeRoTraceMode.RECORD - - def _clean_inflight_param_registry(self) -> None: - for param, handle in self.__inflight_param_registry.items(): - handle.wait() - self.__release_param(param) - self.__inflight_param_registry.clear() - - def _invalidate_trace(self) -> None: - if self.is_invalid_trace(): - raise RuntimeError("attempted to invalidate already invalid trace") - self.__trace_mode = ZeRoTraceMode.INVALID - self._clear_trace_structures() - self._clean_inflight_param_registry() - - def trace_prologue(self, sub_module: Module) -> None: - if self.is_complete_trace(): - # sub_module must match expectation else invalidate trace cache - if len(self.__submodule_order) <= self.__step_id: - print_rank_0( - f"Invalidate trace cache @ step {self.__step_id} and module {sub_module.ds_id}: " - f"cache has only {len(self.__submodule_order)} modules", - force=self.__log_trace_cache_warnings) - self._invalidate_trace() - return - - if sub_module != self.__submodule_order[self.__step_id]: - expected_module_id = self.__submodule_order[self.__step_id].ds_id - print_rank_0( - f"Invalidate trace cache @ step {self.__step_id}: " - f"expected module {expected_module_id}, but got module {sub_module.ds_id}", - force=self.__log_trace_cache_warnings) - self._invalidate_trace() - - @compiler.disable - def record_module(self, sub_module: Module) -> None: - """adds sub module to trace""" - if is_compiling(): - return - - if not self.is_record_trace(): - raise RuntimeError(f"attempted to record trace when status = {self.__trace_mode}") - - self.__submodule_order.append(sub_module) - self.__step_id_module_fetched_for[sub_module.ds_id].append(self.__step_id) - - def record_parameters(self, sub_module: Module) -> None: - if is_compiling(): - return - """adds sub module to trace""" - if not self.is_record_trace(): - raise RuntimeError(f"attempted to record trace when status = {self.__trace_mode}") - - step_id = self.__step_id_module_fetched_for[sub_module.ds_id].popleft() - for param in sorted(set(iter_params(sub_module, recurse=z3_leaf_module(sub_module))), key=lambda p: p.ds_id): - self.__param_order.append(__class__.__ParamInTrace(param=param, step_id_last_used_at=step_id)) - - def construct_parameter_trace_from_module_trace(self): - """use module trace to construct parameter trace""" - self.__param_order = [] - for sub_module in self.__submodule_order: - self.record_parameters(sub_module) - - @compiler.disable - def reset_step(self) -> None: - """indicate that we have completed one fwd+bwd for the model""" - if is_compiling(): - return - - self._clean_inflight_param_registry() - - if not self.is_complete_trace(): # not self.trace_complete: - # Make sure that recorded submodule orders are identical across ranks - assert_ints_same_as_other_ranks([m.ds_id for m in self.__submodule_order]) - - if self.is_record_trace(): - # Successfully recorded a trace - self.construct_parameter_trace_from_module_trace() - # Make sure that recorded parameter orders are identical across ranks - assert_ints_same_as_other_ranks([p.param.ds_id for p in self.__param_order]) - assert_ints_same_as_other_ranks([p.step_id_last_used_at for p in self.__param_order]) - - self.__submodule_order = tuple(self.__submodule_order) # freeze - self.__param_order = tuple(self.__param_order) # freeze - self.__trace_mode = ZeRoTraceMode.COMPLETE - print_rank_0( - f"completed record trace of {len(self.__submodule_order)} sub modules: {[m.ds_id for m in self.__submodule_order]}", - force=False) - else: - # Enable trace recording for next forward/backward pass - self.__trace_mode = ZeRoTraceMode.RECORD - - else: - if self.__profiler is not None: - self.__profiler.log_events() - - self.__param_queue = collections.deque(self.__param_order) # reset fetch queue - self.__most_recent_step_id_param_fetched_for = collections.defaultdict(lambda: int(-1e10)) - self.__step_id_module_fetched_for = collections.defaultdict(lambda: collections.deque()) - self.__step_id = 0 - self.__n_available_params = 0 - self.__profiler.reset_events() - - def _dump_params(self, tag, sub_module, params, step_id=None): - if step_id is None: - step_id = self.__step_id - param_names = [debug_param2name_id_shape(p) for p in params] - print_rank_0(f'{tag} step = {step_id} p_names = {param_names}', force=False) - - def _dump_param_ids(self, tag, mod_id, p_ids, step_id=None): - if step_id is None: - step_id = self.__step_id - print_rank_0(f'{tag} mod = {mod_id}, step = {step_id}, p_ids = {p_ids}', force=False) - - """Fetch and Release - Fetching, prefetching, and releasing parameters - """ - - @compiler.disable - @instrument_w_nvtx - @torch.no_grad() - def fetch_sub_module(self, current_submodule: Module, forward: bool) -> None: - """This method does the following (in order): - 1. kick off fetch for parameters in immediately required sub module - 2. kick off fetch for next few parameters we will need later (prefetch) - 3. block on parameters in immediately required sub module - """ - if logger.isEnabledFor(logging.DEBUG): - debug_rank0( - f"{self.__step_id}: M{current_submodule.ds_id}({type(current_submodule).__name__}) P{[p.ds_id for p in iter_params(current_submodule, recurse=z3_leaf_module(current_submodule))]} " - + str({ - "avail": f"{self.__n_available_params:.1e}", - "queue_sz": f"{len(self.__param_queue or [])}", - "inflight": [p.ds_id for p in self.__inflight_param_registry], - })) - - params_to_fetch = set(iter_params(current_submodule, recurse=z3_leaf_module(current_submodule))) - fetch_numel = sum( - [p.partition_numel() for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]) - - if fetch_numel > 0: - event_name = __class__.FORWARD_FETCH_SUBMIT if forward else __class__.BACKWARD_FETCH_SUBMIT - self._dump_param_ids(event_name, current_submodule.ds_id, - [(p.ds_id, p.ds_shape) - for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]) - # self._dump_params(event_name, current_submodule, [p for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]) - - self.__profiler.start_event(event_name) - # kick off all gather for params in the immediately required submodule - #for param in params_to_fetch: - if logger.isEnabledFor(logging.DEBUG): - for param in params_to_fetch: - debug_rank0(f"-fetch: {param.ds_summary()}") - self.__all_gather_params(params_to_fetch, forward) - self.__profiler.stop_event(event_name, fetch_numel) - - wait_numel = 0 - wait_event_name = __class__.FORWARD_FETCH_WAIT if forward else __class__.BACKWARD_FETCH_WAIT - self.__profiler.start_event(wait_event_name) - fast_fetch = self.fast_sharding_for_leaf_module and z3_leaf_module(current_submodule) - # wait for parameters in the immediately needed submodule to become available - for param in params_to_fetch: - param.ds_active_sub_modules.add(current_submodule.ds_id) - if logger.isEnabledFor(logging.DEBUG): - debug_rank0(f"-wait: {param.ds_summary()}") - if param in self.__inflight_param_registry: - wait_numel += param.partition_numel() - with get_accelerator().stream(self.__allgather_stream): - while self.__ongoing_fetch_events and self.__ongoing_fetch_events[0].query(): - self.__ongoing_fetch_events.popleft() - if len(self.__ongoing_fetch_events) > self.__max_ongoing_fetch_events: - self.__ongoing_fetch_events.popleft().synchronize() - - self.__inflight_param_registry.pop(param).wait(handle_dependency=not fast_fetch) - - if not get_accelerator().handles_memory_backpressure() and not fast_fetch: - event = get_accelerator().Event() - event.record() - self.__ongoing_fetch_events.append(event) - - assert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary() - if not get_accelerator().resolves_data_dependency(): - get_accelerator().current_stream().wait_stream(self.__allgather_stream) - if fast_fetch: - AllGatherCoalescedHandle.free_buffer() - self.__profiler.stop_event(wait_event_name, wait_numel) - - # kick off parameter prefetches for upcoming modules - # don't prefetch if we dont have a completed model trace - if self.is_complete_trace(): - # go through the parameters we need for the current module and pop them - # off the fetch queue so that they aren't prefetched later. - # if params have already been popped off the fetch queue by earlier - # prefetches we won't look for them here - discarded_from_prefetch_queue = set() - params_not_already_fetched = set( - filter(lambda p: self.__most_recent_step_id_param_fetched_for[p] < self.__step_id, params_to_fetch)) - while self.__param_queue and len(discarded_from_prefetch_queue) < len(params_not_already_fetched): - param_in_trace = self.__param_queue.popleft() - self.__most_recent_step_id_param_fetched_for[ - param_in_trace.param] = param_in_trace.step_id_last_used_at - discarded_from_prefetch_queue.add(param_in_trace.param) - - if discarded_from_prefetch_queue != params_not_already_fetched: - raise RuntimeError( - f"tracing error at step {self.__step_id}: \n" - f"module id: {current_submodule.ds_id}, training: {current_submodule.training}\n" - f"expected the next {len(params_not_already_fetched)} parameters in the " - f"parameter fetch queue to be {tuple(p.ds_summary(use_debug_name=True) for p in params_not_already_fetched)} \n" - f"but got \n {tuple(p.ds_summary(use_debug_name=True) for p in discarded_from_prefetch_queue)}.") - - def _is_currently_on_nvme(param): - if param.nvme_swapper is None: - return False - - return param.ds_tensor.final_location == OffloadDeviceEnum.nvme \ - and param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE - - # kick off all gather for params in the next few submodules (prefetch) - if self.__prefetch_bucket_sz > 0: - max_params_to_prefetch = min(self.__max_n_available_params - self.__n_available_params, - self.__prefetch_bucket_sz) - params_to_prefetch = set() - numel_prefetching = 0 - while self.__param_queue and numel_prefetching < max_params_to_prefetch: - param_in_trace: __class__.__ParamInTrace = self.__param_queue.popleft() - - if _is_currently_on_nvme(param_in_trace.param): - # nvme prefetch is handled elsewhere. Need to break here to preserve fetch order - self.__param_queue.appendleft(param_in_trace) - break - - do_prefetch = param_in_trace.param.ds_status == ZeroParamStatus.NOT_AVAILABLE - if param_in_trace.param in params_to_prefetch: - # Avoid duplicates - do_prefetch = False - - self.__most_recent_step_id_param_fetched_for[param_in_trace.param] = \ - max(self.__most_recent_step_id_param_fetched_for[param_in_trace.param], - param_in_trace.step_id_last_used_at) - - if do_prefetch: - params_to_prefetch.add(param_in_trace.param) - numel_prefetching += param_in_trace.param.ds_numel - - if numel_prefetching > 0: - event_name = __class__.FORWARD_PREFETCH_SUBMIT if forward else __class__.BACKWARD_PREFETCH_SUBMIT - self.__profiler.start_event(event_name) - if logger.isEnabledFor(logging.DEBUG): - for param in params_to_prefetch: - debug_rank0(f"-prefetch: {param.ds_summary()}") - self.__all_gather_params(params_to_prefetch, forward) - self.__profiler.stop_event(event_name, numel_prefetching) - - if self.__prefetch_nvme: - self.__prefetch_nvme_param_partitions() - - self.__step_id += 1 - - @instrument_w_nvtx - @torch.no_grad() - def release_sub_module(self, submodule: Module, forward=False) -> None: - """release the parameters of a sub module, assuming they meet conditions to - be released.""" - #print_rank_0(f"release_sub_module {'fwd' if forward else 'bwd'}: {debug_module2name_id(submodule)}", force=False) - params_to_release = (self.__params_to_release(submodule, self.__step_id) if self.is_complete_trace() else set( - p.ds_id for p in iter_params(submodule, recurse=z3_leaf_module(submodule)))) - - free_data = not z3_leaf_module(submodule) or not self.fast_sharding_for_leaf_module - if not free_data: - # wait for the computation to finish and launch as early as possible. - empty_buffer = torch.empty(1, device=get_accelerator().current_device()) - - for param in iter_params(submodule, recurse=z3_leaf_module(submodule)): - param.ds_active_sub_modules.discard(submodule.ds_id) - if param.ds_id in params_to_release and not param.is_external_param: - self.__release_param(param, free_data) - if not free_data: - if param.ds_id in params_to_release and not param.is_external_param: - # empty buffer ensures that all computations are complete - param.data = empty_buffer - - @instrument_w_nvtx - @torch.no_grad() - def release_and_reset_all(self, module: Module) -> None: - """release all module parameters""" - for param in iter_params(module, recurse=True): - if param in self.__inflight_param_registry: - self.__inflight_param_registry.pop(param).wait() - - # TODO. make this throw if if there are still active submodules. currently - # there's a hook execution issue - param.ds_active_sub_modules.clear() - self.__release_param(param) - - for param in iter_params(module, recurse=True): - if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: - raise RuntimeError(f"{param.ds_summary()} expected to be released") - - @instrument_w_nvtx - def __all_gather_params(self, params: Set[Parameter], forward: bool) -> None: - quantized_params = [] - nonquantized_params = [] - for param in params: - if hasattr(param.ds_tensor, 'ds_quant_scale'): - quantized_params.append(param) - else: - nonquantized_params.append(param) - if quantized_params: - self.__all_gather_params_(quantized_params, forward, quantize=True) - if nonquantized_params: - self.__all_gather_params_(nonquantized_params, forward, quantize=self.zero_quantized_weights) - - def __all_gather_params_(self, params: Set[Parameter], forward: bool, quantize: bool = False) -> None: - """for each partitioned parameter, kick off an async allgather and store - the work handle for the in flight parameters.""" - partitioned_params = [] - all_gather_numel = 0 # numel = num of elements - for param in params: - if param.ds_status == ZeroParamStatus.NOT_AVAILABLE: - partitioned_params.append(param) - all_gather_numel += param.ds_numel - - if partitioned_params: - self.__n_available_params += all_gather_numel - # here we need to handle a special case where some of the parameters have a valid hpz secondary tensor (e.g. they are not trainable so their secondary tensor never expire) but others do not. - partitioned_params_with_secondary_tensors = [ - p for p in partitioned_params if p.ds_secondary_tensor is not None - ] - partitioned_params_without_secondary_tensors = [ - p for p in partitioned_params if p.ds_secondary_tensor is None - ] - for param_group in [ - partitioned_params_with_secondary_tensors, partitioned_params_without_secondary_tensors - ]: - if not param_group: - continue - with get_accelerator().stream(self.__allgather_stream): - event_name = __class__.FORWARD_ALL_GATHER if forward else __class__.BACKWARD_ALL_GATHER - self.__profiler.start_event(event_name) - handle = param_group[0].all_gather_coalesced(param_group, quantize=quantize) - self.__profiler.stop_event(event_name, all_gather_numel) - for param in param_group: - assert param.ds_status == ZeroParamStatus.INFLIGHT, param.ds_summary() - self.__inflight_param_registry[param] = handle - - # Release swap buffers for persisted params on nvme since they will never be partitioned or evicted from GPU - swap_persisted_params = [ - p for p in partitioned_params if p.ds_persist and p.ds_tensor.final_location == OffloadDeviceEnum.nvme - ] - if swap_persisted_params: - swap_persisted_params[0].nvme_swapper.remove_partition_and_release_buffers(swap_persisted_params) - - @compiler.disable - @instrument_w_nvtx - def __release_param(self, param: Parameter, free_data: bool = True) -> None: - if param.ds_status == ZeroParamStatus.AVAILABLE and not param.ds_active_sub_modules: - if logger.isEnabledFor(logging.DEBUG): - debug_rank0(f"-release: {param.ds_summary()}") - print_rank_0(f"release: {debug_param2name_id_shape(param)}", force=False) - param.partition(free_data=free_data) - self.__n_available_params -= param.ds_numel - - @instrument_w_nvtx - @functools.lru_cache(maxsize=None) - def __params_to_release(self, submodule_to_release: Module, step_id: int) -> Set[int]: - if not self.is_complete_trace(): - raise RuntimeError("expected trace to be complete") - - params_to_release = set( - p.ds_id for p in iter_params(submodule_to_release, recurse=z3_leaf_module(submodule_to_release)) - if not p.ds_persist) - - # Problem: When prefetcher scans the param trace, it skips AVAILABLE params. - # This creates issues if those params are released before the skipped uses: - # 1) It hurts performance as the skipped uses are never prefetched. - # 2) For nvme params, we run out of swap buffers because the prefetch order - # diverges from the trace. - # Solution: Don't release params whose reuse was skipped by prefetch. This is - # possible because we detect such skips during prefetch and mark those params. - for param in iter_params(submodule_to_release, recurse=z3_leaf_module(submodule_to_release)): - if self.__most_recent_step_id_param_fetched_for[param] > step_id: - params_to_release.discard(param.ds_id) - - # examine all modules within `max_reuse_dist_in_numel` of the current step, - # if we see any of the candidate parameters to be released reoccur while - # doing this, remove them from the set of parameters to release. - params_traversed = 0 - for module in self.__submodule_order[step_id:]: - if params_traversed >= self.__max_reuse_dist_in_numel: - break - for param in iter_params(module, recurse=z3_leaf_module(submodule_to_release)): - params_to_release.discard(param.ds_id) - params_traversed += param.ds_numel - - return params_to_release - - @instrument_w_nvtx - def __prefetch_nvme_param_partitions(self) -> None: - """swap in parameter partitions from nvme for those parameters that will be used - after the ones that are already being prefetched into full parameters - """ - if not self.is_complete_trace(): - return - - numel_in_flight = sum(param.ds_numel for param in self.__inflight_param_registry) - - numel_considered = 0 - swap_in_params = [] - for param_in_trace in self.__param_queue: - param = param_in_trace.param - if param.nvme_swapper is None: - continue - if (numel_considered > 2 * numel_in_flight - or len(swap_in_params) >= param.nvme_swapper.available_swap_in_buffers()): - break - if param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE: - swap_in_params.append(param) - numel_considered += param.ds_numel - - if swap_in_params: - swap_in_params[0].nvme_swapper.swap_in(swap_in_params, async_op=True) diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_profiler.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_profiler.py deleted file mode 100644 index b4ea11f3b83632a5f7f74343ef7092af6c5bf3a7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_profiler.py +++ /dev/null @@ -1,63 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from dataclasses import dataclass -from deepspeed.utils import log_dist - - -class PartitionedParameterProfiler(object): - - @dataclass - class EventCounter: - name: str - count: int - num_elem: int - - def reset(self): - self.count = 0 - self.num_elem = 0 - - def increment(self, numel): - self.count += 1 - self.num_elem += numel - - def __init__(self, timers): - self.timers = timers - self.event_counters = {} - - def reset_events(self): - for event_ctr in self.event_counters.values(): - event_ctr.reset() - - def start_event(self, name): - if self.timers is None: - return - - if name not in self.event_counters: - self.event_counters[name] = __class__.EventCounter(name=name, count=0, num_elem=0) - self.timers(name).start() - - def stop_event(self, name, num_elem): - if self.timers is None: - return - assert name in self.event_counters, f'unknown event {name}' - self.event_counters[name].increment(num_elem) - self.timers(name).stop() - - def _log_timers(self): - if self.timers is None: - return - self.timers.log(names=list(self.event_counters.keys())) - - def _log_event_counters(self): - for event_ctr in self.event_counters.values(): - log_dist( - f'{event_ctr.name}: count = {event_ctr.count}, numel = {event_ctr.num_elem}', - #f'{event_ctr.name}: time = {self._log_timers()},count = {event_ctr.count}, numel = {event_ctr.num_elem}', - ranks=[0]) - - def log_events(self): - self._log_event_counters() - self._log_timers() diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/stage3.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/stage3.py deleted file mode 100644 index 2937191a0ad56136ab2bb9219d487e42f251f3b2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/stage3.py +++ /dev/null @@ -1,3262 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import sys -import gc -import collections -import itertools -from typing import Deque, Dict, Set, List, Tuple, Container, Optional -from contextlib import contextmanager -from dataclasses import dataclass, field - -from deepspeed import comm as dist -from deepspeed.utils import groups, z3_leaf_parameter - -from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors -from deepspeed.runtime.base_optimizer import ZeROOptimizer -from deepspeed.utils import logger -from deepspeed.utils.torch import register_grad_hook -from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler -from deepspeed.runtime.torch_autocast import get_all_autocast_dtypes, is_autocast_initialized, sort_dtypes -from deepspeed.runtime.comm.coalesced_collectives import reduce_scatter_coalesced, all_to_all_quant_reduce, all_to_all_loco_quant_reduce -from deepspeed.runtime.utils import inf, is_model_parallel_parameter, get_only_unique_item, mask_nan_or_inf_with_val_inplace -from deepspeed.runtime.zero.partition_parameters import * -from deepspeed.runtime.zero.config import ZeroStageEnum -from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum, OffloadStateTypeEnum -from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload -from deepspeed.runtime.zero.utils import apply_to_tensors_only, get_mapping_to_flat_buffer -from deepspeed.runtime.zero.offload_states import offload_adam_states, reload_adam_states -from deepspeed.ops.adam import DeepSpeedCPUAdam -from deepspeed.runtime.swap_tensor.partitioned_param_swapper import PartitionedParamStatus -from deepspeed.runtime.swap_tensor.optimizer_utils import OptimizerSwapper -from deepspeed.runtime.swap_tensor.partitioned_optimizer_swapper import PartitionedOptimizerSwapper -from deepspeed.runtime.swap_tensor.pipelined_optimizer_swapper import PipelinedOptimizerSwapper -from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, FP32_FLAT_GROUPS, PARTITION_COUNT, ZERO_STAGE, LOSS_SCALER -from deepspeed.accelerator import get_accelerator - -# Toggle this to true to enable correctness test -# with gradient partitioning and without -pg_correctness_test = False - -OPTIMIZER_SWAP_IN_STATE_TIMER = 'optimizer_swap_in_state' -INIT_OPTIMIZER_TIMER = 'init_optimizer_state' -OPTIMIZER_SWAP_OUT_STATE_TIMER = 'optimizer_swap_out_state' -OPTIMIZER_STEP_TIMER = 'optimizer_step' - - -def print_rank_0(message, debug=False, force=False): - rank = dist.get_rank() - if rank == 0 and (debug or force): - logger.info(message) - # other variations - # - print for all ranks w/o interleaving - # printflock(f"[{rank}] {message}") - # - print to log file per rank - # log_rank_file(rank, message) - - -def input(msg): - return - - -def isclose(a, b, rtol=1e-09, atol=0.0): - return abs(a - b) <= max(rtol * max(abs(a), abs(b)), atol) - - -def lcm(x, y): - from fractions import gcd # or can import gcd from `math` in Python 3 - return x * y // gcd(x, y) - - -def move_to_cpu(tensor_list): - for tensor in tensor_list: - tensor.data = tensor.data.cpu() - - -@contextmanager -def unwrap_model_for_generation(model): - """ - For ZeRO-3 models, we gather the weights once to speed up generation. - """ - with GatheredParameters(model.parameters()): - # Removes the optimizer hooks from a DeepSpeed ZeRO-3 model. - - # Remove hooks - if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"): - optimizer_offload = model.optimizer.parameter_offload - elif model.optimizer is not None: - optimizer_offload = model.optimizer - - for hook in optimizer_offload.forward_hooks: - hook.remove() - for hook in optimizer_offload.backward_hooks: - hook.remove() - - optimizer_offload.forward_hooks = [] - optimizer_offload.backward_hooks = [] - - yield model - - # Adds the optimizer hooks from a DeepSpeed ZeRO-3 model. - if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"): - optimizer_offload = model.optimizer.parameter_offload - elif model.optimizer is not None: - optimizer_offload = model.optimizer - optimizer_offload._register_deepspeed_module(optimizer_offload.module) - return - - -@dataclass -class IPGBucketZ3: - buffer: Optional[torch.Tensor] = None - buffer_meta: Optional[torch.Tensor] = None - params: List[torch.Tensor] = field(default_factory=list) - elements: int = 0 - - def clear(self): - self.buffer = None - self.buffer_meta = None - self.params.clear() - self.elements = 0 - - -INITIAL_MICRO_STEP_ID = -1 - - -class DeepSpeedZeroOptimizer_Stage3(ZeROOptimizer): - """ - DeepSpeedZeroOptimizer designed to reduce the memory footprint - required for training large deep learning models. - - For more details please see ZeRO: Memory Optimization Towards Training A Trillion Parameter Models - https://arxiv.org/abs/1910.02054 - - For usage examples, refer to TODO: DeepSpeed Tutorial - - """ - - def __init__( - self, - module, - init_optimizer, - timers, - ds_config, - static_loss_scale=1.0, - dynamic_loss_scale=False, - dynamic_loss_args=None, - verbose=True, - contiguous_gradients=True, - reduce_bucket_size=500000000, - prefetch_bucket_size=50000000, - max_reuse_distance=1000000000, - max_live_parameters=1000000000, - param_persistence_threshold=100000, - model_persistence_threshold=sys.maxsize, - dp_process_group=None, - reduce_scatter=True, - overlap_comm=False, - offload_optimizer_config=None, - offload_param_config=None, - sub_group_size=1000000000000, - offload_ratio=0.0, - mpu=None, - clip_grad=0.0, - gradient_accumulation_dtype=torch.float32, - communication_data_type=torch.float16, - postscale_gradients=True, - gradient_predivide_factor=1.0, - gradient_accumulation_steps=1, - elastic_checkpoint=False, - aio_config=None, - all2all_process_group=None, - zero_hpz_partition_size=1, - zero_quantized_weights=False, - zero_quantized_nontrainable_weights=False, - zero_module_granularity_threshold=0, - zeropp_loco_param=None, - log_trace_cache_warnings=False, - ): - see_memory_usage("Stage 3 initialize beginning", force=True) - - print_rank_0(f"initialized {__class__.__name__} with args: {locals()}", force=False) - - if dist.get_rank() == 0: - logger.info(f"Reduce bucket size {reduce_bucket_size}") - logger.info(f"Prefetch bucket size {prefetch_bucket_size}") - # The fused optimizer does all the work. We need this layer for two reason: - # 1. maintain same user API from apex.fp16_utils - # 2. keep common stuff here in case we need to add ne552w fused optimizer later - - # differences from apex.fp16_utils: - # - assume all model params in fp16 - # - assume all params requires grad - # - flat by groups, not keeping state. TODO: remove state explicitly? - # - master grad and unflat master weight never exist. TODO: a way to save out unflat master? - if not get_accelerator().is_available(): - raise SystemError("Cannot use fp16 without accelerator.") - - self.optimizer = init_optimizer - - # Use torch (un)flatten ops - self.flatten = _flatten_dense_tensors - self.unflatten = _unflatten_dense_tensors - self.dtype = self.optimizer.param_groups[0]['params'][0].dtype - self.gradient_accumulation_dtype = gradient_accumulation_dtype - self._global_grad_norm = 0. - - self.custom_loss_scaler = False - self.external_loss_scale = None - - self.optimizer_swapper = None - self.swap_optimizer = False - - self.offload_optimizer = False - self.offload_optimizer_pin_memory = False - self.offload_optimizer_fast_init = False - self.offload_param = False - self.offload_param_pin_memory = False - self.params_in_nvme_and_cpu = False - self.max_params_in_cpu = 0 - self.partial_offload = offload_ratio - - #num of ranks in a ZeRO param partitioning group - self.zero_hpz_partition_size = zero_hpz_partition_size - - zero_param_parallel_group = groups._get_zero_param_intra_parallel_group() - print_rank_0( - f"ZeRO Stage 3 param partitioning group {self.zero_hpz_partition_size} {zero_param_parallel_group}", - force=False) - if self.zero_hpz_partition_size > 1 and zero_param_parallel_group is None: - self._set_zero_group_parallelism() - zero_param_parallel_group = groups._get_zero_param_intra_parallel_group() - - self.parameter_offload = self.initialize_ds_offload( - module=module, - timers=timers, - ds_config=ds_config, - overlap_comm=overlap_comm, - prefetch_bucket_size=prefetch_bucket_size, - max_reuse_distance=max_reuse_distance, - max_live_parameters=max_live_parameters, - param_persistence_threshold=param_persistence_threshold, - model_persistence_threshold=model_persistence_threshold, - dp_process_group=dp_process_group, - offload_param_config=offload_param_config, - mpu=mpu, - zero_param_parallel_group=zero_param_parallel_group, - zero_quantized_weights=zero_quantized_weights, - zero_quantized_nontrainable_weights=zero_quantized_nontrainable_weights, - zero_module_granularity_threshold=zero_module_granularity_threshold, - log_trace_cache_warnings=log_trace_cache_warnings, - ) - - self.persistent_parameters = self.parameter_offload.persistent_parameters - self._configure_offloading(offload_optimizer_config, offload_param_config) - - # backup fused_adam optimizer init - if self.offload_optimizer and self.partial_offload != 1.0: - backup_gpu_tensor = torch.randn(1, device=get_accelerator().device_name()).to(self.dtype) - backup_gpu_param = torch.nn.Parameter(backup_gpu_tensor) - assert type(init_optimizer) == DeepSpeedCPUAdam, 'Hybrid Optimizer Only Supports DeepSpeedCPUAdam' - self.backup_optimizer = torch.optim.AdamW([backup_gpu_param], - lr=self.optimizer.param_groups[0]["lr"], - betas=self.optimizer.param_groups[0]["betas"], - eps=self.optimizer.param_groups[0]["eps"], - weight_decay=self.optimizer.param_groups[0]["weight_decay"], - amsgrad=self.optimizer.param_groups[0]["amsgrad"]) - # Multiple param_groups configs for back-up optimizer - if len(self.optimizer.param_groups) > 1: - for i in range(1, len(self.optimizer.param_groups)): - self.backup_optimizer.add_param_group(self.optimizer.param_groups[i]) - - self.module = module - self.elastic_checkpoint = elastic_checkpoint - - self.device = get_accelerator().current_device_name() if not self.offload_optimizer else OffloadDeviceEnum.cpu - - self.inf_or_nan_tracker: Tensor = torch.zeros(1, dtype=torch.bool, device=self.device, requires_grad=False) - - self.deepspeed_adam_offload = (self.offload_optimizer and type(init_optimizer) == DeepSpeedCPUAdam) - - ### streams used for overlapping computation with communication - self.reduce_and_partition_stream = None if get_accelerator().is_synchronized_device() else get_accelerator( - ).Stream() if overlap_comm else get_accelerator().default_stream() - - ############################################################################ - - self.n_caching_allocator_flushes = 0 - - #-------------Stage 3 Setup-------------------# - - self.timers = timers - - self.all2all_process_group = all2all_process_group - - self.reduce_scatter = reduce_scatter - - self.dp_process_group = self.parameter_offload.dp_process_group - self.sequence_parallel_size = groups._get_sequence_parallel_world_size() - - self.all2all_process_group = all2all_process_group - - self.zero_quantized_nontrainable_weights = zero_quantized_nontrainable_weights - - self.partition_count = dist.get_world_size(group=self.dp_process_group) - - self.zeropp_loco_param = zeropp_loco_param - - if mpu is None or hasattr(mpu, 'initialize_sequence_parallel'): - self.model_parallel_group = None - self.model_parallel_rank = 0 - else: - self.model_parallel_group = mpu.get_model_parallel_group() - self.model_parallel_rank = mpu.get_model_parallel_rank() - - self.overflow = False - self.clip_grad = clip_grad - self.communication_data_type = communication_data_type - self.gradient_predivide_factor = gradient_predivide_factor - self.postscale_gradients = postscale_gradients - self.gradient_accumulation_steps = gradient_accumulation_steps - self.micro_step_id = 0 - self.reduce_bucket_size = int(reduce_bucket_size) - - if self.all2all_process_group is not None: - assert self.all2all_process_group is not None and self.reduce_scatter == True, "when enable all_to_all_reduce, reduce_scatter should also be enabled for data type checks." - - if self.reduce_scatter: - valid_reduce_scatter_dtypes = (torch.float16, torch.bfloat16, torch.float32) - assert self.communication_data_type in valid_reduce_scatter_dtypes, f"ZeRO-3 supports {valid_reduce_scatter_dtypes} communication_data_type with reduce scatter enabled. Got: '{self.communication_data_type}'" - assert self.gradient_predivide_factor == 1.0, "gradient_predivide_factor != 1.0 is not yet supported with ZeRO-3 with reduce scatter enabled" - assert self.postscale_gradients, "pre-scale gradients is not yet supported with ZeRO-3 with reduce scatter enabled" - - # Holds the mode parameter - # The param.data may not hold any meaningful data - # when param's status is NOT_AVAILABLE or IN_FLGHT - self.fp16_groups = [] - - # Hold partitioned parameters - self.fp16_partitioned_groups = [] - - # Holds a fused and flattened copy of the parameters - self.fp16_partitioned_groups_flat = [] - self.fp16_partitioned_groups_flat_numel = [] - self.fp16_partitioned_groups_flat_id = [] - - #defragmented pinned memory - self.param_groups_fp16_flat_cpu_memory = [] - - #a single 32-bit partition of the parallel partitioned parameters - #that this process will update - self.fp32_partitioned_groups_flat = [] - self.next_swappable_fp32_partitioned_groups = [] - - # number of elements per partition in each group - self.partition_size = [] - - self.all_reduce_print = False - - self.prefetch_elements = int(prefetch_bucket_size) - - self.contiguous_gradients = contiguous_gradients - - # padding on each partition for alignment purposes - self.groups_padding = [] - - self.sub_group_size = sub_group_size - - self.sub_group_to_group_id = {} - - # Trainable parameters - self.trainable_param_groups = self._get_trainable_parameter_groups() - - see_memory_usage("Before creating fp16 partitions", force=True) - self._create_fp16_partitions_with_defragmentation(self.trainable_param_groups) - num_fp16_subgroups = len(self.fp16_partitioned_groups_flat) - see_memory_usage(f"After creating fp16 partitions: {num_fp16_subgroups}", force=True) - - # Optimizer tensor swapping - if self.swap_optimizer: - self._configure_tensor_swapping(offload_optimizer_config, aio_config) - - self.is_gradient_accumulation_boundary: bool = True - - self.param_reduce_events: Deque[get_accelerator().Event] = collections.deque() - # TODO. make this configurable via JSON - self.max_param_reduce_events: int = 2 - - self.param_dict = {} - - # map between param_id and bool to specify if a param is in this partition - self.is_param_in_current_partition = {} - - if is_autocast_initialized(): - comm_dtypes = get_all_autocast_dtypes([p for params in self.fp16_groups for p in params]) - self.torch_autocast_gradscaler = torch.amp.GradScaler(device=get_accelerator().device_name()) - else: - comm_dtypes = {self.communication_data_type} - self.torch_autocast_gradscaler = None - - self.ipg_buckets: Dict[torch.dtype, IPGBucketZ3] = {dtype: IPGBucketZ3() for dtype in comm_dtypes} - - self.params_already_reduced = {} - self.previous_reduced_grads = None - - # model parameter traversal-based param id that's stable across runs - for params_group in self.fp16_groups: - for param in params_group: - param_id = self.get_param_id(param) - self.param_dict[param_id] = param - self.params_already_reduced[param_id] = False - - #Largest partitioned param - largest_partitioned_param_numel = 0 - for fp16_partitioned_group in self.fp16_partitioned_groups: - if len(fp16_partitioned_group) > 0: - largest_partitioned_param_numel = max( - largest_partitioned_param_numel, - max([max(tensor.numel(), tensor.ds_numel) for tensor in fp16_partitioned_group])) - - print_rank_0(f'Largest partitioned param numel = {largest_partitioned_param_numel}', force=False) - - self._setup_for_real_optimizer() - self.grad_position = {} - self.set_grad_positions() - - if self.offload_optimizer: - self.norm_for_param_grads = {} - - # stores if a partition has been reduced in this step - self.is_partition_reduced = {} - - # stores if a grad in a partition has been computed or not - self.is_grad_computed = {} - - # will store the averaged gradients required by this partition - self.averaged_gradients = {} - - #creates backward hooks for gradient partitioning - ###Calls all gather param - self._grad_acc_hooks = [] - self._leaf_module_hooks = [] - self.create_reduce_and_remove_grad_hooks() - - #exit(0) - - # we may have a way of fusing dynamic scale. Do not support for now - self.loss_scaler = CreateLossScaler(dtype=self.dtype, - static_loss_scale=static_loss_scale, - dynamic_scaling=dynamic_loss_scale, - dynamic_loss_args=dynamic_loss_args) - self.dynamic_loss_scale = self.loss_scaler.dynamic - - self.debug_fp16_grads = [{} for _ in self.fp16_groups] - - self._link_all_hp_params() - - self.offloaded_states: Set[OffloadDeviceEnum] = set() - - if dist.get_rank(group=self.dp_process_group) == 0: - see_memory_usage(f"After initializing ZeRO optimizer", force=True) - - def destroy(self): - self.parameter_offload.destroy() - for hook in self._grad_acc_hooks: - hook.remove() - for hook in self._leaf_module_hooks: - hook.remove() - print_rank_0("Removed grad acc hooks", force=False) - self.ipg_buckets.clear() - - def initialize_ds_offload( - self, - module, - timers, - ds_config, - overlap_comm, - prefetch_bucket_size, - max_reuse_distance, - max_live_parameters, - param_persistence_threshold, - model_persistence_threshold, - dp_process_group, - offload_param_config, - mpu, - zero_param_parallel_group, - zero_quantized_weights, - zero_quantized_nontrainable_weights, - zero_module_granularity_threshold, - log_trace_cache_warnings, - ): - return DeepSpeedZeRoOffload(module=module, - timers=timers, - ds_config=ds_config, - overlap_comm=overlap_comm, - prefetch_bucket_size=prefetch_bucket_size, - max_reuse_distance=max_reuse_distance, - max_live_parameters=max_live_parameters, - param_persistence_threshold=param_persistence_threshold, - model_persistence_threshold=model_persistence_threshold, - dp_process_group=dp_process_group, - offload_param_config=offload_param_config, - mpu=mpu, - zero_param_parallel_group=zero_param_parallel_group, - zero_quantized_weights=zero_quantized_weights, - zero_quantized_nontrainable_weights=zero_quantized_nontrainable_weights, - zero_module_granularity_threshold=zero_module_granularity_threshold, - log_trace_cache_warnings=log_trace_cache_warnings) - - def _get_trainable_parameter_groups(self): - param_groups = [] - PARAMS_KEY = "params" - for param_group in self.optimizer.param_groups: - trainable_params = [p for p in param_group[PARAMS_KEY] if p.requires_grad] - if len(trainable_params) == 0: - continue - - trainable_param_group = {} - for key in param_group.keys(): - if key == PARAMS_KEY: - trainable_param_group[PARAMS_KEY] = trainable_params - else: - trainable_param_group[key] = param_group[key] - param_groups.append(trainable_param_group) - - return param_groups - - def _set_zero_group_parallelism(self): - groups._create_zero_param_parallel_group(self.zero_hpz_partition_size) - - def invalidate_secondary_tensor(self): - for fpg in self.fp16_groups: - for param in fpg: - if param.ds_secondary_tensor is not None: - param.ds_secondary_tensor = None - - def _setup_for_real_optimizer(self): - see_memory_usage("Before creating fp32 partitions", force=True) - self._create_fp32_partitions() - see_memory_usage("After creating fp32 partitions", force=True) - dist.barrier() - - # To support pipelined optimizer swapping - self._create_next_swappable_fp32_groups() - - see_memory_usage("Before initializing optimizer states", force=True) - - self.initialize_optimizer_states() - see_memory_usage("After initializing optimizer states", force=True) - dist.barrier() - - if dist.get_rank() == 0: - logger.info(f"optimizer state initialized") - - # IPG - if self.contiguous_gradients: - for dtype, bucket in self.ipg_buckets.items(): - bucket.buffer = torch.empty(self.reduce_bucket_size, - dtype=dtype, - device=get_accelerator().current_device_name()) - - self.grad_partitions_flat_buffer = None - self.__param_id_to_grad_partition: Dict[int, Tensor] = {} - - all_params = list(itertools.chain.from_iterable(self.fp16_groups)) - - self.grad_partitions_flat_buffer: Tensor = torch.zeros(sum(p.partition_numel() for p in all_params), - dtype=self.gradient_accumulation_dtype, - device=self.device) - if self.offload_optimizer_pin_memory: - self.grad_partitions_flat_buffer = get_accelerator().pin_memory(self.grad_partitions_flat_buffer) - - offset = 0 - for param in all_params: - self.__param_id_to_grad_partition[param.ds_id] = self.grad_partitions_flat_buffer.narrow( - 0, offset, param.partition_numel()) - offset += param.partition_numel() - - def _link_all_hp_params(self): - for p in self.module.parameters(): - p._z3_optimizer = self - - def set_lr(self, lr): - """Set the learning rate.""" - for param_group in self.optimizer.param_groups: - param_group["lr"] = lr - - def get_lr(self): - """Return the current learning rate.""" - return self.optimizer.param_groups[0]["lr"] - - # TODO. factor out to a utility outside of stage3 - @staticmethod - def defragment(tensors: List[Tensor]) -> Tensor: - """move provided tensors into a contiguous flat buffer, with some additional - measures taken to reduce memory fragmentation""" - assert len(set(t.dtype for t in tensors)) == 1 - assert len(set(t.device for t in tensors)) == 1 - - cpu_buffer = torch.empty(sum(p.numel() for p in tensors), - dtype=get_only_unique_item(t.dtype for t in tensors), - device="cpu") - tensor_infos: List[Tuple[Tensor, int, int]] = get_mapping_to_flat_buffer(tensors) - orig_device = get_only_unique_item(t.device for t in tensors) - - offset = 0 - for tensor, offset, tensor_numel in tensor_infos: - # move the tensor from device memory to host memory - cpu_buffer.narrow(0, offset, tensor_numel).copy_(tensor) - tensor.data = torch.empty(0, dtype=tensor.dtype, device=tensor.device) - - gc.collect() - get_accelerator().empty_cache() - - # copy tensors (now flattened and contiguous) back to GPU - device_buffer = cpu_buffer.to(orig_device) - - # restore device tensors - for tensor, offset, tensor_numel in tensor_infos: - tensor.data = device_buffer.narrow(0, offset, tensor_numel) - - return device_buffer - - def _get_param_coordinator(self): - return self.parameter_offload.get_param_coordinator() - - def _configure_offloading(self, offload_optimizer_config, offload_param_config): - ###################### offload optimizer setup ################################## - if offload_optimizer_config is not None and offload_optimizer_config.device != OffloadDeviceEnum.none: - self.offload_optimizer = True - self.offload_optimizer_pin_memory = offload_optimizer_config.pin_memory - self.swap_optimizer = offload_optimizer_config.device == OffloadDeviceEnum.nvme - self.offload_optimizer_fast_init = offload_optimizer_config.fast_init - - ###################### offload param setup ################################## - if offload_param_config is not None and offload_param_config.device != OffloadDeviceEnum.none: - self.offload_param = True - self.offload_param_pin_memory = offload_param_config.pin_memory - self.params_in_nvme_and_cpu = offload_param_config.device == OffloadDeviceEnum.nvme - self.max_params_in_cpu = offload_param_config.max_in_cpu - print_rank_0( - f"FP16 params swapping is {self.params_in_nvme_and_cpu}, Max params in CPU is {self.max_params_in_cpu}", - force=False) - - def _configure_tensor_swapping(self, offload_optimizer_config, aio_config): - nvme_swap_folder = os.path.join(offload_optimizer_config.nvme_path, 'zero_stage_3') - os.makedirs(nvme_swap_folder, exist_ok=True) - if dist.get_rank() == 0: - logger.info(f'Tensor Swapping: Adding optimizer tensors') - - swapper_type = PipelinedOptimizerSwapper if offload_optimizer_config.pipeline else PartitionedOptimizerSwapper - - self.optimizer_swapper = swapper_type(swap_config=offload_optimizer_config, - aio_config=aio_config, - base_folder=nvme_swap_folder, - optimizer=self.optimizer, - largest_numel=max(self.fp16_partitioned_groups_flat_numel), - device=self.device, - dtype=torch.float32, - timers=self.timers) - - def _move_to_flat_buffer(self, param_list, flat_buffer, avoid_copy=False): - '''If flat buffer is None then the parameters in the param_list are - not copied to the flat buffer. This is because they exceed the number of max_params_in_cpu - Some of these parameters may already be in CPU in unflattened buffers - or they maybe in GPU, or they maybe in NVME. If they are in NVME, then - they will be marked as NOT_AVAILABLE, and will be moved to CPU when they are - needed during training.''' - if flat_buffer is None: - # this dst buffer is on NVMe, so skip this - return - - start = 0 - for param in param_list: - src = param.ds_tensor - dest = flat_buffer.narrow(0, start, src.ds_numel) - start = start + src.ds_numel - '''if the parameter was initialized in nvme then bring it to the destination buffer directly''' - if src.status == PartitionedParamStatus.NOT_AVAILABLE: - print_rank_0( - f"Swapping in {param.ds_id} with partition size {param.partition_numel()} permanently to CPU") - param.nvme_swapper.swap_into_buffer(param, dest) - src.data = dest.data - src.status = PartitionedParamStatus.AVAILABLE - else: - assert src.status == PartitionedParamStatus.AVAILABLE, "Partitioned Param must be available here" - if not avoid_copy: - dest.data.copy_(src.data) - src.data = dest.data - - # Final location must be gpu/cpu in this case - param.ds_tensor.final_location = 'not-nvme' - - def _create_param_groups_fp16_flat_cpu_memory(self): - - aggregate_params_count = 0 - - for j, param_group in enumerate(self.trainable_param_groups): - params_in_group = sum([p.partition_numel() for p in param_group['params']]) - - flat_buffer_size = params_in_group - - if self.params_in_nvme_and_cpu and \ - aggregate_params_count + params_in_group > self.max_params_in_cpu: - - flat_buffer_size = max(0, self.max_params_in_cpu - aggregate_params_count) - - aggregate_params_count += params_in_group - - if flat_buffer_size > 0: - print_rank_0(f"group {j} flat buffer size {flat_buffer_size}", force=False) - self.param_groups_fp16_flat_cpu_memory.append(get_accelerator().pin_memory( - torch.empty(int(flat_buffer_size), dtype=self.dtype))) - else: - print_rank_0(f"No flat buffer size. Param group size was {params_in_group}", force=False) - - self.param_groups_fp16_flat_cpu_memory.append(torch.empty(1, dtype=self.dtype)) - - def _create_fp16_partitions_with_defragmentation(self, fp16_param_groups): - dist.barrier() - - param_groups: List[List[Parameter]] = tuple( - self._create_fp16_sub_groups(param_group["params"]) for param_group in fp16_param_groups) - - # bookkeeping related to param groups - for param_group_idx, param_group in enumerate(param_groups): - for sub_group in param_group: - sub_group_idx = len(self.fp16_groups) - - # record sub group and partitions - self.fp16_groups.append(sub_group) - self.fp16_partitioned_groups.append([param.ds_tensor for param in sub_group]) - - # record sub group -> group mapping - self.sub_group_to_group_id[sub_group_idx] = param_group_idx - - # record total elements of parameter partitions in sub group - self.fp16_partitioned_groups_flat_numel.append(sum(p.partition_numel() for p in sub_group)) - - # record ds_ids of parameter partitions in sub group - self.fp16_partitioned_groups_flat_id.append([p.ds_id for p in sub_group]) - - # record padding required to align group to world size (only applies to last rank) - rank_requires_padding = dist.get_rank( - self.dp_process_group) == dist.get_world_size(self.dp_process_group) - 1 - self.groups_padding.append([p.padding_size() if rank_requires_padding else 0 for p in sub_group]) - - # move parameters to flattened buffer - if not self.offload_param: # partitioned params remain in GPU during training - # move parameter partitions into a single contiguous flat buffer - parameter_partitions = self._get_parameter_partitions() - - # We need to keep the reference to this buffer to make sure you can free it in `offload_states` - self.lp_param_buffer = __class__.defragment(parameter_partitions) - self._set_fp16_partitioned_groups_flat() - - else: # partitioned params offloaded to CPU when not in use - # create a flat CPU memory allocation for each param group - self._create_param_groups_fp16_flat_cpu_memory() - for param_group_idx, param_group in enumerate(param_groups): - flat_offset = 0 - for i, sub_group in enumerate(param_group): - total_elements = sum(p.partition_numel() for p in sub_group) - print_rank_0(f"Params in nvme and cpu {self.params_in_nvme_and_cpu}") - #Flat buffer may not be available for parameters that reside in NVME - if not self.params_in_nvme_and_cpu or flat_offset + total_elements <= self.param_groups_fp16_flat_cpu_memory[ - param_group_idx].numel(): - fp16_partitioned_group_flat = self.param_groups_fp16_flat_cpu_memory[param_group_idx].narrow( - 0, flat_offset, total_elements) - print_rank_0( - f"Creating a flat buffer for subgroup {i} requiring {total_elements} elements, and cumulative CPU elements {flat_offset + total_elements}", - force=False) - - elif self.params_in_nvme_and_cpu: - fp16_partitioned_group_flat = None - print_rank_0(f"No flat buffer for sub group {i} of {total_elements} elements", force=False) - else: - assert False, "Either params are in nvme, or they are in CPU memory. This code path should not be triggered. Please see you max_params_in_cpu and params_in_nvme configs" - - self.fp16_partitioned_groups_flat.append(fp16_partitioned_group_flat) - flat_offset += total_elements - - self._move_to_flat_buffer(sub_group, - fp16_partitioned_group_flat, - avoid_copy=not self.offload_param) - - # if necessary, create a pinned memory buffer to be used for swapping out - # params to NVME after optimizer step - should_create_fp16_flat_reuse_buffer = any(flattened_partition_group is None - for flattened_partition_group in self.fp16_partitioned_groups_flat) - if should_create_fp16_flat_reuse_buffer: - max_partition_numel, largest_partition_numel = 0, None - for sub_group in self.fp16_groups: - total_elements = sum(t.partition_numel() for t in sub_group) - if total_elements > max_partition_numel: - largest_partition_numel = [t.ds_numel for t in sub_group] - max_partition_numel = total_elements - - assert len(largest_partition_numel) > 0, f'Unexpected that largest partition is empty' - self.fp16_groups[0][0].nvme_swapper.reserve_partitioned_swap_space(largest_partition_numel) - - def _get_parameter_partitions(self) -> List[Tensor]: - return [param.ds_tensor for sub_group in self.fp16_groups for param in sub_group] - - def _swap_in_sub_group_to_flat_buffer(self, flat_buffer, sub_group_id): - offset = 0 - elements_in_sub_group = sum([t.ds_numel for t in self.fp16_partitioned_groups[sub_group_id]]) - assert (flat_buffer.numel() == elements_in_sub_group) - for param, partitioned_param in zip(self.fp16_groups[sub_group_id], - self.fp16_partitioned_groups[sub_group_id]): - dest = flat_buffer.narrow(0, offset, partitioned_param.ds_numel) - if partitioned_param.status == PartitionedParamStatus.NOT_AVAILABLE: - print_rank_0( - f"Swapping in {param.ds_id} with elements {param.ds_numel} and partition {param.partition_numel()}" - ) - param.nvme_swapper.swap_in([param], async_op=False) - dest.data.copy_(partitioned_param.data) - param.nvme_swapper.remove_partition_and_release_buffers([param]) - print_rank_0(f"Swapping in {param.ds_id} done") - else: - dest.data.copy_(partitioned_param.data) - offset += partitioned_param.ds_numel - - def _create_next_swappable_fp32_groups(self): - reverse_order_indices = [i for i in range(len(self.fp32_partitioned_groups_flat))] - reverse_order_indices.reverse() - - next_group = None - for i in reverse_order_indices: - self.next_swappable_fp32_partitioned_groups.append(next_group) - if self._swappable_optimizer_subgroup(i): - next_group = self.fp32_partitioned_groups_flat[i] - - self.next_swappable_fp32_partitioned_groups.reverse() - - def _get_sub_group_partitions(self, sub_group_id): - sub_group_partitions = [] - for param, partitioned_param in zip(self.fp16_groups[sub_group_id], - self.fp16_partitioned_groups[sub_group_id]): - if partitioned_param.status == PartitionedParamStatus.NOT_AVAILABLE: - swap_path = param.nvme_swapper.get_path(param, True) - sub_group_partitions.append((partitioned_param, param.partition_numel(), swap_path)) - else: - sub_group_partitions.append((partitioned_param, partitioned_param.ds_numel, None)) - - return sub_group_partitions - - def _create_fp32_partitions(self): - cpu_memory_usage = 0 - cpu_memory_sub_groups = 0 - nvme_memory_usage = 0 - num_swappable_partitions = 0 - num_swap_from_nvme_partitions = 0 - num_swap_from_cpu_partitions = 0 - swap_from_nvme_memory_usage = 0 - swap_from_cpu_memory_usage = 0 - GIGA_BYTES = (1024**3) - - swappable_fp32_tensors = [] - swappable_fp16_src_tensors = [] - nvme_fp16_partitions_info = [] - nvme_fp16_num_elems = [] - nvme_fp32_dest_tensors = [] - fp32_element_size = torch.tensor([], dtype=torch.float32).element_size() - - # Assign portion of subgroup to cpu, the other to gpu. - if self.offload_optimizer: - self.subgroup_to_device = {} - sub_group_size = len(self.fp16_partitioned_groups_flat) - # print(f"Partial offload sub_group_size is {sub_group_size}, ratio is {self.partial_offload}\n") - for i in range(sub_group_size): - if i < int(self.partial_offload * sub_group_size): - self.subgroup_to_device[i] = 'cpu' - else: - self.subgroup_to_device[i] = get_accelerator()._name - - for i, tensor in enumerate(self.fp16_partitioned_groups_flat): - num_elements = self.fp16_partitioned_groups_flat_numel[i] - ds_id_begin = str(self.fp16_partitioned_groups_flat_id[i][0]) - ds_id_end = str(self.fp16_partitioned_groups_flat_id[i][-1]) - ds_id = ds_id_begin + '_' + ds_id_end - - # a partition of the fp32 master weights that will be updated by this process - if self._swappable_optimizer_subgroup(i): - self.fp32_partitioned_groups_flat.append(torch.Tensor()) - self.fp32_partitioned_groups_flat[i].ds_id = ds_id - nvme_memory_usage += (fp32_element_size * num_elements) - num_swappable_partitions += 1 - - if self.params_in_nvme_and_cpu and tensor is None: - num_swap_from_nvme_partitions += 1 - swap_from_nvme_memory_usage += (fp32_element_size * num_elements) - if self.offload_optimizer_fast_init: - sub_group_partitions = self._get_sub_group_partitions(i) - nvme_fp16_partitions_info.append(sub_group_partitions) - nvme_fp16_num_elems.append(num_elements) - nvme_fp32_dest_tensors.append(self.fp32_partitioned_groups_flat[i]) - else: - unpinned_fp32_buffer = torch.empty(num_elements, device=self.device, dtype=torch.float) - self._swap_in_sub_group_to_flat_buffer(unpinned_fp32_buffer, i) - self.optimizer_swapper.initialize_parameters(parameters=[self.fp32_partitioned_groups_flat[i]], - src_tensors=[unpinned_fp32_buffer]) - else: - num_swap_from_cpu_partitions += 1 - swap_from_cpu_memory_usage += (fp32_element_size * num_elements) - swappable_fp32_tensors.append(self.fp32_partitioned_groups_flat[i]) - swappable_fp16_src_tensors.append(self.fp16_partitioned_groups_flat[i]) - else: - cpu_memory_usage += (fp32_element_size * num_elements) - cpu_memory_sub_groups += 1 - - if self.params_in_nvme_and_cpu and tensor is None: - unpinned_fp32_buffer = torch.empty(num_elements, device=self.device, dtype=torch.float) - self._swap_in_sub_group_to_flat_buffer(unpinned_fp32_buffer, i) - self.fp32_partitioned_groups_flat.append(unpinned_fp32_buffer) - else: - if self.offload_optimizer: - self.fp32_partitioned_groups_flat.append(self.fp16_partitioned_groups_flat[i].to( - self.subgroup_to_device[i]).clone().float().detach()) - else: - self.fp32_partitioned_groups_flat.append(self.fp16_partitioned_groups_flat[i].to( - self.device).clone().float().detach()) - self.fp32_partitioned_groups_flat[i].ds_id = ds_id - - self.fp32_partitioned_groups_flat[i].requires_grad = True # keep this in case internal optimizer uses it - - if len(swappable_fp32_tensors) > 0: - self.optimizer_swapper.initialize_parameters(parameters=swappable_fp32_tensors, - src_tensors=swappable_fp16_src_tensors) - - if len(nvme_fp32_dest_tensors) > 0: - fp16_pinned_buffers = self.fp16_groups[0][0].nvme_swapper.reserve_available_buffers() - assert len(fp16_pinned_buffers) > 0 - self.optimizer_swapper.initialize_from_swapped_fp16_params(fp16_partitions_info=nvme_fp16_partitions_info, - fp16_num_elems=nvme_fp16_num_elems, - fp16_pinned_buffers=fp16_pinned_buffers, - fp32_parameters=nvme_fp32_dest_tensors) - self.fp16_groups[0][0].nvme_swapper.release_reserved_buffers() - - nvme_gigabytes = nvme_memory_usage / GIGA_BYTES - print_rank_0(f'Swappable FP32 Partitions: count={num_swappable_partitions} size={nvme_gigabytes:5.2f} GB', - force=False) - if self.params_in_nvme_and_cpu: - print_rank_0( - f'Swap from NVMe Partitions: count = {num_swap_from_nvme_partitions}, size = {swap_from_nvme_memory_usage/GIGA_BYTES:5.2f}GB', - force=False) - print_rank_0( - f'Swap from CPU Partitions: count = {num_swap_from_cpu_partitions}, size = {swap_from_cpu_memory_usage/GIGA_BYTES:5.2f}GB', - force=False) - - cpu_memory_gigabytes = cpu_memory_usage / GIGA_BYTES - print_rank_0(f'In-Memory FP32 Partitions: count={cpu_memory_sub_groups} size={cpu_memory_gigabytes:5.2f} GB', - force=False) - - # Clear for on-the-fly population before the optimizer step - for param_group in self.optimizer.param_groups: - param_group['params'] = [] - - def _create_fp16_sub_groups(self, params_group): - - params_group_numel = sum([param.partition_numel() for param in params_group]) - sub_group_size = self.sub_group_size - - if sub_group_size is None or sub_group_size >= params_group_numel: - return [params_group] - - sub_groups = [] - sub_group = [] - local_sub_group_size = 0 - for param in params_group: - - sub_group.append(param) - local_sub_group_size += param.partition_numel() - - if local_sub_group_size >= sub_group_size or id(param) == id(params_group[-1]): - - sub_groups.append(sub_group) - - sub_group = [] - local_sub_group_size = 0 - - return sub_groups - - def _optimizer_step(self, sub_group_id): - param_group_id = self.sub_group_to_group_id[sub_group_id] - fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] - if self.offload_optimizer: - cur_device = self.subgroup_to_device[sub_group_id] - if cur_device == 'cpu': - self.optimizer.param_groups[param_group_id]['params'] = [fp32_param] - cpu_loss = self.optimizer.step() - self.optimizer.param_groups[param_group_id]['params'] = [] - else: - self.backup_optimizer.param_groups[param_group_id]['params'] = [fp32_param] - gpu_loss = self.backup_optimizer.step() - self.backup_optimizer.param_groups[param_group_id]['params'] = [] - else: - self.optimizer.param_groups[param_group_id]['params'] = [fp32_param] - self.optimizer.step() - self.optimizer.param_groups[param_group_id]['params'] = [] - - def _swappable_optimizer_subgroup(self, sub_group_id): - if not self.swap_optimizer: - return False - - return self.optimizer_swapper.is_swappable_tensor(None, - numel=self.fp16_partitioned_groups_flat_numel[sub_group_id]) - - def _partitioned_params_swap_out(self, i): - offset = 0 - fp32_param = self.fp32_partitioned_groups_flat[i] - assert fp32_param is not None, \ - f'fp32 parameters of sub_group {i} is None' - - swap_fp16_params = [] - swap_fp32_params = [] - for param, partitioned_param in zip(self.fp16_groups[i], self.fp16_partitioned_groups[i]): - src = fp32_param.narrow(0, offset, partitioned_param.ds_numel) - if partitioned_param.status == PartitionedParamStatus.AVAILABLE: - partitioned_param.data.copy_(src.data) - else: - swap_fp32_params.append(src) - swap_fp16_params.append(param) - offset += partitioned_param.ds_numel - - if len(swap_fp16_params): - swap_fp16_params[0].nvme_swapper.swap_out_partitioned_params(dst_fp16_params=swap_fp16_params, - src_fp32_params=swap_fp32_params) - - def _set_fp16_partitioned_groups_flat(self): - # setup flat buffers per subgroup, these are each just sections of the - # contiguous flat buffer for all parameters that we created earlier - offset = 0 - for sub_group in self.fp16_groups: - sub_group_numel = sum(param.partition_numel() for param in sub_group) - self.fp16_partitioned_groups_flat.append(self.lp_param_buffer.narrow(0, offset, sub_group_numel)) - offset += sub_group_numel - - def initialize_optimizer_states(self): - num_subgroups = len(self.fp16_groups) - - largest_numel = max([sum([p.ds_numel for p in psg]) for psg in self.fp16_partitioned_groups]) - gradient_dtype = self.fp32_partitioned_groups_flat[0].dtype - gradient_buffer = torch.zeros(int(largest_numel), dtype=gradient_dtype, device=self.device) - - timer_names = set() - - # State initialization for the Adagrad optimizer occurs at construction as opposed to other optimizers - # which do lazy initialization of the state at the first call to step. - is_adagrad = isinstance(self.optimizer, torch.optim.Adagrad) - - if self.swap_optimizer: - self.optimizer_swapper.init_timers() - - timer_names.add(INIT_OPTIMIZER_TIMER) - self.timers(INIT_OPTIMIZER_TIMER).start() - - for i, group in enumerate(self.fp16_groups): - swappable_optimizer_subgroup = self._swappable_optimizer_subgroup(i) - swappable_param_subgroup = self.fp16_partitioned_groups_flat[i] is None - - num_elements = int(self.fp16_partitioned_groups_flat_numel[i]) - - see_memory_usage( - f'[Begin] Initialize optimizer states {i} / {num_subgroups} subgroups, num_elems: {num_elements}, swappable opt/param:{swappable_optimizer_subgroup}/{swappable_param_subgroup}', - force=False) - - if swappable_optimizer_subgroup: - self._optimizer_states_and_gradient_swap_in(i, timer_names) - - if self.offload_optimizer and not swappable_optimizer_subgroup: - subgroup_gradient_buffer = torch.zeros(num_elements, dtype=gradient_dtype, device=self.device) - if self.offload_optimizer_pin_memory: - subgroup_gradient_buffer = get_accelerator().pin_memory(subgroup_gradient_buffer) - - self.fp32_partitioned_groups_flat[i].grad = subgroup_gradient_buffer.to(self.subgroup_to_device[i]) - else: - self.fp32_partitioned_groups_flat[i].grad = gradient_buffer.narrow(0, 0, num_elements) - - if swappable_param_subgroup: - self._partitioned_params_swap_out(i) - - if swappable_optimizer_subgroup: - self._optimizer_states_and_gradient_swap_out(i, timer_names) - - see_memory_usage( - f'[End] Initialize optimizer states {i} / {num_subgroups} subgroups, num_elems: {num_elements}, swappable opt/param:{swappable_optimizer_subgroup}/{swappable_param_subgroup}', - force=False) - - # Initialize the optimizer states with the flattened fp32 partition. - if is_adagrad: - self.optimizer = torch.optim.Adagrad(self.fp32_partitioned_groups_flat, **self.optimizer.defaults) - - self.timers(INIT_OPTIMIZER_TIMER).stop() - self.timers.log(timer_names) - - if self.swap_optimizer: - self.optimizer_swapper.log_timers() - - if not self.offload_optimizer: - for group in self.fp32_partitioned_groups_flat: - group.grad = None - - # Reset steps - return - - ######################################################################### - #########################ZeRO Partition Gradients######################## - ######################################################################### - - def get_first_param_index(self, group_id, param_group, partition_id): - for index, param in enumerate(param_group): - param_id = self.get_param_id(param) - if partition_id in self.param_to_partition_ids[group_id][param_id]: - return index - return None - - def initialize_gradient_partitioning_data_structures(self): - - total_partitions = dist.get_world_size(group=self.dp_process_group) - - for i, param_group in enumerate(self.fp16_groups): - - self.param_to_partition_ids[i] = {} - self.is_partition_reduced[i] = {} - self.total_grads_in_partition[i] = {} - self.remaining_grads_in_partition[i] = {} - self.is_grad_computed[i] = {} - self.grad_partition_insertion_offset[i] = {} - self.grad_start_offset[i] = {} - self.first_param_index_in_partition[i] = {} - - for partition_id in range(total_partitions): - self.is_grad_computed[i][partition_id] = {} - self.grad_partition_insertion_offset[i][partition_id] = {} - self.grad_start_offset[i][partition_id] = {} - self.initialize_gradient_partition(i, param_group, partition_id) - self.is_partition_reduced[i][partition_id] = False - self.first_param_index_in_partition[i][partition_id] = self.get_first_param_index( - i, param_group, partition_id) - - @instrument_w_nvtx - def independent_gradient_partition_epilogue(self): - self.report_ipg_memory_usage(f"In ipg_epilogue before reduce_ipg_grads", 0) - for comm_dtype in sort_dtypes(self.ipg_buckets.keys()): - self.__reduce_and_partition_ipg_grads(comm_dtype) - self.report_ipg_memory_usage(f"In ipg_epilogue after reduce_ipg_grads", 0) - - if not get_accelerator().resolves_data_dependency(): - self.reduce_and_partition_stream.synchronize() - - for param_id in self.params_already_reduced.keys(): - self.params_already_reduced[param_id] = False - - #in case of cpu offload, averaged gradients are already in fp32_partitioned_groups_flat.grad - #TODO: use a similar code path for both cpu_offload and non-cpu offload - if not self.offload_optimizer: - for i, sub_group in enumerate(self.fp16_groups): - #TODO: This is redundant - self.averaged_gradients[i] = [ - self.__param_id_to_grad_partition[param.ds_id] - if param.requires_grad else torch.zeros_like(param.ds_tensor) for param in sub_group - ] - # this method gets called after every backward. need to increment - # here because if it gets incremented in backward() the micro step - # id will be off by one when we do the reduce and partition at the. - # start of this method. - # TODO. make this less error prone - self.micro_step_id += 1 - - def overlapping_partition_gradients_reduce_epilogue(self): - self.independent_gradient_partition_epilogue() - - def create_reduce_and_remove_grad_hooks(self): - print_rank_0(f'[Begin] Create gradient reduction hooks') - self.leaf_parameters = defaultdict(list) - for i, param_group in enumerate(self.fp16_groups): - for param in param_group: - if param.requires_grad: - #print_rank_0(f" Before all gather {param.device}, {param.shape}") - print_rank_0(f"Before all gather {param.device}, {param.shape}", force=False) - - # The hook must be created in un-partitioned parameter - param.all_gather() - - #print(f"After all gather {param.device}, {param.shape}") - def wrapper(param): - - @instrument_w_nvtx - def reduce_partition_and_remove_grads(*notneeded): - self.reduce_ready_partitions_and_remove_grads(param) - - self._grad_acc_hooks.append(register_grad_hook(param, reduce_partition_and_remove_grads)) - - #print(f"param grad fn {param.expand_as(param).grad_fn}") - if z3_leaf_parameter(param): - self.leaf_parameters[param.ds_z3_leaf_module].append(param) - else: - wrapper(param) - - # Partition the parameter after creating the hook - param.partition() - - # We delay reduce-scatter for all gradients in the leaf modules until the backward pass of the leaf module is done - for leaf_module, leaf_parameters in self.leaf_parameters.items(): - - def wrapper_pre_hook(params): - - def forward_pre_hook(module, input): - """Pre-forward hook to set backward hook on input tensors to the leaf module""" - module._leaf_module_inputs_remaining = 0 - - @instrument_w_nvtx - def reduce_leaf_module_grads(grad): - module._leaf_module_inputs_remaining -= 1 - # Make sure everything is done in the leaf module - if module._leaf_module_inputs_remaining == 0: - for param in params: - if param.grad is None: - param.grad = torch.zeros_like(param) - self.reduce_ready_partitions_and_remove_grads(param) - - def set_module_bwd_hook(tensor): - if tensor.requires_grad: - module._leaf_module_inputs_remaining += 1 - tensor.register_hook(reduce_leaf_module_grads) - return tensor - - output = apply_to_tensors_only(set_module_bwd_hook, input) - - return output - - return forward_pre_hook - - def wrapper_post_hook(): - - def forward_post_hook(module, input, output): - """Pre-forward hook to set backward hook on input tensors to the leaf module""" - module._leaf_output_required_grad_num = 0 - - def increment_rg_count_bwd_hook(tensor): - if tensor.requires_grad: - module._leaf_output_required_grad_num += 1 - return tensor - - apply_to_tensors_only(increment_rg_count_bwd_hook, output) - - if module._leaf_module_inputs_remaining == 0 and module._leaf_output_required_grad_num > 0: - raise RuntimeError( - "A module cannot be set as a leaf module when it does not have any input tensors that require gradients and has output tensors that require gradients. This is because the gradient reduction hook will not be called in this case." - ) - - return forward_post_hook - - self._leaf_module_hooks.append(leaf_module.register_forward_pre_hook(wrapper_pre_hook(leaf_parameters))) - self._leaf_module_hooks.append(leaf_module.register_forward_hook(wrapper_post_hook())) - - print_rank_0(f'[End] Create gradient reduction hooks') - - def get_param_id(self, param): - return OptimizerSwapper.parameter_id(param) - - ###############Independent Partition Gradient ######################## - def reduce_independent_p_g_buckets_and_remove_grads(self, param): - #print_rank_0(f"Inside reduce ipg buckets. {debug_param2name_id_shape(param)}, ipg elements {self.elements_in_ipg_bucket}, reduce bucket size {self.reduce_bucket_size}", force=True) - - # Because the ipg bucket is initialized with a random place holder tensor, we must - # explicitly check that the bucket has any real data in it (self.elements_in_ipg_bucket > - # 0). Otherwise if the incoming param.ds_numel is large, this branch may get triggered on a - # garbage data and `self.average_tensor()` will crash because its params_to_reduce will be - # empty, while reduction_list will have that garbage data. - comm_dtype = self.get_param_comm_dtype(param) - bucket = self.ipg_buckets[comm_dtype] - if bucket.elements + param.ds_numel > self.reduce_bucket_size and bucket.elements > 0: - self.report_ipg_memory_usage("In ipg_remove_grads before reduce_ipg_grads", param.ds_numel) - self.__reduce_and_partition_ipg_grads(comm_dtype) - - # deal with a use-case of transient grads that will be generated in a loop for the same computation involving some model params - e.g. when performing a tiled memory calculation that shards the normal single sub-module call into a loop over a shards. - if getattr(param, "ds_grad_is_ready", True): - self.__add_grad_to_ipg_bucket(param) - - @instrument_w_nvtx - @torch.no_grad() - def __add_grad_to_ipg_bucket(self, param: Parameter) -> None: - if not get_accelerator().resolves_data_dependency(): - self.reduce_and_partition_stream.wait_stream(get_accelerator().default_stream()) - - bucket = self.ipg_buckets[self.get_param_comm_dtype(param)] - if self.contiguous_gradients and bucket.elements + param.grad.numel() <= self.reduce_bucket_size: - # move the gradient to a contiguous buffer - with get_accelerator().stream(self.reduce_and_partition_stream): - # move the parameter's gradient to the contiguous flat buffer - new_grad_tensor = bucket.buffer.narrow(0, bucket.elements, param.grad.numel()).view_as(param.grad) - new_grad_tensor.copy_(param.grad, non_blocking=True) - if not get_accelerator().is_synchronized_device(): - param.grad.record_stream(get_accelerator().current_stream()) - param.grad.data = new_grad_tensor - - bucket.params.append(param) - bucket.elements += param.grad.numel() - - @instrument_w_nvtx - @torch.no_grad() - def __reduce_and_partition_ipg_grads(self, communication_data_type: torch.dtype, safe_mode: bool = False) -> None: - bucket = self.ipg_buckets[communication_data_type] - params_in_bucket = bucket.params - - if not params_in_bucket: - return - - for param in params_in_bucket: - if param.grad.numel() != param.ds_numel: - raise RuntimeError(f"{param.grad.numel()} != {param.ds_numel} Cannot reduce scatter " - f"gradients whose size is not same as the params") - - assert len(set(p.ds_id for p in params_in_bucket)) == len(params_in_bucket) - - while self.param_reduce_events and self.param_reduce_events[0].query(): - self.param_reduce_events.popleft() - if len(self.param_reduce_events) > self.max_param_reduce_events: - self.param_reduce_events.popleft().synchronize() - - with get_accelerator().stream(self.reduce_and_partition_stream): - if safe_mode: - assert_ints_same_as_other_ranks([p.ds_id for p in params_in_bucket]) - - if self.contiguous_gradients and bucket.elements <= self.reduce_bucket_size and not self.reduce_scatter: - grad_bucket = bucket.buffer.narrow(0, 0, bucket.elements) - grad_partitions = self.__avg_scatter_contiguous_grads(grad_bucket, communication_data_type) - else: - params_in_bucket.sort(key=lambda p: p.ds_id) - grad_partitions = self.__avg_scatter_grads(params_in_bucket, communication_data_type) - - self.partition_grads(params_in_bucket, grad_partitions) - - params_in_bucket.clear() - - if not get_accelerator().handles_memory_backpressure(): - event = get_accelerator().Event() - event.record() - self.param_reduce_events.append(event) - - @instrument_w_nvtx - def __avg_scatter_contiguous_grads(self, buffer_to_reduce: Tensor, - communication_data_type: torch.dtype) -> List[Tensor]: - dtype = buffer_to_reduce.dtype - if communication_data_type != dtype: - buffer_to_reduce = buffer_to_reduce.to(communication_data_type) - if self.postscale_gradients and self.gradient_predivide_factor != 1.0: - buffer_to_reduce = buffer_to_reduce.div_(self.gradient_predivide_factor) - - world_sz = dist.get_world_size(self.dp_process_group) - rank = dist.get_rank(self.dp_process_group) - buffer_to_reduce.div_(world_sz / float(self.sequence_parallel_size)) - - dist.all_reduce(buffer_to_reduce, group=self.dp_process_group) - - if self.postscale_gradients and self.gradient_predivide_factor != world_sz: - buffer_to_reduce = buffer_to_reduce.mul(self.gradient_predivide_factor) - - if communication_data_type != self.dtype: - buffer_to_reduce = buffer_to_reduce.to(self.dtype) - - grad_partitions = [] - grad_offset_in_buffer = 0 - for param in self.ipg_buckets[communication_data_type].params: - grad = param.grad - chunk_sz = math.ceil(grad.numel() / world_sz) - - start_offset = grad_offset_in_buffer + min(rank * chunk_sz, grad.numel()) - end_offset = grad_offset_in_buffer + min(rank * chunk_sz + chunk_sz, grad.numel()) - - partition = buffer_to_reduce[start_offset:end_offset] - if param.partition_numel() != partition.numel(): - padded_partition = torch.zeros(param.partition_numel(), device=grad.device, dtype=grad.dtype) - if partition.numel() > 0: - padded_partition[:partition.numel()] = partition - grad_partitions.append(padded_partition) - else: - grad_partitions.append(partition) - grad_offset_in_buffer += grad.numel() - - return grad_partitions - - @instrument_w_nvtx - def __avg_scatter_grads(self, params_to_reduce: List[Parameter], - communication_data_type: torch.dtype) -> List[Tensor]: - """average gradients and scatter partitions across ranks""" - - full_grads_for_rank = [p.grad for p in params_to_reduce] - if communication_data_type != self.dtype: - full_grads_for_rank = [g.to(communication_data_type) for g in full_grads_for_rank] - - if self.postscale_gradients and self.gradient_predivide_factor != 1.0: - full_grads_for_rank = [g.div(self.gradient_predivide_factor) for g in full_grads_for_rank] - - local_world_size = get_accelerator().device_count() - global_world_size = dist.get_world_size() - num_nodes = global_world_size // local_world_size - if self.all2all_process_group is not None and num_nodes > 1: - grad_partitions_for_rank = (all_to_all_loco_quant_reduce(params_to_reduce, self.all2all_process_group, - self.zeropp_loco_param) - if self.zeropp_loco_param is not None else all_to_all_quant_reduce( - full_grads_for_rank, self.all2all_process_group)) - else: - grad_partitions_for_rank = reduce_scatter_coalesced(full_grads_for_rank, self.dp_process_group) - - if self.postscale_gradients and self.gradient_predivide_factor != 1.0 and self.gradient_predivide_factor != dist.get_world_size( - self.dp_process_group): - grad_partitions_for_rank = [g.mul(self.gradient_predivide_factor) for g in grad_partitions_for_rank] - - if communication_data_type != self.dtype: - grad_partitions_for_rank = [g.to(self.dtype) for g in grad_partitions_for_rank] - - return grad_partitions_for_rank - - def set_grad_positions(self): - for i, group in enumerate(self.fp16_groups): - current_offset = 0 - for param in group: - param_id = self.get_param_id(param) - num_elements = param.partition_numel() - - self.grad_position[param_id] = [int(i), int(current_offset), int(num_elements)] - #print(f"param id {param_id} i:{i}, ds_tensor {num_elements} numel {param.numel()}") - current_offset += num_elements - see_memory_usage(f"After Set Grad positions", force=False) - - def _constant_buffered_norm2(self, input, buffer_size=250000000): - norm = None - for part in input.view(-1).split(buffer_size): - if norm is None: - norm = part.data.double().norm(2)**2.0 - else: - norm += part.data.double().norm(2)**2.0 - return norm**0.5 - - def set_norm_for_param_grad_in_gpu(self, param): - param_id = self.get_param_id(param) - #self.norm_for_param_grads[param_id] = param.grad.data.double().norm(2) - #Using a more memory efficient version - self.norm_for_param_grads[param_id] = self._constant_buffered_norm2(param.grad) - - def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param, fp32_grad_tensor): - with get_accelerator().stream(self.copy_grad_stream): - param_id = self.get_param_id(param) - src_tensor = param.grad.view(-1).float() - #print(f"src_tensor {src_tensor.size()} and fp32 grad {fp32_grad_tensor.size()}") - fp32_grad_tensor.copy_(src_tensor, non_blocking=True) - param.grad = None - - def complete_grad_norm_calculation_for_cpu_offload(self, params): - total_norm = 0.0 - norm_type = 2.0 - for p in params: - if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0): - param_id = self.get_param_id(p) - if param_id in self.norm_for_param_grads.keys(): - param_norm = self.norm_for_param_grads[param_id] - total_norm += param_norm**2 - - # Sum across all model parallel GPUs. - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=self.dp_process_group) - - self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.SUM) - - total_norm = total_norm_cuda[0]**(1. / norm_type) - - mask_nan_or_inf_with_val_inplace(total_norm, device=total_norm.device) - - return total_norm.cpu() - - @instrument_w_nvtx - def partition_grads(self, params_to_release: List[Parameter], grad_partitions: List[Tensor]) -> None: - offload_fp32_gradients = {} - offload_fp32_offsets = {} - buffers = [] - for param, grad_partition in zip(params_to_release, grad_partitions): - - contains_real_data = param.partition_numel() * dist.get_rank(self.dp_process_group) < param.ds_numel - if not contains_real_data: - # this grad partition is empty - don't need to do anything - param.grad = None - continue - - # move or accumulate gradient partition to target buffer - grad_buffer = self.__param_id_to_grad_partition[param.ds_id].narrow(0, 0, grad_partition.numel()) - buffers.append(grad_buffer) - if self.micro_step_id == 0: # don't accumulate - grad_buffer.copy_(grad_partition, non_blocking=True) - # ensure grad buffer is a CUDA buffer to speed up the next few - # operations and so it can be used asynchronously - grad_buffer = grad_buffer.to(grad_partition.device, non_blocking=True) - elif get_accelerator().on_accelerator(grad_buffer): - grad_buffer.add_(grad_partition.to(self.gradient_accumulation_dtype).view(grad_buffer.shape)) - else: - # if dst is CPU, copy first to src device, do the addition - # there, then move back to dst. adding directly to cpu is very slow - cuda_grad_buffer = grad_buffer.to(grad_partition.device, non_blocking=True) - cuda_grad_buffer.add_(grad_partition.to(self.gradient_accumulation_dtype).view(cuda_grad_buffer.shape)) - grad_buffer.copy_(cuda_grad_buffer, non_blocking=True) - # ensure grad buffer is a CUDA buffer to speed up the next few - # operations and so it can be used asynchronously - grad_buffer = cuda_grad_buffer - - # offload the gradient partition if applicable - if self.offload_optimizer: - i, dest_offset, _ = self.grad_position[self.get_param_id(param)] - - if self.is_gradient_accumulation_boundary: - self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer) - - if self._swappable_optimizer_subgroup(i): - if not i in offload_fp32_gradients.keys(): - offload_fp32_gradients[i] = [] - offload_fp32_offsets[i] = [] - - offload_fp32_gradients[i].append(grad_buffer.float()) - offload_fp32_offsets[i].append(dest_offset) - else: - fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow( - 0, dest_offset, grad_buffer.numel()) - fp32_grad_tensor.copy_(grad_buffer.float()) - - # free the gradient - if not get_accelerator().is_synchronized_device(): - if param.grad is not None: - param.grad.record_stream(get_accelerator().current_stream()) - param.grad = None - - if self.offload_optimizer and self.swap_optimizer: - for i in offload_fp32_gradients.keys(): - self.optimizer_swapper.swap_out_gradients(parameter=self.fp32_partitioned_groups_flat[i], - gradient_offsets=offload_fp32_offsets[i], - gradient_tensors=offload_fp32_gradients[i]) - return buffers - - def reduce_ready_partitions_and_remove_grads(self, param): - #print_rank_0(f"Backward {debug_param2name_id_shape(param)}", force=True) - self.reduce_independent_p_g_buckets_and_remove_grads(param) - - def zero_reduced_gradients(self, partition_id, i): - - def are_all_related_partitions_reduced(params_id): - for partition_id in self.param_to_partition_ids[i][params_id]: - if not self.is_partition_reduced[i][partition_id]: - return False - return True - - for params_id in self.is_grad_computed[i][partition_id]: - if are_all_related_partitions_reduced(params_id): - self.param_dict[params_id].grad = None - - def quantize_nontrainable_params(self): - """ In ZeRO-3, when the zero_quantized_nontrainable_weights flag is set, we quantize the non-trainable weights and also store them in quantized format. However, this check for trainable/non-trainable is done when deepspeed initializes the partitioning. So, if the user changes the trainable/non-trainable status of a parameter after the partitioning is done (e.g. LoRA), the user needs to re-quantize the non-trainable weights by calling this function. - """ - if not self.zero_quantized_nontrainable_weights: - print_rank_0( - f"Warning: quantize_nontrainable_params() called with zero_quantized_nontrainable_weights disabled, return without doing anything", - force=True) - return - quantizer_module = CUDAQuantizer() - - def quantize_dstensor(tensor): - assert tensor.dtype == torch.float16, f"quantize_dstensor() expects tensor.dtype == torch.float16, got {tensor.dtype}" - partition_size = tensor.ds_numel - ds_status = tensor.status - final_location = tensor.final_location - tensor, tensor.ds_quant_scale = quantizer_module.quantize(tensor) - tensor.ds_numel = partition_size - tensor.status = ds_status - tensor.final_location = final_location - tensor.requires_grad = False - return tensor - - for param in self.module.parameters(): - if hasattr(param, "ds_tensor") and (param.ds_tensor.numel() <= 2048 or param.ds_numel <= 500000): - # skip small parameters - continue - if hasattr(param, - "ds_tensor") and not param.requires_grad and not hasattr(param.ds_tensor, "ds_quant_scale"): - param.ds_tensor = quantize_dstensor(param.ds_tensor) - if hasattr(param, "ds_secondary_tensor") and not param.requires_grad and not hasattr( - param.ds_secondary_tensor, "ds_quant_scale") and param.ds_secondary_tensor is not None: - param.ds_secondary_tensor = quantize_dstensor(param.ds_secondary_tensor) - get_accelerator().synchronize() - - def flatten_and_print(self, message, tensors, start=0, n=5): - flatten_tensor = self.flatten(tensors) - - def print_func(): - logger.info(flatten_tensor.contiguous().view(-1).narrow(0, start, n)) - - self.sequential_execution(print_func, message) - - def get_grads_to_reduce(self, i, partition_id): - - def get_reducible_portion(key): - grad = self.param_dict[key].grad - total_elements = grad.numel() - start = self.grad_start_offset[i][partition_id][key] - num_elements = min(total_elements - start, - self.partition_size[i] - self.grad_partition_insertion_offset[i][partition_id][key]) - if not pg_correctness_test: - if num_elements == total_elements: - return grad - else: - return grad.contiguous().view(-1).narrow(0, int(start), int(num_elements)) - else: - if num_elements == total_elements: - return grad.clone() - else: - return grad.clone().contiguous().view(-1).narrow(0, int(start), int(num_elements)) - - grads_to_reduce = [] - for key in self.is_grad_computed[i][partition_id]: - grad = get_reducible_portion(key) - grads_to_reduce.append(grad) - return grads_to_reduce - - def sequential_execution(self, function, message, group=None): - if group is None: - group = self.dp_process_group - if dist.get_rank(group=group) == 0: - logger.info(message) - for id in range(dist.get_world_size(group=group)): - if id == dist.get_rank(group=group): - function() - dist.barrier(group=group) - - def set_none_gradients_to_zero(self, i, partition_id): - for param_id in self.is_grad_computed[i][partition_id]: - param = self.param_dict[param_id] - if param.grad is None: - param.grad = torch.zeros_like(param) - - ######################Reduction Related Methods############################## - - def allreduce_bucket(self, bucket, rank=None, log=None): - rank = None - tensor = self.flatten(bucket) - - tensor_to_allreduce = tensor - - if pg_correctness_test: - communication_data_type = torch.float32 - else: - communication_data_type = self.communication_data_type - - if communication_data_type != tensor.dtype: - tensor_to_allreduce = tensor.to(communication_data_type) - - tensor_to_allreduce.div_(dist.get_world_size(group=self.dp_process_group) / float(self.sequence_parallel_size)) - - if rank is None: - # "All Reducing" - dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group) - else: - global_rank = dist.get_global_rank(self.dp_process_group, rank) - dist.reduce(tensor_to_allreduce, global_rank, group=self.dp_process_group) - - if communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: - if rank is None or rank == dist.get_rank(group=self.dp_process_group): - tensor.copy_(tensor_to_allreduce) - - return tensor - - # if rank is specified do a reduction instead of an allreduce - def allreduce_and_copy(self, small_bucket, rank=None, log=None): - with get_accelerator().stream(self.reduction_stream): - allreduced = self.allreduce_bucket(small_bucket, rank=rank, log=log) - if rank is None or rank == dist.get_rank(group=self.dp_process_group): - for buf, synced in zip(small_bucket, self.unflatten(allreduced, small_bucket)): - buf.copy_(synced) - - def allreduce_no_retain(self, bucket, numel_per_bucket=500000000, rank=None, log=None): - small_bucket = [] - numel = 0 - for tensor in bucket: - small_bucket.append(tensor) - numel = numel + tensor.numel() - if numel > numel_per_bucket: - self.allreduce_and_copy(small_bucket, rank=rank, log=None) - small_bucket = [] - if len(small_bucket) > 0: - self.allreduce_and_copy(small_bucket, rank=rank, log=log) - - ############################################################################# - ############################################################################# - ############################################################################# - - # views the tensor as multiple partitions and returns - # those partitions - def get_data_parallel_partitions(self, tensor): - partitions = [] - - dp = dist.get_world_size(group=self.dp_process_group) - dp_id = dist.get_rank(group=self.dp_process_group) - - total_num_elements = tensor.numel() - - base_size = total_num_elements // dp - remaining = total_num_elements % dp - - start = 0 - for id in range(dp): - partition_size = base_size - if id < remaining: - partition_size = partition_size + 1 - partitions.append(tensor.narrow(0, start, partition_size)) - start = start + partition_size - return partitions - - def get_partition_info(self, tensor_list, partition_size, partition_id): - params_in_partition = [] - params_not_in_partition = [] - - start_index = partition_size * partition_id - end_index = partition_size * (partition_id + 1) - - current_index = 0 - first_offset = 0 - - for tensor in tensor_list: - - tensor_size = tensor.numel() - - if start_index <= current_index < end_index: - params_in_partition.append(tensor) - - elif current_index < start_index < (current_index + tensor_size): - params_in_partition.append(tensor) - - assert (first_offset == 0 - ), "This can happen either zero or only once as this must be the first tensor in the partition" - first_offset = start_index - current_index - - else: - params_not_in_partition.append(tensor) - - current_index = current_index + tensor_size - - return params_in_partition, params_not_in_partition, first_offset - - @instrument_w_nvtx - def zero_grad(self, set_to_none=True): - """ - Zero FP16 parameter grads. - """ - self.micro_step_id = 0 - - # FP32 grad should never exist. - # For speed, set model fp16 grad to None by default - for group in self.fp16_groups: - for p in group: - if set_to_none: - if p.grad is not None and get_accelerator().on_accelerator(p.grad): - p.grad.record_stream(get_accelerator().current_stream()) - p.grad = None - else: - if p.grad is not None: - p.grad.detach_() - p.grad.zero_() - - def _model_parallel_all_reduce(self, tensor, op): - """ Perform all reduce within model parallel group, if any. - """ - if self.model_parallel_group is None: - pass - else: - dist.all_reduce(tensor=tensor, op=op, group=self.model_parallel_group) - - @instrument_w_nvtx - def get_grad_norm_direct(self, gradients, params, norm_type=2): - """Clips gradient norm of an iterable of parameters. - - This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and - added functionality to handle model parallel parameters. Note that - the gradients are modified in place. - - Arguments: - parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a - single Tensor that will have gradients normalized - max_norm (float or int): max norm of the gradients - norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for - infinity norm. - - Returns: - Total norm of the parameters (viewed as a single vector). - """ - norm_type = float(norm_type) - if norm_type == inf: - total_norm = max(g.data.abs().max() for g in gradients) - total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=self.dp_process_group) - - # Take max across all GPUs. - self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.MAX) - total_norm = total_norm_cuda[0] - else: - # if dist.get_rank() == 0: - # logger.info(f"Total Norm beginning {total_norm}") - grad_norms = [] - for g, p in zip(gradients, params): - if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0): - grad_norms.append(g.to(get_accelerator().device_name(), non_blocking=True).double().norm(2)) - - # Sum across all model parallel GPUs. - if len(grad_norms) == 0: - # FIX https://github.com/deepspeedai/DeepSpeed/issues/3564 - total_norm_cuda = torch.tensor(0, - dtype=gradients[0].dtype).to(get_accelerator().device_name()).double() - else: - total_norm_cuda = torch.sum(torch.pow(torch.stack(grad_norms), 2)) - - dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=self.dp_process_group) - - self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.SUM) - - total_norm = total_norm_cuda**(1. / norm_type) - - norm_is_inf = total_norm.isinf() - norm_is_nan = total_norm.isnan() - inf_or_nan = norm_is_nan.logical_or(norm_is_inf) - - err = torch.tensor(-1.0, device=self.device, dtype=torch.float) - total_norm = torch.where(inf_or_nan, err, total_norm) - - return total_norm - - # creates a flat fused tensor from the tensor list starting at the first_offset - # in the first tensor of the list. If there are not enough elements in the tensor - # list then the flat tensor will be padded with zeros - def get_flat_partition(self, tensor_list, first_offset, partition_size, return_tensor_list=False): - flat_tensor_list = [] - current_size = 0 - for i, tensor in enumerate(tensor_list): - if tensor.grad is None: - tensor.grad = torch.zeros_like(tensor) - - tensor = tensor.grad - num_elements = tensor.numel() - tensor_offset = 0 - - # we need to offset to get to the right element - if i == 0 and first_offset > 0: - tensor_offset = first_offset - num_elements = num_elements - tensor_offset - - # we dont need all elements of the tensor - if num_elements > (partition_size - current_size): - num_elements = partition_size - current_size - - # we need a narrow view of the tensor based on the tensor offset and number of elements that - # we need from this tensor - if tensor_offset > 0 or num_elements < tensor.numel(): - flat_tensor_list.append(tensor.contiguous().view(-1).narrow(0, int(tensor_offset), int(num_elements))) - else: - flat_tensor_list.append(tensor) - - current_size = current_size + num_elements - - # this means its the last partition and does not align with the dp boundary. We need to pad before flattening - if current_size < partition_size: - flat_tensor_list.append( - torch.zeros(int(partition_size - current_size), - dtype=tensor_list[0].dtype, - device=tensor_list[0].device)) - - if return_tensor_list: - return flat_tensor_list - - return self.flatten(flat_tensor_list) - - def free_grad_in_param_list(self, param_list): - for p in param_list: - p.grad = None - - def reset_cpu_buffers(self): - self.norm_for_param_grads = {} - - def _pre_step(self): - self.micro_step_id = 0 - - print_rank_0(f"Inside Step function") - see_memory_usage(f"In step before checking overflow", force=False) - - print_rank_0("Finished Tracing at Beginning of Step") - self._get_param_coordinator().hierarchy = 0 - - print_rank_0("Finished Tracing at Beginning of Step") - - @instrument_w_nvtx - def _get_norm_groups(self): - norm_groups = [] - for i, group in enumerate(self.fp16_groups): - if self.offload_optimizer: - norm_groups.append(self.complete_grad_norm_calculation_for_cpu_offload(self.fp16_groups[i])) - else: - norm_groups.append(self.get_grad_norm_direct(self.averaged_gradients[i], self.fp16_groups[i])) - return norm_groups - - @instrument_w_nvtx - def _prepare_fp32_grad_for_sub_group(self, sub_group_id): - partition_id = dist.get_rank(group=self.dp_process_group) - - single_grad_partition = self.flatten(self.averaged_gradients[sub_group_id]).to( - self.fp32_partitioned_groups_flat[sub_group_id].dtype) - - assert single_grad_partition.numel() == self.fp32_partitioned_groups_flat[sub_group_id].numel(), \ - "averaged gradients have different number of elements that partition size {} {} {} {}".format( - single_grad_partition.numel(), self.fp32_partitioned_groups_flat[sub_group_id].numel(), sub_group_id, partition_id) - - self.fp32_partitioned_groups_flat[sub_group_id].grad = single_grad_partition - - # release all the gradient since we have already created a necessary copy in dp_grad_partition - self.zero_grad(set_to_none=True) - - if not get_accelerator().is_synchronized_device(): - for grad in filter(lambda g: get_accelerator().on_accelerator(g), self.averaged_gradients[sub_group_id]): - grad.record_stream(get_accelerator().current_stream()) - - self.averaged_gradients[sub_group_id] = None - - @instrument_w_nvtx - def _prepare_sub_group(self, sub_group_id, timer_names): - see_memory_usage(f'Before prepare optimizer sub group {sub_group_id}', force=False) - if self._swappable_optimizer_subgroup(sub_group_id): - self._optimizer_states_and_gradient_swap_in(sub_group_id, timer_names) - elif not self.offload_optimizer: - self._prepare_fp32_grad_for_sub_group(sub_group_id) - see_memory_usage(f'After prepare optimizer sub group {sub_group_id}', force=False) - - def _optimizer_states_and_gradient_swap_in(self, sub_group_id, timer_names=None): - param_length = self.fp16_partitioned_groups_flat_numel[sub_group_id] - fp32_param_id = self.get_param_id(self.fp32_partitioned_groups_flat[sub_group_id]) - assert self._swappable_optimizer_subgroup(sub_group_id), \ - f'Parameter {fp32_param_id} of numel={param_length} is not swappable' - - see_memory_usage(f'pre-step Before swapping in optimizer tensors {sub_group_id}', force=False) - if timer_names is not None: - timer_names.add(OPTIMIZER_SWAP_IN_STATE_TIMER) - self.timers(OPTIMIZER_SWAP_IN_STATE_TIMER).start() - - self.optimizer_swapper.swap_in_optimizer_state( - parameter=self.fp32_partitioned_groups_flat[sub_group_id], - async_parameter=self.next_swappable_fp32_partitioned_groups[sub_group_id]) - - if timer_names is not None: - self.timers(OPTIMIZER_SWAP_IN_STATE_TIMER).stop() - see_memory_usage(f'pre-step After swapping in optimizer tensors {sub_group_id}', force=False) - - @instrument_w_nvtx - def _release_sub_group(self, sub_group_id, timer_names): - see_memory_usage(f'Before release optimizer sub group {sub_group_id}', force=False) - # get rid of the fp32 gradients. Not needed anymore - if not self.offload_optimizer: - self.fp32_partitioned_groups_flat[sub_group_id].grad = None - - if self._swappable_optimizer_subgroup(sub_group_id): - self._optimizer_states_and_gradient_swap_out(sub_group_id, timer_names) - see_memory_usage(f'After release optimizer sub group {sub_group_id}', force=False) - - # create a flat tensor aligned at the alignment boundary - @instrument_w_nvtx - def flatten_dense_tensors_aligned(self, tensor_list, alignment): - num_elements = 0 - for tens in tensor_list: - num_elements = num_elements + tens.numel() - - remaining = num_elements % alignment - - if remaining: - elements_to_add = alignment - remaining - pad_tensor = torch.zeros(elements_to_add, device=tensor_list[0].device, dtype=tensor_list[0].dtype) - padded_tensor_list = tensor_list + [pad_tensor] - - num_elements = num_elements + elements_to_add - else: - padded_tensor_list = tensor_list - - return self.flatten(padded_tensor_list) - - def _optimizer_states_and_gradient_swap_out(self, sub_group_id, timer_names=None): - param_length = self.fp16_partitioned_groups_flat_numel[sub_group_id] - fp32_param_id = self.get_param_id(self.fp32_partitioned_groups_flat[sub_group_id]) - assert self._swappable_optimizer_subgroup(sub_group_id), \ - f'Parameter {fp32_param_id} of numel={param_length} is not swappable' - - see_memory_usage(f'post-step Before swapping out optimizer tensors {sub_group_id}', force=False) - if timer_names is not None: - timer_names.add(OPTIMIZER_SWAP_OUT_STATE_TIMER) - self.timers(OPTIMIZER_SWAP_OUT_STATE_TIMER).start() - - self.optimizer_swapper.swap_out_optimizer_state( - parameter=self.fp32_partitioned_groups_flat[sub_group_id], - async_swap=self.next_swappable_fp32_partitioned_groups[sub_group_id] is not None) - - if timer_names is not None: - self.timers(OPTIMIZER_SWAP_OUT_STATE_TIMER).stop() - see_memory_usage(f'post-step After swapping out optimizer tensors {sub_group_id}', force=False) - - # get rid of the fp32 gradients. Not needed anymore - self.fp32_partitioned_groups_flat[sub_group_id].grad = None - - def _release_swap_buffers(self, sub_group_id): - self.optimizer_swapper.release_swap_buffers(parameter=self.fp32_partitioned_groups_flat[sub_group_id]) - self.fp32_partitioned_groups_flat[sub_group_id].grad = None - - def _writeback_swap_state(self, sub_group_id, write_opt_state, write_gradients): - self.optimizer_swapper.writeback_optimizer_state_and_gradients(self.fp32_partitioned_groups_flat[sub_group_id], - write_opt_state, write_gradients) - self.fp32_partitioned_groups_flat[sub_group_id].grad = None - - def _unflatten_partitioned_parameters(self, sub_group_id): - updated_params = self.unflatten(self.fp16_partitioned_groups_flat[sub_group_id], - self.fp16_partitioned_groups[sub_group_id]) - - for partitioned_param, q in zip(self.fp16_partitioned_groups[sub_group_id], updated_params): - partitioned_param.data = q.data - - def _overflow_clean_up(self, prev_scale): - see_memory_usage('After overflow before clearing gradients', force=False) - self.zero_grad(set_to_none=True) - - if self.offload_optimizer: - self.reset_cpu_buffers() - else: - self.averaged_gradients = {} - - see_memory_usage('After overflow after clearing gradients', force=False) - - def _loco_err_buf_update(self, overflow: bool, scale=1.0): - """ - Loco Error Buffer update. - """ - if not overflow and scale == 1.0: return - if dist.get_rank() == 0: - logger.info(f"update loco-zero++ error buffer with overflow: {overflow}") - # FP32 grad should never exist. - # For speed, set model fp16 grad to None by default - for group in self.fp16_groups: - for p in group: - if hasattr(p, 'intra_ef_buf'): - if overflow: - del p.intra_ef_buf - del p.inter_ef_buf - else: - p.intra_ef_buf[1] *= scale - p.inter_ef_buf[1] *= scale - - @instrument_w_nvtx - def _overflow_check_and_loss_scale_update(self): - - # First compute norm for all group so we know if there is overflow - if self.dtype == torch.float16: - self.check_overflow() - - #loss scaling related computation - prev_scale = self.loss_scale - self._update_scale(self.overflow) - - if self.overflow: - self._overflow_clean_up(prev_scale) - - #update loco error buf - self._loco_err_buf_update(self.overflow, self.loss_scale / prev_scale) - - return self.overflow - - @instrument_w_nvtx - def _post_step(self, timer_names): - if self.offload_optimizer: - self.reset_cpu_buffers() - - #Gathering persisting parameters - if len(self.persistent_parameters) > 0: - self.persistent_parameters[0].all_gather(self.persistent_parameters) - - if self.swap_optimizer: - self.optimizer_swapper.log_timers() - - self.invalidate_secondary_tensor() - - self.timers.log(timer_names) - - see_memory_usage('After zero_optimizer step', force=False) - print_rank_0(f"------------------Finishing Step-----------------------") - - @instrument_w_nvtx - def _reassign_or_swap_out_partitioned_parameters(self, sub_group_id): - if self.fp16_partitioned_groups_flat[sub_group_id] is not None: - self.fp16_partitioned_groups_flat[sub_group_id].data.copy_( - self.fp32_partitioned_groups_flat[sub_group_id].data) - - #unflatten fp16 parameter subgroup - self._unflatten_partitioned_parameters(sub_group_id) - else: - self._partitioned_params_swap_out(sub_group_id) - - def override_loss_scale(self, loss_scale): - if loss_scale != self.external_loss_scale: - logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}') - self.custom_loss_scaler = True - self.external_loss_scale = loss_scale - - @instrument_w_nvtx - def step(self, closure=None): - """ - Not supporting closure. - """ - self._pre_step() - self._partition_all_parameters() - - #checks for overflow, adjust the loss scale accordingly - if self._overflow_check_and_loss_scale_update(): - if self.swap_optimizer: - self.optimizer_swapper.log_timers() - return - - norm_groups = self._get_norm_groups() - scaled_global_grad_norm = torch.linalg.vector_norm(torch.stack(norm_groups)) - - # Stash unscaled gradient norm - self._global_grad_norm = scaled_global_grad_norm / self.loss_scale - - timer_names = set() - - timer_names.add(OPTIMIZER_STEP_TIMER) - self.timers(OPTIMIZER_STEP_TIMER).start() - - #update parameters one sub group at a time - for sub_group_id, group in enumerate(self.fp16_groups): - - #prepare optimizer states, gradients and fp32 parameters for update - self._prepare_sub_group(sub_group_id, timer_names) - - #scale the fp32 gradients - self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm) - - #apply the optimizer step on the sub group and copy fp32 parameters to fp16 - self._optimizer_step(sub_group_id) - - #put fp16 parameters in appropriate location - self._reassign_or_swap_out_partitioned_parameters(sub_group_id) - - #release memory or swap out optimizer states of fp32 parameters - self._release_sub_group(sub_group_id, timer_names) - - self.timers(OPTIMIZER_STEP_TIMER).stop() - - self._post_step(timer_names) - - # warn user about caching allocator flushes - memory_stats = get_accelerator().memory_stats() - alloc_retries = memory_stats.get("num_alloc_retries") - if alloc_retries is None: - alloc_retries = 0 - if alloc_retries > self.n_caching_allocator_flushes: - if dist.get_rank() == 0: - logger.warning( - "%d pytorch allocator cache flushes since last step. this happens " - "when there is high memory pressure and is detrimental to " - "performance. if this is happening frequently consider adjusting " - "settings to reduce memory consumption. If you are unable to " - "make the cache flushes go away consider adding " - "get_accelerator().empty_cache() calls in your training loop to ensure " - "that all ranks flush their caches at the same time", - alloc_retries - self.n_caching_allocator_flushes) - self.n_caching_allocator_flushes = alloc_retries - - def dump_pre_step_gradients(self, debug_fp32_grads): - # Dump gradient norms for debugging - for i, _ in enumerate(self.fp16_groups): - print(f'Pre-Step Dump Norms for Group {i} FP16P, FP16G, FP32G, FP32GUC') - for fp16_param, fp32_grad in zip(self.fp16_groups[i], debug_fp32_grads[i]): - param_id = self.get_param_id(fp16_param) - fp16_grad_norm = self.debug_fp16_grads[i][param_id] - - fp32_grad_norm = [float(t.data.float().norm(2)) for t in fp32_grad] - norm_list = [fp16_grad_norm, fp32_grad_norm] - print(f'Pre-Step Norms {i} {param_id} = {norm_list}') - - def dump_post_step_gradients(self): - # Dump gradient norms for debugging - for i, group in enumerate(self.fp16_groups): - print(f'Post-Step Dump Norms for Group {i} FP16P, FP16DS, FP16FLAT, FP32FLAT') - unflat_fp16 = self.unflatten(self.fp16_groups_flat[i], self.fp16_groups[i]) - unflat_fp32 = self.unflatten(self.fp32_partitioned_groups_flat[i], self.fp16_groups[i]) - for j, p in enumerate(self.fp16_groups[i]): - param_id = self.get_param_id(p) - param_norm = float(p.data.float().norm(2)) - ds_norm = float(p.ds_tensor.data.float().norm(2)) - - unflat_norm = [float(t.data.float().norm(2)) for t in [unflat_fp16[j], unflat_fp32[j]]] - norm_list = [param_norm, ds_norm] + unflat_norm - print(f'Post-Step Norms {i} {param_id} = {norm_list}') - - @instrument_w_nvtx - def unscale_and_clip_grads(self, sub_group_id, total_norm): - # compute combined scale factor for this group - combined_scale = self.loss_scale - if self.clip_grad > 0.: - # norm is in fact norm*scale - clip = ((total_norm / self.loss_scale) + 1e-6) / self.clip_grad - clip = torch.clamp(clip, min=1.0) - combined_scale = clip * self.loss_scale - - self.fp32_partitioned_groups_flat[sub_group_id].grad.mul_(1. / combined_scale) - - def _check_overflow(self, partition_gradients=True): - self.overflow = self.has_overflow(partition_gradients) - - # `params` is a list / generator of torch.Variable - def has_overflow_serial(self, params, is_grad_list=False): - for p in params: - if p.grad is not None and self._has_inf_or_nan(p.grad.data): - return True - - return False - - def has_overflow_partitioned_grads_serial(self): - for i in range(len(self.fp16_groups)): - for j, grad in enumerate(self.averaged_gradients[i]): - if grad is not None and self._has_inf_or_nan(grad.data, j): - return True - return False - - @instrument_w_nvtx - def has_overflow(self, partition_gradients=True): - if partition_gradients: - with get_accelerator().stream(self.reduce_and_partition_stream): - if hasattr(self.inf_or_nan_tracker, "logical_or_"): - self.inf_or_nan_tracker.logical_or_(torch.isinf(self.grad_partitions_flat_buffer).any()) - self.inf_or_nan_tracker.logical_or_(torch.isnan(self.grad_partitions_flat_buffer).any()) - else: - # logical_or_ not available in older versions of pytorch - self.inf_or_nan_tracker += torch.isinf(self.grad_partitions_flat_buffer).any() - self.inf_or_nan_tracker += torch.isnan(self.grad_partitions_flat_buffer).any() - self.inf_or_nan_tracker = self.inf_or_nan_tracker > 0 - - overflow_gpu = self.inf_or_nan_tracker.clone().to(get_accelerator().current_device_name()).to( - torch.uint8) - self.inf_or_nan_tracker.zero_() - - if not get_accelerator().resolves_data_dependency(): - get_accelerator().default_stream().wait_stream(self.reduce_and_partition_stream) - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.dp_process_group) - - else: - params = [] - for group in self.fp16_groups: - for param in group: - params.append(param) - - overflow = self.has_overflow_serial(params, is_grad_list=partition_gradients) - overflow_gpu = get_accelerator().ByteTensor([overflow]) - - # Since each model parallel GPU carries only part of the model, - # make sure overflow flag is synced across all the model parallel GPUs - self._model_parallel_all_reduce(tensor=overflow_gpu, op=dist.ReduceOp.MAX) - - overflow = overflow_gpu[0].item() - return bool(overflow) - - # `x` is a torch.Tensor - @staticmethod - def _has_inf_or_nan(x, j=None): - try: - # if x is half, the .float() incurs an additional deep copy, but it's necessary if - # Pytorch's .sum() creates a one-element tensor of the same type as x - # (which is true for some recent version of pytorch). - cpu_sum = float(x.float().sum()) - # More efficient version that can be used if .sum() returns a Python scalar - # cpu_sum = float(x.sum()) - except RuntimeError as instance: - # We want to check if inst is actually an overflow exception. - # RuntimeError could come from a different error. - # If so, we still want the exception to propagate. - if "value cannot be converted" not in instance.args[0]: - raise - return True - else: - if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum: - return True - return False - - @instrument_w_nvtx - def backward(self, loss, retain_graph=False): - """ - :attr:`backward` performs the following steps: - - 1. fp32_loss = loss.float() - 2. scaled_loss = fp32_loss*loss_scale - 3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves - """ - if self.swap_optimizer: - self.optimizer_swapper.pre_backward() - - see_memory_usage(f"Before backward", force=False) - - if self.custom_loss_scaler: - scaled_loss = self.external_loss_scale * loss - scaled_loss.backward() - else: - self.loss_scaler.backward(loss.float(), retain_graph=retain_graph) - - if self.swap_optimizer: - self.optimizer_swapper.post_backward() - - def get_fp32_grad_partitions(self) -> Dict[int, Dict[int, Tensor]]: - """get fp32 gradient partition dictionary - accessed as grad_dict[parameter_group_index][parameter_index] - """ - if not get_accelerator().resolves_data_dependency(): - self.reduce_and_partition_stream.synchronize() - grad_dict = collections.defaultdict(dict) - if self.offload_optimizer: - for group in self.fp16_groups: - for param_idx, param in enumerate(group): - group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] - fp32_grad = self.fp32_partitioned_groups_flat[group_idx].grad.narrow(0, dest_offset, num_elements) - grad_dict[group_idx][param_idx] = fp32_grad - else: - for group_idx, group in self.averaged_gradients.items(): - for param_idx, gradient in enumerate(group): - grad_dict[group_idx][param_idx] = gradient.float() - - return grad_dict - - def _fp32_state_allgather(self, param, fp32_state_partition): - reduce_buffer = torch.empty(self.partition_count * fp32_state_partition.numel(), - dtype=torch.float32, - device=param.device) - my_rank = dist.get_rank(group=self.dp_process_group) - partition = reduce_buffer.narrow(0, fp32_state_partition.numel() * my_rank, fp32_state_partition.numel()) - partition.data.copy_(fp32_state_partition.data, non_blocking=False) - dist.all_gather_into_tensor(reduce_buffer, partition, group=self.dp_process_group) - return reduce_buffer.narrow(0, 0, param.ds_numel).view(param.ds_shape) - - def _get_fp32_grad_state_partition(self, param, release_swap_buffers): - if not get_accelerator().resolves_data_dependency(): - self.reduce_and_partition_stream.synchronize() - - group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] - if self.offload_optimizer: - if self._swappable_optimizer_subgroup(group_idx): - self._optimizer_states_and_gradient_swap_in(group_idx) - - fp32_grad = self.fp32_partitioned_groups_flat[group_idx].grad.narrow(0, dest_offset, num_elements) - - if self._swappable_optimizer_subgroup(group_idx) and release_swap_buffers: - self._release_swap_buffers(group_idx) - else: - fp32_grad = self.__param_id_to_grad_partition[param.ds_id] - - return fp32_grad, group_idx - - def get_fp32_grad_for_param(self, param) -> Tensor: - if not param.requires_grad: - return None - - fp32_grad, _ = self._get_fp32_grad_state_partition(param=param, release_swap_buffers=True) - fp32_grad = fp32_grad.to(get_accelerator().current_device_name()).float() - return self._fp32_state_allgather(param, fp32_grad) - - def set_fp32_grad_for_param(self, value, param): - if not param.requires_grad: - return - - # if not get_accelerator().resolves_data_dependency(): - # self.reduce_and_partition_stream.synchronize() - - # if self.offload_optimizer: - # group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] - # fp32_grad = self.fp32_partitioned_groups_flat[group_idx].grad.narrow(0, dest_offset, num_elements) - # else: - # fp32_grad = self.__param_id_to_grad_partition[param.ds_id] - - fp32_grad, group_idx = self._get_fp32_grad_state_partition(param=param, release_swap_buffers=False) - # import pdb; pdb.set_trace() - my_rank = dist.get_rank(group=self.dp_process_group) - value_partition = value.flatten().narrow(0, fp32_grad.numel() * my_rank, fp32_grad.numel()) - fp32_grad.data.copy_(value_partition.data) - - if self._swappable_optimizer_subgroup(group_idx): - self._writeback_swap_state(group_idx, write_opt_state=False, write_gradients=True) - - def _get_fp32_opt_state_partition(self, param, release_swap_buffers, optim_state_key=None): - if not get_accelerator().resolves_data_dependency(): - self.reduce_and_partition_stream.synchronize() - - group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] - - if self._swappable_optimizer_subgroup(group_idx): - self._optimizer_states_and_gradient_swap_in(group_idx) - - fp32_param = self.fp32_partitioned_groups_flat[group_idx] - if optim_state_key is None: - fp32_opt_state = fp32_param.narrow(0, dest_offset, num_elements) - else: - fp32_opt_state = self.optimizer.state[fp32_param][optim_state_key].narrow(0, dest_offset, num_elements) - - if self._swappable_optimizer_subgroup(group_idx) and release_swap_buffers: - self._release_swap_buffers(group_idx) - - return fp32_opt_state, group_idx - - def get_full_hp_param(self, param, optim_state_key=None) -> Tensor: - if not param.requires_grad: - return None - - # import pdb; pdb.set_trace() - fp32_opt_state, group_idx = self._get_fp32_opt_state_partition(param, - release_swap_buffers=True, - optim_state_key=optim_state_key) - fp32_opt_state = fp32_opt_state.to(get_accelerator().current_device_name()) - hp_param = self._fp32_state_allgather(param, fp32_opt_state) - - return hp_param - - def set_full_hp_param(self, value, param, optim_state_key=None): - if not param.requires_grad: - return - - assert value.numel( - ) == param.ds_numel, f" Number of elements do not match: {value.numel()} != {param.ds_numel}" - - fp32_opt_state_partition, group_idx = self._get_fp32_opt_state_partition(param, - release_swap_buffers=False, - optim_state_key=optim_state_key) - # print(f'{dist.get_rank()=} {fp32_opt_state_partition.shape=} -------- {value.shape=}') - # import pdb; pdb.set_trace() - my_rank = dist.get_rank(group=self.dp_process_group) - value_partition = value.flatten().narrow(0, - fp32_opt_state_partition.numel() * my_rank, - fp32_opt_state_partition.numel()) - fp32_opt_state_partition.data.copy_(value_partition.data) - - if self._swappable_optimizer_subgroup(group_idx): - self._optimizer_states_and_gradient_swap_out(group_idx) - - ### Local API START ### - def get_local_fp32_grad_for_param(self, param) -> Tensor: - if not param.requires_grad: - return None - - fp32_grad, _ = self._get_fp32_grad_state_partition(param=param, release_swap_buffers=True) - fp32_grad = fp32_grad.to(get_accelerator().current_device_name()).float() - return fp32_grad - - def set_local_grad_for_param(self, value, param): - if not param.requires_grad: - return - - assert value.numel() == param.ds_tensor.numel( - ), f" Number of elements do not match: {value.numel()} != {param.ds_tensor.ds_numel}" - - # if not get_accelerator().resolves_data_dependency(): - # self.reduce_and_partition_stream.synchronize() - - # if self.offload_optimizer: - # group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] - # fp32_grad = self.fp32_partitioned_groups_flat[group_idx].grad.narrow(0, dest_offset, num_elements) - # else: - # fp32_grad = self.__param_id_to_grad_partition[param.ds_id] - - if self.offload_optimizer: - self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(value) - - fp32_grad, group_idx = self._get_fp32_grad_state_partition(param=param, release_swap_buffers=False) - fp32_grad.data.copy_(value.flatten().data) - - if self._swappable_optimizer_subgroup(group_idx): - self._writeback_swap_state(group_idx, write_opt_state=False, write_gradients=True) - - def get_local_fp32_param(self, param, optim_state_key=None) -> Tensor: - if not param.requires_grad: - return None - fp32_opt_state, group_idx = self._get_fp32_opt_state_partition(param, - release_swap_buffers=True, - optim_state_key=optim_state_key) - fp32_opt_state = fp32_opt_state.to(get_accelerator().current_device_name()) - return fp32_opt_state - - def set_local_hp_param(self, value, param, optim_state_key=None): - if not param.requires_grad: - return - - assert hasattr(param, "ds_tensor"), f" The parameter does not contain the partitioned copy of the tensor." - assert value.numel() == param.ds_tensor.numel( - ), f" Number of elements do not match: {value.numel()} != {param.ds_tensor.ds_numel}" - - fp32_opt_state_partition, group_idx = self._get_fp32_opt_state_partition(param, - release_swap_buffers=False, - optim_state_key=optim_state_key) - value_partition = value.flatten() - fp32_opt_state_partition.data.copy_(value_partition.data) - - if self._swappable_optimizer_subgroup(group_idx): - self._optimizer_states_and_gradient_swap_out(group_idx) - # logger.info(f"[set_local_hp_param][update the params' value successfully]") - - ### Local API END ### - - ### Vectorized API BEGIN ### - def update_fp32_grad_for_param_vectorized(self, update_func, param_list): - params_with_grad = [p for p in param_list if p.requires_grad] - if not params_with_grad: - return - - if not get_accelerator().resolves_data_dependency(): - self.reduce_and_partition_stream.synchronize() - - subgroups = {} - for p in params_with_grad: - group_idx, dest_offset, num_elements = self.grad_position[self.get_param_id(p)] - param_entry = (p, dest_offset, num_elements) - if group_idx in subgroups.keys(): - subgroups[group_idx].append(param_entry) - else: - subgroups[group_idx] = [param_entry] - - for group_idx, group_params in subgroups.items(): - if self._swappable_optimizer_subgroup(group_idx): - self._optimizer_states_and_gradient_swap_in(group_idx) - - for param, dest_offset, num_elements in group_params: - if self.offload_optimizer: - fp32_grad_part = self.fp32_partitioned_groups_flat[group_idx].grad.narrow( - 0, dest_offset, num_elements) - else: - fp32_grad_part = self.__param_id_to_grad_partition[param.ds_id] - - fp32_grad_full = self._fp32_state_allgather(param, fp32_grad_part) - new_fp32_grad_full = update_func(fp32_grad_full, param) - my_rank = dist.get_rank(group=self.dp_process_group) - value_partition = new_fp32_grad_full.flatten().narrow(0, - fp32_grad_part.numel() * my_rank, - fp32_grad_part.numel()) - fp32_grad_part.data.copy_(value_partition.data) - - if self._swappable_optimizer_subgroup(group_idx): - self._writeback_swap_state(sub_group_id=group_idx, write_opt_state=False, write_gradients=True) - - ### Vectorized API END ### - - ### Device API BEGIN ### - def get_hp_param_device(self, param, optim_state_key=None) -> torch.device: - if not param.requires_grad: - return None - - fp32_opt_state, _ = self._get_fp32_opt_state_partition(param, - release_swap_buffers=True, - optim_state_key=optim_state_key) - return fp32_opt_state.device - - ### Device API END ### - - @instrument_w_nvtx - def _partition_all_parameters(self): - self.parameter_offload.partition_all_parameters() - - def check_overflow(self, partition_gradients=True): - self._check_overflow(partition_gradients) - - def _update_scale(self, has_overflow=False): - self.loss_scaler.update_scale(has_overflow) - - # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state" - def _get_state(self): - return self.optimizer.state - - def _set_state(self, value): - self.optimizer.state = value - - state = property(_get_state, _set_state) - - # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups" - # (for example, to adjust the learning rate) - def _get_param_groups(self): - return self.optimizer.param_groups - - def _set_param_groups(self, value): - self.optimizer.param_groups = value - self.trainable_param_groups = self._get_trainable_parameter_groups() - - param_groups = property(_get_param_groups, _set_param_groups) - - # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale" - def _get_loss_scale(self): - if self.custom_loss_scaler: - return self.external_loss_scale - else: - return self.loss_scaler.cur_scale - - def _set_loss_scale(self, value): - self.loss_scaler.cur_scale = value - - loss_scale = property(_get_loss_scale, _set_loss_scale) - cur_scale = property(_get_loss_scale, _set_loss_scale) - - def _get_lean_tensors(self, padded_flattened_tensor, group_tensors, paddings): - # Remove paddings from flattened tensor - individual_tensors = self.unflatten(padded_flattened_tensor, group_tensors) - lean_lengths = [t.numel() - pad for t, pad in zip(group_tensors, paddings)] - lean_tensors = [t[:len] for t, len in zip(individual_tensors, lean_lengths)] - #logger.info(f'rank {dist.get_rank()}: lean_tensors = {[t.numel() for t in lean_tensors]}') - return lean_tensors - - #TODO REVISIT this for stage 3 - def get_lean_optimizer_state(self): - # Return optimizer states after removing paddings. - # This method assumes that each param group contains a single flattened tensor. - optimizer_groups_state = [] - - for i, group in enumerate(self.optimizer.param_groups): - p = group['params'][0] - lean_state = {} - for key, value in self.optimizer.state[p].items(): - if torch.is_tensor(value): - padded_lens = [t.numel() for t in self.fp16_partitioned_groups[i]] - lean_state[key] = self._get_lean_tensors(value, self.fp16_partitioned_groups[i], - self.groups_padding[i]) - lean_flat_len = sum([t.numel() for t in lean_state[key]]) - else: - lean_state[key] = value - - optimizer_groups_state.append(lean_state) - - return optimizer_groups_state - - def get_groups_without_padding(self, groups_with_padding): - # Return group tensor after removing paddings added for alignment to DP world size. - groups_without_padding = [] - for i, group in enumerate(groups_with_padding): - lean_group = self._get_lean_tensors(group, self.fp16_partitioned_groups[i], self.groups_padding[i]) - groups_without_padding.append(lean_group) - - return groups_without_padding - - def _set_fp32_optimizer_param_groups(self): - for sub_group_id, _ in enumerate(self.fp16_groups): - param_group_id = self.sub_group_to_group_id[sub_group_id] - self.optimizer.param_groups[param_group_id]['params'].append( - self.fp32_partitioned_groups_flat[sub_group_id]) - - def _clear_fp32_optimizer_param_groups(self): - for param_group in self.optimizer.param_groups: - param_group['params'] = [] - - def _rigid_state_dict(self): - state_dict = {} - state_dict[ZERO_STAGE] = ZeroStageEnum.weights - state_dict[LOSS_SCALER] = self.loss_scaler - state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale - state_dict['overflow'] = self.overflow - state_dict[PARTITION_COUNT] = self.partition_count - - self._set_fp32_optimizer_param_groups() - state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict() - state_dict[FP32_FLAT_GROUPS] = self.fp32_partitioned_groups_flat - self._clear_fp32_optimizer_param_groups() - - return state_dict - - def state_dict(self): - """ - Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. - This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict - of the contained Pytorch optimizer. - Example:: - checkpoint = {} - checkpoint['model'] = model.state_dict() - checkpoint['optimizer'] = optimizer.state_dict() - torch.save(checkpoint, "saved.pth") - """ - if self.elastic_checkpoint: - raise NotImplementedError("ZeRO-3 does not yet support elastic checkpointing, please disable for now.") - - return self._rigid_state_dict() - - -# Restore base optimizer fp32 weights from checkpoint by: -# 1) Merging fp32 weights from checkpoints of all partitions -# 2) Extracting fp32 weights for current partition from merged weights -# 3) Using extracted weights to update base optimizer weights directly. - - def _restore_from_fp32_weights(self, all_state_dict): - - flat_local_partition = [] - for i in range(len(self.fp32_partitioned_groups_flat)): - merged_partitions = [sd['fp32_groups'][i] for sd in all_state_dict] - flat_local_partition.append(self._get_flattened_partition(merged_partitions)) - - for current, saved in zip(self.fp32_partitioned_groups_flat, flat_local_partition): - current.data.copy_(saved.data) - - # Restore base optimizer fp32 weights from ZeRO fp16 weights - def _restore_from_bit16_weights(self): - for fp16_partitions, fp32_partition in zip(self.fp16_partitioned_groups_flat, - self.fp32_partitioned_groups_flat): - fp32_partition.data.copy_(fp16_partitions.data) - - # Refresh the fp32 master params from the fp16 copies. - def refresh_fp32_params(self): - self._restore_from_bit16_weights() - - # Extract flattened partition for current rank from all partitions - def _get_flattened_partition(self, all_partition_states): - partition_id = dist.get_rank(group=self.dp_process_group) - alignment = dist.get_world_size(group=self.dp_process_group) - - param_partitions = [[] for _ in range(len(all_partition_states[0]))] - for i, partition in enumerate(all_partition_states): - for j, param in enumerate(partition): - param_partitions[j].append(param) - - local_state_partitions = [] - for param_index, param_slices in enumerate(param_partitions): - flattened_merged_tensor = self.flatten_dense_tensors_aligned(param_slices, alignment) - new_partitions = self.get_data_parallel_partitions(flattened_merged_tensor) - local_state_partitions.append(new_partitions[partition_id]) - - if torch.is_tensor(local_state_partitions[0]): - return self.flatten_dense_tensors_aligned(local_state_partitions, alignment) - - # Assume non-tensor states are not partitioned and equal across ranks, so return first one - return local_state_partitions[0] - - # Restore base optimizer state from checkpoint by - # 1) Merging optimizer state from checkpoints of all partitions - # 2) Extracting optimizer state for current partition from the merged state - # 3) Using the extracted value to directly update the base optimizer. - def _restore_base_optimizer_state(self, all_state_dict): - base_optimizer_group_states = [] - for i in range(len(self.optimizer.param_groups)): - partition_states = {} - all_partition_group_states = [sd['base_optimizer_state'][i] for sd in all_state_dict] - for key in all_partition_group_states[0].keys(): - all_partition_states = [all_states[key] for all_states in all_partition_group_states] - partition_states[key] = self._get_flattened_partition(all_partition_states) - base_optimizer_group_states.append(partition_states) - - for i, group in enumerate(self.optimizer.param_groups): - p = group['params'][0] - for key, saved in base_optimizer_group_states[i].items(): - if torch.is_tensor(self.optimizer.state[p][key]): - self.optimizer.state[p][key].data.copy_(saved.data) - else: - self.optimizer.state[p][key] = saved - - def _rigid_load_state_dict(self, state_dict, load_optimizer_states=True): - # I think it should actually be ok to reload the optimizer before the model. - self.loss_scaler = state_dict[LOSS_SCALER] - self.dynamic_loss_scale = state_dict['dynamic_loss_scale'] - self.overflow = state_dict['overflow'] - - if load_optimizer_states: - self._set_fp32_optimizer_param_groups() - self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT]) - self._clear_fp32_optimizer_param_groups() - - if self.swap_optimizer: - # Purge the swapped optimizer state, it was initialized to the freshly created model and not the checkpoint - self.optimizer_swapper.purge_state() - - if self.swap_optimizer: - # Touch all parameters to synchronize all buffers - timer_names = set() - self._partition_all_parameters() - for sub_group_id, group in enumerate(self.fp16_groups): - self._prepare_sub_group(sub_group_id, timer_names) - self._reassign_or_swap_out_partitioned_parameters(sub_group_id) - self._release_sub_group(sub_group_id, timer_names) - self._post_step(timer_names) - - # restore fp32 partitions - for curr_param, saved_param in zip(self.fp32_partitioned_groups_flat, state_dict[FP32_FLAT_GROUPS]): - curr_param.data.copy_(saved_param.data) - - # restore fp16 partitions from fp32 - for sub_group_id in range(len(self.fp32_partitioned_groups_flat)): - fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] - if sum(fp32_param.size()) > 0: - fp16_param = self.fp16_partitioned_groups_flat[sub_group_id] - fp16_param.data.copy_(fp32_param.data) - - # update fp16 unflattened params - for sub_group_id in range(len(self.fp16_partitioned_groups_flat)): - updated_params = self.unflatten(self.fp16_partitioned_groups_flat[sub_group_id], - self.fp16_partitioned_groups[sub_group_id]) - - for partitioned_param, q in zip(self.fp16_partitioned_groups[sub_group_id], updated_params): - partitioned_param.data = q.data - - # TODO: Support different/changing load/save DP degree. - def load_state_dict(self, - state_dict_list, - load_optimizer_states=True, - load_from_fp32_weights=False, - checkpoint_folder=None, - load_serial=None, - param_shapes=None): - r"""Loading a ZeRO checkpoint - Arguments: - state_dict_list: List of all saved ZeRO checkpoints, one for each saved partition. - Note that the number of saved partitions may differ from number of loading partitions to support - changing GPU count, specifically DP world size, between saving and loading checkpoints. - load_optimizer_states: Boolean indicating whether or not to load base optimizer states - load_from_fp32_weights: Boolean indicating whether to initialize fp32 master weights from fp32 - copies in checkpoints (no precision loss) or from model's fp16 copies (with precision loss). - """ - """ - Loads a state_dict created by an earlier call to state_dict(). - If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``, - whose parameters in turn came from ``model``, it is expected that the user - will call ``model.load_state_dict()`` before - ``fp16_optimizer_instance.load_state_dict()`` is called. - Example:: - model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half() - optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) - optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0) - ... - checkpoint = torch.load("saved.pth") - model.load_state_dict(checkpoint['model']) - optimizer.load_state_dict(checkpoint['optimizer']) - """ - - if self.elastic_checkpoint: - raise NotImplementedError("ZeRO-3 does not yet support elastic checkpointing, please disable for now.") - - if checkpoint_folder: - self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights, - param_shapes) - else: - self._rigid_load_state_dict(state_dict_list[dist.get_rank(group=self.dp_process_group)], - load_optimizer_states=load_optimizer_states) - - # when use loading checkpoint serial, after finish loading, we need to - # delete the temp state_dict_list variable to save memory, then trigger - # the next rank's loading - if load_serial is not None: - load_serial += 1 - rank = dist.get_rank(group=self.dp_process_group) - local_rank = dist.get_local_rank() - del state_dict_list[rank] - rank_end = dist.get_world_size() - 1 - if local_rank != rank_end: - dist.send(tensor=load_serial, dst=rank + 1) - - if len(self.persistent_parameters) > 0: - self.persistent_parameters[0].partition(self.persistent_parameters) - # self.persistent_parameters[0].all_gather(self.persistent_parameters) # this will be done in checkpoint_event_epilogue() so remove it to prevent double all_gather - - def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights, - param_shapes): - self.load_hp_checkpoint_state_from_checkpoint_dir_stage3(checkpoint_folder, param_shapes) - - def load_hp_checkpoint_state_from_checkpoint_dir_stage3(self, checkpoint_dir, param_shapes): - """ Load optimizer and model states from the checkpoint directory. """ - checkpoint_dir = os.path.join(checkpoint_dir, "zero") - optim_state_path = os.path.join(checkpoint_dir, "optimizer_state.pt") - assert os.path.isfile( - optim_state_path), f'{optim_state_path} containing optimizer global state is missing! Cannot proceed.' - - optim_sd = torch.load(optim_state_path, weights_only=False) - self._load_global_state_stage3(optim_sd) - - key_list = ["fp32", "exp_avg", "exp_avg_sq"] - - for key in key_list: - key_tensor = torch.empty(0) - for layer in param_shapes[0].keys(): - key_layer_state_partition = self.load_hp_checkpoint_state(os.path.join(checkpoint_dir, layer), key) - key_tensor = torch.cat((key_tensor, key_layer_state_partition)) - if key == "fp32": - self.fp32_partitioned_groups_flat[0].data.copy_(key_tensor) - self.optimizer.param_groups[0]['params'].append(self.fp32_partitioned_groups_flat[0]) - else: - optim_sd[OPTIMIZER_STATE_DICT]['state'][0][key] = key_tensor - - if self.swap_optimizer: - # Purge the swapped optimizer state, it was initialized to the freshly created model and not the checkpoint - self.optimizer_swapper.purge_state() - - if self.swap_optimizer: - # Touch all parameters to synchronize all buffers - timer_names = set() - self._partition_all_parameters() - for sub_group_id, group in enumerate(self.fp16_groups): - self._prepare_sub_group(sub_group_id, timer_names) - self._reassign_or_swap_out_partitioned_parameters(sub_group_id) - self._release_sub_group(sub_group_id, timer_names) - self._post_step(timer_names) - - self.optimizer.load_state_dict(optim_sd[OPTIMIZER_STATE_DICT]) - for param_group in self.optimizer.param_groups: - param_group['params'] = [] - - for sub_group_id in range(len(self.fp32_partitioned_groups_flat)): - fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] - if sum(fp32_param.size()) > 0: - fp16_param = self.fp16_partitioned_groups_flat[sub_group_id] - fp16_param.data.copy_(fp32_param.data) - - for sub_group_id in range(len(self.fp16_partitioned_groups_flat)): - updated_params = self.unflatten(self.fp16_partitioned_groups_flat[sub_group_id], - self.fp16_partitioned_groups[sub_group_id]) - - for partitioned_param, q in zip(self.fp16_partitioned_groups[sub_group_id], updated_params): - partitioned_param.data = q.data - - def _load_global_state_stage3(self, sd): - self.loss_scaler = sd.get(LOSS_SCALER, self.loss_scaler) - self.dynamic_loss_scale = sd.get('dynamic_loss_scale', self.dynamic_loss_scale) - self.overflow = sd.get('overflow', self.overflow) - - def load_hp_checkpoint_state(self, folder, key): - local_rank = dist.get_local_rank() - - # Load tensors from files and reshape them to flat vectors - loaded_checkpoint_state = torch.load(os.path.join(folder, f"{key}.pt"), weights_only=False).view(-1) - - # Partition the loaded data according to the local rank - world_size = dist.get_world_size(group=self.dp_process_group) - unpartitioned_numel = loaded_checkpoint_state.numel() - partitioned_numel = math.ceil(unpartitioned_numel / world_size) - - if world_size * partitioned_numel != unpartitioned_numel: - padding_size = world_size * partitioned_numel - unpartitioned_numel - padding_tensor = torch.zeros(padding_size, dtype=loaded_checkpoint_state.dtype) - loaded_checkpoint_state = torch.cat([loaded_checkpoint_state, padding_tensor]) - checkpoint_state_partition = loaded_checkpoint_state.narrow(0, local_rank * partitioned_numel, - partitioned_numel) - - return checkpoint_state_partition - - def reset_swap_buffers(self): - timer_names = set() - for sub_group_id, group in enumerate(self.fp16_groups): - self._prepare_sub_group(sub_group_id, timer_names) - self._reassign_or_swap_out_partitioned_parameters(sub_group_id) - self._release_sub_group(sub_group_id, timer_names) - - def checkpoint_event_prologue(self): - self._partition_all_parameters() - - def checkpoint_event_epilogue(self): - if len(self.persistent_parameters) > 0: - self.persistent_parameters[0].all_gather(self.persistent_parameters) - - def empty_partition_cache(self): - self.parameter_offload.empty_partition_cache() - - def offload_states(self, - include: Container[OffloadStateTypeEnum] = None, - device: OffloadDeviceEnum = OffloadDeviceEnum.cpu, - pin_memory: bool = True, - non_blocking: bool = False): - device = device.value - - self.empty_partition_cache() - - assert self.optimizer.__class__ == deepspeed.ops.adam.fused_adam.FusedAdam, f"Offloading is supported only for DeepSpeed FusedAdam." - - def needs_offload(target): - # return True - return target not in self.offloaded_states and (include == None or target in include) - - # HP param - if needs_offload(OffloadStateTypeEnum.hp_params): - if pin_memory: - if not hasattr(self, "hp_params_pin_buffers"): - self.hp_params_pin_buffers = [ - get_accelerator().pin_memory(torch.empty_like(t, device=device)) - for t in self.fp32_partitioned_groups_flat - ] - - for src_tensor, dest_buf in zip(self.fp32_partitioned_groups_flat, self.hp_params_pin_buffers): - dest_buf.copy_(src_tensor, non_blocking=non_blocking) - src_tensor.data = dest_buf - else: - for buf in self.fp32_partitioned_groups_flat: - buf.data = buf.data.to(device, non_blocking=non_blocking) - self.offloaded_states.add(OffloadStateTypeEnum.hp_params) - - # LP param - if needs_offload(OffloadStateTypeEnum.lp_params): - if pin_memory: - if not hasattr(self, "lp_param_contiguous_pin_buffer"): - self.lp_param_contiguous_pin_buffer = get_accelerator().pin_memory( - torch.empty_like(self.lp_param_buffer, device=device)) - self.lp_param_contiguous_pin_buffer.copy_(self.lp_param_buffer, non_blocking=non_blocking) - cpu_buffer = self.lp_param_contiguous_pin_buffer - else: - cpu_buffer = self.lp_param_buffer.to(device, non_blocking=non_blocking) - - self.lp_param_buffer.data = cpu_buffer - for tensor, offset, tensor_numel in get_mapping_to_flat_buffer( - [p.ds_tensor for p in self.module.parameters()]): - tensor.data = cpu_buffer.narrow(0, offset, tensor_numel) - - self.fp16_partitioned_groups_flat.clear() - self.offloaded_states.add(OffloadStateTypeEnum.lp_params) - - # LP grad - if needs_offload(OffloadStateTypeEnum.lp_grads): - if pin_memory: - if not hasattr(self, "lp_grad_partitions_flat_pin_buffers"): - self.lp_grad_partitions_flat_pin_buffers = get_accelerator().pin_memory( - torch.empty_like(self.grad_partitions_flat_buffer, device=device)) - self.lp_grad_partitions_flat_pin_buffers.copy_(self.grad_partitions_flat_buffer, - non_blocking=non_blocking) - self.grad_partitions_flat_buffer.data = self.lp_grad_partitions_flat_pin_buffers - else: - self.grad_partitions_flat_buffer.data = self.grad_partitions_flat_buffer.data.to(device) - self.averaged_gradients = {} - - self.__param_id_to_grad_partition = {} - - self.offloaded_states.add(OffloadStateTypeEnum.lp_grads) - - # contiguous bucket - if needs_offload(OffloadStateTypeEnum.contiguous_grad_buffer): - for bucket in self.ipg_buckets.values(): - if bucket.buffer is not None: - # Record properties like shape, strides, etc. as a meta tensor - bucket.buffer_meta = bucket.buffer.to("meta") - bucket.buffer = None - self.offloaded_states.add(OffloadStateTypeEnum.contiguous_grad_buffer) - - # Adam - if needs_offload(OffloadStateTypeEnum.optim_states): - offload_adam_states(self.optimizer, device, pin_memory=pin_memory, non_blocking=non_blocking) - self.offloaded_states.add(OffloadStateTypeEnum.optim_states) - - gc.collect() - get_accelerator().empty_cache() - - def reload_states(self, non_blocking: bool = False): - - device = get_accelerator().current_device_name() - - # HP param - if OffloadStateTypeEnum.hp_params in self.offloaded_states: - if hasattr(self, "hp_params_pin_buffers"): - for src, dest in zip(self.hp_params_pin_buffers, self.fp32_partitioned_groups_flat): - dest.data = src.to(device, non_blocking=non_blocking) - else: - for buf in self.fp32_partitioned_groups_flat: - buf.data = buf.data.to(device, non_blocking=non_blocking) - self.offloaded_states.remove(OffloadStateTypeEnum.hp_params) - - # LP Param - if OffloadStateTypeEnum.lp_params in self.offloaded_states: - cpu_buffer = self.lp_param_contiguous_pin_buffer if hasattr( - self, "lp_param_contiguous_pin_buffer") else self.lp_param_buffer - self.lp_param_buffer.data = cpu_buffer.data.to(device, non_blocking=non_blocking) - self._set_fp16_partitioned_groups_flat() - - parameter_partitions = self._get_parameter_partitions() - for tensor, offset, tensor_numel in get_mapping_to_flat_buffer(parameter_partitions): - tensor.data = self.lp_param_buffer.narrow(0, offset, tensor_numel) - self.offloaded_states.remove(OffloadStateTypeEnum.lp_params) - - # LP grad - if OffloadStateTypeEnum.lp_grads in self.offloaded_states: - if hasattr(self, "lp_grad_partitions_flat_pin_buffers"): - self.grad_partitions_flat_buffer.data = self.lp_grad_partitions_flat_pin_buffers.to( - device, non_blocking=non_blocking) - else: - self.grad_partitions_flat_buffer.data = self.grad_partitions_flat_buffer.data.to( - device, non_blocking=non_blocking) - self.averaged_gradients = {} - - offset = 0 - all_params = list(itertools.chain.from_iterable(self.fp16_groups)) - for param in all_params: - self.__param_id_to_grad_partition[param.ds_id] = self.grad_partitions_flat_buffer.narrow( - 0, offset, param.partition_numel()) - offset += param.partition_numel() - - self.offloaded_states.remove(OffloadStateTypeEnum.lp_grads) - - # contiguous bucket - if OffloadStateTypeEnum.contiguous_grad_buffer in self.offloaded_states: - for bucket in self.ipg_buckets.values(): - if bucket.buffer_meta is not None: - # We don't restore the data - bucket.buffer = torch.empty_like(bucket.buffer_meta, device=device) - - self.offloaded_states.remove(OffloadStateTypeEnum.contiguous_grad_buffer) - - # Adam - if OffloadStateTypeEnum.optim_states in self.offloaded_states: - reload_adam_states(self.optimizer, device, non_blocking=non_blocking) - self.offloaded_states.remove(OffloadStateTypeEnum.optim_states) - - if non_blocking: - get_accelerator().synchronize() - - -def _handle_overflow(cpu_sum, x, i): - import math - rank = dist.get_rank() - if rank == 0: - t_i = -1 - for v_i, v in enumerate(x.data.contiguous().view(-1)): - if not math.isfinite(float(v)): - t_i = v_i - break - logger.info(f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}") - - -def estimate_zero3_model_states_mem_needs(total_params, - largest_layer_params, - num_gpus_per_node=1, - num_nodes=1, - cpu_offload=True, - cpu_offload_params=True, - zero_init=True, - additional_buffer_factor=1.5): - - total_gpus = num_nodes * num_gpus_per_node - gpus_factor = 1 / num_nodes - largest_layer_memory = (4 * largest_layer_params) - - if cpu_offload: - if cpu_offload_params: - gpu_mem = largest_layer_memory - - if zero_init: - cpu_mem = total_params * 18 * gpus_factor * additional_buffer_factor - else: - cpu_mem = total_params * max(4 * num_gpus_per_node, 18 * gpus_factor) * additional_buffer_factor - - else: - gpu_mem = largest_layer_memory + int(2 * total_params / total_gpus) - - if zero_init: - cpu_mem = total_params * 16 * gpus_factor * additional_buffer_factor - else: - cpu_mem = total_params * max(4 * num_gpus_per_node, 16 * gpus_factor) * additional_buffer_factor - else: - gpu_mem = largest_layer_memory + int(18 * total_params / total_gpus) - if zero_init: - cpu_mem = largest_layer_params * 4 * num_gpus_per_node * additional_buffer_factor - else: - cpu_mem = total_params * 4 * num_gpus_per_node * additional_buffer_factor - - return int(cpu_mem), int(gpu_mem), largest_layer_memory - - -def model_to_params(model): - # shared params calculated only once - total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values()) - - largest_layer_params = 0 - for m in model.modules(): - # assuming no shared params within a single layer - layer_params = sum(p.numel() for p in m.parameters(recurse=False)) - largest_layer_params = max(largest_layer_params, layer_params) - - return total_params, largest_layer_params - - -def estimate_zero3_model_states_mem_needs_all_live(model, - num_gpus_per_node=1, - num_nodes=1, - additional_buffer_factor=1.5): - """ - Print out estimates on memory usage requirements for ZeRO 3 params, optim states and gradients - for a given ``model`` and hardware setup. - - If you have an actual model object, use this function and everything will be derived - automatically. - - If it's a hypothetical model, use ``estimate_zero3_model_states_mem_needs_all_cold`` where you have to pass - the ``total_params`` and ``largest_layer_params`` explicitly. - - Args: - - ``model``: ``nn.Module`` object - - ``num_gpus_per_node``: how many gpus per node (defaults to 1) - - ``num_nodes``: how many nodes (defaults to 1), - - ``additional_buffer_factor``: estimation factor (defaults to 1.5): - - """ - - total_params, largest_layer_params = model_to_params(model) - - estimate_zero3_model_states_mem_needs_all_cold(total_params=total_params, - largest_layer_params=largest_layer_params, - num_gpus_per_node=num_gpus_per_node, - num_nodes=num_nodes, - additional_buffer_factor=additional_buffer_factor) - - -def estimate_zero3_model_states_mem_needs_all_cold(total_params, - largest_layer_params, - num_gpus_per_node=1, - num_nodes=1, - additional_buffer_factor=1.5): - """ - Print out estimates on memory usage requirements for ZeRO 3 params, optim states and gradients - for a given ``model`` and hardware setup. - - If it's a hypothetical model, use this function where you have to pass - the ``total_params`` and ``largest_layer_params`` explicitly. - - If you have an actual model object, use ``estimate_zero3_model_states_mem_needs_all_live`` and everything - will be derived automatically. - - Args: - - ``total_params``: total model params - - ``largest_layer_params``: largest layer's params - - ``num_gpus_per_node``: how many gpus per node (defaults to 1) - - ``num_nodes``: how many nodes (defaults to 1), - - ``additional_buffer_factor``: estimation factor (defaults to 1.5): - - """ - - def format_options(cpu_offload, cpu_offload_params, zero_init): - enabled = [] - padded_cpu_str = f'{OffloadDeviceEnum.cpu:4}' - param_device = padded_cpu_str if cpu_offload_params else "none" - enabled.append(f"offload_param={param_device}") - optimizer_device = padded_cpu_str if cpu_offload else "none" - enabled.append(f"offload_optimizer={optimizer_device}") - enabled.append(f"zero_init={1 if zero_init else 0}") - return ", ".join(enabled) - - nodes_str = "nodes" if num_nodes > 1 else "node" - gpus_str = "GPUs" if num_gpus_per_node > 1 else "GPU" - print( - "Estimated memory needed for params, optim states and gradients for a:\n" - f"HW: Setup with {num_nodes} {nodes_str}, {num_gpus_per_node} {gpus_str} per node.\n" - f"SW: Model with {int(total_params/1e6)}M total params, {int(largest_layer_params/1e6)}M largest layer params." - ) - print(" per CPU | per GPU | Options") - for cpu_offload in [True, False]: - for cpu_offload_params in [True, False]: - if not cpu_offload and cpu_offload_params: - continue - for zero_init in [True, False]: - cpu_mem, gpu_mem, largest_layer_memory = estimate_zero3_model_states_mem_needs( - total_params=total_params, - largest_layer_params=largest_layer_params, - num_gpus_per_node=num_gpus_per_node, - num_nodes=num_nodes, - cpu_offload=cpu_offload, - cpu_offload_params=cpu_offload_params, - zero_init=zero_init, - additional_buffer_factor=additional_buffer_factor) - - options_str = format_options(cpu_offload=cpu_offload, - cpu_offload_params=cpu_offload_params, - zero_init=zero_init) - print(f" {cpu_mem/2**30:7.2f}GB | {gpu_mem/2**30:6.2f}GB | {options_str}") diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py deleted file mode 100644 index 7339f538643107a7fd19e476f4f831962a405ff7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py +++ /dev/null @@ -1,2642 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed import comm as dist -from packaging import version as pkg_version -from collections import OrderedDict, defaultdict -from dataclasses import dataclass, field -from typing import List, Dict - -from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors - -from deepspeed.runtime.base_optimizer import ZeROOptimizer -from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler -from deepspeed.runtime.torch_autocast import get_all_autocast_dtypes, is_autocast_initialized, sort_dtypes -from deepspeed.runtime.utils import (empty_cache, see_memory_usage, inf, is_model_parallel_parameter, - align_dense_tensors, all_gather_dp_groups, mask_nan_or_inf_with_val_inplace) -from deepspeed.runtime.zero.config import ZeroStageEnum -from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum -from deepspeed.ops.adam import DeepSpeedCPUAdam -from deepspeed.utils import logger -from deepspeed.utils.torch import register_grad_hook -from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank -from deepspeed.moe.utils import is_moe_param -from deepspeed.git_version_info import version - -from deepspeed.runtime.constants import PIPE_REPLICATED -from deepspeed.accelerator import get_accelerator - -from deepspeed.checkpoint.constants import (DS_VERSION, GROUP_PADDINGS, PARTITION_COUNT, LOSS_SCALER, - SINGLE_PARTITION_OF_FP32_GROUPS, BASE_OPTIMIZER_STATE, - BASE_OPTIMIZER_STATE_STEP, CLIP_GRAD, ZERO_STAGE, PARAM_SLICE_MAPPINGS) -from deepspeed.utils import link_hp_params, lazy_init_hp_params_optimizer_state -from deepspeed.checkpoint import enable_universal_checkpoint - -from deepspeed.utils import groups -# Toggle this to true to enable correctness test -# with gradient partitioning and without -pg_correctness_test = False - -OPTIMIZER_ALLGATHER_TIMER = 'optimizer_allgather' -OPTIMIZER_GRADIENTS_TIMER = 'optimizer_gradients' -OPTIMIZER_STEP_TIMER = 'optimizer_step' -OPTIMIZER_TIMERS = [OPTIMIZER_ALLGATHER_TIMER, OPTIMIZER_GRADIENTS_TIMER, OPTIMIZER_STEP_TIMER] -INITIAL_MICRO_STEP_ID = -1 - - -def input(msg): - return - - -def split_half_float_double(tensors): - device_type = get_accelerator().device_name() - dtypes = [ - "torch.{}.HalfTensor".format(device_type), "torch.{}.FloatTensor".format(device_type), - "torch.{}.DoubleTensor".format(device_type), "torch.{}.BFloat16Tensor".format(device_type) - ] - buckets = [] - for i, dtype in enumerate(dtypes): - bucket = [t for t in tensors if t.type() == dtype] - if bucket: - buckets.append(bucket) - return buckets - - -def isclose(a, b, rtol=1e-09, atol=0.0): - return abs(a - b) <= max(rtol * max(abs(a), abs(b)), atol) - - -def lcm(x, y): - from fractions import gcd # or can import gcd from `math` in Python 3 - return x * y // gcd(x, y) - - -def get_alignment_padding(tensor_list, alignment): - num_elements = sum([tensor.numel() for tensor in tensor_list]) - remainder = num_elements % alignment - return (alignment - remainder) if remainder else remainder - - -def print_rank_msg(msg): - print(f"rank {dist.get_rank()} - {msg}") - - -def _get_padded_tensor(src_tensor, size): - if src_tensor.numel() >= size: - return src_tensor - padded_tensor = torch.zeros(size, dtype=src_tensor.dtype, device=src_tensor.device) - slice_tensor = torch.narrow(padded_tensor, 0, 0, src_tensor.numel()) - slice_tensor.data.copy_(src_tensor.data) - return padded_tensor - - -def _pad_tensor_by_size(src_tensor, pad_size, dtype, device): - padded_tensor = torch.zeros(src_tensor.numel() + pad_size, dtype=dtype, device=device) - padded_tensor.data[:src_tensor.numel()].copy_(src_tensor.data) - return padded_tensor - - -@dataclass -class IPGBucket: - buffer: List[torch.Tensor] = field(default_factory=list) - params: List[torch.Tensor] = field(default_factory=list) - grads: List[torch.Tensor] = field(default_factory=list) - elements: int = 0 - index: int = 0 - has_moe_params: bool = False - - def clear(self): - self.params.clear() - self.grads.clear() - self.elements = 0 - self.index = 0 - self.has_moe_params = False - - -class DeepSpeedZeroOptimizer(ZeROOptimizer): - """ - DeepSpeedZeroOptimizer designed to reduce the memory footprint - required for training large deep learning models. - - For more details please see ZeRO: Memory Optimization Towards Training A Trillion Parameter Models - https://arxiv.org/abs/1910.02054 - - For usage examples, refer to TODO: DeepSpeed Tutorial - - """ - - def __init__(self, - init_optimizer, - param_names, - timers, - static_loss_scale=1.0, - dynamic_loss_scale=False, - dynamic_loss_args=None, - verbose=True, - contiguous_gradients=True, - reduce_bucket_size=500000000, - use_multi_rank_bucket_allreduce=True, - allgather_bucket_size=5000000000, - dp_process_group=None, - expert_parallel_group=None, - expert_data_parallel_group=None, - reduce_scatter=True, - overlap_comm=False, - offload_optimizer_config=None, - mpu=None, - clip_grad=0.0, - gradient_accumulation_dtype=torch.float32, - communication_data_type=torch.float16, - postscale_gradients=True, - gradient_predivide_factor=1.0, - gradient_accumulation_steps=1, - ignore_unused_parameters=True, - partition_grads=True, - round_robin_gradients=False, - has_moe_layers=False, - fp16_master_weights_and_gradients=False, - elastic_checkpoint=False, - check_grad_overflow=True): - - if offload_optimizer_config is not None and offload_optimizer_config.device != OffloadDeviceEnum.none: - self.cpu_offload = True - self.cpu_offload_pin_memory = offload_optimizer_config.pin_memory - else: - self.cpu_offload = False - self.cpu_offload_pin_memory = False - - if dist.get_rank() == 0: - logger.info(f"Reduce bucket size {reduce_bucket_size}") - logger.info(f"Allgather bucket size {allgather_bucket_size}") - logger.info(f"CPU Offload: {self.cpu_offload}") - logger.info(f'Round robin gradient partitioning: {round_robin_gradients}') - # The fused optimizer does all the work. We need this layer for two reason: - # 1. maintain same user API from apex.fp16_utils - # 2. keep common stuff here in case we need to add ne552w fused optimizer later - - self.elastic_checkpoint = elastic_checkpoint - self.check_grad_overflow = check_grad_overflow - self.param_names = param_names - self.mpu = mpu - # differences from apex.fp16_utils: - # - assume all model params in fp16 - # - assume all params requires grad - # - flat by groups, not keeping state. TODO: remove state explicitly? - # - master grad and unflat master weight never exist. TODO: a way to save out unflat master? - if not get_accelerator().is_available(): - raise SystemError("Accelerator is not detected, cannot perform low precision training (e.g., fp16, bf16).") - self.optimizer = init_optimizer - - # Use torch (un)flatten ops - self.flatten = _flatten_dense_tensors - self.unflatten = _unflatten_dense_tensors - - # ZeRO stage 1 (False) or 2 (True) - self.partition_gradients = partition_grads - self.zero_stage_string = "ZeRO-2" if partition_grads else "ZeRO-1" - - self.timers = timers - - self.reduce_scatter = reduce_scatter - - self.overlap_comm = overlap_comm - - self.deepspeed_adam_offload = self.cpu_offload - - self.device = get_accelerator().current_device_name() if not self.cpu_offload else 'cpu' - - self.dp_process_group = dp_process_group - self.sequence_parallel_size = groups._get_sequence_parallel_world_size() - #expert parallel group - self.ep_process_group = expert_parallel_group - - #data parallel group for experts - self.expert_dp_process_group = expert_data_parallel_group - - #data parallel size for non-experts - dp_size = dist.get_world_size(group=self.dp_process_group) - - #For MoE models this maybe different for different param group - #It will be modified during MoE setup later in the init - self.real_dp_process_group = [dp_process_group for i in range(len(self.optimizer.param_groups))] - self.partition_count = [dp_size for i in range(len(self.optimizer.param_groups))] - - self.is_gradient_accumulation_boundary = True - - # CPU-Offload requires contiguous gradients - self.contiguous_gradients = contiguous_gradients or self.cpu_offload - - self.has_moe_layers = has_moe_layers - if self.has_moe_layers: - self._configure_moe_settings() - self._global_grad_norm = 0. - - if mpu is None or hasattr(mpu, 'initialize_sequence_parallel'): - self.model_parallel_group = None - self.model_parallel_world_size = 1 - self.model_parallel_rank = 0 - else: - self.model_parallel_group = mpu.get_model_parallel_group() - self.model_parallel_world_size = mpu.get_model_parallel_world_size() - self.model_parallel_rank = bwc_tensor_model_parallel_rank(mpu) - - self.overflow = False - self.clip_grad = clip_grad - self.communication_data_type = communication_data_type - self.gradient_predivide_factor = gradient_predivide_factor - self.postscale_gradients = postscale_gradients - self.gradient_accumulation_steps = gradient_accumulation_steps - self.micro_step_id = INITIAL_MICRO_STEP_ID - self.ignore_unused_parameters = ignore_unused_parameters - self.round_robin_gradients = round_robin_gradients - - self.extra_large_param_to_reduce: Dict[int, torch.Tensor] = {} - self.fp16_master_weights_and_gradients = fp16_master_weights_and_gradients - - if self.fp16_master_weights_and_gradients: - assert self.cpu_offload and type(self.optimizer) in [DeepSpeedCPUAdam], \ - f"fp16_master_and_gradients requires optimizer to support keeping fp16 master and gradients while keeping the optimizer states in fp32."\ - f"Currently only supported using ZeRO-Offload with DeepSpeedCPUAdam. But current setting is ZeRO-Offload:{self.cpu_offload} and optimizer type {type(self.optimizer)}." \ - f"Either disable fp16_master_weights_and_gradients or enable {self.zero_stage_string} Offload with DeepSpeedCPUAdam." - - if self.reduce_scatter and self.partition_gradients: - valid_reduce_scatter_dtypes = (torch.float16, torch.bfloat16, torch.float32) - assert self.communication_data_type in valid_reduce_scatter_dtypes, f"{self.zero_stage_string} supports {valid_reduce_scatter_dtypes} communication_data_type with reduce scatter enabled. Got: '{self.communication_data_type}'" - assert self.gradient_predivide_factor == 1.0, f"gradient_predivide_factor != 1.0 is not yet supported with {self.zero_stage_string} with reduce scatter enabled" - assert self.postscale_gradients, f"pre-scale gradients is not yet supported with {self.zero_stage_string} with reduce scatter enabled" - - # param flattened by groups - self.bit16_groups = [] - self.bit16_groups_flat = [] - - # param partitioned by data parallel degree - # this will contain a list of equal sized tensors - # each of which will be updated by a different process - self.parallel_partitioned_bit16_groups = [] - - # a single 32-bit partition of the parallel partitioned parameters - # that this process will update - self.single_partition_of_fp32_groups = [] - - # a 16-bit CPU param buffer for cpu offload - if self.cpu_offload: - self.param_buffer_of_bit16_for_cpu_offload_groups = [] - - # param partition info - - # These are the parameters in each group that will not be updated by this process directly - self.params_not_in_partition = [] - - # These are the parameters that will be updated by this process directly - self.params_in_partition = [] - - # Offset from the first parameter in the self.params_in_partition - # the parameter boundaries may not align with partition boundaries - # so we need to keep track of the offset - self.first_offset = [] - - # number of elements per partition in each group - self.partition_size = [] - - # align nccl all-gather send buffers to 4-byte boundary - self.nccl_start_alignment_factor = 2 # 4-byte alignment/sizeof(fp16) = 2 - - assert ( - allgather_bucket_size % self.nccl_start_alignment_factor == 0 - ), f"allgather_bucket_size must be a multiple of nccl_start_alignment_factor, {self.nccl_start_alignment_factor} " - - self.all_reduce_print = False - self.dtype = self.optimizer.param_groups[0]['params'][0].dtype - self.gradient_accumulation_dtype = gradient_accumulation_dtype - - if self.dtype != self.gradient_accumulation_dtype: - self.use_separate_grad_accum = True - else: - self.use_separate_grad_accum = False - if self.use_separate_grad_accum and not self.partition_gradients: - self.use_grad_accum_attribute = True - else: - self.use_grad_accum_attribute = False - - self.round_robin_bit16_groups = [] - self.round_robin_bit16_indices = [] - self.round_robin_bit16_meta = [] - - # Use different parallel to do all_to_all_reduce related things - # padding on each partition for alignment purposes - self.groups_padding = [] - # loop to deal with groups - for i, param_group in enumerate(self.optimizer.param_groups): - partition_id = dist.get_rank(group=self.real_dp_process_group[i]) - - # push this group to list before modify - # TODO: Explore simplification that avoids the extra book-keeping by pushing the reordered group - trainable_parameters = [] - for param in param_group['params']: - if param.requires_grad: - param.grad_accum = None - param.param_idx_in_group = len(trainable_parameters) - trainable_parameters.append(param) - self.bit16_groups.append(trainable_parameters) - - # not sure why apex was cloning the weights before flattening - # removing cloning here - - see_memory_usage(f"Before moving param group {i} to CPU") - # move all the parameters to cpu to free up GPU space for creating flat buffer - - # Create temp CPU param copies, free accelerator tensors - orig_group_numel = 0 - for param in self.bit16_groups[i]: - orig_group_numel += param.numel() - param.cpu_data = param.data.cpu() - param.data = torch.empty(1).to(param.device) - - empty_cache() - see_memory_usage(f"After moving param group {i} to CPU", force=False) - - # Reorder group parameters for load balancing of gradient partitioning during backward among ranks. - # This ensures that gradients are reduced in a fashion such that ownership round robins among the ranks. - # For example, rather than 3 gradients (g_n+2, g_n+1, g_n) that are reduced consecutively belonging - # to the same rank, instead they will belong to 3 ranks (r_m+2, r_m+1, r_m). - if self.round_robin_gradients: - round_robin_tensors, round_robin_indices = self._round_robin_reorder( - self.bit16_groups[i], dist.get_world_size(group=self.real_dp_process_group[i])) - else: - round_robin_tensors = self.bit16_groups[i] - round_robin_indices = list(range(len(self.bit16_groups[i]))) - - self.round_robin_bit16_groups.append(round_robin_tensors) - self.round_robin_bit16_indices.append(round_robin_indices) - - # Create meta tensors list, ordered according to round_robin_tensors - meta_tensors = [] - for param in round_robin_tensors: - meta_tensors.append(torch.zeros_like(param.cpu_data, device="meta")) - self.round_robin_bit16_meta.append(meta_tensors) - - # create flat buffer in CPU - flattened_buffer = self.flatten_dense_tensors_aligned( - self.round_robin_bit16_groups[i], - self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_dp_process_group[i]), - use_cpu_data=True) - - # free temp CPU params - for param in self.bit16_groups[i]: - del param.cpu_data - - # Move CPU flat tensor to the accelerator memory. - self.bit16_groups_flat.append(flattened_buffer.to(get_accelerator().current_device_name())) - del flattened_buffer - - see_memory_usage(f"After flattening and moving param group {i} to GPU", force=False) - - if dist.get_rank(group=self.real_dp_process_group[i]) == 0: - see_memory_usage(f"After Flattening and after emptying param group {i} cache", force=False) - - # set model bit16 weight to slices of flattened buffer - self._update_model_bit16_weights(i) - - # divide the flat weights into near equal partition equal to the data parallel degree - # each process will compute on a different part of the partition - data_parallel_partitions = self.get_data_parallel_partitions(self.bit16_groups_flat[i], i) - self.parallel_partitioned_bit16_groups.append(data_parallel_partitions) - - # Record padding required for alignment - left_boundary = sum([t.numel() for t in data_parallel_partitions[:partition_id]]) - curr_partition_size = data_parallel_partitions[partition_id].numel() - - if orig_group_numel <= left_boundary: - padding = curr_partition_size - elif orig_group_numel < left_boundary + curr_partition_size: - padding = left_boundary + curr_partition_size - orig_group_numel - else: - padding = 0 - self.groups_padding.append(padding) - - # verify that data partition start locations are 4-byte aligned - for partitioned_data in data_parallel_partitions: - assert (partitioned_data.data_ptr() % (2 * self.nccl_start_alignment_factor) == 0) - - # A partition of the fp32 master weights that will be updated by this process. - # Note that the params in single_partition_of_fp32_groups is cloned and detached - # from the origin params of the model. - if not fp16_master_weights_and_gradients: - weights_partition = self.parallel_partitioned_bit16_groups[i][partition_id].to( - self.device).clone().float().detach() - else: - weights_partition = self.parallel_partitioned_bit16_groups[i][partition_id].to( - self.device).clone().half().detach() - - if self.cpu_offload: - weights_partition = get_accelerator().pin_memory(weights_partition) - temp_dtype = self.parallel_partitioned_bit16_groups[i][partition_id].dtype - temp_buffer_bit16 = torch.full(weights_partition.shape, - fill_value=0.0, - dtype=temp_dtype, - device=weights_partition.device) - if self.cpu_offload_pin_memory: - temp_pinned = get_accelerator().pin_memory(temp_buffer_bit16) - self.param_buffer_of_bit16_for_cpu_offload_groups.append(temp_pinned) - else: - self.param_buffer_of_bit16_for_cpu_offload_groups.append(temp_buffer_bit16) - - self.single_partition_of_fp32_groups.append(weights_partition) - - # Set local optimizer to have flat params of its own partition. - # After this, the local optimizer will only contain its own partition of params. - # In that case, the local optimizer only saves the states(momentum, variance, etc.) related to its partition's params(zero stage1). - self.single_partition_of_fp32_groups[ - i].requires_grad = True # keep this in case internal optimizer uses it - param_group['params'] = [self.single_partition_of_fp32_groups[i]] - - partition_size = len(self.bit16_groups_flat[i]) / dist.get_world_size(group=self.real_dp_process_group[i]) - params_in_partition, params_not_in_partition, first_offset = self.get_partition_info( - self.round_robin_bit16_groups[i], partition_size, partition_id) - - self.partition_size.append(partition_size) - self.params_in_partition.append(params_in_partition) - self.params_not_in_partition.append(params_not_in_partition) - self.first_offset.append(first_offset) - - self.reduce_bucket_size = int(reduce_bucket_size) - self.use_multi_rank_bucket_allreduce = use_multi_rank_bucket_allreduce - self.allgather_bucket_size = int(allgather_bucket_size) - - self.reduction_stream = None if get_accelerator().is_synchronized_device() else get_accelerator().Stream() - #self.copy_grad_stream = get_accelerator().Stream() - self.callback_queued = False - - self.param_dict = {} - - # map between param_id and bool to specify if a param is in this partition - self.is_param_in_current_partition = {} - - if is_autocast_initialized(): - comm_dtypes = get_all_autocast_dtypes([p for params in self.bit16_groups for p in params]) - self.torch_autocast_gradscaler = torch.amp.GradScaler(device=get_accelerator().device_name()) - else: - comm_dtypes = {self.communication_data_type} - self.torch_autocast_gradscaler = None - - self.ipg_buckets: Dict[torch.dtype, IPGBucket] = {dtype: IPGBucket() for dtype in comm_dtypes} - - self.params_already_reduced = [] - self._release_ipg_buffers() - self.previous_reduced_grads: Dict[int, List[torch.Tensor]] = defaultdict(list) - - # simplified param id - self.param_id = {} - - #interesting code: unique ids being assigned to individual parameters - largest_param_numel = 0 - count = 0 - for i, params_group in enumerate(self.bit16_groups): - for param in params_group: - unique_id = id(param) - self.param_id[unique_id] = count - self.param_dict[count] = param - self.params_already_reduced.append(False) - if param.numel() > largest_param_numel: - largest_param_numel = param.numel() - count = count + 1 - - for param_group in self.params_in_partition: - for param in param_group: - self.is_param_in_current_partition[self.get_param_id(param)] = True - - for param_group in self.params_not_in_partition: - for param in param_group: - self.is_param_in_current_partition[self.get_param_id(param)] = False - - if self.cpu_offload: - self.accumulated_grads_in_cpu = {} - self.norm_for_param_grads = {} - self.local_overflow = False - self.grad_position = {} - self.temp_grad_buffer_for_cpu_offload = torch.zeros(largest_param_numel, - device=self.device, - dtype=self.dtype) - if self.cpu_offload_pin_memory: - self.temp_grad_buffer_for_cpu_offload = get_accelerator().pin_memory( - self.temp_grad_buffer_for_cpu_offload) - self.temp_grad_buffer_for_gpu_offload = torch.zeros(largest_param_numel, - device=get_accelerator().current_device_name(), - dtype=self.dtype) - for i, params_group in enumerate(self.bit16_groups): - self.get_grad_position(i, self.params_in_partition[i], self.first_offset[i], self.partition_size[i]) - - # mapping from parameter to partition that it belongs to - self.param_to_partition_ids = {} - - # stores if a partition has been reduced in this step - self.is_partition_reduced = {} - - # number of grads in partition that still need to be computed - self.remaining_grads_in_partition = {} - - # total number of grads in partition - self.total_grads_in_partition = {} - - # stores if a grad in a partition has been computed or not - self.is_grad_computed = {} - - # stores the offset at which a parameter gradient needs to be inserted in a partition - self.grad_partition_insertion_offset = {} - - # the offset in the gradient at which it must be inserted at the beginning of the partition - self.grad_start_offset = {} - - # will store the averaged gradients required by this partition - self.averaged_gradients = {} - - # For cpu_offload, will store the averaged gradients required by this partition - self.offload_gradient_dict = {} - - # store index of first parameter in each partition - self.first_param_index_in_partition = {} - - # initializes all data structures for implementing gradient partitioning - self.initialize_gradient_partitioning_data_structures() - - # resets the data structure value for the next backward propagation - self.reset_partition_gradient_structures() - - # creates backward hooks for the following special handling of gradients - # 1. upcasting for fp32 gradient accumulation - # 2. gradient partitioning - # 3. overlapping backward and reduction - self._grad_acc_hooks = [] - - if (self.partition_gradients or self.overlap_comm or self.use_grad_accum_attribute - or self.contiguous_gradients): - self.create_gradient_handling_hooks() - - self.ready_for_gradients = False - self.custom_loss_scaler = False - self.external_loss_scale = None - - # we may have a way of fusing dynamic scale. Do not support for now - self.loss_scaler = CreateLossScaler(dtype=self.dtype, - static_loss_scale=static_loss_scale, - dynamic_scaling=dynamic_loss_scale, - dynamic_loss_args=dynamic_loss_args) - self.dynamic_loss_scale = self.loss_scaler.dynamic - - if self.dtype != torch.float16: - # Only fp16 should use dynamic loss scaling - assert self.loss_scaler.cur_scale == 1.0 - assert not self.dynamic_loss_scale - - see_memory_usage("Before initializing optimizer states", force=True) - self.initialize_optimizer_states() - see_memory_usage("After initializing optimizer states", force=True) - - if dist.get_rank() == 0: - logger.info(f"optimizer state initialized") - - if dist.get_rank(group=self.dp_process_group) == 0: - see_memory_usage(f"After initializing ZeRO optimizer", force=True) - - self._link_all_hp_params() - self._hp_optimizer_states_linked = False - - self._enable_universal_checkpoint() - self._param_slice_mappings = self._create_param_mapping() - if self.cpu_offload: - self._create_optimizer_mapping() - - def destroy(self): - for i, _ in enumerate(self.optimizer.param_groups): - for p in self.bit16_groups[i]: - if getattr(p, '_hp_mapping', None): - p._hp_mapping = None - for hook in self._grad_acc_hooks: - hook.remove() - self.print_rank_0("Removed grad acc hooks") - - def _enable_universal_checkpoint(self): - for lp_param_group in self.bit16_groups: - enable_universal_checkpoint(param_list=lp_param_group) - - def _create_param_mapping(self): - param_mapping = [] - for i, _ in enumerate(self.optimizer.param_groups): - param_mapping_per_group = OrderedDict() - for lp in self.bit16_groups[i]: - if lp._hp_mapping is not None: - lp_name = self.param_names[lp] - param_mapping_per_group[lp_name] = lp._hp_mapping.get_hp_fragment_address() - param_mapping.append(param_mapping_per_group) - - return param_mapping - - def _create_optimizer_mapping(self): - for i, _ in enumerate(self.optimizer.param_groups): - for lp in self.bit16_groups[i]: - if lp._hp_mapping is not None: - lp._zero_optimizer = self - - def _link_all_hp_params(self): - if self.cpu_offload: - self._get_offload_gradient_dict() - - for i, _ in enumerate(self.optimizer.param_groups): - # Link bit16 and fp32 params in partition - partition_id = dist.get_rank(group=self.real_dp_process_group[i]) - partition_size = self.bit16_groups_flat[i].numel() // dist.get_world_size( - group=self.real_dp_process_group[i]) - flat_hp_partition = self.single_partition_of_fp32_groups[i] - link_hp_params(lp_param_list=self.bit16_groups[i], - flat_hp_partition=flat_hp_partition, - gradient_dict=self.averaged_gradients, - offload_gradient_dict=self.offload_gradient_dict, - use_offload=self.cpu_offload, - param_group_index=i, - partition_start=partition_id * partition_size, - partition_size=partition_size, - dp_group=self.real_dp_process_group[i]) - - def _lazy_init_hp_params_optimizer_state(self): - if not self._hp_optimizer_states_linked: - for i, _ in enumerate(self.optimizer.param_groups): - lazy_init_hp_params_optimizer_state(self.bit16_groups[i], self.single_partition_of_fp32_groups[i], - self.optimizer.state) - self._hp_optimizer_states_linked = True - - def is_moe_group(self, group): - return 'moe' in group and group['moe'] - - def _configure_moe_settings(self): - # if we're using ZeRO stage 2, ensure contiguous gradients are used - if self.partition_gradients: - assert self.contiguous_gradients, "Contiguous Gradients in ZeRO Stage 2 must be set to True for MoE. Other code paths are not tested with MoE" - # NOTE: To run ZeRO stage 1 with MoE, we need to set self.contiguous_gradients to True or ignore the assertion - if not self.partition_gradients and not self.contiguous_gradients: - logger.warning( - "ZeRO Stage 1 has not been thoroughly tested with MoE. This configuration is still experimental.") - assert self.reduce_scatter, "Reduce Scatter in ZeRO Stage 2 must be set to True for MoE. Other code paths are not tested with MoE" - - assert any( - [self.is_moe_group(group) for group in self.optimizer.param_groups] - ), "The model has moe layers, but None of the param groups are marked as MoE. Create a param group with 'moe' key set to True before creating optimizer" - self.is_moe_param_group = [] - for i, group in enumerate(self.optimizer.param_groups): - if self.is_moe_group(group): - assert all([is_moe_param(param) - for param in group['params']]), "All params in MoE group must be MoE params" - self.real_dp_process_group[i] = self.expert_dp_process_group[group['name']] - self.partition_count[i] = dist.get_world_size(group=self.expert_dp_process_group[group['name']]) - self.is_moe_param_group.append(True) - else: - self.is_moe_param_group.append(False) - - assert self.expert_dp_process_group is not None, "Expert data parallel group should be configured with MoE" - assert self.ep_process_group is not None, "Expert parallel group should be configured with MoE" - - def _update_model_bit16_weights(self, group_index): - updated_params = self.unflatten(self.bit16_groups_flat[group_index], self.round_robin_bit16_meta[group_index]) - for p, q in zip(self.round_robin_bit16_groups[group_index], updated_params): - p.data = q.data - - # set model fp16 weight to slices of reordered flattened buffer - for param_index, param in enumerate(self.bit16_groups[group_index]): - new_index = self.round_robin_bit16_indices[group_index][param_index] - param.data = self.round_robin_bit16_groups[group_index][new_index].data - - def _round_robin_reorder(self, tensor_list, num_partitions): - - # disable round robin if need to debug something - # return tensor_list, list(range(len(tensor_list))) - - partition_tensors = {} - - for i, tensor in enumerate(tensor_list): - j = i % num_partitions - if not j in partition_tensors: - partition_tensors[j] = [] - partition_tensors[j].append((i, tensor)) - - reordered_tensors = [] - reordered_indices = {} - - for partition_index in partition_tensors.keys(): - for i, (original_index, tensor) in enumerate(partition_tensors[partition_index]): - reordered_indices[original_index] = len(reordered_tensors) - reordered_tensors.append(tensor) - - return reordered_tensors, reordered_indices - - def _release_ipg_buffers(self): - if self.contiguous_gradients: - for bucket in self.ipg_buckets.values(): - bucket.buffer.clear() - - self.grads_in_partition = None - self.grads_in_partition_offset = 0 - self.ready_for_gradients = False - - def initialize_optimizer_states(self): - - for i, group in enumerate(self.bit16_groups): - single_grad_partition = torch.zeros(int(self.partition_size[i]), - dtype=self.single_partition_of_fp32_groups[i].dtype, - device=self.device) - self.single_partition_of_fp32_groups[i].grad = get_accelerator().pin_memory( - single_grad_partition) if self.cpu_offload_pin_memory else single_grad_partition - - # Initialize the optimizer states with the flattened fp32 partition. - # State initialization for the Adagrad optimizer occurs at construction as opposed to other optimizers - # which do lazy initialization of the state at the first call to step. - if isinstance(self.optimizer, torch.optim.Adagrad): - self.optimizer = torch.optim.Adagrad(self.single_partition_of_fp32_groups, **self.optimizer.defaults) - - if not self.cpu_offload: - for group in self.single_partition_of_fp32_groups: - group.grad = None #class init - - return - - ######################################################################### - #################### ZeRO Stage 1 - reduce gradients #################### - ######################################################################### - def reduce_gradients(self, pipeline_parallel=False): - world_size = dist.get_world_size(self.dp_process_group) - my_rank = dist.get_rank(self.dp_process_group) - - # with PP we must create ipg buffer, since backward is handled outside zero - if pipeline_parallel and self.contiguous_gradients: - for dtype, bucket in self.ipg_buckets.items(): - bucket.buffer.append( - torch.empty(int(self.reduce_bucket_size), - dtype=dtype, - device=get_accelerator().current_device_name())) - bucket.index = 0 - - if not self.overlap_comm: - for i, group in enumerate(self.bit16_groups): - for param in group: - grad_reduc = self.get_gradient_for_reduction(param) - if grad_reduc is not None: - self.reduce_ready_partitions_and_remove_grads(param, i) - # reduce any pending grads in either hook/non-hook case - self.overlapping_partition_gradients_reduce_epilogue() - - ######################################################################### - #########################ZeRO Partition Gradients######################## - ######################################################################### - - def get_first_param_index(self, group_id, param_group, partition_id): - for index, param in enumerate(param_group): - param_id = self.get_param_id(param) - if group_id in self.param_to_partition_ids and param_id in self.param_to_partition_ids[group_id]: - if partition_id in self.param_to_partition_ids[group_id][param_id]: - return index - return None - - def initialize_gradient_partitioning_data_structures(self): - - for i, param_group in enumerate(self.round_robin_bit16_groups): - total_partitions = dist.get_world_size(group=self.real_dp_process_group[i]) - - self.param_to_partition_ids[i] = {} - self.is_partition_reduced[i] = {} - self.total_grads_in_partition[i] = {} - self.remaining_grads_in_partition[i] = {} - self.is_grad_computed[i] = {} - self.grad_partition_insertion_offset[i] = {} - self.grad_start_offset[i] = {} - self.first_param_index_in_partition[i] = {} - - for partition_id in range(total_partitions): - self.is_grad_computed[i][partition_id] = {} - self.grad_partition_insertion_offset[i][partition_id] = {} - self.grad_start_offset[i][partition_id] = {} - self.total_grads_in_partition[i][partition_id] = 0 - self.initialize_gradient_partition(i, param_group, partition_id) - self.is_partition_reduced[i][partition_id] = False - self.first_param_index_in_partition[i][partition_id] = self.get_first_param_index( - i, param_group, partition_id) - - def independent_gradient_partition_epilogue(self): - self.report_ipg_memory_usage(f"In ipg_epilogue before reduce_ipg_grads", 0) - self.reduce_ipg_grads() - self.report_ipg_memory_usage(f"In ipg_epilogue after reduce_ipg_grads", 0) - - # if dist.get_rank() == 0: - # logger.info("Params already reduced %s", self.params_already_reduced) - for i in range(len(self.params_already_reduced)): - self.params_already_reduced[i] = False - - if self.overlap_comm: - if not get_accelerator().resolves_data_dependency(): - get_accelerator().synchronize() - # It is safe to clear previously reduced grads of other partitions - self._clear_previous_reduced_grads() - - if self.cpu_offload is False: - for i, _ in enumerate(self.bit16_groups): - - if not i in self.averaged_gradients or self.averaged_gradients[i] is None: - self.averaged_gradients[i] = self.get_flat_partition( - self.params_in_partition[i], - self.first_offset[i], - self.partition_size[i], - dtype=self.gradient_accumulation_dtype, - device=get_accelerator().current_device_name(), - return_tensor_list=True) - else: - avg_new = self.get_flat_partition(self.params_in_partition[i], - self.first_offset[i], - self.partition_size[i], - dtype=self.gradient_accumulation_dtype, - device=get_accelerator().current_device_name(), - return_tensor_list=True) - - for accumulated_grad, new_avg_grad in zip(self.averaged_gradients[i], avg_new): - accumulated_grad.add_(new_avg_grad) - - self._release_ipg_buffers() - - # No need to keep the gradients anymore. - # All gradients required by the step - # are in self.averaged_gradients - self.zero_grad(set_to_none=True) - see_memory_usage(f"End ipg_epilogue") - - # resets all partition to no reduced - # sets remaining grads to the total number of grads in each partition - # set is grad computed to false for all grads in partition - def reset_partition_gradient_structures(self): - for i, _ in enumerate(self.bit16_groups): - total_partitions = dist.get_world_size(group=self.real_dp_process_group[i]) - for partition_id in range(total_partitions): - self.is_partition_reduced[i][partition_id] = False - self.remaining_grads_in_partition[i][partition_id] = self.total_grads_in_partition[i][partition_id] - - for param_id in self.is_grad_computed[i][partition_id]: - self.is_grad_computed[i][partition_id][param_id] = False - - def initialize_gradient_partition(self, i, param_group, partition_id): - - def set_key_value_list(dictionary, key, value): - if key in dictionary: - dictionary[key].append(value) - else: - dictionary[key] = [value] - - def increment_value(dictionary, key): - if key in dictionary: - dictionary[key] += 1 - else: - dictionary[key] = 1 - - partition_size = self.partition_size[i] - - start_index = partition_size * partition_id - end_index = partition_size * (partition_id + 1) - - current_index = 0 - first_offset = 0 - - for param in param_group: - - param_size = param.numel() - param_id = self.get_param_id(param) - - if start_index <= current_index < end_index: - set_key_value_list(self.param_to_partition_ids[i], param_id, partition_id) - increment_value(self.total_grads_in_partition[i], partition_id) - - self.is_grad_computed[i][partition_id][param_id] = False - - self.grad_partition_insertion_offset[i][partition_id][param_id] = current_index - start_index - self.grad_start_offset[i][partition_id][param_id] = 0 - - elif current_index < start_index < (current_index + param_size): - assert (first_offset == 0 - ), "This can happen either zero or only once as this must be the first tensor in the partition" - first_offset = start_index - current_index - - set_key_value_list(self.param_to_partition_ids[i], param_id, partition_id) - increment_value(self.total_grads_in_partition[i], partition_id) - - self.is_grad_computed[i][partition_id][param_id] = False - - self.grad_partition_insertion_offset[i][partition_id][param_id] = 0 - self.grad_start_offset[i][partition_id][param_id] = first_offset - - current_index = current_index + param_size - - def overlapping_partition_gradients_reduce_epilogue(self): - self.independent_gradient_partition_epilogue() - - def _fill_param_grad_accum_attribute(self, param): - if param.grad is not None: - if param.grad_accum is None: - param.grad_accum = param.grad.to(self.gradient_accumulation_dtype) - else: - param.grad_accum.add_(param.grad.to(self.gradient_accumulation_dtype).view(param.grad_accum.shape)) - param.grad = None - - def fill_grad_accum_attribute(self): - for group in self.bit16_groups: - for param in group: - self._fill_param_grad_accum_attribute(param) - - def get_gradient_for_reduction(self, param): - if self.use_grad_accum_attribute: - return param.grad_accum.to(self.dtype) if param.grad_accum is not None else None - else: - return param.grad - - def get_param_gradient_attribute(self, param): - return param.grad_accum if self.use_grad_accum_attribute else param.grad - - # Clear the tensor the reduction gradient attribute is pointing to - def clear_grad_attribute(self, param): - if self.use_grad_accum_attribute: - param.grad_accum = None - else: - param.grad = None - - def create_gradient_handling_hooks(self): - for i, param_group in enumerate(self.bit16_groups): - for param in param_group: - if param.requires_grad: - - def wrapper(param, i): - - def grad_handling_hook(*notneeded): - self.process_gradients(param, i) - - self._grad_acc_hooks.append(register_grad_hook(param, grad_handling_hook)) - - wrapper(param, i) - - def get_param_id(self, param): - unique_id = id(param) - return self.param_id[unique_id] - - # create a flat tensor aligned at the alignment boundary - def flatten_dense_tensors_aligned(self, tensor_list, alignment, use_cpu_data=False): - tensor_list = [param.cpu_data for param in tensor_list] if use_cpu_data else tensor_list - return self.flatten(align_dense_tensors(tensor_list, alignment)) - - ############### Independent Partition Gradient ######################## - def reduce_independent_p_g_buckets_and_remove_grads(self, param, i): - - grad_reduc = self.get_gradient_for_reduction(param) - bucket = self.ipg_buckets[self.get_param_comm_dtype(param)] - if bucket.elements + param.numel() > self.reduce_bucket_size: - self.report_ipg_memory_usage("In ipg_remove_grads before reduce_ipg_grads", param.numel(), param.dtype) - self.reduce_ipg_grads() - if self.contiguous_gradients and self.overlap_comm: - # Swap index between 0 and 1 - bucket.index = 1 - bucket.index - self.report_ipg_memory_usage("In ipg_remove_grads after reduce_ipg_grads", param.numel(), param.dtype) - - param_id = self.get_param_id(param) - assert self.params_already_reduced[param_id] == False, \ - f"The parameter {param_id} has already been reduced. \ - Gradient computed twice for this partition. \ - Multiple gradient reduction is currently not supported" - - if self.contiguous_gradients: - if param.numel() > self.reduce_bucket_size: - self.extra_large_param_to_reduce[param.dtype] = param - else: - # keeping the gradients contiguous to prevent memory fragmentation, and avoid flattening - new_grad_tensor = bucket.buffer[bucket.index].narrow(0, bucket.elements, param.numel()) - new_grad_tensor.copy_(grad_reduc.view(-1)) - grad_reduc.data = new_grad_tensor.data.view_as(grad_reduc) - - bucket.elements += param.numel() - - assert grad_reduc is not None, f"rank {dist.get_rank()} - Invalid to reduce Param {param_id} with None gradient" - - # deal with a use-case of transient grads that will be generated in a loop for the same computation involving some model params - e.g. when performing a tiled memory calculation that shards the normal single sub-module call into a loop over a shards. - if getattr(param, "ds_grad_is_ready", True): - bucket.grads.append(grad_reduc) - bucket.params.append((i, param.param_idx_in_group, param_id)) - - #make sure the average tensor function knows how to average the gradients - if is_moe_param(param): - bucket.has_moe_params = True - - self.report_ipg_memory_usage("End ipg_remove_grads", 0) - - def print_rank_0(self, message): - if dist.get_rank() == 0: - logger.info(message) - - def gradient_reduction_w_predivide(self, tensor, communication_data_type: torch.dtype): - if tensor.size().numel() == 0: - return tensor - - dp_world_size = dist.get_world_size(group=self.dp_process_group) - - tensor_to_allreduce = tensor - - if communication_data_type != tensor.dtype: - tensor_to_allreduce = tensor.to(communication_data_type) - - if self.postscale_gradients: - if self.gradient_predivide_factor != 1.0: - tensor_to_allreduce.mul_(1. / self.gradient_predivide_factor) - - dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group) - - if self.gradient_predivide_factor != dp_world_size: - tensor_to_allreduce.mul_(self.gradient_predivide_factor / - (dp_world_size / float(self.sequence_parallel_size))) - else: - tensor_to_allreduce.div_(dp_world_size / float(self.sequence_parallel_size)) - dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group) - - if communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: - tensor.copy_(tensor_to_allreduce) - - return tensor - - def allreduce_and_copy_with_multiple_ranks(self, - small_bucket, - communication_data_type: torch.dtype, - log=None, - divide=True, - process_group=None, - bucket_ranks=None): - process_group = self.dp_process_group if process_group is None else process_group - allreduced = self.allreduce_bucket(small_bucket, - communication_data_type, - log=log, - divide=divide, - process_group=process_group) - for buf, synced, bucket_rank in zip(small_bucket, self.unflatten(allreduced, small_bucket), bucket_ranks): - if dist.get_rank(group=process_group) == bucket_rank: - buf.copy_(synced) - - def allreduce_and_scatter(self, - bucket, - communication_data_type: torch.dtype, - numel_per_bucket=500000000, - log=None, - divide=True, - process_group=None): - small_bucket = [] - small_bucket_ranks = [] - numel = 0 - allreduce_sizes = [] - - for i, bucket_elem in enumerate(bucket): - rank, tensor = bucket_elem - small_bucket.append(tensor) - small_bucket_ranks.append(rank) - numel = numel + tensor.numel() - if numel > numel_per_bucket: - self.allreduce_and_copy_with_multiple_ranks(small_bucket, - communication_data_type, - log=None, - divide=divide, - process_group=process_group, - bucket_ranks=small_bucket_ranks) - small_bucket = [] - small_bucket_ranks = [] - numel = 0 - - if len(small_bucket) > 0: - self.allreduce_and_copy_with_multiple_ranks(small_bucket, - communication_data_type, - log=None, - divide=divide, - process_group=process_group, - bucket_ranks=small_bucket_ranks) - - def average_tensor(self, tensor: torch.Tensor, communication_data_type: torch.dtype): - if self.overlap_comm: - stream = self.reduction_stream - if not get_accelerator().resolves_data_dependency(): - stream.wait_stream(get_accelerator().current_stream()) - get_accelerator().current_stream().wait_stream(stream) - else: - stream = get_accelerator().current_stream() - - with get_accelerator().stream(stream): - if not self.reduce_scatter: - self.gradient_reduction_w_predivide(tensor, communication_data_type) - return - - # Accumulate destination ranks and bucket offsets for each gradient slice. - # Note: potential future optimization, record access pattern of parameters - # in backward pass and partition gradients w.r.t. access pattern so that our - # bucket is guaranteed to be contiguous w.r.t. ranks - rank_and_offsets = [] - real_dp_process_group = [] - curr_size = 0 - prev_id, prev_process_group = -1, None - - process_group = self.dp_process_group - # count = 0 - bucket = self.ipg_buckets[communication_data_type] - for i, param_idx_in_group, param_id in bucket.params: - param = self.bit16_groups[i][param_idx_in_group] - - process_group = self.dp_process_group - - if bucket.has_moe_params: - process_group = self.expert_dp_process_group[param.group_name] if is_moe_param( - param) else self.dp_process_group - - partition_ids = self.param_to_partition_ids[i][param_id] - assert all([p_id < dist.get_world_size(group=process_group) for p_id in partition_ids - ]), f"world size {dist.get_world_size(group=process_group)} and p_ids: {partition_ids}" - partition_size = self.partition_size[i] - # Get all partition ids + their offsets - partition_ids_w_offsets = [] - for partition_id in partition_ids: - offset = self.grad_start_offset[i][partition_id][param_id] - partition_ids_w_offsets.append((partition_id, offset)) - partition_ids_w_offsets.sort(key=lambda t: t[1]) - - # Calculate rank and offsets for grad slices - for idx in range(len(partition_ids_w_offsets)): - partition_id, offset = partition_ids_w_offsets[idx] - - # if dist.get_rank() == 0 and count < 100: - # print(f"Rank {dist.get_rank()} rank offset id {idx} calculated dp size {dist.get_world_size(group=process_group)} real dp size {dist.get_world_size(self.real_dp_process_group[i])} and dst: {partition_id}") - # count += 1 - - # Calculate numel for grad slice depending on partition location - if idx == len(partition_ids_w_offsets) - 1: - # Last partition_id uses its own offset - numel = param.numel() - offset - else: - # Set numel to next partition's offset - numel = partition_ids_w_offsets[idx + 1][1] - offset - - # Merge bucket ranges if they belong to the same rank - if partition_id == prev_id and process_group == prev_process_group: - prev_pid, prev_size, prev_numel = rank_and_offsets[-1] - rank_and_offsets[-1] = (prev_pid, prev_size, prev_numel + numel) - else: - rank_and_offsets.append((partition_id, curr_size, numel)) - real_dp_process_group.append(process_group) - curr_size += numel - prev_id, prev_process_group = partition_id, process_group - - tensor.div_(dist.get_world_size(group=self.dp_process_group) / float(self.sequence_parallel_size)) - - buckets = {} - for i, (dst, bucket_offset, numel) in enumerate(rank_and_offsets): - grad_slice = tensor.narrow(0, int(bucket_offset), int(numel)) - bucket_key = real_dp_process_group[i] if self.use_multi_rank_bucket_allreduce else ( - dst, real_dp_process_group[i]) - if bucket_key not in buckets: - buckets[bucket_key] = [] - if self.use_multi_rank_bucket_allreduce: - buckets[bucket_key].append((dst, grad_slice)) - else: - buckets[bucket_key].append(grad_slice) - - for bucket_key in buckets: - if self.use_multi_rank_bucket_allreduce: - self.allreduce_and_scatter(buckets[bucket_key], - communication_data_type, - numel_per_bucket=self.reduce_bucket_size, - divide=False, - process_group=bucket_key) - else: - dst, process_group = bucket_key - self.allreduce_no_retain(buckets[bucket_key], - communication_data_type, - numel_per_bucket=self.reduce_bucket_size, - rank=dst, - divide=False, - process_group=process_group) - - ############################################################################## - ############################# CPU Offload Methods############################# - ############################################################################## - def get_grad_position(self, group_id, tensor_list, first_offset, partition_size): - current_offset = 0 - - for i, tensor in enumerate(tensor_list): - param_id = self.get_param_id(tensor) - param_start_offset = 0 - - num_elements = tensor.numel() - - # we need to offset to get to the right element - if i == 0 and first_offset > 0: - tensor_offset = first_offset - num_elements = num_elements - tensor_offset - param_start_offset = first_offset - - # we dont need all elements of the tensor - if num_elements > (partition_size - current_offset): - num_elements = partition_size - current_offset - - self.grad_position[param_id] = [ - int(group_id), int(param_start_offset), - int(current_offset), int(num_elements) - ] - current_offset += num_elements - - def update_offload_overflow_tracker(self, grad): - if grad is not None and self._has_inf_or_nan(grad.data): - self.local_overflow = True - - def update_offload_overflow_tracker_for_param_grad(self, param): - grad_accum = self.get_param_gradient_attribute(param) - self.update_offload_overflow_tracker(grad_accum) - - def _get_offload_gradient_dict(self): - for param_group_index, _ in enumerate(self.optimizer.param_groups): - self.offload_gradient_dict[param_group_index] = [] - for lp_param in self.params_in_partition[param_group_index]: - param_id = self.get_param_id(lp_param) - [_, _, dest_offset, num_elements] = self.grad_position[param_id] - dest_tensor = self.single_partition_of_fp32_groups[param_group_index].grad.view(-1).narrow( - 0, dest_offset, num_elements) - self.offload_gradient_dict[param_group_index].append(dest_tensor) - - def async_accumulate_grad_in_cpu_via_gpu(self, param): - param_id = self.get_param_id(param) - - [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] - - # copy to a preexisiting buffer to avoid memory allocation penalty - dest_buffer = self.temp_grad_buffer_for_gpu_offload.view(-1).narrow(0, 0, param.numel()) - - #buffer for storing gradients for this parameter in CPU - def buffer_to_accumulate_to_in_cpu(): - if not self.fp16_master_weights_and_gradients: - buffer = torch.zeros(param.numel(), dtype=param.dtype, device=self.device) - return get_accelerator().pin_memory(buffer) if self.cpu_offload_pin_memory else buffer - else: - return self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements) - - #accumulate gradients into param.grad_accum or parts of it that belongs to this partition - def accumulate_gradients(): - grad_accum = self.get_param_gradient_attribute(param) - if not self.fp16_master_weights_and_gradients: - dest_buffer.copy_(self.accumulated_grads_in_cpu[param_id].view(-1), non_blocking=True) - grad_accum.data.view(-1).add_(dest_buffer) - else: - dest_buffer.narrow(0, source_offset, - num_elements).copy_(self.accumulated_grads_in_cpu[param_id].view(-1), - non_blocking=True) - grad_accum.data.view(-1).narrow(0, source_offset, - num_elements).add_(dest_buffer.narrow(0, source_offset, num_elements)) - - #move accumulated gradients back to CPU - def copy_gradients_to_cpu(): - grad_accum = self.get_param_gradient_attribute(param) - if not self.fp16_master_weights_and_gradients: - self.accumulated_grads_in_cpu[param_id].data.copy_(grad_accum.data.view(-1), non_blocking=True) - else: - self.accumulated_grads_in_cpu[param_id].data.copy_(grad_accum.data.view(-1).narrow( - 0, source_offset, num_elements), - non_blocking=True) - - if param_id not in self.accumulated_grads_in_cpu: - self.accumulated_grads_in_cpu[param_id] = buffer_to_accumulate_to_in_cpu() - - if self.micro_step_id > 0: - accumulate_gradients() - else: - copy_gradients_to_cpu() - - def set_norm_for_param_grad(self, param): - param_id = self.get_param_id(param) - grad_accum = self.get_param_gradient_attribute(param) - accumulated_grad = self.accumulated_grads_in_cpu[ - param_id] if self.gradient_accumulation_steps > 1 else grad_accum - - [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] - - start = source_offset - accumulated_grad = accumulated_grad.view(-1).narrow(0, start, num_elements) - - self.norm_for_param_grads[param_id] = accumulated_grad.data.double().norm(2) - - def set_norm_for_param_grad_in_gpu(self, param): - param_id = self.get_param_id(param) - grad_accum = self.get_param_gradient_attribute(param) - if grad_accum is None: - accumulated_grad = param.grad - else: - accumulated_grad = grad_accum - - [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] - - start = source_offset - accumulated_grad = accumulated_grad.view(-1).narrow(0, start, num_elements) - - self.norm_for_param_grads[param_id] = accumulated_grad.data.double().norm(2) - - def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param): - param_id = self.get_param_id(param) - - [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id] - - dest_tensor = self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements) - - grad_accum = self.get_param_gradient_attribute(param) - if grad_accum is None: - src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements) - else: - src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements) - if not self.fp16_master_weights_and_gradients: - src_tensor = src_tensor.float() - - dest_tensor.copy_(src_tensor, non_blocking=True) - self.clear_grad_attribute(param) #offload only - - def complete_grad_norm_calculation_for_cpu_offload(self, params): - total_norm = 0.0 - norm_type = 2.0 - for p in params: - # Pipeline parallelism may replicate parameters. Avoid multi-counting. - if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: - continue - - if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0): - param_id = self.get_param_id(p) - # as some model have trainable parameters but skipped in training, - # their backward hooks in self.create_gradient_handling_hooks() will not run, - # so they have no norm_for_param_grads - if param_id in self.norm_for_param_grads: - param_norm = self.norm_for_param_grads[param_id] - total_norm += param_norm.item()**2 - else: - # As unused parameters in modules may not be expected sometimes, - # add an explicit error msg when it occurred and an option to - # avoid the error - assert self.ignore_unused_parameters, """ - This assert indicates that your module has parameters that - were not used in producing loss. - You can avoid this assert by - (1) enable ignore_unused_parameters option in zero_optimization config; - (2) making sure all trainable parameters and `forward` function - outputs participate in calculating loss. - """ - - # Sum across all model parallel GPUs. - total_dev_norm = get_accelerator().FloatTensor([float(total_norm)]) - dist.all_reduce(total_dev_norm, op=dist.ReduceOp.SUM, group=self.dp_process_group) - - self._model_parallel_all_reduce(tensor=total_dev_norm, op=dist.ReduceOp.SUM) - - total_norm = total_dev_norm[0].item()**(1. / norm_type) - - if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm: - total_norm = -1.0 - - return torch.tensor(total_norm, device=self.device, dtype=torch.float) - - ############################################################################################ - def copy_grads_in_partition(self, param): - if self.cpu_offload: - - if self.gradient_accumulation_steps > 1: - self.async_accumulate_grad_in_cpu_via_gpu(param) - - if self.is_gradient_accumulation_boundary: - self.set_norm_for_param_grad_in_gpu(param) - - self.update_offload_overflow_tracker_for_param_grad(param) - - self.async_inplace_copy_grad_to_fp32_buffer_from_gpu(param) - - return - #print(f"ID {self.get_param_id(param)} grad norm {param.grad.norm()}") - if self.grads_in_partition is None: - self.grads_in_partition_offset = 0 - total_size = 0 - for group in self.params_in_partition: - for param_in_partition in group: - total_size += param_in_partition.numel() - - see_memory_usage(f"before copying {total_size} gradients into partition") - self.grads_in_partition = torch.empty(int(total_size), - dtype=self.dtype, - device=get_accelerator().current_device_name()) - see_memory_usage(f"after copying {total_size} gradients into partition") - - grad_reduc = self.get_gradient_for_reduction(param) - # The allreduce buffer will be rewritten. Copy the gradients in partition to a new buffer - new_grad_tensor = self.grads_in_partition.view(-1).narrow(0, self.grads_in_partition_offset, param.numel()) - new_grad_tensor.copy_(grad_reduc.view(-1)) - grad_reduc.data = new_grad_tensor.data.view_as(grad_reduc) - #print(f"Grad norm after copy to contiguous_buffer {param.grad.data.norm()}") - self.grads_in_partition_offset += param.numel() - - def reduce_ipg_grads(self): - for comm_dtype in sort_dtypes(self.ipg_buckets.keys()): - bucket = self.ipg_buckets[comm_dtype] - - if self.contiguous_gradients: - if comm_dtype in self.extra_large_param_to_reduce: - assert len(bucket.params) == 1, "more than 1 param in ipg bucket, this shouldn't happen" - _, _, param_id = bucket.params[0] - assert self.get_param_id(self.extra_large_param_to_reduce[comm_dtype] - ) == param_id, "param in ipg bucket does not match extra-large param" - extra_large_grad_reduc = self.get_gradient_for_reduction( - self.extra_large_param_to_reduce[comm_dtype]) - self.average_tensor(extra_large_grad_reduc.view(-1), comm_dtype) - del self.extra_large_param_to_reduce[comm_dtype] - else: - self.average_tensor(bucket.buffer[bucket.index].narrow(0, 0, bucket.elements), comm_dtype) - else: - self.buffered_reduce_fallback(None, bucket.grads, comm_dtype, elements_per_buffer=bucket.elements) - - if self.overlap_comm: - stream = self.reduction_stream - elif self.cpu_offload: - # TODO: copy_grad_stream is disabled because of race with reduce. This hurts perf and should be fixed. - # get_accelerator().synchronize() - # stream = self.copy_grad_stream - stream = get_accelerator().current_stream() - else: - stream = get_accelerator().current_stream() - - with get_accelerator().stream(stream): - for comm_dtype in sort_dtypes(self.ipg_buckets.keys()): - bucket = self.ipg_buckets[comm_dtype] - - for group_idx, param_idx_in_group, param_id in bucket.params: - param = self.bit16_groups[group_idx][param_idx_in_group] - - assert self.params_already_reduced[param_id] == False, \ - f"The parameter {param_id} has already been reduced. \ - Gradient computed twice for this partition. \ - Multiple gradient reduction is currently not supported" - - self.params_already_reduced[param_id] = True - if self.partition_gradients: - if not self.is_param_in_current_partition[param_id]: - if self.overlap_comm and self.contiguous_gradients is False: - # Clear grads of other partitions during the next reduction - # to avoid clearing them before the reduction is complete. - self.previous_reduced_grads[comm_dtype].append(param) - else: - self.clear_grad_attribute(param) - elif self.contiguous_gradients: - self.copy_grads_in_partition(param) - else: # zero stage 1 - partition only optimizer state - if self.contiguous_gradients and self.is_param_in_current_partition[param_id]: - self.copy_grads_in_partition(param) - bucket.clear() - ##################################################################### - - def process_gradients(self, param, i): - self.backward_prologue() - if self.use_grad_accum_attribute: - self._fill_param_grad_accum_attribute(param) - if self.partition_gradients or self.overlap_comm: - self.reduce_ready_partitions_and_remove_grads(param, i) - - def reduce_ready_partitions_and_remove_grads(self, param, i): - if self.partition_gradients or self.is_gradient_accumulation_boundary: - self.reduce_independent_p_g_buckets_and_remove_grads(param, i) - - def zero_reduced_gradients(self, partition_id, i): - - def are_all_related_partitions_reduced(params_id): - for partition_id in self.param_to_partition_ids[i][params_id]: - if not self.is_partition_reduced[i][partition_id]: - return False - return True - - for params_id in self.is_grad_computed[i][partition_id]: - if are_all_related_partitions_reduced(params_id): - self.param_dict[params_id].grad = None # dead code - - def flatten_and_print(self, message, tensors, start=0, n=5): - flatten_tensor = self.flatten(tensors) - - def print_func(): - logger.info(flatten_tensor.contiguous().view(-1).narrow(0, start, n)) - - self.sequential_execution(print_func, message) - - def get_grads_to_reduce(self, i, partition_id): - - def get_reducible_portion(key): - grad = self.param_dict[key].grad - total_elements = grad.numel() - start = self.grad_start_offset[i][partition_id][key] - num_elements = min(total_elements - start, - self.partition_size[i] - self.grad_partition_insertion_offset[i][partition_id][key]) - if not pg_correctness_test: - if num_elements == total_elements: - return grad - else: - return grad.contiguous().view(-1).narrow(0, int(start), int(num_elements)) - else: - if num_elements == total_elements: - return grad.clone() - else: - return grad.clone().contiguous().view(-1).narrow(0, int(start), int(num_elements)) - - grads_to_reduce = [] - for key in self.is_grad_computed[i][partition_id]: - grad = get_reducible_portion(key) - grads_to_reduce.append(grad) - return grads_to_reduce - - def sequential_execution(self, function, message, group=None): - if group is None: - group = self.dp_process_group - if dist.get_rank(group=group) == 0: - logger.info(message) - for id in range(dist.get_world_size(group=group)): - if id == dist.get_rank(group=group): - function() - dist.barrier(group=group) - - def set_none_gradients_to_zero(self, i, partition_id): - for param_id in self.is_grad_computed[i][partition_id]: - param = self.param_dict[param_id] - if param.grad is None: - param.grad = torch.zeros_like(param) - - ######################Reduction Related Methods############################## - def allreduce_bucket(self, - bucket, - communication_data_type: torch.dtype, - rank=None, - log=None, - divide=True, - process_group=None): - - tensor = self.flatten(bucket) - - process_group = self.dp_process_group if process_group is None else process_group - - tensor_to_allreduce = tensor - - if pg_correctness_test or self.sequence_parallel_size > 1: - communication_data_type = torch.float32 - - if communication_data_type != tensor.dtype: - tensor_to_allreduce = tensor.to(communication_data_type) - - if divide: - tensor_to_allreduce.div_(dist.get_world_size(group=process_group) / float(self.sequence_parallel_size)) - - if rank is None: - # "All Reducing" - dist.all_reduce(tensor_to_allreduce, group=process_group) - else: - global_rank = dist.get_global_rank(process_group, rank) - dist.reduce(tensor_to_allreduce, global_rank, group=process_group) - - if communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: - if rank is None or rank == dist.get_rank(group=process_group): - tensor.copy_(tensor_to_allreduce) - - return tensor - - def _clear_previous_reduced_grads(self): - for dtype in self.previous_reduced_grads: - for param in self.previous_reduced_grads[dtype]: - self.clear_grad_attribute(param) - self.previous_reduced_grads[dtype].clear() - - # if rank is specified do a reduction instead of an allreduce - def allreduce_and_copy(self, - small_bucket, - communication_data_type: torch.dtype, - rank=None, - log=None, - divide=True, - process_group=None): - process_group = self.dp_process_group if process_group is None else process_group - if self.overlap_comm: - if not get_accelerator().resolves_data_dependency(): - get_accelerator().synchronize() - # It is safe to clear the previously reduced grads of other partitions - self._clear_previous_reduced_grads() - stream = self.reduction_stream - else: - stream = get_accelerator().current_stream() - - with get_accelerator().stream(stream): - allreduced = self.allreduce_bucket( - small_bucket, - communication_data_type, - rank=rank, - log=log, - divide=divide, - process_group=process_group, - ) - if rank is None or rank == dist.get_rank(group=self.dp_process_group): - for buf, synced in zip(small_bucket, self.unflatten(allreduced, small_bucket)): - buf.copy_(synced) - - def allreduce_no_retain( - self, - bucket, - communication_data_type: torch.dtype, - numel_per_bucket=500000000, - rank=None, - log=None, - divide=True, - process_group=None, - ): - small_bucket = [] - numel = 0 - for tensor in bucket: - small_bucket.append(tensor) - numel = numel + tensor.numel() - if numel > numel_per_bucket: - self.allreduce_and_copy(small_bucket, - communication_data_type, - rank=rank, - log=None, - divide=divide, - process_group=process_group) - small_bucket = [] - numel = 0 - - if len(small_bucket) > 0: - self.allreduce_and_copy(small_bucket, - communication_data_type, - rank=rank, - log=log, - divide=divide, - process_group=process_group) - - # allows using reduction of gradients instead of using all_reduce - - def buffered_reduce_fallback(self, - rank, - grads, - communication_data_type: torch.dtype, - elements_per_buffer=500000000, - log=None): - split_buckets = split_half_float_double(grads) - - for i, bucket in enumerate(split_buckets): - self.allreduce_no_retain(bucket, - communication_data_type, - numel_per_bucket=elements_per_buffer, - rank=rank, - log=log) - - ############################################################################# - ############################################################################# - ############################################################################# - - # views the tensor as multiple partitions and returns - # those partitions - def get_data_parallel_partitions(self, tensor, group_id): - partitions = [] - - dp = dist.get_world_size(group=self.real_dp_process_group[group_id]) - # dp_id = dist.get_rank(group=self.real_dp_process_group[group_id]) - - total_num_elements = tensor.numel() - - base_size = total_num_elements // dp - remaining = total_num_elements % dp - - start = 0 - for id in range(dp): - partition_size = base_size - if id < remaining: - partition_size = partition_size + 1 - partitions.append(tensor.narrow(0, start, partition_size)) - start = start + partition_size - return partitions - - def get_partition_info(self, tensor_list, partition_size, partition_id): - params_in_partition = [] - params_not_in_partition = [] - - start_index = partition_size * partition_id - end_index = partition_size * (partition_id + 1) - - current_index = 0 - first_offset = 0 - - for tensor in tensor_list: - - tensor_size = tensor.numel() - - if start_index <= current_index < end_index: - params_in_partition.append(tensor) - - elif current_index < start_index < (current_index + tensor_size): - params_in_partition.append(tensor) - - assert (first_offset == 0 - ), "This can happen either zero or only once as this must be the first tensor in the partition" - first_offset = start_index - current_index - - else: - params_not_in_partition.append(tensor) - - current_index = current_index + tensor_size - - return params_in_partition, params_not_in_partition, first_offset - - def zero_grad(self, set_to_none=True): - """ - Zero FP16 parameter grads. - """ - # FP32 grad should never exist. - # For speed, set model fp16 grad to None by default - # zero all pointers to grad tensors - for group in self.bit16_groups: - for p in group: - if set_to_none: - p.grad = None # epilogue and in step - p.grad_accum = None - else: - if p.grad is not None: - p.grad.detach_() - p.grad.zero_() - - def _model_parallel_all_reduce(self, tensor, op): - """ Perform all reduce within model parallel group, if any. - """ - if self.model_parallel_group is None or self.model_parallel_world_size == 1: - pass - else: - dist.all_reduce(tensor=tensor, op=op, group=self.model_parallel_group) - - def get_grad_norm_direct(self, gradients, params, norm_type=2): - """Clips gradient norm of an iterable of parameters. - - This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and - added functionality to handle model parallel parameters. Note that - the gradients are modified in place. - - Arguments: - parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a - single Tensor that will have gradients normalized - max_norm (float or int): max norm of the gradients - norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for - infinity norm. - - Returns: - Total norm of the parameters (viewed as a single vector). - """ - norm_type = float(norm_type) - all_norms = [] - if norm_type == inf: - for g in gradients: - all_norms.append(g.data.abs().max().float()) - total_norm = torch.stack(all_norms).max() - dist.all_reduce(total_norm, op=dist.ReduceOp.MAX, group=self.dp_process_group) - - # Take max across all GPUs. - self._model_parallel_all_reduce(tensor=total_norm, op=dist.ReduceOp.MAX) - else: - # if dist.get_rank() == 0: - # logger.info(f"Total Norm beginning {total_norm}") - for g, p in zip(gradients, params): - # Pipeline parallelism may replicate parameters. Avoid multi-counting. - if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: - continue - if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0): - all_norms.append( - torch.linalg.vector_norm(g.data.double().detach(), - ord=norm_type).to(get_accelerator().current_device_name())) - if len(all_norms) > 0: - total_norm = torch.stack(all_norms).square().sum().float() - else: - total_norm = torch.tensor(0.0, dtype=torch.float32).to(self.device) - # Sum across all model parallel Device. - dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=self.dp_process_group) - - self._model_parallel_all_reduce(tensor=total_norm, op=dist.ReduceOp.SUM) - - total_norm = total_norm.pow(1. / norm_type) - - mask_nan_or_inf_with_val_inplace(total_norm, device=self.device) - - return total_norm - - # creates a flat fused tensor from the tensor list starting at the first_offset - # in the first tensor of the list. If there are not enough elements in the tensor - # list then the flat tensor will be padded with zeros - def get_flat_partition(self, tensor_list, first_offset, partition_size, dtype, device, return_tensor_list=False): - flat_tensor_list = [] - current_size = 0 - - for i, tensor in enumerate(tensor_list): - grad_accum = self.get_param_gradient_attribute(tensor) - if grad_accum is None: - grad_accum = torch.zeros_like(tensor, dtype=dtype) - - tensor = grad_accum - num_elements = tensor.numel() - tensor_offset = 0 - - # we need to offset to get to the right element - if i == 0 and first_offset > 0: - tensor_offset = first_offset - num_elements = num_elements - tensor_offset - - # we dont need all elements of the tensor - if num_elements > (partition_size - current_size): - num_elements = partition_size - current_size - - # we need a narrow view of the tensor based on the tensor offset and number of elements that - # we need from this tensor - if tensor_offset > 0 or num_elements < tensor.numel(): - flat_tensor_list.append(tensor.contiguous().view(-1).narrow(0, int(tensor_offset), int(num_elements))) - else: - flat_tensor_list.append(tensor) - - current_size = current_size + num_elements - - # this means its the last partition and does not align with the dp boundary. We need to pad before flattening - if current_size < partition_size: - flat_tensor_list.append(torch.zeros(int(partition_size - current_size), dtype=dtype, device=device)) - - if return_tensor_list: - return flat_tensor_list - - return self.flatten(flat_tensor_list) - - def free_grad_in_param_list(self, param_list): - for p in param_list: - p.grad = None # in step - p.grad_accum = None - - def reset_cpu_buffers(self): - self.norm_for_param_grads = {} - self.local_overflow = False - - def set_lr(self, lr): - """Set the learning rate.""" - for param_group in self.optimizer.param_groups: - param_group["lr"] = lr - - def get_lr(self): - """Return the current learning rate.""" - return self.optimizer.param_groups[0]["lr"] - - def override_loss_scale(self, loss_scale): - if loss_scale != self.external_loss_scale: - logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}') - self.custom_loss_scaler = True - self.external_loss_scale = loss_scale - - def scaled_global_norm(self, norm_type=2): - assert norm_type == 2, "only L2 norm supported" - norm_groups = [] - for i, group in enumerate(self.bit16_groups): - if self.cpu_offload: - norm = self.complete_grad_norm_calculation_for_cpu_offload(self.params_in_partition[i]) - norm_groups.append(norm) - else: - norm_groups.append(self.get_grad_norm_direct(self.averaged_gradients[i], self.params_in_partition[i])) - - if self.has_moe_layers: - self._average_expert_grad_norms(norm_groups) - - # calculating L2 norm - return torch.linalg.vector_norm(torch.stack(norm_groups), ord=norm_type) - - def get_bit16_param_group(self, group_no): - bit16_partitions = self.parallel_partitioned_bit16_groups[group_no] - partition_id = dist.get_rank(group=self.real_dp_process_group[group_no]) - return [bit16_partitions[dist.get_rank(group=self.real_dp_process_group[group_no])]] - - def _optimizer_step(self, group_no): - original_param_groups = self.optimizer.param_groups - self.optimizer.param_groups = [original_param_groups[group_no]] - # Disabling this as the C++ side copy & synchronize is not working correctly - #from deepspeed.ops.adam import DeepSpeedCPUAdam - #if type(self.optimizer) == DeepSpeedCPUAdam and self.dtype == torch.half: - # self.optimizer.step(fp16_param_groups=[self.get_bit16_param_group(group_no)]) - #else: - # self.optimizer.step() - if self.torch_autocast_gradscaler: - self.torch_autocast_gradscaler.step(self.optimizer) - self.torch_autocast_gradscaler.update() - else: - self.optimizer.step() - self.optimizer.param_groups = original_param_groups - - # We need to link optimizer state after the first step() call - self._lazy_init_hp_params_optimizer_state() - - def step(self, closure=None): - """ - Not supporting closure. - """ - self.micro_step_id = INITIAL_MICRO_STEP_ID - - see_memory_usage(f"In step before checking overflow") - - # First compute norm for all group so we know if there is overflow - if self.check_grad_overflow: - self.check_overflow(partition_gradients=self.partition_gradients) - - prev_scale = self.loss_scale - self._update_scale(self.overflow) - if self.overflow: - see_memory_usage('After overflow before clearing gradients') - self.zero_grad(set_to_none=True) - if self.cpu_offload: - self.reset_cpu_buffers() - else: - for k in self.averaged_gradients.keys(): - self.averaged_gradients[k] = None - - see_memory_usage('After overflow after clearing gradients') - - for timer in OPTIMIZER_TIMERS: - self.timers(timer).start() - self.timers(timer).stop() - return - - # Step 1:- Calculate gradient norm using bit-16 grads - see_memory_usage('Before norm calculation') - scaled_global_grad_norm = self.scaled_global_norm() - self._global_grad_norm = scaled_global_grad_norm / prev_scale - see_memory_usage('After norm before optimizer') - - # Step 2:- run optimizer and upscaling simultaneously - for i, group in enumerate(self.bit16_groups): - self.timers(OPTIMIZER_GRADIENTS_TIMER).start() - partition_id = dist.get_rank(group=self.real_dp_process_group[i]) - if self.cpu_offload: - single_grad_partition = self.single_partition_of_fp32_groups[i].grad - self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm) - - self.timers(OPTIMIZER_GRADIENTS_TIMER).stop() - self.timers(OPTIMIZER_STEP_TIMER).start() - self._optimizer_step(i) - - # Disabled, this is not currently working - #from deepspeed.ops.adam import DeepSpeedCPUAdam - #if not (type(self.optimizer) == DeepSpeedCPUAdam and self.dtype == torch.half): - # bit16_partitions = self.parallel_partitioned_bit16_groups[i] - # fp32_partition = self.single_partition_of_fp32_groups[i] - # bit16_partitions[partition_id].data.copy_(fp32_partition.data) - bit16_partitions = self.parallel_partitioned_bit16_groups[i] - fp32_partition = self.single_partition_of_fp32_groups[i] - bit16_partition_buffer = self.param_buffer_of_bit16_for_cpu_offload_groups[i] - bit16_partition_buffer.data.copy_(fp32_partition.data) - bit16_partitions[partition_id].data.copy_(bit16_partition_buffer.data, non_blocking=True) - - self.timers(OPTIMIZER_STEP_TIMER).stop() - else: - # free gradients for all the parameters that are not updated by this process(ZeRO stage2) - self.free_grad_in_param_list(self.params_not_in_partition[i]) - - # create a flat gradients for parameters updated by this process - # If we are last partition, ensure we have same size grads and partition size, if not pad with zero tensors - if partition_id == dist.get_world_size(group=self.real_dp_process_group[i]) - 1: - single_grad_partition = self.flatten_dense_tensors_aligned( - self.averaged_gradients[i], - int(self.partition_size[i])).to(self.single_partition_of_fp32_groups[i].dtype) - else: - single_grad_partition = self.flatten(self.averaged_gradients[i]).to( - self.single_partition_of_fp32_groups[i].dtype) - assert single_grad_partition.numel() == self.partition_size[i], \ - "averaged gradients have different number of elements that partition size {} {} {} {}".format( - single_grad_partition.numel(), self.partition_size[i], i, partition_id) - - self.single_partition_of_fp32_groups[i].grad = single_grad_partition - # release all the gradient since we have already created a necessary copy in dp_grad_partition(ZeRO stage2) - self.free_grad_in_param_list(self.params_in_partition[i]) - - self.averaged_gradients[i] = None - - self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm) - - self.timers(OPTIMIZER_GRADIENTS_TIMER).stop() - - # Step 3:- run the optimizer if no offloading - self.timers(OPTIMIZER_STEP_TIMER).start() - self._optimizer_step(i) - # Step 4:- get rid of the fp32 gradients. Not needed anymore - self.single_partition_of_fp32_groups[i].grad = None - del single_grad_partition - bit16_partitions = self.parallel_partitioned_bit16_groups[i] - fp32_partition = self.single_partition_of_fp32_groups[i] - bit16_partitions[partition_id].data.copy_(fp32_partition.data) - self.timers(OPTIMIZER_STEP_TIMER).stop() - - see_memory_usage('After optimizer before all-gather') - if self.cpu_offload: - self.reset_cpu_buffers() - - self.timers(OPTIMIZER_ALLGATHER_TIMER).start() - # Gather the updated weights from everyone. - # Then all partitions of the model parameters are updated and ready for next round forward. - all_gather_dp_groups(groups_flat=self.bit16_groups_flat, - partitioned_param_groups=self.parallel_partitioned_bit16_groups, - dp_process_group=self.real_dp_process_group, - start_alignment_factor=self.nccl_start_alignment_factor, - allgather_bucket_size=self.allgather_bucket_size) - self.timers(OPTIMIZER_ALLGATHER_TIMER).stop() - - # TODO: we probably don't need this? just to be safe - for i in range(len(self.bit16_groups)): - self._update_model_bit16_weights(i) - - self.timers.log(OPTIMIZER_TIMERS) - see_memory_usage('After zero_optimizer step') - - return - - @torch.no_grad() - def update_lp_params(self): - for i, (bit16_partitions, fp32_partition) in enumerate( - zip(self.parallel_partitioned_bit16_groups, self.single_partition_of_fp32_groups)): - partition_id = dist.get_rank(group=self.real_dp_process_group[i]) - bit16_partitions[partition_id].data.copy_(fp32_partition.data) - - all_gather_dp_groups(groups_flat=self.bit16_groups_flat, - partitioned_param_groups=self.parallel_partitioned_bit16_groups, - dp_process_group=self.real_dp_process_group, - start_alignment_factor=self.nccl_start_alignment_factor, - allgather_bucket_size=self.allgather_bucket_size) - - def _average_expert_grad_norms(self, norm_groups): - for i, norm in enumerate(norm_groups): - if self.is_moe_param_group[i]: - scaled_norm_tensor = norm * 1.0 / dist.get_world_size(group=self.real_dp_process_group[i]) - if self.device == 'cpu': - scaled_norm_tensor = scaled_norm_tensor.to(get_accelerator().current_device_name()) - dist.all_reduce(scaled_norm_tensor, group=self.real_dp_process_group[i]) - norm_groups[i] = scaled_norm_tensor.to(self.device) - - def unscale_and_clip_grads(self, grad_groups_flat, total_norm): - # compute combined scale factor for this group - combined_scale = self.loss_scale - if self.clip_grad > 0.: - # norm is in fact norm*scale - clip = ((total_norm / self.loss_scale) + 1e-6) / self.clip_grad - clip = torch.clamp(clip, min=1.0) - combined_scale = clip * self.loss_scale - - for grad in grad_groups_flat: - if isinstance(grad, list): - sub_partitions = grad - for g in sub_partitions: - g.data.mul_(1. / combined_scale) - else: - grad.data.mul_(1. / combined_scale) - - def _check_overflow(self, partition_gradients=True): - self.overflow = self.has_overflow(partition_gradients) - - # `params` is a list / generator of torch.Variable - def has_overflow_serial(self, params): - invalid_grad_count = torch.zeros([1], dtype=torch.float, device=get_accelerator().current_device_name()) - for p in params: - if p.grad is not None: - invalid_grad_count += self._has_inf_or_nan(p.grad) - return invalid_grad_count.bool() - - def has_overflow_partitioned_grads_serial(self): - invalid_grad_count = torch.zeros([1], dtype=torch.float, device=get_accelerator().current_device_name()) - for i in range(len(self.bit16_groups)): - for j, grad in enumerate(self.averaged_gradients[i]): - if grad is not None: - invalid_grad_count += self._has_inf_or_nan(grad) - return invalid_grad_count.bool() - - def has_overflow(self, partition_gradients=True): - overflow = self.local_overflow if self.cpu_offload else self.has_overflow_partitioned_grads_serial() - overflow_gpu = get_accelerator().ByteTensor([overflow]) if self.cpu_offload else overflow.byte().to( - get_accelerator().current_device_name()) - - if partition_gradients: - '''This will capture overflow across all data parallel and expert parallel process - Since expert parallel process are a subset of data parallel process''' - dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.dp_process_group) - - # Since each model parallel GPU carries only part of the model, - # make sure overflow flag is synced across all the model parallel GPUs - self._model_parallel_all_reduce(tensor=overflow_gpu, op=dist.ReduceOp.MAX) - - overflow = overflow_gpu[0].item() - return bool(overflow) - - # `x` is a torch.Tensor - @staticmethod - def _has_inf_or_nan(x, j=None): - float_x = x.float() - nan = float_x.isnan() - inf = float_x.isinf() - inf_or_nan = nan.logical_or(inf) - return inf_or_nan.float().max() - - def backward_prologue(self): - if not self.ready_for_gradients: - self.micro_step_id += 1 - - if self.contiguous_gradients: - for _, bucket in self.ipg_buckets.items(): - bucket.buffer.clear() - - # Buffer's dtype is the same as the dtype of optimizer, not dtype for autocast - buf_0 = torch.empty(int(self.reduce_bucket_size), - dtype=self.dtype, - device=get_accelerator().current_device_name()) - bucket.buffer.append(buf_0) - bucket.index = 0 - - # Use double buffers to avoid data access conflict when overlap_comm is enabled. - if self.overlap_comm: - for _, bucket in self.ipg_buckets.items(): - buf_1 = torch.empty(int(self.reduce_bucket_size), - dtype=self.dtype, - device=get_accelerator().current_device_name()) - bucket.buffer.append(buf_1) - - self.ready_for_gradients = True - - def backward_epilogue(self): - # Only for Stage 1, Mode 2 - if self.use_grad_accum_attribute: - self.fill_grad_accum_attribute() - - def backward(self, loss, retain_graph=False): - """ - :attr:`backward` performs the following steps: - - 1. fp32_loss = loss.float() - 2. scaled_loss = fp32_loss*loss_scale - 3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves - """ - if self.custom_loss_scaler: - scaled_loss = self.external_loss_scale * loss - scaled_loss.backward() - elif self.torch_autocast_gradscaler: - self.torch_autocast_gradscaler.scale(loss).backward(retain_graph=retain_graph) - else: - self.loss_scaler.backward(loss.float(), retain_graph=retain_graph) - self.backward_epilogue() - - def check_overflow(self, partition_gradients=True): - self._check_overflow(partition_gradients) - - def _update_scale(self, has_overflow=False): - self.loss_scaler.update_scale(has_overflow) - - # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state" - def _get_state(self): - return self.optimizer.state - - def _set_state(self, value): - self.optimizer.state = value - - state = property(_get_state, _set_state) - - # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups" - # (for example, to adjust the learning rate) - def _get_param_groups(self): - return self.optimizer.param_groups - - def _set_param_groups(self, value): - self.optimizer.param_groups = value - - param_groups = property(_get_param_groups, _set_param_groups) - - # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale" - def _get_loss_scale(self): - if self.custom_loss_scaler: - return self.external_loss_scale - else: - return self.loss_scaler.cur_scale - - def _set_loss_scale(self, value): - self.loss_scaler.cur_scale = value - - loss_scale = property(_get_loss_scale, _set_loss_scale) - cur_scale = property(_get_loss_scale, _set_loss_scale) - - # Return group tensor after removing paddings that are added for alignment to DP world size. - # This method works on the assumption that each group contains a single flattened tensor. - def _get_groups_without_padding(self, groups_with_padding): - groups_without_padding = [] - for i, group in enumerate(groups_with_padding): - lean_length = group.numel() - self.groups_padding[i] - groups_without_padding.append(group[:lean_length]) - - return groups_without_padding - - # Return optimizer state after removing paddings that are added for alignment. - def _get_state_without_padding(self, state_with_padding, padding): - lean_state = {} - for key, value in state_with_padding.items(): - if torch.is_tensor(value): - lean_length = value.numel() - padding - lean_state[key] = value[:lean_length] - else: - lean_state[key] = value - - return lean_state - - # Return base optimizer states. - # This method assumes that each param group contains a single flattened tensor. - def _get_base_optimizer_state(self): - optimizer_groups_state = [] - for i, group in enumerate(self.optimizer.param_groups): - p = group['params'][0] - lean_optimizer_state = self._get_state_without_padding(self.optimizer.state[p], self.groups_padding[i]) - optimizer_groups_state.append(lean_optimizer_state) - - return optimizer_groups_state - - def state_dict(self): - """ - Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. - This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict - of the contained Pytorch optimizer. - Example:: - checkpoint = {} - checkpoint['model'] = model.state_dict() - checkpoint['optimizer'] = optimizer.state_dict() - torch.save(checkpoint, "saved.pth") - """ - state_dict = {} - state_dict[LOSS_SCALER] = self.loss_scaler - state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale - state_dict['overflow'] = self.overflow - state_dict[CLIP_GRAD] = self.clip_grad - - if self.elastic_checkpoint: - state_dict[BASE_OPTIMIZER_STATE] = self._get_base_optimizer_state() - - if "step" in self.optimizer.param_groups[0]: - # Assuming "step" is the only item that changes through training iterations - assert all(group["step"] == self.optimizer.param_groups[0]["step"] - for group in self.optimizer.param_groups), "All param groups must have the same step value" - state_dict[BASE_OPTIMIZER_STATE_STEP] = self.optimizer.param_groups[0]["step"] - else: - state_dict[BASE_OPTIMIZER_STATE] = self.optimizer.state_dict() - - # Remove paddings for DP alignment to enable loading for other alignment values - fp32_groups_without_padding = self._get_groups_without_padding(self.single_partition_of_fp32_groups) - state_dict[SINGLE_PARTITION_OF_FP32_GROUPS] = fp32_groups_without_padding - - state_dict[ - ZERO_STAGE] = ZeroStageEnum.gradients if self.partition_gradients else ZeroStageEnum.optimizer_states - state_dict[GROUP_PADDINGS] = self.groups_padding - state_dict[PARTITION_COUNT] = self.partition_count - - state_dict[DS_VERSION] = version - state_dict[PARAM_SLICE_MAPPINGS] = self._param_slice_mappings - - return state_dict - - # Restore base optimizer fp32 weights from elastic checkpoint by: - # 1) Merging fp32 weights from checkpoints of all partitions - # 2) Extracting fp32 weights for current partition from merged weights - # 3) Using extracted weights to update base optimizer weights directly. - def _restore_from_elastic_fp32_weights(self, all_state_dict): - merged_single_partition_of_fp32_groups = [] - - for i in range(len(self.single_partition_of_fp32_groups)): - partition_id = dist.get_rank(group=self.real_dp_process_group[i]) - merged_partitions = [sd[SINGLE_PARTITION_OF_FP32_GROUPS][i] for sd in all_state_dict] - if self.is_moe_group(self.optimizer.param_groups[i]): - ranks = self.get_ep_ranks(group_name=self.optimizer.param_groups[i]['name']) - merged_partitions = [merged_partitions[i] for i in ranks] - flat_merged_partitions = self.flatten_dense_tensors_aligned( - merged_partitions, - self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_dp_process_group[i])) - dp_partitions = self.get_data_parallel_partitions(flat_merged_partitions, i) - merged_single_partition_of_fp32_groups.append(dp_partitions[partition_id]) - - for current, saved in zip(self.single_partition_of_fp32_groups, merged_single_partition_of_fp32_groups): - current.data.copy_(saved.data) - - # Restore base optimizer fp32 weights from ZeRO fp16 or bfloat16 weights - def _restore_from_bit16_weights(self): - for group_id, (bit16_partitions, fp32_partition) in enumerate( - zip(self.parallel_partitioned_bit16_groups, self.single_partition_of_fp32_groups)): - partition_id = dist.get_rank(group=self.real_dp_process_group[group_id]) - fp32_partition.data.copy_(bit16_partitions[partition_id].data) - - # Refresh the fp32 master params from the fp16 or bfloat16 copies. - def refresh_fp32_params(self): - self._restore_from_bit16_weights() - - # Extract optimizer state for current partition from merged states of all partitions - def _partition_base_optimizer_state(self, state_key, all_partition_states, group_id): - partition_id = dist.get_rank(group=self.real_dp_process_group[group_id]) - alignment = self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_dp_process_group[group_id]) - if torch.is_tensor(all_partition_states[0]): - flat_merged_partitions = self.flatten_dense_tensors_aligned(all_partition_states, alignment) - dp_partitions = self.get_data_parallel_partitions(flat_merged_partitions, group_id) - return dp_partitions[partition_id] - else: - # Assume non-tensor states are not partitioned and equal across ranks, so return first one - return all_partition_states[0] - - def _restore_step_from_elastic_checkpoint(self, all_state_dict): - assert BASE_OPTIMIZER_STATE_STEP in all_state_dict[0] - assert all(sd[BASE_OPTIMIZER_STATE_STEP] == all_state_dict[0][BASE_OPTIMIZER_STATE_STEP] - for sd in all_state_dict), "State dicts of all partitions must have the same step value" - return all_state_dict[0][BASE_OPTIMIZER_STATE_STEP] - - def _restore_base_optimizer_state(self, base_optimizer_group_states, base_optimizer_state_step, group_paddings): - if type(base_optimizer_group_states) == dict: - base_optimizer_group_states = base_optimizer_group_states['state'] - - saved_keys = base_optimizer_group_states[0].keys() - - for i, group in enumerate(self.optimizer.param_groups): - p = group['params'][0] - padding = 0 if group_paddings is None else group_paddings[i] - for key in saved_keys: - saved = base_optimizer_group_states[i][key] - - if torch.is_tensor(saved): - if key in self.optimizer.state[p]: - dst_tensor = self.optimizer.state[p][key] - src_tensor = _get_padded_tensor(saved, dst_tensor.numel()) - self.optimizer.state[p][key].data.copy_(src_tensor.data) - else: - self.optimizer.state[p][key] = _pad_tensor_by_size( - saved, padding, torch.float32, - torch.device('cpu') if self.cpu_offload else self.device) - else: - self.optimizer.state[p][key] = saved - - for param_group in self.optimizer.param_groups: - param_group['step'] = base_optimizer_state_step - - def get_ep_ranks(self, rank=0, group_name=None): - from deepspeed.utils import groups - expert_parallel_size_ = groups._get_expert_parallel_world_size(group_name) - world_size = groups._get_data_parallel_world_size() - rank = groups._get_expert_parallel_rank(group_name) - ranks = range(rank, world_size, expert_parallel_size_) - return list(ranks) - - # Restore base optimizer state from elastic checkpoint by - # 1) Merging optimizer state from checkpoints of all partitions - # 2) Extracting optimizer state for current partition from the merged state - # 3) Using the extracted value to directly update the base optimizer. - def _restore_elastic_base_optimizer_state(self, all_state_dict): - base_optimizer_group_states = [] - for i in range(len(self.optimizer.param_groups)): - partition_states = {} - all_partition_group_states = [sd[BASE_OPTIMIZER_STATE][i] for sd in all_state_dict] - - if self.is_moe_group(self.optimizer.param_groups[i]): - ranks = self.get_ep_ranks(group_name=self.optimizer.param_groups[i]['name']) - all_partition_group_states = [all_partition_group_states[i] for i in ranks] - - for key in all_partition_group_states[0].keys(): - all_partition_states = [all_states[key] for all_states in all_partition_group_states] - partition_states[key] = self._partition_base_optimizer_state(key, all_partition_states, i) - base_optimizer_group_states.append(partition_states) - - self._restore_base_optimizer_state(base_optimizer_group_states, - self._restore_step_from_elastic_checkpoint(all_state_dict), None) - - def load_state_dict(self, - state_dict_list, - load_optimizer_states=True, - load_from_fp32_weights=False, - checkpoint_folder=None, - load_serial=None, - param_shapes=None): - if checkpoint_folder: - self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights) - else: - self._load_legacy_checkpoint(state_dict_list, load_optimizer_states, load_from_fp32_weights) - - def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights): - self.load_hp_checkpoint_state_from_checkpoint_dir("bit16_groups", checkpoint_folder) - - def _load_global_state(self, sd): - self.loss_scaler = sd.get(LOSS_SCALER, self.loss_scaler) - self.dynamic_loss_scale = sd.get('dynamic_loss_scale', self.dynamic_loss_scale) - self.overflow = sd.get('overflow', self.overflow) - self.clip_grad = sd.get(CLIP_GRAD, self.clip_grad) - - ckpt_version = sd.get(DS_VERSION, False) - assert ckpt_version, f"Empty ds_version in checkpoint, not clear how to proceed" - ckpt_version = pkg_version.parse(ckpt_version) - - # zero stage 1 mode - if not self.partition_gradients: - required_version = pkg_version.parse("0.3.17") - error_str = f"ZeRO stage 1 changed in {required_version} and is not backwards compatible " \ - "with older stage 1 checkpoints. If you'd like to load an old ZeRO-1 checkpoint " \ - "please use an older version of DeepSpeed (<= 0.5.8) and set 'legacy_stage1': true in your zero config json." - assert required_version <= ckpt_version, f"Old version: {ckpt_version} {error_str}" - - def _load_legacy_checkpoint(self, state_dict_list, load_optimizer_states=True, load_from_fp32_weights=False): - r"""Loading ZeRO checkpoint - - Arguments: - state_dict_list: List of all saved ZeRO checkpoints, one for each saved partition. - Note that the number of saved partitions may differ from number of loading partitions to support - changing GPU count, specifically DP world size, between saving and loading checkpoints. - load_optimizer_states: Boolean indicating whether or not to load base optimizer states - load_from_fp32_weights: Boolean indicating whether to initialize fp32 master weights from fp32 - copies in checkpoints (no precision loss) or from model's fp16 copies (with precision loss). - """ - """ - Loads a state_dict created by an earlier call to state_dict(). - If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``, - whose parameters in turn came from ``model``, it is expected that the user - will call ``model.load_state_dict()`` before - ``fp16_optimizer_instance.load_state_dict()`` is called. - Example:: - model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half() - optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) - optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0) - ... - checkpoint = torch.load("saved.pth") - model.load_state_dict(checkpoint['model']) - optimizer.load_state_dict(checkpoint['optimizer']) - """ - - # I think it should actually be ok to reload the optimizer before the model. - dp_rank = dist.get_rank(group=self.dp_process_group) - current_rank_sd = state_dict_list[dp_rank] - self._load_global_state(current_rank_sd) - - ckpt_is_rigid = isinstance(current_rank_sd[BASE_OPTIMIZER_STATE], dict) - - # padding is always at the last rank/partition - # if DP=1024 and param-group elems=16 -> padding will be 1024-16 across all but one rank - # scenario-1 (shrink): saving w. 4 gpus -> loading w. 2 gpus - # scenario-2 (expand): saving w. 2 gpus -> loading w. 4 gpus - # if load_optimizer_states: - # if new_dp_size: - # self.strip_padding() - # self.add_padding_w_new_dp_size() - # self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE]) - - if load_optimizer_states: - if ckpt_is_rigid: - # loading rigid ckpt into either rigid or elastic exec - self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE]) - else: - if self.elastic_checkpoint: - # loading elastic into elastic exec - self._restore_elastic_base_optimizer_state(state_dict_list) - else: - # loading an elastic checkpoint into rigid exec - self._restore_base_optimizer_state(current_rank_sd[BASE_OPTIMIZER_STATE], - current_rank_sd[BASE_OPTIMIZER_STATE_STEP], - current_rank_sd[GROUP_PADDINGS]) - - # At this point, the optimizer's references to the model's fp32 parameters are up to date. - # The optimizer's hyperparameters and internal buffers are also up to date. - # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still - # out of date. There are two options. - # 1: Refresh the master params from the model's fp16 params. - # This requires less storage but incurs precision loss. - # 2: Save and restore the fp32 master copies separately. - # We choose option 1 if changing DP degree and option 2 otherwise. - # - # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device - # of their associated parameters, because it's possible those buffers might not exist yet in - # the current optimizer instance. In our case, as long as the current FP16_Optimizer has been - # constructed in the same way as the one whose state_dict we are loading, the same master params - # are guaranteed to exist, so we can just copy_() from the saved master params. - - if load_from_fp32_weights: - # option 2 from above - if self.elastic_checkpoint and not ckpt_is_rigid: - self._restore_from_elastic_fp32_weights(state_dict_list) - else: - # For non-elastic checkpoint, simply copying from saved weights of current rank is sufficient. - for current, saved in zip(self.single_partition_of_fp32_groups, - current_rank_sd[SINGLE_PARTITION_OF_FP32_GROUPS]): - src_tensor = _get_padded_tensor(saved, current.numel()) - current.data.copy_(src_tensor.data) - else: - # option 1 from above - self._restore_from_bit16_weights() - - if load_optimizer_states: - self._link_all_hp_params() - - -def _handle_overflow(cpu_sum, x, i): - import math - rank = dist.get_rank() - if rank == 0: - t_i = -1 - for v_i, v in enumerate(x.data.contiguous().view(-1)): - if not math.isfinite(float(v)): - t_i = v_i - break - logger.info(f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}") - - -def estimate_zero2_model_states_mem_needs(total_params, - num_gpus_per_node=1, - num_nodes=1, - cpu_offload=True, - additional_buffer_factor=1.5): - - total_gpus = num_nodes * num_gpus_per_node - - if cpu_offload: - gpu_mem = 2 * total_params - cpu_mem = total_params * max(4 * total_gpus, 16) * additional_buffer_factor - else: - # GPU's total_params multipliers: 2 = params_16bit, - # 18 = 2_grads_16bit + 4_grads_32bit + 4_params_32bit + 8_optimizer_states_32bit(momentum and variance) - gpu_mem = 2 * total_params + int(18 * total_params / total_gpus) - cpu_mem = total_params * 4 * num_gpus_per_node * additional_buffer_factor - - return int(cpu_mem), int(gpu_mem) - - -def model_to_params(model): - # shared params calculated only once - total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values()) - return total_params - - -def estimate_zero2_model_states_mem_needs_all_live(model, - num_gpus_per_node=1, - num_nodes=1, - additional_buffer_factor=1.5): - """ - Print out estimates on memory usage requirements for ZeRO 2 params, optim states and gradients - for a given ``model`` and hardware setup. - - If you have an actual model object, use this function and everything will be derived - automatically. - - If it's a hypothetical model, use ``estimate_zero2_model_states_mem_needs_all_cold`` where you have to pass - the ``total_params`` explicitly. - - Args: - - ``model``: ``nn.Module`` object - - ``num_gpus_per_node``: how many gpus per node (defaults to 1) - - ``num_nodes``: how many nodes (defaults to 1), - - ``additional_buffer_factor``: estimation factor (defaults to 1.5): - - """ - - total_params = model_to_params(model) - - estimate_zero2_model_states_mem_needs_all_cold(total_params=total_params, - num_gpus_per_node=num_gpus_per_node, - num_nodes=num_nodes, - additional_buffer_factor=additional_buffer_factor) - - -def estimate_zero2_model_states_mem_needs_all_cold(total_params, - num_gpus_per_node=1, - num_nodes=1, - additional_buffer_factor=1.5): - """ - Print out estimates on memory usage requirements for ZeRO 2 params, optim states and gradients - for a given ``model`` and hardware setup. - - If it's a hypothetical model, use this function where you have to pass - the ``total_params`` and ``largest_layer_params`` explicitly. - - If you have an actual model object, use ``estimate_zero2_model_states_mem_needs_all_live`` and everything - will be derived automatically. - - Args: - - ``total_params``: total model params - - ``num_gpus_per_node``: how many gpus per node (defaults to 1) - - ``num_nodes``: how many nodes (defaults to 1), - - ``additional_buffer_factor``: estimation factor (defaults to 1.5): - - """ - - def format_options(cpu_offload): - enabled = [] - device = f'{OffloadDeviceEnum.cpu:4}' if cpu_offload else "none" - enabled.append(f"offload_optimizer={device}") - return ", ".join(enabled) - - nodes_str = "nodes" if num_nodes > 1 else "node" - gpus_str = "GPUs" if num_gpus_per_node > 1 else "GPU" - print("Estimated memory needed for params, optim states and gradients for a:\n" - f"HW: Setup with {num_nodes} {nodes_str}, {num_gpus_per_node} {gpus_str} per node.\n" - f"SW: Model with {int(total_params/1e6)}M total params.") - print(" per CPU | per GPU | Options") - for cpu_offload in [True, False]: - cpu_mem, gpu_mem = estimate_zero2_model_states_mem_needs(total_params=total_params, - num_gpus_per_node=num_gpus_per_node, - num_nodes=num_nodes, - cpu_offload=cpu_offload, - additional_buffer_factor=additional_buffer_factor) - - options_str = format_options(cpu_offload=cpu_offload) - print(f" {cpu_mem/2**30:7.2f}GB | {gpu_mem/2**30:6.2f}GB | {options_str}") diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/test.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/test.py deleted file mode 100644 index 1904a8d395be80a4322bc25bd1008a80bdd8f982..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/test.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from deepspeed.runtime.zero.contiguous_memory_allocator import ContiguousMemoryAllocator - - -def test1(): - mem = ContiguousMemoryAllocator(1024, torch.half, 'cpu') - mem.print_allocation(resolution=100) - a1 = mem.allocate_tensor(64).mul_(0.0).add_(1.0) - mem.print_allocation(resolution=100) - mem.release_tensor(a1) - mem.print_allocation(resolution=100) - a2 = mem.allocate_tensor(64).mul_(0.0).add_(2.0) - a3 = mem.allocate_tensor(256).mul_(0.0).add_(3.0) - a4 = mem.allocate_tensor(128).mul_(0.0).add_(4.0) - mem.print_allocation(resolution=100) - mem.release_tensor(a3) - mem.print_allocation(resolution=100) - a5 = mem.allocate_tensor(64).mul_(0.0).add_(5.0) - a6 = mem.allocate_tensor(256).mul_(0.0).add_(6.0) - a7 = mem.allocate_tensor(128).mul_(0.0).add_(7.0) - mem.print_allocation(resolution=100) - a8 = mem.allocate_tensor(256).mul_(0.0).add_(8.0) - a9 = mem.allocate_tensor(128).mul_(0.0).add_(9.0) - mem.print_allocation(resolution=100) - mem.release_tensor(a9) - mem.release_tensor(a6) - mem.release_tensor(a2) - mem.release_tensor(a5) - - a10 = mem.allocate_tensor(512).mul_(0.0).add_(10.0) - mem.print_allocation(resolution=100) - #print(f"a4:{a4}") - #print(f"a7:{a7}") - #print(f"a8:{a8}") - #print(f"a10:{a10}") - assert (a4.norm() + a7.norm() + a8.norm() + a10.norm()).item() == 474.50, "Test failed" - - -def test2(): - mem = ContiguousMemoryAllocator(512, torch.half, 'cpu') - a1 = mem.allocate_tensor(64).mul_(0.0).add_(1.0) - a2 = mem.allocate_tensor(64).mul_(0.0).add_(2.0) - a3 = mem.allocate_tensor(64).mul_(0.0).add_(3.0) - a4 = mem.allocate_tensor(64).mul_(0.0).add_(4.0) - a5 = mem.allocate_tensor(64).mul_(0.0).add_(5.0) - a6 = mem.allocate_tensor(64).mul_(0.0).add_(6.0) - a7 = mem.allocate_tensor(64).mul_(0.0).add_(7.0) - a8 = mem.allocate_tensor(64).mul_(0.0).add_(8.0) - mem.release_tensor(a2) - mem.release_tensor(a4) - mem.release_tensor(a6) - mem.release_tensor(a8) - mem.print_allocation(resolution=100) - - a9 = mem.allocate_tensor(128).mul_(0.0).add_(9.0) - a10 = mem.allocate_tensor(64).mul_(0.0).add_(10.0) - a11 = mem.allocate_tensor(64).mul_(0.0).add_(11.0) - mem.release_tensor(a1) - mem.release_tensor(a5) - mem.print_allocation(resolution=100) - a12 = mem.allocate_tensor(128).mul_(0.0).add_(12.0) - mem.print_allocation(resolution=100) - print(f"a7:{a7}") - print(f"a9:{a9}") - print(f"a10:{a10}") - print(f"a11:{a11}") - print(f"a12:{a12}") - assert (a7.norm() + a9.norm() + a10.norm() + a11.norm() + a12.norm()) == 460.75, "TestFailed" - - -test1() -test2() diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/tiling.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/tiling.py deleted file mode 100644 index af1c5e54737a75cb66b39736718bf62169952a80..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/tiling.py +++ /dev/null @@ -1,296 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import deepspeed -from deepspeed.runtime.utils import partition_uniform as partition - - -def split_tensor_along_last_dim(tensor, partitions, contiguous_split_chunks=False): - """Split a tensor along its last dimension. Adapted from Megatron-LM. - - Arguments: - tensor: input tensor. - partitions: list of partition sizes to supply to torch.split - contiguous_split_chunks: If True, make each chunk contiguous - in memory. - """ - # Get the size and dimension. - last_dim = tensor.dim() - 1 - - # Split. - tensor_list = torch.split(tensor, partitions, dim=last_dim) - # Note: torch.split does not create contiguous tensors by default. - if contiguous_split_chunks: - return tuple(chunk.contiguous() for chunk in tensor_list) - - return tensor_list - - -class TiledLinear(torch.nn.Module): - - def __init__(self, - in_features, - out_features, - bias=True, - in_splits=1, - out_splits=1, - input_is_already_split=False, - combine_out_splits=True, - linear_cls=torch.nn.Linear, - init_linear=None, - **kwargs): - """A replacement for ``torch.nn.Linear`` that works with ZeRO-3 to reduce - memory requirements via tiling. - - TiledLinear breaks the input and output dimensions of a linear layer - into tiles that are processed in sequence. This class enables huge - linear layers when combined with ZeRO-3 because inactive tiles can be - partitioned and offloaded. - - .. note:: - We recommend using as few tiles as necessary. Tiling - significantly reduces memory usage, but can reduce throughput - for inexpensive layers. This due to the smaller kernels having - less parallelism and lower arithmetic intensity, while - introducing more frequent synchronization and communication. - - Args: - in_features (int): See ``torch.nn.Linear`` - out_features (int): See ``torch.nn.Linear`` - bias (bool, optional): See ``torch.nn.Linear`` - in_splits (int, optional): The number of tiles along the input dimension. Defaults to 1. - out_splits (int, optional): The number of tiles along the output dimension. Defaults to 1. - input_is_already_split (bool, optional): If set to ``True``, assume that the ``input_`` in - to ``forward()`` is already split into ``in_splits`` chunks. Defaults to ``False``. - combine_out_splits (bool, optional): If set to ``False``, do not combine the ``out_splits`` outputs - into a single tensor. Defaults to ``True``. - linear_cls (class, optional): The underlying class to build individual tiles. - Defaults to ``torch.nn.Linear``. - init_linear (``torch.nn.Linear``, optional): If set, copy the parameters of - ``init_linear``. Useful for debugging. Defaults to ``None``. - kwargs (dict, optional): additional keyword arguments to provide to ``linear_cls()``. - - Raises: - RuntimeError: ``in_splits`` must be within the range [1, in_features). - RuntimeError: ``out_splits`` must be within the range of [1, out_features). - """ - - super().__init__() - - if (in_splits < 1) or (in_splits > in_features): - raise RuntimeError('in splits must be in range [1, in_features].') - if (out_splits < 1) or (out_splits > out_features): - raise RuntimeError('out splits must be in range [1, out_features].') - - # global, not necessarily local - self.in_features = in_features - self.out_features = out_features - self.use_bias = bias - - self.out_splits = out_splits - self.in_splits = in_splits - self.input_is_already_split = input_is_already_split - self.combine_out_splits = combine_out_splits - - # Build partition-lists. These are CSR-style splits [0, part0, part1, ..., features] - # For example, row_parts[p] gives the start of partition p and row_parts[p+1] - # is the exclusive end. - self.in_parts = partition(num_items=in_features, num_parts=in_splits) - self.out_parts = partition(num_items=out_features, num_parts=out_splits) - - assert len(self.out_parts) == out_splits + 1 - assert len(self.in_parts) == in_splits + 1 - assert self.out_parts[0] == 0 - assert self.out_parts[out_splits] == out_features - assert self.in_parts[in_splits] == in_features - - self.linears = torch.nn.ModuleList() - for out_id in range(out_splits): - self.linears.append(torch.nn.ModuleList()) - - local_out_dim = self.out_parts[out_id + 1] - self.out_parts[out_id] - - for in_id in range(in_splits): - #if input_size is split, we only need one bias - local_bias = bias if in_id == (in_splits - 1) else False - - local_in_dim = self.in_parts[in_id + 1] - self.in_parts[in_id] - local = linear_cls(local_in_dim, local_out_dim, bias=local_bias, **kwargs) - self.linears[out_id].append(local) - - # Optionally initialize with a known tensor - if init_linear is not None: - self.copy_params_from(init_linear) - - def forward(self, input_): - if self.in_splits > 1 and not self.input_is_already_split: - input_parts = partition(input_.shape[-1], self.in_splits) - split_sizes = [input_parts[p + 1] - input_parts[p] for p in range(self.in_splits)] - inputs = self._split_global_input(input_, split_sizes) - elif self.in_splits > 1: - inputs = input_ - assert len( - inputs) == self.in_splits, f"Col splits {self.in_splits} does not match input splits {len(inputs)}" - else: - # no splits - inputs = [input_] - - outputs = [None] * self.out_splits - for out_id in range(self.out_splits): - for in_id in range(self.in_splits): - local_output = self.linears[out_id][in_id](inputs[in_id]) - - outputs[out_id] = self._reduce_local_output(in_id=in_id, - out_id=out_id, - current_out=outputs[out_id], - new_out=local_output) - - if self.combine_out_splits: - return self._combine_output_splits(outputs) - - return outputs - - def _split_global_input(self, input, split_sizes): - """Partition an input tensor along the last dimension, aligned with given splits. - - Subclasses should override this method to account for new input types. - - Args: - input (List[Tensor]): The tensor to partition along the last dimension. - split_sizes (List[int]): The size of each partition. - - Returns: - List[Any]: A list of the chunks of ``input``. - """ - return split_tensor_along_last_dim(input, split_sizes) - - def _reduce_local_output(self, in_id, out_id, current_out, new_out): - """Reduce (sum) a new local result into the existing local results. - - Subclasses should override this method. - - For a given ``out_id``, this method is called ``in_id-1`` times. The first input - split is a simple assignment. - - Args: - in_id (int): The input split that produced ``new_out``. - out_id (int): The output split that produced ``new_out``. - current_out (Any): The reduced form of all previous ``out_id`` results. - new_out (Any): The local result from forward (``in_id``, ``out_id``)e - - Returns: - Any: The combined result of ``current_out`` and ``new_out``. - """ - - if current_out is None: - #this clone is necessary to preserve auto grad - #there is some issue with inplace update for outputs that are views - return new_out.clone() - else: - return current_out + new_out - - def _combine_output_splits(self, outputs): - """Join the splits of the output into a single result. - - Args: - outputs (List[Any]): The reduced outputs for each output split. - - Returns: - Any: The combined outputs. - """ - assert len(outputs) == self.out_splits - return torch.cat(outputs, dim=-1) - - @torch.no_grad() - def copy_params_from(self, other): - """Copy the weight and bias data from ``other``. - - This is especially useful for reproducible initialization and testing. - - Equivalent to: - - .. code-block:: python - - with torch.no_grad(): - self.weight.copy_(other.weight) - if self.bias is not None: - self.bias.copy_(other.bias) - - .. note:: - If ZeRO-3 is enabled, this is a collective operation and the - updated parameters of data-parallel rank 0 will be visible on all - ranks. See :class:`deepspeed.zero.GatheredParameters` for more - information. - - - Args: - other (``torch.nn.Linear``): the linear layer to copy from. - """ - assert hasattr(other, 'weight') - assert other.weight.size() == (self.out_features, self.in_features) - if self.use_bias: - assert hasattr(other, 'bias') - assert other.bias is not None - assert other.bias.size() == (self.out_features, ) - else: - assert other.bias is None - - for row in range(self.out_splits): - rstart = self.out_parts[row] - rstop = self.out_parts[row + 1] - - for col in range(self.in_splits): - cstart = self.in_parts[col] - cstop = self.in_parts[col + 1] - - local = self.linears[row][col] - global_weight = other.weight[rstart:rstop, cstart:cstop] - with deepspeed.zero.GatheredParameters(local.weight, modifier_rank=0): - local.weight.copy_(global_weight) - - if local.bias is not None: - with deepspeed.zero.GatheredParameters(local.bias, modifier_rank=0): - local.bias.data.copy_(other.bias[rstart:rstop].data) - - -class TiledLinearReturnBias(TiledLinear): - """Wrapper for a Linear class that returns its own bias parameter, such as - used by Megatron-LM. - """ - - def _reduce_local_output(self, in_id, out_id, current_out, new_out): - """Reduces output tensors, but not the returned bias. """ - if current_out is not None: - old_tensor, old_bias = current_out - else: - old_tensor, old_bias = None, None - - assert isinstance(new_out, tuple) - assert len(new_out) == 2 - - tensor, bias = new_out - assert tensor is not None - - tensor = super()._reduce_local_output(in_id=in_id, out_id=out_id, current_out=old_tensor, new_out=tensor) - - if bias is None: - bias = old_bias - - return tensor, bias - - def _combine_output_splits(self, outputs): - # stack output tensors - tensors = [o[0] for o in outputs] - tensor = super()._combine_output_splits(tensors) - - # stack biases if applicable - biases = [o[1] for o in outputs if o[1] is not None] - if len(biases) > 0: - bias = super()._combine_output_splits(biases) - else: - bias = None - - return tensor, bias diff --git a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/utils.py b/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/utils.py deleted file mode 100644 index 2d1cf17962d8a59fee5cb1ac790129fcb6c53a10..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/runtime/zero/utils.py +++ /dev/null @@ -1,175 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import os -from typing import List, Tuple - -import torch -from deepspeed import comm as dist -from deepspeed.utils import logger -from deepspeed.ops.adam import DeepSpeedCPUAdam -from deepspeed.ops.adagrad import DeepSpeedCPUAdagrad -from deepspeed.ops.adam import FusedAdam -from deepspeed.ops.lion import DeepSpeedCPULion, FusedLion -from deepspeed.utils.nvtx import instrument_w_nvtx -from deepspeed.accelerator import get_accelerator - -# ensure we only warn once, otherwise every iteration will trigger a warning -warned = False - - -def _initialize_parameter_parallel_groups(parameter_parallel_size=None): - data_parallel_size = int(dist.get_world_size()) - parameter_parallel_size = parameter_parallel_size or data_parallel_size - logger.info("data_parallel_size: %s, parameter_parallel_size: %s", data_parallel_size, parameter_parallel_size) - assert data_parallel_size % parameter_parallel_size == 0, \ - 'world size should be divisible by parameter parallel size' - rank = dist.get_rank() - my_group = None - for i in range(data_parallel_size // parameter_parallel_size): - ranks = range(i * parameter_parallel_size, (i + 1) * parameter_parallel_size) - group = dist.new_group(ranks) - if rank in ranks: - my_group = group - return my_group - - -class ZeRORuntimeException(Exception): - pass - - -ZERO_SUPPORTED_OPTIMIZERS = [ - torch.optim.Adam, torch.optim.AdamW, FusedAdam, DeepSpeedCPUAdam, torch.optim.Adagrad, DeepSpeedCPUAdagrad, - DeepSpeedCPULion, FusedLion -] - -# Add apex FusedAdam to supported list if apex is installed -try: - import apex - if hasattr(apex, 'optimizers') and hasattr(apex.optimizers, 'FusedAdam'): - ZERO_SUPPORTED_OPTIMIZERS.append(apex.optimizers.FusedAdam) -except ImportError: - pass - - -def is_zero_supported_optimizer(optimizer): - if dist.get_rank() == 0: - logger.info(f'Checking ZeRO support for optimizer={optimizer.__class__.__name__} type={type(optimizer)}') - return type(optimizer) in ZERO_SUPPORTED_OPTIMIZERS - - -def get_lst_from_rank0(lst: List[int]) -> None: - """ - NOTE: creates both communication and synchronization overhead so should be used - sparingly - """ - lst_tensor = torch.tensor( - lst if dist.get_rank() == 0 else [-1] * len(lst), - dtype=int, - device=torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])), - requires_grad=False, - ) - dist.broadcast(lst_tensor, src=0, async_op=False) - - return list(lst_tensor.cpu().numpy()) - - -@instrument_w_nvtx -def assert_ints_same_as_other_ranks(ints: List[int]) -> None: - """ - NOTE: creates both communication and synchronization overhead so should be - used sparingly - - takes a list of ints from each rank and ensures that they are the same - across ranks, throwing an exception if they are not. - """ - rank0_ints = get_lst_from_rank0(ints) - if ints != rank0_ints: - raise RuntimeError(f"disagreement between rank0 and rank{dist.get_rank()}: " - f"rank0: {rank0_ints}, rank{dist.get_rank()}: {ints}") - - -def is_builtin_type(obj): - # https://stackoverflow.com/a/17795199 - return obj.__class__.__module__ == '__builtin__' or obj.__class__.__module__ == "builtins" - - -def isinstance_namedtuple(obj: object) -> bool: - """ - Is this an instance of namedtuple/NamedTuple? - From: https://stackoverflow.com/a/62692640 - - Args: - obj (object): An object. - - Returns: - bool: True if namedtuple/NamedTuple else False. - """ - return isinstance(obj, tuple) and hasattr(obj, '_asdict') and hasattr(obj, '_fields') - - -def is_zero_param(parameter): - if not torch.is_tensor(parameter): - return False - return hasattr(parameter, 'ds_id') - - -def apply_to_tensors_only(function, value, warning_msg_fn=None): - """ - Apply `function` to every Tensor in `value`. - - Args: - functional: The function class to apply. - value (Any): Target object to apply `function` to. - - Returns: - Any: Output of `function`. - """ - if isinstance(value, (tuple, list)): - touched_outputs = [] - for elem in value: - touched_output = apply_to_tensors_only(function, elem) - touched_outputs.append(touched_output) - - if isinstance_namedtuple(value): - # namedtuples require a slightly different syntax. - return value.__class__(*touched_outputs) - - return value.__class__(touched_outputs) - elif isinstance(value, dict): - # apply inplace to avoid recreating dict inherited objects - for key in value.keys(): - value[key] = apply_to_tensors_only(function, value[key]) - return value - - elif isinstance(value, torch.Tensor): - # this also applies to torch.Tensor's subclasses like torch.nn.parameter.Parameter - touched_output = function(value) - - # restore zero param attributes if those get stripped by `backward_function` - if not is_zero_param(touched_output) and is_zero_param(value): - touched_output.ds_param_alias = value - - return touched_output - else: - if not is_builtin_type(value): - global warned - if warning_msg_fn and not warned and dist.get_rank() == 0: - logger.warning(warning_msg_fn(value)) - warned = True - return value - - -def get_mapping_to_flat_buffer(tensors: List[torch.Tensor]) -> List[Tuple[torch.Tensor, int, int]]: - tensor_infos: List[Tuple[torch.Tensor, int, int]] = [] - - offset = 0 - for tensor in tensors: - tensor_numel = tensor.numel() - # record some data so we can restore the device tensor later - tensor_infos.append((tensor, offset, tensor_numel)) - offset += tensor_numel - - return tensor_infos diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/sequence/__init__.py deleted file mode 100644 index 208299fb8c50f73468d293b6fa5dca71649d62e7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/sequence/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 7ee6e3bc5a29137ed25c58a1a576d03bf32623e8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/cross_entropy.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/cross_entropy.cpython-310.pyc deleted file mode 100644 index 605d658129f85945b2b2206640df80c428f3bc6d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/cross_entropy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/fpdt_layer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/fpdt_layer.cpython-310.pyc deleted file mode 100644 index 085469b669ddea5c734839cadd4edaba3b7588b3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/fpdt_layer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/layer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/layer.cpython-310.pyc deleted file mode 100644 index 54065dbe50bd295f3671478eed9b98da3fb26c16..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/sequence/__pycache__/layer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/cross_entropy.py b/venv/lib/python3.10/site-packages/deepspeed/sequence/cross_entropy.py deleted file mode 100644 index baa7bc1ea7a803055f2cca20141ea6888cc3f8c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/sequence/cross_entropy.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -import deepspeed.comm as dist - - -class _VocabSequenceParallelCrossEntropy(torch.autograd.Function): - - @staticmethod - def forward(ctx, vocab_seq_parallel_logits, target, sp_group): - # vocab_seq_parallel_logits: [S/P, B, V] - # target: [S/P, B] - # return: [S, B] - - # Need softmax for backward - softmax = torch.nn.functional.softmax(vocab_seq_parallel_logits, dim=-1) - ctx.vocab_size = vocab_seq_parallel_logits.size(2) - loss = torch.nn.functional.nll_loss(softmax.log().view(-1, ctx.vocab_size), target.view(-1), reduction='none') - - sp_world_size = dist.get_world_size(sp_group) - sp_rank = dist.get_rank(sp_group) - ctx.sp_world_size = sp_world_size - ctx.sp_rank = sp_rank - ctx.seqlen = vocab_seq_parallel_logits.size(0) * sp_world_size - batch_size = vocab_seq_parallel_logits.size(1) - - loss_all = torch.empty(ctx.seqlen, - batch_size, - dtype=vocab_seq_parallel_logits.dtype, - device=vocab_seq_parallel_logits.device) - dist.all_gather_into_tensor(loss_all, loss, group=sp_group) - - ctx.save_for_backward(softmax, target) - - return loss_all - - @staticmethod - def backward(ctx, grad_output): - softmax, target = ctx.saved_tensors - - step_seqlen = ctx.seqlen // ctx.sp_world_size - sp_rank = ctx.sp_rank - grad_output_part = grad_output[step_seqlen * sp_rank:step_seqlen * (sp_rank + 1), :] - - grad_input = softmax - grad_2d = grad_input.view(-1, ctx.vocab_size) - arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device) - - grad_2d[arange_1d, target.view(-1)] -= 1 - grad_input.mul_(grad_output_part.unsqueeze(dim=-1)) - - return grad_input, None, None, None - - -def vocab_sequence_parallel_cross_entropy(vocab_parallel_logits, target, sp_group): - return _VocabSequenceParallelCrossEntropy.apply(vocab_parallel_logits, target, sp_group) diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/fpdt_layer.py b/venv/lib/python3.10/site-packages/deepspeed/sequence/fpdt_layer.py deleted file mode 100644 index 4fa2cc988a19490f91c90867fe9a01d8a120e2e8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/sequence/fpdt_layer.py +++ /dev/null @@ -1,1225 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch - -from typing import Optional, Any, Tuple -from torch import Tensor -from packaging import version -import deepspeed.comm as dist -from deepspeed.accelerator import get_accelerator - -try: - import flash_attn - from flash_attn.flash_attn_interface import _flash_attn_forward, _flash_attn_backward - flash_attn_version = version.parse(flash_attn.__version__) -except ImportError: - _flash_attn_forward = None - _flash_attn_backward = None - -from einops import rearrange -from .layer import single_all_to_all, apply_rotary_pos_emb - - -def _rotate_half_backward(x): - x = rearrange(x, '... (j d) -> ... j d', j=2) - x1, x2 = x.unbind(dim=-2) - return torch.cat((x2, -x1), dim=-1) - - -def apply_rotary_pos_emb_backward(grad_output, freqs_cos, freqs_sin): - rot_dim = freqs_cos.shape[-1] - grad, grad_pass = grad_output[..., :rot_dim], grad_output[..., rot_dim:] - grad_t = (grad * freqs_cos) + (_rotate_half_backward(grad * freqs_sin)) - grad = grad_t if grad_pass.shape[-1] == 0 else torch.cat((grad_t, grad_pass), dim=-1) - return grad - - -def _update_out_and_lse( - out: torch.Tensor, - lse: torch.Tensor, - block_out: torch.Tensor, - block_lse: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - - block_out = block_out.to(torch.float32) - block_lse = block_lse.transpose(-2, -1).unsqueeze(dim=-1) - - new_lse = lse + torch.log1p(torch.exp(block_lse - lse)) - - out = torch.exp(lse - new_lse) * out + torch.exp(block_lse - new_lse) * block_out - - lse = new_lse - return out, lse - - -def update_out_and_lse( - out: Optional[torch.Tensor], - lse: Optional[torch.Tensor], - block_out: torch.Tensor, - block_lse: torch.Tensor, - slice_=None, -) -> Tuple[torch.Tensor, torch.Tensor]: - if out is None: - if slice_ is not None: - raise RuntimeError("first update_out_and_lse should not pass slice_ args") - out = block_out.to(torch.float32) - lse = block_lse.permute(0, 2, 1).contiguous().unsqueeze(dim=-1).contiguous() - elif slice_ is not None: - slice_out, slice_lse = out[slice_], lse[slice_] - slice_out, slice_lse = _update_out_and_lse(slice_out, slice_lse, block_out, block_lse) - out[slice_], lse[slice_] = slice_out, slice_lse - else: - out, lse = _update_out_and_lse(out, lse, block_out, block_lse) - return out, lse - - -class FPDT_InputConstruct(torch.nn.Module): - - def __init__(self, tokens, labels, loss_mask, attention_mask, position_ids, args, sp_size, sp_rank) -> None: - - super(FPDT_InputConstruct, self).__init__() - self.tokens = tokens - self.labels = labels - self.loss_mask = loss_mask - self.attention_mask = attention_mask - self.position_ids = position_ids - global_seq_len = tokens.shape[1] - batch_size = tokens.shape[0] - assert global_seq_len % sp_size == 0 - assert global_seq_len % args.ds_sequence_parallel_fpdt_chunk_size == 0 - num_chunk_per_gpu = global_seq_len // args.ds_sequence_parallel_fpdt_chunk_size - local_seq_len = global_seq_len // sp_size - assert local_seq_len % num_chunk_per_gpu == 0 - - self.num_chunk_per_gpu = num_chunk_per_gpu - self.chunk_size = local_seq_len // num_chunk_per_gpu - self.sp_size = sp_size - self.sp_rank = sp_rank - self.global_seq_len = global_seq_len - self.local_seq_len = local_seq_len - self.batch_size = batch_size - self.device = tokens.device - - def generate(self): - device = self.device - totalChunks = self.global_seq_len // self.chunk_size - token_chunk_idx = torch.arange(self.global_seq_len, device=device, dtype=torch.int) // self.chunk_size - chunk_to_gpu = torch.arange(totalChunks, device=device, dtype=torch.int) - chunk_to_gpu = chunk_to_gpu.reshape(self.num_chunk_per_gpu, -1).t().contiguous() - - gather_chunk = chunk_to_gpu.flatten().unsqueeze(1).contiguous() - mask = gather_chunk == token_chunk_idx - - indices = mask.nonzero(as_tuple=False) - gather_indices = indices[:, 0] - token_chunk_indices = indices[:, 1] - indices = torch.cat([token_chunk_indices[gather_indices == i] for i in range(gather_chunk.shape[0])]) - load_balanced_loss_mask = self.loss_mask[:, indices] if self.loss_mask is not None else self.loss_mask - - indices = indices.reshape(-1, self.chunk_size)[self.num_chunk_per_gpu * self.sp_rank:self.num_chunk_per_gpu * - (self.sp_rank + 1)].flatten().contiguous() - load_balanced_tokens = self.tokens[:, indices] - load_balanced_labels = self.labels[:, indices] if self.labels is not None else self.labels - - load_balanced_attention_mask = self.attention_mask if self.attention_mask is not None else self.attention_mask - load_balanced_position_ids = self.position_ids[:, - indices] if self.position_ids is not None else self.position_ids - - return load_balanced_tokens, load_balanced_labels, load_balanced_loss_mask, load_balanced_attention_mask, load_balanced_position_ids - - -class _FPDTGPUAttentionImpl_(torch.autograd.Function): - generate_vmap_rule = False - - @staticmethod - def forward(ctx: Any, - layernorm_output, - attention_mask, - inference_params, - rotary_pos_emb, - spg, - scatter_idx, - gather_idx, - hidden_size, - projection_size, - hidden_size_per_attention_head, - kv_projection_size, - qkv_linear_weight, - qkv_linear_bias, - dropout, - num_chunks=8, - cpu_offloading=True): - - do_save = layernorm_output.requires_grad - - if rotary_pos_emb is not None: - pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3) - ctx.pos_emb_cos = pos_emb_cos - ctx.pos_emb_sin = pos_emb_sin - else: - ctx.pos_emb_cos = None - ctx.pos_emb_sin = None - - with torch.no_grad(): - per_gpu_seq_len = layernorm_output.shape[0] - chunk_size = per_gpu_seq_len // num_chunks - assert chunk_size * num_chunks == per_gpu_seq_len - assert attention_mask is None - ctx.num_chunks = num_chunks - ctx.cpu_offloading = cpu_offloading - ctx.spg = spg - ctx.scatter_idx = scatter_idx - ctx.gather_idx = gather_idx - - device = get_accelerator().current_device_name() - ctx.device = device - ctx.dtype = layernorm_output.dtype - ctx.projection_size = projection_size - ctx.kv_projection_size = kv_projection_size - - global_q = [] - global_k = [] - global_v = [] - - ctx.softmax_scale = hidden_size_per_attention_head**(-0.5) - - ctx.dropout_p = dropout - ctx.window_size = (-1, -1) - ctx.alibi_slopes = None - - batch_size = layernorm_output.shape[1] - - global_o = [None for _ in range(num_chunks)] - global_lse = [None for _ in range(num_chunks)] - - for i in range(num_chunks): - - st = chunk_size * i - ed = st + chunk_size - - qkv_chunk = torch.matmul(layernorm_output[st:ed], qkv_linear_weight.t()) + qkv_linear_bias - - q_chunk = qkv_chunk[:, :, :projection_size].contiguous().reshape( - qkv_chunk.shape[0], qkv_chunk.shape[1], -1, - hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd - q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg) - global_q_chunk_len = q_chunk.shape[1] - if rotary_pos_emb is not None: - q_chunk = apply_rotary_pos_emb(q_chunk, - pos_emb_cos[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)], - pos_emb_sin[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)]) - global_q.append(q_chunk) - - k_chunk = qkv_chunk[:, :, projection_size:projection_size + kv_projection_size].contiguous().reshape( - qkv_chunk.shape[0], qkv_chunk.shape[1], -1, - hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd - k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg) - if rotary_pos_emb is not None: - k_chunk = apply_rotary_pos_emb(k_chunk, - pos_emb_cos[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)], - pos_emb_sin[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)]) - global_k.append(k_chunk) - - v_chunk = qkv_chunk[:, :, projection_size + kv_projection_size:].contiguous().reshape( - qkv_chunk.shape[0], qkv_chunk.shape[1], -1, - hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd - v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg) - global_v.append(v_chunk) - - for k_i in range(len(global_k)): - causal_chunk = i == k_i - if flash_attn_version >= version.parse("2.6.0"): - block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward(global_q[i], - global_k[k_i], - global_v[k_i], - ctx.dropout_p, - ctx.softmax_scale, - causal=causal_chunk, - window_size=ctx.window_size, - softcap=0.0, - alibi_slopes=ctx.alibi_slopes, - return_softmax=False) - else: - block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward(global_q[i], - global_k[k_i], - global_v[k_i], - ctx.dropout_p, - ctx.softmax_scale, - causal=causal_chunk, - window_size=ctx.window_size, - alibi_slopes=ctx.alibi_slopes, - return_softmax=False) - - global_o[i], global_lse[i] = update_out_and_lse(global_o[i], global_lse[i], block_out, block_lse) - - global_o[i] = global_o[i].to(q_chunk.dtype) - - output = [None for i in range(num_chunks)] - - for i in range(num_chunks): - global_lse[i] = global_lse[i][:, :, :, 0].permute(0, 2, 1).contiguous() - output[i] = single_all_to_all(global_o[i].to(ctx.dtype).contiguous(), gather_idx, scatter_idx, 0, spg) - output = torch.cat(output, dim=1) - - head_dim = output.shape[-1] - - if do_save: - ctx.save_for_backward(layernorm_output) - ctx.global_q = global_q - ctx.global_k = global_k - ctx.global_v = global_v - ctx.attn_output = global_o - ctx.attn_lse = global_lse - ctx.head_dim = head_dim - ctx.batch_size = batch_size - - ctx.qkv_linear_weight = qkv_linear_weight - ctx.qkv_linear_bias = qkv_linear_bias - - return output - - @staticmethod - def backward(ctx, grad_output): - - num_chunks = ctx.num_chunks - device = ctx.device - dtype = ctx.dtype - spg = ctx.spg - scatter_idx = ctx.scatter_idx - gather_idx = ctx.gather_idx - softmax_scale = ctx.softmax_scale - dropout_p = ctx.dropout_p - window_size = ctx.window_size - alibi_slopes = ctx.alibi_slopes - - projection_size = ctx.projection_size - kv_projection_size = ctx.kv_projection_size - - layernorm_output = ctx.saved_tensors[0] - - global_q = ctx.global_q - global_k = ctx.global_k - global_v = ctx.global_v - attn_output = ctx.attn_output - lse = ctx.attn_lse - - qkv_linear_weight = ctx.qkv_linear_weight - qkv_linear_bias = ctx.qkv_linear_bias - - input_chunk_size = layernorm_output.shape[0] // num_chunks - grad_layernorm_output = [ - torch.zeros((input_chunk_size, layernorm_output.shape[1], layernorm_output.shape[2]), - device=device, - dtype=dtype) for _ in range(num_chunks) - ] - - grad_global_attn_output = [] - chunk_size = grad_output.shape[1] // num_chunks - - for i in range(num_chunks): - st = chunk_size * i - ed = st + chunk_size - grad_global_attn_output.append( - single_all_to_all(grad_output[:, st:ed].contiguous(), scatter_idx, gather_idx, 0, spg)) - - del grad_output - - dq = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)] - dk = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)] - dv = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)] - - grad_qkv_linear_weight = torch.zeros(qkv_linear_weight.shape, - device=qkv_linear_weight.device, - dtype=torch.float) - grad_qkv_linear_bias = torch.zeros(qkv_linear_bias.shape, device=qkv_linear_weight.device, dtype=torch.float) - - for i in range(num_chunks): - k_chunk = global_k[i] - v_chunk = global_v[i] - - for q_i in range(num_chunks): - no_computation = q_i < i - if no_computation: - continue - - causal_chunk = q_i == i - - q_chunk = global_q[q_i] - attn_output_chunk = attn_output[q_i] - lse_chunk = lse[q_i] - d_out = grad_global_attn_output[q_i] - - dq_this = torch.zeros(global_q[0].shape, dtype=dtype, device=device) - dk_this = torch.zeros(global_k[0].shape, dtype=dtype, device=device) - dv_this = torch.zeros(global_v[0].shape, dtype=dtype, device=device) - - if flash_attn_version >= version.parse("2.6.0"): - _flash_attn_backward(d_out, - q_chunk, - k_chunk, - v_chunk, - attn_output_chunk, - lse_chunk, - dq_this, - dk_this, - dv_this, - dropout_p, - softmax_scale, - causal_chunk, - window_size, - softcap=0.0, - alibi_slopes=alibi_slopes, - deterministic=False, - rng_state=None) - else: - _flash_attn_backward(d_out, - q_chunk, - k_chunk, - v_chunk, - attn_output_chunk, - lse_chunk, - dq_this, - dk_this, - dv_this, - dropout_p, - softmax_scale, - causal_chunk, - window_size, - alibi_slopes=alibi_slopes, - deterministic=False, - rng_state=None) - - dq[q_i].add_(dq_this.to(torch.float)) - dk[i].add_(dk_this.to(torch.float)) - dv[i].add_(dv_this.to(torch.float)) - - dk_seq_len = dk[i].shape[1] - - if ctx.pos_emb_cos is not None: - dk[i] = apply_rotary_pos_emb_backward(dk[i].to(dtype), - ctx.pos_emb_cos[:, dk_seq_len * i:dk_seq_len * (i + 1)], - ctx.pos_emb_sin[:, dk_seq_len * i:dk_seq_len * (i + 1)]) - else: - dk[i] = dk[i].to(dtype) - dv[i] = dv[i].to(dtype) - dk[i] = single_all_to_all(dk[i].contiguous(), gather_idx, scatter_idx, 0, spg) - dv[i] = single_all_to_all(dv[i].contiguous(), gather_idx, scatter_idx, 0, spg) - - input_st = i * input_chunk_size - input_ed = input_st + input_chunk_size - - input_chunk = layernorm_output[input_st:input_ed].reshape(-1, layernorm_output.shape[-1]) - - dk[i] = dk[i].flatten(2).permute(1, 0, 2) - dv[i] = dv[i].flatten(2).permute(1, 0, 2) - l, b = dk[i].shape[0], dk[i].shape[1] - grad_qkv_linear_weight[projection_size:projection_size + kv_projection_size].add_( - torch.matmul(dk[i].reshape(l * b, -1).t(), input_chunk)) - grad_qkv_linear_weight[projection_size + kv_projection_size:].add_( - torch.matmul(dv[i].reshape(l * b, -1).t(), input_chunk)) - grad_qkv_linear_bias[projection_size:projection_size + kv_projection_size].add_(dk[i].sum(0).sum(0)) - grad_qkv_linear_bias[projection_size + kv_projection_size:].add_(dv[i].sum(0).sum(0)) - - grad_layernorm_output[i].add_( - torch.matmul(dk[i], qkv_linear_weight[projection_size:projection_size + kv_projection_size])) - grad_layernorm_output[i].add_(torch.matmul(dv[i], - qkv_linear_weight[projection_size + kv_projection_size:])) - - dk[i] = None - dv[i] = None - - for i in range(num_chunks): - dq_seq_len = dq[i].shape[1] - if ctx.pos_emb_cos is not None: - dq[i] = apply_rotary_pos_emb_backward(dq[i].to(dtype), - ctx.pos_emb_cos[:, dq_seq_len * i:dq_seq_len * (i + 1)], - ctx.pos_emb_sin[:, dq_seq_len * i:dq_seq_len * (i + 1)]) - else: - dq[i] = dq[i].to(dtype) - dq[i] = single_all_to_all(dq[i].to(dtype).contiguous(), gather_idx, scatter_idx, 0, spg) - - input_chunk = layernorm_output[:input_chunk_size].reshape(-1, layernorm_output.shape[-1]) - layernorm_output = layernorm_output[input_chunk_size:] - - dq[i] = dq[i].flatten(2).permute(1, 0, 2) - l, b = dq[i].shape[0], dq[i].shape[1] - grad_qkv_linear_weight[:projection_size].add_(torch.matmul(dq[i].reshape(l * b, -1).t(), input_chunk)) - grad_qkv_linear_bias[:projection_size].add_(dq[i].sum(0).sum(0)) - - grad_layernorm_output[i].add_(torch.matmul(dq[i], qkv_linear_weight[:projection_size])) - - dq[i] = None - - return torch.cat( - grad_layernorm_output, - dim=0).to(dtype), None, None, None, None, None, None, None, None, None, None, grad_qkv_linear_weight.to( - dtype), grad_qkv_linear_bias.to(dtype), None, None, None - - -class SequenceChunk: - - def __init__(self, chunk: torch.Tensor, device=None, is_in_use=False): - - self.chunk_shape = chunk.shape - self.chunk_dtype = chunk.dtype - self.device = chunk.device if device is None else device - - cpu_chunk = torch.empty(chunk.shape, dtype=chunk.dtype, device='cpu', pin_memory=True) - - if get_accelerator().on_accelerator(chunk): - cpu_chunk.copy_(chunk, non_blocking=True) - else: - cpu_chunk = chunk - - self.cpu_chunk = cpu_chunk - - self.gpu_chunk = chunk if is_in_use else None - - def load_to_gpu(self): - assert self.gpu_chunk is None - if self.gpu_chunk is not None: - pass - else: - gpu_chunk = torch.empty(self.chunk_shape, device=self.device, dtype=self.chunk_dtype) - gpu_chunk.copy_(self.cpu_chunk, non_blocking=True) - self.gpu_chunk = gpu_chunk - - def get_gpu_chunk(self): - assert self.gpu_chunk is not None and self.gpu_chunk.device == self.device - return self.gpu_chunk - - def check_gpu_chunk(self, ): - assert (self.gpu_chunk is not None) and ( - self.gpu_chunk.device == self.device - ), f"gpu_chunk {self.gpu_chunk is not None} shound be on {self.device}, but it is now on {self.gpu_chunk.device}" - return True - - def offload(self): - assert self.gpu_chunk is not None and self.gpu_chunk.device == self.device - del self.gpu_chunk - self.gpu_chunk = None - - def overwrite_to_cpu(self): - assert self.gpu_chunk is not None and self.gpu_chunk.device == self.device - self.cpu_chunk.copy_(self.gpu_chunk, non_blocking=True) - - -class _FPDTGPUOffloadingAttentionImpl_(torch.autograd.Function): - generate_vmap_rule = False - - @staticmethod - def forward(ctx: Any, - layernorm_output, - attention_mask, - inference_params, - rotary_pos_emb, - spg, - scatter_idx, - gather_idx, - hidden_size, - projection_size, - hidden_size_per_attention_head, - kv_projection_size, - qkv_linear_weight, - qkv_linear_bias, - dropout, - num_chunks=8, - cpu_offloading=True): - - do_save = layernorm_output.requires_grad - - if rotary_pos_emb is not None: - pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3) - ctx.pos_emb_cos = pos_emb_cos - ctx.pos_emb_sin = pos_emb_sin - else: - ctx.pos_emb_cos = None - ctx.pos_emb_sin = None - with torch.no_grad(): - per_gpu_seq_len = layernorm_output.shape[0] - chunk_size = per_gpu_seq_len // num_chunks - assert chunk_size * num_chunks == per_gpu_seq_len - assert attention_mask is None - ctx.num_chunks = num_chunks - ctx.cpu_offloading = cpu_offloading - ctx.spg = spg - ctx.scatter_idx = scatter_idx - ctx.gather_idx = gather_idx - - ctx.chunk_size = chunk_size - device = get_accelerator().current_device_name() - ctx.device = device - ctx.dtype = layernorm_output.dtype - ctx.projection_size = projection_size - ctx.kv_projection_size = kv_projection_size - - global_q = [] - global_k = [] - global_v = [] - - ctx.softmax_scale = hidden_size_per_attention_head**(-0.5) - - ctx.dropout_p = dropout - ctx.window_size = (-1, -1) - ctx.alibi_slopes = None - - batch_size = layernorm_output.shape[1] - - global_o = [] - global_lse = [] - - layernorm_output_cpu = [] - final_output = [] - - offload_stream = get_accelerator().Stream() - general_offload_stream = get_accelerator().Stream() - compute_stream = get_accelerator().default_stream() - - q_compute_chunk_idx = 0 - kv_compute_chunk_idx = 0 - for i in range(num_chunks): - - qkv_chunk = torch.matmul(layernorm_output[:chunk_size], - qkv_linear_weight.t()) + qkv_linear_bias # torch.Size([18126, 1, 12288]) - - with get_accelerator().stream(general_offload_stream): - layernorm_output_cpu.append(SequenceChunk(layernorm_output[:chunk_size])) - - layernorm_output = layernorm_output[chunk_size:] - - q_chunk = qkv_chunk[:, :, :projection_size].contiguous().reshape( - qkv_chunk.shape[0], qkv_chunk.shape[1], -1, - hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd - q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg) - global_q_chunk_len = q_chunk.shape[1] - - k_chunk = qkv_chunk[:, :, projection_size:projection_size + kv_projection_size].contiguous().reshape( - qkv_chunk.shape[0], qkv_chunk.shape[1], -1, - hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd - k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg) - - v_chunk = qkv_chunk[:, :, projection_size + kv_projection_size:].contiguous().reshape( - qkv_chunk.shape[0], qkv_chunk.shape[1], -1, - hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd - v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg) - - dist.barrier() - - if ctx.pos_emb_cos is not None: - pos_emb_cos_chunk = pos_emb_cos[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)] - pos_emb_sin_chunk = pos_emb_sin[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)] - - q_chunk = apply_rotary_pos_emb(q_chunk, pos_emb_cos_chunk, pos_emb_sin_chunk) - k_chunk = apply_rotary_pos_emb(k_chunk, pos_emb_cos_chunk, pos_emb_sin_chunk) - - compute_stream.wait_stream(offload_stream) - compute_stream.synchronize() - with get_accelerator().stream(offload_stream): - global_q.append(SequenceChunk(q_chunk, is_in_use=True)) - global_k.append(SequenceChunk(k_chunk, is_in_use=True)) - global_v.append(SequenceChunk(v_chunk, is_in_use=True)) - - del qkv_chunk - - cur_attn_output = None - cur_attn_lse = None - for k_i in range(len(global_k)): - causal_chunk = i == k_i - with get_accelerator().stream(compute_stream): - if flash_attn_version >= version.parse("2.6.0"): - block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward( - global_q[q_compute_chunk_idx].get_gpu_chunk(), - global_k[kv_compute_chunk_idx].get_gpu_chunk(), - global_v[kv_compute_chunk_idx].get_gpu_chunk(), - ctx.dropout_p, - ctx.softmax_scale, - causal=causal_chunk, - window_size=ctx.window_size, - softcap=0.0, - alibi_slopes=ctx.alibi_slopes, - return_softmax=False) - else: - block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward( - global_q[q_compute_chunk_idx].get_gpu_chunk(), - global_k[kv_compute_chunk_idx].get_gpu_chunk(), - global_v[kv_compute_chunk_idx].get_gpu_chunk(), - ctx.dropout_p, - ctx.softmax_scale, - causal=causal_chunk, - window_size=ctx.window_size, - alibi_slopes=ctx.alibi_slopes, - return_softmax=False) - cur_attn_output, cur_attn_lse = update_out_and_lse(cur_attn_output, cur_attn_lse, block_out, - block_lse) - - can_offload_kv = True - if k_i != (len(global_k) - 1) or i != (num_chunks - 1): - if k_i != (len(global_k) - 1): - next_kv_compute_chunk_idx = k_i + 1 - else: - next_kv_compute_chunk_idx = 0 - - if next_kv_compute_chunk_idx == kv_compute_chunk_idx: - can_offload_kv = False - else: - if next_kv_compute_chunk_idx != (len(global_k) - 1): - with get_accelerator().stream(offload_stream): - global_k[next_kv_compute_chunk_idx].load_to_gpu() - global_v[next_kv_compute_chunk_idx].load_to_gpu() - - if i == num_chunks - 1 and k_i == num_chunks - 1: - with get_accelerator().stream(offload_stream): - global_q[0].load_to_gpu() - global_k[0].load_to_gpu() - global_v[0].load_to_gpu() - global_o[0].load_to_gpu() - global_lse[0].load_to_gpu() - - compute_stream.wait_stream(offload_stream) - compute_stream.synchronize() - - if can_offload_kv: - global_k[kv_compute_chunk_idx].offload() - global_v[kv_compute_chunk_idx].offload() - kv_compute_chunk_idx = next_kv_compute_chunk_idx - - global_q[q_compute_chunk_idx].offload() - q_compute_chunk_idx += 1 - - all2all_output = single_all_to_all( - cur_attn_output.to(ctx.dtype).contiguous(), gather_idx, scatter_idx, 0, spg) - final_output.append(all2all_output) - with get_accelerator().stream(general_offload_stream): - global_o.append(SequenceChunk(cur_attn_output.to(ctx.dtype))) - global_lse.append(SequenceChunk(cur_attn_lse[:, :, :, 0].permute(0, 2, 1).contiguous())) - - compute_stream.wait_stream(general_offload_stream) - compute_stream.synchronize() - - final_output = torch.cat(final_output, dim=1) - - head_dim = final_output.shape[-1] - - if do_save: - ctx.layernorm_output = layernorm_output_cpu - ctx.global_q = global_q - ctx.global_k = global_k - ctx.global_v = global_v - ctx.attn_output = global_o - ctx.attn_lse = global_lse - ctx.head_dim = head_dim - ctx.batch_size = batch_size - - ctx.qkv_linear_weight = qkv_linear_weight - ctx.qkv_linear_bias = qkv_linear_bias - - return final_output - - @staticmethod - def backward(ctx, grad_output): - num_chunks = ctx.num_chunks - device = grad_output.device - dtype = ctx.dtype - spg = ctx.spg - scatter_idx = ctx.scatter_idx - gather_idx = ctx.gather_idx - softmax_scale = ctx.softmax_scale - dropout_p = ctx.dropout_p - window_size = ctx.window_size - alibi_slopes = ctx.alibi_slopes - - projection_size = ctx.projection_size - kv_projection_size = ctx.kv_projection_size - - layernorm_output = ctx.layernorm_output - - global_q = ctx.global_q - global_k = ctx.global_k - global_v = ctx.global_v - attn_output = ctx.attn_output - lse = ctx.attn_lse - - qkv_linear_weight = ctx.qkv_linear_weight - qkv_linear_bias = ctx.qkv_linear_bias - - offload_stream = get_accelerator().Stream() - general_offload_stream = get_accelerator().Stream() - compute_stream = get_accelerator().default_stream() - - chunk_size = grad_output.shape[1] // num_chunks - assert chunk_size == layernorm_output[0].cpu_chunk.shape[0] - - grad_layernorm_output = [ - torch.zeros(layernorm_output[0].chunk_shape, device=device, dtype=dtype) for _ in range(num_chunks) - ] - - grad_global_attn_output = [None for _ in range(num_chunks)] - - q_compute_chunk_idx = 0 - kv_compute_chunk_idx = 0 - last_q_accum_idx = 0 - - with get_accelerator().stream(general_offload_stream): - layernorm_output[0].load_to_gpu() - grad_qkv_linear_weight = torch.zeros(qkv_linear_weight.shape, - device=qkv_linear_weight.device, - dtype=torch.float) - grad_qkv_linear_bias = torch.zeros(qkv_linear_bias.shape, - device=qkv_linear_weight.device, - dtype=torch.float) - - grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(), scatter_idx, - gather_idx, 0, spg) - get_accelerator().synchronize() - grad_output = grad_output[:, chunk_size:] - - with get_accelerator().stream(offload_stream): - grad_global_attn_output[0] = SequenceChunk(grad_global_attn_output_chunk, is_in_use=True) - dq = [ - SequenceChunk(torch.zeros(global_q[0].chunk_shape, dtype=torch.float, device=device), is_in_use=True) - ] + [ - SequenceChunk(torch.zeros(global_q[0].chunk_shape, dtype=torch.float, device='cpu', pin_memory=True), - device) for _ in range(num_chunks - 1) - ] - dk_accum = torch.zeros(global_k[0].chunk_shape, dtype=torch.float, device=device) - dv_accum = torch.zeros(global_v[0].chunk_shape, dtype=torch.float, device=device) - - for i in range(num_chunks): - for q_i in range(num_chunks): - no_computation = q_i < i - if no_computation: - continue - - causal_chunk = q_i == i - - dq_this = torch.zeros(global_q[0].chunk_shape, dtype=dtype, device=device) - dk_this = torch.zeros(global_k[0].chunk_shape, dtype=dtype, device=device) - dv_this = torch.zeros(global_v[0].chunk_shape, dtype=dtype, device=device) - - with get_accelerator().stream(compute_stream): - if flash_attn_version >= version.parse("2.6.0"): - _flash_attn_backward(grad_global_attn_output[q_compute_chunk_idx].get_gpu_chunk(), - global_q[q_compute_chunk_idx].get_gpu_chunk(), - global_k[kv_compute_chunk_idx].get_gpu_chunk(), - global_v[kv_compute_chunk_idx].get_gpu_chunk(), - attn_output[q_compute_chunk_idx].get_gpu_chunk(), - lse[q_compute_chunk_idx].get_gpu_chunk(), - dq_this, - dk_this, - dv_this, - dropout_p, - softmax_scale, - causal_chunk, - window_size, - softcap=0.0, - alibi_slopes=alibi_slopes, - deterministic=False, - rng_state=None) - else: - _flash_attn_backward(grad_global_attn_output[q_compute_chunk_idx].get_gpu_chunk(), - global_q[q_compute_chunk_idx].get_gpu_chunk(), - global_k[kv_compute_chunk_idx].get_gpu_chunk(), - global_v[kv_compute_chunk_idx].get_gpu_chunk(), - attn_output[q_compute_chunk_idx].get_gpu_chunk(), - lse[q_compute_chunk_idx].get_gpu_chunk(), - dq_this, - dk_this, - dv_this, - dropout_p, - softmax_scale, - causal_chunk, - window_size, - alibi_slopes=alibi_slopes, - deterministic=False, - rng_state=None) - - if i != (len(global_k) - 1): - if q_i != (len(global_q) - 1): - next_q_compute_chunk_idx = q_i + 1 - else: - next_q_compute_chunk_idx = i + 1 - - can_offload_q = True - - if next_q_compute_chunk_idx == q_compute_chunk_idx: - can_offload_q = False - else: - with get_accelerator().stream(offload_stream): - if i > 0 or q_i > 0: - if can_offload_q and last_q_accum_idx != i: # the first q chunk calculate in the loop will be sent out, therefore we do not offload it - dq[last_q_accum_idx].offload() - dq[next_q_compute_chunk_idx].load_to_gpu() - global_q[next_q_compute_chunk_idx].load_to_gpu() - attn_output[next_q_compute_chunk_idx].load_to_gpu() - lse[next_q_compute_chunk_idx].load_to_gpu() - if grad_global_attn_output[next_q_compute_chunk_idx] is not None: - grad_global_attn_output[next_q_compute_chunk_idx].load_to_gpu() - - if grad_global_attn_output[next_q_compute_chunk_idx] is None: - grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(), - scatter_idx, gather_idx, 0, spg) - dist.barrier() - grad_output = grad_output[:, chunk_size:] - grad_global_attn_output[next_q_compute_chunk_idx] = SequenceChunk( - grad_global_attn_output_chunk, is_in_use=True) - - compute_stream.wait_stream(offload_stream) - compute_stream.synchronize() - - with get_accelerator().stream(compute_stream): - dq[q_compute_chunk_idx].check_gpu_chunk() - dq[q_compute_chunk_idx].gpu_chunk.add_(dq_this) - dk_accum.add_(dk_this) - dv_accum.add_(dv_this) - - offload_stream.wait_stream(compute_stream) - with get_accelerator().stream(offload_stream): - dq[q_compute_chunk_idx].overwrite_to_cpu() - - if can_offload_q: - global_q[q_compute_chunk_idx].offload() - attn_output[q_compute_chunk_idx].offload() - lse[q_compute_chunk_idx].offload() - grad_global_attn_output[q_compute_chunk_idx].offload() - - last_q_accum_idx = q_compute_chunk_idx - q_compute_chunk_idx = next_q_compute_chunk_idx - - compute_stream.wait_stream(offload_stream) - compute_stream.synchronize() - - dk_seq_len = dk_accum.shape[1] - - if ctx.pos_emb_cos is not None: - dq_accum = apply_rotary_pos_emb_backward(dq[kv_compute_chunk_idx].get_gpu_chunk().to(dtype), - ctx.pos_emb_cos[:, dk_seq_len * i:dk_seq_len * (i + 1)], - ctx.pos_emb_sin[:, dk_seq_len * i:dk_seq_len * (i + 1)]) - dk_accum = apply_rotary_pos_emb_backward(dk_accum.to(dtype), - ctx.pos_emb_cos[:, dk_seq_len * i:dk_seq_len * (i + 1)], - ctx.pos_emb_sin[:, dk_seq_len * i:dk_seq_len * (i + 1)]) - else: - dq_accum = dq[kv_compute_chunk_idx].get_gpu_chunk().to(dtype) - dk_accum = dk_accum.to(dtype) - dv_accum = dv_accum.to(dtype) - - dq_accum = single_all_to_all(dq_accum.contiguous(), gather_idx, scatter_idx, 0, spg) - dk_accum = single_all_to_all(dk_accum.contiguous(), gather_idx, scatter_idx, 0, spg) - dv_accum = single_all_to_all(dv_accum.contiguous(), gather_idx, scatter_idx, 0, spg) - - general_offload_stream.synchronize() - compute_stream.wait_stream(general_offload_stream) - dist.barrier() - - with get_accelerator().stream(compute_stream): - input_chunk = layernorm_output[i].get_gpu_chunk().reshape(-1, layernorm_output[i].chunk_shape[-1]) - - dq_accum = dq_accum.flatten(2).permute(1, 0, 2) - dk_accum = dk_accum.flatten(2).permute(1, 0, 2) - dv_accum = dv_accum.flatten(2).permute(1, 0, 2) - - l, b = dk_accum.shape[0], dk_accum.shape[1] - - grad_qkv_linear_weight[:projection_size].add_( - torch.matmul(dq_accum.reshape(l * b, -1).t(), input_chunk)) - grad_qkv_linear_weight[projection_size:projection_size + kv_projection_size].add_( - torch.matmul(dk_accum.reshape(l * b, -1).t(), input_chunk)) - grad_qkv_linear_weight[projection_size + kv_projection_size:].add_( - torch.matmul(dv_accum.reshape(l * b, -1).t(), input_chunk)) - - grad_qkv_linear_bias[:projection_size].add_(dq_accum.sum(0).sum(0)) - grad_qkv_linear_bias[projection_size:projection_size + kv_projection_size].add_(dk_accum.sum(0).sum(0)) - grad_qkv_linear_bias[projection_size + kv_projection_size:].add_(dv_accum.sum(0).sum(0)) - - grad_layernorm_output[i].add_(torch.matmul(dq_accum, qkv_linear_weight[:projection_size])) - grad_layernorm_output[i].add_( - torch.matmul(dk_accum, qkv_linear_weight[projection_size:projection_size + kv_projection_size])) - grad_layernorm_output[i].add_( - torch.matmul(dv_accum, qkv_linear_weight[projection_size + kv_projection_size:])) - - del dq_accum, dk_accum, dv_accum - dk_accum = torch.zeros(global_k[i].chunk_shape, dtype=torch.float, device=device) - dv_accum = torch.zeros(global_v[i].chunk_shape, dtype=torch.float, device=device) - dq[kv_compute_chunk_idx].offload() - dq[kv_compute_chunk_idx] = None - - if i != (len(global_k) - 1): - next_kv_compute_chunk_idx = kv_compute_chunk_idx + 1 - with get_accelerator().stream(offload_stream): - global_k[next_kv_compute_chunk_idx].load_to_gpu() - global_v[next_kv_compute_chunk_idx].load_to_gpu() - - with get_accelerator().stream(general_offload_stream): - layernorm_output[next_kv_compute_chunk_idx].load_to_gpu() - - compute_stream.wait_stream(offload_stream) - compute_stream.synchronize() - - layernorm_output[kv_compute_chunk_idx].offload() - global_k[kv_compute_chunk_idx].offload() - global_v[kv_compute_chunk_idx].offload() - kv_compute_chunk_idx = next_kv_compute_chunk_idx - - return torch.cat( - grad_layernorm_output, - dim=0).to(dtype), None, None, None, None, None, None, None, None, None, None, grad_qkv_linear_weight.to( - dtype), grad_qkv_linear_bias.to(dtype), None, None, None - - -class FPDT_Attention(torch.nn.Module): - - def __init__(self, - config, - first_weight, - first_bias, - second_weight, - second_bias, - sequence_process_group, - gather_idx: int = 0, - scatter_idx: int = 2, - return_bias=True, - chunk_size=65536, - enable_offloading=True) -> None: - - super(FPDT_Attention, self).__init__() - if _flash_attn_forward is None or _flash_attn_backward is None: - raise ImportError( - "DeepSpeed FPDT requires flash-attn 2.6.3. Please install it with `pip install flash-attn --no-build-isolation`." - ) - - self.spg = sequence_process_group - self.scatter_idx = scatter_idx - self.gather_idx = gather_idx - self.config = config - - self.projection_size = config.kv_channels * config.num_attention_heads - self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads - self.kv_projection_size = config.kv_channels * config.num_key_value_heads - self.hidden_size = config.hidden_size - - self.qkv_linear_weight = first_weight - self.qkv_linear_bias = first_bias - self.qkv_dense_weight = second_weight - self.qkv_dense_bias = second_bias - - self.reture_bias = return_bias - self.dropout = config.attention_dropout - - self.chunk_size = chunk_size - self.double_buffer = enable_offloading - - def forward(self, - layernorm_output, - attention_mask, - inference_params, - rotary_pos_emb, - cpu_offloading=True) -> Tensor: - self.num_chunks_attn = layernorm_output.shape[0] * dist.get_world_size(self.spg) // self.chunk_size - - if not cpu_offloading or self.num_chunks_attn == 1: - output = _FPDTGPUAttentionImpl_.apply(layernorm_output, attention_mask, inference_params, rotary_pos_emb, - self.spg, self.scatter_idx, self.gather_idx, self.hidden_size, - self.projection_size, self.hidden_size_per_attention_head, - self.kv_projection_size, self.qkv_linear_weight, - self.qkv_linear_bias, self.dropout, self.num_chunks_attn, - cpu_offloading) - else: - output = _FPDTGPUOffloadingAttentionImpl_.apply( - layernorm_output, attention_mask, inference_params, rotary_pos_emb, self.spg, self.scatter_idx, - self.gather_idx, self.hidden_size, self.projection_size, self.hidden_size_per_attention_head, - self.kv_projection_size, self.qkv_linear_weight, self.qkv_linear_bias, self.dropout, - self.num_chunks_attn, cpu_offloading) - - output = output.flatten(2).permute(1, 0, 2).contiguous() - - output = torch.matmul(output, self.qkv_dense_weight.t()) - if not self.reture_bias: - output += self.qkv_dense_bias - return output, self.qkv_dense_bias if self.reture_bias else None - - -@torch.jit.script -def bias_gelu(x): - return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))) - - -@torch.jit.script -def bias_gelu_back(g, x): - tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) - # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out) - return ff * g - - -class FPDT_FFN(torch.autograd.Function): - generate_vmap_rule = False - - @staticmethod - def forward(ctx: Any, x, w1, b1, w2, b2, add_bias, chunk_size): - do_save = x.requires_grad - ctx.add_bias = add_bias - device = x.device - - with torch.no_grad(): - num_chunk = x.shape[0] // chunk_size - ctx.num_chunk = num_chunk - result = torch.empty(x.shape, device=device, dtype=x.dtype) - assert chunk_size * num_chunk == x.shape[0] - for i in range(num_chunk): - st = i * chunk_size - ed = st + chunk_size - x_ = torch.matmul(x[st:ed], w1.t()) + b1 - x_ = bias_gelu(x_) - if add_bias: - result[st:ed] = torch.matmul(x_, w2.t()) + b2 - else: - result[st:ed] = torch.matmul(x_, w2.t()) - - del x_ - - if do_save: - ctx.device = device - ctx.dtype = x.dtype - ctx.save_for_backward(x, w1, b1, w2, b2) - ctx.grad_x_shape = x.shape - return result.to(x.dtype), b2 if not add_bias else None - - @staticmethod - def backward(ctx, grad_output, grad_bias): - x, w1, b1, w2, b2 = ctx.saved_tensors - device = ctx.device - dtype = ctx.dtype - add_bias = ctx.add_bias - - num_chunk = ctx.num_chunk - chunk_size = x.shape[0] // num_chunk - assert chunk_size * num_chunk == grad_output.shape[0] - - grad_w2 = torch.zeros(w2.shape, device=device, dtype=torch.float) - grad_b2 = torch.zeros(b2.shape, device=device, dtype=torch.float) - grad_w1 = torch.zeros(w1.shape, device=device, dtype=torch.float) - grad_b1 = torch.zeros(b1.shape, device=device, dtype=torch.float) - - for i in range(num_chunk): - st = i * chunk_size - ed = st + chunk_size - x_chunk = x[st:ed] - - before_act = (torch.matmul(x_chunk, w1.t()) + b1) - before_act_2 = before_act**2 - tanh_out = torch.tanh(0.79788456 * before_act * (1 + 0.044715 * before_act_2)) - ff = 0.5 * before_act * ((1 - tanh_out * tanh_out) * - (0.79788456 + 0.1070322243 * before_act_2)) + 0.5 * (1 + tanh_out) - grad_w2.add_( - torch.matmul(grad_output[st:ed].reshape(-1, grad_output.shape[2]).t(), - (before_act * 0.5 * (1 + tanh_out)).reshape(-1, before_act.shape[2]))) - del before_act, before_act_2, tanh_out - - grad_inter = torch.matmul(grad_output[st:ed], w2) * ff - del ff - - grad_w1.add_(torch.matmul( - grad_inter.reshape(-1, grad_inter.shape[2]).t(), x_chunk.reshape(-1, x.shape[2]))) - grad_b1.add_(grad_inter.sum(0).sum(0)) - - x[st:ed].copy_(torch.matmul(grad_inter, w1)) - - del grad_inter - - if add_bias: - grad_b2.add_(grad_output[st:ed].sum(0).sum(0)) - - return x, grad_w1.to(dtype), grad_b1.to(dtype), grad_w2.to(dtype), grad_b2.to(dtype), None, None - - -class FPDT_LogitsLoss(torch.autograd.Function): - generate_vmap_rule = False - - @staticmethod - def forward(ctx: Any, lm_output, labels, logit_weights, rank, spg_size, spg, num_chunk): - labels = labels.t() - chunk_size = lm_output.shape[0] // num_chunk - assert chunk_size * num_chunk == lm_output.shape[0] - batch_size, local_seq_len = lm_output.shape[1], lm_output.shape[0] - loss = torch.empty((batch_size, local_seq_len), dtype=torch.float, device=lm_output.device) - - ctx.num_chunk = num_chunk - ctx.chunk_size = chunk_size - ctx.device = lm_output.device - ctx.dtype = lm_output.dtype - - ctx.rank = rank - ctx.local_seq_len = local_seq_len - with torch.no_grad(): - for i in range(num_chunk): - st = i * chunk_size - ed = st + chunk_size - logits_chunk = torch.matmul(lm_output[st:ed], logit_weights.t()).float() - - vocab_size = logits_chunk.size(2) - # nll - softmax = torch.nn.functional.softmax(logits_chunk, dim=-1) - loss_chunk = torch.nn.functional.nll_loss(softmax.log().reshape(-1, vocab_size).contiguous(), - labels[st:ed, :].reshape(-1).contiguous(), - reduction='none') - loss[:, st:ed] = loss_chunk.reshape(chunk_size, batch_size).t() - - del logits_chunk - ctx.save_for_backward(lm_output.to('cpu'), labels) - ctx.logit_weights = logit_weights - - seqlen = local_seq_len * spg_size - batch_size = loss.size(0) - loss = loss.t().contiguous() - loss_all = torch.empty(seqlen, batch_size, dtype=loss.dtype, device=loss.device).contiguous() - - dist.allgather_fn(loss_all, loss, group=spg) - - return loss_all - - @staticmethod - def backward(ctx, grad_output): - lm_output, labels = ctx.saved_tensors - logit_weights = ctx.logit_weights - device = ctx.device - dtype = ctx.dtype - num_chunk = ctx.num_chunk - chunk_size = ctx.chunk_size - - rank = ctx.rank - local_seq_len = ctx.local_seq_len - - grad_output = grad_output[rank * local_seq_len:(rank + 1) * local_seq_len] - grad_lm_output = [None for _ in range(num_chunk)] - grad_logit_weights = torch.zeros(logit_weights.shape, device=grad_output.device, dtype=torch.float) - for i in range(num_chunk): - st = i * chunk_size - ed = st + chunk_size - lm_output_chunk = lm_output[st:ed].to(device) - logits_chunk = torch.matmul(lm_output_chunk, logit_weights.t()).float() - - # nll - softmax = torch.nn.functional.softmax(logits_chunk, dim=-1) - vocab_size = logits_chunk.size(2) - - grad_input = softmax - grad_2d = grad_input.reshape(-1, vocab_size).contiguous() - arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=device) - - grad_2d[arange_1d, labels[st:ed, :].reshape(-1).contiguous()] -= 1 - grad_input.mul_(grad_output[:chunk_size, :].unsqueeze(dim=-1)) - grad_input = grad_input.to(dtype) - - grad_output = grad_output[chunk_size:].contiguous() - - grad_lm_output_chunk = torch.matmul(grad_input, logit_weights) - grad_lm_output[i] = grad_lm_output_chunk - - grad_logit_weights.add_( - torch.matmul( - grad_input.reshape(-1, grad_input.shape[2]).t(), - lm_output_chunk.reshape(-1, lm_output_chunk.shape[2]))) - - return torch.cat(grad_lm_output, dim=0).to(dtype), None, grad_logit_weights.to(dtype), None, None, None, None diff --git a/venv/lib/python3.10/site-packages/deepspeed/sequence/layer.py b/venv/lib/python3.10/site-packages/deepspeed/sequence/layer.py deleted file mode 100644 index ecbe0d94120e2312d82ce65bf6c12fcd5dbd48d9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/sequence/layer.py +++ /dev/null @@ -1,440 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -import torch - -from typing import Any, Tuple -from torch import Tensor -from torch.nn import Module - -from einops import rearrange - -import deepspeed.comm as dist -from deepspeed.accelerator import get_accelerator -from deepspeed.module_inject.tp_shard import get_shard_size_list, set_num_kv_heads, get_num_kv_heads -from deepspeed.utils import groups - - -def _generate_layout_params(scatter_idx, batch_dim_idx, seq_world_size, input): - """ - This function generates the parameters required for `permute` and `reshape` operations, - which are used to process data before and after `all2all` communication. - """ - if batch_dim_idx == 0: - if scatter_idx < 2: - bs, global_seq_len, num_local_head, head_dim = input.shape - pre_all2all_inp_shape = [bs, seq_world_size, global_seq_len // seq_world_size, num_local_head, head_dim] - pre_all2all_permute_idx = (1, 0, 2, 3, 4) - - post_all2all_permute_idx = (1, 2, 0, 3, 4) - post_all2all_res_shape = [bs, global_seq_len // seq_world_size, seq_world_size * num_local_head, head_dim] - else: - bs, local_seq_len, num_total_head, head_dim = input.shape - assert num_total_head % seq_world_size == 0, f"Number of heads ({num_total_head}) must be divisible by the sequence parallel size ({seq_world_size})!" - pre_all2all_inp_shape = [bs, local_seq_len, seq_world_size, num_total_head // seq_world_size, head_dim] - pre_all2all_permute_idx = (2, 0, 1, 3, 4) - - post_all2all_permute_idx = (1, 0, 2, 3, 4) - post_all2all_res_shape = [bs, seq_world_size * local_seq_len, num_total_head // seq_world_size, head_dim] - else: - if scatter_idx < 2: - global_seq_len, bs, num_local_head, head_dim = input.shape - pre_all2all_inp_shape = [seq_world_size, global_seq_len // seq_world_size, bs, num_local_head, head_dim] - pre_all2all_permute_idx = None - - post_all2all_permute_idx = (1, 2, 0, 3, 4) - post_all2all_res_shape = [bs, seq_world_size * global_seq_len, num_local_head // seq_world_size, head_dim] - else: - local_seq_len, bs, num_total_head, head_dim = input.shape - assert num_total_head % seq_world_size == 0, f"Number of heads ({num_total_head}) must be divisible by the sequence parallel size ({seq_world_size})!" - pre_all2all_inp_shape = [local_seq_len, bs, seq_world_size, num_total_head // seq_world_size, head_dim] - pre_all2all_permute_idx = (2, 0, 1, 3, 4) - post_all2all_permute_idx = None - post_all2all_res_shape = [local_seq_len * seq_world_size, bs, num_total_head // seq_world_size, head_dim] - - return pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape - - -def post_all2all(permute_idx, res_shape): - """ - Post-processing function for `all2all` communication. - """ - - def post_func(input): - if permute_idx is not None: - input = input.permute(permute_idx).contiguous() - output = input.reshape(res_shape).contiguous() - - return output - - return post_func - - -def pre_all2all_fun(permute_idx, inp_shape, input): - """ - Pre-processing function for `all2all` communication. - """ - input_t = input.reshape(inp_shape).contiguous() - if permute_idx is not None: - input_t = input_t.permute(permute_idx).contiguous() - return input_t - - -def _rotate_half(x): - """ - change sign so the last dimension becomes [-odd, +even] - """ - x = rearrange(x, '... (j d) -> ... j d', j=2) - x1, x2 = x.unbind(dim=-2) - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb(t, freqs_cos, freqs_sin): - """ - input tensor t is of shape [seq_length, ..., dim] - rotary positional embeding tensor freqs is of shape [seq_length, ..., dim] - check https://kexue.fm/archives/8265 for detailed formulas - """ - rot_dim = freqs_cos.shape[-1] - # ideally t_pass is empty so rotary pos embedding is applied to all tensor t - t, t_pass = t[..., :rot_dim], t[..., rot_dim:] - - # first part is cosine component - # second part is sine component, need to change signs with _rotate_half method - t = (t * freqs_cos) + (_rotate_half(t) * freqs_sin) - - res = t if t_pass.shape[-1] == 0 else torch.cat((t, t_pass), dim=-1) - return res - - -def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group): - seq_world_size = dist.get_world_size(group) - inp_shape = list(input.shape) - assert batch_dim_idx in [0, 1], "batch_dim_idx must be either 0 or 1" - - if not (scatter_idx < 2): - input_splits = get_shard_size_list(inp_shape[scatter_idx], seq_world_size) - input = input.transpose(0, scatter_idx).contiguous() - local_heads = input_splits[groups._get_sequence_parallel_rank()] - output_splits = [local_heads] * seq_world_size - - output_buffer_shape = [seq_world_size * local_heads] + list(input.shape[1:]) - output = torch.empty(output_buffer_shape, device=input.device, dtype=input.dtype) - dist.all_to_all_single(output,input,output_split_sizes=output_splits,\ - input_split_sizes=input_splits,group=group) - ###[seq_ws*local_heads, ...] to [seq_ws, local_heads, ...] - output = output.view(seq_world_size, local_heads, *output.shape[1:]) - ###[seq_ws,local_heads,b,seq_len,...] to [seq_ws,seq_len,b,local_heads,...] - - ### batch_dim_idx=0 [seq_ws,local_heads,seq_len,b,...] to [b, seq_ws, seq_len, local_heads ...] - ### batch_dim_idx=1 [seq_ws,local_heads,b,seq_len,...] to [seq_ws,seq_len,b,local_heads,...] - if batch_dim_idx == 0: - order = [3, 0, 2, 1] + list(range(4, len(output.shape))) - output = output.permute(order).contiguous() - ###[b, seq_ws*local_seq_len, local_heads,...] - output = output.view(output.shape[0], inp_shape[gather_idx] * seq_world_size, - *output.shape[3:]).contiguous() - elif batch_dim_idx == 1: - output = output.transpose(1, 3).contiguous() - ###[seq_ws*local_seq_len, b, local_heads,...] - output = output.view(inp_shape[gather_idx] * seq_world_size, *output.shape[2:]).contiguous() - else: - # The compatibility handling of 4D and 3D tensors, standardizing to 3D. - input = input.reshape(input.shape[0], input.shape[1], -1) - - if batch_dim_idx == 0: #b,s,h - input = input.permute(1, 2, 0).contiguous() #s,h,b - elif batch_dim_idx == 1: #s,b,h - input = input.transpose(1, 2).contiguous() #s,h,b - seq_len, h, batch_size = input.shape - num_local_heads_list = get_shard_size_list(get_num_kv_heads(), seq_world_size) - local_heads = num_local_heads_list[groups._get_sequence_parallel_rank()] - h_dim = h // local_heads - local_seq_len = seq_len // seq_world_size - - input = input.view(seq_len * h, batch_size) - local_seq_len_with_heads = int(input.shape[0] / seq_world_size) # dim size of local_seq_len*local_heads*hdim - input_splits = [local_seq_len_with_heads] * seq_world_size - coeff = local_seq_len_with_heads // local_heads #per head: dim size of local_seq_len*hdim - - #uneven seq_world_size coeff, total_heads/local_heads. - heads_scale_coeff = get_num_kv_heads() / local_heads - - output_splits = [num_local_heads * coeff for num_local_heads in num_local_heads_list] - output_buff_d1_size = int(heads_scale_coeff * local_seq_len_with_heads) - total_h = int(inp_shape[gather_idx] * heads_scale_coeff) - output = torch.empty(output_buff_d1_size, input.shape[1], device=input.device, dtype=input.dtype) - dist.all_to_all_single(output,input,output_split_sizes=output_splits, \ - input_split_sizes=input_splits,group=group) - ################## - #suppose 7 heads divide into 4 ranks [2,2,2,1] - #chunk_num_heads_small=floor(7/4)=1 - #chunk_num_heads_large=ceil(7/4)=2 - #num_chunk_heads_large=len([2,2,2])=3, all2all_buffer_counts - #num_chunk_heads_small=len([1])=1, all2all_buffer_counts - #total_num_large_heads=sum([2,2,2])=7 - #total_num_small_heads=sum([1])=1 - - chunk_num_heads_small = get_num_kv_heads() // seq_world_size # even heads compatible - chunk_num_heads_large = chunk_num_heads_small + 1 - num_chunk_heads_large = get_num_kv_heads() % seq_world_size - num_chunk_heads_small = seq_world_size - num_chunk_heads_large - total_num_large_heads = num_chunk_heads_large * chunk_num_heads_large - total_num_small_heads = num_chunk_heads_small * chunk_num_heads_small - - heads_large_combine_size = coeff * total_num_large_heads - heads_small_combine_size = coeff * total_num_small_heads - heads_large_chunk, heads_small_chunk = output.split([heads_large_combine_size, heads_small_combine_size], - dim=0) - heads_large_chunk = heads_large_chunk.view(num_chunk_heads_large, local_seq_len, chunk_num_heads_large, h_dim, - batch_size) - heads_small_chunk = heads_small_chunk.view(num_chunk_heads_small, local_seq_len, chunk_num_heads_small, h_dim, - batch_size) - if batch_dim_idx == 0: - #[all2all_buffer_counts, local_seq_len, n_heads,dim,batch]->[batch,local_seq_len,all2all_buffer_counts*n_heads,dim] - order = [4, 1, 0, 2, 3] - heads_large_chunk = heads_large_chunk.permute(order).contiguous().view(batch_size, local_seq_len, - total_num_large_heads, h_dim) - heads_small_chunk = heads_small_chunk.permute(order).contiguous().view(batch_size, local_seq_len, - total_num_small_heads, h_dim) - elif batch_dim_idx == 1: - #[all2all_buffer_counts, local_seq_len, n_heads,dim,batch]->[local_seq_len,batch,all2all_buffer_counts*n_heads,dim] - order = [1, 4, 0, 2, 3] - heads_large_chunk = heads_large_chunk.permute(order).contiguous().view(local_seq_len, batch_size, - total_num_large_heads, h_dim) - heads_small_chunk = heads_small_chunk.permute(order).contiguous().view(local_seq_len, batch_size, - total_num_small_heads, h_dim) - - output = torch.cat([heads_large_chunk, heads_small_chunk], dim=2).contiguous() - - inp_shape[scatter_idx] = inp_shape[scatter_idx] // seq_world_size - output_shape= inp_shape[: gather_idx] + \ - [total_h,] + \ - inp_shape[gather_idx + 1:] - - output = output.view(output_shape) - - return output - - -def single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, async_op=False, handle=None, type=None): - seq_world_size = dist.get_world_size(group) - # we only need num_heads once - num_heads = input.shape[2] - - if get_num_kv_heads() is not None or (num_heads % seq_world_size != 0 and not scatter_idx < 2): - # Assuming here that the number of heads for q is consistent with kv - # If not, additional logic is required for cases like GQA - if get_num_kv_heads() is None: - assert num_heads > seq_world_size, f"Number of heads ({num_heads}) must be larger than sequence parallel size ({seq_world_size})" - # set heads at first call by num_total_heads. - # then use ``get_num_kv_heads() is not None`` to re-entry uneven path. - set_num_kv_heads(num_heads) - assert async_op == False, "uneven head sp does not support async op" - return uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group) - - pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape = _generate_layout_params( - scatter_idx, batch_dim_idx, seq_world_size, input) - - input_t = pre_all2all_fun(pre_all2all_permute_idx, pre_all2all_inp_shape, input) - - post_all2all_fun = post_all2all(post_all2all_permute_idx, post_all2all_res_shape) - output = torch.empty_like(input_t) - work = dist.all_to_all_single(output, input_t, group=group, async_op=async_op) - - if async_op: - if type in ('dq', 'dk'): - handle[type + '_work'] = work - handle[type + '_grad'] = output - handle[type + '_post_all2all_func'] = post_all2all_fun - return output.view(post_all2all_res_shape) - - res = post_all2all_fun(output) - return res - - -class _DimZeroAllToAll(torch.autograd.Function): - """Differentiable All2All across dimension 0.""" - - @staticmethod - def forward(ctx: Any, group: dist.ProcessGroup, input: Tensor) -> Tensor: - world_size = dist.get_world_size(group) - assert input.shape[0] == world_size, f"Dim 0 {input.shape[0]} is not world size" - - ctx.group = group - - output = torch.empty_like(input).contiguous() - # torch.distributed.nn.functional.all_to_all_single(output, input.contiguous(), group=group) - dist.all_to_all_single(output, input.contiguous(), group=group) - return output - - @staticmethod - def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor]: - return (None, _DimZeroAllToAll.apply(ctx.group, *grad_output)) - - -class _SeqAllToAll(torch.autograd.Function): - - @staticmethod - def forward(ctx: Any, - group: dist.ProcessGroup, - input: Tensor, - scatter_idx: int, - gather_idx: int, - batch_dim_idx: int, - stream=None, - handle=None, - type=None, - is_fwd=True) -> Tensor: - ctx.group = group - ctx.scatter_idx = scatter_idx - ctx.gather_idx = gather_idx - ctx.stream = stream - ctx.handle = handle - ctx.type = type - ctx.batch_dim_idx = batch_dim_idx - if ctx.handle is None: - res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False) - - else: - # overlap communication path - if not is_fwd and type == 'o': - assert ctx.stream != None - res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False) - get_accelerator().current_stream().wait_stream(ctx.stream) - # The computation of d o_weight can overlap with the communication of d o_input - - elif not is_fwd and type in ('q', 'k'): - # Achieve communication overlap by pipelining the matrix computation and communication of dq, dk, and dv - type = 'd' + type - res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, True, handle, type) - - elif is_fwd and type in ('q', 'k'): - # Achieve communication overlap by pipelining the matrix computation and communication of q, k, and v - type = 'fwd_' + type - res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False, handle, type) - - else: - res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False) - - return res - - @staticmethod - def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]: - - return (None, - _SeqAllToAll.apply(ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx, ctx.batch_dim_idx, - ctx.stream, ctx.handle, ctx.type, False), None, None, None, None, None, None, None) - - -class DistributedAttention(torch.nn.Module): - """Initialization. - - Arguments: - local_attention (Module): local attention with q,k,v - sequence_process_group (ProcessGroup): sequence parallel process group - scatter_idx (int): scatter_idx for all2all comm - gather_idx (int): gather_idx for all2all comm - """ - - def __init__( - self, - local_attention: Module, - sequence_process_group: dist.ProcessGroup, - scatter_idx: int = 2, - gather_idx: int = 0, - sp_stream=None, - ) -> None: - - super(DistributedAttention, self).__init__() - self.local_attn = local_attention - self.spg = sequence_process_group - self.scatter_idx = scatter_idx - self.gather_idx = gather_idx - self.sp_overlap_comm = False - self.overlap_handles = None - self.sp_stream = sp_stream - if sp_stream is not None: - self.overlap_handles = {} - self.sp_overlap_comm = True - self.default_stream = get_accelerator().default_stream() - - def layer_sync(self, layer): - if self.sp_overlap_comm and hasattr(layer, 'done_event'): - self.default_stream.wait_event(layer.done_event) - - def forward(self, - query: Tensor, - key: Tensor, - value: Tensor, - batch_dim_idx: int, - rotary_pos_emb=None, - *args: Any, - **kwargs) -> Tensor: - """ forward - - Arguments: - query (Tensor): query input to the layer - key (Tensor): key input to the layer - value (Tensor): value input to the layer - batch_dim_idx (int): indicating which dim is batch - args: other args - - Returns: - * output (Tensor): context output - """ - - # TODO Merge three alltoall calls into one - # TODO (Reza): change the api on the megatron-deepspeed side so that we only receive all data (q,k, and v) together! - #in shape : e.g., [s/p:h:] - - def bwd_hook(layer_type): - - def pre_hook_fun(grad): - type = 'd' + layer_type - self.overlap_handles[type + '_work'].wait() - self.sp_stream.wait_stream(self.default_stream) - all2all_output = self.overlap_handles[type + '_grad'] - grad = list(grad) - grad[0] = self.overlap_handles[type + '_post_all2all_func'](all2all_output) - grad = tuple(grad) - - return pre_hook_fun - - self.layer_sync(query) - query_layer = _SeqAllToAll.apply(self.spg, query, self.scatter_idx, self.gather_idx, batch_dim_idx, None, - self.overlap_handles, 'q') - self.layer_sync(key) - key_layer = _SeqAllToAll.apply(self.spg, key, self.scatter_idx, self.gather_idx, batch_dim_idx, None, - self.overlap_handles, 'k') - if self.sp_overlap_comm: - self.default_stream.wait_stream(self.sp_stream) - - value_layer = _SeqAllToAll.apply(self.spg, value, self.scatter_idx, self.gather_idx, batch_dim_idx, None, - self.overlap_handles, 'v') - - if self.sp_overlap_comm: - # Register a hook to synchronize dq and dk after the all-to-all - # operation when the gradient data is used. - # Place this logic after the q, k, v all-to-all operation to - # improve interpreter speed to - # call and launch of the forward all-to-all communication. - grad_fn_q = query.grad_fn.next_functions[0][0] - grad_fn_q.register_prehook(bwd_hook(layer_type='q')) - grad_fn_k = key.grad_fn.next_functions[0][0] - grad_fn_k.register_prehook(bwd_hook(layer_type='k')) - - #out shape : e.g., [s:h/p:] - if rotary_pos_emb is not None: - pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3) - query_layer = apply_rotary_pos_emb(query_layer, pos_emb_cos, pos_emb_sin) - key_layer = apply_rotary_pos_emb(key_layer, pos_emb_cos, pos_emb_sin) - - context_layer = self.local_attn(query_layer, key_layer, value_layer, *args, **kwargs) - - output = _SeqAllToAll.apply(self.spg, context_layer, self.gather_idx, self.scatter_idx, batch_dim_idx, - self.sp_stream, self.overlap_handles, 'o') - - #out e.g., [s/p::h] - return output diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__init__.py b/venv/lib/python3.10/site-packages/deepspeed/utils/__init__.py deleted file mode 100644 index 5d902a883f6eeee72b3c9ccbf5a11abf22de8e33..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from .logging import logger, log_dist -from .comms_logging import get_caller_func -#from .distributed import init_distributed -from .init_on_device import OnDevice -from .groups import * -from .nvtx import instrument_w_nvtx -# TODO: Move tensor fragment and mixed precision to zero utils -from .tensor_fragment import tensor_fragment, get_full_hp_param, get_hp_fragment_mapping, fragment_address, get_full_hp_grad, map_to_flat_opt_states -from .tensor_fragment import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state -from .tensor_fragment import set_full_hp_param, set_full_hp_grad -from .tensor_fragment import safe_set_full_fp32_param, safe_set_full_optimizer_state, safe_set_full_grad -from .tensor_fragment import safe_get_local_fp32_param, safe_get_local_grad, safe_get_local_optimizer_state -from .tensor_fragment import safe_set_local_fp32_param, safe_set_local_grad, safe_set_local_optimizer_state -from .tensor_fragment import safe_update_full_grad_vectorized -from .z3_leaf_module import set_z3_leaf_modules, unset_z3_leaf_modules, get_z3_leaf_modules, z3_leaf_module, z3_leaf_parameter, set_z3_leaf_module -from .mixed_precision_linkage import link_hp_params, lazy_init_hp_params_optimizer_state -from deepspeed.runtime.dataloader import RepeatingLoader -from .numa import get_numactl_cmd diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index e97c4b73da902815fce75bf1af723d1786f36968..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/bwc.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/bwc.cpython-310.pyc deleted file mode 100644 index c0151367e892bb79f161a83cfdf2f6d5a4e9b00e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/bwc.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/comms_logging.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/comms_logging.cpython-310.pyc deleted file mode 100644 index a43e83619e6a7985cce48f87c38affb1a1a42ee3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/comms_logging.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/config.cpython-310.pyc deleted file mode 100644 index 54a36185f40095d6c55fa9d233a798ba225525cd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/debug.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/debug.cpython-310.pyc deleted file mode 100644 index a436ccc447f65cea7fa32e1eefbfa99da08bd058..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/debug.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/exceptions.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/exceptions.cpython-310.pyc deleted file mode 100644 index e599018a9a209d40c28ef7d6c3085cc83cc4bed4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/exceptions.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/groups.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/groups.cpython-310.pyc deleted file mode 100644 index 39a06e80514b1f52db9446477eb72478aaaf88fd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/groups.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/init_on_device.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/init_on_device.cpython-310.pyc deleted file mode 100644 index 0bceb451056c071a9e541771a745f40352f633e3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/init_on_device.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/logging.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/logging.cpython-310.pyc deleted file mode 100644 index 529d98f85d01f99c4f8efa5286dcc8316146e999..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/logging.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/mixed_precision_linkage.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/mixed_precision_linkage.cpython-310.pyc deleted file mode 100644 index 54b800910fe0721cdd2446789d49f3942960c5b9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/mixed_precision_linkage.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/numa.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/numa.cpython-310.pyc deleted file mode 100644 index bf6b7f29152b1aafa5da25845daae72704ac0cb1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/numa.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/nvtx.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/nvtx.cpython-310.pyc deleted file mode 100644 index 5de2f1e3aedc1a6c5a7eff2cc868ad421af7673d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/nvtx.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/tensor_fragment.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/tensor_fragment.cpython-310.pyc deleted file mode 100644 index 35e3ef11987b0ca83954f66b9797cc1e7f2b0284..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/tensor_fragment.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/timer.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/timer.cpython-310.pyc deleted file mode 100644 index 6f49ca5d1010f4aac90cf09db5dfb7042745a49a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/timer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/torch.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/torch.cpython-310.pyc deleted file mode 100644 index ff93fdfd757b4260b69a598f715149ddbc558bbb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/torch.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/types.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/types.cpython-310.pyc deleted file mode 100644 index 150a5bb6412542819eded38ba7553c539245deb9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/types.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/z3_leaf_module.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/z3_leaf_module.cpython-310.pyc deleted file mode 100644 index f98d73b3386d459cbba4536ad17210613fb235a9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/z3_leaf_module.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/zero_to_fp32.cpython-310.pyc b/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/zero_to_fp32.cpython-310.pyc deleted file mode 100644 index 2471b769c3180d410e98b9d6c45ba30509fb4016..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deepspeed/utils/__pycache__/zero_to_fp32.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/bwc.py b/venv/lib/python3.10/site-packages/deepspeed/utils/bwc.py deleted file mode 100644 index 69fcc251a68429a768892b7cf7c6824de7cb5ba2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/bwc.py +++ /dev/null @@ -1,104 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - - -def bwc_tensor_model_parallel_rank(mpu=None): - """Backwards-compatible way of querying the tensor model parallel rank from - an ``mpu`` object. - - *Tensor* model parallelism means that tensors are physically split across - processes. This contrasts with *pipeline* model parallelism, in which the - layers are partitioned but tensors left intact. - - The API for tensor model parallelism has changed across versions and this - helper provides a best-effort implementation across versions of ``mpu`` - objects. The preferred mechanism is - ``mpu.get_tensor_model_parallel_rank()``. - - This should "just work" with both Megatron-LM and DeepSpeed's pipeline - parallelism. - - Args: - mpu (model parallel unit, optional): The tensor model parallel rank. - If ``mpu=None``, returns 0. Defaults to ``None``. - - Returns: - int: the rank - """ - if mpu is None: - # No model parallelism in easy :) - return 0 - - if hasattr(mpu, 'get_tensor_model_parallel_rank'): - # New Megatron and DeepSpeed convention (post pipeline-parallelism release) - return mpu.get_tensor_model_parallel_rank() - elif hasattr(mpu, 'get_slice_parallel_rank'): - # Some DeepSpeed + pipeline parallelism versions - return mpu.get_slice_parallel_rank() - else: - # Deprecated Megatron and DeepSpeed convention - return mpu.get_model_parallel_rank() - - -def bwc_tensor_model_parallel_world_size(mpu=None): - """Backwards-compatible way of querying the tensor model parallel world size. - Similar to bwc_tensor_model_parallel_rank. - """ - if mpu is None: - return 1 - - if hasattr(mpu, 'get_tensor_model_parallel_world_size'): - # New Megatron and DeepSpeed convention (post pipeline-parallelism release) - return mpu.get_tensor_model_parallel_world_size() - elif hasattr(mpu, 'get_slice_parallel_world_size'): - # Some DeepSpeed + pipeline parallelism versions - return mpu.get_slice_parallel_world_size() - else: - # Deprecated Megatron and DeepSpeed convention - return mpu.get_model_parallel_world_size() - - -def bwc_tensor_model_parallel_group(mpu=None): - """Backwards-compatible way of querying the tensor model parallel group. - Similar to bwc_tensor_model_parallel_rank. - """ - if mpu is None: - return None - - if hasattr(mpu, 'get_tensor_model_parallel_group'): - # New Megatron and DeepSpeed convention (post pipeline-parallelism release) - return mpu.get_tensor_model_parallel_group() - elif hasattr(mpu, 'get_slice_parallel_group'): - # Some DeepSpeed + pipeline parallelism versions - return mpu.get_slice_parallel_group() - else: - # Deprecated Megatron and DeepSpeed convention - return mpu.get_model_parallel_group() - - -def bwc_pipeline_parallel_world_size(mpu=None): - """Backwards-compatible way of querying the pipeline parallel world size.""" - world_size = 1 - if mpu is not None: - if hasattr(mpu, 'get_pipeline_model_parallel_world_size'): - # New Megatron and DeepSpeed convention (post pipeline-parallelism release) - world_size = mpu.get_pipeline_model_parallel_world_size() - elif hasattr(mpu, 'get_pipe_parallel_world_size'): - # DeepSpeed Topology - world_size = mpu.get_pipe_parallel_world_size() - return world_size - - -def bwc_pipeline_parallel_group(mpu=None): - """Backwards-compatible way of querying the pipeline parallel group.""" - if mpu is None: - return None - if hasattr(mpu, 'get_pipeline_model_parallel_group'): - # Megatron - return mpu.get_pipeline_model_parallel_group() - elif hasattr(mpu, 'get_pipe_parallel_group'): - # DeepSpeed Topology - return mpu.get_pipe_parallel_group() - assert False, 'mpu does not support pipeline parallel group' diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/comms_logging.py b/venv/lib/python3.10/site-packages/deepspeed/utils/comms_logging.py deleted file mode 100644 index abe0c9d47cd07e124d95d35736ae64f2bc2d85fa..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/comms_logging.py +++ /dev/null @@ -1,181 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import math -from deepspeed.utils import log_dist - - -def get_caller_func(frame=3): - import sys - return sys._getframe(frame).f_code.co_name - - -def print_rank_0(message): - import deepspeed.comm as dist - if dist.get_rank() == 0: - print(message) - - -# Helper function to pretty-print message sizes -def convert_size(size_bytes): - if size_bytes == 0: - return "0B" - size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB") - i = int(math.floor(math.log(size_bytes, 1024))) - p = math.pow(1024, i) - s = round(size_bytes / p, 2) - return "%s %s" % (s, size_name[i]) - - -# Helper function to calculate algbw and busbw. -# See https://gist.github.com/jeffra/b5e80466b4c86be00ea3b6f130fb7a36 and https://github.com/NVIDIA/nccl-tests/blob/master/doc/PERFORMANCE.md -def calc_bw_log(comm_op, size, duration): - import deepspeed.comm as dist - - n = dist.get_world_size() - tput = 0 - busbw = 0 - if comm_op == "all_to_all_single": - tput = (size / duration) - busbw = (size / duration) * ((n - 1) / n) - elif comm_op == "all_gather" or comm_op == "all_gather_into_tensor" or comm_op == "reduce_scatter" or comm_op == "reduce_scatter_tensor": - size *= n - tput = (size / duration) - busbw = (size / duration) * ((n - 1) / n) - elif comm_op == "all_reduce" or comm_op == "all_reduce_coalesced" or comm_op == "inference_all_reduce": - tput = (size * 2 / duration) - busbw = (size / duration) * (2 * (n - 1) / n) - elif comm_op == "send" or comm_op == "recv" or comm_op == "isend" or comm_op == "irecv" or comm_op == "broadcast" or comm_op == "reduce" or comm_op == "gather" or comm_op == "scatter" or comm_op == "barrier": - tput = (size / duration) - busbw = tput - else: - print_rank_0("wrong comm_op specified") # noqa: F821 - exit(0) - - # convert to Gbps - tput *= 8 - busbw *= 8 - - tput /= 1e6 - busbw /= 1e6 - - return tput, busbw - - -class CommsLogger: - - def __init__(self): - from deepspeed.comm.constants import COMMS_LOGGER_VERBOSE_DEFAULT, COMMS_LOGGER_DEBUG_DEFAULT, COMMS_LOGGER_PROF_OPS_DEFAULT, COMMS_LOGGER_PROF_ALL_DEFAULT, COMMS_LOGGER_ENABLED_DEFAULT - self.comms_dict = {} - self.verbose = COMMS_LOGGER_VERBOSE_DEFAULT - self.debug = COMMS_LOGGER_DEBUG_DEFAULT - self.prof_ops = COMMS_LOGGER_PROF_OPS_DEFAULT - self.prof_all = COMMS_LOGGER_PROF_ALL_DEFAULT - self.enabled = COMMS_LOGGER_ENABLED_DEFAULT - - def configure(self, comms_config): - self.enabled = comms_config.comms_logger_enabled - if self.enabled: - self.verbose = comms_config.comms_logger.verbose - self.debug = comms_config.comms_logger.debug - self.prof_ops = comms_config.comms_logger.prof_ops - self.prof_all = comms_config.comms_logger.prof_all - - # There are three settings for the op profiler: - # - Global profiling (profile all comms) - # - Op-type profiling (e.g. profile all all_reduce comms) - # - Op profiling (e.g. profile a specific all_reduce op) - def start_profiling_comms(self): - self.prof_all = True - - def stop_profiling_comms(self): - self.prof_all = True - - # E.g. start_profiling_op('all_reduce') - def start_profiling_op(self, op_name_list): - self.prof_ops = list(set(self.prof_ops) | set(op_name_list)) - - def stop_profiling_op(self, op_name_list): - self.prof_ops = [op for op in self.prof_ops if op not in op_name_list] - - # Add log entry - def append(self, raw_name, record_name, latency, msg_size): - algbw, busbw = calc_bw_log(raw_name, msg_size, latency) - if record_name in self.comms_dict.keys(): - # If this comm_op has already been logged with this message size, just add to existing record - if msg_size in self.comms_dict[record_name].keys(): - self.comms_dict[record_name][msg_size][0] += 1 - self.comms_dict[record_name][msg_size][1].append(latency) - self.comms_dict[record_name][msg_size][2].append(algbw) - self.comms_dict[record_name][msg_size][3].append(busbw) - # If this is a new message size for this comm_op, add new record under existing comm_op - else: - self.comms_dict[record_name][msg_size] = [1, [latency], [algbw], [busbw]] - else: - # Create entirely new record - self.comms_dict[record_name] = {msg_size: [1, [latency], [algbw], [busbw]]} - # If verbose, print every comm op - # TODO: Add to tensorboard - if self.verbose: - log_str = f"comm op: {record_name} | time (ms): {latency:.2f} | msg size: {convert_size(msg_size)} | algbw (Gbps): {algbw:.2f} | busbw (Gbps): {busbw:.2f}" - log_dist(log_str, [0]) - - # Print summary at end of iteration, epoch, or training - def log_all(self, print_log=True, show_straggler=False): - import torch - from deepspeed.utils.timer import trim_mean - import deepspeed.comm as dist - from deepspeed.comm.reduce_op import ReduceOp - from deepspeed.accelerator import get_accelerator - - if print_log: - print( - f"{'Comm. Op': <20}{'Message Size': <20}{'Count': <20}{'Total Latency(ms)': <20}{'Avg Latency(ms)': <20}{'tput_avg (Gbps)': <20}{'busbw_avg (Gbps)': <20}" - ) - for record_name in self.comms_dict.keys(): - if print_log: - print(record_name) - for msg_size, vals in sorted(self.comms_dict[record_name].items()): - # vals[0] is the count for each msg size - count = vals[0] - # vals[1] is a list of latency records for each msg size - total_lat = sum(vals[1]) - # vals[2] and vals[3] are the lists of algbw and busbw, respectively - # Get rid of outliers when we print - avg_lat = trim_mean(vals[1], 0.1) - avg_algbw = trim_mean(vals[2], 0.1) - avg_busbw = trim_mean(vals[3], 0.1) - if print_log: - print( - f"{' ': <20}{convert_size(msg_size): <20}{count: <20}{total_lat: <20.2f}{avg_lat: <20.2f}{avg_algbw: <20.2f}{avg_busbw: <20.2f}" - ) - - if show_straggler: - if print_log: - print("_______________________________") - print("Breakdown with straggler effect") - print("-------------------------------") - print( - f"{'Comm. Op': <20}{'Message Size': <20}{'Count': <20}{'Total comm lat(ms)': <20}{'Total straggler(ms)': <20}{'Avg comm lat(ms)': <20}{'Avg straggler(ms)': <20}" - ) - device = get_accelerator().current_device_name() - for record_name in self.comms_dict.keys(): - if print_log: - print(record_name) - for msg_size, vals in sorted(self.comms_dict[record_name].items()): - # vals[0] is the count for each msg size - count = vals[0] - # vals[1] is a list of latency records for each msg size - lats = torch.tensor(vals[1], device=device) - min_lats = torch.tensor(vals[1], device=device) - dist.all_reduce(min_lats, op=ReduceOp.MIN) - total_lat = min_lats.sum().item() - total_straggler = (lats - min_lats).sum().item() - avg_lat = trim_mean(min_lats.tolist(), 0.1) - avg_straggler = trim_mean((lats - min_lats).tolist(), 0.1) - if print_log: - print( - f"{' ': <20}{convert_size(msg_size): <20}{count: <20}{total_lat: <20.2f}{total_straggler: <20.2f}{avg_lat: <20.2f}{avg_straggler: <20.2f}" - ) diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/config.py b/venv/lib/python3.10/site-packages/deepspeed/utils/config.py deleted file mode 100644 index 15f37ca7d8745cccee562d08b0eae3f10c5d9877..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/config.py +++ /dev/null @@ -1,46 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.runtime.config_utils import DeepSpeedConfigModel - -######################################### -# Timers -######################################### -# Timers. By default, timers are enabled. -# Users can configure in ds_config.json as below example: -TIMERS_FORMAT = ''' -Timers should be enabled as: -"timers": { - "throughput": { - "enabled": true, - "synchronized": true - } -} -''' - -TIMERS = "timers" -TIMERS_THROUGHPUT = "throughput" - - -def get_timers_config(param_dict): - if param_dict and TIMERS in param_dict and TIMERS_THROUGHPUT in param_dict[TIMERS]: - timers_config_dict = param_dict[TIMERS][TIMERS_THROUGHPUT] - else: - timers_config_dict = {} - return DeepSpeedThroughputTimerConfig(**timers_config_dict) - - -class DeepSpeedThroughputTimerConfig(DeepSpeedConfigModel): - """ Configure throughput timers """ - - enabled: bool = True - """ Turn on/off throughput timers """ - - synchronized: bool = True - """ Whether to synchronize a device when measuring the time. - Synchronizing a device is required to produce the most accurate timer measurements. - However, this comes at the expense of performance degradation. The CPU timer provides - sufficient accuracy in many cases. - """ diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/debug.py b/venv/lib/python3.10/site-packages/deepspeed/utils/debug.py deleted file mode 100644 index 6174ba45137e1e70dd004593a69ff41015c0212c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/debug.py +++ /dev/null @@ -1,173 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import deepspeed.comm as dist - -# For lazy import with printflock() -fcntl = None - -# for debug purposes map module and param objects to their fully qualified names -module_names = {} -param_names = {} - - -def debug_clear_module_and_param_names(): - global module_names - global param_names - module_names = {} - param_names = {} - - -def debug_extract_module_and_param_names(model): - # extract the fully qualified names as soon as the model is acquired - global module_names - global param_names - # XXX: can probably make a map of param2module and vice-versa - module_names = {module: name for name, module in model.named_modules()} - param_names = {param: name for name, param in model.named_parameters()} - - -def debug_module2name(module): - if module in module_names: - return module_names[module] - else: - return "unknown" - - -def debug_module2name_id(module): - return f"name={debug_module2name(module)}" - - -def debug_module2name_class(module): - return f"name={debug_module2name(module)} {module.__class__.__name__}" - - -def debug_param2name(param): - if param in param_names: - return param_names[param] - else: - return "unknown" - - -def debug_param2name_id(param): - return f"name={debug_param2name(param)} id={param.ds_id}" - - -def debug_param2name_id_shape(param): - return f"name={debug_param2name(param)} id={param.ds_id} shape={param.ds_shape}" - - -def debug_param2name_id_shape_device(param): - return f"name={debug_param2name(param)} id={param.ds_id} shape={param.ds_shape} device={param.device}" - - -def debug_param2name_id_numel(param): - return f"name={debug_param2name(param)} id={param.ds_id} numel={param.numel()}" - - -def debug_param2name_id_shape_status(param): - return f"name={debug_param2name(param)} id={param.ds_id} shape={param.ds_shape} status={param.ds_status}" - - -def printflock(*msgs): - """ - - For printing messages for all concurrent gpus w/o getting interleaved text. - - This is useful when debugging issues where multi-gpus don't sync. - - 1. Enable the force debug in say partitioning and zero3 files - 2. Override the usual versions with :: - - def print_rank_0(message, debug=False, force=False): - rank = deepspeed.comm.get_rank() - printflock(f"[{rank}] {message}") - 3. run the program and you get both logs non-interleaved - - But this makes it very difficult to make sense of the output, so the ``log_rank_file`` helper - function might be more useful, as it's easier to send each log stream into a separate file and - then compare those. - - """ - global fcntl - if fcntl is None: - import fcntl - - with open(__file__, "r") as fh: - fcntl.flock(fh, fcntl.LOCK_EX) - try: - print(*msgs) - finally: - fcntl.flock(fh, fcntl.LOCK_UN) - - -fh = None - - -def log_rank_file(rank, *msgs): - """ - Print to a log file of the given rank - - This is useful for debugging hanging in sync processes. Here is a possible workflow: - - 1. Enable the force debug in say partitioning and zero3 files - 2. Override the usual versions of print_rank_0 in those files with :: - - def print_rank_0(message, debug=False, force=False): - rank = deepspeed.comm.get_rank() - log_rank_file(rank, message) - - 3. run the program - 4. fix up the expected differences, e.g. different cuda numbers :: - - perl -pi -e 's|cuda:1|cuda:0|' log_rank_* - - 5. now diff and see where names and ids diverge - you will find where the gpus don't do the same - work (e.g. when some layers get conditionally skipped on one gpu but not all) - - diff -u log_rank_0.txt log_rank_1.txt | less - - """ - global fh - if fh is None: - fh = open(f"log_rank_{rank}.txt", "w") - for m in msgs: - fh.write(f"{m}\n") - fh.flush() - - -def print_backward_tensors(tensor): - - def _print_bwd_tensors(grad_fn): - print(f"Backward tensors in {grad_fn}") - for funcs in grad_fn.next_functions: - if funcs[0]: - try: - tensor = getattr(funcs[0], 'variable') - print(funcs[0]) - print(f"Tensor - id: {id(tensor)}, shape: {tensor.shape}, data: {tensor}, grad: {tensor.grad}") - except AttributeError as e: - _print_bwd_tensors(funcs[0]) - - if hasattr(tensor, 'grad_fn'): - _print_bwd_tensors(tensor.grad_fn) - - -def print_rank(*msg, force=False): - """print something on all global ranks with [rank] prefix. - """ - if not force: - return - global_rank = dist.get_rank() - print(f"[{global_rank}]", *msg) - - -def print_rank0(*msg, force=False): - """print something only on rank 0""" - if not force: - return - global_rank = dist.get_rank() - if global_rank == 0: - print(f"[{global_rank}]", *msg) diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/exceptions.py b/venv/lib/python3.10/site-packages/deepspeed/utils/exceptions.py deleted file mode 100644 index 9dba0933cc9e236ec150bfd5439f7ce8e157394a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/exceptions.py +++ /dev/null @@ -1,8 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - - -class DeprecatedException(Exception): - pass diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/groups.py b/venv/lib/python3.10/site-packages/deepspeed/utils/groups.py deleted file mode 100644 index 250aa06b1a7c0610c705313dd1a06c28672f7ef6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/groups.py +++ /dev/null @@ -1,749 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# The file has been adapted from https://github.com/NVIDIA/Megatron-LM and retains the following license from the original file - -# Copyright (c) 2019, NVIDIA CORPORATION. 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. -""" - Support different forms of parallelism in DeepSpeed using multiple process groups. - Given that there are multiple scenarios and use-cases, this file is going to be updated - frequently. For now, the group creation needed for the training scenario is being implemented. - For inference and other new scenarios, the code will be either reused or added to this file. -""" - -from deepspeed import comm as dist -from deepspeed.utils import log_dist -from deepspeed.utils.bwc import bwc_tensor_model_parallel_world_size, bwc_pipeline_parallel_world_size -from deepspeed.utils.exceptions import DeprecatedException -from deepspeed.accelerator import get_accelerator - -# Expert parallel group that the current rank belongs to. -_EXPERT_PARALLEL_GROUP = {} -# Mapping of expert parallel group to ranks -_EXPERT_PARALLEL_GROUP_RANKS = {} -# Expert data parallel group that the current rank belongs to. -_EXPERT_DATA_PARALLEL_GROUP = {} -# Mapping of expert data parallel group to ranks -_EXPERT_DATA_PARALLEL_GROUP_RANKS = {} -# dist world group needs to be cloned for some cases -_WORLD_GROUP = None -# ZeRO parameter partitioning group that the current rank belongs to. -_ZERO_PARAM_INTRA_PARALLEL_GROUP = None -# global object to maintain mpu object if passed by a Megatron client -mpu = None -# global object that stores tensor parallel world size for experts -expert_tensor_parallel_world_size = 1 -# All to All quantized graident communication groups -_ALL_TO_ALL_GROUP = {} - -mesh_device = None - - -# Deprecated groups initialize function. -def initialize(ep_size=1, mpu=None): - """ Deprecated function. Retained to inform the users.""" - raise DeprecatedException( - "Please do not use the groups.initialize() API as it is deprecated. Instead, pass the desired ep_size to deepspeed.moe.layer.MoE(..,ep_size,..)" - ) - - -def _ensure_divisibility(numerator, denominator): - """Ensure that numerator is divisible by the denominator.""" - assert numerator % denominator == 0, '{} is not divisible by {}'.format(numerator, denominator) - - -# ======== Start: Tensor Parallel Group Attributes ======== - -# Intra-layer model parallel group that the current rank belongs to. -_TENSOR_MODEL_PARALLEL_GROUP = None - -# Model parallel group (both intra- and pipeline) that the current rank belongs to. -_MODEL_PARALLEL_GROUP = None -# Data parallel group that the current rank belongs to. -_DATA_PARALLEL_GROUP = None - -# These values enable us to change the mpu sizes on the fly. -_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = None -_MPU_TENSOR_MODEL_PARALLEL_RANK = None - - -def _init_tp_mesh_device(tensor_model_parallel_size=1, data_parallel_size=None): - """Initialize model data parallel groups.""" - - global _DATA_PARALLEL_GROUP - global _MODEL_PARALLEL_GROUP - global _TENSOR_MODEL_PARALLEL_GROUP - - if _TENSOR_MODEL_PARALLEL_GROUP is not None: - return - - if data_parallel_size is None: - data_parallel_size = dist.get_world_size() // tensor_model_parallel_size - - mesh_device = dist.initialize_mesh_device((data_parallel_size, tensor_model_parallel_size), - ("data_parallel", "tensor_parallel")) - _TENSOR_MODEL_PARALLEL_GROUP = mesh_device.get_group(mesh_dim="tensor_parallel") - _DATA_PARALLEL_GROUP = mesh_device.get_group(mesh_dim="data_parallel") - - # They are always equal only in 2D (DP + TP) parallelism. - # _MODEL_PARALLEL_GROUP is assigned the same value as _TENSOR_MODEL_PARALLEL_GROUP - # to allow for future potential changes. - _MODEL_PARALLEL_GROUP = _TENSOR_MODEL_PARALLEL_GROUP - - return _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP - - -def get_tensor_model_parallel_group(): - """Get the tensor model parallel group the caller rank belongs to.""" - - assert _TENSOR_MODEL_PARALLEL_GROUP is not None, \ - 'intra_layer_model parallel group is not initialized' - return _TENSOR_MODEL_PARALLEL_GROUP - - -def get_model_parallel_group(): - """Get the model parallel group the caller rank belongs to.""" - - assert _MODEL_PARALLEL_GROUP is not None, \ - 'model parallel group is not initialized' - return _MODEL_PARALLEL_GROUP - - -def get_data_parallel_group(): - """Get the data parallel group the caller rank belongs to.""" - assert _DATA_PARALLEL_GROUP is not None, \ - 'data parallel group is not initialized' - return _DATA_PARALLEL_GROUP - - -def set_tensor_model_parallel_world_size(world_size): - """Set the tensor model parallel size""" - global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE - _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = world_size - - -def get_tensor_model_parallel_world_size(): - """Return world size for the tensor model parallel group.""" - global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE - if _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE is not None: - return _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE - return dist.get_world_size(group=get_tensor_model_parallel_group()) - - -def get_model_parallel_world_size(): - return get_tensor_model_parallel_world_size() - - -def set_tensor_model_parallel_rank(rank): - """Set tensor model parallel rank.""" - global _MPU_TENSOR_MODEL_PARALLEL_RANK - _MPU_TENSOR_MODEL_PARALLEL_RANK = rank - - -def get_tensor_model_parallel_rank(): - """Return my rank for the tensor model parallel group.""" - global _MPU_TENSOR_MODEL_PARALLEL_RANK - if _MPU_TENSOR_MODEL_PARALLEL_RANK is not None: - return _MPU_TENSOR_MODEL_PARALLEL_RANK - return dist.get_rank(group=get_tensor_model_parallel_group()) - - -def get_model_parallel_rank(): - return get_tensor_model_parallel_rank() - - -def get_tensor_model_parallel_src_rank(): - """Calculate the global rank corresponding to the first local rank - in the tensor model parallel group.""" - global_rank = dist.get_rank() - local_world_size = get_tensor_model_parallel_world_size() - return (global_rank // local_world_size) * local_world_size - - -def get_data_parallel_world_size(): - """Return world size for the data parallel group.""" - return dist.get_world_size(group=get_data_parallel_group()) - - -def get_data_parallel_rank(): - """Return my rank for the data parallel group.""" - return dist.get_rank(group=get_data_parallel_group()) - - -# ======== End: Tensor Parallel Group Attributes ======== - - -# Not currently used. Helper function to create a model (tensor) parallel group. -def _create_model_parallel(model_parallel_size_): - """ - Initialize model data parallel groups. - - Arguments: - model_parallel_size: number of GPUs used to parallelize model. - - Returns: - Tuple of data parallel group and model parallel group - - Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we - use 2 GPUs to parallelize the model. The present function will - create 4 model parallel groups and 2 data parallel groups as: - 4 model parallel groups: - [g0, g1], [g2, g3], [g4, g5], [g6, g7] - 2 data parallel groups: - [g0, g2, g4, g6], [g1, g3, g5, g7] - Note that for efficiency, the caller should make sure adjacent ranks - are on the same DGX box. For example if we are using 2 DGX-1 boxes - with a total of 16 GPUs, rank 0 to 7 belong to the first box and - ranks 8 to 15 belong to the second box. - """ - log_dist(f'Creating model parallel group with size {model_parallel_size_}', ranks=[0]) - # Get world size and rank. Ensure some consistencies. - assert dist.is_initialized() - world_size = dist.get_world_size() - model_parallel_size = min(model_parallel_size_, world_size) - _ensure_divisibility(world_size, model_parallel_size) - rank = dist.get_rank() - - _DATA_PARALLEL_GROUP = None - _MODEL_PARALLEL_GROUP = None - # Build the data parallel groups. - for i in range(model_parallel_size): - ranks = range(i, world_size, model_parallel_size) - group = dist.new_group(ranks) - if i == (rank % model_parallel_size): - _DATA_PARALLEL_GROUP = group - - # Build the model parallel groups. - for i in range(world_size // model_parallel_size): - ranks = range(i * model_parallel_size, (i + 1) * model_parallel_size) - group = dist.new_group(ranks) - if i == (rank // model_parallel_size): - _MODEL_PARALLEL_GROUP = group - - return _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP - - -def _create_expert_and_data_parallel(expert_parallel_size_, use_data_before_expert_parallel_=False): - """ - Create expert and data parallel groups. - - Note: Caller of this function is responsible to check if the groups already exist. - - Example - E + D parallel - world_size = 16 - expert_parallel_size = 2 # number of experts in same group - expert_data_parallel_group = [0,2,4,6,8,10,12,14], [1,3,5,7,9,11,13,15] - all reduce is only on MoE params - expert_parallel_group = [0, 1], [2,3], [4,5], [6,7], [8,9] - no all reduce, but all to all - data_parallel_group = [0,1,...,15] - all reduce is only on non-MoE - use_data_before_expert_parallel_ (bool): Use the D + E instead of E + D topology - """ - assert dist.is_initialized() - - log_dist(f'Creating expert and data parallel groups with size {expert_parallel_size_}', ranks=[0]) - world_size = dist.get_world_size() - pp_world_size = 1 if mpu is None else bwc_pipeline_parallel_world_size(mpu) - rank = dist.get_rank() - - pp_stride = world_size // pp_world_size - _ensure_divisibility(pp_stride, expert_parallel_size_) - - group_name = f"ep_size_{expert_parallel_size_}" - - # Build the expert data parallel groups. - global _EXPERT_DATA_PARALLEL_GROUP - global _EXPERT_DATA_PARALLEL_GROUP_RANKS - - ep_stride = pp_stride // expert_parallel_size_ - - # Only create group if it does not already exist - if group_name not in _EXPERT_DATA_PARALLEL_GROUP: - for pp_stage_start in range(0, world_size, pp_stride): - for i in range(expert_parallel_size_): - if use_data_before_expert_parallel_: - ranks = range(pp_stage_start + i * ep_stride, pp_stage_start + (i + 1) * ep_stride) - else: - ranks = range(pp_stage_start + i, pp_stage_start + pp_stride, expert_parallel_size_) - group = dist.new_group(ranks) - log_dist(f'Creating expert data parallel process group named {group_name} with ranks: {list(ranks)}', - [0]) - if rank in ranks: - _EXPERT_DATA_PARALLEL_GROUP[group_name] = group - _EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name] = ranks - - # Build the expert parallel groups. - global _EXPERT_PARALLEL_GROUP - global _EXPERT_PARALLEL_GROUP_RANKS - - # Only create group if it does not already exist - if group_name not in _EXPERT_PARALLEL_GROUP: - if use_data_before_expert_parallel_: - for pp_stage_start in range(0, world_size, pp_stride): - for i in range(ep_stride): - ranks = range(pp_stage_start + i, pp_stage_start + pp_stride, ep_stride) - group = dist.new_group(ranks) - log_dist( - f'creating expert parallel process group named {group_name} ' - f'with ranks: {list(ranks)}', [0]) - if rank in ranks: - _EXPERT_PARALLEL_GROUP[group_name] = group - _EXPERT_PARALLEL_GROUP_RANKS[group_name] = ranks - else: - for i in range(world_size // expert_parallel_size_): - ranks = range(i * expert_parallel_size_, (i + 1) * expert_parallel_size_) - group = dist.new_group(ranks) - log_dist(f'creating expert parallel process group named {group_name} ' - f'with ranks: {list(ranks)}', [0]) - if rank in ranks: - _EXPERT_PARALLEL_GROUP[group_name] = group - _EXPERT_PARALLEL_GROUP_RANKS[group_name] = ranks - - -def _get_expert_parallel_ranks(world_size, - tensor_parallel_size_, - expert_parallel_size_, - pipeline_parallel_size_=1, - use_data_before_expert_parallel_=False): - """Generate expert parallel and expert data parallel group ranks list. - - Example - E + M + D parallel - world_size = 16 - model_degree = 2 - expert_degree = 4 # number of experts in same group - mp_group = [0, 1], [2,3], [4,5] ... - data_parallel_group =[0,2,4,6,8,10, 12,14], [1,3,5,7,9,11,13,15] - expert_parallel_group = [0,2,4,6], [8,10,12,14] [1,3,5,7], [9,11,13,15] - expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14], [1,9],[3,11],[5,13],[7,15] - - Args: - world_size (int): Distributed world size. - tensor_parallel_size_ (int): Tensor parallel group size. - expert_parallel_size_ (int): Expert parallel group size. - pipeline_parallel_size_ (int): Pipeline parallel group size - use_data_before_expert_parallel_ (bool): Use the D + E instead of E + D topology - Returns: - Expert parallel group ranks and Expert data parallel group ranks list. - """ - _ensure_divisibility(world_size, tensor_parallel_size_ * pipeline_parallel_size_) - dp_world_size = world_size // (tensor_parallel_size_ * pipeline_parallel_size_) - _ensure_divisibility(dp_world_size, expert_parallel_size_) - - # Generate data parallel groups - data_parallel_groups = [] - dp_group_size = tensor_parallel_size_ - pp_stride = world_size // pipeline_parallel_size_ - - if use_data_before_expert_parallel_: - dp_stride = world_size // expert_parallel_size_ // tensor_parallel_size_ // pipeline_parallel_size_ - for pp_stage_start in range(0, world_size, pp_stride): - pp_stage_next = pp_stage_start + pp_stride - for i in range(dp_group_size): - data_parallel_groups.append(list()) - for ds in range(dp_stride): - # [0, 4, 8, 12, 16, 20, 24, 28, 2, 6, 10, 14, 18, 22, 26, 30] - # [1, 5, 9, 13, 17, 21, 25, 29, 3, 7, 11, 15, 19, 23, 27, 31] - data_parallel_groups[-1].extend( - list( - range(pp_stage_start + i + ds * tensor_parallel_size_, pp_stage_next, - dp_stride * tensor_parallel_size_))) - else: - for pp_stage_start in range(0, world_size, pp_stride): - pp_stage_next = pp_stage_start + pp_stride - for i in range(dp_group_size): - data_parallel_groups.append(list(range(pp_stage_start + i, pp_stage_next, dp_group_size))) - - expert_parallel_groups = [] - expert_data_parallel_groups = [] - for dp_ranks in data_parallel_groups: - # partition of expert parallel groups, e.g. [0,2,4,6], [8,10,12,14] - part_ep_groups = [] - for i in range(0, dp_world_size, expert_parallel_size_): - part_ep_groups.append(dp_ranks[i:i + expert_parallel_size_]) - expert_parallel_groups.extend(part_ep_groups) - - # zip part_ep_groups get expert data parallel ranks, e.g [0,8],[2,10],[4,12],[6,14] - for expert_dp_ranks in zip(*part_ep_groups): - expert_data_parallel_groups.append(list(expert_dp_ranks)) - - return expert_parallel_groups, expert_data_parallel_groups - - -def _create_expert_data_and_model_parallel(expert_parallel_size_, mpu, use_data_before_expert_parallel_=False): - """ - Create expert and data parallel groups based on MPU (model parallel) group. - - Note: Caller of this function is responsible to check if the groups already exist. - - Example - E + M + D parallel - world_size = 16 - model_degree = 2 - expert_degree = 4 # number of experts in same group - mp_group = [0, 1], [2,3], [4,5] ... - data_parallel_group =[0,2,4,6,8,10, 12,14], [1,3,5,7,9,11,13,15] - expert_parallel_group = [0,2,4,6], [8,10,12,14] [1,3,5,7], [9,11,13,15] - expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14], [1,9],[3,11],[5,13],[7,15] - """ - assert dist.is_initialized(), "dist is not initialized" - tensor_parallel_size_ = bwc_tensor_model_parallel_world_size(mpu) - - global expert_tensor_parallel_world_size - expert_tensor_parallel_world_size = tensor_parallel_size_ - - world_size = dist.get_world_size() - rank = dist.get_rank() - dp_world_size = mpu.get_data_parallel_world_size() - pp_world_size = 1 if mpu is None else bwc_pipeline_parallel_world_size(mpu) - - _ensure_divisibility(world_size, tensor_parallel_size_) - _ensure_divisibility(dp_world_size, expert_parallel_size_) - - log_dist( - f"Creating deepspeed groups with model parallel size {tensor_parallel_size_}, " - f"pipeline parallel size {pp_world_size}, expert parallel size {expert_parallel_size_}, " - f"world size {world_size}, dp world size {dp_world_size}", [0]) - - global _EXPERT_PARALLEL_GROUP, _EXPERT_DATA_PARALLEL_GROUP - global _EXPERT_PARALLEL_GROUP_RANKS, _EXPERT_DATA_PARALLEL_GROUP_RANKS - - group_name = f"ep_size_{expert_parallel_size_}" - - # Only create groups if they don't already exist - # Need to check conditions outside the group creation loop because of the way torch.dist group creation works - if group_name not in _EXPERT_DATA_PARALLEL_GROUP and group_name not in _EXPERT_PARALLEL_GROUP: - expert_parallel_groups, expert_data_parallel_groups = _get_expert_parallel_ranks( - world_size, tensor_parallel_size_, expert_parallel_size_, pp_world_size, use_data_before_expert_parallel_) - for ranks in expert_parallel_groups: - group = dist.new_group(ranks) - if rank in list(ranks): - _EXPERT_PARALLEL_GROUP[group_name] = group - _EXPERT_PARALLEL_GROUP_RANKS[group_name] = ranks - - for ranks in expert_data_parallel_groups: - group = dist.new_group(ranks) - if rank in list(ranks): - _EXPERT_DATA_PARALLEL_GROUP[group_name] = group - _EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name] = ranks - - -def _get_max_expert_size(): - """Get the maximum ep_size from all the created groups.""" - assert _EXPERT_PARALLEL_GROUP is not None, "Warning! Process group not initialized" - keylist = [] - for key in _EXPERT_PARALLEL_GROUP.keys(): - # index 2 is ep_size in the group name: ep_size_ - index = 2 - keylist.append(int(key.split('_')[index])) - return max(keylist) if len(keylist) > 0 else None - - -def _get_max_expert_size_name(): - """Get the name of the group with max. ep_size""" - return f'ep_size_{_get_max_expert_size()}' - - -def _get_max_expert_parallel_group(): - """Get the max expert parallel size.""" - return _get_expert_parallel_group(_get_max_expert_size_name()) - - -def _get_expert_parallel_group(group_name): - """Get the expert parallel group the caller rank belongs to.""" - assert group_name in _EXPERT_PARALLEL_GROUP, \ - 'expert parallel group is not initialized' - return _EXPERT_PARALLEL_GROUP[group_name] - - -def _get_expert_parallel_group_ranks(group_name): - """Get the ranks of the expert parallel group the caller rank belongs to.""" - assert group_name in _EXPERT_PARALLEL_GROUP_RANKS, \ - 'expert parallel group is not initialized' - return _EXPERT_PARALLEL_GROUP_RANKS[group_name] - - -def _get_expert_parallel_group_dict(): - """Get the expert parallel group dict.""" - return _EXPERT_PARALLEL_GROUP - - -def _get_expert_data_parallel_group(group_name): - """Get the expert data parallel group the caller rank belongs to.""" - assert group_name in _EXPERT_DATA_PARALLEL_GROUP, \ - 'expert data parallel group is not initialized' - return _EXPERT_DATA_PARALLEL_GROUP[group_name] - - -def _get_expert_data_parallel_group_ranks(group_name): - """Get the ranks of the expert data parallel group the caller rank belongs to.""" - assert group_name in _EXPERT_DATA_PARALLEL_GROUP_RANKS, \ - 'expert data parallel group is not initialized' - return _EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name] - - -def _get_expert_data_parallel_group_dict(): - """Get the expert data parallel group dict.""" - return _EXPERT_DATA_PARALLEL_GROUP - - -def _clone_world_group(): - """Create a clone of the world group - Note: We need to clone the dist world group because we - use dist.get_global_rank() utility function in DeepSpeed at many places. - As that function does not work on dist.group.WORLD, we - need to keep a clone of it. - """ - assert dist.is_initialized(), "dist is not initialized" - global _WORLD_GROUP - if _WORLD_GROUP is None: - # If not cloned already, clone the world group - _WORLD_GROUP = dist.new_group(ranks=range(dist.get_world_size())) - return _WORLD_GROUP - - -def _get_local_all_to_all_group(): - assert dist.is_initialized(), 'dist is not initialized' - global _ALL_TO_ALL_GROUP - device_per_node = get_accelerator().device_count() - num_local = dist.get_world_size() // device_per_node - if num_local == 0 and dist.get_world_size() > 0: - assert dist.get_world_size() >= 1, 'num_gpus must >=1, cannot initialize All-To-All' - cur_rank = [] - for i in range(dist.get_world_size()): - cur_rank.append(i) - _ALL_TO_ALL_GROUP['local_0'] = dist.new_group(ranks=cur_rank) - elif num_local == 1: - assert dist.get_world_size( - ) == device_per_node, 'num_gpus not equal to device per node, cannot initialize All-To-All' - _ALL_TO_ALL_GROUP['local_0'] = dist.new_group(ranks=[i for i in range(device_per_node)]) - else: - assert dist.get_world_size() > device_per_node, 'num_nodes<2 cannot initialize All-To-All' - for i in range(num_local): - local_rank = [j + device_per_node * i for j in range(device_per_node)] - _ALL_TO_ALL_GROUP[f"local_{i}"] = dist.new_group(ranks=local_rank) - - for i in range(device_per_node): - cur_rank = [] - for j in range(num_local): - cur_rank.append(i + j * device_per_node) - _ALL_TO_ALL_GROUP[f"global_{i}"] = dist.new_group(ranks=cur_rank) - return _ALL_TO_ALL_GROUP - - -def _get_data_parallel_group(): - """Get the data parallel group the caller rank belongs to.""" - assert dist.is_initialized(), 'dist is not initialized' - global mpu - if mesh_device is not None: - return mesh_device.get_group(mesh_dim="data_parallel") - if mpu is not None: - if hasattr(mpu, 'initialize_sequence_parallel'): - return None - else: - return mpu.get_data_parallel_group() - - # Return the clone of dist world group - return _clone_world_group() - - -def _get_data_parallel_group_ranks(): - """Get the ranks of data parallel group the caller rank belongs to.""" - assert dist.is_initialized(), \ - 'dist is not initialized' - global mpu - if mpu is not None: - return mpu.get_data_parallel_group_ranks() - # Return all ranks - return range(dist.get_world_size()) - - -def _get_broadcast_src_rank(): - return dist.get_global_rank(_get_sequence_data_parallel_group(), 0) - - -def _get_expert_broadcast_src_rank(group_name): - return dist.get_global_rank(_get_expert_data_parallel_group(group_name), 0) - - -def _get_expert_parallel_world_size(group_name): - """Return world size for the expert parallel group.""" - return dist.get_world_size(group=_get_expert_parallel_group(group_name)) - - -def _get_expert_data_parallel_world_size(group_name): - """Return world size for the expert data parallel group.""" - return dist.get_world_size(group=_get_expert_data_parallel_group(group_name)) - - -def _get_expert_parallel_rank(group_name): - """Return my rank for the expert parallel group.""" - return dist.get_rank(group=_get_expert_parallel_group(group_name)) - - -def _get_expert_parallel_src_rank(group_name): - """Calculate the global rank corresponding to a local rank zero - in the expert parallel group.""" - global_rank = dist.get_rank() - local_world_size = _get_expert_parallel_world_size(group_name) - return (global_rank // local_world_size) * local_world_size - - -def _get_expert_data_parallel_rank(group_name): - """Return my rank for the expert data parallel group.""" - return dist.get_rank(group=_get_expert_data_parallel_group(group_name)) - - -def _get_data_parallel_world_size(): - """Return world size for the data parallel group.""" - if mesh_device is not None: - return dist.get_world_size(mesh_device.get_group(mesh_dim="data_parallel")) - global mpu - if mpu is not None: - if hasattr(mpu, 'initialize_sequence_parallel'): - return None - else: - return mpu.get_data_parallel_world_size() - return dist.get_world_size(group=_get_data_parallel_group()) - - -def _get_model_parallel_world_size(): - """Return world size for the model parallel group.""" - global mpu - if mpu is None or hasattr(mpu, 'initialize_sequence_parallel'): - return 1 - return mpu.get_model_parallel_world_size() - - -def _get_data_parallel_rank(): - """Return my rank for the data parallel group.""" - return dist.get_rank(group=_get_data_parallel_group()) - - -def _get_sequence_parallel_world_size(): - """Return world size for the sequence parallel group.""" - global mpu - if mesh_device is not None: - return dist.get_world_size(mesh_device.get_group(mesh_dim="sequence_parallel")) - if mpu is not None and hasattr(mpu, 'get_sequence_parallel_world_size'): - return mpu.get_sequence_parallel_world_size() - return 1 - - -def _get_sequence_parallel_rank(): - """Return my rank for the sequence parallel group.""" - global mpu - if mpu is not None and hasattr(mpu, 'get_sequence_parallel_rank'): - return mpu.get_sequence_parallel_rank() - if mesh_device is not None: - return dist.get_rank(mesh_device.get_group(mesh_dim="sequence_parallel")) - return 0 - - -def _get_sequence_parallel_group(): - global mpu - if mpu is None or not hasattr(mpu, 'get_sequence_parallel_group'): - if mesh_device is None: - raise KeyError("No sequence parallel group found") - return mesh_device.get_group(mesh_dim="sequence_parallel") - return mpu.get_sequence_parallel_group() - - -def _get_sequence_data_parallel_world_size(): - """Return world size for the model parallel group.""" - global mpu - if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_world_size'): - return mpu.get_sequence_data_parallel_world_size() - return _get_data_parallel_world_size() - - -def _get_sequence_data_parallel_rank(): - """Return my rank for the data parallel group.""" - global mpu - if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_rank'): - return mpu.get_sequence_data_parallel_rank() - return _get_data_parallel_rank() - - -def _get_sequence_data_parallel_group(): - global mpu - # When sequence parallelism is enabled, the process group for zero sharding and - # gradient allreduce must be across both dimensions of data and sequence parallelism. - if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_group'): - return mpu.get_sequence_data_parallel_group() - return _get_data_parallel_group() - - -def _get_expert_model_parallel_world_size(): - global expert_tensor_parallel_world_size - return expert_tensor_parallel_world_size - - -def _create_zero_param_parallel_group(group_size): - """ - Create parameter partitioning group within ZeRO data parallel groups. - - Example - ZP + D parallel - world_size = 16 - zero_hpz_partition_size = 2 # number of ranks with replicated params (dual partitioning) - zero_param_intra_parallel_group = [0, 1], [2,3], [4,5], [6,7], [8,9] - segmented (subgroup) with rep partition - data_parallel_group = [0,1,...,15] - all reduce is on ZeRO model - """ - assert dist.is_initialized() - global _ZERO_PARAM_INTRA_PARALLEL_GROUP - # Only create group if it does not already exist - assert _ZERO_PARAM_INTRA_PARALLEL_GROUP is None, \ - 'ZeRO parameter intra parallel group is already initialized' - - world_size = dist.get_world_size() - rank = dist.get_rank() - - zero_param_parallel_size_ = min(group_size, world_size) - _ensure_divisibility(world_size, zero_param_parallel_size_) - - # Build the ZeRO param intra parallel groups. - for i in range(world_size // zero_param_parallel_size_): - ranks = range(i * zero_param_parallel_size_, (i + 1) * zero_param_parallel_size_) - group = dist.new_group(ranks) - if i == (rank // zero_param_parallel_size_): - _ZERO_PARAM_INTRA_PARALLEL_GROUP = group - - -def _get_zero_param_intra_parallel_group(): - """Get the ZeRO parameter partitioning intra parallel group the caller rank belongs to.""" - #assert _ZERO_PARAM_INTRA_PARALLEL_GROUP is not None, \ - # 'ZeRO parameter partitioning group is not initialized' - #TODO: Add warning - return _ZERO_PARAM_INTRA_PARALLEL_GROUP - - -def _zero_param_parallel_is_initialized(): - """Check if ZeRO data parallel with parameter partititioning groups are initialized.""" - ###TODO: assert that MPU is not set - if _ZERO_PARAM_INTRA_PARALLEL_GROUP is None and _DATA_PARALLEL_GROUP is None: - return False - - -def _get_zero_param_intra_parallel_rank_in_mygroup(): - """Return my rank for the ZeRO parameter inter parallel group.""" - return dist.get_rank(group=_get_zero_param_intra_parallel_group()) - - -def _get_zero_param_intra_parallel_group_world_size(): - """Return world size for the ZeRO parameter parallel group.""" - return dist.get_world_size(group=_get_zero_param_intra_parallel_group()) - - -def _get_zero_param_intra_parallel_group_ranks(): - """Return all ranks for the ZeRO parameter intra parallel group.""" - return dist.get_all_ranks_from_group(group=_get_zero_param_intra_parallel_group()) diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/init_on_device.py b/venv/lib/python3.10/site-packages/deepspeed/utils/init_on_device.py deleted file mode 100644 index 52dbf71d9562b73d868f71ee7ec7894a187e00f0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/init_on_device.py +++ /dev/null @@ -1,81 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from typing import Callable -from torch import Tensor -from packaging import version as pkg_version - - -class OnDevice(object): - """ - Create modules/tensors w. specific devices and dtypes. Examples: - - Create MyModule which consists of many different sub-modules and parameters. In this case we can create - MyModule as a collection of 'meta' tensors by passing `device='meta'` or we can create the module _directly_ - on a CUDA device by passing `device=f'cuda:{local_rank}'` (where `local_rank` is the local GPU id. - - with OnDevice(dtype=torch.float16, device='meta'): - model = MyModel() - - with OnDevice(dtype=torch.float16, device=f'cuda:{local_rank}'): - model = MyModel() - - """ - - _orig_torch_empty = torch.empty - _orig_torch_zeros = torch.zeros - _orig_torch_ones = torch.ones - _orig_torch_full = torch.full - - def __init__(self, dtype, device="meta", enabled=True): - self.dtype = dtype - self.enabled = enabled - self.device = device - - if device == "meta": - if pkg_version.parse('1.10') > pkg_version.parse(torch.__version__): - raise NotImplementedError("Meta tensor support is not available, please upgrade to torch 1.10+") - - def fp_tensor_constructor(self, fn: Callable, target_fp_dtype: torch.dtype) -> Callable: - - def wrapped_fn(*args, **kwargs) -> Tensor: - if kwargs.get("device", None) is None: - kwargs['device'] = self.device - tensor: Tensor = fn(*args, **kwargs) - if tensor.is_floating_point(): - tensor = tensor.to(target_fp_dtype) - return tensor - - return wrapped_fn - - def get_new_tensor_fn_for_dtype(self, dtype: torch.dtype) -> Callable: - - def new_tensor(cls, *args) -> Tensor: - tensor = OnDevice._orig_torch_empty(0, device=self.device).new_empty(*args) - if tensor.is_floating_point(): - tensor = tensor.to(dtype) - return tensor - - return new_tensor - - def __enter__(self): - if not self.enabled: - return - torch.Tensor.__old_new__ = torch.Tensor.__new__ - torch.Tensor.__new__ = self.get_new_tensor_fn_for_dtype(self.dtype) - torch.empty = self.fp_tensor_constructor(self._orig_torch_empty, self.dtype) - torch.zeros = self.fp_tensor_constructor(self._orig_torch_zeros, self.dtype) - torch.ones = self.fp_tensor_constructor(self._orig_torch_ones, self.dtype) - torch.full = self.fp_tensor_constructor(self._orig_torch_full, self.dtype) - - def __exit__(self, exc_type, exc_value, traceback): - if not self.enabled: - return - torch.Tensor.__new__ = torch.Tensor.__old_new__ - torch.empty = self._orig_torch_empty - torch.zeros = self._orig_torch_zeros - torch.ones = self._orig_torch_ones - torch.full = self._orig_torch_full diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/logging.py b/venv/lib/python3.10/site-packages/deepspeed/utils/logging.py deleted file mode 100644 index ed691e3985e1f1ca7360fe8dccb30f886013f411..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/logging.py +++ /dev/null @@ -1,163 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import functools -import logging -import sys -import os -import torch -from deepspeed.utils.torch import required_torch_version - -log_levels = { - "debug": logging.DEBUG, - "info": logging.INFO, - "warning": logging.WARNING, - "error": logging.ERROR, - "critical": logging.CRITICAL, -} - - -class LoggerFactory: - - @staticmethod - def create_logger(name=None, level=logging.INFO): - """create a logger - - Args: - name (str): name of the logger - level: level of logger - - Raises: - ValueError is name is None - """ - - if name is None: - raise ValueError("name for logger cannot be None") - - formatter = logging.Formatter("[%(asctime)s] [%(levelname)s] " - "[%(filename)s:%(lineno)d:%(funcName)s] %(message)s") - - logger_ = logging.getLogger(name) - logger_.setLevel(level) - logger_.propagate = False - ch = logging.StreamHandler(stream=sys.stdout) - ch.setLevel(level) - ch.setFormatter(formatter) - logger_.addHandler(ch) - if required_torch_version(min_version=2.6) and os.getenv("DISABLE_LOGS_WHILE_COMPILING", "0") == "1": - excluded_set = { - item.strip() - for item in os.getenv("LOGGER_METHODS_TO_EXCLUDE_FROM_DISABLE", "").split(",") - } - ignore_set = {'info', 'debug', 'error', 'warning', 'critical', 'exception', 'isEnabledFor'} - excluded_set - for method in ignore_set: - original_logger = getattr(logger_, method) - torch._dynamo.config.ignore_logger_methods.add(original_logger) - return logger_ - - -logger = LoggerFactory.create_logger(name="DeepSpeed", level=logging.INFO) - - -@functools.lru_cache(None) -def warning_once(*args, **kwargs): - """ - This method is identical to `logger.warning()`, but will emit the warning with the same message only once - - Note: The cache is for the function arguments, so 2 different callers using the same arguments will hit the cache. - The assumption here is that all warning messages are unique across the code. If they aren't then need to switch to - another type of cache that includes the caller frame information in the hashing function. - """ - logger.warning(*args, **kwargs) - - -logger.warning_once = warning_once - - -def print_configuration(args, name): - logger.info("{}:".format(name)) - for arg in sorted(vars(args)): - dots = "." * (29 - len(arg)) - logger.info(" {} {} {}".format(arg, dots, getattr(args, arg))) - - -def log_dist(message, ranks=None, level=logging.INFO): - from deepspeed import comm as dist - """Log message when one of following condition meets - - + not dist.is_initialized() - + dist.get_rank() in ranks if ranks is not None or ranks = [-1] - - Args: - message (str) - ranks (list) - level (int) - - """ - should_log = not dist.is_initialized() - ranks = ranks or [] - my_rank = dist.get_rank() if dist.is_initialized() else -1 - if ranks and not should_log: - should_log = ranks[0] == -1 - should_log = should_log or (my_rank in set(ranks)) - if should_log: - final_message = "[Rank {}] {}".format(my_rank, message) - logger.log(level, final_message) - - -def print_json_dist(message, ranks=None, path=None): - from deepspeed import comm as dist - """Print message when one of following condition meets - - + not dist.is_initialized() - + dist.get_rank() in ranks if ranks is not None or ranks = [-1] - - Args: - message (str) - ranks (list) - path (str) - - """ - should_log = not dist.is_initialized() - ranks = ranks or [] - my_rank = dist.get_rank() if dist.is_initialized() else -1 - if ranks and not should_log: - should_log = ranks[0] == -1 - should_log = should_log or (my_rank in set(ranks)) - if should_log: - message['rank'] = my_rank - import json - with open(path, 'w') as outfile: - json.dump(message, outfile) - os.fsync(outfile) - - -def get_current_level(): - """ - Return logger's current log level - """ - return logger.getEffectiveLevel() - - -def should_log_le(max_log_level_str): - """ - Args: - max_log_level_str: maximum log level as a string - - Returns ``True`` if the current log_level is less or equal to the specified log level. Otherwise ``False``. - - Example: - - ``should_log_le("info")`` will return ``True`` if the current log level is either ``logging.INFO`` or ``logging.DEBUG`` - """ - - if not isinstance(max_log_level_str, str): - raise ValueError(f"{max_log_level_str} is not a string") - - max_log_level_str = max_log_level_str.lower() - if max_log_level_str not in log_levels: - raise ValueError(f"{max_log_level_str} is not one of the logging levels") - - return get_current_level() <= log_levels[max_log_level_str] diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/mixed_precision_linkage.py b/venv/lib/python3.10/site-packages/deepspeed/utils/mixed_precision_linkage.py deleted file mode 100644 index c97515ca8fef7477e115914e6410bc0e778e3b8a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/mixed_precision_linkage.py +++ /dev/null @@ -1,51 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import types -from deepspeed.utils import get_full_hp_param, get_full_hp_grad, get_hp_fragment_mapping -from deepspeed.utils import set_full_hp_param, set_full_hp_grad - - -def link_hp_params(lp_param_list, flat_hp_partition, gradient_dict, offload_gradient_dict, use_offload, - param_group_index, partition_start, partition_size, dp_group): - local_lp_param_and_offset = _init_lp_to_hp_mapping(lp_param_list, partition_start, partition_size, dp_group) - - for lp_param, lp_start in local_lp_param_and_offset: - lp_param._hp_mapping = get_hp_fragment_mapping(lp_param, lp_start, flat_hp_partition, gradient_dict, - offload_gradient_dict, use_offload, param_group_index, - partition_start, partition_size) - - -def lazy_init_hp_params_optimizer_state(lp_param_list, flat_hp_partition, optimizer_state): - for lp in lp_param_list: - if lp._hp_mapping is not None: - lp._hp_mapping.set_optim_state_fragment(flat_hp_partition, optimizer_state[flat_hp_partition]) - - -def _init_lp_to_hp_mapping(lp_param_list, partition_start, partition_size, dp_group): - current_offset = 0 - param_and_offset_list = [] - partition_end = partition_start + partition_size - index_in_param_group = 0 - for i, lp_param in enumerate(lp_param_list): - lp_param._hp_mapping = None - lp_param._dp_group = dp_group - lp_param.get_full_hp_param = types.MethodType(get_full_hp_param, lp_param) - lp_param.get_full_hp_grad = types.MethodType(get_full_hp_grad, lp_param) - lp_param.set_full_hp_param = types.MethodType(set_full_hp_param, lp_param) - lp_param.set_full_hp_grad = types.MethodType(set_full_hp_grad, lp_param) - - # lp_param overlaps with partition if both are true - # 1) current_offset < partition_end, - # 2) current_offset + lp_param.numel() >= partition_start - lp_param_end = current_offset + lp_param.numel() - if current_offset < partition_end and lp_param_end > partition_start: - param_and_offset_list.append((lp_param, current_offset)) - lp_param._index_in_param_group = index_in_param_group - # Indices for params in this partition/GPU - index_in_param_group += 1 - current_offset += lp_param.numel() - - return param_and_offset_list diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/numa.py b/venv/lib/python3.10/site-packages/deepspeed/utils/numa.py deleted file mode 100644 index 75f0442a04e53042b9593cb72a7b9aae836ba25c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/numa.py +++ /dev/null @@ -1,205 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -# return a list of list for cores to numa mapping -# [ -# [ cores for numa 0 ] -# [ cores belong to numa 1 ] -# ... -# ] - -import os -import psutil -import shutil -import subprocess - - -# return a list of list for cores to numa mapping -# [ -# [ cores for numa 0 ] -# [ cores belong to numa 1 ] -# ... -# ] -def get_numa_cores(): - ret = [] - try: - output = subprocess.check_output(['numactl', '--hardware']).decode("utf-8") - except: - return [] - lines = output.split('\n') - for line in lines: - if line.startswith('available:'): - num_numas = int(line.split(' ')[1]) - break - for numa in range(num_numas): - for line in lines: - if line.startswith(f'node {numa} cpus:'): - cores = line.split(' ')[3:] - ret.append([int(core) for core in cores]) - return ret - - -def check_for_numactl_pkg(): - libs = dict( - dpkg=["-l", "numactl", "apt"], - pacman=["-Q", "numactl", "pacman"], - rpm=["-q", "numactl", "yum"], - ) - - found = False - for pkgmgr, data in libs.items(): - flag, lib, tool = data - path = shutil.which(pkgmgr) - if path is not None: - cmd = [pkgmgr, flag, lib] - result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - if result.wait() == 0: - found = True - else: - print(f"please install the {lib} package with {tool}") - break - return found - - -def parse_range(rng): - try: - value = int(rng) - return range(value, value + 1) - except ValueError: - # value is not a single number - parts = rng.split('-') - if len(parts) != 2: - raise ValueError("Bad range: '%s', range must be either a number or two number separated by dash" % - (rng, )) - start = int(parts[0]) - end = int(parts[1]) - if start > end: - raise ValueError("Bad range: '%s', range end must larger than or equal to start" % (rng, )) - return range(start, end + 1) - - -# parse comma and dash separated range list into list -# i.e. "0,2-4,6" --> [0, 2, 3, 4, 6] -# rules: -# 1. Range list number be comma separated, each item are either a single number, -# or a range marked by two numbers (both number are included in the range) -# 2. Sub ranges must be in ascend order and not overlap with each other -# 3. No space in the range expression -def parse_range_list(range_str): - number_list = [] - last = -1 - range_list = range_str.split(',') - for sub_range in range_list: - sub_number_list = parse_range(sub_range) - if sub_number_list[0] <= last: - raise ValueError( - "Bad range: '%s', sub ranges must not overlap with each other and should be in ascend order" % - (range_str, )) - last = sub_number_list[-1] - number_list.extend(sub_number_list) - return number_list - - -def get_numactl_cmd(bind_core_list, num_local_procs, local_rank): - numactl_cmd = [] - check_for_numactl_pkg() - if 'KMP_AFFINITY' in os.environ.keys(): - raise ValueError("Environment variable KMP_AFFINITY conflicts with numactl " - "because it interfere with how many CPU cores numactl can set. " - "Unset KMP_AFFINITY before launching deepspeed.\n\n" - "\t$ unset KMP_AFFINITY\n" - "\t$ deepspeed ") - if bind_core_list is not None: - core_list = parse_range_list(bind_core_list) - total_cores = len(core_list) - else: - total_cores = psutil.cpu_count(logical=False) - core_list = range(total_cores) - cores_per_rank = total_cores // num_local_procs - assert cores_per_rank >= 1, "At least one core needs to be assigned to each rank" - core_list_for_rank = core_list[cores_per_rank * local_rank:cores_per_rank * (local_rank + 1)] - numactl_cmd.append("numactl") - - # check if all cores belong to same numa, if true, bind process to that numa domain with -m parameter - numa_cores = get_numa_cores() - num_numas = len(numa_cores) - - numa_mode = "normal" - - non_empty_numa_list = [] - empty_numa_list = [] - previous_numa_cores = [] - numa_node_list = [] - numa_node_list_list = [] - for i in range(num_numas): - # look for empty numa which is HBM numa - if numa_cores[i] == []: - empty_numa_list.append(i) - else: - non_empty_numa_list.append(i) - - # check for fakenuma - if numa_cores[i] == previous_numa_cores: - if numa_node_list == []: - #first duplication, add previous node into list - numa_node_list.append(i - 1) - numa_node_list.append(i) - else: - if numa_node_list != []: - numa_node_list_list.append(numa_node_list) - numa_node_list = [] - previous_numa_cores = numa_cores[i] - if numa_node_list != []: - numa_node_list_list.append(numa_node_list) - - if empty_numa_list != [] and len(empty_numa_list) == len(non_empty_numa_list): - numa_mode = "flat_hbm" - numa_dict = dict(zip(non_empty_numa_list, empty_numa_list)) - elif numa_node_list_list != []: - numa_mode = "fake" - - if numa_mode == "normal": - for i in range(num_numas): - if set(core_list_for_rank) <= set(numa_cores[i]): - numactl_cmd.append("-m") - numactl_cmd.append(f"{i}") - break - elif numa_mode == "flat_hbm": - for i in range(num_numas): - if set(core_list_for_rank) <= set(numa_cores[i]): - numactl_cmd.append("-p") - numactl_cmd.append(f"{numa_dict[i]}") - break - elif numa_mode == "fake": - for i in range(num_numas): - if set(core_list_for_rank) <= set(numa_cores[i]): - for nodes in numa_node_list_list: - if i in nodes: - numactl_cmd.append("-m") - numactl_cmd.append(f"{','.join(map(str, nodes))}") - break - # the following construct break the outer loop if inner loop breaks - else: - continue - break - - numactl_cmd.append("-C") - last_core = core_list_for_rank[0] - first_core = last_core - core_list_str = f"{last_core}" - for core_id in core_list_for_rank[1:]: - if core_id == last_core + 1: - last_core = core_id - continue - else: - if first_core == last_core: - core_list_str = f"{core_list_str},{core_id}" - else: - core_list_str = f"{core_list_str}-{last_core},{core_id}" - first_core = core_id - last_core = core_id - if first_core != last_core: - core_list_str = f"{core_list_str}-{last_core}" - numactl_cmd.append(f"{core_list_str}") - return cores_per_rank, numactl_cmd diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/nvtx.py b/venv/lib/python3.10/site-packages/deepspeed/utils/nvtx.py deleted file mode 100644 index 72d7c863a33f5628fd7c327a6b11652ccba17a01..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/nvtx.py +++ /dev/null @@ -1,25 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from deepspeed.accelerator import get_accelerator -from deepspeed.runtime.compiler import is_compiling - -enable_nvtx = True - - -def instrument_w_nvtx(func): - """Decorator that records an NVTX range for the duration of the function call. - Skips NVTX instrumentation when torch.compile is active to avoid graph breaks. - """ - - def wrapped_fn(*args, **kwargs): - if enable_nvtx and not is_compiling(): - get_accelerator().range_push(func.__qualname__) - ret_val = func(*args, **kwargs) - if enable_nvtx and not is_compiling(): - get_accelerator().range_pop() - return ret_val - - return wrapped_fn diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/tensor_fragment.py b/venv/lib/python3.10/site-packages/deepspeed/utils/tensor_fragment.py deleted file mode 100644 index 305f8d56fd978e714a666a7e9617a23324e15ef5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/tensor_fragment.py +++ /dev/null @@ -1,481 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from dataclasses import dataclass -from deepspeed import comm as dist -from typing import Dict, List, Callable - - -@dataclass -class fragment_address: - numel: int - start: int - - -@dataclass -class tensor_fragment: - lp_fragment: torch.Tensor - lp_fragment_address: fragment_address - hp_fragment: torch.Tensor - hp_fragment_address: fragment_address - gradient_dict: Dict - offload_gradient_dict: Dict - use_offload: bool - param_group_index: int - optim_fragment: Dict = None - - def update_hp(self): - self.hp_fragment.data.copy_(self.lp_fragment.data) - - def update_lp(self): - self.lp_fragment.data.copy_(self.hp_fragment.data) - - def get_optim_state_fragment(self, key): - if key in self.optim_fragment: - return self.optim_fragment[key] - else: - raise ValueError(f'{key} not found in optimizer state fragment') - - def set_optim_state_fragment(self, flat_hp_partition, optim_fragment): - self.optim_fragment = { - key: value.narrow(0, self.hp_fragment_address.start, self.hp_fragment_address.numel) - for key, value in optim_fragment.items() - if torch.is_tensor(value) and value.shape == flat_hp_partition.shape - } - - def get_hp_fragment_address(self): - return self.hp_fragment_address - - def get_optim_state_keys(self): - return list(self.optim_fragment.keys()) - - def get_hp_fragment(self, optim_state_key=None): - if optim_state_key is None: - return self.hp_fragment - return self.get_optim_state_fragment(optim_state_key) - - def get_lp_grad_fragment(self, index_in_param_group): - if self.use_offload: - gradient_dict = self.offload_gradient_dict - else: - gradient_dict = self.gradient_dict - - if self.param_group_index not in gradient_dict or gradient_dict[self.param_group_index] is None: - raise ValueError("Gradients are only available immediately after backward and before engine step") - - return gradient_dict[self.param_group_index][index_in_param_group] - - -def map_to_flat_opt_states(flat_hp_tensor, lp_tensors, optim_state, opt_keys): - for key in opt_keys: - hp_param = flat_hp_tensor - buffer = torch.zeros_like(hp_param) - - for lp in lp_tensors: - if lp._hp_mapping is not None: - hp_fragment_address = lp._hp_mapping.get_hp_fragment_address() - hp_fragment = buffer.narrow(0, hp_fragment_address.start, hp_fragment_address.numel) - hp_fragment.data.copy_(lp._hp_mapping.get_hp_fragment(optim_state_key=key).data) - lp._hp_mapping.hp_fragment = hp_fragment - - optim_state[hp_param][key] = buffer - - -def get_full_hp_param(self, optim_state_key=None): - reduce_buffer = torch.zeros_like(self, dtype=torch.float32).flatten() - if self._hp_mapping is not None: - lp_frag_address = self._hp_mapping.lp_fragment_address - reduce_fragment = torch.narrow(reduce_buffer, 0, lp_frag_address.start, lp_frag_address.numel) - hp_fragment = self._hp_mapping.get_hp_fragment(optim_state_key) - reduce_fragment.data.copy_(hp_fragment.data) - dist.all_reduce(reduce_buffer, group=self._dp_group) - return reduce_buffer.reshape_as(self) - - -def set_full_hp_param(self, value, optim_state_key=None): - if self._hp_mapping is not None: - lp_frag_address = self._hp_mapping.lp_fragment_address - value_fragment = torch.narrow(value.flatten(), 0, lp_frag_address.start, lp_frag_address.numel) - hp_fragment = self._hp_mapping.get_hp_fragment(optim_state_key) - hp_fragment.data.copy_(value_fragment.data) - - -def get_full_hp_grad(self): - reduce_buffer = torch.zeros_like(self, dtype=torch.float32).flatten() - if self._hp_mapping is not None: - lp_grad_fragment = self._hp_mapping.get_lp_grad_fragment(self._index_in_param_group) - hp_grad_fragment = lp_grad_fragment.to(torch.float32).flatten() - - lp_frag_address = self._hp_mapping.lp_fragment_address - reduce_fragment = torch.narrow(reduce_buffer, 0, lp_frag_address.start, lp_frag_address.numel) - - if self.view(-1).shape == hp_grad_fragment.shape: - reduce_buffer.data.copy_(hp_grad_fragment.data) - else: - reduce_fragment.data.copy_(hp_grad_fragment.data) - - dist.all_reduce(reduce_buffer, group=self._dp_group) - return reduce_buffer.reshape_as(self) - - -def set_full_hp_grad(self, value): - if self._hp_mapping is not None: - lp_grad_fragment = self._hp_mapping.get_lp_grad_fragment(self._index_in_param_group) - lp_frag_address = self._hp_mapping.lp_fragment_address - value_fragment = torch.narrow(value.flatten(), 0, lp_frag_address.start, lp_frag_address.numel) - lp_grad_fragment.data.copy_(value_fragment.data.reshape_as(lp_grad_fragment.data)) - if hasattr(self, '_zero_optimizer'): - self._zero_optimizer.update_offload_overflow_tracker(value) - - -def safe_get_full_fp32_param(param): - """Assemble and return the fp32 parameter of a low-precision (e.g., fp16) parameter. - - Args: - param (``torch.nn.Parameter``): A model parameter - - Returns: - Union[torch.Tensor, None]: A tensor on accelerator device - """ - # ZeRO stage 3 param - if hasattr(param, 'ds_id'): - return param._z3_optimizer.get_full_hp_param(param) - - # ZeRO stage 1, 2, and bf16_optimizer params - if hasattr(param, '_hp_mapping'): - return param.get_full_hp_param() - return None - - -def safe_set_full_fp32_param(param, value): - """Update the partitioned fp32 parameter of a low-precision (e.g., fp16) parameter. - - Args: - param (``torch.nn.Parameter``): A model parameter - value (``torch.Tensor``): New value - """ - # ZeRO stage 3 param - if hasattr(param, 'ds_id'): - param._z3_optimizer.set_full_hp_param(value, param) - - # ZeRO stage 1, 2, and bf16_optimizer params - if hasattr(param, '_hp_mapping'): - param.set_full_hp_param(value) - - -def safe_get_full_optimizer_state(param, optim_state_key): - """Assemble and return the fp32 optimizer state of a low-precision (e.g., fp16) parameter. - - Args: - param (``torch.nn.Parameter``): A model parameter - optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer) - - Returns: - Union[torch.Tensor, None]: A tensor on accelerator device -""" - # ZeRO stage 3 param - if hasattr(param, 'ds_id'): - return param._z3_optimizer.get_full_hp_param(param, optim_state_key) - - # ZeRO stage 1, 2, and bf16_optimizer params - if hasattr(param, '_hp_mapping'): - return param.get_full_hp_param(optim_state_key) - return None - - -def safe_set_full_optimizer_state(param, value, optim_state_key): - """Update the partitioned fp32 optimizer state of a low-precision (e.g., fp16) parameter. - - Args: - param (``torch.nn.Parameter``): A model parameter - value (``torch.Tensor``): New value - optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer) - """ - # ZeRO stage 3 param - if hasattr(param, 'ds_id'): - param._z3_optimizer.set_full_hp_param(value, param, optim_state_key) - - # ZeRO stage 1, 2, and bf16_optimizer params - if hasattr(param, '_hp_mapping'): - param.set_full_hp_param(value, optim_state_key) - - -# TODO: Figure out the correct return dtype -def safe_get_full_grad(param): - """ - Assemble and return the fp32 gradient of a low-precision (e.g., fp16) parameter. - The return data type is that used for gradient accumulation. This is usually the param data type, - but could also be different (e.g., bf16 param training with fp32 gradient accumulation). - - Args: - param (``torch.nn.Parameter``): A model parameter - - Returns: - Union[torch.Tensor, None]: A tensor on accelerator device - """ - if param.grad is not None: - return param.grad - - # ZeRO stage 3 param - if hasattr(param, 'ds_id'): - return param._z3_optimizer.get_fp32_grad_for_param(param) - - # ZeRO stage 1, 2, and bf16_optimizer params - if hasattr(param, '_hp_mapping'): - return param.get_full_hp_grad() - - return None - - -def safe_set_full_grad(param, value): - """ - Update the partitioned gradient of a low-precision (e.g., fp16) parameter. - To avoid precision issues, the update value should have the data type of - gradient accumulation. - - Args: - param (``torch.nn.Parameter``): A model parameter - value (``torch.Tensor``): The un-partitioned new gradient value. - """ - if param.grad is not None: - param.grad.copy_(value) - elif hasattr(param, 'ds_id'): - # ZeRO stage 3 param - param._z3_optimizer.set_fp32_grad_for_param(value, param) - elif hasattr(param, '_hp_mapping'): - # ZeRO stage 1, 2, and bf16_optimizer params - param.set_full_hp_grad(value) - - -### Local API START ### -def safe_get_local_grad(param): - """ - Get the local gradient partition of a ZeRO-3 partitioned parameter. - The return data type is that used for gradient accumulation. This is usually the param data type, - but could also be different (e.g., bf16 param training with fp32 gradient accumulation). - - Args: - param (``torch.nn.Parameter``): A model parameter - - Returns: - Union[torch.Tensor, None]: A tensor on accelerator device - """ - assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' - return param._z3_optimizer.get_local_fp32_grad_for_param(param) - - -def safe_set_local_grad(param, value): - """ - Update the local gradient partition of a ZeRO-3 partitioned parameter. - To avoid precision issues, the update value should have the data type of - gradient accumulation. - - Args: - param (``torch.nn.Parameter``): A model parameter. - value (``torch.Tensor``): New value of local gradient partition. - """ - assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' - param._z3_optimizer.set_local_grad_for_param(value, param) - - -def safe_get_local_fp32_param(param): - """Get the local partition of a ZeRO-3 partitioned parameter in fp32 precision. - - Args: - param (``torch.nn.Parameter``): A model parameter. - - Returns: - Union[torch.Tensor, None]: A tensor on accelerator device - """ - assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' - return param._z3_optimizer.get_local_fp32_param(param) - - -def safe_get_local_optimizer_state(param, optim_state_key): - """Get the local optimizer state partition of ZeRO-3 partitioned parameter in fp32 precision. - - Args: - param (``torch.nn.Parameter``): A model parameter - optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer) - - Returns: - Union[torch.Tensor, None]: A tensor on accelerator device - """ - assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' - return param._z3_optimizer.get_local_fp32_param(param, optim_state_key) - - -def safe_set_local_optimizer_state(param, value, optim_state_key): - """Update the local optimizer state partition of a ZeRO-3 partitioned parameter. - - Args: - param (``torch.nn.Parameter``): A model parameter. - value (``torch.Tensor``): New value of local optimizer state partition. - optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer). - """ - assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' - param._z3_optimizer.set_local_hp_param(value, param, optim_state_key) - - -def safe_set_local_fp32_param(param, value): - """Update the local partition of ZeRO-3 partitioned parameter. - - Args: - param (``torch.nn.Parameter``): A model parameter. - value (``torch.Tensor``): New value of local parameter partition. - """ - assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' - param._z3_optimizer.set_local_hp_param(value, param) - - -### Local API END ### - - -### VECTORIZED API BEGIN ### -def safe_update_full_grad_vectorized(param_list: List[torch.nn.Parameter], update_func: Callable): - """ - Vectorized update of the partitioned gradients of a list of low-precision (e.g., fp16) parameters. - To avoid precision issues, the update value should have the data type of - gradient accumulation. - - Args: - param_list (``List[torch.nn.Parameter]``): List of model parameters - update_func (``torch.Tensor``): A function that takes current full gradient value and returns new one. - """ - partitioned_grad_params = [] - for p in param_list: - if p.grad is not None: - p.grad.copy_(update_func(p.grad, p)) - elif p.requires_grad: - partitioned_grad_params.append(p) - - if not partitioned_grad_params: - return - - if hasattr(partitioned_grad_params[0], 'ds_id'): - # ZeRO stage 3 param - partitioned_grad_params[0]._z3_optimizer.update_fp32_grad_for_param_vectorized( - update_func, partitioned_grad_params) - elif hasattr(partitioned_grad_params[0], '_hp_mapping'): - # ZeRO stage 1, 2, and bf16_optimizer params - for p in partitioned_grad_params: - old_grad = safe_get_full_grad(p) - new_grad = update_func(old_grad, p) - p.set_full_hp_grad(new_grad) - - -### VECTORIZED API END ### - - -def get_hp_fragment_mapping(lp_param, lp_start, flat_hp_partition, gradient_dict, offload_gradient_dict, use_offload, - param_group_index, partition_start, partition_size): - lp_end = lp_param.numel() + lp_start - hp_start = partition_start - hp_end = partition_start + partition_size - - fragment_start = max(lp_start, hp_start) - fragment_end = min(lp_end, hp_end) - assert fragment_start < fragment_end, \ - f'fragment start {fragment_start} should be < fragment_end {fragment_end}' - - fragment_numel = fragment_end - fragment_start - hp_frag_address = fragment_address(start=fragment_start - hp_start, numel=fragment_numel) - hp_fragment_tensor = flat_hp_partition.narrow(0, hp_frag_address.start, hp_frag_address.numel) - - lp_frag_address = fragment_address(start=fragment_start - lp_start, numel=fragment_numel) - lp_fragment_tensor = lp_param.flatten().narrow(0, lp_frag_address.start, lp_frag_address.numel) - - return tensor_fragment(lp_fragment=lp_fragment_tensor, - lp_fragment_address=lp_frag_address, - hp_fragment=hp_fragment_tensor, - hp_fragment_address=hp_frag_address, - gradient_dict=gradient_dict, - offload_gradient_dict=offload_gradient_dict, - use_offload=use_offload, - param_group_index=param_group_index) - - -''' -Logic for lp_param to hp_param mapping - -lp lp0 lp1 lp2 lp3 lp4 <------- indices/names -lp [ ][ ][ ][ ][ ] <-------- tensors -flat_lp [ ] <-------- flat lp params -flat_hp [ ] <------------------ flat hp partition on current rank -full_hp [ ] <------- full flat hp params - - -lp2 - full numel = 16 - lp_frag - numel = 12 - frag_start = 3 - frag_end = 15 - hp_frag - numel = 12 - frag_start = 0 - frag_end = 11 - - hp_frag.copy_(lp_frag) - - -lp3: - full numel = 4 - lp_frag - numel = 4 - start = 0 - end = 3 - hp_frag - numel = 4 - start = 12 - end = 15 - - -lp4: - full numel = 12 - lp_frag - numel = 4 - start = 0 - end = 3 - hp_frag - numel = 4 - start = 16 - end = 19 - - - -Visual depiction of above -lp { } -flat_lp [ ] -flat_hp ( ) - - -flat_lp [ { ( } ) ] - lx hx ly hy - ly-hx - - -lp { } -flat_lp [ ] -flat_hp ( ) - - -flat_lp [ ( { ) } ] - hx lx hy ly - hy-lx - -lp { } -flat_lp [ ] -flat_hp ( ) - - -flat_lp [ ( { } ) ] - hx lx ly hy - ly-lx - -lp -> (lx, hy) -flat_hp -> (hx, hy) -''' diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/timer.py b/venv/lib/python3.10/site-packages/deepspeed/utils/timer.py deleted file mode 100644 index 64ae8ac0e5b4f73d3bf36e8fdddc9dd3ad69d760..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/timer.py +++ /dev/null @@ -1,313 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import time -from numpy import mean -from deepspeed.utils.logging import log_dist -from deepspeed.accelerator import get_accelerator - -FORWARD_MICRO_TIMER = 'fwd_microstep' -FORWARD_GLOBAL_TIMER = 'fwd' -BACKWARD_MICRO_TIMER = 'bwd_microstep' -BACKWARD_GLOBAL_TIMER = 'bwd' -BACKWARD_INNER_MICRO_TIMER = 'bwd_inner_microstep' -BACKWARD_INNER_GLOBAL_TIMER = 'bwd_inner' -BACKWARD_REDUCE_MICRO_TIMER = 'bwd_allreduce_microstep' -BACKWARD_REDUCE_GLOBAL_TIMER = 'bwd_allreduce' -STEP_MICRO_TIMER = 'step_microstep' -STEP_GLOBAL_TIMER = 'step' -TIME_EPSILON = 1e-6 - -try: - import psutil - - PSUTILS_INSTALLED = True -except ImportError: - PSUTILS_INSTALLED = False - pass - - -class CudaEventTimer(object): - - def __init__(self, start_event: get_accelerator().Event, end_event: get_accelerator().Event): - self.start_event = start_event - self.end_event = end_event - - def get_elapsed_msec(self): - get_accelerator().current_stream().wait_event(self.end_event) - self.end_event.synchronize() - return self.start_event.elapsed_time(self.end_event) - - -class SynchronizedWallClockTimer: - """Group of timers. Borrowed from Nvidia Megatron code""" - - class Timer: - """Timer.""" - - def __init__(self, name): - self.name_ = name - self.started_ = False - self.event_timers = [] - self.use_host_timer = get_accelerator().use_host_timers() - self.start_event = None - self.elapsed_records = None - self.start_time = 0.0 - self.end_time = 0.0 - - def start(self): - """Start the timer.""" - assert not self.started_, f"{self.name_} timer has already been started" - if self.use_host_timer: - self.start_time = time.time() - else: - event_class = get_accelerator().Event - self.start_event = event_class(enable_timing=True) - self.start_event.record() - self.started_ = True - - def stop(self, reset=False, record=False): - """Stop the timer.""" - assert self.started_, "timer is not started" - event_class = get_accelerator().Event - if self.use_host_timer: - self.end_time = time.time() - self.event_timers.append(self.end_time - self.start_time) - else: - event_class = get_accelerator().Event - end_event = event_class(enable_timing=True) - end_event.record() - self.event_timers.append(CudaEventTimer(self.start_event, end_event)) - self.start_event = None - self.started_ = False - - def _get_elapsed_msec(self): - if self.use_host_timer: - self.elapsed_records = [et * 1000.0 for et in self.event_timers] - else: - self.elapsed_records = [et.get_elapsed_msec() for et in self.event_timers] - self.event_timers.clear() - return sum(self.elapsed_records) - - def reset(self): - """Reset timer.""" - self.started_ = False - self.start_event = None - self.elapsed_records = None - self.event_timers.clear() - - def elapsed(self, reset=True): - """Calculate the elapsed time.""" - started_ = self.started_ - # If the timing in progress, end it first. - if self.started_: - self.stop() - # Get the elapsed time. - elapsed_ = self._get_elapsed_msec() - # Reset the elapsed time - if reset: - self.reset() - # If timing was in progress, set it back. - if started_: - self.start() - return elapsed_ - - def mean(self): - self.elapsed(reset=False) - return trim_mean(self.elapsed_records, 0.1) - - def __init__(self): - self.timers = {} - - def get_timers(self): - return self.timers - - def __call__(self, name): - if name not in self.timers: - self.timers[name] = self.Timer(name) - return self.timers[name] - - @staticmethod - def memory_usage(): - alloc = "mem_allocated: {:.4f} GB".format(get_accelerator().memory_allocated() / (1024 * 1024 * 1024)) - max_alloc = "max_mem_allocated: {:.4f} GB".format(get_accelerator().max_memory_allocated() / - (1024 * 1024 * 1024)) - cache = "cache_allocated: {:.4f} GB".format(get_accelerator().memory_cached() / (1024 * 1024 * 1024)) - max_cache = "max_cache_allocated: {:.4f} GB".format(get_accelerator().max_memory_cached() / - (1024 * 1024 * 1024)) - return " | {} | {} | {} | {}".format(alloc, max_alloc, cache, max_cache) - - def log(self, names, normalizer=1.0, reset=True, memory_breakdown=False, ranks=None): - """Log a group of timers.""" - assert normalizer > 0.0 - string = f"time (ms)" - for name in names: - if name in self.timers: - elapsed_time = (self.timers[name].elapsed(reset=reset) / normalizer) - string += " | {}: {:.2f}".format(name, elapsed_time) - - log_dist(string, ranks=ranks or [0]) - - def get_mean(self, names, normalizer=1.0, reset=True): - """Get the mean of a group of timers.""" - assert normalizer > 0.0 - means = {} - for name in names: - if name in self.timers: - elapsed_time = (self.timers[name].mean() * 1000.0 / normalizer) - means[name] = elapsed_time - return means - - -class NoopTimer: - - class Timer: - - def start(self): - ... - - def reset(self): - ... - - def stop(self, **kwargs): - ... - - def elapsed(self, **kwargs): - return 0 - - def mean(self): - return 0 - - def __init__(self): - self.timer = self.Timer() - - def __call__(self, name): - return self.timer - - def get_timers(self): - return {} - - def log(self, names, normalizer=1.0, reset=True, memory_breakdown=False, ranks=None): - ... - - def get_mean(self, names, normalizer=1.0, reset=True): - ... - - -class ThroughputTimer: - - def __init__(self, config, batch_size, start_step=2, steps_per_output=None, monitor_memory=False, logging_fn=None): - from deepspeed.utils import logger - self.config = config - self.start_time = 0 - self.end_time = 0 - self.started = False - self.batch_size = 1 if batch_size is None else batch_size - self.start_step = start_step - self.epoch_count = 0 - self.micro_step_count = 0 - self.global_step_count = 0 - self.total_elapsed_time = 0 - self.step_elapsed_time = 0 - self.steps_per_output = steps_per_output - self.monitor_memory = monitor_memory - self.logging = logging_fn - if self.logging is None: - self.logging = logger.info - self.initialized = False - - if self.monitor_memory and not PSUTILS_INSTALLED: - raise ImportError("Unable to import 'psutils', please install package") - - def update_epoch_count(self): - self.epoch_count += 1 - self.micro_step_count = 0 - - def _init_timer(self): - self.initialized = True - - def start(self): - if not self.config.enabled: - return - self._init_timer() - self.started = True - if self.global_step_count >= self.start_step: - if self.config.synchronized: - get_accelerator().synchronize() - self.start_time = time.time() - - def _is_report_boundary(self): - if self.steps_per_output is None: - return False - return self.global_step_count % self.steps_per_output == 0 - - def stop(self, global_step=False, report_speed=True): - if not self.config.enabled or not self.started: - return - self.started = False - self.micro_step_count += 1 - if global_step: - self.global_step_count += 1 - - if self.start_time > 0: - if self.config.synchronized: - get_accelerator().synchronize() - self.end_time = time.time() - duration = self.end_time - self.start_time - self.total_elapsed_time += duration - self.step_elapsed_time += duration - - if global_step: - if report_speed and self._is_report_boundary(): - self.logging( - "epoch={}/micro_step={}/global_step={}, RunningAvgSamplesPerSec={}, CurrSamplesPerSec={}, " - "MemAllocated={}GB, MaxMemAllocated={}GB".format( - self.epoch_count, - self.micro_step_count, - self.global_step_count, - self.avg_samples_per_sec(), - self.batch_size / (self.step_elapsed_time + TIME_EPSILON), - round(get_accelerator().memory_allocated() / 1024**3, 2), - round(get_accelerator().max_memory_allocated() / 1024**3, 2), - )) - if self.monitor_memory: - virt_mem = psutil.virtual_memory() - swap = psutil.swap_memory() - self.logging("epoch={}/micro_step={}/global_step={}, vm %: {}, swap %: {}".format( - self.epoch_count, - self.micro_step_count, - self.global_step_count, - virt_mem.percent, - swap.percent, - )) - self.step_elapsed_time = 0 - - def avg_samples_per_sec(self): - if self.global_step_count > 0: - total_step_offset = self.global_step_count - self.start_step - avg_time_per_step = self.total_elapsed_time / total_step_offset - # training samples per second - return self.batch_size / avg_time_per_step - return float("-inf") - - -def trim_mean(data, trim_percent): - """Compute the trimmed mean of a list of numbers. - - Args: - data (list): List of numbers. - trim_percent (float): Percentage of data to trim. - - Returns: - float: Trimmed mean. - """ - assert 0.0 <= trim_percent <= 1.0 - n = len(data) - # Account for edge case of empty list - if len(data) == 0: - return 0 - data.sort() - k = int(round(n * (trim_percent))) - return mean(data[k:n - k]) diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/torch.py b/venv/lib/python3.10/site-packages/deepspeed/utils/torch.py deleted file mode 100644 index 1d32775fe64a8e9bf1bff5df6aaa111c2974d53d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/torch.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from packaging import version as pkg_version - -import torch - - -def required_torch_version(min_version=None, max_version=None): - assert min_version or max_version, "Must provide a min_version or max_version argument" - - torch_version = pkg_version.parse(torch.__version__) - - if min_version and pkg_version.parse(str(min_version)) > torch_version: - return False - - if max_version and pkg_version.parse(str(max_version)) < torch_version: - return False - - return True - - -def register_grad_hook(param, hook): - if required_torch_version(min_version=2.1): - return param.register_post_accumulate_grad_hook(hook) - else: - param_tmp = param.expand_as(param) - grad_acc = param_tmp.grad_fn.next_functions[0][0] - return grad_acc.register_hook(hook) diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/types.py b/venv/lib/python3.10/site-packages/deepspeed/utils/types.py deleted file mode 100644 index 96b5df625965401f52f6fd9b60396821943fe177..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/types.py +++ /dev/null @@ -1,27 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from enum import IntEnum - - -class ActivationFuncType(IntEnum): - UNKNOWN = 0 - GELU = 1 - ReLU = 2 - GATED_GELU = 3 - GATED_SILU = 4 - - -GATED_ACTIVATION_TYPES = [ - ActivationFuncType.GATED_GELU, - ActivationFuncType.GATED_SILU, -] - - -class NormType(IntEnum): - UNKNOWN = 0 - LayerNorm = 1 - GroupNorm = 2 - RMSNorm = 3 diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/z3_leaf_module.py b/venv/lib/python3.10/site-packages/deepspeed/utils/z3_leaf_module.py deleted file mode 100644 index 14e8ae2d28235c46016b9030b3492a2eb4e7122a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/z3_leaf_module.py +++ /dev/null @@ -1,93 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -from typing import List, Type, Union - - -def z3_leaf_module(model: torch.nn.Module) -> bool: - """Returns whether a module in `model` has been flagged as a 'leaf' module. - See `set_z3_leaf_modules` for more details. - Args: - model (torch.nn.Module): The model to which the leaf module flag will be applied. - Returns: - bool: Whether the module has been flagged as a 'leaf' module. - """ - return hasattr(model, '_z3_leaf') and model._z3_leaf - - -def z3_leaf_parameter(model: torch.nn.Parameter) -> bool: - """Returns whether a parameter belongs to a leaf module. - See `set_z3_leaf_modules` for more details. - Args: - model (torch.nn.Parameter): The parameter to which the leaf module flag will be applied. - Returns: - bool: Whether the parameter belongs to a leaf module. - """ - return hasattr(model, 'ds_z3_leaf_module') - - -def get_z3_leaf_modules(model: torch.nn.Module) -> List[torch.nn.Module]: - """Returns a list of modules in `model` that have been flagged as 'leaf' modules. - See `set_z3_leaf_modules` for more details. - Args: - model (torch.nn.Module): The model to which the leaf module flag will be applied. - Returns: - List[torch.nn.Module]: A list of modules that have been flagged as 'leaf' modules. - """ - return [module for module in model.modules() if z3_leaf_module(module)] - - -def set_z3_leaf_module(model: torch.nn.Module, flag: bool): - model._z3_leaf = flag - - -def _do_set_z3_leaf_modules(model: torch.nn.Module, leaf_module_classes: Union[List[Type], List[str]], - flag: bool) -> List[torch.nn.Module]: - assert all(isinstance(module_class, (type, str) ) for module_class in leaf_module_classes), \ - f'leaf_module_classes must be a list of types or names, got {leaf_module_classes}' - - leaf_modules = [] - - def _set_z3_leaf_flag(model: torch.nn.Module): - nonlocal leaf_modules - for module in leaf_module_classes: - if (isinstance(module, type) and model.__class__ == module) or \ - (isinstance(module, str) and model.__class__.__name__ == module): - model._z3_leaf = flag - leaf_modules.append(model) - - model.apply(_set_z3_leaf_flag) - - if len(leaf_modules) == 0: - raise ValueError(f'No modules of type {leaf_module_classes} found in model {model}') - - return leaf_modules - - -def set_z3_leaf_modules(model: torch.nn.Module, leaf_module_classes: Union[List[Type], - List[str]]) -> List[torch.nn.Module]: - """Sets a flag within a module in `model` to instruct ZeRO3 to stop setting hooks recursively when it encounters a module class listed in `leaf_module_classes`. - This is particularly useful in the context of Mixture of Experts (MoE) models. In MoE models, the computation order of experts varies across forward passes. This variability can disrupt ZeRO3's functionality, as ZeRO3 relies on tracking the computation order of modules to prefetch parameters efficiently. By designating a module as a 'leaf' node, ZeRO3 will prefetch parameters for all child modules upon entering the module. - Another scenario where this functionality is beneficial is in models with excessively fine-grained nested modules, where it helps to avoid the overhead associated with hooks. - Args: - model (torch.nn.Module): The model to which the leaf module flag will be applied. - leaf_module_classes (Union[List[Type], List[str]]): A list of module classes that should be flagged as 'leaf' modules. - Returns: - List[torch.nn.Module]: A list of modules that match the module classes in `leaf_module_classes`. - """ - return _do_set_z3_leaf_modules(model, leaf_module_classes, True) - - -def unset_z3_leaf_modules(model: torch.nn.Module, leaf_module_classes: List[Type]) -> List[torch.nn.Module]: - """Unsets a flag within a module in `model` to instruct ZeRO3 to resume setting hooks recursively when it encounters a module class listed in `leaf_module_classes`. - See `set_z3_leaf_modules` for more details. - Args: - model (torch.nn.Module): The model to which the leaf module flag will be applied. - leaf_module_classes (Union[List[Type], List[str]]): A list of module classes that should be flagged as 'leaf' modules. - Returns: - List[torch.nn.Module]: A list of modules that match the module classes in `leaf_module_classes`. - """ - return _do_set_z3_leaf_modules(model, leaf_module_classes, False) diff --git a/venv/lib/python3.10/site-packages/deepspeed/utils/zero_to_fp32.py b/venv/lib/python3.10/site-packages/deepspeed/utils/zero_to_fp32.py deleted file mode 100644 index 0e759146cadd92ddfefab3680146c2bd6a2b5c04..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deepspeed/utils/zero_to_fp32.py +++ /dev/null @@ -1,760 +0,0 @@ -#!/usr/bin/env python - -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets -# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in -# the future. Once extracted, the weights don't require DeepSpeed and can be used in any -# application. -# -# example: -# python zero_to_fp32.py . output_dir/ -# or -# python zero_to_fp32.py . output_dir/ --safe_serialization - -import argparse -import torch -import glob -import math -import os -import re -import gc -import json -import numpy as np -from tqdm import tqdm -from collections import OrderedDict -from dataclasses import dataclass - -# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with -# DeepSpeed data structures it has to be available in the current python environment. -from deepspeed.utils import logger -from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS, - FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES, - FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS) - - -@dataclass -class zero_model_state: - buffers: dict() - param_shapes: dict() - shared_params: list - ds_version: int - frozen_param_shapes: dict() - frozen_param_fragments: dict() - - -debug = 0 - -# load to cpu -device = torch.device('cpu') - - -def atoi(text): - return int(text) if text.isdigit() else text - - -def natural_keys(text): - ''' - alist.sort(key=natural_keys) sorts in human order - http://nedbatchelder.com/blog/200712/human_sorting.html - (See Toothy's implementation in the comments) - ''' - return [atoi(c) for c in re.split(r'(\d+)', text)] - - -def get_model_state_file(checkpoint_dir, zero_stage): - if not os.path.isdir(checkpoint_dir): - raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist") - - # there should be only one file - if zero_stage <= 2: - file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt") - elif zero_stage == 3: - file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt") - - if not os.path.exists(file): - raise FileNotFoundError(f"can't find model states file at '{file}'") - - return file - - -def get_checkpoint_files(checkpoint_dir, glob_pattern): - # XXX: need to test that this simple glob rule works for multi-node setup too - ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys) - - if len(ckpt_files) == 0: - raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'") - - return ckpt_files - - -def get_optim_files(checkpoint_dir): - return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt") - - -def get_model_state_files(checkpoint_dir): - return get_checkpoint_files(checkpoint_dir, "*_model_states.pt") - - -def parse_model_states(files): - zero_model_states = [] - for file in files: - state_dict = torch.load(file, map_location=device, weights_only=False) - - if BUFFER_NAMES not in state_dict: - raise ValueError(f"{file} is not a model state checkpoint") - buffer_names = state_dict[BUFFER_NAMES] - if debug: - print("Found buffers:", buffer_names) - - # recover just the buffers while restoring them to fp32 if they were saved in fp16 - buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names} - param_shapes = state_dict[PARAM_SHAPES] - - # collect parameters that are included in param_shapes - param_names = [] - for s in param_shapes: - for name in s.keys(): - param_names.append(name) - - # update with frozen parameters - frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None) - if frozen_param_shapes is not None: - if debug: - print(f"Found frozen_param_shapes: {frozen_param_shapes}") - param_names += list(frozen_param_shapes.keys()) - - # handle shared params - shared_params = [[k, v] for k, v in state_dict["shared_params"].items()] - - ds_version = state_dict.get(DS_VERSION, None) - - frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None) - - z_model_state = zero_model_state(buffers=buffers, - param_shapes=param_shapes, - shared_params=shared_params, - ds_version=ds_version, - frozen_param_shapes=frozen_param_shapes, - frozen_param_fragments=frozen_param_fragments) - zero_model_states.append(z_model_state) - - return zero_model_states - - -def parse_optim_states(files, ds_checkpoint_dir): - total_files = len(files) - state_dicts = [] - for f in tqdm(files, desc='Loading checkpoint shards'): - state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False) - # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights - # and also handle the case where it was already removed by another helper script - state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None) - state_dicts.append(state_dict) - - if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]: - raise ValueError(f"{files[0]} is not a zero checkpoint") - zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE] - world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT] - - # For ZeRO-2 each param group can have different partition_count as data parallelism for expert - # parameters can be different from data parallelism for non-expert parameters. So we can just - # use the max of the partition_count to get the dp world_size. - - if type(world_size) is list: - world_size = max(world_size) - - if world_size != total_files: - raise ValueError( - f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. " - "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes." - ) - - # the groups are named differently in each stage - if zero_stage <= 2: - fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS - elif zero_stage == 3: - fp32_groups_key = FP32_FLAT_GROUPS - else: - raise ValueError(f"unknown zero stage {zero_stage}") - - fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))] - return zero_stage, world_size, fp32_flat_groups - - -def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters): - """ - Returns fp32 state_dict reconstructed from ds checkpoint - - Args: - - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are) - - """ - print(f"Processing zero checkpoint '{ds_checkpoint_dir}'") - - optim_files = get_optim_files(ds_checkpoint_dir) - zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir) - print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}") - - model_files = get_model_state_files(ds_checkpoint_dir) - - zero_model_states = parse_model_states(model_files) - print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}') - - if zero_stage <= 2: - return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states, - exclude_frozen_parameters) - elif zero_stage == 3: - return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states, - exclude_frozen_parameters) - - -def _zero2_merge_frozen_params(state_dict, zero_model_states): - if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0: - return - - frozen_param_shapes = zero_model_states[0].frozen_param_shapes - frozen_param_fragments = zero_model_states[0].frozen_param_fragments - - if debug: - num_elem = sum(s.numel() for s in frozen_param_shapes.values()) - print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}') - - wanted_params = len(frozen_param_shapes) - wanted_numel = sum(s.numel() for s in frozen_param_shapes.values()) - avail_numel = sum([p.numel() for p in frozen_param_fragments.values()]) - print(f'Frozen params: Have {avail_numel} numels to process.') - print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params') - - total_params = 0 - total_numel = 0 - for name, shape in frozen_param_shapes.items(): - total_params += 1 - unpartitioned_numel = shape.numel() - total_numel += unpartitioned_numel - - state_dict[name] = frozen_param_fragments[name] - - if debug: - print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ") - - print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements") - - -def _has_callable(obj, fn): - attr = getattr(obj, fn, None) - return callable(attr) - - -def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states): - param_shapes = zero_model_states[0].param_shapes - - # Reconstruction protocol: - # - # XXX: document this - - if debug: - for i in range(world_size): - for j in range(len(fp32_flat_groups[0])): - print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}") - - # XXX: memory usage doubles here (zero2) - num_param_groups = len(fp32_flat_groups[0]) - merged_single_partition_of_fp32_groups = [] - for i in range(num_param_groups): - merged_partitions = [sd[i] for sd in fp32_flat_groups] - full_single_fp32_vector = torch.cat(merged_partitions, 0) - merged_single_partition_of_fp32_groups.append(full_single_fp32_vector) - avail_numel = sum( - [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups]) - - if debug: - wanted_params = sum([len(shapes) for shapes in param_shapes]) - wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes]) - # not asserting if there is a mismatch due to possible padding - print(f"Have {avail_numel} numels to process.") - print(f"Need {wanted_numel} numels in {wanted_params} params.") - - # params - # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support - # out-of-core computing solution - total_numel = 0 - total_params = 0 - for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups): - offset = 0 - avail_numel = full_single_fp32_vector.numel() - for name, shape in shapes.items(): - - unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape) - total_numel += unpartitioned_numel - total_params += 1 - - if debug: - print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ") - state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape) - offset += unpartitioned_numel - - # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and - # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex - # paddings performed in the code it's almost impossible to predict the exact numbers w/o the - # live optimizer object, so we are checking that the numbers are within the right range - align_to = 2 * world_size - - def zero2_align(x): - return align_to * math.ceil(x / align_to) - - if debug: - print(f"original offset={offset}, avail_numel={avail_numel}") - - offset = zero2_align(offset) - avail_numel = zero2_align(avail_numel) - - if debug: - print(f"aligned offset={offset}, avail_numel={avail_numel}") - - # Sanity check - if offset != avail_numel: - raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong") - - print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements") - - -def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states, - exclude_frozen_parameters): - state_dict = OrderedDict() - - # buffers - buffers = zero_model_states[0].buffers - state_dict.update(buffers) - if debug: - print(f"added {len(buffers)} buffers") - - if not exclude_frozen_parameters: - _zero2_merge_frozen_params(state_dict, zero_model_states) - - _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states) - - # recover shared parameters - for pair in zero_model_states[0].shared_params: - if pair[1] in state_dict: - state_dict[pair[0]] = state_dict[pair[1]] - - return state_dict - - -def zero3_partitioned_param_info(unpartitioned_numel, world_size): - remainder = unpartitioned_numel % world_size - padding_numel = (world_size - remainder) if remainder else 0 - partitioned_numel = math.ceil(unpartitioned_numel / world_size) - return partitioned_numel, padding_numel - - -def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states): - if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0: - return - - if debug: - for i in range(world_size): - num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values()) - print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}') - - frozen_param_shapes = zero_model_states[0].frozen_param_shapes - wanted_params = len(frozen_param_shapes) - wanted_numel = sum(s.numel() for s in frozen_param_shapes.values()) - avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size - print(f'Frozen params: Have {avail_numel} numels to process.') - print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params') - - total_params = 0 - total_numel = 0 - for name, shape in zero_model_states[0].frozen_param_shapes.items(): - total_params += 1 - unpartitioned_numel = shape.numel() - total_numel += unpartitioned_numel - - param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states) - state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape) - - partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size) - - if debug: - print( - f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}" - ) - - print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements") - - -class GatheredTensor: - """ - A pseudo tensor that collects partitioned weights. - It is more memory efficient when there are multiple groups. - """ - - def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape): - self.flat_groups = flat_groups - self.flat_groups_offset = flat_groups_offset - self.offset = offset - self.partitioned_numel = partitioned_numel - self.shape = shape - self.dtype = self.flat_groups[0][0].dtype - - def contiguous(self): - """ - Merge partitioned weights from flat_groups into a single tensor. - """ - end_idx = self.offset + self.partitioned_numel - world_size = len(self.flat_groups) - pad_flat_param_chunks = [] - - for rank_i in range(world_size): - # for each rank, we need to collect weights from related group/groups - flat_groups_at_rank_i = self.flat_groups[rank_i] - start_group_id = None - end_group_id = None - for group_id in range(len(self.flat_groups_offset)): - if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]: - start_group_id = group_id - if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]: - end_group_id = group_id - break - # collect weights from related group/groups - for group_id in range(start_group_id, end_group_id + 1): - flat_tensor = flat_groups_at_rank_i[group_id] - start_offset = self.offset - self.flat_groups_offset[group_id] - end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id] - pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset]) - - # collect weights from all ranks - pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0) - param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous() - return param - - -def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states): - param_shapes = zero_model_states[0].param_shapes - avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size - - # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each - # param, re-consolidating each param, while dealing with padding if any - - # merge list of dicts, preserving order - param_shapes = {k: v for d in param_shapes for k, v in d.items()} - - if debug: - for i in range(world_size): - print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}") - - wanted_params = len(param_shapes) - wanted_numel = sum(shape.numel() for shape in param_shapes.values()) - # not asserting if there is a mismatch due to possible padding - avail_numel = fp32_flat_groups[0].numel() * world_size - print(f"Trainable params: Have {avail_numel} numels to process.") - print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.") - - # params - # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support - # out-of-core computing solution - offset = 0 - total_numel = 0 - total_params = 0 - flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]])) - for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'): - unpartitioned_numel = shape.numel() - total_numel += unpartitioned_numel - total_params += 1 - partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size) - - if debug: - print( - f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}" - ) - - # memory efficient tensor - tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape) - state_dict[name] = tensor - offset += partitioned_numel - - offset *= world_size - - # Sanity check - if offset != avail_numel: - raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong") - - print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements") - - -def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states, - exclude_frozen_parameters): - state_dict = OrderedDict() - - # buffers - buffers = zero_model_states[0].buffers - state_dict.update(buffers) - if debug: - print(f"added {len(buffers)} buffers") - - if not exclude_frozen_parameters: - _zero3_merge_frozen_params(state_dict, world_size, zero_model_states) - - _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states) - - # recover shared parameters - for pair in zero_model_states[0].shared_params: - if pair[1] in state_dict: - state_dict[pair[0]] = state_dict[pair[1]] - - return state_dict - - -def to_torch_tensor(state_dict, return_empty_tensor=False): - """ - Convert state_dict of GatheredTensor to torch tensor - """ - torch_state_dict = {} - converted_tensors = {} - for name, tensor in state_dict.items(): - tensor_id = id(tensor) - if tensor_id in converted_tensors: # shared tensors - shared_tensor = torch_state_dict[converted_tensors[tensor_id]] - torch_state_dict[name] = shared_tensor - else: - converted_tensors[tensor_id] = name - if return_empty_tensor: - torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype) - else: - torch_state_dict[name] = tensor.contiguous() - return torch_state_dict - - -def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, - tag=None, - exclude_frozen_parameters=False, - lazy_mode=False): - """ - Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with - ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example - via a model hub. - - Args: - - ``checkpoint_dir``: path to the desired checkpoint folder - - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14`` - - ``exclude_frozen_parameters``: exclude frozen parameters - - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient. - Convert the pesduo tensor to torch tensor by ``.contiguous()`` - - Returns: - - pytorch ``state_dict`` - - A typical usage might be :: - - from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint - # do the training and checkpoint saving - state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu - model = model.cpu() # move to cpu - model.load_state_dict(state_dict) - # submit to model hub or save the model to share with others - - In this example the ``model`` will no longer be usable in the deepspeed context of the same - application. i.e. you will need to re-initialize the deepspeed engine, since - ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it. - - If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead. - - Note: the above usage may not work if your application doesn't have sufficient free CPU memory. - You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with - the checkpoint. Or you can load state_dict in lazy mode :: - - from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint - state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu - for name, lazy_tensor in state_dict.item(): - tensor = lazy_tensor.contiguous() # to cpu - print(name, tensor) - # del tensor to release memory if it no longer in use - """ - if tag is None: - latest_path = os.path.join(checkpoint_dir, 'latest') - if os.path.isfile(latest_path): - with open(latest_path, 'r') as fd: - tag = fd.read().strip() - else: - raise ValueError(f"Unable to find 'latest' file at {latest_path}") - - ds_checkpoint_dir = os.path.join(checkpoint_dir, tag) - - if not os.path.isdir(ds_checkpoint_dir): - raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist") - - state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters) - if lazy_mode: - return state_dict - else: - return to_torch_tensor(state_dict) - - -def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, - output_dir, - max_shard_size="5GB", - safe_serialization=False, - tag=None, - exclude_frozen_parameters=False): - """ - Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be - loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed. - - Args: - - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``) - - ``output_dir``: directory to the pytorch fp32 state_dict output files - - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB - - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). - - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14`` - - ``exclude_frozen_parameters``: exclude frozen parameters - """ - - # Dependency pre-check - if safe_serialization: - try: - from safetensors.torch import save_file - except ImportError: - print('If you want to use `safe_serialization`, please `pip install safetensors`') - raise - if max_shard_size is not None: - try: - from huggingface_hub import split_torch_state_dict_into_shards - except ImportError: - print('If you want to use `max_shard_size`, please `pip install huggingface_hub`') - raise - - # Convert zero checkpoint to state_dict - state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, - tag, - exclude_frozen_parameters, - lazy_mode=True) - - # Shard the model if it is too big. - weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin" - if max_shard_size is not None: - filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors") - # an memory-efficient approach for sharding - empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True) - state_dict_split = split_torch_state_dict_into_shards(empty_state_dict, - filename_pattern=filename_pattern, - max_shard_size=max_shard_size) - else: - from collections import namedtuple - StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"]) - state_dict_split = StateDictSplit(is_sharded=False, - filename_to_tensors={weights_name: list(state_dict.keys())}) - - # Save the model by shard - os.makedirs(output_dir, exist_ok=True) - filename_to_tensors = state_dict_split.filename_to_tensors.items() - for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"): - shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors} - shard_state_dict = to_torch_tensor(shard_state_dict) - output_path = os.path.join(output_dir, shard_file) - if safe_serialization: - save_file(shard_state_dict, output_path, metadata={"format": "pt"}) - else: - torch.save(shard_state_dict, output_path) - # release the memory of current shard - for tensor_name in list(shard_state_dict.keys()): - del state_dict[tensor_name] - del shard_state_dict[tensor_name] - del shard_state_dict - gc.collect() - - # Save index if sharded - if state_dict_split.is_sharded: - index = { - "metadata": state_dict_split.metadata, - "weight_map": state_dict_split.tensor_to_filename, - } - save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json" - save_index_file = os.path.join(output_dir, save_index_file) - with open(save_index_file, "w", encoding="utf-8") as f: - content = json.dumps(index, indent=2, sort_keys=True) + "\n" - f.write(content) - - -def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None): - """ - 1. Put the provided model to cpu - 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` - 3. Load it into the provided model - - Args: - - ``model``: the model object to update - - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``) - - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14`` - - Returns: - - ``model`: modified model - - Make sure you have plenty of CPU memory available before you call this function. If you don't - have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it - conveniently placed for you in the checkpoint folder. - - A typical usage might be :: - - from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint - model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir) - # submit to model hub or save the model to share with others - - Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context - of the same application. i.e. you will need to re-initialize the deepspeed engine, since - ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it. - - """ - logger.info(f"Extracting fp32 weights") - state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag) - - logger.info(f"Overwriting model with fp32 weights") - model = model.cpu() - model.load_state_dict(state_dict, strict=False) - - return model - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("checkpoint_dir", - type=str, - help="path to the desired checkpoint folder, e.g., path/checkpoint-12") - parser.add_argument("output_dir", - type=str, - help="directory to the pytorch fp32 state_dict output files" - "(e.g. path/checkpoint-12-output/)") - parser.add_argument( - "--max_shard_size", - type=str, - default="5GB", - help="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 `5MB`" - "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances" - "without CPU OOM issues.") - parser.add_argument( - "--safe_serialization", - default=False, - action='store_true', - help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).") - parser.add_argument("-t", - "--tag", - type=str, - default=None, - help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1") - parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters") - parser.add_argument("-d", "--debug", action='store_true', help="enable debug") - args = parser.parse_args() - - debug = args.debug - - convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, - args.output_dir, - max_shard_size=args.max_shard_size, - safe_serialization=args.safe_serialization, - tag=args.tag, - exclude_frozen_parameters=args.exclude_frozen_parameters) diff --git a/venv/lib/python3.10/site-packages/deprecated/__init__.py b/venv/lib/python3.10/site-packages/deprecated/__init__.py deleted file mode 100644 index 4cd847de3ef6d79ecf4612943c4f42aeabf2c883..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deprecated/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -# -*- coding: utf-8 -*- -""" -Deprecated Library -================== - -Python ``@deprecated`` decorator to deprecate old python classes, functions or methods. - -""" - -__version__ = "1.2.18" -__author__ = u"Laurent LAPORTE " -__date__ = "2025-01-27" -__credits__ = "(c) Laurent LAPORTE" - -from deprecated.classic import deprecated diff --git a/venv/lib/python3.10/site-packages/deprecated/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/deprecated/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index d7b7b7a4886aa817e6acf3624e3ae073c7e7302a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deprecated/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deprecated/__pycache__/classic.cpython-310.pyc b/venv/lib/python3.10/site-packages/deprecated/__pycache__/classic.cpython-310.pyc deleted file mode 100644 index fd98600d70dd2083c6e25bf84569cce97a87666e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deprecated/__pycache__/classic.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deprecated/__pycache__/sphinx.cpython-310.pyc b/venv/lib/python3.10/site-packages/deprecated/__pycache__/sphinx.cpython-310.pyc deleted file mode 100644 index cb0227e5fbc4c48abc604356c36c12fd5c1fb32f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/deprecated/__pycache__/sphinx.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/deprecated/classic.py b/venv/lib/python3.10/site-packages/deprecated/classic.py deleted file mode 100644 index 84f683cef469215653730a74f5861ecc52b4d7a2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deprecated/classic.py +++ /dev/null @@ -1,301 +0,0 @@ -# -*- coding: utf-8 -*- -""" -Classic deprecation warning -=========================== - -Classic ``@deprecated`` decorator to deprecate old python classes, functions or methods. - -.. _The Warnings Filter: https://docs.python.org/3/library/warnings.html#the-warnings-filter -""" -import functools -import inspect -import platform -import warnings - -import wrapt - -try: - # If the C extension for wrapt was compiled and wrapt/_wrappers.pyd exists, then the - # stack level that should be passed to warnings.warn should be 2. However, if using - # a pure python wrapt, an extra stacklevel is required. - import wrapt._wrappers - - _routine_stacklevel = 2 - _class_stacklevel = 2 -except ImportError: - _routine_stacklevel = 3 - if platform.python_implementation() == "PyPy": - _class_stacklevel = 2 - else: - _class_stacklevel = 3 - -string_types = (type(b''), type(u'')) - - -class ClassicAdapter(wrapt.AdapterFactory): - """ - Classic adapter -- *for advanced usage only* - - This adapter is used to get the deprecation message according to the wrapped object type: - class, function, standard method, static method, or class method. - - This is the base class of the :class:`~deprecated.sphinx.SphinxAdapter` class - which is used to update the wrapped object docstring. - - You can also inherit this class to change the deprecation message. - - In the following example, we change the message into "The ... is deprecated.": - - .. code-block:: python - - import inspect - - from deprecated.classic import ClassicAdapter - from deprecated.classic import deprecated - - - class MyClassicAdapter(ClassicAdapter): - def get_deprecated_msg(self, wrapped, instance): - if instance is None: - if inspect.isclass(wrapped): - fmt = "The class {name} is deprecated." - else: - fmt = "The function {name} is deprecated." - else: - if inspect.isclass(instance): - fmt = "The class method {name} is deprecated." - else: - fmt = "The method {name} is deprecated." - if self.reason: - fmt += " ({reason})" - if self.version: - fmt += " -- Deprecated since version {version}." - return fmt.format(name=wrapped.__name__, - reason=self.reason or "", - version=self.version or "") - - Then, you can use your ``MyClassicAdapter`` class like this in your source code: - - .. code-block:: python - - @deprecated(reason="use another function", adapter_cls=MyClassicAdapter) - def some_old_function(x, y): - return x + y - """ - - def __init__(self, reason="", version="", action=None, category=DeprecationWarning, extra_stacklevel=0): - """ - Construct a wrapper adapter. - - :type reason: str - :param reason: - Reason message which documents the deprecation in your library (can be omitted). - - :type version: str - :param version: - Version of your project which deprecates this feature. - If you follow the `Semantic Versioning `_, - the version number has the format "MAJOR.MINOR.PATCH". - - :type action: Literal["default", "error", "ignore", "always", "module", "once"] - :param action: - A warning filter used to activate or not the deprecation warning. - Can be one of "error", "ignore", "always", "default", "module", or "once". - If ``None`` or empty, the global filtering mechanism is used. - See: `The Warnings Filter`_ in the Python documentation. - - :type category: Type[Warning] - :param category: - The warning category to use for the deprecation warning. - By default, the category class is :class:`~DeprecationWarning`, - you can inherit this class to define your own deprecation warning category. - - :type extra_stacklevel: int - :param extra_stacklevel: - Number of additional stack levels to consider instrumentation rather than user code. - With the default value of 0, the warning refers to where the class was instantiated - or the function was called. - - .. versionchanged:: 1.2.15 - Add the *extra_stacklevel* parameter. - """ - self.reason = reason or "" - self.version = version or "" - self.action = action - self.category = category - self.extra_stacklevel = extra_stacklevel - super(ClassicAdapter, self).__init__() - - def get_deprecated_msg(self, wrapped, instance): - """ - Get the deprecation warning message for the user. - - :param wrapped: Wrapped class or function. - - :param instance: The object to which the wrapped function was bound when it was called. - - :return: The warning message. - """ - if instance is None: - if inspect.isclass(wrapped): - fmt = "Call to deprecated class {name}." - else: - fmt = "Call to deprecated function (or staticmethod) {name}." - else: - if inspect.isclass(instance): - fmt = "Call to deprecated class method {name}." - else: - fmt = "Call to deprecated method {name}." - if self.reason: - fmt += " ({reason})" - if self.version: - fmt += " -- Deprecated since version {version}." - return fmt.format(name=wrapped.__name__, reason=self.reason or "", version=self.version or "") - - def __call__(self, wrapped): - """ - Decorate your class or function. - - :param wrapped: Wrapped class or function. - - :return: the decorated class or function. - - .. versionchanged:: 1.2.4 - Don't pass arguments to :meth:`object.__new__` (other than *cls*). - - .. versionchanged:: 1.2.8 - The warning filter is not set if the *action* parameter is ``None`` or empty. - """ - if inspect.isclass(wrapped): - old_new1 = wrapped.__new__ - - def wrapped_cls(cls, *args, **kwargs): - msg = self.get_deprecated_msg(wrapped, None) - stacklevel = _class_stacklevel + self.extra_stacklevel - if self.action: - with warnings.catch_warnings(): - warnings.simplefilter(self.action, self.category) - warnings.warn(msg, category=self.category, stacklevel=stacklevel) - else: - warnings.warn(msg, category=self.category, stacklevel=stacklevel) - if old_new1 is object.__new__: - return old_new1(cls) - # actually, we don't know the real signature of *old_new1* - return old_new1(cls, *args, **kwargs) - - wrapped.__new__ = staticmethod(wrapped_cls) - - elif inspect.isroutine(wrapped): - @wrapt.decorator - def wrapper_function(wrapped_, instance_, args_, kwargs_): - msg = self.get_deprecated_msg(wrapped_, instance_) - stacklevel = _routine_stacklevel + self.extra_stacklevel - if self.action: - with warnings.catch_warnings(): - warnings.simplefilter(self.action, self.category) - warnings.warn(msg, category=self.category, stacklevel=stacklevel) - else: - warnings.warn(msg, category=self.category, stacklevel=stacklevel) - return wrapped_(*args_, **kwargs_) - - return wrapper_function(wrapped) - - else: - raise TypeError(repr(type(wrapped))) - - return wrapped - - -def deprecated(*args, **kwargs): - """ - This is a decorator which can be used to mark functions - as deprecated. It will result in a warning being emitted - when the function is used. - - **Classic usage:** - - To use this, decorate your deprecated function with **@deprecated** decorator: - - .. code-block:: python - - from deprecated import deprecated - - - @deprecated - def some_old_function(x, y): - return x + y - - You can also decorate a class or a method: - - .. code-block:: python - - from deprecated import deprecated - - - class SomeClass(object): - @deprecated - def some_old_method(self, x, y): - return x + y - - - @deprecated - class SomeOldClass(object): - pass - - You can give a *reason* message to help the developer to choose another function/class, - and a *version* number to specify the starting version number of the deprecation. - - .. code-block:: python - - from deprecated import deprecated - - - @deprecated(reason="use another function", version='1.2.0') - def some_old_function(x, y): - return x + y - - The *category* keyword argument allow you to specify the deprecation warning class of your choice. - By default, :exc:`DeprecationWarning` is used, but you can choose :exc:`FutureWarning`, - :exc:`PendingDeprecationWarning` or a custom subclass. - - .. code-block:: python - - from deprecated import deprecated - - - @deprecated(category=PendingDeprecationWarning) - def some_old_function(x, y): - return x + y - - The *action* keyword argument allow you to locally change the warning filtering. - *action* can be one of "error", "ignore", "always", "default", "module", or "once". - If ``None``, empty or missing, the global filtering mechanism is used. - See: `The Warnings Filter`_ in the Python documentation. - - .. code-block:: python - - from deprecated import deprecated - - - @deprecated(action="error") - def some_old_function(x, y): - return x + y - - The *extra_stacklevel* keyword argument allows you to specify additional stack levels - to consider instrumentation rather than user code. With the default value of 0, the - warning refers to where the class was instantiated or the function was called. - """ - if args and isinstance(args[0], string_types): - kwargs['reason'] = args[0] - args = args[1:] - - if args and not callable(args[0]): - raise TypeError(repr(type(args[0]))) - - if args: - adapter_cls = kwargs.pop('adapter_cls', ClassicAdapter) - adapter = adapter_cls(**kwargs) - wrapped = args[0] - return adapter(wrapped) - - return functools.partial(deprecated, **kwargs) diff --git a/venv/lib/python3.10/site-packages/deprecated/sphinx.py b/venv/lib/python3.10/site-packages/deprecated/sphinx.py deleted file mode 100644 index 90435c64f38864cbf8ede5e60da2fadc4cf54b44..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/deprecated/sphinx.py +++ /dev/null @@ -1,281 +0,0 @@ -# coding: utf-8 -""" -Sphinx directive integration -============================ - -We usually need to document the life-cycle of functions and classes: -when they are created, modified or deprecated. - -To do that, `Sphinx `_ has a set -of `Paragraph-level markups `_: - -- ``versionadded``: to document the version of the project which added the described feature to the library, -- ``versionchanged``: to document changes of a feature, -- ``deprecated``: to document a deprecated feature. - -The purpose of this module is to defined decorators which adds this Sphinx directives -to the docstring of your function and classes. - -Of course, the ``@deprecated`` decorator will emit a deprecation warning -when the function/method is called or the class is constructed. -""" -import re -import textwrap - -from deprecated.classic import ClassicAdapter -from deprecated.classic import deprecated as _classic_deprecated - - -class SphinxAdapter(ClassicAdapter): - """ - Sphinx adapter -- *for advanced usage only* - - This adapter override the :class:`~deprecated.classic.ClassicAdapter` - in order to add the Sphinx directives to the end of the function/class docstring. - Such a directive is a `Paragraph-level markup `_ - - - The directive can be one of "versionadded", "versionchanged" or "deprecated". - - The version number is added if provided. - - The reason message is obviously added in the directive block if not empty. - """ - - def __init__( - self, - directive, - reason="", - version="", - action=None, - category=DeprecationWarning, - extra_stacklevel=0, - line_length=70, - ): - """ - Construct a wrapper adapter. - - :type directive: str - :param directive: - Sphinx directive: can be one of "versionadded", "versionchanged" or "deprecated". - - :type reason: str - :param reason: - Reason message which documents the deprecation in your library (can be omitted). - - :type version: str - :param version: - Version of your project which deprecates this feature. - If you follow the `Semantic Versioning `_, - the version number has the format "MAJOR.MINOR.PATCH". - - :type action: Literal["default", "error", "ignore", "always", "module", "once"] - :param action: - A warning filter used to activate or not the deprecation warning. - Can be one of "error", "ignore", "always", "default", "module", or "once". - If ``None`` or empty, the global filtering mechanism is used. - See: `The Warnings Filter`_ in the Python documentation. - - :type category: Type[Warning] - :param category: - The warning category to use for the deprecation warning. - By default, the category class is :class:`~DeprecationWarning`, - you can inherit this class to define your own deprecation warning category. - - :type extra_stacklevel: int - :param extra_stacklevel: - Number of additional stack levels to consider instrumentation rather than user code. - With the default value of 0, the warning refers to where the class was instantiated - or the function was called. - - :type line_length: int - :param line_length: - Max line length of the directive text. If non nul, a long text is wrapped in several lines. - - .. versionchanged:: 1.2.15 - Add the *extra_stacklevel* parameter. - """ - if not version: - # https://github.com/laurent-laporte-pro/deprecated/issues/40 - raise ValueError("'version' argument is required in Sphinx directives") - self.directive = directive - self.line_length = line_length - super(SphinxAdapter, self).__init__( - reason=reason, version=version, action=action, category=category, extra_stacklevel=extra_stacklevel - ) - - def __call__(self, wrapped): - """ - Add the Sphinx directive to your class or function. - - :param wrapped: Wrapped class or function. - - :return: the decorated class or function. - """ - # -- build the directive division - fmt = ".. {directive}:: {version}" if self.version else ".. {directive}::" - div_lines = [fmt.format(directive=self.directive, version=self.version)] - width = self.line_length - 3 if self.line_length > 3 else 2**16 - reason = textwrap.dedent(self.reason).strip() - for paragraph in reason.splitlines(): - if paragraph: - div_lines.extend( - textwrap.fill( - paragraph, - width=width, - initial_indent=" ", - subsequent_indent=" ", - ).splitlines() - ) - else: - div_lines.append("") - - # -- get the docstring, normalize the trailing newlines - # keep a consistent behaviour if the docstring starts with newline or directly on the first one - docstring = wrapped.__doc__ or "" - lines = docstring.splitlines(True) or [""] - docstring = textwrap.dedent("".join(lines[1:])) if len(lines) > 1 else "" - docstring = lines[0] + docstring - if docstring: - # An empty line must separate the original docstring and the directive. - docstring = re.sub(r"\n+$", "", docstring, flags=re.DOTALL) + "\n\n" - else: - # Avoid "Explicit markup ends without a blank line" when the decorated function has no docstring - docstring = "\n" - - # -- append the directive division to the docstring - docstring += "".join("{}\n".format(line) for line in div_lines) - - wrapped.__doc__ = docstring - if self.directive in {"versionadded", "versionchanged"}: - return wrapped - return super(SphinxAdapter, self).__call__(wrapped) - - def get_deprecated_msg(self, wrapped, instance): - """ - Get the deprecation warning message (without Sphinx cross-referencing syntax) for the user. - - :param wrapped: Wrapped class or function. - - :param instance: The object to which the wrapped function was bound when it was called. - - :return: The warning message. - - .. versionadded:: 1.2.12 - Strip Sphinx cross-referencing syntax from warning message. - - """ - msg = super(SphinxAdapter, self).get_deprecated_msg(wrapped, instance) - # Strip Sphinx cross-reference syntax (like ":function:", ":py:func:" and ":py:meth:") - # Possible values are ":role:`foo`", ":domain:role:`foo`" - # where ``role`` and ``domain`` should match "[a-zA-Z]+" - msg = re.sub(r"(?: : [a-zA-Z]+ )? : [a-zA-Z]+ : (`[^`]*`)", r"\1", msg, flags=re.X) - return msg - - -def versionadded(reason="", version="", line_length=70): - """ - This decorator can be used to insert a "versionadded" directive - in your function/class docstring in order to document the - version of the project which adds this new functionality in your library. - - :param str reason: - Reason message which documents the addition in your library (can be omitted). - - :param str version: - Version of your project which adds this feature. - If you follow the `Semantic Versioning `_, - the version number has the format "MAJOR.MINOR.PATCH", and, - in the case of a new functionality, the "PATCH" component should be "0". - - :type line_length: int - :param line_length: - Max line length of the directive text. If non nul, a long text is wrapped in several lines. - - :return: the decorated function. - """ - adapter = SphinxAdapter( - 'versionadded', - reason=reason, - version=version, - line_length=line_length, - ) - return adapter - - -def versionchanged(reason="", version="", line_length=70): - """ - This decorator can be used to insert a "versionchanged" directive - in your function/class docstring in order to document the - version of the project which modifies this functionality in your library. - - :param str reason: - Reason message which documents the modification in your library (can be omitted). - - :param str version: - Version of your project which modifies this feature. - If you follow the `Semantic Versioning `_, - the version number has the format "MAJOR.MINOR.PATCH". - - :type line_length: int - :param line_length: - Max line length of the directive text. If non nul, a long text is wrapped in several lines. - - :return: the decorated function. - """ - adapter = SphinxAdapter( - 'versionchanged', - reason=reason, - version=version, - line_length=line_length, - ) - return adapter - - -def deprecated(reason="", version="", line_length=70, **kwargs): - """ - This decorator can be used to insert a "deprecated" directive - in your function/class docstring in order to document the - version of the project which deprecates this functionality in your library. - - :param str reason: - Reason message which documents the deprecation in your library (can be omitted). - - :param str version: - Version of your project which deprecates this feature. - If you follow the `Semantic Versioning `_, - the version number has the format "MAJOR.MINOR.PATCH". - - :type line_length: int - :param line_length: - Max line length of the directive text. If non nul, a long text is wrapped in several lines. - - Keyword arguments can be: - - - "action": - A warning filter used to activate or not the deprecation warning. - Can be one of "error", "ignore", "always", "default", "module", or "once". - If ``None``, empty or missing, the global filtering mechanism is used. - - - "category": - The warning category to use for the deprecation warning. - By default, the category class is :class:`~DeprecationWarning`, - you can inherit this class to define your own deprecation warning category. - - - "extra_stacklevel": - Number of additional stack levels to consider instrumentation rather than user code. - With the default value of 0, the warning refers to where the class was instantiated - or the function was called. - - - :return: a decorator used to deprecate a function. - - .. versionchanged:: 1.2.13 - Change the signature of the decorator to reflect the valid use cases. - - .. versionchanged:: 1.2.15 - Add the *extra_stacklevel* parameter. - """ - directive = kwargs.pop('directive', 'deprecated') - adapter_cls = kwargs.pop('adapter_cls', SphinxAdapter) - kwargs["reason"] = reason - kwargs["version"] = version - kwargs["line_length"] = line_length - return _classic_deprecated(directive=directive, adapter_cls=adapter_cls, **kwargs) diff --git a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/LICENSE b/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/LICENSE deleted file mode 100644 index 261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - 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 - 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 - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - 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. diff --git a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/METADATA b/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/METADATA deleted file mode 100644 index 5abe77c7a776232b22cbd4343c1864141e38a49f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/METADATA +++ /dev/null @@ -1,372 +0,0 @@ -Metadata-Version: 2.2 -Name: diffusers -Version: 0.34.0 -Summary: State-of-the-art diffusion in PyTorch and JAX. -Home-page: https://github.com/huggingface/diffusers -Author: The Hugging Face team (past and future) with the help of all our contributors (https://github.com/huggingface/diffusers/graphs/contributors) -Author-email: diffusers@huggingface.co -License: Apache 2.0 License -Keywords: deep learning diffusion jax pytorch stable diffusion audioldm -Classifier: Development Status :: 5 - Production/Stable -Classifier: Intended Audience :: Developers -Classifier: Intended Audience :: Education -Classifier: Intended Audience :: Science/Research -Classifier: License :: OSI Approved :: Apache Software License -Classifier: Operating System :: OS Independent -Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence -Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3.8 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Programming Language :: Python :: 3.10 -Requires-Python: >=3.8.0 -Description-Content-Type: text/markdown -License-File: LICENSE -Requires-Dist: importlib_metadata -Requires-Dist: filelock -Requires-Dist: huggingface-hub>=0.27.0 -Requires-Dist: numpy -Requires-Dist: regex!=2019.12.17 -Requires-Dist: requests -Requires-Dist: safetensors>=0.3.1 -Requires-Dist: Pillow -Provides-Extra: quality -Requires-Dist: urllib3<=2.0.0; extra == "quality" -Requires-Dist: isort>=5.5.4; extra == "quality" -Requires-Dist: ruff==0.9.10; extra == "quality" -Requires-Dist: hf-doc-builder>=0.3.0; extra == "quality" -Provides-Extra: docs -Requires-Dist: hf-doc-builder>=0.3.0; extra == "docs" -Provides-Extra: training -Requires-Dist: accelerate>=0.31.0; extra == "training" -Requires-Dist: datasets; extra == "training" -Requires-Dist: protobuf<4,>=3.20.3; extra == "training" -Requires-Dist: tensorboard; extra == "training" -Requires-Dist: Jinja2; extra == "training" -Requires-Dist: peft>=0.15.0; extra == "training" -Provides-Extra: test -Requires-Dist: compel==0.1.8; extra == "test" -Requires-Dist: GitPython<3.1.19; extra == "test" -Requires-Dist: datasets; extra == "test" -Requires-Dist: Jinja2; extra == "test" -Requires-Dist: invisible-watermark>=0.2.0; extra == "test" -Requires-Dist: k-diffusion>=0.0.12; extra == "test" -Requires-Dist: librosa; extra == "test" -Requires-Dist: parameterized; extra == "test" -Requires-Dist: pytest; extra == "test" -Requires-Dist: pytest-timeout; extra == "test" -Requires-Dist: pytest-xdist; extra == "test" -Requires-Dist: requests-mock==1.10.0; extra == "test" -Requires-Dist: safetensors>=0.3.1; extra == "test" -Requires-Dist: sentencepiece!=0.1.92,>=0.1.91; extra == "test" -Requires-Dist: scipy; extra == "test" -Requires-Dist: tiktoken>=0.7.0; extra == "test" -Requires-Dist: torchvision; extra == "test" -Requires-Dist: transformers>=4.41.2; extra == "test" -Requires-Dist: phonemizer; extra == "test" -Provides-Extra: torch -Requires-Dist: torch>=1.4; extra == "torch" -Requires-Dist: accelerate>=0.31.0; extra == "torch" -Provides-Extra: bitsandbytes -Requires-Dist: bitsandbytes>=0.43.3; extra == "bitsandbytes" -Requires-Dist: accelerate>=0.31.0; extra == "bitsandbytes" -Provides-Extra: gguf -Requires-Dist: gguf>=0.10.0; extra == "gguf" -Requires-Dist: accelerate>=0.31.0; extra == "gguf" -Provides-Extra: optimum-quanto -Requires-Dist: optimum_quanto>=0.2.6; extra == "optimum-quanto" -Requires-Dist: accelerate>=0.31.0; extra == "optimum-quanto" -Provides-Extra: torchao -Requires-Dist: torchao>=0.7.0; extra == "torchao" -Requires-Dist: accelerate>=0.31.0; extra == "torchao" -Provides-Extra: flax -Requires-Dist: jax>=0.4.1; extra == "flax" -Requires-Dist: jaxlib>=0.4.1; extra == "flax" -Requires-Dist: flax>=0.4.1; extra == "flax" -Provides-Extra: dev -Requires-Dist: urllib3<=2.0.0; extra == "dev" -Requires-Dist: isort>=5.5.4; extra == "dev" -Requires-Dist: ruff==0.9.10; extra == "dev" -Requires-Dist: hf-doc-builder>=0.3.0; extra == "dev" -Requires-Dist: compel==0.1.8; extra == "dev" -Requires-Dist: GitPython<3.1.19; extra == "dev" -Requires-Dist: datasets; extra == "dev" -Requires-Dist: Jinja2; extra == "dev" -Requires-Dist: invisible-watermark>=0.2.0; extra == "dev" -Requires-Dist: k-diffusion>=0.0.12; extra == "dev" -Requires-Dist: librosa; extra == "dev" -Requires-Dist: parameterized; extra == "dev" -Requires-Dist: pytest; extra == "dev" -Requires-Dist: pytest-timeout; extra == "dev" -Requires-Dist: pytest-xdist; extra == "dev" -Requires-Dist: requests-mock==1.10.0; extra == "dev" -Requires-Dist: safetensors>=0.3.1; extra == "dev" -Requires-Dist: sentencepiece!=0.1.92,>=0.1.91; extra == "dev" -Requires-Dist: scipy; extra == "dev" -Requires-Dist: tiktoken>=0.7.0; extra == "dev" -Requires-Dist: torchvision; extra == "dev" -Requires-Dist: transformers>=4.41.2; extra == "dev" -Requires-Dist: phonemizer; extra == "dev" -Requires-Dist: accelerate>=0.31.0; extra == "dev" -Requires-Dist: datasets; extra == "dev" -Requires-Dist: protobuf<4,>=3.20.3; extra == "dev" -Requires-Dist: tensorboard; extra == "dev" -Requires-Dist: Jinja2; extra == "dev" -Requires-Dist: peft>=0.15.0; extra == "dev" -Requires-Dist: hf-doc-builder>=0.3.0; extra == "dev" -Requires-Dist: torch>=1.4; extra == "dev" -Requires-Dist: accelerate>=0.31.0; extra == "dev" -Requires-Dist: jax>=0.4.1; extra == "dev" -Requires-Dist: jaxlib>=0.4.1; extra == "dev" -Requires-Dist: flax>=0.4.1; extra == "dev" -Dynamic: author -Dynamic: author-email -Dynamic: classifier -Dynamic: description -Dynamic: description-content-type -Dynamic: home-page -Dynamic: keywords -Dynamic: license -Dynamic: provides-extra -Dynamic: requires-dist -Dynamic: requires-python -Dynamic: summary - - - -

-
- -
-

-

- GitHub - GitHub release - GitHub release - Contributor Covenant - X account -

- -🤗 Diffusers is the go-to library for state-of-the-art pretrained diffusion models for generating images, audio, and even 3D structures of molecules. Whether you're looking for a simple inference solution or training your own diffusion models, 🤗 Diffusers is a modular toolbox that supports both. Our library is designed with a focus on [usability over performance](https://huggingface.co/docs/diffusers/conceptual/philosophy#usability-over-performance), [simple over easy](https://huggingface.co/docs/diffusers/conceptual/philosophy#simple-over-easy), and [customizability over abstractions](https://huggingface.co/docs/diffusers/conceptual/philosophy#tweakable-contributorfriendly-over-abstraction). - -🤗 Diffusers offers three core components: - -- State-of-the-art [diffusion pipelines](https://huggingface.co/docs/diffusers/api/pipelines/overview) that can be run in inference with just a few lines of code. -- Interchangeable noise [schedulers](https://huggingface.co/docs/diffusers/api/schedulers/overview) for different diffusion speeds and output quality. -- Pretrained [models](https://huggingface.co/docs/diffusers/api/models/overview) that can be used as building blocks, and combined with schedulers, for creating your own end-to-end diffusion systems. - -## Installation - -We recommend installing 🤗 Diffusers in a virtual environment from PyPI or Conda. For more details about installing [PyTorch](https://pytorch.org/get-started/locally/) and [Flax](https://flax.readthedocs.io/en/latest/#installation), please refer to their official documentation. - -### PyTorch - -With `pip` (official package): - -```bash -pip install --upgrade diffusers[torch] -``` - -With `conda` (maintained by the community): - -```sh -conda install -c conda-forge diffusers -``` - -### Flax - -With `pip` (official package): - -```bash -pip install --upgrade diffusers[flax] -``` - -### Apple Silicon (M1/M2) support - -Please refer to the [How to use Stable Diffusion in Apple Silicon](https://huggingface.co/docs/diffusers/optimization/mps) guide. - -## Quickstart - -Generating outputs is super easy with 🤗 Diffusers. To generate an image from text, use the `from_pretrained` method to load any pretrained diffusion model (browse the [Hub](https://huggingface.co/models?library=diffusers&sort=downloads) for 30,000+ checkpoints): - -```python -from diffusers import DiffusionPipeline -import torch - -pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16) -pipeline.to("cuda") -pipeline("An image of a squirrel in Picasso style").images[0] -``` - -You can also dig into the models and schedulers toolbox to build your own diffusion system: - -```python -from diffusers import DDPMScheduler, UNet2DModel -from PIL import Image -import torch - -scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256") -model = UNet2DModel.from_pretrained("google/ddpm-cat-256").to("cuda") -scheduler.set_timesteps(50) - -sample_size = model.config.sample_size -noise = torch.randn((1, 3, sample_size, sample_size), device="cuda") -input = noise - -for t in scheduler.timesteps: - with torch.no_grad(): - noisy_residual = model(input, t).sample - prev_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample - input = prev_noisy_sample - -image = (input / 2 + 0.5).clamp(0, 1) -image = image.cpu().permute(0, 2, 3, 1).numpy()[0] -image = Image.fromarray((image * 255).round().astype("uint8")) -image -``` - -Check out the [Quickstart](https://huggingface.co/docs/diffusers/quicktour) to launch your diffusion journey today! - -## How to navigate the documentation - -| **Documentation** | **What can I learn?** | -|---------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| [Tutorial](https://huggingface.co/docs/diffusers/tutorials/tutorial_overview) | A basic crash course for learning how to use the library's most important features like using models and schedulers to build your own diffusion system, and training your own diffusion model. | -| [Loading](https://huggingface.co/docs/diffusers/using-diffusers/loading) | Guides for how to load and configure all the components (pipelines, models, and schedulers) of the library, as well as how to use different schedulers. | -| [Pipelines for inference](https://huggingface.co/docs/diffusers/using-diffusers/overview_techniques) | Guides for how to use pipelines for different inference tasks, batched generation, controlling generated outputs and randomness, and how to contribute a pipeline to the library. | -| [Optimization](https://huggingface.co/docs/diffusers/optimization/fp16) | Guides for how to optimize your diffusion model to run faster and consume less memory. | -| [Training](https://huggingface.co/docs/diffusers/training/overview) | Guides for how to train a diffusion model for different tasks with different training techniques. | -## Contribution - -We ❤️ contributions from the open-source community! -If you want to contribute to this library, please check out our [Contribution guide](https://github.com/huggingface/diffusers/blob/main/CONTRIBUTING.md). -You can look out for [issues](https://github.com/huggingface/diffusers/issues) you'd like to tackle to contribute to the library. -- See [Good first issues](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22) for general opportunities to contribute -- See [New model/pipeline](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22New+pipeline%2Fmodel%22) to contribute exciting new diffusion models / diffusion pipelines -- See [New scheduler](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22New+scheduler%22) - -Also, say 👋 in our public Discord channel Join us on Discord. We discuss the hottest trends about diffusion models, help each other with contributions, personal projects or just hang out ☕. - - -## Popular Tasks & Pipelines - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
TaskPipeline🤗 Hub
Unconditional Image Generation DDPM google/ddpm-ema-church-256
Text-to-ImageStable Diffusion Text-to-Image stable-diffusion-v1-5/stable-diffusion-v1-5
Text-to-ImageunCLIP kakaobrain/karlo-v1-alpha
Text-to-ImageDeepFloyd IF DeepFloyd/IF-I-XL-v1.0
Text-to-ImageKandinsky kandinsky-community/kandinsky-2-2-decoder
Text-guided Image-to-ImageControlNet lllyasviel/sd-controlnet-canny
Text-guided Image-to-ImageInstructPix2Pix timbrooks/instruct-pix2pix
Text-guided Image-to-ImageStable Diffusion Image-to-Image stable-diffusion-v1-5/stable-diffusion-v1-5
Text-guided Image InpaintingStable Diffusion Inpainting runwayml/stable-diffusion-inpainting
Image VariationStable Diffusion Image Variation lambdalabs/sd-image-variations-diffusers
Super ResolutionStable Diffusion Upscale stabilityai/stable-diffusion-x4-upscaler
Super ResolutionStable Diffusion Latent Upscale stabilityai/sd-x2-latent-upscaler
- -## Popular libraries using 🧨 Diffusers - -- https://github.com/microsoft/TaskMatrix -- https://github.com/invoke-ai/InvokeAI -- https://github.com/InstantID/InstantID -- https://github.com/apple/ml-stable-diffusion -- https://github.com/Sanster/lama-cleaner -- https://github.com/IDEA-Research/Grounded-Segment-Anything -- https://github.com/ashawkey/stable-dreamfusion -- https://github.com/deep-floyd/IF -- https://github.com/bentoml/BentoML -- https://github.com/bmaltais/kohya_ss -- +14,000 other amazing GitHub repositories 💪 - -Thank you for using us ❤️. - -## Credits - -This library concretizes previous work by many different authors and would not have been possible without their great research and implementations. We'd like to thank, in particular, the following implementations which have helped us in our development and without which the API could not have been as polished today: - -- @CompVis' latent diffusion models library, available [here](https://github.com/CompVis/latent-diffusion) -- @hojonathanho original DDPM implementation, available [here](https://github.com/hojonathanho/diffusion) as well as the extremely useful translation into PyTorch by @pesser, available [here](https://github.com/pesser/pytorch_diffusion) -- @ermongroup's DDIM implementation, available [here](https://github.com/ermongroup/ddim) -- @yang-song's Score-VE and Score-VP implementations, available [here](https://github.com/yang-song/score_sde_pytorch) - -We also want to thank @heejkoo for the very helpful overview of papers, code and resources on diffusion models, available [here](https://github.com/heejkoo/Awesome-Diffusion-Models) as well as @crowsonkb and @rromb for useful discussions and insights. - -## Citation - -```bibtex -@misc{von-platen-etal-2022-diffusers, - author = {Patrick von Platen and Suraj Patil and Anton Lozhkov and Pedro Cuenca and Nathan Lambert and Kashif Rasul and Mishig Davaadorj and Dhruv Nair and Sayak Paul and William Berman and Yiyi Xu and Steven Liu and Thomas Wolf}, - title = {Diffusers: State-of-the-art diffusion models}, - year = {2022}, - publisher = {GitHub}, - journal = {GitHub repository}, - howpublished = {\url{https://github.com/huggingface/diffusers}} -} -``` diff --git a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/RECORD b/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/RECORD deleted file mode 100644 index 2065e508ceb59e540ff8fa8be1c0b042147a0c67..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/RECORD +++ /dev/null @@ -1,1273 +0,0 @@ -../../../bin/diffusers-cli,sha256=F0-KZOMEcKsoGXC_Y_BQwRcKF0b8z3YQHhxLuFhDhwE,249 -diffusers-0.34.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -diffusers-0.34.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357 -diffusers-0.34.0.dist-info/METADATA,sha256=9cVvp1z0cfI46dpVFR5F4UK3F2rL9hWyNLol9CH1r_w,20120 -diffusers-0.34.0.dist-info/RECORD,, -diffusers-0.34.0.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -diffusers-0.34.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91 -diffusers-0.34.0.dist-info/entry_points.txt,sha256=_1bvshKV_6_b63_FAkcUs9W6tUKGeIoQ3SHEZsovEWs,72 -diffusers-0.34.0.dist-info/top_level.txt,sha256=axJl2884vMSvhzrFrSoht36QXA_6gZN9cKtg4xOO72o,10 -diffusers/__init__.py,sha256=pDjLfKQ_7dmCkQT0crSvex7rBfURdqh0ZnjZWiLOhBE,46619 -diffusers/__pycache__/__init__.cpython-310.pyc,, -diffusers/__pycache__/callbacks.cpython-310.pyc,, -diffusers/__pycache__/configuration_utils.cpython-310.pyc,, -diffusers/__pycache__/dependency_versions_check.cpython-310.pyc,, -diffusers/__pycache__/dependency_versions_table.cpython-310.pyc,, -diffusers/__pycache__/image_processor.cpython-310.pyc,, -diffusers/__pycache__/optimization.cpython-310.pyc,, -diffusers/__pycache__/training_utils.cpython-310.pyc,, -diffusers/__pycache__/video_processor.cpython-310.pyc,, -diffusers/callbacks.py,sha256=wmLFSUrnEZ9gz6gay2pRWfha5cErTSAbB9PK482_KGY,8749 -diffusers/commands/__init__.py,sha256=KQXlWjcjH6qmwD5B6be8kpniSOPbyWZPTghdXrNquXU,920 -diffusers/commands/__pycache__/__init__.cpython-310.pyc,, -diffusers/commands/__pycache__/diffusers_cli.cpython-310.pyc,, -diffusers/commands/__pycache__/env.cpython-310.pyc,, -diffusers/commands/__pycache__/fp16_safetensors.cpython-310.pyc,, -diffusers/commands/diffusers_cli.py,sha256=bCLOzLKq2s8wj9AifPD6yEbT0gkpmgDM5Kvoiw9DPvQ,1317 -diffusers/commands/env.py,sha256=IZ2e2id8ZjNQ-BNEG3qt58TyEsM6Qs337kL5KdZmqZ0,6224 -diffusers/commands/fp16_safetensors.py,sha256=w6r3FnfiLHKQMpiVcs3vWIZIVENh_Z8um1nPN1HRqus,5423 -diffusers/configuration_utils.py,sha256=bPpagE2tmsaQjoksYORdOaSiMaSgaxoCx7i4iDwlA-g,34480 -diffusers/dependency_versions_check.py,sha256=PcT_deWuvIKrNkjkCnQKi0ZTWCl77tHC02lhttbqQHM,1271 -diffusers/dependency_versions_table.py,sha256=BKcEswVLknflk4QNnPDd9Q7HYAQd7qgj7h6dAtVuJDA,1766 -diffusers/experimental/__init__.py,sha256=0C9ExG0XYiGZuzFJkZuJ53K6Ix5ylF2kWe4PGASchtY,38 -diffusers/experimental/__pycache__/__init__.cpython-310.pyc,, -diffusers/experimental/rl/__init__.py,sha256=Gcoznw9rYjfMvswH0seXekKYDAAN1YXXxZ-RWMdzvrE,57 -diffusers/experimental/rl/__pycache__/__init__.cpython-310.pyc,, -diffusers/experimental/rl/__pycache__/value_guided_sampling.cpython-310.pyc,, -diffusers/experimental/rl/value_guided_sampling.py,sha256=a1F7YC--u1OgoGC8WGgi2pDUhYBI6XYnF2-RUzPD3dI,6033 -diffusers/hooks/__init__.py,sha256=sImnLjLQ2WnTbhCNu_NcawZJk4MUamWMTduczLqCJ7U,439 -diffusers/hooks/__pycache__/__init__.cpython-310.pyc,, -diffusers/hooks/__pycache__/faster_cache.cpython-310.pyc,, -diffusers/hooks/__pycache__/group_offloading.cpython-310.pyc,, -diffusers/hooks/__pycache__/hooks.cpython-310.pyc,, -diffusers/hooks/__pycache__/layerwise_casting.cpython-310.pyc,, -diffusers/hooks/__pycache__/pyramid_attention_broadcast.cpython-310.pyc,, -diffusers/hooks/faster_cache.py,sha256=Xe4m3fUqJItffff5F3g73uS04cd-scRDYiO4GXAp3fo,34950 -diffusers/hooks/group_offloading.py,sha256=T7R7eUSvjh0vp6NIp3-XFBVdNye23ACr5Hv-rPRDdV8,43294 -diffusers/hooks/hooks.py,sha256=rmbl-KsdsYkkoOCY7ySJMXNRvIUtyPmVGAs94h3Vrnc,8795 -diffusers/hooks/layerwise_casting.py,sha256=7vafpuCnKWKWN8hM0BehMMLNIBI8yIkdUge08cuLWrA,10589 -diffusers/hooks/pyramid_attention_broadcast.py,sha256=jGj67TTlk5H0WhtKep8ZkF9EckomPld6P7dQlyBL5kc,15760 -diffusers/image_processor.py,sha256=6qop7filTHdq3L-ZGVWmJMn0-0s4Bg7W-AaAFBNdO0Q,52868 -diffusers/loaders/__init__.py,sha256=6JRHLAuQPsNjEI38uTklH-xy4BvBzUqx6t-mf7ZXLcc,5280 -diffusers/loaders/__pycache__/__init__.cpython-310.pyc,, -diffusers/loaders/__pycache__/ip_adapter.cpython-310.pyc,, -diffusers/loaders/__pycache__/lora_base.cpython-310.pyc,, -diffusers/loaders/__pycache__/lora_conversion_utils.cpython-310.pyc,, -diffusers/loaders/__pycache__/lora_pipeline.cpython-310.pyc,, -diffusers/loaders/__pycache__/peft.cpython-310.pyc,, -diffusers/loaders/__pycache__/single_file.cpython-310.pyc,, -diffusers/loaders/__pycache__/single_file_model.cpython-310.pyc,, -diffusers/loaders/__pycache__/single_file_utils.cpython-310.pyc,, -diffusers/loaders/__pycache__/textual_inversion.cpython-310.pyc,, -diffusers/loaders/__pycache__/transformer_flux.cpython-310.pyc,, -diffusers/loaders/__pycache__/transformer_sd3.cpython-310.pyc,, -diffusers/loaders/__pycache__/unet.cpython-310.pyc,, -diffusers/loaders/__pycache__/unet_loader_utils.cpython-310.pyc,, -diffusers/loaders/__pycache__/utils.cpython-310.pyc,, -diffusers/loaders/ip_adapter.py,sha256=TB1HLgy0u6rKezWbPhXNQttmUNpq7nMNhWnW-ljlTIs,45068 -diffusers/loaders/lora_base.py,sha256=-DCuq6UAGlw0-fmP5vwYZzHNDhZN-NmyuB7VTDaUjEE,43666 -diffusers/loaders/lora_conversion_utils.py,sha256=h2nb6eL8JrRkTNDemOceW2jFPda-P0I95_yPvBk7gzs,90021 -diffusers/loaders/lora_pipeline.py,sha256=02HBNwDA9xMIQkVKeogbM8U7a_lfZvkpWON-iRXQoZ4,289014 -diffusers/loaders/peft.py,sha256=rFBo4wYogICBEsJC6y8euB3TSauI4sNC0QNe828pZws,38003 -diffusers/loaders/single_file.py,sha256=NyDO7tx2ASAENP5rXsWKxpOD320L24wl4QqpPixYYLE,25003 -diffusers/loaders/single_file_model.py,sha256=ZysRXed49QWEqZM2Uvz2jaNZNWIgxwn0BiAcD90Bm9Q,20083 -diffusers/loaders/single_file_utils.py,sha256=pas9vPmcPM0rH51RvrrUCQVffdIGcYoGobEM_tPolD0,156096 -diffusers/loaders/textual_inversion.py,sha256=Yh4Mw-3EAUWaQpDXptIMuNbyAYhNRir-irqoHSPvOx8,26925 -diffusers/loaders/transformer_flux.py,sha256=PKutmTlDU_T7DlssyGgSExeL8f6gw9lRz3JlnneWvSY,7915 -diffusers/loaders/transformer_sd3.py,sha256=LdPxicRUw_YelTZgJ_E9R4C2AlXpije_fgEG08kuWyQ,8482 -diffusers/loaders/unet.py,sha256=TX3a1PHxs0-IktVTU8mvoQZS_ZLQwYd-dbWyiSNecP0,45145 -diffusers/loaders/unet_loader_utils.py,sha256=tVuE5Xd6PiC_GIl_3SHfa_jDxfgn9rgqEZnYUjTCwBI,6220 -diffusers/loaders/utils.py,sha256=tknQHwXnDlc7piArCSDZjFFdgJI3DY6qfffsZlIxLnY,2423 -diffusers/models/__init__.py,sha256=dJpGFKbTJcRRcc7LeYsADK9sv7JGVqgFuY3LJxHYcJU,10601 -diffusers/models/__pycache__/__init__.cpython-310.pyc,, -diffusers/models/__pycache__/activations.cpython-310.pyc,, -diffusers/models/__pycache__/adapter.cpython-310.pyc,, -diffusers/models/__pycache__/attention.cpython-310.pyc,, -diffusers/models/__pycache__/attention_flax.cpython-310.pyc,, -diffusers/models/__pycache__/attention_processor.cpython-310.pyc,, -diffusers/models/__pycache__/auto_model.cpython-310.pyc,, -diffusers/models/__pycache__/cache_utils.cpython-310.pyc,, -diffusers/models/__pycache__/controlnet.cpython-310.pyc,, -diffusers/models/__pycache__/controlnet_flux.cpython-310.pyc,, -diffusers/models/__pycache__/controlnet_sd3.cpython-310.pyc,, -diffusers/models/__pycache__/controlnet_sparsectrl.cpython-310.pyc,, -diffusers/models/__pycache__/downsampling.cpython-310.pyc,, -diffusers/models/__pycache__/embeddings.cpython-310.pyc,, -diffusers/models/__pycache__/embeddings_flax.cpython-310.pyc,, -diffusers/models/__pycache__/lora.cpython-310.pyc,, -diffusers/models/__pycache__/model_loading_utils.cpython-310.pyc,, -diffusers/models/__pycache__/modeling_flax_pytorch_utils.cpython-310.pyc,, -diffusers/models/__pycache__/modeling_flax_utils.cpython-310.pyc,, -diffusers/models/__pycache__/modeling_outputs.cpython-310.pyc,, -diffusers/models/__pycache__/modeling_pytorch_flax_utils.cpython-310.pyc,, -diffusers/models/__pycache__/modeling_utils.cpython-310.pyc,, -diffusers/models/__pycache__/normalization.cpython-310.pyc,, -diffusers/models/__pycache__/resnet.cpython-310.pyc,, -diffusers/models/__pycache__/resnet_flax.cpython-310.pyc,, -diffusers/models/__pycache__/upsampling.cpython-310.pyc,, -diffusers/models/__pycache__/vae_flax.cpython-310.pyc,, -diffusers/models/__pycache__/vq_model.cpython-310.pyc,, -diffusers/models/activations.py,sha256=qxdn6OROfUvxyxgpm6M2VDKeJxH6mDsUI_xP4S3iw6s,6511 -diffusers/models/adapter.py,sha256=NDnqBqD53fg2fWWHt-LUHDyhuH6J-4R7PoStr2ggp-4,24507 -diffusers/models/attention.py,sha256=k3lQlDPoEUtuj-HBkxcMeQslATqCEKRQc_v8X6LkId8,53548 -diffusers/models/attention_flax.py,sha256=NJTCmsALDnRScOy2EG7r6fZGXaNrJGBVTHQxAluWZEs,20392 -diffusers/models/attention_processor.py,sha256=8B2iWzgEYfmgYcDyf-_07N0W_Dm7VsFQrjt_9EqUNcQ,263926 -diffusers/models/auto_model.py,sha256=7bLn0pgTf5TG2txeRLFgUgtYGmrylZB9eZZI6jC401s,11155 -diffusers/models/autoencoders/__init__.py,sha256=uEgP7tgXHRQdw8BSbJ8pEgqk7gnzyHsEgqKGt01txlk,854 -diffusers/models/autoencoders/__pycache__/__init__.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/autoencoder_asym_kl.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/autoencoder_dc.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/autoencoder_kl.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/autoencoder_kl_allegro.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/autoencoder_kl_cogvideox.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/autoencoder_kl_cosmos.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/autoencoder_kl_hunyuan_video.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/autoencoder_kl_ltx.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/autoencoder_kl_magvit.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/autoencoder_kl_mochi.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/autoencoder_kl_temporal_decoder.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/autoencoder_kl_wan.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/autoencoder_oobleck.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/autoencoder_tiny.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/consistency_decoder_vae.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/vae.cpython-310.pyc,, -diffusers/models/autoencoders/__pycache__/vq_model.cpython-310.pyc,, -diffusers/models/autoencoders/autoencoder_asym_kl.py,sha256=_67adzZxCMsV8N7s0JzOCSOz4_Kewb49c0CuA2KZS-I,7788 -diffusers/models/autoencoders/autoencoder_dc.py,sha256=Fvi2jKJQqmhyWgH9XZ58nJzq9WKVlYttL-l8P97I_aU,30503 -diffusers/models/autoencoders/autoencoder_kl.py,sha256=lHKSeG3m5tpfkDdcS0fg1UMqaZ4NVITulBjUJZ602_o,25079 -diffusers/models/autoencoders/autoencoder_kl_allegro.py,sha256=N1tLNJjqtCOGPQREU647RA5LdUsl4fgK2Uiii-mYx2c,45152 -diffusers/models/autoencoders/autoencoder_kl_cogvideox.py,sha256=YDCPc0NugDpHkw4JKBkrNHd0sOgxBbM0QXgovnAToTU,60406 -diffusers/models/autoencoders/autoencoder_kl_cosmos.py,sha256=nHasHRysFvmAbWteH1IscQ7z3rQGSoA15udZPWmSGLs,53706 -diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py,sha256=T4ijgoUTrJou5DMnUBkaCnPA7pMZuxutTSHxE31Y3zk,45481 -diffusers/models/autoencoders/autoencoder_kl_ltx.py,sha256=GQEIy5muXUndwYRH4jlCFUQanbgiUCHLNaEnVcL286g,64433 -diffusers/models/autoencoders/autoencoder_kl_magvit.py,sha256=6jrWs8thq0DERB2zOsGBIQHPtFFyjYuuekFYM-fhE9I,45098 -diffusers/models/autoencoders/autoencoder_kl_mochi.py,sha256=EJBSXMnIE4clpodfDKy4tffZbQkC15iYGmlZRowDQHA,46811 -diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py,sha256=qJ52TwMtQQz7QtczEsFLKXVceS7RLDRSvMMZI-RsjIY,14778 -diffusers/models/autoencoders/autoencoder_kl_wan.py,sha256=1xtFoijZaHBb1040yY9JO1aTGBOGce8yXmjLF5xcOrE,41721 -diffusers/models/autoencoders/autoencoder_oobleck.py,sha256=HV30_00e21-mts-JdT9K9Q-s8Uerk_cMD9_OX9uaNtU,17085 -diffusers/models/autoencoders/autoencoder_tiny.py,sha256=6kLUGAR0fDx7fwy5BiRLaUzkqJg9VQFKUJyisloxF4M,15853 -diffusers/models/autoencoders/consistency_decoder_vae.py,sha256=3uvrzIR2vEgj0etcRvUiBw8YmCQ5RfQ6RdhX4LcCoeM,19761 -diffusers/models/autoencoders/vae.py,sha256=XXm10FfQ9wLhozWomRO2ldfcb6qR85ePF-ShUxYq-6Y,32733 -diffusers/models/autoencoders/vq_model.py,sha256=wCvrranGbiCR8s4Qhq5qNSy_1BP9PGWvWyxvWJEHA8A,7912 -diffusers/models/cache_utils.py,sha256=A88pJJikRVSXZgiyyx4q36xdItvXvmA1DqDT6F1eBXo,4232 -diffusers/models/controlnet.py,sha256=RG8RSUtkCk77Dxw6e1nc2o6vdikTVKtskxgALI7ijsg,5774 -diffusers/models/controlnet_flux.py,sha256=bVej9TQJhszOMJOtHqzh2kwK1PLseuwiGSYdbT6b7f4,3442 -diffusers/models/controlnet_sd3.py,sha256=5cPugxXJrWroX4bjIG_DQGMlOWiFlwl2sMZz6pTSLTc,3377 -diffusers/models/controlnet_sparsectrl.py,sha256=ewfJCGXU-bvthUH2OP93RtMI3zZFNfnioYs8YSa-tmo,5904 -diffusers/models/controlnets/__init__.py,sha256=oKACrY23NZqmM0Hbl-S_nUuDFyQ4NesJMoKBYPzeK0Y,1058 -diffusers/models/controlnets/__pycache__/__init__.cpython-310.pyc,, -diffusers/models/controlnets/__pycache__/controlnet.cpython-310.pyc,, -diffusers/models/controlnets/__pycache__/controlnet_flax.cpython-310.pyc,, -diffusers/models/controlnets/__pycache__/controlnet_flux.cpython-310.pyc,, -diffusers/models/controlnets/__pycache__/controlnet_hunyuan.cpython-310.pyc,, -diffusers/models/controlnets/__pycache__/controlnet_sana.cpython-310.pyc,, -diffusers/models/controlnets/__pycache__/controlnet_sd3.cpython-310.pyc,, -diffusers/models/controlnets/__pycache__/controlnet_sparsectrl.cpython-310.pyc,, -diffusers/models/controlnets/__pycache__/controlnet_union.cpython-310.pyc,, -diffusers/models/controlnets/__pycache__/controlnet_xs.cpython-310.pyc,, -diffusers/models/controlnets/__pycache__/multicontrolnet.cpython-310.pyc,, -diffusers/models/controlnets/__pycache__/multicontrolnet_union.cpython-310.pyc,, -diffusers/models/controlnets/controlnet.py,sha256=dzIASI0h_bNeSjRi3BQTW2K4Jbo8-eQfJ7F5PeOzpsA,43149 -diffusers/models/controlnets/controlnet_flax.py,sha256=DekqQPnow0gSWDCUlIwcP6ezT16ppiYUGFVn1Ixi0zM,16715 -diffusers/models/controlnets/controlnet_flux.py,sha256=u_kd28VbTsro-FMc1j6rrPoTXvzQOoXkwPz6VzxPqx0,22889 -diffusers/models/controlnets/controlnet_hunyuan.py,sha256=UCyQL-GtHIQEeEu6bP04eIm-DDhsqG92I3hAJJE_m2g,16928 -diffusers/models/controlnets/controlnet_sana.py,sha256=2rYVE1EJAE-BavlAWtkkX4zcVk3nrP_BJkG8ewfOYXc,12519 -diffusers/models/controlnets/controlnet_sd3.py,sha256=E-r-_y2M3sN3sT4rQBm3CYKcTNrYTEvnIVUDqpjzj7c,23227 -diffusers/models/controlnets/controlnet_sparsectrl.py,sha256=ui450jzQEfI-OUU90nrmb1Oud4lcTmw7aiVMLtF3uTE,38357 -diffusers/models/controlnets/controlnet_union.py,sha256=mOK2L2xYcGivPqixZbqsj7Ev9JeTUY5BAG4Uv5f39SA,41604 -diffusers/models/controlnets/controlnet_xs.py,sha256=aF13vgEAeCf1UygSzuiktKzA_OrGXPiycRg3YzHKOHg,85416 -diffusers/models/controlnets/multicontrolnet.py,sha256=VoccQ64iwREXC04DseYLZjub2bdph9y8tBdga4br1XM,9332 -diffusers/models/controlnets/multicontrolnet_union.py,sha256=o_5BMCsvXUehw7Jni0Fn6e_qNm_66XnAyN4Ull4txrE,10164 -diffusers/models/downsampling.py,sha256=ss0au04TEj_eMz1vCgi20a1ZGoUEOLgC1zxeQGTN85s,15829 -diffusers/models/embeddings.py,sha256=8BXCFaefAd7LDMqX438v1_sk7i3Qezz8cnlp1hDZzPg,103970 -diffusers/models/embeddings_flax.py,sha256=FMuPD4wNU0wUXOaih81VcFjkKgOg4lyWi7hYIlmnYT4,4361 -diffusers/models/lora.py,sha256=kSZzH-taCGJCSH9PHYEcS1M1LXUGUpIti-taQxRRnS4,18851 -diffusers/models/model_loading_utils.py,sha256=khiNv_O8xz1VG6BvL_tnEzK0FfIpNn_6vQf9nyNrzKQ,21712 -diffusers/models/modeling_flax_pytorch_utils.py,sha256=f-j9Y-AhcrRp9UvLldjZOHPYBluMC5BXwxmAi6gS1rA,5332 -diffusers/models/modeling_flax_utils.py,sha256=j7BYsBBzYlUO7FUlWXs3Y3bWCTo8gtBL0kU0atuvyN4,26954 -diffusers/models/modeling_outputs.py,sha256=XH3sJO34MRW6UuWqqKo05mVqxGSBFRazpap_-YLwO2I,1042 -diffusers/models/modeling_pytorch_flax_utils.py,sha256=Hz5IoBV0vygRekVw8OZc9Jji22gyuEZ84uAWzXpYLvs,6973 -diffusers/models/modeling_utils.py,sha256=p5_UPBrCXIFQGt_Btarq4VN0E_7EO7z5Ja0yFQHsaFI,85234 -diffusers/models/normalization.py,sha256=Hke3BXA1GYAVBKrkJi7X8A0DLT6_NqP2dv2Q7lZFTRU,24616 -diffusers/models/resnet.py,sha256=DIFQHFSh8F0NcrAGjvxaca1fK8ZoArlpTtkTTLm50VU,32254 -diffusers/models/resnet_flax.py,sha256=aXKT-L0yMwFZz-xhSsbLsB5n7M7bNOE2se6RJhPI9Tk,4021 -diffusers/models/transformers/__init__.py,sha256=eSBzRwfhMlec09pYQcMMoH70-ZFn18h-6ZWSoumt9KA,2077 -diffusers/models/transformers/__pycache__/__init__.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/auraflow_transformer_2d.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/cogvideox_transformer_3d.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/consisid_transformer_3d.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/dit_transformer_2d.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/dual_transformer_2d.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/hunyuan_transformer_2d.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/latte_transformer_3d.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/lumina_nextdit2d.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/pixart_transformer_2d.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/prior_transformer.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/sana_transformer.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/stable_audio_transformer.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/t5_film_transformer.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_2d.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_allegro.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_chroma.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_cogview3plus.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_cogview4.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_cosmos.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_easyanimate.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_flux.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_hidream_image.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_hunyuan_video.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_hunyuan_video_framepack.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_ltx.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_lumina2.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_mochi.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_omnigen.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_sd3.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_temporal.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_wan.cpython-310.pyc,, -diffusers/models/transformers/__pycache__/transformer_wan_vace.cpython-310.pyc,, -diffusers/models/transformers/auraflow_transformer_2d.py,sha256=Lnk2i4c0Ay3SZZXiRHIsTvL-IuTyxIhZvyYMK2X0unU,23439 -diffusers/models/transformers/cogvideox_transformer_3d.py,sha256=gCZBGA7R2kFhf0m-nLxg3eISapfABKxVmzWN09J5rCA,22827 -diffusers/models/transformers/consisid_transformer_3d.py,sha256=ZT4plc4SnCJZczuKVhFTjws_3gPU6kCxm0-lIwAQA14,36023 -diffusers/models/transformers/dit_transformer_2d.py,sha256=4c-KTXuQagfdIr0m6VoysBYaNbAOTGlEt_AAoAfhoI0,10544 -diffusers/models/transformers/dual_transformer_2d.py,sha256=4iU7Ljh9_HQa9uiNwwW2xa7FqZOiEdw2TrTBuWdmR4M,7711 -diffusers/models/transformers/hunyuan_transformer_2d.py,sha256=OKYOjAse6jQn0eeimtPgQ5vhnV7720rznzeZWdVvQoA,24311 -diffusers/models/transformers/latte_transformer_3d.py,sha256=30LmMFjvNQyUKB3DHoVQfFph9wQP5i6BWIcS1zWdVhg,15686 -diffusers/models/transformers/lumina_nextdit2d.py,sha256=THVMjGaMNqvZDy9RXusVMH4lTR-e_q3ns6Vmjq5JZf4,14486 -diffusers/models/transformers/pixart_transformer_2d.py,sha256=K-qrAK5OqbgkZzoPqFZz8KiqKTMSyHCjoFOiCbfJ4hU,20940 -diffusers/models/transformers/prior_transformer.py,sha256=PHWVtGDOIo29k7xyb8BukUFbs-yT9s6Vr4Y8hcgYjo0,17467 -diffusers/models/transformers/sana_transformer.py,sha256=pqXzq-npEsee5ateLL4PS2oHsx8Rz_x36VcSyKWNcYg,25256 -diffusers/models/transformers/stable_audio_transformer.py,sha256=ikFDgj4O5rEk2Nxp3sL9RH76MLM5vzQnfMrR0Donlro,18653 -diffusers/models/transformers/t5_film_transformer.py,sha256=YicaHEHbkOkh7HLjzj8jzIjLA9MobcolI_3ZyCKETCE,16040 -diffusers/models/transformers/transformer_2d.py,sha256=9yiEGNq4poMiHaWIh3OfsfB2dsRKlcIqgTNDW_mV3vY,28305 -diffusers/models/transformers/transformer_allegro.py,sha256=b8fdklklJlmepbUgEwpn83L8AvmhycZao17_-sK2Ehg,17186 -diffusers/models/transformers/transformer_chroma.py,sha256=5t8b4LuAt6DTlxM6FFVi8SFzOON2_ueT67Uo-hR2oIo,31322 -diffusers/models/transformers/transformer_cogview3plus.py,sha256=shbvxhjREErZRCl2T16JZC2gnm39PIG6zgyAgMTDNLo,15811 -diffusers/models/transformers/transformer_cogview4.py,sha256=sK5RUMTzQsxMayJ8ZDxnfqTBN7TeYcFaQjO-n1-BTzk,33638 -diffusers/models/transformers/transformer_cosmos.py,sha256=o-iHkWwVOjudOAF0uLZfcCDpKElj34T9uXfycnPm4FA,24311 -diffusers/models/transformers/transformer_easyanimate.py,sha256=Z5innEWRbaOVoC16o-h-uobbYI6tPqzFG-S2H_Dohgo,22001 -diffusers/models/transformers/transformer_flux.py,sha256=LN86DeAqnhX5QxdxYuyml1EvYk9-gpoNowLoMERbKIc,23502 -diffusers/models/transformers/transformer_hidream_image.py,sha256=ciuyPRprws2K77apFzuLHmDZADOJ3AzAwsVze3FmR_E,39322 -diffusers/models/transformers/transformer_hunyuan_video.py,sha256=HN2zgdACpgw29JvZqnQS2X6xMcEX_LAIKRzSs0F6X28,46897 -diffusers/models/transformers/transformer_hunyuan_video_framepack.py,sha256=p9RgBqN5SWWqQU_wvEmaPZ2W5sCsYZoiGblqFKg6GpM,18698 -diffusers/models/transformers/transformer_ltx.py,sha256=iJQ1tR6Ihof0Rzk0-UsO57QnQNJ05xq-ICkVykbJr_I,18854 -diffusers/models/transformers/transformer_lumina2.py,sha256=ACu9X7vatGMoiKSZKkEOcbEQjIvZwmqZ6nmHaSkD9Wo,22075 -diffusers/models/transformers/transformer_mochi.py,sha256=FbNpuQR3MrH7I7CU8tJ_8Nf00Q2VB5hNUpqwuYXS50Y,18521 -diffusers/models/transformers/transformer_omnigen.py,sha256=nxmuNqRXRG51Mw2G8BNu-DZrRKWMDc3G99DFgPS9yZA,20029 -diffusers/models/transformers/transformer_sd3.py,sha256=2Aw1Di240iPngdlg_TZewlgi4WT2fGAOpB5XFa3RmO0,19232 -diffusers/models/transformers/transformer_temporal.py,sha256=GMn5WUbWWX7ZvyqVhO12g6bID7dnrMndYrDN-UZEI0Q,16812 -diffusers/models/transformers/transformer_wan.py,sha256=nUSN30izYESv4ETeL1LD6SkI0G9zN0qK859hd717mf8,20144 -diffusers/models/transformers/transformer_wan_vace.py,sha256=K_n6VSeAxVgl6p0HufXt4JTeYCcc_WvN4eX5MO41gr4,16866 -diffusers/models/unets/__init__.py,sha256=srYFA7zEcDY7LxyUB2jz3TdRgsLz8elrWCpT6Y4YXuU,695 -diffusers/models/unets/__pycache__/__init__.cpython-310.pyc,, -diffusers/models/unets/__pycache__/unet_1d.cpython-310.pyc,, -diffusers/models/unets/__pycache__/unet_1d_blocks.cpython-310.pyc,, -diffusers/models/unets/__pycache__/unet_2d.cpython-310.pyc,, -diffusers/models/unets/__pycache__/unet_2d_blocks.cpython-310.pyc,, -diffusers/models/unets/__pycache__/unet_2d_blocks_flax.cpython-310.pyc,, -diffusers/models/unets/__pycache__/unet_2d_condition.cpython-310.pyc,, -diffusers/models/unets/__pycache__/unet_2d_condition_flax.cpython-310.pyc,, -diffusers/models/unets/__pycache__/unet_3d_blocks.cpython-310.pyc,, -diffusers/models/unets/__pycache__/unet_3d_condition.cpython-310.pyc,, -diffusers/models/unets/__pycache__/unet_i2vgen_xl.cpython-310.pyc,, -diffusers/models/unets/__pycache__/unet_kandinsky3.cpython-310.pyc,, -diffusers/models/unets/__pycache__/unet_motion_model.cpython-310.pyc,, -diffusers/models/unets/__pycache__/unet_spatio_temporal_condition.cpython-310.pyc,, -diffusers/models/unets/__pycache__/unet_stable_cascade.cpython-310.pyc,, -diffusers/models/unets/__pycache__/uvit_2d.cpython-310.pyc,, -diffusers/models/unets/unet_1d.py,sha256=tmSBsH3cPzNj9xCQt1zrV2lApufLEB2xiFTcXulE1Wo,10853 -diffusers/models/unets/unet_1d_blocks.py,sha256=LIuM8MwkcJ0n8wYwS6FBGYSMfl9wYv0-TbQ6FHO6A7k,26829 -diffusers/models/unets/unet_2d.py,sha256=ue2Xbgmra55kSf6D7sq4b_ZCDr4ZmarrZiGWT4qNNXk,16905 -diffusers/models/unets/unet_2d_blocks.py,sha256=r7e-UDzbBRdbx4Iusu-fNGmSm70Vf2YkaDRihltn12w,141758 -diffusers/models/unets/unet_2d_blocks_flax.py,sha256=RWQ-c2gOHFD99nbtWN2TTGvO2khRWNJ6BoQYDg-hqjs,15624 -diffusers/models/unets/unet_2d_condition.py,sha256=zzRzpVkxmP73brYE_QI4n4JK4K6X7ADhkBv_mSiTIMk,67030 -diffusers/models/unets/unet_2d_condition_flax.py,sha256=62kO3PaPNfcPvWbOIL4nws7m_SZ2j4_GMAA1lXVPR4I,22289 -diffusers/models/unets/unet_3d_blocks.py,sha256=OyEIov5NMa_T9kIWzcMoMfmKZVcgDRbe6vH1Wb4ar9o,51676 -diffusers/models/unets/unet_3d_condition.py,sha256=_OscxwREczYttMFghevnaGu-LnIZLlV63LiS8HszdkU,34379 -diffusers/models/unets/unet_i2vgen_xl.py,sha256=BcsJO2K_GDZV3aPt9-VDsJN9hD07xAkFCVf5iE8PA2A,32505 -diffusers/models/unets/unet_kandinsky3.py,sha256=B9kql5-5U7rG7sbZg6Hy49P-4MusRQDwZT78fc2v6FY,20584 -diffusers/models/unets/unet_motion_model.py,sha256=5ct4afFiyEQzq-GhDHbO1HcV0rj-YvaM3gr7jfMpXGo,99370 -diffusers/models/unets/unet_spatio_temporal_condition.py,sha256=LsBj9J6nltqs2vrkIt-OT4weL1GUxJMmcNx3wTMM-vY,23259 -diffusers/models/unets/unet_stable_cascade.py,sha256=6fgr1ksDUhAe54xd_4uRHEqbeGzlVnd7R80GyKAG9_4,27229 -diffusers/models/unets/uvit_2d.py,sha256=0cV3yDVyC1KA9bAMnhcHGGVbT7n9bFl09gX61bcQLeo,17227 -diffusers/models/upsampling.py,sha256=7oUwSN_b8Wsl3nTuxWgZyXiIid5FOBnTpdeGwS64Xfk,19573 -diffusers/models/vae_flax.py,sha256=862gk7zaRCWBHNY1ZDIRJl7rnoHLWJBJyt6wN4anZ4w,31950 -diffusers/models/vq_model.py,sha256=J6dVRtVAQNUtHbl5LY5KjHYYsab55AbuoATHQttDGX4,1524 -diffusers/optimization.py,sha256=CAFI9pabb3C2KD3y_LCr3o_dBISPU_dfDKODtzbkdvs,14741 -diffusers/pipelines/__init__.py,sha256=5tYzTmjhNadauLrwT7LYyNcESP4MQAtFKVPsqpUzUck,33270 -diffusers/pipelines/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/__pycache__/auto_pipeline.cpython-310.pyc,, -diffusers/pipelines/__pycache__/free_init_utils.cpython-310.pyc,, -diffusers/pipelines/__pycache__/free_noise_utils.cpython-310.pyc,, -diffusers/pipelines/__pycache__/onnx_utils.cpython-310.pyc,, -diffusers/pipelines/__pycache__/pipeline_flax_utils.cpython-310.pyc,, -diffusers/pipelines/__pycache__/pipeline_loading_utils.cpython-310.pyc,, -diffusers/pipelines/__pycache__/pipeline_utils.cpython-310.pyc,, -diffusers/pipelines/__pycache__/transformers_loading_utils.cpython-310.pyc,, -diffusers/pipelines/allegro/__init__.py,sha256=T1MLZgDf8Fhh6YunF8a4Ta6NNIqneWsJIvmBhiy1ABM,1290 -diffusers/pipelines/allegro/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/allegro/__pycache__/pipeline_allegro.cpython-310.pyc,, -diffusers/pipelines/allegro/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/allegro/pipeline_allegro.py,sha256=6h-lqIRh_nktZAZeorkOdeMwuaS5UCXskIAVrxhQ9nw,44067 -diffusers/pipelines/allegro/pipeline_output.py,sha256=Q2W_16pT5o5xbzSrCKO8AB3IweqAyvv6BGTlPwAUNhE,722 -diffusers/pipelines/amused/__init__.py,sha256=pzqLeLosNQ29prMLhTxvPpmoIDPB3OFMQMlErOIRkmI,1793 -diffusers/pipelines/amused/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/amused/__pycache__/pipeline_amused.cpython-310.pyc,, -diffusers/pipelines/amused/__pycache__/pipeline_amused_img2img.cpython-310.pyc,, -diffusers/pipelines/amused/__pycache__/pipeline_amused_inpaint.cpython-310.pyc,, -diffusers/pipelines/amused/pipeline_amused.py,sha256=EJbwGd3ePgst-GmA9CSfIOncKaxRozr4zAi1QzGRKks,16018 -diffusers/pipelines/amused/pipeline_amused_img2img.py,sha256=PvInrMAHCOqES6HjltqGbkAmLWa0T_WvZub4IAAAoIE,17495 -diffusers/pipelines/amused/pipeline_amused_inpaint.py,sha256=zmpdSChNiXl8v4C7GKnlEmyGtuPfhYKXVo3OPpFW0uQ,19181 -diffusers/pipelines/animatediff/__init__.py,sha256=8e7xkGr1MrQerNXsfFBaDT8f7ELe5aoPX9v2qRN1hvg,2324 -diffusers/pipelines/animatediff/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff.cpython-310.pyc,, -diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_controlnet.cpython-310.pyc,, -diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_sdxl.cpython-310.pyc,, -diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_sparsectrl.cpython-310.pyc,, -diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_video2video.cpython-310.pyc,, -diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_video2video_controlnet.cpython-310.pyc,, -diffusers/pipelines/animatediff/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/animatediff/pipeline_animatediff.py,sha256=LZWsnSJ0PaWdLcn4hipA3TKSuZxllqYUIAn0M3Uk5oo,42434 -diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py,sha256=wyKBvLx6UMF_vqgmu42HNkj0wMUBjLggXsBrqyi61Jk,55940 -diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py,sha256=O1vn-vdJSyLOxk_o9VO2gPzsKV4WZIS18lcgtpkmeP0,66566 -diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py,sha256=hZJ4Bqy857KN0tbzyjVpRQlUCeJE3z9d0byZcNbYCm8,51514 -diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py,sha256=v-KmxoDFQJnWPcToTTrEScO3UIK2gqhLjOJxuKzA71Q,52258 -diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py,sha256=6NNBLmnPnr9VoB0Uw6plQR_8HEan3a8yJXo8Mx2pTjI,67610 -diffusers/pipelines/animatediff/pipeline_output.py,sha256=Ggp2OfMwdOPjHh4wIEN5aHJHDiSU0ORyzWzfdysrtcA,729 -diffusers/pipelines/audioldm/__init__.py,sha256=HMUjKqEf7OAtgIeV2CQoGIoDE6oY7b26N55yn4qCIpU,1419 -diffusers/pipelines/audioldm/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/audioldm/__pycache__/pipeline_audioldm.cpython-310.pyc,, -diffusers/pipelines/audioldm/pipeline_audioldm.py,sha256=R4WZVEilSXA339bzZWMFzcpVDnf8cn1L1Vd-CGlNNqQ,26375 -diffusers/pipelines/audioldm2/__init__.py,sha256=gR7gTyh-YGI4uxTCPnz_LnCGbErpFGtNMEzM_CQdqgE,1605 -diffusers/pipelines/audioldm2/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/audioldm2/__pycache__/modeling_audioldm2.cpython-310.pyc,, -diffusers/pipelines/audioldm2/__pycache__/pipeline_audioldm2.cpython-310.pyc,, -diffusers/pipelines/audioldm2/modeling_audioldm2.py,sha256=SpKP4UVjYRznsInmq4gvobBnX2uTlA4wILzyWRHWdY0,70483 -diffusers/pipelines/audioldm2/pipeline_audioldm2.py,sha256=2r_aHtDrCijivNAPbGIyxzp_zz8NppHMpzTNQOo9ayg,54856 -diffusers/pipelines/aura_flow/__init__.py,sha256=TOGRbwqwr7j1XIVGAxIBwAp4lM2zt21C_hYm5dFb76o,1296 -diffusers/pipelines/aura_flow/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/aura_flow/__pycache__/pipeline_aura_flow.cpython-310.pyc,, -diffusers/pipelines/aura_flow/pipeline_aura_flow.py,sha256=c2d06A6gdkM09Hciw0An8Hfu5y44bDmggyDptFTe_wU,33217 -diffusers/pipelines/auto_pipeline.py,sha256=7DobcWehSghE8ASytRf292ISW1XrhB7DjjBkWYSBRvc,58301 -diffusers/pipelines/blip_diffusion/__init__.py,sha256=v_PoaUspuKZG54FdKtITSccYo6eIhMnO0d6n7Pf3JJU,697 -diffusers/pipelines/blip_diffusion/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/blip_diffusion/__pycache__/blip_image_processing.cpython-310.pyc,, -diffusers/pipelines/blip_diffusion/__pycache__/modeling_blip2.cpython-310.pyc,, -diffusers/pipelines/blip_diffusion/__pycache__/modeling_ctx_clip.cpython-310.pyc,, -diffusers/pipelines/blip_diffusion/__pycache__/pipeline_blip_diffusion.cpython-310.pyc,, -diffusers/pipelines/blip_diffusion/blip_image_processing.py,sha256=xRyx14zBBKrERosHMm_Pj9HjSfESxFAh6nPmhj42dgE,16742 -diffusers/pipelines/blip_diffusion/modeling_blip2.py,sha256=eBrI0_E7mBN0uctDOHVdcv_tYhHHWpaXGuk7DZoxAQQ,27139 -diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py,sha256=53ggZ2e0gs0Oh7UwpKvSXdrcr70IlYkAUNB-hnkpMYk,9002 -diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py,sha256=mxCrWForS6kYGJ8wkg9Im_Nic0gVwQb2KpfNUdQfM9A,15322 -diffusers/pipelines/chroma/__init__.py,sha256=iRH8V4N01bnqzev2gozrhjVu1cRhLAn61RNFCBxFzgg,1611 -diffusers/pipelines/chroma/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/chroma/__pycache__/pipeline_chroma.cpython-310.pyc,, -diffusers/pipelines/chroma/__pycache__/pipeline_chroma_img2img.cpython-310.pyc,, -diffusers/pipelines/chroma/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/chroma/pipeline_chroma.py,sha256=oqZxPLKh_shVr52SiAd74Ra0v9WTRnup09WyJC_J_WY,44783 -diffusers/pipelines/chroma/pipeline_chroma_img2img.py,sha256=fWvioL2M6BSwU5ewishYX8fr3fNMa7W7waK4mAnEifA,49244 -diffusers/pipelines/chroma/pipeline_output.py,sha256=_LYDjLRSV0vIaNBSwddFVpWwbNGPQ6E5GkNAccpaOQo,600 -diffusers/pipelines/cogvideo/__init__.py,sha256=84bmJbrCvjUtEXFgyCKvX5N4HNtAWorMjQTNvWPh8ZU,1816 -diffusers/pipelines/cogvideo/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox.cpython-310.pyc,, -diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox_fun_control.cpython-310.pyc,, -diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox_image2video.cpython-310.pyc,, -diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox_video2video.cpython-310.pyc,, -diffusers/pipelines/cogvideo/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/cogvideo/pipeline_cogvideox.py,sha256=v4ECnmnyJDebRQvMN-tbNhLBy6sFmav5aaM2NUEhd4U,37686 -diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py,sha256=bephSzLcsNUQVmFw2AI484mAl1BL0RL_D-im0Bs9d3s,40247 -diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py,sha256=7uvUgL33C0enGCqBxj2yZhSEtzplTwhAiZwi0m3kU74,42690 -diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py,sha256=u83mFoVRGOusErNc2dpLMZJYKuRGvOUCLCZ_cGZvRmQ,41301 -diffusers/pipelines/cogvideo/pipeline_output.py,sha256=QOyumhJJERjm7moyxnYzU_X27hvN9p99MIkjT_Vf1x0,616 -diffusers/pipelines/cogview3/__init__.py,sha256=ophRMlB8W7AocUEWUJLbmK1o4yJpHEfKvwyDceyMu00,1497 -diffusers/pipelines/cogview3/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/cogview3/__pycache__/pipeline_cogview3plus.cpython-310.pyc,, -diffusers/pipelines/cogview3/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/cogview3/pipeline_cogview3plus.py,sha256=oCX9IvHfGfY7UbeYWa8jc1-q7xiVhUDjuYNPfzZUV7E,33881 -diffusers/pipelines/cogview3/pipeline_output.py,sha256=qU187W2KZY8KloD6E5EOpkbgJOYrsQFSGHyNHKHqwxs,594 -diffusers/pipelines/cogview4/__init__.py,sha256=DSW0f5XIu2bcirGYm6FE9lhyxADbV5nbiayq1x2ttJg,1633 -diffusers/pipelines/cogview4/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/cogview4/__pycache__/pipeline_cogview4.cpython-310.pyc,, -diffusers/pipelines/cogview4/__pycache__/pipeline_cogview4_control.cpython-310.pyc,, -diffusers/pipelines/cogview4/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/cogview4/pipeline_cogview4.py,sha256=d5ZI1KVrd9b_0Fr6_chIeYWEB7I8jl-k1A8Iz-wMoKE,33662 -diffusers/pipelines/cogview4/pipeline_cogview4_control.py,sha256=oqZJ5XbE3RblxqH5VzIbMpT1AnnQJV-krevmRc4CXyY,35345 -diffusers/pipelines/cogview4/pipeline_output.py,sha256=l3XZTqGTse6epEQn_VZsGyRAyQF9TpfYRaQrqO5O_Dw,594 -diffusers/pipelines/consisid/__init__.py,sha256=wi4mmbsztby5LLgmrtDhz857JWT_4Jbc2sRzRyL0EpY,1367 -diffusers/pipelines/consisid/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/consisid/__pycache__/consisid_utils.cpython-310.pyc,, -diffusers/pipelines/consisid/__pycache__/pipeline_consisid.cpython-310.pyc,, -diffusers/pipelines/consisid/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/consisid/consisid_utils.py,sha256=ZoAiyDYhNV_nMJHr2pkhtHgXVFzdcsIhWvAWWzgpoek,14522 -diffusers/pipelines/consisid/pipeline_consisid.py,sha256=Di5fV28wCWfQrRX8fzqhQj9BnwTTBHat1XLpS9emNwA,46715 -diffusers/pipelines/consisid/pipeline_output.py,sha256=bCTmHqOyRXggYJLwx_nBxNg0-yci9Dhge8Div0k_D1U,615 -diffusers/pipelines/consistency_models/__init__.py,sha256=q_nrLK9DH0_kLcLmRIvgvLP-vDVwloC3lBus776596c,484 -diffusers/pipelines/consistency_models/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/consistency_models/__pycache__/pipeline_consistency_models.cpython-310.pyc,, -diffusers/pipelines/consistency_models/pipeline_consistency_models.py,sha256=-QlWGbE6nflcWLYLY5s7F5drUOUxaR7IXPd35bC3rB8,12596 -diffusers/pipelines/controlnet/__init__.py,sha256=pqndp8HbyQ2D45STcpMp37nO5M4SagpfwADCCOC_2CU,4057 -diffusers/pipelines/controlnet/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/controlnet/__pycache__/multicontrolnet.cpython-310.pyc,, -diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet.cpython-310.pyc,, -diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_blip_diffusion.cpython-310.pyc,, -diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_img2img.cpython-310.pyc,, -diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_inpaint.cpython-310.pyc,, -diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_inpaint_sd_xl.cpython-310.pyc,, -diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_sd_xl.cpython-310.pyc,, -diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_sd_xl_img2img.cpython-310.pyc,, -diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_union_inpaint_sd_xl.cpython-310.pyc,, -diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_union_sd_xl.cpython-310.pyc,, -diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_union_sd_xl_img2img.cpython-310.pyc,, -diffusers/pipelines/controlnet/__pycache__/pipeline_flax_controlnet.cpython-310.pyc,, -diffusers/pipelines/controlnet/multicontrolnet.py,sha256=-sluVPEM3oDV_dscvtnklZf3S_KsbRK98-V1fwGKUtg,684 -diffusers/pipelines/controlnet/pipeline_controlnet.py,sha256=ctdh4Nt3TzEHSi8Hsa5xUz4FDvLK7AIFEw_lxUk4fgk,69320 -diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py,sha256=8ajTqSf96k6OgFyeCCHohAANKk-rxnL7GDHUF1I5nTY,17596 -diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py,sha256=vr7XO-3CfGRVeLlS4QUAtqyekPZYrWVb0pyTXfcHwn8,67610 -diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py,sha256=JxW2wNkNUUY4HiI3IhivKwvG7qEDSgJ3kKk_ILS5YBQ,76660 -diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py,sha256=DsTeSK8OqT9racfudAHSSyPttoXAewKF2hL3hJcgDlA,95313 -diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py,sha256=5yUeJMSfgbQ-po-5ljpfPPeCQwK_m_UEc5Gv1ZiJOs4,82765 -diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py,sha256=lBQ72LPNhz7JB6zLRhkDTsMGOH_OwdCN24iFaoysqGA,87279 -diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py,sha256=26INHgrfnw7QCWEifp6t6aPbefVmXVpOHoy01jgdMJY,90688 -diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py,sha256=twIFMu935FUjrtYNkMDs_B_LBSHQSwEjMcmHFMAdol4,83286 -diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py,sha256=I9PxHFLz81vf7BrP6oVoA72Jfsu81Khsr0TEOdgiQkw,83525 -diffusers/pipelines/controlnet/pipeline_flax_controlnet.py,sha256=gJ54OTBKXNYMk1wPzaaiE7a60F_mp9S4me7DAElgIsQ,22771 -diffusers/pipelines/controlnet_hunyuandit/__init__.py,sha256=LvB-TNhPTnUIdinVZfxzUX40RFWvNWxrjAzsDDiLBfM,1344 -diffusers/pipelines/controlnet_hunyuandit/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/controlnet_hunyuandit/__pycache__/pipeline_hunyuandit_controlnet.cpython-310.pyc,, -diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py,sha256=JmU_qrB_oPBVnu0hnrc5HEgtRCSfWicIZQwSwUFEdRg,51005 -diffusers/pipelines/controlnet_sd3/__init__.py,sha256=_-t5_Jac1hvUKbjACSwVDVWx1lFIBQDCeCE9CVMbsW0,1903 -diffusers/pipelines/controlnet_sd3/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/controlnet_sd3/__pycache__/pipeline_stable_diffusion_3_controlnet.cpython-310.pyc,, -diffusers/pipelines/controlnet_sd3/__pycache__/pipeline_stable_diffusion_3_controlnet_inpainting.cpython-310.pyc,, -diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py,sha256=-2xtwBvccIne5Bau5P7_wOw7SdkOMq5dC9AmEAzGVLI,62633 -diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py,sha256=MOxSarLvTPdiFqFwaVWmAGCqWzOINPSda09taRqPD68,63597 -diffusers/pipelines/controlnet_xs/__init__.py,sha256=TuIgTKgY4MVB6zaoNTduQAEVRsNptBZQZhnxxQ3hpyg,2403 -diffusers/pipelines/controlnet_xs/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/controlnet_xs/__pycache__/pipeline_controlnet_xs.cpython-310.pyc,, -diffusers/pipelines/controlnet_xs/__pycache__/pipeline_controlnet_xs_sd_xl.cpython-310.pyc,, -diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py,sha256=I_wHrWyaxKh8DccRvvOn7SBfjn2AoYTJxEFvOQndyiQ,46056 -diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py,sha256=xnp5qtiDjVZeq4c6DgleZE8a7qJ7Y1VFv-j54ygvAfE,57342 -diffusers/pipelines/cosmos/__init__.py,sha256=pX1QX9zkM2e4Vo_0A3tvOVE0X-vHxYuFOedVtHUPmfs,1820 -diffusers/pipelines/cosmos/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos2_text2image.cpython-310.pyc,, -diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos2_video2world.cpython-310.pyc,, -diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos_text2world.cpython-310.pyc,, -diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos_video2world.cpython-310.pyc,, -diffusers/pipelines/cosmos/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py,sha256=sO54yib9payqZjZpBhvodLdPDkIEtlIonJAgCMJLI5Q,32962 -diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py,sha256=u94ClHRWm87hYfcoDSbN_Hd4WLX9cx0H0mbRyqdfva4,39284 -diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py,sha256=JxExXGuCdoKaahFC1iux5BwLJnWFCGqA0T9JC9undP8,31639 -diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py,sha256=8NWyFa78NXtHxgwSFB0ecA-Qz29EoqU6CrGh8mOX15o,39983 -diffusers/pipelines/cosmos/pipeline_output.py,sha256=iJIxvS8s2YCd790FLEd3Bl3WNm1TS67MGJlgOF6rqjA,1214 -diffusers/pipelines/dance_diffusion/__init__.py,sha256=SOwr8mpuw34oKEUuy4uVLlhjfHuLRCP0kpMjoSPXADU,453 -diffusers/pipelines/dance_diffusion/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/dance_diffusion/__pycache__/pipeline_dance_diffusion.cpython-310.pyc,, -diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py,sha256=VDSk8B05ddFRCTRD3_cjyVH4qCM0MvqdxSIrg_MPDms,6732 -diffusers/pipelines/ddim/__init__.py,sha256=-zCVlqBSKWZdwY5HSsoiRT4nUEuT6dckiD_KIFen3bs,411 -diffusers/pipelines/ddim/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/ddim/__pycache__/pipeline_ddim.cpython-310.pyc,, -diffusers/pipelines/ddim/pipeline_ddim.py,sha256=YkknZ2EFsC7Knf-0KJB-ckEU5-xCHjctxUa7karayAk,6910 -diffusers/pipelines/ddpm/__init__.py,sha256=DAj0i0-iba7KACShx0bzGa9gqAV7yxGgf9sy_Hf095Q,425 -diffusers/pipelines/ddpm/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/ddpm/__pycache__/pipeline_ddpm.cpython-310.pyc,, -diffusers/pipelines/ddpm/pipeline_ddpm.py,sha256=M0hWjZbUbWL2VCiqsJ5oJ2VNDZSjS9jscFUXsF4_bKs,5344 -diffusers/pipelines/deepfloyd_if/__init__.py,sha256=gh1fQ5u6q0d-o3XGExCGD0jPaUK-gWCturfHU-TYIi8,2975 -diffusers/pipelines/deepfloyd_if/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if.cpython-310.pyc,, -diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_img2img.cpython-310.pyc,, -diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_img2img_superresolution.cpython-310.pyc,, -diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_inpainting.cpython-310.pyc,, -diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_inpainting_superresolution.cpython-310.pyc,, -diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_superresolution.cpython-310.pyc,, -diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/deepfloyd_if/__pycache__/safety_checker.cpython-310.pyc,, -diffusers/pipelines/deepfloyd_if/__pycache__/timesteps.cpython-310.pyc,, -diffusers/pipelines/deepfloyd_if/__pycache__/watermark.cpython-310.pyc,, -diffusers/pipelines/deepfloyd_if/pipeline_if.py,sha256=r7908WHOwr0MiDK0IBCQYfDUVncIGZ1nvSkhL84wBXY,35624 -diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py,sha256=K6d5O6eDU8OJ8GcCa5H3ck_PGX9ttRqpgJN_b0IHApw,40033 -diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py,sha256=G6pRC9JmRlWQn5h7jf9-D2URPzWSJdz4ABJ0PrBvBBo,45073 -diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py,sha256=2rT9uoczvSbqvIppIkzZRnJt0KL14gG7X0r9iIJvg-0,45231 -diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py,sha256=2Osy4qE8PfrKPRUC6yDf536tVVnuAePrvvQF_Hwy8ZY,50070 -diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py,sha256=eSFTCAH0GE_Pn6hrMAk_CfQrpZdqHdHw7pm8HDMB-k8,39915 -diffusers/pipelines/deepfloyd_if/pipeline_output.py,sha256=RBxF2VMyXgIiZzJs9ZnmRRPom6GdFzb1DrOnuYmd6uQ,1145 -diffusers/pipelines/deepfloyd_if/safety_checker.py,sha256=zqN0z4Mvf7AtrxlUb6qAoiw_QuxGdDk-6js5YuarxTo,2117 -diffusers/pipelines/deepfloyd_if/timesteps.py,sha256=JO8b-8zlcvk_Tb6s6GGY7MgRPRADs35y0KBcSkqmNDM,5164 -diffusers/pipelines/deepfloyd_if/watermark.py,sha256=d-43jrlsjyJt1NJrXrl7U1LgCPlFD5C1gzJ83GVoijc,1601 -diffusers/pipelines/deprecated/__init__.py,sha256=mXBnea22TkkUdiGxUpZDXTSb1RlURczuRcGeIzn9DcQ,5470 -diffusers/pipelines/deprecated/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/deprecated/alt_diffusion/__init__.py,sha256=1SiGoNJytgnMwGmR48q8erVnU9JP5uz5E6XgHvlFDTc,1783 -diffusers/pipelines/deprecated/alt_diffusion/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/deprecated/alt_diffusion/__pycache__/modeling_roberta_series.cpython-310.pyc,, -diffusers/pipelines/deprecated/alt_diffusion/__pycache__/pipeline_alt_diffusion.cpython-310.pyc,, -diffusers/pipelines/deprecated/alt_diffusion/__pycache__/pipeline_alt_diffusion_img2img.cpython-310.pyc,, -diffusers/pipelines/deprecated/alt_diffusion/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py,sha256=CmrX8Y2bvoTr1_kZYQ-13nXL9ttMsX6gYX_0yPx1F3g,5530 -diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py,sha256=AgwVAwzfqBp0xXay6gTADpLygznVEpCUyPuxVGAd-zU,50256 -diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py,sha256=cGf3Fi3sjqJTRpgjAub3IGU52lPkrq0fc_WOLJ1vwgk,52976 -diffusers/pipelines/deprecated/alt_diffusion/pipeline_output.py,sha256=wtKrIaa_f-rfw5_bbEGi5mdNnQ6qmsaTGcDVBNBnvJ8,928 -diffusers/pipelines/deprecated/audio_diffusion/__init__.py,sha256=SiFqPmeNbqOYTwuTx2WUaMIpMzgSnJ2SZ_97tIDryOE,507 -diffusers/pipelines/deprecated/audio_diffusion/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/deprecated/audio_diffusion/__pycache__/mel.cpython-310.pyc,, -diffusers/pipelines/deprecated/audio_diffusion/__pycache__/pipeline_audio_diffusion.cpython-310.pyc,, -diffusers/pipelines/deprecated/audio_diffusion/mel.py,sha256=BjN1eYahKBBWn3FTSlHpn5VqbC1pym253k7u8WfucXw,5764 -diffusers/pipelines/deprecated/audio_diffusion/pipeline_audio_diffusion.py,sha256=E-e-P_zb5rQQ61QYFP_A_ONMkFLaCSB1FNGrt-C06GQ,13239 -diffusers/pipelines/deprecated/latent_diffusion_uncond/__init__.py,sha256=ZWWt671s-zbWawgtJNoIstZsvOE5ucP2M_vp7OMUMeM,448 -diffusers/pipelines/deprecated/latent_diffusion_uncond/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/deprecated/latent_diffusion_uncond/__pycache__/pipeline_latent_diffusion_uncond.cpython-310.pyc,, -diffusers/pipelines/deprecated/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py,sha256=pl4GRcCIjzva81en87KPzAiX19z3ARqUwrgapzPbKV8,5381 -diffusers/pipelines/deprecated/pndm/__init__.py,sha256=R8RavcZ5QXU-fR4o4HT_xvypifWUcqRKF3bduCgieEI,412 -diffusers/pipelines/deprecated/pndm/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/deprecated/pndm/__pycache__/pipeline_pndm.cpython-310.pyc,, -diffusers/pipelines/deprecated/pndm/pipeline_pndm.py,sha256=H9NvlZdnQ1MfSrsaicaCdvjpVPJdba6yh1a1i9Sa2UY,4662 -diffusers/pipelines/deprecated/repaint/__init__.py,sha256=mlHI_qG20VS7yuags8W0HXpbHkZgObu-jUBuYnOfffo,425 -diffusers/pipelines/deprecated/repaint/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/deprecated/repaint/__pycache__/pipeline_repaint.cpython-310.pyc,, -diffusers/pipelines/deprecated/repaint/pipeline_repaint.py,sha256=nkJmBYTDiF8Znogu769JFGFQDfX1qNY-MjMC201fLl8,10094 -diffusers/pipelines/deprecated/score_sde_ve/__init__.py,sha256=7CLXxU1JqmMFbdm0bLwCHxGUjGJFvS64xueOQdD2X7s,441 -diffusers/pipelines/deprecated/score_sde_ve/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/deprecated/score_sde_ve/__pycache__/pipeline_score_sde_ve.cpython-310.pyc,, -diffusers/pipelines/deprecated/score_sde_ve/pipeline_score_sde_ve.py,sha256=_1vGqiY3cZY-1h_74x7XGUmfT7IZicTl3K6XYlgg49g,4390 -diffusers/pipelines/deprecated/spectrogram_diffusion/__init__.py,sha256=lOJEU-CHJhv0N2BCEM9-dzKmm1Y-HPt1FuF9lGBgIpg,2588 -diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/continuous_encoder.cpython-310.pyc,, -diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/midi_utils.cpython-310.pyc,, -diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/notes_encoder.cpython-310.pyc,, -diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/pipeline_spectrogram_diffusion.cpython-310.pyc,, -diffusers/pipelines/deprecated/spectrogram_diffusion/continuous_encoder.py,sha256=7hqiAFb5gYob4MPvNMDaDJPX1IlJNlHzHRfnC52yhGQ,3100 -diffusers/pipelines/deprecated/spectrogram_diffusion/midi_utils.py,sha256=UY5j5rCm2rzkvX1aGrjqzl-dkvh9nHGaV67PiIDXAqk,25096 -diffusers/pipelines/deprecated/spectrogram_diffusion/notes_encoder.py,sha256=1dphTvTIpodYsic9Rn9DpX9MOchZW_PjyS626MIWw1A,2923 -diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py,sha256=DaKLh1UdQ0jlwo92eU--Js5lL9A027akOx7UaDT5vXs,11528 -diffusers/pipelines/deprecated/stable_diffusion_variants/__init__.py,sha256=mnIQupN59oc3JmKGaQZia7MO92E08wswJrP9QITzWQs,2111 -diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_cycle_diffusion.cpython-310.pyc,, -diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_onnx_stable_diffusion_inpaint_legacy.cpython-310.pyc,, -diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_inpaint_legacy.cpython-310.pyc,, -diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_model_editing.cpython-310.pyc,, -diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_paradigms.cpython-310.pyc,, -diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_pix2pix_zero.cpython-310.pyc,, -diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py,sha256=4iCc0Z2B6RCWsfQ3s03Xj7kOC2RX43RHM5paJwDqztM,48043 -diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py,sha256=d2dTEqsup3V20i9dI8EM1ROGJlBMSrbVLecptK1tBnY,27841 -diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py,sha256=cr9owvhcPbRtB85QQbJ8CK4hVqdJ__TeacljBe6BunI,42556 -diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py,sha256=CDlbEGKY8M5ZOEwSexMxyq0Re_0PfVwzaFmV8OoFjkQ,41530 -diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py,sha256=DzGD_KpDwTModZFDiRxfrWBrgi-goFEglS3A_w6XJZ8,41216 -diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py,sha256=xa7GE5lLiAmb_c_5A4y0BqaLV9yBMTIzjdjEOQx6vPc,63563 -diffusers/pipelines/deprecated/stochastic_karras_ve/__init__.py,sha256=WOKqWaBgVgNkDUUf4ZL1--TauXKeaPqtGf3P2fTFYMw,453 -diffusers/pipelines/deprecated/stochastic_karras_ve/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/deprecated/stochastic_karras_ve/__pycache__/pipeline_stochastic_karras_ve.cpython-310.pyc,, -diffusers/pipelines/deprecated/stochastic_karras_ve/pipeline_stochastic_karras_ve.py,sha256=jRhJFQxTQ18_Kr05blYzF415U1xGMuIVPEycF5UGxX8,5277 -diffusers/pipelines/deprecated/versatile_diffusion/__init__.py,sha256=_CRp2PIJD6loFlES3hMcPigZNOUMf2OgTaRFgoit7hc,2838 -diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/modeling_text_unet.cpython-310.pyc,, -diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion.cpython-310.pyc,, -diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion_dual_guided.cpython-310.pyc,, -diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion_image_variation.cpython-310.pyc,, -diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion_text_to_image.cpython-310.pyc,, -diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py,sha256=s4_K2WJXXXVXXJfrw1pBU14gme12kfgs1Eakjgjm2oc,113024 -diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py,sha256=lBCiSYM8hQrILIEhtxs2bGDQqwsd1RHA2rWHl_xOY2c,21912 -diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py,sha256=Or5-U70TQ716NM2PH5UdEQErlY5trDPj1TsYbJnl5BU,27228 -diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py,sha256=PqrXJWL-ZvPhaF3UMwPgbsYcxYMfJ5liVqkNS1mfwyQ,19742 -diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py,sha256=Ut-bWQmxYf6g6zw7YOf_TRnWR8Om0V3Bx9B7l0fMt5k,22946 -diffusers/pipelines/deprecated/vq_diffusion/__init__.py,sha256=CD0X20a3_61pBaOzDxgU_33PLjxN1W8V46TCAwykUgE,1650 -diffusers/pipelines/deprecated/vq_diffusion/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/deprecated/vq_diffusion/__pycache__/pipeline_vq_diffusion.cpython-310.pyc,, -diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py,sha256=6El_KY-vwVUvYF3LGiwk6K5q5UjwhgT9J9cBw_V5Jv0,15444 -diffusers/pipelines/dit/__init__.py,sha256=w6yUFMbGzaUGPKpLfEfvHlYmrKD0UErczwsHDaDtLuQ,408 -diffusers/pipelines/dit/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/dit/__pycache__/pipeline_dit.cpython-310.pyc,, -diffusers/pipelines/dit/pipeline_dit.py,sha256=D4TI2xAQnUcsKEBput0EDV5p4y2uKsiolRxd-5Us_uQ,10275 -diffusers/pipelines/easyanimate/__init__.py,sha256=CeJHlQus6mhlN2rmFk1LA44ygm4jYJVEyMYWMckcxZI,1634 -diffusers/pipelines/easyanimate/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/easyanimate/__pycache__/pipeline_easyanimate.cpython-310.pyc,, -diffusers/pipelines/easyanimate/__pycache__/pipeline_easyanimate_control.cpython-310.pyc,, -diffusers/pipelines/easyanimate/__pycache__/pipeline_easyanimate_inpaint.cpython-310.pyc,, -diffusers/pipelines/easyanimate/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/easyanimate/pipeline_easyanimate.py,sha256=iNg09pgsMUdZzkvgqxbq5olN0fZBt87jK6Yhiws1xII,35903 -diffusers/pipelines/easyanimate/pipeline_easyanimate_control.py,sha256=evNzX0HHy9HO83ob-yTALTAlDm_ZUyx62-Z92wzKiH8,46131 -diffusers/pipelines/easyanimate/pipeline_easyanimate_inpaint.py,sha256=9OZA8OnhHqaRtxyZFX7SUbrlA6VRMFwPpa7pRYggXdY,58526 -diffusers/pipelines/easyanimate/pipeline_output.py,sha256=cpEvM-GStcMexQOQnPunDqrCOHP_Fk4aNxMvUHYeuxQ,621 -diffusers/pipelines/flux/__init__.py,sha256=wVmWJNnXfm5LIsV0eB1ugtnz_lbRoEDC7jdJZxwWxlE,3143 -diffusers/pipelines/flux/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/flux/__pycache__/modeling_flux.cpython-310.pyc,, -diffusers/pipelines/flux/__pycache__/pipeline_flux.cpython-310.pyc,, -diffusers/pipelines/flux/__pycache__/pipeline_flux_control.cpython-310.pyc,, -diffusers/pipelines/flux/__pycache__/pipeline_flux_control_img2img.cpython-310.pyc,, -diffusers/pipelines/flux/__pycache__/pipeline_flux_control_inpaint.cpython-310.pyc,, -diffusers/pipelines/flux/__pycache__/pipeline_flux_controlnet.cpython-310.pyc,, -diffusers/pipelines/flux/__pycache__/pipeline_flux_controlnet_image_to_image.cpython-310.pyc,, -diffusers/pipelines/flux/__pycache__/pipeline_flux_controlnet_inpainting.cpython-310.pyc,, -diffusers/pipelines/flux/__pycache__/pipeline_flux_fill.cpython-310.pyc,, -diffusers/pipelines/flux/__pycache__/pipeline_flux_img2img.cpython-310.pyc,, -diffusers/pipelines/flux/__pycache__/pipeline_flux_inpaint.cpython-310.pyc,, -diffusers/pipelines/flux/__pycache__/pipeline_flux_prior_redux.cpython-310.pyc,, -diffusers/pipelines/flux/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/flux/modeling_flux.py,sha256=46xvV8iCF77ifCJFPQ9tjtvNxBorVnBrQZw_eHRBujQ,1549 -diffusers/pipelines/flux/pipeline_flux.py,sha256=G9_TO9pX2aEwMPCiMpCgOe6PIq-RApxQhxLHtmOOPaY,47045 -diffusers/pipelines/flux/pipeline_flux_control.py,sha256=rQodOu9pEtTJR-dpoD9RfaoJQ2RORTPxDwBr041747Q,40671 -diffusers/pipelines/flux/pipeline_flux_control_img2img.py,sha256=OBl-7AMaDpApUD09U8CI1AOFsmuxJXNtxN4FBniQ2Rg,44561 -diffusers/pipelines/flux/pipeline_flux_control_inpaint.py,sha256=rx4GqLHlJkZ6P5qIhgyINUYi3U2RdfEp5ca6odp12_8,53655 -diffusers/pipelines/flux/pipeline_flux_controlnet.py,sha256=0qKCt57GyPw0DWTUHuSCxthRZWBH28Ug0aiZXlVp9b0,57915 -diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py,sha256=k3MlFbme1HKJ_2e07E1XL90A0zCRcRjwaAPlD_8TBgE,46497 -diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py,sha256=-1R7oFr-cXRE5uPYBKypIrWG1bw4DDdNlIXMUwBADAg,55525 -diffusers/pipelines/flux/pipeline_flux_fill.py,sha256=L3cEwHTltNl7NCZeEy_VYkA9V6ejxHhzdkvejtC3oDQ,49397 -diffusers/pipelines/flux/pipeline_flux_img2img.py,sha256=9LY5esG__sgSS7c8hAjps37E2Xhsp6jTtj12Dx6d5Gc,52256 -diffusers/pipelines/flux/pipeline_flux_inpaint.py,sha256=rFTDZoCChYfGIflrdhL6trvS5e7rOnQgIGAs1O9bVRs,59241 -diffusers/pipelines/flux/pipeline_flux_prior_redux.py,sha256=ruDOkehjql-hiImJXBiArJ-HxwjFdx_ep4NFEYyvVvI,21808 -diffusers/pipelines/flux/pipeline_output.py,sha256=0w0lxW102Yf5upqYyGO2JLFYygcJNwg9ql92yKZZlTI,1102 -diffusers/pipelines/free_init_utils.py,sha256=SHrGV68cii9sYCKZLbIdBEOB5tANOVbN9vv7KznOlII,7699 -diffusers/pipelines/free_noise_utils.py,sha256=SlcvpUInyDDOOq7CkzXsjbcsC3Z7nY_JBzy6MJorHc4,29691 -diffusers/pipelines/hidream_image/__init__.py,sha256=SeMI0Ae_K8guNBe8doZIX_vJndXz6jWFWhCqtH1s5I0,1499 -diffusers/pipelines/hidream_image/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/hidream_image/__pycache__/pipeline_hidream_image.cpython-310.pyc,, -diffusers/pipelines/hidream_image/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/hidream_image/pipeline_hidream_image.py,sha256=eTkd0terZVFjMRYg43e46iCRuqydOElLXhzXJ_IX4Tk,51391 -diffusers/pipelines/hidream_image/pipeline_output.py,sha256=nfVKQcyZW-RIHpoVd5XUBzP_ekbWvk6xVikTuPbASgU,1229 -diffusers/pipelines/hunyuan_video/__init__.py,sha256=2IhKqUmvwUZpkBWZI1HBbsq15U-_YEC3aSwiNDMogDE,1878 -diffusers/pipelines/hunyuan_video/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_skyreels_image2video.cpython-310.pyc,, -diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_video.cpython-310.pyc,, -diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_video_framepack.cpython-310.pyc,, -diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_video_image2video.cpython-310.pyc,, -diffusers/pipelines/hunyuan_video/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py,sha256=8Z7MjSkajuroRhVI4Vlno578R_eOuMRkxatgk2lR-VU,38826 -diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py,sha256=froIVMQ8BiZxPvH1o69r-Tind4qZKZK_Q6wmyTPNr6s,35591 -diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py,sha256=MKjgjLDy9OEDN32vKWyfqZ7WlHObgpi3s_bpVVSLeG0,53966 -diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py,sha256=yzvQhYFGcHXa0sfP9sKxRCwN0HfsOjIgouK2qJWOF4w,46321 -diffusers/pipelines/hunyuan_video/pipeline_output.py,sha256=xSreG3mREfUvF_dfDzWDLF2pSdqNiUp47whS-tVl7xI,1417 -diffusers/pipelines/hunyuandit/__init__.py,sha256=Zby0yEsLNAoa4cf6W92QXIzyGoijI54xXRVhmrHGHsc,1302 -diffusers/pipelines/hunyuandit/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/hunyuandit/__pycache__/pipeline_hunyuandit.cpython-310.pyc,, -diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py,sha256=Rag59kh4-CYhoUhO3eB38c1ZLTKPaQqt-xdQpguQpCM,43442 -diffusers/pipelines/i2vgen_xl/__init__.py,sha256=5Stj50A-AIJ1pPhilpDRx1PARMs_n8OKTDl64cq0LAY,1307 -diffusers/pipelines/i2vgen_xl/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/i2vgen_xl/__pycache__/pipeline_i2vgen_xl.cpython-310.pyc,, -diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py,sha256=Uhe3P_65v4M0CaoM5cRqO5PcHrYRH2c-uNNYWEHEQcY,37427 -diffusers/pipelines/kandinsky/__init__.py,sha256=wrxuhSw_CunNhm7TdzA_fm__092mibGxp5_ep1boZmQ,2312 -diffusers/pipelines/kandinsky/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky.cpython-310.pyc,, -diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_combined.cpython-310.pyc,, -diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_img2img.cpython-310.pyc,, -diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_inpaint.cpython-310.pyc,, -diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_prior.cpython-310.pyc,, -diffusers/pipelines/kandinsky/__pycache__/text_encoder.cpython-310.pyc,, -diffusers/pipelines/kandinsky/pipeline_kandinsky.py,sha256=9bjqUCjOFgsuFX38fzNTOHvLSQVXWLSUH8-2aVHPf1c,17913 -diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py,sha256=vLWCvPd6Qpv3XNlvpm4t5bFCXaz1Pcz8CA9-hP_GvdA,39610 -diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py,sha256=Gbl8MhTC6cfifzDqS_sO0winbWbi3MfLx9oMOyr7r6Y,21920 -diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py,sha256=7nepe82E_RqUAKaYY76hb8TONq70SKcEomsmKTNEYOE,28740 -diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py,sha256=29gDoi2jh95jyubYbbY_7TDpYiQ3XzmSsrbLonOQhkE,24044 -diffusers/pipelines/kandinsky/text_encoder.py,sha256=zDi1K-p-rPii0ZugI-83D75DR6AW36pkl8SvGBO77bA,1022 -diffusers/pipelines/kandinsky2_2/__init__.py,sha256=WeV8KWoCLj6KTvJ-f3Do87IoX_dR_AZNylBz7_Iu87s,2796 -diffusers/pipelines/kandinsky2_2/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2.cpython-310.pyc,, -diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_combined.cpython-310.pyc,, -diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_controlnet.cpython-310.pyc,, -diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_controlnet_img2img.cpython-310.pyc,, -diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_img2img.cpython-310.pyc,, -diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_inpainting.cpython-310.pyc,, -diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_prior.cpython-310.pyc,, -diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_prior_emb2emb.cpython-310.pyc,, -diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py,sha256=Mk61oIckbX_h2Rz6jn93qR46vNPuNAvpKFdUFqz7zmg,14377 -diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py,sha256=7BW9hMirrEn7QaYjukHgM3C1fKW1lMFOhXIW2URCoX0,44243 -diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py,sha256=TZzib5OqmKlFuWemhCkMxJZubFXsONofoxXNlixqe3M,14342 -diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py,sha256=nbhz22EZYB-S9Qt_srCwwywYSfLskGQOplcAaR8wQR0,16849 -diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py,sha256=9blo8FSUgj4SwVMLJK4YtNNpbphlNeJsNFjZcK85r6o,17167 -diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py,sha256=ljn17ynhnbomb6nJFNJCWX81zhwq7fTkSgNA07D92Xc,25019 -diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py,sha256=2mIAe7glGkd9aGMfTKjS5t4LV4StkcXYbwZ4uyD-BuA,25648 -diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py,sha256=f3dL1Bqr8GZP9MV7Fvb8bWlMsfMo7W_clyeZ3agrSzk,25200 -diffusers/pipelines/kandinsky3/__init__.py,sha256=7Mv8Ov-XstHMLmRQU7psdheFn_e_qXJWWTYV7z7uj4U,1461 -diffusers/pipelines/kandinsky3/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/kandinsky3/__pycache__/convert_kandinsky3_unet.cpython-310.pyc,, -diffusers/pipelines/kandinsky3/__pycache__/pipeline_kandinsky3.cpython-310.pyc,, -diffusers/pipelines/kandinsky3/__pycache__/pipeline_kandinsky3_img2img.cpython-310.pyc,, -diffusers/pipelines/kandinsky3/convert_kandinsky3_unet.py,sha256=FJ8psagvZtQHJupm0hgMUI2mto3IHEXjaoLDXip1LMA,3273 -diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py,sha256=l3_pd1OpfKe34n4HQrHKVqFNcBL3YN01PeW03-3yOqM,27789 -diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py,sha256=cqQnHdUl_nEr_HRNTFUZ3P4_tHA2i-9a6PF3TUW29_Q,30506 -diffusers/pipelines/kolors/__init__.py,sha256=6Xp5M_K6PfByqqnK1HuMD9RKLkOZYekeNNqrGk4HToM,1791 -diffusers/pipelines/kolors/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/kolors/__pycache__/pipeline_kolors.cpython-310.pyc,, -diffusers/pipelines/kolors/__pycache__/pipeline_kolors_img2img.cpython-310.pyc,, -diffusers/pipelines/kolors/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/kolors/__pycache__/text_encoder.cpython-310.pyc,, -diffusers/pipelines/kolors/__pycache__/tokenizer.cpython-310.pyc,, -diffusers/pipelines/kolors/pipeline_kolors.py,sha256=im6XM1e4I13kqVFT_KOZrKLYAQ4GGgDMo2H63XstZSo,56001 -diffusers/pipelines/kolors/pipeline_kolors_img2img.py,sha256=CsA_xfEmCZbzgxB6hl1hOImhs0WQt5y9unkwJxXHrMM,65879 -diffusers/pipelines/kolors/pipeline_output.py,sha256=1POR3PAcQ-ZtBgf8GOwjVI46TCjdtpXPsv1sdslBFA0,590 -diffusers/pipelines/kolors/text_encoder.py,sha256=xpm-7Qv1gLUMTCdh5oQegPq7n1Yh56hbi4kmjlPpXD4,35059 -diffusers/pipelines/kolors/tokenizer.py,sha256=TzqdOwu1gyi8KJ_tb8lAc7_5M6AdvhkvJW2LKOCcDrI,13428 -diffusers/pipelines/latent_consistency_models/__init__.py,sha256=SfUylLTTBCs_wlGOPpW899lgE1E0GOLGu4GhDPFx-Ls,1560 -diffusers/pipelines/latent_consistency_models/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/latent_consistency_models/__pycache__/pipeline_latent_consistency_img2img.cpython-310.pyc,, -diffusers/pipelines/latent_consistency_models/__pycache__/pipeline_latent_consistency_text2img.cpython-310.pyc,, -diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py,sha256=lDaK8Pg1pokNh97zhbPYavsF3HTsErdGntq-IW6xZO0,49548 -diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py,sha256=LzLz1Fkyz65EzxTRE3kFvT2-gXLoEG6Tb6gonb0S3tA,46056 -diffusers/pipelines/latent_diffusion/__init__.py,sha256=iUkMRZY-pteRsvsROOz2Pacm7t02Q6QvbsgQedJt6-E,1542 -diffusers/pipelines/latent_diffusion/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/latent_diffusion/__pycache__/pipeline_latent_diffusion.cpython-310.pyc,, -diffusers/pipelines/latent_diffusion/__pycache__/pipeline_latent_diffusion_superresolution.cpython-310.pyc,, -diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py,sha256=wDR9k8KbxUoKlvbzQ0UQ6bw_gF8An2YXtP0_-OGvq2w,32607 -diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py,sha256=Z50tjLw7uJkngAZ-GHFd92iDT8gnyZBhDFzqTDbRnEg,8294 -diffusers/pipelines/latte/__init__.py,sha256=1XMhkoAvpw2akbDmMTsKJbTU4PsR9H6boq4FEhCGbwo,1282 -diffusers/pipelines/latte/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/latte/__pycache__/pipeline_latte.cpython-310.pyc,, -diffusers/pipelines/latte/pipeline_latte.py,sha256=71-fHckM7jLr84nVlFxZKAnIGN-xILcGCLNsrq-A4ho,42687 -diffusers/pipelines/ledits_pp/__init__.py,sha256=3VaqGS1d39iC5flUifb4vAD_bDJ-sIUFaLIYhBuHbwE,1783 -diffusers/pipelines/ledits_pp/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/ledits_pp/__pycache__/pipeline_leditspp_stable_diffusion.cpython-310.pyc,, -diffusers/pipelines/ledits_pp/__pycache__/pipeline_leditspp_stable_diffusion_xl.cpython-310.pyc,, -diffusers/pipelines/ledits_pp/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py,sha256=vK7RFmkCw5tdrYbgW2JI3q3WB5hj6Pc2qUFj79Q8y1Y,77462 -diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py,sha256=TgulOyUe1gjg_8jCutH1018qMqqTsnJAl9rNxlT_qiw,89461 -diffusers/pipelines/ledits_pp/pipeline_output.py,sha256=xiAplyxGWB6uCHIdryai6UP7ghtFuXhES52ZYpO3k8A,1579 -diffusers/pipelines/ltx/__init__.py,sha256=6STasTKI0b2ALlgCUi9GMmb_aScZzLWcfKKZxdHWOqo,1878 -diffusers/pipelines/ltx/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/ltx/__pycache__/modeling_latent_upsampler.cpython-310.pyc,, -diffusers/pipelines/ltx/__pycache__/pipeline_ltx.cpython-310.pyc,, -diffusers/pipelines/ltx/__pycache__/pipeline_ltx_condition.cpython-310.pyc,, -diffusers/pipelines/ltx/__pycache__/pipeline_ltx_image2video.cpython-310.pyc,, -diffusers/pipelines/ltx/__pycache__/pipeline_ltx_latent_upsample.cpython-310.pyc,, -diffusers/pipelines/ltx/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/ltx/modeling_latent_upsampler.py,sha256=2Fj74GY82S8By8l2mHKq1DIGaYUZQzUSI3Nf_1s00rA,7482 -diffusers/pipelines/ltx/pipeline_ltx.py,sha256=Pxm7U9QGYCuH7KtrpdpJoYIkW_nLOJtRoSsKxp5m_Yg,40849 -diffusers/pipelines/ltx/pipeline_ltx_condition.py,sha256=zMQCxCBCZNtwzQBRXO0ZqoSGU9RFeJrMXk2YgYMC8y4,62329 -diffusers/pipelines/ltx/pipeline_ltx_image2video.py,sha256=wfnWLsiaJER2uGFVNnW9eYZCOJLLdmhrHSNlgvOnIA8,45691 -diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py,sha256=6g5DXLlkMm4U3V4aWjHsGvBpHeO96Y9FSMIrWYlQNJk,11813 -diffusers/pipelines/ltx/pipeline_output.py,sha256=amPUfWP2XV5Qgb0BpfVIbC6L3vIHCMpTWdqgr5nfrvw,605 -diffusers/pipelines/lumina/__init__.py,sha256=AzWsnxikODkQnCxliBN7eDi83TxcSihxfehfLYxRPD4,1336 -diffusers/pipelines/lumina/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/lumina/__pycache__/pipeline_lumina.cpython-310.pyc,, -diffusers/pipelines/lumina/pipeline_lumina.py,sha256=5LTTbsltMVJg83iMMEjwhstpw64b3QbJfcUjYj8RKWE,44964 -diffusers/pipelines/lumina2/__init__.py,sha256=ZnlJglaTqtwptFJ0uelQ4MKg_p-Lxwzu7eQCd1CFxtc,1342 -diffusers/pipelines/lumina2/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/lumina2/__pycache__/pipeline_lumina2.cpython-310.pyc,, -diffusers/pipelines/lumina2/pipeline_lumina2.py,sha256=iiuTXnVG1AsOl--0R2F73VX2EKUh82VDUW9NVwHFs7c,38521 -diffusers/pipelines/marigold/__init__.py,sha256=kAs3DZB4oxiYHqLOq9kgAvBCYw3ptpiQKGr01hM2BDQ,1926 -diffusers/pipelines/marigold/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/marigold/__pycache__/marigold_image_processing.cpython-310.pyc,, -diffusers/pipelines/marigold/__pycache__/pipeline_marigold_depth.cpython-310.pyc,, -diffusers/pipelines/marigold/__pycache__/pipeline_marigold_intrinsics.cpython-310.pyc,, -diffusers/pipelines/marigold/__pycache__/pipeline_marigold_normals.cpython-310.pyc,, -diffusers/pipelines/marigold/marigold_image_processing.py,sha256=veWoIY4QDBKp74-MWHjbrQwgPtD5dOECVPAiTUs2DDE,31409 -diffusers/pipelines/marigold/pipeline_marigold_depth.py,sha256=IQqLiRvuSqFmwr37JlPEMaXZe0DVH8YeCm12MqvewO8,41201 -diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py,sha256=MADfO3Yzh30E5usYZSsf7EurYG3tIar01cy39chDsRY,35861 -diffusers/pipelines/marigold/pipeline_marigold_normals.py,sha256=osc4gL7CVCH5IH530Vej6Q06EGSbDXTeaaAIJBh_oao,34791 -diffusers/pipelines/mochi/__init__.py,sha256=8yDkp3YgOvbC4VhO4Tfin2myNxRlWiX1Mi8rY_UvAh4,1282 -diffusers/pipelines/mochi/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/mochi/__pycache__/pipeline_mochi.cpython-310.pyc,, -diffusers/pipelines/mochi/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/mochi/pipeline_mochi.py,sha256=ONMqBpSi-bzFbWssfEaB1CBcATniiUGoo7YtRY2lYk4,35591 -diffusers/pipelines/mochi/pipeline_output.py,sha256=RyFrgJRzyCNzbHurysrFsN4wtZLdcax8wTarxhUq-50,609 -diffusers/pipelines/musicldm/__init__.py,sha256=l1I5QzvTwMOOltJkcwpTb6nNcr93bWiP_ErHbDdwz6Y,1411 -diffusers/pipelines/musicldm/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/musicldm/__pycache__/pipeline_musicldm.cpython-310.pyc,, -diffusers/pipelines/musicldm/pipeline_musicldm.py,sha256=h5Li1cDVFDi-lOwL5i_fPXy-BfuqfC9kkYXCMEiYK08,30670 -diffusers/pipelines/omnigen/__init__.py,sha256=9596QBScCQfCbrybjkWJ7p0N4CnYB7W3hQJsNlGn3dU,1292 -diffusers/pipelines/omnigen/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/omnigen/__pycache__/pipeline_omnigen.cpython-310.pyc,, -diffusers/pipelines/omnigen/__pycache__/processor_omnigen.cpython-310.pyc,, -diffusers/pipelines/omnigen/pipeline_omnigen.py,sha256=h2-QMBcfcgYXhimDRIwwQ4Vt_dAGSXIVz4SBxxrLkaQ,23752 -diffusers/pipelines/omnigen/processor_omnigen.py,sha256=iGzPTREszduXazqwAHBmXKfoB2bKPBLITwLsnOJ0Wek,14191 -diffusers/pipelines/onnx_utils.py,sha256=oTRc_iLHEKpf_IGFw_ka1bloAI-XUa6_ASMLW2LAH4w,8810 -diffusers/pipelines/pag/__init__.py,sha256=pyv70bIvWZpMgXJgf8I8JH4OWzvdqJqfvxlFax9SzTg,3986 -diffusers/pipelines/pag/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pag_utils.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd_inpaint.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd_xl.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd_xl_img2img.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_hunyuandit.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_kolors.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_pixart_sigma.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_sana.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_sd.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_3.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_3_img2img.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_animatediff.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_img2img.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_inpaint.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_xl.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_xl_img2img.cpython-310.pyc,, -diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_xl_inpaint.cpython-310.pyc,, -diffusers/pipelines/pag/pag_utils.py,sha256=FGN5B2UocvZW7Y_EKwm_0B4BToIzfo0qemoHLW9jPvo,10221 -diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py,sha256=DIkyn-6Q0F0cp9DHx84zTZ9bzkMWYATd-Rh_m3LEc30,68337 -diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py,sha256=aFeuyLzqCZgT9EoVTEyfsHRQxT88Yab_25Hzi2k3024,79030 -diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py,sha256=hs2zkLmgb5eOHM4O6wUXiDFSSLFWteXcyYsY9CxtfLI,83631 -diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl_img2img.py,sha256=jxbc1AoLlYVSZHLp_bvyG68TYV185w0a5TvTpeRLbfU,88473 -diffusers/pipelines/pag/pipeline_pag_hunyuandit.py,sha256=Sszu4cI7T9pUM2CXv35k0-sFs5ecDJX8qnZBwKqjsbA,46649 -diffusers/pipelines/pag/pipeline_pag_kolors.py,sha256=G4ezp_hF7AF5LScEswyNAgt7ziU9XhwRQmfetD541ek,59694 -diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py,sha256=2DneQFcaXGRkurJvgmQ2jkvtI0zwzTcbPN4xj2d7Tv0,42103 -diffusers/pipelines/pag/pipeline_pag_sana.py,sha256=aJzap_XDStpd0kSHLJF136lV7gnqxd3_BBtwPCIhyp0,45728 -diffusers/pipelines/pag/pipeline_pag_sd.py,sha256=g2Fac2_uG1qGr0QzdGGLGPoYVR3rOX0J5BrnxT69OQI,55556 -diffusers/pipelines/pag/pipeline_pag_sd_3.py,sha256=08LmpiuzuMZRzVWyIGiIHvq5Axpwa7UXw0sBq1hRTu4,49479 -diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py,sha256=IFND8xNQbWG2jRxadoIkYdkAqPJHMPhC2z41ZgV-0Fc,53952 -diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py,sha256=GcDAKtLel_OAZ5eoXbSq_Q3iWTRk9SN05BFPFcTwvTQ,43567 -diffusers/pipelines/pag/pipeline_pag_sd_img2img.py,sha256=7v9waEdxFItMC7iJOy3Ac1djLXe0lAgRE_SJ8TEM1hQ,57891 -diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py,sha256=XUjGF65RTE7bylt-mRKjzEng3oIbxC3xF20RBgf5pyo,69978 -diffusers/pipelines/pag/pipeline_pag_sd_xl.py,sha256=gfgXXZw81JPMTgIvMPpw4Vwsf6IstEgXHchJHt-gyRA,70531 -diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py,sha256=Pps9YyXn-4QboB2nIqoXl9a--lrF5IN8zWt5IcZrGUo,81967 -diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py,sha256=001uhtBD-1gykVLFeckJP2Sy1t8DxYcKemzlLut54cw,93720 -diffusers/pipelines/paint_by_example/__init__.py,sha256=EL3EGhjCG7CMzwloJRauSDHc6oArjVsETUCj8mOauRs,1566 -diffusers/pipelines/paint_by_example/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/paint_by_example/__pycache__/image_encoder.cpython-310.pyc,, -diffusers/pipelines/paint_by_example/__pycache__/pipeline_paint_by_example.cpython-310.pyc,, -diffusers/pipelines/paint_by_example/image_encoder.py,sha256=NT3RadNaIVztZWGP628BDHpvaMRPZ-w3J3j096y3L9s,2484 -diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py,sha256=3PrO9c5yVuNajZ-fDmAJBxcMcCvQZB49YQLoQl3fJ8E,31344 -diffusers/pipelines/pia/__init__.py,sha256=md5F8G279iZg4WGSmLP7N8apWkuHkfssjLQFzv6c2zI,1299 -diffusers/pipelines/pia/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/pia/__pycache__/pipeline_pia.cpython-310.pyc,, -diffusers/pipelines/pia/pipeline_pia.py,sha256=TxnQBCeUWggseAABRyspi13LDs2zdyn1GbEOy7SJKlA,46525 -diffusers/pipelines/pipeline_flax_utils.py,sha256=q4emCFj0dF2xQTqeaC4JmwX0eR4JfXxoahY5OJjgk1Y,27034 -diffusers/pipelines/pipeline_loading_utils.py,sha256=VPIua1FpHmlF8bjm8q6VolMYgRTDxoRDbeE8Zdm0GpU,47268 -diffusers/pipelines/pipeline_utils.py,sha256=Bn3oM0fR3cqZL2t4MLeIn8pnaMZ18H1ULZfiB7BnSjg,105342 -diffusers/pipelines/pixart_alpha/__init__.py,sha256=QxcTJF9ryOIejEHQVw3bZAYHn2dah-WPT5pZudE8XxU,1595 -diffusers/pipelines/pixart_alpha/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/pixart_alpha/__pycache__/pipeline_pixart_alpha.cpython-310.pyc,, -diffusers/pipelines/pixart_alpha/__pycache__/pipeline_pixart_sigma.cpython-310.pyc,, -diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py,sha256=9qear7s6XP5Zt8g1fgpmQC4DxhVCMcYg26ZHdu7zvNg,44823 -diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py,sha256=FdbUbM8SKn3StH3pS4mHqvMy-3hRjaXBT_v7Ik77-m8,41217 -diffusers/pipelines/sana/__init__.py,sha256=qkgbJxOAEH4gmyQ4FX4USnOd-PPEDkZGjZ3QO0ID0pA,1719 -diffusers/pipelines/sana/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/sana/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/sana/__pycache__/pipeline_sana.cpython-310.pyc,, -diffusers/pipelines/sana/__pycache__/pipeline_sana_controlnet.cpython-310.pyc,, -diffusers/pipelines/sana/__pycache__/pipeline_sana_sprint.cpython-310.pyc,, -diffusers/pipelines/sana/__pycache__/pipeline_sana_sprint_img2img.cpython-310.pyc,, -diffusers/pipelines/sana/pipeline_output.py,sha256=ErM82CTECCPbaLIpJsEzwl0r_hGNi29IbxKcsJ1mEMM,586 -diffusers/pipelines/sana/pipeline_sana.py,sha256=Y7W4MXlppwU04EWCeE6tkRI31tAystB8mnGiYRdWZ1U,47262 -diffusers/pipelines/sana/pipeline_sana_controlnet.py,sha256=SkMUSXUPVX-Md9mBqtwacqoPWB92ibART5CUE2IPAFo,51742 -diffusers/pipelines/sana/pipeline_sana_sprint.py,sha256=Q4ckz9JXmOngptRQh0up3zKCO8sf5T1FAnIaNB8UpI0,41342 -diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py,sha256=X_PURGwytjQh5S_XRLQEqKE1zERVtHAFNITa775_qE8,45583 -diffusers/pipelines/semantic_stable_diffusion/__init__.py,sha256=4jDvmgpXRVXGeSAcfGN90iQoJJBBRgE7NXzBE_8AYxM,1443 -diffusers/pipelines/semantic_stable_diffusion/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/semantic_stable_diffusion/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/semantic_stable_diffusion/__pycache__/pipeline_semantic_stable_diffusion.cpython-310.pyc,, -diffusers/pipelines/semantic_stable_diffusion/pipeline_output.py,sha256=YBxNQ2JiY3jYW-GB44nzNZxeADAswMQBJfnr2tBX0eY,822 -diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py,sha256=w3qDw26PbhhhPgankdSOnHIACBrrb-9_X5hM8mrmj1s,39024 -diffusers/pipelines/shap_e/__init__.py,sha256=LGToZwsVeVBEsE5eveY0Hc2GgI6UgDz6H_6cB_Snn0Y,2093 -diffusers/pipelines/shap_e/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/shap_e/__pycache__/camera.cpython-310.pyc,, -diffusers/pipelines/shap_e/__pycache__/pipeline_shap_e.cpython-310.pyc,, -diffusers/pipelines/shap_e/__pycache__/pipeline_shap_e_img2img.cpython-310.pyc,, -diffusers/pipelines/shap_e/__pycache__/renderer.cpython-310.pyc,, -diffusers/pipelines/shap_e/camera.py,sha256=oIKN8kcD6gqpeQWABRDfkdm3wW0MXFCZJSW_qdmUMLo,4942 -diffusers/pipelines/shap_e/pipeline_shap_e.py,sha256=YG8TDX2t_2TVseeyS7Rotxorh2prpJsp8ol4ucUxM48,13398 -diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py,sha256=uCfAsOKkrI85NPR3X1lhfK_10D2IhfOVkDrUJBmxpps,13205 -diffusers/pipelines/shap_e/renderer.py,sha256=Wx9GqWWLUmZjOIXdMgPKGyotodqyrdNzbKIkBrmY9ck,39153 -diffusers/pipelines/stable_audio/__init__.py,sha256=R8Tuxx2LsaWWR0lncRJ0faKOmAdaQ0ilvftdBC_07Eo,1561 -diffusers/pipelines/stable_audio/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/stable_audio/__pycache__/modeling_stable_audio.cpython-310.pyc,, -diffusers/pipelines/stable_audio/__pycache__/pipeline_stable_audio.cpython-310.pyc,, -diffusers/pipelines/stable_audio/modeling_stable_audio.py,sha256=vakcN7KjC2PjVWoYbpZSiqW7Fnd63LpvgmYtmoXfSyE,6127 -diffusers/pipelines/stable_audio/pipeline_stable_audio.py,sha256=aph9UJQsqVQs2EJzl2iu6jSY65Nk7TTen7sSd6FWij0,35613 -diffusers/pipelines/stable_cascade/__init__.py,sha256=buKExLbA-qdePd19JSEF29AhOCIaDgqFfLajEmo-Kg4,1672 -diffusers/pipelines/stable_cascade/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/stable_cascade/__pycache__/pipeline_stable_cascade.cpython-310.pyc,, -diffusers/pipelines/stable_cascade/__pycache__/pipeline_stable_cascade_combined.cpython-310.pyc,, -diffusers/pipelines/stable_cascade/__pycache__/pipeline_stable_cascade_prior.cpython-310.pyc,, -diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py,sha256=xXRM34zQrVgayRBVKJrGA11hjU_2u-V-h5slnVcdh0c,26124 -diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py,sha256=CcFI-qOl1VCXQPgbB196gqB12CAcMu23I7m5RjIRnm4,18130 -diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py,sha256=fBpaAE1aPZhvq4EuEg-7Uq6khJLFTP3fYI-hPeIZOFY,31444 -diffusers/pipelines/stable_diffusion/__init__.py,sha256=4DkYPvosNLHumc2wQcV10KkWsGRCRhcyeUogvraHSEM,8555 -diffusers/pipelines/stable_diffusion/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/clip_image_project_model.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/convert_from_ckpt.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_flax_stable_diffusion.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_flax_stable_diffusion_img2img.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_flax_stable_diffusion_inpaint.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion_img2img.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion_inpaint.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion_upscale.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_depth2img.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_image_variation.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_img2img.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_inpaint.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_instruct_pix2pix.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_latent_upscale.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_upscale.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_unclip.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_unclip_img2img.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/safety_checker.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/safety_checker_flax.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/__pycache__/stable_unclip_image_normalizer.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion/clip_image_project_model.py,sha256=9LPlAopNG-vpWhmh8i8pz7CHttqZ203-RnGlvy74xxY,1094 -diffusers/pipelines/stable_diffusion/convert_from_ckpt.py,sha256=7AkKESPubz4W7bYqzGOK87gxsP_HZg5nX5VhRmoCPgg,81729 -diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py,sha256=3SsX3YOJJbZ5El0cLf1jf-9M83ZMenIf7NXiaQA86ho,20756 -diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py,sha256=G6V7Okambalqs8-Agpa0iREwSyc_9GZH5-iRhYP_No8,22554 -diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py,sha256=7F-z40HNgb0a6a_dV_T46nj6ZqSmBDSlsGMPClEPd7U,26130 -diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py,sha256=tVEekoMVfQ1CyF0VSYAU5UnzNvDbv7R44IbvluSU-EQ,24341 -diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py,sha256=J-FbiSOueXBLk153qyUUgAcaAzVLHdHTzXwJIi7BPbc,28552 -diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py,sha256=heCESPus4qrmryvhHK-Rsq6JzpHf-88gJKgSD9gZDRk,29170 -diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py,sha256=UFZ8HlQuuHF8mnT7Fq0xO0FeOaW5SD7X29SMd6QJBh0,27950 -diffusers/pipelines/stable_diffusion/pipeline_output.py,sha256=Io-12AumvYjBOKP4Qq1BJ2rak5pKdtMO-ACcorY6hFE,1496 -diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py,sha256=E9x8m3TsPZnmCguwNBL6-k5MPglX4m4THlXCQwKouF0,55652 -diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py,sha256=sO7DpfNobnuAal0kDM3mD3z2fe5GRlyTHG7u2x_rZIY,44509 -diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py,sha256=XNLerrOA1__emc2DHlPt7-16pxJR_plP3usFH4_DVGQ,22720 -diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py,sha256=Th3gTxeKNQnAHjnPF2DgAB7dKHzLsWNZ6yoGw1XO094,59513 -diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py,sha256=hHJegbumSJD5-Tf213dLcNFy0YY6UTKorL5Y8DuMqbQ,70158 -diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py,sha256=_vsAtsPHV40l1LFYPTD-uJRzaaj0_ACMu2Jb4fT-bFA,45826 -diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py,sha256=jFgn7dOCcGRa_W8ALhRM2fTSq6umDxqTD5fX6LYRY6U,30985 -diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py,sha256=91cjNp53ErZtTQqltnHz6LDHtiOapmCXE_UIXc_6LQY,39757 -diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py,sha256=-ga6Cxugyl9N1OlJiL5W2-BRAcapI0Us4G2XccZcEcU,45544 -diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py,sha256=55kewO9zpxnTmcXrmqwoGDheiRXjSKEdZ5TpedeX3b4,40310 -diffusers/pipelines/stable_diffusion/safety_checker.py,sha256=h3WwFCcuJpWxEebxwgiRs0x7t4-4u4ma0GtLOeERaLA,5759 -diffusers/pipelines/stable_diffusion/safety_checker_flax.py,sha256=RhHxBbrcBrEosONc4uVlA9LgftN3WJv7wn_nV-nSDCk,4476 -diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py,sha256=wgtRMD_dzwanDdbqHTqlc_Ro_1q51SnotP8uqyC7Yr8,1890 -diffusers/pipelines/stable_diffusion_3/__init__.py,sha256=4JrcTgfij4mGbSSnCaHSqRRNhCUry8-HH3zQaUIq3DE,1922 -diffusers/pipelines/stable_diffusion_3/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_stable_diffusion_3.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_stable_diffusion_3_img2img.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_stable_diffusion_3_inpaint.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_3/pipeline_output.py,sha256=empNHoFAmdz6__yOCX2kuJqZtVdtoGAvVmH5mW42-3s,610 -diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py,sha256=3Wdp8VpfehM0pwTGH3jmi_DUqZgPuEJLKuWkMe1BhhY,56846 -diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py,sha256=HDuLdQDpuWnxvSBk_ATg5qcKjbyo7j2RbgkTmD3c-Bc,57765 -diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py,sha256=KvElHFo4l47wHmXx4GpK9EH-J-6PhNGqwMylnZB07FQ,69805 -diffusers/pipelines/stable_diffusion_attend_and_excite/__init__.py,sha256=VpZ5FPx9ACTOT4qiEqun2QYeUtx9Rp0YVDwqhYe28QM,1390 -diffusers/pipelines/stable_diffusion_attend_and_excite/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_attend_and_excite/__pycache__/pipeline_stable_diffusion_attend_and_excite.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py,sha256=RckQAPyM_I7oqlwmA10pXjGG8Jg2vUTSRCAh5SxbRdg,51576 -diffusers/pipelines/stable_diffusion_diffedit/__init__.py,sha256=JlcUNahRBm0uaPzappogqfjyLDsNW6IeyOfuLs4af5M,1358 -diffusers/pipelines/stable_diffusion_diffedit/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_diffedit/__pycache__/pipeline_stable_diffusion_diffedit.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py,sha256=Eq_Wm1HEgtPR5wbLU7mKecM0qlta-nHY8jg_Q0zuTLk,78549 -diffusers/pipelines/stable_diffusion_gligen/__init__.py,sha256=b4dZB5bUuZmEAcg7MmCyWZpyxNmMrlrByEQW_xwGGgI,1568 -diffusers/pipelines/stable_diffusion_gligen/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_gligen/__pycache__/pipeline_stable_diffusion_gligen.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_gligen/__pycache__/pipeline_stable_diffusion_gligen_text_image.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py,sha256=8Z_67jeBdl2ZCrukTkxYkRW73NazIYUWWGJK5WpWZOY,43505 -diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py,sha256=QR2zoPiBxX_6lsfeYyBuiZNoXwrtXgkW0rTfuCOR8I0,52092 -diffusers/pipelines/stable_diffusion_k_diffusion/__init__.py,sha256=EBpyQedEN-jfJ0qeLCFg9t28cFPNbNaniKIGM4ZMF14,1924 -diffusers/pipelines/stable_diffusion_k_diffusion/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_k_diffusion/__pycache__/pipeline_stable_diffusion_k_diffusion.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_k_diffusion/__pycache__/pipeline_stable_diffusion_xl_k_diffusion.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py,sha256=b7pXB2GE2n7W-_he4RoUxiQyvU-JtuUdudhA_tBuv1k,34096 -diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py,sha256=KYwMUrnxseg_O9fLme-wVBOoNYSekS8YYr6K2lIf2aU,45496 -diffusers/pipelines/stable_diffusion_ldm3d/__init__.py,sha256=8p2npGKPPJbPaTa4swOWRMd24x36E563Bhc_mM29va0,1346 -diffusers/pipelines/stable_diffusion_ldm3d/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_ldm3d/__pycache__/pipeline_stable_diffusion_ldm3d.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py,sha256=vrIqC9I1bJvfv79G8SzuU2slJo3CXXmHzvSX9XMv-tE,51801 -diffusers/pipelines/stable_diffusion_panorama/__init__.py,sha256=af52eZSYshuw1d6kqKwx0C5Teopkx8UpO9ph_A4WI0Q,1358 -diffusers/pipelines/stable_diffusion_panorama/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_panorama/__pycache__/pipeline_stable_diffusion_panorama.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py,sha256=l1SmYHbJE-Nz1Iq82WckJEAmeg5xXbwBjxztcsk4X2k,60262 -diffusers/pipelines/stable_diffusion_safe/__init__.py,sha256=rRKtzOjuaHLDqSLSavcy2W8sEljso9MLhmEwrNiJFJ0,2751 -diffusers/pipelines/stable_diffusion_safe/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_safe/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_safe/__pycache__/pipeline_stable_diffusion_safe.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_safe/__pycache__/safety_checker.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_safe/pipeline_output.py,sha256=WGQS6-k9dPH0hYBj_dZMlHFkOvUUti9fjVv0Sf8LCjQ,1459 -diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py,sha256=zTC3C_MaqPnEgJjAqoCqpPU14JkKBtGG95AD5PaDstU,39634 -diffusers/pipelines/stable_diffusion_safe/safety_checker.py,sha256=Ewt9xIW9sKMGL0sbdjBgu4OhuzmjN_xqI87E6CwLBg0,5039 -diffusers/pipelines/stable_diffusion_sag/__init__.py,sha256=06vnWbASiG3o4sQ7CDlDrqEm6dSCerKdLODz1FS-EFE,1338 -diffusers/pipelines/stable_diffusion_sag/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_sag/__pycache__/pipeline_stable_diffusion_sag.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py,sha256=66XjcrcYb2AmD2jNsoqJ1fColkmBWBvSQeojNbGlUeg,48087 -diffusers/pipelines/stable_diffusion_xl/__init__.py,sha256=6lTMI458kVDLzQDeZxEBacdFxpj4xAY9CSZ6Xr_FWoY,3022 -diffusers/pipelines/stable_diffusion_xl/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_flax_stable_diffusion_xl.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl_img2img.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl_inpaint.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl_instruct_pix2pix.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_xl/__pycache__/watermark.cpython-310.pyc,, -diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py,sha256=dGnKmKC_XzmMre6JTtdeTbXEYlenI-7CBmnyULK2BME,11280 -diffusers/pipelines/stable_diffusion_xl/pipeline_output.py,sha256=Isy1wE8hgoScXXHWVel5jRAzgPTelP-aZieugTOTgUc,1037 -diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py,sha256=QgIJOG9PV7jVKgsneQRDGdbdFq0JZM_s6UPT36GgJ5I,67847 -diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py,sha256=B6z9lXb432zNiY3aiwLvln2u8mXFULjQHlP5EuaYMaM,78850 -diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py,sha256=XqCUCWpzy2tEA3rpKjRK2TXjhCL87cO23-hqiTTPRKw,90574 -diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py,sha256=NkET5qoCbXLJYVvOELEacOWYjzBk53JfvVepIrFhPtk,52910 -diffusers/pipelines/stable_diffusion_xl/watermark.py,sha256=LDItvRnZKokIUchP0oIrO2Ew9AARhAP4MMrQY8maQ6Q,1458 -diffusers/pipelines/stable_video_diffusion/__init__.py,sha256=QtcDxzfLJ7loCDspiulKyKU6kd-l3twJyWBDPraD_94,1551 -diffusers/pipelines/stable_video_diffusion/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/stable_video_diffusion/__pycache__/pipeline_stable_video_diffusion.cpython-310.pyc,, -diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py,sha256=x6jEr8eURAGh55q6Iey1pIIhxbWEzxD0NQmzvY_Vd-k,32628 -diffusers/pipelines/t2i_adapter/__init__.py,sha256=PgIg_SzwFAqWOML5BLHvuCTmu4p06MPT66xBpDShx8c,1556 -diffusers/pipelines/t2i_adapter/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/t2i_adapter/__pycache__/pipeline_stable_diffusion_adapter.cpython-310.pyc,, -diffusers/pipelines/t2i_adapter/__pycache__/pipeline_stable_diffusion_xl_adapter.cpython-310.pyc,, -diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py,sha256=izK-JFnaVGq-hbOcPwJVaymM-MQVBJEJHRo-QcT1PiA,47764 -diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py,sha256=qxnmqebGNnOutOPnQowpa8rmGUS-iJmQND4euuO4z9w,69409 -diffusers/pipelines/text_to_video_synthesis/__init__.py,sha256=7-NplGtgnp5GUu4XN_STE9fqAtFCAc6FF3lphjbDBhs,1979 -diffusers/pipelines/text_to_video_synthesis/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_synth.cpython-310.pyc,, -diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_synth_img2img.cpython-310.pyc,, -diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_zero.cpython-310.pyc,, -diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_zero_sdxl.cpython-310.pyc,, -diffusers/pipelines/text_to_video_synthesis/pipeline_output.py,sha256=12i4JmK2TgksR46kwOSw02McNrV7qksA4MFAw6KB6_Q,735 -diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py,sha256=GEH-FNE5tn-dTVB3yt91f_suYhtBzDnrfwEUHnxlvr4,31922 -diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py,sha256=EcVA37GTRP5XWFFmKxoGQzW0BH7KLRkM_4uwB0r-EsM,35440 -diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py,sha256=Y6oKIsoXcLxPOcmeCtZ-YNa7Jg7XX3DU1zocXDWGoIE,45882 -diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py,sha256=Z6hH40KmN_ZD8lMG7xqjW4MeGf1XVF-kyvzNnt7CaV4,64869 -diffusers/pipelines/transformers_loading_utils.py,sha256=98wKUHN89Q1nmmat046hgQxLDlnZNj9Ww4TLB5W52pQ,5281 -diffusers/pipelines/unclip/__init__.py,sha256=jBYZIN7NhTKM_Oq7ipJ4JaMXO-GtdchmFWe07gDerfA,1752 -diffusers/pipelines/unclip/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/unclip/__pycache__/pipeline_unclip.cpython-310.pyc,, -diffusers/pipelines/unclip/__pycache__/pipeline_unclip_image_variation.cpython-310.pyc,, -diffusers/pipelines/unclip/__pycache__/text_proj.cpython-310.pyc,, -diffusers/pipelines/unclip/pipeline_unclip.py,sha256=yR6XocwNUDPzOknyGqq-OepEGMxTJEu9Y99fVN3attc,22445 -diffusers/pipelines/unclip/pipeline_unclip_image_variation.py,sha256=08S5ysw2k0bw0gM3JKU_PEqWWvXcUUBraN02ejB1k7E,19347 -diffusers/pipelines/unclip/text_proj.py,sha256=UkICP2P1vL8tVJr92A6gPFuqrjl03fAs-HiFFkZDvfU,4286 -diffusers/pipelines/unidiffuser/__init__.py,sha256=GvGtf-AToJXNHxv3RAo5_I_9zPQjDFbMTAHICCt-4xY,1814 -diffusers/pipelines/unidiffuser/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/unidiffuser/__pycache__/modeling_text_decoder.cpython-310.pyc,, -diffusers/pipelines/unidiffuser/__pycache__/modeling_uvit.cpython-310.pyc,, -diffusers/pipelines/unidiffuser/__pycache__/pipeline_unidiffuser.cpython-310.pyc,, -diffusers/pipelines/unidiffuser/modeling_text_decoder.py,sha256=pG_QeaB4gBm4sVBgR_LfH_o6Jz91zRqZHrO5ac6xn6o,14111 -diffusers/pipelines/unidiffuser/modeling_uvit.py,sha256=To1JmzOZ1nIKTkHI0Df1qS3Z2ddVuEMECVEjqBhhT8M,54286 -diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py,sha256=7kc_9YFKE4TMoQDQYCZFrgbJc0kUJvxZfG5lAhn8dVo,68891 -diffusers/pipelines/visualcloze/__init__.py,sha256=Cgc6UqhelXaUEYJbmEf_kGBY90hS3wx7eIiO1C5A0g4,1502 -diffusers/pipelines/visualcloze/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/visualcloze/__pycache__/pipeline_visualcloze_combined.cpython-310.pyc,, -diffusers/pipelines/visualcloze/__pycache__/pipeline_visualcloze_generation.cpython-310.pyc,, -diffusers/pipelines/visualcloze/__pycache__/visualcloze_utils.cpython-310.pyc,, -diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py,sha256=bTEZei9dyD81tk0NBsBFV5RyEv-658XQpVvMTZWONs8,23999 -diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py,sha256=16iiLTqGAtUEYeLYQRoHpQyLxiq2dfsp4zUAiCNP5fA,45329 -diffusers/pipelines/visualcloze/visualcloze_utils.py,sha256=4DPWGxtCkC51lPoQTMe3Xlv5Lx8RKo5Uoo6tnDAwr-Y,10864 -diffusers/pipelines/wan/__init__.py,sha256=SOHLX5iUuab80gLF9ZYo54pzFWj0hiSoZU4hgWesuuQ,1677 -diffusers/pipelines/wan/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/wan/__pycache__/pipeline_output.cpython-310.pyc,, -diffusers/pipelines/wan/__pycache__/pipeline_wan.cpython-310.pyc,, -diffusers/pipelines/wan/__pycache__/pipeline_wan_i2v.cpython-310.pyc,, -diffusers/pipelines/wan/__pycache__/pipeline_wan_vace.cpython-310.pyc,, -diffusers/pipelines/wan/__pycache__/pipeline_wan_video2video.cpython-310.pyc,, -diffusers/pipelines/wan/pipeline_output.py,sha256=EjM_BX0gQD1hQ78lbUM40OtvwPqTtncuQdvP5YMBjvE,605 -diffusers/pipelines/wan/pipeline_wan.py,sha256=b2HYq72rpwhlF0E2JfP0w2xzqtDsl1xr7C_GMuO8PnI,27276 -diffusers/pipelines/wan/pipeline_wan_i2v.py,sha256=WTgNIgEC468rpMJFk5hFbxbe_ONDBfdzVjaCIsBFr0Y,34851 -diffusers/pipelines/wan/pipeline_wan_vace.py,sha256=iucIWvDyyiZqluK5nNcRFQ_XmsekaI1pcSlmVgMgBbs,48324 -diffusers/pipelines/wan/pipeline_wan_video2video.py,sha256=tHloW45GFTDBzE1fggRUZDPKqv2ZznT55-5yQ5jdKAE,33331 -diffusers/pipelines/wuerstchen/__init__.py,sha256=JSCoPCwV_rBJiCy4jbILRoAgQSITS4-j77qOPmzy284,2100 -diffusers/pipelines/wuerstchen/__pycache__/__init__.cpython-310.pyc,, -diffusers/pipelines/wuerstchen/__pycache__/modeling_paella_vq_model.cpython-310.pyc,, -diffusers/pipelines/wuerstchen/__pycache__/modeling_wuerstchen_common.cpython-310.pyc,, -diffusers/pipelines/wuerstchen/__pycache__/modeling_wuerstchen_diffnext.cpython-310.pyc,, -diffusers/pipelines/wuerstchen/__pycache__/modeling_wuerstchen_prior.cpython-310.pyc,, -diffusers/pipelines/wuerstchen/__pycache__/pipeline_wuerstchen.cpython-310.pyc,, -diffusers/pipelines/wuerstchen/__pycache__/pipeline_wuerstchen_combined.cpython-310.pyc,, -diffusers/pipelines/wuerstchen/__pycache__/pipeline_wuerstchen_prior.cpython-310.pyc,, -diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py,sha256=lmTIqLgXK7PXwZw3cGBHpEYn0oCZMtFSGsrGrXHkLTI,6925 -diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py,sha256=mx0bj5b87g590UQhoFWY_L0ht_RTIynaPQa9DLk9MTU,2713 -diffusers/pipelines/wuerstchen/modeling_wuerstchen_diffnext.py,sha256=-vKLdD7RkbDlUPpX8sgjKUTpbq7bwBdHP6y8L8kPwS0,10423 -diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py,sha256=VhlC4cLgRwe0r7Ng5vrJad9NwK_3BOFzcklNpvXzSrI,7335 -diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py,sha256=yujVa8eLz7_NHdbrqhPCg02gR0cz_soxAZQXCA0wHeM,20832 -diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py,sha256=zP-zrPDO0_DMT0BRN_kjkiuw2O9eO0f3uOYfvnvGtyo,16686 -diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py,sha256=d4iA2SSyMsLm6RhHCZliks58a2D55kLXuvB5nMAKPcg,24127 -diffusers/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -diffusers/quantizers/__init__.py,sha256=CtSmZkDrIAraPPxpKO-M5Wzeb0d7JPc-IRZnEgNUkg4,8867 -diffusers/quantizers/__pycache__/__init__.cpython-310.pyc,, -diffusers/quantizers/__pycache__/auto.cpython-310.pyc,, -diffusers/quantizers/__pycache__/base.cpython-310.pyc,, -diffusers/quantizers/__pycache__/quantization_config.cpython-310.pyc,, -diffusers/quantizers/auto.py,sha256=_rlfK4_5XQuP3Af0ArlPPWgq3E0_gwxv6z-nSQZTls8,5789 -diffusers/quantizers/base.py,sha256=lx6FKnuNn_v88zcDf4E_IDJd0PDjCD4tlhxe7xlZLi0,9682 -diffusers/quantizers/bitsandbytes/__init__.py,sha256=ILCM6ZopnzrhM_fW1oh4J_YCNsaEQAptcoTuSVgXab8,170 -diffusers/quantizers/bitsandbytes/__pycache__/__init__.cpython-310.pyc,, -diffusers/quantizers/bitsandbytes/__pycache__/bnb_quantizer.cpython-310.pyc,, -diffusers/quantizers/bitsandbytes/__pycache__/utils.cpython-310.pyc,, -diffusers/quantizers/bitsandbytes/bnb_quantizer.py,sha256=hPg52A4NwlYqW1u5M0lhccXeW-sUGAHH0ipzdI7Etj0,26268 -diffusers/quantizers/bitsandbytes/utils.py,sha256=48clXsL2q2NQ41n4dcCAhXsRNxyrkGgXYoJUUfHmHlY,13690 -diffusers/quantizers/gguf/__init__.py,sha256=2bxvfZbFr4xqm953cZaGJMgSCRiGJAWwbwKxNRQIEs4,42 -diffusers/quantizers/gguf/__pycache__/__init__.cpython-310.pyc,, -diffusers/quantizers/gguf/__pycache__/gguf_quantizer.cpython-310.pyc,, -diffusers/quantizers/gguf/__pycache__/utils.cpython-310.pyc,, -diffusers/quantizers/gguf/gguf_quantizer.py,sha256=AP6kUgk9bDh-gBuvLqWhUfaMgrpTF7O_VWEI6QDAfUM,6023 -diffusers/quantizers/gguf/utils.py,sha256=K9yweZftcOz8-4KaSr5gAo8RmXGH7rP8I775Yw9n5ww,16198 -diffusers/quantizers/quantization_config.py,sha256=vkZTAdtuV0vG0L_HOHURI-Qm91Deb__SEdmFPtvP2EE,32541 -diffusers/quantizers/quanto/__init__.py,sha256=ynS7j_VTG-QtimbyxHAaihUmi6eVqEDxA5dnKGjeS5M,46 -diffusers/quantizers/quanto/__pycache__/__init__.cpython-310.pyc,, -diffusers/quantizers/quanto/__pycache__/quanto_quantizer.cpython-310.pyc,, -diffusers/quantizers/quanto/__pycache__/utils.cpython-310.pyc,, -diffusers/quantizers/quanto/quanto_quantizer.py,sha256=nB7c2FHGzZUseLaiwTBKLSUftmLO7kZSOJGvk5I4hGw,6336 -diffusers/quantizers/quanto/utils.py,sha256=6-EaqWTbhb0dkuJ0C8XRDIpMmowPVjJ_C4rPaNNHMkc,2448 -diffusers/quantizers/torchao/__init__.py,sha256=tJimVpSGQsz3owo3yzh2SuCg7NQfiMnrHkAHyYHkmGA,662 -diffusers/quantizers/torchao/__pycache__/__init__.cpython-310.pyc,, -diffusers/quantizers/torchao/__pycache__/torchao_quantizer.cpython-310.pyc,, -diffusers/quantizers/torchao/torchao_quantizer.py,sha256=QaI_oVApbNqh9h81Eb8jlRSJo6MuhR1sO1sc7c-k4nM,14177 -diffusers/schedulers/__init__.py,sha256=HYxAaUCP3JSJXKm4FVJBGxqqwyqqgdDooE0YupPh7V4,11182 -diffusers/schedulers/__pycache__/__init__.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_amused.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_consistency_decoder.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_consistency_models.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_cosine_dpmsolver_multistep.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_ddim.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_ddim_cogvideox.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_ddim_flax.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_ddim_inverse.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_ddim_parallel.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_ddpm.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_ddpm_flax.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_ddpm_parallel.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_ddpm_wuerstchen.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_deis_multistep.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_dpm_cogvideox.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_dpmsolver_multistep.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_dpmsolver_multistep_flax.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_dpmsolver_multistep_inverse.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_dpmsolver_sde.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_dpmsolver_singlestep.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_edm_dpmsolver_multistep.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_edm_euler.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_euler_ancestral_discrete.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_euler_discrete.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_euler_discrete_flax.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_flow_match_euler_discrete.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_flow_match_heun_discrete.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_flow_match_lcm.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_heun_discrete.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_ipndm.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_k_dpm_2_ancestral_discrete.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_k_dpm_2_discrete.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_karras_ve_flax.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_lcm.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_lms_discrete.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_lms_discrete_flax.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_pndm.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_pndm_flax.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_repaint.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_sasolver.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_scm.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_sde_ve.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_sde_ve_flax.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_tcd.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_unclip.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_unipc_multistep.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_utils.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_utils_flax.cpython-310.pyc,, -diffusers/schedulers/__pycache__/scheduling_vq_diffusion.cpython-310.pyc,, -diffusers/schedulers/deprecated/__init__.py,sha256=3QlQ4gSBFu4zUkY3S5KLxd9sukbxLv8Aj4eO0Rymaq0,1349 -diffusers/schedulers/deprecated/__pycache__/__init__.cpython-310.pyc,, -diffusers/schedulers/deprecated/__pycache__/scheduling_karras_ve.cpython-310.pyc,, -diffusers/schedulers/deprecated/__pycache__/scheduling_sde_vp.cpython-310.pyc,, -diffusers/schedulers/deprecated/scheduling_karras_ve.py,sha256=CrJZjLD0QQ1mY_cE80aIDDMLzh9U6sn1iS2eEuCZFCc,9724 -diffusers/schedulers/deprecated/scheduling_sde_vp.py,sha256=4TK-8FOzNWnvlSrFIaOjhPVdNk57zeTDQYjtBdDgr3k,4294 -diffusers/schedulers/scheduling_amused.py,sha256=pioDeoYfXWP2CnAiI3-lczDdfSbkfXRC9m9REN_kmvI,6590 -diffusers/schedulers/scheduling_consistency_decoder.py,sha256=IKNAkeVZIxTPk_hS6QBvFbkjv_h8yWbtU3FmIUSXKVI,6817 -diffusers/schedulers/scheduling_consistency_models.py,sha256=wIgDb_44DIa8phTMn7WPyzAtoWVXRKSkyWLYLF1hI0Q,18699 -diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py,sha256=JZgxQvw4BJHG2RCUe9wA9dwmjKZbJa9fuR0CJVZVu-w,24873 -diffusers/schedulers/scheduling_ddim.py,sha256=a7PIMKeTepas2b_Du6ttdUBM7aAYnEI4mdA8nTTahZk,24919 -diffusers/schedulers/scheduling_ddim_cogvideox.py,sha256=Fa4lj0jnQDms_cpr26bITHnnylU4HLrWpDQkmgqLDsE,21351 -diffusers/schedulers/scheduling_ddim_flax.py,sha256=4w8-PFsXdvtfx7cE8HKwt50xBF5mR30am7BYIaYOIlY,13146 -diffusers/schedulers/scheduling_ddim_inverse.py,sha256=PPEfEduk6yAHGtok30IfG2uAn_3kvHx0H3BvgpcYzr0,17792 -diffusers/schedulers/scheduling_ddim_parallel.py,sha256=vPmnAxJvA_9Fq66rnzgkxr9hdzc1hDF41iryM51BS-E,31598 -diffusers/schedulers/scheduling_ddpm.py,sha256=g9KnAlLfVxUN0uZRcUsLmwSe9MmjSuY_xUIzmr9im6g,26019 -diffusers/schedulers/scheduling_ddpm_flax.py,sha256=IxYvg6SH17czD-N8zHdvuhGxpCI7JePypvJfOTQsLjk,12574 -diffusers/schedulers/scheduling_ddpm_parallel.py,sha256=2yoo64zf6I0RIAJ9aPDpfnbk7qv9t5GEYuHlgXy_Sog,31041 -diffusers/schedulers/scheduling_ddpm_wuerstchen.py,sha256=w7feuL53AY_363GyU1PWVGXfDHcyzJSbWOOtcXaJrAE,8938 -diffusers/schedulers/scheduling_deis_multistep.py,sha256=9ZFbpqO01jlhaKa3REl5X4gcMMChYABxujmhr_q3AG8,39716 -diffusers/schedulers/scheduling_dpm_cogvideox.py,sha256=QQ_3P5LLDrQvltlSuXYVkFl7AHSPXvlwBy_GnrJj2CA,23347 -diffusers/schedulers/scheduling_dpmsolver_multistep.py,sha256=wSYtAGNtxcWmWxp8rIj6yOK6faOi2KY3RZ2J_fRkGBY,55336 -diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py,sha256=UEiHAqI4HWvItbfTqEy90pmDtrig51Owx3I-m66jwc8,28849 -diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py,sha256=fSTExThSvMHA9nbNCq096KkXSGUkJwKeVMDsybvn860,48871 -diffusers/schedulers/scheduling_dpmsolver_sde.py,sha256=MV6XfmTZhUyJed1uJv4S857Wxlz5gZrV1-R4MEdQDig,29491 -diffusers/schedulers/scheduling_dpmsolver_singlestep.py,sha256=ulU5srNvJFpVZWqxpSwXEPmd5CUOTcET9NffFgaWzLk,54259 -diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py,sha256=0b8CYA_EQ0C2CL1AonsiyN6AC5hTBmUi0Venz9vFe90,31915 -diffusers/schedulers/scheduling_edm_euler.py,sha256=lU2uJnAOgIyi2IEr5NI396G049J2Fejm30OpltCEru0,19500 -diffusers/schedulers/scheduling_euler_ancestral_discrete.py,sha256=Z5Azbf97q3XIriLTXtPq92coqTUHHZ3gTST-2aRr5Vg,21098 -diffusers/schedulers/scheduling_euler_discrete.py,sha256=gKhj6feeZ48ORwArAwwaWWmnU5As73MP6BAhkf7zKzY,34944 -diffusers/schedulers/scheduling_euler_discrete_flax.py,sha256=UHWWpKBY-wC6wgXs8f0qMxp1LPH_VXNiTLRrK3vEXcU,10809 -diffusers/schedulers/scheduling_flow_match_euler_discrete.py,sha256=CQ-0P0WfqFHl_w84bsyDThiCbZavW_3MVYKRjwP5Lac,23837 -diffusers/schedulers/scheduling_flow_match_heun_discrete.py,sha256=Yv0bU0HKMhthpzEYULiZs-O8H1Hg934r7T17bgAh8-Y,12154 -diffusers/schedulers/scheduling_flow_match_lcm.py,sha256=gxRWyVmy49gC8Gp7kfWhsnDkTQq6oHs3-4xYO5mpPJ0,23951 -diffusers/schedulers/scheduling_heun_discrete.py,sha256=JkVoEF-vdtBweqjUs57o8JUA3K9uvGW_gHGY5f2tr_8,27708 -diffusers/schedulers/scheduling_ipndm.py,sha256=Ev9x-4E5QnubgRjwLQQimMLXhpEV7NiffKK7RGYWF-c,8768 -diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py,sha256=Zk9Q4-pIu0w0jSkpqPKuDA7UWQmCwN0uA5KYbCeSDr8,27593 -diffusers/schedulers/scheduling_k_dpm_2_discrete.py,sha256=fE24rs9yyxtAR6LAPbtYtgEwYjmv2wG0l5EGWM55vOA,26140 -diffusers/schedulers/scheduling_karras_ve_flax.py,sha256=HFO9Xjg4SBTRFe3e0UgtpDE5dAEHUWOxVA1R3avuXhQ,9630 -diffusers/schedulers/scheduling_lcm.py,sha256=Wcul9ObZdS3JcT2DlW5P9x25AFZLbJYy6Pt4eYdIHx0,32010 -diffusers/schedulers/scheduling_lms_discrete.py,sha256=YbnEAXPU4Rho1-BcMbaNJ0lk4m2dfZ5We1Ih0Cyyq_Y,24372 -diffusers/schedulers/scheduling_lms_discrete_flax.py,sha256=JSeHrj700b-rrDqWpPOUREktzrKwb4S9sHYxUwtlBxc,11077 -diffusers/schedulers/scheduling_pndm.py,sha256=n8P3-s6I6LKA6pU0yZfs6CoopFvArDerMdcSz12zA88,21731 -diffusers/schedulers/scheduling_pndm_flax.py,sha256=o2EiaUcOeeoJKsj4uABXrs3QUH_MwrgQrXCihEXay7U,21555 -diffusers/schedulers/scheduling_repaint.py,sha256=M1dddhIGRnQgnRPMxbOR1m2jgrsBkZo-fPpiSvKbT70,15762 -diffusers/schedulers/scheduling_sasolver.py,sha256=7dYuwAbKLsbr0gjGLwwusqzEhM1eBLa770sC9mNFQMI,55116 -diffusers/schedulers/scheduling_scm.py,sha256=3UgiM11atSskN-LYf9TloBC3fztxmRHu9umhBOtZ698,11402 -diffusers/schedulers/scheduling_sde_ve.py,sha256=_y-9IBk2yxbz1iwTKngBJipWSkiFrdk8ky9LDkLvYo4,13321 -diffusers/schedulers/scheduling_sde_ve_flax.py,sha256=6flCFfSIiGpAwEZK_UhjKmWjzsIXOD5UAORrpkIAbgM,12142 -diffusers/schedulers/scheduling_tcd.py,sha256=tiAvTdewr6CtMe2wn4z7_-iUDcVCvWdK27aiMxq5kJM,34749 -diffusers/schedulers/scheduling_unclip.py,sha256=T7HJXANwbhVCo5FaFkvmEuUl6geBDuz-DxVm5s81Z7o,15056 -diffusers/schedulers/scheduling_unipc_multistep.py,sha256=GNlYBYL4Pk_74A_dv7aHvWfOImW2JHqxZUK-V_0uxUY,46681 -diffusers/schedulers/scheduling_utils.py,sha256=6ArJlV6zTYUwbvcfKH0WuFFr_jBK8Qs7G0e5OZCv-qI,8708 -diffusers/schedulers/scheduling_utils_flax.py,sha256=1vxGnNAjeLqxx0VMOEkvfeDEopQqfAmr2pMMJtaKXNg,12153 -diffusers/schedulers/scheduling_vq_diffusion.py,sha256=xkYjAqC_dF5kfazvKnct6V0Rg19GdBe9Qq2AHlNYi2Y,22954 -diffusers/training_utils.py,sha256=6YUrhwteivj69J0_dqKmvsqaWZt-j-BbgSkBz38pE8E,26769 -diffusers/utils/__init__.py,sha256=9Ilx8KvWskzj3WQiJrUBckDWCGSF0-naGmjZMnwN98U,4587 -diffusers/utils/__pycache__/__init__.cpython-310.pyc,, -diffusers/utils/__pycache__/accelerate_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/constants.cpython-310.pyc,, -diffusers/utils/__pycache__/deprecation_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/doc_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_bitsandbytes_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_flax_and_transformers_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_flax_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_gguf_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_note_seq_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_onnx_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_optimum_quanto_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_pt_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_torch_and_librosa_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_torch_and_scipy_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_torch_and_torchsde_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_torch_and_transformers_and_k_diffusion_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_torch_and_transformers_and_onnx_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_torch_and_transformers_and_opencv_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_torch_and_transformers_and_sentencepiece_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_torch_and_transformers_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_torchao_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dummy_transformers_and_torch_and_note_seq_objects.cpython-310.pyc,, -diffusers/utils/__pycache__/dynamic_modules_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/export_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/hub_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/import_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/loading_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/logging.cpython-310.pyc,, -diffusers/utils/__pycache__/outputs.cpython-310.pyc,, -diffusers/utils/__pycache__/peft_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/pil_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/remote_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/source_code_parsing_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/state_dict_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/testing_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/torch_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/typing_utils.cpython-310.pyc,, -diffusers/utils/__pycache__/versions.cpython-310.pyc,, -diffusers/utils/accelerate_utils.py,sha256=ZkopOK29_6QrrdkIPidHnO078B8BtzzFSts3ul2cP8Q,1839 -diffusers/utils/constants.py,sha256=FIBvtiRBdXih2T03vmKSVYX8X0KAfoMUhZ0Sf-hjGfw,3307 -diffusers/utils/deprecation_utils.py,sha256=WWegSa1ZBX2DNfFp-L2wQuKAyGqXsc275Eua0vw7Z8o,2103 -diffusers/utils/doc_utils.py,sha256=3x-UhJrm7eNcJiQAYnYefoH7Q-V2qPmc04TzbEr25-o,1348 -diffusers/utils/dummy_bitsandbytes_objects.py,sha256=7uoVirIvcuylaDyUas94Wc9AAKJMJS3GEv-in3bEqJA,527 -diffusers/utils/dummy_flax_and_transformers_objects.py,sha256=XyiqnjacRb86sS9F_VwniBrLLEmff2cgJM2X4T_RAg4,2358 -diffusers/utils/dummy_flax_objects.py,sha256=EIyO7jYPH4yjuBIxysZWE0rka3qPLEl1TmMBt5SwXNA,5316 -diffusers/utils/dummy_gguf_objects.py,sha256=H0SYZuOON9cFlkyYSTcUJJk4skYjgjIu-wDausCm0sU,499 -diffusers/utils/dummy_note_seq_objects.py,sha256=DffX40mDzWTMCyYhKudgIeBhtqTSpiSkVzcAMRue8dY,506 -diffusers/utils/dummy_onnx_objects.py,sha256=4Z61m3P9NUwbebsK58wAKs6y32Id6UaiSRyeHXo3ecA,493 -diffusers/utils/dummy_optimum_quanto_objects.py,sha256=_k-3g7WAYcJO0-38rJrHX7aIAzkkICtj2ISiggVFiz8,529 -diffusers/utils/dummy_pt_objects.py,sha256=d6DDSSV-h2-TYWXrzj8t06vdWLhQDOPQ9hYgvcBM1o8,52845 -diffusers/utils/dummy_torch_and_librosa_objects.py,sha256=JUfqU2n3tSKHyWbjSXrpdW_jr-YbMxAvAhLlPa2_Rxs,948 -diffusers/utils/dummy_torch_and_scipy_objects.py,sha256=zOLdmqbtma5nakkdYgoErISV28yaztmBLI3wrC2Z_bU,537 -diffusers/utils/dummy_torch_and_torchsde_objects.py,sha256=EJiExfXva8tnRJEn-VaCkcII31WnPr2HqdTh3PBQ-jk,985 -diffusers/utils/dummy_torch_and_transformers_and_k_diffusion_objects.py,sha256=IMw6Qs9tTdRrMUXyM_Bc_BuJBvw0OVVHNZMOk3suF7g,1151 -diffusers/utils/dummy_torch_and_transformers_and_onnx_objects.py,sha256=SiKni7YZ-pmZrurHU3-lhbDGKOGCCVxSK3GJbrARqgU,3023 -diffusers/utils/dummy_torch_and_transformers_and_opencv_objects.py,sha256=Hvskt_HEoCkRikDyiYWQ95CsSv0fX7jYqxDxjJuNIZc,601 -diffusers/utils/dummy_torch_and_transformers_and_sentencepiece_objects.py,sha256=rauUQkG4sLSyBVeEbfp5fhJFJUGqw273oXbN_KC8NIM,1637 -diffusers/utils/dummy_torch_and_transformers_objects.py,sha256=gXYEKtwJUG1-XZ_6S9Fx_orowRGHZrFNO2CI1xZgcgA,89839 -diffusers/utils/dummy_torchao_objects.py,sha256=XiJAoV11rr_7aSgkb0vgYkdZOGjM1ouTuQZ4YrXzJ4g,502 -diffusers/utils/dummy_transformers_and_torch_and_note_seq_objects.py,sha256=z-JrPgPo2dWv-buMytUqBd6QqEx8Uha6M1cKa6gR4Dc,621 -diffusers/utils/dynamic_modules_utils.py,sha256=u3VLfAM7JpDHjW-MCBePAE-xelL3IEJrhOgaOMnq3Dc,20680 -diffusers/utils/export_utils.py,sha256=YIleDkRb3-m7a2NryRlnvSXLBbjjpkEfR7U5z-cUdO4,7749 -diffusers/utils/hub_utils.py,sha256=kSbSCJ3MHsJOytCuwiqDT-RS09WBhsraRtGmboSxW7g,23904 -diffusers/utils/import_utils.py,sha256=rX4gMytXrMFTBKimJKtoe4Ep1_dxuqIKISA7ppiLy8g,31444 -diffusers/utils/loading_utils.py,sha256=6Qgm5vRGICI7daRewJOKGSBob4tGaoBOApAHhF5waPY,5300 -diffusers/utils/logging.py,sha256=cFr8Bo_UydpL0NEeIwx7CHMUwyTWaFrZztY__nPfP3s,9470 -diffusers/utils/model_card_template.md,sha256=ZhGhzjnMT2oPLbHnC0UYRNEpVC-okH-MLKjvkYsh-Ds,550 -diffusers/utils/outputs.py,sha256=-bBcKV7uRT5nwB8QnqEN83yk4HhoJ9ToA729sFN3A9E,5030 -diffusers/utils/peft_utils.py,sha256=kedTtDw2WtpijASU1_pXepfR06BUmWpzaIvmBJ7H_L8,14131 -diffusers/utils/pil_utils.py,sha256=mNv1FfHvtdW-lpOemxwe-dNoSfSF_sptgpYELP-bA20,1979 -diffusers/utils/remote_utils.py,sha256=DHtS6_2ggX45YltaE9i_o5mb5gJd-VnGjr7UJ18wg2Q,16117 -diffusers/utils/source_code_parsing_utils.py,sha256=Mk4KHfymwXHYmV2zZdVsDSn-VJ97HvUV6fGfYSuYn8g,1863 -diffusers/utils/state_dict_utils.py,sha256=ygt17dCeHowH5QxwQldJHZpc0sS_E-aCuzpjKmAZ3HA,14470 -diffusers/utils/testing_utils.py,sha256=3Vqk3Yvv2gwt0O_5wyJV28nrO4RwmXnk8_3a1BHX7Kc,50706 -diffusers/utils/torch_utils.py,sha256=2q-WsALki6ONS-V1gYnouAK_dAA1J242tjTnv6WyZUw,7250 -diffusers/utils/typing_utils.py,sha256=yeuCJmb1t5n5rG1JRPJo33KO7tg_m9ZwSXQcPKiKyFA,3400 -diffusers/utils/versions.py,sha256=-e7XW1TzZ-tsRo9PMQHp-hNGYHuVDFzLtwg3uAJzqdI,4333 -diffusers/video_processor.py,sha256=jKcIhqCs4EUcEbKEF23u0TJs1nfZrPwdXKr-YzrZ3O0,5398 diff --git a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/REQUESTED deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/WHEEL b/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/WHEEL deleted file mode 100644 index 505164bc02d63fe6b0b3299f849a77c5f1beeb41..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/WHEEL +++ /dev/null @@ -1,5 +0,0 @@ -Wheel-Version: 1.0 -Generator: setuptools (75.8.0) -Root-Is-Purelib: true -Tag: py3-none-any - diff --git a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/entry_points.txt b/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/entry_points.txt deleted file mode 100644 index 4d8cdefa4968bbed7a122fa7a34940bf9a15360b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/entry_points.txt +++ /dev/null @@ -1,2 +0,0 @@ -[console_scripts] -diffusers-cli = diffusers.commands.diffusers_cli:main diff --git a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/top_level.txt deleted file mode 100644 index 6033efb6dbaaff9bc81792fd75a6b39d9f195aeb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers-0.34.0.dist-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -diffusers diff --git a/venv/lib/python3.10/site-packages/diffusers/__init__.py b/venv/lib/python3.10/site-packages/diffusers/__init__.py deleted file mode 100644 index 4c880d794cb96051c42a9201ff78dc4b79ee3ef8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/__init__.py +++ /dev/null @@ -1,1244 +0,0 @@ -__version__ = "0.34.0" - -from typing import TYPE_CHECKING - -from .utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - is_accelerate_available, - is_bitsandbytes_available, - is_flax_available, - is_gguf_available, - is_k_diffusion_available, - is_librosa_available, - is_note_seq_available, - is_onnx_available, - is_opencv_available, - is_optimum_quanto_available, - is_scipy_available, - is_sentencepiece_available, - is_torch_available, - is_torchao_available, - is_torchsde_available, - is_transformers_available, -) - - -# Lazy Import based on -# https://github.com/huggingface/transformers/blob/main/src/transformers/__init__.py - -# When adding a new object to this init, please add it to `_import_structure`. The `_import_structure` is a dictionary submodule to list of object names, -# and is used to defer the actual importing for when the objects are requested. -# This way `import diffusers` provides the names in the namespace without actually importing anything (and especially none of the backends). - -_import_structure = { - "configuration_utils": ["ConfigMixin"], - "hooks": [], - "loaders": ["FromOriginalModelMixin"], - "models": [], - "pipelines": [], - "quantizers.quantization_config": [], - "schedulers": [], - "utils": [ - "OptionalDependencyNotAvailable", - "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", - "is_scipy_available", - "is_torch_available", - "is_torchsde_available", - "is_transformers_available", - "is_transformers_version", - "is_unidecode_available", - "logging", - ], -} - -try: - if not is_torch_available() and not is_accelerate_available() and not is_bitsandbytes_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_bitsandbytes_objects - - _import_structure["utils.dummy_bitsandbytes_objects"] = [ - name for name in dir(dummy_bitsandbytes_objects) if not name.startswith("_") - ] -else: - _import_structure["quantizers.quantization_config"].append("BitsAndBytesConfig") - -try: - if not is_torch_available() and not is_accelerate_available() and not is_gguf_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_gguf_objects - - _import_structure["utils.dummy_gguf_objects"] = [ - name for name in dir(dummy_gguf_objects) if not name.startswith("_") - ] -else: - _import_structure["quantizers.quantization_config"].append("GGUFQuantizationConfig") - -try: - if not is_torch_available() and not is_accelerate_available() and not is_torchao_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_torchao_objects - - _import_structure["utils.dummy_torchao_objects"] = [ - name for name in dir(dummy_torchao_objects) if not name.startswith("_") - ] -else: - _import_structure["quantizers.quantization_config"].append("TorchAoConfig") - -try: - if not is_torch_available() and not is_accelerate_available() and not is_optimum_quanto_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_optimum_quanto_objects - - _import_structure["utils.dummy_optimum_quanto_objects"] = [ - name for name in dir(dummy_optimum_quanto_objects) if not name.startswith("_") - ] -else: - _import_structure["quantizers.quantization_config"].append("QuantoConfig") - -try: - if not is_onnx_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_onnx_objects # noqa F403 - - _import_structure["utils.dummy_onnx_objects"] = [ - name for name in dir(dummy_onnx_objects) if not name.startswith("_") - ] - -else: - _import_structure["pipelines"].extend(["OnnxRuntimeModel"]) - -try: - if not is_torch_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_pt_objects # noqa F403 - - _import_structure["utils.dummy_pt_objects"] = [name for name in dir(dummy_pt_objects) if not name.startswith("_")] - -else: - _import_structure["hooks"].extend( - [ - "FasterCacheConfig", - "HookRegistry", - "PyramidAttentionBroadcastConfig", - "apply_faster_cache", - "apply_pyramid_attention_broadcast", - ] - ) - _import_structure["models"].extend( - [ - "AllegroTransformer3DModel", - "AsymmetricAutoencoderKL", - "AuraFlowTransformer2DModel", - "AutoencoderDC", - "AutoencoderKL", - "AutoencoderKLAllegro", - "AutoencoderKLCogVideoX", - "AutoencoderKLCosmos", - "AutoencoderKLHunyuanVideo", - "AutoencoderKLLTXVideo", - "AutoencoderKLMagvit", - "AutoencoderKLMochi", - "AutoencoderKLTemporalDecoder", - "AutoencoderKLWan", - "AutoencoderOobleck", - "AutoencoderTiny", - "AutoModel", - "CacheMixin", - "ChromaTransformer2DModel", - "CogVideoXTransformer3DModel", - "CogView3PlusTransformer2DModel", - "CogView4Transformer2DModel", - "ConsisIDTransformer3DModel", - "ConsistencyDecoderVAE", - "ControlNetModel", - "ControlNetUnionModel", - "ControlNetXSAdapter", - "CosmosTransformer3DModel", - "DiTTransformer2DModel", - "EasyAnimateTransformer3DModel", - "FluxControlNetModel", - "FluxMultiControlNetModel", - "FluxTransformer2DModel", - "HiDreamImageTransformer2DModel", - "HunyuanDiT2DControlNetModel", - "HunyuanDiT2DModel", - "HunyuanDiT2DMultiControlNetModel", - "HunyuanVideoFramepackTransformer3DModel", - "HunyuanVideoTransformer3DModel", - "I2VGenXLUNet", - "Kandinsky3UNet", - "LatteTransformer3DModel", - "LTXVideoTransformer3DModel", - "Lumina2Transformer2DModel", - "LuminaNextDiT2DModel", - "MochiTransformer3DModel", - "ModelMixin", - "MotionAdapter", - "MultiAdapter", - "MultiControlNetModel", - "OmniGenTransformer2DModel", - "PixArtTransformer2DModel", - "PriorTransformer", - "SanaControlNetModel", - "SanaTransformer2DModel", - "SD3ControlNetModel", - "SD3MultiControlNetModel", - "SD3Transformer2DModel", - "SparseControlNetModel", - "StableAudioDiTModel", - "StableCascadeUNet", - "T2IAdapter", - "T5FilmDecoder", - "Transformer2DModel", - "TransformerTemporalModel", - "UNet1DModel", - "UNet2DConditionModel", - "UNet2DModel", - "UNet3DConditionModel", - "UNetControlNetXSModel", - "UNetMotionModel", - "UNetSpatioTemporalConditionModel", - "UVit2DModel", - "VQModel", - "WanTransformer3DModel", - "WanVACETransformer3DModel", - ] - ) - _import_structure["optimization"] = [ - "get_constant_schedule", - "get_constant_schedule_with_warmup", - "get_cosine_schedule_with_warmup", - "get_cosine_with_hard_restarts_schedule_with_warmup", - "get_linear_schedule_with_warmup", - "get_polynomial_decay_schedule_with_warmup", - "get_scheduler", - ] - _import_structure["pipelines"].extend( - [ - "AudioPipelineOutput", - "AutoPipelineForImage2Image", - "AutoPipelineForInpainting", - "AutoPipelineForText2Image", - "ConsistencyModelPipeline", - "DanceDiffusionPipeline", - "DDIMPipeline", - "DDPMPipeline", - "DiffusionPipeline", - "DiTPipeline", - "ImagePipelineOutput", - "KarrasVePipeline", - "LDMPipeline", - "LDMSuperResolutionPipeline", - "PNDMPipeline", - "RePaintPipeline", - "ScoreSdeVePipeline", - "StableDiffusionMixin", - ] - ) - _import_structure["quantizers"] = ["DiffusersQuantizer"] - _import_structure["schedulers"].extend( - [ - "AmusedScheduler", - "CMStochasticIterativeScheduler", - "CogVideoXDDIMScheduler", - "CogVideoXDPMScheduler", - "DDIMInverseScheduler", - "DDIMParallelScheduler", - "DDIMScheduler", - "DDPMParallelScheduler", - "DDPMScheduler", - "DDPMWuerstchenScheduler", - "DEISMultistepScheduler", - "DPMSolverMultistepInverseScheduler", - "DPMSolverMultistepScheduler", - "DPMSolverSinglestepScheduler", - "EDMDPMSolverMultistepScheduler", - "EDMEulerScheduler", - "EulerAncestralDiscreteScheduler", - "EulerDiscreteScheduler", - "FlowMatchEulerDiscreteScheduler", - "FlowMatchHeunDiscreteScheduler", - "FlowMatchLCMScheduler", - "HeunDiscreteScheduler", - "IPNDMScheduler", - "KarrasVeScheduler", - "KDPM2AncestralDiscreteScheduler", - "KDPM2DiscreteScheduler", - "LCMScheduler", - "PNDMScheduler", - "RePaintScheduler", - "SASolverScheduler", - "SchedulerMixin", - "SCMScheduler", - "ScoreSdeVeScheduler", - "TCDScheduler", - "UnCLIPScheduler", - "UniPCMultistepScheduler", - "VQDiffusionScheduler", - ] - ) - _import_structure["training_utils"] = ["EMAModel"] - -try: - if not (is_torch_available() and is_scipy_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_torch_and_scipy_objects # noqa F403 - - _import_structure["utils.dummy_torch_and_scipy_objects"] = [ - name for name in dir(dummy_torch_and_scipy_objects) if not name.startswith("_") - ] - -else: - _import_structure["schedulers"].extend(["LMSDiscreteScheduler"]) - -try: - if not (is_torch_available() and is_torchsde_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_torch_and_torchsde_objects # noqa F403 - - _import_structure["utils.dummy_torch_and_torchsde_objects"] = [ - name for name in dir(dummy_torch_and_torchsde_objects) if not name.startswith("_") - ] - -else: - _import_structure["schedulers"].extend(["CosineDPMSolverMultistepScheduler", "DPMSolverSDEScheduler"]) - -try: - if not (is_torch_available() and is_transformers_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_torch_and_transformers_objects # noqa F403 - - _import_structure["utils.dummy_torch_and_transformers_objects"] = [ - name for name in dir(dummy_torch_and_transformers_objects) if not name.startswith("_") - ] - -else: - _import_structure["pipelines"].extend( - [ - "AllegroPipeline", - "AltDiffusionImg2ImgPipeline", - "AltDiffusionPipeline", - "AmusedImg2ImgPipeline", - "AmusedInpaintPipeline", - "AmusedPipeline", - "AnimateDiffControlNetPipeline", - "AnimateDiffPAGPipeline", - "AnimateDiffPipeline", - "AnimateDiffSDXLPipeline", - "AnimateDiffSparseControlNetPipeline", - "AnimateDiffVideoToVideoControlNetPipeline", - "AnimateDiffVideoToVideoPipeline", - "AudioLDM2Pipeline", - "AudioLDM2ProjectionModel", - "AudioLDM2UNet2DConditionModel", - "AudioLDMPipeline", - "AuraFlowPipeline", - "BlipDiffusionControlNetPipeline", - "BlipDiffusionPipeline", - "ChromaImg2ImgPipeline", - "ChromaPipeline", - "CLIPImageProjection", - "CogVideoXFunControlPipeline", - "CogVideoXImageToVideoPipeline", - "CogVideoXPipeline", - "CogVideoXVideoToVideoPipeline", - "CogView3PlusPipeline", - "CogView4ControlPipeline", - "CogView4Pipeline", - "ConsisIDPipeline", - "Cosmos2TextToImagePipeline", - "Cosmos2VideoToWorldPipeline", - "CosmosTextToWorldPipeline", - "CosmosVideoToWorldPipeline", - "CycleDiffusionPipeline", - "EasyAnimateControlPipeline", - "EasyAnimateInpaintPipeline", - "EasyAnimatePipeline", - "FluxControlImg2ImgPipeline", - "FluxControlInpaintPipeline", - "FluxControlNetImg2ImgPipeline", - "FluxControlNetInpaintPipeline", - "FluxControlNetPipeline", - "FluxControlPipeline", - "FluxFillPipeline", - "FluxImg2ImgPipeline", - "FluxInpaintPipeline", - "FluxPipeline", - "FluxPriorReduxPipeline", - "HiDreamImagePipeline", - "HunyuanDiTControlNetPipeline", - "HunyuanDiTPAGPipeline", - "HunyuanDiTPipeline", - "HunyuanSkyreelsImageToVideoPipeline", - "HunyuanVideoFramepackPipeline", - "HunyuanVideoImageToVideoPipeline", - "HunyuanVideoPipeline", - "I2VGenXLPipeline", - "IFImg2ImgPipeline", - "IFImg2ImgSuperResolutionPipeline", - "IFInpaintingPipeline", - "IFInpaintingSuperResolutionPipeline", - "IFPipeline", - "IFSuperResolutionPipeline", - "ImageTextPipelineOutput", - "Kandinsky3Img2ImgPipeline", - "Kandinsky3Pipeline", - "KandinskyCombinedPipeline", - "KandinskyImg2ImgCombinedPipeline", - "KandinskyImg2ImgPipeline", - "KandinskyInpaintCombinedPipeline", - "KandinskyInpaintPipeline", - "KandinskyPipeline", - "KandinskyPriorPipeline", - "KandinskyV22CombinedPipeline", - "KandinskyV22ControlnetImg2ImgPipeline", - "KandinskyV22ControlnetPipeline", - "KandinskyV22Img2ImgCombinedPipeline", - "KandinskyV22Img2ImgPipeline", - "KandinskyV22InpaintCombinedPipeline", - "KandinskyV22InpaintPipeline", - "KandinskyV22Pipeline", - "KandinskyV22PriorEmb2EmbPipeline", - "KandinskyV22PriorPipeline", - "LatentConsistencyModelImg2ImgPipeline", - "LatentConsistencyModelPipeline", - "LattePipeline", - "LDMTextToImagePipeline", - "LEditsPPPipelineStableDiffusion", - "LEditsPPPipelineStableDiffusionXL", - "LTXConditionPipeline", - "LTXImageToVideoPipeline", - "LTXLatentUpsamplePipeline", - "LTXPipeline", - "Lumina2Pipeline", - "Lumina2Text2ImgPipeline", - "LuminaPipeline", - "LuminaText2ImgPipeline", - "MarigoldDepthPipeline", - "MarigoldIntrinsicsPipeline", - "MarigoldNormalsPipeline", - "MochiPipeline", - "MusicLDMPipeline", - "OmniGenPipeline", - "PaintByExamplePipeline", - "PIAPipeline", - "PixArtAlphaPipeline", - "PixArtSigmaPAGPipeline", - "PixArtSigmaPipeline", - "ReduxImageEncoder", - "SanaControlNetPipeline", - "SanaPAGPipeline", - "SanaPipeline", - "SanaSprintImg2ImgPipeline", - "SanaSprintPipeline", - "SemanticStableDiffusionPipeline", - "ShapEImg2ImgPipeline", - "ShapEPipeline", - "StableAudioPipeline", - "StableAudioProjectionModel", - "StableCascadeCombinedPipeline", - "StableCascadeDecoderPipeline", - "StableCascadePriorPipeline", - "StableDiffusion3ControlNetInpaintingPipeline", - "StableDiffusion3ControlNetPipeline", - "StableDiffusion3Img2ImgPipeline", - "StableDiffusion3InpaintPipeline", - "StableDiffusion3PAGImg2ImgPipeline", - "StableDiffusion3PAGImg2ImgPipeline", - "StableDiffusion3PAGPipeline", - "StableDiffusion3Pipeline", - "StableDiffusionAdapterPipeline", - "StableDiffusionAttendAndExcitePipeline", - "StableDiffusionControlNetImg2ImgPipeline", - "StableDiffusionControlNetInpaintPipeline", - "StableDiffusionControlNetPAGInpaintPipeline", - "StableDiffusionControlNetPAGPipeline", - "StableDiffusionControlNetPipeline", - "StableDiffusionControlNetXSPipeline", - "StableDiffusionDepth2ImgPipeline", - "StableDiffusionDiffEditPipeline", - "StableDiffusionGLIGENPipeline", - "StableDiffusionGLIGENTextImagePipeline", - "StableDiffusionImageVariationPipeline", - "StableDiffusionImg2ImgPipeline", - "StableDiffusionInpaintPipeline", - "StableDiffusionInpaintPipelineLegacy", - "StableDiffusionInstructPix2PixPipeline", - "StableDiffusionLatentUpscalePipeline", - "StableDiffusionLDM3DPipeline", - "StableDiffusionModelEditingPipeline", - "StableDiffusionPAGImg2ImgPipeline", - "StableDiffusionPAGInpaintPipeline", - "StableDiffusionPAGPipeline", - "StableDiffusionPanoramaPipeline", - "StableDiffusionParadigmsPipeline", - "StableDiffusionPipeline", - "StableDiffusionPipelineSafe", - "StableDiffusionPix2PixZeroPipeline", - "StableDiffusionSAGPipeline", - "StableDiffusionUpscalePipeline", - "StableDiffusionXLAdapterPipeline", - "StableDiffusionXLControlNetImg2ImgPipeline", - "StableDiffusionXLControlNetInpaintPipeline", - "StableDiffusionXLControlNetPAGImg2ImgPipeline", - "StableDiffusionXLControlNetPAGPipeline", - "StableDiffusionXLControlNetPipeline", - "StableDiffusionXLControlNetUnionImg2ImgPipeline", - "StableDiffusionXLControlNetUnionInpaintPipeline", - "StableDiffusionXLControlNetUnionPipeline", - "StableDiffusionXLControlNetXSPipeline", - "StableDiffusionXLImg2ImgPipeline", - "StableDiffusionXLInpaintPipeline", - "StableDiffusionXLInstructPix2PixPipeline", - "StableDiffusionXLPAGImg2ImgPipeline", - "StableDiffusionXLPAGInpaintPipeline", - "StableDiffusionXLPAGPipeline", - "StableDiffusionXLPipeline", - "StableUnCLIPImg2ImgPipeline", - "StableUnCLIPPipeline", - "StableVideoDiffusionPipeline", - "TextToVideoSDPipeline", - "TextToVideoZeroPipeline", - "TextToVideoZeroSDXLPipeline", - "UnCLIPImageVariationPipeline", - "UnCLIPPipeline", - "UniDiffuserModel", - "UniDiffuserPipeline", - "UniDiffuserTextDecoder", - "VersatileDiffusionDualGuidedPipeline", - "VersatileDiffusionImageVariationPipeline", - "VersatileDiffusionPipeline", - "VersatileDiffusionTextToImagePipeline", - "VideoToVideoSDPipeline", - "VisualClozeGenerationPipeline", - "VisualClozePipeline", - "VQDiffusionPipeline", - "WanImageToVideoPipeline", - "WanPipeline", - "WanVACEPipeline", - "WanVideoToVideoPipeline", - "WuerstchenCombinedPipeline", - "WuerstchenDecoderPipeline", - "WuerstchenPriorPipeline", - ] - ) - -try: - if not (is_torch_available() and is_transformers_available() and is_opencv_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_torch_and_transformers_and_opencv_objects # noqa F403 - - _import_structure["utils.dummy_torch_and_transformers_and_opencv_objects"] = [ - name for name in dir(dummy_torch_and_transformers_and_opencv_objects) if not name.startswith("_") - ] - -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() -except OptionalDependencyNotAvailable: - from .utils import dummy_torch_and_transformers_and_sentencepiece_objects # noqa F403 - - _import_structure["utils.dummy_torch_and_transformers_and_sentencepiece_objects"] = [ - name for name in dir(dummy_torch_and_transformers_and_sentencepiece_objects) if not name.startswith("_") - ] - -else: - _import_structure["pipelines"].extend(["KolorsImg2ImgPipeline", "KolorsPAGPipeline", "KolorsPipeline"]) - -try: - if not (is_torch_available() and is_transformers_available() and is_onnx_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_torch_and_transformers_and_onnx_objects # noqa F403 - - _import_structure["utils.dummy_torch_and_transformers_and_onnx_objects"] = [ - name for name in dir(dummy_torch_and_transformers_and_onnx_objects) if not name.startswith("_") - ] - -else: - _import_structure["pipelines"].extend( - [ - "OnnxStableDiffusionImg2ImgPipeline", - "OnnxStableDiffusionInpaintPipeline", - "OnnxStableDiffusionInpaintPipelineLegacy", - "OnnxStableDiffusionPipeline", - "OnnxStableDiffusionUpscalePipeline", - "StableDiffusionOnnxPipeline", - ] - ) - -try: - if not (is_torch_available() and is_librosa_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_torch_and_librosa_objects # noqa F403 - - _import_structure["utils.dummy_torch_and_librosa_objects"] = [ - name for name in dir(dummy_torch_and_librosa_objects) if not name.startswith("_") - ] - -else: - _import_structure["pipelines"].extend(["AudioDiffusionPipeline", "Mel"]) - -try: - if not (is_transformers_available() and is_torch_available() and is_note_seq_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_transformers_and_torch_and_note_seq_objects # noqa F403 - - _import_structure["utils.dummy_transformers_and_torch_and_note_seq_objects"] = [ - name for name in dir(dummy_transformers_and_torch_and_note_seq_objects) if not name.startswith("_") - ] - - -else: - _import_structure["pipelines"].extend(["SpectrogramDiffusionPipeline"]) - -try: - if not is_flax_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_flax_objects # noqa F403 - - _import_structure["utils.dummy_flax_objects"] = [ - name for name in dir(dummy_flax_objects) if not name.startswith("_") - ] - - -else: - _import_structure["models.controlnets.controlnet_flax"] = ["FlaxControlNetModel"] - _import_structure["models.modeling_flax_utils"] = ["FlaxModelMixin"] - _import_structure["models.unets.unet_2d_condition_flax"] = ["FlaxUNet2DConditionModel"] - _import_structure["models.vae_flax"] = ["FlaxAutoencoderKL"] - _import_structure["pipelines"].extend(["FlaxDiffusionPipeline"]) - _import_structure["schedulers"].extend( - [ - "FlaxDDIMScheduler", - "FlaxDDPMScheduler", - "FlaxDPMSolverMultistepScheduler", - "FlaxEulerDiscreteScheduler", - "FlaxKarrasVeScheduler", - "FlaxLMSDiscreteScheduler", - "FlaxPNDMScheduler", - "FlaxSchedulerMixin", - "FlaxScoreSdeVeScheduler", - ] - ) - - -try: - if not (is_flax_available() and is_transformers_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_flax_and_transformers_objects # noqa F403 - - _import_structure["utils.dummy_flax_and_transformers_objects"] = [ - name for name in dir(dummy_flax_and_transformers_objects) if not name.startswith("_") - ] - - -else: - _import_structure["pipelines"].extend( - [ - "FlaxStableDiffusionControlNetPipeline", - "FlaxStableDiffusionImg2ImgPipeline", - "FlaxStableDiffusionInpaintPipeline", - "FlaxStableDiffusionPipeline", - "FlaxStableDiffusionXLPipeline", - ] - ) - -try: - if not (is_note_seq_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_note_seq_objects # noqa F403 - - _import_structure["utils.dummy_note_seq_objects"] = [ - name for name in dir(dummy_note_seq_objects) if not name.startswith("_") - ] - - -else: - _import_structure["pipelines"].extend(["MidiProcessor"]) - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - from .configuration_utils import ConfigMixin - from .quantizers import PipelineQuantizationConfig - - try: - if not is_bitsandbytes_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_bitsandbytes_objects import * - else: - from .quantizers.quantization_config import BitsAndBytesConfig - - try: - if not is_gguf_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_gguf_objects import * - else: - from .quantizers.quantization_config import GGUFQuantizationConfig - - try: - if not is_torchao_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_torchao_objects import * - else: - from .quantizers.quantization_config import TorchAoConfig - - try: - if not is_optimum_quanto_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_optimum_quanto_objects import * - else: - from .quantizers.quantization_config import QuantoConfig - - try: - if not is_onnx_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_onnx_objects import * # noqa F403 - else: - from .pipelines import OnnxRuntimeModel - - try: - if not is_torch_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_pt_objects import * # noqa F403 - else: - from .hooks import ( - FasterCacheConfig, - HookRegistry, - PyramidAttentionBroadcastConfig, - apply_faster_cache, - apply_pyramid_attention_broadcast, - ) - from .models import ( - AllegroTransformer3DModel, - AsymmetricAutoencoderKL, - AuraFlowTransformer2DModel, - AutoencoderDC, - AutoencoderKL, - AutoencoderKLAllegro, - AutoencoderKLCogVideoX, - AutoencoderKLCosmos, - AutoencoderKLHunyuanVideo, - AutoencoderKLLTXVideo, - AutoencoderKLMagvit, - AutoencoderKLMochi, - AutoencoderKLTemporalDecoder, - AutoencoderKLWan, - AutoencoderOobleck, - AutoencoderTiny, - AutoModel, - CacheMixin, - ChromaTransformer2DModel, - CogVideoXTransformer3DModel, - CogView3PlusTransformer2DModel, - CogView4Transformer2DModel, - ConsisIDTransformer3DModel, - ConsistencyDecoderVAE, - ControlNetModel, - ControlNetUnionModel, - ControlNetXSAdapter, - CosmosTransformer3DModel, - DiTTransformer2DModel, - EasyAnimateTransformer3DModel, - FluxControlNetModel, - FluxMultiControlNetModel, - FluxTransformer2DModel, - HiDreamImageTransformer2DModel, - HunyuanDiT2DControlNetModel, - HunyuanDiT2DModel, - HunyuanDiT2DMultiControlNetModel, - HunyuanVideoFramepackTransformer3DModel, - HunyuanVideoTransformer3DModel, - I2VGenXLUNet, - Kandinsky3UNet, - LatteTransformer3DModel, - LTXVideoTransformer3DModel, - Lumina2Transformer2DModel, - LuminaNextDiT2DModel, - MochiTransformer3DModel, - ModelMixin, - MotionAdapter, - MultiAdapter, - MultiControlNetModel, - OmniGenTransformer2DModel, - PixArtTransformer2DModel, - PriorTransformer, - SanaControlNetModel, - SanaTransformer2DModel, - SD3ControlNetModel, - SD3MultiControlNetModel, - SD3Transformer2DModel, - SparseControlNetModel, - StableAudioDiTModel, - T2IAdapter, - T5FilmDecoder, - Transformer2DModel, - TransformerTemporalModel, - UNet1DModel, - UNet2DConditionModel, - UNet2DModel, - UNet3DConditionModel, - UNetControlNetXSModel, - UNetMotionModel, - UNetSpatioTemporalConditionModel, - UVit2DModel, - VQModel, - WanTransformer3DModel, - WanVACETransformer3DModel, - ) - from .optimization import ( - get_constant_schedule, - get_constant_schedule_with_warmup, - get_cosine_schedule_with_warmup, - get_cosine_with_hard_restarts_schedule_with_warmup, - get_linear_schedule_with_warmup, - get_polynomial_decay_schedule_with_warmup, - get_scheduler, - ) - from .pipelines import ( - AudioPipelineOutput, - AutoPipelineForImage2Image, - AutoPipelineForInpainting, - AutoPipelineForText2Image, - BlipDiffusionControlNetPipeline, - BlipDiffusionPipeline, - CLIPImageProjection, - ConsistencyModelPipeline, - DanceDiffusionPipeline, - DDIMPipeline, - DDPMPipeline, - DiffusionPipeline, - DiTPipeline, - ImagePipelineOutput, - KarrasVePipeline, - LDMPipeline, - LDMSuperResolutionPipeline, - PNDMPipeline, - RePaintPipeline, - ScoreSdeVePipeline, - StableDiffusionMixin, - ) - from .quantizers import DiffusersQuantizer - from .schedulers import ( - AmusedScheduler, - CMStochasticIterativeScheduler, - CogVideoXDDIMScheduler, - CogVideoXDPMScheduler, - DDIMInverseScheduler, - DDIMParallelScheduler, - DDIMScheduler, - DDPMParallelScheduler, - DDPMScheduler, - DDPMWuerstchenScheduler, - DEISMultistepScheduler, - DPMSolverMultistepInverseScheduler, - DPMSolverMultistepScheduler, - DPMSolverSinglestepScheduler, - EDMDPMSolverMultistepScheduler, - EDMEulerScheduler, - EulerAncestralDiscreteScheduler, - EulerDiscreteScheduler, - FlowMatchEulerDiscreteScheduler, - FlowMatchHeunDiscreteScheduler, - FlowMatchLCMScheduler, - HeunDiscreteScheduler, - IPNDMScheduler, - KarrasVeScheduler, - KDPM2AncestralDiscreteScheduler, - KDPM2DiscreteScheduler, - LCMScheduler, - PNDMScheduler, - RePaintScheduler, - SASolverScheduler, - SchedulerMixin, - SCMScheduler, - ScoreSdeVeScheduler, - TCDScheduler, - UnCLIPScheduler, - UniPCMultistepScheduler, - VQDiffusionScheduler, - ) - from .training_utils import EMAModel - - try: - if not (is_torch_available() and is_scipy_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_torch_and_scipy_objects import * # noqa F403 - else: - from .schedulers import LMSDiscreteScheduler - - try: - if not (is_torch_available() and is_torchsde_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_torch_and_torchsde_objects import * # noqa F403 - else: - from .schedulers import CosineDPMSolverMultistepScheduler, DPMSolverSDEScheduler - - try: - if not (is_torch_available() and is_transformers_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_torch_and_transformers_objects import * # noqa F403 - else: - from .pipelines import ( - AllegroPipeline, - AltDiffusionImg2ImgPipeline, - AltDiffusionPipeline, - AmusedImg2ImgPipeline, - AmusedInpaintPipeline, - AmusedPipeline, - AnimateDiffControlNetPipeline, - AnimateDiffPAGPipeline, - AnimateDiffPipeline, - AnimateDiffSDXLPipeline, - AnimateDiffSparseControlNetPipeline, - AnimateDiffVideoToVideoControlNetPipeline, - AnimateDiffVideoToVideoPipeline, - AudioLDM2Pipeline, - AudioLDM2ProjectionModel, - AudioLDM2UNet2DConditionModel, - AudioLDMPipeline, - AuraFlowPipeline, - ChromaImg2ImgPipeline, - ChromaPipeline, - CLIPImageProjection, - CogVideoXFunControlPipeline, - CogVideoXImageToVideoPipeline, - CogVideoXPipeline, - CogVideoXVideoToVideoPipeline, - CogView3PlusPipeline, - CogView4ControlPipeline, - CogView4Pipeline, - ConsisIDPipeline, - Cosmos2TextToImagePipeline, - Cosmos2VideoToWorldPipeline, - CosmosTextToWorldPipeline, - CosmosVideoToWorldPipeline, - CycleDiffusionPipeline, - EasyAnimateControlPipeline, - EasyAnimateInpaintPipeline, - EasyAnimatePipeline, - FluxControlImg2ImgPipeline, - FluxControlInpaintPipeline, - FluxControlNetImg2ImgPipeline, - FluxControlNetInpaintPipeline, - FluxControlNetPipeline, - FluxControlPipeline, - FluxFillPipeline, - FluxImg2ImgPipeline, - FluxInpaintPipeline, - FluxPipeline, - FluxPriorReduxPipeline, - HiDreamImagePipeline, - HunyuanDiTControlNetPipeline, - HunyuanDiTPAGPipeline, - HunyuanDiTPipeline, - HunyuanSkyreelsImageToVideoPipeline, - HunyuanVideoFramepackPipeline, - HunyuanVideoImageToVideoPipeline, - HunyuanVideoPipeline, - I2VGenXLPipeline, - IFImg2ImgPipeline, - IFImg2ImgSuperResolutionPipeline, - IFInpaintingPipeline, - IFInpaintingSuperResolutionPipeline, - IFPipeline, - IFSuperResolutionPipeline, - ImageTextPipelineOutput, - Kandinsky3Img2ImgPipeline, - Kandinsky3Pipeline, - KandinskyCombinedPipeline, - KandinskyImg2ImgCombinedPipeline, - KandinskyImg2ImgPipeline, - KandinskyInpaintCombinedPipeline, - KandinskyInpaintPipeline, - KandinskyPipeline, - KandinskyPriorPipeline, - KandinskyV22CombinedPipeline, - KandinskyV22ControlnetImg2ImgPipeline, - KandinskyV22ControlnetPipeline, - KandinskyV22Img2ImgCombinedPipeline, - KandinskyV22Img2ImgPipeline, - KandinskyV22InpaintCombinedPipeline, - KandinskyV22InpaintPipeline, - KandinskyV22Pipeline, - KandinskyV22PriorEmb2EmbPipeline, - KandinskyV22PriorPipeline, - LatentConsistencyModelImg2ImgPipeline, - LatentConsistencyModelPipeline, - LattePipeline, - LDMTextToImagePipeline, - LEditsPPPipelineStableDiffusion, - LEditsPPPipelineStableDiffusionXL, - LTXConditionPipeline, - LTXImageToVideoPipeline, - LTXLatentUpsamplePipeline, - LTXPipeline, - Lumina2Pipeline, - Lumina2Text2ImgPipeline, - LuminaPipeline, - LuminaText2ImgPipeline, - MarigoldDepthPipeline, - MarigoldIntrinsicsPipeline, - MarigoldNormalsPipeline, - MochiPipeline, - MusicLDMPipeline, - OmniGenPipeline, - PaintByExamplePipeline, - PIAPipeline, - PixArtAlphaPipeline, - PixArtSigmaPAGPipeline, - PixArtSigmaPipeline, - ReduxImageEncoder, - SanaControlNetPipeline, - SanaPAGPipeline, - SanaPipeline, - SanaSprintImg2ImgPipeline, - SanaSprintPipeline, - SemanticStableDiffusionPipeline, - ShapEImg2ImgPipeline, - ShapEPipeline, - StableAudioPipeline, - StableAudioProjectionModel, - StableCascadeCombinedPipeline, - StableCascadeDecoderPipeline, - StableCascadePriorPipeline, - StableDiffusion3ControlNetInpaintingPipeline, - StableDiffusion3ControlNetPipeline, - StableDiffusion3Img2ImgPipeline, - StableDiffusion3InpaintPipeline, - StableDiffusion3PAGImg2ImgPipeline, - StableDiffusion3PAGPipeline, - StableDiffusion3Pipeline, - StableDiffusionAdapterPipeline, - StableDiffusionAttendAndExcitePipeline, - StableDiffusionControlNetImg2ImgPipeline, - StableDiffusionControlNetInpaintPipeline, - StableDiffusionControlNetPAGInpaintPipeline, - StableDiffusionControlNetPAGPipeline, - StableDiffusionControlNetPipeline, - StableDiffusionControlNetXSPipeline, - StableDiffusionDepth2ImgPipeline, - StableDiffusionDiffEditPipeline, - StableDiffusionGLIGENPipeline, - StableDiffusionGLIGENTextImagePipeline, - StableDiffusionImageVariationPipeline, - StableDiffusionImg2ImgPipeline, - StableDiffusionInpaintPipeline, - StableDiffusionInpaintPipelineLegacy, - StableDiffusionInstructPix2PixPipeline, - StableDiffusionLatentUpscalePipeline, - StableDiffusionLDM3DPipeline, - StableDiffusionModelEditingPipeline, - StableDiffusionPAGImg2ImgPipeline, - StableDiffusionPAGInpaintPipeline, - StableDiffusionPAGPipeline, - StableDiffusionPanoramaPipeline, - StableDiffusionParadigmsPipeline, - StableDiffusionPipeline, - StableDiffusionPipelineSafe, - StableDiffusionPix2PixZeroPipeline, - StableDiffusionSAGPipeline, - StableDiffusionUpscalePipeline, - StableDiffusionXLAdapterPipeline, - StableDiffusionXLControlNetImg2ImgPipeline, - StableDiffusionXLControlNetInpaintPipeline, - StableDiffusionXLControlNetPAGImg2ImgPipeline, - StableDiffusionXLControlNetPAGPipeline, - StableDiffusionXLControlNetPipeline, - StableDiffusionXLControlNetUnionImg2ImgPipeline, - StableDiffusionXLControlNetUnionInpaintPipeline, - StableDiffusionXLControlNetUnionPipeline, - StableDiffusionXLControlNetXSPipeline, - StableDiffusionXLImg2ImgPipeline, - StableDiffusionXLInpaintPipeline, - StableDiffusionXLInstructPix2PixPipeline, - StableDiffusionXLPAGImg2ImgPipeline, - StableDiffusionXLPAGInpaintPipeline, - StableDiffusionXLPAGPipeline, - StableDiffusionXLPipeline, - StableUnCLIPImg2ImgPipeline, - StableUnCLIPPipeline, - StableVideoDiffusionPipeline, - TextToVideoSDPipeline, - TextToVideoZeroPipeline, - TextToVideoZeroSDXLPipeline, - UnCLIPImageVariationPipeline, - UnCLIPPipeline, - UniDiffuserModel, - UniDiffuserPipeline, - UniDiffuserTextDecoder, - VersatileDiffusionDualGuidedPipeline, - VersatileDiffusionImageVariationPipeline, - VersatileDiffusionPipeline, - VersatileDiffusionTextToImagePipeline, - VideoToVideoSDPipeline, - VisualClozeGenerationPipeline, - VisualClozePipeline, - VQDiffusionPipeline, - WanImageToVideoPipeline, - WanPipeline, - WanVACEPipeline, - WanVideoToVideoPipeline, - WuerstchenCombinedPipeline, - WuerstchenDecoderPipeline, - WuerstchenPriorPipeline, - ) - - 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() - except OptionalDependencyNotAvailable: - from .utils.dummy_torch_and_transformers_and_sentencepiece_objects import * # noqa F403 - else: - from .pipelines import KolorsImg2ImgPipeline, KolorsPAGPipeline, KolorsPipeline - - try: - if not (is_torch_available() and is_transformers_available() and is_opencv_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_torch_and_transformers_and_opencv_objects import * # noqa F403 - else: - from .pipelines import ConsisIDPipeline - - try: - if not (is_torch_available() and is_transformers_available() and is_onnx_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_torch_and_transformers_and_onnx_objects import * # noqa F403 - else: - from .pipelines import ( - OnnxStableDiffusionImg2ImgPipeline, - OnnxStableDiffusionInpaintPipeline, - OnnxStableDiffusionInpaintPipelineLegacy, - OnnxStableDiffusionPipeline, - OnnxStableDiffusionUpscalePipeline, - StableDiffusionOnnxPipeline, - ) - - try: - if not (is_torch_available() and is_librosa_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_torch_and_librosa_objects import * # noqa F403 - else: - from .pipelines import AudioDiffusionPipeline, Mel - - try: - if not (is_transformers_available() and is_torch_available() and is_note_seq_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_transformers_and_torch_and_note_seq_objects import * # noqa F403 - else: - from .pipelines import SpectrogramDiffusionPipeline - - try: - if not is_flax_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_flax_objects import * # noqa F403 - else: - from .models.controlnets.controlnet_flax import FlaxControlNetModel - from .models.modeling_flax_utils import FlaxModelMixin - from .models.unets.unet_2d_condition_flax import FlaxUNet2DConditionModel - from .models.vae_flax import FlaxAutoencoderKL - from .pipelines import FlaxDiffusionPipeline - from .schedulers import ( - FlaxDDIMScheduler, - FlaxDDPMScheduler, - FlaxDPMSolverMultistepScheduler, - FlaxEulerDiscreteScheduler, - FlaxKarrasVeScheduler, - FlaxLMSDiscreteScheduler, - FlaxPNDMScheduler, - FlaxSchedulerMixin, - FlaxScoreSdeVeScheduler, - ) - - try: - if not (is_flax_available() and is_transformers_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_flax_and_transformers_objects import * # noqa F403 - else: - from .pipelines import ( - FlaxStableDiffusionControlNetPipeline, - FlaxStableDiffusionImg2ImgPipeline, - FlaxStableDiffusionInpaintPipeline, - FlaxStableDiffusionPipeline, - FlaxStableDiffusionXLPipeline, - ) - - try: - if not (is_note_seq_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_note_seq_objects import * # noqa F403 - else: - from .pipelines import MidiProcessor - -else: - import sys - - sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, - extra_objects={"__version__": __version__}, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index ea346bbfa4a58f8d3fd9d1a071b60b740540beaf..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/__pycache__/callbacks.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/__pycache__/callbacks.cpython-310.pyc deleted file mode 100644 index 0aa3062a5b223d249fbb02bb5440a326f8164a72..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/__pycache__/callbacks.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/__pycache__/configuration_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/__pycache__/configuration_utils.cpython-310.pyc deleted file mode 100644 index 10e3f1153ca530af55cdd3da0ccdacbe31279190..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/__pycache__/configuration_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/__pycache__/dependency_versions_check.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/__pycache__/dependency_versions_check.cpython-310.pyc deleted file mode 100644 index 191540351e9e59cae17903ef0964cde6a1b35547..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/__pycache__/dependency_versions_check.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/__pycache__/dependency_versions_table.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/__pycache__/dependency_versions_table.cpython-310.pyc deleted file mode 100644 index b566cddf728f542bae1d5374869958c8d25e3279..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/__pycache__/dependency_versions_table.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/__pycache__/image_processor.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/__pycache__/image_processor.cpython-310.pyc deleted file mode 100644 index 92c1a04094591c86ede889256a68f40e6e712b64..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/__pycache__/image_processor.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/__pycache__/optimization.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/__pycache__/optimization.cpython-310.pyc deleted file mode 100644 index 8e53f972c13d66e1b987ee996101ef95ef209664..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/__pycache__/optimization.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/__pycache__/training_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/__pycache__/training_utils.cpython-310.pyc deleted file mode 100644 index 6a3d45a55142c6e3eadee9ce6da712ee28841ba8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/__pycache__/training_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/__pycache__/video_processor.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/__pycache__/video_processor.cpython-310.pyc deleted file mode 100644 index 9c0801bc9c621ddc517c814920f042f6b151d765..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/__pycache__/video_processor.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/callbacks.py b/venv/lib/python3.10/site-packages/diffusers/callbacks.py deleted file mode 100644 index 4b8b15368c4736ce54282e92231a752558663cb6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/callbacks.py +++ /dev/null @@ -1,209 +0,0 @@ -from typing import Any, Dict, List - -from .configuration_utils import ConfigMixin, register_to_config -from .utils import CONFIG_NAME - - -class PipelineCallback(ConfigMixin): - """ - Base class for all the official callbacks used in a pipeline. This class provides a structure for implementing - custom callbacks and ensures that all callbacks have a consistent interface. - - Please implement the following: - `tensor_inputs`: This should return a list of tensor inputs specific to your callback. You will only be able to - include - variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. - `callback_fn`: This method defines the core functionality of your callback. - """ - - config_name = CONFIG_NAME - - @register_to_config - def __init__(self, cutoff_step_ratio=1.0, cutoff_step_index=None): - super().__init__() - - if (cutoff_step_ratio is None and cutoff_step_index is None) or ( - cutoff_step_ratio is not None and cutoff_step_index is not None - ): - raise ValueError("Either cutoff_step_ratio or cutoff_step_index should be provided, not both or none.") - - if cutoff_step_ratio is not None and ( - not isinstance(cutoff_step_ratio, float) or not (0.0 <= cutoff_step_ratio <= 1.0) - ): - raise ValueError("cutoff_step_ratio must be a float between 0.0 and 1.0.") - - @property - 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]: - 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]: - return self.callback_fn(pipeline, step_index, timestep, callback_kwargs) - - -class MultiPipelineCallbacks: - """ - This class is designed to handle multiple pipeline callbacks. It accepts a list of PipelineCallback objects and - provides a unified interface for calling all of them. - """ - - def __init__(self, callbacks: List[PipelineCallback]): - self.callbacks = callbacks - - @property - 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]: - """ - Calls all the callbacks in order with the given arguments and returns the final callback_kwargs. - """ - for callback in self.callbacks: - callback_kwargs = callback(pipeline, step_index, timestep, callback_kwargs) - - return callback_kwargs - - -class SDCFGCutoffCallback(PipelineCallback): - """ - Callback function for Stable Diffusion Pipelines. After certain number of steps (set by `cutoff_step_ratio` or - `cutoff_step_index`), this callback will disable the CFG. - - Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step. - """ - - tensor_inputs = ["prompt_embeds"] - - 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 - - # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio - cutoff_step = ( - cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio) - ) - - if step_index == cutoff_step: - prompt_embeds = callback_kwargs[self.tensor_inputs[0]] - prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens. - - pipeline._guidance_scale = 0.0 - - callback_kwargs[self.tensor_inputs[0]] = prompt_embeds - return callback_kwargs - - -class SDXLCFGCutoffCallback(PipelineCallback): - """ - Callback function for the base Stable Diffusion XL Pipelines. After certain number of steps (set by - `cutoff_step_ratio` or `cutoff_step_index`), this callback will disable the CFG. - - Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step. - """ - - tensor_inputs = [ - "prompt_embeds", - "add_text_embeds", - "add_time_ids", - ] - - 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 - - # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio - cutoff_step = ( - cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio) - ) - - if step_index == cutoff_step: - prompt_embeds = callback_kwargs[self.tensor_inputs[0]] - prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens. - - add_text_embeds = callback_kwargs[self.tensor_inputs[1]] - add_text_embeds = add_text_embeds[-1:] # "-1" denotes the embeddings for conditional pooled text tokens - - add_time_ids = callback_kwargs[self.tensor_inputs[2]] - add_time_ids = add_time_ids[-1:] # "-1" denotes the embeddings for conditional added time vector - - pipeline._guidance_scale = 0.0 - - callback_kwargs[self.tensor_inputs[0]] = prompt_embeds - callback_kwargs[self.tensor_inputs[1]] = add_text_embeds - callback_kwargs[self.tensor_inputs[2]] = add_time_ids - - return callback_kwargs - - -class SDXLControlnetCFGCutoffCallback(PipelineCallback): - """ - Callback function for the Controlnet Stable Diffusion XL Pipelines. After certain number of steps (set by - `cutoff_step_ratio` or `cutoff_step_index`), this callback will disable the CFG. - - Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step. - """ - - tensor_inputs = [ - "prompt_embeds", - "add_text_embeds", - "add_time_ids", - "image", - ] - - 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 - - # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio - cutoff_step = ( - cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio) - ) - - if step_index == cutoff_step: - prompt_embeds = callback_kwargs[self.tensor_inputs[0]] - prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens. - - add_text_embeds = callback_kwargs[self.tensor_inputs[1]] - add_text_embeds = add_text_embeds[-1:] # "-1" denotes the embeddings for conditional pooled text tokens - - add_time_ids = callback_kwargs[self.tensor_inputs[2]] - add_time_ids = add_time_ids[-1:] # "-1" denotes the embeddings for conditional added time vector - - # For Controlnet - image = callback_kwargs[self.tensor_inputs[3]] - image = image[-1:] - - pipeline._guidance_scale = 0.0 - - callback_kwargs[self.tensor_inputs[0]] = prompt_embeds - callback_kwargs[self.tensor_inputs[1]] = add_text_embeds - callback_kwargs[self.tensor_inputs[2]] = add_time_ids - callback_kwargs[self.tensor_inputs[3]] = image - - return callback_kwargs - - -class IPAdapterScaleCutoffCallback(PipelineCallback): - """ - Callback function for any pipeline that inherits `IPAdapterMixin`. After certain number of steps (set by - `cutoff_step_ratio` or `cutoff_step_index`), this callback will set the IP Adapter scale to `0.0`. - - Note: This callback mutates the IP Adapter attention processors by setting the scale to 0.0 after the cutoff step. - """ - - tensor_inputs = [] - - 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 - - # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio - cutoff_step = ( - cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio) - ) - - if step_index == cutoff_step: - pipeline.set_ip_adapter_scale(0.0) - return callback_kwargs diff --git a/venv/lib/python3.10/site-packages/diffusers/commands/__init__.py b/venv/lib/python3.10/site-packages/diffusers/commands/__init__.py deleted file mode 100644 index 3a8c10147e8bdc6be186193fa7775460b474a9ba..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/commands/__init__.py +++ /dev/null @@ -1,27 +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 abc import ABC, abstractmethod -from argparse import ArgumentParser - - -class BaseDiffusersCLICommand(ABC): - @staticmethod - @abstractmethod - def register_subcommand(parser: ArgumentParser): - raise NotImplementedError() - - @abstractmethod - def run(self): - raise NotImplementedError() diff --git a/venv/lib/python3.10/site-packages/diffusers/commands/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/commands/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 8c99b471b8b165a887b4d2b09f44bbfdeb5bece9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/commands/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/commands/__pycache__/diffusers_cli.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/commands/__pycache__/diffusers_cli.cpython-310.pyc deleted file mode 100644 index d7078c44b452e0b6cb1302692af585c23056a8a1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/commands/__pycache__/diffusers_cli.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/commands/__pycache__/env.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/commands/__pycache__/env.cpython-310.pyc deleted file mode 100644 index 35838c564f2bb60ec98aaca6751cad13765b88e1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/commands/__pycache__/env.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/commands/__pycache__/fp16_safetensors.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/commands/__pycache__/fp16_safetensors.cpython-310.pyc deleted file mode 100644 index 03569fc903fa3313bcfe075b19efca692e6db807..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/commands/__pycache__/fp16_safetensors.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/commands/diffusers_cli.py b/venv/lib/python3.10/site-packages/diffusers/commands/diffusers_cli.py deleted file mode 100644 index 3c744c5c4c9580ab367158486f72715cce8d1c8b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/commands/diffusers_cli.py +++ /dev/null @@ -1,43 +0,0 @@ -#!/usr/bin/env python -# 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 argparse import ArgumentParser - -from .env import EnvironmentCommand -from .fp16_safetensors import FP16SafetensorsCommand - - -def main(): - parser = ArgumentParser("Diffusers CLI tool", usage="diffusers-cli []") - commands_parser = parser.add_subparsers(help="diffusers-cli command helpers") - - # Register commands - EnvironmentCommand.register_subcommand(commands_parser) - FP16SafetensorsCommand.register_subcommand(commands_parser) - - # Let's go - args = parser.parse_args() - - if not hasattr(args, "func"): - parser.print_help() - exit(1) - - # Run - service = args.func(args) - service.run() - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/diffusers/commands/env.py b/venv/lib/python3.10/site-packages/diffusers/commands/env.py deleted file mode 100644 index 58f31d478bf35420b3bd324b0aaf57bef934b380..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/commands/env.py +++ /dev/null @@ -1,180 +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 platform -import subprocess -from argparse import ArgumentParser - -import huggingface_hub - -from .. import __version__ as version -from ..utils import ( - is_accelerate_available, - is_bitsandbytes_available, - is_flax_available, - is_google_colab, - is_peft_available, - is_safetensors_available, - is_torch_available, - is_transformers_available, - is_xformers_available, -) -from . import BaseDiffusersCLICommand - - -def info_command_factory(_): - return EnvironmentCommand() - - -class EnvironmentCommand(BaseDiffusersCLICommand): - @staticmethod - def register_subcommand(parser: ArgumentParser) -> None: - download_parser = parser.add_parser("env") - download_parser.set_defaults(func=info_command_factory) - - def run(self) -> dict: - hub_version = huggingface_hub.__version__ - - safetensors_version = "not installed" - if is_safetensors_available(): - import safetensors - - safetensors_version = safetensors.__version__ - - pt_version = "not installed" - pt_cuda_available = "NA" - if is_torch_available(): - import torch - - pt_version = torch.__version__ - pt_cuda_available = torch.cuda.is_available() - - flax_version = "not installed" - jax_version = "not installed" - jaxlib_version = "not installed" - jax_backend = "NA" - if is_flax_available(): - import flax - import jax - import jaxlib - - flax_version = flax.__version__ - jax_version = jax.__version__ - jaxlib_version = jaxlib.__version__ - jax_backend = jax.lib.xla_bridge.get_backend().platform - - transformers_version = "not installed" - if is_transformers_available(): - import transformers - - transformers_version = transformers.__version__ - - accelerate_version = "not installed" - if is_accelerate_available(): - import accelerate - - accelerate_version = accelerate.__version__ - - peft_version = "not installed" - if is_peft_available(): - import peft - - peft_version = peft.__version__ - - bitsandbytes_version = "not installed" - if is_bitsandbytes_available(): - import bitsandbytes - - bitsandbytes_version = bitsandbytes.__version__ - - xformers_version = "not installed" - if is_xformers_available(): - import xformers - - xformers_version = xformers.__version__ - - platform_info = platform.platform() - - is_google_colab_str = "Yes" if is_google_colab() else "No" - - accelerator = "NA" - if platform.system() in {"Linux", "Windows"}: - try: - sp = subprocess.Popen( - ["nvidia-smi", "--query-gpu=gpu_name,memory.total", "--format=csv,noheader"], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - out_str, _ = sp.communicate() - out_str = out_str.decode("utf-8") - - if len(out_str) > 0: - accelerator = out_str.strip() - except FileNotFoundError: - pass - elif platform.system() == "Darwin": # Mac OS - try: - sp = subprocess.Popen( - ["system_profiler", "SPDisplaysDataType"], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - out_str, _ = sp.communicate() - out_str = out_str.decode("utf-8") - - start = out_str.find("Chipset Model:") - if start != -1: - start += len("Chipset Model:") - end = out_str.find("\n", start) - accelerator = out_str[start:end].strip() - - start = out_str.find("VRAM (Total):") - if start != -1: - start += len("VRAM (Total):") - end = out_str.find("\n", start) - accelerator += " VRAM: " + out_str[start:end].strip() - except FileNotFoundError: - pass - else: - print("It seems you are running an unusual OS. Could you fill in the accelerator manually?") - - info = { - "🤗 Diffusers version": version, - "Platform": platform_info, - "Running on Google Colab?": is_google_colab_str, - "Python version": platform.python_version(), - "PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})", - "Flax version (CPU?/GPU?/TPU?)": f"{flax_version} ({jax_backend})", - "Jax version": jax_version, - "JaxLib version": jaxlib_version, - "Huggingface_hub version": hub_version, - "Transformers version": transformers_version, - "Accelerate version": accelerate_version, - "PEFT version": peft_version, - "Bitsandbytes version": bitsandbytes_version, - "Safetensors version": safetensors_version, - "xFormers version": xformers_version, - "Accelerator": accelerator, - "Using GPU in script?": "", - "Using distributed or parallel set-up in script?": "", - } - - print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n") - print(self.format_dict(info)) - - return info - - @staticmethod - def format_dict(d: dict) -> str: - return "\n".join([f"- {prop}: {val}" for prop, val in d.items()]) + "\n" diff --git a/venv/lib/python3.10/site-packages/diffusers/commands/fp16_safetensors.py b/venv/lib/python3.10/site-packages/diffusers/commands/fp16_safetensors.py deleted file mode 100644 index ef60f237ae15abac34ead58427746aff4fe9328a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/commands/fp16_safetensors.py +++ /dev/null @@ -1,132 +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. - -""" -Usage example: - diffusers-cli fp16_safetensors --ckpt_id=openai/shap-e --fp16 --use_safetensors -""" - -import glob -import json -import warnings -from argparse import ArgumentParser, Namespace -from importlib import import_module - -import huggingface_hub -import torch -from huggingface_hub import hf_hub_download -from packaging import version - -from ..utils import logging -from . import BaseDiffusersCLICommand - - -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." - ) - return FP16SafetensorsCommand(args.ckpt_id, args.fp16, args.use_safetensors) - - -class FP16SafetensorsCommand(BaseDiffusersCLICommand): - @staticmethod - def register_subcommand(parser: ArgumentParser): - conversion_parser = parser.add_parser("fp16_safetensors") - conversion_parser.add_argument( - "--ckpt_id", - type=str, - help="Repo id of the checkpoints on which to run the conversion. Example: 'openai/shap-e'.", - ) - conversion_parser.add_argument( - "--fp16", action="store_true", help="If serializing the variables in FP16 precision." - ) - conversion_parser.add_argument( - "--use_safetensors", action="store_true", help="If serializing in the safetensors format." - ) - conversion_parser.add_argument( - "--use_auth_token", - action="store_true", - help="When working with checkpoints having private visibility. When used `huggingface-cli login` needs to be run beforehand.", - ) - conversion_parser.set_defaults(func=conversion_command_factory) - - def __init__(self, ckpt_id: str, fp16: bool, use_safetensors: bool): - self.logger = logging.get_logger("diffusers-cli/fp16_safetensors") - self.ckpt_id = ckpt_id - self.local_ckpt_dir = f"/tmp/{ckpt_id}" - self.fp16 = fp16 - - self.use_safetensors = use_safetensors - - if not self.use_safetensors and not self.fp16: - raise NotImplementedError( - "When `use_safetensors` and `fp16` both are False, then this command is of no use." - ) - - def run(self): - if version.parse(huggingface_hub.__version__) < version.parse("0.9.0"): - raise ImportError( - "The huggingface_hub version must be >= 0.9.0 to use this command. Please update your huggingface_hub" - " installation." - ) - else: - from huggingface_hub import create_commit - from huggingface_hub._commit_api import CommitOperationAdd - - model_index = hf_hub_download(repo_id=self.ckpt_id, filename="model_index.json") - with open(model_index, "r") as f: - pipeline_class_name = json.load(f)["_class_name"] - 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. - pipeline = pipeline_class.from_pretrained( - self.ckpt_id, torch_dtype=torch.float16 if self.fp16 else torch.float32 - ) - pipeline.save_pretrained( - self.local_ckpt_dir, - safe_serialization=True if self.use_safetensors else False, - variant="fp16" if self.fp16 else None, - ) - self.logger.info(f"Pipeline locally saved to {self.local_ckpt_dir}.") - - # Fetch all the paths. - if self.fp16: - modified_paths = glob.glob(f"{self.local_ckpt_dir}/*/*.fp16.*") - elif self.use_safetensors: - modified_paths = glob.glob(f"{self.local_ckpt_dir}/*/*.safetensors") - - # Prepare for the PR. - commit_message = f"Serialize variables with FP16: {self.fp16} and safetensors: {self.use_safetensors}." - operations = [] - for path in modified_paths: - operations.append(CommitOperationAdd(path_in_repo="/".join(path.split("/")[4:]), path_or_fileobj=path)) - - # Open the PR. - commit_description = ( - "Variables converted by the [`diffusers`' `fp16_safetensors`" - " CLI](https://github.com/huggingface/diffusers/blob/main/src/diffusers/commands/fp16_safetensors.py)." - ) - hub_pr_url = create_commit( - repo_id=self.ckpt_id, - operations=operations, - commit_message=commit_message, - commit_description=commit_description, - repo_type="model", - create_pr=True, - ).pr_url - self.logger.info(f"PR created here: {hub_pr_url}.") diff --git a/venv/lib/python3.10/site-packages/diffusers/configuration_utils.py b/venv/lib/python3.10/site-packages/diffusers/configuration_utils.py deleted file mode 100644 index f9b652bbc0216a698bb206f31c139a7e3829226a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/configuration_utils.py +++ /dev/null @@ -1,760 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# Copyright (c) 2022, NVIDIA CORPORATION. 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. -"""ConfigMixin base class and utilities.""" - -import dataclasses -import functools -import importlib -import inspect -import json -import os -import re -from collections import OrderedDict -from pathlib import Path -from typing import Any, Dict, Optional, Tuple, Union - -import numpy as np -from huggingface_hub import DDUFEntry, create_repo, hf_hub_download -from huggingface_hub.utils import ( - EntryNotFoundError, - RepositoryNotFoundError, - RevisionNotFoundError, - validate_hf_hub_args, -) -from requests import HTTPError -from typing_extensions import Self - -from . import __version__ -from .utils import ( - HUGGINGFACE_CO_RESOLVE_ENDPOINT, - DummyObject, - deprecate, - extract_commit_hash, - http_user_agent, - logging, -) - - -logger = logging.get_logger(__name__) - -_re_configuration_file = re.compile(r"config\.(.*)\.json") - - -class FrozenDict(OrderedDict): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - for key, value in self.items(): - setattr(self, key, value) - - self.__frozen = True - - def __delitem__(self, *args, **kwargs): - raise Exception(f"You cannot use ``__delitem__`` on a {self.__class__.__name__} instance.") - - def setdefault(self, *args, **kwargs): - raise Exception(f"You cannot use ``setdefault`` on a {self.__class__.__name__} instance.") - - def pop(self, *args, **kwargs): - raise Exception(f"You cannot use ``pop`` on a {self.__class__.__name__} instance.") - - def update(self, *args, **kwargs): - raise Exception(f"You cannot use ``update`` on a {self.__class__.__name__} instance.") - - def __setattr__(self, name, value): - if hasattr(self, "__frozen") and self.__frozen: - raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.") - super().__setattr__(name, value) - - def __setitem__(self, name, value): - if hasattr(self, "__frozen") and self.__frozen: - raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.") - super().__setitem__(name, value) - - -class ConfigMixin: - r""" - Base class for all configuration classes. All configuration parameters are stored under `self.config`. Also - provides the [`~ConfigMixin.from_config`] and [`~ConfigMixin.save_config`] methods for loading, downloading, and - saving classes that inherit from [`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 - 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 - should only have a `kwargs` argument if at least one argument is deprecated (should be overridden by - subclass). - """ - - config_name = None - ignore_for_config = [] - has_compatibles = False - - _deprecated_kwargs = [] - - def register_to_config(self, **kwargs): - if self.config_name is None: - raise NotImplementedError(f"Make sure that {self.__class__} has defined a class name `config_name`") - # Special case for `kwargs` used in deprecation warning added to schedulers - # TODO: remove this when we remove the deprecation warning, and the `kwargs` argument, - # or solve in a more general way. - kwargs.pop("kwargs", None) - - if not hasattr(self, "_internal_dict"): - internal_dict = kwargs - else: - previous_dict = dict(self._internal_dict) - internal_dict = {**self._internal_dict, **kwargs} - logger.debug(f"Updating config from {previous_dict} to {internal_dict}") - - self._internal_dict = FrozenDict(internal_dict) - - def __getattr__(self, name: str) -> Any: - """The only reason we overwrite `getattr` here is to gracefully deprecate accessing - config attributes directly. See https://github.com/huggingface/diffusers/pull/3129 - - This function is mostly copied from PyTorch's __getattr__ overwrite: - https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module - """ - - is_in_config = "_internal_dict" in self.__dict__ and hasattr(self.__dict__["_internal_dict"], name) - is_attribute = name in self.__dict__ - - if is_in_config and not is_attribute: - deprecation_message = f"Accessing config attribute `{name}` directly via '{type(self).__name__}' object attribute is deprecated. Please access '{name}' over '{type(self).__name__}'s config object instead, e.g. 'scheduler.config.{name}'." - deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False) - return self._internal_dict[name] - - 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): - """ - Save a configuration object to the directory specified in `save_directory` so that it can be reloaded using the - [`~ConfigMixin.from_config`] class method. - - Args: - save_directory (`str` or `os.PathLike`): - Directory where the configuration JSON file is saved (will be created if it does not exist). - push_to_hub (`bool`, *optional*, defaults to `False`): - 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*): - Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. - """ - if os.path.isfile(save_directory): - raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file") - - os.makedirs(save_directory, exist_ok=True) - - # If we save using the predefined names, we can load using `from_config` - output_config_file = os.path.join(save_directory, self.config_name) - - self.to_json_file(output_config_file) - logger.info(f"Configuration saved in {output_config_file}") - - 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]) - repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id - - self._upload_folder( - save_directory, - repo_id, - token=token, - commit_message=commit_message, - create_pr=create_pr, - ) - - @classmethod - def from_config( - cls, config: Union[FrozenDict, Dict[str, Any]] = None, return_unused_kwargs=False, **kwargs - ) -> Union[Self, Tuple[Self, Dict[str, Any]]]: - r""" - Instantiate a Python class from a config dictionary. - - Parameters: - 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`): - Whether kwargs that are not consumed by the Python class should be returned or not. - kwargs (remaining dictionary of keyword arguments, *optional*): - Can be used to update the configuration object (after it is loaded) and initiate the Python class. - `**kwargs` are passed directly to the underlying scheduler/model's `__init__` method and eventually - overwrite the same named arguments in `config`. - - Returns: - [`ModelMixin`] or [`SchedulerMixin`]: - A model or scheduler object instantiated from a config dictionary. - - Examples: - - ```python - >>> from diffusers import DDPMScheduler, DDIMScheduler, PNDMScheduler - - >>> # Download scheduler from huggingface.co and cache. - >>> scheduler = DDPMScheduler.from_pretrained("google/ddpm-cifar10-32") - - >>> # Instantiate DDIM scheduler class with same config as DDPM - >>> scheduler = DDIMScheduler.from_config(scheduler.config) - - >>> # Instantiate PNDM scheduler class with same config as DDPM - >>> scheduler = PNDMScheduler.from_config(scheduler.config) - ``` - """ - # <===== TO BE REMOVED WITH DEPRECATION - # TODO(Patrick) - make sure to remove the following lines when config=="model_path" is deprecated - if "pretrained_model_name_or_path" in kwargs: - config = kwargs.pop("pretrained_model_name_or_path") - - if config is None: - raise ValueError("Please make sure to provide a config as the first positional argument.") - # ======> - - if not isinstance(config, dict): - deprecation_message = "It is deprecated to pass a pretrained model name or path to `from_config`." - if "Scheduler" in cls.__name__: - deprecation_message += ( - f"If you were trying to load a scheduler, please use {cls}.from_pretrained(...) instead." - " Otherwise, please make sure to pass a configuration dictionary instead. This functionality will" - " be removed in v1.0.0." - ) - elif "Model" in cls.__name__: - deprecation_message += ( - f"If you were trying to load a model, please use {cls}.load_config(...) followed by" - f" {cls}.from_config(...) instead. Otherwise, please make sure to pass a configuration dictionary" - " instead. This functionality will be removed in v1.0.0." - ) - deprecate("config-passed-as-path", "1.0.0", deprecation_message, standard_warn=False) - config, kwargs = cls.load_config(pretrained_model_name_or_path=config, return_unused_kwargs=True, **kwargs) - - init_dict, unused_kwargs, hidden_dict = cls.extract_init_dict(config, **kwargs) - - # Allow dtype to be specified on initialization - if "dtype" in unused_kwargs: - init_dict["dtype"] = unused_kwargs.pop("dtype") - - # add possible deprecated kwargs - for deprecated_kwarg in cls._deprecated_kwargs: - if deprecated_kwarg in unused_kwargs: - init_dict[deprecated_kwarg] = unused_kwargs.pop(deprecated_kwarg) - - # Return model and optionally state and/or unused_kwargs - model = cls(**init_dict) - - # make sure to also save config parameters that might be used for compatible classes - # update _class_name - if "_class_name" in hidden_dict: - hidden_dict["_class_name"] = cls.__name__ - - model.register_to_config(**hidden_dict) - - # add hidden kwargs of compatible classes to unused_kwargs - unused_kwargs = {**unused_kwargs, **hidden_dict} - - if return_unused_kwargs: - return (model, unused_kwargs) - else: - return model - - @classmethod - def get_config_dict(cls, *args, **kwargs): - deprecation_message = ( - f" The function get_config_dict is deprecated. Please use {cls}.load_config instead. This function will be" - " removed in version v1.0.0" - ) - deprecate("get_config_dict", "1.0.0", deprecation_message, standard_warn=False) - return cls.load_config(*args, **kwargs) - - @classmethod - @validate_hf_hub_args - def load_config( - cls, - pretrained_model_name_or_path: Union[str, os.PathLike], - return_unused_kwargs=False, - return_commit_hash=False, - **kwargs, - ) -> Tuple[Dict[str, Any], Dict[str, Any]]: - r""" - Load a model or scheduler configuration. - - Parameters: - pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - 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*): - 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*): - 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`): - Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - return_unused_kwargs (`bool`, *optional*, defaults to `False): - Whether unused keyword arguments of the config are returned. - return_commit_hash (`bool`, *optional*, defaults to `False): - Whether the `commit_hash` of the loaded configuration are returned. - - Returns: - `dict`: - A dictionary of all the parameters stored in a JSON configuration file. - - """ - cache_dir = kwargs.pop("cache_dir", None) - local_dir = kwargs.pop("local_dir", None) - local_dir_use_symlinks = kwargs.pop("local_dir_use_symlinks", "auto") - force_download = kwargs.pop("force_download", False) - proxies = kwargs.pop("proxies", None) - token = kwargs.pop("token", None) - local_files_only = kwargs.pop("local_files_only", False) - revision = kwargs.pop("revision", None) - _ = 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) - - user_agent = {**user_agent, "file_type": "config"} - user_agent = http_user_agent(user_agent) - - pretrained_model_name_or_path = str(pretrained_model_name_or_path) - - if cls.config_name is None: - raise ValueError( - "`self.config_name` is not defined. Note that one should not load a config from " - "`ConfigMixin`. Please make sure to define `config_name` in a class inheriting from `ConfigMixin`" - ) - # Custom path for now - if dduf_entries: - if subfolder is not None: - raise ValueError( - "DDUF file only allow for 1 level of directory (e.g transformer/model1/model.safetentors is not allowed). " - "Please check the DDUF structure" - ) - config_file = cls._get_config_file_from_dduf(pretrained_model_name_or_path, dduf_entries) - elif os.path.isfile(pretrained_model_name_or_path): - config_file = pretrained_model_name_or_path - elif os.path.isdir(pretrained_model_name_or_path): - if subfolder is not None and os.path.isfile( - os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name) - ): - config_file = os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name) - elif os.path.isfile(os.path.join(pretrained_model_name_or_path, cls.config_name)): - # Load from a PyTorch checkpoint - config_file = os.path.join(pretrained_model_name_or_path, cls.config_name) - else: - raise EnvironmentError( - f"Error no file named {cls.config_name} found in directory {pretrained_model_name_or_path}." - ) - else: - try: - # Load from URL or cache if already cached - config_file = hf_hub_download( - pretrained_model_name_or_path, - filename=cls.config_name, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - user_agent=user_agent, - subfolder=subfolder, - revision=revision, - local_dir=local_dir, - local_dir_use_symlinks=local_dir_use_symlinks, - ) - except RepositoryNotFoundError: - raise EnvironmentError( - f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier" - " listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a" - " token having permission to this repo with `token` or log in with `huggingface-cli login`." - ) - except RevisionNotFoundError: - raise EnvironmentError( - f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for" - " this model name. Check the model page at" - f" 'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions." - ) - except EntryNotFoundError: - raise EnvironmentError( - f"{pretrained_model_name_or_path} does not appear to have a file named {cls.config_name}." - ) - except HTTPError as err: - raise EnvironmentError( - "There was a specific connection error when trying to load" - f" {pretrained_model_name_or_path}:\n{err}" - ) - except ValueError: - raise EnvironmentError( - f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it" - f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a" - f" directory containing a {cls.config_name} file.\nCheckout your internet connection or see how to" - " run the library in offline mode at" - " 'https://huggingface.co/docs/diffusers/installation#offline-mode'." - ) - except EnvironmentError: - raise EnvironmentError( - f"Can't load config for '{pretrained_model_name_or_path}'. If you were trying to load it from " - "'https://huggingface.co/models', make sure you don't have a local directory with the same name. " - f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory " - f"containing a {cls.config_name} file" - ) - try: - config_dict = cls._dict_from_json_file(config_file, dduf_entries=dduf_entries) - - commit_hash = extract_commit_hash(config_file) - except (json.JSONDecodeError, UnicodeDecodeError): - raise EnvironmentError(f"It looks like the config file at '{config_file}' is not a valid JSON file.") - - if not (return_unused_kwargs or return_commit_hash): - return config_dict - - outputs = (config_dict,) - - if return_unused_kwargs: - outputs += (kwargs,) - - if return_commit_hash: - outputs += (commit_hash,) - - return outputs - - @staticmethod - def _get_init_keys(input_class): - return set(dict(inspect.signature(input_class.__init__).parameters).keys()) - - @classmethod - def extract_init_dict(cls, config_dict, **kwargs): - # Skip keys that were not present in the original config, so default __init__ values were used - used_defaults = config_dict.get("_use_default_values", []) - config_dict = {k: v for k, v in config_dict.items() if k not in used_defaults and k != "_use_default_values"} - - # 0. Copy origin config dict - original_dict = dict(config_dict.items()) - - # 1. Retrieve expected config attributes from __init__ signature - expected_keys = cls._get_init_keys(cls) - expected_keys.remove("self") - # remove general kwargs if present in dict - if "kwargs" in expected_keys: - expected_keys.remove("kwargs") - # remove flax internal keys - if hasattr(cls, "_flax_internal_args"): - for arg in cls._flax_internal_args: - expected_keys.remove(arg) - - # 2. Remove attributes that cannot be expected from expected config attributes - # remove keys to be ignored - if len(cls.ignore_for_config) > 0: - expected_keys = expected_keys - set(cls.ignore_for_config) - - # load diffusers library to import compatible and original scheduler - diffusers_library = importlib.import_module(__name__.split(".")[0]) - - if cls.has_compatibles: - compatible_classes = [c for c in cls._get_compatibles() if not isinstance(c, DummyObject)] - else: - compatible_classes = [] - - expected_keys_comp_cls = set() - for c in compatible_classes: - expected_keys_c = cls._get_init_keys(c) - expected_keys_comp_cls = expected_keys_comp_cls.union(expected_keys_c) - expected_keys_comp_cls = expected_keys_comp_cls - cls._get_init_keys(cls) - config_dict = {k: v for k, v in config_dict.items() if k not in expected_keys_comp_cls} - - # remove attributes from orig class that cannot be expected - orig_cls_name = config_dict.pop("_class_name", cls.__name__) - if ( - isinstance(orig_cls_name, str) - and orig_cls_name != cls.__name__ - and hasattr(diffusers_library, orig_cls_name) - ): - orig_cls = getattr(diffusers_library, orig_cls_name) - unexpected_keys_from_orig = cls._get_init_keys(orig_cls) - expected_keys - config_dict = {k: v for k, v in config_dict.items() if k not in unexpected_keys_from_orig} - elif not isinstance(orig_cls_name, str) and not isinstance(orig_cls_name, (list, tuple)): - raise ValueError( - "Make sure that the `_class_name` is of type string or list of string (for custom pipelines)." - ) - - # remove private attributes - config_dict = {k: v for k, v in config_dict.items() if not k.startswith("_")} - - # remove quantization_config - config_dict = {k: v for k, v in config_dict.items() if k != "quantization_config"} - - # 3. Create keyword arguments that will be passed to __init__ from expected keyword arguments - init_dict = {} - for key in expected_keys: - # if config param is passed to kwarg and is present in config dict - # it should overwrite existing config dict key - if key in kwargs and key in config_dict: - config_dict[key] = kwargs.pop(key) - - if key in kwargs: - # overwrite key - init_dict[key] = kwargs.pop(key) - elif key in config_dict: - # use value from config dict - init_dict[key] = config_dict.pop(key) - - # 4. Give nice warning if unexpected values have been passed - if len(config_dict) > 0: - logger.warning( - f"The config attributes {config_dict} were passed to {cls.__name__}, " - "but are not expected and will be ignored. Please verify your " - f"{cls.config_name} configuration file." - ) - - # 5. Give nice info if config attributes are initialized to default because they have not been passed - passed_keys = set(init_dict.keys()) - if len(expected_keys - passed_keys) > 0: - logger.info( - f"{expected_keys - passed_keys} was not found in config. Values will be initialized to default values." - ) - - # 6. Define unused keyword arguments - unused_kwargs = {**config_dict, **kwargs} - - # 7. Define "hidden" config parameters that were saved for compatible classes - hidden_config_dict = {k: v for k, v in original_dict.items() if k not in init_dict} - - 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 - ): - if dduf_entries: - text = dduf_entries[json_file].read_text() - else: - with open(json_file, "r", encoding="utf-8") as reader: - text = reader.read() - return json.loads(text) - - def __repr__(self): - return f"{self.__class__.__name__} {self.to_json_string()}" - - @property - def config(self) -> Dict[str, Any]: - """ - Returns the config of the class as a frozen dictionary - - Returns: - `Dict[str, Any]`: Config of the class. - """ - return self._internal_dict - - def to_json_string(self) -> str: - """ - Serializes the configuration instance to a JSON string. - - Returns: - `str`: - String containing all the attributes that make up the configuration instance in JSON format. - """ - config_dict = self._internal_dict if hasattr(self, "_internal_dict") else {} - config_dict["_class_name"] = self.__class__.__name__ - config_dict["_diffusers_version"] = __version__ - - def to_json_saveable(value): - if isinstance(value, np.ndarray): - value = value.tolist() - elif isinstance(value, Path): - value = value.as_posix() - return value - - if "quantization_config" in config_dict: - config_dict["quantization_config"] = ( - config_dict.quantization_config.to_dict() - if not isinstance(config_dict.quantization_config, dict) - else config_dict.quantization_config - ) - - config_dict = {k: to_json_saveable(v) for k, v in config_dict.items()} - # Don't save "_ignore_files" or "_use_default_values" - config_dict.pop("_ignore_files", None) - config_dict.pop("_use_default_values", None) - # pop the `_pre_quantization_dtype` as torch.dtypes are not serializable. - _ = config_dict.pop("_pre_quantization_dtype", None) - - return json.dumps(config_dict, indent=2, sort_keys=True) + "\n" - - def to_json_file(self, json_file_path: Union[str, os.PathLike]): - """ - Save the configuration instance's parameters to a JSON file. - - Args: - json_file_path (`str` or `os.PathLike`): - Path to the JSON file to save a configuration instance's parameters. - """ - with open(json_file_path, "w", encoding="utf-8") as writer: - 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]): - # paths inside a DDUF file must always be "/" - config_file = ( - cls.config_name - if pretrained_model_name_or_path == "" - else "/".join([pretrained_model_name_or_path, cls.config_name]) - ) - if config_file not in dduf_entries: - raise ValueError( - f"We did not manage to find the file {config_file} in the dduf file. We only have the following files {dduf_entries.keys()}" - ) - return config_file - - -def register_to_config(init): - r""" - Decorator to apply on the init of classes inheriting from [`ConfigMixin`] so that all the arguments are - automatically sent to `self.register_for_config`. To ignore a specific argument accepted by the init but that - shouldn't be registered in the config, use the `ignore_for_config` class variable - - Warning: Once decorated, all private arguments (beginning with an underscore) are trashed and not sent to the init! - """ - - @functools.wraps(init) - def inner_init(self, *args, **kwargs): - # Ignore private kwargs in the init. - init_kwargs = {k: v for k, v in kwargs.items() if not k.startswith("_")} - config_init_kwargs = {k: v for k, v in kwargs.items() if k.startswith("_")} - if not isinstance(self, ConfigMixin): - raise RuntimeError( - f"`@register_for_config` was applied to {self.__class__.__name__} init method, but this class does " - "not inherit from `ConfigMixin`." - ) - - ignore = getattr(self, "ignore_for_config", []) - # Get positional arguments aligned with kwargs - new_kwargs = {} - signature = inspect.signature(init) - parameters = { - name: p.default for i, (name, p) in enumerate(signature.parameters.items()) if i > 0 and name not in ignore - } - for arg, name in zip(args, parameters.keys()): - new_kwargs[name] = arg - - # Then add all kwargs - new_kwargs.update( - { - k: init_kwargs.get(k, default) - for k, default in parameters.items() - if k not in ignore and k not in new_kwargs - } - ) - - # Take note of the parameters that were not present in the loaded config - if len(set(new_kwargs.keys()) - set(init_kwargs)) > 0: - new_kwargs["_use_default_values"] = list(set(new_kwargs.keys()) - set(init_kwargs)) - - new_kwargs = {**config_init_kwargs, **new_kwargs} - getattr(self, "register_to_config")(**new_kwargs) - init(self, *args, **init_kwargs) - - return inner_init - - -def flax_register_to_config(cls): - original_init = cls.__init__ - - @functools.wraps(original_init) - def init(self, *args, **kwargs): - if not isinstance(self, ConfigMixin): - raise RuntimeError( - f"`@register_for_config` was applied to {self.__class__.__name__} init method, but this class does " - "not inherit from `ConfigMixin`." - ) - - # Ignore private kwargs in the init. Retrieve all passed attributes - init_kwargs = dict(kwargs.items()) - - # Retrieve default values - fields = dataclasses.fields(self) - default_kwargs = {} - for field in fields: - # ignore flax specific attributes - if field.name in self._flax_internal_args: - continue - if type(field.default) == dataclasses._MISSING_TYPE: - default_kwargs[field.name] = None - else: - default_kwargs[field.name] = getattr(self, field.name) - - # Make sure init_kwargs override default kwargs - new_kwargs = {**default_kwargs, **init_kwargs} - # dtype should be part of `init_kwargs`, but not `new_kwargs` - if "dtype" in new_kwargs: - new_kwargs.pop("dtype") - - # Get positional arguments aligned with kwargs - for i, arg in enumerate(args): - name = fields[i].name - new_kwargs[name] = arg - - # Take note of the parameters that were not present in the loaded config - if len(set(new_kwargs.keys()) - set(init_kwargs)) > 0: - new_kwargs["_use_default_values"] = list(set(new_kwargs.keys()) - set(init_kwargs)) - - getattr(self, "register_to_config")(**new_kwargs) - original_init(self, *args, **kwargs) - - cls.__init__ = init - return cls - - -class LegacyConfigMixin(ConfigMixin): - r""" - A subclass of `ConfigMixin` to resolve class mapping from legacy classes (like `Transformer2DModel`) to more - pipeline-specific classes (like `DiTTransformer2DModel`). - """ - - @classmethod - def from_config(cls, config: Union[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 - - # resolve remapping - remapped_class = _fetch_remapped_cls_from_config(config, cls) - - return remapped_class.from_config(config, return_unused_kwargs, **kwargs) diff --git a/venv/lib/python3.10/site-packages/diffusers/dependency_versions_check.py b/venv/lib/python3.10/site-packages/diffusers/dependency_versions_check.py deleted file mode 100644 index e3670b136af45742a64a96b0d68849a2ca0700ed..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/dependency_versions_check.py +++ /dev/null @@ -1,34 +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 .dependency_versions_table import deps -from .utils.versions import require_version, require_version_core - - -# define which module versions we always want to check at run time -# (usually the ones defined in `install_requires` in setup.py) -# -# order specific notes: -# - tqdm must be checked before tokenizers - -pkgs_to_check_at_runtime = "python requests filelock numpy".split() -for pkg in pkgs_to_check_at_runtime: - if pkg in deps: - require_version_core(deps[pkg]) - else: - raise ValueError(f"can't find {pkg} in {deps.keys()}, check dependency_versions_table.py") - - -def dep_version_check(pkg, hint=None): - require_version(deps[pkg], hint) diff --git a/venv/lib/python3.10/site-packages/diffusers/dependency_versions_table.py b/venv/lib/python3.10/site-packages/diffusers/dependency_versions_table.py deleted file mode 100644 index f35353c49eab7f83e3a2670cb3ff0c1923872cfc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/dependency_versions_table.py +++ /dev/null @@ -1,53 +0,0 @@ -# THIS FILE HAS BEEN AUTOGENERATED. To update: -# 1. modify the `_deps` dict in setup.py -# 2. run `make deps_table_update` -deps = { - "Pillow": "Pillow", - "accelerate": "accelerate>=0.31.0", - "compel": "compel==0.1.8", - "datasets": "datasets", - "filelock": "filelock", - "flax": "flax>=0.4.1", - "hf-doc-builder": "hf-doc-builder>=0.3.0", - "huggingface-hub": "huggingface-hub>=0.27.0", - "requests-mock": "requests-mock==1.10.0", - "importlib_metadata": "importlib_metadata", - "invisible-watermark": "invisible-watermark>=0.2.0", - "isort": "isort>=5.5.4", - "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", - "numpy": "numpy", - "parameterized": "parameterized", - "peft": "peft>=0.15.0", - "protobuf": "protobuf>=3.20.3,<4", - "pytest": "pytest", - "pytest-timeout": "pytest-timeout", - "pytest-xdist": "pytest-xdist", - "python": "python>=3.8.0", - "ruff": "ruff==0.9.10", - "safetensors": "safetensors>=0.3.1", - "sentencepiece": "sentencepiece>=0.1.91,!=0.1.92", - "GitPython": "GitPython<3.1.19", - "scipy": "scipy", - "onnx": "onnx", - "optimum_quanto": "optimum_quanto>=0.2.6", - "gguf": "gguf>=0.10.0", - "torchao": "torchao>=0.7.0", - "bitsandbytes": "bitsandbytes>=0.43.3", - "regex": "regex!=2019.12.17", - "requests": "requests", - "tensorboard": "tensorboard", - "tiktoken": "tiktoken>=0.7.0", - "torch": "torch>=1.4", - "torchvision": "torchvision", - "transformers": "transformers>=4.41.2", - "urllib3": "urllib3<=2.0.0", - "black": "black", - "phonemizer": "phonemizer", - "opencv-python": "opencv-python", -} diff --git a/venv/lib/python3.10/site-packages/diffusers/experimental/__init__.py b/venv/lib/python3.10/site-packages/diffusers/experimental/__init__.py deleted file mode 100644 index ebc8155403016dfd8ad7fb78d246f9da9098ac50..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/experimental/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .rl import ValueGuidedRLPipeline diff --git a/venv/lib/python3.10/site-packages/diffusers/experimental/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/experimental/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 1caf15fce0d57419cdac1cf231ccb01aae30daac..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/experimental/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/experimental/rl/__init__.py b/venv/lib/python3.10/site-packages/diffusers/experimental/rl/__init__.py deleted file mode 100644 index 7b338d3173e12d478b6b6d6fd0e50650a0ab5a4c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/experimental/rl/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .value_guided_sampling import ValueGuidedRLPipeline diff --git a/venv/lib/python3.10/site-packages/diffusers/experimental/rl/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/experimental/rl/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 39e964e534470c1d80d5503e2944df9c349f85ad..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/experimental/rl/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/experimental/rl/__pycache__/value_guided_sampling.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/experimental/rl/__pycache__/value_guided_sampling.cpython-310.pyc deleted file mode 100644 index 1fa2f332be2625728f40af8fdbe3f6a2b7651204..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/experimental/rl/__pycache__/value_guided_sampling.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/experimental/rl/value_guided_sampling.py b/venv/lib/python3.10/site-packages/diffusers/experimental/rl/value_guided_sampling.py deleted file mode 100644 index c69d308ecc6894edd39557d7f16fc01914e4bd85..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/experimental/rl/value_guided_sampling.py +++ /dev/null @@ -1,153 +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 numpy as np -import torch -import tqdm - -from ...models.unets.unet_1d import UNet1DModel -from ...pipelines import DiffusionPipeline -from ...utils.dummy_pt_objects import DDPMScheduler -from ...utils.torch_utils import randn_tensor - - -class ValueGuidedRLPipeline(DiffusionPipeline): - r""" - Pipeline for value-guided sampling from a diffusion model trained to predict sequences of states. - - 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.). - - Parameters: - value_function ([`UNet1DModel`]): - A specialized UNet for fine-tuning trajectories base on reward. - unet ([`UNet1DModel`]): - UNet architecture to denoise the encoded trajectories. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded trajectories. Default for this - application is [`DDPMScheduler`]. - env (): - An environment following the OpenAI gym API to act in. For now only Hopper has pretrained models. - """ - - def __init__( - self, - value_function: UNet1DModel, - unet: UNet1DModel, - scheduler: DDPMScheduler, - env, - ): - super().__init__() - - self.register_modules(value_function=value_function, unet=unet, scheduler=scheduler, env=env) - - self.data = env.get_dataset() - self.means = {} - for key in self.data.keys(): - try: - self.means[key] = self.data[key].mean() - except: # noqa: E722 - pass - self.stds = {} - for key in self.data.keys(): - try: - self.stds[key] = self.data[key].std() - except: # noqa: E722 - pass - self.state_dim = env.observation_space.shape[0] - self.action_dim = env.action_space.shape[0] - - def normalize(self, x_in, key): - return (x_in - self.means[key]) / self.stds[key] - - def de_normalize(self, x_in, key): - return x_in * self.stds[key] + self.means[key] - - def to_torch(self, x_in): - if isinstance(x_in, dict): - return {k: self.to_torch(v) for k, v in x_in.items()} - elif torch.is_tensor(x_in): - return x_in.to(self.unet.device) - return torch.tensor(x_in, device=self.unet.device) - - def reset_x0(self, x_in, cond, act_dim): - for key, val in cond.items(): - x_in[:, key, act_dim:] = val.clone() - return x_in - - def run_diffusion(self, x, conditions, n_guide_steps, scale): - batch_size = x.shape[0] - y = None - for i in tqdm.tqdm(self.scheduler.timesteps): - # create batch of timesteps to pass into model - timesteps = torch.full((batch_size,), i, device=self.unet.device, dtype=torch.long) - for _ in range(n_guide_steps): - with torch.enable_grad(): - x.requires_grad_() - - # permute to match dimension for pre-trained models - y = self.value_function(x.permute(0, 2, 1), timesteps).sample - grad = torch.autograd.grad([y.sum()], [x])[0] - - posterior_variance = self.scheduler._get_variance(i) - model_std = torch.exp(0.5 * posterior_variance) - grad = model_std * grad - - grad[timesteps < 2] = 0 - x = x.detach() - x = x + scale * grad - x = self.reset_x0(x, conditions, self.action_dim) - - prev_x = self.unet(x.permute(0, 2, 1), timesteps).sample.permute(0, 2, 1) - - # TODO: verify deprecation of this kwarg - x = self.scheduler.step(prev_x, i, x)["prev_sample"] - - # apply conditions to the trajectory (set the initial state) - x = self.reset_x0(x, conditions, self.action_dim) - x = self.to_torch(x) - return x, y - - def __call__(self, obs, batch_size=64, planning_horizon=32, n_guide_steps=2, scale=0.1): - # normalize the observations and create batch dimension - obs = self.normalize(obs, "observations") - obs = obs[None].repeat(batch_size, axis=0) - - conditions = {0: self.to_torch(obs)} - shape = (batch_size, planning_horizon, self.state_dim + self.action_dim) - - # generate initial noise and apply our conditions (to make the trajectories start at current state) - x1 = randn_tensor(shape, device=self.unet.device) - x = self.reset_x0(x1, conditions, self.action_dim) - x = self.to_torch(x) - - # run the diffusion process - x, y = self.run_diffusion(x, conditions, n_guide_steps, scale) - - # sort output trajectories by value - sorted_idx = y.argsort(0, descending=True).squeeze() - sorted_values = x[sorted_idx] - actions = sorted_values[:, :, : self.action_dim] - actions = actions.detach().cpu().numpy() - denorm_actions = self.de_normalize(actions, key="actions") - - # select the action with the highest value - if y is not None: - selected_index = 0 - else: - # if we didn't run value guiding, select a random action - selected_index = np.random.randint(0, batch_size) - - denorm_actions = denorm_actions[selected_index, 0] - return denorm_actions diff --git a/venv/lib/python3.10/site-packages/diffusers/hooks/__init__.py b/venv/lib/python3.10/site-packages/diffusers/hooks/__init__.py deleted file mode 100644 index 764ceb25b4650d0808aeaacc62e769edc2f6dd77..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/hooks/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -from ..utils import is_torch_available - - -if is_torch_available(): - from .faster_cache import FasterCacheConfig, apply_faster_cache - from .group_offloading import apply_group_offloading - from .hooks import HookRegistry, ModelHook - from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook - from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast diff --git a/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index da1a1a805875109b41b53808bf0b593eb823a949..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/faster_cache.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/faster_cache.cpython-310.pyc deleted file mode 100644 index 2bdf2d6aef814c9884e06453a7d92e9931613220..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/faster_cache.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/group_offloading.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/group_offloading.cpython-310.pyc deleted file mode 100644 index 0fd65d3b1d34ba9231c2ba1d26deccbcd3b81268..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/group_offloading.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/hooks.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/hooks.cpython-310.pyc deleted file mode 100644 index 2633eb22d46d9a7c8d9c090e521742c9f6e6d665..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/hooks.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/layerwise_casting.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/layerwise_casting.cpython-310.pyc deleted file mode 100644 index 2389e2f6e9d7fd60d21862b9457adcf3c18d8fac..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/layerwise_casting.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/pyramid_attention_broadcast.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/pyramid_attention_broadcast.cpython-310.pyc deleted file mode 100644 index 413552a2759d77c242556fc3e2b6325524b50dc3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/hooks/__pycache__/pyramid_attention_broadcast.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/hooks/faster_cache.py b/venv/lib/python3.10/site-packages/diffusers/hooks/faster_cache.py deleted file mode 100644 index 1be5e143629400cd6676e6881c3e594ae30ff930..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/hooks/faster_cache.py +++ /dev/null @@ -1,653 +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 re -from dataclasses import dataclass -from typing import Any, Callable, List, Optional, Tuple - -import torch - -from ..models.attention_processor import Attention, MochiAttention -from ..models.modeling_outputs import Transformer2DModelOutput -from ..utils import logging -from .hooks import HookRegistry, ModelHook - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -_FASTER_CACHE_DENOISER_HOOK = "faster_cache_denoiser" -_FASTER_CACHE_BLOCK_HOOK = "faster_cache_block" -_ATTENTION_CLASSES = (Attention, MochiAttention) -_SPATIAL_ATTENTION_BLOCK_IDENTIFIERS = ( - "^blocks.*attn", - "^transformer_blocks.*attn", - "^single_transformer_blocks.*attn", -) -_TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS = ("^temporal_transformer_blocks.*attn",) -_TRANSFORMER_BLOCK_IDENTIFIERS = _SPATIAL_ATTENTION_BLOCK_IDENTIFIERS + _TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS -_UNCOND_COND_INPUT_KWARGS_IDENTIFIERS = ( - "hidden_states", - "encoder_hidden_states", - "timestep", - "attention_mask", - "encoder_attention_mask", -) - - -@dataclass -class FasterCacheConfig: - r""" - Configuration for [FasterCache](https://huggingface.co/papers/2410.19355). - - Attributes: - spatial_attention_block_skip_range (`int`, defaults to `2`): - 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 re-used) before computing the new attention - states again. - temporal_attention_block_skip_range (`int`, *optional*, defaults to `None`): - 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 re-used) before computing the new attention - states again. - 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 - denoising are in the reversed range of 0 to 1000 (i.e. denoising starts at timestep 1000 and ends at - 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`): - 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)`): - 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)`): - 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. - alpha_low_frequency (`float`, defaults to `1.1`): - The weight to scale the low frequency updates by. This is used to approximate the unconditional branch from - the conditional branch outputs. - alpha_high_frequency (`float`, defaults to `1.1`): - The weight to scale the high frequency updates by. This is used to approximate the unconditional branch - from the conditional branch outputs. - unconditional_batch_skip_range (`int`, defaults to `5`): - 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 re-used) before - computing the new unconditional branch states again. - 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")`): - 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",)`): - 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. - attention_weight_callback (`Callable[[torch.nn.Module], float]`, defaults to `None`): - The callback function to determine the weight to scale the attention outputs by. This function should take - the attention module as input and return a float value. This is used to approximate the unconditional - branch from the conditional branch outputs. If not provided, the default weight is 0.5 for all timesteps. - Typically, as described in the paper, this weight should gradually increase from 0 to 1 as the inference - progresses. Users are encouraged to experiment and provide custom weight schedules that take into account - the number of inference steps and underlying model behaviour as denoising progresses. - low_frequency_weight_callback (`Callable[[torch.nn.Module], float]`, defaults to `None`): - The callback function to determine the weight to scale the low frequency updates by. If not provided, the - default weight is 1.1 for timesteps within the range specified (as described in the paper). - high_frequency_weight_callback (`Callable[[torch.nn.Module], float]`, defaults to `None`): - The callback function to determine the weight to scale the high frequency updates by. If not provided, the - default weight is 1.1 for timesteps within the range specified (as described in the paper). - tensor_format (`str`, defaults to `"BCFHW"`): - The format of the input tensors. This should be one of `"BCFHW"`, `"BFCHW"`, or `"BCHW"`. The format is - used to split individual latent frames in order for low and high frequency components to be computed. - 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")`): - 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 - names that contain the batchwise-concatenated unconditional and conditional inputs. - """ - - # 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 - - 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) - - # ⍺1 and ⍺2 as mentioned in Equation 11 of the paper - alpha_low_frequency: float = 1.1 - alpha_high_frequency: float = 1.1 - - # 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) - - 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 - high_frequency_weight_callback: Callable[[torch.nn.Module], float] = None - - tensor_format: str = "BCFHW" - is_guidance_distilled: bool = False - - current_timestep_callback: Callable[[], int] = None - - _unconditional_conditional_input_kwargs_identifiers: List[str] = _UNCOND_COND_INPUT_KWARGS_IDENTIFIERS - - def __repr__(self) -> str: - return ( - f"FasterCacheConfig(\n" - f" spatial_attention_block_skip_range={self.spatial_attention_block_skip_range},\n" - f" temporal_attention_block_skip_range={self.temporal_attention_block_skip_range},\n" - f" spatial_attention_timestep_skip_range={self.spatial_attention_timestep_skip_range},\n" - f" temporal_attention_timestep_skip_range={self.temporal_attention_timestep_skip_range},\n" - f" low_frequency_weight_update_timestep_range={self.low_frequency_weight_update_timestep_range},\n" - f" high_frequency_weight_update_timestep_range={self.high_frequency_weight_update_timestep_range},\n" - f" alpha_low_frequency={self.alpha_low_frequency},\n" - f" alpha_high_frequency={self.alpha_high_frequency},\n" - f" unconditional_batch_skip_range={self.unconditional_batch_skip_range},\n" - f" unconditional_batch_timestep_skip_range={self.unconditional_batch_timestep_skip_range},\n" - f" spatial_attention_block_identifiers={self.spatial_attention_block_identifiers},\n" - f" temporal_attention_block_identifiers={self.temporal_attention_block_identifiers},\n" - f" tensor_format={self.tensor_format},\n" - f")" - ) - - -class FasterCacheDenoiserState: - r""" - State for [FasterCache](https://huggingface.co/papers/2410.19355) top-level denoiser module. - """ - - def __init__(self) -> None: - self.iteration: int = 0 - self.low_frequency_delta: torch.Tensor = None - self.high_frequency_delta: torch.Tensor = None - - def reset(self): - self.iteration = 0 - self.low_frequency_delta = None - self.high_frequency_delta = None - - -class FasterCacheBlockState: - r""" - State for [FasterCache](https://huggingface.co/papers/2410.19355). Every underlying block that FasterCache is - applied to will have an instance of this state. - """ - - def __init__(self) -> None: - self.iteration: int = 0 - self.batch_size: int = None - self.cache: Tuple[torch.Tensor, torch.Tensor] = None - - def reset(self): - self.iteration = 0 - self.batch_size = None - self.cache = None - - -class FasterCacheDenoiserHook(ModelHook): - _is_stateful = True - - def __init__( - self, - unconditional_batch_skip_range: int, - unconditional_batch_timestep_skip_range: Tuple[int, int], - tensor_format: str, - is_guidance_distilled: bool, - 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], - ) -> None: - super().__init__() - - self.unconditional_batch_skip_range = unconditional_batch_skip_range - self.unconditional_batch_timestep_skip_range = unconditional_batch_timestep_skip_range - # We can't easily detect what args are to be split in unconditional and conditional branches. We - # can only do it for kwargs, hence they are the only ones we split. The args are passed as-is. - # If a model is to be made compatible with FasterCache, the user must ensure that the inputs that - # contain batchwise-concatenated unconditional and conditional inputs are passed as kwargs. - self.uncond_cond_input_kwargs_identifiers = uncond_cond_input_kwargs_identifiers - self.tensor_format = tensor_format - self.is_guidance_distilled = is_guidance_distilled - - self.current_timestep_callback = current_timestep_callback - self.low_frequency_weight_callback = low_frequency_weight_callback - self.high_frequency_weight_callback = high_frequency_weight_callback - - def initialize_hook(self, module): - self.state = FasterCacheDenoiserState() - return module - - @staticmethod - 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) - return cond - - def new_forward(self, module: torch.nn.Module, *args, **kwargs) -> Any: - # Split the unconditional and conditional inputs. We only want to infer the conditional branch if the - # requirements for skipping the unconditional branch are met as described in the paper. - # We skip the unconditional branch only if the following conditions are met: - # 1. We have completed at least one iteration of the denoiser - # 2. The current timestep is within the range specified by the user. This is the optimal timestep range - # where approximating the unconditional branch from the computation of the conditional branch is possible - # without a significant loss in quality. - # 3. The current iteration is not a multiple of the unconditional batch skip range. This is done so that - # we compute the unconditional branch at least once every few iterations to ensure minimal quality loss. - is_within_timestep_range = ( - self.unconditional_batch_timestep_skip_range[0] - < self.current_timestep_callback() - < self.unconditional_batch_timestep_skip_range[1] - ) - should_skip_uncond = ( - self.state.iteration > 0 - and is_within_timestep_range - and self.state.iteration % self.unconditional_batch_skip_range != 0 - and not self.is_guidance_distilled - ) - - if should_skip_uncond: - is_any_kwarg_uncond = any(k in self.uncond_cond_input_kwargs_identifiers for k in kwargs.keys()) - if is_any_kwarg_uncond: - logger.debug("FasterCache - Skipping unconditional branch computation") - args = tuple([self._get_cond_input(arg) if torch.is_tensor(arg) else arg for arg in args]) - kwargs = { - k: v if k not in self.uncond_cond_input_kwargs_identifiers else self._get_cond_input(v) - for k, v in kwargs.items() - } - - output = self.fn_ref.original_forward(*args, **kwargs) - - if self.is_guidance_distilled: - self.state.iteration += 1 - return output - - if torch.is_tensor(output): - hidden_states = output - elif isinstance(output, (tuple, Transformer2DModelOutput)): - hidden_states = output[0] - - batch_size = hidden_states.size(0) - - if should_skip_uncond: - self.state.low_frequency_delta = self.state.low_frequency_delta * self.low_frequency_weight_callback( - module - ) - self.state.high_frequency_delta = self.state.high_frequency_delta * self.high_frequency_weight_callback( - module - ) - - if self.tensor_format == "BCFHW": - hidden_states = hidden_states.permute(0, 2, 1, 3, 4) - if self.tensor_format == "BCFHW" or self.tensor_format == "BFCHW": - hidden_states = hidden_states.flatten(0, 1) - - low_freq_cond, high_freq_cond = _split_low_high_freq(hidden_states.float()) - - # Approximate/compute the unconditional branch outputs as described in Equation 9 and 10 of the paper - low_freq_uncond = self.state.low_frequency_delta + low_freq_cond - high_freq_uncond = self.state.high_frequency_delta + high_freq_cond - uncond_freq = low_freq_uncond + high_freq_uncond - - uncond_states = torch.fft.ifftshift(uncond_freq) - uncond_states = torch.fft.ifft2(uncond_states).real - - if self.tensor_format == "BCFHW" or self.tensor_format == "BFCHW": - uncond_states = uncond_states.unflatten(0, (batch_size, -1)) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)) - if self.tensor_format == "BCFHW": - uncond_states = uncond_states.permute(0, 2, 1, 3, 4) - hidden_states = hidden_states.permute(0, 2, 1, 3, 4) - - # Concatenate the approximated unconditional and predicted conditional branches - uncond_states = uncond_states.to(hidden_states.dtype) - hidden_states = torch.cat([uncond_states, hidden_states], dim=0) - else: - uncond_states, cond_states = hidden_states.chunk(2, dim=0) - if self.tensor_format == "BCFHW": - uncond_states = uncond_states.permute(0, 2, 1, 3, 4) - cond_states = cond_states.permute(0, 2, 1, 3, 4) - if self.tensor_format == "BCFHW" or self.tensor_format == "BFCHW": - uncond_states = uncond_states.flatten(0, 1) - cond_states = cond_states.flatten(0, 1) - - low_freq_uncond, high_freq_uncond = _split_low_high_freq(uncond_states.float()) - low_freq_cond, high_freq_cond = _split_low_high_freq(cond_states.float()) - self.state.low_frequency_delta = low_freq_uncond - low_freq_cond - self.state.high_frequency_delta = high_freq_uncond - high_freq_cond - - self.state.iteration += 1 - if torch.is_tensor(output): - output = hidden_states - elif isinstance(output, tuple): - output = (hidden_states, *output[1:]) - else: - output.sample = hidden_states - - return output - - def reset_state(self, module: torch.nn.Module) -> torch.nn.Module: - self.state.reset() - return module - - -class FasterCacheBlockHook(ModelHook): - _is_stateful = True - - def __init__( - self, - block_skip_range: int, - timestep_skip_range: Tuple[int, int], - is_guidance_distilled: bool, - weight_callback: Callable[[torch.nn.Module], float], - current_timestep_callback: Callable[[], int], - ) -> None: - super().__init__() - - self.block_skip_range = block_skip_range - self.timestep_skip_range = timestep_skip_range - self.is_guidance_distilled = is_guidance_distilled - - self.weight_callback = weight_callback - self.current_timestep_callback = current_timestep_callback - - def initialize_hook(self, module): - self.state = FasterCacheBlockState() - return module - - def _compute_approximated_attention_output( - self, t_2_output: torch.Tensor, t_output: torch.Tensor, weight: float, batch_size: int - ) -> torch.Tensor: - if t_2_output.size(0) != batch_size: - # The cache t_2_output contains both batchwise-concatenated unconditional-conditional branch outputs. Just - # take the conditional branch outputs. - assert t_2_output.size(0) == 2 * batch_size - t_2_output = t_2_output[batch_size:] - if t_output.size(0) != batch_size: - # The cache t_output contains both batchwise-concatenated unconditional-conditional branch outputs. Just - # take the conditional branch outputs. - assert t_output.size(0) == 2 * batch_size - t_output = t_output[batch_size:] - return t_output + (t_output - t_2_output) * weight - - def new_forward(self, module: torch.nn.Module, *args, **kwargs) -> Any: - batch_size = [ - *[arg.size(0) for arg in args if torch.is_tensor(arg)], - *[v.size(0) for v in kwargs.values() if torch.is_tensor(v)], - ][0] - if self.state.batch_size is None: - # Will be updated on first forward pass through the denoiser - self.state.batch_size = batch_size - - # If we have to skip due to the skip conditions, then let's skip as expected. - # But, we can't skip if the denoiser wants to infer both unconditional and conditional branches. This - # is because the expected output shapes of attention layer will not match if we only return values from - # the cache (which only caches conditional branch outputs). So, if state.batch_size (which is the true - # unconditional-conditional batch size) is same as the current batch size, we don't perform the layer - # skip. Otherwise, we conditionally skip the layer based on what state.skip_callback returns. - is_within_timestep_range = ( - self.timestep_skip_range[0] < self.current_timestep_callback() < self.timestep_skip_range[1] - ) - if not is_within_timestep_range: - should_skip_attention = False - else: - should_compute_attention = self.state.iteration > 0 and self.state.iteration % self.block_skip_range == 0 - should_skip_attention = not should_compute_attention - if should_skip_attention: - should_skip_attention = self.is_guidance_distilled or self.state.batch_size != batch_size - - if should_skip_attention: - logger.debug("FasterCache - Skipping attention and using approximation") - if torch.is_tensor(self.state.cache[-1]): - t_2_output, t_output = self.state.cache - weight = self.weight_callback(module) - output = self._compute_approximated_attention_output(t_2_output, t_output, weight, batch_size) - else: - # The cache contains multiple tensors from past N iterations (N=2 for FasterCache). We need to handle all of them. - # Diffusers blocks can return multiple tensors - let's call them [A, B, C, ...] for simplicity. - # In our cache, we would have [[A_1, B_1, C_1, ...], [A_2, B_2, C_2, ...], ...] where each list is the output from - # a forward pass of the block. We need to compute the approximated output for each of these tensors. - # The zip(*state.cache) operation will give us [(A_1, A_2, ...), (B_1, B_2, ...), (C_1, C_2, ...), ...] which - # allows us to compute the approximated attention output for each tensor in the cache. - output = () - for t_2_output, t_output in zip(*self.state.cache): - result = self._compute_approximated_attention_output( - t_2_output, t_output, self.weight_callback(module), batch_size - ) - output += (result,) - else: - logger.debug("FasterCache - Computing attention") - output = self.fn_ref.original_forward(*args, **kwargs) - - # Note that the following condition for getting hidden_states should suffice since Diffusers blocks either return - # a single hidden_states tensor, or a tuple of (hidden_states, encoder_hidden_states) tensors. We need to handle - # both cases. - if torch.is_tensor(output): - cache_output = output - if not self.is_guidance_distilled and cache_output.size(0) == self.state.batch_size: - # The output here can be both unconditional-conditional branch outputs or just conditional branch outputs. - # This is determined at the higher-level denoiser module. We only want to cache the conditional branch outputs. - cache_output = cache_output.chunk(2, dim=0)[1] - else: - # Cache all return values and perform the same operation as above - cache_output = () - for out in output: - if not self.is_guidance_distilled and out.size(0) == self.state.batch_size: - out = out.chunk(2, dim=0)[1] - cache_output += (out,) - - if self.state.cache is None: - self.state.cache = [cache_output, cache_output] - else: - self.state.cache = [self.state.cache[-1], cache_output] - - self.state.iteration += 1 - return output - - def reset_state(self, module: torch.nn.Module) -> torch.nn.Module: - self.state.reset() - return module - - -def apply_faster_cache(module: torch.nn.Module, config: FasterCacheConfig) -> None: - r""" - Applies [FasterCache](https://huggingface.co/papers/2410.19355) to a given pipeline. - - Args: - pipeline (`DiffusionPipeline`): - The diffusion pipeline to apply FasterCache to. - config (`Optional[FasterCacheConfig]`, `optional`, defaults to `None`): - The configuration to use for FasterCache. - - Example: - ```python - >>> import torch - >>> from diffusers import CogVideoXPipeline, FasterCacheConfig, apply_faster_cache - - >>> pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16) - >>> pipe.to("cuda") - - >>> config = FasterCacheConfig( - ... spatial_attention_block_skip_range=2, - ... spatial_attention_timestep_skip_range=(-1, 681), - ... low_frequency_weight_update_timestep_range=(99, 641), - ... high_frequency_weight_update_timestep_range=(-1, 301), - ... spatial_attention_block_identifiers=["transformer_blocks"], - ... attention_weight_callback=lambda _: 0.3, - ... tensor_format="BFCHW", - ... ) - >>> apply_faster_cache(pipe.transformer, config) - ``` - """ - - logger.warning( - "FasterCache is a purely experimental feature and may not work as expected. Not all models support FasterCache. " - "The API is subject to change in future releases, with no guarantee of backward compatibility. Please report any issues at " - "https://github.com/huggingface/diffusers/issues." - ) - - if config.attention_weight_callback is None: - # If the user has not provided a weight callback, we default to 0.5 for all timesteps. - # In the paper, they recommend using a gradually increasing weight from 0 to 1 as the inference progresses, but - # this depends from model-to-model. It is required by the user to provide a weight callback if they want to - # use a different weight function. Defaulting to 0.5 works well in practice for most cases. - logger.warning( - "No `attention_weight_callback` provided when enabling FasterCache. Defaulting to using a weight of 0.5 for all timesteps." - ) - config.attention_weight_callback = lambda _: 0.5 - - if config.low_frequency_weight_callback is None: - logger.debug( - "Low frequency weight callback not provided when enabling FasterCache. Defaulting to behaviour described in the paper." - ) - - def low_frequency_weight_callback(module: torch.nn.Module) -> float: - is_within_range = ( - config.low_frequency_weight_update_timestep_range[0] - < config.current_timestep_callback() - < config.low_frequency_weight_update_timestep_range[1] - ) - return config.alpha_low_frequency if is_within_range else 1.0 - - config.low_frequency_weight_callback = low_frequency_weight_callback - - if config.high_frequency_weight_callback is None: - logger.debug( - "High frequency weight callback not provided when enabling FasterCache. Defaulting to behaviour described in the paper." - ) - - def high_frequency_weight_callback(module: torch.nn.Module) -> float: - is_within_range = ( - config.high_frequency_weight_update_timestep_range[0] - < config.current_timestep_callback() - < config.high_frequency_weight_update_timestep_range[1] - ) - return config.alpha_high_frequency if is_within_range else 1.0 - - config.high_frequency_weight_callback = high_frequency_weight_callback - - supported_tensor_formats = ["BCFHW", "BFCHW", "BCHW"] # TODO(aryan): Support BSC for LTX Video - if config.tensor_format not in supported_tensor_formats: - raise ValueError(f"`tensor_format` must be one of {supported_tensor_formats}, but got {config.tensor_format}.") - - _apply_faster_cache_on_denoiser(module, config) - - for name, submodule in module.named_modules(): - if not isinstance(submodule, _ATTENTION_CLASSES): - continue - if any(re.search(identifier, name) is not None for identifier in _TRANSFORMER_BLOCK_IDENTIFIERS): - _apply_faster_cache_on_attention_class(name, submodule, config) - - -def _apply_faster_cache_on_denoiser(module: torch.nn.Module, config: FasterCacheConfig) -> None: - hook = FasterCacheDenoiserHook( - config.unconditional_batch_skip_range, - config.unconditional_batch_timestep_skip_range, - config.tensor_format, - config.is_guidance_distilled, - config._unconditional_conditional_input_kwargs_identifiers, - config.current_timestep_callback, - config.low_frequency_weight_callback, - config.high_frequency_weight_callback, - ) - registry = HookRegistry.check_if_exists_or_initialize(module) - registry.register_hook(hook, _FASTER_CACHE_DENOISER_HOOK) - - -def _apply_faster_cache_on_attention_class(name: str, module: Attention, config: FasterCacheConfig) -> None: - is_spatial_self_attention = ( - any(re.search(identifier, name) is not None for identifier in config.spatial_attention_block_identifiers) - and config.spatial_attention_block_skip_range is not None - and not getattr(module, "is_cross_attention", False) - ) - is_temporal_self_attention = ( - any(re.search(identifier, name) is not None for identifier in config.temporal_attention_block_identifiers) - and config.temporal_attention_block_skip_range is not None - and not module.is_cross_attention - ) - - block_skip_range, timestep_skip_range, block_type = None, None, None - if is_spatial_self_attention: - block_skip_range = config.spatial_attention_block_skip_range - timestep_skip_range = config.spatial_attention_timestep_skip_range - block_type = "spatial" - elif is_temporal_self_attention: - block_skip_range = config.temporal_attention_block_skip_range - timestep_skip_range = config.temporal_attention_timestep_skip_range - block_type = "temporal" - - if block_skip_range is None or timestep_skip_range is None: - logger.debug( - f'Unable to apply FasterCache to the selected layer: "{name}" because it does ' - f"not match any of the required criteria for spatial or temporal attention layers. Note, " - f"however, that this layer may still be valid for applying PAB. Please specify the correct " - f"block identifiers in the configuration or use the specialized `apply_faster_cache_on_module` " - f"function to apply FasterCache to this layer." - ) - return - - logger.debug(f"Enabling FasterCache ({block_type}) for layer: {name}") - hook = FasterCacheBlockHook( - block_skip_range, - timestep_skip_range, - config.is_guidance_distilled, - config.attention_weight_callback, - config.current_timestep_callback, - ) - registry = HookRegistry.check_if_exists_or_initialize(module) - registry.register_hook(hook, _FASTER_CACHE_BLOCK_HOOK) - - -# Reference: https://github.com/Vchitect/FasterCache/blob/fab32c15014636dc854948319c0a9a8d92c7acb4/scripts/latte/faster_cache_sample_latte.py#L127C1-L143C39 -@torch.no_grad() -def _split_low_high_freq(x): - fft = torch.fft.fft2(x) - fft_shifted = torch.fft.fftshift(fft) - height, width = x.shape[-2:] - radius = min(height, width) // 5 - - y_grid, x_grid = torch.meshgrid(torch.arange(height), torch.arange(width)) - center_x, center_y = width // 2, height // 2 - mask = (x_grid - center_x) ** 2 + (y_grid - center_y) ** 2 <= radius**2 - - low_freq_mask = mask.unsqueeze(0).unsqueeze(0).to(x.device) - high_freq_mask = ~low_freq_mask - - low_freq_fft = fft_shifted * low_freq_mask - high_freq_fft = fft_shifted * high_freq_mask - - return low_freq_fft, high_freq_fft diff --git a/venv/lib/python3.10/site-packages/diffusers/hooks/group_offloading.py b/venv/lib/python3.10/site-packages/diffusers/hooks/group_offloading.py deleted file mode 100644 index 7186cb181aedac651b7977fc9b67bed016d2f0ad..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/hooks/group_offloading.py +++ /dev/null @@ -1,892 +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 os -from contextlib import contextmanager, nullcontext -from typing import Dict, List, Optional, Set, Tuple, Union - -import safetensors.torch -import torch - -from ..utils import get_logger, is_accelerate_available -from .hooks import HookRegistry, ModelHook - - -if is_accelerate_available(): - from accelerate.hooks import AlignDevicesHook, CpuOffload - from accelerate.utils import send_to_device - - -logger = get_logger(__name__) # pylint: disable=invalid-name - - -# fmt: off -_GROUP_OFFLOADING = "group_offloading" -_LAYER_EXECUTION_TRACKER = "layer_execution_tracker" -_LAZY_PREFETCH_GROUP_OFFLOADING = "lazy_prefetch_group_offloading" - -_SUPPORTED_PYTORCH_LAYERS = ( - torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d, - torch.nn.ConvTranspose1d, torch.nn.ConvTranspose2d, torch.nn.ConvTranspose3d, - torch.nn.Linear, - # 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 -) -# fmt: on - - -class ModuleGroup: - def __init__( - self, - 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, - non_blocking: bool = False, - stream: Union[torch.cuda.Stream, torch.Stream, None] = None, - record_stream: Optional[bool] = False, - low_cpu_mem_usage: bool = False, - onload_self: bool = True, - offload_to_disk_path: Optional[str] = None, - ) -> None: - self.modules = modules - self.offload_device = offload_device - self.onload_device = onload_device - self.offload_leader = offload_leader - self.onload_leader = onload_leader - self.parameters = parameters or [] - self.buffers = buffers or [] - self.non_blocking = non_blocking or stream is not None - self.stream = stream - self.record_stream = record_stream - self.onload_self = onload_self - self.low_cpu_mem_usage = low_cpu_mem_usage - - self.offload_to_disk_path = offload_to_disk_path - self._is_offloaded_to_disk = False - - if self.offload_to_disk_path: - self.safetensors_file_path = os.path.join(self.offload_to_disk_path, f"group_{id(self)}.safetensors") - - all_tensors = [] - for module in self.modules: - all_tensors.extend(list(module.parameters())) - all_tensors.extend(list(module.buffers())) - all_tensors.extend(self.parameters) - all_tensors.extend(self.buffers) - all_tensors = list(dict.fromkeys(all_tensors)) # Remove duplicates - - self.tensor_to_key = {tensor: f"tensor_{i}" for i, tensor in enumerate(all_tensors)} - self.key_to_tensor = {v: k for k, v in self.tensor_to_key.items()} - self.cpu_param_dict = {} - else: - self.cpu_param_dict = self._init_cpu_param_dict() - - def _init_cpu_param_dict(self): - cpu_param_dict = {} - if self.stream is None: - return cpu_param_dict - - for module in self.modules: - for param in module.parameters(): - cpu_param_dict[param] = param.data.cpu() if self.low_cpu_mem_usage else param.data.cpu().pin_memory() - for buffer in module.buffers(): - cpu_param_dict[buffer] = ( - buffer.data.cpu() if self.low_cpu_mem_usage else buffer.data.cpu().pin_memory() - ) - - for param in self.parameters: - cpu_param_dict[param] = param.data.cpu() if self.low_cpu_mem_usage else param.data.cpu().pin_memory() - - for buffer in self.buffers: - cpu_param_dict[buffer] = buffer.data.cpu() if self.low_cpu_mem_usage else buffer.data.cpu().pin_memory() - - return cpu_param_dict - - @contextmanager - def _pinned_memory_tensors(self): - pinned_dict = {} - try: - for param, tensor in self.cpu_param_dict.items(): - if not tensor.is_pinned(): - pinned_dict[param] = tensor.pin_memory() - else: - pinned_dict[param] = tensor - - yield pinned_dict - - finally: - pinned_dict = None - - @torch.compiler.disable() - def onload_(self): - r"""Onloads the group of modules to the onload_device.""" - torch_accelerator_module = ( - getattr(torch, torch.accelerator.current_accelerator().type) - if hasattr(torch, "accelerator") - else torch.cuda - ) - context = nullcontext() if self.stream is None else torch_accelerator_module.stream(self.stream) - current_stream = torch_accelerator_module.current_stream() if self.record_stream else None - - if self.offload_to_disk_path: - if self.stream is not None: - # Wait for previous Host->Device transfer to complete - self.stream.synchronize() - - with context: - if self.stream is not None: - # Load to CPU, pin, and async copy to device for overlapping transfer and compute - loaded_cpu_tensors = safetensors.torch.load_file(self.safetensors_file_path, device="cpu") - for key, tensor_obj in self.key_to_tensor.items(): - pinned_tensor = loaded_cpu_tensors[key].pin_memory() - tensor_obj.data = pinned_tensor.to(self.onload_device, non_blocking=self.non_blocking) - if self.record_stream: - tensor_obj.data.record_stream(current_stream) - else: - # Load directly to the target device (synchronous) - onload_device = ( - self.onload_device.type if isinstance(self.onload_device, torch.device) else self.onload_device - ) - loaded_tensors = safetensors.torch.load_file(self.safetensors_file_path, device=onload_device) - for key, tensor_obj in self.key_to_tensor.items(): - tensor_obj.data = loaded_tensors[key] - return - - if self.stream is not None: - # Wait for previous Host->Device transfer to complete - self.stream.synchronize() - - with context: - if self.stream is not None: - with self._pinned_memory_tensors() as pinned_memory: - for group_module in self.modules: - for param in group_module.parameters(): - param.data = pinned_memory[param].to(self.onload_device, non_blocking=self.non_blocking) - if self.record_stream: - param.data.record_stream(current_stream) - for buffer in group_module.buffers(): - buffer.data = pinned_memory[buffer].to(self.onload_device, non_blocking=self.non_blocking) - if self.record_stream: - buffer.data.record_stream(current_stream) - - for param in self.parameters: - param.data = pinned_memory[param].to(self.onload_device, non_blocking=self.non_blocking) - if self.record_stream: - param.data.record_stream(current_stream) - - for buffer in self.buffers: - buffer.data = pinned_memory[buffer].to(self.onload_device, non_blocking=self.non_blocking) - if self.record_stream: - buffer.data.record_stream(current_stream) - - else: - for group_module in self.modules: - for param in group_module.parameters(): - param.data = param.data.to(self.onload_device, non_blocking=self.non_blocking) - for buffer in group_module.buffers(): - buffer.data = buffer.data.to(self.onload_device, non_blocking=self.non_blocking) - - for param in self.parameters: - param.data = param.data.to(self.onload_device, non_blocking=self.non_blocking) - - for buffer in self.buffers: - buffer.data = buffer.data.to(self.onload_device, non_blocking=self.non_blocking) - if self.record_stream: - buffer.data.record_stream(current_stream) - - @torch.compiler.disable() - def offload_(self): - r"""Offloads the group of modules to the offload_device.""" - if self.offload_to_disk_path: - # TODO: we can potentially optimize this code path by checking if the _all_ the desired - # safetensor files exist on the disk and if so, skip this step entirely, reducing IO - # overhead. Currently, we just check if the given `safetensors_file_path` exists and if not - # we perform a write. - # Check if the file has been saved in this session or if it already exists on disk. - if not self._is_offloaded_to_disk and not os.path.exists(self.safetensors_file_path): - os.makedirs(os.path.dirname(self.safetensors_file_path), exist_ok=True) - tensors_to_save = { - key: tensor.data.to(self.offload_device) for tensor, key in self.tensor_to_key.items() - } - safetensors.torch.save_file(tensors_to_save, self.safetensors_file_path) - - # The group is now considered offloaded to disk for the rest of the session. - self._is_offloaded_to_disk = True - - # We do this to free up the RAM which is still holding the up tensor data. - for tensor_obj in self.tensor_to_key.keys(): - tensor_obj.data = torch.empty_like(tensor_obj.data, device=self.offload_device) - return - - torch_accelerator_module = ( - getattr(torch, torch.accelerator.current_accelerator().type) - if hasattr(torch, "accelerator") - else torch.cuda - ) - if self.stream is not None: - if not self.record_stream: - torch_accelerator_module.current_stream().synchronize() - for group_module in self.modules: - for param in group_module.parameters(): - param.data = self.cpu_param_dict[param] - for param in self.parameters: - param.data = self.cpu_param_dict[param] - for buffer in self.buffers: - buffer.data = self.cpu_param_dict[buffer] - - else: - for group_module in self.modules: - group_module.to(self.offload_device, non_blocking=self.non_blocking) - for param in self.parameters: - param.data = param.data.to(self.offload_device, non_blocking=self.non_blocking) - for buffer in self.buffers: - buffer.data = buffer.data.to(self.offload_device, non_blocking=self.non_blocking) - - -class GroupOffloadingHook(ModelHook): - r""" - A hook that offloads groups of torch.nn.Module to the CPU for storage and onloads to accelerator device for - computation. Each group has one "onload leader" module that is responsible for onloading, and an "offload leader" - module that is responsible for offloading. If prefetching is enabled, the onload leader of the previous module - group is responsible for onloading the current module group. - """ - - _is_stateful = False - - def __init__(self, group: ModuleGroup, next_group: Optional[ModuleGroup] = None) -> None: - self.group = group - self.next_group = next_group - - def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module: - if self.group.offload_leader == module: - self.group.offload_() - return module - - def pre_forward(self, module: torch.nn.Module, *args, **kwargs): - # If there wasn't an onload_leader assigned, we assume that the submodule that first called its forward - # method is the onload_leader of the group. - if self.group.onload_leader is None: - self.group.onload_leader = module - - # If the current module is the onload_leader of the group, we onload the group if it is supposed - # to onload itself. In the case of using prefetching with streams, we onload the next group if - # it is not supposed to onload itself. - if self.group.onload_leader == module: - if self.group.onload_self: - self.group.onload_() - if self.next_group is not None and not self.next_group.onload_self: - self.next_group.onload_() - - args = send_to_device(args, self.group.onload_device, non_blocking=self.group.non_blocking) - kwargs = send_to_device(kwargs, self.group.onload_device, non_blocking=self.group.non_blocking) - return args, kwargs - - def post_forward(self, module: torch.nn.Module, output): - if self.group.offload_leader == module: - self.group.offload_() - return output - - -class LazyPrefetchGroupOffloadingHook(ModelHook): - r""" - A hook, used in conjunction with GroupOffloadingHook, that applies lazy prefetching to groups of torch.nn.Module. - This hook is used to determine the order in which the layers are executed during the forward pass. Once the layer - invocation order is known, assignments of the next_group attribute for prefetching can be made, which allows - prefetching groups in the correct order. - """ - - _is_stateful = False - - def __init__(self): - self.execution_order: List[Tuple[str, torch.nn.Module]] = [] - self._layer_execution_tracker_module_names = set() - - def initialize_hook(self, module): - def make_execution_order_update_callback(current_name, current_submodule): - def callback(): - logger.debug(f"Adding {current_name} to the execution order") - self.execution_order.append((current_name, current_submodule)) - - return callback - - # To every submodule that contains a group offloading hook (at this point, no prefetching is enabled for any - # of the groups), we add a layer execution tracker hook that will be used to determine the order in which the - # layers are executed during the forward pass. - for name, submodule in module.named_modules(): - if name == "" or not hasattr(submodule, "_diffusers_hook"): - continue - - registry = HookRegistry.check_if_exists_or_initialize(submodule) - group_offloading_hook = registry.get_hook(_GROUP_OFFLOADING) - - if group_offloading_hook is not None: - # For the first forward pass, we have to load in a blocking manner - group_offloading_hook.group.non_blocking = False - layer_tracker_hook = LayerExecutionTrackerHook(make_execution_order_update_callback(name, submodule)) - registry.register_hook(layer_tracker_hook, _LAYER_EXECUTION_TRACKER) - self._layer_execution_tracker_module_names.add(name) - - return module - - def post_forward(self, module, output): - # At this point, for the current modules' submodules, we know the execution order of the layers. We can now - # remove the layer execution tracker hooks and apply prefetching by setting the next_group attribute for each - # group offloading hook. - num_executed = len(self.execution_order) - execution_order_module_names = {name for name, _ in self.execution_order} - - # It may be possible that some layers were not executed during the forward pass. This can happen if the layer - # is not used in the forward pass, or if the layer is not executed due to some other reason. In such cases, we - # may not be able to apply prefetching in the correct order, which can lead to device-mismatch related errors - # if the missing layers end up being executed in the future. - if execution_order_module_names != self._layer_execution_tracker_module_names: - unexecuted_layers = list(self._layer_execution_tracker_module_names - execution_order_module_names) - logger.warning( - "It seems like some layers were not executed during the forward pass. This may lead to problems when " - "applying lazy prefetching with automatic tracing and lead to device-mismatch related errors. Please " - "make sure that all layers are executed during the forward pass. The following layers were not executed:\n" - f"{unexecuted_layers=}" - ) - - # Remove the layer execution tracker hooks from the submodules - base_module_registry = module._diffusers_hook - registries = [submodule._diffusers_hook for _, submodule in self.execution_order] - group_offloading_hooks = [registry.get_hook(_GROUP_OFFLOADING) for registry in registries] - - for i in range(num_executed): - registries[i].remove_hook(_LAYER_EXECUTION_TRACKER, recurse=False) - - # Remove the current lazy prefetch group offloading hook so that it doesn't interfere with the next forward pass - base_module_registry.remove_hook(_LAZY_PREFETCH_GROUP_OFFLOADING, recurse=False) - - # LazyPrefetchGroupOffloadingHook is only used with streams, so we know that non_blocking should be True. - # We disable non_blocking for the first forward pass, but need to enable it for the subsequent passes to - # see the benefits of prefetching. - for hook in group_offloading_hooks: - hook.group.non_blocking = True - - # Set required attributes for prefetching - if num_executed > 0: - base_module_group_offloading_hook = base_module_registry.get_hook(_GROUP_OFFLOADING) - base_module_group_offloading_hook.next_group = group_offloading_hooks[0].group - base_module_group_offloading_hook.next_group.onload_self = False - - for i in range(num_executed - 1): - name1, _ = self.execution_order[i] - name2, _ = self.execution_order[i + 1] - logger.debug(f"Applying lazy prefetch group offloading from {name1} to {name2}") - group_offloading_hooks[i].next_group = group_offloading_hooks[i + 1].group - group_offloading_hooks[i].next_group.onload_self = False - - return output - - -class LayerExecutionTrackerHook(ModelHook): - r""" - A hook that tracks the order in which the layers are executed during the forward pass by calling back to the - LazyPrefetchGroupOffloadingHook to update the execution order. - """ - - _is_stateful = False - - def __init__(self, execution_order_update_callback): - self.execution_order_update_callback = execution_order_update_callback - - def pre_forward(self, module, *args, **kwargs): - self.execution_order_update_callback() - return args, kwargs - - -def apply_group_offloading( - module: torch.nn.Module, - onload_device: torch.device, - offload_device: torch.device = torch.device("cpu"), - offload_type: str = "block_level", - num_blocks_per_group: Optional[int] = 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, -) -> None: - r""" - Applies group offloading to the internal layers of a torch.nn.Module. To understand what group offloading is, and - where it is beneficial, we need to first provide some context on how other supported offloading methods work. - - Typically, offloading is done at two levels: - - Module-level: In Diffusers, this can be enabled using the `ModelMixin::enable_model_cpu_offload()` method. It - works by offloading each component of a pipeline to the CPU for storage, and onloading to the accelerator device - when needed for computation. This method is more memory-efficient than keeping all components on the accelerator, - but the memory requirements are still quite high. For this method to work, one needs memory equivalent to size of - the model in runtime dtype + size of largest intermediate activation tensors to be able to complete the forward - pass. - - Leaf-level: In Diffusers, this can be enabled using the `ModelMixin::enable_sequential_cpu_offload()` method. It - works by offloading the lowest leaf-level parameters of the computation graph to the CPU for storage, and - onloading only the leafs to the accelerator device for computation. This uses the lowest amount of accelerator - memory, but can be slower due to the excessive number of device synchronizations. - - Group offloading is a middle ground between the two methods. It works by offloading groups of internal layers, - (either `torch.nn.ModuleList` or `torch.nn.Sequential`). This method uses lower memory than module-level - offloading. It is also faster than leaf-level/sequential offloading, as the number of device synchronizations is - reduced. - - Another supported feature (for CUDA devices with support for asynchronous data transfer streams) is the ability to - overlap data transfer and computation to reduce the overall execution time compared to sequential offloading. This - is enabled using layer prefetching with streams, i.e., the layer that is to be executed next starts onloading to - the accelerator device while the current layer is being executed - this increases the memory requirements slightly. - Note that this implementation also supports leaf-level offloading but can be made much faster when using streams. - - Args: - module (`torch.nn.Module`): - The module to which group offloading is applied. - onload_device (`torch.device`): - The device to which the group of modules are onloaded. - offload_device (`torch.device`, defaults to `torch.device("cpu")`): - The device to which the group of modules are offloaded. This should typically be the CPU. Default is CPU. - offload_type (`str`, defaults to "block_level"): - The type of offloading to be applied. Can be one of "block_level" or "leaf_level". Default is - "block_level". - offload_to_disk_path (`str`, *optional*, defaults to `None`): - The path to the directory where parameters will be offloaded. Setting this option can be useful in limited - RAM environment settings where a reasonable speed-memory trade-off is desired. - num_blocks_per_group (`int`, *optional*): - The number of blocks per group when using offload_type="block_level". This is required when using - offload_type="block_level". - non_blocking (`bool`, defaults to `False`): - If True, offloading and onloading is done with non-blocking data transfer. - use_stream (`bool`, defaults to `False`): - If True, offloading and onloading is done asynchronously using a CUDA stream. This can be useful for - overlapping computation and data transfer. - record_stream (`bool`, defaults to `False`): When enabled with `use_stream`, it marks the current tensor - as having been used by this stream. It is faster at the expense of slightly more memory usage. Refer to the - [PyTorch official docs](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) more - details. - low_cpu_mem_usage (`bool`, defaults to `False`): - 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. - - Example: - ```python - >>> from diffusers import CogVideoXTransformer3DModel - >>> from diffusers.hooks import apply_group_offloading - - >>> transformer = CogVideoXTransformer3DModel.from_pretrained( - ... "THUDM/CogVideoX-5b", subfolder="transformer", torch_dtype=torch.bfloat16 - ... ) - - >>> apply_group_offloading( - ... transformer, - ... onload_device=torch.device("cuda"), - ... offload_device=torch.device("cpu"), - ... offload_type="block_level", - ... num_blocks_per_group=2, - ... use_stream=True, - ... ) - ``` - """ - - stream = None - if use_stream: - if torch.cuda.is_available(): - stream = torch.cuda.Stream() - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - stream = torch.Stream() - else: - raise ValueError("Using streams for data transfer requires a CUDA device, or an Intel XPU device.") - - if not use_stream and record_stream: - raise ValueError("`record_stream` cannot be True when `use_stream=False`.") - - _raise_error_if_accelerate_model_or_sequential_hook_present(module) - - if offload_type == "block_level": - if num_blocks_per_group is None: - raise ValueError("num_blocks_per_group must be provided when using offload_type='block_level'.") - - _apply_group_offloading_block_level( - module=module, - num_blocks_per_group=num_blocks_per_group, - offload_device=offload_device, - onload_device=onload_device, - offload_to_disk_path=offload_to_disk_path, - non_blocking=non_blocking, - stream=stream, - record_stream=record_stream, - low_cpu_mem_usage=low_cpu_mem_usage, - ) - elif offload_type == "leaf_level": - _apply_group_offloading_leaf_level( - module=module, - offload_device=offload_device, - onload_device=onload_device, - offload_to_disk_path=offload_to_disk_path, - non_blocking=non_blocking, - stream=stream, - record_stream=record_stream, - low_cpu_mem_usage=low_cpu_mem_usage, - ) - else: - raise ValueError(f"Unsupported offload_type: {offload_type}") - - -def _apply_group_offloading_block_level( - module: torch.nn.Module, - num_blocks_per_group: int, - offload_device: torch.device, - onload_device: torch.device, - non_blocking: bool, - stream: Union[torch.cuda.Stream, torch.Stream, None] = None, - record_stream: Optional[bool] = False, - low_cpu_mem_usage: bool = False, - offload_to_disk_path: Optional[str] = None, -) -> None: - r""" - This function applies offloading to groups of torch.nn.ModuleList or torch.nn.Sequential blocks. In comparison to - the "leaf_level" offloading, which is more fine-grained, this offloading is done at the top-level blocks. - - Args: - module (`torch.nn.Module`): - The module to which group offloading is applied. - offload_device (`torch.device`): - The device to which the group of modules are offloaded. This should typically be the CPU. - offload_to_disk_path (`str`, *optional*, defaults to `None`): - The path to the directory where parameters will be offloaded. Setting this option can be useful in limited - RAM environment settings where a reasonable speed-memory trade-off is desired. - onload_device (`torch.device`): - The device to which the group of modules are onloaded. - non_blocking (`bool`): - If True, offloading and onloading is done asynchronously. This can be useful for overlapping computation - and data transfer. - stream (`torch.cuda.Stream`or `torch.Stream`, *optional*): - If provided, offloading and onloading is done asynchronously using the provided stream. This can be useful - for overlapping computation and data transfer. - record_stream (`bool`, defaults to `False`): When enabled with `use_stream`, it marks the current tensor - as having been used by this stream. It is faster at the expense of slightly more memory usage. Refer to the - [PyTorch official docs](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) more - details. - low_cpu_mem_usage (`bool`, defaults to `False`): - 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. - """ - if stream is not None and num_blocks_per_group != 1: - logger.warning( - f"Using streams is only supported for num_blocks_per_group=1. Got {num_blocks_per_group=}. Setting it to 1." - ) - num_blocks_per_group = 1 - - # Create module groups for ModuleList and Sequential blocks - modules_with_group_offloading = set() - unmatched_modules = [] - matched_module_groups = [] - for name, submodule in module.named_children(): - if not isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)): - unmatched_modules.append((name, submodule)) - modules_with_group_offloading.add(name) - continue - - for i in range(0, len(submodule), num_blocks_per_group): - current_modules = submodule[i : i + num_blocks_per_group] - group = ModuleGroup( - modules=current_modules, - offload_device=offload_device, - onload_device=onload_device, - offload_to_disk_path=offload_to_disk_path, - offload_leader=current_modules[-1], - onload_leader=current_modules[0], - non_blocking=non_blocking, - stream=stream, - record_stream=record_stream, - low_cpu_mem_usage=low_cpu_mem_usage, - onload_self=True, - ) - matched_module_groups.append(group) - for j in range(i, i + len(current_modules)): - modules_with_group_offloading.add(f"{name}.{j}") - - # Apply group offloading hooks to the module groups - for i, group in enumerate(matched_module_groups): - for group_module in group.modules: - _apply_group_offloading_hook(group_module, group, None) - - # Parameters and Buffers of the top-level module need to be offloaded/onloaded separately - # when the forward pass of this module is called. This is because the top-level module is not - # part of any group (as doing so would lead to no VRAM savings). - parameters = _gather_parameters_with_no_group_offloading_parent(module, modules_with_group_offloading) - buffers = _gather_buffers_with_no_group_offloading_parent(module, modules_with_group_offloading) - parameters = [param for _, param in parameters] - buffers = [buffer for _, buffer in buffers] - - # Create a group for the unmatched submodules of the top-level module so that they are on the correct - # device when the forward pass is called. - unmatched_modules = [unmatched_module for _, unmatched_module in unmatched_modules] - unmatched_group = ModuleGroup( - modules=unmatched_modules, - offload_device=offload_device, - onload_device=onload_device, - offload_to_disk_path=offload_to_disk_path, - offload_leader=module, - onload_leader=module, - parameters=parameters, - buffers=buffers, - non_blocking=False, - stream=None, - record_stream=False, - onload_self=True, - ) - if stream is None: - _apply_group_offloading_hook(module, unmatched_group, None) - else: - _apply_lazy_group_offloading_hook(module, unmatched_group, None) - - -def _apply_group_offloading_leaf_level( - module: torch.nn.Module, - offload_device: torch.device, - onload_device: torch.device, - non_blocking: bool, - stream: Union[torch.cuda.Stream, torch.Stream, None] = None, - record_stream: Optional[bool] = False, - low_cpu_mem_usage: bool = False, - offload_to_disk_path: Optional[str] = None, -) -> None: - r""" - This function applies offloading to groups of leaf modules in a torch.nn.Module. This method has minimal memory - requirements. However, it can be slower compared to other offloading methods due to the excessive number of device - synchronizations. When using devices that support streams to overlap data transfer and computation, this method can - reduce memory usage without any performance degradation. - - Args: - module (`torch.nn.Module`): - The module to which group offloading is applied. - offload_device (`torch.device`): - The device to which the group of modules are offloaded. This should typically be the CPU. - onload_device (`torch.device`): - The device to which the group of modules are onloaded. - offload_to_disk_path (`str`, *optional*, defaults to `None`): - The path to the directory where parameters will be offloaded. Setting this option can be useful in limited - RAM environment settings where a reasonable speed-memory trade-off is desired. - non_blocking (`bool`): - If True, offloading and onloading is done asynchronously. This can be useful for overlapping computation - and data transfer. - stream (`torch.cuda.Stream` or `torch.Stream`, *optional*): - If provided, offloading and onloading is done asynchronously using the provided stream. This can be useful - for overlapping computation and data transfer. - record_stream (`bool`, defaults to `False`): When enabled with `use_stream`, it marks the current tensor - as having been used by this stream. It is faster at the expense of slightly more memory usage. Refer to the - [PyTorch official docs](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) more - details. - low_cpu_mem_usage (`bool`, defaults to `False`): - 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. - """ - - # Create module groups for leaf modules and apply group offloading hooks - modules_with_group_offloading = set() - for name, submodule in module.named_modules(): - if not isinstance(submodule, _SUPPORTED_PYTORCH_LAYERS): - continue - group = ModuleGroup( - modules=[submodule], - offload_device=offload_device, - onload_device=onload_device, - offload_to_disk_path=offload_to_disk_path, - offload_leader=submodule, - onload_leader=submodule, - non_blocking=non_blocking, - stream=stream, - record_stream=record_stream, - low_cpu_mem_usage=low_cpu_mem_usage, - onload_self=True, - ) - _apply_group_offloading_hook(submodule, group, None) - modules_with_group_offloading.add(name) - - # Parameters and Buffers at all non-leaf levels need to be offloaded/onloaded separately when the forward pass - # of the module is called - module_dict = dict(module.named_modules()) - parameters = _gather_parameters_with_no_group_offloading_parent(module, modules_with_group_offloading) - buffers = _gather_buffers_with_no_group_offloading_parent(module, modules_with_group_offloading) - - # Find closest module parent for each parameter and buffer, and attach group hooks - parent_to_parameters = {} - for name, param in parameters: - parent_name = _find_parent_module_in_module_dict(name, module_dict) - if parent_name in parent_to_parameters: - parent_to_parameters[parent_name].append(param) - else: - parent_to_parameters[parent_name] = [param] - - parent_to_buffers = {} - for name, buffer in buffers: - parent_name = _find_parent_module_in_module_dict(name, module_dict) - if parent_name in parent_to_buffers: - parent_to_buffers[parent_name].append(buffer) - else: - parent_to_buffers[parent_name] = [buffer] - - parent_names = set(parent_to_parameters.keys()) | set(parent_to_buffers.keys()) - for name in parent_names: - parameters = parent_to_parameters.get(name, []) - buffers = parent_to_buffers.get(name, []) - parent_module = module_dict[name] - assert getattr(parent_module, "_diffusers_hook", None) is None - group = ModuleGroup( - modules=[], - offload_device=offload_device, - onload_device=onload_device, - offload_leader=parent_module, - onload_leader=parent_module, - offload_to_disk_path=offload_to_disk_path, - parameters=parameters, - buffers=buffers, - non_blocking=non_blocking, - stream=stream, - record_stream=record_stream, - low_cpu_mem_usage=low_cpu_mem_usage, - onload_self=True, - ) - _apply_group_offloading_hook(parent_module, group, None) - - if stream is not None: - # When using streams, we need to know the layer execution order for applying prefetching (to overlap data transfer - # and computation). Since we don't know the order beforehand, we apply a lazy prefetching hook that will find the - # execution order and apply prefetching in the correct order. - unmatched_group = ModuleGroup( - modules=[], - offload_device=offload_device, - onload_device=onload_device, - offload_to_disk_path=offload_to_disk_path, - offload_leader=module, - onload_leader=module, - parameters=None, - buffers=None, - non_blocking=False, - stream=None, - record_stream=False, - low_cpu_mem_usage=low_cpu_mem_usage, - onload_self=True, - ) - _apply_lazy_group_offloading_hook(module, unmatched_group, None) - - -def _apply_group_offloading_hook( - module: torch.nn.Module, - group: ModuleGroup, - next_group: Optional[ModuleGroup] = None, -) -> None: - registry = HookRegistry.check_if_exists_or_initialize(module) - - # We may have already registered a group offloading hook if the module had a torch.nn.Parameter whose parent - # is the current module. In such cases, we don't want to overwrite the existing group offloading hook. - if registry.get_hook(_GROUP_OFFLOADING) is None: - hook = GroupOffloadingHook(group, next_group) - registry.register_hook(hook, _GROUP_OFFLOADING) - - -def _apply_lazy_group_offloading_hook( - module: torch.nn.Module, - group: ModuleGroup, - next_group: Optional[ModuleGroup] = None, -) -> None: - registry = HookRegistry.check_if_exists_or_initialize(module) - - # We may have already registered a group offloading hook if the module had a torch.nn.Parameter whose parent - # is the current module. In such cases, we don't want to overwrite the existing group offloading hook. - if registry.get_hook(_GROUP_OFFLOADING) is None: - hook = GroupOffloadingHook(group, next_group) - registry.register_hook(hook, _GROUP_OFFLOADING) - - lazy_prefetch_hook = LazyPrefetchGroupOffloadingHook() - registry.register_hook(lazy_prefetch_hook, _LAZY_PREFETCH_GROUP_OFFLOADING) - - -def _gather_parameters_with_no_group_offloading_parent( - module: torch.nn.Module, modules_with_group_offloading: Set[str] -) -> List[torch.nn.Parameter]: - parameters = [] - for name, parameter in module.named_parameters(): - has_parent_with_group_offloading = False - atoms = name.split(".") - while len(atoms) > 0: - parent_name = ".".join(atoms) - if parent_name in modules_with_group_offloading: - has_parent_with_group_offloading = True - break - atoms.pop() - if not has_parent_with_group_offloading: - parameters.append((name, parameter)) - return parameters - - -def _gather_buffers_with_no_group_offloading_parent( - module: torch.nn.Module, modules_with_group_offloading: Set[str] -) -> List[torch.Tensor]: - buffers = [] - for name, buffer in module.named_buffers(): - has_parent_with_group_offloading = False - atoms = name.split(".") - while len(atoms) > 0: - parent_name = ".".join(atoms) - if parent_name in modules_with_group_offloading: - has_parent_with_group_offloading = True - break - atoms.pop() - if not has_parent_with_group_offloading: - buffers.append((name, buffer)) - return buffers - - -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) - if parent_name in module_dict: - return parent_name - atoms.pop() - return "" - - -def _raise_error_if_accelerate_model_or_sequential_hook_present(module: torch.nn.Module) -> None: - if not is_accelerate_available(): - return - for name, submodule in module.named_modules(): - if not hasattr(submodule, "_hf_hook"): - continue - if isinstance(submodule._hf_hook, (AlignDevicesHook, CpuOffload)): - raise ValueError( - f"Cannot apply group offloading to a module that is already applying an alternative " - f"offloading strategy from Accelerate. If you want to apply group offloading, please " - f"disable the existing offloading strategy first. Offending module: {name} ({type(submodule)})" - ) - - -def _is_group_offload_enabled(module: torch.nn.Module) -> bool: - for submodule in module.modules(): - if hasattr(submodule, "_diffusers_hook") and submodule._diffusers_hook.get_hook(_GROUP_OFFLOADING) is not None: - return True - return False - - -def _get_group_onload_device(module: torch.nn.Module) -> torch.device: - for submodule in module.modules(): - if hasattr(submodule, "_diffusers_hook") and submodule._diffusers_hook.get_hook(_GROUP_OFFLOADING) is not None: - return submodule._diffusers_hook.get_hook(_GROUP_OFFLOADING).group.onload_device - raise ValueError("Group offloading is not enabled for the provided module.") diff --git a/venv/lib/python3.10/site-packages/diffusers/hooks/hooks.py b/venv/lib/python3.10/site-packages/diffusers/hooks/hooks.py deleted file mode 100644 index 96231aadc3f7a068db9586ea363901af649f941c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/hooks/hooks.py +++ /dev/null @@ -1,236 +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 functools -from typing import Any, Dict, Optional, Tuple - -import torch - -from ..utils.logging import get_logger - - -logger = get_logger(__name__) # pylint: disable=invalid-name - - -class ModelHook: - r""" - A hook that contains callbacks to be executed just before and after the forward method of a model. - """ - - _is_stateful = False - - def __init__(self): - self.fn_ref: "HookFunctionReference" = None - - def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module: - r""" - Hook that is executed when a model is initialized. - - Args: - module (`torch.nn.Module`): - The module attached to this hook. - """ - return module - - def deinitalize_hook(self, module: torch.nn.Module) -> torch.nn.Module: - r""" - Hook that is executed when a model is deinitialized. - - Args: - module (`torch.nn.Module`): - The module attached to this hook. - """ - return module - - 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]`): - The positional arguments passed to the module. - kwargs (`Dict[Str, Any]`): - The keyword arguments passed to the module. - Returns: - `Tuple[Tuple[Any], Dict[Str, Any]]`: - A tuple with the treated `args` and `kwargs`. - """ - return args, kwargs - - def post_forward(self, module: torch.nn.Module, output: Any) -> Any: - r""" - Hook that is executed just after the forward method of the model. - - Args: - module (`torch.nn.Module`): - The module whose forward pass been executed just before this event. - output (`Any`): - The output of the module. - Returns: - `Any`: The processed `output`. - """ - return output - - def detach_hook(self, module: torch.nn.Module) -> torch.nn.Module: - r""" - Hook that is executed when the hook is detached from a module. - - Args: - module (`torch.nn.Module`): - The module detached from this hook. - """ - return module - - def reset_state(self, module: torch.nn.Module): - if self._is_stateful: - raise NotImplementedError("This hook is stateful and needs to implement the `reset_state` method.") - return module - - -class HookFunctionReference: - def __init__(self) -> None: - """A container class that maintains mutable references to forward pass functions in a hook chain. - - Its mutable nature allows the hook system to modify the execution chain dynamically without rebuilding the - entire forward pass structure. - - Attributes: - pre_forward: A callable that processes inputs before the main forward pass. - post_forward: A callable that processes outputs after the main forward pass. - forward: The current forward function in the hook chain. - original_forward: The original forward function, stored when a hook provides a custom new_forward. - - The class enables hook removal by allowing updates to the forward chain through reference modification rather - than requiring reconstruction of the entire chain. When a hook is removed, only the relevant references need to - be updated, preserving the execution order of the remaining hooks. - """ - self.pre_forward = None - self.post_forward = None - self.forward = None - self.original_forward = None - - -class HookRegistry: - def __init__(self, module_ref: torch.nn.Module) -> None: - super().__init__() - - self.hooks: Dict[str, ModelHook] = {} - - self._module_ref = module_ref - self._hook_order = [] - self._fn_refs = [] - - def register_hook(self, hook: ModelHook, name: str) -> None: - if name in self.hooks.keys(): - raise ValueError( - f"Hook with name {name} already exists in the registry. Please use a different name or " - f"first remove the existing hook and then add a new one." - ) - - self._module_ref = hook.initialize_hook(self._module_ref) - - def create_new_forward(function_reference: HookFunctionReference): - def new_forward(module, *args, **kwargs): - args, kwargs = function_reference.pre_forward(module, *args, **kwargs) - output = function_reference.forward(*args, **kwargs) - return function_reference.post_forward(module, output) - - return new_forward - - forward = self._module_ref.forward - - fn_ref = HookFunctionReference() - fn_ref.pre_forward = hook.pre_forward - fn_ref.post_forward = hook.post_forward - fn_ref.forward = forward - - if hasattr(hook, "new_forward"): - fn_ref.original_forward = forward - fn_ref.forward = functools.update_wrapper( - functools.partial(hook.new_forward, self._module_ref), hook.new_forward - ) - - rewritten_forward = create_new_forward(fn_ref) - self._module_ref.forward = functools.update_wrapper( - functools.partial(rewritten_forward, self._module_ref), rewritten_forward - ) - - hook.fn_ref = fn_ref - self.hooks[name] = hook - self._hook_order.append(name) - self._fn_refs.append(fn_ref) - - def get_hook(self, name: str) -> Optional[ModelHook]: - return self.hooks.get(name, None) - - def remove_hook(self, name: str, recurse: bool = True) -> None: - if name in self.hooks.keys(): - num_hooks = len(self._hook_order) - hook = self.hooks[name] - index = self._hook_order.index(name) - fn_ref = self._fn_refs[index] - - old_forward = fn_ref.forward - if fn_ref.original_forward is not None: - old_forward = fn_ref.original_forward - - if index == num_hooks - 1: - self._module_ref.forward = old_forward - else: - self._fn_refs[index + 1].forward = old_forward - - self._module_ref = hook.deinitalize_hook(self._module_ref) - del self.hooks[name] - self._hook_order.pop(index) - self._fn_refs.pop(index) - - if recurse: - for module_name, module in self._module_ref.named_modules(): - if module_name == "": - continue - if hasattr(module, "_diffusers_hook"): - module._diffusers_hook.remove_hook(name, recurse=False) - - def reset_stateful_hooks(self, recurse: bool = True) -> None: - for hook_name in reversed(self._hook_order): - hook = self.hooks[hook_name] - if hook._is_stateful: - hook.reset_state(self._module_ref) - - if recurse: - for module_name, module in self._module_ref.named_modules(): - if module_name == "": - continue - if hasattr(module, "_diffusers_hook"): - module._diffusers_hook.reset_stateful_hooks(recurse=False) - - @classmethod - def check_if_exists_or_initialize(cls, module: torch.nn.Module) -> "HookRegistry": - if not hasattr(module, "_diffusers_hook"): - module._diffusers_hook = cls(module) - return module._diffusers_hook - - def __repr__(self) -> str: - registry_repr = "" - for i, hook_name in enumerate(self._hook_order): - if self.hooks[hook_name].__class__.__repr__ is not object.__repr__: - hook_repr = self.hooks[hook_name].__repr__() - else: - hook_repr = self.hooks[hook_name].__class__.__name__ - registry_repr += f" ({i}) {hook_name} - {hook_repr}" - if i < len(self._hook_order) - 1: - registry_repr += "\n" - return f"HookRegistry(\n{registry_repr}\n)" diff --git a/venv/lib/python3.10/site-packages/diffusers/hooks/layerwise_casting.py b/venv/lib/python3.10/site-packages/diffusers/hooks/layerwise_casting.py deleted file mode 100644 index 1747a5c489cba3f34548deee0a44d8c85dabfe6b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/hooks/layerwise_casting.py +++ /dev/null @@ -1,245 +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 re -from typing import Optional, Tuple, Type, Union - -import torch - -from ..utils import get_logger, is_peft_available, is_peft_version -from .hooks import HookRegistry, ModelHook - - -logger = get_logger(__name__) # pylint: disable=invalid-name - - -# fmt: off -_LAYERWISE_CASTING_HOOK = "layerwise_casting" -_PEFT_AUTOCAST_DISABLE_HOOK = "peft_autocast_disable" -SUPPORTED_PYTORCH_LAYERS = ( - torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d, - torch.nn.ConvTranspose1d, torch.nn.ConvTranspose2d, torch.nn.ConvTranspose3d, - torch.nn.Linear, -) - -DEFAULT_SKIP_MODULES_PATTERN = ("pos_embed", "patch_embed", "norm", "^proj_in$", "^proj_out$") -# fmt: on - -_SHOULD_DISABLE_PEFT_INPUT_AUTOCAST = is_peft_available() and is_peft_version(">", "0.14.0") -if _SHOULD_DISABLE_PEFT_INPUT_AUTOCAST: - from peft.helpers import disable_input_dtype_casting - from peft.tuners.tuners_utils import BaseTunerLayer - - -class LayerwiseCastingHook(ModelHook): - r""" - A hook that casts the weights of a module to a high precision dtype for computation, and to a low precision dtype - for storage. This process may lead to quality loss in the output, but can significantly reduce the memory - footprint. - """ - - _is_stateful = False - - def __init__(self, storage_dtype: torch.dtype, compute_dtype: torch.dtype, non_blocking: bool) -> None: - self.storage_dtype = storage_dtype - self.compute_dtype = compute_dtype - self.non_blocking = non_blocking - - def initialize_hook(self, module: torch.nn.Module): - module.to(dtype=self.storage_dtype, non_blocking=self.non_blocking) - return module - - def deinitalize_hook(self, module: torch.nn.Module): - raise NotImplementedError( - "LayerwiseCastingHook does not support deinitialization. A model once enabled with layerwise casting will " - "have casted its weights to a lower precision dtype for storage. Casting this back to the original dtype " - "will lead to precision loss, which might have an impact on the model's generation quality. The model should " - "be re-initialized and loaded in the original dtype." - ) - - def pre_forward(self, module: torch.nn.Module, *args, **kwargs): - module.to(dtype=self.compute_dtype, non_blocking=self.non_blocking) - return args, kwargs - - def post_forward(self, module: torch.nn.Module, output): - module.to(dtype=self.storage_dtype, non_blocking=self.non_blocking) - return output - - -class PeftInputAutocastDisableHook(ModelHook): - r""" - A hook that disables the casting of inputs to the module weight dtype during the forward pass. By default, PEFT - casts the inputs to the weight dtype of the module, which can lead to precision loss. - - The reasons for needing this are: - - If we don't add PEFT layers' weight names to `skip_modules_pattern` when applying layerwise casting, the - inputs will be casted to the, possibly lower precision, storage dtype. Reference: - https://github.com/huggingface/peft/blob/0facdebf6208139cbd8f3586875acb378813dd97/src/peft/tuners/lora/layer.py#L706 - - We can, on our end, use something like accelerate's `send_to_device` but for dtypes. This way, we can ensure - that the inputs are casted to the computation dtype correctly always. However, there are two goals we are - hoping to achieve: - 1. Making forward implementations independent of device/dtype casting operations as much as possible. - 2. Performing inference without losing information from casting to different precisions. With the current - PEFT implementation (as linked in the reference above), and assuming running layerwise casting inference - with storage_dtype=torch.float8_e4m3fn and compute_dtype=torch.bfloat16, inputs are cast to - torch.float8_e4m3fn in the lora layer. We will then upcast back to torch.bfloat16 when we continue the - forward pass in PEFT linear forward or Diffusers layer forward, with a `send_to_dtype` operation from - LayerwiseCastingHook. This will be a lossy operation and result in poorer generation quality. - """ - - def new_forward(self, module: torch.nn.Module, *args, **kwargs): - with disable_input_dtype_casting(module): - return self.fn_ref.original_forward(*args, **kwargs) - - -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, - non_blocking: bool = False, -) -> None: - r""" - Applies layerwise casting to a given module. The module expected here is a Diffusers ModelMixin but it can be any - nn.Module using diffusers layers or pytorch primitives. - - Example: - - ```python - >>> import torch - >>> from diffusers import CogVideoXTransformer3DModel - - >>> transformer = CogVideoXTransformer3DModel.from_pretrained( - ... model_id, subfolder="transformer", torch_dtype=torch.bfloat16 - ... ) - - >>> apply_layerwise_casting( - ... transformer, - ... storage_dtype=torch.float8_e4m3fn, - ... compute_dtype=torch.bfloat16, - ... skip_modules_pattern=["patch_embed", "norm", "proj_out"], - ... non_blocking=True, - ... ) - ``` - - Args: - module (`torch.nn.Module`): - The module whose leaf modules will be cast to a high precision dtype for computation, and to a low - precision dtype for storage. - storage_dtype (`torch.dtype`): - 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"`): - 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`): - 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. - """ - if skip_modules_pattern == "auto": - skip_modules_pattern = DEFAULT_SKIP_MODULES_PATTERN - - if skip_modules_classes is None and skip_modules_pattern is None: - apply_layerwise_casting_hook(module, storage_dtype, compute_dtype, non_blocking) - return - - _apply_layerwise_casting( - module, - storage_dtype, - compute_dtype, - skip_modules_pattern, - skip_modules_classes, - non_blocking, - ) - _disable_peft_input_autocast(module) - - -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, - non_blocking: bool = False, - _prefix: str = "", -) -> None: - should_skip = (skip_modules_classes is not None and isinstance(module, skip_modules_classes)) or ( - skip_modules_pattern is not None and any(re.search(pattern, _prefix) for pattern in skip_modules_pattern) - ) - if should_skip: - logger.debug(f'Skipping layerwise casting for layer "{_prefix}"') - return - - if isinstance(module, SUPPORTED_PYTORCH_LAYERS): - logger.debug(f'Applying layerwise casting to layer "{_prefix}"') - apply_layerwise_casting_hook(module, storage_dtype, compute_dtype, non_blocking) - return - - for name, submodule in module.named_children(): - layer_name = f"{_prefix}.{name}" if _prefix else name - _apply_layerwise_casting( - submodule, - storage_dtype, - compute_dtype, - skip_modules_pattern, - skip_modules_classes, - non_blocking, - _prefix=layer_name, - ) - - -def apply_layerwise_casting_hook( - module: torch.nn.Module, storage_dtype: torch.dtype, compute_dtype: torch.dtype, non_blocking: bool -) -> None: - r""" - Applies a `LayerwiseCastingHook` to a given module. - - Args: - module (`torch.nn.Module`): - The module to attach the hook to. - storage_dtype (`torch.dtype`): - The dtype to cast the module to before the forward pass. - compute_dtype (`torch.dtype`): - The dtype to cast the module to during the forward pass. - non_blocking (`bool`): - If `True`, the weight casting operations are non-blocking. - """ - registry = HookRegistry.check_if_exists_or_initialize(module) - hook = LayerwiseCastingHook(storage_dtype, compute_dtype, non_blocking) - registry.register_hook(hook, _LAYERWISE_CASTING_HOOK) - - -def _is_layerwise_casting_active(module: torch.nn.Module) -> bool: - for submodule in module.modules(): - if ( - hasattr(submodule, "_diffusers_hook") - and submodule._diffusers_hook.get_hook(_LAYERWISE_CASTING_HOOK) is not None - ): - return True - return False - - -def _disable_peft_input_autocast(module: torch.nn.Module) -> None: - if not _SHOULD_DISABLE_PEFT_INPUT_AUTOCAST: - return - for submodule in module.modules(): - if isinstance(submodule, BaseTunerLayer) and _is_layerwise_casting_active(submodule): - registry = HookRegistry.check_if_exists_or_initialize(submodule) - hook = PeftInputAutocastDisableHook() - registry.register_hook(hook, _PEFT_AUTOCAST_DISABLE_HOOK) diff --git a/venv/lib/python3.10/site-packages/diffusers/hooks/pyramid_attention_broadcast.py b/venv/lib/python3.10/site-packages/diffusers/hooks/pyramid_attention_broadcast.py deleted file mode 100644 index bbdd1c3f68d4a0f1d216cf130f533ac708f68ecf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/hooks/pyramid_attention_broadcast.py +++ /dev/null @@ -1,311 +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 re -from dataclasses import dataclass -from typing import Any, Callable, Optional, Tuple, Union - -import torch - -from ..models.attention_processor import Attention, MochiAttention -from ..utils import logging -from .hooks import HookRegistry, ModelHook - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -_PYRAMID_ATTENTION_BROADCAST_HOOK = "pyramid_attention_broadcast" -_ATTENTION_CLASSES = (Attention, MochiAttention) -_SPATIAL_ATTENTION_BLOCK_IDENTIFIERS = ("blocks", "transformer_blocks", "single_transformer_blocks") -_TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS = ("temporal_transformer_blocks",) -_CROSS_ATTENTION_BLOCK_IDENTIFIERS = ("blocks", "transformer_blocks") - - -@dataclass -class PyramidAttentionBroadcastConfig: - r""" - Configuration for Pyramid Attention Broadcast. - - Args: - spatial_attention_block_skip_range (`int`, *optional*, defaults to `None`): - The number of times a specific spatial 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 re-used) before computing the new attention states again. - temporal_attention_block_skip_range (`int`, *optional*, defaults to `None`): - The number of times a specific temporal 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 re-used) before computing the new attention states again. - cross_attention_block_skip_range (`int`, *optional*, defaults to `None`): - 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 re-used) before computing the new attention states again. - 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)`): - 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)`): - 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, ...]`, defaults to `("blocks", "transformer_blocks")`): - The identifiers to match against the layer names to determine if the layer is a spatial attention layer. - temporal_attention_block_identifiers (`Tuple[str, ...]`, defaults to `("temporal_transformer_blocks",)`): - The identifiers to match against the layer names to determine if the layer is a temporal attention layer. - cross_attention_block_identifiers (`Tuple[str, ...]`, defaults to `("blocks", "transformer_blocks")`): - 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_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_ATTENTION_BLOCK_IDENTIFIERS - temporal_attention_block_identifiers: Tuple[str, ...] = _TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS - cross_attention_block_identifiers: Tuple[str, ...] = _CROSS_ATTENTION_BLOCK_IDENTIFIERS - - current_timestep_callback: Callable[[], int] = None - - # TODO(aryan): add PAB for MLP layers (very limited speedup from testing with original codebase - # so not added for now) - - def __repr__(self) -> str: - return ( - f"PyramidAttentionBroadcastConfig(\n" - f" spatial_attention_block_skip_range={self.spatial_attention_block_skip_range},\n" - f" temporal_attention_block_skip_range={self.temporal_attention_block_skip_range},\n" - f" cross_attention_block_skip_range={self.cross_attention_block_skip_range},\n" - f" spatial_attention_timestep_skip_range={self.spatial_attention_timestep_skip_range},\n" - f" temporal_attention_timestep_skip_range={self.temporal_attention_timestep_skip_range},\n" - f" cross_attention_timestep_skip_range={self.cross_attention_timestep_skip_range},\n" - f" spatial_attention_block_identifiers={self.spatial_attention_block_identifiers},\n" - f" temporal_attention_block_identifiers={self.temporal_attention_block_identifiers},\n" - f" cross_attention_block_identifiers={self.cross_attention_block_identifiers},\n" - f" current_timestep_callback={self.current_timestep_callback}\n" - ")" - ) - - -class PyramidAttentionBroadcastState: - r""" - State for Pyramid Attention Broadcast. - - Attributes: - iteration (`int`): - The current iteration of the Pyramid Attention Broadcast. It is necessary to ensure that `reset_state` is - called before starting a new inference forward pass for PAB to work correctly. - cache (`Any`): - The cached output from the previous forward pass. This is used to re-use the attention states when the - attention computation is skipped. It is either a tensor or a tuple of tensors, depending on the module. - """ - - def __init__(self) -> None: - self.iteration = 0 - self.cache = None - - def reset(self): - self.iteration = 0 - self.cache = None - - def __repr__(self): - cache_repr = "" - if self.cache is None: - cache_repr = "None" - else: - cache_repr = f"Tensor(shape={self.cache.shape}, dtype={self.cache.dtype})" - return f"PyramidAttentionBroadcastState(iteration={self.iteration}, cache={cache_repr})" - - -class PyramidAttentionBroadcastHook(ModelHook): - r"""A hook that applies Pyramid Attention Broadcast to a given module.""" - - _is_stateful = True - - def __init__( - self, timestep_skip_range: Tuple[int, int], block_skip_range: int, current_timestep_callback: Callable[[], int] - ) -> None: - super().__init__() - - self.timestep_skip_range = timestep_skip_range - self.block_skip_range = block_skip_range - self.current_timestep_callback = current_timestep_callback - - def initialize_hook(self, module): - self.state = PyramidAttentionBroadcastState() - return module - - def new_forward(self, module: torch.nn.Module, *args, **kwargs) -> Any: - is_within_timestep_range = ( - self.timestep_skip_range[0] < self.current_timestep_callback() < self.timestep_skip_range[1] - ) - should_compute_attention = ( - self.state.cache is None - or self.state.iteration == 0 - or not is_within_timestep_range - or self.state.iteration % self.block_skip_range == 0 - ) - - if should_compute_attention: - output = self.fn_ref.original_forward(*args, **kwargs) - else: - output = self.state.cache - - self.state.cache = output - self.state.iteration += 1 - return output - - def reset_state(self, module: torch.nn.Module) -> None: - self.state.reset() - return module - - -def apply_pyramid_attention_broadcast(module: torch.nn.Module, config: PyramidAttentionBroadcastConfig): - r""" - Apply [Pyramid Attention Broadcast](https://huggingface.co/papers/2408.12588) to a given pipeline. - - PAB is an attention approximation method that leverages the similarity in attention states between timesteps to - reduce the computational cost of attention computation. The key takeaway from the paper is that the attention - similarity in the cross-attention layers between timesteps is high, followed by less similarity in the temporal and - spatial layers. This allows for the skipping of attention computation in the cross-attention layers more frequently - than in the temporal and spatial layers. Applying PAB will, therefore, speedup the inference process. - - Args: - module (`torch.nn.Module`): - The module to apply Pyramid Attention Broadcast to. - config (`Optional[PyramidAttentionBroadcastConfig]`, `optional`, defaults to `None`): - The configuration to use for Pyramid Attention Broadcast. - - Example: - - ```python - >>> import torch - >>> from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast - >>> from diffusers.utils import export_to_video - - >>> pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16) - >>> pipe.to("cuda") - - >>> config = PyramidAttentionBroadcastConfig( - ... spatial_attention_block_skip_range=2, - ... spatial_attention_timestep_skip_range=(100, 800), - ... current_timestep_callback=lambda: pipe.current_timestep, - ... ) - >>> apply_pyramid_attention_broadcast(pipe.transformer, config) - ``` - """ - if config.current_timestep_callback is None: - raise ValueError( - "The `current_timestep_callback` function must be provided in the configuration to apply Pyramid Attention Broadcast." - ) - - if ( - config.spatial_attention_block_skip_range is None - and config.temporal_attention_block_skip_range is None - and config.cross_attention_block_skip_range is None - ): - logger.warning( - "Pyramid Attention Broadcast requires one or more of `spatial_attention_block_skip_range`, `temporal_attention_block_skip_range` " - "or `cross_attention_block_skip_range` parameters to be set to an integer, not `None`. Defaulting to using `spatial_attention_block_skip_range=2`. " - "To avoid this warning, please set one of the above parameters." - ) - config.spatial_attention_block_skip_range = 2 - - for name, submodule in module.named_modules(): - if not isinstance(submodule, _ATTENTION_CLASSES): - # PAB has been implemented specific to Diffusers' Attention classes. However, this does not mean that PAB - # cannot be applied to this layer. For custom layers, users can extend this functionality and implement - # their own PAB logic similar to `_apply_pyramid_attention_broadcast_on_attention_class`. - continue - _apply_pyramid_attention_broadcast_on_attention_class(name, submodule, config) - - -def _apply_pyramid_attention_broadcast_on_attention_class( - name: str, module: Attention, config: PyramidAttentionBroadcastConfig -) -> bool: - is_spatial_self_attention = ( - any(re.search(identifier, name) is not None for identifier in config.spatial_attention_block_identifiers) - and config.spatial_attention_block_skip_range is not None - and not getattr(module, "is_cross_attention", False) - ) - is_temporal_self_attention = ( - any(re.search(identifier, name) is not None for identifier in config.temporal_attention_block_identifiers) - and config.temporal_attention_block_skip_range is not None - and not getattr(module, "is_cross_attention", False) - ) - is_cross_attention = ( - any(re.search(identifier, name) is not None for identifier in config.cross_attention_block_identifiers) - and config.cross_attention_block_skip_range is not None - and getattr(module, "is_cross_attention", False) - ) - - block_skip_range, timestep_skip_range, block_type = None, None, None - if is_spatial_self_attention: - block_skip_range = config.spatial_attention_block_skip_range - timestep_skip_range = config.spatial_attention_timestep_skip_range - block_type = "spatial" - elif is_temporal_self_attention: - block_skip_range = config.temporal_attention_block_skip_range - timestep_skip_range = config.temporal_attention_timestep_skip_range - block_type = "temporal" - elif is_cross_attention: - block_skip_range = config.cross_attention_block_skip_range - timestep_skip_range = config.cross_attention_timestep_skip_range - block_type = "cross" - - if block_skip_range is None or timestep_skip_range is None: - logger.info( - f'Unable to apply Pyramid Attention Broadcast to the selected layer: "{name}" because it does ' - f"not match any of the required criteria for spatial, temporal or cross attention layers. Note, " - f"however, that this layer may still be valid for applying PAB. Please specify the correct " - f"block identifiers in the configuration." - ) - return False - - logger.debug(f"Enabling Pyramid Attention Broadcast ({block_type}) in layer: {name}") - _apply_pyramid_attention_broadcast_hook( - module, timestep_skip_range, block_skip_range, config.current_timestep_callback - ) - return True - - -def _apply_pyramid_attention_broadcast_hook( - module: Union[Attention, MochiAttention], - timestep_skip_range: Tuple[int, int], - block_skip_range: int, - current_timestep_callback: Callable[[], int], -): - r""" - Apply [Pyramid Attention Broadcast](https://huggingface.co/papers/2408.12588) to a given torch.nn.Module. - - Args: - module (`torch.nn.Module`): - The module to apply Pyramid Attention Broadcast to. - 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`): - The number of times a specific 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 re-used) before computing the new attention states again. - current_timestep_callback (`Callable[[], int]`): - A callback function that returns the current inference timestep. - """ - registry = HookRegistry.check_if_exists_or_initialize(module) - hook = PyramidAttentionBroadcastHook(timestep_skip_range, block_skip_range, current_timestep_callback) - registry.register_hook(hook, _PYRAMID_ATTENTION_BROADCAST_HOOK) diff --git a/venv/lib/python3.10/site-packages/diffusers/image_processor.py b/venv/lib/python3.10/site-packages/diffusers/image_processor.py deleted file mode 100644 index 6a3cf77a7df77577c895431ec89366305e176fc8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/image_processor.py +++ /dev/null @@ -1,1319 +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 math -import warnings -from typing import List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from PIL import Image, ImageFilter, ImageOps - -from .configuration_utils import ConfigMixin, register_to_config -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], -] - -PipelineDepthInput = PipelineImageInput - - -def is_valid_image(image) -> bool: - r""" - Checks if the input is a valid image. - - A valid image can be: - - A `PIL.Image.Image`. - - A 2D or 3D `np.ndarray` or `torch.Tensor` (grayscale or color image). - - Args: - image (`Union[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: - `bool`: - `True` if the input is a valid image, `False` otherwise. - """ - return isinstance(image, PIL.Image.Image) or isinstance(image, (np.ndarray, torch.Tensor)) and image.ndim in (2, 3) - - -def is_valid_image_imagelist(images): - r""" - Checks if the input is a valid image or list of images. - - The input can be one of the following formats: - - A 4D tensor or numpy array (batch of images). - - A valid single image: `PIL.Image.Image`, 2D `np.ndarray` or `torch.Tensor` (grayscale image), 3D `np.ndarray` or - `torch.Tensor`. - - A list of valid images. - - Args: - images (`Union[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. - - Returns: - `bool`: - `True` if the input is valid, `False` otherwise. - """ - if isinstance(images, (np.ndarray, torch.Tensor)) and images.ndim == 4: - return True - elif is_valid_image(images): - return True - elif isinstance(images, list): - return all(is_valid_image(image) for image in images) - return False - - -class VaeImageProcessor(ConfigMixin): - """ - Image processor for VAE. - - Args: - do_resize (`bool`, *optional*, defaults to `True`): - Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. Can accept - `height` and `width` arguments from [`image_processor.VaeImageProcessor.preprocess`] method. - vae_scale_factor (`int`, *optional*, defaults to `8`): - VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor. - resample (`str`, *optional*, defaults to `lanczos`): - Resampling filter to use when resizing the image. - do_normalize (`bool`, *optional*, defaults to `True`): - Whether to normalize the image to [-1,1]. - do_binarize (`bool`, *optional*, defaults to `False`): - Whether to binarize the image to 0/1. - do_convert_rgb (`bool`, *optional*, defaults to be `False`): - 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. - """ - - config_name = CONFIG_NAME - - @register_to_config - def __init__( - self, - do_resize: bool = True, - vae_scale_factor: int = 8, - vae_latent_channels: int = 4, - 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, - ): - super().__init__() - if do_convert_rgb and do_convert_grayscale: - raise ValueError( - "`do_convert_rgb` and `do_convert_grayscale` can not both be set to `True`," - " if you intended to convert the image into RGB format, please set `do_convert_grayscale = False`.", - " if you intended to convert the image into grayscale format, please set `do_convert_rgb = False`", - ) - - @staticmethod - 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. - - Args: - images (`np.ndarray`): - The image array to convert to PIL format. - - Returns: - `List[PIL.Image.Image]`: - A list of PIL images. - """ - if images.ndim == 3: - images = images[None, ...] - images = (images * 255).round().astype("uint8") - if images.shape[-1] == 1: - # special case for grayscale (single channel) images - pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images] - else: - pil_images = [Image.fromarray(image) for image in images] - - return pil_images - - @staticmethod - def pil_to_numpy(images: Union[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]`): - The PIL image or list of images to convert to NumPy format. - - Returns: - `np.ndarray`: - A NumPy array representation of the images. - """ - if not isinstance(images, list): - images = [images] - images = [np.array(image).astype(np.float32) / 255.0 for image in images] - images = np.stack(images, axis=0) - - return images - - @staticmethod - def numpy_to_pt(images: np.ndarray) -> torch.Tensor: - r""" - Convert a NumPy image to a PyTorch tensor. - - Args: - images (`np.ndarray`): - The NumPy image array to convert to PyTorch format. - - Returns: - `torch.Tensor`: - A PyTorch tensor representation of the images. - """ - if images.ndim == 3: - images = images[..., None] - - images = torch.from_numpy(images.transpose(0, 3, 1, 2)) - return images - - @staticmethod - def pt_to_numpy(images: torch.Tensor) -> np.ndarray: - r""" - Convert a PyTorch tensor to a NumPy image. - - Args: - images (`torch.Tensor`): - The PyTorch tensor to convert to NumPy format. - - Returns: - `np.ndarray`: - A NumPy array representation of the images. - """ - images = images.cpu().permute(0, 2, 3, 1).float().numpy() - return images - - @staticmethod - def normalize(images: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: - r""" - Normalize an image array to [-1,1]. - - Args: - images (`np.ndarray` or `torch.Tensor`): - The image array to normalize. - - Returns: - `np.ndarray` or `torch.Tensor`: - The normalized image array. - """ - return 2.0 * images - 1.0 - - @staticmethod - def denormalize(images: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: - r""" - Denormalize an image array to [0,1]. - - Args: - images (`np.ndarray` or `torch.Tensor`): - The image array to denormalize. - - Returns: - `np.ndarray` or `torch.Tensor`: - The denormalized image array. - """ - return (images * 0.5 + 0.5).clamp(0, 1) - - @staticmethod - def convert_to_rgb(image: PIL.Image.Image) -> PIL.Image.Image: - r""" - Converts a PIL image to RGB format. - - Args: - image (`PIL.Image.Image`): - The PIL image to convert to RGB. - - Returns: - `PIL.Image.Image`: - The RGB-converted PIL image. - """ - image = image.convert("RGB") - - return image - - @staticmethod - def convert_to_grayscale(image: PIL.Image.Image) -> PIL.Image.Image: - r""" - Converts a given PIL image to grayscale. - - Args: - image (`PIL.Image.Image`): - The input image to convert. - - Returns: - `PIL.Image.Image`: - The image converted to grayscale. - """ - image = image.convert("L") - - return image - - @staticmethod - def blur(image: PIL.Image.Image, blur_factor: int = 4) -> PIL.Image.Image: - r""" - Applies Gaussian blur to an image. - - Args: - image (`PIL.Image.Image`): - The PIL image to convert to grayscale. - - Returns: - `PIL.Image.Image`: - The grayscale-converted PIL image. - """ - image = image.filter(ImageFilter.GaussianBlur(blur_factor)) - - return image - - @staticmethod - def get_crop_region(mask_image: PIL.Image.Image, width: int, height: int, pad=0): - r""" - Finds a rectangular region that contains all masked ares in an image, and expands region to match the aspect - ratio of the original image; for example, if user drew mask in a 128x32 region, and the dimensions for - processing are 512x512, the region will be expanded to 128x128. - - Args: - mask_image (PIL.Image.Image): Mask image. - width (int): Width of the image to be processed. - height (int): Height of the image to be processed. - pad (int, optional): Padding to be added to the crop region. Defaults to 0. - - Returns: - tuple: (x1, y1, x2, y2) represent a rectangular region that contains all masked ares in an image and - matches the original aspect ratio. - """ - - mask_image = mask_image.convert("L") - mask = np.array(mask_image) - - # 1. find a rectangular region that contains all masked ares in an image - h, w = mask.shape - crop_left = 0 - for i in range(w): - if not (mask[:, i] == 0).all(): - break - crop_left += 1 - - crop_right = 0 - for i in reversed(range(w)): - if not (mask[:, i] == 0).all(): - break - crop_right += 1 - - crop_top = 0 - for i in range(h): - if not (mask[i] == 0).all(): - break - crop_top += 1 - - crop_bottom = 0 - for i in reversed(range(h)): - if not (mask[i] == 0).all(): - break - crop_bottom += 1 - - # 2. add padding to the crop region - x1, y1, x2, y2 = ( - int(max(crop_left - pad, 0)), - int(max(crop_top - pad, 0)), - int(min(w - crop_right + pad, w)), - int(min(h - crop_bottom + pad, h)), - ) - - # 3. expands crop region to match the aspect ratio of the image to be processed - ratio_crop_region = (x2 - x1) / (y2 - y1) - ratio_processing = width / height - - if ratio_crop_region > ratio_processing: - desired_height = (x2 - x1) / ratio_processing - desired_height_diff = int(desired_height - (y2 - y1)) - y1 -= desired_height_diff // 2 - y2 += desired_height_diff - desired_height_diff // 2 - if y2 >= mask_image.height: - diff = y2 - mask_image.height - y2 -= diff - y1 -= diff - if y1 < 0: - y2 -= y1 - y1 -= y1 - if y2 >= mask_image.height: - y2 = mask_image.height - else: - desired_width = (y2 - y1) * ratio_processing - desired_width_diff = int(desired_width - (x2 - x1)) - x1 -= desired_width_diff // 2 - x2 += desired_width_diff - desired_width_diff // 2 - if x2 >= mask_image.width: - diff = x2 - mask_image.width - x2 -= diff - x1 -= diff - if x1 < 0: - x2 -= x1 - x1 -= x1 - if x2 >= mask_image.width: - x2 = mask_image.width - - return x1, y1, x2, y2 - - def _resize_and_fill( - self, - image: PIL.Image.Image, - width: int, - height: int, - ) -> PIL.Image.Image: - r""" - Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center - the image within the dimensions, filling empty with data from image. - - Args: - image (`PIL.Image.Image`): - The image to resize and fill. - width (`int`): - The width to resize the image to. - height (`int`): - The height to resize the image to. - - Returns: - `PIL.Image.Image`: - The resized and filled image. - """ - - ratio = width / height - src_ratio = image.width / image.height - - src_w = width if ratio < src_ratio else image.width * height // image.height - src_h = height if ratio >= src_ratio else image.height * width // image.width - - resized = image.resize((src_w, src_h), resample=PIL_INTERPOLATION["lanczos"]) - res = Image.new("RGB", (width, height)) - res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2)) - - if ratio < src_ratio: - fill_height = height // 2 - src_h // 2 - if fill_height > 0: - res.paste(resized.resize((width, fill_height), box=(0, 0, width, 0)), box=(0, 0)) - res.paste( - resized.resize((width, fill_height), box=(0, resized.height, width, resized.height)), - box=(0, fill_height + src_h), - ) - elif ratio > src_ratio: - fill_width = width // 2 - src_w // 2 - if fill_width > 0: - res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0)) - res.paste( - resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)), - box=(fill_width + src_w, 0), - ) - - return res - - def _resize_and_crop( - self, - image: PIL.Image.Image, - width: int, - height: int, - ) -> PIL.Image.Image: - r""" - Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center - the image within the dimensions, cropping the excess. - - Args: - image (`PIL.Image.Image`): - The image to resize and crop. - width (`int`): - The width to resize the image to. - height (`int`): - The height to resize the image to. - - Returns: - `PIL.Image.Image`: - The resized and cropped image. - """ - ratio = width / height - src_ratio = image.width / image.height - - src_w = width if ratio > src_ratio else image.width * height // image.height - src_h = height if ratio <= src_ratio else image.height * width // image.width - - resized = image.resize((src_w, src_h), resample=PIL_INTERPOLATION["lanczos"]) - res = Image.new("RGB", (width, height)) - res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2)) - return res - - def resize( - self, - image: Union[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]: - """ - Resize image. - - Args: - image (`PIL.Image.Image`, `np.ndarray` or `torch.Tensor`): - The image input, can be a PIL image, numpy array or pytorch tensor. - height (`int`): - The height to resize to. - width (`int`): - The width to resize to. - resize_mode (`str`, *optional*, defaults to `default`): - The resize mode to use, can be one of `default` or `fill`. If `default`, will resize the image to fit - within the specified width and height, and it may not maintaining the original aspect ratio. If `fill`, - will resize the image to fit within the specified width and height, maintaining the aspect ratio, and - then center the image within the dimensions, filling empty with data from image. If `crop`, will resize - the 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. - - Returns: - `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`: - The resized image. - """ - if resize_mode != "default" and not isinstance(image, PIL.Image.Image): - raise ValueError(f"Only PIL image input is supported for resize_mode {resize_mode}") - if isinstance(image, PIL.Image.Image): - if resize_mode == "default": - image = image.resize( - (width, height), - resample=PIL_INTERPOLATION[self.config.resample], - reducing_gap=self.config.reducing_gap, - ) - elif resize_mode == "fill": - image = self._resize_and_fill(image, width, height) - elif resize_mode == "crop": - image = self._resize_and_crop(image, width, height) - else: - raise ValueError(f"resize_mode {resize_mode} is not supported") - - elif isinstance(image, torch.Tensor): - image = torch.nn.functional.interpolate( - image, - size=(height, width), - ) - elif isinstance(image, np.ndarray): - image = self.numpy_to_pt(image) - image = torch.nn.functional.interpolate( - image, - size=(height, width), - ) - image = self.pt_to_numpy(image) - return image - - def binarize(self, image: PIL.Image.Image) -> PIL.Image.Image: - """ - Create a mask. - - Args: - image (`PIL.Image.Image`): - The image input, should be a PIL image. - - Returns: - `PIL.Image.Image`: - The binarized image. Values less than 0.5 are set to 0, values greater than 0.5 are set to 1. - """ - image[image < 0.5] = 0 - image[image >= 0.5] = 1 - - return image - - def _denormalize_conditionally( - self, images: torch.Tensor, do_denormalize: Optional[List[bool]] = None - ) -> torch.Tensor: - r""" - Denormalize a batch of images based on a condition list. - - Args: - images (`torch.Tensor`): - The input image tensor. - 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. - """ - if do_denormalize is None: - return self.denormalize(images) if self.config.do_normalize else images - - return torch.stack( - [self.denormalize(images[i]) if do_denormalize[i] else images[i] for i in range(images.shape[0])] - ) - - 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]: - 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]`): - 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`): - The height of the preprocessed image. If `None`, the height of the `image` input will be used. - width (`Optional[int]`, *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]`: - A tuple containing the height and width, both resized to the nearest integer multiple of - `vae_scale_factor`. - """ - - if height is None: - if isinstance(image, PIL.Image.Image): - height = image.height - elif isinstance(image, torch.Tensor): - height = image.shape[2] - else: - height = image.shape[1] - - if width is None: - if isinstance(image, PIL.Image.Image): - width = image.width - elif isinstance(image, torch.Tensor): - width = image.shape[3] - else: - width = image.shape[2] - - width, height = ( - x - x % self.config.vae_scale_factor for x in (width, height) - ) # resize to integer multiple of vae_scale_factor - - return height, width - - def preprocess( - self, - image: PipelineImageInput, - height: Optional[int] = None, - width: Optional[int] = None, - resize_mode: str = "default", # "default", "fill", "crop" - crops_coords: Optional[Tuple[int, int, int, int]] = None, - ) -> torch.Tensor: - """ - Preprocess the image input. - - Args: - image (`PipelineImageInput`): - The image input, accepted formats are PIL images, NumPy arrays, PyTorch tensors; Also accept list of - supported formats. - height (`int`, *optional*): - The height in preprocessed image. If `None`, will use the `get_default_height_width()` to get default - height. - width (`int`, *optional*): - The width in preprocessed. If `None`, will use get_default_height_width()` to get the default width. - resize_mode (`str`, *optional*, defaults to `default`): - The resize mode, can be one of `default` or `fill`. If `default`, will resize the image to fit within - the specified width and height, and it may not maintaining the original aspect ratio. If `fill`, will - resize the image to fit within the specified width and height, maintaining the aspect ratio, and then - center the image within the dimensions, filling empty with data from image. If `crop`, will resize the - 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`): - The crop coordinates for each image in the batch. If `None`, will not crop the image. - - Returns: - `torch.Tensor`: - The preprocessed image. - """ - supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor) - - # Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image - if self.config.do_convert_grayscale and isinstance(image, (torch.Tensor, np.ndarray)) and image.ndim == 3: - if isinstance(image, torch.Tensor): - # if image is a pytorch tensor could have 2 possible shapes: - # 1. batch x height x width: we should insert the channel dimension at position 1 - # 2. channel x height x width: we should insert batch dimension at position 0, - # however, since both channel and batch dimension has same size 1, it is same to insert at position 1 - # for simplicity, we insert a dimension of size 1 at position 1 for both cases - image = image.unsqueeze(1) - else: - # if it is a numpy array, it could have 2 possible shapes: - # 1. batch x height x width: insert channel dimension on last position - # 2. height x width x channel: insert batch dimension on first position - if image.shape[-1] == 1: - image = np.expand_dims(image, axis=0) - else: - image = np.expand_dims(image, axis=-1) - - if isinstance(image, list) and isinstance(image[0], np.ndarray) and image[0].ndim == 4: - warnings.warn( - "Passing `image` as a list of 4d np.ndarray is deprecated." - "Please concatenate the list along the batch dimension and pass it as a single 4d np.ndarray", - FutureWarning, - ) - image = np.concatenate(image, axis=0) - if isinstance(image, list) and isinstance(image[0], torch.Tensor) and image[0].ndim == 4: - warnings.warn( - "Passing `image` as a list of 4d torch.Tensor is deprecated." - "Please concatenate the list along the batch dimension and pass it as a single 4d torch.Tensor", - FutureWarning, - ) - image = torch.cat(image, axis=0) - - if not is_valid_image_imagelist(image): - raise ValueError( - f"Input is in incorrect format. Currently, we only support {', '.join(str(x) for x in supported_formats)}" - ) - if not isinstance(image, list): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - if crops_coords is not None: - image = [i.crop(crops_coords) for i in image] - if self.config.do_resize: - height, width = self.get_default_height_width(image[0], height, width) - image = [self.resize(i, height, width, resize_mode=resize_mode) for i in image] - if self.config.do_convert_rgb: - image = [self.convert_to_rgb(i) for i in image] - elif self.config.do_convert_grayscale: - image = [self.convert_to_grayscale(i) for i in image] - image = self.pil_to_numpy(image) # to np - image = self.numpy_to_pt(image) # to pt - - elif isinstance(image[0], np.ndarray): - image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0) - - image = self.numpy_to_pt(image) - - height, width = self.get_default_height_width(image, height, width) - if self.config.do_resize: - image = self.resize(image, height, width) - - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0) - - if self.config.do_convert_grayscale and image.ndim == 3: - image = image.unsqueeze(1) - - channel = image.shape[1] - # don't need any preprocess if the image is latents - if channel == self.config.vae_latent_channels: - return image - - height, width = self.get_default_height_width(image, height, width) - if self.config.do_resize: - image = self.resize(image, height, width) - - # expected range [0,1], normalize to [-1,1] - do_normalize = self.config.do_normalize - if do_normalize and image.min() < 0: - warnings.warn( - "Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] " - f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{image.min()},{image.max()}]", - FutureWarning, - ) - do_normalize = False - if do_normalize: - image = self.normalize(image) - - if self.config.do_binarize: - image = self.binarize(image) - - return image - - def postprocess( - self, - image: torch.Tensor, - output_type: str = "pil", - do_denormalize: Optional[List[bool]] = None, - ) -> Union[PIL.Image.Image, np.ndarray, torch.Tensor]: - """ - Postprocess the image output from tensor to `output_type`. - - Args: - image (`torch.Tensor`): - 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`): - Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the - `VaeImageProcessor` config. - - Returns: - `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`: - The postprocessed image. - """ - if not isinstance(image, torch.Tensor): - raise ValueError( - f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor" - ) - if output_type not in ["latent", "pt", "np", "pil"]: - deprecation_message = ( - f"the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: " - "`pil`, `np`, `pt`, `latent`" - ) - deprecate("Unsupported output_type", "1.0.0", deprecation_message, standard_warn=False) - output_type = "np" - - if output_type == "latent": - return image - - image = self._denormalize_conditionally(image, do_denormalize) - - if output_type == "pt": - return image - - image = self.pt_to_numpy(image) - - if output_type == "np": - return image - - if output_type == "pil": - return self.numpy_to_pil(image) - - def apply_overlay( - self, - mask: PIL.Image.Image, - init_image: PIL.Image.Image, - image: PIL.Image.Image, - crop_coords: Optional[Tuple[int, int, int, int]] = None, - ) -> PIL.Image.Image: - r""" - Applies an overlay of the mask and the inpainted image on the original image. - - Args: - mask (`PIL.Image.Image`): - The mask image that highlights regions to overlay. - init_image (`PIL.Image.Image`): - 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*): - Coordinates to crop the image. If provided, the image will be cropped accordingly. - - Returns: - `PIL.Image.Image`: - The final image with the overlay applied. - """ - - width, height = init_image.width, init_image.height - - init_image_masked = PIL.Image.new("RGBa", (width, height)) - init_image_masked.paste(init_image.convert("RGBA").convert("RGBa"), mask=ImageOps.invert(mask.convert("L"))) - - init_image_masked = init_image_masked.convert("RGBA") - - if crop_coords is not None: - x, y, x2, y2 = crop_coords - w = x2 - x - h = y2 - y - base_image = PIL.Image.new("RGBA", (width, height)) - image = self.resize(image, height=h, width=w, resize_mode="crop") - base_image.paste(image, (x, y)) - image = base_image.convert("RGB") - - image = image.convert("RGBA") - image.alpha_composite(init_image_masked) - image = image.convert("RGB") - - return image - - -class VaeImageProcessorLDM3D(VaeImageProcessor): - """ - Image processor for VAE LDM3D. - - Args: - do_resize (`bool`, *optional*, defaults to `True`): - Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. - vae_scale_factor (`int`, *optional*, defaults to `8`): - VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor. - resample (`str`, *optional*, defaults to `lanczos`): - Resampling filter to use when resizing the image. - do_normalize (`bool`, *optional*, defaults to `True`): - Whether to normalize the image to [-1,1]. - """ - - config_name = CONFIG_NAME - - @register_to_config - def __init__( - self, - do_resize: bool = True, - vae_scale_factor: int = 8, - resample: str = "lanczos", - do_normalize: bool = True, - ): - super().__init__() - - @staticmethod - 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. - - Args: - images (`np.ndarray`): - The input NumPy array of images, which can be a single image or a batch. - - Returns: - `List[PIL.Image.Image]`: - A list of PIL images converted from the input NumPy array. - """ - if images.ndim == 3: - images = images[None, ...] - images = (images * 255).round().astype("uint8") - if images.shape[-1] == 1: - # special case for grayscale (single channel) images - pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images] - else: - pil_images = [Image.fromarray(image[:, :, :3]) for image in images] - - return pil_images - - @staticmethod - def depth_pil_to_numpy(images: Union[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]`): - The input image or list of images to be converted. - - Returns: - `np.ndarray`: - A NumPy array of the converted images. - """ - if not isinstance(images, list): - images = [images] - - images = [np.array(image).astype(np.float32) / (2**16 - 1) for image in images] - images = np.stack(images, axis=0) - return images - - @staticmethod - def rgblike_to_depthmap(image: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: - r""" - Convert an RGB-like depth image to a depth map. - - Args: - image (`Union[np.ndarray, torch.Tensor]`): - The RGB-like depth image to convert. - - Returns: - `Union[np.ndarray, torch.Tensor]`: - The corresponding depth map. - """ - return image[:, :, 1] * 2**8 + image[:, :, 2] - - 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. - - Args: - images (`np.ndarray`): - The input NumPy array of depth images, which can be a single image or a batch. - - Returns: - `List[PIL.Image.Image]`: - A list of PIL images converted from the input NumPy depth images. - """ - if images.ndim == 3: - images = images[None, ...] - images_depth = images[:, :, :, 3:] - if images.shape[-1] == 6: - images_depth = (images_depth * 255).round().astype("uint8") - pil_images = [ - Image.fromarray(self.rgblike_to_depthmap(image_depth), mode="I;16") for image_depth in images_depth - ] - elif images.shape[-1] == 4: - images_depth = (images_depth * 65535.0).astype(np.uint16) - pil_images = [Image.fromarray(image_depth, mode="I;16") for image_depth in images_depth] - else: - raise Exception("Not supported") - - return pil_images - - def postprocess( - self, - image: torch.Tensor, - output_type: str = "pil", - do_denormalize: Optional[List[bool]] = None, - ) -> Union[PIL.Image.Image, np.ndarray, torch.Tensor]: - """ - Postprocess the image output from tensor to `output_type`. - - Args: - image (`torch.Tensor`): - 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`): - Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the - `VaeImageProcessor` config. - - Returns: - `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`: - The postprocessed image. - """ - if not isinstance(image, torch.Tensor): - raise ValueError( - f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor" - ) - if output_type not in ["latent", "pt", "np", "pil"]: - deprecation_message = ( - f"the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: " - "`pil`, `np`, `pt`, `latent`" - ) - deprecate("Unsupported output_type", "1.0.0", deprecation_message, standard_warn=False) - output_type = "np" - - image = self._denormalize_conditionally(image, do_denormalize) - - image = self.pt_to_numpy(image) - - if output_type == "np": - if image.shape[-1] == 6: - image_depth = np.stack([self.rgblike_to_depthmap(im[:, :, 3:]) for im in image], axis=0) - else: - image_depth = image[:, :, :, 3:] - return image[:, :, :, :3], image_depth - - if output_type == "pil": - return self.numpy_to_pil(image), self.numpy_to_depth(image) - else: - raise Exception(f"This type {output_type} is not supported") - - 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, - ) -> 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]`): - The RGB input image, which can be a single image or a batch. - depth (`Union[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`): - The desired height of the processed image. If `None`, defaults to the height of the input image. - width (`Optional[int]`, *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 resolution for resizing the images. If specified, overrides height and width. - - Returns: - `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) - - # Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image - if self.config.do_convert_grayscale and isinstance(rgb, (torch.Tensor, np.ndarray)) and rgb.ndim == 3: - raise Exception("This is not yet supported") - - if isinstance(rgb, supported_formats): - rgb = [rgb] - depth = [depth] - elif not (isinstance(rgb, list) and all(isinstance(i, supported_formats) for i in rgb)): - raise ValueError( - f"Input is in incorrect format: {[type(i) for i in rgb]}. Currently, we only support {', '.join(supported_formats)}" - ) - - if isinstance(rgb[0], PIL.Image.Image): - if self.config.do_convert_rgb: - raise Exception("This is not yet supported") - # rgb = [self.convert_to_rgb(i) for i in rgb] - # depth = [self.convert_to_depth(i) for i in depth] #TODO define convert_to_depth - if self.config.do_resize or target_res: - height, width = self.get_default_height_width(rgb[0], height, width) if not target_res else target_res - rgb = [self.resize(i, height, width) for i in rgb] - depth = [self.resize(i, height, width) for i in depth] - rgb = self.pil_to_numpy(rgb) # to np - rgb = self.numpy_to_pt(rgb) # to pt - - depth = self.depth_pil_to_numpy(depth) # to np - depth = self.numpy_to_pt(depth) # to pt - - elif isinstance(rgb[0], np.ndarray): - rgb = np.concatenate(rgb, axis=0) if rgb[0].ndim == 4 else np.stack(rgb, axis=0) - rgb = self.numpy_to_pt(rgb) - height, width = self.get_default_height_width(rgb, height, width) - if self.config.do_resize: - rgb = self.resize(rgb, height, width) - - depth = np.concatenate(depth, axis=0) if rgb[0].ndim == 4 else np.stack(depth, axis=0) - depth = self.numpy_to_pt(depth) - height, width = self.get_default_height_width(depth, height, width) - if self.config.do_resize: - depth = self.resize(depth, height, width) - - elif isinstance(rgb[0], torch.Tensor): - raise Exception("This is not yet supported") - # rgb = torch.cat(rgb, axis=0) if rgb[0].ndim == 4 else torch.stack(rgb, axis=0) - - # if self.config.do_convert_grayscale and rgb.ndim == 3: - # rgb = rgb.unsqueeze(1) - - # channel = rgb.shape[1] - - # height, width = self.get_default_height_width(rgb, height, width) - # if self.config.do_resize: - # rgb = self.resize(rgb, height, width) - - # depth = torch.cat(depth, axis=0) if depth[0].ndim == 4 else torch.stack(depth, axis=0) - - # if self.config.do_convert_grayscale and depth.ndim == 3: - # depth = depth.unsqueeze(1) - - # channel = depth.shape[1] - # # don't need any preprocess if the image is latents - # if depth == 4: - # return rgb, depth - - # height, width = self.get_default_height_width(depth, height, width) - # if self.config.do_resize: - # depth = self.resize(depth, height, width) - # expected range [0,1], normalize to [-1,1] - do_normalize = self.config.do_normalize - if rgb.min() < 0 and do_normalize: - warnings.warn( - "Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] " - f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{rgb.min()},{rgb.max()}]", - FutureWarning, - ) - do_normalize = False - - if do_normalize: - rgb = self.normalize(rgb) - depth = self.normalize(depth) - - if self.config.do_binarize: - rgb = self.binarize(rgb) - depth = self.binarize(depth) - - return rgb, depth - - -class IPAdapterMaskProcessor(VaeImageProcessor): - """ - Image processor for IP Adapter image masks. - - Args: - do_resize (`bool`, *optional*, defaults to `True`): - Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. - vae_scale_factor (`int`, *optional*, defaults to `8`): - VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor. - resample (`str`, *optional*, defaults to `lanczos`): - Resampling filter to use when resizing the image. - do_normalize (`bool`, *optional*, defaults to `False`): - Whether to normalize the image to [-1,1]. - do_binarize (`bool`, *optional*, defaults to `True`): - Whether to binarize the image to 0/1. - do_convert_grayscale (`bool`, *optional*, defaults to be `True`): - Whether to convert the images to grayscale format. - - """ - - config_name = CONFIG_NAME - - @register_to_config - def __init__( - self, - do_resize: bool = True, - vae_scale_factor: int = 8, - resample: str = "lanczos", - do_normalize: bool = False, - do_binarize: bool = True, - do_convert_grayscale: bool = True, - ): - super().__init__( - do_resize=do_resize, - vae_scale_factor=vae_scale_factor, - resample=resample, - do_normalize=do_normalize, - do_binarize=do_binarize, - do_convert_grayscale=do_convert_grayscale, - ) - - @staticmethod - def downsample(mask: torch.Tensor, batch_size: int, num_queries: int, value_embed_dim: int): - """ - Downsamples the provided mask tensor to match the expected dimensions for scaled dot-product attention. If the - aspect ratio of the mask does not match the aspect ratio of the output image, a warning is issued. - - Args: - mask (`torch.Tensor`): - The input mask tensor generated with `IPAdapterMaskProcessor.preprocess()`. - batch_size (`int`): - The batch size. - num_queries (`int`): - The number of queries. - value_embed_dim (`int`): - The dimensionality of the value embeddings. - - Returns: - `torch.Tensor`: - The downsampled mask tensor. - - """ - o_h = mask.shape[1] - o_w = mask.shape[2] - ratio = o_w / o_h - mask_h = int(math.sqrt(num_queries / ratio)) - mask_h = int(mask_h) + int((num_queries % int(mask_h)) != 0) - mask_w = num_queries // mask_h - - mask_downsample = F.interpolate(mask.unsqueeze(0), size=(mask_h, mask_w), mode="bicubic").squeeze(0) - - # Repeat batch_size times - if mask_downsample.shape[0] < batch_size: - mask_downsample = mask_downsample.repeat(batch_size, 1, 1) - - mask_downsample = mask_downsample.view(mask_downsample.shape[0], -1) - - downsampled_area = mask_h * mask_w - # If the output image and the mask do not have the same aspect ratio, tensor shapes will not match - # Pad tensor if downsampled_mask.shape[1] is smaller than num_queries - if downsampled_area < num_queries: - warnings.warn( - "The aspect ratio of the mask does not match the aspect ratio of the output image. " - "Please update your masks or adjust the output size for optimal performance.", - UserWarning, - ) - mask_downsample = F.pad(mask_downsample, (0, num_queries - mask_downsample.shape[1]), value=0.0) - # Discard last embeddings if downsampled_mask.shape[1] is bigger than num_queries - if downsampled_area > num_queries: - warnings.warn( - "The aspect ratio of the mask does not match the aspect ratio of the output image. " - "Please update your masks or adjust the output size for optimal performance.", - UserWarning, - ) - mask_downsample = mask_downsample[:, :num_queries] - - # Repeat last dimension to match SDPA output shape - mask_downsample = mask_downsample.view(mask_downsample.shape[0], mask_downsample.shape[1], 1).repeat( - 1, 1, value_embed_dim - ) - - return mask_downsample - - -class PixArtImageProcessor(VaeImageProcessor): - """ - Image processor for PixArt image resize and crop. - - Args: - do_resize (`bool`, *optional*, defaults to `True`): - Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. Can accept - `height` and `width` arguments from [`image_processor.VaeImageProcessor.preprocess`] method. - vae_scale_factor (`int`, *optional*, defaults to `8`): - VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor. - resample (`str`, *optional*, defaults to `lanczos`): - Resampling filter to use when resizing the image. - do_normalize (`bool`, *optional*, defaults to `True`): - Whether to normalize the image to [-1,1]. - do_binarize (`bool`, *optional*, defaults to `False`): - Whether to binarize the image to 0/1. - do_convert_rgb (`bool`, *optional*, defaults to be `False`): - 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. - """ - - @register_to_config - def __init__( - self, - do_resize: bool = True, - vae_scale_factor: int = 8, - resample: str = "lanczos", - do_normalize: bool = True, - do_binarize: bool = False, - do_convert_grayscale: bool = False, - ): - super().__init__( - do_resize=do_resize, - vae_scale_factor=vae_scale_factor, - resample=resample, - do_normalize=do_normalize, - do_binarize=do_binarize, - do_convert_grayscale=do_convert_grayscale, - ) - - @staticmethod - 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. - - Args: - height (`int`): The height of the image. - width (`int`): The width of the image. - 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. - """ - ar = float(height / width) - closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar)) - default_hw = ratios[closest_ratio] - return int(default_hw[0]), int(default_hw[1]) - - @staticmethod - def resize_and_crop_tensor(samples: torch.Tensor, new_width: int, new_height: int) -> torch.Tensor: - r""" - Resizes and crops a tensor of images to the specified dimensions. - - Args: - samples (`torch.Tensor`): - A tensor of shape (N, C, H, W) where N is the batch size, C is the number of channels, H is the height, - and W is the width. - new_width (`int`): The desired width of the output images. - new_height (`int`): The desired height of the output images. - - Returns: - `torch.Tensor`: A tensor containing the resized and cropped images. - """ - orig_height, orig_width = samples.shape[2], samples.shape[3] - - # Check if resizing is needed - if orig_height != new_height or orig_width != new_width: - ratio = max(new_height / orig_height, new_width / orig_width) - resized_width = int(orig_width * ratio) - resized_height = int(orig_height * ratio) - - # Resize - samples = F.interpolate( - samples, size=(resized_height, resized_width), mode="bilinear", align_corners=False - ) - - # Center Crop - start_x = (resized_width - new_width) // 2 - end_x = start_x + new_width - start_y = (resized_height - new_height) // 2 - end_y = start_y + new_height - samples = samples[:, :, start_y:end_y, start_x:end_x] - - return samples diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__init__.py b/venv/lib/python3.10/site-packages/diffusers/loaders/__init__.py deleted file mode 100644 index 84c6d9f32c66793cded3b75989e0c55cea99b681..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/__init__.py +++ /dev/null @@ -1,131 +0,0 @@ -from typing import TYPE_CHECKING - -from ..utils import DIFFUSERS_SLOW_IMPORT, _LazyModule, deprecate -from ..utils.import_utils import is_peft_available, is_torch_available, is_transformers_available - - -def text_encoder_lora_state_dict(text_encoder): - deprecate( - "text_encoder_load_state_dict in `models`", - "0.27.0", - "`text_encoder_lora_state_dict` is deprecated and will be removed in 0.27.0. Make sure to retrieve the weights using `get_peft_model`. See https://huggingface.co/docs/peft/v0.6.2/en/quicktour#peftmodel for more information.", - ) - state_dict = {} - - for name, module in text_encoder_attn_modules(text_encoder): - for k, v in module.q_proj.lora_linear_layer.state_dict().items(): - state_dict[f"{name}.q_proj.lora_linear_layer.{k}"] = v - - for k, v in module.k_proj.lora_linear_layer.state_dict().items(): - state_dict[f"{name}.k_proj.lora_linear_layer.{k}"] = v - - for k, v in module.v_proj.lora_linear_layer.state_dict().items(): - state_dict[f"{name}.v_proj.lora_linear_layer.{k}"] = v - - for k, v in module.out_proj.lora_linear_layer.state_dict().items(): - state_dict[f"{name}.out_proj.lora_linear_layer.{k}"] = v - - return state_dict - - -if is_transformers_available(): - - def text_encoder_attn_modules(text_encoder): - deprecate( - "text_encoder_attn_modules in `models`", - "0.27.0", - "`text_encoder_lora_state_dict` is deprecated and will be removed in 0.27.0. Make sure to retrieve the weights using `get_peft_model`. See https://huggingface.co/docs/peft/v0.6.2/en/quicktour#peftmodel for more information.", - ) - from transformers import CLIPTextModel, CLIPTextModelWithProjection - - attn_modules = [] - - if isinstance(text_encoder, (CLIPTextModel, CLIPTextModelWithProjection)): - for i, layer in enumerate(text_encoder.text_model.encoder.layers): - name = f"text_model.encoder.layers.{i}.self_attn" - mod = layer.self_attn - attn_modules.append((name, mod)) - else: - raise ValueError(f"do not know how to get attention modules for: {text_encoder.__class__.__name__}") - - return attn_modules - - -_import_structure = {} - -if is_torch_available(): - _import_structure["single_file_model"] = ["FromOriginalModelMixin"] - _import_structure["transformer_flux"] = ["FluxTransformer2DLoadersMixin"] - _import_structure["transformer_sd3"] = ["SD3Transformer2DLoadersMixin"] - _import_structure["unet"] = ["UNet2DConditionLoadersMixin"] - _import_structure["utils"] = ["AttnProcsLayers"] - if is_transformers_available(): - _import_structure["single_file"] = ["FromSingleFileMixin"] - _import_structure["lora_pipeline"] = [ - "AmusedLoraLoaderMixin", - "StableDiffusionLoraLoaderMixin", - "SD3LoraLoaderMixin", - "AuraFlowLoraLoaderMixin", - "StableDiffusionXLLoraLoaderMixin", - "LTXVideoLoraLoaderMixin", - "LoraLoaderMixin", - "FluxLoraLoaderMixin", - "CogVideoXLoraLoaderMixin", - "CogView4LoraLoaderMixin", - "Mochi1LoraLoaderMixin", - "HunyuanVideoLoraLoaderMixin", - "SanaLoraLoaderMixin", - "Lumina2LoraLoaderMixin", - "WanLoraLoaderMixin", - "HiDreamImageLoraLoaderMixin", - ] - _import_structure["textual_inversion"] = ["TextualInversionLoaderMixin"] - _import_structure["ip_adapter"] = [ - "IPAdapterMixin", - "FluxIPAdapterMixin", - "SD3IPAdapterMixin", - ] - -_import_structure["peft"] = ["PeftAdapterMixin"] - - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - if is_torch_available(): - from .single_file_model import FromOriginalModelMixin - from .transformer_flux import FluxTransformer2DLoadersMixin - from .transformer_sd3 import SD3Transformer2DLoadersMixin - from .unet import UNet2DConditionLoadersMixin - from .utils import AttnProcsLayers - - if is_transformers_available(): - from .ip_adapter import ( - FluxIPAdapterMixin, - IPAdapterMixin, - SD3IPAdapterMixin, - ) - from .lora_pipeline import ( - AmusedLoraLoaderMixin, - AuraFlowLoraLoaderMixin, - CogVideoXLoraLoaderMixin, - CogView4LoraLoaderMixin, - FluxLoraLoaderMixin, - HiDreamImageLoraLoaderMixin, - HunyuanVideoLoraLoaderMixin, - LoraLoaderMixin, - LTXVideoLoraLoaderMixin, - Lumina2LoraLoaderMixin, - Mochi1LoraLoaderMixin, - SanaLoraLoaderMixin, - SD3LoraLoaderMixin, - StableDiffusionLoraLoaderMixin, - StableDiffusionXLLoraLoaderMixin, - WanLoraLoaderMixin, - ) - from .single_file import FromSingleFileMixin - from .textual_inversion import TextualInversionLoaderMixin - - from .peft import PeftAdapterMixin -else: - import sys - - sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 782403800b0be493d0a0136ebc558ec3af800fbd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/ip_adapter.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/ip_adapter.cpython-310.pyc deleted file mode 100644 index 45c7118f80f2b3440c3fc69460ecfd7c901e40eb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/ip_adapter.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/lora_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/lora_base.cpython-310.pyc deleted file mode 100644 index 4a075984643587bd00c3854a2d576fd792821233..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/lora_base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/lora_conversion_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/lora_conversion_utils.cpython-310.pyc deleted file mode 100644 index 3768cb6d0860ee2c96828d41eeb7fac4c852444f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/lora_conversion_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/lora_pipeline.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/lora_pipeline.cpython-310.pyc deleted file mode 100644 index ea97516bf04c3a6f924efd484135a4054953305f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/lora_pipeline.cpython-310.pyc +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ac4dd539fecb30cfff2282e40e6cc98644817413166da0b767424abadf22d5cb -size 113984 diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/peft.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/peft.cpython-310.pyc deleted file mode 100644 index 652f02f47b822f489a7a0c219e040955e2db8d5c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/peft.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/single_file.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/single_file.cpython-310.pyc deleted file mode 100644 index 13a87c8a3ad8b9ddbea2b3ba2a57c0d68626e542..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/single_file.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/single_file_model.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/single_file_model.cpython-310.pyc deleted file mode 100644 index 1cd06283357acba6b1e7fed15a2a0fb0273aa176..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/single_file_model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/single_file_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/single_file_utils.cpython-310.pyc deleted file mode 100644 index f31bad6bf1a6f8a9e7b65c8d85c428bb6da8be15..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/single_file_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/textual_inversion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/textual_inversion.cpython-310.pyc deleted file mode 100644 index 1613ee78af3f8ffa462e471d1ec675bab49c7ba7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/textual_inversion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/transformer_flux.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/transformer_flux.cpython-310.pyc deleted file mode 100644 index 59843194ed9a8619aeb37a7a52842492d734cd6e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/transformer_flux.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/transformer_sd3.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/transformer_sd3.cpython-310.pyc deleted file mode 100644 index 58543fee85edf0d1d2020030015d4d01841717a3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/transformer_sd3.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/unet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/unet.cpython-310.pyc deleted file mode 100644 index a62b41e64edb1975201b969b6815ae9980eb3273..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/unet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/unet_loader_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/unet_loader_utils.cpython-310.pyc deleted file mode 100644 index 4f556ca56837738c32194d48fd1bd47ba509eb17..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/unet_loader_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index 0b3c75fa46c4e46baf9ba3459a149a98ceb2259c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/loaders/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/ip_adapter.py b/venv/lib/python3.10/site-packages/diffusers/loaders/ip_adapter.py deleted file mode 100644 index 521cb3b6fddffd30c8bb5c4fb40296d8235fb8da..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/ip_adapter.py +++ /dev/null @@ -1,870 +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 pathlib import Path -from typing import Dict, List, Optional, Union - -import torch -import torch.nn.functional as F -from huggingface_hub.utils import validate_hf_hub_args -from safetensors import safe_open - -from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, load_state_dict -from ..utils import ( - USE_PEFT_BACKEND, - _get_detailed_type, - _get_model_file, - _is_valid_type, - is_accelerate_available, - is_torch_version, - is_transformers_available, - logging, -) -from .unet_loader_utils import _maybe_expand_lora_scales - - -if is_transformers_available(): - from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection, SiglipImageProcessor, SiglipVisionModel - -from ..models.attention_processor import ( - AttnProcessor, - AttnProcessor2_0, - FluxAttnProcessor2_0, - FluxIPAdapterJointAttnProcessor2_0, - IPAdapterAttnProcessor, - IPAdapterAttnProcessor2_0, - IPAdapterXFormersAttnProcessor, - JointAttnProcessor2_0, - SD3IPAdapterJointAttnProcessor2_0, -) - - -logger = logging.get_logger(__name__) - - -class IPAdapterMixin: - """Mixin for handling IP Adapters.""" - - @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", - **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]`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - 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]`): - 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]`): - 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`): - The subfolder location of the image encoder within a larger model repository on the Hub or locally. - Pass `None` to not load the image encoder. If the image encoder is located in a folder inside - `subfolder`, you only need to pass the name of the folder that contains image encoder weights, e.g. - `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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): - Speed up model loading only loading the pretrained weights and not initializing the weights. This also - tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. - Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this - argument to `True` will raise an error. - """ - - # handle the list inputs for multiple IP Adapters - if not isinstance(weight_name, list): - weight_name = [weight_name] - - if not isinstance(pretrained_model_name_or_path_or_dict, list): - pretrained_model_name_or_path_or_dict = [pretrained_model_name_or_path_or_dict] - if len(pretrained_model_name_or_path_or_dict) == 1: - pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict * len(weight_name) - - if not isinstance(subfolder, list): - subfolder = [subfolder] - if len(subfolder) == 1: - subfolder = subfolder * len(weight_name) - - if len(weight_name) != len(pretrained_model_name_or_path_or_dict): - raise ValueError("`weight_name` and `pretrained_model_name_or_path_or_dict` must have the same length.") - - if len(weight_name) != len(subfolder): - raise ValueError("`weight_name` and `subfolder` must have the same length.") - - # Load the main state dict first. - 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) - low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) - - if low_cpu_mem_usage and not is_accelerate_available(): - low_cpu_mem_usage = False - logger.warning( - "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the" - " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install" - " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip" - " install accelerate\n```\n." - ) - - if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"): - raise NotImplementedError( - "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set" - " `low_cpu_mem_usage=False`." - ) - - user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"} - state_dicts = [] - for pretrained_model_name_or_path_or_dict, weight_name, subfolder in zip( - pretrained_model_name_or_path_or_dict, weight_name, subfolder - ): - if not isinstance(pretrained_model_name_or_path_or_dict, dict): - model_file = _get_model_file( - pretrained_model_name_or_path_or_dict, - weights_name=weight_name, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - ) - if weight_name.endswith(".safetensors"): - state_dict = {"image_proj": {}, "ip_adapter": {}} - with safe_open(model_file, framework="pt", device="cpu") as f: - for key in f.keys(): - if key.startswith("image_proj."): - state_dict["image_proj"][key.replace("image_proj.", "")] = f.get_tensor(key) - elif key.startswith("ip_adapter."): - state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = f.get_tensor(key) - else: - state_dict = load_state_dict(model_file) - else: - state_dict = pretrained_model_name_or_path_or_dict - - keys = list(state_dict.keys()) - if "image_proj" not in keys and "ip_adapter" not in keys: - raise ValueError("Required keys are (`image_proj` and `ip_adapter`) missing from the state dict.") - - state_dicts.append(state_dict) - - # load CLIP image encoder here if it has not been registered to the pipeline yet - if hasattr(self, "image_encoder") and getattr(self, "image_encoder", None) is None: - if image_encoder_folder is not None: - if not isinstance(pretrained_model_name_or_path_or_dict, dict): - logger.info(f"loading image_encoder from {pretrained_model_name_or_path_or_dict}") - if image_encoder_folder.count("/") == 0: - image_encoder_subfolder = Path(subfolder, image_encoder_folder).as_posix() - else: - image_encoder_subfolder = Path(image_encoder_folder).as_posix() - - image_encoder = CLIPVisionModelWithProjection.from_pretrained( - pretrained_model_name_or_path_or_dict, - subfolder=image_encoder_subfolder, - low_cpu_mem_usage=low_cpu_mem_usage, - cache_dir=cache_dir, - local_files_only=local_files_only, - torch_dtype=self.dtype, - ).to(self.device) - self.register_modules(image_encoder=image_encoder) - else: - raise ValueError( - "`image_encoder` cannot be loaded because `pretrained_model_name_or_path_or_dict` is a state dict." - ) - else: - logger.warning( - "image_encoder is not loaded since `image_encoder_folder=None` passed. You will not be able to use `ip_adapter_image` when calling the pipeline with IP-Adapter." - "Use `ip_adapter_image_embeds` to pass pre-generated image embedding instead." - ) - - # create feature extractor if it has not been registered to the pipeline yet - if hasattr(self, "feature_extractor") and getattr(self, "feature_extractor", None) is None: - # FaceID IP adapters don't need the image encoder so it's not present, in this case we default to 224 - default_clip_size = 224 - clip_image_size = ( - self.image_encoder.config.image_size if self.image_encoder is not None else default_clip_size - ) - feature_extractor = CLIPImageProcessor(size=clip_image_size, crop_size=clip_image_size) - self.register_modules(feature_extractor=feature_extractor) - - # load ip-adapter into unet - unet = getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet - unet._load_ip_adapter_weights(state_dicts, low_cpu_mem_usage=low_cpu_mem_usage) - - extra_loras = unet._load_ip_adapter_loras(state_dicts) - if extra_loras != {}: - if not USE_PEFT_BACKEND: - logger.warning("PEFT backend is required to load these weights.") - else: - # apply the IP Adapter Face ID LoRA weights - peft_config = getattr(unet, "peft_config", {}) - for k, lora in extra_loras.items(): - if f"faceid_{k}" not in peft_config: - self.load_lora_weights(lora, adapter_name=f"faceid_{k}") - self.set_adapters([f"faceid_{k}"], adapter_weights=[1.0]) - - def set_ip_adapter_scale(self, scale): - """ - 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 dictionary. - - Example: - - ```py - # To use original IP-Adapter - scale = 1.0 - pipeline.set_ip_adapter_scale(scale) - - # To use style block only - scale = { - "up": {"block_0": [0.0, 1.0, 0.0]}, - } - pipeline.set_ip_adapter_scale(scale) - - # To use style+layout blocks - scale = { - "down": {"block_2": [0.0, 1.0]}, - "up": {"block_0": [0.0, 1.0, 0.0]}, - } - pipeline.set_ip_adapter_scale(scale) - - # To use style and layout from 2 reference images - scales = [{"down": {"block_2": [0.0, 1.0]}}, {"up": {"block_0": [0.0, 1.0, 0.0]}}] - pipeline.set_ip_adapter_scale(scales) - ``` - """ - unet = getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet - if not isinstance(scale, list): - scale = [scale] - scale_configs = _maybe_expand_lora_scales(unet, scale, default_scale=0.0) - - for attn_name, attn_processor in unet.attn_processors.items(): - if isinstance( - attn_processor, (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0, IPAdapterXFormersAttnProcessor) - ): - if len(scale_configs) != len(attn_processor.scale): - raise ValueError( - f"Cannot assign {len(scale_configs)} scale_configs to {len(attn_processor.scale)} IP-Adapter." - ) - elif len(scale_configs) == 1: - scale_configs = scale_configs * len(attn_processor.scale) - for i, scale_config in enumerate(scale_configs): - if isinstance(scale_config, dict): - for k, s in scale_config.items(): - if attn_name.startswith(k): - attn_processor.scale[i] = s - else: - attn_processor.scale[i] = scale_config - - def unload_ip_adapter(self): - """ - Unloads the IP Adapter weights - - Examples: - - ```python - >>> # Assuming `pipeline` is already loaded with the IP Adapter weights. - >>> pipeline.unload_ip_adapter() - >>> ... - ``` - """ - # remove CLIP image encoder - if hasattr(self, "image_encoder") and getattr(self, "image_encoder", None) is not None: - self.image_encoder = None - self.register_to_config(image_encoder=[None, None]) - - # remove feature extractor only when safety_checker is None as safety_checker uses - # the feature_extractor later - if not hasattr(self, "safety_checker"): - if hasattr(self, "feature_extractor") and getattr(self, "feature_extractor", None) is not None: - self.feature_extractor = None - self.register_to_config(feature_extractor=[None, None]) - - # remove hidden encoder - self.unet.encoder_hid_proj = None - self.unet.config.encoder_hid_dim_type = None - - # Kolors: restore `encoder_hid_proj` with `text_encoder_hid_proj` - if hasattr(self.unet, "text_encoder_hid_proj") and self.unet.text_encoder_hid_proj is not None: - self.unet.encoder_hid_proj = self.unet.text_encoder_hid_proj - self.unet.text_encoder_hid_proj = None - self.unet.config.encoder_hid_dim_type = "text_proj" - - # restore original Unet attention processors layers - attn_procs = {} - for name, value in self.unet.attn_processors.items(): - attn_processor_class = ( - AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnProcessor() - ) - attn_procs[name] = ( - attn_processor_class - if isinstance( - value, (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0, IPAdapterXFormersAttnProcessor) - ) - else value.__class__() - ) - self.unet.set_attn_processor(attn_procs) - - -class FluxIPAdapterMixin: - """Mixin for handling Flux IP Adapters.""" - - @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] = "", - 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]`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - 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]`): - 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]`): - 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`): - Can be either: - - - A string, the *model id* (for example `openai/clip-vit-large-patch14`) of a pretrained model - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): - Speed up model loading only loading the pretrained weights and not initializing the weights. This also - tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. - Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this - argument to `True` will raise an error. - """ - - # handle the list inputs for multiple IP Adapters - if not isinstance(weight_name, list): - weight_name = [weight_name] - - if not isinstance(pretrained_model_name_or_path_or_dict, list): - pretrained_model_name_or_path_or_dict = [pretrained_model_name_or_path_or_dict] - if len(pretrained_model_name_or_path_or_dict) == 1: - pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict * len(weight_name) - - if not isinstance(subfolder, list): - subfolder = [subfolder] - if len(subfolder) == 1: - subfolder = subfolder * len(weight_name) - - if len(weight_name) != len(pretrained_model_name_or_path_or_dict): - raise ValueError("`weight_name` and `pretrained_model_name_or_path_or_dict` must have the same length.") - - if len(weight_name) != len(subfolder): - raise ValueError("`weight_name` and `subfolder` must have the same length.") - - # Load the main state dict first. - 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) - low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) - - if low_cpu_mem_usage and not is_accelerate_available(): - low_cpu_mem_usage = False - logger.warning( - "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the" - " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install" - " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip" - " install accelerate\n```\n." - ) - - if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"): - raise NotImplementedError( - "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set" - " `low_cpu_mem_usage=False`." - ) - - user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"} - state_dicts = [] - for pretrained_model_name_or_path_or_dict, weight_name, subfolder in zip( - pretrained_model_name_or_path_or_dict, weight_name, subfolder - ): - if not isinstance(pretrained_model_name_or_path_or_dict, dict): - model_file = _get_model_file( - pretrained_model_name_or_path_or_dict, - weights_name=weight_name, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - ) - if weight_name.endswith(".safetensors"): - state_dict = {"image_proj": {}, "ip_adapter": {}} - with safe_open(model_file, framework="pt", device="cpu") as f: - image_proj_keys = ["ip_adapter_proj_model.", "image_proj."] - ip_adapter_keys = ["double_blocks.", "ip_adapter."] - for key in f.keys(): - if any(key.startswith(prefix) for prefix in image_proj_keys): - diffusers_name = ".".join(key.split(".")[1:]) - state_dict["image_proj"][diffusers_name] = f.get_tensor(key) - elif any(key.startswith(prefix) for prefix in ip_adapter_keys): - diffusers_name = ( - ".".join(key.split(".")[1:]) - .replace("ip_adapter_double_stream_k_proj", "to_k_ip") - .replace("ip_adapter_double_stream_v_proj", "to_v_ip") - .replace("processor.", "") - ) - state_dict["ip_adapter"][diffusers_name] = f.get_tensor(key) - else: - state_dict = load_state_dict(model_file) - else: - state_dict = pretrained_model_name_or_path_or_dict - - keys = list(state_dict.keys()) - if keys != ["image_proj", "ip_adapter"]: - raise ValueError("Required keys are (`image_proj` and `ip_adapter`) missing from the state dict.") - - state_dicts.append(state_dict) - - # load CLIP image encoder here if it has not been registered to the pipeline yet - if hasattr(self, "image_encoder") and getattr(self, "image_encoder", None) is None: - if image_encoder_pretrained_model_name_or_path is not None: - if not isinstance(pretrained_model_name_or_path_or_dict, dict): - logger.info(f"loading image_encoder from {image_encoder_pretrained_model_name_or_path}") - image_encoder = ( - CLIPVisionModelWithProjection.from_pretrained( - image_encoder_pretrained_model_name_or_path, - subfolder=image_encoder_subfolder, - low_cpu_mem_usage=low_cpu_mem_usage, - cache_dir=cache_dir, - local_files_only=local_files_only, - torch_dtype=image_encoder_dtype, - ) - .to(self.device) - .eval() - ) - self.register_modules(image_encoder=image_encoder) - else: - raise ValueError( - "`image_encoder` cannot be loaded because `pretrained_model_name_or_path_or_dict` is a state dict." - ) - else: - logger.warning( - "image_encoder is not loaded since `image_encoder_folder=None` passed. You will not be able to use `ip_adapter_image` when calling the pipeline with IP-Adapter." - "Use `ip_adapter_image_embeds` to pass pre-generated image embedding instead." - ) - - # create feature extractor if it has not been registered to the pipeline yet - if hasattr(self, "feature_extractor") and getattr(self, "feature_extractor", None) is None: - # FaceID IP adapters don't need the image encoder so it's not present, in this case we default to 224 - default_clip_size = 224 - clip_image_size = ( - self.image_encoder.config.image_size if self.image_encoder is not None else default_clip_size - ) - feature_extractor = CLIPImageProcessor(size=clip_image_size, crop_size=clip_image_size) - self.register_modules(feature_extractor=feature_extractor) - - # 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]]]): - """ - 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 - number of IP adapters and each must match the number of blocks. - - Example: - - ```py - # To use original IP-Adapter - scale = 1.0 - pipeline.set_ip_adapter_scale(scale) - - - def LinearStrengthModel(start, finish, size): - return [(start + (finish - start) * (i / (size - 1))) for i in range(size)] - - - ip_strengths = LinearStrengthModel(0.3, 0.92, 19) - pipeline.set_ip_adapter_scale(ip_strengths) - ``` - """ - - scale_type = Union[int, float] - num_ip_adapters = self.transformer.encoder_hid_proj.num_ip_adapters - num_layers = self.transformer.config.num_layers - - # Single value for all layers of all IP-Adapters - if isinstance(scale, scale_type): - scale = [scale for _ in range(num_ip_adapters)] - # List of per-layer scales for a single IP-Adapter - elif _is_valid_type(scale, List[scale_type]) and num_ip_adapters == 1: - scale = [scale] - # Invalid scale type - elif not _is_valid_type(scale, List[Union[scale_type, List[scale_type]]]): - raise TypeError(f"Unexpected type {_get_detailed_type(scale)} for scale.") - - if len(scale) != num_ip_adapters: - raise ValueError(f"Cannot assign {len(scale)} scales to {num_ip_adapters} IP-Adapters.") - - if any(len(s) != num_layers for s in scale if isinstance(s, list)): - invalid_scale_sizes = {len(s) for s in scale if isinstance(s, list)} - {num_layers} - raise ValueError( - f"Expected list of {num_layers} scales, got {', '.join(str(x) for x in invalid_scale_sizes)}." - ) - - # Scalars are transformed to lists with length num_layers - scale_configs = [[s] * num_layers if isinstance(s, scale_type) else s for s in scale] - - # Set scales. zip over scale_configs prevents going into single transformer layers - for attn_processor, *scale in zip(self.transformer.attn_processors.values(), *scale_configs): - attn_processor.scale = scale - - def unload_ip_adapter(self): - """ - Unloads the IP Adapter weights - - Examples: - - ```python - >>> # Assuming `pipeline` is already loaded with the IP Adapter weights. - >>> pipeline.unload_ip_adapter() - >>> ... - ``` - """ - # remove CLIP image encoder - if hasattr(self, "image_encoder") and getattr(self, "image_encoder", None) is not None: - self.image_encoder = None - self.register_to_config(image_encoder=[None, None]) - - # remove feature extractor only when safety_checker is None as safety_checker uses - # the feature_extractor later - if not hasattr(self, "safety_checker"): - if hasattr(self, "feature_extractor") and getattr(self, "feature_extractor", None) is not None: - self.feature_extractor = None - self.register_to_config(feature_extractor=[None, None]) - - # remove hidden encoder - self.transformer.encoder_hid_proj = None - self.transformer.config.encoder_hid_dim_type = None - - # restore original Transformer attention processors layers - attn_procs = {} - for name, value in self.transformer.attn_processors.items(): - attn_processor_class = FluxAttnProcessor2_0() - attn_procs[name] = ( - attn_processor_class if isinstance(value, (FluxIPAdapterJointAttnProcessor2_0)) else value.__class__() - ) - self.transformer.set_attn_processor(attn_procs) - - -class SD3IPAdapterMixin: - """Mixin for handling StableDiffusion 3 IP Adapters.""" - - @property - def is_ip_adapter_active(self) -> bool: - """Checks if IP-Adapter is loaded and scale > 0. - - IP-Adapter scale controls the influence of the image prompt versus text prompt. When this value is set to 0, - the image context is irrelevant. - - Returns: - `bool`: True when IP-Adapter is loaded and any layer has scale > 0. - """ - scales = [ - attn_proc.scale - for attn_proc in self.transformer.attn_processors.values() - if isinstance(attn_proc, SD3IPAdapterJointAttnProcessor2_0) - ] - - return len(scales) > 0 and any(scale > 0 for scale in scales) - - @validate_hf_hub_args - def load_ip_adapter( - self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - weight_name: str = "ip-adapter.safetensors", - subfolder: Optional[str] = None, - image_encoder_folder: Optional[str] = "image_encoder", - **kwargs, - ) -> None: - """ - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - weight_name (`str`, defaults to "ip-adapter.safetensors"): - The name of the weight file to load. If a list is passed, it should have the same length as - `subfolder`. - subfolder (`str`, *optional*): - 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`. - image_encoder_folder (`str`, *optional*, defaults to `image_encoder`): - The subfolder location of the image encoder within a larger model repository on the Hub or locally. - Pass `None` to not load the image encoder. If the image encoder is located in a folder inside - `subfolder`, you only need to pass the name of the folder that contains image encoder weights, e.g. - `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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): - Speed up model loading only loading the pretrained weights and not initializing the weights. This also - tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. - Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this - argument to `True` will raise an error. - """ - # Load the main state dict first - 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) - low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) - - if low_cpu_mem_usage and not is_accelerate_available(): - low_cpu_mem_usage = False - logger.warning( - "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the" - " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install" - " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip" - " install accelerate\n```\n." - ) - - if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"): - raise NotImplementedError( - "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set" - " `low_cpu_mem_usage=False`." - ) - - user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"} - - if not isinstance(pretrained_model_name_or_path_or_dict, dict): - model_file = _get_model_file( - pretrained_model_name_or_path_or_dict, - weights_name=weight_name, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - ) - if weight_name.endswith(".safetensors"): - state_dict = {"image_proj": {}, "ip_adapter": {}} - with safe_open(model_file, framework="pt", device="cpu") as f: - for key in f.keys(): - if key.startswith("image_proj."): - state_dict["image_proj"][key.replace("image_proj.", "")] = f.get_tensor(key) - elif key.startswith("ip_adapter."): - state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = f.get_tensor(key) - else: - state_dict = load_state_dict(model_file) - else: - state_dict = pretrained_model_name_or_path_or_dict - - keys = list(state_dict.keys()) - if "image_proj" not in keys and "ip_adapter" not in keys: - raise ValueError("Required keys are (`image_proj` and `ip_adapter`) missing from the state dict.") - - # Load image_encoder and feature_extractor here if they haven't been registered to the pipeline yet - if hasattr(self, "image_encoder") and getattr(self, "image_encoder", None) is None: - if image_encoder_folder is not None: - if not isinstance(pretrained_model_name_or_path_or_dict, dict): - logger.info(f"loading image_encoder from {pretrained_model_name_or_path_or_dict}") - if image_encoder_folder.count("/") == 0: - image_encoder_subfolder = Path(subfolder, image_encoder_folder).as_posix() - else: - image_encoder_subfolder = Path(image_encoder_folder).as_posix() - - # Commons args for loading image encoder and image processor - kwargs = { - "low_cpu_mem_usage": low_cpu_mem_usage, - "cache_dir": cache_dir, - "local_files_only": local_files_only, - } - - self.register_modules( - feature_extractor=SiglipImageProcessor.from_pretrained(image_encoder_subfolder, **kwargs), - image_encoder=SiglipVisionModel.from_pretrained( - image_encoder_subfolder, torch_dtype=self.dtype, **kwargs - ).to(self.device), - ) - else: - raise ValueError( - "`image_encoder` cannot be loaded because `pretrained_model_name_or_path_or_dict` is a state dict." - ) - else: - logger.warning( - "image_encoder is not loaded since `image_encoder_folder=None` passed. You will not be able to use `ip_adapter_image` when calling the pipeline with IP-Adapter." - "Use `ip_adapter_image_embeds` to pass pre-generated image embedding instead." - ) - - # Load IP-Adapter into transformer - self.transformer._load_ip_adapter_weights(state_dict, low_cpu_mem_usage=low_cpu_mem_usage) - - def set_ip_adapter_scale(self, scale: float) -> None: - """ - Set IP-Adapter scale, which controls image prompt conditioning. A value of 1.0 means the model is only - conditioned on the image prompt, and 0.0 only conditioned by the text prompt. Lowering this value encourages - the model to produce more diverse images, but they may not be as aligned with the image prompt. - - Example: - - ```python - >>> # Assuming `pipeline` is already loaded with the IP Adapter weights. - >>> pipeline.set_ip_adapter_scale(0.6) - >>> ... - ``` - - Args: - scale (float): - IP-Adapter scale to be set. - - """ - for attn_processor in self.transformer.attn_processors.values(): - if isinstance(attn_processor, SD3IPAdapterJointAttnProcessor2_0): - attn_processor.scale = scale - - def unload_ip_adapter(self) -> None: - """ - Unloads the IP Adapter weights. - - Example: - - ```python - >>> # Assuming `pipeline` is already loaded with the IP Adapter weights. - >>> pipeline.unload_ip_adapter() - >>> ... - ``` - """ - # Remove image encoder - if hasattr(self, "image_encoder") and getattr(self, "image_encoder", None) is not None: - self.image_encoder = None - self.register_to_config(image_encoder=None) - - # Remove feature extractor - if hasattr(self, "feature_extractor") and getattr(self, "feature_extractor", None) is not None: - self.feature_extractor = None - self.register_to_config(feature_extractor=None) - - # Remove image projection - self.transformer.image_proj = None - - # Restore original attention processors layers - attn_procs = { - name: ( - JointAttnProcessor2_0() if isinstance(value, SD3IPAdapterJointAttnProcessor2_0) else value.__class__() - ) - for name, value in self.transformer.attn_processors.items() - } - self.transformer.set_attn_processor(attn_procs) diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/lora_base.py b/venv/lib/python3.10/site-packages/diffusers/loaders/lora_base.py deleted file mode 100644 index 16f0d4836505f3b7f0ef5cca332ca76d907c9e4c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/lora_base.py +++ /dev/null @@ -1,1036 +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 copy -import inspect -import json -import os -from pathlib import Path -from typing import Callable, Dict, List, Optional, Union - -import safetensors -import torch -import torch.nn as nn -from huggingface_hub import model_info -from huggingface_hub.constants import HF_HUB_OFFLINE - -from ..models.modeling_utils import ModelMixin, load_state_dict -from ..utils import ( - USE_PEFT_BACKEND, - _get_model_file, - convert_state_dict_to_diffusers, - convert_state_dict_to_peft, - delete_adapter_layers, - deprecate, - get_adapter_name, - is_accelerate_available, - is_peft_available, - is_peft_version, - is_transformers_available, - is_transformers_version, - logging, - recurse_remove_peft_layers, - scale_lora_layers, - set_adapter_layers, - set_weights_and_activate_adapters, -) -from ..utils.peft_utils import _create_lora_config -from ..utils.state_dict_utils import _load_sft_state_dict_metadata - - -if is_transformers_available(): - from transformers import PreTrainedModel - -if is_peft_available(): - from peft.tuners.tuners_utils import BaseTunerLayer - -if is_accelerate_available(): - from accelerate.hooks import AlignDevicesHook, CpuOffload, remove_hook_from_module - -logger = logging.get_logger(__name__) - -LORA_WEIGHT_NAME = "pytorch_lora_weights.bin" -LORA_WEIGHT_NAME_SAFE = "pytorch_lora_weights.safetensors" -LORA_ADAPTER_METADATA_KEY = "lora_adapter_metadata" - - -def fuse_text_encoder_lora(text_encoder, lora_scale=1.0, safe_fusing=False, adapter_names=None): - """ - Fuses LoRAs for the text encoder. - - Args: - text_encoder (`torch.nn.Module`): - The text encoder module to set the adapter layers for. If `None`, it will try to get the `text_encoder` - attribute. - 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]` or `str`): - The names of the adapters to use. - """ - merge_kwargs = {"safe_merge": safe_fusing} - - for module in text_encoder.modules(): - if isinstance(module, BaseTunerLayer): - if lora_scale != 1.0: - module.scale_layer(lora_scale) - - # For BC with previous PEFT versions, we need to check the signature - # of the `merge` method to see if it supports the `adapter_names` argument. - supported_merge_kwargs = list(inspect.signature(module.merge).parameters) - if "adapter_names" in supported_merge_kwargs: - merge_kwargs["adapter_names"] = adapter_names - elif "adapter_names" not in supported_merge_kwargs and adapter_names is not None: - raise ValueError( - "The `adapter_names` argument is not supported with your PEFT version. " - "Please upgrade to the latest version of PEFT. `pip install -U peft`" - ) - - module.merge(**merge_kwargs) - - -def unfuse_text_encoder_lora(text_encoder): - """ - Unfuses LoRAs for the text encoder. - - Args: - text_encoder (`torch.nn.Module`): - The text encoder module to set the adapter layers for. If `None`, it will try to get the `text_encoder` - attribute. - """ - for module in text_encoder.modules(): - if isinstance(module, BaseTunerLayer): - module.unmerge() - - -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, -): - """ - Sets the adapter layers for the text encoder. - - Args: - 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*): - 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: - raise ValueError( - "The pipeline does not have a default `pipe.text_encoder` class. Please make sure to pass a `text_encoder` instead." - ) - - def process_weights(adapter_names, weights): - # Expand weights into a list, one entry per adapter - # e.g. for 2 adapters: 7 -> [7,7] ; [3, None] -> [3, None] - if not isinstance(weights, list): - weights = [weights] * len(adapter_names) - - if len(adapter_names) != len(weights): - raise ValueError( - f"Length of adapter names {len(adapter_names)} is not equal to the length of the weights {len(weights)}" - ) - - # Set None values to default of 1.0 - # e.g. [7,7] -> [7,7] ; [3, None] -> [3,1] - weights = [w if w is not None else 1.0 for w in weights] - - return weights - - adapter_names = [adapter_names] if isinstance(adapter_names, str) else adapter_names - text_encoder_weights = process_weights(adapter_names, text_encoder_weights) - set_weights_and_activate_adapters(text_encoder, adapter_names, text_encoder_weights) - - -def disable_lora_for_text_encoder(text_encoder: Optional["PreTrainedModel"] = None): - """ - Disables the LoRA layers for the text encoder. - - Args: - text_encoder (`torch.nn.Module`, *optional*): - The text encoder module to disable the LoRA layers for. If `None`, it will try to get the `text_encoder` - attribute. - """ - if text_encoder is None: - raise ValueError("Text Encoder not found.") - set_adapter_layers(text_encoder, enabled=False) - - -def enable_lora_for_text_encoder(text_encoder: Optional["PreTrainedModel"] = None): - """ - Enables the LoRA layers for the text encoder. - - Args: - text_encoder (`torch.nn.Module`, *optional*): - The text encoder module to enable the LoRA layers for. If `None`, it will try to get the `text_encoder` - attribute. - """ - if text_encoder is None: - raise ValueError("Text Encoder not found.") - set_adapter_layers(text_encoder, enabled=True) - - -def _remove_text_encoder_monkey_patch(text_encoder): - recurse_remove_peft_layers(text_encoder) - if getattr(text_encoder, "peft_config", None) is not None: - del text_encoder.peft_config - text_encoder._hf_peft_config_loaded = None - - -def _fetch_state_dict( - pretrained_model_name_or_path_or_dict, - weight_name, - use_safetensors, - local_files_only, - cache_dir, - force_download, - proxies, - token, - revision, - subfolder, - user_agent, - allow_pickle, - metadata=None, -): - model_file = None - if not isinstance(pretrained_model_name_or_path_or_dict, dict): - # Let's first try to load .safetensors weights - if (use_safetensors and weight_name is None) or ( - weight_name is not None and weight_name.endswith(".safetensors") - ): - try: - # Here we're relaxing the loading check to enable more Inference API - # friendliness where sometimes, it's not at all possible to automatically - # determine `weight_name`. - if weight_name is None: - weight_name = _best_guess_weight_name( - pretrained_model_name_or_path_or_dict, - file_extension=".safetensors", - local_files_only=local_files_only, - ) - model_file = _get_model_file( - pretrained_model_name_or_path_or_dict, - weights_name=weight_name or LORA_WEIGHT_NAME_SAFE, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - ) - state_dict = safetensors.torch.load_file(model_file, device="cpu") - metadata = _load_sft_state_dict_metadata(model_file) - - except (IOError, safetensors.SafetensorError) as e: - if not allow_pickle: - raise e - # try loading non-safetensors weights - model_file = None - metadata = None - pass - - if model_file is None: - if weight_name is None: - weight_name = _best_guess_weight_name( - pretrained_model_name_or_path_or_dict, file_extension=".bin", local_files_only=local_files_only - ) - model_file = _get_model_file( - pretrained_model_name_or_path_or_dict, - weights_name=weight_name or LORA_WEIGHT_NAME, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - ) - state_dict = load_state_dict(model_file) - metadata = None - else: - state_dict = pretrained_model_name_or_path_or_dict - - return state_dict, metadata - - -def _best_guess_weight_name( - pretrained_model_name_or_path_or_dict, file_extension=".safetensors", local_files_only=False -): - if local_files_only or HF_HUB_OFFLINE: - raise ValueError("When using the offline mode, you must specify a `weight_name`.") - - targeted_files = [] - - if os.path.isfile(pretrained_model_name_or_path_or_dict): - return - elif os.path.isdir(pretrained_model_name_or_path_or_dict): - targeted_files = [f for f in os.listdir(pretrained_model_name_or_path_or_dict) if f.endswith(file_extension)] - else: - files_in_repo = model_info(pretrained_model_name_or_path_or_dict).siblings - targeted_files = [f.rfilename for f in files_in_repo if f.rfilename.endswith(file_extension)] - if len(targeted_files) == 0: - return - - # "scheduler" does not correspond to a LoRA checkpoint. - # "optimizer" does not correspond to a LoRA checkpoint - # only top-level checkpoints are considered and not the other ones, hence "checkpoint". - unallowed_substrings = {"scheduler", "optimizer", "checkpoint"} - targeted_files = list( - filter(lambda x: all(substring not in x for substring in unallowed_substrings), targeted_files) - ) - - if any(f.endswith(LORA_WEIGHT_NAME) for f in targeted_files): - targeted_files = list(filter(lambda x: x.endswith(LORA_WEIGHT_NAME), targeted_files)) - elif any(f.endswith(LORA_WEIGHT_NAME_SAFE) for f in targeted_files): - targeted_files = list(filter(lambda x: x.endswith(LORA_WEIGHT_NAME_SAFE), targeted_files)) - - if len(targeted_files) > 1: - logger.warning( - f"Provided path contains more than one weights file in the {file_extension} format. `{targeted_files[0]}` is going to be loaded, for precise control, specify a `weight_name` in `load_lora_weights`." - ) - weight_name = targeted_files[0] - return weight_name - - -def _pack_dict_with_prefix(state_dict, prefix): - sd_with_prefix = {f"{prefix}.{key}": value for key, value in state_dict.items()} - return sd_with_prefix - - -def _load_lora_into_text_encoder( - state_dict, - network_alphas, - text_encoder, - prefix=None, - lora_scale=1.0, - text_encoder_name="text_encoder", - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, -): - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for this method.") - - if network_alphas and metadata: - raise ValueError("`network_alphas` and `metadata` cannot be specified both at the same time.") - - peft_kwargs = {} - if low_cpu_mem_usage: - if not is_peft_version(">=", "0.13.1"): - raise ValueError( - "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`." - ) - if not is_transformers_version(">", "4.45.2"): - # Note from sayakpaul: It's not in `transformers` stable yet. - # https://github.com/huggingface/transformers/pull/33725/ - raise ValueError( - "`low_cpu_mem_usage=True` is not compatible with this `transformers` version. Please update it with `pip install -U transformers`." - ) - peft_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage - - # If the serialization format is new (introduced in https://github.com/huggingface/diffusers/pull/2918), - # then the `state_dict` keys should have `unet_name` and/or `text_encoder_name` as - # their prefixes. - prefix = text_encoder_name if prefix is None else prefix - - # Safe prefix to check with. - if hotswap and any(text_encoder_name in key for key in state_dict.keys()): - raise ValueError("At the moment, hotswapping is not supported for text encoders, please pass `hotswap=False`.") - - # Load the layers corresponding to text encoder and make necessary adjustments. - if prefix is not None: - state_dict = {k.removeprefix(f"{prefix}."): v for k, v in state_dict.items() if k.startswith(f"{prefix}.")} - if metadata is not None: - metadata = {k.removeprefix(f"{prefix}."): v for k, v in metadata.items() if k.startswith(f"{prefix}.")} - - if len(state_dict) > 0: - logger.info(f"Loading {prefix}.") - rank = {} - state_dict = convert_state_dict_to_diffusers(state_dict) - - # convert state dict - state_dict = convert_state_dict_to_peft(state_dict) - - for name, _ in text_encoder.named_modules(): - if name.endswith((".q_proj", ".k_proj", ".v_proj", ".out_proj", ".fc1", ".fc2")): - rank_key = f"{name}.lora_B.weight" - if rank_key in state_dict: - rank[rank_key] = state_dict[rank_key].shape[1] - - if network_alphas is not None: - alpha_keys = [k for k in network_alphas.keys() if k.startswith(prefix) and k.split(".")[0] == prefix] - network_alphas = {k.removeprefix(f"{prefix}."): v for k, v in network_alphas.items() if k in alpha_keys} - - # create `LoraConfig` - lora_config = _create_lora_config(state_dict, network_alphas, metadata, rank, is_unet=False) - - # adapter_name - if adapter_name is None: - adapter_name = get_adapter_name(text_encoder) - - # - - if prefix is not None and not state_dict: - model_class_name = text_encoder.__class__.__name__ - logger.warning( - f"No LoRA keys associated to {model_class_name} found with the {prefix=}. " - "This is safe to ignore if LoRA state dict didn't originally have any " - f"{model_class_name} related params. You can also try specifying `prefix=None` " - "to resolve the warning. Otherwise, open an issue if you think it's unexpected: " - "https://github.com/huggingface/diffusers/issues/new" - ) - - -def _func_optionally_disable_offloading(_pipeline): - """ - Optionally removes offloading in case the pipeline has been already sequentially offloaded to CPU. - - Args: - _pipeline (`DiffusionPipeline`): - The pipeline to disable offloading for. - - Returns: - tuple: - A tuple indicating if `is_model_cpu_offload` or `is_sequential_cpu_offload` is True. - """ - is_model_cpu_offload = False - is_sequential_cpu_offload = False - - if _pipeline is not None and _pipeline.hf_device_map is None: - for _, component in _pipeline.components.items(): - if isinstance(component, nn.Module) and hasattr(component, "_hf_hook"): - if not is_model_cpu_offload: - is_model_cpu_offload = isinstance(component._hf_hook, CpuOffload) - if not is_sequential_cpu_offload: - is_sequential_cpu_offload = ( - isinstance(component._hf_hook, AlignDevicesHook) - or hasattr(component._hf_hook, "hooks") - and isinstance(component._hf_hook.hooks[0], AlignDevicesHook) - ) - - logger.info( - "Accelerate hooks detected. Since you have called `load_lora_weights()`, the previous hooks will be first removed. Then the LoRA parameters will be loaded and the hooks will be applied again." - ) - if is_sequential_cpu_offload or is_model_cpu_offload: - remove_hook_from_module(component, recurse=is_sequential_cpu_offload) - - return (is_model_cpu_offload, is_sequential_cpu_offload) - - -class LoraBaseMixin: - """Utility class for handling LoRAs.""" - - _lora_loadable_modules = [] - _merged_adapters = set() - - @property - def lora_scale(self) -> float: - """ - Returns the lora scale which can be set at run time by the pipeline. # if `_lora_scale` has not been set, - return 1. - """ - return self._lora_scale if hasattr(self, "_lora_scale") else 1.0 - - @property - def num_fused_loras(self): - """Returns the number of LoRAs that have been fused.""" - return len(self._merged_adapters) - - @property - def fused_loras(self): - """Returns names of the LoRAs that have been fused.""" - return self._merged_adapters - - def load_lora_weights(self, **kwargs): - raise NotImplementedError("`load_lora_weights()` is not implemented.") - - @classmethod - def save_lora_weights(cls, **kwargs): - raise NotImplementedError("`save_lora_weights()` not implemented.") - - @classmethod - def lora_state_dict(cls, **kwargs): - raise NotImplementedError("`lora_state_dict()` is not implemented.") - - def unload_lora_weights(self): - """ - Unloads the LoRA parameters. - - Examples: - - ```python - >>> # Assuming `pipeline` is already loaded with the LoRA parameters. - >>> pipeline.unload_lora_weights() - >>> ... - ``` - """ - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for this method.") - - for component in self._lora_loadable_modules: - model = getattr(self, component, None) - if model is not None: - if issubclass(model.__class__, ModelMixin): - model.unload_lora() - elif issubclass(model.__class__, PreTrainedModel): - _remove_text_encoder_monkey_patch(model) - - def fuse_lora( - self, - components: List[str] = [], - lora_scale: float = 1.0, - safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - if "fuse_unet" in kwargs: - depr_message = "Passing `fuse_unet` to `fuse_lora()` is deprecated and will be ignored. Please use the `components` argument and provide a list of the components whose LoRAs are to be fused. `fuse_unet` will be removed in a future version." - deprecate( - "fuse_unet", - "1.0.0", - depr_message, - ) - if "fuse_transformer" in kwargs: - depr_message = "Passing `fuse_transformer` to `fuse_lora()` is deprecated and will be ignored. Please use the `components` argument and provide a list of the components whose LoRAs are to be fused. `fuse_transformer` will be removed in a future version." - deprecate( - "fuse_transformer", - "1.0.0", - depr_message, - ) - if "fuse_text_encoder" in kwargs: - depr_message = "Passing `fuse_text_encoder` to `fuse_lora()` is deprecated and will be ignored. Please use the `components` argument and provide a list of the components whose LoRAs are to be fused. `fuse_text_encoder` will be removed in a future version." - deprecate( - "fuse_text_encoder", - "1.0.0", - depr_message, - ) - - if len(components) == 0: - raise ValueError("`components` cannot be an empty list.") - - # Need to retrieve the names as `adapter_names` can be None. So we cannot directly use it - # in `self._merged_adapters = self._merged_adapters | merged_adapter_names`. - merged_adapter_names = set() - for fuse_component in components: - if fuse_component not in self._lora_loadable_modules: - raise ValueError(f"{fuse_component} is not found in {self._lora_loadable_modules=}.") - - model = getattr(self, fuse_component, None) - if model is not None: - # check if diffusers model - if issubclass(model.__class__, ModelMixin): - model.fuse_lora(lora_scale, safe_fusing=safe_fusing, adapter_names=adapter_names) - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - merged_adapter_names.update(set(module.merged_adapters)) - # handle transformers models. - if issubclass(model.__class__, PreTrainedModel): - fuse_text_encoder_lora( - model, lora_scale=lora_scale, safe_fusing=safe_fusing, adapter_names=adapter_names - ) - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - merged_adapter_names.update(set(module.merged_adapters)) - - self._merged_adapters = self._merged_adapters | merged_adapter_names - - 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). - - - - This is an experimental API. - - - - Args: - 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 - LoRA parameters then it won't have any effect. - """ - if "unfuse_unet" in kwargs: - depr_message = "Passing `unfuse_unet` to `unfuse_lora()` is deprecated and will be ignored. Please use the `components` argument. `unfuse_unet` will be removed in a future version." - deprecate( - "unfuse_unet", - "1.0.0", - depr_message, - ) - if "unfuse_transformer" in kwargs: - depr_message = "Passing `unfuse_transformer` to `unfuse_lora()` is deprecated and will be ignored. Please use the `components` argument. `unfuse_transformer` will be removed in a future version." - deprecate( - "unfuse_transformer", - "1.0.0", - depr_message, - ) - if "unfuse_text_encoder" in kwargs: - depr_message = "Passing `unfuse_text_encoder` to `unfuse_lora()` is deprecated and will be ignored. Please use the `components` argument. `unfuse_text_encoder` will be removed in a future version." - deprecate( - "unfuse_text_encoder", - "1.0.0", - depr_message, - ) - - if len(components) == 0: - raise ValueError("`components` cannot be an empty list.") - - for fuse_component in components: - if fuse_component not in self._lora_loadable_modules: - raise ValueError(f"{fuse_component} is not found in {self._lora_loadable_modules=}.") - - model = getattr(self, fuse_component, None) - if model is not None: - if issubclass(model.__class__, (ModelMixin, PreTrainedModel)): - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - for adapter in set(module.merged_adapters): - if adapter and adapter in self._merged_adapters: - self._merged_adapters = self._merged_adapters - {adapter} - module.unmerge() - - def set_adapters( - self, - adapter_names: Union[List[str], str], - adapter_weights: Optional[Union[float, Dict, List[float], List[Dict]]] = None, - ): - """ - Set the currently active adapters for use in the pipeline. - - Args: - adapter_names (`List[str]` or `str`): - The names of the adapters to use. - adapter_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. - - Example: - - ```py - from diffusers import AutoPipelineForText2Image - import torch - - pipeline = AutoPipelineForText2Image.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights( - "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.set_adapters(["cinematic", "pixel"], adapter_weights=[0.5, 0.5]) - ``` - """ - if isinstance(adapter_weights, dict): - components_passed = set(adapter_weights.keys()) - lora_components = set(self._lora_loadable_modules) - - invalid_components = sorted(components_passed - lora_components) - if invalid_components: - logger.warning( - f"The following components in `adapter_weights` are not part of the pipeline: {invalid_components}. " - f"Available components that are LoRA-compatible: {self._lora_loadable_modules}. So, weights belonging " - "to the invalid components will be removed and ignored." - ) - adapter_weights = {k: v for k, v in adapter_weights.items() if k not in invalid_components} - - adapter_names = [adapter_names] if isinstance(adapter_names, str) else adapter_names - adapter_weights = copy.deepcopy(adapter_weights) - - # Expand weights into a list, one entry per adapter - if not isinstance(adapter_weights, list): - adapter_weights = [adapter_weights] * len(adapter_names) - - if len(adapter_names) != len(adapter_weights): - raise ValueError( - f"Length of adapter names {len(adapter_names)} is not equal to the length of the weights {len(adapter_weights)}" - ) - - list_adapters = self.get_list_adapters() # eg {"unet": ["adapter1", "adapter2"], "text_encoder": ["adapter2"]} - # eg ["adapter1", "adapter2"] - all_adapters = {adapter for adapters in list_adapters.values() for adapter in adapters} - missing_adapters = set(adapter_names) - all_adapters - if len(missing_adapters) > 0: - raise ValueError( - f"Adapter name(s) {missing_adapters} not in the list of present adapters: {all_adapters}." - ) - - # eg {"adapter1": ["unet"], "adapter2": ["unet", "text_encoder"]} - invert_list_adapters = { - adapter: [part for part, adapters in list_adapters.items() if adapter in adapters] - for adapter in all_adapters - } - - # Decompose weights into weights for denoiser and text encoders. - _component_adapter_weights = {} - for component in self._lora_loadable_modules: - model = getattr(self, component) - - for adapter_name, weights in zip(adapter_names, adapter_weights): - if isinstance(weights, dict): - component_adapter_weights = weights.pop(component, None) - if component_adapter_weights is not None and component not in invert_list_adapters[adapter_name]: - logger.warning( - ( - f"Lora weight dict for adapter '{adapter_name}' contains {component}," - f"but this will be ignored because {adapter_name} does not contain weights for {component}." - f"Valid parts for {adapter_name} are: {invert_list_adapters[adapter_name]}." - ) - ) - - else: - component_adapter_weights = weights - - _component_adapter_weights.setdefault(component, []) - _component_adapter_weights[component].append(component_adapter_weights) - - if issubclass(model.__class__, ModelMixin): - model.set_adapters(adapter_names, _component_adapter_weights[component]) - elif issubclass(model.__class__, PreTrainedModel): - set_adapters_for_text_encoder(adapter_names, model, _component_adapter_weights[component]) - - def disable_lora(self): - """ - Disables the active LoRA layers of the pipeline. - - Example: - - ```py - from diffusers import AutoPipelineForText2Image - import torch - - pipeline = AutoPipelineForText2Image.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights( - "jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_name="cinematic" - ) - pipeline.disable_lora() - ``` - """ - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for this method.") - - for component in self._lora_loadable_modules: - model = getattr(self, component, None) - if model is not None: - if issubclass(model.__class__, ModelMixin): - model.disable_lora() - elif issubclass(model.__class__, PreTrainedModel): - disable_lora_for_text_encoder(model) - - def enable_lora(self): - """ - Enables the active LoRA layers of the pipeline. - - Example: - - ```py - from diffusers import AutoPipelineForText2Image - import torch - - pipeline = AutoPipelineForText2Image.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights( - "jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_name="cinematic" - ) - pipeline.enable_lora() - ``` - """ - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for this method.") - - for component in self._lora_loadable_modules: - model = getattr(self, component, None) - if model is not None: - if issubclass(model.__class__, ModelMixin): - model.enable_lora() - elif issubclass(model.__class__, PreTrainedModel): - enable_lora_for_text_encoder(model) - - def delete_adapters(self, adapter_names: Union[List[str], str]): - """ - Delete an adapter's LoRA layers from the pipeline. - - Args: - adapter_names (`Union[List[str], str]`): - The names of the adapters to delete. - - Example: - - ```py - from diffusers import AutoPipelineForText2Image - import torch - - pipeline = AutoPipelineForText2Image.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights( - "jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_names="cinematic" - ) - pipeline.delete_adapters("cinematic") - ``` - """ - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for this method.") - - if isinstance(adapter_names, str): - adapter_names = [adapter_names] - - for component in self._lora_loadable_modules: - model = getattr(self, component, None) - if model is not None: - if issubclass(model.__class__, ModelMixin): - model.delete_adapters(adapter_names) - elif issubclass(model.__class__, PreTrainedModel): - for adapter_name in adapter_names: - delete_adapter_layers(model, adapter_name) - - def get_active_adapters(self) -> List[str]: - """ - Gets the list of the current active adapters. - - Example: - - ```python - from diffusers import DiffusionPipeline - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", - ).to("cuda") - pipeline.load_lora_weights("CiroN2022/toy-face", weight_name="toy_face_sdxl.safetensors", adapter_name="toy") - pipeline.get_active_adapters() - ``` - """ - if not USE_PEFT_BACKEND: - raise ValueError( - "PEFT backend is required for this method. Please install the latest version of PEFT `pip install -U peft`" - ) - - active_adapters = [] - - for component in self._lora_loadable_modules: - model = getattr(self, component, None) - if model is not None and issubclass(model.__class__, ModelMixin): - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - active_adapters = module.active_adapters - break - - return active_adapters - - def get_list_adapters(self) -> Dict[str, List[str]]: - """ - Gets the current list of all available adapters in the pipeline. - """ - if not USE_PEFT_BACKEND: - raise ValueError( - "PEFT backend is required for this method. Please install the latest version of PEFT `pip install -U peft`" - ) - - set_adapters = {} - - for component in self._lora_loadable_modules: - model = getattr(self, component, None) - if ( - model is not None - and issubclass(model.__class__, (ModelMixin, PreTrainedModel)) - and hasattr(model, "peft_config") - ): - set_adapters[component] = list(model.peft_config.keys()) - - return set_adapters - - def set_lora_device(self, adapter_names: List[str], device: Union[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. - - Args: - adapter_names (`List[str]`): - List of adapters to send device to. - device (`Union[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: - raise ValueError("PEFT backend is required for this method.") - - for component in self._lora_loadable_modules: - model = getattr(self, component, None) - if model is not None: - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - for adapter_name in adapter_names: - module.lora_A[adapter_name].to(device) - module.lora_B[adapter_name].to(device) - # this is a param, not a module, so device placement is not in-place -> re-assign - if hasattr(module, "lora_magnitude_vector") and module.lora_magnitude_vector is not None: - if adapter_name in module.lora_magnitude_vector: - module.lora_magnitude_vector[adapter_name] = module.lora_magnitude_vector[ - adapter_name - ].to(device) - - def enable_lora_hotswap(self, **kwargs) -> None: - """ - Hotswap adapters without triggering recompilation of a model or if the ranks of the loaded adapters are - different. - - Args: - target_rank (`int`): - The highest rank among all the adapters that will be loaded. - check_compiled (`str`, *optional*, defaults to `"error"`): - How to handle a model that is already compiled. The check can return the following messages: - - "error" (default): raise an error - - "warn": issue a warning - - "ignore": do nothing - """ - for key, component in self.components.items(): - if hasattr(component, "enable_lora_hotswap") and (key in self._lora_loadable_modules): - component.enable_lora_hotswap(**kwargs) - - @staticmethod - def pack_weights(layers, prefix): - layers_weights = layers.state_dict() if isinstance(layers, torch.nn.Module) else layers - return _pack_dict_with_prefix(layers_weights, prefix) - - @staticmethod - def write_lora_layers( - 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, - ): - """Writes the state dict of the LoRA layers (optionally with metadata) to disk.""" - if os.path.isfile(save_directory): - logger.error(f"Provided path ({save_directory}) should be a directory, not a file") - return - - if lora_adapter_metadata and not safe_serialization: - raise ValueError("`lora_adapter_metadata` cannot be specified when not using `safe_serialization`.") - if lora_adapter_metadata and not isinstance(lora_adapter_metadata, dict): - raise TypeError("`lora_adapter_metadata` must be of type `dict`.") - - if save_function is None: - if safe_serialization: - - def save_function(weights, filename): - # Inject framework format. - metadata = {"format": "pt"} - if lora_adapter_metadata: - 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 - ) - - return safetensors.torch.save_file(weights, filename, metadata=metadata) - - else: - save_function = torch.save - - os.makedirs(save_directory, exist_ok=True) - - if weight_name is None: - if safe_serialization: - weight_name = LORA_WEIGHT_NAME_SAFE - else: - weight_name = LORA_WEIGHT_NAME - - save_path = Path(save_directory, weight_name).as_posix() - save_function(state_dict, save_path) - logger.info(f"Model weights saved in {save_path}") - - @classmethod - def _optionally_disable_offloading(cls, _pipeline): - return _func_optionally_disable_offloading(_pipeline=_pipeline) - - @classmethod - def _fetch_state_dict(cls, *args, **kwargs): - deprecation_message = f"Using the `_fetch_state_dict()` method from {cls} has been deprecated and will be removed in a future version. Please use `from diffusers.loaders.lora_base import _fetch_state_dict`." - deprecate("_fetch_state_dict", "0.35.0", deprecation_message) - return _fetch_state_dict(*args, **kwargs) - - @classmethod - def _best_guess_weight_name(cls, *args, **kwargs): - deprecation_message = f"Using the `_best_guess_weight_name()` method from {cls} has been deprecated and will be removed in a future version. Please use `from diffusers.loaders.lora_base import _best_guess_weight_name`." - deprecate("_best_guess_weight_name", "0.35.0", deprecation_message) - return _best_guess_weight_name(*args, **kwargs) diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/lora_conversion_utils.py b/venv/lib/python3.10/site-packages/diffusers/loaders/lora_conversion_utils.py deleted file mode 100644 index 25e06c007fbc7df918aa8381c55df8e2fe68f3e0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/lora_conversion_utils.py +++ /dev/null @@ -1,1851 +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 re -from typing import List - -import torch - -from ..utils import is_peft_version, logging, state_dict_all_zero - - -logger = logging.get_logger(__name__) - - -def swap_scale_shift(weight): - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - return new_weight - - -def _maybe_map_sgm_blocks_to_diffusers(state_dict, unet_config, delimiter="_", block_slice_pos=5): - # 1. get all state_dict_keys - all_keys = list(state_dict.keys()) - sgm_patterns = ["input_blocks", "middle_block", "output_blocks"] - not_sgm_patterns = ["down_blocks", "mid_block", "up_blocks"] - - # check if state_dict contains both patterns - contains_sgm_patterns = False - contains_not_sgm_patterns = False - for key in all_keys: - if any(p in key for p in sgm_patterns): - contains_sgm_patterns = True - elif any(p in key for p in not_sgm_patterns): - contains_not_sgm_patterns = True - - # if state_dict contains both patterns, remove sgm - # we can then return state_dict immediately - if contains_sgm_patterns and contains_not_sgm_patterns: - for key in all_keys: - if any(p in key for p in sgm_patterns): - state_dict.pop(key) - return state_dict - - # 2. check if needs remapping, if not return original dict - is_in_sgm_format = False - for key in all_keys: - if any(p in key for p in sgm_patterns): - is_in_sgm_format = True - break - - if not is_in_sgm_format: - return state_dict - - # 3. Else remap from SGM patterns - new_state_dict = {} - inner_block_map = ["resnets", "attentions", "upsamplers"] - - # Retrieves # of down, mid and up blocks - input_block_ids, middle_block_ids, output_block_ids = set(), set(), set() - - for layer in all_keys: - if "text" in layer: - new_state_dict[layer] = state_dict.pop(layer) - else: - layer_id = int(layer.split(delimiter)[:block_slice_pos][-1]) - if sgm_patterns[0] in layer: - input_block_ids.add(layer_id) - elif sgm_patterns[1] in layer: - middle_block_ids.add(layer_id) - elif sgm_patterns[2] in layer: - output_block_ids.add(layer_id) - else: - raise ValueError(f"Checkpoint not supported because layer {layer} not supported.") - - input_blocks = { - layer_id: [key for key in state_dict if f"input_blocks{delimiter}{layer_id}" in key] - for layer_id in input_block_ids - } - middle_blocks = { - layer_id: [key for key in state_dict if f"middle_block{delimiter}{layer_id}" in key] - for layer_id in middle_block_ids - } - output_blocks = { - layer_id: [key for key in state_dict if f"output_blocks{delimiter}{layer_id}" in key] - for layer_id in output_block_ids - } - - # Rename keys accordingly - for i in input_block_ids: - block_id = (i - 1) // (unet_config.layers_per_block + 1) - layer_in_block_id = (i - 1) % (unet_config.layers_per_block + 1) - - for key in input_blocks[i]: - inner_block_id = int(key.split(delimiter)[block_slice_pos]) - inner_block_key = inner_block_map[inner_block_id] if "op" not in key else "downsamplers" - inner_layers_in_block = str(layer_in_block_id) if "op" not in key else "0" - new_key = delimiter.join( - key.split(delimiter)[: block_slice_pos - 1] - + [str(block_id), inner_block_key, inner_layers_in_block] - + key.split(delimiter)[block_slice_pos + 1 :] - ) - new_state_dict[new_key] = state_dict.pop(key) - - for i in middle_block_ids: - key_part = None - if i == 0: - key_part = [inner_block_map[0], "0"] - elif i == 1: - key_part = [inner_block_map[1], "0"] - elif i == 2: - key_part = [inner_block_map[0], "1"] - else: - raise ValueError(f"Invalid middle block id {i}.") - - for key in middle_blocks[i]: - new_key = delimiter.join( - key.split(delimiter)[: block_slice_pos - 1] + key_part + key.split(delimiter)[block_slice_pos:] - ) - new_state_dict[new_key] = state_dict.pop(key) - - for i in output_block_ids: - block_id = i // (unet_config.layers_per_block + 1) - layer_in_block_id = i % (unet_config.layers_per_block + 1) - - for key in output_blocks[i]: - inner_block_id = int(key.split(delimiter)[block_slice_pos]) - inner_block_key = inner_block_map[inner_block_id] - inner_layers_in_block = str(layer_in_block_id) if inner_block_id < 2 else "0" - new_key = delimiter.join( - key.split(delimiter)[: block_slice_pos - 1] - + [str(block_id), inner_block_key, inner_layers_in_block] - + key.split(delimiter)[block_slice_pos + 1 :] - ) - new_state_dict[new_key] = state_dict.pop(key) - - if state_dict: - raise ValueError("At this point all state dict entries have to be converted.") - - return new_state_dict - - -def _convert_non_diffusers_lora_to_diffusers(state_dict, unet_name="unet", text_encoder_name="text_encoder"): - """ - Converts a non-Diffusers LoRA state dict to a Diffusers compatible state dict. - - Args: - state_dict (`dict`): The state dict to convert. - unet_name (`str`, optional): The name of the U-Net module in the Diffusers model. Defaults to "unet". - text_encoder_name (`str`, optional): The name of the text encoder module in the Diffusers model. Defaults to - "text_encoder". - - Returns: - `tuple`: A tuple containing the converted state dict and a dictionary of alphas. - """ - unet_state_dict = {} - te_state_dict = {} - te2_state_dict = {} - network_alphas = {} - - # Check for DoRA-enabled LoRAs. - dora_present_in_unet = any("dora_scale" in k and "lora_unet_" in k for k in state_dict) - dora_present_in_te = any("dora_scale" in k and ("lora_te_" in k or "lora_te1_" in k) for k in state_dict) - dora_present_in_te2 = any("dora_scale" in k and "lora_te2_" in k for k in state_dict) - if dora_present_in_unet or dora_present_in_te or dora_present_in_te2: - if is_peft_version("<", "0.9.0"): - raise ValueError( - "You need `peft` 0.9.0 at least to use DoRA-enabled LoRAs. Please upgrade your installation of `peft`." - ) - - # Iterate over all LoRA weights. - all_lora_keys = list(state_dict.keys()) - for key in all_lora_keys: - if not key.endswith("lora_down.weight"): - continue - - # Extract LoRA name. - lora_name = key.split(".")[0] - - # Find corresponding up weight and alpha. - lora_name_up = lora_name + ".lora_up.weight" - lora_name_alpha = lora_name + ".alpha" - - # Handle U-Net LoRAs. - if lora_name.startswith("lora_unet_"): - diffusers_name = _convert_unet_lora_key(key) - - # Store down and up weights. - unet_state_dict[diffusers_name] = state_dict.pop(key) - unet_state_dict[diffusers_name.replace(".down.", ".up.")] = state_dict.pop(lora_name_up) - - # Store DoRA scale if present. - if dora_present_in_unet: - dora_scale_key_to_replace = "_lora.down." if "_lora.down." in diffusers_name else ".lora.down." - unet_state_dict[diffusers_name.replace(dora_scale_key_to_replace, ".lora_magnitude_vector.")] = ( - state_dict.pop(key.replace("lora_down.weight", "dora_scale")) - ) - - # Handle text encoder LoRAs. - elif lora_name.startswith(("lora_te_", "lora_te1_", "lora_te2_")): - diffusers_name = _convert_text_encoder_lora_key(key, lora_name) - - # Store down and up weights for te or te2. - if lora_name.startswith(("lora_te_", "lora_te1_")): - te_state_dict[diffusers_name] = state_dict.pop(key) - te_state_dict[diffusers_name.replace(".down.", ".up.")] = state_dict.pop(lora_name_up) - else: - te2_state_dict[diffusers_name] = state_dict.pop(key) - te2_state_dict[diffusers_name.replace(".down.", ".up.")] = state_dict.pop(lora_name_up) - - # Store DoRA scale if present. - if dora_present_in_te or dora_present_in_te2: - dora_scale_key_to_replace_te = ( - "_lora.down." if "_lora.down." in diffusers_name else ".lora_linear_layer." - ) - if lora_name.startswith(("lora_te_", "lora_te1_")): - te_state_dict[diffusers_name.replace(dora_scale_key_to_replace_te, ".lora_magnitude_vector.")] = ( - state_dict.pop(key.replace("lora_down.weight", "dora_scale")) - ) - elif lora_name.startswith("lora_te2_"): - te2_state_dict[diffusers_name.replace(dora_scale_key_to_replace_te, ".lora_magnitude_vector.")] = ( - state_dict.pop(key.replace("lora_down.weight", "dora_scale")) - ) - - # Store alpha if present. - if lora_name_alpha in state_dict: - alpha = state_dict.pop(lora_name_alpha).item() - network_alphas.update(_get_alpha_name(lora_name_alpha, diffusers_name, alpha)) - - # Check if any keys remain. - if len(state_dict) > 0: - raise ValueError(f"The following keys have not been correctly renamed: \n\n {', '.join(state_dict.keys())}") - - logger.info("Non-diffusers checkpoint detected.") - - # Construct final state dict. - unet_state_dict = {f"{unet_name}.{module_name}": params for module_name, params in unet_state_dict.items()} - te_state_dict = {f"{text_encoder_name}.{module_name}": params for module_name, params in te_state_dict.items()} - te2_state_dict = ( - {f"text_encoder_2.{module_name}": params for module_name, params in te2_state_dict.items()} - if len(te2_state_dict) > 0 - else None - ) - if te2_state_dict is not None: - te_state_dict.update(te2_state_dict) - - new_state_dict = {**unet_state_dict, **te_state_dict} - return new_state_dict, network_alphas - - -def _convert_unet_lora_key(key): - """ - Converts a U-Net LoRA key to a Diffusers compatible key. - """ - diffusers_name = key.replace("lora_unet_", "").replace("_", ".") - - # Replace common U-Net naming patterns. - diffusers_name = diffusers_name.replace("input.blocks", "down_blocks") - diffusers_name = diffusers_name.replace("down.blocks", "down_blocks") - diffusers_name = diffusers_name.replace("middle.block", "mid_block") - diffusers_name = diffusers_name.replace("mid.block", "mid_block") - diffusers_name = diffusers_name.replace("output.blocks", "up_blocks") - diffusers_name = diffusers_name.replace("up.blocks", "up_blocks") - diffusers_name = diffusers_name.replace("transformer.blocks", "transformer_blocks") - diffusers_name = diffusers_name.replace("to.q.lora", "to_q_lora") - diffusers_name = diffusers_name.replace("to.k.lora", "to_k_lora") - diffusers_name = diffusers_name.replace("to.v.lora", "to_v_lora") - diffusers_name = diffusers_name.replace("to.out.0.lora", "to_out_lora") - diffusers_name = diffusers_name.replace("proj.in", "proj_in") - diffusers_name = diffusers_name.replace("proj.out", "proj_out") - diffusers_name = diffusers_name.replace("emb.layers", "time_emb_proj") - - # SDXL specific conversions. - if "emb" in diffusers_name and "time.emb.proj" not in diffusers_name: - pattern = r"\.\d+(?=\D*$)" - diffusers_name = re.sub(pattern, "", diffusers_name, count=1) - if ".in." in diffusers_name: - diffusers_name = diffusers_name.replace("in.layers.2", "conv1") - if ".out." in diffusers_name: - diffusers_name = diffusers_name.replace("out.layers.3", "conv2") - if "downsamplers" in diffusers_name or "upsamplers" in diffusers_name: - diffusers_name = diffusers_name.replace("op", "conv") - if "skip" in diffusers_name: - diffusers_name = diffusers_name.replace("skip.connection", "conv_shortcut") - - # LyCORIS specific conversions. - if "time.emb.proj" in diffusers_name: - diffusers_name = diffusers_name.replace("time.emb.proj", "time_emb_proj") - if "conv.shortcut" in diffusers_name: - diffusers_name = diffusers_name.replace("conv.shortcut", "conv_shortcut") - - # General conversions. - if "transformer_blocks" in diffusers_name: - if "attn1" in diffusers_name or "attn2" in diffusers_name: - diffusers_name = diffusers_name.replace("attn1", "attn1.processor") - diffusers_name = diffusers_name.replace("attn2", "attn2.processor") - elif "ff" in diffusers_name: - pass - elif any(key in diffusers_name for key in ("proj_in", "proj_out")): - pass - else: - pass - - return diffusers_name - - -def _convert_text_encoder_lora_key(key, lora_name): - """ - Converts a text encoder LoRA key to a Diffusers compatible key. - """ - if lora_name.startswith(("lora_te_", "lora_te1_")): - key_to_replace = "lora_te_" if lora_name.startswith("lora_te_") else "lora_te1_" - else: - key_to_replace = "lora_te2_" - - diffusers_name = key.replace(key_to_replace, "").replace("_", ".") - diffusers_name = diffusers_name.replace("text.model", "text_model") - diffusers_name = diffusers_name.replace("self.attn", "self_attn") - diffusers_name = diffusers_name.replace("q.proj.lora", "to_q_lora") - diffusers_name = diffusers_name.replace("k.proj.lora", "to_k_lora") - diffusers_name = diffusers_name.replace("v.proj.lora", "to_v_lora") - diffusers_name = diffusers_name.replace("out.proj.lora", "to_out_lora") - diffusers_name = diffusers_name.replace("text.projection", "text_projection") - - if "self_attn" in diffusers_name or "text_projection" in diffusers_name: - pass - elif "mlp" in diffusers_name: - # Be aware that this is the new diffusers convention and the rest of the code might - # not utilize it yet. - diffusers_name = diffusers_name.replace(".lora.", ".lora_linear_layer.") - - return diffusers_name - - -def _get_alpha_name(lora_name_alpha, diffusers_name, alpha): - """ - Gets the correct alpha name for the Diffusers model. - """ - if lora_name_alpha.startswith("lora_unet_"): - prefix = "unet." - elif lora_name_alpha.startswith(("lora_te_", "lora_te1_")): - prefix = "text_encoder." - else: - prefix = "text_encoder_2." - new_name = prefix + diffusers_name.split(".lora.")[0] + ".alpha" - return {new_name: alpha} - - -# The utilities under `_convert_kohya_flux_lora_to_diffusers()` -# are adapted from https://github.com/kohya-ss/sd-scripts/blob/a61cf73a5cb5209c3f4d1a3688dd276a4dfd1ecb/networks/convert_flux_lora.py -def _convert_kohya_flux_lora_to_diffusers(state_dict): - def _convert_to_ai_toolkit(sds_sd, ait_sd, sds_key, ait_key): - if sds_key + ".lora_down.weight" not in sds_sd: - return - down_weight = sds_sd.pop(sds_key + ".lora_down.weight") - - # scale weight by alpha and dim - rank = down_weight.shape[0] - default_alpha = torch.tensor(rank, dtype=down_weight.dtype, device=down_weight.device, requires_grad=False) - alpha = sds_sd.pop(sds_key + ".alpha", default_alpha).item() # alpha is scalar - scale = alpha / rank # LoRA is scaled by 'alpha / rank' in forward pass, so we need to scale it back here - - # calculate scale_down and scale_up to keep the same value. if scale is 4, scale_down is 2 and scale_up is 2 - scale_down = scale - scale_up = 1.0 - while scale_down * 2 < scale_up: - scale_down *= 2 - scale_up /= 2 - - ait_sd[ait_key + ".lora_A.weight"] = down_weight * scale_down - ait_sd[ait_key + ".lora_B.weight"] = sds_sd.pop(sds_key + ".lora_up.weight") * scale_up - - def _convert_to_ai_toolkit_cat(sds_sd, ait_sd, sds_key, ait_keys, dims=None): - if sds_key + ".lora_down.weight" not in sds_sd: - return - down_weight = sds_sd.pop(sds_key + ".lora_down.weight") - up_weight = sds_sd.pop(sds_key + ".lora_up.weight") - sd_lora_rank = down_weight.shape[0] - - # scale weight by alpha and dim - default_alpha = torch.tensor( - sd_lora_rank, dtype=down_weight.dtype, device=down_weight.device, requires_grad=False - ) - alpha = sds_sd.pop(sds_key + ".alpha", default_alpha) - scale = alpha / sd_lora_rank - - # calculate scale_down and scale_up - scale_down = scale - scale_up = 1.0 - while scale_down * 2 < scale_up: - scale_down *= 2 - scale_up /= 2 - - down_weight = down_weight * scale_down - up_weight = up_weight * scale_up - - # calculate dims if not provided - num_splits = len(ait_keys) - if dims is None: - dims = [up_weight.shape[0] // num_splits] * num_splits - else: - assert sum(dims) == up_weight.shape[0] - - # check upweight is sparse or not - is_sparse = False - if sd_lora_rank % num_splits == 0: - ait_rank = sd_lora_rank // num_splits - is_sparse = True - i = 0 - for j in range(len(dims)): - for k in range(len(dims)): - if j == k: - continue - is_sparse = is_sparse and torch.all( - up_weight[i : i + dims[j], k * ait_rank : (k + 1) * ait_rank] == 0 - ) - i += dims[j] - if is_sparse: - logger.info(f"weight is sparse: {sds_key}") - - # make ai-toolkit weight - ait_down_keys = [k + ".lora_A.weight" for k in ait_keys] - ait_up_keys = [k + ".lora_B.weight" for k in ait_keys] - if not is_sparse: - # down_weight is copied to each split - ait_sd.update(dict.fromkeys(ait_down_keys, down_weight)) - - # up_weight is split to each split - ait_sd.update({k: v for k, v in zip(ait_up_keys, torch.split(up_weight, dims, dim=0))}) # noqa: C416 - else: - # down_weight is chunked to each split - ait_sd.update({k: v for k, v in zip(ait_down_keys, torch.chunk(down_weight, num_splits, dim=0))}) # noqa: C416 - - # up_weight is sparse: only non-zero values are copied to each split - i = 0 - for j in range(len(dims)): - ait_sd[ait_up_keys[j]] = up_weight[i : i + dims[j], j * ait_rank : (j + 1) * ait_rank].contiguous() - i += dims[j] - - def _convert_sd_scripts_to_ai_toolkit(sds_sd): - ait_sd = {} - for i in range(19): - _convert_to_ai_toolkit( - sds_sd, - ait_sd, - f"lora_unet_double_blocks_{i}_img_attn_proj", - f"transformer.transformer_blocks.{i}.attn.to_out.0", - ) - _convert_to_ai_toolkit_cat( - sds_sd, - ait_sd, - f"lora_unet_double_blocks_{i}_img_attn_qkv", - [ - f"transformer.transformer_blocks.{i}.attn.to_q", - f"transformer.transformer_blocks.{i}.attn.to_k", - f"transformer.transformer_blocks.{i}.attn.to_v", - ], - ) - _convert_to_ai_toolkit( - sds_sd, - ait_sd, - f"lora_unet_double_blocks_{i}_img_mlp_0", - f"transformer.transformer_blocks.{i}.ff.net.0.proj", - ) - _convert_to_ai_toolkit( - sds_sd, - ait_sd, - f"lora_unet_double_blocks_{i}_img_mlp_2", - f"transformer.transformer_blocks.{i}.ff.net.2", - ) - _convert_to_ai_toolkit( - sds_sd, - ait_sd, - f"lora_unet_double_blocks_{i}_img_mod_lin", - f"transformer.transformer_blocks.{i}.norm1.linear", - ) - _convert_to_ai_toolkit( - sds_sd, - ait_sd, - f"lora_unet_double_blocks_{i}_txt_attn_proj", - f"transformer.transformer_blocks.{i}.attn.to_add_out", - ) - _convert_to_ai_toolkit_cat( - sds_sd, - ait_sd, - f"lora_unet_double_blocks_{i}_txt_attn_qkv", - [ - f"transformer.transformer_blocks.{i}.attn.add_q_proj", - f"transformer.transformer_blocks.{i}.attn.add_k_proj", - f"transformer.transformer_blocks.{i}.attn.add_v_proj", - ], - ) - _convert_to_ai_toolkit( - sds_sd, - ait_sd, - f"lora_unet_double_blocks_{i}_txt_mlp_0", - f"transformer.transformer_blocks.{i}.ff_context.net.0.proj", - ) - _convert_to_ai_toolkit( - sds_sd, - ait_sd, - f"lora_unet_double_blocks_{i}_txt_mlp_2", - f"transformer.transformer_blocks.{i}.ff_context.net.2", - ) - _convert_to_ai_toolkit( - sds_sd, - ait_sd, - f"lora_unet_double_blocks_{i}_txt_mod_lin", - f"transformer.transformer_blocks.{i}.norm1_context.linear", - ) - - for i in range(38): - _convert_to_ai_toolkit_cat( - sds_sd, - ait_sd, - f"lora_unet_single_blocks_{i}_linear1", - [ - f"transformer.single_transformer_blocks.{i}.attn.to_q", - f"transformer.single_transformer_blocks.{i}.attn.to_k", - f"transformer.single_transformer_blocks.{i}.attn.to_v", - f"transformer.single_transformer_blocks.{i}.proj_mlp", - ], - dims=[3072, 3072, 3072, 12288], - ) - _convert_to_ai_toolkit( - sds_sd, - ait_sd, - f"lora_unet_single_blocks_{i}_linear2", - f"transformer.single_transformer_blocks.{i}.proj_out", - ) - _convert_to_ai_toolkit( - sds_sd, - ait_sd, - f"lora_unet_single_blocks_{i}_modulation_lin", - f"transformer.single_transformer_blocks.{i}.norm.linear", - ) - - # TODO: alphas. - def assign_remaining_weights(assignments, source): - for lora_key in ["lora_A", "lora_B"]: - orig_lora_key = "lora_down" if lora_key == "lora_A" else "lora_up" - for target_fmt, source_fmt, transform in assignments: - target_key = target_fmt.format(lora_key=lora_key) - source_key = source_fmt.format(orig_lora_key=orig_lora_key) - value = source.pop(source_key) - if transform: - value = transform(value) - ait_sd[target_key] = value - - if any("guidance_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ( - "time_text_embed.guidance_embedder.linear_1.{lora_key}.weight", - "lora_unet_guidance_in_in_layer.{orig_lora_key}.weight", - None, - ), - ( - "time_text_embed.guidance_embedder.linear_2.{lora_key}.weight", - "lora_unet_guidance_in_out_layer.{orig_lora_key}.weight", - None, - ), - ], - sds_sd, - ) - - if any("img_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ("x_embedder.{lora_key}.weight", "lora_unet_img_in.{orig_lora_key}.weight", None), - ], - sds_sd, - ) - - if any("txt_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ("context_embedder.{lora_key}.weight", "lora_unet_txt_in.{orig_lora_key}.weight", None), - ], - sds_sd, - ) - - if any("time_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ( - "time_text_embed.timestep_embedder.linear_1.{lora_key}.weight", - "lora_unet_time_in_in_layer.{orig_lora_key}.weight", - None, - ), - ( - "time_text_embed.timestep_embedder.linear_2.{lora_key}.weight", - "lora_unet_time_in_out_layer.{orig_lora_key}.weight", - None, - ), - ], - sds_sd, - ) - - if any("vector_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ( - "time_text_embed.text_embedder.linear_1.{lora_key}.weight", - "lora_unet_vector_in_in_layer.{orig_lora_key}.weight", - None, - ), - ( - "time_text_embed.text_embedder.linear_2.{lora_key}.weight", - "lora_unet_vector_in_out_layer.{orig_lora_key}.weight", - None, - ), - ], - sds_sd, - ) - - if any("final_layer" in k for k in sds_sd): - # Notice the swap in processing for "final_layer". - assign_remaining_weights( - [ - ( - "norm_out.linear.{lora_key}.weight", - "lora_unet_final_layer_adaLN_modulation_1.{orig_lora_key}.weight", - swap_scale_shift, - ), - ("proj_out.{lora_key}.weight", "lora_unet_final_layer_linear.{orig_lora_key}.weight", None), - ], - sds_sd, - ) - - remaining_keys = list(sds_sd.keys()) - te_state_dict = {} - if remaining_keys: - if not all(k.startswith(("lora_te", "lora_te1")) for k in remaining_keys): - raise ValueError(f"Incompatible keys detected: \n\n {', '.join(remaining_keys)}") - for key in remaining_keys: - if not key.endswith("lora_down.weight"): - continue - - lora_name = key.split(".")[0] - lora_name_up = f"{lora_name}.lora_up.weight" - lora_name_alpha = f"{lora_name}.alpha" - diffusers_name = _convert_text_encoder_lora_key(key, lora_name) - - if lora_name.startswith(("lora_te_", "lora_te1_")): - down_weight = sds_sd.pop(key) - sd_lora_rank = down_weight.shape[0] - te_state_dict[diffusers_name] = down_weight - te_state_dict[diffusers_name.replace(".down.", ".up.")] = sds_sd.pop(lora_name_up) - - if lora_name_alpha in sds_sd: - alpha = sds_sd.pop(lora_name_alpha).item() - scale = alpha / sd_lora_rank - - scale_down = scale - scale_up = 1.0 - while scale_down * 2 < scale_up: - scale_down *= 2 - scale_up /= 2 - - te_state_dict[diffusers_name] *= scale_down - te_state_dict[diffusers_name.replace(".down.", ".up.")] *= scale_up - - if len(sds_sd) > 0: - logger.warning(f"Unsupported keys for ai-toolkit: {sds_sd.keys()}") - - if te_state_dict: - te_state_dict = {f"text_encoder.{module_name}": params for module_name, params in te_state_dict.items()} - - new_state_dict = {**ait_sd, **te_state_dict} - return new_state_dict - - def _convert_mixture_state_dict_to_diffusers(state_dict): - new_state_dict = {} - - def _convert(original_key, diffusers_key, state_dict, new_state_dict): - down_key = f"{original_key}.lora_down.weight" - down_weight = state_dict.pop(down_key) - lora_rank = down_weight.shape[0] - - up_weight_key = f"{original_key}.lora_up.weight" - up_weight = state_dict.pop(up_weight_key) - - alpha_key = f"{original_key}.alpha" - alpha = state_dict.pop(alpha_key) - - # scale weight by alpha and dim - scale = alpha / lora_rank - # calculate scale_down and scale_up - scale_down = scale - scale_up = 1.0 - while scale_down * 2 < scale_up: - scale_down *= 2 - scale_up /= 2 - down_weight = down_weight * scale_down - up_weight = up_weight * scale_up - - diffusers_down_key = f"{diffusers_key}.lora_A.weight" - new_state_dict[diffusers_down_key] = down_weight - new_state_dict[diffusers_down_key.replace(".lora_A.", ".lora_B.")] = up_weight - - all_unique_keys = { - k.replace(".lora_down.weight", "").replace(".lora_up.weight", "").replace(".alpha", "") - for k in state_dict - if not k.startswith(("lora_unet_")) - } - assert all(k.startswith(("lora_transformer_", "lora_te1_")) for k in all_unique_keys), f"{all_unique_keys=}" - - has_te_keys = False - for k in all_unique_keys: - if k.startswith("lora_transformer_single_transformer_blocks_"): - i = int(k.split("lora_transformer_single_transformer_blocks_")[-1].split("_")[0]) - diffusers_key = f"single_transformer_blocks.{i}" - elif k.startswith("lora_transformer_transformer_blocks_"): - i = int(k.split("lora_transformer_transformer_blocks_")[-1].split("_")[0]) - diffusers_key = f"transformer_blocks.{i}" - elif k.startswith("lora_te1_"): - has_te_keys = True - continue - elif k.startswith("lora_transformer_context_embedder"): - diffusers_key = "context_embedder" - elif k.startswith("lora_transformer_norm_out_linear"): - diffusers_key = "norm_out.linear" - elif k.startswith("lora_transformer_proj_out"): - diffusers_key = "proj_out" - elif k.startswith("lora_transformer_x_embedder"): - diffusers_key = "x_embedder" - elif k.startswith("lora_transformer_time_text_embed_guidance_embedder_linear_"): - i = int(k.split("lora_transformer_time_text_embed_guidance_embedder_linear_")[-1]) - diffusers_key = f"time_text_embed.guidance_embedder.linear_{i}" - elif k.startswith("lora_transformer_time_text_embed_text_embedder_linear_"): - i = int(k.split("lora_transformer_time_text_embed_text_embedder_linear_")[-1]) - diffusers_key = f"time_text_embed.text_embedder.linear_{i}" - elif k.startswith("lora_transformer_time_text_embed_timestep_embedder_linear_"): - i = int(k.split("lora_transformer_time_text_embed_timestep_embedder_linear_")[-1]) - diffusers_key = f"time_text_embed.timestep_embedder.linear_{i}" - else: - raise NotImplementedError(f"Handling for key ({k}) is not implemented.") - - if "attn_" in k: - if "_to_out_0" in k: - diffusers_key += ".attn.to_out.0" - elif "_to_add_out" in k: - diffusers_key += ".attn.to_add_out" - elif any(qkv in k for qkv in ["to_q", "to_k", "to_v"]): - remaining = k.split("attn_")[-1] - diffusers_key += f".attn.{remaining}" - elif any(add_qkv in k for add_qkv in ["add_q_proj", "add_k_proj", "add_v_proj"]): - remaining = k.split("attn_")[-1] - diffusers_key += f".attn.{remaining}" - - _convert(k, diffusers_key, state_dict, new_state_dict) - - if has_te_keys: - layer_pattern = re.compile(r"lora_te1_text_model_encoder_layers_(\d+)") - attn_mapping = { - "q_proj": ".self_attn.q_proj", - "k_proj": ".self_attn.k_proj", - "v_proj": ".self_attn.v_proj", - "out_proj": ".self_attn.out_proj", - } - mlp_mapping = {"fc1": ".mlp.fc1", "fc2": ".mlp.fc2"} - for k in all_unique_keys: - if not k.startswith("lora_te1_"): - continue - - match = layer_pattern.search(k) - if not match: - continue - i = int(match.group(1)) - diffusers_key = f"text_model.encoder.layers.{i}" - - if "attn" in k: - for key_fragment, suffix in attn_mapping.items(): - if key_fragment in k: - diffusers_key += suffix - break - elif "mlp" in k: - for key_fragment, suffix in mlp_mapping.items(): - if key_fragment in k: - diffusers_key += suffix - break - - _convert(k, diffusers_key, state_dict, new_state_dict) - - remaining_all_unet = False - if state_dict: - remaining_all_unet = all(k.startswith("lora_unet_") for k in state_dict) - if remaining_all_unet: - keys = list(state_dict.keys()) - for k in keys: - state_dict.pop(k) - - if len(state_dict) > 0: - raise ValueError( - f"Expected an empty state dict at this point but its has these keys which couldn't be parsed: {list(state_dict.keys())}." - ) - - transformer_state_dict = { - f"transformer.{k}": v for k, v in new_state_dict.items() if not k.startswith("text_model.") - } - te_state_dict = {f"text_encoder.{k}": v for k, v in new_state_dict.items() if k.startswith("text_model.")} - return {**transformer_state_dict, **te_state_dict} - - # This is weird. - # https://huggingface.co/sayakpaul/different-lora-from-civitai/tree/main?show_file_info=sharp_detailed_foot.safetensors - # has both `peft` and non-peft state dict. - has_peft_state_dict = any(k.startswith("transformer.") for k in state_dict) - if has_peft_state_dict: - state_dict = {k: v for k, v in state_dict.items() if k.startswith("transformer.")} - return state_dict - - # Another weird one. - has_mixture = any( - k.startswith("lora_transformer_") and ("lora_down" in k or "lora_up" in k or "alpha" in k) for k in state_dict - ) - - # ComfyUI. - if not has_mixture: - state_dict = {k.replace("diffusion_model.", "lora_unet_"): v for k, v in state_dict.items()} - state_dict = {k.replace("text_encoders.clip_l.transformer.", "lora_te_"): v for k, v in state_dict.items()} - - has_position_embedding = any("position_embedding" in k for k in state_dict) - if has_position_embedding: - zero_status_pe = state_dict_all_zero(state_dict, "position_embedding") - if zero_status_pe: - logger.info( - "The `position_embedding` LoRA params are all zeros which make them ineffective. " - "So, we will purge them out of the current state dict to make loading possible." - ) - - else: - logger.info( - "The state_dict has position_embedding LoRA params and we currently do not support them. " - "Open an issue if you need this supported - https://github.com/huggingface/diffusers/issues/new." - ) - state_dict = {k: v for k, v in state_dict.items() if "position_embedding" not in k} - - has_t5xxl = any(k.startswith("text_encoders.t5xxl.transformer.") for k in state_dict) - if has_t5xxl: - zero_status_t5 = state_dict_all_zero(state_dict, "text_encoders.t5xxl") - if zero_status_t5: - logger.info( - "The `t5xxl` LoRA params are all zeros which make them ineffective. " - "So, we will purge them out of the current state dict to make loading possible." - ) - else: - logger.info( - "T5-xxl keys found in the state dict, which are currently unsupported. We will filter them out." - "Open an issue if this is a problem - https://github.com/huggingface/diffusers/issues/new." - ) - 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) - if has_diffb: - zero_status_diff_b = state_dict_all_zero(state_dict, ".diff_b") - if zero_status_diff_b: - logger.info( - "The `diff_b` LoRA params are all zeros which make them ineffective. " - "So, we will purge them out of the current state dict to make loading possible." - ) - else: - logger.info( - "`diff_b` keys found in the state dict which are currently unsupported. " - "So, we will filter out those keys. Open an issue if this is a problem - " - "https://github.com/huggingface/diffusers/issues/new." - ) - state_dict = {k: v for k, v in state_dict.items() if ".diff_b" not in k} - - has_norm_diff = any(".norm" in k and ".diff" in k for k in state_dict) - if has_norm_diff: - zero_status_diff = state_dict_all_zero(state_dict, ".diff") - if zero_status_diff: - logger.info( - "The `diff` LoRA params are all zeros which make them ineffective. " - "So, we will purge them out of the current state dict to make loading possible." - ) - else: - logger.info( - "Normalization diff keys found in the state dict which are currently unsupported. " - "So, we will filter out those keys. Open an issue if this is a problem - " - "https://github.com/huggingface/diffusers/issues/new." - ) - state_dict = {k: v for k, v in state_dict.items() if ".norm" not in k and ".diff" not in k} - - limit_substrings = ["lora_down", "lora_up"] - if any("alpha" in k for k in state_dict): - limit_substrings.append("alpha") - - state_dict = { - _custom_replace(k, limit_substrings): v - for k, v in state_dict.items() - if k.startswith(("lora_unet_", "lora_te_")) - } - - if any("text_projection" in k for k in state_dict): - logger.info( - "`text_projection` keys found in the `state_dict` which are unexpected. " - "So, we will filter out those keys. Open an issue if this is a problem - " - "https://github.com/huggingface/diffusers/issues/new." - ) - state_dict = {k: v for k, v in state_dict.items() if "text_projection" not in k} - - if has_mixture: - return _convert_mixture_state_dict_to_diffusers(state_dict) - - return _convert_sd_scripts_to_ai_toolkit(state_dict) - - -# Adapted from https://gist.github.com/Leommm-byte/6b331a1e9bd53271210b26543a7065d6 -# Some utilities were reused from -# https://github.com/kohya-ss/sd-scripts/blob/a61cf73a5cb5209c3f4d1a3688dd276a4dfd1ecb/networks/convert_flux_lora.py -def _convert_xlabs_flux_lora_to_diffusers(old_state_dict): - new_state_dict = {} - orig_keys = list(old_state_dict.keys()) - - def handle_qkv(sds_sd, ait_sd, sds_key, ait_keys, dims=None): - down_weight = sds_sd.pop(sds_key) - up_weight = sds_sd.pop(sds_key.replace(".down.weight", ".up.weight")) - - # calculate dims if not provided - num_splits = len(ait_keys) - if dims is None: - dims = [up_weight.shape[0] // num_splits] * num_splits - else: - assert sum(dims) == up_weight.shape[0] - - # make ai-toolkit weight - ait_down_keys = [k + ".lora_A.weight" for k in ait_keys] - ait_up_keys = [k + ".lora_B.weight" for k in ait_keys] - - # down_weight is copied to each split - ait_sd.update(dict.fromkeys(ait_down_keys, down_weight)) - - # up_weight is split to each split - ait_sd.update({k: v for k, v in zip(ait_up_keys, torch.split(up_weight, dims, dim=0))}) # noqa: C416 - - for old_key in orig_keys: - # Handle double_blocks - if old_key.startswith(("diffusion_model.double_blocks", "double_blocks")): - block_num = re.search(r"double_blocks\.(\d+)", old_key).group(1) - new_key = f"transformer.transformer_blocks.{block_num}" - - if "processor.proj_lora1" in old_key: - new_key += ".attn.to_out.0" - elif "processor.proj_lora2" in old_key: - new_key += ".attn.to_add_out" - # Handle text latents. - elif "processor.qkv_lora2" in old_key and "up" not in old_key: - handle_qkv( - old_state_dict, - new_state_dict, - old_key, - [ - f"transformer.transformer_blocks.{block_num}.attn.add_q_proj", - f"transformer.transformer_blocks.{block_num}.attn.add_k_proj", - f"transformer.transformer_blocks.{block_num}.attn.add_v_proj", - ], - ) - # continue - # Handle image latents. - elif "processor.qkv_lora1" in old_key and "up" not in old_key: - handle_qkv( - old_state_dict, - new_state_dict, - old_key, - [ - f"transformer.transformer_blocks.{block_num}.attn.to_q", - f"transformer.transformer_blocks.{block_num}.attn.to_k", - f"transformer.transformer_blocks.{block_num}.attn.to_v", - ], - ) - # continue - - if "down" in old_key: - new_key += ".lora_A.weight" - elif "up" in old_key: - new_key += ".lora_B.weight" - - # Handle single_blocks - elif old_key.startswith(("diffusion_model.single_blocks", "single_blocks")): - block_num = re.search(r"single_blocks\.(\d+)", old_key).group(1) - new_key = f"transformer.single_transformer_blocks.{block_num}" - - if "proj_lora" in old_key: - new_key += ".proj_out" - elif "qkv_lora" in old_key and "up" not in old_key: - handle_qkv( - old_state_dict, - new_state_dict, - old_key, - [ - f"transformer.single_transformer_blocks.{block_num}.attn.to_q", - f"transformer.single_transformer_blocks.{block_num}.attn.to_k", - f"transformer.single_transformer_blocks.{block_num}.attn.to_v", - ], - ) - - if "down" in old_key: - new_key += ".lora_A.weight" - elif "up" in old_key: - new_key += ".lora_B.weight" - - else: - # Handle other potential key patterns here - new_key = old_key - - # Since we already handle qkv above. - if "qkv" not in old_key: - new_state_dict[new_key] = old_state_dict.pop(old_key) - - if len(old_state_dict) > 0: - raise ValueError(f"`old_state_dict` should be at this point but has: {list(old_state_dict.keys())}.") - - return new_state_dict - - -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 - pattern = "(" + "|".join(re.escape(sub) for sub in substrings) + ")" - - match = re.search(pattern, key) - if match: - start_sub = match.start() - if start_sub > 0 and key[start_sub - 1] == ".": - boundary = start_sub - 1 - else: - boundary = start_sub - left = key[:boundary].replace(".", "_") - right = key[boundary:] - return left + right - else: - return key.replace(".", "_") - - -def _convert_bfl_flux_control_lora_to_diffusers(original_state_dict): - converted_state_dict = {} - original_state_dict_keys = list(original_state_dict.keys()) - num_layers = 19 - num_single_layers = 38 - inner_dim = 3072 - mlp_ratio = 4.0 - - for lora_key in ["lora_A", "lora_B"]: - ## time_text_embed.timestep_embedder <- time_in - converted_state_dict[f"time_text_embed.timestep_embedder.linear_1.{lora_key}.weight"] = ( - original_state_dict.pop(f"time_in.in_layer.{lora_key}.weight") - ) - if f"time_in.in_layer.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"time_text_embed.timestep_embedder.linear_1.{lora_key}.bias"] = ( - original_state_dict.pop(f"time_in.in_layer.{lora_key}.bias") - ) - - converted_state_dict[f"time_text_embed.timestep_embedder.linear_2.{lora_key}.weight"] = ( - original_state_dict.pop(f"time_in.out_layer.{lora_key}.weight") - ) - if f"time_in.out_layer.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"time_text_embed.timestep_embedder.linear_2.{lora_key}.bias"] = ( - original_state_dict.pop(f"time_in.out_layer.{lora_key}.bias") - ) - - ## time_text_embed.text_embedder <- vector_in - converted_state_dict[f"time_text_embed.text_embedder.linear_1.{lora_key}.weight"] = original_state_dict.pop( - f"vector_in.in_layer.{lora_key}.weight" - ) - if f"vector_in.in_layer.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"time_text_embed.text_embedder.linear_1.{lora_key}.bias"] = original_state_dict.pop( - f"vector_in.in_layer.{lora_key}.bias" - ) - - converted_state_dict[f"time_text_embed.text_embedder.linear_2.{lora_key}.weight"] = original_state_dict.pop( - f"vector_in.out_layer.{lora_key}.weight" - ) - if f"vector_in.out_layer.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"time_text_embed.text_embedder.linear_2.{lora_key}.bias"] = original_state_dict.pop( - f"vector_in.out_layer.{lora_key}.bias" - ) - - # guidance - has_guidance = any("guidance" in k for k in original_state_dict) - if has_guidance: - converted_state_dict[f"time_text_embed.guidance_embedder.linear_1.{lora_key}.weight"] = ( - original_state_dict.pop(f"guidance_in.in_layer.{lora_key}.weight") - ) - if f"guidance_in.in_layer.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"time_text_embed.guidance_embedder.linear_1.{lora_key}.bias"] = ( - original_state_dict.pop(f"guidance_in.in_layer.{lora_key}.bias") - ) - - converted_state_dict[f"time_text_embed.guidance_embedder.linear_2.{lora_key}.weight"] = ( - original_state_dict.pop(f"guidance_in.out_layer.{lora_key}.weight") - ) - if f"guidance_in.out_layer.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"time_text_embed.guidance_embedder.linear_2.{lora_key}.bias"] = ( - original_state_dict.pop(f"guidance_in.out_layer.{lora_key}.bias") - ) - - # context_embedder - converted_state_dict[f"context_embedder.{lora_key}.weight"] = original_state_dict.pop( - f"txt_in.{lora_key}.weight" - ) - if f"txt_in.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"context_embedder.{lora_key}.bias"] = original_state_dict.pop( - f"txt_in.{lora_key}.bias" - ) - - # x_embedder - converted_state_dict[f"x_embedder.{lora_key}.weight"] = original_state_dict.pop(f"img_in.{lora_key}.weight") - if f"img_in.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"x_embedder.{lora_key}.bias"] = original_state_dict.pop(f"img_in.{lora_key}.bias") - - # double transformer blocks - for i in range(num_layers): - block_prefix = f"transformer_blocks.{i}." - - for lora_key in ["lora_A", "lora_B"]: - # norms - converted_state_dict[f"{block_prefix}norm1.linear.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_mod.lin.{lora_key}.weight" - ) - if f"double_blocks.{i}.img_mod.lin.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"{block_prefix}norm1.linear.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_mod.lin.{lora_key}.bias" - ) - - converted_state_dict[f"{block_prefix}norm1_context.linear.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mod.lin.{lora_key}.weight" - ) - if f"double_blocks.{i}.txt_mod.lin.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"{block_prefix}norm1_context.linear.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mod.lin.{lora_key}.bias" - ) - - # Q, K, V - if lora_key == "lora_A": - sample_lora_weight = original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.{lora_key}.weight") - converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([sample_lora_weight]) - converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([sample_lora_weight]) - converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([sample_lora_weight]) - - context_lora_weight = original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.{lora_key}.weight") - converted_state_dict[f"{block_prefix}attn.add_q_proj.{lora_key}.weight"] = torch.cat( - [context_lora_weight] - ) - converted_state_dict[f"{block_prefix}attn.add_k_proj.{lora_key}.weight"] = torch.cat( - [context_lora_weight] - ) - converted_state_dict[f"{block_prefix}attn.add_v_proj.{lora_key}.weight"] = torch.cat( - [context_lora_weight] - ) - else: - sample_q, sample_k, sample_v = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.{lora_key}.weight"), 3, dim=0 - ) - converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([sample_q]) - converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([sample_k]) - converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([sample_v]) - - context_q, context_k, context_v = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.{lora_key}.weight"), 3, dim=0 - ) - converted_state_dict[f"{block_prefix}attn.add_q_proj.{lora_key}.weight"] = torch.cat([context_q]) - converted_state_dict[f"{block_prefix}attn.add_k_proj.{lora_key}.weight"] = torch.cat([context_k]) - converted_state_dict[f"{block_prefix}attn.add_v_proj.{lora_key}.weight"] = torch.cat([context_v]) - - if f"double_blocks.{i}.img_attn.qkv.{lora_key}.bias" in original_state_dict_keys: - sample_q_bias, sample_k_bias, sample_v_bias = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.{lora_key}.bias"), 3, dim=0 - ) - converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.bias"] = torch.cat([sample_q_bias]) - converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.bias"] = torch.cat([sample_k_bias]) - converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.bias"] = torch.cat([sample_v_bias]) - - if f"double_blocks.{i}.txt_attn.qkv.{lora_key}.bias" in original_state_dict_keys: - context_q_bias, context_k_bias, context_v_bias = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.{lora_key}.bias"), 3, dim=0 - ) - converted_state_dict[f"{block_prefix}attn.add_q_proj.{lora_key}.bias"] = torch.cat([context_q_bias]) - converted_state_dict[f"{block_prefix}attn.add_k_proj.{lora_key}.bias"] = torch.cat([context_k_bias]) - converted_state_dict[f"{block_prefix}attn.add_v_proj.{lora_key}.bias"] = torch.cat([context_v_bias]) - - # ff img_mlp - converted_state_dict[f"{block_prefix}ff.net.0.proj.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.0.{lora_key}.weight" - ) - if f"double_blocks.{i}.img_mlp.0.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"{block_prefix}ff.net.0.proj.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.0.{lora_key}.bias" - ) - - converted_state_dict[f"{block_prefix}ff.net.2.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.2.{lora_key}.weight" - ) - if f"double_blocks.{i}.img_mlp.2.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"{block_prefix}ff.net.2.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.2.{lora_key}.bias" - ) - - converted_state_dict[f"{block_prefix}ff_context.net.0.proj.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.0.{lora_key}.weight" - ) - if f"double_blocks.{i}.txt_mlp.0.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"{block_prefix}ff_context.net.0.proj.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.0.{lora_key}.bias" - ) - - converted_state_dict[f"{block_prefix}ff_context.net.2.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.2.{lora_key}.weight" - ) - if f"double_blocks.{i}.txt_mlp.2.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"{block_prefix}ff_context.net.2.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.2.{lora_key}.bias" - ) - - # output projections. - converted_state_dict[f"{block_prefix}attn.to_out.0.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn.proj.{lora_key}.weight" - ) - if f"double_blocks.{i}.img_attn.proj.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"{block_prefix}attn.to_out.0.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn.proj.{lora_key}.bias" - ) - converted_state_dict[f"{block_prefix}attn.to_add_out.{lora_key}.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn.proj.{lora_key}.weight" - ) - if f"double_blocks.{i}.txt_attn.proj.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"{block_prefix}attn.to_add_out.{lora_key}.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn.proj.{lora_key}.bias" - ) - - # qk_norm - converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn.norm.query_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn.norm.key_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_added_q.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn.norm.query_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_added_k.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn.norm.key_norm.scale" - ) - - # single transformer blocks - for i in range(num_single_layers): - block_prefix = f"single_transformer_blocks.{i}." - - for lora_key in ["lora_A", "lora_B"]: - # norm.linear <- single_blocks.0.modulation.lin - converted_state_dict[f"{block_prefix}norm.linear.{lora_key}.weight"] = original_state_dict.pop( - f"single_blocks.{i}.modulation.lin.{lora_key}.weight" - ) - if f"single_blocks.{i}.modulation.lin.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"{block_prefix}norm.linear.{lora_key}.bias"] = original_state_dict.pop( - f"single_blocks.{i}.modulation.lin.{lora_key}.bias" - ) - - # Q, K, V, mlp - mlp_hidden_dim = int(inner_dim * mlp_ratio) - split_size = (inner_dim, inner_dim, inner_dim, mlp_hidden_dim) - - if lora_key == "lora_A": - lora_weight = original_state_dict.pop(f"single_blocks.{i}.linear1.{lora_key}.weight") - converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([lora_weight]) - converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([lora_weight]) - converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([lora_weight]) - converted_state_dict[f"{block_prefix}proj_mlp.{lora_key}.weight"] = torch.cat([lora_weight]) - - if f"single_blocks.{i}.linear1.{lora_key}.bias" in original_state_dict_keys: - lora_bias = original_state_dict.pop(f"single_blocks.{i}.linear1.{lora_key}.bias") - converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.bias"] = torch.cat([lora_bias]) - converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.bias"] = torch.cat([lora_bias]) - converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.bias"] = torch.cat([lora_bias]) - converted_state_dict[f"{block_prefix}proj_mlp.{lora_key}.bias"] = torch.cat([lora_bias]) - else: - q, k, v, mlp = torch.split( - original_state_dict.pop(f"single_blocks.{i}.linear1.{lora_key}.weight"), split_size, dim=0 - ) - converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.weight"] = torch.cat([q]) - converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.weight"] = torch.cat([k]) - converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.weight"] = torch.cat([v]) - converted_state_dict[f"{block_prefix}proj_mlp.{lora_key}.weight"] = torch.cat([mlp]) - - if f"single_blocks.{i}.linear1.{lora_key}.bias" in original_state_dict_keys: - q_bias, k_bias, v_bias, mlp_bias = torch.split( - original_state_dict.pop(f"single_blocks.{i}.linear1.{lora_key}.bias"), split_size, dim=0 - ) - converted_state_dict[f"{block_prefix}attn.to_q.{lora_key}.bias"] = torch.cat([q_bias]) - converted_state_dict[f"{block_prefix}attn.to_k.{lora_key}.bias"] = torch.cat([k_bias]) - converted_state_dict[f"{block_prefix}attn.to_v.{lora_key}.bias"] = torch.cat([v_bias]) - converted_state_dict[f"{block_prefix}proj_mlp.{lora_key}.bias"] = torch.cat([mlp_bias]) - - # output projections. - converted_state_dict[f"{block_prefix}proj_out.{lora_key}.weight"] = original_state_dict.pop( - f"single_blocks.{i}.linear2.{lora_key}.weight" - ) - if f"single_blocks.{i}.linear2.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"{block_prefix}proj_out.{lora_key}.bias"] = original_state_dict.pop( - f"single_blocks.{i}.linear2.{lora_key}.bias" - ) - - # qk norm - converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = original_state_dict.pop( - f"single_blocks.{i}.norm.query_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = original_state_dict.pop( - f"single_blocks.{i}.norm.key_norm.scale" - ) - - for lora_key in ["lora_A", "lora_B"]: - converted_state_dict[f"proj_out.{lora_key}.weight"] = original_state_dict.pop( - f"final_layer.linear.{lora_key}.weight" - ) - if f"final_layer.linear.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"proj_out.{lora_key}.bias"] = original_state_dict.pop( - f"final_layer.linear.{lora_key}.bias" - ) - - converted_state_dict[f"norm_out.linear.{lora_key}.weight"] = swap_scale_shift( - original_state_dict.pop(f"final_layer.adaLN_modulation.1.{lora_key}.weight") - ) - if f"final_layer.adaLN_modulation.1.{lora_key}.bias" in original_state_dict_keys: - converted_state_dict[f"norm_out.linear.{lora_key}.bias"] = swap_scale_shift( - original_state_dict.pop(f"final_layer.adaLN_modulation.1.{lora_key}.bias") - ) - - if len(original_state_dict) > 0: - raise ValueError(f"`original_state_dict` should be empty at this point but has {original_state_dict.keys()=}.") - - for key in list(converted_state_dict.keys()): - converted_state_dict[f"transformer.{key}"] = converted_state_dict.pop(key) - - return converted_state_dict - - -def _convert_hunyuan_video_lora_to_diffusers(original_state_dict): - converted_state_dict = {k: original_state_dict.pop(k) for k in list(original_state_dict.keys())} - - def remap_norm_scale_shift_(key, state_dict): - weight = state_dict.pop(key) - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - state_dict[key.replace("final_layer.adaLN_modulation.1", "norm_out.linear")] = new_weight - - def remap_txt_in_(key, state_dict): - def rename_key(key): - new_key = key.replace("individual_token_refiner.blocks", "token_refiner.refiner_blocks") - new_key = new_key.replace("adaLN_modulation.1", "norm_out.linear") - new_key = new_key.replace("txt_in", "context_embedder") - new_key = new_key.replace("t_embedder.mlp.0", "time_text_embed.timestep_embedder.linear_1") - new_key = new_key.replace("t_embedder.mlp.2", "time_text_embed.timestep_embedder.linear_2") - new_key = new_key.replace("c_embedder", "time_text_embed.text_embedder") - new_key = new_key.replace("mlp", "ff") - return new_key - - if "self_attn_qkv" in key: - weight = state_dict.pop(key) - to_q, to_k, to_v = weight.chunk(3, dim=0) - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_q"))] = to_q - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_k"))] = to_k - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_v"))] = to_v - else: - state_dict[rename_key(key)] = state_dict.pop(key) - - def remap_img_attn_qkv_(key, state_dict): - weight = state_dict.pop(key) - if "lora_A" in key: - state_dict[key.replace("img_attn_qkv", "attn.to_q")] = weight - state_dict[key.replace("img_attn_qkv", "attn.to_k")] = weight - state_dict[key.replace("img_attn_qkv", "attn.to_v")] = weight - else: - to_q, to_k, to_v = weight.chunk(3, dim=0) - state_dict[key.replace("img_attn_qkv", "attn.to_q")] = to_q - state_dict[key.replace("img_attn_qkv", "attn.to_k")] = to_k - state_dict[key.replace("img_attn_qkv", "attn.to_v")] = to_v - - def remap_txt_attn_qkv_(key, state_dict): - weight = state_dict.pop(key) - if "lora_A" in key: - state_dict[key.replace("txt_attn_qkv", "attn.add_q_proj")] = weight - state_dict[key.replace("txt_attn_qkv", "attn.add_k_proj")] = weight - state_dict[key.replace("txt_attn_qkv", "attn.add_v_proj")] = weight - else: - to_q, to_k, to_v = weight.chunk(3, dim=0) - state_dict[key.replace("txt_attn_qkv", "attn.add_q_proj")] = to_q - state_dict[key.replace("txt_attn_qkv", "attn.add_k_proj")] = to_k - state_dict[key.replace("txt_attn_qkv", "attn.add_v_proj")] = to_v - - def remap_single_transformer_blocks_(key, state_dict): - hidden_size = 3072 - - if "linear1.lora_A.weight" in key or "linear1.lora_B.weight" in key: - linear1_weight = state_dict.pop(key) - if "lora_A" in key: - new_key = key.replace("single_blocks", "single_transformer_blocks").removesuffix( - ".linear1.lora_A.weight" - ) - state_dict[f"{new_key}.attn.to_q.lora_A.weight"] = linear1_weight - state_dict[f"{new_key}.attn.to_k.lora_A.weight"] = linear1_weight - state_dict[f"{new_key}.attn.to_v.lora_A.weight"] = linear1_weight - state_dict[f"{new_key}.proj_mlp.lora_A.weight"] = linear1_weight - else: - split_size = (hidden_size, hidden_size, hidden_size, linear1_weight.size(0) - 3 * hidden_size) - q, k, v, mlp = torch.split(linear1_weight, split_size, dim=0) - new_key = key.replace("single_blocks", "single_transformer_blocks").removesuffix( - ".linear1.lora_B.weight" - ) - state_dict[f"{new_key}.attn.to_q.lora_B.weight"] = q - state_dict[f"{new_key}.attn.to_k.lora_B.weight"] = k - state_dict[f"{new_key}.attn.to_v.lora_B.weight"] = v - state_dict[f"{new_key}.proj_mlp.lora_B.weight"] = mlp - - elif "linear1.lora_A.bias" in key or "linear1.lora_B.bias" in key: - linear1_bias = state_dict.pop(key) - if "lora_A" in key: - new_key = key.replace("single_blocks", "single_transformer_blocks").removesuffix( - ".linear1.lora_A.bias" - ) - state_dict[f"{new_key}.attn.to_q.lora_A.bias"] = linear1_bias - state_dict[f"{new_key}.attn.to_k.lora_A.bias"] = linear1_bias - state_dict[f"{new_key}.attn.to_v.lora_A.bias"] = linear1_bias - state_dict[f"{new_key}.proj_mlp.lora_A.bias"] = linear1_bias - else: - split_size = (hidden_size, hidden_size, hidden_size, linear1_bias.size(0) - 3 * hidden_size) - q_bias, k_bias, v_bias, mlp_bias = torch.split(linear1_bias, split_size, dim=0) - new_key = key.replace("single_blocks", "single_transformer_blocks").removesuffix( - ".linear1.lora_B.bias" - ) - state_dict[f"{new_key}.attn.to_q.lora_B.bias"] = q_bias - state_dict[f"{new_key}.attn.to_k.lora_B.bias"] = k_bias - state_dict[f"{new_key}.attn.to_v.lora_B.bias"] = v_bias - state_dict[f"{new_key}.proj_mlp.lora_B.bias"] = mlp_bias - - else: - new_key = key.replace("single_blocks", "single_transformer_blocks") - new_key = new_key.replace("linear2", "proj_out") - new_key = new_key.replace("q_norm", "attn.norm_q") - new_key = new_key.replace("k_norm", "attn.norm_k") - state_dict[new_key] = state_dict.pop(key) - - TRANSFORMER_KEYS_RENAME_DICT = { - "img_in": "x_embedder", - "time_in.mlp.0": "time_text_embed.timestep_embedder.linear_1", - "time_in.mlp.2": "time_text_embed.timestep_embedder.linear_2", - "guidance_in.mlp.0": "time_text_embed.guidance_embedder.linear_1", - "guidance_in.mlp.2": "time_text_embed.guidance_embedder.linear_2", - "vector_in.in_layer": "time_text_embed.text_embedder.linear_1", - "vector_in.out_layer": "time_text_embed.text_embedder.linear_2", - "double_blocks": "transformer_blocks", - "img_attn_q_norm": "attn.norm_q", - "img_attn_k_norm": "attn.norm_k", - "img_attn_proj": "attn.to_out.0", - "txt_attn_q_norm": "attn.norm_added_q", - "txt_attn_k_norm": "attn.norm_added_k", - "txt_attn_proj": "attn.to_add_out", - "img_mod.linear": "norm1.linear", - "img_norm1": "norm1.norm", - "img_norm2": "norm2", - "img_mlp": "ff", - "txt_mod.linear": "norm1_context.linear", - "txt_norm1": "norm1.norm", - "txt_norm2": "norm2_context", - "txt_mlp": "ff_context", - "self_attn_proj": "attn.to_out.0", - "modulation.linear": "norm.linear", - "pre_norm": "norm.norm", - "final_layer.norm_final": "norm_out.norm", - "final_layer.linear": "proj_out", - "fc1": "net.0.proj", - "fc2": "net.2", - "input_embedder": "proj_in", - } - - TRANSFORMER_SPECIAL_KEYS_REMAP = { - "txt_in": remap_txt_in_, - "img_attn_qkv": remap_img_attn_qkv_, - "txt_attn_qkv": remap_txt_attn_qkv_, - "single_blocks": remap_single_transformer_blocks_, - "final_layer.adaLN_modulation.1": remap_norm_scale_shift_, - } - - # Some folks attempt to make their state dict compatible with diffusers by adding "transformer." prefix to all keys - # and use their custom code. To make sure both "original" and "attempted diffusers" loras work as expected, we make - # sure that both follow the same initial format by stripping off the "transformer." prefix. - for key in list(converted_state_dict.keys()): - if key.startswith("transformer."): - converted_state_dict[key[len("transformer.") :]] = converted_state_dict.pop(key) - if key.startswith("diffusion_model."): - converted_state_dict[key[len("diffusion_model.") :]] = converted_state_dict.pop(key) - - # Rename and remap the state dict keys - for key in list(converted_state_dict.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] = converted_state_dict.pop(key) - - for key in list(converted_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, converted_state_dict) - - # Add back the "transformer." prefix - for key in list(converted_state_dict.keys()): - converted_state_dict[f"transformer.{key}"] = converted_state_dict.pop(key) - - return converted_state_dict - - -def _convert_non_diffusers_lumina2_lora_to_diffusers(state_dict): - # Remove "diffusion_model." prefix from keys. - state_dict = {k[len("diffusion_model.") :]: v for k, v in state_dict.items()} - converted_state_dict = {} - - def get_num_layers(keys, pattern): - layers = set() - for key in keys: - match = re.search(pattern, key) - if match: - layers.add(int(match.group(1))) - return len(layers) - - def process_block(prefix, index, convert_norm): - # Process attention qkv: pop lora_A and lora_B weights. - lora_down = state_dict.pop(f"{prefix}.{index}.attention.qkv.lora_A.weight") - lora_up = state_dict.pop(f"{prefix}.{index}.attention.qkv.lora_B.weight") - for attn_key in ["to_q", "to_k", "to_v"]: - converted_state_dict[f"{prefix}.{index}.attn.{attn_key}.lora_A.weight"] = lora_down - for attn_key, weight in zip(["to_q", "to_k", "to_v"], torch.split(lora_up, [2304, 768, 768], dim=0)): - converted_state_dict[f"{prefix}.{index}.attn.{attn_key}.lora_B.weight"] = weight - - # Process attention out weights. - converted_state_dict[f"{prefix}.{index}.attn.to_out.0.lora_A.weight"] = state_dict.pop( - f"{prefix}.{index}.attention.out.lora_A.weight" - ) - converted_state_dict[f"{prefix}.{index}.attn.to_out.0.lora_B.weight"] = state_dict.pop( - f"{prefix}.{index}.attention.out.lora_B.weight" - ) - - # Process feed-forward weights for layers 1, 2, and 3. - for layer in range(1, 4): - converted_state_dict[f"{prefix}.{index}.feed_forward.linear_{layer}.lora_A.weight"] = state_dict.pop( - f"{prefix}.{index}.feed_forward.w{layer}.lora_A.weight" - ) - converted_state_dict[f"{prefix}.{index}.feed_forward.linear_{layer}.lora_B.weight"] = state_dict.pop( - f"{prefix}.{index}.feed_forward.w{layer}.lora_B.weight" - ) - - if convert_norm: - converted_state_dict[f"{prefix}.{index}.norm1.linear.lora_A.weight"] = state_dict.pop( - f"{prefix}.{index}.adaLN_modulation.1.lora_A.weight" - ) - converted_state_dict[f"{prefix}.{index}.norm1.linear.lora_B.weight"] = state_dict.pop( - f"{prefix}.{index}.adaLN_modulation.1.lora_B.weight" - ) - - noise_refiner_pattern = r"noise_refiner\.(\d+)\." - num_noise_refiner_layers = get_num_layers(state_dict.keys(), noise_refiner_pattern) - for i in range(num_noise_refiner_layers): - process_block("noise_refiner", i, convert_norm=True) - - context_refiner_pattern = r"context_refiner\.(\d+)\." - num_context_refiner_layers = get_num_layers(state_dict.keys(), context_refiner_pattern) - for i in range(num_context_refiner_layers): - process_block("context_refiner", i, convert_norm=False) - - core_transformer_pattern = r"layers\.(\d+)\." - num_core_transformer_layers = get_num_layers(state_dict.keys(), core_transformer_pattern) - for i in range(num_core_transformer_layers): - process_block("layers", i, convert_norm=True) - - if len(state_dict) > 0: - raise ValueError(f"`state_dict` should be empty at this point but has {state_dict.keys()=}") - - for key in list(converted_state_dict.keys()): - converted_state_dict[f"transformer.{key}"] = converted_state_dict.pop(key) - - return converted_state_dict - - -def _convert_non_diffusers_wan_lora_to_diffusers(state_dict): - converted_state_dict = {} - original_state_dict = {k[len("diffusion_model.") :]: v for k, v in state_dict.items()} - - block_numbers = {int(k.split(".")[1]) for k in original_state_dict if k.startswith("blocks.")} - min_block = min(block_numbers) - max_block = max(block_numbers) - - is_i2v_lora = any("k_img" in k for k in original_state_dict) and any("v_img" in k for k in original_state_dict) - lora_down_key = "lora_A" if any("lora_A" in k for k in original_state_dict) else "lora_down" - lora_up_key = "lora_B" if any("lora_B" in k for k in original_state_dict) else "lora_up" - - diff_keys = [k for k in original_state_dict if k.endswith((".diff_b", ".diff"))] - if diff_keys: - for diff_k in diff_keys: - param = original_state_dict[diff_k] - # The magnitudes of the .diff-ending weights are very low (most are below 1e-4, some are upto 1e-3, - # and 2 of them are about 1.6e-2 [the case with AccVideo lora]). The low magnitudes mostly correspond - # to norm layers. Ignoring them is the best option at the moment until a better solution is found. It - # is okay to ignore because they do not affect the model output in a significant manner. - threshold = 1.6e-2 - absdiff = param.abs().max() - param.abs().min() - all_zero = torch.all(param == 0).item() - all_absdiff_lower_than_threshold = absdiff < threshold - if all_zero or all_absdiff_lower_than_threshold: - logger.debug( - f"Removed {diff_k} key from the state dict as it's all zeros, or values lower than hardcoded threshold." - ) - original_state_dict.pop(diff_k) - - # For the `diff_b` keys, we treat them as lora_bias. - # https://huggingface.co/docs/peft/main/en/package_reference/lora#peft.LoraConfig.lora_bias - - for i in range(min_block, max_block + 1): - # Self-attention - for o, c in zip(["q", "k", "v", "o"], ["to_q", "to_k", "to_v", "to_out.0"]): - original_key = f"blocks.{i}.self_attn.{o}.{lora_down_key}.weight" - converted_key = f"blocks.{i}.attn1.{c}.lora_A.weight" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - original_key = f"blocks.{i}.self_attn.{o}.{lora_up_key}.weight" - converted_key = f"blocks.{i}.attn1.{c}.lora_B.weight" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - original_key = f"blocks.{i}.self_attn.{o}.diff_b" - converted_key = f"blocks.{i}.attn1.{c}.lora_B.bias" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - # Cross-attention - for o, c in zip(["q", "k", "v", "o"], ["to_q", "to_k", "to_v", "to_out.0"]): - original_key = f"blocks.{i}.cross_attn.{o}.{lora_down_key}.weight" - converted_key = f"blocks.{i}.attn2.{c}.lora_A.weight" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - original_key = f"blocks.{i}.cross_attn.{o}.{lora_up_key}.weight" - converted_key = f"blocks.{i}.attn2.{c}.lora_B.weight" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - original_key = f"blocks.{i}.cross_attn.{o}.diff_b" - converted_key = f"blocks.{i}.attn2.{c}.lora_B.bias" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - if is_i2v_lora: - for o, c in zip(["k_img", "v_img"], ["add_k_proj", "add_v_proj"]): - original_key = f"blocks.{i}.cross_attn.{o}.{lora_down_key}.weight" - converted_key = f"blocks.{i}.attn2.{c}.lora_A.weight" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - original_key = f"blocks.{i}.cross_attn.{o}.{lora_up_key}.weight" - converted_key = f"blocks.{i}.attn2.{c}.lora_B.weight" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - original_key = f"blocks.{i}.cross_attn.{o}.diff_b" - converted_key = f"blocks.{i}.attn2.{c}.lora_B.bias" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - # FFN - for o, c in zip(["ffn.0", "ffn.2"], ["net.0.proj", "net.2"]): - original_key = f"blocks.{i}.{o}.{lora_down_key}.weight" - converted_key = f"blocks.{i}.ffn.{c}.lora_A.weight" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - original_key = f"blocks.{i}.{o}.{lora_up_key}.weight" - converted_key = f"blocks.{i}.ffn.{c}.lora_B.weight" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - original_key = f"blocks.{i}.{o}.diff_b" - converted_key = f"blocks.{i}.ffn.{c}.lora_B.bias" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - # Remaining. - if original_state_dict: - if any("time_projection" in k for k in original_state_dict): - original_key = f"time_projection.1.{lora_down_key}.weight" - converted_key = "condition_embedder.time_proj.lora_A.weight" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - original_key = f"time_projection.1.{lora_up_key}.weight" - converted_key = "condition_embedder.time_proj.lora_B.weight" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - if "time_projection.1.diff_b" in original_state_dict: - converted_state_dict["condition_embedder.time_proj.lora_B.bias"] = original_state_dict.pop( - "time_projection.1.diff_b" - ) - - if any("head.head" in k for k in state_dict): - converted_state_dict["proj_out.lora_A.weight"] = original_state_dict.pop( - f"head.head.{lora_down_key}.weight" - ) - converted_state_dict["proj_out.lora_B.weight"] = original_state_dict.pop(f"head.head.{lora_up_key}.weight") - if "head.head.diff_b" in original_state_dict: - converted_state_dict["proj_out.lora_B.bias"] = original_state_dict.pop("head.head.diff_b") - - for text_time in ["text_embedding", "time_embedding"]: - if any(text_time in k for k in original_state_dict): - for b_n in [0, 2]: - diffusers_b_n = 1 if b_n == 0 else 2 - diffusers_name = ( - "condition_embedder.text_embedder" - if text_time == "text_embedding" - else "condition_embedder.time_embedder" - ) - if any(f"{text_time}.{b_n}" in k for k in original_state_dict): - converted_state_dict[f"{diffusers_name}.linear_{diffusers_b_n}.lora_A.weight"] = ( - original_state_dict.pop(f"{text_time}.{b_n}.{lora_down_key}.weight") - ) - converted_state_dict[f"{diffusers_name}.linear_{diffusers_b_n}.lora_B.weight"] = ( - original_state_dict.pop(f"{text_time}.{b_n}.{lora_up_key}.weight") - ) - if f"{text_time}.{b_n}.diff_b" in original_state_dict: - converted_state_dict[f"{diffusers_name}.linear_{diffusers_b_n}.lora_B.bias"] = ( - original_state_dict.pop(f"{text_time}.{b_n}.diff_b") - ) - - for img_ours, img_theirs in [ - ("ff.net.0.proj", "img_emb.proj.1"), - ("ff.net.2", "img_emb.proj.3"), - ]: - original_key = f"{img_theirs}.{lora_down_key}.weight" - converted_key = f"condition_embedder.image_embedder.{img_ours}.lora_A.weight" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - original_key = f"{img_theirs}.{lora_up_key}.weight" - converted_key = f"condition_embedder.image_embedder.{img_ours}.lora_B.weight" - if original_key in original_state_dict: - converted_state_dict[converted_key] = original_state_dict.pop(original_key) - - if len(original_state_dict) > 0: - diff = all(".diff" in k for k in original_state_dict) - if diff: - diff_keys = {k for k in original_state_dict if k.endswith(".diff")} - if not all("lora" not in k for k in diff_keys): - raise ValueError - logger.info( - "The remaining `state_dict` contains `diff` keys which we do not handle yet. If you see performance issues, please file an issue: " - "https://github.com/huggingface/diffusers//issues/new" - ) - else: - raise ValueError(f"`state_dict` should be empty at this point but has {original_state_dict.keys()=}") - - for key in list(converted_state_dict.keys()): - converted_state_dict[f"transformer.{key}"] = converted_state_dict.pop(key) - - return converted_state_dict - - -def _convert_musubi_wan_lora_to_diffusers(state_dict): - # https://github.com/kohya-ss/musubi-tuner - converted_state_dict = {} - original_state_dict = {k[len("lora_unet_") :]: v for k, v in state_dict.items()} - - num_blocks = len({k.split("blocks_")[1].split("_")[0] for k in original_state_dict}) - is_i2v_lora = any("k_img" in k for k in original_state_dict) and any("v_img" in k for k in original_state_dict) - - def get_alpha_scales(down_weight, key): - rank = down_weight.shape[0] - alpha = original_state_dict.pop(key + ".alpha").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 - - for i in range(num_blocks): - # Self-attention - for o, c in zip(["q", "k", "v", "o"], ["to_q", "to_k", "to_v", "to_out.0"]): - down_weight = original_state_dict.pop(f"blocks_{i}_self_attn_{o}.lora_down.weight") - up_weight = original_state_dict.pop(f"blocks_{i}_self_attn_{o}.lora_up.weight") - scale_down, scale_up = get_alpha_scales(down_weight, f"blocks_{i}_self_attn_{o}") - converted_state_dict[f"blocks.{i}.attn1.{c}.lora_A.weight"] = down_weight * scale_down - converted_state_dict[f"blocks.{i}.attn1.{c}.lora_B.weight"] = up_weight * scale_up - - # Cross-attention - for o, c in zip(["q", "k", "v", "o"], ["to_q", "to_k", "to_v", "to_out.0"]): - down_weight = original_state_dict.pop(f"blocks_{i}_cross_attn_{o}.lora_down.weight") - up_weight = original_state_dict.pop(f"blocks_{i}_cross_attn_{o}.lora_up.weight") - scale_down, scale_up = get_alpha_scales(down_weight, f"blocks_{i}_cross_attn_{o}") - converted_state_dict[f"blocks.{i}.attn2.{c}.lora_A.weight"] = down_weight * scale_down - converted_state_dict[f"blocks.{i}.attn2.{c}.lora_B.weight"] = up_weight * scale_up - - if is_i2v_lora: - for o, c in zip(["k_img", "v_img"], ["add_k_proj", "add_v_proj"]): - down_weight = original_state_dict.pop(f"blocks_{i}_cross_attn_{o}.lora_down.weight") - up_weight = original_state_dict.pop(f"blocks_{i}_cross_attn_{o}.lora_up.weight") - scale_down, scale_up = get_alpha_scales(down_weight, f"blocks_{i}_cross_attn_{o}") - converted_state_dict[f"blocks.{i}.attn2.{c}.lora_A.weight"] = down_weight * scale_down - converted_state_dict[f"blocks.{i}.attn2.{c}.lora_B.weight"] = up_weight * scale_up - - # FFN - for o, c in zip(["ffn_0", "ffn_2"], ["net.0.proj", "net.2"]): - down_weight = original_state_dict.pop(f"blocks_{i}_{o}.lora_down.weight") - up_weight = original_state_dict.pop(f"blocks_{i}_{o}.lora_up.weight") - scale_down, scale_up = get_alpha_scales(down_weight, f"blocks_{i}_{o}") - converted_state_dict[f"blocks.{i}.ffn.{c}.lora_A.weight"] = down_weight * scale_down - converted_state_dict[f"blocks.{i}.ffn.{c}.lora_B.weight"] = up_weight * scale_up - - if len(original_state_dict) > 0: - raise ValueError(f"`state_dict` should be empty at this point but has {original_state_dict.keys()=}") - - for key in list(converted_state_dict.keys()): - converted_state_dict[f"transformer.{key}"] = converted_state_dict.pop(key) - - return converted_state_dict - - -def _convert_non_diffusers_hidream_lora_to_diffusers(state_dict, non_diffusers_prefix="diffusion_model"): - if not all(k.startswith(non_diffusers_prefix) for k in state_dict): - raise ValueError("Invalid LoRA state dict for HiDream.") - converted_state_dict = {k.removeprefix(f"{non_diffusers_prefix}."): v for k, v in state_dict.items()} - converted_state_dict = {f"transformer.{k}": v for k, v in converted_state_dict.items()} - return converted_state_dict - - -def _convert_non_diffusers_ltxv_lora_to_diffusers(state_dict, non_diffusers_prefix="diffusion_model"): - if not all(k.startswith(f"{non_diffusers_prefix}.") for k in state_dict): - raise ValueError("Invalid LoRA state dict for LTX-Video.") - converted_state_dict = {k.removeprefix(f"{non_diffusers_prefix}."): v for k, v in state_dict.items()} - converted_state_dict = {f"transformer.{k}": v for k, v in converted_state_dict.items()} - return converted_state_dict diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/lora_pipeline.py b/venv/lib/python3.10/site-packages/diffusers/loaders/lora_pipeline.py deleted file mode 100644 index 4fea005cbc397455424e47daa74426728fefca3b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/lora_pipeline.py +++ /dev/null @@ -1,6135 +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 os -from typing import Callable, Dict, List, Optional, Union - -import torch -from huggingface_hub.utils import validate_hf_hub_args - -from ..utils import ( - USE_PEFT_BACKEND, - deprecate, - get_submodule_by_name, - is_bitsandbytes_available, - is_gguf_available, - is_peft_available, - is_peft_version, - is_torch_version, - is_transformers_available, - is_transformers_version, - logging, -) -from .lora_base import ( # noqa - LORA_WEIGHT_NAME, - LORA_WEIGHT_NAME_SAFE, - LoraBaseMixin, - _fetch_state_dict, - _load_lora_into_text_encoder, - _pack_dict_with_prefix, -) -from .lora_conversion_utils import ( - _convert_bfl_flux_control_lora_to_diffusers, - _convert_hunyuan_video_lora_to_diffusers, - _convert_kohya_flux_lora_to_diffusers, - _convert_musubi_wan_lora_to_diffusers, - _convert_non_diffusers_hidream_lora_to_diffusers, - _convert_non_diffusers_lora_to_diffusers, - _convert_non_diffusers_ltxv_lora_to_diffusers, - _convert_non_diffusers_lumina2_lora_to_diffusers, - _convert_non_diffusers_wan_lora_to_diffusers, - _convert_xlabs_flux_lora_to_diffusers, - _maybe_map_sgm_blocks_to_diffusers, -) - - -_LOW_CPU_MEM_USAGE_DEFAULT_LORA = False -if is_torch_version(">=", "1.9.0"): - if ( - is_peft_available() - and is_peft_version(">=", "0.13.1") - and is_transformers_available() - and is_transformers_version(">", "4.45.2") - ): - _LOW_CPU_MEM_USAGE_DEFAULT_LORA = True - - -logger = logging.get_logger(__name__) - -TEXT_ENCODER_NAME = "text_encoder" -UNET_NAME = "unet" -TRANSFORMER_NAME = "transformer" - -_MODULE_NAME_TO_ATTRIBUTE_MAP_FLUX = {"x_embedder": "in_channels"} - - -def _maybe_dequantize_weight_for_expanded_lora(model, module): - if is_bitsandbytes_available(): - from ..quantizers.bitsandbytes import dequantize_bnb_weight - - if is_gguf_available(): - from ..quantizers.gguf.utils import dequantize_gguf_tensor - - is_bnb_4bit_quantized = module.weight.__class__.__name__ == "Params4bit" - is_bnb_8bit_quantized = module.weight.__class__.__name__ == "Int8Params" - is_gguf_quantized = module.weight.__class__.__name__ == "GGUFParameter" - - if is_bnb_4bit_quantized and not is_bitsandbytes_available(): - raise ValueError( - "The checkpoint seems to have been quantized with `bitsandbytes` (4bits). Install `bitsandbytes` to load quantized checkpoints." - ) - if is_bnb_8bit_quantized and not is_bitsandbytes_available(): - raise ValueError( - "The checkpoint seems to have been quantized with `bitsandbytes` (8bits). Install `bitsandbytes` to load quantized checkpoints." - ) - if is_gguf_quantized and not is_gguf_available(): - raise ValueError( - "The checkpoint seems to have been quantized with `gguf`. Install `gguf` to load quantized checkpoints." - ) - - weight_on_cpu = False - if module.weight.device.type == "cpu": - weight_on_cpu = True - - device = torch.accelerator.current_accelerator().type if hasattr(torch, "accelerator") else "cuda" - if is_bnb_4bit_quantized or is_bnb_8bit_quantized: - module_weight = dequantize_bnb_weight( - module.weight.to(device) if weight_on_cpu else module.weight, - state=module.weight.quant_state if is_bnb_4bit_quantized else module.state, - dtype=model.dtype, - ).data - elif is_gguf_quantized: - module_weight = dequantize_gguf_tensor( - module.weight.to(device) if weight_on_cpu else module.weight, - ) - module_weight = module_weight.to(model.dtype) - else: - module_weight = module.weight.data - - if weight_on_cpu: - module_weight = module_weight.cpu() - - return module_weight - - -class StableDiffusionLoraLoaderMixin(LoraBaseMixin): - r""" - Load LoRA layers into Stable Diffusion [`UNet2DConditionModel`] and - [`CLIPTextModel`](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel). - """ - - _lora_loadable_modules = ["unet", "text_encoder"] - unet_name = UNET_NAME - text_encoder_name = TEXT_ENCODER_NAME - - def load_lora_weights( - self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, - hotswap: bool = False, - **kwargs, - ): - """Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.unet` and - `self.text_encoder`. - - All kwargs are forwarded to `self.lora_state_dict`. - - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is - loaded. - - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details on how the state dict is - loaded into `self.unet`. - - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder`] for more details on how the state - dict is loaded into `self.text_encoder`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - Defaults to `False`. Whether to substitute an existing (LoRA) adapter with the newly loaded adapter - in-place. This means that, instead of loading an additional adapter, this will take the existing - adapter weights and replace them with the weights of the new adapter. This can be faster and more - memory efficient. However, the main advantage of hotswapping is that when the model is compiled with - torch.compile, loading the new adapter does not require recompilation of the model. When using - hotswapping, the passed `adapter_name` should be the name of an already loaded adapter. - - If the new adapter and the old adapter have different ranks and/or LoRA alphas (i.e. scaling), you need - to call an additional method before loading the adapter: - - ```py - pipeline = ... # load diffusers pipeline - max_rank = ... # the highest rank among all LoRAs that you want to load - # call *before* compiling and loading the LoRA adapter - pipeline.enable_lora_hotswap(target_rank=max_rank) - pipeline.load_lora_weights(file_name) - # optionally compile the model now - ``` - - Note that hotswapping adapters of the text encoder is not yet supported. There are some further - limitations to this technique, which are documented here: - https://huggingface.co/docs/peft/main/en/package_reference/hotswap - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - """ - 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 not is_peft_version(">=", "0.13.1"): - 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, network_alphas, 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.") - - self.load_lora_into_unet( - state_dict, - network_alphas=network_alphas, - unet=getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet, - adapter_name=adapter_name, - metadata=metadata, - _pipeline=self, - low_cpu_mem_usage=low_cpu_mem_usage, - hotswap=hotswap, - ) - self.load_lora_into_text_encoder( - state_dict, - network_alphas=network_alphas, - text_encoder=getattr(self, self.text_encoder_name) - if not hasattr(self, "text_encoder") - else self.text_encoder, - lora_scale=self.lora_scale, - adapter_name=adapter_name, - _pipeline=self, - metadata=metadata, - low_cpu_mem_usage=low_cpu_mem_usage, - hotswap=hotswap, - ) - - @classmethod - @validate_hf_hub_args - def lora_state_dict( - cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - **kwargs, - ): - r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - weight_name (`str`, *optional*, defaults to None): - Name of the serialized state dict file. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - """ - # Load the main state dict first which has the LoRA layers for either of - # UNet 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) - unet_config = kwargs.pop("unet_config", 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} - - network_alphas = None - # TODO: replace it with a method from `state_dict_utils` - if all( - ( - k.startswith("lora_te_") - or k.startswith("lora_unet_") - or k.startswith("lora_te1_") - or k.startswith("lora_te2_") - ) - for k in state_dict.keys() - ): - # Map SDXL blocks correctly. - if unet_config is not None: - # use unet config to remap block numbers - state_dict = _maybe_map_sgm_blocks_to_diffusers(state_dict, unet_config) - state_dict, network_alphas = _convert_non_diffusers_lora_to_diffusers(state_dict) - - out = (state_dict, network_alphas, metadata) if return_lora_metadata else (state_dict, network_alphas) - return out - - @classmethod - def load_lora_into_unet( - cls, - state_dict, - network_alphas, - unet, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `unet`. - - Parameters: - state_dict (`dict`): - 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]`): - 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). - unet (`UNet2DConditionModel`): - The UNet model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for this method.") - - if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"): - raise ValueError( - "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`." - ) - - # If the serialization format is new (introduced in https://github.com/huggingface/diffusers/pull/2918), - # then the `state_dict` keys should have `cls.unet_name` and/or `cls.text_encoder_name` as - # their prefixes. - logger.info(f"Loading {cls.unet_name}.") - unet.load_lora_adapter( - state_dict, - prefix=cls.unet_name, - network_alphas=network_alphas, - adapter_name=adapter_name, - metadata=metadata, - _pipeline=_pipeline, - low_cpu_mem_usage=low_cpu_mem_usage, - hotswap=hotswap, - ) - - @classmethod - def load_lora_into_text_encoder( - cls, - state_dict, - network_alphas, - text_encoder, - prefix=None, - lora_scale=1.0, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `text_encoder` - - Parameters: - 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]`): - 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). - text_encoder (`CLIPTextModel`): - The text encoder model to load the LoRA layers into. - prefix (`str`): - Expected prefix of the `text_encoder` in the `state_dict`. - lora_scale (`float`): - How much to scale the output of the lora linear layer before it is added with the output of the regular - lora layer. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - _load_lora_into_text_encoder( - state_dict=state_dict, - network_alphas=network_alphas, - lora_scale=lora_scale, - text_encoder=text_encoder, - prefix=prefix, - text_encoder_name=cls.text_encoder_name, - adapter_name=adapter_name, - metadata=metadata, - _pipeline=_pipeline, - low_cpu_mem_usage=low_cpu_mem_usage, - hotswap=hotswap, - ) - - @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, - is_main_process: bool = True, - weight_name: str = None, - save_function: Callable = None, - safe_serialization: bool = True, - unet_lora_adapter_metadata=None, - text_encoder_lora_adapter_metadata=None, - ): - r""" - Save the LoRA parameters corresponding to the UNet and text encoder. - - 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]`): - State dict of the LoRA layers corresponding to the `unet`. - 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`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - unet_lora_adapter_metadata: - LoRA adapter metadata associated with the unet to be serialized with the state dict. - text_encoder_lora_adapter_metadata: - LoRA adapter metadata associated with the text encoder to be serialized with the state dict. - """ - state_dict = {} - lora_adapter_metadata = {} - - if not (unet_lora_layers or text_encoder_lora_layers): - raise ValueError("You must pass at least one of `unet_lora_layers` and `text_encoder_lora_layers`.") - - if unet_lora_layers: - state_dict.update(cls.pack_weights(unet_lora_layers, cls.unet_name)) - - if text_encoder_lora_layers: - state_dict.update(cls.pack_weights(text_encoder_lora_layers, cls.text_encoder_name)) - - if unet_lora_adapter_metadata: - lora_adapter_metadata.update(_pack_dict_with_prefix(unet_lora_adapter_metadata, cls.unet_name)) - - if text_encoder_lora_adapter_metadata: - lora_adapter_metadata.update( - _pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, - ) - - def fuse_lora( - self, - components: List[str] = ["unet", "text_encoder"], - lora_scale: float = 1.0, - safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - super().fuse_lora( - components=components, - lora_scale=lora_scale, - safe_fusing=safe_fusing, - adapter_names=adapter_names, - **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). - - - - This is an experimental API. - - - - Args: - 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 - LoRA parameters then it won't have any effect. - """ - super().unfuse_lora(components=components, **kwargs) - - -class StableDiffusionXLLoraLoaderMixin(LoraBaseMixin): - r""" - Load LoRA layers into Stable Diffusion XL [`UNet2DConditionModel`], - [`CLIPTextModel`](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), and - [`CLIPTextModelWithProjection`](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection). - """ - - _lora_loadable_modules = ["unet", "text_encoder", "text_encoder_2"] - unet_name = UNET_NAME - text_encoder_name = TEXT_ENCODER_NAME - - def load_lora_weights( - self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, - hotswap: bool = False, - **kwargs, - ): - """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.unet` and - `self.text_encoder`. - - All kwargs are forwarded to `self.lora_state_dict`. - - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is - loaded. - - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details on how the state dict is - loaded into `self.unet`. - - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder`] for more details on how the state - dict is loaded into `self.text_encoder`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - """ - 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 not is_peft_version(">=", "0.13.1"): - raise ValueError( - "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`." - ) - - # We could have accessed the unet config from `lora_state_dict()` too. We pass - # it here explicitly to be able to tell that it's coming from an SDXL - # pipeline. - - # 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, network_alphas, metadata = self.lora_state_dict( - pretrained_model_name_or_path_or_dict, - unet_config=self.unet.config, - **kwargs, - ) - - is_correct_format = all("lora" in key for key in state_dict.keys()) - if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") - - self.load_lora_into_unet( - state_dict, - network_alphas=network_alphas, - unet=self.unet, - adapter_name=adapter_name, - metadata=metadata, - _pipeline=self, - low_cpu_mem_usage=low_cpu_mem_usage, - hotswap=hotswap, - ) - self.load_lora_into_text_encoder( - state_dict, - network_alphas=network_alphas, - text_encoder=self.text_encoder, - prefix=self.text_encoder_name, - lora_scale=self.lora_scale, - adapter_name=adapter_name, - metadata=metadata, - _pipeline=self, - low_cpu_mem_usage=low_cpu_mem_usage, - hotswap=hotswap, - ) - self.load_lora_into_text_encoder( - state_dict, - network_alphas=network_alphas, - text_encoder=self.text_encoder_2, - prefix=f"{self.text_encoder_name}_2", - lora_scale=self.lora_scale, - adapter_name=adapter_name, - metadata=metadata, - _pipeline=self, - low_cpu_mem_usage=low_cpu_mem_usage, - hotswap=hotswap, - ) - - @classmethod - @validate_hf_hub_args - # 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]], - **kwargs, - ): - r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - weight_name (`str`, *optional*, defaults to None): - Name of the serialized state dict file. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - """ - # Load the main state dict first which has the LoRA layers for either of - # UNet 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) - unet_config = kwargs.pop("unet_config", 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} - - network_alphas = None - # TODO: replace it with a method from `state_dict_utils` - if all( - ( - k.startswith("lora_te_") - or k.startswith("lora_unet_") - or k.startswith("lora_te1_") - or k.startswith("lora_te2_") - ) - for k in state_dict.keys() - ): - # Map SDXL blocks correctly. - if unet_config is not None: - # use unet config to remap block numbers - state_dict = _maybe_map_sgm_blocks_to_diffusers(state_dict, unet_config) - state_dict, network_alphas = _convert_non_diffusers_lora_to_diffusers(state_dict) - - out = (state_dict, network_alphas, metadata) if return_lora_metadata else (state_dict, network_alphas) - return out - - @classmethod - # Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.load_lora_into_unet - def load_lora_into_unet( - cls, - state_dict, - network_alphas, - unet, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `unet`. - - Parameters: - state_dict (`dict`): - 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]`): - 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). - unet (`UNet2DConditionModel`): - The UNet model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for this method.") - - if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"): - raise ValueError( - "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`." - ) - - # If the serialization format is new (introduced in https://github.com/huggingface/diffusers/pull/2918), - # then the `state_dict` keys should have `cls.unet_name` and/or `cls.text_encoder_name` as - # their prefixes. - logger.info(f"Loading {cls.unet_name}.") - unet.load_lora_adapter( - state_dict, - prefix=cls.unet_name, - network_alphas=network_alphas, - 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.StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder - def load_lora_into_text_encoder( - cls, - state_dict, - network_alphas, - text_encoder, - prefix=None, - lora_scale=1.0, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `text_encoder` - - Parameters: - 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]`): - 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). - text_encoder (`CLIPTextModel`): - The text encoder model to load the LoRA layers into. - prefix (`str`): - Expected prefix of the `text_encoder` in the `state_dict`. - lora_scale (`float`): - How much to scale the output of the lora linear layer before it is added with the output of the regular - lora layer. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - _load_lora_into_text_encoder( - state_dict=state_dict, - network_alphas=network_alphas, - lora_scale=lora_scale, - text_encoder=text_encoder, - prefix=prefix, - text_encoder_name=cls.text_encoder_name, - adapter_name=adapter_name, - metadata=metadata, - _pipeline=_pipeline, - low_cpu_mem_usage=low_cpu_mem_usage, - hotswap=hotswap, - ) - - @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, - is_main_process: bool = True, - weight_name: str = None, - save_function: Callable = None, - safe_serialization: bool = True, - unet_lora_adapter_metadata=None, - text_encoder_lora_adapter_metadata=None, - text_encoder_2_lora_adapter_metadata=None, - ): - r""" - Save the LoRA parameters corresponding to the UNet and text encoder. - - 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]`): - State dict of the LoRA layers corresponding to the `unet`. - 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. - text_encoder_2_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `text_encoder_2`. Must explicitly pass the text - encoder LoRA state dict because it comes from 🤗 Transformers. - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - unet_lora_adapter_metadata: - LoRA adapter metadata associated with the unet to be serialized with the state dict. - text_encoder_lora_adapter_metadata: - LoRA adapter metadata associated with the text encoder to be serialized with the state dict. - text_encoder_2_lora_adapter_metadata: - LoRA adapter metadata associated with the second text encoder to be serialized with the state dict. - """ - state_dict = {} - lora_adapter_metadata = {} - - if not (unet_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers): - raise ValueError( - "You must pass at least one of `unet_lora_layers`, `text_encoder_lora_layers`, `text_encoder_2_lora_layers`." - ) - - if unet_lora_layers: - state_dict.update(cls.pack_weights(unet_lora_layers, cls.unet_name)) - - if text_encoder_lora_layers: - state_dict.update(cls.pack_weights(text_encoder_lora_layers, "text_encoder")) - - if text_encoder_2_lora_layers: - state_dict.update(cls.pack_weights(text_encoder_2_lora_layers, "text_encoder_2")) - - if unet_lora_adapter_metadata is not None: - lora_adapter_metadata.update(_pack_dict_with_prefix(unet_lora_adapter_metadata, cls.unet_name)) - - if text_encoder_lora_adapter_metadata: - lora_adapter_metadata.update( - _pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name) - ) - - if text_encoder_2_lora_adapter_metadata: - lora_adapter_metadata.update( - _pack_dict_with_prefix(text_encoder_2_lora_adapter_metadata, "text_encoder_2") - ) - - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, - ) - - def fuse_lora( - self, - components: List[str] = ["unet", "text_encoder", "text_encoder_2"], - lora_scale: float = 1.0, - safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - super().fuse_lora( - components=components, - lora_scale=lora_scale, - safe_fusing=safe_fusing, - adapter_names=adapter_names, - **kwargs, - ) - - def unfuse_lora(self, components: List[str] = ["unet", "text_encoder", "text_encoder_2"], **kwargs): - r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - 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 - LoRA parameters then it won't have any effect. - """ - super().unfuse_lora(components=components, **kwargs) - - -class SD3LoraLoaderMixin(LoraBaseMixin): - r""" - Load LoRA layers into [`SD3Transformer2DModel`], - [`CLIPTextModel`](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), and - [`CLIPTextModelWithProjection`](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection). - - Specific to [`StableDiffusion3Pipeline`]. - """ - - _lora_loadable_modules = ["transformer", "text_encoder", "text_encoder_2"] - transformer_name = TRANSFORMER_NAME - text_encoder_name = TEXT_ENCODER_NAME - - @classmethod - @validate_hf_hub_args - def lora_state_dict( - cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - **kwargs, - ): - r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - - """ - # 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} - - 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: Union[str, Dict[str, torch.Tensor]], - adapter_name=None, - hotswap: bool = False, - **kwargs, - ): - """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.unet` and - `self.text_encoder`. - - All kwargs are forwarded to `self.lora_state_dict`. - - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is - loaded. - - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - """ - 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.") - - 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, - ) - self.load_lora_into_text_encoder( - state_dict, - network_alphas=None, - text_encoder=self.text_encoder, - prefix=self.text_encoder_name, - lora_scale=self.lora_scale, - adapter_name=adapter_name, - metadata=metadata, - _pipeline=self, - low_cpu_mem_usage=low_cpu_mem_usage, - hotswap=hotswap, - ) - self.load_lora_into_text_encoder( - state_dict, - network_alphas=None, - text_encoder=self.text_encoder_2, - prefix=f"{self.text_encoder_name}_2", - lora_scale=self.lora_scale, - adapter_name=adapter_name, - metadata=metadata, - _pipeline=self, - low_cpu_mem_usage=low_cpu_mem_usage, - hotswap=hotswap, - ) - - @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, - ): - """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - 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. - transformer (`SD3Transformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - 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.StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder - def load_lora_into_text_encoder( - cls, - state_dict, - network_alphas, - text_encoder, - prefix=None, - lora_scale=1.0, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `text_encoder` - - Parameters: - 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]`): - 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). - text_encoder (`CLIPTextModel`): - The text encoder model to load the LoRA layers into. - prefix (`str`): - Expected prefix of the `text_encoder` in the `state_dict`. - lora_scale (`float`): - How much to scale the output of the lora linear layer before it is added with the output of the regular - lora layer. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - _load_lora_into_text_encoder( - state_dict=state_dict, - network_alphas=network_alphas, - lora_scale=lora_scale, - text_encoder=text_encoder, - prefix=prefix, - text_encoder_name=cls.text_encoder_name, - 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.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, - is_main_process: bool = True, - weight_name: str = None, - save_function: Callable = None, - safe_serialization: bool = True, - transformer_lora_adapter_metadata=None, - text_encoder_lora_adapter_metadata=None, - text_encoder_2_lora_adapter_metadata=None, - ): - r""" - Save the LoRA parameters corresponding to the UNet and text encoder. - - 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]`): - State dict of the LoRA layers corresponding to the `transformer`. - 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. - text_encoder_2_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): - State dict of the LoRA layers corresponding to the `text_encoder_2`. Must explicitly pass the text - encoder LoRA state dict because it comes from 🤗 Transformers. - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. - text_encoder_lora_adapter_metadata: - LoRA adapter metadata associated with the text encoder to be serialized with the state dict. - text_encoder_2_lora_adapter_metadata: - LoRA adapter metadata associated with the second text encoder to be serialized with the state dict. - """ - state_dict = {} - lora_adapter_metadata = {} - - if not (transformer_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers): - raise ValueError( - "You must pass at least one of `transformer_lora_layers`, `text_encoder_lora_layers`, `text_encoder_2_lora_layers`." - ) - - if transformer_lora_layers: - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) - - if text_encoder_lora_layers: - state_dict.update(cls.pack_weights(text_encoder_lora_layers, "text_encoder")) - - if text_encoder_2_lora_layers: - state_dict.update(cls.pack_weights(text_encoder_2_lora_layers, "text_encoder_2")) - - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - if text_encoder_lora_adapter_metadata: - lora_adapter_metadata.update( - _pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name) - ) - - if text_encoder_2_lora_adapter_metadata: - lora_adapter_metadata.update( - _pack_dict_with_prefix(text_encoder_2_lora_adapter_metadata, "text_encoder_2") - ) - - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, - ) - - # 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"], - lora_scale: float = 1.0, - safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - super().fuse_lora( - components=components, - lora_scale=lora_scale, - safe_fusing=safe_fusing, - adapter_names=adapter_names, - **kwargs, - ) - - # 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): - r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`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 - LoRA parameters then it won't have any effect. - """ - super().unfuse_lora(components=components, **kwargs) - - -class AuraFlowLoraLoaderMixin(LoraBaseMixin): - r""" - Load LoRA layers into [`AuraFlowTransformer2DModel`] Specific to [`AuraFlowPipeline`]. - """ - - _lora_loadable_modules = ["transformer"] - transformer_name = TRANSFORMER_NAME - - @classmethod - @validate_hf_hub_args - # 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]], - **kwargs, - ): - r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - - """ - # 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} - - out = (state_dict, metadata) if return_lora_metadata else state_dict - return out - - # 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, - hotswap: bool = False, - **kwargs, - ): - """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - """ - 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.") - - 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->AuraFlowTransformer2DModel - def load_lora_into_transformer( - cls, - state_dict, - transformer, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - 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. - transformer (`AuraFlowTransformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - 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: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[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, - ): - r""" - Save the LoRA parameters corresponding to the transformer. - - 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]`): - State dict of the LoRA layers corresponding to the `transformer`. - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. - """ - state_dict = {} - lora_adapter_metadata = {} - - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") - - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) - - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, - ) - - # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora - def fuse_lora( - self, - components: List[str] = ["transformer"], - lora_scale: float = 1.0, - safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - super().fuse_lora( - components=components, - lora_scale=lora_scale, - safe_fusing=safe_fusing, - adapter_names=adapter_names, - **kwargs, - ) - - # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora - 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). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. - """ - super().unfuse_lora(components=components, **kwargs) - - -class FluxLoraLoaderMixin(LoraBaseMixin): - r""" - Load LoRA layers into [`FluxTransformer2DModel`], - [`CLIPTextModel`](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel). - - Specific to [`StableDiffusion3Pipeline`]. - """ - - _lora_loadable_modules = ["transformer", "text_encoder"] - transformer_name = TRANSFORMER_NAME - text_encoder_name = TEXT_ENCODER_NAME - _control_lora_supported_norm_keys = ["norm_q", "norm_k", "norm_added_q", "norm_added_k"] - - @classmethod - @validate_hf_hub_args - def lora_state_dict( - cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - return_alphas: bool = False, - **kwargs, - ): - r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - """ - # 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} - - # TODO (sayakpaul): to a follow-up to clean and try to unify the conditions. - is_kohya = any(".lora_down.weight" in k for k in state_dict) - if is_kohya: - state_dict = _convert_kohya_flux_lora_to_diffusers(state_dict) - # Kohya already takes care of scaling the LoRA parameters with alpha. - return cls._prepare_outputs( - state_dict, - metadata=metadata, - alphas=None, - return_alphas=return_alphas, - return_metadata=return_lora_metadata, - ) - - is_xlabs = any("processor" in k for k in state_dict) - if is_xlabs: - state_dict = _convert_xlabs_flux_lora_to_diffusers(state_dict) - # xlabs doesn't use `alpha`. - return cls._prepare_outputs( - state_dict, - metadata=metadata, - alphas=None, - return_alphas=return_alphas, - return_metadata=return_lora_metadata, - ) - - is_bfl_control = any("query_norm.scale" in k for k in state_dict) - if is_bfl_control: - state_dict = _convert_bfl_flux_control_lora_to_diffusers(state_dict) - return cls._prepare_outputs( - state_dict, - metadata=metadata, - alphas=None, - return_alphas=return_alphas, - return_metadata=return_lora_metadata, - ) - - # For state dicts like - # https://huggingface.co/TheLastBen/Jon_Snow_Flux_LoRA - keys = list(state_dict.keys()) - network_alphas = {} - for k in keys: - if "alpha" in k: - alpha_value = state_dict.get(k) - if (torch.is_tensor(alpha_value) and torch.is_floating_point(alpha_value)) or isinstance( - alpha_value, float - ): - network_alphas[k] = state_dict.pop(k) - else: - raise ValueError( - f"The alpha key ({k}) seems to be incorrect. If you think this error is unexpected, please open as issue." - ) - - if return_alphas or return_lora_metadata: - return cls._prepare_outputs( - state_dict, - metadata=metadata, - alphas=network_alphas, - return_alphas=return_alphas, - return_metadata=return_lora_metadata, - ) - else: - return 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, - hotswap: bool = False, - **kwargs, - ): - """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. - - All kwargs are forwarded to `self.lora_state_dict`. - - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is - loaded. - - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - `Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - """ - 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 not is_peft_version(">=", "0.13.1"): - 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, network_alphas, metadata = self.lora_state_dict( - pretrained_model_name_or_path_or_dict, return_alphas=True, **kwargs - ) - - has_lora_keys = any("lora" in key for key in state_dict.keys()) - - # Flux Control LoRAs also have norm keys - has_norm_keys = any( - norm_key in key for key in state_dict.keys() for norm_key in self._control_lora_supported_norm_keys - ) - - if not (has_lora_keys or has_norm_keys): - raise ValueError("Invalid LoRA checkpoint.") - - transformer_lora_state_dict = { - k: state_dict.get(k) - for k in list(state_dict.keys()) - if k.startswith(f"{self.transformer_name}.") and "lora" in k - } - transformer_norm_state_dict = { - k: state_dict.pop(k) - for k in list(state_dict.keys()) - if k.startswith(f"{self.transformer_name}.") - and any(norm_key in k for norm_key in self._control_lora_supported_norm_keys) - } - - transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer - has_param_with_expanded_shape = False - if len(transformer_lora_state_dict) > 0: - has_param_with_expanded_shape = self._maybe_expand_transformer_param_shape_or_error_( - transformer, transformer_lora_state_dict, transformer_norm_state_dict - ) - - if has_param_with_expanded_shape: - logger.info( - "The LoRA weights contain parameters that have different shapes that expected by the transformer. " - "As a result, the state_dict of the transformer has been expanded to match the LoRA parameter shapes. " - "To get a comprehensive list of parameter names that were modified, enable debug logging." - ) - if len(transformer_lora_state_dict) > 0: - transformer_lora_state_dict = self._maybe_expand_lora_state_dict( - transformer=transformer, lora_state_dict=transformer_lora_state_dict - ) - for k in transformer_lora_state_dict: - state_dict.update({k: transformer_lora_state_dict[k]}) - - self.load_lora_into_transformer( - state_dict, - network_alphas=network_alphas, - transformer=transformer, - adapter_name=adapter_name, - metadata=metadata, - _pipeline=self, - low_cpu_mem_usage=low_cpu_mem_usage, - hotswap=hotswap, - ) - - if len(transformer_norm_state_dict) > 0: - transformer._transformer_norm_layers = self._load_norm_into_transformer( - transformer_norm_state_dict, - transformer=transformer, - discard_original_layers=False, - ) - - self.load_lora_into_text_encoder( - state_dict, - network_alphas=network_alphas, - text_encoder=self.text_encoder, - prefix=self.text_encoder_name, - lora_scale=self.lora_scale, - adapter_name=adapter_name, - metadata=metadata, - _pipeline=self, - low_cpu_mem_usage=low_cpu_mem_usage, - hotswap=hotswap, - ) - - @classmethod - def load_lora_into_transformer( - cls, - state_dict, - network_alphas, - transformer, - adapter_name=None, - metadata=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - ): - """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - 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]`): - 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). - transformer (`FluxTransformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"): - 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=network_alphas, - adapter_name=adapter_name, - metadata=metadata, - _pipeline=_pipeline, - low_cpu_mem_usage=low_cpu_mem_usage, - hotswap=hotswap, - ) - - @classmethod - def _load_norm_into_transformer( - cls, - state_dict, - transformer, - prefix=None, - discard_original_layers=False, - ) -> Dict[str, torch.Tensor]: - # Remove prefix if present - prefix = prefix or cls.transformer_name - for key in list(state_dict.keys()): - if key.split(".")[0] == prefix: - state_dict[key.removeprefix(f"{prefix}.")] = state_dict.pop(key) - - # Find invalid keys - transformer_state_dict = transformer.state_dict() - transformer_keys = set(transformer_state_dict.keys()) - state_dict_keys = set(state_dict.keys()) - extra_keys = list(state_dict_keys - transformer_keys) - - if extra_keys: - logger.warning( - f"Unsupported keys found in state dict when trying to load normalization layers into the transformer. The following keys will be ignored:\n{extra_keys}." - ) - - for key in extra_keys: - state_dict.pop(key) - - # Save the layers that are going to be overwritten so that unload_lora_weights can work as expected - overwritten_layers_state_dict = {} - if not discard_original_layers: - for key in state_dict.keys(): - overwritten_layers_state_dict[key] = transformer_state_dict[key].clone() - - logger.info( - "The provided state dict contains normalization layers in addition to LoRA layers. The normalization layers will directly update the state_dict of the transformer " - 'as opposed to the LoRA layers that will co-exist separately until the "fuse_lora()" method is called. That is to say, the normalization layers will always be directly ' - "fused into the transformer and can only be unfused if `discard_original_layers=True` is passed. This might also have implications when dealing with multiple LoRAs. " - "If you notice something unexpected, please open an issue: https://github.com/huggingface/diffusers/issues." - ) - - # We can't load with strict=True because the current state_dict does not contain all the transformer keys - incompatible_keys = transformer.load_state_dict(state_dict, strict=False) - unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) - - # We shouldn't expect to see the supported norm keys here being present in the unexpected keys. - if unexpected_keys: - if any(norm_key in k for k in unexpected_keys for norm_key in cls._control_lora_supported_norm_keys): - raise ValueError( - f"Found {unexpected_keys} as unexpected keys while trying to load norm layers into the transformer." - ) - - return overwritten_layers_state_dict - - @classmethod - # Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder - def load_lora_into_text_encoder( - cls, - state_dict, - network_alphas, - text_encoder, - prefix=None, - lora_scale=1.0, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `text_encoder` - - Parameters: - 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]`): - 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). - text_encoder (`CLIPTextModel`): - The text encoder model to load the LoRA layers into. - prefix (`str`): - Expected prefix of the `text_encoder` in the `state_dict`. - lora_scale (`float`): - How much to scale the output of the lora linear layer before it is added with the output of the regular - lora layer. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - _load_lora_into_text_encoder( - state_dict=state_dict, - network_alphas=network_alphas, - lora_scale=lora_scale, - text_encoder=text_encoder, - prefix=prefix, - text_encoder_name=cls.text_encoder_name, - 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.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, - is_main_process: bool = True, - weight_name: str = None, - save_function: Callable = None, - safe_serialization: bool = True, - transformer_lora_adapter_metadata=None, - text_encoder_lora_adapter_metadata=None, - ): - r""" - Save the LoRA parameters corresponding to the UNet and text encoder. - - 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]`): - State dict of the LoRA layers corresponding to the `transformer`. - 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`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. - text_encoder_lora_adapter_metadata: - LoRA adapter metadata associated with the text encoder to be serialized with the state dict. - """ - state_dict = {} - lora_adapter_metadata = {} - - if not (transformer_lora_layers or text_encoder_lora_layers): - raise ValueError("You must pass at least one of `transformer_lora_layers` and `text_encoder_lora_layers`.") - - if transformer_lora_layers: - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) - - if text_encoder_lora_layers: - state_dict.update(cls.pack_weights(text_encoder_lora_layers, cls.text_encoder_name)) - - if transformer_lora_adapter_metadata: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - if text_encoder_lora_adapter_metadata: - lora_adapter_metadata.update( - _pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, - ) - - def fuse_lora( - self, - components: List[str] = ["transformer"], - lora_scale: float = 1.0, - safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - - transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer - if ( - hasattr(transformer, "_transformer_norm_layers") - and isinstance(transformer._transformer_norm_layers, dict) - and len(transformer._transformer_norm_layers.keys()) > 0 - ): - logger.info( - "The provided state dict contains normalization layers in addition to LoRA layers. The normalization layers will be directly updated the state_dict of the transformer " - "as opposed to the LoRA layers that will co-exist separately until the 'fuse_lora()' method is called. That is to say, the normalization layers will always be directly " - "fused into the transformer and can only be unfused if `discard_original_layers=True` is passed." - ) - - super().fuse_lora( - components=components, - lora_scale=lora_scale, - safe_fusing=safe_fusing, - adapter_names=adapter_names, - **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). - - - - This is an experimental API. - - - - Args: - 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: - transformer.load_state_dict(transformer._transformer_norm_layers, strict=False) - - super().unfuse_lora(components=components, **kwargs) - - # We override this here account for `_transformer_norm_layers` and `_overwritten_params`. - def unload_lora_weights(self, reset_to_overwritten_params=False): - """ - Unloads the LoRA parameters. - - Args: - reset_to_overwritten_params (`bool`, defaults to `False`): Whether to reset the LoRA-loaded modules - to their original params. Refer to the [Flux - documentation](https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux) to learn more. - - Examples: - - ```python - >>> # Assuming `pipeline` is already loaded with the LoRA parameters. - >>> pipeline.unload_lora_weights() - >>> ... - ``` - """ - super().unload_lora_weights() - - 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: - transformer.load_state_dict(transformer._transformer_norm_layers, strict=False) - transformer._transformer_norm_layers = None - - if reset_to_overwritten_params and getattr(transformer, "_overwritten_params", None) is not None: - overwritten_params = transformer._overwritten_params - module_names = set() - - for param_name in overwritten_params: - if param_name.endswith(".weight"): - module_names.add(param_name.replace(".weight", "")) - - for name, module in transformer.named_modules(): - if isinstance(module, torch.nn.Linear) and name in module_names: - module_weight = module.weight.data - module_bias = module.bias.data if module.bias is not None else None - bias = module_bias is not None - - parent_module_name, _, current_module_name = name.rpartition(".") - parent_module = transformer.get_submodule(parent_module_name) - - current_param_weight = overwritten_params[f"{name}.weight"] - in_features, out_features = current_param_weight.shape[1], current_param_weight.shape[0] - with torch.device("meta"): - original_module = torch.nn.Linear( - in_features, - out_features, - bias=bias, - dtype=module_weight.dtype, - ) - - tmp_state_dict = {"weight": current_param_weight} - if module_bias is not None: - tmp_state_dict.update({"bias": overwritten_params[f"{name}.bias"]}) - original_module.load_state_dict(tmp_state_dict, assign=True, strict=True) - setattr(parent_module, current_module_name, original_module) - - del tmp_state_dict - - if current_module_name in _MODULE_NAME_TO_ATTRIBUTE_MAP_FLUX: - attribute_name = _MODULE_NAME_TO_ATTRIBUTE_MAP_FLUX[current_module_name] - new_value = int(current_param_weight.shape[1]) - old_value = getattr(transformer.config, attribute_name) - setattr(transformer.config, attribute_name, new_value) - logger.info( - f"Set the {attribute_name} attribute of the model to {new_value} from {old_value}." - ) - - @classmethod - def _maybe_expand_transformer_param_shape_or_error_( - cls, - transformer: torch.nn.Module, - lora_state_dict=None, - norm_state_dict=None, - prefix=None, - ) -> bool: - """ - Control LoRA expands the shape of the input layer from (3072, 64) to (3072, 128). This method handles that and - generalizes things a bit so that any parameter that needs expansion receives appropriate treatment. - """ - state_dict = {} - if lora_state_dict is not None: - state_dict.update(lora_state_dict) - if norm_state_dict is not None: - state_dict.update(norm_state_dict) - - # Remove prefix if present - prefix = prefix or cls.transformer_name - for key in list(state_dict.keys()): - if key.split(".")[0] == prefix: - state_dict[key.removeprefix(f"{prefix}.")] = state_dict.pop(key) - - # Expand transformer parameter shapes if they don't match lora - has_param_with_shape_update = False - overwritten_params = {} - - is_peft_loaded = getattr(transformer, "peft_config", None) is not None - is_quantized = hasattr(transformer, "hf_quantizer") - for name, module in transformer.named_modules(): - if isinstance(module, torch.nn.Linear): - module_weight = module.weight.data - module_bias = module.bias.data if module.bias is not None else None - bias = module_bias is not None - - lora_base_name = name.replace(".base_layer", "") if is_peft_loaded else name - lora_A_weight_name = f"{lora_base_name}.lora_A.weight" - lora_B_weight_name = f"{lora_base_name}.lora_B.weight" - if lora_A_weight_name not in state_dict: - continue - - in_features = state_dict[lora_A_weight_name].shape[1] - out_features = state_dict[lora_B_weight_name].shape[0] - - # Model maybe loaded with different quantization schemes which may flatten the params. - # `bitsandbytes`, for example, flatten the weights when using 4bit. 8bit bnb models - # preserve weight shape. - module_weight_shape = cls._calculate_module_shape(model=transformer, base_module=module) - - # This means there's no need for an expansion in the params, so we simply skip. - if tuple(module_weight_shape) == (out_features, in_features): - continue - - module_out_features, module_in_features = module_weight_shape - debug_message = "" - if in_features > module_in_features: - debug_message += ( - f'Expanding the nn.Linear input/output features for module="{name}" because the provided LoRA ' - f"checkpoint contains higher number of features than expected. The number of input_features will be " - f"expanded from {module_in_features} to {in_features}" - ) - if out_features > module_out_features: - debug_message += ( - ", and the number of output features will be " - f"expanded from {module_out_features} to {out_features}." - ) - else: - debug_message += "." - if debug_message: - logger.debug(debug_message) - - if out_features > module_out_features or in_features > module_in_features: - has_param_with_shape_update = True - parent_module_name, _, current_module_name = name.rpartition(".") - parent_module = transformer.get_submodule(parent_module_name) - - if is_quantized: - module_weight = _maybe_dequantize_weight_for_expanded_lora(transformer, module) - - # TODO: consider if this layer needs to be a quantized layer as well if `is_quantized` is True. - with torch.device("meta"): - expanded_module = torch.nn.Linear( - in_features, out_features, bias=bias, dtype=module_weight.dtype - ) - # Only weights are expanded and biases are not. This is because only the input dimensions - # are changed while the output dimensions remain the same. The shape of the weight tensor - # is (out_features, in_features), while the shape of bias tensor is (out_features,), which - # explains the reason why only weights are expanded. - new_weight = torch.zeros_like( - expanded_module.weight.data, device=module_weight.device, dtype=module_weight.dtype - ) - slices = tuple(slice(0, dim) for dim in module_weight_shape) - new_weight[slices] = module_weight - tmp_state_dict = {"weight": new_weight} - if module_bias is not None: - tmp_state_dict["bias"] = module_bias - expanded_module.load_state_dict(tmp_state_dict, strict=True, assign=True) - - setattr(parent_module, current_module_name, expanded_module) - - del tmp_state_dict - - if current_module_name in _MODULE_NAME_TO_ATTRIBUTE_MAP_FLUX: - attribute_name = _MODULE_NAME_TO_ATTRIBUTE_MAP_FLUX[current_module_name] - new_value = int(expanded_module.weight.data.shape[1]) - old_value = getattr(transformer.config, attribute_name) - setattr(transformer.config, attribute_name, new_value) - logger.info( - f"Set the {attribute_name} attribute of the model to {new_value} from {old_value}." - ) - - # For `unload_lora_weights()`. - # TODO: this could lead to more memory overhead if the number of overwritten params - # are large. Should be revisited later and tackled through a `discard_original_layers` arg. - overwritten_params[f"{current_module_name}.weight"] = module_weight - if module_bias is not None: - overwritten_params[f"{current_module_name}.bias"] = module_bias - - if len(overwritten_params) > 0: - transformer._overwritten_params = overwritten_params - - return has_param_with_shape_update - - @classmethod - def _maybe_expand_lora_state_dict(cls, transformer, lora_state_dict): - expanded_module_names = set() - transformer_state_dict = transformer.state_dict() - prefix = f"{cls.transformer_name}." - - lora_module_names = [ - key[: -len(".lora_A.weight")] for key in lora_state_dict if key.endswith(".lora_A.weight") - ] - lora_module_names = [name[len(prefix) :] for name in lora_module_names if name.startswith(prefix)] - lora_module_names = sorted(set(lora_module_names)) - transformer_module_names = sorted({name for name, _ in transformer.named_modules()}) - unexpected_modules = set(lora_module_names) - set(transformer_module_names) - if unexpected_modules: - logger.debug(f"Found unexpected modules: {unexpected_modules}. These will be ignored.") - - for k in lora_module_names: - if k in unexpected_modules: - continue - - base_param_name = ( - f"{k.replace(prefix, '')}.base_layer.weight" - if f"{k.replace(prefix, '')}.base_layer.weight" in transformer_state_dict - else f"{k.replace(prefix, '')}.weight" - ) - base_weight_param = transformer_state_dict[base_param_name] - lora_A_param = lora_state_dict[f"{prefix}{k}.lora_A.weight"] - - # TODO (sayakpaul): Handle the cases when we actually need to expand when using quantization. - base_module_shape = cls._calculate_module_shape(model=transformer, base_weight_param_name=base_param_name) - - if base_module_shape[1] > lora_A_param.shape[1]: - shape = (lora_A_param.shape[0], base_weight_param.shape[1]) - expanded_state_dict_weight = torch.zeros(shape, device=base_weight_param.device) - expanded_state_dict_weight[:, : lora_A_param.shape[1]].copy_(lora_A_param) - lora_state_dict[f"{prefix}{k}.lora_A.weight"] = expanded_state_dict_weight - expanded_module_names.add(k) - elif base_module_shape[1] < lora_A_param.shape[1]: - raise NotImplementedError( - f"This LoRA param ({k}.lora_A.weight) has an incompatible shape {lora_A_param.shape}. Please open an issue to file for a feature request - https://github.com/huggingface/diffusers/issues/new." - ) - - if expanded_module_names: - logger.info( - f"The following LoRA modules were zero padded to match the state dict of {cls.transformer_name}: {expanded_module_names}. Please open an issue if you think this was unexpected - https://github.com/huggingface/diffusers/issues/new." - ) - - return lora_state_dict - - @staticmethod - def _calculate_module_shape( - model: "torch.nn.Module", - base_module: "torch.nn.Linear" = None, - base_weight_param_name: str = None, - ) -> "torch.Size": - def _get_weight_shape(weight: torch.Tensor): - if weight.__class__.__name__ == "Params4bit": - return weight.quant_state.shape - elif weight.__class__.__name__ == "GGUFParameter": - return weight.quant_shape - else: - return weight.shape - - if base_module is not None: - return _get_weight_shape(base_module.weight) - elif base_weight_param_name is not None: - if not base_weight_param_name.endswith(".weight"): - raise ValueError( - f"Invalid `base_weight_param_name` passed as it does not end with '.weight' {base_weight_param_name=}." - ) - module_path = base_weight_param_name.rsplit(".weight", 1)[0] - submodule = get_submodule_by_name(model, module_path) - return _get_weight_shape(submodule.weight) - - raise ValueError("Either `base_module` or `base_weight_param_name` must be provided.") - - @staticmethod - def _prepare_outputs(state_dict, metadata, alphas=None, return_alphas=False, return_metadata=False): - outputs = [state_dict] - if return_alphas: - outputs.append(alphas) - if return_metadata: - outputs.append(metadata) - return tuple(outputs) if (return_alphas or return_metadata) else state_dict - - -# The reason why we subclass from `StableDiffusionLoraLoaderMixin` here is because Amused initially -# relied on `StableDiffusionLoraLoaderMixin` for its LoRA support. -class AmusedLoraLoaderMixin(StableDiffusionLoraLoaderMixin): - _lora_loadable_modules = ["transformer", "text_encoder"] - transformer_name = TRANSFORMER_NAME - text_encoder_name = TEXT_ENCODER_NAME - - @classmethod - # Copied from diffusers.loaders.lora_pipeline.FluxLoraLoaderMixin.load_lora_into_transformer with FluxTransformer2DModel->UVit2DModel - def load_lora_into_transformer( - cls, - state_dict, - network_alphas, - transformer, - adapter_name=None, - metadata=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - ): - """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - 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]`): - 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). - transformer (`UVit2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"): - 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=network_alphas, - 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.StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder - def load_lora_into_text_encoder( - cls, - state_dict, - network_alphas, - text_encoder, - prefix=None, - lora_scale=1.0, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `text_encoder` - - Parameters: - 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]`): - 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). - text_encoder (`CLIPTextModel`): - The text encoder model to load the LoRA layers into. - prefix (`str`): - Expected prefix of the `text_encoder` in the `state_dict`. - lora_scale (`float`): - How much to scale the output of the lora linear layer before it is added with the output of the regular - lora layer. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - _load_lora_into_text_encoder( - state_dict=state_dict, - network_alphas=network_alphas, - lora_scale=lora_scale, - text_encoder=text_encoder, - prefix=prefix, - text_encoder_name=cls.text_encoder_name, - adapter_name=adapter_name, - metadata=metadata, - _pipeline=_pipeline, - low_cpu_mem_usage=low_cpu_mem_usage, - hotswap=hotswap, - ) - - @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, - is_main_process: bool = True, - weight_name: str = None, - save_function: Callable = None, - safe_serialization: bool = True, - ): - r""" - Save the LoRA parameters corresponding to the UNet and text encoder. - - 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]`): - State dict of the LoRA layers corresponding to the `unet`. - 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`): - Whether the process calling this is the main process or not. Useful during distributed training and you - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - """ - state_dict = {} - - if not (transformer_lora_layers or text_encoder_lora_layers): - raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") - - if transformer_lora_layers: - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) - - if text_encoder_lora_layers: - state_dict.update(cls.pack_weights(text_encoder_lora_layers, cls.text_encoder_name)) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - ) - - -class CogVideoXLoraLoaderMixin(LoraBaseMixin): - r""" - Load LoRA layers into [`CogVideoXTransformer3DModel`]. Specific to [`CogVideoXPipeline`]. - """ - - _lora_loadable_modules = ["transformer"] - transformer_name = TRANSFORMER_NAME - - @classmethod - @validate_hf_hub_args - # 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]], - **kwargs, - ): - r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - - """ - # 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} - - 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: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, - hotswap: bool = False, - **kwargs, - ): - """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - """ - 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.") - - 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->CogVideoXTransformer3DModel - def load_lora_into_transformer( - cls, - state_dict, - transformer, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - 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. - transformer (`CogVideoXTransformer3DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - 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 - # Adapted from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.save_lora_weights without support for text encoder - def save_lora_weights( - cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[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, - ): - r""" - Save the LoRA parameters corresponding to the transformer. - - 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]`): - State dict of the LoRA layers corresponding to the `transformer`. - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. - """ - state_dict = {} - lora_adapter_metadata = {} - - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") - - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) - - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, - ) - - def fuse_lora( - self, - components: List[str] = ["transformer"], - lora_scale: float = 1.0, - safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - super().fuse_lora( - components=components, - lora_scale=lora_scale, - safe_fusing=safe_fusing, - adapter_names=adapter_names, - **kwargs, - ) - - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): - r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. - """ - super().unfuse_lora(components=components, **kwargs) - - -class Mochi1LoraLoaderMixin(LoraBaseMixin): - r""" - Load LoRA layers into [`MochiTransformer3DModel`]. Specific to [`MochiPipeline`]. - """ - - _lora_loadable_modules = ["transformer"] - transformer_name = TRANSFORMER_NAME - - @classmethod - @validate_hf_hub_args - # 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]], - **kwargs, - ): - r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - - """ - # 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} - - out = (state_dict, metadata) if return_lora_metadata else state_dict - return out - - # 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, - hotswap: bool = False, - **kwargs, - ): - """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - """ - 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.") - - 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->MochiTransformer3DModel - def load_lora_into_transformer( - cls, - state_dict, - transformer, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - 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. - transformer (`MochiTransformer3DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - 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: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[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, - ): - r""" - Save the LoRA parameters corresponding to the transformer. - - 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]`): - State dict of the LoRA layers corresponding to the `transformer`. - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. - """ - state_dict = {} - lora_adapter_metadata = {} - - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") - - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) - - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, - ) - - # 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: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - 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""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. - """ - super().unfuse_lora(components=components, **kwargs) - - -class LTXVideoLoraLoaderMixin(LoraBaseMixin): - r""" - Load LoRA layers into [`LTXVideoTransformer3DModel`]. Specific to [`LTXPipeline`]. - """ - - _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: Union[str, Dict[str, torch.Tensor]], - **kwargs, - ): - r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - """ - # 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} - - is_non_diffusers_format = any(k.startswith("diffusion_model.") for k in state_dict) - if is_non_diffusers_format: - state_dict = _convert_non_diffusers_ltxv_lora_to_diffusers(state_dict) - - out = (state_dict, metadata) if return_lora_metadata else state_dict - return out - - # 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, - hotswap: bool = False, - **kwargs, - ): - """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - """ - 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.") - - 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->LTXVideoTransformer3DModel - def load_lora_into_transformer( - cls, - state_dict, - transformer, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - 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. - transformer (`LTXVideoTransformer3DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - 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: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[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, - ): - r""" - Save the LoRA parameters corresponding to the transformer. - - 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]`): - State dict of the LoRA layers corresponding to the `transformer`. - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. - """ - state_dict = {} - lora_adapter_metadata = {} - - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") - - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) - - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, - ) - - # 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: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - 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""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. - """ - super().unfuse_lora(components=components, **kwargs) - - -class SanaLoraLoaderMixin(LoraBaseMixin): - r""" - Load LoRA layers into [`SanaTransformer2DModel`]. Specific to [`SanaPipeline`]. - """ - - _lora_loadable_modules = ["transformer"] - transformer_name = TRANSFORMER_NAME - - @classmethod - @validate_hf_hub_args - # 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]], - **kwargs, - ): - r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - - """ - # 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} - - out = (state_dict, metadata) if return_lora_metadata else state_dict - return out - - # 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, - hotswap: bool = False, - **kwargs, - ): - """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - """ - 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.") - - 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->SanaTransformer2DModel - def load_lora_into_transformer( - cls, - state_dict, - transformer, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - 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. - transformer (`SanaTransformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - 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: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[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, - ): - r""" - Save the LoRA parameters corresponding to the transformer. - - 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]`): - State dict of the LoRA layers corresponding to the `transformer`. - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. - """ - state_dict = {} - lora_adapter_metadata = {} - - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") - - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) - - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, - ) - - # 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: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - 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""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. - """ - super().unfuse_lora(components=components, **kwargs) - - -class HunyuanVideoLoraLoaderMixin(LoraBaseMixin): - r""" - Load LoRA layers into [`HunyuanVideoTransformer3DModel`]. Specific to [`HunyuanVideoPipeline`]. - """ - - _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: Union[str, Dict[str, torch.Tensor]], - **kwargs, - ): - r""" - Return state dict for lora weights and the network alphas. - - - - We support loading original format HunyuanVideo LoRA checkpoints. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - """ - # 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} - - is_original_hunyuan_video = any("img_attn_qkv" in k for k in state_dict) - if is_original_hunyuan_video: - state_dict = _convert_hunyuan_video_lora_to_diffusers(state_dict) - - out = (state_dict, metadata) if return_lora_metadata else state_dict - return out - - # 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, - hotswap: bool = False, - **kwargs, - ): - """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - """ - 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.") - - 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->HunyuanVideoTransformer3DModel - def load_lora_into_transformer( - cls, - state_dict, - transformer, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - 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. - transformer (`HunyuanVideoTransformer3DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - 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: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[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, - ): - r""" - Save the LoRA parameters corresponding to the transformer. - - 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]`): - State dict of the LoRA layers corresponding to the `transformer`. - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. - """ - state_dict = {} - lora_adapter_metadata = {} - - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") - - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) - - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, - ) - - # 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: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - 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""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. - """ - super().unfuse_lora(components=components, **kwargs) - - -class Lumina2LoraLoaderMixin(LoraBaseMixin): - r""" - Load LoRA layers into [`Lumina2Transformer2DModel`]. Specific to [`Lumina2Text2ImgPipeline`]. - """ - - _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: Union[str, Dict[str, torch.Tensor]], - **kwargs, - ): - r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - """ - # 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} - - # conversion. - non_diffusers = any(k.startswith("diffusion_model.") for k in state_dict) - if non_diffusers: - state_dict = _convert_non_diffusers_lumina2_lora_to_diffusers(state_dict) - - out = (state_dict, metadata) if return_lora_metadata else state_dict - return out - - # 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, - hotswap: bool = False, - **kwargs, - ): - """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - """ - 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.") - - 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->Lumina2Transformer2DModel - def load_lora_into_transformer( - cls, - state_dict, - transformer, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - 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. - transformer (`Lumina2Transformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - 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: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[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, - ): - r""" - Save the LoRA parameters corresponding to the transformer. - - 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]`): - State dict of the LoRA layers corresponding to the `transformer`. - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. - """ - state_dict = {} - lora_adapter_metadata = {} - - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") - - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) - - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, - ) - - # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora - def fuse_lora( - self, - components: List[str] = ["transformer"], - lora_scale: float = 1.0, - safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - super().fuse_lora( - components=components, - lora_scale=lora_scale, - safe_fusing=safe_fusing, - adapter_names=adapter_names, - **kwargs, - ) - - # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): - r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. - """ - super().unfuse_lora(components=components, **kwargs) - - -class WanLoraLoaderMixin(LoraBaseMixin): - r""" - Load LoRA layers into [`WanTransformer3DModel`]. Specific to [`WanPipeline`] and `[WanImageToVideoPipeline`]. - """ - - _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: Union[str, Dict[str, torch.Tensor]], - **kwargs, - ): - r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - """ - # 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 - - @classmethod - def _maybe_expand_t2v_lora_for_i2v( - cls, - transformer: torch.nn.Module, - state_dict, - ): - if transformer.config.image_dim is None: - return state_dict - - target_device = transformer.device - - if any(k.startswith("transformer.blocks.") for k in state_dict): - num_blocks = len({k.split("blocks.")[1].split(".")[0] for k in state_dict if "blocks." in k}) - is_i2v_lora = any("add_k_proj" in k for k in state_dict) and any("add_v_proj" in k for k in state_dict) - has_bias = any(".lora_B.bias" in k for k in state_dict) - - if is_i2v_lora: - return state_dict - - for i in range(num_blocks): - for o, c in zip(["k_img", "v_img"], ["add_k_proj", "add_v_proj"]): - # These keys should exist if the block `i` was part of the T2V LoRA. - ref_key_lora_A = f"transformer.blocks.{i}.attn2.to_k.lora_A.weight" - ref_key_lora_B = f"transformer.blocks.{i}.attn2.to_k.lora_B.weight" - - if ref_key_lora_A not in state_dict or ref_key_lora_B not in state_dict: - continue - - state_dict[f"transformer.blocks.{i}.attn2.{c}.lora_A.weight"] = torch.zeros_like( - state_dict[f"transformer.blocks.{i}.attn2.to_k.lora_A.weight"], device=target_device - ) - state_dict[f"transformer.blocks.{i}.attn2.{c}.lora_B.weight"] = torch.zeros_like( - state_dict[f"transformer.blocks.{i}.attn2.to_k.lora_B.weight"], device=target_device - ) - - # If the original LoRA had biases (indicated by has_bias) - # AND the specific reference bias key exists for this block. - - ref_key_lora_B_bias = f"transformer.blocks.{i}.attn2.to_k.lora_B.bias" - if has_bias and ref_key_lora_B_bias in state_dict: - ref_lora_B_bias_tensor = state_dict[ref_key_lora_B_bias] - state_dict[f"transformer.blocks.{i}.attn2.{c}.lora_B.bias"] = torch.zeros_like( - ref_lora_B_bias_tensor, - device=target_device, - ) - - return 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, - hotswap: bool = False, - **kwargs, - ): - """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - """ - 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) - # convert T2V LoRA to I2V LoRA (when loaded to Wan I2V) by adding zeros for the additional (missing) _img layers - state_dict = self._maybe_expand_t2v_lora_for_i2v( - transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer, - state_dict=state_dict, - ) - is_correct_format = all("lora" in key for key in state_dict.keys()) - if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") - - 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, - ): - """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - 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. - transformer (`WanTransformer3DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - 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: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[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, - ): - r""" - Save the LoRA parameters corresponding to the transformer. - - 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]`): - State dict of the LoRA layers corresponding to the `transformer`. - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. - """ - state_dict = {} - lora_adapter_metadata = {} - - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") - - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) - - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, - ) - - # 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: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - 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""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. - """ - super().unfuse_lora(components=components, **kwargs) - - -class CogView4LoraLoaderMixin(LoraBaseMixin): - r""" - Load LoRA layers into [`WanTransformer3DModel`]. Specific to [`CogView4Pipeline`]. - """ - - _lora_loadable_modules = ["transformer"] - transformer_name = TRANSFORMER_NAME - - @classmethod - @validate_hf_hub_args - # 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]], - **kwargs, - ): - r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - - """ - # 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} - - out = (state_dict, metadata) if return_lora_metadata else state_dict - return out - - # 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, - hotswap: bool = False, - **kwargs, - ): - """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - """ - 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.") - - 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->CogView4Transformer2DModel - def load_lora_into_transformer( - cls, - state_dict, - transformer, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - 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. - transformer (`CogView4Transformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - 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: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[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, - ): - r""" - Save the LoRA parameters corresponding to the transformer. - - 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]`): - State dict of the LoRA layers corresponding to the `transformer`. - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. - """ - state_dict = {} - lora_adapter_metadata = {} - - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") - - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) - - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, - ) - - # 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: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - 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""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. - """ - super().unfuse_lora(components=components, **kwargs) - - -class HiDreamImageLoraLoaderMixin(LoraBaseMixin): - r""" - Load LoRA layers into [`HiDreamImageTransformer2DModel`]. Specific to [`HiDreamImagePipeline`]. - """ - - _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: Union[str, Dict[str, torch.Tensor]], - **kwargs, - ): - r""" - Return state dict for lora weights and the network alphas. - - - - We support loading A1111 formatted LoRA checkpoints in a limited capacity. - - This function is experimental and might change in the future. - - - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - return_lora_metadata (`bool`, *optional*, defaults to False): - When enabled, additionally return the LoRA adapter metadata, typically found in the state dict. - """ - # 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} - - is_non_diffusers_format = any("diffusion_model" in k for k in state_dict) - if is_non_diffusers_format: - state_dict = _convert_non_diffusers_hidream_lora_to_diffusers(state_dict) - - out = (state_dict, metadata) if return_lora_metadata else state_dict - return out - - # 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, - hotswap: bool = False, - **kwargs, - ): - """ - Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and - `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See - [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded. - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state - dict is loaded into `self.transformer`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - kwargs (`dict`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`]. - """ - 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.") - - 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->HiDreamImageTransformer2DModel - def load_lora_into_transformer( - cls, - state_dict, - transformer, - adapter_name=None, - _pipeline=None, - low_cpu_mem_usage=False, - hotswap: bool = False, - metadata=None, - ): - """ - This will load the LoRA layers specified in `state_dict` into `transformer`. - - Parameters: - state_dict (`dict`): - 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. - transformer (`HiDreamImageTransformer2DModel`): - The Transformer model to load the LoRA layers into. - adapter_name (`str`, *optional*): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap (`bool`, *optional*): - See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]. - metadata (`dict`): - Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived - from the state dict. - """ - 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: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[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, - ): - r""" - Save the LoRA parameters corresponding to the transformer. - - 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]`): - State dict of the LoRA layers corresponding to the `transformer`. - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - transformer_lora_adapter_metadata: - LoRA adapter metadata associated with the transformer to be serialized with the state dict. - """ - state_dict = {} - lora_adapter_metadata = {} - - if not transformer_lora_layers: - raise ValueError("You must pass `transformer_lora_layers`.") - - state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name)) - - if transformer_lora_adapter_metadata is not None: - lora_adapter_metadata.update( - _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name) - ) - - # Save the model - cls.write_lora_layers( - state_dict=state_dict, - save_directory=save_directory, - is_main_process=is_main_process, - weight_name=weight_name, - save_function=save_function, - safe_serialization=safe_serialization, - lora_adapter_metadata=lora_adapter_metadata, - ) - - # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora - def fuse_lora( - self, - components: List[str] = ["transformer"], - lora_scale: float = 1.0, - safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, - **kwargs, - ): - r""" - Fuses the LoRA parameters into the original parameters of the corresponding blocks. - - - - This is an experimental API. - - - - Args: - 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 to be used for fusing. If nothing is passed, all active adapters will be fused. - - Example: - - ```py - from diffusers import DiffusionPipeline - import torch - - pipeline = DiffusionPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.fuse_lora(lora_scale=0.7) - ``` - """ - super().fuse_lora( - components=components, - lora_scale=lora_scale, - safe_fusing=safe_fusing, - adapter_names=adapter_names, - **kwargs, - ) - - # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): - r""" - Reverses the effect of - [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). - - - - This is an experimental API. - - - - Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. - unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. - """ - super().unfuse_lora(components=components, **kwargs) - - -class LoraLoaderMixin(StableDiffusionLoraLoaderMixin): - def __init__(self, *args, **kwargs): - deprecation_message = "LoraLoaderMixin is deprecated and this will be removed in a future version. Please use `StableDiffusionLoraLoaderMixin`, instead." - deprecate("LoraLoaderMixin", "1.0.0", deprecation_message) - super().__init__(*args, **kwargs) diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/peft.py b/venv/lib/python3.10/site-packages/diffusers/loaders/peft.py deleted file mode 100644 index 343623071340ffef7db20ca11e58fa437d122a4c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/peft.py +++ /dev/null @@ -1,803 +0,0 @@ -# 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. -import inspect -import json -import os -from functools import partial -from pathlib import Path -from typing import Dict, List, Literal, Optional, Union - -import safetensors -import torch - -from ..utils import ( - MIN_PEFT_VERSION, - USE_PEFT_BACKEND, - check_peft_version, - convert_unet_state_dict_to_peft, - delete_adapter_layers, - get_adapter_name, - is_peft_available, - is_peft_version, - logging, - set_adapter_layers, - set_weights_and_activate_adapters, -) -from ..utils.peft_utils import _create_lora_config, _maybe_warn_for_unhandled_keys -from .lora_base import _fetch_state_dict, _func_optionally_disable_offloading -from .unet_loader_utils import _maybe_expand_lora_scales - - -logger = logging.get_logger(__name__) - -_SET_ADAPTER_SCALE_FN_MAPPING = { - "UNet2DConditionModel": _maybe_expand_lora_scales, - "UNetMotionModel": _maybe_expand_lora_scales, - "SD3Transformer2DModel": lambda model_cls, weights: weights, - "FluxTransformer2DModel": lambda model_cls, weights: weights, - "CogVideoXTransformer3DModel": lambda model_cls, weights: weights, - "ConsisIDTransformer3DModel": lambda model_cls, weights: weights, - "MochiTransformer3DModel": lambda model_cls, weights: weights, - "HunyuanVideoTransformer3DModel": lambda model_cls, weights: weights, - "LTXVideoTransformer3DModel": lambda model_cls, weights: weights, - "SanaTransformer2DModel": lambda model_cls, weights: weights, - "AuraFlowTransformer2DModel": lambda model_cls, weights: weights, - "Lumina2Transformer2DModel": lambda model_cls, weights: weights, - "WanTransformer3DModel": lambda model_cls, weights: weights, - "CogView4Transformer2DModel": lambda model_cls, weights: weights, - "HiDreamImageTransformer2DModel": lambda model_cls, weights: weights, - "HunyuanVideoFramepackTransformer3DModel": lambda model_cls, weights: weights, - "WanVACETransformer3DModel": lambda model_cls, weights: weights, - "ChromaTransformer2DModel": lambda model_cls, weights: weights, -} - - -class PeftAdapterMixin: - """ - A class containing all functions for loading and using adapters weights that are supported in PEFT library. For - more details about adapters and injecting them in a base model, check out the PEFT - [documentation](https://huggingface.co/docs/peft/index). - - Install the latest version of PEFT, and use this mixin to: - - - Attach new adapters in the model. - - Attach multiple adapters and iteratively activate/deactivate them. - - Activate/deactivate all adapters from the model. - - Get a list of the active adapters. - """ - - _hf_peft_config_loaded = False - # kwargs for prepare_model_for_compiled_hotswap, if required - _prepare_lora_hotswap_kwargs: Optional[dict] = None - - @classmethod - # Copied from diffusers.loaders.lora_base.LoraBaseMixin._optionally_disable_offloading - def _optionally_disable_offloading(cls, _pipeline): - return _func_optionally_disable_offloading(_pipeline=_pipeline) - - def load_lora_adapter( - self, pretrained_model_name_or_path_or_dict, prefix="transformer", hotswap: bool = False, **kwargs - ): - r""" - Loads a LoRA adapter into the underlying model. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - prefix (`str`, *optional*): Prefix to filter the state dict. - - 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 weights and configuration files, overriding the - cached versions if they exist. - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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]`): - 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). - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - hotswap : (`bool`, *optional*) - Defaults to `False`. Whether to substitute an existing (LoRA) adapter with the newly loaded adapter - in-place. This means that, instead of loading an additional adapter, this will take the existing - adapter weights and replace them with the weights of the new adapter. This can be faster and more - memory efficient. However, the main advantage of hotswapping is that when the model is compiled with - torch.compile, loading the new adapter does not require recompilation of the model. When using - hotswapping, the passed `adapter_name` should be the name of an already loaded adapter. - - If the new adapter and the old adapter have different ranks and/or LoRA alphas (i.e. scaling), you need - to call an additional method before loading the adapter: - - ```py - pipeline = ... # load diffusers pipeline - max_rank = ... # the highest rank among all LoRAs that you want to load - # call *before* compiling and loading the LoRA adapter - pipeline.enable_lora_hotswap(target_rank=max_rank) - pipeline.load_lora_weights(file_name) - # optionally compile the model now - ``` - - Note that hotswapping adapters of the text encoder is not yet supported. There are some further - limitations to this technique, which are documented here: - https://huggingface.co/docs/peft/main/en/package_reference/hotswap - metadata: - LoRA adapter metadata. When supplied, the metadata inferred through the state dict isn't used to - initialize `LoraConfig`. - """ - from peft import inject_adapter_in_model, set_peft_model_state_dict - from peft.tuners.tuners_utils import BaseTunerLayer - - 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) - adapter_name = kwargs.pop("adapter_name", None) - network_alphas = kwargs.pop("network_alphas", None) - _pipeline = kwargs.pop("_pipeline", None) - low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", False) - metadata = kwargs.pop("metadata", None) - allow_pickle = False - - 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`." - ) - - 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, - metadata=metadata, - ) - if network_alphas is not None and prefix is None: - raise ValueError("`network_alphas` cannot be None when `prefix` is None.") - if network_alphas and metadata: - raise ValueError("Both `network_alphas` and `metadata` cannot be specified.") - - if prefix is not None: - state_dict = {k.removeprefix(f"{prefix}."): v for k, v in state_dict.items() if k.startswith(f"{prefix}.")} - if metadata is not None: - metadata = {k.removeprefix(f"{prefix}."): v for k, v in metadata.items() if k.startswith(f"{prefix}.")} - - if len(state_dict) > 0: - if adapter_name in getattr(self, "peft_config", {}) and not hotswap: - raise ValueError( - f"Adapter name {adapter_name} already in use in the model - please select a new adapter name." - ) - elif adapter_name not in getattr(self, "peft_config", {}) and hotswap: - raise ValueError( - f"Trying to hotswap LoRA adapter '{adapter_name}' but there is no existing adapter by that name. " - "Please choose an existing adapter name or set `hotswap=False` to prevent hotswapping." - ) - - # check with first key if is not in peft format - first_key = next(iter(state_dict.keys())) - if "lora_A" not in first_key: - state_dict = convert_unet_state_dict_to_peft(state_dict) - - rank = {} - for key, val in state_dict.items(): - # Cannot figure out rank from lora layers that don't have at least 2 dimensions. - # Bias layers in LoRA only have a single dimension - if "lora_B" in key and val.ndim > 1: - # Check out https://github.com/huggingface/peft/pull/2419 for the `^` symbol. - # We may run into some ambiguous configuration values when a model has module - # names, sharing a common prefix (`proj_out.weight` and `blocks.transformer.proj_out.weight`, - # for example) and they have different LoRA ranks. - rank[f"^{key}"] = val.shape[1] - - if network_alphas is not None and len(network_alphas) >= 1: - alpha_keys = [k for k in network_alphas.keys() if k.startswith(f"{prefix}.")] - network_alphas = { - k.removeprefix(f"{prefix}."): v for k, v in network_alphas.items() if k in alpha_keys - } - - # create LoraConfig - lora_config = _create_lora_config(state_dict, network_alphas, metadata, rank) - - # adapter_name - if adapter_name is None: - adapter_name = get_adapter_name(self) - - # =", "0.13.1"): - peft_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage - - if hotswap or (self._prepare_lora_hotswap_kwargs is not None): - if is_peft_version(">", "0.14.0"): - from peft.utils.hotswap import ( - check_hotswap_configs_compatible, - hotswap_adapter_from_state_dict, - prepare_model_for_compiled_hotswap, - ) - else: - msg = ( - "Hotswapping requires PEFT > v0.14. Please upgrade PEFT to a higher version or install it " - "from source." - ) - raise ImportError(msg) - - if hotswap: - - def map_state_dict_for_hotswap(sd): - # For hotswapping, we need the adapter name to be present in the state dict keys - new_sd = {} - for k, v in sd.items(): - if k.endswith("lora_A.weight") or key.endswith("lora_B.weight"): - k = k[: -len(".weight")] + f".{adapter_name}.weight" - elif k.endswith("lora_B.bias"): # lora_bias=True option - k = k[: -len(".bias")] + f".{adapter_name}.bias" - new_sd[k] = v - return new_sd - - # To handle scenarios where we cannot successfully set state dict. If it's unsuccessful, - # we should also delete the `peft_config` associated to the `adapter_name`. - try: - if hotswap: - state_dict = map_state_dict_for_hotswap(state_dict) - check_hotswap_configs_compatible(self.peft_config[adapter_name], lora_config) - try: - hotswap_adapter_from_state_dict( - model=self, - state_dict=state_dict, - adapter_name=adapter_name, - config=lora_config, - ) - except Exception as e: - logger.error(f"Hotswapping {adapter_name} was unsuccessful with the following error: \n{e}") - raise - # the hotswap function raises if there are incompatible keys, so if we reach this point we can set - # it to None - incompatible_keys = None - else: - inject_adapter_in_model(lora_config, self, adapter_name=adapter_name, **peft_kwargs) - incompatible_keys = set_peft_model_state_dict(self, state_dict, adapter_name, **peft_kwargs) - - if self._prepare_lora_hotswap_kwargs is not None: - # For hotswapping of compiled models or adapters with different ranks. - # If the user called enable_lora_hotswap, we need to ensure it is called: - # - after the first adapter was loaded - # - before the model is compiled and the 2nd adapter is being hotswapped in - # Therefore, it needs to be called here - prepare_model_for_compiled_hotswap( - self, config=lora_config, **self._prepare_lora_hotswap_kwargs - ) - # We only want to call prepare_model_for_compiled_hotswap once - self._prepare_lora_hotswap_kwargs = None - - # Set peft config loaded flag to True if module has been successfully injected and incompatible keys retrieved - if not self._hf_peft_config_loaded: - self._hf_peft_config_loaded = True - except Exception as e: - # In case `inject_adapter_in_model()` was unsuccessful even before injecting the `peft_config`. - if hasattr(self, "peft_config"): - for module in self.modules(): - if isinstance(module, BaseTunerLayer): - active_adapters = module.active_adapters - for active_adapter in active_adapters: - if adapter_name in active_adapter: - module.delete_adapter(adapter_name) - - self.peft_config.pop(adapter_name) - logger.error(f"Loading {adapter_name} was unsuccessful with the following error: \n{e}") - raise - - _maybe_warn_for_unhandled_keys(incompatible_keys, adapter_name) - - # Offload back. - if is_model_cpu_offload: - _pipeline.enable_model_cpu_offload() - elif is_sequential_cpu_offload: - _pipeline.enable_sequential_cpu_offload() - # Unsafe code /> - - if prefix is not None and not state_dict: - model_class_name = self.__class__.__name__ - logger.warning( - f"No LoRA keys associated to {model_class_name} found with the {prefix=}. " - "This is safe to ignore if LoRA state dict didn't originally have any " - f"{model_class_name} related params. You can also try specifying `prefix=None` " - "to resolve the warning. Otherwise, open an issue if you think it's unexpected: " - "https://github.com/huggingface/diffusers/issues/new" - ) - - def save_lora_adapter( - self, - save_directory, - adapter_name: str = "default", - upcast_before_saving: bool = False, - safe_serialization: bool = True, - weight_name: Optional[str] = None, - ): - """ - Save the LoRA parameters corresponding to the underlying model. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save LoRA parameters to. Will be created if it doesn't exist. - adapter_name: (`str`, defaults to "default"): The name of the adapter to serialize. Useful when the - underlying model has multiple adapters loaded. - upcast_before_saving (`bool`, defaults to `False`): - Whether to cast the underlying model to `torch.float32` before serialization. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. - weight_name: (`str`, *optional*, defaults to `None`): Name of the file to serialize the state dict with. - """ - from peft.utils import get_peft_model_state_dict - - from .lora_base import LORA_ADAPTER_METADATA_KEY, LORA_WEIGHT_NAME, LORA_WEIGHT_NAME_SAFE - - if adapter_name is None: - adapter_name = get_adapter_name(self) - - if adapter_name not in getattr(self, "peft_config", {}): - raise ValueError(f"Adapter name {adapter_name} not found in the model.") - - lora_adapter_metadata = self.peft_config[adapter_name].to_dict() - - lora_layers_to_save = get_peft_model_state_dict( - self.to(dtype=torch.float32 if upcast_before_saving else None), adapter_name=adapter_name - ) - if os.path.isfile(save_directory): - raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file") - - if safe_serialization: - - def save_function(weights, filename): - # Inject framework format. - metadata = {"format": "pt"} - if lora_adapter_metadata is not None: - 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) - - return safetensors.torch.save_file(weights, filename, metadata=metadata) - - else: - save_function = torch.save - - os.makedirs(save_directory, exist_ok=True) - - if weight_name is None: - if safe_serialization: - weight_name = LORA_WEIGHT_NAME_SAFE - else: - weight_name = LORA_WEIGHT_NAME - - save_path = Path(save_directory, weight_name).as_posix() - save_function(lora_layers_to_save, save_path) - logger.info(f"Model weights saved in {save_path}") - - def set_adapters( - self, - adapter_names: Union[List[str], str], - weights: Optional[Union[float, Dict, List[float], List[Dict], List[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`): - The names of the adapters to use. - adapter_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. - - Example: - - ```py - from diffusers import AutoPipelineForText2Image - import torch - - pipeline = AutoPipelineForText2Image.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights( - "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]) - ``` - """ - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for `set_adapters()`.") - - adapter_names = [adapter_names] if isinstance(adapter_names, str) else adapter_names - - # Expand weights into a list, one entry per adapter - # examples for e.g. 2 adapters: [{...}, 7] -> [7,7] ; None -> [None, None] - if not isinstance(weights, list): - weights = [weights] * len(adapter_names) - - if len(adapter_names) != len(weights): - raise ValueError( - f"Length of adapter names {len(adapter_names)} is not equal to the length of their weights {len(weights)}." - ) - - # Set None values to default of 1.0 - # e.g. [{...}, 7] -> [{...}, 7] ; [None, None] -> [1.0, 1.0] - weights = [w if w is not None else 1.0 for w in weights] - - # e.g. [{...}, 7] -> [{expanded dict...}, 7] - scale_expansion_fn = _SET_ADAPTER_SCALE_FN_MAPPING[self.__class__.__name__] - weights = scale_expansion_fn(self, weights) - - set_weights_and_activate_adapters(self, adapter_names, weights) - - def add_adapter(self, adapter_config, adapter_name: str = "default") -> None: - r""" - Adds a new adapter to the current model for training. If no adapter name is passed, a default name is assigned - to the adapter to follow the convention of the PEFT library. - - If you are not familiar with adapters and PEFT methods, we invite you to read more about them in the PEFT - [documentation](https://huggingface.co/docs/peft). - - Args: - adapter_config (`[~peft.PeftConfig]`): - The configuration of the adapter to add; supported adapters are non-prefix tuning and adaption prompt - methods. - adapter_name (`str`, *optional*, defaults to `"default"`): - The name of the adapter to add. If no name is passed, a default name is assigned to the adapter. - """ - check_peft_version(min_version=MIN_PEFT_VERSION) - - if not is_peft_available(): - raise ImportError("PEFT is not available. Please install PEFT to use this function: `pip install peft`.") - - from peft import PeftConfig, inject_adapter_in_model - - if not self._hf_peft_config_loaded: - self._hf_peft_config_loaded = True - elif adapter_name in self.peft_config: - raise ValueError(f"Adapter with name {adapter_name} already exists. Please use a different name.") - - if not isinstance(adapter_config, PeftConfig): - raise ValueError( - f"adapter_config should be an instance of PeftConfig. Got {type(adapter_config)} instead." - ) - - # Unlike transformers, here we don't need to retrieve the name_or_path of the unet as the loading logic is - # handled by the `load_lora_layers` or `StableDiffusionLoraLoaderMixin`. Therefore we set it to `None` here. - adapter_config.base_model_name_or_path = 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: - """ - Sets a specific adapter by forcing the model to only use that adapter and disables the other adapters. - - If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT - [documentation](https://huggingface.co/docs/peft). - - Args: - adapter_name (Union[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) - - if not self._hf_peft_config_loaded: - raise ValueError("No adapter loaded. Please load an adapter first.") - - if isinstance(adapter_name, str): - adapter_name = [adapter_name] - - missing = set(adapter_name) - set(self.peft_config) - if len(missing) > 0: - raise ValueError( - f"Following adapter(s) could not be found: {', '.join(missing)}. Make sure you are passing the correct adapter name(s)." - f" current loaded adapters are: {list(self.peft_config.keys())}" - ) - - from peft.tuners.tuners_utils import BaseTunerLayer - - _adapters_has_been_set = False - - for _, module in self.named_modules(): - if isinstance(module, BaseTunerLayer): - if hasattr(module, "set_adapter"): - module.set_adapter(adapter_name) - # Previous versions of PEFT does not support multi-adapter inference - elif not hasattr(module, "set_adapter") and len(adapter_name) != 1: - raise ValueError( - "You are trying to set multiple adapters and you have a PEFT version that does not support multi-adapter inference. Please upgrade to the latest version of PEFT." - " `pip install -U peft` or `pip install -U git+https://github.com/huggingface/peft.git`" - ) - else: - module.active_adapter = adapter_name - _adapters_has_been_set = True - - if not _adapters_has_been_set: - raise ValueError( - "Did not succeeded in setting the adapter. Please make sure you are using a model that supports adapters." - ) - - def disable_adapters(self) -> None: - r""" - Disable all adapters attached to the model and fallback to inference with the base model only. - - If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT - [documentation](https://huggingface.co/docs/peft). - """ - check_peft_version(min_version=MIN_PEFT_VERSION) - - if not self._hf_peft_config_loaded: - raise ValueError("No adapter loaded. Please load an adapter first.") - - from peft.tuners.tuners_utils import BaseTunerLayer - - for _, module in self.named_modules(): - if isinstance(module, BaseTunerLayer): - if hasattr(module, "enable_adapters"): - module.enable_adapters(enabled=False) - else: - # support for older PEFT versions - module.disable_adapters = True - - def enable_adapters(self) -> None: - """ - Enable adapters that are attached to the model. The model uses `self.active_adapters()` to retrieve the list of - adapters to enable. - - If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT - [documentation](https://huggingface.co/docs/peft). - """ - check_peft_version(min_version=MIN_PEFT_VERSION) - - if not self._hf_peft_config_loaded: - raise ValueError("No adapter loaded. Please load an adapter first.") - - from peft.tuners.tuners_utils import BaseTunerLayer - - for _, module in self.named_modules(): - if isinstance(module, BaseTunerLayer): - if hasattr(module, "enable_adapters"): - module.enable_adapters(enabled=True) - else: - # support for older PEFT versions - module.disable_adapters = False - - def active_adapters(self) -> List[str]: - """ - Gets the current list of active adapters of the model. - - If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT - [documentation](https://huggingface.co/docs/peft). - """ - check_peft_version(min_version=MIN_PEFT_VERSION) - - if not is_peft_available(): - raise ImportError("PEFT is not available. Please install PEFT to use this function: `pip install peft`.") - - if not self._hf_peft_config_loaded: - raise ValueError("No adapter loaded. Please load an adapter first.") - - from peft.tuners.tuners_utils import BaseTunerLayer - - for _, module in self.named_modules(): - if isinstance(module, BaseTunerLayer): - return module.active_adapter - - def fuse_lora(self, lora_scale=1.0, safe_fusing=False, adapter_names=None): - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for `fuse_lora()`.") - - self.lora_scale = lora_scale - self._safe_fusing = safe_fusing - self.apply(partial(self._fuse_lora_apply, adapter_names=adapter_names)) - - def _fuse_lora_apply(self, module, adapter_names=None): - from peft.tuners.tuners_utils import BaseTunerLayer - - merge_kwargs = {"safe_merge": self._safe_fusing} - - if isinstance(module, BaseTunerLayer): - if self.lora_scale != 1.0: - module.scale_layer(self.lora_scale) - - # For BC with previous PEFT versions, we need to check the signature - # of the `merge` method to see if it supports the `adapter_names` argument. - supported_merge_kwargs = list(inspect.signature(module.merge).parameters) - if "adapter_names" in supported_merge_kwargs: - merge_kwargs["adapter_names"] = adapter_names - elif "adapter_names" not in supported_merge_kwargs and adapter_names is not None: - raise ValueError( - "The `adapter_names` argument is not supported with your PEFT version. Please upgrade" - " to the latest version of PEFT. `pip install -U peft`" - ) - - module.merge(**merge_kwargs) - - def unfuse_lora(self): - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for `unfuse_lora()`.") - self.apply(self._unfuse_lora_apply) - - def _unfuse_lora_apply(self, module): - from peft.tuners.tuners_utils import BaseTunerLayer - - if isinstance(module, BaseTunerLayer): - module.unmerge() - - def unload_lora(self): - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for `unload_lora()`.") - - from ..utils import recurse_remove_peft_layers - - recurse_remove_peft_layers(self) - if hasattr(self, "peft_config"): - del self.peft_config - - def disable_lora(self): - """ - Disables the active LoRA layers of the underlying model. - - Example: - - ```py - from diffusers import AutoPipelineForText2Image - import torch - - pipeline = AutoPipelineForText2Image.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights( - "jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_name="cinematic" - ) - pipeline.unet.disable_lora() - ``` - """ - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for this method.") - set_adapter_layers(self, enabled=False) - - def enable_lora(self): - """ - Enables the active LoRA layers of the underlying model. - - Example: - - ```py - from diffusers import AutoPipelineForText2Image - import torch - - pipeline = AutoPipelineForText2Image.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights( - "jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_name="cinematic" - ) - pipeline.unet.enable_lora() - ``` - """ - if not USE_PEFT_BACKEND: - 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]): - """ - Delete an adapter's LoRA layers from the underlying model. - - Args: - adapter_names (`Union[List[str], str]`): - The names (single string or list of strings) of the adapter to delete. - - Example: - - ```py - from diffusers import AutoPipelineForText2Image - import torch - - pipeline = AutoPipelineForText2Image.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.load_lora_weights( - "jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_names="cinematic" - ) - pipeline.unet.delete_adapters("cinematic") - ``` - """ - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for this method.") - - if isinstance(adapter_names, str): - adapter_names = [adapter_names] - - for adapter_name in adapter_names: - delete_adapter_layers(self, adapter_name) - - # Pop also the corresponding adapter from the config - if hasattr(self, "peft_config"): - self.peft_config.pop(adapter_name, None) - - def enable_lora_hotswap( - self, target_rank: int = 128, check_compiled: Literal["error", "warn", "ignore"] = "error" - ) -> None: - """Enables the possibility to hotswap LoRA adapters. - - Calling this method is only required when hotswapping adapters and if the model is compiled or if the ranks of - the loaded adapters differ. - - Args: - target_rank (`int`, *optional*, defaults to `128`): - The highest rank among all the adapters that will be loaded. - - check_compiled (`str`, *optional*, defaults to `"error"`): - How to handle the case when the model is already compiled, which should generally be avoided. The - options are: - - "error" (default): raise an error - - "warn": issue a warning - - "ignore": do nothing - """ - if getattr(self, "peft_config", {}): - if check_compiled == "error": - raise RuntimeError("Call `enable_lora_hotswap` before loading the first adapter.") - elif check_compiled == "warn": - logger.warning( - "It is recommended to call `enable_lora_hotswap` before loading the first adapter to avoid recompilation." - ) - elif check_compiled != "ignore": - raise ValueError( - f"check_compiles should be one of 'error', 'warn', or 'ignore', got '{check_compiled}' instead." - ) - - self._prepare_lora_hotswap_kwargs = {"target_rank": target_rank, "check_compiled": check_compiled} diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/single_file.py b/venv/lib/python3.10/site-packages/diffusers/loaders/single_file.py deleted file mode 100644 index 667f79437985c3e9b2e0fe73e05de8caeb351860..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/single_file.py +++ /dev/null @@ -1,565 +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 -import os - -import torch -from huggingface_hub import snapshot_download -from huggingface_hub.utils import LocalEntryNotFoundError, validate_hf_hub_args -from packaging import version -from typing_extensions import Self - -from ..utils import deprecate, is_transformers_available, logging -from .single_file_utils import ( - SingleFileComponentError, - _is_legacy_scheduler_kwargs, - _is_model_weights_in_cached_folder, - _legacy_load_clip_tokenizer, - _legacy_load_safety_checker, - _legacy_load_scheduler, - create_diffusers_clip_model_from_ldm, - create_diffusers_t5_model_from_checkpoint, - fetch_diffusers_config, - fetch_original_config, - is_clip_model_in_single_file, - is_t5_in_single_file, - load_single_file_checkpoint, -) - - -logger = logging.get_logger(__name__) - -# Legacy behaviour. `from_single_file` does not load the safety checker unless explicitly provided -SINGLE_FILE_OPTIONAL_COMPONENTS = ["safety_checker"] - -if is_transformers_available(): - import transformers - from transformers import PreTrainedModel, PreTrainedTokenizer - - -def load_single_file_sub_model( - library_name, - class_name, - name, - checkpoint, - pipelines, - is_pipeline_module, - cached_model_config_path, - original_config=None, - local_files_only=False, - torch_dtype=None, - is_legacy_loading=False, - disable_mmap=False, - **kwargs, -): - if is_pipeline_module: - pipeline_module = getattr(pipelines, library_name) - class_obj = getattr(pipeline_module, class_name) - else: - # else we just import it from the library. - library = importlib.import_module(library_name) - class_obj = getattr(library, class_name) - - if is_transformers_available(): - transformers_version = version.parse(version.parse(transformers.__version__).base_version) - else: - transformers_version = "N/A" - - is_transformers_model = ( - is_transformers_available() - and issubclass(class_obj, PreTrainedModel) - and transformers_version >= version.parse("4.20.0") - ) - is_tokenizer = ( - is_transformers_available() - and issubclass(class_obj, PreTrainedTokenizer) - and transformers_version >= version.parse("4.20.0") - ) - - diffusers_module = importlib.import_module(__name__.split(".")[0]) - is_diffusers_single_file_model = issubclass(class_obj, diffusers_module.FromOriginalModelMixin) - is_diffusers_model = issubclass(class_obj, diffusers_module.ModelMixin) - is_diffusers_scheduler = issubclass(class_obj, diffusers_module.SchedulerMixin) - - if is_diffusers_single_file_model: - load_method = getattr(class_obj, "from_single_file") - - # We cannot provide two different config options to the `from_single_file` method - # Here we have to ignore loading the config from `cached_model_config_path` if `original_config` is provided - if original_config: - cached_model_config_path = None - - loaded_sub_model = load_method( - pretrained_model_link_or_path_or_dict=checkpoint, - original_config=original_config, - config=cached_model_config_path, - subfolder=name, - torch_dtype=torch_dtype, - local_files_only=local_files_only, - disable_mmap=disable_mmap, - **kwargs, - ) - - elif is_transformers_model and is_clip_model_in_single_file(class_obj, checkpoint): - loaded_sub_model = create_diffusers_clip_model_from_ldm( - class_obj, - checkpoint=checkpoint, - config=cached_model_config_path, - subfolder=name, - torch_dtype=torch_dtype, - local_files_only=local_files_only, - is_legacy_loading=is_legacy_loading, - ) - - elif is_transformers_model and is_t5_in_single_file(checkpoint): - loaded_sub_model = create_diffusers_t5_model_from_checkpoint( - class_obj, - checkpoint=checkpoint, - config=cached_model_config_path, - subfolder=name, - torch_dtype=torch_dtype, - local_files_only=local_files_only, - ) - - elif is_tokenizer and is_legacy_loading: - loaded_sub_model = _legacy_load_clip_tokenizer( - class_obj, checkpoint=checkpoint, config=cached_model_config_path, local_files_only=local_files_only - ) - - elif is_diffusers_scheduler and (is_legacy_loading or _is_legacy_scheduler_kwargs(kwargs)): - loaded_sub_model = _legacy_load_scheduler( - class_obj, checkpoint=checkpoint, component_name=name, original_config=original_config, **kwargs - ) - - else: - if not hasattr(class_obj, "from_pretrained"): - raise ValueError( - ( - f"The component {class_obj.__name__} cannot be loaded as it does not seem to have" - " a supported loading method." - ) - ) - - loading_kwargs = {} - loading_kwargs.update( - { - "pretrained_model_name_or_path": cached_model_config_path, - "subfolder": name, - "local_files_only": local_files_only, - } - ) - - # Schedulers and Tokenizers don't make use of torch_dtype - # Skip passing it to those objects - if issubclass(class_obj, torch.nn.Module): - loading_kwargs.update({"torch_dtype": torch_dtype}) - - if is_diffusers_model or is_transformers_model: - if not _is_model_weights_in_cached_folder(cached_model_config_path, name): - raise SingleFileComponentError( - f"Failed to load {class_name}. Weights for this component appear to be missing in the checkpoint." - ) - - load_method = getattr(class_obj, "from_pretrained") - loaded_sub_model = load_method(**loading_kwargs) - - return loaded_sub_model - - -def _map_component_types_to_config_dict(component_types): - diffusers_module = importlib.import_module(__name__.split(".")[0]) - config_dict = {} - component_types.pop("self", None) - - if is_transformers_available(): - transformers_version = version.parse(version.parse(transformers.__version__).base_version) - else: - transformers_version = "N/A" - - for component_name, component_value in component_types.items(): - is_diffusers_model = issubclass(component_value[0], diffusers_module.ModelMixin) - is_scheduler_enum = component_value[0].__name__ == "KarrasDiffusionSchedulers" - is_scheduler = issubclass(component_value[0], diffusers_module.SchedulerMixin) - - is_transformers_model = ( - is_transformers_available() - and issubclass(component_value[0], PreTrainedModel) - and transformers_version >= version.parse("4.20.0") - ) - is_transformers_tokenizer = ( - is_transformers_available() - and issubclass(component_value[0], PreTrainedTokenizer) - and transformers_version >= version.parse("4.20.0") - ) - - if is_diffusers_model and component_name not in SINGLE_FILE_OPTIONAL_COMPONENTS: - config_dict[component_name] = ["diffusers", component_value[0].__name__] - - elif is_scheduler_enum or is_scheduler: - if is_scheduler_enum: - # Since we cannot fetch a scheduler config from the hub, we default to DDIMScheduler - # if the type hint is a KarrassDiffusionSchedulers enum - config_dict[component_name] = ["diffusers", "DDIMScheduler"] - - elif is_scheduler: - config_dict[component_name] = ["diffusers", component_value[0].__name__] - - elif ( - is_transformers_model or is_transformers_tokenizer - ) and component_name not in SINGLE_FILE_OPTIONAL_COMPONENTS: - config_dict[component_name] = ["transformers", component_value[0].__name__] - - else: - config_dict[component_name] = [None, None] - - return config_dict - - -def _infer_pipeline_config_dict(pipeline_class): - parameters = inspect.signature(pipeline_class.__init__).parameters - required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty} - component_types = pipeline_class._get_signature_types() - - # Ignore parameters that are not required for the pipeline - component_types = {k: v for k, v in component_types.items() if k in required_parameters} - config_dict = _map_component_types_to_config_dict(component_types) - - return config_dict - - -def _download_diffusers_model_config_from_hub( - pretrained_model_name_or_path, - cache_dir, - revision, - proxies, - force_download=None, - local_files_only=None, - token=None, -): - allow_patterns = ["**/*.json", "*.json", "*.txt", "**/*.txt", "**/*.model"] - cached_model_path = snapshot_download( - pretrained_model_name_or_path, - cache_dir=cache_dir, - revision=revision, - proxies=proxies, - force_download=force_download, - local_files_only=local_files_only, - token=token, - allow_patterns=allow_patterns, - ) - - return cached_model_path - - -class FromSingleFileMixin: - """ - Load model weights saved in the `.ckpt` format into a [`DiffusionPipeline`]. - """ - - @classmethod - @validate_hf_hub_args - def from_single_file(cls, pretrained_model_link_or_path, **kwargs) -> Self: - r""" - Instantiate a [`DiffusionPipeline`] from pretrained pipeline weights saved in the `.ckpt` or `.safetensors` - format. The pipeline is set in evaluation mode (`model.eval()`) by default. - - Parameters: - pretrained_model_link_or_path (`str` or `os.PathLike`, *optional*): - Can be either: - - A link to the `.ckpt` file (for example - `"https://huggingface.co//blob/main/.ckpt"`) on the Hub. - - A path to a *file* containing all pipeline weights. - torch_dtype (`str` or `torch.dtype`, *optional*): - Override the default `torch.dtype` and load the model with another dtype. - 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*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - original_config_file (`str`, *optional*): - The path to the original config file that was used to train the model. If not provided, the config file - will be inferred from the checkpoint file. - config (`str`, *optional*): - Can be either: - - A string, the *repo id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained pipeline - hosted on the Hub. - - A path to a *directory* (for example `./my_pipeline_directory/`) containing the pipeline - component configs in Diffusers format. - 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. - kwargs (remaining dictionary of keyword arguments, *optional*): - Can be used to overwrite load and saveable variables (the pipeline components of the specific pipeline - class). The overwritten components are passed directly to the pipelines `__init__` method. See example - below for more information. - - Examples: - - ```py - >>> from diffusers import StableDiffusionPipeline - - >>> # Download pipeline from huggingface.co and cache. - >>> pipeline = StableDiffusionPipeline.from_single_file( - ... "https://huggingface.co/WarriorMama777/OrangeMixs/blob/main/Models/AbyssOrangeMix/AbyssOrangeMix.safetensors" - ... ) - - >>> # Download pipeline from local file - >>> # file is downloaded under ./v1-5-pruned-emaonly.ckpt - >>> pipeline = StableDiffusionPipeline.from_single_file("./v1-5-pruned-emaonly.ckpt") - - >>> # Enable float16 and move to GPU - >>> pipeline = StableDiffusionPipeline.from_single_file( - ... "https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/blob/main/v1-5-pruned-emaonly.ckpt", - ... torch_dtype=torch.float16, - ... ) - >>> pipeline.to("cuda") - ``` - - """ - original_config_file = kwargs.pop("original_config_file", None) - config = kwargs.pop("config", None) - original_config = kwargs.pop("original_config", None) - - if original_config_file is not None: - deprecation_message = ( - "`original_config_file` argument is deprecated and will be removed in future versions." - "please use the `original_config` argument instead." - ) - deprecate("original_config_file", "1.0.0", deprecation_message) - original_config = original_config_file - - force_download = kwargs.pop("force_download", False) - proxies = kwargs.pop("proxies", None) - token = kwargs.pop("token", None) - cache_dir = kwargs.pop("cache_dir", None) - local_files_only = kwargs.pop("local_files_only", False) - revision = kwargs.pop("revision", None) - torch_dtype = kwargs.pop("torch_dtype", None) - disable_mmap = kwargs.pop("disable_mmap", False) - - is_legacy_loading = False - - if torch_dtype is not None and not isinstance(torch_dtype, torch.dtype): - torch_dtype = torch.float32 - logger.warning( - f"Passed `torch_dtype` {torch_dtype} is not a `torch.dtype`. Defaulting to `torch.float32`." - ) - - # We shouldn't allow configuring individual models components through a Pipeline creation method - # These model kwargs should be deprecated - scaling_factor = kwargs.get("scaling_factor", None) - if scaling_factor is not None: - deprecation_message = ( - "Passing the `scaling_factor` argument to `from_single_file is deprecated " - "and will be ignored in future versions." - ) - deprecate("scaling_factor", "1.0.0", deprecation_message) - - if original_config is not None: - original_config = fetch_original_config(original_config, local_files_only=local_files_only) - - from ..pipelines.pipeline_utils import _get_pipeline_class - - pipeline_class = _get_pipeline_class(cls, config=None) - - checkpoint = load_single_file_checkpoint( - pretrained_model_link_or_path, - force_download=force_download, - proxies=proxies, - token=token, - cache_dir=cache_dir, - local_files_only=local_files_only, - revision=revision, - disable_mmap=disable_mmap, - ) - - if config is None: - config = fetch_diffusers_config(checkpoint) - default_pretrained_model_config_name = config["pretrained_model_name_or_path"] - else: - default_pretrained_model_config_name = config - - if not os.path.isdir(default_pretrained_model_config_name): - # Provided config is a repo_id - if default_pretrained_model_config_name.count("/") > 1: - raise ValueError( - f'The provided config "{config}"' - " is neither a valid local path nor a valid repo id. Please check the parameter." - ) - try: - # Attempt to download the config files for the pipeline - cached_model_config_path = _download_diffusers_model_config_from_hub( - default_pretrained_model_config_name, - cache_dir=cache_dir, - revision=revision, - proxies=proxies, - force_download=force_download, - local_files_only=local_files_only, - token=token, - ) - config_dict = pipeline_class.load_config(cached_model_config_path) - - except LocalEntryNotFoundError: - # `local_files_only=True` but a local diffusers format model config is not available in the cache - # If `original_config` is not provided, we need override `local_files_only` to False - # to fetch the config files from the hub so that we have a way - # to configure the pipeline components. - - if original_config is None: - logger.warning( - "`local_files_only` is True but no local configs were found for this checkpoint.\n" - "Attempting to download the necessary config files for this pipeline.\n" - ) - cached_model_config_path = _download_diffusers_model_config_from_hub( - default_pretrained_model_config_name, - cache_dir=cache_dir, - revision=revision, - proxies=proxies, - force_download=force_download, - local_files_only=False, - token=token, - ) - config_dict = pipeline_class.load_config(cached_model_config_path) - - else: - # For backwards compatibility - # If `original_config` is provided, then we need to assume we are using legacy loading for pipeline components - logger.warning( - "Detected legacy `from_single_file` loading behavior. Attempting to create the pipeline based on inferred components.\n" - "This may lead to errors if the model components are not correctly inferred. \n" - "To avoid this warning, please explicitly pass the `config` argument to `from_single_file` with a path to a local diffusers model repo \n" - "e.g. `from_single_file(, config=) \n" - "or run `from_single_file` with `local_files_only=False` first to update the local cache directory with " - "the necessary config files.\n" - ) - is_legacy_loading = True - cached_model_config_path = None - - config_dict = _infer_pipeline_config_dict(pipeline_class) - config_dict["_class_name"] = pipeline_class.__name__ - - else: - # Provided config is a path to a local directory attempt to load directly. - cached_model_config_path = default_pretrained_model_config_name - config_dict = pipeline_class.load_config(cached_model_config_path) - - # pop out "_ignore_files" as it is only needed for download - config_dict.pop("_ignore_files", None) - - expected_modules, optional_kwargs = pipeline_class._get_signature_keys(cls) - passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs} - passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs} - - init_dict, unused_kwargs, _ = pipeline_class.extract_init_dict(config_dict, **kwargs) - init_kwargs = {k: init_dict.pop(k) for k in optional_kwargs if k in init_dict} - init_kwargs = {**init_kwargs, **passed_pipe_kwargs} - - from diffusers import pipelines - - # remove `null` components - def load_module(name, value): - if value[0] is None: - return False - if name in passed_class_obj and passed_class_obj[name] is None: - return False - if name in SINGLE_FILE_OPTIONAL_COMPONENTS: - return False - - return True - - init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)} - - for name, (library_name, class_name) in logging.tqdm( - sorted(init_dict.items()), desc="Loading pipeline components..." - ): - loaded_sub_model = None - is_pipeline_module = hasattr(pipelines, library_name) - - if name in passed_class_obj: - loaded_sub_model = passed_class_obj[name] - - else: - try: - loaded_sub_model = load_single_file_sub_model( - library_name=library_name, - class_name=class_name, - name=name, - checkpoint=checkpoint, - is_pipeline_module=is_pipeline_module, - cached_model_config_path=cached_model_config_path, - pipelines=pipelines, - torch_dtype=torch_dtype, - original_config=original_config, - local_files_only=local_files_only, - is_legacy_loading=is_legacy_loading, - disable_mmap=disable_mmap, - **kwargs, - ) - except SingleFileComponentError as e: - raise SingleFileComponentError( - ( - f"{e.message}\n" - f"Please load the component before passing it in as an argument to `from_single_file`.\n" - f"\n" - f"{name} = {class_name}.from_pretrained('...')\n" - f"pipe = {pipeline_class.__name__}.from_single_file(, {name}={name})\n" - f"\n" - ) - ) - - init_kwargs[name] = loaded_sub_model - - missing_modules = set(expected_modules) - set(init_kwargs.keys()) - passed_modules = list(passed_class_obj.keys()) - optional_modules = pipeline_class._optional_components - - if len(missing_modules) > 0 and missing_modules <= set(passed_modules + optional_modules): - for module in missing_modules: - init_kwargs[module] = passed_class_obj.get(module, None) - elif len(missing_modules) > 0: - passed_modules = set(list(init_kwargs.keys()) + list(passed_class_obj.keys())) - optional_kwargs - raise ValueError( - f"Pipeline {pipeline_class} expected {expected_modules}, but only {passed_modules} were passed." - ) - - # deprecated kwargs - load_safety_checker = kwargs.pop("load_safety_checker", None) - if load_safety_checker is not None: - deprecation_message = ( - "Please pass instances of `StableDiffusionSafetyChecker` and `AutoImageProcessor`" - "using the `safety_checker` and `feature_extractor` arguments in `from_single_file`" - ) - deprecate("load_safety_checker", "1.0.0", deprecation_message) - - safety_checker_components = _legacy_load_safety_checker(local_files_only, torch_dtype) - init_kwargs.update(safety_checker_components) - - pipe = pipeline_class(**init_kwargs) - - return pipe diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/single_file_model.py b/venv/lib/python3.10/site-packages/diffusers/loaders/single_file_model.py deleted file mode 100644 index 0c6f3cda666eba0be19525df26dae3ec9c08b140..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/single_file_model.py +++ /dev/null @@ -1,449 +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 -import re -from contextlib import nullcontext -from typing import Optional - -import torch -from huggingface_hub.utils import validate_hf_hub_args -from typing_extensions import Self - -from .. import __version__ -from ..quantizers import DiffusersAutoQuantizer -from ..utils import deprecate, is_accelerate_available, logging -from .single_file_utils import ( - SingleFileComponentError, - convert_animatediff_checkpoint_to_diffusers, - convert_auraflow_transformer_checkpoint_to_diffusers, - convert_autoencoder_dc_checkpoint_to_diffusers, - convert_chroma_transformer_checkpoint_to_diffusers, - convert_controlnet_checkpoint, - convert_flux_transformer_checkpoint_to_diffusers, - convert_hidream_transformer_to_diffusers, - convert_hunyuan_video_transformer_to_diffusers, - convert_ldm_unet_checkpoint, - convert_ldm_vae_checkpoint, - convert_ltx_transformer_checkpoint_to_diffusers, - convert_ltx_vae_checkpoint_to_diffusers, - convert_lumina2_to_diffusers, - convert_mochi_transformer_checkpoint_to_diffusers, - convert_sana_transformer_to_diffusers, - convert_sd3_transformer_checkpoint_to_diffusers, - convert_stable_cascade_unet_single_file_to_diffusers, - convert_wan_transformer_to_diffusers, - convert_wan_vae_to_diffusers, - create_controlnet_diffusers_config_from_ldm, - create_unet_diffusers_config_from_ldm, - create_vae_diffusers_config_from_ldm, - fetch_diffusers_config, - fetch_original_config, - load_single_file_checkpoint, -) - - -logger = logging.get_logger(__name__) - - -if is_accelerate_available(): - from accelerate import dispatch_model, init_empty_weights - - from ..models.modeling_utils import load_model_dict_into_meta - - -SINGLE_FILE_LOADABLE_CLASSES = { - "StableCascadeUNet": { - "checkpoint_mapping_fn": convert_stable_cascade_unet_single_file_to_diffusers, - }, - "UNet2DConditionModel": { - "checkpoint_mapping_fn": convert_ldm_unet_checkpoint, - "config_mapping_fn": create_unet_diffusers_config_from_ldm, - "default_subfolder": "unet", - "legacy_kwargs": { - "num_in_channels": "in_channels", # Legacy kwargs supported by `from_single_file` mapped to new args - }, - }, - "AutoencoderKL": { - "checkpoint_mapping_fn": convert_ldm_vae_checkpoint, - "config_mapping_fn": create_vae_diffusers_config_from_ldm, - "default_subfolder": "vae", - }, - "ControlNetModel": { - "checkpoint_mapping_fn": convert_controlnet_checkpoint, - "config_mapping_fn": create_controlnet_diffusers_config_from_ldm, - }, - "SD3Transformer2DModel": { - "checkpoint_mapping_fn": convert_sd3_transformer_checkpoint_to_diffusers, - "default_subfolder": "transformer", - }, - "MotionAdapter": { - "checkpoint_mapping_fn": convert_animatediff_checkpoint_to_diffusers, - }, - "SparseControlNetModel": { - "checkpoint_mapping_fn": convert_animatediff_checkpoint_to_diffusers, - }, - "FluxTransformer2DModel": { - "checkpoint_mapping_fn": convert_flux_transformer_checkpoint_to_diffusers, - "default_subfolder": "transformer", - }, - "ChromaTransformer2DModel": { - "checkpoint_mapping_fn": convert_chroma_transformer_checkpoint_to_diffusers, - "default_subfolder": "transformer", - }, - "LTXVideoTransformer3DModel": { - "checkpoint_mapping_fn": convert_ltx_transformer_checkpoint_to_diffusers, - "default_subfolder": "transformer", - }, - "AutoencoderKLLTXVideo": { - "checkpoint_mapping_fn": convert_ltx_vae_checkpoint_to_diffusers, - "default_subfolder": "vae", - }, - "AutoencoderDC": {"checkpoint_mapping_fn": convert_autoencoder_dc_checkpoint_to_diffusers}, - "MochiTransformer3DModel": { - "checkpoint_mapping_fn": convert_mochi_transformer_checkpoint_to_diffusers, - "default_subfolder": "transformer", - }, - "HunyuanVideoTransformer3DModel": { - "checkpoint_mapping_fn": convert_hunyuan_video_transformer_to_diffusers, - "default_subfolder": "transformer", - }, - "AuraFlowTransformer2DModel": { - "checkpoint_mapping_fn": convert_auraflow_transformer_checkpoint_to_diffusers, - "default_subfolder": "transformer", - }, - "Lumina2Transformer2DModel": { - "checkpoint_mapping_fn": convert_lumina2_to_diffusers, - "default_subfolder": "transformer", - }, - "SanaTransformer2DModel": { - "checkpoint_mapping_fn": convert_sana_transformer_to_diffusers, - "default_subfolder": "transformer", - }, - "WanTransformer3DModel": { - "checkpoint_mapping_fn": convert_wan_transformer_to_diffusers, - "default_subfolder": "transformer", - }, - "AutoencoderKLWan": { - "checkpoint_mapping_fn": convert_wan_vae_to_diffusers, - "default_subfolder": "vae", - }, - "HiDreamImageTransformer2DModel": { - "checkpoint_mapping_fn": convert_hidream_transformer_to_diffusers, - "default_subfolder": "transformer", - }, -} - - -def _get_single_file_loadable_mapping_class(cls): - diffusers_module = importlib.import_module(__name__.split(".")[0]) - for loadable_class_str in SINGLE_FILE_LOADABLE_CLASSES: - loadable_class = getattr(diffusers_module, loadable_class_str) - - if issubclass(cls, loadable_class): - return loadable_class_str - - return None - - -def _get_mapping_function_kwargs(mapping_fn, **kwargs): - parameters = inspect.signature(mapping_fn).parameters - - mapping_kwargs = {} - for parameter in parameters: - if parameter in kwargs: - mapping_kwargs[parameter] = kwargs[parameter] - - return mapping_kwargs - - -class FromOriginalModelMixin: - """ - Load pretrained weights saved in the `.ckpt` or `.safetensors` format into a model. - """ - - @classmethod - @validate_hf_hub_args - def from_single_file(cls, pretrained_model_link_or_path_or_dict: Optional[str] = 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. - - Parameters: - pretrained_model_link_or_path_or_dict (`str`, *optional*): - Can be either: - - A link to the `.safetensors` or `.ckpt` file (for example - `"https://huggingface.co//blob/main/.safetensors"`) on the Hub. - - A path to a local *file* containing the weights of the component model. - - A state dict containing the component model weights. - config (`str`, *optional*): - - A string, the *repo id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained pipeline hosted - on the Hub. - - A path to a *directory* (for example `./my_pipeline_directory/`) containing the pipeline component - configs in Diffusers format. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - original_config (`str`, *optional*): - Dict or path to a yaml file containing the configuration for the model in its original format. - If a dict is provided, it will be used to initialize the model configuration. - torch_dtype (`torch.dtype`, *optional*): - Override the default `torch.dtype` and load the model with another dtype. - 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*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - - 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`): - Whether to only load local model weights and configuration files or not. If set to True, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - 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. - kwargs (remaining dictionary of keyword arguments, *optional*): - Can be used to overwrite load and saveable variables (for example the pipeline components of the - specific pipeline class). The overwritten components are directly passed to the pipelines `__init__` - method. See example below for more information. - - ```py - >>> from diffusers import StableCascadeUNet - - >>> ckpt_path = "https://huggingface.co/stabilityai/stable-cascade/blob/main/stage_b_lite.safetensors" - >>> model = StableCascadeUNet.from_single_file(ckpt_path) - ``` - """ - - mapping_class_name = _get_single_file_loadable_mapping_class(cls) - # if class_name not in SINGLE_FILE_LOADABLE_CLASSES: - if mapping_class_name is None: - raise ValueError( - f"FromOriginalModelMixin is currently only compatible with {', '.join(SINGLE_FILE_LOADABLE_CLASSES.keys())}" - ) - - pretrained_model_link_or_path = kwargs.get("pretrained_model_link_or_path", None) - if pretrained_model_link_or_path is not None: - deprecation_message = ( - "Please use `pretrained_model_link_or_path_or_dict` argument instead for model classes" - ) - deprecate("pretrained_model_link_or_path", "1.0.0", deprecation_message) - pretrained_model_link_or_path_or_dict = pretrained_model_link_or_path - - config = kwargs.pop("config", None) - original_config = kwargs.pop("original_config", None) - - if config is not None and original_config is not None: - raise ValueError( - "`from_single_file` cannot accept both `config` and `original_config` arguments. Please provide only one of these arguments" - ) - - force_download = kwargs.pop("force_download", False) - proxies = kwargs.pop("proxies", None) - token = kwargs.pop("token", None) - cache_dir = kwargs.pop("cache_dir", None) - local_files_only = kwargs.pop("local_files_only", None) - subfolder = kwargs.pop("subfolder", None) - revision = kwargs.pop("revision", None) - config_revision = kwargs.pop("config_revision", None) - torch_dtype = kwargs.pop("torch_dtype", None) - quantization_config = kwargs.pop("quantization_config", None) - device = kwargs.pop("device", None) - disable_mmap = kwargs.pop("disable_mmap", False) - - user_agent = {"diffusers": __version__, "file_type": "single_file", "framework": "pytorch"} - # In order to ensure popular quantization methods are supported. Can be disable with `disable_telemetry` - if quantization_config is not None: - user_agent["quant"] = quantization_config.quant_method.value - - if torch_dtype is not None and not isinstance(torch_dtype, torch.dtype): - torch_dtype = torch.float32 - logger.warning( - f"Passed `torch_dtype` {torch_dtype} is not a `torch.dtype`. Defaulting to `torch.float32`." - ) - - if isinstance(pretrained_model_link_or_path_or_dict, dict): - checkpoint = pretrained_model_link_or_path_or_dict - else: - checkpoint = load_single_file_checkpoint( - pretrained_model_link_or_path_or_dict, - force_download=force_download, - proxies=proxies, - token=token, - cache_dir=cache_dir, - local_files_only=local_files_only, - revision=revision, - disable_mmap=disable_mmap, - user_agent=user_agent, - ) - if quantization_config is not None: - hf_quantizer = DiffusersAutoQuantizer.from_config(quantization_config) - hf_quantizer.validate_environment() - torch_dtype = hf_quantizer.update_torch_dtype(torch_dtype) - - else: - hf_quantizer = None - - mapping_functions = SINGLE_FILE_LOADABLE_CLASSES[mapping_class_name] - - checkpoint_mapping_fn = mapping_functions["checkpoint_mapping_fn"] - if original_config is not None: - if "config_mapping_fn" in mapping_functions: - config_mapping_fn = mapping_functions["config_mapping_fn"] - else: - config_mapping_fn = None - - if config_mapping_fn is None: - raise ValueError( - ( - f"`original_config` has been provided for {mapping_class_name} but no mapping function" - "was found to convert the original config to a Diffusers config in" - "`diffusers.loaders.single_file_utils`" - ) - ) - - if isinstance(original_config, str): - # If original_config is a URL or filepath fetch the original_config dict - original_config = fetch_original_config(original_config, local_files_only=local_files_only) - - config_mapping_kwargs = _get_mapping_function_kwargs(config_mapping_fn, **kwargs) - diffusers_model_config = config_mapping_fn( - original_config=original_config, checkpoint=checkpoint, **config_mapping_kwargs - ) - else: - if config is not None: - if isinstance(config, str): - default_pretrained_model_config_name = config - else: - raise ValueError( - ( - "Invalid `config` argument. Please provide a string representing a repo id" - "or path to a local Diffusers model repo." - ) - ) - - else: - config = fetch_diffusers_config(checkpoint) - default_pretrained_model_config_name = config["pretrained_model_name_or_path"] - - if "default_subfolder" in mapping_functions: - subfolder = mapping_functions["default_subfolder"] - - subfolder = subfolder or config.pop( - "subfolder", None - ) # some configs contain a subfolder key, e.g. StableCascadeUNet - - diffusers_model_config = cls.load_config( - pretrained_model_name_or_path=default_pretrained_model_config_name, - subfolder=subfolder, - local_files_only=local_files_only, - token=token, - revision=config_revision, - ) - expected_kwargs, optional_kwargs = cls._get_signature_keys(cls) - - # Map legacy kwargs to new kwargs - if "legacy_kwargs" in mapping_functions: - legacy_kwargs = mapping_functions["legacy_kwargs"] - for legacy_key, new_key in legacy_kwargs.items(): - if legacy_key in kwargs: - kwargs[new_key] = kwargs.pop(legacy_key) - - model_kwargs = {k: kwargs.get(k) for k in kwargs if k in expected_kwargs or k in optional_kwargs} - diffusers_model_config.update(model_kwargs) - - checkpoint_mapping_kwargs = _get_mapping_function_kwargs(checkpoint_mapping_fn, **kwargs) - diffusers_format_checkpoint = checkpoint_mapping_fn( - config=diffusers_model_config, checkpoint=checkpoint, **checkpoint_mapping_kwargs - ) - if not diffusers_format_checkpoint: - raise SingleFileComponentError( - f"Failed to load {mapping_class_name}. Weights for this component appear to be missing in the checkpoint." - ) - - ctx = init_empty_weights if is_accelerate_available() else nullcontext - with ctx(): - model = cls.from_config(diffusers_model_config) - - # Check if `_keep_in_fp32_modules` is not None - use_keep_in_fp32_modules = (cls._keep_in_fp32_modules is not None) and ( - (torch_dtype == torch.float16) or hasattr(hf_quantizer, "use_keep_in_fp32_modules") - ) - if use_keep_in_fp32_modules: - keep_in_fp32_modules = cls._keep_in_fp32_modules - if not isinstance(keep_in_fp32_modules, list): - keep_in_fp32_modules = [keep_in_fp32_modules] - - else: - keep_in_fp32_modules = [] - - if hf_quantizer is not None: - hf_quantizer.preprocess_model( - model=model, - device_map=None, - state_dict=diffusers_format_checkpoint, - keep_in_fp32_modules=keep_in_fp32_modules, - ) - - device_map = None - if is_accelerate_available(): - param_device = torch.device(device) if device else torch.device("cpu") - empty_state_dict = model.state_dict() - unexpected_keys = [ - param_name for param_name in diffusers_format_checkpoint if param_name not in empty_state_dict - ] - device_map = {"": param_device} - load_model_dict_into_meta( - model, - diffusers_format_checkpoint, - dtype=torch_dtype, - device_map=device_map, - hf_quantizer=hf_quantizer, - keep_in_fp32_modules=keep_in_fp32_modules, - unexpected_keys=unexpected_keys, - ) - else: - _, unexpected_keys = model.load_state_dict(diffusers_format_checkpoint, strict=False) - - if model._keys_to_ignore_on_load_unexpected is not None: - for pat in model._keys_to_ignore_on_load_unexpected: - unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None] - - if len(unexpected_keys) > 0: - logger.warning( - f"Some weights of the model checkpoint were not used when initializing {cls.__name__}: \n {[', '.join(unexpected_keys)]}" - ) - - if hf_quantizer is not None: - hf_quantizer.postprocess_model(model) - model.hf_quantizer = hf_quantizer - - if torch_dtype is not None and hf_quantizer is None: - model.to(torch_dtype) - - model.eval() - - if device_map is not None: - device_map_kwargs = {"device_map": device_map} - dispatch_model(model, **device_map_kwargs) - - return model diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/single_file_utils.py b/venv/lib/python3.10/site-packages/diffusers/loaders/single_file_utils.py deleted file mode 100644 index d8d183304e9a2ae46e81d8225bc7e4d766375a78..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/single_file_utils.py +++ /dev/null @@ -1,3481 +0,0 @@ -# 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. -"""Conversion script for the Stable Diffusion checkpoints.""" - -import copy -import os -import re -from contextlib import nullcontext -from io import BytesIO -from urllib.parse import urlparse - -import requests -import torch -import yaml - -from ..models.modeling_utils import load_state_dict -from ..schedulers import ( - DDIMScheduler, - DPMSolverMultistepScheduler, - EDMDPMSolverMultistepScheduler, - EulerAncestralDiscreteScheduler, - EulerDiscreteScheduler, - HeunDiscreteScheduler, - LMSDiscreteScheduler, - PNDMScheduler, -) -from ..utils import ( - SAFETENSORS_WEIGHTS_NAME, - WEIGHTS_NAME, - deprecate, - is_accelerate_available, - is_transformers_available, - logging, -) -from ..utils.constants import DIFFUSERS_REQUEST_TIMEOUT -from ..utils.hub_utils import _get_model_file - - -if is_transformers_available(): - from transformers import AutoImageProcessor - -if is_accelerate_available(): - from accelerate import init_empty_weights - - from ..models.modeling_utils import load_model_dict_into_meta - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -CHECKPOINT_KEY_NAMES = { - "v2": "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_k.weight", - "xl_base": "conditioner.embedders.1.model.transformer.resblocks.9.mlp.c_proj.bias", - "xl_refiner": "conditioner.embedders.0.model.transformer.resblocks.9.mlp.c_proj.bias", - "upscale": "model.diffusion_model.input_blocks.10.0.skip_connection.bias", - "controlnet": [ - "control_model.time_embed.0.weight", - "controlnet_cond_embedding.conv_in.weight", - ], - # TODO: find non-Diffusers keys for controlnet_xl - "controlnet_xl": "add_embedding.linear_1.weight", - "controlnet_xl_large": "down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_k.weight", - "controlnet_xl_mid": "down_blocks.1.attentions.0.norm.weight", - "playground-v2-5": "edm_mean", - "inpainting": "model.diffusion_model.input_blocks.0.0.weight", - "clip": "cond_stage_model.transformer.text_model.embeddings.position_embedding.weight", - "clip_sdxl": "conditioner.embedders.0.transformer.text_model.embeddings.position_embedding.weight", - "clip_sd3": "text_encoders.clip_l.transformer.text_model.embeddings.position_embedding.weight", - "open_clip": "cond_stage_model.model.token_embedding.weight", - "open_clip_sdxl": "conditioner.embedders.1.model.positional_embedding", - "open_clip_sdxl_refiner": "conditioner.embedders.0.model.text_projection", - "open_clip_sd3": "text_encoders.clip_g.transformer.text_model.embeddings.position_embedding.weight", - "stable_cascade_stage_b": "down_blocks.1.0.channelwise.0.weight", - "stable_cascade_stage_c": "clip_txt_mapper.weight", - "sd3": [ - "joint_blocks.0.context_block.adaLN_modulation.1.bias", - "model.diffusion_model.joint_blocks.0.context_block.adaLN_modulation.1.bias", - ], - "sd35_large": [ - "joint_blocks.37.x_block.mlp.fc1.weight", - "model.diffusion_model.joint_blocks.37.x_block.mlp.fc1.weight", - ], - "animatediff": "down_blocks.0.motion_modules.0.temporal_transformer.transformer_blocks.0.attention_blocks.0.pos_encoder.pe", - "animatediff_v2": "mid_block.motion_modules.0.temporal_transformer.norm.bias", - "animatediff_sdxl_beta": "up_blocks.2.motion_modules.0.temporal_transformer.norm.weight", - "animatediff_scribble": "controlnet_cond_embedding.conv_in.weight", - "animatediff_rgb": "controlnet_cond_embedding.weight", - "auraflow": [ - "double_layers.0.attn.w2q.weight", - "double_layers.0.attn.w1q.weight", - "cond_seq_linear.weight", - "t_embedder.mlp.0.weight", - ], - "flux": [ - "double_blocks.0.img_attn.norm.key_norm.scale", - "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale", - ], - "ltx-video": [ - "model.diffusion_model.patchify_proj.weight", - "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", - ], - "autoencoder-dc": "decoder.stages.1.op_list.0.main.conv.conv.bias", - "autoencoder-dc-sana": "encoder.project_in.conv.bias", - "mochi-1-preview": ["model.diffusion_model.blocks.0.attn.qkv_x.weight", "blocks.0.attn.qkv_x.weight"], - "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"], - "sana": [ - "blocks.0.cross_attn.q_linear.weight", - "blocks.0.cross_attn.q_linear.bias", - "blocks.0.cross_attn.kv_linear.weight", - "blocks.0.cross_attn.kv_linear.bias", - ], - "wan": ["model.diffusion_model.head.modulation", "head.modulation"], - "wan_vae": "decoder.middle.0.residual.0.gamma", - "hidream": "double_stream_blocks.0.block.adaLN_modulation.1.bias", -} - -DIFFUSERS_DEFAULT_PIPELINE_PATHS = { - "xl_base": {"pretrained_model_name_or_path": "stabilityai/stable-diffusion-xl-base-1.0"}, - "xl_refiner": {"pretrained_model_name_or_path": "stabilityai/stable-diffusion-xl-refiner-1.0"}, - "xl_inpaint": {"pretrained_model_name_or_path": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1"}, - "playground-v2-5": {"pretrained_model_name_or_path": "playgroundai/playground-v2.5-1024px-aesthetic"}, - "upscale": {"pretrained_model_name_or_path": "stabilityai/stable-diffusion-x4-upscaler"}, - "inpainting": {"pretrained_model_name_or_path": "stable-diffusion-v1-5/stable-diffusion-inpainting"}, - "inpainting_v2": {"pretrained_model_name_or_path": "stabilityai/stable-diffusion-2-inpainting"}, - "controlnet": {"pretrained_model_name_or_path": "lllyasviel/control_v11p_sd15_canny"}, - "controlnet_xl_large": {"pretrained_model_name_or_path": "diffusers/controlnet-canny-sdxl-1.0"}, - "controlnet_xl_mid": {"pretrained_model_name_or_path": "diffusers/controlnet-canny-sdxl-1.0-mid"}, - "controlnet_xl_small": {"pretrained_model_name_or_path": "diffusers/controlnet-canny-sdxl-1.0-small"}, - "v2": {"pretrained_model_name_or_path": "stabilityai/stable-diffusion-2-1"}, - "v1": {"pretrained_model_name_or_path": "stable-diffusion-v1-5/stable-diffusion-v1-5"}, - "stable_cascade_stage_b": {"pretrained_model_name_or_path": "stabilityai/stable-cascade", "subfolder": "decoder"}, - "stable_cascade_stage_b_lite": { - "pretrained_model_name_or_path": "stabilityai/stable-cascade", - "subfolder": "decoder_lite", - }, - "stable_cascade_stage_c": { - "pretrained_model_name_or_path": "stabilityai/stable-cascade-prior", - "subfolder": "prior", - }, - "stable_cascade_stage_c_lite": { - "pretrained_model_name_or_path": "stabilityai/stable-cascade-prior", - "subfolder": "prior_lite", - }, - "sd3": { - "pretrained_model_name_or_path": "stabilityai/stable-diffusion-3-medium-diffusers", - }, - "sd35_large": { - "pretrained_model_name_or_path": "stabilityai/stable-diffusion-3.5-large", - }, - "sd35_medium": { - "pretrained_model_name_or_path": "stabilityai/stable-diffusion-3.5-medium", - }, - "animatediff_v1": {"pretrained_model_name_or_path": "guoyww/animatediff-motion-adapter-v1-5"}, - "animatediff_v2": {"pretrained_model_name_or_path": "guoyww/animatediff-motion-adapter-v1-5-2"}, - "animatediff_v3": {"pretrained_model_name_or_path": "guoyww/animatediff-motion-adapter-v1-5-3"}, - "animatediff_sdxl_beta": {"pretrained_model_name_or_path": "guoyww/animatediff-motion-adapter-sdxl-beta"}, - "animatediff_scribble": {"pretrained_model_name_or_path": "guoyww/animatediff-sparsectrl-scribble"}, - "animatediff_rgb": {"pretrained_model_name_or_path": "guoyww/animatediff-sparsectrl-rgb"}, - "auraflow": {"pretrained_model_name_or_path": "fal/AuraFlow-v0.3"}, - "flux-dev": {"pretrained_model_name_or_path": "black-forest-labs/FLUX.1-dev"}, - "flux-fill": {"pretrained_model_name_or_path": "black-forest-labs/FLUX.1-Fill-dev"}, - "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"}, - "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"}, - "ltx-video-0.9.7": {"pretrained_model_name_or_path": "Lightricks/LTX-Video-0.9.7-dev"}, - "autoencoder-dc-f128c512": {"pretrained_model_name_or_path": "mit-han-lab/dc-ae-f128c512-mix-1.0-diffusers"}, - "autoencoder-dc-f64c128": {"pretrained_model_name_or_path": "mit-han-lab/dc-ae-f64c128-mix-1.0-diffusers"}, - "autoencoder-dc-f32c32": {"pretrained_model_name_or_path": "mit-han-lab/dc-ae-f32c32-mix-1.0-diffusers"}, - "autoencoder-dc-f32c32-sana": {"pretrained_model_name_or_path": "mit-han-lab/dc-ae-f32c32-sana-1.0-diffusers"}, - "mochi-1-preview": {"pretrained_model_name_or_path": "genmo/mochi-1-preview"}, - "hunyuan-video": {"pretrained_model_name_or_path": "hunyuanvideo-community/HunyuanVideo"}, - "instruct-pix2pix": {"pretrained_model_name_or_path": "timbrooks/instruct-pix2pix"}, - "lumina2": {"pretrained_model_name_or_path": "Alpha-VLLM/Lumina-Image-2.0"}, - "sana": {"pretrained_model_name_or_path": "Efficient-Large-Model/Sana_1600M_1024px_diffusers"}, - "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"}, - "hidream": {"pretrained_model_name_or_path": "HiDream-ai/HiDream-I1-Dev"}, -} - -# Use to configure model sample size when original config is provided -DIFFUSERS_TO_LDM_DEFAULT_IMAGE_SIZE_MAP = { - "xl_base": 1024, - "xl_refiner": 1024, - "xl_inpaint": 1024, - "playground-v2-5": 1024, - "upscale": 512, - "inpainting": 512, - "inpainting_v2": 512, - "controlnet": 512, - "instruct-pix2pix": 512, - "v2": 768, - "v1": 512, -} - - -DIFFUSERS_TO_LDM_MAPPING = { - "unet": { - "layers": { - "time_embedding.linear_1.weight": "time_embed.0.weight", - "time_embedding.linear_1.bias": "time_embed.0.bias", - "time_embedding.linear_2.weight": "time_embed.2.weight", - "time_embedding.linear_2.bias": "time_embed.2.bias", - "conv_in.weight": "input_blocks.0.0.weight", - "conv_in.bias": "input_blocks.0.0.bias", - "conv_norm_out.weight": "out.0.weight", - "conv_norm_out.bias": "out.0.bias", - "conv_out.weight": "out.2.weight", - "conv_out.bias": "out.2.bias", - }, - "class_embed_type": { - "class_embedding.linear_1.weight": "label_emb.0.0.weight", - "class_embedding.linear_1.bias": "label_emb.0.0.bias", - "class_embedding.linear_2.weight": "label_emb.0.2.weight", - "class_embedding.linear_2.bias": "label_emb.0.2.bias", - }, - "addition_embed_type": { - "add_embedding.linear_1.weight": "label_emb.0.0.weight", - "add_embedding.linear_1.bias": "label_emb.0.0.bias", - "add_embedding.linear_2.weight": "label_emb.0.2.weight", - "add_embedding.linear_2.bias": "label_emb.0.2.bias", - }, - }, - "controlnet": { - "layers": { - "time_embedding.linear_1.weight": "time_embed.0.weight", - "time_embedding.linear_1.bias": "time_embed.0.bias", - "time_embedding.linear_2.weight": "time_embed.2.weight", - "time_embedding.linear_2.bias": "time_embed.2.bias", - "conv_in.weight": "input_blocks.0.0.weight", - "conv_in.bias": "input_blocks.0.0.bias", - "controlnet_cond_embedding.conv_in.weight": "input_hint_block.0.weight", - "controlnet_cond_embedding.conv_in.bias": "input_hint_block.0.bias", - "controlnet_cond_embedding.conv_out.weight": "input_hint_block.14.weight", - "controlnet_cond_embedding.conv_out.bias": "input_hint_block.14.bias", - }, - "class_embed_type": { - "class_embedding.linear_1.weight": "label_emb.0.0.weight", - "class_embedding.linear_1.bias": "label_emb.0.0.bias", - "class_embedding.linear_2.weight": "label_emb.0.2.weight", - "class_embedding.linear_2.bias": "label_emb.0.2.bias", - }, - "addition_embed_type": { - "add_embedding.linear_1.weight": "label_emb.0.0.weight", - "add_embedding.linear_1.bias": "label_emb.0.0.bias", - "add_embedding.linear_2.weight": "label_emb.0.2.weight", - "add_embedding.linear_2.bias": "label_emb.0.2.bias", - }, - }, - "vae": { - "encoder.conv_in.weight": "encoder.conv_in.weight", - "encoder.conv_in.bias": "encoder.conv_in.bias", - "encoder.conv_out.weight": "encoder.conv_out.weight", - "encoder.conv_out.bias": "encoder.conv_out.bias", - "encoder.conv_norm_out.weight": "encoder.norm_out.weight", - "encoder.conv_norm_out.bias": "encoder.norm_out.bias", - "decoder.conv_in.weight": "decoder.conv_in.weight", - "decoder.conv_in.bias": "decoder.conv_in.bias", - "decoder.conv_out.weight": "decoder.conv_out.weight", - "decoder.conv_out.bias": "decoder.conv_out.bias", - "decoder.conv_norm_out.weight": "decoder.norm_out.weight", - "decoder.conv_norm_out.bias": "decoder.norm_out.bias", - "quant_conv.weight": "quant_conv.weight", - "quant_conv.bias": "quant_conv.bias", - "post_quant_conv.weight": "post_quant_conv.weight", - "post_quant_conv.bias": "post_quant_conv.bias", - }, - "openclip": { - "layers": { - "text_model.embeddings.position_embedding.weight": "positional_embedding", - "text_model.embeddings.token_embedding.weight": "token_embedding.weight", - "text_model.final_layer_norm.weight": "ln_final.weight", - "text_model.final_layer_norm.bias": "ln_final.bias", - "text_projection.weight": "text_projection", - }, - "transformer": { - "text_model.encoder.layers.": "resblocks.", - "layer_norm1": "ln_1", - "layer_norm2": "ln_2", - ".fc1.": ".c_fc.", - ".fc2.": ".c_proj.", - ".self_attn": ".attn", - "transformer.text_model.final_layer_norm.": "ln_final.", - "transformer.text_model.embeddings.token_embedding.weight": "token_embedding.weight", - "transformer.text_model.embeddings.position_embedding.weight": "positional_embedding", - }, - }, -} - -SD_2_TEXT_ENCODER_KEYS_TO_IGNORE = [ - "cond_stage_model.model.transformer.resblocks.23.attn.in_proj_bias", - "cond_stage_model.model.transformer.resblocks.23.attn.in_proj_weight", - "cond_stage_model.model.transformer.resblocks.23.attn.out_proj.bias", - "cond_stage_model.model.transformer.resblocks.23.attn.out_proj.weight", - "cond_stage_model.model.transformer.resblocks.23.ln_1.bias", - "cond_stage_model.model.transformer.resblocks.23.ln_1.weight", - "cond_stage_model.model.transformer.resblocks.23.ln_2.bias", - "cond_stage_model.model.transformer.resblocks.23.ln_2.weight", - "cond_stage_model.model.transformer.resblocks.23.mlp.c_fc.bias", - "cond_stage_model.model.transformer.resblocks.23.mlp.c_fc.weight", - "cond_stage_model.model.transformer.resblocks.23.mlp.c_proj.bias", - "cond_stage_model.model.transformer.resblocks.23.mlp.c_proj.weight", - "cond_stage_model.model.text_projection", -] - -# To support legacy scheduler_type argument -SCHEDULER_DEFAULT_CONFIG = { - "beta_schedule": "scaled_linear", - "beta_start": 0.00085, - "beta_end": 0.012, - "interpolation_type": "linear", - "num_train_timesteps": 1000, - "prediction_type": "epsilon", - "sample_max_value": 1.0, - "set_alpha_to_one": False, - "skip_prk_steps": True, - "steps_offset": 1, - "timestep_spacing": "leading", -} - -LDM_VAE_KEYS = ["first_stage_model.", "vae."] -LDM_VAE_DEFAULT_SCALING_FACTOR = 0.18215 -PLAYGROUND_VAE_SCALING_FACTOR = 0.5 -LDM_UNET_KEY = "model.diffusion_model." -LDM_CONTROLNET_KEY = "control_model." -LDM_CLIP_PREFIX_TO_REMOVE = [ - "cond_stage_model.transformer.", - "conditioner.embedders.0.transformer.", -] -LDM_OPEN_CLIP_TEXT_PROJECTION_DIM = 1024 -SCHEDULER_LEGACY_KWARGS = ["prediction_type", "scheduler_type"] - -VALID_URL_PREFIXES = ["https://huggingface.co/", "huggingface.co/", "hf.co/", "https://hf.co/"] - - -class SingleFileComponentError(Exception): - def __init__(self, message=None): - self.message = message - super().__init__(self.message) - - -def is_valid_url(url): - result = urlparse(url) - if result.scheme and result.netloc: - return True - - return False - - -def _extract_repo_id_and_weights_name(pretrained_model_name_or_path): - if not is_valid_url(pretrained_model_name_or_path): - raise ValueError("Invalid `pretrained_model_name_or_path` provided. Please set it to a valid URL.") - - pattern = r"([^/]+)/([^/]+)/(?:blob/main/)?(.+)" - weights_name = None - repo_id = (None,) - for prefix in VALID_URL_PREFIXES: - pretrained_model_name_or_path = pretrained_model_name_or_path.replace(prefix, "") - match = re.match(pattern, pretrained_model_name_or_path) - if not match: - logger.warning("Unable to identify the repo_id and weights_name from the provided URL.") - return repo_id, weights_name - - repo_id = f"{match.group(1)}/{match.group(2)}" - weights_name = match.group(3) - - return repo_id, weights_name - - -def _is_model_weights_in_cached_folder(cached_folder, name): - pretrained_model_name_or_path = os.path.join(cached_folder, name) - weights_exist = False - - for weights_name in [WEIGHTS_NAME, SAFETENSORS_WEIGHTS_NAME]: - if os.path.isfile(os.path.join(pretrained_model_name_or_path, weights_name)): - weights_exist = True - - return weights_exist - - -def _is_legacy_scheduler_kwargs(kwargs): - return any(k in SCHEDULER_LEGACY_KWARGS for k in kwargs.keys()) - - -def load_single_file_checkpoint( - pretrained_model_link_or_path, - force_download=False, - proxies=None, - token=None, - cache_dir=None, - local_files_only=None, - revision=None, - disable_mmap=False, - user_agent=None, -): - if user_agent is None: - user_agent = {"file_type": "single_file", "framework": "pytorch"} - - if os.path.isfile(pretrained_model_link_or_path): - pretrained_model_link_or_path = pretrained_model_link_or_path - - else: - repo_id, weights_name = _extract_repo_id_and_weights_name(pretrained_model_link_or_path) - pretrained_model_link_or_path = _get_model_file( - repo_id, - weights_name=weights_name, - force_download=force_download, - cache_dir=cache_dir, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - user_agent=user_agent, - ) - - checkpoint = load_state_dict(pretrained_model_link_or_path, disable_mmap=disable_mmap) - - # some checkpoints contain the model state dict under a "state_dict" key - while "state_dict" in checkpoint: - checkpoint = checkpoint["state_dict"] - - return checkpoint - - -def fetch_original_config(original_config_file, local_files_only=False): - if os.path.isfile(original_config_file): - with open(original_config_file, "r") as fp: - original_config_file = fp.read() - - elif is_valid_url(original_config_file): - if local_files_only: - raise ValueError( - "`local_files_only` is set to True, but a URL was provided as `original_config_file`. " - "Please provide a valid local file path." - ) - - original_config_file = BytesIO(requests.get(original_config_file, timeout=DIFFUSERS_REQUEST_TIMEOUT).content) - - else: - raise ValueError("Invalid `original_config_file` provided. Please set it to a valid file path or URL.") - - original_config = yaml.safe_load(original_config_file) - - return original_config - - -def is_clip_model(checkpoint): - if CHECKPOINT_KEY_NAMES["clip"] in checkpoint: - return True - - return False - - -def is_clip_sdxl_model(checkpoint): - if CHECKPOINT_KEY_NAMES["clip_sdxl"] in checkpoint: - return True - - return False - - -def is_clip_sd3_model(checkpoint): - if CHECKPOINT_KEY_NAMES["clip_sd3"] in checkpoint: - return True - - return False - - -def is_open_clip_model(checkpoint): - if CHECKPOINT_KEY_NAMES["open_clip"] in checkpoint: - return True - - return False - - -def is_open_clip_sdxl_model(checkpoint): - if CHECKPOINT_KEY_NAMES["open_clip_sdxl"] in checkpoint: - return True - - return False - - -def is_open_clip_sd3_model(checkpoint): - if CHECKPOINT_KEY_NAMES["open_clip_sd3"] in checkpoint: - return True - - return False - - -def is_open_clip_sdxl_refiner_model(checkpoint): - if CHECKPOINT_KEY_NAMES["open_clip_sdxl_refiner"] in checkpoint: - return True - - return False - - -def is_clip_model_in_single_file(class_obj, checkpoint): - is_clip_in_checkpoint = any( - [ - is_clip_model(checkpoint), - is_clip_sd3_model(checkpoint), - is_open_clip_model(checkpoint), - is_open_clip_sdxl_model(checkpoint), - is_open_clip_sdxl_refiner_model(checkpoint), - is_open_clip_sd3_model(checkpoint), - ] - ) - if ( - class_obj.__name__ == "CLIPTextModel" or class_obj.__name__ == "CLIPTextModelWithProjection" - ) and is_clip_in_checkpoint: - return True - - return False - - -def infer_diffusers_model_type(checkpoint): - if ( - CHECKPOINT_KEY_NAMES["inpainting"] in checkpoint - and checkpoint[CHECKPOINT_KEY_NAMES["inpainting"]].shape[1] == 9 - ): - if CHECKPOINT_KEY_NAMES["v2"] in checkpoint and checkpoint[CHECKPOINT_KEY_NAMES["v2"]].shape[-1] == 1024: - model_type = "inpainting_v2" - elif CHECKPOINT_KEY_NAMES["xl_base"] in checkpoint: - model_type = "xl_inpaint" - else: - model_type = "inpainting" - - elif CHECKPOINT_KEY_NAMES["v2"] in checkpoint and checkpoint[CHECKPOINT_KEY_NAMES["v2"]].shape[-1] == 1024: - model_type = "v2" - - elif CHECKPOINT_KEY_NAMES["playground-v2-5"] in checkpoint: - model_type = "playground-v2-5" - - elif CHECKPOINT_KEY_NAMES["xl_base"] in checkpoint: - model_type = "xl_base" - - elif CHECKPOINT_KEY_NAMES["xl_refiner"] in checkpoint: - model_type = "xl_refiner" - - elif CHECKPOINT_KEY_NAMES["upscale"] in checkpoint: - model_type = "upscale" - - elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["controlnet"]): - if CHECKPOINT_KEY_NAMES["controlnet_xl"] in checkpoint: - if CHECKPOINT_KEY_NAMES["controlnet_xl_large"] in checkpoint: - model_type = "controlnet_xl_large" - elif CHECKPOINT_KEY_NAMES["controlnet_xl_mid"] in checkpoint: - model_type = "controlnet_xl_mid" - else: - model_type = "controlnet_xl_small" - else: - model_type = "controlnet" - - elif ( - CHECKPOINT_KEY_NAMES["stable_cascade_stage_c"] in checkpoint - and checkpoint[CHECKPOINT_KEY_NAMES["stable_cascade_stage_c"]].shape[0] == 1536 - ): - model_type = "stable_cascade_stage_c_lite" - - elif ( - CHECKPOINT_KEY_NAMES["stable_cascade_stage_c"] in checkpoint - and checkpoint[CHECKPOINT_KEY_NAMES["stable_cascade_stage_c"]].shape[0] == 2048 - ): - model_type = "stable_cascade_stage_c" - - elif ( - CHECKPOINT_KEY_NAMES["stable_cascade_stage_b"] in checkpoint - and checkpoint[CHECKPOINT_KEY_NAMES["stable_cascade_stage_b"]].shape[-1] == 576 - ): - model_type = "stable_cascade_stage_b_lite" - - elif ( - CHECKPOINT_KEY_NAMES["stable_cascade_stage_b"] in checkpoint - and checkpoint[CHECKPOINT_KEY_NAMES["stable_cascade_stage_b"]].shape[-1] == 640 - ): - model_type = "stable_cascade_stage_b" - - elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["sd3"]) and any( - checkpoint[key].shape[-1] == 9216 if key in checkpoint else False for key in CHECKPOINT_KEY_NAMES["sd3"] - ): - if "model.diffusion_model.pos_embed" in checkpoint: - key = "model.diffusion_model.pos_embed" - else: - key = "pos_embed" - - if checkpoint[key].shape[1] == 36864: - model_type = "sd3" - elif checkpoint[key].shape[1] == 147456: - model_type = "sd35_medium" - - elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["sd35_large"]): - model_type = "sd35_large" - - elif CHECKPOINT_KEY_NAMES["animatediff"] in checkpoint: - if CHECKPOINT_KEY_NAMES["animatediff_scribble"] in checkpoint: - model_type = "animatediff_scribble" - - elif CHECKPOINT_KEY_NAMES["animatediff_rgb"] in checkpoint: - model_type = "animatediff_rgb" - - elif CHECKPOINT_KEY_NAMES["animatediff_v2"] in checkpoint: - model_type = "animatediff_v2" - - elif checkpoint[CHECKPOINT_KEY_NAMES["animatediff_sdxl_beta"]].shape[-1] == 320: - model_type = "animatediff_sdxl_beta" - - elif checkpoint[CHECKPOINT_KEY_NAMES["animatediff"]].shape[1] == 24: - model_type = "animatediff_v1" - - else: - model_type = "animatediff_v3" - - elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["flux"]): - if 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: - key = "model.diffusion_model.img_in.weight" - else: - key = "img_in.weight" - - if checkpoint[key].shape[1] == 384: - model_type = "flux-fill" - elif checkpoint[key].shape[1] == 128: - model_type = "flux-depth" - else: - model_type = "flux-dev" - else: - model_type = "flux-schnell" - - elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["ltx-video"]): - has_vae = "vae.encoder.conv_in.conv.bias" in checkpoint - if any(key.endswith("transformer_blocks.47.scale_shift_table") for key in checkpoint): - model_type = "ltx-video-0.9.7" - elif has_vae and checkpoint["vae.encoder.conv_out.conv.weight"].shape[1] == 2048: - model_type = "ltx-video-0.9.5" - elif "vae.decoder.last_time_embedder.timestep_embedder.linear_1.weight" in checkpoint: - model_type = "ltx-video-0.9.1" - else: - model_type = "ltx-video" - - elif CHECKPOINT_KEY_NAMES["autoencoder-dc"] in checkpoint: - encoder_key = "encoder.project_in.conv.conv.bias" - decoder_key = "decoder.project_in.main.conv.weight" - - if CHECKPOINT_KEY_NAMES["autoencoder-dc-sana"] in checkpoint: - model_type = "autoencoder-dc-f32c32-sana" - - elif checkpoint[encoder_key].shape[-1] == 64 and checkpoint[decoder_key].shape[1] == 32: - model_type = "autoencoder-dc-f32c32" - - elif checkpoint[encoder_key].shape[-1] == 64 and checkpoint[decoder_key].shape[1] == 128: - model_type = "autoencoder-dc-f64c128" - - else: - model_type = "autoencoder-dc-f128c512" - - elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["mochi-1-preview"]): - model_type = "mochi-1-preview" - - elif CHECKPOINT_KEY_NAMES["hunyuan-video"] in checkpoint: - model_type = "hunyuan-video" - - elif all(key in checkpoint for key in CHECKPOINT_KEY_NAMES["auraflow"]): - model_type = "auraflow" - - elif ( - CHECKPOINT_KEY_NAMES["instruct-pix2pix"] in checkpoint - and checkpoint[CHECKPOINT_KEY_NAMES["instruct-pix2pix"]].shape[1] == 8 - ): - model_type = "instruct-pix2pix" - - elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["lumina2"]): - model_type = "lumina2" - - elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["sana"]): - model_type = "sana" - - elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["wan"]): - if "model.diffusion_model.patch_embedding.weight" in checkpoint: - target_key = "model.diffusion_model.patch_embedding.weight" - else: - target_key = "patch_embedding.weight" - - if 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: - model_type = "wan-t2v-14B" - else: - model_type = "wan-i2v-14B" - elif CHECKPOINT_KEY_NAMES["wan_vae"] in checkpoint: - # All Wan models use the same VAE so we can use the same default model repo to fetch the config - model_type = "wan-t2v-14B" - elif CHECKPOINT_KEY_NAMES["hidream"] in checkpoint: - model_type = "hidream" - else: - model_type = "v1" - - return model_type - - -def fetch_diffusers_config(checkpoint): - model_type = infer_diffusers_model_type(checkpoint) - model_path = DIFFUSERS_DEFAULT_PIPELINE_PATHS[model_type] - model_path = copy.deepcopy(model_path) - - return model_path - - -def set_image_size(checkpoint, image_size=None): - if image_size: - return image_size - - model_type = infer_diffusers_model_type(checkpoint) - image_size = DIFFUSERS_TO_LDM_DEFAULT_IMAGE_SIZE_MAP[model_type] - - return image_size - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.conv_attn_to_linear -def conv_attn_to_linear(checkpoint): - keys = list(checkpoint.keys()) - attn_keys = ["query.weight", "key.weight", "value.weight"] - for key in keys: - if ".".join(key.split(".")[-2:]) in attn_keys: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0, 0] - elif "proj_attn.weight" in key: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0] - - -def create_unet_diffusers_config_from_ldm( - original_config, checkpoint, image_size=None, upcast_attention=None, num_in_channels=None -): - """ - Creates a config for the diffusers based on the config of the LDM model. - """ - if image_size is not None: - deprecation_message = ( - "Configuring UNet2DConditionModel with the `image_size` argument to `from_single_file`" - "is deprecated and will be ignored in future versions." - ) - deprecate("image_size", "1.0.0", deprecation_message) - - image_size = set_image_size(checkpoint, image_size=image_size) - - if ( - "unet_config" in original_config["model"]["params"] - and original_config["model"]["params"]["unet_config"] is not None - ): - unet_params = original_config["model"]["params"]["unet_config"]["params"] - else: - unet_params = original_config["model"]["params"]["network_config"]["params"] - - if num_in_channels is not None: - deprecation_message = ( - "Configuring UNet2DConditionModel with the `num_in_channels` argument to `from_single_file`" - "is deprecated and will be ignored in future versions." - ) - deprecate("image_size", "1.0.0", deprecation_message) - in_channels = num_in_channels - else: - in_channels = unet_params["in_channels"] - - vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] - block_out_channels = [unet_params["model_channels"] * mult for mult in unet_params["channel_mult"]] - - down_block_types = [] - resolution = 1 - for i in range(len(block_out_channels)): - block_type = "CrossAttnDownBlock2D" if resolution in unet_params["attention_resolutions"] else "DownBlock2D" - down_block_types.append(block_type) - if i != len(block_out_channels) - 1: - resolution *= 2 - - up_block_types = [] - for i in range(len(block_out_channels)): - block_type = "CrossAttnUpBlock2D" if resolution in unet_params["attention_resolutions"] else "UpBlock2D" - up_block_types.append(block_type) - resolution //= 2 - - if unet_params["transformer_depth"] is not None: - transformer_layers_per_block = ( - unet_params["transformer_depth"] - if isinstance(unet_params["transformer_depth"], int) - else list(unet_params["transformer_depth"]) - ) - else: - transformer_layers_per_block = 1 - - vae_scale_factor = 2 ** (len(vae_params["ch_mult"]) - 1) - - head_dim = unet_params["num_heads"] if "num_heads" in unet_params else None - use_linear_projection = ( - unet_params["use_linear_in_transformer"] if "use_linear_in_transformer" in unet_params else False - ) - if use_linear_projection: - # stable diffusion 2-base-512 and 2-768 - if head_dim is None: - head_dim_mult = unet_params["model_channels"] // unet_params["num_head_channels"] - head_dim = [head_dim_mult * c for c in list(unet_params["channel_mult"])] - - class_embed_type = None - addition_embed_type = None - addition_time_embed_dim = None - projection_class_embeddings_input_dim = None - context_dim = None - - if unet_params["context_dim"] is not None: - context_dim = ( - unet_params["context_dim"] - if isinstance(unet_params["context_dim"], int) - else unet_params["context_dim"][0] - ) - - if "num_classes" in unet_params: - if unet_params["num_classes"] == "sequential": - if context_dim in [2048, 1280]: - # SDXL - addition_embed_type = "text_time" - addition_time_embed_dim = 256 - else: - class_embed_type = "projection" - assert "adm_in_channels" in unet_params - projection_class_embeddings_input_dim = unet_params["adm_in_channels"] - - config = { - "sample_size": image_size // vae_scale_factor, - "in_channels": in_channels, - "down_block_types": down_block_types, - "block_out_channels": block_out_channels, - "layers_per_block": unet_params["num_res_blocks"], - "cross_attention_dim": context_dim, - "attention_head_dim": head_dim, - "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, - "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim, - "transformer_layers_per_block": transformer_layers_per_block, - } - - if upcast_attention is not None: - deprecation_message = ( - "Configuring UNet2DConditionModel with the `upcast_attention` argument to `from_single_file`" - "is deprecated and will be ignored in future versions." - ) - deprecate("image_size", "1.0.0", deprecation_message) - config["upcast_attention"] = upcast_attention - - if "disable_self_attentions" in unet_params: - config["only_cross_attention"] = unet_params["disable_self_attentions"] - - if "num_classes" in unet_params and isinstance(unet_params["num_classes"], int): - config["num_class_embeds"] = unet_params["num_classes"] - - config["out_channels"] = unet_params["out_channels"] - config["up_block_types"] = up_block_types - - return config - - -def create_controlnet_diffusers_config_from_ldm(original_config, checkpoint, image_size=None, **kwargs): - if image_size is not None: - deprecation_message = ( - "Configuring ControlNetModel with the `image_size` argument" - "is deprecated and will be ignored in future versions." - ) - deprecate("image_size", "1.0.0", deprecation_message) - - image_size = set_image_size(checkpoint, image_size=image_size) - - unet_params = original_config["model"]["params"]["control_stage_config"]["params"] - diffusers_unet_config = create_unet_diffusers_config_from_ldm(original_config, image_size=image_size) - - controlnet_config = { - "conditioning_channels": unet_params["hint_channels"], - "in_channels": diffusers_unet_config["in_channels"], - "down_block_types": diffusers_unet_config["down_block_types"], - "block_out_channels": diffusers_unet_config["block_out_channels"], - "layers_per_block": diffusers_unet_config["layers_per_block"], - "cross_attention_dim": diffusers_unet_config["cross_attention_dim"], - "attention_head_dim": diffusers_unet_config["attention_head_dim"], - "use_linear_projection": diffusers_unet_config["use_linear_projection"], - "class_embed_type": diffusers_unet_config["class_embed_type"], - "addition_embed_type": diffusers_unet_config["addition_embed_type"], - "addition_time_embed_dim": diffusers_unet_config["addition_time_embed_dim"], - "projection_class_embeddings_input_dim": diffusers_unet_config["projection_class_embeddings_input_dim"], - "transformer_layers_per_block": diffusers_unet_config["transformer_layers_per_block"], - } - - return controlnet_config - - -def create_vae_diffusers_config_from_ldm(original_config, checkpoint, image_size=None, scaling_factor=None): - """ - Creates a config for the diffusers based on the config of the LDM model. - """ - if image_size is not None: - deprecation_message = ( - "Configuring AutoencoderKL with the `image_size` argument" - "is deprecated and will be ignored in future versions." - ) - deprecate("image_size", "1.0.0", deprecation_message) - - image_size = set_image_size(checkpoint, image_size=image_size) - - if "edm_mean" in checkpoint and "edm_std" in checkpoint: - latents_mean = checkpoint["edm_mean"] - latents_std = checkpoint["edm_std"] - else: - latents_mean = None - latents_std = None - - vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] - if (scaling_factor is None) and (latents_mean is not None) and (latents_std is not None): - scaling_factor = PLAYGROUND_VAE_SCALING_FACTOR - - elif (scaling_factor is None) and ("scale_factor" in original_config["model"]["params"]): - scaling_factor = original_config["model"]["params"]["scale_factor"] - - elif scaling_factor is None: - scaling_factor = LDM_VAE_DEFAULT_SCALING_FACTOR - - block_out_channels = [vae_params["ch"] * mult for mult in vae_params["ch_mult"]] - down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) - up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) - - config = { - "sample_size": image_size, - "in_channels": vae_params["in_channels"], - "out_channels": vae_params["out_ch"], - "down_block_types": down_block_types, - "up_block_types": up_block_types, - "block_out_channels": block_out_channels, - "latent_channels": vae_params["z_channels"], - "layers_per_block": vae_params["num_res_blocks"], - "scaling_factor": scaling_factor, - } - if latents_mean is not None and latents_std is not None: - config.update({"latents_mean": latents_mean, "latents_std": latents_std}) - - return config - - -def update_unet_resnet_ldm_to_diffusers(ldm_keys, new_checkpoint, checkpoint, mapping=None): - for ldm_key in ldm_keys: - diffusers_key = ( - ldm_key.replace("in_layers.0", "norm1") - .replace("in_layers.2", "conv1") - .replace("out_layers.0", "norm2") - .replace("out_layers.3", "conv2") - .replace("emb_layers.1", "time_emb_proj") - .replace("skip_connection", "conv_shortcut") - ) - if mapping: - diffusers_key = diffusers_key.replace(mapping["old"], mapping["new"]) - new_checkpoint[diffusers_key] = checkpoint.get(ldm_key) - - -def update_unet_attention_ldm_to_diffusers(ldm_keys, new_checkpoint, checkpoint, mapping): - for ldm_key in ldm_keys: - diffusers_key = ldm_key.replace(mapping["old"], mapping["new"]) - new_checkpoint[diffusers_key] = checkpoint.get(ldm_key) - - -def update_vae_resnet_ldm_to_diffusers(keys, new_checkpoint, checkpoint, mapping): - for ldm_key in keys: - diffusers_key = ldm_key.replace(mapping["old"], mapping["new"]).replace("nin_shortcut", "conv_shortcut") - new_checkpoint[diffusers_key] = checkpoint.get(ldm_key) - - -def update_vae_attentions_ldm_to_diffusers(keys, new_checkpoint, checkpoint, mapping): - for ldm_key in keys: - diffusers_key = ( - ldm_key.replace(mapping["old"], mapping["new"]) - .replace("norm.weight", "group_norm.weight") - .replace("norm.bias", "group_norm.bias") - .replace("q.weight", "to_q.weight") - .replace("q.bias", "to_q.bias") - .replace("k.weight", "to_k.weight") - .replace("k.bias", "to_k.bias") - .replace("v.weight", "to_v.weight") - .replace("v.bias", "to_v.bias") - .replace("proj_out.weight", "to_out.0.weight") - .replace("proj_out.bias", "to_out.0.bias") - ) - new_checkpoint[diffusers_key] = checkpoint.get(ldm_key) - - # proj_attn.weight has to be converted from conv 1D to linear - shape = new_checkpoint[diffusers_key].shape - - if len(shape) == 3: - new_checkpoint[diffusers_key] = new_checkpoint[diffusers_key][:, :, 0] - elif len(shape) == 4: - new_checkpoint[diffusers_key] = new_checkpoint[diffusers_key][:, :, 0, 0] - - -def convert_stable_cascade_unet_single_file_to_diffusers(checkpoint, **kwargs): - is_stage_c = "clip_txt_mapper.weight" in checkpoint - - if is_stage_c: - state_dict = {} - for key in checkpoint.keys(): - if key.endswith("in_proj_weight"): - weights = checkpoint[key].chunk(3, 0) - state_dict[key.replace("attn.in_proj_weight", "to_q.weight")] = weights[0] - state_dict[key.replace("attn.in_proj_weight", "to_k.weight")] = weights[1] - state_dict[key.replace("attn.in_proj_weight", "to_v.weight")] = weights[2] - elif key.endswith("in_proj_bias"): - weights = checkpoint[key].chunk(3, 0) - state_dict[key.replace("attn.in_proj_bias", "to_q.bias")] = weights[0] - state_dict[key.replace("attn.in_proj_bias", "to_k.bias")] = weights[1] - state_dict[key.replace("attn.in_proj_bias", "to_v.bias")] = weights[2] - elif key.endswith("out_proj.weight"): - weights = checkpoint[key] - state_dict[key.replace("attn.out_proj.weight", "to_out.0.weight")] = weights - elif key.endswith("out_proj.bias"): - weights = checkpoint[key] - state_dict[key.replace("attn.out_proj.bias", "to_out.0.bias")] = weights - else: - state_dict[key] = checkpoint[key] - else: - state_dict = {} - for key in checkpoint.keys(): - if key.endswith("in_proj_weight"): - weights = checkpoint[key].chunk(3, 0) - state_dict[key.replace("attn.in_proj_weight", "to_q.weight")] = weights[0] - state_dict[key.replace("attn.in_proj_weight", "to_k.weight")] = weights[1] - state_dict[key.replace("attn.in_proj_weight", "to_v.weight")] = weights[2] - elif key.endswith("in_proj_bias"): - weights = checkpoint[key].chunk(3, 0) - state_dict[key.replace("attn.in_proj_bias", "to_q.bias")] = weights[0] - state_dict[key.replace("attn.in_proj_bias", "to_k.bias")] = weights[1] - state_dict[key.replace("attn.in_proj_bias", "to_v.bias")] = weights[2] - elif key.endswith("out_proj.weight"): - weights = checkpoint[key] - state_dict[key.replace("attn.out_proj.weight", "to_out.0.weight")] = weights - elif key.endswith("out_proj.bias"): - weights = checkpoint[key] - state_dict[key.replace("attn.out_proj.bias", "to_out.0.bias")] = weights - # rename clip_mapper to clip_txt_pooled_mapper - elif key.endswith("clip_mapper.weight"): - weights = checkpoint[key] - state_dict[key.replace("clip_mapper.weight", "clip_txt_pooled_mapper.weight")] = weights - elif key.endswith("clip_mapper.bias"): - weights = checkpoint[key] - state_dict[key.replace("clip_mapper.bias", "clip_txt_pooled_mapper.bias")] = weights - else: - state_dict[key] = checkpoint[key] - - return state_dict - - -def convert_ldm_unet_checkpoint(checkpoint, config, extract_ema=False, **kwargs): - """ - Takes a state dict and a config, and returns a converted checkpoint. - """ - # extract state_dict for UNet - unet_state_dict = {} - keys = list(checkpoint.keys()) - unet_key = LDM_UNET_KEY - - # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA - if sum(k.startswith("model_ema") for k in keys) > 100 and extract_ema: - logger.warning("Checkpoint has both EMA and non-EMA weights.") - logger.warning( - "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA" - " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag." - ) - for key in keys: - if key.startswith("model.diffusion_model"): - flat_ema_key = "model_ema." + "".join(key.split(".")[1:]) - unet_state_dict[key.replace(unet_key, "")] = checkpoint.get(flat_ema_key) - else: - if sum(k.startswith("model_ema") for k in keys) > 100: - logger.warning( - "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA" - " weights (usually better for inference), please make sure to add the `--extract_ema` flag." - ) - for key in keys: - if key.startswith(unet_key): - unet_state_dict[key.replace(unet_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - ldm_unet_keys = DIFFUSERS_TO_LDM_MAPPING["unet"]["layers"] - for diffusers_key, ldm_key in ldm_unet_keys.items(): - if ldm_key not in unet_state_dict: - continue - new_checkpoint[diffusers_key] = unet_state_dict[ldm_key] - - if ("class_embed_type" in config) and (config["class_embed_type"] in ["timestep", "projection"]): - class_embed_keys = DIFFUSERS_TO_LDM_MAPPING["unet"]["class_embed_type"] - for diffusers_key, ldm_key in class_embed_keys.items(): - new_checkpoint[diffusers_key] = unet_state_dict[ldm_key] - - if ("addition_embed_type" in config) and (config["addition_embed_type"] == "text_time"): - addition_embed_keys = DIFFUSERS_TO_LDM_MAPPING["unet"]["addition_embed_type"] - for diffusers_key, ldm_key in addition_embed_keys.items(): - new_checkpoint[diffusers_key] = unet_state_dict[ldm_key] - - # Relevant to StableDiffusionUpscalePipeline - if "num_class_embeds" in config: - if (config["num_class_embeds"] is not None) and ("label_emb.weight" in unet_state_dict): - new_checkpoint["class_embedding.weight"] = unet_state_dict["label_emb.weight"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key] - for layer_id in range(num_output_blocks) - } - - # Down blocks - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - update_unet_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - unet_state_dict, - {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"}, - ) - - if f"input_blocks.{i}.0.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.get( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.get( - f"input_blocks.{i}.0.op.bias" - ) - - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - if attentions: - update_unet_attention_ldm_to_diffusers( - attentions, - new_checkpoint, - unet_state_dict, - {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"}, - ) - - # Mid blocks - for key in middle_blocks.keys(): - diffusers_key = max(key - 1, 0) - if key % 2 == 0: - update_unet_resnet_ldm_to_diffusers( - middle_blocks[key], - new_checkpoint, - unet_state_dict, - mapping={"old": f"middle_block.{key}", "new": f"mid_block.resnets.{diffusers_key}"}, - ) - else: - update_unet_attention_ldm_to_diffusers( - middle_blocks[key], - new_checkpoint, - unet_state_dict, - mapping={"old": f"middle_block.{key}", "new": f"mid_block.attentions.{diffusers_key}"}, - ) - - # Up Blocks - for i in range(num_output_blocks): - block_id = i // (config["layers_per_block"] + 1) - layer_in_block_id = i % (config["layers_per_block"] + 1) - - resnets = [ - key for key in output_blocks[i] if f"output_blocks.{i}.0" in key and f"output_blocks.{i}.0.op" not in key - ] - update_unet_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - unet_state_dict, - {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"}, - ) - - attentions = [ - key for key in output_blocks[i] if f"output_blocks.{i}.1" in key and f"output_blocks.{i}.1.conv" not in key - ] - if attentions: - update_unet_attention_ldm_to_diffusers( - attentions, - new_checkpoint, - unet_state_dict, - {"old": f"output_blocks.{i}.1", "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}"}, - ) - - if f"output_blocks.{i}.1.conv.weight" in unet_state_dict: - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.1.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.1.conv.bias" - ] - if f"output_blocks.{i}.2.conv.weight" in unet_state_dict: - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.2.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.2.conv.bias" - ] - - return new_checkpoint - - -def convert_controlnet_checkpoint( - checkpoint, - config, - **kwargs, -): - # Return checkpoint if it's already been converted - if "time_embedding.linear_1.weight" in checkpoint: - return checkpoint - # Some controlnet ckpt files are distributed independently from the rest of the - # model components i.e. https://huggingface.co/thibaud/controlnet-sd21/ - if "time_embed.0.weight" in checkpoint: - controlnet_state_dict = checkpoint - - else: - controlnet_state_dict = {} - keys = list(checkpoint.keys()) - controlnet_key = LDM_CONTROLNET_KEY - for key in keys: - if key.startswith(controlnet_key): - controlnet_state_dict[key.replace(controlnet_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - ldm_controlnet_keys = DIFFUSERS_TO_LDM_MAPPING["controlnet"]["layers"] - for diffusers_key, ldm_key in ldm_controlnet_keys.items(): - if ldm_key not in controlnet_state_dict: - continue - new_checkpoint[diffusers_key] = controlnet_state_dict[ldm_key] - - # Retrieves the keys for the input blocks only - num_input_blocks = len( - {".".join(layer.split(".")[:2]) for layer in controlnet_state_dict if "input_blocks" in layer} - ) - input_blocks = { - layer_id: [key for key in controlnet_state_dict if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Down blocks - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - update_unet_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - controlnet_state_dict, - {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"}, - ) - - if f"input_blocks.{i}.0.op.weight" in controlnet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = controlnet_state_dict.get( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = controlnet_state_dict.get( - f"input_blocks.{i}.0.op.bias" - ) - - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - if attentions: - update_unet_attention_ldm_to_diffusers( - attentions, - new_checkpoint, - controlnet_state_dict, - {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"}, - ) - - # controlnet down blocks - for i in range(num_input_blocks): - new_checkpoint[f"controlnet_down_blocks.{i}.weight"] = controlnet_state_dict.get(f"zero_convs.{i}.0.weight") - new_checkpoint[f"controlnet_down_blocks.{i}.bias"] = controlnet_state_dict.get(f"zero_convs.{i}.0.bias") - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len( - {".".join(layer.split(".")[:2]) for layer in controlnet_state_dict if "middle_block" in layer} - ) - middle_blocks = { - layer_id: [key for key in controlnet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Mid blocks - for key in middle_blocks.keys(): - diffusers_key = max(key - 1, 0) - if key % 2 == 0: - update_unet_resnet_ldm_to_diffusers( - middle_blocks[key], - new_checkpoint, - controlnet_state_dict, - mapping={"old": f"middle_block.{key}", "new": f"mid_block.resnets.{diffusers_key}"}, - ) - else: - update_unet_attention_ldm_to_diffusers( - middle_blocks[key], - new_checkpoint, - controlnet_state_dict, - mapping={"old": f"middle_block.{key}", "new": f"mid_block.attentions.{diffusers_key}"}, - ) - - # mid block - new_checkpoint["controlnet_mid_block.weight"] = controlnet_state_dict.get("middle_block_out.0.weight") - new_checkpoint["controlnet_mid_block.bias"] = controlnet_state_dict.get("middle_block_out.0.bias") - - # controlnet cond embedding blocks - cond_embedding_blocks = { - ".".join(layer.split(".")[:2]) - for layer in controlnet_state_dict - if "input_hint_block" in layer and ("input_hint_block.0" not in layer) and ("input_hint_block.14" not in layer) - } - num_cond_embedding_blocks = len(cond_embedding_blocks) - - for idx in range(1, num_cond_embedding_blocks + 1): - diffusers_idx = idx - 1 - cond_block_id = 2 * idx - - new_checkpoint[f"controlnet_cond_embedding.blocks.{diffusers_idx}.weight"] = controlnet_state_dict.get( - f"input_hint_block.{cond_block_id}.weight" - ) - new_checkpoint[f"controlnet_cond_embedding.blocks.{diffusers_idx}.bias"] = controlnet_state_dict.get( - f"input_hint_block.{cond_block_id}.bias" - ) - - return new_checkpoint - - -def convert_ldm_vae_checkpoint(checkpoint, config): - # extract state dict for VAE - # remove the LDM_VAE_KEY prefix from the ldm checkpoint keys so that it is easier to map them to diffusers keys - vae_state_dict = {} - keys = list(checkpoint.keys()) - vae_key = "" - for ldm_vae_key in LDM_VAE_KEYS: - if any(k.startswith(ldm_vae_key) for k in keys): - vae_key = ldm_vae_key - - for key in keys: - if key.startswith(vae_key): - vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - vae_diffusers_ldm_map = DIFFUSERS_TO_LDM_MAPPING["vae"] - for diffusers_key, ldm_key in vae_diffusers_ldm_map.items(): - if ldm_key not in vae_state_dict: - continue - new_checkpoint[diffusers_key] = vae_state_dict[ldm_key] - - # Retrieves the keys for the encoder down blocks only - num_down_blocks = len(config["down_block_types"]) - down_blocks = { - layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - for i in range(num_down_blocks): - resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] - update_vae_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - vae_state_dict, - mapping={"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}, - ) - if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.get( - f"encoder.down.{i}.downsample.conv.weight" - ) - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.get( - f"encoder.down.{i}.downsample.conv.bias" - ) - - mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] - update_vae_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - vae_state_dict, - mapping={"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}, - ) - - mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] - update_vae_attentions_ldm_to_diffusers( - mid_attentions, new_checkpoint, vae_state_dict, mapping={"old": "mid.attn_1", "new": "mid_block.attentions.0"} - ) - - # Retrieves the keys for the decoder up blocks only - num_up_blocks = len(config["up_block_types"]) - up_blocks = { - layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) - } - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - resnets = [ - key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key - ] - update_vae_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - vae_state_dict, - mapping={"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}, - ) - if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.bias" - ] - - mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] - update_vae_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - vae_state_dict, - mapping={"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}, - ) - - mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] - update_vae_attentions_ldm_to_diffusers( - mid_attentions, new_checkpoint, vae_state_dict, mapping={"old": "mid.attn_1", "new": "mid_block.attentions.0"} - ) - conv_attn_to_linear(new_checkpoint) - - return new_checkpoint - - -def convert_ldm_clip_checkpoint(checkpoint, remove_prefix=None): - keys = list(checkpoint.keys()) - text_model_dict = {} - - remove_prefixes = [] - remove_prefixes.extend(LDM_CLIP_PREFIX_TO_REMOVE) - if remove_prefix: - remove_prefixes.append(remove_prefix) - - for key in keys: - for prefix in remove_prefixes: - if key.startswith(prefix): - diffusers_key = key.replace(prefix, "") - text_model_dict[diffusers_key] = checkpoint.get(key) - - return text_model_dict - - -def convert_open_clip_checkpoint( - text_model, - checkpoint, - prefix="cond_stage_model.model.", -): - text_model_dict = {} - text_proj_key = prefix + "text_projection" - - if text_proj_key in checkpoint: - text_proj_dim = int(checkpoint[text_proj_key].shape[0]) - elif hasattr(text_model.config, "hidden_size"): - text_proj_dim = text_model.config.hidden_size - else: - text_proj_dim = LDM_OPEN_CLIP_TEXT_PROJECTION_DIM - - keys = list(checkpoint.keys()) - keys_to_ignore = SD_2_TEXT_ENCODER_KEYS_TO_IGNORE - - openclip_diffusers_ldm_map = DIFFUSERS_TO_LDM_MAPPING["openclip"]["layers"] - for diffusers_key, ldm_key in openclip_diffusers_ldm_map.items(): - ldm_key = prefix + ldm_key - if ldm_key not in checkpoint: - continue - if ldm_key in keys_to_ignore: - continue - if ldm_key.endswith("text_projection"): - text_model_dict[diffusers_key] = checkpoint[ldm_key].T.contiguous() - else: - text_model_dict[diffusers_key] = checkpoint[ldm_key] - - for key in keys: - if key in keys_to_ignore: - continue - - if not key.startswith(prefix + "transformer."): - continue - - diffusers_key = key.replace(prefix + "transformer.", "") - transformer_diffusers_to_ldm_map = DIFFUSERS_TO_LDM_MAPPING["openclip"]["transformer"] - for new_key, old_key in transformer_diffusers_to_ldm_map.items(): - diffusers_key = ( - diffusers_key.replace(old_key, new_key).replace(".in_proj_weight", "").replace(".in_proj_bias", "") - ) - - if key.endswith(".in_proj_weight"): - weight_value = checkpoint.get(key) - - text_model_dict[diffusers_key + ".q_proj.weight"] = weight_value[:text_proj_dim, :].clone().detach() - text_model_dict[diffusers_key + ".k_proj.weight"] = ( - weight_value[text_proj_dim : text_proj_dim * 2, :].clone().detach() - ) - text_model_dict[diffusers_key + ".v_proj.weight"] = weight_value[text_proj_dim * 2 :, :].clone().detach() - - elif key.endswith(".in_proj_bias"): - weight_value = checkpoint.get(key) - text_model_dict[diffusers_key + ".q_proj.bias"] = weight_value[:text_proj_dim].clone().detach() - text_model_dict[diffusers_key + ".k_proj.bias"] = ( - weight_value[text_proj_dim : text_proj_dim * 2].clone().detach() - ) - text_model_dict[diffusers_key + ".v_proj.bias"] = weight_value[text_proj_dim * 2 :].clone().detach() - else: - text_model_dict[diffusers_key] = checkpoint.get(key) - - return text_model_dict - - -def create_diffusers_clip_model_from_ldm( - cls, - checkpoint, - subfolder="", - config=None, - torch_dtype=None, - local_files_only=None, - is_legacy_loading=False, -): - if config: - config = {"pretrained_model_name_or_path": config} - else: - config = fetch_diffusers_config(checkpoint) - - # For backwards compatibility - # Older versions of `from_single_file` expected CLIP configs to be placed in their original transformers model repo - # in the cache_dir, rather than in a subfolder of the Diffusers model - if is_legacy_loading: - logger.warning( - ( - "Detected legacy CLIP loading behavior. Please run `from_single_file` with `local_files_only=False once to update " - "the local cache directory with the necessary CLIP model config files. " - "Attempting to load CLIP model from legacy cache directory." - ) - ) - - if is_clip_model(checkpoint) or is_clip_sdxl_model(checkpoint): - clip_config = "openai/clip-vit-large-patch14" - config["pretrained_model_name_or_path"] = clip_config - subfolder = "" - - elif is_open_clip_model(checkpoint): - clip_config = "stabilityai/stable-diffusion-2" - config["pretrained_model_name_or_path"] = clip_config - subfolder = "text_encoder" - - else: - clip_config = "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k" - config["pretrained_model_name_or_path"] = clip_config - subfolder = "" - - model_config = cls.config_class.from_pretrained(**config, subfolder=subfolder, local_files_only=local_files_only) - ctx = init_empty_weights if is_accelerate_available() else nullcontext - with ctx(): - model = cls(model_config) - - position_embedding_dim = model.text_model.embeddings.position_embedding.weight.shape[-1] - - if is_clip_model(checkpoint): - diffusers_format_checkpoint = convert_ldm_clip_checkpoint(checkpoint) - - elif ( - is_clip_sdxl_model(checkpoint) - and checkpoint[CHECKPOINT_KEY_NAMES["clip_sdxl"]].shape[-1] == position_embedding_dim - ): - diffusers_format_checkpoint = convert_ldm_clip_checkpoint(checkpoint) - - elif ( - is_clip_sd3_model(checkpoint) - and checkpoint[CHECKPOINT_KEY_NAMES["clip_sd3"]].shape[-1] == position_embedding_dim - ): - diffusers_format_checkpoint = convert_ldm_clip_checkpoint(checkpoint, "text_encoders.clip_l.transformer.") - diffusers_format_checkpoint["text_projection.weight"] = torch.eye(position_embedding_dim) - - elif is_open_clip_model(checkpoint): - prefix = "cond_stage_model.model." - diffusers_format_checkpoint = convert_open_clip_checkpoint(model, checkpoint, prefix=prefix) - - elif ( - is_open_clip_sdxl_model(checkpoint) - and checkpoint[CHECKPOINT_KEY_NAMES["open_clip_sdxl"]].shape[-1] == position_embedding_dim - ): - prefix = "conditioner.embedders.1.model." - diffusers_format_checkpoint = convert_open_clip_checkpoint(model, checkpoint, prefix=prefix) - - elif is_open_clip_sdxl_refiner_model(checkpoint): - prefix = "conditioner.embedders.0.model." - diffusers_format_checkpoint = convert_open_clip_checkpoint(model, checkpoint, prefix=prefix) - - elif ( - is_open_clip_sd3_model(checkpoint) - and checkpoint[CHECKPOINT_KEY_NAMES["open_clip_sd3"]].shape[-1] == position_embedding_dim - ): - diffusers_format_checkpoint = convert_ldm_clip_checkpoint(checkpoint, "text_encoders.clip_g.transformer.") - - else: - raise ValueError("The provided checkpoint does not seem to contain a valid CLIP model.") - - if is_accelerate_available(): - load_model_dict_into_meta(model, diffusers_format_checkpoint, dtype=torch_dtype) - else: - model.load_state_dict(diffusers_format_checkpoint, strict=False) - - if torch_dtype is not None: - model.to(torch_dtype) - - model.eval() - - return model - - -def _legacy_load_scheduler( - cls, - checkpoint, - component_name, - original_config=None, - **kwargs, -): - scheduler_type = kwargs.get("scheduler_type", None) - prediction_type = kwargs.get("prediction_type", None) - - if scheduler_type is not None: - deprecation_message = ( - "Please pass an instance of a Scheduler object directly to the `scheduler` argument in `from_single_file`\n\n" - "Example:\n\n" - "from diffusers import StableDiffusionPipeline, DDIMScheduler\n\n" - "scheduler = DDIMScheduler()\n" - "pipe = StableDiffusionPipeline.from_single_file(, scheduler=scheduler)\n" - ) - deprecate("scheduler_type", "1.0.0", deprecation_message) - - if prediction_type is not None: - deprecation_message = ( - "Please configure an instance of a Scheduler with the appropriate `prediction_type` and " - "pass the object directly to the `scheduler` argument in `from_single_file`.\n\n" - "Example:\n\n" - "from diffusers import StableDiffusionPipeline, DDIMScheduler\n\n" - 'scheduler = DDIMScheduler(prediction_type="v_prediction")\n' - "pipe = StableDiffusionPipeline.from_single_file(, scheduler=scheduler)\n" - ) - deprecate("prediction_type", "1.0.0", deprecation_message) - - scheduler_config = SCHEDULER_DEFAULT_CONFIG - model_type = infer_diffusers_model_type(checkpoint=checkpoint) - - global_step = checkpoint["global_step"] if "global_step" in checkpoint else None - - if original_config: - num_train_timesteps = getattr(original_config["model"]["params"], "timesteps", 1000) - else: - num_train_timesteps = 1000 - - scheduler_config["num_train_timesteps"] = num_train_timesteps - - if model_type == "v2": - if prediction_type is None: - # NOTE: For stable diffusion 2 base it is recommended to pass `prediction_type=="epsilon"` # as it relies on a brittle global step parameter here - prediction_type = "epsilon" if global_step == 875000 else "v_prediction" - - else: - prediction_type = prediction_type or "epsilon" - - scheduler_config["prediction_type"] = prediction_type - - if model_type in ["xl_base", "xl_refiner"]: - scheduler_type = "euler" - elif model_type == "playground": - scheduler_type = "edm_dpm_solver_multistep" - else: - if original_config: - beta_start = original_config["model"]["params"].get("linear_start") - beta_end = original_config["model"]["params"].get("linear_end") - - else: - beta_start = 0.02 - beta_end = 0.085 - - scheduler_config["beta_start"] = beta_start - scheduler_config["beta_end"] = beta_end - scheduler_config["beta_schedule"] = "scaled_linear" - scheduler_config["clip_sample"] = False - scheduler_config["set_alpha_to_one"] = False - - # to deal with an edge case StableDiffusionUpscale pipeline has two schedulers - if component_name == "low_res_scheduler": - return cls.from_config( - { - "beta_end": 0.02, - "beta_schedule": "scaled_linear", - "beta_start": 0.0001, - "clip_sample": True, - "num_train_timesteps": 1000, - "prediction_type": "epsilon", - "trained_betas": None, - "variance_type": "fixed_small", - } - ) - - if scheduler_type is None: - return cls.from_config(scheduler_config) - - elif scheduler_type == "pndm": - scheduler_config["skip_prk_steps"] = True - scheduler = PNDMScheduler.from_config(scheduler_config) - - elif scheduler_type == "lms": - scheduler = LMSDiscreteScheduler.from_config(scheduler_config) - - elif scheduler_type == "heun": - scheduler = HeunDiscreteScheduler.from_config(scheduler_config) - - elif scheduler_type == "euler": - scheduler = EulerDiscreteScheduler.from_config(scheduler_config) - - elif scheduler_type == "euler-ancestral": - scheduler = EulerAncestralDiscreteScheduler.from_config(scheduler_config) - - elif scheduler_type == "dpm": - scheduler = DPMSolverMultistepScheduler.from_config(scheduler_config) - - elif scheduler_type == "ddim": - scheduler = DDIMScheduler.from_config(scheduler_config) - - elif scheduler_type == "edm_dpm_solver_multistep": - scheduler_config = { - "algorithm_type": "dpmsolver++", - "dynamic_thresholding_ratio": 0.995, - "euler_at_final": False, - "final_sigmas_type": "zero", - "lower_order_final": True, - "num_train_timesteps": 1000, - "prediction_type": "epsilon", - "rho": 7.0, - "sample_max_value": 1.0, - "sigma_data": 0.5, - "sigma_max": 80.0, - "sigma_min": 0.002, - "solver_order": 2, - "solver_type": "midpoint", - "thresholding": False, - } - scheduler = EDMDPMSolverMultistepScheduler(**scheduler_config) - - else: - raise ValueError(f"Scheduler of type {scheduler_type} doesn't exist!") - - return scheduler - - -def _legacy_load_clip_tokenizer(cls, checkpoint, config=None, local_files_only=False): - if config: - config = {"pretrained_model_name_or_path": config} - else: - config = fetch_diffusers_config(checkpoint) - - if is_clip_model(checkpoint) or is_clip_sdxl_model(checkpoint): - clip_config = "openai/clip-vit-large-patch14" - config["pretrained_model_name_or_path"] = clip_config - subfolder = "" - - elif is_open_clip_model(checkpoint): - clip_config = "stabilityai/stable-diffusion-2" - config["pretrained_model_name_or_path"] = clip_config - subfolder = "tokenizer" - - else: - clip_config = "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k" - config["pretrained_model_name_or_path"] = clip_config - subfolder = "" - - tokenizer = cls.from_pretrained(**config, subfolder=subfolder, local_files_only=local_files_only) - - return tokenizer - - -def _legacy_load_safety_checker(local_files_only, torch_dtype): - # Support for loading safety checker components using the deprecated - # `load_safety_checker` argument. - - from ..pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - feature_extractor = AutoImageProcessor.from_pretrained( - "CompVis/stable-diffusion-safety-checker", local_files_only=local_files_only, torch_dtype=torch_dtype - ) - safety_checker = StableDiffusionSafetyChecker.from_pretrained( - "CompVis/stable-diffusion-safety-checker", local_files_only=local_files_only, torch_dtype=torch_dtype - ) - - return {"safety_checker": safety_checker, "feature_extractor": feature_extractor} - - -# 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, dim): - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - return new_weight - - -def swap_proj_gate(weight): - proj, gate = weight.chunk(2, dim=0) - new_weight = torch.cat([gate, proj], dim=0) - return new_weight - - -def get_attn2_layers(state_dict): - attn2_layers = [] - for key in state_dict.keys(): - if "attn2." in key: - # Extract the layer number from the key - layer_num = int(key.split(".")[1]) - attn2_layers.append(layer_num) - - return tuple(sorted(set(attn2_layers))) - - -def get_caption_projection_dim(state_dict): - caption_projection_dim = state_dict["context_embedder.weight"].shape[0] - return caption_projection_dim - - -def convert_sd3_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 "joint_blocks" in k))[-1] + 1 # noqa: C401 - dual_attention_layers = get_attn2_layers(checkpoint) - - caption_projection_dim = get_caption_projection_dim(checkpoint) - has_qk_norm = any("ln_q" in key for key in checkpoint.keys()) - - # Positional and patch embeddings. - converted_state_dict["pos_embed.pos_embed"] = checkpoint.pop("pos_embed") - converted_state_dict["pos_embed.proj.weight"] = checkpoint.pop("x_embedder.proj.weight") - converted_state_dict["pos_embed.proj.bias"] = checkpoint.pop("x_embedder.proj.bias") - - # Timestep embeddings. - converted_state_dict["time_text_embed.timestep_embedder.linear_1.weight"] = checkpoint.pop( - "t_embedder.mlp.0.weight" - ) - converted_state_dict["time_text_embed.timestep_embedder.linear_1.bias"] = checkpoint.pop("t_embedder.mlp.0.bias") - converted_state_dict["time_text_embed.timestep_embedder.linear_2.weight"] = checkpoint.pop( - "t_embedder.mlp.2.weight" - ) - converted_state_dict["time_text_embed.timestep_embedder.linear_2.bias"] = checkpoint.pop("t_embedder.mlp.2.bias") - - # Context projections. - converted_state_dict["context_embedder.weight"] = checkpoint.pop("context_embedder.weight") - converted_state_dict["context_embedder.bias"] = checkpoint.pop("context_embedder.bias") - - # Pooled context projection. - converted_state_dict["time_text_embed.text_embedder.linear_1.weight"] = checkpoint.pop("y_embedder.mlp.0.weight") - converted_state_dict["time_text_embed.text_embedder.linear_1.bias"] = checkpoint.pop("y_embedder.mlp.0.bias") - converted_state_dict["time_text_embed.text_embedder.linear_2.weight"] = checkpoint.pop("y_embedder.mlp.2.weight") - converted_state_dict["time_text_embed.text_embedder.linear_2.bias"] = checkpoint.pop("y_embedder.mlp.2.bias") - - # Transformer blocks 🎸. - for i in range(num_layers): - # Q, K, V - sample_q, sample_k, sample_v = torch.chunk( - checkpoint.pop(f"joint_blocks.{i}.x_block.attn.qkv.weight"), 3, dim=0 - ) - context_q, context_k, context_v = torch.chunk( - checkpoint.pop(f"joint_blocks.{i}.context_block.attn.qkv.weight"), 3, dim=0 - ) - sample_q_bias, sample_k_bias, sample_v_bias = torch.chunk( - checkpoint.pop(f"joint_blocks.{i}.x_block.attn.qkv.bias"), 3, dim=0 - ) - context_q_bias, context_k_bias, context_v_bias = torch.chunk( - checkpoint.pop(f"joint_blocks.{i}.context_block.attn.qkv.bias"), 3, dim=0 - ) - - converted_state_dict[f"transformer_blocks.{i}.attn.to_q.weight"] = torch.cat([sample_q]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_q.bias"] = torch.cat([sample_q_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_k.weight"] = torch.cat([sample_k]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_k.bias"] = torch.cat([sample_k_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_v.weight"] = torch.cat([sample_v]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_v.bias"] = torch.cat([sample_v_bias]) - - converted_state_dict[f"transformer_blocks.{i}.attn.add_q_proj.weight"] = torch.cat([context_q]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_q_proj.bias"] = torch.cat([context_q_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_k_proj.weight"] = torch.cat([context_k]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_k_proj.bias"] = torch.cat([context_k_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_v_proj.weight"] = torch.cat([context_v]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_v_proj.bias"] = torch.cat([context_v_bias]) - - # qk norm - if has_qk_norm: - converted_state_dict[f"transformer_blocks.{i}.attn.norm_q.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn.ln_q.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.norm_k.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn.ln_k.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.norm_added_q.weight"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.attn.ln_q.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.norm_added_k.weight"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.attn.ln_k.weight" - ) - - # output projections. - converted_state_dict[f"transformer_blocks.{i}.attn.to_out.0.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.to_out.0.bias"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn.proj.bias" - ) - if not (i == num_layers - 1): - converted_state_dict[f"transformer_blocks.{i}.attn.to_add_out.weight"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.to_add_out.bias"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.attn.proj.bias" - ) - - if i in dual_attention_layers: - # Q, K, V - sample_q2, sample_k2, sample_v2 = torch.chunk( - checkpoint.pop(f"joint_blocks.{i}.x_block.attn2.qkv.weight"), 3, dim=0 - ) - sample_q2_bias, sample_k2_bias, sample_v2_bias = torch.chunk( - checkpoint.pop(f"joint_blocks.{i}.x_block.attn2.qkv.bias"), 3, dim=0 - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_q.weight"] = torch.cat([sample_q2]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_q.bias"] = torch.cat([sample_q2_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_k.weight"] = torch.cat([sample_k2]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_k.bias"] = torch.cat([sample_k2_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_v.weight"] = torch.cat([sample_v2]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_v.bias"] = torch.cat([sample_v2_bias]) - - # qk norm - if has_qk_norm: - converted_state_dict[f"transformer_blocks.{i}.attn2.norm_q.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn2.ln_q.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.norm_k.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn2.ln_k.weight" - ) - - # output projections. - converted_state_dict[f"transformer_blocks.{i}.attn2.to_out.0.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn2.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_out.0.bias"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn2.proj.bias" - ) - - # norms. - converted_state_dict[f"transformer_blocks.{i}.norm1.linear.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.adaLN_modulation.1.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.norm1.linear.bias"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.adaLN_modulation.1.bias" - ) - if not (i == num_layers - 1): - converted_state_dict[f"transformer_blocks.{i}.norm1_context.linear.weight"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.adaLN_modulation.1.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.norm1_context.linear.bias"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.adaLN_modulation.1.bias" - ) - else: - converted_state_dict[f"transformer_blocks.{i}.norm1_context.linear.weight"] = swap_scale_shift( - checkpoint.pop(f"joint_blocks.{i}.context_block.adaLN_modulation.1.weight"), - dim=caption_projection_dim, - ) - converted_state_dict[f"transformer_blocks.{i}.norm1_context.linear.bias"] = swap_scale_shift( - checkpoint.pop(f"joint_blocks.{i}.context_block.adaLN_modulation.1.bias"), - dim=caption_projection_dim, - ) - - # ffs. - converted_state_dict[f"transformer_blocks.{i}.ff.net.0.proj.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.mlp.fc1.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.net.0.proj.bias"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.mlp.fc1.bias" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.net.2.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.mlp.fc2.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.net.2.bias"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.mlp.fc2.bias" - ) - if not (i == num_layers - 1): - converted_state_dict[f"transformer_blocks.{i}.ff_context.net.0.proj.weight"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.mlp.fc1.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff_context.net.0.proj.bias"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.mlp.fc1.bias" - ) - converted_state_dict[f"transformer_blocks.{i}.ff_context.net.2.weight"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.mlp.fc2.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff_context.net.2.bias"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.mlp.fc2.bias" - ) - - # Final blocks. - converted_state_dict["proj_out.weight"] = checkpoint.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = checkpoint.pop("final_layer.linear.bias") - converted_state_dict["norm_out.linear.weight"] = swap_scale_shift( - checkpoint.pop("final_layer.adaLN_modulation.1.weight"), dim=caption_projection_dim - ) - converted_state_dict["norm_out.linear.bias"] = swap_scale_shift( - checkpoint.pop("final_layer.adaLN_modulation.1.bias"), dim=caption_projection_dim - ) - - return converted_state_dict - - -def is_t5_in_single_file(checkpoint): - if "text_encoders.t5xxl.transformer.shared.weight" in checkpoint: - return True - - return False - - -def convert_sd3_t5_checkpoint_to_diffusers(checkpoint): - keys = list(checkpoint.keys()) - text_model_dict = {} - - remove_prefixes = ["text_encoders.t5xxl.transformer."] - - for key in keys: - for prefix in remove_prefixes: - if key.startswith(prefix): - diffusers_key = key.replace(prefix, "") - text_model_dict[diffusers_key] = checkpoint.get(key) - - return text_model_dict - - -def create_diffusers_t5_model_from_checkpoint( - cls, - checkpoint, - subfolder="", - config=None, - torch_dtype=None, - local_files_only=None, -): - if config: - config = {"pretrained_model_name_or_path": config} - else: - config = fetch_diffusers_config(checkpoint) - - model_config = cls.config_class.from_pretrained(**config, subfolder=subfolder, local_files_only=local_files_only) - ctx = init_empty_weights if is_accelerate_available() else nullcontext - with ctx(): - model = cls(model_config) - - diffusers_format_checkpoint = convert_sd3_t5_checkpoint_to_diffusers(checkpoint) - - if is_accelerate_available(): - load_model_dict_into_meta(model, diffusers_format_checkpoint, dtype=torch_dtype) - else: - model.load_state_dict(diffusers_format_checkpoint) - - use_keep_in_fp32_modules = (cls._keep_in_fp32_modules is not None) and (torch_dtype == torch.float16) - if use_keep_in_fp32_modules: - keep_in_fp32_modules = model._keep_in_fp32_modules - else: - keep_in_fp32_modules = [] - - if keep_in_fp32_modules is not None: - for name, param in model.named_parameters(): - if any(module_to_keep_in_fp32 in name.split(".") for module_to_keep_in_fp32 in keep_in_fp32_modules): - # param = param.to(torch.float32) does not work here as only in the local scope. - param.data = param.data.to(torch.float32) - - return model - - -def convert_animatediff_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {} - for k, v in checkpoint.items(): - if "pos_encoder" in k: - continue - - else: - converted_state_dict[ - k.replace(".norms.0", ".norm1") - .replace(".norms.1", ".norm2") - .replace(".ff_norm", ".norm3") - .replace(".attention_blocks.0", ".attn1") - .replace(".attention_blocks.1", ".attn2") - .replace(".temporal_transformer", "") - ] = v - - return converted_state_dict - - -def convert_flux_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 - 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 - - ## time_text_embed.timestep_embedder <- time_in - converted_state_dict["time_text_embed.timestep_embedder.linear_1.weight"] = checkpoint.pop( - "time_in.in_layer.weight" - ) - converted_state_dict["time_text_embed.timestep_embedder.linear_1.bias"] = checkpoint.pop("time_in.in_layer.bias") - converted_state_dict["time_text_embed.timestep_embedder.linear_2.weight"] = checkpoint.pop( - "time_in.out_layer.weight" - ) - converted_state_dict["time_text_embed.timestep_embedder.linear_2.bias"] = checkpoint.pop("time_in.out_layer.bias") - - ## time_text_embed.text_embedder <- vector_in - converted_state_dict["time_text_embed.text_embedder.linear_1.weight"] = checkpoint.pop("vector_in.in_layer.weight") - converted_state_dict["time_text_embed.text_embedder.linear_1.bias"] = checkpoint.pop("vector_in.in_layer.bias") - converted_state_dict["time_text_embed.text_embedder.linear_2.weight"] = checkpoint.pop( - "vector_in.out_layer.weight" - ) - converted_state_dict["time_text_embed.text_embedder.linear_2.bias"] = checkpoint.pop("vector_in.out_layer.bias") - - # guidance - has_guidance = any("guidance" in k for k in checkpoint) - if has_guidance: - converted_state_dict["time_text_embed.guidance_embedder.linear_1.weight"] = checkpoint.pop( - "guidance_in.in_layer.weight" - ) - converted_state_dict["time_text_embed.guidance_embedder.linear_1.bias"] = checkpoint.pop( - "guidance_in.in_layer.bias" - ) - converted_state_dict["time_text_embed.guidance_embedder.linear_2.weight"] = checkpoint.pop( - "guidance_in.out_layer.weight" - ) - converted_state_dict["time_text_embed.guidance_embedder.linear_2.bias"] = checkpoint.pop( - "guidance_in.out_layer.bias" - ) - - # context_embedder - converted_state_dict["context_embedder.weight"] = checkpoint.pop("txt_in.weight") - converted_state_dict["context_embedder.bias"] = checkpoint.pop("txt_in.bias") - - # x_embedder - converted_state_dict["x_embedder.weight"] = checkpoint.pop("img_in.weight") - converted_state_dict["x_embedder.bias"] = checkpoint.pop("img_in.bias") - - # double transformer blocks - for i in range(num_layers): - block_prefix = f"transformer_blocks.{i}." - # norms. - ## norm1 - converted_state_dict[f"{block_prefix}norm1.linear.weight"] = checkpoint.pop( - f"double_blocks.{i}.img_mod.lin.weight" - ) - converted_state_dict[f"{block_prefix}norm1.linear.bias"] = checkpoint.pop( - f"double_blocks.{i}.img_mod.lin.bias" - ) - ## norm1_context - converted_state_dict[f"{block_prefix}norm1_context.linear.weight"] = checkpoint.pop( - f"double_blocks.{i}.txt_mod.lin.weight" - ) - converted_state_dict[f"{block_prefix}norm1_context.linear.bias"] = checkpoint.pop( - f"double_blocks.{i}.txt_mod.lin.bias" - ) - # 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}." - # norm.linear <- single_blocks.0.modulation.lin - converted_state_dict[f"{block_prefix}norm.linear.weight"] = checkpoint.pop( - f"single_blocks.{i}.modulation.lin.weight" - ) - converted_state_dict[f"{block_prefix}norm.linear.bias"] = checkpoint.pop( - f"single_blocks.{i}.modulation.lin.bias" - ) - # 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") - - converted_state_dict["proj_out.weight"] = checkpoint.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = checkpoint.pop("final_layer.linear.bias") - converted_state_dict["norm_out.linear.weight"] = swap_scale_shift( - checkpoint.pop("final_layer.adaLN_modulation.1.weight") - ) - converted_state_dict["norm_out.linear.bias"] = swap_scale_shift( - checkpoint.pop("final_layer.adaLN_modulation.1.bias") - ) - - return converted_state_dict - - -def convert_ltx_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys()) if "vae" not in key} - - TRANSFORMER_KEYS_RENAME_DICT = { - "model.diffusion_model.": "", - "patchify_proj": "proj_in", - "adaln_single": "time_embed", - "q_norm": "norm_q", - "k_norm": "norm_k", - } - - TRANSFORMER_SPECIAL_KEYS_REMAP = {} - - for key in list(converted_state_dict.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] = converted_state_dict.pop(key) - - for key in list(converted_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, converted_state_dict) - - return converted_state_dict - - -def convert_ltx_vae_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys()) if "vae." in key} - - def remove_keys_(key: str, state_dict): - state_dict.pop(key) - - VAE_KEYS_RENAME_DICT = { - # common - "vae.": "", - # decoder - "up_blocks.0": "mid_block", - "up_blocks.1": "up_blocks.0", - "up_blocks.2": "up_blocks.1.upsamplers.0", - "up_blocks.3": "up_blocks.1", - "up_blocks.4": "up_blocks.2.conv_in", - "up_blocks.5": "up_blocks.2.upsamplers.0", - "up_blocks.6": "up_blocks.2", - "up_blocks.7": "up_blocks.3.conv_in", - "up_blocks.8": "up_blocks.3.upsamplers.0", - "up_blocks.9": "up_blocks.3", - # encoder - "down_blocks.0": "down_blocks.0", - "down_blocks.1": "down_blocks.0.downsamplers.0", - "down_blocks.2": "down_blocks.0.conv_out", - "down_blocks.3": "down_blocks.1", - "down_blocks.4": "down_blocks.1.downsamplers.0", - "down_blocks.5": "down_blocks.1.conv_out", - "down_blocks.6": "down_blocks.2", - "down_blocks.7": "down_blocks.2.downsamplers.0", - "down_blocks.8": "down_blocks.3", - "down_blocks.9": "mid_block", - # common - "conv_shortcut": "conv_shortcut.conv", - "res_blocks": "resnets", - "norm3.norm": "norm3", - "per_channel_statistics.mean-of-means": "latents_mean", - "per_channel_statistics.std-of-means": "latents_std", - } - - VAE_091_RENAME_DICT = { - # 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", - "up_blocks.7": "up_blocks.3.upsamplers.0", - "up_blocks.8": "up_blocks.3", - # common - "last_time_embedder": "time_embedder", - "last_scale_shift_table": "scale_shift_table", - } - - VAE_095_RENAME_DICT = { - # 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", - "up_blocks.7": "up_blocks.3.upsamplers.0", - "up_blocks.8": "up_blocks.3", - # 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", - # common - "last_time_embedder": "time_embedder", - "last_scale_shift_table": "scale_shift_table", - } - - VAE_SPECIAL_KEYS_REMAP = { - "per_channel_statistics.channel": remove_keys_, - "per_channel_statistics.mean-of-means": remove_keys_, - "per_channel_statistics.mean-of-stds": remove_keys_, - } - - if converted_state_dict["vae.encoder.conv_out.conv.weight"].shape[1] == 2048: - VAE_KEYS_RENAME_DICT.update(VAE_095_RENAME_DICT) - elif "vae.decoder.last_time_embedder.timestep_embedder.linear_1.weight" in converted_state_dict: - VAE_KEYS_RENAME_DICT.update(VAE_091_RENAME_DICT) - - for key in list(converted_state_dict.keys()): - new_key = key - for replace_key, rename_key in VAE_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - converted_state_dict[new_key] = converted_state_dict.pop(key) - - for key in list(converted_state_dict.keys()): - for special_key, handler_fn_inplace in 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_autoencoder_dc_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} - - def remap_qkv_(key: str, state_dict): - qkv = state_dict.pop(key) - q, k, v = torch.chunk(qkv, 3, dim=0) - parent_module, _, _ = key.rpartition(".qkv.conv.weight") - state_dict[f"{parent_module}.to_q.weight"] = q.squeeze() - state_dict[f"{parent_module}.to_k.weight"] = k.squeeze() - state_dict[f"{parent_module}.to_v.weight"] = v.squeeze() - - def remap_proj_conv_(key: str, state_dict): - parent_module, _, _ = key.rpartition(".proj.conv.weight") - state_dict[f"{parent_module}.to_out.weight"] = state_dict.pop(key).squeeze() - - AE_KEYS_RENAME_DICT = { - # common - "main.": "", - "op_list.": "", - "context_module": "attn", - "local_module": "conv_out", - # NOTE: The below two lines work because scales in the available configs only have a tuple length of 1 - # If there were more scales, there would be more layers, so a loop would be better to handle this - "aggreg.0.0": "to_qkv_multiscale.0.proj_in", - "aggreg.0.1": "to_qkv_multiscale.0.proj_out", - "depth_conv.conv": "conv_depth", - "inverted_conv.conv": "conv_inverted", - "point_conv.conv": "conv_point", - "point_conv.norm": "norm", - "conv.conv.": "conv.", - "conv1.conv": "conv1", - "conv2.conv": "conv2", - "conv2.norm": "norm", - "proj.norm": "norm_out", - # encoder - "encoder.project_in.conv": "encoder.conv_in", - "encoder.project_out.0.conv": "encoder.conv_out", - "encoder.stages": "encoder.down_blocks", - # decoder - "decoder.project_in.conv": "decoder.conv_in", - "decoder.project_out.0": "decoder.norm_out", - "decoder.project_out.2.conv": "decoder.conv_out", - "decoder.stages": "decoder.up_blocks", - } - - AE_F32C32_F64C128_F128C512_KEYS = { - "encoder.project_in.conv": "encoder.conv_in.conv", - "decoder.project_out.2.conv": "decoder.conv_out.conv", - } - - AE_SPECIAL_KEYS_REMAP = { - "qkv.conv.weight": remap_qkv_, - "proj.conv.weight": remap_proj_conv_, - } - if "encoder.project_in.conv.bias" not in converted_state_dict: - AE_KEYS_RENAME_DICT.update(AE_F32C32_F64C128_F128C512_KEYS) - - for key in list(converted_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in AE_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - converted_state_dict[new_key] = converted_state_dict.pop(key) - - for key in list(converted_state_dict.keys()): - for special_key, handler_fn_inplace in AE_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, converted_state_dict) - - return converted_state_dict - - -def convert_mochi_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {} - - # Comfy checkpoints add this prefix - keys = list(checkpoint.keys()) - for k in keys: - if "model.diffusion_model." in k: - checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) - - # Convert patch_embed - converted_state_dict["patch_embed.proj.weight"] = checkpoint.pop("x_embedder.proj.weight") - converted_state_dict["patch_embed.proj.bias"] = checkpoint.pop("x_embedder.proj.bias") - - # Convert time_embed - converted_state_dict["time_embed.timestep_embedder.linear_1.weight"] = checkpoint.pop("t_embedder.mlp.0.weight") - converted_state_dict["time_embed.timestep_embedder.linear_1.bias"] = checkpoint.pop("t_embedder.mlp.0.bias") - converted_state_dict["time_embed.timestep_embedder.linear_2.weight"] = checkpoint.pop("t_embedder.mlp.2.weight") - converted_state_dict["time_embed.timestep_embedder.linear_2.bias"] = checkpoint.pop("t_embedder.mlp.2.bias") - converted_state_dict["time_embed.pooler.to_kv.weight"] = checkpoint.pop("t5_y_embedder.to_kv.weight") - converted_state_dict["time_embed.pooler.to_kv.bias"] = checkpoint.pop("t5_y_embedder.to_kv.bias") - converted_state_dict["time_embed.pooler.to_q.weight"] = checkpoint.pop("t5_y_embedder.to_q.weight") - converted_state_dict["time_embed.pooler.to_q.bias"] = checkpoint.pop("t5_y_embedder.to_q.bias") - converted_state_dict["time_embed.pooler.to_out.weight"] = checkpoint.pop("t5_y_embedder.to_out.weight") - converted_state_dict["time_embed.pooler.to_out.bias"] = checkpoint.pop("t5_y_embedder.to_out.bias") - converted_state_dict["time_embed.caption_proj.weight"] = checkpoint.pop("t5_yproj.weight") - converted_state_dict["time_embed.caption_proj.bias"] = checkpoint.pop("t5_yproj.bias") - - # Convert transformer blocks - num_layers = 48 - for i in range(num_layers): - block_prefix = f"transformer_blocks.{i}." - old_prefix = f"blocks.{i}." - - # norm1 - converted_state_dict[block_prefix + "norm1.linear.weight"] = checkpoint.pop(old_prefix + "mod_x.weight") - converted_state_dict[block_prefix + "norm1.linear.bias"] = checkpoint.pop(old_prefix + "mod_x.bias") - if i < num_layers - 1: - converted_state_dict[block_prefix + "norm1_context.linear.weight"] = checkpoint.pop( - old_prefix + "mod_y.weight" - ) - converted_state_dict[block_prefix + "norm1_context.linear.bias"] = checkpoint.pop( - old_prefix + "mod_y.bias" - ) - else: - converted_state_dict[block_prefix + "norm1_context.linear_1.weight"] = checkpoint.pop( - old_prefix + "mod_y.weight" - ) - converted_state_dict[block_prefix + "norm1_context.linear_1.bias"] = checkpoint.pop( - old_prefix + "mod_y.bias" - ) - - # Visual attention - qkv_weight = checkpoint.pop(old_prefix + "attn.qkv_x.weight") - q, k, v = qkv_weight.chunk(3, dim=0) - - converted_state_dict[block_prefix + "attn1.to_q.weight"] = q - converted_state_dict[block_prefix + "attn1.to_k.weight"] = k - converted_state_dict[block_prefix + "attn1.to_v.weight"] = v - converted_state_dict[block_prefix + "attn1.norm_q.weight"] = checkpoint.pop( - old_prefix + "attn.q_norm_x.weight" - ) - converted_state_dict[block_prefix + "attn1.norm_k.weight"] = checkpoint.pop( - old_prefix + "attn.k_norm_x.weight" - ) - converted_state_dict[block_prefix + "attn1.to_out.0.weight"] = checkpoint.pop( - old_prefix + "attn.proj_x.weight" - ) - converted_state_dict[block_prefix + "attn1.to_out.0.bias"] = checkpoint.pop(old_prefix + "attn.proj_x.bias") - - # Context attention - qkv_weight = checkpoint.pop(old_prefix + "attn.qkv_y.weight") - q, k, v = qkv_weight.chunk(3, dim=0) - - converted_state_dict[block_prefix + "attn1.add_q_proj.weight"] = q - converted_state_dict[block_prefix + "attn1.add_k_proj.weight"] = k - converted_state_dict[block_prefix + "attn1.add_v_proj.weight"] = v - converted_state_dict[block_prefix + "attn1.norm_added_q.weight"] = checkpoint.pop( - old_prefix + "attn.q_norm_y.weight" - ) - converted_state_dict[block_prefix + "attn1.norm_added_k.weight"] = checkpoint.pop( - old_prefix + "attn.k_norm_y.weight" - ) - if i < num_layers - 1: - converted_state_dict[block_prefix + "attn1.to_add_out.weight"] = checkpoint.pop( - old_prefix + "attn.proj_y.weight" - ) - converted_state_dict[block_prefix + "attn1.to_add_out.bias"] = checkpoint.pop( - old_prefix + "attn.proj_y.bias" - ) - - # MLP - converted_state_dict[block_prefix + "ff.net.0.proj.weight"] = swap_proj_gate( - checkpoint.pop(old_prefix + "mlp_x.w1.weight") - ) - converted_state_dict[block_prefix + "ff.net.2.weight"] = checkpoint.pop(old_prefix + "mlp_x.w2.weight") - if i < num_layers - 1: - converted_state_dict[block_prefix + "ff_context.net.0.proj.weight"] = swap_proj_gate( - checkpoint.pop(old_prefix + "mlp_y.w1.weight") - ) - converted_state_dict[block_prefix + "ff_context.net.2.weight"] = checkpoint.pop( - old_prefix + "mlp_y.w2.weight" - ) - - # Output layers - converted_state_dict["norm_out.linear.weight"] = swap_scale_shift(checkpoint.pop("final_layer.mod.weight"), dim=0) - converted_state_dict["norm_out.linear.bias"] = swap_scale_shift(checkpoint.pop("final_layer.mod.bias"), dim=0) - converted_state_dict["proj_out.weight"] = checkpoint.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = checkpoint.pop("final_layer.linear.bias") - - converted_state_dict["pos_frequencies"] = checkpoint.pop("pos_frequencies") - - return converted_state_dict - - -def convert_hunyuan_video_transformer_to_diffusers(checkpoint, **kwargs): - def remap_norm_scale_shift_(key, state_dict): - weight = state_dict.pop(key) - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - state_dict[key.replace("final_layer.adaLN_modulation.1", "norm_out.linear")] = new_weight - - def remap_txt_in_(key, state_dict): - def rename_key(key): - new_key = key.replace("individual_token_refiner.blocks", "token_refiner.refiner_blocks") - new_key = new_key.replace("adaLN_modulation.1", "norm_out.linear") - new_key = new_key.replace("txt_in", "context_embedder") - new_key = new_key.replace("t_embedder.mlp.0", "time_text_embed.timestep_embedder.linear_1") - new_key = new_key.replace("t_embedder.mlp.2", "time_text_embed.timestep_embedder.linear_2") - new_key = new_key.replace("c_embedder", "time_text_embed.text_embedder") - new_key = new_key.replace("mlp", "ff") - return new_key - - if "self_attn_qkv" in key: - weight = state_dict.pop(key) - to_q, to_k, to_v = weight.chunk(3, dim=0) - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_q"))] = to_q - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_k"))] = to_k - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_v"))] = to_v - else: - state_dict[rename_key(key)] = state_dict.pop(key) - - def remap_img_attn_qkv_(key, state_dict): - weight = state_dict.pop(key) - to_q, to_k, to_v = weight.chunk(3, dim=0) - state_dict[key.replace("img_attn_qkv", "attn.to_q")] = to_q - state_dict[key.replace("img_attn_qkv", "attn.to_k")] = to_k - state_dict[key.replace("img_attn_qkv", "attn.to_v")] = to_v - - def remap_txt_attn_qkv_(key, state_dict): - weight = state_dict.pop(key) - to_q, to_k, to_v = weight.chunk(3, dim=0) - state_dict[key.replace("txt_attn_qkv", "attn.add_q_proj")] = to_q - state_dict[key.replace("txt_attn_qkv", "attn.add_k_proj")] = to_k - state_dict[key.replace("txt_attn_qkv", "attn.add_v_proj")] = to_v - - def remap_single_transformer_blocks_(key, state_dict): - hidden_size = 3072 - - if "linear1.weight" in key: - linear1_weight = state_dict.pop(key) - split_size = (hidden_size, hidden_size, hidden_size, linear1_weight.size(0) - 3 * hidden_size) - q, k, v, mlp = torch.split(linear1_weight, split_size, dim=0) - new_key = key.replace("single_blocks", "single_transformer_blocks").removesuffix(".linear1.weight") - state_dict[f"{new_key}.attn.to_q.weight"] = q - state_dict[f"{new_key}.attn.to_k.weight"] = k - state_dict[f"{new_key}.attn.to_v.weight"] = v - state_dict[f"{new_key}.proj_mlp.weight"] = mlp - - elif "linear1.bias" in key: - linear1_bias = state_dict.pop(key) - split_size = (hidden_size, hidden_size, hidden_size, linear1_bias.size(0) - 3 * hidden_size) - q_bias, k_bias, v_bias, mlp_bias = torch.split(linear1_bias, split_size, dim=0) - new_key = key.replace("single_blocks", "single_transformer_blocks").removesuffix(".linear1.bias") - state_dict[f"{new_key}.attn.to_q.bias"] = q_bias - state_dict[f"{new_key}.attn.to_k.bias"] = k_bias - state_dict[f"{new_key}.attn.to_v.bias"] = v_bias - state_dict[f"{new_key}.proj_mlp.bias"] = mlp_bias - - else: - new_key = key.replace("single_blocks", "single_transformer_blocks") - new_key = new_key.replace("linear2", "proj_out") - new_key = new_key.replace("q_norm", "attn.norm_q") - new_key = new_key.replace("k_norm", "attn.norm_k") - state_dict[new_key] = state_dict.pop(key) - - TRANSFORMER_KEYS_RENAME_DICT = { - "img_in": "x_embedder", - "time_in.mlp.0": "time_text_embed.timestep_embedder.linear_1", - "time_in.mlp.2": "time_text_embed.timestep_embedder.linear_2", - "guidance_in.mlp.0": "time_text_embed.guidance_embedder.linear_1", - "guidance_in.mlp.2": "time_text_embed.guidance_embedder.linear_2", - "vector_in.in_layer": "time_text_embed.text_embedder.linear_1", - "vector_in.out_layer": "time_text_embed.text_embedder.linear_2", - "double_blocks": "transformer_blocks", - "img_attn_q_norm": "attn.norm_q", - "img_attn_k_norm": "attn.norm_k", - "img_attn_proj": "attn.to_out.0", - "txt_attn_q_norm": "attn.norm_added_q", - "txt_attn_k_norm": "attn.norm_added_k", - "txt_attn_proj": "attn.to_add_out", - "img_mod.linear": "norm1.linear", - "img_norm1": "norm1.norm", - "img_norm2": "norm2", - "img_mlp": "ff", - "txt_mod.linear": "norm1_context.linear", - "txt_norm1": "norm1.norm", - "txt_norm2": "norm2_context", - "txt_mlp": "ff_context", - "self_attn_proj": "attn.to_out.0", - "modulation.linear": "norm.linear", - "pre_norm": "norm.norm", - "final_layer.norm_final": "norm_out.norm", - "final_layer.linear": "proj_out", - "fc1": "net.0.proj", - "fc2": "net.2", - "input_embedder": "proj_in", - } - - TRANSFORMER_SPECIAL_KEYS_REMAP = { - "txt_in": remap_txt_in_, - "img_attn_qkv": remap_img_attn_qkv_, - "txt_attn_qkv": remap_txt_attn_qkv_, - "single_blocks": remap_single_transformer_blocks_, - "final_layer.adaLN_modulation.1": remap_norm_scale_shift_, - } - - def update_state_dict_(state_dict, old_key, new_key): - state_dict[new_key] = state_dict.pop(old_key) - - 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) - update_state_dict_(checkpoint, key, new_key) - - for key in list(checkpoint.keys()): - for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, checkpoint) - - return checkpoint - - -def convert_auraflow_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {} - state_dict_keys = list(checkpoint.keys()) - - # Handle register tokens and positional embeddings - converted_state_dict["register_tokens"] = checkpoint.pop("register_tokens", None) - - # Handle time step projection - converted_state_dict["time_step_proj.linear_1.weight"] = checkpoint.pop("t_embedder.mlp.0.weight", None) - converted_state_dict["time_step_proj.linear_1.bias"] = checkpoint.pop("t_embedder.mlp.0.bias", None) - converted_state_dict["time_step_proj.linear_2.weight"] = checkpoint.pop("t_embedder.mlp.2.weight", None) - converted_state_dict["time_step_proj.linear_2.bias"] = checkpoint.pop("t_embedder.mlp.2.bias", None) - - # Handle context embedder - converted_state_dict["context_embedder.weight"] = checkpoint.pop("cond_seq_linear.weight", None) - - # Calculate the number of layers - def calculate_layers(keys, key_prefix): - layers = set() - for k in keys: - if key_prefix in k: - layer_num = int(k.split(".")[1]) # get the layer number - layers.add(layer_num) - return len(layers) - - mmdit_layers = calculate_layers(state_dict_keys, key_prefix="double_layers") - single_dit_layers = calculate_layers(state_dict_keys, key_prefix="single_layers") - - # MMDiT blocks - for i in range(mmdit_layers): - # Feed-forward - path_mapping = {"mlpX": "ff", "mlpC": "ff_context"} - weight_mapping = {"c_fc1": "linear_1", "c_fc2": "linear_2", "c_proj": "out_projection"} - for orig_k, diffuser_k in path_mapping.items(): - for k, v in weight_mapping.items(): - converted_state_dict[f"joint_transformer_blocks.{i}.{diffuser_k}.{v}.weight"] = checkpoint.pop( - f"double_layers.{i}.{orig_k}.{k}.weight", None - ) - - # Norms - path_mapping = {"modX": "norm1", "modC": "norm1_context"} - for orig_k, diffuser_k in path_mapping.items(): - converted_state_dict[f"joint_transformer_blocks.{i}.{diffuser_k}.linear.weight"] = checkpoint.pop( - f"double_layers.{i}.{orig_k}.1.weight", None - ) - - # Attentions - x_attn_mapping = {"w2q": "to_q", "w2k": "to_k", "w2v": "to_v", "w2o": "to_out.0"} - context_attn_mapping = {"w1q": "add_q_proj", "w1k": "add_k_proj", "w1v": "add_v_proj", "w1o": "to_add_out"} - for attn_mapping in [x_attn_mapping, context_attn_mapping]: - for k, v in attn_mapping.items(): - converted_state_dict[f"joint_transformer_blocks.{i}.attn.{v}.weight"] = checkpoint.pop( - f"double_layers.{i}.attn.{k}.weight", None - ) - - # Single-DiT blocks - for i in range(single_dit_layers): - # Feed-forward - mapping = {"c_fc1": "linear_1", "c_fc2": "linear_2", "c_proj": "out_projection"} - for k, v in mapping.items(): - converted_state_dict[f"single_transformer_blocks.{i}.ff.{v}.weight"] = checkpoint.pop( - f"single_layers.{i}.mlp.{k}.weight", None - ) - - # Norms - converted_state_dict[f"single_transformer_blocks.{i}.norm1.linear.weight"] = checkpoint.pop( - f"single_layers.{i}.modCX.1.weight", None - ) - - # Attentions - x_attn_mapping = {"w1q": "to_q", "w1k": "to_k", "w1v": "to_v", "w1o": "to_out.0"} - for k, v in x_attn_mapping.items(): - converted_state_dict[f"single_transformer_blocks.{i}.attn.{v}.weight"] = checkpoint.pop( - f"single_layers.{i}.attn.{k}.weight", None - ) - # Final blocks - converted_state_dict["proj_out.weight"] = checkpoint.pop("final_linear.weight", None) - - # Handle the final norm layer - norm_weight = checkpoint.pop("modF.1.weight", None) - if norm_weight is not None: - converted_state_dict["norm_out.linear.weight"] = swap_scale_shift(norm_weight, dim=None) - else: - converted_state_dict["norm_out.linear.weight"] = None - - converted_state_dict["pos_embed.pos_embed"] = checkpoint.pop("positional_encoding") - converted_state_dict["pos_embed.proj.weight"] = checkpoint.pop("init_x_linear.weight") - converted_state_dict["pos_embed.proj.bias"] = checkpoint.pop("init_x_linear.bias") - - return converted_state_dict - - -def convert_lumina2_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {} - - # Original Lumina-Image-2 has an extra norm parameter that is unused - # We just remove it here - checkpoint.pop("norm_final.weight", None) - - # Comfy checkpoints add this prefix - keys = list(checkpoint.keys()) - for k in keys: - if "model.diffusion_model." in k: - checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) - - LUMINA_KEY_MAP = { - "cap_embedder": "time_caption_embed.caption_embedder", - "t_embedder.mlp.0": "time_caption_embed.timestep_embedder.linear_1", - "t_embedder.mlp.2": "time_caption_embed.timestep_embedder.linear_2", - "attention": "attn", - ".out.": ".to_out.0.", - "k_norm": "norm_k", - "q_norm": "norm_q", - "w1": "linear_1", - "w2": "linear_2", - "w3": "linear_3", - "adaLN_modulation.1": "norm1.linear", - } - ATTENTION_NORM_MAP = { - "attention_norm1": "norm1.norm", - "attention_norm2": "norm2", - } - CONTEXT_REFINER_MAP = { - "context_refiner.0.attention_norm1": "context_refiner.0.norm1", - "context_refiner.0.attention_norm2": "context_refiner.0.norm2", - "context_refiner.1.attention_norm1": "context_refiner.1.norm1", - "context_refiner.1.attention_norm2": "context_refiner.1.norm2", - } - FINAL_LAYER_MAP = { - "final_layer.adaLN_modulation.1": "norm_out.linear_1", - "final_layer.linear": "norm_out.linear_2", - } - - def convert_lumina_attn_to_diffusers(tensor, diffusers_key): - q_dim = 2304 - k_dim = v_dim = 768 - - to_q, to_k, to_v = torch.split(tensor, [q_dim, k_dim, v_dim], dim=0) - - return { - diffusers_key.replace("qkv", "to_q"): to_q, - diffusers_key.replace("qkv", "to_k"): to_k, - diffusers_key.replace("qkv", "to_v"): to_v, - } - - for key in keys: - diffusers_key = key - for k, v in CONTEXT_REFINER_MAP.items(): - diffusers_key = diffusers_key.replace(k, v) - for k, v in FINAL_LAYER_MAP.items(): - diffusers_key = diffusers_key.replace(k, v) - for k, v in ATTENTION_NORM_MAP.items(): - diffusers_key = diffusers_key.replace(k, v) - for k, v in LUMINA_KEY_MAP.items(): - diffusers_key = diffusers_key.replace(k, v) - - if "qkv" in diffusers_key: - converted_state_dict.update(convert_lumina_attn_to_diffusers(checkpoint.pop(key), diffusers_key)) - else: - converted_state_dict[diffusers_key] = checkpoint.pop(key) - - return converted_state_dict - - -def convert_sana_transformer_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 "blocks" in k))[-1] + 1 # noqa: C401 - - # Positional and patch embeddings. - checkpoint.pop("pos_embed") - converted_state_dict["patch_embed.proj.weight"] = checkpoint.pop("x_embedder.proj.weight") - converted_state_dict["patch_embed.proj.bias"] = checkpoint.pop("x_embedder.proj.bias") - - # Timestep embeddings. - converted_state_dict["time_embed.emb.timestep_embedder.linear_1.weight"] = checkpoint.pop( - "t_embedder.mlp.0.weight" - ) - converted_state_dict["time_embed.emb.timestep_embedder.linear_1.bias"] = checkpoint.pop("t_embedder.mlp.0.bias") - converted_state_dict["time_embed.emb.timestep_embedder.linear_2.weight"] = checkpoint.pop( - "t_embedder.mlp.2.weight" - ) - converted_state_dict["time_embed.emb.timestep_embedder.linear_2.bias"] = checkpoint.pop("t_embedder.mlp.2.bias") - converted_state_dict["time_embed.linear.weight"] = checkpoint.pop("t_block.1.weight") - converted_state_dict["time_embed.linear.bias"] = checkpoint.pop("t_block.1.bias") - - # Caption Projection. - checkpoint.pop("y_embedder.y_embedding") - converted_state_dict["caption_projection.linear_1.weight"] = checkpoint.pop("y_embedder.y_proj.fc1.weight") - converted_state_dict["caption_projection.linear_1.bias"] = checkpoint.pop("y_embedder.y_proj.fc1.bias") - converted_state_dict["caption_projection.linear_2.weight"] = checkpoint.pop("y_embedder.y_proj.fc2.weight") - converted_state_dict["caption_projection.linear_2.bias"] = checkpoint.pop("y_embedder.y_proj.fc2.bias") - converted_state_dict["caption_norm.weight"] = checkpoint.pop("attention_y_norm.weight") - - for i in range(num_layers): - converted_state_dict[f"transformer_blocks.{i}.scale_shift_table"] = checkpoint.pop( - f"blocks.{i}.scale_shift_table" - ) - - # Self-Attention - sample_q, sample_k, sample_v = torch.chunk(checkpoint.pop(f"blocks.{i}.attn.qkv.weight"), 3, dim=0) - converted_state_dict[f"transformer_blocks.{i}.attn1.to_q.weight"] = torch.cat([sample_q]) - converted_state_dict[f"transformer_blocks.{i}.attn1.to_k.weight"] = torch.cat([sample_k]) - converted_state_dict[f"transformer_blocks.{i}.attn1.to_v.weight"] = torch.cat([sample_v]) - - # Output Projections - converted_state_dict[f"transformer_blocks.{i}.attn1.to_out.0.weight"] = checkpoint.pop( - f"blocks.{i}.attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn1.to_out.0.bias"] = checkpoint.pop( - f"blocks.{i}.attn.proj.bias" - ) - - # Cross-Attention - converted_state_dict[f"transformer_blocks.{i}.attn2.to_q.weight"] = checkpoint.pop( - f"blocks.{i}.cross_attn.q_linear.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_q.bias"] = checkpoint.pop( - f"blocks.{i}.cross_attn.q_linear.bias" - ) - - linear_sample_k, linear_sample_v = torch.chunk( - checkpoint.pop(f"blocks.{i}.cross_attn.kv_linear.weight"), 2, dim=0 - ) - linear_sample_k_bias, linear_sample_v_bias = torch.chunk( - checkpoint.pop(f"blocks.{i}.cross_attn.kv_linear.bias"), 2, dim=0 - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_k.weight"] = linear_sample_k - converted_state_dict[f"transformer_blocks.{i}.attn2.to_v.weight"] = linear_sample_v - converted_state_dict[f"transformer_blocks.{i}.attn2.to_k.bias"] = linear_sample_k_bias - converted_state_dict[f"transformer_blocks.{i}.attn2.to_v.bias"] = linear_sample_v_bias - - # Output Projections - converted_state_dict[f"transformer_blocks.{i}.attn2.to_out.0.weight"] = checkpoint.pop( - f"blocks.{i}.cross_attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_out.0.bias"] = checkpoint.pop( - f"blocks.{i}.cross_attn.proj.bias" - ) - - # MLP - converted_state_dict[f"transformer_blocks.{i}.ff.conv_inverted.weight"] = checkpoint.pop( - f"blocks.{i}.mlp.inverted_conv.conv.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.conv_inverted.bias"] = checkpoint.pop( - f"blocks.{i}.mlp.inverted_conv.conv.bias" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.conv_depth.weight"] = checkpoint.pop( - f"blocks.{i}.mlp.depth_conv.conv.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.conv_depth.bias"] = checkpoint.pop( - f"blocks.{i}.mlp.depth_conv.conv.bias" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.conv_point.weight"] = checkpoint.pop( - f"blocks.{i}.mlp.point_conv.conv.weight" - ) - - # Final layer - converted_state_dict["proj_out.weight"] = checkpoint.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = checkpoint.pop("final_layer.linear.bias") - converted_state_dict["scale_shift_table"] = checkpoint.pop("final_layer.scale_shift_table") - - return converted_state_dict - - -def convert_wan_transformer_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) - - TRANSFORMER_KEYS_RENAME_DICT = { - "time_embedding.0": "condition_embedder.time_embedder.linear_1", - "time_embedding.2": "condition_embedder.time_embedder.linear_2", - "text_embedding.0": "condition_embedder.text_embedder.linear_1", - "text_embedding.2": "condition_embedder.text_embedder.linear_2", - "time_projection.1": "condition_embedder.time_proj", - "cross_attn": "attn2", - "self_attn": "attn1", - ".o.": ".to_out.0.", - ".q.": ".to_q.", - ".k.": ".to_k.", - ".v.": ".to_v.", - ".k_img.": ".add_k_proj.", - ".v_img.": ".add_v_proj.", - ".norm_k_img.": ".norm_added_k.", - "head.modulation": "scale_shift_table", - "head.head": "proj_out", - "modulation": "scale_shift_table", - "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 - "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 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) - - return converted_state_dict - - -def convert_wan_vae_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {} - - # Create mappings for specific components - middle_key_mapping = { - # Encoder middle block - "encoder.middle.0.residual.0.gamma": "encoder.mid_block.resnets.0.norm1.gamma", - "encoder.middle.0.residual.2.bias": "encoder.mid_block.resnets.0.conv1.bias", - "encoder.middle.0.residual.2.weight": "encoder.mid_block.resnets.0.conv1.weight", - "encoder.middle.0.residual.3.gamma": "encoder.mid_block.resnets.0.norm2.gamma", - "encoder.middle.0.residual.6.bias": "encoder.mid_block.resnets.0.conv2.bias", - "encoder.middle.0.residual.6.weight": "encoder.mid_block.resnets.0.conv2.weight", - "encoder.middle.2.residual.0.gamma": "encoder.mid_block.resnets.1.norm1.gamma", - "encoder.middle.2.residual.2.bias": "encoder.mid_block.resnets.1.conv1.bias", - "encoder.middle.2.residual.2.weight": "encoder.mid_block.resnets.1.conv1.weight", - "encoder.middle.2.residual.3.gamma": "encoder.mid_block.resnets.1.norm2.gamma", - "encoder.middle.2.residual.6.bias": "encoder.mid_block.resnets.1.conv2.bias", - "encoder.middle.2.residual.6.weight": "encoder.mid_block.resnets.1.conv2.weight", - # Decoder middle block - "decoder.middle.0.residual.0.gamma": "decoder.mid_block.resnets.0.norm1.gamma", - "decoder.middle.0.residual.2.bias": "decoder.mid_block.resnets.0.conv1.bias", - "decoder.middle.0.residual.2.weight": "decoder.mid_block.resnets.0.conv1.weight", - "decoder.middle.0.residual.3.gamma": "decoder.mid_block.resnets.0.norm2.gamma", - "decoder.middle.0.residual.6.bias": "decoder.mid_block.resnets.0.conv2.bias", - "decoder.middle.0.residual.6.weight": "decoder.mid_block.resnets.0.conv2.weight", - "decoder.middle.2.residual.0.gamma": "decoder.mid_block.resnets.1.norm1.gamma", - "decoder.middle.2.residual.2.bias": "decoder.mid_block.resnets.1.conv1.bias", - "decoder.middle.2.residual.2.weight": "decoder.mid_block.resnets.1.conv1.weight", - "decoder.middle.2.residual.3.gamma": "decoder.mid_block.resnets.1.norm2.gamma", - "decoder.middle.2.residual.6.bias": "decoder.mid_block.resnets.1.conv2.bias", - "decoder.middle.2.residual.6.weight": "decoder.mid_block.resnets.1.conv2.weight", - } - - # Create a mapping for attention blocks - attention_mapping = { - # Encoder middle attention - "encoder.middle.1.norm.gamma": "encoder.mid_block.attentions.0.norm.gamma", - "encoder.middle.1.to_qkv.weight": "encoder.mid_block.attentions.0.to_qkv.weight", - "encoder.middle.1.to_qkv.bias": "encoder.mid_block.attentions.0.to_qkv.bias", - "encoder.middle.1.proj.weight": "encoder.mid_block.attentions.0.proj.weight", - "encoder.middle.1.proj.bias": "encoder.mid_block.attentions.0.proj.bias", - # Decoder middle attention - "decoder.middle.1.norm.gamma": "decoder.mid_block.attentions.0.norm.gamma", - "decoder.middle.1.to_qkv.weight": "decoder.mid_block.attentions.0.to_qkv.weight", - "decoder.middle.1.to_qkv.bias": "decoder.mid_block.attentions.0.to_qkv.bias", - "decoder.middle.1.proj.weight": "decoder.mid_block.attentions.0.proj.weight", - "decoder.middle.1.proj.bias": "decoder.mid_block.attentions.0.proj.bias", - } - - # Create a mapping for the head components - head_mapping = { - # Encoder head - "encoder.head.0.gamma": "encoder.norm_out.gamma", - "encoder.head.2.bias": "encoder.conv_out.bias", - "encoder.head.2.weight": "encoder.conv_out.weight", - # Decoder head - "decoder.head.0.gamma": "decoder.norm_out.gamma", - "decoder.head.2.bias": "decoder.conv_out.bias", - "decoder.head.2.weight": "decoder.conv_out.weight", - } - - # Create a mapping for the quant components - quant_mapping = { - "conv1.weight": "quant_conv.weight", - "conv1.bias": "quant_conv.bias", - "conv2.weight": "post_quant_conv.weight", - "conv2.bias": "post_quant_conv.bias", - } - - # Process each key in the state dict - for key, value in checkpoint.items(): - # Handle middle block keys using the mapping - if key in middle_key_mapping: - new_key = middle_key_mapping[key] - converted_state_dict[new_key] = value - # Handle attention blocks using the mapping - elif key in attention_mapping: - new_key = attention_mapping[key] - converted_state_dict[new_key] = value - # Handle head keys using the mapping - elif key in head_mapping: - new_key = head_mapping[key] - converted_state_dict[new_key] = value - # Handle quant keys using the mapping - elif key in quant_mapping: - new_key = quant_mapping[key] - converted_state_dict[new_key] = value - # Handle encoder conv1 - elif key == "encoder.conv1.weight": - converted_state_dict["encoder.conv_in.weight"] = value - elif key == "encoder.conv1.bias": - converted_state_dict["encoder.conv_in.bias"] = value - # Handle decoder conv1 - elif key == "decoder.conv1.weight": - converted_state_dict["decoder.conv_in.weight"] = value - elif key == "decoder.conv1.bias": - converted_state_dict["decoder.conv_in.bias"] = value - # Handle encoder downsamples - elif key.startswith("encoder.downsamples."): - # Convert to down_blocks - new_key = key.replace("encoder.downsamples.", "encoder.down_blocks.") - - # Convert residual block naming but keep the original structure - if ".residual.0.gamma" in new_key: - new_key = new_key.replace(".residual.0.gamma", ".norm1.gamma") - elif ".residual.2.bias" in new_key: - new_key = new_key.replace(".residual.2.bias", ".conv1.bias") - elif ".residual.2.weight" in new_key: - new_key = new_key.replace(".residual.2.weight", ".conv1.weight") - elif ".residual.3.gamma" in new_key: - new_key = new_key.replace(".residual.3.gamma", ".norm2.gamma") - elif ".residual.6.bias" in new_key: - new_key = new_key.replace(".residual.6.bias", ".conv2.bias") - elif ".residual.6.weight" in new_key: - new_key = new_key.replace(".residual.6.weight", ".conv2.weight") - elif ".shortcut.bias" in new_key: - new_key = new_key.replace(".shortcut.bias", ".conv_shortcut.bias") - elif ".shortcut.weight" in new_key: - new_key = new_key.replace(".shortcut.weight", ".conv_shortcut.weight") - - converted_state_dict[new_key] = value - - # Handle decoder upsamples - elif key.startswith("decoder.upsamples."): - # Convert to up_blocks - parts = key.split(".") - block_idx = int(parts[2]) - - # Group residual blocks - if "residual" in key: - if block_idx in [0, 1, 2]: - new_block_idx = 0 - resnet_idx = block_idx - elif block_idx in [4, 5, 6]: - new_block_idx = 1 - resnet_idx = block_idx - 4 - elif block_idx in [8, 9, 10]: - new_block_idx = 2 - resnet_idx = block_idx - 8 - elif block_idx in [12, 13, 14]: - new_block_idx = 3 - resnet_idx = block_idx - 12 - else: - # Keep as is for other blocks - converted_state_dict[key] = value - continue - - # Convert residual block naming - if ".residual.0.gamma" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.norm1.gamma" - elif ".residual.2.bias" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv1.bias" - elif ".residual.2.weight" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv1.weight" - elif ".residual.3.gamma" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.norm2.gamma" - elif ".residual.6.bias" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv2.bias" - elif ".residual.6.weight" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv2.weight" - else: - new_key = key - - converted_state_dict[new_key] = value - - # Handle shortcut connections - elif ".shortcut." in key: - if block_idx == 4: - new_key = key.replace(".shortcut.", ".resnets.0.conv_shortcut.") - new_key = new_key.replace("decoder.upsamples.4", "decoder.up_blocks.1") - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - new_key = new_key.replace(".shortcut.", ".conv_shortcut.") - - converted_state_dict[new_key] = value - - # Handle upsamplers - elif ".resample." in key or ".time_conv." in key: - if block_idx == 3: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.0.upsamplers.0") - elif block_idx == 7: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.1.upsamplers.0") - elif block_idx == 11: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.2.upsamplers.0") - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - - converted_state_dict[new_key] = value - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - converted_state_dict[new_key] = value - else: - # Keep other keys unchanged - converted_state_dict[key] = value - - return converted_state_dict - - -def convert_hidream_transformer_to_diffusers(checkpoint, **kwargs): - keys = list(checkpoint.keys()) - for k in keys: - if "model.diffusion_model." in k: - checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) - - return checkpoint - - -def convert_chroma_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 - ) - 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") - - # x_embedder - converted_state_dict["x_embedder.weight"] = checkpoint.pop("img_in.weight") - converted_state_dict["x_embedder.bias"] = checkpoint.pop("img_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") - - converted_state_dict["proj_out.weight"] = checkpoint.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = checkpoint.pop("final_layer.linear.bias") - - return converted_state_dict diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/textual_inversion.py b/venv/lib/python3.10/site-packages/diffusers/loaders/textual_inversion.py deleted file mode 100644 index 63fc97ed431f3689e5cd3f14218894c587f61f63..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/textual_inversion.py +++ /dev/null @@ -1,581 +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 Dict, List, Optional, Union - -import safetensors -import torch -from huggingface_hub.utils import validate_hf_hub_args -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 - - -if is_transformers_available(): - from transformers import PreTrainedModel, PreTrainedTokenizer - -if is_accelerate_available(): - from accelerate.hooks import AlignDevicesHook, CpuOffload, remove_hook_from_module - -logger = logging.get_logger(__name__) - -TEXT_INVERSION_NAME = "learned_embeds.bin" -TEXT_INVERSION_NAME_SAFE = "learned_embeds.safetensors" - - -@validate_hf_hub_args -def load_textual_inversion_state_dicts(pretrained_model_name_or_paths, **kwargs): - 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) - hf_token = kwargs.pop("hf_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) - - allow_pickle = False - if use_safetensors is None: - use_safetensors = True - allow_pickle = True - - user_agent = { - "file_type": "text_inversion", - "framework": "pytorch", - } - state_dicts = [] - for pretrained_model_name_or_path in pretrained_model_name_or_paths: - if not isinstance(pretrained_model_name_or_path, (dict, torch.Tensor)): - # 3.1. Load textual inversion file - model_file = None - - # Let's first try to load .safetensors weights - if (use_safetensors and weight_name is None) or ( - weight_name is not None and weight_name.endswith(".safetensors") - ): - try: - model_file = _get_model_file( - pretrained_model_name_or_path, - weights_name=weight_name or TEXT_INVERSION_NAME_SAFE, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=hf_token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - ) - state_dict = safetensors.torch.load_file(model_file, device="cpu") - except Exception as e: - if not allow_pickle: - raise e - - model_file = None - - if model_file is None: - model_file = _get_model_file( - pretrained_model_name_or_path, - weights_name=weight_name or TEXT_INVERSION_NAME, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=hf_token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - ) - state_dict = load_state_dict(model_file) - else: - state_dict = pretrained_model_name_or_path - - state_dicts.append(state_dict) - - return state_dicts - - -class TextualInversionLoaderMixin: - r""" - 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 - 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 - inversion token or if the textual inversion token is a single vector, the input prompt is returned. - - Parameters: - prompt (`str` or list of `str`): - The prompt or prompts to guide the image generation. - tokenizer (`PreTrainedTokenizer`): - The tokenizer responsible for encoding the prompt into input tokens. - - Returns: - `str` or list of `str`: The converted prompt - """ - 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): - return prompts[0] - - return prompts - - def _maybe_convert_prompt(self, prompt: str, tokenizer: "PreTrainedTokenizer"): # noqa: F821 - r""" - Maybe convert a prompt into a "multi vector"-compatible prompt. If the prompt includes a token that corresponds - to a multi-vector textual inversion embedding, this function will process the prompt so that the special token - is replaced with multiple special tokens each corresponding to one of the vectors. If the prompt has no textual - inversion token or a textual inversion token that is a single vector, the input prompt is simply returned. - - Parameters: - prompt (`str`): - The prompt to guide the image generation. - tokenizer (`PreTrainedTokenizer`): - The tokenizer responsible for encoding the prompt into input tokens. - - Returns: - `str`: The converted prompt - """ - tokens = tokenizer.tokenize(prompt) - unique_tokens = set(tokens) - for token in unique_tokens: - if token in tokenizer.added_tokens_encoder: - replacement = token - i = 1 - while f"{token}_{i}" in tokenizer.added_tokens_encoder: - replacement += f" {token}_{i}" - i += 1 - - prompt = prompt.replace(token, replacement) - - return prompt - - def _check_text_inv_inputs(self, tokenizer, text_encoder, pretrained_model_name_or_paths, tokens): - if tokenizer is None: - raise ValueError( - f"{self.__class__.__name__} requires `self.tokenizer` or passing a `tokenizer` of type `PreTrainedTokenizer` for calling" - f" `{self.load_textual_inversion.__name__}`" - ) - - if text_encoder is None: - raise ValueError( - f"{self.__class__.__name__} requires `self.text_encoder` or passing a `text_encoder` of type `PreTrainedModel` for calling" - f" `{self.load_textual_inversion.__name__}`" - ) - - if len(pretrained_model_name_or_paths) > 1 and len(pretrained_model_name_or_paths) != len(tokens): - raise ValueError( - f"You have passed a list of models of length {len(pretrained_model_name_or_paths)}, and list of tokens of length {len(tokens)} " - f"Make sure both lists have the same length." - ) - - valid_tokens = [t for t in tokens if t is not None] - if len(set(valid_tokens)) < len(valid_tokens): - raise ValueError(f"You have passed a list of tokens that contains duplicates: {tokens}") - - @staticmethod - def _retrieve_tokens_and_embeddings(tokens, state_dicts, tokenizer): - all_tokens = [] - all_embeddings = [] - for state_dict, token in zip(state_dicts, tokens): - if isinstance(state_dict, torch.Tensor): - if token is None: - raise ValueError( - "You are trying to load a textual inversion embedding that has been saved as a PyTorch tensor. Make sure to pass the name of the corresponding token in this case: `token=...`." - ) - loaded_token = token - embedding = state_dict - elif len(state_dict) == 1: - # diffusers - loaded_token, embedding = next(iter(state_dict.items())) - elif "string_to_param" in state_dict: - # A1111 - loaded_token = state_dict["name"] - embedding = state_dict["string_to_param"]["*"] - else: - raise ValueError( - f"Loaded state dictionary is incorrect: {state_dict}. \n\n" - "Please verify that the loaded state dictionary of the textual embedding either only has a single key or includes the `string_to_param`" - " input key." - ) - - if token is not None and loaded_token != token: - logger.info(f"The loaded token: {loaded_token} is overwritten by the passed token {token}.") - else: - token = loaded_token - - if token in tokenizer.get_vocab(): - raise ValueError( - f"Token {token} already in tokenizer vocabulary. Please choose a different token name or remove {token} and embedding from the tokenizer and text encoder." - ) - - all_tokens.append(token) - all_embeddings.append(embedding) - - return all_tokens, all_embeddings - - @staticmethod - def _extend_tokens_and_embeddings(tokens, embeddings, tokenizer): - all_tokens = [] - all_embeddings = [] - - for embedding, token in zip(embeddings, tokens): - if f"{token}_1" in tokenizer.get_vocab(): - multi_vector_tokens = [token] - i = 1 - while f"{token}_{i}" in tokenizer.added_tokens_encoder: - multi_vector_tokens.append(f"{token}_{i}") - i += 1 - - raise ValueError( - f"Multi-vector Token {multi_vector_tokens} already in tokenizer vocabulary. Please choose a different token name or remove the {multi_vector_tokens} and embedding from the tokenizer and text encoder." - ) - - is_multi_vector = len(embedding.shape) > 1 and embedding.shape[0] > 1 - if is_multi_vector: - all_tokens += [token] + [f"{token}_{i}" for i in range(1, embedding.shape[0])] - all_embeddings += [e for e in embedding] # noqa: C416 - else: - all_tokens += [token] - all_embeddings += [embedding[0]] if len(embedding.shape) > 1 else [embedding] - - return all_tokens, all_embeddings - - @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 - **kwargs, - ): - r""" - Load Textual Inversion embeddings into the text encoder of [`StableDiffusionPipeline`] (both 🤗 Diffusers and - 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]`): - 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 - pretrained model hosted on the Hub. - - A path to a *directory* (for example `./my_text_inversion_directory/`) containing the textual - inversion weights. - - A path to a *file* (for example `./my_text_inversions.pt`) containing textual inversion weights. - - A [torch state - dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - - 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*): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - If not specified, function will take self.tokenizer. - tokenizer ([`~transformers.CLIPTokenizer`], *optional*): - A `CLIPTokenizer` to tokenize text. If not specified, function will take self.tokenizer. - weight_name (`str`, *optional*): - Name of a custom weight file. This should be used when: - - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - hf_token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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. - mirror (`str`, *optional*): - 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. - - Example: - - To load a Textual Inversion embedding vector in 🤗 Diffusers format: - - ```py - from diffusers import StableDiffusionPipeline - import torch - - model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" - pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda") - - pipe.load_textual_inversion("sd-concepts-library/cat-toy") - - prompt = "A backpack" - - image = pipe(prompt, num_inference_steps=50).images[0] - image.save("cat-backpack.png") - ``` - - To load a Textual Inversion embedding vector in Automatic1111 format, make sure to download the vector first - (for example from [civitAI](https://civitai.com/models/3036?modelVersionId=9857)) and then load the vector - locally: - - ```py - from diffusers import StableDiffusionPipeline - import torch - - model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" - pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda") - - pipe.load_textual_inversion("./charturnerv2.pt", token="charturnerv2") - - prompt = "charturnerv2, multiple views of the same character in the same outfit, a character turnaround of a woman wearing a black jacket and red shirt, best quality, intricate details." - - image = pipe(prompt, num_inference_steps=50).images[0] - image.save("character.png") - ``` - - """ - # 1. Set correct tokenizer and text encoder - tokenizer = tokenizer or getattr(self, "tokenizer", None) - text_encoder = text_encoder or getattr(self, "text_encoder", None) - - # 2. Normalize inputs - pretrained_model_name_or_paths = ( - [pretrained_model_name_or_path] - if not isinstance(pretrained_model_name_or_path, list) - else pretrained_model_name_or_path - ) - tokens = [token] if not isinstance(token, list) else token - if tokens[0] is None: - tokens = tokens * len(pretrained_model_name_or_paths) - - # 3. Check inputs - self._check_text_inv_inputs(tokenizer, text_encoder, pretrained_model_name_or_paths, tokens) - - # 4. Load state dicts of textual embeddings - state_dicts = load_textual_inversion_state_dicts(pretrained_model_name_or_paths, **kwargs) - - # 4.1 Handle the special case when state_dict is a tensor that contains n embeddings for n tokens - if len(tokens) > 1 and len(state_dicts) == 1: - if isinstance(state_dicts[0], torch.Tensor): - state_dicts = list(state_dicts[0]) - if len(tokens) != len(state_dicts): - raise ValueError( - f"You have passed a state_dict contains {len(state_dicts)} embeddings, and list of tokens of length {len(tokens)} " - f"Make sure both have the same length." - ) - - # 4. Retrieve tokens and embeddings - tokens, embeddings = self._retrieve_tokens_and_embeddings(tokens, state_dicts, tokenizer) - - # 5. Extend tokens and embeddings for multi vector - tokens, embeddings = self._extend_tokens_and_embeddings(tokens, embeddings, tokenizer) - - # 6. Make sure all embeddings have the correct size - expected_emb_dim = text_encoder.get_input_embeddings().weight.shape[-1] - if any(expected_emb_dim != emb.shape[-1] for emb in embeddings): - raise ValueError( - "Loaded embeddings are of incorrect shape. Expected each textual inversion embedding " - "to be of shape {input_embeddings.shape[-1]}, but are {embeddings.shape[-1]} " - ) - - # 7. Now we can be sure that loading the embedding matrix works - # < Unsafe code: - - # 7.1 Offload all hooks in case the pipeline was cpu offloaded before make sure, we offload and onload again - is_model_cpu_offload = False - is_sequential_cpu_offload = False - if self.hf_device_map is None: - for _, component in self.components.items(): - if isinstance(component, nn.Module): - if hasattr(component, "_hf_hook"): - is_model_cpu_offload = isinstance(getattr(component, "_hf_hook"), CpuOffload) - is_sequential_cpu_offload = ( - isinstance(getattr(component, "_hf_hook"), AlignDevicesHook) - or hasattr(component._hf_hook, "hooks") - and isinstance(component._hf_hook.hooks[0], AlignDevicesHook) - ) - logger.info( - "Accelerate hooks detected. Since you have called `load_textual_inversion()`, the previous hooks will be first removed. Then the textual inversion parameters will be loaded and the hooks will be applied again." - ) - if is_sequential_cpu_offload or is_model_cpu_offload: - remove_hook_from_module(component, recurse=is_sequential_cpu_offload) - - # 7.2 save expected device and dtype - device = text_encoder.device - dtype = text_encoder.dtype - - # 7.3 Increase token embedding matrix - text_encoder.resize_token_embeddings(len(tokenizer) + len(tokens)) - input_embeddings = text_encoder.get_input_embeddings().weight - - # 7.4 Load token and embedding - for token, embedding in zip(tokens, embeddings): - # add tokens and get ids - tokenizer.add_tokens(token) - token_id = tokenizer.convert_tokens_to_ids(token) - input_embeddings.data[token_id] = embedding - logger.info(f"Loaded textual inversion embedding for {token}.") - - input_embeddings.to(dtype=dtype, device=device) - - # 7.5 Offload the model again - if is_model_cpu_offload: - self.enable_model_cpu_offload(device=device) - elif is_sequential_cpu_offload: - self.enable_sequential_cpu_offload(device=device) - - # / Unsafe Code > - - def unload_textual_inversion( - self, - tokens: Optional[Union[str, List[str]]] = None, - tokenizer: Optional["PreTrainedTokenizer"] = None, - text_encoder: Optional["PreTrainedModel"] = None, - ): - r""" - Unload Textual Inversion embeddings from the text encoder of [`StableDiffusionPipeline`] - - Example: - ```py - from diffusers import AutoPipelineForText2Image - import torch - - pipeline = AutoPipelineForText2Image.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") - - # Example 1 - pipeline.load_textual_inversion("sd-concepts-library/gta5-artwork") - pipeline.load_textual_inversion("sd-concepts-library/moeb-style") - - # Remove all token embeddings - pipeline.unload_textual_inversion() - - # Example 2 - pipeline.load_textual_inversion("sd-concepts-library/moeb-style") - pipeline.load_textual_inversion("sd-concepts-library/gta5-artwork") - - # Remove just one token - pipeline.unload_textual_inversion("") - - # Example 3: unload from SDXL - pipeline = AutoPipelineForText2Image.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0") - embedding_path = hf_hub_download( - repo_id="linoyts/web_y2k", filename="web_y2k_emb.safetensors", repo_type="model" - ) - - # load embeddings to the text encoders - state_dict = load_file(embedding_path) - - # load embeddings of text_encoder 1 (CLIP ViT-L/14) - pipeline.load_textual_inversion( - state_dict["clip_l"], - tokens=["", ""], - text_encoder=pipeline.text_encoder, - tokenizer=pipeline.tokenizer, - ) - # load embeddings of text_encoder 2 (CLIP ViT-G/14) - pipeline.load_textual_inversion( - state_dict["clip_g"], - tokens=["", ""], - text_encoder=pipeline.text_encoder_2, - tokenizer=pipeline.tokenizer_2, - ) - - # Unload explicitly from both text encoders and tokenizers - pipeline.unload_textual_inversion( - tokens=["", ""], text_encoder=pipeline.text_encoder, tokenizer=pipeline.tokenizer - ) - pipeline.unload_textual_inversion( - tokens=["", ""], text_encoder=pipeline.text_encoder_2, tokenizer=pipeline.tokenizer_2 - ) - ``` - """ - - tokenizer = tokenizer or getattr(self, "tokenizer", None) - text_encoder = text_encoder or getattr(self, "text_encoder", None) - - # Get textual inversion tokens and ids - token_ids = [] - last_special_token_id = None - - if tokens: - if isinstance(tokens, str): - tokens = [tokens] - for added_token_id, added_token in tokenizer.added_tokens_decoder.items(): - if not added_token.special: - if added_token.content in tokens: - token_ids.append(added_token_id) - else: - last_special_token_id = added_token_id - if len(token_ids) == 0: - raise ValueError("No tokens to remove found") - else: - tokens = [] - for added_token_id, added_token in tokenizer.added_tokens_decoder.items(): - if not added_token.special: - token_ids.append(added_token_id) - tokens.append(added_token.content) - 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 - 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() - - # Delete from text encoder - text_embedding_dim = text_encoder.get_input_embeddings().embedding_dim - temp_text_embedding_weights = text_encoder.get_input_embeddings().weight - text_embedding_weights = temp_text_embedding_weights[: last_special_token_id + 1] - to_append = [] - for i in range(last_special_token_id + 1, temp_text_embedding_weights.shape[0]): - if i not in token_ids: - to_append.append(temp_text_embedding_weights[i].unsqueeze(0)) - if len(to_append) > 0: - to_append = torch.cat(to_append, dim=0) - text_embedding_weights = torch.cat([text_embedding_weights, to_append], dim=0) - text_embeddings_filtered = nn.Embedding(text_embedding_weights.shape[0], text_embedding_dim) - text_embeddings_filtered.weight.data = text_embedding_weights - text_encoder.set_input_embeddings(text_embeddings_filtered) diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/transformer_flux.py b/venv/lib/python3.10/site-packages/diffusers/loaders/transformer_flux.py deleted file mode 100644 index c7d81a8baebd080da2e44b84e46c266af21bd7c0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/transformer_flux.py +++ /dev/null @@ -1,180 +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 contextlib import nullcontext - -from ..models.embeddings import ( - ImageProjection, - MultiIPAdapterImageProjection, -) -from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, load_model_dict_into_meta -from ..utils import ( - is_accelerate_available, - is_torch_version, - logging, -) - - -if is_accelerate_available(): - pass - -logger = logging.get_logger(__name__) - - -class FluxTransformer2DLoadersMixin: - """ - Load layers into a [`FluxTransformer2DModel`]. - """ - - def _convert_ip_adapter_image_proj_to_diffusers(self, state_dict, low_cpu_mem_usage=_LOW_CPU_MEM_USAGE_DEFAULT): - if low_cpu_mem_usage: - if is_accelerate_available(): - from accelerate import init_empty_weights - - else: - low_cpu_mem_usage = False - logger.warning( - "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the" - " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install" - " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip" - " install accelerate\n```\n." - ) - - if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"): - raise NotImplementedError( - "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set" - " `low_cpu_mem_usage=False`." - ) - - updated_state_dict = {} - image_projection = None - init_context = init_empty_weights if low_cpu_mem_usage else nullcontext - - if "proj.weight" in state_dict: - # IP-Adapter - num_image_text_embeds = 4 - if state_dict["proj.weight"].shape[0] == 65536: - num_image_text_embeds = 16 - clip_embeddings_dim = state_dict["proj.weight"].shape[-1] - cross_attention_dim = state_dict["proj.weight"].shape[0] // num_image_text_embeds - - with init_context(): - image_projection = ImageProjection( - cross_attention_dim=cross_attention_dim, - image_embed_dim=clip_embeddings_dim, - num_image_text_embeds=num_image_text_embeds, - ) - - for key, value in state_dict.items(): - diffusers_name = key.replace("proj", "image_embeds") - updated_state_dict[diffusers_name] = value - - if not low_cpu_mem_usage: - image_projection.load_state_dict(updated_state_dict, strict=True) - else: - device_map = {"": self.device} - load_model_dict_into_meta(image_projection, updated_state_dict, device_map=device_map, dtype=self.dtype) - - return image_projection - - def _convert_ip_adapter_attn_to_diffusers(self, state_dicts, low_cpu_mem_usage=_LOW_CPU_MEM_USAGE_DEFAULT): - from ..models.attention_processor import ( - FluxIPAdapterJointAttnProcessor2_0, - ) - - if low_cpu_mem_usage: - if is_accelerate_available(): - from accelerate import init_empty_weights - - else: - low_cpu_mem_usage = False - logger.warning( - "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the" - " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install" - " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip" - " install accelerate\n```\n." - ) - - if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"): - raise NotImplementedError( - "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set" - " `low_cpu_mem_usage=False`." - ) - - # set ip-adapter cross-attention processors & load state_dict - attn_procs = {} - key_id = 0 - init_context = init_empty_weights if low_cpu_mem_usage else nullcontext - for name in self.attn_processors.keys(): - if name.startswith("single_transformer_blocks"): - attn_processor_class = self.attn_processors[name].__class__ - attn_procs[name] = attn_processor_class() - else: - cross_attention_dim = self.config.joint_attention_dim - hidden_size = self.inner_dim - attn_processor_class = FluxIPAdapterJointAttnProcessor2_0 - num_image_text_embeds = [] - for state_dict in state_dicts: - if "proj.weight" in state_dict["image_proj"]: - num_image_text_embed = 4 - if state_dict["image_proj"]["proj.weight"].shape[0] == 65536: - num_image_text_embed = 16 - # IP-Adapter - num_image_text_embeds += [num_image_text_embed] - - with init_context(): - attn_procs[name] = attn_processor_class( - hidden_size=hidden_size, - cross_attention_dim=cross_attention_dim, - scale=1.0, - num_tokens=num_image_text_embeds, - dtype=self.dtype, - device=self.device, - ) - - value_dict = {} - for i, state_dict in enumerate(state_dicts): - value_dict.update({f"to_k_ip.{i}.weight": state_dict["ip_adapter"][f"{key_id}.to_k_ip.weight"]}) - value_dict.update({f"to_v_ip.{i}.weight": state_dict["ip_adapter"][f"{key_id}.to_v_ip.weight"]}) - value_dict.update({f"to_k_ip.{i}.bias": state_dict["ip_adapter"][f"{key_id}.to_k_ip.bias"]}) - value_dict.update({f"to_v_ip.{i}.bias": state_dict["ip_adapter"][f"{key_id}.to_v_ip.bias"]}) - - if not low_cpu_mem_usage: - attn_procs[name].load_state_dict(value_dict) - else: - device_map = {"": self.device} - dtype = self.dtype - load_model_dict_into_meta(attn_procs[name], value_dict, device_map=device_map, dtype=dtype) - - key_id += 1 - - return attn_procs - - def _load_ip_adapter_weights(self, state_dicts, low_cpu_mem_usage=_LOW_CPU_MEM_USAGE_DEFAULT): - if not isinstance(state_dicts, list): - state_dicts = [state_dicts] - - self.encoder_hid_proj = None - - attn_procs = self._convert_ip_adapter_attn_to_diffusers(state_dicts, low_cpu_mem_usage=low_cpu_mem_usage) - self.set_attn_processor(attn_procs) - - image_projection_layers = [] - for state_dict in state_dicts: - image_projection_layer = self._convert_ip_adapter_image_proj_to_diffusers( - state_dict["image_proj"], low_cpu_mem_usage=low_cpu_mem_usage - ) - image_projection_layers.append(image_projection_layer) - - self.encoder_hid_proj = MultiIPAdapterImageProjection(image_projection_layers) - self.config.encoder_hid_dim_type = "ip_image_proj" diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/transformer_sd3.py b/venv/lib/python3.10/site-packages/diffusers/loaders/transformer_sd3.py deleted file mode 100644 index c58d3280cfe1f756204bcaf851b535fddccb92ec..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/transformer_sd3.py +++ /dev/null @@ -1,170 +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 contextlib import nullcontext -from typing import Dict - -from ..models.attention_processor import SD3IPAdapterJointAttnProcessor2_0 -from ..models.embeddings import IPAdapterTimeImageProjection -from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, load_model_dict_into_meta -from ..utils import is_accelerate_available, is_torch_version, logging - - -logger = logging.get_logger(__name__) - - -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: - if low_cpu_mem_usage: - if is_accelerate_available(): - from accelerate import init_empty_weights - - else: - low_cpu_mem_usage = False - logger.warning( - "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the" - " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install" - " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip" - " install accelerate\n```\n." - ) - - if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"): - raise NotImplementedError( - "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set" - " `low_cpu_mem_usage=False`." - ) - - # IP-Adapter cross attention parameters - hidden_size = self.config.attention_head_dim * self.config.num_attention_heads - ip_hidden_states_dim = self.config.attention_head_dim * self.config.num_attention_heads - timesteps_emb_dim = state_dict["0.norm_ip.linear.weight"].shape[1] - - # Dict where key is transformer layer index, value is attention processor's state dict - # ip_adapter state dict keys example: "0.norm_ip.linear.weight" - layer_state_dict = {idx: {} for idx in range(len(self.attn_processors))} - for key, weights in state_dict.items(): - idx, name = key.split(".", maxsplit=1) - layer_state_dict[int(idx)][name] = weights - - # Create IP-Adapter attention processor & load state_dict - attn_procs = {} - init_context = init_empty_weights if low_cpu_mem_usage else nullcontext - for idx, name in enumerate(self.attn_processors.keys()): - with init_context(): - attn_procs[name] = SD3IPAdapterJointAttnProcessor2_0( - hidden_size=hidden_size, - ip_hidden_states_dim=ip_hidden_states_dim, - head_dim=self.config.attention_head_dim, - timesteps_emb_dim=timesteps_emb_dim, - ) - - if not low_cpu_mem_usage: - attn_procs[name].load_state_dict(layer_state_dict[idx], strict=True) - else: - device_map = {"": self.device} - load_model_dict_into_meta( - attn_procs[name], layer_state_dict[idx], device_map=device_map, dtype=self.dtype - ) - - 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 - ) -> IPAdapterTimeImageProjection: - if low_cpu_mem_usage: - if is_accelerate_available(): - from accelerate import init_empty_weights - - else: - low_cpu_mem_usage = False - logger.warning( - "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the" - " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install" - " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip" - " install accelerate\n```\n." - ) - - if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"): - raise NotImplementedError( - "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set" - " `low_cpu_mem_usage=False`." - ) - - init_context = init_empty_weights if low_cpu_mem_usage else nullcontext - - # Convert to diffusers - updated_state_dict = {} - for key, value in state_dict.items(): - # InstantX/SD3.5-Large-IP-Adapter - if key.startswith("layers."): - idx = key.split(".")[1] - key = key.replace(f"layers.{idx}.0.norm1", f"layers.{idx}.ln0") - key = key.replace(f"layers.{idx}.0.norm2", f"layers.{idx}.ln1") - key = key.replace(f"layers.{idx}.0.to_q", f"layers.{idx}.attn.to_q") - key = key.replace(f"layers.{idx}.0.to_kv", f"layers.{idx}.attn.to_kv") - key = key.replace(f"layers.{idx}.0.to_out", f"layers.{idx}.attn.to_out.0") - key = key.replace(f"layers.{idx}.1.0", f"layers.{idx}.adaln_norm") - key = key.replace(f"layers.{idx}.1.1", f"layers.{idx}.ff.net.0.proj") - key = key.replace(f"layers.{idx}.1.3", f"layers.{idx}.ff.net.2") - key = key.replace(f"layers.{idx}.2.1", f"layers.{idx}.adaln_proj") - updated_state_dict[key] = value - - # Image projection parameters - embed_dim = updated_state_dict["proj_in.weight"].shape[1] - output_dim = updated_state_dict["proj_out.weight"].shape[0] - hidden_dim = updated_state_dict["proj_in.weight"].shape[0] - heads = updated_state_dict["layers.0.attn.to_q.weight"].shape[0] // 64 - num_queries = updated_state_dict["latents"].shape[1] - timestep_in_dim = updated_state_dict["time_embedding.linear_1.weight"].shape[1] - - # Image projection - with init_context(): - image_proj = IPAdapterTimeImageProjection( - embed_dim=embed_dim, - output_dim=output_dim, - hidden_dim=hidden_dim, - heads=heads, - num_queries=num_queries, - timestep_in_dim=timestep_in_dim, - ) - - if not low_cpu_mem_usage: - image_proj.load_state_dict(updated_state_dict, strict=True) - else: - device_map = {"": self.device} - load_model_dict_into_meta(image_proj, updated_state_dict, device_map=device_map, dtype=self.dtype) - - return image_proj - - 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: - state_dict (`Dict`): - State dict with keys "ip_adapter", which contains parameters for attention processors, and - "image_proj", which contains parameters for image projection net. - low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): - Speed up model loading only loading the pretrained weights and not initializing the weights. This also - tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. - Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this - argument to `True` will raise an error. - """ - - attn_procs = self._convert_ip_adapter_attn_to_diffusers(state_dict["ip_adapter"], low_cpu_mem_usage) - self.set_attn_processor(attn_procs) - - self.image_proj = self._convert_ip_adapter_image_proj_to_diffusers(state_dict["image_proj"], low_cpu_mem_usage) diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/unet.py b/venv/lib/python3.10/site-packages/diffusers/loaders/unet.py deleted file mode 100644 index 68be841191778a9b7e3cc0c0e8b522d3ecf4aeca..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/unet.py +++ /dev/null @@ -1,926 +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 os -from collections import defaultdict -from contextlib import nullcontext -from pathlib import Path -from typing import Callable, Dict, Union - -import safetensors -import torch -import torch.nn.functional as F -from huggingface_hub.utils import validate_hf_hub_args - -from ..models.embeddings import ( - ImageProjection, - IPAdapterFaceIDImageProjection, - IPAdapterFaceIDPlusImageProjection, - IPAdapterFullImageProjection, - IPAdapterPlusImageProjection, - MultiIPAdapterImageProjection, -) -from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, load_model_dict_into_meta, load_state_dict -from ..utils import ( - USE_PEFT_BACKEND, - _get_model_file, - convert_unet_state_dict_to_peft, - deprecate, - get_adapter_name, - get_peft_kwargs, - is_accelerate_available, - is_peft_version, - is_torch_version, - logging, -) -from .lora_base import _func_optionally_disable_offloading -from .lora_pipeline import LORA_WEIGHT_NAME, LORA_WEIGHT_NAME_SAFE, TEXT_ENCODER_NAME, UNET_NAME -from .utils import AttnProcsLayers - - -logger = logging.get_logger(__name__) - - -CUSTOM_DIFFUSION_WEIGHT_NAME = "pytorch_custom_diffusion_weights.bin" -CUSTOM_DIFFUSION_WEIGHT_NAME_SAFE = "pytorch_custom_diffusion_weights.safetensors" - - -class UNet2DConditionLoadersMixin: - """ - Load LoRA layers into a [`UNet2DCondtionModel`]. - """ - - text_encoder_name = TEXT_ENCODER_NAME - 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): - r""" - Load pretrained attention processor layers into [`UNet2DConditionModel`]. Attention processor layers have to be - defined in - [`attention_processor.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py) - and be a `torch.nn.Module` class. Currently supported: LoRA, Custom Diffusion. For LoRA, one must install - `peft`: `pip install -U peft`. - - Parameters: - pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): - Can be either: - - - A string, the model id (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a directory (for example `./my_model_directory`) containing the model weights saved - with [`ModelMixin.save_pretrained`]. - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - 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]`): - 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). - adapter_name (`str`, *optional*, defaults to None): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. - weight_name (`str`, *optional*, defaults to None): - Name of the serialized state dict file. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. - - Example: - - ```py - from diffusers import AutoPipelineForText2Image - import torch - - pipeline = AutoPipelineForText2Image.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ).to("cuda") - pipeline.unet.load_attn_procs( - "jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_name="cinematic" - ) - ``` - """ - 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) - adapter_name = kwargs.pop("adapter_name", None) - _pipeline = kwargs.pop("_pipeline", None) - network_alphas = kwargs.pop("network_alphas", None) - low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) - allow_pickle = False - - 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 use_safetensors is None: - use_safetensors = True - allow_pickle = True - - user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"} - - model_file = None - if not isinstance(pretrained_model_name_or_path_or_dict, dict): - # Let's first try to load .safetensors weights - if (use_safetensors and weight_name is None) or ( - weight_name is not None and weight_name.endswith(".safetensors") - ): - try: - model_file = _get_model_file( - pretrained_model_name_or_path_or_dict, - weights_name=weight_name or LORA_WEIGHT_NAME_SAFE, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - ) - state_dict = safetensors.torch.load_file(model_file, device="cpu") - except IOError as e: - if not allow_pickle: - raise e - # try loading non-safetensors weights - pass - if model_file is None: - model_file = _get_model_file( - pretrained_model_name_or_path_or_dict, - weights_name=weight_name or LORA_WEIGHT_NAME, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - ) - state_dict = load_state_dict(model_file) - else: - state_dict = pretrained_model_name_or_path_or_dict - - is_custom_diffusion = any("custom_diffusion" in k for k in state_dict.keys()) - is_lora = all(("lora" in k or k.endswith(".alpha")) for k in state_dict.keys()) - is_model_cpu_offload = False - is_sequential_cpu_offload = False - - if is_lora: - deprecation_message = "Using the `load_attn_procs()` method has been deprecated and will be removed in a future version. Please use `load_lora_adapter()`." - deprecate("load_attn_procs", "0.40.0", deprecation_message) - - if is_custom_diffusion: - attn_processors = self._process_custom_diffusion(state_dict=state_dict) - elif is_lora: - is_model_cpu_offload, is_sequential_cpu_offload = self._process_lora( - state_dict=state_dict, - unet_identifier_key=self.unet_name, - network_alphas=network_alphas, - adapter_name=adapter_name, - _pipeline=_pipeline, - low_cpu_mem_usage=low_cpu_mem_usage, - ) - else: - raise ValueError( - f"{model_file} does not seem to be in the correct format expected by Custom Diffusion training." - ) - - # - - def _process_custom_diffusion(self, state_dict): - from ..models.attention_processor import CustomDiffusionAttnProcessor - - attn_processors = {} - custom_diffusion_grouped_dict = defaultdict(dict) - for key, value in state_dict.items(): - if len(value) == 0: - custom_diffusion_grouped_dict[key] = {} - else: - if "to_out" in key: - attn_processor_key, sub_key = ".".join(key.split(".")[:-3]), ".".join(key.split(".")[-3:]) - else: - attn_processor_key, sub_key = ".".join(key.split(".")[:-2]), ".".join(key.split(".")[-2:]) - custom_diffusion_grouped_dict[attn_processor_key][sub_key] = value - - for key, value_dict in custom_diffusion_grouped_dict.items(): - if len(value_dict) == 0: - attn_processors[key] = CustomDiffusionAttnProcessor( - train_kv=False, train_q_out=False, hidden_size=None, cross_attention_dim=None - ) - else: - cross_attention_dim = value_dict["to_k_custom_diffusion.weight"].shape[1] - hidden_size = value_dict["to_k_custom_diffusion.weight"].shape[0] - train_q_out = True if "to_q_custom_diffusion.weight" in value_dict else False - attn_processors[key] = CustomDiffusionAttnProcessor( - train_kv=True, - train_q_out=train_q_out, - hidden_size=hidden_size, - cross_attention_dim=cross_attention_dim, - ) - attn_processors[key].load_state_dict(value_dict) - - return attn_processors - - def _process_lora( - self, state_dict, unet_identifier_key, network_alphas, adapter_name, _pipeline, low_cpu_mem_usage - ): - # This method does the following things: - # 1. Filters the `state_dict` with keys matching `unet_identifier_key` when using the non-legacy - # format. For legacy format no filtering is applied. - # 2. Converts the `state_dict` to the `peft` compatible format. - # 3. Creates a `LoraConfig` and then injects the converted `state_dict` into the UNet per the - # `LoraConfig` specs. - # 4. It also reports if the underlying `_pipeline` has any kind of offloading inside of it. - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for this method.") - - from peft import LoraConfig, inject_adapter_in_model, set_peft_model_state_dict - - keys = list(state_dict.keys()) - - unet_keys = [k for k in keys if k.startswith(unet_identifier_key)] - unet_state_dict = { - k.replace(f"{unet_identifier_key}.", ""): v for k, v in state_dict.items() if k in unet_keys - } - - if network_alphas is not None: - alpha_keys = [k for k in network_alphas.keys() if k.startswith(unet_identifier_key)] - network_alphas = { - k.replace(f"{unet_identifier_key}.", ""): v for k, v in network_alphas.items() if k in alpha_keys - } - - is_model_cpu_offload = False - is_sequential_cpu_offload = False - state_dict_to_be_used = unet_state_dict if len(unet_state_dict) > 0 else state_dict - - if len(state_dict_to_be_used) > 0: - if adapter_name in getattr(self, "peft_config", {}): - raise ValueError( - f"Adapter name {adapter_name} already in use in the Unet - please select a new adapter name." - ) - - state_dict = convert_unet_state_dict_to_peft(state_dict_to_be_used) - - if network_alphas is not None: - # The alphas state dict have the same structure as Unet, thus we convert it to peft format using - # `convert_unet_state_dict_to_peft` method. - network_alphas = convert_unet_state_dict_to_peft(network_alphas) - - rank = {} - for key, val in state_dict.items(): - if "lora_B" in key: - rank[key] = val.shape[1] - - lora_config_kwargs = get_peft_kwargs(rank, network_alphas, state_dict, is_unet=True) - if "use_dora" in lora_config_kwargs: - if lora_config_kwargs["use_dora"]: - if is_peft_version("<", "0.9.0"): - raise ValueError( - "You need `peft` 0.9.0 at least to use DoRA-enabled LoRAs. Please upgrade your installation of `peft`." - ) - else: - if is_peft_version("<", "0.9.0"): - lora_config_kwargs.pop("use_dora") - - if "lora_bias" in lora_config_kwargs: - if lora_config_kwargs["lora_bias"]: - if is_peft_version("<=", "0.13.2"): - raise ValueError( - "You need `peft` 0.14.0 at least to use `bias` in LoRAs. Please upgrade your installation of `peft`." - ) - else: - if is_peft_version("<=", "0.13.2"): - lora_config_kwargs.pop("lora_bias") - - lora_config = LoraConfig(**lora_config_kwargs) - - # adapter_name - if adapter_name is None: - adapter_name = get_adapter_name(self) - - # In case the pipeline has been already offloaded to CPU - temporarily remove the hooks - # otherwise loading LoRA weights will lead to an error - is_model_cpu_offload, is_sequential_cpu_offload = self._optionally_disable_offloading(_pipeline) - peft_kwargs = {} - if is_peft_version(">=", "0.13.1"): - peft_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage - - inject_adapter_in_model(lora_config, self, adapter_name=adapter_name, **peft_kwargs) - incompatible_keys = set_peft_model_state_dict(self, state_dict, adapter_name, **peft_kwargs) - - warn_msg = "" - if incompatible_keys is not None: - # Check only for unexpected keys. - unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) - if unexpected_keys: - lora_unexpected_keys = [k for k in unexpected_keys if "lora_" in k and adapter_name in k] - if lora_unexpected_keys: - warn_msg = ( - f"Loading adapter weights from state_dict led to unexpected keys found in the model:" - f" {', '.join(lora_unexpected_keys)}. " - ) - - # Filter missing keys specific to the current adapter. - missing_keys = getattr(incompatible_keys, "missing_keys", None) - if missing_keys: - lora_missing_keys = [k for k in missing_keys if "lora_" in k and adapter_name in k] - if lora_missing_keys: - warn_msg += ( - f"Loading adapter weights from state_dict led to missing keys in the model:" - f" {', '.join(lora_missing_keys)}." - ) - - if warn_msg: - logger.warning(warn_msg) - - return is_model_cpu_offload, is_sequential_cpu_offload - - @classmethod - # Copied from diffusers.loaders.lora_base.LoraBaseMixin._optionally_disable_offloading - def _optionally_disable_offloading(cls, _pipeline): - return _func_optionally_disable_offloading(_pipeline=_pipeline) - - def save_attn_procs( - self, - save_directory: Union[str, os.PathLike], - is_main_process: bool = True, - weight_name: str = None, - save_function: Callable = None, - safe_serialization: bool = True, - **kwargs, - ): - r""" - Save attention processor layers to a directory so that it can be reloaded with the - [`~loaders.UNet2DConditionLoadersMixin.load_attn_procs`] method. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save an attention processor to (will be created if it doesn't exist). - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or with `pickle`. - - Example: - - ```py - import torch - from diffusers import DiffusionPipeline - - pipeline = DiffusionPipeline.from_pretrained( - "CompVis/stable-diffusion-v1-4", - torch_dtype=torch.float16, - ).to("cuda") - pipeline.unet.load_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin") - pipeline.unet.save_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin") - ``` - """ - from ..models.attention_processor import ( - CustomDiffusionAttnProcessor, - CustomDiffusionAttnProcessor2_0, - CustomDiffusionXFormersAttnProcessor, - ) - - if os.path.isfile(save_directory): - logger.error(f"Provided path ({save_directory}) should be a directory, not a file") - return - - is_custom_diffusion = any( - isinstance( - x, - (CustomDiffusionAttnProcessor, CustomDiffusionAttnProcessor2_0, CustomDiffusionXFormersAttnProcessor), - ) - for (_, x) in self.attn_processors.items() - ) - if is_custom_diffusion: - state_dict = self._get_custom_diffusion_state_dict() - if save_function is None and safe_serialization: - # safetensors does not support saving dicts with non-tensor values - empty_state_dict = {k: v for k, v in state_dict.items() if not isinstance(v, torch.Tensor)} - if len(empty_state_dict) > 0: - logger.warning( - f"Safetensors does not support saving dicts with non-tensor values. " - f"The following keys will be ignored: {empty_state_dict.keys()}" - ) - state_dict = {k: v for k, v in state_dict.items() if isinstance(v, torch.Tensor)} - else: - deprecation_message = "Using the `save_attn_procs()` method has been deprecated and will be removed in a future version. Please use `save_lora_adapter()`." - deprecate("save_attn_procs", "0.40.0", deprecation_message) - - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for saving LoRAs using the `save_attn_procs()` method.") - - from peft.utils import get_peft_model_state_dict - - state_dict = get_peft_model_state_dict(self) - - if save_function is None: - if safe_serialization: - - def save_function(weights, filename): - return safetensors.torch.save_file(weights, filename, metadata={"format": "pt"}) - - else: - save_function = torch.save - - os.makedirs(save_directory, exist_ok=True) - - if weight_name is None: - if safe_serialization: - weight_name = CUSTOM_DIFFUSION_WEIGHT_NAME_SAFE if is_custom_diffusion else LORA_WEIGHT_NAME_SAFE - else: - weight_name = CUSTOM_DIFFUSION_WEIGHT_NAME if is_custom_diffusion else LORA_WEIGHT_NAME - - # Save the model - save_path = Path(save_directory, weight_name).as_posix() - save_function(state_dict, save_path) - logger.info(f"Model weights saved in {save_path}") - - def _get_custom_diffusion_state_dict(self): - from ..models.attention_processor import ( - CustomDiffusionAttnProcessor, - CustomDiffusionAttnProcessor2_0, - CustomDiffusionXFormersAttnProcessor, - ) - - model_to_save = AttnProcsLayers( - { - y: x - for (y, x) in self.attn_processors.items() - if isinstance( - x, - ( - CustomDiffusionAttnProcessor, - CustomDiffusionAttnProcessor2_0, - CustomDiffusionXFormersAttnProcessor, - ), - ) - } - ) - state_dict = model_to_save.state_dict() - for name, attn in self.attn_processors.items(): - if len(attn.state_dict()) == 0: - state_dict[name] = {} - - return state_dict - - def _convert_ip_adapter_image_proj_to_diffusers(self, state_dict, low_cpu_mem_usage=_LOW_CPU_MEM_USAGE_DEFAULT): - if low_cpu_mem_usage: - if is_accelerate_available(): - from accelerate import init_empty_weights - - else: - low_cpu_mem_usage = False - logger.warning( - "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the" - " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install" - " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip" - " install accelerate\n```\n." - ) - - if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"): - raise NotImplementedError( - "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set" - " `low_cpu_mem_usage=False`." - ) - - updated_state_dict = {} - image_projection = None - init_context = init_empty_weights if low_cpu_mem_usage else nullcontext - - if "proj.weight" in state_dict: - # IP-Adapter - num_image_text_embeds = 4 - clip_embeddings_dim = state_dict["proj.weight"].shape[-1] - cross_attention_dim = state_dict["proj.weight"].shape[0] // 4 - - with init_context(): - image_projection = ImageProjection( - cross_attention_dim=cross_attention_dim, - image_embed_dim=clip_embeddings_dim, - num_image_text_embeds=num_image_text_embeds, - ) - - for key, value in state_dict.items(): - diffusers_name = key.replace("proj", "image_embeds") - updated_state_dict[diffusers_name] = value - - elif "proj.3.weight" in state_dict: - # IP-Adapter Full - clip_embeddings_dim = state_dict["proj.0.weight"].shape[0] - cross_attention_dim = state_dict["proj.3.weight"].shape[0] - - with init_context(): - image_projection = IPAdapterFullImageProjection( - cross_attention_dim=cross_attention_dim, image_embed_dim=clip_embeddings_dim - ) - - for key, value in state_dict.items(): - diffusers_name = key.replace("proj.0", "ff.net.0.proj") - diffusers_name = diffusers_name.replace("proj.2", "ff.net.2") - diffusers_name = diffusers_name.replace("proj.3", "norm") - updated_state_dict[diffusers_name] = value - - elif "perceiver_resampler.proj_in.weight" in state_dict: - # IP-Adapter Face ID Plus - id_embeddings_dim = state_dict["proj.0.weight"].shape[1] - embed_dims = state_dict["perceiver_resampler.proj_in.weight"].shape[0] - hidden_dims = state_dict["perceiver_resampler.proj_in.weight"].shape[1] - output_dims = state_dict["perceiver_resampler.proj_out.weight"].shape[0] - heads = state_dict["perceiver_resampler.layers.0.0.to_q.weight"].shape[0] // 64 - - with init_context(): - image_projection = IPAdapterFaceIDPlusImageProjection( - embed_dims=embed_dims, - output_dims=output_dims, - hidden_dims=hidden_dims, - heads=heads, - id_embeddings_dim=id_embeddings_dim, - ) - - for key, value in state_dict.items(): - diffusers_name = key.replace("perceiver_resampler.", "") - diffusers_name = diffusers_name.replace("0.to", "attn.to") - diffusers_name = diffusers_name.replace("0.1.0.", "0.ff.0.") - diffusers_name = diffusers_name.replace("0.1.1.weight", "0.ff.1.net.0.proj.weight") - diffusers_name = diffusers_name.replace("0.1.3.weight", "0.ff.1.net.2.weight") - diffusers_name = diffusers_name.replace("1.1.0.", "1.ff.0.") - diffusers_name = diffusers_name.replace("1.1.1.weight", "1.ff.1.net.0.proj.weight") - diffusers_name = diffusers_name.replace("1.1.3.weight", "1.ff.1.net.2.weight") - diffusers_name = diffusers_name.replace("2.1.0.", "2.ff.0.") - diffusers_name = diffusers_name.replace("2.1.1.weight", "2.ff.1.net.0.proj.weight") - diffusers_name = diffusers_name.replace("2.1.3.weight", "2.ff.1.net.2.weight") - diffusers_name = diffusers_name.replace("3.1.0.", "3.ff.0.") - diffusers_name = diffusers_name.replace("3.1.1.weight", "3.ff.1.net.0.proj.weight") - diffusers_name = diffusers_name.replace("3.1.3.weight", "3.ff.1.net.2.weight") - diffusers_name = diffusers_name.replace("layers.0.0", "layers.0.ln0") - diffusers_name = diffusers_name.replace("layers.0.1", "layers.0.ln1") - diffusers_name = diffusers_name.replace("layers.1.0", "layers.1.ln0") - diffusers_name = diffusers_name.replace("layers.1.1", "layers.1.ln1") - diffusers_name = diffusers_name.replace("layers.2.0", "layers.2.ln0") - diffusers_name = diffusers_name.replace("layers.2.1", "layers.2.ln1") - diffusers_name = diffusers_name.replace("layers.3.0", "layers.3.ln0") - diffusers_name = diffusers_name.replace("layers.3.1", "layers.3.ln1") - - if "norm1" in diffusers_name: - updated_state_dict[diffusers_name.replace("0.norm1", "0")] = value - elif "norm2" in diffusers_name: - updated_state_dict[diffusers_name.replace("0.norm2", "1")] = value - elif "to_kv" in diffusers_name: - v_chunk = value.chunk(2, dim=0) - updated_state_dict[diffusers_name.replace("to_kv", "to_k")] = v_chunk[0] - updated_state_dict[diffusers_name.replace("to_kv", "to_v")] = v_chunk[1] - elif "to_out" in diffusers_name: - updated_state_dict[diffusers_name.replace("to_out", "to_out.0")] = value - elif "proj.0.weight" == diffusers_name: - updated_state_dict["proj.net.0.proj.weight"] = value - elif "proj.0.bias" == diffusers_name: - updated_state_dict["proj.net.0.proj.bias"] = value - elif "proj.2.weight" == diffusers_name: - updated_state_dict["proj.net.2.weight"] = value - elif "proj.2.bias" == diffusers_name: - updated_state_dict["proj.net.2.bias"] = value - else: - updated_state_dict[diffusers_name] = value - - elif "norm.weight" in state_dict: - # IP-Adapter Face ID - id_embeddings_dim_in = state_dict["proj.0.weight"].shape[1] - id_embeddings_dim_out = state_dict["proj.0.weight"].shape[0] - multiplier = id_embeddings_dim_out // id_embeddings_dim_in - norm_layer = "norm.weight" - cross_attention_dim = state_dict[norm_layer].shape[0] - num_tokens = state_dict["proj.2.weight"].shape[0] // cross_attention_dim - - with init_context(): - image_projection = IPAdapterFaceIDImageProjection( - cross_attention_dim=cross_attention_dim, - image_embed_dim=id_embeddings_dim_in, - mult=multiplier, - num_tokens=num_tokens, - ) - - for key, value in state_dict.items(): - diffusers_name = key.replace("proj.0", "ff.net.0.proj") - diffusers_name = diffusers_name.replace("proj.2", "ff.net.2") - updated_state_dict[diffusers_name] = value - - else: - # IP-Adapter Plus - num_image_text_embeds = state_dict["latents"].shape[1] - embed_dims = state_dict["proj_in.weight"].shape[1] - output_dims = state_dict["proj_out.weight"].shape[0] - hidden_dims = state_dict["latents"].shape[2] - attn_key_present = any("attn" in k for k in state_dict) - heads = ( - state_dict["layers.0.attn.to_q.weight"].shape[0] // 64 - if attn_key_present - else state_dict["layers.0.0.to_q.weight"].shape[0] // 64 - ) - - with init_context(): - image_projection = IPAdapterPlusImageProjection( - embed_dims=embed_dims, - output_dims=output_dims, - hidden_dims=hidden_dims, - heads=heads, - num_queries=num_image_text_embeds, - ) - - for key, value in state_dict.items(): - diffusers_name = key.replace("0.to", "2.to") - - diffusers_name = diffusers_name.replace("0.0.norm1", "0.ln0") - diffusers_name = diffusers_name.replace("0.0.norm2", "0.ln1") - diffusers_name = diffusers_name.replace("1.0.norm1", "1.ln0") - diffusers_name = diffusers_name.replace("1.0.norm2", "1.ln1") - diffusers_name = diffusers_name.replace("2.0.norm1", "2.ln0") - diffusers_name = diffusers_name.replace("2.0.norm2", "2.ln1") - diffusers_name = diffusers_name.replace("3.0.norm1", "3.ln0") - diffusers_name = diffusers_name.replace("3.0.norm2", "3.ln1") - - if "to_kv" in diffusers_name: - parts = diffusers_name.split(".") - parts[2] = "attn" - diffusers_name = ".".join(parts) - v_chunk = value.chunk(2, dim=0) - updated_state_dict[diffusers_name.replace("to_kv", "to_k")] = v_chunk[0] - updated_state_dict[diffusers_name.replace("to_kv", "to_v")] = v_chunk[1] - elif "to_q" in diffusers_name: - parts = diffusers_name.split(".") - parts[2] = "attn" - diffusers_name = ".".join(parts) - updated_state_dict[diffusers_name] = value - elif "to_out" in diffusers_name: - parts = diffusers_name.split(".") - parts[2] = "attn" - diffusers_name = ".".join(parts) - updated_state_dict[diffusers_name.replace("to_out", "to_out.0")] = value - else: - diffusers_name = diffusers_name.replace("0.1.0", "0.ff.0") - diffusers_name = diffusers_name.replace("0.1.1", "0.ff.1.net.0.proj") - diffusers_name = diffusers_name.replace("0.1.3", "0.ff.1.net.2") - - diffusers_name = diffusers_name.replace("1.1.0", "1.ff.0") - diffusers_name = diffusers_name.replace("1.1.1", "1.ff.1.net.0.proj") - diffusers_name = diffusers_name.replace("1.1.3", "1.ff.1.net.2") - - diffusers_name = diffusers_name.replace("2.1.0", "2.ff.0") - diffusers_name = diffusers_name.replace("2.1.1", "2.ff.1.net.0.proj") - diffusers_name = diffusers_name.replace("2.1.3", "2.ff.1.net.2") - - diffusers_name = diffusers_name.replace("3.1.0", "3.ff.0") - diffusers_name = diffusers_name.replace("3.1.1", "3.ff.1.net.0.proj") - diffusers_name = diffusers_name.replace("3.1.3", "3.ff.1.net.2") - updated_state_dict[diffusers_name] = value - - if not low_cpu_mem_usage: - image_projection.load_state_dict(updated_state_dict, strict=True) - else: - device_map = {"": self.device} - load_model_dict_into_meta(image_projection, updated_state_dict, device_map=device_map, dtype=self.dtype) - - return image_projection - - def _convert_ip_adapter_attn_to_diffusers(self, state_dicts, low_cpu_mem_usage=_LOW_CPU_MEM_USAGE_DEFAULT): - from ..models.attention_processor import ( - IPAdapterAttnProcessor, - IPAdapterAttnProcessor2_0, - IPAdapterXFormersAttnProcessor, - ) - - if low_cpu_mem_usage: - if is_accelerate_available(): - from accelerate import init_empty_weights - - else: - low_cpu_mem_usage = False - logger.warning( - "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the" - " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install" - " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip" - " install accelerate\n```\n." - ) - - if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"): - raise NotImplementedError( - "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set" - " `low_cpu_mem_usage=False`." - ) - - # set ip-adapter cross-attention processors & load state_dict - attn_procs = {} - key_id = 1 - init_context = init_empty_weights if low_cpu_mem_usage else nullcontext - for name in self.attn_processors.keys(): - cross_attention_dim = None if name.endswith("attn1.processor") else self.config.cross_attention_dim - if name.startswith("mid_block"): - hidden_size = self.config.block_out_channels[-1] - elif name.startswith("up_blocks"): - block_id = int(name[len("up_blocks.")]) - hidden_size = list(reversed(self.config.block_out_channels))[block_id] - elif name.startswith("down_blocks"): - block_id = int(name[len("down_blocks.")]) - hidden_size = self.config.block_out_channels[block_id] - - if cross_attention_dim is None or "motion_modules" in name: - attn_processor_class = self.attn_processors[name].__class__ - attn_procs[name] = attn_processor_class() - else: - if "XFormers" in str(self.attn_processors[name].__class__): - attn_processor_class = IPAdapterXFormersAttnProcessor - else: - attn_processor_class = ( - IPAdapterAttnProcessor2_0 - if hasattr(F, "scaled_dot_product_attention") - else IPAdapterAttnProcessor - ) - num_image_text_embeds = [] - for state_dict in state_dicts: - if "proj.weight" in state_dict["image_proj"]: - # IP-Adapter - num_image_text_embeds += [4] - elif "proj.3.weight" in state_dict["image_proj"]: - # IP-Adapter Full Face - num_image_text_embeds += [257] # 256 CLIP tokens + 1 CLS token - elif "perceiver_resampler.proj_in.weight" in state_dict["image_proj"]: - # IP-Adapter Face ID Plus - num_image_text_embeds += [4] - elif "norm.weight" in state_dict["image_proj"]: - # IP-Adapter Face ID - num_image_text_embeds += [4] - else: - # IP-Adapter Plus - num_image_text_embeds += [state_dict["image_proj"]["latents"].shape[1]] - - with init_context(): - attn_procs[name] = attn_processor_class( - hidden_size=hidden_size, - cross_attention_dim=cross_attention_dim, - scale=1.0, - num_tokens=num_image_text_embeds, - ) - - value_dict = {} - for i, state_dict in enumerate(state_dicts): - value_dict.update({f"to_k_ip.{i}.weight": state_dict["ip_adapter"][f"{key_id}.to_k_ip.weight"]}) - value_dict.update({f"to_v_ip.{i}.weight": state_dict["ip_adapter"][f"{key_id}.to_v_ip.weight"]}) - - if not low_cpu_mem_usage: - attn_procs[name].load_state_dict(value_dict) - else: - device = next(iter(value_dict.values())).device - dtype = next(iter(value_dict.values())).dtype - device_map = {"": device} - load_model_dict_into_meta(attn_procs[name], value_dict, device_map=device_map, dtype=dtype) - - key_id += 2 - - return attn_procs - - def _load_ip_adapter_weights(self, state_dicts, low_cpu_mem_usage=_LOW_CPU_MEM_USAGE_DEFAULT): - if not isinstance(state_dicts, list): - state_dicts = [state_dicts] - - # Kolors Unet already has a `encoder_hid_proj` - if ( - self.encoder_hid_proj is not None - and self.config.encoder_hid_dim_type == "text_proj" - and not hasattr(self, "text_encoder_hid_proj") - ): - self.text_encoder_hid_proj = self.encoder_hid_proj - - # Set encoder_hid_proj after loading ip_adapter weights, - # because `IPAdapterPlusImageProjection` also has `attn_processors`. - self.encoder_hid_proj = None - - attn_procs = self._convert_ip_adapter_attn_to_diffusers(state_dicts, low_cpu_mem_usage=low_cpu_mem_usage) - self.set_attn_processor(attn_procs) - - # convert IP-Adapter Image Projection layers to diffusers - image_projection_layers = [] - for state_dict in state_dicts: - image_projection_layer = self._convert_ip_adapter_image_proj_to_diffusers( - state_dict["image_proj"], low_cpu_mem_usage=low_cpu_mem_usage - ) - image_projection_layers.append(image_projection_layer) - - self.encoder_hid_proj = MultiIPAdapterImageProjection(image_projection_layers) - self.config.encoder_hid_dim_type = "ip_image_proj" - - self.to(dtype=self.dtype, device=self.device) - - def _load_ip_adapter_loras(self, state_dicts): - lora_dicts = {} - for key_id, name in enumerate(self.attn_processors.keys()): - for i, state_dict in enumerate(state_dicts): - if f"{key_id}.to_k_lora.down.weight" in state_dict["ip_adapter"]: - if i not in lora_dicts: - lora_dicts[i] = {} - lora_dicts[i].update( - { - f"unet.{name}.to_k_lora.down.weight": state_dict["ip_adapter"][ - f"{key_id}.to_k_lora.down.weight" - ] - } - ) - lora_dicts[i].update( - { - f"unet.{name}.to_q_lora.down.weight": state_dict["ip_adapter"][ - f"{key_id}.to_q_lora.down.weight" - ] - } - ) - lora_dicts[i].update( - { - f"unet.{name}.to_v_lora.down.weight": state_dict["ip_adapter"][ - f"{key_id}.to_v_lora.down.weight" - ] - } - ) - lora_dicts[i].update( - { - f"unet.{name}.to_out_lora.down.weight": state_dict["ip_adapter"][ - f"{key_id}.to_out_lora.down.weight" - ] - } - ) - lora_dicts[i].update( - {f"unet.{name}.to_k_lora.up.weight": state_dict["ip_adapter"][f"{key_id}.to_k_lora.up.weight"]} - ) - lora_dicts[i].update( - {f"unet.{name}.to_q_lora.up.weight": state_dict["ip_adapter"][f"{key_id}.to_q_lora.up.weight"]} - ) - lora_dicts[i].update( - {f"unet.{name}.to_v_lora.up.weight": state_dict["ip_adapter"][f"{key_id}.to_v_lora.up.weight"]} - ) - lora_dicts[i].update( - { - f"unet.{name}.to_out_lora.up.weight": state_dict["ip_adapter"][ - f"{key_id}.to_out_lora.up.weight" - ] - } - ) - return lora_dicts diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/unet_loader_utils.py b/venv/lib/python3.10/site-packages/diffusers/loaders/unet_loader_utils.py deleted file mode 100644 index 274665204de063892f3ce8221f36e176de811505..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/unet_loader_utils.py +++ /dev/null @@ -1,163 +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 copy -from typing import TYPE_CHECKING, Dict, List, Union - -from ..utils import logging - - -if TYPE_CHECKING: - # import here to avoid circular imports - from ..models import UNet2DConditionModel - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -def _translate_into_actual_layer_name(name): - """Translate user-friendly name (e.g. 'mid') into actual layer name (e.g. 'mid_block.attentions.0')""" - if name == "mid": - return "mid_block.attentions.0" - - updown, block, attn = name.split(".") - - updown = updown.replace("down", "down_blocks").replace("up", "up_blocks") - block = block.replace("block_", "") - attn = "attentions." + attn - - return ".".join((updown, block, attn)) - - -def _maybe_expand_lora_scales( - unet: "UNet2DConditionModel", weight_scales: List[Union[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")], - } - transformer_per_block = {"down": unet.config.layers_per_block, "up": unet.config.layers_per_block + 1} - - expanded_weight_scales = [ - _maybe_expand_lora_scales_for_one_adapter( - weight_for_adapter, - blocks_with_transformer, - transformer_per_block, - unet.state_dict(), - default_scale=default_scale, - ) - for weight_for_adapter in weight_scales - ] - - return expanded_weight_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], - state_dict: None, - default_scale: float = 1.0, -): - """ - Expands the inputs into a more granular dictionary. See the example below for more details. - - Parameters: - scales (`Union[float, Dict]`): - Scales dict to expand. - blocks_with_transformer (`Dict[str, int]`): - Dict with keys 'up' and 'down', showing which blocks have transformer layers - transformer_per_block (`Dict[str, int]`): - Dict with keys 'up' and 'down', showing how many transformer layers each block has - - E.g. turns - ```python - scales = {"down": 2, "mid": 3, "up": {"block_0": 4, "block_1": [5, 6, 7]}} - blocks_with_transformer = {"down": [1, 2], "up": [0, 1]} - transformer_per_block = {"down": 2, "up": 3} - ``` - into - ```python - { - "down.block_1.0": 2, - "down.block_1.1": 2, - "down.block_2.0": 2, - "down.block_2.1": 2, - "mid": 3, - "up.block_0.0": 4, - "up.block_0.1": 4, - "up.block_0.2": 4, - "up.block_1.0": 5, - "up.block_1.1": 6, - "up.block_1.2": 7, - } - ``` - """ - if sorted(blocks_with_transformer.keys()) != ["down", "up"]: - raise ValueError("blocks_with_transformer needs to be a dict with keys `'down' and `'up'`") - - if sorted(transformer_per_block.keys()) != ["down", "up"]: - raise ValueError("transformer_per_block needs to be a dict with keys `'down' and `'up'`") - - if not isinstance(scales, dict): - # don't expand if scales is a single number - return scales - - scales = copy.deepcopy(scales) - - if "mid" not in scales: - scales["mid"] = default_scale - elif isinstance(scales["mid"], list): - if len(scales["mid"]) == 1: - scales["mid"] = scales["mid"][0] - else: - raise ValueError(f"Expected 1 scales for mid, got {len(scales['mid'])}.") - - for updown in ["up", "down"]: - if updown not in scales: - scales[updown] = default_scale - - # eg {"down": 1} to {"down": {"block_1": 1, "block_2": 1}}} - if not isinstance(scales[updown], dict): - scales[updown] = {f"block_{i}": copy.deepcopy(scales[updown]) for i in blocks_with_transformer[updown]} - - # eg {"down": {"block_1": 1}} to {"down": {"block_1": [1, 1]}} - for i in blocks_with_transformer[updown]: - block = f"block_{i}" - # set not assigned blocks to default scale - if block not in scales[updown]: - scales[updown][block] = default_scale - if not isinstance(scales[updown][block], list): - scales[updown][block] = [scales[updown][block] for _ in range(transformer_per_block[updown])] - elif len(scales[updown][block]) == 1: - # a list specifying scale to each masked IP input - scales[updown][block] = scales[updown][block] * transformer_per_block[updown] - elif len(scales[updown][block]) != transformer_per_block[updown]: - raise ValueError( - f"Expected {transformer_per_block[updown]} scales for {updown}.{block}, got {len(scales[updown][block])}." - ) - - # eg {"down": "block_1": [1, 1]}} to {"down.block_1.0": 1, "down.block_1.1": 1} - for i in blocks_with_transformer[updown]: - block = f"block_{i}" - for tf_idx, value in enumerate(scales[updown][block]): - scales[f"{updown}.{block}.{tf_idx}"] = value - - del scales[updown] - - for layer in scales.keys(): - if not any(_translate_into_actual_layer_name(layer) in module for module in state_dict.keys()): - raise ValueError( - f"Can't set lora scale for layer {layer}. It either doesn't exist in this unet or it has no attentions." - ) - - return {_translate_into_actual_layer_name(name): weight for name, weight in scales.items()} diff --git a/venv/lib/python3.10/site-packages/diffusers/loaders/utils.py b/venv/lib/python3.10/site-packages/diffusers/loaders/utils.py deleted file mode 100644 index 2d39e7bfb7d2087fd5b5fa9423f31334d6d3321f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/loaders/utils.py +++ /dev/null @@ -1,59 +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 Dict - -import torch - - -class AttnProcsLayers(torch.nn.Module): - 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())) - self.rev_mapping = {v: k for k, v in enumerate(state_dict.keys())} - - # .processor for unet, .self_attn for text encoder - self.split_keys = [".processor", ".self_attn"] - - # we add a hook to state_dict() and load_state_dict() so that the - # naming fits with `unet.attn_processors` - def map_to(module, state_dict, *args, **kwargs): - new_state_dict = {} - for key, value in state_dict.items(): - num = int(key.split(".")[1]) # 0 is always "layers" - new_key = key.replace(f"layers.{num}", module.mapping[num]) - new_state_dict[new_key] = value - - return new_state_dict - - def remap_key(key, state_dict): - for k in self.split_keys: - if k in key: - return key.split(k)[0] + k - - raise ValueError( - f"There seems to be a problem with the state_dict: {set(state_dict.keys())}. {key} has to have one of {self.split_keys}." - ) - - def map_from(module, state_dict, *args, **kwargs): - all_keys = list(state_dict.keys()) - for key in all_keys: - replace_key = remap_key(key, state_dict) - new_key = key.replace(replace_key, f"layers.{module.rev_mapping[replace_key]}") - state_dict[new_key] = state_dict[key] - del state_dict[key] - - self._register_state_dict_hook(map_to) - self._register_load_state_dict_pre_hook(map_from, with_module=True) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__init__.py b/venv/lib/python3.10/site-packages/diffusers/models/__init__.py deleted file mode 100644 index 73903a62741593ad77e385810d9f51bf3ec69a8c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/__init__.py +++ /dev/null @@ -1,208 +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 TYPE_CHECKING - -from ..utils import ( - DIFFUSERS_SLOW_IMPORT, - _LazyModule, - is_flax_available, - is_torch_available, -) - - -_import_structure = {} - -if is_torch_available(): - _import_structure["adapter"] = ["MultiAdapter", "T2IAdapter"] - _import_structure["auto_model"] = ["AutoModel"] - _import_structure["autoencoders.autoencoder_asym_kl"] = ["AsymmetricAutoencoderKL"] - _import_structure["autoencoders.autoencoder_dc"] = ["AutoencoderDC"] - _import_structure["autoencoders.autoencoder_kl"] = ["AutoencoderKL"] - _import_structure["autoencoders.autoencoder_kl_allegro"] = ["AutoencoderKLAllegro"] - _import_structure["autoencoders.autoencoder_kl_cogvideox"] = ["AutoencoderKLCogVideoX"] - _import_structure["autoencoders.autoencoder_kl_cosmos"] = ["AutoencoderKLCosmos"] - _import_structure["autoencoders.autoencoder_kl_hunyuan_video"] = ["AutoencoderKLHunyuanVideo"] - _import_structure["autoencoders.autoencoder_kl_ltx"] = ["AutoencoderKLLTXVideo"] - _import_structure["autoencoders.autoencoder_kl_magvit"] = ["AutoencoderKLMagvit"] - _import_structure["autoencoders.autoencoder_kl_mochi"] = ["AutoencoderKLMochi"] - _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_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_flux"] = ["FluxControlNetModel", "FluxMultiControlNetModel"] - _import_structure["controlnets.controlnet_hunyuan"] = [ - "HunyuanDiT2DControlNetModel", - "HunyuanDiT2DMultiControlNetModel", - ] - _import_structure["controlnets.controlnet_sana"] = ["SanaControlNetModel"] - _import_structure["controlnets.controlnet_sd3"] = ["SD3ControlNetModel", "SD3MultiControlNetModel"] - _import_structure["controlnets.controlnet_sparsectrl"] = ["SparseControlNetModel"] - _import_structure["controlnets.controlnet_union"] = ["ControlNetUnionModel"] - _import_structure["controlnets.controlnet_xs"] = ["ControlNetXSAdapter", "UNetControlNetXSModel"] - _import_structure["controlnets.multicontrolnet"] = ["MultiControlNetModel"] - _import_structure["controlnets.multicontrolnet_union"] = ["MultiControlNetUnionModel"] - _import_structure["embeddings"] = ["ImageProjection"] - _import_structure["modeling_utils"] = ["ModelMixin"] - _import_structure["transformers.auraflow_transformer_2d"] = ["AuraFlowTransformer2DModel"] - _import_structure["transformers.cogvideox_transformer_3d"] = ["CogVideoXTransformer3DModel"] - _import_structure["transformers.consisid_transformer_3d"] = ["ConsisIDTransformer3DModel"] - _import_structure["transformers.dit_transformer_2d"] = ["DiTTransformer2DModel"] - _import_structure["transformers.dual_transformer_2d"] = ["DualTransformer2DModel"] - _import_structure["transformers.hunyuan_transformer_2d"] = ["HunyuanDiT2DModel"] - _import_structure["transformers.latte_transformer_3d"] = ["LatteTransformer3DModel"] - _import_structure["transformers.lumina_nextdit2d"] = ["LuminaNextDiT2DModel"] - _import_structure["transformers.pixart_transformer_2d"] = ["PixArtTransformer2DModel"] - _import_structure["transformers.prior_transformer"] = ["PriorTransformer"] - _import_structure["transformers.sana_transformer"] = ["SanaTransformer2DModel"] - _import_structure["transformers.stable_audio_transformer"] = ["StableAudioDiTModel"] - _import_structure["transformers.t5_film_transformer"] = ["T5FilmDecoder"] - _import_structure["transformers.transformer_2d"] = ["Transformer2DModel"] - _import_structure["transformers.transformer_allegro"] = ["AllegroTransformer3DModel"] - _import_structure["transformers.transformer_chroma"] = ["ChromaTransformer2DModel"] - _import_structure["transformers.transformer_cogview3plus"] = ["CogView3PlusTransformer2DModel"] - _import_structure["transformers.transformer_cogview4"] = ["CogView4Transformer2DModel"] - _import_structure["transformers.transformer_cosmos"] = ["CosmosTransformer3DModel"] - _import_structure["transformers.transformer_easyanimate"] = ["EasyAnimateTransformer3DModel"] - _import_structure["transformers.transformer_flux"] = ["FluxTransformer2DModel"] - _import_structure["transformers.transformer_hidream_image"] = ["HiDreamImageTransformer2DModel"] - _import_structure["transformers.transformer_hunyuan_video"] = ["HunyuanVideoTransformer3DModel"] - _import_structure["transformers.transformer_hunyuan_video_framepack"] = ["HunyuanVideoFramepackTransformer3DModel"] - _import_structure["transformers.transformer_ltx"] = ["LTXVideoTransformer3DModel"] - _import_structure["transformers.transformer_lumina2"] = ["Lumina2Transformer2DModel"] - _import_structure["transformers.transformer_mochi"] = ["MochiTransformer3DModel"] - _import_structure["transformers.transformer_omnigen"] = ["OmniGenTransformer2DModel"] - _import_structure["transformers.transformer_sd3"] = ["SD3Transformer2DModel"] - _import_structure["transformers.transformer_temporal"] = ["TransformerTemporalModel"] - _import_structure["transformers.transformer_wan"] = ["WanTransformer3DModel"] - _import_structure["transformers.transformer_wan_vace"] = ["WanVACETransformer3DModel"] - _import_structure["unets.unet_1d"] = ["UNet1DModel"] - _import_structure["unets.unet_2d"] = ["UNet2DModel"] - _import_structure["unets.unet_2d_condition"] = ["UNet2DConditionModel"] - _import_structure["unets.unet_3d_condition"] = ["UNet3DConditionModel"] - _import_structure["unets.unet_i2vgen_xl"] = ["I2VGenXLUNet"] - _import_structure["unets.unet_kandinsky3"] = ["Kandinsky3UNet"] - _import_structure["unets.unet_motion_model"] = ["MotionAdapter", "UNetMotionModel"] - _import_structure["unets.unet_spatio_temporal_condition"] = ["UNetSpatioTemporalConditionModel"] - _import_structure["unets.unet_stable_cascade"] = ["StableCascadeUNet"] - _import_structure["unets.uvit_2d"] = ["UVit2DModel"] - -if is_flax_available(): - _import_structure["controlnets.controlnet_flax"] = ["FlaxControlNetModel"] - _import_structure["unets.unet_2d_condition_flax"] = ["FlaxUNet2DConditionModel"] - _import_structure["vae_flax"] = ["FlaxAutoencoderKL"] - - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - if is_torch_available(): - from .adapter import MultiAdapter, T2IAdapter - from .auto_model import AutoModel - from .autoencoders import ( - AsymmetricAutoencoderKL, - AutoencoderDC, - AutoencoderKL, - AutoencoderKLAllegro, - AutoencoderKLCogVideoX, - AutoencoderKLCosmos, - AutoencoderKLHunyuanVideo, - AutoencoderKLLTXVideo, - AutoencoderKLMagvit, - AutoencoderKLMochi, - AutoencoderKLTemporalDecoder, - AutoencoderKLWan, - AutoencoderOobleck, - AutoencoderTiny, - ConsistencyDecoderVAE, - VQModel, - ) - from .cache_utils import CacheMixin - from .controlnets import ( - ControlNetModel, - ControlNetUnionModel, - ControlNetXSAdapter, - FluxControlNetModel, - FluxMultiControlNetModel, - HunyuanDiT2DControlNetModel, - HunyuanDiT2DMultiControlNetModel, - MultiControlNetModel, - MultiControlNetUnionModel, - SanaControlNetModel, - SD3ControlNetModel, - SD3MultiControlNetModel, - SparseControlNetModel, - UNetControlNetXSModel, - ) - from .embeddings import ImageProjection - from .modeling_utils import ModelMixin - from .transformers import ( - AllegroTransformer3DModel, - AuraFlowTransformer2DModel, - ChromaTransformer2DModel, - CogVideoXTransformer3DModel, - CogView3PlusTransformer2DModel, - CogView4Transformer2DModel, - ConsisIDTransformer3DModel, - CosmosTransformer3DModel, - DiTTransformer2DModel, - DualTransformer2DModel, - EasyAnimateTransformer3DModel, - FluxTransformer2DModel, - HiDreamImageTransformer2DModel, - HunyuanDiT2DModel, - HunyuanVideoFramepackTransformer3DModel, - HunyuanVideoTransformer3DModel, - LatteTransformer3DModel, - LTXVideoTransformer3DModel, - Lumina2Transformer2DModel, - LuminaNextDiT2DModel, - MochiTransformer3DModel, - OmniGenTransformer2DModel, - PixArtTransformer2DModel, - PriorTransformer, - SanaTransformer2DModel, - SD3Transformer2DModel, - StableAudioDiTModel, - T5FilmDecoder, - Transformer2DModel, - TransformerTemporalModel, - WanTransformer3DModel, - WanVACETransformer3DModel, - ) - from .unets import ( - I2VGenXLUNet, - Kandinsky3UNet, - MotionAdapter, - StableCascadeUNet, - UNet1DModel, - UNet2DConditionModel, - UNet2DModel, - UNet3DConditionModel, - UNetMotionModel, - UNetSpatioTemporalConditionModel, - UVit2DModel, - ) - - if is_flax_available(): - from .controlnets import FlaxControlNetModel - from .unets import FlaxUNet2DConditionModel - from .vae_flax import FlaxAutoencoderKL - -else: - import sys - - sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index fb077d55132377877d81cadffa75f97990be1c34..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/activations.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/activations.cpython-310.pyc deleted file mode 100644 index 77f590a1db4ae340755d2b4b6e148aabfd7a4752..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/activations.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/adapter.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/adapter.cpython-310.pyc deleted file mode 100644 index 5b87b5504a9abe1be7f237b18077c835fd40825d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/adapter.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/attention.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/attention.cpython-310.pyc deleted file mode 100644 index 72ab8ff033fb919641bd6015ccb2f79739af74ef..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/attention.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/attention_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/attention_flax.cpython-310.pyc deleted file mode 100644 index 44c16b71bb07e11a13dbbff19618c27dd208babe..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/attention_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/attention_processor.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/attention_processor.cpython-310.pyc deleted file mode 100644 index 42873018b06d1b48ec7a75c01c51c79fced31e1b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/attention_processor.cpython-310.pyc +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:3e969a49ca989c663cf1d81a98abc89c333c7cedc2e6cf1f3ef724084b7d4965 -size 125754 diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/auto_model.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/auto_model.cpython-310.pyc deleted file mode 100644 index b310b4c6b3364dd4f22df60091d8cbfa3136e007..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/auto_model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/cache_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/cache_utils.cpython-310.pyc deleted file mode 100644 index 8b59895b53eca0c1aef1465130076cf8f032f463..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/cache_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/controlnet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/controlnet.cpython-310.pyc deleted file mode 100644 index 87a38ed5d95859480e5e913fc18d27140959f4ff..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/controlnet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/controlnet_flux.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/controlnet_flux.cpython-310.pyc deleted file mode 100644 index 0c05341e43c17c37142aba3c3bce7bd110720965..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/controlnet_flux.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/controlnet_sd3.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/controlnet_sd3.cpython-310.pyc deleted file mode 100644 index 25645f959edde499d77d235637f34aa04f0762e9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/controlnet_sd3.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/controlnet_sparsectrl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/controlnet_sparsectrl.cpython-310.pyc deleted file mode 100644 index 7b71f2470206d59a858bf169e4b3995fe35c9c84..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/controlnet_sparsectrl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/downsampling.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/downsampling.cpython-310.pyc deleted file mode 100644 index 73607cb7e2ca9a9f58493a87c7bb4bd67958b733..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/downsampling.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/embeddings.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/embeddings.cpython-310.pyc deleted file mode 100644 index 0bf23db4b82ea90f015beb6d961a7b2ed7946264..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/embeddings.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/embeddings_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/embeddings_flax.cpython-310.pyc deleted file mode 100644 index 459e16d89b396790dbe1e3a7fe769a024ffec2be..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/embeddings_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/lora.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/lora.cpython-310.pyc deleted file mode 100644 index e31587e32d4b8e09b0aecf40382f0f3e9de914ed..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/lora.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/model_loading_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/model_loading_utils.cpython-310.pyc deleted file mode 100644 index 492ccc8753df892934c6e50d40cdaff538c85649..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/model_loading_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_flax_pytorch_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_flax_pytorch_utils.cpython-310.pyc deleted file mode 100644 index 46b2154b8ff42da62f9c507c1af5789cd7d62105..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_flax_pytorch_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_flax_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_flax_utils.cpython-310.pyc deleted file mode 100644 index 9d6c1dd02271776b21afe627d49a94458b463c6e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_flax_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_outputs.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_outputs.cpython-310.pyc deleted file mode 100644 index a181c7c18ef21c1002d4ecffe5c9621f1cfcb251..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_outputs.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_pytorch_flax_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_pytorch_flax_utils.cpython-310.pyc deleted file mode 100644 index af1ae495e4b4dab28b9018f8defdd0145fe1b676..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_pytorch_flax_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_utils.cpython-310.pyc deleted file mode 100644 index 4194ec021d950d3513d27164b3c2e90098c8299f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/modeling_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/normalization.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/normalization.cpython-310.pyc deleted file mode 100644 index 88ed711387fa6b14d8545235ab3fbd514fd21ed6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/normalization.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/resnet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/resnet.cpython-310.pyc deleted file mode 100644 index 08f8de4cfdb3bf1c02bfb4a8563ad39a0429d315..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/resnet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/resnet_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/resnet_flax.cpython-310.pyc deleted file mode 100644 index e03829e6d089c716d31b09c10f746b25a2db939c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/resnet_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/upsampling.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/upsampling.cpython-310.pyc deleted file mode 100644 index 5f08169872b008d90a9f785d53a0a785146dfaee..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/upsampling.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/vae_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/vae_flax.cpython-310.pyc deleted file mode 100644 index 49c0370b7224dba5d483e791aae33d11d9649eb1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/vae_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/vq_model.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/vq_model.cpython-310.pyc deleted file mode 100644 index 138c9e2b6602028d0889b8e309dd3c64d5c5a689..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/__pycache__/vq_model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/activations.py b/venv/lib/python3.10/site-packages/diffusers/models/activations.py deleted file mode 100644 index 2d1fdb5f7d8303ae605afe4c6905af38950e4acb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/activations.py +++ /dev/null @@ -1,178 +0,0 @@ -# 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 -import torch.nn.functional as F -from torch import nn - -from ..utils import deprecate -from ..utils.import_utils import is_torch_npu_available, is_torch_version - - -if is_torch_npu_available(): - import torch_npu - -ACT2CLS = { - "swish": nn.SiLU, - "silu": nn.SiLU, - "mish": nn.Mish, - "gelu": nn.GELU, - "relu": nn.ReLU, -} - - -def get_activation(act_fn: str) -> nn.Module: - """Helper function to get activation function from string. - - Args: - act_fn (str): Name of activation function. - - Returns: - nn.Module: Activation function. - """ - - act_fn = act_fn.lower() - if act_fn in ACT2CLS: - return ACT2CLS[act_fn]() - else: - raise ValueError(f"activation function {act_fn} not found in ACT2FN mapping {list(ACT2CLS.keys())}") - - -class FP32SiLU(nn.Module): - r""" - SiLU activation function with input upcasted to torch.float32. - """ - - def __init__(self): - super().__init__() - - def forward(self, inputs: torch.Tensor) -> torch.Tensor: - return F.silu(inputs.float(), inplace=False).to(inputs.dtype) - - -class GELU(nn.Module): - r""" - GELU activation function with tanh approximation support with `approximate="tanh"`. - - Parameters: - dim_in (`int`): The number of channels in the input. - dim_out (`int`): The number of channels in the output. - approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation. - bias (`bool`, defaults to True): Whether to use a bias in the linear layer. - """ - - def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True): - super().__init__() - self.proj = nn.Linear(dim_in, dim_out, bias=bias) - self.approximate = approximate - - def gelu(self, gate: torch.Tensor) -> torch.Tensor: - if gate.device.type == "mps" and is_torch_version("<", "2.0.0"): - # fp16 gelu not supported on mps before torch 2.0 - return F.gelu(gate.to(dtype=torch.float32), approximate=self.approximate).to(dtype=gate.dtype) - return F.gelu(gate, approximate=self.approximate) - - def forward(self, hidden_states): - hidden_states = self.proj(hidden_states) - hidden_states = self.gelu(hidden_states) - return hidden_states - - -class GEGLU(nn.Module): - r""" - A [variant](https://huggingface.co/papers/2002.05202) of the gated linear unit activation function. - - Parameters: - dim_in (`int`): The number of channels in the input. - dim_out (`int`): The number of channels in the output. - bias (`bool`, defaults to True): Whether to use a bias in the linear layer. - """ - - def __init__(self, dim_in: int, dim_out: int, bias: bool = True): - super().__init__() - self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias) - - def gelu(self, gate: torch.Tensor) -> torch.Tensor: - if gate.device.type == "mps" and is_torch_version("<", "2.0.0"): - # fp16 gelu not supported on mps before torch 2.0 - return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype) - return F.gelu(gate) - - def forward(self, hidden_states, *args, **kwargs): - 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) - hidden_states = self.proj(hidden_states) - if is_torch_npu_available(): - # using torch_npu.npu_geglu can run faster and save memory on NPU. - return torch_npu.npu_geglu(hidden_states, dim=-1, approximate=1)[0] - else: - hidden_states, gate = hidden_states.chunk(2, dim=-1) - return hidden_states * self.gelu(gate) - - -class SwiGLU(nn.Module): - r""" - A [variant](https://huggingface.co/papers/2002.05202) of the gated linear unit activation function. It's similar to - `GEGLU` but uses SiLU / Swish instead of GeLU. - - Parameters: - dim_in (`int`): The number of channels in the input. - dim_out (`int`): The number of channels in the output. - bias (`bool`, defaults to True): Whether to use a bias in the linear layer. - """ - - def __init__(self, dim_in: int, dim_out: int, bias: bool = True): - super().__init__() - - self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias) - self.activation = nn.SiLU() - - def forward(self, hidden_states): - hidden_states = self.proj(hidden_states) - hidden_states, gate = hidden_states.chunk(2, dim=-1) - return hidden_states * self.activation(gate) - - -class ApproximateGELU(nn.Module): - r""" - The approximate form of the Gaussian Error Linear Unit (GELU). For more details, see section 2 of this - [paper](https://huggingface.co/papers/1606.08415). - - Parameters: - dim_in (`int`): The number of channels in the input. - dim_out (`int`): The number of channels in the output. - bias (`bool`, defaults to True): Whether to use a bias in the linear layer. - """ - - def __init__(self, dim_in: int, dim_out: int, bias: bool = True): - super().__init__() - self.proj = nn.Linear(dim_in, dim_out, bias=bias) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.proj(x) - return x * torch.sigmoid(1.702 * x) - - -class LinearActivation(nn.Module): - def __init__(self, dim_in: int, dim_out: int, bias: bool = True, activation: str = "silu"): - super().__init__() - - self.proj = nn.Linear(dim_in, dim_out, bias=bias) - self.activation = get_activation(activation) - - def forward(self, hidden_states): - hidden_states = self.proj(hidden_states) - return self.activation(hidden_states) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/adapter.py b/venv/lib/python3.10/site-packages/diffusers/models/adapter.py deleted file mode 100644 index e475fe6bee88e791b8b377ce01bb4c38a0badf1e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/adapter.py +++ /dev/null @@ -1,583 +0,0 @@ -# Copyright 2022 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 os -from typing import Callable, List, Optional, Union - -import torch -import torch.nn as nn - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import logging -from .modeling_utils import ModelMixin - - -logger = logging.get_logger(__name__) - - -class MultiAdapter(ModelMixin): - r""" - MultiAdapter is a wrapper model that contains multiple adapter models and merges their outputs according to - user-assigned weighting. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for common methods such as downloading - or saving. - - Args: - adapters (`List[T2IAdapter]`, *optional*, defaults to None): - A list of `T2IAdapter` model instances. - """ - - def __init__(self, adapters: List["T2IAdapter"]): - super(MultiAdapter, self).__init__() - - self.num_adapter = len(adapters) - self.adapters = nn.ModuleList(adapters) - - if len(adapters) == 0: - raise ValueError("Expecting at least one adapter") - - if len(adapters) == 1: - raise ValueError("For a single adapter, please use the `T2IAdapter` class instead of `MultiAdapter`") - - # The outputs from each adapter are added together with a weight. - # This means that the change in dimensions from downsampling must - # be the same for all adapters. Inductively, it also means the - # downscale_factor and total_downscale_factor must be the same for all - # adapters. - first_adapter_total_downscale_factor = adapters[0].total_downscale_factor - first_adapter_downscale_factor = adapters[0].downscale_factor - for idx in range(1, len(adapters)): - if ( - adapters[idx].total_downscale_factor != first_adapter_total_downscale_factor - or adapters[idx].downscale_factor != first_adapter_downscale_factor - ): - raise ValueError( - f"Expecting all adapters to have the same downscaling behavior, but got:\n" - f"adapters[0].total_downscale_factor={first_adapter_total_downscale_factor}\n" - f"adapters[0].downscale_factor={first_adapter_downscale_factor}\n" - f"adapter[`{idx}`].total_downscale_factor={adapters[idx].total_downscale_factor}\n" - f"adapter[`{idx}`].downscale_factor={adapters[idx].downscale_factor}" - ) - - 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]: - r""" - Args: - xs (`torch.Tensor`): - A tensor of shape (batch, channel, height, width) representing input images for multiple adapter - 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): - 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. - """ - if adapter_weights is None: - adapter_weights = torch.tensor([1 / self.num_adapter] * self.num_adapter) - else: - adapter_weights = torch.tensor(adapter_weights) - - accume_state = None - for x, w, adapter in zip(xs, adapter_weights, self.adapters): - features = adapter(x) - if accume_state is None: - accume_state = features - for i in range(len(accume_state)): - accume_state[i] = w * accume_state[i] - else: - for i in range(len(features)): - accume_state[i] += w * features[i] - return accume_state - - def save_pretrained( - self, - save_directory: Union[str, os.PathLike], - is_main_process: bool = True, - save_function: Callable = None, - safe_serialization: bool = True, - variant: Optional[str] = None, - ): - """ - Save a model and its configuration file to a specified directory, allowing it to be re-loaded with the - `[`~models.adapter.MultiAdapter.from_pretrained`]` class method. - - Args: - save_directory (`str` or `os.PathLike`): - The directory where the model will be saved. If the directory does not exist, it will be created. - is_main_process (`bool`, optional, defaults=True): - Indicates whether current process is the main process or not. Useful for distributed training (e.g., - TPUs) and need to call this function on all processes. In this case, set `is_main_process=True` only - for the main process to avoid race conditions. - save_function (`Callable`): - Function used to save the state dictionary. Useful for distributed training (e.g., TPUs) to replace - `torch.save` with another method. Can also be configured using`DIFFUSERS_SAVE_MODE` environment - variable. - safe_serialization (`bool`, optional, defaults=True): - If `True`, save the model using `safetensors`. If `False`, save the model with `pickle`. - variant (`str`, *optional*): - If specified, weights are saved in the format `pytorch_model..bin`. - """ - idx = 0 - model_path_to_save = save_directory - for adapter in self.adapters: - adapter.save_pretrained( - model_path_to_save, - is_main_process=is_main_process, - save_function=save_function, - safe_serialization=safe_serialization, - variant=variant, - ) - - idx += 1 - model_path_to_save = model_path_to_save + f"_{idx}" - - @classmethod - def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os.PathLike]], **kwargs): - r""" - Instantiate a pretrained `MultiAdapter` model from multiple pre-trained adapter models. - - The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train - the model, set it back to training mode using `model.train()`. - - Warnings: - *Weights from XXX not initialized from pretrained model* means that the weights of XXX are not pretrained - with the rest of the model. It is up to you to train those weights with a downstream fine-tuning. *Weights - from XXX not used in YYY* means that the layer XXX is not used by YYY, so those weights are discarded. - - Args: - pretrained_model_path (`os.PathLike`): - A path to a *directory* containing model weights saved using - [`~diffusers.models.adapter.MultiAdapter.save_pretrained`], e.g., `./my_model_directory/adapter`. - torch_dtype (`torch.dtype`, *optional*): - 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*): - 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. - - To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For - more information about each option see [designing a device - map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). - max_memory (`Dict`, *optional*): - A dictionary mapping device identifiers to their maximum memory. Default to the maximum memory - available for each GPU and the available CPU RAM if unset. - low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): - Speed up model loading by not initializing the weights and only loading the pre-trained weights. This - also tries to not use more than 1x model size in CPU memory (including peak memory) while loading the - model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch, - setting this argument to `True` will raise an error. - variant (`str`, *optional*): - If specified, load weights from a `variant` file (*e.g.* pytorch_model..bin). `variant` will - be ignored when using `from_flax`. - use_safetensors (`bool`, *optional*, defaults to `None`): - If `None`, the `safetensors` weights will be downloaded if available **and** if`safetensors` library is - installed. If `True`, the model will be forcibly loaded from`safetensors` weights. If `False`, - `safetensors` is not used. - """ - idx = 0 - adapters = [] - - # load adapter and append to list until no adapter directory exists anymore - # first adapter has to be saved under `./mydirectory/adapter` to be compliant with `DiffusionPipeline.from_pretrained` - # second, third, ... adapters have to be saved under `./mydirectory/adapter_1`, `./mydirectory/adapter_2`, ... - model_path_to_load = pretrained_model_path - while os.path.isdir(model_path_to_load): - adapter = T2IAdapter.from_pretrained(model_path_to_load, **kwargs) - adapters.append(adapter) - - idx += 1 - model_path_to_load = pretrained_model_path + f"_{idx}" - - logger.info(f"{len(adapters)} adapters loaded from {pretrained_model_path}.") - - if len(adapters) == 0: - raise ValueError( - f"No T2IAdapters found under {os.path.dirname(pretrained_model_path)}. Expected at least {pretrained_model_path + '_0'}." - ) - - return cls(adapters) - - -class T2IAdapter(ModelMixin, ConfigMixin): - r""" - A simple ResNet-like model that accepts images containing control signals such as keyposes and depth. The model - generates multiple feature maps that are used as additional conditioning in [`UNet2DConditionModel`]. The model's - architecture follows the original implementation of - [Adapter](https://github.com/TencentARC/T2I-Adapter/blob/686de4681515662c0ac2ffa07bf5dda83af1038a/ldm/modules/encoders/adapter.py#L97) - and - [AdapterLight](https://github.com/TencentARC/T2I-Adapter/blob/686de4681515662c0ac2ffa07bf5dda83af1038a/ldm/modules/encoders/adapter.py#L235). - - This model inherits from [`ModelMixin`]. Check the superclass documentation for the common methods, such as - downloading or saving. - - Args: - 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)`): - 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`): - Number of ResNet blocks in each downsample block. - downscale_factor (`int`, *optional*, defaults to `8`): - A factor that determines the total downscale factor of the Adapter. - adapter_type (`str`, *optional*, defaults to `full_adapter`): - Adapter type (`full_adapter` or `full_adapter_xl` or `light_adapter`) to use. - """ - - @register_to_config - def __init__( - self, - in_channels: int = 3, - channels: List[int] = [320, 640, 1280, 1280], - num_res_blocks: int = 2, - downscale_factor: int = 8, - adapter_type: str = "full_adapter", - ): - super().__init__() - - if adapter_type == "full_adapter": - self.adapter = FullAdapter(in_channels, channels, num_res_blocks, downscale_factor) - elif adapter_type == "full_adapter_xl": - self.adapter = FullAdapterXL(in_channels, channels, num_res_blocks, downscale_factor) - elif adapter_type == "light_adapter": - self.adapter = LightAdapter(in_channels, channels, num_res_blocks, downscale_factor) - else: - raise ValueError( - f"Unsupported adapter_type: '{adapter_type}'. Choose either 'full_adapter' or " - "'full_adapter_xl' or 'light_adapter'." - ) - - 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 - determined by the number of downsample blocks in the Adapter, as specified by the `channels` and - `num_res_blocks` parameters during initialization. - """ - return self.adapter(x) - - @property - def total_downscale_factor(self): - return self.adapter.total_downscale_factor - - @property - def downscale_factor(self): - """The downscale factor applied in the T2I-Adapter's initial pixel unshuffle operation. If an input image's dimensions are - not evenly divisible by the downscale_factor then an exception will be raised. - """ - return self.adapter.unshuffle.downscale_factor - - -# full adapter - - -class FullAdapter(nn.Module): - r""" - See [`T2IAdapter`] for more information. - """ - - def __init__( - self, - in_channels: int = 3, - channels: List[int] = [320, 640, 1280, 1280], - num_res_blocks: int = 2, - downscale_factor: int = 8, - ): - super().__init__() - - in_channels = in_channels * downscale_factor**2 - - self.unshuffle = nn.PixelUnshuffle(downscale_factor) - self.conv_in = nn.Conv2d(in_channels, channels[0], kernel_size=3, padding=1) - - self.body = nn.ModuleList( - [ - AdapterBlock(channels[0], channels[0], num_res_blocks), - *[ - AdapterBlock(channels[i - 1], channels[i], num_res_blocks, down=True) - for i in range(1, len(channels)) - ], - ] - ) - - self.total_downscale_factor = downscale_factor * 2 ** (len(channels) - 1) - - 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 - capturing information at a different stage of processing within the FullAdapter model. The number of feature - tensors in the list is determined by the number of downsample blocks specified during initialization. - """ - x = self.unshuffle(x) - x = self.conv_in(x) - - features = [] - - for block in self.body: - x = block(x) - features.append(x) - - return features - - -class FullAdapterXL(nn.Module): - r""" - See [`T2IAdapter`] for more information. - """ - - def __init__( - self, - in_channels: int = 3, - channels: List[int] = [320, 640, 1280, 1280], - num_res_blocks: int = 2, - downscale_factor: int = 16, - ): - super().__init__() - - in_channels = in_channels * downscale_factor**2 - - self.unshuffle = nn.PixelUnshuffle(downscale_factor) - self.conv_in = nn.Conv2d(in_channels, channels[0], kernel_size=3, padding=1) - - self.body = [] - # blocks to extract XL features with dimensions of [320, 64, 64], [640, 64, 64], [1280, 32, 32], [1280, 32, 32] - for i in range(len(channels)): - if i == 1: - self.body.append(AdapterBlock(channels[i - 1], channels[i], num_res_blocks)) - elif i == 2: - self.body.append(AdapterBlock(channels[i - 1], channels[i], num_res_blocks, down=True)) - else: - self.body.append(AdapterBlock(channels[i], channels[i], num_res_blocks)) - - self.body = nn.ModuleList(self.body) - # XL has only one downsampling AdapterBlock. - self.total_downscale_factor = downscale_factor * 2 - - 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. - """ - x = self.unshuffle(x) - x = self.conv_in(x) - - features = [] - - for block in self.body: - x = block(x) - features.append(x) - - return features - - -class AdapterBlock(nn.Module): - r""" - An AdapterBlock is a helper model that contains multiple ResNet-like blocks. It is used in the `FullAdapter` and - `FullAdapterXL` models. - - Args: - in_channels (`int`): - Number of channels of AdapterBlock's input. - out_channels (`int`): - Number of channels of AdapterBlock's output. - num_res_blocks (`int`): - Number of ResNet blocks in the AdapterBlock. - down (`bool`, *optional*, defaults to `False`): - If `True`, perform downsampling on AdapterBlock's input. - """ - - def __init__(self, in_channels: int, out_channels: int, num_res_blocks: int, down: bool = False): - super().__init__() - - self.downsample = None - if down: - self.downsample = nn.AvgPool2d(kernel_size=2, stride=2, ceil_mode=True) - - self.in_conv = None - if in_channels != out_channels: - self.in_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1) - - self.resnets = nn.Sequential( - *[AdapterResnetBlock(out_channels) for _ in range(num_res_blocks)], - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - r""" - This method takes tensor x as input and performs operations downsampling and convolutional layers if the - self.downsample and self.in_conv properties of AdapterBlock model are specified. Then it applies a series of - residual blocks to the input tensor. - """ - if self.downsample is not None: - x = self.downsample(x) - - if self.in_conv is not None: - x = self.in_conv(x) - - x = self.resnets(x) - - return x - - -class AdapterResnetBlock(nn.Module): - r""" - An `AdapterResnetBlock` is a helper model that implements a ResNet-like block. - - Args: - channels (`int`): - Number of channels of AdapterResnetBlock's input and output. - """ - - def __init__(self, channels: int): - super().__init__() - self.block1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1) - self.act = nn.ReLU() - self.block2 = nn.Conv2d(channels, channels, kernel_size=1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - r""" - This method takes input tensor x and applies a convolutional layer, ReLU activation, and another convolutional - layer on the input tensor. It returns addition with the input tensor. - """ - - h = self.act(self.block1(x)) - h = self.block2(h) - - return h + x - - -# light adapter - - -class LightAdapter(nn.Module): - r""" - See [`T2IAdapter`] for more information. - """ - - def __init__( - self, - in_channels: int = 3, - channels: List[int] = [320, 640, 1280], - num_res_blocks: int = 4, - downscale_factor: int = 8, - ): - super().__init__() - - in_channels = in_channels * downscale_factor**2 - - self.unshuffle = nn.PixelUnshuffle(downscale_factor) - - self.body = nn.ModuleList( - [ - LightAdapterBlock(in_channels, channels[0], num_res_blocks), - *[ - LightAdapterBlock(channels[i], channels[i + 1], num_res_blocks, down=True) - for i in range(len(channels) - 1) - ], - LightAdapterBlock(channels[-1], channels[-1], num_res_blocks, down=True), - ] - ) - - self.total_downscale_factor = downscale_factor * (2 ** len(channels)) - - 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. - """ - x = self.unshuffle(x) - - features = [] - - for block in self.body: - x = block(x) - features.append(x) - - return features - - -class LightAdapterBlock(nn.Module): - r""" - A `LightAdapterBlock` is a helper model that contains multiple `LightAdapterResnetBlocks`. It is used in the - `LightAdapter` model. - - Args: - in_channels (`int`): - Number of channels of LightAdapterBlock's input. - out_channels (`int`): - Number of channels of LightAdapterBlock's output. - num_res_blocks (`int`): - Number of LightAdapterResnetBlocks in the LightAdapterBlock. - down (`bool`, *optional*, defaults to `False`): - If `True`, perform downsampling on LightAdapterBlock's input. - """ - - def __init__(self, in_channels: int, out_channels: int, num_res_blocks: int, down: bool = False): - super().__init__() - mid_channels = out_channels // 4 - - self.downsample = None - if down: - self.downsample = nn.AvgPool2d(kernel_size=2, stride=2, ceil_mode=True) - - self.in_conv = nn.Conv2d(in_channels, mid_channels, kernel_size=1) - self.resnets = nn.Sequential(*[LightAdapterResnetBlock(mid_channels) for _ in range(num_res_blocks)]) - self.out_conv = nn.Conv2d(mid_channels, out_channels, kernel_size=1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - r""" - This method takes tensor x as input and performs downsampling if required. Then it applies in convolution - layer, a sequence of residual blocks, and out convolutional layer. - """ - if self.downsample is not None: - x = self.downsample(x) - - x = self.in_conv(x) - x = self.resnets(x) - x = self.out_conv(x) - - return x - - -class LightAdapterResnetBlock(nn.Module): - """ - A `LightAdapterResnetBlock` is a helper model that implements a ResNet-like block with a slightly different - architecture than `AdapterResnetBlock`. - - Args: - channels (`int`): - Number of channels of LightAdapterResnetBlock's input and output. - """ - - def __init__(self, channels: int): - super().__init__() - self.block1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1) - self.act = nn.ReLU() - self.block2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - r""" - This function takes input tensor x and processes it through one convolutional layer, ReLU activation, and - another convolutional layer and adds it to input tensor. - """ - - h = self.act(self.block1(x)) - h = self.block2(h) - - return h + x diff --git a/venv/lib/python3.10/site-packages/diffusers/models/attention.py b/venv/lib/python3.10/site-packages/diffusers/models/attention.py deleted file mode 100644 index ae51d3ab134912b01288ff9337d7b13ba43a1127..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/attention.py +++ /dev/null @@ -1,1251 +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 Any, Dict, List, Optional, Tuple - -import torch -import torch.nn.functional as F -from torch import nn - -from ..utils import deprecate, logging -from ..utils.torch_utils import maybe_allow_in_graph -from .activations import GEGLU, GELU, ApproximateGELU, FP32SiLU, LinearActivation, SwiGLU -from .attention_processor import Attention, JointAttnProcessor2_0 -from .embeddings import SinusoidalPositionalEmbedding -from .normalization import AdaLayerNorm, AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm, SD35AdaLayerNormZeroX - - -logger = logging.get_logger(__name__) - - -def _chunked_feed_forward(ff: nn.Module, hidden_states: torch.Tensor, chunk_dim: int, chunk_size: int): - # "feed_forward_chunk_size" can be used to save memory - if hidden_states.shape[chunk_dim] % chunk_size != 0: - raise ValueError( - f"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]} has to be divisible by chunk size: {chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." - ) - - num_chunks = hidden_states.shape[chunk_dim] // chunk_size - ff_output = torch.cat( - [ff(hid_slice) for hid_slice in hidden_states.chunk(num_chunks, dim=chunk_dim)], - dim=chunk_dim, - ) - return ff_output - - -@maybe_allow_in_graph -class GatedSelfAttentionDense(nn.Module): - r""" - A gated self-attention dense layer that combines visual features and object features. - - Parameters: - query_dim (`int`): The number of channels in the query. - context_dim (`int`): The number of channels in the context. - n_heads (`int`): The number of heads to use for attention. - d_head (`int`): The number of channels in each head. - """ - - def __init__(self, query_dim: int, context_dim: int, n_heads: int, d_head: int): - super().__init__() - - # we need a linear projection since we need cat visual feature and obj feature - self.linear = nn.Linear(context_dim, query_dim) - - self.attn = Attention(query_dim=query_dim, heads=n_heads, dim_head=d_head) - self.ff = FeedForward(query_dim, activation_fn="geglu") - - self.norm1 = nn.LayerNorm(query_dim) - self.norm2 = nn.LayerNorm(query_dim) - - self.register_parameter("alpha_attn", nn.Parameter(torch.tensor(0.0))) - self.register_parameter("alpha_dense", nn.Parameter(torch.tensor(0.0))) - - self.enabled = True - - def forward(self, x: torch.Tensor, objs: torch.Tensor) -> torch.Tensor: - if not self.enabled: - return x - - n_visual = x.shape[1] - objs = self.linear(objs) - - x = x + self.alpha_attn.tanh() * self.attn(self.norm1(torch.cat([x, objs], dim=1)))[:, :n_visual, :] - x = x + self.alpha_dense.tanh() * self.ff(self.norm2(x)) - - return x - - -@maybe_allow_in_graph -class JointTransformerBlock(nn.Module): - r""" - A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3. - - Reference: https://huggingface.co/papers/2403.03206 - - Parameters: - 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. - context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the - processing of `context` conditions. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - context_pre_only: bool = False, - qk_norm: Optional[str] = None, - use_dual_attention: bool = False, - ): - super().__init__() - - self.use_dual_attention = use_dual_attention - self.context_pre_only = context_pre_only - context_norm_type = "ada_norm_continous" if context_pre_only else "ada_norm_zero" - - if use_dual_attention: - self.norm1 = SD35AdaLayerNormZeroX(dim) - else: - self.norm1 = AdaLayerNormZero(dim) - - if context_norm_type == "ada_norm_continous": - self.norm1_context = AdaLayerNormContinuous( - dim, dim, elementwise_affine=False, eps=1e-6, bias=True, norm_type="layer_norm" - ) - elif context_norm_type == "ada_norm_zero": - self.norm1_context = AdaLayerNormZero(dim) - else: - raise ValueError( - f"Unknown context_norm_type: {context_norm_type}, currently only support `ada_norm_continous`, `ada_norm_zero`" - ) - - if hasattr(F, "scaled_dot_product_attention"): - processor = JointAttnProcessor2_0() - else: - raise ValueError( - "The current PyTorch version does not support the `scaled_dot_product_attention` function." - ) - - self.attn = Attention( - query_dim=dim, - cross_attention_dim=None, - added_kv_proj_dim=dim, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=dim, - context_pre_only=context_pre_only, - bias=True, - processor=processor, - qk_norm=qk_norm, - eps=1e-6, - ) - - if use_dual_attention: - self.attn2 = Attention( - query_dim=dim, - cross_attention_dim=None, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=dim, - bias=True, - processor=processor, - qk_norm=qk_norm, - eps=1e-6, - ) - else: - self.attn2 = None - - self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) - self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") - - if not context_pre_only: - self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) - self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") - else: - self.norm2_context = None - self.ff_context = None - - # let chunk size default to None - self._chunk_size = None - 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): - # Sets chunk feed-forward - self._chunk_size = chunk_size - self._chunk_dim = dim - - def forward( - self, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor, - temb: torch.FloatTensor, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - ): - 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( - hidden_states, emb=temb - ) - else: - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) - - if self.context_pre_only: - norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, temb) - else: - norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( - encoder_hidden_states, emb=temb - ) - - # Attention. - attn_output, context_attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - **joint_attention_kwargs, - ) - - # Process attention outputs for the `hidden_states`. - attn_output = gate_msa.unsqueeze(1) * attn_output - hidden_states = hidden_states + attn_output - - if self.use_dual_attention: - attn_output2 = self.attn2(hidden_states=norm_hidden_states2, **joint_attention_kwargs) - attn_output2 = gate_msa2.unsqueeze(1) * attn_output2 - hidden_states = hidden_states + attn_output2 - - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - if self._chunk_size is not None: - # "feed_forward_chunk_size" can be used to save memory - ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size) - else: - ff_output = self.ff(norm_hidden_states) - ff_output = gate_mlp.unsqueeze(1) * ff_output - - hidden_states = hidden_states + ff_output - - # Process attention outputs for the `encoder_hidden_states`. - if self.context_pre_only: - encoder_hidden_states = None - else: - context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output - encoder_hidden_states = encoder_hidden_states + context_attn_output - - norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) - norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] - if self._chunk_size is not None: - # "feed_forward_chunk_size" can be used to save memory - context_ff_output = _chunked_feed_forward( - self.ff_context, norm_encoder_hidden_states, self._chunk_dim, self._chunk_size - ) - else: - context_ff_output = self.ff_context(norm_encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output - - return encoder_hidden_states, hidden_states - - -@maybe_allow_in_graph -class BasicTransformerBlock(nn.Module): - r""" - A basic Transformer block. - - Parameters: - 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. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. - activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. - num_embeds_ada_norm (: - obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. - attention_bias (: - obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. - only_cross_attention (`bool`, *optional*): - Whether to use only cross-attention layers. In this case two cross attention layers are used. - double_self_attention (`bool`, *optional*): - Whether to use two self-attention layers. In this case no cross attention layers are used. - upcast_attention (`bool`, *optional*): - Whether to upcast the attention computation to float32. This is useful for mixed precision training. - norm_elementwise_affine (`bool`, *optional*, defaults to `True`): - Whether to use learnable elementwise affine parameters for normalization. - norm_type (`str`, *optional*, defaults to `"layer_norm"`): - The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`. - final_dropout (`bool` *optional*, defaults to False): - Whether to apply a final dropout after the last feed-forward layer. - attention_type (`str`, *optional*, defaults to `"default"`): - The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`. - positional_embeddings (`str`, *optional*, defaults to `None`): - The type of positional embeddings to apply to. - num_positional_embeddings (`int`, *optional*, defaults to `None`): - The maximum number of positional embeddings to apply. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - dropout=0.0, - cross_attention_dim: Optional[int] = None, - activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, - attention_bias: bool = False, - only_cross_attention: bool = False, - double_self_attention: bool = False, - upcast_attention: bool = False, - norm_elementwise_affine: bool = True, - norm_type: str = "layer_norm", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen' - 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, - ff_bias: bool = True, - attention_out_bias: bool = True, - ): - super().__init__() - self.dim = dim - self.num_attention_heads = num_attention_heads - self.attention_head_dim = attention_head_dim - self.dropout = dropout - self.cross_attention_dim = cross_attention_dim - self.activation_fn = activation_fn - self.attention_bias = attention_bias - self.double_self_attention = double_self_attention - self.norm_elementwise_affine = norm_elementwise_affine - self.positional_embeddings = positional_embeddings - self.num_positional_embeddings = num_positional_embeddings - self.only_cross_attention = only_cross_attention - - # We keep these boolean flags for backward-compatibility. - self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" - self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" - self.use_ada_layer_norm_single = norm_type == "ada_norm_single" - self.use_layer_norm = norm_type == "layer_norm" - self.use_ada_layer_norm_continuous = norm_type == "ada_norm_continuous" - - if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: - raise ValueError( - f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" - f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." - ) - - self.norm_type = norm_type - self.num_embeds_ada_norm = num_embeds_ada_norm - - if positional_embeddings and (num_positional_embeddings is None): - raise ValueError( - "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined." - ) - - if positional_embeddings == "sinusoidal": - self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings) - else: - self.pos_embed = None - - # Define 3 blocks. Each block has its own normalization layer. - # 1. Self-Attn - if norm_type == "ada_norm": - self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) - elif norm_type == "ada_norm_zero": - self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) - elif norm_type == "ada_norm_continuous": - self.norm1 = AdaLayerNormContinuous( - dim, - ada_norm_continous_conditioning_embedding_dim, - norm_elementwise_affine, - norm_eps, - ada_norm_bias, - "rms_norm", - ) - else: - self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) - - self.attn1 = Attention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - cross_attention_dim=cross_attention_dim if only_cross_attention else None, - upcast_attention=upcast_attention, - out_bias=attention_out_bias, - ) - - # 2. Cross-Attn - if cross_attention_dim is not None or double_self_attention: - # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. - # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during - # the second cross attention block. - if norm_type == "ada_norm": - self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm) - elif norm_type == "ada_norm_continuous": - self.norm2 = AdaLayerNormContinuous( - dim, - ada_norm_continous_conditioning_embedding_dim, - norm_elementwise_affine, - norm_eps, - ada_norm_bias, - "rms_norm", - ) - else: - self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine) - - self.attn2 = Attention( - query_dim=dim, - cross_attention_dim=cross_attention_dim if not double_self_attention else None, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - upcast_attention=upcast_attention, - out_bias=attention_out_bias, - ) # is self-attn if encoder_hidden_states is none - else: - if norm_type == "ada_norm_single": # For Latte - self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine) - else: - self.norm2 = None - self.attn2 = None - - # 3. Feed-forward - if norm_type == "ada_norm_continuous": - self.norm3 = AdaLayerNormContinuous( - dim, - ada_norm_continous_conditioning_embedding_dim, - norm_elementwise_affine, - norm_eps, - ada_norm_bias, - "layer_norm", - ) - - elif norm_type in ["ada_norm_zero", "ada_norm", "layer_norm"]: - self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine) - elif norm_type == "layer_norm_i2vgen": - self.norm3 = None - - self.ff = FeedForward( - dim, - dropout=dropout, - activation_fn=activation_fn, - final_dropout=final_dropout, - inner_dim=ff_inner_dim, - bias=ff_bias, - ) - - # 4. Fuser - if attention_type == "gated" or attention_type == "gated-text-image": - self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim) - - # 5. Scale-shift for PixArt-Alpha. - if norm_type == "ada_norm_single": - self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) - - # let chunk size default to None - self._chunk_size = None - self._chunk_dim = 0 - - def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): - # Sets chunk feed-forward - self._chunk_size = chunk_size - self._chunk_dim = dim - - 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, - ) -> 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.") - - # Notice that normalization is always applied before the real computation in the following blocks. - # 0. Self-Attention - batch_size = hidden_states.shape[0] - - if self.norm_type == "ada_norm": - norm_hidden_states = self.norm1(hidden_states, timestep) - elif self.norm_type == "ada_norm_zero": - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( - hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype - ) - elif self.norm_type in ["layer_norm", "layer_norm_i2vgen"]: - norm_hidden_states = self.norm1(hidden_states) - elif self.norm_type == "ada_norm_continuous": - norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs["pooled_text_emb"]) - elif self.norm_type == "ada_norm_single": - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( - self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) - ).chunk(6, dim=1) - norm_hidden_states = self.norm1(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa - else: - raise ValueError("Incorrect norm used") - - if self.pos_embed is not None: - norm_hidden_states = self.pos_embed(norm_hidden_states) - - # 1. Prepare GLIGEN inputs - cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} - gligen_kwargs = cross_attention_kwargs.pop("gligen", None) - - attn_output = self.attn1( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, - attention_mask=attention_mask, - **cross_attention_kwargs, - ) - - if self.norm_type == "ada_norm_zero": - attn_output = gate_msa.unsqueeze(1) * attn_output - elif self.norm_type == "ada_norm_single": - attn_output = gate_msa * attn_output - - hidden_states = attn_output + hidden_states - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - # 1.2 GLIGEN Control - if gligen_kwargs is not None: - hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"]) - - # 3. Cross-Attention - if self.attn2 is not None: - if self.norm_type == "ada_norm": - norm_hidden_states = self.norm2(hidden_states, timestep) - elif self.norm_type in ["ada_norm_zero", "layer_norm", "layer_norm_i2vgen"]: - norm_hidden_states = self.norm2(hidden_states) - elif self.norm_type == "ada_norm_single": - # For PixArt norm2 isn't applied here: - # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103 - norm_hidden_states = hidden_states - elif self.norm_type == "ada_norm_continuous": - norm_hidden_states = self.norm2(hidden_states, added_cond_kwargs["pooled_text_emb"]) - else: - raise ValueError("Incorrect norm") - - if self.pos_embed is not None and self.norm_type != "ada_norm_single": - norm_hidden_states = self.pos_embed(norm_hidden_states) - - attn_output = self.attn2( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - **cross_attention_kwargs, - ) - hidden_states = attn_output + hidden_states - - # 4. Feed-forward - # i2vgen doesn't have this norm 🤷‍♂️ - if self.norm_type == "ada_norm_continuous": - norm_hidden_states = self.norm3(hidden_states, added_cond_kwargs["pooled_text_emb"]) - elif not self.norm_type == "ada_norm_single": - norm_hidden_states = self.norm3(hidden_states) - - if self.norm_type == "ada_norm_zero": - norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - - if self.norm_type == "ada_norm_single": - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp - - if self._chunk_size is not None: - # "feed_forward_chunk_size" can be used to save memory - ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size) - else: - ff_output = self.ff(norm_hidden_states) - - if self.norm_type == "ada_norm_zero": - ff_output = gate_mlp.unsqueeze(1) * ff_output - elif self.norm_type == "ada_norm_single": - ff_output = gate_mlp * ff_output - - hidden_states = ff_output + hidden_states - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - return hidden_states - - -class LuminaFeedForward(nn.Module): - r""" - A feed-forward layer. - - Parameters: - hidden_size (`int`): - The dimensionality of the hidden layers in the model. This parameter determines the width of the model's - hidden representations. - intermediate_size (`int`): The intermediate dimension of the feedforward layer. - multiple_of (`int`, *optional*): Value to ensure hidden dimension is a multiple - of this value. - ffn_dim_multiplier (float, *optional*): Custom multiplier for hidden - dimension. Defaults to None. - """ - - def __init__( - self, - dim: int, - inner_dim: int, - multiple_of: Optional[int] = 256, - ffn_dim_multiplier: Optional[float] = None, - ): - super().__init__() - # custom hidden_size factor multiplier - if ffn_dim_multiplier is not None: - inner_dim = int(ffn_dim_multiplier * inner_dim) - inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of) - - self.linear_1 = nn.Linear( - dim, - inner_dim, - bias=False, - ) - self.linear_2 = nn.Linear( - inner_dim, - dim, - bias=False, - ) - self.linear_3 = nn.Linear( - dim, - inner_dim, - bias=False, - ) - self.silu = FP32SiLU() - - def forward(self, x): - return self.linear_2(self.silu(self.linear_1(x)) * self.linear_3(x)) - - -@maybe_allow_in_graph -class TemporalBasicTransformerBlock(nn.Module): - r""" - A basic Transformer block for video like data. - - Parameters: - dim (`int`): The number of channels in the input and output. - time_mix_inner_dim (`int`): The number of channels for temporal attention. - 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. - cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. - """ - - def __init__( - self, - dim: int, - time_mix_inner_dim: int, - num_attention_heads: int, - attention_head_dim: int, - cross_attention_dim: Optional[int] = None, - ): - super().__init__() - self.is_res = dim == time_mix_inner_dim - - self.norm_in = nn.LayerNorm(dim) - - # Define 3 blocks. Each block has its own normalization layer. - # 1. Self-Attn - self.ff_in = FeedForward( - dim, - dim_out=time_mix_inner_dim, - activation_fn="geglu", - ) - - self.norm1 = nn.LayerNorm(time_mix_inner_dim) - self.attn1 = Attention( - query_dim=time_mix_inner_dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - cross_attention_dim=None, - ) - - # 2. Cross-Attn - if cross_attention_dim is not None: - # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. - # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during - # the second cross attention block. - self.norm2 = nn.LayerNorm(time_mix_inner_dim) - self.attn2 = Attention( - query_dim=time_mix_inner_dim, - cross_attention_dim=cross_attention_dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - ) # is self-attn if encoder_hidden_states is none - else: - self.norm2 = None - self.attn2 = None - - # 3. Feed-forward - self.norm3 = nn.LayerNorm(time_mix_inner_dim) - self.ff = FeedForward(time_mix_inner_dim, activation_fn="geglu") - - # let chunk size default to None - self._chunk_size = None - self._chunk_dim = None - - def set_chunk_feed_forward(self, chunk_size: Optional[int], **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 - self._chunk_dim = 1 - - def forward( - self, - hidden_states: torch.Tensor, - num_frames: int, - encoder_hidden_states: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - # Notice that normalization is always applied before the real computation in the following blocks. - # 0. Self-Attention - batch_size = hidden_states.shape[0] - - batch_frames, seq_length, channels = hidden_states.shape - batch_size = batch_frames // num_frames - - hidden_states = hidden_states[None, :].reshape(batch_size, num_frames, seq_length, channels) - hidden_states = hidden_states.permute(0, 2, 1, 3) - hidden_states = hidden_states.reshape(batch_size * seq_length, num_frames, channels) - - residual = hidden_states - hidden_states = self.norm_in(hidden_states) - - if self._chunk_size is not None: - hidden_states = _chunked_feed_forward(self.ff_in, hidden_states, self._chunk_dim, self._chunk_size) - else: - hidden_states = self.ff_in(hidden_states) - - if self.is_res: - hidden_states = hidden_states + residual - - norm_hidden_states = self.norm1(hidden_states) - attn_output = self.attn1(norm_hidden_states, encoder_hidden_states=None) - hidden_states = attn_output + hidden_states - - # 3. Cross-Attention - if self.attn2 is not None: - norm_hidden_states = self.norm2(hidden_states) - attn_output = self.attn2(norm_hidden_states, encoder_hidden_states=encoder_hidden_states) - hidden_states = attn_output + hidden_states - - # 4. Feed-forward - norm_hidden_states = self.norm3(hidden_states) - - if self._chunk_size is not None: - ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size) - else: - ff_output = self.ff(norm_hidden_states) - - if self.is_res: - hidden_states = ff_output + hidden_states - else: - hidden_states = ff_output - - hidden_states = hidden_states[None, :].reshape(batch_size, seq_length, num_frames, channels) - hidden_states = hidden_states.permute(0, 2, 1, 3) - hidden_states = hidden_states.reshape(batch_size * num_frames, seq_length, channels) - - return hidden_states - - -class SkipFFTransformerBlock(nn.Module): - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - kv_input_dim: int, - kv_input_dim_proj_use_bias: bool, - dropout=0.0, - cross_attention_dim: Optional[int] = None, - attention_bias: bool = False, - attention_out_bias: bool = True, - ): - super().__init__() - if kv_input_dim != dim: - self.kv_mapper = nn.Linear(kv_input_dim, dim, kv_input_dim_proj_use_bias) - else: - self.kv_mapper = None - - self.norm1 = RMSNorm(dim, 1e-06) - - self.attn1 = Attention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - cross_attention_dim=cross_attention_dim, - out_bias=attention_out_bias, - ) - - self.norm2 = RMSNorm(dim, 1e-06) - - self.attn2 = Attention( - query_dim=dim, - cross_attention_dim=cross_attention_dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - out_bias=attention_out_bias, - ) - - def forward(self, hidden_states, encoder_hidden_states, cross_attention_kwargs): - cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} - - if self.kv_mapper is not None: - encoder_hidden_states = self.kv_mapper(F.silu(encoder_hidden_states)) - - norm_hidden_states = self.norm1(hidden_states) - - attn_output = self.attn1( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states, - **cross_attention_kwargs, - ) - - hidden_states = attn_output + hidden_states - - norm_hidden_states = self.norm2(hidden_states) - - attn_output = self.attn2( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states, - **cross_attention_kwargs, - ) - - hidden_states = attn_output + hidden_states - - return hidden_states - - -@maybe_allow_in_graph -class FreeNoiseTransformerBlock(nn.Module): - r""" - A FreeNoise Transformer block. - - Parameters: - 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. - dropout (`float`, *optional*, defaults to 0.0): - The dropout probability to use. - cross_attention_dim (`int`, *optional*): - The size of the encoder_hidden_states vector for cross attention. - activation_fn (`str`, *optional*, defaults to `"geglu"`): - Activation function to be used in feed-forward. - num_embeds_ada_norm (`int`, *optional*): - The number of diffusion steps used during training. See `Transformer2DModel`. - attention_bias (`bool`, defaults to `False`): - Configure if the attentions should contain a bias parameter. - only_cross_attention (`bool`, defaults to `False`): - Whether to use only cross-attention layers. In this case two cross attention layers are used. - double_self_attention (`bool`, defaults to `False`): - Whether to use two self-attention layers. In this case no cross attention layers are used. - upcast_attention (`bool`, defaults to `False`): - Whether to upcast the attention computation to float32. This is useful for mixed precision training. - norm_elementwise_affine (`bool`, defaults to `True`): - Whether to use learnable elementwise affine parameters for normalization. - norm_type (`str`, defaults to `"layer_norm"`): - The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`. - final_dropout (`bool` defaults to `False`): - Whether to apply a final dropout after the last feed-forward layer. - attention_type (`str`, defaults to `"default"`): - The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`. - positional_embeddings (`str`, *optional*): - The type of positional embeddings to apply to. - num_positional_embeddings (`int`, *optional*, defaults to `None`): - The maximum number of positional embeddings to apply. - ff_inner_dim (`int`, *optional*): - Hidden dimension of feed-forward MLP. - ff_bias (`bool`, defaults to `True`): - Whether or not to use bias in feed-forward MLP. - attention_out_bias (`bool`, defaults to `True`): - Whether or not to use bias in attention output project layer. - context_length (`int`, defaults to `16`): - The maximum number of frames that the FreeNoise block processes at once. - context_stride (`int`, defaults to `4`): - The number of frames to be skipped before starting to process a new batch of `context_length` frames. - weighting_scheme (`str`, defaults to `"pyramid"`): - The weighting scheme to use for weighting averaging of processed latent frames. As described in the - Equation 9. of the [FreeNoise](https://huggingface.co/papers/2310.15169) paper, "pyramid" is the default - setting used. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - dropout: float = 0.0, - cross_attention_dim: Optional[int] = None, - activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, - attention_bias: bool = False, - only_cross_attention: bool = False, - double_self_attention: bool = False, - upcast_attention: bool = False, - norm_elementwise_affine: bool = True, - 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, - ff_bias: bool = True, - attention_out_bias: bool = True, - context_length: int = 16, - context_stride: int = 4, - weighting_scheme: str = "pyramid", - ): - super().__init__() - self.dim = dim - self.num_attention_heads = num_attention_heads - self.attention_head_dim = attention_head_dim - self.dropout = dropout - self.cross_attention_dim = cross_attention_dim - self.activation_fn = activation_fn - self.attention_bias = attention_bias - self.double_self_attention = double_self_attention - self.norm_elementwise_affine = norm_elementwise_affine - self.positional_embeddings = positional_embeddings - self.num_positional_embeddings = num_positional_embeddings - self.only_cross_attention = only_cross_attention - - self.set_free_noise_properties(context_length, context_stride, weighting_scheme) - - # We keep these boolean flags for backward-compatibility. - self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" - self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" - self.use_ada_layer_norm_single = norm_type == "ada_norm_single" - self.use_layer_norm = norm_type == "layer_norm" - self.use_ada_layer_norm_continuous = norm_type == "ada_norm_continuous" - - if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: - raise ValueError( - f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" - f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." - ) - - self.norm_type = norm_type - self.num_embeds_ada_norm = num_embeds_ada_norm - - if positional_embeddings and (num_positional_embeddings is None): - raise ValueError( - "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined." - ) - - if positional_embeddings == "sinusoidal": - self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings) - else: - self.pos_embed = None - - # Define 3 blocks. Each block has its own normalization layer. - # 1. Self-Attn - self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) - - self.attn1 = Attention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - cross_attention_dim=cross_attention_dim if only_cross_attention else None, - upcast_attention=upcast_attention, - out_bias=attention_out_bias, - ) - - # 2. Cross-Attn - if cross_attention_dim is not None or double_self_attention: - self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine) - - self.attn2 = Attention( - query_dim=dim, - cross_attention_dim=cross_attention_dim if not double_self_attention else None, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - upcast_attention=upcast_attention, - out_bias=attention_out_bias, - ) # is self-attn if encoder_hidden_states is none - - # 3. Feed-forward - self.ff = FeedForward( - dim, - dropout=dropout, - activation_fn=activation_fn, - final_dropout=final_dropout, - inner_dim=ff_inner_dim, - bias=ff_bias, - ) - - self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine) - - # let chunk size default to None - self._chunk_size = None - self._chunk_dim = 0 - - 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 - window_end = min(num_frames, i + self.context_length) - frame_indices.append((window_start, window_end)) - return frame_indices - - def _get_frame_weights(self, num_frames: int, weighting_scheme: str = "pyramid") -> List[float]: - if weighting_scheme == "flat": - weights = [1.0] * num_frames - - elif weighting_scheme == "pyramid": - if num_frames % 2 == 0: - # num_frames = 4 => [1, 2, 2, 1] - mid = num_frames // 2 - weights = list(range(1, mid + 1)) - weights = weights + weights[::-1] - else: - # num_frames = 5 => [1, 2, 3, 2, 1] - mid = (num_frames + 1) // 2 - weights = list(range(1, mid)) - weights = weights + [mid] + weights[::-1] - - elif weighting_scheme == "delayed_reverse_sawtooth": - if num_frames % 2 == 0: - # num_frames = 4 => [0.01, 2, 2, 1] - mid = num_frames // 2 - weights = [0.01] * (mid - 1) + [mid] - weights = weights + list(range(mid, 0, -1)) - else: - # num_frames = 5 => [0.01, 0.01, 3, 2, 1] - mid = (num_frames + 1) // 2 - weights = [0.01] * mid - weights = weights + list(range(mid, 0, -1)) - else: - raise ValueError(f"Unsupported value for weighting_scheme={weighting_scheme}") - - return weights - - def set_free_noise_properties( - self, context_length: int, context_stride: int, weighting_scheme: str = "pyramid" - ) -> None: - self.context_length = context_length - self.context_stride = context_stride - self.weighting_scheme = weighting_scheme - - def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0) -> None: - # Sets chunk feed-forward - self._chunk_size = chunk_size - self._chunk_dim = dim - - 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, - *args, - **kwargs, - ) -> 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.") - - cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} - - # hidden_states: [B x H x W, F, C] - device = hidden_states.device - dtype = hidden_states.dtype - - num_frames = hidden_states.size(1) - frame_indices = self._get_frame_indices(num_frames) - frame_weights = self._get_frame_weights(self.context_length, self.weighting_scheme) - frame_weights = torch.tensor(frame_weights, device=device, dtype=dtype).unsqueeze(0).unsqueeze(-1) - is_last_frame_batch_complete = frame_indices[-1][1] == num_frames - - # Handle out-of-bounds case if num_frames isn't perfectly divisible by context_length - # For example, num_frames=25, context_length=16, context_stride=4, then we expect the ranges: - # [(0, 16), (4, 20), (8, 24), (10, 26)] - if not is_last_frame_batch_complete: - if num_frames < self.context_length: - raise ValueError(f"Expected {num_frames=} to be greater or equal than {self.context_length=}") - last_frame_batch_length = num_frames - frame_indices[-1][1] - frame_indices.append((num_frames - self.context_length, num_frames)) - - num_times_accumulated = torch.zeros((1, num_frames, 1), device=device) - accumulated_values = torch.zeros_like(hidden_states) - - for i, (frame_start, frame_end) in enumerate(frame_indices): - # The reason for slicing here is to ensure that if (frame_end - frame_start) is to handle - # cases like frame_indices=[(0, 16), (16, 20)], if the user provided a video with 19 frames, or - # essentially a non-multiple of `context_length`. - weights = torch.ones_like(num_times_accumulated[:, frame_start:frame_end]) - weights *= frame_weights - - hidden_states_chunk = hidden_states[:, frame_start:frame_end] - - # Notice that normalization is always applied before the real computation in the following blocks. - # 1. Self-Attention - norm_hidden_states = self.norm1(hidden_states_chunk) - - if self.pos_embed is not None: - norm_hidden_states = self.pos_embed(norm_hidden_states) - - attn_output = self.attn1( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, - attention_mask=attention_mask, - **cross_attention_kwargs, - ) - - hidden_states_chunk = attn_output + hidden_states_chunk - if hidden_states_chunk.ndim == 4: - hidden_states_chunk = hidden_states_chunk.squeeze(1) - - # 2. Cross-Attention - if self.attn2 is not None: - norm_hidden_states = self.norm2(hidden_states_chunk) - - if self.pos_embed is not None and self.norm_type != "ada_norm_single": - norm_hidden_states = self.pos_embed(norm_hidden_states) - - attn_output = self.attn2( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - **cross_attention_kwargs, - ) - hidden_states_chunk = attn_output + hidden_states_chunk - - if i == len(frame_indices) - 1 and not is_last_frame_batch_complete: - accumulated_values[:, -last_frame_batch_length:] += ( - hidden_states_chunk[:, -last_frame_batch_length:] * weights[:, -last_frame_batch_length:] - ) - num_times_accumulated[:, -last_frame_batch_length:] += weights[:, -last_frame_batch_length] - else: - accumulated_values[:, frame_start:frame_end] += hidden_states_chunk * weights - num_times_accumulated[:, frame_start:frame_end] += weights - - # TODO(aryan): Maybe this could be done in a better way. - # - # Previously, this was: - # hidden_states = torch.where( - # num_times_accumulated > 0, accumulated_values / num_times_accumulated, accumulated_values - # ) - # - # The reasoning for the change here is `torch.where` became a bottleneck at some point when golfing memory - # spikes. It is particularly noticeable when the number of frames is high. My understanding is that this comes - # from tensors being copied - which is why we resort to spliting and concatenating here. I've not particularly - # looked into this deeply because other memory optimizations led to more pronounced reductions. - hidden_states = torch.cat( - [ - torch.where(num_times_split > 0, accumulated_split / num_times_split, accumulated_split) - for accumulated_split, num_times_split in zip( - accumulated_values.split(self.context_length, dim=1), - num_times_accumulated.split(self.context_length, dim=1), - ) - ], - dim=1, - ).to(dtype) - - # 3. Feed-forward - norm_hidden_states = self.norm3(hidden_states) - - if self._chunk_size is not None: - ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size) - else: - ff_output = self.ff(norm_hidden_states) - - hidden_states = ff_output + hidden_states - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - return hidden_states - - -class FeedForward(nn.Module): - r""" - A feed-forward layer. - - Parameters: - dim (`int`): The number of channels in the input. - dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`. - mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. - final_dropout (`bool` *optional*, defaults to False): Apply a final dropout. - bias (`bool`, defaults to True): Whether to use a bias in the linear layer. - """ - - def __init__( - self, - dim: int, - dim_out: Optional[int] = None, - mult: int = 4, - dropout: float = 0.0, - activation_fn: str = "geglu", - final_dropout: bool = False, - inner_dim=None, - bias: bool = True, - ): - super().__init__() - if inner_dim is None: - inner_dim = int(dim * mult) - dim_out = dim_out if dim_out is not None else dim - - if activation_fn == "gelu": - act_fn = GELU(dim, inner_dim, bias=bias) - if activation_fn == "gelu-approximate": - act_fn = GELU(dim, inner_dim, approximate="tanh", bias=bias) - elif activation_fn == "geglu": - act_fn = GEGLU(dim, inner_dim, bias=bias) - elif activation_fn == "geglu-approximate": - act_fn = ApproximateGELU(dim, inner_dim, bias=bias) - elif activation_fn == "swiglu": - act_fn = SwiGLU(dim, inner_dim, bias=bias) - elif activation_fn == "linear-silu": - act_fn = LinearActivation(dim, inner_dim, bias=bias, activation="silu") - - self.net = nn.ModuleList([]) - # project in - self.net.append(act_fn) - # project dropout - self.net.append(nn.Dropout(dropout)) - # project out - self.net.append(nn.Linear(inner_dim, dim_out, bias=bias)) - # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout - if final_dropout: - self.net.append(nn.Dropout(dropout)) - - def forward(self, hidden_states: torch.Tensor, *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) - for module in self.net: - hidden_states = module(hidden_states) - return hidden_states diff --git a/venv/lib/python3.10/site-packages/diffusers/models/attention_flax.py b/venv/lib/python3.10/site-packages/diffusers/models/attention_flax.py deleted file mode 100644 index 17e6f33df0510fe309ff7b2ffd9f6d832dfb11d2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/attention_flax.py +++ /dev/null @@ -1,494 +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 functools -import math - -import flax.linen as nn -import jax -import jax.numpy as jnp - - -def _query_chunk_attention(query, key, value, precision, key_chunk_size: int = 4096): - """Multi-head dot product attention with a limited number of queries.""" - num_kv, num_heads, k_features = key.shape[-3:] - v_features = value.shape[-1] - key_chunk_size = min(key_chunk_size, num_kv) - query = query / jnp.sqrt(k_features) - - @functools.partial(jax.checkpoint, prevent_cse=False) - def summarize_chunk(query, key, value): - attn_weights = jnp.einsum("...qhd,...khd->...qhk", query, key, precision=precision) - - max_score = jnp.max(attn_weights, axis=-1, keepdims=True) - max_score = jax.lax.stop_gradient(max_score) - exp_weights = jnp.exp(attn_weights - max_score) - - exp_values = jnp.einsum("...vhf,...qhv->...qhf", value, exp_weights, precision=precision) - max_score = jnp.einsum("...qhk->...qh", max_score) - - return (exp_values, exp_weights.sum(axis=-1), max_score) - - def chunk_scanner(chunk_idx): - # julienne key array - key_chunk = jax.lax.dynamic_slice( - operand=key, - start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0], # [...,k,h,d] - slice_sizes=list(key.shape[:-3]) + [key_chunk_size, num_heads, k_features], # [...,k,h,d] - ) - - # julienne value array - value_chunk = jax.lax.dynamic_slice( - operand=value, - start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0], # [...,v,h,d] - slice_sizes=list(value.shape[:-3]) + [key_chunk_size, num_heads, v_features], # [...,v,h,d] - ) - - return summarize_chunk(query, key_chunk, value_chunk) - - chunk_values, chunk_weights, chunk_max = jax.lax.map(f=chunk_scanner, xs=jnp.arange(0, num_kv, key_chunk_size)) - - global_max = jnp.max(chunk_max, axis=0, keepdims=True) - max_diffs = jnp.exp(chunk_max - global_max) - - chunk_values *= jnp.expand_dims(max_diffs, axis=-1) - chunk_weights *= max_diffs - - all_values = chunk_values.sum(axis=0) - all_weights = jnp.expand_dims(chunk_weights, -1).sum(axis=0) - - return all_values / all_weights - - -def jax_memory_efficient_attention( - query, key, value, precision=jax.lax.Precision.HIGHEST, query_chunk_size: int = 1024, key_chunk_size: int = 4096 -): - r""" - Flax Memory-efficient multi-head dot product attention. https://huggingface.co/papers/2112.05682v2 - https://github.com/AminRezaei0x443/memory-efficient-attention - - Args: - query (`jnp.ndarray`): (batch..., query_length, head, query_key_depth_per_head) - key (`jnp.ndarray`): (batch..., key_value_length, head, query_key_depth_per_head) - value (`jnp.ndarray`): (batch..., key_value_length, head, value_depth_per_head) - precision (`jax.lax.Precision`, *optional*, defaults to `jax.lax.Precision.HIGHEST`): - numerical precision for computation - query_chunk_size (`int`, *optional*, defaults to 1024): - chunk size to divide query array value must divide query_length equally without remainder - key_chunk_size (`int`, *optional*, defaults to 4096): - chunk size to divide key and value array value must divide key_value_length equally without remainder - - Returns: - (`jnp.ndarray`) with shape of (batch..., query_length, head, value_depth_per_head) - """ - num_q, num_heads, q_features = query.shape[-3:] - - def chunk_scanner(chunk_idx, _): - # julienne query array - query_chunk = jax.lax.dynamic_slice( - operand=query, - start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0], # [...,q,h,d] - slice_sizes=list(query.shape[:-3]) + [min(query_chunk_size, num_q), num_heads, q_features], # [...,q,h,d] - ) - - return ( - chunk_idx + query_chunk_size, # unused ignore it - _query_chunk_attention( - query=query_chunk, key=key, value=value, precision=precision, key_chunk_size=key_chunk_size - ), - ) - - _, res = jax.lax.scan( - f=chunk_scanner, - init=0, - xs=None, - length=math.ceil(num_q / query_chunk_size), # start counter # stop counter - ) - - return jnp.concatenate(res, axis=-3) # fuse the chunked result back - - -class FlaxAttention(nn.Module): - r""" - A Flax multi-head attention module as described in: https://huggingface.co/papers/1706.03762 - - Parameters: - query_dim (:obj:`int`): - Input hidden states dimension - heads (:obj:`int`, *optional*, defaults to 8): - Number of heads - dim_head (:obj:`int`, *optional*, defaults to 64): - Hidden states dimension inside each head - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - use_memory_efficient_attention (`bool`, *optional*, defaults to `False`): - enable memory efficient attention https://huggingface.co/papers/2112.05682 - split_head_dim (`bool`, *optional*, defaults to `False`): - Whether to split the head dimension into a new axis for the self-attention computation. In most cases, - enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL. - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - - """ - - query_dim: int - heads: int = 8 - dim_head: int = 64 - dropout: float = 0.0 - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - dtype: jnp.dtype = jnp.float32 - - def setup(self): - inner_dim = self.dim_head * self.heads - self.scale = self.dim_head**-0.5 - - # Weights were exported with old names {to_q, to_k, to_v, to_out} - self.query = nn.Dense(inner_dim, use_bias=False, dtype=self.dtype, name="to_q") - self.key = nn.Dense(inner_dim, use_bias=False, dtype=self.dtype, name="to_k") - self.value = nn.Dense(inner_dim, use_bias=False, dtype=self.dtype, name="to_v") - - self.proj_attn = nn.Dense(self.query_dim, dtype=self.dtype, name="to_out_0") - self.dropout_layer = nn.Dropout(rate=self.dropout) - - def reshape_heads_to_batch_dim(self, tensor): - batch_size, seq_len, dim = tensor.shape - head_size = self.heads - tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size) - tensor = jnp.transpose(tensor, (0, 2, 1, 3)) - tensor = tensor.reshape(batch_size * head_size, seq_len, dim // head_size) - return tensor - - def reshape_batch_dim_to_heads(self, tensor): - batch_size, seq_len, dim = tensor.shape - head_size = self.heads - tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim) - tensor = jnp.transpose(tensor, (0, 2, 1, 3)) - tensor = tensor.reshape(batch_size // head_size, seq_len, dim * head_size) - return tensor - - def __call__(self, hidden_states, context=None, deterministic=True): - context = hidden_states if context is None else context - - query_proj = self.query(hidden_states) - key_proj = self.key(context) - value_proj = self.value(context) - - if self.split_head_dim: - b = hidden_states.shape[0] - query_states = jnp.reshape(query_proj, (b, -1, self.heads, self.dim_head)) - key_states = jnp.reshape(key_proj, (b, -1, self.heads, self.dim_head)) - value_states = jnp.reshape(value_proj, (b, -1, self.heads, self.dim_head)) - else: - query_states = self.reshape_heads_to_batch_dim(query_proj) - key_states = self.reshape_heads_to_batch_dim(key_proj) - value_states = self.reshape_heads_to_batch_dim(value_proj) - - if self.use_memory_efficient_attention: - query_states = query_states.transpose(1, 0, 2) - key_states = key_states.transpose(1, 0, 2) - value_states = value_states.transpose(1, 0, 2) - - # this if statement create a chunk size for each layer of the unet - # the chunk size is equal to the query_length dimension of the deepest layer of the unet - - flatten_latent_dim = query_states.shape[-3] - if flatten_latent_dim % 64 == 0: - query_chunk_size = int(flatten_latent_dim / 64) - elif flatten_latent_dim % 16 == 0: - query_chunk_size = int(flatten_latent_dim / 16) - elif flatten_latent_dim % 4 == 0: - query_chunk_size = int(flatten_latent_dim / 4) - else: - query_chunk_size = int(flatten_latent_dim) - - hidden_states = jax_memory_efficient_attention( - query_states, key_states, value_states, query_chunk_size=query_chunk_size, key_chunk_size=4096 * 4 - ) - hidden_states = hidden_states.transpose(1, 0, 2) - hidden_states = self.reshape_batch_dim_to_heads(hidden_states) - else: - # compute attentions - if self.split_head_dim: - attention_scores = jnp.einsum("b t n h, b f n h -> b n f t", key_states, query_states) - else: - attention_scores = jnp.einsum("b i d, b j d->b i j", query_states, key_states) - - attention_scores = attention_scores * self.scale - attention_probs = nn.softmax(attention_scores, axis=-1 if self.split_head_dim else 2) - - # attend to values - if self.split_head_dim: - hidden_states = jnp.einsum("b n f t, b t n h -> b f n h", attention_probs, value_states) - b = hidden_states.shape[0] - hidden_states = jnp.reshape(hidden_states, (b, -1, self.heads * self.dim_head)) - else: - hidden_states = jnp.einsum("b i j, b j d -> b i d", attention_probs, value_states) - hidden_states = self.reshape_batch_dim_to_heads(hidden_states) - - hidden_states = self.proj_attn(hidden_states) - return self.dropout_layer(hidden_states, deterministic=deterministic) - - -class FlaxBasicTransformerBlock(nn.Module): - r""" - A Flax transformer block layer with `GLU` (Gated Linear Unit) activation function as described in: - https://huggingface.co/papers/1706.03762 - - - Parameters: - dim (:obj:`int`): - Inner hidden states dimension - n_heads (:obj:`int`): - Number of heads - d_head (:obj:`int`): - Hidden states dimension inside each head - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - only_cross_attention (`bool`, defaults to `False`): - Whether to only apply cross attention. - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - use_memory_efficient_attention (`bool`, *optional*, defaults to `False`): - enable memory efficient attention https://huggingface.co/papers/2112.05682 - split_head_dim (`bool`, *optional*, defaults to `False`): - Whether to split the head dimension into a new axis for the self-attention computation. In most cases, - enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL. - """ - - dim: int - n_heads: int - d_head: int - dropout: float = 0.0 - only_cross_attention: bool = False - dtype: jnp.dtype = jnp.float32 - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - - def setup(self): - # self attention (or cross_attention if only_cross_attention is True) - self.attn1 = FlaxAttention( - self.dim, - self.n_heads, - self.d_head, - self.dropout, - self.use_memory_efficient_attention, - self.split_head_dim, - dtype=self.dtype, - ) - # cross attention - self.attn2 = FlaxAttention( - self.dim, - self.n_heads, - self.d_head, - self.dropout, - self.use_memory_efficient_attention, - self.split_head_dim, - dtype=self.dtype, - ) - self.ff = FlaxFeedForward(dim=self.dim, dropout=self.dropout, dtype=self.dtype) - self.norm1 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype) - self.norm2 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype) - self.norm3 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype) - self.dropout_layer = nn.Dropout(rate=self.dropout) - - def __call__(self, hidden_states, context, deterministic=True): - # self attention - residual = hidden_states - if self.only_cross_attention: - hidden_states = self.attn1(self.norm1(hidden_states), context, deterministic=deterministic) - else: - hidden_states = self.attn1(self.norm1(hidden_states), deterministic=deterministic) - hidden_states = hidden_states + residual - - # cross attention - residual = hidden_states - hidden_states = self.attn2(self.norm2(hidden_states), context, deterministic=deterministic) - hidden_states = hidden_states + residual - - # feed forward - residual = hidden_states - hidden_states = self.ff(self.norm3(hidden_states), deterministic=deterministic) - hidden_states = hidden_states + residual - - return self.dropout_layer(hidden_states, deterministic=deterministic) - - -class FlaxTransformer2DModel(nn.Module): - r""" - A Spatial Transformer layer with Gated Linear Unit (GLU) activation function as described in: - https://huggingface.co/papers/1506.02025 - - - Parameters: - in_channels (:obj:`int`): - Input number of channels - n_heads (:obj:`int`): - Number of heads - d_head (:obj:`int`): - Hidden states dimension inside each head - depth (:obj:`int`, *optional*, defaults to 1): - Number of transformers block - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - use_linear_projection (`bool`, defaults to `False`): tbd - only_cross_attention (`bool`, defaults to `False`): tbd - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - use_memory_efficient_attention (`bool`, *optional*, defaults to `False`): - enable memory efficient attention https://huggingface.co/papers/2112.05682 - split_head_dim (`bool`, *optional*, defaults to `False`): - Whether to split the head dimension into a new axis for the self-attention computation. In most cases, - enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL. - """ - - in_channels: int - n_heads: int - d_head: int - depth: int = 1 - dropout: float = 0.0 - use_linear_projection: bool = False - only_cross_attention: bool = False - dtype: jnp.dtype = jnp.float32 - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - - def setup(self): - self.norm = nn.GroupNorm(num_groups=32, epsilon=1e-5) - - inner_dim = self.n_heads * self.d_head - if self.use_linear_projection: - self.proj_in = nn.Dense(inner_dim, dtype=self.dtype) - else: - self.proj_in = nn.Conv( - inner_dim, - kernel_size=(1, 1), - strides=(1, 1), - padding="VALID", - dtype=self.dtype, - ) - - self.transformer_blocks = [ - FlaxBasicTransformerBlock( - inner_dim, - self.n_heads, - self.d_head, - dropout=self.dropout, - only_cross_attention=self.only_cross_attention, - dtype=self.dtype, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - ) - for _ in range(self.depth) - ] - - if self.use_linear_projection: - self.proj_out = nn.Dense(inner_dim, dtype=self.dtype) - else: - self.proj_out = nn.Conv( - inner_dim, - kernel_size=(1, 1), - strides=(1, 1), - padding="VALID", - dtype=self.dtype, - ) - - self.dropout_layer = nn.Dropout(rate=self.dropout) - - def __call__(self, hidden_states, context, deterministic=True): - batch, height, width, channels = hidden_states.shape - residual = hidden_states - hidden_states = self.norm(hidden_states) - if self.use_linear_projection: - hidden_states = hidden_states.reshape(batch, height * width, channels) - hidden_states = self.proj_in(hidden_states) - else: - hidden_states = self.proj_in(hidden_states) - hidden_states = hidden_states.reshape(batch, height * width, channels) - - for transformer_block in self.transformer_blocks: - hidden_states = transformer_block(hidden_states, context, deterministic=deterministic) - - if self.use_linear_projection: - hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.reshape(batch, height, width, channels) - else: - hidden_states = hidden_states.reshape(batch, height, width, channels) - hidden_states = self.proj_out(hidden_states) - - hidden_states = hidden_states + residual - return self.dropout_layer(hidden_states, deterministic=deterministic) - - -class FlaxFeedForward(nn.Module): - r""" - Flax module that encapsulates two Linear layers separated by a non-linearity. It is the counterpart of PyTorch's - [`FeedForward`] class, with the following simplifications: - - The activation function is currently hardcoded to a gated linear unit from: - https://huggingface.co/papers/2002.05202 - - `dim_out` is equal to `dim`. - - The number of hidden dimensions is hardcoded to `dim * 4` in [`FlaxGELU`]. - - Parameters: - dim (:obj:`int`): - Inner hidden states dimension - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - dim: int - dropout: float = 0.0 - dtype: jnp.dtype = jnp.float32 - - def setup(self): - # The second linear layer needs to be called - # net_2 for now to match the index of the Sequential layer - self.net_0 = FlaxGEGLU(self.dim, self.dropout, self.dtype) - self.net_2 = nn.Dense(self.dim, dtype=self.dtype) - - def __call__(self, hidden_states, deterministic=True): - hidden_states = self.net_0(hidden_states, deterministic=deterministic) - hidden_states = self.net_2(hidden_states) - return hidden_states - - -class FlaxGEGLU(nn.Module): - r""" - Flax implementation of a Linear layer followed by the variant of the gated linear unit activation function from - https://huggingface.co/papers/2002.05202. - - Parameters: - dim (:obj:`int`): - Input hidden states dimension - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - dim: int - dropout: float = 0.0 - dtype: jnp.dtype = jnp.float32 - - def setup(self): - inner_dim = self.dim * 4 - self.proj = nn.Dense(inner_dim * 2, dtype=self.dtype) - self.dropout_layer = nn.Dropout(rate=self.dropout) - - def __call__(self, hidden_states, deterministic=True): - hidden_states = self.proj(hidden_states) - hidden_linear, hidden_gelu = jnp.split(hidden_states, 2, axis=2) - return self.dropout_layer(hidden_linear * nn.gelu(hidden_gelu), deterministic=deterministic) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/attention_processor.py b/venv/lib/python3.10/site-packages/diffusers/models/attention_processor.py deleted file mode 100644 index 4760cfd40b3cd2393df2062b63bdc58bf2007972..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/attention_processor.py +++ /dev/null @@ -1,6239 +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 inspect -import math -from typing import Callable, List, Optional, Tuple, Union - -import torch -import torch.nn.functional as F -from torch import nn - -from ..image_processor import IPAdapterMaskProcessor -from ..utils import deprecate, is_torch_xla_available, logging -from ..utils.import_utils import is_torch_npu_available, is_torch_xla_version, is_xformers_available -from ..utils.torch_utils import is_torch_version, maybe_allow_in_graph - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -if is_torch_npu_available(): - import torch_npu - -if is_xformers_available(): - import xformers - import xformers.ops -else: - xformers = None - -if is_torch_xla_available(): - # flash attention pallas kernel is introduced in the torch_xla 2.3 release. - if is_torch_xla_version(">", "2.2"): - from torch_xla.experimental.custom_kernel import flash_attention - from torch_xla.runtime import is_spmd - XLA_AVAILABLE = True -else: - XLA_AVAILABLE = False - - -@maybe_allow_in_graph -class Attention(nn.Module): - r""" - A cross attention layer. - - Parameters: - query_dim (`int`): - The number of channels in the query. - cross_attention_dim (`int`, *optional*): - The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`. - heads (`int`, *optional*, defaults to 8): - The number of heads to use for multi-head attention. - kv_heads (`int`, *optional*, defaults to `None`): - The number of key and value heads to use for multi-head attention. Defaults to `heads`. If - `kv_heads=heads`, the model will use Multi Head Attention (MHA), if `kv_heads=1` the model will use Multi - Query Attention (MQA) otherwise GQA is used. - dim_head (`int`, *optional*, defaults to 64): - The number of channels in each head. - dropout (`float`, *optional*, defaults to 0.0): - The dropout probability to use. - bias (`bool`, *optional*, defaults to False): - Set to `True` for the query, key, and value linear layers to contain a bias parameter. - upcast_attention (`bool`, *optional*, defaults to False): - Set to `True` to upcast the attention computation to `float32`. - upcast_softmax (`bool`, *optional*, defaults to False): - Set to `True` to upcast the softmax computation to `float32`. - cross_attention_norm (`str`, *optional*, defaults to `None`): - The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`. - cross_attention_norm_num_groups (`int`, *optional*, defaults to 32): - The number of groups to use for the group norm in the cross attention. - 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. - norm_num_groups (`int`, *optional*, defaults to `None`): - The number of groups to use for the group norm in the attention. - spatial_norm_dim (`int`, *optional*, defaults to `None`): - The number of channels to use for the spatial normalization. - out_bias (`bool`, *optional*, defaults to `True`): - Set to `True` to use a bias in the output linear layer. - scale_qk (`bool`, *optional*, defaults to `True`): - Set to `True` to scale the query and key by `1 / sqrt(dim_head)`. - only_cross_attention (`bool`, *optional*, defaults to `False`): - Set to `True` to only use cross attention and not added_kv_proj_dim. Can only be set to `True` if - `added_kv_proj_dim` is not `None`. - eps (`float`, *optional*, defaults to 1e-5): - An additional value added to the denominator in group normalization that is used for numerical stability. - rescale_output_factor (`float`, *optional*, defaults to 1.0): - A factor to rescale the output by dividing it with this value. - residual_connection (`bool`, *optional*, defaults to `False`): - Set to `True` to add the residual connection to the output. - _from_deprecated_attn_block (`bool`, *optional*, defaults to `False`): - Set to `True` if the attention block is loaded from a deprecated state dict. - processor (`AttnProcessor`, *optional*, defaults to `None`): - The attention processor to use. If `None`, defaults to `AttnProcessor2_0` if `torch 2.x` is used and - `AttnProcessor` otherwise. - """ - - def __init__( - self, - query_dim: int, - cross_attention_dim: Optional[int] = None, - heads: int = 8, - kv_heads: Optional[int] = 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_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, - out_bias: bool = True, - scale_qk: bool = True, - only_cross_attention: bool = False, - eps: float = 1e-5, - rescale_output_factor: float = 1.0, - residual_connection: bool = False, - _from_deprecated_attn_block: bool = False, - processor: Optional["AttnProcessor"] = None, - out_dim: int = None, - out_context_dim: int = None, - context_pre_only=None, - pre_only=False, - elementwise_affine: bool = True, - is_causal: bool = False, - ): - super().__init__() - - # To prevent circular import. - from .normalization import FP32LayerNorm, LpNorm, RMSNorm - - self.inner_dim = out_dim if out_dim is not None else 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.use_bias = bias - self.is_cross_attention = cross_attention_dim is not None - self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim - self.upcast_attention = upcast_attention - self.upcast_softmax = upcast_softmax - self.rescale_output_factor = rescale_output_factor - self.residual_connection = residual_connection - self.dropout = dropout - self.fused_projections = False - self.out_dim = out_dim if out_dim is not None else query_dim - self.out_context_dim = out_context_dim if out_context_dim is not None else query_dim - self.context_pre_only = context_pre_only - self.pre_only = pre_only - self.is_causal = is_causal - - # we make use of this private variable to know whether this class is loaded - # with an deprecated state dict so that we can convert it on the fly - self._from_deprecated_attn_block = _from_deprecated_attn_block - - self.scale_qk = scale_qk - self.scale = dim_head**-0.5 if self.scale_qk else 1.0 - - self.heads = out_dim // dim_head if out_dim is not None else heads - # for slice_size > 0 the attention score computation - # is split across the batch axis to save memory - # You can set slice_size with `set_attention_slice` - self.sliceable_head_dim = heads - - self.added_kv_proj_dim = added_kv_proj_dim - self.only_cross_attention = only_cross_attention - - if self.added_kv_proj_dim is None and self.only_cross_attention: - raise ValueError( - "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None. Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`." - ) - - if norm_num_groups is not None: - self.group_norm = nn.GroupNorm(num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True) - else: - self.group_norm = None - - if spatial_norm_dim is not None: - self.spatial_norm = SpatialNorm(f_channels=query_dim, zq_channels=spatial_norm_dim) - else: - self.spatial_norm = None - - if qk_norm is None: - self.norm_q = None - self.norm_k = None - elif qk_norm == "layer_norm": - self.norm_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) - self.norm_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) - elif qk_norm == "fp32_layer_norm": - self.norm_q = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps) - self.norm_k = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps) - elif qk_norm == "layer_norm_across_heads": - # Lumina applies qk norm across all heads - self.norm_q = nn.LayerNorm(dim_head * heads, eps=eps) - self.norm_k = nn.LayerNorm(dim_head * kv_heads, eps=eps) - elif qk_norm == "rms_norm": - self.norm_q = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) - self.norm_k = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) - elif qk_norm == "rms_norm_across_heads": - # LTX applies qk norm across all heads - self.norm_q = RMSNorm(dim_head * heads, eps=eps) - self.norm_k = RMSNorm(dim_head * kv_heads, eps=eps) - elif qk_norm == "l2": - self.norm_q = LpNorm(p=2, dim=-1, eps=eps) - self.norm_k = LpNorm(p=2, dim=-1, eps=eps) - else: - raise ValueError( - f"unknown qk_norm: {qk_norm}. Should be one of None, 'layer_norm', 'fp32_layer_norm', 'layer_norm_across_heads', 'rms_norm', 'rms_norm_across_heads', 'l2'." - ) - - if cross_attention_norm is None: - self.norm_cross = None - elif cross_attention_norm == "layer_norm": - self.norm_cross = nn.LayerNorm(self.cross_attention_dim) - elif cross_attention_norm == "group_norm": - if self.added_kv_proj_dim is not None: - # The given `encoder_hidden_states` are initially of shape - # (batch_size, seq_len, added_kv_proj_dim) before being projected - # to (batch_size, seq_len, cross_attention_dim). The norm is applied - # before the projection, so we need to use `added_kv_proj_dim` as - # the number of channels for the group norm. - norm_cross_num_channels = added_kv_proj_dim - else: - norm_cross_num_channels = self.cross_attention_dim - - self.norm_cross = nn.GroupNorm( - num_channels=norm_cross_num_channels, num_groups=cross_attention_norm_num_groups, eps=1e-5, affine=True - ) - else: - raise ValueError( - f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'" - ) - - self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias) - - if not self.only_cross_attention: - # only relevant for the `AddedKVProcessor` classes - self.to_k = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias) - self.to_v = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias) - else: - self.to_k = None - self.to_v = None - - self.added_proj_bias = added_proj_bias - if self.added_kv_proj_dim is not None: - self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias) - self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias) - if self.context_pre_only is not None: - self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) - else: - self.add_q_proj = None - self.add_k_proj = None - self.add_v_proj = None - - if not self.pre_only: - self.to_out = nn.ModuleList([]) - self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) - self.to_out.append(nn.Dropout(dropout)) - else: - self.to_out = None - - if self.context_pre_only is not None and not self.context_pre_only: - self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias) - else: - self.to_add_out = None - - if qk_norm is not None and added_kv_proj_dim is not None: - if qk_norm == "layer_norm": - self.norm_added_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) - self.norm_added_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) - elif qk_norm == "fp32_layer_norm": - self.norm_added_q = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps) - self.norm_added_k = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps) - elif qk_norm == "rms_norm": - self.norm_added_q = RMSNorm(dim_head, eps=eps) - self.norm_added_k = RMSNorm(dim_head, eps=eps) - elif qk_norm == "rms_norm_across_heads": - # Wan applies qk norm across all heads - # Wan also doesn't apply a q norm - self.norm_added_q = None - self.norm_added_k = RMSNorm(dim_head * kv_heads, eps=eps) - else: - raise ValueError( - f"unknown qk_norm: {qk_norm}. Should be one of `None,'layer_norm','fp32_layer_norm','rms_norm'`" - ) - else: - self.norm_added_q = None - self.norm_added_k = None - - # set attention processor - # We use the AttnProcessor2_0 by default when torch 2.x is used which uses - # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention - # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1 - if processor is None: - processor = ( - AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor() - ) - self.set_processor(processor) - - def set_use_xla_flash_attention( - self, - use_xla_flash_attention: bool, - partition_spec: Optional[Tuple[Optional[str], ...]] = None, - is_flux=False, - ) -> None: - r""" - Set whether to use xla flash attention from `torch_xla` or not. - - Args: - use_xla_flash_attention (`bool`): - Whether to use pallas flash attention kernel from `torch_xla` or not. - partition_spec (`Tuple[]`, *optional*): - Specify the partition specification if using SPMD. Otherwise None. - """ - if use_xla_flash_attention: - if not is_torch_xla_available: - raise "torch_xla is not available" - elif is_torch_xla_version("<", "2.3"): - raise "flash attention pallas kernel is supported from torch_xla version 2.3" - elif is_spmd() and is_torch_xla_version("<", "2.4"): - raise "flash attention pallas kernel using SPMD is supported from torch_xla version 2.4" - else: - if is_flux: - processor = XLAFluxFlashAttnProcessor2_0(partition_spec) - else: - processor = XLAFlashAttnProcessor2_0(partition_spec) - else: - processor = ( - AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor() - ) - self.set_processor(processor) - - def set_use_npu_flash_attention(self, use_npu_flash_attention: bool) -> None: - r""" - Set whether to use npu flash attention from `torch_npu` or not. - - """ - if use_npu_flash_attention: - processor = AttnProcessorNPU() - else: - # set attention processor - # We use the AttnProcessor2_0 by default when torch 2.x is used which uses - # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention - # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1 - processor = ( - AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor() - ) - self.set_processor(processor) - - def set_use_memory_efficient_attention_xformers( - self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None - ) -> None: - r""" - Set whether to use memory efficient attention from `xformers` or not. - - Args: - use_memory_efficient_attention_xformers (`bool`): - Whether to use memory efficient attention from `xformers` or not. - attention_op (`Callable`, *optional*): - The attention operation to use. Defaults to `None` which uses the default attention operation from - `xformers`. - """ - is_custom_diffusion = hasattr(self, "processor") and isinstance( - self.processor, - (CustomDiffusionAttnProcessor, CustomDiffusionXFormersAttnProcessor, CustomDiffusionAttnProcessor2_0), - ) - is_added_kv_processor = hasattr(self, "processor") and isinstance( - self.processor, - ( - AttnAddedKVProcessor, - AttnAddedKVProcessor2_0, - SlicedAttnAddedKVProcessor, - XFormersAttnAddedKVProcessor, - ), - ) - is_ip_adapter = hasattr(self, "processor") and isinstance( - self.processor, - (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0, IPAdapterXFormersAttnProcessor), - ) - is_joint_processor = hasattr(self, "processor") and isinstance( - self.processor, - ( - JointAttnProcessor2_0, - XFormersJointAttnProcessor, - ), - ) - - if use_memory_efficient_attention_xformers: - if is_added_kv_processor and is_custom_diffusion: - raise NotImplementedError( - f"Memory efficient attention is currently not supported for custom diffusion for attention processor type {self.processor}" - ) - if not is_xformers_available(): - raise ModuleNotFoundError( - ( - "Refer to https://github.com/facebookresearch/xformers for more information on how to install" - " xformers" - ), - name="xformers", - ) - elif not torch.cuda.is_available(): - raise ValueError( - "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is" - " only available for GPU " - ) - else: - try: - # Make sure we can run the memory efficient attention - dtype = None - if attention_op is not None: - op_fw, op_bw = attention_op - dtype, *_ = op_fw.SUPPORTED_DTYPES - q = torch.randn((1, 2, 40), device="cuda", dtype=dtype) - _ = xformers.ops.memory_efficient_attention(q, q, q) - except Exception as e: - raise e - - if is_custom_diffusion: - processor = CustomDiffusionXFormersAttnProcessor( - train_kv=self.processor.train_kv, - train_q_out=self.processor.train_q_out, - hidden_size=self.processor.hidden_size, - cross_attention_dim=self.processor.cross_attention_dim, - attention_op=attention_op, - ) - processor.load_state_dict(self.processor.state_dict()) - if hasattr(self.processor, "to_k_custom_diffusion"): - processor.to(self.processor.to_k_custom_diffusion.weight.device) - elif is_added_kv_processor: - # TODO(Patrick, Suraj, William) - currently xformers doesn't work for UnCLIP - # which uses this type of cross attention ONLY because the attention mask of format - # [0, ..., -10.000, ..., 0, ...,] is not supported - # throw warning - logger.info( - "Memory efficient attention with `xformers` might currently not work correctly if an attention mask is required for the attention operation." - ) - processor = XFormersAttnAddedKVProcessor(attention_op=attention_op) - elif is_ip_adapter: - processor = IPAdapterXFormersAttnProcessor( - hidden_size=self.processor.hidden_size, - cross_attention_dim=self.processor.cross_attention_dim, - num_tokens=self.processor.num_tokens, - scale=self.processor.scale, - attention_op=attention_op, - ) - processor.load_state_dict(self.processor.state_dict()) - if hasattr(self.processor, "to_k_ip"): - processor.to( - device=self.processor.to_k_ip[0].weight.device, dtype=self.processor.to_k_ip[0].weight.dtype - ) - elif is_joint_processor: - processor = XFormersJointAttnProcessor(attention_op=attention_op) - else: - processor = XFormersAttnProcessor(attention_op=attention_op) - else: - if is_custom_diffusion: - attn_processor_class = ( - CustomDiffusionAttnProcessor2_0 - if hasattr(F, "scaled_dot_product_attention") - else CustomDiffusionAttnProcessor - ) - processor = attn_processor_class( - train_kv=self.processor.train_kv, - train_q_out=self.processor.train_q_out, - hidden_size=self.processor.hidden_size, - cross_attention_dim=self.processor.cross_attention_dim, - ) - processor.load_state_dict(self.processor.state_dict()) - if hasattr(self.processor, "to_k_custom_diffusion"): - processor.to(self.processor.to_k_custom_diffusion.weight.device) - elif is_ip_adapter: - processor = IPAdapterAttnProcessor2_0( - hidden_size=self.processor.hidden_size, - cross_attention_dim=self.processor.cross_attention_dim, - num_tokens=self.processor.num_tokens, - scale=self.processor.scale, - ) - processor.load_state_dict(self.processor.state_dict()) - if hasattr(self.processor, "to_k_ip"): - processor.to( - device=self.processor.to_k_ip[0].weight.device, dtype=self.processor.to_k_ip[0].weight.dtype - ) - else: - # set attention processor - # We use the AttnProcessor2_0 by default when torch 2.x is used which uses - # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention - # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1 - processor = ( - AttnProcessor2_0() - if hasattr(F, "scaled_dot_product_attention") and self.scale_qk - else AttnProcessor() - ) - - self.set_processor(processor) - - def set_attention_slice(self, slice_size: int) -> None: - r""" - Set the slice size for attention computation. - - Args: - slice_size (`int`): - The slice size for attention computation. - """ - if slice_size is not None and slice_size > self.sliceable_head_dim: - raise ValueError(f"slice_size {slice_size} has to be smaller or equal to {self.sliceable_head_dim}.") - - if slice_size is not None and self.added_kv_proj_dim is not None: - processor = SlicedAttnAddedKVProcessor(slice_size) - elif slice_size is not None: - processor = SlicedAttnProcessor(slice_size) - elif self.added_kv_proj_dim is not None: - processor = AttnAddedKVProcessor() - else: - # set attention processor - # We use the AttnProcessor2_0 by default when torch 2.x is used which uses - # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention - # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1 - processor = ( - AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor() - ) - - self.set_processor(processor) - - def set_processor(self, processor: "AttnProcessor") -> None: - r""" - Set the attention processor to use. - - Args: - processor (`AttnProcessor`): - The attention processor to use. - """ - # if current processor is in `self._modules` and if passed `processor` is not, we need to - # pop `processor` from `self._modules` - if ( - hasattr(self, "processor") - and isinstance(self.processor, torch.nn.Module) - and not isinstance(processor, torch.nn.Module) - ): - logger.info(f"You are removing possibly trained weights of {self.processor} with {processor}") - self._modules.pop("processor") - - self.processor = processor - - def get_processor(self, return_deprecated_lora: bool = False) -> "AttentionProcessor": - r""" - Get the attention processor in use. - - Args: - return_deprecated_lora (`bool`, *optional*, defaults to `False`): - Set to `True` to return the deprecated LoRA attention processor. - - Returns: - "AttentionProcessor": The attention processor in use. - """ - if not return_deprecated_lora: - return self.processor - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - **cross_attention_kwargs, - ) -> torch.Tensor: - r""" - The forward method of the `Attention` class. - - Args: - hidden_states (`torch.Tensor`): - The hidden states of the query. - encoder_hidden_states (`torch.Tensor`, *optional*): - The hidden states of the encoder. - attention_mask (`torch.Tensor`, *optional*): - The attention mask to use. If `None`, no mask is applied. - **cross_attention_kwargs: - Additional keyword arguments to pass along to the cross attention. - - Returns: - `torch.Tensor`: The output of the attention layer. - """ - # The `Attention` class can call different attention processors / attention functions - # here we simply pass along all tensors to the selected processor class - # For standard processors that are defined here, `**cross_attention_kwargs` is empty - - attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) - quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"} - unused_kwargs = [ - k for k, _ in cross_attention_kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters - ] - if len(unused_kwargs) > 0: - logger.warning( - f"cross_attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored." - ) - cross_attention_kwargs = {k: w for k, w in cross_attention_kwargs.items() if k in attn_parameters} - - return self.processor( - self, - hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - **cross_attention_kwargs, - ) - - def batch_to_head_dim(self, tensor: torch.Tensor) -> torch.Tensor: - r""" - Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size // heads, seq_len, dim * heads]`. `heads` - is the number of heads initialized while constructing the `Attention` class. - - Args: - tensor (`torch.Tensor`): The tensor to reshape. - - Returns: - `torch.Tensor`: The reshaped tensor. - """ - head_size = self.heads - batch_size, seq_len, dim = tensor.shape - tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim) - tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size) - return tensor - - def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Tensor: - r""" - Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size, seq_len, heads, dim // heads]` `heads` is - the number of heads initialized while constructing the `Attention` class. - - Args: - tensor (`torch.Tensor`): The tensor to reshape. - out_dim (`int`, *optional*, defaults to `3`): The output dimension of the tensor. If `3`, the tensor is - reshaped to `[batch_size * heads, seq_len, dim // heads]`. - - Returns: - `torch.Tensor`: The reshaped tensor. - """ - head_size = self.heads - if tensor.ndim == 3: - batch_size, seq_len, dim = tensor.shape - extra_dim = 1 - else: - batch_size, extra_dim, seq_len, dim = tensor.shape - tensor = tensor.reshape(batch_size, seq_len * extra_dim, head_size, dim // head_size) - tensor = tensor.permute(0, 2, 1, 3) - - if out_dim == 3: - tensor = tensor.reshape(batch_size * head_size, seq_len * extra_dim, dim // head_size) - - return tensor - - def get_attention_scores( - self, query: torch.Tensor, key: torch.Tensor, attention_mask: Optional[torch.Tensor] = None - ) -> torch.Tensor: - r""" - Compute the attention scores. - - Args: - query (`torch.Tensor`): The query tensor. - key (`torch.Tensor`): The key tensor. - attention_mask (`torch.Tensor`, *optional*): The attention mask to use. If `None`, no mask is applied. - - Returns: - `torch.Tensor`: The attention probabilities/scores. - """ - dtype = query.dtype - if self.upcast_attention: - query = query.float() - key = key.float() - - if attention_mask is None: - baddbmm_input = torch.empty( - query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device - ) - beta = 0 - else: - baddbmm_input = attention_mask - beta = 1 - - attention_scores = torch.baddbmm( - baddbmm_input, - query, - key.transpose(-1, -2), - beta=beta, - alpha=self.scale, - ) - del baddbmm_input - - if self.upcast_softmax: - attention_scores = attention_scores.float() - - attention_probs = attention_scores.softmax(dim=-1) - del attention_scores - - attention_probs = attention_probs.to(dtype) - - return attention_probs - - def prepare_attention_mask( - self, attention_mask: torch.Tensor, target_length: int, batch_size: int, out_dim: int = 3 - ) -> torch.Tensor: - r""" - Prepare the attention mask for the attention computation. - - Args: - attention_mask (`torch.Tensor`): - The attention mask to prepare. - target_length (`int`): - The target length of the attention mask. This is the length of the attention mask after padding. - batch_size (`int`): - The batch size, which is used to repeat the attention mask. - out_dim (`int`, *optional*, defaults to `3`): - The output dimension of the attention mask. Can be either `3` or `4`. - - Returns: - `torch.Tensor`: The prepared attention mask. - """ - head_size = self.heads - if attention_mask is None: - return attention_mask - - current_length: int = attention_mask.shape[-1] - if current_length != target_length: - if attention_mask.device.type == "mps": - # HACK: MPS: Does not support padding by greater than dimension of input tensor. - # Instead, we can manually construct the padding tensor. - padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length) - padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device) - attention_mask = torch.cat([attention_mask, padding], dim=2) - else: - # TODO: for pipelines such as stable-diffusion, padding cross-attn mask: - # we want to instead pad by (0, remaining_length), where remaining_length is: - # remaining_length: int = target_length - current_length - # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding - attention_mask = F.pad(attention_mask, (0, target_length), value=0.0) - - if out_dim == 3: - if attention_mask.shape[0] < batch_size * head_size: - attention_mask = attention_mask.repeat_interleave( - head_size, dim=0, output_size=attention_mask.shape[0] * head_size - ) - elif out_dim == 4: - attention_mask = attention_mask.unsqueeze(1) - attention_mask = attention_mask.repeat_interleave( - head_size, dim=1, output_size=attention_mask.shape[1] * head_size - ) - - return attention_mask - - def norm_encoder_hidden_states(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor: - r""" - Normalize the encoder hidden states. Requires `self.norm_cross` to be specified when constructing the - `Attention` class. - - Args: - encoder_hidden_states (`torch.Tensor`): Hidden states of the encoder. - - Returns: - `torch.Tensor`: The normalized encoder hidden states. - """ - assert self.norm_cross is not None, "self.norm_cross must be defined to call self.norm_encoder_hidden_states" - - if isinstance(self.norm_cross, nn.LayerNorm): - encoder_hidden_states = self.norm_cross(encoder_hidden_states) - elif isinstance(self.norm_cross, nn.GroupNorm): - # Group norm norms along the channels dimension and expects - # input to be in the shape of (N, C, *). In this case, we want - # to norm along the hidden dimension, so we need to move - # (batch_size, sequence_length, hidden_size) -> - # (batch_size, hidden_size, sequence_length) - encoder_hidden_states = encoder_hidden_states.transpose(1, 2) - encoder_hidden_states = self.norm_cross(encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states.transpose(1, 2) - else: - assert False - - return encoder_hidden_states - - @torch.no_grad() - def fuse_projections(self, fuse=True): - device = self.to_q.weight.data.device - dtype = self.to_q.weight.data.dtype - - if not self.is_cross_attention: - # fetch weight matrices. - concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data]) - in_features = concatenated_weights.shape[1] - out_features = concatenated_weights.shape[0] - - # create a new single projection layer and copy over the weights. - self.to_qkv = nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype) - self.to_qkv.weight.copy_(concatenated_weights) - if self.use_bias: - concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data]) - self.to_qkv.bias.copy_(concatenated_bias) - - else: - concatenated_weights = torch.cat([self.to_k.weight.data, self.to_v.weight.data]) - in_features = concatenated_weights.shape[1] - out_features = concatenated_weights.shape[0] - - self.to_kv = nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype) - self.to_kv.weight.copy_(concatenated_weights) - if self.use_bias: - concatenated_bias = torch.cat([self.to_k.bias.data, self.to_v.bias.data]) - self.to_kv.bias.copy_(concatenated_bias) - - # handle added projections for SD3 and others. - if ( - getattr(self, "add_q_proj", None) is not None - and getattr(self, "add_k_proj", None) is not None - and getattr(self, "add_v_proj", None) is not None - ): - concatenated_weights = torch.cat( - [self.add_q_proj.weight.data, self.add_k_proj.weight.data, self.add_v_proj.weight.data] - ) - in_features = concatenated_weights.shape[1] - out_features = concatenated_weights.shape[0] - - self.to_added_qkv = nn.Linear( - in_features, out_features, bias=self.added_proj_bias, device=device, dtype=dtype - ) - self.to_added_qkv.weight.copy_(concatenated_weights) - if self.added_proj_bias: - concatenated_bias = torch.cat( - [self.add_q_proj.bias.data, self.add_k_proj.bias.data, self.add_v_proj.bias.data] - ) - self.to_added_qkv.bias.copy_(concatenated_bias) - - self.fused_projections = fuse - - -class SanaMultiscaleAttentionProjection(nn.Module): - def __init__( - self, - in_channels: int, - num_attention_heads: int, - kernel_size: int, - ) -> None: - super().__init__() - - channels = 3 * in_channels - self.proj_in = nn.Conv2d( - channels, - channels, - kernel_size, - padding=kernel_size // 2, - groups=channels, - bias=False, - ) - self.proj_out = nn.Conv2d(channels, channels, 1, 1, 0, groups=3 * num_attention_heads, bias=False) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.proj_in(hidden_states) - hidden_states = self.proj_out(hidden_states) - return hidden_states - - -class SanaMultiscaleLinearAttention(nn.Module): - r"""Lightweight multi-scale linear attention""" - - def __init__( - self, - in_channels: int, - out_channels: int, - num_attention_heads: Optional[int] = None, - attention_head_dim: int = 8, - mult: float = 1.0, - norm_type: str = "batch_norm", - kernel_sizes: Tuple[int, ...] = (5,), - eps: float = 1e-15, - residual_connection: bool = False, - ): - super().__init__() - - # To prevent circular import - from .normalization import get_normalization - - self.eps = eps - self.attention_head_dim = attention_head_dim - self.norm_type = norm_type - self.residual_connection = residual_connection - - num_attention_heads = ( - int(in_channels // attention_head_dim * mult) if num_attention_heads is None else num_attention_heads - ) - inner_dim = num_attention_heads * attention_head_dim - - self.to_q = nn.Linear(in_channels, inner_dim, bias=False) - self.to_k = nn.Linear(in_channels, inner_dim, bias=False) - self.to_v = nn.Linear(in_channels, inner_dim, bias=False) - - self.to_qkv_multiscale = nn.ModuleList() - for kernel_size in kernel_sizes: - self.to_qkv_multiscale.append( - SanaMultiscaleAttentionProjection(inner_dim, num_attention_heads, kernel_size) - ) - - self.nonlinearity = nn.ReLU() - self.to_out = nn.Linear(inner_dim * (1 + len(kernel_sizes)), out_channels, bias=False) - self.norm_out = get_normalization(norm_type, num_features=out_channels) - - self.processor = SanaMultiscaleAttnProcessor2_0() - - def apply_linear_attention(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: - value = F.pad(value, (0, 0, 0, 1), mode="constant", value=1) # Adds padding - scores = torch.matmul(value, key.transpose(-1, -2)) - hidden_states = torch.matmul(scores, query) - - hidden_states = hidden_states.to(dtype=torch.float32) - hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + self.eps) - return hidden_states - - def apply_quadratic_attention(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: - scores = torch.matmul(key.transpose(-1, -2), query) - scores = scores.to(dtype=torch.float32) - scores = scores / (torch.sum(scores, dim=2, keepdim=True) + self.eps) - hidden_states = torch.matmul(value, scores.to(value.dtype)) - return hidden_states - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.processor(self, hidden_states) - - -class MochiAttention(nn.Module): - def __init__( - self, - query_dim: int, - added_kv_proj_dim: int, - processor: "MochiAttnProcessor2_0", - heads: int = 8, - dim_head: int = 64, - dropout: float = 0.0, - bias: bool = False, - added_proj_bias: bool = True, - out_dim: Optional[int] = None, - out_context_dim: Optional[int] = None, - out_bias: bool = True, - context_pre_only: bool = False, - eps: float = 1e-5, - ): - super().__init__() - from .normalization import MochiRMSNorm - - self.inner_dim = out_dim if out_dim is not None else dim_head * heads - self.out_dim = out_dim if out_dim is not None else query_dim - self.out_context_dim = out_context_dim if out_context_dim else query_dim - self.context_pre_only = context_pre_only - - self.heads = out_dim // dim_head if out_dim is not None else heads - - self.norm_q = MochiRMSNorm(dim_head, eps, True) - self.norm_k = MochiRMSNorm(dim_head, eps, True) - self.norm_added_q = MochiRMSNorm(dim_head, eps, True) - self.norm_added_k = MochiRMSNorm(dim_head, eps, True) - - self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias) - self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias) - self.to_v = nn.Linear(query_dim, self.inner_dim, bias=bias) - - self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) - self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) - if self.context_pre_only is not None: - self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) - - self.to_out = nn.ModuleList([]) - self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) - self.to_out.append(nn.Dropout(dropout)) - - if not self.context_pre_only: - self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias) - - self.processor = processor - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - **kwargs, - ): - return self.processor( - self, - hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - **kwargs, - ) - - -class MochiAttnProcessor2_0: - """Attention processor used in Mochi.""" - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("MochiAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.") - - def __call__( - self, - attn: "MochiAttention", - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - attention_mask: torch.Tensor, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - 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)) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - encoder_query = attn.add_q_proj(encoder_hidden_states) - encoder_key = attn.add_k_proj(encoder_hidden_states) - encoder_value = attn.add_v_proj(encoder_hidden_states) - - encoder_query = encoder_query.unflatten(2, (attn.heads, -1)) - encoder_key = encoder_key.unflatten(2, (attn.heads, -1)) - encoder_value = encoder_value.unflatten(2, (attn.heads, -1)) - - if attn.norm_added_q is not None: - encoder_query = attn.norm_added_q(encoder_query) - if attn.norm_added_k is not None: - encoder_key = attn.norm_added_k(encoder_key) - - if image_rotary_emb is not None: - - def apply_rotary_emb(x, freqs_cos, freqs_sin): - x_even = x[..., 0::2].float() - x_odd = x[..., 1::2].float() - - cos = (x_even * freqs_cos - x_odd * freqs_sin).to(x.dtype) - sin = (x_even * freqs_sin + x_odd * freqs_cos).to(x.dtype) - - return torch.stack([cos, sin], dim=-1).flatten(-2) - - query = apply_rotary_emb(query, *image_rotary_emb) - key = apply_rotary_emb(key, *image_rotary_emb) - - query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) - encoder_query, encoder_key, encoder_value = ( - encoder_query.transpose(1, 2), - encoder_key.transpose(1, 2), - encoder_value.transpose(1, 2), - ) - - sequence_length = query.size(2) - encoder_sequence_length = encoder_query.size(2) - total_length = sequence_length + encoder_sequence_length - - batch_size, heads, _, dim = query.shape - attn_outputs = [] - for idx in range(batch_size): - mask = attention_mask[idx][None, :] - valid_prompt_token_indices = torch.nonzero(mask.flatten(), as_tuple=False).flatten() - - valid_encoder_query = encoder_query[idx : idx + 1, :, valid_prompt_token_indices, :] - valid_encoder_key = encoder_key[idx : idx + 1, :, valid_prompt_token_indices, :] - valid_encoder_value = encoder_value[idx : idx + 1, :, valid_prompt_token_indices, :] - - valid_query = torch.cat([query[idx : idx + 1], valid_encoder_query], dim=2) - valid_key = torch.cat([key[idx : idx + 1], valid_encoder_key], dim=2) - valid_value = torch.cat([value[idx : idx + 1], valid_encoder_value], dim=2) - - attn_output = F.scaled_dot_product_attention( - valid_query, valid_key, valid_value, dropout_p=0.0, is_causal=False - ) - valid_sequence_length = attn_output.size(2) - attn_output = F.pad(attn_output, (0, 0, 0, total_length - valid_sequence_length)) - attn_outputs.append(attn_output) - - hidden_states = torch.cat(attn_outputs, dim=0) - hidden_states = hidden_states.transpose(1, 2).flatten(2, 3) - - hidden_states, encoder_hidden_states = hidden_states.split_with_sizes( - (sequence_length, encoder_sequence_length), dim=1 - ) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if hasattr(attn, "to_add_out"): - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - return hidden_states, encoder_hidden_states - - -class AttnProcessor: - r""" - Default processor for performing attention-related computations. - """ - - 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, - *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) - - residual = hidden_states - - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query = attn.head_to_batch_dim(query) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - - attention_probs = attn.get_attention_scores(query, key, attention_mask) - hidden_states = torch.bmm(attention_probs, value) - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class CustomDiffusionAttnProcessor(nn.Module): - r""" - Processor for implementing attention for the Custom Diffusion method. - - Args: - train_kv (`bool`, defaults to `True`): - Whether to newly train the key and value matrices corresponding to the text features. - train_q_out (`bool`, defaults to `True`): - Whether to newly train query matrices corresponding to the latent image features. - hidden_size (`int`, *optional*, defaults to `None`): - The hidden size of the attention layer. - cross_attention_dim (`int`, *optional*, defaults to `None`): - The number of channels in the `encoder_hidden_states`. - out_bias (`bool`, defaults to `True`): - Whether to include the bias parameter in `train_q_out`. - dropout (`float`, *optional*, defaults to 0.0): - The dropout probability to use. - """ - - def __init__( - self, - train_kv: bool = True, - train_q_out: bool = True, - hidden_size: Optional[int] = None, - cross_attention_dim: Optional[int] = None, - out_bias: bool = True, - dropout: float = 0.0, - ): - super().__init__() - self.train_kv = train_kv - self.train_q_out = train_q_out - - self.hidden_size = hidden_size - self.cross_attention_dim = cross_attention_dim - - # `_custom_diffusion` id for easy serialization and loading. - if self.train_kv: - self.to_k_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) - self.to_v_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) - if self.train_q_out: - self.to_q_custom_diffusion = nn.Linear(hidden_size, hidden_size, bias=False) - self.to_out_custom_diffusion = nn.ModuleList([]) - self.to_out_custom_diffusion.append(nn.Linear(hidden_size, hidden_size, bias=out_bias)) - self.to_out_custom_diffusion.append(nn.Dropout(dropout)) - - def __call__( - self, - attn: Attention, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - batch_size, sequence_length, _ = hidden_states.shape - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - if self.train_q_out: - query = self.to_q_custom_diffusion(hidden_states).to(attn.to_q.weight.dtype) - else: - query = attn.to_q(hidden_states.to(attn.to_q.weight.dtype)) - - if encoder_hidden_states is None: - crossattn = False - encoder_hidden_states = hidden_states - else: - crossattn = True - if attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - if self.train_kv: - key = self.to_k_custom_diffusion(encoder_hidden_states.to(self.to_k_custom_diffusion.weight.dtype)) - value = self.to_v_custom_diffusion(encoder_hidden_states.to(self.to_v_custom_diffusion.weight.dtype)) - key = key.to(attn.to_q.weight.dtype) - value = value.to(attn.to_q.weight.dtype) - else: - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - if crossattn: - detach = torch.ones_like(key) - detach[:, :1, :] = detach[:, :1, :] * 0.0 - key = detach * key + (1 - detach) * key.detach() - value = detach * value + (1 - detach) * value.detach() - - query = attn.head_to_batch_dim(query) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - - attention_probs = attn.get_attention_scores(query, key, attention_mask) - hidden_states = torch.bmm(attention_probs, value) - hidden_states = attn.batch_to_head_dim(hidden_states) - - if self.train_q_out: - # linear proj - hidden_states = self.to_out_custom_diffusion[0](hidden_states) - # dropout - hidden_states = self.to_out_custom_diffusion[1](hidden_states) - else: - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - return hidden_states - - -class AttnAddedKVProcessor: - r""" - Processor for performing attention-related computations with extra learnable key and value matrices for the text - encoder. - """ - - def __call__( - self, - attn: Attention, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = 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) - - residual = hidden_states - - hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2) - batch_size, sequence_length, _ = hidden_states.shape - - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - query = attn.head_to_batch_dim(query) - - encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) - encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) - encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj) - encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj) - - if not attn.only_cross_attention: - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - key = torch.cat([encoder_hidden_states_key_proj, key], dim=1) - value = torch.cat([encoder_hidden_states_value_proj, value], dim=1) - else: - key = encoder_hidden_states_key_proj - value = encoder_hidden_states_value_proj - - attention_probs = attn.get_attention_scores(query, key, attention_mask) - hidden_states = torch.bmm(attention_probs, value) - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape) - hidden_states = hidden_states + residual - - return hidden_states - - -class AttnAddedKVProcessor2_0: - r""" - Processor for performing scaled dot-product attention (enabled by default if you're using PyTorch 2.0), with extra - learnable key and value matrices for the text encoder. - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "AttnAddedKVProcessor2_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, - *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) - - residual = hidden_states - - hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2) - batch_size, sequence_length, _ = hidden_states.shape - - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size, out_dim=4) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - query = attn.head_to_batch_dim(query, out_dim=4) - - encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) - encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) - encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj, out_dim=4) - encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj, out_dim=4) - - if not attn.only_cross_attention: - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - key = attn.head_to_batch_dim(key, out_dim=4) - value = attn.head_to_batch_dim(value, out_dim=4) - key = torch.cat([encoder_hidden_states_key_proj, key], dim=2) - value = torch.cat([encoder_hidden_states_value_proj, value], dim=2) - else: - key = encoder_hidden_states_key_proj - value = encoder_hidden_states_value_proj - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, residual.shape[1]) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape) - hidden_states = hidden_states + residual - - return hidden_states - - -class JointAttnProcessor2_0: - """Attention processor used typically in processing the SD3-like self-attention projections.""" - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("JointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, - *args, - **kwargs, - ) -> torch.FloatTensor: - residual = hidden_states - - batch_size = hidden_states.shape[0] - - # `sample` projections. - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # `context` projections. - if encoder_hidden_states is not None: - encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states) - encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) - encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) - - encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - - if attn.norm_added_q is not None: - encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj) - if attn.norm_added_k is not None: - encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj) - - query = torch.cat([query, encoder_hidden_states_query_proj], dim=2) - key = torch.cat([key, encoder_hidden_states_key_proj], dim=2) - value = torch.cat([value, encoder_hidden_states_value_proj], dim=2) - - hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - if encoder_hidden_states is not None: - # Split the attention outputs. - hidden_states, encoder_hidden_states = ( - hidden_states[:, : residual.shape[1]], - hidden_states[:, residual.shape[1] :], - ) - if not attn.context_pre_only: - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if encoder_hidden_states is not None: - return hidden_states, encoder_hidden_states - else: - return hidden_states - - -class PAGJointAttnProcessor2_0: - """Attention processor used typically in processing the SD3-like self-attention projections.""" - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "PAGJointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." - ) - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, - ) -> torch.FloatTensor: - residual = hidden_states - - input_ndim = hidden_states.ndim - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - context_input_ndim = encoder_hidden_states.ndim - if context_input_ndim == 4: - batch_size, channel, height, width = encoder_hidden_states.shape - encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - # store the length of image patch sequences to create a mask that prevents interaction between patches - # similar to making the self-attention map an identity matrix - identity_block_size = hidden_states.shape[1] - - # chunk - hidden_states_org, hidden_states_ptb = hidden_states.chunk(2) - encoder_hidden_states_org, encoder_hidden_states_ptb = encoder_hidden_states.chunk(2) - - ################## original path ################## - batch_size = encoder_hidden_states_org.shape[0] - - # `sample` projections. - query_org = attn.to_q(hidden_states_org) - key_org = attn.to_k(hidden_states_org) - value_org = attn.to_v(hidden_states_org) - - # `context` projections. - encoder_hidden_states_org_query_proj = attn.add_q_proj(encoder_hidden_states_org) - encoder_hidden_states_org_key_proj = attn.add_k_proj(encoder_hidden_states_org) - encoder_hidden_states_org_value_proj = attn.add_v_proj(encoder_hidden_states_org) - - # attention - query_org = torch.cat([query_org, encoder_hidden_states_org_query_proj], dim=1) - key_org = torch.cat([key_org, encoder_hidden_states_org_key_proj], dim=1) - value_org = torch.cat([value_org, encoder_hidden_states_org_value_proj], dim=1) - - inner_dim = key_org.shape[-1] - head_dim = inner_dim // attn.heads - query_org = query_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key_org = key_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value_org = value_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - hidden_states_org = F.scaled_dot_product_attention( - query_org, key_org, value_org, dropout_p=0.0, is_causal=False - ) - hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states_org = hidden_states_org.to(query_org.dtype) - - # Split the attention outputs. - hidden_states_org, encoder_hidden_states_org = ( - hidden_states_org[:, : residual.shape[1]], - hidden_states_org[:, residual.shape[1] :], - ) - - # linear proj - hidden_states_org = attn.to_out[0](hidden_states_org) - # dropout - hidden_states_org = attn.to_out[1](hidden_states_org) - if not attn.context_pre_only: - encoder_hidden_states_org = attn.to_add_out(encoder_hidden_states_org) - - if input_ndim == 4: - hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) - if context_input_ndim == 4: - encoder_hidden_states_org = encoder_hidden_states_org.transpose(-1, -2).reshape( - batch_size, channel, height, width - ) - - ################## perturbed path ################## - - batch_size = encoder_hidden_states_ptb.shape[0] - - # `sample` projections. - query_ptb = attn.to_q(hidden_states_ptb) - key_ptb = attn.to_k(hidden_states_ptb) - value_ptb = attn.to_v(hidden_states_ptb) - - # `context` projections. - encoder_hidden_states_ptb_query_proj = attn.add_q_proj(encoder_hidden_states_ptb) - encoder_hidden_states_ptb_key_proj = attn.add_k_proj(encoder_hidden_states_ptb) - encoder_hidden_states_ptb_value_proj = attn.add_v_proj(encoder_hidden_states_ptb) - - # attention - query_ptb = torch.cat([query_ptb, encoder_hidden_states_ptb_query_proj], dim=1) - key_ptb = torch.cat([key_ptb, encoder_hidden_states_ptb_key_proj], dim=1) - value_ptb = torch.cat([value_ptb, encoder_hidden_states_ptb_value_proj], dim=1) - - inner_dim = key_ptb.shape[-1] - head_dim = inner_dim // attn.heads - query_ptb = query_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key_ptb = key_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value_ptb = value_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # create a full mask with all entries set to 0 - seq_len = query_ptb.size(2) - full_mask = torch.zeros((seq_len, seq_len), device=query_ptb.device, dtype=query_ptb.dtype) - - # set the attention value between image patches to -inf - full_mask[:identity_block_size, :identity_block_size] = float("-inf") - - # set the diagonal of the attention value between image patches to 0 - full_mask[:identity_block_size, :identity_block_size].fill_diagonal_(0) - - # expand the mask to match the attention weights shape - full_mask = full_mask.unsqueeze(0).unsqueeze(0) # Add batch and num_heads dimensions - - hidden_states_ptb = F.scaled_dot_product_attention( - query_ptb, key_ptb, value_ptb, attn_mask=full_mask, dropout_p=0.0, is_causal=False - ) - hidden_states_ptb = hidden_states_ptb.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states_ptb = hidden_states_ptb.to(query_ptb.dtype) - - # split the attention outputs. - hidden_states_ptb, encoder_hidden_states_ptb = ( - hidden_states_ptb[:, : residual.shape[1]], - hidden_states_ptb[:, residual.shape[1] :], - ) - - # linear proj - hidden_states_ptb = attn.to_out[0](hidden_states_ptb) - # dropout - hidden_states_ptb = attn.to_out[1](hidden_states_ptb) - if not attn.context_pre_only: - encoder_hidden_states_ptb = attn.to_add_out(encoder_hidden_states_ptb) - - if input_ndim == 4: - hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width) - if context_input_ndim == 4: - encoder_hidden_states_ptb = encoder_hidden_states_ptb.transpose(-1, -2).reshape( - batch_size, channel, height, width - ) - - ################ concat ############### - hidden_states = torch.cat([hidden_states_org, hidden_states_ptb]) - encoder_hidden_states = torch.cat([encoder_hidden_states_org, encoder_hidden_states_ptb]) - - return hidden_states, encoder_hidden_states - - -class PAGCFGJointAttnProcessor2_0: - """Attention processor used typically in processing the SD3-like self-attention projections.""" - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "PAGCFGJointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." - ) - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, - *args, - **kwargs, - ) -> torch.FloatTensor: - residual = hidden_states - - input_ndim = hidden_states.ndim - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - context_input_ndim = encoder_hidden_states.ndim - if context_input_ndim == 4: - batch_size, channel, height, width = encoder_hidden_states.shape - encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - identity_block_size = hidden_states.shape[ - 1 - ] # patch embeddings width * height (correspond to self-attention map width or height) - - # chunk - hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3) - hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org]) - - ( - encoder_hidden_states_uncond, - encoder_hidden_states_org, - encoder_hidden_states_ptb, - ) = encoder_hidden_states.chunk(3) - encoder_hidden_states_org = torch.cat([encoder_hidden_states_uncond, encoder_hidden_states_org]) - - ################## original path ################## - batch_size = encoder_hidden_states_org.shape[0] - - # `sample` projections. - query_org = attn.to_q(hidden_states_org) - key_org = attn.to_k(hidden_states_org) - value_org = attn.to_v(hidden_states_org) - - # `context` projections. - encoder_hidden_states_org_query_proj = attn.add_q_proj(encoder_hidden_states_org) - encoder_hidden_states_org_key_proj = attn.add_k_proj(encoder_hidden_states_org) - encoder_hidden_states_org_value_proj = attn.add_v_proj(encoder_hidden_states_org) - - # attention - query_org = torch.cat([query_org, encoder_hidden_states_org_query_proj], dim=1) - key_org = torch.cat([key_org, encoder_hidden_states_org_key_proj], dim=1) - value_org = torch.cat([value_org, encoder_hidden_states_org_value_proj], dim=1) - - inner_dim = key_org.shape[-1] - head_dim = inner_dim // attn.heads - query_org = query_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key_org = key_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value_org = value_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - hidden_states_org = F.scaled_dot_product_attention( - query_org, key_org, value_org, dropout_p=0.0, is_causal=False - ) - hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states_org = hidden_states_org.to(query_org.dtype) - - # Split the attention outputs. - hidden_states_org, encoder_hidden_states_org = ( - hidden_states_org[:, : residual.shape[1]], - hidden_states_org[:, residual.shape[1] :], - ) - - # linear proj - hidden_states_org = attn.to_out[0](hidden_states_org) - # dropout - hidden_states_org = attn.to_out[1](hidden_states_org) - if not attn.context_pre_only: - encoder_hidden_states_org = attn.to_add_out(encoder_hidden_states_org) - - if input_ndim == 4: - hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) - if context_input_ndim == 4: - encoder_hidden_states_org = encoder_hidden_states_org.transpose(-1, -2).reshape( - batch_size, channel, height, width - ) - - ################## perturbed path ################## - - batch_size = encoder_hidden_states_ptb.shape[0] - - # `sample` projections. - query_ptb = attn.to_q(hidden_states_ptb) - key_ptb = attn.to_k(hidden_states_ptb) - value_ptb = attn.to_v(hidden_states_ptb) - - # `context` projections. - encoder_hidden_states_ptb_query_proj = attn.add_q_proj(encoder_hidden_states_ptb) - encoder_hidden_states_ptb_key_proj = attn.add_k_proj(encoder_hidden_states_ptb) - encoder_hidden_states_ptb_value_proj = attn.add_v_proj(encoder_hidden_states_ptb) - - # attention - query_ptb = torch.cat([query_ptb, encoder_hidden_states_ptb_query_proj], dim=1) - key_ptb = torch.cat([key_ptb, encoder_hidden_states_ptb_key_proj], dim=1) - value_ptb = torch.cat([value_ptb, encoder_hidden_states_ptb_value_proj], dim=1) - - inner_dim = key_ptb.shape[-1] - head_dim = inner_dim // attn.heads - query_ptb = query_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key_ptb = key_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value_ptb = value_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # create a full mask with all entries set to 0 - seq_len = query_ptb.size(2) - full_mask = torch.zeros((seq_len, seq_len), device=query_ptb.device, dtype=query_ptb.dtype) - - # set the attention value between image patches to -inf - full_mask[:identity_block_size, :identity_block_size] = float("-inf") - - # set the diagonal of the attention value between image patches to 0 - full_mask[:identity_block_size, :identity_block_size].fill_diagonal_(0) - - # expand the mask to match the attention weights shape - full_mask = full_mask.unsqueeze(0).unsqueeze(0) # Add batch and num_heads dimensions - - hidden_states_ptb = F.scaled_dot_product_attention( - query_ptb, key_ptb, value_ptb, attn_mask=full_mask, dropout_p=0.0, is_causal=False - ) - hidden_states_ptb = hidden_states_ptb.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states_ptb = hidden_states_ptb.to(query_ptb.dtype) - - # split the attention outputs. - hidden_states_ptb, encoder_hidden_states_ptb = ( - hidden_states_ptb[:, : residual.shape[1]], - hidden_states_ptb[:, residual.shape[1] :], - ) - - # linear proj - hidden_states_ptb = attn.to_out[0](hidden_states_ptb) - # dropout - hidden_states_ptb = attn.to_out[1](hidden_states_ptb) - if not attn.context_pre_only: - encoder_hidden_states_ptb = attn.to_add_out(encoder_hidden_states_ptb) - - if input_ndim == 4: - hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width) - if context_input_ndim == 4: - encoder_hidden_states_ptb = encoder_hidden_states_ptb.transpose(-1, -2).reshape( - batch_size, channel, height, width - ) - - ################ concat ############### - hidden_states = torch.cat([hidden_states_org, hidden_states_ptb]) - encoder_hidden_states = torch.cat([encoder_hidden_states_org, encoder_hidden_states_ptb]) - - return hidden_states, encoder_hidden_states - - -class FusedJointAttnProcessor2_0: - """Attention processor used typically in processing the SD3-like self-attention projections.""" - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, - *args, - **kwargs, - ) -> torch.FloatTensor: - residual = hidden_states - - input_ndim = hidden_states.ndim - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - context_input_ndim = encoder_hidden_states.ndim - if context_input_ndim == 4: - batch_size, channel, height, width = encoder_hidden_states.shape - encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - batch_size = encoder_hidden_states.shape[0] - - # `sample` projections. - qkv = attn.to_qkv(hidden_states) - split_size = qkv.shape[-1] // 3 - query, key, value = torch.split(qkv, split_size, dim=-1) - - # `context` projections. - encoder_qkv = attn.to_added_qkv(encoder_hidden_states) - split_size = encoder_qkv.shape[-1] // 3 - ( - encoder_hidden_states_query_proj, - encoder_hidden_states_key_proj, - encoder_hidden_states_value_proj, - ) = torch.split(encoder_qkv, split_size, dim=-1) - - # attention - query = torch.cat([query, encoder_hidden_states_query_proj], dim=1) - key = torch.cat([key, encoder_hidden_states_key_proj], dim=1) - value = torch.cat([value, encoder_hidden_states_value_proj], dim=1) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # Split the attention outputs. - hidden_states, encoder_hidden_states = ( - hidden_states[:, : residual.shape[1]], - hidden_states[:, residual.shape[1] :], - ) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - if not attn.context_pre_only: - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - if context_input_ndim == 4: - encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - return hidden_states, encoder_hidden_states - - -class XFormersJointAttnProcessor: - r""" - Processor for implementing memory efficient attention using xFormers. - - Args: - attention_op (`Callable`, *optional*, defaults to `None`): - The base - [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to - use as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best - operator. - """ - - def __init__(self, attention_op: Optional[Callable] = None): - self.attention_op = attention_op - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, - *args, - **kwargs, - ) -> torch.FloatTensor: - residual = hidden_states - - # `sample` projections. - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - query = attn.head_to_batch_dim(query).contiguous() - key = attn.head_to_batch_dim(key).contiguous() - value = attn.head_to_batch_dim(value).contiguous() - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # `context` projections. - if encoder_hidden_states is not None: - encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states) - encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) - encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) - - encoder_hidden_states_query_proj = attn.head_to_batch_dim(encoder_hidden_states_query_proj).contiguous() - encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj).contiguous() - encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj).contiguous() - - if attn.norm_added_q is not None: - encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj) - if attn.norm_added_k is not None: - encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj) - - query = torch.cat([query, encoder_hidden_states_query_proj], dim=1) - key = torch.cat([key, encoder_hidden_states_key_proj], dim=1) - value = torch.cat([value, encoder_hidden_states_value_proj], dim=1) - - hidden_states = xformers.ops.memory_efficient_attention( - query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale - ) - hidden_states = hidden_states.to(query.dtype) - hidden_states = attn.batch_to_head_dim(hidden_states) - - if encoder_hidden_states is not None: - # Split the attention outputs. - hidden_states, encoder_hidden_states = ( - hidden_states[:, : residual.shape[1]], - hidden_states[:, residual.shape[1] :], - ) - if not attn.context_pre_only: - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if encoder_hidden_states is not None: - return hidden_states, encoder_hidden_states - else: - return hidden_states - - -class AllegroAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is - used in the Allegro model. It applies a normalization layer and rotary embedding on the query and key vector. - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "AllegroAttnProcessor2_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, - temb: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - residual = hidden_states - - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - 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) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # Apply RoPE if needed - if image_rotary_emb is not None and not attn.is_cross_attention: - from .embeddings import apply_rotary_emb_allegro - - query = apply_rotary_emb_allegro(query, image_rotary_emb[0], image_rotary_emb[1]) - key = apply_rotary_emb_allegro(key, image_rotary_emb[0], image_rotary_emb[1]) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class AuraFlowAttnProcessor2_0: - """Attention processor used typically in processing Aura Flow.""" - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention") and is_torch_version("<", "2.1"): - raise ImportError( - "AuraFlowAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to at least 2.1 or above as we use `scale` in `F.scaled_dot_product_attention()`. " - ) - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - *args, - **kwargs, - ) -> torch.FloatTensor: - batch_size = hidden_states.shape[0] - - # `sample` projections. - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - # `context` projections. - if encoder_hidden_states is not None: - encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states) - encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) - encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) - - # Reshape. - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - query = query.view(batch_size, -1, attn.heads, head_dim) - key = key.view(batch_size, -1, attn.heads, head_dim) - value = value.view(batch_size, -1, attn.heads, head_dim) - - # Apply QK norm. - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # Concatenate the projections. - if encoder_hidden_states is not None: - encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( - batch_size, -1, attn.heads, head_dim - ) - encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(batch_size, -1, attn.heads, head_dim) - encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( - batch_size, -1, attn.heads, head_dim - ) - - if attn.norm_added_q is not None: - encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj) - if attn.norm_added_k is not None: - encoder_hidden_states_key_proj = attn.norm_added_q(encoder_hidden_states_key_proj) - - query = torch.cat([encoder_hidden_states_query_proj, query], dim=1) - key = torch.cat([encoder_hidden_states_key_proj, key], dim=1) - value = torch.cat([encoder_hidden_states_value_proj, value], dim=1) - - query = query.transpose(1, 2) - key = key.transpose(1, 2) - value = value.transpose(1, 2) - - # Attention. - hidden_states = F.scaled_dot_product_attention( - query, key, value, dropout_p=0.0, scale=attn.scale, is_causal=False - ) - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # Split the attention outputs. - if encoder_hidden_states is not None: - hidden_states, encoder_hidden_states = ( - hidden_states[:, encoder_hidden_states.shape[1] :], - hidden_states[:, : encoder_hidden_states.shape[1]], - ) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - if encoder_hidden_states is not None: - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - if encoder_hidden_states is not None: - return hidden_states, encoder_hidden_states - else: - return hidden_states - - -class FusedAuraFlowAttnProcessor2_0: - """Attention processor used typically in processing Aura Flow with fused projections.""" - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention") and is_torch_version("<", "2.1"): - raise ImportError( - "FusedAuraFlowAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to at least 2.1 or above as we use `scale` in `F.scaled_dot_product_attention()`. " - ) - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - *args, - **kwargs, - ) -> torch.FloatTensor: - batch_size = hidden_states.shape[0] - - # `sample` projections. - qkv = attn.to_qkv(hidden_states) - split_size = qkv.shape[-1] // 3 - query, key, value = torch.split(qkv, split_size, dim=-1) - - # `context` projections. - if encoder_hidden_states is not None: - encoder_qkv = attn.to_added_qkv(encoder_hidden_states) - split_size = encoder_qkv.shape[-1] // 3 - ( - encoder_hidden_states_query_proj, - encoder_hidden_states_key_proj, - encoder_hidden_states_value_proj, - ) = torch.split(encoder_qkv, split_size, dim=-1) - - # Reshape. - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - query = query.view(batch_size, -1, attn.heads, head_dim) - key = key.view(batch_size, -1, attn.heads, head_dim) - value = value.view(batch_size, -1, attn.heads, head_dim) - - # Apply QK norm. - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # Concatenate the projections. - if encoder_hidden_states is not None: - encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( - batch_size, -1, attn.heads, head_dim - ) - encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(batch_size, -1, attn.heads, head_dim) - encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( - batch_size, -1, attn.heads, head_dim - ) - - if attn.norm_added_q is not None: - encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj) - if attn.norm_added_k is not None: - encoder_hidden_states_key_proj = attn.norm_added_q(encoder_hidden_states_key_proj) - - query = torch.cat([encoder_hidden_states_query_proj, query], dim=1) - key = torch.cat([encoder_hidden_states_key_proj, key], dim=1) - value = torch.cat([encoder_hidden_states_value_proj, value], dim=1) - - query = query.transpose(1, 2) - key = key.transpose(1, 2) - value = value.transpose(1, 2) - - # Attention. - hidden_states = F.scaled_dot_product_attention( - query, key, value, dropout_p=0.0, scale=attn.scale, is_causal=False - ) - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # Split the attention outputs. - if encoder_hidden_states is not None: - hidden_states, encoder_hidden_states = ( - hidden_states[:, encoder_hidden_states.shape[1] :], - hidden_states[:, : encoder_hidden_states.shape[1]], - ) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - if encoder_hidden_states is not None: - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - if encoder_hidden_states is not None: - return hidden_states, encoder_hidden_states - else: - return hidden_states - - -class FluxAttnProcessor2_0: - """Attention processor used typically in processing the SD3-like self-attention projections.""" - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("FluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.FloatTensor: - batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - - # `sample` projections. - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states` - if encoder_hidden_states is not None: - # `context` projections. - encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states) - encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) - encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) - - encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - - if attn.norm_added_q is not None: - encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj) - if attn.norm_added_k is not None: - encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj) - - # attention - query = torch.cat([encoder_hidden_states_query_proj, query], dim=2) - key = torch.cat([encoder_hidden_states_key_proj, key], dim=2) - value = torch.cat([encoder_hidden_states_value_proj, value], dim=2) - - if image_rotary_emb is not None: - from .embeddings import apply_rotary_emb - - query = apply_rotary_emb(query, image_rotary_emb) - key = apply_rotary_emb(key, image_rotary_emb) - - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - if encoder_hidden_states is not None: - encoder_hidden_states, hidden_states = ( - hidden_states[:, : encoder_hidden_states.shape[1]], - hidden_states[:, encoder_hidden_states.shape[1] :], - ) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - return hidden_states, encoder_hidden_states - else: - return hidden_states - - -class FluxAttnProcessor2_0_NPU: - """Attention processor used typically in processing the SD3-like self-attention projections.""" - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "FluxAttnProcessor2_0_NPU requires PyTorch 2.0 and torch NPU, to use it, please upgrade PyTorch to 2.0 and install torch NPU" - ) - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.FloatTensor: - batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - - # `sample` projections. - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states` - if encoder_hidden_states is not None: - # `context` projections. - encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states) - encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) - encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) - - encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - - if attn.norm_added_q is not None: - encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj) - if attn.norm_added_k is not None: - encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj) - - # attention - query = torch.cat([encoder_hidden_states_query_proj, query], dim=2) - key = torch.cat([encoder_hidden_states_key_proj, key], dim=2) - value = torch.cat([encoder_hidden_states_value_proj, value], dim=2) - - if image_rotary_emb is not None: - from .embeddings import apply_rotary_emb - - query = apply_rotary_emb(query, image_rotary_emb) - key = apply_rotary_emb(key, image_rotary_emb) - - if query.dtype in (torch.float16, torch.bfloat16): - hidden_states = torch_npu.npu_fusion_attention( - query, - key, - value, - attn.heads, - input_layout="BNSD", - pse=None, - scale=1.0 / math.sqrt(query.shape[-1]), - pre_tockens=65536, - next_tockens=65536, - keep_prob=1.0, - sync=False, - inner_precise=0, - )[0] - else: - hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - if encoder_hidden_states is not None: - encoder_hidden_states, hidden_states = ( - hidden_states[:, : encoder_hidden_states.shape[1]], - hidden_states[:, encoder_hidden_states.shape[1] :], - ) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - return hidden_states, encoder_hidden_states - else: - return hidden_states - - -class FusedFluxAttnProcessor2_0: - """Attention processor used typically in processing the SD3-like self-attention projections.""" - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "FusedFluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." - ) - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.FloatTensor: - batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - - # `sample` projections. - qkv = attn.to_qkv(hidden_states) - split_size = qkv.shape[-1] // 3 - query, key, value = torch.split(qkv, split_size, dim=-1) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states` - # `context` projections. - if encoder_hidden_states is not None: - encoder_qkv = attn.to_added_qkv(encoder_hidden_states) - split_size = encoder_qkv.shape[-1] // 3 - ( - encoder_hidden_states_query_proj, - encoder_hidden_states_key_proj, - encoder_hidden_states_value_proj, - ) = torch.split(encoder_qkv, split_size, dim=-1) - - encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - - if attn.norm_added_q is not None: - encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj) - if attn.norm_added_k is not None: - encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj) - - # attention - query = torch.cat([encoder_hidden_states_query_proj, query], dim=2) - key = torch.cat([encoder_hidden_states_key_proj, key], dim=2) - value = torch.cat([encoder_hidden_states_value_proj, value], dim=2) - - if image_rotary_emb is not None: - from .embeddings import apply_rotary_emb - - query = apply_rotary_emb(query, image_rotary_emb) - key = apply_rotary_emb(key, image_rotary_emb) - - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - if encoder_hidden_states is not None: - encoder_hidden_states, hidden_states = ( - hidden_states[:, : encoder_hidden_states.shape[1]], - hidden_states[:, encoder_hidden_states.shape[1] :], - ) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - return hidden_states, encoder_hidden_states - else: - return hidden_states - - -class FusedFluxAttnProcessor2_0_NPU: - """Attention processor used typically in processing the SD3-like self-attention projections.""" - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "FluxAttnProcessor2_0_NPU requires PyTorch 2.0 and torch NPU, to use it, please upgrade PyTorch to 2.0, and install torch NPU" - ) - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.FloatTensor: - batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - - # `sample` projections. - qkv = attn.to_qkv(hidden_states) - split_size = qkv.shape[-1] // 3 - query, key, value = torch.split(qkv, split_size, dim=-1) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states` - # `context` projections. - if encoder_hidden_states is not None: - encoder_qkv = attn.to_added_qkv(encoder_hidden_states) - split_size = encoder_qkv.shape[-1] // 3 - ( - encoder_hidden_states_query_proj, - encoder_hidden_states_key_proj, - encoder_hidden_states_value_proj, - ) = torch.split(encoder_qkv, split_size, dim=-1) - - encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - - if attn.norm_added_q is not None: - encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj) - if attn.norm_added_k is not None: - encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj) - - # attention - query = torch.cat([encoder_hidden_states_query_proj, query], dim=2) - key = torch.cat([encoder_hidden_states_key_proj, key], dim=2) - value = torch.cat([encoder_hidden_states_value_proj, value], dim=2) - - if image_rotary_emb is not None: - from .embeddings import apply_rotary_emb - - query = apply_rotary_emb(query, image_rotary_emb) - key = apply_rotary_emb(key, image_rotary_emb) - - if query.dtype in (torch.float16, torch.bfloat16): - hidden_states = torch_npu.npu_fusion_attention( - query, - key, - value, - attn.heads, - input_layout="BNSD", - pse=None, - scale=1.0 / math.sqrt(query.shape[-1]), - pre_tockens=65536, - next_tockens=65536, - keep_prob=1.0, - sync=False, - inner_precise=0, - )[0] - else: - hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - if encoder_hidden_states is not None: - encoder_hidden_states, hidden_states = ( - hidden_states[:, : encoder_hidden_states.shape[1]], - hidden_states[:, encoder_hidden_states.shape[1] :], - ) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - return hidden_states, encoder_hidden_states - else: - return hidden_states - - -class FluxIPAdapterJointAttnProcessor2_0(torch.nn.Module): - """Flux Attention processor for IP-Adapter.""" - - def __init__( - self, hidden_size: int, cross_attention_dim: int, num_tokens=(4,), scale=1.0, device=None, dtype=None - ): - super().__init__() - - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - f"{self.__class__.__name__} requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." - ) - - self.hidden_size = hidden_size - self.cross_attention_dim = cross_attention_dim - - if not isinstance(num_tokens, (tuple, list)): - num_tokens = [num_tokens] - - if not isinstance(scale, list): - scale = [scale] * len(num_tokens) - if len(scale) != len(num_tokens): - raise ValueError("`scale` should be a list of integers with the same length as `num_tokens`.") - self.scale = scale - - self.to_k_ip = nn.ModuleList( - [ - nn.Linear(cross_attention_dim, hidden_size, bias=True, device=device, dtype=dtype) - for _ in range(len(num_tokens)) - ] - ) - self.to_v_ip = nn.ModuleList( - [ - nn.Linear(cross_attention_dim, hidden_size, bias=True, device=device, dtype=dtype) - for _ in range(len(num_tokens)) - ] - ) - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ip_hidden_states: Optional[List[torch.Tensor]] = None, - ip_adapter_masks: Optional[torch.Tensor] = None, - ) -> torch.FloatTensor: - batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - - # `sample` projections. - hidden_states_query_proj = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - hidden_states_query_proj = hidden_states_query_proj.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - hidden_states_query_proj = attn.norm_q(hidden_states_query_proj) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states` - if encoder_hidden_states is not None: - # `context` projections. - encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states) - encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) - encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) - - encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - - if attn.norm_added_q is not None: - encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj) - if attn.norm_added_k is not None: - encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj) - - # attention - query = torch.cat([encoder_hidden_states_query_proj, hidden_states_query_proj], dim=2) - key = torch.cat([encoder_hidden_states_key_proj, key], dim=2) - value = torch.cat([encoder_hidden_states_value_proj, value], dim=2) - - if image_rotary_emb is not None: - from .embeddings import apply_rotary_emb - - query = apply_rotary_emb(query, image_rotary_emb) - key = apply_rotary_emb(key, image_rotary_emb) - - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - if encoder_hidden_states is not None: - encoder_hidden_states, hidden_states = ( - hidden_states[:, : encoder_hidden_states.shape[1]], - hidden_states[:, encoder_hidden_states.shape[1] :], - ) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - # IP-adapter - ip_query = hidden_states_query_proj - ip_attn_output = torch.zeros_like(hidden_states) - - for current_ip_hidden_states, scale, to_k_ip, to_v_ip in zip( - ip_hidden_states, self.scale, self.to_k_ip, self.to_v_ip - ): - ip_key = to_k_ip(current_ip_hidden_states) - ip_value = to_v_ip(current_ip_hidden_states) - - ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - current_ip_hidden_states = F.scaled_dot_product_attention( - ip_query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False - ) - current_ip_hidden_states = current_ip_hidden_states.transpose(1, 2).reshape( - batch_size, -1, attn.heads * head_dim - ) - current_ip_hidden_states = current_ip_hidden_states.to(ip_query.dtype) - ip_attn_output += scale * current_ip_hidden_states - - return hidden_states, encoder_hidden_states, ip_attn_output - else: - return hidden_states - - -class CogVideoXAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on - query and key vectors, but does not include spatial normalization. - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("CogVideoXAttnProcessor 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: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - text_seq_length = encoder_hidden_states.size(1) - - hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) - - batch_size, sequence_length, _ = 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]) - - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # Apply RoPE if needed - 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) - if not attn.is_cross_attention: - key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb) - - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - 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 - - -class FusedCogVideoXAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on - query and key vectors, but does not include spatial normalization. - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("CogVideoXAttnProcessor 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: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - text_seq_length = encoder_hidden_states.size(1) - - hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) - - 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]) - - qkv = attn.to_qkv(hidden_states) - split_size = qkv.shape[-1] // 3 - query, key, value = torch.split(qkv, split_size, dim=-1) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # Apply RoPE if needed - 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) - if not attn.is_cross_attention: - key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb) - - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - 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 - - -class XFormersAttnAddedKVProcessor: - r""" - Processor for implementing memory efficient attention using xFormers. - - Args: - attention_op (`Callable`, *optional*, defaults to `None`): - The base - [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to - use as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best - operator. - """ - - def __init__(self, attention_op: Optional[Callable] = 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, - ) -> torch.Tensor: - residual = hidden_states - hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2) - batch_size, sequence_length, _ = hidden_states.shape - - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - query = attn.head_to_batch_dim(query) - - encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) - encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) - encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj) - encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj) - - if not attn.only_cross_attention: - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - key = torch.cat([encoder_hidden_states_key_proj, key], dim=1) - value = torch.cat([encoder_hidden_states_value_proj, value], dim=1) - else: - key = encoder_hidden_states_key_proj - value = encoder_hidden_states_value_proj - - hidden_states = xformers.ops.memory_efficient_attention( - query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale - ) - hidden_states = hidden_states.to(query.dtype) - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape) - hidden_states = hidden_states + residual - - return hidden_states - - -class XFormersAttnProcessor: - r""" - Processor for implementing memory efficient attention using xFormers. - - Args: - attention_op (`Callable`, *optional*, defaults to `None`): - The base - [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to - use as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best - operator. - """ - - def __init__(self, attention_op: Optional[Callable] = 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, - *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) - - residual = hidden_states - - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - batch_size, key_tokens, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - - attention_mask = attn.prepare_attention_mask(attention_mask, key_tokens, batch_size) - if attention_mask is not None: - # expand our mask's singleton query_tokens dimension: - # [batch*heads, 1, key_tokens] -> - # [batch*heads, query_tokens, key_tokens] - # so that it can be added as a bias onto the attention scores that xformers computes: - # [batch*heads, query_tokens, key_tokens] - # we do this explicitly because xformers doesn't broadcast the singleton dimension for us. - _, query_tokens, _ = hidden_states.shape - attention_mask = attention_mask.expand(-1, query_tokens, -1) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query = attn.head_to_batch_dim(query).contiguous() - key = attn.head_to_batch_dim(key).contiguous() - value = attn.head_to_batch_dim(value).contiguous() - - hidden_states = xformers.ops.memory_efficient_attention( - query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale - ) - hidden_states = hidden_states.to(query.dtype) - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class AttnProcessorNPU: - r""" - Processor for implementing flash attention using torch_npu. Torch_npu supports only fp16 and bf16 data types. If - fp32 is used, F.scaled_dot_product_attention will be used for computation, but the acceleration effect on NPU is - not significant. - - """ - - def __init__(self): - if not is_torch_npu_available(): - raise ImportError("AttnProcessorNPU requires torch_npu extensions and is supported only on npu devices.") - - 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, - *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) - - residual = hidden_states - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - 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) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - attention_mask = attention_mask.repeat(1, 1, hidden_states.shape[1], 1) - if attention_mask.dtype == torch.bool: - attention_mask = torch.logical_not(attention_mask.bool()) - else: - attention_mask = attention_mask.bool() - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - if query.dtype in (torch.float16, torch.bfloat16): - hidden_states = torch_npu.npu_fusion_attention( - query, - key, - value, - attn.heads, - input_layout="BNSD", - pse=None, - atten_mask=attention_mask, - scale=1.0 / math.sqrt(query.shape[-1]), - pre_tockens=65536, - next_tockens=65536, - keep_prob=1.0, - sync=False, - inner_precise=0, - )[0] - else: - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class AttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_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, - temb: Optional[torch.Tensor] = 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) - - residual = hidden_states - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - 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) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class XLAFlashAttnProcessor2_0: - r""" - 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): - 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." - ) - if is_torch_xla_version("<", "2.3"): - raise ImportError("XLA flash attention requires torch_xla version >= 2.3.") - if is_spmd() and is_torch_xla_version("<", "2.4"): - raise ImportError("SPMD support for XLA flash attention needs torch_xla version >= 2.4.") - self.partition_spec = partition_spec - - 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, - *args, - **kwargs, - ) -> torch.Tensor: - residual = hidden_states - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - 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) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - if all(tensor.shape[2] >= 4096 for tensor in [query, key, value]): - if attention_mask is not None: - attention_mask = attention_mask.view(batch_size, 1, 1, attention_mask.shape[-1]) - # Convert mask to float and replace 0s with -inf and 1s with 0 - attention_mask = ( - attention_mask.float() - .masked_fill(attention_mask == 0, float("-inf")) - .masked_fill(attention_mask == 1, float(0.0)) - ) - - # Apply attention mask to key - key = key + attention_mask - query /= math.sqrt(query.shape[3]) - partition_spec = self.partition_spec if is_spmd() else None - hidden_states = flash_attention(query, key, value, causal=False, partition_spec=partition_spec) - else: - logger.warning( - "Unable to use the flash attention pallas kernel API call due to QKV sequence length < 4096." - ) - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class XLAFluxFlashAttnProcessor2_0: - r""" - 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): - 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." - ) - if is_torch_xla_version("<", "2.3"): - raise ImportError("XLA flash attention requires torch_xla version >= 2.3.") - if is_spmd() and is_torch_xla_version("<", "2.4"): - raise ImportError("SPMD support for XLA flash attention needs torch_xla version >= 2.4.") - self.partition_spec = partition_spec - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.FloatTensor: - batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - - # `sample` projections. - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states` - if encoder_hidden_states is not None: - # `context` projections. - encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states) - encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) - encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) - - encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - - if attn.norm_added_q is not None: - encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj) - if attn.norm_added_k is not None: - encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj) - - # attention - query = torch.cat([encoder_hidden_states_query_proj, query], dim=2) - key = torch.cat([encoder_hidden_states_key_proj, key], dim=2) - value = torch.cat([encoder_hidden_states_value_proj, value], dim=2) - - if image_rotary_emb is not None: - from .embeddings import apply_rotary_emb - - query = apply_rotary_emb(query, image_rotary_emb) - key = apply_rotary_emb(key, image_rotary_emb) - - query /= math.sqrt(head_dim) - hidden_states = flash_attention(query, key, value, causal=False) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - if encoder_hidden_states is not None: - encoder_hidden_states, hidden_states = ( - hidden_states[:, : encoder_hidden_states.shape[1]], - hidden_states[:, encoder_hidden_states.shape[1] :], - ) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - return hidden_states, encoder_hidden_states - else: - return hidden_states - - -class MochiVaeAttnProcessor2_0: - r""" - Attention processor used in Mochi VAE. - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_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, - ) -> torch.Tensor: - residual = hidden_states - is_single_frame = hidden_states.shape[1] == 1 - - 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) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - if is_single_frame: - hidden_states = attn.to_v(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - return hidden_states - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=attn.is_causal - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class StableAudioAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is - used in the Stable Audio model. It applies rotary embedding on query and key vector, and allows MHA, GQA or MQA. - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "StableAudioAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." - ) - - def apply_partial_rotary_emb( - self, - x: torch.Tensor, - freqs_cis: Tuple[torch.Tensor], - ) -> torch.Tensor: - from .embeddings import apply_rotary_emb - - rot_dim = freqs_cis[0].shape[-1] - x_to_rotate, x_unrotated = x[..., :rot_dim], x[..., rot_dim:] - - x_rotated = apply_rotary_emb(x_to_rotate, freqs_cis, use_real=True, use_real_unbind_dim=-2) - - out = torch.cat((x_rotated, x_unrotated), dim=-1) - return out - - 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, - ) -> torch.Tensor: - from .embeddings import apply_rotary_emb - - residual = hidden_states - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - 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) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - head_dim = query.shape[-1] // attn.heads - kv_heads = key.shape[-1] // head_dim - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - key = key.view(batch_size, -1, kv_heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, kv_heads, head_dim).transpose(1, 2) - - if kv_heads != attn.heads: - # if GQA or MQA, repeat the key/value heads to reach the number of query heads. - heads_per_kv_head = attn.heads // kv_heads - key = torch.repeat_interleave(key, heads_per_kv_head, dim=1, output_size=key.shape[1] * heads_per_kv_head) - value = torch.repeat_interleave( - value, heads_per_kv_head, dim=1, output_size=value.shape[1] * heads_per_kv_head - ) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # Apply RoPE if needed - if rotary_emb is not None: - query_dtype = query.dtype - key_dtype = key.dtype - query = query.to(torch.float32) - key = key.to(torch.float32) - - rot_dim = rotary_emb[0].shape[-1] - query_to_rotate, query_unrotated = query[..., :rot_dim], query[..., rot_dim:] - query_rotated = apply_rotary_emb(query_to_rotate, rotary_emb, use_real=True, use_real_unbind_dim=-2) - - query = torch.cat((query_rotated, query_unrotated), dim=-1) - - if not attn.is_cross_attention: - key_to_rotate, key_unrotated = key[..., :rot_dim], key[..., rot_dim:] - key_rotated = apply_rotary_emb(key_to_rotate, rotary_emb, use_real=True, use_real_unbind_dim=-2) - - key = torch.cat((key_rotated, key_unrotated), dim=-1) - - query = query.to(query_dtype) - key = key.to(key_dtype) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class HunyuanAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is - used in the HunyuanDiT model. It applies a s normalization layer and rotary embedding on query and key vector. - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_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, - temb: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - from .embeddings import apply_rotary_emb - - residual = hidden_states - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - 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) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # Apply RoPE if needed - if image_rotary_emb is not None: - query = apply_rotary_emb(query, image_rotary_emb) - if not attn.is_cross_attention: - key = apply_rotary_emb(key, image_rotary_emb) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class FusedHunyuanAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0) with fused - projection layers. This is used in the HunyuanDiT model. It applies a s normalization layer and rotary embedding on - query and key vector. - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "FusedHunyuanAttnProcessor2_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, - temb: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - from .embeddings import apply_rotary_emb - - residual = hidden_states - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - 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) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - if encoder_hidden_states is None: - qkv = attn.to_qkv(hidden_states) - split_size = qkv.shape[-1] // 3 - query, key, value = torch.split(qkv, split_size, dim=-1) - else: - if attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - query = attn.to_q(hidden_states) - - kv = attn.to_kv(encoder_hidden_states) - split_size = kv.shape[-1] // 2 - key, value = torch.split(kv, split_size, dim=-1) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # Apply RoPE if needed - if image_rotary_emb is not None: - query = apply_rotary_emb(query, image_rotary_emb) - if not attn.is_cross_attention: - key = apply_rotary_emb(key, image_rotary_emb) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class PAGHunyuanAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is - used in the HunyuanDiT model. It applies a normalization layer and rotary embedding on query and key vector. This - variant of the processor employs [Pertubed Attention Guidance](https://huggingface.co/papers/2403.17377). - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "PAGHunyuanAttnProcessor2_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, - temb: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - from .embeddings import apply_rotary_emb - - residual = hidden_states - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - # chunk - hidden_states_org, hidden_states_ptb = hidden_states.chunk(2) - - # 1. Original Path - batch_size, sequence_length, _ = ( - hidden_states_org.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) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - if attn.group_norm is not None: - hidden_states_org = attn.group_norm(hidden_states_org.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states_org) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states_org - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # Apply RoPE if needed - if image_rotary_emb is not None: - query = apply_rotary_emb(query, image_rotary_emb) - if not attn.is_cross_attention: - key = apply_rotary_emb(key, image_rotary_emb) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states_org = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states_org = hidden_states_org.to(query.dtype) - - # linear proj - hidden_states_org = attn.to_out[0](hidden_states_org) - # dropout - hidden_states_org = attn.to_out[1](hidden_states_org) - - if input_ndim == 4: - hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) - - # 2. Perturbed Path - if attn.group_norm is not None: - hidden_states_ptb = attn.group_norm(hidden_states_ptb.transpose(1, 2)).transpose(1, 2) - - hidden_states_ptb = attn.to_v(hidden_states_ptb) - hidden_states_ptb = hidden_states_ptb.to(query.dtype) - - # linear proj - hidden_states_ptb = attn.to_out[0](hidden_states_ptb) - # dropout - hidden_states_ptb = attn.to_out[1](hidden_states_ptb) - - if input_ndim == 4: - hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width) - - # cat - hidden_states = torch.cat([hidden_states_org, hidden_states_ptb]) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class PAGCFGHunyuanAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is - used in the HunyuanDiT model. It applies a normalization layer and rotary embedding on query and key vector. This - variant of the processor employs [Pertubed Attention Guidance](https://huggingface.co/papers/2403.17377). - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "PAGCFGHunyuanAttnProcessor2_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, - temb: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - from .embeddings import apply_rotary_emb - - residual = hidden_states - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - # chunk - hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3) - hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org]) - - # 1. Original Path - batch_size, sequence_length, _ = ( - hidden_states_org.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) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - if attn.group_norm is not None: - hidden_states_org = attn.group_norm(hidden_states_org.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states_org) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states_org - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # Apply RoPE if needed - if image_rotary_emb is not None: - query = apply_rotary_emb(query, image_rotary_emb) - if not attn.is_cross_attention: - key = apply_rotary_emb(key, image_rotary_emb) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states_org = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states_org = hidden_states_org.to(query.dtype) - - # linear proj - hidden_states_org = attn.to_out[0](hidden_states_org) - # dropout - hidden_states_org = attn.to_out[1](hidden_states_org) - - if input_ndim == 4: - hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) - - # 2. Perturbed Path - if attn.group_norm is not None: - hidden_states_ptb = attn.group_norm(hidden_states_ptb.transpose(1, 2)).transpose(1, 2) - - hidden_states_ptb = attn.to_v(hidden_states_ptb) - hidden_states_ptb = hidden_states_ptb.to(query.dtype) - - # linear proj - hidden_states_ptb = attn.to_out[0](hidden_states_ptb) - # dropout - hidden_states_ptb = attn.to_out[1](hidden_states_ptb) - - if input_ndim == 4: - hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width) - - # cat - hidden_states = torch.cat([hidden_states_org, hidden_states_ptb]) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class LuminaAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is - used in the LuminaNextDiT model. It applies a s normalization layer and rotary embedding on query and key vector. - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_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: 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, - ) -> torch.Tensor: - from .embeddings import apply_rotary_emb - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - batch_size, sequence_length, _ = hidden_states.shape - - # Get Query-Key-Value Pair - query = attn.to_q(hidden_states) - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query_dim = query.shape[-1] - inner_dim = key.shape[-1] - head_dim = query_dim // attn.heads - dtype = query.dtype - - # Get key-value heads - kv_heads = inner_dim // head_dim - - # Apply Query-Key Norm if needed - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - query = query.view(batch_size, -1, attn.heads, head_dim) - - key = key.view(batch_size, -1, kv_heads, head_dim) - value = value.view(batch_size, -1, kv_heads, head_dim) - - # Apply RoPE if needed - if query_rotary_emb is not None: - query = apply_rotary_emb(query, query_rotary_emb, use_real=False) - if key_rotary_emb is not None: - key = apply_rotary_emb(key, key_rotary_emb, use_real=False) - - query, key = query.to(dtype), key.to(dtype) - - # Apply proportional attention if true - if key_rotary_emb is None: - softmax_scale = None - else: - if base_sequence_length is not None: - softmax_scale = math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale - else: - softmax_scale = attn.scale - - # perform Grouped-qurey Attention (GQA) - n_rep = attn.heads // kv_heads - if n_rep >= 1: - key = key.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3) - value = value.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3) - - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.bool().view(batch_size, 1, 1, -1) - attention_mask = attention_mask.expand(-1, attn.heads, sequence_length, -1) - - query = query.transpose(1, 2) - key = key.transpose(1, 2) - value = value.transpose(1, 2) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, scale=softmax_scale - ) - hidden_states = hidden_states.transpose(1, 2).to(dtype) - - return hidden_states - - -class FusedAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). It uses - fused projection layers. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. - For cross-attention modules, key and value projection matrices are fused. - - - - This API is currently 🧪 experimental in nature and can change in future. - - - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "FusedAttnProcessor2_0 requires at least 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, - temb: Optional[torch.Tensor] = 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) - - residual = hidden_states - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - 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) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - if encoder_hidden_states is None: - qkv = attn.to_qkv(hidden_states) - split_size = qkv.shape[-1] // 3 - query, key, value = torch.split(qkv, split_size, dim=-1) - else: - if attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - query = attn.to_q(hidden_states) - - kv = attn.to_kv(encoder_hidden_states) - split_size = kv.shape[-1] // 2 - key, value = torch.split(kv, split_size, dim=-1) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class CustomDiffusionXFormersAttnProcessor(nn.Module): - r""" - Processor for implementing memory efficient attention using xFormers for the Custom Diffusion method. - - Args: - train_kv (`bool`, defaults to `True`): - Whether to newly train the key and value matrices corresponding to the text features. - train_q_out (`bool`, defaults to `True`): - Whether to newly train query matrices corresponding to the latent image features. - hidden_size (`int`, *optional*, defaults to `None`): - The hidden size of the attention layer. - cross_attention_dim (`int`, *optional*, defaults to `None`): - The number of channels in the `encoder_hidden_states`. - out_bias (`bool`, defaults to `True`): - Whether to include the bias parameter in `train_q_out`. - dropout (`float`, *optional*, defaults to 0.0): - The dropout probability to use. - attention_op (`Callable`, *optional*, defaults to `None`): - The base - [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to use - as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best operator. - """ - - def __init__( - self, - train_kv: bool = True, - train_q_out: bool = False, - hidden_size: Optional[int] = None, - cross_attention_dim: Optional[int] = None, - out_bias: bool = True, - dropout: float = 0.0, - attention_op: Optional[Callable] = None, - ): - super().__init__() - self.train_kv = train_kv - self.train_q_out = train_q_out - - self.hidden_size = hidden_size - self.cross_attention_dim = cross_attention_dim - self.attention_op = attention_op - - # `_custom_diffusion` id for easy serialization and loading. - if self.train_kv: - self.to_k_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) - self.to_v_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) - if self.train_q_out: - self.to_q_custom_diffusion = nn.Linear(hidden_size, hidden_size, bias=False) - self.to_out_custom_diffusion = nn.ModuleList([]) - self.to_out_custom_diffusion.append(nn.Linear(hidden_size, hidden_size, bias=out_bias)) - self.to_out_custom_diffusion.append(nn.Dropout(dropout)) - - def __call__( - self, - attn: Attention, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - - if self.train_q_out: - query = self.to_q_custom_diffusion(hidden_states).to(attn.to_q.weight.dtype) - else: - query = attn.to_q(hidden_states.to(attn.to_q.weight.dtype)) - - if encoder_hidden_states is None: - crossattn = False - encoder_hidden_states = hidden_states - else: - crossattn = True - if attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - if self.train_kv: - key = self.to_k_custom_diffusion(encoder_hidden_states.to(self.to_k_custom_diffusion.weight.dtype)) - value = self.to_v_custom_diffusion(encoder_hidden_states.to(self.to_v_custom_diffusion.weight.dtype)) - key = key.to(attn.to_q.weight.dtype) - value = value.to(attn.to_q.weight.dtype) - else: - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - if crossattn: - detach = torch.ones_like(key) - detach[:, :1, :] = detach[:, :1, :] * 0.0 - key = detach * key + (1 - detach) * key.detach() - value = detach * value + (1 - detach) * value.detach() - - query = attn.head_to_batch_dim(query).contiguous() - key = attn.head_to_batch_dim(key).contiguous() - value = attn.head_to_batch_dim(value).contiguous() - - hidden_states = xformers.ops.memory_efficient_attention( - query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale - ) - hidden_states = hidden_states.to(query.dtype) - hidden_states = attn.batch_to_head_dim(hidden_states) - - if self.train_q_out: - # linear proj - hidden_states = self.to_out_custom_diffusion[0](hidden_states) - # dropout - hidden_states = self.to_out_custom_diffusion[1](hidden_states) - else: - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - return hidden_states - - -class CustomDiffusionAttnProcessor2_0(nn.Module): - r""" - Processor for implementing attention for the Custom Diffusion method using PyTorch 2.0’s memory-efficient scaled - dot-product attention. - - Args: - train_kv (`bool`, defaults to `True`): - Whether to newly train the key and value matrices corresponding to the text features. - train_q_out (`bool`, defaults to `True`): - Whether to newly train query matrices corresponding to the latent image features. - hidden_size (`int`, *optional*, defaults to `None`): - The hidden size of the attention layer. - cross_attention_dim (`int`, *optional*, defaults to `None`): - The number of channels in the `encoder_hidden_states`. - out_bias (`bool`, defaults to `True`): - Whether to include the bias parameter in `train_q_out`. - dropout (`float`, *optional*, defaults to 0.0): - The dropout probability to use. - """ - - def __init__( - self, - train_kv: bool = True, - train_q_out: bool = True, - hidden_size: Optional[int] = None, - cross_attention_dim: Optional[int] = None, - out_bias: bool = True, - dropout: float = 0.0, - ): - super().__init__() - self.train_kv = train_kv - self.train_q_out = train_q_out - - self.hidden_size = hidden_size - self.cross_attention_dim = cross_attention_dim - - # `_custom_diffusion` id for easy serialization and loading. - if self.train_kv: - self.to_k_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) - self.to_v_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) - if self.train_q_out: - self.to_q_custom_diffusion = nn.Linear(hidden_size, hidden_size, bias=False) - self.to_out_custom_diffusion = nn.ModuleList([]) - self.to_out_custom_diffusion.append(nn.Linear(hidden_size, hidden_size, bias=out_bias)) - self.to_out_custom_diffusion.append(nn.Dropout(dropout)) - - def __call__( - self, - attn: Attention, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - batch_size, sequence_length, _ = hidden_states.shape - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - if self.train_q_out: - query = self.to_q_custom_diffusion(hidden_states) - else: - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - crossattn = False - encoder_hidden_states = hidden_states - else: - crossattn = True - if attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - if self.train_kv: - key = self.to_k_custom_diffusion(encoder_hidden_states.to(self.to_k_custom_diffusion.weight.dtype)) - value = self.to_v_custom_diffusion(encoder_hidden_states.to(self.to_v_custom_diffusion.weight.dtype)) - key = key.to(attn.to_q.weight.dtype) - value = value.to(attn.to_q.weight.dtype) - - else: - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - if crossattn: - detach = torch.ones_like(key) - detach[:, :1, :] = detach[:, :1, :] * 0.0 - key = detach * key + (1 - detach) * key.detach() - value = detach * value + (1 - detach) * value.detach() - - inner_dim = hidden_states.shape[-1] - - head_dim = inner_dim // attn.heads - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - if self.train_q_out: - # linear proj - hidden_states = self.to_out_custom_diffusion[0](hidden_states) - # dropout - hidden_states = self.to_out_custom_diffusion[1](hidden_states) - else: - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - return hidden_states - - -class SlicedAttnProcessor: - r""" - Processor for implementing sliced attention. - - Args: - slice_size (`int`, *optional*): - The number of steps to compute attention. Uses as many slices as `attention_head_dim // slice_size`, and - `attention_head_dim` must be a multiple of the `slice_size`. - """ - - def __init__(self, slice_size: int): - self.slice_size = slice_size - - def __call__( - self, - attn: Attention, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - residual = hidden_states - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - dim = query.shape[-1] - query = attn.head_to_batch_dim(query) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - - batch_size_attention, query_tokens, _ = query.shape - hidden_states = torch.zeros( - (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype - ) - - for i in range((batch_size_attention - 1) // self.slice_size + 1): - start_idx = i * self.slice_size - end_idx = (i + 1) * self.slice_size - - query_slice = query[start_idx:end_idx] - key_slice = key[start_idx:end_idx] - attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None - - attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice) - - attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx]) - - hidden_states[start_idx:end_idx] = attn_slice - - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class SlicedAttnAddedKVProcessor: - r""" - Processor for implementing sliced attention with extra learnable key and value matrices for the text encoder. - - Args: - slice_size (`int`, *optional*): - The number of steps to compute attention. Uses as many slices as `attention_head_dim // slice_size`, and - `attention_head_dim` must be a multiple of the `slice_size`. - """ - - def __init__(self, slice_size): - self.slice_size = slice_size - - 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, - ) -> torch.Tensor: - residual = hidden_states - - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2) - - batch_size, sequence_length, _ = hidden_states.shape - - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - dim = query.shape[-1] - query = attn.head_to_batch_dim(query) - - encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) - encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) - - encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj) - encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj) - - if not attn.only_cross_attention: - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - key = torch.cat([encoder_hidden_states_key_proj, key], dim=1) - value = torch.cat([encoder_hidden_states_value_proj, value], dim=1) - else: - key = encoder_hidden_states_key_proj - value = encoder_hidden_states_value_proj - - batch_size_attention, query_tokens, _ = query.shape - hidden_states = torch.zeros( - (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype - ) - - for i in range((batch_size_attention - 1) // self.slice_size + 1): - start_idx = i * self.slice_size - end_idx = (i + 1) * self.slice_size - - query_slice = query[start_idx:end_idx] - key_slice = key[start_idx:end_idx] - attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None - - attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice) - - attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx]) - - hidden_states[start_idx:end_idx] = attn_slice - - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape) - hidden_states = hidden_states + residual - - return hidden_states - - -class SpatialNorm(nn.Module): - """ - Spatially conditioned normalization as defined in https://huggingface.co/papers/2209.09002. - - Args: - f_channels (`int`): - The number of channels for input to group normalization layer, and output of the spatial norm layer. - zq_channels (`int`): - The number of channels for the quantized vector as described in the paper. - """ - - def __init__( - self, - f_channels: int, - zq_channels: int, - ): - super().__init__() - self.norm_layer = nn.GroupNorm(num_channels=f_channels, num_groups=32, eps=1e-6, affine=True) - self.conv_y = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) - self.conv_b = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) - - def forward(self, f: torch.Tensor, zq: torch.Tensor) -> torch.Tensor: - f_size = f.shape[-2:] - zq = F.interpolate(zq, size=f_size, mode="nearest") - norm_f = self.norm_layer(f) - new_f = norm_f * self.conv_y(zq) + self.conv_b(zq) - return new_f - - -class IPAdapterAttnProcessor(nn.Module): - r""" - Attention processor for Multiple IP-Adapters. - - Args: - hidden_size (`int`): - 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,)`): - The context length of the image features. - scale (`float` or List[`float`], defaults to 1.0): - the weight scale of image prompt. - """ - - def __init__(self, hidden_size, cross_attention_dim=None, num_tokens=(4,), scale=1.0): - super().__init__() - - self.hidden_size = hidden_size - self.cross_attention_dim = cross_attention_dim - - if not isinstance(num_tokens, (tuple, list)): - num_tokens = [num_tokens] - self.num_tokens = num_tokens - - if not isinstance(scale, list): - scale = [scale] * len(num_tokens) - if len(scale) != len(num_tokens): - raise ValueError("`scale` should be a list of integers with the same length as `num_tokens`.") - self.scale = scale - - self.to_k_ip = nn.ModuleList( - [nn.Linear(cross_attention_dim, hidden_size, bias=False) for _ in range(len(num_tokens))] - ) - self.to_v_ip = nn.ModuleList( - [nn.Linear(cross_attention_dim, hidden_size, bias=False) for _ in range(len(num_tokens))] - ) - - 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, - scale: float = 1.0, - ip_adapter_masks: Optional[torch.Tensor] = None, - ): - residual = hidden_states - - # separate ip_hidden_states from encoder_hidden_states - if encoder_hidden_states is not None: - if isinstance(encoder_hidden_states, tuple): - encoder_hidden_states, ip_hidden_states = encoder_hidden_states - else: - deprecation_message = ( - "You have passed a tensor as `encoder_hidden_states`. This is deprecated and will be removed in a future release." - " Please make sure to update your script to pass `encoder_hidden_states` as a tuple to suppress this warning." - ) - deprecate("encoder_hidden_states not a tuple", "1.0.0", deprecation_message, standard_warn=False) - end_pos = encoder_hidden_states.shape[1] - self.num_tokens[0] - encoder_hidden_states, ip_hidden_states = ( - encoder_hidden_states[:, :end_pos, :], - [encoder_hidden_states[:, end_pos:, :]], - ) - - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query = attn.head_to_batch_dim(query) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - - attention_probs = attn.get_attention_scores(query, key, attention_mask) - hidden_states = torch.bmm(attention_probs, value) - hidden_states = attn.batch_to_head_dim(hidden_states) - - if ip_adapter_masks is not None: - 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)): - raise ValueError( - f"Length of ip_adapter_masks array ({len(ip_adapter_masks)}) must match " - f"length of self.scale array ({len(self.scale)}) and number of ip_hidden_states " - f"({len(ip_hidden_states)})" - ) - else: - for index, (mask, scale, ip_state) in enumerate(zip(ip_adapter_masks, self.scale, ip_hidden_states)): - if mask is None: - continue - if not isinstance(mask, torch.Tensor) or mask.ndim != 4: - raise ValueError( - "Each element of the ip_adapter_masks array should be a tensor with shape " - "[1, num_images_for_ip_adapter, height, width]." - " Please use `IPAdapterMaskProcessor` to preprocess your mask" - ) - if mask.shape[1] != ip_state.shape[1]: - raise ValueError( - f"Number of masks ({mask.shape[1]}) does not match " - f"number of ip images ({ip_state.shape[1]}) at index {index}" - ) - if isinstance(scale, list) and not len(scale) == mask.shape[1]: - raise ValueError( - f"Number of masks ({mask.shape[1]}) does not match " - f"number of scales ({len(scale)}) at index {index}" - ) - else: - ip_adapter_masks = [None] * len(self.scale) - - # for ip-adapter - for current_ip_hidden_states, scale, to_k_ip, to_v_ip, mask in zip( - ip_hidden_states, self.scale, self.to_k_ip, self.to_v_ip, ip_adapter_masks - ): - skip = False - if isinstance(scale, list): - if all(s == 0 for s in scale): - skip = True - elif scale == 0: - skip = True - if not skip: - if mask is not None: - if not isinstance(scale, list): - scale = [scale] * mask.shape[1] - - current_num_images = mask.shape[1] - for i in range(current_num_images): - ip_key = to_k_ip(current_ip_hidden_states[:, i, :, :]) - ip_value = to_v_ip(current_ip_hidden_states[:, i, :, :]) - - ip_key = attn.head_to_batch_dim(ip_key) - ip_value = attn.head_to_batch_dim(ip_value) - - ip_attention_probs = attn.get_attention_scores(query, ip_key, None) - _current_ip_hidden_states = torch.bmm(ip_attention_probs, ip_value) - _current_ip_hidden_states = attn.batch_to_head_dim(_current_ip_hidden_states) - - mask_downsample = IPAdapterMaskProcessor.downsample( - mask[:, i, :, :], - batch_size, - _current_ip_hidden_states.shape[1], - _current_ip_hidden_states.shape[2], - ) - - mask_downsample = mask_downsample.to(dtype=query.dtype, device=query.device) - - hidden_states = hidden_states + scale[i] * (_current_ip_hidden_states * mask_downsample) - else: - ip_key = to_k_ip(current_ip_hidden_states) - ip_value = to_v_ip(current_ip_hidden_states) - - ip_key = attn.head_to_batch_dim(ip_key) - ip_value = attn.head_to_batch_dim(ip_value) - - ip_attention_probs = attn.get_attention_scores(query, ip_key, None) - current_ip_hidden_states = torch.bmm(ip_attention_probs, ip_value) - current_ip_hidden_states = attn.batch_to_head_dim(current_ip_hidden_states) - - hidden_states = hidden_states + scale * current_ip_hidden_states - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class IPAdapterAttnProcessor2_0(torch.nn.Module): - r""" - Attention processor for IP-Adapter for PyTorch 2.0. - - Args: - hidden_size (`int`): - 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,)`): - The context length of the image features. - scale (`float` or `List[float]`, defaults to 1.0): - the weight scale of image prompt. - """ - - def __init__(self, hidden_size, cross_attention_dim=None, num_tokens=(4,), scale=1.0): - super().__init__() - - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - f"{self.__class__.__name__} requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." - ) - - self.hidden_size = hidden_size - self.cross_attention_dim = cross_attention_dim - - if not isinstance(num_tokens, (tuple, list)): - num_tokens = [num_tokens] - self.num_tokens = num_tokens - - if not isinstance(scale, list): - scale = [scale] * len(num_tokens) - if len(scale) != len(num_tokens): - raise ValueError("`scale` should be a list of integers with the same length as `num_tokens`.") - self.scale = scale - - self.to_k_ip = nn.ModuleList( - [nn.Linear(cross_attention_dim, hidden_size, bias=False) for _ in range(len(num_tokens))] - ) - self.to_v_ip = nn.ModuleList( - [nn.Linear(cross_attention_dim, hidden_size, bias=False) for _ in range(len(num_tokens))] - ) - - 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, - scale: float = 1.0, - ip_adapter_masks: Optional[torch.Tensor] = None, - ): - residual = hidden_states - - # separate ip_hidden_states from encoder_hidden_states - if encoder_hidden_states is not None: - if isinstance(encoder_hidden_states, tuple): - encoder_hidden_states, ip_hidden_states = encoder_hidden_states - else: - deprecation_message = ( - "You have passed a tensor as `encoder_hidden_states`. This is deprecated and will be removed in a future release." - " Please make sure to update your script to pass `encoder_hidden_states` as a tuple to suppress this warning." - ) - deprecate("encoder_hidden_states not a tuple", "1.0.0", deprecation_message, standard_warn=False) - end_pos = encoder_hidden_states.shape[1] - self.num_tokens[0] - encoder_hidden_states, ip_hidden_states = ( - encoder_hidden_states[:, :end_pos, :], - [encoder_hidden_states[:, end_pos:, :]], - ) - - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - 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) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - if ip_adapter_masks is not None: - 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)): - raise ValueError( - f"Length of ip_adapter_masks array ({len(ip_adapter_masks)}) must match " - f"length of self.scale array ({len(self.scale)}) and number of ip_hidden_states " - f"({len(ip_hidden_states)})" - ) - else: - for index, (mask, scale, ip_state) in enumerate(zip(ip_adapter_masks, self.scale, ip_hidden_states)): - if mask is None: - continue - if not isinstance(mask, torch.Tensor) or mask.ndim != 4: - raise ValueError( - "Each element of the ip_adapter_masks array should be a tensor with shape " - "[1, num_images_for_ip_adapter, height, width]." - " Please use `IPAdapterMaskProcessor` to preprocess your mask" - ) - if mask.shape[1] != ip_state.shape[1]: - raise ValueError( - f"Number of masks ({mask.shape[1]}) does not match " - f"number of ip images ({ip_state.shape[1]}) at index {index}" - ) - if isinstance(scale, list) and not len(scale) == mask.shape[1]: - raise ValueError( - f"Number of masks ({mask.shape[1]}) does not match " - f"number of scales ({len(scale)}) at index {index}" - ) - else: - ip_adapter_masks = [None] * len(self.scale) - - # for ip-adapter - for current_ip_hidden_states, scale, to_k_ip, to_v_ip, mask in zip( - ip_hidden_states, self.scale, self.to_k_ip, self.to_v_ip, ip_adapter_masks - ): - skip = False - if isinstance(scale, list): - if all(s == 0 for s in scale): - skip = True - elif scale == 0: - skip = True - if not skip: - if mask is not None: - if not isinstance(scale, list): - scale = [scale] * mask.shape[1] - - current_num_images = mask.shape[1] - for i in range(current_num_images): - ip_key = to_k_ip(current_ip_hidden_states[:, i, :, :]) - ip_value = to_v_ip(current_ip_hidden_states[:, i, :, :]) - - ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - _current_ip_hidden_states = F.scaled_dot_product_attention( - query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False - ) - - _current_ip_hidden_states = _current_ip_hidden_states.transpose(1, 2).reshape( - batch_size, -1, attn.heads * head_dim - ) - _current_ip_hidden_states = _current_ip_hidden_states.to(query.dtype) - - mask_downsample = IPAdapterMaskProcessor.downsample( - mask[:, i, :, :], - batch_size, - _current_ip_hidden_states.shape[1], - _current_ip_hidden_states.shape[2], - ) - - mask_downsample = mask_downsample.to(dtype=query.dtype, device=query.device) - hidden_states = hidden_states + scale[i] * (_current_ip_hidden_states * mask_downsample) - else: - ip_key = to_k_ip(current_ip_hidden_states) - ip_value = to_v_ip(current_ip_hidden_states) - - ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - current_ip_hidden_states = F.scaled_dot_product_attention( - query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False - ) - - current_ip_hidden_states = current_ip_hidden_states.transpose(1, 2).reshape( - batch_size, -1, attn.heads * head_dim - ) - current_ip_hidden_states = current_ip_hidden_states.to(query.dtype) - - hidden_states = hidden_states + scale * current_ip_hidden_states - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class IPAdapterXFormersAttnProcessor(torch.nn.Module): - r""" - Attention processor for IP-Adapter using xFormers. - - Args: - hidden_size (`int`): - 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,)`): - The context length of the image features. - scale (`float` or `List[float]`, defaults to 1.0): - the weight scale of image prompt. - attention_op (`Callable`, *optional*, defaults to `None`): - The base - [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to - use as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best - operator. - """ - - def __init__( - self, - hidden_size, - cross_attention_dim=None, - num_tokens=(4,), - scale=1.0, - attention_op: Optional[Callable] = None, - ): - super().__init__() - - self.hidden_size = hidden_size - self.cross_attention_dim = cross_attention_dim - self.attention_op = attention_op - - if not isinstance(num_tokens, (tuple, list)): - num_tokens = [num_tokens] - self.num_tokens = num_tokens - - if not isinstance(scale, list): - scale = [scale] * len(num_tokens) - if len(scale) != len(num_tokens): - raise ValueError("`scale` should be a list of integers with the same length as `num_tokens`.") - self.scale = scale - - self.to_k_ip = nn.ModuleList( - [nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) for _ in range(len(num_tokens))] - ) - self.to_v_ip = nn.ModuleList( - [nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) for _ in range(len(num_tokens))] - ) - - 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, - scale: float = 1.0, - ip_adapter_masks: Optional[torch.FloatTensor] = None, - ): - residual = hidden_states - - # separate ip_hidden_states from encoder_hidden_states - if encoder_hidden_states is not None: - if isinstance(encoder_hidden_states, tuple): - encoder_hidden_states, ip_hidden_states = encoder_hidden_states - else: - deprecation_message = ( - "You have passed a tensor as `encoder_hidden_states`. This is deprecated and will be removed in a future release." - " Please make sure to update your script to pass `encoder_hidden_states` as a tuple to suppress this warning." - ) - deprecate("encoder_hidden_states not a tuple", "1.0.0", deprecation_message, standard_warn=False) - end_pos = encoder_hidden_states.shape[1] - self.num_tokens[0] - encoder_hidden_states, ip_hidden_states = ( - encoder_hidden_states[:, :end_pos, :], - [encoder_hidden_states[:, end_pos:, :]], - ) - - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - 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) - # expand our mask's singleton query_tokens dimension: - # [batch*heads, 1, key_tokens] -> - # [batch*heads, query_tokens, key_tokens] - # so that it can be added as a bias onto the attention scores that xformers computes: - # [batch*heads, query_tokens, key_tokens] - # we do this explicitly because xformers doesn't broadcast the singleton dimension for us. - _, query_tokens, _ = hidden_states.shape - attention_mask = attention_mask.expand(-1, query_tokens, -1) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query = attn.head_to_batch_dim(query).contiguous() - key = attn.head_to_batch_dim(key).contiguous() - value = attn.head_to_batch_dim(value).contiguous() - - hidden_states = xformers.ops.memory_efficient_attention( - query, key, value, attn_bias=attention_mask, op=self.attention_op - ) - hidden_states = hidden_states.to(query.dtype) - hidden_states = attn.batch_to_head_dim(hidden_states) - - if ip_hidden_states: - if ip_adapter_masks is not None: - 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)): - raise ValueError( - f"Length of ip_adapter_masks array ({len(ip_adapter_masks)}) must match " - f"length of self.scale array ({len(self.scale)}) and number of ip_hidden_states " - f"({len(ip_hidden_states)})" - ) - else: - for index, (mask, scale, ip_state) in enumerate( - zip(ip_adapter_masks, self.scale, ip_hidden_states) - ): - if mask is None: - continue - if not isinstance(mask, torch.Tensor) or mask.ndim != 4: - raise ValueError( - "Each element of the ip_adapter_masks array should be a tensor with shape " - "[1, num_images_for_ip_adapter, height, width]." - " Please use `IPAdapterMaskProcessor` to preprocess your mask" - ) - if mask.shape[1] != ip_state.shape[1]: - raise ValueError( - f"Number of masks ({mask.shape[1]}) does not match " - f"number of ip images ({ip_state.shape[1]}) at index {index}" - ) - if isinstance(scale, list) and not len(scale) == mask.shape[1]: - raise ValueError( - f"Number of masks ({mask.shape[1]}) does not match " - f"number of scales ({len(scale)}) at index {index}" - ) - else: - ip_adapter_masks = [None] * len(self.scale) - - # for ip-adapter - for current_ip_hidden_states, scale, to_k_ip, to_v_ip, mask in zip( - ip_hidden_states, self.scale, self.to_k_ip, self.to_v_ip, ip_adapter_masks - ): - skip = False - if isinstance(scale, list): - if all(s == 0 for s in scale): - skip = True - elif scale == 0: - skip = True - if not skip: - if mask is not None: - mask = mask.to(torch.float16) - if not isinstance(scale, list): - scale = [scale] * mask.shape[1] - - current_num_images = mask.shape[1] - for i in range(current_num_images): - ip_key = to_k_ip(current_ip_hidden_states[:, i, :, :]) - ip_value = to_v_ip(current_ip_hidden_states[:, i, :, :]) - - ip_key = attn.head_to_batch_dim(ip_key).contiguous() - ip_value = attn.head_to_batch_dim(ip_value).contiguous() - - _current_ip_hidden_states = xformers.ops.memory_efficient_attention( - query, ip_key, ip_value, op=self.attention_op - ) - _current_ip_hidden_states = _current_ip_hidden_states.to(query.dtype) - _current_ip_hidden_states = attn.batch_to_head_dim(_current_ip_hidden_states) - - mask_downsample = IPAdapterMaskProcessor.downsample( - mask[:, i, :, :], - batch_size, - _current_ip_hidden_states.shape[1], - _current_ip_hidden_states.shape[2], - ) - - mask_downsample = mask_downsample.to(dtype=query.dtype, device=query.device) - hidden_states = hidden_states + scale[i] * (_current_ip_hidden_states * mask_downsample) - else: - ip_key = to_k_ip(current_ip_hidden_states) - ip_value = to_v_ip(current_ip_hidden_states) - - ip_key = attn.head_to_batch_dim(ip_key).contiguous() - ip_value = attn.head_to_batch_dim(ip_value).contiguous() - - current_ip_hidden_states = xformers.ops.memory_efficient_attention( - query, ip_key, ip_value, op=self.attention_op - ) - current_ip_hidden_states = current_ip_hidden_states.to(query.dtype) - current_ip_hidden_states = attn.batch_to_head_dim(current_ip_hidden_states) - - hidden_states = hidden_states + scale * current_ip_hidden_states - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class SD3IPAdapterJointAttnProcessor2_0(torch.nn.Module): - """ - Attention processor for IP-Adapter used typically in processing the SD3-like self-attention projections, with - additional image-based information and timestep embeddings. - - Args: - hidden_size (`int`): - The number of hidden channels. - ip_hidden_states_dim (`int`): - The image feature dimension. - head_dim (`int`): - The number of head channels. - timesteps_emb_dim (`int`, defaults to 1280): - The number of input channels for timestep embedding. - scale (`float`, defaults to 0.5): - IP-Adapter scale. - """ - - def __init__( - self, - hidden_size: int, - ip_hidden_states_dim: int, - head_dim: int, - timesteps_emb_dim: int = 1280, - scale: float = 0.5, - ): - super().__init__() - - # To prevent circular import - from .normalization import AdaLayerNorm, RMSNorm - - self.norm_ip = AdaLayerNorm(timesteps_emb_dim, output_dim=ip_hidden_states_dim * 2, norm_eps=1e-6, chunk_dim=1) - self.to_k_ip = nn.Linear(ip_hidden_states_dim, hidden_size, bias=False) - self.to_v_ip = nn.Linear(ip_hidden_states_dim, hidden_size, bias=False) - self.norm_q = RMSNorm(head_dim, 1e-6) - self.norm_k = RMSNorm(head_dim, 1e-6) - self.norm_ip_k = RMSNorm(head_dim, 1e-6) - self.scale = scale - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, - ip_hidden_states: torch.FloatTensor = None, - temb: torch.FloatTensor = None, - ) -> torch.FloatTensor: - """ - Perform the attention computation, integrating image features (if provided) and timestep embeddings. - - If `ip_hidden_states` is `None`, this is equivalent to using JointAttnProcessor2_0. - - Args: - attn (`Attention`): - Attention instance. - hidden_states (`torch.FloatTensor`): - Input `hidden_states`. - encoder_hidden_states (`torch.FloatTensor`, *optional*): - The encoder hidden states. - attention_mask (`torch.FloatTensor`, *optional*): - Attention mask. - ip_hidden_states (`torch.FloatTensor`, *optional*): - Image embeddings. - temb (`torch.FloatTensor`, *optional*): - Timestep embeddings. - - Returns: - `torch.FloatTensor`: Output hidden states. - """ - residual = hidden_states - - batch_size = hidden_states.shape[0] - - # `sample` projections. - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - img_query = query - img_key = key - img_value = value - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # `context` projections. - if encoder_hidden_states is not None: - encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states) - encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states) - encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states) - - encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( - batch_size, -1, attn.heads, head_dim - ).transpose(1, 2) - - if attn.norm_added_q is not None: - encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj) - if attn.norm_added_k is not None: - encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj) - - query = torch.cat([query, encoder_hidden_states_query_proj], dim=2) - key = torch.cat([key, encoder_hidden_states_key_proj], dim=2) - value = torch.cat([value, encoder_hidden_states_value_proj], dim=2) - - hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - if encoder_hidden_states is not None: - # Split the attention outputs. - hidden_states, encoder_hidden_states = ( - hidden_states[:, : residual.shape[1]], - hidden_states[:, residual.shape[1] :], - ) - if not attn.context_pre_only: - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - # IP Adapter - if self.scale != 0 and ip_hidden_states is not None: - # Norm image features - norm_ip_hidden_states = self.norm_ip(ip_hidden_states, temb=temb) - - # To k and v - ip_key = self.to_k_ip(norm_ip_hidden_states) - ip_value = self.to_v_ip(norm_ip_hidden_states) - - # Reshape - ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # Norm - query = self.norm_q(img_query) - img_key = self.norm_k(img_key) - ip_key = self.norm_ip_k(ip_key) - - # cat img - key = torch.cat([img_key, ip_key], dim=2) - value = torch.cat([img_value, ip_value], dim=2) - - ip_hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) - ip_hidden_states = ip_hidden_states.transpose(1, 2).view(batch_size, -1, attn.heads * head_dim) - ip_hidden_states = ip_hidden_states.to(query.dtype) - - hidden_states = hidden_states + ip_hidden_states * self.scale - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if encoder_hidden_states is not None: - return hidden_states, encoder_hidden_states - else: - return hidden_states - - -class PAGIdentitySelfAttnProcessor2_0: - r""" - Processor for implementing PAG using scaled dot-product attention (enabled by default if you're using PyTorch 2.0). - PAG reference: https://huggingface.co/papers/2403.17377 - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "PAGIdentitySelfAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." - ) - - 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, - ) -> torch.Tensor: - residual = hidden_states - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - # chunk - hidden_states_org, hidden_states_ptb = hidden_states.chunk(2) - - # original path - batch_size, sequence_length, _ = hidden_states_org.shape - - if attention_mask is not None: - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - if attn.group_norm is not None: - hidden_states_org = attn.group_norm(hidden_states_org.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states_org) - key = attn.to_k(hidden_states_org) - value = attn.to_v(hidden_states_org) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states_org = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states_org = hidden_states_org.to(query.dtype) - - # linear proj - hidden_states_org = attn.to_out[0](hidden_states_org) - # dropout - hidden_states_org = attn.to_out[1](hidden_states_org) - - if input_ndim == 4: - hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) - - # perturbed path (identity attention) - batch_size, sequence_length, _ = hidden_states_ptb.shape - - if attn.group_norm is not None: - hidden_states_ptb = attn.group_norm(hidden_states_ptb.transpose(1, 2)).transpose(1, 2) - - hidden_states_ptb = attn.to_v(hidden_states_ptb) - hidden_states_ptb = hidden_states_ptb.to(query.dtype) - - # linear proj - hidden_states_ptb = attn.to_out[0](hidden_states_ptb) - # dropout - hidden_states_ptb = attn.to_out[1](hidden_states_ptb) - - if input_ndim == 4: - hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width) - - # cat - hidden_states = torch.cat([hidden_states_org, hidden_states_ptb]) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class PAGCFGIdentitySelfAttnProcessor2_0: - r""" - Processor for implementing PAG using scaled dot-product attention (enabled by default if you're using PyTorch 2.0). - PAG reference: https://huggingface.co/papers/2403.17377 - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "PAGCFGIdentitySelfAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." - ) - - 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, - ) -> torch.Tensor: - residual = hidden_states - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - # chunk - hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3) - hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org]) - - # original path - batch_size, sequence_length, _ = hidden_states_org.shape - - if attention_mask is not None: - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - if attn.group_norm is not None: - hidden_states_org = attn.group_norm(hidden_states_org.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states_org) - key = attn.to_k(hidden_states_org) - value = attn.to_v(hidden_states_org) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states_org = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states_org = hidden_states_org.to(query.dtype) - - # linear proj - hidden_states_org = attn.to_out[0](hidden_states_org) - # dropout - hidden_states_org = attn.to_out[1](hidden_states_org) - - if input_ndim == 4: - hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width) - - # perturbed path (identity attention) - batch_size, sequence_length, _ = hidden_states_ptb.shape - - if attn.group_norm is not None: - hidden_states_ptb = attn.group_norm(hidden_states_ptb.transpose(1, 2)).transpose(1, 2) - - value = attn.to_v(hidden_states_ptb) - hidden_states_ptb = value - hidden_states_ptb = hidden_states_ptb.to(query.dtype) - - # linear proj - hidden_states_ptb = attn.to_out[0](hidden_states_ptb) - # dropout - hidden_states_ptb = attn.to_out[1](hidden_states_ptb) - - if input_ndim == 4: - hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width) - - # cat - hidden_states = torch.cat([hidden_states_org, hidden_states_ptb]) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class SanaMultiscaleAttnProcessor2_0: - r""" - Processor for implementing multiscale quadratic attention. - """ - - def __call__(self, attn: SanaMultiscaleLinearAttention, hidden_states: torch.Tensor) -> torch.Tensor: - height, width = hidden_states.shape[-2:] - if height * width > attn.attention_head_dim: - use_linear_attention = True - else: - use_linear_attention = False - - residual = hidden_states - - batch_size, _, height, width = list(hidden_states.size()) - original_dtype = hidden_states.dtype - - hidden_states = hidden_states.movedim(1, -1) - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - hidden_states = torch.cat([query, key, value], dim=3) - hidden_states = hidden_states.movedim(-1, 1) - - multi_scale_qkv = [hidden_states] - for block in attn.to_qkv_multiscale: - multi_scale_qkv.append(block(hidden_states)) - - hidden_states = torch.cat(multi_scale_qkv, dim=1) - - if use_linear_attention: - # for linear attention upcast hidden_states to float32 - hidden_states = hidden_states.to(dtype=torch.float32) - - hidden_states = hidden_states.reshape(batch_size, -1, 3 * attn.attention_head_dim, height * width) - - query, key, value = hidden_states.chunk(3, dim=2) - query = attn.nonlinearity(query) - key = attn.nonlinearity(key) - - if use_linear_attention: - hidden_states = attn.apply_linear_attention(query, key, value) - hidden_states = hidden_states.to(dtype=original_dtype) - else: - hidden_states = attn.apply_quadratic_attention(query, key, value) - - hidden_states = torch.reshape(hidden_states, (batch_size, -1, height, width)) - hidden_states = attn.to_out(hidden_states.movedim(1, -1)).movedim(-1, 1) - - if attn.norm_type == "rms_norm": - hidden_states = attn.norm_out(hidden_states.movedim(1, -1)).movedim(-1, 1) - else: - hidden_states = attn.norm_out(hidden_states) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - return hidden_states - - -class LoRAAttnProcessor: - r""" - Processor for implementing attention with LoRA. - """ - - def __init__(self): - pass - - -class LoRAAttnProcessor2_0: - r""" - Processor for implementing attention with LoRA (enabled by default if you're using PyTorch 2.0). - """ - - def __init__(self): - pass - - -class LoRAXFormersAttnProcessor: - r""" - Processor for implementing attention with LoRA using xFormers. - """ - - def __init__(self): - pass - - -class LoRAAttnAddedKVProcessor: - r""" - Processor for implementing attention with LoRA with extra learnable key and value matrices for the text encoder. - """ - - def __init__(self): - pass - - -class FluxSingleAttnProcessor2_0(FluxAttnProcessor2_0): - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). - """ - - def __init__(self): - deprecation_message = "`FluxSingleAttnProcessor2_0` is deprecated and will be removed in a future version. Please use `FluxAttnProcessor2_0` instead." - deprecate("FluxSingleAttnProcessor2_0", "0.32.0", deprecation_message) - super().__init__() - - -class SanaLinearAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product linear attention. - """ - - def __call__( - self, - attn: Attention, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - original_dtype = hidden_states.dtype - - 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) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - query = query.transpose(1, 2).unflatten(1, (attn.heads, -1)) - key = key.transpose(1, 2).unflatten(1, (attn.heads, -1)).transpose(2, 3) - value = value.transpose(1, 2).unflatten(1, (attn.heads, -1)) - - query = F.relu(query) - key = F.relu(key) - - query, key, value = query.float(), key.float(), value.float() - - value = F.pad(value, (0, 0, 0, 1), mode="constant", value=1.0) - scores = torch.matmul(value, key) - hidden_states = torch.matmul(scores, query) - - hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + 1e-15) - hidden_states = hidden_states.flatten(1, 2).transpose(1, 2) - hidden_states = hidden_states.to(original_dtype) - - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - - if original_dtype == torch.float16: - hidden_states = hidden_states.clip(-65504, 65504) - - return hidden_states - - -class PAGCFGSanaLinearAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product linear attention. - """ - - def __call__( - self, - attn: Attention, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - original_dtype = hidden_states.dtype - - hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3) - hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org]) - - query = attn.to_q(hidden_states_org) - key = attn.to_k(hidden_states_org) - value = attn.to_v(hidden_states_org) - - query = query.transpose(1, 2).unflatten(1, (attn.heads, -1)) - key = key.transpose(1, 2).unflatten(1, (attn.heads, -1)).transpose(2, 3) - value = value.transpose(1, 2).unflatten(1, (attn.heads, -1)) - - query = F.relu(query) - key = F.relu(key) - - query, key, value = query.float(), key.float(), value.float() - - value = F.pad(value, (0, 0, 0, 1), mode="constant", value=1.0) - scores = torch.matmul(value, key) - hidden_states_org = torch.matmul(scores, query) - - hidden_states_org = hidden_states_org[:, :, :-1] / (hidden_states_org[:, :, -1:] + 1e-15) - hidden_states_org = hidden_states_org.flatten(1, 2).transpose(1, 2) - hidden_states_org = hidden_states_org.to(original_dtype) - - hidden_states_org = attn.to_out[0](hidden_states_org) - hidden_states_org = attn.to_out[1](hidden_states_org) - - # perturbed path (identity attention) - hidden_states_ptb = attn.to_v(hidden_states_ptb).to(original_dtype) - - hidden_states_ptb = attn.to_out[0](hidden_states_ptb) - hidden_states_ptb = attn.to_out[1](hidden_states_ptb) - - hidden_states = torch.cat([hidden_states_org, hidden_states_ptb]) - - if original_dtype == torch.float16: - hidden_states = hidden_states.clip(-65504, 65504) - - return hidden_states - - -class PAGIdentitySanaLinearAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product linear attention. - """ - - def __call__( - self, - attn: Attention, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - original_dtype = hidden_states.dtype - - hidden_states_org, hidden_states_ptb = hidden_states.chunk(2) - - query = attn.to_q(hidden_states_org) - key = attn.to_k(hidden_states_org) - value = attn.to_v(hidden_states_org) - - query = query.transpose(1, 2).unflatten(1, (attn.heads, -1)) - key = key.transpose(1, 2).unflatten(1, (attn.heads, -1)).transpose(2, 3) - value = value.transpose(1, 2).unflatten(1, (attn.heads, -1)) - - query = F.relu(query) - key = F.relu(key) - - query, key, value = query.float(), key.float(), value.float() - - value = F.pad(value, (0, 0, 0, 1), mode="constant", value=1.0) - scores = torch.matmul(value, key) - hidden_states_org = torch.matmul(scores, query) - - if hidden_states_org.dtype in [torch.float16, torch.bfloat16]: - hidden_states_org = hidden_states_org.float() - - hidden_states_org = hidden_states_org[:, :, :-1] / (hidden_states_org[:, :, -1:] + 1e-15) - hidden_states_org = hidden_states_org.flatten(1, 2).transpose(1, 2) - hidden_states_org = hidden_states_org.to(original_dtype) - - hidden_states_org = attn.to_out[0](hidden_states_org) - hidden_states_org = attn.to_out[1](hidden_states_org) - - # perturbed path (identity attention) - hidden_states_ptb = attn.to_v(hidden_states_ptb).to(original_dtype) - - hidden_states_ptb = attn.to_out[0](hidden_states_ptb) - hidden_states_ptb = attn.to_out[1](hidden_states_ptb) - - hidden_states = torch.cat([hidden_states_org, hidden_states_ptb]) - - if original_dtype == torch.float16: - hidden_states = hidden_states.clip(-65504, 65504) - - return hidden_states - - -ADDED_KV_ATTENTION_PROCESSORS = ( - AttnAddedKVProcessor, - SlicedAttnAddedKVProcessor, - AttnAddedKVProcessor2_0, - XFormersAttnAddedKVProcessor, -) - -CROSS_ATTENTION_PROCESSORS = ( - AttnProcessor, - AttnProcessor2_0, - XFormersAttnProcessor, - SlicedAttnProcessor, - IPAdapterAttnProcessor, - IPAdapterAttnProcessor2_0, - 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, -] diff --git a/venv/lib/python3.10/site-packages/diffusers/models/auto_model.py b/venv/lib/python3.10/site-packages/diffusers/models/auto_model.py deleted file mode 100644 index 96785ce6f5c3b030a4f091c774d83aae829baa86..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/auto_model.py +++ /dev/null @@ -1,206 +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 os -from typing import Optional, Union - -from huggingface_hub.utils import validate_hf_hub_args - -from ..configuration_utils import ConfigMixin -from ..utils import logging - - -logger = logging.get_logger(__name__) - - -class AutoModel(ConfigMixin): - config_name = "config.json" - - 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"`{self.__class__.__name__}.from_pipe(pipeline)` methods." - ) - - @classmethod - @validate_hf_hub_args - def from_pretrained(cls, pretrained_model_or_path: Optional[Union[str, os.PathLike]] = None, **kwargs): - r""" - Instantiate a pretrained PyTorch model from a pretrained model configuration. - - The model is set in evaluation mode - `model.eval()` - by default, and dropout modules are deactivated. To - train the model, set it back in training mode with `model.train()`. - - Parameters: - pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model 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*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - torch_dtype (`torch.dtype`, *optional*): - Override the default `torch.dtype` and load the model with another dtype. - 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*): - 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`): - Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - local_files_only(`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - from_flax (`bool`, *optional*, defaults to `False`): - Load the model weights from a Flax checkpoint save file. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - mirror (`str`, *optional*): - 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*): - 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. - - Set `device_map="auto"` to have 🤗 Accelerate automatically compute the most optimized `device_map`. For - more information about each option see [designing a device - map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). - max_memory (`Dict`, *optional*): - A dictionary device identifier for the maximum memory. Will default to the maximum memory available for - each GPU and the available CPU RAM if unset. - offload_folder (`str` or `os.PathLike`, *optional*): - The path to offload weights if `device_map` contains the value `"disk"`. - offload_state_dict (`bool`, *optional*): - If `True`, temporarily offloads the CPU state dict to the hard drive to avoid running out of CPU RAM if - the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to `True` - when there is some disk offload. - low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): - Speed up model loading only loading the pretrained weights and not initializing the weights. This also - tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. - Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this - argument to `True` will raise an error. - variant (`str`, *optional*): - Load weights from a specified `variant` filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. - use_safetensors (`bool`, *optional*, defaults to `None`): - If set to `None`, the `safetensors` weights are downloaded if they're available **and** if the - `safetensors` library is installed. If set to `True`, the model is forcibly loaded from `safetensors` - weights. If set to `False`, `safetensors` weights are not loaded. - 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. - - - - To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with - `huggingface-cli login`. You can also activate the special - ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a - firewalled environment. - - - - Example: - - ```py - from diffusers import AutoModel - - unet = AutoModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet") - ``` - - If you get the error message below, you need to finetune the weights for your downstream task: - - ```bash - Some weights of UNet2DConditionModel were not initialized from the model checkpoint at runwayml/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: - - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated - You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. - ``` - """ - cache_dir = kwargs.pop("cache_dir", None) - force_download = kwargs.pop("force_download", False) - proxies = kwargs.pop("proxies", None) - token = kwargs.pop("token", None) - local_files_only = kwargs.pop("local_files_only", False) - revision = kwargs.pop("revision", None) - subfolder = kwargs.pop("subfolder", None) - - load_config_kwargs = { - "cache_dir": cache_dir, - "force_download": force_download, - "proxies": proxies, - "token": token, - "local_files_only": local_files_only, - "revision": revision, - } - - library = None - orig_class_name = None - - # Always attempt to fetch model_index.json first - try: - cls.config_name = "model_index.json" - config = cls.load_config(pretrained_model_or_path, **load_config_kwargs) - - if subfolder is not None and subfolder in config: - library, orig_class_name = config[subfolder] - load_config_kwargs.update({"subfolder": subfolder}) - - except EnvironmentError as e: - logger.debug(e) - - # Unable to load from model_index.json so fallback to loading from config - if library is None and orig_class_name is None: - cls.config_name = "config.json" - config = cls.load_config(pretrained_model_or_path, subfolder=subfolder, **load_config_kwargs) - - if "_class_name" in config: - # If we find a class name in the config, we can try to load the model as a diffusers model - orig_class_name = config["_class_name"] - library = "diffusers" - load_config_kwargs.update({"subfolder": subfolder}) - elif "model_type" in config: - orig_class_name = "AutoModel" - library = "transformers" - load_config_kwargs.update({"subfolder": "" if subfolder is None else subfolder}) - else: - raise ValueError(f"Couldn't find model associated with the config file at {pretrained_model_or_path}.") - - 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=orig_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 {orig_class_name}.") - - kwargs = {**load_config_kwargs, **kwargs} - return model_cls.from_pretrained(pretrained_model_or_path, **kwargs) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__init__.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__init__.py deleted file mode 100644 index 742d747ae25e701b31cba0f413201aa283b9c2d9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -from .autoencoder_asym_kl import AsymmetricAutoencoderKL -from .autoencoder_dc import AutoencoderDC -from .autoencoder_kl import AutoencoderKL -from .autoencoder_kl_allegro import AutoencoderKLAllegro -from .autoencoder_kl_cogvideox import AutoencoderKLCogVideoX -from .autoencoder_kl_cosmos import AutoencoderKLCosmos -from .autoencoder_kl_hunyuan_video import AutoencoderKLHunyuanVideo -from .autoencoder_kl_ltx import AutoencoderKLLTXVideo -from .autoencoder_kl_magvit import AutoencoderKLMagvit -from .autoencoder_kl_mochi import AutoencoderKLMochi -from .autoencoder_kl_temporal_decoder import AutoencoderKLTemporalDecoder -from .autoencoder_kl_wan import AutoencoderKLWan -from .autoencoder_oobleck import AutoencoderOobleck -from .autoencoder_tiny import AutoencoderTiny -from .consistency_decoder_vae import ConsistencyDecoderVAE -from .vq_model import VQModel diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 87f85b4a29d5921a68b6561a848d5c404df8a037..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_asym_kl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_asym_kl.cpython-310.pyc deleted file mode 100644 index 0375f6be5e3271a869cec9d3f3498f06271f2f99..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_asym_kl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_dc.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_dc.cpython-310.pyc deleted file mode 100644 index 8a2e2754e4fdeb6465193e07ca5c7c5b7428855d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_dc.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl.cpython-310.pyc deleted file mode 100644 index 1486e6dbb35d808c896d78a67385d4c907b8bfe2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_allegro.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_allegro.cpython-310.pyc deleted file mode 100644 index d99d7c5a2d98d98d2eb3079487af1437607bd620..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_allegro.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_cogvideox.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_cogvideox.cpython-310.pyc deleted file mode 100644 index dd9eb22481961f6b7b654eb465cfe0233a520c94..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_cogvideox.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_cosmos.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_cosmos.cpython-310.pyc deleted file mode 100644 index 835d985078ea0ae1b7f1d1de15f6a800d610dea4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_cosmos.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_hunyuan_video.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_hunyuan_video.cpython-310.pyc deleted file mode 100644 index 4a88701a6be67c2b9b0b709498bc75b2f0e8106f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_hunyuan_video.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_ltx.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_ltx.cpython-310.pyc deleted file mode 100644 index b1f4e4d927eb482045fe2b63559c59cbe193ab55..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_ltx.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_magvit.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_magvit.cpython-310.pyc deleted file mode 100644 index 608e3b5f1176642a9d24c5646434b965fa4b4d3b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_magvit.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_mochi.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_mochi.cpython-310.pyc deleted file mode 100644 index 495be79d0ee94d3c2cf6e7a6a9f72a0a76279b55..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_mochi.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_temporal_decoder.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_temporal_decoder.cpython-310.pyc deleted file mode 100644 index a910680789aa0a445c6640848e16c46e2383321b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_temporal_decoder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_wan.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_wan.cpython-310.pyc deleted file mode 100644 index 0fa368a687be03a93a6c6d39e0e460f2f35f20f9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_kl_wan.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_oobleck.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_oobleck.cpython-310.pyc deleted file mode 100644 index 83e9e12785d50c686ad0a0dc8993b39bdbeef010..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_oobleck.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_tiny.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_tiny.cpython-310.pyc deleted file mode 100644 index 91463e34b1d6c19138f5d7d5d1ee3a1a07f8b471..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/autoencoder_tiny.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/consistency_decoder_vae.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/consistency_decoder_vae.cpython-310.pyc deleted file mode 100644 index 8b597f9d91b09910917b34c1ed668358e4763e08..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/consistency_decoder_vae.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/vae.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/vae.cpython-310.pyc deleted file mode 100644 index 82065af6ded060fb7a186228f9700a299efb7494..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/vae.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/vq_model.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/vq_model.cpython-310.pyc deleted file mode 100644 index 9815477da6e1de1917903dfea606086ee2559058..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/__pycache__/vq_model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_asym_kl.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_asym_kl.py deleted file mode 100644 index 54b1fc6771885a96f1042b4f9a3af22ea335c555..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_asym_kl.py +++ /dev/null @@ -1,186 +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 - -import torch -import torch.nn as nn - -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 DecoderOutput, DiagonalGaussianDistribution, Encoder, MaskConditionDecoder - - -class AsymmetricAutoencoderKL(ModelMixin, ConfigMixin): - r""" - Designing a Better Asymmetric VQGAN for StableDiffusion https://huggingface.co/papers/2306.04632 . A VAE model with - KL loss for encoding images into latents and decoding latent representations into images. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - 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. - 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. - 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. - latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space. - sample_size (`int`, *optional*, defaults to `32`): Sample input size. - norm_num_groups (`int`, *optional*, defaults to `32`): - Number of groups to use for the first normalization layer in ResNet blocks. - scaling_factor (`float`, *optional*, defaults to 0.18215): - 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. - """ - - _skip_layerwise_casting_patterns = ["decoder"] - - @register_to_config - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - 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,), - layers_per_up_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, - ) -> None: - super().__init__() - - # pass init params to Encoder - self.encoder = Encoder( - in_channels=in_channels, - out_channels=latent_channels, - down_block_types=down_block_types, - block_out_channels=down_block_out_channels, - layers_per_block=layers_per_down_block, - act_fn=act_fn, - norm_num_groups=norm_num_groups, - double_z=True, - ) - - # pass init params to Decoder - self.decoder = MaskConditionDecoder( - in_channels=latent_channels, - out_channels=out_channels, - up_block_types=up_block_types, - block_out_channels=up_block_out_channels, - layers_per_block=layers_per_up_block, - act_fn=act_fn, - norm_num_groups=norm_num_groups, - ) - - self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) - self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1) - - self.use_slicing = False - self.use_tiling = False - - self.register_to_config(block_out_channels=up_block_out_channels) - 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]]: - h = self.encoder(x) - moments = self.quant_conv(h) - posterior = DiagonalGaussianDistribution(moments) - - if not return_dict: - return (posterior,) - - return AutoencoderKLOutput(latent_dist=posterior) - - def _decode( - self, - z: torch.Tensor, - image: Optional[torch.Tensor] = None, - mask: Optional[torch.Tensor] = None, - return_dict: bool = True, - ) -> Union[DecoderOutput, Tuple[torch.Tensor]]: - z = self.post_quant_conv(z) - dec = self.decoder(z, image, mask) - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) - - @apply_forward_hook - def decode( - self, - z: torch.Tensor, - generator: Optional[torch.Generator] = None, - image: Optional[torch.Tensor] = None, - mask: Optional[torch.Tensor] = None, - return_dict: bool = True, - ) -> Union[DecoderOutput, Tuple[torch.Tensor]]: - decoded = self._decode(z, image, mask).sample - - if not return_dict: - return (decoded,) - - return DecoderOutput(sample=decoded) - - def forward( - self, - sample: torch.Tensor, - mask: Optional[torch.Tensor] = None, - sample_posterior: bool = False, - return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, Tuple[torch.Tensor]]: - r""" - Args: - sample (`torch.Tensor`): Input sample. - mask (`torch.Tensor`, *optional*, defaults to `None`): Optional inpainting mask. - sample_posterior (`bool`, *optional*, defaults to `False`): - Whether to sample from the posterior. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`DecoderOutput`] instead of a plain tuple. - """ - x = sample - posterior = self.encode(x).latent_dist - if sample_posterior: - z = posterior.sample(generator=generator) - else: - z = posterior.mode() - dec = self.decode(z, generator, sample, mask).sample - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_dc.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_dc.py deleted file mode 100644 index 9c7d6360e06e7fc7e11abafe1b2bd4b3e2da24aa..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_dc.py +++ /dev/null @@ -1,722 +0,0 @@ -# Copyright 2025 MIT, Tsinghua University, NVIDIA CORPORATION 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 Optional, Tuple, Union - -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.accelerate_utils import apply_forward_hook -from ..activations import get_activation -from ..attention_processor import SanaMultiscaleLinearAttention -from ..modeling_utils import ModelMixin -from ..normalization import RMSNorm, get_normalization -from ..transformers.sana_transformer import GLUMBConv -from .vae import DecoderOutput, EncoderOutput - - -class ResBlock(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - norm_type: str = "batch_norm", - act_fn: str = "relu6", - ) -> None: - super().__init__() - - self.norm_type = norm_type - - self.nonlinearity = get_activation(act_fn) if act_fn is not None else nn.Identity() - self.conv1 = nn.Conv2d(in_channels, in_channels, 3, 1, 1) - self.conv2 = nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=False) - self.norm = get_normalization(norm_type, out_channels) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - residual = hidden_states - hidden_states = self.conv1(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - hidden_states = self.conv2(hidden_states) - - if self.norm_type == "rms_norm": - # move channel to the last dimension so we apply RMSnorm across channel dimension - hidden_states = self.norm(hidden_states.movedim(1, -1)).movedim(-1, 1) - else: - hidden_states = self.norm(hidden_states) - - return hidden_states + residual - - -class EfficientViTBlock(nn.Module): - def __init__( - self, - in_channels: int, - mult: float = 1.0, - attention_head_dim: int = 32, - qkv_multiscales: Tuple[int, ...] = (5,), - norm_type: str = "batch_norm", - ) -> None: - super().__init__() - - self.attn = SanaMultiscaleLinearAttention( - in_channels=in_channels, - out_channels=in_channels, - mult=mult, - attention_head_dim=attention_head_dim, - norm_type=norm_type, - kernel_sizes=qkv_multiscales, - residual_connection=True, - ) - - self.conv_out = GLUMBConv( - in_channels=in_channels, - out_channels=in_channels, - norm_type="rms_norm", - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.attn(x) - x = self.conv_out(x) - return x - - -def get_block( - block_type: str, - in_channels: int, - out_channels: int, - attention_head_dim: int, - norm_type: str, - act_fn: str, - qkv_mutliscales: 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 - ) - - else: - raise ValueError(f"Block with {block_type=} is not supported.") - - return block - - -class DCDownBlock2d(nn.Module): - def __init__(self, in_channels: int, out_channels: int, downsample: bool = False, shortcut: bool = True) -> None: - super().__init__() - - self.downsample = downsample - self.factor = 2 - self.stride = 1 if downsample else 2 - self.group_size = in_channels * self.factor**2 // out_channels - self.shortcut = shortcut - - out_ratio = self.factor**2 - if downsample: - assert out_channels % out_ratio == 0 - out_channels = out_channels // out_ratio - - self.conv = nn.Conv2d( - in_channels, - out_channels, - kernel_size=3, - stride=self.stride, - padding=1, - ) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - x = self.conv(hidden_states) - if self.downsample: - x = F.pixel_unshuffle(x, self.factor) - - if self.shortcut: - y = F.pixel_unshuffle(hidden_states, self.factor) - y = y.unflatten(1, (-1, self.group_size)) - y = y.mean(dim=2) - hidden_states = x + y - else: - hidden_states = x - - return hidden_states - - -class DCUpBlock2d(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - interpolate: bool = False, - shortcut: bool = True, - interpolation_mode: str = "nearest", - ) -> None: - super().__init__() - - self.interpolate = interpolate - self.interpolation_mode = interpolation_mode - self.shortcut = shortcut - self.factor = 2 - self.repeats = out_channels * self.factor**2 // in_channels - - out_ratio = self.factor**2 - - if not interpolate: - out_channels = out_channels * out_ratio - - self.conv = nn.Conv2d(in_channels, out_channels, 3, 1, 1) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.interpolate: - x = F.interpolate(hidden_states, scale_factor=self.factor, mode=self.interpolation_mode) - x = self.conv(x) - else: - x = self.conv(hidden_states) - x = F.pixel_shuffle(x, self.factor) - - if self.shortcut: - y = hidden_states.repeat_interleave(self.repeats, dim=1, output_size=hidden_states.shape[1] * self.repeats) - y = F.pixel_shuffle(y, self.factor) - hidden_states = x + y - else: - hidden_states = x - - return hidden_states - - -class Encoder(nn.Module): - def __init__( - self, - 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,)), - downsample_block_type: str = "pixel_unshuffle", - out_shortcut: bool = True, - ): - super().__init__() - - num_blocks = len(block_out_channels) - - if isinstance(block_type, str): - block_type = (block_type,) * num_blocks - - if layers_per_block[0] > 0: - self.conv_in = nn.Conv2d( - in_channels, - block_out_channels[0] if layers_per_block[0] > 0 else block_out_channels[1], - kernel_size=3, - stride=1, - padding=1, - ) - else: - self.conv_in = DCDownBlock2d( - in_channels=in_channels, - out_channels=block_out_channels[0] if layers_per_block[0] > 0 else block_out_channels[1], - downsample=downsample_block_type == "pixel_unshuffle", - shortcut=False, - ) - - down_blocks = [] - for i, (out_channel, num_layers) in enumerate(zip(block_out_channels, layers_per_block)): - down_block_list = [] - - for _ in range(num_layers): - block = get_block( - block_type[i], - out_channel, - out_channel, - attention_head_dim=attention_head_dim, - norm_type="rms_norm", - act_fn="silu", - qkv_mutliscales=qkv_multiscales[i], - ) - down_block_list.append(block) - - if i < num_blocks - 1 and num_layers > 0: - downsample_block = DCDownBlock2d( - in_channels=out_channel, - out_channels=block_out_channels[i + 1], - downsample=downsample_block_type == "pixel_unshuffle", - shortcut=True, - ) - down_block_list.append(downsample_block) - - down_blocks.append(nn.Sequential(*down_block_list)) - - self.down_blocks = nn.ModuleList(down_blocks) - - self.conv_out = nn.Conv2d(block_out_channels[-1], latent_channels, 3, 1, 1) - - self.out_shortcut = out_shortcut - if out_shortcut: - self.out_shortcut_average_group_size = block_out_channels[-1] // latent_channels - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.conv_in(hidden_states) - for down_block in self.down_blocks: - hidden_states = down_block(hidden_states) - - if self.out_shortcut: - x = hidden_states.unflatten(1, (-1, self.out_shortcut_average_group_size)) - x = x.mean(dim=2) - hidden_states = self.conv_out(hidden_states) + x - else: - hidden_states = self.conv_out(hidden_states) - - return hidden_states - - -class Decoder(nn.Module): - def __init__( - self, - 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", - upsample_block_type: str = "pixel_shuffle", - in_shortcut: bool = True, - ): - super().__init__() - - num_blocks = len(block_out_channels) - - if isinstance(block_type, str): - block_type = (block_type,) * num_blocks - if isinstance(norm_type, str): - norm_type = (norm_type,) * num_blocks - if isinstance(act_fn, str): - act_fn = (act_fn,) * num_blocks - - self.conv_in = nn.Conv2d(latent_channels, block_out_channels[-1], 3, 1, 1) - - self.in_shortcut = in_shortcut - if in_shortcut: - self.in_shortcut_repeats = block_out_channels[-1] // latent_channels - - up_blocks = [] - for i, (out_channel, num_layers) in reversed(list(enumerate(zip(block_out_channels, layers_per_block)))): - up_block_list = [] - - if i < num_blocks - 1 and num_layers > 0: - upsample_block = DCUpBlock2d( - block_out_channels[i + 1], - out_channel, - interpolate=upsample_block_type == "interpolate", - shortcut=True, - ) - up_block_list.append(upsample_block) - - for _ in range(num_layers): - block = get_block( - block_type[i], - out_channel, - out_channel, - attention_head_dim=attention_head_dim, - norm_type=norm_type[i], - act_fn=act_fn[i], - qkv_mutliscales=qkv_multiscales[i], - ) - up_block_list.append(block) - - up_blocks.insert(0, nn.Sequential(*up_block_list)) - - self.up_blocks = nn.ModuleList(up_blocks) - - channels = block_out_channels[0] if layers_per_block[0] > 0 else block_out_channels[1] - - self.norm_out = RMSNorm(channels, 1e-5, elementwise_affine=True, bias=True) - self.conv_act = nn.ReLU() - self.conv_out = None - - if layers_per_block[0] > 0: - self.conv_out = nn.Conv2d(channels, in_channels, 3, 1, 1) - else: - self.conv_out = DCUpBlock2d( - channels, in_channels, interpolate=upsample_block_type == "interpolate", shortcut=False - ) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.in_shortcut: - x = hidden_states.repeat_interleave( - self.in_shortcut_repeats, dim=1, output_size=hidden_states.shape[1] * self.in_shortcut_repeats - ) - hidden_states = self.conv_in(hidden_states) + x - else: - hidden_states = self.conv_in(hidden_states) - - for up_block in reversed(self.up_blocks): - hidden_states = up_block(hidden_states) - - hidden_states = self.norm_out(hidden_states.movedim(1, -1)).movedim(-1, 1) - hidden_states = self.conv_act(hidden_states) - hidden_states = self.conv_out(hidden_states) - return hidden_states - - -class AutoencoderDC(ModelMixin, ConfigMixin, FromOriginalModelMixin): - r""" - An Autoencoder model introduced in [DCAE](https://huggingface.co/papers/2410.10733) and used in - [SANA](https://huggingface.co/papers/2410.10629). - - 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`): - 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"`): - The type(s) of block to use in the encoder. - decoder_block_types (`Union[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)`): - 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)`): - 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)`): - The number of layers per block in the encoder. - 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,))`): - Multi-scale configurations for the encoder's QKV (query-key-value) transformations. - 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"`): - The normalization type(s) to use in the decoder. - decoder_act_fns (`Union[str, Tuple[str]]`, defaults to `"silu"`): - The activation function(s) to use in the decoder. - scaling_factor (`float`, defaults to `1.0`): - The multiplicative inverse of the root mean square of the latent features. 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`. - """ - - _supports_gradient_checkpointing = False - - @register_to_config - def __init__( - self, - 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,)), - 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", - scaling_factor: float = 1.0, - ) -> None: - super().__init__() - - self.encoder = Encoder( - in_channels=in_channels, - latent_channels=latent_channels, - attention_head_dim=attention_head_dim, - block_type=encoder_block_types, - block_out_channels=encoder_block_out_channels, - layers_per_block=encoder_layers_per_block, - qkv_multiscales=encoder_qkv_multiscales, - downsample_block_type=downsample_block_type, - ) - self.decoder = Decoder( - in_channels=in_channels, - latent_channels=latent_channels, - attention_head_dim=attention_head_dim, - block_type=decoder_block_types, - block_out_channels=decoder_block_out_channels, - layers_per_block=decoder_layers_per_block, - qkv_multiscales=decoder_qkv_multiscales, - norm_type=decoder_norm_types, - act_fn=decoder_act_fns, - upsample_block_type=upsample_block_type, - ) - - self.spatial_compression_ratio = 2 ** (len(encoder_block_out_channels) - 1) - self.temporal_compression_ratio = 1 - - # 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 - - # The minimal tile height and width for spatial tiling to be used - self.tile_sample_min_height = 512 - self.tile_sample_min_width = 512 - - # The minimal distance between two spatial tiles - self.tile_sample_stride_height = 448 - self.tile_sample_stride_width = 448 - - self.tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio - self.tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio - - 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, - ) -> None: - r""" - Enable tiled AE decoding. When this option is enabled, the AE 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_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_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio - self.tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio - - def disable_tiling(self) -> None: - r""" - Disable tiled AE 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 AE decoding. When this option is enabled, the AE 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 AE 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: - batch_size, num_channels, height, width = x.shape - - if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): - return self.tiled_encode(x, return_dict=False)[0] - - encoded = self.encoder(x) - - return encoded - - @apply_forward_hook - def encode(self, x: torch.Tensor, return_dict: bool = True) -> Union[EncoderOutput, Tuple[torch.Tensor]]: - r""" - Encode a batch of images into latents. - - Args: - x (`torch.Tensor`): Input batch of images. - return_dict (`bool`, defaults to `True`): - Whether to return a [`~models.vae.EncoderOutput`] instead of a plain tuple. - - Returns: - The latent representations of the encoded videos. If `return_dict` is True, a - [`~models.vae.EncoderOutput`] is returned, otherwise a plain `tuple` is returned. - """ - if self.use_slicing and x.shape[0] > 1: - encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] - encoded = torch.cat(encoded_slices) - else: - encoded = self._encode(x) - - if not return_dict: - return (encoded,) - return EncoderOutput(latent=encoded) - - def _decode(self, z: torch.Tensor) -> torch.Tensor: - batch_size, num_channels, height, width = z.shape - - if self.use_tiling and (width > self.tile_latent_min_width or height > self.tile_latent_min_height): - return self.tiled_decode(z, return_dict=False)[0] - - decoded = self.decoder(z) - - return decoded - - @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, Tuple[torch.Tensor]]: - r""" - Decode a batch of images. - - Args: - z (`torch.Tensor`): Input batch of latent vectors. - return_dict (`bool`, 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.size(0) > 1: - decoded_slices = [self._decode(z_slice).sample 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 blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: - blend_extent = min(a.shape[2], b.shape[2], 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[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, return_dict: bool = True) -> torch.Tensor: - batch_size, num_channels, 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, x.shape[2], self.tile_sample_stride_height): - row = [] - for j in range(0, x.shape[3], self.tile_sample_stride_width): - tile = x[:, :, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width] - if ( - tile.shape[2] % self.spatial_compression_ratio != 0 - or tile.shape[3] % self.spatial_compression_ratio != 0 - ): - pad_h = (self.spatial_compression_ratio - tile.shape[2]) % self.spatial_compression_ratio - pad_w = (self.spatial_compression_ratio - tile.shape[3]) % self.spatial_compression_ratio - tile = F.pad(tile, (0, pad_w, 0, pad_h)) - tile = self.encoder(tile) - row.append(tile) - 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=3)) - - encoded = torch.cat(result_rows, dim=2)[:, :, :latent_height, :latent_width] - - if not return_dict: - return (encoded,) - return EncoderOutput(latent=encoded) - - def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: - batch_size, num_channels, 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_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): - tile = z[:, :, i : i + tile_latent_min_height, j : j + tile_latent_min_width] - decoded = self.decoder(tile) - row.append(decoded) - 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=3)) - - decoded = torch.cat(result_rows, dim=2) - - if not return_dict: - return (decoded,) - return DecoderOutput(sample=decoded) - - def forward(self, sample: torch.Tensor, return_dict: bool = True) -> torch.Tensor: - encoded = self.encode(sample, return_dict=False)[0] - decoded = self.decode(encoded, return_dict=False)[0] - if not return_dict: - return (decoded,) - return DecoderOutput(sample=decoded) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl.py deleted file mode 100644 index 640ee3492807c2f884261cd968079cfefdc2691c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl.py +++ /dev/null @@ -1,567 +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 Dict, Optional, Tuple, Union - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin -from ...loaders.single_file_model import FromOriginalModelMixin -from ...utils import deprecate -from ...utils.accelerate_utils import apply_forward_hook -from ..attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - Attention, - AttentionProcessor, - AttnAddedKVProcessor, - AttnProcessor, - FusedAttnProcessor2_0, -) -from ..modeling_outputs import AutoencoderKLOutput -from ..modeling_utils import ModelMixin -from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder - - -class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin): - r""" - A VAE model with KL loss for encoding images into latents and decoding latent representations into images. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - 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. - 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. - scaling_factor (`float`, *optional*, defaults to 0.18215): - 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. - force_upcast (`bool`, *optional*, default to `True`): - If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE - can be fine-tuned / trained to a lower range without losing too much precision in which case `force_upcast` - can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix - mid_block_add_attention (`bool`, *optional*, default to `True`): - If enabled, the mid_block of the Encoder and Decoder will have attention blocks. If set to false, the - mid_block will only have resnet blocks - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"] - - @register_to_config - 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,), - 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, - force_upcast: float = True, - use_quant_conv: bool = True, - use_post_quant_conv: bool = True, - mid_block_add_attention: bool = True, - ): - super().__init__() - - # pass init params to Encoder - self.encoder = Encoder( - in_channels=in_channels, - out_channels=latent_channels, - down_block_types=down_block_types, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - act_fn=act_fn, - norm_num_groups=norm_num_groups, - double_z=True, - mid_block_add_attention=mid_block_add_attention, - ) - - # pass init params to Decoder - self.decoder = Decoder( - in_channels=latent_channels, - out_channels=out_channels, - up_block_types=up_block_types, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - norm_num_groups=norm_num_groups, - act_fn=act_fn, - mid_block_add_attention=mid_block_add_attention, - ) - - self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) if use_quant_conv else None - self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1) if use_post_quant_conv else None - - self.use_slicing = False - self.use_tiling = False - - # only relevant if vae tiling is enabled - self.tile_sample_min_size = self.config.sample_size - sample_size = ( - self.config.sample_size[0] - if isinstance(self.config.sample_size, (list, tuple)) - else self.config.sample_size - ) - self.tile_latent_min_size = int(sample_size / (2 ** (len(self.config.block_out_channels) - 1))) - self.tile_overlap_factor = 0.25 - - def enable_tiling(self, use_tiling: bool = True): - 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. - """ - self.use_tiling = use_tiling - - def disable_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.enable_tiling(False) - - def enable_slicing(self): - 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): - 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 - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - self.set_attn_processor(processor) - - def _encode(self, x: torch.Tensor) -> torch.Tensor: - batch_size, num_channels, height, width = x.shape - - if self.use_tiling and (width > self.tile_sample_min_size or height > self.tile_sample_min_size): - return self._tiled_encode(x) - - enc = self.encoder(x) - if self.quant_conv is not None: - enc = self.quant_conv(enc) - - return enc - - @apply_forward_hook - def encode( - self, x: torch.Tensor, return_dict: bool = True - ) -> Union[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 images. 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) 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, return_dict: bool = True) -> Union[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) - - if self.post_quant_conv is not None: - z = self.post_quant_conv(z) - - dec = self.decoder(z) - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) - - @apply_forward_hook - def decode( - self, z: torch.FloatTensor, return_dict: bool = True, generator=None - ) -> Union[DecoderOutput, torch.FloatTensor]: - """ - 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: - decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] - decoded = torch.cat(decoded_slices) - else: - decoded = self._decode(z).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[2], b.shape[2], 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[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) -> torch.Tensor: - 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 - steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is - different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the - tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the - output, but they should be much less noticeable. - - Args: - x (`torch.Tensor`): Input batch of images. - - Returns: - `torch.Tensor`: - The latent representation of the encoded videos. - """ - - overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) - blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) - row_limit = self.tile_latent_min_size - blend_extent - - # Split the image into 512x512 tiles and encode them separately. - rows = [] - for i in range(0, x.shape[2], overlap_size): - row = [] - for j in range(0, x.shape[3], overlap_size): - tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size] - tile = self.encoder(tile) - if self.config.use_quant_conv: - tile = self.quant_conv(tile) - row.append(tile) - 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_extent) - if j > 0: - tile = self.blend_h(row[j - 1], tile, blend_extent) - result_row.append(tile[:, :, :row_limit, :row_limit]) - result_rows.append(torch.cat(result_row, dim=3)) - - enc = torch.cat(result_rows, dim=2) - return enc - - def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutput: - 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 - steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is - different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the - tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the - output, but they should be much less noticeable. - - Args: - x (`torch.Tensor`): Input batch of images. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. - - Returns: - [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`: - If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain - `tuple` is returned. - """ - deprecation_message = ( - "The tiled_encode implementation supporting the `return_dict` parameter is deprecated. In the future, the " - "implementation of this method will be replaced with that of `_tiled_encode` and you will no longer be able " - "to pass `return_dict`. You will also have to create a `DiagonalGaussianDistribution()` from the returned value." - ) - deprecate("tiled_encode", "1.0.0", deprecation_message, standard_warn=False) - - overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) - blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) - row_limit = self.tile_latent_min_size - blend_extent - - # Split the image into 512x512 tiles and encode them separately. - rows = [] - for i in range(0, x.shape[2], overlap_size): - row = [] - for j in range(0, x.shape[3], overlap_size): - tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size] - tile = self.encoder(tile) - if self.config.use_quant_conv: - tile = self.quant_conv(tile) - row.append(tile) - 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_extent) - if j > 0: - tile = self.blend_h(row[j - 1], tile, blend_extent) - result_row.append(tile[:, :, :row_limit, :row_limit]) - result_rows.append(torch.cat(result_row, dim=3)) - - moments = torch.cat(result_rows, dim=2) - posterior = DiagonalGaussianDistribution(moments) - - if not return_dict: - return (posterior,) - - return AutoencoderKLOutput(latent_dist=posterior) - - def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[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. - """ - overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) - blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) - row_limit = self.tile_sample_min_size - blend_extent - - # Split z into overlapping 64x64 tiles and decode them separately. - # The tiles have an overlap to avoid seams between tiles. - rows = [] - for i in range(0, z.shape[2], overlap_size): - row = [] - for j in range(0, z.shape[3], overlap_size): - tile = z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size] - if self.config.use_post_quant_conv: - tile = self.post_quant_conv(tile) - decoded = self.decoder(tile) - row.append(decoded) - 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_extent) - if j > 0: - tile = self.blend_h(row[j - 1], tile, blend_extent) - result_row.append(tile[:, :, :row_limit, :row_limit]) - result_rows.append(torch.cat(result_row, dim=3)) - - dec = torch.cat(result_rows, dim=2) - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) - - def forward( - self, - sample: torch.Tensor, - sample_posterior: bool = False, - return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: - r""" - Args: - sample (`torch.Tensor`): Input sample. - sample_posterior (`bool`, *optional*, defaults to `False`): - Whether to sample from the posterior. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`DecoderOutput`] instead of a plain tuple. - """ - x = sample - posterior = self.encode(x).latent_dist - if sample_posterior: - z = posterior.sample(generator=generator) - else: - z = posterior.mode() - dec = self.decode(z).sample - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - self.set_attn_processor(FusedAttnProcessor2_0()) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - if self.original_attn_processors is not None: - self.set_attn_processor(self.original_attn_processors) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_allegro.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_allegro.py deleted file mode 100644 index c24b8f42aca4addd26ec421bbf48f1cd737b52c5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_allegro.py +++ /dev/null @@ -1,1131 +0,0 @@ -# Copyright 2025 The RhymesAI 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 typing import Optional, Tuple, Union - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils.accelerate_utils import apply_forward_hook -from ..attention_processor import Attention, SpatialNorm -from ..autoencoders.vae import DecoderOutput, DiagonalGaussianDistribution -from ..downsampling import Downsample2D -from ..modeling_outputs import AutoencoderKLOutput -from ..modeling_utils import ModelMixin -from ..resnet import ResnetBlock2D -from ..upsampling import Upsample2D - - -class AllegroTemporalConvLayer(nn.Module): - r""" - Temporal convolutional layer that can be used for video (sequence of images) input. Code adapted from: - https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/models/multi_modal/video_synthesis/unet_sd.py#L1016 - """ - - def __init__( - self, - in_dim: int, - out_dim: Optional[int] = None, - dropout: float = 0.0, - norm_num_groups: int = 32, - up_sample: bool = False, - down_sample: bool = False, - stride: int = 1, - ) -> None: - super().__init__() - - out_dim = out_dim or in_dim - pad_h = pad_w = int((stride - 1) * 0.5) - pad_t = 0 - - self.down_sample = down_sample - self.up_sample = up_sample - - if down_sample: - self.conv1 = nn.Sequential( - nn.GroupNorm(norm_num_groups, in_dim), - nn.SiLU(), - nn.Conv3d(in_dim, out_dim, (2, stride, stride), stride=(2, 1, 1), padding=(0, pad_h, pad_w)), - ) - elif up_sample: - self.conv1 = nn.Sequential( - nn.GroupNorm(norm_num_groups, in_dim), - nn.SiLU(), - nn.Conv3d(in_dim, out_dim * 2, (1, stride, stride), padding=(0, pad_h, pad_w)), - ) - else: - self.conv1 = nn.Sequential( - nn.GroupNorm(norm_num_groups, in_dim), - nn.SiLU(), - nn.Conv3d(in_dim, out_dim, (3, stride, stride), padding=(pad_t, pad_h, pad_w)), - ) - self.conv2 = nn.Sequential( - nn.GroupNorm(norm_num_groups, out_dim), - nn.SiLU(), - nn.Dropout(dropout), - nn.Conv3d(out_dim, in_dim, (3, stride, stride), padding=(pad_t, pad_h, pad_w)), - ) - self.conv3 = nn.Sequential( - nn.GroupNorm(norm_num_groups, out_dim), - nn.SiLU(), - nn.Dropout(dropout), - nn.Conv3d(out_dim, in_dim, (3, stride, stride), padding=(pad_t, pad_h, pad_h)), - ) - self.conv4 = nn.Sequential( - nn.GroupNorm(norm_num_groups, out_dim), - nn.SiLU(), - nn.Conv3d(out_dim, in_dim, (3, stride, stride), padding=(pad_t, pad_h, pad_h)), - ) - - @staticmethod - def _pad_temporal_dim(hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = torch.cat((hidden_states[:, :, 0:1], hidden_states), dim=2) - hidden_states = torch.cat((hidden_states, hidden_states[:, :, -1:]), dim=2) - return hidden_states - - def forward(self, hidden_states: torch.Tensor, batch_size: int) -> torch.Tensor: - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - - if self.down_sample: - identity = hidden_states[:, :, ::2] - elif self.up_sample: - identity = hidden_states.repeat_interleave(2, dim=2, output_size=hidden_states.shape[2] * 2) - else: - identity = hidden_states - - if self.down_sample or self.up_sample: - hidden_states = self.conv1(hidden_states) - else: - hidden_states = self._pad_temporal_dim(hidden_states) - hidden_states = self.conv1(hidden_states) - - if self.up_sample: - hidden_states = hidden_states.unflatten(1, (2, -1)).permute(0, 2, 3, 1, 4, 5).flatten(2, 3) - - hidden_states = self._pad_temporal_dim(hidden_states) - hidden_states = self.conv2(hidden_states) - - hidden_states = self._pad_temporal_dim(hidden_states) - hidden_states = self.conv3(hidden_states) - - hidden_states = self._pad_temporal_dim(hidden_states) - hidden_states = self.conv4(hidden_states) - - hidden_states = identity + hidden_states - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) - - return hidden_states - - -class AllegroDownBlock3D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor: float = 1.0, - spatial_downsample: bool = True, - temporal_downsample: bool = False, - downsample_padding: int = 1, - ): - super().__init__() - - resnets = [] - temp_convs = [] - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=None, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - temp_convs.append( - AllegroTemporalConvLayer( - out_channels, - out_channels, - dropout=0.1, - norm_num_groups=resnet_groups, - ) - ) - - self.resnets = nn.ModuleList(resnets) - self.temp_convs = nn.ModuleList(temp_convs) - - if temporal_downsample: - self.temp_convs_down = AllegroTemporalConvLayer( - out_channels, out_channels, dropout=0.1, norm_num_groups=resnet_groups, down_sample=True, stride=3 - ) - self.add_temp_downsample = temporal_downsample - - if spatial_downsample: - self.downsamplers = nn.ModuleList( - [ - Downsample2D( - out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" - ) - ] - ) - else: - self.downsamplers = None - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - batch_size = hidden_states.shape[0] - - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) - - for resnet, temp_conv in zip(self.resnets, self.temp_convs): - hidden_states = resnet(hidden_states, temb=None) - hidden_states = temp_conv(hidden_states, batch_size=batch_size) - - if self.add_temp_downsample: - hidden_states = self.temp_convs_down(hidden_states, batch_size=batch_size) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - return hidden_states - - -class AllegroUpBlock3D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", # default, spatial - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor: float = 1.0, - spatial_upsample: bool = True, - temporal_upsample: bool = False, - temb_channels: Optional[int] = None, - ): - super().__init__() - - resnets = [] - temp_convs = [] - - for i in range(num_layers): - input_channels = in_channels if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=input_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - temp_convs.append( - AllegroTemporalConvLayer( - out_channels, - out_channels, - dropout=0.1, - norm_num_groups=resnet_groups, - ) - ) - - self.resnets = nn.ModuleList(resnets) - self.temp_convs = nn.ModuleList(temp_convs) - - self.add_temp_upsample = temporal_upsample - if temporal_upsample: - self.temp_conv_up = AllegroTemporalConvLayer( - out_channels, out_channels, dropout=0.1, norm_num_groups=resnet_groups, up_sample=True, stride=3 - ) - - if spatial_upsample: - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - batch_size = hidden_states.shape[0] - - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) - - for resnet, temp_conv in zip(self.resnets, self.temp_convs): - hidden_states = resnet(hidden_states, temb=None) - hidden_states = temp_conv(hidden_states, batch_size=batch_size) - - if self.add_temp_upsample: - hidden_states = self.temp_conv_up(hidden_states, batch_size=batch_size) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states) - - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - return hidden_states - - -class AllegroMidBlock3DConv(nn.Module): - def __init__( - self, - in_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", # default, spatial - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - add_attention: bool = True, - attention_head_dim: int = 1, - output_scale_factor: float = 1.0, - ): - super().__init__() - - # there is always at least one resnet - resnets = [ - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ] - temp_convs = [ - AllegroTemporalConvLayer( - in_channels, - in_channels, - dropout=0.1, - norm_num_groups=resnet_groups, - ) - ] - attentions = [] - - if attention_head_dim is None: - attention_head_dim = in_channels - - for _ in range(num_layers): - if add_attention: - attentions.append( - Attention( - in_channels, - heads=in_channels // attention_head_dim, - dim_head=attention_head_dim, - rescale_output_factor=output_scale_factor, - eps=resnet_eps, - norm_num_groups=resnet_groups if resnet_time_scale_shift == "default" else None, - spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None, - residual_connection=True, - bias=True, - upcast_softmax=True, - _from_deprecated_attn_block=True, - ) - ) - else: - attentions.append(None) - - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - temp_convs.append( - AllegroTemporalConvLayer( - in_channels, - in_channels, - dropout=0.1, - norm_num_groups=resnet_groups, - ) - ) - - self.resnets = nn.ModuleList(resnets) - self.temp_convs = nn.ModuleList(temp_convs) - self.attentions = nn.ModuleList(attentions) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - batch_size = hidden_states.shape[0] - - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) - hidden_states = self.resnets[0](hidden_states, temb=None) - - hidden_states = self.temp_convs[0](hidden_states, batch_size=batch_size) - - for attn, resnet, temp_conv in zip(self.attentions, self.resnets[1:], self.temp_convs[1:]): - hidden_states = attn(hidden_states) - hidden_states = resnet(hidden_states, temb=None) - hidden_states = temp_conv(hidden_states, batch_size=batch_size) - - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - return hidden_states - - -class AllegroEncoder3D(nn.Module): - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - 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], - layers_per_block: int = 2, - norm_num_groups: int = 32, - act_fn: str = "silu", - double_z: bool = True, - ): - super().__init__() - - self.conv_in = nn.Conv2d( - in_channels, - block_out_channels[0], - kernel_size=3, - stride=1, - padding=1, - ) - - self.temp_conv_in = nn.Conv3d( - in_channels=block_out_channels[0], - out_channels=block_out_channels[0], - kernel_size=(3, 1, 1), - padding=(1, 0, 0), - ) - - self.down_blocks = nn.ModuleList([]) - - # down - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - if down_block_type == "AllegroDownBlock3D": - down_block = AllegroDownBlock3D( - num_layers=layers_per_block, - in_channels=input_channel, - out_channels=output_channel, - spatial_downsample=not is_final_block, - temporal_downsample=temporal_downsample_blocks[i], - resnet_eps=1e-6, - downsample_padding=0, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - ) - else: - raise ValueError("Invalid `down_block_type` encountered. Must be `AllegroDownBlock3D`") - - self.down_blocks.append(down_block) - - # mid - self.mid_block = AllegroMidBlock3DConv( - in_channels=block_out_channels[-1], - resnet_eps=1e-6, - resnet_act_fn=act_fn, - output_scale_factor=1, - resnet_time_scale_shift="default", - attention_head_dim=block_out_channels[-1], - resnet_groups=norm_num_groups, - temb_channels=None, - ) - - # out - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) - self.conv_act = nn.SiLU() - - conv_out_channels = 2 * out_channels if double_z else out_channels - - self.temp_conv_out = nn.Conv3d(block_out_channels[-1], block_out_channels[-1], (3, 1, 1), padding=(1, 0, 0)) - self.conv_out = nn.Conv2d(block_out_channels[-1], conv_out_channels, 3, padding=1) - - self.gradient_checkpointing = False - - def forward(self, sample: torch.Tensor) -> torch.Tensor: - batch_size = sample.shape[0] - - sample = sample.permute(0, 2, 1, 3, 4).flatten(0, 1) - sample = self.conv_in(sample) - - sample = sample.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - residual = sample - sample = self.temp_conv_in(sample) - sample = sample + residual - - if torch.is_grad_enabled() and self.gradient_checkpointing: - # Down blocks - for down_block in self.down_blocks: - sample = self._gradient_checkpointing_func(down_block, sample) - - # Mid block - sample = self._gradient_checkpointing_func(self.mid_block, sample) - else: - # Down blocks - for down_block in self.down_blocks: - sample = down_block(sample) - - # Mid block - sample = self.mid_block(sample) - - # Post process - sample = sample.permute(0, 2, 1, 3, 4).flatten(0, 1) - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - - sample = sample.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - residual = sample - sample = self.temp_conv_out(sample) - sample = sample + residual - - sample = sample.permute(0, 2, 1, 3, 4).flatten(0, 1) - sample = self.conv_out(sample) - - sample = sample.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - return sample - - -class AllegroDecoder3D(nn.Module): - def __init__( - self, - in_channels: int = 4, - out_channels: int = 3, - 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), - layers_per_block: int = 2, - norm_num_groups: int = 32, - act_fn: str = "silu", - norm_type: str = "group", # group, spatial - ): - super().__init__() - - self.conv_in = nn.Conv2d( - in_channels, - block_out_channels[-1], - kernel_size=3, - stride=1, - padding=1, - ) - - self.temp_conv_in = nn.Conv3d(block_out_channels[-1], block_out_channels[-1], (3, 1, 1), padding=(1, 0, 0)) - - self.mid_block = None - self.up_blocks = nn.ModuleList([]) - - temb_channels = in_channels if norm_type == "spatial" else None - - # mid - self.mid_block = AllegroMidBlock3DConv( - in_channels=block_out_channels[-1], - resnet_eps=1e-6, - resnet_act_fn=act_fn, - output_scale_factor=1, - resnet_time_scale_shift="default" if norm_type == "group" else norm_type, - attention_head_dim=block_out_channels[-1], - resnet_groups=norm_num_groups, - temb_channels=temb_channels, - ) - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - - is_final_block = i == len(block_out_channels) - 1 - - if up_block_type == "AllegroUpBlock3D": - up_block = AllegroUpBlock3D( - num_layers=layers_per_block + 1, - in_channels=prev_output_channel, - out_channels=output_channel, - spatial_upsample=not is_final_block, - temporal_upsample=temporal_upsample_blocks[i], - resnet_eps=1e-6, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - temb_channels=temb_channels, - resnet_time_scale_shift=norm_type, - ) - else: - raise ValueError("Invalid `UP_block_type` encountered. Must be `AllegroUpBlock3D`") - - self.up_blocks.append(up_block) - prev_output_channel = output_channel - - # out - if norm_type == "spatial": - self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels) - else: - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) - - self.conv_act = nn.SiLU() - - self.temp_conv_out = nn.Conv3d(block_out_channels[0], block_out_channels[0], (3, 1, 1), padding=(1, 0, 0)) - self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1) - - self.gradient_checkpointing = False - - def forward(self, sample: torch.Tensor) -> torch.Tensor: - batch_size = sample.shape[0] - - sample = sample.permute(0, 2, 1, 3, 4).flatten(0, 1) - sample = self.conv_in(sample) - - sample = sample.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - residual = sample - sample = self.temp_conv_in(sample) - sample = sample + residual - - upscale_dtype = next(iter(self.up_blocks.parameters())).dtype - - if torch.is_grad_enabled() and self.gradient_checkpointing: - # Mid block - sample = self._gradient_checkpointing_func(self.mid_block, sample) - - # Up blocks - for up_block in self.up_blocks: - sample = self._gradient_checkpointing_func(up_block, sample) - - else: - # Mid block - sample = self.mid_block(sample) - sample = sample.to(upscale_dtype) - - # Up blocks - for up_block in self.up_blocks: - sample = up_block(sample) - - # Post process - sample = sample.permute(0, 2, 1, 3, 4).flatten(0, 1) - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - - sample = sample.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - residual = sample - sample = self.temp_conv_out(sample) - sample = sample + residual - - sample = sample.permute(0, 2, 1, 3, 4).flatten(0, 1) - sample = self.conv_out(sample) - - sample = sample.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - return sample - - -class AutoencoderKLAllegro(ModelMixin, ConfigMixin): - r""" - A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Used in - [Allegro](https://github.com/rhymes-ai/Allegro). - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - Parameters: - in_channels (int, defaults to `3`): - 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. - latent_channels (`int`, defaults to `4`): - Number of channels in latents. - layers_per_block (`int`, defaults to `2`): - Number of resnet or attention or temporal convolution layers per down/up block. - act_fn (`str`, defaults to `"silu"`): - The activation function to use. - norm_num_groups (`int`, defaults to `32`): - Number of groups to use in normalization layers. - temporal_compression_ratio (`int`, defaults to `4`): - Ratio by which temporal dimension of samples are compressed. - sample_size (`int`, defaults to `320`): - Default latent size. - scaling_factor (`float`, defaults to `0.13235`): - 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. - force_upcast (`bool`, default to `True`): - If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE - can be fine-tuned / trained to a lower range without losing too much precision in which case `force_upcast` - can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix - """ - - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - down_block_types: Tuple[str, ...] = ( - "AllegroDownBlock3D", - "AllegroDownBlock3D", - "AllegroDownBlock3D", - "AllegroDownBlock3D", - ), - 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), - latent_channels: int = 4, - layers_per_block: int = 2, - act_fn: str = "silu", - norm_num_groups: int = 32, - temporal_compression_ratio: float = 4, - sample_size: int = 320, - scaling_factor: float = 0.13, - force_upcast: bool = True, - ) -> None: - super().__init__() - - self.encoder = AllegroEncoder3D( - in_channels=in_channels, - out_channels=latent_channels, - down_block_types=down_block_types, - temporal_downsample_blocks=temporal_downsample_blocks, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - act_fn=act_fn, - norm_num_groups=norm_num_groups, - double_z=True, - ) - self.decoder = AllegroDecoder3D( - in_channels=latent_channels, - out_channels=out_channels, - up_block_types=up_block_types, - temporal_upsample_blocks=temporal_upsample_blocks, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - norm_num_groups=norm_num_groups, - act_fn=act_fn, - ) - self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) - self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1) - - # TODO(aryan): For the 1.0.0 refactor, `temporal_compression_ratio` can be inferred directly and we don't need - # to use a specific parameter here or in other VAEs. - - self.use_slicing = False - self.use_tiling = False - - self.spatial_compression_ratio = 2 ** (len(block_out_channels) - 1) - self.tile_overlap_t = 8 - self.tile_overlap_h = 120 - self.tile_overlap_w = 80 - sample_frames = 24 - - self.kernel = (sample_frames, sample_size, sample_size) - self.stride = ( - sample_frames - self.tile_overlap_t, - sample_size - self.tile_overlap_h, - sample_size - self.tile_overlap_w, - ) - - def enable_tiling(self) -> 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. - """ - self.use_tiling = True - - 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: - # TODO(aryan) - # if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): - if self.use_tiling: - return self.tiled_encode(x) - - raise NotImplementedError("Encoding without tiling has not been implemented yet.") - - @apply_forward_hook - def encode( - self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: - r""" - Encode a batch of videos into latents. - - Args: - x (`torch.Tensor`): - Input batch of videos. - return_dict (`bool`, 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) 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: - # TODO(aryan): refactor tiling implementation - # if self.use_tiling and (width > self.tile_latent_min_width or height > self.tile_latent_min_height): - if self.use_tiling: - return self.tiled_decode(z) - - 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]: - """ - Decode a batch of videos. - - Args: - z (`torch.Tensor`): - Input batch of latent vectors. - return_dict (`bool`, 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: - 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 tiled_encode(self, x: torch.Tensor) -> torch.Tensor: - local_batch_size = 1 - rs = self.spatial_compression_ratio - rt = self.config.temporal_compression_ratio - - batch_size, num_channels, num_frames, height, width = x.shape - - output_num_frames = math.floor((num_frames - self.kernel[0]) / self.stride[0]) + 1 - output_height = math.floor((height - self.kernel[1]) / self.stride[1]) + 1 - output_width = math.floor((width - self.kernel[2]) / self.stride[2]) + 1 - - count = 0 - output_latent = x.new_zeros( - ( - output_num_frames * output_height * output_width, - 2 * self.config.latent_channels, - self.kernel[0] // rt, - self.kernel[1] // rs, - self.kernel[2] // rs, - ) - ) - vae_batch_input = x.new_zeros((local_batch_size, num_channels, self.kernel[0], self.kernel[1], self.kernel[2])) - - for i in range(output_num_frames): - for j in range(output_height): - for k in range(output_width): - n_start, n_end = i * self.stride[0], i * self.stride[0] + self.kernel[0] - h_start, h_end = j * self.stride[1], j * self.stride[1] + self.kernel[1] - w_start, w_end = k * self.stride[2], k * self.stride[2] + self.kernel[2] - - video_cube = x[:, :, n_start:n_end, h_start:h_end, w_start:w_end] - vae_batch_input[count % local_batch_size] = video_cube - - if ( - count % local_batch_size == local_batch_size - 1 - or count == output_num_frames * output_height * output_width - 1 - ): - latent = self.encoder(vae_batch_input) - - if ( - count == output_num_frames * output_height * output_width - 1 - and count % local_batch_size != local_batch_size - 1 - ): - output_latent[count - count % local_batch_size :] = latent[: count % local_batch_size + 1] - else: - output_latent[count - local_batch_size + 1 : count + 1] = latent - - vae_batch_input = x.new_zeros( - (local_batch_size, num_channels, self.kernel[0], self.kernel[1], self.kernel[2]) - ) - - count += 1 - - latent = x.new_zeros( - (batch_size, 2 * self.config.latent_channels, num_frames // rt, height // rs, width // rs) - ) - output_kernel = self.kernel[0] // rt, self.kernel[1] // rs, self.kernel[2] // rs - output_stride = self.stride[0] // rt, self.stride[1] // rs, self.stride[2] // rs - output_overlap = ( - output_kernel[0] - output_stride[0], - output_kernel[1] - output_stride[1], - output_kernel[2] - output_stride[2], - ) - - for i in range(output_num_frames): - n_start, n_end = i * output_stride[0], i * output_stride[0] + output_kernel[0] - for j in range(output_height): - h_start, h_end = j * output_stride[1], j * output_stride[1] + output_kernel[1] - for k in range(output_width): - w_start, w_end = k * output_stride[2], k * output_stride[2] + output_kernel[2] - latent_mean = _prepare_for_blend( - (i, output_num_frames, output_overlap[0]), - (j, output_height, output_overlap[1]), - (k, output_width, output_overlap[2]), - output_latent[i * output_height * output_width + j * output_width + k].unsqueeze(0), - ) - latent[:, :, n_start:n_end, h_start:h_end, w_start:w_end] += latent_mean - - latent = latent.permute(0, 2, 1, 3, 4).flatten(0, 1) - latent = self.quant_conv(latent) - latent = latent.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - return latent - - def tiled_decode(self, z: torch.Tensor) -> torch.Tensor: - local_batch_size = 1 - rs = self.spatial_compression_ratio - rt = self.config.temporal_compression_ratio - - latent_kernel = self.kernel[0] // rt, self.kernel[1] // rs, self.kernel[2] // rs - latent_stride = self.stride[0] // rt, self.stride[1] // rs, self.stride[2] // rs - - batch_size, num_channels, num_frames, height, width = z.shape - - ## post quant conv (a mapping) - z = z.permute(0, 2, 1, 3, 4).flatten(0, 1) - z = self.post_quant_conv(z) - z = z.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - - output_num_frames = math.floor((num_frames - latent_kernel[0]) / latent_stride[0]) + 1 - output_height = math.floor((height - latent_kernel[1]) / latent_stride[1]) + 1 - output_width = math.floor((width - latent_kernel[2]) / latent_stride[2]) + 1 - - count = 0 - decoded_videos = z.new_zeros( - ( - output_num_frames * output_height * output_width, - self.config.out_channels, - self.kernel[0], - self.kernel[1], - self.kernel[2], - ) - ) - vae_batch_input = z.new_zeros( - (local_batch_size, num_channels, latent_kernel[0], latent_kernel[1], latent_kernel[2]) - ) - - for i in range(output_num_frames): - for j in range(output_height): - for k in range(output_width): - n_start, n_end = i * latent_stride[0], i * latent_stride[0] + latent_kernel[0] - h_start, h_end = j * latent_stride[1], j * latent_stride[1] + latent_kernel[1] - w_start, w_end = k * latent_stride[2], k * latent_stride[2] + latent_kernel[2] - - current_latent = z[:, :, n_start:n_end, h_start:h_end, w_start:w_end] - vae_batch_input[count % local_batch_size] = current_latent - - if ( - count % local_batch_size == local_batch_size - 1 - or count == output_num_frames * output_height * output_width - 1 - ): - current_video = self.decoder(vae_batch_input) - - if ( - count == output_num_frames * output_height * output_width - 1 - and count % local_batch_size != local_batch_size - 1 - ): - decoded_videos[count - count % local_batch_size :] = current_video[ - : count % local_batch_size + 1 - ] - else: - decoded_videos[count - local_batch_size + 1 : count + 1] = current_video - - vae_batch_input = z.new_zeros( - (local_batch_size, num_channels, latent_kernel[0], latent_kernel[1], latent_kernel[2]) - ) - - count += 1 - - video = z.new_zeros((batch_size, self.config.out_channels, num_frames * rt, height * rs, width * rs)) - video_overlap = ( - self.kernel[0] - self.stride[0], - self.kernel[1] - self.stride[1], - self.kernel[2] - self.stride[2], - ) - - for i in range(output_num_frames): - n_start, n_end = i * self.stride[0], i * self.stride[0] + self.kernel[0] - for j in range(output_height): - h_start, h_end = j * self.stride[1], j * self.stride[1] + self.kernel[1] - for k in range(output_width): - w_start, w_end = k * self.stride[2], k * self.stride[2] + self.kernel[2] - out_video_blend = _prepare_for_blend( - (i, output_num_frames, video_overlap[0]), - (j, output_height, video_overlap[1]), - (k, output_width, video_overlap[2]), - decoded_videos[i * output_height * output_width + j * output_width + k].unsqueeze(0), - ) - video[:, :, n_start:n_end, h_start:h_end, w_start:w_end] += out_video_blend - - video = video.permute(0, 2, 1, 3, 4).contiguous() - return video - - def forward( - self, - sample: torch.Tensor, - sample_posterior: bool = False, - return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: - r""" - Args: - sample (`torch.Tensor`): Input sample. - sample_posterior (`bool`, *optional*, defaults to `False`): - Whether to sample from the posterior. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`DecoderOutput`] instead of a plain tuple. - generator (`torch.Generator`, *optional*): - PyTorch random number generator. - """ - x = sample - posterior = self.encode(x).latent_dist - if sample_posterior: - z = posterior.sample(generator=generator) - else: - z = posterior.mode() - dec = self.decode(z).sample - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) - - -def _prepare_for_blend(n_param, h_param, w_param, x): - # TODO(aryan): refactor - n, n_max, overlap_n = n_param - h, h_max, overlap_h = h_param - w, w_max, overlap_w = w_param - if overlap_n > 0: - if n > 0: # the head overlap part decays from 0 to 1 - x[:, :, 0:overlap_n, :, :] = x[:, :, 0:overlap_n, :, :] * ( - torch.arange(0, overlap_n).float().to(x.device) / overlap_n - ).reshape(overlap_n, 1, 1) - if n < n_max - 1: # the tail overlap part decays from 1 to 0 - x[:, :, -overlap_n:, :, :] = x[:, :, -overlap_n:, :, :] * ( - 1 - torch.arange(0, overlap_n).float().to(x.device) / overlap_n - ).reshape(overlap_n, 1, 1) - if h > 0: - x[:, :, :, 0:overlap_h, :] = x[:, :, :, 0:overlap_h, :] * ( - torch.arange(0, overlap_h).float().to(x.device) / overlap_h - ).reshape(overlap_h, 1) - if h < h_max - 1: - x[:, :, :, -overlap_h:, :] = x[:, :, :, -overlap_h:, :] * ( - 1 - torch.arange(0, overlap_h).float().to(x.device) / overlap_h - ).reshape(overlap_h, 1) - if w > 0: - x[:, :, :, :, 0:overlap_w] = x[:, :, :, :, 0:overlap_w] * ( - torch.arange(0, overlap_w).float().to(x.device) / overlap_w - ) - if w < w_max - 1: - x[:, :, :, :, -overlap_w:] = x[:, :, :, :, -overlap_w:] * ( - 1 - torch.arange(0, overlap_w).float().to(x.device) / overlap_w - ) - return x diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py deleted file mode 100644 index e0e9436e8984fb9ea89b9c0042b6f0a866e4b6ef..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py +++ /dev/null @@ -1,1444 +0,0 @@ -# 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. - -from typing import Dict, Optional, Tuple, Union - -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.single_file_model import FromOriginalModelMixin -from ...utils import logging -from ...utils.accelerate_utils import apply_forward_hook -from ..activations import get_activation -from ..downsampling import CogVideoXDownsample3D -from ..modeling_outputs import AutoencoderKLOutput -from ..modeling_utils import ModelMixin -from ..upsampling import CogVideoXUpsample3D -from .vae import DecoderOutput, DiagonalGaussianDistribution - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class CogVideoXSafeConv3d(nn.Conv3d): - r""" - A 3D convolution layer that splits the input tensor into smaller parts to avoid OOM in CogVideoX Model. - """ - - def forward(self, input: torch.Tensor) -> torch.Tensor: - memory_count = ( - (input.shape[0] * input.shape[1] * input.shape[2] * input.shape[3] * input.shape[4]) * 2 / 1024**3 - ) - - # Set to 2GB, suitable for CuDNN - if memory_count > 2: - kernel_size = self.kernel_size[0] - part_num = int(memory_count / 2) + 1 - input_chunks = torch.chunk(input, part_num, dim=2) - - if kernel_size > 1: - input_chunks = [input_chunks[0]] + [ - torch.cat((input_chunks[i - 1][:, :, -kernel_size + 1 :], input_chunks[i]), dim=2) - for i in range(1, len(input_chunks)) - ] - - output_chunks = [] - for input_chunk in input_chunks: - output_chunks.append(super().forward(input_chunk)) - output = torch.cat(output_chunks, dim=2) - return output - else: - return super().forward(input) - - -class CogVideoXCausalConv3d(nn.Module): - r"""A 3D causal convolution layer that pads the input tensor to ensure causality in CogVideoX Model. - - 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. - 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. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: Union[int, Tuple[int, int, int]], - stride: int = 1, - dilation: int = 1, - pad_mode: str = "constant", - ): - super().__init__() - - if isinstance(kernel_size, int): - kernel_size = (kernel_size,) * 3 - - time_kernel_size, height_kernel_size, width_kernel_size = kernel_size - - # TODO(aryan): configure calculation based on stride and dilation in the future. - # Since CogVideoX does not use it, it is currently tailored to "just work" with Mochi - time_pad = time_kernel_size - 1 - height_pad = (height_kernel_size - 1) // 2 - width_pad = (width_kernel_size - 1) // 2 - - self.pad_mode = pad_mode - self.height_pad = height_pad - self.width_pad = width_pad - self.time_pad = time_pad - self.time_causal_padding = (width_pad, width_pad, height_pad, height_pad, time_pad, 0) - self.const_padding_conv3d = (0, self.width_pad, self.height_pad) - - self.temporal_dim = 2 - self.time_kernel_size = time_kernel_size - - stride = stride if isinstance(stride, tuple) else (stride, 1, 1) - dilation = (dilation, 1, 1) - self.conv = CogVideoXSafeConv3d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - stride=stride, - dilation=dilation, - padding=0 if self.pad_mode == "replicate" else self.const_padding_conv3d, - padding_mode="zeros", - ) - - def fake_context_parallel_forward( - self, inputs: torch.Tensor, conv_cache: Optional[torch.Tensor] = None - ) -> torch.Tensor: - if self.pad_mode == "replicate": - inputs = F.pad(inputs, self.time_causal_padding, mode="replicate") - else: - kernel_size = self.time_kernel_size - if kernel_size > 1: - cached_inputs = [conv_cache] if conv_cache is not None else [inputs[:, :, :1]] * (kernel_size - 1) - inputs = torch.cat(cached_inputs + [inputs], dim=2) - return inputs - - def forward(self, inputs: torch.Tensor, conv_cache: Optional[torch.Tensor] = None) -> torch.Tensor: - inputs = self.fake_context_parallel_forward(inputs, conv_cache) - - if self.pad_mode == "replicate": - conv_cache = None - else: - conv_cache = inputs[:, :, -self.time_kernel_size + 1 :].clone() - - output = self.conv(inputs) - return output, conv_cache - - -class CogVideoXSpatialNorm3D(nn.Module): - r""" - Spatially conditioned normalization as defined in https://huggingface.co/papers/2209.09002. This implementation is - specific to 3D-video like data. - - CogVideoXSafeConv3d is used instead of nn.Conv3d to avoid OOM in CogVideoX Model. - - Args: - f_channels (`int`): - The number of channels for input to group normalization layer, and output of the spatial norm layer. - zq_channels (`int`): - The number of channels for the quantized vector as described in the paper. - groups (`int`): - Number of groups to separate the channels into for group normalization. - """ - - def __init__( - self, - f_channels: int, - zq_channels: int, - groups: int = 32, - ): - super().__init__() - self.norm_layer = nn.GroupNorm(num_channels=f_channels, num_groups=groups, eps=1e-6, affine=True) - self.conv_y = CogVideoXCausalConv3d(zq_channels, f_channels, kernel_size=1, stride=1) - 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 - ) -> torch.Tensor: - new_conv_cache = {} - conv_cache = conv_cache or {} - - if f.shape[2] > 1 and f.shape[2] % 2 == 1: - f_first, f_rest = f[:, :, :1], f[:, :, 1:] - f_first_size, f_rest_size = f_first.shape[-3:], f_rest.shape[-3:] - z_first, z_rest = zq[:, :, :1], zq[:, :, 1:] - z_first = F.interpolate(z_first, size=f_first_size) - z_rest = F.interpolate(z_rest, size=f_rest_size) - zq = torch.cat([z_first, z_rest], dim=2) - else: - zq = F.interpolate(zq, size=f.shape[-3:]) - - conv_y, new_conv_cache["conv_y"] = self.conv_y(zq, conv_cache=conv_cache.get("conv_y")) - conv_b, new_conv_cache["conv_b"] = self.conv_b(zq, conv_cache=conv_cache.get("conv_b")) - - norm_f = self.norm_layer(f) - new_f = norm_f * conv_y + conv_b - return new_f, new_conv_cache - - -class CogVideoXResnetBlock3D(nn.Module): - r""" - A 3D ResNet block used in the CogVideoX 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. - temb_channels (`int`, defaults to `512`): - Number of time embedding channels. - groups (`int`, defaults to `32`): - Number of groups to separate the channels into for group normalization. - eps (`float`, defaults to `1e-6`): - Epsilon value for 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. - spatial_norm_dim (`int`, *optional*): - The dimension to use for spatial norm if it is to be used instead of group norm. - pad_mode (str, defaults to `"first"`): - Padding mode. - """ - - def __init__( - self, - in_channels: int, - out_channels: Optional[int] = 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, - pad_mode: str = "first", - ): - super().__init__() - - out_channels = out_channels or in_channels - - self.in_channels = in_channels - self.out_channels = out_channels - self.nonlinearity = get_activation(non_linearity) - self.use_conv_shortcut = conv_shortcut - self.spatial_norm_dim = spatial_norm_dim - - if spatial_norm_dim is None: - self.norm1 = nn.GroupNorm(num_channels=in_channels, num_groups=groups, eps=eps) - self.norm2 = nn.GroupNorm(num_channels=out_channels, num_groups=groups, eps=eps) - else: - self.norm1 = CogVideoXSpatialNorm3D( - f_channels=in_channels, - zq_channels=spatial_norm_dim, - groups=groups, - ) - self.norm2 = CogVideoXSpatialNorm3D( - f_channels=out_channels, - zq_channels=spatial_norm_dim, - groups=groups, - ) - - self.conv1 = CogVideoXCausalConv3d( - in_channels=in_channels, out_channels=out_channels, kernel_size=3, pad_mode=pad_mode - ) - - if temb_channels > 0: - self.temb_proj = nn.Linear(in_features=temb_channels, out_features=out_channels) - - self.dropout = nn.Dropout(dropout) - self.conv2 = CogVideoXCausalConv3d( - in_channels=out_channels, out_channels=out_channels, kernel_size=3, pad_mode=pad_mode - ) - - if self.in_channels != self.out_channels: - if self.use_conv_shortcut: - self.conv_shortcut = CogVideoXCausalConv3d( - in_channels=in_channels, out_channels=out_channels, kernel_size=3, pad_mode=pad_mode - ) - else: - self.conv_shortcut = CogVideoXSafeConv3d( - in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0 - ) - - def forward( - self, - inputs: torch.Tensor, - temb: Optional[torch.Tensor] = None, - zq: Optional[torch.Tensor] = None, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, - ) -> torch.Tensor: - new_conv_cache = {} - conv_cache = conv_cache or {} - - hidden_states = inputs - - if zq is not None: - hidden_states, new_conv_cache["norm1"] = self.norm1(hidden_states, zq, conv_cache=conv_cache.get("norm1")) - else: - hidden_states = self.norm1(hidden_states) - - hidden_states = self.nonlinearity(hidden_states) - hidden_states, new_conv_cache["conv1"] = self.conv1(hidden_states, conv_cache=conv_cache.get("conv1")) - - if temb is not None: - hidden_states = hidden_states + self.temb_proj(self.nonlinearity(temb))[:, :, None, None, None] - - if zq is not None: - hidden_states, new_conv_cache["norm2"] = self.norm2(hidden_states, zq, conv_cache=conv_cache.get("norm2")) - else: - hidden_states = self.norm2(hidden_states) - - hidden_states = self.nonlinearity(hidden_states) - hidden_states = self.dropout(hidden_states) - hidden_states, new_conv_cache["conv2"] = self.conv2(hidden_states, conv_cache=conv_cache.get("conv2")) - - if self.in_channels != self.out_channels: - if self.use_conv_shortcut: - inputs, new_conv_cache["conv_shortcut"] = self.conv_shortcut( - inputs, conv_cache=conv_cache.get("conv_shortcut") - ) - else: - inputs = self.conv_shortcut(inputs) - - hidden_states = hidden_states + inputs - return hidden_states, new_conv_cache - - -class CogVideoXDownBlock3D(nn.Module): - r""" - A downsampling block used in the CogVideoX model. - - Args: - in_channels (`int`): - Number of input channels. - out_channels (`int`, *optional*): - Number of output channels. If None, defaults to `in_channels`. - temb_channels (`int`, defaults to `512`): - Number of time embedding 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. - resnet_groups (`int`, defaults to `32`): - Number of groups to separate the channels into for group normalization. - add_downsample (`bool`, defaults to `True`): - Whether or not to use a downsampling layer. If not used, output dimension would be same as input dimension. - compress_time (`bool`, defaults to `False`): - Whether or not to downsample across temporal dimension. - pad_mode (str, defaults to `"first"`): - Padding mode. - """ - - _supports_gradient_checkpointing = True - - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - add_downsample: bool = True, - downsample_padding: int = 0, - compress_time: bool = False, - pad_mode: str = "first", - ): - super().__init__() - - resnets = [] - for i in range(num_layers): - in_channel = in_channels if i == 0 else out_channels - resnets.append( - CogVideoXResnetBlock3D( - in_channels=in_channel, - out_channels=out_channels, - dropout=dropout, - temb_channels=temb_channels, - groups=resnet_groups, - eps=resnet_eps, - non_linearity=resnet_act_fn, - pad_mode=pad_mode, - ) - ) - - self.resnets = nn.ModuleList(resnets) - self.downsamplers = None - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - CogVideoXDownsample3D( - out_channels, out_channels, padding=downsample_padding, compress_time=compress_time - ) - ] - ) - - self.gradient_checkpointing = False - - 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, - ) -> torch.Tensor: - r"""Forward method of the `CogVideoXDownBlock3D` class.""" - - new_conv_cache = {} - conv_cache = conv_cache or {} - - for i, resnet in enumerate(self.resnets): - conv_cache_key = f"resnet_{i}" - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, new_conv_cache[conv_cache_key] = self._gradient_checkpointing_func( - resnet, - hidden_states, - temb, - zq, - conv_cache.get(conv_cache_key), - ) - else: - hidden_states, new_conv_cache[conv_cache_key] = resnet( - hidden_states, temb, zq, conv_cache=conv_cache.get(conv_cache_key) - ) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - return hidden_states, new_conv_cache - - -class CogVideoXMidBlock3D(nn.Module): - r""" - A middle block used in the CogVideoX model. - - Args: - in_channels (`int`): - Number of input channels. - temb_channels (`int`, defaults to `512`): - Number of time embedding channels. - dropout (`float`, defaults to `0.0`): - Dropout rate. - num_layers (`int`, defaults to `1`): - Number of resnet layers. - resnet_eps (`float`, defaults to `1e-6`): - Epsilon value for normalization layers. - resnet_act_fn (`str`, defaults to `"swish"`): - Activation function to use. - resnet_groups (`int`, defaults to `32`): - Number of groups to separate the channels into for group normalization. - spatial_norm_dim (`int`, *optional*): - The dimension to use for spatial norm if it is to be used instead of group norm. - pad_mode (str, defaults to `"first"`): - Padding mode. - """ - - _supports_gradient_checkpointing = True - - def __init__( - self, - in_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - spatial_norm_dim: Optional[int] = None, - pad_mode: str = "first", - ): - super().__init__() - - resnets = [] - for _ in range(num_layers): - resnets.append( - CogVideoXResnetBlock3D( - in_channels=in_channels, - out_channels=in_channels, - dropout=dropout, - temb_channels=temb_channels, - groups=resnet_groups, - eps=resnet_eps, - spatial_norm_dim=spatial_norm_dim, - non_linearity=resnet_act_fn, - pad_mode=pad_mode, - ) - ) - self.resnets = nn.ModuleList(resnets) - - self.gradient_checkpointing = False - - 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, - ) -> torch.Tensor: - r"""Forward method of the `CogVideoXMidBlock3D` class.""" - - new_conv_cache = {} - conv_cache = conv_cache or {} - - for i, resnet in enumerate(self.resnets): - conv_cache_key = f"resnet_{i}" - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, new_conv_cache[conv_cache_key] = self._gradient_checkpointing_func( - resnet, hidden_states, temb, zq, conv_cache.get(conv_cache_key) - ) - else: - hidden_states, new_conv_cache[conv_cache_key] = resnet( - hidden_states, temb, zq, conv_cache=conv_cache.get(conv_cache_key) - ) - - return hidden_states, new_conv_cache - - -class CogVideoXUpBlock3D(nn.Module): - r""" - An upsampling block used in the CogVideoX model. - - Args: - in_channels (`int`): - Number of input channels. - out_channels (`int`, *optional*): - Number of output channels. If None, defaults to `in_channels`. - temb_channels (`int`, defaults to `512`): - Number of time embedding channels. - dropout (`float`, defaults to `0.0`): - Dropout rate. - num_layers (`int`, defaults to `1`): - Number of resnet layers. - resnet_eps (`float`, defaults to `1e-6`): - Epsilon value for normalization layers. - resnet_act_fn (`str`, defaults to `"swish"`): - Activation function to use. - resnet_groups (`int`, defaults to `32`): - Number of groups to separate the channels into for group normalization. - spatial_norm_dim (`int`, defaults to `16`): - The dimension to use for spatial norm if it is to be used instead of group norm. - add_upsample (`bool`, defaults to `True`): - Whether or not to use a upsampling layer. If not used, output dimension would be same as input dimension. - compress_time (`bool`, defaults to `False`): - Whether or not to downsample across temporal dimension. - pad_mode (str, defaults to `"first"`): - Padding mode. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - spatial_norm_dim: int = 16, - add_upsample: bool = True, - upsample_padding: int = 1, - compress_time: bool = False, - pad_mode: str = "first", - ): - super().__init__() - - resnets = [] - for i in range(num_layers): - in_channel = in_channels if i == 0 else out_channels - resnets.append( - CogVideoXResnetBlock3D( - in_channels=in_channel, - out_channels=out_channels, - dropout=dropout, - temb_channels=temb_channels, - groups=resnet_groups, - eps=resnet_eps, - non_linearity=resnet_act_fn, - spatial_norm_dim=spatial_norm_dim, - pad_mode=pad_mode, - ) - ) - - self.resnets = nn.ModuleList(resnets) - self.upsamplers = None - - if add_upsample: - self.upsamplers = nn.ModuleList( - [ - CogVideoXUpsample3D( - out_channels, out_channels, padding=upsample_padding, compress_time=compress_time - ) - ] - ) - - self.gradient_checkpointing = False - - 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, - ) -> torch.Tensor: - r"""Forward method of the `CogVideoXUpBlock3D` class.""" - - new_conv_cache = {} - conv_cache = conv_cache or {} - - for i, resnet in enumerate(self.resnets): - conv_cache_key = f"resnet_{i}" - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, new_conv_cache[conv_cache_key] = self._gradient_checkpointing_func( - resnet, - hidden_states, - temb, - zq, - conv_cache.get(conv_cache_key), - ) - else: - hidden_states, new_conv_cache[conv_cache_key] = resnet( - hidden_states, temb, zq, conv_cache=conv_cache.get(conv_cache_key) - ) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states) - - return hidden_states, new_conv_cache - - -class CogVideoXEncoder3D(nn.Module): - r""" - The `CogVideoXEncoder3D` layer of a variational autoencoder that encodes its input into a latent representation. - - Args: - in_channels (`int`, *optional*, defaults to 3): - 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",)`): - 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,)`): - 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. - 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 for normalization. - """ - - _supports_gradient_checkpointing = True - - def __init__( - self, - in_channels: int = 3, - out_channels: int = 16, - down_block_types: Tuple[str, ...] = ( - "CogVideoXDownBlock3D", - "CogVideoXDownBlock3D", - "CogVideoXDownBlock3D", - "CogVideoXDownBlock3D", - ), - block_out_channels: Tuple[int, ...] = (128, 256, 256, 512), - layers_per_block: int = 3, - act_fn: str = "silu", - norm_eps: float = 1e-6, - norm_num_groups: int = 32, - dropout: float = 0.0, - pad_mode: str = "first", - temporal_compression_ratio: float = 4, - ): - super().__init__() - - # log2 of temporal_compress_times - temporal_compress_level = int(np.log2(temporal_compression_ratio)) - - self.conv_in = CogVideoXCausalConv3d(in_channels, block_out_channels[0], kernel_size=3, pad_mode=pad_mode) - self.down_blocks = nn.ModuleList([]) - - # down blocks - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - compress_time = i < temporal_compress_level - - if down_block_type == "CogVideoXDownBlock3D": - down_block = CogVideoXDownBlock3D( - in_channels=input_channel, - out_channels=output_channel, - temb_channels=0, - dropout=dropout, - num_layers=layers_per_block, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - add_downsample=not is_final_block, - compress_time=compress_time, - ) - else: - raise ValueError("Invalid `down_block_type` encountered. Must be `CogVideoXDownBlock3D`") - - self.down_blocks.append(down_block) - - # mid block - self.mid_block = CogVideoXMidBlock3D( - in_channels=block_out_channels[-1], - temb_channels=0, - dropout=dropout, - num_layers=2, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - pad_mode=pad_mode, - ) - - self.norm_out = nn.GroupNorm(norm_num_groups, block_out_channels[-1], eps=1e-6) - self.conv_act = nn.SiLU() - self.conv_out = CogVideoXCausalConv3d( - block_out_channels[-1], 2 * out_channels, kernel_size=3, pad_mode=pad_mode - ) - - self.gradient_checkpointing = False - - def forward( - self, - sample: torch.Tensor, - temb: Optional[torch.Tensor] = None, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, - ) -> torch.Tensor: - r"""The forward method of the `CogVideoXEncoder3D` class.""" - - new_conv_cache = {} - conv_cache = conv_cache or {} - - hidden_states, new_conv_cache["conv_in"] = self.conv_in(sample, conv_cache=conv_cache.get("conv_in")) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - # 1. Down - for i, down_block in enumerate(self.down_blocks): - conv_cache_key = f"down_block_{i}" - hidden_states, new_conv_cache[conv_cache_key] = self._gradient_checkpointing_func( - down_block, - hidden_states, - temb, - None, - conv_cache.get(conv_cache_key), - ) - - # 2. Mid - hidden_states, new_conv_cache["mid_block"] = self._gradient_checkpointing_func( - self.mid_block, - hidden_states, - temb, - None, - conv_cache.get("mid_block"), - ) - else: - # 1. Down - for i, down_block in enumerate(self.down_blocks): - conv_cache_key = f"down_block_{i}" - hidden_states, new_conv_cache[conv_cache_key] = down_block( - hidden_states, temb, None, conv_cache.get(conv_cache_key) - ) - - # 2. Mid - hidden_states, new_conv_cache["mid_block"] = self.mid_block( - hidden_states, temb, None, conv_cache=conv_cache.get("mid_block") - ) - - # 3. Post-process - hidden_states = self.norm_out(hidden_states) - hidden_states = self.conv_act(hidden_states) - - hidden_states, new_conv_cache["conv_out"] = self.conv_out(hidden_states, conv_cache=conv_cache.get("conv_out")) - - return hidden_states, new_conv_cache - - -class CogVideoXDecoder3D(nn.Module): - r""" - The `CogVideoXDecoder3D` layer of a variational autoencoder that decodes its latent representation into an output - sample. - - Args: - in_channels (`int`, *optional*, defaults to 3): - 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",)`): - 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,)`): - 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. - 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 for normalization. - """ - - _supports_gradient_checkpointing = True - - def __init__( - self, - in_channels: int = 16, - out_channels: int = 3, - up_block_types: Tuple[str, ...] = ( - "CogVideoXUpBlock3D", - "CogVideoXUpBlock3D", - "CogVideoXUpBlock3D", - "CogVideoXUpBlock3D", - ), - block_out_channels: Tuple[int, ...] = (128, 256, 256, 512), - layers_per_block: int = 3, - act_fn: str = "silu", - norm_eps: float = 1e-6, - norm_num_groups: int = 32, - dropout: float = 0.0, - pad_mode: str = "first", - temporal_compression_ratio: float = 4, - ): - super().__init__() - - reversed_block_out_channels = list(reversed(block_out_channels)) - - self.conv_in = CogVideoXCausalConv3d( - in_channels, reversed_block_out_channels[0], kernel_size=3, pad_mode=pad_mode - ) - - # mid block - self.mid_block = CogVideoXMidBlock3D( - in_channels=reversed_block_out_channels[0], - temb_channels=0, - num_layers=2, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - spatial_norm_dim=in_channels, - pad_mode=pad_mode, - ) - - # up blocks - self.up_blocks = nn.ModuleList([]) - - output_channel = reversed_block_out_channels[0] - temporal_compress_level = int(np.log2(temporal_compression_ratio)) - - for i, up_block_type in enumerate(up_block_types): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - compress_time = i < temporal_compress_level - - if up_block_type == "CogVideoXUpBlock3D": - up_block = CogVideoXUpBlock3D( - in_channels=prev_output_channel, - out_channels=output_channel, - temb_channels=0, - dropout=dropout, - num_layers=layers_per_block + 1, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - spatial_norm_dim=in_channels, - add_upsample=not is_final_block, - compress_time=compress_time, - pad_mode=pad_mode, - ) - prev_output_channel = output_channel - else: - raise ValueError("Invalid `up_block_type` encountered. Must be `CogVideoXUpBlock3D`") - - self.up_blocks.append(up_block) - - self.norm_out = CogVideoXSpatialNorm3D(reversed_block_out_channels[-1], in_channels, groups=norm_num_groups) - self.conv_act = nn.SiLU() - self.conv_out = CogVideoXCausalConv3d( - reversed_block_out_channels[-1], out_channels, kernel_size=3, pad_mode=pad_mode - ) - - self.gradient_checkpointing = False - - def forward( - self, - sample: torch.Tensor, - temb: Optional[torch.Tensor] = None, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, - ) -> torch.Tensor: - r"""The forward method of the `CogVideoXDecoder3D` class.""" - - new_conv_cache = {} - conv_cache = conv_cache or {} - - hidden_states, new_conv_cache["conv_in"] = self.conv_in(sample, conv_cache=conv_cache.get("conv_in")) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - # 1. Mid - hidden_states, new_conv_cache["mid_block"] = self._gradient_checkpointing_func( - self.mid_block, - hidden_states, - temb, - sample, - conv_cache.get("mid_block"), - ) - - # 2. Up - for i, up_block in enumerate(self.up_blocks): - conv_cache_key = f"up_block_{i}" - hidden_states, new_conv_cache[conv_cache_key] = self._gradient_checkpointing_func( - up_block, - hidden_states, - temb, - sample, - conv_cache.get(conv_cache_key), - ) - else: - # 1. Mid - hidden_states, new_conv_cache["mid_block"] = self.mid_block( - hidden_states, temb, sample, conv_cache=conv_cache.get("mid_block") - ) - - # 2. Up - for i, up_block in enumerate(self.up_blocks): - conv_cache_key = f"up_block_{i}" - hidden_states, new_conv_cache[conv_cache_key] = up_block( - hidden_states, temb, sample, conv_cache=conv_cache.get(conv_cache_key) - ) - - # 3. Post-process - hidden_states, new_conv_cache["norm_out"] = self.norm_out( - hidden_states, sample, conv_cache=conv_cache.get("norm_out") - ) - hidden_states = self.conv_act(hidden_states) - hidden_states, new_conv_cache["conv_out"] = self.conv_out(hidden_states, conv_cache=conv_cache.get("conv_out")) - - return hidden_states, new_conv_cache - - -class AutoencoderKLCogVideoX(ModelMixin, ConfigMixin, FromOriginalModelMixin): - r""" - A VAE model with KL loss for encoding images into latents and decoding latent representations into images. Used in - [CogVideoX](https://github.com/THUDM/CogVideo). - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - 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. - 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`): - 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. - force_upcast (`bool`, *optional*, default to `True`): - If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE - can be fine-tuned / trained to a lower range without losing too much precision in which case `force_upcast` - can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["CogVideoXResnetBlock3D"] - - @register_to_config - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - down_block_types: Tuple[str] = ( - "CogVideoXDownBlock3D", - "CogVideoXDownBlock3D", - "CogVideoXDownBlock3D", - "CogVideoXDownBlock3D", - ), - up_block_types: Tuple[str] = ( - "CogVideoXUpBlock3D", - "CogVideoXUpBlock3D", - "CogVideoXUpBlock3D", - "CogVideoXUpBlock3D", - ), - block_out_channels: Tuple[int] = (128, 256, 256, 512), - latent_channels: int = 16, - layers_per_block: int = 3, - act_fn: str = "silu", - norm_eps: float = 1e-6, - norm_num_groups: int = 32, - temporal_compression_ratio: float = 4, - 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, - force_upcast: float = True, - use_quant_conv: bool = False, - use_post_quant_conv: bool = False, - invert_scale_latents: bool = False, - ): - super().__init__() - - self.encoder = CogVideoXEncoder3D( - in_channels=in_channels, - out_channels=latent_channels, - down_block_types=down_block_types, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - act_fn=act_fn, - norm_eps=norm_eps, - norm_num_groups=norm_num_groups, - temporal_compression_ratio=temporal_compression_ratio, - ) - self.decoder = CogVideoXDecoder3D( - in_channels=latent_channels, - out_channels=out_channels, - up_block_types=up_block_types, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - act_fn=act_fn, - norm_eps=norm_eps, - norm_num_groups=norm_num_groups, - temporal_compression_ratio=temporal_compression_ratio, - ) - self.quant_conv = CogVideoXSafeConv3d(2 * out_channels, 2 * out_channels, 1) if use_quant_conv else None - self.post_quant_conv = CogVideoXSafeConv3d(out_channels, out_channels, 1) if use_post_quant_conv else None - - self.use_slicing = False - self.use_tiling = False - - # Can be increased to decode more latent frames at once, but comes at a reasonable memory cost and it is not - # recommended because the temporal parts of the VAE, here, are tricky to understand. - # If you decode X latent frames together, the number of output frames is: - # (X + (2 conv cache) + (2 time upscale_1) + (4 time upscale_2) - (2 causal conv downscale)) => X + 6 frames - # - # Example with num_latent_frames_batch_size = 2: - # - 12 latent frames: (0, 1), (2, 3), (4, 5), (6, 7), (8, 9), (10, 11) are processed together - # => (12 // 2 frame slices) * ((2 num_latent_frames_batch_size) + (2 conv cache) + (2 time upscale_1) + (4 time upscale_2) - (2 causal conv downscale)) - # => 6 * 8 = 48 frames - # - 13 latent frames: (0, 1, 2) (special case), (3, 4), (5, 6), (7, 8), (9, 10), (11, 12) are processed together - # => (1 frame slice) * ((3 num_latent_frames_batch_size) + (2 conv cache) + (2 time upscale_1) + (4 time upscale_2) - (2 causal conv downscale)) + - # ((13 - 3) // 2) * ((2 num_latent_frames_batch_size) + (2 conv cache) + (2 time upscale_1) + (4 time upscale_2) - (2 causal conv downscale)) - # => 1 * 9 + 5 * 8 = 49 frames - # It has been implemented this way so as to not have "magic values" in the code base that would be hard to explain. Note that - # setting it to anything other than 2 would give poor results because the VAE hasn't been trained to be adaptive with different - # number of temporal frames. - self.num_latent_frames_batch_size = 2 - self.num_sample_frames_batch_size = 8 - - # We make the minimum height and width of sample for tiling half that of the generally supported - self.tile_sample_min_height = sample_height // 2 - self.tile_sample_min_width = sample_width // 2 - self.tile_latent_min_height = int( - self.tile_sample_min_height / (2 ** (len(self.config.block_out_channels) - 1)) - ) - self.tile_latent_min_width = int(self.tile_sample_min_width / (2 ** (len(self.config.block_out_channels) - 1))) - - # These are experimental overlap factors that were chosen based on experimentation and seem to work best for - # 720x480 (WxH) resolution. The above resolution is the strongly recommended generation resolution in CogVideoX - # and so the tiling implementation has only been tested on those specific resolutions. - self.tile_overlap_factor_height = 1 / 6 - self.tile_overlap_factor_width = 1 / 5 - - 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, - ) -> 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_overlap_factor_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. Must be between 0 and 1. Setting a higher - value might cause more tiles to be processed leading to slow down of the decoding process. - tile_overlap_factor_width (`int`, *optional*): - The minimum amount of overlap between two consecutive horizontal tiles. This is to ensure that there - are no tiling artifacts produced across the width dimension. Must be between 0 and 1. Setting a higher - value might cause more tiles to be processed leading to slow down of the decoding process. - """ - 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_latent_min_height = int( - self.tile_sample_min_height / (2 ** (len(self.config.block_out_channels) - 1)) - ) - self.tile_latent_min_width = int(self.tile_sample_min_width / (2 ** (len(self.config.block_out_channels) - 1))) - self.tile_overlap_factor_height = tile_overlap_factor_height or self.tile_overlap_factor_height - self.tile_overlap_factor_width = tile_overlap_factor_width or self.tile_overlap_factor_width - - 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: - batch_size, num_channels, num_frames, height, width = x.shape - - if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): - return self.tiled_encode(x) - - frame_batch_size = self.num_sample_frames_batch_size - # Note: We expect the number of frames to be either `1` or `frame_batch_size * k` or `frame_batch_size * k + 1` for some k. - # As the extra single frame is handled inside the loop, it is not required to round up here. - num_batches = max(num_frames // frame_batch_size, 1) - conv_cache = None - enc = [] - - for i in range(num_batches): - remaining_frames = num_frames % frame_batch_size - start_frame = frame_batch_size * i + (0 if i == 0 else remaining_frames) - end_frame = frame_batch_size * (i + 1) + remaining_frames - x_intermediate = x[:, :, start_frame:end_frame] - x_intermediate, conv_cache = self.encoder(x_intermediate, conv_cache=conv_cache) - if self.quant_conv is not None: - x_intermediate = self.quant_conv(x_intermediate) - enc.append(x_intermediate) - - enc = torch.cat(enc, dim=2) - return enc - - @apply_forward_hook - def encode( - self, x: torch.Tensor, return_dict: bool = True - ) -> Union[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) 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, return_dict: bool = True) -> Union[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): - return self.tiled_decode(z, return_dict=return_dict) - - frame_batch_size = self.num_latent_frames_batch_size - num_batches = max(num_frames // frame_batch_size, 1) - conv_cache = None - dec = [] - - for i in range(num_batches): - remaining_frames = num_frames % frame_batch_size - start_frame = frame_batch_size * i + (0 if i == 0 else remaining_frames) - end_frame = frame_batch_size * (i + 1) + remaining_frames - z_intermediate = z[:, :, start_frame:end_frame] - if self.post_quant_conv is not None: - z_intermediate = self.post_quant_conv(z_intermediate) - z_intermediate, conv_cache = self.decoder(z_intermediate, conv_cache=conv_cache) - dec.append(z_intermediate) - - dec = torch.cat(dec, dim=2) - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) - - @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[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: - decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] - decoded = torch.cat(decoded_slices) - else: - decoded = self._decode(z).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 tiled_encode(self, x: torch.Tensor) -> torch.Tensor: - 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 - steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is - different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the - tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the - output, but they should be much less noticeable. - - Args: - x (`torch.Tensor`): Input batch of videos. - - Returns: - `torch.Tensor`: - The latent representation of the encoded videos. - """ - # For a rough memory estimate, take a look at the `tiled_decode` method. - batch_size, num_channels, num_frames, height, width = x.shape - - overlap_height = int(self.tile_sample_min_height * (1 - self.tile_overlap_factor_height)) - overlap_width = int(self.tile_sample_min_width * (1 - self.tile_overlap_factor_width)) - blend_extent_height = int(self.tile_latent_min_height * self.tile_overlap_factor_height) - blend_extent_width = int(self.tile_latent_min_width * self.tile_overlap_factor_width) - row_limit_height = self.tile_latent_min_height - blend_extent_height - row_limit_width = self.tile_latent_min_width - blend_extent_width - frame_batch_size = self.num_sample_frames_batch_size - - # 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, overlap_height): - row = [] - for j in range(0, width, overlap_width): - # Note: We expect the number of frames to be either `1` or `frame_batch_size * k` or `frame_batch_size * k + 1` for some k. - # As the extra single frame is handled inside the loop, it is not required to round up here. - num_batches = max(num_frames // frame_batch_size, 1) - conv_cache = None - time = [] - - for k in range(num_batches): - remaining_frames = num_frames % frame_batch_size - start_frame = frame_batch_size * k + (0 if k == 0 else remaining_frames) - end_frame = frame_batch_size * (k + 1) + remaining_frames - tile = x[ - :, - :, - start_frame:end_frame, - i : i + self.tile_sample_min_height, - j : j + self.tile_sample_min_width, - ] - tile, conv_cache = self.encoder(tile, conv_cache=conv_cache) - if self.quant_conv is not None: - tile = self.quant_conv(tile) - time.append(tile) - - row.append(torch.cat(time, dim=2)) - 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_extent_height) - if j > 0: - tile = self.blend_h(row[j - 1], tile, blend_extent_width) - result_row.append(tile[:, :, :, :row_limit_height, :row_limit_width]) - result_rows.append(torch.cat(result_row, dim=4)) - - enc = torch.cat(result_rows, dim=3) - return enc - - def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[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. - """ - # Rough memory assessment: - # - In CogVideoX-2B, there are a total of 24 CausalConv3d layers. - # - The biggest intermediate dimensions are: [1, 128, 9, 480, 720]. - # - Assume fp16 (2 bytes per value). - # Memory required: 1 * 128 * 9 * 480 * 720 * 24 * 2 / 1024**3 = 17.8 GB - # - # Memory assessment when using tiling: - # - Assume everything as above but now HxW is 240x360 by tiling in half - # Memory required: 1 * 128 * 9 * 240 * 360 * 24 * 2 / 1024**3 = 4.5 GB - - batch_size, num_channels, num_frames, height, width = z.shape - - overlap_height = int(self.tile_latent_min_height * (1 - self.tile_overlap_factor_height)) - overlap_width = int(self.tile_latent_min_width * (1 - self.tile_overlap_factor_width)) - blend_extent_height = int(self.tile_sample_min_height * self.tile_overlap_factor_height) - blend_extent_width = int(self.tile_sample_min_width * self.tile_overlap_factor_width) - row_limit_height = self.tile_sample_min_height - blend_extent_height - row_limit_width = self.tile_sample_min_width - blend_extent_width - frame_batch_size = self.num_latent_frames_batch_size - - # 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, overlap_height): - row = [] - for j in range(0, width, overlap_width): - num_batches = max(num_frames // frame_batch_size, 1) - conv_cache = None - time = [] - - for k in range(num_batches): - remaining_frames = num_frames % frame_batch_size - start_frame = frame_batch_size * k + (0 if k == 0 else remaining_frames) - end_frame = frame_batch_size * (k + 1) + remaining_frames - tile = z[ - :, - :, - start_frame:end_frame, - i : i + self.tile_latent_min_height, - j : j + self.tile_latent_min_width, - ] - if self.post_quant_conv is not None: - tile = self.post_quant_conv(tile) - tile, conv_cache = self.decoder(tile, conv_cache=conv_cache) - time.append(tile) - - row.append(torch.cat(time, dim=2)) - 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_extent_height) - if j > 0: - tile = self.blend_h(row[j - 1], tile, blend_extent_width) - result_row.append(tile[:, :, :, :row_limit_height, :row_limit_width]) - result_rows.append(torch.cat(result_row, dim=4)) - - dec = torch.cat(result_rows, dim=3) - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) - - def forward( - self, - sample: torch.Tensor, - sample_posterior: bool = False, - return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[torch.Tensor, torch.Tensor]: - x = sample - posterior = self.encode(x).latent_dist - if sample_posterior: - z = posterior.sample(generator=generator) - else: - z = posterior.mode() - dec = self.decode(z).sample - if not return_dict: - return (dec,) - return DecoderOutput(sample=dec) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_cosmos.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_cosmos.py deleted file mode 100644 index 7ab79a0bb857bb8616377182067e2ab0b9f3d0b6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_cosmos.py +++ /dev/null @@ -1,1108 +0,0 @@ -# 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. - -import math -from typing import List, Optional, Tuple, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import get_logger -from ...utils.accelerate_utils import apply_forward_hook -from ..modeling_outputs import AutoencoderKLOutput -from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, IdentityDistribution - - -logger = get_logger(__name__) - - -# fmt: off -# These latents and means are from CV8x8x8-1.0. Each checkpoint has different values, but since this is the main VAE used, -# we will default to these values. -LATENTS_MEAN = [0.11362758, -0.0171717, 0.03071163, 0.02046862, 0.01931456, 0.02138567, 0.01999342, 0.02189187, 0.02011935, 0.01872694, 0.02168613, 0.02207148, 0.01986941, 0.01770413, 0.02067643, 0.02028245, 0.19125476, 0.04556972, 0.0595558, 0.05315534, 0.05496629, 0.05356264, 0.04856596, 0.05327453, 0.05410472, 0.05597149, 0.05524866, 0.05181874, 0.05071663, 0.05204537, 0.0564108, 0.05518042, 0.01306714, 0.03341161, 0.03847246, 0.02810185, 0.02790166, 0.02920026, 0.02823597, 0.02631033, 0.0278531, 0.02880507, 0.02977769, 0.03145441, 0.02888389, 0.03280773, 0.03484927, 0.03049198, -0.00197727, 0.07534957, 0.04963879, 0.05530893, 0.05410828, 0.05252541, 0.05029899, 0.05321025, 0.05149245, 0.0511921, 0.04643495, 0.04604527, 0.04631618, 0.04404101, 0.04403536, 0.04499495, -0.02994183, -0.04787003, -0.01064558, -0.01779824, -0.01490502, -0.02157517, -0.0204778, -0.02180816, -0.01945375, -0.02062863, -0.02192209, -0.02520639, -0.02246656, -0.02427533, -0.02683363, -0.02762006, 0.08019473, -0.13005368, -0.07568636, -0.06082374, -0.06036175, -0.05875364, -0.05921887, -0.05869788, -0.05273941, -0.052565, -0.05346428, -0.05456541, -0.053657, -0.05656897, -0.05728589, -0.05321847, 0.16718403, -0.00390146, 0.0379406, 0.0356561, 0.03554131, 0.03924074, 0.03873615, 0.04187329, 0.04226924, 0.04378717, 0.04684274, 0.05117614, 0.04547792, 0.05251586, 0.05048339, 0.04950784, 0.09564418, 0.0547128, 0.08183969, 0.07978633, 0.08076023, 0.08108605, 0.08011818, 0.07965573, 0.08187773, 0.08350263, 0.08101469, 0.0786941, 0.0774442, 0.07724521, 0.07830418, 0.07599796, -0.04987567, 0.05923908, -0.01058746, -0.01177603, -0.01116162, -0.01364149, -0.01546014, -0.0117213, -0.01780043, -0.01648314, -0.02100247, -0.02104417, -0.02482123, -0.02611689, -0.02561143, -0.02597336, -0.05364667, 0.08211684, 0.04686937, 0.04605641, 0.04304186, 0.0397355, 0.03686767, 0.04087112, 0.03704741, 0.03706401, 0.03120073, 0.03349091, 0.03319963, 0.03205781, 0.03195127, 0.03180481, 0.16427967, -0.11048453, -0.04595276, -0.04982893, -0.05213465, -0.04809378, -0.05080318, -0.04992863, -0.04493337, -0.0467619, -0.04884703, -0.04627892, -0.04913311, -0.04955709, -0.04533982, -0.04570218, -0.10612928, -0.05121198, -0.06761009, -0.07251801, -0.07265285, -0.07417855, -0.07202412, -0.07499027, -0.07625481, -0.07535747, -0.07638787, -0.07920305, -0.07596069, -0.07959418, -0.08265036, -0.07955471, -0.16888915, 0.0753242, 0.04062594, 0.03375093, 0.03337452, 0.03699376, 0.03651138, 0.03611023, 0.03555622, 0.03378554, 0.0300498, 0.03395559, 0.02941847, 0.03156432, 0.03431173, 0.03016853, -0.03415358, -0.01699573, -0.04029295, -0.04912157, -0.0498858, -0.04917918, -0.04918056, -0.0525189, -0.05325506, -0.05341973, -0.04983329, -0.04883146, -0.04985548, -0.04736718, -0.0462027, -0.04836091, 0.02055675, 0.03419799, -0.02907669, -0.04350509, -0.04156144, -0.04234421, -0.04446109, -0.04461774, -0.04882839, -0.04822346, -0.04502493, -0.0506244, -0.05146913, -0.04655267, -0.04862994, -0.04841615, 0.20312774, -0.07208502, -0.03635615, -0.03556088, -0.04246174, -0.04195838, -0.04293778, -0.04071276, -0.04240569, -0.04125213, -0.04395144, -0.03959096, -0.04044993, -0.04015875, -0.04088107, -0.03885176] -LATENTS_STD = [0.56700271, 0.65488982, 0.65589428, 0.66524369, 0.66619784, 0.6666382, 0.6720838, 0.66955978, 0.66928875, 0.67108786, 0.67092526, 0.67397463, 0.67894882, 0.67668313, 0.67769569, 0.67479557, 0.85245121, 0.8688373, 0.87348086, 0.88459337, 0.89135885, 0.8910504, 0.89714909, 0.89947474, 0.90201765, 0.90411824, 0.90692616, 0.90847772, 0.90648711, 0.91006982, 0.91033435, 0.90541548, 0.84960359, 0.85863352, 0.86895317, 0.88460612, 0.89245003, 0.89451706, 0.89931005, 0.90647358, 0.90338236, 0.90510076, 0.91008312, 0.90961218, 0.9123717, 0.91313171, 0.91435546, 0.91565102, 0.91877103, 0.85155135, 0.857804, 0.86998034, 0.87365264, 0.88161767, 0.88151032, 0.88758916, 0.89015514, 0.89245576, 0.89276224, 0.89450496, 0.90054202, 0.89994133, 0.90136105, 0.90114892, 0.77755755, 0.81456852, 0.81911844, 0.83137071, 0.83820474, 0.83890373, 0.84401101, 0.84425181, 0.84739357, 0.84798753, 0.85249585, 0.85114998, 0.85160935, 0.85626358, 0.85677862, 0.85641026, 0.69903517, 0.71697885, 0.71696913, 0.72583169, 0.72931731, 0.73254126, 0.73586977, 0.73734969, 0.73664582, 0.74084908, 0.74399322, 0.74471819, 0.74493188, 0.74824578, 0.75024873, 0.75274801, 0.8187142, 0.82251883, 0.82616025, 0.83164483, 0.84072375, 0.8396467, 0.84143305, 0.84880769, 0.8503468, 0.85196948, 0.85211051, 0.85386664, 0.85410017, 0.85439342, 0.85847849, 0.85385275, 0.67583984, 0.68259847, 0.69198853, 0.69928843, 0.70194328, 0.70467001, 0.70755547, 0.70917857, 0.71007699, 0.70963502, 0.71064079, 0.71027333, 0.71291167, 0.71537536, 0.71902508, 0.71604162, 0.72450989, 0.71979928, 0.72057378, 0.73035461, 0.73329622, 0.73660028, 0.73891461, 0.74279994, 0.74105692, 0.74002433, 0.74257588, 0.74416119, 0.74543899, 0.74694443, 0.74747062, 0.74586403, 0.90176988, 0.90990674, 0.91106802, 0.92163783, 0.92390233, 0.93056196, 0.93482202, 0.93642414, 0.93858379, 0.94064975, 0.94078934, 0.94325715, 0.94955301, 0.94814706, 0.95144123, 0.94923073, 0.49853548, 0.64968109, 0.6427654, 0.64966393, 0.6487664, 0.65203559, 0.6584242, 0.65351611, 0.65464371, 0.6574859, 0.65626335, 0.66123748, 0.66121179, 0.66077942, 0.66040152, 0.66474909, 0.61986589, 0.69138134, 0.6884557, 0.6955843, 0.69765401, 0.70015347, 0.70529598, 0.70468754, 0.70399523, 0.70479989, 0.70887572, 0.71126866, 0.7097227, 0.71249932, 0.71231949, 0.71175605, 0.35586974, 0.68723857, 0.68973219, 0.69958478, 0.6943453, 0.6995818, 0.70980215, 0.69899458, 0.70271689, 0.70095056, 0.69912851, 0.70522696, 0.70392174, 0.70916915, 0.70585734, 0.70373541, 0.98101336, 0.89024764, 0.89607251, 0.90678179, 0.91308665, 0.91812348, 0.91980827, 0.92480654, 0.92635667, 0.92887944, 0.93338072, 0.93468094, 0.93619436, 0.93906063, 0.94191772, 0.94471723, 0.83202779, 0.84106231, 0.84463632, 0.85829508, 0.86319661, 0.86751342, 0.86914337, 0.87085921, 0.87286359, 0.87537396, 0.87931138, 0.88054478, 0.8811838, 0.88872558, 0.88942474, 0.88934827, 0.44025335, 0.63061613, 0.63110614, 0.63601959, 0.6395812, 0.64104342, 0.65019929, 0.6502797, 0.64355946, 0.64657205, 0.64847094, 0.64728117, 0.64972943, 0.65162975, 0.65328044, 0.64914775] -_WAVELETS = { - "haar": torch.tensor([0.7071067811865476, 0.7071067811865476]), - "rearrange": torch.tensor([1.0, 1.0]), -} -# fmt: on - - -class CosmosCausalConv3d(nn.Conv3d): - 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), - padding: int = 1, - pad_mode: str = "constant", - ) -> None: - kernel_size = (kernel_size, kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size - dilation = (dilation, dilation, dilation) if isinstance(dilation, int) else dilation - stride = (stride, stride, stride) if isinstance(stride, int) else stride - - _, height_kernel_size, width_kernel_size = kernel_size - assert height_kernel_size % 2 == 1 and width_kernel_size % 2 == 1 - - super().__init__( - in_channels, - out_channels, - kernel_size, - stride=stride, - dilation=dilation, - ) - - self.pad_mode = pad_mode - self.temporal_pad = dilation[0] * (kernel_size[0] - 1) + (1 - stride[0]) - self.spatial_pad = (padding, padding, padding, padding) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states_prev = hidden_states[:, :, :1, ...].repeat(1, 1, self.temporal_pad, 1, 1) - hidden_states = torch.cat([hidden_states_prev, hidden_states], dim=2) - hidden_states = F.pad(hidden_states, (*self.spatial_pad, 0, 0), mode=self.pad_mode, value=0.0) - return super().forward(hidden_states) - - -class CosmosCausalGroupNorm(torch.nn.Module): - def __init__(self, in_channels: int, num_groups: int = 1): - super().__init__() - self.norm = nn.GroupNorm( - num_groups=num_groups, - num_channels=in_channels, - eps=1e-6, - affine=True, - ) - self.num_groups = num_groups - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.num_groups == 1: - batch_size = hidden_states.size(0) - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) # [B, C, T, H, W] -> [B * T, C, H, W] - hidden_states = self.norm(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute( - 0, 2, 1, 3, 4 - ) # [B * T, C, H, W] -> [B, C, T, H, W] - else: - hidden_states = self.norm(hidden_states) - return hidden_states - - -class CosmosPatchEmbed3d(nn.Module): - def __init__(self, patch_size: int = 1, patch_method: str = "haar") -> None: - super().__init__() - - self.patch_size = patch_size - self.patch_method = patch_method - - wavelets = _WAVELETS.get(patch_method).clone() - arange = torch.arange(wavelets.shape[0]) - - self.register_buffer("wavelets", wavelets, persistent=False) - self.register_buffer("_arange", arange, persistent=False) - - def _dwt(self, hidden_states: torch.Tensor, mode: str = "reflect", rescale=False) -> torch.Tensor: - dtype = hidden_states.dtype - wavelets = self.wavelets - - n = wavelets.shape[0] - g = hidden_states.shape[1] - hl = wavelets.flip(0).reshape(1, 1, -1).repeat(g, 1, 1) - hh = (wavelets * ((-1) ** self._arange)).reshape(1, 1, -1).repeat(g, 1, 1) - hh = hh.to(dtype=dtype) - hl = hl.to(dtype=dtype) - - # Handles temporal axis - hidden_states = F.pad(hidden_states, pad=(max(0, n - 2), n - 1, n - 2, n - 1, n - 2, n - 1), mode=mode).to( - dtype - ) - xl = F.conv3d(hidden_states, hl.unsqueeze(3).unsqueeze(4), groups=g, stride=(2, 1, 1)) - xh = F.conv3d(hidden_states, hh.unsqueeze(3).unsqueeze(4), groups=g, stride=(2, 1, 1)) - - # Handles spatial axes - xll = F.conv3d(xl, hl.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1)) - xlh = F.conv3d(xl, hh.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1)) - xhl = F.conv3d(xh, hl.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1)) - xhh = F.conv3d(xh, hh.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1)) - - xlll = F.conv3d(xll, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) - xllh = F.conv3d(xll, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) - xlhl = F.conv3d(xlh, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) - xlhh = F.conv3d(xlh, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) - xhll = F.conv3d(xhl, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) - xhlh = F.conv3d(xhl, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) - xhhl = F.conv3d(xhh, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) - xhhh = F.conv3d(xhh, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) - - hidden_states = torch.cat([xlll, xllh, xlhl, xlhh, xhll, xhlh, xhhl, xhhh], dim=1) - if rescale: - hidden_states = hidden_states / 8**0.5 - return hidden_states - - def _haar(self, hidden_states: torch.Tensor) -> torch.Tensor: - xi, xv = torch.split(hidden_states, [1, hidden_states.shape[2] - 1], dim=2) - hidden_states = torch.cat([xi.repeat_interleave(self.patch_size, dim=2), xv], dim=2) - for _ in range(int(math.log2(self.patch_size))): - hidden_states = self._dwt(hidden_states, rescale=True) - return hidden_states - - def _arrange(self, hidden_states: torch.Tensor) -> torch.Tensor: - xi, xv = torch.split(hidden_states, [1, hidden_states.shape[2] - 1], dim=2) - hidden_states = torch.cat([xi.repeat_interleave(self.patch_size, dim=2), xv], dim=2) - - batch_size, num_channels, num_frames, height, width = hidden_states.shape - p = self.patch_size - - hidden_states = torch.reshape(batch_size, num_channels, num_frames // p, p, height // p, p, width // p, p) - hidden_states = hidden_states.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(1, 4).contiguous() - return hidden_states - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.patch_method == "haar": - return self._haar(hidden_states) - elif self.patch_method == "rearrange": - return self._arrange(hidden_states) - else: - raise ValueError(f"Unsupported patch method: {self.patch_method}") - - -class CosmosUnpatcher3d(nn.Module): - def __init__(self, patch_size: int = 1, patch_method: str = "haar"): - super().__init__() - - self.patch_size = patch_size - self.patch_method = patch_method - - wavelets = _WAVELETS.get(patch_method).clone() - arange = torch.arange(wavelets.shape[0]) - - self.register_buffer("wavelets", wavelets, persistent=False) - self.register_buffer("_arange", arange, persistent=False) - - def _idwt(self, hidden_states: torch.Tensor, rescale: bool = False) -> torch.Tensor: - device = hidden_states.device - dtype = hidden_states.dtype - h = self.wavelets.to(device) - - g = hidden_states.shape[1] // 8 # split into 8 spatio-temporal filtered tesnors. - hl = h.flip([0]).reshape(1, 1, -1).repeat([g, 1, 1]) - hh = (h * ((-1) ** self._arange.to(device))).reshape(1, 1, -1).repeat(g, 1, 1) - hl = hl.to(dtype=dtype) - hh = hh.to(dtype=dtype) - - xlll, xllh, xlhl, xlhh, xhll, xhlh, xhhl, xhhh = torch.chunk(hidden_states, 8, dim=1) - - # Handle height transposed convolutions - xll = F.conv_transpose3d(xlll, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) - xll = F.conv_transpose3d(xllh, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) + xll - - xlh = F.conv_transpose3d(xlhl, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) - xlh = F.conv_transpose3d(xlhh, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) + xlh - - xhl = F.conv_transpose3d(xhll, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) - xhl = F.conv_transpose3d(xhlh, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) + xhl - - xhh = F.conv_transpose3d(xhhl, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) - xhh = F.conv_transpose3d(xhhh, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2)) + xhh - - # Handles width transposed convolutions - xl = F.conv_transpose3d(xll, hl.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1)) - xl = F.conv_transpose3d(xlh, hh.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1)) + xl - xh = F.conv_transpose3d(xhl, hl.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1)) - xh = F.conv_transpose3d(xhh, hh.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1)) + xh - - # Handles time axis transposed convolutions - hidden_states = F.conv_transpose3d(xl, hl.unsqueeze(3).unsqueeze(4), groups=g, stride=(2, 1, 1)) - hidden_states = ( - F.conv_transpose3d(xh, hh.unsqueeze(3).unsqueeze(4), groups=g, stride=(2, 1, 1)) + hidden_states - ) - - if rescale: - hidden_states = hidden_states * 8**0.5 - - return hidden_states - - def _ihaar(self, hidden_states: torch.Tensor) -> torch.Tensor: - for _ in range(int(math.log2(self.patch_size))): - hidden_states = self._idwt(hidden_states, rescale=True) - hidden_states = hidden_states[:, :, self.patch_size - 1 :, ...] - return hidden_states - - def _irearrange(self, hidden_states: torch.Tensor) -> torch.Tensor: - p = self.patch_size - hidden_states = hidden_states.unflatten(1, (-1, p, p, p)) - hidden_states = hidden_states.permute(0, 1, 5, 2, 6, 3, 7, 4) - hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - hidden_states = hidden_states[:, :, p - 1 :, ...] - return hidden_states - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.patch_method == "haar": - return self._ihaar(hidden_states) - elif self.patch_method == "rearrange": - return self._irearrange(hidden_states) - else: - raise ValueError("Unknown patch method: " + self.patch_method) - - -class CosmosConvProjection3d(nn.Module): - def __init__(self, in_channels: int, out_channels: int) -> None: - super().__init__() - - self.conv_s = CosmosCausalConv3d(in_channels, out_channels, kernel_size=(1, 3, 3), stride=1, padding=1) - self.conv_t = CosmosCausalConv3d(out_channels, out_channels, kernel_size=(3, 1, 1), stride=1, padding=0) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.conv_s(hidden_states) - hidden_states = self.conv_t(hidden_states) - return hidden_states - - -class CosmosResnetBlock3d(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - dropout: float = 0.0, - num_groups: int = 1, - ) -> None: - super().__init__() - out_channels = out_channels or in_channels - - self.norm1 = CosmosCausalGroupNorm(in_channels, num_groups) - self.conv1 = CosmosConvProjection3d(in_channels, out_channels) - - self.norm2 = CosmosCausalGroupNorm(out_channels, num_groups) - self.dropout = nn.Dropout(dropout) - self.conv2 = CosmosConvProjection3d(out_channels, out_channels) - - if in_channels != out_channels: - self.conv_shortcut = CosmosCausalConv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) - else: - self.conv_shortcut = nn.Identity() - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - residual = hidden_states - residual = self.conv_shortcut(residual) - - hidden_states = self.norm1(hidden_states) - hidden_states = F.silu(hidden_states) - hidden_states = self.conv1(hidden_states) - - hidden_states = self.norm2(hidden_states) - hidden_states = F.silu(hidden_states) - hidden_states = self.dropout(hidden_states) - hidden_states = self.conv2(hidden_states) - - return hidden_states + residual - - -class CosmosDownsample3d(nn.Module): - def __init__( - self, - in_channels: int, - spatial_downsample: bool = True, - temporal_downsample: bool = True, - ) -> None: - super().__init__() - - self.spatial_downsample = spatial_downsample - self.temporal_downsample = temporal_downsample - - self.conv1 = nn.Identity() - self.conv2 = nn.Identity() - self.conv3 = nn.Identity() - - if spatial_downsample: - self.conv1 = CosmosCausalConv3d( - in_channels, in_channels, kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=0 - ) - if temporal_downsample: - self.conv2 = CosmosCausalConv3d( - in_channels, in_channels, kernel_size=(3, 1, 1), stride=(2, 1, 1), padding=0 - ) - if spatial_downsample or temporal_downsample: - self.conv3 = CosmosCausalConv3d( - in_channels, in_channels, kernel_size=(1, 1, 1), stride=(1, 1, 1), padding=0 - ) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if not self.spatial_downsample and not self.temporal_downsample: - return hidden_states - - if self.spatial_downsample: - pad = (0, 1, 0, 1, 0, 0) - hidden_states = F.pad(hidden_states, pad, mode="constant", value=0) - conv_out = self.conv1(hidden_states) - pool_out = F.avg_pool3d(hidden_states, kernel_size=(1, 2, 2), stride=(1, 2, 2)) - hidden_states = conv_out + pool_out - - if self.temporal_downsample: - hidden_states = torch.cat([hidden_states[:, :, :1, ...], hidden_states], dim=2) - conv_out = self.conv2(hidden_states) - pool_out = F.avg_pool3d(hidden_states, kernel_size=(2, 1, 1), stride=(2, 1, 1)) - hidden_states = conv_out + pool_out - - hidden_states = self.conv3(hidden_states) - return hidden_states - - -class CosmosUpsample3d(nn.Module): - def __init__( - self, - in_channels: int, - spatial_upsample: bool = True, - temporal_upsample: bool = True, - ) -> None: - super().__init__() - - self.spatial_upsample = spatial_upsample - self.temporal_upsample = temporal_upsample - - self.conv1 = nn.Identity() - self.conv2 = nn.Identity() - self.conv3 = nn.Identity() - - if temporal_upsample: - self.conv1 = CosmosCausalConv3d( - in_channels, in_channels, kernel_size=(3, 1, 1), stride=(1, 1, 1), padding=0 - ) - if spatial_upsample: - self.conv2 = CosmosCausalConv3d( - in_channels, in_channels, kernel_size=(1, 3, 3), stride=(1, 1, 1), padding=1 - ) - if spatial_upsample or temporal_upsample: - self.conv3 = CosmosCausalConv3d( - in_channels, in_channels, kernel_size=(1, 1, 1), stride=(1, 1, 1), padding=0 - ) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if not self.spatial_upsample and not self.temporal_upsample: - return hidden_states - - if self.temporal_upsample: - num_frames = hidden_states.size(2) - time_factor = int(1.0 + 1.0 * (num_frames > 1)) - hidden_states = hidden_states.repeat_interleave(int(time_factor), dim=2) - hidden_states = hidden_states[..., time_factor - 1 :, :, :] - hidden_states = self.conv1(hidden_states) + hidden_states - - if self.spatial_upsample: - hidden_states = hidden_states.repeat_interleave(2, dim=3).repeat_interleave(2, dim=4) - hidden_states = self.conv2(hidden_states) + hidden_states - - hidden_states = self.conv3(hidden_states) - return hidden_states - - -class CosmosCausalAttention(nn.Module): - def __init__( - self, - num_attention_heads: int, - attention_head_dim: int, - num_groups: int = 1, - dropout: float = 0.0, - processor: Union["CosmosSpatialAttentionProcessor2_0", "CosmosTemporalAttentionProcessor2_0"] = None, - ) -> None: - super().__init__() - self.num_attention_heads = num_attention_heads - - self.norm = CosmosCausalGroupNorm(attention_head_dim, num_groups=num_groups) - self.to_q = CosmosCausalConv3d(attention_head_dim, attention_head_dim, kernel_size=1, stride=1, padding=0) - self.to_k = CosmosCausalConv3d(attention_head_dim, attention_head_dim, kernel_size=1, stride=1, padding=0) - self.to_v = CosmosCausalConv3d(attention_head_dim, attention_head_dim, kernel_size=1, stride=1, padding=0) - self.to_out = nn.ModuleList([]) - self.to_out.append( - CosmosCausalConv3d(attention_head_dim, attention_head_dim, kernel_size=1, stride=1, padding=0) - ) - self.to_out.append(nn.Dropout(dropout)) - - self.processor = processor - 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: - return self.processor(self, hidden_states=hidden_states, attention_mask=attention_mask) - - -class CosmosSpatialAttentionProcessor2_0: - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "CosmosSpatialAttentionProcessor2_0 requires PyTorch 2.0 or higher. To use it, please upgrade PyTorch." - ) - - def __call__( - self, attn: CosmosCausalAttention, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None - ) -> torch.Tensor: - batch_size, num_channels, num_frames, height, width = hidden_states.shape - residual = hidden_states - - hidden_states = attn.norm(hidden_states) - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - # [B, C, T, H, W] -> [B * T, H * W, C] - query = query.permute(0, 2, 3, 4, 1).flatten(2, 3).flatten(0, 1) - key = key.permute(0, 2, 3, 4, 1).flatten(2, 3).flatten(0, 1) - value = value.permute(0, 2, 3, 4, 1).flatten(2, 3).flatten(0, 1) - - # [B * T, H * W, C] -> [B * T, N, H * W, C // N] - query = query.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2) - key = key.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2) - value = value.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2) - - hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask) - hidden_states = hidden_states.transpose(1, 2).flatten(2, 3).type_as(query) - hidden_states = hidden_states.unflatten(1, (height, width)).unflatten(0, (batch_size, num_frames)) - hidden_states = hidden_states.permute(0, 4, 1, 2, 3) - - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - - return hidden_states + residual - - -class CosmosTemporalAttentionProcessor2_0: - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "CosmosSpatialAttentionProcessor2_0 requires PyTorch 2.0 or higher. To use it, please upgrade PyTorch." - ) - - def __call__( - self, attn: CosmosCausalAttention, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None - ) -> torch.Tensor: - batch_size, num_channels, num_frames, height, width = hidden_states.shape - residual = hidden_states - - hidden_states = attn.norm(hidden_states) - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - # [B, C, T, H, W] -> [B * T, H * W, C] - query = query.permute(0, 3, 4, 2, 1).flatten(0, 2) - key = key.permute(0, 3, 4, 2, 1).flatten(0, 2) - value = value.permute(0, 3, 4, 2, 1).flatten(0, 2) - - # [B * T, H * W, C] -> [B * T, N, H * W, C // N] - query = query.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2) - key = key.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2) - value = value.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2) - - hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask) - hidden_states = hidden_states.transpose(1, 2).flatten(2, 3).type_as(query) - hidden_states = hidden_states.unflatten(0, (batch_size, height, width)) - hidden_states = hidden_states.permute(0, 4, 3, 1, 2) - - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - - return hidden_states + residual - - -class CosmosDownBlock3d(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - num_layers: int, - dropout: float, - use_attention: bool, - use_downsample: bool, - spatial_downsample: bool, - temporal_downsample: bool, - ) -> None: - super().__init__() - - resnets, attentions, temp_attentions = [], [], [] - in_channel, out_channel = in_channels, out_channels - - for _ in range(num_layers): - resnets.append(CosmosResnetBlock3d(in_channel, out_channel, dropout, num_groups=1)) - in_channel = out_channel - - if use_attention: - attentions.append( - CosmosCausalAttention( - num_attention_heads=1, - attention_head_dim=out_channel, - num_groups=1, - dropout=dropout, - processor=CosmosSpatialAttentionProcessor2_0(), - ) - ) - temp_attentions.append( - CosmosCausalAttention( - num_attention_heads=1, - attention_head_dim=out_channel, - num_groups=1, - dropout=dropout, - processor=CosmosTemporalAttentionProcessor2_0(), - ) - ) - else: - attentions.append(None) - temp_attentions.append(None) - - self.resnets = nn.ModuleList(resnets) - self.attentions = nn.ModuleList(attentions) - self.temp_attentions = nn.ModuleList(temp_attentions) - - self.downsamplers = None - if use_downsample: - self.downsamplers = nn.ModuleList([]) - self.downsamplers.append(CosmosDownsample3d(out_channel, spatial_downsample, temporal_downsample)) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - for resnet, attention, temp_attention in zip(self.resnets, self.attentions, self.temp_attentions): - hidden_states = resnet(hidden_states) - if attention is not None: - hidden_states = attention(hidden_states) - if temp_attention is not None: - num_frames = hidden_states.size(2) - attention_mask = torch.tril(hidden_states.new_ones(num_frames, num_frames)).bool() - hidden_states = temp_attention(hidden_states, attention_mask) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - return hidden_states - - -class CosmosMidBlock3d(nn.Module): - def __init__(self, in_channels: int, num_layers: int, dropout: float, num_groups: int = 1) -> None: - super().__init__() - - resnets, attentions, temp_attentions = [], [], [] - - resnets.append(CosmosResnetBlock3d(in_channels, in_channels, dropout, num_groups)) - for _ in range(num_layers): - attentions.append( - CosmosCausalAttention( - num_attention_heads=1, - attention_head_dim=in_channels, - num_groups=num_groups, - dropout=dropout, - processor=CosmosSpatialAttentionProcessor2_0(), - ) - ) - temp_attentions.append( - CosmosCausalAttention( - num_attention_heads=1, - attention_head_dim=in_channels, - num_groups=num_groups, - dropout=dropout, - processor=CosmosTemporalAttentionProcessor2_0(), - ) - ) - resnets.append(CosmosResnetBlock3d(in_channels, in_channels, dropout, num_groups)) - - self.resnets = nn.ModuleList(resnets) - self.attentions = nn.ModuleList(attentions) - self.temp_attentions = nn.ModuleList(temp_attentions) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.resnets[0](hidden_states) - - for attention, temp_attention, resnet in zip(self.attentions, self.temp_attentions, self.resnets[1:]): - num_frames = hidden_states.size(2) - attention_mask = torch.tril(hidden_states.new_ones(num_frames, num_frames)).bool() - - hidden_states = attention(hidden_states) - hidden_states = temp_attention(hidden_states, attention_mask) - hidden_states = resnet(hidden_states) - - return hidden_states - - -class CosmosUpBlock3d(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - num_layers: int, - dropout: float, - use_attention: bool, - use_upsample: bool, - spatial_upsample: bool, - temporal_upsample: bool, - ) -> None: - super().__init__() - - resnets, attention, temp_attentions = [], [], [] - in_channel, out_channel = in_channels, out_channels - - for _ in range(num_layers): - resnets.append(CosmosResnetBlock3d(in_channel, out_channel, dropout, num_groups=1)) - in_channel = out_channel - - if use_attention: - attention.append( - CosmosCausalAttention( - num_attention_heads=1, - attention_head_dim=out_channel, - num_groups=1, - dropout=dropout, - processor=CosmosSpatialAttentionProcessor2_0(), - ) - ) - temp_attentions.append( - CosmosCausalAttention( - num_attention_heads=1, - attention_head_dim=out_channel, - num_groups=1, - dropout=dropout, - processor=CosmosTemporalAttentionProcessor2_0(), - ) - ) - else: - attention.append(None) - temp_attentions.append(None) - - self.resnets = nn.ModuleList(resnets) - self.attentions = nn.ModuleList(attention) - self.temp_attentions = nn.ModuleList(temp_attentions) - - self.upsamplers = None - if use_upsample: - self.upsamplers = nn.ModuleList([]) - self.upsamplers.append(CosmosUpsample3d(out_channel, spatial_upsample, temporal_upsample)) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - for resnet, attention, temp_attention in zip(self.resnets, self.attentions, self.temp_attentions): - hidden_states = resnet(hidden_states) - if attention is not None: - hidden_states = attention(hidden_states) - if temp_attention is not None: - num_frames = hidden_states.size(2) - attention_mask = torch.tril(hidden_states.new_ones(num_frames, num_frames)).bool() - hidden_states = temp_attention(hidden_states, attention_mask) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states) - - return hidden_states - - -class CosmosEncoder3d(nn.Module): - def __init__( - self, - in_channels: int = 3, - out_channels: int = 16, - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), - num_resnet_blocks: int = 2, - attention_resolutions: Tuple[int, ...] = (32,), - resolution: int = 1024, - patch_size: int = 4, - patch_type: str = "haar", - dropout: float = 0.0, - spatial_compression_ratio: int = 8, - temporal_compression_ratio: int = 8, - ) -> None: - super().__init__() - inner_dim = in_channels * patch_size**3 - num_spatial_layers = int(math.log2(spatial_compression_ratio)) - int(math.log2(patch_size)) - num_temporal_layers = int(math.log2(temporal_compression_ratio)) - int(math.log2(patch_size)) - - # 1. Input patching & projection - self.patch_embed = CosmosPatchEmbed3d(patch_size, patch_type) - - self.conv_in = CosmosConvProjection3d(inner_dim, block_out_channels[0]) - - # 2. Down blocks - current_resolution = resolution // patch_size - down_blocks = [] - for i in range(len(block_out_channels) - 1): - in_channel = block_out_channels[i] - out_channel = block_out_channels[i + 1] - - use_attention = current_resolution in attention_resolutions - spatial_downsample = temporal_downsample = False - if i < len(block_out_channels) - 2: - use_downsample = True - spatial_downsample = i < num_spatial_layers - temporal_downsample = i < num_temporal_layers - current_resolution = current_resolution // 2 - else: - use_downsample = False - - down_blocks.append( - CosmosDownBlock3d( - in_channel, - out_channel, - num_resnet_blocks, - dropout, - use_attention, - use_downsample, - spatial_downsample, - temporal_downsample, - ) - ) - self.down_blocks = nn.ModuleList(down_blocks) - - # 3. Mid block - self.mid_block = CosmosMidBlock3d(block_out_channels[-1], num_layers=1, dropout=dropout, num_groups=1) - - # 4. Output norm & projection - self.norm_out = CosmosCausalGroupNorm(block_out_channels[-1], num_groups=1) - self.conv_out = CosmosConvProjection3d(block_out_channels[-1], out_channels) - - self.gradient_checkpointing = False - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.patch_embed(hidden_states) - hidden_states = self.conv_in(hidden_states) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - for block in self.down_blocks: - hidden_states = self._gradient_checkpointing_func(block, hidden_states) - hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states) - else: - for block in self.down_blocks: - hidden_states = block(hidden_states) - hidden_states = self.mid_block(hidden_states) - - hidden_states = self.norm_out(hidden_states) - hidden_states = F.silu(hidden_states) - hidden_states = self.conv_out(hidden_states) - return hidden_states - - -class CosmosDecoder3d(nn.Module): - def __init__( - self, - in_channels: int = 16, - out_channels: int = 3, - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), - num_resnet_blocks: int = 2, - attention_resolutions: Tuple[int, ...] = (32,), - resolution: int = 1024, - patch_size: int = 4, - patch_type: str = "haar", - dropout: float = 0.0, - spatial_compression_ratio: int = 8, - temporal_compression_ratio: int = 8, - ) -> None: - super().__init__() - inner_dim = out_channels * patch_size**3 - num_spatial_layers = int(math.log2(spatial_compression_ratio)) - int(math.log2(patch_size)) - num_temporal_layers = int(math.log2(temporal_compression_ratio)) - int(math.log2(patch_size)) - reversed_block_out_channels = list(reversed(block_out_channels)) - - # 1. Input projection - self.conv_in = CosmosConvProjection3d(in_channels, reversed_block_out_channels[0]) - - # 2. Mid block - self.mid_block = CosmosMidBlock3d(reversed_block_out_channels[0], num_layers=1, dropout=dropout, num_groups=1) - - # 3. Up blocks - current_resolution = (resolution // patch_size) // 2 ** (len(block_out_channels) - 2) - up_blocks = [] - for i in range(len(block_out_channels) - 1): - in_channel = reversed_block_out_channels[i] - out_channel = reversed_block_out_channels[i + 1] - - use_attention = current_resolution in attention_resolutions - spatial_upsample = temporal_upsample = False - if i < len(block_out_channels) - 2: - use_upsample = True - temporal_upsample = 0 < i < num_temporal_layers + 1 - spatial_upsample = temporal_upsample or ( - i < num_spatial_layers and num_spatial_layers > num_temporal_layers - ) - current_resolution = current_resolution * 2 - else: - use_upsample = False - - up_blocks.append( - CosmosUpBlock3d( - in_channel, - out_channel, - num_resnet_blocks + 1, - dropout, - use_attention, - use_upsample, - spatial_upsample, - temporal_upsample, - ) - ) - self.up_blocks = nn.ModuleList(up_blocks) - - # 4. Output norm & projection & unpatching - self.norm_out = CosmosCausalGroupNorm(reversed_block_out_channels[-1], num_groups=1) - self.conv_out = CosmosConvProjection3d(reversed_block_out_channels[-1], inner_dim) - - self.unpatch_embed = CosmosUnpatcher3d(patch_size, patch_type) - - self.gradient_checkpointing = False - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.conv_in(hidden_states) - hidden_states = self.mid_block(hidden_states) - - for block in self.up_blocks: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(block, hidden_states) - else: - hidden_states = block(hidden_states) - - hidden_states = self.norm_out(hidden_states) - hidden_states = F.silu(hidden_states) - hidden_states = self.conv_out(hidden_states) - hidden_states = self.unpatch_embed(hidden_states) - return hidden_states - - -class AutoencoderKLCosmos(ModelMixin, ConfigMixin): - r""" - Autoencoder used in [Cosmos](https://huggingface.co/papers/2501.03575). - - 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 `16`): - Number of latent channels. - 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)`): - 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. - 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`): - Number of resnet blocks in each encoder/decoder block. - patch_size (`int`, defaults to `4`): - Patch size used for patching the input image/video. - patch_type (`str`, defaults to `haar`): - Patch type used for patching the input image/video. Can be either `haar` or `rearrange`. - scaling_factor (`float`, 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. Not applicable in - Cosmos, but we default to 1.0 for consistency. - spatial_compression_ratio (`int`, defaults to `8`): - The spatial compression ratio to apply in the VAE. The number of downsample blocks is determined using - this. - temporal_compression_ratio (`int`, defaults to `8`): - The temporal compression ratio to apply in the VAE. The number of downsample blocks is determined using - this. - """ - - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - 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,), - resolution: int = 1024, - num_layers: int = 2, - patch_size: int = 4, - patch_type: str = "haar", - 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, - ) -> None: - super().__init__() - - self.encoder = CosmosEncoder3d( - in_channels=in_channels, - out_channels=latent_channels, - block_out_channels=encoder_block_out_channels, - num_resnet_blocks=num_layers, - attention_resolutions=attention_resolutions, - resolution=resolution, - patch_size=patch_size, - patch_type=patch_type, - spatial_compression_ratio=spatial_compression_ratio, - temporal_compression_ratio=temporal_compression_ratio, - ) - self.decoder = CosmosDecoder3d( - in_channels=latent_channels, - out_channels=out_channels, - block_out_channels=decode_block_out_channels, - num_resnet_blocks=num_layers, - attention_resolutions=attention_resolutions, - resolution=resolution, - patch_size=patch_size, - patch_type=patch_type, - spatial_compression_ratio=spatial_compression_ratio, - temporal_compression_ratio=temporal_compression_ratio, - ) - - self.quant_conv = CosmosCausalConv3d(latent_channels, latent_channels, kernel_size=1, padding=0) - self.post_quant_conv = CosmosCausalConv3d(latent_channels, latent_channels, kernel_size=1, padding=0) - - # 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: 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, - ) -> 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 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: - x = self.encoder(x) - enc = self.quant_conv(x) - return enc - - @apply_forward_hook - def encode(self, x: torch.Tensor, return_dict: bool = True) -> torch.Tensor: - 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 = IdentityDistribution(h) - - if not return_dict: - return (posterior,) - return AutoencoderKLOutput(latent_dist=posterior) - - def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, Tuple[torch.Tensor]]: - z = self.post_quant_conv(z) - dec = self.decoder(z) - - if not return_dict: - return (dec,) - return DecoderOutput(sample=dec) - - @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[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) - else: - decoded = self._decode(z).sample - - 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: Optional[torch.Generator] = None, - ) -> Union[Tuple[torch.Tensor], DecoderOutput]: - x = sample - posterior = self.encode(x).latent_dist - if sample_posterior: - z = posterior.sample(generator=generator) - else: - z = posterior.mode() - dec = self.decode(z).sample - if not return_dict: - return (dec,) - return DecoderOutput(sample=dec) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py deleted file mode 100644 index 7b0f9889a52fc93463dba4831c3de3924945a95c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py +++ /dev/null @@ -1,1096 +0,0 @@ -# Copyright 2025 The Hunyuan 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 Optional, Tuple, Union - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch.utils.checkpoint - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import logging -from ...utils.accelerate_utils import apply_forward_hook -from ..activations import get_activation -from ..attention_processor import Attention -from ..modeling_outputs import AutoencoderKLOutput -from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, DiagonalGaussianDistribution - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -def prepare_causal_attention_mask( - num_frames: int, height_width: int, dtype: torch.dtype, device: torch.device, batch_size: int = None -) -> torch.Tensor: - indices = torch.arange(1, num_frames + 1, dtype=torch.int32, device=device) - indices_blocks = indices.repeat_interleave(height_width) - x, y = torch.meshgrid(indices_blocks, indices_blocks, indexing="xy") - mask = torch.where(x <= y, 0, -float("inf")).to(dtype=dtype) - - if batch_size is not None: - mask = mask.unsqueeze(0).expand(batch_size, -1, -1) - return mask - - -class HunyuanVideoCausalConv3d(nn.Module): - 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, - bias: bool = True, - pad_mode: str = "replicate", - ) -> None: - super().__init__() - - kernel_size = (kernel_size, kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size - - self.pad_mode = pad_mode - self.time_causal_padding = ( - kernel_size[0] // 2, - kernel_size[0] // 2, - kernel_size[1] // 2, - kernel_size[1] // 2, - kernel_size[2] - 1, - 0, - ) - - self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = F.pad(hidden_states, self.time_causal_padding, mode=self.pad_mode) - return self.conv(hidden_states) - - -class HunyuanVideoUpsampleCausal3D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: Optional[int] = None, - kernel_size: int = 3, - stride: int = 1, - bias: bool = True, - upsample_factor: Tuple[float, float, float] = (2, 2, 2), - ) -> None: - super().__init__() - - out_channels = out_channels or in_channels - self.upsample_factor = upsample_factor - - self.conv = HunyuanVideoCausalConv3d(in_channels, out_channels, kernel_size, stride, bias=bias) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - num_frames = hidden_states.size(2) - - first_frame, other_frames = hidden_states.split((1, num_frames - 1), dim=2) - first_frame = F.interpolate( - first_frame.squeeze(2), scale_factor=self.upsample_factor[1:], mode="nearest" - ).unsqueeze(2) - - if num_frames > 1: - # See: https://github.com/pytorch/pytorch/issues/81665 - # Unless you have a version of pytorch where non-contiguous implementation of F.interpolate - # is fixed, this will raise either a runtime error, or fail silently with bad outputs. - # If you are encountering an error here, make sure to try running encoding/decoding with - # `vae.enable_tiling()` first. If that doesn't work, open an issue at: - # https://github.com/huggingface/diffusers/issues - other_frames = other_frames.contiguous() - other_frames = F.interpolate(other_frames, scale_factor=self.upsample_factor, mode="nearest") - hidden_states = torch.cat((first_frame, other_frames), dim=2) - else: - hidden_states = first_frame - - hidden_states = self.conv(hidden_states) - return hidden_states - - -class HunyuanVideoDownsampleCausal3D(nn.Module): - def __init__( - self, - channels: int, - out_channels: Optional[int] = None, - padding: int = 1, - kernel_size: int = 3, - bias: bool = True, - stride=2, - ) -> None: - super().__init__() - out_channels = out_channels or channels - - self.conv = HunyuanVideoCausalConv3d(channels, out_channels, kernel_size, stride, padding, bias=bias) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.conv(hidden_states) - return hidden_states - - -class HunyuanVideoResnetBlockCausal3D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: Optional[int] = None, - dropout: float = 0.0, - groups: int = 32, - eps: float = 1e-6, - non_linearity: str = "swish", - ) -> None: - super().__init__() - out_channels = out_channels or in_channels - - self.nonlinearity = get_activation(non_linearity) - - self.norm1 = nn.GroupNorm(groups, in_channels, eps=eps, affine=True) - self.conv1 = HunyuanVideoCausalConv3d(in_channels, out_channels, 3, 1, 0) - - self.norm2 = nn.GroupNorm(groups, out_channels, eps=eps, affine=True) - self.dropout = nn.Dropout(dropout) - self.conv2 = HunyuanVideoCausalConv3d(out_channels, out_channels, 3, 1, 0) - - self.conv_shortcut = None - if in_channels != out_channels: - self.conv_shortcut = HunyuanVideoCausalConv3d(in_channels, out_channels, 1, 1, 0) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = hidden_states.contiguous() - residual = hidden_states - - hidden_states = self.norm1(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - hidden_states = self.conv1(hidden_states) - - hidden_states = self.norm2(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - hidden_states = self.dropout(hidden_states) - hidden_states = self.conv2(hidden_states) - - if self.conv_shortcut is not None: - residual = self.conv_shortcut(residual) - - hidden_states = hidden_states + residual - return hidden_states - - -class HunyuanVideoMidBlock3D(nn.Module): - def __init__( - self, - in_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - add_attention: bool = True, - attention_head_dim: int = 1, - ) -> None: - super().__init__() - resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) - self.add_attention = add_attention - - # There is always at least one resnet - resnets = [ - HunyuanVideoResnetBlockCausal3D( - in_channels=in_channels, - out_channels=in_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - non_linearity=resnet_act_fn, - ) - ] - attentions = [] - - for _ in range(num_layers): - if self.add_attention: - attentions.append( - Attention( - in_channels, - heads=in_channels // attention_head_dim, - dim_head=attention_head_dim, - eps=resnet_eps, - norm_num_groups=resnet_groups, - residual_connection=True, - bias=True, - upcast_softmax=True, - _from_deprecated_attn_block=True, - ) - ) - else: - attentions.append(None) - - resnets.append( - HunyuanVideoResnetBlockCausal3D( - in_channels=in_channels, - out_channels=in_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - non_linearity=resnet_act_fn, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - self.gradient_checkpointing = False - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(self.resnets[0], hidden_states) - - for attn, resnet in zip(self.attentions, self.resnets[1:]): - if attn is not None: - batch_size, num_channels, num_frames, height, width = hidden_states.shape - hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3) - attention_mask = prepare_causal_attention_mask( - num_frames, height * width, hidden_states.dtype, hidden_states.device, batch_size=batch_size - ) - hidden_states = attn(hidden_states, attention_mask=attention_mask) - hidden_states = hidden_states.unflatten(1, (num_frames, height, width)).permute(0, 4, 1, 2, 3) - - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) - - else: - hidden_states = self.resnets[0](hidden_states) - - for attn, resnet in zip(self.attentions, self.resnets[1:]): - if attn is not None: - batch_size, num_channels, num_frames, height, width = hidden_states.shape - hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3) - attention_mask = prepare_causal_attention_mask( - num_frames, height * width, hidden_states.dtype, hidden_states.device, batch_size=batch_size - ) - hidden_states = attn(hidden_states, attention_mask=attention_mask) - hidden_states = hidden_states.unflatten(1, (num_frames, height, width)).permute(0, 4, 1, 2, 3) - - hidden_states = resnet(hidden_states) - - return hidden_states - - -class HunyuanVideoDownBlock3D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - add_downsample: bool = True, - downsample_stride: int = 2, - downsample_padding: int = 1, - ) -> None: - super().__init__() - resnets = [] - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - HunyuanVideoResnetBlockCausal3D( - in_channels=in_channels, - out_channels=out_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - non_linearity=resnet_act_fn, - ) - ) - - self.resnets = nn.ModuleList(resnets) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - HunyuanVideoDownsampleCausal3D( - out_channels, - out_channels=out_channels, - padding=downsample_padding, - stride=downsample_stride, - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if torch.is_grad_enabled() and self.gradient_checkpointing: - for resnet in self.resnets: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) - else: - for resnet in self.resnets: - hidden_states = resnet(hidden_states) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - return hidden_states - - -class HunyuanVideoUpBlock3D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - add_upsample: bool = True, - upsample_scale_factor: Tuple[int, int, int] = (2, 2, 2), - ) -> None: - super().__init__() - resnets = [] - - for i in range(num_layers): - input_channels = in_channels if i == 0 else out_channels - - resnets.append( - HunyuanVideoResnetBlockCausal3D( - in_channels=input_channels, - out_channels=out_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - non_linearity=resnet_act_fn, - ) - ) - - self.resnets = nn.ModuleList(resnets) - - if add_upsample: - self.upsamplers = nn.ModuleList( - [ - HunyuanVideoUpsampleCausal3D( - out_channels, - out_channels=out_channels, - upsample_factor=upsample_scale_factor, - ) - ] - ) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if torch.is_grad_enabled() and self.gradient_checkpointing: - for resnet in self.resnets: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) - - else: - for resnet in self.resnets: - hidden_states = resnet(hidden_states) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states) - - return hidden_states - - -class HunyuanVideoEncoder3D(nn.Module): - r""" - Causal encoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603). - """ - - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - down_block_types: Tuple[str, ...] = ( - "HunyuanVideoDownBlock3D", - "HunyuanVideoDownBlock3D", - "HunyuanVideoDownBlock3D", - "HunyuanVideoDownBlock3D", - ), - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), - layers_per_block: int = 2, - norm_num_groups: int = 32, - act_fn: str = "silu", - double_z: bool = True, - mid_block_add_attention=True, - temporal_compression_ratio: int = 4, - spatial_compression_ratio: int = 8, - ) -> None: - super().__init__() - - self.conv_in = HunyuanVideoCausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1) - self.mid_block = None - self.down_blocks = nn.ModuleList([]) - - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - if down_block_type != "HunyuanVideoDownBlock3D": - raise ValueError(f"Unsupported down_block_type: {down_block_type}") - - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio)) - num_time_downsample_layers = int(np.log2(temporal_compression_ratio)) - - if temporal_compression_ratio == 4: - add_spatial_downsample = bool(i < num_spatial_downsample_layers) - add_time_downsample = bool( - i >= (len(block_out_channels) - 1 - num_time_downsample_layers) and not is_final_block - ) - elif temporal_compression_ratio == 8: - add_spatial_downsample = bool(i < num_spatial_downsample_layers) - add_time_downsample = bool(i < num_time_downsample_layers) - else: - raise ValueError(f"Unsupported time_compression_ratio: {temporal_compression_ratio}") - - downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1) - downsample_stride_T = (2,) if add_time_downsample else (1,) - downsample_stride = tuple(downsample_stride_T + downsample_stride_HW) - - down_block = HunyuanVideoDownBlock3D( - num_layers=layers_per_block, - in_channels=input_channel, - out_channels=output_channel, - add_downsample=bool(add_spatial_downsample or add_time_downsample), - resnet_eps=1e-6, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - downsample_stride=downsample_stride, - downsample_padding=0, - ) - - self.down_blocks.append(down_block) - - self.mid_block = HunyuanVideoMidBlock3D( - in_channels=block_out_channels[-1], - resnet_eps=1e-6, - resnet_act_fn=act_fn, - attention_head_dim=block_out_channels[-1], - resnet_groups=norm_num_groups, - add_attention=mid_block_add_attention, - ) - - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) - self.conv_act = nn.SiLU() - - conv_out_channels = 2 * out_channels if double_z else out_channels - self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3) - - self.gradient_checkpointing = False - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.conv_in(hidden_states) - - 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) - - hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states) - else: - for down_block in self.down_blocks: - hidden_states = down_block(hidden_states) - - hidden_states = self.mid_block(hidden_states) - - hidden_states = self.conv_norm_out(hidden_states) - hidden_states = self.conv_act(hidden_states) - hidden_states = self.conv_out(hidden_states) - - return hidden_states - - -class HunyuanVideoDecoder3D(nn.Module): - r""" - Causal decoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603). - """ - - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - up_block_types: Tuple[str, ...] = ( - "HunyuanVideoUpBlock3D", - "HunyuanVideoUpBlock3D", - "HunyuanVideoUpBlock3D", - "HunyuanVideoUpBlock3D", - ), - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), - layers_per_block: int = 2, - norm_num_groups: int = 32, - act_fn: str = "silu", - mid_block_add_attention=True, - time_compression_ratio: int = 4, - spatial_compression_ratio: int = 8, - ): - super().__init__() - self.layers_per_block = layers_per_block - - self.conv_in = HunyuanVideoCausalConv3d(in_channels, block_out_channels[-1], kernel_size=3, stride=1) - self.up_blocks = nn.ModuleList([]) - - # mid - self.mid_block = HunyuanVideoMidBlock3D( - in_channels=block_out_channels[-1], - resnet_eps=1e-6, - resnet_act_fn=act_fn, - attention_head_dim=block_out_channels[-1], - resnet_groups=norm_num_groups, - add_attention=mid_block_add_attention, - ) - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - if up_block_type != "HunyuanVideoUpBlock3D": - raise ValueError(f"Unsupported up_block_type: {up_block_type}") - - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio)) - num_time_upsample_layers = int(np.log2(time_compression_ratio)) - - if time_compression_ratio == 4: - add_spatial_upsample = bool(i < num_spatial_upsample_layers) - add_time_upsample = bool( - i >= len(block_out_channels) - 1 - num_time_upsample_layers and not is_final_block - ) - else: - raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}") - - upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1) - upsample_scale_factor_T = (2,) if add_time_upsample else (1,) - upsample_scale_factor = tuple(upsample_scale_factor_T + upsample_scale_factor_HW) - - up_block = HunyuanVideoUpBlock3D( - num_layers=self.layers_per_block + 1, - in_channels=prev_output_channel, - out_channels=output_channel, - add_upsample=bool(add_spatial_upsample or add_time_upsample), - upsample_scale_factor=upsample_scale_factor, - resnet_eps=1e-6, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - ) - - self.up_blocks.append(up_block) - prev_output_channel = output_channel - - # out - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) - self.conv_act = nn.SiLU() - self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[0], out_channels, kernel_size=3) - - self.gradient_checkpointing = False - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.conv_in(hidden_states) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states) - - for up_block in self.up_blocks: - hidden_states = self._gradient_checkpointing_func(up_block, hidden_states) - else: - hidden_states = self.mid_block(hidden_states) - - for up_block in self.up_blocks: - hidden_states = up_block(hidden_states) - - # post-process - hidden_states = self.conv_norm_out(hidden_states) - hidden_states = self.conv_act(hidden_states) - hidden_states = self.conv_out(hidden_states) - - return hidden_states - - -class AutoencoderKLHunyuanVideo(ModelMixin, ConfigMixin): - r""" - A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. - Introduced in [HunyuanVideo](https://huggingface.co/papers/2412.03603). - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - """ - - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - latent_channels: int = 16, - down_block_types: Tuple[str, ...] = ( - "HunyuanVideoDownBlock3D", - "HunyuanVideoDownBlock3D", - "HunyuanVideoDownBlock3D", - "HunyuanVideoDownBlock3D", - ), - up_block_types: Tuple[str, ...] = ( - "HunyuanVideoUpBlock3D", - "HunyuanVideoUpBlock3D", - "HunyuanVideoUpBlock3D", - "HunyuanVideoUpBlock3D", - ), - block_out_channels: Tuple[int] = (128, 256, 512, 512), - layers_per_block: int = 2, - act_fn: str = "silu", - norm_num_groups: int = 32, - scaling_factor: float = 0.476986, - spatial_compression_ratio: int = 8, - temporal_compression_ratio: int = 4, - mid_block_add_attention: bool = True, - ) -> None: - super().__init__() - - self.time_compression_ratio = temporal_compression_ratio - - self.encoder = HunyuanVideoEncoder3D( - in_channels=in_channels, - out_channels=latent_channels, - down_block_types=down_block_types, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - norm_num_groups=norm_num_groups, - act_fn=act_fn, - double_z=True, - mid_block_add_attention=mid_block_add_attention, - temporal_compression_ratio=temporal_compression_ratio, - spatial_compression_ratio=spatial_compression_ratio, - ) - - self.decoder = HunyuanVideoDecoder3D( - in_channels=latent_channels, - out_channels=out_channels, - up_block_types=up_block_types, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - norm_num_groups=norm_num_groups, - act_fn=act_fn, - time_compression_ratio=temporal_compression_ratio, - spatial_compression_ratio=spatial_compression_ratio, - mid_block_add_attention=mid_block_add_attention, - ) - - self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1) - self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1) - - self.spatial_compression_ratio = spatial_compression_ratio - self.temporal_compression_ratio = 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.tile_sample_min_num_frames`), the memory requirement can be lowered. - self.use_framewise_encoding = True - self.use_framewise_decoding = True - - # The minimal tile height and width for spatial tiling to be used - self.tile_sample_min_height = 256 - self.tile_sample_min_width = 256 - self.tile_sample_min_num_frames = 16 - - # The minimal distance between two spatial tiles - self.tile_sample_stride_height = 192 - self.tile_sample_stride_width = 192 - self.tile_sample_stride_num_frames = 12 - - 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, - ) -> 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_min_num_frames (`int`, *optional*): - The minimum number of frames required for a sample to be separated into tiles across the frame - 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. - tile_sample_stride_num_frames (`int`, *optional*): - The stride between two consecutive frame tiles. This is to ensure that there are no tiling artifacts - produced across the frame 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 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: - batch_size, num_channels, num_frames, height, width = x.shape - - if self.use_framewise_encoding and num_frames > self.tile_sample_min_num_frames: - return self._temporal_tiled_encode(x) - - if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): - return self.tiled_encode(x) - - x = self.encoder(x) - enc = self.quant_conv(x) - return enc - - @apply_forward_hook - def encode( - self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: - r""" - 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) 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, return_dict: bool = True) -> Union[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, 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, return_dict=return_dict) - - z = self.post_quant_conv(z) - dec = self.decoder(z) - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) - - @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: - r""" - 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: - decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] - decoded = torch.cat(decoded_slices) - else: - decoded = self._decode(z).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[-2], b.shape[-2], 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[-1], b.shape[-1], 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) -> AutoencoderKLOutput: - 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): - tile = x[:, :, :, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width] - tile = self.encoder(tile) - tile = self.quant_conv(tile) - row.append(tile) - 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, return_dict: bool = True) -> Union[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): - tile = z[:, :, :, i : i + tile_latent_min_height, j : j + tile_latent_min_width] - tile = self.post_quant_conv(tile) - decoded = self.decoder(tile) - row.append(decoded) - 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=-1)) - - 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) -> 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) - else: - tile = self.encoder(tile) - tile = self.quant_conv(tile) - 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, return_dict: bool = True) -> Union[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, return_dict=True).sample - else: - tile = self.post_quant_conv(tile) - decoded = self.decoder(tile) - 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) - result_row.append(tile[:, :, : self.tile_sample_stride_num_frames, :, :]) - 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, - sample_posterior: bool = False, - return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: - r""" - Args: - sample (`torch.Tensor`): Input sample. - sample_posterior (`bool`, *optional*, defaults to `False`): - Whether to sample from the posterior. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`DecoderOutput`] instead of a plain tuple. - """ - x = sample - posterior = self.encode(x).latent_dist - if sample_posterior: - z = posterior.sample(generator=generator) - else: - z = posterior.mode() - dec = self.decode(z, return_dict=return_dict) - return dec diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_ltx.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_ltx.py deleted file mode 100644 index 51c600a4e915bd9a2393b1dcea264c23b2a2bfb2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_ltx.py +++ /dev/null @@ -1,1557 +0,0 @@ -# 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 typing import Optional, Tuple, Union - -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 ..normalization import RMSNorm -from .vae import DecoderOutput, DiagonalGaussianDistribution - - -class LTXVideoCausalConv3d(nn.Module): - 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, - groups: int = 1, - padding_mode: str = "zeros", - is_causal: bool = True, - ): - super().__init__() - - self.in_channels = in_channels - self.out_channels = out_channels - self.is_causal = is_causal - 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=padding_mode, - ) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - time_kernel_size = self.kernel_size[0] - - if self.is_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 - - -class LTXVideoResnetBlock3d(nn.Module): - r""" - A 3D ResNet 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`. - 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: Optional[int] = None, - dropout: float = 0.0, - eps: float = 1e-6, - elementwise_affine: bool = False, - non_linearity: str = "swish", - is_causal: bool = True, - inject_noise: bool = False, - timestep_conditioning: bool = False, - ) -> None: - super().__init__() - - out_channels = out_channels or in_channels - - self.nonlinearity = get_activation(non_linearity) - - self.norm1 = RMSNorm(in_channels, eps=1e-8, elementwise_affine=elementwise_affine) - self.conv1 = LTXVideoCausalConv3d( - in_channels=in_channels, out_channels=out_channels, kernel_size=3, is_causal=is_causal - ) - - self.norm2 = RMSNorm(out_channels, eps=1e-8, elementwise_affine=elementwise_affine) - self.dropout = nn.Dropout(dropout) - self.conv2 = LTXVideoCausalConv3d( - in_channels=out_channels, out_channels=out_channels, kernel_size=3, is_causal=is_causal - ) - - 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) - self.conv_shortcut = LTXVideoCausalConv3d( - in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, is_causal=is_causal - ) - - 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: Optional[torch.Tensor] = None, generator: Optional[torch.Generator] = None - ) -> torch.Tensor: - hidden_states = inputs - - hidden_states = self.norm1(hidden_states.movedim(1, -1)).movedim(-1, 1) - - 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) - - 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.movedim(1, -1)).movedim(-1, 1) - - 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) - - 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 - - -class LTXVideoDownsampler3d(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - stride: Union[int, Tuple[int, int, int]] = 1, - is_causal: bool = True, - 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 = LTXVideoCausalConv3d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=3, - stride=1, - is_causal=is_causal, - padding_mode=padding_mode, - ) - - def forward(self, hidden_states: torch.Tensor) -> 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) - 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 - - -class LTXVideoUpsampler3d(nn.Module): - def __init__( - self, - in_channels: int, - stride: Union[int, Tuple[int, int, int]] = 1, - is_causal: bool = True, - residual: bool = False, - upscale_factor: int = 1, - 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 = LTXVideoCausalConv3d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=3, - stride=1, - is_causal=is_causal, - padding_mode=padding_mode, - ) - - def forward(self, hidden_states: torch.Tensor) -> 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) - 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 - - -class LTXVideoDownBlock3D(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: Optional[int] = None, - num_layers: int = 1, - dropout: float = 0.0, - resnet_eps: float = 1e-6, - resnet_act_fn: str = "swish", - spatio_temporal_scale: bool = True, - is_causal: bool = True, - ): - super().__init__() - - out_channels = out_channels or in_channels - - resnets = [] - for _ in range(num_layers): - resnets.append( - LTXVideoResnetBlock3d( - in_channels=in_channels, - out_channels=in_channels, - dropout=dropout, - eps=resnet_eps, - non_linearity=resnet_act_fn, - is_causal=is_causal, - ) - ) - self.resnets = nn.ModuleList(resnets) - - self.downsamplers = None - if spatio_temporal_scale: - self.downsamplers = nn.ModuleList( - [ - LTXVideoCausalConv3d( - in_channels=in_channels, - out_channels=in_channels, - kernel_size=3, - stride=(2, 2, 2), - is_causal=is_causal, - ) - ] - ) - - self.conv_out = None - if in_channels != out_channels: - self.conv_out = LTXVideoResnetBlock3d( - in_channels=in_channels, - out_channels=out_channels, - dropout=dropout, - eps=resnet_eps, - non_linearity=resnet_act_fn, - is_causal=is_causal, - ) - - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, - ) -> 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) - else: - hidden_states = resnet(hidden_states, temb, generator) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - if self.conv_out is not None: - hidden_states = self.conv_out(hidden_states, temb, generator) - - return hidden_states - - -class LTXVideo095DownBlock3D(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: Optional[int] = None, - num_layers: int = 1, - dropout: float = 0.0, - resnet_eps: float = 1e-6, - resnet_act_fn: str = "swish", - spatio_temporal_scale: bool = True, - is_causal: bool = True, - downsample_type: str = "conv", - ): - super().__init__() - - out_channels = out_channels or in_channels - - resnets = [] - for _ in range(num_layers): - resnets.append( - LTXVideoResnetBlock3d( - in_channels=in_channels, - out_channels=in_channels, - dropout=dropout, - eps=resnet_eps, - non_linearity=resnet_act_fn, - is_causal=is_causal, - ) - ) - self.resnets = nn.ModuleList(resnets) - - self.downsamplers = None - if spatio_temporal_scale: - self.downsamplers = nn.ModuleList() - - if downsample_type == "conv": - self.downsamplers.append( - LTXVideoCausalConv3d( - in_channels=in_channels, - out_channels=in_channels, - kernel_size=3, - stride=(2, 2, 2), - is_causal=is_causal, - ) - ) - elif downsample_type == "spatial": - self.downsamplers.append( - LTXVideoDownsampler3d( - in_channels=in_channels, out_channels=out_channels, stride=(1, 2, 2), is_causal=is_causal - ) - ) - elif downsample_type == "temporal": - self.downsamplers.append( - LTXVideoDownsampler3d( - in_channels=in_channels, out_channels=out_channels, stride=(2, 1, 1), is_causal=is_causal - ) - ) - elif downsample_type == "spatiotemporal": - self.downsamplers.append( - LTXVideoDownsampler3d( - in_channels=in_channels, out_channels=out_channels, stride=(2, 2, 2), is_causal=is_causal - ) - ) - - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, - ) -> 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) - else: - hidden_states = resnet(hidden_states, temb, generator) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - return hidden_states - - -# Adapted from diffusers.models.autoencoders.autoencoder_kl_cogvideox.CogVideoMidBlock3d -class LTXVideoMidBlock3d(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", - is_causal: bool = True, - inject_noise: bool = False, - timestep_conditioning: bool = False, - ) -> 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( - LTXVideoResnetBlock3d( - in_channels=in_channels, - out_channels=in_channels, - dropout=dropout, - eps=resnet_eps, - non_linearity=resnet_act_fn, - is_causal=is_causal, - inject_noise=inject_noise, - timestep_conditioning=timestep_conditioning, - ) - ) - self.resnets = nn.ModuleList(resnets) - - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, - ) -> 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) - else: - hidden_states = resnet(hidden_states, temb, generator) - - return hidden_states - - -class LTXVideoUpBlock3d(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: Optional[int] = None, - num_layers: int = 1, - dropout: float = 0.0, - resnet_eps: float = 1e-6, - resnet_act_fn: str = "swish", - spatio_temporal_scale: bool = True, - is_causal: bool = True, - inject_noise: bool = False, - timestep_conditioning: bool = False, - upsample_residual: bool = False, - upscale_factor: int = 1, - ): - 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 = LTXVideoResnetBlock3d( - in_channels=in_channels, - out_channels=out_channels, - dropout=dropout, - eps=resnet_eps, - non_linearity=resnet_act_fn, - is_causal=is_causal, - inject_noise=inject_noise, - timestep_conditioning=timestep_conditioning, - ) - - self.upsamplers = None - if spatio_temporal_scale: - self.upsamplers = nn.ModuleList( - [ - LTXVideoUpsampler3d( - out_channels * upscale_factor, - stride=(2, 2, 2), - is_causal=is_causal, - residual=upsample_residual, - upscale_factor=upscale_factor, - ) - ] - ) - - resnets = [] - for _ in range(num_layers): - resnets.append( - LTXVideoResnetBlock3d( - in_channels=out_channels, - out_channels=out_channels, - dropout=dropout, - eps=resnet_eps, - non_linearity=resnet_act_fn, - is_causal=is_causal, - inject_noise=inject_noise, - timestep_conditioning=timestep_conditioning, - ) - ) - self.resnets = nn.ModuleList(resnets) - - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, - ) -> torch.Tensor: - if self.conv_in is not None: - hidden_states = self.conv_in(hidden_states, temb, generator) - - 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) - - 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) - else: - hidden_states = resnet(hidden_states, temb, generator) - - return hidden_states - - -class LTXVideoEncoder3d(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 `(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 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 `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, ...] = (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"), - patch_size: int = 4, - patch_size_t: int = 1, - resnet_norm_eps: float = 1e-6, - is_causal: bool = True, - ): - super().__init__() - - self.patch_size = patch_size - self.patch_size_t = patch_size_t - self.in_channels = in_channels * patch_size**2 - - output_channel = block_out_channels[0] - - self.conv_in = LTXVideoCausalConv3d( - in_channels=self.in_channels, - out_channels=output_channel, - kernel_size=3, - stride=1, - is_causal=is_causal, - ) - - # down blocks - is_ltx_095 = down_block_types[-1] == "LTXVideo095DownBlock3D" - num_block_out_channels = len(block_out_channels) - (1 if is_ltx_095 else 0) - self.down_blocks = nn.ModuleList([]) - for i in range(num_block_out_channels): - input_channel = output_channel - if not is_ltx_095: - output_channel = block_out_channels[i + 1] if i + 1 < num_block_out_channels else block_out_channels[i] - else: - output_channel = block_out_channels[i + 1] - - if down_block_types[i] == "LTXVideoDownBlock3D": - down_block = LTXVideoDownBlock3D( - 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], - is_causal=is_causal, - ) - elif down_block_types[i] == "LTXVideo095DownBlock3D": - down_block = LTXVideo095DownBlock3D( - 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], - is_causal=is_causal, - downsample_type=downsample_type[i], - ) - else: - raise ValueError(f"Unknown down block type: {down_block_types[i]}") - - self.down_blocks.append(down_block) - - # mid block - self.mid_block = LTXVideoMidBlock3d( - in_channels=output_channel, - num_layers=layers_per_block[-1], - resnet_eps=resnet_norm_eps, - is_causal=is_causal, - ) - - # out - self.norm_out = RMSNorm(out_channels, eps=1e-8, elementwise_affine=False) - self.conv_act = nn.SiLU() - self.conv_out = LTXVideoCausalConv3d( - in_channels=output_channel, out_channels=out_channels + 1, kernel_size=3, stride=1, is_causal=is_causal - ) - - self.gradient_checkpointing = False - - def forward(self, hidden_states: torch.Tensor) -> 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 - - 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) - - 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) - - hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states) - else: - for down_block in self.down_blocks: - hidden_states = down_block(hidden_states) - - hidden_states = self.mid_block(hidden_states) - - hidden_states = self.norm_out(hidden_states.movedim(1, -1)).movedim(-1, 1) - hidden_states = self.conv_act(hidden_states) - hidden_states = self.conv_out(hidden_states) - - 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 - - -class LTXVideoDecoder3d(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, ...] = (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), - timestep_conditioning: bool = False, - upsample_residual: Tuple[bool, ...] = (False, False, False, False), - upsample_factor: Tuple[bool, ...] = (1, 1, 1, 1), - ) -> None: - super().__init__() - - self.patch_size = patch_size - self.patch_size_t = patch_size_t - self.out_channels = out_channels * patch_size**2 - - 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 = LTXVideoCausalConv3d( - in_channels=in_channels, out_channels=output_channel, kernel_size=3, stride=1, is_causal=is_causal - ) - - self.mid_block = LTXVideoMidBlock3d( - in_channels=output_channel, - num_layers=layers_per_block[0], - resnet_eps=resnet_norm_eps, - is_causal=is_causal, - inject_noise=inject_noise[0], - timestep_conditioning=timestep_conditioning, - ) - - # 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 = LTXVideoUpBlock3d( - 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], - is_causal=is_causal, - inject_noise=inject_noise[i + 1], - timestep_conditioning=timestep_conditioning, - upsample_residual=upsample_residual[i], - upscale_factor=upsample_factor[i], - ) - - self.up_blocks.append(up_block) - - # out - self.norm_out = RMSNorm(out_channels, eps=1e-8, elementwise_affine=False) - self.conv_act = nn.SiLU() - self.conv_out = LTXVideoCausalConv3d( - in_channels=output_channel, out_channels=self.out_channels, kernel_size=3, stride=1, is_causal=is_causal - ) - - # 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: Optional[torch.Tensor] = None) -> torch.Tensor: - hidden_states = self.conv_in(hidden_states) - - 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) - - for up_block in self.up_blocks: - hidden_states = self._gradient_checkpointing_func(up_block, hidden_states, temb) - else: - hidden_states = self.mid_block(hidden_states, temb) - - for up_block in self.up_blocks: - hidden_states = up_block(hidden_states, temb) - - hidden_states = self.norm_out(hidden_states.movedim(1, -1)).movedim(-1, 1) - - 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) - - 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 AutoencoderKLLTXVideo(ModelMixin, ConfigMixin, FromOriginalModelMixin): - r""" - A VAE model with KL loss for encoding images into latents and decoding latent representations into images. Used in - [LTX](https://huggingface.co/Lightricks/LTX-Video). - - 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, ...] = (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), - 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 = False, - spatial_compression_ratio: int = None, - temporal_compression_ratio: int = None, - ) -> None: - super().__init__() - - self.encoder = LTXVideoEncoder3d( - 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, - ) - self.decoder = LTXVideoDecoder3d( - 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, - ) - - 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: 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, - ) -> 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 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: - 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) - - if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): - return self.tiled_encode(x) - - enc = self.encoder(x) - - return enc - - @apply_forward_hook - def encode( - self, x: torch.Tensor, return_dict: bool = True - ) -> Union[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) 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, temb: Optional[torch.Tensor] = None, return_dict: bool = True - ) -> Union[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, 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, return_dict=return_dict) - - dec = self.decoder(z, temb) - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) - - @apply_forward_hook - def decode( - self, z: torch.Tensor, temb: Optional[torch.Tensor] = None, return_dict: bool = True - ) -> Union[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).sample for z_slice, t_slice in (z.split(1), temb.split(1)) - ] - else: - decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] - decoded = torch.cat(decoded_slices) - else: - decoded = self._decode(z, temb).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) -> 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] - ) - - 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: Optional[torch.Tensor], return_dict: bool = True - ) -> Union[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) - - 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) -> 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) - else: - tile = self.encoder(tile) - 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: Optional[torch.Tensor], return_dict: bool = True - ) -> Union[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, return_dict=True).sample - else: - decoded = self.decoder(tile, temb) - 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: Optional[torch.Tensor] = None, - sample_posterior: bool = False, - return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[torch.Tensor, torch.Tensor]: - x = sample - posterior = self.encode(x).latent_dist - if sample_posterior: - z = posterior.sample(generator=generator) - else: - z = posterior.mode() - dec = self.decode(z, temb) - if not return_dict: - return (dec.sample,) - return dec diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_magvit.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_magvit.py deleted file mode 100644 index 43294a901f02c57f38a88efebfddf2b817f6828e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_magvit.py +++ /dev/null @@ -1,1094 +0,0 @@ -# Copyright 2025 The EasyAnimate 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 typing import Optional, Tuple, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import logging -from ...utils.accelerate_utils import apply_forward_hook -from ..activations import get_activation -from ..modeling_outputs import AutoencoderKLOutput -from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, DiagonalGaussianDistribution - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class EasyAnimateCausalConv3d(nn.Conv3d): - 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, - groups: int = 1, - bias: bool = True, - padding_mode: str = "zeros", - ): - # Ensure kernel_size, stride, and dilation are tuples of length 3 - kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size,) * 3 - assert len(kernel_size) == 3, f"Kernel size must be a 3-tuple, got {kernel_size} instead." - - stride = stride if isinstance(stride, tuple) else (stride,) * 3 - assert len(stride) == 3, f"Stride must be a 3-tuple, got {stride} instead." - - dilation = dilation if isinstance(dilation, tuple) else (dilation,) * 3 - assert len(dilation) == 3, f"Dilation must be a 3-tuple, got {dilation} instead." - - # Unpack kernel size, stride, and dilation for temporal, height, and width dimensions - t_ks, h_ks, w_ks = kernel_size - self.t_stride, h_stride, w_stride = stride - t_dilation, h_dilation, w_dilation = dilation - - # Calculate padding for temporal dimension to maintain causality - t_pad = (t_ks - 1) * t_dilation - - # Calculate padding for height and width dimensions based on the padding parameter - if padding is None: - h_pad = math.ceil(((h_ks - 1) * h_dilation + (1 - h_stride)) / 2) - w_pad = math.ceil(((w_ks - 1) * w_dilation + (1 - w_stride)) / 2) - elif isinstance(padding, int): - h_pad = w_pad = padding - else: - assert NotImplementedError - - # Store temporal padding and initialize flags and previous features cache - self.temporal_padding = t_pad - self.temporal_padding_origin = math.ceil(((t_ks - 1) * w_dilation + (1 - w_stride)) / 2) - - self.prev_features = None - - # Initialize the parent class with modified padding - super().__init__( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - stride=stride, - dilation=dilation, - padding=(0, h_pad, w_pad), - groups=groups, - bias=bias, - padding_mode=padding_mode, - ) - - def _clear_conv_cache(self): - del self.prev_features - self.prev_features = None - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - # Ensure input tensor is of the correct type - dtype = hidden_states.dtype - if self.prev_features is None: - # Pad the input tensor in the temporal dimension to maintain causality - hidden_states = F.pad( - hidden_states, - pad=(0, 0, 0, 0, self.temporal_padding, 0), - mode="replicate", # TODO: check if this is necessary - ) - hidden_states = hidden_states.to(dtype=dtype) - - # Clear cache before processing and store previous features for causality - self._clear_conv_cache() - self.prev_features = hidden_states[:, :, -self.temporal_padding :].clone() - - # Process the input tensor in chunks along the temporal dimension - num_frames = hidden_states.size(2) - outputs = [] - i = 0 - while i + self.temporal_padding + 1 <= num_frames: - out = super().forward(hidden_states[:, :, i : i + self.temporal_padding + 1]) - i += self.t_stride - outputs.append(out) - return torch.concat(outputs, 2) - else: - # Concatenate previous features with the input tensor for continuous temporal processing - if self.t_stride == 2: - hidden_states = torch.concat( - [self.prev_features[:, :, -(self.temporal_padding - 1) :], hidden_states], dim=2 - ) - else: - hidden_states = torch.concat([self.prev_features, hidden_states], dim=2) - hidden_states = hidden_states.to(dtype=dtype) - - # Clear cache and update previous features - self._clear_conv_cache() - self.prev_features = hidden_states[:, :, -self.temporal_padding :].clone() - - # Process the concatenated tensor in chunks along the temporal dimension - num_frames = hidden_states.size(2) - outputs = [] - i = 0 - while i + self.temporal_padding + 1 <= num_frames: - out = super().forward(hidden_states[:, :, i : i + self.temporal_padding + 1]) - i += self.t_stride - outputs.append(out) - return torch.concat(outputs, 2) - - -class EasyAnimateResidualBlock3D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - non_linearity: str = "silu", - norm_num_groups: int = 32, - norm_eps: float = 1e-6, - spatial_group_norm: bool = True, - dropout: float = 0.0, - output_scale_factor: float = 1.0, - ): - super().__init__() - - self.output_scale_factor = output_scale_factor - - # Group normalization for input tensor - self.norm1 = nn.GroupNorm( - num_groups=norm_num_groups, - num_channels=in_channels, - eps=norm_eps, - affine=True, - ) - self.nonlinearity = get_activation(non_linearity) - self.conv1 = EasyAnimateCausalConv3d(in_channels, out_channels, kernel_size=3) - - self.norm2 = nn.GroupNorm(num_groups=norm_num_groups, num_channels=out_channels, eps=norm_eps, affine=True) - self.dropout = nn.Dropout(dropout) - self.conv2 = EasyAnimateCausalConv3d(out_channels, out_channels, kernel_size=3) - - if in_channels != out_channels: - self.shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1) - else: - self.shortcut = nn.Identity() - - self.spatial_group_norm = spatial_group_norm - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - shortcut = self.shortcut(hidden_states) - - if self.spatial_group_norm: - batch_size = hidden_states.size(0) - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) # [B, C, T, H, W] -> [B * T, C, H, W] - hidden_states = self.norm1(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute( - 0, 2, 1, 3, 4 - ) # [B * T, C, H, W] -> [B, C, T, H, W] - else: - hidden_states = self.norm1(hidden_states) - - hidden_states = self.nonlinearity(hidden_states) - hidden_states = self.conv1(hidden_states) - - if self.spatial_group_norm: - batch_size = hidden_states.size(0) - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) # [B, C, T, H, W] -> [B * T, C, H, W] - hidden_states = self.norm2(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute( - 0, 2, 1, 3, 4 - ) # [B * T, C, H, W] -> [B, C, T, H, W] - else: - hidden_states = self.norm2(hidden_states) - - hidden_states = self.nonlinearity(hidden_states) - hidden_states = self.dropout(hidden_states) - hidden_states = self.conv2(hidden_states) - - return (hidden_states + shortcut) / self.output_scale_factor - - -class EasyAnimateDownsampler3D(nn.Module): - def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3, stride: tuple = (2, 2, 2)): - super().__init__() - - self.conv = EasyAnimateCausalConv3d( - in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=0 - ) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = F.pad(hidden_states, (0, 1, 0, 1)) - hidden_states = self.conv(hidden_states) - return hidden_states - - -class EasyAnimateUpsampler3D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int = 3, - temporal_upsample: bool = False, - spatial_group_norm: bool = True, - ): - super().__init__() - out_channels = out_channels or in_channels - - self.temporal_upsample = temporal_upsample - self.spatial_group_norm = spatial_group_norm - - self.conv = EasyAnimateCausalConv3d( - in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size - ) - self.prev_features = None - - def _clear_conv_cache(self): - del self.prev_features - self.prev_features = None - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = F.interpolate(hidden_states, scale_factor=(1, 2, 2), mode="nearest") - hidden_states = self.conv(hidden_states) - - if self.temporal_upsample: - if self.prev_features is None: - self.prev_features = hidden_states - else: - hidden_states = F.interpolate( - hidden_states, - scale_factor=(2, 1, 1), - mode="trilinear" if not self.spatial_group_norm else "nearest", - ) - return hidden_states - - -class EasyAnimateDownBlock3D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - num_layers: int = 1, - act_fn: str = "silu", - norm_num_groups: int = 32, - norm_eps: float = 1e-6, - spatial_group_norm: bool = True, - dropout: float = 0.0, - output_scale_factor: float = 1.0, - add_downsample: bool = True, - add_temporal_downsample: bool = True, - ): - super().__init__() - - self.convs = nn.ModuleList([]) - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - self.convs.append( - EasyAnimateResidualBlock3D( - in_channels=in_channels, - out_channels=out_channels, - non_linearity=act_fn, - norm_num_groups=norm_num_groups, - norm_eps=norm_eps, - spatial_group_norm=spatial_group_norm, - dropout=dropout, - output_scale_factor=output_scale_factor, - ) - ) - - if add_downsample and add_temporal_downsample: - self.downsampler = EasyAnimateDownsampler3D(out_channels, out_channels, kernel_size=3, stride=(2, 2, 2)) - self.spatial_downsample_factor = 2 - self.temporal_downsample_factor = 2 - elif add_downsample and not add_temporal_downsample: - self.downsampler = EasyAnimateDownsampler3D(out_channels, out_channels, kernel_size=3, stride=(1, 2, 2)) - self.spatial_downsample_factor = 2 - self.temporal_downsample_factor = 1 - else: - self.downsampler = None - self.spatial_downsample_factor = 1 - self.temporal_downsample_factor = 1 - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - for conv in self.convs: - hidden_states = conv(hidden_states) - if self.downsampler is not None: - hidden_states = self.downsampler(hidden_states) - return hidden_states - - -class EasyAnimateUpBlock3d(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - num_layers: int = 1, - act_fn: str = "silu", - norm_num_groups: int = 32, - norm_eps: float = 1e-6, - spatial_group_norm: bool = False, - dropout: float = 0.0, - output_scale_factor: float = 1.0, - add_upsample: bool = True, - add_temporal_upsample: bool = True, - ): - super().__init__() - - self.convs = nn.ModuleList([]) - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - self.convs.append( - EasyAnimateResidualBlock3D( - in_channels=in_channels, - out_channels=out_channels, - non_linearity=act_fn, - norm_num_groups=norm_num_groups, - norm_eps=norm_eps, - spatial_group_norm=spatial_group_norm, - dropout=dropout, - output_scale_factor=output_scale_factor, - ) - ) - - if add_upsample: - self.upsampler = EasyAnimateUpsampler3D( - in_channels, - in_channels, - temporal_upsample=add_temporal_upsample, - spatial_group_norm=spatial_group_norm, - ) - else: - self.upsampler = None - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - for conv in self.convs: - hidden_states = conv(hidden_states) - if self.upsampler is not None: - hidden_states = self.upsampler(hidden_states) - return hidden_states - - -class EasyAnimateMidBlock3d(nn.Module): - def __init__( - self, - in_channels: int, - num_layers: int = 1, - act_fn: str = "silu", - norm_num_groups: int = 32, - norm_eps: float = 1e-6, - spatial_group_norm: bool = True, - dropout: float = 0.0, - output_scale_factor: float = 1.0, - ): - super().__init__() - - norm_num_groups = norm_num_groups if norm_num_groups is not None else min(in_channels // 4, 32) - - self.convs = nn.ModuleList( - [ - EasyAnimateResidualBlock3D( - in_channels=in_channels, - out_channels=in_channels, - non_linearity=act_fn, - norm_num_groups=norm_num_groups, - norm_eps=norm_eps, - spatial_group_norm=spatial_group_norm, - dropout=dropout, - output_scale_factor=output_scale_factor, - ) - ] - ) - - for _ in range(num_layers - 1): - self.convs.append( - EasyAnimateResidualBlock3D( - in_channels=in_channels, - out_channels=in_channels, - non_linearity=act_fn, - norm_num_groups=norm_num_groups, - norm_eps=norm_eps, - spatial_group_norm=spatial_group_norm, - dropout=dropout, - output_scale_factor=output_scale_factor, - ) - ) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.convs[0](hidden_states) - for resnet in self.convs[1:]: - hidden_states = resnet(hidden_states) - return hidden_states - - -class EasyAnimateEncoder(nn.Module): - r""" - Causal encoder for 3D video-like data used in [EasyAnimate](https://huggingface.co/papers/2405.18991). - """ - - _supports_gradient_checkpointing = True - - def __init__( - self, - in_channels: int = 3, - out_channels: int = 8, - down_block_types: Tuple[str, ...] = ( - "SpatialDownBlock3D", - "SpatialTemporalDownBlock3D", - "SpatialTemporalDownBlock3D", - "SpatialTemporalDownBlock3D", - ), - block_out_channels: Tuple[int, ...] = [128, 256, 512, 512], - layers_per_block: int = 2, - norm_num_groups: int = 32, - act_fn: str = "silu", - double_z: bool = True, - spatial_group_norm: bool = False, - ): - super().__init__() - - # 1. Input convolution - self.conv_in = EasyAnimateCausalConv3d(in_channels, block_out_channels[0], kernel_size=3) - - # 2. Down blocks - self.down_blocks = nn.ModuleList([]) - output_channels = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channels = output_channels - output_channels = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - if down_block_type == "SpatialDownBlock3D": - down_block = EasyAnimateDownBlock3D( - in_channels=input_channels, - out_channels=output_channels, - num_layers=layers_per_block, - act_fn=act_fn, - norm_num_groups=norm_num_groups, - norm_eps=1e-6, - spatial_group_norm=spatial_group_norm, - add_downsample=not is_final_block, - add_temporal_downsample=False, - ) - elif down_block_type == "SpatialTemporalDownBlock3D": - down_block = EasyAnimateDownBlock3D( - in_channels=input_channels, - out_channels=output_channels, - num_layers=layers_per_block, - act_fn=act_fn, - norm_num_groups=norm_num_groups, - norm_eps=1e-6, - spatial_group_norm=spatial_group_norm, - add_downsample=not is_final_block, - add_temporal_downsample=True, - ) - else: - raise ValueError(f"Unknown up block type: {down_block_type}") - self.down_blocks.append(down_block) - - # 3. Middle block - self.mid_block = EasyAnimateMidBlock3d( - in_channels=block_out_channels[-1], - num_layers=layers_per_block, - act_fn=act_fn, - spatial_group_norm=spatial_group_norm, - norm_num_groups=norm_num_groups, - norm_eps=1e-6, - dropout=0, - output_scale_factor=1, - ) - - # 4. Output normalization & convolution - self.spatial_group_norm = spatial_group_norm - self.conv_norm_out = nn.GroupNorm( - num_channels=block_out_channels[-1], - num_groups=norm_num_groups, - eps=1e-6, - ) - self.conv_act = get_activation(act_fn) - - # Initialize the output convolution layer - conv_out_channels = 2 * out_channels if double_z else out_channels - self.conv_out = EasyAnimateCausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3) - - self.gradient_checkpointing = False - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - # hidden_states: (B, C, T, H, W) - hidden_states = self.conv_in(hidden_states) - - for down_block in self.down_blocks: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(down_block, hidden_states) - else: - hidden_states = down_block(hidden_states) - - hidden_states = self.mid_block(hidden_states) - - if self.spatial_group_norm: - batch_size = hidden_states.size(0) - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) - hidden_states = self.conv_norm_out(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - else: - hidden_states = self.conv_norm_out(hidden_states) - - hidden_states = self.conv_act(hidden_states) - hidden_states = self.conv_out(hidden_states) - return hidden_states - - -class EasyAnimateDecoder(nn.Module): - r""" - Causal decoder for 3D video-like data used in [EasyAnimate](https://huggingface.co/papers/2405.18991). - """ - - _supports_gradient_checkpointing = True - - def __init__( - self, - in_channels: int = 8, - out_channels: int = 3, - up_block_types: Tuple[str, ...] = ( - "SpatialUpBlock3D", - "SpatialTemporalUpBlock3D", - "SpatialTemporalUpBlock3D", - "SpatialTemporalUpBlock3D", - ), - block_out_channels: Tuple[int, ...] = [128, 256, 512, 512], - layers_per_block: int = 2, - norm_num_groups: int = 32, - act_fn: str = "silu", - spatial_group_norm: bool = False, - ): - super().__init__() - - # 1. Input convolution - self.conv_in = EasyAnimateCausalConv3d(in_channels, block_out_channels[-1], kernel_size=3) - - # 2. Middle block - self.mid_block = EasyAnimateMidBlock3d( - in_channels=block_out_channels[-1], - num_layers=layers_per_block, - act_fn=act_fn, - norm_num_groups=norm_num_groups, - norm_eps=1e-6, - dropout=0, - output_scale_factor=1, - ) - - # 3. Up blocks - self.up_blocks = nn.ModuleList([]) - reversed_block_out_channels = list(reversed(block_out_channels)) - output_channels = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - input_channels = output_channels - output_channels = reversed_block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - # Create and append up block to up_blocks - if up_block_type == "SpatialUpBlock3D": - up_block = EasyAnimateUpBlock3d( - in_channels=input_channels, - out_channels=output_channels, - num_layers=layers_per_block + 1, - act_fn=act_fn, - norm_num_groups=norm_num_groups, - norm_eps=1e-6, - spatial_group_norm=spatial_group_norm, - add_upsample=not is_final_block, - add_temporal_upsample=False, - ) - elif up_block_type == "SpatialTemporalUpBlock3D": - up_block = EasyAnimateUpBlock3d( - in_channels=input_channels, - out_channels=output_channels, - num_layers=layers_per_block + 1, - act_fn=act_fn, - norm_num_groups=norm_num_groups, - norm_eps=1e-6, - spatial_group_norm=spatial_group_norm, - add_upsample=not is_final_block, - add_temporal_upsample=True, - ) - else: - raise ValueError(f"Unknown up block type: {up_block_type}") - self.up_blocks.append(up_block) - - # Output normalization and activation - self.spatial_group_norm = spatial_group_norm - self.conv_norm_out = nn.GroupNorm( - num_channels=block_out_channels[0], - num_groups=norm_num_groups, - eps=1e-6, - ) - self.conv_act = get_activation(act_fn) - - # Output convolution layer - self.conv_out = EasyAnimateCausalConv3d(block_out_channels[0], out_channels, kernel_size=3) - - self.gradient_checkpointing = False - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - # hidden_states: (B, C, T, H, W) - hidden_states = self.conv_in(hidden_states) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states) - else: - hidden_states = self.mid_block(hidden_states) - - for up_block in self.up_blocks: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(up_block, hidden_states) - else: - hidden_states = up_block(hidden_states) - - if self.spatial_group_norm: - batch_size = hidden_states.size(0) - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) # [B, C, T, H, W] -> [B * T, C, H, W] - hidden_states = self.conv_norm_out(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute( - 0, 2, 1, 3, 4 - ) # [B * T, C, H, W] -> [B, C, T, H, W] - else: - hidden_states = self.conv_norm_out(hidden_states) - - hidden_states = self.conv_act(hidden_states) - hidden_states = self.conv_out(hidden_states) - return hidden_states - - -class AutoencoderKLMagvit(ModelMixin, ConfigMixin): - r""" - A VAE model with KL loss for encoding images into latents and decoding latent representations into images. This - model is used in [EasyAnimate](https://huggingface.co/papers/2405.18991). - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - """ - - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - 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, ...] = [ - "SpatialDownBlock3D", - "SpatialTemporalDownBlock3D", - "SpatialTemporalDownBlock3D", - "SpatialTemporalDownBlock3D", - ], - up_block_types: Tuple[str, ...] = [ - "SpatialUpBlock3D", - "SpatialTemporalUpBlock3D", - "SpatialTemporalUpBlock3D", - "SpatialTemporalUpBlock3D", - ], - layers_per_block: int = 2, - act_fn: str = "silu", - norm_num_groups: int = 32, - scaling_factor: float = 0.7125, - spatial_group_norm: bool = True, - ): - super().__init__() - - # Initialize the encoder - self.encoder = EasyAnimateEncoder( - in_channels=in_channels, - out_channels=latent_channels, - down_block_types=down_block_types, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - norm_num_groups=norm_num_groups, - act_fn=act_fn, - double_z=True, - spatial_group_norm=spatial_group_norm, - ) - - # Initialize the decoder - self.decoder = EasyAnimateDecoder( - in_channels=latent_channels, - out_channels=out_channels, - up_block_types=up_block_types, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - norm_num_groups=norm_num_groups, - act_fn=act_fn, - spatial_group_norm=spatial_group_norm, - ) - - # Initialize convolution layers for quantization and post-quantization - self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1) - self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1) - - self.spatial_compression_ratio = 2 ** (len(block_out_channels) - 1) - self.temporal_compression_ratio = 2 ** (len(block_out_channels) - 2) - - # 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_size`), the memory requirement can be lowered. - self.use_framewise_encoding = False - self.use_framewise_decoding = False - - # Assign mini-batch sizes for encoder and decoder - self.num_sample_frames_batch_size = 4 - self.num_latent_frames_batch_size = 1 - - # 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 = 4 - - # 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 _clear_conv_cache(self): - # Clear cache for convolutional layers if needed - for name, module in self.named_modules(): - if isinstance(module, EasyAnimateCausalConv3d): - module._clear_conv_cache() - if isinstance(module, EasyAnimateUpsampler3D): - module._clear_conv_cache() - - 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, - ) -> 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.use_framewise_decoding = True - self.use_framewise_encoding = 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 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 - - @apply_forward_hook - def _encode( - self, x: torch.Tensor, return_dict: bool = True - ) -> Union[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 images. If `return_dict` is True, a - [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. - """ - if self.use_tiling and (x.shape[-1] > self.tile_sample_min_height or x.shape[-2] > self.tile_sample_min_width): - return self.tiled_encode(x, return_dict=return_dict) - - first_frames = self.encoder(x[:, :, :1, :, :]) - h = [first_frames] - for i in range(1, x.shape[2], self.num_sample_frames_batch_size): - next_frames = self.encoder(x[:, :, i : i + self.num_sample_frames_batch_size, :, :]) - h.append(next_frames) - h = torch.cat(h, dim=2) - moments = self.quant_conv(h) - - self._clear_conv_cache() - return moments - - @apply_forward_hook - def encode( - self, x: torch.Tensor, return_dict: bool = True - ) -> Union[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) 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, return_dict: bool = True) -> Union[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 - - if self.use_tiling and (z.shape[-1] > tile_latent_min_height or z.shape[-2] > tile_latent_min_width): - return self.tiled_decode(z, return_dict=return_dict) - - z = self.post_quant_conv(z) - - # Process the first frame and save the result - first_frames = self.decoder(z[:, :, :1, :, :]) - # Initialize the list to store the processed frames, starting with the first frame - dec = [first_frames] - # Process the remaining frames, with the number of frames processed at a time determined by mini_batch_decoder - for i in range(1, z.shape[2], self.num_latent_frames_batch_size): - next_frames = self.decoder(z[:, :, i : i + self.num_latent_frames_batch_size, :, :]) - dec.append(next_frames) - # Concatenate all processed frames along the channel dimension - dec = torch.cat(dec, dim=2) - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) - - @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[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: - decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] - decoded = torch.cat(decoded_slices) - else: - decoded = self._decode(z).sample - - self._clear_conv_cache() - 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 tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutput: - 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 the image into 512x512 tiles and encode them separately. - rows = [] - for i in range(0, height, self.tile_sample_stride_height): - row = [] - for j in range(0, width, self.tile_sample_stride_width): - tile = x[ - :, - :, - :, - i : i + self.tile_sample_min_height, - j : j + self.tile_sample_min_width, - ] - - first_frames = self.encoder(tile[:, :, 0:1, :, :]) - tile_h = [first_frames] - for k in range(1, num_frames, self.num_sample_frames_batch_size): - next_frames = self.encoder(tile[:, :, k : k + self.num_sample_frames_batch_size, :, :]) - tile_h.append(next_frames) - tile = torch.cat(tile_h, dim=2) - tile = self.quant_conv(tile) - self._clear_conv_cache() - row.append(tile) - 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[:, :, :, :latent_height, :latent_width]) - result_rows.append(torch.cat(result_row, dim=4)) - - 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]: - 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 64x64 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): - tile = z[ - :, - :, - :, - i : i + tile_latent_min_height, - j : j + tile_latent_min_width, - ] - tile = self.post_quant_conv(tile) - - # Process the first frame and save the result - first_frames = self.decoder(tile[:, :, :1, :, :]) - # Initialize the list to store the processed frames, starting with the first frame - tile_dec = [first_frames] - # Process the remaining frames, with the number of frames processed at a time determined by mini_batch_decoder - for k in range(1, num_frames, self.num_latent_frames_batch_size): - next_frames = self.decoder(tile[:, :, k : k + self.num_latent_frames_batch_size, :, :]) - tile_dec.append(next_frames) - # Concatenate all processed frames along the channel dimension - decoded = torch.cat(tile_dec, dim=2) - self._clear_conv_cache() - row.append(decoded) - 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 forward( - self, - sample: torch.Tensor, - sample_posterior: bool = False, - return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: - r""" - Args: - sample (`torch.Tensor`): Input sample. - sample_posterior (`bool`, *optional*, defaults to `False`): - Whether to sample from the posterior. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`DecoderOutput`] instead of a plain tuple. - """ - x = sample - posterior = self.encode(x).latent_dist - if sample_posterior: - z = posterior.sample(generator=generator) - else: - z = posterior.mode() - dec = self.decode(z).sample - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_mochi.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_mochi.py deleted file mode 100644 index 404d2f6d860ad69b8965e78153f6c6a12eb3b417..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_mochi.py +++ /dev/null @@ -1,1131 +0,0 @@ -# Copyright 2025 The Mochi 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 functools -from typing import Dict, Optional, Tuple, Union - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import logging -from ...utils.accelerate_utils import apply_forward_hook -from ..activations import get_activation -from ..attention_processor import Attention, MochiVaeAttnProcessor2_0 -from ..modeling_outputs import AutoencoderKLOutput -from ..modeling_utils import ModelMixin -from .autoencoder_kl_cogvideox import CogVideoXCausalConv3d -from .vae import DecoderOutput, DiagonalGaussianDistribution - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class MochiChunkedGroupNorm3D(nn.Module): - r""" - Applies per-frame group normalization for 5D video inputs. It also supports memory-efficient chunked group - normalization. - - Args: - num_channels (int): Number of channels expected in input - num_groups (int, optional): Number of groups to separate the channels into. Default: 32 - affine (bool, optional): If True, this module has learnable affine parameters. Default: True - chunk_size (int, optional): Size of each chunk for processing. Default: 8 - - """ - - def __init__( - self, - num_channels: int, - num_groups: int = 32, - affine: bool = True, - chunk_size: int = 8, - ): - super().__init__() - self.norm_layer = nn.GroupNorm(num_channels=num_channels, num_groups=num_groups, affine=affine) - self.chunk_size = chunk_size - - def forward(self, x: torch.Tensor = None) -> torch.Tensor: - batch_size = x.size(0) - - x = x.permute(0, 2, 1, 3, 4).flatten(0, 1) - output = torch.cat([self.norm_layer(chunk) for chunk in x.split(self.chunk_size, dim=0)], dim=0) - output = output.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - - return output - - -class MochiResnetBlock3D(nn.Module): - r""" - A 3D ResNet block used in the Mochi model. - - Args: - in_channels (`int`): - Number of input channels. - out_channels (`int`, *optional*): - Number of output channels. If None, defaults to `in_channels`. - non_linearity (`str`, defaults to `"swish"`): - Activation function to use. - """ - - def __init__( - self, - in_channels: int, - out_channels: Optional[int] = None, - act_fn: str = "swish", - ): - super().__init__() - - out_channels = out_channels or in_channels - - self.in_channels = in_channels - self.out_channels = out_channels - self.nonlinearity = get_activation(act_fn) - - self.norm1 = MochiChunkedGroupNorm3D(num_channels=in_channels) - self.conv1 = CogVideoXCausalConv3d( - in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, pad_mode="replicate" - ) - self.norm2 = MochiChunkedGroupNorm3D(num_channels=out_channels) - self.conv2 = CogVideoXCausalConv3d( - in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, pad_mode="replicate" - ) - - def forward( - self, - inputs: torch.Tensor, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, - ) -> torch.Tensor: - new_conv_cache = {} - conv_cache = conv_cache or {} - - hidden_states = inputs - - hidden_states = self.norm1(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - hidden_states, new_conv_cache["conv1"] = self.conv1(hidden_states, conv_cache=conv_cache.get("conv1")) - - hidden_states = self.norm2(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - hidden_states, new_conv_cache["conv2"] = self.conv2(hidden_states, conv_cache=conv_cache.get("conv2")) - - hidden_states = hidden_states + inputs - return hidden_states, new_conv_cache - - -class MochiDownBlock3D(nn.Module): - r""" - An downsampling block used in the Mochi 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 blocks in the block. - temporal_expansion (`int`, defaults to `2`): - Temporal expansion factor. - spatial_expansion (`int`, defaults to `2`): - Spatial expansion factor. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - num_layers: int = 1, - temporal_expansion: int = 2, - spatial_expansion: int = 2, - add_attention: bool = True, - ): - super().__init__() - self.temporal_expansion = temporal_expansion - self.spatial_expansion = spatial_expansion - - self.conv_in = CogVideoXCausalConv3d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=(temporal_expansion, spatial_expansion, spatial_expansion), - stride=(temporal_expansion, spatial_expansion, spatial_expansion), - pad_mode="replicate", - ) - - resnets = [] - norms = [] - attentions = [] - for _ in range(num_layers): - resnets.append(MochiResnetBlock3D(in_channels=out_channels)) - if add_attention: - norms.append(MochiChunkedGroupNorm3D(num_channels=out_channels)) - attentions.append( - Attention( - query_dim=out_channels, - heads=out_channels // 32, - dim_head=32, - qk_norm="l2", - is_causal=True, - processor=MochiVaeAttnProcessor2_0(), - ) - ) - else: - norms.append(None) - attentions.append(None) - - self.resnets = nn.ModuleList(resnets) - self.norms = nn.ModuleList(norms) - self.attentions = nn.ModuleList(attentions) - - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, - chunk_size: int = 2**15, - ) -> torch.Tensor: - r"""Forward method of the `MochiUpBlock3D` class.""" - - new_conv_cache = {} - conv_cache = conv_cache or {} - - hidden_states, new_conv_cache["conv_in"] = self.conv_in(hidden_states) - - for i, (resnet, norm, attn) in enumerate(zip(self.resnets, self.norms, self.attentions)): - conv_cache_key = f"resnet_{i}" - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, new_conv_cache[conv_cache_key] = self._gradient_checkpointing_func( - resnet, - hidden_states, - conv_cache.get(conv_cache_key), - ) - else: - hidden_states, new_conv_cache[conv_cache_key] = resnet( - hidden_states, conv_cache=conv_cache.get(conv_cache_key) - ) - - if attn is not None: - residual = hidden_states - hidden_states = norm(hidden_states) - - batch_size, num_channels, num_frames, height, width = hidden_states.shape - hidden_states = hidden_states.permute(0, 3, 4, 2, 1).flatten(0, 2).contiguous() - - # Perform attention in chunks to avoid following error: - # RuntimeError: CUDA error: invalid configuration argument - if hidden_states.size(0) <= chunk_size: - hidden_states = attn(hidden_states) - else: - hidden_states_chunks = [] - for i in range(0, hidden_states.size(0), chunk_size): - hidden_states_chunk = hidden_states[i : i + chunk_size] - hidden_states_chunk = attn(hidden_states_chunk) - hidden_states_chunks.append(hidden_states_chunk) - hidden_states = torch.cat(hidden_states_chunks) - - hidden_states = hidden_states.unflatten(0, (batch_size, height, width)).permute(0, 4, 3, 1, 2) - - hidden_states = residual + hidden_states - - return hidden_states, new_conv_cache - - -class MochiMidBlock3D(nn.Module): - r""" - A middle block used in the Mochi model. - - Args: - in_channels (`int`): - Number of input channels. - num_layers (`int`, defaults to `3`): - Number of resnet blocks in the block. - """ - - def __init__( - self, - in_channels: int, # 768 - num_layers: int = 3, - add_attention: bool = True, - ): - super().__init__() - - resnets = [] - norms = [] - attentions = [] - - for _ in range(num_layers): - resnets.append(MochiResnetBlock3D(in_channels=in_channels)) - - if add_attention: - norms.append(MochiChunkedGroupNorm3D(num_channels=in_channels)) - attentions.append( - Attention( - query_dim=in_channels, - heads=in_channels // 32, - dim_head=32, - qk_norm="l2", - is_causal=True, - processor=MochiVaeAttnProcessor2_0(), - ) - ) - else: - norms.append(None) - attentions.append(None) - - self.resnets = nn.ModuleList(resnets) - self.norms = nn.ModuleList(norms) - self.attentions = nn.ModuleList(attentions) - - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, - ) -> torch.Tensor: - r"""Forward method of the `MochiMidBlock3D` class.""" - - new_conv_cache = {} - conv_cache = conv_cache or {} - - for i, (resnet, norm, attn) in enumerate(zip(self.resnets, self.norms, self.attentions)): - conv_cache_key = f"resnet_{i}" - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, new_conv_cache[conv_cache_key] = self._gradient_checkpointing_func( - resnet, hidden_states, conv_cache.get(conv_cache_key) - ) - else: - hidden_states, new_conv_cache[conv_cache_key] = resnet( - hidden_states, conv_cache=conv_cache.get(conv_cache_key) - ) - - if attn is not None: - residual = hidden_states - hidden_states = norm(hidden_states) - - batch_size, num_channels, num_frames, height, width = hidden_states.shape - hidden_states = hidden_states.permute(0, 3, 4, 2, 1).flatten(0, 2).contiguous() - hidden_states = attn(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, height, width)).permute(0, 4, 3, 1, 2) - - hidden_states = residual + hidden_states - - return hidden_states, new_conv_cache - - -class MochiUpBlock3D(nn.Module): - r""" - An upsampling block used in the Mochi 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 blocks in the block. - temporal_expansion (`int`, defaults to `2`): - Temporal expansion factor. - spatial_expansion (`int`, defaults to `2`): - Spatial expansion factor. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - num_layers: int = 1, - temporal_expansion: int = 2, - spatial_expansion: int = 2, - ): - super().__init__() - self.temporal_expansion = temporal_expansion - self.spatial_expansion = spatial_expansion - - resnets = [] - for _ in range(num_layers): - resnets.append(MochiResnetBlock3D(in_channels=in_channels)) - self.resnets = nn.ModuleList(resnets) - - self.proj = nn.Linear(in_channels, out_channels * temporal_expansion * spatial_expansion**2) - - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, - ) -> torch.Tensor: - r"""Forward method of the `MochiUpBlock3D` class.""" - - new_conv_cache = {} - conv_cache = conv_cache or {} - - for i, resnet in enumerate(self.resnets): - conv_cache_key = f"resnet_{i}" - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, new_conv_cache[conv_cache_key] = self._gradient_checkpointing_func( - resnet, - hidden_states, - conv_cache.get(conv_cache_key), - ) - else: - hidden_states, new_conv_cache[conv_cache_key] = resnet( - hidden_states, conv_cache=conv_cache.get(conv_cache_key) - ) - - hidden_states = hidden_states.permute(0, 2, 3, 4, 1) - hidden_states = self.proj(hidden_states) - hidden_states = hidden_states.permute(0, 4, 1, 2, 3) - - batch_size, num_channels, num_frames, height, width = hidden_states.shape - st = self.temporal_expansion - sh = self.spatial_expansion - sw = self.spatial_expansion - - # Reshape and unpatchify - hidden_states = hidden_states.view(batch_size, -1, st, sh, sw, num_frames, height, width) - hidden_states = hidden_states.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() - hidden_states = hidden_states.view(batch_size, -1, num_frames * st, height * sh, width * sw) - - return hidden_states, new_conv_cache - - -class FourierFeatures(nn.Module): - def __init__(self, start: int = 6, stop: int = 8, step: int = 1): - super().__init__() - - self.start = start - self.stop = stop - self.step = step - - def forward(self, inputs: torch.Tensor) -> torch.Tensor: - r"""Forward method of the `FourierFeatures` class.""" - original_dtype = inputs.dtype - inputs = inputs.to(torch.float32) - num_channels = inputs.shape[1] - num_freqs = (self.stop - self.start) // self.step - - freqs = torch.arange(self.start, self.stop, self.step, dtype=inputs.dtype, device=inputs.device) - w = torch.pow(2.0, freqs) * (2 * torch.pi) # [num_freqs] - w = w.repeat(num_channels)[None, :, None, None, None] # [1, num_channels * num_freqs, 1, 1, 1] - - # Interleaved repeat of input channels to match w - h = inputs.repeat_interleave( - num_freqs, dim=1, output_size=inputs.shape[1] * num_freqs - ) # [B, C * num_freqs, T, H, W] - # Scale channels by frequency. - h = w * h - - return torch.cat([inputs, torch.sin(h), torch.cos(h)], dim=1).to(original_dtype) - - -class MochiEncoder3D(nn.Module): - r""" - The `MochiEncoder3D` layer of a variational autoencoder that encodes input video samples to its latent - representation. - - Args: - in_channels (`int`, *optional*): - 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)`): - The number of output channels for each block. - 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)`): - The temporal expansion factor for each of the up blocks. - 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. - """ - - 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), - act_fn: str = "swish", - ): - super().__init__() - - self.nonlinearity = get_activation(act_fn) - - self.fourier_features = FourierFeatures() - self.proj_in = nn.Linear(in_channels, block_out_channels[0]) - self.block_in = MochiMidBlock3D( - in_channels=block_out_channels[0], num_layers=layers_per_block[0], add_attention=add_attention_block[0] - ) - - down_blocks = [] - for i in range(len(block_out_channels) - 1): - down_block = MochiDownBlock3D( - in_channels=block_out_channels[i], - out_channels=block_out_channels[i + 1], - num_layers=layers_per_block[i + 1], - temporal_expansion=temporal_expansions[i], - spatial_expansion=spatial_expansions[i], - add_attention=add_attention_block[i + 1], - ) - down_blocks.append(down_block) - self.down_blocks = nn.ModuleList(down_blocks) - - self.block_out = MochiMidBlock3D( - in_channels=block_out_channels[-1], num_layers=layers_per_block[-1], add_attention=add_attention_block[-1] - ) - self.norm_out = MochiChunkedGroupNorm3D(block_out_channels[-1]) - self.proj_out = nn.Linear(block_out_channels[-1], 2 * out_channels, bias=False) - - self.gradient_checkpointing = False - - def forward( - self, hidden_states: torch.Tensor, conv_cache: Optional[Dict[str, torch.Tensor]] = None - ) -> torch.Tensor: - r"""Forward method of the `MochiEncoder3D` class.""" - - new_conv_cache = {} - conv_cache = conv_cache or {} - - hidden_states = self.fourier_features(hidden_states) - - hidden_states = hidden_states.permute(0, 2, 3, 4, 1) - hidden_states = self.proj_in(hidden_states) - hidden_states = hidden_states.permute(0, 4, 1, 2, 3) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, new_conv_cache["block_in"] = self._gradient_checkpointing_func( - self.block_in, hidden_states, conv_cache.get("block_in") - ) - - for i, down_block in enumerate(self.down_blocks): - conv_cache_key = f"down_block_{i}" - hidden_states, new_conv_cache[conv_cache_key] = self._gradient_checkpointing_func( - down_block, hidden_states, conv_cache.get(conv_cache_key) - ) - else: - hidden_states, new_conv_cache["block_in"] = self.block_in( - hidden_states, conv_cache=conv_cache.get("block_in") - ) - - for i, down_block in enumerate(self.down_blocks): - conv_cache_key = f"down_block_{i}" - hidden_states, new_conv_cache[conv_cache_key] = down_block( - hidden_states, conv_cache=conv_cache.get(conv_cache_key) - ) - - hidden_states, new_conv_cache["block_out"] = self.block_out( - hidden_states, conv_cache=conv_cache.get("block_out") - ) - - hidden_states = self.norm_out(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - - hidden_states = hidden_states.permute(0, 2, 3, 4, 1) - hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.permute(0, 4, 1, 2, 3) - - return hidden_states, new_conv_cache - - -class MochiDecoder3D(nn.Module): - r""" - The `MochiDecoder3D` layer of a variational autoencoder that decodes its latent representation into an output - sample. - - Args: - in_channels (`int`, *optional*): - 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)`): - The number of output channels for each block. - 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)`): - The temporal expansion factor for each of the up blocks. - 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. - """ - - 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), - act_fn: str = "swish", - ): - super().__init__() - - self.nonlinearity = get_activation(act_fn) - - self.conv_in = nn.Conv3d(in_channels, block_out_channels[-1], kernel_size=(1, 1, 1)) - self.block_in = MochiMidBlock3D( - in_channels=block_out_channels[-1], - num_layers=layers_per_block[-1], - add_attention=False, - ) - - up_blocks = [] - for i in range(len(block_out_channels) - 1): - up_block = MochiUpBlock3D( - in_channels=block_out_channels[-i - 1], - out_channels=block_out_channels[-i - 2], - num_layers=layers_per_block[-i - 2], - temporal_expansion=temporal_expansions[-i - 1], - spatial_expansion=spatial_expansions[-i - 1], - ) - up_blocks.append(up_block) - self.up_blocks = nn.ModuleList(up_blocks) - - self.block_out = MochiMidBlock3D( - in_channels=block_out_channels[0], - num_layers=layers_per_block[0], - add_attention=False, - ) - self.proj_out = nn.Linear(block_out_channels[0], out_channels) - - self.gradient_checkpointing = False - - def forward( - self, hidden_states: torch.Tensor, conv_cache: Optional[Dict[str, torch.Tensor]] = None - ) -> torch.Tensor: - r"""Forward method of the `MochiDecoder3D` class.""" - - new_conv_cache = {} - conv_cache = conv_cache or {} - - hidden_states = self.conv_in(hidden_states) - - # 1. Mid - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, new_conv_cache["block_in"] = self._gradient_checkpointing_func( - self.block_in, hidden_states, conv_cache.get("block_in") - ) - - for i, up_block in enumerate(self.up_blocks): - conv_cache_key = f"up_block_{i}" - hidden_states, new_conv_cache[conv_cache_key] = self._gradient_checkpointing_func( - up_block, hidden_states, conv_cache.get(conv_cache_key) - ) - else: - hidden_states, new_conv_cache["block_in"] = self.block_in( - hidden_states, conv_cache=conv_cache.get("block_in") - ) - - for i, up_block in enumerate(self.up_blocks): - conv_cache_key = f"up_block_{i}" - hidden_states, new_conv_cache[conv_cache_key] = up_block( - hidden_states, conv_cache=conv_cache.get(conv_cache_key) - ) - - hidden_states, new_conv_cache["block_out"] = self.block_out( - hidden_states, conv_cache=conv_cache.get("block_out") - ) - - hidden_states = self.nonlinearity(hidden_states) - - hidden_states = hidden_states.permute(0, 2, 3, 4, 1) - hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.permute(0, 4, 1, 2, 3) - - return hidden_states, new_conv_cache - - -class AutoencoderKLMochi(ModelMixin, ConfigMixin): - r""" - A VAE model with KL loss for encoding images into latents and decoding latent representations into images. Used in - [Mochi 1 preview](https://github.com/genmoai/models). - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - 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. - 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 - 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. - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["MochiResnetBlock3D"] - - @register_to_config - 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), - latent_channels: int = 12, - 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, ...] = ( - -0.06730895953510081, - -0.038011381506090416, - -0.07477820912866141, - -0.05565264470995561, - 0.012767231469026969, - -0.04703542746246419, - 0.043896967884726704, - -0.09346305707025976, - -0.09918314763016893, - -0.008729793427399178, - -0.011931556316503654, - -0.0321993391887285, - ), - latents_std: Tuple[float, ...] = ( - 0.9263795028493863, - 0.9248894543193766, - 0.9393059390890617, - 0.959253732819592, - 0.8244560132752793, - 0.917259975397747, - 0.9294154431013696, - 1.3720942357788521, - 0.881393668867029, - 0.9168315692124348, - 0.9185249279345552, - 0.9274757570805041, - ), - scaling_factor: float = 1.0, - ): - super().__init__() - - self.encoder = MochiEncoder3D( - in_channels=in_channels, - out_channels=latent_channels, - block_out_channels=encoder_block_out_channels, - layers_per_block=layers_per_block, - temporal_expansions=temporal_expansions, - spatial_expansions=spatial_expansions, - add_attention_block=add_attention_block, - act_fn=act_fn, - ) - self.decoder = MochiDecoder3D( - in_channels=latent_channels, - out_channels=out_channels, - block_out_channels=decoder_block_out_channels, - layers_per_block=layers_per_block, - temporal_expansions=temporal_expansions, - spatial_expansions=spatial_expansions, - act_fn=act_fn, - ) - - self.spatial_compression_ratio = functools.reduce(lambda x, y: x * y, spatial_expansions, 1) - self.temporal_compression_ratio = functools.reduce(lambda x, y: x * y, temporal_expansions, 1) - - # 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 used to determine how the number of output frames in the final decoded video. To maintain consistency with - # the original implementation, this defaults to `True`. - # - Original implementation (drop_last_temporal_frames=True): - # Output frames = (latent_frames - 1) * temporal_compression_ratio + 1 - # - Without dropping additional temporal upscaled frames (drop_last_temporal_frames=False): - # Output frames = latent_frames * temporal_compression_ratio - # The latter case is useful for frame packing and some training/finetuning scenarios where the additional. - self.drop_last_temporal_frames = True - - # This can be configured based on the amount of GPU memory available. - # `12` 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 = 12 - self.num_latent_frames_batch_size = 2 - - # The minimal tile height and width for spatial tiling to be used - self.tile_sample_min_height = 256 - self.tile_sample_min_width = 256 - - # The minimal distance between two spatial tiles - self.tile_sample_stride_height = 192 - self.tile_sample_stride_width = 192 - - 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, - ) -> 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_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 - - 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 _enable_framewise_encoding(self): - r""" - Enables the framewise VAE encoding implementation with past latent padding. By default, Diffusers uses the - oneshot encoding implementation without current latent replicate padding. - - Warning: Framewise encoding may not work as expected due to the causal attention layers. If you enable - framewise encoding, encode a video, and try to decode it, there will be noticeable jittering effect. - """ - self.use_framewise_encoding = True - for name, module in self.named_modules(): - if isinstance(module, CogVideoXCausalConv3d): - module.pad_mode = "constant" - - def _enable_framewise_decoding(self): - r""" - Enables the framewise VAE decoding implementation with past latent padding. By default, Diffusers uses the - oneshot decoding implementation without current latent replicate padding. - """ - self.use_framewise_decoding = True - for name, module in self.named_modules(): - if isinstance(module, CogVideoXCausalConv3d): - module.pad_mode = "constant" - - def _encode(self, x: torch.Tensor) -> torch.Tensor: - batch_size, num_channels, num_frames, height, width = x.shape - - if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): - return self.tiled_encode(x) - - if self.use_framewise_encoding: - raise NotImplementedError( - "Frame-wise encoding does not work with the Mochi VAE Encoder due to the presence of attention layers. " - "As intermediate frames are not independent from each other, they cannot be encoded frame-wise." - ) - else: - enc, _ = self.encoder(x) - - return enc - - @apply_forward_hook - def encode( - self, x: torch.Tensor, return_dict: bool = True - ) -> Union[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) 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, return_dict: bool = True) -> Union[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 - - if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height): - return self.tiled_decode(z, return_dict=return_dict) - - if self.use_framewise_decoding: - conv_cache = None - dec = [] - - for i in range(0, num_frames, self.num_latent_frames_batch_size): - z_intermediate = z[:, :, i : i + self.num_latent_frames_batch_size] - z_intermediate, conv_cache = self.decoder(z_intermediate, conv_cache=conv_cache) - dec.append(z_intermediate) - - dec = torch.cat(dec, dim=2) - else: - dec, _ = self.decoder(z) - - if self.drop_last_temporal_frames and dec.size(2) >= self.temporal_compression_ratio: - dec = dec[:, :, self.temporal_compression_ratio - 1 :] - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) - - @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[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: - decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] - decoded = torch.cat(decoded_slices) - else: - decoded = self._decode(z).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 tiled_encode(self, x: torch.Tensor) -> 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): - if self.use_framewise_encoding: - raise NotImplementedError( - "Frame-wise encoding does not work with the Mochi VAE Encoder due to the presence of attention layers. " - "As intermediate frames are not independent from each other, they cannot be encoded frame-wise." - ) - else: - time, _ = self.encoder( - x[:, :, :, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width] - ) - - 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, return_dict: bool = True) -> Union[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): - if self.use_framewise_decoding: - time = [] - conv_cache = None - - for k in range(0, num_frames, self.num_latent_frames_batch_size): - tile = z[ - :, - :, - k : k + self.num_latent_frames_batch_size, - i : i + tile_latent_min_height, - j : j + tile_latent_min_width, - ] - tile, conv_cache = self.decoder(tile, conv_cache=conv_cache) - time.append(tile) - - time = torch.cat(time, dim=2) - else: - time, _ = self.decoder(z[:, :, :, i : i + tile_latent_min_height, j : j + tile_latent_min_width]) - - if self.drop_last_temporal_frames and time.size(2) >= self.temporal_compression_ratio: - time = time[:, :, self.temporal_compression_ratio - 1 :] - - 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 forward( - self, - sample: torch.Tensor, - sample_posterior: bool = False, - return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[torch.Tensor, torch.Tensor]: - x = sample - posterior = self.encode(x).latent_dist - if sample_posterior: - z = posterior.sample(generator=generator) - else: - z = posterior.mode() - dec = self.decode(z) - if not return_dict: - return (dec,) - return dec diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py deleted file mode 100644 index cf46e52564bfc3bd58549b183bfba944ec8a81ad..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py +++ /dev/null @@ -1,363 +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 itertools -from typing import Dict, Optional, Tuple, Union - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils.accelerate_utils import apply_forward_hook -from ..attention_processor import CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnProcessor -from ..modeling_outputs import AutoencoderKLOutput -from ..modeling_utils import ModelMixin -from ..unets.unet_3d_blocks import MidBlockTemporalDecoder, UpBlockTemporalDecoder -from .vae import DecoderOutput, DiagonalGaussianDistribution, Encoder - - -class TemporalDecoder(nn.Module): - def __init__( - self, - in_channels: int = 4, - out_channels: int = 3, - block_out_channels: Tuple[int] = (128, 256, 512, 512), - layers_per_block: int = 2, - ): - super().__init__() - self.layers_per_block = layers_per_block - - self.conv_in = nn.Conv2d(in_channels, block_out_channels[-1], kernel_size=3, stride=1, padding=1) - self.mid_block = MidBlockTemporalDecoder( - num_layers=self.layers_per_block, - in_channels=block_out_channels[-1], - out_channels=block_out_channels[-1], - attention_head_dim=block_out_channels[-1], - ) - - # up - self.up_blocks = nn.ModuleList([]) - reversed_block_out_channels = list(reversed(block_out_channels)) - output_channel = reversed_block_out_channels[0] - for i in range(len(block_out_channels)): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - - is_final_block = i == len(block_out_channels) - 1 - up_block = UpBlockTemporalDecoder( - num_layers=self.layers_per_block + 1, - in_channels=prev_output_channel, - out_channels=output_channel, - add_upsample=not is_final_block, - ) - self.up_blocks.append(up_block) - prev_output_channel = output_channel - - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=32, eps=1e-6) - - self.conv_act = nn.SiLU() - self.conv_out = torch.nn.Conv2d( - in_channels=block_out_channels[0], - out_channels=out_channels, - kernel_size=3, - padding=1, - ) - - conv_out_kernel_size = (3, 1, 1) - padding = [int(k // 2) for k in conv_out_kernel_size] - self.time_conv_out = torch.nn.Conv3d( - in_channels=out_channels, - out_channels=out_channels, - kernel_size=conv_out_kernel_size, - padding=padding, - ) - - self.gradient_checkpointing = False - - def forward( - self, - sample: torch.Tensor, - image_only_indicator: torch.Tensor, - num_frames: int = 1, - ) -> torch.Tensor: - r"""The forward method of the `Decoder` class.""" - - sample = self.conv_in(sample) - - upscale_dtype = next(itertools.chain(self.up_blocks.parameters(), self.up_blocks.buffers())).dtype - if torch.is_grad_enabled() and self.gradient_checkpointing: - # middle - sample = self._gradient_checkpointing_func( - self.mid_block, - sample, - image_only_indicator, - ) - sample = sample.to(upscale_dtype) - - # up - for up_block in self.up_blocks: - sample = self._gradient_checkpointing_func( - up_block, - sample, - image_only_indicator, - ) - else: - # middle - sample = self.mid_block(sample, image_only_indicator=image_only_indicator) - sample = sample.to(upscale_dtype) - - # up - for up_block in self.up_blocks: - sample = up_block(sample, image_only_indicator=image_only_indicator) - - # post-process - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - sample = self.conv_out(sample) - - batch_frames, channels, height, width = sample.shape - batch_size = batch_frames // num_frames - sample = sample[None, :].reshape(batch_size, num_frames, channels, height, width).permute(0, 2, 1, 3, 4) - sample = self.time_conv_out(sample) - - sample = sample.permute(0, 2, 1, 3, 4).reshape(batch_frames, channels, height, width) - - return sample - - -class AutoencoderKLTemporalDecoder(ModelMixin, ConfigMixin): - r""" - A VAE model with KL loss for encoding images into latents and decoding latent representations into images. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - 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. - 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. - scaling_factor (`float`, *optional*, defaults to 0.18215): - 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. - force_upcast (`bool`, *optional*, default to `True`): - If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE - can be fine-tuned / trained to a lower range without losing too much precision in which case `force_upcast` - can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix - """ - - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - 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, - scaling_factor: float = 0.18215, - force_upcast: float = True, - ): - super().__init__() - - # pass init params to Encoder - self.encoder = Encoder( - in_channels=in_channels, - out_channels=latent_channels, - down_block_types=down_block_types, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - double_z=True, - ) - - # pass init params to Decoder - self.decoder = TemporalDecoder( - in_channels=latent_channels, - out_channels=out_channels, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - ) - - self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - self.set_attn_processor(processor) - - @apply_forward_hook - def encode( - self, x: torch.Tensor, return_dict: bool = True - ) -> Union[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.autoencoders.autoencoder_kl.AutoencoderKLOutput`] instead of a plain - tuple. - - Returns: - The latent representations of the encoded images. If `return_dict` is True, a - [`~models.autoencoders.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is - returned. - """ - h = self.encoder(x) - moments = self.quant_conv(h) - posterior = DiagonalGaussianDistribution(moments) - - if not return_dict: - return (posterior,) - - return AutoencoderKLOutput(latent_dist=posterior) - - @apply_forward_hook - def decode( - self, - z: torch.Tensor, - num_frames: int, - return_dict: bool = True, - ) -> Union[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. - - """ - batch_size = z.shape[0] // num_frames - image_only_indicator = torch.zeros(batch_size, num_frames, dtype=z.dtype, device=z.device) - decoded = self.decoder(z, num_frames=num_frames, image_only_indicator=image_only_indicator) - - 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: Optional[torch.Generator] = None, - num_frames: int = 1, - ) -> Union[DecoderOutput, torch.Tensor]: - r""" - Args: - sample (`torch.Tensor`): Input sample. - sample_posterior (`bool`, *optional*, defaults to `False`): - Whether to sample from the posterior. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`DecoderOutput`] instead of a plain tuple. - """ - x = sample - posterior = self.encode(x).latent_dist - if sample_posterior: - z = posterior.sample(generator=generator) - else: - z = posterior.mode() - - dec = self.decode(z, num_frames=num_frames).sample - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_wan.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_wan.py deleted file mode 100644 index 49cefcd8a142f5575fb28b0773f4873083429363..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl_wan.py +++ /dev/null @@ -1,1089 +0,0 @@ -# Copyright 2025 The Wan 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 List, Optional, Tuple, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch.utils.checkpoint - -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import FromOriginalModelMixin -from ...utils import logging -from ...utils.accelerate_utils import apply_forward_hook -from ..activations import get_activation -from ..modeling_outputs import AutoencoderKLOutput -from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, DiagonalGaussianDistribution - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -CACHE_T = 2 - - -class WanCausalConv3d(nn.Conv3d): - r""" - A custom 3D causal convolution layer with feature caching support. - - This layer extends the standard Conv3D layer by ensuring causality in the time dimension and handling feature - caching for efficient inference. - - Args: - in_channels (int): Number of channels in the input image - out_channels (int): Number of channels produced by the convolution - kernel_size (int or tuple): Size of the convolving kernel - stride (int or tuple, optional): Stride of the convolution. Default: 1 - padding (int or tuple, optional): Zero-padding added to all three sides of the input. Default: 0 - """ - - 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, - ) -> None: - super().__init__( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - stride=stride, - padding=padding, - ) - - # Set up causal padding - self._padding = (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 2 * self.padding[0], 0) - self.padding = (0, 0, 0) - - def forward(self, x, cache_x=None): - padding = list(self._padding) - if cache_x is not None and self._padding[4] > 0: - cache_x = cache_x.to(x.device) - x = torch.cat([cache_x, x], dim=2) - padding[4] -= cache_x.shape[2] - x = F.pad(x, padding) - return super().forward(x) - - -class WanRMS_norm(nn.Module): - r""" - A custom RMS normalization layer. - - Args: - dim (int): The number of dimensions to normalize over. - channel_first (bool, optional): Whether the input tensor has channels as the first dimension. - Default is True. - images (bool, optional): Whether the input represents image data. Default is True. - bias (bool, optional): Whether to include a learnable bias term. Default is False. - """ - - def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None: - super().__init__() - broadcastable_dims = (1, 1, 1) if not images else (1, 1) - shape = (dim, *broadcastable_dims) if channel_first else (dim,) - - self.channel_first = channel_first - self.scale = dim**0.5 - self.gamma = nn.Parameter(torch.ones(shape)) - self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 - - def forward(self, x): - return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias - - -class WanUpsample(nn.Upsample): - r""" - Perform upsampling while ensuring the output tensor has the same data type as the input. - - Args: - x (torch.Tensor): Input tensor to be upsampled. - - Returns: - torch.Tensor: Upsampled tensor with the same data type as the input. - """ - - def forward(self, x): - return super().forward(x.float()).type_as(x) - - -class WanResample(nn.Module): - r""" - A custom resampling module for 2D and 3D data. - - Args: - dim (int): The number of input/output channels. - mode (str): The resampling mode. Must be one of: - - 'none': No resampling (identity operation). - - 'upsample2d': 2D upsampling with nearest-exact interpolation and convolution. - - 'upsample3d': 3D upsampling with nearest-exact interpolation, convolution, and causal 3D convolution. - - 'downsample2d': 2D downsampling with zero-padding and convolution. - - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution. - """ - - def __init__(self, dim: int, mode: str) -> None: - super().__init__() - self.dim = dim - self.mode = mode - - # layers - if mode == "upsample2d": - self.resample = nn.Sequential( - WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim // 2, 3, padding=1) - ) - elif mode == "upsample3d": - self.resample = nn.Sequential( - WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim // 2, 3, padding=1) - ) - self.time_conv = WanCausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) - - elif mode == "downsample2d": - self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) - elif mode == "downsample3d": - self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) - self.time_conv = WanCausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) - - else: - self.resample = nn.Identity() - - def forward(self, x, feat_cache=None, feat_idx=[0]): - b, c, t, h, w = x.size() - if self.mode == "upsample3d": - if feat_cache is not None: - idx = feat_idx[0] - if feat_cache[idx] is None: - feat_cache[idx] = "Rep" - feat_idx[0] += 1 - else: - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep": - # cache last frame of last two chunk - cache_x = torch.cat( - [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 - ) - if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep": - cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) - if feat_cache[idx] == "Rep": - x = self.time_conv(x) - else: - x = self.time_conv(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - - x = x.reshape(b, 2, c, t, h, w) - x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) - x = x.reshape(b, c, t * 2, h, w) - t = x.shape[2] - x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) - x = self.resample(x) - x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4) - - if self.mode == "downsample3d": - if feat_cache is not None: - idx = feat_idx[0] - if feat_cache[idx] is None: - feat_cache[idx] = x.clone() - feat_idx[0] += 1 - else: - cache_x = x[:, :, -1:, :, :].clone() - x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - return x - - -class WanResidualBlock(nn.Module): - r""" - A custom residual block module. - - Args: - in_dim (int): Number of input channels. - out_dim (int): Number of output channels. - dropout (float, optional): Dropout rate for the dropout layer. Default is 0.0. - non_linearity (str, optional): Type of non-linearity to use. Default is "silu". - """ - - def __init__( - self, - in_dim: int, - out_dim: int, - dropout: float = 0.0, - non_linearity: str = "silu", - ) -> None: - super().__init__() - self.in_dim = in_dim - self.out_dim = out_dim - self.nonlinearity = get_activation(non_linearity) - - # layers - self.norm1 = WanRMS_norm(in_dim, images=False) - self.conv1 = WanCausalConv3d(in_dim, out_dim, 3, padding=1) - self.norm2 = WanRMS_norm(out_dim, images=False) - self.dropout = nn.Dropout(dropout) - self.conv2 = WanCausalConv3d(out_dim, out_dim, 3, padding=1) - self.conv_shortcut = WanCausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() - - def forward(self, x, feat_cache=None, feat_idx=[0]): - # Apply shortcut connection - h = self.conv_shortcut(x) - - # First normalization and activation - x = self.norm1(x) - x = self.nonlinearity(x) - - if feat_cache is not None: - idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) - - x = self.conv1(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - else: - x = self.conv1(x) - - # Second normalization and activation - x = self.norm2(x) - x = self.nonlinearity(x) - - # Dropout - x = self.dropout(x) - - if feat_cache is not None: - idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) - - x = self.conv2(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - else: - x = self.conv2(x) - - # Add residual connection - return x + h - - -class WanAttentionBlock(nn.Module): - r""" - Causal self-attention with a single head. - - Args: - dim (int): The number of channels in the input tensor. - """ - - def __init__(self, dim): - super().__init__() - self.dim = dim - - # layers - self.norm = WanRMS_norm(dim) - self.to_qkv = nn.Conv2d(dim, dim * 3, 1) - self.proj = nn.Conv2d(dim, dim, 1) - - def forward(self, x): - identity = x - batch_size, channels, time, height, width = x.size() - - x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, height, width) - x = self.norm(x) - - # compute query, key, value - qkv = self.to_qkv(x) - qkv = qkv.reshape(batch_size * time, 1, channels * 3, -1) - qkv = qkv.permute(0, 1, 3, 2).contiguous() - q, k, v = qkv.chunk(3, dim=-1) - - # apply attention - x = F.scaled_dot_product_attention(q, k, v) - - x = x.squeeze(1).permute(0, 2, 1).reshape(batch_size * time, channels, height, width) - - # output projection - x = self.proj(x) - - # Reshape back: [(b*t), c, h, w] -> [b, c, t, h, w] - x = x.view(batch_size, time, channels, height, width) - x = x.permute(0, 2, 1, 3, 4) - - return x + identity - - -class WanMidBlock(nn.Module): - """ - Middle block for WanVAE encoder and decoder. - - Args: - dim (int): Number of input/output channels. - dropout (float): Dropout rate. - non_linearity (str): Type of non-linearity to use. - """ - - def __init__(self, dim: int, dropout: float = 0.0, non_linearity: str = "silu", num_layers: int = 1): - super().__init__() - self.dim = dim - - # Create the components - resnets = [WanResidualBlock(dim, dim, dropout, non_linearity)] - attentions = [] - for _ in range(num_layers): - attentions.append(WanAttentionBlock(dim)) - resnets.append(WanResidualBlock(dim, dim, dropout, non_linearity)) - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - self.gradient_checkpointing = False - - def forward(self, x, feat_cache=None, feat_idx=[0]): - # First residual block - x = self.resnets[0](x, feat_cache, feat_idx) - - # Process through attention and residual blocks - for attn, resnet in zip(self.attentions, self.resnets[1:]): - if attn is not None: - x = attn(x) - - x = resnet(x, feat_cache, feat_idx) - - return x - - -class WanEncoder3d(nn.Module): - r""" - A 3D encoder module. - - Args: - dim (int): The base number of channels in the first layer. - z_dim (int): The dimensionality of the latent space. - dim_mult (list of int): Multipliers for the number of channels in each block. - num_res_blocks (int): Number of residual blocks in each block. - attn_scales (list of float): Scales at which to apply attention mechanisms. - temperal_downsample (list of bool): Whether to downsample temporally in each block. - dropout (float): Dropout rate for the dropout layers. - non_linearity (str): Type of non-linearity to use. - """ - - def __init__( - self, - dim=128, - z_dim=4, - dim_mult=[1, 2, 4, 4], - num_res_blocks=2, - attn_scales=[], - temperal_downsample=[True, True, False], - dropout=0.0, - non_linearity: str = "silu", - ): - super().__init__() - self.dim = dim - self.z_dim = z_dim - self.dim_mult = dim_mult - self.num_res_blocks = num_res_blocks - self.attn_scales = attn_scales - self.temperal_downsample = temperal_downsample - self.nonlinearity = get_activation(non_linearity) - - # dimensions - dims = [dim * u for u in [1] + dim_mult] - scale = 1.0 - - # init block - self.conv_in = WanCausalConv3d(3, dims[0], 3, padding=1) - - # downsample blocks - self.down_blocks = nn.ModuleList([]) - for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): - # residual (+attention) blocks - for _ in range(num_res_blocks): - self.down_blocks.append(WanResidualBlock(in_dim, out_dim, dropout)) - if scale in attn_scales: - self.down_blocks.append(WanAttentionBlock(out_dim)) - in_dim = out_dim - - # downsample block - if i != len(dim_mult) - 1: - mode = "downsample3d" if temperal_downsample[i] else "downsample2d" - self.down_blocks.append(WanResample(out_dim, mode=mode)) - scale /= 2.0 - - # middle blocks - self.mid_block = WanMidBlock(out_dim, dropout, non_linearity, num_layers=1) - - # output blocks - self.norm_out = WanRMS_norm(out_dim, images=False) - self.conv_out = WanCausalConv3d(out_dim, z_dim, 3, padding=1) - - self.gradient_checkpointing = False - - def forward(self, x, feat_cache=None, feat_idx=[0]): - if feat_cache is not None: - idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) - x = self.conv_in(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - else: - x = self.conv_in(x) - - ## downsamples - for layer in self.down_blocks: - if feat_cache is not None: - x = layer(x, feat_cache, feat_idx) - else: - x = layer(x) - - ## middle - x = self.mid_block(x, feat_cache, feat_idx) - - ## head - x = self.norm_out(x) - x = self.nonlinearity(x) - if feat_cache is not None: - idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) - x = self.conv_out(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - else: - x = self.conv_out(x) - return x - - -class WanUpBlock(nn.Module): - """ - A block that handles upsampling for the WanVAE decoder. - - Args: - in_dim (int): Input dimension - out_dim (int): Output dimension - num_res_blocks (int): Number of residual blocks - dropout (float): Dropout rate - upsample_mode (str, optional): Mode for upsampling ('upsample2d' or 'upsample3d') - non_linearity (str): Type of non-linearity to use - """ - - def __init__( - self, - in_dim: int, - out_dim: int, - num_res_blocks: int, - dropout: float = 0.0, - upsample_mode: Optional[str] = None, - non_linearity: str = "silu", - ): - super().__init__() - self.in_dim = in_dim - self.out_dim = out_dim - - # Create layers list - resnets = [] - # Add residual blocks and attention if needed - current_dim = in_dim - for _ in range(num_res_blocks + 1): - resnets.append(WanResidualBlock(current_dim, out_dim, dropout, non_linearity)) - current_dim = out_dim - - self.resnets = nn.ModuleList(resnets) - - # Add upsampling layer if needed - self.upsamplers = None - if upsample_mode is not None: - self.upsamplers = nn.ModuleList([WanResample(out_dim, mode=upsample_mode)]) - - self.gradient_checkpointing = False - - def forward(self, x, feat_cache=None, feat_idx=[0]): - """ - Forward pass through the upsampling block. - - Args: - x (torch.Tensor): Input tensor - feat_cache (list, optional): Feature cache for causal convolutions - feat_idx (list, optional): Feature index for cache management - - Returns: - torch.Tensor: Output tensor - """ - for resnet in self.resnets: - if feat_cache is not None: - x = resnet(x, feat_cache, feat_idx) - else: - x = resnet(x) - - if self.upsamplers is not None: - if feat_cache is not None: - x = self.upsamplers[0](x, feat_cache, feat_idx) - else: - x = self.upsamplers[0](x) - return x - - -class WanDecoder3d(nn.Module): - r""" - A 3D decoder module. - - Args: - dim (int): The base number of channels in the first layer. - z_dim (int): The dimensionality of the latent space. - dim_mult (list of int): Multipliers for the number of channels in each block. - num_res_blocks (int): Number of residual blocks in each block. - attn_scales (list of float): Scales at which to apply attention mechanisms. - temperal_upsample (list of bool): Whether to upsample temporally in each block. - dropout (float): Dropout rate for the dropout layers. - non_linearity (str): Type of non-linearity to use. - """ - - def __init__( - self, - dim=128, - z_dim=4, - dim_mult=[1, 2, 4, 4], - num_res_blocks=2, - attn_scales=[], - temperal_upsample=[False, True, True], - dropout=0.0, - non_linearity: str = "silu", - ): - super().__init__() - self.dim = dim - self.z_dim = z_dim - self.dim_mult = dim_mult - self.num_res_blocks = num_res_blocks - self.attn_scales = attn_scales - self.temperal_upsample = temperal_upsample - - self.nonlinearity = get_activation(non_linearity) - - # dimensions - dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] - scale = 1.0 / 2 ** (len(dim_mult) - 2) - - # init block - self.conv_in = WanCausalConv3d(z_dim, dims[0], 3, padding=1) - - # middle blocks - self.mid_block = WanMidBlock(dims[0], dropout, non_linearity, num_layers=1) - - # upsample blocks - self.up_blocks = nn.ModuleList([]) - for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): - # residual (+attention) blocks - if i > 0: - in_dim = in_dim // 2 - - # Determine if we need upsampling - upsample_mode = None - if i != len(dim_mult) - 1: - upsample_mode = "upsample3d" if temperal_upsample[i] else "upsample2d" - - # Create and add the upsampling block - up_block = WanUpBlock( - in_dim=in_dim, - out_dim=out_dim, - num_res_blocks=num_res_blocks, - dropout=dropout, - upsample_mode=upsample_mode, - non_linearity=non_linearity, - ) - self.up_blocks.append(up_block) - - # Update scale for next iteration - if upsample_mode is not None: - scale *= 2.0 - - # output blocks - self.norm_out = WanRMS_norm(out_dim, images=False) - self.conv_out = WanCausalConv3d(out_dim, 3, 3, padding=1) - - self.gradient_checkpointing = False - - def forward(self, x, feat_cache=None, feat_idx=[0]): - ## conv1 - if feat_cache is not None: - idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) - x = self.conv_in(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - else: - x = self.conv_in(x) - - ## middle - x = self.mid_block(x, feat_cache, feat_idx) - - ## upsamples - for up_block in self.up_blocks: - x = up_block(x, feat_cache, feat_idx) - - ## head - x = self.norm_out(x) - x = self.nonlinearity(x) - if feat_cache is not None: - idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) - x = self.conv_out(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - else: - x = self.conv_out(x) - return x - - -class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): - r""" - A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. - Introduced in [Wan 2.1]. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - """ - - _supports_gradient_checkpointing = False - - @register_to_config - def __init__( - self, - base_dim: int = 96, - z_dim: int = 16, - dim_mult: Tuple[int] = [1, 2, 4, 4], - num_res_blocks: int = 2, - attn_scales: List[float] = [], - temperal_downsample: List[bool] = [False, True, True], - dropout: float = 0.0, - 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: - super().__init__() - - self.z_dim = z_dim - self.temperal_downsample = temperal_downsample - self.temperal_upsample = temperal_downsample[::-1] - - self.encoder = WanEncoder3d( - base_dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout - ) - self.quant_conv = WanCausalConv3d(z_dim * 2, z_dim * 2, 1) - self.post_quant_conv = WanCausalConv3d(z_dim, z_dim, 1) - - self.decoder = WanDecoder3d( - base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout - ) - - self.spatial_compression_ratio = 2 ** len(self.temperal_downsample) - - # 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 - - # The minimal tile height and width for spatial tiling to be used - self.tile_sample_min_height = 256 - self.tile_sample_min_width = 256 - - # The minimal distance between two spatial tiles - self.tile_sample_stride_height = 192 - self.tile_sample_stride_width = 192 - - # Precompute and cache conv counts for encoder and decoder for clear_cache speedup - self._cached_conv_counts = { - "decoder": sum(isinstance(m, WanCausalConv3d) for m in self.decoder.modules()) - if self.decoder is not None - else 0, - "encoder": sum(isinstance(m, WanCausalConv3d) for m in self.encoder.modules()) - if self.encoder is not None - else 0, - } - - 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, - ) -> 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_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 - - 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 clear_cache(self): - # Use cached conv counts for decoder and encoder to avoid re-iterating modules each call - self._conv_num = self._cached_conv_counts["decoder"] - self._conv_idx = [0] - self._feat_map = [None] * self._conv_num - # cache encode - self._enc_conv_num = self._cached_conv_counts["encoder"] - self._enc_conv_idx = [0] - self._enc_feat_map = [None] * self._enc_conv_num - - def _encode(self, x: torch.Tensor): - _, _, num_frame, height, width = x.shape - - if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): - return self.tiled_encode(x) - - self.clear_cache() - iter_ = 1 + (num_frame - 1) // 4 - for i in range(iter_): - self._enc_conv_idx = [0] - if i == 0: - out = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) - else: - out_ = self.encoder( - x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], - feat_cache=self._enc_feat_map, - feat_idx=self._enc_conv_idx, - ) - out = torch.cat([out, out_], 2) - - enc = self.quant_conv(out) - self.clear_cache() - return enc - - @apply_forward_hook - def encode( - self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: - r""" - 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) 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, return_dict: bool = True): - _, _, num_frame, 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 - - if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height): - return self.tiled_decode(z, return_dict=return_dict) - - self.clear_cache() - x = self.post_quant_conv(z) - for i in range(num_frame): - self._conv_idx = [0] - if i == 0: - out = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) - else: - out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) - out = torch.cat([out, out_], 2) - - out = torch.clamp(out, min=-1.0, max=1.0) - self.clear_cache() - if not return_dict: - return (out,) - - return DecoderOutput(sample=out) - - @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: - r""" - 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: - decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] - decoded = torch.cat(decoded_slices) - else: - decoded = self._decode(z).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[-2], b.shape[-2], 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[-1], b.shape[-1], 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) -> AutoencoderKLOutput: - 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. - """ - _, _, 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): - self.clear_cache() - time = [] - frame_range = 1 + (num_frames - 1) // 4 - for k in range(frame_range): - self._enc_conv_idx = [0] - if k == 0: - tile = x[:, :, :1, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width] - else: - tile = x[ - :, - :, - 1 + 4 * (k - 1) : 1 + 4 * k, - i : i + self.tile_sample_min_height, - j : j + self.tile_sample_min_width, - ] - tile = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) - tile = self.quant_conv(tile) - time.append(tile) - row.append(torch.cat(time, dim=2)) - rows.append(row) - self.clear_cache() - - 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=-1)) - - 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]: - 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. - """ - _, _, 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): - self.clear_cache() - time = [] - for k in range(num_frames): - self._conv_idx = [0] - tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width] - tile = self.post_quant_conv(tile) - decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx) - time.append(decoded) - row.append(torch.cat(time, dim=2)) - rows.append(row) - self.clear_cache() - - 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=-1)) - - dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width] - - if not return_dict: - return (dec,) - return DecoderOutput(sample=dec) - - def forward( - self, - sample: torch.Tensor, - sample_posterior: bool = False, - return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: - """ - Args: - sample (`torch.Tensor`): Input sample. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`DecoderOutput`] instead of a plain tuple. - """ - x = sample - posterior = self.encode(x).latent_dist - if sample_posterior: - z = posterior.sample(generator=generator) - else: - z = posterior.mode() - dec = self.decode(z, return_dict=return_dict) - return dec diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_oobleck.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_oobleck.py deleted file mode 100644 index a10b616b4e25444c5998b81cc0d137e7dfdd5fd0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_oobleck.py +++ /dev/null @@ -1,465 +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 math -from dataclasses import dataclass -from typing import Optional, Tuple, Union - -import numpy as np -import torch -import torch.nn as nn -from torch.nn.utils import weight_norm - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import BaseOutput -from ...utils.accelerate_utils import apply_forward_hook -from ...utils.torch_utils import randn_tensor -from ..modeling_utils import ModelMixin - - -class Snake1d(nn.Module): - """ - A 1-dimensional Snake activation function module. - """ - - def __init__(self, hidden_dim, logscale=True): - super().__init__() - self.alpha = nn.Parameter(torch.zeros(1, hidden_dim, 1)) - self.beta = nn.Parameter(torch.zeros(1, hidden_dim, 1)) - - self.alpha.requires_grad = True - self.beta.requires_grad = True - self.logscale = logscale - - def forward(self, hidden_states): - shape = hidden_states.shape - - alpha = self.alpha if not self.logscale else torch.exp(self.alpha) - beta = self.beta if not self.logscale else torch.exp(self.beta) - - hidden_states = hidden_states.reshape(shape[0], shape[1], -1) - hidden_states = hidden_states + (beta + 1e-9).reciprocal() * torch.sin(alpha * hidden_states).pow(2) - hidden_states = hidden_states.reshape(shape) - return hidden_states - - -class OobleckResidualUnit(nn.Module): - """ - A residual unit composed of Snake1d and weight-normalized Conv1d layers with dilations. - """ - - def __init__(self, dimension: int = 16, dilation: int = 1): - super().__init__() - pad = ((7 - 1) * dilation) // 2 - - self.snake1 = Snake1d(dimension) - self.conv1 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=7, dilation=dilation, padding=pad)) - self.snake2 = Snake1d(dimension) - self.conv2 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=1)) - - def forward(self, hidden_state): - """ - Forward pass through the residual unit. - - Args: - hidden_state (`torch.Tensor` of shape `(batch_size, channels, time_steps)`): - Input tensor . - - Returns: - output_tensor (`torch.Tensor` of shape `(batch_size, channels, time_steps)`) - Input tensor after passing through the residual unit. - """ - output_tensor = hidden_state - output_tensor = self.conv1(self.snake1(output_tensor)) - output_tensor = self.conv2(self.snake2(output_tensor)) - - padding = (hidden_state.shape[-1] - output_tensor.shape[-1]) // 2 - if padding > 0: - hidden_state = hidden_state[..., padding:-padding] - output_tensor = hidden_state + output_tensor - return output_tensor - - -class OobleckEncoderBlock(nn.Module): - """Encoder block used in Oobleck encoder.""" - - def __init__(self, input_dim, output_dim, stride: int = 1): - super().__init__() - - self.res_unit1 = OobleckResidualUnit(input_dim, dilation=1) - self.res_unit2 = OobleckResidualUnit(input_dim, dilation=3) - self.res_unit3 = OobleckResidualUnit(input_dim, dilation=9) - self.snake1 = Snake1d(input_dim) - self.conv1 = weight_norm( - nn.Conv1d(input_dim, output_dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)) - ) - - def forward(self, hidden_state): - hidden_state = self.res_unit1(hidden_state) - hidden_state = self.res_unit2(hidden_state) - hidden_state = self.snake1(self.res_unit3(hidden_state)) - hidden_state = self.conv1(hidden_state) - - return hidden_state - - -class OobleckDecoderBlock(nn.Module): - """Decoder block used in Oobleck decoder.""" - - def __init__(self, input_dim, output_dim, stride: int = 1): - super().__init__() - - self.snake1 = Snake1d(input_dim) - self.conv_t1 = weight_norm( - nn.ConvTranspose1d( - input_dim, - output_dim, - kernel_size=2 * stride, - stride=stride, - padding=math.ceil(stride / 2), - ) - ) - self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1) - self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3) - self.res_unit3 = OobleckResidualUnit(output_dim, dilation=9) - - def forward(self, hidden_state): - hidden_state = self.snake1(hidden_state) - hidden_state = self.conv_t1(hidden_state) - hidden_state = self.res_unit1(hidden_state) - hidden_state = self.res_unit2(hidden_state) - hidden_state = self.res_unit3(hidden_state) - - return hidden_state - - -class OobleckDiagonalGaussianDistribution(object): - def __init__(self, parameters: torch.Tensor, deterministic: bool = False): - self.parameters = parameters - self.mean, self.scale = parameters.chunk(2, dim=1) - self.std = nn.functional.softplus(self.scale) + 1e-4 - self.var = self.std * self.std - self.logvar = torch.log(self.var) - self.deterministic = deterministic - - def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor: - # make sure sample is on the same device as the parameters and has same dtype - sample = randn_tensor( - self.mean.shape, - generator=generator, - device=self.parameters.device, - dtype=self.parameters.dtype, - ) - x = self.mean + self.std * sample - return x - - def kl(self, other: "OobleckDiagonalGaussianDistribution" = None) -> torch.Tensor: - if self.deterministic: - return torch.Tensor([0.0]) - else: - if other is None: - return (self.mean * self.mean + self.var - self.logvar - 1.0).sum(1).mean() - else: - normalized_diff = torch.pow(self.mean - other.mean, 2) / other.var - var_ratio = self.var / other.var - logvar_diff = self.logvar - other.logvar - - kl = normalized_diff + var_ratio + logvar_diff - 1 - - kl = kl.sum(1).mean() - return kl - - def mode(self) -> torch.Tensor: - return self.mean - - -@dataclass -class AutoencoderOobleckOutput(BaseOutput): - """ - Output of AutoencoderOobleck encoding method. - - Args: - latent_dist (`OobleckDiagonalGaussianDistribution`): - Encoded outputs of `Encoder` represented as the mean and standard deviation of - `OobleckDiagonalGaussianDistribution`. `OobleckDiagonalGaussianDistribution` allows for sampling latents - from the distribution. - """ - - latent_dist: "OobleckDiagonalGaussianDistribution" # noqa: F821 - - -@dataclass -class OobleckDecoderOutput(BaseOutput): - r""" - Output of decoding method. - - Args: - sample (`torch.Tensor` of shape `(batch_size, audio_channels, sequence_length)`): - The decoded output sample from the last layer of the model. - """ - - sample: torch.Tensor - - -class OobleckEncoder(nn.Module): - """Oobleck Encoder""" - - def __init__(self, encoder_hidden_size, audio_channels, downsampling_ratios, channel_multiples): - super().__init__() - - strides = downsampling_ratios - channel_multiples = [1] + channel_multiples - - # Create first convolution - self.conv1 = weight_norm(nn.Conv1d(audio_channels, encoder_hidden_size, kernel_size=7, padding=3)) - - self.block = [] - # Create EncoderBlocks that double channels as they downsample by `stride` - for stride_index, stride in enumerate(strides): - self.block += [ - OobleckEncoderBlock( - input_dim=encoder_hidden_size * channel_multiples[stride_index], - output_dim=encoder_hidden_size * channel_multiples[stride_index + 1], - stride=stride, - ) - ] - - self.block = nn.ModuleList(self.block) - d_model = encoder_hidden_size * channel_multiples[-1] - self.snake1 = Snake1d(d_model) - self.conv2 = weight_norm(nn.Conv1d(d_model, encoder_hidden_size, kernel_size=3, padding=1)) - - def forward(self, hidden_state): - hidden_state = self.conv1(hidden_state) - - for module in self.block: - hidden_state = module(hidden_state) - - hidden_state = self.snake1(hidden_state) - hidden_state = self.conv2(hidden_state) - - return hidden_state - - -class OobleckDecoder(nn.Module): - """Oobleck Decoder""" - - def __init__(self, channels, input_channels, audio_channels, upsampling_ratios, channel_multiples): - super().__init__() - - strides = upsampling_ratios - channel_multiples = [1] + channel_multiples - - # Add first conv layer - self.conv1 = weight_norm(nn.Conv1d(input_channels, channels * channel_multiples[-1], kernel_size=7, padding=3)) - - # Add upsampling + MRF blocks - block = [] - for stride_index, stride in enumerate(strides): - block += [ - OobleckDecoderBlock( - input_dim=channels * channel_multiples[len(strides) - stride_index], - output_dim=channels * channel_multiples[len(strides) - stride_index - 1], - stride=stride, - ) - ] - - self.block = nn.ModuleList(block) - output_dim = channels - self.snake1 = Snake1d(output_dim) - self.conv2 = weight_norm(nn.Conv1d(channels, audio_channels, kernel_size=7, padding=3, bias=False)) - - def forward(self, hidden_state): - hidden_state = self.conv1(hidden_state) - - for layer in self.block: - hidden_state = layer(hidden_state) - - hidden_state = self.snake1(hidden_state) - hidden_state = self.conv2(hidden_state) - - return hidden_state - - -class AutoencoderOobleck(ModelMixin, ConfigMixin): - r""" - An autoencoder for encoding waveforms into latents and decoding latent representations into waveforms. First - introduced in Stable Audio. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - 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]`): - 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]`): - Multiples used to determine the hidden sizes of the hidden layers. - decoder_channels (`int`, *optional*, defaults to 128): - Intermediate representation dimension for the decoder. - decoder_input_channels (`int`, *optional*, defaults to 64): - Input dimension for the decoder. Corresponds to the latent dimension. - audio_channels (`int`, *optional*, defaults to 2): - Number of channels in the audio data. Either 1 for mono or 2 for stereo. - sampling_rate (`int`, *optional*, defaults to 44100): - The sampling rate at which the audio waveform should be digitalized expressed in hertz (Hz). - """ - - _supports_gradient_checkpointing = False - _supports_group_offloading = False - - @register_to_config - def __init__( - self, - encoder_hidden_size=128, - downsampling_ratios=[2, 4, 4, 8, 8], - channel_multiples=[1, 2, 4, 8, 16], - decoder_channels=128, - decoder_input_channels=64, - audio_channels=2, - sampling_rate=44100, - ): - super().__init__() - - self.encoder_hidden_size = encoder_hidden_size - self.downsampling_ratios = downsampling_ratios - self.decoder_channels = decoder_channels - self.upsampling_ratios = downsampling_ratios[::-1] - self.hop_length = int(np.prod(downsampling_ratios)) - self.sampling_rate = sampling_rate - - self.encoder = OobleckEncoder( - encoder_hidden_size=encoder_hidden_size, - audio_channels=audio_channels, - downsampling_ratios=downsampling_ratios, - channel_multiples=channel_multiples, - ) - - self.decoder = OobleckDecoder( - channels=decoder_channels, - input_channels=decoder_input_channels, - audio_channels=audio_channels, - upsampling_ratios=self.upsampling_ratios, - channel_multiples=channel_multiples, - ) - - self.use_slicing = False - - def enable_slicing(self): - 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): - 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 - - @apply_forward_hook - def encode( - self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderOobleckOutput, Tuple[OobleckDiagonalGaussianDistribution]]: - """ - 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 images. 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.encoder(x_slice) for x_slice in x.split(1)] - h = torch.cat(encoded_slices) - else: - h = self.encoder(x) - - posterior = OobleckDiagonalGaussianDistribution(h) - - if not return_dict: - return (posterior,) - - return AutoencoderOobleckOutput(latent_dist=posterior) - - def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[OobleckDecoderOutput, torch.Tensor]: - dec = self.decoder(z) - - if not return_dict: - return (dec,) - - return OobleckDecoderOutput(sample=dec) - - @apply_forward_hook - def decode( - self, z: torch.FloatTensor, return_dict: bool = True, generator=None - ) -> Union[OobleckDecoderOutput, torch.FloatTensor]: - """ - 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.OobleckDecoderOutput`] instead of a plain tuple. - - Returns: - [`~models.vae.OobleckDecoderOutput`] or `tuple`: - If return_dict is True, a [`~models.vae.OobleckDecoderOutput`] is returned, otherwise a plain `tuple` - is returned. - - """ - 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) - else: - decoded = self._decode(z).sample - - if not return_dict: - return (decoded,) - - return OobleckDecoderOutput(sample=decoded) - - def forward( - self, - sample: torch.Tensor, - sample_posterior: bool = False, - return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[OobleckDecoderOutput, torch.Tensor]: - r""" - Args: - sample (`torch.Tensor`): Input sample. - sample_posterior (`bool`, *optional*, defaults to `False`): - Whether to sample from the posterior. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`OobleckDecoderOutput`] instead of a plain tuple. - """ - x = sample - posterior = self.encode(x).latent_dist - if sample_posterior: - z = posterior.sample(generator=generator) - else: - z = posterior.mode() - dec = self.decode(z).sample - - if not return_dict: - return (dec,) - - return OobleckDecoderOutput(sample=dec) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_tiny.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_tiny.py deleted file mode 100644 index 3e2b28606e29bb5ee4024e6b705e0cfe0e68ab45..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_tiny.py +++ /dev/null @@ -1,346 +0,0 @@ -# Copyright 2025 Ollin Boer Bohan 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 dataclasses import dataclass -from typing import Optional, Tuple, Union - -import torch - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import BaseOutput -from ...utils.accelerate_utils import apply_forward_hook -from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, DecoderTiny, EncoderTiny - - -@dataclass -class AutoencoderTinyOutput(BaseOutput): - """ - Output of AutoencoderTiny encoding method. - - Args: - latents (`torch.Tensor`): Encoded outputs of the `Encoder`. - - """ - - latents: torch.Tensor - - -class AutoencoderTiny(ModelMixin, ConfigMixin): - r""" - A tiny distilled VAE model for encoding images into latents and decoding latent representations into images. - - [`AutoencoderTiny`] is a wrapper around the original implementation of `TAESD`. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for its generic methods implemented for - all models (such as downloading or saving). - - 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 - 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 - 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. - latent_channels (`int`, *optional*, defaults to 4): - Number of channels in the latent representation. The latent space acts as a compressed representation of - the input image. - 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 - 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 - 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): - Magnitude of the latent representation. This parameter scales the latent representation values to control - the extent of information preservation. - latent_shift (float, *optional*, defaults to 0.5): - Shift applied to the latent representation. This parameter controls the center of the latent space. - 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. For this - Autoencoder, however, no such scaling factor was used, hence the value of 1.0 as the default. - force_upcast (`bool`, *optional*, default to `False`): - If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE - can be fine-tuned / trained to a lower range without losing too much precision, in which case - `force_upcast` can be set to `False` (see this fp16-friendly - [AutoEncoder](https://huggingface.co/madebyollin/sdxl-vae-fp16-fix)). - """ - - _supports_gradient_checkpointing = True - - @register_to_config - 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), - 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), - latent_magnitude: int = 3, - latent_shift: float = 0.5, - force_upcast: bool = False, - scaling_factor: float = 1.0, - shift_factor: float = 0.0, - ): - super().__init__() - - if len(encoder_block_out_channels) != len(num_encoder_blocks): - raise ValueError("`encoder_block_out_channels` should have the same length as `num_encoder_blocks`.") - if len(decoder_block_out_channels) != len(num_decoder_blocks): - raise ValueError("`decoder_block_out_channels` should have the same length as `num_decoder_blocks`.") - - self.encoder = EncoderTiny( - in_channels=in_channels, - out_channels=latent_channels, - num_blocks=num_encoder_blocks, - block_out_channels=encoder_block_out_channels, - act_fn=act_fn, - ) - - self.decoder = DecoderTiny( - in_channels=latent_channels, - out_channels=out_channels, - num_blocks=num_decoder_blocks, - block_out_channels=decoder_block_out_channels, - upsampling_scaling_factor=upsampling_scaling_factor, - act_fn=act_fn, - upsample_fn=upsample_fn, - ) - - self.latent_magnitude = latent_magnitude - self.latent_shift = latent_shift - self.scaling_factor = scaling_factor - - self.use_slicing = False - self.use_tiling = False - - # only relevant if vae tiling is enabled - self.spatial_scale_factor = 2**out_channels - self.tile_overlap_factor = 0.125 - self.tile_sample_min_size = 512 - self.tile_latent_min_size = self.tile_sample_min_size // self.spatial_scale_factor - - self.register_to_config(block_out_channels=decoder_block_out_channels) - self.register_to_config(force_upcast=False) - - def scale_latents(self, x: torch.Tensor) -> torch.Tensor: - """raw latents -> [0, 1]""" - return x.div(2 * self.latent_magnitude).add(self.latent_shift).clamp(0, 1) - - def unscale_latents(self, x: torch.Tensor) -> torch.Tensor: - """[0, 1] -> raw latents""" - return x.sub(self.latent_shift).mul(2 * self.latent_magnitude) - - 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 enable_tiling(self, use_tiling: bool = True) -> 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. - """ - self.use_tiling = use_tiling - - 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.enable_tiling(False) - - def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor: - 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 - steps. This is useful to keep memory use constant regardless of image size. To avoid tiling artifacts, the - tiles overlap and are blended together to form a smooth output. - - Args: - x (`torch.Tensor`): Input batch of images. - - Returns: - `torch.Tensor`: Encoded batch of images. - """ - # scale of encoder output relative to input - sf = self.spatial_scale_factor - tile_size = self.tile_sample_min_size - - # number of pixels to blend and to traverse between tile - blend_size = int(tile_size * self.tile_overlap_factor) - traverse_size = tile_size - blend_size - - # tiles index (up/left) - ti = range(0, x.shape[-2], traverse_size) - tj = range(0, x.shape[-1], traverse_size) - - # mask for blending - blend_masks = torch.stack( - torch.meshgrid([torch.arange(tile_size / sf) / (blend_size / sf - 1)] * 2, indexing="ij") - ) - blend_masks = blend_masks.clamp(0, 1).to(x.device) - - # output array - out = torch.zeros(x.shape[0], 4, x.shape[-2] // sf, x.shape[-1] // sf, device=x.device) - for i in ti: - for j in tj: - tile_in = x[..., i : i + tile_size, j : j + tile_size] - # tile result - tile_out = out[..., i // sf : (i + tile_size) // sf, j // sf : (j + tile_size) // sf] - tile = self.encoder(tile_in) - h, w = tile.shape[-2], tile.shape[-1] - # blend tile result into output - blend_mask_i = torch.ones_like(blend_masks[0]) if i == 0 else blend_masks[0] - blend_mask_j = torch.ones_like(blend_masks[1]) if j == 0 else blend_masks[1] - blend_mask = blend_mask_i * blend_mask_j - tile, blend_mask = tile[..., :h, :w], blend_mask[..., :h, :w] - tile_out.copy_(blend_mask * tile + (1 - blend_mask) * tile_out) - return out - - def _tiled_decode(self, x: torch.Tensor) -> torch.Tensor: - 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 - steps. This is useful to keep memory use constant regardless of image size. To avoid tiling artifacts, the - tiles overlap and are blended together to form a smooth output. - - Args: - x (`torch.Tensor`): Input batch of images. - - Returns: - `torch.Tensor`: Encoded batch of images. - """ - # scale of decoder output relative to input - sf = self.spatial_scale_factor - tile_size = self.tile_latent_min_size - - # number of pixels to blend and to traverse between tiles - blend_size = int(tile_size * self.tile_overlap_factor) - traverse_size = tile_size - blend_size - - # tiles index (up/left) - ti = range(0, x.shape[-2], traverse_size) - tj = range(0, x.shape[-1], traverse_size) - - # mask for blending - blend_masks = torch.stack( - torch.meshgrid([torch.arange(tile_size * sf) / (blend_size * sf - 1)] * 2, indexing="ij") - ) - blend_masks = blend_masks.clamp(0, 1).to(x.device) - - # output array - out = torch.zeros(x.shape[0], 3, x.shape[-2] * sf, x.shape[-1] * sf, device=x.device) - for i in ti: - for j in tj: - tile_in = x[..., i : i + tile_size, j : j + tile_size] - # tile result - tile_out = out[..., i * sf : (i + tile_size) * sf, j * sf : (j + tile_size) * sf] - tile = self.decoder(tile_in) - h, w = tile.shape[-2], tile.shape[-1] - # blend tile result into output - blend_mask_i = torch.ones_like(blend_masks[0]) if i == 0 else blend_masks[0] - blend_mask_j = torch.ones_like(blend_masks[1]) if j == 0 else blend_masks[1] - blend_mask = (blend_mask_i * blend_mask_j)[..., :h, :w] - tile_out.copy_(blend_mask * tile + (1 - blend_mask) * tile_out) - return out - - @apply_forward_hook - def encode(self, x: torch.Tensor, return_dict: bool = True) -> Union[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) - ] - output = torch.cat(output) - else: - output = self._tiled_encode(x) if self.use_tiling else self.encoder(x) - - if not return_dict: - return (output,) - - return AutoencoderTinyOutput(latents=output) - - @apply_forward_hook - def decode( - self, x: torch.Tensor, generator: Optional[torch.Generator] = None, return_dict: bool = True - ) -> Union[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) - ] - output = torch.cat(output) - else: - output = self._tiled_decode(x) if self.use_tiling else self.decoder(x) - - if not return_dict: - return (output,) - - return DecoderOutput(sample=output) - - def forward( - self, - sample: torch.Tensor, - return_dict: bool = True, - ) -> Union[DecoderOutput, Tuple[torch.Tensor]]: - r""" - Args: - sample (`torch.Tensor`): Input sample. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`DecoderOutput`] instead of a plain tuple. - """ - enc = self.encode(sample).latents - - # scale latents to be in [0, 1], then quantize latents to a byte tensor, - # as if we were storing the latents in an RGBA uint8 image. - scaled_enc = self.scale_latents(enc).mul_(255).round_().byte() - - # unquantize latents back into [0, 1], then unscale latents back to their original range, - # as if we were loading the latents from an RGBA uint8 image. - unscaled_enc = self.unscale_latents(scaled_enc / 255.0) - - dec = self.decode(unscaled_enc).sample - - if not return_dict: - return (dec,) - return DecoderOutput(sample=dec) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/consistency_decoder_vae.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/consistency_decoder_vae.py deleted file mode 100644 index b3017a878092eba21d07cf78a6595565439def41..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/consistency_decoder_vae.py +++ /dev/null @@ -1,462 +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 dataclasses import dataclass -from typing import Dict, Optional, Tuple, Union - -import torch -import torch.nn.functional as F -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...schedulers import ConsistencyDecoderScheduler -from ...utils import BaseOutput -from ...utils.accelerate_utils import apply_forward_hook -from ...utils.torch_utils import randn_tensor -from ..attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - AttentionProcessor, - AttnAddedKVProcessor, - AttnProcessor, -) -from ..modeling_utils import ModelMixin -from ..unets.unet_2d import UNet2DModel -from .vae import DecoderOutput, DiagonalGaussianDistribution, Encoder - - -@dataclass -class ConsistencyDecoderVAEOutput(BaseOutput): - """ - Output of encoding method. - - Args: - latent_dist (`DiagonalGaussianDistribution`): - Encoded outputs of `Encoder` represented as the mean and logvar of `DiagonalGaussianDistribution`. - `DiagonalGaussianDistribution` allows for sampling latents from the distribution. - """ - - latent_dist: "DiagonalGaussianDistribution" - - -class ConsistencyDecoderVAE(ModelMixin, ConfigMixin): - r""" - The consistency decoder used with DALL-E 3. - - Examples: - ```py - >>> import torch - >>> from diffusers import StableDiffusionPipeline, ConsistencyDecoderVAE - - >>> vae = ConsistencyDecoderVAE.from_pretrained("openai/consistency-decoder", torch_dtype=torch.float16) - >>> pipe = StableDiffusionPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", vae=vae, torch_dtype=torch.float16 - ... ).to("cuda") - - >>> image = pipe("horse", generator=torch.manual_seed(0)).images[0] - >>> image - ``` - """ - - _supports_group_offloading = False - - @register_to_config - def __init__( - self, - scaling_factor: float = 0.18215, - latent_channels: int = 4, - sample_size: int = 32, - encoder_act_fn: str = "silu", - encoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), - encoder_double_z: bool = True, - encoder_down_block_types: Tuple[str, ...] = ( - "DownEncoderBlock2D", - "DownEncoderBlock2D", - "DownEncoderBlock2D", - "DownEncoderBlock2D", - ), - encoder_in_channels: int = 3, - encoder_layers_per_block: int = 2, - 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, ...] = ( - "ResnetDownsampleBlock2D", - "ResnetDownsampleBlock2D", - "ResnetDownsampleBlock2D", - "ResnetDownsampleBlock2D", - ), - decoder_downsample_padding: int = 1, - decoder_in_channels: int = 7, - decoder_layers_per_block: int = 3, - decoder_norm_eps: float = 1e-05, - decoder_norm_num_groups: int = 32, - decoder_num_train_timesteps: int = 1024, - 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, ...] = ( - "ResnetUpsampleBlock2D", - "ResnetUpsampleBlock2D", - "ResnetUpsampleBlock2D", - "ResnetUpsampleBlock2D", - ), - ): - super().__init__() - self.encoder = Encoder( - act_fn=encoder_act_fn, - block_out_channels=encoder_block_out_channels, - double_z=encoder_double_z, - down_block_types=encoder_down_block_types, - in_channels=encoder_in_channels, - layers_per_block=encoder_layers_per_block, - norm_num_groups=encoder_norm_num_groups, - out_channels=encoder_out_channels, - ) - - self.decoder_unet = UNet2DModel( - add_attention=decoder_add_attention, - block_out_channels=decoder_block_out_channels, - down_block_types=decoder_down_block_types, - downsample_padding=decoder_downsample_padding, - in_channels=decoder_in_channels, - layers_per_block=decoder_layers_per_block, - norm_eps=decoder_norm_eps, - norm_num_groups=decoder_norm_num_groups, - num_train_timesteps=decoder_num_train_timesteps, - out_channels=decoder_out_channels, - resnet_time_scale_shift=decoder_resnet_time_scale_shift, - time_embedding_type=decoder_time_embedding_type, - up_block_types=decoder_up_block_types, - ) - self.decoder_scheduler = ConsistencyDecoderScheduler() - self.register_to_config(block_out_channels=encoder_block_out_channels) - self.register_to_config(force_upcast=False) - self.register_buffer( - "means", - torch.tensor([0.38862467, 0.02253063, 0.07381133, -0.0171294])[None, :, None, None], - persistent=False, - ) - self.register_buffer( - "stds", torch.tensor([0.9654121, 1.0440036, 0.76147926, 0.77022034])[None, :, None, None], persistent=False - ) - - self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) - - self.use_slicing = False - self.use_tiling = False - - # only relevant if vae tiling is enabled - self.tile_sample_min_size = self.config.sample_size - sample_size = ( - self.config.sample_size[0] - if isinstance(self.config.sample_size, (list, tuple)) - else self.config.sample_size - ) - self.tile_latent_min_size = int(sample_size / (2 ** (len(self.config.block_out_channels) - 1))) - self.tile_overlap_factor = 0.25 - - # Copied from diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL.enable_tiling - def enable_tiling(self, use_tiling: bool = True): - 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. - """ - self.use_tiling = use_tiling - - # Copied from diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL.disable_tiling - def disable_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.enable_tiling(False) - - # Copied from diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL.enable_slicing - def enable_slicing(self): - 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 - - # Copied from diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL.disable_slicing - def disable_slicing(self): - 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 - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - self.set_attn_processor(processor) - - @apply_forward_hook - def encode( - self, x: torch.Tensor, return_dict: bool = True - ) -> Union[ConsistencyDecoderVAEOutput, 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.autoencoders.consistency_decoder_vae.ConsistencyDecoderVAEOutput`] - instead of a plain tuple. - - Returns: - The latent representations of the encoded images. If `return_dict` is True, a - [`~models.autoencoders.consistency_decoder_vae.ConsistencyDecoderVAEOutput`] is returned, otherwise a - plain `tuple` is returned. - """ - if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size): - return self.tiled_encode(x, return_dict=return_dict) - - if self.use_slicing and x.shape[0] > 1: - encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)] - h = torch.cat(encoded_slices) - else: - h = self.encoder(x) - - moments = self.quant_conv(h) - posterior = DiagonalGaussianDistribution(moments) - - if not return_dict: - return (posterior,) - - return ConsistencyDecoderVAEOutput(latent_dist=posterior) - - @apply_forward_hook - def decode( - self, - z: torch.Tensor, - generator: Optional[torch.Generator] = None, - return_dict: bool = True, - num_inference_steps: int = 2, - ) -> Union[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. - 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. - - """ - z = (z * self.config.scaling_factor - self.means) / self.stds - - scale_factor = 2 ** (len(self.config.block_out_channels) - 1) - z = F.interpolate(z, mode="nearest", scale_factor=scale_factor) - - batch_size, _, height, width = z.shape - - self.decoder_scheduler.set_timesteps(num_inference_steps, device=self.device) - - x_t = self.decoder_scheduler.init_noise_sigma * randn_tensor( - (batch_size, 3, height, width), generator=generator, dtype=z.dtype, device=z.device - ) - - for t in self.decoder_scheduler.timesteps: - model_input = torch.concat([self.decoder_scheduler.scale_model_input(x_t, t), z], dim=1) - model_output = self.decoder_unet(model_input, t).sample[:, :3, :, :] - prev_sample = self.decoder_scheduler.step(model_output, t, x_t, generator).prev_sample - x_t = prev_sample - - x_0 = x_t - - if not return_dict: - return (x_0,) - - return DecoderOutput(sample=x_0) - - # Copied from diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL.blend_v - def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: - blend_extent = min(a.shape[2], b.shape[2], 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 - - # Copied from diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL.blend_h - def blend_h(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, return_dict: bool = True) -> Union[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 - steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is - different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the - tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the - output, but they should be much less noticeable. - - Args: - x (`torch.Tensor`): Input batch of images. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.autoencoders.consistency_decoder_vae.ConsistencyDecoderVAEOutput`] - instead of a plain tuple. - - Returns: - [`~models.autoencoders.consistency_decoder_vae.ConsistencyDecoderVAEOutput`] or `tuple`: - If return_dict is True, a [`~models.autoencoders.consistency_decoder_vae.ConsistencyDecoderVAEOutput`] - is returned, otherwise a plain `tuple` is returned. - """ - overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) - blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) - row_limit = self.tile_latent_min_size - blend_extent - - # Split the image into 512x512 tiles and encode them separately. - rows = [] - for i in range(0, x.shape[2], overlap_size): - row = [] - for j in range(0, x.shape[3], overlap_size): - tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size] - tile = self.encoder(tile) - tile = self.quant_conv(tile) - row.append(tile) - 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_extent) - if j > 0: - tile = self.blend_h(row[j - 1], tile, blend_extent) - result_row.append(tile[:, :, :row_limit, :row_limit]) - result_rows.append(torch.cat(result_row, dim=3)) - - moments = torch.cat(result_rows, dim=2) - posterior = DiagonalGaussianDistribution(moments) - - if not return_dict: - return (posterior,) - - return ConsistencyDecoderVAEOutput(latent_dist=posterior) - - def forward( - self, - sample: torch.Tensor, - sample_posterior: bool = False, - return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, Tuple[torch.Tensor]]: - r""" - Args: - sample (`torch.Tensor`): Input sample. - sample_posterior (`bool`, *optional*, defaults to `False`): - Whether to sample from the posterior. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`DecoderOutput`] instead of a plain tuple. - generator (`torch.Generator`, *optional*, defaults to `None`): - Generator to use for sampling. - - Returns: - [`DecoderOutput`] or `tuple`: - If return_dict is True, a [`DecoderOutput`] is returned, otherwise a plain `tuple` is returned. - """ - x = sample - posterior = self.encode(x).latent_dist - if sample_posterior: - z = posterior.sample(generator=generator) - else: - z = posterior.mode() - dec = self.decode(z, generator=generator).sample - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/vae.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/vae.py deleted file mode 100644 index 1d74d4f472d7f02a78cadd9e610cb40c91078473..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/vae.py +++ /dev/null @@ -1,896 +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 dataclasses import dataclass -from typing import Optional, Tuple - -import numpy as np -import torch -import torch.nn as nn - -from ...utils import BaseOutput -from ...utils.torch_utils import randn_tensor -from ..activations import get_activation -from ..attention_processor import SpatialNorm -from ..unets.unet_2d_blocks import ( - AutoencoderTinyBlock, - UNetMidBlock2D, - get_down_block, - get_up_block, -) - - -@dataclass -class EncoderOutput(BaseOutput): - r""" - Output of encoding method. - - Args: - latent (`torch.Tensor` of shape `(batch_size, num_channels, latent_height, latent_width)`): - The encoded latent. - """ - - latent: torch.Tensor - - -@dataclass -class DecoderOutput(BaseOutput): - r""" - Output of decoding method. - - Args: - sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)`): - The decoded output sample from the last layer of the model. - """ - - sample: torch.Tensor - commit_loss: Optional[torch.FloatTensor] = None - - -class Encoder(nn.Module): - r""" - The `Encoder` layer of a variational autoencoder that encodes its input into a latent representation. - - Args: - in_channels (`int`, *optional*, defaults to 3): - 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",)`): - 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,)`): - The number 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 for normalization. - act_fn (`str`, *optional*, defaults to `"silu"`): - The activation function to use. See `~diffusers.models.activations.get_activation` for available options. - double_z (`bool`, *optional*, defaults to `True`): - Whether to double the number of output channels for the last block. - """ - - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - 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", - double_z: bool = True, - mid_block_add_attention=True, - ): - super().__init__() - self.layers_per_block = layers_per_block - - self.conv_in = nn.Conv2d( - in_channels, - block_out_channels[0], - kernel_size=3, - stride=1, - padding=1, - ) - - self.down_blocks = nn.ModuleList([]) - - # down - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = get_down_block( - down_block_type, - num_layers=self.layers_per_block, - in_channels=input_channel, - out_channels=output_channel, - add_downsample=not is_final_block, - resnet_eps=1e-6, - downsample_padding=0, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - attention_head_dim=output_channel, - temb_channels=None, - ) - self.down_blocks.append(down_block) - - # mid - self.mid_block = UNetMidBlock2D( - in_channels=block_out_channels[-1], - resnet_eps=1e-6, - resnet_act_fn=act_fn, - output_scale_factor=1, - resnet_time_scale_shift="default", - attention_head_dim=block_out_channels[-1], - resnet_groups=norm_num_groups, - temb_channels=None, - add_attention=mid_block_add_attention, - ) - - # out - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) - self.conv_act = nn.SiLU() - - conv_out_channels = 2 * out_channels if double_z else out_channels - self.conv_out = nn.Conv2d(block_out_channels[-1], conv_out_channels, 3, padding=1) - - self.gradient_checkpointing = False - - def forward(self, sample: torch.Tensor) -> torch.Tensor: - r"""The forward method of the `Encoder` class.""" - - sample = self.conv_in(sample) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - # down - for down_block in self.down_blocks: - sample = self._gradient_checkpointing_func(down_block, sample) - # middle - sample = self._gradient_checkpointing_func(self.mid_block, sample) - - else: - # down - for down_block in self.down_blocks: - sample = down_block(sample) - - # middle - sample = self.mid_block(sample) - - # post-process - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - sample = self.conv_out(sample) - - return sample - - -class Decoder(nn.Module): - r""" - The `Decoder` layer of a variational autoencoder that decodes its latent representation into an output sample. - - Args: - in_channels (`int`, *optional*, defaults to 3): - 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",)`): - 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,)`): - The number 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 for normalization. - act_fn (`str`, *optional*, defaults to `"silu"`): - The activation function to use. See `~diffusers.models.activations.get_activation` for available options. - norm_type (`str`, *optional*, defaults to `"group"`): - The normalization type to use. Can be either `"group"` or `"spatial"`. - """ - - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - 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", - norm_type: str = "group", # group, spatial - mid_block_add_attention=True, - ): - super().__init__() - self.layers_per_block = layers_per_block - - self.conv_in = nn.Conv2d( - in_channels, - block_out_channels[-1], - kernel_size=3, - stride=1, - padding=1, - ) - - self.up_blocks = nn.ModuleList([]) - - temb_channels = in_channels if norm_type == "spatial" else None - - # mid - self.mid_block = UNetMidBlock2D( - in_channels=block_out_channels[-1], - resnet_eps=1e-6, - resnet_act_fn=act_fn, - output_scale_factor=1, - resnet_time_scale_shift="default" if norm_type == "group" else norm_type, - attention_head_dim=block_out_channels[-1], - resnet_groups=norm_num_groups, - temb_channels=temb_channels, - add_attention=mid_block_add_attention, - ) - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - - is_final_block = i == len(block_out_channels) - 1 - - up_block = get_up_block( - up_block_type, - num_layers=self.layers_per_block + 1, - in_channels=prev_output_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - add_upsample=not is_final_block, - resnet_eps=1e-6, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - attention_head_dim=output_channel, - temb_channels=temb_channels, - resnet_time_scale_shift=norm_type, - ) - self.up_blocks.append(up_block) - prev_output_channel = output_channel - - # out - if norm_type == "spatial": - self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels) - else: - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) - self.conv_act = nn.SiLU() - self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1) - - self.gradient_checkpointing = False - - def forward( - self, - sample: torch.Tensor, - latent_embeds: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - r"""The forward method of the `Decoder` class.""" - - sample = self.conv_in(sample) - - upscale_dtype = next(iter(self.up_blocks.parameters())).dtype - if torch.is_grad_enabled() and self.gradient_checkpointing: - # middle - sample = self._gradient_checkpointing_func(self.mid_block, sample, latent_embeds) - sample = sample.to(upscale_dtype) - - # up - for up_block in self.up_blocks: - sample = self._gradient_checkpointing_func(up_block, sample, latent_embeds) - else: - # middle - sample = self.mid_block(sample, latent_embeds) - sample = sample.to(upscale_dtype) - - # up - for up_block in self.up_blocks: - sample = up_block(sample, latent_embeds) - - # post-process - if latent_embeds is None: - sample = self.conv_norm_out(sample) - else: - sample = self.conv_norm_out(sample, latent_embeds) - sample = self.conv_act(sample) - sample = self.conv_out(sample) - - return sample - - -class UpSample(nn.Module): - r""" - The `UpSample` layer of a variational autoencoder that upsamples its input. - - Args: - in_channels (`int`, *optional*, defaults to 3): - The number of input channels. - out_channels (`int`, *optional*, defaults to 3): - The number of output channels. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - ) -> None: - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.deconv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - r"""The forward method of the `UpSample` class.""" - x = torch.relu(x) - x = self.deconv(x) - return x - - -class MaskConditionEncoder(nn.Module): - """ - used in AsymmetricAutoencoderKL - """ - - def __init__( - self, - in_ch: int, - out_ch: int = 192, - res_ch: int = 768, - stride: int = 16, - ) -> None: - super().__init__() - - channels = [] - while stride > 1: - stride = stride // 2 - in_ch_ = out_ch * 2 - if out_ch > res_ch: - out_ch = res_ch - if stride == 1: - in_ch_ = res_ch - channels.append((in_ch_, out_ch)) - out_ch *= 2 - - out_channels = [] - for _in_ch, _out_ch in channels: - out_channels.append(_out_ch) - out_channels.append(channels[-1][0]) - - layers = [] - in_ch_ = in_ch - for l in range(len(out_channels)): - out_ch_ = out_channels[l] - if l == 0 or l == 1: - layers.append(nn.Conv2d(in_ch_, out_ch_, kernel_size=3, stride=1, padding=1)) - else: - layers.append(nn.Conv2d(in_ch_, out_ch_, kernel_size=4, stride=2, padding=1)) - in_ch_ = out_ch_ - - self.layers = nn.Sequential(*layers) - - def forward(self, x: torch.Tensor, mask=None) -> torch.Tensor: - r"""The forward method of the `MaskConditionEncoder` class.""" - out = {} - for l in range(len(self.layers)): - layer = self.layers[l] - x = layer(x) - out[str(tuple(x.shape))] = x - x = torch.relu(x) - return out - - -class MaskConditionDecoder(nn.Module): - r"""The `MaskConditionDecoder` should be used in combination with [`AsymmetricAutoencoderKL`] to enhance the model's - decoder with a conditioner on the mask and masked image. - - Args: - in_channels (`int`, *optional*, defaults to 3): - 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",)`): - 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,)`): - The number 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 for normalization. - act_fn (`str`, *optional*, defaults to `"silu"`): - The activation function to use. See `~diffusers.models.activations.get_activation` for available options. - norm_type (`str`, *optional*, defaults to `"group"`): - The normalization type to use. Can be either `"group"` or `"spatial"`. - """ - - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - 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", - norm_type: str = "group", # group, spatial - ): - super().__init__() - self.layers_per_block = layers_per_block - - self.conv_in = nn.Conv2d( - in_channels, - block_out_channels[-1], - kernel_size=3, - stride=1, - padding=1, - ) - - self.up_blocks = nn.ModuleList([]) - - temb_channels = in_channels if norm_type == "spatial" else None - - # mid - self.mid_block = UNetMidBlock2D( - in_channels=block_out_channels[-1], - resnet_eps=1e-6, - resnet_act_fn=act_fn, - output_scale_factor=1, - resnet_time_scale_shift="default" if norm_type == "group" else norm_type, - attention_head_dim=block_out_channels[-1], - resnet_groups=norm_num_groups, - temb_channels=temb_channels, - ) - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - - is_final_block = i == len(block_out_channels) - 1 - - up_block = get_up_block( - up_block_type, - num_layers=self.layers_per_block + 1, - in_channels=prev_output_channel, - out_channels=output_channel, - prev_output_channel=None, - add_upsample=not is_final_block, - resnet_eps=1e-6, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - attention_head_dim=output_channel, - temb_channels=temb_channels, - resnet_time_scale_shift=norm_type, - ) - self.up_blocks.append(up_block) - prev_output_channel = output_channel - - # condition encoder - self.condition_encoder = MaskConditionEncoder( - in_ch=out_channels, - out_ch=block_out_channels[0], - res_ch=block_out_channels[-1], - ) - - # out - if norm_type == "spatial": - self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels) - else: - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) - self.conv_act = nn.SiLU() - self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1) - - self.gradient_checkpointing = False - - def forward( - self, - z: torch.Tensor, - image: Optional[torch.Tensor] = None, - mask: Optional[torch.Tensor] = None, - latent_embeds: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - r"""The forward method of the `MaskConditionDecoder` class.""" - sample = z - sample = self.conv_in(sample) - - upscale_dtype = next(iter(self.up_blocks.parameters())).dtype - if torch.is_grad_enabled() and self.gradient_checkpointing: - # middle - sample = self._gradient_checkpointing_func(self.mid_block, sample, latent_embeds) - sample = sample.to(upscale_dtype) - - # condition encoder - if image is not None and mask is not None: - masked_image = (1 - mask) * image - im_x = self._gradient_checkpointing_func( - self.condition_encoder, - masked_image, - mask, - ) - - # up - for up_block in self.up_blocks: - if image is not None and mask is not None: - sample_ = im_x[str(tuple(sample.shape))] - mask_ = nn.functional.interpolate(mask, size=sample.shape[-2:], mode="nearest") - sample = sample * mask_ + sample_ * (1 - mask_) - sample = self._gradient_checkpointing_func(up_block, sample, latent_embeds) - if image is not None and mask is not None: - sample = sample * mask + im_x[str(tuple(sample.shape))] * (1 - mask) - else: - # middle - sample = self.mid_block(sample, latent_embeds) - sample = sample.to(upscale_dtype) - - # condition encoder - if image is not None and mask is not None: - masked_image = (1 - mask) * image - im_x = self.condition_encoder(masked_image, mask) - - # up - for up_block in self.up_blocks: - if image is not None and mask is not None: - sample_ = im_x[str(tuple(sample.shape))] - mask_ = nn.functional.interpolate(mask, size=sample.shape[-2:], mode="nearest") - sample = sample * mask_ + sample_ * (1 - mask_) - sample = up_block(sample, latent_embeds) - if image is not None and mask is not None: - sample = sample * mask + im_x[str(tuple(sample.shape))] * (1 - mask) - - # post-process - if latent_embeds is None: - sample = self.conv_norm_out(sample) - else: - sample = self.conv_norm_out(sample, latent_embeds) - sample = self.conv_act(sample) - sample = self.conv_out(sample) - - return sample - - -class VectorQuantizer(nn.Module): - """ - Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly avoids costly matrix - multiplications and allows for post-hoc remapping of indices. - """ - - # NOTE: due to a bug the beta term was applied to the wrong term. for - # backwards compatibility we use the buggy version by default, but you can - # specify legacy=False to fix it. - def __init__( - self, - n_e: int, - vq_embed_dim: int, - beta: float, - remap=None, - unknown_index: str = "random", - sane_index_shape: bool = False, - legacy: bool = True, - ): - super().__init__() - self.n_e = n_e - self.vq_embed_dim = vq_embed_dim - self.beta = beta - self.legacy = legacy - - self.embedding = nn.Embedding(self.n_e, self.vq_embed_dim) - self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) - - self.remap = remap - if self.remap is not None: - self.register_buffer("used", torch.tensor(np.load(self.remap))) - self.used: torch.Tensor - self.re_embed = self.used.shape[0] - self.unknown_index = unknown_index # "random" or "extra" or integer - if self.unknown_index == "extra": - self.unknown_index = self.re_embed - self.re_embed = self.re_embed + 1 - print( - f"Remapping {self.n_e} indices to {self.re_embed} indices. " - f"Using {self.unknown_index} for unknown indices." - ) - else: - self.re_embed = n_e - - self.sane_index_shape = sane_index_shape - - def remap_to_used(self, inds: torch.LongTensor) -> torch.LongTensor: - ishape = inds.shape - assert len(ishape) > 1 - inds = inds.reshape(ishape[0], -1) - used = self.used.to(inds) - match = (inds[:, :, None] == used[None, None, ...]).long() - new = match.argmax(-1) - unknown = match.sum(2) < 1 - if self.unknown_index == "random": - new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to(device=new.device) - else: - new[unknown] = self.unknown_index - return new.reshape(ishape) - - def unmap_to_all(self, inds: torch.LongTensor) -> torch.LongTensor: - ishape = inds.shape - assert len(ishape) > 1 - inds = inds.reshape(ishape[0], -1) - used = self.used.to(inds) - if self.re_embed > self.used.shape[0]: # extra token - inds[inds >= self.used.shape[0]] = 0 # simply set to zero - 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]: - # 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) - - # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z - min_encoding_indices = torch.argmin(torch.cdist(z_flattened, self.embedding.weight), dim=1) - - z_q = self.embedding(min_encoding_indices).view(z.shape) - perplexity = None - min_encodings = None - - # compute loss for embedding - if not self.legacy: - loss = self.beta * torch.mean((z_q.detach() - z) ** 2) + torch.mean((z_q - z.detach()) ** 2) - else: - loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean((z_q - z.detach()) ** 2) - - # preserve gradients - z_q: torch.Tensor = z + (z_q - z).detach() - - # reshape back to match original input shape - z_q = z_q.permute(0, 3, 1, 2).contiguous() - - if self.remap is not None: - min_encoding_indices = min_encoding_indices.reshape(z.shape[0], -1) # add batch axis - min_encoding_indices = self.remap_to_used(min_encoding_indices) - min_encoding_indices = min_encoding_indices.reshape(-1, 1) # flatten - - if self.sane_index_shape: - min_encoding_indices = min_encoding_indices.reshape(z_q.shape[0], z_q.shape[2], z_q.shape[3]) - - return z_q, loss, (perplexity, min_encodings, min_encoding_indices) - - 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 - indices = self.unmap_to_all(indices) - indices = indices.reshape(-1) # flatten again - - # get quantized latent vectors - z_q: torch.Tensor = self.embedding(indices) - - if shape is not None: - z_q = z_q.view(shape) - # reshape back to match original input shape - z_q = z_q.permute(0, 3, 1, 2).contiguous() - - return z_q - - -class DiagonalGaussianDistribution(object): - def __init__(self, parameters: torch.Tensor, deterministic: bool = False): - self.parameters = parameters - self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) - self.logvar = torch.clamp(self.logvar, -30.0, 20.0) - self.deterministic = deterministic - self.std = torch.exp(0.5 * self.logvar) - self.var = torch.exp(self.logvar) - if self.deterministic: - self.var = self.std = torch.zeros_like( - self.mean, device=self.parameters.device, dtype=self.parameters.dtype - ) - - def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor: - # make sure sample is on the same device as the parameters and has same dtype - sample = randn_tensor( - self.mean.shape, - generator=generator, - device=self.parameters.device, - dtype=self.parameters.dtype, - ) - x = self.mean + self.std * sample - return x - - def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor: - if self.deterministic: - return torch.Tensor([0.0]) - else: - if other is None: - return 0.5 * torch.sum( - torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, - dim=[1, 2, 3], - ) - else: - return 0.5 * torch.sum( - torch.pow(self.mean - other.mean, 2) / other.var - + self.var / other.var - - 1.0 - - self.logvar - + other.logvar, - dim=[1, 2, 3], - ) - - 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) - return 0.5 * torch.sum( - logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, - dim=dims, - ) - - def mode(self) -> torch.Tensor: - return self.mean - - -class IdentityDistribution(object): - def __init__(self, parameters: torch.Tensor): - self.parameters = parameters - - def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor: - return self.parameters - - def mode(self) -> torch.Tensor: - return self.parameters - - -class EncoderTiny(nn.Module): - r""" - The `EncoderTiny` layer is a simpler version of the `Encoder` layer. - - Args: - in_channels (`int`): - The number of input channels. - out_channels (`int`): - The number of output channels. - 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, ...]`): - 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. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - num_blocks: Tuple[int, ...], - block_out_channels: Tuple[int, ...], - act_fn: str, - ): - super().__init__() - - layers = [] - for i, num_block in enumerate(num_blocks): - num_channels = block_out_channels[i] - - if i == 0: - layers.append(nn.Conv2d(in_channels, num_channels, kernel_size=3, padding=1)) - else: - layers.append( - nn.Conv2d( - num_channels, - num_channels, - kernel_size=3, - padding=1, - stride=2, - bias=False, - ) - ) - - for _ in range(num_block): - layers.append(AutoencoderTinyBlock(num_channels, num_channels, act_fn)) - - layers.append(nn.Conv2d(block_out_channels[-1], out_channels, kernel_size=3, padding=1)) - - self.layers = nn.Sequential(*layers) - self.gradient_checkpointing = False - - def forward(self, x: torch.Tensor) -> torch.Tensor: - r"""The forward method of the `EncoderTiny` class.""" - if torch.is_grad_enabled() and self.gradient_checkpointing: - x = self._gradient_checkpointing_func(self.layers, x) - - else: - # scale image from [-1, 1] to [0, 1] to match TAESD convention - x = self.layers(x.add(1).div(2)) - - return x - - -class DecoderTiny(nn.Module): - r""" - The `DecoderTiny` layer is a simpler version of the `Decoder` layer. - - Args: - in_channels (`int`): - The number of input channels. - out_channels (`int`): - The number of output channels. - 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, ...]`): - The number of output channels for each block. - upsampling_scaling_factor (`int`): - The scaling factor to use for upsampling. - act_fn (`str`): - The activation function to use. See `~diffusers.models.activations.get_activation` for available options. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - num_blocks: Tuple[int, ...], - block_out_channels: Tuple[int, ...], - upsampling_scaling_factor: int, - act_fn: str, - upsample_fn: str, - ): - super().__init__() - - layers = [ - nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, padding=1), - get_activation(act_fn), - ] - - for i, num_block in enumerate(num_blocks): - is_final_block = i == (len(num_blocks) - 1) - num_channels = block_out_channels[i] - - for _ in range(num_block): - layers.append(AutoencoderTinyBlock(num_channels, num_channels, act_fn)) - - if not is_final_block: - layers.append(nn.Upsample(scale_factor=upsampling_scaling_factor, mode=upsample_fn)) - - conv_out_channel = num_channels if not is_final_block else out_channels - layers.append( - nn.Conv2d( - num_channels, - conv_out_channel, - kernel_size=3, - padding=1, - bias=is_final_block, - ) - ) - - self.layers = nn.Sequential(*layers) - self.gradient_checkpointing = False - - def forward(self, x: torch.Tensor) -> torch.Tensor: - r"""The forward method of the `DecoderTiny` class.""" - # Clamp. - x = torch.tanh(x / 3) * 3 - - if torch.is_grad_enabled() and self.gradient_checkpointing: - x = self._gradient_checkpointing_func(self.layers, x) - else: - x = self.layers(x) - - # scale image from [0, 1] to [-1, 1] to match diffusers convention - return x.mul(2).sub(1) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/vq_model.py b/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/vq_model.py deleted file mode 100644 index c1094e62f7ecbdef12026b783734957b9270c03f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/autoencoders/vq_model.py +++ /dev/null @@ -1,185 +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 dataclasses import dataclass -from typing import Optional, Tuple, Union - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import BaseOutput -from ...utils.accelerate_utils import apply_forward_hook -from ..autoencoders.vae import Decoder, DecoderOutput, Encoder, VectorQuantizer -from ..modeling_utils import ModelMixin - - -@dataclass -class VQEncoderOutput(BaseOutput): - """ - Output of VQModel encoding method. - - Args: - latents (`torch.Tensor` of shape `(batch_size, num_channels, height, width)`): - The encoded output sample from the last layer of the model. - """ - - latents: torch.Tensor - - -class VQModel(ModelMixin, ConfigMixin): - r""" - A VQ-VAE model for decoding latent representations. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - 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. - 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. - sample_size (`int`, *optional*, defaults to `32`): Sample input size. - num_vq_embeddings (`int`, *optional*, defaults to `256`): Number of codebook vectors in the VQ-VAE. - norm_num_groups (`int`, *optional*, defaults to `32`): Number of groups for normalization layers. - vq_embed_dim (`int`, *optional*): Hidden dim of codebook vectors in the VQ-VAE. - scaling_factor (`float`, *optional*, defaults to `0.18215`): - 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. - norm_type (`str`, *optional*, defaults to `"group"`): - Type of normalization layer to use. Can be one of `"group"` or `"spatial"`. - """ - - _skip_layerwise_casting_patterns = ["quantize"] - _supports_group_offloading = False - - @register_to_config - 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,), - 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, - scaling_factor: float = 0.18215, - norm_type: str = "group", # group, spatial - mid_block_add_attention=True, - lookup_from_codebook=False, - force_upcast=False, - ): - super().__init__() - - # pass init params to Encoder - self.encoder = Encoder( - in_channels=in_channels, - out_channels=latent_channels, - down_block_types=down_block_types, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - act_fn=act_fn, - norm_num_groups=norm_num_groups, - double_z=False, - mid_block_add_attention=mid_block_add_attention, - ) - - vq_embed_dim = vq_embed_dim if vq_embed_dim is not None else latent_channels - - self.quant_conv = nn.Conv2d(latent_channels, vq_embed_dim, 1) - self.quantize = VectorQuantizer(num_vq_embeddings, vq_embed_dim, beta=0.25, remap=None, sane_index_shape=False) - self.post_quant_conv = nn.Conv2d(vq_embed_dim, latent_channels, 1) - - # pass init params to Decoder - self.decoder = Decoder( - in_channels=latent_channels, - out_channels=out_channels, - up_block_types=up_block_types, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - act_fn=act_fn, - norm_num_groups=norm_num_groups, - norm_type=norm_type, - mid_block_add_attention=mid_block_add_attention, - ) - - @apply_forward_hook - def encode(self, x: torch.Tensor, return_dict: bool = True) -> VQEncoderOutput: - h = self.encoder(x) - h = self.quant_conv(h) - - if not return_dict: - return (h,) - - return VQEncoderOutput(latents=h) - - @apply_forward_hook - def decode( - self, h: torch.Tensor, force_not_quantize: bool = False, return_dict: bool = True, shape=None - ) -> Union[DecoderOutput, torch.Tensor]: - # also go through quantization layer - if not force_not_quantize: - quant, commit_loss, _ = self.quantize(h) - elif self.config.lookup_from_codebook: - quant = self.quantize.get_codebook_entry(h, shape) - commit_loss = torch.zeros((h.shape[0])).to(h.device, dtype=h.dtype) - else: - quant = h - commit_loss = torch.zeros((h.shape[0])).to(h.device, dtype=h.dtype) - quant2 = self.post_quant_conv(quant) - dec = self.decoder(quant2, quant if self.config.norm_type == "spatial" else None) - - if not return_dict: - return dec, commit_loss - - return DecoderOutput(sample=dec, commit_loss=commit_loss) - - def forward( - self, sample: torch.Tensor, return_dict: bool = True - ) -> Union[DecoderOutput, Tuple[torch.Tensor, ...]]: - r""" - The [`VQModel`] forward method. - - Args: - sample (`torch.Tensor`): Input sample. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`models.autoencoders.vq_model.VQEncoderOutput`] instead of a plain tuple. - - Returns: - [`~models.autoencoders.vq_model.VQEncoderOutput`] or `tuple`: - If return_dict is True, a [`~models.autoencoders.vq_model.VQEncoderOutput`] is returned, otherwise a - plain `tuple` is returned. - """ - - h = self.encode(sample).latents - dec = self.decode(h) - - if not return_dict: - return dec.sample, dec.commit_loss - return dec diff --git a/venv/lib/python3.10/site-packages/diffusers/models/cache_utils.py b/venv/lib/python3.10/site-packages/diffusers/models/cache_utils.py deleted file mode 100644 index 3fd1ca6e9d3d50e149e70dbd483a3109cd5bc861..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/cache_utils.py +++ /dev/null @@ -1,108 +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.logging import get_logger - - -logger = get_logger(__name__) # pylint: disable=invalid-name - - -class CacheMixin: - r""" - A class for enable/disabling caching techniques on diffusion models. - - Supported caching techniques: - - [Pyramid Attention Broadcast](https://huggingface.co/papers/2408.12588) - - [FasterCache](https://huggingface.co/papers/2410.19355) - """ - - _cache_config = None - - @property - def is_cache_enabled(self) -> bool: - return self._cache_config is not None - - def enable_cache(self, config) -> None: - r""" - Enable caching techniques on the model. - - Args: - config (`Union[PyramidAttentionBroadcastConfig]`): - The configuration for applying the caching technique. Currently supported caching techniques are: - - [`~hooks.PyramidAttentionBroadcastConfig`] - - Example: - - ```python - >>> import torch - >>> from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig - - >>> pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16) - >>> pipe.to("cuda") - - >>> config = PyramidAttentionBroadcastConfig( - ... spatial_attention_block_skip_range=2, - ... spatial_attention_timestep_skip_range=(100, 800), - ... current_timestep_callback=lambda: pipe.current_timestep, - ... ) - >>> pipe.transformer.enable_cache(config) - ``` - """ - - from ..hooks import ( - FasterCacheConfig, - PyramidAttentionBroadcastConfig, - apply_faster_cache, - apply_pyramid_attention_broadcast, - ) - - if self.is_cache_enabled: - raise ValueError( - f"Caching has already been enabled with {type(self._cache_config)}. To apply a new caching technique, please disable the existing one first." - ) - - if isinstance(config, PyramidAttentionBroadcastConfig): - apply_pyramid_attention_broadcast(self, config) - elif isinstance(config, FasterCacheConfig): - apply_faster_cache(self, config) - else: - raise ValueError(f"Cache config {type(config)} is not supported.") - - self._cache_config = config - - def disable_cache(self) -> None: - from ..hooks import FasterCacheConfig, HookRegistry, PyramidAttentionBroadcastConfig - from ..hooks.faster_cache import _FASTER_CACHE_BLOCK_HOOK, _FASTER_CACHE_DENOISER_HOOK - from ..hooks.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK - - if self._cache_config is None: - logger.warning("Caching techniques have not been enabled, so there's nothing to disable.") - return - - if isinstance(self._cache_config, PyramidAttentionBroadcastConfig): - registry = HookRegistry.check_if_exists_or_initialize(self) - registry.remove_hook(_PYRAMID_ATTENTION_BROADCAST_HOOK, recurse=True) - elif isinstance(self._cache_config, FasterCacheConfig): - registry = HookRegistry.check_if_exists_or_initialize(self) - registry.remove_hook(_FASTER_CACHE_DENOISER_HOOK, recurse=True) - registry.remove_hook(_FASTER_CACHE_BLOCK_HOOK, recurse=True) - else: - raise ValueError(f"Cache config {type(self._cache_config)} is not supported.") - - self._cache_config = None - - def _reset_stateful_cache(self, recurse: bool = True) -> None: - from ..hooks import HookRegistry - - HookRegistry.check_if_exists_or_initialize(self).reset_stateful_hooks(recurse=recurse) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnet.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnet.py deleted file mode 100644 index c18bd8751dcb84744ee44a37a922c62fd5b16caa..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/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/venv/lib/python3.10/site-packages/diffusers/models/controlnet_flux.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnet_flux.py deleted file mode 100644 index e82748436d86250180a42f9d075dc51dc074690d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/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/venv/lib/python3.10/site-packages/diffusers/models/controlnet_sd3.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnet_sd3.py deleted file mode 100644 index d239ad4eb3e891c42880a9fc6bf2c7b2ec1220d4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/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/venv/lib/python3.10/site-packages/diffusers/models/controlnet_sparsectrl.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnet_sparsectrl.py deleted file mode 100644 index 5c67af4fe9c1d2fc21c69f2ae4ff86f35fb0ea03..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/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/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__init__.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__init__.py deleted file mode 100644 index 90ef438d2533f680dfb66d90a7b7997663999251..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -from ...utils import is_flax_available, is_torch_available - - -if is_torch_available(): - from .controlnet import ControlNetModel, ControlNetOutput - from .controlnet_flux import FluxControlNetModel, FluxControlNetOutput, FluxMultiControlNetModel - from .controlnet_hunyuan import ( - HunyuanControlNetOutput, - HunyuanDiT2DControlNetModel, - HunyuanDiT2DMultiControlNetModel, - ) - from .controlnet_sana import SanaControlNetModel - from .controlnet_sd3 import SD3ControlNetModel, SD3ControlNetOutput, SD3MultiControlNetModel - from .controlnet_sparsectrl import ( - SparseControlNetConditioningEmbedding, - SparseControlNetModel, - SparseControlNetOutput, - ) - from .controlnet_union import ControlNetUnionModel - from .controlnet_xs import ControlNetXSAdapter, ControlNetXSOutput, UNetControlNetXSModel - from .multicontrolnet import MultiControlNetModel - from .multicontrolnet_union import MultiControlNetUnionModel - -if is_flax_available(): - from .controlnet_flax import FlaxControlNetModel diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 8616ebba240a53c1176614223115798a7592a467..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet.cpython-310.pyc deleted file mode 100644 index 38822b3c3d4de119e25a357e22d1679db7c51a3a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_flax.cpython-310.pyc deleted file mode 100644 index 8153065a896a26bba50904cd51584bbdeb67a919..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_flux.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_flux.cpython-310.pyc deleted file mode 100644 index a382779368f85733cb5fe7d8b8d178c459d0f94d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_flux.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_hunyuan.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_hunyuan.cpython-310.pyc deleted file mode 100644 index 0cc2ded3547eaf6d9dab6e9a8b6063b18a46686a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_hunyuan.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_sana.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_sana.cpython-310.pyc deleted file mode 100644 index 61c0370e274ff12cf87b61f40014cc243f8716f1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_sana.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_sd3.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_sd3.cpython-310.pyc deleted file mode 100644 index 995b6a51365d4bfaa6196ca684d04535ae824cb6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_sd3.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_sparsectrl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_sparsectrl.cpython-310.pyc deleted file mode 100644 index 1dafb297e1ea5e5c00374fe493898114009a967f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_sparsectrl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_union.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_union.cpython-310.pyc deleted file mode 100644 index 689f896e897a90dff479ea17ad11aa32ac167838..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_union.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_xs.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_xs.cpython-310.pyc deleted file mode 100644 index e084935d1205d0a16f6760ebaaa6b26a94421bc5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/controlnet_xs.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/multicontrolnet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/multicontrolnet.cpython-310.pyc deleted file mode 100644 index 68ca25fb0ccc000ccc9ecfc432410e446beabd65..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/multicontrolnet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/multicontrolnet_union.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/multicontrolnet_union.cpython-310.pyc deleted file mode 100644 index ec602976bb5ba15820e595c0cc4907f543fb9664..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/__pycache__/multicontrolnet_union.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet.py deleted file mode 100644 index 3f330999d4c8366d7bd66fdd58256b49d9ebab49..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet.py +++ /dev/null @@ -1,867 +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 dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union - -import torch -from torch import nn -from torch.nn import functional as F - -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders.single_file_model import FromOriginalModelMixin -from ...utils import BaseOutput, logging -from ..attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - AttentionProcessor, - AttnAddedKVProcessor, - AttnProcessor, -) -from ..embeddings import TextImageProjection, TextImageTimeEmbedding, TextTimeEmbedding, TimestepEmbedding, Timesteps -from ..modeling_utils import ModelMixin -from ..unets.unet_2d_blocks import ( - UNetMidBlock2D, - UNetMidBlock2DCrossAttn, - get_down_block, -) -from ..unets.unet_2d_condition import UNet2DConditionModel - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class ControlNetOutput(BaseOutput): - """ - The output of [`ControlNetModel`]. - - Args: - down_block_res_samples (`tuple[torch.Tensor]`): - A tuple of downsample activations at different resolutions for each downsampling block. Each tensor should - be of shape `(batch_size, channel * resolution, height //resolution, width // resolution)`. Output can be - used to condition the original UNet's downsampling activations. - mid_down_block_re_sample (`torch.Tensor`): - The activation of the middle block (the lowest sample resolution). Each tensor should be of shape - `(batch_size, channel * lowest_resolution, height // lowest_resolution, width // lowest_resolution)`. - Output can be used to condition the original UNet's middle block activation. - """ - - down_block_res_samples: Tuple[torch.Tensor] - mid_block_res_sample: torch.Tensor - - -class ControlNetConditioningEmbedding(nn.Module): - """ - Quoting from https://huggingface.co/papers/2302.05543: "Stable Diffusion uses a pre-processing method similar to - VQ-GAN [11] to convert the entire dataset of 512 × 512 images into smaller 64 × 64 “latent images” for stabilized - training. This requires ControlNets to convert image-based conditions to 64 × 64 feature space to match the - convolution size. We use a tiny network E(·) of four convolution layers with 4 × 4 kernels and 2 × 2 strides - (activated by ReLU, channels are 16, 32, 64, 128, initialized with Gaussian weights, trained jointly with the full - model) to encode image-space conditions ... into feature maps ..." - """ - - def __init__( - self, - conditioning_embedding_channels: int, - conditioning_channels: int = 3, - block_out_channels: Tuple[int, ...] = (16, 32, 96, 256), - ): - super().__init__() - - self.conv_in = nn.Conv2d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1) - - self.blocks = nn.ModuleList([]) - - for i in range(len(block_out_channels) - 1): - channel_in = block_out_channels[i] - channel_out = block_out_channels[i + 1] - self.blocks.append(nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1)) - self.blocks.append(nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1, stride=2)) - - self.conv_out = zero_module( - nn.Conv2d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1) - ) - - def forward(self, conditioning): - embedding = self.conv_in(conditioning) - embedding = F.silu(embedding) - - for block in self.blocks: - embedding = block(embedding) - embedding = F.silu(embedding) - - embedding = self.conv_out(embedding) - - return embedding - - -class ControlNetModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): - """ - A ControlNet model. - - Args: - in_channels (`int`, defaults to 4): - The number of channels in the input sample. - flip_sin_to_cos (`bool`, defaults to `True`): - Whether to flip the sin to cos in the time embedding. - freq_shift (`int`, defaults to 0): - 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`): - 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): - The number of layers per block. - downsample_padding (`int`, defaults to 1): - The padding to use for the downsampling convolution. - mid_block_scale_factor (`float`, defaults to 1): - The scale factor to use for the mid block. - act_fn (`str`, defaults to "silu"): - The activation function to use. - 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`, defaults to 1e-5): - 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): - 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`]. - encoder_hid_dim (`int`, *optional*, defaults to None): - If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim` - dimension to `cross_attention_dim`. - 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): - The dimension of the attention heads. - use_linear_projection (`bool`, defaults to `False`): - class_embed_type (`str`, *optional*, defaults to `None`): - The type of class embedding to use which is ultimately summed with the time embeddings. Choose from None, - `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`. - addition_embed_type (`str`, *optional*, defaults to `None`): - Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or - "text". "text" will use the `TextTimeEmbedding` layer. - num_class_embeds (`int`, *optional*, defaults to 0): - Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing - class conditioning with `class_embed_type` equal to `None`. - 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`. - projection_class_embeddings_input_dim (`int`, *optional*, defaults to `None`): - The dimension of the `class_labels` input when `class_embed_type="projection"`. Required when - `class_embed_type="projection"`. - controlnet_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]`, *optional*, 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. - addition_embed_type_num_heads (`int`, defaults to 64): - The number of heads to use for the `TextTimeEmbedding` layer. - """ - - _supports_gradient_checkpointing = True - - @register_to_config - 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, - ): - super().__init__() - - # If `num_attention_heads` is not defined (which is the case for most models) - # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. - # The reason for this behavior is to correct for incorrectly named variables that were introduced - # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 - # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking - # which is why we correct for the naming here. - num_attention_heads = num_attention_heads or attention_head_dim - - # Check inputs - if len(block_out_channels) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." - ) - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) - - # input - conv_in_kernel = 3 - conv_in_padding = (conv_in_kernel - 1) // 2 - self.conv_in = nn.Conv2d( - in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding - ) - - # time - time_embed_dim = block_out_channels[0] * 4 - self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) - timestep_input_dim = block_out_channels[0] - self.time_embedding = TimestepEmbedding( - timestep_input_dim, - time_embed_dim, - act_fn=act_fn, - ) - - if encoder_hid_dim_type is None and encoder_hid_dim is not None: - encoder_hid_dim_type = "text_proj" - self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type) - logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.") - - if encoder_hid_dim is None and encoder_hid_dim_type is not None: - raise ValueError( - f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}." - ) - - if encoder_hid_dim_type == "text_proj": - self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim) - elif encoder_hid_dim_type == "text_image_proj": - # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much - # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use - # case when `addition_embed_type == "text_image_proj"` (Kandinsky 2.1)` - self.encoder_hid_proj = TextImageProjection( - text_embed_dim=encoder_hid_dim, - image_embed_dim=cross_attention_dim, - cross_attention_dim=cross_attention_dim, - ) - - elif encoder_hid_dim_type is not None: - raise ValueError( - f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'." - ) - else: - self.encoder_hid_proj = None - - # class embedding - if class_embed_type is None and num_class_embeds is not None: - self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) - elif class_embed_type == "timestep": - self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim) - elif class_embed_type == "identity": - self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) - elif class_embed_type == "projection": - if projection_class_embeddings_input_dim is None: - raise ValueError( - "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set" - ) - # The projection `class_embed_type` is the same as the timestep `class_embed_type` except - # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings - # 2. it projects from an arbitrary input dimension. - # - # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations. - # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings. - # As a result, `TimestepEmbedding` can be passed arbitrary vectors. - self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) - else: - self.class_embedding = None - - if addition_embed_type == "text": - if encoder_hid_dim is not None: - text_time_embedding_from_dim = encoder_hid_dim - else: - text_time_embedding_from_dim = cross_attention_dim - - self.add_embedding = TextTimeEmbedding( - text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads - ) - elif addition_embed_type == "text_image": - # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much - # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use - # case when `addition_embed_type == "text_image"` (Kandinsky 2.1)` - self.add_embedding = TextImageTimeEmbedding( - text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim - ) - elif addition_embed_type == "text_time": - self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift) - self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) - - elif addition_embed_type is not None: - raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.") - - # control net conditioning embedding - self.controlnet_cond_embedding = ControlNetConditioningEmbedding( - conditioning_embedding_channels=block_out_channels[0], - block_out_channels=conditioning_embedding_out_channels, - conditioning_channels=conditioning_channels, - ) - - self.down_blocks = nn.ModuleList([]) - self.controlnet_down_blocks = nn.ModuleList([]) - - if isinstance(only_cross_attention, bool): - only_cross_attention = [only_cross_attention] * len(down_block_types) - - if isinstance(attention_head_dim, int): - attention_head_dim = (attention_head_dim,) * len(down_block_types) - - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(down_block_types) - - # down - output_channel = block_out_channels[0] - - controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1) - controlnet_block = zero_module(controlnet_block) - self.controlnet_down_blocks.append(controlnet_block) - - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = get_down_block( - down_block_type, - num_layers=layers_per_block, - transformer_layers_per_block=transformer_layers_per_block[i], - in_channels=input_channel, - out_channels=output_channel, - temb_channels=time_embed_dim, - add_downsample=not is_final_block, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads[i], - attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, - downsample_padding=downsample_padding, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention[i], - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - self.down_blocks.append(down_block) - - for _ in range(layers_per_block): - controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1) - controlnet_block = zero_module(controlnet_block) - self.controlnet_down_blocks.append(controlnet_block) - - if not is_final_block: - controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1) - controlnet_block = zero_module(controlnet_block) - self.controlnet_down_blocks.append(controlnet_block) - - # mid - mid_block_channel = block_out_channels[-1] - - controlnet_block = nn.Conv2d(mid_block_channel, mid_block_channel, kernel_size=1) - controlnet_block = zero_module(controlnet_block) - self.controlnet_mid_block = controlnet_block - - if mid_block_type == "UNetMidBlock2DCrossAttn": - self.mid_block = UNetMidBlock2DCrossAttn( - transformer_layers_per_block=transformer_layers_per_block[-1], - in_channels=mid_block_channel, - temb_channels=time_embed_dim, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - output_scale_factor=mid_block_scale_factor, - resnet_time_scale_shift=resnet_time_scale_shift, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads[-1], - resnet_groups=norm_num_groups, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - ) - elif mid_block_type == "UNetMidBlock2D": - self.mid_block = UNetMidBlock2D( - in_channels=block_out_channels[-1], - temb_channels=time_embed_dim, - num_layers=0, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - output_scale_factor=mid_block_scale_factor, - resnet_groups=norm_num_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - add_attention=False, - ) - else: - raise ValueError(f"unknown mid_block_type : {mid_block_type}") - - @classmethod - def from_unet( - cls, - unet: UNet2DConditionModel, - controlnet_conditioning_channel_order: str = "rgb", - conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256), - load_weights_from_unet: bool = True, - conditioning_channels: int = 3, - ): - r""" - Instantiate a [`ControlNetModel`] from [`UNet2DConditionModel`]. - - Parameters: - unet (`UNet2DConditionModel`): - The UNet model weights to copy to the [`ControlNetModel`]. All configuration options are also copied - where applicable. - """ - transformer_layers_per_block = ( - unet.config.transformer_layers_per_block if "transformer_layers_per_block" in unet.config else 1 - ) - encoder_hid_dim = unet.config.encoder_hid_dim if "encoder_hid_dim" in unet.config else None - encoder_hid_dim_type = unet.config.encoder_hid_dim_type if "encoder_hid_dim_type" in unet.config else None - addition_embed_type = unet.config.addition_embed_type if "addition_embed_type" in unet.config else None - addition_time_embed_dim = ( - unet.config.addition_time_embed_dim if "addition_time_embed_dim" in unet.config else None - ) - - controlnet = cls( - encoder_hid_dim=encoder_hid_dim, - encoder_hid_dim_type=encoder_hid_dim_type, - addition_embed_type=addition_embed_type, - addition_time_embed_dim=addition_time_embed_dim, - transformer_layers_per_block=transformer_layers_per_block, - in_channels=unet.config.in_channels, - flip_sin_to_cos=unet.config.flip_sin_to_cos, - freq_shift=unet.config.freq_shift, - down_block_types=unet.config.down_block_types, - only_cross_attention=unet.config.only_cross_attention, - block_out_channels=unet.config.block_out_channels, - layers_per_block=unet.config.layers_per_block, - downsample_padding=unet.config.downsample_padding, - mid_block_scale_factor=unet.config.mid_block_scale_factor, - act_fn=unet.config.act_fn, - norm_num_groups=unet.config.norm_num_groups, - norm_eps=unet.config.norm_eps, - cross_attention_dim=unet.config.cross_attention_dim, - attention_head_dim=unet.config.attention_head_dim, - num_attention_heads=unet.config.num_attention_heads, - use_linear_projection=unet.config.use_linear_projection, - class_embed_type=unet.config.class_embed_type, - num_class_embeds=unet.config.num_class_embeds, - upcast_attention=unet.config.upcast_attention, - resnet_time_scale_shift=unet.config.resnet_time_scale_shift, - projection_class_embeddings_input_dim=unet.config.projection_class_embeddings_input_dim, - mid_block_type=unet.config.mid_block_type, - controlnet_conditioning_channel_order=controlnet_conditioning_channel_order, - conditioning_embedding_out_channels=conditioning_embedding_out_channels, - conditioning_channels=conditioning_channels, - ) - - if load_weights_from_unet: - controlnet.conv_in.load_state_dict(unet.conv_in.state_dict()) - controlnet.time_proj.load_state_dict(unet.time_proj.state_dict()) - controlnet.time_embedding.load_state_dict(unet.time_embedding.state_dict()) - - if controlnet.class_embedding: - controlnet.class_embedding.load_state_dict(unet.class_embedding.state_dict()) - - if hasattr(controlnet, "add_embedding"): - controlnet.add_embedding.load_state_dict(unet.add_embedding.state_dict()) - - controlnet.down_blocks.load_state_dict(unet.down_blocks.state_dict()) - controlnet.mid_block.load_state_dict(unet.mid_block.state_dict()) - - return controlnet - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - 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: - 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. This is useful for saving some memory in exchange for a small decrease in speed. - - Args: - slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): - When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If - `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is - provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` - must be a multiple of `slice_size`. - """ - sliceable_head_dims = [] - - def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): - if hasattr(module, "set_attention_slice"): - sliceable_head_dims.append(module.sliceable_head_dim) - - for child in module.children(): - fn_recursive_retrieve_sliceable_dims(child) - - # retrieve number of attention layers - for module in self.children(): - fn_recursive_retrieve_sliceable_dims(module) - - num_sliceable_layers = len(sliceable_head_dims) - - if slice_size == "auto": - # half the attention head size is usually a good trade-off between - # speed and memory - slice_size = [dim // 2 for dim in sliceable_head_dims] - elif slice_size == "max": - # make smallest slice possible - slice_size = num_sliceable_layers * [1] - - slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size - - if len(slice_size) != len(sliceable_head_dims): - raise ValueError( - f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" - f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." - ) - - for i in range(len(slice_size)): - size = slice_size[i] - dim = sliceable_head_dims[i] - if size is not None and size > dim: - raise ValueError(f"size {size} has to be smaller or equal to {dim}.") - - # 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]): - if hasattr(module, "set_attention_slice"): - module.set_attention_slice(slice_size.pop()) - - for child in module.children(): - fn_recursive_set_attention_slice(child, slice_size) - - reversed_slice_size = list(reversed(slice_size)) - for module in self.children(): - fn_recursive_set_attention_slice(module, reversed_slice_size) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[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, - guess_mode: bool = False, - return_dict: bool = True, - ) -> Union[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]`): - The number of timesteps to denoise an input. - encoder_hidden_states (`torch.Tensor`): - The encoder hidden states. - controlnet_cond (`torch.Tensor`): - The conditional input tensor of shape `(batch_size, sequence_length, hidden_size)`. - conditioning_scale (`float`, defaults to `1.0`): - The scale factor for ControlNet outputs. - class_labels (`torch.Tensor`, *optional*, defaults to `None`): - Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. - timestep_cond (`torch.Tensor`, *optional*, defaults to `None`): - Additional conditional embeddings for timestep. If provided, the embeddings will be summed with the - timestep_embedding passed through the `self.time_embedding` layer to obtain the final timestep - embeddings. - attention_mask (`torch.Tensor`, *optional*, defaults to `None`): - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask - is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large - negative values to the attention scores corresponding to "discard" tokens. - added_cond_kwargs (`dict`): - Additional conditions for the Stable Diffusion XL UNet. - cross_attention_kwargs (`dict[str]`, *optional*, defaults to `None`): - A kwargs dictionary that if specified is passed along to the `AttnProcessor`. - guess_mode (`bool`, defaults to `False`): - In this mode, the ControlNet encoder tries its best to recognize the input content of the input even if - you remove all prompts. A `guidance_scale` between 3.0 and 5.0 is recommended. - return_dict (`bool`, defaults to `True`): - Whether or not to return a [`~models.controlnets.controlnet.ControlNetOutput`] instead of a plain - tuple. - - Returns: - [`~models.controlnets.controlnet.ControlNetOutput`] **or** `tuple`: - If `return_dict` is `True`, a [`~models.controlnets.controlnet.ControlNetOutput`] is returned, - otherwise a tuple is returned where the first element is the sample tensor. - """ - # check channel order - channel_order = self.config.controlnet_conditioning_channel_order - - if channel_order == "rgb": - # in rgb order by default - ... - elif channel_order == "bgr": - controlnet_cond = torch.flip(controlnet_cond, dims=[1]) - else: - raise ValueError(f"unknown `controlnet_conditioning_channel_order`: {channel_order}") - - # prepare attention_mask - if attention_mask is not None: - attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 - attention_mask = attention_mask.unsqueeze(1) - - # 1. time - 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 - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = sample.device.type == "mps" - is_npu = sample.device.type == "npu" - if isinstance(timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) - elif len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps.expand(sample.shape[0]) - - t_emb = self.time_proj(timesteps) - - # timesteps does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - t_emb = t_emb.to(dtype=sample.dtype) - - emb = self.time_embedding(t_emb, timestep_cond) - aug_emb = None - - if self.class_embedding is not None: - if class_labels is None: - raise ValueError("class_labels should be provided when num_class_embeds > 0") - - if self.config.class_embed_type == "timestep": - class_labels = self.time_proj(class_labels) - - class_emb = self.class_embedding(class_labels).to(dtype=self.dtype) - emb = emb + class_emb - - if self.config.addition_embed_type is not None: - if self.config.addition_embed_type == "text": - aug_emb = self.add_embedding(encoder_hidden_states) - - elif self.config.addition_embed_type == "text_time": - if "text_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" - ) - text_embeds = added_cond_kwargs.get("text_embeds") - if "time_ids" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" - ) - time_ids = added_cond_kwargs.get("time_ids") - time_embeds = self.add_time_proj(time_ids.flatten()) - time_embeds = time_embeds.reshape((text_embeds.shape[0], -1)) - - add_embeds = torch.concat([text_embeds, time_embeds], dim=-1) - add_embeds = add_embeds.to(emb.dtype) - aug_emb = self.add_embedding(add_embeds) - - emb = emb + aug_emb if aug_emb is not None else emb - - # 2. pre-process - sample = self.conv_in(sample) - - controlnet_cond = self.controlnet_cond_embedding(controlnet_cond) - sample = sample + controlnet_cond - - # 3. down - down_block_res_samples = (sample,) - for downsample_block in self.down_blocks: - if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: - sample, res_samples = downsample_block( - hidden_states=sample, - temb=emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - ) - else: - sample, res_samples = downsample_block(hidden_states=sample, temb=emb) - - down_block_res_samples += res_samples - - # 4. mid - if self.mid_block is not None: - if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention: - sample = self.mid_block( - sample, - emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - ) - else: - sample = self.mid_block(sample, emb) - - # 5. Control net blocks - - controlnet_down_block_res_samples = () - - for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks): - down_block_res_sample = controlnet_block(down_block_res_sample) - controlnet_down_block_res_samples = controlnet_down_block_res_samples + (down_block_res_sample,) - - down_block_res_samples = controlnet_down_block_res_samples - - mid_block_res_sample = self.controlnet_mid_block(sample) - - # 6. scaling - if guess_mode and not self.config.global_pool_conditions: - scales = torch.logspace(-1, 0, len(down_block_res_samples) + 1, device=sample.device) # 0.1 to 1.0 - scales = scales * conditioning_scale - down_block_res_samples = [sample * scale for sample, scale in zip(down_block_res_samples, scales)] - mid_block_res_sample = mid_block_res_sample * scales[-1] # last one - else: - down_block_res_samples = [sample * conditioning_scale for sample in down_block_res_samples] - mid_block_res_sample = mid_block_res_sample * conditioning_scale - - if self.config.global_pool_conditions: - down_block_res_samples = [ - torch.mean(sample, dim=(2, 3), keepdim=True) for sample in down_block_res_samples - ] - mid_block_res_sample = torch.mean(mid_block_res_sample, dim=(2, 3), keepdim=True) - - if not return_dict: - return (down_block_res_samples, mid_block_res_sample) - - return ControlNetOutput( - down_block_res_samples=down_block_res_samples, mid_block_res_sample=mid_block_res_sample - ) - - -def zero_module(module): - for p in module.parameters(): - nn.init.zeros_(p) - return module diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_flax.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_flax.py deleted file mode 100644 index 4b2148666ebf507d4badbeb8af2e8599fdfb6fb6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_flax.py +++ /dev/null @@ -1,395 +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 - -import flax -import flax.linen as nn -import jax -import jax.numpy as jnp -from flax.core.frozen_dict import FrozenDict - -from ...configuration_utils import ConfigMixin, flax_register_to_config -from ...utils import BaseOutput -from ..embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps -from ..modeling_flax_utils import FlaxModelMixin -from ..unets.unet_2d_blocks_flax import ( - FlaxCrossAttnDownBlock2D, - FlaxDownBlock2D, - FlaxUNetMidBlock2DCrossAttn, -) - - -@flax.struct.dataclass -class FlaxControlNetOutput(BaseOutput): - """ - The output of [`FlaxControlNetModel`]. - - Args: - down_block_res_samples (`jnp.ndarray`): - mid_block_res_sample (`jnp.ndarray`): - """ - - down_block_res_samples: jnp.ndarray - mid_block_res_sample: jnp.ndarray - - -class FlaxControlNetConditioningEmbedding(nn.Module): - conditioning_embedding_channels: int - block_out_channels: Tuple[int, ...] = (16, 32, 96, 256) - dtype: jnp.dtype = jnp.float32 - - def setup(self) -> None: - self.conv_in = nn.Conv( - self.block_out_channels[0], - kernel_size=(3, 3), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - blocks = [] - for i in range(len(self.block_out_channels) - 1): - channel_in = self.block_out_channels[i] - channel_out = self.block_out_channels[i + 1] - conv1 = nn.Conv( - channel_in, - kernel_size=(3, 3), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - blocks.append(conv1) - conv2 = nn.Conv( - channel_out, - kernel_size=(3, 3), - strides=(2, 2), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - blocks.append(conv2) - self.blocks = blocks - - self.conv_out = nn.Conv( - self.conditioning_embedding_channels, - kernel_size=(3, 3), - padding=((1, 1), (1, 1)), - kernel_init=nn.initializers.zeros_init(), - bias_init=nn.initializers.zeros_init(), - dtype=self.dtype, - ) - - def __call__(self, conditioning: jnp.ndarray) -> jnp.ndarray: - embedding = self.conv_in(conditioning) - embedding = nn.silu(embedding) - - for block in self.blocks: - embedding = block(embedding) - embedding = nn.silu(embedding) - - embedding = self.conv_out(embedding) - - return embedding - - -@flax_register_to_config -class FlaxControlNetModel(nn.Module, FlaxModelMixin, ConfigMixin): - r""" - A ControlNet model. - - This model inherits from [`FlaxModelMixin`]. Check the superclass documentation for it’s generic methods - implemented for all models (such as downloading or saving). - - This model is also a Flax Linen [`flax.linen.Module`](https://flax.readthedocs.io/en/latest/flax.linen.html#module) - subclass. Use it as a regular Flax Linen module and refer to the Flax documentation for all matters related to its - general usage and behavior. - - Inherent JAX features such as the following are supported: - - - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit) - - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation) - - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap) - - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap) - - Parameters: - sample_size (`int`, *optional*): - 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")`): - The tuple of downsample blocks to use. - 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): - The dimension of the attention heads. - 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. - dropout (`float`, *optional*, defaults to 0): - Dropout probability for down, up and bottleneck blocks. - 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. - controlnet_conditioning_channel_order (`str`, *optional*, defaults to `rgb`): - The channel order of conditional image. Will convert to `rgb` if it's `bgr`. - conditioning_embedding_out_channels (`tuple`, *optional*, defaults to `(16, 32, 96, 256)`): - The tuple of output channel for each block in the `conditioning_embedding` layer. - """ - - sample_size: int = 32 - in_channels: int = 4 - 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) - layers_per_block: int = 2 - attention_head_dim: Union[int, Tuple[int, ...]] = 8 - num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None - cross_attention_dim: int = 1280 - dropout: float = 0.0 - use_linear_projection: bool = False - dtype: jnp.dtype = jnp.float32 - 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) - - def init_weights(self, rng: jax.Array) -> FrozenDict: - # init input tensors - sample_shape = (1, self.in_channels, self.sample_size, self.sample_size) - sample = jnp.zeros(sample_shape, dtype=jnp.float32) - timesteps = jnp.ones((1,), dtype=jnp.int32) - encoder_hidden_states = jnp.zeros((1, 1, self.cross_attention_dim), dtype=jnp.float32) - controlnet_cond_shape = (1, 3, self.sample_size * 8, self.sample_size * 8) - controlnet_cond = jnp.zeros(controlnet_cond_shape, dtype=jnp.float32) - - params_rng, dropout_rng = jax.random.split(rng) - rngs = {"params": params_rng, "dropout": dropout_rng} - - return self.init(rngs, sample, timesteps, encoder_hidden_states, controlnet_cond)["params"] - - def setup(self) -> None: - block_out_channels = self.block_out_channels - time_embed_dim = block_out_channels[0] * 4 - - # If `num_attention_heads` is not defined (which is the case for most models) - # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. - # The reason for this behavior is to correct for incorrectly named variables that were introduced - # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 - # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking - # which is why we correct for the naming here. - num_attention_heads = self.num_attention_heads or self.attention_head_dim - - # input - self.conv_in = nn.Conv( - block_out_channels[0], - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - # time - self.time_proj = FlaxTimesteps( - block_out_channels[0], flip_sin_to_cos=self.flip_sin_to_cos, freq_shift=self.config.freq_shift - ) - self.time_embedding = FlaxTimestepEmbedding(time_embed_dim, dtype=self.dtype) - - self.controlnet_cond_embedding = FlaxControlNetConditioningEmbedding( - conditioning_embedding_channels=block_out_channels[0], - block_out_channels=self.conditioning_embedding_out_channels, - ) - - only_cross_attention = self.only_cross_attention - if isinstance(only_cross_attention, bool): - only_cross_attention = (only_cross_attention,) * len(self.down_block_types) - - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(self.down_block_types) - - # down - down_blocks = [] - controlnet_down_blocks = [] - - output_channel = block_out_channels[0] - - controlnet_block = nn.Conv( - output_channel, - kernel_size=(1, 1), - padding="VALID", - kernel_init=nn.initializers.zeros_init(), - bias_init=nn.initializers.zeros_init(), - dtype=self.dtype, - ) - controlnet_down_blocks.append(controlnet_block) - - for i, down_block_type in enumerate(self.down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - if down_block_type == "CrossAttnDownBlock2D": - down_block = FlaxCrossAttnDownBlock2D( - in_channels=input_channel, - out_channels=output_channel, - dropout=self.dropout, - num_layers=self.layers_per_block, - num_attention_heads=num_attention_heads[i], - add_downsample=not is_final_block, - use_linear_projection=self.use_linear_projection, - only_cross_attention=only_cross_attention[i], - dtype=self.dtype, - ) - else: - down_block = FlaxDownBlock2D( - in_channels=input_channel, - out_channels=output_channel, - dropout=self.dropout, - num_layers=self.layers_per_block, - add_downsample=not is_final_block, - dtype=self.dtype, - ) - - down_blocks.append(down_block) - - for _ in range(self.layers_per_block): - controlnet_block = nn.Conv( - output_channel, - kernel_size=(1, 1), - padding="VALID", - kernel_init=nn.initializers.zeros_init(), - bias_init=nn.initializers.zeros_init(), - dtype=self.dtype, - ) - controlnet_down_blocks.append(controlnet_block) - - if not is_final_block: - controlnet_block = nn.Conv( - output_channel, - kernel_size=(1, 1), - padding="VALID", - kernel_init=nn.initializers.zeros_init(), - bias_init=nn.initializers.zeros_init(), - dtype=self.dtype, - ) - controlnet_down_blocks.append(controlnet_block) - - self.down_blocks = down_blocks - self.controlnet_down_blocks = controlnet_down_blocks - - # mid - mid_block_channel = block_out_channels[-1] - self.mid_block = FlaxUNetMidBlock2DCrossAttn( - in_channels=mid_block_channel, - dropout=self.dropout, - num_attention_heads=num_attention_heads[-1], - use_linear_projection=self.use_linear_projection, - dtype=self.dtype, - ) - - self.controlnet_mid_block = nn.Conv( - mid_block_channel, - kernel_size=(1, 1), - padding="VALID", - kernel_init=nn.initializers.zeros_init(), - bias_init=nn.initializers.zeros_init(), - dtype=self.dtype, - ) - - def __call__( - self, - sample: jnp.ndarray, - timesteps: Union[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]]: - r""" - Args: - sample (`jnp.ndarray`): (batch, channel, height, width) noisy inputs tensor - timestep (`jnp.ndarray` or `float` or `int`): timesteps - encoder_hidden_states (`jnp.ndarray`): (batch_size, sequence_length, hidden_size) encoder hidden states - controlnet_cond (`jnp.ndarray`): (batch, channel, height, width) the conditional input tensor - conditioning_scale (`float`, *optional*, defaults to `1.0`): the scale factor for controlnet outputs - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`models.unets.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] instead of - a plain tuple. - train (`bool`, *optional*, defaults to `False`): - Use deterministic functions and disable dropout when not training. - - Returns: - [`~models.unets.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] or `tuple`: - [`~models.unets.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] if `return_dict` is True, otherwise - a `tuple`. When returning a tuple, the first element is the sample tensor. - """ - channel_order = self.controlnet_conditioning_channel_order - if channel_order == "bgr": - controlnet_cond = jnp.flip(controlnet_cond, axis=1) - - # 1. time - if not isinstance(timesteps, jnp.ndarray): - timesteps = jnp.array([timesteps], dtype=jnp.int32) - elif isinstance(timesteps, jnp.ndarray) and len(timesteps.shape) == 0: - timesteps = timesteps.astype(dtype=jnp.float32) - timesteps = jnp.expand_dims(timesteps, 0) - - t_emb = self.time_proj(timesteps) - t_emb = self.time_embedding(t_emb) - - # 2. pre-process - sample = jnp.transpose(sample, (0, 2, 3, 1)) - sample = self.conv_in(sample) - - controlnet_cond = jnp.transpose(controlnet_cond, (0, 2, 3, 1)) - controlnet_cond = self.controlnet_cond_embedding(controlnet_cond) - sample += controlnet_cond - - # 3. down - down_block_res_samples = (sample,) - for down_block in self.down_blocks: - if isinstance(down_block, FlaxCrossAttnDownBlock2D): - sample, res_samples = down_block(sample, t_emb, encoder_hidden_states, deterministic=not train) - else: - sample, res_samples = down_block(sample, t_emb, deterministic=not train) - down_block_res_samples += res_samples - - # 4. mid - sample = self.mid_block(sample, t_emb, encoder_hidden_states, deterministic=not train) - - # 5. contronet blocks - controlnet_down_block_res_samples = () - for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks): - down_block_res_sample = controlnet_block(down_block_res_sample) - controlnet_down_block_res_samples += (down_block_res_sample,) - - down_block_res_samples = controlnet_down_block_res_samples - - mid_block_res_sample = self.controlnet_mid_block(sample) - - # 6. scaling - down_block_res_samples = [sample * conditioning_scale for sample in down_block_res_samples] - mid_block_res_sample *= conditioning_scale - - if not return_dict: - return (down_block_res_samples, mid_block_res_sample) - - return FlaxControlNetOutput( - down_block_res_samples=down_block_res_samples, mid_block_res_sample=mid_block_res_sample - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_flux.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_flux.py deleted file mode 100644 index d8e99ee45eb64a8dca76bf432b2d6fb47e49418b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_flux.py +++ /dev/null @@ -1,509 +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 dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union - -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 ..attention_processor import AttentionProcessor -from ..controlnets.controlnet import ControlNetConditioningEmbedding, zero_module -from ..embeddings import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, FluxPosEmbed -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..transformers.transformer_flux import FluxSingleTransformerBlock, FluxTransformerBlock - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class FluxControlNetOutput(BaseOutput): - controlnet_block_samples: Tuple[torch.Tensor] - controlnet_single_block_samples: Tuple[torch.Tensor] - - -class FluxControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin): - _supports_gradient_checkpointing = True - - @register_to_config - 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, - ): - super().__init__() - self.out_channels = in_channels - self.inner_dim = num_attention_heads * attention_head_dim - - self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope) - text_time_guidance_cls = ( - CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings - ) - self.time_text_embed = text_time_guidance_cls( - embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim - ) - - self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim) - self.x_embedder = torch.nn.Linear(in_channels, self.inner_dim) - - self.transformer_blocks = nn.ModuleList( - [ - FluxTransformerBlock( - dim=self.inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - ) - for i in range(num_layers) - ] - ) - - self.single_transformer_blocks = nn.ModuleList( - [ - FluxSingleTransformerBlock( - dim=self.inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - ) - for i in range(num_single_layers) - ] - ) - - # controlnet_blocks - self.controlnet_blocks = nn.ModuleList([]) - for _ in range(len(self.transformer_blocks)): - self.controlnet_blocks.append(zero_module(nn.Linear(self.inner_dim, self.inner_dim))) - - self.controlnet_single_blocks = nn.ModuleList([]) - for _ in range(len(self.single_transformer_blocks)): - self.controlnet_single_blocks.append(zero_module(nn.Linear(self.inner_dim, self.inner_dim))) - - self.union = num_mode is not None - if self.union: - self.controlnet_mode_embedder = nn.Embedding(num_mode, self.inner_dim) - - if conditioning_embedding_channels is not None: - self.input_hint_block = ControlNetConditioningEmbedding( - conditioning_embedding_channels=conditioning_embedding_channels, block_out_channels=(16, 16, 16, 16) - ) - self.controlnet_x_embedder = torch.nn.Linear(in_channels, self.inner_dim) - else: - self.input_hint_block = None - self.controlnet_x_embedder = zero_module(torch.nn.Linear(in_channels, self.inner_dim)) - - self.gradient_checkpointing = False - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self): - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - @classmethod - def from_transformer( - cls, - transformer, - num_layers: int = 4, - num_single_layers: int = 10, - attention_head_dim: int = 128, - num_attention_heads: int = 24, - load_weights_from_transformer=True, - ): - config = dict(transformer.config) - config["num_layers"] = num_layers - config["num_single_layers"] = num_single_layers - config["attention_head_dim"] = attention_head_dim - config["num_attention_heads"] = num_attention_heads - - controlnet = cls.from_config(config) - - if load_weights_from_transformer: - controlnet.pos_embed.load_state_dict(transformer.pos_embed.state_dict()) - controlnet.time_text_embed.load_state_dict(transformer.time_text_embed.state_dict()) - controlnet.context_embedder.load_state_dict(transformer.context_embedder.state_dict()) - controlnet.x_embedder.load_state_dict(transformer.x_embedder.state_dict()) - controlnet.transformer_blocks.load_state_dict(transformer.transformer_blocks.state_dict(), strict=False) - controlnet.single_transformer_blocks.load_state_dict( - transformer.single_transformer_blocks.state_dict(), strict=False - ) - - controlnet.controlnet_x_embedder = zero_module(controlnet.controlnet_x_embedder) - - return controlnet - - def forward( - self, - hidden_states: torch.Tensor, - controlnet_cond: torch.Tensor, - controlnet_mode: torch.Tensor = None, - conditioning_scale: float = 1.0, - 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, - return_dict: bool = True, - ) -> Union[torch.FloatTensor, Transformer2DModelOutput]: - """ - The [`FluxTransformer2DModel`] forward method. - - Args: - hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`): - Input `hidden_states`. - controlnet_cond (`torch.Tensor`): - The conditional input tensor of shape `(batch_size, sequence_length, hidden_size)`. - controlnet_mode (`torch.Tensor`): - The mode tensor of shape `(batch_size, 1)`. - conditioning_scale (`float`, defaults to `1.0`): - 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. - 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. - """ - 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: - controlnet_cond = self.input_hint_block(controlnet_cond) - batch_size, channels, height_pw, width_pw = controlnet_cond.shape - height = height_pw // self.config.patch_size - width = width_pw // self.config.patch_size - controlnet_cond = controlnet_cond.reshape( - batch_size, channels, height, self.config.patch_size, width, self.config.patch_size - ) - controlnet_cond = controlnet_cond.permute(0, 2, 4, 1, 3, 5) - controlnet_cond = controlnet_cond.reshape(batch_size, height * width, -1) - # add - hidden_states = hidden_states + self.controlnet_x_embedder(controlnet_cond) - - timestep = timestep.to(hidden_states.dtype) * 1000 - if guidance is not None: - guidance = guidance.to(hidden_states.dtype) * 1000 - else: - guidance = None - temb = ( - self.time_text_embed(timestep, pooled_projections) - if guidance is None - else self.time_text_embed(timestep, guidance, pooled_projections) - ) - 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] - - if self.union: - # union mode - if controlnet_mode is None: - raise ValueError("`controlnet_mode` cannot be `None` when applying ControlNet-Union") - # union mode emb - controlnet_mode_emb = self.controlnet_mode_embedder(controlnet_mode) - encoder_hidden_states = torch.cat([controlnet_mode_emb, encoder_hidden_states], dim=1) - txt_ids = torch.cat([txt_ids[:1], txt_ids], dim=0) - - ids = torch.cat((txt_ids, img_ids), dim=0) - image_rotary_emb = self.pos_embed(ids) - - block_samples = () - for index_block, block in enumerate(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, - temb, - image_rotary_emb, - ) - - 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, - ) - block_samples = block_samples + (hidden_states,) - - hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) - - single_block_samples = () - for index_block, block in enumerate(self.single_transformer_blocks): - 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, - ) - single_block_samples = single_block_samples + (hidden_states[:, encoder_hidden_states.shape[1] :],) - - # controlnet block - controlnet_block_samples = () - for block_sample, controlnet_block in zip(block_samples, self.controlnet_blocks): - block_sample = controlnet_block(block_sample) - controlnet_block_samples = controlnet_block_samples + (block_sample,) - - controlnet_single_block_samples = () - for single_block_sample, controlnet_block in zip(single_block_samples, self.controlnet_single_blocks): - single_block_sample = controlnet_block(single_block_sample) - controlnet_single_block_samples = controlnet_single_block_samples + (single_block_sample,) - - # scaling - controlnet_block_samples = [sample * conditioning_scale for sample in controlnet_block_samples] - controlnet_single_block_samples = [sample * conditioning_scale for sample in controlnet_single_block_samples] - - controlnet_block_samples = None if len(controlnet_block_samples) == 0 else controlnet_block_samples - controlnet_single_block_samples = ( - 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) - - return FluxControlNetOutput( - controlnet_block_samples=controlnet_block_samples, - controlnet_single_block_samples=controlnet_single_block_samples, - ) - - -class FluxMultiControlNetModel(ModelMixin): - r""" - `FluxMultiControlNetModel` wrapper class for Multi-FluxControlNetModel - - This module is a wrapper for multiple instances of the `FluxControlNetModel`. The `forward()` API is designed to be - compatible with `FluxControlNetModel`. - - Args: - controlnets (`List[FluxControlNetModel]`): - Provides additional conditioning to the unet during the denoising process. You must set multiple - `FluxControlNetModel` as a list. - """ - - def __init__(self, controlnets): - super().__init__() - self.nets = nn.ModuleList(controlnets) - - def forward( - self, - hidden_states: torch.FloatTensor, - 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, - return_dict: bool = True, - ) -> Union[FluxControlNetOutput, Tuple]: - # ControlNet-Union with multiple conditions - # only load one ControlNet for saving memories - if len(self.nets) == 1: - controlnet = self.nets[0] - - for i, (image, mode, scale) in enumerate(zip(controlnet_cond, controlnet_mode, conditioning_scale)): - block_samples, single_block_samples = controlnet( - hidden_states=hidden_states, - controlnet_cond=image, - controlnet_mode=mode[:, None], - conditioning_scale=scale, - timestep=timestep, - guidance=guidance, - pooled_projections=pooled_projections, - encoder_hidden_states=encoder_hidden_states, - txt_ids=txt_ids, - img_ids=img_ids, - joint_attention_kwargs=joint_attention_kwargs, - return_dict=return_dict, - ) - - # merge samples - if i == 0: - control_block_samples = block_samples - control_single_block_samples = single_block_samples - else: - if block_samples is not None and control_block_samples is not None: - control_block_samples = [ - control_block_sample + block_sample - for control_block_sample, block_sample in zip(control_block_samples, block_samples) - ] - if single_block_samples is not None and control_single_block_samples is not None: - control_single_block_samples = [ - control_single_block_sample + block_sample - for control_single_block_sample, block_sample in zip( - control_single_block_samples, single_block_samples - ) - ] - - # Regular Multi-ControlNets - # load all ControlNets into memories - else: - for i, (image, mode, scale, controlnet) in enumerate( - zip(controlnet_cond, controlnet_mode, conditioning_scale, self.nets) - ): - block_samples, single_block_samples = controlnet( - hidden_states=hidden_states, - controlnet_cond=image, - controlnet_mode=mode[:, None], - conditioning_scale=scale, - timestep=timestep, - guidance=guidance, - pooled_projections=pooled_projections, - encoder_hidden_states=encoder_hidden_states, - txt_ids=txt_ids, - img_ids=img_ids, - joint_attention_kwargs=joint_attention_kwargs, - return_dict=return_dict, - ) - - # merge samples - if i == 0: - control_block_samples = block_samples - control_single_block_samples = single_block_samples - else: - if block_samples is not None and control_block_samples is not None: - control_block_samples = [ - control_block_sample + block_sample - for control_block_sample, block_sample in zip(control_block_samples, block_samples) - ] - if single_block_samples is not None and control_single_block_samples is not None: - control_single_block_samples = [ - control_single_block_sample + block_sample - for control_single_block_sample, block_sample in zip( - control_single_block_samples, single_block_samples - ) - ] - - return control_block_samples, control_single_block_samples diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_hunyuan.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_hunyuan.py deleted file mode 100644 index d17d5692aa40d9da7ca280cb1ee55c07542c8679..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_hunyuan.py +++ /dev/null @@ -1,401 +0,0 @@ -# Copyright 2025 HunyuanDiT Authors, Qixun Wang 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 dataclasses import dataclass -from typing import Dict, Optional, Union - -import torch -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import BaseOutput, logging -from ..attention_processor import AttentionProcessor -from ..embeddings import ( - HunyuanCombinedTimestepTextSizeStyleEmbedding, - PatchEmbed, - PixArtAlphaTextProjection, -) -from ..modeling_utils import ModelMixin -from ..transformers.hunyuan_transformer_2d import HunyuanDiTBlock -from .controlnet import Tuple, zero_module - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class HunyuanControlNetOutput(BaseOutput): - controlnet_block_samples: Tuple[torch.Tensor] - - -class HunyuanDiT2DControlNetModel(ModelMixin, ConfigMixin): - @register_to_config - def __init__( - self, - conditioning_channels: int = 3, - num_attention_heads: int = 16, - attention_head_dim: int = 88, - in_channels: Optional[int] = None, - patch_size: Optional[int] = None, - activation_fn: str = "gelu-approximate", - sample_size=32, - hidden_size=1152, - transformer_num_layers: int = 40, - mlp_ratio: float = 4.0, - cross_attention_dim: int = 1024, - cross_attention_dim_t5: int = 2048, - pooled_projection_dim: int = 1024, - text_len: int = 77, - text_len_t5: int = 256, - use_style_cond_and_image_meta_size: bool = True, - ): - super().__init__() - self.num_heads = num_attention_heads - self.inner_dim = num_attention_heads * attention_head_dim - - self.text_embedder = PixArtAlphaTextProjection( - in_features=cross_attention_dim_t5, - hidden_size=cross_attention_dim_t5 * 4, - out_features=cross_attention_dim, - act_fn="silu_fp32", - ) - - self.text_embedding_padding = nn.Parameter( - torch.randn(text_len + text_len_t5, cross_attention_dim, dtype=torch.float32) - ) - - self.pos_embed = PatchEmbed( - height=sample_size, - width=sample_size, - in_channels=in_channels, - embed_dim=hidden_size, - patch_size=patch_size, - pos_embed_type=None, - ) - - self.time_extra_emb = HunyuanCombinedTimestepTextSizeStyleEmbedding( - hidden_size, - pooled_projection_dim=pooled_projection_dim, - seq_len=text_len_t5, - cross_attention_dim=cross_attention_dim_t5, - use_style_cond_and_image_meta_size=use_style_cond_and_image_meta_size, - ) - - # controlnet_blocks - self.controlnet_blocks = nn.ModuleList([]) - - # HunyuanDiT Blocks - self.blocks = nn.ModuleList( - [ - HunyuanDiTBlock( - dim=self.inner_dim, - num_attention_heads=self.config.num_attention_heads, - activation_fn=activation_fn, - ff_inner_dim=int(self.inner_dim * mlp_ratio), - cross_attention_dim=cross_attention_dim, - qk_norm=True, # See https://huggingface.co/papers/2302.05442 for details. - skip=False, # always False as it is the first half of the model - ) - for layer in range(transformer_num_layers // 2 - 1) - ] - ) - self.input_block = zero_module(nn.Linear(hidden_size, hidden_size)) - for _ in range(len(self.blocks)): - controlnet_block = nn.Linear(hidden_size, hidden_size) - controlnet_block = zero_module(controlnet_block) - self.controlnet_blocks.append(controlnet_block) - - @property - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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) - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. If `processor` is a dict, the key needs to define the path to the - corresponding cross attention processor. This is strongly recommended when setting trainable attention - processors. - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - @classmethod - def from_transformer( - cls, transformer, conditioning_channels=3, transformer_num_layers=None, load_weights_from_transformer=True - ): - config = transformer.config - activation_fn = config.activation_fn - attention_head_dim = config.attention_head_dim - cross_attention_dim = config.cross_attention_dim - cross_attention_dim_t5 = config.cross_attention_dim_t5 - hidden_size = config.hidden_size - in_channels = config.in_channels - mlp_ratio = config.mlp_ratio - num_attention_heads = config.num_attention_heads - patch_size = config.patch_size - sample_size = config.sample_size - text_len = config.text_len - text_len_t5 = config.text_len_t5 - - conditioning_channels = conditioning_channels - transformer_num_layers = transformer_num_layers or config.transformer_num_layers - - controlnet = cls( - conditioning_channels=conditioning_channels, - transformer_num_layers=transformer_num_layers, - activation_fn=activation_fn, - attention_head_dim=attention_head_dim, - cross_attention_dim=cross_attention_dim, - cross_attention_dim_t5=cross_attention_dim_t5, - hidden_size=hidden_size, - in_channels=in_channels, - mlp_ratio=mlp_ratio, - num_attention_heads=num_attention_heads, - patch_size=patch_size, - sample_size=sample_size, - text_len=text_len, - text_len_t5=text_len_t5, - ) - if load_weights_from_transformer: - key = controlnet.load_state_dict(transformer.state_dict(), strict=False) - logger.warning(f"controlnet load from Hunyuan-DiT. missing_keys: {key[0]}") - return controlnet - - def forward( - self, - hidden_states, - timestep, - controlnet_cond: torch.Tensor, - conditioning_scale: float = 1.0, - encoder_hidden_states=None, - text_embedding_mask=None, - encoder_hidden_states_t5=None, - text_embedding_mask_t5=None, - image_meta_size=None, - style=None, - image_rotary_emb=None, - return_dict=True, - ): - """ - The [`HunyuanDiT2DControlNetModel`] forward method. - - Args: - hidden_states (`torch.Tensor` of shape `(batch size, dim, height, width)`): - The input tensor. - timestep ( `torch.LongTensor`, *optional*): - Used to indicate denoising step. - controlnet_cond ( `torch.Tensor` ): - The conditioning input to ControlNet. - conditioning_scale ( `float` ): - Indicate the conditioning scale. - encoder_hidden_states ( `torch.Tensor` of shape `(batch size, sequence len, embed dims)`, *optional*): - Conditional embeddings for cross attention layer. This is the output of `BertModel`. - text_embedding_mask: torch.Tensor - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. This is the output - of `BertModel`. - encoder_hidden_states_t5 ( `torch.Tensor` of shape `(batch size, sequence len, embed dims)`, *optional*): - Conditional embeddings for cross attention layer. This is the output of T5 Text Encoder. - text_embedding_mask_t5: torch.Tensor - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. This is the output - of T5 Text Encoder. - image_meta_size (torch.Tensor): - Conditional embedding indicate the image sizes - style: torch.Tensor: - Conditional embedding indicate the style - image_rotary_emb (`torch.Tensor`): - The image rotary embeddings to apply on query and key tensors during attention calculation. - return_dict: bool - Whether to return a dictionary. - """ - - height, width = hidden_states.shape[-2:] - - hidden_states = self.pos_embed(hidden_states) # b,c,H,W -> b, N, C - - # 2. pre-process - hidden_states = hidden_states + self.input_block(self.pos_embed(controlnet_cond)) - - temb = self.time_extra_emb( - timestep, encoder_hidden_states_t5, image_meta_size, style, hidden_dtype=timestep.dtype - ) # [B, D] - - # text projection - batch_size, sequence_length, _ = encoder_hidden_states_t5.shape - encoder_hidden_states_t5 = self.text_embedder( - encoder_hidden_states_t5.view(-1, encoder_hidden_states_t5.shape[-1]) - ) - encoder_hidden_states_t5 = encoder_hidden_states_t5.view(batch_size, sequence_length, -1) - - encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states_t5], dim=1) - text_embedding_mask = torch.cat([text_embedding_mask, text_embedding_mask_t5], dim=-1) - text_embedding_mask = text_embedding_mask.unsqueeze(2).bool() - - encoder_hidden_states = torch.where(text_embedding_mask, encoder_hidden_states, self.text_embedding_padding) - - block_res_samples = () - for layer, block in enumerate(self.blocks): - hidden_states = block( - hidden_states, - temb=temb, - encoder_hidden_states=encoder_hidden_states, - image_rotary_emb=image_rotary_emb, - ) # (N, L, D) - - block_res_samples = block_res_samples + (hidden_states,) - - controlnet_block_res_samples = () - for block_res_sample, controlnet_block in zip(block_res_samples, self.controlnet_blocks): - block_res_sample = controlnet_block(block_res_sample) - controlnet_block_res_samples = controlnet_block_res_samples + (block_res_sample,) - - # 6. scaling - controlnet_block_res_samples = [sample * conditioning_scale for sample in controlnet_block_res_samples] - - if not return_dict: - return (controlnet_block_res_samples,) - - return HunyuanControlNetOutput(controlnet_block_samples=controlnet_block_res_samples) - - -class HunyuanDiT2DMultiControlNetModel(ModelMixin): - r""" - `HunyuanDiT2DMultiControlNetModel` wrapper class for Multi-HunyuanDiT2DControlNetModel - - This module is a wrapper for multiple instances of the `HunyuanDiT2DControlNetModel`. The `forward()` API is - designed to be compatible with `HunyuanDiT2DControlNetModel`. - - Args: - controlnets (`List[HunyuanDiT2DControlNetModel]`): - Provides additional conditioning to the unet during the denoising process. You must set multiple - `HunyuanDiT2DControlNetModel` as a list. - """ - - def __init__(self, controlnets): - super().__init__() - self.nets = nn.ModuleList(controlnets) - - def forward( - self, - hidden_states, - timestep, - controlnet_cond: torch.Tensor, - conditioning_scale: float = 1.0, - encoder_hidden_states=None, - text_embedding_mask=None, - encoder_hidden_states_t5=None, - text_embedding_mask_t5=None, - image_meta_size=None, - style=None, - image_rotary_emb=None, - return_dict=True, - ): - """ - The [`HunyuanDiT2DControlNetModel`] forward method. - - Args: - hidden_states (`torch.Tensor` of shape `(batch size, dim, height, width)`): - The input tensor. - timestep ( `torch.LongTensor`, *optional*): - Used to indicate denoising step. - controlnet_cond ( `torch.Tensor` ): - The conditioning input to ControlNet. - conditioning_scale ( `float` ): - Indicate the conditioning scale. - encoder_hidden_states ( `torch.Tensor` of shape `(batch size, sequence len, embed dims)`, *optional*): - Conditional embeddings for cross attention layer. This is the output of `BertModel`. - text_embedding_mask: torch.Tensor - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. This is the output - of `BertModel`. - encoder_hidden_states_t5 ( `torch.Tensor` of shape `(batch size, sequence len, embed dims)`, *optional*): - Conditional embeddings for cross attention layer. This is the output of T5 Text Encoder. - text_embedding_mask_t5: torch.Tensor - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. This is the output - of T5 Text Encoder. - image_meta_size (torch.Tensor): - Conditional embedding indicate the image sizes - style: torch.Tensor: - Conditional embedding indicate the style - image_rotary_emb (`torch.Tensor`): - The image rotary embeddings to apply on query and key tensors during attention calculation. - return_dict: bool - Whether to return a dictionary. - """ - for i, (image, scale, controlnet) in enumerate(zip(controlnet_cond, conditioning_scale, self.nets)): - block_samples = controlnet( - hidden_states=hidden_states, - timestep=timestep, - controlnet_cond=image, - conditioning_scale=scale, - encoder_hidden_states=encoder_hidden_states, - text_embedding_mask=text_embedding_mask, - encoder_hidden_states_t5=encoder_hidden_states_t5, - text_embedding_mask_t5=text_embedding_mask_t5, - image_meta_size=image_meta_size, - style=style, - image_rotary_emb=image_rotary_emb, - return_dict=return_dict, - ) - - # merge samples - if i == 0: - control_block_samples = block_samples - else: - control_block_samples = [ - control_block_sample + block_sample - for control_block_sample, block_sample in zip(control_block_samples[0], block_samples[0]) - ] - control_block_samples = (control_block_samples,) - - return control_block_samples diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_sana.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_sana.py deleted file mode 100644 index ed521adbedda391e8e6a3c07c6084a9a19566ced..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_sana.py +++ /dev/null @@ -1,290 +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 dataclasses import dataclass -from typing import Any, Dict, Optional, Tuple, Union - -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 ..attention_processor import AttentionProcessor -from ..embeddings import PatchEmbed, PixArtAlphaTextProjection -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormSingle, RMSNorm -from ..transformers.sana_transformer import SanaTransformerBlock -from .controlnet import zero_module - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class SanaControlNetOutput(BaseOutput): - controlnet_block_samples: Tuple[torch.Tensor] - - -class SanaControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin): - _supports_gradient_checkpointing = True - _no_split_modules = ["SanaTransformerBlock", "PatchEmbed"] - _skip_layerwise_casting_patterns = ["patch_embed", "norm"] - - @register_to_config - def __init__( - self, - in_channels: int = 32, - out_channels: Optional[int] = 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, - caption_channels: int = 2304, - mlp_ratio: float = 2.5, - dropout: float = 0.0, - attention_bias: bool = False, - sample_size: int = 32, - patch_size: int = 1, - norm_elementwise_affine: bool = False, - norm_eps: float = 1e-6, - interpolation_scale: Optional[int] = None, - ) -> None: - super().__init__() - - out_channels = out_channels or in_channels - inner_dim = num_attention_heads * attention_head_dim - - # 1. Patch Embedding - self.patch_embed = PatchEmbed( - height=sample_size, - width=sample_size, - patch_size=patch_size, - in_channels=in_channels, - embed_dim=inner_dim, - interpolation_scale=interpolation_scale, - pos_embed_type="sincos" if interpolation_scale is not None else None, - ) - - # 2. Additional condition embeddings - self.time_embed = AdaLayerNormSingle(inner_dim) - - self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) - self.caption_norm = RMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) - - # 3. Transformer blocks - self.transformer_blocks = nn.ModuleList( - [ - SanaTransformerBlock( - inner_dim, - num_attention_heads, - attention_head_dim, - dropout=dropout, - num_cross_attention_heads=num_cross_attention_heads, - cross_attention_head_dim=cross_attention_head_dim, - cross_attention_dim=cross_attention_dim, - attention_bias=attention_bias, - norm_elementwise_affine=norm_elementwise_affine, - norm_eps=norm_eps, - mlp_ratio=mlp_ratio, - ) - for _ in range(num_layers) - ] - ) - - # controlnet_blocks - self.controlnet_blocks = nn.ModuleList([]) - - self.input_block = zero_module(nn.Linear(inner_dim, inner_dim)) - for _ in range(len(self.transformer_blocks)): - controlnet_block = nn.Linear(inner_dim, inner_dim) - controlnet_block = zero_module(controlnet_block) - self.controlnet_blocks.append(controlnet_block) - - self.gradient_checkpointing = False - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - 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, - 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." - ) - - # 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. - # expects mask of shape: - # [batch, key_tokens] - # adds singleton query_tokens dimension: - # [batch, 1, key_tokens] - # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: - # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) - # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) - if attention_mask is not None and attention_mask.ndim == 2: - # assume that mask is expressed as: - # (1 = keep, 0 = discard) - # convert mask into a bias that can be added to attention scores: - # (keep = +0, discard = -10000.0) - attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0 - attention_mask = attention_mask.unsqueeze(1) - - # 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) - - # 1. Input - batch_size, num_channels, height, width = hidden_states.shape - p = self.config.patch_size - post_patch_height, post_patch_width = height // p, width // p - - hidden_states = self.patch_embed(hidden_states) - hidden_states = hidden_states + self.input_block(self.patch_embed(controlnet_cond.to(hidden_states.dtype))) - - timestep, embedded_timestep = self.time_embed( - timestep, batch_size=batch_size, hidden_dtype=hidden_states.dtype - ) - - encoder_hidden_states = self.caption_projection(encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1]) - - encoder_hidden_states = self.caption_norm(encoder_hidden_states) - - # 2. Transformer blocks - block_res_samples = () - if torch.is_grad_enabled() and self.gradient_checkpointing: - for block in self.transformer_blocks: - hidden_states = self._gradient_checkpointing_func( - block, - hidden_states, - attention_mask, - encoder_hidden_states, - encoder_attention_mask, - timestep, - post_patch_height, - post_patch_width, - ) - block_res_samples = block_res_samples + (hidden_states,) - else: - for block in self.transformer_blocks: - hidden_states = block( - hidden_states, - attention_mask, - encoder_hidden_states, - encoder_attention_mask, - timestep, - post_patch_height, - post_patch_width, - ) - block_res_samples = block_res_samples + (hidden_states,) - - # 3. ControlNet blocks - controlnet_block_res_samples = () - for block_res_sample, controlnet_block in zip(block_res_samples, self.controlnet_blocks): - 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: - return (controlnet_block_res_samples,) - - return SanaControlNetOutput(controlnet_block_samples=controlnet_block_res_samples) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_sd3.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_sd3.py deleted file mode 100644 index 8d892cb3b697c94285d0c970c8bbeca6f3cf0947..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_sd3.py +++ /dev/null @@ -1,513 +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 dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union - -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 ..attention import JointTransformerBlock -from ..attention_processor import Attention, AttentionProcessor, FusedJointAttnProcessor2_0 -from ..embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..transformers.transformer_sd3 import SD3SingleTransformerBlock -from .controlnet import BaseOutput, zero_module - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class SD3ControlNetOutput(BaseOutput): - controlnet_block_samples: Tuple[torch.Tensor] - - -class SD3ControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): - r""" - ControlNet model for [Stable Diffusion 3](https://huggingface.co/papers/2403.03206). - - Parameters: - sample_size (`int`, defaults to `128`): - The width/height of the latents. This is fixed during training since it is used to learn a number of - position embeddings. - patch_size (`int`, defaults to `2`): - Patch size to turn the input data into small patches. - in_channels (`int`, defaults to `16`): - The number of latent channels in the input. - num_layers (`int`, defaults to `18`): - The number of layers of transformer blocks to use. - attention_head_dim (`int`, defaults to `64`): - The number of channels in each head. - num_attention_heads (`int`, defaults to `18`): - The number of heads to use for multi-head attention. - joint_attention_dim (`int`, defaults to `4096`): - The embedding dimension to use for joint text-image attention. - caption_projection_dim (`int`, defaults to `1152`): - The embedding dimension of caption embeddings. - pooled_projection_dim (`int`, defaults to `2048`): - The embedding dimension of pooled text projections. - out_channels (`int`, defaults to `16`): - The number of latent channels in the output. - pos_embed_max_size (`int`, defaults to `96`): - 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 `()`): - 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. - pos_embed_type (`str`, defaults to `"sincos"`): - The type of positional embedding to use. Choose between `"sincos"` and `None`. - use_pos_embed (`bool`, defaults to `True`): - Whether to use positional embeddings. - force_zeros_for_pooled_projection (`bool`, defaults to `True`): - Whether to force zeros for pooled projection embeddings. This is handled in the pipelines by reading the - config value of the ControlNet model. - """ - - _supports_gradient_checkpointing = True - - @register_to_config - 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, - dual_attention_layers: Tuple[int, ...] = (), - qk_norm: Optional[str] = None, - pos_embed_type: Optional[str] = "sincos", - use_pos_embed: bool = True, - force_zeros_for_pooled_projection: bool = True, - ): - super().__init__() - default_out_channels = in_channels - self.out_channels = out_channels if out_channels is not None else default_out_channels - self.inner_dim = num_attention_heads * attention_head_dim - - if use_pos_embed: - self.pos_embed = PatchEmbed( - height=sample_size, - width=sample_size, - patch_size=patch_size, - in_channels=in_channels, - embed_dim=self.inner_dim, - pos_embed_max_size=pos_embed_max_size, - pos_embed_type=pos_embed_type, - ) - else: - self.pos_embed = None - self.time_text_embed = CombinedTimestepTextProjEmbeddings( - embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim - ) - if joint_attention_dim is not None: - self.context_embedder = nn.Linear(joint_attention_dim, caption_projection_dim) - - # `attention_head_dim` is doubled to account for the mixing. - # It needs to crafted when we get the actual checkpoints. - self.transformer_blocks = nn.ModuleList( - [ - JointTransformerBlock( - dim=self.inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - context_pre_only=False, - qk_norm=qk_norm, - use_dual_attention=True if i in dual_attention_layers else False, - ) - for i in range(num_layers) - ] - ) - else: - self.context_embedder = None - self.transformer_blocks = nn.ModuleList( - [ - SD3SingleTransformerBlock( - dim=self.inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - ) - for _ in range(num_layers) - ] - ) - - # controlnet_blocks - self.controlnet_blocks = nn.ModuleList([]) - for _ in range(len(self.transformer_blocks)): - controlnet_block = nn.Linear(self.inner_dim, self.inner_dim) - controlnet_block = zero_module(controlnet_block) - self.controlnet_blocks.append(controlnet_block) - pos_embed_input = PatchEmbed( - height=sample_size, - width=sample_size, - patch_size=patch_size, - in_channels=in_channels + extra_conditioning_channels, - embed_dim=self.inner_dim, - pos_embed_type=None, - ) - self.pos_embed_input = zero_module(pos_embed_input) - - 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: - """ - Sets the attention processor to use [feed forward - chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). - - Parameters: - chunk_size (`int`, *optional*): - The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually - over each tensor of dim=`dim`. - dim (`int`, *optional*, defaults to `0`): - The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch) - or dim=1 (sequence length). - """ - if dim not in [0, 1]: - raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}") - - # By default chunk size is 1 - chunk_size = chunk_size or 1 - - def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int): - if hasattr(module, "set_chunk_feed_forward"): - module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) - - for child in module.children(): - fn_recursive_feed_forward(child, chunk_size, dim) - - for module in self.children(): - fn_recursive_feed_forward(module, chunk_size, dim) - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.transformers.transformer_sd3.SD3Transformer2DModel.fuse_qkv_projections - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - self.set_attn_processor(FusedJointAttnProcessor2_0()) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - if self.original_attn_processors is not None: - self.set_attn_processor(self.original_attn_processors) - - # Notes: This is for SD3.5 8b controlnet, which shares the pos_embed with the transformer - # we should have handled this in conversion script - def _get_pos_embed_from_transformer(self, transformer): - pos_embed = PatchEmbed( - height=transformer.config.sample_size, - width=transformer.config.sample_size, - patch_size=transformer.config.patch_size, - in_channels=transformer.config.in_channels, - embed_dim=transformer.inner_dim, - pos_embed_max_size=transformer.config.pos_embed_max_size, - ) - pos_embed.load_state_dict(transformer.pos_embed.state_dict(), strict=True) - return pos_embed - - @classmethod - def from_transformer( - cls, transformer, num_layers=12, num_extra_conditioning_channels=1, load_weights_from_transformer=True - ): - config = transformer.config - config["num_layers"] = num_layers or config.num_layers - config["extra_conditioning_channels"] = num_extra_conditioning_channels - controlnet = cls.from_config(config) - - if load_weights_from_transformer: - controlnet.pos_embed.load_state_dict(transformer.pos_embed.state_dict()) - controlnet.time_text_embed.load_state_dict(transformer.time_text_embed.state_dict()) - controlnet.context_embedder.load_state_dict(transformer.context_embedder.state_dict()) - controlnet.transformer_blocks.load_state_dict(transformer.transformer_blocks.state_dict(), strict=False) - - controlnet.pos_embed_input = zero_module(controlnet.pos_embed_input) - - return controlnet - - def forward( - self, - hidden_states: torch.Tensor, - controlnet_cond: torch.Tensor, - conditioning_scale: float = 1.0, - encoder_hidden_states: torch.Tensor = None, - pooled_projections: torch.Tensor = None, - timestep: torch.LongTensor = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - return_dict: bool = True, - ) -> Union[torch.Tensor, Transformer2DModelOutput]: - """ - The [`SD3Transformer2DModel`] forward method. - - Args: - hidden_states (`torch.Tensor` of shape `(batch size, channel, height, width)`): - Input `hidden_states`. - controlnet_cond (`torch.Tensor`): - The conditional input tensor of shape `(batch_size, sequence_length, hidden_size)`. - conditioning_scale (`float`, defaults to `1.0`): - The scale factor for ControlNet outputs. - encoder_hidden_states (`torch.Tensor` of shape `(batch size, sequence_len, embed_dims)`): - Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. - pooled_projections (`torch.Tensor` of shape `(batch_size, projection_dim)`): Embeddings projected - from the embeddings of input conditions. - timestep ( `torch.LongTensor`): - Used to indicate denoising step. - 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. - """ - 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") - - # SD3.5 8b controlnet does not have a `pos_embed`, - # it use the `pos_embed` from the transformer to process input before passing to controlnet - elif self.pos_embed is None and hidden_states.ndim != 3: - raise ValueError("hidden_states must be 3D when pos_embed is not used") - - if self.context_embedder is not None and encoder_hidden_states is None: - raise ValueError("encoder_hidden_states must be provided when context_embedder is used") - # SD3.5 8b controlnet does not have a `context_embedder`, it does not use `encoder_hidden_states` - elif self.context_embedder is None and encoder_hidden_states is not None: - raise ValueError("encoder_hidden_states should not be provided when context_embedder is not used") - - if self.pos_embed is not None: - hidden_states = self.pos_embed(hidden_states) # takes care of adding positional embeddings too. - - temb = self.time_text_embed(timestep, pooled_projections) - - if self.context_embedder is not None: - encoder_hidden_states = self.context_embedder(encoder_hidden_states) - - # add - hidden_states = hidden_states + self.pos_embed_input(controlnet_cond) - - block_res_samples = () - - for block in self.transformer_blocks: - if torch.is_grad_enabled() and self.gradient_checkpointing: - if self.context_embedder is not None: - encoder_hidden_states, hidden_states = self._gradient_checkpointing_func( - block, - hidden_states, - encoder_hidden_states, - temb, - ) - else: - # SD3.5 8b controlnet use single transformer block, which does not use `encoder_hidden_states` - hidden_states = self._gradient_checkpointing_func(block, hidden_states, temb) - - else: - if self.context_embedder is not None: - encoder_hidden_states, hidden_states = block( - hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb - ) - else: - # SD3.5 8b controlnet use single transformer block, which does not use `encoder_hidden_states` - hidden_states = block(hidden_states, temb) - - block_res_samples = block_res_samples + (hidden_states,) - - controlnet_block_res_samples = () - for block_res_sample, controlnet_block in zip(block_res_samples, self.controlnet_blocks): - block_res_sample = controlnet_block(block_res_sample) - controlnet_block_res_samples = controlnet_block_res_samples + (block_res_sample,) - - # 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,) - - return SD3ControlNetOutput(controlnet_block_samples=controlnet_block_res_samples) - - -class SD3MultiControlNetModel(ModelMixin): - r""" - `SD3ControlNetModel` wrapper class for Multi-SD3ControlNet - - This module is a wrapper for multiple instances of the `SD3ControlNetModel`. The `forward()` API is designed to be - compatible with `SD3ControlNetModel`. - - Args: - controlnets (`List[SD3ControlNetModel]`): - Provides additional conditioning to the unet during the denoising process. You must set multiple - `SD3ControlNetModel` as a list. - """ - - def __init__(self, controlnets): - super().__init__() - self.nets = nn.ModuleList(controlnets) - - def forward( - self, - hidden_states: torch.Tensor, - 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, - return_dict: bool = True, - ) -> Union[SD3ControlNetOutput, Tuple]: - for i, (image, scale, controlnet) in enumerate(zip(controlnet_cond, conditioning_scale, self.nets)): - block_samples = controlnet( - hidden_states=hidden_states, - timestep=timestep, - encoder_hidden_states=encoder_hidden_states, - pooled_projections=pooled_projections, - controlnet_cond=image, - conditioning_scale=scale, - joint_attention_kwargs=joint_attention_kwargs, - return_dict=return_dict, - ) - - # merge samples - if i == 0: - control_block_samples = block_samples - else: - control_block_samples = [ - control_block_sample + block_sample - for control_block_sample, block_sample in zip(control_block_samples[0], block_samples[0]) - ] - control_block_samples = (tuple(control_block_samples),) - - return control_block_samples diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_sparsectrl.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_sparsectrl.py deleted file mode 100644 index b04be5726234303b1ad192c881a308407c584376..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_sparsectrl.py +++ /dev/null @@ -1,785 +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 dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union - -import torch -from torch import nn -from torch.nn import functional as F - -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import FromOriginalModelMixin -from ...utils import BaseOutput, logging -from ..attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - AttentionProcessor, - AttnAddedKVProcessor, - AttnProcessor, -) -from ..embeddings import TimestepEmbedding, Timesteps -from ..modeling_utils import ModelMixin -from ..unets.unet_2d_blocks import UNetMidBlock2DCrossAttn -from ..unets.unet_2d_condition import UNet2DConditionModel -from ..unets.unet_motion_model import CrossAttnDownBlockMotion, DownBlockMotion - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class SparseControlNetOutput(BaseOutput): - """ - The output of [`SparseControlNetModel`]. - - Args: - down_block_res_samples (`tuple[torch.Tensor]`): - A tuple of downsample activations at different resolutions for each downsampling block. Each tensor should - be of shape `(batch_size, channel * resolution, height //resolution, width // resolution)`. Output can be - used to condition the original UNet's downsampling activations. - mid_down_block_re_sample (`torch.Tensor`): - The activation of the middle block (the lowest sample resolution). Each tensor should be of shape - `(batch_size, channel * lowest_resolution, height // lowest_resolution, width // lowest_resolution)`. - Output can be used to condition the original UNet's middle block activation. - """ - - down_block_res_samples: Tuple[torch.Tensor] - mid_block_res_sample: torch.Tensor - - -class SparseControlNetConditioningEmbedding(nn.Module): - def __init__( - self, - conditioning_embedding_channels: int, - conditioning_channels: int = 3, - block_out_channels: Tuple[int, ...] = (16, 32, 96, 256), - ): - super().__init__() - - self.conv_in = nn.Conv2d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1) - self.blocks = nn.ModuleList([]) - - for i in range(len(block_out_channels) - 1): - channel_in = block_out_channels[i] - channel_out = block_out_channels[i + 1] - self.blocks.append(nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1)) - self.blocks.append(nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1, stride=2)) - - self.conv_out = zero_module( - nn.Conv2d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1) - ) - - def forward(self, conditioning: torch.Tensor) -> torch.Tensor: - embedding = self.conv_in(conditioning) - embedding = F.silu(embedding) - - for block in self.blocks: - embedding = block(embedding) - embedding = F.silu(embedding) - - embedding = self.conv_out(embedding) - return embedding - - -class SparseControlNetModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): - """ - A SparseControlNet model as described in [SparseCtrl: Adding Sparse Controls to Text-to-Video Diffusion - Models](https://huggingface.co/papers/2311.16933). - - Args: - in_channels (`int`, defaults to 4): - The number of channels in the input sample. - conditioning_channels (`int`, defaults to 4): - The number of input channels in the controlnet conditional embedding module. If - `concat_condition_embedding` is True, the value provided here is incremented by 1. - flip_sin_to_cos (`bool`, defaults to `True`): - Whether to flip the sin to cos in the time embedding. - freq_shift (`int`, defaults to 0): - 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`): - 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): - The number of layers per block. - downsample_padding (`int`, defaults to 1): - The padding to use for the downsampling convolution. - mid_block_scale_factor (`float`, defaults to 1): - The scale factor to use for the mid block. - act_fn (`str`, defaults to "silu"): - The activation function to use. - 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`, defaults to 1e-5): - 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): - 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): - The number of transformer layers to use in each layer in the middle block. - attention_head_dim (`int` or `Tuple[int]`, defaults to 8): - The dimension of the attention heads. - 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)`): - 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`): - 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`): - """ - - _supports_gradient_checkpointing = True - - @register_to_config - 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, - ): - super().__init__() - self.use_simplified_condition_embedding = use_simplified_condition_embedding - - # If `num_attention_heads` is not defined (which is the case for most models) - # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. - # The reason for this behavior is to correct for incorrectly named variables that were introduced - # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 - # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking - # which is why we correct for the naming here. - num_attention_heads = num_attention_heads or attention_head_dim - - # Check inputs - if len(block_out_channels) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." - ) - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) - if isinstance(temporal_transformer_layers_per_block, int): - temporal_transformer_layers_per_block = [temporal_transformer_layers_per_block] * len(down_block_types) - - # input - conv_in_kernel = 3 - conv_in_padding = (conv_in_kernel - 1) // 2 - self.conv_in = nn.Conv2d( - in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding - ) - - if concat_conditioning_mask: - conditioning_channels = conditioning_channels + 1 - - self.concat_conditioning_mask = concat_conditioning_mask - - # control net conditioning embedding - if use_simplified_condition_embedding: - self.controlnet_cond_embedding = zero_module( - nn.Conv2d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1) - ) - else: - self.controlnet_cond_embedding = SparseControlNetConditioningEmbedding( - conditioning_embedding_channels=block_out_channels[0], - block_out_channels=conditioning_embedding_out_channels, - conditioning_channels=conditioning_channels, - ) - - # time - time_embed_dim = block_out_channels[0] * 4 - self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) - timestep_input_dim = block_out_channels[0] - - self.time_embedding = TimestepEmbedding( - timestep_input_dim, - time_embed_dim, - act_fn=act_fn, - ) - - self.down_blocks = nn.ModuleList([]) - self.controlnet_down_blocks = nn.ModuleList([]) - - if isinstance(cross_attention_dim, int): - cross_attention_dim = (cross_attention_dim,) * len(down_block_types) - - if isinstance(only_cross_attention, bool): - only_cross_attention = [only_cross_attention] * len(down_block_types) - - if isinstance(attention_head_dim, int): - attention_head_dim = (attention_head_dim,) * len(down_block_types) - - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(down_block_types) - - if isinstance(motion_num_attention_heads, int): - motion_num_attention_heads = (motion_num_attention_heads,) * len(down_block_types) - - # down - output_channel = block_out_channels[0] - - controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1) - controlnet_block = zero_module(controlnet_block) - self.controlnet_down_blocks.append(controlnet_block) - - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - if down_block_type == "CrossAttnDownBlockMotion": - down_block = CrossAttnDownBlockMotion( - in_channels=input_channel, - out_channels=output_channel, - temb_channels=time_embed_dim, - dropout=0, - num_layers=layers_per_block, - transformer_layers_per_block=transformer_layers_per_block[i], - resnet_eps=norm_eps, - resnet_time_scale_shift=resnet_time_scale_shift, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - resnet_pre_norm=True, - num_attention_heads=num_attention_heads[i], - cross_attention_dim=cross_attention_dim[i], - add_downsample=not is_final_block, - dual_cross_attention=False, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention[i], - upcast_attention=upcast_attention, - temporal_num_attention_heads=motion_num_attention_heads[i], - temporal_max_seq_length=motion_max_seq_length, - temporal_transformer_layers_per_block=temporal_transformer_layers_per_block[i], - temporal_double_self_attention=False, - ) - elif down_block_type == "DownBlockMotion": - down_block = DownBlockMotion( - in_channels=input_channel, - out_channels=output_channel, - temb_channels=time_embed_dim, - dropout=0, - num_layers=layers_per_block, - resnet_eps=norm_eps, - resnet_time_scale_shift=resnet_time_scale_shift, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - resnet_pre_norm=True, - add_downsample=not is_final_block, - temporal_num_attention_heads=motion_num_attention_heads[i], - temporal_max_seq_length=motion_max_seq_length, - temporal_transformer_layers_per_block=temporal_transformer_layers_per_block[i], - temporal_double_self_attention=False, - ) - else: - raise ValueError( - "Invalid `block_type` encountered. Must be one of `CrossAttnDownBlockMotion` or `DownBlockMotion`" - ) - - self.down_blocks.append(down_block) - - for _ in range(layers_per_block): - controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1) - controlnet_block = zero_module(controlnet_block) - self.controlnet_down_blocks.append(controlnet_block) - - if not is_final_block: - controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1) - controlnet_block = zero_module(controlnet_block) - self.controlnet_down_blocks.append(controlnet_block) - - # mid - mid_block_channels = block_out_channels[-1] - - controlnet_block = nn.Conv2d(mid_block_channels, mid_block_channels, kernel_size=1) - controlnet_block = zero_module(controlnet_block) - self.controlnet_mid_block = controlnet_block - - if transformer_layers_per_mid_block is None: - transformer_layers_per_mid_block = ( - transformer_layers_per_block[-1] if isinstance(transformer_layers_per_block[-1], int) else 1 - ) - - self.mid_block = UNetMidBlock2DCrossAttn( - in_channels=mid_block_channels, - temb_channels=time_embed_dim, - dropout=0, - num_layers=1, - transformer_layers_per_block=transformer_layers_per_mid_block, - resnet_eps=norm_eps, - resnet_time_scale_shift=resnet_time_scale_shift, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - resnet_pre_norm=True, - num_attention_heads=num_attention_heads[-1], - output_scale_factor=mid_block_scale_factor, - cross_attention_dim=cross_attention_dim[-1], - dual_cross_attention=False, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - attention_type="default", - ) - - @classmethod - def from_unet( - cls, - unet: UNet2DConditionModel, - controlnet_conditioning_channel_order: str = "rgb", - conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256), - load_weights_from_unet: bool = True, - conditioning_channels: int = 3, - ) -> "SparseControlNetModel": - r""" - Instantiate a [`SparseControlNetModel`] from [`UNet2DConditionModel`]. - - Parameters: - unet (`UNet2DConditionModel`): - The UNet model weights to copy to the [`SparseControlNetModel`]. All configuration options are also - copied where applicable. - """ - transformer_layers_per_block = ( - unet.config.transformer_layers_per_block if "transformer_layers_per_block" in unet.config else 1 - ) - down_block_types = unet.config.down_block_types - - for i in range(len(down_block_types)): - if "CrossAttn" in down_block_types[i]: - down_block_types[i] = "CrossAttnDownBlockMotion" - elif "Down" in down_block_types[i]: - down_block_types[i] = "DownBlockMotion" - else: - raise ValueError("Invalid `block_type` encountered. Must be a cross-attention or down block") - - controlnet = cls( - in_channels=unet.config.in_channels, - conditioning_channels=conditioning_channels, - flip_sin_to_cos=unet.config.flip_sin_to_cos, - freq_shift=unet.config.freq_shift, - down_block_types=unet.config.down_block_types, - only_cross_attention=unet.config.only_cross_attention, - block_out_channels=unet.config.block_out_channels, - layers_per_block=unet.config.layers_per_block, - downsample_padding=unet.config.downsample_padding, - mid_block_scale_factor=unet.config.mid_block_scale_factor, - act_fn=unet.config.act_fn, - norm_num_groups=unet.config.norm_num_groups, - norm_eps=unet.config.norm_eps, - cross_attention_dim=unet.config.cross_attention_dim, - transformer_layers_per_block=transformer_layers_per_block, - attention_head_dim=unet.config.attention_head_dim, - num_attention_heads=unet.config.num_attention_heads, - use_linear_projection=unet.config.use_linear_projection, - upcast_attention=unet.config.upcast_attention, - resnet_time_scale_shift=unet.config.resnet_time_scale_shift, - conditioning_embedding_out_channels=conditioning_embedding_out_channels, - controlnet_conditioning_channel_order=controlnet_conditioning_channel_order, - ) - - if load_weights_from_unet: - controlnet.conv_in.load_state_dict(unet.conv_in.state_dict(), strict=False) - controlnet.time_proj.load_state_dict(unet.time_proj.state_dict(), strict=False) - controlnet.time_embedding.load_state_dict(unet.time_embedding.state_dict(), strict=False) - controlnet.down_blocks.load_state_dict(unet.down_blocks.state_dict(), strict=False) - controlnet.mid_block.load_state_dict(unet.mid_block.state_dict(), strict=False) - - return controlnet - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - 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: - 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. This is useful for saving some memory in exchange for a small decrease in speed. - - Args: - slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): - When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If - `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is - provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` - must be a multiple of `slice_size`. - """ - sliceable_head_dims = [] - - def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): - if hasattr(module, "set_attention_slice"): - sliceable_head_dims.append(module.sliceable_head_dim) - - for child in module.children(): - fn_recursive_retrieve_sliceable_dims(child) - - # retrieve number of attention layers - for module in self.children(): - fn_recursive_retrieve_sliceable_dims(module) - - num_sliceable_layers = len(sliceable_head_dims) - - if slice_size == "auto": - # half the attention head size is usually a good trade-off between - # speed and memory - slice_size = [dim // 2 for dim in sliceable_head_dims] - elif slice_size == "max": - # make smallest slice possible - slice_size = num_sliceable_layers * [1] - - slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size - - if len(slice_size) != len(sliceable_head_dims): - raise ValueError( - f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" - f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." - ) - - for i in range(len(slice_size)): - size = slice_size[i] - dim = sliceable_head_dims[i] - if size is not None and size > dim: - raise ValueError(f"size {size} has to be smaller or equal to {dim}.") - - # 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]): - if hasattr(module, "set_attention_slice"): - module.set_attention_slice(slice_size.pop()) - - for child in module.children(): - fn_recursive_set_attention_slice(child, slice_size) - - reversed_slice_size = list(reversed(slice_size)) - for module in self.children(): - fn_recursive_set_attention_slice(module, reversed_slice_size) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[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, - guess_mode: bool = False, - return_dict: bool = True, - ) -> Union[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]`): - The number of timesteps to denoise an input. - encoder_hidden_states (`torch.Tensor`): - The encoder hidden states. - controlnet_cond (`torch.Tensor`): - The conditional input tensor of shape `(batch_size, sequence_length, hidden_size)`. - conditioning_scale (`float`, defaults to `1.0`): - The scale factor for ControlNet outputs. - class_labels (`torch.Tensor`, *optional*, defaults to `None`): - Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. - timestep_cond (`torch.Tensor`, *optional*, defaults to `None`): - Additional conditional embeddings for timestep. If provided, the embeddings will be summed with the - timestep_embedding passed through the `self.time_embedding` layer to obtain the final timestep - embeddings. - attention_mask (`torch.Tensor`, *optional*, defaults to `None`): - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask - is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large - negative values to the attention scores corresponding to "discard" tokens. - added_cond_kwargs (`dict`): - Additional conditions for the Stable Diffusion XL UNet. - cross_attention_kwargs (`dict[str]`, *optional*, defaults to `None`): - A kwargs dictionary that if specified is passed along to the `AttnProcessor`. - guess_mode (`bool`, defaults to `False`): - In this mode, the ControlNet encoder tries its best to recognize the input content of the input even if - you remove all prompts. A `guidance_scale` between 3.0 and 5.0 is recommended. - return_dict (`bool`, defaults to `True`): - Whether or not to return a [`~models.controlnet.ControlNetOutput`] instead of a plain tuple. - Returns: - [`~models.controlnet.ControlNetOutput`] **or** `tuple`: - If `return_dict` is `True`, a [`~models.controlnet.ControlNetOutput`] is returned, otherwise a tuple is - returned where the first element is the sample tensor. - """ - sample_batch_size, sample_channels, sample_num_frames, sample_height, sample_width = sample.shape - sample = torch.zeros_like(sample) - - # check channel order - channel_order = self.config.controlnet_conditioning_channel_order - - if channel_order == "rgb": - # in rgb order by default - ... - elif channel_order == "bgr": - controlnet_cond = torch.flip(controlnet_cond, dims=[1]) - else: - raise ValueError(f"unknown `controlnet_conditioning_channel_order`: {channel_order}") - - # prepare attention_mask - if attention_mask is not None: - attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 - attention_mask = attention_mask.unsqueeze(1) - - # 1. time - 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 - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = sample.device.type == "mps" - is_npu = sample.device.type == "npu" - if isinstance(timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) - elif len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps.expand(sample.shape[0]) - - t_emb = self.time_proj(timesteps) - - # timesteps does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - t_emb = t_emb.to(dtype=sample.dtype) - - emb = self.time_embedding(t_emb, timestep_cond) - emb = emb.repeat_interleave(sample_num_frames, dim=0, output_size=emb.shape[0] * sample_num_frames) - - # 2. pre-process - batch_size, channels, num_frames, height, width = sample.shape - - sample = sample.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) - sample = self.conv_in(sample) - - batch_frames, channels, height, width = sample.shape - sample = sample[:, None].reshape(sample_batch_size, sample_num_frames, channels, height, width) - - if self.concat_conditioning_mask: - controlnet_cond = torch.cat([controlnet_cond, conditioning_mask], dim=1) - - batch_size, channels, num_frames, height, width = controlnet_cond.shape - controlnet_cond = controlnet_cond.permute(0, 2, 1, 3, 4).reshape( - batch_size * num_frames, channels, height, width - ) - controlnet_cond = self.controlnet_cond_embedding(controlnet_cond) - batch_frames, channels, height, width = controlnet_cond.shape - controlnet_cond = controlnet_cond[:, None].reshape(batch_size, num_frames, channels, height, width) - - sample = sample + controlnet_cond - - batch_size, num_frames, channels, height, width = sample.shape - sample = sample.reshape(sample_batch_size * sample_num_frames, channels, height, width) - - # 3. down - down_block_res_samples = (sample,) - for downsample_block in self.down_blocks: - if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: - sample, res_samples = downsample_block( - hidden_states=sample, - temb=emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - ) - else: - sample, res_samples = downsample_block(hidden_states=sample, temb=emb, num_frames=num_frames) - - down_block_res_samples += res_samples - - # 4. mid - if self.mid_block is not None: - if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention: - sample = self.mid_block( - sample, - emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - ) - else: - sample = self.mid_block(sample, emb) - - # 5. Control net blocks - controlnet_down_block_res_samples = () - - for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks): - down_block_res_sample = controlnet_block(down_block_res_sample) - controlnet_down_block_res_samples = controlnet_down_block_res_samples + (down_block_res_sample,) - - down_block_res_samples = controlnet_down_block_res_samples - mid_block_res_sample = self.controlnet_mid_block(sample) - - # 6. scaling - if guess_mode and not self.config.global_pool_conditions: - scales = torch.logspace(-1, 0, len(down_block_res_samples) + 1, device=sample.device) # 0.1 to 1.0 - scales = scales * conditioning_scale - down_block_res_samples = [sample * scale for sample, scale in zip(down_block_res_samples, scales)] - mid_block_res_sample = mid_block_res_sample * scales[-1] # last one - else: - down_block_res_samples = [sample * conditioning_scale for sample in down_block_res_samples] - mid_block_res_sample = mid_block_res_sample * conditioning_scale - - if self.config.global_pool_conditions: - down_block_res_samples = [ - torch.mean(sample, dim=(2, 3), keepdim=True) for sample in down_block_res_samples - ] - mid_block_res_sample = torch.mean(mid_block_res_sample, dim=(2, 3), keepdim=True) - - if not return_dict: - return (down_block_res_samples, mid_block_res_sample) - - return SparseControlNetOutput( - down_block_res_samples=down_block_res_samples, mid_block_res_sample=mid_block_res_sample - ) - - -# Copied from diffusers.models.controlnets.controlnet.zero_module -def zero_module(module: nn.Module) -> nn.Module: - for p in module.parameters(): - nn.init.zeros_(p) - return module diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_union.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_union.py deleted file mode 100644 index d731ef137b6cecf17e55c2823cd74280b2363052..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_union.py +++ /dev/null @@ -1,841 +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 Any, Dict, List, Optional, Tuple, Union - -import torch -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders.single_file_model import FromOriginalModelMixin -from ...utils import logging -from ..attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - AttentionProcessor, - AttnAddedKVProcessor, - AttnProcessor, -) -from ..embeddings import TextImageTimeEmbedding, TextTimeEmbedding, TimestepEmbedding, Timesteps -from ..modeling_utils import ModelMixin -from ..unets.unet_2d_blocks import ( - UNetMidBlock2DCrossAttn, - get_down_block, -) -from ..unets.unet_2d_condition import UNet2DConditionModel -from .controlnet import ControlNetConditioningEmbedding, ControlNetOutput, zero_module - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class QuickGELU(nn.Module): - """ - Applies GELU approximation that is fast but somewhat inaccurate. See: https://github.com/hendrycks/GELUs - """ - - def forward(self, input: torch.Tensor) -> torch.Tensor: - return input * torch.sigmoid(1.702 * input) - - -class ResidualAttentionMlp(nn.Module): - def __init__(self, d_model: int): - super().__init__() - self.c_fc = nn.Linear(d_model, d_model * 4) - self.gelu = QuickGELU() - self.c_proj = nn.Linear(d_model * 4, d_model) - - def forward(self, x: torch.Tensor): - x = self.c_fc(x) - x = self.gelu(x) - x = self.c_proj(x) - return x - - -class ResidualAttentionBlock(nn.Module): - def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None): - super().__init__() - self.attn = nn.MultiheadAttention(d_model, n_head) - self.ln_1 = nn.LayerNorm(d_model) - self.mlp = ResidualAttentionMlp(d_model) - self.ln_2 = nn.LayerNorm(d_model) - self.attn_mask = attn_mask - - def attention(self, x: torch.Tensor): - self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None - return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0] - - def forward(self, x: torch.Tensor): - x = x + self.attention(self.ln_1(x)) - x = x + self.mlp(self.ln_2(x)) - return x - - -class ControlNetUnionModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): - """ - A ControlNetUnion model. - - Args: - in_channels (`int`, defaults to 4): - The number of channels in the input sample. - flip_sin_to_cos (`bool`, defaults to `True`): - Whether to flip the sin to cos in the time embedding. - freq_shift (`int`, defaults to 0): - 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`): - 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): - The number of layers per block. - downsample_padding (`int`, defaults to 1): - The padding to use for the downsampling convolution. - mid_block_scale_factor (`float`, defaults to 1): - The scale factor to use for the mid block. - act_fn (`str`, defaults to "silu"): - The activation function to use. - 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`, defaults to 1e-5): - 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): - 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`]. - encoder_hid_dim (`int`, *optional*, defaults to None): - If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim` - dimension to `cross_attention_dim`. - 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): - The dimension of the attention heads. - use_linear_projection (`bool`, defaults to `False`): - class_embed_type (`str`, *optional*, defaults to `None`): - The type of class embedding to use which is ultimately summed with the time embeddings. Choose from None, - `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`. - addition_embed_type (`str`, *optional*, defaults to `None`): - Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or - "text". "text" will use the `TextTimeEmbedding` layer. - num_class_embeds (`int`, *optional*, defaults to 0): - Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing - class conditioning with `class_embed_type` equal to `None`. - 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`. - projection_class_embeddings_input_dim (`int`, *optional*, defaults to `None`): - The dimension of the `class_labels` input when `class_embed_type="projection"`. Required when - `class_embed_type="projection"`. - controlnet_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]`, *optional*, defaults to `(48, 96, 192, 384)`): - The tuple of output channel for each block in the `conditioning_embedding` layer. - global_pool_conditions (`bool`, defaults to `False`): - """ - - _supports_gradient_checkpointing = True - - @register_to_config - 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", - ), - 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, ...]] = (48, 96, 192, 384), - global_pool_conditions: bool = False, - addition_embed_type_num_heads: int = 64, - num_control_type: int = 6, - num_trans_channel: int = 320, - num_trans_head: int = 8, - num_trans_layer: int = 1, - num_proj_channel: int = 320, - ): - super().__init__() - - # If `num_attention_heads` is not defined (which is the case for most models) - # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. - # The reason for this behavior is to correct for incorrectly named variables that were introduced - # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 - # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking - # which is why we correct for the naming here. - num_attention_heads = num_attention_heads or attention_head_dim - - # Check inputs - if len(block_out_channels) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." - ) - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) - - # input - conv_in_kernel = 3 - conv_in_padding = (conv_in_kernel - 1) // 2 - self.conv_in = nn.Conv2d( - in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding - ) - - # time - time_embed_dim = block_out_channels[0] * 4 - self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) - timestep_input_dim = block_out_channels[0] - self.time_embedding = TimestepEmbedding( - timestep_input_dim, - time_embed_dim, - act_fn=act_fn, - ) - - if encoder_hid_dim_type is not None: - raise ValueError(f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None.") - else: - self.encoder_hid_proj = None - - # class embedding - if class_embed_type is None and num_class_embeds is not None: - self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) - elif class_embed_type == "timestep": - self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim) - elif class_embed_type == "identity": - self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) - elif class_embed_type == "projection": - if projection_class_embeddings_input_dim is None: - raise ValueError( - "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set" - ) - # The projection `class_embed_type` is the same as the timestep `class_embed_type` except - # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings - # 2. it projects from an arbitrary input dimension. - # - # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations. - # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings. - # As a result, `TimestepEmbedding` can be passed arbitrary vectors. - self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) - else: - self.class_embedding = None - - if addition_embed_type == "text": - if encoder_hid_dim is not None: - text_time_embedding_from_dim = encoder_hid_dim - else: - text_time_embedding_from_dim = cross_attention_dim - - self.add_embedding = TextTimeEmbedding( - text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads - ) - elif addition_embed_type == "text_image": - # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much - # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use - # case when `addition_embed_type == "text_image"` (Kandinsky 2.1)` - self.add_embedding = TextImageTimeEmbedding( - text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim - ) - elif addition_embed_type == "text_time": - self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift) - self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) - - elif addition_embed_type is not None: - raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.") - - # control net conditioning embedding - self.controlnet_cond_embedding = ControlNetConditioningEmbedding( - conditioning_embedding_channels=block_out_channels[0], - block_out_channels=conditioning_embedding_out_channels, - conditioning_channels=conditioning_channels, - ) - - task_scale_factor = num_trans_channel**0.5 - self.task_embedding = nn.Parameter(task_scale_factor * torch.randn(num_control_type, num_trans_channel)) - self.transformer_layes = nn.ModuleList( - [ResidualAttentionBlock(num_trans_channel, num_trans_head) for _ in range(num_trans_layer)] - ) - self.spatial_ch_projs = zero_module(nn.Linear(num_trans_channel, num_proj_channel)) - self.control_type_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift) - self.control_add_embedding = TimestepEmbedding(addition_time_embed_dim * num_control_type, time_embed_dim) - - self.down_blocks = nn.ModuleList([]) - self.controlnet_down_blocks = nn.ModuleList([]) - - if isinstance(only_cross_attention, bool): - only_cross_attention = [only_cross_attention] * len(down_block_types) - - if isinstance(attention_head_dim, int): - attention_head_dim = (attention_head_dim,) * len(down_block_types) - - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(down_block_types) - - # down - output_channel = block_out_channels[0] - - controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1) - controlnet_block = zero_module(controlnet_block) - self.controlnet_down_blocks.append(controlnet_block) - - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = get_down_block( - down_block_type, - num_layers=layers_per_block, - transformer_layers_per_block=transformer_layers_per_block[i], - in_channels=input_channel, - out_channels=output_channel, - temb_channels=time_embed_dim, - add_downsample=not is_final_block, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads[i], - attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, - downsample_padding=downsample_padding, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention[i], - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - self.down_blocks.append(down_block) - - for _ in range(layers_per_block): - controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1) - controlnet_block = zero_module(controlnet_block) - self.controlnet_down_blocks.append(controlnet_block) - - if not is_final_block: - controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1) - controlnet_block = zero_module(controlnet_block) - self.controlnet_down_blocks.append(controlnet_block) - - # mid - mid_block_channel = block_out_channels[-1] - - controlnet_block = nn.Conv2d(mid_block_channel, mid_block_channel, kernel_size=1) - controlnet_block = zero_module(controlnet_block) - self.controlnet_mid_block = controlnet_block - - self.mid_block = UNetMidBlock2DCrossAttn( - transformer_layers_per_block=transformer_layers_per_block[-1], - in_channels=mid_block_channel, - temb_channels=time_embed_dim, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - output_scale_factor=mid_block_scale_factor, - resnet_time_scale_shift=resnet_time_scale_shift, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads[-1], - resnet_groups=norm_num_groups, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - ) - - @classmethod - def from_unet( - cls, - unet: UNet2DConditionModel, - controlnet_conditioning_channel_order: str = "rgb", - conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256), - load_weights_from_unet: bool = True, - ): - r""" - Instantiate a [`ControlNetUnionModel`] from [`UNet2DConditionModel`]. - - Parameters: - unet (`UNet2DConditionModel`): - The UNet model weights to copy to the [`ControlNetUnionModel`]. All configuration options are also - copied where applicable. - """ - transformer_layers_per_block = ( - unet.config.transformer_layers_per_block if "transformer_layers_per_block" in unet.config else 1 - ) - encoder_hid_dim = unet.config.encoder_hid_dim if "encoder_hid_dim" in unet.config else None - encoder_hid_dim_type = unet.config.encoder_hid_dim_type if "encoder_hid_dim_type" in unet.config else None - addition_embed_type = unet.config.addition_embed_type if "addition_embed_type" in unet.config else None - addition_time_embed_dim = ( - unet.config.addition_time_embed_dim if "addition_time_embed_dim" in unet.config else None - ) - - controlnet = cls( - encoder_hid_dim=encoder_hid_dim, - encoder_hid_dim_type=encoder_hid_dim_type, - addition_embed_type=addition_embed_type, - addition_time_embed_dim=addition_time_embed_dim, - transformer_layers_per_block=transformer_layers_per_block, - in_channels=unet.config.in_channels, - flip_sin_to_cos=unet.config.flip_sin_to_cos, - freq_shift=unet.config.freq_shift, - down_block_types=unet.config.down_block_types, - only_cross_attention=unet.config.only_cross_attention, - block_out_channels=unet.config.block_out_channels, - layers_per_block=unet.config.layers_per_block, - downsample_padding=unet.config.downsample_padding, - mid_block_scale_factor=unet.config.mid_block_scale_factor, - act_fn=unet.config.act_fn, - norm_num_groups=unet.config.norm_num_groups, - norm_eps=unet.config.norm_eps, - cross_attention_dim=unet.config.cross_attention_dim, - attention_head_dim=unet.config.attention_head_dim, - num_attention_heads=unet.config.num_attention_heads, - use_linear_projection=unet.config.use_linear_projection, - class_embed_type=unet.config.class_embed_type, - num_class_embeds=unet.config.num_class_embeds, - upcast_attention=unet.config.upcast_attention, - resnet_time_scale_shift=unet.config.resnet_time_scale_shift, - projection_class_embeddings_input_dim=unet.config.projection_class_embeddings_input_dim, - controlnet_conditioning_channel_order=controlnet_conditioning_channel_order, - conditioning_embedding_out_channels=conditioning_embedding_out_channels, - ) - - if load_weights_from_unet: - controlnet.conv_in.load_state_dict(unet.conv_in.state_dict()) - controlnet.time_proj.load_state_dict(unet.time_proj.state_dict()) - controlnet.time_embedding.load_state_dict(unet.time_embedding.state_dict()) - - if controlnet.class_embedding: - controlnet.class_embedding.load_state_dict(unet.class_embedding.state_dict()) - - controlnet.down_blocks.load_state_dict(unet.down_blocks.state_dict(), strict=False) - controlnet.mid_block.load_state_dict(unet.mid_block.state_dict(), strict=False) - - return controlnet - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - 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: - 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. This is useful for saving some memory in exchange for a small decrease in speed. - - Args: - slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): - When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If - `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is - provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` - must be a multiple of `slice_size`. - """ - sliceable_head_dims = [] - - def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): - if hasattr(module, "set_attention_slice"): - sliceable_head_dims.append(module.sliceable_head_dim) - - for child in module.children(): - fn_recursive_retrieve_sliceable_dims(child) - - # retrieve number of attention layers - for module in self.children(): - fn_recursive_retrieve_sliceable_dims(module) - - num_sliceable_layers = len(sliceable_head_dims) - - if slice_size == "auto": - # half the attention head size is usually a good trade-off between - # speed and memory - slice_size = [dim // 2 for dim in sliceable_head_dims] - elif slice_size == "max": - # make smallest slice possible - slice_size = num_sliceable_layers * [1] - - slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size - - if len(slice_size) != len(sliceable_head_dims): - raise ValueError( - f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" - f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." - ) - - for i in range(len(slice_size)): - size = slice_size[i] - dim = sliceable_head_dims[i] - if size is not None and size > dim: - raise ValueError(f"size {size} has to be smaller or equal to {dim}.") - - # 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]): - if hasattr(module, "set_attention_slice"): - module.set_attention_slice(slice_size.pop()) - - for child in module.children(): - fn_recursive_set_attention_slice(child, slice_size) - - reversed_slice_size = list(reversed(slice_size)) - for module in self.children(): - fn_recursive_set_attention_slice(module, reversed_slice_size) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], - encoder_hidden_states: 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, - from_multi: bool = False, - guess_mode: bool = False, - return_dict: bool = True, - ) -> Union[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]`): - The number of timesteps to denoise an input. - encoder_hidden_states (`torch.Tensor`): - The encoder hidden states. - 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]`): - The indices of `control_type`. - conditioning_scale (`float`, defaults to `1.0`): - The scale factor for ControlNet outputs. - class_labels (`torch.Tensor`, *optional*, defaults to `None`): - Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. - timestep_cond (`torch.Tensor`, *optional*, defaults to `None`): - Additional conditional embeddings for timestep. If provided, the embeddings will be summed with the - timestep_embedding passed through the `self.time_embedding` layer to obtain the final timestep - embeddings. - attention_mask (`torch.Tensor`, *optional*, defaults to `None`): - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask - is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large - negative values to the attention scores corresponding to "discard" tokens. - added_cond_kwargs (`dict`): - Additional conditions for the Stable Diffusion XL UNet. - cross_attention_kwargs (`dict[str]`, *optional*, defaults to `None`): - A kwargs dictionary that if specified is passed along to the `AttnProcessor`. - from_multi (`bool`, defaults to `False`): - Use standard scaling when called from `MultiControlNetUnionModel`. - guess_mode (`bool`, defaults to `False`): - In this mode, the ControlNet encoder tries its best to recognize the input content of the input even if - you remove all prompts. A `guidance_scale` between 3.0 and 5.0 is recommended. - return_dict (`bool`, defaults to `True`): - Whether or not to return a [`~models.controlnet.ControlNetOutput`] instead of a plain tuple. - - Returns: - [`~models.controlnet.ControlNetOutput`] **or** `tuple`: - If `return_dict` is `True`, a [`~models.controlnet.ControlNetOutput`] is returned, otherwise a tuple is - returned where the first element is the sample tensor. - """ - if isinstance(conditioning_scale, float): - conditioning_scale = [conditioning_scale] * len(controlnet_cond) - - # check channel order - channel_order = self.config.controlnet_conditioning_channel_order - - if channel_order != "rgb": - raise ValueError(f"unknown `controlnet_conditioning_channel_order`: {channel_order}") - - # prepare attention_mask - if attention_mask is not None: - attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 - attention_mask = attention_mask.unsqueeze(1) - - # 1. time - 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 - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = sample.device.type == "mps" - is_npu = sample.device.type == "npu" - if isinstance(timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) - elif len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps.expand(sample.shape[0]) - - t_emb = self.time_proj(timesteps) - - # timesteps does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - t_emb = t_emb.to(dtype=sample.dtype) - - emb = self.time_embedding(t_emb, timestep_cond) - aug_emb = None - - if self.class_embedding is not None: - if class_labels is None: - raise ValueError("class_labels should be provided when num_class_embeds > 0") - - if self.config.class_embed_type == "timestep": - class_labels = self.time_proj(class_labels) - - class_emb = self.class_embedding(class_labels).to(dtype=self.dtype) - emb = emb + class_emb - - if self.config.addition_embed_type is not None: - if self.config.addition_embed_type == "text": - aug_emb = self.add_embedding(encoder_hidden_states) - - elif self.config.addition_embed_type == "text_time": - if "text_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" - ) - text_embeds = added_cond_kwargs.get("text_embeds") - if "time_ids" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" - ) - time_ids = added_cond_kwargs.get("time_ids") - time_embeds = self.add_time_proj(time_ids.flatten()) - time_embeds = time_embeds.reshape((text_embeds.shape[0], -1)) - - add_embeds = torch.concat([text_embeds, time_embeds], dim=-1) - add_embeds = add_embeds.to(emb.dtype) - aug_emb = self.add_embedding(add_embeds) - - control_embeds = self.control_type_proj(control_type.flatten()) - control_embeds = control_embeds.reshape((t_emb.shape[0], -1)) - control_embeds = control_embeds.to(emb.dtype) - control_emb = self.control_add_embedding(control_embeds) - emb = emb + control_emb - emb = emb + aug_emb if aug_emb is not None else emb - - # 2. pre-process - sample = self.conv_in(sample) - - inputs = [] - condition_list = [] - - for cond, control_idx, scale in zip(controlnet_cond, control_type_idx, conditioning_scale): - condition = self.controlnet_cond_embedding(cond) - feat_seq = torch.mean(condition, dim=(2, 3)) - feat_seq = feat_seq + self.task_embedding[control_idx] - if from_multi: - inputs.append(feat_seq.unsqueeze(1)) - condition_list.append(condition) - else: - inputs.append(feat_seq.unsqueeze(1) * scale) - condition_list.append(condition * scale) - - condition = sample - feat_seq = torch.mean(condition, dim=(2, 3)) - inputs.append(feat_seq.unsqueeze(1)) - condition_list.append(condition) - - x = torch.cat(inputs, dim=1) - for layer in self.transformer_layes: - x = layer(x) - - controlnet_cond_fuser = sample * 0.0 - for (idx, condition), scale in zip(enumerate(condition_list[:-1]), conditioning_scale): - alpha = self.spatial_ch_projs(x[:, idx]) - alpha = alpha.unsqueeze(-1).unsqueeze(-1) - if from_multi: - controlnet_cond_fuser += condition + alpha - else: - controlnet_cond_fuser += condition + alpha * scale - - sample = sample + controlnet_cond_fuser - - # 3. down - down_block_res_samples = (sample,) - for downsample_block in self.down_blocks: - if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: - sample, res_samples = downsample_block( - hidden_states=sample, - temb=emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - ) - else: - sample, res_samples = downsample_block(hidden_states=sample, temb=emb) - - down_block_res_samples += res_samples - - # 4. mid - if self.mid_block is not None: - sample = self.mid_block( - sample, - emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - ) - - # 5. Control net blocks - controlnet_down_block_res_samples = () - - for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks): - down_block_res_sample = controlnet_block(down_block_res_sample) - controlnet_down_block_res_samples = controlnet_down_block_res_samples + (down_block_res_sample,) - - down_block_res_samples = controlnet_down_block_res_samples - - mid_block_res_sample = self.controlnet_mid_block(sample) - - # 6. scaling - if guess_mode and not self.config.global_pool_conditions: - scales = torch.logspace(-1, 0, len(down_block_res_samples) + 1, device=sample.device) # 0.1 to 1.0 - if from_multi: - scales = scales * conditioning_scale[0] - down_block_res_samples = [sample * scale for sample, scale in zip(down_block_res_samples, scales)] - mid_block_res_sample = mid_block_res_sample * scales[-1] # last one - elif from_multi: - down_block_res_samples = [sample * conditioning_scale[0] for sample in down_block_res_samples] - mid_block_res_sample = mid_block_res_sample * conditioning_scale[0] - - if self.config.global_pool_conditions: - down_block_res_samples = [ - torch.mean(sample, dim=(2, 3), keepdim=True) for sample in down_block_res_samples - ] - mid_block_res_sample = torch.mean(mid_block_res_sample, dim=(2, 3), keepdim=True) - - if not return_dict: - return (down_block_res_samples, mid_block_res_sample) - - return ControlNetOutput( - down_block_res_samples=down_block_res_samples, mid_block_res_sample=mid_block_res_sample - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_xs.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_xs.py deleted file mode 100644 index aabae709e9886661a1251977371362eac9327c3b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/controlnet_xs.py +++ /dev/null @@ -1,1907 +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 dataclasses import dataclass -from math import gcd -from typing import Any, Dict, List, Optional, Tuple, Union - -import torch -import torch.utils.checkpoint -from torch import Tensor, nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import BaseOutput, logging -from ...utils.torch_utils import apply_freeu -from ..attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - Attention, - AttentionProcessor, - AttnAddedKVProcessor, - AttnProcessor, - FusedAttnProcessor2_0, -) -from ..embeddings import TimestepEmbedding, Timesteps -from ..modeling_utils import ModelMixin -from ..unets.unet_2d_blocks import ( - CrossAttnDownBlock2D, - CrossAttnUpBlock2D, - Downsample2D, - ResnetBlock2D, - Transformer2DModel, - UNetMidBlock2DCrossAttn, - Upsample2D, -) -from ..unets.unet_2d_condition import UNet2DConditionModel -from .controlnet import ControlNetConditioningEmbedding - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class ControlNetXSOutput(BaseOutput): - """ - The output of [`UNetControlNetXSModel`]. - - Args: - sample (`Tensor` of shape `(batch_size, num_channels, height, width)`): - The output of the `UNetControlNetXSModel`. Unlike `ControlNetOutput` this is NOT to be added to the base - model output, but is already the final output. - """ - - sample: Tensor = None - - -class DownBlockControlNetXSAdapter(nn.Module): - """Components that together with corresponding components from the base model will form a - `ControlNetXSCrossAttnDownBlock2D`""" - - def __init__( - self, - resnets: nn.ModuleList, - base_to_ctrl: nn.ModuleList, - ctrl_to_base: nn.ModuleList, - attentions: Optional[nn.ModuleList] = None, - downsampler: Optional[nn.Conv2d] = None, - ): - super().__init__() - self.resnets = resnets - self.base_to_ctrl = base_to_ctrl - self.ctrl_to_base = ctrl_to_base - self.attentions = attentions - self.downsamplers = downsampler - - -class MidBlockControlNetXSAdapter(nn.Module): - """Components that together with corresponding components from the base model will form a - `ControlNetXSCrossAttnMidBlock2D`""" - - def __init__(self, midblock: UNetMidBlock2DCrossAttn, base_to_ctrl: nn.ModuleList, ctrl_to_base: nn.ModuleList): - super().__init__() - self.midblock = midblock - self.base_to_ctrl = base_to_ctrl - self.ctrl_to_base = ctrl_to_base - - -class UpBlockControlNetXSAdapter(nn.Module): - """Components that together with corresponding components from the base model will form a `ControlNetXSCrossAttnUpBlock2D`""" - - def __init__(self, ctrl_to_base: nn.ModuleList): - super().__init__() - self.ctrl_to_base = ctrl_to_base - - -def get_down_block_adapter( - base_in_channels: int, - base_out_channels: int, - ctrl_in_channels: int, - ctrl_out_channels: int, - temb_channels: int, - max_norm_num_groups: Optional[int] = 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, - add_downsample: bool = True, - upcast_attention: Optional[bool] = False, - use_linear_projection: Optional[bool] = True, -): - num_layers = 2 # only support sd + sdxl - - resnets = [] - attentions = [] - ctrl_to_base = [] - base_to_ctrl = [] - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * num_layers - - for i in range(num_layers): - base_in_channels = base_in_channels if i == 0 else base_out_channels - ctrl_in_channels = ctrl_in_channels if i == 0 else ctrl_out_channels - - # Before the resnet/attention application, information is concatted from base to control. - # Concat doesn't require change in number of channels - base_to_ctrl.append(make_zero_conv(base_in_channels, base_in_channels)) - - resnets.append( - ResnetBlock2D( - in_channels=ctrl_in_channels + base_in_channels, # information from base is concatted to ctrl - out_channels=ctrl_out_channels, - temb_channels=temb_channels, - groups=find_largest_factor(ctrl_in_channels + base_in_channels, max_factor=max_norm_num_groups), - groups_out=find_largest_factor(ctrl_out_channels, max_factor=max_norm_num_groups), - eps=1e-5, - ) - ) - - if has_crossattn: - attentions.append( - Transformer2DModel( - num_attention_heads, - ctrl_out_channels // num_attention_heads, - in_channels=ctrl_out_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - norm_num_groups=find_largest_factor(ctrl_out_channels, max_factor=max_norm_num_groups), - ) - ) - - # After the resnet/attention application, information is added from control to base - # Addition requires change in number of channels - ctrl_to_base.append(make_zero_conv(ctrl_out_channels, base_out_channels)) - - if add_downsample: - # Before the downsampler application, information is concatted from base to control - # Concat doesn't require change in number of channels - base_to_ctrl.append(make_zero_conv(base_out_channels, base_out_channels)) - - downsamplers = Downsample2D( - ctrl_out_channels + base_out_channels, use_conv=True, out_channels=ctrl_out_channels, name="op" - ) - - # After the downsampler application, information is added from control to base - # Addition requires change in number of channels - ctrl_to_base.append(make_zero_conv(ctrl_out_channels, base_out_channels)) - else: - downsamplers = None - - down_block_components = DownBlockControlNetXSAdapter( - resnets=nn.ModuleList(resnets), - base_to_ctrl=nn.ModuleList(base_to_ctrl), - ctrl_to_base=nn.ModuleList(ctrl_to_base), - ) - - if has_crossattn: - down_block_components.attentions = nn.ModuleList(attentions) - if downsamplers is not None: - down_block_components.downsamplers = downsamplers - - return down_block_components - - -def get_mid_block_adapter( - base_channels: int, - ctrl_channels: int, - temb_channels: Optional[int] = None, - max_norm_num_groups: Optional[int] = 32, - transformer_layers_per_block: int = 1, - num_attention_heads: Optional[int] = 1, - cross_attention_dim: Optional[int] = 1024, - upcast_attention: bool = False, - use_linear_projection: bool = True, -): - # Before the midblock application, information is concatted from base to control. - # Concat doesn't require change in number of channels - base_to_ctrl = make_zero_conv(base_channels, base_channels) - - midblock = UNetMidBlock2DCrossAttn( - transformer_layers_per_block=transformer_layers_per_block, - in_channels=ctrl_channels + base_channels, - out_channels=ctrl_channels, - temb_channels=temb_channels, - # number or norm groups must divide both in_channels and out_channels - resnet_groups=find_largest_factor(gcd(ctrl_channels, ctrl_channels + base_channels), max_norm_num_groups), - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - ) - - # After the midblock application, information is added from control to base - # Addition requires change in number of channels - ctrl_to_base = make_zero_conv(ctrl_channels, base_channels) - - return MidBlockControlNetXSAdapter(base_to_ctrl=base_to_ctrl, midblock=midblock, ctrl_to_base=ctrl_to_base) - - -def get_up_block_adapter( - out_channels: int, - prev_output_channel: int, - ctrl_skip_channels: List[int], -): - ctrl_to_base = [] - num_layers = 3 # only support sd + sdxl - for i in range(num_layers): - resnet_in_channels = prev_output_channel if i == 0 else out_channels - ctrl_to_base.append(make_zero_conv(ctrl_skip_channels[i], resnet_in_channels)) - - return UpBlockControlNetXSAdapter(ctrl_to_base=nn.ModuleList(ctrl_to_base)) - - -class ControlNetXSAdapter(ModelMixin, ConfigMixin): - r""" - A `ControlNetXSAdapter` model. To use it, pass it into a `UNetControlNetXSModel` (together with a - `UNet2DConditionModel` base model). - - This model inherits from [`ModelMixin`] and [`ConfigMixin`]. Check the superclass documentation for it's generic - methods implemented for all models (such as downloading or saving). - - Like `UNetControlNetXSModel`, `ControlNetXSAdapter` is compatible with StableDiffusion and StableDiffusion-XL. It's - default parameters are compatible with StableDiffusion. - - Parameters: - conditioning_channels (`int`, defaults to 3): - 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)`): - The tuple of output channels for each block in the `controlnet_cond_embedding` layer. - time_embedding_mix (`float`, defaults to 1.0): - If 0, then only the control adapters's time embedding is used. If 1, then only the base unet's time - embedding is used. Otherwise, both are combined. - learn_time_embedding (`bool`, defaults to `False`): - Whether a time embedding should be learned. If yes, `UNetControlNetXSModel` will combine the time - embeddings of the base model and the control adapter. If no, `UNetControlNetXSModel` will use the base - model's time embedding. - num_attention_heads (`list[int]`, defaults to `[4]`): - The number of attention heads. - block_out_channels (`list[int]`, defaults to `[4, 8, 16, 16]`): - The tuple of output channels for each block. - base_block_out_channels (`list[int]`, defaults to `[320, 640, 1280, 1280]`): - The tuple of output channels for each block in the base unet. - cross_attention_dim (`int`, defaults to 1024): - The dimension of the cross attention features. - down_block_types (`list[str]`, defaults to `["CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D"]`): - 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): - 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`): - Whether the attention computation should always be upcasted. - max_norm_num_groups (`int`, defaults to 32): - Maximum number of groups in group normal. The actual number will be the largest divisor of the respective - channels, that is <= max_norm_num_groups. - """ - - @register_to_config - def __init__( - self, - conditioning_channels: int = 3, - conditioning_channel_order: str = "rgb", - 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), - cross_attention_dim: int = 1024, - down_block_types: Tuple[str] = ( - "CrossAttnDownBlock2D", - "CrossAttnDownBlock2D", - "CrossAttnDownBlock2D", - "DownBlock2D", - ), - sample_size: Optional[int] = 96, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - upcast_attention: bool = True, - max_norm_num_groups: int = 32, - use_linear_projection: bool = True, - ): - super().__init__() - - time_embedding_input_dim = base_block_out_channels[0] - time_embedding_dim = base_block_out_channels[0] * 4 - - # Check inputs - if conditioning_channel_order not in ["rgb", "bgr"]: - raise ValueError(f"unknown `conditioning_channel_order`: {conditioning_channel_order}") - - if len(block_out_channels) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(transformer_layers_per_block, (list, tuple)): - transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) - if not isinstance(cross_attention_dim, (list, tuple)): - cross_attention_dim = [cross_attention_dim] * len(down_block_types) - # see https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 for why `ControlNetXSAdapter` takes `num_attention_heads` instead of `attention_head_dim` - if not isinstance(num_attention_heads, (list, tuple)): - num_attention_heads = [num_attention_heads] * len(down_block_types) - - if len(num_attention_heads) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." - ) - - # 5 - Create conditioning hint embedding - self.controlnet_cond_embedding = ControlNetConditioningEmbedding( - conditioning_embedding_channels=block_out_channels[0], - block_out_channels=conditioning_embedding_out_channels, - conditioning_channels=conditioning_channels, - ) - - # time - if learn_time_embedding: - self.time_embedding = TimestepEmbedding(time_embedding_input_dim, time_embedding_dim) - else: - self.time_embedding = None - - self.down_blocks = nn.ModuleList([]) - self.up_connections = nn.ModuleList([]) - - # input - self.conv_in = nn.Conv2d(4, block_out_channels[0], kernel_size=3, padding=1) - self.control_to_base_for_conv_in = make_zero_conv(block_out_channels[0], base_block_out_channels[0]) - - # down - base_out_channels = base_block_out_channels[0] - ctrl_out_channels = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - base_in_channels = base_out_channels - base_out_channels = base_block_out_channels[i] - ctrl_in_channels = ctrl_out_channels - ctrl_out_channels = block_out_channels[i] - has_crossattn = "CrossAttn" in down_block_type - is_final_block = i == len(down_block_types) - 1 - - self.down_blocks.append( - get_down_block_adapter( - base_in_channels=base_in_channels, - base_out_channels=base_out_channels, - ctrl_in_channels=ctrl_in_channels, - ctrl_out_channels=ctrl_out_channels, - temb_channels=time_embedding_dim, - max_norm_num_groups=max_norm_num_groups, - has_crossattn=has_crossattn, - transformer_layers_per_block=transformer_layers_per_block[i], - num_attention_heads=num_attention_heads[i], - cross_attention_dim=cross_attention_dim[i], - add_downsample=not is_final_block, - upcast_attention=upcast_attention, - use_linear_projection=use_linear_projection, - ) - ) - - # mid - self.mid_block = get_mid_block_adapter( - base_channels=base_block_out_channels[-1], - ctrl_channels=block_out_channels[-1], - temb_channels=time_embedding_dim, - transformer_layers_per_block=transformer_layers_per_block[-1], - num_attention_heads=num_attention_heads[-1], - cross_attention_dim=cross_attention_dim[-1], - upcast_attention=upcast_attention, - use_linear_projection=use_linear_projection, - ) - - # up - # The skip connection channels are the output of the conv_in and of all the down subblocks - ctrl_skip_channels = [block_out_channels[0]] - for i, out_channels in enumerate(block_out_channels): - number_of_subblocks = ( - 3 if i < len(block_out_channels) - 1 else 2 - ) # every block has 3 subblocks, except last one, which has 2 as it has no downsampler - ctrl_skip_channels.extend([out_channels] * number_of_subblocks) - - reversed_base_block_out_channels = list(reversed(base_block_out_channels)) - - base_out_channels = reversed_base_block_out_channels[0] - for i in range(len(down_block_types)): - prev_base_output_channel = base_out_channels - base_out_channels = reversed_base_block_out_channels[i] - ctrl_skip_channels_ = [ctrl_skip_channels.pop() for _ in range(3)] - - self.up_connections.append( - get_up_block_adapter( - out_channels=base_out_channels, - prev_output_channel=prev_base_output_channel, - ctrl_skip_channels=ctrl_skip_channels_, - ) - ) - - @classmethod - 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, - 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), - ): - r""" - Instantiate a [`ControlNetXSAdapter`] from a [`UNet2DConditionModel`]. - - Parameters: - unet (`UNet2DConditionModel`): - The UNet model we want to control. The dimensions of the ControlNetXSAdapter will be adapted to it. - 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`): - Down blocks output channels in control model. Either this or `size_ratio` must be given. - 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`): - Whether the `ControlNetXSAdapter` should learn a time embedding. - time_embedding_mix (`float`, defaults to 1.0): - If 0, then only the control adapter's time embedding is used. If 1, then only the base unet's time - embedding is used. Otherwise, both are combined. - conditioning_channels (`int`, defaults to 3): - 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)`): - The tuple of output channel for each block in the `controlnet_cond_embedding` layer. - """ - - # Check input - fixed_size = block_out_channels is not None - relative_size = size_ratio is not None - if not (fixed_size ^ relative_size): - raise ValueError( - "Pass exactly one of `block_out_channels` (for absolute sizing) or `size_ratio` (for relative sizing)." - ) - - # Create model - block_out_channels = block_out_channels or [int(b * size_ratio) for b in unet.config.block_out_channels] - if num_attention_heads is None: - # The naming seems a bit confusing and it is, see https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 for why. - num_attention_heads = unet.config.attention_head_dim - - model = cls( - conditioning_channels=conditioning_channels, - conditioning_channel_order=conditioning_channel_order, - conditioning_embedding_out_channels=conditioning_embedding_out_channels, - time_embedding_mix=time_embedding_mix, - learn_time_embedding=learn_time_embedding, - num_attention_heads=num_attention_heads, - block_out_channels=block_out_channels, - base_block_out_channels=unet.config.block_out_channels, - cross_attention_dim=unet.config.cross_attention_dim, - down_block_types=unet.config.down_block_types, - sample_size=unet.config.sample_size, - transformer_layers_per_block=unet.config.transformer_layers_per_block, - upcast_attention=unet.config.upcast_attention, - max_norm_num_groups=unet.config.norm_num_groups, - use_linear_projection=unet.config.use_linear_projection, - ) - - # ensure that the ControlNetXSAdapter is the same dtype as the UNet2DConditionModel - model.to(unet.dtype) - - return model - - def forward(self, *args, **kwargs): - raise ValueError( - "A ControlNetXSAdapter cannot be run by itself. Use it together with a UNet2DConditionModel to instantiate a UNetControlNetXSModel." - ) - - -class UNetControlNetXSModel(ModelMixin, ConfigMixin): - r""" - A UNet fused with a ControlNet-XS adapter model - - This model inherits from [`ModelMixin`] and [`ConfigMixin`]. Check the superclass documentation for it's generic - methods implemented for all models (such as downloading or saving). - - `UNetControlNetXSModel` is compatible with StableDiffusion and StableDiffusion-XL. It's default parameters are - compatible with StableDiffusion. - - It's parameters are either passed to the underlying `UNet2DConditionModel` or used exactly like in - `ControlNetXSAdapter` . See their documentation for details. - """ - - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - # unet configs - sample_size: Optional[int] = 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, - upcast_attention: bool = True, - use_linear_projection: bool = True, - time_cond_proj_dim: Optional[int] = None, - projection_class_embeddings_input_dim: Optional[int] = 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_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_max_norm_num_groups: int = 32, - ): - super().__init__() - - if time_embedding_mix < 0 or time_embedding_mix > 1: - raise ValueError("`time_embedding_mix` needs to be between 0 and 1.") - if time_embedding_mix < 1 and not ctrl_learn_time_embedding: - raise ValueError("To use `time_embedding_mix` < 1, `ctrl_learn_time_embedding` must be `True`") - - if addition_embed_type is not None and addition_embed_type != "text_time": - raise ValueError( - "As `UNetControlNetXSModel` currently only supports StableDiffusion and StableDiffusion-XL, `addition_embed_type` must be `None` or `'text_time'`." - ) - - if not isinstance(transformer_layers_per_block, (list, tuple)): - transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) - if not isinstance(cross_attention_dim, (list, tuple)): - cross_attention_dim = [cross_attention_dim] * len(down_block_types) - if not isinstance(num_attention_heads, (list, tuple)): - num_attention_heads = [num_attention_heads] * len(down_block_types) - if not isinstance(ctrl_num_attention_heads, (list, tuple)): - ctrl_num_attention_heads = [ctrl_num_attention_heads] * len(down_block_types) - - base_num_attention_heads = num_attention_heads - - self.in_channels = 4 - - # # Input - self.base_conv_in = nn.Conv2d(4, block_out_channels[0], kernel_size=3, padding=1) - self.controlnet_cond_embedding = ControlNetConditioningEmbedding( - conditioning_embedding_channels=ctrl_block_out_channels[0], - block_out_channels=ctrl_conditioning_embedding_out_channels, - conditioning_channels=ctrl_conditioning_channels, - ) - self.ctrl_conv_in = nn.Conv2d(4, ctrl_block_out_channels[0], kernel_size=3, padding=1) - self.control_to_base_for_conv_in = make_zero_conv(ctrl_block_out_channels[0], block_out_channels[0]) - - # # Time - time_embed_input_dim = block_out_channels[0] - time_embed_dim = block_out_channels[0] * 4 - - self.base_time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos=True, downscale_freq_shift=0) - self.base_time_embedding = TimestepEmbedding( - time_embed_input_dim, - time_embed_dim, - cond_proj_dim=time_cond_proj_dim, - ) - if ctrl_learn_time_embedding: - self.ctrl_time_embedding = TimestepEmbedding( - in_channels=time_embed_input_dim, time_embed_dim=time_embed_dim - ) - else: - self.ctrl_time_embedding = None - - if addition_embed_type is None: - self.base_add_time_proj = None - self.base_add_embedding = None - else: - self.base_add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos=True, downscale_freq_shift=0) - self.base_add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) - - # # Create down blocks - down_blocks = [] - base_out_channels = block_out_channels[0] - ctrl_out_channels = ctrl_block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - base_in_channels = base_out_channels - base_out_channels = block_out_channels[i] - ctrl_in_channels = ctrl_out_channels - ctrl_out_channels = ctrl_block_out_channels[i] - has_crossattn = "CrossAttn" in down_block_type - is_final_block = i == len(down_block_types) - 1 - - down_blocks.append( - ControlNetXSCrossAttnDownBlock2D( - base_in_channels=base_in_channels, - base_out_channels=base_out_channels, - ctrl_in_channels=ctrl_in_channels, - ctrl_out_channels=ctrl_out_channels, - temb_channels=time_embed_dim, - norm_num_groups=norm_num_groups, - ctrl_max_norm_num_groups=ctrl_max_norm_num_groups, - has_crossattn=has_crossattn, - transformer_layers_per_block=transformer_layers_per_block[i], - base_num_attention_heads=base_num_attention_heads[i], - ctrl_num_attention_heads=ctrl_num_attention_heads[i], - cross_attention_dim=cross_attention_dim[i], - add_downsample=not is_final_block, - upcast_attention=upcast_attention, - use_linear_projection=use_linear_projection, - ) - ) - - # # Create mid block - self.mid_block = ControlNetXSCrossAttnMidBlock2D( - base_channels=block_out_channels[-1], - ctrl_channels=ctrl_block_out_channels[-1], - temb_channels=time_embed_dim, - norm_num_groups=norm_num_groups, - ctrl_max_norm_num_groups=ctrl_max_norm_num_groups, - transformer_layers_per_block=transformer_layers_per_block[-1], - base_num_attention_heads=base_num_attention_heads[-1], - ctrl_num_attention_heads=ctrl_num_attention_heads[-1], - cross_attention_dim=cross_attention_dim[-1], - upcast_attention=upcast_attention, - use_linear_projection=use_linear_projection, - ) - - # # Create up blocks - up_blocks = [] - rev_transformer_layers_per_block = list(reversed(transformer_layers_per_block)) - rev_num_attention_heads = list(reversed(base_num_attention_heads)) - rev_cross_attention_dim = list(reversed(cross_attention_dim)) - - # The skip connection channels are the output of the conv_in and of all the down subblocks - ctrl_skip_channels = [ctrl_block_out_channels[0]] - for i, out_channels in enumerate(ctrl_block_out_channels): - number_of_subblocks = ( - 3 if i < len(ctrl_block_out_channels) - 1 else 2 - ) # every block has 3 subblocks, except last one, which has 2 as it has no downsampler - ctrl_skip_channels.extend([out_channels] * number_of_subblocks) - - reversed_block_out_channels = list(reversed(block_out_channels)) - - out_channels = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - prev_output_channel = out_channels - out_channels = reversed_block_out_channels[i] - in_channels = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] - ctrl_skip_channels_ = [ctrl_skip_channels.pop() for _ in range(3)] - - has_crossattn = "CrossAttn" in up_block_type - is_final_block = i == len(block_out_channels) - 1 - - up_blocks.append( - ControlNetXSCrossAttnUpBlock2D( - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - ctrl_skip_channels=ctrl_skip_channels_, - temb_channels=time_embed_dim, - resolution_idx=i, - has_crossattn=has_crossattn, - transformer_layers_per_block=rev_transformer_layers_per_block[i], - num_attention_heads=rev_num_attention_heads[i], - cross_attention_dim=rev_cross_attention_dim[i], - add_upsample=not is_final_block, - upcast_attention=upcast_attention, - norm_num_groups=norm_num_groups, - use_linear_projection=use_linear_projection, - ) - ) - - self.down_blocks = nn.ModuleList(down_blocks) - self.up_blocks = nn.ModuleList(up_blocks) - - self.base_conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups) - self.base_conv_act = nn.SiLU() - self.base_conv_out = nn.Conv2d(block_out_channels[0], 4, kernel_size=3, padding=1) - - @classmethod - 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, - ): - r""" - Instantiate a [`UNetControlNetXSModel`] from a [`UNet2DConditionModel`] and an optional [`ControlNetXSAdapter`] - . - - Parameters: - unet (`UNet2DConditionModel`): - The UNet model we want to control. - controlnet (`ControlNetXSAdapter`): - The ControlNet-XS adapter with which the UNet will be fused. If none is given, a new ControlNet-XS - 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`): - 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): - Used to construct the controlnet if none is given. See [`ControlNetXSAdapter.from_unet`] for details. - ctrl_optional_kwargs (`Dict`, *optional*, defaults to `None`): - Passed to the `init` of the new controlnet if no controlnet was given. - """ - if controlnet is None: - controlnet = ControlNetXSAdapter.from_unet( - unet, size_ratio, ctrl_block_out_channels, **ctrl_optional_kwargs - ) - else: - if any( - o is not None for o in (size_ratio, ctrl_block_out_channels, time_embedding_mix, ctrl_optional_kwargs) - ): - raise ValueError( - "When a controlnet is passed, none of these parameters should be passed: size_ratio, ctrl_block_out_channels, time_embedding_mix, ctrl_optional_kwargs." - ) - - # # get params - params_for_unet = [ - "sample_size", - "down_block_types", - "up_block_types", - "block_out_channels", - "norm_num_groups", - "cross_attention_dim", - "transformer_layers_per_block", - "addition_embed_type", - "addition_time_embed_dim", - "upcast_attention", - "use_linear_projection", - "time_cond_proj_dim", - "projection_class_embeddings_input_dim", - ] - params_for_unet = {k: v for k, v in unet.config.items() if k in params_for_unet} - # The naming seems a bit confusing and it is, see https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 for why. - params_for_unet["num_attention_heads"] = unet.config.attention_head_dim - - params_for_controlnet = [ - "conditioning_channels", - "conditioning_embedding_out_channels", - "conditioning_channel_order", - "learn_time_embedding", - "block_out_channels", - "num_attention_heads", - "max_norm_num_groups", - ] - params_for_controlnet = {"ctrl_" + k: v for k, v in controlnet.config.items() if k in params_for_controlnet} - params_for_controlnet["time_embedding_mix"] = controlnet.config.time_embedding_mix - - # # create model - model = cls.from_config({**params_for_unet, **params_for_controlnet}) - - # # load weights - # from unet - modules_from_unet = [ - "time_embedding", - "conv_in", - "conv_norm_out", - "conv_out", - ] - for m in modules_from_unet: - getattr(model, "base_" + m).load_state_dict(getattr(unet, m).state_dict()) - - optional_modules_from_unet = [ - "add_time_proj", - "add_embedding", - ] - for m in optional_modules_from_unet: - if hasattr(unet, m) and getattr(unet, m) is not None: - getattr(model, "base_" + m).load_state_dict(getattr(unet, m).state_dict()) - - # from controlnet - model.controlnet_cond_embedding.load_state_dict(controlnet.controlnet_cond_embedding.state_dict()) - model.ctrl_conv_in.load_state_dict(controlnet.conv_in.state_dict()) - if controlnet.time_embedding is not None: - model.ctrl_time_embedding.load_state_dict(controlnet.time_embedding.state_dict()) - model.control_to_base_for_conv_in.load_state_dict(controlnet.control_to_base_for_conv_in.state_dict()) - - # from both - model.down_blocks = nn.ModuleList( - ControlNetXSCrossAttnDownBlock2D.from_modules(b, c) - for b, c in zip(unet.down_blocks, controlnet.down_blocks) - ) - model.mid_block = ControlNetXSCrossAttnMidBlock2D.from_modules(unet.mid_block, controlnet.mid_block) - model.up_blocks = nn.ModuleList( - ControlNetXSCrossAttnUpBlock2D.from_modules(b, c) - for b, c in zip(unet.up_blocks, controlnet.up_connections) - ) - - # ensure that the UNetControlNetXSModel is the same dtype as the UNet2DConditionModel - model.to(unet.dtype) - - return model - - def freeze_unet_params(self) -> None: - """Freeze the weights of the parts belonging to the base UNet2DConditionModel, and leave everything else unfrozen for fine - tuning.""" - # Freeze everything - for param in self.parameters(): - param.requires_grad = True - - # Unfreeze ControlNetXSAdapter - base_parts = [ - "base_time_proj", - "base_time_embedding", - "base_add_time_proj", - "base_add_embedding", - "base_conv_in", - "base_conv_norm_out", - "base_conv_act", - "base_conv_out", - ] - base_parts = [getattr(self, part) for part in base_parts if getattr(self, part) is not None] - for part in base_parts: - for param in part.parameters(): - param.requires_grad = False - - for d in self.down_blocks: - d.freeze_base_params() - self.mid_block.freeze_base_params() - for u in self.up_blocks: - u.freeze_base_params() - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - self.set_attn_processor(processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.enable_freeu - def enable_freeu(self, s1: float, s2: float, b1: float, b2: float): - r"""Enables the FreeU mechanism from https://huggingface.co/papers/2309.11497. - - The suffixes after the scaling factors represent the stage blocks where they are being applied. - - Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that - are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL. - - Args: - s1 (`float`): - Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to - mitigate the "oversmoothing effect" in the enhanced denoising process. - s2 (`float`): - Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to - mitigate the "oversmoothing effect" in the enhanced denoising process. - b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features. - b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features. - """ - for i, upsample_block in enumerate(self.up_blocks): - setattr(upsample_block, "s1", s1) - setattr(upsample_block, "s2", s2) - setattr(upsample_block, "b1", b1) - setattr(upsample_block, "b2", b2) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.disable_freeu - def disable_freeu(self): - """Disables the FreeU mechanism.""" - freeu_keys = {"s1", "s2", "b1", "b2"} - for i, upsample_block in enumerate(self.up_blocks): - for k in freeu_keys: - if hasattr(upsample_block, k) or getattr(upsample_block, k, None) is not None: - setattr(upsample_block, k, None) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - self.set_attn_processor(FusedAttnProcessor2_0()) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - if self.original_attn_processors is not None: - self.set_attn_processor(self.original_attn_processors) - - def forward( - self, - sample: Tensor, - timestep: Union[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, - return_dict: bool = True, - apply_control: bool = True, - ) -> Union[ControlNetXSOutput, Tuple]: - """ - The [`ControlNetXSModel`] forward method. - - Args: - sample (`Tensor`): - The noisy input tensor. - timestep (`Union[torch.Tensor, float, int]`): - The number of timesteps to denoise an input. - encoder_hidden_states (`torch.Tensor`): - The encoder hidden states. - controlnet_cond (`Tensor`): - The conditional input tensor of shape `(batch_size, sequence_length, hidden_size)`. - conditioning_scale (`float`, defaults to `1.0`): - How much the control model affects the base model outputs. - class_labels (`torch.Tensor`, *optional*, defaults to `None`): - Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. - timestep_cond (`torch.Tensor`, *optional*, defaults to `None`): - Additional conditional embeddings for timestep. If provided, the embeddings will be summed with the - timestep_embedding passed through the `self.time_embedding` layer to obtain the final timestep - embeddings. - attention_mask (`torch.Tensor`, *optional*, defaults to `None`): - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask - is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large - negative values to the attention scores corresponding to "discard" tokens. - cross_attention_kwargs (`dict[str]`, *optional*, defaults to `None`): - A kwargs dictionary that if specified is passed along to the `AttnProcessor`. - added_cond_kwargs (`dict`): - Additional conditions for the Stable Diffusion XL UNet. - return_dict (`bool`, defaults to `True`): - Whether or not to return a [`~models.controlnets.controlnet.ControlNetOutput`] instead of a plain - tuple. - apply_control (`bool`, defaults to `True`): - If `False`, the input is run only through the base model. - - Returns: - [`~models.controlnetxs.ControlNetXSOutput`] **or** `tuple`: - If `return_dict` is `True`, a [`~models.controlnetxs.ControlNetXSOutput`] is returned, otherwise a - tuple is returned where the first element is the sample tensor. - """ - - # check channel order - if self.config.ctrl_conditioning_channel_order == "bgr": - controlnet_cond = torch.flip(controlnet_cond, dims=[1]) - - # prepare attention_mask - if attention_mask is not None: - attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 - attention_mask = attention_mask.unsqueeze(1) - - # 1. time - 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 - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = sample.device.type == "mps" - is_npu = sample.device.type == "npu" - if isinstance(timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) - elif len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps.expand(sample.shape[0]) - - t_emb = self.base_time_proj(timesteps) - - # timesteps does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - t_emb = t_emb.to(dtype=sample.dtype) - - if self.config.ctrl_learn_time_embedding and apply_control: - ctrl_temb = self.ctrl_time_embedding(t_emb, timestep_cond) - base_temb = self.base_time_embedding(t_emb, timestep_cond) - interpolation_param = self.config.time_embedding_mix**0.3 - - temb = ctrl_temb * interpolation_param + base_temb * (1 - interpolation_param) - else: - temb = self.base_time_embedding(t_emb) - - # added time & text embeddings - aug_emb = None - - if self.config.addition_embed_type is None: - pass - elif self.config.addition_embed_type == "text_time": - # SDXL - style - if "text_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" - ) - text_embeds = added_cond_kwargs.get("text_embeds") - if "time_ids" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" - ) - time_ids = added_cond_kwargs.get("time_ids") - time_embeds = self.base_add_time_proj(time_ids.flatten()) - time_embeds = time_embeds.reshape((text_embeds.shape[0], -1)) - add_embeds = torch.concat([text_embeds, time_embeds], dim=-1) - add_embeds = add_embeds.to(temb.dtype) - aug_emb = self.base_add_embedding(add_embeds) - else: - raise ValueError( - f"ControlNet-XS currently only supports StableDiffusion and StableDiffusion-XL, so addition_embed_type = {self.config.addition_embed_type} is currently not supported." - ) - - temb = temb + aug_emb if aug_emb is not None else temb - - # text embeddings - cemb = encoder_hidden_states - - # Preparation - h_ctrl = h_base = sample - hs_base, hs_ctrl = [], [] - - # Cross Control - guided_hint = self.controlnet_cond_embedding(controlnet_cond) - - # 1 - conv in & down - - h_base = self.base_conv_in(h_base) - h_ctrl = self.ctrl_conv_in(h_ctrl) - if guided_hint is not None: - h_ctrl += guided_hint - if apply_control: - h_base = h_base + self.control_to_base_for_conv_in(h_ctrl) * conditioning_scale # add ctrl -> base - - hs_base.append(h_base) - hs_ctrl.append(h_ctrl) - - for down in self.down_blocks: - h_base, h_ctrl, residual_hb, residual_hc = down( - hidden_states_base=h_base, - hidden_states_ctrl=h_ctrl, - temb=temb, - encoder_hidden_states=cemb, - conditioning_scale=conditioning_scale, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - apply_control=apply_control, - ) - hs_base.extend(residual_hb) - hs_ctrl.extend(residual_hc) - - # 2 - mid - h_base, h_ctrl = self.mid_block( - hidden_states_base=h_base, - hidden_states_ctrl=h_ctrl, - temb=temb, - encoder_hidden_states=cemb, - conditioning_scale=conditioning_scale, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - apply_control=apply_control, - ) - - # 3 - up - for up in self.up_blocks: - n_resnets = len(up.resnets) - skips_hb = hs_base[-n_resnets:] - skips_hc = hs_ctrl[-n_resnets:] - hs_base = hs_base[:-n_resnets] - hs_ctrl = hs_ctrl[:-n_resnets] - h_base = up( - hidden_states=h_base, - res_hidden_states_tuple_base=skips_hb, - res_hidden_states_tuple_ctrl=skips_hc, - temb=temb, - encoder_hidden_states=cemb, - conditioning_scale=conditioning_scale, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - apply_control=apply_control, - ) - - # 4 - conv out - h_base = self.base_conv_norm_out(h_base) - h_base = self.base_conv_act(h_base) - h_base = self.base_conv_out(h_base) - - if not return_dict: - return (h_base,) - - return ControlNetXSOutput(sample=h_base) - - -class ControlNetXSCrossAttnDownBlock2D(nn.Module): - def __init__( - self, - base_in_channels: int, - base_out_channels: int, - ctrl_in_channels: int, - ctrl_out_channels: int, - temb_channels: int, - 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, - add_downsample: bool = True, - upcast_attention: Optional[bool] = False, - use_linear_projection: Optional[bool] = True, - ): - super().__init__() - base_resnets = [] - base_attentions = [] - ctrl_resnets = [] - ctrl_attentions = [] - ctrl_to_base = [] - base_to_ctrl = [] - - num_layers = 2 # only support sd + sdxl - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * num_layers - - for i in range(num_layers): - base_in_channels = base_in_channels if i == 0 else base_out_channels - ctrl_in_channels = ctrl_in_channels if i == 0 else ctrl_out_channels - - # Before the resnet/attention application, information is concatted from base to control. - # Concat doesn't require change in number of channels - base_to_ctrl.append(make_zero_conv(base_in_channels, base_in_channels)) - - base_resnets.append( - ResnetBlock2D( - in_channels=base_in_channels, - out_channels=base_out_channels, - temb_channels=temb_channels, - groups=norm_num_groups, - ) - ) - ctrl_resnets.append( - ResnetBlock2D( - in_channels=ctrl_in_channels + base_in_channels, # information from base is concatted to ctrl - out_channels=ctrl_out_channels, - temb_channels=temb_channels, - groups=find_largest_factor( - ctrl_in_channels + base_in_channels, max_factor=ctrl_max_norm_num_groups - ), - groups_out=find_largest_factor(ctrl_out_channels, max_factor=ctrl_max_norm_num_groups), - eps=1e-5, - ) - ) - - if has_crossattn: - base_attentions.append( - Transformer2DModel( - base_num_attention_heads, - base_out_channels // base_num_attention_heads, - in_channels=base_out_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - norm_num_groups=norm_num_groups, - ) - ) - ctrl_attentions.append( - Transformer2DModel( - ctrl_num_attention_heads, - ctrl_out_channels // ctrl_num_attention_heads, - in_channels=ctrl_out_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - norm_num_groups=find_largest_factor(ctrl_out_channels, max_factor=ctrl_max_norm_num_groups), - ) - ) - - # After the resnet/attention application, information is added from control to base - # Addition requires change in number of channels - ctrl_to_base.append(make_zero_conv(ctrl_out_channels, base_out_channels)) - - if add_downsample: - # Before the downsampler application, information is concatted from base to control - # Concat doesn't require change in number of channels - base_to_ctrl.append(make_zero_conv(base_out_channels, base_out_channels)) - - self.base_downsamplers = Downsample2D( - base_out_channels, use_conv=True, out_channels=base_out_channels, name="op" - ) - self.ctrl_downsamplers = Downsample2D( - ctrl_out_channels + base_out_channels, use_conv=True, out_channels=ctrl_out_channels, name="op" - ) - - # After the downsampler application, information is added from control to base - # Addition requires change in number of channels - ctrl_to_base.append(make_zero_conv(ctrl_out_channels, base_out_channels)) - else: - self.base_downsamplers = None - self.ctrl_downsamplers = None - - self.base_resnets = nn.ModuleList(base_resnets) - self.ctrl_resnets = nn.ModuleList(ctrl_resnets) - self.base_attentions = nn.ModuleList(base_attentions) if has_crossattn else [None] * num_layers - self.ctrl_attentions = nn.ModuleList(ctrl_attentions) if has_crossattn else [None] * num_layers - self.base_to_ctrl = nn.ModuleList(base_to_ctrl) - self.ctrl_to_base = nn.ModuleList(ctrl_to_base) - - self.gradient_checkpointing = False - - @classmethod - def from_modules(cls, base_downblock: CrossAttnDownBlock2D, ctrl_downblock: DownBlockControlNetXSAdapter): - # get params - def get_first_cross_attention(block): - return block.attentions[0].transformer_blocks[0].attn2 - - base_in_channels = base_downblock.resnets[0].in_channels - base_out_channels = base_downblock.resnets[0].out_channels - ctrl_in_channels = ( - ctrl_downblock.resnets[0].in_channels - base_in_channels - ) # base channels are concatted to ctrl channels in init - ctrl_out_channels = ctrl_downblock.resnets[0].out_channels - temb_channels = base_downblock.resnets[0].time_emb_proj.in_features - num_groups = base_downblock.resnets[0].norm1.num_groups - ctrl_num_groups = ctrl_downblock.resnets[0].norm1.num_groups - if hasattr(base_downblock, "attentions"): - has_crossattn = True - transformer_layers_per_block = len(base_downblock.attentions[0].transformer_blocks) - base_num_attention_heads = get_first_cross_attention(base_downblock).heads - ctrl_num_attention_heads = get_first_cross_attention(ctrl_downblock).heads - cross_attention_dim = get_first_cross_attention(base_downblock).cross_attention_dim - upcast_attention = get_first_cross_attention(base_downblock).upcast_attention - use_linear_projection = base_downblock.attentions[0].use_linear_projection - else: - has_crossattn = False - transformer_layers_per_block = None - base_num_attention_heads = None - ctrl_num_attention_heads = None - cross_attention_dim = None - upcast_attention = None - use_linear_projection = None - add_downsample = base_downblock.downsamplers is not None - - # create model - model = cls( - base_in_channels=base_in_channels, - base_out_channels=base_out_channels, - ctrl_in_channels=ctrl_in_channels, - ctrl_out_channels=ctrl_out_channels, - temb_channels=temb_channels, - norm_num_groups=num_groups, - ctrl_max_norm_num_groups=ctrl_num_groups, - has_crossattn=has_crossattn, - transformer_layers_per_block=transformer_layers_per_block, - base_num_attention_heads=base_num_attention_heads, - ctrl_num_attention_heads=ctrl_num_attention_heads, - cross_attention_dim=cross_attention_dim, - add_downsample=add_downsample, - upcast_attention=upcast_attention, - use_linear_projection=use_linear_projection, - ) - - # # load weights - model.base_resnets.load_state_dict(base_downblock.resnets.state_dict()) - model.ctrl_resnets.load_state_dict(ctrl_downblock.resnets.state_dict()) - if has_crossattn: - model.base_attentions.load_state_dict(base_downblock.attentions.state_dict()) - model.ctrl_attentions.load_state_dict(ctrl_downblock.attentions.state_dict()) - if add_downsample: - model.base_downsamplers.load_state_dict(base_downblock.downsamplers[0].state_dict()) - model.ctrl_downsamplers.load_state_dict(ctrl_downblock.downsamplers.state_dict()) - model.base_to_ctrl.load_state_dict(ctrl_downblock.base_to_ctrl.state_dict()) - model.ctrl_to_base.load_state_dict(ctrl_downblock.ctrl_to_base.state_dict()) - - return model - - def freeze_base_params(self) -> None: - """Freeze the weights of the parts belonging to the base UNet2DConditionModel, and leave everything else unfrozen for fine - tuning.""" - # Unfreeze everything - for param in self.parameters(): - param.requires_grad = True - - # Freeze base part - base_parts = [self.base_resnets] - if isinstance(self.base_attentions, nn.ModuleList): # attentions can be a list of Nones - base_parts.append(self.base_attentions) - if self.base_downsamplers is not None: - base_parts.append(self.base_downsamplers) - for part in base_parts: - for param in part.parameters(): - param.requires_grad = False - - 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, - apply_control: bool = True, - ) -> 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.") - - h_base = hidden_states_base - h_ctrl = hidden_states_ctrl - - base_output_states = () - ctrl_output_states = () - - base_blocks = list(zip(self.base_resnets, self.base_attentions)) - ctrl_blocks = list(zip(self.ctrl_resnets, self.ctrl_attentions)) - - for (b_res, b_attn), (c_res, c_attn), b2c, c2b in zip( - base_blocks, ctrl_blocks, self.base_to_ctrl, self.ctrl_to_base - ): - # concat base -> ctrl - if apply_control: - h_ctrl = torch.cat([h_ctrl, b2c(h_base)], dim=1) - - # apply base subblock - if torch.is_grad_enabled() and self.gradient_checkpointing: - h_base = self._gradient_checkpointing_func(b_res, h_base, temb) - else: - h_base = b_res(h_base, temb) - - if b_attn is not None: - h_base = b_attn( - h_base, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - - # apply ctrl subblock - if apply_control: - if torch.is_grad_enabled() and self.gradient_checkpointing: - h_ctrl = self._gradient_checkpointing_func(c_res, h_ctrl, temb) - else: - h_ctrl = c_res(h_ctrl, temb) - if c_attn is not None: - h_ctrl = c_attn( - h_ctrl, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - - # add ctrl -> base - if apply_control: - h_base = h_base + c2b(h_ctrl) * conditioning_scale - - base_output_states = base_output_states + (h_base,) - ctrl_output_states = ctrl_output_states + (h_ctrl,) - - if self.base_downsamplers is not None: # if we have a base_downsampler, then also a ctrl_downsampler - b2c = self.base_to_ctrl[-1] - c2b = self.ctrl_to_base[-1] - - # concat base -> ctrl - if apply_control: - h_ctrl = torch.cat([h_ctrl, b2c(h_base)], dim=1) - # apply base subblock - h_base = self.base_downsamplers(h_base) - # apply ctrl subblock - if apply_control: - h_ctrl = self.ctrl_downsamplers(h_ctrl) - # add ctrl -> base - if apply_control: - h_base = h_base + c2b(h_ctrl) * conditioning_scale - - base_output_states = base_output_states + (h_base,) - ctrl_output_states = ctrl_output_states + (h_ctrl,) - - return h_base, h_ctrl, base_output_states, ctrl_output_states - - -class ControlNetXSCrossAttnMidBlock2D(nn.Module): - def __init__( - self, - base_channels: int, - ctrl_channels: int, - temb_channels: Optional[int] = 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, - upcast_attention: bool = False, - use_linear_projection: Optional[bool] = True, - ): - super().__init__() - - # Before the midblock application, information is concatted from base to control. - # Concat doesn't require change in number of channels - self.base_to_ctrl = make_zero_conv(base_channels, base_channels) - - self.base_midblock = UNetMidBlock2DCrossAttn( - transformer_layers_per_block=transformer_layers_per_block, - in_channels=base_channels, - temb_channels=temb_channels, - resnet_groups=norm_num_groups, - cross_attention_dim=cross_attention_dim, - num_attention_heads=base_num_attention_heads, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - ) - - self.ctrl_midblock = UNetMidBlock2DCrossAttn( - transformer_layers_per_block=transformer_layers_per_block, - in_channels=ctrl_channels + base_channels, - out_channels=ctrl_channels, - temb_channels=temb_channels, - # number or norm groups must divide both in_channels and out_channels - resnet_groups=find_largest_factor( - gcd(ctrl_channels, ctrl_channels + base_channels), ctrl_max_norm_num_groups - ), - cross_attention_dim=cross_attention_dim, - num_attention_heads=ctrl_num_attention_heads, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - ) - - # After the midblock application, information is added from control to base - # Addition requires change in number of channels - self.ctrl_to_base = make_zero_conv(ctrl_channels, base_channels) - - self.gradient_checkpointing = False - - @classmethod - def from_modules( - cls, - base_midblock: UNetMidBlock2DCrossAttn, - ctrl_midblock: MidBlockControlNetXSAdapter, - ): - base_to_ctrl = ctrl_midblock.base_to_ctrl - ctrl_to_base = ctrl_midblock.ctrl_to_base - ctrl_midblock = ctrl_midblock.midblock - - # get params - def get_first_cross_attention(midblock): - return midblock.attentions[0].transformer_blocks[0].attn2 - - base_channels = ctrl_to_base.out_channels - ctrl_channels = ctrl_to_base.in_channels - transformer_layers_per_block = len(base_midblock.attentions[0].transformer_blocks) - temb_channels = base_midblock.resnets[0].time_emb_proj.in_features - num_groups = base_midblock.resnets[0].norm1.num_groups - ctrl_num_groups = ctrl_midblock.resnets[0].norm1.num_groups - base_num_attention_heads = get_first_cross_attention(base_midblock).heads - ctrl_num_attention_heads = get_first_cross_attention(ctrl_midblock).heads - cross_attention_dim = get_first_cross_attention(base_midblock).cross_attention_dim - upcast_attention = get_first_cross_attention(base_midblock).upcast_attention - use_linear_projection = base_midblock.attentions[0].use_linear_projection - - # create model - model = cls( - base_channels=base_channels, - ctrl_channels=ctrl_channels, - temb_channels=temb_channels, - norm_num_groups=num_groups, - ctrl_max_norm_num_groups=ctrl_num_groups, - transformer_layers_per_block=transformer_layers_per_block, - base_num_attention_heads=base_num_attention_heads, - ctrl_num_attention_heads=ctrl_num_attention_heads, - cross_attention_dim=cross_attention_dim, - upcast_attention=upcast_attention, - use_linear_projection=use_linear_projection, - ) - - # load weights - model.base_to_ctrl.load_state_dict(base_to_ctrl.state_dict()) - model.base_midblock.load_state_dict(base_midblock.state_dict()) - model.ctrl_midblock.load_state_dict(ctrl_midblock.state_dict()) - model.ctrl_to_base.load_state_dict(ctrl_to_base.state_dict()) - - return model - - def freeze_base_params(self) -> None: - """Freeze the weights of the parts belonging to the base UNet2DConditionModel, and leave everything else unfrozen for fine - tuning.""" - # Unfreeze everything - for param in self.parameters(): - param.requires_grad = True - - # Freeze base part - for param in self.base_midblock.parameters(): - param.requires_grad = False - - def forward( - self, - 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, - apply_control: bool = True, - ) -> 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.") - - h_base = hidden_states_base - h_ctrl = hidden_states_ctrl - - joint_args = { - "temb": temb, - "encoder_hidden_states": encoder_hidden_states, - "attention_mask": attention_mask, - "cross_attention_kwargs": cross_attention_kwargs, - "encoder_attention_mask": encoder_attention_mask, - } - - if apply_control: - h_ctrl = torch.cat([h_ctrl, self.base_to_ctrl(h_base)], dim=1) # concat base -> ctrl - h_base = self.base_midblock(h_base, **joint_args) # apply base mid block - if apply_control: - h_ctrl = self.ctrl_midblock(h_ctrl, **joint_args) # apply ctrl mid block - h_base = h_base + self.ctrl_to_base(h_ctrl) * conditioning_scale # add ctrl -> base - - return h_base, h_ctrl - - -class ControlNetXSCrossAttnUpBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - prev_output_channel: int, - ctrl_skip_channels: List[int], - temb_channels: int, - norm_num_groups: int = 32, - resolution_idx: Optional[int] = 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, - ): - super().__init__() - resnets = [] - attentions = [] - ctrl_to_base = [] - - num_layers = 3 # only support sd + sdxl - - self.has_cross_attention = has_crossattn - self.num_attention_heads = num_attention_heads - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * num_layers - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - ctrl_to_base.append(make_zero_conv(ctrl_skip_channels[i], resnet_in_channels)) - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - groups=norm_num_groups, - ) - ) - - if has_crossattn: - attentions.append( - Transformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - norm_num_groups=norm_num_groups, - ) - ) - - self.resnets = nn.ModuleList(resnets) - self.attentions = nn.ModuleList(attentions) if has_crossattn else [None] * num_layers - self.ctrl_to_base = nn.ModuleList(ctrl_to_base) - - if add_upsample: - self.upsamplers = Upsample2D(out_channels, use_conv=True, out_channels=out_channels) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - @classmethod - def from_modules(cls, base_upblock: CrossAttnUpBlock2D, ctrl_upblock: UpBlockControlNetXSAdapter): - ctrl_to_base_skip_connections = ctrl_upblock.ctrl_to_base - - # get params - def get_first_cross_attention(block): - return block.attentions[0].transformer_blocks[0].attn2 - - out_channels = base_upblock.resnets[0].out_channels - in_channels = base_upblock.resnets[-1].in_channels - out_channels - prev_output_channels = base_upblock.resnets[0].in_channels - out_channels - ctrl_skip_channelss = [c.in_channels for c in ctrl_to_base_skip_connections] - temb_channels = base_upblock.resnets[0].time_emb_proj.in_features - num_groups = base_upblock.resnets[0].norm1.num_groups - resolution_idx = base_upblock.resolution_idx - if hasattr(base_upblock, "attentions"): - has_crossattn = True - transformer_layers_per_block = len(base_upblock.attentions[0].transformer_blocks) - num_attention_heads = get_first_cross_attention(base_upblock).heads - cross_attention_dim = get_first_cross_attention(base_upblock).cross_attention_dim - upcast_attention = get_first_cross_attention(base_upblock).upcast_attention - use_linear_projection = base_upblock.attentions[0].use_linear_projection - else: - has_crossattn = False - transformer_layers_per_block = None - num_attention_heads = None - cross_attention_dim = None - upcast_attention = None - use_linear_projection = None - add_upsample = base_upblock.upsamplers is not None - - # create model - model = cls( - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channels, - ctrl_skip_channels=ctrl_skip_channelss, - temb_channels=temb_channels, - norm_num_groups=num_groups, - resolution_idx=resolution_idx, - has_crossattn=has_crossattn, - transformer_layers_per_block=transformer_layers_per_block, - num_attention_heads=num_attention_heads, - cross_attention_dim=cross_attention_dim, - add_upsample=add_upsample, - upcast_attention=upcast_attention, - use_linear_projection=use_linear_projection, - ) - - # load weights - model.resnets.load_state_dict(base_upblock.resnets.state_dict()) - if has_crossattn: - model.attentions.load_state_dict(base_upblock.attentions.state_dict()) - if add_upsample: - model.upsamplers.load_state_dict(base_upblock.upsamplers[0].state_dict()) - model.ctrl_to_base.load_state_dict(ctrl_to_base_skip_connections.state_dict()) - - return model - - def freeze_base_params(self) -> None: - """Freeze the weights of the parts belonging to the base UNet2DConditionModel, and leave everything else unfrozen for fine - tuning.""" - # Unfreeze everything - for param in self.parameters(): - param.requires_grad = True - - # Freeze base part - base_parts = [self.resnets] - if isinstance(self.attentions, nn.ModuleList): # attentions can be a list of Nones - base_parts.append(self.attentions) - if self.upsamplers is not None: - base_parts.append(self.upsamplers) - for part in base_parts: - for param in part.parameters(): - param.requires_grad = False - - def forward( - self, - hidden_states: 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, - apply_control: bool = True, - ) -> 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.") - - is_freeu_enabled = ( - getattr(self, "s1", None) - and getattr(self, "s2", None) - and getattr(self, "b1", None) - and getattr(self, "b2", None) - ) - - def maybe_apply_freeu_to_subblock(hidden_states, res_h_base): - # FreeU: Only operate on the first two stages - if is_freeu_enabled: - return apply_freeu( - self.resolution_idx, - hidden_states, - res_h_base, - s1=self.s1, - s2=self.s2, - b1=self.b1, - b2=self.b2, - ) - else: - return hidden_states, res_h_base - - for resnet, attn, c2b, res_h_base, res_h_ctrl in zip( - self.resnets, - self.attentions, - self.ctrl_to_base, - reversed(res_hidden_states_tuple_base), - reversed(res_hidden_states_tuple_ctrl), - ): - if apply_control: - hidden_states += c2b(res_h_ctrl) * conditioning_scale - - hidden_states, res_h_base = maybe_apply_freeu_to_subblock(hidden_states, res_h_base) - hidden_states = torch.cat([hidden_states, res_h_base], dim=1) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = resnet(hidden_states, temb) - - if attn is not None: - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - - if self.upsamplers is not None: - hidden_states = self.upsamplers(hidden_states, upsample_size) - - return hidden_states - - -def make_zero_conv(in_channels, out_channels=None): - return zero_module(nn.Conv2d(in_channels, out_channels, 1, padding=0)) - - -def zero_module(module): - for p in module.parameters(): - nn.init.zeros_(p) - return module - - -def find_largest_factor(number, max_factor): - factor = max_factor - if factor >= number: - return number - while factor != 0: - residual = number % factor - if residual == 0: - return factor - factor -= 1 diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/multicontrolnet.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/multicontrolnet.py deleted file mode 100644 index 87a952294997d3d8e2c67a91fb7774b690889355..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/multicontrolnet.py +++ /dev/null @@ -1,182 +0,0 @@ -import os -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from torch import nn - -from ...utils import logging -from ..controlnets.controlnet import ControlNetModel, ControlNetOutput -from ..modeling_utils import ModelMixin - - -logger = logging.get_logger(__name__) - - -class MultiControlNetModel(ModelMixin): - r""" - Multiple `ControlNetModel` wrapper class for Multi-ControlNet - - This module is a wrapper for multiple instances of the `ControlNetModel`. The `forward()` API is designed to be - compatible with `ControlNetModel`. - - Args: - 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]]): - super().__init__() - self.nets = nn.ModuleList(controlnets) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[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, - guess_mode: bool = False, - return_dict: bool = True, - ) -> Union[ControlNetOutput, Tuple]: - for i, (image, scale, controlnet) in enumerate(zip(controlnet_cond, conditioning_scale, self.nets)): - down_samples, mid_sample = controlnet( - sample=sample, - timestep=timestep, - encoder_hidden_states=encoder_hidden_states, - controlnet_cond=image, - conditioning_scale=scale, - class_labels=class_labels, - timestep_cond=timestep_cond, - attention_mask=attention_mask, - added_cond_kwargs=added_cond_kwargs, - cross_attention_kwargs=cross_attention_kwargs, - guess_mode=guess_mode, - return_dict=return_dict, - ) - - # merge samples - if i == 0: - down_block_res_samples, mid_block_res_sample = down_samples, mid_sample - else: - down_block_res_samples = [ - samples_prev + samples_curr - for samples_prev, samples_curr in zip(down_block_res_samples, down_samples) - ] - mid_block_res_sample += mid_sample - - return down_block_res_samples, mid_block_res_sample - - def save_pretrained( - self, - save_directory: Union[str, os.PathLike], - is_main_process: bool = True, - save_function: Callable = None, - safe_serialization: bool = True, - variant: Optional[str] = None, - ): - """ - Save a model and its configuration file to a directory, so that it can be re-loaded using the - `[`~models.controlnets.multicontrolnet.MultiControlNetModel.from_pretrained`]` class method. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to which to save. Will be created if it doesn't exist. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful when in distributed training like - TPUs and need to call this function on all processes. In this case, set `is_main_process=True` only on - the main process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful on distributed training like TPUs when one - need to replace `torch.save` by another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). - variant (`str`, *optional*): - If specified, weights are saved in the format pytorch_model..bin. - """ - for idx, controlnet in enumerate(self.nets): - suffix = "" if idx == 0 else f"_{idx}" - controlnet.save_pretrained( - save_directory + suffix, - is_main_process=is_main_process, - save_function=save_function, - safe_serialization=safe_serialization, - variant=variant, - ) - - @classmethod - def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os.PathLike]], **kwargs): - r""" - Instantiate a pretrained MultiControlNet model from multiple pre-trained controlnet models. - - The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train - the model, you should first set it back in training mode with `model.train()`. - - The warning *Weights from XXX not initialized from pretrained model* means that the weights of XXX do not come - pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning - task. - - The warning *Weights from XXX not used in YYY* means that the layer XXX is not used by YYY, therefore those - weights are discarded. - - Parameters: - pretrained_model_path (`os.PathLike`): - A path to a *directory* containing model weights saved using - [`~models.controlnets.multicontrolnet.MultiControlNetModel.save_pretrained`], e.g., - `./my_model_directory/controlnet`. - torch_dtype (`torch.dtype`, *optional*): - 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*): - 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. - - To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For - more information about each option see [designing a device - map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). - max_memory (`Dict`, *optional*): - A dictionary device identifier to maximum memory. Will default to the maximum memory available for each - GPU and the available CPU RAM if unset. - low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): - Speed up model loading by not initializing the weights and only loading the pre-trained weights. This - also tries to not use more than 1x model size in CPU memory (including peak memory) while loading the - model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch, - setting this argument to `True` will raise an error. - variant (`str`, *optional*): - If specified load weights from `variant` filename, *e.g.* pytorch_model..bin. `variant` is - ignored when using `from_flax`. - use_safetensors (`bool`, *optional*, defaults to `None`): - If set to `None`, the `safetensors` weights will be downloaded if they're available **and** if the - `safetensors` library is installed. If set to `True`, the model will be forcibly loaded from - `safetensors` weights. If set to `False`, loading will *not* use `safetensors`. - """ - idx = 0 - controlnets = [] - - # load controlnet and append to list until no controlnet directory exists anymore - # first controlnet has to be saved under `./mydirectory/controlnet` to be compliant with `DiffusionPipeline.from_prertained` - # second, third, ... controlnets have to be saved under `./mydirectory/controlnet_1`, `./mydirectory/controlnet_2`, ... - model_path_to_load = pretrained_model_path - while os.path.isdir(model_path_to_load): - controlnet = ControlNetModel.from_pretrained(model_path_to_load, **kwargs) - controlnets.append(controlnet) - - idx += 1 - model_path_to_load = pretrained_model_path + f"_{idx}" - - logger.info(f"{len(controlnets)} controlnets loaded from {pretrained_model_path}.") - - if len(controlnets) == 0: - raise ValueError( - f"No ControlNets found under {os.path.dirname(pretrained_model_path)}. Expected at least {pretrained_model_path + '_0'}." - ) - - return cls(controlnets) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/multicontrolnet_union.py b/venv/lib/python3.10/site-packages/diffusers/models/controlnets/multicontrolnet_union.py deleted file mode 100644 index d5506dc186e37d9b5d89c231488b2286d422e17e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/controlnets/multicontrolnet_union.py +++ /dev/null @@ -1,195 +0,0 @@ -import os -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from torch import nn - -from ...utils import logging -from ..controlnets.controlnet import ControlNetOutput -from ..controlnets.controlnet_union import ControlNetUnionModel -from ..modeling_utils import ModelMixin - - -logger = logging.get_logger(__name__) - - -class MultiControlNetUnionModel(ModelMixin): - r""" - Multiple `ControlNetUnionModel` wrapper class for Multi-ControlNet-Union. - - This module is a wrapper for multiple instances of the `ControlNetUnionModel`. The `forward()` API is designed to - be compatible with `ControlNetUnionModel`. - - Args: - 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]]): - super().__init__() - self.nets = nn.ModuleList(controlnets) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[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, - guess_mode: bool = False, - return_dict: bool = True, - ) -> Union[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) - ): - if scale == 0.0: - continue - down_samples, mid_sample = controlnet( - sample=sample, - timestep=timestep, - encoder_hidden_states=encoder_hidden_states, - controlnet_cond=image, - control_type=ctype, - control_type_idx=ctype_idx, - conditioning_scale=scale, - class_labels=class_labels, - timestep_cond=timestep_cond, - attention_mask=attention_mask, - added_cond_kwargs=added_cond_kwargs, - cross_attention_kwargs=cross_attention_kwargs, - from_multi=True, - guess_mode=guess_mode, - return_dict=return_dict, - ) - - # merge samples - if down_block_res_samples is None and mid_block_res_sample is None: - down_block_res_samples, mid_block_res_sample = down_samples, mid_sample - else: - down_block_res_samples = [ - samples_prev + samples_curr - for samples_prev, samples_curr in zip(down_block_res_samples, down_samples) - ] - mid_block_res_sample += mid_sample - - return down_block_res_samples, mid_block_res_sample - - # Copied from diffusers.models.controlnets.multicontrolnet.MultiControlNetModel.save_pretrained with ControlNet->ControlNetUnion - def save_pretrained( - self, - save_directory: Union[str, os.PathLike], - is_main_process: bool = True, - save_function: Callable = None, - safe_serialization: bool = True, - variant: Optional[str] = None, - ): - """ - Save a model and its configuration file to a directory, so that it can be re-loaded using the - `[`~models.controlnets.multicontrolnet.MultiControlNetUnionModel.from_pretrained`]` class method. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to which to save. Will be created if it doesn't exist. - is_main_process (`bool`, *optional*, defaults to `True`): - Whether the process calling this is the main process or not. Useful when in distributed training like - TPUs and need to call this function on all processes. In this case, set `is_main_process=True` only on - the main process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful on distributed training like TPUs when one - need to replace `torch.save` by another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). - variant (`str`, *optional*): - If specified, weights are saved in the format pytorch_model..bin. - """ - for idx, controlnet in enumerate(self.nets): - suffix = "" if idx == 0 else f"_{idx}" - controlnet.save_pretrained( - save_directory + suffix, - is_main_process=is_main_process, - save_function=save_function, - safe_serialization=safe_serialization, - variant=variant, - ) - - @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): - r""" - Instantiate a pretrained MultiControlNetUnion model from multiple pre-trained controlnet models. - - The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train - the model, you should first set it back in training mode with `model.train()`. - - The warning *Weights from XXX not initialized from pretrained model* means that the weights of XXX do not come - pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning - task. - - The warning *Weights from XXX not used in YYY* means that the layer XXX is not used by YYY, therefore those - weights are discarded. - - Parameters: - pretrained_model_path (`os.PathLike`): - A path to a *directory* containing model weights saved using - [`~models.controlnets.multicontrolnet.MultiControlNetUnionModel.save_pretrained`], e.g., - `./my_model_directory/controlnet`. - torch_dtype (`torch.dtype`, *optional*): - 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*): - 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. - - To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For - more information about each option see [designing a device - map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). - max_memory (`Dict`, *optional*): - A dictionary device identifier to maximum memory. Will default to the maximum memory available for each - GPU and the available CPU RAM if unset. - low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): - Speed up model loading by not initializing the weights and only loading the pre-trained weights. This - also tries to not use more than 1x model size in CPU memory (including peak memory) while loading the - model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch, - setting this argument to `True` will raise an error. - variant (`str`, *optional*): - If specified load weights from `variant` filename, *e.g.* pytorch_model..bin. `variant` is - ignored when using `from_flax`. - use_safetensors (`bool`, *optional*, defaults to `None`): - If set to `None`, the `safetensors` weights will be downloaded if they're available **and** if the - `safetensors` library is installed. If set to `True`, the model will be forcibly loaded from - `safetensors` weights. If set to `False`, loading will *not* use `safetensors`. - """ - idx = 0 - controlnets = [] - - # load controlnet and append to list until no controlnet directory exists anymore - # first controlnet has to be saved under `./mydirectory/controlnet` to be compliant with `DiffusionPipeline.from_prertained` - # second, third, ... controlnets have to be saved under `./mydirectory/controlnet_1`, `./mydirectory/controlnet_2`, ... - model_path_to_load = pretrained_model_path - while os.path.isdir(model_path_to_load): - controlnet = ControlNetUnionModel.from_pretrained(model_path_to_load, **kwargs) - controlnets.append(controlnet) - - idx += 1 - model_path_to_load = pretrained_model_path + f"_{idx}" - - logger.info(f"{len(controlnets)} controlnets loaded from {pretrained_model_path}.") - - if len(controlnets) == 0: - raise ValueError( - f"No ControlNetUnions found under {os.path.dirname(pretrained_model_path)}. Expected at least {pretrained_model_path + '_0'}." - ) - - return cls(controlnets) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/downsampling.py b/venv/lib/python3.10/site-packages/diffusers/models/downsampling.py deleted file mode 100644 index 505816422b2a54cfb605b86741ec1c725b5a4711..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/downsampling.py +++ /dev/null @@ -1,401 +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 - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from ..utils import deprecate -from .normalization import RMSNorm -from .upsampling import upfirdn2d_native - - -class Downsample1D(nn.Module): - """A 1D downsampling layer with an optional convolution. - - Parameters: - channels (`int`): - number of channels in the inputs and outputs. - use_conv (`bool`, default `False`): - option to use a convolution. - out_channels (`int`, optional): - number of output channels. Defaults to `channels`. - padding (`int`, default `1`): - padding for the convolution. - name (`str`, default `conv`): - name of the downsampling 1D layer. - """ - - def __init__( - self, - channels: int, - use_conv: bool = False, - out_channels: Optional[int] = None, - padding: int = 1, - name: str = "conv", - ): - super().__init__() - self.channels = channels - self.out_channels = out_channels or channels - self.use_conv = use_conv - self.padding = padding - stride = 2 - self.name = name - - if use_conv: - self.conv = nn.Conv1d(self.channels, self.out_channels, 3, stride=stride, padding=padding) - else: - assert self.channels == self.out_channels - self.conv = nn.AvgPool1d(kernel_size=stride, stride=stride) - - def forward(self, inputs: torch.Tensor) -> torch.Tensor: - assert inputs.shape[1] == self.channels - return self.conv(inputs) - - -class Downsample2D(nn.Module): - """A 2D downsampling layer with an optional convolution. - - Parameters: - channels (`int`): - number of channels in the inputs and outputs. - use_conv (`bool`, default `False`): - option to use a convolution. - out_channels (`int`, optional): - number of output channels. Defaults to `channels`. - padding (`int`, default `1`): - padding for the convolution. - name (`str`, default `conv`): - name of the downsampling 2D layer. - """ - - def __init__( - self, - channels: int, - use_conv: bool = False, - out_channels: Optional[int] = None, - padding: int = 1, - name: str = "conv", - kernel_size=3, - norm_type=None, - eps=None, - elementwise_affine=None, - bias=True, - ): - super().__init__() - self.channels = channels - self.out_channels = out_channels or channels - self.use_conv = use_conv - self.padding = padding - stride = 2 - self.name = name - - if norm_type == "ln_norm": - self.norm = nn.LayerNorm(channels, eps, elementwise_affine) - elif norm_type == "rms_norm": - self.norm = RMSNorm(channels, eps, elementwise_affine) - elif norm_type is None: - self.norm = None - else: - raise ValueError(f"unknown norm_type: {norm_type}") - - if use_conv: - conv = nn.Conv2d( - self.channels, self.out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias - ) - else: - assert self.channels == self.out_channels - conv = nn.AvgPool2d(kernel_size=stride, stride=stride) - - # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed - if name == "conv": - self.Conv2d_0 = conv - self.conv = conv - elif name == "Conv2d_0": - self.conv = conv - else: - self.conv = conv - - def forward(self, hidden_states: torch.Tensor, *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) - assert hidden_states.shape[1] == self.channels - - if self.norm is not None: - hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) - - if self.use_conv and self.padding == 0: - pad = (0, 1, 0, 1) - hidden_states = F.pad(hidden_states, pad, mode="constant", value=0) - - assert hidden_states.shape[1] == self.channels - - hidden_states = self.conv(hidden_states) - - return hidden_states - - -class FirDownsample2D(nn.Module): - """A 2D FIR downsampling layer with an optional convolution. - - Parameters: - channels (`int`): - number of channels in the inputs and outputs. - use_conv (`bool`, default `False`): - option to use a convolution. - out_channels (`int`, optional): - number of output channels. Defaults to `channels`. - fir_kernel (`tuple`, default `(1, 3, 3, 1)`): - kernel for the FIR filter. - """ - - def __init__( - self, - channels: Optional[int] = None, - out_channels: Optional[int] = None, - use_conv: bool = False, - fir_kernel: Tuple[int, int, int, int] = (1, 3, 3, 1), - ): - super().__init__() - out_channels = out_channels if out_channels else channels - if use_conv: - self.Conv2d_0 = nn.Conv2d(channels, out_channels, kernel_size=3, stride=1, padding=1) - self.fir_kernel = fir_kernel - self.use_conv = use_conv - self.out_channels = out_channels - - def _downsample_2d( - self, - hidden_states: torch.Tensor, - weight: Optional[torch.Tensor] = None, - kernel: Optional[torch.Tensor] = None, - factor: int = 2, - gain: float = 1, - ) -> torch.Tensor: - """Fused `Conv2d()` followed by `downsample_2d()`. - Padding is performed only once at the beginning, not between the operations. The fused op is considerably more - efficient than performing the same calculation using standard TensorFlow ops. It supports gradients of - arbitrary order. - - Args: - hidden_states (`torch.Tensor`): - Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`. - weight (`torch.Tensor`, *optional*): - Weight tensor of the shape `[filterH, filterW, inChannels, outChannels]`. Grouped convolution can be - performed by `inChannels = x.shape[0] // numGroups`. - kernel (`torch.Tensor`, *optional*): - FIR filter of the shape `[firH, firW]` or `[firN]` (separable). The default is `[1] * factor`, which - corresponds to average pooling. - factor (`int`, *optional*, default to `2`): - Integer downsampling factor. - gain (`float`, *optional*, default to `1.0`): - Scaling factor for signal magnitude. - - Returns: - output (`torch.Tensor`): - Tensor of the shape `[N, C, H // factor, W // factor]` or `[N, H // factor, W // factor, C]`, and same - datatype as `x`. - """ - - assert isinstance(factor, int) and factor >= 1 - if kernel is None: - kernel = [1] * factor - - # setup kernel - kernel = torch.tensor(kernel, dtype=torch.float32) - if kernel.ndim == 1: - kernel = torch.outer(kernel, kernel) - kernel /= torch.sum(kernel) - - kernel = kernel * gain - - if self.use_conv: - _, _, convH, convW = weight.shape - pad_value = (kernel.shape[0] - factor) + (convW - 1) - stride_value = [factor, factor] - upfirdn_input = upfirdn2d_native( - hidden_states, - torch.tensor(kernel, device=hidden_states.device), - pad=((pad_value + 1) // 2, pad_value // 2), - ) - output = F.conv2d(upfirdn_input, weight, stride=stride_value, padding=0) - else: - pad_value = kernel.shape[0] - factor - output = upfirdn2d_native( - hidden_states, - torch.tensor(kernel, device=hidden_states.device), - down=factor, - pad=((pad_value + 1) // 2, pad_value // 2), - ) - - return output - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.use_conv: - downsample_input = self._downsample_2d(hidden_states, weight=self.Conv2d_0.weight, kernel=self.fir_kernel) - hidden_states = downsample_input + self.Conv2d_0.bias.reshape(1, -1, 1, 1) - else: - hidden_states = self._downsample_2d(hidden_states, kernel=self.fir_kernel, factor=2) - - return hidden_states - - -# downsample/upsample layer used in k-upscaler, might be able to use FirDownsample2D/DirUpsample2D instead -class KDownsample2D(nn.Module): - r"""A 2D K-downsampling layer. - - Parameters: - pad_mode (`str`, *optional*, default to `"reflect"`): the padding mode to use. - """ - - def __init__(self, pad_mode: str = "reflect"): - super().__init__() - self.pad_mode = pad_mode - kernel_1d = torch.tensor([[1 / 8, 3 / 8, 3 / 8, 1 / 8]]) - self.pad = kernel_1d.shape[1] // 2 - 1 - self.register_buffer("kernel", kernel_1d.T @ kernel_1d, persistent=False) - - def forward(self, inputs: torch.Tensor) -> torch.Tensor: - inputs = F.pad(inputs, (self.pad,) * 4, self.pad_mode) - weight = inputs.new_zeros( - [ - inputs.shape[1], - inputs.shape[1], - self.kernel.shape[0], - self.kernel.shape[1], - ] - ) - indices = torch.arange(inputs.shape[1], device=inputs.device) - kernel = self.kernel.to(weight)[None, :].expand(inputs.shape[1], -1, -1) - weight[indices, indices] = kernel - return F.conv2d(inputs, weight, stride=2) - - -class CogVideoXDownsample3D(nn.Module): - # Todo: Wait for paper release. - r""" - A 3D Downsampling layer using in [CogVideoX]() by Tsinghua University & ZhipuAI - - Args: - in_channels (`int`): - Number of channels in the input image. - out_channels (`int`): - Number of channels produced by the convolution. - kernel_size (`int`, defaults to `3`): - Size of the convolving kernel. - stride (`int`, defaults to `2`): - Stride of the convolution. - padding (`int`, defaults to `0`): - Padding added to all four sides of the input. - compress_time (`bool`, defaults to `False`): - Whether or not to compress the time dimension. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int = 3, - stride: int = 2, - padding: int = 0, - compress_time: bool = False, - ): - super().__init__() - - self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding) - self.compress_time = compress_time - - def forward(self, x: torch.Tensor) -> torch.Tensor: - if self.compress_time: - batch_size, channels, frames, height, width = x.shape - - # (batch_size, channels, frames, height, width) -> (batch_size, height, width, channels, frames) -> (batch_size * height * width, channels, frames) - x = x.permute(0, 3, 4, 1, 2).reshape(batch_size * height * width, channels, frames) - - if x.shape[-1] % 2 == 1: - x_first, x_rest = x[..., 0], x[..., 1:] - if x_rest.shape[-1] > 0: - # (batch_size * height * width, channels, frames - 1) -> (batch_size * height * width, channels, (frames - 1) // 2) - x_rest = F.avg_pool1d(x_rest, kernel_size=2, stride=2) - - x = torch.cat([x_first[..., None], x_rest], dim=-1) - # (batch_size * height * width, channels, (frames // 2) + 1) -> (batch_size, height, width, channels, (frames // 2) + 1) -> (batch_size, channels, (frames // 2) + 1, height, width) - x = x.reshape(batch_size, height, width, channels, x.shape[-1]).permute(0, 3, 4, 1, 2) - else: - # (batch_size * height * width, channels, frames) -> (batch_size * height * width, channels, frames // 2) - x = F.avg_pool1d(x, kernel_size=2, stride=2) - # (batch_size * height * width, channels, frames // 2) -> (batch_size, height, width, channels, frames // 2) -> (batch_size, channels, frames // 2, height, width) - x = x.reshape(batch_size, height, width, channels, x.shape[-1]).permute(0, 3, 4, 1, 2) - - # Pad the tensor - pad = (0, 1, 0, 1) - x = F.pad(x, pad, mode="constant", value=0) - batch_size, channels, frames, height, width = x.shape - # (batch_size, channels, frames, height, width) -> (batch_size, frames, channels, height, width) -> (batch_size * frames, channels, height, width) - x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * frames, channels, height, width) - x = self.conv(x) - # (batch_size * frames, channels, height, width) -> (batch_size, frames, channels, height, width) -> (batch_size, channels, frames, height, width) - x = x.reshape(batch_size, frames, x.shape[1], x.shape[2], x.shape[3]).permute(0, 2, 1, 3, 4) - return x - - -def downsample_2d( - hidden_states: torch.Tensor, - kernel: Optional[torch.Tensor] = None, - factor: int = 2, - gain: float = 1, -) -> torch.Tensor: - r"""Downsample2D a batch of 2D images with the given filter. - Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]` and downsamples each image with the - given filter. The filter is normalized so that if the input pixels are constant, they will be scaled by the - specified `gain`. Pixels outside the image are assumed to be zero, and the filter is padded with zeros so that its - shape is a multiple of the downsampling factor. - - Args: - hidden_states (`torch.Tensor`) - Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`. - kernel (`torch.Tensor`, *optional*): - FIR filter of the shape `[firH, firW]` or `[firN]` (separable). The default is `[1] * factor`, which - corresponds to average pooling. - factor (`int`, *optional*, default to `2`): - Integer downsampling factor. - gain (`float`, *optional*, default to `1.0`): - Scaling factor for signal magnitude. - - Returns: - output (`torch.Tensor`): - Tensor of the shape `[N, C, H // factor, W // factor]` - """ - - assert isinstance(factor, int) and factor >= 1 - if kernel is None: - kernel = [1] * factor - - kernel = torch.tensor(kernel, dtype=torch.float32) - if kernel.ndim == 1: - kernel = torch.outer(kernel, kernel) - kernel /= torch.sum(kernel) - - kernel = kernel * gain - pad_value = kernel.shape[0] - factor - output = upfirdn2d_native( - hidden_states, - kernel.to(device=hidden_states.device), - down=factor, - pad=((pad_value + 1) // 2, pad_value // 2), - ) - return output diff --git a/venv/lib/python3.10/site-packages/diffusers/models/embeddings.py b/venv/lib/python3.10/site-packages/diffusers/models/embeddings.py deleted file mode 100644 index 4f268bfa018fedf4f6f1d97e4bd9b64990ab9530..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/embeddings.py +++ /dev/null @@ -1,2621 +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 math -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch -import torch.nn.functional as F -from torch import nn - -from ..utils import deprecate -from .activations import FP32SiLU, get_activation -from .attention_processor import Attention - - -def get_timestep_embedding( - timesteps: torch.Tensor, - embedding_dim: int, - flip_sin_to_cos: bool = False, - downscale_freq_shift: float = 1, - scale: float = 1, - max_period: int = 10000, -) -> torch.Tensor: - """ - This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. - - Args - timesteps (torch.Tensor): - a 1-D Tensor of N indices, one per batch element. These may be fractional. - embedding_dim (int): - the dimension of the output. - flip_sin_to_cos (bool): - Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False) - downscale_freq_shift (float): - Controls the delta between frequencies between dimensions - scale (float): - Scaling factor applied to the embeddings. - max_period (int): - Controls the maximum frequency of the embeddings - Returns - torch.Tensor: an [N x dim] Tensor of positional embeddings. - """ - assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" - - half_dim = embedding_dim // 2 - exponent = -math.log(max_period) * torch.arange( - start=0, end=half_dim, dtype=torch.float32, device=timesteps.device - ) - exponent = exponent / (half_dim - downscale_freq_shift) - - emb = torch.exp(exponent) - emb = timesteps[:, None].float() * emb[None, :] - - # scale embeddings - emb = scale * emb - - # concat sine and cosine embeddings - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) - - # flip sine and cosine embeddings - if flip_sin_to_cos: - emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) - - # zero pad - if embedding_dim % 2 == 1: - emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) - return emb - - -def get_3d_sincos_pos_embed( - embed_dim: int, - spatial_size: Union[int, Tuple[int, int]], - temporal_size: int, - spatial_interpolation_scale: float = 1.0, - temporal_interpolation_scale: float = 1.0, - device: Optional[torch.device] = None, - output_type: str = "np", -) -> torch.Tensor: - r""" - Creates 3D sinusoidal positional embeddings. - - Args: - embed_dim (`int`): - The embedding dimension of inputs. It must be divisible by 16. - 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`): - The temporal dimension of positional embeddings (number of frames). - spatial_interpolation_scale (`float`, defaults to 1.0): - Scale factor for spatial grid interpolation. - temporal_interpolation_scale (`float`, defaults to 1.0): - Scale factor for temporal grid interpolation. - - Returns: - `torch.Tensor`: - The 3D sinusoidal positional embeddings of shape `[temporal_size, spatial_size[0] * spatial_size[1], - embed_dim]`. - """ - if output_type == "np": - return _get_3d_sincos_pos_embed_np( - embed_dim=embed_dim, - spatial_size=spatial_size, - temporal_size=temporal_size, - spatial_interpolation_scale=spatial_interpolation_scale, - temporal_interpolation_scale=temporal_interpolation_scale, - ) - if embed_dim % 4 != 0: - raise ValueError("`embed_dim` must be divisible by 4") - if isinstance(spatial_size, int): - spatial_size = (spatial_size, spatial_size) - - embed_dim_spatial = 3 * embed_dim // 4 - embed_dim_temporal = embed_dim // 4 - - # 1. Spatial - grid_h = torch.arange(spatial_size[1], device=device, dtype=torch.float32) / spatial_interpolation_scale - grid_w = torch.arange(spatial_size[0], device=device, dtype=torch.float32) / spatial_interpolation_scale - grid = torch.meshgrid(grid_w, grid_h, indexing="xy") # here w goes first - grid = torch.stack(grid, dim=0) - - grid = grid.reshape([2, 1, spatial_size[1], spatial_size[0]]) - pos_embed_spatial = get_2d_sincos_pos_embed_from_grid(embed_dim_spatial, grid, output_type="pt") - - # 2. Temporal - grid_t = torch.arange(temporal_size, device=device, dtype=torch.float32) / temporal_interpolation_scale - pos_embed_temporal = get_1d_sincos_pos_embed_from_grid(embed_dim_temporal, grid_t, output_type="pt") - - # 3. Concat - pos_embed_spatial = pos_embed_spatial[None, :, :] - pos_embed_spatial = pos_embed_spatial.repeat_interleave( - temporal_size, dim=0, output_size=pos_embed_spatial.shape[0] * temporal_size - ) # [T, H*W, D // 4 * 3] - - pos_embed_temporal = pos_embed_temporal[:, None, :] - pos_embed_temporal = pos_embed_temporal.repeat_interleave( - spatial_size[0] * spatial_size[1], dim=1 - ) # [T, H*W, D // 4] - - pos_embed = torch.concat([pos_embed_temporal, pos_embed_spatial], dim=-1) # [T, H*W, D] - return pos_embed - - -def _get_3d_sincos_pos_embed_np( - embed_dim: int, - spatial_size: Union[int, Tuple[int, int]], - temporal_size: int, - spatial_interpolation_scale: float = 1.0, - temporal_interpolation_scale: float = 1.0, -) -> np.ndarray: - r""" - Creates 3D sinusoidal positional embeddings. - - Args: - embed_dim (`int`): - The embedding dimension of inputs. It must be divisible by 16. - 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`): - The temporal dimension of positional embeddings (number of frames). - spatial_interpolation_scale (`float`, defaults to 1.0): - Scale factor for spatial grid interpolation. - temporal_interpolation_scale (`float`, defaults to 1.0): - Scale factor for temporal grid interpolation. - - Returns: - `np.ndarray`: - The 3D sinusoidal positional embeddings of shape `[temporal_size, spatial_size[0] * spatial_size[1], - embed_dim]`. - """ - deprecation_message = ( - "`get_3d_sincos_pos_embed` uses `torch` and supports `device`." - " `from_numpy` is no longer required." - " Pass `output_type='pt' to use the new version now." - ) - deprecate("output_type=='np'", "0.33.0", deprecation_message, standard_warn=False) - if embed_dim % 4 != 0: - raise ValueError("`embed_dim` must be divisible by 4") - if isinstance(spatial_size, int): - spatial_size = (spatial_size, spatial_size) - - embed_dim_spatial = 3 * embed_dim // 4 - embed_dim_temporal = embed_dim // 4 - - # 1. Spatial - grid_h = np.arange(spatial_size[1], dtype=np.float32) / spatial_interpolation_scale - grid_w = np.arange(spatial_size[0], dtype=np.float32) / spatial_interpolation_scale - grid = np.meshgrid(grid_w, grid_h) # here w goes first - grid = np.stack(grid, axis=0) - - grid = grid.reshape([2, 1, spatial_size[1], spatial_size[0]]) - pos_embed_spatial = get_2d_sincos_pos_embed_from_grid(embed_dim_spatial, grid) - - # 2. Temporal - grid_t = np.arange(temporal_size, dtype=np.float32) / temporal_interpolation_scale - pos_embed_temporal = get_1d_sincos_pos_embed_from_grid(embed_dim_temporal, grid_t) - - # 3. Concat - pos_embed_spatial = pos_embed_spatial[np.newaxis, :, :] - pos_embed_spatial = np.repeat(pos_embed_spatial, temporal_size, axis=0) # [T, H*W, D // 4 * 3] - - pos_embed_temporal = pos_embed_temporal[:, np.newaxis, :] - pos_embed_temporal = np.repeat(pos_embed_temporal, spatial_size[0] * spatial_size[1], axis=1) # [T, H*W, D // 4] - - pos_embed = np.concatenate([pos_embed_temporal, pos_embed_spatial], axis=-1) # [T, H*W, D] - return pos_embed - - -def get_2d_sincos_pos_embed( - embed_dim, - grid_size, - cls_token=False, - extra_tokens=0, - interpolation_scale=1.0, - base_size=16, - device: Optional[torch.device] = None, - output_type: str = "np", -): - """ - Creates 2D sinusoidal positional embeddings. - - Args: - embed_dim (`int`): - The embedding dimension. - grid_size (`int`): - The size of the grid height and width. - cls_token (`bool`, defaults to `False`): - Whether or not to add a classification token. - extra_tokens (`int`, defaults to `0`): - The number of extra tokens to add. - interpolation_scale (`float`, defaults to `1.0`): - The scale of the interpolation. - - Returns: - pos_embed (`torch.Tensor`): - Shape is either `[grid_size * grid_size, embed_dim]` if not using cls_token, or `[1 + grid_size*grid_size, - embed_dim]` if using cls_token - """ - if output_type == "np": - deprecation_message = ( - "`get_2d_sincos_pos_embed` uses `torch` and supports `device`." - " `from_numpy` is no longer required." - " Pass `output_type='pt' to use the new version now." - ) - deprecate("output_type=='np'", "0.33.0", deprecation_message, standard_warn=False) - return get_2d_sincos_pos_embed_np( - embed_dim=embed_dim, - grid_size=grid_size, - cls_token=cls_token, - extra_tokens=extra_tokens, - interpolation_scale=interpolation_scale, - base_size=base_size, - ) - if isinstance(grid_size, int): - grid_size = (grid_size, grid_size) - - grid_h = ( - torch.arange(grid_size[0], device=device, dtype=torch.float32) - / (grid_size[0] / base_size) - / interpolation_scale - ) - grid_w = ( - torch.arange(grid_size[1], device=device, dtype=torch.float32) - / (grid_size[1] / base_size) - / interpolation_scale - ) - grid = torch.meshgrid(grid_w, grid_h, indexing="xy") # here w goes first - grid = torch.stack(grid, dim=0) - - grid = grid.reshape([2, 1, grid_size[1], grid_size[0]]) - pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid, output_type=output_type) - if cls_token and extra_tokens > 0: - pos_embed = torch.concat([torch.zeros([extra_tokens, embed_dim]), pos_embed], dim=0) - return pos_embed - - -def get_2d_sincos_pos_embed_from_grid(embed_dim, grid, output_type="np"): - r""" - This function generates 2D sinusoidal positional embeddings from a grid. - - Args: - embed_dim (`int`): The embedding dimension. - grid (`torch.Tensor`): Grid of positions with shape `(H * W,)`. - - Returns: - `torch.Tensor`: The 2D sinusoidal positional embeddings with shape `(H * W, embed_dim)` - """ - if output_type == "np": - deprecation_message = ( - "`get_2d_sincos_pos_embed_from_grid` uses `torch` and supports `device`." - " `from_numpy` is no longer required." - " Pass `output_type='pt' to use the new version now." - ) - deprecate("output_type=='np'", "0.33.0", deprecation_message, standard_warn=False) - return get_2d_sincos_pos_embed_from_grid_np( - embed_dim=embed_dim, - grid=grid, - ) - if embed_dim % 2 != 0: - raise ValueError("embed_dim must be divisible by 2") - - # use half of dimensions to encode grid_h - emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0], output_type=output_type) # (H*W, D/2) - emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1], output_type=output_type) # (H*W, D/2) - - emb = torch.concat([emb_h, emb_w], dim=1) # (H*W, D) - return emb - - -def get_1d_sincos_pos_embed_from_grid(embed_dim, pos, output_type="np"): - """ - This function generates 1D positional embeddings from a grid. - - Args: - embed_dim (`int`): The embedding dimension `D` - pos (`torch.Tensor`): 1D tensor of positions with shape `(M,)` - - Returns: - `torch.Tensor`: Sinusoidal positional embeddings of shape `(M, D)`. - """ - if output_type == "np": - deprecation_message = ( - "`get_1d_sincos_pos_embed_from_grid` uses `torch` and supports `device`." - " `from_numpy` is no longer required." - " Pass `output_type='pt' to use the new version now." - ) - deprecate("output_type=='np'", "0.34.0", deprecation_message, standard_warn=False) - return get_1d_sincos_pos_embed_from_grid_np(embed_dim=embed_dim, pos=pos) - if embed_dim % 2 != 0: - raise ValueError("embed_dim must be divisible by 2") - - omega = torch.arange(embed_dim // 2, device=pos.device, dtype=torch.float64) - omega /= embed_dim / 2.0 - omega = 1.0 / 10000**omega # (D/2,) - - pos = pos.reshape(-1) # (M,) - out = torch.outer(pos, omega) # (M, D/2), outer product - - emb_sin = torch.sin(out) # (M, D/2) - emb_cos = torch.cos(out) # (M, D/2) - - emb = torch.concat([emb_sin, emb_cos], dim=1) # (M, D) - return emb - - -def get_2d_sincos_pos_embed_np( - embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16 -): - """ - Creates 2D sinusoidal positional embeddings. - - Args: - embed_dim (`int`): - The embedding dimension. - grid_size (`int`): - The size of the grid height and width. - cls_token (`bool`, defaults to `False`): - Whether or not to add a classification token. - extra_tokens (`int`, defaults to `0`): - The number of extra tokens to add. - interpolation_scale (`float`, defaults to `1.0`): - The scale of the interpolation. - - Returns: - pos_embed (`np.ndarray`): - Shape is either `[grid_size * grid_size, embed_dim]` if not using cls_token, or `[1 + grid_size*grid_size, - embed_dim]` if using cls_token - """ - if isinstance(grid_size, int): - grid_size = (grid_size, grid_size) - - grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size) / interpolation_scale - grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size) / interpolation_scale - grid = np.meshgrid(grid_w, grid_h) # here w goes first - grid = np.stack(grid, axis=0) - - grid = grid.reshape([2, 1, grid_size[1], grid_size[0]]) - pos_embed = get_2d_sincos_pos_embed_from_grid_np(embed_dim, grid) - if cls_token and extra_tokens > 0: - pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0) - return pos_embed - - -def get_2d_sincos_pos_embed_from_grid_np(embed_dim, grid): - r""" - This function generates 2D sinusoidal positional embeddings from a grid. - - Args: - embed_dim (`int`): The embedding dimension. - grid (`np.ndarray`): Grid of positions with shape `(H * W,)`. - - Returns: - `np.ndarray`: The 2D sinusoidal positional embeddings with shape `(H * W, embed_dim)` - """ - if embed_dim % 2 != 0: - raise ValueError("embed_dim must be divisible by 2") - - # use half of dimensions to encode grid_h - emb_h = get_1d_sincos_pos_embed_from_grid_np(embed_dim // 2, grid[0]) # (H*W, D/2) - emb_w = get_1d_sincos_pos_embed_from_grid_np(embed_dim // 2, grid[1]) # (H*W, D/2) - - emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D) - return emb - - -def get_1d_sincos_pos_embed_from_grid_np(embed_dim, pos): - """ - This function generates 1D positional embeddings from a grid. - - Args: - embed_dim (`int`): The embedding dimension `D` - pos (`numpy.ndarray`): 1D tensor of positions with shape `(M,)` - - Returns: - `numpy.ndarray`: Sinusoidal positional embeddings of shape `(M, D)`. - """ - if embed_dim % 2 != 0: - raise ValueError("embed_dim must be divisible by 2") - - omega = np.arange(embed_dim // 2, dtype=np.float64) - omega /= embed_dim / 2.0 - omega = 1.0 / 10000**omega # (D/2,) - - pos = pos.reshape(-1) # (M,) - out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product - - emb_sin = np.sin(out) # (M, D/2) - emb_cos = np.cos(out) # (M, D/2) - - emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) - return emb - - -class PatchEmbed(nn.Module): - """ - 2D Image to Patch Embedding with support for SD3 cropping. - - Args: - height (`int`, defaults to `224`): The height of the image. - width (`int`, defaults to `224`): The width of the image. - patch_size (`int`, defaults to `16`): The size of the patches. - in_channels (`int`, defaults to `3`): The number of input channels. - embed_dim (`int`, defaults to `768`): The output dimension of the embedding. - layer_norm (`bool`, defaults to `False`): Whether or not to use layer normalization. - flatten (`bool`, defaults to `True`): Whether or not to flatten the output. - bias (`bool`, defaults to `True`): Whether or not to use bias. - interpolation_scale (`float`, defaults to `1`): The scale of the interpolation. - pos_embed_type (`str`, defaults to `"sincos"`): The type of positional embedding. - pos_embed_max_size (`int`, defaults to `None`): The maximum size of the positional embedding. - """ - - def __init__( - self, - height=224, - width=224, - patch_size=16, - in_channels=3, - embed_dim=768, - layer_norm=False, - flatten=True, - bias=True, - interpolation_scale=1, - pos_embed_type="sincos", - pos_embed_max_size=None, # For SD3 cropping - ): - super().__init__() - - num_patches = (height // patch_size) * (width // patch_size) - self.flatten = flatten - self.layer_norm = layer_norm - self.pos_embed_max_size = pos_embed_max_size - - self.proj = nn.Conv2d( - in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias - ) - if layer_norm: - self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6) - else: - self.norm = None - - self.patch_size = patch_size - self.height, self.width = height // patch_size, width // patch_size - self.base_size = height // patch_size - self.interpolation_scale = interpolation_scale - - # Calculate positional embeddings based on max size or default - if pos_embed_max_size: - grid_size = pos_embed_max_size - else: - grid_size = int(num_patches**0.5) - - if pos_embed_type is None: - self.pos_embed = None - elif pos_embed_type == "sincos": - pos_embed = get_2d_sincos_pos_embed( - embed_dim, - grid_size, - base_size=self.base_size, - interpolation_scale=self.interpolation_scale, - output_type="pt", - ) - persistent = True if pos_embed_max_size else False - self.register_buffer("pos_embed", pos_embed.float().unsqueeze(0), persistent=persistent) - else: - raise ValueError(f"Unsupported pos_embed_type: {pos_embed_type}") - - def cropped_pos_embed(self, height, width): - """Crops positional embeddings for SD3 compatibility.""" - if self.pos_embed_max_size is None: - raise ValueError("`pos_embed_max_size` must be set for cropping.") - - height = height // self.patch_size - width = width // self.patch_size - if height > self.pos_embed_max_size: - raise ValueError( - f"Height ({height}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}." - ) - if width > self.pos_embed_max_size: - raise ValueError( - f"Width ({width}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}." - ) - - top = (self.pos_embed_max_size - height) // 2 - left = (self.pos_embed_max_size - width) // 2 - spatial_pos_embed = self.pos_embed.reshape(1, self.pos_embed_max_size, self.pos_embed_max_size, -1) - spatial_pos_embed = spatial_pos_embed[:, top : top + height, left : left + width, :] - spatial_pos_embed = spatial_pos_embed.reshape(1, -1, spatial_pos_embed.shape[-1]) - return spatial_pos_embed - - def forward(self, latent): - if self.pos_embed_max_size is not None: - height, width = latent.shape[-2:] - else: - height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size - latent = self.proj(latent) - if self.flatten: - latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC - if self.layer_norm: - latent = self.norm(latent) - if self.pos_embed is None: - return latent.to(latent.dtype) - # Interpolate or crop positional embeddings as needed - if self.pos_embed_max_size: - pos_embed = self.cropped_pos_embed(height, width) - else: - if self.height != height or self.width != width: - pos_embed = get_2d_sincos_pos_embed( - embed_dim=self.pos_embed.shape[-1], - grid_size=(height, width), - base_size=self.base_size, - interpolation_scale=self.interpolation_scale, - device=latent.device, - output_type="pt", - ) - pos_embed = pos_embed.float().unsqueeze(0) - else: - pos_embed = self.pos_embed - - return (latent + pos_embed).to(latent.dtype) - - -class LuminaPatchEmbed(nn.Module): - """ - 2D Image to Patch Embedding with support for Lumina-T2X - - Args: - patch_size (`int`, defaults to `2`): The size of the patches. - in_channels (`int`, defaults to `4`): The number of input channels. - embed_dim (`int`, defaults to `768`): The output dimension of the embedding. - bias (`bool`, defaults to `True`): Whether or not to use bias. - """ - - def __init__(self, patch_size=2, in_channels=4, embed_dim=768, bias=True): - super().__init__() - self.patch_size = patch_size - self.proj = nn.Linear( - in_features=patch_size * patch_size * in_channels, - out_features=embed_dim, - bias=bias, - ) - - 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. - - Returns: - 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). - """ - freqs_cis = freqs_cis.to(x[0].device) - patch_height = patch_width = self.patch_size - batch_size, channel, height, width = x.size() - height_tokens, width_tokens = height // patch_height, width // patch_width - - x = x.view(batch_size, channel, height_tokens, patch_height, width_tokens, patch_width).permute( - 0, 2, 4, 1, 3, 5 - ) - x = x.flatten(3) - x = self.proj(x) - x = x.flatten(1, 2) - - mask = torch.ones(x.shape[0], x.shape[1], dtype=torch.int32, device=x.device) - - return ( - x, - mask, - [(height, width)] * batch_size, - freqs_cis[:height_tokens, :width_tokens].flatten(0, 1).unsqueeze(0), - ) - - -class CogVideoXPatchEmbed(nn.Module): - def __init__( - self, - patch_size: int = 2, - patch_size_t: Optional[int] = None, - in_channels: int = 16, - embed_dim: int = 1920, - text_embed_dim: int = 4096, - bias: bool = True, - sample_width: int = 90, - sample_height: int = 60, - sample_frames: int = 49, - temporal_compression_ratio: int = 4, - max_text_seq_length: int = 226, - spatial_interpolation_scale: float = 1.875, - temporal_interpolation_scale: float = 1.0, - use_positional_embeddings: bool = True, - use_learned_positional_embeddings: bool = True, - ) -> None: - super().__init__() - - self.patch_size = patch_size - self.patch_size_t = patch_size_t - self.embed_dim = embed_dim - self.sample_height = sample_height - self.sample_width = sample_width - self.sample_frames = sample_frames - self.temporal_compression_ratio = temporal_compression_ratio - self.max_text_seq_length = max_text_seq_length - self.spatial_interpolation_scale = spatial_interpolation_scale - self.temporal_interpolation_scale = temporal_interpolation_scale - self.use_positional_embeddings = use_positional_embeddings - self.use_learned_positional_embeddings = use_learned_positional_embeddings - - if patch_size_t is None: - # CogVideoX 1.0 checkpoints - self.proj = nn.Conv2d( - in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias - ) - else: - # CogVideoX 1.5 checkpoints - self.proj = nn.Linear(in_channels * patch_size * patch_size * patch_size_t, embed_dim) - - self.text_proj = nn.Linear(text_embed_dim, embed_dim) - - if use_positional_embeddings or use_learned_positional_embeddings: - persistent = use_learned_positional_embeddings - pos_embedding = self._get_positional_embeddings(sample_height, sample_width, sample_frames) - 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 - ) -> torch.Tensor: - post_patch_height = sample_height // self.patch_size - post_patch_width = sample_width // self.patch_size - post_time_compression_frames = (sample_frames - 1) // self.temporal_compression_ratio + 1 - num_patches = post_patch_height * post_patch_width * post_time_compression_frames - - pos_embedding = get_3d_sincos_pos_embed( - self.embed_dim, - (post_patch_width, post_patch_height), - post_time_compression_frames, - self.spatial_interpolation_scale, - self.temporal_interpolation_scale, - device=device, - output_type="pt", - ) - pos_embedding = pos_embedding.flatten(0, 1) - joint_pos_embedding = pos_embedding.new_zeros( - 1, self.max_text_seq_length + num_patches, self.embed_dim, requires_grad=False - ) - joint_pos_embedding.data[:, self.max_text_seq_length :].copy_(pos_embedding) - - return joint_pos_embedding - - def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor): - r""" - Args: - text_embeds (`torch.Tensor`): - Input text embeddings. Expected shape: (batch_size, seq_length, embedding_dim). - image_embeds (`torch.Tensor`): - Input image embeddings. Expected shape: (batch_size, num_frames, channels, height, width). - """ - text_embeds = self.text_proj(text_embeds) - - batch_size, num_frames, channels, height, width = image_embeds.shape - - if self.patch_size_t is None: - image_embeds = image_embeds.reshape(-1, channels, height, width) - image_embeds = self.proj(image_embeds) - image_embeds = image_embeds.view(batch_size, num_frames, *image_embeds.shape[1:]) - image_embeds = image_embeds.flatten(3).transpose(2, 3) # [batch, num_frames, height x width, channels] - image_embeds = image_embeds.flatten(1, 2) # [batch, num_frames x height x width, channels] - else: - p = self.patch_size - p_t = self.patch_size_t - - image_embeds = image_embeds.permute(0, 1, 3, 4, 2) - image_embeds = image_embeds.reshape( - batch_size, num_frames // p_t, p_t, height // p, p, width // p, p, channels - ) - image_embeds = image_embeds.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(4, 7).flatten(1, 3) - image_embeds = self.proj(image_embeds) - - embeds = torch.cat( - [text_embeds, image_embeds], dim=1 - ).contiguous() # [batch, seq_length + num_frames x height x width, channels] - - if self.use_positional_embeddings or self.use_learned_positional_embeddings: - if self.use_learned_positional_embeddings and (self.sample_width != width or self.sample_height != height): - raise ValueError( - "It is currently not possible to generate videos at a different resolution that the defaults. This should only be the case with 'THUDM/CogVideoX-5b-I2V'." - "If you think this is incorrect, please open an issue at https://github.com/huggingface/diffusers/issues." - ) - - pre_time_compression_frames = (num_frames - 1) * self.temporal_compression_ratio + 1 - - if ( - self.sample_height != height - or self.sample_width != width - or self.sample_frames != pre_time_compression_frames - ): - pos_embedding = self._get_positional_embeddings( - height, width, pre_time_compression_frames, device=embeds.device - ) - else: - pos_embedding = self.pos_embedding - - pos_embedding = pos_embedding.to(dtype=embeds.dtype) - embeds = embeds + pos_embedding - - return embeds - - -class CogView3PlusPatchEmbed(nn.Module): - def __init__( - self, - in_channels: int = 16, - hidden_size: int = 2560, - patch_size: int = 2, - text_hidden_size: int = 4096, - pos_embed_max_size: int = 128, - ): - super().__init__() - self.in_channels = in_channels - self.hidden_size = hidden_size - self.patch_size = patch_size - self.text_hidden_size = text_hidden_size - self.pos_embed_max_size = pos_embed_max_size - # Linear projection for image patches - self.proj = nn.Linear(in_channels * patch_size**2, hidden_size) - - # Linear projection for text embeddings - self.text_proj = nn.Linear(text_hidden_size, hidden_size) - - pos_embed = get_2d_sincos_pos_embed( - hidden_size, pos_embed_max_size, base_size=pos_embed_max_size, output_type="pt" - ) - pos_embed = pos_embed.reshape(pos_embed_max_size, pos_embed_max_size, hidden_size) - self.register_buffer("pos_embed", pos_embed.float(), persistent=False) - - def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor) -> torch.Tensor: - batch_size, channel, height, width = hidden_states.shape - - if height % self.patch_size != 0 or width % self.patch_size != 0: - raise ValueError("Height and width must be divisible by patch size") - - height = height // self.patch_size - width = width // self.patch_size - hidden_states = hidden_states.view(batch_size, channel, height, self.patch_size, width, self.patch_size) - hidden_states = hidden_states.permute(0, 2, 4, 1, 3, 5).contiguous() - hidden_states = hidden_states.view(batch_size, height * width, channel * self.patch_size * self.patch_size) - - # Project the patches - hidden_states = self.proj(hidden_states) - encoder_hidden_states = self.text_proj(encoder_hidden_states) - hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) - - # Calculate text_length - text_length = encoder_hidden_states.shape[1] - - image_pos_embed = self.pos_embed[:height, :width].reshape(height * width, -1) - text_pos_embed = torch.zeros( - (text_length, self.hidden_size), dtype=image_pos_embed.dtype, device=image_pos_embed.device - ) - pos_embed = torch.cat([text_pos_embed, image_pos_embed], dim=0)[None, ...] - - return (hidden_states + pos_embed).to(hidden_states.dtype) - - -def get_3d_rotary_pos_embed( - embed_dim, - crops_coords, - grid_size, - temporal_size, - 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]]: - """ - RoPE for video tokens with 3D structure. - - Args: - embed_dim: (`int`): - The embedding dimension size, corresponding to hidden_size_head. - crops_coords (`Tuple[int]`): - The top-left and bottom-right coordinates of the crop. - grid_size (`Tuple[int]`): - The grid size of the spatial positional embedding (height, width). - temporal_size (`int`): - The size of the temporal dimension. - theta (`float`): - Scaling factor for frequency computation. - grid_type (`str`): - Whether to use "linspace" or "slice" to compute grids. - - Returns: - `torch.Tensor`: positional embedding with shape `(temporal_size * grid_size[0] * grid_size[1], embed_dim/2)`. - """ - if use_real is not True: - raise ValueError(" `use_real = False` is not currently supported for get_3d_rotary_pos_embed") - - if grid_type == "linspace": - start, stop = crops_coords - grid_size_h, grid_size_w = grid_size - grid_h = torch.linspace( - start[0], stop[0] * (grid_size_h - 1) / grid_size_h, grid_size_h, device=device, dtype=torch.float32 - ) - grid_w = torch.linspace( - start[1], stop[1] * (grid_size_w - 1) / grid_size_w, grid_size_w, device=device, dtype=torch.float32 - ) - grid_t = torch.arange(temporal_size, device=device, dtype=torch.float32) - grid_t = torch.linspace( - 0, temporal_size * (temporal_size - 1) / temporal_size, temporal_size, device=device, dtype=torch.float32 - ) - elif grid_type == "slice": - max_h, max_w = max_size - grid_size_h, grid_size_w = grid_size - grid_h = torch.arange(max_h, device=device, dtype=torch.float32) - grid_w = torch.arange(max_w, device=device, dtype=torch.float32) - grid_t = torch.arange(temporal_size, device=device, dtype=torch.float32) - else: - raise ValueError("Invalid value passed for `grid_type`.") - - # Compute dimensions for each axis - dim_t = embed_dim // 4 - dim_h = embed_dim // 8 * 3 - dim_w = embed_dim // 8 * 3 - - # Temporal frequencies - freqs_t = get_1d_rotary_pos_embed(dim_t, grid_t, theta=theta, use_real=True) - # Spatial frequencies for height and width - freqs_h = get_1d_rotary_pos_embed(dim_h, grid_h, theta=theta, use_real=True) - freqs_w = get_1d_rotary_pos_embed(dim_w, grid_w, theta=theta, use_real=True) - - # BroadCast and concatenate temporal and spaial frequencie (height and width) into a 3d tensor - def combine_time_height_width(freqs_t, freqs_h, freqs_w): - freqs_t = freqs_t[:, None, None, :].expand( - -1, grid_size_h, grid_size_w, -1 - ) # temporal_size, grid_size_h, grid_size_w, dim_t - freqs_h = freqs_h[None, :, None, :].expand( - temporal_size, -1, grid_size_w, -1 - ) # temporal_size, grid_size_h, grid_size_2, dim_h - freqs_w = freqs_w[None, None, :, :].expand( - temporal_size, grid_size_h, -1, -1 - ) # temporal_size, grid_size_h, grid_size_2, dim_w - - freqs = torch.cat( - [freqs_t, freqs_h, freqs_w], dim=-1 - ) # temporal_size, grid_size_h, grid_size_w, (dim_t + dim_h + dim_w) - freqs = freqs.view( - temporal_size * grid_size_h * grid_size_w, -1 - ) # (temporal_size * grid_size_h * grid_size_w), (dim_t + dim_h + dim_w) - return freqs - - t_cos, t_sin = freqs_t # both t_cos and t_sin has shape: temporal_size, dim_t - h_cos, h_sin = freqs_h # both h_cos and h_sin has shape: grid_size_h, dim_h - w_cos, w_sin = freqs_w # both w_cos and w_sin has shape: grid_size_w, dim_w - - if grid_type == "slice": - t_cos, t_sin = t_cos[:temporal_size], t_sin[:temporal_size] - h_cos, h_sin = h_cos[:grid_size_h], h_sin[:grid_size_h] - w_cos, w_sin = w_cos[:grid_size_w], w_sin[:grid_size_w] - - cos = combine_time_height_width(t_cos, h_cos, w_cos) - sin = combine_time_height_width(t_sin, h_sin, w_sin) - return cos, sin - - -def get_3d_rotary_pos_embed_allegro( - embed_dim, - crops_coords, - grid_size, - temporal_size, - 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]]: - # TODO(aryan): docs - start, stop = crops_coords - grid_size_h, grid_size_w = grid_size - interpolation_scale_t, interpolation_scale_h, interpolation_scale_w = interpolation_scale - grid_t = torch.linspace( - 0, temporal_size * (temporal_size - 1) / temporal_size, temporal_size, device=device, dtype=torch.float32 - ) - grid_h = torch.linspace( - start[0], stop[0] * (grid_size_h - 1) / grid_size_h, grid_size_h, device=device, dtype=torch.float32 - ) - grid_w = torch.linspace( - start[1], stop[1] * (grid_size_w - 1) / grid_size_w, grid_size_w, device=device, dtype=torch.float32 - ) - - # Compute dimensions for each axis - dim_t = embed_dim // 3 - dim_h = embed_dim // 3 - dim_w = embed_dim // 3 - - # Temporal frequencies - freqs_t = get_1d_rotary_pos_embed( - dim_t, grid_t / interpolation_scale_t, theta=theta, use_real=True, repeat_interleave_real=False - ) - # Spatial frequencies for height and width - freqs_h = get_1d_rotary_pos_embed( - dim_h, grid_h / interpolation_scale_h, theta=theta, use_real=True, repeat_interleave_real=False - ) - freqs_w = get_1d_rotary_pos_embed( - dim_w, grid_w / interpolation_scale_w, theta=theta, use_real=True, repeat_interleave_real=False - ) - - return freqs_t, freqs_h, freqs_w, grid_t, grid_h, grid_w - - -def get_2d_rotary_pos_embed( - embed_dim, crops_coords, grid_size, use_real=True, device: Optional[torch.device] = None, output_type: str = "np" -): - """ - RoPE for image tokens with 2d structure. - - Args: - embed_dim: (`int`): - The embedding dimension size - crops_coords (`Tuple[int]`) - The top-left and bottom-right coordinates of the crop. - 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. - device: (`torch.device`, **optional**): - The device used to create tensors. - - Returns: - `torch.Tensor`: positional embedding with shape `( grid_size * grid_size, embed_dim/2)`. - """ - if output_type == "np": - deprecation_message = ( - "`get_2d_sincos_pos_embed` uses `torch` and supports `device`." - " `from_numpy` is no longer required." - " Pass `output_type='pt' to use the new version now." - ) - deprecate("output_type=='np'", "0.33.0", deprecation_message, standard_warn=False) - return _get_2d_rotary_pos_embed_np( - embed_dim=embed_dim, - crops_coords=crops_coords, - grid_size=grid_size, - use_real=use_real, - ) - start, stop = crops_coords - # scale end by (steps−1)/steps matches np.linspace(..., endpoint=False) - grid_h = torch.linspace( - start[0], stop[0] * (grid_size[0] - 1) / grid_size[0], grid_size[0], device=device, dtype=torch.float32 - ) - grid_w = torch.linspace( - start[1], stop[1] * (grid_size[1] - 1) / grid_size[1], grid_size[1], device=device, dtype=torch.float32 - ) - grid = torch.meshgrid(grid_w, grid_h, indexing="xy") - grid = torch.stack(grid, dim=0) # [2, W, H] - - grid = grid.reshape([2, 1, *grid.shape[1:]]) - pos_embed = get_2d_rotary_pos_embed_from_grid(embed_dim, grid, use_real=use_real) - return pos_embed - - -def _get_2d_rotary_pos_embed_np(embed_dim, crops_coords, grid_size, use_real=True): - """ - RoPE for image tokens with 2d structure. - - Args: - embed_dim: (`int`): - The embedding dimension size - crops_coords (`Tuple[int]`) - The top-left and bottom-right coordinates of the crop. - 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. - - Returns: - `torch.Tensor`: positional embedding with shape `( grid_size * grid_size, embed_dim/2)`. - """ - start, stop = crops_coords - grid_h = np.linspace(start[0], stop[0], grid_size[0], endpoint=False, dtype=np.float32) - grid_w = np.linspace(start[1], stop[1], grid_size[1], endpoint=False, dtype=np.float32) - grid = np.meshgrid(grid_w, grid_h) # here w goes first - grid = np.stack(grid, axis=0) # [2, W, H] - - grid = grid.reshape([2, 1, *grid.shape[1:]]) - pos_embed = get_2d_rotary_pos_embed_from_grid(embed_dim, grid, use_real=use_real) - return pos_embed - - -def get_2d_rotary_pos_embed_from_grid(embed_dim, grid, use_real=False): - """ - Get 2D RoPE from grid. - - Args: - embed_dim: (`int`): - The embedding dimension size, corresponding to hidden_size_head. - grid (`np.ndarray`): - The grid of the positional embedding. - use_real (`bool`): - If True, return real part and imaginary part separately. Otherwise, return complex numbers. - - Returns: - `torch.Tensor`: positional embedding with shape `( grid_size * grid_size, embed_dim/2)`. - """ - assert embed_dim % 4 == 0 - - # use half of dimensions to encode grid_h - emb_h = get_1d_rotary_pos_embed( - embed_dim // 2, grid[0].reshape(-1), use_real=use_real - ) # (H*W, D/2) if use_real else (H*W, D/4) - emb_w = get_1d_rotary_pos_embed( - embed_dim // 2, grid[1].reshape(-1), use_real=use_real - ) # (H*W, D/2) if use_real else (H*W, D/4) - - if use_real: - cos = torch.cat([emb_h[0], emb_w[0]], dim=1) # (H*W, D) - sin = torch.cat([emb_h[1], emb_w[1]], dim=1) # (H*W, D) - return cos, sin - else: - emb = torch.cat([emb_h, emb_w], dim=1) # (H*W, D/2) - return emb - - -def get_2d_rotary_pos_embed_lumina(embed_dim, len_h, len_w, linear_factor=1.0, ntk_factor=1.0): - """ - Get 2D RoPE from grid. - - Args: - embed_dim: (`int`): - The embedding dimension size, corresponding to hidden_size_head. - grid (`np.ndarray`): - The grid of the positional embedding. - linear_factor (`float`): - The linear factor of the positional embedding, which is used to scale the positional embedding in the linear - layer. - ntk_factor (`float`): - The ntk factor of the positional embedding, which is used to scale the positional embedding in the ntk layer. - - Returns: - `torch.Tensor`: positional embedding with shape `( grid_size * grid_size, embed_dim/2)`. - """ - assert embed_dim % 4 == 0 - - emb_h = get_1d_rotary_pos_embed( - embed_dim // 2, len_h, linear_factor=linear_factor, ntk_factor=ntk_factor - ) # (H, D/4) - emb_w = get_1d_rotary_pos_embed( - embed_dim // 2, len_w, linear_factor=linear_factor, ntk_factor=ntk_factor - ) # (W, D/4) - emb_h = emb_h.view(len_h, 1, embed_dim // 4, 1).repeat(1, len_w, 1, 1) # (H, W, D/4, 1) - emb_w = emb_w.view(1, len_w, embed_dim // 4, 1).repeat(len_h, 1, 1, 1) # (H, W, D/4, 1) - - emb = torch.cat([emb_h, emb_w], dim=-1).flatten(2) # (H, W, D/2) - return emb - - -def get_1d_rotary_pos_embed( - dim: int, - pos: Union[np.ndarray, int], - theta: float = 10000.0, - use_real=False, - linear_factor=1.0, - ntk_factor=1.0, - repeat_interleave_real=True, - freqs_dtype=torch.float32, # torch.float32, torch.float64 (flux) -): - """ - Precompute the frequency tensor for complex exponentials (cis) with given dimensions. - - This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' and the end - index 'end'. The 'theta' parameter scales the frequencies. The returned tensor contains complex values in complex64 - data type. - - Args: - dim (`int`): Dimension of the frequency tensor. - pos (`np.ndarray` or `int`): Position indices for the frequency tensor. [S] or scalar - theta (`float`, *optional*, defaults to 10000.0): - Scaling factor for frequency computation. Defaults to 10000.0. - use_real (`bool`, *optional*): - If True, return real part and imaginary part separately. Otherwise, return complex numbers. - linear_factor (`float`, *optional*, defaults to 1.0): - Scaling factor for the context extrapolation. Defaults to 1.0. - ntk_factor (`float`, *optional*, defaults to 1.0): - Scaling factor for the NTK-Aware RoPE. Defaults to 1.0. - repeat_interleave_real (`bool`, *optional*, defaults to `True`): - If `True` and `use_real`, real part and imaginary part are each interleaved with themselves to reach `dim`. - Otherwise, they are concateanted with themselves. - freqs_dtype (`torch.float32` or `torch.float64`, *optional*, defaults to `torch.float32`): - the dtype of the frequency tensor. - Returns: - `torch.Tensor`: Precomputed frequency tensor with complex exponentials. [S, D/2] - """ - assert dim % 2 == 0 - - if isinstance(pos, int): - pos = torch.arange(pos) - if isinstance(pos, np.ndarray): - pos = torch.from_numpy(pos) # type: ignore # [S] - - theta = theta * ntk_factor - freqs = ( - 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype, device=pos.device) / dim)) / linear_factor - ) # [D/2] - freqs = torch.outer(pos, freqs) # type: ignore # [S, D/2] - is_npu = freqs.device.type == "npu" - if is_npu: - freqs = freqs.float() - if use_real and repeat_interleave_real: - # flux, hunyuan-dit, cogvideox - freqs_cos = freqs.cos().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float() # [S, D] - freqs_sin = freqs.sin().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float() # [S, D] - return freqs_cos, freqs_sin - elif use_real: - # stable audio, allegro - freqs_cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).float() # [S, D] - freqs_sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).float() # [S, D] - return freqs_cos, freqs_sin - else: - # lumina - freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 # [S, D/2] - return freqs_cis - - -def apply_rotary_emb( - x: torch.Tensor, - freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], - use_real: bool = True, - use_real_unbind_dim: int = -1, -) -> 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 - reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting - tensors contain rotary embeddings and are returned as real tensors. - - 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],) - - Returns: - 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] - cos = cos[None, None] - sin = sin[None, None] - cos, sin = cos.to(x.device), sin.to(x.device) - - if use_real_unbind_dim == -1: - # Used for flux, cogvideox, hunyuan-dit - x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, H, S, D//2] - x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) - elif use_real_unbind_dim == -2: - # Used for Stable Audio, OmniGen, CogView4 and Cosmos - x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, H, S, D//2] - x_rotated = torch.cat([-x_imag, x_real], dim=-1) - else: - raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.") - - out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) - - return out - else: - # used for lumina - x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) - freqs_cis = freqs_cis.unsqueeze(2) - x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) - - return x_out.type_as(x) - - -def apply_rotary_emb_allegro(x: torch.Tensor, freqs_cis, positions): - # TODO(aryan): rewrite - def apply_1d_rope(tokens, pos, cos, sin): - cos = F.embedding(pos, cos)[:, None, :, :] - sin = F.embedding(pos, sin)[:, None, :, :] - x1, x2 = tokens[..., : tokens.shape[-1] // 2], tokens[..., tokens.shape[-1] // 2 :] - tokens_rotated = torch.cat((-x2, x1), dim=-1) - return (tokens.float() * cos + tokens_rotated.float() * sin).to(tokens.dtype) - - (t_cos, t_sin), (h_cos, h_sin), (w_cos, w_sin) = freqs_cis - t, h, w = x.chunk(3, dim=-1) - t = apply_1d_rope(t, positions[0], t_cos, t_sin) - h = apply_1d_rope(h, positions[1], h_cos, h_sin) - w = apply_1d_rope(w, positions[2], w_cos, w_sin) - x = torch.cat([t, h, w], dim=-1) - return x - - -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]): - super().__init__() - self.theta = theta - self.axes_dim = axes_dim - - def forward(self, ids: torch.Tensor) -> torch.Tensor: - n_axes = ids.shape[-1] - cos_out = [] - sin_out = [] - pos = ids.float() - is_mps = ids.device.type == "mps" - is_npu = ids.device.type == "npu" - freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - for i in range(n_axes): - cos, sin = get_1d_rotary_pos_embed( - self.axes_dim[i], - pos[:, i], - theta=self.theta, - repeat_interleave_real=True, - use_real=True, - freqs_dtype=freqs_dtype, - ) - cos_out.append(cos) - sin_out.append(sin) - freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) - freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) - return freqs_cos, freqs_sin - - -class TimestepEmbedding(nn.Module): - def __init__( - self, - in_channels: int, - time_embed_dim: int, - act_fn: str = "silu", - out_dim: int = None, - post_act_fn: Optional[str] = None, - cond_proj_dim=None, - sample_proj_bias=True, - ): - super().__init__() - - self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias) - - if cond_proj_dim is not None: - self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False) - else: - self.cond_proj = None - - self.act = get_activation(act_fn) - - if out_dim is not None: - time_embed_dim_out = out_dim - else: - time_embed_dim_out = time_embed_dim - self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias) - - if post_act_fn is None: - self.post_act = None - else: - self.post_act = get_activation(post_act_fn) - - def forward(self, sample, condition=None): - if condition is not None: - sample = sample + self.cond_proj(condition) - sample = self.linear_1(sample) - - if self.act is not None: - sample = self.act(sample) - - sample = self.linear_2(sample) - - if self.post_act is not None: - sample = self.post_act(sample) - return sample - - -class Timesteps(nn.Module): - def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1): - super().__init__() - self.num_channels = num_channels - self.flip_sin_to_cos = flip_sin_to_cos - self.downscale_freq_shift = downscale_freq_shift - self.scale = scale - - def forward(self, timesteps: torch.Tensor) -> torch.Tensor: - t_emb = get_timestep_embedding( - timesteps, - self.num_channels, - flip_sin_to_cos=self.flip_sin_to_cos, - downscale_freq_shift=self.downscale_freq_shift, - scale=self.scale, - ) - return t_emb - - -class GaussianFourierProjection(nn.Module): - """Gaussian Fourier embeddings for noise levels.""" - - def __init__( - self, embedding_size: int = 256, scale: float = 1.0, set_W_to_weight=True, log=True, flip_sin_to_cos=False - ): - super().__init__() - self.weight = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False) - self.log = log - self.flip_sin_to_cos = flip_sin_to_cos - - if set_W_to_weight: - # to delete later - del self.weight - self.W = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False) - self.weight = self.W - del self.W - - def forward(self, x): - if self.log: - x = torch.log(x) - - x_proj = x[:, None] * self.weight[None, :] * 2 * np.pi - - if self.flip_sin_to_cos: - out = torch.cat([torch.cos(x_proj), torch.sin(x_proj)], dim=-1) - else: - out = torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1) - return out - - -class SinusoidalPositionalEmbedding(nn.Module): - """Apply positional information to a sequence of embeddings. - - Takes in a sequence of embeddings with shape (batch_size, seq_length, embed_dim) and adds positional embeddings to - them - - Args: - embed_dim: (int): Dimension of the positional embedding. - max_seq_length: Maximum sequence length to apply positional embeddings - - """ - - def __init__(self, embed_dim: int, max_seq_length: int = 32): - super().__init__() - position = torch.arange(max_seq_length).unsqueeze(1) - div_term = torch.exp(torch.arange(0, embed_dim, 2) * (-math.log(10000.0) / embed_dim)) - pe = torch.zeros(1, max_seq_length, embed_dim) - pe[0, :, 0::2] = torch.sin(position * div_term) - pe[0, :, 1::2] = torch.cos(position * div_term) - self.register_buffer("pe", pe) - - def forward(self, x): - _, seq_length, _ = x.shape - x = x + self.pe[:, :seq_length] - return x - - -class ImagePositionalEmbeddings(nn.Module): - """ - Converts latent image classes into vector embeddings. Sums the vector embeddings with positional embeddings for the - height and width of the latent space. - - For more details, see figure 10 of the dall-e paper: https://huggingface.co/papers/2102.12092 - - For VQ-diffusion: - - Output vector embeddings are used as input for the transformer. - - Note that the vector embeddings for the transformer are different than the vector embeddings from the VQVAE. - - Args: - num_embed (`int`): - Number of embeddings for the latent pixels embeddings. - height (`int`): - Height of the latent image i.e. the number of height embeddings. - width (`int`): - Width of the latent image i.e. the number of width embeddings. - embed_dim (`int`): - Dimension of the produced vector embeddings. Used for the latent pixel, height, and width embeddings. - """ - - def __init__( - self, - num_embed: int, - height: int, - width: int, - embed_dim: int, - ): - super().__init__() - - self.height = height - self.width = width - self.num_embed = num_embed - self.embed_dim = embed_dim - - self.emb = nn.Embedding(self.num_embed, embed_dim) - self.height_emb = nn.Embedding(self.height, embed_dim) - self.width_emb = nn.Embedding(self.width, embed_dim) - - def forward(self, index): - emb = self.emb(index) - - height_emb = self.height_emb(torch.arange(self.height, device=index.device).view(1, self.height)) - - # 1 x H x D -> 1 x H x 1 x D - height_emb = height_emb.unsqueeze(2) - - width_emb = self.width_emb(torch.arange(self.width, device=index.device).view(1, self.width)) - - # 1 x W x D -> 1 x 1 x W x D - width_emb = width_emb.unsqueeze(1) - - pos_emb = height_emb + width_emb - - # 1 x H x W x D -> 1 x L xD - pos_emb = pos_emb.view(1, self.height * self.width, -1) - - emb = emb + pos_emb[:, : emb.shape[1], :] - - return emb - - -class LabelEmbedding(nn.Module): - """ - Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance. - - Args: - num_classes (`int`): The number of classes. - hidden_size (`int`): The size of the vector embeddings. - dropout_prob (`float`): The probability of dropping a label. - """ - - def __init__(self, num_classes, hidden_size, dropout_prob): - super().__init__() - use_cfg_embedding = dropout_prob > 0 - self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size) - self.num_classes = num_classes - self.dropout_prob = dropout_prob - - def token_drop(self, labels, force_drop_ids=None): - """ - Drops labels to enable classifier-free guidance. - """ - if force_drop_ids is None: - drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob - else: - drop_ids = torch.tensor(force_drop_ids == 1) - labels = torch.where(drop_ids, self.num_classes, labels) - return labels - - def forward(self, labels: torch.LongTensor, force_drop_ids=None): - use_dropout = self.dropout_prob > 0 - if (self.training and use_dropout) or (force_drop_ids is not None): - labels = self.token_drop(labels, force_drop_ids) - embeddings = self.embedding_table(labels) - return embeddings - - -class TextImageProjection(nn.Module): - def __init__( - self, - text_embed_dim: int = 1024, - image_embed_dim: int = 768, - cross_attention_dim: int = 768, - num_image_text_embeds: int = 10, - ): - super().__init__() - - self.num_image_text_embeds = num_image_text_embeds - self.image_embeds = nn.Linear(image_embed_dim, self.num_image_text_embeds * cross_attention_dim) - self.text_proj = nn.Linear(text_embed_dim, cross_attention_dim) - - def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor): - batch_size = text_embeds.shape[0] - - # image - image_text_embeds = self.image_embeds(image_embeds) - image_text_embeds = image_text_embeds.reshape(batch_size, self.num_image_text_embeds, -1) - - # text - text_embeds = self.text_proj(text_embeds) - - return torch.cat([image_text_embeds, text_embeds], dim=1) - - -class ImageProjection(nn.Module): - def __init__( - self, - image_embed_dim: int = 768, - cross_attention_dim: int = 768, - num_image_text_embeds: int = 32, - ): - super().__init__() - - self.num_image_text_embeds = num_image_text_embeds - self.image_embeds = nn.Linear(image_embed_dim, self.num_image_text_embeds * cross_attention_dim) - self.norm = nn.LayerNorm(cross_attention_dim) - - def forward(self, image_embeds: torch.Tensor): - batch_size = image_embeds.shape[0] - - # image - image_embeds = self.image_embeds(image_embeds.to(self.image_embeds.weight.dtype)) - image_embeds = image_embeds.reshape(batch_size, self.num_image_text_embeds, -1) - image_embeds = self.norm(image_embeds) - return image_embeds - - -class IPAdapterFullImageProjection(nn.Module): - def __init__(self, image_embed_dim=1024, cross_attention_dim=1024): - super().__init__() - from .attention import FeedForward - - self.ff = FeedForward(image_embed_dim, cross_attention_dim, mult=1, activation_fn="gelu") - self.norm = nn.LayerNorm(cross_attention_dim) - - def forward(self, image_embeds: torch.Tensor): - return self.norm(self.ff(image_embeds)) - - -class IPAdapterFaceIDImageProjection(nn.Module): - def __init__(self, image_embed_dim=1024, cross_attention_dim=1024, mult=1, num_tokens=1): - super().__init__() - from .attention import FeedForward - - self.num_tokens = num_tokens - self.cross_attention_dim = cross_attention_dim - self.ff = FeedForward(image_embed_dim, cross_attention_dim * num_tokens, mult=mult, activation_fn="gelu") - self.norm = nn.LayerNorm(cross_attention_dim) - - def forward(self, image_embeds: torch.Tensor): - x = self.ff(image_embeds) - x = x.reshape(-1, self.num_tokens, self.cross_attention_dim) - return self.norm(x) - - -class CombinedTimestepLabelEmbeddings(nn.Module): - def __init__(self, num_classes, embedding_dim, class_dropout_prob=0.1): - super().__init__() - - self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1) - self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) - self.class_embedder = LabelEmbedding(num_classes, embedding_dim, class_dropout_prob) - - def forward(self, timestep, class_labels, hidden_dtype=None): - timesteps_proj = self.time_proj(timestep) - timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) - - class_labels = self.class_embedder(class_labels) # (N, D) - - conditioning = timesteps_emb + class_labels # (N, D) - - return conditioning - - -class CombinedTimestepTextProjEmbeddings(nn.Module): - def __init__(self, embedding_dim, pooled_projection_dim): - super().__init__() - - self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) - self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) - self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu") - - def forward(self, timestep, pooled_projection): - 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 - - return conditioning - - -class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module): - def __init__(self, embedding_dim, pooled_projection_dim): - super().__init__() - - self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) - self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) - self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) - self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu") - - def forward(self, timestep, guidance, pooled_projection): - timesteps_proj = self.time_proj(timestep) - timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype)) # (N, D) - - guidance_proj = self.time_proj(guidance) - guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype)) # (N, D) - - time_guidance_emb = timesteps_emb + guidance_emb - - pooled_projections = self.text_embedder(pooled_projection) - conditioning = time_guidance_emb + pooled_projections - - return conditioning - - -class CogView3CombinedTimestepSizeEmbeddings(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, - original_size: torch.Tensor, - target_size: torch.Tensor, - crop_coords: torch.Tensor, - hidden_dtype: torch.dtype, - ) -> torch.Tensor: - timesteps_proj = self.time_proj(timestep) - - original_size_proj = self.condition_proj(original_size.flatten()).view(original_size.size(0), -1) - 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, 3 * condition_dim) - condition_proj = torch.cat([original_size_proj, 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 - return conditioning - - -class HunyuanDiTAttentionPool(nn.Module): - # Copied from https://github.com/Tencent/HunyuanDiT/blob/cb709308d92e6c7e8d59d0dff41b74d35088db6a/hydit/modules/poolers.py#L6 - - def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None): - super().__init__() - self.positional_embedding = nn.Parameter(torch.randn(spacial_dim + 1, embed_dim) / embed_dim**0.5) - self.k_proj = nn.Linear(embed_dim, embed_dim) - self.q_proj = nn.Linear(embed_dim, embed_dim) - self.v_proj = nn.Linear(embed_dim, embed_dim) - self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim) - self.num_heads = num_heads - - def forward(self, x): - x = x.permute(1, 0, 2) # NLC -> LNC - x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (L+1)NC - x = x + self.positional_embedding[:, None, :].to(x.dtype) # (L+1)NC - x, _ = F.multi_head_attention_forward( - query=x[:1], - key=x, - value=x, - embed_dim_to_check=x.shape[-1], - num_heads=self.num_heads, - q_proj_weight=self.q_proj.weight, - k_proj_weight=self.k_proj.weight, - v_proj_weight=self.v_proj.weight, - in_proj_weight=None, - in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]), - bias_k=None, - bias_v=None, - add_zero_attn=False, - dropout_p=0, - out_proj_weight=self.c_proj.weight, - out_proj_bias=self.c_proj.bias, - use_separate_proj_weight=True, - training=self.training, - need_weights=False, - ) - return x.squeeze(0) - - -class HunyuanCombinedTimestepTextSizeStyleEmbedding(nn.Module): - def __init__( - self, - embedding_dim, - pooled_projection_dim=1024, - seq_len=256, - cross_attention_dim=2048, - use_style_cond_and_image_meta_size=True, - ): - super().__init__() - - self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) - self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) - - self.size_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) - - self.pooler = HunyuanDiTAttentionPool( - seq_len, cross_attention_dim, num_heads=8, output_dim=pooled_projection_dim - ) - - # Here we use a default learned embedder layer for future extension. - self.use_style_cond_and_image_meta_size = use_style_cond_and_image_meta_size - if use_style_cond_and_image_meta_size: - self.style_embedder = nn.Embedding(1, embedding_dim) - extra_in_dim = 256 * 6 + embedding_dim + pooled_projection_dim - else: - extra_in_dim = pooled_projection_dim - - self.extra_embedder = PixArtAlphaTextProjection( - in_features=extra_in_dim, - hidden_size=embedding_dim * 4, - out_features=embedding_dim, - act_fn="silu_fp32", - ) - - def forward(self, timestep, encoder_hidden_states, image_meta_size, style, hidden_dtype=None): - timesteps_proj = self.time_proj(timestep) - timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, 256) - - # extra condition1: text - pooled_projections = self.pooler(encoder_hidden_states) # (N, 1024) - - if self.use_style_cond_and_image_meta_size: - # extra condition2: image meta size embedding - image_meta_size = self.size_proj(image_meta_size.view(-1)) - image_meta_size = image_meta_size.to(dtype=hidden_dtype) - image_meta_size = image_meta_size.view(-1, 6 * 256) # (N, 1536) - - # extra condition3: style embedding - style_embedding = self.style_embedder(style) # (N, embedding_dim) - - # Concatenate all extra vectors - extra_cond = torch.cat([pooled_projections, image_meta_size, style_embedding], dim=1) - else: - extra_cond = torch.cat([pooled_projections], dim=1) - - conditioning = timesteps_emb + self.extra_embedder(extra_cond) # [B, D] - - return conditioning - - -class LuminaCombinedTimestepCaptionEmbedding(nn.Module): - def __init__(self, hidden_size=4096, cross_attention_dim=2048, frequency_embedding_size=256): - super().__init__() - self.time_proj = Timesteps( - num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0.0 - ) - - self.timestep_embedder = TimestepEmbedding(in_channels=frequency_embedding_size, time_embed_dim=hidden_size) - - self.caption_embedder = nn.Sequential( - nn.LayerNorm(cross_attention_dim), - nn.Linear( - cross_attention_dim, - hidden_size, - bias=True, - ), - ) - - def forward(self, timestep, caption_feat, caption_mask): - # timestep embedding: - time_freq = self.time_proj(timestep) - time_embed = self.timestep_embedder(time_freq.to(dtype=caption_feat.dtype)) - - # caption condition embedding: - caption_mask_float = caption_mask.float().unsqueeze(-1) - caption_feats_pool = (caption_feat * caption_mask_float).sum(dim=1) / caption_mask_float.sum(dim=1) - caption_feats_pool = caption_feats_pool.to(caption_feat) - caption_embed = self.caption_embedder(caption_feats_pool) - - conditioning = time_embed + caption_embed - - return conditioning - - -class MochiCombinedTimestepCaptionEmbedding(nn.Module): - def __init__( - self, - embedding_dim: int, - pooled_projection_dim: int, - text_embed_dim: int, - time_embed_dim: int = 256, - num_attention_heads: int = 8, - ) -> None: - super().__init__() - - self.time_proj = Timesteps(num_channels=time_embed_dim, flip_sin_to_cos=True, downscale_freq_shift=0.0) - self.timestep_embedder = TimestepEmbedding(in_channels=time_embed_dim, time_embed_dim=embedding_dim) - self.pooler = MochiAttentionPool( - num_attention_heads=num_attention_heads, embed_dim=text_embed_dim, output_dim=embedding_dim - ) - self.caption_proj = nn.Linear(text_embed_dim, pooled_projection_dim) - - def forward( - self, - timestep: torch.LongTensor, - encoder_hidden_states: torch.Tensor, - encoder_attention_mask: torch.Tensor, - hidden_dtype: Optional[torch.dtype] = None, - ): - time_proj = self.time_proj(timestep) - time_emb = self.timestep_embedder(time_proj.to(dtype=hidden_dtype)) - - pooled_projections = self.pooler(encoder_hidden_states, encoder_attention_mask) - caption_proj = self.caption_proj(encoder_hidden_states) - - conditioning = time_emb + pooled_projections - return conditioning, caption_proj - - -class TextTimeEmbedding(nn.Module): - def __init__(self, encoder_dim: int, time_embed_dim: int, num_heads: int = 64): - super().__init__() - self.norm1 = nn.LayerNorm(encoder_dim) - self.pool = AttentionPooling(num_heads, encoder_dim) - self.proj = nn.Linear(encoder_dim, time_embed_dim) - self.norm2 = nn.LayerNorm(time_embed_dim) - - def forward(self, hidden_states): - hidden_states = self.norm1(hidden_states) - hidden_states = self.pool(hidden_states) - hidden_states = self.proj(hidden_states) - hidden_states = self.norm2(hidden_states) - return hidden_states - - -class TextImageTimeEmbedding(nn.Module): - def __init__(self, text_embed_dim: int = 768, image_embed_dim: int = 768, time_embed_dim: int = 1536): - super().__init__() - self.text_proj = nn.Linear(text_embed_dim, time_embed_dim) - self.text_norm = nn.LayerNorm(time_embed_dim) - self.image_proj = nn.Linear(image_embed_dim, time_embed_dim) - - def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor): - # text - time_text_embeds = self.text_proj(text_embeds) - time_text_embeds = self.text_norm(time_text_embeds) - - # image - time_image_embeds = self.image_proj(image_embeds) - - return time_image_embeds + time_text_embeds - - -class ImageTimeEmbedding(nn.Module): - def __init__(self, image_embed_dim: int = 768, time_embed_dim: int = 1536): - super().__init__() - self.image_proj = nn.Linear(image_embed_dim, time_embed_dim) - self.image_norm = nn.LayerNorm(time_embed_dim) - - def forward(self, image_embeds: torch.Tensor): - # image - time_image_embeds = self.image_proj(image_embeds) - time_image_embeds = self.image_norm(time_image_embeds) - return time_image_embeds - - -class ImageHintTimeEmbedding(nn.Module): - def __init__(self, image_embed_dim: int = 768, time_embed_dim: int = 1536): - super().__init__() - self.image_proj = nn.Linear(image_embed_dim, time_embed_dim) - self.image_norm = nn.LayerNorm(time_embed_dim) - self.input_hint_block = nn.Sequential( - nn.Conv2d(3, 16, 3, padding=1), - nn.SiLU(), - nn.Conv2d(16, 16, 3, padding=1), - nn.SiLU(), - nn.Conv2d(16, 32, 3, padding=1, stride=2), - nn.SiLU(), - nn.Conv2d(32, 32, 3, padding=1), - nn.SiLU(), - nn.Conv2d(32, 96, 3, padding=1, stride=2), - nn.SiLU(), - nn.Conv2d(96, 96, 3, padding=1), - nn.SiLU(), - nn.Conv2d(96, 256, 3, padding=1, stride=2), - nn.SiLU(), - nn.Conv2d(256, 4, 3, padding=1), - ) - - def forward(self, image_embeds: torch.Tensor, hint: torch.Tensor): - # image - time_image_embeds = self.image_proj(image_embeds) - time_image_embeds = self.image_norm(time_image_embeds) - hint = self.input_hint_block(hint) - return time_image_embeds, hint - - -class AttentionPooling(nn.Module): - # Copied from https://github.com/deep-floyd/IF/blob/2f91391f27dd3c468bf174be5805b4cc92980c0b/deepfloyd_if/model/nn.py#L54 - - def __init__(self, num_heads, embed_dim, dtype=None): - super().__init__() - self.dtype = dtype - self.positional_embedding = nn.Parameter(torch.randn(1, embed_dim) / embed_dim**0.5) - self.k_proj = nn.Linear(embed_dim, embed_dim, dtype=self.dtype) - self.q_proj = nn.Linear(embed_dim, embed_dim, dtype=self.dtype) - self.v_proj = nn.Linear(embed_dim, embed_dim, dtype=self.dtype) - self.num_heads = num_heads - self.dim_per_head = embed_dim // self.num_heads - - def forward(self, x): - bs, length, width = x.size() - - def shape(x): - # (bs, length, width) --> (bs, length, n_heads, dim_per_head) - x = x.view(bs, -1, self.num_heads, self.dim_per_head) - # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head) - x = x.transpose(1, 2) - # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head) - x = x.reshape(bs * self.num_heads, -1, self.dim_per_head) - # (bs*n_heads, length, dim_per_head) --> (bs*n_heads, dim_per_head, length) - x = x.transpose(1, 2) - return x - - class_token = x.mean(dim=1, keepdim=True) + self.positional_embedding.to(x.dtype) - x = torch.cat([class_token, x], dim=1) # (bs, length+1, width) - - # (bs*n_heads, class_token_length, dim_per_head) - q = shape(self.q_proj(class_token)) - # (bs*n_heads, length+class_token_length, dim_per_head) - k = shape(self.k_proj(x)) - v = shape(self.v_proj(x)) - - # (bs*n_heads, class_token_length, length+class_token_length): - scale = 1 / math.sqrt(math.sqrt(self.dim_per_head)) - weight = torch.einsum("bct,bcs->bts", q * scale, k * scale) # More stable with f16 than dividing afterwards - weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) - - # (bs*n_heads, dim_per_head, class_token_length) - a = torch.einsum("bts,bcs->bct", weight, v) - - # (bs, length+1, width) - a = a.reshape(bs, -1, 1).transpose(1, 2) - - return a[:, 0, :] # cls_token - - -class MochiAttentionPool(nn.Module): - def __init__( - self, - num_attention_heads: int, - embed_dim: int, - output_dim: Optional[int] = None, - ) -> None: - super().__init__() - - self.output_dim = output_dim or embed_dim - self.num_attention_heads = num_attention_heads - - self.to_kv = nn.Linear(embed_dim, 2 * embed_dim) - self.to_q = nn.Linear(embed_dim, embed_dim) - self.to_out = nn.Linear(embed_dim, self.output_dim) - - @staticmethod - def pool_tokens(x: torch.Tensor, mask: torch.Tensor, *, keepdim=False) -> torch.Tensor: - """ - Pool tokens in x using mask. - - NOTE: We assume x does not require gradients. - - Args: - x: (B, L, D) tensor of tokens. - mask: (B, L) boolean tensor indicating which tokens are not padding. - - Returns: - pooled: (B, D) tensor of pooled tokens. - """ - assert x.size(1) == mask.size(1) # Expected mask to have same length as tokens. - assert x.size(0) == mask.size(0) # Expected mask to have same batch size as tokens. - mask = mask[:, :, None].to(dtype=x.dtype) - mask = mask / mask.sum(dim=1, keepdim=True).clamp(min=1) - pooled = (x * mask).sum(dim=1, keepdim=keepdim) - return pooled - - def forward(self, x: torch.Tensor, mask: torch.BoolTensor) -> torch.Tensor: - r""" - Args: - x (`torch.Tensor`): - Tensor of shape `(B, S, D)` of input tokens. - mask (`torch.Tensor`): - Boolean ensor of shape `(B, S)` indicating which tokens are not padding. - - Returns: - `torch.Tensor`: - `(B, D)` tensor of pooled tokens. - """ - D = x.size(2) - - # Construct attention mask, shape: (B, 1, num_queries=1, num_keys=1+L). - attn_mask = mask[:, None, None, :].bool() # (B, 1, 1, L). - attn_mask = F.pad(attn_mask, (1, 0), value=True) # (B, 1, 1, 1+L). - - # Average non-padding token features. These will be used as the query. - x_pool = self.pool_tokens(x, mask, keepdim=True) # (B, 1, D) - - # Concat pooled features to input sequence. - x = torch.cat([x_pool, x], dim=1) # (B, L+1, D) - - # Compute queries, keys, values. Only the mean token is used to create a query. - kv = self.to_kv(x) # (B, L+1, 2 * D) - q = self.to_q(x[:, 0]) # (B, D) - - # Extract heads. - head_dim = D // self.num_attention_heads - kv = kv.unflatten(2, (2, self.num_attention_heads, head_dim)) # (B, 1+L, 2, H, head_dim) - kv = kv.transpose(1, 3) # (B, H, 2, 1+L, head_dim) - k, v = kv.unbind(2) # (B, H, 1+L, head_dim) - q = q.unflatten(1, (self.num_attention_heads, head_dim)) # (B, H, head_dim) - q = q.unsqueeze(2) # (B, H, 1, head_dim) - - # Compute attention. - x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=0.0) # (B, H, 1, head_dim) - - # Concatenate heads and run output. - x = x.squeeze(2).flatten(1, 2) # (B, D = H * head_dim) - x = self.to_out(x) - return x - - -def get_fourier_embeds_from_boundingbox(embed_dim, box): - """ - Args: - embed_dim: int - box: a 3-D tensor [B x N x 4] representing the bounding boxes for GLIGEN pipeline - Returns: - [B x N x embed_dim] tensor of positional embeddings - """ - - batch_size, num_boxes = box.shape[:2] - - emb = 100 ** (torch.arange(embed_dim) / embed_dim) - emb = emb[None, None, None].to(device=box.device, dtype=box.dtype) - emb = emb * box.unsqueeze(-1) - - emb = torch.stack((emb.sin(), emb.cos()), dim=-1) - emb = emb.permute(0, 1, 3, 4, 2).reshape(batch_size, num_boxes, embed_dim * 2 * 4) - - return emb - - -class GLIGENTextBoundingboxProjection(nn.Module): - def __init__(self, positive_len, out_dim, feature_type="text-only", fourier_freqs=8): - super().__init__() - self.positive_len = positive_len - self.out_dim = out_dim - - self.fourier_embedder_dim = fourier_freqs - self.position_dim = fourier_freqs * 2 * 4 # 2: sin/cos, 4: xyxy - - if isinstance(out_dim, tuple): - out_dim = out_dim[0] - - if feature_type == "text-only": - self.linears = nn.Sequential( - nn.Linear(self.positive_len + self.position_dim, 512), - nn.SiLU(), - nn.Linear(512, 512), - nn.SiLU(), - nn.Linear(512, out_dim), - ) - self.null_positive_feature = torch.nn.Parameter(torch.zeros([self.positive_len])) - - elif feature_type == "text-image": - self.linears_text = nn.Sequential( - nn.Linear(self.positive_len + self.position_dim, 512), - nn.SiLU(), - nn.Linear(512, 512), - nn.SiLU(), - nn.Linear(512, out_dim), - ) - self.linears_image = nn.Sequential( - nn.Linear(self.positive_len + self.position_dim, 512), - nn.SiLU(), - nn.Linear(512, 512), - nn.SiLU(), - nn.Linear(512, out_dim), - ) - self.null_text_feature = torch.nn.Parameter(torch.zeros([self.positive_len])) - self.null_image_feature = torch.nn.Parameter(torch.zeros([self.positive_len])) - - self.null_position_feature = torch.nn.Parameter(torch.zeros([self.position_dim])) - - def forward( - self, - boxes, - masks, - positive_embeddings=None, - phrases_masks=None, - image_masks=None, - phrases_embeddings=None, - image_embeddings=None, - ): - masks = masks.unsqueeze(-1) - - # embedding position (it may includes padding as placeholder) - xyxy_embedding = get_fourier_embeds_from_boundingbox(self.fourier_embedder_dim, boxes) # B*N*4 -> B*N*C - - # learnable null embedding - xyxy_null = self.null_position_feature.view(1, 1, -1) - - # replace padding with learnable null embedding - xyxy_embedding = xyxy_embedding * masks + (1 - masks) * xyxy_null - - # positionet with text only information - if positive_embeddings is not None: - # learnable null embedding - positive_null = self.null_positive_feature.view(1, 1, -1) - - # replace padding with learnable null embedding - positive_embeddings = positive_embeddings * masks + (1 - masks) * positive_null - - objs = self.linears(torch.cat([positive_embeddings, xyxy_embedding], dim=-1)) - - # positionet with text and image information - else: - phrases_masks = phrases_masks.unsqueeze(-1) - image_masks = image_masks.unsqueeze(-1) - - # learnable null embedding - text_null = self.null_text_feature.view(1, 1, -1) - image_null = self.null_image_feature.view(1, 1, -1) - - # replace padding with learnable null embedding - phrases_embeddings = phrases_embeddings * phrases_masks + (1 - phrases_masks) * text_null - image_embeddings = image_embeddings * image_masks + (1 - image_masks) * image_null - - objs_text = self.linears_text(torch.cat([phrases_embeddings, xyxy_embedding], dim=-1)) - objs_image = self.linears_image(torch.cat([image_embeddings, xyxy_embedding], dim=-1)) - objs = torch.cat([objs_text, objs_image], dim=1) - - return objs - - -class PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module): - """ - For PixArt-Alpha. - - Reference: - https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L164C9-L168C29 - """ - - def __init__(self, embedding_dim, size_emb_dim, use_additional_conditions: bool = False): - super().__init__() - - self.outdim = size_emb_dim - self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) - self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) - - self.use_additional_conditions = use_additional_conditions - if use_additional_conditions: - self.additional_condition_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) - self.resolution_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=size_emb_dim) - self.aspect_ratio_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=size_emb_dim) - - def forward(self, timestep, resolution, aspect_ratio, batch_size, hidden_dtype): - timesteps_proj = self.time_proj(timestep) - timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) - - if self.use_additional_conditions: - resolution_emb = self.additional_condition_proj(resolution.flatten()).to(hidden_dtype) - resolution_emb = self.resolution_embedder(resolution_emb).reshape(batch_size, -1) - aspect_ratio_emb = self.additional_condition_proj(aspect_ratio.flatten()).to(hidden_dtype) - aspect_ratio_emb = self.aspect_ratio_embedder(aspect_ratio_emb).reshape(batch_size, -1) - conditioning = timesteps_emb + torch.cat([resolution_emb, aspect_ratio_emb], dim=1) - else: - conditioning = timesteps_emb - - return conditioning - - -class PixArtAlphaTextProjection(nn.Module): - """ - Projects caption embeddings. Also handles dropout for classifier-free guidance. - - Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py - """ - - def __init__(self, in_features, hidden_size, out_features=None, act_fn="gelu_tanh"): - super().__init__() - if out_features is None: - out_features = hidden_size - self.linear_1 = nn.Linear(in_features=in_features, out_features=hidden_size, bias=True) - if act_fn == "gelu_tanh": - self.act_1 = nn.GELU(approximate="tanh") - elif act_fn == "silu": - self.act_1 = nn.SiLU() - elif act_fn == "silu_fp32": - self.act_1 = FP32SiLU() - else: - raise ValueError(f"Unknown activation function: {act_fn}") - self.linear_2 = nn.Linear(in_features=hidden_size, out_features=out_features, bias=True) - - def forward(self, caption): - hidden_states = self.linear_1(caption) - hidden_states = self.act_1(hidden_states) - hidden_states = self.linear_2(hidden_states) - return hidden_states - - -class IPAdapterPlusImageProjectionBlock(nn.Module): - def __init__( - self, - embed_dims: int = 768, - dim_head: int = 64, - heads: int = 16, - ffn_ratio: float = 4, - ) -> None: - super().__init__() - from .attention import FeedForward - - self.ln0 = nn.LayerNorm(embed_dims) - self.ln1 = nn.LayerNorm(embed_dims) - self.attn = Attention( - query_dim=embed_dims, - dim_head=dim_head, - heads=heads, - out_bias=False, - ) - self.ff = nn.Sequential( - nn.LayerNorm(embed_dims), - FeedForward(embed_dims, embed_dims, activation_fn="gelu", mult=ffn_ratio, bias=False), - ) - - def forward(self, x, latents, residual): - encoder_hidden_states = self.ln0(x) - latents = self.ln1(latents) - encoder_hidden_states = torch.cat([encoder_hidden_states, latents], dim=-2) - latents = self.attn(latents, encoder_hidden_states) + residual - latents = self.ff(latents) + latents - return latents - - -class IPAdapterPlusImageProjection(nn.Module): - """Resampler of IP-Adapter Plus. - - Args: - embed_dims (int): The feature dimension. Defaults to 768. output_dims (int): The number of output channels, - that is the same - number of the channels in the `unet.config.cross_attention_dim`. Defaults to 1024. - hidden_dims (int): - The number of hidden channels. Defaults to 1280. depth (int): The number of blocks. Defaults - to 8. dim_head (int): The number of head channels. Defaults to 64. heads (int): Parallel attention heads. - Defaults to 16. num_queries (int): - The number of queries. Defaults to 8. ffn_ratio (float): The expansion ratio - of feedforward network hidden - layer channels. Defaults to 4. - """ - - def __init__( - self, - embed_dims: int = 768, - output_dims: int = 1024, - hidden_dims: int = 1280, - depth: int = 4, - dim_head: int = 64, - heads: int = 16, - num_queries: int = 8, - ffn_ratio: float = 4, - ) -> None: - super().__init__() - self.latents = nn.Parameter(torch.randn(1, num_queries, hidden_dims) / hidden_dims**0.5) - - self.proj_in = nn.Linear(embed_dims, hidden_dims) - - self.proj_out = nn.Linear(hidden_dims, output_dims) - self.norm_out = nn.LayerNorm(output_dims) - - self.layers = nn.ModuleList( - [IPAdapterPlusImageProjectionBlock(hidden_dims, dim_head, heads, ffn_ratio) for _ in range(depth)] - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass. - - Args: - x (torch.Tensor): Input Tensor. - Returns: - torch.Tensor: Output Tensor. - """ - latents = self.latents.repeat(x.size(0), 1, 1) - - x = self.proj_in(x) - - for block in self.layers: - residual = latents - latents = block(x, latents, residual) - - latents = self.proj_out(latents) - return self.norm_out(latents) - - -class IPAdapterFaceIDPlusImageProjection(nn.Module): - """FacePerceiverResampler of IP-Adapter Plus. - - Args: - embed_dims (int): The feature dimension. Defaults to 768. output_dims (int): The number of output channels, - that is the same - number of the channels in the `unet.config.cross_attention_dim`. Defaults to 1024. - hidden_dims (int): - The number of hidden channels. Defaults to 1280. depth (int): The number of blocks. Defaults - to 8. dim_head (int): The number of head channels. Defaults to 64. heads (int): Parallel attention heads. - Defaults to 16. num_tokens (int): Number of tokens num_queries (int): The number of queries. Defaults to 8. - ffn_ratio (float): The expansion ratio of feedforward network hidden - layer channels. Defaults to 4. - ffproj_ratio (float): The expansion ratio of feedforward network hidden - layer channels (for ID embeddings). Defaults to 4. - """ - - def __init__( - self, - embed_dims: int = 768, - output_dims: int = 768, - hidden_dims: int = 1280, - id_embeddings_dim: int = 512, - depth: int = 4, - dim_head: int = 64, - heads: int = 16, - num_tokens: int = 4, - num_queries: int = 8, - ffn_ratio: float = 4, - ffproj_ratio: int = 2, - ) -> None: - super().__init__() - from .attention import FeedForward - - self.num_tokens = num_tokens - self.embed_dim = embed_dims - self.clip_embeds = None - self.shortcut = False - self.shortcut_scale = 1.0 - - self.proj = FeedForward(id_embeddings_dim, embed_dims * num_tokens, activation_fn="gelu", mult=ffproj_ratio) - self.norm = nn.LayerNorm(embed_dims) - - self.proj_in = nn.Linear(hidden_dims, embed_dims) - - self.proj_out = nn.Linear(embed_dims, output_dims) - self.norm_out = nn.LayerNorm(output_dims) - - self.layers = nn.ModuleList( - [IPAdapterPlusImageProjectionBlock(embed_dims, dim_head, heads, ffn_ratio) for _ in range(depth)] - ) - - def forward(self, id_embeds: torch.Tensor) -> torch.Tensor: - """Forward pass. - - Args: - id_embeds (torch.Tensor): Input Tensor (ID embeds). - Returns: - torch.Tensor: Output Tensor. - """ - id_embeds = id_embeds.to(self.clip_embeds.dtype) - id_embeds = self.proj(id_embeds) - id_embeds = id_embeds.reshape(-1, self.num_tokens, self.embed_dim) - id_embeds = self.norm(id_embeds) - latents = id_embeds - - clip_embeds = self.proj_in(self.clip_embeds) - x = clip_embeds.reshape(-1, clip_embeds.shape[2], clip_embeds.shape[3]) - - for block in self.layers: - residual = latents - latents = block(x, latents, residual) - - latents = self.proj_out(latents) - out = self.norm_out(latents) - if self.shortcut: - out = id_embeds + self.shortcut_scale * out - return out - - -class IPAdapterTimeImageProjectionBlock(nn.Module): - """Block for IPAdapterTimeImageProjection. - - Args: - hidden_dim (`int`, defaults to 1280): - The number of hidden channels. - dim_head (`int`, defaults to 64): - The number of head channels. - heads (`int`, defaults to 20): - Parallel attention heads. - ffn_ratio (`int`, defaults to 4): - The expansion ratio of feedforward network hidden layer channels. - """ - - def __init__( - self, - hidden_dim: int = 1280, - dim_head: int = 64, - heads: int = 20, - ffn_ratio: int = 4, - ) -> None: - super().__init__() - from .attention import FeedForward - - self.ln0 = nn.LayerNorm(hidden_dim) - self.ln1 = nn.LayerNorm(hidden_dim) - self.attn = Attention( - query_dim=hidden_dim, - cross_attention_dim=hidden_dim, - dim_head=dim_head, - heads=heads, - bias=False, - out_bias=False, - ) - self.ff = FeedForward(hidden_dim, hidden_dim, activation_fn="gelu", mult=ffn_ratio, bias=False) - - # AdaLayerNorm - self.adaln_silu = nn.SiLU() - self.adaln_proj = nn.Linear(hidden_dim, 4 * hidden_dim) - self.adaln_norm = nn.LayerNorm(hidden_dim) - - # Set attention scale and fuse KV - self.attn.scale = 1 / math.sqrt(math.sqrt(dim_head)) - self.attn.fuse_projections() - self.attn.to_k = None - self.attn.to_v = None - - def forward(self, x: torch.Tensor, latents: torch.Tensor, timestep_emb: torch.Tensor) -> torch.Tensor: - """Forward pass. - - Args: - x (`torch.Tensor`): - Image features. - latents (`torch.Tensor`): - Latent features. - timestep_emb (`torch.Tensor`): - Timestep embedding. - - Returns: - `torch.Tensor`: Output latent features. - """ - - # Shift and scale for AdaLayerNorm - emb = self.adaln_proj(self.adaln_silu(timestep_emb)) - shift_msa, scale_msa, shift_mlp, scale_mlp = emb.chunk(4, dim=1) - - # Fused Attention - residual = latents - x = self.ln0(x) - latents = self.ln1(latents) * (1 + scale_msa[:, None]) + shift_msa[:, None] - - batch_size = latents.shape[0] - - query = self.attn.to_q(latents) - kv_input = torch.cat((x, latents), dim=-2) - key, value = self.attn.to_kv(kv_input).chunk(2, dim=-1) - - inner_dim = key.shape[-1] - head_dim = inner_dim // self.attn.heads - - query = query.view(batch_size, -1, self.attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, self.attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, self.attn.heads, head_dim).transpose(1, 2) - - weight = (query * self.attn.scale) @ (key * self.attn.scale).transpose(-2, -1) - weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) - latents = weight @ value - - latents = latents.transpose(1, 2).reshape(batch_size, -1, self.attn.heads * head_dim) - latents = self.attn.to_out[0](latents) - latents = self.attn.to_out[1](latents) - latents = latents + residual - - ## FeedForward - residual = latents - latents = self.adaln_norm(latents) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - return self.ff(latents) + residual - - -# Modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py -class IPAdapterTimeImageProjection(nn.Module): - """Resampler of SD3 IP-Adapter with timestep embedding. - - Args: - embed_dim (`int`, defaults to 1152): - The feature dimension. - output_dim (`int`, defaults to 2432): - The number of output channels. - hidden_dim (`int`, defaults to 1280): - The number of hidden channels. - depth (`int`, defaults to 4): - The number of blocks. - dim_head (`int`, defaults to 64): - The number of head channels. - heads (`int`, defaults to 20): - Parallel attention heads. - num_queries (`int`, defaults to 64): - The number of queries. - ffn_ratio (`int`, defaults to 4): - The expansion ratio of feedforward network hidden layer channels. - timestep_in_dim (`int`, defaults to 320): - The number of input channels for timestep embedding. - timestep_flip_sin_to_cos (`bool`, defaults to True): - Flip the timestep embedding order to `cos, sin` (if True) or `sin, cos` (if False). - timestep_freq_shift (`int`, defaults to 0): - Controls the timestep delta between frequencies between dimensions. - """ - - def __init__( - self, - embed_dim: int = 1152, - output_dim: int = 2432, - hidden_dim: int = 1280, - depth: int = 4, - dim_head: int = 64, - heads: int = 20, - num_queries: int = 64, - ffn_ratio: int = 4, - timestep_in_dim: int = 320, - timestep_flip_sin_to_cos: bool = True, - timestep_freq_shift: int = 0, - ) -> None: - super().__init__() - self.latents = nn.Parameter(torch.randn(1, num_queries, hidden_dim) / hidden_dim**0.5) - self.proj_in = nn.Linear(embed_dim, hidden_dim) - self.proj_out = nn.Linear(hidden_dim, output_dim) - self.norm_out = nn.LayerNorm(output_dim) - self.layers = nn.ModuleList( - [IPAdapterTimeImageProjectionBlock(hidden_dim, dim_head, heads, ffn_ratio) for _ in range(depth)] - ) - 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]: - """Forward pass. - - Args: - x (`torch.Tensor`): - Image features. - timestep (`torch.Tensor`): - Timestep in denoising process. - Returns: - `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) - - latents = self.latents.repeat(x.size(0), 1, 1) - - x = self.proj_in(x) - x = x + timestep_emb[:, None] - - for block in self.layers: - latents = block(x, latents, timestep_emb) - - latents = self.proj_out(latents) - latents = self.norm_out(latents) - - return latents, timestep_emb - - -class MultiIPAdapterImageProjection(nn.Module): - def __init__(self, IPAdapterImageProjectionLayers: Union[List[nn.Module], Tuple[nn.Module]]): - super().__init__() - self.image_projection_layers = nn.ModuleList(IPAdapterImageProjectionLayers) - - @property - 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]): - projected_image_embeds = [] - - # currently, we accept `image_embeds` as - # 1. a tensor (deprecated) with shape [batch_size, embed_dim] or [batch_size, sequence_length, embed_dim] - # 2. list of `n` tensors where `n` is number of ip-adapters, each tensor can hae shape [batch_size, num_images, embed_dim] or [batch_size, num_images, sequence_length, embed_dim] - if not isinstance(image_embeds, list): - deprecation_message = ( - "You have passed a tensor as `image_embeds`.This is deprecated and will be removed in a future release." - " Please make sure to update your script to pass `image_embeds` as a list of tensors to suppress this warning." - ) - deprecate("image_embeds not a list", "1.0.0", deprecation_message, standard_warn=False) - image_embeds = [image_embeds.unsqueeze(1)] - - if len(image_embeds) != len(self.image_projection_layers): - raise ValueError( - f"image_embeds must have the same length as image_projection_layers, got {len(image_embeds)} and {len(self.image_projection_layers)}" - ) - - for image_embed, image_projection_layer in zip(image_embeds, self.image_projection_layers): - batch_size, num_images = image_embed.shape[0], image_embed.shape[1] - image_embed = image_embed.reshape((batch_size * num_images,) + image_embed.shape[2:]) - image_embed = image_projection_layer(image_embed) - image_embed = image_embed.reshape((batch_size, num_images) + image_embed.shape[1:]) - - projected_image_embeds.append(image_embed) - - return projected_image_embeds diff --git a/venv/lib/python3.10/site-packages/diffusers/models/embeddings_flax.py b/venv/lib/python3.10/site-packages/diffusers/models/embeddings_flax.py deleted file mode 100644 index 1e7e84edeaeb5753aece1f67443a02a20a8c87ef..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/embeddings_flax.py +++ /dev/null @@ -1,111 +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 math - -import flax.linen as nn -import jax.numpy as jnp - - -def get_sinusoidal_embeddings( - timesteps: jnp.ndarray, - embedding_dim: int, - freq_shift: float = 1, - min_timescale: float = 1, - max_timescale: float = 1.0e4, - flip_sin_to_cos: bool = False, - scale: float = 1.0, -) -> jnp.ndarray: - """Returns the positional encoding (same as Tensor2Tensor). - - Args: - timesteps (`jnp.ndarray` of shape `(N,)`): - A 1-D array of N indices, one per batch element. These may be fractional. - embedding_dim (`int`): - The number of output channels. - freq_shift (`float`, *optional*, defaults to `1`): - Shift applied to the frequency scaling of the embeddings. - min_timescale (`float`, *optional*, defaults to `1`): - The smallest time unit used in the sinusoidal calculation (should probably be 0.0). - max_timescale (`float`, *optional*, defaults to `1.0e4`): - The largest time unit used in the sinusoidal calculation. - flip_sin_to_cos (`bool`, *optional*, defaults to `False`): - Whether to flip the order of sinusoidal components to cosine first. - scale (`float`, *optional*, defaults to `1.0`): - A scaling factor applied to the positional embeddings. - - Returns: - a Tensor of timing signals [N, num_channels] - """ - assert timesteps.ndim == 1, "Timesteps should be a 1d-array" - assert embedding_dim % 2 == 0, f"Embedding dimension {embedding_dim} should be even" - num_timescales = float(embedding_dim // 2) - log_timescale_increment = math.log(max_timescale / min_timescale) / (num_timescales - freq_shift) - inv_timescales = min_timescale * jnp.exp(jnp.arange(num_timescales, dtype=jnp.float32) * -log_timescale_increment) - emb = jnp.expand_dims(timesteps, 1) * jnp.expand_dims(inv_timescales, 0) - - # scale embeddings - scaled_time = scale * emb - - if flip_sin_to_cos: - signal = jnp.concatenate([jnp.cos(scaled_time), jnp.sin(scaled_time)], axis=1) - else: - signal = jnp.concatenate([jnp.sin(scaled_time), jnp.cos(scaled_time)], axis=1) - signal = jnp.reshape(signal, [jnp.shape(timesteps)[0], embedding_dim]) - return signal - - -class FlaxTimestepEmbedding(nn.Module): - r""" - Time step Embedding Module. Learns embeddings for input time steps. - - Args: - time_embed_dim (`int`, *optional*, defaults to `32`): - Time step embedding dimension. - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - The data type for the embedding parameters. - """ - - time_embed_dim: int = 32 - dtype: jnp.dtype = jnp.float32 - - @nn.compact - def __call__(self, temb): - temb = nn.Dense(self.time_embed_dim, dtype=self.dtype, name="linear_1")(temb) - temb = nn.silu(temb) - temb = nn.Dense(self.time_embed_dim, dtype=self.dtype, name="linear_2")(temb) - return temb - - -class FlaxTimesteps(nn.Module): - r""" - Wrapper Module for sinusoidal Time step Embeddings as described in https://huggingface.co/papers/2006.11239 - - Args: - dim (`int`, *optional*, defaults to `32`): - Time step embedding dimension. - flip_sin_to_cos (`bool`, *optional*, defaults to `False`): - Whether to flip the sinusoidal function from sine to cosine. - freq_shift (`float`, *optional*, defaults to `1`): - Frequency shift applied to the sinusoidal embeddings. - """ - - dim: int = 32 - flip_sin_to_cos: bool = False - freq_shift: float = 1 - - @nn.compact - def __call__(self, timesteps): - return get_sinusoidal_embeddings( - timesteps, embedding_dim=self.dim, flip_sin_to_cos=self.flip_sin_to_cos, freq_shift=self.freq_shift - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/lora.py b/venv/lib/python3.10/site-packages/diffusers/models/lora.py deleted file mode 100644 index 85d61d6d7cdf5a8cc2cde89852491d6beb95ba53..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/lora.py +++ /dev/null @@ -1,457 +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. - - -# IMPORTANT: # -################################################################### -# ----------------------------------------------------------------# -# This file is deprecated and will be removed soon # -# (as soon as PEFT will become a required dependency for LoRA) # -# ----------------------------------------------------------------# -################################################################### - -from typing import Optional, Tuple, Union - -import torch -import torch.nn.functional as F -from torch import nn - -from ..utils import deprecate, logging -from ..utils.import_utils import is_transformers_available - - -if is_transformers_available(): - from transformers import CLIPTextModel, CLIPTextModelWithProjection - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -def text_encoder_attn_modules(text_encoder: nn.Module): - attn_modules = [] - - if isinstance(text_encoder, (CLIPTextModel, CLIPTextModelWithProjection)): - for i, layer in enumerate(text_encoder.text_model.encoder.layers): - name = f"text_model.encoder.layers.{i}.self_attn" - mod = layer.self_attn - attn_modules.append((name, mod)) - else: - raise ValueError(f"do not know how to get attention modules for: {text_encoder.__class__.__name__}") - - return attn_modules - - -def text_encoder_mlp_modules(text_encoder: nn.Module): - mlp_modules = [] - - if isinstance(text_encoder, (CLIPTextModel, CLIPTextModelWithProjection)): - for i, layer in enumerate(text_encoder.text_model.encoder.layers): - mlp_mod = layer.mlp - name = f"text_model.encoder.layers.{i}.mlp" - mlp_modules.append((name, mlp_mod)) - else: - raise ValueError(f"do not know how to get mlp modules for: {text_encoder.__class__.__name__}") - - return mlp_modules - - -def adjust_lora_scale_text_encoder(text_encoder, lora_scale: float = 1.0): - for _, attn_module in text_encoder_attn_modules(text_encoder): - if isinstance(attn_module.q_proj, PatchedLoraProjection): - attn_module.q_proj.lora_scale = lora_scale - attn_module.k_proj.lora_scale = lora_scale - attn_module.v_proj.lora_scale = lora_scale - attn_module.out_proj.lora_scale = lora_scale - - for _, mlp_module in text_encoder_mlp_modules(text_encoder): - if isinstance(mlp_module.fc1, PatchedLoraProjection): - mlp_module.fc1.lora_scale = lora_scale - mlp_module.fc2.lora_scale = lora_scale - - -class PatchedLoraProjection(torch.nn.Module): - def __init__(self, regular_linear_layer, lora_scale=1, network_alpha=None, rank=4, dtype=None): - deprecation_message = "Use of `PatchedLoraProjection` is deprecated. Please switch to PEFT backend by installing PEFT: `pip install peft`." - deprecate("PatchedLoraProjection", "1.0.0", deprecation_message) - - super().__init__() - from ..models.lora import LoRALinearLayer - - self.regular_linear_layer = regular_linear_layer - - device = self.regular_linear_layer.weight.device - - if dtype is None: - dtype = self.regular_linear_layer.weight.dtype - - self.lora_linear_layer = LoRALinearLayer( - self.regular_linear_layer.in_features, - self.regular_linear_layer.out_features, - network_alpha=network_alpha, - device=device, - dtype=dtype, - rank=rank, - ) - - self.lora_scale = lora_scale - - # overwrite PyTorch's `state_dict` to be sure that only the 'regular_linear_layer' weights are saved - # when saving the whole text encoder model and when LoRA is unloaded or fused - def state_dict(self, *args, destination=None, prefix="", keep_vars=False): - if self.lora_linear_layer is None: - return self.regular_linear_layer.state_dict( - *args, destination=destination, prefix=prefix, keep_vars=keep_vars - ) - - return super().state_dict(*args, destination=destination, prefix=prefix, keep_vars=keep_vars) - - def _fuse_lora(self, lora_scale=1.0, safe_fusing=False): - if self.lora_linear_layer is None: - return - - dtype, device = self.regular_linear_layer.weight.data.dtype, self.regular_linear_layer.weight.data.device - - w_orig = self.regular_linear_layer.weight.data.float() - w_up = self.lora_linear_layer.up.weight.data.float() - w_down = self.lora_linear_layer.down.weight.data.float() - - if self.lora_linear_layer.network_alpha is not None: - w_up = w_up * self.lora_linear_layer.network_alpha / self.lora_linear_layer.rank - - fused_weight = w_orig + (lora_scale * torch.bmm(w_up[None, :], w_down[None, :])[0]) - - if safe_fusing and torch.isnan(fused_weight).any().item(): - raise ValueError( - "This LoRA weight seems to be broken. " - f"Encountered NaN values when trying to fuse LoRA weights for {self}." - "LoRA weights will not be fused." - ) - - self.regular_linear_layer.weight.data = fused_weight.to(device=device, dtype=dtype) - - # we can drop the lora layer now - self.lora_linear_layer = None - - # offload the up and down matrices to CPU to not blow the memory - self.w_up = w_up.cpu() - self.w_down = w_down.cpu() - self.lora_scale = lora_scale - - def _unfuse_lora(self): - if not (getattr(self, "w_up", None) is not None and getattr(self, "w_down", None) is not None): - return - - fused_weight = self.regular_linear_layer.weight.data - dtype, device = fused_weight.dtype, fused_weight.device - - w_up = self.w_up.to(device=device).float() - w_down = self.w_down.to(device).float() - - unfused_weight = fused_weight.float() - (self.lora_scale * torch.bmm(w_up[None, :], w_down[None, :])[0]) - self.regular_linear_layer.weight.data = unfused_weight.to(device=device, dtype=dtype) - - self.w_up = None - self.w_down = None - - def forward(self, input): - if self.lora_scale is None: - self.lora_scale = 1.0 - if self.lora_linear_layer is None: - return self.regular_linear_layer(input) - return self.regular_linear_layer(input) + (self.lora_scale * self.lora_linear_layer(input)) - - -class LoRALinearLayer(nn.Module): - r""" - A linear layer that is used with LoRA. - - Parameters: - in_features (`int`): - Number of input features. - out_features (`int`): - Number of output features. - rank (`int`, `optional`, defaults to 4): - The rank of the LoRA layer. - network_alpha (`float`, `optional`, defaults to `None`): - 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. See - https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning - device (`torch.device`, `optional`, defaults to `None`): - The device to use for the layer's weights. - dtype (`torch.dtype`, `optional`, defaults to `None`): - The dtype to use for the layer's weights. - """ - - def __init__( - self, - 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, - ): - super().__init__() - - deprecation_message = "Use of `LoRALinearLayer` is deprecated. Please switch to PEFT backend by installing PEFT: `pip install peft`." - deprecate("LoRALinearLayer", "1.0.0", deprecation_message) - - self.down = nn.Linear(in_features, rank, bias=False, device=device, dtype=dtype) - self.up = nn.Linear(rank, out_features, bias=False, device=device, dtype=dtype) - # This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. - # See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning - self.network_alpha = network_alpha - self.rank = rank - self.out_features = out_features - self.in_features = in_features - - nn.init.normal_(self.down.weight, std=1 / rank) - nn.init.zeros_(self.up.weight) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - orig_dtype = hidden_states.dtype - dtype = self.down.weight.dtype - - down_hidden_states = self.down(hidden_states.to(dtype)) - up_hidden_states = self.up(down_hidden_states) - - if self.network_alpha is not None: - up_hidden_states *= self.network_alpha / self.rank - - return up_hidden_states.to(orig_dtype) - - -class LoRAConv2dLayer(nn.Module): - r""" - A convolutional layer that is used with LoRA. - - Parameters: - in_features (`int`): - Number of input features. - out_features (`int`): - Number of output features. - rank (`int`, `optional`, defaults to 4): - The rank of the LoRA layer. - kernel_size (`int` or `tuple` of two `int`, `optional`, defaults to 1): - The kernel size of the convolution. - stride (`int` or `tuple` of two `int`, `optional`, defaults to 1): - The stride of the convolution. - padding (`int` or `tuple` of two `int` or `str`, `optional`, defaults to 0): - The padding of the convolution. - network_alpha (`float`, `optional`, defaults to `None`): - 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. See - https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning - """ - - def __init__( - self, - 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, - ): - super().__init__() - - deprecation_message = "Use of `LoRAConv2dLayer` is deprecated. Please switch to PEFT backend by installing PEFT: `pip install peft`." - deprecate("LoRAConv2dLayer", "1.0.0", deprecation_message) - - self.down = nn.Conv2d(in_features, rank, kernel_size=kernel_size, stride=stride, padding=padding, bias=False) - # according to the official kohya_ss trainer kernel_size are always fixed for the up layer - # # see: https://github.com/bmaltais/kohya_ss/blob/2accb1305979ba62f5077a23aabac23b4c37e935/networks/lora_diffusers.py#L129 - self.up = nn.Conv2d(rank, out_features, kernel_size=(1, 1), stride=(1, 1), bias=False) - - # This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. - # See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning - self.network_alpha = network_alpha - self.rank = rank - - nn.init.normal_(self.down.weight, std=1 / rank) - nn.init.zeros_(self.up.weight) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - orig_dtype = hidden_states.dtype - dtype = self.down.weight.dtype - - down_hidden_states = self.down(hidden_states.to(dtype)) - up_hidden_states = self.up(down_hidden_states) - - if self.network_alpha is not None: - up_hidden_states *= self.network_alpha / self.rank - - return up_hidden_states.to(orig_dtype) - - -class LoRACompatibleConv(nn.Conv2d): - """ - A convolutional layer that can be used with LoRA. - """ - - def __init__(self, *args, lora_layer: Optional[LoRAConv2dLayer] = 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]): - 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 - - def _fuse_lora(self, lora_scale: float = 1.0, safe_fusing: bool = False): - if self.lora_layer is None: - return - - dtype, device = self.weight.data.dtype, self.weight.data.device - - w_orig = self.weight.data.float() - w_up = self.lora_layer.up.weight.data.float() - w_down = self.lora_layer.down.weight.data.float() - - if self.lora_layer.network_alpha is not None: - w_up = w_up * self.lora_layer.network_alpha / self.lora_layer.rank - - fusion = torch.mm(w_up.flatten(start_dim=1), w_down.flatten(start_dim=1)) - fusion = fusion.reshape((w_orig.shape)) - fused_weight = w_orig + (lora_scale * fusion) - - if safe_fusing and torch.isnan(fused_weight).any().item(): - raise ValueError( - "This LoRA weight seems to be broken. " - f"Encountered NaN values when trying to fuse LoRA weights for {self}." - "LoRA weights will not be fused." - ) - - self.weight.data = fused_weight.to(device=device, dtype=dtype) - - # we can drop the lora layer now - self.lora_layer = None - - # offload the up and down matrices to CPU to not blow the memory - self.w_up = w_up.cpu() - self.w_down = w_down.cpu() - self._lora_scale = lora_scale - - def _unfuse_lora(self): - if not (getattr(self, "w_up", None) is not None and getattr(self, "w_down", None) is not None): - return - - fused_weight = self.weight.data - dtype, device = fused_weight.data.dtype, fused_weight.data.device - - self.w_up = self.w_up.to(device=device).float() - self.w_down = self.w_down.to(device).float() - - fusion = torch.mm(self.w_up.flatten(start_dim=1), self.w_down.flatten(start_dim=1)) - fusion = fusion.reshape((fused_weight.shape)) - unfused_weight = fused_weight.float() - (self._lora_scale * fusion) - self.weight.data = unfused_weight.to(device=device, dtype=dtype) - - self.w_up = None - self.w_down = None - - def forward(self, hidden_states: torch.Tensor, scale: float = 1.0) -> torch.Tensor: - if self.padding_mode != "zeros": - hidden_states = F.pad(hidden_states, self._reversed_padding_repeated_twice, mode=self.padding_mode) - padding = (0, 0) - else: - padding = self.padding - - original_outputs = F.conv2d( - hidden_states, self.weight, self.bias, self.stride, padding, self.dilation, self.groups - ) - - if self.lora_layer is None: - return original_outputs - else: - return original_outputs + (scale * self.lora_layer(hidden_states)) - - -class LoRACompatibleLinear(nn.Linear): - """ - A Linear layer that can be used with LoRA. - """ - - def __init__(self, *args, lora_layer: Optional[LoRALinearLayer] = 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]): - 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 - - def _fuse_lora(self, lora_scale: float = 1.0, safe_fusing: bool = False): - if self.lora_layer is None: - return - - dtype, device = self.weight.data.dtype, self.weight.data.device - - w_orig = self.weight.data.float() - w_up = self.lora_layer.up.weight.data.float() - w_down = self.lora_layer.down.weight.data.float() - - if self.lora_layer.network_alpha is not None: - w_up = w_up * self.lora_layer.network_alpha / self.lora_layer.rank - - fused_weight = w_orig + (lora_scale * torch.bmm(w_up[None, :], w_down[None, :])[0]) - - if safe_fusing and torch.isnan(fused_weight).any().item(): - raise ValueError( - "This LoRA weight seems to be broken. " - f"Encountered NaN values when trying to fuse LoRA weights for {self}." - "LoRA weights will not be fused." - ) - - self.weight.data = fused_weight.to(device=device, dtype=dtype) - - # we can drop the lora layer now - self.lora_layer = None - - # offload the up and down matrices to CPU to not blow the memory - self.w_up = w_up.cpu() - self.w_down = w_down.cpu() - self._lora_scale = lora_scale - - def _unfuse_lora(self): - if not (getattr(self, "w_up", None) is not None and getattr(self, "w_down", None) is not None): - return - - fused_weight = self.weight.data - dtype, device = fused_weight.dtype, fused_weight.device - - w_up = self.w_up.to(device=device).float() - w_down = self.w_down.to(device).float() - - unfused_weight = fused_weight.float() - (self._lora_scale * torch.bmm(w_up[None, :], w_down[None, :])[0]) - self.weight.data = unfused_weight.to(device=device, dtype=dtype) - - self.w_up = None - self.w_down = None - - def forward(self, hidden_states: torch.Tensor, scale: float = 1.0) -> torch.Tensor: - if self.lora_layer is None: - out = super().forward(hidden_states) - return out - else: - out = super().forward(hidden_states) + (scale * self.lora_layer(hidden_states)) - return out diff --git a/venv/lib/python3.10/site-packages/diffusers/models/model_loading_utils.py b/venv/lib/python3.10/site-packages/diffusers/models/model_loading_utils.py deleted file mode 100644 index ebc7d79aeb281f1bb55a08fc6620511223dbdd7f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/model_loading_utils.py +++ /dev/null @@ -1,522 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# Copyright (c) 2022, NVIDIA CORPORATION. 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 -import os -from array import array -from collections import OrderedDict -from pathlib import Path -from typing import Dict, List, Optional, Union -from zipfile import is_zipfile - -import safetensors -import torch -from huggingface_hub import DDUFEntry -from huggingface_hub.utils import EntryNotFoundError - -from ..quantizers import DiffusersQuantizer -from ..utils import ( - GGUF_FILE_EXTENSION, - SAFE_WEIGHTS_INDEX_NAME, - SAFETENSORS_FILE_EXTENSION, - WEIGHTS_INDEX_NAME, - _add_variant, - _get_model_file, - deprecate, - is_accelerate_available, - is_gguf_available, - is_torch_available, - is_torch_version, - logging, -) - - -logger = logging.get_logger(__name__) - -_CLASS_REMAPPING_DICT = { - "Transformer2DModel": { - "ada_norm_zero": "DiTTransformer2DModel", - "ada_norm_single": "PixArtTransformer2DModel", - } -} - - -if is_accelerate_available(): - from accelerate import infer_auto_device_map - from accelerate.utils import get_balanced_memory, get_max_memory, offload_weight, set_module_tensor_to_device - - -# Adapted from `transformers` (see modeling_utils.py) -def _determine_device_map( - model: torch.nn.Module, device_map, max_memory, torch_dtype, keep_in_fp32_modules=[], hf_quantizer=None -): - if isinstance(device_map, str): - special_dtypes = {} - if hf_quantizer is not None: - special_dtypes.update(hf_quantizer.get_special_dtypes_update(model, torch_dtype)) - special_dtypes.update( - { - name: torch.float32 - for name, _ in model.named_parameters() - if any(m in name for m in keep_in_fp32_modules) - } - ) - - target_dtype = torch_dtype - if hf_quantizer is not None: - target_dtype = hf_quantizer.adjust_target_dtype(target_dtype) - - no_split_modules = model._get_no_split_modules(device_map) - device_map_kwargs = {"no_split_module_classes": no_split_modules} - - if "special_dtypes" in inspect.signature(infer_auto_device_map).parameters: - device_map_kwargs["special_dtypes"] = special_dtypes - elif len(special_dtypes) > 0: - logger.warning( - "This model has some weights that should be kept in higher precision, you need to upgrade " - "`accelerate` to properly deal with them (`pip install --upgrade accelerate`)." - ) - - if device_map != "sequential": - max_memory = get_balanced_memory( - model, - dtype=torch_dtype, - low_zero=(device_map == "balanced_low_0"), - max_memory=max_memory, - **device_map_kwargs, - ) - else: - max_memory = get_max_memory(max_memory) - - if hf_quantizer is not None: - max_memory = hf_quantizer.adjust_max_memory(max_memory) - - device_map_kwargs["max_memory"] = max_memory - device_map = infer_auto_device_map(model, dtype=target_dtype, **device_map_kwargs) - - if hf_quantizer is not None: - hf_quantizer.validate_environment(device_map=device_map) - - return device_map - - -def _fetch_remapped_cls_from_config(config, old_class): - previous_class_name = old_class.__name__ - remapped_class_name = _CLASS_REMAPPING_DICT.get(previous_class_name).get(config["norm_type"], None) - - # Details: - # https://github.com/huggingface/diffusers/pull/7647#discussion_r1621344818 - if remapped_class_name: - # load diffusers library to import compatible and original scheduler - diffusers_library = importlib.import_module(__name__.split(".")[0]) - remapped_class = getattr(diffusers_library, remapped_class_name) - logger.info( - f"Changing class object to be of `{remapped_class_name}` type from `{previous_class_name}` type." - f"This is because `{previous_class_name}` is scheduled to be deprecated in a future version. Note that this" - " DOESN'T affect the final results." - ) - return remapped_class - else: - return old_class - - -def _determine_param_device(param_name: str, device_map: Optional[Dict[str, Union[int, str, torch.device]]]): - """ - Find the device of param_name from the device_map. - """ - if device_map is None: - return "cpu" - else: - module_name = param_name - # find next higher level module that is defined in device_map: - # bert.lm_head.weight -> bert.lm_head -> bert -> '' - while len(module_name) > 0 and module_name not in device_map: - module_name = ".".join(module_name.split(".")[:-1]) - if module_name == "" and "" not in device_map: - raise ValueError(f"{param_name} doesn't have any device set.") - return device_map[module_name] - - -def load_state_dict( - checkpoint_file: Union[str, os.PathLike], - dduf_entries: Optional[Dict[str, DDUFEntry]] = None, - disable_mmap: bool = False, - map_location: Union[str, torch.device] = "cpu", -): - """ - Reads a checkpoint file, returning properly formatted errors if they arise. - """ - # TODO: maybe refactor a bit this part where we pass a dict here - if isinstance(checkpoint_file, dict): - return checkpoint_file - try: - file_extension = os.path.basename(checkpoint_file).split(".")[-1] - if file_extension == SAFETENSORS_FILE_EXTENSION: - if dduf_entries: - # tensors are loaded on cpu - with dduf_entries[checkpoint_file].as_mmap() as mm: - return safetensors.torch.load(mm) - if disable_mmap: - return safetensors.torch.load(open(checkpoint_file, "rb").read()) - else: - return safetensors.torch.load_file(checkpoint_file, device=map_location) - elif file_extension == GGUF_FILE_EXTENSION: - return load_gguf_checkpoint(checkpoint_file) - else: - extra_args = {} - weights_only_kwarg = {"weights_only": True} if is_torch_version(">=", "1.13") else {} - # mmap can only be used with files serialized with zipfile-based format. - if ( - isinstance(checkpoint_file, str) - and map_location != "meta" - and is_torch_version(">=", "2.1.0") - and is_zipfile(checkpoint_file) - and not disable_mmap - ): - extra_args = {"mmap": True} - return torch.load(checkpoint_file, map_location=map_location, **weights_only_kwarg, **extra_args) - except Exception as e: - try: - with open(checkpoint_file) as f: - if f.read().startswith("version"): - raise OSError( - "You seem to have cloned a repository without having git-lfs installed. Please install " - "git-lfs and run `git lfs install` followed by `git lfs pull` in the folder " - "you cloned." - ) - else: - raise ValueError( - f"Unable to locate the file {checkpoint_file} which is necessary to load this pretrained " - "model. Make sure you have saved the model properly." - ) from e - except (UnicodeDecodeError, ValueError): - raise OSError( - f"Unable to load weights from checkpoint file for '{checkpoint_file}' at '{checkpoint_file}'. " - ) - - -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]: - """ - 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` - """ - - is_quantized = hf_quantizer is not None - empty_state_dict = model.state_dict() - - for param_name, param in state_dict.items(): - if param_name not in empty_state_dict: - continue - - set_module_kwargs = {} - # We convert floating dtypes to the `dtype` passed. We also want to keep the buffers/params - # in int/uint/bool and not cast them. - # TODO: revisit cases when param.dtype == torch.float8_e4m3fn - if dtype is not None and torch.is_floating_point(param): - if keep_in_fp32_modules is not None and any( - module_to_keep_in_fp32 in param_name.split(".") for module_to_keep_in_fp32 in keep_in_fp32_modules - ): - param = param.to(torch.float32) - set_module_kwargs["dtype"] = torch.float32 - # For quantizers have save weights using torch.float8_e4m3fn - elif hf_quantizer is not None and param.dtype == getattr(torch, "float8_e4m3fn", None): - pass - else: - param = param.to(dtype) - set_module_kwargs["dtype"] = dtype - - # For compatibility with PyTorch load_state_dict which converts state dict dtype to existing dtype in model, and which - # uses `param.copy_(input_param)` that preserves the contiguity of the parameter in the model. - # Reference: https://github.com/pytorch/pytorch/blob/db79ceb110f6646523019a59bbd7b838f43d4a86/torch/nn/modules/module.py#L2040C29-L2040C29 - old_param = model - splits = param_name.split(".") - for split in splits: - old_param = getattr(old_param, split) - - if not isinstance(old_param, (torch.nn.Parameter, torch.Tensor)): - old_param = None - - if old_param is not None: - if dtype is None: - param = param.to(old_param.dtype) - - if old_param.is_contiguous(): - param = param.contiguous() - - param_device = _determine_param_device(param_name, device_map) - - # bnb params are flattened. - # gguf quants have a different shape based on the type of quantization applied - if empty_state_dict[param_name].shape != param.shape: - if ( - is_quantized - and hf_quantizer.pre_quantized - and hf_quantizer.check_if_quantized_param( - model, param, param_name, state_dict, param_device=param_device - ) - ): - hf_quantizer.check_quantized_param_shape(param_name, empty_state_dict[param_name], param) - else: - model_name_or_path_str = f"{model_name_or_path} " if model_name_or_path is not None else "" - raise ValueError( - f"Cannot load {model_name_or_path_str} because {param_name} expected shape {empty_state_dict[param_name].shape}, but got {param.shape}. If you want to instead overwrite randomly initialized weights, please make sure to pass both `low_cpu_mem_usage=False` and `ignore_mismatched_sizes=True`. For more information, see also: https://github.com/huggingface/diffusers/issues/1619#issuecomment-1345604389 as an example." - ) - if param_device == "disk": - offload_index = offload_weight(param, param_name, offload_folder, offload_index) - elif param_device == "cpu" and state_dict_index is not None: - state_dict_index = offload_weight(param, param_name, state_dict_folder, state_dict_index) - elif is_quantized and ( - hf_quantizer.check_if_quantized_param(model, param, param_name, state_dict, param_device=param_device) - ): - hf_quantizer.create_quantized_param( - model, param, param_name, param_device, state_dict, unexpected_keys, dtype=dtype - ) - else: - set_module_tensor_to_device(model, param_name, param_device, value=param, **set_module_kwargs) - - return offload_index, state_dict_index - - -def _load_state_dict_into_model( - model_to_load, state_dict: OrderedDict, assign_to_params_buffers: bool = False -) -> 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() - error_msgs = [] - - # PyTorch's `_load_from_state_dict` does not copy parameters in a module's descendants - # so we need to apply the function recursively. - def load(module: torch.nn.Module, prefix: str = "", assign_to_params_buffers: bool = False): - local_metadata = {} - local_metadata["assign_to_params_buffers"] = assign_to_params_buffers - if assign_to_params_buffers and not is_torch_version(">=", "2.1"): - logger.info("You need to have torch>=2.1 in order to load the model with assign_to_params_buffers=True") - args = (state_dict, prefix, local_metadata, True, [], [], error_msgs) - module._load_from_state_dict(*args) - - for name, child in module._modules.items(): - if child is not None: - load(child, prefix + name + ".", assign_to_params_buffers) - - load(model_to_load, assign_to_params_buffers=assign_to_params_buffers) - - return error_msgs - - -def _fetch_index_file( - is_local, - pretrained_model_name_or_path, - subfolder, - use_safetensors, - cache_dir, - variant, - force_download, - proxies, - local_files_only, - token, - revision, - user_agent, - commit_hash, - dduf_entries: Optional[Dict[str, DDUFEntry]] = None, -): - if is_local: - index_file = Path( - pretrained_model_name_or_path, - subfolder or "", - _add_variant(SAFE_WEIGHTS_INDEX_NAME if use_safetensors else WEIGHTS_INDEX_NAME, variant), - ) - else: - index_file_in_repo = Path( - subfolder or "", - _add_variant(SAFE_WEIGHTS_INDEX_NAME if use_safetensors else WEIGHTS_INDEX_NAME, variant), - ).as_posix() - try: - index_file = _get_model_file( - pretrained_model_name_or_path, - weights_name=index_file_in_repo, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=None, - user_agent=user_agent, - commit_hash=commit_hash, - dduf_entries=dduf_entries, - ) - if not dduf_entries: - index_file = Path(index_file) - except (EntryNotFoundError, EnvironmentError): - index_file = None - - return index_file - - -def _fetch_index_file_legacy( - is_local, - pretrained_model_name_or_path, - subfolder, - use_safetensors, - cache_dir, - variant, - force_download, - proxies, - local_files_only, - token, - revision, - user_agent, - commit_hash, - dduf_entries: Optional[Dict[str, DDUFEntry]] = None, -): - if is_local: - index_file = Path( - pretrained_model_name_or_path, - subfolder or "", - SAFE_WEIGHTS_INDEX_NAME if use_safetensors else WEIGHTS_INDEX_NAME, - ).as_posix() - splits = index_file.split(".") - split_index = -3 if ".cache" in index_file else -2 - splits = splits[:-split_index] + [variant] + splits[-split_index:] - index_file = ".".join(splits) - if os.path.exists(index_file): - deprecation_message = f"This serialization format is now deprecated to standardize the serialization format between `transformers` and `diffusers`. We recommend you to remove the existing files associated with the current variant ({variant}) and re-obtain them by running a `save_pretrained()`." - deprecate("legacy_sharded_ckpts_with_variant", "1.0.0", deprecation_message, standard_warn=False) - index_file = Path(index_file) - else: - index_file = None - else: - if variant is not None: - index_file_in_repo = Path( - subfolder or "", - SAFE_WEIGHTS_INDEX_NAME if use_safetensors else WEIGHTS_INDEX_NAME, - ).as_posix() - splits = index_file_in_repo.split(".") - split_index = -2 - splits = splits[:-split_index] + [variant] + splits[-split_index:] - index_file_in_repo = ".".join(splits) - try: - index_file = _get_model_file( - pretrained_model_name_or_path, - weights_name=index_file_in_repo, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=None, - user_agent=user_agent, - commit_hash=commit_hash, - dduf_entries=dduf_entries, - ) - index_file = Path(index_file) - deprecation_message = f"This serialization format is now deprecated to standardize the serialization format between `transformers` and `diffusers`. We recommend you to remove the existing files associated with the current variant ({variant}) and re-obtain them by running a `save_pretrained()`." - deprecate("legacy_sharded_ckpts_with_variant", "1.0.0", deprecation_message, standard_warn=False) - except (EntryNotFoundError, EnvironmentError): - index_file = None - - return index_file - - -def _gguf_parse_value(_value, data_type): - if not isinstance(data_type, list): - data_type = [data_type] - if len(data_type) == 1: - data_type = data_type[0] - array_data_type = None - else: - if data_type[0] != 9: - raise ValueError("Received multiple types, therefore expected the first type to indicate an array.") - data_type, array_data_type = data_type - - if data_type in [0, 1, 2, 3, 4, 5, 10, 11]: - _value = int(_value[0]) - elif data_type in [6, 12]: - _value = float(_value[0]) - elif data_type in [7]: - _value = bool(_value[0]) - elif data_type in [8]: - _value = array("B", list(_value)).tobytes().decode() - elif data_type in [9]: - _value = _gguf_parse_value(_value, array_data_type) - return _value - - -def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False): - """ - Load a GGUF file and return a dictionary of parsed parameters containing tensors, the parsed tokenizer and config - attributes. - - Args: - gguf_checkpoint_path (`str`): - The path the to GGUF file to load - return_tensors (`bool`, defaults to `True`): - Whether to read the tensors from the file and return them. Not doing so is faster and only loads the - metadata in memory. - """ - - if is_gguf_available() and is_torch_available(): - import gguf - from gguf import GGUFReader - - from ..quantizers.gguf.utils import SUPPORTED_GGUF_QUANT_TYPES, GGUFParameter - else: - logger.error( - "Loading a GGUF checkpoint in PyTorch, requires both PyTorch and GGUF>=0.10.0 to be installed. Please see " - "https://pytorch.org/ and https://github.com/ggerganov/llama.cpp/tree/master/gguf-py for installation instructions." - ) - raise ImportError("Please install torch and gguf>=0.10.0 to load a GGUF checkpoint in PyTorch.") - - reader = GGUFReader(gguf_checkpoint_path) - - parsed_parameters = {} - for tensor in reader.tensors: - name = tensor.name - quant_type = tensor.tensor_type - - # if the tensor is a torch supported dtype do not use GGUFParameter - is_gguf_quant = quant_type not in [gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16] - if is_gguf_quant and quant_type not in SUPPORTED_GGUF_QUANT_TYPES: - _supported_quants_str = "\n".join([str(type) for type in SUPPORTED_GGUF_QUANT_TYPES]) - raise ValueError( - ( - f"{name} has a quantization type: {str(quant_type)} which is unsupported." - "\n\nCurrently the following quantization types are supported: \n\n" - f"{_supported_quants_str}" - "\n\nTo request support for this quantization type please open an issue here: https://github.com/huggingface/diffusers" - ) - ) - - weights = torch.from_numpy(tensor.data.copy()) - parsed_parameters[name] = GGUFParameter(weights, quant_type=quant_type) if is_gguf_quant else weights - - return parsed_parameters diff --git a/venv/lib/python3.10/site-packages/diffusers/models/modeling_flax_pytorch_utils.py b/venv/lib/python3.10/site-packages/diffusers/models/modeling_flax_pytorch_utils.py deleted file mode 100644 index d64c48a9601ec75afc6685174ede20673df5d01f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/modeling_flax_pytorch_utils.py +++ /dev/null @@ -1,135 +0,0 @@ -# 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. -"""PyTorch - Flax general utilities.""" - -import re - -import jax.numpy as jnp -from flax.traverse_util import flatten_dict, unflatten_dict -from jax.random import PRNGKey - -from ..utils import logging - - -logger = logging.get_logger(__name__) - - -def rename_key(key): - regex = r"\w+[.]\d+" - pats = re.findall(regex, key) - for pat in pats: - key = key.replace(pat, "_".join(pat.split("."))) - return key - - -##################### -# PyTorch => Flax # -##################### - - -# Adapted from https://github.com/huggingface/transformers/blob/c603c80f46881ae18b2ca50770ef65fa4033eacd/src/transformers/modeling_flax_pytorch_utils.py#L69 -# and https://github.com/patil-suraj/stable-diffusion-jax/blob/main/stable_diffusion_jax/convert_diffusers_to_jax.py -def rename_key_and_reshape_tensor(pt_tuple_key, pt_tensor, random_flax_state_dict): - """Rename PT weight names to corresponding Flax weight names and reshape tensor if necessary""" - # conv norm or layer norm - renamed_pt_tuple_key = pt_tuple_key[:-1] + ("scale",) - - # rename attention layers - if len(pt_tuple_key) > 1: - for rename_from, rename_to in ( - ("to_out_0", "proj_attn"), - ("to_k", "key"), - ("to_v", "value"), - ("to_q", "query"), - ): - if pt_tuple_key[-2] == rename_from: - weight_name = pt_tuple_key[-1] - weight_name = "kernel" if weight_name == "weight" else weight_name - renamed_pt_tuple_key = pt_tuple_key[:-2] + (rename_to, weight_name) - if renamed_pt_tuple_key in random_flax_state_dict: - assert random_flax_state_dict[renamed_pt_tuple_key].shape == pt_tensor.T.shape - return renamed_pt_tuple_key, pt_tensor.T - - if ( - any("norm" in str_ for str_ in pt_tuple_key) - and (pt_tuple_key[-1] == "bias") - and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict) - and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict) - ): - renamed_pt_tuple_key = pt_tuple_key[:-1] + ("scale",) - return renamed_pt_tuple_key, pt_tensor - elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict: - renamed_pt_tuple_key = pt_tuple_key[:-1] + ("scale",) - return renamed_pt_tuple_key, pt_tensor - - # embedding - if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict: - pt_tuple_key = pt_tuple_key[:-1] + ("embedding",) - return renamed_pt_tuple_key, pt_tensor - - # conv layer - renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",) - if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4: - pt_tensor = pt_tensor.transpose(2, 3, 1, 0) - return renamed_pt_tuple_key, pt_tensor - - # linear layer - renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",) - if pt_tuple_key[-1] == "weight": - pt_tensor = pt_tensor.T - return renamed_pt_tuple_key, pt_tensor - - # old PyTorch layer norm weight - renamed_pt_tuple_key = pt_tuple_key[:-1] + ("weight",) - if pt_tuple_key[-1] == "gamma": - return renamed_pt_tuple_key, pt_tensor - - # old PyTorch layer norm bias - renamed_pt_tuple_key = pt_tuple_key[:-1] + ("bias",) - if pt_tuple_key[-1] == "beta": - return renamed_pt_tuple_key, pt_tensor - - return pt_tuple_key, pt_tensor - - -def convert_pytorch_state_dict_to_flax(pt_state_dict, flax_model, init_key=42): - # Step 1: Convert pytorch tensor to numpy - pt_state_dict = {k: v.numpy() for k, v in pt_state_dict.items()} - - # Step 2: Since the model is stateless, get random Flax params - random_flax_params = flax_model.init_weights(PRNGKey(init_key)) - - random_flax_state_dict = flatten_dict(random_flax_params) - flax_state_dict = {} - - # Need to change some parameters name to match Flax names - for pt_key, pt_tensor in pt_state_dict.items(): - renamed_pt_key = rename_key(pt_key) - pt_tuple_key = tuple(renamed_pt_key.split(".")) - - # Correctly rename weight parameters - flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, pt_tensor, random_flax_state_dict) - - if flax_key in random_flax_state_dict: - if flax_tensor.shape != random_flax_state_dict[flax_key].shape: - raise ValueError( - f"PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape " - f"{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}." - ) - - # also add unexpected weight so that warning is thrown - flax_state_dict[flax_key] = jnp.asarray(flax_tensor) - - return unflatten_dict(flax_state_dict) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/modeling_flax_utils.py b/venv/lib/python3.10/site-packages/diffusers/models/modeling_flax_utils.py deleted file mode 100644 index 52f004f6f93fd32360d43997ea0ea40de7dd9a08..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/modeling_flax_utils.py +++ /dev/null @@ -1,561 +0,0 @@ -# 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. - -import os -from pickle import UnpicklingError -from typing import Any, Dict, Union - -import jax -import jax.numpy as jnp -import msgpack.exceptions -from flax.core.frozen_dict import FrozenDict, unfreeze -from flax.serialization import from_bytes, to_bytes -from flax.traverse_util import flatten_dict, unflatten_dict -from huggingface_hub import create_repo, hf_hub_download -from huggingface_hub.utils import ( - EntryNotFoundError, - RepositoryNotFoundError, - RevisionNotFoundError, - validate_hf_hub_args, -) -from requests import HTTPError - -from .. import __version__, is_torch_available -from ..utils import ( - CONFIG_NAME, - FLAX_WEIGHTS_NAME, - HUGGINGFACE_CO_RESOLVE_ENDPOINT, - WEIGHTS_NAME, - PushToHubMixin, - logging, -) -from .modeling_flax_pytorch_utils import convert_pytorch_state_dict_to_flax - - -logger = logging.get_logger(__name__) - - -class FlaxModelMixin(PushToHubMixin): - r""" - Base class for all Flax models. - - [`FlaxModelMixin`] takes care of storing the model configuration and provides methods for loading, downloading and - saving models. - - - **config_name** ([`str`]) -- Filename to save a model to when calling [`~FlaxModelMixin.save_pretrained`]. - """ - - config_name = CONFIG_NAME - _automatically_saved_args = ["_diffusers_version", "_class_name", "_name_or_path"] - _flax_internal_args = ["name", "parent", "dtype"] - - @classmethod - def _from_config(cls, config, **kwargs): - """ - All context managers that the model should be initialized under go here. - """ - return cls(config, **kwargs) - - def _cast_floating_to(self, params: Union[Dict, FrozenDict], dtype: jnp.dtype, mask: Any = None) -> Any: - """ - Helper method to cast floating-point values of given parameter `PyTree` to given `dtype`. - """ - - # taken from https://github.com/deepmind/jmp/blob/3a8318abc3292be38582794dbf7b094e6583b192/jmp/_src/policy.py#L27 - def conditional_cast(param): - if isinstance(param, jnp.ndarray) and jnp.issubdtype(param.dtype, jnp.floating): - param = param.astype(dtype) - return param - - if mask is None: - return jax.tree_map(conditional_cast, params) - - flat_params = flatten_dict(params) - flat_mask, _ = jax.tree_flatten(mask) - - for masked, key in zip(flat_mask, flat_params.keys()): - if masked: - param = flat_params[key] - flat_params[key] = conditional_cast(param) - - return unflatten_dict(flat_params) - - def to_bf16(self, params: Union[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. - - This method can be used on a TPU to explicitly convert the model parameters to bfloat16 precision to do full - half-precision training or to save weights in bfloat16 for inference in order to save memory and improve speed. - - Arguments: - params (`Union[Dict, FrozenDict]`): - A `PyTree` of model parameters. - mask (`Union[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. - - Examples: - - ```python - >>> from diffusers import FlaxUNet2DConditionModel - - >>> # load model - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5") - >>> # By default, the model parameters will be in fp32 precision, to cast these to bfloat16 precision - >>> params = model.to_bf16(params) - >>> # If you don't want to cast certain parameters (for example layer norm bias and scale) - >>> # then pass the mask as follows - >>> from flax import traverse_util - - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5") - >>> flat_params = traverse_util.flatten_dict(params) - >>> mask = { - ... path: (path[-2] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale")) - ... for path in flat_params - ... } - >>> mask = traverse_util.unflatten_dict(mask) - >>> params = model.to_bf16(params, mask) - ```""" - return self._cast_floating_to(params, jnp.bfloat16, mask) - - def to_fp32(self, params: Union[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]`): - A `PyTree` of model parameters. - mask (`Union[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. - - Examples: - - ```python - >>> from diffusers import FlaxUNet2DConditionModel - - >>> # Download model and configuration from huggingface.co - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5") - >>> # By default, the model params will be in fp32, to illustrate the use of this method, - >>> # we'll first cast to fp16 and back to fp32 - >>> params = model.to_f16(params) - >>> # now cast back to fp32 - >>> params = model.to_fp32(params) - ```""" - return self._cast_floating_to(params, jnp.float32, mask) - - def to_fp16(self, params: Union[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. - - This method can be used on a GPU to explicitly convert the model parameters to float16 precision to do full - half-precision training or to save weights in float16 for inference in order to save memory and improve speed. - - Arguments: - params (`Union[Dict, FrozenDict]`): - A `PyTree` of model parameters. - mask (`Union[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. - - Examples: - - ```python - >>> from diffusers import FlaxUNet2DConditionModel - - >>> # load model - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5") - >>> # By default, the model params will be in fp32, to cast these to float16 - >>> params = model.to_fp16(params) - >>> # If you want don't want to cast certain parameters (for example layer norm bias and scale) - >>> # then pass the mask as follows - >>> from flax import traverse_util - - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5") - >>> flat_params = traverse_util.flatten_dict(params) - >>> mask = { - ... path: (path[-2] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale")) - ... for path in flat_params - ... } - >>> mask = traverse_util.unflatten_dict(mask) - >>> params = model.to_fp16(params, mask) - ```""" - return self._cast_floating_to(params, jnp.float16, mask) - - 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], - dtype: jnp.dtype = jnp.float32, - *model_args, - **kwargs, - ): - r""" - Instantiate a pretrained Flax model from a pretrained model configuration. - - Parameters: - pretrained_model_name_or_path (`str` or `os.PathLike`): - Can be either: - - - A string, the *model id* (for example `runwayml/stable-diffusion-v1-5`) of a pretrained model - hosted on the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - using [`~FlaxModelMixin.save_pretrained`]. - dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`): - The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and - `jax.numpy.bfloat16` (on TPUs). - - This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If - specified, all the computation will be performed with the given `dtype`. - - - - This only specifies the dtype of the *computation* and does not influence the dtype of model - parameters. - - If you wish to change the dtype of the model parameters, see [`~FlaxModelMixin.to_fp16`] and - [`~FlaxModelMixin.to_bf16`]. - - - - 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*): - 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*): - 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`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - from_pt (`bool`, *optional*, defaults to `False`): - Load the model weights from a PyTorch checkpoint save file. - kwargs (remaining dictionary of keyword arguments, *optional*): - Can be used to update the configuration object (after it is loaded) and initiate the model (for - example, `output_attentions=True`). Behaves differently depending on whether a `config` is provided or - automatically loaded: - - - If a configuration is provided with `config`, `kwargs` are directly passed to the underlying - model's `__init__` method (we assume all relevant updates to the configuration have already been - done). - - If a configuration is not provided, `kwargs` are first passed to the configuration class - initialization function [`~ConfigMixin.from_config`]. Each key of the `kwargs` that corresponds - to a configuration attribute is used to override said attribute with the supplied `kwargs` value. - Remaining keys that do not correspond to any configuration attribute are passed to the underlying - model's `__init__` function. - - Examples: - - ```python - >>> from diffusers import FlaxUNet2DConditionModel - - >>> # Download model and configuration from huggingface.co and cache. - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5") - >>> # Model was saved using *save_pretrained('./test/saved_model/')* (for example purposes, not runnable). - >>> model, params = FlaxUNet2DConditionModel.from_pretrained("./test/saved_model/") - ``` - - If you get the error message below, you need to finetune the weights for your downstream task: - - ```bash - Some weights of UNet2DConditionModel were not initialized from the model checkpoint at runwayml/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: - - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated - You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. - ``` - """ - config = kwargs.pop("config", None) - cache_dir = kwargs.pop("cache_dir", None) - force_download = kwargs.pop("force_download", False) - from_pt = kwargs.pop("from_pt", False) - proxies = kwargs.pop("proxies", None) - local_files_only = kwargs.pop("local_files_only", False) - token = kwargs.pop("token", None) - revision = kwargs.pop("revision", None) - subfolder = kwargs.pop("subfolder", None) - - user_agent = { - "diffusers": __version__, - "file_type": "model", - "framework": "flax", - } - - # Load config if we don't provide one - if config is None: - config, unused_kwargs = cls.load_config( - pretrained_model_name_or_path, - cache_dir=cache_dir, - return_unused_kwargs=True, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - **kwargs, - ) - - model, model_kwargs = cls.from_config(config, dtype=dtype, return_unused_kwargs=True, **unused_kwargs) - - # Load model - pretrained_path_with_subfolder = ( - pretrained_model_name_or_path - if subfolder is None - else os.path.join(pretrained_model_name_or_path, subfolder) - ) - if os.path.isdir(pretrained_path_with_subfolder): - if from_pt: - if not os.path.isfile(os.path.join(pretrained_path_with_subfolder, WEIGHTS_NAME)): - raise EnvironmentError( - f"Error no file named {WEIGHTS_NAME} found in directory {pretrained_path_with_subfolder} " - ) - model_file = os.path.join(pretrained_path_with_subfolder, WEIGHTS_NAME) - elif os.path.isfile(os.path.join(pretrained_path_with_subfolder, FLAX_WEIGHTS_NAME)): - # Load from a Flax checkpoint - model_file = os.path.join(pretrained_path_with_subfolder, FLAX_WEIGHTS_NAME) - # Check if pytorch weights exist instead - elif os.path.isfile(os.path.join(pretrained_path_with_subfolder, WEIGHTS_NAME)): - raise EnvironmentError( - f"{WEIGHTS_NAME} file found in directory {pretrained_path_with_subfolder}. Please load the model" - " using `from_pt=True`." - ) - else: - raise EnvironmentError( - f"Error no file named {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME} found in directory " - f"{pretrained_path_with_subfolder}." - ) - else: - try: - model_file = hf_hub_download( - pretrained_model_name_or_path, - filename=FLAX_WEIGHTS_NAME if not from_pt else WEIGHTS_NAME, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - user_agent=user_agent, - subfolder=subfolder, - revision=revision, - ) - - except RepositoryNotFoundError: - raise EnvironmentError( - f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier " - "listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a " - "token having permission to this repo with `token` or log in with `huggingface-cli " - "login`." - ) - except RevisionNotFoundError: - raise EnvironmentError( - f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for " - "this model name. Check the model page at " - f"'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions." - ) - except EntryNotFoundError: - raise EnvironmentError( - f"{pretrained_model_name_or_path} does not appear to have a file named {FLAX_WEIGHTS_NAME}." - ) - except HTTPError as err: - raise EnvironmentError( - f"There was a specific connection error when trying to load {pretrained_model_name_or_path}:\n" - f"{err}" - ) - except ValueError: - raise EnvironmentError( - f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it" - f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a" - f" directory containing a file named {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME}.\nCheckout your" - " internet connection or see how to run the library in offline mode at" - " 'https://huggingface.co/docs/transformers/installation#offline-mode'." - ) - except EnvironmentError: - raise EnvironmentError( - f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it from " - "'https://huggingface.co/models', make sure you don't have a local directory with the same name. " - f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory " - f"containing a file named {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME}." - ) - - if from_pt: - if is_torch_available(): - from .modeling_utils import load_state_dict - else: - raise EnvironmentError( - "Can't load the model in PyTorch format because PyTorch is not installed. " - "Please, install PyTorch or use native Flax weights." - ) - - # Step 1: Get the pytorch file - pytorch_model_file = load_state_dict(model_file) - - # Step 2: Convert the weights - state = convert_pytorch_state_dict_to_flax(pytorch_model_file, model) - else: - try: - with open(model_file, "rb") as state_f: - state = from_bytes(cls, state_f.read()) - except (UnpicklingError, msgpack.exceptions.ExtraData) as e: - try: - with open(model_file) as f: - if f.read().startswith("version"): - raise OSError( - "You seem to have cloned a repository without having git-lfs installed. Please" - " install git-lfs and run `git lfs install` followed by `git lfs pull` in the" - " folder you cloned." - ) - else: - raise ValueError from e - except (UnicodeDecodeError, ValueError): - raise EnvironmentError(f"Unable to convert {model_file} to Flax deserializable object. ") - # make sure all arrays are stored as jnp.ndarray - # NOTE: This is to prevent a bug this will be fixed in Flax >= v0.3.4: - # https://github.com/google/flax/issues/1261 - state = jax.tree_util.tree_map(lambda x: jax.device_put(x, jax.local_devices(backend="cpu")[0]), state) - - # flatten dicts - state = flatten_dict(state) - - params_shape_tree = jax.eval_shape(model.init_weights, rng=jax.random.PRNGKey(0)) - required_params = set(flatten_dict(unfreeze(params_shape_tree)).keys()) - - shape_state = flatten_dict(unfreeze(params_shape_tree)) - - missing_keys = required_params - set(state.keys()) - unexpected_keys = set(state.keys()) - required_params - - if missing_keys: - logger.warning( - f"The checkpoint {pretrained_model_name_or_path} is missing required keys: {missing_keys}. " - "Make sure to call model.init_weights to initialize the missing weights." - ) - cls._missing_keys = missing_keys - - for key in state.keys(): - if key in shape_state and state[key].shape != shape_state[key].shape: - raise ValueError( - f"Trying to load the pretrained weight for {key} failed: checkpoint has shape " - f"{state[key].shape} which is incompatible with the model shape {shape_state[key].shape}. " - ) - - # remove unexpected keys to not be saved again - for unexpected_key in unexpected_keys: - del state[unexpected_key] - - if len(unexpected_keys) > 0: - logger.warning( - f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when" - f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are" - f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or" - " with another architecture." - ) - else: - logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n") - - if len(missing_keys) > 0: - logger.warning( - f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" - f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably" - " TRAIN this model on a down-stream task to be able to use it for predictions and inference." - ) - else: - logger.info( - f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at" - f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the checkpoint" - f" was trained on, you can already use {model.__class__.__name__} for predictions without further" - " training." - ) - - return model, unflatten_dict(state) - - def save_pretrained( - self, - save_directory: Union[str, os.PathLike], - params: Union[Dict, FrozenDict], - is_main_process: bool = True, - push_to_hub: bool = False, - **kwargs, - ): - """ - Save a model and its configuration file to a directory so that it can be reloaded using the - [`~FlaxModelMixin.from_pretrained`] class method. - - 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]`): - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - push_to_hub (`bool`, *optional*, defaults to `False`): - 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*): - Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. - """ - if os.path.isfile(save_directory): - logger.error(f"Provided path ({save_directory}) should be a directory, not a file") - return - - os.makedirs(save_directory, exist_ok=True) - - 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]) - repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id - - model_to_save = self - - # Attach architecture to the config - # Save the config - if is_main_process: - model_to_save.save_config(save_directory) - - # save model - output_model_file = os.path.join(save_directory, FLAX_WEIGHTS_NAME) - with open(output_model_file, "wb") as f: - model_bytes = to_bytes(params) - f.write(model_bytes) - - logger.info(f"Model weights saved in {output_model_file}") - - if push_to_hub: - self._upload_folder( - save_directory, - repo_id, - token=token, - commit_message=commit_message, - create_pr=create_pr, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/modeling_outputs.py b/venv/lib/python3.10/site-packages/diffusers/models/modeling_outputs.py deleted file mode 100644 index 0120a34d9052fe6b499b519ba4366e7a088f7910..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/modeling_outputs.py +++ /dev/null @@ -1,31 +0,0 @@ -from dataclasses import dataclass - -from ..utils import BaseOutput - - -@dataclass -class AutoencoderKLOutput(BaseOutput): - """ - Output of AutoencoderKL encoding method. - - Args: - latent_dist (`DiagonalGaussianDistribution`): - Encoded outputs of `Encoder` represented as the mean and logvar of `DiagonalGaussianDistribution`. - `DiagonalGaussianDistribution` allows for sampling latents from the distribution. - """ - - latent_dist: "DiagonalGaussianDistribution" # noqa: F821 - - -@dataclass -class Transformer2DModelOutput(BaseOutput): - """ - The output of [`Transformer2DModel`]. - - Args: - sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete): - The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability - distributions for the unnoised latent pixels. - """ - - sample: "torch.Tensor" # noqa: F821 diff --git a/venv/lib/python3.10/site-packages/diffusers/models/modeling_pytorch_flax_utils.py b/venv/lib/python3.10/site-packages/diffusers/models/modeling_pytorch_flax_utils.py deleted file mode 100644 index ada55073dd5555084913b1bffbe747329509a182..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/modeling_pytorch_flax_utils.py +++ /dev/null @@ -1,161 +0,0 @@ -# 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. -"""PyTorch - Flax general utilities.""" - -from pickle import UnpicklingError - -import jax -import jax.numpy as jnp -import numpy as np -from flax.serialization import from_bytes -from flax.traverse_util import flatten_dict - -from ..utils import logging - - -logger = logging.get_logger(__name__) - - -##################### -# Flax => PyTorch # -##################### - - -# from https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_flax_pytorch_utils.py#L224-L352 -def load_flax_checkpoint_in_pytorch_model(pt_model, model_file): - try: - with open(model_file, "rb") as flax_state_f: - flax_state = from_bytes(None, flax_state_f.read()) - except UnpicklingError as e: - try: - with open(model_file) as f: - if f.read().startswith("version"): - raise OSError( - "You seem to have cloned a repository without having git-lfs installed. Please" - " install git-lfs and run `git lfs install` followed by `git lfs pull` in the" - " folder you cloned." - ) - else: - raise ValueError from e - except (UnicodeDecodeError, ValueError): - raise EnvironmentError(f"Unable to convert {model_file} to Flax deserializable object. ") - - return load_flax_weights_in_pytorch_model(pt_model, flax_state) - - -def load_flax_weights_in_pytorch_model(pt_model, flax_state): - """Load flax checkpoints in a PyTorch model""" - - try: - import torch # noqa: F401 - except ImportError: - logger.error( - "Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see" - " https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation" - " instructions." - ) - raise - - # check if we have bf16 weights - is_type_bf16 = flatten_dict(jax.tree_util.tree_map(lambda x: x.dtype == jnp.bfloat16, flax_state)).values() - if any(is_type_bf16): - # convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16 - - # and bf16 is not fully supported in PT yet. - logger.warning( - "Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` " - "before loading those in PyTorch model." - ) - flax_state = jax.tree_util.tree_map( - lambda params: params.astype(np.float32) if params.dtype == jnp.bfloat16 else params, flax_state - ) - - pt_model.base_model_prefix = "" - - flax_state_dict = flatten_dict(flax_state, sep=".") - pt_model_dict = pt_model.state_dict() - - # keep track of unexpected & missing keys - unexpected_keys = [] - missing_keys = set(pt_model_dict.keys()) - - for flax_key_tuple, flax_tensor in flax_state_dict.items(): - flax_key_tuple_array = flax_key_tuple.split(".") - - if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4: - flax_key_tuple_array = flax_key_tuple_array[:-1] + ["weight"] - flax_tensor = jnp.transpose(flax_tensor, (3, 2, 0, 1)) - elif flax_key_tuple_array[-1] == "kernel": - flax_key_tuple_array = flax_key_tuple_array[:-1] + ["weight"] - flax_tensor = flax_tensor.T - elif flax_key_tuple_array[-1] == "scale": - flax_key_tuple_array = flax_key_tuple_array[:-1] + ["weight"] - - if "time_embedding" not in flax_key_tuple_array: - for i, flax_key_tuple_string in enumerate(flax_key_tuple_array): - flax_key_tuple_array[i] = ( - flax_key_tuple_string.replace("_0", ".0") - .replace("_1", ".1") - .replace("_2", ".2") - .replace("_3", ".3") - .replace("_4", ".4") - .replace("_5", ".5") - .replace("_6", ".6") - .replace("_7", ".7") - .replace("_8", ".8") - .replace("_9", ".9") - ) - - flax_key = ".".join(flax_key_tuple_array) - - if flax_key in pt_model_dict: - if flax_tensor.shape != pt_model_dict[flax_key].shape: - raise ValueError( - f"Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected " - f"to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}." - ) - else: - # add weight to pytorch dict - flax_tensor = np.asarray(flax_tensor) if not isinstance(flax_tensor, np.ndarray) else flax_tensor - pt_model_dict[flax_key] = torch.from_numpy(flax_tensor) - # remove from missing keys - missing_keys.remove(flax_key) - else: - # weight is not expected by PyTorch model - unexpected_keys.append(flax_key) - - pt_model.load_state_dict(pt_model_dict) - - # re-transform missing_keys to list - missing_keys = list(missing_keys) - - if len(unexpected_keys) > 0: - logger.warning( - "Some weights of the Flax model were not used when initializing the PyTorch model" - f" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing" - f" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture" - " (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This" - f" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect" - " to be exactly identical (e.g. initializing a BertForSequenceClassification model from a" - " FlaxBertForSequenceClassification model)." - ) - if len(missing_keys) > 0: - logger.warning( - f"Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly" - f" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to" - " use it for predictions and inference." - ) - - return pt_model diff --git a/venv/lib/python3.10/site-packages/diffusers/models/modeling_utils.py b/venv/lib/python3.10/site-packages/diffusers/models/modeling_utils.py deleted file mode 100644 index 5fa04fb2606f46ff1edf099ba1af7d7588ca5a0c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/modeling_utils.py +++ /dev/null @@ -1,1861 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# Copyright (c) 2022, NVIDIA CORPORATION. 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 -import itertools -import json -import os -import re -import shutil -import tempfile -from collections import OrderedDict -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 - -import safetensors -import torch -import torch.utils.checkpoint -from huggingface_hub import DDUFEntry, create_repo, split_torch_state_dict_into_shards -from huggingface_hub.utils import validate_hf_hub_args -from torch import Tensor, nn -from typing_extensions import Self - -from .. import __version__ -from ..quantizers import DiffusersAutoQuantizer, DiffusersQuantizer -from ..quantizers.quantization_config import QuantizationMethod -from ..utils import ( - CONFIG_NAME, - FLAX_WEIGHTS_NAME, - SAFE_WEIGHTS_INDEX_NAME, - SAFETENSORS_WEIGHTS_NAME, - WEIGHTS_INDEX_NAME, - WEIGHTS_NAME, - _add_variant, - _get_checkpoint_shard_files, - _get_model_file, - deprecate, - is_accelerate_available, - is_bitsandbytes_available, - is_bitsandbytes_version, - is_peft_available, - is_torch_version, - logging, -) -from ..utils.hub_utils import ( - PushToHubMixin, - load_or_create_model_card, - populate_model_card, -) -from .model_loading_utils import ( - _determine_device_map, - _fetch_index_file, - _fetch_index_file_legacy, - _load_state_dict_into_model, - load_model_dict_into_meta, - load_state_dict, -) - - -class ContextManagers: - """ - Wrapper for `contextlib.ExitStack` which enters a collection of context managers. Adaptation of `ContextManagers` - in the `fastcore` library. - """ - - def __init__(self, context_managers: List[ContextManager]): - self.context_managers = context_managers - self.stack = ExitStack() - - def __enter__(self): - for context_manager in self.context_managers: - self.stack.enter_context(context_manager) - - def __exit__(self, *args, **kwargs): - self.stack.__exit__(*args, **kwargs) - - -logger = logging.get_logger(__name__) - -_REGEX_SHARD = re.compile(r"(.*?)-\d{5}-of-\d{5}") - -TORCH_INIT_FUNCTIONS = { - "uniform_": nn.init.uniform_, - "normal_": nn.init.normal_, - "trunc_normal_": nn.init.trunc_normal_, - "constant_": nn.init.constant_, - "xavier_uniform_": nn.init.xavier_uniform_, - "xavier_normal_": nn.init.xavier_normal_, - "kaiming_uniform_": nn.init.kaiming_uniform_, - "kaiming_normal_": nn.init.kaiming_normal_, - "uniform": nn.init.uniform, - "normal": nn.init.normal, - "xavier_uniform": nn.init.xavier_uniform, - "xavier_normal": nn.init.xavier_normal, - "kaiming_uniform": nn.init.kaiming_uniform, - "kaiming_normal": nn.init.kaiming_normal, -} - -if is_torch_version(">=", "1.9.0"): - _LOW_CPU_MEM_USAGE_DEFAULT = True -else: - _LOW_CPU_MEM_USAGE_DEFAULT = False - - -if is_accelerate_available(): - import accelerate - from accelerate import dispatch_model - from accelerate.utils import load_offloaded_weights, save_offload_index - - -def get_parameter_device(parameter: torch.nn.Module) -> torch.device: - from ..hooks.group_offloading import _get_group_onload_device - - try: - # Try to get the onload device from the group offloading hook - return _get_group_onload_device(parameter) - except ValueError: - pass - - try: - # If the onload device is not available due to no group offloading hooks, try to get the device - # from the first parameter or buffer - parameters_and_buffers = itertools.chain(parameter.parameters(), parameter.buffers()) - return next(parameters_and_buffers).device - except StopIteration: - # For torch.nn.DataParallel compatibility in PyTorch 1.5 - - 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 - - gen = parameter._named_members(get_members_fn=find_tensor_attributes) - first_tuple = next(gen) - return first_tuple[1].device - - -def get_parameter_dtype(parameter: torch.nn.Module) -> torch.dtype: - """ - Returns the first found floating dtype in parameters if there is one, otherwise returns the last dtype it found. - """ - # 1. Check if we have attached any dtype modifying hooks (eg. layerwise casting) - if isinstance(parameter, nn.Module): - for name, submodule in parameter.named_modules(): - if not hasattr(submodule, "_diffusers_hook"): - continue - registry = submodule._diffusers_hook - hook = registry.get_hook("layerwise_casting") - if hook is not None: - return hook.compute_dtype - - # 2. If no dtype modifying hooks are attached, return the dtype of the first floating point parameter/buffer - last_dtype = None - - for name, param in parameter.named_parameters(): - last_dtype = param.dtype - if parameter._keep_in_fp32_modules and any(m in name for m in parameter._keep_in_fp32_modules): - continue - - if param.is_floating_point(): - return param.dtype - - for buffer in parameter.buffers(): - last_dtype = buffer.dtype - if buffer.is_floating_point(): - return buffer.dtype - - if last_dtype is not None: - # if no floating dtype was found return whatever the first dtype is - return last_dtype - - # For nn.DataParallel compatibility in PyTorch > 1.5 - 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 - - gen = parameter._named_members(get_members_fn=find_tensor_attributes) - last_tuple = None - for tuple in gen: - last_tuple = tuple - if tuple[1].is_floating_point(): - return tuple[1].dtype - - if last_tuple is not None: - # fallback to the last dtype - return last_tuple[1].dtype - - -def check_support_param_buffer_assignment(model_to_load, state_dict, start_prefix=""): - """ - Checks if `model_to_load` supports param buffer assignment (such as when loading in empty weights) by first - checking if the model explicitly disables it, then by ensuring that the state dict keys are a subset of the model's - parameters. - - """ - if model_to_load.device.type == "meta": - return False - - if len([key for key in state_dict if key.startswith(start_prefix)]) == 0: - return False - - # Some models explicitly do not support param buffer assignment - if not getattr(model_to_load, "_supports_param_buffer_assignment", True): - logger.debug( - f"{model_to_load.__class__.__name__} does not support param buffer assignment, loading will be slower" - ) - return False - - # If the model does, the incoming `state_dict` and the `model_to_load` must be the same dtype - first_key = next(iter(model_to_load.state_dict().keys())) - if start_prefix + first_key in state_dict: - return state_dict[start_prefix + first_key].dtype == model_to_load.state_dict()[first_key].dtype - - return False - - -@contextmanager -def no_init_weights(): - """ - Context manager to globally disable weight initialization to speed up loading large models. To do that, all the - torch.nn.init function are all replaced with skip. - """ - - def _skip_init(*args, **kwargs): - pass - - for name, init_func in TORCH_INIT_FUNCTIONS.items(): - setattr(torch.nn.init, name, _skip_init) - try: - yield - finally: - # Restore the original initialization functions - for name, init_func in TORCH_INIT_FUNCTIONS.items(): - setattr(torch.nn.init, name, init_func) - - -class ModelMixin(torch.nn.Module, PushToHubMixin): - r""" - Base class for all models. - - [`ModelMixin`] takes care of storing the model configuration and provides methods for loading, downloading and - saving models. - - - **config_name** ([`str`]) -- Filename to save a model to when calling [`~models.ModelMixin.save_pretrained`]. - """ - - config_name = CONFIG_NAME - _automatically_saved_args = ["_diffusers_version", "_class_name", "_name_or_path"] - _supports_gradient_checkpointing = False - _keys_to_ignore_on_load_unexpected = None - _no_split_modules = None - _keep_in_fp32_modules = None - _skip_layerwise_casting_patterns = None - _supports_group_offloading = True - - def __init__(self): - super().__init__() - - self._gradient_checkpointing_func = None - - def __getattr__(self, name: str) -> Any: - """The only reason we overwrite `getattr` here is to gracefully deprecate accessing - config attributes directly. See https://github.com/huggingface/diffusers/pull/3129 We need to overwrite - __getattr__ here in addition so that we don't trigger `torch.nn.Module`'s __getattr__': - https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module - """ - - is_in_config = "_internal_dict" in self.__dict__ and hasattr(self.__dict__["_internal_dict"], name) - is_attribute = name in self.__dict__ - - if is_in_config and not is_attribute: - deprecation_message = f"Accessing config attribute `{name}` directly via '{type(self).__name__}' object attribute is deprecated. Please access '{name}' over '{type(self).__name__}'s config object instead, e.g. 'unet.config.{name}'." - deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False, stacklevel=3) - return self._internal_dict[name] - - # call PyTorch's https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module - return super().__getattr__(name) - - @property - def is_gradient_checkpointing(self) -> bool: - """ - Whether gradient checkpointing is activated for this model or not. - """ - 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: - """ - Activates gradient checkpointing for the current model (may be referred to as *activation checkpointing* or - *checkpoint activations* in other frameworks). - - Args: - gradient_checkpointing_func (`Callable`, *optional*): - The function to use for gradient checkpointing. If `None`, the default PyTorch checkpointing function - is used (`torch.utils.checkpoint.checkpoint`). - """ - if not self._supports_gradient_checkpointing: - raise ValueError( - f"{self.__class__.__name__} does not support gradient checkpointing. Please make sure to set the boolean attribute " - f"`_supports_gradient_checkpointing` to `True` in the class definition." - ) - - if gradient_checkpointing_func is None: - - def _gradient_checkpointing_func(module, *args): - ckpt_kwargs = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} - return torch.utils.checkpoint.checkpoint( - module.__call__, - *args, - **ckpt_kwargs, - ) - - gradient_checkpointing_func = _gradient_checkpointing_func - - self._set_gradient_checkpointing(enable=True, gradient_checkpointing_func=gradient_checkpointing_func) - - def disable_gradient_checkpointing(self) -> None: - """ - Deactivates gradient checkpointing for the current model (may be referred to as *activation checkpointing* or - *checkpoint activations* in other frameworks). - """ - if self._supports_gradient_checkpointing: - self._set_gradient_checkpointing(enable=False) - - def set_use_npu_flash_attention(self, valid: bool) -> None: - r""" - Set the switch for the npu flash attention. - """ - - def fn_recursive_set_npu_flash_attention(module: torch.nn.Module): - if hasattr(module, "set_use_npu_flash_attention"): - module.set_use_npu_flash_attention(valid) - - for child in module.children(): - fn_recursive_set_npu_flash_attention(child) - - for module in self.children(): - if isinstance(module, torch.nn.Module): - fn_recursive_set_npu_flash_attention(module) - - def enable_npu_flash_attention(self) -> None: - r""" - Enable npu flash attention from torch_npu - - """ - self.set_use_npu_flash_attention(True) - - def disable_npu_flash_attention(self) -> None: - r""" - disable npu flash attention from torch_npu - - """ - 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 - ) -> None: - # Recursively walk through all the children. - # Any children which exposes the set_use_xla_flash_attention method - # gets the message - def fn_recursive_set_flash_attention(module: torch.nn.Module): - if hasattr(module, "set_use_xla_flash_attention"): - module.set_use_xla_flash_attention(use_xla_flash_attention, partition_spec, **kwargs) - - for child in module.children(): - fn_recursive_set_flash_attention(child) - - for module in self.children(): - if isinstance(module, torch.nn.Module): - fn_recursive_set_flash_attention(module) - - def enable_xla_flash_attention(self, partition_spec: Optional[Callable] = None, **kwargs): - r""" - Enable the flash attention pallals kernel for torch_xla. - """ - self.set_use_xla_flash_attention(True, partition_spec, **kwargs) - - def disable_xla_flash_attention(self): - r""" - Disable the flash attention pallals kernel for torch_xla. - """ - self.set_use_xla_flash_attention(False) - - def set_use_memory_efficient_attention_xformers( - self, valid: bool, attention_op: Optional[Callable] = None - ) -> None: - # Recursively walk through all the children. - # Any children which exposes the set_use_memory_efficient_attention_xformers method - # gets the message - def fn_recursive_set_mem_eff(module: torch.nn.Module): - if hasattr(module, "set_use_memory_efficient_attention_xformers"): - module.set_use_memory_efficient_attention_xformers(valid, attention_op) - - for child in module.children(): - fn_recursive_set_mem_eff(child) - - for module in self.children(): - 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: - 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 up during training is not guaranteed. - - - - ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient attention takes - precedent. - - - - Parameters: - attention_op (`Callable`, *optional*): - Override the default `None` operator for use as `op` argument to the - [`memory_efficient_attention()`](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.memory_efficient_attention) - function of xFormers. - - Examples: - - ```py - >>> import torch - >>> from diffusers import UNet2DConditionModel - >>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp - - >>> model = UNet2DConditionModel.from_pretrained( - ... "stabilityai/stable-diffusion-2-1", subfolder="unet", torch_dtype=torch.float16 - ... ) - >>> model = model.to("cuda") - >>> model.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp) - ``` - """ - self.set_use_memory_efficient_attention_xformers(True, attention_op) - - def disable_xformers_memory_efficient_attention(self) -> None: - r""" - Disable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). - """ - self.set_use_memory_efficient_attention_xformers(False) - - 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, - non_blocking: bool = False, - ) -> None: - r""" - Activates layerwise casting for the current model. - - Layerwise casting is a technique that casts the model weights to a lower precision dtype for storage but - upcasts them on-the-fly to a higher precision dtype for computation. This process can significantly reduce the - memory footprint from model weights, but may lead to some quality degradation in the outputs. Most degradations - are negligible, mostly stemming from weight casting in normalization and modulation layers. - - By default, most models in diffusers set the `_skip_layerwise_casting_patterns` attribute to ignore patch - embedding, positional embedding and normalization layers. This is because these layers are most likely - precision-critical for quality. If you wish to change this behavior, you can set the - `_skip_layerwise_casting_patterns` attribute to `None`, or call - [`~hooks.layerwise_casting.apply_layerwise_casting`] with custom arguments. - - Example: - Using [`~models.ModelMixin.enable_layerwise_casting`]: - - ```python - >>> from diffusers import CogVideoXTransformer3DModel - - >>> transformer = CogVideoXTransformer3DModel.from_pretrained( - ... "THUDM/CogVideoX-5b", subfolder="transformer", torch_dtype=torch.bfloat16 - ... ) - - >>> # Enable layerwise casting via the model, which ignores certain modules by default - >>> transformer.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16) - ``` - - Args: - storage_dtype (`torch.dtype`): - 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*): - 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*): - 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. - """ - from ..hooks import apply_layerwise_casting - - user_provided_patterns = True - if skip_modules_pattern is None: - from ..hooks.layerwise_casting import DEFAULT_SKIP_MODULES_PATTERN - - skip_modules_pattern = DEFAULT_SKIP_MODULES_PATTERN - user_provided_patterns = False - if self._keep_in_fp32_modules is not None: - skip_modules_pattern += tuple(self._keep_in_fp32_modules) - if self._skip_layerwise_casting_patterns is not None: - skip_modules_pattern += tuple(self._skip_layerwise_casting_patterns) - skip_modules_pattern = tuple(set(skip_modules_pattern)) - - if is_peft_available() and not user_provided_patterns: - # By default, we want to skip all peft layers because they have a very low memory footprint. - # If users want to apply layerwise casting on peft layers as well, they can utilize the - # `~diffusers.hooks.layerwise_casting.apply_layerwise_casting` function which provides - # them with more flexibility and control. - - from peft.tuners.loha.layer import LoHaLayer - from peft.tuners.lokr.layer import LoKrLayer - from peft.tuners.lora.layer import LoraLayer - - for layer in (LoHaLayer, LoKrLayer, LoraLayer): - skip_modules_pattern += tuple(layer.adapter_layer_names) - - if compute_dtype is None: - logger.info("`compute_dtype` not provided when enabling layerwise casting. Using dtype of the model.") - compute_dtype = self.dtype - - apply_layerwise_casting( - self, storage_dtype, compute_dtype, skip_modules_pattern, skip_modules_classes, non_blocking - ) - - def enable_group_offload( - self, - onload_device: torch.device, - offload_device: torch.device = torch.device("cpu"), - offload_type: str = "block_level", - num_blocks_per_group: Optional[int] = 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, - ) -> None: - r""" - Activates group offloading for the current model. - - See [`~hooks.group_offloading.apply_group_offloading`] for more information. - - Example: - - ```python - >>> from diffusers import CogVideoXTransformer3DModel - - >>> transformer = CogVideoXTransformer3DModel.from_pretrained( - ... "THUDM/CogVideoX-5b", subfolder="transformer", torch_dtype=torch.bfloat16 - ... ) - - >>> transformer.enable_group_offload( - ... onload_device=torch.device("cuda"), - ... offload_device=torch.device("cpu"), - ... offload_type="leaf_level", - ... use_stream=True, - ... ) - ``` - """ - from ..hooks import apply_group_offloading - - if getattr(self, "enable_tiling", None) is not None and getattr(self, "use_tiling", False) and use_stream: - msg = ( - "Applying group offloading on autoencoders, with CUDA streams, may not work as expected if the first " - "forward pass is executed with tiling enabled. Please make sure to either:\n" - "1. Run a forward pass with small input shapes.\n" - "2. Or, run a forward pass with tiling disabled (can still use small dummy inputs)." - ) - logger.warning(msg) - if not self._supports_group_offloading: - raise ValueError( - f"{self.__class__.__name__} does not support group offloading. Please make sure to set the boolean attribute " - f"`_supports_group_offloading` to `True` in the class definition. If you believe this is a mistake, please " - f"open an issue at https://github.com/huggingface/diffusers/issues." - ) - apply_group_offloading( - module=self, - onload_device=onload_device, - offload_device=offload_device, - offload_type=offload_type, - num_blocks_per_group=num_blocks_per_group, - non_blocking=non_blocking, - use_stream=use_stream, - record_stream=record_stream, - low_cpu_mem_usage=low_cpu_mem_usage, - offload_to_disk_path=offload_to_disk_path, - ) - - def save_pretrained( - self, - save_directory: Union[str, os.PathLike], - is_main_process: bool = True, - save_function: Optional[Callable] = None, - safe_serialization: bool = True, - variant: Optional[str] = None, - max_shard_size: Union[int, str] = "10GB", - push_to_hub: bool = False, - **kwargs, - ): - """ - Save a model and its configuration file to a directory so that it can be reloaded using the - [`~models.ModelMixin.from_pretrained`] class method. - - 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. - 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 - need to call this function on all processes. In this case, set `is_main_process=True` only on the main - process to avoid race conditions. - save_function (`Callable`): - The function to use to save the state dictionary. Useful during distributed training when you need to - replace `torch.save` with another method. Can be configured with the environment variable - `DIFFUSERS_SAVE_MODE`. - 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 `"10GB"`): - 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. Note that this limit will be decreased after a certain - period of time (starting from Oct 2024) to allow users to upgrade to the latest version of `diffusers`. - This is to establish a common default size for this argument across different libraries in the Hugging - Face ecosystem (`transformers`, and `accelerate`, for example). - push_to_hub (`bool`, *optional*, defaults to `False`): - 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*): - Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. - """ - if os.path.isfile(save_directory): - logger.error(f"Provided path ({save_directory}) should be a directory, not a file") - return - - hf_quantizer = getattr(self, "hf_quantizer", None) - if hf_quantizer is not None: - quantization_serializable = ( - hf_quantizer is not None - and isinstance(hf_quantizer, DiffusersQuantizer) - and hf_quantizer.is_serializable - ) - if not quantization_serializable: - raise ValueError( - f"The model is quantized with {hf_quantizer.quantization_config.quant_method} and is not serializable - check out the warnings from" - " the logger on the traceback to understand the reason why the quantized model is not serializable." - ) - - weights_name = SAFETENSORS_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME - weights_name = _add_variant(weights_name, variant) - weights_name_pattern = weights_name.replace(".bin", "{suffix}.bin").replace( - ".safetensors", "{suffix}.safetensors" - ) - - os.makedirs(save_directory, exist_ok=True) - - 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]) - repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id - - # Only save the model itself if we are using distributed training - model_to_save = self - - # Attach architecture to the config - # Save the config - if is_main_process: - model_to_save.save_config(save_directory) - - # Save the model - state_dict = model_to_save.state_dict() - - # Save the model - state_dict_split = split_torch_state_dict_into_shards( - state_dict, max_shard_size=max_shard_size, filename_pattern=weights_name_pattern - ) - - # Clean the folder from a previous save - if is_main_process: - for filename in os.listdir(save_directory): - if filename in state_dict_split.filename_to_tensors.keys(): - continue - full_filename = os.path.join(save_directory, filename) - if not os.path.isfile(full_filename): - continue - weights_without_ext = weights_name_pattern.replace(".bin", "").replace(".safetensors", "") - weights_without_ext = weights_without_ext.replace("{suffix}", "") - filename_without_ext = filename.replace(".bin", "").replace(".safetensors", "") - # make sure that file to be deleted matches format of sharded file, e.g. pytorch_model-00001-of-00005 - if ( - filename.startswith(weights_without_ext) - and _REGEX_SHARD.fullmatch(filename_without_ext) is not None - ): - os.remove(full_filename) - - for filename, tensors in state_dict_split.filename_to_tensors.items(): - shard = {tensor: state_dict[tensor].contiguous() for tensor in tensors} - filepath = os.path.join(save_directory, filename) - if safe_serialization: - # At some point we will need to deal better with save_function (used for TPU and other distributed - # joyfulness), but for now this enough. - safetensors.torch.save_file(shard, filepath, metadata={"format": "pt"}) - else: - torch.save(shard, filepath) - - if state_dict_split.is_sharded: - index = { - "metadata": state_dict_split.metadata, - "weight_map": state_dict_split.tensor_to_filename, - } - save_index_file = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME - save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant)) - # Save the index as well - with open(save_index_file, "w", encoding="utf-8") as f: - content = json.dumps(index, indent=2, sort_keys=True) + "\n" - f.write(content) - logger.info( - f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be " - f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the " - f"index located at {save_index_file}." - ) - else: - path_to_weights = os.path.join(save_directory, weights_name) - logger.info(f"Model weights saved in {path_to_weights}") - - if push_to_hub: - # Create a new empty model card and eventually tag it - model_card = load_or_create_model_card(repo_id, token=token) - model_card = populate_model_card(model_card) - model_card.save(Path(save_directory, "README.md").as_posix()) - - self._upload_folder( - save_directory, - repo_id, - token=token, - commit_message=commit_message, - create_pr=create_pr, - ) - - def dequantize(self): - """ - Potentially dequantize the model in case it has been quantized by a quantization method that support - dequantization. - """ - hf_quantizer = getattr(self, "hf_quantizer", None) - - if hf_quantizer is None: - raise ValueError("You need to first quantize your model in order to dequantize it") - - return hf_quantizer.dequantize(self) - - @classmethod - @validate_hf_hub_args - def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs) -> Self: - r""" - Instantiate a pretrained PyTorch model from a pretrained model configuration. - - The model is set in evaluation mode - `model.eval()` - by default, and dropout modules are deactivated. To - train the model, set it back in training mode with `model.train()`. - - Parameters: - pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model 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*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - torch_dtype (`torch.dtype`, *optional*): - Override the default `torch.dtype` and load the model with another dtype. - 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*): - 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`): - Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - local_files_only(`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - from_flax (`bool`, *optional*, defaults to `False`): - Load the model weights from a Flax checkpoint save file. - subfolder (`str`, *optional*, defaults to `""`): - The subfolder location of a model file within a larger model repository on the Hub or locally. - mirror (`str`, *optional*): - 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*): - 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. - - Examples: - - ```py - >>> from diffusers import AutoModel - >>> import torch - - >>> # This works. - >>> model = AutoModel.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", device_map="cuda" - ... ) - >>> # This also works (integer accelerator device ID). - >>> model = AutoModel.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", device_map=0 - ... ) - >>> # Specifying a supported offloading strategy like "auto" also works. - >>> model = AutoModel.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", device_map="auto" - ... ) - >>> # Specifying a dictionary as `device_map` also works. - >>> model = AutoModel.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", - ... subfolder="unet", - ... device_map={"": torch.device("cuda")}, - ... ) - ``` - - Set `device_map="auto"` to have 🤗 Accelerate automatically compute the most optimized `device_map`. For - more information about each option see [designing a device - map](https://huggingface.co/docs/accelerate/en/concept_guides/big_model_inference#the-devicemap). You - can also refer to the [Diffusers-specific - documentation](https://huggingface.co/docs/diffusers/main/en/training/distributed_inference#model-sharding) - for more concrete examples. - max_memory (`Dict`, *optional*): - A dictionary device identifier for the maximum memory. Will default to the maximum memory available for - each GPU and the available CPU RAM if unset. - offload_folder (`str` or `os.PathLike`, *optional*): - The path to offload weights if `device_map` contains the value `"disk"`. - offload_state_dict (`bool`, *optional*): - If `True`, temporarily offloads the CPU state dict to the hard drive to avoid running out of CPU RAM if - the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to `True` - when there is some disk offload. - low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): - Speed up model loading only loading the pretrained weights and not initializing the weights. This also - tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. - Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this - argument to `True` will raise an error. - variant (`str`, *optional*): - Load weights from a specified `variant` filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. - use_safetensors (`bool`, *optional*, defaults to `None`): - If set to `None`, the `safetensors` weights are downloaded if they're available **and** if the - `safetensors` library is installed. If set to `True`, the model is forcibly loaded from `safetensors` - weights. If set to `False`, `safetensors` weights are not loaded. - 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. - - - - To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with - `huggingface-cli login`. You can also activate the special - ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a - firewalled environment. - - - - Example: - - ```py - from diffusers import UNet2DConditionModel - - unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet") - ``` - - If you get the error message below, you need to finetune the weights for your downstream task: - - ```bash - Some weights of UNet2DConditionModel were not initialized from the model checkpoint at runwayml/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: - - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated - You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. - ``` - """ - cache_dir = kwargs.pop("cache_dir", None) - ignore_mismatched_sizes = kwargs.pop("ignore_mismatched_sizes", False) - force_download = kwargs.pop("force_download", False) - from_flax = kwargs.pop("from_flax", False) - proxies = kwargs.pop("proxies", None) - output_loading_info = kwargs.pop("output_loading_info", False) - local_files_only = kwargs.pop("local_files_only", None) - token = kwargs.pop("token", None) - revision = kwargs.pop("revision", None) - torch_dtype = kwargs.pop("torch_dtype", None) - subfolder = kwargs.pop("subfolder", None) - device_map = kwargs.pop("device_map", None) - max_memory = kwargs.pop("max_memory", None) - offload_folder = kwargs.pop("offload_folder", None) - offload_state_dict = kwargs.pop("offload_state_dict", None) - low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) - 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) - disable_mmap = kwargs.pop("disable_mmap", False) - - if torch_dtype is not None and not isinstance(torch_dtype, torch.dtype): - torch_dtype = torch.float32 - logger.warning( - f"Passed `torch_dtype` {torch_dtype} is not a `torch.dtype`. Defaulting to `torch.float32`." - ) - - allow_pickle = False - if use_safetensors is None: - use_safetensors = True - allow_pickle = True - - if low_cpu_mem_usage and not is_accelerate_available(): - low_cpu_mem_usage = False - logger.warning( - "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the" - " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install" - " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip" - " install accelerate\n```\n." - ) - - if device_map is not None and not is_accelerate_available(): - raise NotImplementedError( - "Loading and dispatching requires `accelerate`. Please make sure to install accelerate or set" - " `device_map=None`. You can install accelerate with `pip install accelerate`." - ) - - # Check if we can handle device_map and dispatching the weights - if device_map is not None and not is_torch_version(">=", "1.9.0"): - raise NotImplementedError( - "Loading and dispatching requires torch >= 1.9.0. Please either update your PyTorch version or set" - " `device_map=None`." - ) - - if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"): - raise NotImplementedError( - "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set" - " `low_cpu_mem_usage=False`." - ) - - if low_cpu_mem_usage is False and device_map is not None: - raise ValueError( - f"You cannot set `low_cpu_mem_usage` to `False` while using device_map={device_map} for loading and" - " dispatching. Please make sure to set `low_cpu_mem_usage=True`." - ) - - # change device_map into a map if we passed an int, a str or a torch.device - if isinstance(device_map, torch.device): - device_map = {"": device_map} - elif isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]: - try: - device_map = {"": torch.device(device_map)} - except RuntimeError: - raise ValueError( - "When passing device_map as a string, the value needs to be a device name (e.g. cpu, cuda:0) or " - f"'auto', 'balanced', 'balanced_low_0', 'sequential' but found {device_map}." - ) - elif isinstance(device_map, int): - if device_map < 0: - raise ValueError( - "You can't pass device_map as a negative int. If you want to put the model on the cpu, pass device_map = 'cpu' " - ) - else: - device_map = {"": device_map} - - if device_map is not None: - if low_cpu_mem_usage is None: - low_cpu_mem_usage = True - elif not low_cpu_mem_usage: - raise ValueError("Passing along a `device_map` requires `low_cpu_mem_usage=True`") - - if low_cpu_mem_usage: - if device_map is not None and not is_torch_version(">=", "1.10"): - # The max memory utils require PyTorch >= 1.10 to have torch.cuda.mem_get_info. - raise ValueError("`low_cpu_mem_usage` and `device_map` require PyTorch >= 1.10.") - - user_agent = { - "diffusers": __version__, - "file_type": "model", - "framework": "pytorch", - } - unused_kwargs = {} - - # Load config if we don't provide a configuration - config_path = pretrained_model_name_or_path - - # load config - config, unused_kwargs, commit_hash = cls.load_config( - config_path, - cache_dir=cache_dir, - return_unused_kwargs=True, - return_commit_hash=True, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - dduf_entries=dduf_entries, - **kwargs, - ) - # no in-place modification of the original config. - config = copy.deepcopy(config) - - # determine initial quantization config. - ####################################### - pre_quantized = "quantization_config" in config and config["quantization_config"] is not None - if pre_quantized or quantization_config is not None: - if pre_quantized: - config["quantization_config"] = DiffusersAutoQuantizer.merge_quantization_configs( - config["quantization_config"], quantization_config - ) - else: - config["quantization_config"] = quantization_config - hf_quantizer = DiffusersAutoQuantizer.from_config( - config["quantization_config"], pre_quantized=pre_quantized - ) - else: - hf_quantizer = None - - if hf_quantizer is not None: - hf_quantizer.validate_environment(torch_dtype=torch_dtype, from_flax=from_flax, device_map=device_map) - torch_dtype = hf_quantizer.update_torch_dtype(torch_dtype) - device_map = hf_quantizer.update_device_map(device_map) - - # In order to ensure popular quantization methods are supported. Can be disable with `disable_telemetry` - user_agent["quant"] = hf_quantizer.quantization_config.quant_method.value - - # Force-set to `True` for more mem efficiency - if low_cpu_mem_usage is None: - low_cpu_mem_usage = True - logger.info("Set `low_cpu_mem_usage` to True as `hf_quantizer` is not None.") - elif not low_cpu_mem_usage: - raise ValueError("`low_cpu_mem_usage` cannot be False or None when using quantization.") - - # Check if `_keep_in_fp32_modules` is not None - use_keep_in_fp32_modules = cls._keep_in_fp32_modules is not None and ( - hf_quantizer is None or getattr(hf_quantizer, "use_keep_in_fp32_modules", False) - ) - - if use_keep_in_fp32_modules: - keep_in_fp32_modules = cls._keep_in_fp32_modules - if not isinstance(keep_in_fp32_modules, list): - keep_in_fp32_modules = [keep_in_fp32_modules] - - if low_cpu_mem_usage is None: - low_cpu_mem_usage = True - logger.info("Set `low_cpu_mem_usage` to True as `_keep_in_fp32_modules` is not None.") - elif not low_cpu_mem_usage: - raise ValueError("`low_cpu_mem_usage` cannot be False when `keep_in_fp32_modules` is True.") - else: - keep_in_fp32_modules = [] - - is_sharded = False - resolved_model_file = None - - # Determine if we're loading from a directory of sharded checkpoints. - sharded_metadata = None - index_file = None - is_local = os.path.isdir(pretrained_model_name_or_path) - index_file_kwargs = { - "is_local": is_local, - "pretrained_model_name_or_path": pretrained_model_name_or_path, - "subfolder": subfolder or "", - "use_safetensors": use_safetensors, - "cache_dir": cache_dir, - "variant": variant, - "force_download": force_download, - "proxies": proxies, - "local_files_only": local_files_only, - "token": token, - "revision": revision, - "user_agent": user_agent, - "commit_hash": commit_hash, - "dduf_entries": dduf_entries, - } - index_file = _fetch_index_file(**index_file_kwargs) - # In case the index file was not found we still have to consider the legacy format. - # this becomes applicable when the variant is not None. - if variant is not None and (index_file is None or not os.path.exists(index_file)): - index_file = _fetch_index_file_legacy(**index_file_kwargs) - if index_file is not None and (dduf_entries or index_file.is_file()): - is_sharded = True - - if is_sharded and from_flax: - raise ValueError("Loading of sharded checkpoints is not supported when `from_flax=True`.") - - # load model - if from_flax: - resolved_model_file = _get_model_file( - pretrained_model_name_or_path, - weights_name=FLAX_WEIGHTS_NAME, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - commit_hash=commit_hash, - ) - model = cls.from_config(config, **unused_kwargs) - - # Convert the weights - from .modeling_pytorch_flax_utils import load_flax_checkpoint_in_pytorch_model - - model = load_flax_checkpoint_in_pytorch_model(model, resolved_model_file) - else: - # in the case it is sharded, we have already the index - if is_sharded: - resolved_model_file, sharded_metadata = _get_checkpoint_shard_files( - pretrained_model_name_or_path, - index_file, - cache_dir=cache_dir, - proxies=proxies, - local_files_only=local_files_only, - token=token, - user_agent=user_agent, - revision=revision, - subfolder=subfolder or "", - dduf_entries=dduf_entries, - ) - elif use_safetensors: - try: - resolved_model_file = _get_model_file( - pretrained_model_name_or_path, - weights_name=_add_variant(SAFETENSORS_WEIGHTS_NAME, variant), - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - commit_hash=commit_hash, - dduf_entries=dduf_entries, - ) - - except IOError as e: - logger.error(f"An error occurred while trying to fetch {pretrained_model_name_or_path}: {e}") - if not allow_pickle: - raise - logger.warning( - "Defaulting to unsafe serialization. Pass `allow_pickle=False` to raise an error instead." - ) - - if resolved_model_file is None and not is_sharded: - resolved_model_file = _get_model_file( - pretrained_model_name_or_path, - weights_name=_add_variant(WEIGHTS_NAME, variant), - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - commit_hash=commit_hash, - dduf_entries=dduf_entries, - ) - - if not isinstance(resolved_model_file, list): - resolved_model_file = [resolved_model_file] - - # set dtype to instantiate the model under: - # 1. If torch_dtype is not None, we use that dtype - # 2. If torch_dtype is float8, we don't use _set_default_torch_dtype and we downcast after loading the model - dtype_orig = None - if torch_dtype is not None and not torch_dtype == getattr(torch, "float8_e4m3fn", None): - if not isinstance(torch_dtype, torch.dtype): - raise ValueError( - f"{torch_dtype} needs to be of type `torch.dtype`, e.g. `torch.float16`, but is {type(torch_dtype)}." - ) - dtype_orig = cls._set_default_torch_dtype(torch_dtype) - - init_contexts = [no_init_weights()] - - if low_cpu_mem_usage: - init_contexts.append(accelerate.init_empty_weights()) - - with ContextManagers(init_contexts): - model = cls.from_config(config, **unused_kwargs) - - if dtype_orig is not None: - torch.set_default_dtype(dtype_orig) - - state_dict = None - if not is_sharded: - # Time to load the checkpoint - state_dict = load_state_dict(resolved_model_file[0], disable_mmap=disable_mmap, dduf_entries=dduf_entries) - # We only fix it for non sharded checkpoints as we don't need it yet for sharded one. - model._fix_state_dict_keys_on_load(state_dict) - - if is_sharded: - loaded_keys = sharded_metadata["all_checkpoint_keys"] - else: - loaded_keys = list(state_dict.keys()) - - if hf_quantizer is not None: - hf_quantizer.preprocess_model( - model=model, device_map=device_map, keep_in_fp32_modules=keep_in_fp32_modules - ) - - # Now that the model is loaded, we can determine the device_map - device_map = _determine_device_map( - model, device_map, max_memory, torch_dtype, keep_in_fp32_modules, hf_quantizer - ) - if hf_quantizer is not None: - hf_quantizer.validate_environment(device_map=device_map) - - ( - model, - missing_keys, - unexpected_keys, - mismatched_keys, - offload_index, - error_msgs, - ) = cls._load_pretrained_model( - model, - state_dict, - resolved_model_file, - pretrained_model_name_or_path, - loaded_keys, - ignore_mismatched_sizes=ignore_mismatched_sizes, - low_cpu_mem_usage=low_cpu_mem_usage, - device_map=device_map, - offload_folder=offload_folder, - offload_state_dict=offload_state_dict, - dtype=torch_dtype, - hf_quantizer=hf_quantizer, - keep_in_fp32_modules=keep_in_fp32_modules, - dduf_entries=dduf_entries, - ) - loading_info = { - "missing_keys": missing_keys, - "unexpected_keys": unexpected_keys, - "mismatched_keys": mismatched_keys, - "error_msgs": error_msgs, - } - - # Dispatch model with hooks on all devices if necessary - if device_map is not None: - device_map_kwargs = { - "device_map": device_map, - "offload_dir": offload_folder, - "offload_index": offload_index, - } - dispatch_model(model, **device_map_kwargs) - - if hf_quantizer is not None: - hf_quantizer.postprocess_model(model) - model.hf_quantizer = hf_quantizer - - if ( - torch_dtype is not None - and torch_dtype == getattr(torch, "float8_e4m3fn", None) - and hf_quantizer is None - and not use_keep_in_fp32_modules - ): - model = model.to(torch_dtype) - - if hf_quantizer is not None: - # We also make sure to purge `_pre_quantization_dtype` when we serialize - # the model config because `_pre_quantization_dtype` is `torch.dtype`, not JSON serializable. - model.register_to_config(_name_or_path=pretrained_model_name_or_path, _pre_quantization_dtype=torch_dtype) - else: - model.register_to_config(_name_or_path=pretrained_model_name_or_path) - - # Set model in evaluation mode to deactivate DropOut modules by default - model.eval() - - if output_loading_info: - return model, loading_info - - return model - - # Adapted from `transformers`. - @wraps(torch.nn.Module.cuda) - def cuda(self, *args, **kwargs): - from ..hooks.group_offloading import _is_group_offload_enabled - - # 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): - 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." - ) - elif 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." - ) - - # Checks if group offloading is enabled - if _is_group_offload_enabled(self): - logger.warning( - f"The module '{self.__class__.__name__}' is group offloaded and moving it using `.cuda()` is not supported." - ) - return self - - return super().cuda(*args, **kwargs) - - # Adapted from `transformers`. - @wraps(torch.nn.Module.to) - def to(self, *args, **kwargs): - from ..hooks.group_offloading import _is_group_offload_enabled - - device_arg_or_kwarg_present = any(isinstance(arg, torch.device) for arg in args) or "device" in kwargs - dtype_present_in_args = "dtype" in kwargs - - # Try converting arguments to torch.device in case they are passed as strings - for arg in args: - if not isinstance(arg, str): - continue - try: - torch.device(arg) - device_arg_or_kwarg_present = True - except RuntimeError: - pass - - if not dtype_present_in_args: - for arg in args: - if isinstance(arg, torch.dtype): - dtype_present_in_args = True - break - - if getattr(self, "is_quantized", False): - if dtype_present_in_args: - raise ValueError( - "Casting a quantized model to a new `dtype` is unsupported. To set the dtype of unquantized layers, please " - "use the `torch_dtype` argument when loading the model using `from_pretrained` or `from_single_file`" - ) - - if getattr(self, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES: - if getattr(self, "is_loaded_in_8bit", False): - 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`." - ) - elif 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." - ) - return self - - return super().to(*args, **kwargs) - - # Taken from `transformers`. - def half(self, *args): - # Checks if the model is quantized - if getattr(self, "is_quantized", False): - raise ValueError( - "`.half()` is not supported for quantized model. Please use the model as it is, since the" - " model has already been cast to the correct `dtype`." - ) - else: - return super().half(*args) - - # Taken from `transformers`. - def float(self, *args): - # Checks if the model is quantized - if getattr(self, "is_quantized", False): - raise ValueError( - "`.float()` is not supported for quantized model. Please use the model as it is, since the" - " model has already been cast to the correct `dtype`." - ) - else: - return super().float(*args) - - @classmethod - 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], - ignore_mismatched_sizes: bool = False, - assign_to_params_buffers: bool = False, - hf_quantizer: Optional[DiffusersQuantizer] = 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, - ): - model_state_dict = model.state_dict() - expected_keys = list(model_state_dict.keys()) - missing_keys = list(set(expected_keys) - set(loaded_keys)) - if hf_quantizer is not None: - missing_keys = hf_quantizer.update_missing_keys(model, missing_keys, prefix="") - unexpected_keys = list(set(loaded_keys) - set(expected_keys)) - # Some models may have keys that are not in the state by design, removing them before needlessly warning - # the user. - if cls._keys_to_ignore_on_load_unexpected is not None: - for pat in cls._keys_to_ignore_on_load_unexpected: - unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None] - - mismatched_keys = [] - - assign_to_params_buffers = None - error_msgs = [] - - # Deal with offload - if device_map is not None and "disk" in device_map.values(): - if offload_folder is None: - raise ValueError( - "The current `device_map` had weights offloaded to the disk. Please provide an `offload_folder`" - " for them. Alternatively, make sure you have `safetensors` installed if the model you are using" - " offers the weights in this format." - ) - if offload_folder is not None: - os.makedirs(offload_folder, exist_ok=True) - if offload_state_dict is None: - offload_state_dict = True - - offload_index = {} if device_map is not None and "disk" in device_map.values() else None - if offload_state_dict: - state_dict_folder = tempfile.mkdtemp() - state_dict_index = {} - else: - state_dict_folder = None - state_dict_index = None - - if state_dict is not None: - # load_state_dict will manage the case where we pass a dict instead of a file - # if state dict is not None, it means that we don't need to read the files from resolved_model_file also - resolved_model_file = [state_dict] - - if len(resolved_model_file) > 1: - resolved_model_file = logging.tqdm(resolved_model_file, desc="Loading checkpoint shards") - - for shard_file in resolved_model_file: - state_dict = load_state_dict(shard_file, dduf_entries=dduf_entries) - - def _find_mismatched_keys( - state_dict, - model_state_dict, - loaded_keys, - ignore_mismatched_sizes, - ): - mismatched_keys = [] - if ignore_mismatched_sizes: - for checkpoint_key in loaded_keys: - model_key = checkpoint_key - # If the checkpoint is sharded, we may not have the key here. - if checkpoint_key not in state_dict: - continue - - if ( - model_key in model_state_dict - and state_dict[checkpoint_key].shape != model_state_dict[model_key].shape - ): - mismatched_keys.append( - (checkpoint_key, state_dict[checkpoint_key].shape, model_state_dict[model_key].shape) - ) - del state_dict[checkpoint_key] - return mismatched_keys - - mismatched_keys += _find_mismatched_keys( - state_dict, - model_state_dict, - loaded_keys, - ignore_mismatched_sizes, - ) - - if low_cpu_mem_usage: - offload_index, state_dict_index = load_model_dict_into_meta( - model, - state_dict, - device_map=device_map, - dtype=dtype, - hf_quantizer=hf_quantizer, - keep_in_fp32_modules=keep_in_fp32_modules, - unexpected_keys=unexpected_keys, - offload_folder=offload_folder, - offload_index=offload_index, - state_dict_index=state_dict_index, - state_dict_folder=state_dict_folder, - ) - else: - if assign_to_params_buffers is None: - assign_to_params_buffers = check_support_param_buffer_assignment(model, state_dict) - - error_msgs += _load_state_dict_into_model(model, state_dict, assign_to_params_buffers) - - if offload_index is not None and len(offload_index) > 0: - save_offload_index(offload_index, offload_folder) - offload_index = None - - if offload_state_dict: - load_offloaded_weights(model, state_dict_index, state_dict_folder) - shutil.rmtree(state_dict_folder) - - if len(error_msgs) > 0: - error_msg = "\n\t".join(error_msgs) - if "size mismatch" in error_msg: - error_msg += ( - "\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method." - ) - raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}") - - if len(unexpected_keys) > 0: - logger.warning( - f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when initializing {cls.__name__}: \n {[', '.join(unexpected_keys)]}" - ) - else: - logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n") - - if len(missing_keys) > 0: - logger.warning( - f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" - f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably" - " TRAIN this model on a down-stream task to be able to use it for predictions and inference." - ) - elif len(mismatched_keys) == 0: - logger.info( - f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at" - f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the" - f" checkpoint was trained on, you can already use {model.__class__.__name__} for predictions" - " without further training." - ) - if len(mismatched_keys) > 0: - mismatched_warning = "\n".join( - [ - f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated" - for key, shape1, shape2 in mismatched_keys - ] - ) - logger.warning( - f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" - f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not" - f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be" - " able to use it for predictions and inference." - ) - - return model, missing_keys, unexpected_keys, mismatched_keys, offload_index, error_msgs - - @classmethod - def _get_signature_keys(cls, obj): - parameters = inspect.signature(obj.__init__).parameters - required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty} - optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty}) - expected_modules = set(required_parameters.keys()) - {"self"} - - return expected_modules, optional_parameters - - # Adapted from `transformers` modeling_utils.py - def _get_no_split_modules(self, device_map: str): - """ - Get the modules of the model that should not be split when using device_map. We iterate through the modules to - get the underlying `_no_split_modules`. - - Args: - 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 - """ - _no_split_modules = set() - modules_to_check = [self] - while len(modules_to_check) > 0: - module = modules_to_check.pop(-1) - # if the module does not appear in _no_split_modules, we also check the children - if module.__class__.__name__ not in _no_split_modules: - if isinstance(module, ModelMixin): - if module._no_split_modules is None: - raise ValueError( - f"{module.__class__.__name__} does not support `device_map='{device_map}'`. To implement support, the model " - "class needs to implement the `_no_split_modules` attribute." - ) - else: - _no_split_modules = _no_split_modules | set(module._no_split_modules) - modules_to_check += list(module.children()) - return list(_no_split_modules) - - @classmethod - def _set_default_torch_dtype(cls, dtype: torch.dtype) -> torch.dtype: - """ - Change the default dtype and return the previous one. This is needed when wanting to instantiate the model - under specific dtype. - - Args: - dtype (`torch.dtype`): - a floating dtype to set to. - - Returns: - `torch.dtype`: the original `dtype` that can be used to restore `torch.set_default_dtype(dtype)` if it was - modified. If it wasn't, returns `None`. - - Note `set_default_dtype` currently only works with floating-point types and asserts if for example, - `torch.int64` is passed. So if a non-float `dtype` is passed this functions will throw an exception. - """ - if not dtype.is_floating_point: - raise ValueError( - f"Can't instantiate {cls.__name__} model under dtype={dtype} since it is not a floating point dtype" - ) - - logger.info(f"Instantiating {cls.__name__} model under default dtype {dtype}.") - dtype_orig = torch.get_default_dtype() - torch.set_default_dtype(dtype) - return dtype_orig - - @property - def device(self) -> torch.device: - """ - `torch.device`: The device on which the module is (assuming that all the module parameters are on the same - device). - """ - return get_parameter_device(self) - - @property - def dtype(self) -> torch.dtype: - """ - `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype). - """ - return get_parameter_dtype(self) - - def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool = False) -> int: - """ - Get number of (trainable or non-embedding) parameters in the module. - - Args: - only_trainable (`bool`, *optional*, defaults to `False`): - Whether or not to return only the number of trainable parameters. - exclude_embeddings (`bool`, *optional*, defaults to `False`): - Whether or not to return only the number of non-embedding parameters. - - Returns: - `int`: The number of parameters. - - Example: - - ```py - from diffusers import UNet2DConditionModel - - model_id = "runwayml/stable-diffusion-v1-5" - unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet") - unet.num_parameters(only_trainable=True) - 859520964 - ``` - """ - is_loaded_in_4bit = getattr(self, "is_loaded_in_4bit", False) - - if is_loaded_in_4bit: - if is_bitsandbytes_available(): - import bitsandbytes as bnb - else: - raise ValueError( - "bitsandbytes is not installed but it seems that the model has been loaded in 4bit precision, something went wrong" - " make sure to install bitsandbytes with `pip install bitsandbytes`. You also need a GPU. " - ) - - if exclude_embeddings: - embedding_param_names = [ - f"{name}.weight" for name, module_type in self.named_modules() if isinstance(module_type, nn.Embedding) - ] - total_parameters = [ - parameter for name, parameter in self.named_parameters() if name not in embedding_param_names - ] - else: - total_parameters = list(self.parameters()) - - total_numel = [] - - for param in total_parameters: - if param.requires_grad or not only_trainable: - # For 4bit models, we need to multiply the number of parameters by 2 as half of the parameters are - # used for the 4bit quantization (uint8 tensors are stored) - if is_loaded_in_4bit and isinstance(param, bnb.nn.Params4bit): - if hasattr(param, "element_size"): - num_bytes = param.element_size() - elif hasattr(param, "quant_storage"): - num_bytes = param.quant_storage.itemsize - else: - num_bytes = 1 - total_numel.append(param.numel() * 2 * num_bytes) - else: - total_numel.append(param.numel()) - - return sum(total_numel) - - def get_memory_footprint(self, return_buffers=True): - r""" - Get the memory footprint of a model. This will return the memory footprint of the current model in bytes. - Useful to benchmark the memory footprint of the current model and design some tests. Solution inspired from the - PyTorch discussions: https://discuss.pytorch.org/t/gpu-memory-that-model-uses/56822/2 - - Arguments: - return_buffers (`bool`, *optional*, defaults to `True`): - Whether to return the size of the buffer tensors in the computation of the memory footprint. Buffers - are tensors that do not require gradients and not registered as parameters. E.g. mean and std in batch - norm layers. Please see: https://discuss.pytorch.org/t/what-pytorch-means-by-buffers/120266/2 - """ - mem = sum([param.nelement() * param.element_size() for param in self.parameters()]) - if return_buffers: - mem_bufs = sum([buf.nelement() * buf.element_size() for buf in self.buffers()]) - mem = mem + mem_bufs - return mem - - def _set_gradient_checkpointing( - self, enable: bool = True, gradient_checkpointing_func: Callable = torch.utils.checkpoint.checkpoint - ) -> None: - is_gradient_checkpointing_set = False - - for name, module in self.named_modules(): - if hasattr(module, "gradient_checkpointing"): - logger.debug(f"Setting `gradient_checkpointing={enable}` for '{name}'") - module._gradient_checkpointing_func = gradient_checkpointing_func - module.gradient_checkpointing = enable - is_gradient_checkpointing_set = True - - if not is_gradient_checkpointing_set: - raise ValueError( - f"The module {self.__class__.__name__} does not support gradient checkpointing. Please make sure to " - f"use a module that supports gradient checkpointing by creating a boolean attribute `gradient_checkpointing`." - ) - - def _fix_state_dict_keys_on_load(self, state_dict: OrderedDict) -> None: - """ - This function fix the state dict of the model to take into account some changes that were made in the model - architecture: - - deprecated attention blocks (happened before we introduced sharded checkpoint, - so this is why we apply this method only when loading non sharded checkpoints for now) - """ - deprecated_attention_block_paths = [] - - def recursive_find_attn_block(name, module): - if hasattr(module, "_from_deprecated_attn_block") and module._from_deprecated_attn_block: - deprecated_attention_block_paths.append(name) - - for sub_name, sub_module in module.named_children(): - sub_name = sub_name if name == "" else f"{name}.{sub_name}" - recursive_find_attn_block(sub_name, sub_module) - - recursive_find_attn_block("", self) - - # NOTE: we have to check if the deprecated parameters are in the state dict - # because it is possible we are loading from a state dict that was already - # converted - - for path in deprecated_attention_block_paths: - # group_norm path stays the same - - # query -> to_q - if f"{path}.query.weight" in state_dict: - state_dict[f"{path}.to_q.weight"] = state_dict.pop(f"{path}.query.weight") - if f"{path}.query.bias" in state_dict: - state_dict[f"{path}.to_q.bias"] = state_dict.pop(f"{path}.query.bias") - - # key -> to_k - if f"{path}.key.weight" in state_dict: - state_dict[f"{path}.to_k.weight"] = state_dict.pop(f"{path}.key.weight") - if f"{path}.key.bias" in state_dict: - state_dict[f"{path}.to_k.bias"] = state_dict.pop(f"{path}.key.bias") - - # value -> to_v - if f"{path}.value.weight" in state_dict: - state_dict[f"{path}.to_v.weight"] = state_dict.pop(f"{path}.value.weight") - if f"{path}.value.bias" in state_dict: - state_dict[f"{path}.to_v.bias"] = state_dict.pop(f"{path}.value.bias") - - # proj_attn -> to_out.0 - if f"{path}.proj_attn.weight" in state_dict: - state_dict[f"{path}.to_out.0.weight"] = state_dict.pop(f"{path}.proj_attn.weight") - if f"{path}.proj_attn.bias" in state_dict: - state_dict[f"{path}.to_out.0.bias"] = state_dict.pop(f"{path}.proj_attn.bias") - return state_dict - - -class LegacyModelMixin(ModelMixin): - r""" - A subclass of `ModelMixin` to resolve class mapping from legacy classes (like `Transformer2DModel`) to more - pipeline-specific classes (like `DiTTransformer2DModel`). - """ - - @classmethod - @validate_hf_hub_args - def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs): - # To prevent dependency import problem. - from .model_loading_utils import _fetch_remapped_cls_from_config - - # Create a copy of the kwargs so that we don't mess with the keyword arguments in the downstream calls. - kwargs_copy = kwargs.copy() - - 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) - - # Load config if we don't provide a configuration - config_path = pretrained_model_name_or_path - - user_agent = { - "diffusers": __version__, - "file_type": "model", - "framework": "pytorch", - } - - # load config - config, _, _ = cls.load_config( - config_path, - cache_dir=cache_dir, - return_unused_kwargs=True, - return_commit_hash=True, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - **kwargs, - ) - # resolve remapping - remapped_class = _fetch_remapped_cls_from_config(config, cls) - - return remapped_class.from_pretrained(pretrained_model_name_or_path, **kwargs_copy) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/normalization.py b/venv/lib/python3.10/site-packages/diffusers/models/normalization.py deleted file mode 100644 index ae2a6298f5f7a3c9691def0357cf40f1f4de83c9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/normalization.py +++ /dev/null @@ -1,648 +0,0 @@ -# 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 numbers -from typing import Dict, Optional, Tuple - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from ..utils import is_torch_npu_available, is_torch_version -from .activations import get_activation -from .embeddings import CombinedTimestepLabelEmbeddings, PixArtAlphaCombinedTimestepSizeEmbeddings - - -class AdaLayerNorm(nn.Module): - r""" - Norm layer modified to incorporate timestep embeddings. - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - num_embeddings (`int`, *optional*): The size of the embeddings dictionary. - output_dim (`int`, *optional*): - norm_elementwise_affine (`bool`, defaults to `False): - norm_eps (`bool`, defaults to `False`): - chunk_dim (`int`, defaults to `0`): - """ - - def __init__( - self, - embedding_dim: int, - num_embeddings: Optional[int] = None, - output_dim: Optional[int] = None, - norm_elementwise_affine: bool = False, - norm_eps: float = 1e-5, - chunk_dim: int = 0, - ): - super().__init__() - - self.chunk_dim = chunk_dim - output_dim = output_dim or embedding_dim * 2 - - if num_embeddings is not None: - self.emb = nn.Embedding(num_embeddings, embedding_dim) - else: - self.emb = None - - self.silu = nn.SiLU() - self.linear = nn.Linear(embedding_dim, output_dim) - 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 - ) -> torch.Tensor: - if self.emb is not None: - temb = self.emb(timestep) - - temb = self.linear(self.silu(temb)) - - if self.chunk_dim == 1: - # This is a bit weird why we have the order of "shift, scale" here and "scale, shift" in the - # other if-branch. This branch is specific to CogVideoX and OmniGen for now. - shift, scale = temb.chunk(2, dim=1) - shift = shift[:, None, :] - scale = scale[:, None, :] - else: - scale, shift = temb.chunk(2, dim=0) - - x = self.norm(x) * (1 + scale) + shift - return x - - -class FP32LayerNorm(nn.LayerNorm): - def forward(self, inputs: torch.Tensor) -> torch.Tensor: - origin_dtype = inputs.dtype - return F.layer_norm( - inputs.float(), - self.normalized_shape, - self.weight.float() if self.weight is not None else None, - self.bias.float() if self.bias is not None else None, - self.eps, - ).to(origin_dtype) - - -class SD35AdaLayerNormZeroX(nn.Module): - r""" - Norm layer adaptive layer norm zero (AdaLN-Zero). - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - num_embeddings (`int`): The size of the embeddings dictionary. - """ - - def __init__(self, embedding_dim: int, norm_type: str = "layer_norm", bias: bool = True) -> None: - super().__init__() - - self.silu = nn.SiLU() - self.linear = nn.Linear(embedding_dim, 9 * embedding_dim, bias=bias) - if norm_type == "layer_norm": - self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) - else: - raise ValueError(f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm'.") - - def forward( - self, - hidden_states: torch.Tensor, - emb: Optional[torch.Tensor] = 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 - ) - norm_hidden_states = self.norm(hidden_states) - hidden_states = norm_hidden_states * (1 + scale_msa[:, None]) + shift_msa[:, None] - norm_hidden_states2 = norm_hidden_states * (1 + scale_msa2[:, None]) + shift_msa2[:, None] - return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 - - -class AdaLayerNormZero(nn.Module): - r""" - Norm layer adaptive layer norm zero (adaLN-Zero). - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - 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): - super().__init__() - if num_embeddings is not None: - self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, embedding_dim) - else: - self.emb = None - - self.silu = nn.SiLU() - self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=bias) - if norm_type == "layer_norm": - self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) - elif norm_type == "fp32_layer_norm": - self.norm = FP32LayerNorm(embedding_dim, elementwise_affine=False, bias=False) - else: - raise ValueError( - f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'." - ) - - 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]: - if self.emb is not None: - emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype) - emb = self.linear(self.silu(emb)) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1) - x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] - return x, gate_msa, shift_mlp, scale_mlp, gate_mlp - - -class AdaLayerNormZeroSingle(nn.Module): - r""" - Norm layer adaptive layer norm zero (adaLN-Zero). - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - num_embeddings (`int`): The size of the embeddings dictionary. - """ - - def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True): - super().__init__() - - self.silu = nn.SiLU() - self.linear = nn.Linear(embedding_dim, 3 * embedding_dim, bias=bias) - if norm_type == "layer_norm": - self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) - else: - raise ValueError( - f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'." - ) - - def forward( - self, - x: torch.Tensor, - emb: Optional[torch.Tensor] = 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] - return x, gate_msa - - -class LuminaRMSNormZero(nn.Module): - """ - Norm layer adaptive RMS normalization zero. - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - """ - - def __init__(self, embedding_dim: int, norm_eps: float, norm_elementwise_affine: bool): - super().__init__() - self.silu = nn.SiLU() - self.linear = nn.Linear( - min(embedding_dim, 1024), - 4 * embedding_dim, - bias=True, - ) - self.norm = RMSNorm(embedding_dim, eps=norm_eps) - - def forward( - self, - x: torch.Tensor, - emb: Optional[torch.Tensor] = 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]) - - return x, gate_msa, scale_mlp, gate_mlp - - -class AdaLayerNormSingle(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). - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - use_additional_conditions (`bool`): To use additional conditions for normalization or not. - """ - - def __init__(self, embedding_dim: int, use_additional_conditions: bool = False): - super().__init__() - - 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, 6 * embedding_dim, bias=True) - - 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]: - # 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 AdaGroupNorm(nn.Module): - r""" - GroupNorm layer modified to incorporate timestep embeddings. - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - num_embeddings (`int`): The size of the embeddings dictionary. - num_groups (`int`): The number of groups to separate the channels into. - act_fn (`str`, *optional*, defaults to `None`): The activation function to use. - eps (`float`, *optional*, defaults to `1e-5`): The epsilon value to use for numerical stability. - """ - - def __init__( - self, embedding_dim: int, out_dim: int, num_groups: int, act_fn: Optional[str] = None, eps: float = 1e-5 - ): - super().__init__() - self.num_groups = num_groups - self.eps = eps - - if act_fn is None: - self.act = None - else: - self.act = get_activation(act_fn) - - self.linear = nn.Linear(embedding_dim, out_dim * 2) - - def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: - if self.act: - emb = self.act(emb) - emb = self.linear(emb) - emb = emb[:, :, None, None] - scale, shift = emb.chunk(2, dim=1) - - x = F.group_norm(x, self.num_groups, eps=self.eps) - x = x * (1 + scale) + shift - return x - - -class AdaLayerNormContinuous(nn.Module): - r""" - Adaptive normalization layer with a norm layer (layer_norm or rms_norm). - - Args: - embedding_dim (`int`): Embedding dimension to use during projection. - conditioning_embedding_dim (`int`): Dimension of the input condition. - elementwise_affine (`bool`, defaults to `True`): - Boolean flag to denote if affine transformation should be applied. - eps (`float`, defaults to 1e-5): Epsilon factor. - bias (`bias`, defaults to `True`): Boolean flag to denote if bias should be use. - norm_type (`str`, defaults to `"layer_norm"`): - Normalization layer to use. Values supported: "layer_norm", "rms_norm". - """ - - def __init__( - self, - embedding_dim: int, - conditioning_embedding_dim: int, - # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters - # because the output is immediately scaled and shifted by the projected conditioning embeddings. - # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters. - # However, this is how it was implemented in the original code, and it's rather likely you should - # set `elementwise_affine` to False. - elementwise_affine=True, - eps=1e-5, - bias=True, - norm_type="layer_norm", - ): - super().__init__() - self.silu = nn.SiLU() - 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: - # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT) - emb = self.linear(self.silu(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 LuminaLayerNormContinuous(nn.Module): - def __init__( - self, - embedding_dim: int, - conditioning_embedding_dim: int, - # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters - # because the output is immediately scaled and shifted by the projected conditioning embeddings. - # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters. - # However, this is how it was implemented in the original code, and it's rather likely you should - # set `elementwise_affine` to False. - elementwise_affine=True, - eps=1e-5, - bias=True, - norm_type="layer_norm", - out_dim: Optional[int] = None, - ): - super().__init__() - - # AdaLN - self.silu = nn.SiLU() - self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, 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=eps, elementwise_affine=elementwise_affine) - else: - raise ValueError(f"unknown norm_type {norm_type}") - - self.linear_2 = None - if out_dim is not None: - self.linear_2 = nn.Linear(embedding_dim, out_dim, bias=bias) - - def forward( - self, - x: torch.Tensor, - conditioning_embedding: torch.Tensor, - ) -> torch.Tensor: - # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT) - emb = self.linear_1(self.silu(conditioning_embedding).to(x.dtype)) - scale = emb - x = self.norm(x) * (1 + scale)[:, None, :] - - if self.linear_2 is not None: - x = self.linear_2(x) - - return x - - -class CogView3PlusAdaLayerNormZeroTextImage(nn.Module): - r""" - Norm layer adaptive layer norm zero (adaLN-Zero). - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - num_embeddings (`int`): The size of the embeddings dictionary. - """ - - def __init__(self, embedding_dim: int, dim: int): - super().__init__() - - self.silu = nn.SiLU() - self.linear = nn.Linear(embedding_dim, 12 * dim, bias=True) - self.norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) - self.norm_c = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) - - 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 = self.linear(self.silu(emb)) - ( - shift_msa, - scale_msa, - gate_msa, - shift_mlp, - scale_mlp, - gate_mlp, - c_shift_msa, - c_scale_msa, - c_gate_msa, - c_shift_mlp, - c_scale_mlp, - c_gate_mlp, - ) = emb.chunk(12, dim=1) - normed_x = self.norm_x(x) - normed_context = self.norm_c(context) - x = normed_x * (1 + scale_msa[:, None]) + shift_msa[:, None] - context = normed_context * (1 + c_scale_msa[:, None]) + c_shift_msa[:, None] - return x, gate_msa, shift_mlp, scale_mlp, gate_mlp, context, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp - - -class CogVideoXLayerNormZero(nn.Module): - def __init__( - self, - conditioning_dim: int, - embedding_dim: int, - elementwise_affine: bool = True, - eps: float = 1e-5, - bias: bool = True, - ) -> None: - super().__init__() - - self.silu = nn.SiLU() - self.linear = nn.Linear(conditioning_dim, 6 * embedding_dim, bias=bias) - self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine) - - def forward( - self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: 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, :] - return hidden_states, encoder_hidden_states, gate[:, None, :], enc_gate[:, None, :] - - -if is_torch_version(">=", "2.1.0"): - LayerNorm = nn.LayerNorm -else: - # Has optional bias parameter compared to torch layer norm - # TODO: replace with torch layernorm once min required torch version >= 2.1 - class LayerNorm(nn.Module): - r""" - LayerNorm with the bias parameter. - - Args: - dim (`int`): Dimensionality to use for the parameters. - eps (`float`, defaults to 1e-5): Epsilon factor. - elementwise_affine (`bool`, defaults to `True`): - Boolean flag to denote if affine transformation should be applied. - bias (`bias`, defaults to `True`): Boolean flag to denote if bias should be use. - """ - - def __init__(self, dim, eps: float = 1e-5, elementwise_affine: bool = True, bias: bool = True): - super().__init__() - - self.eps = eps - - if isinstance(dim, numbers.Integral): - dim = (dim,) - - self.dim = torch.Size(dim) - - if elementwise_affine: - self.weight = nn.Parameter(torch.ones(dim)) - self.bias = nn.Parameter(torch.zeros(dim)) if bias else None - else: - self.weight = None - self.bias = None - - def forward(self, input): - return F.layer_norm(input, self.dim, self.weight, self.bias, self.eps) - - -class RMSNorm(nn.Module): - r""" - RMS Norm as introduced in https://huggingface.co/papers/1910.07467 by Zhang et al. - - Args: - dim (`int`): Number of dimensions to use for `weights`. Only effective when `elementwise_affine` is True. - eps (`float`): Small value to use when calculating the reciprocal of the square-root. - elementwise_affine (`bool`, defaults to `True`): - Boolean flag to denote if affine transformation should be applied. - bias (`bool`, defaults to False): If also training the `bias` param. - """ - - def __init__(self, dim, eps: float, elementwise_affine: bool = True, bias: bool = False): - super().__init__() - - self.eps = eps - self.elementwise_affine = elementwise_affine - - if isinstance(dim, numbers.Integral): - dim = (dim,) - - self.dim = torch.Size(dim) - - self.weight = None - self.bias = None - - if elementwise_affine: - self.weight = nn.Parameter(torch.ones(dim)) - if bias: - self.bias = nn.Parameter(torch.zeros(dim)) - - def forward(self, hidden_states): - if is_torch_npu_available(): - import torch_npu - - if self.weight is not None: - # convert into half-precision if necessary - if self.weight.dtype in [torch.float16, torch.bfloat16]: - hidden_states = hidden_states.to(self.weight.dtype) - hidden_states = torch_npu.npu_rms_norm(hidden_states, self.weight, epsilon=self.eps)[0] - if self.bias is not None: - hidden_states = hidden_states + self.bias - else: - input_dtype = hidden_states.dtype - variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.eps) - - if self.weight is not None: - # convert into half-precision if necessary - if self.weight.dtype in [torch.float16, torch.bfloat16]: - hidden_states = hidden_states.to(self.weight.dtype) - hidden_states = hidden_states * self.weight - if self.bias is not None: - hidden_states = hidden_states + self.bias - else: - hidden_states = hidden_states.to(input_dtype) - - return hidden_states - - -# TODO: (Dhruv) This can be replaced with regular RMSNorm in Mochi once `_keep_in_fp32_modules` is supported -# for sharded checkpoints, see: https://github.com/huggingface/diffusers/issues/10013 -class MochiRMSNorm(nn.Module): - def __init__(self, dim, eps: float, elementwise_affine: bool = True): - super().__init__() - - self.eps = eps - - if isinstance(dim, numbers.Integral): - dim = (dim,) - - self.dim = torch.Size(dim) - - if elementwise_affine: - self.weight = nn.Parameter(torch.ones(dim)) - else: - self.weight = None - - def forward(self, hidden_states): - input_dtype = hidden_states.dtype - variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.eps) - - if self.weight is not None: - hidden_states = hidden_states * self.weight - hidden_states = hidden_states.to(input_dtype) - - return hidden_states - - -class GlobalResponseNorm(nn.Module): - r""" - Global response normalization as introduced in ConvNeXt-v2 (https://huggingface.co/papers/2301.00808). - - Args: - dim (`int`): Number of dimensions to use for the `gamma` and `beta`. - """ - - # Taken from https://github.com/facebookresearch/ConvNeXt-V2/blob/3608f67cc1dae164790c5d0aead7bf2d73d9719b/models/utils.py#L105 - def __init__(self, dim): - super().__init__() - self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) - self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) - - def forward(self, x): - gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) - nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-6) - return self.gamma * (x * nx) + self.beta + x - - -class LpNorm(nn.Module): - def __init__(self, p: int = 2, dim: int = -1, eps: float = 1e-12): - super().__init__() - - self.p = p - self.dim = dim - self.eps = eps - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return F.normalize(hidden_states, p=self.p, dim=self.dim, eps=self.eps) - - -def get_normalization( - norm_type: str = "batch_norm", - num_features: Optional[int] = None, - eps: float = 1e-5, - elementwise_affine: bool = True, - bias: bool = True, -) -> nn.Module: - if norm_type == "rms_norm": - norm = RMSNorm(num_features, eps=eps, elementwise_affine=elementwise_affine, bias=bias) - elif norm_type == "layer_norm": - norm = nn.LayerNorm(num_features, eps=eps, elementwise_affine=elementwise_affine, bias=bias) - elif norm_type == "batch_norm": - norm = nn.BatchNorm2d(num_features, eps=eps, affine=elementwise_affine) - else: - raise ValueError(f"{norm_type=} is not supported.") - return norm diff --git a/venv/lib/python3.10/site-packages/diffusers/models/resnet.py b/venv/lib/python3.10/site-packages/diffusers/models/resnet.py deleted file mode 100644 index c0b4ad40055a4961ec874b50e821b221f594cd2f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/resnet.py +++ /dev/null @@ -1,797 +0,0 @@ -# Copyright 2025 The HuggingFace Team. All rights reserved. -# `TemporalConvLayer` Copyright 2025 Alibaba DAMO-VILAB, The ModelScope 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 functools import partial -from typing import Optional, Tuple, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from ..utils import deprecate -from .activations import get_activation -from .attention_processor import SpatialNorm -from .downsampling import ( # noqa - Downsample1D, - Downsample2D, - FirDownsample2D, - KDownsample2D, - downsample_2d, -) -from .normalization import AdaGroupNorm -from .upsampling import ( # noqa - FirUpsample2D, - KUpsample2D, - Upsample1D, - Upsample2D, - upfirdn2d_native, - upsample_2d, -) - - -class ResnetBlockCondNorm2D(nn.Module): - r""" - A Resnet block that use normalization layer that incorporate conditioning information. - - Parameters: - in_channels (`int`): The number of channels in the input. - out_channels (`int`, *optional*, default to be `None`): - The number of output channels for the first conv2d layer. If None, same as `in_channels`. - dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use. - temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding. - groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer. - groups_out (`int`, *optional*, default to None): - The number of groups to use for the second normalization layer. if set to None, same as `groups`. - eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization. - non_linearity (`str`, *optional*, default to `"swish"`): the activation function to use. - time_embedding_norm (`str`, *optional*, default to `"ada_group"` ): - The normalization layer for time embedding `temb`. Currently only support "ada_group" or "spatial". - kernel (`torch.Tensor`, optional, default to None): FIR filter, see - [`~models.resnet.FirUpsample2D`] and [`~models.resnet.FirDownsample2D`]. - output_scale_factor (`float`, *optional*, default to be `1.0`): the scale factor to use for the output. - use_in_shortcut (`bool`, *optional*, default to `True`): - If `True`, add a 1x1 nn.conv2d layer for skip-connection. - up (`bool`, *optional*, default to `False`): If `True`, add an upsample layer. - down (`bool`, *optional*, default to `False`): If `True`, add a downsample layer. - conv_shortcut_bias (`bool`, *optional*, default to `True`): If `True`, adds a learnable bias to the - `conv_shortcut` output. - conv_2d_out_channels (`int`, *optional*, default to `None`): the number of channels in the output. - If None, same as `out_channels`. - """ - - def __init__( - self, - *, - in_channels: int, - out_channels: Optional[int] = None, - conv_shortcut: bool = False, - dropout: float = 0.0, - temb_channels: int = 512, - groups: int = 32, - groups_out: Optional[int] = 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, - up: bool = False, - down: bool = False, - conv_shortcut_bias: bool = True, - conv_2d_out_channels: Optional[int] = None, - ): - super().__init__() - 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 - self.up = up - self.down = down - self.output_scale_factor = output_scale_factor - self.time_embedding_norm = time_embedding_norm - - if groups_out is None: - groups_out = groups - - if self.time_embedding_norm == "ada_group": # ada_group - self.norm1 = AdaGroupNorm(temb_channels, in_channels, groups, eps=eps) - elif self.time_embedding_norm == "spatial": - self.norm1 = SpatialNorm(in_channels, temb_channels) - else: - raise ValueError(f" unsupported time_embedding_norm: {self.time_embedding_norm}") - - self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) - - if self.time_embedding_norm == "ada_group": # ada_group - self.norm2 = AdaGroupNorm(temb_channels, out_channels, groups_out, eps=eps) - elif self.time_embedding_norm == "spatial": # spatial - self.norm2 = SpatialNorm(out_channels, temb_channels) - else: - raise ValueError(f" unsupported time_embedding_norm: {self.time_embedding_norm}") - - self.dropout = torch.nn.Dropout(dropout) - - conv_2d_out_channels = conv_2d_out_channels or out_channels - self.conv2 = nn.Conv2d(out_channels, conv_2d_out_channels, kernel_size=3, stride=1, padding=1) - - self.nonlinearity = get_activation(non_linearity) - - self.upsample = self.downsample = None - if self.up: - self.upsample = Upsample2D(in_channels, use_conv=False) - elif self.down: - self.downsample = Downsample2D(in_channels, use_conv=False, padding=1, name="op") - - self.use_in_shortcut = self.in_channels != conv_2d_out_channels if use_in_shortcut is None else use_in_shortcut - - self.conv_shortcut = None - if self.use_in_shortcut: - self.conv_shortcut = nn.Conv2d( - in_channels, - conv_2d_out_channels, - kernel_size=1, - stride=1, - padding=0, - bias=conv_shortcut_bias, - ) - - def forward(self, input_tensor: torch.Tensor, temb: torch.Tensor, *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) - - hidden_states = input_tensor - - hidden_states = self.norm1(hidden_states, temb) - - hidden_states = self.nonlinearity(hidden_states) - - if self.upsample is not None: - # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 - if hidden_states.shape[0] >= 64: - input_tensor = input_tensor.contiguous() - hidden_states = hidden_states.contiguous() - input_tensor = self.upsample(input_tensor) - hidden_states = self.upsample(hidden_states) - - elif self.downsample is not None: - input_tensor = self.downsample(input_tensor) - hidden_states = self.downsample(hidden_states) - - hidden_states = self.conv1(hidden_states) - - hidden_states = self.norm2(hidden_states, temb) - - hidden_states = self.nonlinearity(hidden_states) - - hidden_states = self.dropout(hidden_states) - hidden_states = self.conv2(hidden_states) - - if self.conv_shortcut is not None: - input_tensor = self.conv_shortcut(input_tensor) - - output_tensor = (input_tensor + hidden_states) / self.output_scale_factor - - return output_tensor - - -class ResnetBlock2D(nn.Module): - r""" - A Resnet block. - - Parameters: - in_channels (`int`): The number of channels in the input. - out_channels (`int`, *optional*, default to be `None`): - The number of output channels for the first conv2d layer. If None, same as `in_channels`. - dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use. - temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding. - groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer. - groups_out (`int`, *optional*, default to None): - The number of groups to use for the second normalization layer. if set to None, same as `groups`. - eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization. - non_linearity (`str`, *optional*, default to `"swish"`): the activation function to use. - time_embedding_norm (`str`, *optional*, default to `"default"` ): Time scale shift config. - By default, apply timestep embedding conditioning with a simple shift mechanism. Choose "scale_shift" for a - stronger conditioning with scale and shift. - kernel (`torch.Tensor`, optional, default to None): FIR filter, see - [`~models.resnet.FirUpsample2D`] and [`~models.resnet.FirDownsample2D`]. - output_scale_factor (`float`, *optional*, default to be `1.0`): the scale factor to use for the output. - use_in_shortcut (`bool`, *optional*, default to `True`): - If `True`, add a 1x1 nn.conv2d layer for skip-connection. - up (`bool`, *optional*, default to `False`): If `True`, add an upsample layer. - down (`bool`, *optional*, default to `False`): If `True`, add a downsample layer. - conv_shortcut_bias (`bool`, *optional*, default to `True`): If `True`, adds a learnable bias to the - `conv_shortcut` output. - conv_2d_out_channels (`int`, *optional*, default to `None`): the number of channels in the output. - If None, same as `out_channels`. - """ - - def __init__( - self, - *, - in_channels: int, - out_channels: Optional[int] = None, - conv_shortcut: bool = False, - dropout: float = 0.0, - temb_channels: int = 512, - groups: int = 32, - groups_out: Optional[int] = 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, - output_scale_factor: float = 1.0, - use_in_shortcut: Optional[bool] = None, - up: bool = False, - down: bool = False, - conv_shortcut_bias: bool = True, - conv_2d_out_channels: Optional[int] = None, - ): - super().__init__() - if time_embedding_norm == "ada_group": - raise ValueError( - "This class cannot be used with `time_embedding_norm==ada_group`, please use `ResnetBlockCondNorm2D` instead", - ) - if time_embedding_norm == "spatial": - raise ValueError( - "This class cannot be used with `time_embedding_norm==spatial`, please use `ResnetBlockCondNorm2D` instead", - ) - - self.pre_norm = True - 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 - self.up = up - self.down = down - self.output_scale_factor = output_scale_factor - self.time_embedding_norm = time_embedding_norm - self.skip_time_act = skip_time_act - - if groups_out is None: - groups_out = groups - - self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) - - self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) - - if temb_channels is not None: - if self.time_embedding_norm == "default": - self.time_emb_proj = nn.Linear(temb_channels, out_channels) - elif self.time_embedding_norm == "scale_shift": - self.time_emb_proj = nn.Linear(temb_channels, 2 * out_channels) - else: - raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm} ") - else: - self.time_emb_proj = None - - self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True) - - self.dropout = torch.nn.Dropout(dropout) - conv_2d_out_channels = conv_2d_out_channels or out_channels - self.conv2 = nn.Conv2d(out_channels, conv_2d_out_channels, kernel_size=3, stride=1, padding=1) - - self.nonlinearity = get_activation(non_linearity) - - self.upsample = self.downsample = None - if self.up: - if kernel == "fir": - fir_kernel = (1, 3, 3, 1) - self.upsample = lambda x: upsample_2d(x, kernel=fir_kernel) - elif kernel == "sde_vp": - self.upsample = partial(F.interpolate, scale_factor=2.0, mode="nearest") - else: - self.upsample = Upsample2D(in_channels, use_conv=False) - elif self.down: - if kernel == "fir": - fir_kernel = (1, 3, 3, 1) - self.downsample = lambda x: downsample_2d(x, kernel=fir_kernel) - elif kernel == "sde_vp": - self.downsample = partial(F.avg_pool2d, kernel_size=2, stride=2) - else: - self.downsample = Downsample2D(in_channels, use_conv=False, padding=1, name="op") - - self.use_in_shortcut = self.in_channels != conv_2d_out_channels if use_in_shortcut is None else use_in_shortcut - - self.conv_shortcut = None - if self.use_in_shortcut: - self.conv_shortcut = nn.Conv2d( - in_channels, - conv_2d_out_channels, - kernel_size=1, - stride=1, - padding=0, - bias=conv_shortcut_bias, - ) - - def forward(self, input_tensor: torch.Tensor, temb: torch.Tensor, *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) - - hidden_states = input_tensor - - hidden_states = self.norm1(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - - if self.upsample is not None: - # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 - if hidden_states.shape[0] >= 64: - input_tensor = input_tensor.contiguous() - hidden_states = hidden_states.contiguous() - input_tensor = self.upsample(input_tensor) - hidden_states = self.upsample(hidden_states) - elif self.downsample is not None: - input_tensor = self.downsample(input_tensor) - hidden_states = self.downsample(hidden_states) - - hidden_states = self.conv1(hidden_states) - - if self.time_emb_proj is not None: - if not self.skip_time_act: - temb = self.nonlinearity(temb) - temb = self.time_emb_proj(temb)[:, :, None, None] - - if self.time_embedding_norm == "default": - if temb is not None: - hidden_states = hidden_states + temb - hidden_states = self.norm2(hidden_states) - elif self.time_embedding_norm == "scale_shift": - if temb is None: - raise ValueError( - f" `temb` should not be None when `time_embedding_norm` is {self.time_embedding_norm}" - ) - time_scale, time_shift = torch.chunk(temb, 2, dim=1) - hidden_states = self.norm2(hidden_states) - hidden_states = hidden_states * (1 + time_scale) + time_shift - else: - hidden_states = self.norm2(hidden_states) - - hidden_states = self.nonlinearity(hidden_states) - - hidden_states = self.dropout(hidden_states) - hidden_states = self.conv2(hidden_states) - - if self.conv_shortcut is not None: - input_tensor = self.conv_shortcut(input_tensor.contiguous()) - - output_tensor = (input_tensor + hidden_states) / self.output_scale_factor - - return output_tensor - - -# unet_rl.py -def rearrange_dims(tensor: torch.Tensor) -> torch.Tensor: - if len(tensor.shape) == 2: - return tensor[:, :, None] - if len(tensor.shape) == 3: - return tensor[:, :, None, :] - elif len(tensor.shape) == 4: - return tensor[:, :, 0, :] - else: - raise ValueError(f"`len(tensor)`: {len(tensor)} has to be 2, 3 or 4.") - - -class Conv1dBlock(nn.Module): - """ - Conv1d --> GroupNorm --> Mish - - Parameters: - inp_channels (`int`): Number of input channels. - out_channels (`int`): Number of output channels. - kernel_size (`int` or `tuple`): Size of the convolving kernel. - n_groups (`int`, default `8`): Number of groups to separate the channels into. - activation (`str`, defaults to `mish`): Name of the activation function. - """ - - def __init__( - self, - inp_channels: int, - out_channels: int, - kernel_size: Union[int, Tuple[int, int]], - n_groups: int = 8, - activation: str = "mish", - ): - super().__init__() - - self.conv1d = nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2) - self.group_norm = nn.GroupNorm(n_groups, out_channels) - self.mish = get_activation(activation) - - def forward(self, inputs: torch.Tensor) -> torch.Tensor: - intermediate_repr = self.conv1d(inputs) - intermediate_repr = rearrange_dims(intermediate_repr) - intermediate_repr = self.group_norm(intermediate_repr) - intermediate_repr = rearrange_dims(intermediate_repr) - output = self.mish(intermediate_repr) - return output - - -# unet_rl.py -class ResidualTemporalBlock1D(nn.Module): - """ - Residual 1D block with temporal convolutions. - - Parameters: - inp_channels (`int`): Number of input channels. - out_channels (`int`): Number of output channels. - embed_dim (`int`): Embedding dimension. - kernel_size (`int` or `tuple`): Size of the convolving kernel. - activation (`str`, defaults `mish`): It is possible to choose the right activation function. - """ - - def __init__( - self, - inp_channels: int, - out_channels: int, - embed_dim: int, - kernel_size: Union[int, Tuple[int, int]] = 5, - activation: str = "mish", - ): - super().__init__() - self.conv_in = Conv1dBlock(inp_channels, out_channels, kernel_size) - self.conv_out = Conv1dBlock(out_channels, out_channels, kernel_size) - - self.time_emb_act = get_activation(activation) - self.time_emb = nn.Linear(embed_dim, out_channels) - - self.residual_conv = ( - nn.Conv1d(inp_channels, out_channels, 1) if inp_channels != out_channels else nn.Identity() - ) - - def forward(self, inputs: torch.Tensor, t: torch.Tensor) -> torch.Tensor: - """ - Args: - inputs : [ batch_size x inp_channels x horizon ] - t : [ batch_size x embed_dim ] - - returns: - out : [ batch_size x out_channels x horizon ] - """ - t = self.time_emb_act(t) - t = self.time_emb(t) - out = self.conv_in(inputs) + rearrange_dims(t) - out = self.conv_out(out) - return out + self.residual_conv(inputs) - - -class TemporalConvLayer(nn.Module): - """ - Temporal convolutional layer that can be used for video (sequence of images) input Code mostly copied from: - https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/models/multi_modal/video_synthesis/unet_sd.py#L1016 - - Parameters: - in_dim (`int`): Number of input channels. - out_dim (`int`): Number of output channels. - dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use. - """ - - def __init__( - self, - in_dim: int, - out_dim: Optional[int] = None, - dropout: float = 0.0, - norm_num_groups: int = 32, - ): - super().__init__() - out_dim = out_dim or in_dim - self.in_dim = in_dim - self.out_dim = out_dim - - # conv layers - self.conv1 = nn.Sequential( - nn.GroupNorm(norm_num_groups, in_dim), - nn.SiLU(), - nn.Conv3d(in_dim, out_dim, (3, 1, 1), padding=(1, 0, 0)), - ) - self.conv2 = nn.Sequential( - nn.GroupNorm(norm_num_groups, out_dim), - nn.SiLU(), - nn.Dropout(dropout), - nn.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0)), - ) - self.conv3 = nn.Sequential( - nn.GroupNorm(norm_num_groups, out_dim), - nn.SiLU(), - nn.Dropout(dropout), - nn.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0)), - ) - self.conv4 = nn.Sequential( - nn.GroupNorm(norm_num_groups, out_dim), - nn.SiLU(), - nn.Dropout(dropout), - nn.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0)), - ) - - # zero out the last layer params,so the conv block is identity - nn.init.zeros_(self.conv4[-1].weight) - nn.init.zeros_(self.conv4[-1].bias) - - def forward(self, hidden_states: torch.Tensor, num_frames: int = 1) -> torch.Tensor: - hidden_states = ( - hidden_states[None, :].reshape((-1, num_frames) + hidden_states.shape[1:]).permute(0, 2, 1, 3, 4) - ) - - identity = hidden_states - hidden_states = self.conv1(hidden_states) - hidden_states = self.conv2(hidden_states) - hidden_states = self.conv3(hidden_states) - hidden_states = self.conv4(hidden_states) - - hidden_states = identity + hidden_states - - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).reshape( - (hidden_states.shape[0] * hidden_states.shape[2], -1) + hidden_states.shape[3:] - ) - return hidden_states - - -class TemporalResnetBlock(nn.Module): - r""" - A Resnet block. - - Parameters: - in_channels (`int`): The number of channels in the input. - out_channels (`int`, *optional*, default to be `None`): - The number of output channels for the first conv2d layer. If None, same as `in_channels`. - temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding. - eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization. - """ - - def __init__( - self, - in_channels: int, - out_channels: Optional[int] = None, - temb_channels: int = 512, - eps: float = 1e-6, - ): - super().__init__() - self.in_channels = in_channels - out_channels = in_channels if out_channels is None else out_channels - self.out_channels = out_channels - - kernel_size = (3, 1, 1) - padding = [k // 2 for k in kernel_size] - - self.norm1 = torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=eps, affine=True) - self.conv1 = nn.Conv3d( - in_channels, - out_channels, - kernel_size=kernel_size, - stride=1, - padding=padding, - ) - - if temb_channels is not None: - self.time_emb_proj = nn.Linear(temb_channels, out_channels) - else: - self.time_emb_proj = None - - self.norm2 = torch.nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=eps, affine=True) - - self.dropout = torch.nn.Dropout(0.0) - self.conv2 = nn.Conv3d( - out_channels, - out_channels, - kernel_size=kernel_size, - stride=1, - padding=padding, - ) - - self.nonlinearity = get_activation("silu") - - self.use_in_shortcut = self.in_channels != out_channels - - self.conv_shortcut = None - if self.use_in_shortcut: - self.conv_shortcut = nn.Conv3d( - in_channels, - out_channels, - kernel_size=1, - stride=1, - padding=0, - ) - - def forward(self, input_tensor: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: - hidden_states = input_tensor - - hidden_states = self.norm1(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - hidden_states = self.conv1(hidden_states) - - if self.time_emb_proj is not None: - temb = self.nonlinearity(temb) - temb = self.time_emb_proj(temb)[:, :, :, None, None] - temb = temb.permute(0, 2, 1, 3, 4) - hidden_states = hidden_states + temb - - hidden_states = self.norm2(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - hidden_states = self.dropout(hidden_states) - hidden_states = self.conv2(hidden_states) - - if self.conv_shortcut is not None: - input_tensor = self.conv_shortcut(input_tensor) - - output_tensor = input_tensor + hidden_states - - return output_tensor - - -# VideoResBlock -class SpatioTemporalResBlock(nn.Module): - r""" - A SpatioTemporal Resnet block. - - Parameters: - in_channels (`int`): The number of channels in the input. - out_channels (`int`, *optional*, default to be `None`): - The number of output channels for the first conv2d layer. If None, same as `in_channels`. - temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding. - eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the spatial resenet. - temporal_eps (`float`, *optional*, defaults to `eps`): The epsilon to use for the temporal resnet. - merge_factor (`float`, *optional*, defaults to `0.5`): The merge factor to use for the temporal mixing. - merge_strategy (`str`, *optional*, defaults to `learned_with_images`): - The merge strategy to use for the temporal mixing. - switch_spatial_to_temporal_mix (`bool`, *optional*, defaults to `False`): - If `True`, switch the spatial and temporal mixing. - """ - - def __init__( - self, - in_channels: int, - out_channels: Optional[int] = None, - temb_channels: int = 512, - eps: float = 1e-6, - temporal_eps: Optional[float] = None, - merge_factor: float = 0.5, - merge_strategy="learned_with_images", - switch_spatial_to_temporal_mix: bool = False, - ): - super().__init__() - - self.spatial_res_block = ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=eps, - ) - - self.temporal_res_block = TemporalResnetBlock( - in_channels=out_channels if out_channels is not None else in_channels, - out_channels=out_channels if out_channels is not None else in_channels, - temb_channels=temb_channels, - eps=temporal_eps if temporal_eps is not None else eps, - ) - - self.time_mixer = AlphaBlender( - alpha=merge_factor, - merge_strategy=merge_strategy, - switch_spatial_to_temporal_mix=switch_spatial_to_temporal_mix, - ) - - def forward( - self, - hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - image_only_indicator: Optional[torch.Tensor] = None, - ): - num_frames = image_only_indicator.shape[-1] - hidden_states = self.spatial_res_block(hidden_states, temb) - - batch_frames, channels, height, width = hidden_states.shape - batch_size = batch_frames // num_frames - - hidden_states_mix = ( - hidden_states[None, :].reshape(batch_size, num_frames, channels, height, width).permute(0, 2, 1, 3, 4) - ) - hidden_states = ( - hidden_states[None, :].reshape(batch_size, num_frames, channels, height, width).permute(0, 2, 1, 3, 4) - ) - - if temb is not None: - temb = temb.reshape(batch_size, num_frames, -1) - - hidden_states = self.temporal_res_block(hidden_states, temb) - hidden_states = self.time_mixer( - x_spatial=hidden_states_mix, - x_temporal=hidden_states, - image_only_indicator=image_only_indicator, - ) - - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).reshape(batch_frames, channels, height, width) - return hidden_states - - -class AlphaBlender(nn.Module): - r""" - A module to blend spatial and temporal features. - - Parameters: - alpha (`float`): The initial value of the blending factor. - merge_strategy (`str`, *optional*, defaults to `learned_with_images`): - The merge strategy to use for the temporal mixing. - switch_spatial_to_temporal_mix (`bool`, *optional*, defaults to `False`): - If `True`, switch the spatial and temporal mixing. - """ - - strategies = ["learned", "fixed", "learned_with_images"] - - def __init__( - self, - alpha: float, - merge_strategy: str = "learned_with_images", - switch_spatial_to_temporal_mix: bool = False, - ): - super().__init__() - self.merge_strategy = merge_strategy - self.switch_spatial_to_temporal_mix = switch_spatial_to_temporal_mix # For TemporalVAE - - if merge_strategy not in self.strategies: - raise ValueError(f"merge_strategy needs to be in {self.strategies}") - - if self.merge_strategy == "fixed": - self.register_buffer("mix_factor", torch.Tensor([alpha])) - elif self.merge_strategy == "learned" or self.merge_strategy == "learned_with_images": - self.register_parameter("mix_factor", torch.nn.Parameter(torch.Tensor([alpha]))) - else: - raise ValueError(f"Unknown merge strategy {self.merge_strategy}") - - def get_alpha(self, image_only_indicator: torch.Tensor, ndims: int) -> torch.Tensor: - if self.merge_strategy == "fixed": - alpha = self.mix_factor - - elif self.merge_strategy == "learned": - alpha = torch.sigmoid(self.mix_factor) - - elif self.merge_strategy == "learned_with_images": - if image_only_indicator is None: - raise ValueError("Please provide image_only_indicator to use learned_with_images merge strategy") - - alpha = torch.where( - image_only_indicator.bool(), - torch.ones(1, 1, device=image_only_indicator.device), - torch.sigmoid(self.mix_factor)[..., None], - ) - - # (batch, channel, frames, height, width) - if ndims == 5: - alpha = alpha[:, None, :, None, None] - # (batch*frames, height*width, channels) - elif ndims == 3: - alpha = alpha.reshape(-1)[:, None, None] - else: - raise ValueError(f"Unexpected ndims {ndims}. Dimensions should be 3 or 5") - - else: - raise NotImplementedError - - return alpha - - def forward( - self, - x_spatial: torch.Tensor, - x_temporal: torch.Tensor, - image_only_indicator: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - alpha = self.get_alpha(image_only_indicator, x_spatial.ndim) - alpha = alpha.to(x_spatial.dtype) - - if self.switch_spatial_to_temporal_mix: - alpha = 1.0 - alpha - - x = alpha * x_spatial + (1.0 - alpha) * x_temporal - return x diff --git a/venv/lib/python3.10/site-packages/diffusers/models/resnet_flax.py b/venv/lib/python3.10/site-packages/diffusers/models/resnet_flax.py deleted file mode 100644 index 9c80932c5c5d7b4867d6f9d77849ab499d301cdb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/resnet_flax.py +++ /dev/null @@ -1,124 +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 flax.linen as nn -import jax -import jax.numpy as jnp - - -class FlaxUpsample2D(nn.Module): - out_channels: int - dtype: jnp.dtype = jnp.float32 - - def setup(self): - self.conv = nn.Conv( - self.out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - def __call__(self, hidden_states): - batch, height, width, channels = hidden_states.shape - hidden_states = jax.image.resize( - hidden_states, - shape=(batch, height * 2, width * 2, channels), - method="nearest", - ) - hidden_states = self.conv(hidden_states) - return hidden_states - - -class FlaxDownsample2D(nn.Module): - out_channels: int - dtype: jnp.dtype = jnp.float32 - - def setup(self): - self.conv = nn.Conv( - self.out_channels, - kernel_size=(3, 3), - strides=(2, 2), - padding=((1, 1), (1, 1)), # padding="VALID", - dtype=self.dtype, - ) - - def __call__(self, hidden_states): - # pad = ((0, 0), (0, 1), (0, 1), (0, 0)) # pad height and width dim - # hidden_states = jnp.pad(hidden_states, pad_width=pad) - hidden_states = self.conv(hidden_states) - return hidden_states - - -class FlaxResnetBlock2D(nn.Module): - in_channels: int - out_channels: int = None - dropout_prob: float = 0.0 - use_nin_shortcut: bool = None - dtype: jnp.dtype = jnp.float32 - - def setup(self): - out_channels = self.in_channels if self.out_channels is None else self.out_channels - - self.norm1 = nn.GroupNorm(num_groups=32, epsilon=1e-5) - self.conv1 = nn.Conv( - out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - self.time_emb_proj = nn.Dense(out_channels, dtype=self.dtype) - - self.norm2 = nn.GroupNorm(num_groups=32, epsilon=1e-5) - self.dropout = nn.Dropout(self.dropout_prob) - self.conv2 = nn.Conv( - out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - use_nin_shortcut = self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut - - self.conv_shortcut = None - if use_nin_shortcut: - self.conv_shortcut = nn.Conv( - out_channels, - kernel_size=(1, 1), - strides=(1, 1), - padding="VALID", - dtype=self.dtype, - ) - - def __call__(self, hidden_states, temb, deterministic=True): - residual = hidden_states - hidden_states = self.norm1(hidden_states) - hidden_states = nn.swish(hidden_states) - hidden_states = self.conv1(hidden_states) - - temb = self.time_emb_proj(nn.swish(temb)) - temb = jnp.expand_dims(jnp.expand_dims(temb, 1), 1) - hidden_states = hidden_states + temb - - hidden_states = self.norm2(hidden_states) - hidden_states = nn.swish(hidden_states) - hidden_states = self.dropout(hidden_states, deterministic) - hidden_states = self.conv2(hidden_states) - - if self.conv_shortcut is not None: - residual = self.conv_shortcut(residual) - - return hidden_states + residual diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__init__.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__init__.py deleted file mode 100644 index cc03a0ccbcdf12c4df4280c3ec3450df3e5a1d87..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__init__.py +++ /dev/null @@ -1,36 +0,0 @@ -from ...utils import is_torch_available - - -if is_torch_available(): - from .auraflow_transformer_2d import AuraFlowTransformer2DModel - from .cogvideox_transformer_3d import CogVideoXTransformer3DModel - from .consisid_transformer_3d import ConsisIDTransformer3DModel - from .dit_transformer_2d import DiTTransformer2DModel - from .dual_transformer_2d import DualTransformer2DModel - from .hunyuan_transformer_2d import HunyuanDiT2DModel - from .latte_transformer_3d import LatteTransformer3DModel - from .lumina_nextdit2d import LuminaNextDiT2DModel - from .pixart_transformer_2d import PixArtTransformer2DModel - from .prior_transformer import PriorTransformer - from .sana_transformer import SanaTransformer2DModel - from .stable_audio_transformer import StableAudioDiTModel - from .t5_film_transformer import T5FilmDecoder - from .transformer_2d import Transformer2DModel - from .transformer_allegro import AllegroTransformer3DModel - from .transformer_chroma import ChromaTransformer2DModel - from .transformer_cogview3plus import CogView3PlusTransformer2DModel - from .transformer_cogview4 import CogView4Transformer2DModel - from .transformer_cosmos import CosmosTransformer3DModel - from .transformer_easyanimate import EasyAnimateTransformer3DModel - from .transformer_flux import FluxTransformer2DModel - from .transformer_hidream_image import HiDreamImageTransformer2DModel - from .transformer_hunyuan_video import HunyuanVideoTransformer3DModel - from .transformer_hunyuan_video_framepack import HunyuanVideoFramepackTransformer3DModel - from .transformer_ltx import LTXVideoTransformer3DModel - from .transformer_lumina2 import Lumina2Transformer2DModel - from .transformer_mochi import MochiTransformer3DModel - from .transformer_omnigen import OmniGenTransformer2DModel - from .transformer_sd3 import SD3Transformer2DModel - from .transformer_temporal import TransformerTemporalModel - from .transformer_wan import WanTransformer3DModel - from .transformer_wan_vace import WanVACETransformer3DModel diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 3d2faa56910e614161c694a4b5aa64dc83e089b2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/auraflow_transformer_2d.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/auraflow_transformer_2d.cpython-310.pyc deleted file mode 100644 index 1ffe39a305adeeae446513ee5e1f4440e6c54f08..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/auraflow_transformer_2d.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/cogvideox_transformer_3d.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/cogvideox_transformer_3d.cpython-310.pyc deleted file mode 100644 index 1bb742a223ee727f3deb1be089ce329886ca317d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/cogvideox_transformer_3d.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/consisid_transformer_3d.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/consisid_transformer_3d.cpython-310.pyc deleted file mode 100644 index 20f1b69a9c4fa523ec09885d239850400b733022..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/consisid_transformer_3d.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/dit_transformer_2d.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/dit_transformer_2d.cpython-310.pyc deleted file mode 100644 index 802607f8aaa080ad036ea599fecf1b183eb85546..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/dit_transformer_2d.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/dual_transformer_2d.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/dual_transformer_2d.cpython-310.pyc deleted file mode 100644 index 835c3f8d313445a79cb0711cd96a71c93f462bdf..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/dual_transformer_2d.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/hunyuan_transformer_2d.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/hunyuan_transformer_2d.cpython-310.pyc deleted file mode 100644 index fa2cdef50ce5a6bd44aa0491bd25642a188db851..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/hunyuan_transformer_2d.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/latte_transformer_3d.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/latte_transformer_3d.cpython-310.pyc deleted file mode 100644 index c3457a3bd15a692d01962a0c0f73a100a95d3edb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/latte_transformer_3d.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/lumina_nextdit2d.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/lumina_nextdit2d.cpython-310.pyc deleted file mode 100644 index 7f559cbc967885a345bfc8f8afe5f421d2627b70..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/lumina_nextdit2d.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/pixart_transformer_2d.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/pixart_transformer_2d.cpython-310.pyc deleted file mode 100644 index 0fc63f149e9a5b8a8671f186086f32430fc997d4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/pixart_transformer_2d.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/prior_transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/prior_transformer.cpython-310.pyc deleted file mode 100644 index 1cd7dcb71ee712e067797fcf5463a96551bfa82e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/prior_transformer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/sana_transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/sana_transformer.cpython-310.pyc deleted file mode 100644 index f8445683fc90730ae61c129bbad2e1abf7f6d100..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/sana_transformer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/stable_audio_transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/stable_audio_transformer.cpython-310.pyc deleted file mode 100644 index 2d384f8539a6af34fb993060085343da310e0836..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/stable_audio_transformer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/t5_film_transformer.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/t5_film_transformer.cpython-310.pyc deleted file mode 100644 index 8ca82e8c470e0db3c30b4a1582a87a01231ed02b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/t5_film_transformer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_2d.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_2d.cpython-310.pyc deleted file mode 100644 index 66482a84cabd3d02e96a821a92723bfca6c327c5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_2d.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_allegro.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_allegro.cpython-310.pyc deleted file mode 100644 index 3ac40c64035c4a43e541c01c90ec45ece082281f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_allegro.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_chroma.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_chroma.cpython-310.pyc deleted file mode 100644 index f180aa0725464e51b94666c484c4697777951ca2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_chroma.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_cogview3plus.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_cogview3plus.cpython-310.pyc deleted file mode 100644 index ef2108477549ed159f88fed8c9674c5f8299614c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_cogview3plus.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_cogview4.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_cogview4.cpython-310.pyc deleted file mode 100644 index d7d0f37a20ba114ec87c90f647dc9c769a1aa17a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_cogview4.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_cosmos.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_cosmos.cpython-310.pyc deleted file mode 100644 index e18a19189a6180e5227809055b601233e78c7dfd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_cosmos.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_easyanimate.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_easyanimate.cpython-310.pyc deleted file mode 100644 index 9492888c22c2aee80efbcf2e8360a3d9b7ed7be0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_easyanimate.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_flux.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_flux.cpython-310.pyc deleted file mode 100644 index 6ee05b031ec7ec04dd9aac08d79e2a61c82b9532..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_flux.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_hidream_image.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_hidream_image.cpython-310.pyc deleted file mode 100644 index 9b50d9eca2fea20697c3075404852ee12102791d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_hidream_image.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_hunyuan_video.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_hunyuan_video.cpython-310.pyc deleted file mode 100644 index ea31057e7253d3aff2908d3a9afb42246ae43078..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_hunyuan_video.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_hunyuan_video_framepack.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_hunyuan_video_framepack.cpython-310.pyc deleted file mode 100644 index 13aa6aee5507285ee8160cbfa7f957f7618e840f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_hunyuan_video_framepack.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_ltx.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_ltx.cpython-310.pyc deleted file mode 100644 index fd7768f655437f358ce7d4cf384b571927276701..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_ltx.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_lumina2.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_lumina2.cpython-310.pyc deleted file mode 100644 index 76931fbf74f857caa909fa4c0e84b5a3d7335844..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_lumina2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_mochi.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_mochi.cpython-310.pyc deleted file mode 100644 index 97eabe2817a23c30b9ad94a7e67026234297c2e3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_mochi.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_omnigen.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_omnigen.cpython-310.pyc deleted file mode 100644 index eb5d24b0baee0a8837814a4feda98f9965c3c04a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_omnigen.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_sd3.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_sd3.cpython-310.pyc deleted file mode 100644 index af19a141bc3a85e18f98138b989e44e0958e059c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_sd3.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_temporal.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_temporal.cpython-310.pyc deleted file mode 100644 index 633b36877a163afdf056e8e72a29c524811742ae..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_temporal.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_wan.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_wan.cpython-310.pyc deleted file mode 100644 index 998ccdad6a0b9954b1b1f09dc4d5a8d3e67c7968..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_wan.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_wan_vace.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_wan_vace.cpython-310.pyc deleted file mode 100644 index 0f679edc0965436ed3820dbcb3d173c8748a30b5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/transformers/__pycache__/transformer_wan_vace.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/auraflow_transformer_2d.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/auraflow_transformer_2d.py deleted file mode 100644 index a8d275d14214977eb92fe17fa49cb2e6b4d891f1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/auraflow_transformer_2d.py +++ /dev/null @@ -1,564 +0,0 @@ -# Copyright 2025 AuraFlow Authors, 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, Dict, Optional, Union - -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 USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers -from ...utils.torch_utils import maybe_allow_in_graph -from ..attention_processor import ( - Attention, - AttentionProcessor, - AuraFlowAttnProcessor2_0, - FusedAuraFlowAttnProcessor2_0, -) -from ..embeddings import TimestepEmbedding, Timesteps -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormZero, FP32LayerNorm - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -# Taken from the original aura flow inference code. -def find_multiple(n: int, k: int) -> int: - if n % k == 0: - return n - return n + k - (n % k) - - -# Aura Flow patch embed doesn't use convs for projections. -# Additionally, it uses learned positional embeddings. -class AuraFlowPatchEmbed(nn.Module): - def __init__( - self, - height=224, - width=224, - patch_size=16, - in_channels=3, - embed_dim=768, - pos_embed_max_size=None, - ): - super().__init__() - - self.num_patches = (height // patch_size) * (width // patch_size) - self.pos_embed_max_size = pos_embed_max_size - - self.proj = nn.Linear(patch_size * patch_size * in_channels, embed_dim) - self.pos_embed = nn.Parameter(torch.randn(1, pos_embed_max_size, embed_dim) * 0.1) - - self.patch_size = patch_size - self.height, self.width = height // patch_size, width // patch_size - self.base_size = height // patch_size - - def pe_selection_index_based_on_dim(self, h, w): - # select subset of positional embedding based on H, W, where H, W is size of latent - # PE will be viewed as 2d-grid, and H/p x W/p of the PE will be selected - # because original input are in flattened format, we have to flatten this 2d grid as well. - h_p, w_p = h // self.patch_size, w // self.patch_size - h_max, w_max = int(self.pos_embed_max_size**0.5), int(self.pos_embed_max_size**0.5) - - # Calculate the top-left corner indices for the centered patch grid - starth = h_max // 2 - h_p // 2 - startw = w_max // 2 - w_p // 2 - - # Generate the row and column indices for the desired patch grid - rows = torch.arange(starth, starth + h_p, device=self.pos_embed.device) - cols = torch.arange(startw, startw + w_p, device=self.pos_embed.device) - - # Create a 2D grid of indices - row_indices, col_indices = torch.meshgrid(rows, cols, indexing="ij") - - # Convert the 2D grid indices to flattened 1D indices - selected_indices = (row_indices * w_max + col_indices).flatten() - - return selected_indices - - def forward(self, latent): - batch_size, num_channels, height, width = latent.size() - latent = latent.view( - batch_size, - num_channels, - height // self.patch_size, - self.patch_size, - width // self.patch_size, - self.patch_size, - ) - latent = latent.permute(0, 2, 4, 1, 3, 5).flatten(-3).flatten(1, 2) - latent = self.proj(latent) - pe_index = self.pe_selection_index_based_on_dim(height, width) - return latent + self.pos_embed[:, pe_index] - - -# Taken from the original Aura flow inference code. -# Our feedforward only has GELU but Aura uses SiLU. -class AuraFlowFeedForward(nn.Module): - def __init__(self, dim, hidden_dim=None) -> None: - super().__init__() - if hidden_dim is None: - hidden_dim = 4 * dim - - final_hidden_dim = int(2 * hidden_dim / 3) - final_hidden_dim = find_multiple(final_hidden_dim, 256) - - self.linear_1 = nn.Linear(dim, final_hidden_dim, bias=False) - self.linear_2 = nn.Linear(dim, final_hidden_dim, bias=False) - self.out_projection = nn.Linear(final_hidden_dim, dim, bias=False) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = F.silu(self.linear_1(x)) * self.linear_2(x) - x = self.out_projection(x) - return x - - -class AuraFlowPreFinalBlock(nn.Module): - def __init__(self, embedding_dim: int, conditioning_embedding_dim: int): - super().__init__() - - self.silu = nn.SiLU() - self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=False) - - def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor: - emb = self.linear(self.silu(conditioning_embedding).to(x.dtype)) - scale, shift = torch.chunk(emb, 2, dim=1) - x = x * (1 + scale)[:, None, :] + shift[:, None, :] - return x - - -@maybe_allow_in_graph -class AuraFlowSingleTransformerBlock(nn.Module): - """Similar to `AuraFlowJointTransformerBlock` with a single DiT instead of an MMDiT.""" - - def __init__(self, dim, num_attention_heads, attention_head_dim): - super().__init__() - - self.norm1 = AdaLayerNormZero(dim, bias=False, norm_type="fp32_layer_norm") - - processor = AuraFlowAttnProcessor2_0() - self.attn = Attention( - query_dim=dim, - cross_attention_dim=None, - dim_head=attention_head_dim, - heads=num_attention_heads, - qk_norm="fp32_layer_norm", - out_dim=dim, - bias=False, - out_bias=False, - processor=processor, - ) - - self.norm2 = FP32LayerNorm(dim, elementwise_affine=False, bias=False) - self.ff = AuraFlowFeedForward(dim, dim * 4) - - def forward( - self, - hidden_states: torch.FloatTensor, - temb: torch.FloatTensor, - attention_kwargs: Optional[Dict[str, Any]] = None, - ): - residual = hidden_states - attention_kwargs = attention_kwargs or {} - - # Norm + Projection. - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) - - # Attention. - attn_output = self.attn(hidden_states=norm_hidden_states, **attention_kwargs) - - # Process attention outputs for the `hidden_states`. - hidden_states = self.norm2(residual + gate_msa.unsqueeze(1) * attn_output) - hidden_states = hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - ff_output = self.ff(hidden_states) - hidden_states = gate_mlp.unsqueeze(1) * ff_output - hidden_states = residual + hidden_states - - return hidden_states - - -@maybe_allow_in_graph -class AuraFlowJointTransformerBlock(nn.Module): - r""" - Transformer block for Aura Flow. Similar to SD3 MMDiT. Differences (non-exhaustive): - - * QK Norm in the attention blocks - * No bias in the attention blocks - * Most LayerNorms are in FP32 - - Parameters: - 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. - is_last (`bool`): Boolean to determine if this is the last block in the model. - """ - - def __init__(self, dim, num_attention_heads, attention_head_dim): - super().__init__() - - self.norm1 = AdaLayerNormZero(dim, bias=False, norm_type="fp32_layer_norm") - self.norm1_context = AdaLayerNormZero(dim, bias=False, norm_type="fp32_layer_norm") - - processor = AuraFlowAttnProcessor2_0() - self.attn = Attention( - query_dim=dim, - cross_attention_dim=None, - added_kv_proj_dim=dim, - added_proj_bias=False, - dim_head=attention_head_dim, - heads=num_attention_heads, - qk_norm="fp32_layer_norm", - out_dim=dim, - bias=False, - out_bias=False, - processor=processor, - context_pre_only=False, - ) - - self.norm2 = FP32LayerNorm(dim, elementwise_affine=False, bias=False) - self.ff = AuraFlowFeedForward(dim, dim * 4) - self.norm2_context = FP32LayerNorm(dim, elementwise_affine=False, bias=False) - self.ff_context = AuraFlowFeedForward(dim, dim * 4) - - def forward( - self, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor, - temb: torch.FloatTensor, - attention_kwargs: Optional[Dict[str, Any]] = None, - ): - residual = hidden_states - residual_context = encoder_hidden_states - attention_kwargs = attention_kwargs or {} - - # Norm + Projection. - 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( - encoder_hidden_states, emb=temb - ) - - # Attention. - attn_output, context_attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - **attention_kwargs, - ) - - # Process attention outputs for the `hidden_states`. - hidden_states = self.norm2(residual + gate_msa.unsqueeze(1) * attn_output) - hidden_states = hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - hidden_states = gate_mlp.unsqueeze(1) * self.ff(hidden_states) - hidden_states = residual + hidden_states - - # Process attention outputs for the `encoder_hidden_states`. - encoder_hidden_states = self.norm2_context(residual_context + c_gate_msa.unsqueeze(1) * context_attn_output) - encoder_hidden_states = encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] - encoder_hidden_states = c_gate_mlp.unsqueeze(1) * self.ff_context(encoder_hidden_states) - encoder_hidden_states = residual_context + encoder_hidden_states - - return encoder_hidden_states, hidden_states - - -class AuraFlowTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): - r""" - A 2D Transformer model as introduced in AuraFlow (https://blog.fal.ai/auraflow/). - - Parameters: - sample_size (`int`): The width of the latent images. This is fixed during training since - it is used to learn a number of position embeddings. - patch_size (`int`): Patch size to turn the input data into small patches. - in_channels (`int`, *optional*, defaults to 4): The number of channels in the input. - num_mmdit_layers (`int`, *optional*, defaults to 4): The number of layers of MMDiT Transformer blocks to use. - num_single_dit_layers (`int`, *optional*, defaults to 32): - The number of layers of Transformer blocks to use. These blocks use concatenated image and text - representations. - attention_head_dim (`int`, *optional*, defaults to 256): The number of channels in each head. - num_attention_heads (`int`, *optional*, defaults to 12): The number of heads to use for multi-head attention. - joint_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use. - caption_projection_dim (`int`): Number of dimensions to use when projecting the `encoder_hidden_states`. - out_channels (`int`, defaults to 4): Number of output channels. - pos_embed_max_size (`int`, defaults to 1024): Maximum positions to embed from the image latents. - """ - - _no_split_modules = ["AuraFlowJointTransformerBlock", "AuraFlowSingleTransformerBlock", "AuraFlowPatchEmbed"] - _skip_layerwise_casting_patterns = ["pos_embed", "norm"] - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - sample_size: int = 64, - patch_size: int = 2, - in_channels: int = 4, - num_mmdit_layers: int = 4, - num_single_dit_layers: int = 32, - attention_head_dim: int = 256, - num_attention_heads: int = 12, - joint_attention_dim: int = 2048, - caption_projection_dim: int = 3072, - out_channels: int = 4, - pos_embed_max_size: int = 1024, - ): - super().__init__() - default_out_channels = in_channels - self.out_channels = out_channels if out_channels is not None else default_out_channels - self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim - - self.pos_embed = AuraFlowPatchEmbed( - height=self.config.sample_size, - width=self.config.sample_size, - patch_size=self.config.patch_size, - in_channels=self.config.in_channels, - embed_dim=self.inner_dim, - pos_embed_max_size=pos_embed_max_size, - ) - - self.context_embedder = nn.Linear( - self.config.joint_attention_dim, self.config.caption_projection_dim, bias=False - ) - self.time_step_embed = Timesteps(num_channels=256, downscale_freq_shift=0, scale=1000, flip_sin_to_cos=True) - self.time_step_proj = TimestepEmbedding(in_channels=256, time_embed_dim=self.inner_dim) - - self.joint_transformer_blocks = nn.ModuleList( - [ - AuraFlowJointTransformerBlock( - dim=self.inner_dim, - num_attention_heads=self.config.num_attention_heads, - attention_head_dim=self.config.attention_head_dim, - ) - for i in range(self.config.num_mmdit_layers) - ] - ) - self.single_transformer_blocks = nn.ModuleList( - [ - AuraFlowSingleTransformerBlock( - dim=self.inner_dim, - num_attention_heads=self.config.num_attention_heads, - attention_head_dim=self.config.attention_head_dim, - ) - for _ in range(self.config.num_single_dit_layers) - ] - ) - - self.norm_out = AuraFlowPreFinalBlock(self.inner_dim, self.inner_dim) - self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=False) - - # https://huggingface.co/papers/2309.16588 - # prevents artifacts in the attention maps - self.register_tokens = nn.Parameter(torch.randn(1, 8, self.inner_dim) * 0.02) - - self.gradient_checkpointing = False - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedAuraFlowAttnProcessor2_0 - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - self.set_attn_processor(FusedAuraFlowAttnProcessor2_0()) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - if self.original_attn_processors is not None: - self.set_attn_processor(self.original_attn_processors) - - def forward( - self, - hidden_states: torch.FloatTensor, - encoder_hidden_states: torch.FloatTensor = None, - timestep: torch.LongTensor = None, - attention_kwargs: Optional[Dict[str, Any]] = None, - return_dict: bool = True, - ) -> Union[torch.FloatTensor, 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." - ) - - height, width = hidden_states.shape[-2:] - - # Apply patch embedding, timestep embedding, and project the caption embeddings. - hidden_states = self.pos_embed(hidden_states) # takes care of adding positional embeddings too. - temb = self.time_step_embed(timestep).to(dtype=next(self.parameters()).dtype) - temb = self.time_step_proj(temb) - encoder_hidden_states = self.context_embedder(encoder_hidden_states) - encoder_hidden_states = torch.cat( - [self.register_tokens.repeat(encoder_hidden_states.size(0), 1, 1), encoder_hidden_states], dim=1 - ) - - # MMDiT blocks. - for index_block, block in enumerate(self.joint_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, - temb, - ) - - else: - encoder_hidden_states, hidden_states = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - temb=temb, - attention_kwargs=attention_kwargs, - ) - - # Single DiT blocks that combine the `hidden_states` (image) and `encoder_hidden_states` (text) - if len(self.single_transformer_blocks) > 0: - encoder_seq_len = encoder_hidden_states.size(1) - combined_hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) - - for index_block, block in enumerate(self.single_transformer_blocks): - if torch.is_grad_enabled() and self.gradient_checkpointing: - combined_hidden_states = self._gradient_checkpointing_func( - block, - combined_hidden_states, - temb, - ) - - else: - combined_hidden_states = block( - hidden_states=combined_hidden_states, temb=temb, attention_kwargs=attention_kwargs - ) - - hidden_states = combined_hidden_states[:, encoder_seq_len:] - - hidden_states = self.norm_out(hidden_states, temb) - hidden_states = self.proj_out(hidden_states) - - # unpatchify - patch_size = self.config.patch_size - out_channels = self.config.out_channels - height = height // patch_size - width = width // patch_size - - hidden_states = hidden_states.reshape( - shape=(hidden_states.shape[0], height, width, patch_size, patch_size, out_channels) - ) - hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) - output = hidden_states.reshape( - 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,) - - return Transformer2DModelOutput(sample=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/cogvideox_transformer_3d.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/cogvideox_transformer_3d.py deleted file mode 100644 index a8c98bccb86ce445b7ce8bf8b29245dad3178518..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/cogvideox_transformer_3d.py +++ /dev/null @@ -1,531 +0,0 @@ -# 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. - -from typing import Any, Dict, Optional, Tuple, Union - -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.torch_utils import maybe_allow_in_graph -from ..attention import Attention, FeedForward -from ..attention_processor import AttentionProcessor, CogVideoXAttnProcessor2_0, FusedCogVideoXAttnProcessor2_0 -from ..cache_utils import CacheMixin -from ..embeddings import CogVideoXPatchEmbed, TimestepEmbedding, Timesteps -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNorm, CogVideoXLayerNormZero - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@maybe_allow_in_graph -class CogVideoXBlock(nn.Module): - r""" - Transformer block used in [CogVideoX](https://github.com/THUDM/CogVideo) model. - - Parameters: - 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. - time_embed_dim (`int`): - The number of channels in timestep embedding. - dropout (`float`, defaults to `0.0`): - The dropout probability to use. - activation_fn (`str`, defaults to `"gelu-approximate"`): - Activation function to be used in feed-forward. - attention_bias (`bool`, defaults to `False`): - Whether or not to use bias in attention projection layers. - qk_norm (`bool`, defaults to `True`): - Whether or not to use normalization after query and key projections in Attention. - norm_elementwise_affine (`bool`, defaults to `True`): - Whether to use learnable elementwise affine parameters for normalization. - norm_eps (`float`, defaults to `1e-5`): - Epsilon value for normalization layers. - final_dropout (`bool` defaults to `False`): - Whether to apply a final dropout after the last feed-forward layer. - ff_inner_dim (`int`, *optional*, defaults to `None`): - Custom hidden dimension of Feed-forward layer. If not provided, `4 * dim` is used. - ff_bias (`bool`, defaults to `True`): - Whether or not to use bias in Feed-forward layer. - attention_out_bias (`bool`, defaults to `True`): - Whether or not to use bias in Attention output projection layer. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - time_embed_dim: int, - dropout: float = 0.0, - activation_fn: str = "gelu-approximate", - attention_bias: bool = False, - qk_norm: bool = True, - norm_elementwise_affine: bool = True, - norm_eps: float = 1e-5, - final_dropout: bool = True, - ff_inner_dim: Optional[int] = None, - ff_bias: bool = True, - attention_out_bias: bool = True, - ): - super().__init__() - - # 1. Self Attention - self.norm1 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True) - - self.attn1 = Attention( - query_dim=dim, - dim_head=attention_head_dim, - heads=num_attention_heads, - qk_norm="layer_norm" if qk_norm else None, - eps=1e-6, - bias=attention_bias, - out_bias=attention_out_bias, - processor=CogVideoXAttnProcessor2_0(), - ) - - # 2. Feed Forward - self.norm2 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True) - - self.ff = FeedForward( - dim, - dropout=dropout, - activation_fn=activation_fn, - final_dropout=final_dropout, - inner_dim=ff_inner_dim, - bias=ff_bias, - ) - - def forward( - self, - 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, - ) -> torch.Tensor: - text_seq_length = encoder_hidden_states.size(1) - attention_kwargs = attention_kwargs or {} - - # norm & modulate - norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1( - hidden_states, encoder_hidden_states, temb - ) - - # attention - 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_kwargs, - ) - - hidden_states = hidden_states + gate_msa * attn_hidden_states - encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states - - # norm & modulate - norm_hidden_states, norm_encoder_hidden_states, gate_ff, enc_gate_ff = self.norm2( - hidden_states, encoder_hidden_states, temb - ) - - # feed-forward - norm_hidden_states = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1) - ff_output = self.ff(norm_hidden_states) - - hidden_states = hidden_states + gate_ff * ff_output[:, text_seq_length:] - encoder_hidden_states = encoder_hidden_states + enc_gate_ff * ff_output[:, :text_seq_length] - - return hidden_states, encoder_hidden_states - - -class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, CacheMixin): - """ - A Transformer model for video-like data in [CogVideoX](https://github.com/THUDM/CogVideo). - - Parameters: - num_attention_heads (`int`, defaults to `30`): - The number of heads to use for multi-head attention. - attention_head_dim (`int`, defaults to `64`): - The number of channels in each head. - in_channels (`int`, defaults to `16`): - The number of channels in the input. - out_channels (`int`, *optional*, defaults to `16`): - The number of channels in the output. - flip_sin_to_cos (`bool`, defaults to `True`): - Whether to flip the sin to cos in the time embedding. - time_embed_dim (`int`, defaults to `512`): - Output dimension of timestep embeddings. - ofs_embed_dim (`int`, defaults to `512`): - Output dimension of "ofs" embeddings used in CogVideoX-5b-I2B in version 1.5 - text_embed_dim (`int`, defaults to `4096`): - Input dimension of text embeddings from the text encoder. - num_layers (`int`, defaults to `30`): - The number of layers of Transformer blocks to use. - dropout (`float`, defaults to `0.0`): - The dropout probability to use. - attention_bias (`bool`, defaults to `True`): - Whether to use bias in the attention projection layers. - sample_width (`int`, defaults to `90`): - The width of the input latents. - sample_height (`int`, defaults to `60`): - The height of the input latents. - sample_frames (`int`, defaults to `49`): - The number of frames in the input latents. Note that this parameter was incorrectly initialized to 49 - instead of 13 because CogVideoX processed 13 latent frames at once in its default and recommended settings, - but cannot be changed to the correct value to ensure backwards compatibility. To create a transformer with - K latent frames, the correct value to pass here would be: ((K - 1) * temporal_compression_ratio + 1). - patch_size (`int`, defaults to `2`): - The size of the patches to use in the patch embedding layer. - temporal_compression_ratio (`int`, defaults to `4`): - The compression ratio across the temporal dimension. See documentation for `sample_frames`. - max_text_seq_length (`int`, defaults to `226`): - The maximum sequence length of the input text embeddings. - activation_fn (`str`, defaults to `"gelu-approximate"`): - Activation function to use in feed-forward. - timestep_activation_fn (`str`, defaults to `"silu"`): - Activation function to use when generating the timestep embeddings. - norm_elementwise_affine (`bool`, defaults to `True`): - Whether to use elementwise affine in normalization layers. - norm_eps (`float`, defaults to `1e-5`): - The epsilon value to use in normalization layers. - spatial_interpolation_scale (`float`, defaults to `1.875`): - Scaling factor to apply in 3D positional embeddings across spatial dimensions. - temporal_interpolation_scale (`float`, defaults to `1.0`): - Scaling factor to apply in 3D positional embeddings across temporal dimensions. - """ - - _skip_layerwise_casting_patterns = ["patch_embed", "norm"] - _supports_gradient_checkpointing = True - _no_split_modules = ["CogVideoXBlock", "CogVideoXPatchEmbed"] - - @register_to_config - def __init__( - self, - num_attention_heads: int = 30, - attention_head_dim: int = 64, - in_channels: int = 16, - out_channels: Optional[int] = 16, - flip_sin_to_cos: bool = True, - freq_shift: int = 0, - time_embed_dim: int = 512, - ofs_embed_dim: Optional[int] = None, - text_embed_dim: int = 4096, - num_layers: int = 30, - dropout: float = 0.0, - attention_bias: bool = True, - sample_width: int = 90, - sample_height: int = 60, - sample_frames: int = 49, - patch_size: int = 2, - patch_size_t: Optional[int] = None, - temporal_compression_ratio: int = 4, - max_text_seq_length: int = 226, - activation_fn: str = "gelu-approximate", - timestep_activation_fn: str = "silu", - norm_elementwise_affine: bool = True, - norm_eps: float = 1e-5, - spatial_interpolation_scale: float = 1.875, - temporal_interpolation_scale: float = 1.0, - use_rotary_positional_embeddings: bool = False, - use_learned_positional_embeddings: bool = False, - patch_bias: bool = True, - ): - super().__init__() - inner_dim = num_attention_heads * attention_head_dim - - if not use_rotary_positional_embeddings and use_learned_positional_embeddings: - raise ValueError( - "There are no CogVideoX checkpoints available with disable rotary embeddings and learned positional " - "embeddings. If you're using a custom model and/or believe this should be supported, please open an " - "issue at https://github.com/huggingface/diffusers/issues." - ) - - # 1. Patch embedding - self.patch_embed = CogVideoXPatchEmbed( - patch_size=patch_size, - patch_size_t=patch_size_t, - in_channels=in_channels, - embed_dim=inner_dim, - text_embed_dim=text_embed_dim, - bias=patch_bias, - sample_width=sample_width, - sample_height=sample_height, - sample_frames=sample_frames, - temporal_compression_ratio=temporal_compression_ratio, - max_text_seq_length=max_text_seq_length, - spatial_interpolation_scale=spatial_interpolation_scale, - temporal_interpolation_scale=temporal_interpolation_scale, - use_positional_embeddings=not use_rotary_positional_embeddings, - use_learned_positional_embeddings=use_learned_positional_embeddings, - ) - self.embedding_dropout = nn.Dropout(dropout) - - # 2. Time embeddings and ofs embedding(Only CogVideoX1.5-5B I2V have) - - self.time_proj = Timesteps(inner_dim, flip_sin_to_cos, freq_shift) - self.time_embedding = TimestepEmbedding(inner_dim, time_embed_dim, timestep_activation_fn) - - self.ofs_proj = None - self.ofs_embedding = None - if ofs_embed_dim: - self.ofs_proj = Timesteps(ofs_embed_dim, flip_sin_to_cos, freq_shift) - self.ofs_embedding = TimestepEmbedding( - ofs_embed_dim, ofs_embed_dim, timestep_activation_fn - ) # same as time embeddings, for ofs - - # 3. Define spatio-temporal transformers blocks - self.transformer_blocks = nn.ModuleList( - [ - CogVideoXBlock( - dim=inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - time_embed_dim=time_embed_dim, - dropout=dropout, - activation_fn=activation_fn, - attention_bias=attention_bias, - norm_elementwise_affine=norm_elementwise_affine, - norm_eps=norm_eps, - ) - for _ in range(num_layers) - ] - ) - self.norm_final = nn.LayerNorm(inner_dim, norm_eps, norm_elementwise_affine) - - # 4. Output blocks - self.norm_out = AdaLayerNorm( - embedding_dim=time_embed_dim, - output_dim=2 * inner_dim, - norm_elementwise_affine=norm_elementwise_affine, - norm_eps=norm_eps, - chunk_dim=1, - ) - - if patch_size_t is None: - # For CogVideox 1.0 - output_dim = patch_size * patch_size * out_channels - else: - # For CogVideoX 1.5 - output_dim = patch_size * patch_size * patch_size_t * out_channels - - self.proj_out = nn.Linear(inner_dim, output_dim) - - self.gradient_checkpointing = False - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedCogVideoXAttnProcessor2_0 - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - self.set_attn_processor(FusedCogVideoXAttnProcessor2_0()) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - if self.original_attn_processors is not None: - self.set_attn_processor(self.original_attn_processors) - - 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, - return_dict: bool = True, - ): - 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_frames, channels, height, width = hidden_states.shape - - # 1. Time embedding - timesteps = timestep - t_emb = self.time_proj(timesteps) - - # timesteps does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - t_emb = t_emb.to(dtype=hidden_states.dtype) - emb = self.time_embedding(t_emb, timestep_cond) - - if self.ofs_embedding is not None: - ofs_emb = self.ofs_proj(ofs) - ofs_emb = ofs_emb.to(dtype=hidden_states.dtype) - ofs_emb = self.ofs_embedding(ofs_emb) - emb = emb + ofs_emb - - # 2. Patch embedding - hidden_states = self.patch_embed(encoder_hidden_states, hidden_states) - hidden_states = self.embedding_dropout(hidden_states) - - text_seq_length = encoder_hidden_states.shape[1] - encoder_hidden_states = hidden_states[:, :text_seq_length] - hidden_states = hidden_states[:, text_seq_length:] - - # 3. Transformer blocks - for i, 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, - emb, - image_rotary_emb, - attention_kwargs, - ) - else: - hidden_states, encoder_hidden_states = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - temb=emb, - image_rotary_emb=image_rotary_emb, - attention_kwargs=attention_kwargs, - ) - - hidden_states = self.norm_final(hidden_states) - - # 4. Final block - hidden_states = self.norm_out(hidden_states, temb=emb) - hidden_states = self.proj_out(hidden_states) - - # 5. Unpatchify - p = self.config.patch_size - p_t = self.config.patch_size_t - - if p_t is None: - 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) - else: - output = hidden_states.reshape( - batch_size, (num_frames + p_t - 1) // p_t, height // p, width // p, -1, p_t, p, p - ) - 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/venv/lib/python3.10/site-packages/diffusers/models/transformers/consisid_transformer_3d.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/consisid_transformer_3d.py deleted file mode 100644 index 41632dbd475119db9793a298ec9887e7a16aaffa..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/consisid_transformer_3d.py +++ /dev/null @@ -1,789 +0,0 @@ -# Copyright 2025 ConsisID Authors 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 typing import Any, Dict, List, Optional, Tuple, Union - -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.torch_utils import maybe_allow_in_graph -from ..attention import Attention, FeedForward -from ..attention_processor import AttentionProcessor, CogVideoXAttnProcessor2_0 -from ..embeddings import CogVideoXPatchEmbed, TimestepEmbedding, Timesteps -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNorm, CogVideoXLayerNormZero - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class PerceiverAttention(nn.Module): - def __init__(self, dim: int, dim_head: int = 64, heads: int = 8, kv_dim: Optional[int] = None): - super().__init__() - - self.scale = dim_head**-0.5 - self.dim_head = dim_head - self.heads = heads - inner_dim = dim_head * heads - - self.norm1 = nn.LayerNorm(dim if kv_dim is None else kv_dim) - self.norm2 = nn.LayerNorm(dim) - - self.to_q = nn.Linear(dim, inner_dim, bias=False) - self.to_kv = nn.Linear(dim if kv_dim is None else kv_dim, inner_dim * 2, bias=False) - self.to_out = nn.Linear(inner_dim, dim, bias=False) - - def forward(self, image_embeds: torch.Tensor, latents: torch.Tensor) -> torch.Tensor: - # Apply normalization - image_embeds = self.norm1(image_embeds) - latents = self.norm2(latents) - - batch_size, seq_len, _ = latents.shape # Get batch size and sequence length - - # Compute query, key, and value matrices - query = self.to_q(latents) - kv_input = torch.cat((image_embeds, latents), dim=-2) - key, value = self.to_kv(kv_input).chunk(2, dim=-1) - - # Reshape the tensors for multi-head attention - query = query.reshape(query.size(0), -1, self.heads, self.dim_head).transpose(1, 2) - key = key.reshape(key.size(0), -1, self.heads, self.dim_head).transpose(1, 2) - value = value.reshape(value.size(0), -1, self.heads, self.dim_head).transpose(1, 2) - - # attention - scale = 1 / math.sqrt(math.sqrt(self.dim_head)) - weight = (query * scale) @ (key * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards - weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) - output = weight @ value - - # Reshape and return the final output - output = output.permute(0, 2, 1, 3).reshape(batch_size, seq_len, -1) - - return self.to_out(output) - - -class LocalFacialExtractor(nn.Module): - def __init__( - self, - id_dim: int = 1280, - vit_dim: int = 1024, - depth: int = 10, - dim_head: int = 64, - heads: int = 16, - num_id_token: int = 5, - num_queries: int = 32, - output_dim: int = 2048, - ff_mult: int = 4, - num_scale: int = 5, - ): - super().__init__() - - # Storing identity token and query information - self.num_id_token = num_id_token - self.vit_dim = vit_dim - self.num_queries = num_queries - assert depth % num_scale == 0 - self.depth = depth // num_scale - self.num_scale = num_scale - scale = vit_dim**-0.5 - - # Learnable latent query embeddings - self.latents = nn.Parameter(torch.randn(1, num_queries, vit_dim) * scale) - # Projection layer to map the latent output to the desired dimension - self.proj_out = nn.Parameter(scale * torch.randn(vit_dim, output_dim)) - - # Attention and ConsisIDFeedForward layer stack - self.layers = nn.ModuleList([]) - for _ in range(depth): - self.layers.append( - nn.ModuleList( - [ - PerceiverAttention(dim=vit_dim, dim_head=dim_head, heads=heads), # Perceiver Attention layer - nn.Sequential( - nn.LayerNorm(vit_dim), - nn.Linear(vit_dim, vit_dim * ff_mult, bias=False), - nn.GELU(), - nn.Linear(vit_dim * ff_mult, vit_dim, bias=False), - ), # ConsisIDFeedForward layer - ] - ) - ) - - # Mappings for each of the 5 different ViT features - for i in range(num_scale): - setattr( - self, - f"mapping_{i}", - nn.Sequential( - nn.Linear(vit_dim, vit_dim), - nn.LayerNorm(vit_dim), - nn.LeakyReLU(), - nn.Linear(vit_dim, vit_dim), - nn.LayerNorm(vit_dim), - nn.LeakyReLU(), - nn.Linear(vit_dim, vit_dim), - ), - ) - - # Mapping for identity embedding vectors - self.id_embedding_mapping = nn.Sequential( - nn.Linear(id_dim, vit_dim), - nn.LayerNorm(vit_dim), - nn.LeakyReLU(), - nn.Linear(vit_dim, vit_dim), - nn.LayerNorm(vit_dim), - nn.LeakyReLU(), - nn.Linear(vit_dim, vit_dim * num_id_token), - ) - - 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) - - # Map the identity embedding to tokens - id_embeds = self.id_embedding_mapping(id_embeds) - id_embeds = id_embeds.reshape(-1, self.num_id_token, self.vit_dim) - - # Concatenate identity tokens with the latent queries - latents = torch.cat((latents, id_embeds), dim=1) - - # Process each of the num_scale visual feature inputs - for i in range(self.num_scale): - vit_feature = getattr(self, f"mapping_{i}")(vit_hidden_states[i]) - ctx_feature = torch.cat((id_embeds, vit_feature), dim=1) - - # Pass through the PerceiverAttention and ConsisIDFeedForward layers - for attn, ff in self.layers[i * self.depth : (i + 1) * self.depth]: - latents = attn(ctx_feature, latents) + latents - latents = ff(latents) + latents - - # Retain only the query latents - latents = latents[:, : self.num_queries] - # Project the latents to the output dimension - latents = latents @ self.proj_out - return latents - - -class PerceiverCrossAttention(nn.Module): - def __init__(self, dim: int = 3072, dim_head: int = 128, heads: int = 16, kv_dim: int = 2048): - super().__init__() - - self.scale = dim_head**-0.5 - self.dim_head = dim_head - self.heads = heads - inner_dim = dim_head * heads - - # Layer normalization to stabilize training - self.norm1 = nn.LayerNorm(dim if kv_dim is None else kv_dim) - self.norm2 = nn.LayerNorm(dim) - - # Linear transformations to produce queries, keys, and values - self.to_q = nn.Linear(dim, inner_dim, bias=False) - self.to_kv = nn.Linear(dim if kv_dim is None else kv_dim, inner_dim * 2, bias=False) - self.to_out = nn.Linear(inner_dim, dim, bias=False) - - def forward(self, image_embeds: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: - # Apply layer normalization to the input image and latent features - image_embeds = self.norm1(image_embeds) - hidden_states = self.norm2(hidden_states) - - batch_size, seq_len, _ = hidden_states.shape - - # Compute queries, keys, and values - query = self.to_q(hidden_states) - key, value = self.to_kv(image_embeds).chunk(2, dim=-1) - - # Reshape tensors to split into attention heads - query = query.reshape(query.size(0), -1, self.heads, self.dim_head).transpose(1, 2) - key = key.reshape(key.size(0), -1, self.heads, self.dim_head).transpose(1, 2) - value = value.reshape(value.size(0), -1, self.heads, self.dim_head).transpose(1, 2) - - # Compute attention weights - scale = 1 / math.sqrt(math.sqrt(self.dim_head)) - weight = (query * scale) @ (key * scale).transpose(-2, -1) # More stable scaling than post-division - weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) - - # Compute the output via weighted combination of values - out = weight @ value - - # Reshape and permute to prepare for final linear transformation - out = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, -1) - - return self.to_out(out) - - -@maybe_allow_in_graph -class ConsisIDBlock(nn.Module): - r""" - Transformer block used in [ConsisID](https://github.com/PKU-YuanGroup/ConsisID) model. - - Parameters: - 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. - time_embed_dim (`int`): - The number of channels in timestep embedding. - dropout (`float`, defaults to `0.0`): - The dropout probability to use. - activation_fn (`str`, defaults to `"gelu-approximate"`): - Activation function to be used in feed-forward. - attention_bias (`bool`, defaults to `False`): - Whether or not to use bias in attention projection layers. - qk_norm (`bool`, defaults to `True`): - Whether or not to use normalization after query and key projections in Attention. - norm_elementwise_affine (`bool`, defaults to `True`): - Whether to use learnable elementwise affine parameters for normalization. - norm_eps (`float`, defaults to `1e-5`): - Epsilon value for normalization layers. - final_dropout (`bool` defaults to `False`): - Whether to apply a final dropout after the last feed-forward layer. - ff_inner_dim (`int`, *optional*, defaults to `None`): - Custom hidden dimension of Feed-forward layer. If not provided, `4 * dim` is used. - ff_bias (`bool`, defaults to `True`): - Whether or not to use bias in Feed-forward layer. - attention_out_bias (`bool`, defaults to `True`): - Whether or not to use bias in Attention output projection layer. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - time_embed_dim: int, - dropout: float = 0.0, - activation_fn: str = "gelu-approximate", - attention_bias: bool = False, - qk_norm: bool = True, - norm_elementwise_affine: bool = True, - norm_eps: float = 1e-5, - final_dropout: bool = True, - ff_inner_dim: Optional[int] = None, - ff_bias: bool = True, - attention_out_bias: bool = True, - ): - super().__init__() - - # 1. Self Attention - self.norm1 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True) - - self.attn1 = Attention( - query_dim=dim, - dim_head=attention_head_dim, - heads=num_attention_heads, - qk_norm="layer_norm" if qk_norm else None, - eps=1e-6, - bias=attention_bias, - out_bias=attention_out_bias, - processor=CogVideoXAttnProcessor2_0(), - ) - - # 2. Feed Forward - self.norm2 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True) - - self.ff = FeedForward( - dim, - dropout=dropout, - activation_fn=activation_fn, - final_dropout=final_dropout, - inner_dim=ff_inner_dim, - bias=ff_bias, - ) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - ) -> torch.Tensor: - text_seq_length = encoder_hidden_states.size(1) - - # norm & modulate - norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1( - hidden_states, encoder_hidden_states, temb - ) - - # attention - 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, - ) - - hidden_states = hidden_states + gate_msa * attn_hidden_states - encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states - - # norm & modulate - norm_hidden_states, norm_encoder_hidden_states, gate_ff, enc_gate_ff = self.norm2( - hidden_states, encoder_hidden_states, temb - ) - - # feed-forward - norm_hidden_states = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1) - ff_output = self.ff(norm_hidden_states) - - hidden_states = hidden_states + gate_ff * ff_output[:, text_seq_length:] - encoder_hidden_states = encoder_hidden_states + enc_gate_ff * ff_output[:, :text_seq_length] - - return hidden_states, encoder_hidden_states - - -class ConsisIDTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin): - """ - A Transformer model for video-like data in [ConsisID](https://github.com/PKU-YuanGroup/ConsisID). - - Parameters: - num_attention_heads (`int`, defaults to `30`): - The number of heads to use for multi-head attention. - attention_head_dim (`int`, defaults to `64`): - The number of channels in each head. - in_channels (`int`, defaults to `16`): - The number of channels in the input. - out_channels (`int`, *optional*, defaults to `16`): - The number of channels in the output. - flip_sin_to_cos (`bool`, defaults to `True`): - Whether to flip the sin to cos in the time embedding. - time_embed_dim (`int`, defaults to `512`): - Output dimension of timestep embeddings. - text_embed_dim (`int`, defaults to `4096`): - Input dimension of text embeddings from the text encoder. - num_layers (`int`, defaults to `30`): - The number of layers of Transformer blocks to use. - dropout (`float`, defaults to `0.0`): - The dropout probability to use. - attention_bias (`bool`, defaults to `True`): - Whether to use bias in the attention projection layers. - sample_width (`int`, defaults to `90`): - The width of the input latents. - sample_height (`int`, defaults to `60`): - The height of the input latents. - sample_frames (`int`, defaults to `49`): - The number of frames in the input latents. Note that this parameter was incorrectly initialized to 49 - instead of 13 because ConsisID processed 13 latent frames at once in its default and recommended settings, - but cannot be changed to the correct value to ensure backwards compatibility. To create a transformer with - K latent frames, the correct value to pass here would be: ((K - 1) * temporal_compression_ratio + 1). - patch_size (`int`, defaults to `2`): - The size of the patches to use in the patch embedding layer. - temporal_compression_ratio (`int`, defaults to `4`): - The compression ratio across the temporal dimension. See documentation for `sample_frames`. - max_text_seq_length (`int`, defaults to `226`): - The maximum sequence length of the input text embeddings. - activation_fn (`str`, defaults to `"gelu-approximate"`): - Activation function to use in feed-forward. - timestep_activation_fn (`str`, defaults to `"silu"`): - Activation function to use when generating the timestep embeddings. - norm_elementwise_affine (`bool`, defaults to `True`): - Whether to use elementwise affine in normalization layers. - norm_eps (`float`, defaults to `1e-5`): - The epsilon value to use in normalization layers. - spatial_interpolation_scale (`float`, defaults to `1.875`): - Scaling factor to apply in 3D positional embeddings across spatial dimensions. - temporal_interpolation_scale (`float`, defaults to `1.0`): - Scaling factor to apply in 3D positional embeddings across temporal dimensions. - is_train_face (`bool`, defaults to `False`): - Whether to use enable the identity-preserving module during the training process. When set to `True`, the - model will focus on identity-preserving tasks. - is_kps (`bool`, defaults to `False`): - Whether to enable keypoint for global facial extractor. If `True`, keypoints will be in the model. - cross_attn_interval (`int`, defaults to `2`): - The interval between cross-attention layers in the Transformer architecture. A larger value may reduce the - frequency of cross-attention computations, which can help reduce computational overhead. - cross_attn_dim_head (`int`, optional, defaults to `128`): - The dimensionality of each attention head in the cross-attention layers of the Transformer architecture. A - larger value increases the capacity to attend to more complex patterns, but also increases memory and - computation costs. - cross_attn_num_heads (`int`, optional, defaults to `16`): - The number of attention heads in the cross-attention layers. More heads allow for more parallel attention - mechanisms, capturing diverse relationships between different components of the input, but can also - increase computational requirements. - LFE_id_dim (`int`, optional, defaults to `1280`): - The dimensionality of the identity vector used in the Local Facial Extractor (LFE). This vector represents - the identity features of a face, which are important for tasks like face recognition and identity - preservation across different frames. - LFE_vit_dim (`int`, optional, defaults to `1024`): - The dimension of the vision transformer (ViT) output used in the Local Facial Extractor (LFE). This value - dictates the size of the transformer-generated feature vectors that will be processed for facial feature - extraction. - LFE_depth (`int`, optional, defaults to `10`): - The number of layers in the Local Facial Extractor (LFE). Increasing the depth allows the model to capture - more complex representations of facial features, but also increases the computational load. - LFE_dim_head (`int`, optional, defaults to `64`): - The dimensionality of each attention head in the Local Facial Extractor (LFE). This parameter affects how - finely the model can process and focus on different parts of the facial features during the extraction - process. - LFE_num_heads (`int`, optional, defaults to `16`): - The number of attention heads in the Local Facial Extractor (LFE). More heads can improve the model's - ability to capture diverse facial features, but at the cost of increased computational complexity. - LFE_num_id_token (`int`, optional, defaults to `5`): - The number of identity tokens used in the Local Facial Extractor (LFE). This defines how many - identity-related tokens the model will process to ensure face identity preservation during feature - extraction. - LFE_num_querie (`int`, optional, defaults to `32`): - The number of query tokens used in the Local Facial Extractor (LFE). These tokens are used to capture - high-frequency face-related information that aids in accurate facial feature extraction. - LFE_output_dim (`int`, optional, defaults to `2048`): - The output dimension of the Local Facial Extractor (LFE). This dimension determines the size of the feature - vectors produced by the LFE module, which will be used for subsequent tasks such as face recognition or - tracking. - LFE_ff_mult (`int`, optional, defaults to `4`): - The multiplication factor applied to the feed-forward network's hidden layer size in the Local Facial - Extractor (LFE). A higher value increases the model's capacity to learn more complex facial feature - transformations, but also increases the computation and memory requirements. - LFE_num_scale (`int`, optional, defaults to `5`): - The number of different scales visual feature. A higher value increases the model's capacity to learn more - complex facial feature transformations, but also increases the computation and memory requirements. - local_face_scale (`float`, defaults to `1.0`): - A scaling factor used to adjust the importance of local facial features in the model. This can influence - how strongly the model focuses on high frequency face-related content. - """ - - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - num_attention_heads: int = 30, - attention_head_dim: int = 64, - in_channels: int = 16, - out_channels: Optional[int] = 16, - flip_sin_to_cos: bool = True, - freq_shift: int = 0, - time_embed_dim: int = 512, - text_embed_dim: int = 4096, - num_layers: int = 30, - dropout: float = 0.0, - attention_bias: bool = True, - sample_width: int = 90, - sample_height: int = 60, - sample_frames: int = 49, - patch_size: int = 2, - temporal_compression_ratio: int = 4, - max_text_seq_length: int = 226, - activation_fn: str = "gelu-approximate", - timestep_activation_fn: str = "silu", - norm_elementwise_affine: bool = True, - norm_eps: float = 1e-5, - spatial_interpolation_scale: float = 1.875, - temporal_interpolation_scale: float = 1.0, - use_rotary_positional_embeddings: bool = False, - use_learned_positional_embeddings: bool = False, - is_train_face: bool = False, - is_kps: bool = False, - cross_attn_interval: int = 2, - cross_attn_dim_head: int = 128, - cross_attn_num_heads: int = 16, - LFE_id_dim: int = 1280, - LFE_vit_dim: int = 1024, - LFE_depth: int = 10, - LFE_dim_head: int = 64, - LFE_num_heads: int = 16, - LFE_num_id_token: int = 5, - LFE_num_querie: int = 32, - LFE_output_dim: int = 2048, - LFE_ff_mult: int = 4, - LFE_num_scale: int = 5, - local_face_scale: float = 1.0, - ): - super().__init__() - inner_dim = num_attention_heads * attention_head_dim - - if not use_rotary_positional_embeddings and use_learned_positional_embeddings: - raise ValueError( - "There are no ConsisID checkpoints available with disable rotary embeddings and learned positional " - "embeddings. If you're using a custom model and/or believe this should be supported, please open an " - "issue at https://github.com/huggingface/diffusers/issues." - ) - - # 1. Patch embedding - self.patch_embed = CogVideoXPatchEmbed( - patch_size=patch_size, - in_channels=in_channels, - embed_dim=inner_dim, - text_embed_dim=text_embed_dim, - bias=True, - sample_width=sample_width, - sample_height=sample_height, - sample_frames=sample_frames, - temporal_compression_ratio=temporal_compression_ratio, - max_text_seq_length=max_text_seq_length, - spatial_interpolation_scale=spatial_interpolation_scale, - temporal_interpolation_scale=temporal_interpolation_scale, - use_positional_embeddings=not use_rotary_positional_embeddings, - use_learned_positional_embeddings=use_learned_positional_embeddings, - ) - self.embedding_dropout = nn.Dropout(dropout) - - # 2. Time embeddings - self.time_proj = Timesteps(inner_dim, flip_sin_to_cos, freq_shift) - self.time_embedding = TimestepEmbedding(inner_dim, time_embed_dim, timestep_activation_fn) - - # 3. Define spatio-temporal transformers blocks - self.transformer_blocks = nn.ModuleList( - [ - ConsisIDBlock( - dim=inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - time_embed_dim=time_embed_dim, - dropout=dropout, - activation_fn=activation_fn, - attention_bias=attention_bias, - norm_elementwise_affine=norm_elementwise_affine, - norm_eps=norm_eps, - ) - for _ in range(num_layers) - ] - ) - self.norm_final = nn.LayerNorm(inner_dim, norm_eps, norm_elementwise_affine) - - # 4. Output blocks - self.norm_out = AdaLayerNorm( - embedding_dim=time_embed_dim, - output_dim=2 * inner_dim, - norm_elementwise_affine=norm_elementwise_affine, - norm_eps=norm_eps, - chunk_dim=1, - ) - self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels) - - self.is_train_face = is_train_face - self.is_kps = is_kps - - # 5. Define identity-preserving config - if is_train_face: - # LFE configs - self.LFE_id_dim = LFE_id_dim - self.LFE_vit_dim = LFE_vit_dim - self.LFE_depth = LFE_depth - self.LFE_dim_head = LFE_dim_head - self.LFE_num_heads = LFE_num_heads - self.LFE_num_id_token = LFE_num_id_token - self.LFE_num_querie = LFE_num_querie - self.LFE_output_dim = LFE_output_dim - self.LFE_ff_mult = LFE_ff_mult - self.LFE_num_scale = LFE_num_scale - # cross configs - self.inner_dim = inner_dim - self.cross_attn_interval = cross_attn_interval - self.num_cross_attn = num_layers // cross_attn_interval - self.cross_attn_dim_head = cross_attn_dim_head - self.cross_attn_num_heads = cross_attn_num_heads - self.cross_attn_kv_dim = int(self.inner_dim / 3 * 2) - self.local_face_scale = local_face_scale - # face modules - self._init_face_inputs() - - self.gradient_checkpointing = False - - def _init_face_inputs(self): - self.local_facial_extractor = LocalFacialExtractor( - id_dim=self.LFE_id_dim, - vit_dim=self.LFE_vit_dim, - depth=self.LFE_depth, - dim_head=self.LFE_dim_head, - heads=self.LFE_num_heads, - num_id_token=self.LFE_num_id_token, - num_queries=self.LFE_num_querie, - output_dim=self.LFE_output_dim, - ff_mult=self.LFE_ff_mult, - num_scale=self.LFE_num_scale, - ) - self.perceiver_cross_attention = nn.ModuleList( - [ - PerceiverCrossAttention( - dim=self.inner_dim, - dim_head=self.cross_attn_dim_head, - heads=self.cross_attn_num_heads, - kv_dim=self.cross_attn_kv_dim, - ) - for _ in range(self.num_cross_attn) - ] - ) - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - 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, - return_dict: bool = True, - ): - 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." - ) - - # fuse clip and insightface - valid_face_emb = None - if self.is_train_face: - id_cond = id_cond.to(device=hidden_states.device, dtype=hidden_states.dtype) - id_vit_hidden = [ - tensor.to(device=hidden_states.device, dtype=hidden_states.dtype) for tensor in id_vit_hidden - ] - valid_face_emb = self.local_facial_extractor( - id_cond, id_vit_hidden - ) # torch.Size([1, 1280]), list[5](torch.Size([1, 577, 1024])) -> torch.Size([1, 32, 2048]) - - batch_size, num_frames, channels, height, width = hidden_states.shape - - # 1. Time embedding - timesteps = timestep - t_emb = self.time_proj(timesteps) - - # timesteps does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - t_emb = t_emb.to(dtype=hidden_states.dtype) - emb = self.time_embedding(t_emb, timestep_cond) - - # 2. Patch embedding - # torch.Size([1, 226, 4096]) torch.Size([1, 13, 32, 60, 90]) - hidden_states = self.patch_embed(encoder_hidden_states, hidden_states) # torch.Size([1, 17776, 3072]) - hidden_states = self.embedding_dropout(hidden_states) # torch.Size([1, 17776, 3072]) - - text_seq_length = encoder_hidden_states.shape[1] - encoder_hidden_states = hidden_states[:, :text_seq_length] # torch.Size([1, 226, 3072]) - hidden_states = hidden_states[:, text_seq_length:] # torch.Size([1, 17550, 3072]) - - # 3. Transformer blocks - ca_idx = 0 - for i, 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, - emb, - image_rotary_emb, - ) - else: - hidden_states, encoder_hidden_states = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - temb=emb, - image_rotary_emb=image_rotary_emb, - ) - - if self.is_train_face: - if i % self.cross_attn_interval == 0 and valid_face_emb is not None: - hidden_states = hidden_states + self.local_face_scale * self.perceiver_cross_attention[ca_idx]( - valid_face_emb, hidden_states - ) # torch.Size([2, 32, 2048]) torch.Size([2, 17550, 3072]) - ca_idx += 1 - - hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) - hidden_states = self.norm_final(hidden_states) - hidden_states = hidden_states[:, text_seq_length:] - - # 4. Final block - hidden_states = self.norm_out(hidden_states, temb=emb) - hidden_states = self.proj_out(hidden_states) - - # 5. Unpatchify - # Note: we use `-1` instead of `channels`: - # - It is okay to `channels` use for ConsisID (number of input channels is equal to output channels) - p = self.config.patch_size - 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/venv/lib/python3.10/site-packages/diffusers/models/transformers/dit_transformer_2d.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/dit_transformer_2d.py deleted file mode 100644 index 68f6f769436e8a40e7b62df42fd851d56062be46..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/dit_transformer_2d.py +++ /dev/null @@ -1,226 +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 Any, Dict, Optional - -import torch -import torch.nn.functional as F -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import logging -from ..attention import BasicTransformerBlock -from ..embeddings import PatchEmbed -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class DiTTransformer2DModel(ModelMixin, ConfigMixin): - r""" - A 2D Transformer model as introduced in DiT (https://huggingface.co/papers/2212.09748). - - Parameters: - num_attention_heads (int, optional, defaults to 16): The number of heads to use for multi-head attention. - attention_head_dim (int, optional, defaults to 72): The number of channels in each head. - in_channels (int, defaults to 4): The number of channels in the input. - out_channels (int, optional): - The number of channels in the output. Specify this parameter if the output channel number differs from the - input. - num_layers (int, optional, defaults to 28): The number of layers of Transformer blocks to use. - dropout (float, optional, defaults to 0.0): The dropout probability to use within the Transformer blocks. - norm_num_groups (int, optional, defaults to 32): - Number of groups for group normalization within Transformer blocks. - attention_bias (bool, optional, defaults to True): - Configure if the Transformer blocks' attention should contain a bias parameter. - sample_size (int, defaults to 32): - The width of the latent images. This parameter is fixed during training. - patch_size (int, defaults to 2): - Size of the patches the model processes, relevant for architectures working on non-sequential data. - activation_fn (str, optional, defaults to "gelu-approximate"): - Activation function to use in feed-forward networks within Transformer blocks. - num_embeds_ada_norm (int, optional, defaults to 1000): - Number of embeddings for AdaLayerNorm, fixed during training and affects the maximum denoising steps during - inference. - upcast_attention (bool, optional, defaults to False): - If true, upcasts the attention mechanism dimensions for potentially improved performance. - norm_type (str, optional, defaults to "ada_norm_zero"): - Specifies the type of normalization used, can be 'ada_norm_zero'. - norm_elementwise_affine (bool, optional, defaults to False): - If true, enables element-wise affine parameters in the normalization layers. - norm_eps (float, optional, defaults to 1e-5): - A small constant added to the denominator in normalization layers to prevent division by zero. - """ - - _skip_layerwise_casting_patterns = ["pos_embed", "norm"] - _supports_gradient_checkpointing = True - _supports_group_offloading = False - - @register_to_config - def __init__( - self, - num_attention_heads: int = 16, - attention_head_dim: int = 72, - in_channels: int = 4, - out_channels: Optional[int] = None, - num_layers: int = 28, - dropout: float = 0.0, - norm_num_groups: int = 32, - attention_bias: bool = True, - sample_size: int = 32, - patch_size: int = 2, - activation_fn: str = "gelu-approximate", - num_embeds_ada_norm: Optional[int] = 1000, - upcast_attention: bool = False, - norm_type: str = "ada_norm_zero", - norm_elementwise_affine: bool = False, - norm_eps: float = 1e-5, - ): - super().__init__() - - # Validate inputs. - if norm_type != "ada_norm_zero": - raise NotImplementedError( - f"Forward pass is not implemented when `patch_size` is not None and `norm_type` is '{norm_type}'." - ) - elif norm_type == "ada_norm_zero" and num_embeds_ada_norm is None: - raise ValueError( - f"When using a `patch_size` and this `norm_type` ({norm_type}), `num_embeds_ada_norm` cannot be None." - ) - - # Set some common variables used across the board. - self.attention_head_dim = attention_head_dim - self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim - self.out_channels = in_channels if out_channels is None else out_channels - self.gradient_checkpointing = False - - # 2. Initialize the position embedding and transformer blocks. - self.height = self.config.sample_size - self.width = self.config.sample_size - - self.patch_size = self.config.patch_size - self.pos_embed = PatchEmbed( - height=self.config.sample_size, - width=self.config.sample_size, - patch_size=self.config.patch_size, - in_channels=self.config.in_channels, - embed_dim=self.inner_dim, - ) - - self.transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - self.inner_dim, - self.config.num_attention_heads, - self.config.attention_head_dim, - dropout=self.config.dropout, - activation_fn=self.config.activation_fn, - num_embeds_ada_norm=self.config.num_embeds_ada_norm, - attention_bias=self.config.attention_bias, - upcast_attention=self.config.upcast_attention, - norm_type=norm_type, - norm_elementwise_affine=self.config.norm_elementwise_affine, - norm_eps=self.config.norm_eps, - ) - for _ in range(self.config.num_layers) - ] - ) - - # 3. Output blocks. - self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6) - self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim) - self.proj_out_2 = nn.Linear( - self.inner_dim, self.config.patch_size * self.config.patch_size * self.out_channels - ) - - 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, - return_dict: bool = True, - ): - """ - The [`DiTTransformer2DModel`] forward method. - - Args: - hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous): - Input `hidden_states`. - timestep ( `torch.LongTensor`, *optional*): - Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`. - 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*): - 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.unets.unet_2d_condition.UNet2DConditionOutput`] 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. - """ - # 1. Input - height, width = hidden_states.shape[-2] // self.patch_size, hidden_states.shape[-1] // self.patch_size - hidden_states = self.pos_embed(hidden_states) - - # 2. 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, - None, - None, - None, - timestep, - cross_attention_kwargs, - class_labels, - ) - else: - hidden_states = block( - hidden_states, - attention_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - timestep=timestep, - cross_attention_kwargs=cross_attention_kwargs, - class_labels=class_labels, - ) - - # 3. Output - conditioning = self.transformer_blocks[0].norm1.emb(timestep, class_labels, hidden_dtype=hidden_states.dtype) - shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) - hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] - hidden_states = self.proj_out_2(hidden_states) - - # unpatchify - height = width = int(hidden_states.shape[1] ** 0.5) - hidden_states = hidden_states.reshape( - shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels) - ) - hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) - output = hidden_states.reshape( - shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size) - ) - - if not return_dict: - return (output,) - - return Transformer2DModelOutput(sample=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/dual_transformer_2d.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/dual_transformer_2d.py deleted file mode 100644 index 24eed2168229c896e1c7012d4f8fc045dc64a7c8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/dual_transformer_2d.py +++ /dev/null @@ -1,156 +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 - -from torch import nn - -from ..modeling_outputs import Transformer2DModelOutput -from .transformer_2d import Transformer2DModel - - -class DualTransformer2DModel(nn.Module): - """ - Dual transformer wrapper that combines two `Transformer2DModel`s for mixed inference. - - Parameters: - num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. - attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. - in_channels (`int`, *optional*): - Pass if the input is continuous. The number of channels in the input and output. - num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. - dropout (`float`, *optional*, defaults to 0.1): The dropout probability to use. - cross_attention_dim (`int`, *optional*): The number of encoder_hidden_states dimensions to use. - sample_size (`int`, *optional*): Pass if the input is discrete. The width of the latent images. - Note that this is fixed at training time as it is used for learning a number of position embeddings. See - `ImagePositionalEmbeddings`. - num_vector_embeds (`int`, *optional*): - Pass if the input is discrete. The number of classes of the vector embeddings of the latent pixels. - Includes the class for the masked latent pixel. - activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. - num_embeds_ada_norm ( `int`, *optional*): Pass if at least one of the norm_layers is `AdaLayerNorm`. - The number of diffusion steps used during training. Note that this is fixed at training time as it is used - to learn a number of embeddings that are added to the hidden states. During inference, you can denoise for - up to but not more than steps than `num_embeds_ada_norm`. - attention_bias (`bool`, *optional*): - Configure if the TransformerBlocks' attention should contain a bias parameter. - """ - - def __init__( - self, - num_attention_heads: int = 16, - attention_head_dim: int = 88, - in_channels: Optional[int] = None, - num_layers: int = 1, - dropout: float = 0.0, - norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = None, - attention_bias: bool = False, - sample_size: Optional[int] = None, - num_vector_embeds: Optional[int] = None, - activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, - ): - super().__init__() - self.transformers = nn.ModuleList( - [ - Transformer2DModel( - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - in_channels=in_channels, - num_layers=num_layers, - dropout=dropout, - norm_num_groups=norm_num_groups, - cross_attention_dim=cross_attention_dim, - attention_bias=attention_bias, - sample_size=sample_size, - num_vector_embeds=num_vector_embeds, - activation_fn=activation_fn, - num_embeds_ada_norm=num_embeds_ada_norm, - ) - for _ in range(2) - ] - ) - - # Variables that can be set by a pipeline: - - # The ratio of transformer1 to transformer2's output states to be combined during inference - self.mix_ratio = 0.5 - - # The shape of `encoder_hidden_states` is expected to be - # `(batch_size, condition_lengths[0]+condition_lengths[1], num_features)` - self.condition_lengths = [77, 257] - - # Which transformer to use to encode which condition. - # E.g. `(1, 0)` means that we'll use `transformers[1](conditions[0])` and `transformers[0](conditions[1])` - self.transformer_index_for_condition = [1, 0] - - def forward( - self, - hidden_states, - encoder_hidden_states, - timestep=None, - attention_mask=None, - cross_attention_kwargs=None, - return_dict: bool = True, - ): - """ - Args: - hidden_states ( When discrete, `torch.LongTensor` of shape `(batch size, num latent pixels)`. - When continuous, `torch.Tensor` of shape `(batch size, channel, height, width)`): Input hidden_states. - encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*): - Conditional embeddings for cross attention layer. If not given, cross-attention defaults to - self-attention. - timestep ( `torch.long`, *optional*): - Optional timestep to be applied as an embedding in AdaLayerNorm's. Used to indicate denoising step. - attention_mask (`torch.Tensor`, *optional*): - Optional attention mask to be applied in Attention. - cross_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.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain - tuple. - - Returns: - [`~models.transformers.transformer_2d.Transformer2DModelOutput`] or `tuple`: - [`~models.transformers.transformer_2d.Transformer2DModelOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - """ - input_states = hidden_states - - encoded_states = [] - tokens_start = 0 - # attention_mask is not used yet - for i in range(2): - # for each of the two transformers, pass the corresponding condition tokens - condition_state = encoder_hidden_states[:, tokens_start : tokens_start + self.condition_lengths[i]] - transformer_index = self.transformer_index_for_condition[i] - encoded_state = self.transformers[transformer_index]( - input_states, - encoder_hidden_states=condition_state, - timestep=timestep, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - encoded_states.append(encoded_state - input_states) - tokens_start += self.condition_lengths[i] - - output_states = encoded_states[0] * self.mix_ratio + encoded_states[1] * (1 - self.mix_ratio) - output_states = output_states + input_states - - if not return_dict: - return (output_states,) - - return Transformer2DModelOutput(sample=output_states) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/hunyuan_transformer_2d.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/hunyuan_transformer_2d.py deleted file mode 100644 index f634718788572c663be3ca2ed3d627e71a6c6a66..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/hunyuan_transformer_2d.py +++ /dev/null @@ -1,579 +0,0 @@ -# Copyright 2025 HunyuanDiT Authors, Qixun Wang 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 Dict, Optional, Union - -import torch -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import logging -from ...utils.torch_utils import maybe_allow_in_graph -from ..attention import FeedForward -from ..attention_processor import Attention, AttentionProcessor, FusedHunyuanAttnProcessor2_0, HunyuanAttnProcessor2_0 -from ..embeddings import ( - HunyuanCombinedTimestepTextSizeStyleEmbedding, - PatchEmbed, - PixArtAlphaTextProjection, -) -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormContinuous, FP32LayerNorm - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class AdaLayerNormShift(nn.Module): - r""" - Norm layer modified to incorporate timestep embeddings. - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - num_embeddings (`int`): The size of the embeddings dictionary. - """ - - def __init__(self, embedding_dim: int, elementwise_affine=True, eps=1e-6): - super().__init__() - self.silu = nn.SiLU() - self.linear = nn.Linear(embedding_dim, embedding_dim) - self.norm = FP32LayerNorm(embedding_dim, elementwise_affine=elementwise_affine, eps=eps) - - def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: - shift = self.linear(self.silu(emb.to(torch.float32)).to(emb.dtype)) - x = self.norm(x) + shift.unsqueeze(dim=1) - return x - - -@maybe_allow_in_graph -class HunyuanDiTBlock(nn.Module): - r""" - Transformer block used in Hunyuan-DiT model (https://github.com/Tencent/HunyuanDiT). Allow skip connection and - QKNorm - - Parameters: - dim (`int`): - The number of channels in the input and output. - num_attention_heads (`int`): - The number of headsto use for multi-head attention. - cross_attention_dim (`int`,*optional*): - The size of the encoder_hidden_states vector for cross attention. - dropout(`float`, *optional*, defaults to 0.0): - The dropout probability to use. - activation_fn (`str`,*optional*, defaults to `"geglu"`): - Activation function to be used in feed-forward. . - norm_elementwise_affine (`bool`, *optional*, defaults to `True`): - Whether to use learnable elementwise affine parameters for normalization. - norm_eps (`float`, *optional*, defaults to 1e-6): - A small constant added to the denominator in normalization layers to prevent division by zero. - final_dropout (`bool` *optional*, defaults to False): - Whether to apply a final dropout after the last feed-forward layer. - ff_inner_dim (`int`, *optional*): - The size of the hidden layer in the feed-forward block. Defaults to `None`. - ff_bias (`bool`, *optional*, defaults to `True`): - Whether to use bias in the feed-forward block. - skip (`bool`, *optional*, defaults to `False`): - Whether to use skip connection. Defaults to `False` for down-blocks and mid-blocks. - qk_norm (`bool`, *optional*, defaults to `True`): - Whether to use normalization in QK calculation. Defaults to `True`. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - cross_attention_dim: int = 1024, - dropout=0.0, - activation_fn: str = "geglu", - norm_elementwise_affine: bool = True, - norm_eps: float = 1e-6, - final_dropout: bool = False, - ff_inner_dim: Optional[int] = None, - ff_bias: bool = True, - skip: bool = False, - qk_norm: bool = True, - ): - super().__init__() - - # Define 3 blocks. Each block has its own normalization layer. - # NOTE: when new version comes, check norm2 and norm 3 - # 1. Self-Attn - self.norm1 = AdaLayerNormShift(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) - - self.attn1 = Attention( - query_dim=dim, - cross_attention_dim=None, - dim_head=dim // num_attention_heads, - heads=num_attention_heads, - qk_norm="layer_norm" if qk_norm else None, - eps=1e-6, - bias=True, - processor=HunyuanAttnProcessor2_0(), - ) - - # 2. Cross-Attn - self.norm2 = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) - - self.attn2 = Attention( - query_dim=dim, - cross_attention_dim=cross_attention_dim, - dim_head=dim // num_attention_heads, - heads=num_attention_heads, - qk_norm="layer_norm" if qk_norm else None, - eps=1e-6, - bias=True, - processor=HunyuanAttnProcessor2_0(), - ) - # 3. Feed-forward - self.norm3 = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) - - self.ff = FeedForward( - dim, - dropout=dropout, ### 0.0 - activation_fn=activation_fn, ### approx GeLU - final_dropout=final_dropout, ### 0.0 - inner_dim=ff_inner_dim, ### int(dim * mlp_ratio) - bias=ff_bias, - ) - - # 4. Skip Connection - if skip: - self.skip_norm = FP32LayerNorm(2 * dim, norm_eps, elementwise_affine=True) - self.skip_linear = nn.Linear(2 * dim, dim) - else: - self.skip_linear = None - - # let chunk size default to None - self._chunk_size = None - 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): - # Sets chunk feed-forward - self._chunk_size = chunk_size - self._chunk_dim = dim - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, - image_rotary_emb=None, - skip=None, - ) -> torch.Tensor: - # Notice that normalization is always applied before the real computation in the following blocks. - # 0. Long Skip Connection - if self.skip_linear is not None: - cat = torch.cat([hidden_states, skip], dim=-1) - cat = self.skip_norm(cat) - hidden_states = self.skip_linear(cat) - - # 1. Self-Attention - norm_hidden_states = self.norm1(hidden_states, temb) ### checked: self.norm1 is correct - attn_output = self.attn1( - norm_hidden_states, - image_rotary_emb=image_rotary_emb, - ) - hidden_states = hidden_states + attn_output - - # 2. Cross-Attention - hidden_states = hidden_states + self.attn2( - self.norm2(hidden_states), - encoder_hidden_states=encoder_hidden_states, - image_rotary_emb=image_rotary_emb, - ) - - # FFN Layer ### TODO: switch norm2 and norm3 in the state dict - mlp_inputs = self.norm3(hidden_states) - hidden_states = hidden_states + self.ff(mlp_inputs) - - return hidden_states - - -class HunyuanDiT2DModel(ModelMixin, ConfigMixin): - """ - HunYuanDiT: Diffusion model with a Transformer backbone. - - Inherit ModelMixin and ConfigMixin to be compatible with the sampler StableDiffusionPipeline of diffusers. - - Parameters: - num_attention_heads (`int`, *optional*, defaults to 16): - The number of heads to use for multi-head attention. - attention_head_dim (`int`, *optional*, defaults to 88): - The number of channels in each head. - in_channels (`int`, *optional*): - The number of channels in the input and output (specify if the input is **continuous**). - patch_size (`int`, *optional*): - The size of the patch to use for the input. - activation_fn (`str`, *optional*, defaults to `"geglu"`): - Activation function to use in feed-forward. - sample_size (`int`, *optional*): - The width of the latent images. This is fixed during training since it is used to learn a number of - position embeddings. - dropout (`float`, *optional*, defaults to 0.0): - The dropout probability to use. - cross_attention_dim (`int`, *optional*): - The number of dimension in the clip text embedding. - hidden_size (`int`, *optional*): - The size of hidden layer in the conditioning embedding layers. - num_layers (`int`, *optional*, defaults to 1): - The number of layers of Transformer blocks to use. - mlp_ratio (`float`, *optional*, defaults to 4.0): - The ratio of the hidden layer size to the input size. - learn_sigma (`bool`, *optional*, defaults to `True`): - Whether to predict variance. - cross_attention_dim_t5 (`int`, *optional*): - The number dimensions in t5 text embedding. - pooled_projection_dim (`int`, *optional*): - The size of the pooled projection. - text_len (`int`, *optional*): - The length of the clip text embedding. - text_len_t5 (`int`, *optional*): - The length of the T5 text embedding. - use_style_cond_and_image_meta_size (`bool`, *optional*): - Whether or not to use style condition and image meta size. True for version <=1.1, False for version >= 1.2 - """ - - _skip_layerwise_casting_patterns = ["pos_embed", "norm", "pooler"] - _supports_group_offloading = False - - @register_to_config - def __init__( - self, - num_attention_heads: int = 16, - attention_head_dim: int = 88, - in_channels: Optional[int] = None, - patch_size: Optional[int] = None, - activation_fn: str = "gelu-approximate", - sample_size=32, - hidden_size=1152, - num_layers: int = 28, - mlp_ratio: float = 4.0, - learn_sigma: bool = True, - cross_attention_dim: int = 1024, - norm_type: str = "layer_norm", - cross_attention_dim_t5: int = 2048, - pooled_projection_dim: int = 1024, - text_len: int = 77, - text_len_t5: int = 256, - use_style_cond_and_image_meta_size: bool = True, - ): - super().__init__() - self.out_channels = in_channels * 2 if learn_sigma else in_channels - self.num_heads = num_attention_heads - self.inner_dim = num_attention_heads * attention_head_dim - - self.text_embedder = PixArtAlphaTextProjection( - in_features=cross_attention_dim_t5, - hidden_size=cross_attention_dim_t5 * 4, - out_features=cross_attention_dim, - act_fn="silu_fp32", - ) - - self.text_embedding_padding = nn.Parameter(torch.randn(text_len + text_len_t5, cross_attention_dim)) - - self.pos_embed = PatchEmbed( - height=sample_size, - width=sample_size, - in_channels=in_channels, - embed_dim=hidden_size, - patch_size=patch_size, - pos_embed_type=None, - ) - - self.time_extra_emb = HunyuanCombinedTimestepTextSizeStyleEmbedding( - hidden_size, - pooled_projection_dim=pooled_projection_dim, - seq_len=text_len_t5, - cross_attention_dim=cross_attention_dim_t5, - use_style_cond_and_image_meta_size=use_style_cond_and_image_meta_size, - ) - - # HunyuanDiT Blocks - self.blocks = nn.ModuleList( - [ - HunyuanDiTBlock( - dim=self.inner_dim, - num_attention_heads=self.config.num_attention_heads, - activation_fn=activation_fn, - ff_inner_dim=int(self.inner_dim * mlp_ratio), - cross_attention_dim=cross_attention_dim, - qk_norm=True, # See https://huggingface.co/papers/2302.05442 for details. - skip=layer > num_layers // 2, - ) - for layer in range(num_layers) - ] - ) - - self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6) - self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedHunyuanAttnProcessor2_0 - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - self.set_attn_processor(FusedHunyuanAttnProcessor2_0()) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - if self.original_attn_processors is not None: - self.set_attn_processor(self.original_attn_processors) - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - self.set_attn_processor(HunyuanAttnProcessor2_0()) - - def forward( - self, - hidden_states, - timestep, - encoder_hidden_states=None, - text_embedding_mask=None, - encoder_hidden_states_t5=None, - text_embedding_mask_t5=None, - image_meta_size=None, - style=None, - image_rotary_emb=None, - controlnet_block_samples=None, - return_dict=True, - ): - """ - The [`HunyuanDiT2DModel`] forward method. - - Args: - hidden_states (`torch.Tensor` of shape `(batch size, dim, height, width)`): - The input tensor. - timestep ( `torch.LongTensor`, *optional*): - Used to indicate denoising step. - encoder_hidden_states ( `torch.Tensor` of shape `(batch size, sequence len, embed dims)`, *optional*): - Conditional embeddings for cross attention layer. This is the output of `BertModel`. - text_embedding_mask: torch.Tensor - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. This is the output - of `BertModel`. - encoder_hidden_states_t5 ( `torch.Tensor` of shape `(batch size, sequence len, embed dims)`, *optional*): - Conditional embeddings for cross attention layer. This is the output of T5 Text Encoder. - text_embedding_mask_t5: torch.Tensor - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. This is the output - of T5 Text Encoder. - image_meta_size (torch.Tensor): - Conditional embedding indicate the image sizes - style: torch.Tensor: - Conditional embedding indicate the style - image_rotary_emb (`torch.Tensor`): - The image rotary embeddings to apply on query and key tensors during attention calculation. - return_dict: bool - Whether to return a dictionary. - """ - - height, width = hidden_states.shape[-2:] - - hidden_states = self.pos_embed(hidden_states) - - temb = self.time_extra_emb( - timestep, encoder_hidden_states_t5, image_meta_size, style, hidden_dtype=timestep.dtype - ) # [B, D] - - # text projection - batch_size, sequence_length, _ = encoder_hidden_states_t5.shape - encoder_hidden_states_t5 = self.text_embedder( - encoder_hidden_states_t5.view(-1, encoder_hidden_states_t5.shape[-1]) - ) - encoder_hidden_states_t5 = encoder_hidden_states_t5.view(batch_size, sequence_length, -1) - - encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states_t5], dim=1) - text_embedding_mask = torch.cat([text_embedding_mask, text_embedding_mask_t5], dim=-1) - text_embedding_mask = text_embedding_mask.unsqueeze(2).bool() - - encoder_hidden_states = torch.where(text_embedding_mask, encoder_hidden_states, self.text_embedding_padding) - - skips = [] - for layer, block in enumerate(self.blocks): - if layer > self.config.num_layers // 2: - if controlnet_block_samples is not None: - skip = skips.pop() + controlnet_block_samples.pop() - else: - skip = skips.pop() - hidden_states = block( - hidden_states, - temb=temb, - encoder_hidden_states=encoder_hidden_states, - image_rotary_emb=image_rotary_emb, - skip=skip, - ) # (N, L, D) - else: - hidden_states = block( - hidden_states, - temb=temb, - encoder_hidden_states=encoder_hidden_states, - image_rotary_emb=image_rotary_emb, - ) # (N, L, D) - - if layer < (self.config.num_layers // 2 - 1): - skips.append(hidden_states) - - if controlnet_block_samples is not None and len(controlnet_block_samples) != 0: - raise ValueError("The number of controls is not equal to the number of skip connections.") - - # final layer - hidden_states = self.norm_out(hidden_states, temb.to(torch.float32)) - hidden_states = self.proj_out(hidden_states) - # (N, L, patch_size ** 2 * out_channels) - - # unpatchify: (N, out_channels, H, W) - patch_size = self.pos_embed.patch_size - height = height // patch_size - width = width // patch_size - - hidden_states = hidden_states.reshape( - shape=(hidden_states.shape[0], height, width, patch_size, patch_size, self.out_channels) - ) - hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) - output = hidden_states.reshape( - shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size) - ) - if not return_dict: - return (output,) - 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: - """ - Sets the attention processor to use [feed forward - chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). - - Parameters: - chunk_size (`int`, *optional*): - The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually - over each tensor of dim=`dim`. - dim (`int`, *optional*, defaults to `0`): - The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch) - or dim=1 (sequence length). - """ - if dim not in [0, 1]: - raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}") - - # By default chunk size is 1 - chunk_size = chunk_size or 1 - - def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int): - if hasattr(module, "set_chunk_feed_forward"): - module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) - - for child in module.children(): - fn_recursive_feed_forward(child, chunk_size, dim) - - for module in self.children(): - fn_recursive_feed_forward(module, chunk_size, dim) - - # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.disable_forward_chunking - def disable_forward_chunking(self): - def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int): - if hasattr(module, "set_chunk_feed_forward"): - module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) - - for child in module.children(): - fn_recursive_feed_forward(child, chunk_size, dim) - - for module in self.children(): - fn_recursive_feed_forward(module, None, 0) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/latte_transformer_3d.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/latte_transformer_3d.py deleted file mode 100644 index 990c90512e39753186de6b16a77b7265010375b4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/latte_transformer_3d.py +++ /dev/null @@ -1,331 +0,0 @@ -# Copyright 2025 the Latte 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 Optional - -import torch -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ..attention import BasicTransformerBlock -from ..cache_utils import CacheMixin -from ..embeddings import PatchEmbed, PixArtAlphaTextProjection, get_1d_sincos_pos_embed_from_grid -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormSingle - - -class LatteTransformer3DModel(ModelMixin, ConfigMixin, CacheMixin): - _supports_gradient_checkpointing = True - - """ - A 3D Transformer model for video-like data, paper: https://huggingface.co/papers/2401.03048, official code: - https://github.com/Vchitect/Latte - - Parameters: - num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. - attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. - in_channels (`int`, *optional*): - The number of channels in the input. - out_channels (`int`, *optional*): - The number of channels in the output. - num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use. - attention_bias (`bool`, *optional*): - Configure if the `TransformerBlocks` attention should contain a bias parameter. - sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**). - This is fixed during training since it is used to learn a number of position embeddings. - patch_size (`int`, *optional*): - The size of the patches to use in the patch embedding layer. - activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward. - num_embeds_ada_norm ( `int`, *optional*): - The number of diffusion steps used during training. Pass if at least one of the norm_layers is - `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are - added to the hidden states. During inference, you can denoise for up to but not more steps than - `num_embeds_ada_norm`. - norm_type (`str`, *optional*, defaults to `"layer_norm"`): - The type of normalization to use. Options are `"layer_norm"` or `"ada_layer_norm"`. - norm_elementwise_affine (`bool`, *optional*, defaults to `True`): - Whether or not to use elementwise affine in normalization layers. - norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon value to use in normalization layers. - caption_channels (`int`, *optional*): - The number of channels in the caption embeddings. - video_length (`int`, *optional*): - The number of frames in the video-like data. - """ - - _skip_layerwise_casting_patterns = ["pos_embed", "norm"] - - @register_to_config - def __init__( - self, - num_attention_heads: int = 16, - attention_head_dim: int = 88, - in_channels: Optional[int] = None, - out_channels: Optional[int] = None, - num_layers: int = 1, - dropout: float = 0.0, - cross_attention_dim: Optional[int] = None, - attention_bias: bool = False, - sample_size: int = 64, - patch_size: Optional[int] = None, - activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, - norm_type: str = "layer_norm", - norm_elementwise_affine: bool = True, - norm_eps: float = 1e-5, - caption_channels: int = None, - video_length: int = 16, - ): - super().__init__() - inner_dim = num_attention_heads * attention_head_dim - - # 1. Define input layers - self.height = sample_size - self.width = sample_size - - interpolation_scale = self.config.sample_size // 64 - interpolation_scale = max(interpolation_scale, 1) - self.pos_embed = PatchEmbed( - height=sample_size, - width=sample_size, - patch_size=patch_size, - in_channels=in_channels, - embed_dim=inner_dim, - interpolation_scale=interpolation_scale, - ) - - # 2. Define spatial transformers blocks - self.transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - inner_dim, - num_attention_heads, - attention_head_dim, - dropout=dropout, - cross_attention_dim=cross_attention_dim, - activation_fn=activation_fn, - num_embeds_ada_norm=num_embeds_ada_norm, - attention_bias=attention_bias, - norm_type=norm_type, - norm_elementwise_affine=norm_elementwise_affine, - norm_eps=norm_eps, - ) - for d in range(num_layers) - ] - ) - - # 3. Define temporal transformers blocks - self.temporal_transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - inner_dim, - num_attention_heads, - attention_head_dim, - dropout=dropout, - cross_attention_dim=None, - activation_fn=activation_fn, - num_embeds_ada_norm=num_embeds_ada_norm, - attention_bias=attention_bias, - norm_type=norm_type, - norm_elementwise_affine=norm_elementwise_affine, - norm_eps=norm_eps, - ) - for d in range(num_layers) - ] - ) - - # 4. Define output layers - self.out_channels = in_channels if out_channels is None else out_channels - self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6) - self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5) - self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels) - - # 5. Latte other blocks. - self.adaln_single = AdaLayerNormSingle(inner_dim, use_additional_conditions=False) - self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) - - # define temporal positional embedding - temp_pos_embed = get_1d_sincos_pos_embed_from_grid( - inner_dim, torch.arange(0, video_length).unsqueeze(1), output_type="pt" - ) # 1152 hidden size - self.register_buffer("temp_pos_embed", temp_pos_embed.float().unsqueeze(0), persistent=False) - - self.gradient_checkpointing = False - - 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, - enable_temporal_attentions: bool = True, - return_dict: bool = True, - ): - """ - The [`LatteTransformer3DModel`] forward method. - - Args: - hidden_states shape `(batch size, channel, num_frame, height, width)`: - Input `hidden_states`. - timestep ( `torch.LongTensor`, *optional*): - Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`. - encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*): - Conditional embeddings for cross attention layer. If not given, cross-attention defaults to - self-attention. - encoder_attention_mask ( `torch.Tensor`, *optional*): - Cross-attention mask applied to `encoder_hidden_states`. Two formats supported: - - * Mask `(batcheight, sequence_length)` True = keep, False = discard. - * Bias `(batcheight, 1, sequence_length)` 0 = keep, -10000 = discard. - - If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format - above. This bias will be added to the cross-attention scores. - enable_temporal_attentions: - (`bool`, *optional*, defaults to `True`): Whether to enable temporal attentions. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] 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. - """ - - # Reshape hidden states - batch_size, channels, num_frame, height, width = hidden_states.shape - # batch_size channels num_frame height width -> (batch_size * num_frame) channels height width - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).reshape(-1, channels, height, width) - - # Input - height, width = ( - hidden_states.shape[-2] // self.config.patch_size, - hidden_states.shape[-1] // self.config.patch_size, - ) - num_patches = height * width - - hidden_states = self.pos_embed(hidden_states) # already add positional embeddings - - added_cond_kwargs = {"resolution": None, "aspect_ratio": None} - timestep, embedded_timestep = self.adaln_single( - timestep, added_cond_kwargs=added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_states.dtype - ) - - # Prepare text embeddings for spatial block - # batch_size num_tokens hidden_size -> (batch_size * num_frame) num_tokens hidden_size - encoder_hidden_states = self.caption_projection(encoder_hidden_states) # 3 120 1152 - encoder_hidden_states_spatial = encoder_hidden_states.repeat_interleave( - num_frame, dim=0, output_size=encoder_hidden_states.shape[0] * num_frame - ).view(-1, encoder_hidden_states.shape[-2], encoder_hidden_states.shape[-1]) - - # Prepare timesteps for spatial and temporal block - timestep_spatial = timestep.repeat_interleave( - num_frame, dim=0, output_size=timestep.shape[0] * num_frame - ).view(-1, timestep.shape[-1]) - timestep_temp = timestep.repeat_interleave( - num_patches, dim=0, output_size=timestep.shape[0] * num_patches - ).view(-1, timestep.shape[-1]) - - # Spatial and temporal transformer blocks - for i, (spatial_block, temp_block) in enumerate( - zip(self.transformer_blocks, self.temporal_transformer_blocks) - ): - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - spatial_block, - hidden_states, - None, # attention_mask - encoder_hidden_states_spatial, - encoder_attention_mask, - timestep_spatial, - None, # cross_attention_kwargs - None, # class_labels - ) - else: - hidden_states = spatial_block( - hidden_states, - None, # attention_mask - encoder_hidden_states_spatial, - encoder_attention_mask, - timestep_spatial, - None, # cross_attention_kwargs - None, # class_labels - ) - - if enable_temporal_attentions: - # (batch_size * num_frame) num_tokens hidden_size -> (batch_size * num_tokens) num_frame hidden_size - hidden_states = hidden_states.reshape( - batch_size, -1, hidden_states.shape[-2], hidden_states.shape[-1] - ).permute(0, 2, 1, 3) - hidden_states = hidden_states.reshape(-1, hidden_states.shape[-2], hidden_states.shape[-1]) - - if i == 0 and num_frame > 1: - hidden_states = hidden_states + self.temp_pos_embed.to(hidden_states.dtype) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - temp_block, - hidden_states, - None, # attention_mask - None, # encoder_hidden_states - None, # encoder_attention_mask - timestep_temp, - None, # cross_attention_kwargs - None, # class_labels - ) - else: - hidden_states = temp_block( - hidden_states, - None, # attention_mask - None, # encoder_hidden_states - None, # encoder_attention_mask - timestep_temp, - None, # cross_attention_kwargs - None, # class_labels - ) - - # (batch_size * num_tokens) num_frame hidden_size -> (batch_size * num_frame) num_tokens hidden_size - hidden_states = hidden_states.reshape( - batch_size, -1, hidden_states.shape[-2], hidden_states.shape[-1] - ).permute(0, 2, 1, 3) - hidden_states = hidden_states.reshape(-1, hidden_states.shape[-2], hidden_states.shape[-1]) - - embedded_timestep = embedded_timestep.repeat_interleave( - num_frame, dim=0, output_size=embedded_timestep.shape[0] * num_frame - ).view(-1, embedded_timestep.shape[-1]) - shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1) - hidden_states = self.norm_out(hidden_states) - # Modulation - hidden_states = hidden_states * (1 + scale) + shift - hidden_states = self.proj_out(hidden_states) - - # unpatchify - if self.adaln_single is None: - height = width = int(hidden_states.shape[1] ** 0.5) - hidden_states = hidden_states.reshape( - shape=(-1, height, width, self.config.patch_size, self.config.patch_size, self.out_channels) - ) - hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) - output = hidden_states.reshape( - shape=(-1, self.out_channels, height * self.config.patch_size, width * self.config.patch_size) - ) - output = output.reshape(batch_size, -1, output.shape[-3], output.shape[-2], output.shape[-1]).permute( - 0, 2, 1, 3, 4 - ) - - if not return_dict: - return (output,) - - return Transformer2DModelOutput(sample=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/lumina_nextdit2d.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/lumina_nextdit2d.py deleted file mode 100644 index 84b1175386b0a72faa80d5a690dc9963317ae198..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/lumina_nextdit2d.py +++ /dev/null @@ -1,342 +0,0 @@ -# Copyright 2025 Alpha-VLLM Authors 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, Dict, Optional - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import logging -from ..attention import LuminaFeedForward -from ..attention_processor import Attention, LuminaAttnProcessor2_0 -from ..embeddings import ( - LuminaCombinedTimestepCaptionEmbedding, - LuminaPatchEmbed, -) -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import LuminaLayerNormContinuous, LuminaRMSNormZero, RMSNorm - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class LuminaNextDiTBlock(nn.Module): - """ - A LuminaNextDiTBlock for LuminaNextDiT2DModel. - - Parameters: - dim (`int`): Embedding dimension of the input features. - num_attention_heads (`int`): Number of attention heads. - num_kv_heads (`int`): - Number of attention heads in key and value features (if using GQA), or set to None for the same as query. - multiple_of (`int`): The number of multiple of ffn layer. - ffn_dim_multiplier (`float`): The multiplier factor of ffn layer dimension. - norm_eps (`float`): The eps for norm layer. - qk_norm (`bool`): normalization for query and key. - cross_attention_dim (`int`): Cross attention embedding dimension of the input text prompt hidden_states. - norm_elementwise_affine (`bool`, *optional*, defaults to True), - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - num_kv_heads: int, - multiple_of: int, - ffn_dim_multiplier: float, - norm_eps: float, - qk_norm: bool, - cross_attention_dim: int, - norm_elementwise_affine: bool = True, - ) -> None: - super().__init__() - self.head_dim = dim // num_attention_heads - - self.gate = nn.Parameter(torch.zeros([num_attention_heads])) - - # Self-attention - self.attn1 = Attention( - query_dim=dim, - cross_attention_dim=None, - dim_head=dim // num_attention_heads, - qk_norm="layer_norm_across_heads" if qk_norm else None, - heads=num_attention_heads, - kv_heads=num_kv_heads, - eps=1e-5, - bias=False, - out_bias=False, - processor=LuminaAttnProcessor2_0(), - ) - self.attn1.to_out = nn.Identity() - - # Cross-attention - self.attn2 = Attention( - query_dim=dim, - cross_attention_dim=cross_attention_dim, - dim_head=dim // num_attention_heads, - qk_norm="layer_norm_across_heads" if qk_norm else None, - heads=num_attention_heads, - kv_heads=num_kv_heads, - eps=1e-5, - bias=False, - out_bias=False, - processor=LuminaAttnProcessor2_0(), - ) - - self.feed_forward = LuminaFeedForward( - dim=dim, - inner_dim=int(4 * 2 * dim / 3), - multiple_of=multiple_of, - ffn_dim_multiplier=ffn_dim_multiplier, - ) - - self.norm1 = LuminaRMSNormZero( - embedding_dim=dim, - norm_eps=norm_eps, - norm_elementwise_affine=norm_elementwise_affine, - ) - self.ffn_norm1 = RMSNorm(dim, eps=norm_eps, elementwise_affine=norm_elementwise_affine) - - self.norm2 = RMSNorm(dim, eps=norm_eps, elementwise_affine=norm_elementwise_affine) - self.ffn_norm2 = RMSNorm(dim, eps=norm_eps, elementwise_affine=norm_elementwise_affine) - - self.norm1_context = RMSNorm(cross_attention_dim, eps=norm_eps, elementwise_affine=norm_elementwise_affine) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor, - image_rotary_emb: torch.Tensor, - encoder_hidden_states: torch.Tensor, - encoder_mask: torch.Tensor, - temb: torch.Tensor, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - ): - """ - Perform a forward pass through the LuminaNextDiTBlock. - - Parameters: - hidden_states (`torch.Tensor`): The input of hidden_states for LuminaNextDiTBlock. - attention_mask (`torch.Tensor): The input of hidden_states corresponse attention mask. - image_rotary_emb (`torch.Tensor`): Precomputed cosine and sine frequencies. - 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. - """ - residual = hidden_states - - # Self-attention - norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb) - self_attn_output = self.attn1( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_hidden_states, - attention_mask=attention_mask, - query_rotary_emb=image_rotary_emb, - key_rotary_emb=image_rotary_emb, - **cross_attention_kwargs, - ) - - # Cross-attention - norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states) - cross_attn_output = self.attn2( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - attention_mask=encoder_mask, - query_rotary_emb=image_rotary_emb, - key_rotary_emb=None, - **cross_attention_kwargs, - ) - cross_attn_output = cross_attn_output * self.gate.tanh().view(1, 1, -1, 1) - mixed_attn_output = self_attn_output + cross_attn_output - mixed_attn_output = mixed_attn_output.flatten(-2) - # linear proj - hidden_states = self.attn2.to_out[0](mixed_attn_output) - - hidden_states = residual + gate_msa.unsqueeze(1).tanh() * self.norm2(hidden_states) - - mlp_output = self.feed_forward(self.ffn_norm1(hidden_states) * (1 + scale_mlp.unsqueeze(1))) - - hidden_states = hidden_states + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(mlp_output) - - return hidden_states - - -class LuminaNextDiT2DModel(ModelMixin, ConfigMixin): - """ - LuminaNextDiT: Diffusion model with a Transformer backbone. - - Inherit ModelMixin and ConfigMixin to be compatible with the sampler StableDiffusionPipeline of diffusers. - - Parameters: - sample_size (`int`): The width of the latent images. This is fixed during training since - it is used to learn a number of position embeddings. - patch_size (`int`, *optional*, (`int`, *optional*, defaults to 2): - The size of each patch in the image. This parameter defines the resolution of patches fed into the model. - in_channels (`int`, *optional*, defaults to 4): - The number of input channels for the model. Typically, this matches the number of channels in the input - images. - hidden_size (`int`, *optional*, defaults to 4096): - The dimensionality of the hidden layers in the model. This parameter determines the width of the model's - hidden representations. - num_layers (`int`, *optional*, default to 32): - The number of layers in the model. This defines the depth of the neural network. - num_attention_heads (`int`, *optional*, defaults to 32): - The number of attention heads in each attention layer. This parameter specifies how many separate attention - mechanisms are used. - num_kv_heads (`int`, *optional*, defaults to 8): - The number of key-value heads in the attention mechanism, if different from the number of attention heads. - If None, it defaults to num_attention_heads. - multiple_of (`int`, *optional*, defaults to 256): - A factor that the hidden size should be a multiple of. This can help optimize certain hardware - configurations. - ffn_dim_multiplier (`float`, *optional*): - A multiplier for the dimensionality of the feed-forward network. If None, it uses a default value based on - the model configuration. - norm_eps (`float`, *optional*, defaults to 1e-5): - A small value added to the denominator for numerical stability in normalization layers. - learn_sigma (`bool`, *optional*, defaults to True): - Whether the model should learn the sigma parameter, which might be related to uncertainty or variance in - predictions. - qk_norm (`bool`, *optional*, defaults to True): - Indicates if the queries and keys in the attention mechanism should be normalized. - cross_attention_dim (`int`, *optional*, defaults to 2048): - The dimensionality of the text embeddings. This parameter defines the size of the text representations used - in the model. - scaling_factor (`float`, *optional*, defaults to 1.0): - A scaling factor applied to certain parameters or layers in the model. This can be used for adjusting the - overall scale of the model's operations. - """ - - _skip_layerwise_casting_patterns = ["patch_embedder", "norm", "ffn_norm"] - - @register_to_config - 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, - ) -> None: - super().__init__() - self.sample_size = sample_size - self.patch_size = patch_size - self.in_channels = in_channels - self.out_channels = in_channels * 2 if learn_sigma else in_channels - self.hidden_size = hidden_size - self.num_attention_heads = num_attention_heads - self.head_dim = hidden_size // num_attention_heads - self.scaling_factor = scaling_factor - - self.patch_embedder = LuminaPatchEmbed( - patch_size=patch_size, in_channels=in_channels, embed_dim=hidden_size, bias=True - ) - - self.pad_token = nn.Parameter(torch.empty(hidden_size)) - - self.time_caption_embed = LuminaCombinedTimestepCaptionEmbedding( - hidden_size=min(hidden_size, 1024), cross_attention_dim=cross_attention_dim - ) - - self.layers = nn.ModuleList( - [ - LuminaNextDiTBlock( - hidden_size, - num_attention_heads, - num_kv_heads, - multiple_of, - ffn_dim_multiplier, - norm_eps, - qk_norm, - cross_attention_dim, - ) - for _ in range(num_layers) - ] - ) - self.norm_out = LuminaLayerNormContinuous( - embedding_dim=hidden_size, - conditioning_embedding_dim=min(hidden_size, 1024), - elementwise_affine=False, - eps=1e-6, - bias=True, - out_dim=patch_size * patch_size * self.out_channels, - ) - # self.final_layer = LuminaFinalLayer(hidden_size, patch_size, self.out_channels) - - assert (hidden_size // num_attention_heads) % 4 == 0, "2d rope needs head dim to be divisible by 4" - - def forward( - self, - hidden_states: torch.Tensor, - timestep: torch.Tensor, - encoder_hidden_states: torch.Tensor, - encoder_mask: torch.Tensor, - image_rotary_emb: torch.Tensor, - cross_attention_kwargs: Dict[str, Any] = None, - return_dict=True, - ) -> torch.Tensor: - """ - Forward pass of LuminaNextDiT. - - Parameters: - hidden_states (torch.Tensor): Input tensor of shape (N, C, H, W). - timestep (torch.Tensor): Tensor of diffusion timesteps of shape (N,). - encoder_hidden_states (torch.Tensor): Tensor of caption features of shape (N, D). - encoder_mask (torch.Tensor): Tensor of caption masks of shape (N, L). - """ - hidden_states, mask, img_size, image_rotary_emb = self.patch_embedder(hidden_states, image_rotary_emb) - image_rotary_emb = image_rotary_emb.to(hidden_states.device) - - temb = self.time_caption_embed(timestep, encoder_hidden_states, encoder_mask) - - encoder_mask = encoder_mask.bool() - for layer in self.layers: - hidden_states = layer( - hidden_states, - mask, - image_rotary_emb, - encoder_hidden_states, - encoder_mask, - temb=temb, - cross_attention_kwargs=cross_attention_kwargs, - ) - - hidden_states = self.norm_out(hidden_states, temb) - - # unpatchify - height_tokens = width_tokens = self.patch_size - height, width = img_size[0] - batch_size = hidden_states.size(0) - sequence_length = (height // height_tokens) * (width // width_tokens) - hidden_states = hidden_states[:, :sequence_length].view( - batch_size, height // height_tokens, width // width_tokens, height_tokens, width_tokens, self.out_channels - ) - output = hidden_states.permute(0, 5, 1, 3, 2, 4).flatten(4, 5).flatten(2, 3) - - if not return_dict: - return (output,) - - return Transformer2DModelOutput(sample=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/pixart_transformer_2d.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/pixart_transformer_2d.py deleted file mode 100644 index 40a14bfd9b2734f9be857ed6adaeb6dd2301ee78..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/pixart_transformer_2d.py +++ /dev/null @@ -1,430 +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 Any, Dict, Optional, Union - -import torch -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import logging -from ..attention import BasicTransformerBlock -from ..attention_processor import Attention, AttentionProcessor, AttnProcessor, FusedAttnProcessor2_0 -from ..embeddings import PatchEmbed, PixArtAlphaTextProjection -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormSingle - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class PixArtTransformer2DModel(ModelMixin, ConfigMixin): - r""" - A 2D Transformer model as introduced in PixArt family of models (https://huggingface.co/papers/2310.00426, - https://huggingface.co/papers/2403.04692). - - Parameters: - num_attention_heads (int, optional, defaults to 16): The number of heads to use for multi-head attention. - attention_head_dim (int, optional, defaults to 72): The number of channels in each head. - in_channels (int, defaults to 4): The number of channels in the input. - out_channels (int, optional): - The number of channels in the output. Specify this parameter if the output channel number differs from the - input. - num_layers (int, optional, defaults to 28): The number of layers of Transformer blocks to use. - dropout (float, optional, defaults to 0.0): The dropout probability to use within the Transformer blocks. - norm_num_groups (int, optional, defaults to 32): - Number of groups for group normalization within Transformer blocks. - cross_attention_dim (int, optional): - The dimensionality for cross-attention layers, typically matching the encoder's hidden dimension. - attention_bias (bool, optional, defaults to True): - Configure if the Transformer blocks' attention should contain a bias parameter. - sample_size (int, defaults to 128): - The width of the latent images. This parameter is fixed during training. - patch_size (int, defaults to 2): - Size of the patches the model processes, relevant for architectures working on non-sequential data. - activation_fn (str, optional, defaults to "gelu-approximate"): - Activation function to use in feed-forward networks within Transformer blocks. - num_embeds_ada_norm (int, optional, defaults to 1000): - Number of embeddings for AdaLayerNorm, fixed during training and affects the maximum denoising steps during - inference. - upcast_attention (bool, optional, defaults to False): - If true, upcasts the attention mechanism dimensions for potentially improved performance. - norm_type (str, optional, defaults to "ada_norm_zero"): - Specifies the type of normalization used, can be 'ada_norm_zero'. - norm_elementwise_affine (bool, optional, defaults to False): - If true, enables element-wise affine parameters in the normalization layers. - norm_eps (float, optional, defaults to 1e-6): - A small constant added to the denominator in normalization layers to prevent division by zero. - interpolation_scale (int, optional): Scale factor to use during interpolating the position embeddings. - use_additional_conditions (bool, optional): If we're using additional conditions as inputs. - attention_type (str, optional, defaults to "default"): Kind of attention mechanism to be used. - caption_channels (int, optional, defaults to None): - Number of channels to use for projecting the caption embeddings. - use_linear_projection (bool, optional, defaults to False): - Deprecated argument. Will be removed in a future version. - num_vector_embeds (bool, optional, defaults to False): - Deprecated argument. Will be removed in a future version. - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["BasicTransformerBlock", "PatchEmbed"] - _skip_layerwise_casting_patterns = ["pos_embed", "norm", "adaln_single"] - - @register_to_config - def __init__( - self, - num_attention_heads: int = 16, - attention_head_dim: int = 72, - in_channels: int = 4, - out_channels: Optional[int] = 8, - num_layers: int = 28, - dropout: float = 0.0, - norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = 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, - 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", - ): - super().__init__() - - # Validate inputs. - if norm_type != "ada_norm_single": - raise NotImplementedError( - f"Forward pass is not implemented when `patch_size` is not None and `norm_type` is '{norm_type}'." - ) - elif norm_type == "ada_norm_single" and num_embeds_ada_norm is None: - raise ValueError( - f"When using a `patch_size` and this `norm_type` ({norm_type}), `num_embeds_ada_norm` cannot be None." - ) - - # Set some common variables used across the board. - self.attention_head_dim = attention_head_dim - self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim - self.out_channels = in_channels if out_channels is None else out_channels - if use_additional_conditions is None: - if sample_size == 128: - use_additional_conditions = True - else: - use_additional_conditions = False - self.use_additional_conditions = use_additional_conditions - - self.gradient_checkpointing = False - - # 2. Initialize the position embedding and transformer blocks. - self.height = self.config.sample_size - self.width = self.config.sample_size - - interpolation_scale = ( - self.config.interpolation_scale - if self.config.interpolation_scale is not None - else max(self.config.sample_size // 64, 1) - ) - self.pos_embed = PatchEmbed( - height=self.config.sample_size, - width=self.config.sample_size, - patch_size=self.config.patch_size, - in_channels=self.config.in_channels, - embed_dim=self.inner_dim, - interpolation_scale=interpolation_scale, - ) - - self.transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - self.inner_dim, - self.config.num_attention_heads, - self.config.attention_head_dim, - dropout=self.config.dropout, - cross_attention_dim=self.config.cross_attention_dim, - activation_fn=self.config.activation_fn, - num_embeds_ada_norm=self.config.num_embeds_ada_norm, - attention_bias=self.config.attention_bias, - upcast_attention=self.config.upcast_attention, - norm_type=norm_type, - norm_elementwise_affine=self.config.norm_elementwise_affine, - norm_eps=self.config.norm_eps, - attention_type=self.config.attention_type, - ) - for _ in range(self.config.num_layers) - ] - ) - - # 3. Output blocks. - self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6) - self.scale_shift_table = nn.Parameter(torch.randn(2, self.inner_dim) / self.inner_dim**0.5) - self.proj_out = nn.Linear(self.inner_dim, self.config.patch_size * self.config.patch_size * self.out_channels) - - self.adaln_single = AdaLayerNormSingle( - self.inner_dim, use_additional_conditions=self.use_additional_conditions - ) - self.caption_projection = None - if self.config.caption_channels is not None: - self.caption_projection = PixArtAlphaTextProjection( - in_features=self.config.caption_channels, hidden_size=self.inner_dim - ) - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - - Safe to just use `AttnProcessor()` as PixArt doesn't have any exotic attention processors in default model. - """ - self.set_attn_processor(AttnProcessor()) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - self.set_attn_processor(FusedAttnProcessor2_0()) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - if self.original_attn_processors is not None: - self.set_attn_processor(self.original_attn_processors) - - 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, - return_dict: bool = True, - ): - """ - The [`PixArtTransformer2DModel`] forward method. - - Args: - hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`): - Input `hidden_states`. - encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*): - Conditional embeddings for cross attention layer. If not given, cross-attention defaults to - 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*): - 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). - attention_mask ( `torch.Tensor`, *optional*): - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask - is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large - negative values to the attention scores corresponding to "discard" tokens. - encoder_attention_mask ( `torch.Tensor`, *optional*): - Cross-attention mask applied to `encoder_hidden_states`. Two formats supported: - - * Mask `(batch, sequence_length)` True = keep, False = discard. - * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard. - - If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format - above. This bias will be added to the cross-attention scores. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] 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 self.use_additional_conditions and added_cond_kwargs is None: - raise ValueError("`added_cond_kwargs` cannot be None when using additional conditions for `adaln_single`.") - - # 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. - # expects mask of shape: - # [batch, key_tokens] - # adds singleton query_tokens dimension: - # [batch, 1, key_tokens] - # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: - # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) - # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) - if attention_mask is not None and attention_mask.ndim == 2: - # assume that mask is expressed as: - # (1 = keep, 0 = discard) - # convert mask into a bias that can be added to attention scores: - # (keep = +0, discard = -10000.0) - attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0 - attention_mask = attention_mask.unsqueeze(1) - - # 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) - - # 1. Input - batch_size = hidden_states.shape[0] - height, width = ( - hidden_states.shape[-2] // self.config.patch_size, - hidden_states.shape[-1] // self.config.patch_size, - ) - hidden_states = self.pos_embed(hidden_states) - - timestep, embedded_timestep = self.adaln_single( - timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_states.dtype - ) - - if self.caption_projection is not None: - encoder_hidden_states = self.caption_projection(encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1]) - - # 2. 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, - encoder_hidden_states, - encoder_attention_mask, - timestep, - cross_attention_kwargs, - None, - ) - else: - hidden_states = block( - hidden_states, - attention_mask=attention_mask, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - timestep=timestep, - cross_attention_kwargs=cross_attention_kwargs, - class_labels=None, - ) - - # 3. Output - shift, scale = ( - self.scale_shift_table[None] + embedded_timestep[:, None].to(self.scale_shift_table.device) - ).chunk(2, dim=1) - hidden_states = self.norm_out(hidden_states) - # Modulation - hidden_states = hidden_states * (1 + scale.to(hidden_states.device)) + shift.to(hidden_states.device) - hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.squeeze(1) - - # unpatchify - hidden_states = hidden_states.reshape( - shape=(-1, height, width, self.config.patch_size, self.config.patch_size, self.out_channels) - ) - hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) - output = hidden_states.reshape( - shape=(-1, self.out_channels, height * self.config.patch_size, width * self.config.patch_size) - ) - - if not return_dict: - return (output,) - - return Transformer2DModelOutput(sample=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/prior_transformer.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/prior_transformer.py deleted file mode 100644 index 565da0da8b6eefddc3653976dba21a4cf687c68a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/prior_transformer.py +++ /dev/null @@ -1,384 +0,0 @@ -from dataclasses import dataclass -from typing import Dict, Optional, Union - -import torch -import torch.nn.functional as F -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin -from ...utils import BaseOutput -from ..attention import BasicTransformerBlock -from ..attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - AttentionProcessor, - AttnAddedKVProcessor, - AttnProcessor, -) -from ..embeddings import TimestepEmbedding, Timesteps -from ..modeling_utils import ModelMixin - - -@dataclass -class PriorTransformerOutput(BaseOutput): - """ - The output of [`PriorTransformer`]. - - Args: - predicted_image_embedding (`torch.Tensor` of shape `(batch_size, embedding_dim)`): - The predicted CLIP image embedding conditioned on the CLIP text embedding input. - """ - - predicted_image_embedding: torch.Tensor - - -class PriorTransformer(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin): - """ - A Prior Transformer model. - - Parameters: - num_attention_heads (`int`, *optional*, defaults to 32): The number of heads to use for multi-head attention. - attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head. - num_layers (`int`, *optional*, defaults to 20): The number of layers of Transformer blocks to use. - embedding_dim (`int`, *optional*, defaults to 768): The dimension of the model input `hidden_states` - num_embeddings (`int`, *optional*, defaults to 77): - The number of embeddings of the model input `hidden_states` - additional_embeddings (`int`, *optional*, defaults to 4): The number of additional tokens appended to the - projected `hidden_states`. The actual length of the used `hidden_states` is `num_embeddings + - additional_embeddings`. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - time_embed_act_fn (`str`, *optional*, defaults to 'silu'): - The activation function to use to create timestep embeddings. - norm_in_type (`str`, *optional*, defaults to None): The normalization layer to apply on hidden states before - passing to Transformer blocks. Set it to `None` if normalization is not needed. - embedding_proj_norm_type (`str`, *optional*, defaults to None): - The normalization layer to apply on the input `proj_embedding`. Set it to `None` if normalization is not - needed. - encoder_hid_proj_type (`str`, *optional*, defaults to `linear`): - The projection layer to apply on the input `encoder_hidden_states`. Set it to `None` if - `encoder_hidden_states` is `None`. - added_emb_type (`str`, *optional*, defaults to `prd`): Additional embeddings to condition the model. - Choose from `prd` or `None`. if choose `prd`, it will prepend a token indicating the (quantized) dot - product between the text embedding and image embedding as proposed in the unclip paper - https://huggingface.co/papers/2204.06125 If it is `None`, no additional embeddings will be prepended. - time_embed_dim (`int, *optional*, defaults to None): The dimension of timestep embeddings. - If None, will be set to `num_attention_heads * attention_head_dim` - embedding_proj_dim (`int`, *optional*, default to None): - The dimension of `proj_embedding`. If None, will be set to `embedding_dim`. - clip_embed_dim (`int`, *optional*, default to None): - The dimension of the output. If None, will be set to `embedding_dim`. - """ - - @register_to_config - def __init__( - self, - num_attention_heads: int = 32, - attention_head_dim: int = 64, - num_layers: int = 20, - embedding_dim: int = 768, - num_embeddings=77, - 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, - ): - super().__init__() - self.num_attention_heads = num_attention_heads - self.attention_head_dim = attention_head_dim - inner_dim = num_attention_heads * attention_head_dim - self.additional_embeddings = additional_embeddings - - time_embed_dim = time_embed_dim or inner_dim - embedding_proj_dim = embedding_proj_dim or embedding_dim - clip_embed_dim = clip_embed_dim or embedding_dim - - self.time_proj = Timesteps(inner_dim, True, 0) - self.time_embedding = TimestepEmbedding(inner_dim, time_embed_dim, out_dim=inner_dim, act_fn=time_embed_act_fn) - - self.proj_in = nn.Linear(embedding_dim, inner_dim) - - if embedding_proj_norm_type is None: - self.embedding_proj_norm = None - elif embedding_proj_norm_type == "layer": - self.embedding_proj_norm = nn.LayerNorm(embedding_proj_dim) - else: - raise ValueError(f"unsupported embedding_proj_norm_type: {embedding_proj_norm_type}") - - self.embedding_proj = nn.Linear(embedding_proj_dim, inner_dim) - - if encoder_hid_proj_type is None: - self.encoder_hidden_states_proj = None - elif encoder_hid_proj_type == "linear": - self.encoder_hidden_states_proj = nn.Linear(embedding_dim, inner_dim) - else: - raise ValueError(f"unsupported encoder_hid_proj_type: {encoder_hid_proj_type}") - - self.positional_embedding = nn.Parameter(torch.zeros(1, num_embeddings + additional_embeddings, inner_dim)) - - if added_emb_type == "prd": - self.prd_embedding = nn.Parameter(torch.zeros(1, 1, inner_dim)) - elif added_emb_type is None: - self.prd_embedding = None - else: - raise ValueError( - f"`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `'prd'` or `None`." - ) - - self.transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - inner_dim, - num_attention_heads, - attention_head_dim, - dropout=dropout, - activation_fn="gelu", - attention_bias=True, - ) - for d in range(num_layers) - ] - ) - - if norm_in_type == "layer": - self.norm_in = nn.LayerNorm(inner_dim) - elif norm_in_type is None: - self.norm_in = None - else: - raise ValueError(f"Unsupported norm_in_type: {norm_in_type}.") - - self.norm_out = nn.LayerNorm(inner_dim) - - self.proj_to_clip_embeddings = nn.Linear(inner_dim, clip_embed_dim) - - causal_attention_mask = torch.full( - [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings], -10000.0 - ) - causal_attention_mask.triu_(1) - causal_attention_mask = causal_attention_mask[None, ...] - self.register_buffer("causal_attention_mask", causal_attention_mask, persistent=False) - - self.clip_mean = nn.Parameter(torch.zeros(1, clip_embed_dim)) - self.clip_std = nn.Parameter(torch.zeros(1, clip_embed_dim)) - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - self.set_attn_processor(processor) - - def forward( - self, - hidden_states, - timestep: Union[torch.Tensor, float, int], - proj_embedding: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.BoolTensor] = None, - return_dict: bool = True, - ): - """ - The [`PriorTransformer`] forward method. - - Args: - hidden_states (`torch.Tensor` of shape `(batch_size, embedding_dim)`): - The currently predicted image embeddings. - timestep (`torch.LongTensor`): - Current denoising step. - proj_embedding (`torch.Tensor` of shape `(batch_size, embedding_dim)`): - Projected embedding vector the denoising process is conditioned on. - encoder_hidden_states (`torch.Tensor` of shape `(batch_size, num_embeddings, embedding_dim)`): - Hidden states of the text embeddings the denoising process is conditioned on. - attention_mask (`torch.BoolTensor` of shape `(batch_size, num_embeddings)`): - Text mask for the text embeddings. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.transformers.prior_transformer.PriorTransformerOutput`] instead of - a plain tuple. - - Returns: - [`~models.transformers.prior_transformer.PriorTransformerOutput`] or `tuple`: - If return_dict is True, a [`~models.transformers.prior_transformer.PriorTransformerOutput`] is - returned, otherwise a tuple is returned where the first element is the sample tensor. - """ - batch_size = hidden_states.shape[0] - - timesteps = timestep - if not torch.is_tensor(timesteps): - timesteps = torch.tensor([timesteps], dtype=torch.long, device=hidden_states.device) - elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0: - timesteps = timesteps[None].to(hidden_states.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps * torch.ones(batch_size, dtype=timesteps.dtype, device=timesteps.device) - - timesteps_projected = self.time_proj(timesteps) - - # timesteps does not contain any weights and will always return f32 tensors - # but time_embedding might be fp16, so we need to cast here. - timesteps_projected = timesteps_projected.to(dtype=self.dtype) - time_embeddings = self.time_embedding(timesteps_projected) - - if self.embedding_proj_norm is not None: - proj_embedding = self.embedding_proj_norm(proj_embedding) - - proj_embeddings = self.embedding_proj(proj_embedding) - if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None: - encoder_hidden_states = self.encoder_hidden_states_proj(encoder_hidden_states) - elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None: - raise ValueError("`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set") - - hidden_states = self.proj_in(hidden_states) - - positional_embeddings = self.positional_embedding.to(hidden_states.dtype) - - additional_embeds = [] - additional_embeddings_len = 0 - - if encoder_hidden_states is not None: - additional_embeds.append(encoder_hidden_states) - additional_embeddings_len += encoder_hidden_states.shape[1] - - if len(proj_embeddings.shape) == 2: - proj_embeddings = proj_embeddings[:, None, :] - - if len(hidden_states.shape) == 2: - hidden_states = hidden_states[:, None, :] - - additional_embeds = additional_embeds + [ - proj_embeddings, - time_embeddings[:, None, :], - hidden_states, - ] - - if self.prd_embedding is not None: - prd_embedding = self.prd_embedding.to(hidden_states.dtype).expand(batch_size, -1, -1) - additional_embeds.append(prd_embedding) - - hidden_states = torch.cat( - additional_embeds, - dim=1, - ) - - # Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens - additional_embeddings_len = additional_embeddings_len + proj_embeddings.shape[1] + 1 - if positional_embeddings.shape[1] < hidden_states.shape[1]: - positional_embeddings = F.pad( - positional_embeddings, - ( - 0, - 0, - additional_embeddings_len, - self.prd_embedding.shape[1] if self.prd_embedding is not None else 0, - ), - value=0.0, - ) - - hidden_states = hidden_states + positional_embeddings - - if attention_mask is not None: - attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0 - attention_mask = F.pad(attention_mask, (0, self.additional_embeddings), value=0.0) - attention_mask = (attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype) - attention_mask = attention_mask.repeat_interleave( - self.config.num_attention_heads, - dim=0, - output_size=attention_mask.shape[0] * self.config.num_attention_heads, - ) - - if self.norm_in is not None: - hidden_states = self.norm_in(hidden_states) - - for block in self.transformer_blocks: - hidden_states = block(hidden_states, attention_mask=attention_mask) - - hidden_states = self.norm_out(hidden_states) - - if self.prd_embedding is not None: - hidden_states = hidden_states[:, -1] - else: - hidden_states = hidden_states[:, additional_embeddings_len:] - - predicted_image_embedding = self.proj_to_clip_embeddings(hidden_states) - - if not return_dict: - return (predicted_image_embedding,) - - return PriorTransformerOutput(predicted_image_embedding=predicted_image_embedding) - - def post_process_latents(self, prior_latents): - prior_latents = (prior_latents * self.clip_std) + self.clip_mean - return prior_latents diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/sana_transformer.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/sana_transformer.py deleted file mode 100644 index 1e02ac32e86c94280b1be9b8a88a6180fd1edc41..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/sana_transformer.py +++ /dev/null @@ -1,597 +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 Any, Dict, Optional, Tuple, Union - -import torch -import torch.nn.functional as F -from torch import 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 ..attention_processor import ( - Attention, - AttentionProcessor, - SanaLinearAttnProcessor2_0, -) -from ..embeddings import PatchEmbed, PixArtAlphaTextProjection, TimestepEmbedding, Timesteps -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormSingle, RMSNorm - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class GLUMBConv(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - expand_ratio: float = 4, - norm_type: Optional[str] = None, - residual_connection: bool = True, - ) -> None: - super().__init__() - - hidden_channels = int(expand_ratio * in_channels) - self.norm_type = norm_type - self.residual_connection = residual_connection - - self.nonlinearity = nn.SiLU() - self.conv_inverted = nn.Conv2d(in_channels, hidden_channels * 2, 1, 1, 0) - self.conv_depth = nn.Conv2d(hidden_channels * 2, hidden_channels * 2, 3, 1, 1, groups=hidden_channels * 2) - self.conv_point = nn.Conv2d(hidden_channels, out_channels, 1, 1, 0, bias=False) - - self.norm = None - if norm_type == "rms_norm": - self.norm = RMSNorm(out_channels, eps=1e-5, elementwise_affine=True, bias=True) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.residual_connection: - residual = hidden_states - - hidden_states = self.conv_inverted(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - - hidden_states = self.conv_depth(hidden_states) - hidden_states, gate = torch.chunk(hidden_states, 2, dim=1) - hidden_states = hidden_states * self.nonlinearity(gate) - - hidden_states = self.conv_point(hidden_states) - - if self.norm_type == "rms_norm": - # move channel to the last dimension so we apply RMSnorm across channel dimension - hidden_states = self.norm(hidden_states.movedim(1, -1)).movedim(-1, 1) - - if self.residual_connection: - hidden_states = hidden_states + residual - - return hidden_states - - -class SanaModulatedNorm(nn.Module): - def __init__(self, dim: int, elementwise_affine: bool = False, eps: float = 1e-6): - super().__init__() - self.norm = nn.LayerNorm(dim, elementwise_affine=elementwise_affine, eps=eps) - - def forward( - self, hidden_states: torch.Tensor, temb: torch.Tensor, scale_shift_table: torch.Tensor - ) -> torch.Tensor: - hidden_states = self.norm(hidden_states) - shift, scale = (scale_shift_table[None] + temb[:, None].to(scale_shift_table.device)).chunk(2, dim=1) - hidden_states = hidden_states * (1 + scale) + shift - return hidden_states - - -class SanaCombinedTimestepGuidanceEmbeddings(nn.Module): - def __init__(self, embedding_dim): - super().__init__() - self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) - self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) - - self.guidance_condition_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) - self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) - - self.silu = nn.SiLU() - self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) - - def forward(self, timestep: torch.Tensor, guidance: torch.Tensor = None, hidden_dtype: torch.dtype = None): - timesteps_proj = self.time_proj(timestep) - timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) - - guidance_proj = self.guidance_condition_proj(guidance) - guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=hidden_dtype)) - conditioning = timesteps_emb + guidance_emb - - return self.linear(self.silu(conditioning)), conditioning - - -class SanaAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("SanaAttnProcessor2_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, - ) -> 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) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class SanaTransformerBlock(nn.Module): - r""" - Transformer block introduced in [Sana](https://huggingface.co/papers/2410.10629). - """ - - def __init__( - self, - dim: int = 2240, - 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, - 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, - ) -> None: - super().__init__() - - # 1. Self Attention - self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps) - self.attn1 = Attention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - kv_heads=num_attention_heads if qk_norm is not None else None, - qk_norm=qk_norm, - dropout=dropout, - bias=attention_bias, - cross_attention_dim=None, - processor=SanaLinearAttnProcessor2_0(), - ) - - # 2. Cross Attention - if cross_attention_dim is not None: - self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) - self.attn2 = Attention( - query_dim=dim, - qk_norm=qk_norm, - kv_heads=num_cross_attention_heads if qk_norm is not None else None, - cross_attention_dim=cross_attention_dim, - heads=num_cross_attention_heads, - dim_head=cross_attention_head_dim, - dropout=dropout, - bias=True, - out_bias=attention_out_bias, - processor=SanaAttnProcessor2_0(), - ) - - # 3. Feed-forward - self.ff = GLUMBConv(dim, dim, mlp_ratio, norm_type=None, residual_connection=False) - - self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) - - 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, - height: int = None, - width: int = None, - ) -> torch.Tensor: - batch_size = hidden_states.shape[0] - - # 1. Modulation - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( - self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) - ).chunk(6, dim=1) - - # 2. Self Attention - norm_hidden_states = self.norm1(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa - norm_hidden_states = norm_hidden_states.to(hidden_states.dtype) - - attn_output = self.attn1(norm_hidden_states) - hidden_states = hidden_states + gate_msa * attn_output - - # 3. Cross Attention - if self.attn2 is not None: - attn_output = self.attn2( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - ) - hidden_states = attn_output + hidden_states - - # 4. Feed-forward - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp - - norm_hidden_states = norm_hidden_states.unflatten(1, (height, width)).permute(0, 3, 1, 2) - ff_output = self.ff(norm_hidden_states) - ff_output = ff_output.flatten(2, 3).permute(0, 2, 1) - hidden_states = hidden_states + gate_mlp * ff_output - - return hidden_states - - -class SanaTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): - r""" - A 2D Transformer model introduced in [Sana](https://huggingface.co/papers/2410.10629) family of models. - - Args: - in_channels (`int`, defaults to `32`): - The number of channels in the input. - out_channels (`int`, *optional*, defaults to `32`): - The number of channels in the output. - num_attention_heads (`int`, defaults to `70`): - The number of heads to use for multi-head attention. - attention_head_dim (`int`, defaults to `32`): - The number of channels in each head. - num_layers (`int`, defaults to `20`): - The number of layers of Transformer blocks to use. - num_cross_attention_heads (`int`, *optional*, defaults to `20`): - The number of heads to use for cross-attention. - cross_attention_head_dim (`int`, *optional*, defaults to `112`): - The number of channels in each head for cross-attention. - cross_attention_dim (`int`, *optional*, defaults to `2240`): - The number of channels in the cross-attention output. - caption_channels (`int`, defaults to `2304`): - The number of channels in the caption embeddings. - mlp_ratio (`float`, defaults to `2.5`): - The expansion ratio to use in the GLUMBConv layer. - dropout (`float`, defaults to `0.0`): - The dropout probability. - attention_bias (`bool`, defaults to `False`): - Whether to use bias in the attention layer. - sample_size (`int`, defaults to `32`): - The base size of the input latent. - patch_size (`int`, defaults to `1`): - The size of the patches to use in the patch embedding layer. - norm_elementwise_affine (`bool`, defaults to `False`): - Whether to use elementwise affinity in the normalization layer. - norm_eps (`float`, defaults to `1e-6`): - The epsilon value for the normalization layer. - qk_norm (`str`, *optional*, defaults to `None`): - The normalization to use for the query and key. - timestep_scale (`float`, defaults to `1.0`): - The scale to use for the timesteps. - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["SanaTransformerBlock", "PatchEmbed", "SanaModulatedNorm"] - _skip_layerwise_casting_patterns = ["patch_embed", "norm"] - - @register_to_config - def __init__( - self, - in_channels: int = 32, - out_channels: Optional[int] = 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, - caption_channels: int = 2304, - mlp_ratio: float = 2.5, - dropout: float = 0.0, - attention_bias: bool = False, - sample_size: int = 32, - patch_size: int = 1, - norm_elementwise_affine: bool = False, - norm_eps: float = 1e-6, - interpolation_scale: Optional[int] = None, - guidance_embeds: bool = False, - guidance_embeds_scale: float = 0.1, - qk_norm: Optional[str] = None, - timestep_scale: float = 1.0, - ) -> None: - super().__init__() - - out_channels = out_channels or in_channels - inner_dim = num_attention_heads * attention_head_dim - - # 1. Patch Embedding - self.patch_embed = PatchEmbed( - height=sample_size, - width=sample_size, - patch_size=patch_size, - in_channels=in_channels, - embed_dim=inner_dim, - interpolation_scale=interpolation_scale, - pos_embed_type="sincos" if interpolation_scale is not None else None, - ) - - # 2. Additional condition embeddings - if guidance_embeds: - self.time_embed = SanaCombinedTimestepGuidanceEmbeddings(inner_dim) - else: - self.time_embed = AdaLayerNormSingle(inner_dim) - - self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) - self.caption_norm = RMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) - - # 3. Transformer blocks - self.transformer_blocks = nn.ModuleList( - [ - SanaTransformerBlock( - inner_dim, - num_attention_heads, - attention_head_dim, - dropout=dropout, - num_cross_attention_heads=num_cross_attention_heads, - cross_attention_head_dim=cross_attention_head_dim, - cross_attention_dim=cross_attention_dim, - attention_bias=attention_bias, - norm_elementwise_affine=norm_elementwise_affine, - norm_eps=norm_eps, - mlp_ratio=mlp_ratio, - qk_norm=qk_norm, - ) - for _ in range(num_layers) - ] - ) - - # 4. Output blocks - self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5) - self.norm_out = SanaModulatedNorm(inner_dim, elementwise_affine=False, eps=1e-6) - self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels) - - self.gradient_checkpointing = False - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - 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, - 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." - ) - - # 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. - # expects mask of shape: - # [batch, key_tokens] - # adds singleton query_tokens dimension: - # [batch, 1, key_tokens] - # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: - # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) - # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) - if attention_mask is not None and attention_mask.ndim == 2: - # assume that mask is expressed as: - # (1 = keep, 0 = discard) - # convert mask into a bias that can be added to attention scores: - # (keep = +0, discard = -10000.0) - attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0 - attention_mask = attention_mask.unsqueeze(1) - - # 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) - - # 1. Input - batch_size, num_channels, height, width = hidden_states.shape - p = self.config.patch_size - post_patch_height, post_patch_width = height // p, width // p - - hidden_states = self.patch_embed(hidden_states) - - if guidance is not None: - timestep, embedded_timestep = self.time_embed( - timestep, guidance=guidance, hidden_dtype=hidden_states.dtype - ) - else: - timestep, embedded_timestep = self.time_embed( - timestep, batch_size=batch_size, hidden_dtype=hidden_states.dtype - ) - - encoder_hidden_states = self.caption_projection(encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1]) - - encoder_hidden_states = self.caption_norm(encoder_hidden_states) - - # 2. Transformer blocks - if torch.is_grad_enabled() and self.gradient_checkpointing: - for index_block, block in enumerate(self.transformer_blocks): - hidden_states = self._gradient_checkpointing_func( - block, - hidden_states, - attention_mask, - encoder_hidden_states, - encoder_attention_mask, - timestep, - post_patch_height, - post_patch_width, - ) - if controlnet_block_samples is not None and 0 < index_block <= len(controlnet_block_samples): - hidden_states = hidden_states + controlnet_block_samples[index_block - 1] - - else: - for index_block, block in enumerate(self.transformer_blocks): - hidden_states = block( - hidden_states, - attention_mask, - encoder_hidden_states, - encoder_attention_mask, - timestep, - post_patch_height, - post_patch_width, - ) - if controlnet_block_samples is not None and 0 < index_block <= len(controlnet_block_samples): - hidden_states = hidden_states + controlnet_block_samples[index_block - 1] - - # 3. Normalization - hidden_states = self.norm_out(hidden_states, embedded_timestep, self.scale_shift_table) - - hidden_states = self.proj_out(hidden_states) - - # 5. Unpatchify - hidden_states = hidden_states.reshape( - batch_size, post_patch_height, post_patch_width, self.config.patch_size, self.config.patch_size, -1 - ) - 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,) - - return Transformer2DModelOutput(sample=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/stable_audio_transformer.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/stable_audio_transformer.py deleted file mode 100644 index 969e6db122d9373fb599bfd111062889901fb42d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/stable_audio_transformer.py +++ /dev/null @@ -1,439 +0,0 @@ -# Copyright 2025 Stability 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. - - -from typing import Dict, Optional, Union - -import numpy as np -import torch -import torch.nn as nn -import torch.utils.checkpoint - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import logging -from ...utils.torch_utils import maybe_allow_in_graph -from ..attention import FeedForward -from ..attention_processor import Attention, AttentionProcessor, StableAudioAttnProcessor2_0 -from ..modeling_utils import ModelMixin -from ..transformers.transformer_2d import Transformer2DModelOutput - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class StableAudioGaussianFourierProjection(nn.Module): - """Gaussian Fourier embeddings for noise levels.""" - - # Copied from diffusers.models.embeddings.GaussianFourierProjection.__init__ - def __init__( - self, embedding_size: int = 256, scale: float = 1.0, set_W_to_weight=True, log=True, flip_sin_to_cos=False - ): - super().__init__() - self.weight = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False) - self.log = log - self.flip_sin_to_cos = flip_sin_to_cos - - if set_W_to_weight: - # to delete later - del self.weight - self.W = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False) - self.weight = self.W - del self.W - - def forward(self, x): - if self.log: - x = torch.log(x) - - x_proj = 2 * np.pi * x[:, None] @ self.weight[None, :] - - if self.flip_sin_to_cos: - out = torch.cat([torch.cos(x_proj), torch.sin(x_proj)], dim=-1) - else: - out = torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1) - return out - - -@maybe_allow_in_graph -class StableAudioDiTBlock(nn.Module): - r""" - Transformer block used in Stable Audio model (https://github.com/Stability-AI/stable-audio-tools). Allow skip - connection and QKNorm - - Parameters: - dim (`int`): The number of channels in the input and output. - num_attention_heads (`int`): The number of heads to use for the query states. - num_key_value_attention_heads (`int`): The number of heads to use for the key and value states. - attention_head_dim (`int`): The number of channels in each head. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. - upcast_attention (`bool`, *optional*): - Whether to upcast the attention computation to float32. This is useful for mixed precision training. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - num_key_value_attention_heads: int, - attention_head_dim: int, - dropout=0.0, - cross_attention_dim: Optional[int] = None, - upcast_attention: bool = False, - norm_eps: float = 1e-5, - ff_inner_dim: Optional[int] = None, - ): - super().__init__() - # Define 3 blocks. Each block has its own normalization layer. - # 1. Self-Attn - self.norm1 = nn.LayerNorm(dim, elementwise_affine=True, eps=norm_eps) - self.attn1 = Attention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=False, - upcast_attention=upcast_attention, - out_bias=False, - processor=StableAudioAttnProcessor2_0(), - ) - - # 2. Cross-Attn - self.norm2 = nn.LayerNorm(dim, norm_eps, True) - - self.attn2 = Attention( - query_dim=dim, - cross_attention_dim=cross_attention_dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - kv_heads=num_key_value_attention_heads, - dropout=dropout, - bias=False, - upcast_attention=upcast_attention, - out_bias=False, - processor=StableAudioAttnProcessor2_0(), - ) # is self-attn if encoder_hidden_states is none - - # 3. Feed-forward - self.norm3 = nn.LayerNorm(dim, norm_eps, True) - self.ff = FeedForward( - dim, - dropout=dropout, - activation_fn="swiglu", - final_dropout=False, - inner_dim=ff_inner_dim, - bias=True, - ) - - # let chunk size default to None - self._chunk_size = None - self._chunk_dim = 0 - - def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): - # Sets chunk feed-forward - self._chunk_size = chunk_size - self._chunk_dim = dim - - 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, - ) -> torch.Tensor: - # Notice that normalization is always applied before the real computation in the following blocks. - # 0. Self-Attention - norm_hidden_states = self.norm1(hidden_states) - - attn_output = self.attn1( - norm_hidden_states, - attention_mask=attention_mask, - rotary_emb=rotary_embedding, - ) - - hidden_states = attn_output + hidden_states - - # 2. Cross-Attention - norm_hidden_states = self.norm2(hidden_states) - - attn_output = self.attn2( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - ) - hidden_states = attn_output + hidden_states - - # 3. Feed-forward - norm_hidden_states = self.norm3(hidden_states) - ff_output = self.ff(norm_hidden_states) - - hidden_states = ff_output + hidden_states - - return hidden_states - - -class StableAudioDiTModel(ModelMixin, ConfigMixin): - """ - The Diffusion Transformer model introduced in Stable Audio. - - Reference: https://github.com/Stability-AI/stable-audio-tools - - Parameters: - sample_size ( `int`, *optional*, defaults to 1024): The size of the input sample. - in_channels (`int`, *optional*, defaults to 64): The number of channels in the input. - num_layers (`int`, *optional*, defaults to 24): The number of layers of Transformer blocks to use. - attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head. - num_attention_heads (`int`, *optional*, defaults to 24): The number of heads to use for the query states. - num_key_value_attention_heads (`int`, *optional*, defaults to 12): - The number of heads to use for the key and value states. - out_channels (`int`, defaults to 64): Number of output channels. - cross_attention_dim ( `int`, *optional*, defaults to 768): Dimension of the cross-attention projection. - time_proj_dim ( `int`, *optional*, defaults to 256): Dimension of the timestep inner projection. - global_states_input_dim ( `int`, *optional*, defaults to 1536): - Input dimension of the global hidden states projection. - cross_attention_input_dim ( `int`, *optional*, defaults to 768): - Input dimension of the cross-attention projection - """ - - _supports_gradient_checkpointing = True - _skip_layerwise_casting_patterns = ["preprocess_conv", "postprocess_conv", "^proj_in$", "^proj_out$", "norm"] - - @register_to_config - def __init__( - self, - sample_size: int = 1024, - in_channels: int = 64, - num_layers: int = 24, - attention_head_dim: int = 64, - num_attention_heads: int = 24, - num_key_value_attention_heads: int = 12, - out_channels: int = 64, - cross_attention_dim: int = 768, - time_proj_dim: int = 256, - global_states_input_dim: int = 1536, - cross_attention_input_dim: int = 768, - ): - super().__init__() - self.sample_size = sample_size - self.out_channels = out_channels - self.inner_dim = num_attention_heads * attention_head_dim - - self.time_proj = StableAudioGaussianFourierProjection( - embedding_size=time_proj_dim // 2, - flip_sin_to_cos=True, - log=False, - set_W_to_weight=False, - ) - - self.timestep_proj = nn.Sequential( - nn.Linear(time_proj_dim, self.inner_dim, bias=True), - nn.SiLU(), - nn.Linear(self.inner_dim, self.inner_dim, bias=True), - ) - - self.global_proj = nn.Sequential( - nn.Linear(global_states_input_dim, self.inner_dim, bias=False), - nn.SiLU(), - nn.Linear(self.inner_dim, self.inner_dim, bias=False), - ) - - self.cross_attention_proj = nn.Sequential( - nn.Linear(cross_attention_input_dim, cross_attention_dim, bias=False), - nn.SiLU(), - nn.Linear(cross_attention_dim, cross_attention_dim, bias=False), - ) - - self.preprocess_conv = nn.Conv1d(in_channels, in_channels, 1, bias=False) - self.proj_in = nn.Linear(in_channels, self.inner_dim, bias=False) - - self.transformer_blocks = nn.ModuleList( - [ - StableAudioDiTBlock( - dim=self.inner_dim, - num_attention_heads=num_attention_heads, - num_key_value_attention_heads=num_key_value_attention_heads, - attention_head_dim=attention_head_dim, - cross_attention_dim=cross_attention_dim, - ) - for i in range(num_layers) - ] - ) - - self.proj_out = nn.Linear(self.inner_dim, self.out_channels, bias=False) - self.postprocess_conv = nn.Conv1d(self.out_channels, self.out_channels, 1, bias=False) - - self.gradient_checkpointing = False - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.transformers.hunyuan_transformer_2d.HunyuanDiT2DModel.set_default_attn_processor with Hunyuan->StableAudio - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - self.set_attn_processor(StableAudioAttnProcessor2_0()) - - def forward( - self, - hidden_states: torch.FloatTensor, - timestep: torch.LongTensor = None, - encoder_hidden_states: torch.FloatTensor = None, - 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]: - """ - The [`StableAudioDiTModel`] forward method. - - Args: - hidden_states (`torch.FloatTensor` of shape `(batch size, in_channels, sequence_len)`): - Input `hidden_states`. - timestep ( `torch.LongTensor`): - Used to indicate denoising step. - encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, encoder_sequence_len, cross_attention_input_dim)`): - Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. - global_hidden_states (`torch.FloatTensor` of shape `(batch size, global_sequence_len, global_states_input_dim)`): - Global embeddings that will be prepended to the hidden states. - rotary_embedding (`torch.Tensor`): - The rotary embeddings to apply on query and key tensors during attention calculation. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain - tuple. - attention_mask (`torch.Tensor` of shape `(batch_size, sequence_len)`, *optional*): - Mask to avoid performing attention on padding token indices, formed by concatenating the attention - masks - for the two text encoders together. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - encoder_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_len)`, *optional*): - Mask to avoid performing attention on padding token cross-attention indices, formed by concatenating - the attention masks - for the two text encoders together. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - Returns: - If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a - `tuple` where the first element is the sample tensor. - """ - cross_attention_hidden_states = self.cross_attention_proj(encoder_hidden_states) - global_hidden_states = self.global_proj(global_hidden_states) - time_hidden_states = self.timestep_proj(self.time_proj(timestep.to(self.dtype))) - - global_hidden_states = global_hidden_states + time_hidden_states.unsqueeze(1) - - hidden_states = self.preprocess_conv(hidden_states) + hidden_states - # (batch_size, dim, sequence_length) -> (batch_size, sequence_length, dim) - hidden_states = hidden_states.transpose(1, 2) - - hidden_states = self.proj_in(hidden_states) - - # prepend global states to hidden states - hidden_states = torch.cat([global_hidden_states, hidden_states], dim=-2) - if attention_mask is not None: - prepend_mask = torch.ones((hidden_states.shape[0], 1), device=hidden_states.device, dtype=torch.bool) - attention_mask = torch.cat([prepend_mask, attention_mask], dim=-1) - - 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, - cross_attention_hidden_states, - encoder_attention_mask, - rotary_embedding, - ) - - else: - hidden_states = block( - hidden_states=hidden_states, - attention_mask=attention_mask, - encoder_hidden_states=cross_attention_hidden_states, - encoder_attention_mask=encoder_attention_mask, - rotary_embedding=rotary_embedding, - ) - - hidden_states = self.proj_out(hidden_states) - - # (batch_size, sequence_length, dim) -> (batch_size, dim, sequence_length) - # remove prepend length that has been added by global hidden states - hidden_states = hidden_states.transpose(1, 2)[:, :, 1:] - hidden_states = self.postprocess_conv(hidden_states) + hidden_states - - if not return_dict: - return (hidden_states,) - - return Transformer2DModelOutput(sample=hidden_states) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/t5_film_transformer.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/t5_film_transformer.py deleted file mode 100644 index 7a9608735e3298a5f9ceaa940c291e98b9a7409a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/t5_film_transformer.py +++ /dev/null @@ -1,436 +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 math -from typing import Optional, Tuple - -import torch -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ..attention_processor import Attention -from ..embeddings import get_timestep_embedding -from ..modeling_utils import ModelMixin - - -class T5FilmDecoder(ModelMixin, ConfigMixin): - r""" - T5 style decoder with FiLM conditioning. - - Args: - input_dims (`int`, *optional*, defaults to `128`): - The number of input dimensions. - targets_length (`int`, *optional*, defaults to `256`): - The length of the targets. - d_model (`int`, *optional*, defaults to `768`): - Size of the input hidden states. - num_layers (`int`, *optional*, defaults to `12`): - The number of `DecoderLayer`'s to use. - num_heads (`int`, *optional*, defaults to `12`): - The number of attention heads to use. - d_kv (`int`, *optional*, defaults to `64`): - Size of the key-value projection vectors. - d_ff (`int`, *optional*, defaults to `2048`): - The number of dimensions in the intermediate feed-forward layer of `DecoderLayer`'s. - dropout_rate (`float`, *optional*, defaults to `0.1`): - Dropout probability. - """ - - @register_to_config - def __init__( - self, - input_dims: int = 128, - targets_length: int = 256, - max_decoder_noise_time: float = 2000.0, - d_model: int = 768, - num_layers: int = 12, - num_heads: int = 12, - d_kv: int = 64, - d_ff: int = 2048, - dropout_rate: float = 0.1, - ): - super().__init__() - - self.conditioning_emb = nn.Sequential( - nn.Linear(d_model, d_model * 4, bias=False), - nn.SiLU(), - nn.Linear(d_model * 4, d_model * 4, bias=False), - nn.SiLU(), - ) - - self.position_encoding = nn.Embedding(targets_length, d_model) - self.position_encoding.weight.requires_grad = False - - self.continuous_inputs_projection = nn.Linear(input_dims, d_model, bias=False) - - self.dropout = nn.Dropout(p=dropout_rate) - - self.decoders = nn.ModuleList() - for lyr_num in range(num_layers): - # FiLM conditional T5 decoder - lyr = DecoderLayer(d_model=d_model, d_kv=d_kv, num_heads=num_heads, d_ff=d_ff, dropout_rate=dropout_rate) - self.decoders.append(lyr) - - self.decoder_norm = T5LayerNorm(d_model) - - self.post_dropout = nn.Dropout(p=dropout_rate) - self.spec_out = nn.Linear(d_model, input_dims, bias=False) - - def encoder_decoder_mask(self, query_input: torch.Tensor, key_input: torch.Tensor) -> torch.Tensor: - mask = torch.mul(query_input.unsqueeze(-1), key_input.unsqueeze(-2)) - return mask.unsqueeze(-3) - - def forward(self, encodings_and_masks, decoder_input_tokens, decoder_noise_time): - batch, _, _ = decoder_input_tokens.shape - assert decoder_noise_time.shape == (batch,) - - # decoder_noise_time is in [0, 1), so rescale to expected timing range. - time_steps = get_timestep_embedding( - decoder_noise_time * self.config.max_decoder_noise_time, - embedding_dim=self.config.d_model, - max_period=self.config.max_decoder_noise_time, - ).to(dtype=self.dtype) - - conditioning_emb = self.conditioning_emb(time_steps).unsqueeze(1) - - assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4) - - seq_length = decoder_input_tokens.shape[1] - - # If we want to use relative positions for audio context, we can just offset - # this sequence by the length of encodings_and_masks. - decoder_positions = torch.broadcast_to( - torch.arange(seq_length, device=decoder_input_tokens.device), - (batch, seq_length), - ) - - position_encodings = self.position_encoding(decoder_positions) - - inputs = self.continuous_inputs_projection(decoder_input_tokens) - inputs += position_encodings - y = self.dropout(inputs) - - # decoder: No padding present. - decoder_mask = torch.ones( - decoder_input_tokens.shape[:2], device=decoder_input_tokens.device, dtype=inputs.dtype - ) - - # Translate encoding masks to encoder-decoder masks. - encodings_and_encdec_masks = [(x, self.encoder_decoder_mask(decoder_mask, y)) for x, y in encodings_and_masks] - - # cross attend style: concat encodings - encoded = torch.cat([x[0] for x in encodings_and_encdec_masks], dim=1) - encoder_decoder_mask = torch.cat([x[1] for x in encodings_and_encdec_masks], dim=-1) - - for lyr in self.decoders: - y = lyr( - y, - conditioning_emb=conditioning_emb, - encoder_hidden_states=encoded, - encoder_attention_mask=encoder_decoder_mask, - )[0] - - y = self.decoder_norm(y) - y = self.post_dropout(y) - - spec_out = self.spec_out(y) - return spec_out - - -class DecoderLayer(nn.Module): - r""" - T5 decoder layer. - - Args: - d_model (`int`): - Size of the input hidden states. - d_kv (`int`): - Size of the key-value projection vectors. - num_heads (`int`): - Number of attention heads. - d_ff (`int`): - Size of the intermediate feed-forward layer. - dropout_rate (`float`): - Dropout probability. - layer_norm_epsilon (`float`, *optional*, defaults to `1e-6`): - A small value used for numerical stability to avoid dividing by zero. - """ - - def __init__( - self, d_model: int, d_kv: int, num_heads: int, d_ff: int, dropout_rate: float, layer_norm_epsilon: float = 1e-6 - ): - super().__init__() - self.layer = nn.ModuleList() - - # cond self attention: layer 0 - self.layer.append( - T5LayerSelfAttentionCond(d_model=d_model, d_kv=d_kv, num_heads=num_heads, dropout_rate=dropout_rate) - ) - - # cross attention: layer 1 - self.layer.append( - T5LayerCrossAttention( - d_model=d_model, - d_kv=d_kv, - num_heads=num_heads, - dropout_rate=dropout_rate, - layer_norm_epsilon=layer_norm_epsilon, - ) - ) - - # Film Cond MLP + dropout: last layer - self.layer.append( - T5LayerFFCond(d_model=d_model, d_ff=d_ff, dropout_rate=dropout_rate, layer_norm_epsilon=layer_norm_epsilon) - ) - - 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, - encoder_decoder_position_bias=None, - ) -> Tuple[torch.Tensor]: - hidden_states = self.layer[0]( - hidden_states, - conditioning_emb=conditioning_emb, - attention_mask=attention_mask, - ) - - if encoder_hidden_states is not None: - encoder_extended_attention_mask = torch.where(encoder_attention_mask > 0, 0, -1e10).to( - encoder_hidden_states.dtype - ) - - hidden_states = self.layer[1]( - hidden_states, - key_value_states=encoder_hidden_states, - attention_mask=encoder_extended_attention_mask, - ) - - # Apply Film Conditional Feed Forward layer - hidden_states = self.layer[-1](hidden_states, conditioning_emb) - - return (hidden_states,) - - -class T5LayerSelfAttentionCond(nn.Module): - r""" - T5 style self-attention layer with conditioning. - - Args: - d_model (`int`): - Size of the input hidden states. - d_kv (`int`): - Size of the key-value projection vectors. - num_heads (`int`): - Number of attention heads. - dropout_rate (`float`): - Dropout probability. - """ - - def __init__(self, d_model: int, d_kv: int, num_heads: int, dropout_rate: float): - super().__init__() - self.layer_norm = T5LayerNorm(d_model) - self.FiLMLayer = T5FiLMLayer(in_features=d_model * 4, out_features=d_model) - self.attention = Attention(query_dim=d_model, heads=num_heads, dim_head=d_kv, out_bias=False, scale_qk=False) - self.dropout = nn.Dropout(dropout_rate) - - def forward( - self, - hidden_states: torch.Tensor, - conditioning_emb: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - # pre_self_attention_layer_norm - normed_hidden_states = self.layer_norm(hidden_states) - - if conditioning_emb is not None: - normed_hidden_states = self.FiLMLayer(normed_hidden_states, conditioning_emb) - - # Self-attention block - attention_output = self.attention(normed_hidden_states) - - hidden_states = hidden_states + self.dropout(attention_output) - - return hidden_states - - -class T5LayerCrossAttention(nn.Module): - r""" - T5 style cross-attention layer. - - Args: - d_model (`int`): - Size of the input hidden states. - d_kv (`int`): - Size of the key-value projection vectors. - num_heads (`int`): - Number of attention heads. - dropout_rate (`float`): - Dropout probability. - layer_norm_epsilon (`float`): - A small value used for numerical stability to avoid dividing by zero. - """ - - def __init__(self, d_model: int, d_kv: int, num_heads: int, dropout_rate: float, layer_norm_epsilon: float): - super().__init__() - self.attention = Attention(query_dim=d_model, heads=num_heads, dim_head=d_kv, out_bias=False, scale_qk=False) - self.layer_norm = T5LayerNorm(d_model, eps=layer_norm_epsilon) - self.dropout = nn.Dropout(dropout_rate) - - def forward( - self, - hidden_states: torch.Tensor, - key_value_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.attention( - normed_hidden_states, - encoder_hidden_states=key_value_states, - attention_mask=attention_mask.squeeze(1), - ) - layer_output = hidden_states + self.dropout(attention_output) - return layer_output - - -class T5LayerFFCond(nn.Module): - r""" - T5 style feed-forward conditional layer. - - Args: - d_model (`int`): - Size of the input hidden states. - d_ff (`int`): - Size of the intermediate feed-forward layer. - dropout_rate (`float`): - Dropout probability. - layer_norm_epsilon (`float`): - A small value used for numerical stability to avoid dividing by zero. - """ - - def __init__(self, d_model: int, d_ff: int, dropout_rate: float, layer_norm_epsilon: float): - super().__init__() - self.DenseReluDense = T5DenseGatedActDense(d_model=d_model, d_ff=d_ff, dropout_rate=dropout_rate) - self.film = T5FiLMLayer(in_features=d_model * 4, out_features=d_model) - 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: - forwarded_states = self.layer_norm(hidden_states) - if conditioning_emb is not None: - forwarded_states = self.film(forwarded_states, conditioning_emb) - - forwarded_states = self.DenseReluDense(forwarded_states) - hidden_states = hidden_states + self.dropout(forwarded_states) - return hidden_states - - -class T5DenseGatedActDense(nn.Module): - r""" - T5 style feed-forward layer with gated activations and dropout. - - Args: - d_model (`int`): - Size of the input hidden states. - d_ff (`int`): - Size of the intermediate feed-forward layer. - dropout_rate (`float`): - Dropout probability. - """ - - def __init__(self, d_model: int, d_ff: int, dropout_rate: float): - super().__init__() - self.wi_0 = nn.Linear(d_model, d_ff, bias=False) - self.wi_1 = nn.Linear(d_model, d_ff, bias=False) - self.wo = nn.Linear(d_ff, d_model, bias=False) - self.dropout = nn.Dropout(dropout_rate) - self.act = NewGELUActivation() - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_gelu = self.act(self.wi_0(hidden_states)) - hidden_linear = self.wi_1(hidden_states) - hidden_states = hidden_gelu * hidden_linear - hidden_states = self.dropout(hidden_states) - - hidden_states = self.wo(hidden_states) - return hidden_states - - -class T5LayerNorm(nn.Module): - r""" - T5 style layer normalization module. - - Args: - hidden_size (`int`): - Size of the input hidden states. - eps (`float`, `optional`, defaults to `1e-6`): - A small value used for numerical stability to avoid dividing by zero. - """ - - def __init__(self, hidden_size: int, eps: float = 1e-6): - """ - Construct a layernorm module in the T5 style. No bias and no subtraction of mean. - """ - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - # T5 uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean - # Square Layer Normalization https://huggingface.co/papers/1910.07467 thus variance is calculated - # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for - # half-precision inputs is done in fp32 - - variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - - # convert into half-precision if necessary - if self.weight.dtype in [torch.float16, torch.bfloat16]: - hidden_states = hidden_states.to(self.weight.dtype) - - return self.weight * hidden_states - - -class NewGELUActivation(nn.Module): - """ - Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT). Also see - the Gaussian Error Linear Units paper: https://huggingface.co/papers/1606.08415 - """ - - def forward(self, input: torch.Tensor) -> torch.Tensor: - return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (input + 0.044715 * torch.pow(input, 3.0)))) - - -class T5FiLMLayer(nn.Module): - """ - T5 style FiLM Layer. - - Args: - in_features (`int`): - Number of input features. - out_features (`int`): - Number of output features. - """ - - def __init__(self, in_features: int, out_features: int): - super().__init__() - self.scale_bias = nn.Linear(in_features, out_features * 2, bias=False) - - def forward(self, x: torch.Tensor, conditioning_emb: torch.Tensor) -> torch.Tensor: - emb = self.scale_bias(conditioning_emb) - scale, shift = torch.chunk(emb, 2, -1) - x = x * (1 + scale) + shift - return x diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_2d.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_2d.py deleted file mode 100644 index 67fe9a33109b55ab291d2ec73ef9edc345326eda..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_2d.py +++ /dev/null @@ -1,551 +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 Any, Dict, Optional - -import torch -import torch.nn.functional as F -from torch import nn - -from ...configuration_utils import LegacyConfigMixin, register_to_config -from ...utils import deprecate, logging -from ..attention import BasicTransformerBlock -from ..embeddings import ImagePositionalEmbeddings, PatchEmbed, PixArtAlphaTextProjection -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import LegacyModelMixin -from ..normalization import AdaLayerNormSingle - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class Transformer2DModelOutput(Transformer2DModelOutput): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `Transformer2DModelOutput` from `diffusers.models.transformer_2d` is deprecated and this will be removed in a future version. Please use `from diffusers.models.modeling_outputs import Transformer2DModelOutput`, instead." - deprecate("Transformer2DModelOutput", "1.0.0", deprecation_message) - super().__init__(*args, **kwargs) - - -class Transformer2DModel(LegacyModelMixin, LegacyConfigMixin): - """ - A 2D Transformer model for image-like data. - - Parameters: - num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. - attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. - in_channels (`int`, *optional*): - The number of channels in the input and output (specify if the input is **continuous**). - num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use. - sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**). - This is fixed during training since it is used to learn a number of position embeddings. - num_vector_embeds (`int`, *optional*): - The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**). - Includes the class for the masked latent pixel. - activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward. - num_embeds_ada_norm ( `int`, *optional*): - The number of diffusion steps used during training. Pass if at least one of the norm_layers is - `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are - added to the hidden states. - - During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`. - attention_bias (`bool`, *optional*): - Configure if the `TransformerBlocks` attention should contain a bias parameter. - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["BasicTransformerBlock"] - _skip_layerwise_casting_patterns = ["latent_image_embedding", "norm"] - - @register_to_config - def __init__( - self, - num_attention_heads: int = 16, - attention_head_dim: int = 88, - in_channels: Optional[int] = None, - out_channels: Optional[int] = None, - num_layers: int = 1, - dropout: float = 0.0, - norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = None, - attention_bias: bool = False, - sample_size: Optional[int] = None, - num_vector_embeds: Optional[int] = None, - patch_size: Optional[int] = None, - activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, - use_linear_projection: bool = False, - only_cross_attention: bool = False, - double_self_attention: bool = False, - upcast_attention: bool = False, - norm_type: str = "layer_norm", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen' - norm_elementwise_affine: bool = True, - norm_eps: float = 1e-5, - attention_type: str = "default", - caption_channels: int = None, - interpolation_scale: float = None, - use_additional_conditions: Optional[bool] = None, - ): - super().__init__() - - # Validate inputs. - if patch_size is not None: - if norm_type not in ["ada_norm", "ada_norm_zero", "ada_norm_single"]: - raise NotImplementedError( - f"Forward pass is not implemented when `patch_size` is not None and `norm_type` is '{norm_type}'." - ) - elif norm_type in ["ada_norm", "ada_norm_zero"] and num_embeds_ada_norm is None: - raise ValueError( - f"When using a `patch_size` and this `norm_type` ({norm_type}), `num_embeds_ada_norm` cannot be None." - ) - - # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)` - # Define whether input is continuous or discrete depending on configuration - self.is_input_continuous = (in_channels is not None) and (patch_size is None) - self.is_input_vectorized = num_vector_embeds is not None - self.is_input_patches = in_channels is not None and patch_size is not None - - if self.is_input_continuous and self.is_input_vectorized: - raise ValueError( - f"Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make" - " sure that either `in_channels` or `num_vector_embeds` is None." - ) - elif self.is_input_vectorized and self.is_input_patches: - raise ValueError( - f"Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make" - " sure that either `num_vector_embeds` or `num_patches` is None." - ) - elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches: - raise ValueError( - f"Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:" - f" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None." - ) - - if norm_type == "layer_norm" and num_embeds_ada_norm is not None: - deprecation_message = ( - f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or" - " incorrectly set to `'layer_norm'`. Make sure to set `norm_type` to `'ada_norm'` in the config." - " Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect" - " results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it" - " would be very nice if you could open a Pull request for the `transformer/config.json` file" - ) - deprecate("norm_type!=num_embeds_ada_norm", "1.0.0", deprecation_message, standard_warn=False) - norm_type = "ada_norm" - - # Set some common variables used across the board. - self.use_linear_projection = use_linear_projection - self.interpolation_scale = interpolation_scale - self.caption_channels = caption_channels - self.num_attention_heads = num_attention_heads - self.attention_head_dim = attention_head_dim - self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim - self.in_channels = in_channels - self.out_channels = in_channels if out_channels is None else out_channels - self.gradient_checkpointing = False - - if use_additional_conditions is None: - if norm_type == "ada_norm_single" and sample_size == 128: - use_additional_conditions = True - else: - use_additional_conditions = False - self.use_additional_conditions = use_additional_conditions - - # 2. Initialize the right blocks. - # These functions follow a common structure: - # a. Initialize the input blocks. b. Initialize the transformer blocks. - # c. Initialize the output blocks and other projection blocks when necessary. - if self.is_input_continuous: - self._init_continuous_input(norm_type=norm_type) - elif self.is_input_vectorized: - self._init_vectorized_inputs(norm_type=norm_type) - elif self.is_input_patches: - self._init_patched_inputs(norm_type=norm_type) - - def _init_continuous_input(self, norm_type): - self.norm = torch.nn.GroupNorm( - num_groups=self.config.norm_num_groups, num_channels=self.in_channels, eps=1e-6, affine=True - ) - if self.use_linear_projection: - self.proj_in = torch.nn.Linear(self.in_channels, self.inner_dim) - else: - self.proj_in = torch.nn.Conv2d(self.in_channels, self.inner_dim, kernel_size=1, stride=1, padding=0) - - self.transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - self.inner_dim, - self.config.num_attention_heads, - self.config.attention_head_dim, - dropout=self.config.dropout, - cross_attention_dim=self.config.cross_attention_dim, - activation_fn=self.config.activation_fn, - num_embeds_ada_norm=self.config.num_embeds_ada_norm, - attention_bias=self.config.attention_bias, - only_cross_attention=self.config.only_cross_attention, - double_self_attention=self.config.double_self_attention, - upcast_attention=self.config.upcast_attention, - norm_type=norm_type, - norm_elementwise_affine=self.config.norm_elementwise_affine, - norm_eps=self.config.norm_eps, - attention_type=self.config.attention_type, - ) - for _ in range(self.config.num_layers) - ] - ) - - if self.use_linear_projection: - self.proj_out = torch.nn.Linear(self.inner_dim, self.out_channels) - else: - self.proj_out = torch.nn.Conv2d(self.inner_dim, self.out_channels, kernel_size=1, stride=1, padding=0) - - def _init_vectorized_inputs(self, norm_type): - assert self.config.sample_size is not None, "Transformer2DModel over discrete input must provide sample_size" - assert self.config.num_vector_embeds is not None, ( - "Transformer2DModel over discrete input must provide num_embed" - ) - - self.height = self.config.sample_size - self.width = self.config.sample_size - self.num_latent_pixels = self.height * self.width - - self.latent_image_embedding = ImagePositionalEmbeddings( - num_embed=self.config.num_vector_embeds, embed_dim=self.inner_dim, height=self.height, width=self.width - ) - - self.transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - self.inner_dim, - self.config.num_attention_heads, - self.config.attention_head_dim, - dropout=self.config.dropout, - cross_attention_dim=self.config.cross_attention_dim, - activation_fn=self.config.activation_fn, - num_embeds_ada_norm=self.config.num_embeds_ada_norm, - attention_bias=self.config.attention_bias, - only_cross_attention=self.config.only_cross_attention, - double_self_attention=self.config.double_self_attention, - upcast_attention=self.config.upcast_attention, - norm_type=norm_type, - norm_elementwise_affine=self.config.norm_elementwise_affine, - norm_eps=self.config.norm_eps, - attention_type=self.config.attention_type, - ) - for _ in range(self.config.num_layers) - ] - ) - - self.norm_out = nn.LayerNorm(self.inner_dim) - self.out = nn.Linear(self.inner_dim, self.config.num_vector_embeds - 1) - - def _init_patched_inputs(self, norm_type): - assert self.config.sample_size is not None, "Transformer2DModel over patched input must provide sample_size" - - self.height = self.config.sample_size - self.width = self.config.sample_size - - self.patch_size = self.config.patch_size - interpolation_scale = ( - self.config.interpolation_scale - if self.config.interpolation_scale is not None - else max(self.config.sample_size // 64, 1) - ) - self.pos_embed = PatchEmbed( - height=self.config.sample_size, - width=self.config.sample_size, - patch_size=self.config.patch_size, - in_channels=self.in_channels, - embed_dim=self.inner_dim, - interpolation_scale=interpolation_scale, - ) - - self.transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - self.inner_dim, - self.config.num_attention_heads, - self.config.attention_head_dim, - dropout=self.config.dropout, - cross_attention_dim=self.config.cross_attention_dim, - activation_fn=self.config.activation_fn, - num_embeds_ada_norm=self.config.num_embeds_ada_norm, - attention_bias=self.config.attention_bias, - only_cross_attention=self.config.only_cross_attention, - double_self_attention=self.config.double_self_attention, - upcast_attention=self.config.upcast_attention, - norm_type=norm_type, - norm_elementwise_affine=self.config.norm_elementwise_affine, - norm_eps=self.config.norm_eps, - attention_type=self.config.attention_type, - ) - for _ in range(self.config.num_layers) - ] - ) - - if self.config.norm_type != "ada_norm_single": - self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6) - self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim) - self.proj_out_2 = nn.Linear( - self.inner_dim, self.config.patch_size * self.config.patch_size * self.out_channels - ) - elif self.config.norm_type == "ada_norm_single": - self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6) - self.scale_shift_table = nn.Parameter(torch.randn(2, self.inner_dim) / self.inner_dim**0.5) - self.proj_out = nn.Linear( - self.inner_dim, self.config.patch_size * self.config.patch_size * self.out_channels - ) - - # PixArt-Alpha blocks. - self.adaln_single = None - if self.config.norm_type == "ada_norm_single": - # TODO(Sayak, PVP) clean this, for now we use sample size to determine whether to use - # additional conditions until we find better name - self.adaln_single = AdaLayerNormSingle( - self.inner_dim, use_additional_conditions=self.use_additional_conditions - ) - - self.caption_projection = None - if self.caption_channels is not None: - self.caption_projection = PixArtAlphaTextProjection( - in_features=self.caption_channels, hidden_size=self.inner_dim - ) - - 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, - return_dict: bool = True, - ): - """ - The [`Transformer2DModel`] forward method. - - Args: - hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.Tensor` of shape `(batch size, channel, height, width)` if continuous): - Input `hidden_states`. - encoder_hidden_states ( `torch.Tensor` of shape `(batch size, sequence len, embed dims)`, *optional*): - Conditional embeddings for cross attention layer. If not given, cross-attention defaults to - self-attention. - timestep ( `torch.LongTensor`, *optional*): - Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`. - 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*): - 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). - attention_mask ( `torch.Tensor`, *optional*): - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask - is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large - negative values to the attention scores corresponding to "discard" tokens. - encoder_attention_mask ( `torch.Tensor`, *optional*): - Cross-attention mask applied to `encoder_hidden_states`. Two formats supported: - - * Mask `(batch, sequence_length)` True = keep, False = discard. - * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard. - - If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format - above. This bias will be added to the cross-attention scores. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain - tuple. - - Returns: - If `return_dict` is True, an [`~models.transformers.transformer_2d.Transformer2DModelOutput`] is returned, - otherwise a `tuple` where the first element is the sample 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.") - # 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. - # expects mask of shape: - # [batch, key_tokens] - # adds singleton query_tokens dimension: - # [batch, 1, key_tokens] - # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: - # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) - # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) - if attention_mask is not None and attention_mask.ndim == 2: - # assume that mask is expressed as: - # (1 = keep, 0 = discard) - # convert mask into a bias that can be added to attention scores: - # (keep = +0, discard = -10000.0) - attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0 - attention_mask = attention_mask.unsqueeze(1) - - # 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) - - # 1. Input - if self.is_input_continuous: - batch_size, _, height, width = hidden_states.shape - residual = hidden_states - hidden_states, inner_dim = self._operate_on_continuous_inputs(hidden_states) - elif self.is_input_vectorized: - hidden_states = self.latent_image_embedding(hidden_states) - elif self.is_input_patches: - height, width = hidden_states.shape[-2] // self.patch_size, hidden_states.shape[-1] // self.patch_size - hidden_states, encoder_hidden_states, timestep, embedded_timestep = self._operate_on_patched_inputs( - hidden_states, encoder_hidden_states, timestep, added_cond_kwargs - ) - - # 2. 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, - encoder_hidden_states, - encoder_attention_mask, - timestep, - cross_attention_kwargs, - class_labels, - ) - else: - hidden_states = block( - hidden_states, - attention_mask=attention_mask, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - timestep=timestep, - cross_attention_kwargs=cross_attention_kwargs, - class_labels=class_labels, - ) - - # 3. Output - if self.is_input_continuous: - output = self._get_output_for_continuous_inputs( - hidden_states=hidden_states, - residual=residual, - batch_size=batch_size, - height=height, - width=width, - inner_dim=inner_dim, - ) - elif self.is_input_vectorized: - output = self._get_output_for_vectorized_inputs(hidden_states) - elif self.is_input_patches: - output = self._get_output_for_patched_inputs( - hidden_states=hidden_states, - timestep=timestep, - class_labels=class_labels, - embedded_timestep=embedded_timestep, - height=height, - width=width, - ) - - if not return_dict: - return (output,) - - return Transformer2DModelOutput(sample=output) - - def _operate_on_continuous_inputs(self, hidden_states): - batch, _, height, width = hidden_states.shape - hidden_states = self.norm(hidden_states) - - if not self.use_linear_projection: - hidden_states = self.proj_in(hidden_states) - inner_dim = hidden_states.shape[1] - hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim) - else: - inner_dim = hidden_states.shape[1] - hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim) - hidden_states = self.proj_in(hidden_states) - - return hidden_states, inner_dim - - def _operate_on_patched_inputs(self, hidden_states, encoder_hidden_states, timestep, added_cond_kwargs): - batch_size = hidden_states.shape[0] - hidden_states = self.pos_embed(hidden_states) - embedded_timestep = None - - if self.adaln_single is not None: - if self.use_additional_conditions and added_cond_kwargs is None: - raise ValueError( - "`added_cond_kwargs` cannot be None when using additional conditions for `adaln_single`." - ) - timestep, embedded_timestep = self.adaln_single( - timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_states.dtype - ) - - if self.caption_projection is not None: - encoder_hidden_states = self.caption_projection(encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1]) - - return hidden_states, encoder_hidden_states, timestep, embedded_timestep - - def _get_output_for_continuous_inputs(self, hidden_states, residual, batch_size, height, width, inner_dim): - if not self.use_linear_projection: - hidden_states = ( - hidden_states.reshape(batch_size, height, width, inner_dim).permute(0, 3, 1, 2).contiguous() - ) - hidden_states = self.proj_out(hidden_states) - else: - hidden_states = self.proj_out(hidden_states) - hidden_states = ( - hidden_states.reshape(batch_size, height, width, inner_dim).permute(0, 3, 1, 2).contiguous() - ) - - output = hidden_states + residual - return output - - def _get_output_for_vectorized_inputs(self, hidden_states): - hidden_states = self.norm_out(hidden_states) - logits = self.out(hidden_states) - # (batch, self.num_vector_embeds - 1, self.num_latent_pixels) - logits = logits.permute(0, 2, 1) - # log(p(x_0)) - output = F.log_softmax(logits.double(), dim=1).float() - return output - - def _get_output_for_patched_inputs( - self, hidden_states, timestep, class_labels, embedded_timestep, height=None, width=None - ): - if self.config.norm_type != "ada_norm_single": - conditioning = self.transformer_blocks[0].norm1.emb( - timestep, class_labels, hidden_dtype=hidden_states.dtype - ) - shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) - hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] - hidden_states = self.proj_out_2(hidden_states) - elif self.config.norm_type == "ada_norm_single": - shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1) - hidden_states = self.norm_out(hidden_states) - # Modulation - hidden_states = hidden_states * (1 + scale) + shift - hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.squeeze(1) - - # unpatchify - if self.adaln_single is None: - height = width = int(hidden_states.shape[1] ** 0.5) - hidden_states = hidden_states.reshape( - shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels) - ) - hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) - output = hidden_states.reshape( - shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size) - ) - return output diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_allegro.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_allegro.py deleted file mode 100644 index 5fa59a71d977e5f0516ae314c0da332d789ff492..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_allegro.py +++ /dev/null @@ -1,414 +0,0 @@ -# Copyright 2025 The RhymesAI 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 Optional, Tuple - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import logging -from ...utils.torch_utils import maybe_allow_in_graph -from ..attention import FeedForward -from ..attention_processor import AllegroAttnProcessor2_0, Attention -from ..cache_utils import CacheMixin -from ..embeddings import PatchEmbed, PixArtAlphaTextProjection -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormSingle - - -logger = logging.get_logger(__name__) - - -@maybe_allow_in_graph -class AllegroTransformerBlock(nn.Module): - r""" - Transformer block used in [Allegro](https://github.com/rhymes-ai/Allegro) model. - - 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. - dropout (`float`, defaults to `0.0`): - The dropout probability to use. - cross_attention_dim (`int`, defaults to `2304`): - The dimension of the cross attention features. - activation_fn (`str`, defaults to `"gelu-approximate"`): - Activation function to be used in feed-forward. - attention_bias (`bool`, defaults to `False`): - Whether or not to use bias in attention projection layers. - only_cross_attention (`bool`, defaults to `False`): - norm_elementwise_affine (`bool`, defaults to `True`): - Whether to use learnable elementwise affine parameters for normalization. - norm_eps (`float`, defaults to `1e-5`): - Epsilon value for normalization layers. - final_dropout (`bool` defaults to `False`): - Whether to apply a final dropout after the last feed-forward layer. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - dropout=0.0, - cross_attention_dim: Optional[int] = None, - activation_fn: str = "geglu", - attention_bias: bool = False, - norm_elementwise_affine: bool = True, - norm_eps: float = 1e-5, - ): - super().__init__() - - # 1. Self Attention - self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) - - self.attn1 = Attention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - cross_attention_dim=None, - processor=AllegroAttnProcessor2_0(), - ) - - # 2. Cross Attention - self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) - self.attn2 = Attention( - query_dim=dim, - cross_attention_dim=cross_attention_dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - processor=AllegroAttnProcessor2_0(), - ) - - # 3. Feed Forward - self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) - - self.ff = FeedForward( - dim, - dropout=dropout, - activation_fn=activation_fn, - ) - - # 4. Scale-shift - self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) - - 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, - image_rotary_emb=None, - ) -> torch.Tensor: - # 0. Self-Attention - batch_size = hidden_states.shape[0] - - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( - self.scale_shift_table[None] + temb.reshape(batch_size, 6, -1) - ).chunk(6, dim=1) - norm_hidden_states = self.norm1(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa - norm_hidden_states = norm_hidden_states.squeeze(1) - - attn_output = self.attn1( - norm_hidden_states, - encoder_hidden_states=None, - attention_mask=attention_mask, - image_rotary_emb=image_rotary_emb, - ) - attn_output = gate_msa * attn_output - - hidden_states = attn_output + hidden_states - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - # 1. Cross-Attention - if self.attn2 is not None: - norm_hidden_states = hidden_states - - attn_output = self.attn2( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - image_rotary_emb=None, - ) - hidden_states = attn_output + hidden_states - - # 2. Feed-forward - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp - - ff_output = self.ff(norm_hidden_states) - ff_output = gate_mlp * ff_output - - hidden_states = ff_output + hidden_states - - # TODO(aryan): maybe following line is not required - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - return hidden_states - - -class AllegroTransformer3DModel(ModelMixin, ConfigMixin, CacheMixin): - _supports_gradient_checkpointing = True - - """ - A 3D Transformer model for video-like data. - - Args: - patch_size (`int`, defaults to `2`): - The size of spatial patches to use in the patch embedding layer. - patch_size_t (`int`, defaults to `1`): - The size of temporal patches to use in the patch embedding layer. - num_attention_heads (`int`, defaults to `24`): - The number of heads to use for multi-head attention. - attention_head_dim (`int`, defaults to `96`): - The number of channels in each head. - in_channels (`int`, defaults to `4`): - The number of channels in the input. - out_channels (`int`, *optional*, defaults to `4`): - The number of channels in the output. - num_layers (`int`, defaults to `32`): - The number of layers of Transformer blocks to use. - dropout (`float`, defaults to `0.0`): - The dropout probability to use. - cross_attention_dim (`int`, defaults to `2304`): - The dimension of the cross attention features. - attention_bias (`bool`, defaults to `True`): - Whether or not to use bias in the attention projection layers. - sample_height (`int`, defaults to `90`): - The height of the input latents. - sample_width (`int`, defaults to `160`): - The width of the input latents. - sample_frames (`int`, defaults to `22`): - The number of frames in the input latents. - activation_fn (`str`, defaults to `"gelu-approximate"`): - Activation function to use in feed-forward. - norm_elementwise_affine (`bool`, defaults to `False`): - Whether or not to use elementwise affine in normalization layers. - norm_eps (`float`, defaults to `1e-6`): - The epsilon value to use in normalization layers. - caption_channels (`int`, defaults to `4096`): - Number of channels to use for projecting the caption embeddings. - interpolation_scale_h (`float`, defaults to `2.0`): - Scaling factor to apply in 3D positional embeddings across height dimension. - interpolation_scale_w (`float`, defaults to `2.0`): - Scaling factor to apply in 3D positional embeddings across width dimension. - interpolation_scale_t (`float`, defaults to `2.2`): - Scaling factor to apply in 3D positional embeddings across time dimension. - """ - - _supports_gradient_checkpointing = True - _skip_layerwise_casting_patterns = ["pos_embed", "norm", "adaln_single"] - - @register_to_config - def __init__( - self, - patch_size: int = 2, - patch_size_t: int = 1, - num_attention_heads: int = 24, - attention_head_dim: int = 96, - in_channels: int = 4, - out_channels: int = 4, - num_layers: int = 32, - dropout: float = 0.0, - cross_attention_dim: int = 2304, - attention_bias: bool = True, - sample_height: int = 90, - sample_width: int = 160, - sample_frames: int = 22, - activation_fn: str = "gelu-approximate", - norm_elementwise_affine: bool = False, - norm_eps: float = 1e-6, - caption_channels: int = 4096, - interpolation_scale_h: float = 2.0, - interpolation_scale_w: float = 2.0, - interpolation_scale_t: float = 2.2, - ): - super().__init__() - - self.inner_dim = num_attention_heads * attention_head_dim - - interpolation_scale_t = ( - interpolation_scale_t - if interpolation_scale_t is not None - else ((sample_frames - 1) // 16 + 1) - if sample_frames % 2 == 1 - else sample_frames // 16 - ) - interpolation_scale_h = interpolation_scale_h if interpolation_scale_h is not None else sample_height / 30 - interpolation_scale_w = interpolation_scale_w if interpolation_scale_w is not None else sample_width / 40 - - # 1. Patch embedding - self.pos_embed = PatchEmbed( - height=sample_height, - width=sample_width, - patch_size=patch_size, - in_channels=in_channels, - embed_dim=self.inner_dim, - pos_embed_type=None, - ) - - # 2. Transformer blocks - self.transformer_blocks = nn.ModuleList( - [ - AllegroTransformerBlock( - self.inner_dim, - num_attention_heads, - attention_head_dim, - dropout=dropout, - cross_attention_dim=cross_attention_dim, - activation_fn=activation_fn, - attention_bias=attention_bias, - norm_elementwise_affine=norm_elementwise_affine, - norm_eps=norm_eps, - ) - for _ in range(num_layers) - ] - ) - - # 3. Output projection & norm - self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6) - self.scale_shift_table = nn.Parameter(torch.randn(2, self.inner_dim) / self.inner_dim**0.5) - self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * out_channels) - - # 4. Timestep embeddings - self.adaln_single = AdaLayerNormSingle(self.inner_dim, use_additional_conditions=False) - - # 5. Caption projection - self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=self.inner_dim) - - self.gradient_checkpointing = False - - def forward( - self, - 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, - return_dict: bool = True, - ): - batch_size, num_channels, num_frames, height, width = hidden_states.shape - p_t = self.config.patch_size_t - p = self.config.patch_size - - post_patch_num_frames = num_frames // p_t - post_patch_height = height // p - post_patch_width = width // p - - # 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. - # expects mask of shape: - # [batch, key_tokens] - # adds singleton query_tokens dimension: - # [batch, 1, key_tokens] - # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: - # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) - # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) attention_mask_vid, attention_mask_img = None, None - if attention_mask is not None and attention_mask.ndim == 4: - # assume that mask is expressed as: - # (1 = keep, 0 = discard) - # convert mask into a bias that can be added to attention scores: - # (keep = +0, discard = -10000.0) - # b, frame+use_image_num, h, w -> a video with images - # b, 1, h, w -> only images - attention_mask = attention_mask.to(hidden_states.dtype) - attention_mask = attention_mask[:, :num_frames] # [batch_size, num_frames, height, width] - - if attention_mask.numel() > 0: - attention_mask = attention_mask.unsqueeze(1) # [batch_size, 1, num_frames, height, width] - attention_mask = F.max_pool3d(attention_mask, kernel_size=(p_t, p, p), stride=(p_t, p, p)) - attention_mask = attention_mask.flatten(1).view(batch_size, 1, -1) - - attention_mask = ( - (1 - attention_mask.bool().to(hidden_states.dtype)) * -10000.0 if attention_mask.numel() > 0 else None - ) - - # 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(self.dtype)) * -10000.0 - encoder_attention_mask = encoder_attention_mask.unsqueeze(1) - - # 1. Timestep embeddings - timestep, embedded_timestep = self.adaln_single( - timestep, batch_size=batch_size, hidden_dtype=hidden_states.dtype - ) - - # 2. Patch embeddings - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) - hidden_states = self.pos_embed(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).flatten(1, 2) - - encoder_hidden_states = self.caption_projection(encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, encoder_hidden_states.shape[-1]) - - # 3. Transformer blocks - for i, block in enumerate(self.transformer_blocks): - # TODO(aryan): Implement gradient checkpointing - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - block, - hidden_states, - encoder_hidden_states, - timestep, - attention_mask, - encoder_attention_mask, - image_rotary_emb, - ) - else: - hidden_states = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - temb=timestep, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - image_rotary_emb=image_rotary_emb, - ) - - # 4. Output normalization & projection - shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1) - hidden_states = self.norm_out(hidden_states) - - # Modulation - hidden_states = hidden_states * (1 + scale) + shift - hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.squeeze(1) - - # 5. Unpatchify - hidden_states = hidden_states.reshape( - batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p, p, -1 - ) - hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) - output = hidden_states.reshape(batch_size, -1, num_frames, height, width) - - if not return_dict: - return (output,) - - return Transformer2DModelOutput(sample=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_chroma.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_chroma.py deleted file mode 100644 index d11f6c2a5e25780015bedc48f8b5c5cc2c4703cc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_chroma.py +++ /dev/null @@ -1,742 +0,0 @@ -# Copyright 2025 Black Forest Labs, The HuggingFace Team and loadstone-rock . 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, Dict, Optional, Tuple, Union - -import numpy as np -import torch -import torch.nn as nn - -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_utils import is_torch_npu_available -from ...utils.torch_utils import maybe_allow_in_graph -from ..attention import FeedForward -from ..attention_processor import ( - Attention, - AttentionProcessor, - FluxAttnProcessor2_0, - FluxAttnProcessor2_0_NPU, - FusedFluxAttnProcessor2_0, -) -from ..cache_utils import CacheMixin -from ..embeddings import FluxPosEmbed, PixArtAlphaTextProjection, Timesteps, get_timestep_embedding -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import CombinedTimestepLabelEmbeddings, FP32LayerNorm, RMSNorm - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class ChromaAdaLayerNormZeroPruned(nn.Module): - r""" - Norm layer adaptive layer norm zero (adaLN-Zero). - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - 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): - super().__init__() - if num_embeddings is not None: - self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, embedding_dim) - else: - self.emb = None - - if norm_type == "layer_norm": - self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) - elif norm_type == "fp32_layer_norm": - self.norm = FP32LayerNorm(embedding_dim, elementwise_affine=False, bias=False) - else: - raise ValueError( - f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'." - ) - - 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]: - 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) - x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] - return x, gate_msa, shift_mlp, scale_mlp, gate_mlp - - -class ChromaAdaLayerNormZeroSinglePruned(nn.Module): - r""" - Norm layer adaptive layer norm zero (adaLN-Zero). - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - num_embeddings (`int`): The size of the embeddings dictionary. - """ - - def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True): - super().__init__() - - if norm_type == "layer_norm": - self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) - else: - raise ValueError( - f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'." - ) - - def forward( - self, - x: torch.Tensor, - emb: Optional[torch.Tensor] = 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 - - -class ChromaAdaLayerNormContinuousPruned(nn.Module): - r""" - Adaptive normalization layer with a norm layer (layer_norm or rms_norm). - - Args: - embedding_dim (`int`): Embedding dimension to use during projection. - conditioning_embedding_dim (`int`): Dimension of the input condition. - elementwise_affine (`bool`, defaults to `True`): - Boolean flag to denote if affine transformation should be applied. - eps (`float`, defaults to 1e-5): Epsilon factor. - bias (`bias`, defaults to `True`): Boolean flag to denote if bias should be use. - norm_type (`str`, defaults to `"layer_norm"`): - Normalization layer to use. Values supported: "layer_norm", "rms_norm". - """ - - def __init__( - self, - embedding_dim: int, - conditioning_embedding_dim: int, - # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters - # because the output is immediately scaled and shifted by the projected conditioning embeddings. - # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters. - # However, this is how it was implemented in the original code, and it's rather likely you should - # set `elementwise_affine` to False. - elementwise_affine=True, - eps=1e-5, - bias=True, - norm_type="layer_norm", - ): - super().__init__() - if norm_type == "layer_norm": - self.norm = nn.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, emb: torch.Tensor) -> torch.Tensor: - # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT) - shift, scale = torch.chunk(emb.flatten(1, 2).to(x.dtype), 2, dim=1) - x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] - return x - - -class ChromaCombinedTimestepTextProjEmbeddings(nn.Module): - def __init__(self, num_channels: int, out_dim: int): - super().__init__() - - self.time_proj = Timesteps(num_channels=num_channels, flip_sin_to_cos=True, downscale_freq_shift=0) - self.guidance_proj = Timesteps(num_channels=num_channels, flip_sin_to_cos=True, downscale_freq_shift=0) - - self.register_buffer( - "mod_proj", - get_timestep_embedding( - torch.arange(out_dim) * 1000, 2 * num_channels, flip_sin_to_cos=True, downscale_freq_shift=0 - ), - persistent=False, - ) - - def forward(self, timestep: torch.Tensor) -> torch.Tensor: - mod_index_length = self.mod_proj.shape[0] - batch_size = timestep.shape[0] - - timesteps_proj = self.time_proj(timestep).to(dtype=timestep.dtype) - guidance_proj = self.guidance_proj(torch.tensor([0] * batch_size)).to( - dtype=timestep.dtype, device=timestep.device - ) - - mod_proj = self.mod_proj.to(dtype=timesteps_proj.dtype, device=timesteps_proj.device).repeat(batch_size, 1, 1) - timestep_guidance = ( - torch.cat([timesteps_proj, guidance_proj], dim=1).unsqueeze(1).repeat(1, mod_index_length, 1) - ) - input_vec = torch.cat([timestep_guidance, mod_proj], dim=-1) - return input_vec.to(timestep.dtype) - - -class ChromaApproximator(nn.Module): - def __init__(self, in_dim: int, out_dim: int, hidden_dim: int, n_layers: int = 5): - super().__init__() - self.in_proj = nn.Linear(in_dim, hidden_dim, bias=True) - self.layers = nn.ModuleList( - [PixArtAlphaTextProjection(hidden_dim, hidden_dim, act_fn="silu") for _ in range(n_layers)] - ) - self.norms = nn.ModuleList([nn.RMSNorm(hidden_dim) for _ in range(n_layers)]) - self.out_proj = nn.Linear(hidden_dim, out_dim) - - def forward(self, x): - x = self.in_proj(x) - - for layer, norms in zip(self.layers, self.norms): - x = x + layer(norms(x)) - - return self.out_proj(x) - - -@maybe_allow_in_graph -class ChromaSingleTransformerBlock(nn.Module): - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - mlp_ratio: float = 4.0, - ): - super().__init__() - self.mlp_hidden_dim = int(dim * mlp_ratio) - self.norm = ChromaAdaLayerNormZeroSinglePruned(dim) - self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim) - self.act_mlp = nn.GELU(approximate="tanh") - self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim) - - if is_torch_npu_available(): - deprecation_message = ( - "Defaulting to FluxAttnProcessor2_0_NPU for NPU devices will be removed. Attention processors " - "should be set explicitly using the `set_attn_processor` method." - ) - deprecate("npu_processor", "0.34.0", deprecation_message) - processor = FluxAttnProcessor2_0_NPU() - else: - processor = FluxAttnProcessor2_0() - - self.attn = Attention( - query_dim=dim, - cross_attention_dim=None, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=dim, - bias=True, - processor=processor, - qk_norm="rms_norm", - eps=1e-6, - pre_only=True, - ) - - 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, - ) -> torch.Tensor: - residual = hidden_states - norm_hidden_states, gate = self.norm(hidden_states, emb=temb) - mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)) - joint_attention_kwargs = joint_attention_kwargs or {} - - if attention_mask is not None: - attention_mask = attention_mask[:, None, None, :] * attention_mask[:, None, :, None] - - attn_output = self.attn( - hidden_states=norm_hidden_states, - image_rotary_emb=image_rotary_emb, - attention_mask=attention_mask, - **joint_attention_kwargs, - ) - - hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) - gate = gate.unsqueeze(1) - hidden_states = gate * self.proj_out(hidden_states) - hidden_states = residual + hidden_states - if hidden_states.dtype == torch.float16: - hidden_states = hidden_states.clip(-65504, 65504) - - return hidden_states - - -@maybe_allow_in_graph -class ChromaTransformerBlock(nn.Module): - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - qk_norm: str = "rms_norm", - eps: float = 1e-6, - ): - super().__init__() - self.norm1 = ChromaAdaLayerNormZeroPruned(dim) - self.norm1_context = ChromaAdaLayerNormZeroPruned(dim) - - self.attn = Attention( - query_dim=dim, - cross_attention_dim=None, - added_kv_proj_dim=dim, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=dim, - context_pre_only=False, - bias=True, - processor=FluxAttnProcessor2_0(), - qk_norm=qk_norm, - eps=eps, - ) - - self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) - self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") - - self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) - self.ff_context = 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, - 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]: - 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) - - norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( - encoder_hidden_states, emb=temb_txt - ) - joint_attention_kwargs = joint_attention_kwargs or {} - if attention_mask is not None: - attention_mask = attention_mask[:, None, None, :] * attention_mask[:, None, :, None] - - # Attention. - attention_outputs = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - image_rotary_emb=image_rotary_emb, - attention_mask=attention_mask, - **joint_attention_kwargs, - ) - - if len(attention_outputs) == 2: - attn_output, context_attn_output = attention_outputs - elif len(attention_outputs) == 3: - attn_output, context_attn_output, ip_attn_output = attention_outputs - - # Process attention outputs for the `hidden_states`. - attn_output = gate_msa.unsqueeze(1) * attn_output - hidden_states = hidden_states + attn_output - - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - - ff_output = self.ff(norm_hidden_states) - ff_output = gate_mlp.unsqueeze(1) * ff_output - - hidden_states = hidden_states + ff_output - if len(attention_outputs) == 3: - hidden_states = hidden_states + ip_attn_output - - # Process attention outputs for the `encoder_hidden_states`. - - context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output - encoder_hidden_states = encoder_hidden_states + context_attn_output - - norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) - norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] - - context_ff_output = self.ff_context(norm_encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output - if encoder_hidden_states.dtype == torch.float16: - encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) - - return encoder_hidden_states, hidden_states - - -class ChromaTransformer2DModel( - ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, FluxTransformer2DLoadersMixin, CacheMixin -): - """ - The Transformer model introduced in Flux, modified for Chroma. - - Reference: https://huggingface.co/lodestones/Chroma - - Args: - patch_size (`int`, defaults to `1`): - Patch size to turn the input data into small patches. - in_channels (`int`, defaults to `64`): - 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. - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["ChromaTransformerBlock", "ChromaSingleTransformerBlock"] - _skip_layerwise_casting_patterns = ["pos_embed", "norm"] - - @register_to_config - def __init__( - self, - patch_size: int = 1, - in_channels: int = 64, - 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, - ): - 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.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim) - self.x_embedder = nn.Linear(in_channels, 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.norm_out = ChromaAdaLayerNormContinuousPruned( - self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6 - ) - self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True) - - self.gradient_checkpointing = False - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedFluxAttnProcessor2_0 - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - self.set_attn_processor(FusedFluxAttnProcessor2_0()) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - if self.original_attn_processors is not None: - self.set_attn_processor(self.original_attn_processors) - - 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: Optional[Dict[str, Any]] = None, - controlnet_block_samples=None, - controlnet_single_block_samples=None, - return_dict: bool = True, - controlnet_blocks_repeat: bool = False, - ) -> Union[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. - """ - 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) * 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 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/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_cogview3plus.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_cogview3plus.py deleted file mode 100644 index 77f15f6ca6f1ff51ce7ad60c9c1451cd3965cb91..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_cogview3plus.py +++ /dev/null @@ -1,370 +0,0 @@ -# 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 Dict, Union - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import logging -from ..attention import FeedForward -from ..attention_processor import Attention, AttentionProcessor, CogVideoXAttnProcessor2_0 -from ..embeddings import CogView3CombinedTimestepSizeEmbeddings, CogView3PlusPatchEmbed -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormContinuous, CogView3PlusAdaLayerNormZeroTextImage - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class CogView3PlusTransformerBlock(nn.Module): - r""" - Transformer block used in [CogView](https://github.com/THUDM/CogView3) model. - - 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. - time_embed_dim (`int`): - The number of channels in timestep embedding. - """ - - def __init__( - self, - dim: int = 2560, - num_attention_heads: int = 64, - attention_head_dim: int = 40, - time_embed_dim: int = 512, - ): - super().__init__() - - self.norm1 = CogView3PlusAdaLayerNormZeroTextImage(embedding_dim=time_embed_dim, 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-6, - processor=CogVideoXAttnProcessor2_0(), - ) - - 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, - emb: torch.Tensor, - ) -> torch.Tensor: - text_seq_length = encoder_hidden_states.size(1) - - # norm & modulate - ( - 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, emb) - - # attention - attn_hidden_states, attn_encoder_hidden_states = self.attn1( - hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states - ) - - hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_hidden_states - encoder_hidden_states = encoder_hidden_states + c_gate_msa.unsqueeze(1) * attn_encoder_hidden_states - - # norm & modulate - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - - norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) - norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] - - # feed-forward - norm_hidden_states = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1) - ff_output = self.ff(norm_hidden_states) - - hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output[:, text_seq_length:] - encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * ff_output[:, :text_seq_length] - - if hidden_states.dtype == torch.float16: - hidden_states = hidden_states.clip(-65504, 65504) - if encoder_hidden_states.dtype == torch.float16: - encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) - return hidden_states, encoder_hidden_states - - -class CogView3PlusTransformer2DModel(ModelMixin, ConfigMixin): - r""" - The Transformer model introduced in [CogView3: Finer and Faster Text-to-Image Generation via Relay - Diffusion](https://huggingface.co/papers/2403.05121). - - 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 `4096`): - 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 - _skip_layerwise_casting_patterns = ["patch_embed", "norm"] - _no_split_modules = ["CogView3PlusTransformerBlock", "CogView3PlusPatchEmbed"] - - @register_to_config - def __init__( - self, - patch_size: int = 2, - in_channels: int = 16, - num_layers: int = 30, - attention_head_dim: int = 40, - num_attention_heads: int = 64, - out_channels: int = 16, - text_embed_dim: int = 4096, - time_embed_dim: int = 512, - condition_dim: int = 256, - pos_embed_max_size: int = 128, - sample_size: int = 128, - ): - super().__init__() - self.out_channels = out_channels - self.inner_dim = num_attention_heads * attention_head_dim - - # CogView3 uses 3 additional SDXL-like conditions - original_size, target_size, crop_coords - # Each of these are sincos embeddings of shape 2 * condition_dim - self.pooled_projection_dim = 3 * 2 * condition_dim - - self.patch_embed = CogView3PlusPatchEmbed( - in_channels=in_channels, - hidden_size=self.inner_dim, - patch_size=patch_size, - text_hidden_size=text_embed_dim, - pos_embed_max_size=pos_embed_max_size, - ) - - self.time_condition_embed = CogView3CombinedTimestepSizeEmbeddings( - embedding_dim=time_embed_dim, - condition_dim=condition_dim, - pooled_projection_dim=self.pooled_projection_dim, - timesteps_dim=self.inner_dim, - ) - - self.transformer_blocks = nn.ModuleList( - [ - CogView3PlusTransformerBlock( - dim=self.inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - time_embed_dim=time_embed_dim, - ) - for _ in range(num_layers) - ] - ) - - self.norm_out = AdaLayerNormContinuous( - embedding_dim=self.inner_dim, - conditioning_embedding_dim=time_embed_dim, - elementwise_affine=False, - eps=1e-6, - ) - self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True) - - self.gradient_checkpointing = False - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - timestep: torch.LongTensor, - original_size: torch.Tensor, - target_size: torch.Tensor, - crop_coords: torch.Tensor, - return_dict: bool = True, - ) -> Union[torch.Tensor, Transformer2DModelOutput]: - """ - The [`CogView3PlusTransformer2DModel`] forward method. - - Args: - hidden_states (`torch.Tensor`): - Input `hidden_states` of shape `(batch size, channel, height, width)`. - encoder_hidden_states (`torch.Tensor`): - Conditional embeddings (embeddings computed from the input conditions such as prompts) of shape - `(batch_size, sequence_len, text_embed_dim)` - timestep (`torch.LongTensor`): - Used to indicate denoising step. - original_size (`torch.Tensor`): - CogView3 uses SDXL-like micro-conditioning for original image size as explained in section 2.2 of - [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`torch.Tensor`): - CogView3 uses SDXL-like micro-conditioning for target image size as explained in section 2.2 of - [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crop_coords (`torch.Tensor`): - CogView3 uses SDXL-like micro-conditioning for crop coordinates as explained in section 2.2 of - [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain - tuple. - - Returns: - `torch.Tensor` or [`~models.transformer_2d.Transformer2DModelOutput`]: - The denoised latents using provided inputs as conditioning. - """ - height, width = hidden_states.shape[-2:] - text_seq_length = encoder_hidden_states.shape[1] - - hidden_states = self.patch_embed( - hidden_states, encoder_hidden_states - ) # takes care of adding positional embeddings too. - emb = self.time_condition_embed(timestep, original_size, target_size, crop_coords, hidden_states.dtype) - - encoder_hidden_states = hidden_states[:, :text_seq_length] - hidden_states = hidden_states[:, text_seq_length:] - - for index_block, 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, - emb, - ) - else: - hidden_states, encoder_hidden_states = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - emb=emb, - ) - - hidden_states = self.norm_out(hidden_states, emb) - hidden_states = self.proj_out(hidden_states) # (batch_size, height*width, patch_size*patch_size*out_channels) - - # unpatchify - patch_size = self.config.patch_size - height = height // patch_size - width = width // patch_size - - hidden_states = hidden_states.reshape( - shape=(hidden_states.shape[0], height, width, self.out_channels, patch_size, patch_size) - ) - hidden_states = torch.einsum("nhwcpq->nchpwq", hidden_states) - output = hidden_states.reshape( - shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size) - ) - - if not return_dict: - return (output,) - - return Transformer2DModelOutput(sample=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_cogview4.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_cogview4.py deleted file mode 100644 index e4144d0c8e5750765dc186c54883acd8fa50f91c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_cogview4.py +++ /dev/null @@ -1,754 +0,0 @@ -# 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, Dict, List, Optional, Tuple, Union - -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 USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers -from ..attention import FeedForward -from ..attention_processor import Attention -from ..cache_utils import CacheMixin -from ..embeddings import CogView3CombinedTimestepSizeEmbeddings -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormContinuous - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class CogView4PatchEmbed(nn.Module): - def __init__( - self, - in_channels: int = 16, - hidden_size: int = 2560, - patch_size: int = 2, - text_hidden_size: int = 4096, - ): - super().__init__() - self.patch_size = patch_size - - self.proj = nn.Linear(in_channels * patch_size**2, hidden_size) - self.text_proj = nn.Linear(text_hidden_size, hidden_size) - - def forward(self, hidden_states: torch.Tensor, encoder_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) - encoder_hidden_states = self.text_proj(encoder_hidden_states) - - return hidden_states, encoder_hidden_states - - -class CogView4AdaLayerNormZero(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 CogView4AttnProcessor: - """ - Processor for implementing scaled dot-product attention for the CogView4 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. - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("CogView4AttnProcessor 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: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = 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)).transpose(1, 2) - key = key.unflatten(2, (attn.heads, -1)).transpose(1, 2) - value = value.unflatten(2, (attn.heads, -1)).transpose(1, 2) - - # 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, use_real_unbind_dim=-2 - ) - key[:, :, text_seq_length:, :] = apply_rotary_emb( - key[:, :, text_seq_length:, :], image_rotary_emb, use_real_unbind_dim=-2 - ) - - # 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 = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - hidden_states = hidden_states.transpose(1, 2).flatten(2, 3) - hidden_states = hidden_states.type_as(query) - - # 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 - - -class CogView4TrainingAttnProcessor: - """ - Training Processor for implementing scaled dot-product attention for the CogView4 model. It applies a rotary - embedding on query and key vectors, but does not include spatial normalization. - - This processor differs from CogView4AttnProcessor in several important ways: - 1. It supports attention masking with variable sequence lengths for multi-resolution training - 2. It unpacks and repacks sequences for efficient training with variable sequence lengths when batch_flag is - provided - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("CogView4AttnProcessor 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, - 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, - **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Args: - attn (`Attention`): - The attention module. - hidden_states (`torch.Tensor`): - The input hidden states. - encoder_hidden_states (`torch.Tensor`): - The encoder hidden states for cross-attention. - latent_attn_mask (`torch.Tensor`, *optional*): - Mask for latent tokens where 0 indicates pad token and 1 indicates non-pad token. If None, full - attention is used for all latent tokens. Note: the shape of latent_attn_mask is (batch_size, - num_latent_tokens). - text_attn_mask (`torch.Tensor`, *optional*): - Mask for text tokens where 0 indicates pad token and 1 indicates non-pad token. If None, full attention - is used for all text tokens. - batch_flag (`torch.Tensor`, *optional*): - 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*): - 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. - """ - - # Get dimensions and device info - batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape - batch_size, image_seq_length, embed_dim = hidden_states.shape - dtype = encoder_hidden_states.dtype - device = encoder_hidden_states.device - latent_hidden_states = hidden_states - # Combine text and image streams for joint processing - mixed_hidden_states = torch.cat([encoder_hidden_states, latent_hidden_states], dim=1) - - # 1. Construct attention mask and maybe packing input - # Create default masks if not provided - if text_attn_mask is None: - text_attn_mask = torch.ones((batch_size, text_seq_length), dtype=torch.int32, device=device) - if latent_attn_mask is None: - latent_attn_mask = torch.ones((batch_size, image_seq_length), dtype=torch.int32, device=device) - - # Validate mask shapes and types - assert text_attn_mask.dim() == 2, "the shape of text_attn_mask should be (batch_size, text_seq_length)" - assert text_attn_mask.dtype == torch.int32, "the dtype of text_attn_mask should be torch.int32" - assert latent_attn_mask.dim() == 2, "the shape of latent_attn_mask should be (batch_size, num_latent_tokens)" - assert latent_attn_mask.dtype == torch.int32, "the dtype of latent_attn_mask should be torch.int32" - - # Create combined mask for text and image tokens - mixed_attn_mask = torch.ones( - (batch_size, text_seq_length + image_seq_length), dtype=torch.int32, device=device - ) - mixed_attn_mask[:, :text_seq_length] = text_attn_mask - mixed_attn_mask[:, text_seq_length:] = latent_attn_mask - - # Convert mask to attention matrix format (where 1 means attend, 0 means don't attend) - mixed_attn_mask_input = mixed_attn_mask.unsqueeze(2).to(dtype=dtype) - attn_mask_matrix = mixed_attn_mask_input @ mixed_attn_mask_input.transpose(1, 2) - - # Handle batch packing if enabled - if batch_flag is not None: - assert batch_flag.dim() == 1 - # Determine packed batch size based on batch_flag - packing_batch_size = torch.max(batch_flag).item() + 1 - - # Calculate actual sequence lengths for each sample based on masks - text_seq_length = torch.sum(text_attn_mask, dim=1) - latent_seq_length = torch.sum(latent_attn_mask, dim=1) - mixed_seq_length = text_seq_length + latent_seq_length - - # Calculate packed sequence lengths for each packed batch - mixed_seq_length_packed = [ - torch.sum(mixed_attn_mask[batch_flag == batch_idx]).item() for batch_idx in range(packing_batch_size) - ] - - assert len(mixed_seq_length_packed) == packing_batch_size - - # Pack sequences by removing padding tokens - mixed_attn_mask_flatten = mixed_attn_mask.flatten(0, 1) - mixed_hidden_states_flatten = mixed_hidden_states.flatten(0, 1) - mixed_hidden_states_unpad = mixed_hidden_states_flatten[mixed_attn_mask_flatten == 1] - assert torch.sum(mixed_seq_length) == mixed_hidden_states_unpad.shape[0] - - # Split the unpadded sequence into packed batches - mixed_hidden_states_packed = torch.split(mixed_hidden_states_unpad, mixed_seq_length_packed) - - # Re-pad to create packed batches with right-side padding - mixed_hidden_states_packed_padded = torch.nn.utils.rnn.pad_sequence( - mixed_hidden_states_packed, - batch_first=True, - padding_value=0.0, - padding_side="right", - ) - - # Create attention mask for packed batches - l = mixed_hidden_states_packed_padded.shape[1] - attn_mask_matrix = torch.zeros( - (packing_batch_size, l, l), - dtype=dtype, - device=device, - ) - - # Fill attention mask with block diagonal matrices - # This ensures that tokens can only attend to other tokens within the same original sample - for idx, mask in enumerate(attn_mask_matrix): - seq_lengths = mixed_seq_length[batch_flag == idx] - offset = 0 - for length in seq_lengths: - # Create a block of 1s for each sample in the packed batch - mask[offset : offset + length, offset : offset + length] = 1 - offset += length - - attn_mask_matrix = attn_mask_matrix.to(dtype=torch.bool) - attn_mask_matrix = attn_mask_matrix.unsqueeze(1) # Add attention head dim - attention_mask = attn_mask_matrix - - # Prepare hidden states for attention computation - if batch_flag is None: - # If no packing, just combine text and image tokens - hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) - else: - # If packing, use the packed sequence - hidden_states = mixed_hidden_states_packed_padded - - # 2. QKV projections - convert hidden states to query, key, value - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - - # Reshape for multi-head attention: [batch, seq_len, heads*dim] -> [batch, heads, seq_len, dim] - 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) - - # 3. QK normalization - apply layer norm to queries and keys if configured - 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) - - # 4. Apply rotary positional embeddings to image tokens only - if image_rotary_emb is not None: - from ..embeddings import apply_rotary_emb - - if batch_flag is None: - # Apply RoPE only to image tokens (after text tokens) - query[:, :, text_seq_length:, :] = apply_rotary_emb( - query[:, :, text_seq_length:, :], image_rotary_emb, use_real_unbind_dim=-2 - ) - key[:, :, text_seq_length:, :] = apply_rotary_emb( - key[:, :, text_seq_length:, :], image_rotary_emb, use_real_unbind_dim=-2 - ) - else: - # For packed batches, need to carefully apply RoPE to appropriate tokens - assert query.shape[0] == packing_batch_size - assert key.shape[0] == packing_batch_size - assert len(image_rotary_emb) == batch_size - - rope_idx = 0 - for idx in range(packing_batch_size): - offset = 0 - # Get text and image sequence lengths for samples in this packed batch - text_seq_length_bi = text_seq_length[batch_flag == idx] - latent_seq_length_bi = latent_seq_length[batch_flag == idx] - - # Apply RoPE to each image segment in the packed sequence - for tlen, llen in zip(text_seq_length_bi, latent_seq_length_bi): - mlen = tlen + llen - # Apply RoPE only to image tokens (after text tokens) - query[idx, :, offset + tlen : offset + mlen, :] = apply_rotary_emb( - query[idx, :, offset + tlen : offset + mlen, :], - image_rotary_emb[rope_idx], - use_real_unbind_dim=-2, - ) - key[idx, :, offset + tlen : offset + mlen, :] = apply_rotary_emb( - key[idx, :, offset + tlen : offset + mlen, :], - image_rotary_emb[rope_idx], - use_real_unbind_dim=-2, - ) - offset += mlen - rope_idx += 1 - - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - # Reshape back: [batch, heads, seq_len, dim] -> [batch, seq_len, heads*dim] - hidden_states = hidden_states.transpose(1, 2).flatten(2, 3) - hidden_states = hidden_states.type_as(query) - - # 5. Output projection - project attention output to model dimension - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - - # Split the output back into text and image streams - if batch_flag is None: - # Simple split for non-packed case - encoder_hidden_states, hidden_states = hidden_states.split( - [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1 - ) - else: - # For packed case: need to unpack, split text/image, then restore to original shapes - # First, unpad the sequence based on the packed sequence lengths - hidden_states_unpad = torch.nn.utils.rnn.unpad_sequence( - hidden_states, - lengths=torch.tensor(mixed_seq_length_packed), - batch_first=True, - ) - # Concatenate all unpadded sequences - hidden_states_flatten = torch.cat(hidden_states_unpad, dim=0) - # Split by original sample sequence lengths - hidden_states_unpack = torch.split(hidden_states_flatten, mixed_seq_length.tolist()) - assert len(hidden_states_unpack) == batch_size - - # Further split each sample's sequence into text and image parts - hidden_states_unpack = [ - torch.split(h, [tlen, llen]) - for h, tlen, llen in zip(hidden_states_unpack, text_seq_length, latent_seq_length) - ] - # Separate text and image sequences - encoder_hidden_states_unpad = [h[0] for h in hidden_states_unpack] - hidden_states_unpad = [h[1] for h in hidden_states_unpack] - - # Update the original tensors with the processed values, respecting the attention masks - for idx in range(batch_size): - # Place unpacked text tokens back in the encoder_hidden_states tensor - encoder_hidden_states[idx][text_attn_mask[idx] == 1] = encoder_hidden_states_unpad[idx] - # Place unpacked image tokens back in the latent_hidden_states tensor - latent_hidden_states[idx][latent_attn_mask[idx] == 1] = hidden_states_unpad[idx] - - # Update the output hidden states - hidden_states = latent_hidden_states - - return hidden_states, encoder_hidden_states - - -class CogView4TransformerBlock(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 = CogView4AdaLayerNormZero(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=CogView4AttnProcessor(), - ) - - # 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: 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, - ) -> 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 - if attention_kwargs is None: - attention_kwargs = {} - 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, - **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 CogView4RotaryPosEmbed(nn.Module): - def __init__(self, dim: int, patch_size: int, rope_axes_dim: Tuple[int, int], theta: float = 10000.0) -> None: - super().__init__() - - self.dim = dim - self.patch_size = patch_size - self.rope_axes_dim = rope_axes_dim - 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(self.rope_axes_dim[0]) - w_seq = torch.arange(self.rope_axes_dim[1]) - freqs_h = torch.outer(h_seq, h_inv_freq) - freqs_w = torch.outer(w_seq, w_inv_freq) - - h_idx = torch.arange(height, device=freqs_h.device) - w_idx = torch.arange(width, device=freqs_w.device) - inner_h_idx = h_idx * self.rope_axes_dim[0] // height - inner_w_idx = w_idx * self.rope_axes_dim[1] // width - - freqs_h = freqs_h[inner_h_idx] - freqs_w = freqs_w[inner_w_idx] - - # 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 CogView4Transformer2DModel(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 `4096`): - 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 = ["CogView4TransformerBlock", "CogView4PatchEmbed", "CogView4PatchEmbed"] - _skip_layerwise_casting_patterns = ["patch_embed", "norm", "proj_out"] - - @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 = 4096, - time_embed_dim: int = 512, - condition_dim: int = 256, - pos_embed_max_size: int = 128, - sample_size: int = 128, - rope_axes_dim: Tuple[int, int] = (256, 256), - ): - super().__init__() - - # CogView4 uses 3 additional SDXL-like conditions - original_size, target_size, crop_coords - # Each of these are sincos embeddings of shape 2 * condition_dim - pooled_projection_dim = 3 * 2 * condition_dim - inner_dim = num_attention_heads * attention_head_dim - out_channels = out_channels - - # 1. RoPE - self.rope = CogView4RotaryPosEmbed(attention_head_dim, patch_size, rope_axes_dim, theta=10000.0) - - # 2. Patch & Text-timestep embedding - self.patch_embed = CogView4PatchEmbed(in_channels, inner_dim, patch_size, text_embed_dim) - - self.time_condition_embed = CogView3CombinedTimestepSizeEmbeddings( - embedding_dim=time_embed_dim, - condition_dim=condition_dim, - pooled_projection_dim=pooled_projection_dim, - timesteps_dim=inner_dim, - ) - - # 3. Transformer blocks - self.transformer_blocks = nn.ModuleList( - [ - CogView4TransformerBlock(inner_dim, num_attention_heads, attention_head_dim, time_embed_dim) - for _ in range(num_layers) - ] - ) - - # 4. Output projection - self.norm_out = AdaLayerNormContinuous(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, - timestep: torch.LongTensor, - original_size: torch.Tensor, - target_size: torch.Tensor, - crop_coords: torch.Tensor, - attention_kwargs: Optional[Dict[str, Any]] = 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[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." - ) - - 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, encoder_hidden_states = self.patch_embed(hidden_states, encoder_hidden_states) - - temb = self.time_condition_embed(timestep, original_size, target_size, crop_coords, hidden_states.dtype) - temb = F.silu(temb) - - # 3. Transformer blocks - for block in 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, - ) - else: - hidden_states, encoder_hidden_states = block( - hidden_states, - encoder_hidden_states, - temb, - image_rotary_emb, - attention_mask, - attention_kwargs, - ) - - # 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) - 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/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_cosmos.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_cosmos.py deleted file mode 100644 index 6c312b7a5a3fea3beb042f7b2c0b5f6d3b919c5a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_cosmos.py +++ /dev/null @@ -1,579 +0,0 @@ -# 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 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 ...utils import is_torchvision_available -from ..attention import FeedForward -from ..attention_processor import Attention -from ..embeddings import Timesteps -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import RMSNorm - - -if is_torchvision_available(): - from torchvision import transforms - - -class CosmosPatchEmbed(nn.Module): - def __init__( - self, in_channels: int, out_channels: int, patch_size: Tuple[int, int, int], bias: bool = True - ) -> None: - super().__init__() - self.patch_size = patch_size - - self.proj = nn.Linear(in_channels * patch_size[0] * patch_size[1] * patch_size[2], out_channels, bias=bias) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - batch_size, num_channels, num_frames, height, width = hidden_states.shape - p_t, p_h, p_w = self.patch_size - hidden_states = hidden_states.reshape( - batch_size, num_channels, num_frames // p_t, p_t, height // p_h, p_h, width // p_w, p_w - ) - hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7) - hidden_states = self.proj(hidden_states) - return hidden_states - - -class CosmosTimestepEmbedding(nn.Module): - def __init__(self, in_features: int, out_features: int) -> None: - super().__init__() - self.linear_1 = nn.Linear(in_features, out_features, bias=False) - self.activation = nn.SiLU() - self.linear_2 = nn.Linear(out_features, 3 * out_features, bias=False) - - def forward(self, timesteps: torch.Tensor) -> torch.Tensor: - emb = self.linear_1(timesteps) - emb = self.activation(emb) - emb = self.linear_2(emb) - return emb - - -class CosmosEmbedding(nn.Module): - def __init__(self, embedding_dim: int, condition_dim: int) -> None: - super().__init__() - - self.time_proj = Timesteps(embedding_dim, flip_sin_to_cos=True, downscale_freq_shift=0.0) - self.t_embedder = CosmosTimestepEmbedding(embedding_dim, condition_dim) - self.norm = RMSNorm(embedding_dim, eps=1e-6, elementwise_affine=True) - - def forward(self, hidden_states: torch.Tensor, timestep: torch.LongTensor) -> torch.Tensor: - timesteps_proj = self.time_proj(timestep).type_as(hidden_states) - temb = self.t_embedder(timesteps_proj) - embedded_timestep = self.norm(timesteps_proj) - return temb, embedded_timestep - - -class CosmosAdaLayerNorm(nn.Module): - def __init__(self, in_features: int, hidden_features: int) -> None: - super().__init__() - self.embedding_dim = in_features - - self.activation = nn.SiLU() - self.norm = nn.LayerNorm(in_features, elementwise_affine=False, eps=1e-6) - self.linear_1 = nn.Linear(in_features, hidden_features, bias=False) - 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 - ) -> torch.Tensor: - embedded_timestep = self.activation(embedded_timestep) - embedded_timestep = self.linear_1(embedded_timestep) - embedded_timestep = self.linear_2(embedded_timestep) - - if temb is not None: - embedded_timestep = embedded_timestep + temb[..., : 2 * self.embedding_dim] - - shift, scale = embedded_timestep.chunk(2, dim=-1) - hidden_states = self.norm(hidden_states) - - if embedded_timestep.ndim == 2: - shift, scale = (x.unsqueeze(1) for x in (shift, scale)) - - hidden_states = hidden_states * (1 + scale) + shift - return hidden_states - - -class CosmosAdaLayerNormZero(nn.Module): - def __init__(self, in_features: int, hidden_features: Optional[int] = None) -> None: - super().__init__() - - self.norm = nn.LayerNorm(in_features, elementwise_affine=False, eps=1e-6) - self.activation = nn.SiLU() - - if hidden_features is None: - self.linear_1 = nn.Identity() - else: - self.linear_1 = nn.Linear(in_features, hidden_features, bias=False) - - self.linear_2 = nn.Linear(hidden_features, 3 * in_features, bias=False) - - def forward( - self, - hidden_states: torch.Tensor, - embedded_timestep: torch.Tensor, - temb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - embedded_timestep = self.activation(embedded_timestep) - embedded_timestep = self.linear_1(embedded_timestep) - embedded_timestep = self.linear_2(embedded_timestep) - - if temb is not None: - embedded_timestep = embedded_timestep + temb - - shift, scale, gate = embedded_timestep.chunk(3, dim=-1) - hidden_states = self.norm(hidden_states) - - if embedded_timestep.ndim == 2: - shift, scale, gate = (x.unsqueeze(1) for x in (shift, scale, gate)) - - hidden_states = hidden_states * (1 + scale) + shift - return hidden_states, gate - - -class CosmosAttnProcessor2_0: - def __init__(self): - if not hasattr(F, "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, - ) -> torch.Tensor: - # 1. QKV projections - 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 = 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) - - # 2. QK normalization - query = attn.norm_q(query) - key = attn.norm_k(key) - - # 3. Apply RoPE - 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) - - # 4. Prepare for GQA - 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) - key = key.repeat_interleave(query_idx // key_idx, dim=3) - 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 = hidden_states.transpose(1, 2).flatten(2, 3).type_as(query) - - # 6. Output projection - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - - return hidden_states - - -class CosmosTransformerBlock(nn.Module): - def __init__( - self, - num_attention_heads: int, - attention_head_dim: int, - cross_attention_dim: int, - mlp_ratio: float = 4.0, - adaln_lora_dim: int = 256, - qk_norm: str = "rms_norm", - out_bias: 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.attn1 = Attention( - query_dim=hidden_size, - cross_attention_dim=None, - heads=num_attention_heads, - dim_head=attention_head_dim, - qk_norm=qk_norm, - elementwise_affine=True, - out_bias=out_bias, - processor=CosmosAttnProcessor2_0(), - ) - - 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(), - ) - - 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) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - 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: - if extra_pos_emb is not None: - hidden_states = hidden_states + extra_pos_emb - - # 1. Self Attention - norm_hidden_states, gate = self.norm1(hidden_states, embedded_timestep, temb) - attn_output = self.attn1(norm_hidden_states, image_rotary_emb=image_rotary_emb) - hidden_states = hidden_states + gate * attn_output - - # 2. Cross Attention - norm_hidden_states, gate = self.norm2(hidden_states, embedded_timestep, temb) - attn_output = self.attn2( - norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask - ) - hidden_states = hidden_states + gate * attn_output - - # 3. Feed Forward - norm_hidden_states, gate = self.norm3(hidden_states, embedded_timestep, temb) - ff_output = self.ff(norm_hidden_states) - hidden_states = hidden_states + gate * ff_output - - return hidden_states - - -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), - base_fps: int = 24, - rope_scale: Tuple[float, float, float] = (2.0, 1.0, 1.0), - ) -> None: - super().__init__() - - self.max_size = [size // patch for size, patch in zip(max_size, patch_size)] - self.patch_size = patch_size - self.base_fps = base_fps - - self.dim_h = hidden_size // 6 * 2 - self.dim_w = hidden_size // 6 * 2 - self.dim_t = hidden_size - self.dim_h - self.dim_w - - self.h_ntk_factor = rope_scale[1] ** (self.dim_h / (self.dim_h - 2)) - 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]: - 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 - - h_theta = 10000.0 * self.h_ntk_factor - w_theta = 10000.0 * self.w_ntk_factor - t_theta = 10000.0 * self.t_ntk_factor - - seq = torch.arange(max(self.max_size), device=device, dtype=torch.float32) - dim_h_range = ( - torch.arange(0, self.dim_h, 2, device=device, dtype=torch.float32)[: (self.dim_h // 2)] / self.dim_h - ) - dim_w_range = ( - torch.arange(0, self.dim_w, 2, device=device, dtype=torch.float32)[: (self.dim_w // 2)] / self.dim_w - ) - dim_t_range = ( - torch.arange(0, self.dim_t, 2, device=device, dtype=torch.float32)[: (self.dim_t // 2)] / self.dim_t - ) - h_spatial_freqs = 1.0 / (h_theta**dim_h_range) - w_spatial_freqs = 1.0 / (w_theta**dim_w_range) - temporal_freqs = 1.0 / (t_theta**dim_t_range) - - emb_h = torch.outer(seq[: pe_size[1]], h_spatial_freqs)[None, :, None, :].repeat(pe_size[0], 1, pe_size[2], 1) - emb_w = torch.outer(seq[: pe_size[2]], w_spatial_freqs)[None, None, :, :].repeat(pe_size[0], pe_size[1], 1, 1) - - # Apply sequence scaling in temporal dimension - if fps is None: - # Images - emb_t = torch.outer(seq[: pe_size[0]], temporal_freqs) - else: - # Videos - emb_t = torch.outer(seq[: pe_size[0]] / fps * self.base_fps, temporal_freqs) - - emb_t = emb_t[:, None, None, :].repeat(1, pe_size[1], pe_size[2], 1) - freqs = torch.cat([emb_t, emb_h, emb_w] * 2, dim=-1).flatten(0, 2).float() - cos = torch.cos(freqs) - sin = torch.sin(freqs) - return cos, sin - - -class CosmosLearnablePositionalEmbed(nn.Module): - def __init__( - self, - hidden_size: int, - max_size: Tuple[int, int, int], - patch_size: Tuple[int, int, int], - eps: float = 1e-6, - ) -> None: - super().__init__() - - self.max_size = [size // patch for size, patch in zip(max_size, patch_size)] - self.patch_size = patch_size - self.eps = eps - - self.pos_emb_t = nn.Parameter(torch.zeros(self.max_size[0], hidden_size)) - self.pos_emb_h = nn.Parameter(torch.zeros(self.max_size[1], hidden_size)) - self.pos_emb_w = nn.Parameter(torch.zeros(self.max_size[2], hidden_size)) - - def forward(self, hidden_states: 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]] - - emb_t = self.pos_emb_t[: pe_size[0]][None, :, None, None, :].repeat(batch_size, 1, pe_size[1], pe_size[2], 1) - emb_h = self.pos_emb_h[: pe_size[1]][None, None, :, None, :].repeat(batch_size, pe_size[0], 1, pe_size[2], 1) - emb_w = self.pos_emb_w[: pe_size[2]][None, None, None, :, :].repeat(batch_size, pe_size[0], pe_size[1], 1, 1) - emb = emb_t + emb_h + emb_w - emb = emb.flatten(1, 3) - - norm = torch.linalg.vector_norm(emb, dim=-1, keepdim=True, dtype=torch.float32) - norm = torch.add(self.eps, norm, alpha=np.sqrt(norm.numel() / emb.numel())) - return (emb / norm).type_as(hidden_states) - - -class CosmosTransformer3DModel(ModelMixin, ConfigMixin): - r""" - A Transformer model for video-like data used in [Cosmos](https://github.com/NVIDIA/Cosmos). - - Args: - 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. - num_attention_heads (`int`, defaults to `32`): - The number of heads to use for multi-head attention. - attention_head_dim (`int`, defaults to `128`): - The number of channels in each attention head. - num_layers (`int`, defaults to `28`): - The number of layers of transformer blocks to use. - mlp_ratio (`float`, defaults to `4.0`): - The ratio of the hidden layer size to the input size in the feedforward network. - text_embed_dim (`int`, defaults to `4096`): - 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)`): - 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)`): - 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)`): - 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`. - """ - - _supports_gradient_checkpointing = True - _skip_layerwise_casting_patterns = ["patch_embed", "final_layer", "norm"] - _no_split_modules = ["CosmosTransformerBlock"] - _keep_in_fp32_modules = ["learnable_pos_embed"] - - @register_to_config - def __init__( - self, - in_channels: int = 16, - out_channels: int = 16, - num_attention_heads: int = 32, - attention_head_dim: int = 128, - num_layers: int = 28, - 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), - concat_padding_mask: bool = True, - extra_pos_embed_type: Optional[str] = "learnable", - ) -> None: - super().__init__() - hidden_size = num_attention_heads * attention_head_dim - - # 1. Patch Embedding - patch_embed_in_channels = in_channels + 1 if concat_padding_mask else in_channels - self.patch_embed = CosmosPatchEmbed(patch_embed_in_channels, hidden_size, patch_size, bias=False) - - # 2. Positional Embedding - self.rope = CosmosRotaryPosEmbed( - hidden_size=attention_head_dim, max_size=max_size, patch_size=patch_size, rope_scale=rope_scale - ) - - self.learnable_pos_embed = None - if extra_pos_embed_type == "learnable": - self.learnable_pos_embed = CosmosLearnablePositionalEmbed( - hidden_size=hidden_size, - max_size=max_size, - patch_size=patch_size, - ) - - # 3. Time Embedding - self.time_embed = CosmosEmbedding(hidden_size, hidden_size) - - # 4. Transformer Blocks - self.transformer_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, - ) - for _ in range(num_layers) - ] - ) - - # 5. Output norm & projection - self.norm_out = CosmosAdaLayerNorm(hidden_size, adaln_lora_dim) - self.proj_out = nn.Linear( - hidden_size, patch_size[0] * patch_size[1] * patch_size[2] * out_channels, bias=False - ) - - self.gradient_checkpointing = False - - 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, - return_dict: bool = True, - ) -> torch.Tensor: - batch_size, num_channels, num_frames, height, width = hidden_states.shape - - # 1. Concatenate padding mask if needed & prepare attention mask - if condition_mask is not None: - hidden_states = torch.cat([hidden_states, condition_mask], dim=1) - - if self.config.concat_padding_mask: - padding_mask = 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 - ) - - if attention_mask is not None: - attention_mask = attention_mask.unsqueeze(1).unsqueeze(1) # [B, 1, 1, S] - - # 2. Generate positional embeddings - image_rotary_emb = self.rope(hidden_states, fps=fps) - extra_pos_emb = self.learnable_pos_embed(hidden_states) if self.config.extra_pos_embed_type else None - - # 3. Patchify input - p_t, p_h, p_w = self.config.patch_size - post_patch_num_frames = num_frames // p_t - post_patch_height = height // p_h - post_patch_width = width // p_w - hidden_states = self.patch_embed(hidden_states) - hidden_states = hidden_states.flatten(1, 3) # [B, T, H, W, C] -> [B, THW, C] - - # 4. Timestep embeddings - if timestep.ndim == 1: - temb, embedded_timestep = self.time_embed(hidden_states, timestep) - elif timestep.ndim == 5: - 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 = timestep.flatten() - temb, embedded_timestep = self.time_embed(hidden_states, timestep) - # We can do this because num_frames == post_patch_num_frames, as p_t is 1 - 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) - ) # [BT, C] -> [B, T, 1, 1, C] -> [B, T, H, W, C] -> [B, THW, C] - else: - assert False - - # 5. 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, - encoder_hidden_states, - embedded_timestep, - temb, - image_rotary_emb, - extra_pos_emb, - attention_mask, - ) - 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, - ) - - # 6. 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)) - hidden_states = hidden_states.unflatten(1, (post_patch_num_frames, post_patch_height, post_patch_width)) - # NOTE: The permutation order here is not the inverse operation of what happens when patching as usually expected. - # It might be a source of confusion to the reader, but this is correct - hidden_states = hidden_states.permute(0, 7, 1, 6, 2, 4, 3, 5) - hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - - if not return_dict: - return (hidden_states,) - - return Transformer2DModelOutput(sample=hidden_states) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_easyanimate.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_easyanimate.py deleted file mode 100644 index 545fa29730db226f27296b6318b1c13018bee63b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_easyanimate.py +++ /dev/null @@ -1,527 +0,0 @@ -# Copyright 2025 The EasyAnimate 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 List, Optional, Tuple, Union - -import torch -import torch.nn.functional as F -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import logging -from ...utils.torch_utils import maybe_allow_in_graph -from ..attention import Attention, FeedForward -from ..embeddings import TimestepEmbedding, Timesteps, get_3d_rotary_pos_embed -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNorm, FP32LayerNorm, RMSNorm - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class EasyAnimateLayerNormZero(nn.Module): - def __init__( - self, - conditioning_dim: int, - embedding_dim: int, - elementwise_affine: bool = True, - eps: float = 1e-5, - bias: bool = True, - norm_type: str = "fp32_layer_norm", - ) -> None: - super().__init__() - - self.silu = nn.SiLU() - self.linear = nn.Linear(conditioning_dim, 6 * embedding_dim, bias=bias) - - if norm_type == "layer_norm": - self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=elementwise_affine, eps=eps) - elif norm_type == "fp32_layer_norm": - self.norm = FP32LayerNorm(embedding_dim, elementwise_affine=elementwise_affine, eps=eps) - else: - raise ValueError( - f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'." - ) - - def forward( - self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: 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( - 1 - ) - return hidden_states, encoder_hidden_states, gate, enc_gate - - -class EasyAnimateRotaryPosEmbed(nn.Module): - def __init__(self, patch_size: int, rope_dim: List[int]) -> None: - super().__init__() - - self.patch_size = patch_size - self.rope_dim = rope_dim - - def get_resize_crop_region_for_grid(self, src, tgt_width, tgt_height): - tw = tgt_width - th = tgt_height - h, w = src - r = h / w - if r > (th / tw): - resize_height = th - resize_width = int(round(th / h * w)) - else: - resize_width = tw - resize_height = int(round(tw / w * h)) - - crop_top = int(round((th - resize_height) / 2.0)) - crop_left = int(round((tw - resize_width) / 2.0)) - - return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - bs, c, num_frames, grid_height, grid_width = hidden_states.size() - grid_height = grid_height // self.patch_size - grid_width = grid_width // self.patch_size - base_size_width = 90 // self.patch_size - base_size_height = 60 // self.patch_size - - grid_crops_coords = self.get_resize_crop_region_for_grid( - (grid_height, grid_width), base_size_width, base_size_height - ) - image_rotary_emb = get_3d_rotary_pos_embed( - self.rope_dim, - grid_crops_coords, - grid_size=(grid_height, grid_width), - temporal_size=hidden_states.size(2), - use_real=True, - ) - return image_rotary_emb - - -class EasyAnimateAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is - used in the EasyAnimateTransformer3DModel model. - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "EasyAnimateAttnProcessor2_0 requires PyTorch 2.0 or above. To use it, please install PyTorch 2.0." - ) - - def __call__( - self, - attn: Attention, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = 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) - - # 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)).transpose(1, 2) - key = key.unflatten(2, (attn.heads, -1)).transpose(1, 2) - value = value.unflatten(2, (attn.heads, -1)).transpose(1, 2) - - # 2. QK normalization - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # 3. Encoder condition QKV projection and normalization - if attn.add_q_proj is not None and encoder_hidden_states is not None: - encoder_query = attn.add_q_proj(encoder_hidden_states) - encoder_key = attn.add_k_proj(encoder_hidden_states) - encoder_value = attn.add_v_proj(encoder_hidden_states) - - encoder_query = encoder_query.unflatten(2, (attn.heads, -1)).transpose(1, 2) - encoder_key = encoder_key.unflatten(2, (attn.heads, -1)).transpose(1, 2) - encoder_value = encoder_value.unflatten(2, (attn.heads, -1)).transpose(1, 2) - - if attn.norm_added_q is not None: - encoder_query = attn.norm_added_q(encoder_query) - if attn.norm_added_k is not None: - encoder_key = attn.norm_added_k(encoder_key) - - query = torch.cat([encoder_query, query], dim=2) - key = torch.cat([encoder_key, key], dim=2) - value = torch.cat([encoder_value, value], dim=2) - - if image_rotary_emb is not None: - from ..embeddings import apply_rotary_emb - - query[:, :, encoder_hidden_states.shape[1] :] = apply_rotary_emb( - query[:, :, encoder_hidden_states.shape[1] :], image_rotary_emb - ) - if not attn.is_cross_attention: - key[:, :, encoder_hidden_states.shape[1] :] = apply_rotary_emb( - key[:, :, encoder_hidden_states.shape[1] :], image_rotary_emb - ) - - # 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 = hidden_states.transpose(1, 2).flatten(2, 3) - hidden_states = hidden_states.to(query.dtype) - - # 6. Output projection - if encoder_hidden_states is not None: - encoder_hidden_states, hidden_states = ( - hidden_states[:, : encoder_hidden_states.shape[1]], - hidden_states[:, encoder_hidden_states.shape[1] :], - ) - - if getattr(attn, "to_out", None) is not None: - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - - if getattr(attn, "to_add_out", None) is not None: - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - else: - if getattr(attn, "to_out", None) is not None: - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - - return hidden_states, encoder_hidden_states - - -@maybe_allow_in_graph -class EasyAnimateTransformerBlock(nn.Module): - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - time_embed_dim: int, - dropout: float = 0.0, - activation_fn: str = "gelu-approximate", - norm_elementwise_affine: bool = True, - norm_eps: float = 1e-6, - final_dropout: bool = True, - ff_inner_dim: Optional[int] = None, - ff_bias: bool = True, - qk_norm: bool = True, - after_norm: bool = False, - norm_type: str = "fp32_layer_norm", - is_mmdit_block: bool = True, - ): - super().__init__() - - # Attention Part - self.norm1 = EasyAnimateLayerNormZero( - time_embed_dim, dim, norm_elementwise_affine, norm_eps, norm_type=norm_type, bias=True - ) - - self.attn1 = Attention( - query_dim=dim, - dim_head=attention_head_dim, - heads=num_attention_heads, - qk_norm="layer_norm" if qk_norm else None, - eps=1e-6, - bias=True, - added_proj_bias=True, - added_kv_proj_dim=dim if is_mmdit_block else None, - context_pre_only=False if is_mmdit_block else None, - processor=EasyAnimateAttnProcessor2_0(), - ) - - # FFN Part - self.norm2 = EasyAnimateLayerNormZero( - time_embed_dim, dim, norm_elementwise_affine, norm_eps, norm_type=norm_type, bias=True - ) - self.ff = FeedForward( - dim, - dropout=dropout, - activation_fn=activation_fn, - final_dropout=final_dropout, - inner_dim=ff_inner_dim, - bias=ff_bias, - ) - - self.txt_ff = None - if is_mmdit_block: - self.txt_ff = FeedForward( - dim, - dropout=dropout, - activation_fn=activation_fn, - final_dropout=final_dropout, - inner_dim=ff_inner_dim, - bias=ff_bias, - ) - - self.norm3 = None - if after_norm: - self.norm3 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) - - def forward( - self, - 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]: - # 1. Attention - norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1( - hidden_states, encoder_hidden_states, temb - ) - 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, - ) - hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_hidden_states - encoder_hidden_states = encoder_hidden_states + enc_gate_msa.unsqueeze(1) * attn_encoder_hidden_states - - # 2. Feed-forward - norm_hidden_states, norm_encoder_hidden_states, gate_ff, enc_gate_ff = self.norm2( - hidden_states, encoder_hidden_states, temb - ) - if self.norm3 is not None: - norm_hidden_states = self.norm3(self.ff(norm_hidden_states)) - if self.txt_ff is not None: - norm_encoder_hidden_states = self.norm3(self.txt_ff(norm_encoder_hidden_states)) - else: - norm_encoder_hidden_states = self.norm3(self.ff(norm_encoder_hidden_states)) - else: - norm_hidden_states = self.ff(norm_hidden_states) - if self.txt_ff is not None: - norm_encoder_hidden_states = self.txt_ff(norm_encoder_hidden_states) - else: - norm_encoder_hidden_states = self.ff(norm_encoder_hidden_states) - hidden_states = hidden_states + gate_ff.unsqueeze(1) * norm_hidden_states - encoder_hidden_states = encoder_hidden_states + enc_gate_ff.unsqueeze(1) * norm_encoder_hidden_states - return hidden_states, encoder_hidden_states - - -class EasyAnimateTransformer3DModel(ModelMixin, ConfigMixin): - """ - A Transformer model for video-like data in [EasyAnimate](https://github.com/aigc-apps/EasyAnimate). - - Parameters: - num_attention_heads (`int`, defaults to `48`): - The number of heads to use for multi-head attention. - attention_head_dim (`int`, defaults to `64`): - The number of channels in each head. - in_channels (`int`, defaults to `16`): - The number of channels in the input. - out_channels (`int`, *optional*, defaults to `16`): - The number of channels in the output. - patch_size (`int`, defaults to `2`): - The size of the patches to use in the patch embedding layer. - sample_width (`int`, defaults to `90`): - The width of the input latents. - sample_height (`int`, defaults to `60`): - The height of the input latents. - activation_fn (`str`, defaults to `"gelu-approximate"`): - Activation function to use in feed-forward. - timestep_activation_fn (`str`, defaults to `"silu"`): - Activation function to use when generating the timestep embeddings. - num_layers (`int`, defaults to `30`): - The number of layers of Transformer blocks to use. - mmdit_layers (`int`, defaults to `1000`): - The number of layers of Multi Modal Transformer blocks to use. - dropout (`float`, defaults to `0.0`): - The dropout probability to use. - time_embed_dim (`int`, defaults to `512`): - Output dimension of timestep embeddings. - text_embed_dim (`int`, defaults to `4096`): - Input dimension of text embeddings from the text encoder. - norm_eps (`float`, defaults to `1e-5`): - The epsilon value to use in normalization layers. - norm_elementwise_affine (`bool`, defaults to `True`): - Whether to use elementwise affine in normalization layers. - flip_sin_to_cos (`bool`, defaults to `True`): - Whether to flip the sin to cos in the time embedding. - time_position_encoding_type (`str`, defaults to `3d_rope`): - Type of time position encoding. - after_norm (`bool`, defaults to `False`): - Flag to apply normalization after. - resize_inpaint_mask_directly (`bool`, defaults to `True`): - Flag to resize inpaint mask directly. - enable_text_attention_mask (`bool`, defaults to `True`): - Flag to enable text attention mask. - add_noise_in_inpaint_model (`bool`, defaults to `False`): - Flag to add noise in inpaint model. - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["EasyAnimateTransformerBlock"] - _skip_layerwise_casting_patterns = ["^proj$", "norm", "^proj_out$"] - - @register_to_config - 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, - sample_width: int = 90, - sample_height: int = 60, - activation_fn: str = "gelu-approximate", - timestep_activation_fn: str = "silu", - freq_shift: int = 0, - num_layers: int = 48, - mmdit_layers: int = 48, - dropout: float = 0.0, - time_embed_dim: int = 512, - add_norm_text_encoder: bool = False, - text_embed_dim: int = 3584, - text_embed_dim_t5: int = None, - norm_eps: float = 1e-5, - norm_elementwise_affine: bool = True, - flip_sin_to_cos: bool = True, - time_position_encoding_type: str = "3d_rope", - after_norm=False, - resize_inpaint_mask_directly: bool = True, - enable_text_attention_mask: bool = True, - add_noise_in_inpaint_model: bool = True, - ): - super().__init__() - inner_dim = num_attention_heads * attention_head_dim - - # 1. Timestep embedding - self.time_proj = Timesteps(inner_dim, flip_sin_to_cos, freq_shift) - self.time_embedding = TimestepEmbedding(inner_dim, time_embed_dim, timestep_activation_fn) - self.rope_embedding = EasyAnimateRotaryPosEmbed(patch_size, attention_head_dim) - - # 2. Patch embedding - self.proj = nn.Conv2d( - in_channels, inner_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=True - ) - - # 3. Text refined embedding - self.text_proj = None - self.text_proj_t5 = None - if not add_norm_text_encoder: - self.text_proj = nn.Linear(text_embed_dim, inner_dim) - if text_embed_dim_t5 is not None: - self.text_proj_t5 = nn.Linear(text_embed_dim_t5, inner_dim) - else: - self.text_proj = nn.Sequential( - RMSNorm(text_embed_dim, 1e-6, elementwise_affine=True), nn.Linear(text_embed_dim, inner_dim) - ) - if text_embed_dim_t5 is not None: - self.text_proj_t5 = nn.Sequential( - RMSNorm(text_embed_dim, 1e-6, elementwise_affine=True), nn.Linear(text_embed_dim_t5, inner_dim) - ) - - # 4. Transformer blocks - self.transformer_blocks = nn.ModuleList( - [ - EasyAnimateTransformerBlock( - dim=inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - time_embed_dim=time_embed_dim, - dropout=dropout, - activation_fn=activation_fn, - norm_elementwise_affine=norm_elementwise_affine, - norm_eps=norm_eps, - after_norm=after_norm, - is_mmdit_block=True if _ < mmdit_layers else False, - ) - for _ in range(num_layers) - ] - ) - self.norm_final = nn.LayerNorm(inner_dim, norm_eps, norm_elementwise_affine) - - # 5. Output norm & projection - self.norm_out = AdaLayerNorm( - embedding_dim=time_embed_dim, - output_dim=2 * inner_dim, - norm_elementwise_affine=norm_elementwise_affine, - norm_eps=norm_eps, - chunk_dim=1, - ) - self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels) - - self.gradient_checkpointing = False - - 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, - return_dict: bool = True, - ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: - batch_size, channels, video_length, height, width = hidden_states.size() - p = self.config.patch_size - post_patch_height = height // p - post_patch_width = width // p - - # 1. Time embedding - temb = self.time_proj(timestep).to(dtype=hidden_states.dtype) - temb = self.time_embedding(temb, timestep_cond) - image_rotary_emb = self.rope_embedding(hidden_states) - - # 2. Patch embedding - if inpaint_latents is not None: - hidden_states = torch.concat([hidden_states, inpaint_latents], 1) - if control_latents is not None: - hidden_states = torch.concat([hidden_states, control_latents], 1) - - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) # [B, C, F, H, W] -> [BF, C, H, W] - hidden_states = self.proj(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute( - 0, 2, 1, 3, 4 - ) # [BF, C, H, W] -> [B, F, C, H, W] - hidden_states = hidden_states.flatten(2, 4).transpose(1, 2) # [B, F, C, H, W] -> [B, FHW, C] - - # 3. Text embedding - encoder_hidden_states = self.text_proj(encoder_hidden_states) - if encoder_hidden_states_t5 is not None: - encoder_hidden_states_t5 = self.text_proj_t5(encoder_hidden_states_t5) - encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states_t5], dim=1).contiguous() - - # 4. Transformer blocks - for block in 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 - ) - else: - hidden_states, encoder_hidden_states = block( - hidden_states, encoder_hidden_states, temb, image_rotary_emb - ) - - hidden_states = self.norm_final(hidden_states) - - # 5. Output norm & projection - hidden_states = self.norm_out(hidden_states, temb=temb) - hidden_states = self.proj_out(hidden_states) - - # 6. Unpatchify - p = self.config.patch_size - output = hidden_states.reshape(batch_size, video_length, post_patch_height, post_patch_width, channels, p, p) - output = output.permute(0, 4, 1, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4) - - if not return_dict: - return (output,) - return Transformer2DModelOutput(sample=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_flux.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_flux.py deleted file mode 100644 index ab579a0eb53155404a92e709c22f76e09db938f6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_flux.py +++ /dev/null @@ -1,549 +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 Any, Dict, Optional, Tuple, Union - -import numpy as np -import torch -import torch.nn as nn - -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_utils import is_torch_npu_available -from ...utils.torch_utils import maybe_allow_in_graph -from ..attention import FeedForward -from ..attention_processor import ( - Attention, - AttentionProcessor, - FluxAttnProcessor2_0, - FluxAttnProcessor2_0_NPU, - FusedFluxAttnProcessor2_0, -) -from ..cache_utils import CacheMixin -from ..embeddings import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, FluxPosEmbed -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@maybe_allow_in_graph -class FluxSingleTransformerBlock(nn.Module): - def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 4.0): - super().__init__() - self.mlp_hidden_dim = int(dim * mlp_ratio) - - self.norm = AdaLayerNormZeroSingle(dim) - self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim) - self.act_mlp = nn.GELU(approximate="tanh") - self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim) - - if is_torch_npu_available(): - deprecation_message = ( - "Defaulting to FluxAttnProcessor2_0_NPU for NPU devices will be removed. Attention processors " - "should be set explicitly using the `set_attn_processor` method." - ) - deprecate("npu_processor", "0.34.0", deprecation_message) - processor = FluxAttnProcessor2_0_NPU() - else: - processor = FluxAttnProcessor2_0() - - self.attn = Attention( - query_dim=dim, - cross_attention_dim=None, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=dim, - bias=True, - processor=processor, - qk_norm="rms_norm", - eps=1e-6, - pre_only=True, - ) - - 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, - ) -> torch.Tensor: - residual = hidden_states - norm_hidden_states, gate = self.norm(hidden_states, emb=temb) - mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)) - joint_attention_kwargs = joint_attention_kwargs or {} - attn_output = self.attn( - hidden_states=norm_hidden_states, - image_rotary_emb=image_rotary_emb, - **joint_attention_kwargs, - ) - - hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) - gate = gate.unsqueeze(1) - hidden_states = gate * self.proj_out(hidden_states) - hidden_states = residual + hidden_states - if hidden_states.dtype == torch.float16: - hidden_states = hidden_states.clip(-65504, 65504) - - return hidden_states - - -@maybe_allow_in_graph -class FluxTransformerBlock(nn.Module): - def __init__( - self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6 - ): - super().__init__() - - self.norm1 = AdaLayerNormZero(dim) - self.norm1_context = AdaLayerNormZero(dim) - - self.attn = Attention( - query_dim=dim, - cross_attention_dim=None, - added_kv_proj_dim=dim, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=dim, - context_pre_only=False, - bias=True, - processor=FluxAttnProcessor2_0(), - qk_norm=qk_norm, - eps=eps, - ) - - self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) - self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") - - self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) - self.ff_context = 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, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = 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( - encoder_hidden_states, emb=temb - ) - joint_attention_kwargs = joint_attention_kwargs or {} - # Attention. - attention_outputs = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - image_rotary_emb=image_rotary_emb, - **joint_attention_kwargs, - ) - - if len(attention_outputs) == 2: - attn_output, context_attn_output = attention_outputs - elif len(attention_outputs) == 3: - attn_output, context_attn_output, ip_attn_output = attention_outputs - - # Process attention outputs for the `hidden_states`. - attn_output = gate_msa.unsqueeze(1) * attn_output - hidden_states = hidden_states + attn_output - - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - - ff_output = self.ff(norm_hidden_states) - ff_output = gate_mlp.unsqueeze(1) * ff_output - - hidden_states = hidden_states + ff_output - if len(attention_outputs) == 3: - hidden_states = hidden_states + ip_attn_output - - # Process attention outputs for the `encoder_hidden_states`. - - context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output - encoder_hidden_states = encoder_hidden_states + context_attn_output - - norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) - norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] - - context_ff_output = self.ff_context(norm_encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output - if encoder_hidden_states.dtype == torch.float16: - encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) - - return encoder_hidden_states, hidden_states - - -class FluxTransformer2DModel( - ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, FluxTransformer2DLoadersMixin, CacheMixin -): - """ - The Transformer model introduced in Flux. - - Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ - - Args: - patch_size (`int`, defaults to `1`): - Patch size to turn the input data into small patches. - in_channels (`int`, defaults to `64`): - 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`). - pooled_projection_dim (`int`, defaults to `768`): - 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)`): - The dimensions to use for the rotary positional embeddings. - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["FluxTransformerBlock", "FluxSingleTransformerBlock"] - _skip_layerwise_casting_patterns = ["pos_embed", "norm"] - - @register_to_config - def __init__( - self, - patch_size: int = 1, - in_channels: int = 64, - 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, - pooled_projection_dim: int = 768, - guidance_embeds: bool = False, - axes_dims_rope: Tuple[int, int, int] = (16, 56, 56), - ): - 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) - - text_time_guidance_cls = ( - CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings - ) - self.time_text_embed = text_time_guidance_cls( - embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim - ) - - self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim) - self.x_embedder = nn.Linear(in_channels, self.inner_dim) - - self.transformer_blocks = nn.ModuleList( - [ - FluxTransformerBlock( - 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( - [ - FluxSingleTransformerBlock( - dim=self.inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - ) - for _ in range(num_single_layers) - ] - ) - - self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6) - self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True) - - self.gradient_checkpointing = False - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedFluxAttnProcessor2_0 - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - self.set_attn_processor(FusedFluxAttnProcessor2_0()) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - if self.original_attn_processors is not None: - self.set_attn_processor(self.original_attn_processors) - - def forward( - self, - hidden_states: torch.Tensor, - 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, - controlnet_block_samples=None, - controlnet_single_block_samples=None, - return_dict: bool = True, - controlnet_blocks_repeat: bool = False, - ) -> Union[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. - pooled_projections (`torch.Tensor` of shape `(batch_size, projection_dim)`): Embeddings projected - from the embeddings of input conditions. - 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. - """ - 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) * 1000 - if guidance is not None: - guidance = guidance.to(hidden_states.dtype) * 1000 - - temb = ( - self.time_text_embed(timestep, pooled_projections) - if guidance is None - else self.time_text_embed(timestep, guidance, pooled_projections) - ) - 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): - 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, - ) - - 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, - 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): - 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, - 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] :, ...] - - 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,) - - return Transformer2DModelOutput(sample=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_hidream_image.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_hidream_image.py deleted file mode 100644 index 77902dcf585282514b15734ee4183e954b5bb570..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_hidream_image.py +++ /dev/null @@ -1,942 +0,0 @@ -from typing import Any, Dict, List, Optional, Tuple, Union - -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 ...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.torch_utils import maybe_allow_in_graph -from ..attention import Attention -from ..embeddings import TimestepEmbedding, Timesteps - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class HiDreamImageFeedForwardSwiGLU(nn.Module): - def __init__( - self, - dim: int, - hidden_dim: int, - multiple_of: int = 256, - ffn_dim_multiplier: Optional[float] = None, - ): - super().__init__() - hidden_dim = int(2 * hidden_dim / 3) - # custom dim factor multiplier - if ffn_dim_multiplier is not None: - hidden_dim = int(ffn_dim_multiplier * hidden_dim) - hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) - - self.w1 = nn.Linear(dim, hidden_dim, bias=False) - self.w2 = nn.Linear(hidden_dim, dim, bias=False) - self.w3 = nn.Linear(dim, hidden_dim, bias=False) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.w2(torch.nn.functional.silu(self.w1(x)) * self.w3(x)) - - -class HiDreamImagePooledEmbed(nn.Module): - def __init__(self, text_emb_dim, hidden_size): - super().__init__() - self.pooled_embedder = TimestepEmbedding(in_channels=text_emb_dim, time_embed_dim=hidden_size) - - def forward(self, pooled_embed: torch.Tensor) -> torch.Tensor: - return self.pooled_embedder(pooled_embed) - - -class HiDreamImageTimestepEmbed(nn.Module): - def __init__(self, hidden_size, frequency_embedding_size=256): - super().__init__() - 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): - t_emb = self.time_proj(timesteps).to(dtype=wdtype) - t_emb = self.timestep_embedder(t_emb) - return t_emb - - -class HiDreamImageOutEmbed(nn.Module): - def __init__(self, hidden_size, patch_size, out_channels): - super().__init__() - self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) - self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)) - - def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: - shift, scale = self.adaLN_modulation(temb).chunk(2, dim=1) - hidden_states = self.norm_final(hidden_states) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) - hidden_states = self.linear(hidden_states) - return hidden_states - - -class HiDreamImagePatchEmbed(nn.Module): - def __init__( - self, - patch_size=2, - in_channels=4, - out_channels=1024, - ): - super().__init__() - self.patch_size = patch_size - self.out_channels = out_channels - self.proj = nn.Linear(in_channels * patch_size * patch_size, out_channels, bias=True) - - def forward(self, latent): - latent = self.proj(latent) - return latent - - -def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor: - assert dim % 2 == 0, "The dimension must be even." - - is_mps = pos.device.type == "mps" - is_npu = pos.device.type == "npu" - - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - - scale = torch.arange(0, dim, 2, dtype=dtype, device=pos.device) / dim - omega = 1.0 / (theta**scale) - - batch_size, seq_length = pos.shape - out = torch.einsum("...n,d->...nd", pos, omega) - cos_out = torch.cos(out) - sin_out = torch.sin(out) - - stacked_out = torch.stack([cos_out, -sin_out, sin_out, cos_out], dim=-1) - out = stacked_out.view(batch_size, -1, dim // 2, 2, 2) - return out.float() - - -class HiDreamImageEmbedND(nn.Module): - def __init__(self, theta: int, axes_dim: List[int]): - super().__init__() - self.theta = theta - self.axes_dim = axes_dim - - def forward(self, ids: torch.Tensor) -> torch.Tensor: - n_axes = ids.shape[-1] - emb = torch.cat( - [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], - dim=-3, - ) - return emb.unsqueeze(2) - - -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] - xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1] - return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk) - - -@maybe_allow_in_graph -class HiDreamAttention(Attention): - def __init__( - self, - query_dim: int, - heads: int = 8, - dim_head: int = 64, - upcast_attention: bool = False, - upcast_softmax: bool = False, - scale_qk: bool = True, - eps: float = 1e-5, - processor=None, - out_dim: int = None, - single: bool = False, - ): - super(Attention, self).__init__() - self.inner_dim = out_dim if out_dim is not None else dim_head * heads - self.query_dim = query_dim - self.upcast_attention = upcast_attention - self.upcast_softmax = upcast_softmax - self.out_dim = out_dim if out_dim is not None else query_dim - - self.scale_qk = scale_qk - self.scale = dim_head**-0.5 if self.scale_qk else 1.0 - - self.heads = out_dim // dim_head if out_dim is not None else heads - self.sliceable_head_dim = heads - self.single = single - - self.to_q = nn.Linear(query_dim, self.inner_dim) - self.to_k = nn.Linear(self.inner_dim, self.inner_dim) - self.to_v = nn.Linear(self.inner_dim, self.inner_dim) - self.to_out = nn.Linear(self.inner_dim, self.out_dim) - self.q_rms_norm = nn.RMSNorm(self.inner_dim, eps) - self.k_rms_norm = nn.RMSNorm(self.inner_dim, eps) - - if not single: - self.to_q_t = nn.Linear(query_dim, self.inner_dim) - self.to_k_t = nn.Linear(self.inner_dim, self.inner_dim) - self.to_v_t = nn.Linear(self.inner_dim, self.inner_dim) - self.to_out_t = nn.Linear(self.inner_dim, self.out_dim) - self.q_rms_norm_t = nn.RMSNorm(self.inner_dim, eps) - self.k_rms_norm_t = nn.RMSNorm(self.inner_dim, eps) - - self.set_processor(processor) - - def forward( - self, - norm_hidden_states: torch.Tensor, - hidden_states_masks: torch.Tensor = None, - norm_encoder_hidden_states: torch.Tensor = None, - image_rotary_emb: torch.Tensor = None, - ) -> torch.Tensor: - return self.processor( - self, - hidden_states=norm_hidden_states, - hidden_states_masks=hidden_states_masks, - encoder_hidden_states=norm_encoder_hidden_states, - image_rotary_emb=image_rotary_emb, - ) - - -class HiDreamAttnProcessor: - """Attention processor used typically in processing the SD3-like self-attention projections.""" - - def __call__( - self, - attn: HiDreamAttention, - hidden_states: torch.Tensor, - hidden_states_masks: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - image_rotary_emb: torch.Tensor = None, - *args, - **kwargs, - ) -> torch.Tensor: - dtype = hidden_states.dtype - batch_size = hidden_states.shape[0] - - query_i = attn.q_rms_norm(attn.to_q(hidden_states)).to(dtype=dtype) - key_i = attn.k_rms_norm(attn.to_k(hidden_states)).to(dtype=dtype) - value_i = attn.to_v(hidden_states) - - inner_dim = key_i.shape[-1] - head_dim = inner_dim // attn.heads - - query_i = query_i.view(batch_size, -1, attn.heads, head_dim) - key_i = key_i.view(batch_size, -1, attn.heads, head_dim) - value_i = value_i.view(batch_size, -1, attn.heads, head_dim) - if hidden_states_masks is not None: - key_i = key_i * hidden_states_masks.view(batch_size, -1, 1, 1) - - if not attn.single: - query_t = attn.q_rms_norm_t(attn.to_q_t(encoder_hidden_states)).to(dtype=dtype) - key_t = attn.k_rms_norm_t(attn.to_k_t(encoder_hidden_states)).to(dtype=dtype) - value_t = attn.to_v_t(encoder_hidden_states) - - query_t = query_t.view(batch_size, -1, attn.heads, head_dim) - key_t = key_t.view(batch_size, -1, attn.heads, head_dim) - value_t = value_t.view(batch_size, -1, attn.heads, head_dim) - - num_image_tokens = query_i.shape[1] - num_text_tokens = query_t.shape[1] - query = torch.cat([query_i, query_t], dim=1) - key = torch.cat([key_i, key_t], dim=1) - value = torch.cat([value_i, value_t], dim=1) - else: - query = query_i - key = key_i - value = value_i - - if query.shape[-1] == image_rotary_emb.shape[-3] * 2: - query, key = apply_rope(query, key, image_rotary_emb) - - else: - query_1, query_2 = query.chunk(2, dim=-1) - key_1, key_2 = key.chunk(2, dim=-1) - query_1, key_1 = apply_rope(query_1, key_1, image_rotary_emb) - query = torch.cat([query_1, query_2], dim=-1) - key = torch.cat([key_1, key_2], dim=-1) - - hidden_states = F.scaled_dot_product_attention( - query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - if not attn.single: - hidden_states_i, hidden_states_t = torch.split(hidden_states, [num_image_tokens, num_text_tokens], dim=1) - hidden_states_i = attn.to_out(hidden_states_i) - hidden_states_t = attn.to_out_t(hidden_states_t) - return hidden_states_i, hidden_states_t - else: - hidden_states = attn.to_out(hidden_states) - return hidden_states - - -# Modified from https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py -class MoEGate(nn.Module): - def __init__( - self, - embed_dim, - num_routed_experts=4, - num_activated_experts=2, - aux_loss_alpha=0.01, - _force_inference_output=False, - ): - super().__init__() - self.top_k = num_activated_experts - self.n_routed_experts = num_routed_experts - - self.scoring_func = "softmax" - self.alpha = aux_loss_alpha - self.seq_aux = False - - # topk selection algorithm - self.norm_topk_prob = False - self.gating_dim = embed_dim - self.weight = nn.Parameter(torch.randn(self.n_routed_experts, self.gating_dim) / embed_dim**0.5) - - self._force_inference_output = _force_inference_output - - def forward(self, hidden_states): - bsz, seq_len, h = hidden_states.shape - ### compute gating score - hidden_states = hidden_states.view(-1, h) - logits = F.linear(hidden_states, self.weight, None) - if self.scoring_func == "softmax": - scores = logits.softmax(dim=-1) - else: - raise NotImplementedError(f"insupportable scoring function for MoE gating: {self.scoring_func}") - - ### select top-k experts - topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False) - - ### norm gate to sum 1 - if self.top_k > 1 and self.norm_topk_prob: - denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 - topk_weight = topk_weight / denominator - - ### expert-level computation auxiliary loss - if self.training and self.alpha > 0.0 and not self._force_inference_output: - scores_for_aux = scores - aux_topk = self.top_k - # always compute aux loss based on the naive greedy topk method - topk_idx_for_aux_loss = topk_idx.view(bsz, -1) - if self.seq_aux: - scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1) - ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device) - ce.scatter_add_( - 1, topk_idx_for_aux_loss, torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device) - ).div_(seq_len * aux_topk / self.n_routed_experts) - aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha - else: - mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts) - ce = mask_ce.float().mean(0) - - Pi = scores_for_aux.mean(0) - fi = ce * self.n_routed_experts - aux_loss = (Pi * fi).sum() * self.alpha - else: - aux_loss = None - return topk_idx, topk_weight, aux_loss - - -# Modified from https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py -class MOEFeedForwardSwiGLU(nn.Module): - def __init__( - self, - dim: int, - hidden_dim: int, - num_routed_experts: int, - num_activated_experts: int, - _force_inference_output: bool = False, - ): - super().__init__() - self.shared_experts = HiDreamImageFeedForwardSwiGLU(dim, hidden_dim // 2) - self.experts = nn.ModuleList( - [HiDreamImageFeedForwardSwiGLU(dim, hidden_dim) for i in range(num_routed_experts)] - ) - self._force_inference_output = _force_inference_output - self.gate = MoEGate( - embed_dim=dim, - num_routed_experts=num_routed_experts, - num_activated_experts=num_activated_experts, - _force_inference_output=_force_inference_output, - ) - self.num_activated_experts = num_activated_experts - - def forward(self, x): - wtype = x.dtype - identity = x - orig_shape = x.shape - topk_idx, topk_weight, aux_loss = self.gate(x) - x = x.view(-1, x.shape[-1]) - flat_topk_idx = topk_idx.view(-1) - if self.training and not self._force_inference_output: - x = x.repeat_interleave(self.num_activated_experts, dim=0) - y = torch.empty_like(x, dtype=wtype) - for i, expert in enumerate(self.experts): - y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(dtype=wtype) - y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1) - y = y.view(*orig_shape).to(dtype=wtype) - # y = AddAuxiliaryLoss.apply(y, aux_loss) - else: - y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape) - y = y + self.shared_experts(identity) - return y - - @torch.no_grad() - def moe_infer(self, x, flat_expert_indices, flat_expert_weights): - expert_cache = torch.zeros_like(x) - idxs = flat_expert_indices.argsort() - tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0) - token_idxs = idxs // self.num_activated_experts - for i, end_idx in enumerate(tokens_per_expert): - start_idx = 0 if i == 0 else tokens_per_expert[i - 1] - if start_idx == end_idx: - continue - expert = self.experts[i] - exp_token_idx = token_idxs[start_idx:end_idx] - expert_tokens = x[exp_token_idx] - expert_out = expert(expert_tokens) - expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]]) - - # for fp16 and other dtype - expert_cache = expert_cache.to(expert_out.dtype) - expert_cache.scatter_reduce_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out, reduce="sum") - return expert_cache - - -class TextProjection(nn.Module): - def __init__(self, in_features, hidden_size): - super().__init__() - self.linear = nn.Linear(in_features=in_features, out_features=hidden_size, bias=False) - - def forward(self, caption): - hidden_states = self.linear(caption) - return hidden_states - - -@maybe_allow_in_graph -class HiDreamImageSingleTransformerBlock(nn.Module): - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - num_routed_experts: int = 4, - num_activated_experts: int = 2, - _force_inference_output: bool = False, - ): - super().__init__() - self.num_attention_heads = num_attention_heads - self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True)) - - # 1. Attention - self.norm1_i = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False) - self.attn1 = HiDreamAttention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - processor=HiDreamAttnProcessor(), - single=True, - ) - - # 3. Feed-forward - self.norm3_i = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False) - if num_routed_experts > 0: - self.ff_i = MOEFeedForwardSwiGLU( - dim=dim, - hidden_dim=4 * dim, - num_routed_experts=num_routed_experts, - num_activated_experts=num_activated_experts, - _force_inference_output=_force_inference_output, - ) - else: - self.ff_i = HiDreamImageFeedForwardSwiGLU(dim=dim, hidden_dim=4 * dim) - - 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, - image_rotary_emb: torch.Tensor = None, - ) -> torch.Tensor: - wtype = hidden_states.dtype - shift_msa_i, scale_msa_i, gate_msa_i, shift_mlp_i, scale_mlp_i, gate_mlp_i = self.adaLN_modulation(temb)[ - :, None - ].chunk(6, dim=-1) - - # 1. MM-Attention - norm_hidden_states = self.norm1_i(hidden_states).to(dtype=wtype) - norm_hidden_states = norm_hidden_states * (1 + scale_msa_i) + shift_msa_i - attn_output_i = self.attn1( - norm_hidden_states, - hidden_states_masks, - image_rotary_emb=image_rotary_emb, - ) - hidden_states = gate_msa_i * attn_output_i + hidden_states - - # 2. Feed-forward - norm_hidden_states = self.norm3_i(hidden_states).to(dtype=wtype) - norm_hidden_states = norm_hidden_states * (1 + scale_mlp_i) + shift_mlp_i - ff_output_i = gate_mlp_i * self.ff_i(norm_hidden_states.to(dtype=wtype)) - hidden_states = ff_output_i + hidden_states - return hidden_states - - -@maybe_allow_in_graph -class HiDreamImageTransformerBlock(nn.Module): - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - num_routed_experts: int = 4, - num_activated_experts: int = 2, - _force_inference_output: bool = False, - ): - super().__init__() - self.num_attention_heads = num_attention_heads - self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 12 * dim, bias=True)) - - # 1. Attention - self.norm1_i = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False) - self.norm1_t = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False) - self.attn1 = HiDreamAttention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - processor=HiDreamAttnProcessor(), - single=False, - ) - - # 3. Feed-forward - self.norm3_i = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False) - if num_routed_experts > 0: - self.ff_i = MOEFeedForwardSwiGLU( - dim=dim, - hidden_dim=4 * dim, - num_routed_experts=num_routed_experts, - num_activated_experts=num_activated_experts, - _force_inference_output=_force_inference_output, - ) - else: - self.ff_i = HiDreamImageFeedForwardSwiGLU(dim=dim, hidden_dim=4 * dim) - self.norm3_t = nn.LayerNorm(dim, eps=1e-06, elementwise_affine=False) - self.ff_t = HiDreamImageFeedForwardSwiGLU(dim=dim, hidden_dim=4 * dim) - - 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, - image_rotary_emb: torch.Tensor = None, - ) -> torch.Tensor: - wtype = hidden_states.dtype - ( - shift_msa_i, - scale_msa_i, - gate_msa_i, - shift_mlp_i, - scale_mlp_i, - gate_mlp_i, - shift_msa_t, - scale_msa_t, - gate_msa_t, - shift_mlp_t, - scale_mlp_t, - gate_mlp_t, - ) = self.adaLN_modulation(temb)[:, None].chunk(12, dim=-1) - - # 1. MM-Attention - norm_hidden_states = self.norm1_i(hidden_states).to(dtype=wtype) - norm_hidden_states = norm_hidden_states * (1 + scale_msa_i) + shift_msa_i - norm_encoder_hidden_states = self.norm1_t(encoder_hidden_states).to(dtype=wtype) - norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + scale_msa_t) + shift_msa_t - - attn_output_i, attn_output_t = self.attn1( - norm_hidden_states, - hidden_states_masks, - norm_encoder_hidden_states, - image_rotary_emb=image_rotary_emb, - ) - - hidden_states = gate_msa_i * attn_output_i + hidden_states - encoder_hidden_states = gate_msa_t * attn_output_t + encoder_hidden_states - - # 2. Feed-forward - norm_hidden_states = self.norm3_i(hidden_states).to(dtype=wtype) - norm_hidden_states = norm_hidden_states * (1 + scale_mlp_i) + shift_mlp_i - norm_encoder_hidden_states = self.norm3_t(encoder_hidden_states).to(dtype=wtype) - norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + scale_mlp_t) + shift_mlp_t - - ff_output_i = gate_mlp_i * self.ff_i(norm_hidden_states) - ff_output_t = gate_mlp_t * self.ff_t(norm_encoder_hidden_states) - hidden_states = ff_output_i + hidden_states - encoder_hidden_states = ff_output_t + encoder_hidden_states - return hidden_states, encoder_hidden_states - - -class HiDreamBlock(nn.Module): - def __init__(self, block: Union[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, - image_rotary_emb: torch.Tensor = None, - ) -> torch.Tensor: - return self.block( - hidden_states=hidden_states, - hidden_states_masks=hidden_states_masks, - encoder_hidden_states=encoder_hidden_states, - temb=temb, - image_rotary_emb=image_rotary_emb, - ) - - -class HiDreamImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): - _supports_gradient_checkpointing = True - _no_split_modules = ["HiDreamImageTransformerBlock", "HiDreamImageSingleTransformerBlock"] - - @register_to_config - def __init__( - self, - patch_size: Optional[int] = None, - in_channels: int = 64, - out_channels: Optional[int] = 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, - 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, - force_inference_output: bool = False, - ): - super().__init__() - self.out_channels = out_channels or in_channels - self.inner_dim = num_attention_heads * attention_head_dim - - self.t_embedder = HiDreamImageTimestepEmbed(self.inner_dim) - self.p_embedder = HiDreamImagePooledEmbed(text_emb_dim, self.inner_dim) - self.x_embedder = HiDreamImagePatchEmbed( - patch_size=patch_size, - in_channels=in_channels, - out_channels=self.inner_dim, - ) - self.pe_embedder = HiDreamImageEmbedND(theta=10000, axes_dim=axes_dims_rope) - - self.double_stream_blocks = nn.ModuleList( - [ - HiDreamBlock( - HiDreamImageTransformerBlock( - dim=self.inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - num_routed_experts=num_routed_experts, - num_activated_experts=num_activated_experts, - _force_inference_output=force_inference_output, - ) - ) - for _ in range(num_layers) - ] - ) - - self.single_stream_blocks = nn.ModuleList( - [ - HiDreamBlock( - HiDreamImageSingleTransformerBlock( - dim=self.inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - num_routed_experts=num_routed_experts, - num_activated_experts=num_activated_experts, - _force_inference_output=force_inference_output, - ) - ) - for _ in range(num_single_layers) - ] - ) - - self.final_layer = HiDreamImageOutEmbed(self.inner_dim, patch_size, self.out_channels) - - caption_channels = [caption_channels[1]] * (num_layers + num_single_layers) + [caption_channels[0]] - caption_projection = [] - for caption_channel in caption_channels: - caption_projection.append(TextProjection(in_features=caption_channel, hidden_size=self.inner_dim)) - self.caption_projection = nn.ModuleList(caption_projection) - self.max_seq = max_resolution[0] * max_resolution[1] // (patch_size * patch_size) - - self.gradient_checkpointing = False - - 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) - x = ( - x.reshape(B, S, self.config.patch_size, self.config.patch_size, C) - .permute(0, 4, 1, 2, 3) - .reshape(B, C, S, self.config.patch_size * self.config.patch_size) - ) - else: - x_arr = [] - p1 = self.config.patch_size - p2 = self.config.patch_size - for i, img_size in enumerate(img_sizes): - pH, pW = img_size - t = x[i, : pH * pW].reshape(1, pH, pW, -1) - F_token = t.shape[-1] - C = F_token // (p1 * p2) - t = t.reshape(1, pH, pW, p1, p2, C) - t = t.permute(0, 5, 1, 3, 2, 4) - t = t.reshape(1, C, pH * p1, pW * p2) - x_arr.append(t) - x = torch.cat(x_arr, dim=0) - return x - - def patchify(self, hidden_states): - batch_size, channels, height, width = hidden_states.shape - patch_size = self.config.patch_size - patch_height, patch_width = height // patch_size, width // patch_size - device = hidden_states.device - dtype = hidden_states.dtype - - # create img_sizes - img_sizes = torch.tensor([patch_height, patch_width], dtype=torch.int64, device=device).reshape(-1) - img_sizes = img_sizes.unsqueeze(0).repeat(batch_size, 1) - - # create hidden_states_masks - if hidden_states.shape[-2] != hidden_states.shape[-1]: - hidden_states_masks = torch.zeros((batch_size, self.max_seq), dtype=dtype, device=device) - hidden_states_masks[:, : patch_height * patch_width] = 1.0 - else: - hidden_states_masks = None - - # create img_ids - img_ids = torch.zeros(patch_height, patch_width, 3, device=device) - row_indices = torch.arange(patch_height, device=device)[:, None] - col_indices = torch.arange(patch_width, device=device)[None, :] - img_ids[..., 1] = img_ids[..., 1] + row_indices - img_ids[..., 2] = img_ids[..., 2] + col_indices - img_ids = img_ids.reshape(patch_height * patch_width, -1) - - if hidden_states.shape[-2] != hidden_states.shape[-1]: - # Handle non-square latents - img_ids_pad = torch.zeros(self.max_seq, 3, device=device) - img_ids_pad[: patch_height * patch_width, :] = img_ids - img_ids = img_ids_pad.unsqueeze(0).repeat(batch_size, 1, 1) - else: - img_ids = img_ids.unsqueeze(0).repeat(batch_size, 1, 1) - - # patchify hidden_states - if hidden_states.shape[-2] != hidden_states.shape[-1]: - # Handle non-square latents - out = torch.zeros( - (batch_size, channels, self.max_seq, patch_size * patch_size), - dtype=dtype, - device=device, - ) - hidden_states = hidden_states.reshape( - batch_size, channels, patch_height, patch_size, patch_width, patch_size - ) - hidden_states = hidden_states.permute(0, 1, 2, 4, 3, 5) - hidden_states = hidden_states.reshape( - batch_size, channels, patch_height * patch_width, patch_size * patch_size - ) - out[:, :, 0 : patch_height * patch_width] = hidden_states - hidden_states = out - hidden_states = hidden_states.permute(0, 2, 3, 1).reshape( - batch_size, self.max_seq, patch_size * patch_size * channels - ) - - else: - # Handle square latents - hidden_states = hidden_states.reshape( - batch_size, channels, patch_height, patch_size, patch_width, patch_size - ) - hidden_states = hidden_states.permute(0, 2, 4, 3, 5, 1) - hidden_states = hidden_states.reshape( - batch_size, patch_height * patch_width, patch_size * patch_size * channels - ) - - return hidden_states, hidden_states_masks, img_sizes, img_ids - - def forward( - self, - hidden_states: torch.Tensor, - timesteps: torch.LongTensor = None, - 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, - return_dict: bool = True, - **kwargs, - ): - encoder_hidden_states = kwargs.get("encoder_hidden_states", None) - - if encoder_hidden_states is not None: - deprecation_message = "The `encoder_hidden_states` argument is deprecated. Please use `encoder_hidden_states_t5` and `encoder_hidden_states_llama3` instead." - deprecate("encoder_hidden_states", "0.35.0", deprecation_message) - encoder_hidden_states_t5 = encoder_hidden_states[0] - encoder_hidden_states_llama3 = encoder_hidden_states[1] - - if img_ids is not None and img_sizes is not None and hidden_states_masks is None: - deprecation_message = ( - "Passing `img_ids` and `img_sizes` with unpachified `hidden_states` is deprecated and will be ignored." - ) - deprecate("img_ids", "0.35.0", deprecation_message) - - if hidden_states_masks is not None and (img_ids is None or img_sizes is None): - raise ValueError("if `hidden_states_masks` is passed, `img_ids` and `img_sizes` must also be passed.") - elif hidden_states_masks is not None and hidden_states.ndim != 3: - raise ValueError( - "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 - - # Patchify the input - if hidden_states_masks is None: - hidden_states, hidden_states_masks, img_sizes, img_ids = self.patchify(hidden_states) - - # Embed the hidden states - hidden_states = self.x_embedder(hidden_states) - - # 0. time - timesteps = self.t_embedder(timesteps, hidden_states_type) - p_embedder = self.p_embedder(pooled_embeds) - temb = timesteps + p_embedder - - encoder_hidden_states = [encoder_hidden_states_llama3[k] for k in self.config.llama_layers] - - if self.caption_projection is not None: - new_encoder_hidden_states = [] - for i, enc_hidden_state in enumerate(encoder_hidden_states): - enc_hidden_state = self.caption_projection[i](enc_hidden_state) - enc_hidden_state = enc_hidden_state.view(batch_size, -1, hidden_states.shape[-1]) - new_encoder_hidden_states.append(enc_hidden_state) - encoder_hidden_states = new_encoder_hidden_states - encoder_hidden_states_t5 = self.caption_projection[-1](encoder_hidden_states_t5) - encoder_hidden_states_t5 = encoder_hidden_states_t5.view(batch_size, -1, hidden_states.shape[-1]) - encoder_hidden_states.append(encoder_hidden_states_t5) - - txt_ids = torch.zeros( - batch_size, - encoder_hidden_states[-1].shape[1] - + encoder_hidden_states[-2].shape[1] - + encoder_hidden_states[0].shape[1], - 3, - device=img_ids.device, - dtype=img_ids.dtype, - ) - ids = torch.cat((img_ids, txt_ids), dim=1) - image_rotary_emb = self.pe_embedder(ids) - - # 2. Blocks - block_id = 0 - initial_encoder_hidden_states = torch.cat([encoder_hidden_states[-1], encoder_hidden_states[-2]], dim=1) - initial_encoder_hidden_states_seq_len = initial_encoder_hidden_states.shape[1] - for bid, block in enumerate(self.double_stream_blocks): - cur_llama31_encoder_hidden_states = encoder_hidden_states[block_id] - cur_encoder_hidden_states = torch.cat( - [initial_encoder_hidden_states, cur_llama31_encoder_hidden_states], dim=1 - ) - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, initial_encoder_hidden_states = self._gradient_checkpointing_func( - block, - hidden_states, - hidden_states_masks, - cur_encoder_hidden_states, - temb, - image_rotary_emb, - ) - else: - hidden_states, initial_encoder_hidden_states = block( - hidden_states=hidden_states, - hidden_states_masks=hidden_states_masks, - encoder_hidden_states=cur_encoder_hidden_states, - temb=temb, - image_rotary_emb=image_rotary_emb, - ) - initial_encoder_hidden_states = initial_encoder_hidden_states[:, :initial_encoder_hidden_states_seq_len] - block_id += 1 - - image_tokens_seq_len = hidden_states.shape[1] - hidden_states = torch.cat([hidden_states, initial_encoder_hidden_states], dim=1) - hidden_states_seq_len = hidden_states.shape[1] - if hidden_states_masks is not None: - encoder_attention_mask_ones = torch.ones( - (batch_size, initial_encoder_hidden_states.shape[1] + cur_llama31_encoder_hidden_states.shape[1]), - device=hidden_states_masks.device, - dtype=hidden_states_masks.dtype, - ) - hidden_states_masks = torch.cat([hidden_states_masks, encoder_attention_mask_ones], dim=1) - - for bid, block in enumerate(self.single_stream_blocks): - cur_llama31_encoder_hidden_states = encoder_hidden_states[block_id] - hidden_states = torch.cat([hidden_states, cur_llama31_encoder_hidden_states], dim=1) - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - block, - hidden_states, - hidden_states_masks, - None, - temb, - image_rotary_emb, - ) - else: - hidden_states = block( - hidden_states=hidden_states, - hidden_states_masks=hidden_states_masks, - encoder_hidden_states=None, - temb=temb, - image_rotary_emb=image_rotary_emb, - ) - hidden_states = hidden_states[:, :hidden_states_seq_len] - block_id += 1 - - hidden_states = hidden_states[:, :image_tokens_seq_len, ...] - output = self.final_layer(hidden_states, temb) - output = self.unpatchify(output, img_sizes, self.training) - 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/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_hunyuan_video.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_hunyuan_video.py deleted file mode 100644 index c48c586a28def7ae34e7321193b0101b779749dd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_hunyuan_video.py +++ /dev/null @@ -1,1147 +0,0 @@ -# Copyright 2025 The Hunyuan 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 Any, Dict, List, Optional, Tuple, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from diffusers.loaders import FromOriginalModelMixin - -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 ..attention import FeedForward -from ..attention_processor import Attention, AttentionProcessor -from ..cache_utils import CacheMixin -from ..embeddings import ( - CombinedTimestepTextProjEmbeddings, - PixArtAlphaTextProjection, - TimestepEmbedding, - Timesteps, - get_1d_rotary_pos_embed, -) -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle, FP32LayerNorm - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class HunyuanVideoAttnProcessor2_0: - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "HunyuanVideoAttnProcessor2_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, - ) -> 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) - - # 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)).transpose(1, 2) - key = key.unflatten(2, (attn.heads, -1)).transpose(1, 2) - value = value.unflatten(2, (attn.heads, -1)).transpose(1, 2) - - # 2. QK normalization - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # 3. Rotational positional embeddings applied to latent stream - if image_rotary_emb is not None: - from ..embeddings import apply_rotary_emb - - if attn.add_q_proj is None and encoder_hidden_states is not None: - query = torch.cat( - [ - apply_rotary_emb(query[:, :, : -encoder_hidden_states.shape[1]], image_rotary_emb), - query[:, :, -encoder_hidden_states.shape[1] :], - ], - dim=2, - ) - key = torch.cat( - [ - apply_rotary_emb(key[:, :, : -encoder_hidden_states.shape[1]], image_rotary_emb), - key[:, :, -encoder_hidden_states.shape[1] :], - ], - dim=2, - ) - else: - query = apply_rotary_emb(query, image_rotary_emb) - key = apply_rotary_emb(key, image_rotary_emb) - - # 4. Encoder condition QKV projection and normalization - if attn.add_q_proj is not None and encoder_hidden_states is not None: - encoder_query = attn.add_q_proj(encoder_hidden_states) - encoder_key = attn.add_k_proj(encoder_hidden_states) - encoder_value = attn.add_v_proj(encoder_hidden_states) - - encoder_query = encoder_query.unflatten(2, (attn.heads, -1)).transpose(1, 2) - encoder_key = encoder_key.unflatten(2, (attn.heads, -1)).transpose(1, 2) - encoder_value = encoder_value.unflatten(2, (attn.heads, -1)).transpose(1, 2) - - if attn.norm_added_q is not None: - encoder_query = attn.norm_added_q(encoder_query) - if attn.norm_added_k is not None: - encoder_key = attn.norm_added_k(encoder_key) - - query = torch.cat([query, encoder_query], dim=2) - key = torch.cat([key, encoder_key], dim=2) - value = torch.cat([value, encoder_value], dim=2) - - # 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 = hidden_states.transpose(1, 2).flatten(2, 3) - hidden_states = hidden_states.to(query.dtype) - - # 6. Output projection - if encoder_hidden_states is not None: - hidden_states, encoder_hidden_states = ( - hidden_states[:, : -encoder_hidden_states.shape[1]], - hidden_states[:, -encoder_hidden_states.shape[1] :], - ) - - if getattr(attn, "to_out", None) is not None: - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - - if getattr(attn, "to_add_out", None) is not None: - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) - - return hidden_states, encoder_hidden_states - - -class HunyuanVideoPatchEmbed(nn.Module): - def __init__( - self, - patch_size: Union[int, Tuple[int, int, int]] = 16, - in_chans: int = 3, - embed_dim: int = 768, - ) -> None: - super().__init__() - - patch_size = (patch_size, patch_size, patch_size) if isinstance(patch_size, int) else patch_size - self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.proj(hidden_states) - hidden_states = hidden_states.flatten(2).transpose(1, 2) # BCFHW -> BNC - return hidden_states - - -class HunyuanVideoAdaNorm(nn.Module): - def __init__(self, in_features: int, out_features: Optional[int] = None) -> None: - super().__init__() - - out_features = out_features or 2 * in_features - self.linear = nn.Linear(in_features, out_features) - self.nonlinearity = nn.SiLU() - - def forward( - self, temb: 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) - return gate_msa, gate_mlp - - -class HunyuanVideoTokenReplaceAdaLayerNormZero(nn.Module): - def __init__(self, embedding_dim: int, norm_type: str = "layer_norm", bias: bool = True): - super().__init__() - - self.silu = nn.SiLU() - self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=bias) - - if norm_type == "layer_norm": - self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) - elif norm_type == "fp32_layer_norm": - self.norm = FP32LayerNorm(embedding_dim, elementwise_affine=False, bias=False) - else: - raise ValueError( - f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'." - ) - - def forward( - self, - hidden_states: torch.Tensor, - emb: torch.Tensor, - token_replace_emb: torch.Tensor, - first_frame_num_tokens: int, - ) -> 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)) - - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1) - tr_shift_msa, tr_scale_msa, tr_gate_msa, tr_shift_mlp, tr_scale_mlp, tr_gate_mlp = token_replace_emb.chunk( - 6, dim=1 - ) - - norm_hidden_states = self.norm(hidden_states) - hidden_states_zero = ( - norm_hidden_states[:, :first_frame_num_tokens] * (1 + tr_scale_msa[:, None]) + tr_shift_msa[:, None] - ) - hidden_states_orig = ( - norm_hidden_states[:, first_frame_num_tokens:] * (1 + scale_msa[:, None]) + shift_msa[:, None] - ) - hidden_states = torch.cat([hidden_states_zero, hidden_states_orig], dim=1) - - return ( - hidden_states, - gate_msa, - shift_mlp, - scale_mlp, - gate_mlp, - tr_gate_msa, - tr_shift_mlp, - tr_scale_mlp, - tr_gate_mlp, - ) - - -class HunyuanVideoTokenReplaceAdaLayerNormZeroSingle(nn.Module): - def __init__(self, embedding_dim: int, norm_type: str = "layer_norm", bias: bool = True): - super().__init__() - - self.silu = nn.SiLU() - self.linear = nn.Linear(embedding_dim, 3 * embedding_dim, bias=bias) - - if norm_type == "layer_norm": - self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) - else: - raise ValueError( - f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'." - ) - - def forward( - self, - hidden_states: torch.Tensor, - emb: torch.Tensor, - token_replace_emb: torch.Tensor, - first_frame_num_tokens: int, - ) -> 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)) - - shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1) - tr_shift_msa, tr_scale_msa, tr_gate_msa = token_replace_emb.chunk(3, dim=1) - - norm_hidden_states = self.norm(hidden_states) - hidden_states_zero = ( - norm_hidden_states[:, :first_frame_num_tokens] * (1 + tr_scale_msa[:, None]) + tr_shift_msa[:, None] - ) - hidden_states_orig = ( - norm_hidden_states[:, first_frame_num_tokens:] * (1 + scale_msa[:, None]) + shift_msa[:, None] - ) - hidden_states = torch.cat([hidden_states_zero, hidden_states_orig], dim=1) - - return hidden_states, gate_msa, tr_gate_msa - - -class HunyuanVideoConditionEmbedding(nn.Module): - def __init__( - self, - embedding_dim: int, - pooled_projection_dim: int, - guidance_embeds: bool, - image_condition_type: Optional[str] = None, - ): - super().__init__() - - self.image_condition_type = image_condition_type - - self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) - self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) - self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu") - - self.guidance_embedder = None - if guidance_embeds: - 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]: - 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": - token_replace_timestep = torch.zeros_like(timestep) - token_replace_proj = self.time_proj(token_replace_timestep) - token_replace_emb = self.timestep_embedder(token_replace_proj.to(dtype=pooled_projection.dtype)) - token_replace_emb = token_replace_emb + pooled_projections - - 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 - - return conditioning, token_replace_emb - - -class HunyuanVideoIndividualTokenRefinerBlock(nn.Module): - def __init__( - self, - num_attention_heads: int, - attention_head_dim: int, - mlp_width_ratio: str = 4.0, - mlp_drop_rate: float = 0.0, - attention_bias: bool = True, - ) -> None: - super().__init__() - - hidden_size = num_attention_heads * attention_head_dim - - self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6) - self.attn = Attention( - query_dim=hidden_size, - cross_attention_dim=None, - heads=num_attention_heads, - dim_head=attention_head_dim, - bias=attention_bias, - ) - - self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6) - self.ff = FeedForward(hidden_size, mult=mlp_width_ratio, activation_fn="linear-silu", dropout=mlp_drop_rate) - - self.norm_out = HunyuanVideoAdaNorm(hidden_size, 2 * hidden_size) - - def forward( - self, - hidden_states: torch.Tensor, - temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - norm_hidden_states = self.norm1(hidden_states) - - attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=None, - attention_mask=attention_mask, - ) - - gate_msa, gate_mlp = self.norm_out(temb) - hidden_states = hidden_states + attn_output * gate_msa - - ff_output = self.ff(self.norm2(hidden_states)) - hidden_states = hidden_states + ff_output * gate_mlp - - return hidden_states - - -class HunyuanVideoIndividualTokenRefiner(nn.Module): - def __init__( - self, - num_attention_heads: int, - attention_head_dim: int, - num_layers: int, - mlp_width_ratio: float = 4.0, - mlp_drop_rate: float = 0.0, - attention_bias: bool = True, - ) -> None: - super().__init__() - - self.refiner_blocks = nn.ModuleList( - [ - HunyuanVideoIndividualTokenRefinerBlock( - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - mlp_width_ratio=mlp_width_ratio, - mlp_drop_rate=mlp_drop_rate, - attention_bias=attention_bias, - ) - for _ in range(num_layers) - ] - ) - - def forward( - self, - hidden_states: torch.Tensor, - temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - ) -> None: - self_attn_mask = None - if attention_mask is not None: - batch_size = attention_mask.shape[0] - seq_len = attention_mask.shape[1] - attention_mask = attention_mask.to(hidden_states.device).bool() - self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1) - self_attn_mask_2 = self_attn_mask_1.transpose(2, 3) - self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool() - self_attn_mask[:, :, :, 0] = True - - for block in self.refiner_blocks: - hidden_states = block(hidden_states, temb, self_attn_mask) - - return hidden_states - - -class HunyuanVideoTokenRefiner(nn.Module): - def __init__( - self, - in_channels: int, - num_attention_heads: int, - attention_head_dim: int, - num_layers: int, - mlp_ratio: float = 4.0, - mlp_drop_rate: float = 0.0, - attention_bias: bool = True, - ) -> None: - super().__init__() - - hidden_size = num_attention_heads * attention_head_dim - - self.time_text_embed = CombinedTimestepTextProjEmbeddings( - embedding_dim=hidden_size, pooled_projection_dim=in_channels - ) - self.proj_in = nn.Linear(in_channels, hidden_size, bias=True) - self.token_refiner = HunyuanVideoIndividualTokenRefiner( - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - num_layers=num_layers, - mlp_width_ratio=mlp_ratio, - mlp_drop_rate=mlp_drop_rate, - attention_bias=attention_bias, - ) - - def forward( - self, - hidden_states: torch.Tensor, - timestep: torch.LongTensor, - attention_mask: Optional[torch.LongTensor] = None, - ) -> torch.Tensor: - if attention_mask is None: - pooled_projections = hidden_states.mean(dim=1) - else: - original_dtype = hidden_states.dtype - mask_float = attention_mask.float().unsqueeze(-1) - pooled_projections = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1) - pooled_projections = pooled_projections.to(original_dtype) - - temb = self.time_text_embed(timestep, pooled_projections) - hidden_states = self.proj_in(hidden_states) - hidden_states = self.token_refiner(hidden_states, temb, attention_mask) - - return hidden_states - - -class HunyuanVideoRotaryPosEmbed(nn.Module): - 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 - self.patch_size_t = patch_size_t - self.rope_dim = rope_dim - self.theta = theta - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - batch_size, num_channels, num_frames, height, width = hidden_states.shape - rope_sizes = [num_frames // self.patch_size_t, height // self.patch_size, width // self.patch_size] - - axes_grids = [] - for i in range(3): - # Note: The following line diverges from original behaviour. We create the grid on the device, whereas - # original implementation creates it on CPU and then moves it to device. This results in numerical - # differences in layerwise debugging outputs, but visually it is the same. - grid = torch.arange(0, rope_sizes[i], device=hidden_states.device, dtype=torch.float32) - axes_grids.append(grid) - grid = torch.meshgrid(*axes_grids, indexing="ij") # [W, H, T] - grid = torch.stack(grid, dim=0) # [3, W, H, T] - - freqs = [] - for i in range(3): - freq = get_1d_rotary_pos_embed(self.rope_dim[i], grid[i].reshape(-1), self.theta, use_real=True) - freqs.append(freq) - - freqs_cos = torch.cat([f[0] for f in freqs], dim=1) # (W * H * T, D / 2) - freqs_sin = torch.cat([f[1] for f in freqs], dim=1) # (W * H * T, D / 2) - return freqs_cos, freqs_sin - - -class HunyuanVideoSingleTransformerBlock(nn.Module): - def __init__( - self, - num_attention_heads: int, - attention_head_dim: int, - mlp_ratio: float = 4.0, - qk_norm: str = "rms_norm", - ) -> None: - super().__init__() - - hidden_size = num_attention_heads * attention_head_dim - mlp_dim = int(hidden_size * mlp_ratio) - - self.attn = Attention( - query_dim=hidden_size, - cross_attention_dim=None, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=hidden_size, - bias=True, - processor=HunyuanVideoAttnProcessor2_0(), - qk_norm=qk_norm, - eps=1e-6, - pre_only=True, - ) - - self.norm = AdaLayerNormZeroSingle(hidden_size, norm_type="layer_norm") - self.proj_mlp = nn.Linear(hidden_size, mlp_dim) - self.act_mlp = nn.GELU(approximate="tanh") - self.proj_out = nn.Linear(hidden_size + mlp_dim, hidden_size) - - def forward( - self, - 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, - *args, - **kwargs, - ) -> torch.Tensor: - text_seq_length = encoder_hidden_states.shape[1] - hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) - - residual = hidden_states - - # 1. Input normalization - norm_hidden_states, gate = self.norm(hidden_states, emb=temb) - mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)) - - norm_hidden_states, norm_encoder_hidden_states = ( - norm_hidden_states[:, :-text_seq_length, :], - norm_hidden_states[:, -text_seq_length:, :], - ) - - # 2. Attention - attn_output, context_attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - attention_mask=attention_mask, - image_rotary_emb=image_rotary_emb, - ) - attn_output = torch.cat([attn_output, context_attn_output], dim=1) - - # 3. Modulation and residual connection - hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) - hidden_states = gate.unsqueeze(1) * self.proj_out(hidden_states) - hidden_states = hidden_states + residual - - hidden_states, encoder_hidden_states = ( - hidden_states[:, :-text_seq_length, :], - hidden_states[:, -text_seq_length:, :], - ) - return hidden_states, encoder_hidden_states - - -class HunyuanVideoTransformerBlock(nn.Module): - def __init__( - self, - num_attention_heads: int, - attention_head_dim: int, - mlp_ratio: float, - qk_norm: str = "rms_norm", - ) -> None: - super().__init__() - - hidden_size = num_attention_heads * attention_head_dim - - self.norm1 = AdaLayerNormZero(hidden_size, norm_type="layer_norm") - self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm") - - self.attn = Attention( - query_dim=hidden_size, - cross_attention_dim=None, - added_kv_proj_dim=hidden_size, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=hidden_size, - context_pre_only=False, - bias=True, - processor=HunyuanVideoAttnProcessor2_0(), - qk_norm=qk_norm, - eps=1e-6, - ) - - self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate") - - self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.ff_context = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate") - - def forward( - self, - 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, - *args, - **kwargs, - ) -> 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( - encoder_hidden_states, emb=temb - ) - - # 2. Joint attention - attn_output, context_attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - attention_mask=attention_mask, - image_rotary_emb=freqs_cis, - ) - - # 3. Modulation and residual connection - hidden_states = hidden_states + attn_output * gate_msa.unsqueeze(1) - encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa.unsqueeze(1) - - norm_hidden_states = self.norm2(hidden_states) - norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) - - norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] - - # 4. Feed-forward - ff_output = self.ff(norm_hidden_states) - context_ff_output = self.ff_context(norm_encoder_hidden_states) - - hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output - encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output - - return hidden_states, encoder_hidden_states - - -class HunyuanVideoTokenReplaceSingleTransformerBlock(nn.Module): - def __init__( - self, - num_attention_heads: int, - attention_head_dim: int, - mlp_ratio: float = 4.0, - qk_norm: str = "rms_norm", - ) -> None: - super().__init__() - - hidden_size = num_attention_heads * attention_head_dim - mlp_dim = int(hidden_size * mlp_ratio) - - self.attn = Attention( - query_dim=hidden_size, - cross_attention_dim=None, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=hidden_size, - bias=True, - processor=HunyuanVideoAttnProcessor2_0(), - qk_norm=qk_norm, - eps=1e-6, - pre_only=True, - ) - - self.norm = HunyuanVideoTokenReplaceAdaLayerNormZeroSingle(hidden_size, norm_type="layer_norm") - self.proj_mlp = nn.Linear(hidden_size, mlp_dim) - self.act_mlp = nn.GELU(approximate="tanh") - self.proj_out = nn.Linear(hidden_size + mlp_dim, hidden_size) - - def forward( - self, - 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, - token_replace_emb: torch.Tensor = None, - num_tokens: int = None, - ) -> torch.Tensor: - text_seq_length = encoder_hidden_states.shape[1] - hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) - - residual = hidden_states - - # 1. Input normalization - norm_hidden_states, gate, tr_gate = self.norm(hidden_states, temb, token_replace_emb, num_tokens) - mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)) - - norm_hidden_states, norm_encoder_hidden_states = ( - norm_hidden_states[:, :-text_seq_length, :], - norm_hidden_states[:, -text_seq_length:, :], - ) - - # 2. Attention - attn_output, context_attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - attention_mask=attention_mask, - image_rotary_emb=image_rotary_emb, - ) - attn_output = torch.cat([attn_output, context_attn_output], dim=1) - - # 3. Modulation and residual connection - hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) - - proj_output = self.proj_out(hidden_states) - hidden_states_zero = proj_output[:, :num_tokens] * tr_gate.unsqueeze(1) - hidden_states_orig = proj_output[:, num_tokens:] * gate.unsqueeze(1) - hidden_states = torch.cat([hidden_states_zero, hidden_states_orig], dim=1) - hidden_states = hidden_states + residual - - hidden_states, encoder_hidden_states = ( - hidden_states[:, :-text_seq_length, :], - hidden_states[:, -text_seq_length:, :], - ) - return hidden_states, encoder_hidden_states - - -class HunyuanVideoTokenReplaceTransformerBlock(nn.Module): - def __init__( - self, - num_attention_heads: int, - attention_head_dim: int, - mlp_ratio: float, - qk_norm: str = "rms_norm", - ) -> None: - super().__init__() - - hidden_size = num_attention_heads * attention_head_dim - - self.norm1 = HunyuanVideoTokenReplaceAdaLayerNormZero(hidden_size, norm_type="layer_norm") - self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm") - - self.attn = Attention( - query_dim=hidden_size, - cross_attention_dim=None, - added_kv_proj_dim=hidden_size, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=hidden_size, - context_pre_only=False, - bias=True, - processor=HunyuanVideoAttnProcessor2_0(), - qk_norm=qk_norm, - eps=1e-6, - ) - - self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate") - - self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.ff_context = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate") - - def forward( - self, - 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, - token_replace_emb: torch.Tensor = None, - num_tokens: int = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: - # 1. Input normalization - ( - norm_hidden_states, - gate_msa, - shift_mlp, - scale_mlp, - gate_mlp, - tr_gate_msa, - tr_shift_mlp, - tr_scale_mlp, - tr_gate_mlp, - ) = self.norm1(hidden_states, temb, token_replace_emb, num_tokens) - norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( - encoder_hidden_states, emb=temb - ) - - # 2. Joint attention - attn_output, context_attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - attention_mask=attention_mask, - image_rotary_emb=freqs_cis, - ) - - # 3. Modulation and residual connection - hidden_states_zero = hidden_states[:, :num_tokens] + attn_output[:, :num_tokens] * tr_gate_msa.unsqueeze(1) - hidden_states_orig = hidden_states[:, num_tokens:] + attn_output[:, num_tokens:] * gate_msa.unsqueeze(1) - hidden_states = torch.cat([hidden_states_zero, hidden_states_orig], dim=1) - encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa.unsqueeze(1) - - norm_hidden_states = self.norm2(hidden_states) - norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) - - hidden_states_zero = norm_hidden_states[:, :num_tokens] * (1 + tr_scale_mlp[:, None]) + tr_shift_mlp[:, None] - hidden_states_orig = norm_hidden_states[:, num_tokens:] * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - norm_hidden_states = torch.cat([hidden_states_zero, hidden_states_orig], dim=1) - norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] - - # 4. Feed-forward - ff_output = self.ff(norm_hidden_states) - context_ff_output = self.ff_context(norm_encoder_hidden_states) - - hidden_states_zero = hidden_states[:, :num_tokens] + ff_output[:, :num_tokens] * tr_gate_mlp.unsqueeze(1) - hidden_states_orig = hidden_states[:, num_tokens:] + ff_output[:, num_tokens:] * gate_mlp.unsqueeze(1) - hidden_states = torch.cat([hidden_states_zero, hidden_states_orig], dim=1) - encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output - - return hidden_states, encoder_hidden_states - - -class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin): - r""" - A Transformer model for video-like data used in [HunyuanVideo](https://huggingface.co/tencent/HunyuanVideo). - - Args: - 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. - num_attention_heads (`int`, defaults to `24`): - The number of heads to use for multi-head attention. - attention_head_dim (`int`, defaults to `128`): - The number of channels in each head. - num_layers (`int`, defaults to `20`): - The number of layers of dual-stream blocks to use. - num_single_layers (`int`, defaults to `40`): - The number of layers of single-stream blocks to use. - num_refiner_layers (`int`, defaults to `2`): - The number of layers of refiner blocks to use. - mlp_ratio (`float`, defaults to `4.0`): - The ratio of the hidden layer size to the input size in the feedforward network. - patch_size (`int`, defaults to `2`): - 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. - qk_norm (`str`, defaults to `rms_norm`): - The normalization to use for the query and key projections in the attention layers. - guidance_embeds (`bool`, defaults to `True`): - Whether to use guidance embeddings in the model. - text_embed_dim (`int`, defaults to `4096`): - Input dimension of text embeddings from the text encoder. - pooled_projection_dim (`int`, defaults to `768`): - 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)`): - 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 - image is concatenated to the latent stream. If `token_replace`, the image is used to replace first-frame - tokens in the latent stream and apply conditioning. - """ - - _supports_gradient_checkpointing = True - _skip_layerwise_casting_patterns = ["x_embedder", "context_embedder", "norm"] - _no_split_modules = [ - "HunyuanVideoTransformerBlock", - "HunyuanVideoSingleTransformerBlock", - "HunyuanVideoPatchEmbed", - "HunyuanVideoTokenRefiner", - ] - - @register_to_config - def __init__( - self, - in_channels: int = 16, - out_channels: int = 16, - num_attention_heads: int = 24, - attention_head_dim: int = 128, - num_layers: int = 20, - num_single_layers: int = 40, - num_refiner_layers: int = 2, - mlp_ratio: float = 4.0, - patch_size: int = 2, - patch_size_t: int = 1, - qk_norm: str = "rms_norm", - guidance_embeds: bool = True, - 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, - ) -> None: - super().__init__() - - supported_image_condition_types = ["latent_concat", "token_replace"] - if image_condition_type is not None and image_condition_type not in supported_image_condition_types: - raise ValueError( - f"Invalid `image_condition_type` ({image_condition_type}). Supported ones are: {supported_image_condition_types}" - ) - - inner_dim = num_attention_heads * attention_head_dim - out_channels = out_channels or in_channels - - # 1. Latent and condition embedders - self.x_embedder = HunyuanVideoPatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim) - self.context_embedder = HunyuanVideoTokenRefiner( - text_embed_dim, num_attention_heads, attention_head_dim, num_layers=num_refiner_layers - ) - - self.time_text_embed = HunyuanVideoConditionEmbedding( - inner_dim, pooled_projection_dim, guidance_embeds, image_condition_type - ) - - # 2. RoPE - self.rope = HunyuanVideoRotaryPosEmbed(patch_size, patch_size_t, rope_axes_dim, rope_theta) - - # 3. Dual stream transformer blocks - if image_condition_type == "token_replace": - self.transformer_blocks = nn.ModuleList( - [ - HunyuanVideoTokenReplaceTransformerBlock( - num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm - ) - for _ in range(num_layers) - ] - ) - else: - self.transformer_blocks = nn.ModuleList( - [ - HunyuanVideoTransformerBlock( - num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm - ) - for _ in range(num_layers) - ] - ) - - # 4. Single stream transformer blocks - if image_condition_type == "token_replace": - self.single_transformer_blocks = nn.ModuleList( - [ - HunyuanVideoTokenReplaceSingleTransformerBlock( - num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm - ) - for _ in range(num_single_layers) - ] - ) - else: - self.single_transformer_blocks = nn.ModuleList( - [ - HunyuanVideoSingleTransformerBlock( - num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm - ) - for _ in range(num_single_layers) - ] - ) - - # 5. Output projection - self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6) - self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * out_channels) - - self.gradient_checkpointing = False - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - def forward( - self, - hidden_states: torch.Tensor, - timestep: torch.LongTensor, - encoder_hidden_states: torch.Tensor, - encoder_attention_mask: torch.Tensor, - pooled_projections: torch.Tensor, - guidance: torch.Tensor = None, - attention_kwargs: Optional[Dict[str, Any]] = 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." - ) - - 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 - post_patch_height = height // p - post_patch_width = width // p - first_frame_num_tokens = 1 * post_patch_height * post_patch_width - - # 1. RoPE - image_rotary_emb = self.rope(hidden_states) - - # 2. Conditional embeddings - temb, token_replace_emb = self.time_text_embed(timestep, pooled_projections, guidance) - - hidden_states = self.x_embedder(hidden_states) - encoder_hidden_states = self.context_embedder(encoder_hidden_states, timestep, encoder_attention_mask) - - # 3. Attention mask preparation - latent_sequence_length = hidden_states.shape[1] - condition_sequence_length = encoder_hidden_states.shape[1] - sequence_length = latent_sequence_length + condition_sequence_length - attention_mask = torch.ones( - batch_size, sequence_length, device=hidden_states.device, dtype=torch.bool - ) # [B, N] - effective_condition_sequence_length = encoder_attention_mask.sum(dim=1, dtype=torch.int) # [B,] - effective_sequence_length = latent_sequence_length + effective_condition_sequence_length - indices = torch.arange(sequence_length, device=hidden_states.device).unsqueeze(0) # [1, N] - mask_indices = indices >= effective_sequence_length.unsqueeze(1) # [B, N] - attention_mask = attention_mask.masked_fill(mask_indices, False) - attention_mask = attention_mask.unsqueeze(1).unsqueeze(1) # [B, 1, 1, N] - - # 4. Transformer blocks - if torch.is_grad_enabled() and self.gradient_checkpointing: - for block in self.transformer_blocks: - hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( - block, - hidden_states, - encoder_hidden_states, - temb, - attention_mask, - image_rotary_emb, - token_replace_emb, - first_frame_num_tokens, - ) - - for block in self.single_transformer_blocks: - hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( - block, - hidden_states, - encoder_hidden_states, - temb, - attention_mask, - image_rotary_emb, - token_replace_emb, - first_frame_num_tokens, - ) - - else: - for block in self.transformer_blocks: - hidden_states, encoder_hidden_states = block( - hidden_states, - encoder_hidden_states, - temb, - attention_mask, - image_rotary_emb, - token_replace_emb, - first_frame_num_tokens, - ) - - for block in self.single_transformer_blocks: - hidden_states, encoder_hidden_states = block( - hidden_states, - encoder_hidden_states, - temb, - attention_mask, - image_rotary_emb, - token_replace_emb, - first_frame_num_tokens, - ) - - # 5. Output projection - hidden_states = self.norm_out(hidden_states, temb) - hidden_states = self.proj_out(hidden_states) - - hidden_states = hidden_states.reshape( - batch_size, post_patch_num_frames, post_patch_height, post_patch_width, -1, p_t, p, p - ) - 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) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_hunyuan_video_framepack.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_hunyuan_video_framepack.py deleted file mode 100644 index c2eb7fd2a705257294357629ddd1a79fde965a0e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_hunyuan_video_framepack.py +++ /dev/null @@ -1,416 +0,0 @@ -# Copyright 2025 The Framepack Team, The Hunyuan 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 Any, Dict, List, Optional, Tuple - -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 USE_PEFT_BACKEND, get_logger, scale_lora_layers, unscale_lora_layers -from ..cache_utils import CacheMixin -from ..embeddings import get_1d_rotary_pos_embed -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormContinuous -from .transformer_hunyuan_video import ( - HunyuanVideoConditionEmbedding, - HunyuanVideoPatchEmbed, - HunyuanVideoSingleTransformerBlock, - HunyuanVideoTokenRefiner, - HunyuanVideoTransformerBlock, -) - - -logger = get_logger(__name__) # pylint: disable=invalid-name - - -class HunyuanVideoFramepackRotaryPosEmbed(nn.Module): - 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 - self.patch_size_t = patch_size_t - self.rope_dim = rope_dim - self.theta = theta - - def forward(self, frame_indices: torch.Tensor, height: int, width: int, device: torch.device): - height = height // self.patch_size - width = width // self.patch_size - grid = torch.meshgrid( - frame_indices.to(device=device, dtype=torch.float32), - torch.arange(0, height, device=device, dtype=torch.float32), - torch.arange(0, width, device=device, dtype=torch.float32), - indexing="ij", - ) # 3 * [W, H, T] - grid = torch.stack(grid, dim=0) # [3, W, H, T] - - freqs = [] - for i in range(3): - freq = get_1d_rotary_pos_embed(self.rope_dim[i], grid[i].reshape(-1), self.theta, use_real=True) - freqs.append(freq) - - freqs_cos = torch.cat([f[0] for f in freqs], dim=1) # (W * H * T, D / 2) - freqs_sin = torch.cat([f[1] for f in freqs], dim=1) # (W * H * T, D / 2) - - return freqs_cos, freqs_sin - - -class FramepackClipVisionProjection(nn.Module): - def __init__(self, in_channels: int, out_channels: int): - super().__init__() - self.up = nn.Linear(in_channels, out_channels * 3) - self.down = nn.Linear(out_channels * 3, out_channels) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.up(hidden_states) - hidden_states = F.silu(hidden_states) - hidden_states = self.down(hidden_states) - return hidden_states - - -class HunyuanVideoHistoryPatchEmbed(nn.Module): - def __init__(self, in_channels: int, inner_dim: int): - super().__init__() - self.proj = nn.Conv3d(in_channels, inner_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2)) - self.proj_2x = nn.Conv3d(in_channels, inner_dim, kernel_size=(2, 4, 4), stride=(2, 4, 4)) - self.proj_4x = nn.Conv3d(in_channels, inner_dim, kernel_size=(4, 8, 8), stride=(4, 8, 8)) - - def forward( - self, - latents_clean: Optional[torch.Tensor] = None, - latents_clean_2x: Optional[torch.Tensor] = None, - latents_clean_4x: Optional[torch.Tensor] = None, - ): - if latents_clean is not None: - latents_clean = self.proj(latents_clean) - latents_clean = latents_clean.flatten(2).transpose(1, 2) - if latents_clean_2x is not None: - latents_clean_2x = _pad_for_3d_conv(latents_clean_2x, (2, 4, 4)) - latents_clean_2x = self.proj_2x(latents_clean_2x) - latents_clean_2x = latents_clean_2x.flatten(2).transpose(1, 2) - if latents_clean_4x is not None: - latents_clean_4x = _pad_for_3d_conv(latents_clean_4x, (4, 8, 8)) - latents_clean_4x = self.proj_4x(latents_clean_4x) - latents_clean_4x = latents_clean_4x.flatten(2).transpose(1, 2) - return latents_clean, latents_clean_2x, latents_clean_4x - - -class HunyuanVideoFramepackTransformer3DModel( - ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin -): - _supports_gradient_checkpointing = True - _skip_layerwise_casting_patterns = ["x_embedder", "context_embedder", "norm"] - _no_split_modules = [ - "HunyuanVideoTransformerBlock", - "HunyuanVideoSingleTransformerBlock", - "HunyuanVideoHistoryPatchEmbed", - "HunyuanVideoTokenRefiner", - ] - - @register_to_config - def __init__( - self, - in_channels: int = 16, - out_channels: int = 16, - num_attention_heads: int = 24, - attention_head_dim: int = 128, - num_layers: int = 20, - num_single_layers: int = 40, - num_refiner_layers: int = 2, - mlp_ratio: float = 4.0, - patch_size: int = 2, - patch_size_t: int = 1, - qk_norm: str = "rms_norm", - guidance_embeds: bool = True, - 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, - has_image_proj: int = False, - image_proj_dim: int = 1152, - has_clean_x_embedder: int = False, - ) -> None: - super().__init__() - - inner_dim = num_attention_heads * attention_head_dim - out_channels = out_channels or in_channels - - # 1. Latent and condition embedders - self.x_embedder = HunyuanVideoPatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim) - - # Framepack history projection embedder - self.clean_x_embedder = None - if has_clean_x_embedder: - self.clean_x_embedder = HunyuanVideoHistoryPatchEmbed(in_channels, inner_dim) - - self.context_embedder = HunyuanVideoTokenRefiner( - text_embed_dim, num_attention_heads, attention_head_dim, num_layers=num_refiner_layers - ) - - # Framepack image-conditioning embedder - self.image_projection = FramepackClipVisionProjection(image_proj_dim, inner_dim) if has_image_proj else None - - self.time_text_embed = HunyuanVideoConditionEmbedding( - inner_dim, pooled_projection_dim, guidance_embeds, image_condition_type - ) - - # 2. RoPE - self.rope = HunyuanVideoFramepackRotaryPosEmbed(patch_size, patch_size_t, rope_axes_dim, rope_theta) - - # 3. Dual stream transformer blocks - self.transformer_blocks = nn.ModuleList( - [ - HunyuanVideoTransformerBlock( - num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm - ) - for _ in range(num_layers) - ] - ) - - # 4. Single stream transformer blocks - self.single_transformer_blocks = nn.ModuleList( - [ - HunyuanVideoSingleTransformerBlock( - num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm - ) - for _ in range(num_single_layers) - ] - ) - - # 5. Output projection - self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6) - self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * out_channels) - - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - timestep: torch.LongTensor, - encoder_hidden_states: torch.Tensor, - encoder_attention_mask: torch.Tensor, - 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, - return_dict: bool = True, - ): - 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, p_t = self.config.patch_size, self.config.patch_size_t - post_patch_num_frames = num_frames // p_t - post_patch_height = height // p - post_patch_width = width // p - original_context_length = post_patch_num_frames * post_patch_height * post_patch_width - - if indices_latents is None: - indices_latents = torch.arange(0, num_frames).unsqueeze(0).expand(batch_size, -1) - - hidden_states = self.x_embedder(hidden_states) - image_rotary_emb = self.rope( - frame_indices=indices_latents, height=height, width=width, device=hidden_states.device - ) - - latents_clean, latents_history_2x, latents_history_4x = self.clean_x_embedder( - latents_clean, latents_history_2x, latents_history_4x - ) - - if latents_clean is not None and indices_latents_clean is not None: - image_rotary_emb_clean = self.rope( - frame_indices=indices_latents_clean, height=height, width=width, device=hidden_states.device - ) - if latents_history_2x is not None and indices_latents_history_2x is not None: - image_rotary_emb_history_2x = self.rope( - frame_indices=indices_latents_history_2x, height=height, width=width, device=hidden_states.device - ) - if latents_history_4x is not None and indices_latents_history_4x is not None: - image_rotary_emb_history_4x = self.rope( - frame_indices=indices_latents_history_4x, height=height, width=width, device=hidden_states.device - ) - - hidden_states, image_rotary_emb = self._pack_history_states( - hidden_states, - latents_clean, - latents_history_2x, - latents_history_4x, - image_rotary_emb, - image_rotary_emb_clean, - image_rotary_emb_history_2x, - image_rotary_emb_history_4x, - post_patch_height, - post_patch_width, - ) - - temb, _ = self.time_text_embed(timestep, pooled_projections, guidance) - encoder_hidden_states = self.context_embedder(encoder_hidden_states, timestep, encoder_attention_mask) - - encoder_hidden_states_image = self.image_projection(image_embeds) - attention_mask_image = encoder_attention_mask.new_ones((batch_size, encoder_hidden_states_image.shape[1])) - - # must cat before (not after) encoder_hidden_states, due to attn masking - encoder_hidden_states = torch.cat([encoder_hidden_states_image, encoder_hidden_states], dim=1) - encoder_attention_mask = torch.cat([attention_mask_image, encoder_attention_mask], dim=1) - - latent_sequence_length = hidden_states.shape[1] - condition_sequence_length = encoder_hidden_states.shape[1] - sequence_length = latent_sequence_length + condition_sequence_length - attention_mask = torch.zeros( - batch_size, sequence_length, device=hidden_states.device, dtype=torch.bool - ) # [B, N] - effective_condition_sequence_length = encoder_attention_mask.sum(dim=1, dtype=torch.int) # [B,] - effective_sequence_length = latent_sequence_length + effective_condition_sequence_length - - if batch_size == 1: - encoder_hidden_states = encoder_hidden_states[:, : effective_condition_sequence_length[0]] - attention_mask = None - else: - for i in range(batch_size): - attention_mask[i, : effective_sequence_length[i]] = True - # [B, 1, 1, N], for broadcasting across attention heads - attention_mask = attention_mask.unsqueeze(1).unsqueeze(1) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - for block in self.transformer_blocks: - hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, temb, attention_mask, image_rotary_emb - ) - - for block in self.single_transformer_blocks: - hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, temb, attention_mask, image_rotary_emb - ) - - else: - for block in self.transformer_blocks: - hidden_states, encoder_hidden_states = block( - hidden_states, encoder_hidden_states, temb, attention_mask, image_rotary_emb - ) - - for block in self.single_transformer_blocks: - hidden_states, encoder_hidden_states = block( - hidden_states, encoder_hidden_states, temb, attention_mask, image_rotary_emb - ) - - hidden_states = hidden_states[:, -original_context_length:] - hidden_states = self.norm_out(hidden_states, temb) - hidden_states = self.proj_out(hidden_states) - - hidden_states = hidden_states.reshape( - batch_size, post_patch_num_frames, post_patch_height, post_patch_width, -1, p_t, p, p - ) - 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) - - 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, - height: int = None, - width: int = None, - ): - image_rotary_emb = list(image_rotary_emb) # convert tuple to list for in-place modification - - if latents_clean is not None and image_rotary_emb_clean is not None: - hidden_states = torch.cat([latents_clean, hidden_states], dim=1) - image_rotary_emb[0] = torch.cat([image_rotary_emb_clean[0], image_rotary_emb[0]], dim=0) - image_rotary_emb[1] = torch.cat([image_rotary_emb_clean[1], image_rotary_emb[1]], dim=0) - - if latents_history_2x is not None and image_rotary_emb_history_2x is not None: - hidden_states = torch.cat([latents_history_2x, hidden_states], dim=1) - image_rotary_emb_history_2x = self._pad_rotary_emb(image_rotary_emb_history_2x, height, width, (2, 2, 2)) - image_rotary_emb[0] = torch.cat([image_rotary_emb_history_2x[0], image_rotary_emb[0]], dim=0) - image_rotary_emb[1] = torch.cat([image_rotary_emb_history_2x[1], image_rotary_emb[1]], dim=0) - - if latents_history_4x is not None and image_rotary_emb_history_4x is not None: - hidden_states = torch.cat([latents_history_4x, hidden_states], dim=1) - image_rotary_emb_history_4x = self._pad_rotary_emb(image_rotary_emb_history_4x, height, width, (4, 4, 4)) - image_rotary_emb[0] = torch.cat([image_rotary_emb_history_4x[0], image_rotary_emb[0]], dim=0) - image_rotary_emb[1] = torch.cat([image_rotary_emb_history_4x[1], image_rotary_emb[1]], dim=0) - - return hidden_states, tuple(image_rotary_emb) - - def _pad_rotary_emb( - self, - image_rotary_emb: Tuple[torch.Tensor], - height: int, - width: 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 - freqs_cos = freqs_cos.unsqueeze(0).permute(0, 2, 1).unflatten(2, (-1, height, width)) - freqs_sin = freqs_sin.unsqueeze(0).permute(0, 2, 1).unflatten(2, (-1, height, width)) - freqs_cos = _pad_for_3d_conv(freqs_cos, kernel_size) - freqs_sin = _pad_for_3d_conv(freqs_sin, kernel_size) - freqs_cos = _center_down_sample_3d(freqs_cos, kernel_size) - freqs_sin = _center_down_sample_3d(freqs_sin, kernel_size) - freqs_cos = freqs_cos.flatten(2).permute(0, 2, 1).squeeze(0) - freqs_sin = freqs_sin.flatten(2).permute(0, 2, 1).squeeze(0) - return freqs_cos, freqs_sin - - -def _pad_for_3d_conv(x, kernel_size): - if isinstance(x, (tuple, list)): - return tuple(_pad_for_3d_conv(i, kernel_size) for i in x) - 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): - if isinstance(x, (tuple, list)): - return tuple(_center_down_sample_3d(i, kernel_size) for i in x) - return torch.nn.functional.avg_pool3d(x, kernel_size, stride=kernel_size) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_ltx.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_ltx.py deleted file mode 100644 index 38b7b6af50f9221f5fd635d69cf7aa13ee8a6923..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_ltx.py +++ /dev/null @@ -1,487 +0,0 @@ -# Copyright 2025 The Genmo 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 typing import Any, Dict, Optional, Tuple, Union - -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 USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers -from ...utils.torch_utils import maybe_allow_in_graph -from ..attention import FeedForward -from ..attention_processor import Attention -from ..cache_utils import CacheMixin -from ..embeddings import PixArtAlphaTextProjection -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormSingle, RMSNorm - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class LTXVideoAttentionProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is - used in the LTX model. It applies a normalization layer and rotary embedding on the query and key vector. - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError( - "LTXVideoAttentionProcessor2_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, - ) -> 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 image_rotary_emb is not None: - query = apply_rotary_emb(query, image_rotary_emb) - key = apply_rotary_emb(key, image_rotary_emb) - - 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) - - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - hidden_states = hidden_states.transpose(1, 2).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 LTXVideoRotaryPosEmbed(nn.Module): - def __init__( - self, - dim: int, - base_num_frames: int = 20, - base_height: int = 2048, - base_width: int = 2048, - patch_size: int = 1, - patch_size_t: int = 1, - theta: float = 10000.0, - ) -> None: - super().__init__() - - self.dim = dim - self.base_num_frames = base_num_frames - self.base_height = base_height - self.base_width = base_width - self.patch_size = patch_size - self.patch_size_t = patch_size_t - self.theta = theta - - def _prepare_video_coords( - self, - batch_size: int, - num_frames: int, - height: int, - width: int, - rope_interpolation_scale: Tuple[torch.Tensor, float, float], - device: torch.device, - ) -> torch.Tensor: - # Always compute rope in fp32 - grid_h = torch.arange(height, dtype=torch.float32, device=device) - grid_w = torch.arange(width, dtype=torch.float32, device=device) - grid_f = torch.arange(num_frames, dtype=torch.float32, device=device) - grid = torch.meshgrid(grid_f, grid_h, grid_w, indexing="ij") - grid = torch.stack(grid, dim=0) - grid = grid.unsqueeze(0).repeat(batch_size, 1, 1, 1, 1) - - if rope_interpolation_scale is not None: - grid[:, 0:1] = grid[:, 0:1] * rope_interpolation_scale[0] * self.patch_size_t / self.base_num_frames - grid[:, 1:2] = grid[:, 1:2] * rope_interpolation_scale[1] * self.patch_size / self.base_height - grid[:, 2:3] = grid[:, 2:3] * rope_interpolation_scale[2] * self.patch_size / self.base_width - - grid = grid.flatten(2, 4).transpose(1, 2) - - return grid - - 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]: - batch_size = hidden_states.size(0) - - if video_coords is None: - grid = self._prepare_video_coords( - batch_size, - num_frames, - height, - width, - rope_interpolation_scale=rope_interpolation_scale, - device=hidden_states.device, - ) - else: - grid = torch.stack( - [ - video_coords[:, 0] / self.base_num_frames, - video_coords[:, 1] / self.base_height, - video_coords[:, 2] / self.base_width, - ], - dim=-1, - ) - - start = 1.0 - end = self.theta - freqs = self.theta ** torch.linspace( - math.log(start, self.theta), - math.log(end, self.theta), - self.dim // 6, - device=hidden_states.device, - dtype=torch.float32, - ) - freqs = freqs * math.pi / 2.0 - freqs = freqs * (grid.unsqueeze(-1) * 2 - 1) - freqs = freqs.transpose(-1, -2).flatten(2) - - cos_freqs = freqs.cos().repeat_interleave(2, dim=-1) - sin_freqs = freqs.sin().repeat_interleave(2, dim=-1) - - if self.dim % 6 != 0: - cos_padding = torch.ones_like(cos_freqs[:, :, : self.dim % 6]) - sin_padding = torch.zeros_like(cos_freqs[:, :, : self.dim % 6]) - cos_freqs = torch.cat([cos_padding, cos_freqs], dim=-1) - sin_freqs = torch.cat([sin_padding, sin_freqs], dim=-1) - - return cos_freqs, sin_freqs - - -@maybe_allow_in_graph -class LTXVideoTransformerBlock(nn.Module): - r""" - Transformer block used in [LTX](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, - 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, - ): - super().__init__() - - self.norm1 = RMSNorm(dim, eps=eps, elementwise_affine=elementwise_affine) - self.attn1 = Attention( - 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, - processor=LTXVideoAttentionProcessor2_0(), - ) - - self.norm2 = RMSNorm(dim, eps=eps, elementwise_affine=elementwise_affine) - self.attn2 = Attention( - 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, - processor=LTXVideoAttentionProcessor2_0(), - ) - - self.ff = FeedForward(dim, activation_fn=activation_fn) - - self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) - - def forward( - self, - 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, - ) -> torch.Tensor: - batch_size = hidden_states.size(0) - norm_hidden_states = self.norm1(hidden_states) - - num_ada_params = self.scale_shift_table.shape[0] - ada_values = self.scale_shift_table[None, None] + 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, - image_rotary_emb=image_rotary_emb, - ) - hidden_states = hidden_states + attn_hidden_states * gate_msa - - attn_hidden_states = self.attn2( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - image_rotary_emb=None, - attention_mask=encoder_attention_mask, - ) - hidden_states = hidden_states + attn_hidden_states - norm_hidden_states = self.norm2(hidden_states) * (1 + scale_mlp) + shift_mlp - - ff_output = self.ff(norm_hidden_states) - hidden_states = hidden_states + ff_output * gate_mlp - - return hidden_states - - -@maybe_allow_in_graph -class LTXVideoTransformer3DModel(ModelMixin, ConfigMixin, 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"] - - @register_to_config - def __init__( - self, - in_channels: int = 128, - out_channels: int = 128, - patch_size: int = 1, - patch_size_t: int = 1, - num_attention_heads: int = 32, - attention_head_dim: int = 64, - cross_attention_dim: int = 2048, - num_layers: int = 28, - 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 = 4096, - attention_bias: bool = True, - attention_out_bias: bool = True, - ) -> None: - super().__init__() - - out_channels = out_channels or in_channels - inner_dim = num_attention_heads * attention_head_dim - - self.proj_in = nn.Linear(in_channels, inner_dim) - - self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5) - self.time_embed = AdaLayerNormSingle(inner_dim, use_additional_conditions=False) - - self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) - - self.rope = LTXVideoRotaryPosEmbed( - dim=inner_dim, - base_num_frames=20, - base_height=2048, - base_width=2048, - patch_size=patch_size, - patch_size_t=patch_size_t, - theta=10000.0, - ) - - self.transformer_blocks = nn.ModuleList( - [ - LTXVideoTransformerBlock( - dim=inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - cross_attention_dim=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, - ) - for _ in range(num_layers) - ] - ) - - self.norm_out = nn.LayerNorm(inner_dim, eps=1e-6, elementwise_affine=False) - self.proj_out = nn.Linear(inner_dim, out_channels) - - self.gradient_checkpointing = False - - 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, - 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 - 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) - - batch_size = hidden_states.size(0) - hidden_states = self.proj_in(hidden_states) - - 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)) - - encoder_hidden_states = self.caption_projection(encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1)) - - for block in self.transformer_blocks: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - block, - hidden_states, - encoder_hidden_states, - temb, - image_rotary_emb, - encoder_attention_mask, - ) - else: - hidden_states = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - temb=temb, - image_rotary_emb=image_rotary_emb, - encoder_attention_mask=encoder_attention_mask, - ) - - 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) - - 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) - - -def apply_rotary_emb(x, freqs): - 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 diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_lumina2.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_lumina2.py deleted file mode 100644 index 77121edb9fc96bac29e0caa9cf6a20db3370dc19..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_lumina2.py +++ /dev/null @@ -1,548 +0,0 @@ -# Copyright 2025 Alpha-VLLM Authors 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 typing import Any, Dict, List, Optional, Tuple, Union - -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 ...loaders.single_file_model import FromOriginalModelMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers -from ..attention import LuminaFeedForward -from ..attention_processor import Attention -from ..embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import LuminaLayerNormContinuous, LuminaRMSNormZero, RMSNorm - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class Lumina2CombinedTimestepCaptionEmbedding(nn.Module): - def __init__( - self, - hidden_size: int = 4096, - cap_feat_dim: int = 2048, - frequency_embedding_size: int = 256, - norm_eps: float = 1e-5, - ) -> None: - super().__init__() - - self.time_proj = Timesteps( - num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0.0 - ) - - self.timestep_embedder = TimestepEmbedding( - in_channels=frequency_embedding_size, time_embed_dim=min(hidden_size, 1024) - ) - - self.caption_embedder = nn.Sequential( - RMSNorm(cap_feat_dim, eps=norm_eps), nn.Linear(cap_feat_dim, hidden_size, bias=True) - ) - - def forward( - self, hidden_states: torch.Tensor, timestep: torch.Tensor, encoder_hidden_states: 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) - return time_embed, caption_embed - - -class Lumina2AttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is - used in the Lumina2Transformer2DModel model. It applies normalization and RoPE on query and key vectors. - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_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: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - base_sequence_length: Optional[int] = None, - ) -> torch.Tensor: - batch_size, sequence_length, _ = hidden_states.shape - - # Get Query-Key-Value Pair - query = attn.to_q(hidden_states) - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query_dim = query.shape[-1] - inner_dim = key.shape[-1] - head_dim = query_dim // attn.heads - dtype = query.dtype - - # Get key-value heads - kv_heads = inner_dim // head_dim - - query = query.view(batch_size, -1, attn.heads, head_dim) - key = key.view(batch_size, -1, kv_heads, head_dim) - value = value.view(batch_size, -1, kv_heads, head_dim) - - # Apply Query-Key Norm if needed - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # Apply RoPE if needed - if image_rotary_emb is not None: - query = apply_rotary_emb(query, image_rotary_emb, use_real=False) - key = apply_rotary_emb(key, image_rotary_emb, use_real=False) - - query, key = query.to(dtype), key.to(dtype) - - # Apply proportional attention if true - if base_sequence_length is not None: - softmax_scale = math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale - else: - softmax_scale = attn.scale - - # perform Grouped-qurey Attention (GQA) - n_rep = attn.heads // kv_heads - if n_rep >= 1: - key = key.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3) - value = value.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3) - - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - if attention_mask is not None: - attention_mask = attention_mask.bool().view(batch_size, 1, 1, -1) - - query = query.transpose(1, 2) - key = key.transpose(1, 2) - value = value.transpose(1, 2) - - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, scale=softmax_scale - ) - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.type_as(query) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - return hidden_states - - -class Lumina2TransformerBlock(nn.Module): - def __init__( - self, - dim: int, - num_attention_heads: int, - num_kv_heads: int, - multiple_of: int, - ffn_dim_multiplier: float, - norm_eps: float, - modulation: bool = True, - ) -> None: - super().__init__() - self.head_dim = dim // num_attention_heads - self.modulation = modulation - - self.attn = Attention( - query_dim=dim, - cross_attention_dim=None, - dim_head=dim // num_attention_heads, - qk_norm="rms_norm", - heads=num_attention_heads, - kv_heads=num_kv_heads, - eps=1e-5, - bias=False, - out_bias=False, - processor=Lumina2AttnProcessor2_0(), - ) - - self.feed_forward = LuminaFeedForward( - dim=dim, - inner_dim=4 * dim, - multiple_of=multiple_of, - ffn_dim_multiplier=ffn_dim_multiplier, - ) - - if modulation: - self.norm1 = LuminaRMSNormZero( - embedding_dim=dim, - norm_eps=norm_eps, - norm_elementwise_affine=True, - ) - else: - self.norm1 = RMSNorm(dim, eps=norm_eps) - self.ffn_norm1 = RMSNorm(dim, eps=norm_eps) - - self.norm2 = RMSNorm(dim, eps=norm_eps) - self.ffn_norm2 = RMSNorm(dim, eps=norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor, - image_rotary_emb: torch.Tensor, - temb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - if self.modulation: - norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb) - attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_hidden_states, - attention_mask=attention_mask, - image_rotary_emb=image_rotary_emb, - ) - hidden_states = hidden_states + gate_msa.unsqueeze(1).tanh() * self.norm2(attn_output) - mlp_output = self.feed_forward(self.ffn_norm1(hidden_states) * (1 + scale_mlp.unsqueeze(1))) - hidden_states = hidden_states + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(mlp_output) - else: - norm_hidden_states = self.norm1(hidden_states) - attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_hidden_states, - attention_mask=attention_mask, - image_rotary_emb=image_rotary_emb, - ) - hidden_states = hidden_states + self.norm2(attn_output) - mlp_output = self.feed_forward(self.ffn_norm1(hidden_states)) - hidden_states = hidden_states + self.ffn_norm2(mlp_output) - - return hidden_states - - -class Lumina2RotaryPosEmbed(nn.Module): - 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 - self.axes_lens = axes_lens - self.patch_size = patch_size - - 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]: - 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)): - emb = get_1d_rotary_pos_embed(d, e, theta=self.theta, freqs_dtype=freqs_dtype) - freqs_cis.append(emb) - return freqs_cis - - def _get_freqs_cis(self, ids: torch.Tensor) -> torch.Tensor: - device = ids.device - if ids.device.type == "mps": - ids = ids.to("cpu") - - result = [] - for i in range(len(self.axes_dim)): - freqs = self.freqs_cis[i].to(ids.device) - index = ids[:, :, i : i + 1].repeat(1, 1, freqs.shape[-1]).to(torch.int64) - result.append(torch.gather(freqs.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index)) - return torch.cat(result, dim=-1).to(device) - - def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor): - batch_size, channels, height, width = hidden_states.shape - p = self.patch_size - post_patch_height, post_patch_width = height // p, width // p - image_seq_len = post_patch_height * post_patch_width - device = hidden_states.device - - encoder_seq_len = attention_mask.shape[1] - l_effective_cap_len = attention_mask.sum(dim=1).tolist() - seq_lengths = [cap_seq_len + image_seq_len for cap_seq_len in l_effective_cap_len] - max_seq_len = max(seq_lengths) - - # Create position IDs - position_ids = torch.zeros(batch_size, max_seq_len, 3, dtype=torch.int32, device=device) - - for i, (cap_seq_len, seq_len) in enumerate(zip(l_effective_cap_len, seq_lengths)): - # add caption position ids - position_ids[i, :cap_seq_len, 0] = torch.arange(cap_seq_len, dtype=torch.int32, device=device) - position_ids[i, cap_seq_len:seq_len, 0] = cap_seq_len - - # add image position ids - row_ids = ( - torch.arange(post_patch_height, dtype=torch.int32, device=device) - .view(-1, 1) - .repeat(1, post_patch_width) - .flatten() - ) - col_ids = ( - torch.arange(post_patch_width, dtype=torch.int32, device=device) - .view(1, -1) - .repeat(post_patch_height, 1) - .flatten() - ) - position_ids[i, cap_seq_len:seq_len, 1] = row_ids - position_ids[i, cap_seq_len:seq_len, 2] = col_ids - - # Get combined rotary embeddings - freqs_cis = self._get_freqs_cis(position_ids) - - # create separate rotary embeddings for captions and images - cap_freqs_cis = torch.zeros( - batch_size, encoder_seq_len, freqs_cis.shape[-1], device=device, dtype=freqs_cis.dtype - ) - img_freqs_cis = torch.zeros( - batch_size, image_seq_len, freqs_cis.shape[-1], device=device, dtype=freqs_cis.dtype - ) - - for i, (cap_seq_len, seq_len) in enumerate(zip(l_effective_cap_len, seq_lengths)): - cap_freqs_cis[i, :cap_seq_len] = freqs_cis[i, :cap_seq_len] - img_freqs_cis[i, :image_seq_len] = freqs_cis[i, cap_seq_len:seq_len] - - # image patch embeddings - hidden_states = ( - hidden_states.view(batch_size, channels, post_patch_height, p, post_patch_width, p) - .permute(0, 2, 4, 3, 5, 1) - .flatten(3) - .flatten(1, 2) - ) - - return hidden_states, cap_freqs_cis, img_freqs_cis, freqs_cis, l_effective_cap_len, seq_lengths - - -class Lumina2Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): - r""" - Lumina2NextDiT: Diffusion model with a Transformer backbone. - - Parameters: - sample_size (`int`): The width of the latent images. This is fixed during training since - it is used to learn a number of position embeddings. - patch_size (`int`, *optional*, (`int`, *optional*, defaults to 2): - The size of each patch in the image. This parameter defines the resolution of patches fed into the model. - in_channels (`int`, *optional*, defaults to 4): - The number of input channels for the model. Typically, this matches the number of channels in the input - images. - hidden_size (`int`, *optional*, defaults to 4096): - The dimensionality of the hidden layers in the model. This parameter determines the width of the model's - hidden representations. - num_layers (`int`, *optional*, default to 32): - The number of layers in the model. This defines the depth of the neural network. - num_attention_heads (`int`, *optional*, defaults to 32): - The number of attention heads in each attention layer. This parameter specifies how many separate attention - mechanisms are used. - num_kv_heads (`int`, *optional*, defaults to 8): - The number of key-value heads in the attention mechanism, if different from the number of attention heads. - If None, it defaults to num_attention_heads. - multiple_of (`int`, *optional*, defaults to 256): - A factor that the hidden size should be a multiple of. This can help optimize certain hardware - configurations. - ffn_dim_multiplier (`float`, *optional*): - A multiplier for the dimensionality of the feed-forward network. If None, it uses a default value based on - the model configuration. - norm_eps (`float`, *optional*, defaults to 1e-5): - A small value added to the denominator for numerical stability in normalization layers. - scaling_factor (`float`, *optional*, defaults to 1.0): - A scaling factor applied to certain parameters or layers in the model. This can be used for adjusting the - overall scale of the model's operations. - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["Lumina2TransformerBlock"] - _skip_layerwise_casting_patterns = ["x_embedder", "norm"] - - @register_to_config - def __init__( - self, - sample_size: int = 128, - patch_size: int = 2, - in_channels: int = 16, - out_channels: Optional[int] = 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, - 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), - cap_feat_dim: int = 1024, - ) -> None: - super().__init__() - self.out_channels = out_channels or in_channels - - # 1. Positional, patch & conditional embeddings - self.rope_embedder = Lumina2RotaryPosEmbed( - theta=10000, axes_dim=axes_dim_rope, axes_lens=axes_lens, patch_size=patch_size - ) - - self.x_embedder = nn.Linear(in_features=patch_size * patch_size * in_channels, out_features=hidden_size) - - self.time_caption_embed = Lumina2CombinedTimestepCaptionEmbedding( - hidden_size=hidden_size, cap_feat_dim=cap_feat_dim, norm_eps=norm_eps - ) - - # 2. Noise and context refinement blocks - self.noise_refiner = nn.ModuleList( - [ - Lumina2TransformerBlock( - hidden_size, - num_attention_heads, - num_kv_heads, - multiple_of, - ffn_dim_multiplier, - norm_eps, - modulation=True, - ) - for _ in range(num_refiner_layers) - ] - ) - - self.context_refiner = nn.ModuleList( - [ - Lumina2TransformerBlock( - hidden_size, - num_attention_heads, - num_kv_heads, - multiple_of, - ffn_dim_multiplier, - norm_eps, - modulation=False, - ) - for _ in range(num_refiner_layers) - ] - ) - - # 3. Transformer blocks - self.layers = nn.ModuleList( - [ - Lumina2TransformerBlock( - hidden_size, - num_attention_heads, - num_kv_heads, - multiple_of, - ffn_dim_multiplier, - norm_eps, - modulation=True, - ) - for _ in range(num_layers) - ] - ) - - # 4. Output norm & projection - self.norm_out = LuminaLayerNormContinuous( - embedding_dim=hidden_size, - conditioning_embedding_dim=min(hidden_size, 1024), - elementwise_affine=False, - eps=1e-6, - bias=True, - out_dim=patch_size * patch_size * self.out_channels, - ) - - self.gradient_checkpointing = False - - 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, - 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." - ) - - # 1. Condition, positional & patch embedding - batch_size, _, height, width = hidden_states.shape - - temb, encoder_hidden_states = self.time_caption_embed(hidden_states, timestep, encoder_hidden_states) - - ( - hidden_states, - context_rotary_emb, - noise_rotary_emb, - rotary_emb, - encoder_seq_lengths, - seq_lengths, - ) = self.rope_embedder(hidden_states, encoder_attention_mask) - - hidden_states = self.x_embedder(hidden_states) - - # 2. Context & noise refinement - for layer in self.context_refiner: - encoder_hidden_states = layer(encoder_hidden_states, encoder_attention_mask, context_rotary_emb) - - for layer in self.noise_refiner: - hidden_states = layer(hidden_states, None, noise_rotary_emb, temb) - - # 3. Joint Transformer blocks - max_seq_len = max(seq_lengths) - use_mask = len(set(seq_lengths)) > 1 - - attention_mask = hidden_states.new_zeros(batch_size, max_seq_len, dtype=torch.bool) - joint_hidden_states = hidden_states.new_zeros(batch_size, max_seq_len, self.config.hidden_size) - for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)): - attention_mask[i, :seq_len] = True - joint_hidden_states[i, :encoder_seq_len] = encoder_hidden_states[i, :encoder_seq_len] - joint_hidden_states[i, encoder_seq_len:seq_len] = hidden_states[i] - - hidden_states = joint_hidden_states - - for layer in self.layers: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - layer, hidden_states, attention_mask if use_mask else None, rotary_emb, temb - ) - else: - hidden_states = layer(hidden_states, attention_mask if use_mask else None, rotary_emb, temb) - - # 4. Output norm & projection - hidden_states = self.norm_out(hidden_states, temb) - - # 5. Unpatchify - p = self.config.patch_size - output = [] - for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)): - output.append( - hidden_states[i][encoder_seq_len:seq_len] - .view(height // p, width // p, p, p, self.out_channels) - .permute(4, 0, 2, 1, 3) - .flatten(3, 4) - .flatten(1, 2) - ) - 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/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_mochi.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_mochi.py deleted file mode 100644 index 63911fe7c10d037e4d97b8ca95a7e9ea151391ee..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_mochi.py +++ /dev/null @@ -1,488 +0,0 @@ -# Copyright 2025 The Genmo 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 Any, Dict, Optional, Tuple - -import torch -import torch.nn as nn - -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.torch_utils import maybe_allow_in_graph -from ..attention import FeedForward -from ..attention_processor import MochiAttention, MochiAttnProcessor2_0 -from ..cache_utils import CacheMixin -from ..embeddings import MochiCombinedTimestepCaptionEmbedding, PatchEmbed -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormContinuous, RMSNorm - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class MochiModulatedRMSNorm(nn.Module): - def __init__(self, eps: float): - super().__init__() - - self.eps = eps - self.norm = RMSNorm(0, eps, False) - - def forward(self, hidden_states, scale=None): - hidden_states_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - - hidden_states = self.norm(hidden_states) - - if scale is not None: - hidden_states = hidden_states * scale - - hidden_states = hidden_states.to(hidden_states_dtype) - - return hidden_states - - -class MochiLayerNormContinuous(nn.Module): - def __init__( - self, - embedding_dim: int, - conditioning_embedding_dim: int, - eps=1e-5, - bias=True, - ): - super().__init__() - - # AdaLN - self.silu = nn.SiLU() - self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias) - self.norm = MochiModulatedRMSNorm(eps=eps) - - def forward( - self, - x: torch.Tensor, - conditioning_embedding: torch.Tensor, - ) -> torch.Tensor: - input_dtype = x.dtype - - # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT) - scale = self.linear_1(self.silu(conditioning_embedding).to(x.dtype)) - x = self.norm(x, (1 + scale.unsqueeze(1).to(torch.float32))) - - return x.to(input_dtype) - - -class MochiRMSNormZero(nn.Module): - r""" - Adaptive RMS Norm used in Mochi. - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - """ - - def __init__( - self, embedding_dim: int, hidden_dim: int, eps: float = 1e-5, elementwise_affine: bool = False - ) -> None: - super().__init__() - - self.silu = nn.SiLU() - self.linear = nn.Linear(embedding_dim, hidden_dim) - self.norm = RMSNorm(0, eps, False) - - def forward( - self, hidden_states: torch.Tensor, emb: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - hidden_states_dtype = hidden_states.dtype - - emb = self.linear(self.silu(emb)) - scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1) - hidden_states = self.norm(hidden_states.to(torch.float32)) * (1 + scale_msa[:, None].to(torch.float32)) - hidden_states = hidden_states.to(hidden_states_dtype) - - return hidden_states, gate_msa, scale_mlp, gate_mlp - - -@maybe_allow_in_graph -class MochiTransformerBlock(nn.Module): - r""" - Transformer block used in [Mochi](https://huggingface.co/genmo/mochi-1-preview). - - 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 `"swiglu"`): - Activation function to use in feed-forward. - context_pre_only (`bool`, defaults to `False`): - Whether or not to process context-related conditions with additional layers. - eps (`float`, defaults to `1e-6`): - Epsilon value for normalization layers. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - pooled_projection_dim: int, - qk_norm: str = "rms_norm", - activation_fn: str = "swiglu", - context_pre_only: bool = False, - eps: float = 1e-6, - ) -> None: - super().__init__() - - self.context_pre_only = context_pre_only - self.ff_inner_dim = (4 * dim * 2) // 3 - self.ff_context_inner_dim = (4 * pooled_projection_dim * 2) // 3 - - self.norm1 = MochiRMSNormZero(dim, 4 * dim, eps=eps, elementwise_affine=False) - - if not context_pre_only: - self.norm1_context = MochiRMSNormZero(dim, 4 * pooled_projection_dim, eps=eps, elementwise_affine=False) - else: - self.norm1_context = MochiLayerNormContinuous( - embedding_dim=pooled_projection_dim, - conditioning_embedding_dim=dim, - eps=eps, - ) - - self.attn1 = MochiAttention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - bias=False, - added_kv_proj_dim=pooled_projection_dim, - added_proj_bias=False, - out_dim=dim, - out_context_dim=pooled_projection_dim, - context_pre_only=context_pre_only, - processor=MochiAttnProcessor2_0(), - eps=1e-5, - ) - - # TODO(aryan): norm_context layers are not needed when `context_pre_only` is True - self.norm2 = MochiModulatedRMSNorm(eps=eps) - self.norm2_context = MochiModulatedRMSNorm(eps=eps) if not self.context_pre_only else None - - self.norm3 = MochiModulatedRMSNorm(eps) - self.norm3_context = MochiModulatedRMSNorm(eps=eps) if not self.context_pre_only else None - - self.ff = FeedForward(dim, inner_dim=self.ff_inner_dim, activation_fn=activation_fn, bias=False) - self.ff_context = None - if not context_pre_only: - self.ff_context = FeedForward( - pooled_projection_dim, - inner_dim=self.ff_context_inner_dim, - activation_fn=activation_fn, - bias=False, - ) - - self.norm4 = MochiModulatedRMSNorm(eps=eps) - self.norm4_context = MochiModulatedRMSNorm(eps=eps) - - def forward( - self, - hidden_states: torch.Tensor, - 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]: - norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb) - - if not self.context_pre_only: - norm_encoder_hidden_states, enc_gate_msa, enc_scale_mlp, enc_gate_mlp = self.norm1_context( - encoder_hidden_states, temb - ) - else: - norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, temb) - - attn_hidden_states, context_attn_hidden_states = self.attn1( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - image_rotary_emb=image_rotary_emb, - attention_mask=encoder_attention_mask, - ) - - hidden_states = hidden_states + self.norm2(attn_hidden_states, torch.tanh(gate_msa).unsqueeze(1)) - norm_hidden_states = self.norm3(hidden_states, (1 + scale_mlp.unsqueeze(1).to(torch.float32))) - ff_output = self.ff(norm_hidden_states) - hidden_states = hidden_states + self.norm4(ff_output, torch.tanh(gate_mlp).unsqueeze(1)) - - if not self.context_pre_only: - encoder_hidden_states = encoder_hidden_states + self.norm2_context( - context_attn_hidden_states, torch.tanh(enc_gate_msa).unsqueeze(1) - ) - norm_encoder_hidden_states = self.norm3_context( - encoder_hidden_states, (1 + enc_scale_mlp.unsqueeze(1).to(torch.float32)) - ) - context_ff_output = self.ff_context(norm_encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states + self.norm4_context( - context_ff_output, torch.tanh(enc_gate_mlp).unsqueeze(1) - ) - - return hidden_states, encoder_hidden_states - - -class MochiRoPE(nn.Module): - r""" - RoPE implementation used in [Mochi](https://huggingface.co/genmo/mochi-1-preview). - - Args: - base_height (`int`, defaults to `192`): - Base height used to compute interpolation scale for rotary positional embeddings. - base_width (`int`, defaults to `192`): - Base width used to compute interpolation scale for rotary positional embeddings. - """ - - def __init__(self, base_height: int = 192, base_width: int = 192) -> None: - super().__init__() - - self.target_area = base_height * base_width - - def _centers(self, start, stop, num, device, dtype) -> torch.Tensor: - edges = torch.linspace(start, stop, num + 1, device=device, dtype=dtype) - return (edges[:-1] + edges[1:]) / 2 - - def _get_positions( - self, - num_frames: int, - height: int, - width: int, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - ) -> torch.Tensor: - scale = (self.target_area / (height * width)) ** 0.5 - - t = torch.arange(num_frames, device=device, dtype=dtype) - h = self._centers(-height * scale / 2, height * scale / 2, height, device, dtype) - w = self._centers(-width * scale / 2, width * scale / 2, width, device, dtype) - - grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij") - - positions = torch.stack([grid_t, grid_h, grid_w], dim=-1).view(-1, 3) - return positions - - def _create_rope(self, freqs: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: - with torch.autocast(freqs.device.type, torch.float32): - # Always run ROPE freqs computation in FP32 - freqs = torch.einsum("nd,dhf->nhf", pos.to(torch.float32), freqs.to(torch.float32)) - - freqs_cos = torch.cos(freqs) - freqs_sin = torch.sin(freqs) - return freqs_cos, freqs_sin - - def forward( - self, - pos_frequencies: torch.Tensor, - num_frames: int, - height: int, - width: int, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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 - - -@maybe_allow_in_graph -class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin): - r""" - A Transformer model for video-like data introduced in [Mochi](https://huggingface.co/genmo/mochi-1-preview). - - Args: - patch_size (`int`, defaults to `2`): - The size of the patches to use in the patch embedding layer. - num_attention_heads (`int`, defaults to `24`): - The number of heads to use for multi-head attention. - attention_head_dim (`int`, defaults to `128`): - The number of channels in each head. - num_layers (`int`, defaults to `48`): - The number of layers of Transformer blocks to use. - in_channels (`int`, defaults to `12`): - The number of channels in the input. - out_channels (`int`, *optional*, defaults to `None`): - The number of channels in the output. - qk_norm (`str`, defaults to `"rms_norm"`): - The normalization layer to use. - text_embed_dim (`int`, defaults to `4096`): - Input dimension of text embeddings from the text encoder. - time_embed_dim (`int`, defaults to `256`): - Output dimension of timestep embeddings. - activation_fn (`str`, defaults to `"swiglu"`): - Activation function to use in feed-forward. - max_sequence_length (`int`, defaults to `256`): - The maximum sequence length of text embeddings supported. - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["MochiTransformerBlock"] - _skip_layerwise_casting_patterns = ["patch_embed", "norm"] - - @register_to_config - def __init__( - self, - patch_size: int = 2, - num_attention_heads: int = 24, - attention_head_dim: int = 128, - num_layers: int = 48, - pooled_projection_dim: int = 1536, - in_channels: int = 12, - out_channels: Optional[int] = None, - qk_norm: str = "rms_norm", - text_embed_dim: int = 4096, - time_embed_dim: int = 256, - activation_fn: str = "swiglu", - max_sequence_length: int = 256, - ) -> None: - super().__init__() - - inner_dim = num_attention_heads * attention_head_dim - out_channels = out_channels or in_channels - - self.patch_embed = PatchEmbed( - patch_size=patch_size, - in_channels=in_channels, - embed_dim=inner_dim, - pos_embed_type=None, - ) - - self.time_embed = MochiCombinedTimestepCaptionEmbedding( - embedding_dim=inner_dim, - pooled_projection_dim=pooled_projection_dim, - text_embed_dim=text_embed_dim, - time_embed_dim=time_embed_dim, - num_attention_heads=8, - ) - - self.pos_frequencies = nn.Parameter(torch.full((3, num_attention_heads, attention_head_dim // 2), 0.0)) - self.rope = MochiRoPE() - - self.transformer_blocks = nn.ModuleList( - [ - MochiTransformerBlock( - dim=inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - pooled_projection_dim=pooled_projection_dim, - qk_norm=qk_norm, - activation_fn=activation_fn, - context_pre_only=i == num_layers - 1, - ) - for i in range(num_layers) - ] - ) - - self.norm_out = AdaLayerNormContinuous( - inner_dim, - inner_dim, - elementwise_affine=False, - eps=1e-6, - norm_type="layer_norm", - ) - self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels) - - self.gradient_checkpointing = False - - 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, - 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 - - post_patch_height = height // p - post_patch_width = width // p - - temb, encoder_hidden_states = self.time_embed( - timestep, - encoder_hidden_states, - encoder_attention_mask, - hidden_dtype=hidden_states.dtype, - ) - - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) - hidden_states = self.patch_embed(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).flatten(1, 2) - - image_rotary_emb = self.rope( - self.pos_frequencies, - num_frames, - post_patch_height, - post_patch_width, - device=hidden_states.device, - dtype=torch.float32, - ) - - for i, 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, - encoder_attention_mask, - image_rotary_emb, - ) - else: - hidden_states, encoder_hidden_states = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - temb=temb, - encoder_attention_mask=encoder_attention_mask, - image_rotary_emb=image_rotary_emb, - ) - hidden_states = self.norm_out(hidden_states, temb) - hidden_states = self.proj_out(hidden_states) - - hidden_states = hidden_states.reshape(batch_size, num_frames, post_patch_height, post_patch_width, p, p, -1) - 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/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_omnigen.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_omnigen.py deleted file mode 100644 index 6939cac0a3a705dee4912c18ab6d94f229f49859..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_omnigen.py +++ /dev/null @@ -1,469 +0,0 @@ -# Copyright 2025 OmniGen 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 typing import Dict, List, Optional, Tuple, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import logging -from ..attention_processor import Attention -from ..embeddings import TimestepEmbedding, Timesteps, get_2d_sincos_pos_embed -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNorm, RMSNorm - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class OmniGenFeedForward(nn.Module): - def __init__(self, hidden_size: int, intermediate_size: int): - super().__init__() - - self.gate_up_proj = nn.Linear(hidden_size, 2 * intermediate_size, bias=False) - self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) - self.activation_fn = nn.SiLU() - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - up_states = self.gate_up_proj(hidden_states) - gate, up_states = up_states.chunk(2, dim=-1) - up_states = up_states * self.activation_fn(gate) - return self.down_proj(up_states) - - -class OmniGenPatchEmbed(nn.Module): - def __init__( - self, - patch_size: int = 2, - in_channels: int = 4, - embed_dim: int = 768, - bias: bool = True, - interpolation_scale: float = 1, - pos_embed_max_size: int = 192, - base_size: int = 64, - ): - super().__init__() - - self.output_image_proj = nn.Conv2d( - in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias - ) - self.input_image_proj = nn.Conv2d( - in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias - ) - - self.patch_size = patch_size - self.interpolation_scale = interpolation_scale - self.pos_embed_max_size = pos_embed_max_size - - pos_embed = get_2d_sincos_pos_embed( - embed_dim, - self.pos_embed_max_size, - base_size=base_size, - interpolation_scale=self.interpolation_scale, - output_type="pt", - ) - self.register_buffer("pos_embed", pos_embed.float().unsqueeze(0), persistent=True) - - def _cropped_pos_embed(self, height, width): - """Crops positional embeddings for SD3 compatibility.""" - if self.pos_embed_max_size is None: - raise ValueError("`pos_embed_max_size` must be set for cropping.") - - height = height // self.patch_size - width = width // self.patch_size - if height > self.pos_embed_max_size: - raise ValueError( - f"Height ({height}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}." - ) - if width > self.pos_embed_max_size: - raise ValueError( - f"Width ({width}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}." - ) - - top = (self.pos_embed_max_size - height) // 2 - left = (self.pos_embed_max_size - width) // 2 - spatial_pos_embed = self.pos_embed.reshape(1, self.pos_embed_max_size, self.pos_embed_max_size, -1) - spatial_pos_embed = spatial_pos_embed[:, top : top + height, left : left + width, :] - spatial_pos_embed = spatial_pos_embed.reshape(1, -1, spatial_pos_embed.shape[-1]) - return spatial_pos_embed - - def _patch_embeddings(self, hidden_states: torch.Tensor, is_input_image: bool) -> torch.Tensor: - if is_input_image: - hidden_states = self.input_image_proj(hidden_states) - else: - hidden_states = self.output_image_proj(hidden_states) - hidden_states = hidden_states.flatten(2).transpose(1, 2) - return hidden_states - - def forward( - self, hidden_states: torch.Tensor, is_input_image: bool, padding_latent: torch.Tensor = None - ) -> torch.Tensor: - if isinstance(hidden_states, list): - if padding_latent is None: - padding_latent = [None] * len(hidden_states) - patched_latents = [] - for sub_latent, padding in zip(hidden_states, padding_latent): - height, width = sub_latent.shape[-2:] - sub_latent = self._patch_embeddings(sub_latent, is_input_image) - pos_embed = self._cropped_pos_embed(height, width) - sub_latent = sub_latent + pos_embed - if padding is not None: - sub_latent = torch.cat([sub_latent, padding.to(sub_latent.device)], dim=-2) - patched_latents.append(sub_latent) - else: - height, width = hidden_states.shape[-2:] - pos_embed = self._cropped_pos_embed(height, width) - hidden_states = self._patch_embeddings(hidden_states, is_input_image) - patched_latents = hidden_states + pos_embed - - return patched_latents - - -class OmniGenSuScaledRotaryEmbedding(nn.Module): - def __init__( - self, dim, max_position_embeddings=131072, original_max_position_embeddings=4096, base=10000, rope_scaling=None - ): - super().__init__() - - self.dim = dim - self.max_position_embeddings = max_position_embeddings - self.base = base - - inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim)) - self.register_buffer("inv_freq", tensor=inv_freq, persistent=False) - - self.short_factor = rope_scaling["short_factor"] - self.long_factor = rope_scaling["long_factor"] - self.original_max_position_embeddings = original_max_position_embeddings - - def forward(self, hidden_states, position_ids): - seq_len = torch.max(position_ids) + 1 - if seq_len > self.original_max_position_embeddings: - ext_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=hidden_states.device) - else: - ext_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=hidden_states.device) - - inv_freq_shape = ( - torch.arange(0, self.dim, 2, dtype=torch.int64, device=hidden_states.device).float() / self.dim - ) - self.inv_freq = 1.0 / (ext_factors * self.base**inv_freq_shape) - - inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) - position_ids_expanded = position_ids[:, None, :].float() - - # Force float32 since bfloat16 loses precision on long contexts - # See https://github.com/huggingface/transformers/pull/29285 - device_type = hidden_states.device.type - device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" - with torch.autocast(device_type=device_type, enabled=False): - freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) - emb = torch.cat((freqs, freqs), dim=-1)[0] - - scale = self.max_position_embeddings / self.original_max_position_embeddings - if scale <= 1.0: - scaling_factor = 1.0 - else: - scaling_factor = math.sqrt(1 + math.log(scale) / math.log(self.original_max_position_embeddings)) - - cos = emb.cos() * scaling_factor - sin = emb.sin() * scaling_factor - return cos, sin - - -class OmniGenAttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is - used in the OmniGen model. - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_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: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - batch_size, sequence_length, _ = hidden_states.shape - - # Get Query-Key-Value Pair - query = attn.to_q(hidden_states) - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - bsz, q_len, query_dim = query.size() - inner_dim = key.shape[-1] - head_dim = query_dim // attn.heads - - # Get key-value heads - kv_heads = inner_dim // head_dim - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, kv_heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, kv_heads, head_dim).transpose(1, 2) - - # Apply RoPE if needed - if image_rotary_emb is not None: - from ..embeddings import apply_rotary_emb - - query = apply_rotary_emb(query, image_rotary_emb, use_real_unbind_dim=-2) - key = apply_rotary_emb(key, image_rotary_emb, use_real_unbind_dim=-2) - - hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask) - hidden_states = hidden_states.transpose(1, 2).type_as(query) - hidden_states = hidden_states.reshape(bsz, q_len, attn.out_dim) - hidden_states = attn.to_out[0](hidden_states) - return hidden_states - - -class OmniGenBlock(nn.Module): - def __init__( - self, - hidden_size: int, - num_attention_heads: int, - num_key_value_heads: int, - intermediate_size: int, - rms_norm_eps: float, - ) -> None: - super().__init__() - - self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) - self.self_attn = Attention( - query_dim=hidden_size, - cross_attention_dim=hidden_size, - dim_head=hidden_size // num_attention_heads, - heads=num_attention_heads, - kv_heads=num_key_value_heads, - bias=False, - out_dim=hidden_size, - out_bias=False, - processor=OmniGenAttnProcessor2_0(), - ) - self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) - self.mlp = OmniGenFeedForward(hidden_size, intermediate_size) - - def forward( - self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, image_rotary_emb: torch.Tensor - ) -> torch.Tensor: - # 1. Attention - norm_hidden_states = self.input_layernorm(hidden_states) - attn_output = self.self_attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_hidden_states, - attention_mask=attention_mask, - image_rotary_emb=image_rotary_emb, - ) - hidden_states = hidden_states + attn_output - - # 2. Feed Forward - norm_hidden_states = self.post_attention_layernorm(hidden_states) - ff_output = self.mlp(norm_hidden_states) - hidden_states = hidden_states + ff_output - return hidden_states - - -class OmniGenTransformer2DModel(ModelMixin, ConfigMixin): - """ - The Transformer model introduced in OmniGen (https://huggingface.co/papers/2409.11340). - - Parameters: - in_channels (`int`, defaults to `4`): - The number of channels in the input. - patch_size (`int`, defaults to `2`): - The size of the spatial patches to use in the patch embedding layer. - hidden_size (`int`, defaults to `3072`): - The dimensionality of the hidden layers in the model. - rms_norm_eps (`float`, defaults to `1e-5`): - Eps for RMSNorm layer. - num_attention_heads (`int`, defaults to `32`): - The number of heads to use for multi-head attention. - num_key_value_heads (`int`, defaults to `32`): - The number of heads to use for keys and values in multi-head attention. - intermediate_size (`int`, defaults to `8192`): - Dimension of the hidden layer in FeedForward layers. - num_layers (`int`, default to `32`): - The number of layers of transformer blocks to use. - pad_token_id (`int`, default to `32000`): - The id of the padding token. - vocab_size (`int`, default to `32064`): - 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): - 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. - time_step_dim (`int`, default to `256`): - Output dimension of timestep embeddings. - flip_sin_to_cos (`bool`, default to `True`): - Whether to flip the sin and cos in the positional embeddings when preparing timestep embeddings. - downscale_freq_shift (`int`, default to `0`): - The frequency shift to use when downscaling the timestep embeddings. - timestep_activation_fn (`str`, default to `silu`): - The activation function to use for the timestep embeddings. - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["OmniGenBlock"] - _skip_layerwise_casting_patterns = ["patch_embedding", "embed_tokens", "norm"] - - @register_to_config - def __init__( - self, - in_channels: int = 4, - patch_size: int = 2, - hidden_size: int = 3072, - rms_norm_eps: float = 1e-5, - num_attention_heads: int = 32, - num_key_value_heads: int = 32, - intermediate_size: int = 8192, - num_layers: int = 32, - pad_token_id: int = 32000, - vocab_size: int = 32064, - max_position_embeddings: int = 131072, - original_max_position_embeddings: int = 4096, - rope_base: int = 10000, - rope_scaling: Dict = None, - pos_embed_max_size: int = 192, - time_step_dim: int = 256, - flip_sin_to_cos: bool = True, - downscale_freq_shift: int = 0, - timestep_activation_fn: str = "silu", - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = in_channels - - self.patch_embedding = OmniGenPatchEmbed( - patch_size=patch_size, - in_channels=in_channels, - embed_dim=hidden_size, - pos_embed_max_size=pos_embed_max_size, - ) - - self.time_proj = Timesteps(time_step_dim, flip_sin_to_cos, downscale_freq_shift) - self.time_token = TimestepEmbedding(time_step_dim, hidden_size, timestep_activation_fn) - self.t_embedder = TimestepEmbedding(time_step_dim, hidden_size, timestep_activation_fn) - - self.embed_tokens = nn.Embedding(vocab_size, hidden_size, pad_token_id) - self.rope = OmniGenSuScaledRotaryEmbedding( - hidden_size // num_attention_heads, - max_position_embeddings=max_position_embeddings, - original_max_position_embeddings=original_max_position_embeddings, - base=rope_base, - rope_scaling=rope_scaling, - ) - - self.layers = nn.ModuleList( - [ - OmniGenBlock(hidden_size, num_attention_heads, num_key_value_heads, intermediate_size, rms_norm_eps) - for _ in range(num_layers) - ] - ) - - self.norm = RMSNorm(hidden_size, eps=rms_norm_eps) - self.norm_out = AdaLayerNorm(hidden_size, norm_elementwise_affine=False, norm_eps=1e-6, chunk_dim=1) - self.proj_out = nn.Linear(hidden_size, patch_size * patch_size * self.out_channels, bias=True) - - 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]: - if input_ids is None: - return None - - input_img_latents = [x.to(self.dtype) for x in input_img_latents] - condition_tokens = self.embed_tokens(input_ids) - input_img_inx = 0 - input_image_tokens = self.patch_embedding(input_img_latents, is_input_image=True) - for b_inx in input_image_sizes.keys(): - for start_inx, end_inx in input_image_sizes[b_inx]: - # replace the placeholder in text tokens with the image embedding. - condition_tokens[b_inx, start_inx:end_inx] = input_image_tokens[input_img_inx].to( - condition_tokens.dtype - ) - input_img_inx += 1 - return condition_tokens - - def forward( - self, - hidden_states: torch.Tensor, - timestep: Union[int, float, torch.FloatTensor], - input_ids: torch.Tensor, - 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]]: - batch_size, num_channels, height, width = hidden_states.shape - p = self.config.patch_size - post_patch_height, post_patch_width = height // p, width // p - - # 1. Patch & Timestep & Conditional Embedding - hidden_states = self.patch_embedding(hidden_states, is_input_image=False) - num_tokens_for_output_image = hidden_states.size(1) - - timestep_proj = self.time_proj(timestep).type_as(hidden_states) - time_token = self.time_token(timestep_proj).unsqueeze(1) - temb = self.t_embedder(timestep_proj) - - condition_tokens = self._get_multimodal_embeddings(input_ids, input_img_latents, input_image_sizes) - if condition_tokens is not None: - hidden_states = torch.cat([condition_tokens, time_token, hidden_states], dim=1) - else: - hidden_states = torch.cat([time_token, hidden_states], dim=1) - - seq_length = hidden_states.size(1) - position_ids = position_ids.view(-1, seq_length).long() - - # 2. Attention mask preprocessing - if attention_mask is not None and attention_mask.dim() == 3: - dtype = hidden_states.dtype - min_dtype = torch.finfo(dtype).min - attention_mask = (1 - attention_mask) * min_dtype - attention_mask = attention_mask.unsqueeze(1).type_as(hidden_states) - - # 3. Rotary position embedding - image_rotary_emb = self.rope(hidden_states, position_ids) - - # 4. Transformer blocks - for block in self.layers: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - block, hidden_states, attention_mask, image_rotary_emb - ) - else: - hidden_states = block(hidden_states, attention_mask=attention_mask, image_rotary_emb=image_rotary_emb) - - # 5. Output norm & projection - hidden_states = self.norm(hidden_states) - hidden_states = hidden_states[:, -num_tokens_for_output_image:] - hidden_states = self.norm_out(hidden_states, temb=temb) - hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.reshape(batch_size, post_patch_height, post_patch_width, p, p, -1) - output = hidden_states.permute(0, 5, 1, 3, 2, 4).flatten(4, 5).flatten(2, 3) - - if not return_dict: - return (output,) - return Transformer2DModelOutput(sample=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_sd3.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_sd3.py deleted file mode 100644 index edf77a7df793c7fbaa112860278d0b47d9ad7a7e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_sd3.py +++ /dev/null @@ -1,431 +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 typing import Any, Dict, List, Optional, Tuple, Union - -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.torch_utils import maybe_allow_in_graph -from ..attention import FeedForward, JointTransformerBlock -from ..attention_processor import ( - Attention, - AttentionProcessor, - FusedJointAttnProcessor2_0, - JointAttnProcessor2_0, -) -from ..embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import AdaLayerNormContinuous, AdaLayerNormZero - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@maybe_allow_in_graph -class SD3SingleTransformerBlock(nn.Module): - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - ): - super().__init__() - - self.norm1 = AdaLayerNormZero(dim) - self.attn = Attention( - query_dim=dim, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=dim, - bias=True, - processor=JointAttnProcessor2_0(), - eps=1e-6, - ) - - self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) - self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") - - def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor): - # 1. Attention - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) - attn_output = self.attn(hidden_states=norm_hidden_states, encoder_hidden_states=None) - attn_output = gate_msa.unsqueeze(1) * attn_output - hidden_states = hidden_states + attn_output - - # 2. Feed Forward - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1) - ff_output = self.ff(norm_hidden_states) - ff_output = gate_mlp.unsqueeze(1) * ff_output - hidden_states = hidden_states + ff_output - - return hidden_states - - -class SD3Transformer2DModel( - ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, SD3Transformer2DLoadersMixin -): - """ - The Transformer model introduced in [Stable Diffusion 3](https://huggingface.co/papers/2403.03206). - - Parameters: - sample_size (`int`, defaults to `128`): - The width/height of the latents. This is fixed during training since it is used to learn a number of - position embeddings. - patch_size (`int`, defaults to `2`): - Patch size to turn the input data into small patches. - in_channels (`int`, defaults to `16`): - The number of latent channels in the input. - num_layers (`int`, defaults to `18`): - The number of layers of transformer blocks to use. - attention_head_dim (`int`, defaults to `64`): - The number of channels in each head. - num_attention_heads (`int`, defaults to `18`): - The number of heads to use for multi-head attention. - joint_attention_dim (`int`, defaults to `4096`): - The embedding dimension to use for joint text-image attention. - caption_projection_dim (`int`, defaults to `1152`): - The embedding dimension of caption embeddings. - pooled_projection_dim (`int`, defaults to `2048`): - The embedding dimension of pooled text projections. - out_channels (`int`, defaults to `16`): - 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 `()`): - 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. - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["JointTransformerBlock"] - _skip_layerwise_casting_patterns = ["pos_embed", "norm"] - - @register_to_config - 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, - 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, - ): - super().__init__() - self.out_channels = out_channels if out_channels is not None else in_channels - self.inner_dim = num_attention_heads * attention_head_dim - - self.pos_embed = PatchEmbed( - height=sample_size, - width=sample_size, - patch_size=patch_size, - in_channels=in_channels, - embed_dim=self.inner_dim, - pos_embed_max_size=pos_embed_max_size, # hard-code for now. - ) - self.time_text_embed = CombinedTimestepTextProjEmbeddings( - embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim - ) - self.context_embedder = nn.Linear(joint_attention_dim, caption_projection_dim) - - self.transformer_blocks = nn.ModuleList( - [ - JointTransformerBlock( - dim=self.inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - context_pre_only=i == num_layers - 1, - qk_norm=qk_norm, - use_dual_attention=True if i in dual_attention_layers else False, - ) - for i in range(num_layers) - ] - ) - - self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6) - self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True) - - 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: - """ - Sets the attention processor to use [feed forward - chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). - - Parameters: - chunk_size (`int`, *optional*): - The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually - over each tensor of dim=`dim`. - dim (`int`, *optional*, defaults to `0`): - The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch) - or dim=1 (sequence length). - """ - if dim not in [0, 1]: - raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}") - - # By default chunk size is 1 - chunk_size = chunk_size or 1 - - def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int): - if hasattr(module, "set_chunk_feed_forward"): - module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) - - for child in module.children(): - fn_recursive_feed_forward(child, chunk_size, dim) - - for module in self.children(): - fn_recursive_feed_forward(module, chunk_size, dim) - - # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.disable_forward_chunking - def disable_forward_chunking(self): - def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int): - if hasattr(module, "set_chunk_feed_forward"): - module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) - - for child in module.children(): - fn_recursive_feed_forward(child, chunk_size, dim) - - for module in self.children(): - fn_recursive_feed_forward(module, None, 0) - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedJointAttnProcessor2_0 - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - self.set_attn_processor(FusedJointAttnProcessor2_0()) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - if self.original_attn_processors is not None: - self.set_attn_processor(self.original_attn_processors) - - 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, - return_dict: bool = True, - skip_layers: Optional[List[int]] = None, - ) -> Union[torch.Tensor, Transformer2DModelOutput]: - """ - The [`SD3Transformer2DModel`] forward method. - - Args: - hidden_states (`torch.Tensor` of shape `(batch size, channel, height, width)`): - Input `hidden_states`. - encoder_hidden_states (`torch.Tensor` of shape `(batch size, sequence_len, embed_dims)`): - Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. - pooled_projections (`torch.Tensor` of shape `(batch_size, projection_dim)`): - Embeddings projected from the embeddings of input conditions. - 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. - skip_layers (`list` of `int`, *optional*): - A list of layer indices to skip during the forward pass. - - 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 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:] - - hidden_states = self.pos_embed(hidden_states) # takes care of adding positional embeddings too. - temb = self.time_text_embed(timestep, pooled_projections) - encoder_hidden_states = self.context_embedder(encoder_hidden_states) - - 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, ip_temb = self.image_proj(ip_adapter_image_embeds, timestep) - - joint_attention_kwargs.update(ip_hidden_states=ip_hidden_states, temb=ip_temb) - - for index_block, block in enumerate(self.transformer_blocks): - # Skip specified layers - is_skip = True if skip_layers is not None and index_block in skip_layers else False - - if torch.is_grad_enabled() and self.gradient_checkpointing and not is_skip: - encoder_hidden_states, hidden_states = self._gradient_checkpointing_func( - block, - hidden_states, - encoder_hidden_states, - temb, - joint_attention_kwargs, - ) - elif not is_skip: - encoder_hidden_states, hidden_states = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - temb=temb, - joint_attention_kwargs=joint_attention_kwargs, - ) - - # controlnet residual - if block_controlnet_hidden_states is not None and block.context_pre_only is False: - interval_control = len(self.transformer_blocks) / len(block_controlnet_hidden_states) - hidden_states = hidden_states + block_controlnet_hidden_states[int(index_block / interval_control)] - - hidden_states = self.norm_out(hidden_states, temb) - hidden_states = self.proj_out(hidden_states) - - # unpatchify - patch_size = self.config.patch_size - height = height // patch_size - width = width // patch_size - - hidden_states = hidden_states.reshape( - shape=(hidden_states.shape[0], height, width, patch_size, patch_size, self.out_channels) - ) - hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) - output = hidden_states.reshape( - 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,) - - return Transformer2DModelOutput(sample=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_temporal.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_temporal.py deleted file mode 100644 index ffaf31d04570825767db665f7b79f01722d366af..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_temporal.py +++ /dev/null @@ -1,375 +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 dataclasses import dataclass -from typing import Any, Dict, Optional - -import torch -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import BaseOutput -from ..attention import BasicTransformerBlock, TemporalBasicTransformerBlock -from ..embeddings import TimestepEmbedding, Timesteps -from ..modeling_utils import ModelMixin -from ..resnet import AlphaBlender - - -@dataclass -class TransformerTemporalModelOutput(BaseOutput): - """ - The output of [`TransformerTemporalModel`]. - - Args: - sample (`torch.Tensor` of shape `(batch_size x num_frames, num_channels, height, width)`): - The hidden states output conditioned on `encoder_hidden_states` input. - """ - - sample: torch.Tensor - - -class TransformerTemporalModel(ModelMixin, ConfigMixin): - """ - A Transformer model for video-like data. - - Parameters: - num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. - attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. - in_channels (`int`, *optional*): - The number of channels in the input and output (specify if the input is **continuous**). - num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use. - attention_bias (`bool`, *optional*): - Configure if the `TransformerBlock` attention should contain a bias parameter. - sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**). - This is fixed during training since it is used to learn a number of position embeddings. - activation_fn (`str`, *optional*, defaults to `"geglu"`): - Activation function to use in feed-forward. See `diffusers.models.activations.get_activation` for supported - activation functions. - norm_elementwise_affine (`bool`, *optional*): - Configure if the `TransformerBlock` should use learnable elementwise affine parameters for normalization. - double_self_attention (`bool`, *optional*): - Configure if each `TransformerBlock` should contain two self-attention layers. - positional_embeddings: (`str`, *optional*): - The type of positional embeddings to apply to the sequence input before passing use. - num_positional_embeddings: (`int`, *optional*): - The maximum length of the sequence over which to apply positional embeddings. - """ - - _skip_layerwise_casting_patterns = ["norm"] - - @register_to_config - def __init__( - self, - num_attention_heads: int = 16, - attention_head_dim: int = 88, - in_channels: Optional[int] = None, - out_channels: Optional[int] = None, - num_layers: int = 1, - dropout: float = 0.0, - norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = None, - attention_bias: bool = False, - sample_size: Optional[int] = 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, - ): - super().__init__() - self.num_attention_heads = num_attention_heads - self.attention_head_dim = attention_head_dim - inner_dim = num_attention_heads * attention_head_dim - - self.in_channels = in_channels - - self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True) - self.proj_in = nn.Linear(in_channels, inner_dim) - - # 3. Define transformers blocks - self.transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - inner_dim, - num_attention_heads, - attention_head_dim, - dropout=dropout, - cross_attention_dim=cross_attention_dim, - activation_fn=activation_fn, - attention_bias=attention_bias, - double_self_attention=double_self_attention, - norm_elementwise_affine=norm_elementwise_affine, - positional_embeddings=positional_embeddings, - num_positional_embeddings=num_positional_embeddings, - ) - for d in range(num_layers) - ] - ) - - self.proj_out = nn.Linear(inner_dim, in_channels) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.LongTensor] = None, - timestep: Optional[torch.LongTensor] = None, - class_labels: torch.LongTensor = None, - num_frames: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - return_dict: bool = True, - ) -> TransformerTemporalModelOutput: - """ - The [`TransformerTemporal`] forward method. - - Args: - hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.Tensor` of shape `(batch size, channel, height, width)` if continuous): - Input hidden_states. - encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*): - Conditional embeddings for cross attention layer. If not given, cross-attention defaults to - self-attention. - timestep ( `torch.LongTensor`, *optional*): - Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`. - 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`. - num_frames (`int`, *optional*, defaults to 1): - The number of frames to be processed per batch. This is used to reshape the hidden states. - cross_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.transformers.transformer_temporal.TransformerTemporalModelOutput`] - instead of a plain tuple. - - Returns: - [`~models.transformers.transformer_temporal.TransformerTemporalModelOutput`] or `tuple`: - If `return_dict` is True, an - [`~models.transformers.transformer_temporal.TransformerTemporalModelOutput`] is returned, otherwise a - `tuple` where the first element is the sample tensor. - """ - # 1. Input - batch_frames, channel, height, width = hidden_states.shape - batch_size = batch_frames // num_frames - - residual = hidden_states - - hidden_states = hidden_states[None, :].reshape(batch_size, num_frames, channel, height, width) - hidden_states = hidden_states.permute(0, 2, 1, 3, 4) - - hidden_states = self.norm(hidden_states) - hidden_states = hidden_states.permute(0, 3, 4, 2, 1).reshape(batch_size * height * width, num_frames, channel) - - hidden_states = self.proj_in(hidden_states) - - # 2. Blocks - for block in self.transformer_blocks: - hidden_states = block( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=timestep, - cross_attention_kwargs=cross_attention_kwargs, - class_labels=class_labels, - ) - - # 3. Output - hidden_states = self.proj_out(hidden_states) - hidden_states = ( - hidden_states[None, None, :] - .reshape(batch_size, height, width, num_frames, channel) - .permute(0, 3, 4, 1, 2) - .contiguous() - ) - hidden_states = hidden_states.reshape(batch_frames, channel, height, width) - - output = hidden_states + residual - - if not return_dict: - return (output,) - - return TransformerTemporalModelOutput(sample=output) - - -class TransformerSpatioTemporalModel(nn.Module): - """ - A Transformer model for video-like data. - - Parameters: - num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. - attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. - in_channels (`int`, *optional*): - The number of channels in the input and output (specify if the input is **continuous**). - out_channels (`int`, *optional*): - The number of channels in the output (specify if the input is **continuous**). - num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. - cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use. - """ - - def __init__( - self, - num_attention_heads: int = 16, - attention_head_dim: int = 88, - in_channels: int = 320, - out_channels: Optional[int] = None, - num_layers: int = 1, - cross_attention_dim: Optional[int] = None, - ): - super().__init__() - self.num_attention_heads = num_attention_heads - self.attention_head_dim = attention_head_dim - - inner_dim = num_attention_heads * attention_head_dim - self.inner_dim = inner_dim - - # 2. Define input layers - self.in_channels = in_channels - self.norm = torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6) - self.proj_in = nn.Linear(in_channels, inner_dim) - - # 3. Define transformers blocks - self.transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - inner_dim, - num_attention_heads, - attention_head_dim, - cross_attention_dim=cross_attention_dim, - ) - for d in range(num_layers) - ] - ) - - time_mix_inner_dim = inner_dim - self.temporal_transformer_blocks = nn.ModuleList( - [ - TemporalBasicTransformerBlock( - inner_dim, - time_mix_inner_dim, - num_attention_heads, - attention_head_dim, - cross_attention_dim=cross_attention_dim, - ) - for _ in range(num_layers) - ] - ) - - time_embed_dim = in_channels * 4 - self.time_pos_embed = TimestepEmbedding(in_channels, time_embed_dim, out_dim=in_channels) - self.time_proj = Timesteps(in_channels, True, 0) - self.time_mixer = AlphaBlender(alpha=0.5, merge_strategy="learned_with_images") - - # 4. Define output layers - self.out_channels = in_channels if out_channels is None else out_channels - # TODO: should use out_channels for continuous projections - self.proj_out = nn.Linear(inner_dim, in_channels) - - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - image_only_indicator: Optional[torch.Tensor] = None, - return_dict: bool = True, - ): - """ - Args: - hidden_states (`torch.Tensor` of shape `(batch size, channel, height, width)`): - Input hidden_states. - num_frames (`int`): - The number of frames to be processed per batch. This is used to reshape the hidden states. - encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*): - Conditional embeddings for cross attention layer. If not given, cross-attention defaults to - self-attention. - image_only_indicator (`torch.LongTensor` of shape `(batch size, num_frames)`, *optional*): - A tensor indicating whether the input contains only images. 1 indicates that the input contains only - images, 0 indicates that the input contains video frames. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.transformers.transformer_temporal.TransformerTemporalModelOutput`] - instead of a plain tuple. - - Returns: - [`~models.transformers.transformer_temporal.TransformerTemporalModelOutput`] or `tuple`: - If `return_dict` is True, an - [`~models.transformers.transformer_temporal.TransformerTemporalModelOutput`] is returned, otherwise a - `tuple` where the first element is the sample tensor. - """ - # 1. Input - batch_frames, _, height, width = hidden_states.shape - num_frames = image_only_indicator.shape[-1] - batch_size = batch_frames // num_frames - - time_context = encoder_hidden_states - time_context_first_timestep = time_context[None, :].reshape( - batch_size, num_frames, -1, time_context.shape[-1] - )[:, 0] - time_context = time_context_first_timestep[:, None].broadcast_to( - batch_size, height * width, time_context.shape[-2], time_context.shape[-1] - ) - time_context = time_context.reshape(batch_size * height * width, -1, time_context.shape[-1]) - - residual = hidden_states - - hidden_states = self.norm(hidden_states) - inner_dim = hidden_states.shape[1] - hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_frames, height * width, inner_dim) - hidden_states = self.proj_in(hidden_states) - - num_frames_emb = torch.arange(num_frames, device=hidden_states.device) - num_frames_emb = num_frames_emb.repeat(batch_size, 1) - num_frames_emb = num_frames_emb.reshape(-1) - t_emb = self.time_proj(num_frames_emb) - - # `Timesteps` does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - t_emb = t_emb.to(dtype=hidden_states.dtype) - - emb = self.time_pos_embed(t_emb) - emb = emb[:, None, :] - - # 2. Blocks - for block, temporal_block in zip(self.transformer_blocks, self.temporal_transformer_blocks): - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - block, hidden_states, None, encoder_hidden_states, None - ) - else: - hidden_states = block(hidden_states, encoder_hidden_states=encoder_hidden_states) - - hidden_states_mix = hidden_states - hidden_states_mix = hidden_states_mix + emb - - hidden_states_mix = temporal_block( - hidden_states_mix, - num_frames=num_frames, - encoder_hidden_states=time_context, - ) - hidden_states = self.time_mixer( - x_spatial=hidden_states, - x_temporal=hidden_states_mix, - image_only_indicator=image_only_indicator, - ) - - # 3. Output - hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.reshape(batch_frames, height, width, inner_dim).permute(0, 3, 1, 2).contiguous() - - output = hidden_states + residual - - if not return_dict: - return (output,) - - return TransformerTemporalModelOutput(sample=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_wan.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_wan.py deleted file mode 100644 index baa0ede4184ecb226ea43c3cec26914492a6ed6c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_wan.py +++ /dev/null @@ -1,485 +0,0 @@ -# Copyright 2025 The Wan 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 typing import Any, Dict, Optional, Tuple, Union - -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 USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers -from ..attention import FeedForward -from ..attention_processor import Attention -from ..cache_utils import CacheMixin -from ..embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps, get_1d_rotary_pos_embed -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import FP32LayerNorm - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class WanAttnProcessor2_0: - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("WanAttnProcessor2_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, - rotary_emb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - encoder_hidden_states_img = None - if attn.add_k_proj is not None: - # 512 is the context length of the text encoder, hardcoded for now - image_context_length = encoder_hidden_states.shape[1] - 512 - encoder_hidden_states_img = encoder_hidden_states[:, :image_context_length] - encoder_hidden_states = encoder_hidden_states[:, image_context_length:] - 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) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - 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) - - if rotary_emb is not None: - - def apply_rotary_emb(hidden_states: torch.Tensor, freqs: torch.Tensor): - dtype = torch.float32 if hidden_states.device.type == "mps" else torch.float64 - x_rotated = torch.view_as_complex(hidden_states.to(dtype).unflatten(3, (-1, 2))) - x_out = torch.view_as_real(x_rotated * freqs).flatten(3, 4) - return x_out.type_as(hidden_states) - - query = apply_rotary_emb(query, rotary_emb) - key = apply_rotary_emb(key, rotary_emb) - - # I2V task - hidden_states_img = None - if encoder_hidden_states_img is not None: - key_img = attn.add_k_proj(encoder_hidden_states_img) - key_img = attn.norm_added_k(key_img) - value_img = attn.add_v_proj(encoder_hidden_states_img) - - key_img = key_img.unflatten(2, (attn.heads, -1)).transpose(1, 2) - value_img = value_img.unflatten(2, (attn.heads, -1)).transpose(1, 2) - - hidden_states_img = F.scaled_dot_product_attention( - query, key_img, value_img, attn_mask=None, dropout_p=0.0, is_causal=False - ) - hidden_states_img = hidden_states_img.transpose(1, 2).flatten(2, 3) - hidden_states_img = hidden_states_img.type_as(query) - - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - hidden_states = hidden_states.transpose(1, 2).flatten(2, 3) - hidden_states = hidden_states.type_as(query) - - if hidden_states_img is not None: - hidden_states = hidden_states + hidden_states_img - - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - return hidden_states - - -class WanImageEmbedding(torch.nn.Module): - def __init__(self, in_features: int, out_features: int, pos_embed_seq_len=None): - super().__init__() - - self.norm1 = FP32LayerNorm(in_features) - self.ff = FeedForward(in_features, out_features, mult=1, activation_fn="gelu") - self.norm2 = FP32LayerNorm(out_features) - if pos_embed_seq_len is not None: - self.pos_embed = nn.Parameter(torch.zeros(1, pos_embed_seq_len, in_features)) - else: - self.pos_embed = None - - def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: - if self.pos_embed is not None: - batch_size, seq_len, embed_dim = encoder_hidden_states_image.shape - encoder_hidden_states_image = encoder_hidden_states_image.view(-1, 2 * seq_len, embed_dim) - encoder_hidden_states_image = encoder_hidden_states_image + self.pos_embed - - hidden_states = self.norm1(encoder_hidden_states_image) - hidden_states = self.ff(hidden_states) - hidden_states = self.norm2(hidden_states) - return hidden_states - - -class WanTimeTextImageEmbedding(nn.Module): - def __init__( - self, - dim: int, - 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, - ): - 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") - - self.image_embedder = None - if image_embed_dim is not None: - self.image_embedder = WanImageEmbedding(image_embed_dim, dim, pos_embed_seq_len=pos_embed_seq_len) - - def forward( - self, - timestep: torch.Tensor, - encoder_hidden_states: torch.Tensor, - encoder_hidden_states_image: Optional[torch.Tensor] = None, - ): - 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)) - - encoder_hidden_states = self.text_embedder(encoder_hidden_states) - if encoder_hidden_states_image is not None: - encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image) - - return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image - - -class WanRotaryPosEmbed(nn.Module): - def __init__( - self, attention_head_dim: int, patch_size: Tuple[int, int, int], max_seq_len: int, theta: float = 10000.0 - ): - super().__init__() - - self.attention_head_dim = attention_head_dim - self.patch_size = patch_size - self.max_seq_len = max_seq_len - - h_dim = w_dim = 2 * (attention_head_dim // 6) - t_dim = attention_head_dim - h_dim - w_dim - - freqs = [] - freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64 - for dim in [t_dim, h_dim, w_dim]: - freq = get_1d_rotary_pos_embed( - dim, max_seq_len, theta, use_real=False, repeat_interleave_real=False, freqs_dtype=freqs_dtype - ) - freqs.append(freq) - self.freqs = torch.cat(freqs, dim=1) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - batch_size, num_channels, num_frames, height, width = hidden_states.shape - p_t, p_h, p_w = self.patch_size - ppf, pph, ppw = num_frames // p_t, height // p_h, width // p_w - - freqs = self.freqs.to(hidden_states.device) - freqs = freqs.split_with_sizes( - [ - self.attention_head_dim // 2 - 2 * (self.attention_head_dim // 6), - self.attention_head_dim // 6, - self.attention_head_dim // 6, - ], - dim=1, - ) - - freqs_f = freqs[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) - freqs_h = freqs[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) - freqs_w = freqs[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) - freqs = torch.cat([freqs_f, freqs_h, freqs_w], dim=-1).reshape(1, 1, ppf * pph * ppw, -1) - return freqs - - -class WanTransformerBlock(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: Optional[int] = None, - ): - super().__init__() - - # 1. Self-attention - self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) - self.attn1 = Attention( - query_dim=dim, - heads=num_heads, - kv_heads=num_heads, - dim_head=dim // num_heads, - qk_norm=qk_norm, - eps=eps, - bias=True, - cross_attention_dim=None, - out_bias=True, - processor=WanAttnProcessor2_0(), - ) - - # 2. Cross-attention - self.attn2 = Attention( - query_dim=dim, - heads=num_heads, - kv_heads=num_heads, - dim_head=dim // num_heads, - qk_norm=qk_norm, - eps=eps, - bias=True, - cross_attention_dim=None, - out_bias=True, - added_kv_proj_dim=added_kv_proj_dim, - added_proj_bias=True, - processor=WanAttnProcessor2_0(), - ) - 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) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - temb: torch.Tensor, - rotary_emb: torch.Tensor, - ) -> torch.Tensor: - 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(hidden_states=norm_hidden_states, rotary_emb=rotary_emb) - hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states) - - # 2. Cross-attention - norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states) - attn_output = self.attn2(hidden_states=norm_hidden_states, encoder_hidden_states=encoder_hidden_states) - 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 WanTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin): - r""" - A Transformer model for video-like data used in the Wan 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", "condition_embedder", "norm"] - _no_split_modules = ["WanTransformerBlock"] - _keep_in_fp32_modules = ["time_embedder", "scale_shift_table", "norm1", "norm2", "norm3"] - _keys_to_ignore_on_load_unexpected = ["norm_added_q"] - - @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: Optional[str] = "rms_norm_across_heads", - eps: float = 1e-6, - image_dim: Optional[int] = None, - added_kv_proj_dim: Optional[int] = None, - rope_max_seq_len: int = 1024, - pos_embed_seq_len: Optional[int] = None, - ) -> None: - super().__init__() - - inner_dim = num_attention_heads * attention_head_dim - out_channels = out_channels or in_channels - - # 1. Patch & position embedding - self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) - self.patch_embedding = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size) - - # 2. Condition embeddings - # image_embedding_dim=1280 for I2V model - self.condition_embedder = WanTimeTextImageEmbedding( - dim=inner_dim, - time_freq_dim=freq_dim, - time_proj_dim=inner_dim * 6, - text_embed_dim=text_dim, - image_embed_dim=image_dim, - pos_embed_seq_len=pos_embed_seq_len, - ) - - # 3. Transformer blocks - self.blocks = nn.ModuleList( - [ - WanTransformerBlock( - inner_dim, ffn_dim, num_attention_heads, qk_norm, cross_attn_norm, eps, added_kv_proj_dim - ) - for _ in range(num_layers) - ] - ) - - # 4. Output norm & projection - self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False) - self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) - self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) - - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - timestep: torch.LongTensor, - encoder_hidden_states: torch.Tensor, - encoder_hidden_states_image: Optional[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." - ) - - 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 - post_patch_height = height // p_h - post_patch_width = width // p_w - - rotary_emb = self.rope(hidden_states) - - hidden_states = self.patch_embedding(hidden_states) - hidden_states = hidden_states.flatten(2).transpose(1, 2) - - temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( - timestep, encoder_hidden_states, encoder_hidden_states_image - ) - timestep_proj = timestep_proj.unflatten(1, (6, -1)) - - if encoder_hidden_states_image is not None: - encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) - - # 4. Transformer blocks - 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 - ) - else: - for block in self.blocks: - hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb) - - # 5. Output norm, projection & unpatchify - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) - - # Move the shift and scale tensors to the same device as hidden_states. - # When using multi-GPU inference via accelerate these will be on the - # first device rather than the last device, which hidden_states ends up - # on. - shift = shift.to(hidden_states.device) - scale = scale.to(hidden_states.device) - - hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states) - hidden_states = self.proj_out(hidden_states) - - 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 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/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_wan_vace.py b/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_wan_vace.py deleted file mode 100644 index 1a6f2af59a87c79a6be3087c28efed1616dd777b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_wan_vace.py +++ /dev/null @@ -1,393 +0,0 @@ -# Copyright 2025 The Wan 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 typing import Any, Dict, List, Optional, Tuple, Union - -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 ..attention import FeedForward -from ..attention_processor import Attention -from ..cache_utils import CacheMixin -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin -from ..normalization import FP32LayerNorm -from .transformer_wan import WanAttnProcessor2_0, WanRotaryPosEmbed, WanTimeTextImageEmbedding, WanTransformerBlock - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class WanVACETransformerBlock(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: Optional[int] = None, - apply_input_projection: bool = False, - apply_output_projection: bool = False, - ): - super().__init__() - - # 1. Input projection - self.proj_in = None - if apply_input_projection: - self.proj_in = nn.Linear(dim, dim) - - # 2. Self-attention - self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) - self.attn1 = Attention( - query_dim=dim, - heads=num_heads, - kv_heads=num_heads, - dim_head=dim // num_heads, - qk_norm=qk_norm, - eps=eps, - bias=True, - cross_attention_dim=None, - out_bias=True, - processor=WanAttnProcessor2_0(), - ) - - # 3. Cross-attention - self.attn2 = Attention( - query_dim=dim, - heads=num_heads, - kv_heads=num_heads, - dim_head=dim // num_heads, - qk_norm=qk_norm, - eps=eps, - bias=True, - cross_attention_dim=None, - out_bias=True, - added_kv_proj_dim=added_kv_proj_dim, - added_proj_bias=True, - processor=WanAttnProcessor2_0(), - ) - self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() - - # 4. Feed-forward - self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate") - self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) - - # 5. Output projection - self.proj_out = None - if apply_output_projection: - self.proj_out = nn.Linear(dim, dim) - - self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - control_hidden_states: torch.Tensor, - temb: torch.Tensor, - rotary_emb: torch.Tensor, - ) -> torch.Tensor: - if self.proj_in is not None: - control_hidden_states = self.proj_in(control_hidden_states) - control_hidden_states = control_hidden_states + hidden_states - - 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(control_hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as( - control_hidden_states - ) - attn_output = self.attn1(hidden_states=norm_hidden_states, rotary_emb=rotary_emb) - control_hidden_states = (control_hidden_states.float() + attn_output * gate_msa).type_as(control_hidden_states) - - # 2. Cross-attention - norm_hidden_states = self.norm2(control_hidden_states.float()).type_as(control_hidden_states) - attn_output = self.attn2(hidden_states=norm_hidden_states, encoder_hidden_states=encoder_hidden_states) - control_hidden_states = control_hidden_states + attn_output - - # 3. Feed-forward - norm_hidden_states = (self.norm3(control_hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as( - control_hidden_states - ) - ff_output = self.ffn(norm_hidden_states) - control_hidden_states = (control_hidden_states.float() + ff_output.float() * c_gate_msa).type_as( - control_hidden_states - ) - - conditioning_states = None - if self.proj_out is not None: - conditioning_states = self.proj_out(control_hidden_states) - - return conditioning_states, control_hidden_states - - -class WanVACETransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin): - r""" - A Transformer model for video-like data used in the Wan 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", "vace_patch_embedding", "condition_embedder", "norm"] - _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"] - - @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: Optional[str] = "rms_norm_across_heads", - eps: float = 1e-6, - image_dim: Optional[int] = None, - added_kv_proj_dim: Optional[int] = 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], - vace_in_channels: int = 96, - ) -> None: - super().__init__() - - inner_dim = num_attention_heads * attention_head_dim - out_channels = out_channels or in_channels - - if max(vace_layers) >= num_layers: - raise ValueError(f"VACE layers {vace_layers} exceed the number of transformer layers {num_layers}.") - if 0 not in vace_layers: - raise ValueError("VACE layers must include layer 0.") - - # 1. Patch & position embedding - self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) - self.patch_embedding = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size) - self.vace_patch_embedding = nn.Conv3d(vace_in_channels, inner_dim, kernel_size=patch_size, stride=patch_size) - - # 2. Condition embeddings - # image_embedding_dim=1280 for I2V model - self.condition_embedder = WanTimeTextImageEmbedding( - dim=inner_dim, - time_freq_dim=freq_dim, - time_proj_dim=inner_dim * 6, - text_embed_dim=text_dim, - image_embed_dim=image_dim, - pos_embed_seq_len=pos_embed_seq_len, - ) - - # 3. Transformer blocks - self.blocks = nn.ModuleList( - [ - WanTransformerBlock( - inner_dim, ffn_dim, num_attention_heads, qk_norm, cross_attn_norm, eps, added_kv_proj_dim - ) - for _ in range(num_layers) - ] - ) - - self.vace_blocks = nn.ModuleList( - [ - WanVACETransformerBlock( - inner_dim, - ffn_dim, - num_attention_heads, - qk_norm, - cross_attn_norm, - eps, - added_kv_proj_dim, - apply_input_projection=i == 0, # Layer 0 always has input projection and is in vace_layers - apply_output_projection=True, - ) - for i in range(len(vace_layers)) - ] - ) - - # 4. Output norm & projection - self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False) - self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) - self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) - - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - timestep: torch.LongTensor, - encoder_hidden_states: torch.Tensor, - encoder_hidden_states_image: Optional[torch.Tensor] = 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." - ) - - 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 - post_patch_height = height // p_h - post_patch_width = width // p_w - - if control_hidden_states_scale is None: - control_hidden_states_scale = control_hidden_states.new_ones(len(self.config.vace_layers)) - control_hidden_states_scale = torch.unbind(control_hidden_states_scale) - if len(control_hidden_states_scale) != len(self.config.vace_layers): - raise ValueError( - f"Length of `control_hidden_states_scale` {len(control_hidden_states_scale)} should be " - f"equal to {len(self.config.vace_layers)}." - ) - - # 1. Rotary position embedding - rotary_emb = self.rope(hidden_states) - - # 2. Patch embedding - hidden_states = self.patch_embedding(hidden_states) - hidden_states = hidden_states.flatten(2).transpose(1, 2) - - control_hidden_states = self.vace_patch_embedding(control_hidden_states) - control_hidden_states = control_hidden_states.flatten(2).transpose(1, 2) - control_hidden_states_padding = control_hidden_states.new_zeros( - batch_size, hidden_states.size(1) - control_hidden_states.size(1), control_hidden_states.size(2) - ) - control_hidden_states = torch.cat([control_hidden_states, control_hidden_states_padding], dim=1) - - # 3. Time embedding - temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( - timestep, encoder_hidden_states, encoder_hidden_states_image - ) - timestep_proj = timestep_proj.unflatten(1, (6, -1)) - - # 4. Image embedding - if encoder_hidden_states_image is not None: - encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) - - # 5. Transformer blocks - if torch.is_grad_enabled() and self.gradient_checkpointing: - # Prepare VACE hints - control_hidden_states_list = [] - for i, block in enumerate(self.vace_blocks): - conditioning_states, control_hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, control_hidden_states, timestep_proj, rotary_emb - ) - control_hidden_states_list.append((conditioning_states, control_hidden_states_scale[i])) - control_hidden_states_list = control_hidden_states_list[::-1] - - for i, block in enumerate(self.blocks): - hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, timestep_proj, rotary_emb - ) - if i in self.config.vace_layers: - control_hint, scale = control_hidden_states_list.pop() - hidden_states = hidden_states + control_hint * scale - else: - # Prepare VACE hints - control_hidden_states_list = [] - for i, block in enumerate(self.vace_blocks): - conditioning_states, control_hidden_states = block( - hidden_states, encoder_hidden_states, control_hidden_states, timestep_proj, rotary_emb - ) - control_hidden_states_list.append((conditioning_states, control_hidden_states_scale[i])) - control_hidden_states_list = control_hidden_states_list[::-1] - - for i, block in enumerate(self.blocks): - hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb) - if i in self.config.vace_layers: - control_hint, scale = control_hidden_states_list.pop() - hidden_states = hidden_states + control_hint * scale - - # 6. Output norm, projection & unpatchify - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) - - # Move the shift and scale tensors to the same device as hidden_states. - # When using multi-GPU inference via accelerate these will be on the - # first device rather than the last device, which hidden_states ends up - # on. - shift = shift.to(hidden_states.device) - scale = scale.to(hidden_states.device) - - hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states) - hidden_states = self.proj_out(hidden_states) - - 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 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/venv/lib/python3.10/site-packages/diffusers/models/unets/__init__.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/__init__.py deleted file mode 100644 index 9ef04fb62606ed30699be8ac6cb38b011a692070..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -from ...utils import is_flax_available, is_torch_available - - -if is_torch_available(): - from .unet_1d import UNet1DModel - from .unet_2d import UNet2DModel - from .unet_2d_condition import UNet2DConditionModel - from .unet_3d_condition import UNet3DConditionModel - from .unet_i2vgen_xl import I2VGenXLUNet - from .unet_kandinsky3 import Kandinsky3UNet - from .unet_motion_model import MotionAdapter, UNetMotionModel - from .unet_spatio_temporal_condition import UNetSpatioTemporalConditionModel - from .unet_stable_cascade import StableCascadeUNet - from .uvit_2d import UVit2DModel - - -if is_flax_available(): - from .unet_2d_condition_flax import FlaxUNet2DConditionModel diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 548e7c659e427f7b0c4cc2a233cebbc83caa0aed..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_1d.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_1d.cpython-310.pyc deleted file mode 100644 index 09c4a8361bc80fb225438ea4271902823227aa74..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_1d.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_1d_blocks.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_1d_blocks.cpython-310.pyc deleted file mode 100644 index 77029d6d16d8d2f2d234ef9019bedfd32d198045..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_1d_blocks.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d.cpython-310.pyc deleted file mode 100644 index 36177ba3e72d200b892cfaa5c102d06df1a48ede..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d_blocks.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d_blocks.cpython-310.pyc deleted file mode 100644 index 346beb7ff660e6f731851556458869110d01c77d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d_blocks.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d_blocks_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d_blocks_flax.cpython-310.pyc deleted file mode 100644 index 5a30f948ebe4ceb5b44e31c11f8553f9a9191506..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d_blocks_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d_condition.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d_condition.cpython-310.pyc deleted file mode 100644 index 634571787af7d9c96c1ed17ff3d88bdf878d953f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d_condition.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d_condition_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d_condition_flax.cpython-310.pyc deleted file mode 100644 index bcabb2ab7e68680ce7cf522875b4bcf147fb2f44..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_2d_condition_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_3d_blocks.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_3d_blocks.cpython-310.pyc deleted file mode 100644 index ce7363487510a9d99f652ab41e05fab25fc052db..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_3d_blocks.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_3d_condition.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_3d_condition.cpython-310.pyc deleted file mode 100644 index 9dbd61bc708c445d44fa76a7b5b9bfa2143f44c8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_3d_condition.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_i2vgen_xl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_i2vgen_xl.cpython-310.pyc deleted file mode 100644 index 246303ae1c1df28304441774522ca9fcf450c107..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_i2vgen_xl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_kandinsky3.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_kandinsky3.cpython-310.pyc deleted file mode 100644 index e14b00d036d1463642254b1f913011c04bc5483d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_kandinsky3.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_motion_model.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_motion_model.cpython-310.pyc deleted file mode 100644 index eed9cae43f237a0dd21846134291e9c90f59f55e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_motion_model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_spatio_temporal_condition.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_spatio_temporal_condition.cpython-310.pyc deleted file mode 100644 index c6db6ff9677adf1ecefb2f502f6d6e8182307698..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_spatio_temporal_condition.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_stable_cascade.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_stable_cascade.cpython-310.pyc deleted file mode 100644 index 9504f6a7eb7224ee67182ddabe29e9c39ae41f1e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/unet_stable_cascade.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/uvit_2d.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/uvit_2d.cpython-310.pyc deleted file mode 100644 index de23075cc43144e122082fbcfc86967e2282daa0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/models/unets/__pycache__/uvit_2d.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_1d.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_1d.py deleted file mode 100644 index 4f57f3349bc47ff5ddd0f91baa3ac5c79a65a523..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_1d.py +++ /dev/null @@ -1,257 +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 dataclasses import dataclass -from typing import Optional, Tuple, Union - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import BaseOutput -from ..embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps -from ..modeling_utils import ModelMixin -from .unet_1d_blocks import get_down_block, get_mid_block, get_out_block, get_up_block - - -@dataclass -class UNet1DOutput(BaseOutput): - """ - The output of [`UNet1DModel`]. - - Args: - sample (`torch.Tensor` of shape `(batch_size, num_channels, sample_size)`): - The hidden states output from the last layer of the model. - """ - - sample: torch.Tensor - - -class UNet1DModel(ModelMixin, ConfigMixin): - r""" - A 1D UNet model that takes a noisy sample and a timestep and returns a sample shaped output. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - Parameters: - sample_size (`int`, *optional*): Default length of sample. Should be adaptable at runtime. - in_channels (`int`, *optional*, defaults to 2): Number of channels in the input sample. - out_channels (`int`, *optional*, defaults to 2): Number of channels in the output. - extra_in_channels (`int`, *optional*, defaults to 0): - Number of additional channels to be added to the input of the first down block. Useful for cases where the - input data has more channels than what the model was initially designed for. - time_embedding_type (`str`, *optional*, defaults to `"fourier"`): Type of time embedding to use. - 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. - 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. - norm_num_groups (`int`, *optional*, defaults to 8): The number of groups for normalization. - layers_per_block (`int`, *optional*, defaults to 1): The number of layers per block. - downsample_each_block (`int`, *optional*, defaults to `False`): - Experimental feature for using a UNet without upsampling. - """ - - _skip_layerwise_casting_patterns = ["norm"] - - @register_to_config - def __init__( - self, - sample_size: int = 65536, - sample_rate: Optional[int] = None, - in_channels: int = 2, - out_channels: int = 2, - extra_in_channels: int = 0, - time_embedding_type: str = "fourier", - 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"), - mid_block_type: Tuple[str] = "UNetMidBlock1D", - out_block_type: str = None, - block_out_channels: Tuple[int] = (32, 32, 64), - act_fn: str = None, - norm_num_groups: int = 8, - layers_per_block: int = 1, - downsample_each_block: bool = False, - ): - super().__init__() - self.sample_size = sample_size - - # time - if time_embedding_type == "fourier": - self.time_proj = GaussianFourierProjection( - embedding_size=8, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos - ) - timestep_input_dim = 2 * block_out_channels[0] - elif time_embedding_type == "positional": - self.time_proj = Timesteps( - block_out_channels[0], flip_sin_to_cos=flip_sin_to_cos, downscale_freq_shift=freq_shift - ) - timestep_input_dim = block_out_channels[0] - - if use_timestep_embedding: - time_embed_dim = block_out_channels[0] * 4 - self.time_mlp = TimestepEmbedding( - in_channels=timestep_input_dim, - time_embed_dim=time_embed_dim, - act_fn=act_fn, - out_dim=block_out_channels[0], - ) - - self.down_blocks = nn.ModuleList([]) - self.mid_block = None - self.up_blocks = nn.ModuleList([]) - self.out_block = None - - # down - output_channel = in_channels - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - - if i == 0: - input_channel += extra_in_channels - - is_final_block = i == len(block_out_channels) - 1 - - down_block = get_down_block( - down_block_type, - num_layers=layers_per_block, - in_channels=input_channel, - out_channels=output_channel, - temb_channels=block_out_channels[0], - add_downsample=not is_final_block or downsample_each_block, - ) - self.down_blocks.append(down_block) - - # mid - self.mid_block = get_mid_block( - mid_block_type, - in_channels=block_out_channels[-1], - mid_channels=block_out_channels[-1], - out_channels=block_out_channels[-1], - embed_dim=block_out_channels[0], - num_layers=layers_per_block, - add_downsample=downsample_each_block, - ) - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - output_channel = reversed_block_out_channels[0] - if out_block_type is None: - final_upsample_channels = out_channels - else: - final_upsample_channels = block_out_channels[0] - - for i, up_block_type in enumerate(up_block_types): - prev_output_channel = output_channel - output_channel = ( - reversed_block_out_channels[i + 1] if i < len(up_block_types) - 1 else final_upsample_channels - ) - - is_final_block = i == len(block_out_channels) - 1 - - up_block = get_up_block( - up_block_type, - num_layers=layers_per_block, - in_channels=prev_output_channel, - out_channels=output_channel, - temb_channels=block_out_channels[0], - add_upsample=not is_final_block, - ) - self.up_blocks.append(up_block) - prev_output_channel = output_channel - - # out - num_groups_out = norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4, 32) - self.out_block = get_out_block( - out_block_type=out_block_type, - num_groups_out=num_groups_out, - embed_dim=block_out_channels[0], - out_channels=out_channels, - act_fn=act_fn, - fc_dim=block_out_channels[-1] // 4, - ) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], - return_dict: bool = True, - ) -> Union[UNet1DOutput, Tuple]: - r""" - The [`UNet1DModel`] forward method. - - Args: - sample (`torch.Tensor`): - The noisy input tensor with the following shape `(batch_size, num_channels, sample_size)`. - timestep (`torch.Tensor` or `float` or `int`): The number of timesteps to denoise an input. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.unets.unet_1d.UNet1DOutput`] instead of a plain tuple. - - Returns: - [`~models.unets.unet_1d.UNet1DOutput`] or `tuple`: - If `return_dict` is True, an [`~models.unets.unet_1d.UNet1DOutput`] is returned, otherwise a `tuple` is - returned where the first element is the sample tensor. - """ - - # 1. time - timesteps = timestep - if not torch.is_tensor(timesteps): - timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device) - elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - - timestep_embed = self.time_proj(timesteps) - if self.config.use_timestep_embedding: - timestep_embed = self.time_mlp(timestep_embed.to(sample.dtype)) - else: - timestep_embed = timestep_embed[..., None] - timestep_embed = timestep_embed.repeat([1, 1, sample.shape[2]]).to(sample.dtype) - timestep_embed = timestep_embed.broadcast_to((sample.shape[:1] + timestep_embed.shape[1:])) - - # 2. down - down_block_res_samples = () - for downsample_block in self.down_blocks: - sample, res_samples = downsample_block(hidden_states=sample, temb=timestep_embed) - down_block_res_samples += res_samples - - # 3. mid - if self.mid_block: - sample = self.mid_block(sample, timestep_embed) - - # 4. up - for i, upsample_block in enumerate(self.up_blocks): - res_samples = down_block_res_samples[-1:] - down_block_res_samples = down_block_res_samples[:-1] - sample = upsample_block(sample, res_hidden_states_tuple=res_samples, temb=timestep_embed) - - # 5. post-process - if self.out_block: - sample = self.out_block(sample, timestep_embed) - - if not return_dict: - return (sample,) - - return UNet1DOutput(sample=sample) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_1d_blocks.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_1d_blocks.py deleted file mode 100644 index 58cbdfd005b62216acf74a915ffd0b3ff1821d29..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_1d_blocks.py +++ /dev/null @@ -1,702 +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 math -from typing import Optional, Tuple, Union - -import torch -import torch.nn.functional as F -from torch import nn - -from ..activations import get_activation -from ..resnet import Downsample1D, ResidualTemporalBlock1D, Upsample1D, rearrange_dims - - -class DownResnetBlock1D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: Optional[int] = 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, - time_embedding_norm: str = "default", - output_scale_factor: float = 1.0, - add_downsample: bool = True, - ): - super().__init__() - 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 - self.time_embedding_norm = time_embedding_norm - self.add_downsample = add_downsample - self.output_scale_factor = output_scale_factor - - if groups_out is None: - groups_out = groups - - # there will always be at least one resnet - resnets = [ResidualTemporalBlock1D(in_channels, out_channels, embed_dim=temb_channels)] - - for _ in range(num_layers): - resnets.append(ResidualTemporalBlock1D(out_channels, out_channels, embed_dim=temb_channels)) - - self.resnets = nn.ModuleList(resnets) - - if non_linearity is None: - self.nonlinearity = None - else: - self.nonlinearity = get_activation(non_linearity) - - self.downsample = None - 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: - output_states = () - - hidden_states = self.resnets[0](hidden_states, temb) - for resnet in self.resnets[1:]: - hidden_states = resnet(hidden_states, temb) - - output_states += (hidden_states,) - - if self.nonlinearity is not None: - hidden_states = self.nonlinearity(hidden_states) - - if self.downsample is not None: - hidden_states = self.downsample(hidden_states) - - return hidden_states, output_states - - -class UpResnetBlock1D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: Optional[int] = None, - num_layers: int = 1, - temb_channels: int = 32, - groups: int = 32, - groups_out: Optional[int] = None, - non_linearity: Optional[str] = None, - time_embedding_norm: str = "default", - output_scale_factor: float = 1.0, - add_upsample: bool = True, - ): - super().__init__() - self.in_channels = in_channels - out_channels = in_channels if out_channels is None else out_channels - self.out_channels = out_channels - self.time_embedding_norm = time_embedding_norm - self.add_upsample = add_upsample - self.output_scale_factor = output_scale_factor - - if groups_out is None: - groups_out = groups - - # there will always be at least one resnet - resnets = [ResidualTemporalBlock1D(2 * in_channels, out_channels, embed_dim=temb_channels)] - - for _ in range(num_layers): - resnets.append(ResidualTemporalBlock1D(out_channels, out_channels, embed_dim=temb_channels)) - - self.resnets = nn.ModuleList(resnets) - - if non_linearity is None: - self.nonlinearity = None - else: - self.nonlinearity = get_activation(non_linearity) - - self.upsample = None - if add_upsample: - self.upsample = Upsample1D(out_channels, use_conv_transpose=True) - - def forward( - self, - hidden_states: torch.Tensor, - res_hidden_states_tuple: Optional[Tuple[torch.Tensor, ...]] = None, - temb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - if res_hidden_states_tuple is not None: - res_hidden_states = res_hidden_states_tuple[-1] - hidden_states = torch.cat((hidden_states, res_hidden_states), dim=1) - - hidden_states = self.resnets[0](hidden_states, temb) - for resnet in self.resnets[1:]: - hidden_states = resnet(hidden_states, temb) - - if self.nonlinearity is not None: - hidden_states = self.nonlinearity(hidden_states) - - if self.upsample is not None: - hidden_states = self.upsample(hidden_states) - - return hidden_states - - -class ValueFunctionMidBlock1D(nn.Module): - def __init__(self, in_channels: int, out_channels: int, embed_dim: int): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.embed_dim = embed_dim - - self.res1 = ResidualTemporalBlock1D(in_channels, in_channels // 2, embed_dim=embed_dim) - self.down1 = Downsample1D(out_channels // 2, use_conv=True) - 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: - x = self.res1(x, temb) - x = self.down1(x) - x = self.res2(x, temb) - x = self.down2(x) - return x - - -class MidResTemporalBlock1D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - embed_dim: int, - num_layers: int = 1, - add_downsample: bool = False, - add_upsample: bool = False, - non_linearity: Optional[str] = None, - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.add_downsample = add_downsample - - # there will always be at least one resnet - resnets = [ResidualTemporalBlock1D(in_channels, out_channels, embed_dim=embed_dim)] - - for _ in range(num_layers): - resnets.append(ResidualTemporalBlock1D(out_channels, out_channels, embed_dim=embed_dim)) - - self.resnets = nn.ModuleList(resnets) - - if non_linearity is None: - self.nonlinearity = None - else: - self.nonlinearity = get_activation(non_linearity) - - self.upsample = None - if add_upsample: - self.upsample = Upsample1D(out_channels, use_conv=True) - - self.downsample = None - if add_downsample: - self.downsample = Downsample1D(out_channels, use_conv=True) - - if self.upsample and self.downsample: - raise ValueError("Block cannot downsample and upsample") - - def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: - hidden_states = self.resnets[0](hidden_states, temb) - for resnet in self.resnets[1:]: - hidden_states = resnet(hidden_states, temb) - - if self.upsample: - hidden_states = self.upsample(hidden_states) - if self.downsample: - hidden_states = self.downsample(hidden_states) - - return hidden_states - - -class OutConv1DBlock(nn.Module): - def __init__(self, num_groups_out: int, out_channels: int, embed_dim: int, act_fn: str): - super().__init__() - self.final_conv1d_1 = nn.Conv1d(embed_dim, embed_dim, 5, padding=2) - self.final_conv1d_gn = nn.GroupNorm(num_groups_out, embed_dim) - 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: - hidden_states = self.final_conv1d_1(hidden_states) - hidden_states = rearrange_dims(hidden_states) - hidden_states = self.final_conv1d_gn(hidden_states) - hidden_states = rearrange_dims(hidden_states) - hidden_states = self.final_conv1d_act(hidden_states) - hidden_states = self.final_conv1d_2(hidden_states) - return hidden_states - - -class OutValueFunctionBlock(nn.Module): - def __init__(self, fc_dim: int, embed_dim: int, act_fn: str = "mish"): - super().__init__() - self.final_block = nn.ModuleList( - [ - nn.Linear(fc_dim + embed_dim, fc_dim // 2), - get_activation(act_fn), - nn.Linear(fc_dim // 2, 1), - ] - ) - - def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: - hidden_states = hidden_states.view(hidden_states.shape[0], -1) - hidden_states = torch.cat((hidden_states, temb), dim=-1) - for layer in self.final_block: - hidden_states = layer(hidden_states) - - return hidden_states - - -_kernels = { - "linear": [1 / 8, 3 / 8, 3 / 8, 1 / 8], - "cubic": [-0.01171875, -0.03515625, 0.11328125, 0.43359375, 0.43359375, 0.11328125, -0.03515625, -0.01171875], - "lanczos3": [ - 0.003689131001010537, - 0.015056144446134567, - -0.03399861603975296, - -0.066637322306633, - 0.13550527393817902, - 0.44638532400131226, - 0.44638532400131226, - 0.13550527393817902, - -0.066637322306633, - -0.03399861603975296, - 0.015056144446134567, - 0.003689131001010537, - ], -} - - -class Downsample1d(nn.Module): - def __init__(self, kernel: str = "linear", pad_mode: str = "reflect"): - super().__init__() - self.pad_mode = pad_mode - kernel_1d = torch.tensor(_kernels[kernel]) - self.pad = kernel_1d.shape[0] // 2 - 1 - self.register_buffer("kernel", kernel_1d) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = F.pad(hidden_states, (self.pad,) * 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) - kernel = self.kernel.to(weight)[None, :].expand(hidden_states.shape[1], -1) - weight[indices, indices] = kernel - return F.conv1d(hidden_states, weight, stride=2) - - -class Upsample1d(nn.Module): - def __init__(self, kernel: str = "linear", pad_mode: str = "reflect"): - super().__init__() - self.pad_mode = pad_mode - kernel_1d = torch.tensor(_kernels[kernel]) * 2 - 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: - 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) - kernel = self.kernel.to(weight)[None, :].expand(hidden_states.shape[1], -1) - weight[indices, indices] = kernel - return F.conv_transpose1d(hidden_states, weight, stride=2, padding=self.pad * 2 + 1) - - -class SelfAttention1d(nn.Module): - def __init__(self, in_channels: int, n_head: int = 1, dropout_rate: float = 0.0): - super().__init__() - self.channels = in_channels - self.group_norm = nn.GroupNorm(1, num_channels=in_channels) - self.num_heads = n_head - - self.query = nn.Linear(self.channels, self.channels) - self.key = nn.Linear(self.channels, self.channels) - self.value = nn.Linear(self.channels, self.channels) - - self.proj_attn = nn.Linear(self.channels, self.channels, bias=True) - - self.dropout = nn.Dropout(dropout_rate, inplace=True) - - def transpose_for_scores(self, projection: torch.Tensor) -> torch.Tensor: - new_projection_shape = projection.size()[:-1] + (self.num_heads, -1) - # move heads to 2nd position (B, T, H * D) -> (B, T, H, D) -> (B, H, T, D) - new_projection = projection.view(new_projection_shape).permute(0, 2, 1, 3) - return new_projection - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - residual = hidden_states - batch, channel_dim, seq = hidden_states.shape - - hidden_states = self.group_norm(hidden_states) - hidden_states = hidden_states.transpose(1, 2) - - query_proj = self.query(hidden_states) - key_proj = self.key(hidden_states) - value_proj = self.value(hidden_states) - - query_states = self.transpose_for_scores(query_proj) - key_states = self.transpose_for_scores(key_proj) - value_states = self.transpose_for_scores(value_proj) - - scale = 1 / math.sqrt(math.sqrt(key_states.shape[-1])) - - attention_scores = torch.matmul(query_states * scale, key_states.transpose(-1, -2) * scale) - attention_probs = torch.softmax(attention_scores, dim=-1) - - # compute attention output - hidden_states = torch.matmul(attention_probs, value_states) - - hidden_states = hidden_states.permute(0, 2, 1, 3).contiguous() - new_hidden_states_shape = hidden_states.size()[:-2] + (self.channels,) - hidden_states = hidden_states.view(new_hidden_states_shape) - - # compute next hidden_states - hidden_states = self.proj_attn(hidden_states) - hidden_states = hidden_states.transpose(1, 2) - hidden_states = self.dropout(hidden_states) - - output = hidden_states + residual - - return output - - -class ResConvBlock(nn.Module): - def __init__(self, in_channels: int, mid_channels: int, out_channels: int, is_last: bool = False): - super().__init__() - self.is_last = is_last - self.has_conv_skip = in_channels != out_channels - - if self.has_conv_skip: - self.conv_skip = nn.Conv1d(in_channels, out_channels, 1, bias=False) - - self.conv_1 = nn.Conv1d(in_channels, mid_channels, 5, padding=2) - self.group_norm_1 = nn.GroupNorm(1, mid_channels) - self.gelu_1 = nn.GELU() - self.conv_2 = nn.Conv1d(mid_channels, out_channels, 5, padding=2) - - if not self.is_last: - self.group_norm_2 = nn.GroupNorm(1, out_channels) - self.gelu_2 = nn.GELU() - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - residual = self.conv_skip(hidden_states) if self.has_conv_skip else hidden_states - - hidden_states = self.conv_1(hidden_states) - hidden_states = self.group_norm_1(hidden_states) - hidden_states = self.gelu_1(hidden_states) - hidden_states = self.conv_2(hidden_states) - - if not self.is_last: - hidden_states = self.group_norm_2(hidden_states) - hidden_states = self.gelu_2(hidden_states) - - output = hidden_states + residual - return output - - -class UNetMidBlock1D(nn.Module): - def __init__(self, mid_channels: int, in_channels: int, out_channels: Optional[int] = None): - super().__init__() - - out_channels = in_channels if out_channels is None else out_channels - - # there is always at least one resnet - self.down = Downsample1d("cubic") - resnets = [ - ResConvBlock(in_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, out_channels), - ] - attentions = [ - SelfAttention1d(mid_channels, mid_channels // 32), - SelfAttention1d(mid_channels, mid_channels // 32), - SelfAttention1d(mid_channels, mid_channels // 32), - SelfAttention1d(mid_channels, mid_channels // 32), - SelfAttention1d(mid_channels, mid_channels // 32), - SelfAttention1d(out_channels, out_channels // 32), - ] - self.up = Upsample1d(kernel="cubic") - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: - hidden_states = self.down(hidden_states) - for attn, resnet in zip(self.attentions, self.resnets): - hidden_states = resnet(hidden_states) - hidden_states = attn(hidden_states) - - hidden_states = self.up(hidden_states) - - return hidden_states - - -class AttnDownBlock1D(nn.Module): - def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[int] = None): - super().__init__() - mid_channels = out_channels if mid_channels is None else mid_channels - - self.down = Downsample1d("cubic") - resnets = [ - ResConvBlock(in_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, out_channels), - ] - attentions = [ - SelfAttention1d(mid_channels, mid_channels // 32), - SelfAttention1d(mid_channels, mid_channels // 32), - SelfAttention1d(out_channels, out_channels // 32), - ] - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: - hidden_states = self.down(hidden_states) - - for resnet, attn in zip(self.resnets, self.attentions): - hidden_states = resnet(hidden_states) - hidden_states = attn(hidden_states) - - return hidden_states, (hidden_states,) - - -class DownBlock1D(nn.Module): - def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[int] = None): - super().__init__() - mid_channels = out_channels if mid_channels is None else mid_channels - - self.down = Downsample1d("cubic") - resnets = [ - ResConvBlock(in_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, out_channels), - ] - - self.resnets = nn.ModuleList(resnets) - - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: - hidden_states = self.down(hidden_states) - - for resnet in self.resnets: - hidden_states = resnet(hidden_states) - - return hidden_states, (hidden_states,) - - -class DownBlock1DNoSkip(nn.Module): - def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[int] = None): - super().__init__() - mid_channels = out_channels if mid_channels is None else mid_channels - - resnets = [ - ResConvBlock(in_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, out_channels), - ] - - self.resnets = nn.ModuleList(resnets) - - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: - hidden_states = torch.cat([hidden_states, temb], dim=1) - for resnet in self.resnets: - hidden_states = resnet(hidden_states) - - return hidden_states, (hidden_states,) - - -class AttnUpBlock1D(nn.Module): - def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[int] = None): - super().__init__() - mid_channels = out_channels if mid_channels is None else mid_channels - - resnets = [ - ResConvBlock(2 * in_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, out_channels), - ] - attentions = [ - SelfAttention1d(mid_channels, mid_channels // 32), - SelfAttention1d(mid_channels, mid_channels // 32), - SelfAttention1d(out_channels, out_channels // 32), - ] - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - self.up = Upsample1d(kernel="cubic") - - def forward( - self, - hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - res_hidden_states = res_hidden_states_tuple[-1] - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - for resnet, attn in zip(self.resnets, self.attentions): - hidden_states = resnet(hidden_states) - hidden_states = attn(hidden_states) - - hidden_states = self.up(hidden_states) - - return hidden_states - - -class UpBlock1D(nn.Module): - def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[int] = None): - super().__init__() - mid_channels = in_channels if mid_channels is None else mid_channels - - resnets = [ - ResConvBlock(2 * in_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, out_channels), - ] - - self.resnets = nn.ModuleList(resnets) - self.up = Upsample1d(kernel="cubic") - - def forward( - self, - hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - res_hidden_states = res_hidden_states_tuple[-1] - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - for resnet in self.resnets: - hidden_states = resnet(hidden_states) - - hidden_states = self.up(hidden_states) - - return hidden_states - - -class UpBlock1DNoSkip(nn.Module): - def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[int] = None): - super().__init__() - mid_channels = in_channels if mid_channels is None else mid_channels - - resnets = [ - ResConvBlock(2 * in_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, mid_channels), - ResConvBlock(mid_channels, mid_channels, out_channels, is_last=True), - ] - - self.resnets = nn.ModuleList(resnets) - - def forward( - self, - hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - res_hidden_states = res_hidden_states_tuple[-1] - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - for resnet in self.resnets: - hidden_states = resnet(hidden_states) - - return hidden_states - - -DownBlockType = Union[DownResnetBlock1D, DownBlock1D, AttnDownBlock1D, DownBlock1DNoSkip] -MidBlockType = Union[MidResTemporalBlock1D, ValueFunctionMidBlock1D, UNetMidBlock1D] -OutBlockType = Union[OutConv1DBlock, OutValueFunctionBlock] -UpBlockType = Union[UpResnetBlock1D, UpBlock1D, AttnUpBlock1D, UpBlock1DNoSkip] - - -def get_down_block( - down_block_type: str, - num_layers: int, - in_channels: int, - out_channels: int, - temb_channels: int, - add_downsample: bool, -) -> DownBlockType: - if down_block_type == "DownResnetBlock1D": - return DownResnetBlock1D( - in_channels=in_channels, - num_layers=num_layers, - out_channels=out_channels, - temb_channels=temb_channels, - add_downsample=add_downsample, - ) - elif down_block_type == "DownBlock1D": - return DownBlock1D(out_channels=out_channels, in_channels=in_channels) - elif down_block_type == "AttnDownBlock1D": - return AttnDownBlock1D(out_channels=out_channels, in_channels=in_channels) - elif down_block_type == "DownBlock1DNoSkip": - return DownBlock1DNoSkip(out_channels=out_channels, in_channels=in_channels) - raise ValueError(f"{down_block_type} does not exist.") - - -def get_up_block( - up_block_type: str, num_layers: int, in_channels: int, out_channels: int, temb_channels: int, add_upsample: bool -) -> UpBlockType: - if up_block_type == "UpResnetBlock1D": - return UpResnetBlock1D( - in_channels=in_channels, - num_layers=num_layers, - out_channels=out_channels, - temb_channels=temb_channels, - add_upsample=add_upsample, - ) - elif up_block_type == "UpBlock1D": - return UpBlock1D(in_channels=in_channels, out_channels=out_channels) - elif up_block_type == "AttnUpBlock1D": - return AttnUpBlock1D(in_channels=in_channels, out_channels=out_channels) - elif up_block_type == "UpBlock1DNoSkip": - return UpBlock1DNoSkip(in_channels=in_channels, out_channels=out_channels) - raise ValueError(f"{up_block_type} does not exist.") - - -def get_mid_block( - mid_block_type: str, - num_layers: int, - in_channels: int, - mid_channels: int, - out_channels: int, - embed_dim: int, - add_downsample: bool, -) -> MidBlockType: - if mid_block_type == "MidResTemporalBlock1D": - return MidResTemporalBlock1D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - embed_dim=embed_dim, - add_downsample=add_downsample, - ) - elif mid_block_type == "ValueFunctionMidBlock1D": - return ValueFunctionMidBlock1D(in_channels=in_channels, out_channels=out_channels, embed_dim=embed_dim) - elif mid_block_type == "UNetMidBlock1D": - return UNetMidBlock1D(in_channels=in_channels, mid_channels=mid_channels, out_channels=out_channels) - raise ValueError(f"{mid_block_type} does not exist.") - - -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]: - if out_block_type == "OutConv1DBlock": - return OutConv1DBlock(num_groups_out, out_channels, embed_dim, act_fn) - elif out_block_type == "ValueFunction": - return OutValueFunctionBlock(fc_dim, embed_dim, act_fn) - return None diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d.py deleted file mode 100644 index 2588a9c518bdb9bd7073cdd3c38da05e10fa781e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d.py +++ /dev/null @@ -1,354 +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 dataclasses import dataclass -from typing import Optional, Tuple, Union - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import BaseOutput -from ..embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps -from ..modeling_utils import ModelMixin -from .unet_2d_blocks import UNetMidBlock2D, get_down_block, get_up_block - - -@dataclass -class UNet2DOutput(BaseOutput): - """ - The output of [`UNet2DModel`]. - - Args: - sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)`): - The hidden states output from the last layer of the model. - """ - - sample: torch.Tensor - - -class UNet2DModel(ModelMixin, ConfigMixin): - r""" - A 2D UNet model that takes a noisy sample and a timestep and returns a sample shaped output. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - Parameters: - 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. - out_channels (`int`, *optional*, defaults to 3): Number of channels in the output. - center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample. - time_embedding_type (`str`, *optional*, defaults to `"positional"`): Type of time embedding to use. - 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. - 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. - 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. - downsample_type (`str`, *optional*, defaults to `conv`): - The downsample type for downsampling layers. Choose between "conv" and "resnet" - upsample_type (`str`, *optional*, defaults to `conv`): - The upsample type for upsampling layers. Choose between "conv" and "resnet" - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. - attention_head_dim (`int`, *optional*, defaults to `8`): The attention head dimension. - norm_num_groups (`int`, *optional*, defaults to `32`): The number of groups for normalization. - attn_norm_num_groups (`int`, *optional*, defaults to `None`): - If set to an integer, a group norm layer will be created in the mid block's [`Attention`] layer with the - given number of groups. If left as `None`, the group norm layer will only be created if - `resnet_time_scale_shift` is set to `default`, and if created will have `norm_num_groups` groups. - norm_eps (`float`, *optional*, defaults to `1e-5`): The epsilon for normalization. - resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config - for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`. - class_embed_type (`str`, *optional*, defaults to `None`): - The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`, - `"timestep"`, or `"identity"`. - num_class_embeds (`int`, *optional*, defaults to `None`): - Input dimension of the learnable embedding matrix to be projected to `time_embed_dim` when performing class - conditioning with `class_embed_type` equal to `None`. - """ - - _supports_gradient_checkpointing = True - _skip_layerwise_casting_patterns = ["norm"] - - @register_to_config - def __init__( - self, - sample_size: Optional[Union[int, Tuple[int, int]]] = 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, - 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), - layers_per_block: int = 2, - mid_block_scale_factor: float = 1, - downsample_padding: int = 1, - downsample_type: str = "conv", - upsample_type: str = "conv", - dropout: float = 0.0, - act_fn: str = "silu", - attention_head_dim: Optional[int] = 8, - norm_num_groups: int = 32, - attn_norm_num_groups: Optional[int] = 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, - ): - super().__init__() - - self.sample_size = sample_size - time_embed_dim = time_embedding_dim or block_out_channels[0] * 4 - - # Check inputs - if len(down_block_types) != len(up_block_types): - raise ValueError( - f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}." - ) - - if len(block_out_channels) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." - ) - - # input - self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1)) - - # time - if time_embedding_type == "fourier": - self.time_proj = GaussianFourierProjection(embedding_size=block_out_channels[0], scale=16) - timestep_input_dim = 2 * block_out_channels[0] - elif time_embedding_type == "positional": - self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) - timestep_input_dim = block_out_channels[0] - elif time_embedding_type == "learned": - self.time_proj = nn.Embedding(num_train_timesteps, block_out_channels[0]) - timestep_input_dim = block_out_channels[0] - - self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim) - - # class embedding - if class_embed_type is None and num_class_embeds is not None: - self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) - elif class_embed_type == "timestep": - self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim) - elif class_embed_type == "identity": - self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) - else: - self.class_embedding = None - - self.down_blocks = nn.ModuleList([]) - self.mid_block = None - self.up_blocks = nn.ModuleList([]) - - # down - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = get_down_block( - down_block_type, - num_layers=layers_per_block, - in_channels=input_channel, - out_channels=output_channel, - temb_channels=time_embed_dim, - add_downsample=not is_final_block, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel, - downsample_padding=downsample_padding, - resnet_time_scale_shift=resnet_time_scale_shift, - downsample_type=downsample_type, - dropout=dropout, - ) - self.down_blocks.append(down_block) - - # mid - if mid_block_type is None: - self.mid_block = None - else: - self.mid_block = UNetMidBlock2D( - in_channels=block_out_channels[-1], - temb_channels=time_embed_dim, - dropout=dropout, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - output_scale_factor=mid_block_scale_factor, - resnet_time_scale_shift=resnet_time_scale_shift, - attention_head_dim=attention_head_dim if attention_head_dim is not None else block_out_channels[-1], - resnet_groups=norm_num_groups, - attn_groups=attn_norm_num_groups, - add_attention=add_attention, - ) - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] - - is_final_block = i == len(block_out_channels) - 1 - - up_block = get_up_block( - up_block_type, - num_layers=layers_per_block + 1, - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - temb_channels=time_embed_dim, - add_upsample=not is_final_block, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel, - resnet_time_scale_shift=resnet_time_scale_shift, - upsample_type=upsample_type, - dropout=dropout, - ) - self.up_blocks.append(up_block) - - # out - num_groups_out = norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4, 32) - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=num_groups_out, eps=norm_eps) - self.conv_act = nn.SiLU() - self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=3, padding=1) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], - class_labels: Optional[torch.Tensor] = None, - return_dict: bool = True, - ) -> Union[UNet2DOutput, Tuple]: - r""" - The [`UNet2DModel`] forward method. - - Args: - sample (`torch.Tensor`): - The noisy input tensor with the following shape `(batch, channel, height, width)`. - timestep (`torch.Tensor` or `float` or `int`): The number of timesteps to denoise an input. - class_labels (`torch.Tensor`, *optional*, defaults to `None`): - Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.unets.unet_2d.UNet2DOutput`] instead of a plain tuple. - - Returns: - [`~models.unets.unet_2d.UNet2DOutput`] or `tuple`: - If `return_dict` is True, an [`~models.unets.unet_2d.UNet2DOutput`] is returned, otherwise a `tuple` is - returned where the first element is the sample tensor. - """ - # 0. center input if necessary - if self.config.center_input_sample: - sample = 2 * sample - 1.0 - - # 1. time - timesteps = timestep - if not torch.is_tensor(timesteps): - timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device) - elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps * torch.ones(sample.shape[0], dtype=timesteps.dtype, device=timesteps.device) - - t_emb = self.time_proj(timesteps) - - # timesteps does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - t_emb = t_emb.to(dtype=self.dtype) - emb = self.time_embedding(t_emb) - - if self.class_embedding is not None: - if class_labels is None: - raise ValueError("class_labels should be provided when doing class conditioning") - - if self.config.class_embed_type == "timestep": - class_labels = self.time_proj(class_labels) - - class_emb = self.class_embedding(class_labels).to(dtype=self.dtype) - emb = emb + class_emb - elif self.class_embedding is None and class_labels is not None: - raise ValueError("class_embedding needs to be initialized in order to use class conditioning") - - # 2. pre-process - skip_sample = sample - sample = self.conv_in(sample) - - # 3. down - down_block_res_samples = (sample,) - for downsample_block in self.down_blocks: - if hasattr(downsample_block, "skip_conv"): - sample, res_samples, skip_sample = downsample_block( - hidden_states=sample, temb=emb, skip_sample=skip_sample - ) - else: - sample, res_samples = downsample_block(hidden_states=sample, temb=emb) - - down_block_res_samples += res_samples - - # 4. mid - if self.mid_block is not None: - sample = self.mid_block(sample, emb) - - # 5. up - skip_sample = None - for upsample_block in self.up_blocks: - res_samples = down_block_res_samples[-len(upsample_block.resnets) :] - down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] - - if hasattr(upsample_block, "skip_conv"): - sample, skip_sample = upsample_block(sample, res_samples, emb, skip_sample) - else: - sample = upsample_block(sample, res_samples, emb) - - # 6. post-process - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - sample = self.conv_out(sample) - - if skip_sample is not None: - sample += skip_sample - - if self.config.time_embedding_type == "fourier": - timesteps = timesteps.reshape((sample.shape[0], *([1] * len(sample.shape[1:])))) - sample = sample / timesteps - - if not return_dict: - return (sample,) - - return UNet2DOutput(sample=sample) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d_blocks.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d_blocks.py deleted file mode 100644 index 94a9245e567c88fa2a9c81c3b81959137db4beca..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d_blocks.py +++ /dev/null @@ -1,3583 +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 Any, Dict, Optional, Tuple, Union - -import numpy as np -import torch -import torch.nn.functional as F -from torch import nn - -from ...utils import deprecate, logging -from ...utils.torch_utils import apply_freeu -from ..activations import get_activation -from ..attention_processor import Attention, AttnAddedKVProcessor, AttnAddedKVProcessor2_0 -from ..normalization import AdaGroupNorm -from ..resnet import ( - Downsample2D, - FirDownsample2D, - FirUpsample2D, - KDownsample2D, - KUpsample2D, - ResnetBlock2D, - ResnetBlockCondNorm2D, - Upsample2D, -) -from ..transformers.dual_transformer_2d import DualTransformer2DModel -from ..transformers.transformer_2d import Transformer2DModel - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -def get_down_block( - down_block_type: str, - num_layers: int, - in_channels: int, - out_channels: int, - temb_channels: int, - add_downsample: bool, - 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, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - only_cross_attention: bool = False, - upcast_attention: bool = False, - resnet_time_scale_shift: str = "default", - 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, - dropout: float = 0.0, -): - # If attn head dim is not defined, we default it to the number of heads - if attention_head_dim is None: - logger.warning( - f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}." - ) - attention_head_dim = num_attention_heads - - down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type - if down_block_type == "DownBlock2D": - return DownBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - dropout=dropout, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - downsample_padding=downsample_padding, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif down_block_type == "ResnetDownsampleBlock2D": - return ResnetDownsampleBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - dropout=dropout, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - skip_time_act=resnet_skip_time_act, - output_scale_factor=resnet_out_scale_factor, - ) - elif down_block_type == "AttnDownBlock2D": - if add_downsample is False: - downsample_type = None - else: - downsample_type = downsample_type or "conv" # default to 'conv' - return AttnDownBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - dropout=dropout, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - downsample_padding=downsample_padding, - attention_head_dim=attention_head_dim, - resnet_time_scale_shift=resnet_time_scale_shift, - downsample_type=downsample_type, - ) - elif down_block_type == "CrossAttnDownBlock2D": - if cross_attention_dim is None: - raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock2D") - return CrossAttnDownBlock2D( - num_layers=num_layers, - transformer_layers_per_block=transformer_layers_per_block, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - dropout=dropout, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - downsample_padding=downsample_padding, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads, - dual_cross_attention=dual_cross_attention, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - attention_type=attention_type, - ) - elif down_block_type == "SimpleCrossAttnDownBlock2D": - if cross_attention_dim is None: - raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnDownBlock2D") - return SimpleCrossAttnDownBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - dropout=dropout, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - cross_attention_dim=cross_attention_dim, - attention_head_dim=attention_head_dim, - resnet_time_scale_shift=resnet_time_scale_shift, - skip_time_act=resnet_skip_time_act, - output_scale_factor=resnet_out_scale_factor, - only_cross_attention=only_cross_attention, - cross_attention_norm=cross_attention_norm, - ) - elif down_block_type == "SkipDownBlock2D": - return SkipDownBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - dropout=dropout, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - downsample_padding=downsample_padding, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif down_block_type == "AttnSkipDownBlock2D": - return AttnSkipDownBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - dropout=dropout, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - attention_head_dim=attention_head_dim, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif down_block_type == "DownEncoderBlock2D": - return DownEncoderBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - dropout=dropout, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - downsample_padding=downsample_padding, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif down_block_type == "AttnDownEncoderBlock2D": - return AttnDownEncoderBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - dropout=dropout, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - downsample_padding=downsample_padding, - attention_head_dim=attention_head_dim, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif down_block_type == "KDownBlock2D": - return KDownBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - dropout=dropout, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - ) - elif down_block_type == "KCrossAttnDownBlock2D": - return KCrossAttnDownBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - dropout=dropout, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - cross_attention_dim=cross_attention_dim, - attention_head_dim=attention_head_dim, - add_self_attention=True if not add_downsample else False, - ) - raise ValueError(f"{down_block_type} does not exist.") - - -def get_mid_block( - mid_block_type: str, - temb_channels: int, - in_channels: int, - resnet_eps: float, - resnet_act_fn: str, - 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, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - mid_block_only_cross_attention: bool = False, - upcast_attention: bool = False, - 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, - dropout: float = 0.0, -): - if mid_block_type == "UNetMidBlock2DCrossAttn": - return UNetMidBlock2DCrossAttn( - transformer_layers_per_block=transformer_layers_per_block, - in_channels=in_channels, - temb_channels=temb_channels, - dropout=dropout, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - output_scale_factor=output_scale_factor, - resnet_time_scale_shift=resnet_time_scale_shift, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads, - resnet_groups=resnet_groups, - dual_cross_attention=dual_cross_attention, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - attention_type=attention_type, - ) - elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn": - return UNetMidBlock2DSimpleCrossAttn( - in_channels=in_channels, - temb_channels=temb_channels, - dropout=dropout, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - output_scale_factor=output_scale_factor, - cross_attention_dim=cross_attention_dim, - attention_head_dim=attention_head_dim, - resnet_groups=resnet_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - skip_time_act=resnet_skip_time_act, - only_cross_attention=mid_block_only_cross_attention, - cross_attention_norm=cross_attention_norm, - ) - elif mid_block_type == "UNetMidBlock2D": - return UNetMidBlock2D( - in_channels=in_channels, - temb_channels=temb_channels, - dropout=dropout, - num_layers=0, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - output_scale_factor=output_scale_factor, - resnet_groups=resnet_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - add_attention=False, - ) - elif mid_block_type is None: - return None - else: - raise ValueError(f"unknown mid_block_type : {mid_block_type}") - - -def get_up_block( - up_block_type: str, - num_layers: int, - in_channels: int, - out_channels: int, - prev_output_channel: int, - temb_channels: int, - add_upsample: bool, - resnet_eps: float, - resnet_act_fn: str, - resolution_idx: Optional[int] = None, - transformer_layers_per_block: int = 1, - num_attention_heads: Optional[int] = None, - resnet_groups: Optional[int] = None, - cross_attention_dim: Optional[int] = None, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - only_cross_attention: bool = False, - upcast_attention: bool = False, - resnet_time_scale_shift: str = "default", - 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, - dropout: float = 0.0, -) -> nn.Module: - # If attn head dim is not defined, we default it to the number of heads - if attention_head_dim is None: - logger.warning( - f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}." - ) - attention_head_dim = num_attention_heads - - up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type - if up_block_type == "UpBlock2D": - return UpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif up_block_type == "ResnetUpsampleBlock2D": - return ResnetUpsampleBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - skip_time_act=resnet_skip_time_act, - output_scale_factor=resnet_out_scale_factor, - ) - elif up_block_type == "CrossAttnUpBlock2D": - if cross_attention_dim is None: - raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D") - return CrossAttnUpBlock2D( - num_layers=num_layers, - transformer_layers_per_block=transformer_layers_per_block, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads, - dual_cross_attention=dual_cross_attention, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - attention_type=attention_type, - ) - elif up_block_type == "SimpleCrossAttnUpBlock2D": - if cross_attention_dim is None: - raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnUpBlock2D") - return SimpleCrossAttnUpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - cross_attention_dim=cross_attention_dim, - attention_head_dim=attention_head_dim, - resnet_time_scale_shift=resnet_time_scale_shift, - skip_time_act=resnet_skip_time_act, - output_scale_factor=resnet_out_scale_factor, - only_cross_attention=only_cross_attention, - cross_attention_norm=cross_attention_norm, - ) - elif up_block_type == "AttnUpBlock2D": - if add_upsample is False: - upsample_type = None - else: - upsample_type = upsample_type or "conv" # default to 'conv' - - return AttnUpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - attention_head_dim=attention_head_dim, - resnet_time_scale_shift=resnet_time_scale_shift, - upsample_type=upsample_type, - ) - elif up_block_type == "SkipUpBlock2D": - return SkipUpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif up_block_type == "AttnSkipUpBlock2D": - return AttnSkipUpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - attention_head_dim=attention_head_dim, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif up_block_type == "UpDecoderBlock2D": - return UpDecoderBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - temb_channels=temb_channels, - ) - elif up_block_type == "AttnUpDecoderBlock2D": - return AttnUpDecoderBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - attention_head_dim=attention_head_dim, - resnet_time_scale_shift=resnet_time_scale_shift, - temb_channels=temb_channels, - ) - elif up_block_type == "KUpBlock2D": - return KUpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - ) - elif up_block_type == "KCrossAttnUpBlock2D": - return KCrossAttnUpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - cross_attention_dim=cross_attention_dim, - attention_head_dim=attention_head_dim, - ) - - raise ValueError(f"{up_block_type} does not exist.") - - -class AutoencoderTinyBlock(nn.Module): - """ - Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU - blocks. - - Args: - in_channels (`int`): The number of input channels. - out_channels (`int`): The number of output channels. - act_fn (`str`): - ` The activation function to use. Supported values are `"swish"`, `"mish"`, `"gelu"`, and `"relu"`. - - Returns: - `torch.Tensor`: A tensor with the same shape as the input tensor, but with the number of channels equal to - `out_channels`. - """ - - def __init__(self, in_channels: int, out_channels: int, act_fn: str): - super().__init__() - act_fn = get_activation(act_fn) - self.conv = nn.Sequential( - nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), - act_fn, - nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), - act_fn, - nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), - ) - self.skip = ( - nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False) - if in_channels != out_channels - else nn.Identity() - ) - self.fuse = nn.ReLU() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.fuse(self.conv(x) + self.skip(x)) - - -class UNetMidBlock2D(nn.Module): - """ - A 2D UNet mid-block [`UNetMidBlock2D`] with multiple residual blocks and optional attention blocks. - - Args: - in_channels (`int`): The number of input channels. - temb_channels (`int`): The number of temporal embedding channels. - dropout (`float`, *optional*, defaults to 0.0): The dropout rate. - num_layers (`int`, *optional*, defaults to 1): The number of residual blocks. - resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks. - resnet_time_scale_shift (`str`, *optional*, defaults to `default`): - The type of normalization to apply to the time embeddings. This can help to improve the performance of the - model on tasks with long-range temporal dependencies. - 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. - 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. - attention_head_dim (`int`, *optional*, defaults to 1): - Dimension of a single attention head. The number of attention heads is determined based on this value and - the number of input channels. - output_scale_factor (`float`, *optional*, defaults to 1.0): The output scale factor. - - Returns: - `torch.Tensor`: The output of the last residual block, which is a tensor of shape `(batch_size, in_channels, - height, width)`. - - """ - - def __init__( - self, - in_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", # default, spatial - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - attn_groups: Optional[int] = None, - resnet_pre_norm: bool = True, - add_attention: bool = True, - attention_head_dim: int = 1, - output_scale_factor: float = 1.0, - ): - super().__init__() - resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) - self.add_attention = add_attention - - if attn_groups is None: - attn_groups = resnet_groups if resnet_time_scale_shift == "default" else None - - # there is always at least one resnet - if resnet_time_scale_shift == "spatial": - resnets = [ - ResnetBlockCondNorm2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm="spatial", - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - ) - ] - else: - resnets = [ - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ] - attentions = [] - - if attention_head_dim is None: - logger.warning( - f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}." - ) - attention_head_dim = in_channels - - for _ in range(num_layers): - if self.add_attention: - attentions.append( - Attention( - in_channels, - heads=in_channels // attention_head_dim, - dim_head=attention_head_dim, - rescale_output_factor=output_scale_factor, - eps=resnet_eps, - norm_num_groups=attn_groups, - spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None, - residual_connection=True, - bias=True, - upcast_softmax=True, - _from_deprecated_attn_block=True, - ) - ) - else: - attentions.append(None) - - if resnet_time_scale_shift == "spatial": - resnets.append( - ResnetBlockCondNorm2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm="spatial", - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - ) - ) - else: - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - self.gradient_checkpointing = False - - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = 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: - if attn is not None: - hidden_states = attn(hidden_states, temb=temb) - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - if attn is not None: - hidden_states = attn(hidden_states, temb=temb) - hidden_states = resnet(hidden_states, temb) - - return hidden_states - - -class UNetMidBlock2DCrossAttn(nn.Module): - def __init__( - self, - in_channels: int, - temb_channels: int, - out_channels: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[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_pre_norm: bool = True, - num_attention_heads: int = 1, - output_scale_factor: float = 1.0, - cross_attention_dim: int = 1280, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - upcast_attention: bool = False, - attention_type: str = "default", - ): - super().__init__() - - out_channels = out_channels or in_channels - self.in_channels = in_channels - self.out_channels = out_channels - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) - - # support for variable transformer layers per block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * num_layers - - resnet_groups_out = resnet_groups_out or resnet_groups - - # there is always at least one resnet - resnets = [ - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - groups_out=resnet_groups_out, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ] - attentions = [] - - for i in range(num_layers): - if not dual_cross_attention: - attentions.append( - Transformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups_out, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - attention_type=attention_type, - ) - ) - else: - attentions.append( - DualTransformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - ) - ) - resnets.append( - ResnetBlock2D( - in_channels=out_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups_out, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - self.gradient_checkpointing = False - - 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, - ) -> 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.") - - 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: - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - hidden_states = resnet(hidden_states, temb) - - return hidden_states - - -class UNetMidBlock2DSimpleCrossAttn(nn.Module): - def __init__( - self, - in_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attention_head_dim: int = 1, - output_scale_factor: float = 1.0, - cross_attention_dim: int = 1280, - skip_time_act: bool = False, - only_cross_attention: bool = False, - cross_attention_norm: Optional[str] = None, - ): - super().__init__() - - self.has_cross_attention = True - - self.attention_head_dim = attention_head_dim - resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) - - self.num_heads = in_channels // self.attention_head_dim - - # there is always at least one resnet - resnets = [ - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - skip_time_act=skip_time_act, - ) - ] - attentions = [] - - for _ in range(num_layers): - processor = ( - AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor() - ) - - attentions.append( - Attention( - query_dim=in_channels, - cross_attention_dim=in_channels, - heads=self.num_heads, - dim_head=self.attention_head_dim, - added_kv_proj_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - bias=True, - upcast_softmax=True, - only_cross_attention=only_cross_attention, - cross_attention_norm=cross_attention_norm, - processor=processor, - ) - ) - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - skip_time_act=skip_time_act, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - 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, - ) -> 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.") - - if attention_mask is None: - # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask. - mask = None if encoder_hidden_states is None else encoder_attention_mask - else: - # when attention_mask is defined: we don't even check for encoder_attention_mask. - # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks. - # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask. - # then we can simplify this whole if/else block to: - # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask - mask = attention_mask - - hidden_states = self.resnets[0](hidden_states, temb) - for attn, resnet in zip(self.attentions, self.resnets[1:]): - # attn - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=mask, - **cross_attention_kwargs, - ) - - # resnet - hidden_states = resnet(hidden_states, temb) - - return hidden_states - - -class AttnDownBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attention_head_dim: int = 1, - output_scale_factor: float = 1.0, - downsample_padding: int = 1, - downsample_type: str = "conv", - ): - super().__init__() - resnets = [] - attentions = [] - self.downsample_type = downsample_type - - if attention_head_dim is None: - logger.warning( - f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." - ) - attention_head_dim = out_channels - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - attentions.append( - Attention( - out_channels, - heads=out_channels // attention_head_dim, - dim_head=attention_head_dim, - rescale_output_factor=output_scale_factor, - eps=resnet_eps, - norm_num_groups=resnet_groups, - residual_connection=True, - bias=True, - upcast_softmax=True, - _from_deprecated_attn_block=True, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - if downsample_type == "conv": - self.downsamplers = nn.ModuleList( - [ - Downsample2D( - out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" - ) - ] - ) - elif downsample_type == "resnet": - self.downsamplers = nn.ModuleList( - [ - ResnetBlock2D( - in_channels=out_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - down=True, - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - 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, ...]]: - 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.") - - output_states = () - - for resnet, attn in zip(self.resnets, self.attentions): - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - hidden_states = attn(hidden_states, **cross_attention_kwargs) - output_states = output_states + (hidden_states,) - else: - hidden_states = resnet(hidden_states, temb) - hidden_states = attn(hidden_states, **cross_attention_kwargs) - output_states = output_states + (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - if self.downsample_type == "resnet": - hidden_states = downsampler(hidden_states, temb=temb) - else: - hidden_states = downsampler(hidden_states) - - output_states += (hidden_states,) - - return hidden_states, output_states - - -class CrossAttnDownBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[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_pre_norm: bool = True, - num_attention_heads: int = 1, - cross_attention_dim: int = 1280, - output_scale_factor: float = 1.0, - downsample_padding: int = 1, - add_downsample: bool = True, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - only_cross_attention: bool = False, - upcast_attention: bool = False, - attention_type: str = "default", - ): - super().__init__() - resnets = [] - attentions = [] - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * num_layers - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - if not dual_cross_attention: - attentions.append( - Transformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - attention_type=attention_type, - ) - ) - else: - attentions.append( - DualTransformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - ) - ) - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - Downsample2D( - out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - 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, ...]]: - 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.") - - output_states = () - - blocks = list(zip(self.resnets, self.attentions)) - - for i, (resnet, attn) in enumerate(blocks): - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - else: - hidden_states = resnet(hidden_states, temb) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - - # apply additional residuals to the output of the last pair of resnet and attention blocks - if i == len(blocks) - 1 and additional_residuals is not None: - hidden_states = hidden_states + additional_residuals - - output_states = output_states + (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - output_states = output_states + (hidden_states,) - - return hidden_states, output_states - - -class DownBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor: float = 1.0, - add_downsample: bool = True, - downsample_padding: int = 1, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - self.resnets = nn.ModuleList(resnets) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - Downsample2D( - out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - def forward( - self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = 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) - - output_states = () - - for resnet in self.resnets: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = resnet(hidden_states, temb) - - output_states = output_states + (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - output_states = output_states + (hidden_states,) - - return hidden_states, output_states - - -class DownEncoderBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor: float = 1.0, - add_downsample: bool = True, - downsample_padding: int = 1, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - if resnet_time_scale_shift == "spatial": - resnets.append( - ResnetBlockCondNorm2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=None, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm="spatial", - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - ) - ) - else: - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=None, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - self.resnets = nn.ModuleList(resnets) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - Downsample2D( - out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" - ) - ] - ) - else: - self.downsamplers = None - - def forward(self, hidden_states: torch.Tensor, *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) - - for resnet in self.resnets: - hidden_states = resnet(hidden_states, temb=None) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - return hidden_states - - -class AttnDownEncoderBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attention_head_dim: int = 1, - output_scale_factor: float = 1.0, - add_downsample: bool = True, - downsample_padding: int = 1, - ): - super().__init__() - resnets = [] - attentions = [] - - if attention_head_dim is None: - logger.warning( - f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." - ) - attention_head_dim = out_channels - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - if resnet_time_scale_shift == "spatial": - resnets.append( - ResnetBlockCondNorm2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=None, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm="spatial", - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - ) - ) - else: - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=None, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - attentions.append( - Attention( - out_channels, - heads=out_channels // attention_head_dim, - dim_head=attention_head_dim, - rescale_output_factor=output_scale_factor, - eps=resnet_eps, - norm_num_groups=resnet_groups, - residual_connection=True, - bias=True, - upcast_softmax=True, - _from_deprecated_attn_block=True, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - Downsample2D( - out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" - ) - ] - ) - else: - self.downsamplers = None - - def forward(self, hidden_states: torch.Tensor, *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) - - for resnet, attn in zip(self.resnets, self.attentions): - hidden_states = resnet(hidden_states, temb=None) - hidden_states = attn(hidden_states) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - return hidden_states - - -class AttnSkipDownBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_pre_norm: bool = True, - attention_head_dim: int = 1, - output_scale_factor: float = np.sqrt(2.0), - add_downsample: bool = True, - ): - super().__init__() - self.attentions = nn.ModuleList([]) - self.resnets = nn.ModuleList([]) - - if attention_head_dim is None: - logger.warning( - f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." - ) - attention_head_dim = out_channels - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - self.resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=min(in_channels // 4, 32), - groups_out=min(out_channels // 4, 32), - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - self.attentions.append( - Attention( - out_channels, - heads=out_channels // attention_head_dim, - dim_head=attention_head_dim, - rescale_output_factor=output_scale_factor, - eps=resnet_eps, - norm_num_groups=32, - residual_connection=True, - bias=True, - upcast_softmax=True, - _from_deprecated_attn_block=True, - ) - ) - - if add_downsample: - self.resnet_down = ResnetBlock2D( - in_channels=out_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=min(out_channels // 4, 32), - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - use_in_shortcut=True, - down=True, - kernel="fir", - ) - self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)]) - self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1)) - else: - self.resnet_down = None - self.downsamplers = None - self.skip_conv = None - - def forward( - self, - hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - skip_sample: Optional[torch.Tensor] = None, - *args, - **kwargs, - ) -> 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) - - output_states = () - - for resnet, attn in zip(self.resnets, self.attentions): - hidden_states = resnet(hidden_states, temb) - hidden_states = attn(hidden_states) - output_states += (hidden_states,) - - if self.downsamplers is not None: - hidden_states = self.resnet_down(hidden_states, temb) - for downsampler in self.downsamplers: - skip_sample = downsampler(skip_sample) - - hidden_states = self.skip_conv(skip_sample) + hidden_states - - output_states += (hidden_states,) - - return hidden_states, output_states, skip_sample - - -class SkipDownBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_pre_norm: bool = True, - output_scale_factor: float = np.sqrt(2.0), - add_downsample: bool = True, - downsample_padding: int = 1, - ): - super().__init__() - self.resnets = nn.ModuleList([]) - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - self.resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=min(in_channels // 4, 32), - groups_out=min(out_channels // 4, 32), - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - if add_downsample: - self.resnet_down = ResnetBlock2D( - in_channels=out_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=min(out_channels // 4, 32), - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - use_in_shortcut=True, - down=True, - kernel="fir", - ) - self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)]) - self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1)) - else: - self.resnet_down = None - self.downsamplers = None - self.skip_conv = None - - def forward( - self, - hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - skip_sample: Optional[torch.Tensor] = None, - *args, - **kwargs, - ) -> 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) - - output_states = () - - for resnet in self.resnets: - hidden_states = resnet(hidden_states, temb) - output_states += (hidden_states,) - - if self.downsamplers is not None: - hidden_states = self.resnet_down(hidden_states, temb) - for downsampler in self.downsamplers: - skip_sample = downsampler(skip_sample) - - hidden_states = self.skip_conv(skip_sample) + hidden_states - - output_states += (hidden_states,) - - return hidden_states, output_states, skip_sample - - -class ResnetDownsampleBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor: float = 1.0, - add_downsample: bool = True, - skip_time_act: bool = False, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - skip_time_act=skip_time_act, - ) - ) - - self.resnets = nn.ModuleList(resnets) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - ResnetBlock2D( - in_channels=out_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - skip_time_act=skip_time_act, - down=True, - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - def forward( - self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = 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) - - output_states = () - - for resnet in self.resnets: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = resnet(hidden_states, temb) - - output_states = output_states + (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states, temb) - - output_states = output_states + (hidden_states,) - - return hidden_states, output_states - - -class SimpleCrossAttnDownBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attention_head_dim: int = 1, - cross_attention_dim: int = 1280, - output_scale_factor: float = 1.0, - add_downsample: bool = True, - skip_time_act: bool = False, - only_cross_attention: bool = False, - cross_attention_norm: Optional[str] = None, - ): - super().__init__() - - self.has_cross_attention = True - - resnets = [] - attentions = [] - - self.attention_head_dim = attention_head_dim - self.num_heads = out_channels // self.attention_head_dim - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - skip_time_act=skip_time_act, - ) - ) - - processor = ( - AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor() - ) - - attentions.append( - Attention( - query_dim=out_channels, - cross_attention_dim=out_channels, - heads=self.num_heads, - dim_head=attention_head_dim, - added_kv_proj_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - bias=True, - upcast_softmax=True, - only_cross_attention=only_cross_attention, - cross_attention_norm=cross_attention_norm, - processor=processor, - ) - ) - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - ResnetBlock2D( - in_channels=out_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - skip_time_act=skip_time_act, - down=True, - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - 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, ...]]: - 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.") - - output_states = () - - if attention_mask is None: - # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask. - mask = None if encoder_hidden_states is None else encoder_attention_mask - else: - # when attention_mask is defined: we don't even check for encoder_attention_mask. - # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks. - # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask. - # then we can simplify this whole if/else block to: - # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask - mask = attention_mask - - for resnet, attn in zip(self.resnets, self.attentions): - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=mask, - **cross_attention_kwargs, - ) - else: - hidden_states = resnet(hidden_states, temb) - - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=mask, - **cross_attention_kwargs, - ) - - output_states = output_states + (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states, temb) - - output_states = output_states + (hidden_states,) - - return hidden_states, output_states - - -class KDownBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 4, - resnet_eps: float = 1e-5, - resnet_act_fn: str = "gelu", - resnet_group_size: int = 32, - add_downsample: bool = False, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - groups = in_channels // resnet_group_size - groups_out = out_channels // resnet_group_size - - resnets.append( - ResnetBlockCondNorm2D( - in_channels=in_channels, - out_channels=out_channels, - dropout=dropout, - temb_channels=temb_channels, - groups=groups, - groups_out=groups_out, - eps=resnet_eps, - non_linearity=resnet_act_fn, - time_embedding_norm="ada_group", - conv_shortcut_bias=False, - ) - ) - - self.resnets = nn.ModuleList(resnets) - - if add_downsample: - # YiYi's comments- might be able to use FirDownsample2D, look into details later - self.downsamplers = nn.ModuleList([KDownsample2D()]) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - def forward( - self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = 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) - - output_states = () - - for resnet in self.resnets: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = resnet(hidden_states, temb) - - output_states += (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - return hidden_states, output_states - - -class KCrossAttnDownBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - cross_attention_dim: int, - dropout: float = 0.0, - num_layers: int = 4, - resnet_group_size: int = 32, - add_downsample: bool = True, - attention_head_dim: int = 64, - add_self_attention: bool = False, - resnet_eps: float = 1e-5, - resnet_act_fn: str = "gelu", - ): - super().__init__() - resnets = [] - attentions = [] - - self.has_cross_attention = True - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - groups = in_channels // resnet_group_size - groups_out = out_channels // resnet_group_size - - resnets.append( - ResnetBlockCondNorm2D( - in_channels=in_channels, - out_channels=out_channels, - dropout=dropout, - temb_channels=temb_channels, - groups=groups, - groups_out=groups_out, - eps=resnet_eps, - non_linearity=resnet_act_fn, - time_embedding_norm="ada_group", - conv_shortcut_bias=False, - ) - ) - attentions.append( - KAttentionBlock( - out_channels, - out_channels // attention_head_dim, - attention_head_dim, - cross_attention_dim=cross_attention_dim, - temb_channels=temb_channels, - attention_bias=True, - add_self_attention=add_self_attention, - cross_attention_norm="layer_norm", - group_size=resnet_group_size, - ) - ) - - self.resnets = nn.ModuleList(resnets) - self.attentions = nn.ModuleList(attentions) - - if add_downsample: - self.downsamplers = nn.ModuleList([KDownsample2D()]) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - 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, ...]]: - 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.") - - output_states = () - - for resnet, attn in zip(self.resnets, self.attentions): - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - resnet, - hidden_states, - temb, - ) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - emb=temb, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - encoder_attention_mask=encoder_attention_mask, - ) - else: - hidden_states = resnet(hidden_states, temb) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - emb=temb, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - encoder_attention_mask=encoder_attention_mask, - ) - - if self.downsamplers is None: - output_states += (None,) - else: - output_states += (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - return hidden_states, output_states - - -class AttnUpBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - prev_output_channel: int, - out_channels: int, - temb_channels: int, - resolution_idx: int = None, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attention_head_dim: int = 1, - output_scale_factor: float = 1.0, - upsample_type: str = "conv", - ): - super().__init__() - resnets = [] - attentions = [] - - self.upsample_type = upsample_type - - if attention_head_dim is None: - logger.warning( - f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." - ) - attention_head_dim = out_channels - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - attentions.append( - Attention( - out_channels, - heads=out_channels // attention_head_dim, - dim_head=attention_head_dim, - rescale_output_factor=output_scale_factor, - eps=resnet_eps, - norm_num_groups=resnet_groups, - residual_connection=True, - bias=True, - upcast_softmax=True, - _from_deprecated_attn_block=True, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - if upsample_type == "conv": - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) - elif upsample_type == "resnet": - self.upsamplers = nn.ModuleList( - [ - ResnetBlock2D( - in_channels=out_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - up=True, - ) - ] - ) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - def forward( - self, - hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = 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) - - for resnet, attn in zip(self.resnets, self.attentions): - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - hidden_states = attn(hidden_states) - else: - hidden_states = resnet(hidden_states, temb) - hidden_states = attn(hidden_states) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - if self.upsample_type == "resnet": - hidden_states = upsampler(hidden_states, temb=temb) - else: - hidden_states = upsampler(hidden_states) - - return hidden_states - - -class CrossAttnUpBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - prev_output_channel: int, - temb_channels: int, - resolution_idx: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[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_pre_norm: bool = True, - num_attention_heads: int = 1, - cross_attention_dim: int = 1280, - output_scale_factor: float = 1.0, - add_upsample: bool = True, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - only_cross_attention: bool = False, - upcast_attention: bool = False, - attention_type: str = "default", - ): - super().__init__() - resnets = [] - attentions = [] - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * num_layers - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - if not dual_cross_attention: - attentions.append( - Transformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - attention_type=attention_type, - ) - ) - else: - attentions.append( - DualTransformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - ) - ) - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - if add_upsample: - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - 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, - ) -> 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.") - - is_freeu_enabled = ( - getattr(self, "s1", None) - and getattr(self, "s2", None) - and getattr(self, "b1", None) - and getattr(self, "b2", None) - ) - - for resnet, attn in zip(self.resnets, self.attentions): - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - - # FreeU: Only operate on the first two stages - if is_freeu_enabled: - hidden_states, res_hidden_states = apply_freeu( - self.resolution_idx, - hidden_states, - res_hidden_states, - s1=self.s1, - s2=self.s2, - b1=self.b1, - b2=self.b2, - ) - - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - else: - hidden_states = resnet(hidden_states, temb) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) - - return hidden_states - - -class UpBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - prev_output_channel: int, - out_channels: int, - temb_channels: int, - resolution_idx: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor: float = 1.0, - add_upsample: bool = True, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - self.resnets = nn.ModuleList(resnets) - - if add_upsample: - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - def forward( - self, - hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = 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) - - is_freeu_enabled = ( - getattr(self, "s1", None) - and getattr(self, "s2", None) - and getattr(self, "b1", None) - and getattr(self, "b2", None) - ) - - for resnet in self.resnets: - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - - # FreeU: Only operate on the first two stages - if is_freeu_enabled: - hidden_states, res_hidden_states = apply_freeu( - self.resolution_idx, - hidden_states, - res_hidden_states, - s1=self.s1, - s2=self.s2, - b1=self.b1, - b2=self.b2, - ) - - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = resnet(hidden_states, temb) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) - - return hidden_states - - -class UpDecoderBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - resolution_idx: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", # default, spatial - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor: float = 1.0, - add_upsample: bool = True, - temb_channels: Optional[int] = None, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - input_channels = in_channels if i == 0 else out_channels - - if resnet_time_scale_shift == "spatial": - resnets.append( - ResnetBlockCondNorm2D( - in_channels=input_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm="spatial", - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - ) - ) - else: - resnets.append( - ResnetBlock2D( - in_channels=input_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - self.resnets = nn.ModuleList(resnets) - - if add_upsample: - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - self.resolution_idx = resolution_idx - - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: - for resnet in self.resnets: - hidden_states = resnet(hidden_states, temb=temb) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states) - - return hidden_states - - -class AttnUpDecoderBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - resolution_idx: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attention_head_dim: int = 1, - output_scale_factor: float = 1.0, - add_upsample: bool = True, - temb_channels: Optional[int] = None, - ): - super().__init__() - resnets = [] - attentions = [] - - if attention_head_dim is None: - logger.warning( - f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `out_channels`: {out_channels}." - ) - attention_head_dim = out_channels - - for i in range(num_layers): - input_channels = in_channels if i == 0 else out_channels - - if resnet_time_scale_shift == "spatial": - resnets.append( - ResnetBlockCondNorm2D( - in_channels=input_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm="spatial", - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - ) - ) - else: - resnets.append( - ResnetBlock2D( - in_channels=input_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - attentions.append( - Attention( - out_channels, - heads=out_channels // attention_head_dim, - dim_head=attention_head_dim, - rescale_output_factor=output_scale_factor, - eps=resnet_eps, - norm_num_groups=resnet_groups if resnet_time_scale_shift != "spatial" else None, - spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None, - residual_connection=True, - bias=True, - upcast_softmax=True, - _from_deprecated_attn_block=True, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - if add_upsample: - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - self.resolution_idx = resolution_idx - - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = 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) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states) - - return hidden_states - - -class AttnSkipUpBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - prev_output_channel: int, - out_channels: int, - temb_channels: int, - resolution_idx: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_pre_norm: bool = True, - attention_head_dim: int = 1, - output_scale_factor: float = np.sqrt(2.0), - add_upsample: bool = True, - ): - super().__init__() - self.attentions = nn.ModuleList([]) - self.resnets = nn.ModuleList([]) - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - self.resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=min(resnet_in_channels + res_skip_channels // 4, 32), - groups_out=min(out_channels // 4, 32), - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - if attention_head_dim is None: - logger.warning( - f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `out_channels`: {out_channels}." - ) - attention_head_dim = out_channels - - self.attentions.append( - Attention( - out_channels, - heads=out_channels // attention_head_dim, - dim_head=attention_head_dim, - rescale_output_factor=output_scale_factor, - eps=resnet_eps, - norm_num_groups=32, - residual_connection=True, - bias=True, - upcast_softmax=True, - _from_deprecated_attn_block=True, - ) - ) - - self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels) - if add_upsample: - self.resnet_up = ResnetBlock2D( - in_channels=out_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=min(out_channels // 4, 32), - groups_out=min(out_channels // 4, 32), - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - use_in_shortcut=True, - up=True, - kernel="fir", - ) - self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - self.skip_norm = torch.nn.GroupNorm( - num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True - ) - self.act = nn.SiLU() - else: - self.resnet_up = None - self.skip_conv = None - self.skip_norm = None - self.act = None - - self.resolution_idx = resolution_idx - - def forward( - self, - hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - skip_sample=None, - *args, - **kwargs, - ) -> 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) - - for resnet in self.resnets: - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - hidden_states = resnet(hidden_states, temb) - - hidden_states = self.attentions[0](hidden_states) - - if skip_sample is not None: - skip_sample = self.upsampler(skip_sample) - else: - skip_sample = 0 - - if self.resnet_up is not None: - skip_sample_states = self.skip_norm(hidden_states) - skip_sample_states = self.act(skip_sample_states) - skip_sample_states = self.skip_conv(skip_sample_states) - - skip_sample = skip_sample + skip_sample_states - - hidden_states = self.resnet_up(hidden_states, temb) - - return hidden_states, skip_sample - - -class SkipUpBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - prev_output_channel: int, - out_channels: int, - temb_channels: int, - resolution_idx: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_pre_norm: bool = True, - output_scale_factor: float = np.sqrt(2.0), - add_upsample: bool = True, - upsample_padding: int = 1, - ): - super().__init__() - self.resnets = nn.ModuleList([]) - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - self.resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=min((resnet_in_channels + res_skip_channels) // 4, 32), - groups_out=min(out_channels // 4, 32), - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels) - if add_upsample: - self.resnet_up = ResnetBlock2D( - in_channels=out_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=min(out_channels // 4, 32), - groups_out=min(out_channels // 4, 32), - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - use_in_shortcut=True, - up=True, - kernel="fir", - ) - self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) - self.skip_norm = torch.nn.GroupNorm( - num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True - ) - self.act = nn.SiLU() - else: - self.resnet_up = None - self.skip_conv = None - self.skip_norm = None - self.act = None - - self.resolution_idx = resolution_idx - - def forward( - self, - hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - skip_sample=None, - *args, - **kwargs, - ) -> 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) - - for resnet in self.resnets: - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - hidden_states = resnet(hidden_states, temb) - - if skip_sample is not None: - skip_sample = self.upsampler(skip_sample) - else: - skip_sample = 0 - - if self.resnet_up is not None: - skip_sample_states = self.skip_norm(hidden_states) - skip_sample_states = self.act(skip_sample_states) - skip_sample_states = self.skip_conv(skip_sample_states) - - skip_sample = skip_sample + skip_sample_states - - hidden_states = self.resnet_up(hidden_states, temb) - - return hidden_states, skip_sample - - -class ResnetUpsampleBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - prev_output_channel: int, - out_channels: int, - temb_channels: int, - resolution_idx: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor: float = 1.0, - add_upsample: bool = True, - skip_time_act: bool = False, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - skip_time_act=skip_time_act, - ) - ) - - self.resnets = nn.ModuleList(resnets) - - if add_upsample: - self.upsamplers = nn.ModuleList( - [ - ResnetBlock2D( - in_channels=out_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - skip_time_act=skip_time_act, - up=True, - ) - ] - ) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - def forward( - self, - hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = 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) - - for resnet in self.resnets: - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = resnet(hidden_states, temb) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, temb) - - return hidden_states - - -class SimpleCrossAttnUpBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - prev_output_channel: int, - temb_channels: int, - resolution_idx: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attention_head_dim: int = 1, - cross_attention_dim: int = 1280, - output_scale_factor: float = 1.0, - add_upsample: bool = True, - skip_time_act: bool = False, - only_cross_attention: bool = False, - cross_attention_norm: Optional[str] = None, - ): - super().__init__() - resnets = [] - attentions = [] - - self.has_cross_attention = True - self.attention_head_dim = attention_head_dim - - self.num_heads = out_channels // self.attention_head_dim - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - skip_time_act=skip_time_act, - ) - ) - - processor = ( - AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor() - ) - - attentions.append( - Attention( - query_dim=out_channels, - cross_attention_dim=out_channels, - heads=self.num_heads, - dim_head=self.attention_head_dim, - added_kv_proj_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - bias=True, - upcast_softmax=True, - only_cross_attention=only_cross_attention, - cross_attention_norm=cross_attention_norm, - processor=processor, - ) - ) - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - if add_upsample: - self.upsamplers = nn.ModuleList( - [ - ResnetBlock2D( - in_channels=out_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - skip_time_act=skip_time_act, - up=True, - ) - ] - ) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - 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, - ) -> 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.") - - if attention_mask is None: - # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask. - mask = None if encoder_hidden_states is None else encoder_attention_mask - else: - # when attention_mask is defined: we don't even check for encoder_attention_mask. - # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks. - # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask. - # then we can simplify this whole if/else block to: - # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask - mask = attention_mask - - for resnet, attn in zip(self.resnets, self.attentions): - # resnet - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=mask, - **cross_attention_kwargs, - ) - else: - hidden_states = resnet(hidden_states, temb) - - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=mask, - **cross_attention_kwargs, - ) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, temb) - - return hidden_states - - -class KUpBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - resolution_idx: int, - dropout: float = 0.0, - num_layers: int = 5, - resnet_eps: float = 1e-5, - resnet_act_fn: str = "gelu", - resnet_group_size: Optional[int] = 32, - add_upsample: bool = True, - ): - super().__init__() - resnets = [] - k_in_channels = 2 * out_channels - k_out_channels = in_channels - num_layers = num_layers - 1 - - for i in range(num_layers): - in_channels = k_in_channels if i == 0 else out_channels - groups = in_channels // resnet_group_size - groups_out = out_channels // resnet_group_size - - resnets.append( - ResnetBlockCondNorm2D( - in_channels=in_channels, - out_channels=k_out_channels if (i == num_layers - 1) else out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=groups, - groups_out=groups_out, - dropout=dropout, - non_linearity=resnet_act_fn, - time_embedding_norm="ada_group", - conv_shortcut_bias=False, - ) - ) - - self.resnets = nn.ModuleList(resnets) - - if add_upsample: - self.upsamplers = nn.ModuleList([KUpsample2D()]) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - def forward( - self, - hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = 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) - - res_hidden_states_tuple = res_hidden_states_tuple[-1] - if res_hidden_states_tuple is not None: - hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], dim=1) - - for resnet in self.resnets: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = resnet(hidden_states, temb) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states) - - return hidden_states - - -class KCrossAttnUpBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - resolution_idx: int, - dropout: float = 0.0, - num_layers: int = 4, - resnet_eps: float = 1e-5, - resnet_act_fn: str = "gelu", - resnet_group_size: int = 32, - attention_head_dim: int = 1, # attention dim_head - cross_attention_dim: int = 768, - add_upsample: bool = True, - upcast_attention: bool = False, - ): - super().__init__() - resnets = [] - attentions = [] - - is_first_block = in_channels == out_channels == temb_channels - is_middle_block = in_channels != out_channels - add_self_attention = True if is_first_block else False - - self.has_cross_attention = True - self.attention_head_dim = attention_head_dim - - # in_channels, and out_channels for the block (k-unet) - k_in_channels = out_channels if is_first_block else 2 * out_channels - k_out_channels = in_channels - - num_layers = num_layers - 1 - - for i in range(num_layers): - in_channels = k_in_channels if i == 0 else out_channels - groups = in_channels // resnet_group_size - groups_out = out_channels // resnet_group_size - - if is_middle_block and (i == num_layers - 1): - conv_2d_out_channels = k_out_channels - else: - conv_2d_out_channels = None - - resnets.append( - ResnetBlockCondNorm2D( - in_channels=in_channels, - out_channels=out_channels, - conv_2d_out_channels=conv_2d_out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=groups, - groups_out=groups_out, - dropout=dropout, - non_linearity=resnet_act_fn, - time_embedding_norm="ada_group", - conv_shortcut_bias=False, - ) - ) - attentions.append( - KAttentionBlock( - k_out_channels if (i == num_layers - 1) else out_channels, - k_out_channels // attention_head_dim - if (i == num_layers - 1) - else out_channels // attention_head_dim, - attention_head_dim, - cross_attention_dim=cross_attention_dim, - temb_channels=temb_channels, - attention_bias=True, - add_self_attention=add_self_attention, - cross_attention_norm="layer_norm", - upcast_attention=upcast_attention, - ) - ) - - self.resnets = nn.ModuleList(resnets) - self.attentions = nn.ModuleList(attentions) - - if add_upsample: - self.upsamplers = nn.ModuleList([KUpsample2D()]) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - 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, - ) -> torch.Tensor: - res_hidden_states_tuple = res_hidden_states_tuple[-1] - if res_hidden_states_tuple is not None: - hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], dim=1) - - for resnet, attn in zip(self.resnets, self.attentions): - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - resnet, - hidden_states, - temb, - ) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - emb=temb, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - encoder_attention_mask=encoder_attention_mask, - ) - else: - hidden_states = resnet(hidden_states, temb) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - emb=temb, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - encoder_attention_mask=encoder_attention_mask, - ) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states) - - return hidden_states - - -# can potentially later be renamed to `No-feed-forward` attention -class KAttentionBlock(nn.Module): - r""" - A basic Transformer block. - - Parameters: - 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. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. - attention_bias (`bool`, *optional*, defaults to `False`): - Configure if the attention layers should contain a bias parameter. - upcast_attention (`bool`, *optional*, defaults to `False`): - Set to `True` to upcast the attention computation to `float32`. - temb_channels (`int`, *optional*, defaults to 768): - The number of channels in the token embedding. - add_self_attention (`bool`, *optional*, defaults to `False`): - Set to `True` to add self-attention to the block. - cross_attention_norm (`str`, *optional*, defaults to `None`): - The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`. - group_size (`int`, *optional*, defaults to 32): - The number of groups to separate the channels into for group normalization. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - dropout: float = 0.0, - cross_attention_dim: Optional[int] = 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, - group_size: int = 32, - ): - super().__init__() - self.add_self_attention = add_self_attention - - # 1. Self-Attn - if add_self_attention: - self.norm1 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size)) - self.attn1 = Attention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - cross_attention_dim=None, - cross_attention_norm=None, - ) - - # 2. Cross-Attn - self.norm2 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size)) - self.attn2 = Attention( - query_dim=dim, - cross_attention_dim=cross_attention_dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - upcast_attention=upcast_attention, - cross_attention_norm=cross_attention_norm, - ) - - def _to_3d(self, hidden_states: torch.Tensor, height: int, weight: int) -> torch.Tensor: - return hidden_states.permute(0, 2, 3, 1).reshape(hidden_states.shape[0], height * weight, -1) - - def _to_4d(self, hidden_states: torch.Tensor, height: int, weight: int) -> torch.Tensor: - return hidden_states.permute(0, 2, 1).reshape(hidden_states.shape[0], -1, height, weight) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = 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, - ) -> 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.") - - # 1. Self-Attention - if self.add_self_attention: - norm_hidden_states = self.norm1(hidden_states, emb) - - height, weight = norm_hidden_states.shape[2:] - norm_hidden_states = self._to_3d(norm_hidden_states, height, weight) - - attn_output = self.attn1( - norm_hidden_states, - encoder_hidden_states=None, - attention_mask=attention_mask, - **cross_attention_kwargs, - ) - attn_output = self._to_4d(attn_output, height, weight) - - hidden_states = attn_output + hidden_states - - # 2. Cross-Attention/None - norm_hidden_states = self.norm2(hidden_states, emb) - - height, weight = norm_hidden_states.shape[2:] - norm_hidden_states = self._to_3d(norm_hidden_states, height, weight) - attn_output = self.attn2( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask if encoder_hidden_states is None else encoder_attention_mask, - **cross_attention_kwargs, - ) - attn_output = self._to_4d(attn_output, height, weight) - - hidden_states = attn_output + hidden_states - - return hidden_states diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d_blocks_flax.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d_blocks_flax.py deleted file mode 100644 index abd025165ecf6411215df54679a28cb17d4e9c9d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d_blocks_flax.py +++ /dev/null @@ -1,401 +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 flax.linen as nn -import jax.numpy as jnp - -from ..attention_flax import FlaxTransformer2DModel -from ..resnet_flax import FlaxDownsample2D, FlaxResnetBlock2D, FlaxUpsample2D - - -class FlaxCrossAttnDownBlock2D(nn.Module): - r""" - Cross Attention 2D Downsizing block - original architecture from Unet transformers: - https://huggingface.co/papers/2103.06104 - - Parameters: - in_channels (:obj:`int`): - Input channels - out_channels (:obj:`int`): - Output channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of attention blocks layers - num_attention_heads (:obj:`int`, *optional*, defaults to 1): - Number of attention heads of each spatial transformer block - add_downsample (:obj:`bool`, *optional*, defaults to `True`): - Whether to add downsampling layer before each final output - use_memory_efficient_attention (`bool`, *optional*, defaults to `False`): - enable memory efficient attention https://huggingface.co/papers/2112.05682 - split_head_dim (`bool`, *optional*, defaults to `False`): - Whether to split the head dimension into a new axis for the self-attention computation. In most cases, - enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL. - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - out_channels: int - dropout: float = 0.0 - num_layers: int = 1 - num_attention_heads: int = 1 - add_downsample: bool = True - use_linear_projection: bool = False - only_cross_attention: bool = False - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - dtype: jnp.dtype = jnp.float32 - transformer_layers_per_block: int = 1 - - def setup(self): - resnets = [] - attentions = [] - - for i in range(self.num_layers): - in_channels = self.in_channels if i == 0 else self.out_channels - - res_block = FlaxResnetBlock2D( - in_channels=in_channels, - out_channels=self.out_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - - attn_block = FlaxTransformer2DModel( - in_channels=self.out_channels, - n_heads=self.num_attention_heads, - d_head=self.out_channels // self.num_attention_heads, - depth=self.transformer_layers_per_block, - use_linear_projection=self.use_linear_projection, - only_cross_attention=self.only_cross_attention, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - attentions.append(attn_block) - - self.resnets = resnets - self.attentions = attentions - - if self.add_downsample: - self.downsamplers_0 = FlaxDownsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, temb, encoder_hidden_states, deterministic=True): - output_states = () - - for resnet, attn in zip(self.resnets, self.attentions): - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - hidden_states = attn(hidden_states, encoder_hidden_states, deterministic=deterministic) - output_states += (hidden_states,) - - if self.add_downsample: - hidden_states = self.downsamplers_0(hidden_states) - output_states += (hidden_states,) - - return hidden_states, output_states - - -class FlaxDownBlock2D(nn.Module): - r""" - Flax 2D downsizing block - - Parameters: - in_channels (:obj:`int`): - Input channels - out_channels (:obj:`int`): - Output channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of attention blocks layers - add_downsample (:obj:`bool`, *optional*, defaults to `True`): - Whether to add downsampling layer before each final output - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - out_channels: int - dropout: float = 0.0 - num_layers: int = 1 - add_downsample: bool = True - dtype: jnp.dtype = jnp.float32 - - def setup(self): - resnets = [] - - for i in range(self.num_layers): - in_channels = self.in_channels if i == 0 else self.out_channels - - res_block = FlaxResnetBlock2D( - in_channels=in_channels, - out_channels=self.out_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - self.resnets = resnets - - if self.add_downsample: - self.downsamplers_0 = FlaxDownsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, temb, deterministic=True): - output_states = () - - for resnet in self.resnets: - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - output_states += (hidden_states,) - - if self.add_downsample: - hidden_states = self.downsamplers_0(hidden_states) - output_states += (hidden_states,) - - return hidden_states, output_states - - -class FlaxCrossAttnUpBlock2D(nn.Module): - r""" - Cross Attention 2D Upsampling block - original architecture from Unet transformers: - https://huggingface.co/papers/2103.06104 - - Parameters: - in_channels (:obj:`int`): - Input channels - out_channels (:obj:`int`): - Output channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of attention blocks layers - num_attention_heads (:obj:`int`, *optional*, defaults to 1): - Number of attention heads of each spatial transformer block - add_upsample (:obj:`bool`, *optional*, defaults to `True`): - Whether to add upsampling layer before each final output - use_memory_efficient_attention (`bool`, *optional*, defaults to `False`): - enable memory efficient attention https://huggingface.co/papers/2112.05682 - split_head_dim (`bool`, *optional*, defaults to `False`): - Whether to split the head dimension into a new axis for the self-attention computation. In most cases, - enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL. - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - out_channels: int - prev_output_channel: int - dropout: float = 0.0 - num_layers: int = 1 - num_attention_heads: int = 1 - add_upsample: bool = True - use_linear_projection: bool = False - only_cross_attention: bool = False - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - dtype: jnp.dtype = jnp.float32 - transformer_layers_per_block: int = 1 - - def setup(self): - resnets = [] - attentions = [] - - for i in range(self.num_layers): - res_skip_channels = self.in_channels if (i == self.num_layers - 1) else self.out_channels - resnet_in_channels = self.prev_output_channel if i == 0 else self.out_channels - - res_block = FlaxResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=self.out_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - - attn_block = FlaxTransformer2DModel( - in_channels=self.out_channels, - n_heads=self.num_attention_heads, - d_head=self.out_channels // self.num_attention_heads, - depth=self.transformer_layers_per_block, - use_linear_projection=self.use_linear_projection, - only_cross_attention=self.only_cross_attention, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - attentions.append(attn_block) - - self.resnets = resnets - self.attentions = attentions - - if self.add_upsample: - self.upsamplers_0 = FlaxUpsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, res_hidden_states_tuple, temb, encoder_hidden_states, deterministic=True): - for resnet, attn in zip(self.resnets, self.attentions): - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = jnp.concatenate((hidden_states, res_hidden_states), axis=-1) - - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - hidden_states = attn(hidden_states, encoder_hidden_states, deterministic=deterministic) - - if self.add_upsample: - hidden_states = self.upsamplers_0(hidden_states) - - return hidden_states - - -class FlaxUpBlock2D(nn.Module): - r""" - Flax 2D upsampling block - - Parameters: - in_channels (:obj:`int`): - Input channels - out_channels (:obj:`int`): - Output channels - prev_output_channel (:obj:`int`): - Output channels from the previous block - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of attention blocks layers - add_downsample (:obj:`bool`, *optional*, defaults to `True`): - Whether to add downsampling layer before each final output - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - out_channels: int - prev_output_channel: int - dropout: float = 0.0 - num_layers: int = 1 - add_upsample: bool = True - dtype: jnp.dtype = jnp.float32 - - def setup(self): - resnets = [] - - for i in range(self.num_layers): - res_skip_channels = self.in_channels if (i == self.num_layers - 1) else self.out_channels - resnet_in_channels = self.prev_output_channel if i == 0 else self.out_channels - - res_block = FlaxResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=self.out_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - - self.resnets = resnets - - if self.add_upsample: - self.upsamplers_0 = FlaxUpsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, res_hidden_states_tuple, temb, deterministic=True): - for resnet in self.resnets: - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = jnp.concatenate((hidden_states, res_hidden_states), axis=-1) - - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - - if self.add_upsample: - hidden_states = self.upsamplers_0(hidden_states) - - return hidden_states - - -class FlaxUNetMidBlock2DCrossAttn(nn.Module): - r""" - Cross Attention 2D Mid-level block - original architecture from Unet transformers: - https://huggingface.co/papers/2103.06104 - - Parameters: - in_channels (:obj:`int`): - Input channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of attention blocks layers - num_attention_heads (:obj:`int`, *optional*, defaults to 1): - Number of attention heads of each spatial transformer block - use_memory_efficient_attention (`bool`, *optional*, defaults to `False`): - enable memory efficient attention https://huggingface.co/papers/2112.05682 - split_head_dim (`bool`, *optional*, defaults to `False`): - Whether to split the head dimension into a new axis for the self-attention computation. In most cases, - enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL. - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - dropout: float = 0.0 - num_layers: int = 1 - num_attention_heads: int = 1 - use_linear_projection: bool = False - use_memory_efficient_attention: bool = False - split_head_dim: bool = False - dtype: jnp.dtype = jnp.float32 - transformer_layers_per_block: int = 1 - - def setup(self): - # there is always at least one resnet - resnets = [ - FlaxResnetBlock2D( - in_channels=self.in_channels, - out_channels=self.in_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - ] - - attentions = [] - - for _ in range(self.num_layers): - attn_block = FlaxTransformer2DModel( - in_channels=self.in_channels, - n_heads=self.num_attention_heads, - d_head=self.in_channels // self.num_attention_heads, - depth=self.transformer_layers_per_block, - use_linear_projection=self.use_linear_projection, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - attentions.append(attn_block) - - res_block = FlaxResnetBlock2D( - in_channels=self.in_channels, - out_channels=self.in_channels, - dropout_prob=self.dropout, - dtype=self.dtype, - ) - resnets.append(res_block) - - self.resnets = resnets - self.attentions = attentions - - def __call__(self, hidden_states, temb, encoder_hidden_states, deterministic=True): - hidden_states = self.resnets[0](hidden_states, temb) - for attn, resnet in zip(self.attentions, self.resnets[1:]): - hidden_states = attn(hidden_states, encoder_hidden_states, deterministic=deterministic) - hidden_states = resnet(hidden_states, temb, deterministic=deterministic) - - return hidden_states diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d_condition.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d_condition.py deleted file mode 100644 index 0cf5133c5405067c0cec8f51652cfa01e595552b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d_condition.py +++ /dev/null @@ -1,1310 +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 dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union - -import torch -import torch.nn as nn -import torch.utils.checkpoint - -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 ..activations import get_activation -from ..attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - Attention, - AttentionProcessor, - AttnAddedKVProcessor, - AttnProcessor, - FusedAttnProcessor2_0, -) -from ..embeddings import ( - GaussianFourierProjection, - GLIGENTextBoundingboxProjection, - ImageHintTimeEmbedding, - ImageProjection, - ImageTimeEmbedding, - TextImageProjection, - TextImageTimeEmbedding, - TextTimeEmbedding, - TimestepEmbedding, - Timesteps, -) -from ..modeling_utils import ModelMixin -from .unet_2d_blocks import ( - get_down_block, - get_mid_block, - get_up_block, -) - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class UNet2DConditionOutput(BaseOutput): - """ - The output of [`UNet2DConditionModel`]. - - Args: - sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)`): - The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model. - """ - - sample: torch.Tensor = None - - -class UNet2DConditionModel( - ModelMixin, ConfigMixin, FromOriginalModelMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin -): - r""" - A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample - shaped output. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - Parameters: - 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. - center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample. - 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")`): - 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")`): - The tuple of upsample blocks to use. - 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)`): - 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. - mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. - 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): - The dimension of the cross attention features. - 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): - 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 - [`~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): - If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim` - dimension to `cross_attention_dim`. - 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 (`int`, *optional*, defaults to 8): The dimension of the attention heads. - num_attention_heads (`int`, *optional*): - The number of attention heads. If not defined, defaults to `attention_head_dim` - resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config - for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`. - class_embed_type (`str`, *optional*, defaults to `None`): - The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`, - `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`. - addition_embed_type (`str`, *optional*, defaults to `None`): - Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or - "text". "text" will use the `TextTimeEmbedding` layer. - addition_time_embed_dim: (`int`, *optional*, defaults to `None`): - Dimension for the timestep embeddings. - num_class_embeds (`int`, *optional*, defaults to `None`): - Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing - class conditioning with `class_embed_type` equal to `None`. - time_embedding_type (`str`, *optional*, defaults to `positional`): - The type of position embedding to use for timesteps. Choose from `positional` or `fourier`. - time_embedding_dim (`int`, *optional*, defaults to `None`): - An optional override for the dimension of the projected time embedding. - time_embedding_act_fn (`str`, *optional*, defaults to `None`): - Optional activation function to use only once on the time embeddings before they are passed to the rest of - the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`. - timestep_post_act (`str`, *optional*, defaults to `None`): - The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`. - time_cond_proj_dim (`int`, *optional*, defaults to `None`): - The dimension of `cond_proj` layer in the timestep embedding. - conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer. - conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer. - projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when - `class_embed_type="projection"`. Required when `class_embed_type="projection"`. - class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time - embeddings with the class embeddings. - mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`): - Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If - `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the - `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False` - otherwise. - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D", "CrossAttnUpBlock2D"] - _skip_layerwise_casting_patterns = ["norm"] - - @register_to_config - def __init__( - self, - sample_size: Optional[Union[int, Tuple[int, int]]] = 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] = ( - "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, - 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_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, - 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, - 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, - conv_in_kernel: int = 3, - conv_out_kernel: int = 3, - projection_class_embeddings_input_dim: Optional[int] = None, - attention_type: str = "default", - class_embeddings_concat: bool = False, - mid_block_only_cross_attention: Optional[bool] = None, - cross_attention_norm: Optional[str] = None, - addition_embed_type_num_heads: int = 64, - ): - super().__init__() - - self.sample_size = sample_size - - if num_attention_heads is not None: - raise ValueError( - "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19." - ) - - # If `num_attention_heads` is not defined (which is the case for most models) - # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. - # The reason for this behavior is to correct for incorrectly named variables that were introduced - # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 - # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking - # which is why we correct for the naming here. - num_attention_heads = num_attention_heads or attention_head_dim - - # Check inputs - self._check_config( - down_block_types=down_block_types, - up_block_types=up_block_types, - only_cross_attention=only_cross_attention, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - cross_attention_dim=cross_attention_dim, - transformer_layers_per_block=transformer_layers_per_block, - reverse_transformer_layers_per_block=reverse_transformer_layers_per_block, - attention_head_dim=attention_head_dim, - num_attention_heads=num_attention_heads, - ) - - # input - conv_in_padding = (conv_in_kernel - 1) // 2 - self.conv_in = nn.Conv2d( - in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding - ) - - # time - time_embed_dim, timestep_input_dim = self._set_time_proj( - time_embedding_type, - block_out_channels=block_out_channels, - flip_sin_to_cos=flip_sin_to_cos, - freq_shift=freq_shift, - time_embedding_dim=time_embedding_dim, - ) - - self.time_embedding = TimestepEmbedding( - timestep_input_dim, - time_embed_dim, - act_fn=act_fn, - post_act_fn=timestep_post_act, - cond_proj_dim=time_cond_proj_dim, - ) - - self._set_encoder_hid_proj( - encoder_hid_dim_type, - cross_attention_dim=cross_attention_dim, - encoder_hid_dim=encoder_hid_dim, - ) - - # class embedding - self._set_class_embedding( - class_embed_type, - act_fn=act_fn, - num_class_embeds=num_class_embeds, - projection_class_embeddings_input_dim=projection_class_embeddings_input_dim, - time_embed_dim=time_embed_dim, - timestep_input_dim=timestep_input_dim, - ) - - self._set_add_embedding( - addition_embed_type, - addition_embed_type_num_heads=addition_embed_type_num_heads, - addition_time_embed_dim=addition_time_embed_dim, - cross_attention_dim=cross_attention_dim, - encoder_hid_dim=encoder_hid_dim, - flip_sin_to_cos=flip_sin_to_cos, - freq_shift=freq_shift, - projection_class_embeddings_input_dim=projection_class_embeddings_input_dim, - time_embed_dim=time_embed_dim, - ) - - if time_embedding_act_fn is None: - self.time_embed_act = None - else: - self.time_embed_act = get_activation(time_embedding_act_fn) - - self.down_blocks = nn.ModuleList([]) - self.up_blocks = nn.ModuleList([]) - - if isinstance(only_cross_attention, bool): - if mid_block_only_cross_attention is None: - mid_block_only_cross_attention = only_cross_attention - - only_cross_attention = [only_cross_attention] * len(down_block_types) - - if mid_block_only_cross_attention is None: - mid_block_only_cross_attention = False - - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(down_block_types) - - if isinstance(attention_head_dim, int): - attention_head_dim = (attention_head_dim,) * len(down_block_types) - - if isinstance(cross_attention_dim, int): - cross_attention_dim = (cross_attention_dim,) * len(down_block_types) - - if isinstance(layers_per_block, int): - layers_per_block = [layers_per_block] * len(down_block_types) - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) - - if class_embeddings_concat: - # The time embeddings are concatenated with the class embeddings. The dimension of the - # time embeddings passed to the down, middle, and up blocks is twice the dimension of the - # regular time embeddings - blocks_time_embed_dim = time_embed_dim * 2 - else: - blocks_time_embed_dim = time_embed_dim - - # down - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = get_down_block( - down_block_type, - num_layers=layers_per_block[i], - transformer_layers_per_block=transformer_layers_per_block[i], - in_channels=input_channel, - out_channels=output_channel, - temb_channels=blocks_time_embed_dim, - add_downsample=not is_final_block, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - cross_attention_dim=cross_attention_dim[i], - num_attention_heads=num_attention_heads[i], - downsample_padding=downsample_padding, - dual_cross_attention=dual_cross_attention, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention[i], - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - attention_type=attention_type, - resnet_skip_time_act=resnet_skip_time_act, - resnet_out_scale_factor=resnet_out_scale_factor, - cross_attention_norm=cross_attention_norm, - attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, - dropout=dropout, - ) - self.down_blocks.append(down_block) - - # mid - self.mid_block = get_mid_block( - mid_block_type, - temb_channels=blocks_time_embed_dim, - in_channels=block_out_channels[-1], - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - output_scale_factor=mid_block_scale_factor, - transformer_layers_per_block=transformer_layers_per_block[-1], - num_attention_heads=num_attention_heads[-1], - cross_attention_dim=cross_attention_dim[-1], - dual_cross_attention=dual_cross_attention, - use_linear_projection=use_linear_projection, - mid_block_only_cross_attention=mid_block_only_cross_attention, - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - attention_type=attention_type, - resnet_skip_time_act=resnet_skip_time_act, - cross_attention_norm=cross_attention_norm, - attention_head_dim=attention_head_dim[-1], - dropout=dropout, - ) - - # count how many layers upsample the images - self.num_upsamplers = 0 - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - reversed_num_attention_heads = list(reversed(num_attention_heads)) - reversed_layers_per_block = list(reversed(layers_per_block)) - reversed_cross_attention_dim = list(reversed(cross_attention_dim)) - reversed_transformer_layers_per_block = ( - list(reversed(transformer_layers_per_block)) - if reverse_transformer_layers_per_block is None - else reverse_transformer_layers_per_block - ) - only_cross_attention = list(reversed(only_cross_attention)) - - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - is_final_block = i == len(block_out_channels) - 1 - - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] - - # add upsample block for all BUT final layer - if not is_final_block: - add_upsample = True - self.num_upsamplers += 1 - else: - add_upsample = False - - up_block = get_up_block( - up_block_type, - num_layers=reversed_layers_per_block[i] + 1, - transformer_layers_per_block=reversed_transformer_layers_per_block[i], - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - temb_channels=blocks_time_embed_dim, - add_upsample=add_upsample, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resolution_idx=i, - resnet_groups=norm_num_groups, - cross_attention_dim=reversed_cross_attention_dim[i], - num_attention_heads=reversed_num_attention_heads[i], - dual_cross_attention=dual_cross_attention, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention[i], - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - attention_type=attention_type, - resnet_skip_time_act=resnet_skip_time_act, - resnet_out_scale_factor=resnet_out_scale_factor, - cross_attention_norm=cross_attention_norm, - attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, - dropout=dropout, - ) - self.up_blocks.append(up_block) - - # out - if norm_num_groups is not None: - self.conv_norm_out = nn.GroupNorm( - num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps - ) - - self.conv_act = get_activation(act_fn) - - else: - self.conv_norm_out = None - self.conv_act = None - - conv_out_padding = (conv_out_kernel - 1) // 2 - self.conv_out = nn.Conv2d( - block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding - ) - - self._set_pos_net_if_use_gligen(attention_type=attention_type, cross_attention_dim=cross_attention_dim) - - 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]]], - reverse_transformer_layers_per_block: bool, - attention_head_dim: int, - num_attention_heads: Optional[Union[int, Tuple[int]]], - ): - if len(down_block_types) != len(up_block_types): - raise ValueError( - f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}." - ) - - if len(block_out_channels) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}." - ) - - if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}." - ) - if isinstance(transformer_layers_per_block, list) and reverse_transformer_layers_per_block is None: - for layer_number_per_block in transformer_layers_per_block: - if isinstance(layer_number_per_block, list): - raise ValueError("Must provide 'reverse_transformer_layers_per_block` if using asymmetrical UNet.") - - def _set_time_proj( - self, - time_embedding_type: str, - block_out_channels: int, - flip_sin_to_cos: bool, - freq_shift: float, - time_embedding_dim: 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: - raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.") - self.time_proj = GaussianFourierProjection( - time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos - ) - timestep_input_dim = time_embed_dim - elif time_embedding_type == "positional": - time_embed_dim = time_embedding_dim or block_out_channels[0] * 4 - - self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) - timestep_input_dim = block_out_channels[0] - else: - raise ValueError( - f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`." - ) - - return time_embed_dim, timestep_input_dim - - def _set_encoder_hid_proj( - self, - encoder_hid_dim_type: Optional[str], - cross_attention_dim: Union[int, Tuple[int]], - encoder_hid_dim: Optional[int], - ): - if encoder_hid_dim_type is None and encoder_hid_dim is not None: - encoder_hid_dim_type = "text_proj" - self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type) - logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.") - - if encoder_hid_dim is None and encoder_hid_dim_type is not None: - raise ValueError( - f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}." - ) - - if encoder_hid_dim_type == "text_proj": - self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim) - elif encoder_hid_dim_type == "text_image_proj": - # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much - # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use - # case when `addition_embed_type == "text_image_proj"` (Kandinsky 2.1)` - self.encoder_hid_proj = TextImageProjection( - text_embed_dim=encoder_hid_dim, - image_embed_dim=cross_attention_dim, - cross_attention_dim=cross_attention_dim, - ) - elif encoder_hid_dim_type == "image_proj": - # Kandinsky 2.2 - self.encoder_hid_proj = ImageProjection( - image_embed_dim=encoder_hid_dim, - cross_attention_dim=cross_attention_dim, - ) - elif encoder_hid_dim_type is not None: - raise ValueError( - f"`encoder_hid_dim_type`: {encoder_hid_dim_type} must be None, 'text_proj', 'text_image_proj', or 'image_proj'." - ) - else: - self.encoder_hid_proj = None - - def _set_class_embedding( - self, - class_embed_type: Optional[str], - act_fn: str, - num_class_embeds: Optional[int], - projection_class_embeddings_input_dim: Optional[int], - time_embed_dim: int, - timestep_input_dim: int, - ): - if class_embed_type is None and num_class_embeds is not None: - self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) - elif class_embed_type == "timestep": - self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn=act_fn) - elif class_embed_type == "identity": - self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) - elif class_embed_type == "projection": - if projection_class_embeddings_input_dim is None: - raise ValueError( - "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set" - ) - # The projection `class_embed_type` is the same as the timestep `class_embed_type` except - # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings - # 2. it projects from an arbitrary input dimension. - # - # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations. - # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings. - # As a result, `TimestepEmbedding` can be passed arbitrary vectors. - self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) - elif class_embed_type == "simple_projection": - if projection_class_embeddings_input_dim is None: - raise ValueError( - "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set" - ) - self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim) - else: - self.class_embedding = None - - def _set_add_embedding( - self, - addition_embed_type: str, - addition_embed_type_num_heads: int, - addition_time_embed_dim: Optional[int], - 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], - time_embed_dim: int, - ): - if addition_embed_type == "text": - if encoder_hid_dim is not None: - text_time_embedding_from_dim = encoder_hid_dim - else: - text_time_embedding_from_dim = cross_attention_dim - - self.add_embedding = TextTimeEmbedding( - text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads - ) - elif addition_embed_type == "text_image": - # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much - # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use - # case when `addition_embed_type == "text_image"` (Kandinsky 2.1)` - self.add_embedding = TextImageTimeEmbedding( - text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim - ) - elif addition_embed_type == "text_time": - self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift) - self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) - elif addition_embed_type == "image": - # Kandinsky 2.2 - self.add_embedding = ImageTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim) - elif addition_embed_type == "image_hint": - # Kandinsky 2.2 ControlNet - self.add_embedding = ImageHintTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim) - elif addition_embed_type is not None: - raise ValueError( - f"`addition_embed_type`: {addition_embed_type} must be None, 'text', 'text_image', 'text_time', 'image', or 'image_hint'." - ) - - def _set_pos_net_if_use_gligen(self, attention_type: str, cross_attention_dim: int): - if attention_type in ["gated", "gated-text-image"]: - positive_len = 768 - if isinstance(cross_attention_dim, int): - positive_len = cross_attention_dim - elif isinstance(cross_attention_dim, (list, tuple)): - positive_len = cross_attention_dim[0] - - feature_type = "text-only" if attention_type == "gated" else "text-image" - self.position_net = GLIGENTextBoundingboxProjection( - positive_len=positive_len, out_dim=cross_attention_dim, feature_type=feature_type - ) - - @property - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - self.set_attn_processor(processor) - - def set_attention_slice(self, slice_size: Union[str, int, List[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. This is useful for saving some memory in exchange for a small decrease in speed. - - Args: - slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): - When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If - `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is - provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` - must be a multiple of `slice_size`. - """ - sliceable_head_dims = [] - - def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): - if hasattr(module, "set_attention_slice"): - sliceable_head_dims.append(module.sliceable_head_dim) - - for child in module.children(): - fn_recursive_retrieve_sliceable_dims(child) - - # retrieve number of attention layers - for module in self.children(): - fn_recursive_retrieve_sliceable_dims(module) - - num_sliceable_layers = len(sliceable_head_dims) - - if slice_size == "auto": - # half the attention head size is usually a good trade-off between - # speed and memory - slice_size = [dim // 2 for dim in sliceable_head_dims] - elif slice_size == "max": - # make smallest slice possible - slice_size = num_sliceable_layers * [1] - - slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size - - if len(slice_size) != len(sliceable_head_dims): - raise ValueError( - f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" - f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." - ) - - for i in range(len(slice_size)): - size = slice_size[i] - dim = sliceable_head_dims[i] - if size is not None and size > dim: - raise ValueError(f"size {size} has to be smaller or equal to {dim}.") - - # 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]): - if hasattr(module, "set_attention_slice"): - module.set_attention_slice(slice_size.pop()) - - for child in module.children(): - fn_recursive_set_attention_slice(child, slice_size) - - reversed_slice_size = list(reversed(slice_size)) - for module in self.children(): - fn_recursive_set_attention_slice(module, reversed_slice_size) - - def enable_freeu(self, s1: float, s2: float, b1: float, b2: float): - r"""Enables the FreeU mechanism from https://huggingface.co/papers/2309.11497. - - The suffixes after the scaling factors represent the stage blocks where they are being applied. - - Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that - are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL. - - Args: - s1 (`float`): - Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to - mitigate the "oversmoothing effect" in the enhanced denoising process. - s2 (`float`): - Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to - mitigate the "oversmoothing effect" in the enhanced denoising process. - b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features. - b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features. - """ - for i, upsample_block in enumerate(self.up_blocks): - setattr(upsample_block, "s1", s1) - setattr(upsample_block, "s2", s2) - setattr(upsample_block, "b1", b1) - setattr(upsample_block, "b2", b2) - - def disable_freeu(self): - """Disables the FreeU mechanism.""" - freeu_keys = {"s1", "s2", "b1", "b2"} - for i, upsample_block in enumerate(self.up_blocks): - for k in freeu_keys: - if hasattr(upsample_block, k) or getattr(upsample_block, k, None) is not None: - setattr(upsample_block, k, None) - - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - self.set_attn_processor(FusedAttnProcessor2_0()) - - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - 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]: - 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 - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = sample.device.type == "mps" - is_npu = sample.device.type == "npu" - if isinstance(timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) - elif len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps.expand(sample.shape[0]) - - t_emb = self.time_proj(timesteps) - # `Timesteps` does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - 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]: - class_emb = None - if self.class_embedding is not None: - if class_labels is None: - raise ValueError("class_labels should be provided when num_class_embeds > 0") - - if self.config.class_embed_type == "timestep": - class_labels = self.time_proj(class_labels) - - # `Timesteps` does not contain any weights and will always return f32 tensors - # there might be better ways to encapsulate this. - class_labels = class_labels.to(dtype=sample.dtype) - - class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype) - 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]: - aug_emb = None - if self.config.addition_embed_type == "text": - aug_emb = self.add_embedding(encoder_hidden_states) - elif self.config.addition_embed_type == "text_image": - # Kandinsky 2.1 - style - if "image_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" - ) - - image_embs = added_cond_kwargs.get("image_embeds") - text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states) - aug_emb = self.add_embedding(text_embs, image_embs) - elif self.config.addition_embed_type == "text_time": - # SDXL - style - if "text_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" - ) - text_embeds = added_cond_kwargs.get("text_embeds") - if "time_ids" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" - ) - time_ids = added_cond_kwargs.get("time_ids") - time_embeds = self.add_time_proj(time_ids.flatten()) - time_embeds = time_embeds.reshape((text_embeds.shape[0], -1)) - add_embeds = torch.concat([text_embeds, time_embeds], dim=-1) - add_embeds = add_embeds.to(emb.dtype) - aug_emb = self.add_embedding(add_embeds) - elif self.config.addition_embed_type == "image": - # Kandinsky 2.2 - style - if "image_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" - ) - image_embs = added_cond_kwargs.get("image_embeds") - aug_emb = self.add_embedding(image_embs) - elif self.config.addition_embed_type == "image_hint": - # Kandinsky 2.2 ControlNet - style - if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`" - ) - image_embs = added_cond_kwargs.get("image_embeds") - hint = added_cond_kwargs.get("hint") - aug_emb = self.add_embedding(image_embs, hint) - return aug_emb - - def process_encoder_hidden_states( - 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) - elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj": - # Kandinsky 2.1 - style - if "image_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" - ) - - image_embeds = added_cond_kwargs.get("image_embeds") - encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds) - elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj": - # Kandinsky 2.2 - style - if "image_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" - ) - image_embeds = added_cond_kwargs.get("image_embeds") - encoder_hidden_states = self.encoder_hid_proj(image_embeds) - elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "ip_image_proj": - if "image_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'ip_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" - ) - - if hasattr(self, "text_encoder_hid_proj") and self.text_encoder_hid_proj is not None: - encoder_hidden_states = self.text_encoder_hid_proj(encoder_hidden_states) - - image_embeds = added_cond_kwargs.get("image_embeds") - image_embeds = self.encoder_hid_proj(image_embeds) - encoder_hidden_states = (encoder_hidden_states, image_embeds) - return encoder_hidden_states - - def forward( - self, - sample: torch.Tensor, - timestep: Union[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, - return_dict: bool = True, - ) -> Union[UNet2DConditionOutput, Tuple]: - r""" - The [`UNet2DConditionModel`] forward method. - - Args: - sample (`torch.Tensor`): - The noisy input tensor with the following shape `(batch, channel, height, width)`. - timestep (`torch.Tensor` or `float` or `int`): The number of timesteps to denoise an input. - encoder_hidden_states (`torch.Tensor`): - The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. - class_labels (`torch.Tensor`, *optional*, defaults to `None`): - Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. - timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`): - Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed - through the `self.time_embedding` layer to obtain the timestep embeddings. - attention_mask (`torch.Tensor`, *optional*, defaults to `None`): - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask - is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large - negative values to the attention scores corresponding to "discard" tokens. - cross_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). - added_cond_kwargs: (`dict`, *optional*): - A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that - are passed along to the UNet blocks. - down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*): - A tuple of tensors that if specified are added to the residuals of down unet blocks. - mid_block_additional_residual: (`torch.Tensor`, *optional*): - A tensor that if specified is added to the residual of the middle unet block. - down_intrablock_additional_residuals (`tuple` of `torch.Tensor`, *optional*): - additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s) - encoder_attention_mask (`torch.Tensor`): - A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If - `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias, - which adds large negative values to the attention scores corresponding to "discard" tokens. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain - tuple. - - Returns: - [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: - If `return_dict` is True, an [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] is returned, - otherwise a `tuple` is returned where the first element is the sample tensor. - """ - # By default samples have to be AT least a multiple of the overall upsampling factor. - # The overall upsampling factor is equal to 2 ** (# num of upsampling layers). - # However, the upsampling interpolation output size can be forced to fit any upsampling size - # on the fly if necessary. - default_overall_up_factor = 2**self.num_upsamplers - - # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` - forward_upsample_size = False - upsample_size = None - - for dim in sample.shape[-2:]: - if dim % default_overall_up_factor != 0: - # Forward upsample size to force interpolation output size. - forward_upsample_size = True - break - - # ensure attention_mask is a bias, and give it a singleton query_tokens dimension - # expects mask of shape: - # [batch, key_tokens] - # adds singleton query_tokens dimension: - # [batch, 1, key_tokens] - # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: - # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) - # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) - if attention_mask is not None: - # assume that mask is expressed as: - # (1 = keep, 0 = discard) - # convert mask into a bias that can be added to attention scores: - # (keep = +0, discard = -10000.0) - attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 - attention_mask = attention_mask.unsqueeze(1) - - # convert encoder_attention_mask to a bias the same way we do for attention_mask - if encoder_attention_mask is not None: - encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0 - encoder_attention_mask = encoder_attention_mask.unsqueeze(1) - - # 0. center input if necessary - if self.config.center_input_sample: - sample = 2 * sample - 1.0 - - # 1. time - t_emb = self.get_time_embed(sample=sample, timestep=timestep) - emb = self.time_embedding(t_emb, timestep_cond) - - class_emb = self.get_class_embed(sample=sample, class_labels=class_labels) - if class_emb is not None: - if self.config.class_embeddings_concat: - emb = torch.cat([emb, class_emb], dim=-1) - else: - emb = emb + class_emb - - aug_emb = self.get_aug_embed( - emb=emb, encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs - ) - if self.config.addition_embed_type == "image_hint": - aug_emb, hint = aug_emb - sample = torch.cat([sample, hint], dim=1) - - emb = emb + aug_emb if aug_emb is not None else emb - - if self.time_embed_act is not None: - emb = self.time_embed_act(emb) - - encoder_hidden_states = self.process_encoder_hidden_states( - encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs - ) - - # 2. pre-process - sample = self.conv_in(sample) - - # 2.5 GLIGEN position net - if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None: - cross_attention_kwargs = cross_attention_kwargs.copy() - gligen_args = cross_attention_kwargs.pop("gligen") - 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 - # maintain backward compatibility for legacy usage, where - # T2I-Adapter and ControlNet both use down_block_additional_residuals arg - # but can only use one or the other - if not is_adapter and mid_block_additional_residual is None and down_block_additional_residuals is not None: - deprecate( - "T2I should not use down_block_additional_residuals", - "1.3.0", - "Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \ - and will be removed in diffusers 1.3.0. `down_block_additional_residuals` should only be used \ - for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ", - standard_warn=False, - ) - down_intrablock_additional_residuals = down_block_additional_residuals - is_adapter = True - - down_block_res_samples = (sample,) - for downsample_block in self.down_blocks: - if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: - # For t2i-adapter CrossAttnDownBlock2D - additional_residuals = {} - if is_adapter and len(down_intrablock_additional_residuals) > 0: - additional_residuals["additional_residuals"] = down_intrablock_additional_residuals.pop(0) - - sample, res_samples = downsample_block( - hidden_states=sample, - temb=emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - encoder_attention_mask=encoder_attention_mask, - **additional_residuals, - ) - else: - sample, res_samples = downsample_block(hidden_states=sample, temb=emb) - if is_adapter and len(down_intrablock_additional_residuals) > 0: - sample += down_intrablock_additional_residuals.pop(0) - - down_block_res_samples += res_samples - - if is_controlnet: - new_down_block_res_samples = () - - for down_block_res_sample, down_block_additional_residual in zip( - down_block_res_samples, down_block_additional_residuals - ): - down_block_res_sample = down_block_res_sample + down_block_additional_residual - new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,) - - down_block_res_samples = new_down_block_res_samples - - # 4. mid - if self.mid_block is not None: - if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention: - sample = self.mid_block( - sample, - emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - encoder_attention_mask=encoder_attention_mask, - ) - else: - sample = self.mid_block(sample, emb) - - # To support T2I-Adapter-XL - if ( - is_adapter - and len(down_intrablock_additional_residuals) > 0 - and sample.shape == down_intrablock_additional_residuals[0].shape - ): - sample += down_intrablock_additional_residuals.pop(0) - - if is_controlnet: - sample = sample + mid_block_additional_residual - - # 5. up - for i, upsample_block in enumerate(self.up_blocks): - is_final_block = i == len(self.up_blocks) - 1 - - res_samples = down_block_res_samples[-len(upsample_block.resnets) :] - down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] - - # if we have not reached the final block and need to forward the - # upsample size, we do it here - if not is_final_block and forward_upsample_size: - upsample_size = down_block_res_samples[-1].shape[2:] - - if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - upsample_size=upsample_size, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - ) - else: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - upsample_size=upsample_size, - ) - - # 6. post-process - if self.conv_norm_out: - sample = self.conv_norm_out(sample) - 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,) - - return UNet2DConditionOutput(sample=sample) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d_condition_flax.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d_condition_flax.py deleted file mode 100644 index 7c21ddb690ae770b13e72e77cb137cd3b607cb89..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d_condition_flax.py +++ /dev/null @@ -1,454 +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 Dict, Optional, Tuple, Union - -import flax -import flax.linen as nn -import jax -import jax.numpy as jnp -from flax.core.frozen_dict import FrozenDict - -from ...configuration_utils import ConfigMixin, flax_register_to_config -from ...utils import BaseOutput -from ..embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps -from ..modeling_flax_utils import FlaxModelMixin -from .unet_2d_blocks_flax import ( - FlaxCrossAttnDownBlock2D, - FlaxCrossAttnUpBlock2D, - FlaxDownBlock2D, - FlaxUNetMidBlock2DCrossAttn, - FlaxUpBlock2D, -) - - -@flax.struct.dataclass -class FlaxUNet2DConditionOutput(BaseOutput): - """ - The output of [`FlaxUNet2DConditionModel`]. - - Args: - sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)`): - The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model. - """ - - sample: jnp.ndarray - - -@flax_register_to_config -class FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin): - r""" - A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample - shaped output. - - This model inherits from [`FlaxModelMixin`]. Check the superclass documentation for it's generic methods - implemented for all models (such as downloading or saving). - - This model is also a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module) - subclass. Use it as a regular Flax Linen module and refer to the Flax documentation for all matters related to its - general usage and behavior. - - Inherent JAX features such as the following are supported: - - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit) - - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation) - - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap) - - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap) - - Parameters: - sample_size (`int`, *optional*): - The size of the input 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 `("FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D")`): - The tuple of downsample blocks to use. - 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)`): - 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): - The dimension of the attention heads. - 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. - dropout (`float`, *optional*, defaults to 0): - Dropout probability for down, up and bottleneck blocks. - 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. - use_memory_efficient_attention (`bool`, *optional*, defaults to `False`): - Enable memory efficient attention as described [here](https://huggingface.co/papers/2112.05682). - split_head_dim (`bool`, *optional*, defaults to `False`): - Whether to split the head dimension into a new axis for the self-attention computation. In most cases, - enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL. - """ - - sample_size: int = 32 - in_channels: int = 4 - out_channels: int = 4 - 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) - layers_per_block: int = 2 - attention_head_dim: Union[int, Tuple[int, ...]] = 8 - num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None - cross_attention_dim: int = 1280 - dropout: float = 0.0 - use_linear_projection: bool = False - dtype: jnp.dtype = jnp.float32 - flip_sin_to_cos: bool = True - 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 - addition_embed_type_num_heads: int = 64 - projection_class_embeddings_input_dim: Optional[int] = None - - def init_weights(self, rng: jax.Array) -> FrozenDict: - # init input tensors - sample_shape = (1, self.in_channels, self.sample_size, self.sample_size) - sample = jnp.zeros(sample_shape, dtype=jnp.float32) - timesteps = jnp.ones((1,), dtype=jnp.int32) - encoder_hidden_states = jnp.zeros((1, 1, self.cross_attention_dim), dtype=jnp.float32) - - params_rng, dropout_rng = jax.random.split(rng) - rngs = {"params": params_rng, "dropout": dropout_rng} - - added_cond_kwargs = None - if self.addition_embed_type == "text_time": - # we retrieve the expected `text_embeds_dim` by first checking if the architecture is a refiner - # or non-refiner architecture and then by "reverse-computing" from `projection_class_embeddings_input_dim` - is_refiner = ( - 5 * self.config.addition_time_embed_dim + self.config.cross_attention_dim - == self.config.projection_class_embeddings_input_dim - ) - num_micro_conditions = 5 if is_refiner else 6 - - text_embeds_dim = self.config.projection_class_embeddings_input_dim - ( - num_micro_conditions * self.config.addition_time_embed_dim - ) - - time_ids_channels = self.projection_class_embeddings_input_dim - text_embeds_dim - time_ids_dims = time_ids_channels // self.addition_time_embed_dim - added_cond_kwargs = { - "text_embeds": jnp.zeros((1, text_embeds_dim), dtype=jnp.float32), - "time_ids": jnp.zeros((1, time_ids_dims), dtype=jnp.float32), - } - return self.init(rngs, sample, timesteps, encoder_hidden_states, added_cond_kwargs)["params"] - - def setup(self) -> None: - block_out_channels = self.block_out_channels - time_embed_dim = block_out_channels[0] * 4 - - if self.num_attention_heads is not None: - raise ValueError( - "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19." - ) - - # If `num_attention_heads` is not defined (which is the case for most models) - # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. - # The reason for this behavior is to correct for incorrectly named variables that were introduced - # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 - # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking - # which is why we correct for the naming here. - num_attention_heads = self.num_attention_heads or self.attention_head_dim - - # input - self.conv_in = nn.Conv( - block_out_channels[0], - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - # time - self.time_proj = FlaxTimesteps( - block_out_channels[0], flip_sin_to_cos=self.flip_sin_to_cos, freq_shift=self.config.freq_shift - ) - self.time_embedding = FlaxTimestepEmbedding(time_embed_dim, dtype=self.dtype) - - only_cross_attention = self.only_cross_attention - if isinstance(only_cross_attention, bool): - only_cross_attention = (only_cross_attention,) * len(self.down_block_types) - - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(self.down_block_types) - - # transformer layers per block - transformer_layers_per_block = self.transformer_layers_per_block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * len(self.down_block_types) - - # addition embed types - if self.addition_embed_type is None: - self.add_embedding = None - elif self.addition_embed_type == "text_time": - if self.addition_time_embed_dim is None: - raise ValueError( - f"addition_embed_type {self.addition_embed_type} requires `addition_time_embed_dim` to not be None" - ) - self.add_time_proj = FlaxTimesteps(self.addition_time_embed_dim, self.flip_sin_to_cos, self.freq_shift) - self.add_embedding = FlaxTimestepEmbedding(time_embed_dim, dtype=self.dtype) - else: - raise ValueError(f"addition_embed_type: {self.addition_embed_type} must be None or `text_time`.") - - # down - down_blocks = [] - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(self.down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - if down_block_type == "CrossAttnDownBlock2D": - down_block = FlaxCrossAttnDownBlock2D( - in_channels=input_channel, - out_channels=output_channel, - dropout=self.dropout, - num_layers=self.layers_per_block, - transformer_layers_per_block=transformer_layers_per_block[i], - num_attention_heads=num_attention_heads[i], - add_downsample=not is_final_block, - use_linear_projection=self.use_linear_projection, - only_cross_attention=only_cross_attention[i], - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - else: - down_block = FlaxDownBlock2D( - in_channels=input_channel, - out_channels=output_channel, - dropout=self.dropout, - num_layers=self.layers_per_block, - add_downsample=not is_final_block, - dtype=self.dtype, - ) - - down_blocks.append(down_block) - self.down_blocks = down_blocks - - # mid - if self.config.mid_block_type == "UNetMidBlock2DCrossAttn": - self.mid_block = FlaxUNetMidBlock2DCrossAttn( - in_channels=block_out_channels[-1], - dropout=self.dropout, - num_attention_heads=num_attention_heads[-1], - transformer_layers_per_block=transformer_layers_per_block[-1], - use_linear_projection=self.use_linear_projection, - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - elif self.config.mid_block_type is None: - self.mid_block = None - else: - raise ValueError(f"Unexpected mid_block_type {self.config.mid_block_type}") - - # up - up_blocks = [] - reversed_block_out_channels = list(reversed(block_out_channels)) - reversed_num_attention_heads = list(reversed(num_attention_heads)) - only_cross_attention = list(reversed(only_cross_attention)) - output_channel = reversed_block_out_channels[0] - reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block)) - for i, up_block_type in enumerate(self.up_block_types): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] - - is_final_block = i == len(block_out_channels) - 1 - - if up_block_type == "CrossAttnUpBlock2D": - up_block = FlaxCrossAttnUpBlock2D( - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - num_layers=self.layers_per_block + 1, - transformer_layers_per_block=reversed_transformer_layers_per_block[i], - num_attention_heads=reversed_num_attention_heads[i], - add_upsample=not is_final_block, - dropout=self.dropout, - use_linear_projection=self.use_linear_projection, - only_cross_attention=only_cross_attention[i], - use_memory_efficient_attention=self.use_memory_efficient_attention, - split_head_dim=self.split_head_dim, - dtype=self.dtype, - ) - else: - up_block = FlaxUpBlock2D( - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - num_layers=self.layers_per_block + 1, - add_upsample=not is_final_block, - dropout=self.dropout, - dtype=self.dtype, - ) - - up_blocks.append(up_block) - prev_output_channel = output_channel - self.up_blocks = up_blocks - - # out - self.conv_norm_out = nn.GroupNorm(num_groups=32, epsilon=1e-5) - self.conv_out = nn.Conv( - self.out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - def __call__( - self, - sample: jnp.ndarray, - timesteps: Union[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, - return_dict: bool = True, - train: bool = False, - ) -> Union[FlaxUNet2DConditionOutput, Tuple[jnp.ndarray]]: - r""" - Args: - sample (`jnp.ndarray`): (batch, channel, height, width) noisy inputs tensor - timestep (`jnp.ndarray` or `float` or `int`): timesteps - encoder_hidden_states (`jnp.ndarray`): (batch_size, sequence_length, hidden_size) encoder hidden states - added_cond_kwargs: (`dict`, *optional*): - A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that - are passed along to the UNet blocks. - down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*): - A tuple of tensors that if specified are added to the residuals of down unet blocks. - mid_block_additional_residual: (`torch.Tensor`, *optional*): - A tensor that if specified is added to the residual of the middle unet block. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`models.unets.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] instead of - a plain tuple. - train (`bool`, *optional*, defaults to `False`): - Use deterministic functions and disable dropout when not training. - - Returns: - [`~models.unets.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] or `tuple`: - [`~models.unets.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - """ - # 1. time - if not isinstance(timesteps, jnp.ndarray): - timesteps = jnp.array([timesteps], dtype=jnp.int32) - elif isinstance(timesteps, jnp.ndarray) and len(timesteps.shape) == 0: - timesteps = timesteps.astype(dtype=jnp.float32) - timesteps = jnp.expand_dims(timesteps, 0) - - t_emb = self.time_proj(timesteps) - t_emb = self.time_embedding(t_emb) - - # additional embeddings - aug_emb = None - if self.addition_embed_type == "text_time": - if added_cond_kwargs is None: - raise ValueError( - f"Need to provide argument `added_cond_kwargs` for {self.__class__} when using `addition_embed_type={self.addition_embed_type}`" - ) - text_embeds = added_cond_kwargs.get("text_embeds") - if text_embeds is None: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" - ) - time_ids = added_cond_kwargs.get("time_ids") - if time_ids is None: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" - ) - # compute time embeds - time_embeds = self.add_time_proj(jnp.ravel(time_ids)) # (1, 6) => (6,) => (6, 256) - time_embeds = jnp.reshape(time_embeds, (text_embeds.shape[0], -1)) - add_embeds = jnp.concatenate([text_embeds, time_embeds], axis=-1) - aug_emb = self.add_embedding(add_embeds) - - t_emb = t_emb + aug_emb if aug_emb is not None else t_emb - - # 2. pre-process - sample = jnp.transpose(sample, (0, 2, 3, 1)) - sample = self.conv_in(sample) - - # 3. down - down_block_res_samples = (sample,) - for down_block in self.down_blocks: - if isinstance(down_block, FlaxCrossAttnDownBlock2D): - sample, res_samples = down_block(sample, t_emb, encoder_hidden_states, deterministic=not train) - else: - sample, res_samples = down_block(sample, t_emb, deterministic=not train) - down_block_res_samples += res_samples - - if down_block_additional_residuals is not None: - new_down_block_res_samples = () - - for down_block_res_sample, down_block_additional_residual in zip( - down_block_res_samples, down_block_additional_residuals - ): - down_block_res_sample += down_block_additional_residual - new_down_block_res_samples += (down_block_res_sample,) - - down_block_res_samples = new_down_block_res_samples - - # 4. mid - if self.mid_block is not None: - sample = self.mid_block(sample, t_emb, encoder_hidden_states, deterministic=not train) - - if mid_block_additional_residual is not None: - sample += mid_block_additional_residual - - # 5. up - for up_block in self.up_blocks: - res_samples = down_block_res_samples[-(self.layers_per_block + 1) :] - down_block_res_samples = down_block_res_samples[: -(self.layers_per_block + 1)] - if isinstance(up_block, FlaxCrossAttnUpBlock2D): - sample = up_block( - sample, - temb=t_emb, - encoder_hidden_states=encoder_hidden_states, - res_hidden_states_tuple=res_samples, - deterministic=not train, - ) - else: - sample = up_block(sample, temb=t_emb, res_hidden_states_tuple=res_samples, deterministic=not train) - - # 6. post-process - sample = self.conv_norm_out(sample) - sample = nn.silu(sample) - sample = self.conv_out(sample) - sample = jnp.transpose(sample, (0, 3, 1, 2)) - - if not return_dict: - return (sample,) - - return FlaxUNet2DConditionOutput(sample=sample) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_3d_blocks.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_3d_blocks.py deleted file mode 100644 index 53c0f4bae38bd91017ed06c5bf8f572903f01911..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_3d_blocks.py +++ /dev/null @@ -1,1427 +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 Any, Dict, Optional, Tuple, Union - -import torch -from torch import nn - -from ...utils import deprecate, logging -from ...utils.torch_utils import apply_freeu -from ..attention import Attention -from ..resnet import ( - Downsample2D, - ResnetBlock2D, - SpatioTemporalResBlock, - TemporalConvLayer, - Upsample2D, -) -from ..transformers.transformer_2d import Transformer2DModel -from ..transformers.transformer_temporal import ( - TransformerSpatioTemporalModel, - TransformerTemporalModel, -) -from .unet_motion_model import ( - CrossAttnDownBlockMotion, - CrossAttnUpBlockMotion, - DownBlockMotion, - UNetMidBlockCrossAttnMotion, - UpBlockMotion, -) - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class DownBlockMotion(DownBlockMotion): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `DownBlockMotion` from `diffusers.models.unets.unet_3d_blocks` is deprecated and this will be removed in a future version. Please use `from diffusers.models.unets.unet_motion_model import DownBlockMotion` instead." - deprecate("DownBlockMotion", "1.0.0", deprecation_message) - super().__init__(*args, **kwargs) - - -class CrossAttnDownBlockMotion(CrossAttnDownBlockMotion): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `CrossAttnDownBlockMotion` from `diffusers.models.unets.unet_3d_blocks` is deprecated and this will be removed in a future version. Please use `from diffusers.models.unets.unet_motion_model import CrossAttnDownBlockMotion` instead." - deprecate("CrossAttnDownBlockMotion", "1.0.0", deprecation_message) - super().__init__(*args, **kwargs) - - -class UpBlockMotion(UpBlockMotion): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `UpBlockMotion` from `diffusers.models.unets.unet_3d_blocks` is deprecated and this will be removed in a future version. Please use `from diffusers.models.unets.unet_motion_model import UpBlockMotion` instead." - deprecate("UpBlockMotion", "1.0.0", deprecation_message) - super().__init__(*args, **kwargs) - - -class CrossAttnUpBlockMotion(CrossAttnUpBlockMotion): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `CrossAttnUpBlockMotion` from `diffusers.models.unets.unet_3d_blocks` is deprecated and this will be removed in a future version. Please use `from diffusers.models.unets.unet_motion_model import CrossAttnUpBlockMotion` instead." - deprecate("CrossAttnUpBlockMotion", "1.0.0", deprecation_message) - super().__init__(*args, **kwargs) - - -class UNetMidBlockCrossAttnMotion(UNetMidBlockCrossAttnMotion): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `UNetMidBlockCrossAttnMotion` from `diffusers.models.unets.unet_3d_blocks` is deprecated and this will be removed in a future version. Please use `from diffusers.models.unets.unet_motion_model import UNetMidBlockCrossAttnMotion` instead." - deprecate("UNetMidBlockCrossAttnMotion", "1.0.0", deprecation_message) - super().__init__(*args, **kwargs) - - -def get_down_block( - down_block_type: str, - num_layers: int, - in_channels: int, - out_channels: int, - temb_channels: int, - add_downsample: bool, - 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, - 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_max_seq_length: int = 32, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1, - dropout: float = 0.0, -) -> Union[ - "DownBlock3D", - "CrossAttnDownBlock3D", - "DownBlockSpatioTemporal", - "CrossAttnDownBlockSpatioTemporal", -]: - if down_block_type == "DownBlock3D": - return DownBlock3D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - downsample_padding=downsample_padding, - resnet_time_scale_shift=resnet_time_scale_shift, - dropout=dropout, - ) - elif down_block_type == "CrossAttnDownBlock3D": - if cross_attention_dim is None: - raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock3D") - return CrossAttnDownBlock3D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - downsample_padding=downsample_padding, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads, - dual_cross_attention=dual_cross_attention, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - dropout=dropout, - ) - elif down_block_type == "DownBlockSpatioTemporal": - # added for SDV - return DownBlockSpatioTemporal( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - add_downsample=add_downsample, - ) - elif down_block_type == "CrossAttnDownBlockSpatioTemporal": - # added for SDV - if cross_attention_dim is None: - raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlockSpatioTemporal") - return CrossAttnDownBlockSpatioTemporal( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - num_layers=num_layers, - transformer_layers_per_block=transformer_layers_per_block, - add_downsample=add_downsample, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads, - ) - - raise ValueError(f"{down_block_type} does not exist.") - - -def get_up_block( - up_block_type: str, - num_layers: int, - in_channels: int, - out_channels: int, - prev_output_channel: int, - temb_channels: int, - add_upsample: bool, - 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, - 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_max_seq_length: int = 32, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1, - dropout: float = 0.0, -) -> Union[ - "UpBlock3D", - "CrossAttnUpBlock3D", - "UpBlockSpatioTemporal", - "CrossAttnUpBlockSpatioTemporal", -]: - if up_block_type == "UpBlock3D": - return UpBlock3D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - resolution_idx=resolution_idx, - dropout=dropout, - ) - elif up_block_type == "CrossAttnUpBlock3D": - if cross_attention_dim is None: - raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock3D") - return CrossAttnUpBlock3D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads, - dual_cross_attention=dual_cross_attention, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - resolution_idx=resolution_idx, - dropout=dropout, - ) - elif up_block_type == "UpBlockSpatioTemporal": - # added for SDV - return UpBlockSpatioTemporal( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - resolution_idx=resolution_idx, - add_upsample=add_upsample, - ) - elif up_block_type == "CrossAttnUpBlockSpatioTemporal": - # added for SDV - if cross_attention_dim is None: - raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlockSpatioTemporal") - return CrossAttnUpBlockSpatioTemporal( - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - num_layers=num_layers, - transformer_layers_per_block=transformer_layers_per_block, - add_upsample=add_upsample, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads, - resolution_idx=resolution_idx, - ) - - raise ValueError(f"{up_block_type} does not exist.") - - -class UNetMidBlock3DCrossAttn(nn.Module): - def __init__( - self, - in_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - num_attention_heads: int = 1, - output_scale_factor: float = 1.0, - cross_attention_dim: int = 1280, - dual_cross_attention: bool = False, - use_linear_projection: bool = True, - upcast_attention: bool = False, - ): - super().__init__() - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) - - # there is always at least one resnet - resnets = [ - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ] - temp_convs = [ - TemporalConvLayer( - in_channels, - in_channels, - dropout=0.1, - norm_num_groups=resnet_groups, - ) - ] - attentions = [] - temp_attentions = [] - - for _ in range(num_layers): - attentions.append( - Transformer2DModel( - in_channels // num_attention_heads, - num_attention_heads, - in_channels=in_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - ) - ) - temp_attentions.append( - TransformerTemporalModel( - in_channels // num_attention_heads, - num_attention_heads, - in_channels=in_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - ) - ) - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - temp_convs.append( - TemporalConvLayer( - in_channels, - in_channels, - dropout=0.1, - norm_num_groups=resnet_groups, - ) - ) - - self.resnets = nn.ModuleList(resnets) - self.temp_convs = nn.ModuleList(temp_convs) - self.attentions = nn.ModuleList(attentions) - self.temp_attentions = nn.ModuleList(temp_attentions) - - 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, - num_frames: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> torch.Tensor: - hidden_states = self.resnets[0](hidden_states, temb) - hidden_states = self.temp_convs[0](hidden_states, num_frames=num_frames) - for attn, temp_attn, resnet, temp_conv in zip( - self.attentions, self.temp_attentions, self.resnets[1:], self.temp_convs[1:] - ): - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - hidden_states = temp_attn( - hidden_states, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - hidden_states = resnet(hidden_states, temb) - hidden_states = temp_conv(hidden_states, num_frames=num_frames) - - return hidden_states - - -class CrossAttnDownBlock3D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - num_attention_heads: int = 1, - cross_attention_dim: int = 1280, - output_scale_factor: float = 1.0, - downsample_padding: int = 1, - add_downsample: bool = True, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - only_cross_attention: bool = False, - upcast_attention: bool = False, - ): - super().__init__() - resnets = [] - attentions = [] - temp_attentions = [] - temp_convs = [] - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - temp_convs.append( - TemporalConvLayer( - out_channels, - out_channels, - dropout=0.1, - norm_num_groups=resnet_groups, - ) - ) - attentions.append( - Transformer2DModel( - out_channels // num_attention_heads, - num_attention_heads, - in_channels=out_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - ) - ) - temp_attentions.append( - TransformerTemporalModel( - out_channels // num_attention_heads, - num_attention_heads, - in_channels=out_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - ) - ) - self.resnets = nn.ModuleList(resnets) - self.temp_convs = nn.ModuleList(temp_convs) - self.attentions = nn.ModuleList(attentions) - self.temp_attentions = nn.ModuleList(temp_attentions) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - Downsample2D( - out_channels, - use_conv=True, - out_channels=out_channels, - padding=downsample_padding, - name="op", - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - 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, - num_frames: int = 1, - cross_attention_kwargs: Dict[str, Any] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: - # TODO(Patrick, William) - attention mask is not used - output_states = () - - for resnet, temp_conv, attn, temp_attn in zip( - self.resnets, self.temp_convs, self.attentions, self.temp_attentions - ): - hidden_states = resnet(hidden_states, temb) - hidden_states = temp_conv(hidden_states, num_frames=num_frames) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - hidden_states = temp_attn( - hidden_states, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - output_states += (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - output_states += (hidden_states,) - - return hidden_states, output_states - - -class DownBlock3D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor: float = 1.0, - add_downsample: bool = True, - downsample_padding: int = 1, - ): - super().__init__() - resnets = [] - temp_convs = [] - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - temp_convs.append( - TemporalConvLayer( - out_channels, - out_channels, - dropout=0.1, - norm_num_groups=resnet_groups, - ) - ) - - self.resnets = nn.ModuleList(resnets) - self.temp_convs = nn.ModuleList(temp_convs) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - Downsample2D( - out_channels, - use_conv=True, - out_channels=out_channels, - padding=downsample_padding, - name="op", - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - num_frames: int = 1, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: - output_states = () - - for resnet, temp_conv in zip(self.resnets, self.temp_convs): - hidden_states = resnet(hidden_states, temb) - hidden_states = temp_conv(hidden_states, num_frames=num_frames) - - output_states += (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - output_states += (hidden_states,) - - return hidden_states, output_states - - -class CrossAttnUpBlock3D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - prev_output_channel: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - num_attention_heads: int = 1, - cross_attention_dim: int = 1280, - output_scale_factor: float = 1.0, - add_upsample: bool = True, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - only_cross_attention: bool = False, - upcast_attention: bool = False, - resolution_idx: Optional[int] = None, - ): - super().__init__() - resnets = [] - temp_convs = [] - attentions = [] - temp_attentions = [] - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - temp_convs.append( - TemporalConvLayer( - out_channels, - out_channels, - dropout=0.1, - norm_num_groups=resnet_groups, - ) - ) - attentions.append( - Transformer2DModel( - out_channels // num_attention_heads, - num_attention_heads, - in_channels=out_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - ) - ) - temp_attentions.append( - TransformerTemporalModel( - out_channels // num_attention_heads, - num_attention_heads, - in_channels=out_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - ) - ) - self.resnets = nn.ModuleList(resnets) - self.temp_convs = nn.ModuleList(temp_convs) - self.attentions = nn.ModuleList(attentions) - self.temp_attentions = nn.ModuleList(temp_attentions) - - if add_upsample: - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - 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, - num_frames: int = 1, - cross_attention_kwargs: Dict[str, Any] = None, - ) -> torch.Tensor: - is_freeu_enabled = ( - getattr(self, "s1", None) - and getattr(self, "s2", None) - and getattr(self, "b1", None) - and getattr(self, "b2", None) - ) - - # TODO(Patrick, William) - attention mask is not used - for resnet, temp_conv, attn, temp_attn in zip( - self.resnets, self.temp_convs, self.attentions, self.temp_attentions - ): - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - - # FreeU: Only operate on the first two stages - if is_freeu_enabled: - hidden_states, res_hidden_states = apply_freeu( - self.resolution_idx, - hidden_states, - res_hidden_states, - s1=self.s1, - s2=self.s2, - b1=self.b1, - b2=self.b2, - ) - - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - hidden_states = resnet(hidden_states, temb) - hidden_states = temp_conv(hidden_states, num_frames=num_frames) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - hidden_states = temp_attn( - hidden_states, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) - - return hidden_states - - -class UpBlock3D(nn.Module): - def __init__( - self, - in_channels: int, - prev_output_channel: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor: float = 1.0, - add_upsample: bool = True, - resolution_idx: Optional[int] = None, - ): - super().__init__() - resnets = [] - temp_convs = [] - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - temp_convs.append( - TemporalConvLayer( - out_channels, - out_channels, - dropout=0.1, - norm_num_groups=resnet_groups, - ) - ) - - self.resnets = nn.ModuleList(resnets) - self.temp_convs = nn.ModuleList(temp_convs) - - if add_upsample: - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - def forward( - self, - hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = None, - num_frames: int = 1, - ) -> torch.Tensor: - is_freeu_enabled = ( - getattr(self, "s1", None) - and getattr(self, "s2", None) - and getattr(self, "b1", None) - and getattr(self, "b2", None) - ) - for resnet, temp_conv in zip(self.resnets, self.temp_convs): - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - - # FreeU: Only operate on the first two stages - if is_freeu_enabled: - hidden_states, res_hidden_states = apply_freeu( - self.resolution_idx, - hidden_states, - res_hidden_states, - s1=self.s1, - s2=self.s2, - b1=self.b1, - b2=self.b2, - ) - - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - hidden_states = resnet(hidden_states, temb) - hidden_states = temp_conv(hidden_states, num_frames=num_frames) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) - - return hidden_states - - -class MidBlockTemporalDecoder(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - attention_head_dim: int = 512, - num_layers: int = 1, - upcast_attention: bool = False, - ): - super().__init__() - - resnets = [] - attentions = [] - for i in range(num_layers): - input_channels = in_channels if i == 0 else out_channels - resnets.append( - SpatioTemporalResBlock( - in_channels=input_channels, - out_channels=out_channels, - temb_channels=None, - eps=1e-6, - temporal_eps=1e-5, - merge_factor=0.0, - merge_strategy="learned", - switch_spatial_to_temporal_mix=True, - ) - ) - - attentions.append( - Attention( - query_dim=in_channels, - heads=in_channels // attention_head_dim, - dim_head=attention_head_dim, - eps=1e-6, - upcast_attention=upcast_attention, - norm_num_groups=32, - bias=True, - residual_connection=True, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - def forward( - self, - hidden_states: torch.Tensor, - image_only_indicator: torch.Tensor, - ): - hidden_states = self.resnets[0]( - hidden_states, - image_only_indicator=image_only_indicator, - ) - for resnet, attn in zip(self.resnets[1:], self.attentions): - hidden_states = attn(hidden_states) - hidden_states = resnet( - hidden_states, - image_only_indicator=image_only_indicator, - ) - - return hidden_states - - -class UpBlockTemporalDecoder(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - num_layers: int = 1, - add_upsample: bool = True, - ): - super().__init__() - resnets = [] - for i in range(num_layers): - input_channels = in_channels if i == 0 else out_channels - - resnets.append( - SpatioTemporalResBlock( - in_channels=input_channels, - out_channels=out_channels, - temb_channels=None, - eps=1e-6, - temporal_eps=1e-5, - merge_factor=0.0, - merge_strategy="learned", - switch_spatial_to_temporal_mix=True, - ) - ) - self.resnets = nn.ModuleList(resnets) - - if add_upsample: - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - def forward( - self, - hidden_states: torch.Tensor, - image_only_indicator: torch.Tensor, - ) -> torch.Tensor: - for resnet in self.resnets: - hidden_states = resnet( - hidden_states, - image_only_indicator=image_only_indicator, - ) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states) - - return hidden_states - - -class UNetMidBlockSpatioTemporal(nn.Module): - def __init__( - self, - in_channels: int, - temb_channels: int, - num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - num_attention_heads: int = 1, - cross_attention_dim: int = 1280, - ): - super().__init__() - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - - # support for variable transformer layers per block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * num_layers - - # there is always at least one resnet - resnets = [ - SpatioTemporalResBlock( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=1e-5, - ) - ] - attentions = [] - - for i in range(num_layers): - attentions.append( - TransformerSpatioTemporalModel( - num_attention_heads, - in_channels // num_attention_heads, - in_channels=in_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - ) - ) - - resnets.append( - SpatioTemporalResBlock( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=1e-5, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - self.gradient_checkpointing = False - - 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, - ) -> torch.Tensor: - hidden_states = self.resnets[0]( - hidden_states, - temb, - image_only_indicator=image_only_indicator, - ) - - for attn, resnet in zip(self.attentions, self.resnets[1:]): - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - image_only_indicator=image_only_indicator, - return_dict=False, - )[0] - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb, image_only_indicator) - else: - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - image_only_indicator=image_only_indicator, - return_dict=False, - )[0] - hidden_states = resnet(hidden_states, temb, image_only_indicator=image_only_indicator) - - return hidden_states - - -class DownBlockSpatioTemporal(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - num_layers: int = 1, - add_downsample: bool = True, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - SpatioTemporalResBlock( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=1e-5, - ) - ) - - self.resnets = nn.ModuleList(resnets) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - Downsample2D( - out_channels, - use_conv=True, - out_channels=out_channels, - name="op", - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - 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, ...]]: - output_states = () - for resnet in self.resnets: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb, image_only_indicator) - else: - hidden_states = resnet(hidden_states, temb, image_only_indicator=image_only_indicator) - - output_states = output_states + (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - output_states = output_states + (hidden_states,) - - return hidden_states, output_states - - -class CrossAttnDownBlockSpatioTemporal(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - num_attention_heads: int = 1, - cross_attention_dim: int = 1280, - add_downsample: bool = True, - ): - super().__init__() - resnets = [] - attentions = [] - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * num_layers - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - SpatioTemporalResBlock( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=1e-6, - ) - ) - attentions.append( - TransformerSpatioTemporalModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - Downsample2D( - out_channels, - use_conv=True, - out_channels=out_channels, - padding=1, - name="op", - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - 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, ...]]: - output_states = () - - blocks = list(zip(self.resnets, self.attentions)) - for resnet, attn in blocks: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb, image_only_indicator) - - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - image_only_indicator=image_only_indicator, - return_dict=False, - )[0] - else: - hidden_states = resnet(hidden_states, temb, image_only_indicator=image_only_indicator) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - image_only_indicator=image_only_indicator, - return_dict=False, - )[0] - - output_states = output_states + (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - output_states = output_states + (hidden_states,) - - return hidden_states, output_states - - -class UpBlockSpatioTemporal(nn.Module): - def __init__( - self, - in_channels: int, - prev_output_channel: int, - out_channels: int, - temb_channels: int, - resolution_idx: Optional[int] = None, - num_layers: int = 1, - resnet_eps: float = 1e-6, - add_upsample: bool = True, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - SpatioTemporalResBlock( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - ) - ) - - self.resnets = nn.ModuleList(resnets) - - if add_upsample: - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - 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, - ) -> torch.Tensor: - for resnet in self.resnets: - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb, image_only_indicator) - else: - hidden_states = resnet(hidden_states, temb, image_only_indicator=image_only_indicator) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) - - return hidden_states - - -class CrossAttnUpBlockSpatioTemporal(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - prev_output_channel: int, - temb_channels: int, - resolution_idx: Optional[int] = None, - num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - resnet_eps: float = 1e-6, - num_attention_heads: int = 1, - cross_attention_dim: int = 1280, - add_upsample: bool = True, - ): - super().__init__() - resnets = [] - attentions = [] - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * num_layers - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - SpatioTemporalResBlock( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - ) - ) - attentions.append( - TransformerSpatioTemporalModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - if add_upsample: - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - 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, - ) -> torch.Tensor: - for resnet, attn in zip(self.resnets, self.attentions): - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb, image_only_indicator) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - image_only_indicator=image_only_indicator, - return_dict=False, - )[0] - else: - hidden_states = resnet(hidden_states, temb, image_only_indicator=image_only_indicator) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - image_only_indicator=image_only_indicator, - return_dict=False, - )[0] - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) - - return hidden_states diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_3d_condition.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_3d_condition.py deleted file mode 100644 index bd67ea414ab81c55dd1938237534c623f3a95664..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_3d_condition.py +++ /dev/null @@ -1,744 +0,0 @@ -# Copyright 2025 Alibaba DAMO-VILAB and The HuggingFace Team. All rights reserved. -# Copyright 2025 The ModelScope 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. - -from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union - -import torch -import torch.nn as nn -import torch.utils.checkpoint - -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import UNet2DConditionLoadersMixin -from ...utils import BaseOutput, logging -from ..activations import get_activation -from ..attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - Attention, - AttentionProcessor, - AttnAddedKVProcessor, - AttnProcessor, - FusedAttnProcessor2_0, -) -from ..embeddings import TimestepEmbedding, Timesteps -from ..modeling_utils import ModelMixin -from ..transformers.transformer_temporal import TransformerTemporalModel -from .unet_3d_blocks import ( - UNetMidBlock3DCrossAttn, - get_down_block, - get_up_block, -) - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class UNet3DConditionOutput(BaseOutput): - """ - The output of [`UNet3DConditionModel`]. - - Args: - sample (`torch.Tensor` of shape `(batch_size, num_channels, num_frames, height, width)`): - The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model. - """ - - sample: torch.Tensor - - -class UNet3DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): - r""" - A conditional 3D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample - shaped output. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - Parameters: - 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")`): - The tuple of downsample blocks to use. - 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)`): - 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. - mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block. - act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. - 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`, *optional*, defaults to 1024): The dimension of the cross attention features. - attention_head_dim (`int`, *optional*, defaults to 64): The dimension of the attention heads. - num_attention_heads (`int`, *optional*): The number of attention heads. - time_cond_proj_dim (`int`, *optional*, defaults to `None`): - The dimension of `cond_proj` layer in the timestep embedding. - """ - - _supports_gradient_checkpointing = False - _skip_layerwise_casting_patterns = ["norm", "time_embedding"] - - @register_to_config - def __init__( - self, - sample_size: Optional[int] = None, - in_channels: int = 4, - out_channels: int = 4, - down_block_types: Tuple[str, ...] = ( - "CrossAttnDownBlock3D", - "CrossAttnDownBlock3D", - "CrossAttnDownBlock3D", - "DownBlock3D", - ), - up_block_types: Tuple[str, ...] = ( - "UpBlock3D", - "CrossAttnUpBlock3D", - "CrossAttnUpBlock3D", - "CrossAttnUpBlock3D", - ), - 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 = 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, - ): - super().__init__() - - self.sample_size = sample_size - - if num_attention_heads is not None: - raise NotImplementedError( - "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19." - ) - - # If `num_attention_heads` is not defined (which is the case for most models) - # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. - # The reason for this behavior is to correct for incorrectly named variables that were introduced - # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 - # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking - # which is why we correct for the naming here. - num_attention_heads = num_attention_heads or attention_head_dim - - # Check inputs - if len(down_block_types) != len(up_block_types): - raise ValueError( - f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}." - ) - - if len(block_out_channels) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." - ) - - # input - conv_in_kernel = 3 - conv_out_kernel = 3 - conv_in_padding = (conv_in_kernel - 1) // 2 - self.conv_in = nn.Conv2d( - in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding - ) - - # time - time_embed_dim = block_out_channels[0] * 4 - self.time_proj = Timesteps(block_out_channels[0], True, 0) - timestep_input_dim = block_out_channels[0] - - self.time_embedding = TimestepEmbedding( - timestep_input_dim, - time_embed_dim, - act_fn=act_fn, - cond_proj_dim=time_cond_proj_dim, - ) - - self.transformer_in = TransformerTemporalModel( - num_attention_heads=8, - attention_head_dim=attention_head_dim, - in_channels=block_out_channels[0], - num_layers=1, - norm_num_groups=norm_num_groups, - ) - - # class embedding - self.down_blocks = nn.ModuleList([]) - self.up_blocks = nn.ModuleList([]) - - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(down_block_types) - - # down - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = get_down_block( - down_block_type, - num_layers=layers_per_block, - in_channels=input_channel, - out_channels=output_channel, - temb_channels=time_embed_dim, - add_downsample=not is_final_block, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads[i], - downsample_padding=downsample_padding, - dual_cross_attention=False, - ) - self.down_blocks.append(down_block) - - # mid - self.mid_block = UNetMidBlock3DCrossAttn( - in_channels=block_out_channels[-1], - temb_channels=time_embed_dim, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - output_scale_factor=mid_block_scale_factor, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads[-1], - resnet_groups=norm_num_groups, - dual_cross_attention=False, - ) - - # count how many layers upsample the images - self.num_upsamplers = 0 - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - reversed_num_attention_heads = list(reversed(num_attention_heads)) - - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - is_final_block = i == len(block_out_channels) - 1 - - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] - - # add upsample block for all BUT final layer - if not is_final_block: - add_upsample = True - self.num_upsamplers += 1 - else: - add_upsample = False - - up_block = get_up_block( - up_block_type, - num_layers=layers_per_block + 1, - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - temb_channels=time_embed_dim, - add_upsample=add_upsample, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - cross_attention_dim=cross_attention_dim, - num_attention_heads=reversed_num_attention_heads[i], - dual_cross_attention=False, - resolution_idx=i, - ) - self.up_blocks.append(up_block) - prev_output_channel = output_channel - - # out - if norm_num_groups is not None: - self.conv_norm_out = nn.GroupNorm( - num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps - ) - self.conv_act = get_activation("silu") - else: - self.conv_norm_out = None - self.conv_act = None - - conv_out_padding = (conv_out_kernel - 1) // 2 - self.conv_out = nn.Conv2d( - block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding - ) - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # 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: - 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. This is useful for saving some memory in exchange for a small decrease in speed. - - Args: - slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): - When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If - `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is - provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` - must be a multiple of `slice_size`. - """ - sliceable_head_dims = [] - - def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): - if hasattr(module, "set_attention_slice"): - sliceable_head_dims.append(module.sliceable_head_dim) - - for child in module.children(): - fn_recursive_retrieve_sliceable_dims(child) - - # retrieve number of attention layers - for module in self.children(): - fn_recursive_retrieve_sliceable_dims(module) - - num_sliceable_layers = len(sliceable_head_dims) - - if slice_size == "auto": - # half the attention head size is usually a good trade-off between - # speed and memory - slice_size = [dim // 2 for dim in sliceable_head_dims] - elif slice_size == "max": - # make smallest slice possible - slice_size = num_sliceable_layers * [1] - - slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size - - if len(slice_size) != len(sliceable_head_dims): - raise ValueError( - f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" - f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." - ) - - for i in range(len(slice_size)): - size = slice_size[i] - dim = sliceable_head_dims[i] - if size is not None and size > dim: - raise ValueError(f"size {size} has to be smaller or equal to {dim}.") - - # 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]): - if hasattr(module, "set_attention_slice"): - module.set_attention_slice(slice_size.pop()) - - for child in module.children(): - fn_recursive_set_attention_slice(child, slice_size) - - reversed_slice_size = list(reversed(slice_size)) - for module in self.children(): - fn_recursive_set_attention_slice(module, reversed_slice_size) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None: - """ - Sets the attention processor to use [feed forward - chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). - - Parameters: - chunk_size (`int`, *optional*): - The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually - over each tensor of dim=`dim`. - dim (`int`, *optional*, defaults to `0`): - The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch) - or dim=1 (sequence length). - """ - if dim not in [0, 1]: - raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}") - - # By default chunk size is 1 - chunk_size = chunk_size or 1 - - def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int): - if hasattr(module, "set_chunk_feed_forward"): - module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) - - for child in module.children(): - fn_recursive_feed_forward(child, chunk_size, dim) - - for module in self.children(): - fn_recursive_feed_forward(module, chunk_size, dim) - - def disable_forward_chunking(self): - def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int): - if hasattr(module, "set_chunk_feed_forward"): - module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) - - for child in module.children(): - fn_recursive_feed_forward(child, chunk_size, dim) - - for module in self.children(): - fn_recursive_feed_forward(module, None, 0) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - self.set_attn_processor(processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.enable_freeu - def enable_freeu(self, s1, s2, b1, b2): - r"""Enables the FreeU mechanism from https://huggingface.co/papers/2309.11497. - - The suffixes after the scaling factors represent the stage blocks where they are being applied. - - Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that - are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL. - - Args: - s1 (`float`): - Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to - mitigate the "oversmoothing effect" in the enhanced denoising process. - s2 (`float`): - Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to - mitigate the "oversmoothing effect" in the enhanced denoising process. - b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features. - b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features. - """ - for i, upsample_block in enumerate(self.up_blocks): - setattr(upsample_block, "s1", s1) - setattr(upsample_block, "s2", s2) - setattr(upsample_block, "b1", b1) - setattr(upsample_block, "b2", b2) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.disable_freeu - def disable_freeu(self): - """Disables the FreeU mechanism.""" - freeu_keys = {"s1", "s2", "b1", "b2"} - for i, upsample_block in enumerate(self.up_blocks): - for k in freeu_keys: - if hasattr(upsample_block, k) or getattr(upsample_block, k, None) is not None: - setattr(upsample_block, k, None) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - self.set_attn_processor(FusedAttnProcessor2_0()) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - if self.original_attn_processors is not None: - self.set_attn_processor(self.original_attn_processors) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[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, - return_dict: bool = True, - ) -> Union[UNet3DConditionOutput, Tuple[torch.Tensor]]: - r""" - The [`UNet3DConditionModel`] forward method. - - Args: - sample (`torch.Tensor`): - The noisy input tensor with the following shape `(batch, num_channels, num_frames, height, width`. - timestep (`torch.Tensor` or `float` or `int`): The number of timesteps to denoise an input. - encoder_hidden_states (`torch.Tensor`): - The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. - class_labels (`torch.Tensor`, *optional*, defaults to `None`): - Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. - timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`): - Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed - through the `self.time_embedding` layer to obtain the timestep embeddings. - attention_mask (`torch.Tensor`, *optional*, defaults to `None`): - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask - is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large - negative values to the attention scores corresponding to "discard" tokens. - cross_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). - down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*): - A tuple of tensors that if specified are added to the residuals of down unet blocks. - mid_block_additional_residual: (`torch.Tensor`, *optional*): - A tensor that if specified is added to the residual of the middle unet block. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.unets.unet_3d_condition.UNet3DConditionOutput`] instead of a plain - tuple. - cross_attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the [`AttnProcessor`]. - - Returns: - [`~models.unets.unet_3d_condition.UNet3DConditionOutput`] or `tuple`: - If `return_dict` is True, an [`~models.unets.unet_3d_condition.UNet3DConditionOutput`] is returned, - otherwise a `tuple` is returned where the first element is the sample tensor. - """ - # By default samples have to be AT least a multiple of the overall upsampling factor. - # The overall upsampling factor is equal to 2 ** (# num of upsampling layears). - # However, the upsampling interpolation output size can be forced to fit any upsampling size - # on the fly if necessary. - default_overall_up_factor = 2**self.num_upsamplers - - # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` - forward_upsample_size = False - upsample_size = None - - if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): - logger.info("Forward upsample size to force interpolation output size.") - forward_upsample_size = True - - # prepare attention_mask - if attention_mask is not None: - attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 - attention_mask = attention_mask.unsqueeze(1) - - # 1. time - 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 - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = sample.device.type == "mps" - is_npu = sample.device.type == "npu" - if isinstance(timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) - elif len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - num_frames = sample.shape[2] - timesteps = timesteps.expand(sample.shape[0]) - - t_emb = self.time_proj(timesteps) - - # timesteps does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - t_emb = t_emb.to(dtype=self.dtype) - - emb = self.time_embedding(t_emb, timestep_cond) - emb = emb.repeat_interleave(num_frames, dim=0, output_size=emb.shape[0] * num_frames) - encoder_hidden_states = encoder_hidden_states.repeat_interleave( - num_frames, dim=0, output_size=encoder_hidden_states.shape[0] * num_frames - ) - - # 2. pre-process - sample = sample.permute(0, 2, 1, 3, 4).reshape((sample.shape[0] * num_frames, -1) + sample.shape[3:]) - sample = self.conv_in(sample) - - sample = self.transformer_in( - sample, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # 3. down - down_block_res_samples = (sample,) - for downsample_block in self.down_blocks: - if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: - sample, res_samples = downsample_block( - hidden_states=sample, - temb=emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - ) - else: - sample, res_samples = downsample_block(hidden_states=sample, temb=emb, num_frames=num_frames) - - down_block_res_samples += res_samples - - if down_block_additional_residuals is not None: - new_down_block_res_samples = () - - for down_block_res_sample, down_block_additional_residual in zip( - down_block_res_samples, down_block_additional_residuals - ): - down_block_res_sample = down_block_res_sample + down_block_additional_residual - new_down_block_res_samples += (down_block_res_sample,) - - down_block_res_samples = new_down_block_res_samples - - # 4. mid - if self.mid_block is not None: - sample = self.mid_block( - sample, - emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - ) - - if mid_block_additional_residual is not None: - sample = sample + mid_block_additional_residual - - # 5. up - for i, upsample_block in enumerate(self.up_blocks): - is_final_block = i == len(self.up_blocks) - 1 - - res_samples = down_block_res_samples[-len(upsample_block.resnets) :] - down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] - - # if we have not reached the final block and need to forward the - # upsample size, we do it here - if not is_final_block and forward_upsample_size: - upsample_size = down_block_res_samples[-1].shape[2:] - - if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - encoder_hidden_states=encoder_hidden_states, - upsample_size=upsample_size, - attention_mask=attention_mask, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - ) - else: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - upsample_size=upsample_size, - num_frames=num_frames, - ) - - # 6. post-process - if self.conv_norm_out: - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - - sample = self.conv_out(sample) - - # reshape to (batch, channel, framerate, width, height) - sample = sample[None, :].reshape((-1, num_frames) + sample.shape[1:]).permute(0, 2, 1, 3, 4) - - if not return_dict: - return (sample,) - - return UNet3DConditionOutput(sample=sample) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_i2vgen_xl.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_i2vgen_xl.py deleted file mode 100644 index 8449bf894cc9b18a2d21486dc6e9f5954c1c4300..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_i2vgen_xl.py +++ /dev/null @@ -1,720 +0,0 @@ -# Copyright 2025 Alibaba DAMO-VILAB 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, Dict, Optional, Tuple, Union - -import torch -import torch.nn as nn -import torch.utils.checkpoint - -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import UNet2DConditionLoadersMixin -from ...utils import logging -from ..activations import get_activation -from ..attention import Attention, FeedForward -from ..attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - AttentionProcessor, - AttnAddedKVProcessor, - AttnProcessor, - FusedAttnProcessor2_0, -) -from ..embeddings import TimestepEmbedding, Timesteps -from ..modeling_utils import ModelMixin -from ..transformers.transformer_temporal import TransformerTemporalModel -from .unet_3d_blocks import ( - UNetMidBlock3DCrossAttn, - get_down_block, - get_up_block, -) -from .unet_3d_condition import UNet3DConditionOutput - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class I2VGenXLTransformerTemporalEncoder(nn.Module): - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - activation_fn: str = "geglu", - upcast_attention: bool = False, - ff_inner_dim: Optional[int] = None, - dropout: int = 0.0, - ): - super().__init__() - self.norm1 = nn.LayerNorm(dim, elementwise_affine=True, eps=1e-5) - self.attn1 = Attention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=False, - upcast_attention=upcast_attention, - out_bias=True, - ) - self.ff = FeedForward( - dim, - dropout=dropout, - activation_fn=activation_fn, - final_dropout=False, - inner_dim=ff_inner_dim, - bias=True, - ) - - def forward( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor: - norm_hidden_states = self.norm1(hidden_states) - attn_output = self.attn1(norm_hidden_states, encoder_hidden_states=None) - hidden_states = attn_output + hidden_states - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - ff_output = self.ff(hidden_states) - hidden_states = ff_output + hidden_states - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - return hidden_states - - -class I2VGenXLUNet(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): - r""" - I2VGenXL UNet. It is a conditional 3D UNet model that takes a noisy sample, conditional state, and a timestep and - returns a sample-shaped output. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - Parameters: - 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")`): - The tuple of downsample blocks to use. - 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)`): - 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. - If `None`, normalization and activation layers is skipped in post-processing. - cross_attention_dim (`int`, *optional*, defaults to 1280): The dimension of the cross attention features. - attention_head_dim (`int`, *optional*, defaults to 64): Attention head dim. - num_attention_heads (`int`, *optional*): The number of attention heads. - """ - - _supports_gradient_checkpointing = False - - @register_to_config - def __init__( - self, - sample_size: Optional[int] = None, - in_channels: int = 4, - out_channels: int = 4, - down_block_types: Tuple[str, ...] = ( - "CrossAttnDownBlock3D", - "CrossAttnDownBlock3D", - "CrossAttnDownBlock3D", - "DownBlock3D", - ), - up_block_types: Tuple[str, ...] = ( - "UpBlock3D", - "CrossAttnUpBlock3D", - "CrossAttnUpBlock3D", - "CrossAttnUpBlock3D", - ), - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), - layers_per_block: int = 2, - norm_num_groups: Optional[int] = 32, - cross_attention_dim: int = 1024, - attention_head_dim: Union[int, Tuple[int]] = 64, - num_attention_heads: Optional[Union[int, Tuple[int]]] = None, - ): - super().__init__() - - # When we first integrated the UNet into the library, we didn't have `attention_head_dim`. As a consequence - # of that, we used `num_attention_heads` for arguments that actually denote attention head dimension. This - # is why we ignore `num_attention_heads` and calculate it from `attention_head_dims` below. - # This is still an incorrect way of calculating `num_attention_heads` but we need to stick to it - # without running proper deprecation cycles for the {down,mid,up} blocks which are a - # part of the public API. - num_attention_heads = attention_head_dim - - # Check inputs - if len(down_block_types) != len(up_block_types): - raise ValueError( - f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}." - ) - - if len(block_out_channels) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." - ) - - # input - self.conv_in = nn.Conv2d(in_channels + in_channels, block_out_channels[0], kernel_size=3, padding=1) - - self.transformer_in = TransformerTemporalModel( - num_attention_heads=8, - attention_head_dim=num_attention_heads, - in_channels=block_out_channels[0], - num_layers=1, - norm_num_groups=norm_num_groups, - ) - - # image embedding - self.image_latents_proj_in = nn.Sequential( - nn.Conv2d(4, in_channels * 4, 3, padding=1), - nn.SiLU(), - nn.Conv2d(in_channels * 4, in_channels * 4, 3, stride=1, padding=1), - nn.SiLU(), - nn.Conv2d(in_channels * 4, in_channels, 3, stride=1, padding=1), - ) - self.image_latents_temporal_encoder = I2VGenXLTransformerTemporalEncoder( - dim=in_channels, - num_attention_heads=2, - ff_inner_dim=in_channels * 4, - attention_head_dim=in_channels, - activation_fn="gelu", - ) - self.image_latents_context_embedding = nn.Sequential( - nn.Conv2d(4, in_channels * 8, 3, padding=1), - nn.SiLU(), - nn.AdaptiveAvgPool2d((32, 32)), - nn.Conv2d(in_channels * 8, in_channels * 16, 3, stride=2, padding=1), - nn.SiLU(), - nn.Conv2d(in_channels * 16, cross_attention_dim, 3, stride=2, padding=1), - ) - - # other embeddings -- time, context, fps, etc. - time_embed_dim = block_out_channels[0] * 4 - self.time_proj = Timesteps(block_out_channels[0], True, 0) - timestep_input_dim = block_out_channels[0] - - self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn="silu") - self.context_embedding = nn.Sequential( - nn.Linear(cross_attention_dim, time_embed_dim), - nn.SiLU(), - nn.Linear(time_embed_dim, cross_attention_dim * in_channels), - ) - self.fps_embedding = nn.Sequential( - nn.Linear(timestep_input_dim, time_embed_dim), nn.SiLU(), nn.Linear(time_embed_dim, time_embed_dim) - ) - - # blocks - self.down_blocks = nn.ModuleList([]) - self.up_blocks = nn.ModuleList([]) - - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(down_block_types) - - # down - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = get_down_block( - down_block_type, - num_layers=layers_per_block, - in_channels=input_channel, - out_channels=output_channel, - temb_channels=time_embed_dim, - add_downsample=not is_final_block, - resnet_eps=1e-05, - resnet_act_fn="silu", - resnet_groups=norm_num_groups, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads[i], - downsample_padding=1, - dual_cross_attention=False, - ) - self.down_blocks.append(down_block) - - # mid - self.mid_block = UNetMidBlock3DCrossAttn( - in_channels=block_out_channels[-1], - temb_channels=time_embed_dim, - resnet_eps=1e-05, - resnet_act_fn="silu", - output_scale_factor=1, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads[-1], - resnet_groups=norm_num_groups, - dual_cross_attention=False, - ) - - # count how many layers upsample the images - self.num_upsamplers = 0 - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - reversed_num_attention_heads = list(reversed(num_attention_heads)) - - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - is_final_block = i == len(block_out_channels) - 1 - - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] - - # add upsample block for all BUT final layer - if not is_final_block: - add_upsample = True - self.num_upsamplers += 1 - else: - add_upsample = False - - up_block = get_up_block( - up_block_type, - num_layers=layers_per_block + 1, - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - temb_channels=time_embed_dim, - add_upsample=add_upsample, - resnet_eps=1e-05, - resnet_act_fn="silu", - resnet_groups=norm_num_groups, - cross_attention_dim=cross_attention_dim, - num_attention_heads=reversed_num_attention_heads[i], - dual_cross_attention=False, - resolution_idx=i, - ) - self.up_blocks.append(up_block) - prev_output_channel = output_channel - - # out - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-05) - self.conv_act = get_activation("silu") - self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=3, padding=1) - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, 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: - """ - Sets the attention processor to use [feed forward - chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). - - Parameters: - chunk_size (`int`, *optional*): - The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually - over each tensor of dim=`dim`. - dim (`int`, *optional*, defaults to `0`): - The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch) - or dim=1 (sequence length). - """ - if dim not in [0, 1]: - raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}") - - # By default chunk size is 1 - chunk_size = chunk_size or 1 - - def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int): - if hasattr(module, "set_chunk_feed_forward"): - module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) - - for child in module.children(): - fn_recursive_feed_forward(child, chunk_size, dim) - - for module in self.children(): - fn_recursive_feed_forward(module, chunk_size, dim) - - # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.disable_forward_chunking - def disable_forward_chunking(self): - def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int): - if hasattr(module, "set_chunk_feed_forward"): - module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) - - for child in module.children(): - fn_recursive_feed_forward(child, chunk_size, dim) - - for module in self.children(): - fn_recursive_feed_forward(module, None, 0) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - self.set_attn_processor(processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.enable_freeu - def enable_freeu(self, s1, s2, b1, b2): - r"""Enables the FreeU mechanism from https://huggingface.co/papers/2309.11497. - - The suffixes after the scaling factors represent the stage blocks where they are being applied. - - Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that - are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL. - - Args: - s1 (`float`): - Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to - mitigate the "oversmoothing effect" in the enhanced denoising process. - s2 (`float`): - Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to - mitigate the "oversmoothing effect" in the enhanced denoising process. - b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features. - b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features. - """ - for i, upsample_block in enumerate(self.up_blocks): - setattr(upsample_block, "s1", s1) - setattr(upsample_block, "s2", s2) - setattr(upsample_block, "b1", b1) - setattr(upsample_block, "b2", b2) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.disable_freeu - def disable_freeu(self): - """Disables the FreeU mechanism.""" - freeu_keys = {"s1", "s2", "b1", "b2"} - for i, upsample_block in enumerate(self.up_blocks): - for k in freeu_keys: - if hasattr(upsample_block, k) or getattr(upsample_block, k, None) is not None: - setattr(upsample_block, k, None) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - self.set_attn_processor(FusedAttnProcessor2_0()) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - if self.original_attn_processors is not None: - self.set_attn_processor(self.original_attn_processors) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[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, - return_dict: bool = True, - ) -> Union[UNet3DConditionOutput, Tuple[torch.Tensor]]: - r""" - The [`I2VGenXLUNet`] forward method. - - Args: - sample (`torch.Tensor`): - The noisy input tensor with the following shape `(batch, num_frames, channel, height, width`. - timestep (`torch.Tensor` or `float` or `int`): The number of timesteps to denoise an input. - fps (`torch.Tensor`): Frames per second for the video being generated. Used as a "micro-condition". - image_latents (`torch.Tensor`): Image encodings from the VAE. - image_embeddings (`torch.Tensor`): - Projection embeddings of the conditioning image computed with a vision encoder. - encoder_hidden_states (`torch.Tensor`): - The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. - cross_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.unets.unet_3d_condition.UNet3DConditionOutput`] instead of a plain - tuple. - - Returns: - [`~models.unets.unet_3d_condition.UNet3DConditionOutput`] or `tuple`: - If `return_dict` is True, an [`~models.unets.unet_3d_condition.UNet3DConditionOutput`] is returned, - otherwise a `tuple` is returned where the first element is the sample tensor. - """ - batch_size, channels, num_frames, height, width = sample.shape - - # By default samples have to be AT least a multiple of the overall upsampling factor. - # The overall upsampling factor is equal to 2 ** (# num of upsampling layears). - # However, the upsampling interpolation output size can be forced to fit any upsampling size - # on the fly if necessary. - default_overall_up_factor = 2**self.num_upsamplers - - # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` - forward_upsample_size = False - upsample_size = None - - if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): - logger.info("Forward upsample size to force interpolation output size.") - forward_upsample_size = True - - # 1. time - 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 - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = sample.device.type == "mps" - is_npu = sample.device.type == "npu" - if isinstance(timesteps, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) - elif len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps.expand(sample.shape[0]) - t_emb = self.time_proj(timesteps) - - # timesteps does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - t_emb = t_emb.to(dtype=self.dtype) - t_emb = self.time_embedding(t_emb, timestep_cond) - - # 2. FPS - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - fps = fps.expand(fps.shape[0]) - fps_emb = self.fps_embedding(self.time_proj(fps).to(dtype=self.dtype)) - - # 3. time + FPS embeddings. - emb = t_emb + fps_emb - emb = emb.repeat_interleave(num_frames, dim=0, output_size=emb.shape[0] * num_frames) - - # 4. context embeddings. - # The context embeddings consist of both text embeddings from the input prompt - # AND the image embeddings from the input image. For images, both VAE encodings - # and the CLIP image embeddings are incorporated. - # So the final `context_embeddings` becomes the query for cross-attention. - context_emb = sample.new_zeros(batch_size, 0, self.config.cross_attention_dim) - context_emb = torch.cat([context_emb, encoder_hidden_states], dim=1) - - image_latents_for_context_embds = image_latents[:, :, :1, :] - image_latents_context_embs = image_latents_for_context_embds.permute(0, 2, 1, 3, 4).reshape( - image_latents_for_context_embds.shape[0] * image_latents_for_context_embds.shape[2], - image_latents_for_context_embds.shape[1], - image_latents_for_context_embds.shape[3], - image_latents_for_context_embds.shape[4], - ) - image_latents_context_embs = self.image_latents_context_embedding(image_latents_context_embs) - - _batch_size, _channels, _height, _width = image_latents_context_embs.shape - image_latents_context_embs = image_latents_context_embs.permute(0, 2, 3, 1).reshape( - _batch_size, _height * _width, _channels - ) - context_emb = torch.cat([context_emb, image_latents_context_embs], dim=1) - - image_emb = self.context_embedding(image_embeddings) - image_emb = image_emb.view(-1, self.config.in_channels, self.config.cross_attention_dim) - context_emb = torch.cat([context_emb, image_emb], dim=1) - context_emb = context_emb.repeat_interleave(num_frames, dim=0, output_size=context_emb.shape[0] * num_frames) - - image_latents = image_latents.permute(0, 2, 1, 3, 4).reshape( - image_latents.shape[0] * image_latents.shape[2], - image_latents.shape[1], - image_latents.shape[3], - image_latents.shape[4], - ) - image_latents = self.image_latents_proj_in(image_latents) - image_latents = ( - image_latents[None, :] - .reshape(batch_size, num_frames, channels, height, width) - .permute(0, 3, 4, 1, 2) - .reshape(batch_size * height * width, num_frames, channels) - ) - image_latents = self.image_latents_temporal_encoder(image_latents) - image_latents = image_latents.reshape(batch_size, height, width, num_frames, channels).permute(0, 4, 3, 1, 2) - - # 5. pre-process - sample = torch.cat([sample, image_latents], dim=1) - sample = sample.permute(0, 2, 1, 3, 4).reshape((sample.shape[0] * num_frames, -1) + sample.shape[3:]) - sample = self.conv_in(sample) - sample = self.transformer_in( - sample, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # 6. down - down_block_res_samples = (sample,) - for downsample_block in self.down_blocks: - if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: - sample, res_samples = downsample_block( - hidden_states=sample, - temb=emb, - encoder_hidden_states=context_emb, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - ) - else: - sample, res_samples = downsample_block(hidden_states=sample, temb=emb, num_frames=num_frames) - - down_block_res_samples += res_samples - - # 7. mid - if self.mid_block is not None: - sample = self.mid_block( - sample, - emb, - encoder_hidden_states=context_emb, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - ) - # 8. up - for i, upsample_block in enumerate(self.up_blocks): - is_final_block = i == len(self.up_blocks) - 1 - - res_samples = down_block_res_samples[-len(upsample_block.resnets) :] - down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] - - # if we have not reached the final block and need to forward the - # upsample size, we do it here - if not is_final_block and forward_upsample_size: - upsample_size = down_block_res_samples[-1].shape[2:] - - if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - encoder_hidden_states=context_emb, - upsample_size=upsample_size, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - ) - else: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - upsample_size=upsample_size, - num_frames=num_frames, - ) - - # 9. post-process - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - - sample = self.conv_out(sample) - - # reshape to (batch, channel, framerate, width, height) - sample = sample[None, :].reshape((-1, num_frames) + sample.shape[1:]).permute(0, 2, 1, 3, 4) - - if not return_dict: - return (sample,) - - return UNet3DConditionOutput(sample=sample) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_kandinsky3.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_kandinsky3.py deleted file mode 100644 index 423669a22fa2be3ae4f1c4ea899eb4b1d9bcfcb3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_kandinsky3.py +++ /dev/null @@ -1,531 +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 dataclasses import dataclass -from typing import Dict, Tuple, Union - -import torch -import torch.utils.checkpoint -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import BaseOutput, logging -from ..attention_processor import Attention, AttentionProcessor, AttnProcessor -from ..embeddings import TimestepEmbedding, Timesteps -from ..modeling_utils import ModelMixin - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class Kandinsky3UNetOutput(BaseOutput): - sample: torch.Tensor = None - - -class Kandinsky3EncoderProj(nn.Module): - def __init__(self, encoder_hid_dim, cross_attention_dim): - super().__init__() - self.projection_linear = nn.Linear(encoder_hid_dim, cross_attention_dim, bias=False) - self.projection_norm = nn.LayerNorm(cross_attention_dim) - - def forward(self, x): - x = self.projection_linear(x) - x = self.projection_norm(x) - return x - - -class Kandinsky3UNet(ModelMixin, ConfigMixin): - @register_to_config - def __init__( - self, - in_channels: int = 4, - 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, - encoder_hid_dim: int = 4096, - ): - super().__init__() - - # TODO(Yiyi): Give better name and put into config for the following 4 parameters - expansion_ratio = 4 - compression_ratio = 2 - add_cross_attention = (False, True, True, True) - add_self_attention = (False, True, True, True) - - out_channels = in_channels - init_channels = block_out_channels[0] // 2 - self.time_proj = Timesteps(init_channels, flip_sin_to_cos=False, downscale_freq_shift=1) - - self.time_embedding = TimestepEmbedding( - init_channels, - time_embedding_dim, - ) - - self.add_time_condition = Kandinsky3AttentionPooling( - time_embedding_dim, cross_attention_dim, attention_head_dim - ) - - self.conv_in = nn.Conv2d(in_channels, init_channels, kernel_size=3, padding=1) - - self.encoder_hid_proj = Kandinsky3EncoderProj(encoder_hid_dim, cross_attention_dim) - - hidden_dims = [init_channels] + list(block_out_channels) - in_out_dims = list(zip(hidden_dims[:-1], hidden_dims[1:])) - text_dims = [cross_attention_dim if is_exist else None for is_exist in add_cross_attention] - num_blocks = len(block_out_channels) * [layers_per_block] - layer_params = [num_blocks, text_dims, add_self_attention] - rev_layer_params = map(reversed, layer_params) - - cat_dims = [] - self.num_levels = len(in_out_dims) - self.down_blocks = nn.ModuleList([]) - for level, ((in_dim, out_dim), res_block_num, text_dim, self_attention) in enumerate( - zip(in_out_dims, *layer_params) - ): - down_sample = level != (self.num_levels - 1) - cat_dims.append(out_dim if level != (self.num_levels - 1) else 0) - self.down_blocks.append( - Kandinsky3DownSampleBlock( - in_dim, - out_dim, - time_embedding_dim, - text_dim, - res_block_num, - groups, - attention_head_dim, - expansion_ratio, - compression_ratio, - down_sample, - self_attention, - ) - ) - - self.up_blocks = nn.ModuleList([]) - for level, ((out_dim, in_dim), res_block_num, text_dim, self_attention) in enumerate( - zip(reversed(in_out_dims), *rev_layer_params) - ): - up_sample = level != 0 - self.up_blocks.append( - Kandinsky3UpSampleBlock( - in_dim, - cat_dims.pop(), - out_dim, - time_embedding_dim, - text_dim, - res_block_num, - groups, - attention_head_dim, - expansion_ratio, - compression_ratio, - up_sample, - self_attention, - ) - ) - - self.conv_norm_out = nn.GroupNorm(groups, init_channels) - self.conv_act_out = nn.SiLU() - self.conv_out = nn.Conv2d(init_channels, out_channels, kernel_size=3, padding=1) - - @property - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]): - if hasattr(module, "set_processor"): - processors[f"{name}.processor"] = module.processor - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - self.set_attn_processor(AttnProcessor()) - - def forward(self, sample, timestep, encoder_hidden_states=None, encoder_attention_mask=None, return_dict=True): - if encoder_attention_mask is not None: - encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0 - encoder_attention_mask = encoder_attention_mask.unsqueeze(1) - - if not torch.is_tensor(timestep): - dtype = torch.float32 if isinstance(timestep, float) else torch.int32 - timestep = torch.tensor([timestep], dtype=dtype, device=sample.device) - elif len(timestep.shape) == 0: - timestep = timestep[None].to(sample.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = timestep.expand(sample.shape[0]) - time_embed_input = self.time_proj(timestep).to(sample.dtype) - time_embed = self.time_embedding(time_embed_input) - - encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states) - - if encoder_hidden_states is not None: - time_embed = self.add_time_condition(time_embed, encoder_hidden_states, encoder_attention_mask) - - hidden_states = [] - sample = self.conv_in(sample) - for level, down_sample in enumerate(self.down_blocks): - sample = down_sample(sample, time_embed, encoder_hidden_states, encoder_attention_mask) - if level != self.num_levels - 1: - hidden_states.append(sample) - - for level, up_sample in enumerate(self.up_blocks): - if level != 0: - sample = torch.cat([sample, hidden_states.pop()], dim=1) - sample = up_sample(sample, time_embed, encoder_hidden_states, encoder_attention_mask) - - sample = self.conv_norm_out(sample) - sample = self.conv_act_out(sample) - sample = self.conv_out(sample) - - if not return_dict: - return (sample,) - return Kandinsky3UNetOutput(sample=sample) - - -class Kandinsky3UpSampleBlock(nn.Module): - def __init__( - self, - in_channels, - cat_dim, - out_channels, - time_embed_dim, - context_dim=None, - num_blocks=3, - groups=32, - head_dim=64, - expansion_ratio=4, - compression_ratio=2, - up_sample=True, - self_attention=True, - ): - super().__init__() - up_resolutions = [[None, True if up_sample else None, None, None]] + [[None] * 4] * (num_blocks - 1) - hidden_channels = ( - [(in_channels + cat_dim, in_channels)] - + [(in_channels, in_channels)] * (num_blocks - 2) - + [(in_channels, out_channels)] - ) - attentions = [] - resnets_in = [] - resnets_out = [] - - self.self_attention = self_attention - self.context_dim = context_dim - - if self_attention: - attentions.append( - Kandinsky3AttentionBlock(out_channels, time_embed_dim, None, groups, head_dim, expansion_ratio) - ) - else: - attentions.append(nn.Identity()) - - for (in_channel, out_channel), up_resolution in zip(hidden_channels, up_resolutions): - resnets_in.append( - Kandinsky3ResNetBlock(in_channel, in_channel, time_embed_dim, groups, compression_ratio, up_resolution) - ) - - if context_dim is not None: - attentions.append( - Kandinsky3AttentionBlock( - in_channel, time_embed_dim, context_dim, groups, head_dim, expansion_ratio - ) - ) - else: - attentions.append(nn.Identity()) - - resnets_out.append( - Kandinsky3ResNetBlock(in_channel, out_channel, time_embed_dim, groups, compression_ratio) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets_in = nn.ModuleList(resnets_in) - self.resnets_out = nn.ModuleList(resnets_out) - - def forward(self, x, time_embed, context=None, context_mask=None, image_mask=None): - for attention, resnet_in, resnet_out in zip(self.attentions[1:], self.resnets_in, self.resnets_out): - x = resnet_in(x, time_embed) - if self.context_dim is not None: - x = attention(x, time_embed, context, context_mask, image_mask) - x = resnet_out(x, time_embed) - - if self.self_attention: - x = self.attentions[0](x, time_embed, image_mask=image_mask) - return x - - -class Kandinsky3DownSampleBlock(nn.Module): - def __init__( - self, - in_channels, - out_channels, - time_embed_dim, - context_dim=None, - num_blocks=3, - groups=32, - head_dim=64, - expansion_ratio=4, - compression_ratio=2, - down_sample=True, - self_attention=True, - ): - super().__init__() - attentions = [] - resnets_in = [] - resnets_out = [] - - self.self_attention = self_attention - self.context_dim = context_dim - - if self_attention: - attentions.append( - Kandinsky3AttentionBlock(in_channels, time_embed_dim, None, groups, head_dim, expansion_ratio) - ) - else: - attentions.append(nn.Identity()) - - up_resolutions = [[None] * 4] * (num_blocks - 1) + [[None, None, False if down_sample else None, None]] - hidden_channels = [(in_channels, out_channels)] + [(out_channels, out_channels)] * (num_blocks - 1) - for (in_channel, out_channel), up_resolution in zip(hidden_channels, up_resolutions): - resnets_in.append( - Kandinsky3ResNetBlock(in_channel, out_channel, time_embed_dim, groups, compression_ratio) - ) - - if context_dim is not None: - attentions.append( - Kandinsky3AttentionBlock( - out_channel, time_embed_dim, context_dim, groups, head_dim, expansion_ratio - ) - ) - else: - attentions.append(nn.Identity()) - - resnets_out.append( - Kandinsky3ResNetBlock( - out_channel, out_channel, time_embed_dim, groups, compression_ratio, up_resolution - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets_in = nn.ModuleList(resnets_in) - self.resnets_out = nn.ModuleList(resnets_out) - - def forward(self, x, time_embed, context=None, context_mask=None, image_mask=None): - if self.self_attention: - x = self.attentions[0](x, time_embed, image_mask=image_mask) - - for attention, resnet_in, resnet_out in zip(self.attentions[1:], self.resnets_in, self.resnets_out): - x = resnet_in(x, time_embed) - if self.context_dim is not None: - x = attention(x, time_embed, context, context_mask, image_mask) - x = resnet_out(x, time_embed) - return x - - -class Kandinsky3ConditionalGroupNorm(nn.Module): - def __init__(self, groups, normalized_shape, context_dim): - super().__init__() - self.norm = nn.GroupNorm(groups, normalized_shape, affine=False) - self.context_mlp = nn.Sequential(nn.SiLU(), nn.Linear(context_dim, 2 * normalized_shape)) - self.context_mlp[1].weight.data.zero_() - self.context_mlp[1].bias.data.zero_() - - def forward(self, x, context): - context = self.context_mlp(context) - - for _ in range(len(x.shape[2:])): - context = context.unsqueeze(-1) - - scale, shift = context.chunk(2, dim=1) - x = self.norm(x) * (scale + 1.0) + shift - return x - - -class Kandinsky3Block(nn.Module): - def __init__(self, in_channels, out_channels, time_embed_dim, kernel_size=3, norm_groups=32, up_resolution=None): - super().__init__() - self.group_norm = Kandinsky3ConditionalGroupNorm(norm_groups, in_channels, time_embed_dim) - self.activation = nn.SiLU() - if up_resolution is not None and up_resolution: - self.up_sample = nn.ConvTranspose2d(in_channels, in_channels, kernel_size=2, stride=2) - else: - self.up_sample = nn.Identity() - - padding = int(kernel_size > 1) - self.projection = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=padding) - - if up_resolution is not None and not up_resolution: - self.down_sample = nn.Conv2d(out_channels, out_channels, kernel_size=2, stride=2) - else: - self.down_sample = nn.Identity() - - def forward(self, x, time_embed): - x = self.group_norm(x, time_embed) - x = self.activation(x) - x = self.up_sample(x) - x = self.projection(x) - x = self.down_sample(x) - return x - - -class Kandinsky3ResNetBlock(nn.Module): - def __init__( - self, in_channels, out_channels, time_embed_dim, norm_groups=32, compression_ratio=2, up_resolutions=4 * [None] - ): - super().__init__() - kernel_sizes = [1, 3, 3, 1] - hidden_channel = max(in_channels, out_channels) // compression_ratio - hidden_channels = ( - [(in_channels, hidden_channel)] + [(hidden_channel, hidden_channel)] * 2 + [(hidden_channel, out_channels)] - ) - self.resnet_blocks = nn.ModuleList( - [ - Kandinsky3Block(in_channel, out_channel, time_embed_dim, kernel_size, norm_groups, up_resolution) - for (in_channel, out_channel), kernel_size, up_resolution in zip( - hidden_channels, kernel_sizes, up_resolutions - ) - ] - ) - self.shortcut_up_sample = ( - nn.ConvTranspose2d(in_channels, in_channels, kernel_size=2, stride=2) - if True in up_resolutions - else nn.Identity() - ) - self.shortcut_projection = ( - nn.Conv2d(in_channels, out_channels, kernel_size=1) if in_channels != out_channels else nn.Identity() - ) - self.shortcut_down_sample = ( - nn.Conv2d(out_channels, out_channels, kernel_size=2, stride=2) - if False in up_resolutions - else nn.Identity() - ) - - def forward(self, x, time_embed): - out = x - for resnet_block in self.resnet_blocks: - out = resnet_block(out, time_embed) - - x = self.shortcut_up_sample(x) - x = self.shortcut_projection(x) - x = self.shortcut_down_sample(x) - x = x + out - return x - - -class Kandinsky3AttentionPooling(nn.Module): - def __init__(self, num_channels, context_dim, head_dim=64): - super().__init__() - self.attention = Attention( - context_dim, - context_dim, - dim_head=head_dim, - out_dim=num_channels, - out_bias=False, - ) - - def forward(self, x, context, context_mask=None): - context_mask = context_mask.to(dtype=context.dtype) - context = self.attention(context.mean(dim=1, keepdim=True), context, context_mask) - return x + context.squeeze(1) - - -class Kandinsky3AttentionBlock(nn.Module): - def __init__(self, num_channels, time_embed_dim, context_dim=None, norm_groups=32, head_dim=64, expansion_ratio=4): - super().__init__() - self.in_norm = Kandinsky3ConditionalGroupNorm(norm_groups, num_channels, time_embed_dim) - self.attention = Attention( - num_channels, - context_dim or num_channels, - dim_head=head_dim, - out_dim=num_channels, - out_bias=False, - ) - - hidden_channels = expansion_ratio * num_channels - self.out_norm = Kandinsky3ConditionalGroupNorm(norm_groups, num_channels, time_embed_dim) - self.feed_forward = nn.Sequential( - nn.Conv2d(num_channels, hidden_channels, kernel_size=1, bias=False), - nn.SiLU(), - nn.Conv2d(hidden_channels, num_channels, kernel_size=1, bias=False), - ) - - def forward(self, x, time_embed, context=None, context_mask=None, image_mask=None): - height, width = x.shape[-2:] - out = self.in_norm(x, time_embed) - out = out.reshape(x.shape[0], -1, height * width).permute(0, 2, 1) - context = context if context is not None else out - if context_mask is not None: - context_mask = context_mask.to(dtype=context.dtype) - - out = self.attention(out, context, context_mask) - out = out.permute(0, 2, 1).unsqueeze(-1).reshape(out.shape[0], -1, height, width) - x = x + out - - out = self.out_norm(x, time_embed) - out = self.feed_forward(out) - x = x + out - return x diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_motion_model.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_motion_model.py deleted file mode 100644 index 0a112b524911af26d711f2290dbef68aa5278366..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_motion_model.py +++ /dev/null @@ -1,2178 +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 dataclasses import dataclass -from typing import Any, Dict, Optional, Tuple, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch.utils.checkpoint - -from ...configuration_utils import ConfigMixin, FrozenDict, register_to_config -from ...loaders import FromOriginalModelMixin, PeftAdapterMixin, UNet2DConditionLoadersMixin -from ...utils import BaseOutput, deprecate, logging -from ...utils.torch_utils import apply_freeu -from ..attention import BasicTransformerBlock -from ..attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - Attention, - AttentionProcessor, - AttnAddedKVProcessor, - AttnProcessor, - AttnProcessor2_0, - FusedAttnProcessor2_0, - IPAdapterAttnProcessor, - IPAdapterAttnProcessor2_0, -) -from ..embeddings import TimestepEmbedding, Timesteps -from ..modeling_utils import ModelMixin -from ..resnet import Downsample2D, ResnetBlock2D, Upsample2D -from ..transformers.dual_transformer_2d import DualTransformer2DModel -from ..transformers.transformer_2d import Transformer2DModel -from .unet_2d_blocks import UNetMidBlock2DCrossAttn -from .unet_2d_condition import UNet2DConditionModel - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class UNetMotionOutput(BaseOutput): - """ - The output of [`UNetMotionOutput`]. - - Args: - sample (`torch.Tensor` of shape `(batch_size, num_channels, num_frames, height, width)`): - The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model. - """ - - sample: torch.Tensor - - -class AnimateDiffTransformer3D(nn.Module): - """ - A Transformer model for video-like data. - - Parameters: - num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. - attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. - in_channels (`int`, *optional*): - The number of channels in the input and output (specify if the input is **continuous**). - num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use. - attention_bias (`bool`, *optional*): - Configure if the `TransformerBlock` attention should contain a bias parameter. - sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**). - This is fixed during training since it is used to learn a number of position embeddings. - activation_fn (`str`, *optional*, defaults to `"geglu"`): - Activation function to use in feed-forward. See `diffusers.models.activations.get_activation` for supported - activation functions. - norm_elementwise_affine (`bool`, *optional*): - Configure if the `TransformerBlock` should use learnable elementwise affine parameters for normalization. - double_self_attention (`bool`, *optional*): - Configure if each `TransformerBlock` should contain two self-attention layers. - positional_embeddings: (`str`, *optional*): - The type of positional embeddings to apply to the sequence input before passing use. - num_positional_embeddings: (`int`, *optional*): - The maximum length of the sequence over which to apply positional embeddings. - """ - - def __init__( - self, - num_attention_heads: int = 16, - attention_head_dim: int = 88, - in_channels: Optional[int] = None, - out_channels: Optional[int] = None, - num_layers: int = 1, - dropout: float = 0.0, - norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = None, - attention_bias: bool = False, - sample_size: Optional[int] = 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, - ): - super().__init__() - self.num_attention_heads = num_attention_heads - self.attention_head_dim = attention_head_dim - inner_dim = num_attention_heads * attention_head_dim - - self.in_channels = in_channels - - self.norm = nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True) - self.proj_in = nn.Linear(in_channels, inner_dim) - - # 3. Define transformers blocks - self.transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - inner_dim, - num_attention_heads, - attention_head_dim, - dropout=dropout, - cross_attention_dim=cross_attention_dim, - activation_fn=activation_fn, - attention_bias=attention_bias, - double_self_attention=double_self_attention, - norm_elementwise_affine=norm_elementwise_affine, - positional_embeddings=positional_embeddings, - num_positional_embeddings=num_positional_embeddings, - ) - for _ in range(num_layers) - ] - ) - - self.proj_out = nn.Linear(inner_dim, in_channels) - - 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, - num_frames: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> torch.Tensor: - """ - The [`AnimateDiffTransformer3D`] forward method. - - Args: - hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.Tensor` of shape `(batch size, channel, height, width)` if continuous): - Input hidden_states. - encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*): - Conditional embeddings for cross attention layer. If not given, cross-attention defaults to - self-attention. - timestep ( `torch.LongTensor`, *optional*): - Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`. - 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`. - num_frames (`int`, *optional*, defaults to 1): - The number of frames to be processed per batch. This is used to reshape the hidden states. - cross_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). - - Returns: - torch.Tensor: - The output tensor. - """ - # 1. Input - batch_frames, channel, height, width = hidden_states.shape - batch_size = batch_frames // num_frames - - residual = hidden_states - - hidden_states = hidden_states[None, :].reshape(batch_size, num_frames, channel, height, width) - hidden_states = hidden_states.permute(0, 2, 1, 3, 4) - - hidden_states = self.norm(hidden_states) - hidden_states = hidden_states.permute(0, 3, 4, 2, 1).reshape(batch_size * height * width, num_frames, channel) - - hidden_states = self.proj_in(input=hidden_states) - - # 2. Blocks - for block in self.transformer_blocks: - hidden_states = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=timestep, - cross_attention_kwargs=cross_attention_kwargs, - class_labels=class_labels, - ) - - # 3. Output - hidden_states = self.proj_out(input=hidden_states) - hidden_states = ( - hidden_states[None, None, :] - .reshape(batch_size, height, width, num_frames, channel) - .permute(0, 3, 4, 1, 2) - .contiguous() - ) - hidden_states = hidden_states.reshape(batch_frames, channel, height, width) - - output = hidden_states + residual - return output - - -class DownBlockMotion(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - 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_max_seq_length: int = 32, - temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1, - temporal_double_self_attention: bool = True, - ): - super().__init__() - resnets = [] - motion_modules = [] - - # support for variable transformer layers per temporal block - if isinstance(temporal_transformer_layers_per_block, int): - temporal_transformer_layers_per_block = (temporal_transformer_layers_per_block,) * num_layers - elif len(temporal_transformer_layers_per_block) != num_layers: - raise ValueError( - f"`temporal_transformer_layers_per_block` must be an integer or a tuple of integers of length {num_layers}" - ) - - # support for variable number of attention head per temporal layers - if isinstance(temporal_num_attention_heads, int): - temporal_num_attention_heads = (temporal_num_attention_heads,) * num_layers - elif len(temporal_num_attention_heads) != num_layers: - raise ValueError( - f"`temporal_num_attention_heads` must be an integer or a tuple of integers of length {num_layers}" - ) - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - motion_modules.append( - AnimateDiffTransformer3D( - num_attention_heads=temporal_num_attention_heads[i], - in_channels=out_channels, - num_layers=temporal_transformer_layers_per_block[i], - norm_num_groups=resnet_groups, - cross_attention_dim=temporal_cross_attention_dim, - attention_bias=False, - activation_fn="geglu", - positional_embeddings="sinusoidal", - num_positional_embeddings=temporal_max_seq_length, - attention_head_dim=out_channels // temporal_num_attention_heads[i], - double_self_attention=temporal_double_self_attention, - ) - ) - - self.resnets = nn.ModuleList(resnets) - self.motion_modules = nn.ModuleList(motion_modules) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - Downsample2D( - out_channels, - use_conv=True, - out_channels=out_channels, - padding=downsample_padding, - name="op", - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - num_frames: int = 1, - *args, - **kwargs, - ) -> Union[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) - - output_states = () - - blocks = zip(self.resnets, self.motion_modules) - for resnet, motion_module in blocks: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = resnet(input_tensor=hidden_states, temb=temb) - - hidden_states = motion_module(hidden_states, num_frames=num_frames) - - output_states = output_states + (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states=hidden_states) - - output_states = output_states + (hidden_states,) - - return hidden_states, output_states - - -class CrossAttnDownBlockMotion(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[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_pre_norm: bool = True, - num_attention_heads: int = 1, - cross_attention_dim: int = 1280, - output_scale_factor: float = 1.0, - downsample_padding: int = 1, - add_downsample: bool = True, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - only_cross_attention: bool = False, - upcast_attention: bool = False, - attention_type: str = "default", - temporal_cross_attention_dim: Optional[int] = None, - temporal_num_attention_heads: int = 8, - temporal_max_seq_length: int = 32, - temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1, - temporal_double_self_attention: bool = True, - ): - super().__init__() - resnets = [] - attentions = [] - motion_modules = [] - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - - # support for variable transformer layers per block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = (transformer_layers_per_block,) * num_layers - elif len(transformer_layers_per_block) != num_layers: - raise ValueError( - f"transformer_layers_per_block must be an integer or a list of integers of length {num_layers}" - ) - - # support for variable transformer layers per temporal block - if isinstance(temporal_transformer_layers_per_block, int): - temporal_transformer_layers_per_block = (temporal_transformer_layers_per_block,) * num_layers - elif len(temporal_transformer_layers_per_block) != num_layers: - raise ValueError( - f"temporal_transformer_layers_per_block must be an integer or a list of integers of length {num_layers}" - ) - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - if not dual_cross_attention: - attentions.append( - Transformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - attention_type=attention_type, - ) - ) - else: - attentions.append( - DualTransformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - ) - ) - - motion_modules.append( - AnimateDiffTransformer3D( - num_attention_heads=temporal_num_attention_heads, - in_channels=out_channels, - num_layers=temporal_transformer_layers_per_block[i], - norm_num_groups=resnet_groups, - cross_attention_dim=temporal_cross_attention_dim, - attention_bias=False, - activation_fn="geglu", - positional_embeddings="sinusoidal", - num_positional_embeddings=temporal_max_seq_length, - attention_head_dim=out_channels // temporal_num_attention_heads, - double_self_attention=temporal_double_self_attention, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - self.motion_modules = nn.ModuleList(motion_modules) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - Downsample2D( - out_channels, - use_conv=True, - out_channels=out_channels, - padding=downsample_padding, - name="op", - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - 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, - 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, - ): - 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.") - - output_states = () - - blocks = list(zip(self.resnets, self.attentions, self.motion_modules)) - for i, (resnet, attn, motion_module) in enumerate(blocks): - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = resnet(input_tensor=hidden_states, temb=temb) - - hidden_states = attn( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - - hidden_states = motion_module(hidden_states, num_frames=num_frames) - - # apply additional residuals to the output of the last pair of resnet and attention blocks - if i == len(blocks) - 1 and additional_residuals is not None: - hidden_states = hidden_states + additional_residuals - - output_states = output_states + (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states=hidden_states) - - output_states = output_states + (hidden_states,) - - return hidden_states, output_states - - -class CrossAttnUpBlockMotion(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - prev_output_channel: int, - temb_channels: int, - resolution_idx: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[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_pre_norm: bool = True, - num_attention_heads: int = 1, - cross_attention_dim: int = 1280, - output_scale_factor: float = 1.0, - add_upsample: bool = True, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - only_cross_attention: bool = False, - upcast_attention: bool = False, - attention_type: str = "default", - temporal_cross_attention_dim: Optional[int] = None, - temporal_num_attention_heads: int = 8, - temporal_max_seq_length: int = 32, - temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1, - ): - super().__init__() - resnets = [] - attentions = [] - motion_modules = [] - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - - # support for variable transformer layers per block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = (transformer_layers_per_block,) * num_layers - elif len(transformer_layers_per_block) != num_layers: - raise ValueError( - f"transformer_layers_per_block must be an integer or a list of integers of length {num_layers}, got {len(transformer_layers_per_block)}" - ) - - # support for variable transformer layers per temporal block - if isinstance(temporal_transformer_layers_per_block, int): - temporal_transformer_layers_per_block = (temporal_transformer_layers_per_block,) * num_layers - elif len(temporal_transformer_layers_per_block) != num_layers: - raise ValueError( - f"temporal_transformer_layers_per_block must be an integer or a list of integers of length {num_layers}, got {len(temporal_transformer_layers_per_block)}" - ) - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - if not dual_cross_attention: - attentions.append( - Transformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - attention_type=attention_type, - ) - ) - else: - attentions.append( - DualTransformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - ) - ) - motion_modules.append( - AnimateDiffTransformer3D( - num_attention_heads=temporal_num_attention_heads, - in_channels=out_channels, - num_layers=temporal_transformer_layers_per_block[i], - norm_num_groups=resnet_groups, - cross_attention_dim=temporal_cross_attention_dim, - attention_bias=False, - activation_fn="geglu", - positional_embeddings="sinusoidal", - num_positional_embeddings=temporal_max_seq_length, - attention_head_dim=out_channels // temporal_num_attention_heads, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - self.motion_modules = nn.ModuleList(motion_modules) - - if add_upsample: - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - 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, - num_frames: int = 1, - ) -> 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.") - - is_freeu_enabled = ( - getattr(self, "s1", None) - and getattr(self, "s2", None) - and getattr(self, "b1", None) - and getattr(self, "b2", None) - ) - - blocks = zip(self.resnets, self.attentions, self.motion_modules) - for resnet, attn, motion_module in blocks: - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - - # FreeU: Only operate on the first two stages - if is_freeu_enabled: - hidden_states, res_hidden_states = apply_freeu( - self.resolution_idx, - hidden_states, - res_hidden_states, - s1=self.s1, - s2=self.s2, - b1=self.b1, - b2=self.b2, - ) - - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = resnet(input_tensor=hidden_states, temb=temb) - - hidden_states = attn( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - - hidden_states = motion_module(hidden_states, num_frames=num_frames) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states=hidden_states, output_size=upsample_size) - - return hidden_states - - -class UpBlockMotion(nn.Module): - def __init__( - self, - in_channels: int, - prev_output_channel: int, - out_channels: int, - temb_channels: int, - resolution_idx: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor: float = 1.0, - add_upsample: bool = True, - temporal_cross_attention_dim: Optional[int] = None, - temporal_num_attention_heads: int = 8, - temporal_max_seq_length: int = 32, - temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1, - ): - super().__init__() - resnets = [] - motion_modules = [] - - # support for variable transformer layers per temporal block - if isinstance(temporal_transformer_layers_per_block, int): - temporal_transformer_layers_per_block = (temporal_transformer_layers_per_block,) * num_layers - elif len(temporal_transformer_layers_per_block) != num_layers: - raise ValueError( - f"temporal_transformer_layers_per_block must be an integer or a list of integers of length {num_layers}" - ) - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - motion_modules.append( - AnimateDiffTransformer3D( - num_attention_heads=temporal_num_attention_heads, - in_channels=out_channels, - num_layers=temporal_transformer_layers_per_block[i], - norm_num_groups=resnet_groups, - cross_attention_dim=temporal_cross_attention_dim, - attention_bias=False, - activation_fn="geglu", - positional_embeddings="sinusoidal", - num_positional_embeddings=temporal_max_seq_length, - attention_head_dim=out_channels // temporal_num_attention_heads, - ) - ) - - self.resnets = nn.ModuleList(resnets) - self.motion_modules = nn.ModuleList(motion_modules) - - if add_upsample: - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - def forward( - self, - hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - upsample_size=None, - num_frames: int = 1, - *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) - - is_freeu_enabled = ( - getattr(self, "s1", None) - and getattr(self, "s2", None) - and getattr(self, "b1", None) - and getattr(self, "b2", None) - ) - - blocks = zip(self.resnets, self.motion_modules) - - for resnet, motion_module in blocks: - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - - # FreeU: Only operate on the first two stages - if is_freeu_enabled: - hidden_states, res_hidden_states = apply_freeu( - self.resolution_idx, - hidden_states, - res_hidden_states, - s1=self.s1, - s2=self.s2, - b1=self.b1, - b2=self.b2, - ) - - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = resnet(input_tensor=hidden_states, temb=temb) - - hidden_states = motion_module(hidden_states, num_frames=num_frames) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states=hidden_states, output_size=upsample_size) - - return hidden_states - - -class UNetMidBlockCrossAttnMotion(nn.Module): - def __init__( - self, - in_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[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_pre_norm: bool = True, - num_attention_heads: int = 1, - output_scale_factor: float = 1.0, - cross_attention_dim: int = 1280, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - upcast_attention: bool = False, - attention_type: str = "default", - temporal_num_attention_heads: int = 1, - temporal_cross_attention_dim: Optional[int] = None, - temporal_max_seq_length: int = 32, - temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1, - ): - super().__init__() - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) - - # support for variable transformer layers per block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = (transformer_layers_per_block,) * num_layers - elif len(transformer_layers_per_block) != num_layers: - raise ValueError( - f"`transformer_layers_per_block` should be an integer or a list of integers of length {num_layers}." - ) - - # support for variable transformer layers per temporal block - if isinstance(temporal_transformer_layers_per_block, int): - temporal_transformer_layers_per_block = (temporal_transformer_layers_per_block,) * num_layers - elif len(temporal_transformer_layers_per_block) != num_layers: - raise ValueError( - f"`temporal_transformer_layers_per_block` should be an integer or a list of integers of length {num_layers}." - ) - - # there is always at least one resnet - resnets = [ - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ] - attentions = [] - motion_modules = [] - - for i in range(num_layers): - if not dual_cross_attention: - attentions.append( - Transformer2DModel( - num_attention_heads, - in_channels // num_attention_heads, - in_channels=in_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - attention_type=attention_type, - ) - ) - else: - attentions.append( - DualTransformer2DModel( - num_attention_heads, - in_channels // num_attention_heads, - in_channels=in_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - ) - ) - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - motion_modules.append( - AnimateDiffTransformer3D( - num_attention_heads=temporal_num_attention_heads, - attention_head_dim=in_channels // temporal_num_attention_heads, - in_channels=in_channels, - num_layers=temporal_transformer_layers_per_block[i], - norm_num_groups=resnet_groups, - cross_attention_dim=temporal_cross_attention_dim, - attention_bias=False, - positional_embeddings="sinusoidal", - num_positional_embeddings=temporal_max_seq_length, - activation_fn="geglu", - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - self.motion_modules = nn.ModuleList(motion_modules) - - self.gradient_checkpointing = False - - 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, - num_frames: int = 1, - ) -> 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.") - - hidden_states = self.resnets[0](input_tensor=hidden_states, temb=temb) - - blocks = zip(self.attentions, self.resnets[1:], self.motion_modules) - for attn, resnet, motion_module in blocks: - hidden_states = attn( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - motion_module, hidden_states, None, None, None, num_frames, None - ) - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = motion_module(hidden_states, None, None, None, num_frames, None) - hidden_states = resnet(input_tensor=hidden_states, temb=temb) - - return hidden_states - - -class MotionModules(nn.Module): - 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, - attention_bias: bool = False, - cross_attention_dim: Optional[int] = None, - activation_fn: str = "geglu", - norm_num_groups: int = 32, - max_seq_length: int = 32, - ): - super().__init__() - self.motion_modules = nn.ModuleList([]) - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = (transformer_layers_per_block,) * layers_per_block - elif len(transformer_layers_per_block) != layers_per_block: - raise ValueError( - f"The number of transformer layers per block must match the number of layers per block, " - f"got {layers_per_block} and {len(transformer_layers_per_block)}" - ) - - for i in range(layers_per_block): - self.motion_modules.append( - AnimateDiffTransformer3D( - in_channels=in_channels, - num_layers=transformer_layers_per_block[i], - norm_num_groups=norm_num_groups, - cross_attention_dim=cross_attention_dim, - activation_fn=activation_fn, - attention_bias=attention_bias, - num_attention_heads=num_attention_heads, - attention_head_dim=in_channels // num_attention_heads, - positional_embeddings="sinusoidal", - num_positional_embeddings=max_seq_length, - ) - ) - - -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, - 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_norm_num_groups: int = 32, - motion_max_seq_length: int = 32, - use_motion_mid_block: bool = True, - conv_in_channels: Optional[int] = None, - ): - """Container to store AnimateDiff Motion Modules - - Args: - 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): - The number of motion layers per UNet block. - 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): - 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): - 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. - motion_max_seq_length (`int`, *optional*, defaults to 32): - The maximum sequence length to use in the motion module. - use_motion_mid_block (`bool`, *optional*, defaults to True): - Whether to use a motion module in the middle of the UNet. - """ - - super().__init__() - down_blocks = [] - up_blocks = [] - - if isinstance(motion_layers_per_block, int): - motion_layers_per_block = (motion_layers_per_block,) * len(block_out_channels) - elif len(motion_layers_per_block) != len(block_out_channels): - raise ValueError( - f"The number of motion layers per block must match the number of blocks, " - f"got {len(block_out_channels)} and {len(motion_layers_per_block)}" - ) - - if isinstance(motion_transformer_layers_per_block, int): - motion_transformer_layers_per_block = (motion_transformer_layers_per_block,) * len(block_out_channels) - - if isinstance(motion_transformer_layers_per_mid_block, int): - motion_transformer_layers_per_mid_block = ( - motion_transformer_layers_per_mid_block, - ) * motion_mid_block_layers_per_block - elif len(motion_transformer_layers_per_mid_block) != motion_mid_block_layers_per_block: - raise ValueError( - f"The number of layers per mid block ({motion_mid_block_layers_per_block}) " - f"must match the length of motion_transformer_layers_per_mid_block ({len(motion_transformer_layers_per_mid_block)})" - ) - - if isinstance(motion_num_attention_heads, int): - motion_num_attention_heads = (motion_num_attention_heads,) * len(block_out_channels) - elif len(motion_num_attention_heads) != len(block_out_channels): - raise ValueError( - f"The length of the attention head number tuple in the motion module must match the " - f"number of block, got {len(motion_num_attention_heads)} and {len(block_out_channels)}" - ) - - if conv_in_channels: - # input - self.conv_in = nn.Conv2d(conv_in_channels, block_out_channels[0], kernel_size=3, padding=1) - else: - self.conv_in = None - - for i, channel in enumerate(block_out_channels): - output_channel = block_out_channels[i] - down_blocks.append( - MotionModules( - in_channels=output_channel, - norm_num_groups=motion_norm_num_groups, - cross_attention_dim=None, - activation_fn="geglu", - attention_bias=False, - num_attention_heads=motion_num_attention_heads[i], - max_seq_length=motion_max_seq_length, - layers_per_block=motion_layers_per_block[i], - transformer_layers_per_block=motion_transformer_layers_per_block[i], - ) - ) - - if use_motion_mid_block: - self.mid_block = MotionModules( - in_channels=block_out_channels[-1], - norm_num_groups=motion_norm_num_groups, - cross_attention_dim=None, - activation_fn="geglu", - attention_bias=False, - num_attention_heads=motion_num_attention_heads[-1], - max_seq_length=motion_max_seq_length, - layers_per_block=motion_mid_block_layers_per_block, - transformer_layers_per_block=motion_transformer_layers_per_mid_block, - ) - else: - self.mid_block = None - - reversed_block_out_channels = list(reversed(block_out_channels)) - output_channel = reversed_block_out_channels[0] - - reversed_motion_layers_per_block = list(reversed(motion_layers_per_block)) - reversed_motion_transformer_layers_per_block = list(reversed(motion_transformer_layers_per_block)) - reversed_motion_num_attention_heads = list(reversed(motion_num_attention_heads)) - for i, channel in enumerate(reversed_block_out_channels): - output_channel = reversed_block_out_channels[i] - up_blocks.append( - MotionModules( - in_channels=output_channel, - norm_num_groups=motion_norm_num_groups, - cross_attention_dim=None, - activation_fn="geglu", - attention_bias=False, - num_attention_heads=reversed_motion_num_attention_heads[i], - max_seq_length=motion_max_seq_length, - layers_per_block=reversed_motion_layers_per_block[i] + 1, - transformer_layers_per_block=reversed_motion_transformer_layers_per_block[i], - ) - ) - - self.down_blocks = nn.ModuleList(down_blocks) - self.up_blocks = nn.ModuleList(up_blocks) - - def forward(self, sample): - pass - - -class UNetMotionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin): - r""" - A modified conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a - sample shaped output. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - """ - - _supports_gradient_checkpointing = True - _skip_layerwise_casting_patterns = ["norm"] - - @register_to_config - def __init__( - self, - sample_size: Optional[int] = None, - in_channels: int = 4, - out_channels: int = 4, - down_block_types: Tuple[str, ...] = ( - "CrossAttnDownBlockMotion", - "CrossAttnDownBlockMotion", - "CrossAttnDownBlockMotion", - "DownBlockMotion", - ), - 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, - 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, - use_linear_projection: bool = False, - num_attention_heads: Union[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, - 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, - ): - super().__init__() - - self.sample_size = sample_size - - # Check inputs - if len(down_block_types) != len(up_block_types): - raise ValueError( - f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}." - ) - - if len(block_out_channels) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." - ) - - if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}." - ) - - if isinstance(transformer_layers_per_block, list) and reverse_transformer_layers_per_block is None: - for layer_number_per_block in transformer_layers_per_block: - if isinstance(layer_number_per_block, list): - raise ValueError("Must provide 'reverse_transformer_layers_per_block` if using asymmetrical UNet.") - - if ( - isinstance(temporal_transformer_layers_per_block, list) - and reverse_temporal_transformer_layers_per_block is None - ): - for layer_number_per_block in temporal_transformer_layers_per_block: - if isinstance(layer_number_per_block, list): - raise ValueError( - "Must provide 'reverse_temporal_transformer_layers_per_block` if using asymmetrical motion module in UNet." - ) - - # input - conv_in_kernel = 3 - conv_out_kernel = 3 - conv_in_padding = (conv_in_kernel - 1) // 2 - self.conv_in = nn.Conv2d( - in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding - ) - - # time - time_embed_dim = block_out_channels[0] * 4 - self.time_proj = Timesteps(block_out_channels[0], True, 0) - timestep_input_dim = block_out_channels[0] - - self.time_embedding = TimestepEmbedding( - timestep_input_dim, time_embed_dim, act_fn=act_fn, cond_proj_dim=time_cond_proj_dim - ) - - if encoder_hid_dim_type is None: - self.encoder_hid_proj = None - - if addition_embed_type == "text_time": - self.add_time_proj = Timesteps(addition_time_embed_dim, True, 0) - self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) - - # class embedding - self.down_blocks = nn.ModuleList([]) - self.up_blocks = nn.ModuleList([]) - - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(down_block_types) - - if isinstance(cross_attention_dim, int): - cross_attention_dim = (cross_attention_dim,) * len(down_block_types) - - if isinstance(layers_per_block, int): - layers_per_block = [layers_per_block] * len(down_block_types) - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) - - if isinstance(reverse_transformer_layers_per_block, int): - reverse_transformer_layers_per_block = [reverse_transformer_layers_per_block] * len(down_block_types) - - if isinstance(temporal_transformer_layers_per_block, int): - temporal_transformer_layers_per_block = [temporal_transformer_layers_per_block] * len(down_block_types) - - if isinstance(reverse_temporal_transformer_layers_per_block, int): - reverse_temporal_transformer_layers_per_block = [reverse_temporal_transformer_layers_per_block] * len( - down_block_types - ) - - if isinstance(motion_num_attention_heads, int): - motion_num_attention_heads = (motion_num_attention_heads,) * len(down_block_types) - - # down - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - if down_block_type == "CrossAttnDownBlockMotion": - down_block = CrossAttnDownBlockMotion( - in_channels=input_channel, - out_channels=output_channel, - temb_channels=time_embed_dim, - num_layers=layers_per_block[i], - transformer_layers_per_block=transformer_layers_per_block[i], - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - num_attention_heads=num_attention_heads[i], - cross_attention_dim=cross_attention_dim[i], - downsample_padding=downsample_padding, - add_downsample=not is_final_block, - use_linear_projection=use_linear_projection, - temporal_num_attention_heads=motion_num_attention_heads[i], - temporal_max_seq_length=motion_max_seq_length, - temporal_transformer_layers_per_block=temporal_transformer_layers_per_block[i], - ) - elif down_block_type == "DownBlockMotion": - down_block = DownBlockMotion( - in_channels=input_channel, - out_channels=output_channel, - temb_channels=time_embed_dim, - num_layers=layers_per_block[i], - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - add_downsample=not is_final_block, - downsample_padding=downsample_padding, - temporal_num_attention_heads=motion_num_attention_heads[i], - temporal_max_seq_length=motion_max_seq_length, - temporal_transformer_layers_per_block=temporal_transformer_layers_per_block[i], - ) - else: - raise ValueError( - "Invalid `down_block_type` encountered. Must be one of `CrossAttnDownBlockMotion` or `DownBlockMotion`" - ) - - self.down_blocks.append(down_block) - - # mid - if transformer_layers_per_mid_block is None: - transformer_layers_per_mid_block = ( - transformer_layers_per_block[-1] if isinstance(transformer_layers_per_block[-1], int) else 1 - ) - - if use_motion_mid_block: - self.mid_block = UNetMidBlockCrossAttnMotion( - in_channels=block_out_channels[-1], - temb_channels=time_embed_dim, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - output_scale_factor=mid_block_scale_factor, - cross_attention_dim=cross_attention_dim[-1], - num_attention_heads=num_attention_heads[-1], - resnet_groups=norm_num_groups, - dual_cross_attention=False, - use_linear_projection=use_linear_projection, - num_layers=mid_block_layers, - temporal_num_attention_heads=motion_num_attention_heads[-1], - temporal_max_seq_length=motion_max_seq_length, - transformer_layers_per_block=transformer_layers_per_mid_block, - temporal_transformer_layers_per_block=temporal_transformer_layers_per_mid_block, - ) - - else: - self.mid_block = UNetMidBlock2DCrossAttn( - in_channels=block_out_channels[-1], - temb_channels=time_embed_dim, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - output_scale_factor=mid_block_scale_factor, - cross_attention_dim=cross_attention_dim[-1], - num_attention_heads=num_attention_heads[-1], - resnet_groups=norm_num_groups, - dual_cross_attention=False, - use_linear_projection=use_linear_projection, - num_layers=mid_block_layers, - transformer_layers_per_block=transformer_layers_per_mid_block, - ) - - # count how many layers upsample the images - self.num_upsamplers = 0 - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - reversed_num_attention_heads = list(reversed(num_attention_heads)) - reversed_layers_per_block = list(reversed(layers_per_block)) - reversed_cross_attention_dim = list(reversed(cross_attention_dim)) - reversed_motion_num_attention_heads = list(reversed(motion_num_attention_heads)) - - if reverse_transformer_layers_per_block is None: - reverse_transformer_layers_per_block = list(reversed(transformer_layers_per_block)) - - if reverse_temporal_transformer_layers_per_block is None: - reverse_temporal_transformer_layers_per_block = list(reversed(temporal_transformer_layers_per_block)) - - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - is_final_block = i == len(block_out_channels) - 1 - - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] - - # add upsample block for all BUT final layer - if not is_final_block: - add_upsample = True - self.num_upsamplers += 1 - else: - add_upsample = False - - if up_block_type == "CrossAttnUpBlockMotion": - up_block = CrossAttnUpBlockMotion( - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - temb_channels=time_embed_dim, - resolution_idx=i, - num_layers=reversed_layers_per_block[i] + 1, - transformer_layers_per_block=reverse_transformer_layers_per_block[i], - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - num_attention_heads=reversed_num_attention_heads[i], - cross_attention_dim=reversed_cross_attention_dim[i], - add_upsample=add_upsample, - use_linear_projection=use_linear_projection, - temporal_num_attention_heads=reversed_motion_num_attention_heads[i], - temporal_max_seq_length=motion_max_seq_length, - temporal_transformer_layers_per_block=reverse_temporal_transformer_layers_per_block[i], - ) - elif up_block_type == "UpBlockMotion": - up_block = UpBlockMotion( - in_channels=input_channel, - prev_output_channel=prev_output_channel, - out_channels=output_channel, - temb_channels=time_embed_dim, - resolution_idx=i, - num_layers=reversed_layers_per_block[i] + 1, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - add_upsample=add_upsample, - temporal_num_attention_heads=reversed_motion_num_attention_heads[i], - temporal_max_seq_length=motion_max_seq_length, - temporal_transformer_layers_per_block=reverse_temporal_transformer_layers_per_block[i], - ) - else: - raise ValueError( - "Invalid `up_block_type` encountered. Must be one of `CrossAttnUpBlockMotion` or `UpBlockMotion`" - ) - - self.up_blocks.append(up_block) - prev_output_channel = output_channel - - # out - if norm_num_groups is not None: - self.conv_norm_out = nn.GroupNorm( - num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps - ) - self.conv_act = nn.SiLU() - else: - self.conv_norm_out = None - self.conv_act = None - - conv_out_padding = (conv_out_kernel - 1) // 2 - self.conv_out = nn.Conv2d( - block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding - ) - - @classmethod - def from_unet2d( - cls, - unet: UNet2DConditionModel, - motion_adapter: Optional[MotionAdapter] = None, - load_weights: bool = True, - ): - has_motion_adapter = motion_adapter is not None - - if has_motion_adapter: - motion_adapter.to(device=unet.device) - - # check compatibility of number of blocks - if len(unet.config["down_block_types"]) != len(motion_adapter.config["block_out_channels"]): - raise ValueError("Incompatible Motion Adapter, got different number of blocks") - - # check layers compatibility for each block - if isinstance(unet.config["layers_per_block"], int): - expanded_layers_per_block = [unet.config["layers_per_block"]] * len(unet.config["down_block_types"]) - else: - expanded_layers_per_block = list(unet.config["layers_per_block"]) - if isinstance(motion_adapter.config["motion_layers_per_block"], int): - expanded_adapter_layers_per_block = [motion_adapter.config["motion_layers_per_block"]] * len( - motion_adapter.config["block_out_channels"] - ) - else: - expanded_adapter_layers_per_block = list(motion_adapter.config["motion_layers_per_block"]) - if expanded_layers_per_block != expanded_adapter_layers_per_block: - raise ValueError("Incompatible Motion Adapter, got different number of layers per block") - - # based on https://github.com/guoyww/AnimateDiff/blob/895f3220c06318ea0760131ec70408b466c49333/animatediff/models/unet.py#L459 - config = dict(unet.config) - config["_class_name"] = cls.__name__ - - down_blocks = [] - for down_blocks_type in config["down_block_types"]: - if "CrossAttn" in down_blocks_type: - down_blocks.append("CrossAttnDownBlockMotion") - else: - down_blocks.append("DownBlockMotion") - config["down_block_types"] = down_blocks - - up_blocks = [] - for down_blocks_type in config["up_block_types"]: - if "CrossAttn" in down_blocks_type: - up_blocks.append("CrossAttnUpBlockMotion") - else: - up_blocks.append("UpBlockMotion") - config["up_block_types"] = up_blocks - - if has_motion_adapter: - config["motion_num_attention_heads"] = motion_adapter.config["motion_num_attention_heads"] - config["motion_max_seq_length"] = motion_adapter.config["motion_max_seq_length"] - config["use_motion_mid_block"] = motion_adapter.config["use_motion_mid_block"] - config["layers_per_block"] = motion_adapter.config["motion_layers_per_block"] - config["temporal_transformer_layers_per_mid_block"] = motion_adapter.config[ - "motion_transformer_layers_per_mid_block" - ] - config["temporal_transformer_layers_per_block"] = motion_adapter.config[ - "motion_transformer_layers_per_block" - ] - config["motion_num_attention_heads"] = motion_adapter.config["motion_num_attention_heads"] - - # For PIA UNets we need to set the number input channels to 9 - if motion_adapter.config["conv_in_channels"]: - config["in_channels"] = motion_adapter.config["conv_in_channels"] - - # Need this for backwards compatibility with UNet2DConditionModel checkpoints - if not config.get("num_attention_heads"): - config["num_attention_heads"] = config["attention_head_dim"] - - expected_kwargs, optional_kwargs = cls._get_signature_keys(cls) - config = FrozenDict({k: config.get(k) for k in config if k in expected_kwargs or k in optional_kwargs}) - config["_class_name"] = cls.__name__ - model = cls.from_config(config) - - if not load_weights: - return model - - # Logic for loading PIA UNets which allow the first 4 channels to be any UNet2DConditionModel conv_in weight - # while the last 5 channels must be PIA conv_in weights. - if has_motion_adapter and motion_adapter.config["conv_in_channels"]: - model.conv_in = motion_adapter.conv_in - updated_conv_in_weight = torch.cat( - [unet.conv_in.weight, motion_adapter.conv_in.weight[:, 4:, :, :]], dim=1 - ) - model.conv_in.load_state_dict({"weight": updated_conv_in_weight, "bias": unet.conv_in.bias}) - else: - model.conv_in.load_state_dict(unet.conv_in.state_dict()) - - model.time_proj.load_state_dict(unet.time_proj.state_dict()) - model.time_embedding.load_state_dict(unet.time_embedding.state_dict()) - - if any( - isinstance(proc, (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0)) - for proc in unet.attn_processors.values() - ): - attn_procs = {} - for name, processor in unet.attn_processors.items(): - if name.endswith("attn1.processor"): - attn_processor_class = ( - AttnProcessor2_0 if hasattr(F, "scaled_dot_product_attention") else AttnProcessor - ) - attn_procs[name] = attn_processor_class() - else: - attn_processor_class = ( - IPAdapterAttnProcessor2_0 - if hasattr(F, "scaled_dot_product_attention") - else IPAdapterAttnProcessor - ) - attn_procs[name] = attn_processor_class( - hidden_size=processor.hidden_size, - cross_attention_dim=processor.cross_attention_dim, - scale=processor.scale, - num_tokens=processor.num_tokens, - ) - for name, processor in model.attn_processors.items(): - if name not in attn_procs: - attn_procs[name] = processor.__class__() - model.set_attn_processor(attn_procs) - model.config.encoder_hid_dim_type = "ip_image_proj" - model.encoder_hid_proj = unet.encoder_hid_proj - - for i, down_block in enumerate(unet.down_blocks): - model.down_blocks[i].resnets.load_state_dict(down_block.resnets.state_dict()) - if hasattr(model.down_blocks[i], "attentions"): - model.down_blocks[i].attentions.load_state_dict(down_block.attentions.state_dict()) - if model.down_blocks[i].downsamplers: - model.down_blocks[i].downsamplers.load_state_dict(down_block.downsamplers.state_dict()) - - for i, up_block in enumerate(unet.up_blocks): - model.up_blocks[i].resnets.load_state_dict(up_block.resnets.state_dict()) - if hasattr(model.up_blocks[i], "attentions"): - model.up_blocks[i].attentions.load_state_dict(up_block.attentions.state_dict()) - if model.up_blocks[i].upsamplers: - model.up_blocks[i].upsamplers.load_state_dict(up_block.upsamplers.state_dict()) - - model.mid_block.resnets.load_state_dict(unet.mid_block.resnets.state_dict()) - model.mid_block.attentions.load_state_dict(unet.mid_block.attentions.state_dict()) - - if unet.conv_norm_out is not None: - model.conv_norm_out.load_state_dict(unet.conv_norm_out.state_dict()) - if unet.conv_act is not None: - model.conv_act.load_state_dict(unet.conv_act.state_dict()) - model.conv_out.load_state_dict(unet.conv_out.state_dict()) - - if has_motion_adapter: - model.load_motion_modules(motion_adapter) - - # ensure that the Motion UNet is the same dtype as the UNet2DConditionModel - model.to(unet.dtype) - - return model - - def freeze_unet2d_params(self) -> None: - """Freeze the weights of just the UNet2DConditionModel, and leave the motion modules - unfrozen for fine tuning. - """ - # Freeze everything - for param in self.parameters(): - param.requires_grad = False - - # Unfreeze Motion Modules - for down_block in self.down_blocks: - motion_modules = down_block.motion_modules - for param in motion_modules.parameters(): - param.requires_grad = True - - for up_block in self.up_blocks: - motion_modules = up_block.motion_modules - for param in motion_modules.parameters(): - param.requires_grad = True - - if hasattr(self.mid_block, "motion_modules"): - motion_modules = self.mid_block.motion_modules - for param in motion_modules.parameters(): - param.requires_grad = True - - def load_motion_modules(self, motion_adapter: Optional[MotionAdapter]) -> 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): - self.up_blocks[i].motion_modules.load_state_dict(up_block.motion_modules.state_dict()) - - # to support older motion modules that don't have a mid_block - if hasattr(self.mid_block, "motion_modules"): - self.mid_block.motion_modules.load_state_dict(motion_adapter.mid_block.motion_modules.state_dict()) - - def save_motion_modules( - self, - save_directory: str, - is_main_process: bool = True, - safe_serialization: bool = True, - variant: Optional[str] = None, - push_to_hub: bool = False, - **kwargs, - ) -> None: - state_dict = self.state_dict() - - # Extract all motion modules - motion_state_dict = {} - for k, v in state_dict.items(): - if "motion_modules" in k: - motion_state_dict[k] = v - - adapter = MotionAdapter( - block_out_channels=self.config["block_out_channels"], - motion_layers_per_block=self.config["layers_per_block"], - motion_norm_num_groups=self.config["norm_num_groups"], - motion_num_attention_heads=self.config["motion_num_attention_heads"], - motion_max_seq_length=self.config["motion_max_seq_length"], - use_motion_mid_block=self.config["use_motion_mid_block"], - ) - adapter.load_state_dict(motion_state_dict) - adapter.save_pretrained( - save_directory=save_directory, - is_main_process=is_main_process, - safe_serialization=safe_serialization, - variant=variant, - push_to_hub=push_to_hub, - **kwargs, - ) - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None: - """ - Sets the attention processor to use [feed forward - chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). - - Parameters: - chunk_size (`int`, *optional*): - The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually - over each tensor of dim=`dim`. - dim (`int`, *optional*, defaults to `0`): - The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch) - or dim=1 (sequence length). - """ - if dim not in [0, 1]: - raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}") - - # By default chunk size is 1 - chunk_size = chunk_size or 1 - - def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int): - if hasattr(module, "set_chunk_feed_forward"): - module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) - - for child in module.children(): - fn_recursive_feed_forward(child, chunk_size, dim) - - for module in self.children(): - fn_recursive_feed_forward(module, chunk_size, dim) - - def disable_forward_chunking(self) -> None: - def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int): - if hasattr(module, "set_chunk_feed_forward"): - module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) - - for child in module.children(): - fn_recursive_feed_forward(child, chunk_size, dim) - - for module in self.children(): - fn_recursive_feed_forward(module, None, 0) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor - def set_default_attn_processor(self) -> None: - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - self.set_attn_processor(processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.enable_freeu - def enable_freeu(self, s1: float, s2: float, b1: float, b2: float) -> None: - r"""Enables the FreeU mechanism from https://huggingface.co/papers/2309.11497. - - The suffixes after the scaling factors represent the stage blocks where they are being applied. - - Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that - are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL. - - Args: - s1 (`float`): - Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to - mitigate the "oversmoothing effect" in the enhanced denoising process. - s2 (`float`): - Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to - mitigate the "oversmoothing effect" in the enhanced denoising process. - b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features. - b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features. - """ - for i, upsample_block in enumerate(self.up_blocks): - setattr(upsample_block, "s1", s1) - setattr(upsample_block, "s2", s2) - setattr(upsample_block, "b1", b1) - setattr(upsample_block, "b2", b2) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.disable_freeu - def disable_freeu(self) -> None: - """Disables the FreeU mechanism.""" - freeu_keys = {"s1", "s2", "b1", "b2"} - for i, upsample_block in enumerate(self.up_blocks): - for k in freeu_keys: - if hasattr(upsample_block, k) or getattr(upsample_block, k, None) is not None: - setattr(upsample_block, k, None) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - self.set_attn_processor(FusedAttnProcessor2_0()) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - if self.original_attn_processors is not None: - self.set_attn_processor(self.original_attn_processors) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[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, - return_dict: bool = True, - ) -> Union[UNetMotionOutput, Tuple[torch.Tensor]]: - r""" - The [`UNetMotionModel`] forward method. - - Args: - sample (`torch.Tensor`): - The noisy input tensor with the following shape `(batch, num_frames, channel, height, width`. - timestep (`torch.Tensor` or `float` or `int`): The number of timesteps to denoise an input. - encoder_hidden_states (`torch.Tensor`): - The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. - timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`): - Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed - through the `self.time_embedding` layer to obtain the timestep embeddings. - attention_mask (`torch.Tensor`, *optional*, defaults to `None`): - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask - is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large - negative values to the attention scores corresponding to "discard" tokens. - cross_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). - down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*): - A tuple of tensors that if specified are added to the residuals of down unet blocks. - mid_block_additional_residual: (`torch.Tensor`, *optional*): - A tensor that if specified is added to the residual of the middle unet block. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.unets.unet_motion_model.UNetMotionOutput`] instead of a plain - tuple. - - Returns: - [`~models.unets.unet_motion_model.UNetMotionOutput`] or `tuple`: - If `return_dict` is True, an [`~models.unets.unet_motion_model.UNetMotionOutput`] is returned, - otherwise a `tuple` is returned where the first element is the sample tensor. - """ - # By default samples have to be AT least a multiple of the overall upsampling factor. - # The overall upsampling factor is equal to 2 ** (# num of upsampling layears). - # However, the upsampling interpolation output size can be forced to fit any upsampling size - # on the fly if necessary. - default_overall_up_factor = 2**self.num_upsamplers - - # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` - forward_upsample_size = False - upsample_size = None - - if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): - logger.info("Forward upsample size to force interpolation output size.") - forward_upsample_size = True - - # prepare attention_mask - if attention_mask is not None: - attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 - attention_mask = attention_mask.unsqueeze(1) - - # 1. time - 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 - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = sample.device.type == "mps" - is_npu = sample.device.type == "npu" - if isinstance(timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) - elif len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - num_frames = sample.shape[2] - timesteps = timesteps.expand(sample.shape[0]) - - t_emb = self.time_proj(timesteps) - - # timesteps does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - t_emb = t_emb.to(dtype=self.dtype) - - emb = self.time_embedding(t_emb, timestep_cond) - aug_emb = None - - if self.config.addition_embed_type == "text_time": - if "text_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" - ) - - text_embeds = added_cond_kwargs.get("text_embeds") - if "time_ids" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" - ) - time_ids = added_cond_kwargs.get("time_ids") - time_embeds = self.add_time_proj(time_ids.flatten()) - time_embeds = time_embeds.reshape((text_embeds.shape[0], -1)) - - add_embeds = torch.concat([text_embeds, time_embeds], dim=-1) - add_embeds = add_embeds.to(emb.dtype) - aug_emb = self.add_embedding(add_embeds) - - emb = emb if aug_emb is None else emb + aug_emb - emb = emb.repeat_interleave(num_frames, dim=0, output_size=emb.shape[0] * num_frames) - - if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "ip_image_proj": - if "image_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'ip_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" - ) - image_embeds = added_cond_kwargs.get("image_embeds") - image_embeds = self.encoder_hid_proj(image_embeds) - image_embeds = [ - image_embed.repeat_interleave(num_frames, dim=0, output_size=image_embed.shape[0] * num_frames) - for image_embed in image_embeds - ] - encoder_hidden_states = (encoder_hidden_states, image_embeds) - - # 2. pre-process - sample = sample.permute(0, 2, 1, 3, 4).reshape((sample.shape[0] * num_frames, -1) + sample.shape[3:]) - sample = self.conv_in(sample) - - # 3. down - down_block_res_samples = (sample,) - for downsample_block in self.down_blocks: - if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: - sample, res_samples = downsample_block( - hidden_states=sample, - temb=emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - ) - else: - sample, res_samples = downsample_block(hidden_states=sample, temb=emb, num_frames=num_frames) - - down_block_res_samples += res_samples - - if down_block_additional_residuals is not None: - new_down_block_res_samples = () - - for down_block_res_sample, down_block_additional_residual in zip( - down_block_res_samples, down_block_additional_residuals - ): - down_block_res_sample = down_block_res_sample + down_block_additional_residual - new_down_block_res_samples += (down_block_res_sample,) - - down_block_res_samples = new_down_block_res_samples - - # 4. mid - if self.mid_block is not None: - # To support older versions of motion modules that don't have a mid_block - if hasattr(self.mid_block, "motion_modules"): - sample = self.mid_block( - sample, - emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - ) - else: - sample = self.mid_block( - sample, - emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - ) - - if mid_block_additional_residual is not None: - sample = sample + mid_block_additional_residual - - # 5. up - for i, upsample_block in enumerate(self.up_blocks): - is_final_block = i == len(self.up_blocks) - 1 - - res_samples = down_block_res_samples[-len(upsample_block.resnets) :] - down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] - - # if we have not reached the final block and need to forward the - # upsample size, we do it here - if not is_final_block and forward_upsample_size: - upsample_size = down_block_res_samples[-1].shape[2:] - - if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - encoder_hidden_states=encoder_hidden_states, - upsample_size=upsample_size, - attention_mask=attention_mask, - num_frames=num_frames, - cross_attention_kwargs=cross_attention_kwargs, - ) - else: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - upsample_size=upsample_size, - num_frames=num_frames, - ) - - # 6. post-process - if self.conv_norm_out: - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - - sample = self.conv_out(sample) - - # reshape to (batch, channel, framerate, width, height) - sample = sample[None, :].reshape((-1, num_frames) + sample.shape[1:]).permute(0, 2, 1, 3, 4) - - if not return_dict: - return (sample,) - - return UNetMotionOutput(sample=sample) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_spatio_temporal_condition.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_spatio_temporal_condition.py deleted file mode 100644 index 059a6e807c8e57eb8d4bc3a8846ee047d9eea84e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_spatio_temporal_condition.py +++ /dev/null @@ -1,512 +0,0 @@ -from dataclasses import dataclass -from typing import Dict, Optional, Tuple, Union - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import UNet2DConditionLoadersMixin -from ...utils import BaseOutput, logging -from ..attention_processor import CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnProcessor -from ..embeddings import TimestepEmbedding, Timesteps -from ..modeling_utils import ModelMixin -from .unet_3d_blocks import UNetMidBlockSpatioTemporal, get_down_block, get_up_block - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class UNetSpatioTemporalConditionOutput(BaseOutput): - """ - The output of [`UNetSpatioTemporalConditionModel`]. - - Args: - sample (`torch.Tensor` of shape `(batch_size, num_frames, num_channels, height, width)`): - The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model. - """ - - sample: torch.Tensor = None - - -class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): - r""" - A conditional Spatio-Temporal UNet model that takes a noisy video frames, conditional state, and a timestep and - returns a sample shaped output. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - Parameters: - 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")`): - The tuple of downsample blocks to use. - 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)`): - 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): - The dimension of the cross attention features. - 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)`): - The number of attention heads. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - """ - - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - sample_size: Optional[int] = None, - in_channels: int = 8, - out_channels: int = 4, - down_block_types: Tuple[str] = ( - "CrossAttnDownBlockSpatioTemporal", - "CrossAttnDownBlockSpatioTemporal", - "CrossAttnDownBlockSpatioTemporal", - "DownBlockSpatioTemporal", - ), - up_block_types: Tuple[str] = ( - "UpBlockSpatioTemporal", - "CrossAttnUpBlockSpatioTemporal", - "CrossAttnUpBlockSpatioTemporal", - "CrossAttnUpBlockSpatioTemporal", - ), - 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), - num_frames: int = 25, - ): - super().__init__() - - self.sample_size = sample_size - - # Check inputs - if len(down_block_types) != len(up_block_types): - raise ValueError( - f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}." - ) - - if len(block_out_channels) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." - ) - - if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}." - ) - - # input - self.conv_in = nn.Conv2d( - in_channels, - block_out_channels[0], - kernel_size=3, - padding=1, - ) - - # time - time_embed_dim = block_out_channels[0] * 4 - - self.time_proj = Timesteps(block_out_channels[0], True, downscale_freq_shift=0) - timestep_input_dim = block_out_channels[0] - - self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim) - - self.add_time_proj = Timesteps(addition_time_embed_dim, True, downscale_freq_shift=0) - self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) - - self.down_blocks = nn.ModuleList([]) - self.up_blocks = nn.ModuleList([]) - - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(down_block_types) - - if isinstance(cross_attention_dim, int): - cross_attention_dim = (cross_attention_dim,) * len(down_block_types) - - if isinstance(layers_per_block, int): - layers_per_block = [layers_per_block] * len(down_block_types) - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) - - blocks_time_embed_dim = time_embed_dim - - # down - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = get_down_block( - down_block_type, - num_layers=layers_per_block[i], - transformer_layers_per_block=transformer_layers_per_block[i], - in_channels=input_channel, - out_channels=output_channel, - temb_channels=blocks_time_embed_dim, - add_downsample=not is_final_block, - resnet_eps=1e-5, - cross_attention_dim=cross_attention_dim[i], - num_attention_heads=num_attention_heads[i], - resnet_act_fn="silu", - ) - self.down_blocks.append(down_block) - - # mid - self.mid_block = UNetMidBlockSpatioTemporal( - block_out_channels[-1], - temb_channels=blocks_time_embed_dim, - transformer_layers_per_block=transformer_layers_per_block[-1], - cross_attention_dim=cross_attention_dim[-1], - num_attention_heads=num_attention_heads[-1], - ) - - # count how many layers upsample the images - self.num_upsamplers = 0 - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - reversed_num_attention_heads = list(reversed(num_attention_heads)) - reversed_layers_per_block = list(reversed(layers_per_block)) - reversed_cross_attention_dim = list(reversed(cross_attention_dim)) - reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block)) - - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - is_final_block = i == len(block_out_channels) - 1 - - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] - - # add upsample block for all BUT final layer - if not is_final_block: - add_upsample = True - self.num_upsamplers += 1 - else: - add_upsample = False - - up_block = get_up_block( - up_block_type, - num_layers=reversed_layers_per_block[i] + 1, - transformer_layers_per_block=reversed_transformer_layers_per_block[i], - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - temb_channels=blocks_time_embed_dim, - add_upsample=add_upsample, - resnet_eps=1e-5, - resolution_idx=i, - cross_attention_dim=reversed_cross_attention_dim[i], - num_attention_heads=reversed_num_attention_heads[i], - resnet_act_fn="silu", - ) - self.up_blocks.append(up_block) - prev_output_channel = output_channel - - # out - self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=32, eps=1e-5) - self.conv_act = nn.SiLU() - - self.conv_out = nn.Conv2d( - block_out_channels[0], - out_channels, - kernel_size=3, - padding=1, - ) - - @property - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - 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: - """ - Sets the attention processor to use [feed forward - chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). - - Parameters: - chunk_size (`int`, *optional*): - The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually - over each tensor of dim=`dim`. - dim (`int`, *optional*, defaults to `0`): - The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch) - or dim=1 (sequence length). - """ - if dim not in [0, 1]: - raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}") - - # By default chunk size is 1 - chunk_size = chunk_size or 1 - - def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int): - if hasattr(module, "set_chunk_feed_forward"): - module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) - - for child in module.children(): - fn_recursive_feed_forward(child, chunk_size, dim) - - for module in self.children(): - fn_recursive_feed_forward(module, chunk_size, dim) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], - encoder_hidden_states: torch.Tensor, - added_time_ids: torch.Tensor, - return_dict: bool = True, - ) -> Union[UNetSpatioTemporalConditionOutput, Tuple]: - r""" - The [`UNetSpatioTemporalConditionModel`] forward method. - - Args: - sample (`torch.Tensor`): - The noisy input tensor with the following shape `(batch, num_frames, channel, height, width)`. - timestep (`torch.Tensor` or `float` or `int`): The number of timesteps to denoise an input. - encoder_hidden_states (`torch.Tensor`): - The encoder hidden states with shape `(batch, sequence_length, cross_attention_dim)`. - added_time_ids: (`torch.Tensor`): - The additional time ids with shape `(batch, num_additional_ids)`. These are encoded with sinusoidal - embeddings and added to the time embeddings. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] instead - of a plain tuple. - Returns: - [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] or `tuple`: - If `return_dict` is True, an [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] is - returned, otherwise a `tuple` is returned where the first element is the sample tensor. - """ - # By default samples have to be AT least a multiple of the overall upsampling factor. - # The overall upsampling factor is equal to 2 ** (# num of upsampling layears). - # However, the upsampling interpolation output size can be forced to fit any upsampling size - # on the fly if necessary. - default_overall_up_factor = 2**self.num_upsamplers - - # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` - forward_upsample_size = False - upsample_size = None - - if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): - logger.info("Forward upsample size to force interpolation output size.") - forward_upsample_size = True - - # 1. time - 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 - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = sample.device.type == "mps" - is_npu = sample.device.type == "npu" - if isinstance(timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) - elif len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - batch_size, num_frames = sample.shape[:2] - timesteps = timesteps.expand(batch_size) - - t_emb = self.time_proj(timesteps) - - # `Timesteps` does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - t_emb = t_emb.to(dtype=sample.dtype) - - emb = self.time_embedding(t_emb) - - time_embeds = self.add_time_proj(added_time_ids.flatten()) - time_embeds = time_embeds.reshape((batch_size, -1)) - time_embeds = time_embeds.to(emb.dtype) - aug_emb = self.add_embedding(time_embeds) - emb = emb + aug_emb - - # Flatten the batch and frames dimensions - # sample: [batch, frames, channels, height, width] -> [batch * frames, channels, height, width] - sample = sample.flatten(0, 1) - # Repeat the embeddings num_video_frames times - # emb: [batch, channels] -> [batch * frames, channels] - emb = emb.repeat_interleave(num_frames, dim=0, output_size=emb.shape[0] * num_frames) - # encoder_hidden_states: [batch, 1, channels] -> [batch * frames, 1, channels] - encoder_hidden_states = encoder_hidden_states.repeat_interleave( - num_frames, dim=0, output_size=encoder_hidden_states.shape[0] * num_frames - ) - - # 2. pre-process - sample = self.conv_in(sample) - - image_only_indicator = torch.zeros(batch_size, num_frames, dtype=sample.dtype, device=sample.device) - - down_block_res_samples = (sample,) - for downsample_block in self.down_blocks: - if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: - sample, res_samples = downsample_block( - hidden_states=sample, - temb=emb, - encoder_hidden_states=encoder_hidden_states, - image_only_indicator=image_only_indicator, - ) - else: - sample, res_samples = downsample_block( - hidden_states=sample, - temb=emb, - image_only_indicator=image_only_indicator, - ) - - down_block_res_samples += res_samples - - # 4. mid - sample = self.mid_block( - hidden_states=sample, - temb=emb, - encoder_hidden_states=encoder_hidden_states, - image_only_indicator=image_only_indicator, - ) - - # 5. up - for i, upsample_block in enumerate(self.up_blocks): - is_final_block = i == len(self.up_blocks) - 1 - - res_samples = down_block_res_samples[-len(upsample_block.resnets) :] - down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] - - # if we have not reached the final block and need to forward the - # upsample size, we do it here - if not is_final_block and forward_upsample_size: - upsample_size = down_block_res_samples[-1].shape[2:] - - if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - encoder_hidden_states=encoder_hidden_states, - upsample_size=upsample_size, - image_only_indicator=image_only_indicator, - ) - else: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - upsample_size=upsample_size, - image_only_indicator=image_only_indicator, - ) - - # 6. post-process - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - sample = self.conv_out(sample) - - # 7. Reshape back to original shape - sample = sample.reshape(batch_size, num_frames, *sample.shape[1:]) - - if not return_dict: - return (sample,) - - return UNetSpatioTemporalConditionOutput(sample=sample) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_stable_cascade.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_stable_cascade.py deleted file mode 100644 index e79ce8ac1da4038d90841fa2a69fc8eb89d37f3a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/unet_stable_cascade.py +++ /dev/null @@ -1,584 +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 math -from dataclasses import dataclass -from typing import Optional, Tuple, Union - -import numpy as np -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import FromOriginalModelMixin -from ...utils import BaseOutput -from ..attention_processor import Attention -from ..modeling_utils import ModelMixin - - -# Copied from diffusers.pipelines.wuerstchen.modeling_wuerstchen_common.WuerstchenLayerNorm with WuerstchenLayerNorm -> SDCascadeLayerNorm -class SDCascadeLayerNorm(nn.LayerNorm): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def forward(self, x): - x = x.permute(0, 2, 3, 1) - x = super().forward(x) - return x.permute(0, 3, 1, 2) - - -class SDCascadeTimestepBlock(nn.Module): - def __init__(self, c, c_timestep, conds=[]): - super().__init__() - - self.mapper = nn.Linear(c_timestep, c * 2) - self.conds = conds - for cname in conds: - setattr(self, f"mapper_{cname}", nn.Linear(c_timestep, c * 2)) - - def forward(self, x, t): - t = t.chunk(len(self.conds) + 1, dim=1) - a, b = self.mapper(t[0])[:, :, None, None].chunk(2, dim=1) - for i, c in enumerate(self.conds): - ac, bc = getattr(self, f"mapper_{c}")(t[i + 1])[:, :, None, None].chunk(2, dim=1) - a, b = a + ac, b + bc - return x * (1 + a) + b - - -class SDCascadeResBlock(nn.Module): - def __init__(self, c, c_skip=0, kernel_size=3, dropout=0.0): - super().__init__() - self.depthwise = nn.Conv2d(c, c, kernel_size=kernel_size, padding=kernel_size // 2, groups=c) - self.norm = SDCascadeLayerNorm(c, elementwise_affine=False, eps=1e-6) - self.channelwise = nn.Sequential( - nn.Linear(c + c_skip, c * 4), - nn.GELU(), - GlobalResponseNorm(c * 4), - nn.Dropout(dropout), - nn.Linear(c * 4, c), - ) - - def forward(self, x, x_skip=None): - x_res = x - x = self.norm(self.depthwise(x)) - if x_skip is not None: - x = torch.cat([x, x_skip], dim=1) - x = self.channelwise(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) - return x + x_res - - -# from https://github.com/facebookresearch/ConvNeXt-V2/blob/3608f67cc1dae164790c5d0aead7bf2d73d9719b/models/utils.py#L105 -class GlobalResponseNorm(nn.Module): - def __init__(self, dim): - super().__init__() - self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) - self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) - - def forward(self, x): - agg_norm = torch.norm(x, p=2, dim=(1, 2), keepdim=True) - stand_div_norm = agg_norm / (agg_norm.mean(dim=-1, keepdim=True) + 1e-6) - return self.gamma * (x * stand_div_norm) + self.beta + x - - -class SDCascadeAttnBlock(nn.Module): - def __init__(self, c, c_cond, nhead, self_attn=True, dropout=0.0): - super().__init__() - - self.self_attn = self_attn - self.norm = SDCascadeLayerNorm(c, elementwise_affine=False, eps=1e-6) - self.attention = Attention(query_dim=c, heads=nhead, dim_head=c // nhead, dropout=dropout, bias=True) - self.kv_mapper = nn.Sequential(nn.SiLU(), nn.Linear(c_cond, c)) - - def forward(self, x, kv): - kv = self.kv_mapper(kv) - norm_x = self.norm(x) - if self.self_attn: - batch_size, channel, _, _ = x.shape - kv = torch.cat([norm_x.view(batch_size, channel, -1).transpose(1, 2), kv], dim=1) - x = x + self.attention(norm_x, encoder_hidden_states=kv) - return x - - -class UpDownBlock2d(nn.Module): - def __init__(self, in_channels, out_channels, mode, enabled=True): - super().__init__() - if mode not in ["up", "down"]: - raise ValueError(f"{mode} not supported") - interpolation = ( - nn.Upsample(scale_factor=2 if mode == "up" else 0.5, mode="bilinear", align_corners=True) - if enabled - else nn.Identity() - ) - mapping = nn.Conv2d(in_channels, out_channels, kernel_size=1) - self.blocks = nn.ModuleList([interpolation, mapping] if mode == "up" else [mapping, interpolation]) - - def forward(self, x): - for block in self.blocks: - x = block(x) - return x - - -@dataclass -class StableCascadeUNetOutput(BaseOutput): - sample: torch.Tensor = None - - -class StableCascadeUNet(ModelMixin, ConfigMixin, FromOriginalModelMixin): - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - in_channels: int = 16, - out_channels: int = 16, - 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]] = ( - 1, - 1, - ), - up_blocks_repeat_mappers: Optional[Tuple[int]] = (1, 1), - block_types_per_layer: Tuple[Tuple[str]] = ( - ("SDCascadeResBlock", "SDCascadeTimestepBlock", "SDCascadeAttnBlock"), - ("SDCascadeResBlock", "SDCascadeTimestepBlock", "SDCascadeAttnBlock"), - ), - clip_text_in_channels: Optional[int] = None, - clip_text_pooled_in_channels=1280, - clip_image_in_channels: Optional[int] = None, - clip_seq=4, - effnet_in_channels: Optional[int] = None, - pixel_mapper_in_channels: Optional[int] = 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, - ): - """ - - Parameters: - in_channels (`int`, defaults to 16): - Number of channels in the input sample. - out_channels (`int`, defaults to 16): - Number of channels in the output sample. - timestep_ratio_embedding_dim (`int`, defaults to 64): - Dimension of the projected time embedding. - patch_size (`int`, defaults to 1): - 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)): - 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]): - Number of layers in each down block. - 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]): - Number of 1x1 Convolutional layers to repeat in each down block. - 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, - defaults to ( - ("SDCascadeResBlock", "SDCascadeTimestepBlock", "SDCascadeAttnBlock"), ("SDCascadeResBlock", - "SDCascadeTimestepBlock", "SDCascadeAttnBlock") - ): Block types used in each layer of the up/down blocks. - clip_text_in_channels (`int`, *optional*, defaults to `None`): - Number of input channels for CLIP based text conditioning. - clip_text_pooled_in_channels (`int`, *optional*, defaults to 1280): - Number of input channels for pooled CLIP text embeddings. - clip_image_in_channels (`int`, *optional*): - Number of input channels for CLIP based image conditioning. - clip_seq (`int`, *optional*, defaults to 4): - effnet_in_channels (`int`, *optional*, defaults to `None`): - Number of input channels for effnet conditioning. - pixel_mapper_in_channels (`int`, defaults to `None`): - 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 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")): - 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 - """ - - super().__init__() - - if len(block_out_channels) != len(down_num_layers_per_block): - raise ValueError( - f"Number of elements in `down_num_layers_per_block` must match the length of `block_out_channels`: {len(block_out_channels)}" - ) - - elif len(block_out_channels) != len(up_num_layers_per_block): - raise ValueError( - f"Number of elements in `up_num_layers_per_block` must match the length of `block_out_channels`: {len(block_out_channels)}" - ) - - elif len(block_out_channels) != len(down_blocks_repeat_mappers): - raise ValueError( - f"Number of elements in `down_blocks_repeat_mappers` must match the length of `block_out_channels`: {len(block_out_channels)}" - ) - - elif len(block_out_channels) != len(up_blocks_repeat_mappers): - raise ValueError( - f"Number of elements in `up_blocks_repeat_mappers` must match the length of `block_out_channels`: {len(block_out_channels)}" - ) - - elif len(block_out_channels) != len(block_types_per_layer): - raise ValueError( - f"Number of elements in `block_types_per_layer` must match the length of `block_out_channels`: {len(block_out_channels)}" - ) - - if isinstance(dropout, float): - dropout = (dropout,) * len(block_out_channels) - if isinstance(self_attn, bool): - self_attn = (self_attn,) * len(block_out_channels) - - # CONDITIONING - if effnet_in_channels is not None: - self.effnet_mapper = nn.Sequential( - nn.Conv2d(effnet_in_channels, block_out_channels[0] * 4, kernel_size=1), - nn.GELU(), - nn.Conv2d(block_out_channels[0] * 4, block_out_channels[0], kernel_size=1), - SDCascadeLayerNorm(block_out_channels[0], elementwise_affine=False, eps=1e-6), - ) - if pixel_mapper_in_channels is not None: - self.pixels_mapper = nn.Sequential( - nn.Conv2d(pixel_mapper_in_channels, block_out_channels[0] * 4, kernel_size=1), - nn.GELU(), - nn.Conv2d(block_out_channels[0] * 4, block_out_channels[0], kernel_size=1), - SDCascadeLayerNorm(block_out_channels[0], elementwise_affine=False, eps=1e-6), - ) - - self.clip_txt_pooled_mapper = nn.Linear(clip_text_pooled_in_channels, conditioning_dim * clip_seq) - if clip_text_in_channels is not None: - self.clip_txt_mapper = nn.Linear(clip_text_in_channels, conditioning_dim) - if clip_image_in_channels is not None: - self.clip_img_mapper = nn.Linear(clip_image_in_channels, conditioning_dim * clip_seq) - self.clip_norm = nn.LayerNorm(conditioning_dim, elementwise_affine=False, eps=1e-6) - - self.embedding = nn.Sequential( - nn.PixelUnshuffle(patch_size), - nn.Conv2d(in_channels * (patch_size**2), block_out_channels[0], kernel_size=1), - SDCascadeLayerNorm(block_out_channels[0], elementwise_affine=False, eps=1e-6), - ) - - def get_block(block_type, in_channels, nhead, c_skip=0, dropout=0, self_attn=True): - if block_type == "SDCascadeResBlock": - return SDCascadeResBlock(in_channels, c_skip, kernel_size=kernel_size, dropout=dropout) - elif block_type == "SDCascadeAttnBlock": - return SDCascadeAttnBlock(in_channels, conditioning_dim, nhead, self_attn=self_attn, dropout=dropout) - elif block_type == "SDCascadeTimestepBlock": - return SDCascadeTimestepBlock( - in_channels, timestep_ratio_embedding_dim, conds=timestep_conditioning_type - ) - else: - raise ValueError(f"Block type {block_type} not supported") - - # BLOCKS - # -- down blocks - self.down_blocks = nn.ModuleList() - self.down_downscalers = nn.ModuleList() - self.down_repeat_mappers = nn.ModuleList() - for i in range(len(block_out_channels)): - if i > 0: - self.down_downscalers.append( - nn.Sequential( - SDCascadeLayerNorm(block_out_channels[i - 1], elementwise_affine=False, eps=1e-6), - UpDownBlock2d( - block_out_channels[i - 1], block_out_channels[i], mode="down", enabled=switch_level[i - 1] - ) - if switch_level is not None - else nn.Conv2d(block_out_channels[i - 1], block_out_channels[i], kernel_size=2, stride=2), - ) - ) - else: - self.down_downscalers.append(nn.Identity()) - - down_block = nn.ModuleList() - for _ in range(down_num_layers_per_block[i]): - for block_type in block_types_per_layer[i]: - block = get_block( - block_type, - block_out_channels[i], - num_attention_heads[i], - dropout=dropout[i], - self_attn=self_attn[i], - ) - down_block.append(block) - self.down_blocks.append(down_block) - - if down_blocks_repeat_mappers is not None: - block_repeat_mappers = nn.ModuleList() - for _ in range(down_blocks_repeat_mappers[i] - 1): - block_repeat_mappers.append(nn.Conv2d(block_out_channels[i], block_out_channels[i], kernel_size=1)) - self.down_repeat_mappers.append(block_repeat_mappers) - - # -- up blocks - self.up_blocks = nn.ModuleList() - self.up_upscalers = nn.ModuleList() - self.up_repeat_mappers = nn.ModuleList() - for i in reversed(range(len(block_out_channels))): - if i > 0: - self.up_upscalers.append( - nn.Sequential( - SDCascadeLayerNorm(block_out_channels[i], elementwise_affine=False, eps=1e-6), - UpDownBlock2d( - block_out_channels[i], block_out_channels[i - 1], mode="up", enabled=switch_level[i - 1] - ) - if switch_level is not None - else nn.ConvTranspose2d( - block_out_channels[i], block_out_channels[i - 1], kernel_size=2, stride=2 - ), - ) - ) - else: - self.up_upscalers.append(nn.Identity()) - - up_block = nn.ModuleList() - for j in range(up_num_layers_per_block[::-1][i]): - for k, block_type in enumerate(block_types_per_layer[i]): - c_skip = block_out_channels[i] if i < len(block_out_channels) - 1 and j == k == 0 else 0 - block = get_block( - block_type, - block_out_channels[i], - num_attention_heads[i], - c_skip=c_skip, - dropout=dropout[i], - self_attn=self_attn[i], - ) - up_block.append(block) - self.up_blocks.append(up_block) - - if up_blocks_repeat_mappers is not None: - block_repeat_mappers = nn.ModuleList() - for _ in range(up_blocks_repeat_mappers[::-1][i] - 1): - block_repeat_mappers.append(nn.Conv2d(block_out_channels[i], block_out_channels[i], kernel_size=1)) - self.up_repeat_mappers.append(block_repeat_mappers) - - # OUTPUT - self.clf = nn.Sequential( - SDCascadeLayerNorm(block_out_channels[0], elementwise_affine=False, eps=1e-6), - nn.Conv2d(block_out_channels[0], out_channels * (patch_size**2), kernel_size=1), - nn.PixelShuffle(patch_size), - ) - - self.gradient_checkpointing = False - - def _init_weights(self, m): - if isinstance(m, (nn.Conv2d, nn.Linear)): - torch.nn.init.xavier_uniform_(m.weight) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - - nn.init.normal_(self.clip_txt_pooled_mapper.weight, std=0.02) - nn.init.normal_(self.clip_txt_mapper.weight, std=0.02) if hasattr(self, "clip_txt_mapper") else None - nn.init.normal_(self.clip_img_mapper.weight, std=0.02) if hasattr(self, "clip_img_mapper") else None - - if hasattr(self, "effnet_mapper"): - nn.init.normal_(self.effnet_mapper[0].weight, std=0.02) # conditionings - nn.init.normal_(self.effnet_mapper[2].weight, std=0.02) # conditionings - - if hasattr(self, "pixels_mapper"): - nn.init.normal_(self.pixels_mapper[0].weight, std=0.02) # conditionings - nn.init.normal_(self.pixels_mapper[2].weight, std=0.02) # conditionings - - torch.nn.init.xavier_uniform_(self.embedding[1].weight, 0.02) # inputs - nn.init.constant_(self.clf[1].weight, 0) # outputs - - # blocks - for level_block in self.down_blocks + self.up_blocks: - for block in level_block: - if isinstance(block, SDCascadeResBlock): - block.channelwise[-1].weight.data *= np.sqrt(1 / sum(self.config.blocks[0])) - elif isinstance(block, SDCascadeTimestepBlock): - nn.init.constant_(block.mapper.weight, 0) - - def get_timestep_ratio_embedding(self, timestep_ratio, max_positions=10000): - r = timestep_ratio * max_positions - half_dim = self.config.timestep_ratio_embedding_dim // 2 - - emb = math.log(max_positions) / (half_dim - 1) - emb = torch.arange(half_dim, device=r.device).float().mul(-emb).exp() - emb = r[:, None] * emb[None, :] - emb = torch.cat([emb.sin(), emb.cos()], dim=1) - - if self.config.timestep_ratio_embedding_dim % 2 == 1: # zero pad - emb = nn.functional.pad(emb, (0, 1), mode="constant") - - return emb.to(dtype=r.dtype) - - def get_clip_embeddings(self, clip_txt_pooled, clip_txt=None, clip_img=None): - if len(clip_txt_pooled.shape) == 2: - clip_txt_pool = clip_txt_pooled.unsqueeze(1) - clip_txt_pool = self.clip_txt_pooled_mapper(clip_txt_pooled).view( - clip_txt_pooled.size(0), clip_txt_pooled.size(1) * self.config.clip_seq, -1 - ) - if clip_txt is not None and clip_img is not None: - clip_txt = self.clip_txt_mapper(clip_txt) - if len(clip_img.shape) == 2: - clip_img = clip_img.unsqueeze(1) - clip_img = self.clip_img_mapper(clip_img).view( - clip_img.size(0), clip_img.size(1) * self.config.clip_seq, -1 - ) - clip = torch.cat([clip_txt, clip_txt_pool, clip_img], dim=1) - else: - clip = clip_txt_pool - return self.clip_norm(clip) - - def _down_encode(self, x, r_embed, clip): - level_outputs = [] - block_group = zip(self.down_blocks, self.down_downscalers, self.down_repeat_mappers) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - for down_block, downscaler, repmap in block_group: - x = downscaler(x) - for i in range(len(repmap) + 1): - for block in down_block: - if isinstance(block, SDCascadeResBlock): - x = self._gradient_checkpointing_func(block, x) - elif isinstance(block, SDCascadeAttnBlock): - x = self._gradient_checkpointing_func(block, x, clip) - elif isinstance(block, SDCascadeTimestepBlock): - x = self._gradient_checkpointing_func(block, x, r_embed) - else: - x = self._gradient_checkpointing_func(block) - if i < len(repmap): - x = repmap[i](x) - level_outputs.insert(0, x) - else: - for down_block, downscaler, repmap in block_group: - x = downscaler(x) - for i in range(len(repmap) + 1): - for block in down_block: - if isinstance(block, SDCascadeResBlock): - x = block(x) - elif isinstance(block, SDCascadeAttnBlock): - x = block(x, clip) - elif isinstance(block, SDCascadeTimestepBlock): - x = block(x, r_embed) - else: - x = block(x) - if i < len(repmap): - x = repmap[i](x) - level_outputs.insert(0, x) - return level_outputs - - def _up_decode(self, level_outputs, r_embed, clip): - x = level_outputs[0] - block_group = zip(self.up_blocks, self.up_upscalers, self.up_repeat_mappers) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - for i, (up_block, upscaler, repmap) in enumerate(block_group): - for j in range(len(repmap) + 1): - for k, block in enumerate(up_block): - if isinstance(block, SDCascadeResBlock): - skip = level_outputs[i] if k == 0 and i > 0 else None - if skip is not None and (x.size(-1) != skip.size(-1) or x.size(-2) != skip.size(-2)): - orig_type = x.dtype - x = torch.nn.functional.interpolate( - x.float(), skip.shape[-2:], mode="bilinear", align_corners=True - ) - x = x.to(orig_type) - x = self._gradient_checkpointing_func(block, x, skip) - elif isinstance(block, SDCascadeAttnBlock): - x = self._gradient_checkpointing_func(block, x, clip) - elif isinstance(block, SDCascadeTimestepBlock): - x = self._gradient_checkpointing_func(block, x, r_embed) - else: - x = self._gradient_checkpointing_func(block, x) - if j < len(repmap): - x = repmap[j](x) - x = upscaler(x) - else: - for i, (up_block, upscaler, repmap) in enumerate(block_group): - for j in range(len(repmap) + 1): - for k, block in enumerate(up_block): - if isinstance(block, SDCascadeResBlock): - skip = level_outputs[i] if k == 0 and i > 0 else None - if skip is not None and (x.size(-1) != skip.size(-1) or x.size(-2) != skip.size(-2)): - orig_type = x.dtype - x = torch.nn.functional.interpolate( - x.float(), skip.shape[-2:], mode="bilinear", align_corners=True - ) - x = x.to(orig_type) - x = block(x, skip) - elif isinstance(block, SDCascadeAttnBlock): - x = block(x, clip) - elif isinstance(block, SDCascadeTimestepBlock): - x = block(x, r_embed) - else: - x = block(x) - if j < len(repmap): - x = repmap[j](x) - x = upscaler(x) - return x - - def forward( - self, - sample, - timestep_ratio, - clip_text_pooled, - clip_text=None, - clip_img=None, - effnet=None, - pixels=None, - sca=None, - crp=None, - return_dict=True, - ): - if pixels is None: - pixels = sample.new_zeros(sample.size(0), 3, 8, 8) - - # Process the conditioning embeddings - timestep_ratio_embed = self.get_timestep_ratio_embedding(timestep_ratio) - for c in self.config.timestep_conditioning_type: - if c == "sca": - cond = sca - elif c == "crp": - cond = crp - else: - cond = None - t_cond = cond or torch.zeros_like(timestep_ratio) - timestep_ratio_embed = torch.cat([timestep_ratio_embed, self.get_timestep_ratio_embedding(t_cond)], dim=1) - clip = self.get_clip_embeddings(clip_txt_pooled=clip_text_pooled, clip_txt=clip_text, clip_img=clip_img) - - # Model Blocks - x = self.embedding(sample) - if hasattr(self, "effnet_mapper") and effnet is not None: - x = x + self.effnet_mapper( - nn.functional.interpolate(effnet, size=x.shape[-2:], mode="bilinear", align_corners=True) - ) - if hasattr(self, "pixels_mapper"): - x = x + nn.functional.interpolate( - self.pixels_mapper(pixels), size=x.shape[-2:], mode="bilinear", align_corners=True - ) - level_outputs = self._down_encode(x, timestep_ratio_embed, clip) - x = self._up_decode(level_outputs, timestep_ratio_embed, clip) - sample = self.clf(x) - - if not return_dict: - return (sample,) - return StableCascadeUNetOutput(sample=sample) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/unets/uvit_2d.py b/venv/lib/python3.10/site-packages/diffusers/models/unets/uvit_2d.py deleted file mode 100644 index 94b39c84f0550369d432346fd8ec48e4c778f2ed..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/unets/uvit_2d.py +++ /dev/null @@ -1,467 +0,0 @@ -# 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. - -from typing import Dict, Union - -import torch -import torch.nn.functional as F -from torch import nn -from torch.utils.checkpoint import checkpoint - -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin -from ..attention import BasicTransformerBlock, SkipFFTransformerBlock -from ..attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - AttentionProcessor, - AttnAddedKVProcessor, - AttnProcessor, -) -from ..embeddings import TimestepEmbedding, get_timestep_embedding -from ..modeling_utils import ModelMixin -from ..normalization import GlobalResponseNorm, RMSNorm -from ..resnet import Downsample2D, Upsample2D - - -class UVit2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin): - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - # global config - hidden_size: int = 1024, - use_bias: bool = False, - hidden_dropout: float = 0.0, - # conditioning dimensions - cond_embed_dim: int = 768, - micro_cond_encode_dim: int = 256, - micro_cond_embed_dim: int = 1280, - encoder_hidden_size: int = 768, - # num tokens - vocab_size: int = 8256, # codebook_size + 1 (for the mask token) rounded - codebook_size: int = 8192, - # `UVit2DConvEmbed` - in_channels: int = 768, - block_out_channels: int = 768, - num_res_blocks: int = 3, - downsample: bool = False, - upsample: bool = False, - block_num_heads: int = 12, - # `TransformerLayer` - num_hidden_layers: int = 22, - num_attention_heads: int = 16, - # `Attention` - attention_dropout: float = 0.0, - # `FeedForward` - intermediate_size: int = 2816, - # `Norm` - layer_norm_eps: float = 1e-6, - ln_elementwise_affine: bool = True, - sample_size: int = 64, - ): - super().__init__() - - self.encoder_proj = nn.Linear(encoder_hidden_size, hidden_size, bias=use_bias) - self.encoder_proj_layer_norm = RMSNorm(hidden_size, layer_norm_eps, ln_elementwise_affine) - - self.embed = UVit2DConvEmbed( - in_channels, block_out_channels, vocab_size, ln_elementwise_affine, layer_norm_eps, use_bias - ) - - self.cond_embed = TimestepEmbedding( - micro_cond_embed_dim + cond_embed_dim, hidden_size, sample_proj_bias=use_bias - ) - - self.down_block = UVitBlock( - block_out_channels, - num_res_blocks, - hidden_size, - hidden_dropout, - ln_elementwise_affine, - layer_norm_eps, - use_bias, - block_num_heads, - attention_dropout, - downsample, - False, - ) - - self.project_to_hidden_norm = RMSNorm(block_out_channels, layer_norm_eps, ln_elementwise_affine) - self.project_to_hidden = nn.Linear(block_out_channels, hidden_size, bias=use_bias) - - self.transformer_layers = nn.ModuleList( - [ - BasicTransformerBlock( - dim=hidden_size, - num_attention_heads=num_attention_heads, - attention_head_dim=hidden_size // num_attention_heads, - dropout=hidden_dropout, - cross_attention_dim=hidden_size, - attention_bias=use_bias, - norm_type="ada_norm_continuous", - ada_norm_continous_conditioning_embedding_dim=hidden_size, - norm_elementwise_affine=ln_elementwise_affine, - norm_eps=layer_norm_eps, - ada_norm_bias=use_bias, - ff_inner_dim=intermediate_size, - ff_bias=use_bias, - attention_out_bias=use_bias, - ) - for _ in range(num_hidden_layers) - ] - ) - - self.project_from_hidden_norm = RMSNorm(hidden_size, layer_norm_eps, ln_elementwise_affine) - self.project_from_hidden = nn.Linear(hidden_size, block_out_channels, bias=use_bias) - - self.up_block = UVitBlock( - block_out_channels, - num_res_blocks, - hidden_size, - hidden_dropout, - ln_elementwise_affine, - layer_norm_eps, - use_bias, - block_num_heads, - attention_dropout, - downsample=False, - upsample=upsample, - ) - - self.mlm_layer = ConvMlmLayer( - block_out_channels, in_channels, use_bias, ln_elementwise_affine, layer_norm_eps, codebook_size - ) - - self.gradient_checkpointing = False - - 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) - - micro_cond_embeds = get_timestep_embedding( - micro_conds.flatten(), self.config.micro_cond_encode_dim, flip_sin_to_cos=True, downscale_freq_shift=0 - ) - - micro_cond_embeds = micro_cond_embeds.reshape((input_ids.shape[0], -1)) - - pooled_text_emb = torch.cat([pooled_text_emb, micro_cond_embeds], dim=1) - pooled_text_emb = pooled_text_emb.to(dtype=self.dtype) - pooled_text_emb = self.cond_embed(pooled_text_emb).to(encoder_hidden_states.dtype) - - hidden_states = self.embed(input_ids) - - hidden_states = self.down_block( - hidden_states, - pooled_text_emb=pooled_text_emb, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - ) - - batch_size, channels, height, width = hidden_states.shape - hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels) - - hidden_states = self.project_to_hidden_norm(hidden_states) - hidden_states = self.project_to_hidden(hidden_states) - - for layer in self.transformer_layers: - if torch.is_grad_enabled() and self.gradient_checkpointing: - - def layer_(*args): - return checkpoint(layer, *args) - - else: - layer_ = layer - - hidden_states = layer_( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs={"pooled_text_emb": pooled_text_emb}, - ) - - hidden_states = self.project_from_hidden_norm(hidden_states) - hidden_states = self.project_from_hidden(hidden_states) - - hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2) - - hidden_states = self.up_block( - hidden_states, - pooled_text_emb=pooled_text_emb, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - ) - - logits = self.mlm_layer(hidden_states) - - return logits - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - self.set_attn_processor(processor) - - -class UVit2DConvEmbed(nn.Module): - def __init__(self, in_channels, block_out_channels, vocab_size, elementwise_affine, eps, bias): - super().__init__() - self.embeddings = nn.Embedding(vocab_size, in_channels) - self.layer_norm = RMSNorm(in_channels, eps, elementwise_affine) - self.conv = nn.Conv2d(in_channels, block_out_channels, kernel_size=1, bias=bias) - - def forward(self, input_ids): - embeddings = self.embeddings(input_ids) - embeddings = self.layer_norm(embeddings) - embeddings = embeddings.permute(0, 3, 1, 2) - embeddings = self.conv(embeddings) - return embeddings - - -class UVitBlock(nn.Module): - def __init__( - self, - channels, - num_res_blocks: int, - hidden_size, - hidden_dropout, - ln_elementwise_affine, - layer_norm_eps, - use_bias, - block_num_heads, - attention_dropout, - downsample: bool, - upsample: bool, - ): - super().__init__() - - if downsample: - self.downsample = Downsample2D( - channels, - use_conv=True, - padding=0, - name="Conv2d_0", - kernel_size=2, - norm_type="rms_norm", - eps=layer_norm_eps, - elementwise_affine=ln_elementwise_affine, - bias=use_bias, - ) - else: - self.downsample = None - - self.res_blocks = nn.ModuleList( - [ - ConvNextBlock( - channels, - layer_norm_eps, - ln_elementwise_affine, - use_bias, - hidden_dropout, - hidden_size, - ) - for i in range(num_res_blocks) - ] - ) - - self.attention_blocks = nn.ModuleList( - [ - SkipFFTransformerBlock( - channels, - block_num_heads, - channels // block_num_heads, - hidden_size, - use_bias, - attention_dropout, - channels, - attention_bias=use_bias, - attention_out_bias=use_bias, - ) - for _ in range(num_res_blocks) - ] - ) - - if upsample: - self.upsample = Upsample2D( - channels, - use_conv_transpose=True, - kernel_size=2, - padding=0, - name="conv", - norm_type="rms_norm", - eps=layer_norm_eps, - elementwise_affine=ln_elementwise_affine, - bias=use_bias, - interpolate=False, - ) - else: - self.upsample = None - - def forward(self, x, pooled_text_emb, encoder_hidden_states, cross_attention_kwargs): - if self.downsample is not None: - x = self.downsample(x) - - for res_block, attention_block in zip(self.res_blocks, self.attention_blocks): - x = res_block(x, pooled_text_emb) - - batch_size, channels, height, width = x.shape - x = x.view(batch_size, channels, height * width).permute(0, 2, 1) - x = attention_block( - x, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs - ) - x = x.permute(0, 2, 1).view(batch_size, channels, height, width) - - if self.upsample is not None: - x = self.upsample(x) - - return x - - -class ConvNextBlock(nn.Module): - def __init__( - self, channels, layer_norm_eps, ln_elementwise_affine, use_bias, hidden_dropout, hidden_size, res_ffn_factor=4 - ): - super().__init__() - self.depthwise = nn.Conv2d( - channels, - channels, - kernel_size=3, - padding=1, - groups=channels, - bias=use_bias, - ) - self.norm = RMSNorm(channels, layer_norm_eps, ln_elementwise_affine) - self.channelwise_linear_1 = nn.Linear(channels, int(channels * res_ffn_factor), bias=use_bias) - self.channelwise_act = nn.GELU() - self.channelwise_norm = GlobalResponseNorm(int(channels * res_ffn_factor)) - self.channelwise_linear_2 = nn.Linear(int(channels * res_ffn_factor), channels, bias=use_bias) - self.channelwise_dropout = nn.Dropout(hidden_dropout) - self.cond_embeds_mapper = nn.Linear(hidden_size, channels * 2, use_bias) - - def forward(self, x, cond_embeds): - x_res = x - - x = self.depthwise(x) - - x = x.permute(0, 2, 3, 1) - x = self.norm(x) - - x = self.channelwise_linear_1(x) - x = self.channelwise_act(x) - x = self.channelwise_norm(x) - x = self.channelwise_linear_2(x) - x = self.channelwise_dropout(x) - - x = x.permute(0, 3, 1, 2) - - x = x + x_res - - scale, shift = self.cond_embeds_mapper(F.silu(cond_embeds)).chunk(2, dim=1) - x = x * (1 + scale[:, :, None, None]) + shift[:, :, None, None] - - return x - - -class ConvMlmLayer(nn.Module): - def __init__( - self, - block_out_channels: int, - in_channels: int, - use_bias: bool, - ln_elementwise_affine: bool, - layer_norm_eps: float, - codebook_size: int, - ): - super().__init__() - self.conv1 = nn.Conv2d(block_out_channels, in_channels, kernel_size=1, bias=use_bias) - self.layer_norm = RMSNorm(in_channels, layer_norm_eps, ln_elementwise_affine) - self.conv2 = nn.Conv2d(in_channels, codebook_size, kernel_size=1, bias=use_bias) - - def forward(self, hidden_states): - hidden_states = self.conv1(hidden_states) - hidden_states = self.layer_norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) - logits = self.conv2(hidden_states) - return logits diff --git a/venv/lib/python3.10/site-packages/diffusers/models/upsampling.py b/venv/lib/python3.10/site-packages/diffusers/models/upsampling.py deleted file mode 100644 index 8a47c69f1264647afc15ae74d4aaa381ac3b8883..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/upsampling.py +++ /dev/null @@ -1,517 +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 - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from ..utils import deprecate -from ..utils.import_utils import is_torch_version -from .normalization import RMSNorm - - -class Upsample1D(nn.Module): - """A 1D upsampling layer with an optional convolution. - - Parameters: - channels (`int`): - number of channels in the inputs and outputs. - use_conv (`bool`, default `False`): - option to use a convolution. - use_conv_transpose (`bool`, default `False`): - option to use a convolution transpose. - out_channels (`int`, optional): - number of output channels. Defaults to `channels`. - name (`str`, default `conv`): - name of the upsampling 1D layer. - """ - - def __init__( - self, - channels: int, - use_conv: bool = False, - use_conv_transpose: bool = False, - out_channels: Optional[int] = None, - name: str = "conv", - ): - super().__init__() - self.channels = channels - self.out_channels = out_channels or channels - self.use_conv = use_conv - self.use_conv_transpose = use_conv_transpose - self.name = name - - self.conv = None - if use_conv_transpose: - self.conv = nn.ConvTranspose1d(channels, self.out_channels, 4, 2, 1) - elif use_conv: - self.conv = nn.Conv1d(self.channels, self.out_channels, 3, padding=1) - - def forward(self, inputs: torch.Tensor) -> torch.Tensor: - assert inputs.shape[1] == self.channels - if self.use_conv_transpose: - return self.conv(inputs) - - outputs = F.interpolate(inputs, scale_factor=2.0, mode="nearest") - - if self.use_conv: - outputs = self.conv(outputs) - - return outputs - - -class Upsample2D(nn.Module): - """A 2D upsampling layer with an optional convolution. - - Parameters: - channels (`int`): - number of channels in the inputs and outputs. - use_conv (`bool`, default `False`): - option to use a convolution. - use_conv_transpose (`bool`, default `False`): - option to use a convolution transpose. - out_channels (`int`, optional): - number of output channels. Defaults to `channels`. - name (`str`, default `conv`): - name of the upsampling 2D layer. - """ - - def __init__( - self, - channels: int, - use_conv: bool = False, - use_conv_transpose: bool = False, - out_channels: Optional[int] = None, - name: str = "conv", - kernel_size: Optional[int] = None, - padding=1, - norm_type=None, - eps=None, - elementwise_affine=None, - bias=True, - interpolate=True, - ): - super().__init__() - self.channels = channels - self.out_channels = out_channels or channels - self.use_conv = use_conv - self.use_conv_transpose = use_conv_transpose - self.name = name - self.interpolate = interpolate - - if norm_type == "ln_norm": - self.norm = nn.LayerNorm(channels, eps, elementwise_affine) - elif norm_type == "rms_norm": - self.norm = RMSNorm(channels, eps, elementwise_affine) - elif norm_type is None: - self.norm = None - else: - raise ValueError(f"unknown norm_type: {norm_type}") - - conv = None - if use_conv_transpose: - if kernel_size is None: - kernel_size = 4 - conv = nn.ConvTranspose2d( - channels, self.out_channels, kernel_size=kernel_size, stride=2, padding=padding, bias=bias - ) - elif use_conv: - if kernel_size is None: - kernel_size = 3 - conv = nn.Conv2d(self.channels, self.out_channels, kernel_size=kernel_size, padding=padding, bias=bias) - - # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed - if name == "conv": - self.conv = conv - else: - self.Conv2d_0 = conv - - def forward(self, hidden_states: torch.Tensor, output_size: Optional[int] = 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) - - assert hidden_states.shape[1] == self.channels - - if self.norm is not None: - hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) - - if self.use_conv_transpose: - return self.conv(hidden_states) - - # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16 until PyTorch 2.1 - # https://github.com/pytorch/pytorch/issues/86679#issuecomment-1783978767 - dtype = hidden_states.dtype - if dtype == torch.bfloat16 and is_torch_version("<", "2.1"): - hidden_states = hidden_states.to(torch.float32) - - # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 - if hidden_states.shape[0] >= 64: - hidden_states = hidden_states.contiguous() - - # if `output_size` is passed we force the interpolation output - # size and do not make use of `scale_factor=2` - if self.interpolate: - # upsample_nearest_nhwc also fails when the number of output elements is large - # https://github.com/pytorch/pytorch/issues/141831 - scale_factor = ( - 2 if output_size is None else max([f / s for f, s in zip(output_size, hidden_states.shape[-2:])]) - ) - if hidden_states.numel() * scale_factor > pow(2, 31): - hidden_states = hidden_states.contiguous() - - if output_size is None: - hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest") - else: - hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest") - - # Cast back to original dtype - if dtype == torch.bfloat16 and is_torch_version("<", "2.1"): - hidden_states = hidden_states.to(dtype) - - # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed - if self.use_conv: - if self.name == "conv": - hidden_states = self.conv(hidden_states) - else: - hidden_states = self.Conv2d_0(hidden_states) - - return hidden_states - - -class FirUpsample2D(nn.Module): - """A 2D FIR upsampling layer with an optional convolution. - - Parameters: - channels (`int`, optional): - number of channels in the inputs and outputs. - use_conv (`bool`, default `False`): - option to use a convolution. - out_channels (`int`, optional): - number of output channels. Defaults to `channels`. - fir_kernel (`tuple`, default `(1, 3, 3, 1)`): - kernel for the FIR filter. - """ - - def __init__( - self, - channels: Optional[int] = None, - out_channels: Optional[int] = None, - use_conv: bool = False, - fir_kernel: Tuple[int, int, int, int] = (1, 3, 3, 1), - ): - super().__init__() - out_channels = out_channels if out_channels else channels - if use_conv: - self.Conv2d_0 = nn.Conv2d(channels, out_channels, kernel_size=3, stride=1, padding=1) - self.use_conv = use_conv - self.fir_kernel = fir_kernel - self.out_channels = out_channels - - def _upsample_2d( - self, - hidden_states: torch.Tensor, - weight: Optional[torch.Tensor] = None, - kernel: Optional[torch.Tensor] = None, - factor: int = 2, - gain: float = 1, - ) -> torch.Tensor: - """Fused `upsample_2d()` followed by `Conv2d()`. - - Padding is performed only once at the beginning, not between the operations. The fused op is considerably more - efficient than performing the same calculation using standard TensorFlow ops. It supports gradients of - arbitrary order. - - Args: - hidden_states (`torch.Tensor`): - Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`. - weight (`torch.Tensor`, *optional*): - Weight tensor of the shape `[filterH, filterW, inChannels, outChannels]`. Grouped convolution can be - performed by `inChannels = x.shape[0] // numGroups`. - kernel (`torch.Tensor`, *optional*): - FIR filter of the shape `[firH, firW]` or `[firN]` (separable). The default is `[1] * factor`, which - corresponds to nearest-neighbor upsampling. - factor (`int`, *optional*): Integer upsampling factor (default: 2). - gain (`float`, *optional*): Scaling factor for signal magnitude (default: 1.0). - - Returns: - output (`torch.Tensor`): - Tensor of the shape `[N, C, H * factor, W * factor]` or `[N, H * factor, W * factor, C]`, and same - datatype as `hidden_states`. - """ - - assert isinstance(factor, int) and factor >= 1 - - # Setup filter kernel. - if kernel is None: - kernel = [1] * factor - - # setup kernel - kernel = torch.tensor(kernel, dtype=torch.float32) - if kernel.ndim == 1: - kernel = torch.outer(kernel, kernel) - kernel /= torch.sum(kernel) - - kernel = kernel * (gain * (factor**2)) - - if self.use_conv: - convH = weight.shape[2] - convW = weight.shape[3] - inC = weight.shape[1] - - pad_value = (kernel.shape[0] - factor) - (convW - 1) - - stride = (factor, factor) - # Determine data dimensions. - output_shape = ( - (hidden_states.shape[2] - 1) * factor + convH, - (hidden_states.shape[3] - 1) * factor + convW, - ) - output_padding = ( - output_shape[0] - (hidden_states.shape[2] - 1) * stride[0] - convH, - output_shape[1] - (hidden_states.shape[3] - 1) * stride[1] - convW, - ) - assert output_padding[0] >= 0 and output_padding[1] >= 0 - num_groups = hidden_states.shape[1] // inC - - # Transpose weights. - weight = torch.reshape(weight, (num_groups, -1, inC, convH, convW)) - weight = torch.flip(weight, dims=[3, 4]).permute(0, 2, 1, 3, 4) - weight = torch.reshape(weight, (num_groups * inC, -1, convH, convW)) - - inverse_conv = F.conv_transpose2d( - hidden_states, - weight, - stride=stride, - output_padding=output_padding, - padding=0, - ) - - output = upfirdn2d_native( - inverse_conv, - torch.tensor(kernel, device=inverse_conv.device), - pad=((pad_value + 1) // 2 + factor - 1, pad_value // 2 + 1), - ) - else: - pad_value = kernel.shape[0] - factor - output = upfirdn2d_native( - hidden_states, - torch.tensor(kernel, device=hidden_states.device), - up=factor, - pad=((pad_value + 1) // 2 + factor - 1, pad_value // 2), - ) - - return output - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.use_conv: - height = self._upsample_2d(hidden_states, self.Conv2d_0.weight, kernel=self.fir_kernel) - height = height + self.Conv2d_0.bias.reshape(1, -1, 1, 1) - else: - height = self._upsample_2d(hidden_states, kernel=self.fir_kernel, factor=2) - - return height - - -class KUpsample2D(nn.Module): - r"""A 2D K-upsampling layer. - - Parameters: - pad_mode (`str`, *optional*, default to `"reflect"`): the padding mode to use. - """ - - def __init__(self, pad_mode: str = "reflect"): - super().__init__() - self.pad_mode = pad_mode - kernel_1d = torch.tensor([[1 / 8, 3 / 8, 3 / 8, 1 / 8]]) * 2 - self.pad = kernel_1d.shape[1] // 2 - 1 - self.register_buffer("kernel", kernel_1d.T @ kernel_1d, persistent=False) - - def forward(self, inputs: torch.Tensor) -> torch.Tensor: - inputs = F.pad(inputs, ((self.pad + 1) // 2,) * 4, self.pad_mode) - weight = inputs.new_zeros( - [ - inputs.shape[1], - inputs.shape[1], - self.kernel.shape[0], - self.kernel.shape[1], - ] - ) - indices = torch.arange(inputs.shape[1], device=inputs.device) - kernel = self.kernel.to(weight)[None, :].expand(inputs.shape[1], -1, -1) - weight[indices, indices] = kernel - return F.conv_transpose2d(inputs, weight, stride=2, padding=self.pad * 2 + 1) - - -class CogVideoXUpsample3D(nn.Module): - r""" - A 3D Upsample layer using in CogVideoX by Tsinghua University & ZhipuAI # Todo: Wait for paper release. - - Args: - in_channels (`int`): - Number of channels in the input image. - out_channels (`int`): - Number of channels produced by the convolution. - kernel_size (`int`, defaults to `3`): - Size of the convolving kernel. - stride (`int`, defaults to `1`): - Stride of the convolution. - padding (`int`, defaults to `1`): - Padding added to all four sides of the input. - compress_time (`bool`, defaults to `False`): - Whether or not to compress the time dimension. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int = 3, - stride: int = 1, - padding: int = 1, - compress_time: bool = False, - ) -> None: - super().__init__() - - self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding) - self.compress_time = compress_time - - def forward(self, inputs: torch.Tensor) -> torch.Tensor: - if self.compress_time: - if inputs.shape[2] > 1 and inputs.shape[2] % 2 == 1: - # split first frame - x_first, x_rest = inputs[:, :, 0], inputs[:, :, 1:] - - x_first = F.interpolate(x_first, scale_factor=2.0) - x_rest = F.interpolate(x_rest, scale_factor=2.0) - x_first = x_first[:, :, None, :, :] - inputs = torch.cat([x_first, x_rest], dim=2) - elif inputs.shape[2] > 1: - inputs = F.interpolate(inputs, scale_factor=2.0) - else: - inputs = inputs.squeeze(2) - inputs = F.interpolate(inputs, scale_factor=2.0) - inputs = inputs[:, :, None, :, :] - else: - # only interpolate 2D - b, c, t, h, w = inputs.shape - inputs = inputs.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) - inputs = F.interpolate(inputs, scale_factor=2.0) - inputs = inputs.reshape(b, t, c, *inputs.shape[2:]).permute(0, 2, 1, 3, 4) - - b, c, t, h, w = inputs.shape - inputs = inputs.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) - inputs = self.conv(inputs) - inputs = inputs.reshape(b, t, *inputs.shape[1:]).permute(0, 2, 1, 3, 4) - - return inputs - - -def upfirdn2d_native( - tensor: torch.Tensor, - kernel: torch.Tensor, - up: int = 1, - down: int = 1, - pad: Tuple[int, int] = (0, 0), -) -> torch.Tensor: - up_x = up_y = up - down_x = down_y = down - pad_x0 = pad_y0 = pad[0] - pad_x1 = pad_y1 = pad[1] - - _, channel, in_h, in_w = tensor.shape - tensor = tensor.reshape(-1, in_h, in_w, 1) - - _, in_h, in_w, minor = tensor.shape - kernel_h, kernel_w = kernel.shape - - out = tensor.view(-1, in_h, 1, in_w, 1, minor) - out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1]) - out = out.view(-1, in_h * up_y, in_w * up_x, minor) - - out = F.pad(out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]) - out = out.to(tensor.device) # Move back to mps if necessary - out = out[ - :, - max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0), - max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0), - :, - ] - - out = out.permute(0, 3, 1, 2) - out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]) - w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w) - out = F.conv2d(out, w) - out = out.reshape( - -1, - minor, - in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1, - in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1, - ) - out = out.permute(0, 2, 3, 1) - out = out[:, ::down_y, ::down_x, :] - - out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1 - out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1 - - return out.view(-1, channel, out_h, out_w) - - -def upsample_2d( - hidden_states: torch.Tensor, - kernel: Optional[torch.Tensor] = None, - factor: int = 2, - gain: float = 1, -) -> torch.Tensor: - r"""Upsample2D a batch of 2D images with the given filter. - Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]` and upsamples each image with the given - filter. The filter is normalized so that if the input pixels are constant, they will be scaled by the specified - `gain`. Pixels outside the image are assumed to be zero, and the filter is padded with zeros so that its shape is - a: multiple of the upsampling factor. - - Args: - hidden_states (`torch.Tensor`): - Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`. - kernel (`torch.Tensor`, *optional*): - FIR filter of the shape `[firH, firW]` or `[firN]` (separable). The default is `[1] * factor`, which - corresponds to nearest-neighbor upsampling. - factor (`int`, *optional*, default to `2`): - Integer upsampling factor. - gain (`float`, *optional*, default to `1.0`): - Scaling factor for signal magnitude (default: 1.0). - - Returns: - output (`torch.Tensor`): - Tensor of the shape `[N, C, H * factor, W * factor]` - """ - assert isinstance(factor, int) and factor >= 1 - if kernel is None: - kernel = [1] * factor - - kernel = torch.tensor(kernel, dtype=torch.float32) - if kernel.ndim == 1: - kernel = torch.outer(kernel, kernel) - kernel /= torch.sum(kernel) - - kernel = kernel * (gain * (factor**2)) - pad_value = kernel.shape[0] - factor - output = upfirdn2d_native( - hidden_states, - kernel.to(device=hidden_states.device), - up=factor, - pad=((pad_value + 1) // 2 + factor - 1, pad_value // 2), - ) - return output diff --git a/venv/lib/python3.10/site-packages/diffusers/models/vae_flax.py b/venv/lib/python3.10/site-packages/diffusers/models/vae_flax.py deleted file mode 100644 index 93398a51eac71212253ca9c6751033e8759c555b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/vae_flax.py +++ /dev/null @@ -1,876 +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. - -# JAX implementation of VQGAN from taming-transformers https://github.com/CompVis/taming-transformers - -import math -from functools import partial -from typing import Tuple - -import flax -import flax.linen as nn -import jax -import jax.numpy as jnp -from flax.core.frozen_dict import FrozenDict - -from ..configuration_utils import ConfigMixin, flax_register_to_config -from ..utils import BaseOutput -from .modeling_flax_utils import FlaxModelMixin - - -@flax.struct.dataclass -class FlaxDecoderOutput(BaseOutput): - """ - Output of decoding method. - - Args: - sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)`): - The decoded output sample from the last layer of the model. - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - The `dtype` of the parameters. - """ - - sample: jnp.ndarray - - -@flax.struct.dataclass -class FlaxAutoencoderKLOutput(BaseOutput): - """ - Output of AutoencoderKL encoding method. - - Args: - latent_dist (`FlaxDiagonalGaussianDistribution`): - Encoded outputs of `Encoder` represented as the mean and logvar of `FlaxDiagonalGaussianDistribution`. - `FlaxDiagonalGaussianDistribution` allows for sampling latents from the distribution. - """ - - latent_dist: "FlaxDiagonalGaussianDistribution" - - -class FlaxUpsample2D(nn.Module): - """ - Flax implementation of 2D Upsample layer - - Args: - in_channels (`int`): - Input channels - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - dtype: jnp.dtype = jnp.float32 - - def setup(self): - self.conv = nn.Conv( - self.in_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - def __call__(self, hidden_states): - batch, height, width, channels = hidden_states.shape - hidden_states = jax.image.resize( - hidden_states, - shape=(batch, height * 2, width * 2, channels), - method="nearest", - ) - hidden_states = self.conv(hidden_states) - return hidden_states - - -class FlaxDownsample2D(nn.Module): - """ - Flax implementation of 2D Downsample layer - - Args: - in_channels (`int`): - Input channels - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - dtype: jnp.dtype = jnp.float32 - - def setup(self): - self.conv = nn.Conv( - self.in_channels, - kernel_size=(3, 3), - strides=(2, 2), - padding="VALID", - dtype=self.dtype, - ) - - def __call__(self, hidden_states): - pad = ((0, 0), (0, 1), (0, 1), (0, 0)) # pad height and width dim - hidden_states = jnp.pad(hidden_states, pad_width=pad) - hidden_states = self.conv(hidden_states) - return hidden_states - - -class FlaxResnetBlock2D(nn.Module): - """ - Flax implementation of 2D Resnet Block. - - Args: - in_channels (`int`): - Input channels - out_channels (`int`): - Output channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - groups (:obj:`int`, *optional*, defaults to `32`): - The number of groups to use for group norm. - use_nin_shortcut (:obj:`bool`, *optional*, defaults to `None`): - Whether to use `nin_shortcut`. This activates a new layer inside ResNet block - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - out_channels: int = None - dropout: float = 0.0 - groups: int = 32 - use_nin_shortcut: bool = None - dtype: jnp.dtype = jnp.float32 - - def setup(self): - out_channels = self.in_channels if self.out_channels is None else self.out_channels - - self.norm1 = nn.GroupNorm(num_groups=self.groups, epsilon=1e-6) - self.conv1 = nn.Conv( - out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - self.norm2 = nn.GroupNorm(num_groups=self.groups, epsilon=1e-6) - self.dropout_layer = nn.Dropout(self.dropout) - self.conv2 = nn.Conv( - out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - use_nin_shortcut = self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut - - self.conv_shortcut = None - if use_nin_shortcut: - self.conv_shortcut = nn.Conv( - out_channels, - kernel_size=(1, 1), - strides=(1, 1), - padding="VALID", - dtype=self.dtype, - ) - - def __call__(self, hidden_states, deterministic=True): - residual = hidden_states - hidden_states = self.norm1(hidden_states) - hidden_states = nn.swish(hidden_states) - hidden_states = self.conv1(hidden_states) - - hidden_states = self.norm2(hidden_states) - hidden_states = nn.swish(hidden_states) - hidden_states = self.dropout_layer(hidden_states, deterministic) - hidden_states = self.conv2(hidden_states) - - if self.conv_shortcut is not None: - residual = self.conv_shortcut(residual) - - return hidden_states + residual - - -class FlaxAttentionBlock(nn.Module): - r""" - Flax Convolutional based multi-head attention block for diffusion-based VAE. - - Parameters: - channels (:obj:`int`): - Input channels - num_head_channels (:obj:`int`, *optional*, defaults to `None`): - Number of attention heads - num_groups (:obj:`int`, *optional*, defaults to `32`): - The number of groups to use for group norm - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - - """ - - channels: int - num_head_channels: int = None - num_groups: int = 32 - dtype: jnp.dtype = jnp.float32 - - def setup(self): - self.num_heads = self.channels // self.num_head_channels if self.num_head_channels is not None else 1 - - dense = partial(nn.Dense, self.channels, dtype=self.dtype) - - self.group_norm = nn.GroupNorm(num_groups=self.num_groups, epsilon=1e-6) - self.query, self.key, self.value = dense(), dense(), dense() - self.proj_attn = dense() - - def transpose_for_scores(self, projection): - new_projection_shape = projection.shape[:-1] + (self.num_heads, -1) - # move heads to 2nd position (B, T, H * D) -> (B, T, H, D) - new_projection = projection.reshape(new_projection_shape) - # (B, T, H, D) -> (B, H, T, D) - new_projection = jnp.transpose(new_projection, (0, 2, 1, 3)) - return new_projection - - def __call__(self, hidden_states): - residual = hidden_states - batch, height, width, channels = hidden_states.shape - - hidden_states = self.group_norm(hidden_states) - - hidden_states = hidden_states.reshape((batch, height * width, channels)) - - query = self.query(hidden_states) - key = self.key(hidden_states) - value = self.value(hidden_states) - - # transpose - query = self.transpose_for_scores(query) - key = self.transpose_for_scores(key) - value = self.transpose_for_scores(value) - - # compute attentions - scale = 1 / math.sqrt(math.sqrt(self.channels / self.num_heads)) - attn_weights = jnp.einsum("...qc,...kc->...qk", query * scale, key * scale) - attn_weights = nn.softmax(attn_weights, axis=-1) - - # attend to values - hidden_states = jnp.einsum("...kc,...qk->...qc", value, attn_weights) - - hidden_states = jnp.transpose(hidden_states, (0, 2, 1, 3)) - new_hidden_states_shape = hidden_states.shape[:-2] + (self.channels,) - hidden_states = hidden_states.reshape(new_hidden_states_shape) - - hidden_states = self.proj_attn(hidden_states) - hidden_states = hidden_states.reshape((batch, height, width, channels)) - hidden_states = hidden_states + residual - return hidden_states - - -class FlaxDownEncoderBlock2D(nn.Module): - r""" - Flax Resnet blocks-based Encoder block for diffusion-based VAE. - - Parameters: - in_channels (:obj:`int`): - Input channels - out_channels (:obj:`int`): - Output channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of Resnet layer block - resnet_groups (:obj:`int`, *optional*, defaults to `32`): - The number of groups to use for the Resnet block group norm - add_downsample (:obj:`bool`, *optional*, defaults to `True`): - Whether to add downsample layer - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - out_channels: int - dropout: float = 0.0 - num_layers: int = 1 - resnet_groups: int = 32 - add_downsample: bool = True - dtype: jnp.dtype = jnp.float32 - - def setup(self): - resnets = [] - for i in range(self.num_layers): - in_channels = self.in_channels if i == 0 else self.out_channels - - res_block = FlaxResnetBlock2D( - in_channels=in_channels, - out_channels=self.out_channels, - dropout=self.dropout, - groups=self.resnet_groups, - dtype=self.dtype, - ) - resnets.append(res_block) - self.resnets = resnets - - if self.add_downsample: - self.downsamplers_0 = FlaxDownsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, deterministic=True): - for resnet in self.resnets: - hidden_states = resnet(hidden_states, deterministic=deterministic) - - if self.add_downsample: - hidden_states = self.downsamplers_0(hidden_states) - - return hidden_states - - -class FlaxUpDecoderBlock2D(nn.Module): - r""" - Flax Resnet blocks-based Decoder block for diffusion-based VAE. - - Parameters: - in_channels (:obj:`int`): - Input channels - out_channels (:obj:`int`): - Output channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of Resnet layer block - resnet_groups (:obj:`int`, *optional*, defaults to `32`): - The number of groups to use for the Resnet block group norm - add_upsample (:obj:`bool`, *optional*, defaults to `True`): - Whether to add upsample layer - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - out_channels: int - dropout: float = 0.0 - num_layers: int = 1 - resnet_groups: int = 32 - add_upsample: bool = True - dtype: jnp.dtype = jnp.float32 - - def setup(self): - resnets = [] - for i in range(self.num_layers): - in_channels = self.in_channels if i == 0 else self.out_channels - res_block = FlaxResnetBlock2D( - in_channels=in_channels, - out_channels=self.out_channels, - dropout=self.dropout, - groups=self.resnet_groups, - dtype=self.dtype, - ) - resnets.append(res_block) - - self.resnets = resnets - - if self.add_upsample: - self.upsamplers_0 = FlaxUpsample2D(self.out_channels, dtype=self.dtype) - - def __call__(self, hidden_states, deterministic=True): - for resnet in self.resnets: - hidden_states = resnet(hidden_states, deterministic=deterministic) - - if self.add_upsample: - hidden_states = self.upsamplers_0(hidden_states) - - return hidden_states - - -class FlaxUNetMidBlock2D(nn.Module): - r""" - Flax Unet Mid-Block module. - - Parameters: - in_channels (:obj:`int`): - Input channels - dropout (:obj:`float`, *optional*, defaults to 0.0): - Dropout rate - num_layers (:obj:`int`, *optional*, defaults to 1): - Number of Resnet layer block - resnet_groups (:obj:`int`, *optional*, defaults to `32`): - The number of groups to use for the Resnet and Attention block group norm - num_attention_heads (:obj:`int`, *optional*, defaults to `1`): - Number of attention heads for each attention block - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int - dropout: float = 0.0 - num_layers: int = 1 - resnet_groups: int = 32 - num_attention_heads: int = 1 - dtype: jnp.dtype = jnp.float32 - - def setup(self): - resnet_groups = self.resnet_groups if self.resnet_groups is not None else min(self.in_channels // 4, 32) - - # there is always at least one resnet - resnets = [ - FlaxResnetBlock2D( - in_channels=self.in_channels, - out_channels=self.in_channels, - dropout=self.dropout, - groups=resnet_groups, - dtype=self.dtype, - ) - ] - - attentions = [] - - for _ in range(self.num_layers): - attn_block = FlaxAttentionBlock( - channels=self.in_channels, - num_head_channels=self.num_attention_heads, - num_groups=resnet_groups, - dtype=self.dtype, - ) - attentions.append(attn_block) - - res_block = FlaxResnetBlock2D( - in_channels=self.in_channels, - out_channels=self.in_channels, - dropout=self.dropout, - groups=resnet_groups, - dtype=self.dtype, - ) - resnets.append(res_block) - - self.resnets = resnets - self.attentions = attentions - - def __call__(self, hidden_states, deterministic=True): - hidden_states = self.resnets[0](hidden_states, deterministic=deterministic) - for attn, resnet in zip(self.attentions, self.resnets[1:]): - hidden_states = attn(hidden_states) - hidden_states = resnet(hidden_states, deterministic=deterministic) - - return hidden_states - - -class FlaxEncoder(nn.Module): - r""" - Flax Implementation of VAE Encoder. - - This model is a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module) - subclass. Use it as a regular Flax linen Module and refer to the Flax documentation for all matter related to - general usage and behavior. - - Finally, this model supports inherent JAX features such as: - - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit) - - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation) - - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap) - - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap) - - Parameters: - in_channels (:obj:`int`, *optional*, defaults to 3): - Input channels - out_channels (:obj:`int`, *optional*, defaults to 3): - Output channels - 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 - 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`): - norm num group - act_fn (:obj:`str`, *optional*, defaults to `silu`): - Activation function - double_z (:obj:`bool`, *optional*, defaults to `False`): - Whether to double the last output channels - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - Parameters `dtype` - """ - - in_channels: int = 3 - out_channels: int = 3 - 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" - double_z: bool = False - dtype: jnp.dtype = jnp.float32 - - def setup(self): - block_out_channels = self.block_out_channels - # in - self.conv_in = nn.Conv( - block_out_channels[0], - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - # downsampling - down_blocks = [] - output_channel = block_out_channels[0] - for i, _ in enumerate(self.down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = FlaxDownEncoderBlock2D( - in_channels=input_channel, - out_channels=output_channel, - num_layers=self.layers_per_block, - resnet_groups=self.norm_num_groups, - add_downsample=not is_final_block, - dtype=self.dtype, - ) - down_blocks.append(down_block) - self.down_blocks = down_blocks - - # middle - self.mid_block = FlaxUNetMidBlock2D( - in_channels=block_out_channels[-1], - resnet_groups=self.norm_num_groups, - num_attention_heads=None, - dtype=self.dtype, - ) - - # end - conv_out_channels = 2 * self.out_channels if self.double_z else self.out_channels - self.conv_norm_out = nn.GroupNorm(num_groups=self.norm_num_groups, epsilon=1e-6) - self.conv_out = nn.Conv( - conv_out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - def __call__(self, sample, deterministic: bool = True): - # in - sample = self.conv_in(sample) - - # downsampling - for block in self.down_blocks: - sample = block(sample, deterministic=deterministic) - - # middle - sample = self.mid_block(sample, deterministic=deterministic) - - # end - sample = self.conv_norm_out(sample) - sample = nn.swish(sample) - sample = self.conv_out(sample) - - return sample - - -class FlaxDecoder(nn.Module): - r""" - Flax Implementation of VAE Decoder. - - This model is a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module) - subclass. Use it as a regular Flax linen Module and refer to the Flax documentation for all matter related to - general usage and behavior. - - Finally, this model supports inherent JAX features such as: - - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit) - - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation) - - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap) - - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap) - - Parameters: - in_channels (:obj:`int`, *optional*, defaults to 3): - Input channels - out_channels (:obj:`int`, *optional*, defaults to 3): - Output channels - 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 - 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`): - norm num group - act_fn (:obj:`str`, *optional*, defaults to `silu`): - Activation function - double_z (:obj:`bool`, *optional*, defaults to `False`): - Whether to double the last output channels - dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): - parameters `dtype` - """ - - in_channels: int = 3 - out_channels: int = 3 - up_block_types: Tuple[str] = ("UpDecoderBlock2D",) - block_out_channels: int = (64,) - layers_per_block: int = 2 - norm_num_groups: int = 32 - act_fn: str = "silu" - dtype: jnp.dtype = jnp.float32 - - def setup(self): - block_out_channels = self.block_out_channels - - # z to block_in - self.conv_in = nn.Conv( - block_out_channels[-1], - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - # middle - self.mid_block = FlaxUNetMidBlock2D( - in_channels=block_out_channels[-1], - resnet_groups=self.norm_num_groups, - num_attention_heads=None, - dtype=self.dtype, - ) - - # upsampling - reversed_block_out_channels = list(reversed(block_out_channels)) - output_channel = reversed_block_out_channels[0] - up_blocks = [] - for i, _ in enumerate(self.up_block_types): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - - is_final_block = i == len(block_out_channels) - 1 - - up_block = FlaxUpDecoderBlock2D( - in_channels=prev_output_channel, - out_channels=output_channel, - num_layers=self.layers_per_block + 1, - resnet_groups=self.norm_num_groups, - add_upsample=not is_final_block, - dtype=self.dtype, - ) - up_blocks.append(up_block) - prev_output_channel = output_channel - - self.up_blocks = up_blocks - - # end - self.conv_norm_out = nn.GroupNorm(num_groups=self.norm_num_groups, epsilon=1e-6) - self.conv_out = nn.Conv( - self.out_channels, - kernel_size=(3, 3), - strides=(1, 1), - padding=((1, 1), (1, 1)), - dtype=self.dtype, - ) - - def __call__(self, sample, deterministic: bool = True): - # z to block_in - sample = self.conv_in(sample) - - # middle - sample = self.mid_block(sample, deterministic=deterministic) - - # upsampling - for block in self.up_blocks: - sample = block(sample, deterministic=deterministic) - - sample = self.conv_norm_out(sample) - sample = nn.swish(sample) - sample = self.conv_out(sample) - - return sample - - -class FlaxDiagonalGaussianDistribution(object): - def __init__(self, parameters, deterministic=False): - # Last axis to account for channels-last - self.mean, self.logvar = jnp.split(parameters, 2, axis=-1) - self.logvar = jnp.clip(self.logvar, -30.0, 20.0) - self.deterministic = deterministic - self.std = jnp.exp(0.5 * self.logvar) - self.var = jnp.exp(self.logvar) - if self.deterministic: - self.var = self.std = jnp.zeros_like(self.mean) - - def sample(self, key): - return self.mean + self.std * jax.random.normal(key, self.mean.shape) - - def kl(self, other=None): - if self.deterministic: - return jnp.array([0.0]) - - if other is None: - return 0.5 * jnp.sum(self.mean**2 + self.var - 1.0 - self.logvar, axis=[1, 2, 3]) - - return 0.5 * jnp.sum( - jnp.square(self.mean - other.mean) / other.var + self.var / other.var - 1.0 - self.logvar + other.logvar, - axis=[1, 2, 3], - ) - - def nll(self, sample, axis=[1, 2, 3]): - if self.deterministic: - return jnp.array([0.0]) - - logtwopi = jnp.log(2.0 * jnp.pi) - return 0.5 * jnp.sum(logtwopi + self.logvar + jnp.square(sample - self.mean) / self.var, axis=axis) - - def mode(self): - return self.mean - - -@flax_register_to_config -class FlaxAutoencoderKL(nn.Module, FlaxModelMixin, ConfigMixin): - r""" - Flax implementation of a VAE model with KL loss for decoding latent representations. - - This model inherits from [`FlaxModelMixin`]. Check the superclass documentation for it's generic methods - implemented for all models (such as downloading or saving). - - This model is a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module) - subclass. Use it as a regular Flax Linen module and refer to the Flax documentation for all matter related to its - general usage and behavior. - - Inherent JAX features such as the following are supported: - - - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit) - - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation) - - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap) - - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap) - - 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[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`): - The activation function to use. - latent_channels (`int`, *optional*, defaults to `4`): - Number of channels in the latent space. - norm_num_groups (`int`, *optional*, defaults to `32`): - The number of groups for normalization. - sample_size (`int`, *optional*, defaults to 32): - Sample input size. - scaling_factor (`float`, *optional*, defaults to 0.18215): - 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. - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - The `dtype` of the parameters. - """ - - 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,) - 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 - dtype: jnp.dtype = jnp.float32 - - def setup(self): - self.encoder = FlaxEncoder( - in_channels=self.config.in_channels, - out_channels=self.config.latent_channels, - down_block_types=self.config.down_block_types, - block_out_channels=self.config.block_out_channels, - layers_per_block=self.config.layers_per_block, - act_fn=self.config.act_fn, - norm_num_groups=self.config.norm_num_groups, - double_z=True, - dtype=self.dtype, - ) - self.decoder = FlaxDecoder( - in_channels=self.config.latent_channels, - out_channels=self.config.out_channels, - up_block_types=self.config.up_block_types, - block_out_channels=self.config.block_out_channels, - layers_per_block=self.config.layers_per_block, - norm_num_groups=self.config.norm_num_groups, - act_fn=self.config.act_fn, - dtype=self.dtype, - ) - self.quant_conv = nn.Conv( - 2 * self.config.latent_channels, - kernel_size=(1, 1), - strides=(1, 1), - padding="VALID", - dtype=self.dtype, - ) - self.post_quant_conv = nn.Conv( - self.config.latent_channels, - kernel_size=(1, 1), - strides=(1, 1), - padding="VALID", - dtype=self.dtype, - ) - - def init_weights(self, rng: jax.Array) -> FrozenDict: - # init input tensors - sample_shape = (1, self.in_channels, self.sample_size, self.sample_size) - sample = jnp.zeros(sample_shape, dtype=jnp.float32) - - params_rng, dropout_rng, gaussian_rng = jax.random.split(rng, 3) - rngs = {"params": params_rng, "dropout": dropout_rng, "gaussian": gaussian_rng} - - return self.init(rngs, sample)["params"] - - def encode(self, sample, deterministic: bool = True, return_dict: bool = True): - sample = jnp.transpose(sample, (0, 2, 3, 1)) - - hidden_states = self.encoder(sample, deterministic=deterministic) - moments = self.quant_conv(hidden_states) - posterior = FlaxDiagonalGaussianDistribution(moments) - - if not return_dict: - return (posterior,) - - return FlaxAutoencoderKLOutput(latent_dist=posterior) - - def decode(self, latents, deterministic: bool = True, return_dict: bool = True): - if latents.shape[-1] != self.config.latent_channels: - latents = jnp.transpose(latents, (0, 2, 3, 1)) - - hidden_states = self.post_quant_conv(latents) - hidden_states = self.decoder(hidden_states, deterministic=deterministic) - - hidden_states = jnp.transpose(hidden_states, (0, 3, 1, 2)) - - if not return_dict: - return (hidden_states,) - - return FlaxDecoderOutput(sample=hidden_states) - - def __call__(self, sample, sample_posterior=False, deterministic: bool = True, return_dict: bool = True): - posterior = self.encode(sample, deterministic=deterministic, return_dict=return_dict) - if sample_posterior: - rng = self.make_rng("gaussian") - hidden_states = posterior.latent_dist.sample(rng) - else: - hidden_states = posterior.latent_dist.mode() - - sample = self.decode(hidden_states, return_dict=return_dict).sample - - if not return_dict: - return (sample,) - - return FlaxDecoderOutput(sample=sample) diff --git a/venv/lib/python3.10/site-packages/diffusers/models/vq_model.py b/venv/lib/python3.10/site-packages/diffusers/models/vq_model.py deleted file mode 100644 index 63dee5eec5543d495cbe631600679088320a6f34..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/models/vq_model.py +++ /dev/null @@ -1,29 +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 deprecate -from .autoencoders.vq_model import VQEncoderOutput, VQModel - - -class VQEncoderOutput(VQEncoderOutput): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `VQEncoderOutput` from `diffusers.models.vq_model` is deprecated and this will be removed in a future version. Please use `from diffusers.models.autoencoders.vq_model import VQEncoderOutput`, instead." - deprecate("VQEncoderOutput", "0.31", deprecation_message) - super().__init__(*args, **kwargs) - - -class VQModel(VQModel): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `VQModel` from `diffusers.models.vq_model` is deprecated and this will be removed in a future version. Please use `from diffusers.models.autoencoders.vq_model import VQModel`, instead." - deprecate("VQModel", "0.31", deprecation_message) - super().__init__(*args, **kwargs) diff --git a/venv/lib/python3.10/site-packages/diffusers/optimization.py b/venv/lib/python3.10/site-packages/diffusers/optimization.py deleted file mode 100644 index e0b3576e44265d501e4a9ddcf3cfca64ae3b9bee..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/optimization.py +++ /dev/null @@ -1,361 +0,0 @@ -# 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. -"""PyTorch optimization for diffusion models.""" - -import math -from enum import Enum -from typing import Optional, Union - -from torch.optim import Optimizer -from torch.optim.lr_scheduler import LambdaLR - -from .utils import logging - - -logger = logging.get_logger(__name__) - - -class SchedulerType(Enum): - LINEAR = "linear" - COSINE = "cosine" - COSINE_WITH_RESTARTS = "cosine_with_restarts" - POLYNOMIAL = "polynomial" - CONSTANT = "constant" - CONSTANT_WITH_WARMUP = "constant_with_warmup" - PIECEWISE_CONSTANT = "piecewise_constant" - - -def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1) -> LambdaLR: - """ - Create a schedule with a constant learning rate, using the learning rate set in optimizer. - - Args: - optimizer ([`~torch.optim.Optimizer`]): - The optimizer for which to schedule the learning rate. - last_epoch (`int`, *optional*, defaults to -1): - The index of the last epoch when resuming training. - - Return: - `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. - """ - return LambdaLR(optimizer, lambda _: 1, last_epoch=last_epoch) - - -def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1) -> LambdaLR: - """ - Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate - increases linearly between 0 and the initial lr set in the optimizer. - - Args: - optimizer ([`~torch.optim.Optimizer`]): - The optimizer for which to schedule the learning rate. - num_warmup_steps (`int`): - The number of steps for the warmup phase. - last_epoch (`int`, *optional*, defaults to -1): - The index of the last epoch when resuming training. - - Return: - `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. - """ - - def lr_lambda(current_step: int): - if current_step < num_warmup_steps: - return float(current_step) / float(max(1.0, num_warmup_steps)) - return 1.0 - - return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch) - - -def get_piecewise_constant_schedule(optimizer: Optimizer, step_rules: str, last_epoch: int = -1) -> LambdaLR: - """ - Create a schedule with a constant learning rate, using the learning rate set in optimizer. - - Args: - optimizer ([`~torch.optim.Optimizer`]): - The optimizer for which to schedule the learning rate. - step_rules (`string`): - The rules for the learning rate. ex: rule_steps="1:10,0.1:20,0.01:30,0.005" it means that the learning rate - if multiple 1 for the first 10 steps, multiple 0.1 for the next 20 steps, multiple 0.01 for the next 30 - steps and multiple 0.005 for the other steps. - last_epoch (`int`, *optional*, defaults to -1): - The index of the last epoch when resuming training. - - Return: - `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. - """ - - rules_dict = {} - rule_list = step_rules.split(",") - for rule_str in rule_list[:-1]: - value_str, steps_str = rule_str.split(":") - steps = int(steps_str) - value = float(value_str) - rules_dict[steps] = value - last_lr_multiple = float(rule_list[-1]) - - def create_rules_function(rules_dict, last_lr_multiple): - def rule_func(steps: int) -> float: - sorted_steps = sorted(rules_dict.keys()) - for i, sorted_step in enumerate(sorted_steps): - if steps < sorted_step: - return rules_dict[sorted_steps[i]] - return last_lr_multiple - - return rule_func - - rules_func = create_rules_function(rules_dict, last_lr_multiple) - - return LambdaLR(optimizer, rules_func, last_epoch=last_epoch) - - -def get_linear_schedule_with_warmup( - optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, last_epoch: int = -1 -) -> LambdaLR: - """ - Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after - a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer. - - Args: - optimizer ([`~torch.optim.Optimizer`]): - The optimizer for which to schedule the learning rate. - num_warmup_steps (`int`): - The number of steps for the warmup phase. - num_training_steps (`int`): - The total number of training steps. - last_epoch (`int`, *optional*, defaults to -1): - The index of the last epoch when resuming training. - - Return: - `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. - """ - - def lr_lambda(current_step: int): - if current_step < num_warmup_steps: - return float(current_step) / float(max(1, num_warmup_steps)) - return max( - 0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps)) - ) - - return LambdaLR(optimizer, lr_lambda, last_epoch) - - -def get_cosine_schedule_with_warmup( - optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1 -) -> LambdaLR: - """ - Create a schedule with a learning rate that decreases following the values of the cosine function between the - initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the - initial lr set in the optimizer. - - Args: - optimizer ([`~torch.optim.Optimizer`]): - The optimizer for which to schedule the learning rate. - num_warmup_steps (`int`): - The number of steps for the warmup phase. - num_training_steps (`int`): - The total number of training steps. - num_periods (`float`, *optional*, defaults to 0.5): - The number of periods of the cosine function in a schedule (the default is to just decrease from the max - value to 0 following a half-cosine). - last_epoch (`int`, *optional*, defaults to -1): - The index of the last epoch when resuming training. - - Return: - `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. - """ - - def lr_lambda(current_step): - if current_step < num_warmup_steps: - return float(current_step) / float(max(1, num_warmup_steps)) - progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) - return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))) - - return LambdaLR(optimizer, lr_lambda, last_epoch) - - -def get_cosine_with_hard_restarts_schedule_with_warmup( - optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1 -) -> LambdaLR: - """ - Create a schedule with a learning rate that decreases following the values of the cosine function between the - initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases - linearly between 0 and the initial lr set in the optimizer. - - Args: - optimizer ([`~torch.optim.Optimizer`]): - The optimizer for which to schedule the learning rate. - num_warmup_steps (`int`): - The number of steps for the warmup phase. - num_training_steps (`int`): - The total number of training steps. - num_cycles (`int`, *optional*, defaults to 1): - The number of hard restarts to use. - last_epoch (`int`, *optional*, defaults to -1): - The index of the last epoch when resuming training. - - Return: - `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. - """ - - def lr_lambda(current_step): - if current_step < num_warmup_steps: - return float(current_step) / float(max(1, num_warmup_steps)) - progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) - if progress >= 1.0: - return 0.0 - return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0)))) - - return LambdaLR(optimizer, lr_lambda, last_epoch) - - -def get_polynomial_decay_schedule_with_warmup( - optimizer: Optimizer, - num_warmup_steps: int, - num_training_steps: int, - lr_end: float = 1e-7, - power: float = 1.0, - last_epoch: int = -1, -) -> LambdaLR: - """ - Create a schedule with a learning rate that decreases as a polynomial decay from the initial lr set in the - optimizer to end lr defined by *lr_end*, after a warmup period during which it increases linearly from 0 to the - initial lr set in the optimizer. - - Args: - optimizer ([`~torch.optim.Optimizer`]): - The optimizer for which to schedule the learning rate. - num_warmup_steps (`int`): - The number of steps for the warmup phase. - num_training_steps (`int`): - The total number of training steps. - lr_end (`float`, *optional*, defaults to 1e-7): - The end LR. - power (`float`, *optional*, defaults to 1.0): - Power factor. - last_epoch (`int`, *optional*, defaults to -1): - The index of the last epoch when resuming training. - - Note: *power* defaults to 1.0 as in the fairseq implementation, which in turn is based on the original BERT - implementation at - https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37 - - Return: - `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. - - """ - - lr_init = optimizer.defaults["lr"] - if not (lr_init > lr_end): - raise ValueError(f"lr_end ({lr_end}) must be smaller than initial lr ({lr_init})") - - def lr_lambda(current_step: int): - if current_step < num_warmup_steps: - return float(current_step) / float(max(1, num_warmup_steps)) - elif current_step > num_training_steps: - return lr_end / lr_init # as LambdaLR multiplies by lr_init - else: - lr_range = lr_init - lr_end - decay_steps = num_training_steps - num_warmup_steps - pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps - decay = lr_range * pct_remaining**power + lr_end - return decay / lr_init # as LambdaLR multiplies by lr_init - - return LambdaLR(optimizer, lr_lambda, last_epoch) - - -TYPE_TO_SCHEDULER_FUNCTION = { - SchedulerType.LINEAR: get_linear_schedule_with_warmup, - SchedulerType.COSINE: get_cosine_schedule_with_warmup, - SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup, - SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup, - SchedulerType.CONSTANT: get_constant_schedule, - SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup, - SchedulerType.PIECEWISE_CONSTANT: get_piecewise_constant_schedule, -} - - -def get_scheduler( - name: Union[str, SchedulerType], - optimizer: Optimizer, - step_rules: Optional[str] = None, - num_warmup_steps: Optional[int] = None, - num_training_steps: Optional[int] = None, - num_cycles: int = 1, - power: float = 1.0, - last_epoch: int = -1, -) -> LambdaLR: - """ - Unified API to get any scheduler from its name. - - Args: - name (`str` or `SchedulerType`): - The name of the scheduler to use. - optimizer (`torch.optim.Optimizer`): - The optimizer that will be used during training. - step_rules (`str`, *optional*): - A string representing the step rules to use. This is only used by the `PIECEWISE_CONSTANT` scheduler. - num_warmup_steps (`int`, *optional*): - The number of warmup steps to do. This is not required by all schedulers (hence the argument being - optional), the function will raise an error if it's unset and the scheduler type requires it. - num_training_steps (`int``, *optional*): - The number of training steps to do. This is not required by all schedulers (hence the argument being - optional), the function will raise an error if it's unset and the scheduler type requires it. - num_cycles (`int`, *optional*): - The number of hard restarts used in `COSINE_WITH_RESTARTS` scheduler. - power (`float`, *optional*, defaults to 1.0): - Power factor. See `POLYNOMIAL` scheduler - last_epoch (`int`, *optional*, defaults to -1): - The index of the last epoch when resuming training. - """ - name = SchedulerType(name) - schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name] - if name == SchedulerType.CONSTANT: - return schedule_func(optimizer, last_epoch=last_epoch) - - if name == SchedulerType.PIECEWISE_CONSTANT: - return schedule_func(optimizer, step_rules=step_rules, last_epoch=last_epoch) - - # All other schedulers require `num_warmup_steps` - if num_warmup_steps is None: - raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.") - - if name == SchedulerType.CONSTANT_WITH_WARMUP: - return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, last_epoch=last_epoch) - - # All other schedulers require `num_training_steps` - if num_training_steps is None: - raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.") - - if name == SchedulerType.COSINE_WITH_RESTARTS: - return schedule_func( - optimizer, - num_warmup_steps=num_warmup_steps, - num_training_steps=num_training_steps, - num_cycles=num_cycles, - last_epoch=last_epoch, - ) - - if name == SchedulerType.POLYNOMIAL: - return schedule_func( - optimizer, - num_warmup_steps=num_warmup_steps, - num_training_steps=num_training_steps, - power=power, - last_epoch=last_epoch, - ) - - return schedule_func( - optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, last_epoch=last_epoch - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/__init__.py deleted file mode 100644 index b32d55bd51715ccf4c7a809eb924b577a2072ac9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/__init__.py +++ /dev/null @@ -1,860 +0,0 @@ -from typing import TYPE_CHECKING - -from ..utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - get_objects_from_module, - is_flax_available, - is_k_diffusion_available, - is_librosa_available, - is_note_seq_available, - is_onnx_available, - is_opencv_available, - is_sentencepiece_available, - is_torch_available, - is_torch_npu_available, - is_transformers_available, -) - - -# These modules contain pipelines from multiple libraries/frameworks -_dummy_objects = {} -_import_structure = { - "controlnet": [], - "controlnet_hunyuandit": [], - "controlnet_sd3": [], - "controlnet_xs": [], - "deprecated": [], - "latent_diffusion": [], - "ledits_pp": [], - "marigold": [], - "pag": [], - "stable_diffusion": [], - "stable_diffusion_xl": [], -} - -try: - if not is_torch_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import dummy_pt_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_pt_objects)) -else: - _import_structure["auto_pipeline"] = [ - "AutoPipelineForImage2Image", - "AutoPipelineForInpainting", - "AutoPipelineForText2Image", - ] - _import_structure["consistency_models"] = ["ConsistencyModelPipeline"] - _import_structure["dance_diffusion"] = ["DanceDiffusionPipeline"] - _import_structure["ddim"] = ["DDIMPipeline"] - _import_structure["ddpm"] = ["DDPMPipeline"] - _import_structure["dit"] = ["DiTPipeline"] - _import_structure["latent_diffusion"].extend(["LDMSuperResolutionPipeline"]) - _import_structure["pipeline_utils"] = [ - "AudioPipelineOutput", - "DiffusionPipeline", - "StableDiffusionMixin", - "ImagePipelineOutput", - ] - _import_structure["deprecated"].extend( - [ - "PNDMPipeline", - "LDMPipeline", - "RePaintPipeline", - "ScoreSdeVePipeline", - "KarrasVePipeline", - ] - ) -try: - if not (is_torch_available() and is_librosa_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import dummy_torch_and_librosa_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_librosa_objects)) -else: - _import_structure["deprecated"].extend(["AudioDiffusionPipeline", "Mel"]) - -try: - if not (is_transformers_available() and is_torch_available() and is_note_seq_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import dummy_transformers_and_torch_and_note_seq_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_transformers_and_torch_and_note_seq_objects)) -else: - _import_structure["deprecated"].extend( - [ - "MidiProcessor", - "SpectrogramDiffusionPipeline", - ] - ) - -try: - if not (is_torch_available() and is_transformers_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["deprecated"].extend( - [ - "VQDiffusionPipeline", - "AltDiffusionPipeline", - "AltDiffusionImg2ImgPipeline", - "CycleDiffusionPipeline", - "StableDiffusionInpaintPipelineLegacy", - "StableDiffusionPix2PixZeroPipeline", - "StableDiffusionParadigmsPipeline", - "StableDiffusionModelEditingPipeline", - "VersatileDiffusionDualGuidedPipeline", - "VersatileDiffusionImageVariationPipeline", - "VersatileDiffusionPipeline", - "VersatileDiffusionTextToImagePipeline", - ] - ) - _import_structure["allegro"] = ["AllegroPipeline"] - _import_structure["amused"] = ["AmusedImg2ImgPipeline", "AmusedInpaintPipeline", "AmusedPipeline"] - _import_structure["animatediff"] = [ - "AnimateDiffPipeline", - "AnimateDiffControlNetPipeline", - "AnimateDiffSDXLPipeline", - "AnimateDiffSparseControlNetPipeline", - "AnimateDiffVideoToVideoPipeline", - "AnimateDiffVideoToVideoControlNetPipeline", - ] - _import_structure["flux"] = [ - "FluxControlPipeline", - "FluxControlInpaintPipeline", - "FluxControlImg2ImgPipeline", - "FluxControlNetPipeline", - "FluxControlNetImg2ImgPipeline", - "FluxControlNetInpaintPipeline", - "FluxImg2ImgPipeline", - "FluxInpaintPipeline", - "FluxPipeline", - "FluxFillPipeline", - "FluxPriorReduxPipeline", - "ReduxImageEncoder", - ] - _import_structure["audioldm"] = ["AudioLDMPipeline"] - _import_structure["audioldm2"] = [ - "AudioLDM2Pipeline", - "AudioLDM2ProjectionModel", - "AudioLDM2UNet2DConditionModel", - ] - _import_structure["blip_diffusion"] = ["BlipDiffusionPipeline"] - _import_structure["chroma"] = ["ChromaPipeline", "ChromaImg2ImgPipeline"] - _import_structure["cogvideo"] = [ - "CogVideoXPipeline", - "CogVideoXImageToVideoPipeline", - "CogVideoXVideoToVideoPipeline", - "CogVideoXFunControlPipeline", - ] - _import_structure["cogview3"] = ["CogView3PlusPipeline"] - _import_structure["cogview4"] = ["CogView4Pipeline", "CogView4ControlPipeline"] - _import_structure["consisid"] = ["ConsisIDPipeline"] - _import_structure["cosmos"] = [ - "Cosmos2TextToImagePipeline", - "CosmosTextToWorldPipeline", - "CosmosVideoToWorldPipeline", - "Cosmos2VideoToWorldPipeline", - ] - _import_structure["controlnet"].extend( - [ - "BlipDiffusionControlNetPipeline", - "StableDiffusionControlNetImg2ImgPipeline", - "StableDiffusionControlNetInpaintPipeline", - "StableDiffusionControlNetPipeline", - "StableDiffusionXLControlNetImg2ImgPipeline", - "StableDiffusionXLControlNetInpaintPipeline", - "StableDiffusionXLControlNetPipeline", - "StableDiffusionXLControlNetUnionPipeline", - "StableDiffusionXLControlNetUnionInpaintPipeline", - "StableDiffusionXLControlNetUnionImg2ImgPipeline", - ] - ) - _import_structure["pag"].extend( - [ - "StableDiffusionControlNetPAGInpaintPipeline", - "AnimateDiffPAGPipeline", - "KolorsPAGPipeline", - "HunyuanDiTPAGPipeline", - "StableDiffusion3PAGPipeline", - "StableDiffusion3PAGImg2ImgPipeline", - "StableDiffusionPAGPipeline", - "StableDiffusionPAGImg2ImgPipeline", - "StableDiffusionPAGInpaintPipeline", - "StableDiffusionControlNetPAGPipeline", - "StableDiffusionXLPAGPipeline", - "StableDiffusionXLPAGInpaintPipeline", - "StableDiffusionXLControlNetPAGImg2ImgPipeline", - "StableDiffusionXLControlNetPAGPipeline", - "StableDiffusionXLPAGImg2ImgPipeline", - "PixArtSigmaPAGPipeline", - "SanaPAGPipeline", - ] - ) - _import_structure["controlnet_xs"].extend( - [ - "StableDiffusionControlNetXSPipeline", - "StableDiffusionXLControlNetXSPipeline", - ] - ) - _import_structure["controlnet_hunyuandit"].extend( - [ - "HunyuanDiTControlNetPipeline", - ] - ) - _import_structure["controlnet_sd3"].extend( - [ - "StableDiffusion3ControlNetPipeline", - "StableDiffusion3ControlNetInpaintingPipeline", - ] - ) - _import_structure["deepfloyd_if"] = [ - "IFImg2ImgPipeline", - "IFImg2ImgSuperResolutionPipeline", - "IFInpaintingPipeline", - "IFInpaintingSuperResolutionPipeline", - "IFPipeline", - "IFSuperResolutionPipeline", - ] - _import_structure["easyanimate"] = [ - "EasyAnimatePipeline", - "EasyAnimateInpaintPipeline", - "EasyAnimateControlPipeline", - ] - _import_structure["hidream_image"] = ["HiDreamImagePipeline"] - _import_structure["hunyuandit"] = ["HunyuanDiTPipeline"] - _import_structure["hunyuan_video"] = [ - "HunyuanVideoPipeline", - "HunyuanSkyreelsImageToVideoPipeline", - "HunyuanVideoImageToVideoPipeline", - "HunyuanVideoFramepackPipeline", - ] - _import_structure["kandinsky"] = [ - "KandinskyCombinedPipeline", - "KandinskyImg2ImgCombinedPipeline", - "KandinskyImg2ImgPipeline", - "KandinskyInpaintCombinedPipeline", - "KandinskyInpaintPipeline", - "KandinskyPipeline", - "KandinskyPriorPipeline", - ] - _import_structure["kandinsky2_2"] = [ - "KandinskyV22CombinedPipeline", - "KandinskyV22ControlnetImg2ImgPipeline", - "KandinskyV22ControlnetPipeline", - "KandinskyV22Img2ImgCombinedPipeline", - "KandinskyV22Img2ImgPipeline", - "KandinskyV22InpaintCombinedPipeline", - "KandinskyV22InpaintPipeline", - "KandinskyV22Pipeline", - "KandinskyV22PriorEmb2EmbPipeline", - "KandinskyV22PriorPipeline", - ] - _import_structure["kandinsky3"] = [ - "Kandinsky3Img2ImgPipeline", - "Kandinsky3Pipeline", - ] - _import_structure["latent_consistency_models"] = [ - "LatentConsistencyModelImg2ImgPipeline", - "LatentConsistencyModelPipeline", - ] - _import_structure["latent_diffusion"].extend(["LDMTextToImagePipeline"]) - _import_structure["ledits_pp"].extend( - [ - "LEditsPPPipelineStableDiffusion", - "LEditsPPPipelineStableDiffusionXL", - ] - ) - _import_structure["latte"] = ["LattePipeline"] - _import_structure["ltx"] = [ - "LTXPipeline", - "LTXImageToVideoPipeline", - "LTXConditionPipeline", - "LTXLatentUpsamplePipeline", - ] - _import_structure["lumina"] = ["LuminaPipeline", "LuminaText2ImgPipeline"] - _import_structure["lumina2"] = ["Lumina2Pipeline", "Lumina2Text2ImgPipeline"] - _import_structure["marigold"].extend( - [ - "MarigoldDepthPipeline", - "MarigoldIntrinsicsPipeline", - "MarigoldNormalsPipeline", - ] - ) - _import_structure["mochi"] = ["MochiPipeline"] - _import_structure["musicldm"] = ["MusicLDMPipeline"] - _import_structure["omnigen"] = ["OmniGenPipeline"] - _import_structure["visualcloze"] = ["VisualClozePipeline", "VisualClozeGenerationPipeline"] - _import_structure["paint_by_example"] = ["PaintByExamplePipeline"] - _import_structure["pia"] = ["PIAPipeline"] - _import_structure["pixart_alpha"] = ["PixArtAlphaPipeline", "PixArtSigmaPipeline"] - _import_structure["sana"] = [ - "SanaPipeline", - "SanaSprintPipeline", - "SanaControlNetPipeline", - "SanaSprintImg2ImgPipeline", - ] - _import_structure["semantic_stable_diffusion"] = ["SemanticStableDiffusionPipeline"] - _import_structure["shap_e"] = ["ShapEImg2ImgPipeline", "ShapEPipeline"] - _import_structure["stable_audio"] = [ - "StableAudioProjectionModel", - "StableAudioPipeline", - ] - _import_structure["stable_cascade"] = [ - "StableCascadeCombinedPipeline", - "StableCascadeDecoderPipeline", - "StableCascadePriorPipeline", - ] - _import_structure["stable_diffusion"].extend( - [ - "CLIPImageProjection", - "StableDiffusionDepth2ImgPipeline", - "StableDiffusionImageVariationPipeline", - "StableDiffusionImg2ImgPipeline", - "StableDiffusionInpaintPipeline", - "StableDiffusionInstructPix2PixPipeline", - "StableDiffusionLatentUpscalePipeline", - "StableDiffusionPipeline", - "StableDiffusionUpscalePipeline", - "StableUnCLIPImg2ImgPipeline", - "StableUnCLIPPipeline", - "StableDiffusionLDM3DPipeline", - ] - ) - _import_structure["aura_flow"] = ["AuraFlowPipeline"] - _import_structure["stable_diffusion_3"] = [ - "StableDiffusion3Pipeline", - "StableDiffusion3Img2ImgPipeline", - "StableDiffusion3InpaintPipeline", - ] - _import_structure["stable_diffusion_attend_and_excite"] = ["StableDiffusionAttendAndExcitePipeline"] - _import_structure["stable_diffusion_safe"] = ["StableDiffusionPipelineSafe"] - _import_structure["stable_diffusion_sag"] = ["StableDiffusionSAGPipeline"] - _import_structure["stable_diffusion_gligen"] = [ - "StableDiffusionGLIGENPipeline", - "StableDiffusionGLIGENTextImagePipeline", - ] - _import_structure["stable_video_diffusion"] = ["StableVideoDiffusionPipeline"] - _import_structure["stable_diffusion_xl"].extend( - [ - "StableDiffusionXLImg2ImgPipeline", - "StableDiffusionXLInpaintPipeline", - "StableDiffusionXLInstructPix2PixPipeline", - "StableDiffusionXLPipeline", - ] - ) - _import_structure["stable_diffusion_diffedit"] = ["StableDiffusionDiffEditPipeline"] - _import_structure["stable_diffusion_ldm3d"] = ["StableDiffusionLDM3DPipeline"] - _import_structure["stable_diffusion_panorama"] = ["StableDiffusionPanoramaPipeline"] - _import_structure["t2i_adapter"] = [ - "StableDiffusionAdapterPipeline", - "StableDiffusionXLAdapterPipeline", - ] - _import_structure["text_to_video_synthesis"] = [ - "TextToVideoSDPipeline", - "TextToVideoZeroPipeline", - "TextToVideoZeroSDXLPipeline", - "VideoToVideoSDPipeline", - ] - _import_structure["i2vgen_xl"] = ["I2VGenXLPipeline"] - _import_structure["unclip"] = ["UnCLIPImageVariationPipeline", "UnCLIPPipeline"] - _import_structure["unidiffuser"] = [ - "ImageTextPipelineOutput", - "UniDiffuserModel", - "UniDiffuserPipeline", - "UniDiffuserTextDecoder", - ] - _import_structure["wuerstchen"] = [ - "WuerstchenCombinedPipeline", - "WuerstchenDecoderPipeline", - "WuerstchenPriorPipeline", - ] - _import_structure["wan"] = ["WanPipeline", "WanImageToVideoPipeline", "WanVideoToVideoPipeline", "WanVACEPipeline"] -try: - if not is_onnx_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import dummy_onnx_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_onnx_objects)) -else: - _import_structure["onnx_utils"] = ["OnnxRuntimeModel"] -try: - if not (is_torch_available() and is_transformers_available() and is_onnx_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import dummy_torch_and_transformers_and_onnx_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_and_onnx_objects)) -else: - _import_structure["stable_diffusion"].extend( - [ - "OnnxStableDiffusionImg2ImgPipeline", - "OnnxStableDiffusionInpaintPipeline", - "OnnxStableDiffusionPipeline", - "OnnxStableDiffusionUpscalePipeline", - "StableDiffusionOnnxPipeline", - ] - ) - -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() -except OptionalDependencyNotAvailable: - from ..utils import ( - dummy_torch_and_transformers_and_sentencepiece_objects, - ) - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_and_sentencepiece_objects)) -else: - _import_structure["kolors"] = [ - "KolorsPipeline", - "KolorsImg2ImgPipeline", - ] - -try: - if not (is_torch_available() and is_transformers_available() and is_opencv_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import ( - dummy_torch_and_transformers_and_opencv_objects, - ) - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_and_opencv_objects)) -else: - _import_structure["consisid"] = ["ConsisIDPipeline"] - -try: - if not is_flax_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import dummy_flax_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_flax_objects)) -else: - _import_structure["pipeline_flax_utils"] = ["FlaxDiffusionPipeline"] -try: - if not (is_flax_available() and is_transformers_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import dummy_flax_and_transformers_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_flax_and_transformers_objects)) -else: - _import_structure["controlnet"].extend(["FlaxStableDiffusionControlNetPipeline"]) - _import_structure["stable_diffusion"].extend( - [ - "FlaxStableDiffusionImg2ImgPipeline", - "FlaxStableDiffusionInpaintPipeline", - "FlaxStableDiffusionPipeline", - ] - ) - _import_structure["stable_diffusion_xl"].extend( - [ - "FlaxStableDiffusionXLPipeline", - ] - ) - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - try: - if not is_torch_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_pt_objects import * # noqa F403 - - else: - from .auto_pipeline import ( - AutoPipelineForImage2Image, - AutoPipelineForInpainting, - AutoPipelineForText2Image, - ) - from .consistency_models import ConsistencyModelPipeline - from .dance_diffusion import DanceDiffusionPipeline - from .ddim import DDIMPipeline - from .ddpm import DDPMPipeline - from .deprecated import KarrasVePipeline, LDMPipeline, PNDMPipeline, RePaintPipeline, ScoreSdeVePipeline - from .dit import DiTPipeline - from .latent_diffusion import LDMSuperResolutionPipeline - from .pipeline_utils import ( - AudioPipelineOutput, - DiffusionPipeline, - ImagePipelineOutput, - StableDiffusionMixin, - ) - - try: - if not (is_torch_available() and is_librosa_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_torch_and_librosa_objects import * - else: - from .deprecated import AudioDiffusionPipeline, Mel - - try: - if not (is_torch_available() and is_transformers_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_torch_and_transformers_objects import * - else: - from .allegro import AllegroPipeline - from .amused import AmusedImg2ImgPipeline, AmusedInpaintPipeline, AmusedPipeline - from .animatediff import ( - AnimateDiffControlNetPipeline, - AnimateDiffPipeline, - AnimateDiffSDXLPipeline, - AnimateDiffSparseControlNetPipeline, - AnimateDiffVideoToVideoControlNetPipeline, - AnimateDiffVideoToVideoPipeline, - ) - from .audioldm import AudioLDMPipeline - from .audioldm2 import ( - AudioLDM2Pipeline, - AudioLDM2ProjectionModel, - AudioLDM2UNet2DConditionModel, - ) - from .aura_flow import AuraFlowPipeline - from .blip_diffusion import BlipDiffusionPipeline - from .chroma import ChromaImg2ImgPipeline, ChromaPipeline - from .cogvideo import ( - CogVideoXFunControlPipeline, - CogVideoXImageToVideoPipeline, - CogVideoXPipeline, - CogVideoXVideoToVideoPipeline, - ) - from .cogview3 import CogView3PlusPipeline - from .cogview4 import CogView4ControlPipeline, CogView4Pipeline - from .controlnet import ( - BlipDiffusionControlNetPipeline, - StableDiffusionControlNetImg2ImgPipeline, - StableDiffusionControlNetInpaintPipeline, - StableDiffusionControlNetPipeline, - StableDiffusionXLControlNetImg2ImgPipeline, - StableDiffusionXLControlNetInpaintPipeline, - StableDiffusionXLControlNetPipeline, - StableDiffusionXLControlNetUnionImg2ImgPipeline, - StableDiffusionXLControlNetUnionInpaintPipeline, - StableDiffusionXLControlNetUnionPipeline, - ) - from .controlnet_hunyuandit import ( - HunyuanDiTControlNetPipeline, - ) - from .controlnet_sd3 import StableDiffusion3ControlNetInpaintingPipeline, StableDiffusion3ControlNetPipeline - from .controlnet_xs import ( - StableDiffusionControlNetXSPipeline, - StableDiffusionXLControlNetXSPipeline, - ) - from .cosmos import ( - Cosmos2TextToImagePipeline, - Cosmos2VideoToWorldPipeline, - CosmosTextToWorldPipeline, - CosmosVideoToWorldPipeline, - ) - from .deepfloyd_if import ( - IFImg2ImgPipeline, - IFImg2ImgSuperResolutionPipeline, - IFInpaintingPipeline, - IFInpaintingSuperResolutionPipeline, - IFPipeline, - IFSuperResolutionPipeline, - ) - from .deprecated import ( - AltDiffusionImg2ImgPipeline, - AltDiffusionPipeline, - CycleDiffusionPipeline, - StableDiffusionInpaintPipelineLegacy, - StableDiffusionModelEditingPipeline, - StableDiffusionParadigmsPipeline, - StableDiffusionPix2PixZeroPipeline, - VersatileDiffusionDualGuidedPipeline, - VersatileDiffusionImageVariationPipeline, - VersatileDiffusionPipeline, - VersatileDiffusionTextToImagePipeline, - VQDiffusionPipeline, - ) - from .easyanimate import ( - EasyAnimateControlPipeline, - EasyAnimateInpaintPipeline, - EasyAnimatePipeline, - ) - from .flux import ( - FluxControlImg2ImgPipeline, - FluxControlInpaintPipeline, - FluxControlNetImg2ImgPipeline, - FluxControlNetInpaintPipeline, - FluxControlNetPipeline, - FluxControlPipeline, - FluxFillPipeline, - FluxImg2ImgPipeline, - FluxInpaintPipeline, - FluxPipeline, - FluxPriorReduxPipeline, - ReduxImageEncoder, - ) - from .hidream_image import HiDreamImagePipeline - from .hunyuan_video import ( - HunyuanSkyreelsImageToVideoPipeline, - HunyuanVideoFramepackPipeline, - HunyuanVideoImageToVideoPipeline, - HunyuanVideoPipeline, - ) - from .hunyuandit import HunyuanDiTPipeline - from .i2vgen_xl import I2VGenXLPipeline - from .kandinsky import ( - KandinskyCombinedPipeline, - KandinskyImg2ImgCombinedPipeline, - KandinskyImg2ImgPipeline, - KandinskyInpaintCombinedPipeline, - KandinskyInpaintPipeline, - KandinskyPipeline, - KandinskyPriorPipeline, - ) - from .kandinsky2_2 import ( - KandinskyV22CombinedPipeline, - KandinskyV22ControlnetImg2ImgPipeline, - KandinskyV22ControlnetPipeline, - KandinskyV22Img2ImgCombinedPipeline, - KandinskyV22Img2ImgPipeline, - KandinskyV22InpaintCombinedPipeline, - KandinskyV22InpaintPipeline, - KandinskyV22Pipeline, - KandinskyV22PriorEmb2EmbPipeline, - KandinskyV22PriorPipeline, - ) - from .kandinsky3 import ( - Kandinsky3Img2ImgPipeline, - Kandinsky3Pipeline, - ) - from .latent_consistency_models import ( - LatentConsistencyModelImg2ImgPipeline, - LatentConsistencyModelPipeline, - ) - from .latent_diffusion import LDMTextToImagePipeline - from .latte import LattePipeline - from .ledits_pp import ( - LEditsPPDiffusionPipelineOutput, - LEditsPPInversionPipelineOutput, - LEditsPPPipelineStableDiffusion, - LEditsPPPipelineStableDiffusionXL, - ) - from .ltx import LTXConditionPipeline, LTXImageToVideoPipeline, LTXLatentUpsamplePipeline, LTXPipeline - from .lumina import LuminaPipeline, LuminaText2ImgPipeline - from .lumina2 import Lumina2Pipeline, Lumina2Text2ImgPipeline - from .marigold import ( - MarigoldDepthPipeline, - MarigoldIntrinsicsPipeline, - MarigoldNormalsPipeline, - ) - from .mochi import MochiPipeline - from .musicldm import MusicLDMPipeline - from .omnigen import OmniGenPipeline - from .pag import ( - AnimateDiffPAGPipeline, - HunyuanDiTPAGPipeline, - KolorsPAGPipeline, - PixArtSigmaPAGPipeline, - SanaPAGPipeline, - StableDiffusion3PAGImg2ImgPipeline, - StableDiffusion3PAGPipeline, - StableDiffusionControlNetPAGInpaintPipeline, - StableDiffusionControlNetPAGPipeline, - StableDiffusionPAGImg2ImgPipeline, - StableDiffusionPAGInpaintPipeline, - StableDiffusionPAGPipeline, - StableDiffusionXLControlNetPAGImg2ImgPipeline, - StableDiffusionXLControlNetPAGPipeline, - StableDiffusionXLPAGImg2ImgPipeline, - StableDiffusionXLPAGInpaintPipeline, - StableDiffusionXLPAGPipeline, - ) - from .paint_by_example import PaintByExamplePipeline - from .pia import PIAPipeline - from .pixart_alpha import PixArtAlphaPipeline, PixArtSigmaPipeline - from .sana import SanaControlNetPipeline, SanaPipeline, SanaSprintImg2ImgPipeline, SanaSprintPipeline - from .semantic_stable_diffusion import SemanticStableDiffusionPipeline - from .shap_e import ShapEImg2ImgPipeline, ShapEPipeline - from .stable_audio import StableAudioPipeline, StableAudioProjectionModel - from .stable_cascade import ( - StableCascadeCombinedPipeline, - StableCascadeDecoderPipeline, - StableCascadePriorPipeline, - ) - from .stable_diffusion import ( - CLIPImageProjection, - StableDiffusionDepth2ImgPipeline, - StableDiffusionImageVariationPipeline, - StableDiffusionImg2ImgPipeline, - StableDiffusionInpaintPipeline, - StableDiffusionInstructPix2PixPipeline, - StableDiffusionLatentUpscalePipeline, - StableDiffusionPipeline, - StableDiffusionUpscalePipeline, - StableUnCLIPImg2ImgPipeline, - StableUnCLIPPipeline, - ) - from .stable_diffusion_3 import ( - StableDiffusion3Img2ImgPipeline, - StableDiffusion3InpaintPipeline, - StableDiffusion3Pipeline, - ) - from .stable_diffusion_attend_and_excite import StableDiffusionAttendAndExcitePipeline - from .stable_diffusion_diffedit import StableDiffusionDiffEditPipeline - from .stable_diffusion_gligen import StableDiffusionGLIGENPipeline, StableDiffusionGLIGENTextImagePipeline - from .stable_diffusion_ldm3d import StableDiffusionLDM3DPipeline - from .stable_diffusion_panorama import StableDiffusionPanoramaPipeline - from .stable_diffusion_safe import StableDiffusionPipelineSafe - from .stable_diffusion_sag import StableDiffusionSAGPipeline - from .stable_diffusion_xl import ( - StableDiffusionXLImg2ImgPipeline, - StableDiffusionXLInpaintPipeline, - StableDiffusionXLInstructPix2PixPipeline, - StableDiffusionXLPipeline, - ) - from .stable_video_diffusion import StableVideoDiffusionPipeline - from .t2i_adapter import ( - StableDiffusionAdapterPipeline, - StableDiffusionXLAdapterPipeline, - ) - from .text_to_video_synthesis import ( - TextToVideoSDPipeline, - TextToVideoZeroPipeline, - TextToVideoZeroSDXLPipeline, - VideoToVideoSDPipeline, - ) - from .unclip import UnCLIPImageVariationPipeline, UnCLIPPipeline - from .unidiffuser import ( - ImageTextPipelineOutput, - UniDiffuserModel, - UniDiffuserPipeline, - UniDiffuserTextDecoder, - ) - from .visualcloze import VisualClozeGenerationPipeline, VisualClozePipeline - from .wan import WanImageToVideoPipeline, WanPipeline, WanVACEPipeline, WanVideoToVideoPipeline - from .wuerstchen import ( - WuerstchenCombinedPipeline, - WuerstchenDecoderPipeline, - WuerstchenPriorPipeline, - ) - - try: - if not is_onnx_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_onnx_objects import * # noqa F403 - - else: - from .onnx_utils import OnnxRuntimeModel - - try: - if not (is_torch_available() and is_transformers_available() and is_onnx_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_torch_and_transformers_and_onnx_objects import * - else: - from .stable_diffusion import ( - OnnxStableDiffusionImg2ImgPipeline, - OnnxStableDiffusionInpaintPipeline, - OnnxStableDiffusionPipeline, - OnnxStableDiffusionUpscalePipeline, - 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() - except OptionalDependencyNotAvailable: - from ..utils.dummy_torch_and_transformers_and_sentencepiece_objects import * - else: - from .kolors import ( - KolorsImg2ImgPipeline, - KolorsPipeline, - ) - - try: - if not (is_torch_available() and is_transformers_available() and is_opencv_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_torch_and_transformers_and_opencv_objects import * - else: - from .consisid import ConsisIDPipeline - - try: - if not is_flax_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_flax_objects import * # noqa F403 - else: - from .pipeline_flax_utils import FlaxDiffusionPipeline - - try: - if not (is_flax_available() and is_transformers_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_flax_and_transformers_objects import * - else: - from .controlnet import FlaxStableDiffusionControlNetPipeline - from .stable_diffusion import ( - FlaxStableDiffusionImg2ImgPipeline, - FlaxStableDiffusionInpaintPipeline, - FlaxStableDiffusionPipeline, - ) - from .stable_diffusion_xl import ( - FlaxStableDiffusionXLPipeline, - ) - - try: - if not (is_transformers_available() and is_torch_available() and is_note_seq_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_transformers_and_torch_and_note_seq_objects import * # noqa F403 - - else: - from .deprecated import ( - MidiProcessor, - SpectrogramDiffusionPipeline, - ) - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 351156a139d2cca9c2caa4f937e26a13ed944fb0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/auto_pipeline.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/auto_pipeline.cpython-310.pyc deleted file mode 100644 index f08629bc99b97036e9bae10aca8cd4640349c464..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/auto_pipeline.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/free_init_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/free_init_utils.cpython-310.pyc deleted file mode 100644 index 1101fcf19bc341a23bd049aaf21acff8c7c87ff0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/free_init_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/free_noise_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/free_noise_utils.cpython-310.pyc deleted file mode 100644 index cec013d9a9968c164c4acb870549e579e7f67563..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/free_noise_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/onnx_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/onnx_utils.cpython-310.pyc deleted file mode 100644 index bd3ad7cbbb19086658a9a761d4d6a7f0df8876a1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/onnx_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/pipeline_flax_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/pipeline_flax_utils.cpython-310.pyc deleted file mode 100644 index 2c99a4c42369b49a44010d6ac59667e5eda8b7c4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/pipeline_flax_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/pipeline_loading_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/pipeline_loading_utils.cpython-310.pyc deleted file mode 100644 index 9e1664e1cbf582331f2b372ba717f4da1ef8e871..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/pipeline_loading_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/pipeline_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/pipeline_utils.cpython-310.pyc deleted file mode 100644 index 30859d5d33518ee6d66df95c367488d55fa93c62..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/pipeline_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/transformers_loading_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/transformers_loading_utils.cpython-310.pyc deleted file mode 100644 index 81f308a00604033b2b7357e674a64bfb1cd80219..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/__pycache__/transformers_loading_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/__init__.py deleted file mode 100644 index 2162b825e0a2d59aa79430fc0b66f664b6feccb5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -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_allegro"] = ["AllegroPipeline"] - -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_allegro import AllegroPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 2580b1d4c7718eed4ef270b2494e01eccad67178..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/__pycache__/pipeline_allegro.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/__pycache__/pipeline_allegro.cpython-310.pyc deleted file mode 100644 index 2fa2b9ebb24bc6a74f039fce05d49f62c84f624d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/__pycache__/pipeline_allegro.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index 94bf4ac61c539e308f013da4efc37177467f3cc8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/pipeline_allegro.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/pipeline_allegro.py deleted file mode 100644 index 0993c8b912b05876ea7b2aaf282a0def8fd283ce..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/pipeline_allegro.py +++ /dev/null @@ -1,958 +0,0 @@ -# Copyright 2025 The RhymesAI 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 inspect -import math -import re -import urllib.parse as ul -from typing import Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import T5EncoderModel, T5Tokenizer - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...models import AllegroTransformer3DModel, AutoencoderKLAllegro -from ...models.embeddings import get_3d_rotary_pos_embed_allegro -from ...pipelines.pipeline_utils import DiffusionPipeline -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - BACKENDS_MAPPING, - deprecate, - is_bs4_available, - 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_output import AllegroPipelineOutput - - -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__) - -if is_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import AutoencoderKLAllegro, AllegroPipeline - >>> from diffusers.utils import export_to_video - - >>> vae = AutoencoderKLAllegro.from_pretrained("rhymes-ai/Allegro", subfolder="vae", torch_dtype=torch.float32) - >>> pipe = AllegroPipeline.from_pretrained("rhymes-ai/Allegro", vae=vae, torch_dtype=torch.bfloat16).to("cuda") - >>> pipe.enable_vae_tiling() - - >>> prompt = ( - ... "A seaside harbor with bright sunlight and sparkling seawater, with many boats in the water. From an aerial view, " - ... "the boats vary in size and color, some moving and some stationary. Fishing boats in the water suggest that this " - ... "location might be a popular spot for docking fishing boats." - ... ) - >>> video = pipe(prompt, guidance_scale=7.5, max_sequence_length=512).frames[0] - >>> export_to_video(video, "output.mp4", fps=15) - ``` -""" - - -# 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 AllegroPipeline(DiffusionPipeline): - r""" - Pipeline for text-to-video generation using Allegro. - - 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.) - - Args: - vae ([`AllegroAutoEncoderKL3D`]): - Variational Auto-Encoder (VAE) Model to encode and decode video to and from latent representations. - text_encoder ([`T5EncoderModel`]): - Frozen text-encoder. PixArt-Alpha uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the - [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant. - tokenizer (`T5Tokenizer`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - transformer ([`AllegroTransformer3DModel`]): - A text conditioned `AllegroTransformer3DModel` to denoise the encoded video latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `transformer` to denoise the encoded video latents. - """ - - bad_punct_regex = re.compile( - r"[" - + "#®•©™&@·º½¾¿¡§~" - + r"\)" - + r"\(" - + r"\]" - + r"\[" - + r"\}" - + r"\{" - + r"\|" - + "\\" - + r"\/" - + r"\*" - + r"]{1,}" - ) # noqa - - _optional_components = [] - model_cpu_offload_seq = "text_encoder->transformer->vae" - - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - ] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - vae: AutoencoderKLAllegro, - transformer: AllegroTransformer3DModel, - scheduler: KarrasDiffusionSchedulers, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler - ) - self.vae_scale_factor_spatial = ( - 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - ) - self.vae_scale_factor_temporal = ( - self.vae.config.temporal_compression_ratio if getattr(self, "vae", None) else 4 - ) - - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - - # 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]], - 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, - clean_caption: bool = False, - max_sequence_length: int = 512, - **kwargs, - ): - 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 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 "". - 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 images that should be generated per prompt - device: (`torch.device`, *optional*): - 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. For PixArt-Alpha, it's should be the embeddings of the "" - string. - clean_caption (`bool`, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - max_sequence_length (`int`, defaults to 512): Maximum sequence length to use for the prompt. - """ - - if "mask_feature" in kwargs: - deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version." - deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False) - - if device is None: - device = self._execution_device - - # See Section 3.1. of the paper. - max_length = max_sequence_length - - if prompt_embeds is None: - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - add_special_tokens=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[:, max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because T5 can only handle sequences up to" - f" {max_length} tokens: {removed_text}" - ) - - prompt_attention_mask = text_inputs.attention_mask - prompt_attention_mask = prompt_attention_mask.to(device) - - prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask) - prompt_embeds = prompt_embeds[0] - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - elif self.transformer is not None: - dtype = self.transformer.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - - bs_embed, 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_videos_per_prompt, 1) - prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, seq_len, -1) - prompt_attention_mask = prompt_attention_mask.repeat(1, num_videos_per_prompt) - prompt_attention_mask = prompt_attention_mask.view(bs_embed * num_videos_per_prompt, -1) - - # get unconditional embeddings for classifier free guidance - if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens = [negative_prompt] * bs_embed if isinstance(negative_prompt, str) else negative_prompt - uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption) - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - negative_prompt_attention_mask = uncond_input.attention_mask - negative_prompt_attention_mask = negative_prompt_attention_mask.to(device) - - negative_prompt_embeds = self.text_encoder( - uncond_input.input_ids.to(device), attention_mask=negative_prompt_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=dtype, device=device) - - negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(bs_embed * num_videos_per_prompt, seq_len, -1) - - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(1, num_videos_per_prompt) - negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed * num_videos_per_prompt, -1) - else: - negative_prompt_embeds = None - negative_prompt_attention_mask = None - - return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - num_frames, - height, - width, - callback_on_step_end_tensor_inputs, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - ): - if num_frames <= 0: - raise ValueError(f"`num_frames` have to be positive but is {num_frames}.") - 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_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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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("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}." - ) - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - def prepare_latents( - self, batch_size, num_channels_latents, num_frames, 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." - ) - - if num_frames % 2 == 0: - num_frames = math.ceil(num_frames / self.vae_scale_factor_temporal) - else: - num_frames = math.ceil((num_frames - 1) / self.vae_scale_factor_temporal) + 1 - - shape = ( - batch_size, - num_channels_latents, - num_frames, - height // self.vae_scale_factor_spatial, - width // self.vae_scale_factor_spatial, - ) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def decode_latents(self, latents: torch.Tensor) -> torch.Tensor: - latents = 1 / self.vae.config.scaling_factor * latents - frames = self.vae.decode(latents).sample - frames = frames.permute(0, 2, 1, 3, 4) # [batch_size, channels, num_frames, height, width] - return frames - - def _prepare_rotary_positional_embeddings( - self, - batch_size: int, - height: int, - width: int, - num_frames: int, - device: torch.device, - ): - 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) - - start, stop = (0, 0), (grid_height, grid_width) - freqs_t, freqs_h, freqs_w, grid_t, grid_h, grid_w = get_3d_rotary_pos_embed_allegro( - embed_dim=self.transformer.config.attention_head_dim, - crops_coords=(start, stop), - grid_size=(grid_height, grid_width), - temporal_size=num_frames, - interpolation_scale=( - self.transformer.config.interpolation_scale_t, - self.transformer.config.interpolation_scale_h, - self.transformer.config.interpolation_scale_w, - ), - device=device, - ) - - grid_t = grid_t.to(dtype=torch.long) - grid_h = grid_h.to(dtype=torch.long) - grid_w = grid_w.to(dtype=torch.long) - - pos = torch.cartesian_prod(grid_t, grid_h, grid_w) - pos = pos.reshape(-1, 3).transpose(0, 1).reshape(3, 1, -1).contiguous() - grid_t, grid_h, grid_w = pos - - return (freqs_t, freqs_h, freqs_w), (grid_t, grid_h, grid_w) - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - @property - def guidance_scale(self): - return self._guidance_scale - - @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: str = "", - num_inference_steps: int = 100, - timesteps: List[int] = None, - guidance_scale: float = 7.5, - num_frames: Optional[int] = None, - height: Optional[int] = None, - width: Optional[int] = 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", - 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"], - clean_caption: bool = True, - max_sequence_length: int = 512, - ) -> Union[AllegroPipelineOutput, Tuple]: - """ - Function invoked when calling the pipeline for generation. - - Args: - 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*): - 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*): - 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): - 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 videos that are closely linked to - the text `prompt`, usually at the expense of lower video quality. - num_videos_per_prompt (`int`, *optional*, defaults to 1): - The number of videos to generate per prompt. - num_frames: (`int`, *optional*, defaults to 88): - The number controls the generated video frames. - height (`int`, *optional*, defaults to self.unet.config.sample_size): - The height in pixels of the generated video. - width (`int`, *optional*, defaults to self.unet.config.sample_size): - The width in pixels of the generated video. - 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*): - 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*): - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for video - tensor will ge 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.Tensor`, *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.Tensor`, *optional*): - Pre-generated attention mask for negative text embeddings. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate video. 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.IFPipelineOutput`] 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. - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - max_sequence_length (`int` defaults to `512`): - Maximum sequence length to use with the `prompt`. - - Examples: - - Returns: - [`~pipelines.allegro.pipeline_output.AllegroPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.allegro.pipeline_output.AllegroPipelineOutput`] is returned, - otherwise a `tuple` is returned where the first element is a list with the generated videos. - """ - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - num_videos_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - num_frames = num_frames or self.transformer.config.sample_frames * self.vae_scale_factor_temporal - height = height or self.transformer.config.sample_height * self.vae_scale_factor_spatial - width = width or self.transformer.config.sample_width * self.vae_scale_factor_spatial - - self.check_inputs( - prompt, - num_frames, - height, - width, - callback_on_step_end_tensor_inputs, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - ) - self._guidance_scale = guidance_scale - self._current_timestep = None - self._interrupt = False - - # 2. Default height and width to transformer - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt, - do_classifier_free_guidance, - negative_prompt=negative_prompt, - num_videos_per_prompt=num_videos_per_prompt, - device=device, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - ) - if 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) - if prompt_embeds.ndim == 3: - 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) - self.scheduler.set_timesteps(num_inference_steps, device=device) - - # 5. Prepare latents. - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - latent_channels, - num_frames, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Prepare rotary embeddings - image_rotary_emb = self._prepare_rotary_positional_embeddings( - batch_size, height, width, latents.size(2), device - ) - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - 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 - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - # predict noise model_output - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - encoder_attention_mask=prompt_attention_mask, - timestep=timestep, - image_rotary_emb=image_rotary_emb, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute previous image: x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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 == 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 not output_type == "latent": - latents = latents.to(self.vae.dtype) - video = self.decode_latents(latents) - video = video[:, :, :num_frames, :height, :width] - video = self.video_processor.postprocess_video(video=video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return AllegroPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/pipeline_output.py deleted file mode 100644 index 6a721783ca863864bd8b81fd829891ddd0a2c055..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/allegro/pipeline_output.py +++ /dev/null @@ -1,23 +0,0 @@ -from dataclasses import dataclass -from typing import List, Union - -import numpy as np -import PIL -import torch - -from diffusers.utils import BaseOutput - - -@dataclass -class AllegroPipelineOutput(BaseOutput): - r""" - 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 - 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]]] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__init__.py deleted file mode 100644 index 3c4d07a426b54fabfcdf35bfb8e4486cd828b3b3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__init__.py +++ /dev/null @@ -1,62 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - 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.dummy_torch_and_transformers_objects import ( - AmusedImg2ImgPipeline, - AmusedInpaintPipeline, - AmusedPipeline, - ) - - _dummy_objects.update( - { - "AmusedPipeline": AmusedPipeline, - "AmusedImg2ImgPipeline": AmusedImg2ImgPipeline, - "AmusedInpaintPipeline": AmusedInpaintPipeline, - } - ) -else: - _import_structure["pipeline_amused"] = ["AmusedPipeline"] - _import_structure["pipeline_amused_img2img"] = ["AmusedImg2ImgPipeline"] - _import_structure["pipeline_amused_inpaint"] = ["AmusedInpaintPipeline"] - - -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 ( - AmusedPipeline, - ) - else: - from .pipeline_amused import AmusedPipeline - from .pipeline_amused_img2img import AmusedImg2ImgPipeline - from .pipeline_amused_inpaint import AmusedInpaintPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 199a4fb9459834282dd1fe5bff971a5cfc7584cb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__pycache__/pipeline_amused.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__pycache__/pipeline_amused.cpython-310.pyc deleted file mode 100644 index c50d65391f1b02efc3a794df09c27283b250fdc6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__pycache__/pipeline_amused.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__pycache__/pipeline_amused_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__pycache__/pipeline_amused_img2img.cpython-310.pyc deleted file mode 100644 index f341044e64be6be5f9648a07317cc8b2a14b55ad..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__pycache__/pipeline_amused_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__pycache__/pipeline_amused_inpaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__pycache__/pipeline_amused_inpaint.cpython-310.pyc deleted file mode 100644 index 1050569caa045aaa8fa0b5f09205d90e592d1b28..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/__pycache__/pipeline_amused_inpaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/pipeline_amused.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/pipeline_amused.py deleted file mode 100644 index 131e34d1a4a138e9982125ef2225f09fb8b86fd2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/pipeline_amused.py +++ /dev/null @@ -1,342 +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 Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import CLIPTextModelWithProjection, CLIPTokenizer - -from ...image_processor import VaeImageProcessor -from ...models import UVit2DModel, VQModel -from ...schedulers import AmusedScheduler -from ...utils import is_torch_xla_available, replace_example_docstring -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, ImagePipelineOutput - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - - XLA_AVAILABLE = True -else: - XLA_AVAILABLE = False - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import AmusedPipeline - - >>> pipe = AmusedPipeline.from_pretrained("amused/amused-512", variant="fp16", torch_dtype=torch.float16) - >>> pipe = pipe.to("cuda") - - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> image = pipe(prompt).images[0] - ``` -""" - - -class AmusedPipeline(DeprecatedPipelineMixin, DiffusionPipeline): - _last_supported_version = "0.33.1" - image_processor: VaeImageProcessor - vqvae: VQModel - tokenizer: CLIPTokenizer - text_encoder: CLIPTextModelWithProjection - transformer: UVit2DModel - scheduler: AmusedScheduler - - model_cpu_offload_seq = "text_encoder->transformer->vqvae" - - def __init__( - self, - vqvae: VQModel, - tokenizer: CLIPTokenizer, - text_encoder: CLIPTextModelWithProjection, - transformer: UVit2DModel, - scheduler: AmusedScheduler, - ): - super().__init__() - - self.register_modules( - vqvae=vqvae, - tokenizer=tokenizer, - text_encoder=text_encoder, - transformer=transformer, - scheduler=scheduler, - ) - self.vae_scale_factor = ( - 2 ** (len(self.vqvae.config.block_out_channels) - 1) if getattr(self, "vqvae", None) else 8 - ) - self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_normalize=False) - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Optional[Union[List[str], str]] = None, - height: Optional[int] = None, - width: Optional[int] = 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, - output_type="pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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), - ): - """ - The call function to the pipeline for generation. - - Args: - 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. - 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 16): - 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 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*): - 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): - The number of images to generate per prompt. - generator (`torch.Generator`, *optional*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - latents (`torch.IntTensor`, *optional*): - Pre-generated tokens representing latent vectors in `self.vqvae`, to be used as inputs for image - generation. If not provided, the starting latents will be completely masked. - 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. A single vector from the - pooled and projected final hidden states. - encoder_hidden_states (`torch.Tensor`, *optional*): - Pre-generated penultimate hidden states from the text encoder providing additional text conditioning. - 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. - negative_encoder_hidden_states (`torch.Tensor`, *optional*): - Analogous to `encoder_hidden_states` for the positive prompt. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - micro_conditioning_aesthetic_score (`int`, *optional*, defaults to 6): - 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)): - 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)): - Configures the temperature scheduler on `self.scheduler` see `AmusedScheduler#set_timesteps`. - - Examples: - - Returns: - [`~pipelines.pipeline_utils.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.pipeline_utils.ImagePipelineOutput`] is returned, otherwise a - `tuple` is returned where the first element is a list with the generated images. - """ - if (prompt_embeds is not None and encoder_hidden_states is None) or ( - prompt_embeds is None and encoder_hidden_states is not None - ): - raise ValueError("pass either both `prompt_embeds` and `encoder_hidden_states` or neither") - - if (negative_prompt_embeds is not None and negative_encoder_hidden_states is None) or ( - negative_prompt_embeds is None and negative_encoder_hidden_states is not None - ): - raise ValueError( - "pass either both `negatve_prompt_embeds` and `negative_encoder_hidden_states` or neither" - ) - - if (prompt is None and prompt_embeds is None) or (prompt is not None and prompt_embeds is not None): - raise ValueError("pass only one of `prompt` or `prompt_embeds`") - - if isinstance(prompt, str): - prompt = [prompt] - - if prompt is not None: - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - batch_size = batch_size * num_images_per_prompt - - if height is None: - height = self.transformer.config.sample_size * self.vae_scale_factor - - if width is None: - width = self.transformer.config.sample_size * self.vae_scale_factor - - if prompt_embeds is None: - input_ids = self.tokenizer( - prompt, - return_tensors="pt", - padding="max_length", - truncation=True, - max_length=self.tokenizer.model_max_length, - ).input_ids.to(self._execution_device) - - outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True) - prompt_embeds = outputs.text_embeds - encoder_hidden_states = outputs.hidden_states[-2] - - prompt_embeds = prompt_embeds.repeat(num_images_per_prompt, 1) - encoder_hidden_states = encoder_hidden_states.repeat(num_images_per_prompt, 1, 1) - - if guidance_scale > 1.0: - if negative_prompt_embeds is None: - if negative_prompt is None: - negative_prompt = [""] * len(prompt) - - if isinstance(negative_prompt, str): - negative_prompt = [negative_prompt] - - input_ids = self.tokenizer( - negative_prompt, - return_tensors="pt", - padding="max_length", - truncation=True, - max_length=self.tokenizer.model_max_length, - ).input_ids.to(self._execution_device) - - outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True) - negative_prompt_embeds = outputs.text_embeds - negative_encoder_hidden_states = outputs.hidden_states[-2] - - negative_prompt_embeds = negative_prompt_embeds.repeat(num_images_per_prompt, 1) - negative_encoder_hidden_states = negative_encoder_hidden_states.repeat(num_images_per_prompt, 1, 1) - - prompt_embeds = torch.concat([negative_prompt_embeds, prompt_embeds]) - encoder_hidden_states = torch.concat([negative_encoder_hidden_states, encoder_hidden_states]) - - # Note that the micro conditionings _do_ flip the order of width, height for the original size - # and the crop coordinates. This is how it was done in the original code base - micro_conds = torch.tensor( - [ - width, - height, - micro_conditioning_crop_coord[0], - micro_conditioning_crop_coord[1], - micro_conditioning_aesthetic_score, - ], - device=self._execution_device, - dtype=encoder_hidden_states.dtype, - ) - micro_conds = micro_conds.unsqueeze(0) - micro_conds = micro_conds.expand(2 * batch_size if guidance_scale > 1.0 else batch_size, -1) - - shape = (batch_size, height // self.vae_scale_factor, width // self.vae_scale_factor) - - if latents is None: - latents = torch.full( - shape, self.scheduler.config.mask_token_id, dtype=torch.long, device=self._execution_device - ) - - self.scheduler.set_timesteps(num_inference_steps, temperature, self._execution_device) - - num_warmup_steps = len(self.scheduler.timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, timestep in enumerate(self.scheduler.timesteps): - if guidance_scale > 1.0: - model_input = torch.cat([latents] * 2) - else: - model_input = latents - - model_output = self.transformer( - model_input, - micro_conds=micro_conds, - pooled_text_emb=prompt_embeds, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - ) - - if guidance_scale > 1.0: - uncond_logits, cond_logits = model_output.chunk(2) - model_output = uncond_logits + guidance_scale * (cond_logits - uncond_logits) - - latents = self.scheduler.step( - model_output=model_output, - timestep=timestep, - sample=latents, - generator=generator, - ).prev_sample - - if i == len(self.scheduler.timesteps) - 1 or ( - (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 - ): - progress_bar.update() - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, timestep, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if output_type == "latent": - output = latents - else: - needs_upcasting = self.vqvae.dtype == torch.float16 and self.vqvae.config.force_upcast - - if needs_upcasting: - self.vqvae.float() - - output = self.vqvae.decode( - latents, - force_not_quantize=True, - shape=( - batch_size, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - self.vqvae.config.latent_channels, - ), - ).sample.clip(0, 1) - output = self.image_processor.postprocess(output, output_type) - - if needs_upcasting: - self.vqvae.half() - - self.maybe_free_model_hooks() - - if not return_dict: - return (output,) - - return ImagePipelineOutput(output) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/pipeline_amused_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/pipeline_amused_img2img.py deleted file mode 100644 index a122c12236dd0f620fe2d83ce182db119b88d796..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/pipeline_amused_img2img.py +++ /dev/null @@ -1,363 +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 Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import CLIPTextModelWithProjection, CLIPTokenizer - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...models import UVit2DModel, VQModel -from ...schedulers import AmusedScheduler -from ...utils import is_torch_xla_available, replace_example_docstring -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, ImagePipelineOutput - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - - XLA_AVAILABLE = True -else: - XLA_AVAILABLE = False - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import AmusedImg2ImgPipeline - >>> from diffusers.utils import load_image - - >>> pipe = AmusedImg2ImgPipeline.from_pretrained( - ... "amused/amused-512", variant="fp16", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = "winter mountains" - >>> input_image = ( - ... load_image( - ... "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/open_muse/mountains.jpg" - ... ) - ... .resize((512, 512)) - ... .convert("RGB") - ... ) - >>> image = pipe(prompt, input_image).images[0] - ``` -""" - - -class AmusedImg2ImgPipeline(DeprecatedPipelineMixin, DiffusionPipeline): - _last_supported_version = "0.33.1" - image_processor: VaeImageProcessor - vqvae: VQModel - tokenizer: CLIPTokenizer - text_encoder: CLIPTextModelWithProjection - transformer: UVit2DModel - scheduler: AmusedScheduler - - model_cpu_offload_seq = "text_encoder->transformer->vqvae" - - # TODO - when calling self.vqvae.quantize, it uses self.vqvae.quantize.embedding.weight before - # the forward method of self.vqvae.quantize, so the hook doesn't get called to move the parameter - # off the meta device. There should be a way to fix this instead of just not offloading it - _exclude_from_cpu_offload = ["vqvae"] - - def __init__( - self, - vqvae: VQModel, - tokenizer: CLIPTokenizer, - text_encoder: CLIPTextModelWithProjection, - transformer: UVit2DModel, - scheduler: AmusedScheduler, - ): - super().__init__() - - self.register_modules( - vqvae=vqvae, - tokenizer=tokenizer, - text_encoder=text_encoder, - transformer=transformer, - scheduler=scheduler, - ) - self.vae_scale_factor = ( - 2 ** (len(self.vqvae.config.block_out_channels) - 1) if getattr(self, "vqvae", None) else 8 - ) - self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_normalize=False) - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Optional[Union[List[str], str]] = 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, - output_type="pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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), - ): - """ - The call function to the pipeline for generation. - - Args: - 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`, 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. - strength (`float`, *optional*, defaults to 0.5): - 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`. - 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. - 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*): - 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): - The number of images to generate per prompt. - generator (`torch.Generator`, *optional*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - 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. A single vector from the - pooled and projected final hidden states. - encoder_hidden_states (`torch.Tensor`, *optional*): - Pre-generated penultimate hidden states from the text encoder providing additional text conditioning. - 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. - negative_encoder_hidden_states (`torch.Tensor`, *optional*): - Analogous to `encoder_hidden_states` for the positive prompt. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - micro_conditioning_aesthetic_score (`int`, *optional*, defaults to 6): - 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)): - 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)): - Configures the temperature scheduler on `self.scheduler` see `AmusedScheduler#set_timesteps`. - - Examples: - - Returns: - [`~pipelines.pipeline_utils.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.pipeline_utils.ImagePipelineOutput`] is returned, otherwise a - `tuple` is returned where the first element is a list with the generated images. - """ - - if (prompt_embeds is not None and encoder_hidden_states is None) or ( - prompt_embeds is None and encoder_hidden_states is not None - ): - raise ValueError("pass either both `prompt_embeds` and `encoder_hidden_states` or neither") - - if (negative_prompt_embeds is not None and negative_encoder_hidden_states is None) or ( - negative_prompt_embeds is None and negative_encoder_hidden_states is not None - ): - raise ValueError( - "pass either both `negative_prompt_embeds` and `negative_encoder_hidden_states` or neither" - ) - - if (prompt is None and prompt_embeds is None) or (prompt is not None and prompt_embeds is not None): - raise ValueError("pass only one of `prompt` or `prompt_embeds`") - - if isinstance(prompt, str): - prompt = [prompt] - - if prompt is not None: - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - batch_size = batch_size * num_images_per_prompt - - if prompt_embeds is None: - input_ids = self.tokenizer( - prompt, - return_tensors="pt", - padding="max_length", - truncation=True, - max_length=self.tokenizer.model_max_length, - ).input_ids.to(self._execution_device) - - outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True) - prompt_embeds = outputs.text_embeds - encoder_hidden_states = outputs.hidden_states[-2] - - prompt_embeds = prompt_embeds.repeat(num_images_per_prompt, 1) - encoder_hidden_states = encoder_hidden_states.repeat(num_images_per_prompt, 1, 1) - - if guidance_scale > 1.0: - if negative_prompt_embeds is None: - if negative_prompt is None: - negative_prompt = [""] * len(prompt) - - if isinstance(negative_prompt, str): - negative_prompt = [negative_prompt] - - input_ids = self.tokenizer( - negative_prompt, - return_tensors="pt", - padding="max_length", - truncation=True, - max_length=self.tokenizer.model_max_length, - ).input_ids.to(self._execution_device) - - outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True) - negative_prompt_embeds = outputs.text_embeds - negative_encoder_hidden_states = outputs.hidden_states[-2] - - negative_prompt_embeds = negative_prompt_embeds.repeat(num_images_per_prompt, 1) - negative_encoder_hidden_states = negative_encoder_hidden_states.repeat(num_images_per_prompt, 1, 1) - - prompt_embeds = torch.concat([negative_prompt_embeds, prompt_embeds]) - encoder_hidden_states = torch.concat([negative_encoder_hidden_states, encoder_hidden_states]) - - image = self.image_processor.preprocess(image) - - height, width = image.shape[-2:] - - # Note that the micro conditionings _do_ flip the order of width, height for the original size - # and the crop coordinates. This is how it was done in the original code base - micro_conds = torch.tensor( - [ - width, - height, - micro_conditioning_crop_coord[0], - micro_conditioning_crop_coord[1], - micro_conditioning_aesthetic_score, - ], - device=self._execution_device, - dtype=encoder_hidden_states.dtype, - ) - - micro_conds = micro_conds.unsqueeze(0) - micro_conds = micro_conds.expand(2 * batch_size if guidance_scale > 1.0 else batch_size, -1) - - self.scheduler.set_timesteps(num_inference_steps, temperature, self._execution_device) - num_inference_steps = int(len(self.scheduler.timesteps) * strength) - start_timestep_idx = len(self.scheduler.timesteps) - num_inference_steps - - needs_upcasting = self.vqvae.dtype == torch.float16 and self.vqvae.config.force_upcast - - if needs_upcasting: - self.vqvae.float() - - latents = self.vqvae.encode(image.to(dtype=self.vqvae.dtype, device=self._execution_device)).latents - latents_bsz, channels, latents_height, latents_width = latents.shape - latents = self.vqvae.quantize(latents)[2][2].reshape(latents_bsz, latents_height, latents_width) - latents = self.scheduler.add_noise( - latents, self.scheduler.timesteps[start_timestep_idx - 1], generator=generator - ) - latents = latents.repeat(num_images_per_prompt, 1, 1) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i in range(start_timestep_idx, len(self.scheduler.timesteps)): - timestep = self.scheduler.timesteps[i] - - if guidance_scale > 1.0: - model_input = torch.cat([latents] * 2) - else: - model_input = latents - - model_output = self.transformer( - model_input, - micro_conds=micro_conds, - pooled_text_emb=prompt_embeds, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - ) - - if guidance_scale > 1.0: - uncond_logits, cond_logits = model_output.chunk(2) - model_output = uncond_logits + guidance_scale * (cond_logits - uncond_logits) - - latents = self.scheduler.step( - model_output=model_output, - timestep=timestep, - sample=latents, - generator=generator, - ).prev_sample - - if i == len(self.scheduler.timesteps) - 1 or ((i + 1) % self.scheduler.order == 0): - progress_bar.update() - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, timestep, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if output_type == "latent": - output = latents - else: - output = self.vqvae.decode( - latents, - force_not_quantize=True, - shape=( - batch_size, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - self.vqvae.config.latent_channels, - ), - ).sample.clip(0, 1) - output = self.image_processor.postprocess(output, output_type) - - if needs_upcasting: - self.vqvae.half() - - self.maybe_free_model_hooks() - - if not return_dict: - return (output,) - - return ImagePipelineOutput(output) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/pipeline_amused_inpaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/pipeline_amused_inpaint.py deleted file mode 100644 index f4bd4944ff9a2772176c94bc044856b91de31fa2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/amused/pipeline_amused_inpaint.py +++ /dev/null @@ -1,394 +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 Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import CLIPTextModelWithProjection, CLIPTokenizer - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...models import UVit2DModel, VQModel -from ...schedulers import AmusedScheduler -from ...utils import is_torch_xla_available, replace_example_docstring -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, ImagePipelineOutput - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - - XLA_AVAILABLE = True -else: - XLA_AVAILABLE = False - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import AmusedInpaintPipeline - >>> from diffusers.utils import load_image - - >>> pipe = AmusedInpaintPipeline.from_pretrained( - ... "amused/amused-512", variant="fp16", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = "fall mountains" - >>> input_image = ( - ... load_image( - ... "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/open_muse/mountains_1.jpg" - ... ) - ... .resize((512, 512)) - ... .convert("RGB") - ... ) - >>> mask = ( - ... load_image( - ... "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/open_muse/mountains_1_mask.png" - ... ) - ... .resize((512, 512)) - ... .convert("L") - ... ) - >>> pipe(prompt, input_image, mask).images[0].save("out.png") - ``` -""" - - -class AmusedInpaintPipeline(DeprecatedPipelineMixin, DiffusionPipeline): - _last_supported_version = "0.33.1" - image_processor: VaeImageProcessor - vqvae: VQModel - tokenizer: CLIPTokenizer - text_encoder: CLIPTextModelWithProjection - transformer: UVit2DModel - scheduler: AmusedScheduler - - model_cpu_offload_seq = "text_encoder->transformer->vqvae" - - # TODO - when calling self.vqvae.quantize, it uses self.vqvae.quantize.embedding.weight before - # the forward method of self.vqvae.quantize, so the hook doesn't get called to move the parameter - # off the meta device. There should be a way to fix this instead of just not offloading it - _exclude_from_cpu_offload = ["vqvae"] - - def __init__( - self, - vqvae: VQModel, - tokenizer: CLIPTokenizer, - text_encoder: CLIPTextModelWithProjection, - transformer: UVit2DModel, - scheduler: AmusedScheduler, - ): - super().__init__() - - self.register_modules( - vqvae=vqvae, - tokenizer=tokenizer, - text_encoder=text_encoder, - transformer=transformer, - scheduler=scheduler, - ) - self.vae_scale_factor = ( - 2 ** (len(self.vqvae.config.block_out_channels) - 1) if getattr(self, "vqvae", None) else 8 - ) - self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_normalize=False) - self.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, - do_normalize=False, - do_binarize=True, - do_convert_grayscale=True, - do_resize=True, - ) - self.scheduler.register_to_config(masking_schedule="linear") - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Optional[Union[List[str], str]] = 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, - output_type="pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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), - ): - """ - The call function to the pipeline for generation. - - Args: - 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`, 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]`): - `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)`. - 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`. - num_inference_steps (`int`, *optional*, defaults to 16): - 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 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*): - 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): - The number of images to generate per prompt. - generator (`torch.Generator`, *optional*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - 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. A single vector from the - pooled and projected final hidden states. - encoder_hidden_states (`torch.Tensor`, *optional*): - Pre-generated penultimate hidden states from the text encoder providing additional text conditioning. - 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. - negative_encoder_hidden_states (`torch.Tensor`, *optional*): - Analogous to `encoder_hidden_states` for the positive prompt. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - micro_conditioning_aesthetic_score (`int`, *optional*, defaults to 6): - 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)): - 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)): - Configures the temperature scheduler on `self.scheduler` see `AmusedScheduler#set_timesteps`. - - Examples: - - Returns: - [`~pipelines.pipeline_utils.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.pipeline_utils.ImagePipelineOutput`] is returned, otherwise a - `tuple` is returned where the first element is a list with the generated images. - """ - - if (prompt_embeds is not None and encoder_hidden_states is None) or ( - prompt_embeds is None and encoder_hidden_states is not None - ): - raise ValueError("pass either both `prompt_embeds` and `encoder_hidden_states` or neither") - - if (negative_prompt_embeds is not None and negative_encoder_hidden_states is None) or ( - negative_prompt_embeds is None and negative_encoder_hidden_states is not None - ): - raise ValueError( - "pass either both `negatve_prompt_embeds` and `negative_encoder_hidden_states` or neither" - ) - - if (prompt is None and prompt_embeds is None) or (prompt is not None and prompt_embeds is not None): - raise ValueError("pass only one of `prompt` or `prompt_embeds`") - - if isinstance(prompt, str): - prompt = [prompt] - - if prompt is not None: - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - batch_size = batch_size * num_images_per_prompt - - if prompt_embeds is None: - input_ids = self.tokenizer( - prompt, - return_tensors="pt", - padding="max_length", - truncation=True, - max_length=self.tokenizer.model_max_length, - ).input_ids.to(self._execution_device) - - outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True) - prompt_embeds = outputs.text_embeds - encoder_hidden_states = outputs.hidden_states[-2] - - prompt_embeds = prompt_embeds.repeat(num_images_per_prompt, 1) - encoder_hidden_states = encoder_hidden_states.repeat(num_images_per_prompt, 1, 1) - - if guidance_scale > 1.0: - if negative_prompt_embeds is None: - if negative_prompt is None: - negative_prompt = [""] * len(prompt) - - if isinstance(negative_prompt, str): - negative_prompt = [negative_prompt] - - input_ids = self.tokenizer( - negative_prompt, - return_tensors="pt", - padding="max_length", - truncation=True, - max_length=self.tokenizer.model_max_length, - ).input_ids.to(self._execution_device) - - outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True) - negative_prompt_embeds = outputs.text_embeds - negative_encoder_hidden_states = outputs.hidden_states[-2] - - negative_prompt_embeds = negative_prompt_embeds.repeat(num_images_per_prompt, 1) - negative_encoder_hidden_states = negative_encoder_hidden_states.repeat(num_images_per_prompt, 1, 1) - - prompt_embeds = torch.concat([negative_prompt_embeds, prompt_embeds]) - encoder_hidden_states = torch.concat([negative_encoder_hidden_states, encoder_hidden_states]) - - image = self.image_processor.preprocess(image) - - height, width = image.shape[-2:] - - # Note that the micro conditionings _do_ flip the order of width, height for the original size - # and the crop coordinates. This is how it was done in the original code base - micro_conds = torch.tensor( - [ - width, - height, - micro_conditioning_crop_coord[0], - micro_conditioning_crop_coord[1], - micro_conditioning_aesthetic_score, - ], - device=self._execution_device, - dtype=encoder_hidden_states.dtype, - ) - - micro_conds = micro_conds.unsqueeze(0) - micro_conds = micro_conds.expand(2 * batch_size if guidance_scale > 1.0 else batch_size, -1) - - self.scheduler.set_timesteps(num_inference_steps, temperature, self._execution_device) - num_inference_steps = int(len(self.scheduler.timesteps) * strength) - start_timestep_idx = len(self.scheduler.timesteps) - num_inference_steps - - needs_upcasting = self.vqvae.dtype == torch.float16 and self.vqvae.config.force_upcast - - if needs_upcasting: - self.vqvae.float() - - latents = self.vqvae.encode(image.to(dtype=self.vqvae.dtype, device=self._execution_device)).latents - latents_bsz, channels, latents_height, latents_width = latents.shape - latents = self.vqvae.quantize(latents)[2][2].reshape(latents_bsz, latents_height, latents_width) - - mask = self.mask_processor.preprocess( - mask_image, height // self.vae_scale_factor, width // self.vae_scale_factor - ) - mask = mask.reshape(mask.shape[0], latents_height, latents_width).bool().to(latents.device) - latents[mask] = self.scheduler.config.mask_token_id - - starting_mask_ratio = mask.sum() / latents.numel() - - latents = latents.repeat(num_images_per_prompt, 1, 1) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i in range(start_timestep_idx, len(self.scheduler.timesteps)): - timestep = self.scheduler.timesteps[i] - - if guidance_scale > 1.0: - model_input = torch.cat([latents] * 2) - else: - model_input = latents - - model_output = self.transformer( - model_input, - micro_conds=micro_conds, - pooled_text_emb=prompt_embeds, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - ) - - if guidance_scale > 1.0: - uncond_logits, cond_logits = model_output.chunk(2) - model_output = uncond_logits + guidance_scale * (cond_logits - uncond_logits) - - latents = self.scheduler.step( - model_output=model_output, - timestep=timestep, - sample=latents, - generator=generator, - starting_mask_ratio=starting_mask_ratio, - ).prev_sample - - if i == len(self.scheduler.timesteps) - 1 or ((i + 1) % self.scheduler.order == 0): - progress_bar.update() - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, timestep, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if output_type == "latent": - output = latents - else: - output = self.vqvae.decode( - latents, - force_not_quantize=True, - shape=( - batch_size, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - self.vqvae.config.latent_channels, - ), - ).sample.clip(0, 1) - output = self.image_processor.postprocess(output, output_type) - - if needs_upcasting: - self.vqvae.half() - - self.maybe_free_model_hooks() - - if not return_dict: - return (output,) - - return ImagePipelineOutput(output) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__init__.py deleted file mode 100644 index d916abf2d85dfeee217479f14711dc2d33002ae6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__init__.py +++ /dev/null @@ -1,57 +0,0 @@ -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 = {"pipeline_output": ["AnimateDiffPipelineOutput"]} - -try: - if not (is_transformers_available() and is_torch_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_torch_and_transformers_objects - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) -else: - _import_structure["pipeline_animatediff"] = ["AnimateDiffPipeline"] - _import_structure["pipeline_animatediff_controlnet"] = ["AnimateDiffControlNetPipeline"] - _import_structure["pipeline_animatediff_sdxl"] = ["AnimateDiffSDXLPipeline"] - _import_structure["pipeline_animatediff_sparsectrl"] = ["AnimateDiffSparseControlNetPipeline"] - _import_structure["pipeline_animatediff_video2video"] = ["AnimateDiffVideoToVideoPipeline"] - _import_structure["pipeline_animatediff_video2video_controlnet"] = ["AnimateDiffVideoToVideoControlNetPipeline"] - -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_animatediff import AnimateDiffPipeline - from .pipeline_animatediff_controlnet import AnimateDiffControlNetPipeline - from .pipeline_animatediff_sdxl import AnimateDiffSDXLPipeline - from .pipeline_animatediff_sparsectrl import AnimateDiffSparseControlNetPipeline - from .pipeline_animatediff_video2video import AnimateDiffVideoToVideoPipeline - from .pipeline_animatediff_video2video_controlnet import AnimateDiffVideoToVideoControlNetPipeline - from .pipeline_output import AnimateDiffPipelineOutput - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 4b098ff1837dddf7ac7e9fbf299e95b6b83a97f4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff.cpython-310.pyc deleted file mode 100644 index 21cff018a1ac01d74132b959c18b7fba770c4b8c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_controlnet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_controlnet.cpython-310.pyc deleted file mode 100644 index a8811680e5d65b5cfbf97d1a9c135adea86a64bf..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_controlnet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_sdxl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_sdxl.cpython-310.pyc deleted file mode 100644 index 4cafb685b73b86f704549c59e7a89ac71a455205..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_sdxl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_sparsectrl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_sparsectrl.cpython-310.pyc deleted file mode 100644 index 2afe763e232d985badf2cb6a7985529511df7b51..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_sparsectrl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_video2video.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_video2video.cpython-310.pyc deleted file mode 100644 index a78be5cf4f228916cb5748dbc99b42d10d1fbf13..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_video2video.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_video2video_controlnet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_video2video_controlnet.cpython-310.pyc deleted file mode 100644 index 26225ea01d11ca2c489c01c6298daa574f76f19f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_animatediff_video2video_controlnet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index 98a5dffb6b7447e2b5eac3db00cc27216f1c89f8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff.py deleted file mode 100644 index 091b6db713bade2a405966deafd6b8dcbe547948..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff.py +++ /dev/null @@ -1,873 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...image_processor import PipelineImageInput -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel, UNetMotionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...models.unets.unet_motion_model import MotionAdapter -from ...schedulers import ( - DDIMScheduler, - DPMSolverMultistepScheduler, - EulerAncestralDiscreteScheduler, - EulerDiscreteScheduler, - LMSDiscreteScheduler, - PNDMScheduler, -) -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import randn_tensor -from ...video_processor import VideoProcessor -from ..free_init_utils import FreeInitMixin -from ..free_noise_utils import AnimateDiffFreeNoiseMixin -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from .pipeline_output import AnimateDiffPipelineOutput - - -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 MotionAdapter, AnimateDiffPipeline, DDIMScheduler - >>> from diffusers.utils import export_to_gif - - >>> adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-2") - >>> pipe = AnimateDiffPipeline.from_pretrained("frankjoshua/toonyou_beta6", motion_adapter=adapter) - >>> pipe.scheduler = DDIMScheduler(beta_schedule="linear", steps_offset=1, clip_sample=False) - >>> output = pipe(prompt="A corgi walking in the park") - >>> frames = output.frames[0] - >>> export_to_gif(frames, "animation.gif") - ``` -""" - - -class AnimateDiffPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - FreeInitMixin, - AnimateDiffFreeNoiseMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-to-video generation. - - 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.). - - 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 - - [`~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 ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer (`CLIPTokenizer`): - A [`~transformers.CLIPTokenizer`] to tokenize text. - unet ([`UNet2DConditionModel`]): - 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. - 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`]. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["feature_extractor", "image_encoder", "motion_adapter"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], - motion_adapter: MotionAdapter, - scheduler: Union[ - DDIMScheduler, - PNDMScheduler, - LMSDiscreteScheduler, - EulerDiscreteScheduler, - EulerAncestralDiscreteScheduler, - DPMSolverMultistepScheduler, - ], - feature_extractor: CLIPImageProcessor = None, - image_encoder: CLIPVisionModelWithProjection = None, - ): - super().__init__() - if isinstance(unet, UNet2DConditionModel): - unet = UNetMotionModel.from_unet2d(unet, motion_adapter) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - motion_adapter=motion_adapter, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - self.video_processor = VideoProcessor(do_resize=False, vae_scale_factor=self.vae_scale_factor) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt with num_images_per_prompt -> num_videos_per_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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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 decode_latents(self, latents, decode_chunk_size: int = 16): - latents = 1 / self.vae.config.scaling_factor * latents - - batch_size, channels, num_frames, height, width = latents.shape - latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) - - video = [] - for i in range(0, latents.shape[0], decode_chunk_size): - batch_latents = latents[i : i + decode_chunk_size] - batch_latents = self.vae.decode(batch_latents).sample - video.append(batch_latents) - - video = torch.cat(video) - video = video[None, :].reshape((batch_size, num_frames, -1) + video.shape[2:]).permute(0, 2, 1, 3, 4) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - video = video.float() - return video - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_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, list, dict)): - raise ValueError(f"`prompt` has to be of type `str`, `list` 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 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" - ) - - def prepare_latents( - self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None - ): - # If FreeNoise is enabled, generate latents as described in Equation (7) of [FreeNoise](https://huggingface.co/papers/2310.15169) - if self.free_noise_enabled: - latents = self._prepare_latents_free_noise( - batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents - ) - - 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." - ) - - shape = ( - batch_size, - num_channels_latents, - num_frames, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_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: Optional[Union[str, List[str]]] = None, - num_frames: Optional[int] = 16, - 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_videos_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, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "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"], - decode_chunk_size: int = 16, - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The width in pixels of the generated video. - num_frames (`int`, *optional*, defaults to 16): - The number of video frames that are generated. Defaults to 16 frames which at 8 frames per seconds - amounts to 2 seconds of video. - 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. - 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*): - 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*): - 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 video - 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`. Latents should be of shape - `(batch_size, num_channel, num_frames, height, width)`. - 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. - 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. - 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)`. 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. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated video. Choose between `torch.Tensor`, `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.text_to_video_synthesis.TextToVideoSDPipelineOutput`] instead - of a plain tuple. - 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). - 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. - 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. - decode_chunk_size (`int`, defaults to `16`): - The number of frames to decode at a time when calling `decode_latents` method. - - Examples: - - Returns: - [`~pipelines.animatediff.pipeline_output.AnimateDiffPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.animatediff.pipeline_output.AnimateDiffPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is a list with the generated frames. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - - # 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 - - num_videos_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._interrupt = False - - # 2. Define call parameters - if prompt is not None and isinstance(prompt, (str, dict)): - 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 - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - if self.free_noise_enabled: - prompt_embeds, negative_prompt_embeds = self._encode_prompt_free_noise( - prompt=prompt, - num_frames=num_frames, - device=device, - num_videos_per_prompt=num_videos_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - else: - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_videos_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - prompt_embeds = prompt_embeds.repeat_interleave(repeats=num_frames, dim=0) - - 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_videos_per_prompt, - self.do_classifier_free_guidance, - ) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - num_channels_latents, - num_frames, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - num_free_init_iters = self._free_init_num_iters if self.free_init_enabled else 1 - for free_init_iter in range(num_free_init_iters): - if self.free_init_enabled: - latents, timesteps = self._apply_free_init( - latents, free_init_iter, num_inference_steps, device, latents.dtype, generator - ) - - self._num_timesteps = len(timesteps) - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - # 8. Denoising loop - with self.progress_bar(total=self._num_timesteps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - ).sample - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - 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 callback is not None and i % callback_steps == 0: - callback(i, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # 9. Post processing - if output_type == "latent": - video = latents - else: - video_tensor = self.decode_latents(latents, decode_chunk_size) - video = self.video_processor.postprocess_video(video=video_tensor, output_type=output_type) - - # 10. Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return AnimateDiffPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py deleted file mode 100644 index 70180ccf06505000949091e69cb869f8cb486314..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py +++ /dev/null @@ -1,1118 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -import torch.nn.functional as F -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...image_processor import PipelineImageInput -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import ( - AutoencoderKL, - ControlNetModel, - ImageProjection, - MultiControlNetModel, - UNet2DConditionModel, - UNetMotionModel, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...models.unets.unet_motion_model import MotionAdapter -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import USE_PEFT_BACKEND, is_torch_xla_available, logging, scale_lora_layers, unscale_lora_layers -from ...utils.torch_utils import is_compiled_module, randn_tensor -from ...video_processor import VideoProcessor -from ..free_init_utils import FreeInitMixin -from ..free_noise_utils import AnimateDiffFreeNoiseMixin -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from .pipeline_output import AnimateDiffPipelineOutput - - -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 ( - ... AnimateDiffControlNetPipeline, - ... AutoencoderKL, - ... ControlNetModel, - ... MotionAdapter, - ... LCMScheduler, - ... ) - >>> from diffusers.utils import export_to_gif, load_video - - >>> # Additionally, you will need a preprocess videos before they can be used with the ControlNet - >>> # HF maintains just the right package for it: `pip install controlnet_aux` - >>> from controlnet_aux.processor import ZoeDetector - - >>> # Download controlnets from https://huggingface.co/lllyasviel/ControlNet-v1-1 to use .from_single_file - >>> # Download Diffusers-format controlnets, such as https://huggingface.co/lllyasviel/sd-controlnet-depth, to use .from_pretrained() - >>> controlnet = ControlNetModel.from_single_file("control_v11f1p_sd15_depth.pth", torch_dtype=torch.float16) - - >>> # We use AnimateLCM for this example but one can use the original motion adapters as well (for example, https://huggingface.co/guoyww/animatediff-motion-adapter-v1-5-3) - >>> motion_adapter = MotionAdapter.from_pretrained("wangfuyun/AnimateLCM") - - >>> vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16) - >>> pipe: AnimateDiffControlNetPipeline = AnimateDiffControlNetPipeline.from_pretrained( - ... "SG161222/Realistic_Vision_V5.1_noVAE", - ... motion_adapter=motion_adapter, - ... controlnet=controlnet, - ... vae=vae, - ... ).to(device="cuda", dtype=torch.float16) - >>> pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, beta_schedule="linear") - >>> pipe.load_lora_weights( - ... "wangfuyun/AnimateLCM", weight_name="AnimateLCM_sd15_t2v_lora.safetensors", adapter_name="lcm-lora" - ... ) - >>> pipe.set_adapters(["lcm-lora"], [0.8]) - - >>> depth_detector = ZoeDetector.from_pretrained("lllyasviel/Annotators").to("cuda") - >>> video = load_video( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/animatediff-vid2vid-input-1.gif" - ... ) - >>> conditioning_frames = [] - - >>> with pipe.progress_bar(total=len(video)) as progress_bar: - ... for frame in video: - ... conditioning_frames.append(depth_detector(frame)) - ... progress_bar.update() - - >>> prompt = "a panda, playing a guitar, sitting in a pink boat, in the ocean, mountains in background, realistic, high quality" - >>> negative_prompt = "bad quality, worst quality" - - >>> video = pipe( - ... prompt=prompt, - ... negative_prompt=negative_prompt, - ... num_frames=len(video), - ... num_inference_steps=10, - ... guidance_scale=2.0, - ... conditioning_frames=conditioning_frames, - ... generator=torch.Generator().manual_seed(42), - ... ).frames[0] - - >>> export_to_gif(video, "animatediff_controlnet.gif", fps=8) - ``` -""" - - -class AnimateDiffControlNetPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - FreeInitMixin, - AnimateDiffFreeNoiseMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-to-video generation with ControlNet guidance. - - 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.). - - 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 - - [`~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 ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer (`CLIPTokenizer`): - A [`~transformers.CLIPTokenizer`] to tokenize text. - unet ([`UNet2DConditionModel`]): - 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. - 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`]. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["feature_extractor", "image_encoder"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], - motion_adapter: MotionAdapter, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], - scheduler: KarrasDiffusionSchedulers, - feature_extractor: Optional[CLIPImageProcessor] = None, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, - ): - super().__init__() - if isinstance(unet, UNet2DConditionModel): - unet = UNetMotionModel.from_unet2d(unet, motion_adapter) - - if isinstance(controlnet, (list, tuple)): - controlnet = MultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - motion_adapter=motion_adapter, - controlnet=controlnet, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor) - self.control_video_processor = VideoProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt with num_images_per_prompt -> num_videos_per_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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.animatediff.pipeline_animatediff.AnimateDiffPipeline.decode_latents - def decode_latents(self, latents, decode_chunk_size: int = 16): - latents = 1 / self.vae.config.scaling_factor * latents - - batch_size, channels, num_frames, height, width = latents.shape - latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) - - video = [] - for i in range(0, latents.shape[0], decode_chunk_size): - batch_latents = latents[i : i + decode_chunk_size] - batch_latents = self.vae.decode(batch_latents).sample - video.append(batch_latents) - - video = torch.cat(video) - video = video[None, :].reshape((batch_size, num_frames, -1) + video.shape[2:]).permute(0, 2, 1, 3, 4) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - video = video.float() - return video - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - num_frames, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - video=None, - controlnet_conditioning_scale=1.0, - control_guidance_start=0.0, - control_guidance_end=1.0, - ): - 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_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, list, dict)): - raise ValueError(f"`prompt` has to be of type `str`, `list` 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}." - ) - - # `prompt` needs more sophisticated handling when there are multiple - # conditionings. - if isinstance(self.controlnet, MultiControlNetModel): - if isinstance(prompt, list): - 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` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - if not isinstance(video, list): - raise TypeError(f"For single controlnet, `image` must be of type `list` but got {type(video)}") - if len(video) != num_frames: - raise ValueError(f"Excepted image to have length {num_frames} but got {len(video)=}") - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if not isinstance(video, list) or not isinstance(video[0], list): - raise TypeError(f"For multiple controlnets: `image` must be type list of lists but got {type(video)=}") - if len(video[0]) != num_frames: - raise ValueError(f"Expected length of image sublist as {num_frames} but got {len(video[0])=}") - if any(len(img) != len(video[0]) for img in video): - raise ValueError("All conditioning frame batches for multicontrolnet must be same size") - else: - assert False - - # Check `controlnet_conditioning_scale` - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if isinstance(controlnet_conditioning_scale, list): - if any(isinstance(i, list) for i in controlnet_conditioning_scale): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - elif 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" - ) - else: - assert False - - if not isinstance(control_guidance_start, (tuple, list)): - control_guidance_start = [control_guidance_start] - - if not isinstance(control_guidance_end, (tuple, list)): - control_guidance_end = [control_guidance_end] - - 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(self.controlnet, MultiControlNetModel): - 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.") - - # Copied from diffusers.pipelines.animatediff.pipeline_animatediff.AnimateDiffPipeline.prepare_latents - def prepare_latents( - self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None - ): - # If FreeNoise is enabled, generate latents as described in Equation (7) of [FreeNoise](https://huggingface.co/papers/2310.15169) - if self.free_noise_enabled: - latents = self._prepare_latents_free_noise( - batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents - ) - - 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." - ) - - shape = ( - batch_size, - num_channels_latents, - num_frames, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def prepare_video( - self, - video, - width, - height, - batch_size, - num_videos_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - video = self.control_video_processor.preprocess_video(video, height=height, width=width).to( - dtype=torch.float32 - ) - video = video.permute(0, 2, 1, 3, 4).flatten(0, 1) - video_batch_size = video.shape[0] - - if video_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_videos_per_prompt - - video = video.repeat_interleave(repeat_by, dim=0) - video = video.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - video = torch.cat([video] * 2) - - return video - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @property - def interrupt(self): - return self._interrupt - - @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, - 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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"], - decode_chunk_size: int = 16, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The width in pixels of the generated video. - num_frames (`int`, *optional*, defaults to 16): - The number of video frames that are generated. Defaults to 16 frames which at 8 frames per seconds - amounts to 2 seconds of video. - 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. - 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*): - 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*): - 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 video - 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`. Latents should be of shape - `(batch_size, num_channel, num_frames, height, width)`. - 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. - 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. - 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)`. 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*): - 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. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated video. Choose between `torch.Tensor`, `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.text_to_video_synthesis.TextToVideoSDPipelineOutput`] instead - of a plain tuple. - 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): - 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): - The percentage of total steps at which the ControlNet starts applying. - 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 - the output of the pre-final layer will be used for computing the prompt embeddings. - 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. - - Examples: - - Returns: - [`~pipelines.animatediff.pipeline_output.AnimateDiffPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.animatediff.pipeline_output.AnimateDiffPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is a list with the generated frames. - """ - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 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 - - num_videos_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt=prompt, - height=height, - width=width, - num_frames=num_frames, - negative_prompt=negative_prompt, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - video=conditioning_frames, - controlnet_conditioning_scale=controlnet_conditioning_scale, - control_guidance_start=control_guidance_start, - control_guidance_end=control_guidance_end, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._interrupt = False - - # 2. Define call parameters - if prompt is not None and isinstance(prompt, (str, dict)): - 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 isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): - controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) - - global_pool_conditions = ( - controlnet.config.global_pool_conditions - if isinstance(controlnet, ControlNetModel) - else controlnet.nets[0].config.global_pool_conditions - ) - guess_mode = guess_mode or global_pool_conditions - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - if self.free_noise_enabled: - prompt_embeds, negative_prompt_embeds = self._encode_prompt_free_noise( - prompt=prompt, - num_frames=num_frames, - device=device, - num_videos_per_prompt=num_videos_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - else: - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_videos_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - prompt_embeds = prompt_embeds.repeat_interleave(repeats=num_frames, dim=0) - - 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_videos_per_prompt, - self.do_classifier_free_guidance, - ) - - if isinstance(controlnet, ControlNetModel): - conditioning_frames = self.prepare_video( - video=conditioning_frames, - width=width, - height=height, - batch_size=batch_size * num_videos_per_prompt * num_frames, - num_videos_per_prompt=num_videos_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - elif isinstance(controlnet, MultiControlNetModel): - cond_prepared_videos = [] - for frame_ in conditioning_frames: - prepared_video = self.prepare_video( - video=frame_, - width=width, - height=height, - batch_size=batch_size * num_videos_per_prompt * num_frames, - num_videos_per_prompt=num_videos_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - cond_prepared_videos.append(prepared_video) - conditioning_frames = cond_prepared_videos - else: - assert False - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - num_channels_latents, - num_frames, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - # 7.1 Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) - - num_free_init_iters = self._free_init_num_iters if self.free_init_enabled else 1 - for free_init_iter in range(num_free_init_iters): - if self.free_init_enabled: - latents, timesteps = self._apply_free_init( - latents, free_init_iter, num_inference_steps, device, latents.dtype, generator - ) - - self._num_timesteps = len(timesteps) - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - # 8. Denoising loop - with self.progress_bar(total=self._num_timesteps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - if guess_mode and self.do_classifier_free_guidance: - # Infer ControlNet only for the conditional batch. - control_model_input = latents - control_model_input = self.scheduler.scale_model_input(control_model_input, t) - controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] - else: - control_model_input = latent_model_input - controlnet_prompt_embeds = prompt_embeds - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - control_model_input = torch.transpose(control_model_input, 1, 2) - control_model_input = control_model_input.reshape( - (-1, control_model_input.shape[2], control_model_input.shape[3], control_model_input.shape[4]) - ) - - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=conditioning_frames, - conditioning_scale=cond_scale, - guess_mode=guess_mode, - return_dict=False, - ) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - ).sample - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - 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() - - # 9. Post processing - if output_type == "latent": - video = latents - else: - video_tensor = self.decode_latents(latents, decode_chunk_size) - video = self.video_processor.postprocess_video(video=video_tensor, output_type=output_type) - - # 10. Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return AnimateDiffPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py deleted file mode 100644 index 260669ddaf51c49251c909ea91f2420d419c7840..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py +++ /dev/null @@ -1,1309 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ...image_processor import PipelineImageInput -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ImageProjection, MotionAdapter, UNet2DConditionModel, UNetMotionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - FusedAttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import ( - DDIMScheduler, - DPMSolverMultistepScheduler, - EulerAncestralDiscreteScheduler, - EulerDiscreteScheduler, - LMSDiscreteScheduler, - PNDMScheduler, -) -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 ...video_processor import VideoProcessor -from ..free_init_utils import FreeInitMixin -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from .pipeline_output import AnimateDiffPipelineOutput - - -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.models import MotionAdapter - >>> from diffusers import AnimateDiffSDXLPipeline, DDIMScheduler - >>> from diffusers.utils import export_to_gif - - >>> adapter = MotionAdapter.from_pretrained( - ... "a-r-r-o-w/animatediff-motion-adapter-sdxl-beta", torch_dtype=torch.float16 - ... ) - - >>> model_id = "stabilityai/stable-diffusion-xl-base-1.0" - >>> scheduler = DDIMScheduler.from_pretrained( - ... model_id, - ... subfolder="scheduler", - ... clip_sample=False, - ... timestep_spacing="linspace", - ... beta_schedule="linear", - ... steps_offset=1, - ... ) - >>> pipe = AnimateDiffSDXLPipeline.from_pretrained( - ... model_id, - ... motion_adapter=adapter, - ... scheduler=scheduler, - ... torch_dtype=torch.float16, - ... variant="fp16", - ... ).to("cuda") - - >>> # enable memory savings - >>> pipe.enable_vae_slicing() - >>> pipe.enable_vae_tiling() - - >>> output = pipe( - ... prompt="a panda surfing in the ocean, realistic, high quality", - ... negative_prompt="low quality, worst quality", - ... num_inference_steps=20, - ... guidance_scale=8, - ... width=1024, - ... height=1024, - ... num_frames=16, - ... ) - - >>> frames = output.frames[0] - >>> export_to_gif(frames, "animation.gif") - ``` -""" - - -# 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 - - -# 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 AnimateDiffSDXLPipeline( - DiffusionPipeline, - StableDiffusionMixin, - FromSingleFileMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - FreeInitMixin, -): - r""" - Pipeline for text-to-video generation using Stable Diffusion XL. - - 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->image_encoder->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "image_encoder", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "negative_add_time_ids", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], - motion_adapter: MotionAdapter, - scheduler: Union[ - DDIMScheduler, - PNDMScheduler, - LMSDiscreteScheduler, - EulerDiscreteScheduler, - EulerAncestralDiscreteScheduler, - DPMSolverMultistepScheduler, - ], - image_encoder: CLIPVisionModelWithProjection = None, - feature_extractor: CLIPImageProcessor = None, - force_zeros_for_empty_prompt: bool = True, - ): - super().__init__() - - if isinstance(unet, UNet2DConditionModel): - unet = UNetMotionModel.from_unet2d(unet, motion_adapter) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - motion_adapter=motion_adapter, - scheduler=scheduler, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - ) - 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.video_processor = VideoProcessor(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 - ) - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.encode_prompt with num_images_per_prompt->num_videos_per_prompt - def encode_prompt( - self, - prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = 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, - ): - 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_videos_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_videos_per_prompt, 1) - prompt_embeds = prompt_embeds.view(bs_embed * num_videos_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_videos_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) - - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_videos_per_prompt).view( - bs_embed * num_videos_per_prompt, -1 - ) - if do_classifier_free_guidance: - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_videos_per_prompt).view( - bs_embed * num_videos_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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.text_to_video_synthesis/pipeline_text_to_video_synth.TextToVideoSDPipeline.decode_latents - def decode_latents(self, latents): - latents = 1 / self.vae.config.scaling_factor * latents - - batch_size, channels, num_frames, height, width = latents.shape - latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) - - image = self.vae.decode(latents).sample - video = image[None, :].reshape((batch_size, num_frames, -1) + image.shape[2:]).permute(0, 2, 1, 3, 4) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - video = video.float() - return video - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - 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, - 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_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_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`." - ) - - # Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_synth.TextToVideoSDPipeline.prepare_latents - def prepare_latents( - self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None - ): - shape = ( - batch_size, - num_channels_latents, - num_frames, - height // self.vae_scale_factor, - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - 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 - - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - FusedAttnProcessor2_0, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_end(self): - return self._denoising_end - - @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, - prompt_2: Optional[Union[str, List[str]]] = None, - num_frames: int = 16, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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"], - ): - r""" - Function invoked when calling the pipeline for generation. - - Args: - 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*): - 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: - The number of video frames that are generated. Defaults to 16 frames which at 8 frames per seconds - amounts to 2 seconds of video. - height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The height in pixels of the generated video. 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 video. 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 video 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. - 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. - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output) - 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 video quality. - 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*): - 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): - The number of videos 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` 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 ge 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. - 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. If not provided, embeddings are computed from the - `ip_adapter_image` input argument. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated video. 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.AnimateDiffPipelineOutput`] instead of a - plain tuple. - cross_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). - 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. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.animatediff.pipeline_output.AnimateDiffPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.animatediff.pipeline_output.AnimateDiffPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is a list with the generated frames. - """ - - # 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 - - num_videos_per_prompt = 1 - - 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, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - 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. Encode input prompt - lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else 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_videos_per_prompt=num_videos_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 - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 5. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - num_channels_latents, - num_frames, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 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.repeat_interleave(repeats=num_frames, 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_videos_per_prompt, 1) - - 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_videos_per_prompt, - self.do_classifier_free_guidance, - ) - - # 7.1 Apply denoising_end - if ( - self.denoising_end is not None - and isinstance(self.denoising_end, float) - and self.denoising_end > 0 - and self.denoising_end < 1 - ): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - # 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_videos_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) - - num_free_init_iters = self._free_init_num_iters if self.free_init_enabled else 1 - for free_init_iter in range(num_free_init_iters): - if self.free_init_enabled: - latents, timesteps = self._apply_free_init( - latents, free_init_iter, num_inference_steps, device, latents.dtype, generator - ) - - self._num_timesteps = len(timesteps) - - # 9. Denoising loop - with self.progress_bar(total=self._num_timesteps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - if ip_adapter_image is not None or ip_adapter_image_embeds: - added_cond_kwargs["image_embeds"] = image_embeds - - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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.do_classifier_free_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg( - noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale - ) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds - ) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids) - - progress_bar.update() - - if XLA_AVAILABLE: - xm.mark_step() - - # 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) - - # 10. Post processing - if output_type == "latent": - video = latents - else: - video_tensor = self.decode_latents(latents) - video = self.video_processor.postprocess_video(video=video_tensor, output_type=output_type) - - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - - # 11. Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return AnimateDiffPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py deleted file mode 100644 index 46d650efe8b6ee01eaa005b073c21b25fe2b6105..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py +++ /dev/null @@ -1,1023 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL -import torch -import torch.nn.functional as F -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel, UNetMotionModel -from ...models.controlnets.controlnet_sparsectrl import SparseControlNetModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...models.unets.unet_motion_model import MotionAdapter -from ...schedulers import 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 is_compiled_module, randn_tensor -from ...video_processor import VideoProcessor -from ..free_init_utils import FreeInitMixin -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from .pipeline_output import AnimateDiffPipelineOutput - - -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 AnimateDiffSparseControlNetPipeline - >>> from diffusers.models import AutoencoderKL, MotionAdapter, SparseControlNetModel - >>> from diffusers.schedulers import DPMSolverMultistepScheduler - >>> from diffusers.utils import export_to_gif, load_image - - >>> model_id = "SG161222/Realistic_Vision_V5.1_noVAE" - >>> motion_adapter_id = "guoyww/animatediff-motion-adapter-v1-5-3" - >>> controlnet_id = "guoyww/animatediff-sparsectrl-scribble" - >>> lora_adapter_id = "guoyww/animatediff-motion-lora-v1-5-3" - >>> vae_id = "stabilityai/sd-vae-ft-mse" - >>> device = "cuda" - - >>> motion_adapter = MotionAdapter.from_pretrained(motion_adapter_id, torch_dtype=torch.float16).to(device) - >>> controlnet = SparseControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch.float16).to(device) - >>> vae = AutoencoderKL.from_pretrained(vae_id, torch_dtype=torch.float16).to(device) - >>> scheduler = DPMSolverMultistepScheduler.from_pretrained( - ... model_id, - ... subfolder="scheduler", - ... beta_schedule="linear", - ... algorithm_type="dpmsolver++", - ... use_karras_sigmas=True, - ... ) - >>> pipe = AnimateDiffSparseControlNetPipeline.from_pretrained( - ... model_id, - ... motion_adapter=motion_adapter, - ... controlnet=controlnet, - ... vae=vae, - ... scheduler=scheduler, - ... torch_dtype=torch.float16, - ... ).to(device) - >>> pipe.load_lora_weights(lora_adapter_id, adapter_name="motion_lora") - >>> pipe.fuse_lora(lora_scale=1.0) - - >>> prompt = "an aerial view of a cyberpunk city, night time, neon lights, masterpiece, high quality" - >>> negative_prompt = "low quality, worst quality, letterboxed" - - >>> image_files = [ - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/animatediff-scribble-1.png", - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/animatediff-scribble-2.png", - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/animatediff-scribble-3.png", - ... ] - >>> condition_frame_indices = [0, 8, 15] - >>> conditioning_frames = [load_image(img_file) for img_file in image_files] - - >>> video = pipe( - ... prompt=prompt, - ... negative_prompt=negative_prompt, - ... num_inference_steps=25, - ... conditioning_frames=conditioning_frames, - ... controlnet_conditioning_scale=1.0, - ... controlnet_frame_indices=condition_frame_indices, - ... generator=torch.Generator().manual_seed(1337), - ... ).frames[0] - >>> export_to_gif(video, "output.gif") - ``` -""" - - -# 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" -): - 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 AnimateDiffSparseControlNetPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - FreeInitMixin, - FromSingleFileMixin, -): - r""" - Pipeline for controlled text-to-video generation using the method described in [SparseCtrl: Adding Sparse Controls - to Text-to-Video Diffusion Models](https://huggingface.co/papers/2311.16933). - - 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.). - - 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 - - [`~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 ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer (`CLIPTokenizer`): - A [`~transformers.CLIPTokenizer`] to tokenize text. - unet ([`UNet2DConditionModel`]): - 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. - 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`]. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["feature_extractor", "image_encoder", "motion_adapter"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], - motion_adapter: MotionAdapter, - controlnet: SparseControlNetModel, - scheduler: KarrasDiffusionSchedulers, - feature_extractor: CLIPImageProcessor = None, - image_encoder: CLIPVisionModelWithProjection = None, - ): - super().__init__() - if isinstance(unet, UNet2DConditionModel): - unet = UNetMotionModel.from_unet2d(unet, motion_adapter) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - motion_adapter=motion_adapter, - controlnet=controlnet, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - self.video_processor = VideoProcessor(do_resize=False, vae_scale_factor=self.vae_scale_factor) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt with num_images_per_prompt -> num_videos_per_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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.text_to_video_synthesis/pipeline_text_to_video_synth.TextToVideoSDPipeline.decode_latents - def decode_latents(self, latents): - latents = 1 / self.vae.config.scaling_factor * latents - - batch_size, channels, num_frames, height, width = latents.shape - latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) - - image = self.vae.decode(latents).sample - video = image[None, :].reshape((batch_size, num_frames, -1) + image.shape[2:]).permute(0, 2, 1, 3, 4) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - video = video.float() - return video - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - image=None, - controlnet_conditioning_scale: float = 1.0, - ): - 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_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 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" - ) - - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - - # check `image` - if ( - isinstance(self.controlnet, SparseControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, SparseControlNetModel) - ): - if isinstance(image, list): - for image_ in image: - self.check_image(image_, prompt, prompt_embeds) - else: - self.check_image(image, prompt, prompt_embeds) - else: - assert False - - # Check `controlnet_conditioning_scale` - if ( - isinstance(self.controlnet, SparseControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, SparseControlNetModel) - ): - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - else: - assert False - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.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}" - ) - - # Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_synth.TextToVideoSDPipeline.prepare_latents - def prepare_latents( - self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None - ): - shape = ( - batch_size, - num_channels_latents, - num_frames, - height // self.vae_scale_factor, - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def prepare_image(self, image, width, height, device, dtype): - image = self.control_image_processor.preprocess(image, height=height, width=width) - controlnet_images = image.unsqueeze(0).to(device, dtype) - batch_size, num_frames, channels, height, width = controlnet_images.shape - - # TODO: remove below line - assert controlnet_images.min() >= 0 and controlnet_images.max() <= 1 - - if self.controlnet.use_simplified_condition_embedding: - controlnet_images = controlnet_images.reshape(batch_size * num_frames, channels, height, width) - controlnet_images = 2 * controlnet_images - 1 - conditioning_frames = retrieve_latents(self.vae.encode(controlnet_images)) * self.vae.config.scaling_factor - conditioning_frames = conditioning_frames.reshape( - batch_size, num_frames, 4, height // self.vae_scale_factor, width // self.vae_scale_factor - ) - else: - conditioning_frames = controlnet_images - - conditioning_frames = conditioning_frames.permute(0, 2, 1, 3, 4) # [b, c, f, h, w] - return conditioning_frames - - def prepare_sparse_control_conditioning( - self, - conditioning_frames: torch.Tensor, - num_frames: int, - controlnet_frame_indices: int, - device: torch.device, - dtype: torch.dtype, - ) -> Tuple[torch.Tensor, torch.Tensor]: - assert conditioning_frames.shape[2] >= len(controlnet_frame_indices) - - batch_size, channels, _, height, width = conditioning_frames.shape - controlnet_cond = torch.zeros((batch_size, channels, num_frames, height, width), dtype=dtype, device=device) - controlnet_cond_mask = torch.zeros((batch_size, 1, num_frames, height, width), dtype=dtype, device=device) - controlnet_cond[:, :, controlnet_frame_indices] = conditioning_frames[:, :, : len(controlnet_frame_indices)] - controlnet_cond_mask[:, :, controlnet_frame_indices] = 1 - - return controlnet_cond, controlnet_cond_mask - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_frames: int = 16, - num_inference_steps: int = 50, - guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = 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, - 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], - 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"], - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The width in pixels of the generated video. - num_frames (`int`, *optional*, defaults to 16): - The number of video frames that are generated. Defaults to 16 frames which at 8 frames per seconds - amounts to 2 seconds of video. - 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. - 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*): - 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*): - 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 video - 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`. Latents should be of shape - `(batch_size, num_channel, num_frames, height, width)`. - 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. - 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. - 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)`. 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*): - 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`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.text_to_video_synthesis.TextToVideoSDPipelineOutput`] instead - of a plain tuple. - 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): - 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]`): - 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 - structure. Must have the same length as `conditioning_frames`. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.animatediff.pipeline_output.AnimateDiffPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.animatediff.pipeline_output.AnimateDiffPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is a list with the generated frames. - """ - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # 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 - num_videos_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt=prompt, - height=height, - width=width, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ip_adapter_image=ip_adapter_image, - ip_adapter_image_embeds=ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - image=conditioning_frames, - controlnet_conditioning_scale=controlnet_conditioning_scale, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - - # 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 - - global_pool_conditions = ( - controlnet.config.global_pool_conditions - if isinstance(controlnet, SparseControlNetModel) - else controlnet.nets[0].config.global_pool_conditions - ) - guess_mode = guess_mode or global_pool_conditions - - # 3. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_videos_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - prompt_embeds = prompt_embeds.repeat_interleave(repeats=num_frames, dim=0) - - # 4. Prepare IP-Adapter embeddings - 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_videos_per_prompt, - self.do_classifier_free_guidance, - ) - - # 5. Prepare controlnet conditioning - conditioning_frames = self.prepare_image(conditioning_frames, width, height, device, controlnet.dtype) - controlnet_cond, controlnet_cond_mask = self.prepare_sparse_control_conditioning( - conditioning_frames, num_frames, controlnet_frame_indices, device, controlnet.dtype - ) - - # 6. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 7. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - num_channels_latents, - num_frames, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 9. Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - num_free_init_iters = self._free_init_num_iters if self.free_init_enabled else 1 - for free_init_iter in range(num_free_init_iters): - if self.free_init_enabled: - latents, timesteps = self._apply_free_init( - latents, free_init_iter, num_inference_steps, device, latents.dtype, generator - ) - - self._num_timesteps = len(timesteps) - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - # 10. Denoising loop - with self.progress_bar(total=self._num_timesteps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - if guess_mode and self.do_classifier_free_guidance: - # Infer SparseControlNetModel only for the conditional batch. - control_model_input = latents - control_model_input = self.scheduler.scale_model_input(control_model_input, t) - controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] - else: - control_model_input = latent_model_input - controlnet_prompt_embeds = prompt_embeds - - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=controlnet_cond, - conditioning_mask=controlnet_cond_mask, - conditioning_scale=controlnet_conditioning_scale, - guess_mode=guess_mode, - return_dict=False, - ) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - ).sample - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - 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() - - # 11. Post processing - if output_type == "latent": - video = latents - else: - video_tensor = self.decode_latents(latents) - video = self.video_processor.postprocess_video(video=video_tensor, output_type=output_type) - - # 12. Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return AnimateDiffPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py deleted file mode 100644 index 6f3a609aba4a6ac7ca13993f62ad4a40e101da1c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py +++ /dev/null @@ -1,1065 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...image_processor import PipelineImageInput -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel, UNetMotionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...models.unets.unet_motion_model import MotionAdapter -from ...schedulers import ( - DDIMScheduler, - DPMSolverMultistepScheduler, - EulerAncestralDiscreteScheduler, - EulerDiscreteScheduler, - LMSDiscreteScheduler, - PNDMScheduler, -) -from ...utils import USE_PEFT_BACKEND, is_torch_xla_available, logging, scale_lora_layers, unscale_lora_layers -from ...utils.torch_utils import randn_tensor -from ...video_processor import VideoProcessor -from ..free_init_utils import FreeInitMixin -from ..free_noise_utils import AnimateDiffFreeNoiseMixin -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from .pipeline_output import AnimateDiffPipelineOutput - - -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 imageio - >>> import requests - >>> import torch - >>> from diffusers import AnimateDiffVideoToVideoPipeline, DDIMScheduler, MotionAdapter - >>> from diffusers.utils import export_to_gif - >>> from io import BytesIO - >>> from PIL import Image - - >>> adapter = MotionAdapter.from_pretrained( - ... "guoyww/animatediff-motion-adapter-v1-5-2", torch_dtype=torch.float16 - ... ) - >>> pipe = AnimateDiffVideoToVideoPipeline.from_pretrained( - ... "SG161222/Realistic_Vision_V5.1_noVAE", motion_adapter=adapter - ... ).to("cuda") - >>> pipe.scheduler = DDIMScheduler( - ... beta_schedule="linear", steps_offset=1, clip_sample=False, timespace_spacing="linspace" - ... ) - - - >>> def load_video(file_path: str): - ... images = [] - - ... if file_path.startswith(("http://", "https://")): - ... # If the file_path is a URL - ... response = requests.get(file_path) - ... response.raise_for_status() - ... content = BytesIO(response.content) - ... vid = imageio.get_reader(content) - ... else: - ... # Assuming it's a local file path - ... vid = imageio.get_reader(file_path) - - ... for frame in vid: - ... pil_image = Image.fromarray(frame) - ... images.append(pil_image) - - ... return images - - - >>> video = load_video( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/animatediff-vid2vid-input-1.gif" - ... ) - >>> output = pipe( - ... video=video, prompt="panda playing a guitar, on a boat, in the ocean, high quality", strength=0.5 - ... ) - >>> frames = output.frames[0] - >>> export_to_gif(frames, "animation.gif") - ``` -""" - - -# 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" -): - 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 AnimateDiffVideoToVideoPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - FreeInitMixin, - AnimateDiffFreeNoiseMixin, - FromSingleFileMixin, -): - r""" - Pipeline for video-to-video generation. - - 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.). - - 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 - - [`~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 ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer (`CLIPTokenizer`): - A [`~transformers.CLIPTokenizer`] to tokenize text. - unet ([`UNet2DConditionModel`]): - 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. - 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`]. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["feature_extractor", "image_encoder", "motion_adapter"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], - motion_adapter: MotionAdapter, - scheduler: Union[ - DDIMScheduler, - PNDMScheduler, - LMSDiscreteScheduler, - EulerDiscreteScheduler, - EulerAncestralDiscreteScheduler, - DPMSolverMultistepScheduler, - ], - feature_extractor: CLIPImageProcessor = None, - image_encoder: CLIPVisionModelWithProjection = None, - ): - super().__init__() - if isinstance(unet, UNet2DConditionModel): - unet = UNetMotionModel.from_unet2d(unet, motion_adapter) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - motion_adapter=motion_adapter, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor) - - 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, dict)): - 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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 encode_video(self, video, generator, decode_chunk_size: int = 16) -> torch.Tensor: - latents = [] - for i in range(0, len(video), decode_chunk_size): - batch_video = video[i : i + decode_chunk_size] - batch_video = retrieve_latents(self.vae.encode(batch_video), generator=generator) - latents.append(batch_video) - return torch.cat(latents) - - # Copied from diffusers.pipelines.animatediff.pipeline_animatediff.AnimateDiffPipeline.decode_latents - def decode_latents(self, latents, decode_chunk_size: int = 16): - latents = 1 / self.vae.config.scaling_factor * latents - - batch_size, channels, num_frames, height, width = latents.shape - latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) - - video = [] - for i in range(0, latents.shape[0], decode_chunk_size): - batch_latents = latents[i : i + decode_chunk_size] - batch_latents = self.vae.decode(batch_latents).sample - video.append(batch_latents) - - video = torch.cat(video) - video = video[None, :].reshape((batch_size, num_frames, -1) + video.shape[2:]).permute(0, 2, 1, 3, 4) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - video = video.float() - return video - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - strength, - height, - width, - video=None, - latents=None, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=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 % 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_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, list, dict)): - raise ValueError(f"`prompt` has to be of type `str`, `list` 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 video is not None and latents is not None: - raise ValueError("Only one of `video` or `latents` should be provided") - - 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" - ) - - def get_timesteps(self, num_inference_steps, timesteps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - timesteps = timesteps[t_start * self.scheduler.order :] - - return timesteps, num_inference_steps - t_start - - def prepare_latents( - self, - video: Optional[torch.Tensor] = 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, - decode_chunk_size: int = 16, - add_noise: bool = False, - ) -> torch.Tensor: - num_frames = video.shape[1] if latents is None else latents.shape[2] - shape = ( - batch_size, - num_channels_latents, - num_frames, - height // self.vae_scale_factor, - 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: - # make sure the VAE is in float32 mode, as it overflows in float16 - if self.vae.config.force_upcast: - video = video.float() - self.vae.to(dtype=torch.float32) - - if isinstance(generator, list): - init_latents = [ - self.encode_video(video[i], generator[i], decode_chunk_size).unsqueeze(0) - for i in range(batch_size) - ] - else: - init_latents = [self.encode_video(vid, generator, decode_chunk_size).unsqueeze(0) for vid in video] - - init_latents = torch.cat(init_latents, dim=0) - - # restore vae to original dtype - if self.vae.config.force_upcast: - self.vae.to(dtype) - - init_latents = init_latents.to(dtype) - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - error_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial" - " images (`image`). Please make sure to update your script to pass as many initial images as text prompts" - ) - raise ValueError(error_message) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=dtype) - latents = self.scheduler.add_noise(init_latents, noise, timestep).permute(0, 2, 1, 3, 4) - else: - if shape != latents.shape: - # [B, C, F, H, W] - raise ValueError(f"`latents` expected to have {shape=}, but found {latents.shape=}") - - latents = latents.to(device, dtype=dtype) - - if add_noise: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = self.scheduler.add_noise(latents, noise, timestep) - - return latents - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @property - def interrupt(self): - return self._interrupt - - @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, - num_inference_steps: int = 50, - enforce_inference_steps: bool = False, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = 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, - 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", - 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"], - decode_chunk_size: int = 16, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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*): - 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. - width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The width in pixels of the generated video. - 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*): - 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*): - 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. - strength (`float`, *optional*, defaults to 0.8): - Higher strength leads to more differences between original video and generated video. - 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*): - 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*): - 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 video - 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`. Latents should be of shape - `(batch_size, num_channel, num_frames, height, width)`. - 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. - 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. - 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)`. 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. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated video. Choose between `torch.Tensor`, `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`AnimateDiffPipelineOutput`] instead of a plain tuple. - 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). - 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. - 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. - decode_chunk_size (`int`, defaults to `16`): - The number of frames to decode at a time when calling `decode_latents` method. - - Examples: - - Returns: - [`pipelines.animatediff.pipeline_output.AnimateDiffPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`pipelines.animatediff.pipeline_output.AnimateDiffPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is a list with the generated frames. - """ - - # 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 - - num_videos_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt=prompt, - strength=strength, - height=height, - width=width, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - video=video, - latents=latents, - ip_adapter_image=ip_adapter_image, - ip_adapter_image_embeds=ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._interrupt = False - - # 2. Define call parameters - if prompt is not None and isinstance(prompt, (str, dict)): - 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 - dtype = self.dtype - - # 3. Prepare timesteps - if not enforce_inference_steps: - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, 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 - ) - timesteps = timesteps[-num_inference_steps:] - latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt) - - # 4. Prepare latent variables - if latents is None: - video = self.video_processor.preprocess_video(video, height=height, width=width) - # Move the number of frames before the number of channels. - video = video.permute(0, 2, 1, 3, 4) - video = video.to(device=device, dtype=dtype) - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - video=video, - height=height, - width=width, - num_channels_latents=num_channels_latents, - batch_size=batch_size * num_videos_per_prompt, - timestep=latent_timestep, - dtype=dtype, - device=device, - generator=generator, - latents=latents, - decode_chunk_size=decode_chunk_size, - add_noise=enforce_inference_steps, - ) - - # 5. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - num_frames = latents.shape[2] - if self.free_noise_enabled: - prompt_embeds, negative_prompt_embeds = self._encode_prompt_free_noise( - prompt=prompt, - num_frames=num_frames, - device=device, - num_videos_per_prompt=num_videos_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - else: - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_videos_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - prompt_embeds = prompt_embeds.repeat_interleave(repeats=num_frames, dim=0) - - # 6. Prepare IP-Adapter embeddings - 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_videos_per_prompt, - self.do_classifier_free_guidance, - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 8. Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - num_free_init_iters = self._free_init_num_iters if self.free_init_enabled else 1 - for free_init_iter in range(num_free_init_iters): - if self.free_init_enabled: - latents, timesteps = self._apply_free_init( - latents, free_init_iter, num_inference_steps, device, latents.dtype, generator - ) - num_inference_steps = len(timesteps) - # make sure to readjust timesteps based on strength - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, timesteps, strength, device) - - self._num_timesteps = len(timesteps) - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - # 9. Denoising loop - with self.progress_bar(total=self._num_timesteps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - ).sample - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - 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() - - # 10. Post-processing - if output_type == "latent": - video = latents - else: - video_tensor = self.decode_latents(latents, decode_chunk_size) - video = self.video_processor.postprocess_video(video=video_tensor, output_type=output_type) - - # 11. Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return AnimateDiffPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py deleted file mode 100644 index b00f344598ad26431c3d49a75a3f1627b240c618..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py +++ /dev/null @@ -1,1353 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -import torch.nn.functional as F -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...image_processor import PipelineImageInput -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import ( - AutoencoderKL, - ControlNetModel, - ImageProjection, - MultiControlNetModel, - UNet2DConditionModel, - UNetMotionModel, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...models.unets.unet_motion_model import MotionAdapter -from ...schedulers import ( - DDIMScheduler, - DPMSolverMultistepScheduler, - EulerAncestralDiscreteScheduler, - EulerDiscreteScheduler, - LMSDiscreteScheduler, - PNDMScheduler, -) -from ...utils import USE_PEFT_BACKEND, is_torch_xla_available, logging, scale_lora_layers, unscale_lora_layers -from ...utils.torch_utils import is_compiled_module, randn_tensor -from ...video_processor import VideoProcessor -from ..free_init_utils import FreeInitMixin -from ..free_noise_utils import AnimateDiffFreeNoiseMixin -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from .pipeline_output import AnimateDiffPipelineOutput - - -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 PIL import Image - >>> from tqdm.auto import tqdm - - >>> from diffusers import AnimateDiffVideoToVideoControlNetPipeline - >>> from diffusers.utils import export_to_gif, load_video - >>> from diffusers import AutoencoderKL, ControlNetModel, MotionAdapter, LCMScheduler - - >>> controlnet = ControlNetModel.from_pretrained( - ... "lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16 - ... ) - >>> motion_adapter = MotionAdapter.from_pretrained("wangfuyun/AnimateLCM") - >>> vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16) - - >>> pipe = AnimateDiffVideoToVideoControlNetPipeline.from_pretrained( - ... "SG161222/Realistic_Vision_V5.1_noVAE", - ... motion_adapter=motion_adapter, - ... controlnet=controlnet, - ... vae=vae, - ... ).to(device="cuda", dtype=torch.float16) - - >>> pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, beta_schedule="linear") - >>> pipe.load_lora_weights( - ... "wangfuyun/AnimateLCM", weight_name="AnimateLCM_sd15_t2v_lora.safetensors", adapter_name="lcm-lora" - ... ) - >>> pipe.set_adapters(["lcm-lora"], [0.8]) - - >>> video = load_video( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/dance.gif" - ... ) - >>> video = [frame.convert("RGB") for frame in video] - - >>> from controlnet_aux.processor import OpenposeDetector - - >>> open_pose = OpenposeDetector.from_pretrained("lllyasviel/Annotators").to("cuda") - >>> for frame in tqdm(video): - ... conditioning_frames.append(open_pose(frame)) - - >>> prompt = "astronaut in space, dancing" - >>> negative_prompt = "bad quality, worst quality, jpeg artifacts, ugly" - - >>> strength = 0.8 - >>> with torch.inference_mode(): - ... video = pipe( - ... video=video, - ... prompt=prompt, - ... negative_prompt=negative_prompt, - ... num_inference_steps=10, - ... guidance_scale=2.0, - ... controlnet_conditioning_scale=0.75, - ... conditioning_frames=conditioning_frames, - ... strength=strength, - ... generator=torch.Generator().manual_seed(42), - ... ).frames[0] - - >>> video = [frame.resize(conditioning_frames[0].size) for frame in video] - >>> export_to_gif(video, f"animatediff_vid2vid_controlnet.gif", fps=8) - ``` -""" - - -# 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" -): - 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 AnimateDiffVideoToVideoControlNetPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - FreeInitMixin, - AnimateDiffFreeNoiseMixin, - FromSingleFileMixin, -): - r""" - Pipeline for video-to-video generation with ControlNet guidance. - - 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.). - - 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 - - [`~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 ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer (`CLIPTokenizer`): - A [`~transformers.CLIPTokenizer`] to tokenize text. - unet ([`UNet2DConditionModel`]): - 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`): - 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. - 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`]. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["feature_extractor", "image_encoder", "motion_adapter"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], - motion_adapter: MotionAdapter, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], - scheduler: Union[ - DDIMScheduler, - PNDMScheduler, - LMSDiscreteScheduler, - EulerDiscreteScheduler, - EulerAncestralDiscreteScheduler, - DPMSolverMultistepScheduler, - ], - feature_extractor: CLIPImageProcessor = None, - image_encoder: CLIPVisionModelWithProjection = None, - ): - super().__init__() - if isinstance(unet, UNet2DConditionModel): - unet = UNetMotionModel.from_unet2d(unet, motion_adapter) - - if isinstance(controlnet, (list, tuple)): - controlnet = MultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - motion_adapter=motion_adapter, - controlnet=controlnet, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor) - self.control_video_processor = VideoProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - - # Copied from diffusers.pipelines.animatediff.pipeline_animatediff_video2video.AnimateDiffVideoToVideoPipeline.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, dict)): - 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.animatediff.pipeline_animatediff_video2video.AnimateDiffVideoToVideoPipeline.encode_video - def encode_video(self, video, generator, decode_chunk_size: int = 16) -> torch.Tensor: - latents = [] - for i in range(0, len(video), decode_chunk_size): - batch_video = video[i : i + decode_chunk_size] - batch_video = retrieve_latents(self.vae.encode(batch_video), generator=generator) - latents.append(batch_video) - return torch.cat(latents) - - # Copied from diffusers.pipelines.animatediff.pipeline_animatediff.AnimateDiffPipeline.decode_latents - def decode_latents(self, latents, decode_chunk_size: int = 16): - latents = 1 / self.vae.config.scaling_factor * latents - - batch_size, channels, num_frames, height, width = latents.shape - latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) - - video = [] - for i in range(0, latents.shape[0], decode_chunk_size): - batch_latents = latents[i : i + decode_chunk_size] - batch_latents = self.vae.decode(batch_latents).sample - video.append(batch_latents) - - video = torch.cat(video) - video = video[None, :].reshape((batch_size, num_frames, -1) + video.shape[2:]).permute(0, 2, 1, 3, 4) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - video = video.float() - return video - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - strength, - height, - width, - video=None, - conditioning_frames=None, - latents=None, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - controlnet_conditioning_scale=1.0, - control_guidance_start=0.0, - control_guidance_end=1.0, - ): - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - - 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_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, list, dict)): - raise ValueError(f"`prompt` has to be of type `str`, `list` 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 video is not None and latents is not None: - raise ValueError("Only one of `video` or `latents` should be provided") - - 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" - ) - - if isinstance(self.controlnet, MultiControlNetModel): - if isinstance(prompt, list): - 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." - ) - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - - num_frames = len(video) if latents is None else latents.shape[2] - - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - if not isinstance(conditioning_frames, list): - raise TypeError( - f"For single controlnet, `image` must be of type `list` but got {type(conditioning_frames)}" - ) - if len(conditioning_frames) != num_frames: - raise ValueError(f"Excepted image to have length {num_frames} but got {len(conditioning_frames)=}") - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if not isinstance(conditioning_frames, list) or not isinstance(conditioning_frames[0], list): - raise TypeError( - f"For multiple controlnets: `image` must be type list of lists but got {type(conditioning_frames)=}" - ) - if len(conditioning_frames[0]) != num_frames: - raise ValueError( - f"Expected length of image sublist as {num_frames} but got {len(conditioning_frames)=}" - ) - if any(len(img) != len(conditioning_frames[0]) for img in conditioning_frames): - raise ValueError("All conditioning frame batches for multicontrolnet must be same size") - else: - assert False - - # Check `controlnet_conditioning_scale` - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if isinstance(controlnet_conditioning_scale, list): - if any(isinstance(i, list) for i in controlnet_conditioning_scale): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - elif 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" - ) - else: - assert False - - if not isinstance(control_guidance_start, (tuple, list)): - control_guidance_start = [control_guidance_start] - - if not isinstance(control_guidance_end, (tuple, list)): - control_guidance_end = [control_guidance_end] - - 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(self.controlnet, MultiControlNetModel): - 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.") - - # Copied from diffusers.pipelines.animatediff.pipeline_animatediff_video2video.AnimateDiffVideoToVideoPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, timesteps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - timesteps = timesteps[t_start * self.scheduler.order :] - - return timesteps, num_inference_steps - t_start - - # Copied from diffusers.pipelines.animatediff.pipeline_animatediff_video2video.AnimateDiffVideoToVideoPipeline.prepare_latents - def prepare_latents( - self, - video: Optional[torch.Tensor] = 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, - decode_chunk_size: int = 16, - add_noise: bool = False, - ) -> torch.Tensor: - num_frames = video.shape[1] if latents is None else latents.shape[2] - shape = ( - batch_size, - num_channels_latents, - num_frames, - height // self.vae_scale_factor, - 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: - # make sure the VAE is in float32 mode, as it overflows in float16 - if self.vae.config.force_upcast: - video = video.float() - self.vae.to(dtype=torch.float32) - - if isinstance(generator, list): - init_latents = [ - self.encode_video(video[i], generator[i], decode_chunk_size).unsqueeze(0) - for i in range(batch_size) - ] - else: - init_latents = [self.encode_video(vid, generator, decode_chunk_size).unsqueeze(0) for vid in video] - - init_latents = torch.cat(init_latents, dim=0) - - # restore vae to original dtype - if self.vae.config.force_upcast: - self.vae.to(dtype) - - init_latents = init_latents.to(dtype) - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - error_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial" - " images (`image`). Please make sure to update your script to pass as many initial images as text prompts" - ) - raise ValueError(error_message) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=dtype) - latents = self.scheduler.add_noise(init_latents, noise, timestep).permute(0, 2, 1, 3, 4) - else: - if shape != latents.shape: - # [B, C, F, H, W] - raise ValueError(f"`latents` expected to have {shape=}, but found {latents.shape=}") - - latents = latents.to(device, dtype=dtype) - - if add_noise: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = self.scheduler.add_noise(latents, noise, timestep) - - return latents - - # Copied from diffusers.pipelines.animatediff.pipeline_animatediff_controlnet.AnimateDiffControlNetPipeline.prepare_video - def prepare_conditioning_frames( - self, - video, - width, - height, - batch_size, - num_videos_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - video = self.control_video_processor.preprocess_video(video, height=height, width=width).to( - dtype=torch.float32 - ) - video = video.permute(0, 2, 1, 3, 4).flatten(0, 1) - video_batch_size = video.shape[0] - - if video_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_videos_per_prompt - - video = video.repeat_interleave(repeat_by, dim=0) - video = video.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - video = torch.cat([video] * 2) - - return video - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @property - def interrupt(self): - return self._interrupt - - @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, - num_inference_steps: int = 50, - enforce_inference_steps: bool = False, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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"], - decode_chunk_size: int = 16, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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*): - 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. - width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The width in pixels of the generated video. - 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*): - 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*): - 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. - strength (`float`, *optional*, defaults to 0.8): - Higher strength leads to more differences between original video and generated video. - 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*): - 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*): - 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 video - 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`. Latents should be of shape - `(batch_size, num_channel, num_frames, height, width)`. - 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. - 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. - 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)`. 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*): - 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. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated video. Choose between `torch.Tensor`, `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`AnimateDiffPipelineOutput`] instead of a plain tuple. - 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): - 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): - The percentage of total steps at which the ControlNet starts applying. - 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 - the output of the pre-final layer will be used for computing the prompt embeddings. - 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. - decode_chunk_size (`int`, defaults to `16`): - The number of frames to decode at a time when calling `decode_latents` method. - - Examples: - - Returns: - [`pipelines.animatediff.pipeline_output.AnimateDiffPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`pipelines.animatediff.pipeline_output.AnimateDiffPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is a list with the generated frames. - """ - - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 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 - - num_videos_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt=prompt, - strength=strength, - height=height, - width=width, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - video=video, - conditioning_frames=conditioning_frames, - latents=latents, - ip_adapter_image=ip_adapter_image, - ip_adapter_image_embeds=ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - controlnet_conditioning_scale=controlnet_conditioning_scale, - control_guidance_start=control_guidance_start, - control_guidance_end=control_guidance_end, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._interrupt = False - - # 2. Define call parameters - if prompt is not None and isinstance(prompt, (str, dict)): - 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 - dtype = self.dtype - - # 3. Prepare timesteps - if not enforce_inference_steps: - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, 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 - ) - timesteps = timesteps[-num_inference_steps:] - latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt) - - # 4. Prepare latent variables - if latents is None: - video = self.video_processor.preprocess_video(video, height=height, width=width) - # Move the number of frames before the number of channels. - video = video.permute(0, 2, 1, 3, 4) - video = video.to(device=device, dtype=dtype) - - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - video=video, - height=height, - width=width, - num_channels_latents=num_channels_latents, - batch_size=batch_size * num_videos_per_prompt, - timestep=latent_timestep, - dtype=dtype, - device=device, - generator=generator, - latents=latents, - decode_chunk_size=decode_chunk_size, - add_noise=enforce_inference_steps, - ) - - # 5. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - num_frames = latents.shape[2] - if self.free_noise_enabled: - prompt_embeds, negative_prompt_embeds = self._encode_prompt_free_noise( - prompt=prompt, - num_frames=num_frames, - device=device, - num_videos_per_prompt=num_videos_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - else: - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_videos_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - prompt_embeds = prompt_embeds.repeat_interleave(repeats=num_frames, dim=0) - - # 6. Prepare IP-Adapter embeddings - 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_videos_per_prompt, - self.do_classifier_free_guidance, - ) - - # 7. Prepare ControlNet conditions - if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): - controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) - - global_pool_conditions = ( - controlnet.config.global_pool_conditions - if isinstance(controlnet, ControlNetModel) - else controlnet.nets[0].config.global_pool_conditions - ) - guess_mode = guess_mode or global_pool_conditions - - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) - - if isinstance(controlnet, ControlNetModel): - conditioning_frames = self.prepare_conditioning_frames( - video=conditioning_frames, - width=width, - height=height, - batch_size=batch_size * num_videos_per_prompt * num_frames, - num_videos_per_prompt=num_videos_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - elif isinstance(controlnet, MultiControlNetModel): - cond_prepared_videos = [] - for frame_ in conditioning_frames: - prepared_video = self.prepare_conditioning_frames( - video=frame_, - width=width, - height=height, - batch_size=batch_size * num_videos_per_prompt * num_frames, - num_videos_per_prompt=num_videos_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - cond_prepared_videos.append(prepared_video) - conditioning_frames = cond_prepared_videos - else: - assert False - - # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 9. Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - num_free_init_iters = self._free_init_num_iters if self.free_init_enabled else 1 - for free_init_iter in range(num_free_init_iters): - if self.free_init_enabled: - latents, timesteps = self._apply_free_init( - latents, free_init_iter, num_inference_steps, device, latents.dtype, generator - ) - num_inference_steps = len(timesteps) - # make sure to readjust timesteps based on strength - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, timesteps, strength, device) - - self._num_timesteps = len(timesteps) - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - # 10. Denoising loop - with self.progress_bar(total=self._num_timesteps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - if guess_mode and self.do_classifier_free_guidance: - # Infer ControlNet only for the conditional batch. - control_model_input = latents - control_model_input = self.scheduler.scale_model_input(control_model_input, t) - controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] - else: - control_model_input = latent_model_input - controlnet_prompt_embeds = prompt_embeds - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - control_model_input = torch.transpose(control_model_input, 1, 2) - control_model_input = control_model_input.reshape( - (-1, control_model_input.shape[2], control_model_input.shape[3], control_model_input.shape[4]) - ) - - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=conditioning_frames, - conditioning_scale=cond_scale, - guess_mode=guess_mode, - return_dict=False, - ) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - ).sample - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - 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() - - # 11. Post-processing - if output_type == "latent": - video = latents - else: - video_tensor = self.decode_latents(latents, decode_chunk_size) - video = self.video_processor.postprocess_video(video=video_tensor, output_type=output_type) - - # 12. Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return AnimateDiffPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_output.py deleted file mode 100644 index 2417223cf95e682aa9b86f4c6013c5285c9afce4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/animatediff/pipeline_output.py +++ /dev/null @@ -1,24 +0,0 @@ -from dataclasses import dataclass -from typing import List, Union - -import numpy as np -import PIL.Image -import torch - -from ...utils import BaseOutput - - -@dataclass -class AnimateDiffPipelineOutput(BaseOutput): - r""" - 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 - 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]]] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm/__init__.py deleted file mode 100644 index a002b4aa72e0a180c7042c406667d37122d6e4cc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm/__init__.py +++ /dev/null @@ -1,51 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - is_torch_available, - is_transformers_available, - is_transformers_version, -) - - -_dummy_objects = {} -_import_structure = {} - -try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_objects import ( - AudioLDMPipeline, - ) - - _dummy_objects.update({"AudioLDMPipeline": AudioLDMPipeline}) -else: - _import_structure["pipeline_audioldm"] = ["AudioLDMPipeline"] - - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_objects import ( - AudioLDMPipeline, - ) - - else: - from .pipeline_audioldm import AudioLDMPipeline -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/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index e2f6ec5134418c43f70897362ded2fab65d5fc7b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm/__pycache__/pipeline_audioldm.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm/__pycache__/pipeline_audioldm.cpython-310.pyc deleted file mode 100644 index 0783774d0c52fa395d573b1ea28ce395fb552af2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm/__pycache__/pipeline_audioldm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm/pipeline_audioldm.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm/pipeline_audioldm.py deleted file mode 100644 index 6a70f00c76c785c1c0e94a80624e1113b3a31cb9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm/pipeline_audioldm.py +++ /dev/null @@ -1,558 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import torch -import torch.nn.functional as F -from transformers import ClapTextModelWithProjection, RobertaTokenizer, RobertaTokenizerFast, SpeechT5HifiGan - -from ...models import AutoencoderKL, UNet2DConditionModel -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import is_torch_xla_available, logging, replace_example_docstring -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import AudioPipelineOutput, DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin - - -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 - >>> from diffusers import AudioLDMPipeline - >>> import torch - >>> import scipy - - >>> repo_id = "cvssp/audioldm-s-full-v2" - >>> pipe = AudioLDMPipeline.from_pretrained(repo_id, torch_dtype=torch.float16) - >>> pipe = pipe.to("cuda") - - >>> prompt = "Techno music with a strong, upbeat tempo and high melodic riffs" - >>> audio = pipe(prompt, num_inference_steps=10, audio_length_in_s=5.0).audios[0] - - >>> # save the audio sample as a .wav file - >>> scipy.io.wavfile.write("techno.wav", rate=16000, data=audio) - ``` -""" - - -class AudioLDMPipeline(DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin): - r""" - Pipeline for text-to-audio generation using AudioLDM. - - 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: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.ClapTextModelWithProjection`]): - Frozen text-encoder (`ClapTextModelWithProjection`, specifically the - [laion/clap-htsat-unfused](https://huggingface.co/laion/clap-htsat-unfused) variant. - tokenizer ([`PreTrainedTokenizer`]): - A [`~transformers.RobertaTokenizer`] to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded audio latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded audio latents. Can be one of - [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. - vocoder ([`~transformers.SpeechT5HifiGan`]): - Vocoder of class `SpeechT5HifiGan`. - """ - - _last_supported_version = "0.33.1" - model_cpu_offload_seq = "text_encoder->unet->vae" - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: ClapTextModelWithProjection, - tokenizer: Union[RobertaTokenizer, RobertaTokenizerFast], - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - vocoder: SpeechT5HifiGan, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - vocoder=vocoder, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - - def _encode_prompt( - self, - prompt, - device, - num_waveforms_per_prompt, - do_classifier_free_guidance, - negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = 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_waveforms_per_prompt (`int`): - 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*): - 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`). - 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. - """ - 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: - 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 - attention_mask = text_inputs.attention_mask - 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 CLAP can only handle sequences up to" - f" {self.tokenizer.model_max_length} tokens: {removed_text}" - ) - - prompt_embeds = self.text_encoder( - text_input_ids.to(device), - attention_mask=attention_mask.to(device), - ) - prompt_embeds = prompt_embeds.text_embeds - # additional L_2 normalization over each hidden-state - prompt_embeds = F.normalize(prompt_embeds, dim=-1) - - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.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_waveforms_per_prompt) - prompt_embeds = prompt_embeds.view(bs_embed * num_waveforms_per_prompt, seq_len) - - # 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 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 - - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - - uncond_input_ids = uncond_input.input_ids.to(device) - attention_mask = uncond_input.attention_mask.to(device) - - negative_prompt_embeds = self.text_encoder( - uncond_input_ids, - attention_mask=attention_mask, - ) - negative_prompt_embeds = negative_prompt_embeds.text_embeds - # additional L_2 normalization over each hidden-state - negative_prompt_embeds = F.normalize(negative_prompt_embeds, dim=-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] - - negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_waveforms_per_prompt) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_waveforms_per_prompt, seq_len) - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - return prompt_embeds - - def decode_latents(self, latents): - latents = 1 / self.vae.config.scaling_factor * latents - mel_spectrogram = self.vae.decode(latents).sample - return mel_spectrogram - - def mel_spectrogram_to_waveform(self, mel_spectrogram): - if mel_spectrogram.dim() == 4: - mel_spectrogram = mel_spectrogram.squeeze(1) - - waveform = self.vocoder(mel_spectrogram) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - waveform = waveform.cpu().float() - return waveform - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - audio_length_in_s, - vocoder_upsample_factor, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - min_audio_length_in_s = vocoder_upsample_factor * self.vae_scale_factor - if audio_length_in_s < min_audio_length_in_s: - raise ValueError( - f"`audio_length_in_s` has to be a positive value greater than or equal to {min_audio_length_in_s}, but " - f"is {audio_length_in_s}." - ) - - if self.vocoder.config.model_in_dim % self.vae_scale_factor != 0: - raise ValueError( - f"The number of frequency bins in the vocoder's log-mel spectrogram has to be divisible by the " - f"VAE scale factor, but got {self.vocoder.config.model_in_dim} bins and a scale factor of " - f"{self.vae_scale_factor}." - ) - - if (callback_steps is None) or ( - 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 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}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents with width->self.vocoder.config.model_in_dim - def prepare_latents(self, batch_size, num_channels_latents, height, dtype, device, generator, latents=None): - shape = ( - batch_size, - num_channels_latents, - int(height) // self.vae_scale_factor, - int(self.vocoder.config.model_in_dim) // 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - audio_length_in_s: Optional[float] = 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, - 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, - 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", - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - num_inference_steps (`int`, *optional*, defaults to 10): - The number of denoising steps. More denoising steps usually lead to a higher quality audio at the - expense of slower inference. - 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*): - 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): - The number of waveforms to generate per prompt. - 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*): - 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. - 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. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.AudioPipelineOutput`] instead of a plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - output_type (`str`, *optional*, defaults to `"np"`): - The output format of the generated image. Choose between `"np"` to return a NumPy `np.ndarray` or - `"pt"` to return a PyTorch `torch.Tensor` object. - - Examples: - - Returns: - [`~pipelines.AudioPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.AudioPipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated audio. - """ - # 0. Convert audio input length from seconds to spectrogram height - vocoder_upsample_factor = np.prod(self.vocoder.config.upsample_rates) / self.vocoder.config.sampling_rate - - if audio_length_in_s is None: - audio_length_in_s = self.unet.config.sample_size * self.vae_scale_factor * vocoder_upsample_factor - - height = int(audio_length_in_s / vocoder_upsample_factor) - - original_waveform_length = int(audio_length_in_s * self.vocoder.config.sampling_rate) - if height % self.vae_scale_factor != 0: - height = int(np.ceil(height / self.vae_scale_factor)) * self.vae_scale_factor - logger.info( - f"Audio length in seconds {audio_length_in_s} is increased to {height * vocoder_upsample_factor} " - f"so that it can be handled by the model. It will be cut to {audio_length_in_s} after the " - f"denoising process." - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - audio_length_in_s, - vocoder_upsample_factor, - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds = self._encode_prompt( - prompt, - device, - num_waveforms_per_prompt, - do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_waveforms_per_prompt, - num_channels_latents, - height, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=None, - class_labels=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - ).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # 8. Post-processing - mel_spectrogram = self.decode_latents(latents) - - audio = self.mel_spectrogram_to_waveform(mel_spectrogram) - - audio = audio[:, :original_waveform_length] - - if output_type == "np": - audio = audio.numpy() - - if not return_dict: - return (audio,) - - return AudioPipelineOutput(audios=audio) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/__init__.py deleted file mode 100644 index 23cd0e44f89217b8391d0ce236070271db9aaf83..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/__init__.py +++ /dev/null @@ -1,50 +0,0 @@ -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 = {} -_import_structure = {} - -try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_torch_and_transformers_objects - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) -else: - _import_structure["modeling_audioldm2"] = ["AudioLDM2ProjectionModel", "AudioLDM2UNet2DConditionModel"] - _import_structure["pipeline_audioldm2"] = ["AudioLDM2Pipeline"] - - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_objects import * - - else: - from .modeling_audioldm2 import AudioLDM2ProjectionModel, AudioLDM2UNet2DConditionModel - from .pipeline_audioldm2 import AudioLDM2Pipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index d7f607aea6f210edf47e5eca61a76d45af9a95c9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/__pycache__/modeling_audioldm2.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/__pycache__/modeling_audioldm2.cpython-310.pyc deleted file mode 100644 index 7c561f2c1b004113c03cac5562e3d229226bef57..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/__pycache__/modeling_audioldm2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/__pycache__/pipeline_audioldm2.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/__pycache__/pipeline_audioldm2.cpython-310.pyc deleted file mode 100644 index ce0610b86da861a37d8ab78a0d694e2a47d103b5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/__pycache__/pipeline_audioldm2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/modeling_audioldm2.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/modeling_audioldm2.py deleted file mode 100644 index 546ae9239a02e54da68954d7c5b2abf98be04763..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/modeling_audioldm2.py +++ /dev/null @@ -1,1475 +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 dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union - -import torch -import torch.nn as nn -import torch.utils.checkpoint - -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import UNet2DConditionLoadersMixin -from ...models.activations import get_activation -from ...models.attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - AttentionProcessor, - AttnAddedKVProcessor, - AttnProcessor, -) -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 -from ...models.unets.unet_2d_blocks import DownBlock2D, UpBlock2D -from ...models.unets.unet_2d_condition import UNet2DConditionOutput -from ...utils import BaseOutput, logging - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -def add_special_tokens(hidden_states, attention_mask, sos_token, eos_token): - batch_size = hidden_states.shape[0] - - if attention_mask is not None: - # Add two more steps to attn mask - new_attn_mask_step = attention_mask.new_ones((batch_size, 1)) - attention_mask = torch.concat([new_attn_mask_step, attention_mask, new_attn_mask_step], dim=-1) - - # Add the SOS / EOS tokens at the start / end of the sequence respectively - sos_token = sos_token.expand(batch_size, 1, -1) - eos_token = eos_token.expand(batch_size, 1, -1) - hidden_states = torch.concat([sos_token, hidden_states, eos_token], dim=1) - return hidden_states, attention_mask - - -@dataclass -class AudioLDM2ProjectionModelOutput(BaseOutput): - """ - Args: - Class for AudioLDM2 projection layer's outputs. - hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): - Sequence of hidden-states obtained by linearly projecting the hidden-states for each of the text - encoders and subsequently concatenating them together. - attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Mask to avoid performing attention on padding token indices, formed by concatenating the attention masks - for the two text encoders together. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - """ - - hidden_states: torch.Tensor - attention_mask: Optional[torch.LongTensor] = None - - -class AudioLDM2ProjectionModel(ModelMixin, ConfigMixin): - """ - A simple linear projection model to map two text embeddings to a shared latent space. It also inserts learned - embedding vectors at the start and end of each text embedding sequence respectively. Each variable appended with - `_1` refers to that corresponding to the second text encoder. Otherwise, it is from the first. - - Args: - text_encoder_dim (`int`): - Dimensionality of the text embeddings from the first text encoder (CLAP). - text_encoder_1_dim (`int`): - Dimensionality of the text embeddings from the second text encoder (T5 or VITS). - langauge_model_dim (`int`): - Dimensionality of the text embeddings from the language model (GPT2). - """ - - @register_to_config - def __init__( - self, - text_encoder_dim, - text_encoder_1_dim, - langauge_model_dim, - use_learned_position_embedding=None, - max_seq_length=None, - ): - super().__init__() - # additional projection layers for each text encoder - self.projection = nn.Linear(text_encoder_dim, langauge_model_dim) - self.projection_1 = nn.Linear(text_encoder_1_dim, langauge_model_dim) - - # learnable SOS / EOS token embeddings for each text encoder - self.sos_embed = nn.Parameter(torch.ones(langauge_model_dim)) - self.eos_embed = nn.Parameter(torch.ones(langauge_model_dim)) - - self.sos_embed_1 = nn.Parameter(torch.ones(langauge_model_dim)) - self.eos_embed_1 = nn.Parameter(torch.ones(langauge_model_dim)) - - self.use_learned_position_embedding = use_learned_position_embedding - - # learable positional embedding for vits encoder - if self.use_learned_position_embedding is not None: - self.learnable_positional_embedding = torch.nn.Parameter( - torch.zeros((1, text_encoder_1_dim, max_seq_length)) - ) - - 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 = self.projection(hidden_states) - hidden_states, attention_mask = add_special_tokens( - hidden_states, attention_mask, sos_token=self.sos_embed, eos_token=self.eos_embed - ) - - # Add positional embedding for Vits hidden state - if self.use_learned_position_embedding is not None: - hidden_states_1 = (hidden_states_1.permute(0, 2, 1) + self.learnable_positional_embedding).permute(0, 2, 1) - - hidden_states_1 = self.projection_1(hidden_states_1) - hidden_states_1, attention_mask_1 = add_special_tokens( - hidden_states_1, attention_mask_1, sos_token=self.sos_embed_1, eos_token=self.eos_embed_1 - ) - - # concatenate clap and t5 text encoding - hidden_states = torch.cat([hidden_states, hidden_states_1], dim=1) - - # concatenate attention masks - if attention_mask is None and attention_mask_1 is not None: - attention_mask = attention_mask_1.new_ones((hidden_states[:2])) - elif attention_mask is not None and attention_mask_1 is None: - attention_mask_1 = attention_mask.new_ones((hidden_states_1[:2])) - - if attention_mask is not None and attention_mask_1 is not None: - attention_mask = torch.cat([attention_mask, attention_mask_1], dim=-1) - else: - attention_mask = None - - return AudioLDM2ProjectionModelOutput( - hidden_states=hidden_states, - attention_mask=attention_mask, - ) - - -class AudioLDM2UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): - r""" - A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample - shaped output. Compared to the vanilla [`UNet2DConditionModel`], this variant optionally includes an additional - self-attention layer in each Transformer block, as well as multiple cross-attention layers. It also allows for up - to two cross-attention embeddings, `encoder_hidden_states` and `encoder_hidden_states_1`. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - Parameters: - 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")`): - 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")`): - The tuple of upsample blocks to use. - 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)`): - 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. - mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block. - act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. - 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): - The dimension of the cross attention features. - 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`]. - attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads. - num_attention_heads (`int`, *optional*): - The number of attention heads. If not defined, defaults to `attention_head_dim` - resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config - for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`. - class_embed_type (`str`, *optional*, defaults to `None`): - The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`, - `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`. - num_class_embeds (`int`, *optional*, defaults to `None`): - Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing - class conditioning with `class_embed_type` equal to `None`. - time_embedding_type (`str`, *optional*, defaults to `positional`): - The type of position embedding to use for timesteps. Choose from `positional` or `fourier`. - time_embedding_dim (`int`, *optional*, defaults to `None`): - An optional override for the dimension of the projected time embedding. - time_embedding_act_fn (`str`, *optional*, defaults to `None`): - Optional activation function to use only once on the time embeddings before they are passed to the rest of - the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`. - timestep_post_act (`str`, *optional*, defaults to `None`): - The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`. - time_cond_proj_dim (`int`, *optional*, defaults to `None`): - The dimension of `cond_proj` layer in the timestep embedding. - conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer. - conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer. - projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when - `class_embed_type="projection"`. Required when `class_embed_type="projection"`. - class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time - embeddings with the class embeddings. - """ - - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - sample_size: Optional[int] = None, - in_channels: int = 4, - out_channels: int = 4, - 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", - 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, - 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: 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, - use_linear_projection: bool = False, - class_embed_type: Optional[str] = None, - num_class_embeds: Optional[int] = 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, - conv_in_kernel: int = 3, - conv_out_kernel: int = 3, - projection_class_embeddings_input_dim: Optional[int] = None, - class_embeddings_concat: bool = False, - ): - super().__init__() - - self.sample_size = sample_size - - if num_attention_heads is not None: - raise ValueError( - "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19." - ) - - # If `num_attention_heads` is not defined (which is the case for most models) - # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. - # The reason for this behavior is to correct for incorrectly named variables that were introduced - # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 - # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking - # which is why we correct for the naming here. - num_attention_heads = num_attention_heads or attention_head_dim - - # Check inputs - if len(down_block_types) != len(up_block_types): - raise ValueError( - f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}." - ) - - if len(block_out_channels) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}." - ) - - if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}." - ) - - # input - conv_in_padding = (conv_in_kernel - 1) // 2 - self.conv_in = nn.Conv2d( - in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding - ) - - # time - if time_embedding_type == "positional": - time_embed_dim = time_embedding_dim or block_out_channels[0] * 4 - - self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) - timestep_input_dim = block_out_channels[0] - else: - raise ValueError(f"{time_embedding_type} does not exist. Please make sure to use `positional`.") - - self.time_embedding = TimestepEmbedding( - timestep_input_dim, - time_embed_dim, - act_fn=act_fn, - post_act_fn=timestep_post_act, - cond_proj_dim=time_cond_proj_dim, - ) - - # class embedding - if class_embed_type is None and num_class_embeds is not None: - self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) - elif class_embed_type == "timestep": - self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn=act_fn) - elif class_embed_type == "identity": - self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) - elif class_embed_type == "projection": - if projection_class_embeddings_input_dim is None: - raise ValueError( - "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set" - ) - # The projection `class_embed_type` is the same as the timestep `class_embed_type` except - # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings - # 2. it projects from an arbitrary input dimension. - # - # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations. - # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings. - # As a result, `TimestepEmbedding` can be passed arbitrary vectors. - self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) - elif class_embed_type == "simple_projection": - if projection_class_embeddings_input_dim is None: - raise ValueError( - "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set" - ) - self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim) - else: - self.class_embedding = None - - if time_embedding_act_fn is None: - self.time_embed_act = None - else: - self.time_embed_act = get_activation(time_embedding_act_fn) - - self.down_blocks = nn.ModuleList([]) - self.up_blocks = nn.ModuleList([]) - - if isinstance(only_cross_attention, bool): - only_cross_attention = [only_cross_attention] * len(down_block_types) - - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(down_block_types) - - if isinstance(cross_attention_dim, int): - cross_attention_dim = (cross_attention_dim,) * len(down_block_types) - - if isinstance(layers_per_block, int): - layers_per_block = [layers_per_block] * len(down_block_types) - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) - - if class_embeddings_concat: - # The time embeddings are concatenated with the class embeddings. The dimension of the - # time embeddings passed to the down, middle, and up blocks is twice the dimension of the - # regular time embeddings - blocks_time_embed_dim = time_embed_dim * 2 - else: - blocks_time_embed_dim = time_embed_dim - - # down - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = get_down_block( - down_block_type, - num_layers=layers_per_block[i], - transformer_layers_per_block=transformer_layers_per_block[i], - in_channels=input_channel, - out_channels=output_channel, - temb_channels=blocks_time_embed_dim, - add_downsample=not is_final_block, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - cross_attention_dim=cross_attention_dim[i], - num_attention_heads=num_attention_heads[i], - downsample_padding=downsample_padding, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention[i], - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - self.down_blocks.append(down_block) - - # mid - if mid_block_type == "UNetMidBlock2DCrossAttn": - self.mid_block = UNetMidBlock2DCrossAttn( - transformer_layers_per_block=transformer_layers_per_block[-1], - in_channels=block_out_channels[-1], - temb_channels=blocks_time_embed_dim, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - output_scale_factor=mid_block_scale_factor, - resnet_time_scale_shift=resnet_time_scale_shift, - cross_attention_dim=cross_attention_dim[-1], - num_attention_heads=num_attention_heads[-1], - resnet_groups=norm_num_groups, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - ) - else: - raise ValueError( - f"unknown mid_block_type : {mid_block_type}. Should be `UNetMidBlock2DCrossAttn` for AudioLDM2." - ) - - # count how many layers upsample the images - self.num_upsamplers = 0 - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - reversed_num_attention_heads = list(reversed(num_attention_heads)) - reversed_layers_per_block = list(reversed(layers_per_block)) - reversed_cross_attention_dim = list(reversed(cross_attention_dim)) - reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block)) - only_cross_attention = list(reversed(only_cross_attention)) - - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - is_final_block = i == len(block_out_channels) - 1 - - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] - - # add upsample block for all BUT final layer - if not is_final_block: - add_upsample = True - self.num_upsamplers += 1 - else: - add_upsample = False - - up_block = get_up_block( - up_block_type, - num_layers=reversed_layers_per_block[i] + 1, - transformer_layers_per_block=reversed_transformer_layers_per_block[i], - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - temb_channels=blocks_time_embed_dim, - add_upsample=add_upsample, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - cross_attention_dim=reversed_cross_attention_dim[i], - num_attention_heads=reversed_num_attention_heads[i], - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention[i], - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - self.up_blocks.append(up_block) - prev_output_channel = output_channel - - # out - if norm_num_groups is not None: - self.conv_norm_out = nn.GroupNorm( - num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps - ) - - self.conv_act = get_activation(act_fn) - - else: - self.conv_norm_out = None - self.conv_act = None - - conv_out_padding = (conv_out_kernel - 1) // 2 - self.conv_out = nn.Conv2d( - block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding - ) - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - self.set_attn_processor(processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attention_slice - def set_attention_slice(self, slice_size): - 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. This is useful for saving some memory in exchange for a small decrease in speed. - - Args: - slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): - When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If - `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is - provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` - must be a multiple of `slice_size`. - """ - sliceable_head_dims = [] - - def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): - if hasattr(module, "set_attention_slice"): - sliceable_head_dims.append(module.sliceable_head_dim) - - for child in module.children(): - fn_recursive_retrieve_sliceable_dims(child) - - # retrieve number of attention layers - for module in self.children(): - fn_recursive_retrieve_sliceable_dims(module) - - num_sliceable_layers = len(sliceable_head_dims) - - if slice_size == "auto": - # half the attention head size is usually a good trade-off between - # speed and memory - slice_size = [dim // 2 for dim in sliceable_head_dims] - elif slice_size == "max": - # make smallest slice possible - slice_size = num_sliceable_layers * [1] - - slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size - - if len(slice_size) != len(sliceable_head_dims): - raise ValueError( - f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" - f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." - ) - - for i in range(len(slice_size)): - size = slice_size[i] - dim = sliceable_head_dims[i] - if size is not None and size > dim: - raise ValueError(f"size {size} has to be smaller or equal to {dim}.") - - # 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]): - if hasattr(module, "set_attention_slice"): - module.set_attention_slice(slice_size.pop()) - - for child in module.children(): - fn_recursive_set_attention_slice(child, slice_size) - - reversed_slice_size = list(reversed(slice_size)) - for module in self.children(): - fn_recursive_set_attention_slice(module, reversed_slice_size) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[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, - return_dict: bool = True, - encoder_hidden_states_1: Optional[torch.Tensor] = None, - encoder_attention_mask_1: Optional[torch.Tensor] = None, - ) -> Union[UNet2DConditionOutput, Tuple]: - r""" - The [`AudioLDM2UNet2DConditionModel`] forward method. - - Args: - sample (`torch.Tensor`): - The noisy input tensor with the following shape `(batch, channel, height, width)`. - timestep (`torch.Tensor` or `float` or `int`): The number of timesteps to denoise an input. - encoder_hidden_states (`torch.Tensor`): - The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. - encoder_attention_mask (`torch.Tensor`): - A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If - `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias, - which adds large negative values to the attention scores corresponding to "discard" tokens. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain - tuple. - cross_attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the [`AttnProcessor`]. - encoder_hidden_states_1 (`torch.Tensor`, *optional*): - A second set of encoder hidden states with shape `(batch, sequence_length_2, feature_dim_2)`. Can be - used to condition the model on a different set of embeddings to `encoder_hidden_states`. - encoder_attention_mask_1 (`torch.Tensor`, *optional*): - A cross-attention mask of shape `(batch, sequence_length_2)` is applied to `encoder_hidden_states_1`. - If `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias, - which adds large negative values to the attention scores corresponding to "discard" tokens. - - Returns: - [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: - If `return_dict` is True, an [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] is returned, - otherwise a `tuple` is returned where the first element is the sample tensor. - """ - # By default samples have to be AT least a multiple of the overall upsampling factor. - # The overall upsampling factor is equal to 2 ** (# num of upsampling layers). - # However, the upsampling interpolation output size can be forced to fit any upsampling size - # on the fly if necessary. - default_overall_up_factor = 2**self.num_upsamplers - - # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` - forward_upsample_size = False - upsample_size = None - - if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): - logger.info("Forward upsample size to force interpolation output size.") - forward_upsample_size = True - - # ensure attention_mask is a bias, and give it a singleton query_tokens dimension - # expects mask of shape: - # [batch, key_tokens] - # adds singleton query_tokens dimension: - # [batch, 1, key_tokens] - # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: - # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) - # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) - if attention_mask is not None: - # assume that mask is expressed as: - # (1 = keep, 0 = discard) - # convert mask into a bias that can be added to attention scores: - # (keep = +0, discard = -10000.0) - attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 - attention_mask = attention_mask.unsqueeze(1) - - # convert encoder_attention_mask to a bias the same way we do for attention_mask - if encoder_attention_mask is not None: - encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0 - encoder_attention_mask = encoder_attention_mask.unsqueeze(1) - - if encoder_attention_mask_1 is not None: - encoder_attention_mask_1 = (1 - encoder_attention_mask_1.to(sample.dtype)) * -10000.0 - encoder_attention_mask_1 = encoder_attention_mask_1.unsqueeze(1) - - # 1. time - 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 - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = sample.device.type == "mps" - is_npu = sample.device.type == "npu" - if isinstance(timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) - elif len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps.expand(sample.shape[0]) - - t_emb = self.time_proj(timesteps) - - # `Timesteps` does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - t_emb = t_emb.to(dtype=sample.dtype) - - emb = self.time_embedding(t_emb, timestep_cond) - aug_emb = None - - if self.class_embedding is not None: - if class_labels is None: - raise ValueError("class_labels should be provided when num_class_embeds > 0") - - if self.config.class_embed_type == "timestep": - class_labels = self.time_proj(class_labels) - - # `Timesteps` does not contain any weights and will always return f32 tensors - # there might be better ways to encapsulate this. - class_labels = class_labels.to(dtype=sample.dtype) - - class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype) - - if self.config.class_embeddings_concat: - emb = torch.cat([emb, class_emb], dim=-1) - else: - emb = emb + class_emb - - emb = emb + aug_emb if aug_emb is not None else emb - - if self.time_embed_act is not None: - emb = self.time_embed_act(emb) - - # 2. pre-process - sample = self.conv_in(sample) - - # 3. down - down_block_res_samples = (sample,) - for downsample_block in self.down_blocks: - if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: - sample, res_samples = downsample_block( - hidden_states=sample, - temb=emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - encoder_attention_mask=encoder_attention_mask, - encoder_hidden_states_1=encoder_hidden_states_1, - encoder_attention_mask_1=encoder_attention_mask_1, - ) - else: - sample, res_samples = downsample_block(hidden_states=sample, temb=emb) - - down_block_res_samples += res_samples - - # 4. mid - if self.mid_block is not None: - sample = self.mid_block( - sample, - emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - encoder_attention_mask=encoder_attention_mask, - encoder_hidden_states_1=encoder_hidden_states_1, - encoder_attention_mask_1=encoder_attention_mask_1, - ) - - # 5. up - for i, upsample_block in enumerate(self.up_blocks): - is_final_block = i == len(self.up_blocks) - 1 - - res_samples = down_block_res_samples[-len(upsample_block.resnets) :] - down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] - - # if we have not reached the final block and need to forward the - # upsample size, we do it here - if not is_final_block and forward_upsample_size: - upsample_size = down_block_res_samples[-1].shape[2:] - - if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - upsample_size=upsample_size, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - encoder_hidden_states_1=encoder_hidden_states_1, - encoder_attention_mask_1=encoder_attention_mask_1, - ) - else: - sample = upsample_block( - hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size - ) - - # 6. post-process - if self.conv_norm_out: - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - sample = self.conv_out(sample) - - if not return_dict: - return (sample,) - - return UNet2DConditionOutput(sample=sample) - - -def get_down_block( - down_block_type, - num_layers, - in_channels, - out_channels, - temb_channels, - add_downsample, - resnet_eps, - resnet_act_fn, - transformer_layers_per_block=1, - num_attention_heads=None, - resnet_groups=None, - cross_attention_dim=None, - downsample_padding=None, - use_linear_projection=False, - only_cross_attention=False, - upcast_attention=False, - resnet_time_scale_shift="default", -): - down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type - if down_block_type == "DownBlock2D": - return DownBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - downsample_padding=downsample_padding, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif down_block_type == "CrossAttnDownBlock2D": - if cross_attention_dim is None: - raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock2D") - return CrossAttnDownBlock2D( - num_layers=num_layers, - transformer_layers_per_block=transformer_layers_per_block, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - downsample_padding=downsample_padding, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - raise ValueError(f"{down_block_type} does not exist.") - - -def get_up_block( - up_block_type, - num_layers, - in_channels, - out_channels, - prev_output_channel, - temb_channels, - add_upsample, - resnet_eps, - resnet_act_fn, - transformer_layers_per_block=1, - num_attention_heads=None, - resnet_groups=None, - cross_attention_dim=None, - use_linear_projection=False, - only_cross_attention=False, - upcast_attention=False, - resnet_time_scale_shift="default", -): - up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type - if up_block_type == "UpBlock2D": - return UpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif up_block_type == "CrossAttnUpBlock2D": - if cross_attention_dim is None: - raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D") - return CrossAttnUpBlock2D( - num_layers=num_layers, - transformer_layers_per_block=transformer_layers_per_block, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - raise ValueError(f"{up_block_type} does not exist.") - - -class CrossAttnDownBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - num_attention_heads=1, - cross_attention_dim=1280, - output_scale_factor=1.0, - downsample_padding=1, - add_downsample=True, - use_linear_projection=False, - only_cross_attention=False, - upcast_attention=False, - ): - super().__init__() - resnets = [] - attentions = [] - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - - if isinstance(cross_attention_dim, int): - cross_attention_dim = (cross_attention_dim,) - if isinstance(cross_attention_dim, (list, tuple)) and len(cross_attention_dim) > 4: - raise ValueError( - "Only up to 4 cross-attention layers are supported. Ensure that the length of cross-attention " - f"dims is less than or equal to 4. Got cross-attention dims {cross_attention_dim} of length {len(cross_attention_dim)}" - ) - self.cross_attention_dim = cross_attention_dim - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - for j in range(len(cross_attention_dim)): - attentions.append( - Transformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=transformer_layers_per_block, - cross_attention_dim=cross_attention_dim[j], - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - double_self_attention=True if cross_attention_dim[j] is None else False, - ) - ) - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - Downsample2D( - out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - 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, - ): - output_states = () - num_layers = len(self.resnets) - num_attention_per_layer = len(self.attentions) // num_layers - - encoder_hidden_states_1 = ( - encoder_hidden_states_1 if encoder_hidden_states_1 is not None else encoder_hidden_states - ) - encoder_attention_mask_1 = ( - encoder_attention_mask_1 if encoder_hidden_states_1 is not None else encoder_attention_mask - ) - - for i in range(num_layers): - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(self.resnets[i], hidden_states, temb) - for idx, cross_attention_dim in enumerate(self.cross_attention_dim): - if cross_attention_dim is not None and idx <= 1: - forward_encoder_hidden_states = encoder_hidden_states - forward_encoder_attention_mask = encoder_attention_mask - elif cross_attention_dim is not None and idx > 1: - forward_encoder_hidden_states = encoder_hidden_states_1 - forward_encoder_attention_mask = encoder_attention_mask_1 - else: - forward_encoder_hidden_states = None - forward_encoder_attention_mask = None - hidden_states = self._gradient_checkpointing_func( - self.attentions[i * num_attention_per_layer + idx], - hidden_states, - forward_encoder_hidden_states, - None, # timestep - None, # class_labels - cross_attention_kwargs, - attention_mask, - forward_encoder_attention_mask, - )[0] - else: - hidden_states = self.resnets[i](hidden_states, temb) - for idx, cross_attention_dim in enumerate(self.cross_attention_dim): - if cross_attention_dim is not None and idx <= 1: - forward_encoder_hidden_states = encoder_hidden_states - forward_encoder_attention_mask = encoder_attention_mask - elif cross_attention_dim is not None and idx > 1: - forward_encoder_hidden_states = encoder_hidden_states_1 - forward_encoder_attention_mask = encoder_attention_mask_1 - else: - forward_encoder_hidden_states = None - forward_encoder_attention_mask = None - hidden_states = self.attentions[i * num_attention_per_layer + idx]( - hidden_states, - attention_mask=attention_mask, - encoder_hidden_states=forward_encoder_hidden_states, - encoder_attention_mask=forward_encoder_attention_mask, - return_dict=False, - )[0] - - output_states = output_states + (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - output_states = output_states + (hidden_states,) - - return hidden_states, output_states - - -class UNetMidBlock2DCrossAttn(nn.Module): - def __init__( - self, - in_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - num_attention_heads=1, - output_scale_factor=1.0, - cross_attention_dim=1280, - use_linear_projection=False, - upcast_attention=False, - ): - super().__init__() - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) - - if isinstance(cross_attention_dim, int): - cross_attention_dim = (cross_attention_dim,) - if isinstance(cross_attention_dim, (list, tuple)) and len(cross_attention_dim) > 4: - raise ValueError( - "Only up to 4 cross-attention layers are supported. Ensure that the length of cross-attention " - f"dims is less than or equal to 4. Got cross-attention dims {cross_attention_dim} of length {len(cross_attention_dim)}" - ) - self.cross_attention_dim = cross_attention_dim - - # there is always at least one resnet - resnets = [ - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ] - attentions = [] - - for i in range(num_layers): - for j in range(len(cross_attention_dim)): - attentions.append( - Transformer2DModel( - num_attention_heads, - in_channels // num_attention_heads, - in_channels=in_channels, - num_layers=transformer_layers_per_block, - cross_attention_dim=cross_attention_dim[j], - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - double_self_attention=True if cross_attention_dim[j] is None else False, - ) - ) - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - self.gradient_checkpointing = False - - 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, - ) -> torch.Tensor: - hidden_states = self.resnets[0](hidden_states, temb) - num_attention_per_layer = len(self.attentions) // (len(self.resnets) - 1) - - encoder_hidden_states_1 = ( - encoder_hidden_states_1 if encoder_hidden_states_1 is not None else encoder_hidden_states - ) - encoder_attention_mask_1 = ( - encoder_attention_mask_1 if encoder_hidden_states_1 is not None else encoder_attention_mask - ) - - for i in range(len(self.resnets[1:])): - if torch.is_grad_enabled() and self.gradient_checkpointing: - for idx, cross_attention_dim in enumerate(self.cross_attention_dim): - if cross_attention_dim is not None and idx <= 1: - forward_encoder_hidden_states = encoder_hidden_states - forward_encoder_attention_mask = encoder_attention_mask - elif cross_attention_dim is not None and idx > 1: - forward_encoder_hidden_states = encoder_hidden_states_1 - forward_encoder_attention_mask = encoder_attention_mask_1 - else: - forward_encoder_hidden_states = None - forward_encoder_attention_mask = None - hidden_states = self._gradient_checkpointing_func( - self.attentions[i * num_attention_per_layer + idx], - hidden_states, - forward_encoder_hidden_states, - None, # timestep - None, # class_labels - cross_attention_kwargs, - attention_mask, - forward_encoder_attention_mask, - )[0] - hidden_states = self._gradient_checkpointing_func(self.resnets[i + 1], hidden_states, temb) - else: - for idx, cross_attention_dim in enumerate(self.cross_attention_dim): - if cross_attention_dim is not None and idx <= 1: - forward_encoder_hidden_states = encoder_hidden_states - forward_encoder_attention_mask = encoder_attention_mask - elif cross_attention_dim is not None and idx > 1: - forward_encoder_hidden_states = encoder_hidden_states_1 - forward_encoder_attention_mask = encoder_attention_mask_1 - else: - forward_encoder_hidden_states = None - forward_encoder_attention_mask = None - hidden_states = self.attentions[i * num_attention_per_layer + idx]( - hidden_states, - attention_mask=attention_mask, - encoder_hidden_states=forward_encoder_hidden_states, - encoder_attention_mask=forward_encoder_attention_mask, - return_dict=False, - )[0] - - hidden_states = self.resnets[i + 1](hidden_states, temb) - - return hidden_states - - -class CrossAttnUpBlock2D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - prev_output_channel: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - num_attention_heads=1, - cross_attention_dim=1280, - output_scale_factor=1.0, - add_upsample=True, - use_linear_projection=False, - only_cross_attention=False, - upcast_attention=False, - ): - super().__init__() - resnets = [] - attentions = [] - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - - if isinstance(cross_attention_dim, int): - cross_attention_dim = (cross_attention_dim,) - if isinstance(cross_attention_dim, (list, tuple)) and len(cross_attention_dim) > 4: - raise ValueError( - "Only up to 4 cross-attention layers are supported. Ensure that the length of cross-attention " - f"dims is less than or equal to 4. Got cross-attention dims {cross_attention_dim} of length {len(cross_attention_dim)}" - ) - self.cross_attention_dim = cross_attention_dim - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - for j in range(len(cross_attention_dim)): - attentions.append( - Transformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=transformer_layers_per_block, - cross_attention_dim=cross_attention_dim[j], - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - double_self_attention=True if cross_attention_dim[j] is None else False, - ) - ) - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - if add_upsample: - self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - - 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, - ): - num_layers = len(self.resnets) - num_attention_per_layer = len(self.attentions) // num_layers - - encoder_hidden_states_1 = ( - encoder_hidden_states_1 if encoder_hidden_states_1 is not None else encoder_hidden_states - ) - encoder_attention_mask_1 = ( - encoder_attention_mask_1 if encoder_hidden_states_1 is not None else encoder_attention_mask - ) - - for i in range(num_layers): - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(self.resnets[i], hidden_states, temb) - for idx, cross_attention_dim in enumerate(self.cross_attention_dim): - if cross_attention_dim is not None and idx <= 1: - forward_encoder_hidden_states = encoder_hidden_states - forward_encoder_attention_mask = encoder_attention_mask - elif cross_attention_dim is not None and idx > 1: - forward_encoder_hidden_states = encoder_hidden_states_1 - forward_encoder_attention_mask = encoder_attention_mask_1 - else: - forward_encoder_hidden_states = None - forward_encoder_attention_mask = None - hidden_states = self._gradient_checkpointing_func( - self.attentions[i * num_attention_per_layer + idx], - hidden_states, - forward_encoder_hidden_states, - None, # timestep - None, # class_labels - cross_attention_kwargs, - attention_mask, - forward_encoder_attention_mask, - )[0] - else: - hidden_states = self.resnets[i](hidden_states, temb) - for idx, cross_attention_dim in enumerate(self.cross_attention_dim): - if cross_attention_dim is not None and idx <= 1: - forward_encoder_hidden_states = encoder_hidden_states - forward_encoder_attention_mask = encoder_attention_mask - elif cross_attention_dim is not None and idx > 1: - forward_encoder_hidden_states = encoder_hidden_states_1 - forward_encoder_attention_mask = encoder_attention_mask_1 - else: - forward_encoder_hidden_states = None - forward_encoder_attention_mask = None - hidden_states = self.attentions[i * num_attention_per_layer + idx]( - hidden_states, - attention_mask=attention_mask, - encoder_hidden_states=forward_encoder_hidden_states, - encoder_attention_mask=forward_encoder_attention_mask, - return_dict=False, - )[0] - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) - - return hidden_states diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/pipeline_audioldm2.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/pipeline_audioldm2.py deleted file mode 100644 index 2a376013235671baac916dedd8d7cd26747c4e1d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/audioldm2/pipeline_audioldm2.py +++ /dev/null @@ -1,1105 +0,0 @@ -# Copyright 2025 CVSSP, ByteDance 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 ( - ClapFeatureExtractor, - ClapModel, - GPT2LMHeadModel, - RobertaTokenizer, - RobertaTokenizerFast, - SpeechT5HifiGan, - T5EncoderModel, - T5Tokenizer, - T5TokenizerFast, - VitsModel, - VitsTokenizer, -) - -from ...models import AutoencoderKL -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - is_accelerate_available, - is_accelerate_version, - is_librosa_available, - logging, - replace_example_docstring, -) -from ...utils.import_utils import is_transformers_version -from ...utils.torch_utils import empty_device_cache, randn_tensor -from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline -from .modeling_audioldm2 import AudioLDM2ProjectionModel, AudioLDM2UNet2DConditionModel - - -if is_librosa_available(): - import librosa - - -from ...utils import is_torch_xla_available - - -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 scipy - >>> import torch - >>> from diffusers import AudioLDM2Pipeline - - >>> repo_id = "cvssp/audioldm2" - >>> pipe = AudioLDM2Pipeline.from_pretrained(repo_id, torch_dtype=torch.float16) - >>> pipe = pipe.to("cuda") - - >>> # define the prompts - >>> prompt = "The sound of a hammer hitting a wooden surface." - >>> negative_prompt = "Low quality." - - >>> # set the seed for generator - >>> generator = torch.Generator("cuda").manual_seed(0) - - >>> # run the generation - >>> audio = pipe( - ... prompt, - ... negative_prompt=negative_prompt, - ... num_inference_steps=200, - ... audio_length_in_s=10.0, - ... num_waveforms_per_prompt=3, - ... generator=generator, - ... ).audios - - >>> # save the best audio sample (index 0) as a .wav file - >>> scipy.io.wavfile.write("techno.wav", rate=16000, data=audio[0]) - ``` - ``` - #Using AudioLDM2 for Text To Speech - >>> import scipy - >>> import torch - >>> from diffusers import AudioLDM2Pipeline - - >>> repo_id = "anhnct/audioldm2_gigaspeech" - >>> pipe = AudioLDM2Pipeline.from_pretrained(repo_id, torch_dtype=torch.float16) - >>> pipe = pipe.to("cuda") - - >>> # define the prompts - >>> prompt = "A female reporter is speaking" - >>> transcript = "wish you have a good day" - - >>> # set the seed for generator - >>> generator = torch.Generator("cuda").manual_seed(0) - - >>> # run the generation - >>> audio = pipe( - ... prompt, - ... transcription=transcript, - ... num_inference_steps=200, - ... audio_length_in_s=10.0, - ... num_waveforms_per_prompt=2, - ... generator=generator, - ... max_new_tokens=512, #Must set max_new_tokens equa to 512 for TTS - ... ).audios - - >>> # save the best audio sample (index 0) as a .wav file - >>> scipy.io.wavfile.write("tts.wav", rate=16000, data=audio[0]) - ``` -""" - - -def prepare_inputs_for_generation( - inputs_embeds, - attention_mask=None, - past_key_values=None, - **kwargs, -): - if past_key_values is not None: - # only last token for inputs_embeds if past is defined in kwargs - inputs_embeds = inputs_embeds[:, -1:] - - return { - "inputs_embeds": inputs_embeds, - "attention_mask": attention_mask, - "past_key_values": past_key_values, - "use_cache": kwargs.get("use_cache"), - } - - -class AudioLDM2Pipeline(DiffusionPipeline): - r""" - Pipeline for text-to-audio generation using AudioLDM2. - - 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: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.ClapModel`]): - First frozen text-encoder. AudioLDM2 uses the joint audio-text embedding model - [CLAP](https://huggingface.co/docs/transformers/model_doc/clap#transformers.CLAPTextModelWithProjection), - specifically the [laion/clap-htsat-unfused](https://huggingface.co/laion/clap-htsat-unfused) variant. The - text branch is used to encode the text prompt to a prompt embedding. The full audio-text model is used to - rank generated waveforms against the text prompt by computing similarity scores. - text_encoder_2 ([`~transformers.T5EncoderModel`, `~transformers.VitsModel`]): - Second frozen text-encoder. AudioLDM2 uses the encoder of - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the - [google/flan-t5-large](https://huggingface.co/google/flan-t5-large) variant. Second frozen text-encoder use - for TTS. AudioLDM2 uses the encoder of - [Vits](https://huggingface.co/docs/transformers/model_doc/vits#transformers.VitsModel). - projection_model ([`AudioLDM2ProjectionModel`]): - A trained model used to linearly project the hidden-states from the first and second text encoder models - and insert learned SOS and EOS token embeddings. The projected hidden-states from the two text encoders are - concatenated to give the input to the language model. A Learned Position Embedding for the Vits - hidden-states - language_model ([`~transformers.GPT2Model`]): - An auto-regressive language model used to generate a sequence of hidden-states conditioned on the projected - outputs from the two text encoders. - tokenizer ([`~transformers.RobertaTokenizer`]): - Tokenizer to tokenize text for the first frozen text-encoder. - tokenizer_2 ([`~transformers.T5Tokenizer`, `~transformers.VitsTokenizer`]): - Tokenizer to tokenize text for the second frozen text-encoder. - feature_extractor ([`~transformers.ClapFeatureExtractor`]): - Feature extractor to pre-process generated audio waveforms to log-mel spectrograms for automatic scoring. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded audio latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded audio latents. Can be one of - [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. - vocoder ([`~transformers.SpeechT5HifiGan`]): - Vocoder of class `SpeechT5HifiGan` to convert the mel-spectrogram latents to the final audio waveform. - """ - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: ClapModel, - text_encoder_2: Union[T5EncoderModel, VitsModel], - projection_model: AudioLDM2ProjectionModel, - language_model: GPT2LMHeadModel, - tokenizer: Union[RobertaTokenizer, RobertaTokenizerFast], - tokenizer_2: Union[T5Tokenizer, T5TokenizerFast, VitsTokenizer], - feature_extractor: ClapFeatureExtractor, - unet: AudioLDM2UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - vocoder: SpeechT5HifiGan, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - projection_model=projection_model, - language_model=language_model, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - feature_extractor=feature_extractor, - unet=unet, - scheduler=scheduler, - vocoder=vocoder, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.enable_vae_slicing - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.disable_vae_slicing - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[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` - method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with - `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`. - """ - if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"): - from accelerate import cpu_offload_with_hook - else: - raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.") - - torch_device = torch.device(device) - device_index = torch_device.index - - if gpu_id is not None and device_index is not None: - raise ValueError( - f"You have passed both `gpu_id`={gpu_id} and an index as part of the passed device `device`={device}" - f"Cannot pass both. Please make sure to either not define `gpu_id` or not pass the index as part of the device: `device`={torch_device.type}" - ) - - device_type = torch_device.type - device_str = device_type - if gpu_id or torch_device.index: - device_str = f"{device_str}:{gpu_id or torch_device.index}" - device = torch.device(device_str) - - if self.device.type != "cpu": - self.to("cpu", silence_dtype_warnings=True) - empty_device_cache(device.type) - - model_sequence = [ - self.text_encoder.text_model, - self.text_encoder.text_projection, - self.text_encoder_2, - self.projection_model, - self.language_model, - self.unet, - self.vae, - self.vocoder, - self.text_encoder, - ] - - hook = None - for cpu_offloaded_model in model_sequence: - _, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook) - - # We'll offload the last model manually. - self.final_offload_hook = hook - - def generate_language_model( - self, - inputs_embeds: torch.Tensor = None, - max_new_tokens: int = 8, - **model_kwargs, - ): - """ - - Generates a sequence of hidden-states from the language model, conditioned on the embedding inputs. - - Parameters: - inputs_embeds (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): - 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*): - Ad hoc parametrization of additional model-specific kwargs that will be forwarded to the `forward` - function of the model. - - Return: - `inputs_embeds (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): - The sequence of generated hidden-states. - """ - cache_position_kwargs = {} - if is_transformers_version("<", "4.52.0.dev0"): - cache_position_kwargs["input_ids"] = inputs_embeds - cache_position_kwargs["model_kwargs"] = model_kwargs - else: - cache_position_kwargs["seq_length"] = inputs_embeds.shape[0] - cache_position_kwargs["device"] = ( - self.language_model.device if getattr(self, "language_model", None) is not None else self.device - ) - cache_position_kwargs["model_kwargs"] = model_kwargs - max_new_tokens = max_new_tokens if max_new_tokens is not None else self.language_model.config.max_new_tokens - model_kwargs = self.language_model._get_initial_cache_position(**cache_position_kwargs) - - for _ in range(max_new_tokens): - # prepare model inputs - model_inputs = prepare_inputs_for_generation(inputs_embeds, **model_kwargs) - - # forward pass to get next hidden states - output = self.language_model(**model_inputs, output_hidden_states=True, return_dict=True) - - next_hidden_states = output.hidden_states[-1] - - # Update the model input - inputs_embeds = torch.cat([inputs_embeds, next_hidden_states[:, -1:, :]], dim=1) - - # Update generated hidden states, model inputs, and length for next step - model_kwargs = self.language_model._update_model_kwargs_for_generation(output, model_kwargs) - - return inputs_embeds[:, -max_new_tokens:, :] - - def encode_prompt( - self, - prompt, - device, - num_waveforms_per_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, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`, *optional*): - prompt to be encoded - transcription (`str` or `List[str]`): - transcription of text to speech - device (`torch.device`): - torch device - num_waveforms_per_prompt (`int`): - 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*): - 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`). - prompt_embeds (`torch.Tensor`, *optional*): - Pre-computed text embeddings from the Flan T5 model. Can be used to easily tweak text inputs, *e.g.* - prompt weighting. If not provided, text embeddings will be computed from `prompt` input argument. - negative_prompt_embeds (`torch.Tensor`, *optional*): - Pre-computed negative text embeddings from the Flan T5 model. Can be used to easily tweak text inputs, - *e.g.* prompt weighting. If not provided, negative_prompt_embeds will be computed from - `negative_prompt` input argument. - generated_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated text embeddings from the GPT2 language model. 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_generated_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated negative text embeddings from the GPT2 language model. Can be used to easily tweak text - inputs, *e.g.* prompt weighting. If not provided, negative_prompt_embeds will be computed from - `negative_prompt` input argument. - attention_mask (`torch.LongTensor`, *optional*): - Pre-computed attention mask to be applied to the `prompt_embeds`. If not provided, attention mask will - be computed from `prompt` input argument. - negative_attention_mask (`torch.LongTensor`, *optional*): - Pre-computed attention mask to be applied to the `negative_prompt_embeds`. If not provided, attention - mask will be computed from `negative_prompt` input argument. - max_new_tokens (`int`, *optional*, defaults to None): - The number of new tokens to generate with the GPT2 language model. - Returns: - prompt_embeds (`torch.Tensor`): - Text embeddings from the Flan T5 model. - attention_mask (`torch.LongTensor`): - Attention mask to be applied to the `prompt_embeds`. - generated_prompt_embeds (`torch.Tensor`): - Text embeddings generated from the GPT2 language model. - - Example: - - ```python - >>> import scipy - >>> import torch - >>> from diffusers import AudioLDM2Pipeline - - >>> repo_id = "cvssp/audioldm2" - >>> pipe = AudioLDM2Pipeline.from_pretrained(repo_id, torch_dtype=torch.float16) - >>> pipe = pipe.to("cuda") - - >>> # Get text embedding vectors - >>> prompt_embeds, attention_mask, generated_prompt_embeds = pipe.encode_prompt( - ... prompt="Techno music with a strong, upbeat tempo and high melodic riffs", - ... device="cuda", - ... do_classifier_free_guidance=True, - ... ) - - >>> # Pass text embeddings to pipeline for text-conditional audio generation - >>> audio = pipe( - ... prompt_embeds=prompt_embeds, - ... attention_mask=attention_mask, - ... generated_prompt_embeds=generated_prompt_embeds, - ... num_inference_steps=200, - ... audio_length_in_s=10.0, - ... ).audios[0] - - >>> # save generated audio sample - >>> scipy.io.wavfile.write("techno.wav", rate=16000, data=audio) - ```""" - 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] - - # Define tokenizers and text encoders - tokenizers = [self.tokenizer, self.tokenizer_2] - is_vits_text_encoder = isinstance(self.text_encoder_2, VitsModel) - - if is_vits_text_encoder: - text_encoders = [self.text_encoder, self.text_encoder_2.text_encoder] - else: - text_encoders = [self.text_encoder, self.text_encoder_2] - - if prompt_embeds is None: - prompt_embeds_list = [] - attention_mask_list = [] - - for tokenizer, text_encoder in zip(tokenizers, text_encoders): - use_prompt = isinstance( - tokenizer, (RobertaTokenizer, RobertaTokenizerFast, T5Tokenizer, T5TokenizerFast) - ) - text_inputs = tokenizer( - prompt if use_prompt else transcription, - padding="max_length" - if isinstance(tokenizer, (RobertaTokenizer, RobertaTokenizerFast, VitsTokenizer)) - else True, - max_length=tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - attention_mask = text_inputs.attention_mask - 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( - f"The following part of your input was truncated because {text_encoder.config.model_type} can " - f"only handle sequences up to {tokenizer.model_max_length} tokens: {removed_text}" - ) - - text_input_ids = text_input_ids.to(device) - attention_mask = attention_mask.to(device) - - if text_encoder.config.model_type == "clap": - prompt_embeds = text_encoder.get_text_features( - text_input_ids, - attention_mask=attention_mask, - ) - # 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 - attention_mask = attention_mask.new_ones((batch_size, 1)) - elif is_vits_text_encoder: - # Add end_token_id and attention mask in the end of sequence phonemes - for text_input_id, text_attention_mask in zip(text_input_ids, attention_mask): - for idx, phoneme_id in enumerate(text_input_id): - if phoneme_id == 0: - text_input_id[idx] = 182 - text_attention_mask[idx] = 1 - break - prompt_embeds = text_encoder( - text_input_ids, attention_mask=attention_mask, padding_mask=attention_mask.unsqueeze(-1) - ) - prompt_embeds = prompt_embeds[0] - else: - prompt_embeds = text_encoder( - text_input_ids, - attention_mask=attention_mask, - ) - prompt_embeds = prompt_embeds[0] - - prompt_embeds_list.append(prompt_embeds) - attention_mask_list.append(attention_mask) - - projection_output = self.projection_model( - hidden_states=prompt_embeds_list[0], - hidden_states_1=prompt_embeds_list[1], - attention_mask=attention_mask_list[0], - attention_mask_1=attention_mask_list[1], - ) - projected_prompt_embeds = projection_output.hidden_states - projected_attention_mask = projection_output.attention_mask - - generated_prompt_embeds = self.generate_language_model( - projected_prompt_embeds, - attention_mask=projected_attention_mask, - max_new_tokens=max_new_tokens, - ) - - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) - attention_mask = ( - attention_mask.to(device=device) - if attention_mask is not None - else torch.ones(prompt_embeds.shape[:2], dtype=torch.long, device=device) - ) - generated_prompt_embeds = generated_prompt_embeds.to(dtype=self.language_model.dtype, device=device) - - bs_embed, seq_len, hidden_size = prompt_embeds.shape - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_waveforms_per_prompt, 1) - prompt_embeds = prompt_embeds.view(bs_embed * num_waveforms_per_prompt, seq_len, hidden_size) - - # duplicate attention mask for each generation per prompt - attention_mask = attention_mask.repeat(1, num_waveforms_per_prompt) - attention_mask = attention_mask.view(bs_embed * num_waveforms_per_prompt, seq_len) - - bs_embed, seq_len, hidden_size = generated_prompt_embeds.shape - # duplicate generated embeddings for each generation per prompt, using mps friendly method - generated_prompt_embeds = generated_prompt_embeds.repeat(1, num_waveforms_per_prompt, 1) - generated_prompt_embeds = generated_prompt_embeds.view( - bs_embed * num_waveforms_per_prompt, seq_len, hidden_size - ) - - # 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 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 - - negative_prompt_embeds_list = [] - negative_attention_mask_list = [] - max_length = prompt_embeds.shape[1] - for tokenizer, text_encoder in zip(tokenizers, text_encoders): - uncond_input = tokenizer( - uncond_tokens, - padding="max_length", - max_length=tokenizer.model_max_length - if isinstance(tokenizer, (RobertaTokenizer, RobertaTokenizerFast, VitsTokenizer)) - else max_length, - truncation=True, - return_tensors="pt", - ) - - uncond_input_ids = uncond_input.input_ids.to(device) - negative_attention_mask = uncond_input.attention_mask.to(device) - - if text_encoder.config.model_type == "clap": - negative_prompt_embeds = text_encoder.get_text_features( - uncond_input_ids, - attention_mask=negative_attention_mask, - ) - # 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 - negative_attention_mask = negative_attention_mask.new_ones((batch_size, 1)) - elif is_vits_text_encoder: - negative_prompt_embeds = torch.zeros( - batch_size, - tokenizer.model_max_length, - text_encoder.config.hidden_size, - ).to(dtype=self.text_encoder_2.dtype, device=device) - negative_attention_mask = torch.zeros(batch_size, tokenizer.model_max_length).to( - dtype=self.text_encoder_2.dtype, device=device - ) - else: - negative_prompt_embeds = text_encoder( - uncond_input_ids, - attention_mask=negative_attention_mask, - ) - negative_prompt_embeds = negative_prompt_embeds[0] - - negative_prompt_embeds_list.append(negative_prompt_embeds) - negative_attention_mask_list.append(negative_attention_mask) - - projection_output = self.projection_model( - hidden_states=negative_prompt_embeds_list[0], - hidden_states_1=negative_prompt_embeds_list[1], - attention_mask=negative_attention_mask_list[0], - attention_mask_1=negative_attention_mask_list[1], - ) - negative_projected_prompt_embeds = projection_output.hidden_states - negative_projected_attention_mask = projection_output.attention_mask - - negative_generated_prompt_embeds = self.generate_language_model( - negative_projected_prompt_embeds, - attention_mask=negative_projected_attention_mask, - max_new_tokens=max_new_tokens, - ) - - if do_classifier_free_guidance: - seq_len = negative_prompt_embeds.shape[1] - - negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) - negative_attention_mask = ( - negative_attention_mask.to(device=device) - if negative_attention_mask is not None - else torch.ones(negative_prompt_embeds.shape[:2], dtype=torch.long, device=device) - ) - negative_generated_prompt_embeds = negative_generated_prompt_embeds.to( - dtype=self.language_model.dtype, device=device - ) - - # duplicate unconditional embeddings for each generation per prompt, using mps friendly method - negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_waveforms_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_waveforms_per_prompt, seq_len, -1) - - # duplicate unconditional attention mask for each generation per prompt - negative_attention_mask = negative_attention_mask.repeat(1, num_waveforms_per_prompt) - negative_attention_mask = negative_attention_mask.view(batch_size * num_waveforms_per_prompt, seq_len) - - # duplicate unconditional generated embeddings for each generation per prompt - seq_len = negative_generated_prompt_embeds.shape[1] - negative_generated_prompt_embeds = negative_generated_prompt_embeds.repeat(1, num_waveforms_per_prompt, 1) - negative_generated_prompt_embeds = negative_generated_prompt_embeds.view( - batch_size * num_waveforms_per_prompt, seq_len, -1 - ) - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - attention_mask = torch.cat([negative_attention_mask, attention_mask]) - generated_prompt_embeds = torch.cat([negative_generated_prompt_embeds, generated_prompt_embeds]) - - return prompt_embeds, attention_mask, generated_prompt_embeds - - # Copied from diffusers.pipelines.audioldm.pipeline_audioldm.AudioLDMPipeline.mel_spectrogram_to_waveform - def mel_spectrogram_to_waveform(self, mel_spectrogram): - if mel_spectrogram.dim() == 4: - mel_spectrogram = mel_spectrogram.squeeze(1) - - waveform = self.vocoder(mel_spectrogram) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - waveform = waveform.cpu().float() - return waveform - - def score_waveforms(self, text, audio, num_waveforms_per_prompt, device, dtype): - if not is_librosa_available(): - logger.info( - "Automatic scoring of the generated audio waveforms against the input prompt text requires the " - "`librosa` package to resample the generated waveforms. Returning the audios in the order they were " - "generated. To enable automatic scoring, install `librosa` with: `pip install librosa`." - ) - return audio - inputs = self.tokenizer(text, return_tensors="pt", padding=True) - resampled_audio = librosa.resample( - audio.numpy(), orig_sr=self.vocoder.config.sampling_rate, target_sr=self.feature_extractor.sampling_rate - ) - inputs["input_features"] = self.feature_extractor( - list(resampled_audio), return_tensors="pt", sampling_rate=self.feature_extractor.sampling_rate - ).input_features.type(dtype) - inputs = inputs.to(device) - - # compute the audio-text similarity score using the CLAP model - logits_per_text = self.text_encoder(**inputs).logits_per_text - # sort by the highest matching generations per prompt - indices = torch.argsort(logits_per_text, dim=1, descending=True)[:, :num_waveforms_per_prompt] - audio = torch.index_select(audio, 0, indices.reshape(-1).cpu()) - return audio - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - audio_length_in_s, - vocoder_upsample_factor, - callback_steps, - transcription=None, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - generated_prompt_embeds=None, - negative_generated_prompt_embeds=None, - attention_mask=None, - negative_attention_mask=None, - ): - min_audio_length_in_s = vocoder_upsample_factor * self.vae_scale_factor - if audio_length_in_s < min_audio_length_in_s: - raise ValueError( - f"`audio_length_in_s` has to be a positive value greater than or equal to {min_audio_length_in_s}, but " - f"is {audio_length_in_s}." - ) - - if self.vocoder.config.model_in_dim % self.vae_scale_factor != 0: - raise ValueError( - f"The number of frequency bins in the vocoder's log-mel spectrogram has to be divisible by the " - f"VAE scale factor, but got {self.vocoder.config.model_in_dim} bins and a scale factor of " - f"{self.vae_scale_factor}." - ) - - if (callback_steps is None) or ( - 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 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 or generated_prompt_embeds is None): - raise ValueError( - "Provide either `prompt`, or `prompt_embeds` and `generated_prompt_embeds`. Cannot leave " - "`prompt` undefined without specifying both `prompt_embeds` and `generated_prompt_embeds`." - ) - 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." - ) - elif negative_prompt_embeds is not None and negative_generated_prompt_embeds is None: - raise ValueError( - "Cannot forward `negative_prompt_embeds` without `negative_generated_prompt_embeds`. Ensure that" - "both arguments are specified" - ) - - 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 attention_mask is not None and attention_mask.shape != prompt_embeds.shape[:2]: - raise ValueError( - "`attention_mask should have the same batch size and sequence length as `prompt_embeds`, but got:" - f"`attention_mask: {attention_mask.shape} != `prompt_embeds` {prompt_embeds.shape}" - ) - - if transcription is None: - if self.text_encoder_2.config.model_type == "vits": - raise ValueError("Cannot forward without transcription. Please make sure to have transcription") - elif transcription is not None and ( - not isinstance(transcription, str) and not isinstance(transcription, list) - ): - raise ValueError(f"`transcription` has to be of type `str` or `list` but is {type(transcription)}") - - if generated_prompt_embeds is not None and negative_generated_prompt_embeds is not None: - if generated_prompt_embeds.shape != negative_generated_prompt_embeds.shape: - raise ValueError( - "`generated_prompt_embeds` and `negative_generated_prompt_embeds` must have the same shape when " - f"passed directly, but got: `generated_prompt_embeds` {generated_prompt_embeds.shape} != " - f"`negative_generated_prompt_embeds` {negative_generated_prompt_embeds.shape}." - ) - if ( - negative_attention_mask is not None - and negative_attention_mask.shape != negative_prompt_embeds.shape[:2] - ): - raise ValueError( - "`attention_mask should have the same batch size and sequence length as `prompt_embeds`, but got:" - f"`attention_mask: {negative_attention_mask.shape} != `prompt_embeds` {negative_prompt_embeds.shape}" - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents with width->self.vocoder.config.model_in_dim - def prepare_latents(self, batch_size, num_channels_latents, height, dtype, device, generator, latents=None): - shape = ( - batch_size, - num_channels_latents, - int(height) // self.vae_scale_factor, - int(self.vocoder.config.model_in_dim) // 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @torch.no_grad() - @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, - 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, - 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, - 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", - ): - r""" - The call function to the pipeline for generation. - - Args: - 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*):\ - 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. - num_inference_steps (`int`, *optional*, defaults to 200): - The number of denoising steps. More denoising steps usually lead to a higher quality audio at the - expense of slower inference. - 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*): - 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): - The number of waveforms to generate per prompt. If `num_waveforms_per_prompt > 1`, then automatic - scoring is performed between the generated outputs and the text prompt. This scoring ranks the - generated waveforms based on their cosine similarity with the text input in the joint text-audio - embedding space. - 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*): - 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 spectrogram - 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. - 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. - generated_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated text embeddings from the GPT2 language model. 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_generated_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated negative text embeddings from the GPT2 language model. Can be used to easily tweak text - inputs, *e.g.* prompt weighting. If not provided, negative_prompt_embeds will be computed from - `negative_prompt` input argument. - attention_mask (`torch.LongTensor`, *optional*): - Pre-computed attention mask to be applied to the `prompt_embeds`. If not provided, attention mask will - be computed from `prompt` input argument. - negative_attention_mask (`torch.LongTensor`, *optional*): - Pre-computed attention mask to be applied to the `negative_prompt_embeds`. If not provided, attention - mask will be computed from `negative_prompt` input argument. - max_new_tokens (`int`, *optional*, defaults to None): - Number of new tokens to generate with the GPT2 language model. If not provided, number of tokens will - be taken from the config of the model. - 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 calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - output_type (`str`, *optional*, defaults to `"np"`): - The output format of the generated audio. Choose between `"np"` to return a NumPy `np.ndarray` or - `"pt"` to return a PyTorch `torch.Tensor` object. Set to `"latent"` to return the latent diffusion - model (LDM) output. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned, - otherwise a `tuple` is returned where the first element is a list with the generated audio. - """ - # 0. Convert audio input length from seconds to spectrogram height - vocoder_upsample_factor = np.prod(self.vocoder.config.upsample_rates) / self.vocoder.config.sampling_rate - - if audio_length_in_s is None: - audio_length_in_s = self.unet.config.sample_size * self.vae_scale_factor * vocoder_upsample_factor - - height = int(audio_length_in_s / vocoder_upsample_factor) - - original_waveform_length = int(audio_length_in_s * self.vocoder.config.sampling_rate) - if height % self.vae_scale_factor != 0: - height = int(np.ceil(height / self.vae_scale_factor)) * self.vae_scale_factor - logger.info( - f"Audio length in seconds {audio_length_in_s} is increased to {height * vocoder_upsample_factor} " - f"so that it can be handled by the model. It will be cut to {audio_length_in_s} after the " - f"denoising process." - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - audio_length_in_s, - vocoder_upsample_factor, - callback_steps, - transcription, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - generated_prompt_embeds, - negative_generated_prompt_embeds, - attention_mask, - negative_attention_mask, - ) - - # 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds, attention_mask, generated_prompt_embeds = self.encode_prompt( - prompt, - device, - num_waveforms_per_prompt, - do_classifier_free_guidance, - transcription, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - generated_prompt_embeds=generated_prompt_embeds, - negative_generated_prompt_embeds=negative_generated_prompt_embeds, - attention_mask=attention_mask, - negative_attention_mask=negative_attention_mask, - max_new_tokens=max_new_tokens, - ) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_waveforms_per_prompt, - num_channels_latents, - height, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=generated_prompt_embeds, - encoder_hidden_states_1=prompt_embeds, - encoder_attention_mask_1=attention_mask, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - self.maybe_free_model_hooks() - - # 8. Post-processing - if not output_type == "latent": - latents = 1 / self.vae.config.scaling_factor * latents - mel_spectrogram = self.vae.decode(latents).sample - else: - return AudioPipelineOutput(audios=latents) - - audio = self.mel_spectrogram_to_waveform(mel_spectrogram) - - audio = audio[:, :original_waveform_length] - - # 9. Automatic scoring - if num_waveforms_per_prompt > 1 and prompt is not None: - audio = self.score_waveforms( - text=prompt, - audio=audio, - num_waveforms_per_prompt=num_waveforms_per_prompt, - device=device, - dtype=prompt_embeds.dtype, - ) - - if output_type == "np": - audio = audio.numpy() - - if not return_dict: - return (audio,) - - return AudioPipelineOutput(audios=audio) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/aura_flow/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/aura_flow/__init__.py deleted file mode 100644 index e1917baa61e28abf0970080c90b5864579cbe8f9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/aura_flow/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -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_aura_flow"] = ["AuraFlowPipeline"] - -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_aura_flow import AuraFlowPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/aura_flow/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/aura_flow/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 74bbe95be46e200e4fc715b4d98a6477d8e0edb7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/aura_flow/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/aura_flow/__pycache__/pipeline_aura_flow.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/aura_flow/__pycache__/pipeline_aura_flow.cpython-310.pyc deleted file mode 100644 index c71091e50bef8d9e763b2aeb162a8aa234b1c45c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/aura_flow/__pycache__/pipeline_aura_flow.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/aura_flow/pipeline_aura_flow.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/aura_flow/pipeline_aura_flow.py deleted file mode 100644 index 7ff9925c452d74ef3e1980f88a7324083d95dde5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/aura_flow/pipeline_aura_flow.py +++ /dev/null @@ -1,677 +0,0 @@ -# Copyright 2025 AuraFlow Authors 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, Tuple, Union - -import torch -from transformers import T5Tokenizer, UMT5EncoderModel - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import VaeImageProcessor -from ...loaders import AuraFlowLoraLoaderMixin -from ...models import AuraFlowTransformer2DModel, AutoencoderKL -from ...models.attention_processor import AttnProcessor2_0, FusedAttnProcessor2_0, XFormersAttnProcessor -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, ImagePipelineOutput - - -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 AuraFlowPipeline - - >>> pipe = AuraFlowPipeline.from_pretrained("fal/AuraFlow", torch_dtype=torch.float16) - >>> pipe = pipe.to("cuda") - >>> prompt = "A cat holding a sign that says hello world" - >>> image = pipe(prompt).images[0] - >>> image.save("aura_flow.png") - ``` -""" - - -# 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 AuraFlowPipeline(DiffusionPipeline, AuraFlowLoraLoaderMixin): - r""" - Args: - tokenizer (`T5TokenizerFast`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - text_encoder ([`T5EncoderModel`]): - Frozen text-encoder. AuraFlow uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the - [EleutherAI/pile-t5-xl](https://huggingface.co/EleutherAI/pile-t5-xl) variant. - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - transformer ([`AuraFlowTransformer2DModel`]): - Conditional Transformer (MMDiT and DiT) architecture to denoise the encoded image latents. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - """ - - _optional_components = [] - model_cpu_offload_seq = "text_encoder->transformer->vae" - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - ] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: UMT5EncoderModel, - vae: AutoencoderKL, - transformer: AuraFlowTransformer2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_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) - - def check_inputs( - self, - prompt, - height, - width, - negative_prompt, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - callback_on_step_end_tensor_inputs=None, - ): - if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0: - raise ValueError( - f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} 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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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("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}." - ) - - def encode_prompt( - self, - prompt: Union[str, List[str]], - negative_prompt: Union[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, - max_sequence_length: int = 256, - lora_scale: Optional[float] = 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 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_images_per_prompt (`int`, *optional*, defaults to 1): - number of images that should be generated per prompt - device: (`torch.device`, *optional*): - 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. - prompt_attention_mask (`torch.Tensor`, *optional*): - Pre-generated attention mask for text embeddings. - negative_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated negative text embeddings. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Pre-generated attention mask for negative text embeddings. - max_sequence_length (`int`, defaults to 256): Maximum sequence length to use for the prompt. - lora_scale (`float`, *optional*): - A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. - """ - # 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, AuraFlowLoraLoaderMixin): - 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) - - if device is None: - device = self._execution_device - 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] - - max_length = max_sequence_length - if prompt_embeds is None: - text_inputs = self.tokenizer( - prompt, - truncation=True, - max_length=max_length, - padding="max_length", - 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[:, max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because T5 can only handle sequences up to" - f" {max_length} tokens: {removed_text}" - ) - - text_inputs = {k: v.to(device) for k, v in text_inputs.items()} - prompt_embeds = self.text_encoder(**text_inputs)[0] - prompt_attention_mask = text_inputs["attention_mask"].unsqueeze(-1).expand(prompt_embeds.shape) - prompt_embeds = prompt_embeds * prompt_attention_mask - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - elif self.transformer is not None: - dtype = self.transformer.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - - bs_embed, 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(bs_embed * num_images_per_prompt, seq_len, -1) - prompt_attention_mask = prompt_attention_mask.reshape(bs_embed, -1) - prompt_attention_mask = prompt_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: - negative_prompt = negative_prompt or "" - uncond_tokens = [negative_prompt] * batch_size if isinstance(negative_prompt, str) else negative_prompt - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - truncation=True, - max_length=max_length, - padding="max_length", - return_tensors="pt", - ) - uncond_input = {k: v.to(device) for k, v in uncond_input.items()} - negative_prompt_embeds = self.text_encoder(**uncond_input)[0] - negative_prompt_attention_mask = ( - uncond_input["attention_mask"].unsqueeze(-1).expand(negative_prompt_embeds.shape) - ) - negative_prompt_embeds = negative_prompt_embeds * negative_prompt_attention_mask - - 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=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) - - negative_prompt_attention_mask = negative_prompt_attention_mask.reshape(bs_embed, -1) - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1) - else: - negative_prompt_embeds = None - negative_prompt_attention_mask = None - - if self.text_encoder is not None: - if isinstance(self, AuraFlowLoraLoaderMixin) 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, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_latents - 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=device, dtype=dtype) - - 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 - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - FusedAttnProcessor2_0, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def attention_kwargs(self): - return self._attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @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, - num_inference_steps: int = 50, - 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, - max_sequence_length: int = 256, - output_type: Optional[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"], - ) -> Union[ImagePipelineOutput, Tuple]: - 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 - less than `1`). - height (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor): - The height in pixels of the generated image. This is set to 1024 by default for best results. - width (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor): - The width in pixels of the generated image. This is set to 1024 by default for 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 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): - 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.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. - prompt_attention_mask (`torch.Tensor`, *optional*): - Pre-generated attention mask for text embeddings. - 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. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Pre-generated attention mask for negative text embeddings. - 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. - 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.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is returned - where the first element is a list with the generated images. - """ - # 1. Check inputs. Raise error if not correct - height = height or self.transformer.config.sample_size * self.vae_scale_factor - width = width or self.transformer.config.sample_size * self.vae_scale_factor - - self.check_inputs( - prompt, - height, - width, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - - # 2. Determine batch size. - 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.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None - - # 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - ( - 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=do_classifier_free_guidance, - num_images_per_prompt=num_images_per_prompt, - device=device, - 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, - lora_scale=lora_scale, - ) - if do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - - # 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) - - # 5. Prepare latents. - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - # aura use timestep value between 0 and 1, with t=1 as noise and t=0 as the image - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = torch.tensor([t / 1000]).expand(latent_model_input.shape[0]) - timestep = timestep.to(latents.device, dtype=latents.dtype) - - # predict noise model_output - noise_pred = self.transformer( - latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=timestep, - return_dict=False, - attention_kwargs=self.attention_kwargs, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - 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, 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() - - if output_type == "latent": - image = latents - else: - # 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] - 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 ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/auto_pipeline.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/auto_pipeline.py deleted file mode 100644 index b1a7ffaaea9cef1268e9e793bc702ef2aad68022..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/auto_pipeline.py +++ /dev/null @@ -1,1175 +0,0 @@ -# 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. - -from collections import OrderedDict - -from huggingface_hub.utils import validate_hf_hub_args - -from ..configuration_utils import ConfigMixin -from ..models.controlnets import ControlNetUnionModel -from ..utils import is_sentencepiece_available -from .aura_flow import AuraFlowPipeline -from .chroma import ChromaPipeline -from .cogview3 import CogView3PlusPipeline -from .cogview4 import CogView4ControlPipeline, CogView4Pipeline -from .controlnet import ( - StableDiffusionControlNetImg2ImgPipeline, - StableDiffusionControlNetInpaintPipeline, - StableDiffusionControlNetPipeline, - StableDiffusionXLControlNetImg2ImgPipeline, - StableDiffusionXLControlNetInpaintPipeline, - StableDiffusionXLControlNetPipeline, - StableDiffusionXLControlNetUnionImg2ImgPipeline, - StableDiffusionXLControlNetUnionInpaintPipeline, - StableDiffusionXLControlNetUnionPipeline, -) -from .controlnet_sd3 import ( - StableDiffusion3ControlNetInpaintingPipeline, - StableDiffusion3ControlNetPipeline, -) -from .deepfloyd_if import IFImg2ImgPipeline, IFInpaintingPipeline, IFPipeline -from .flux import ( - FluxControlImg2ImgPipeline, - FluxControlInpaintPipeline, - FluxControlNetImg2ImgPipeline, - FluxControlNetInpaintPipeline, - FluxControlNetPipeline, - FluxControlPipeline, - FluxImg2ImgPipeline, - FluxInpaintPipeline, - FluxPipeline, -) -from .hunyuandit import HunyuanDiTPipeline -from .kandinsky import ( - KandinskyCombinedPipeline, - KandinskyImg2ImgCombinedPipeline, - KandinskyImg2ImgPipeline, - KandinskyInpaintCombinedPipeline, - KandinskyInpaintPipeline, - KandinskyPipeline, -) -from .kandinsky2_2 import ( - KandinskyV22CombinedPipeline, - KandinskyV22Img2ImgCombinedPipeline, - KandinskyV22Img2ImgPipeline, - KandinskyV22InpaintCombinedPipeline, - KandinskyV22InpaintPipeline, - KandinskyV22Pipeline, -) -from .kandinsky3 import Kandinsky3Img2ImgPipeline, Kandinsky3Pipeline -from .latent_consistency_models import LatentConsistencyModelImg2ImgPipeline, LatentConsistencyModelPipeline -from .lumina import LuminaPipeline -from .lumina2 import Lumina2Pipeline -from .pag import ( - HunyuanDiTPAGPipeline, - PixArtSigmaPAGPipeline, - SanaPAGPipeline, - StableDiffusion3PAGImg2ImgPipeline, - StableDiffusion3PAGPipeline, - StableDiffusionControlNetPAGInpaintPipeline, - StableDiffusionControlNetPAGPipeline, - StableDiffusionPAGImg2ImgPipeline, - StableDiffusionPAGInpaintPipeline, - StableDiffusionPAGPipeline, - StableDiffusionXLControlNetPAGImg2ImgPipeline, - StableDiffusionXLControlNetPAGPipeline, - StableDiffusionXLPAGImg2ImgPipeline, - StableDiffusionXLPAGInpaintPipeline, - StableDiffusionXLPAGPipeline, -) -from .pixart_alpha import PixArtAlphaPipeline, PixArtSigmaPipeline -from .sana import SanaPipeline -from .stable_cascade import StableCascadeCombinedPipeline, StableCascadeDecoderPipeline -from .stable_diffusion import ( - StableDiffusionImg2ImgPipeline, - StableDiffusionInpaintPipeline, - StableDiffusionPipeline, -) -from .stable_diffusion_3 import ( - StableDiffusion3Img2ImgPipeline, - StableDiffusion3InpaintPipeline, - StableDiffusion3Pipeline, -) -from .stable_diffusion_xl import ( - StableDiffusionXLImg2ImgPipeline, - StableDiffusionXLInpaintPipeline, - StableDiffusionXLPipeline, -) -from .wuerstchen import WuerstchenCombinedPipeline, WuerstchenDecoderPipeline - - -AUTO_TEXT2IMAGE_PIPELINES_MAPPING = OrderedDict( - [ - ("stable-diffusion", StableDiffusionPipeline), - ("stable-diffusion-xl", StableDiffusionXLPipeline), - ("stable-diffusion-3", StableDiffusion3Pipeline), - ("stable-diffusion-3-pag", StableDiffusion3PAGPipeline), - ("if", IFPipeline), - ("hunyuan", HunyuanDiTPipeline), - ("hunyuan-pag", HunyuanDiTPAGPipeline), - ("kandinsky", KandinskyCombinedPipeline), - ("kandinsky22", KandinskyV22CombinedPipeline), - ("kandinsky3", Kandinsky3Pipeline), - ("stable-diffusion-controlnet", StableDiffusionControlNetPipeline), - ("stable-diffusion-xl-controlnet", StableDiffusionXLControlNetPipeline), - ("stable-diffusion-xl-controlnet-union", StableDiffusionXLControlNetUnionPipeline), - ("stable-diffusion-3-controlnet", StableDiffusion3ControlNetPipeline), - ("wuerstchen", WuerstchenCombinedPipeline), - ("cascade", StableCascadeCombinedPipeline), - ("lcm", LatentConsistencyModelPipeline), - ("pixart-alpha", PixArtAlphaPipeline), - ("pixart-sigma", PixArtSigmaPipeline), - ("sana", SanaPipeline), - ("sana-pag", SanaPAGPipeline), - ("stable-diffusion-pag", StableDiffusionPAGPipeline), - ("stable-diffusion-controlnet-pag", StableDiffusionControlNetPAGPipeline), - ("stable-diffusion-xl-pag", StableDiffusionXLPAGPipeline), - ("stable-diffusion-xl-controlnet-pag", StableDiffusionXLControlNetPAGPipeline), - ("pixart-sigma-pag", PixArtSigmaPAGPipeline), - ("auraflow", AuraFlowPipeline), - ("flux", FluxPipeline), - ("flux-control", FluxControlPipeline), - ("flux-controlnet", FluxControlNetPipeline), - ("lumina", LuminaPipeline), - ("lumina2", Lumina2Pipeline), - ("chroma", ChromaPipeline), - ("cogview3", CogView3PlusPipeline), - ("cogview4", CogView4Pipeline), - ("cogview4-control", CogView4ControlPipeline), - ] -) - -AUTO_IMAGE2IMAGE_PIPELINES_MAPPING = OrderedDict( - [ - ("stable-diffusion", StableDiffusionImg2ImgPipeline), - ("stable-diffusion-xl", StableDiffusionXLImg2ImgPipeline), - ("stable-diffusion-3", StableDiffusion3Img2ImgPipeline), - ("stable-diffusion-3-pag", StableDiffusion3PAGImg2ImgPipeline), - ("if", IFImg2ImgPipeline), - ("kandinsky", KandinskyImg2ImgCombinedPipeline), - ("kandinsky22", KandinskyV22Img2ImgCombinedPipeline), - ("kandinsky3", Kandinsky3Img2ImgPipeline), - ("stable-diffusion-controlnet", StableDiffusionControlNetImg2ImgPipeline), - ("stable-diffusion-pag", StableDiffusionPAGImg2ImgPipeline), - ("stable-diffusion-xl-controlnet", StableDiffusionXLControlNetImg2ImgPipeline), - ("stable-diffusion-xl-controlnet-union", StableDiffusionXLControlNetUnionImg2ImgPipeline), - ("stable-diffusion-xl-pag", StableDiffusionXLPAGImg2ImgPipeline), - ("stable-diffusion-xl-controlnet-pag", StableDiffusionXLControlNetPAGImg2ImgPipeline), - ("lcm", LatentConsistencyModelImg2ImgPipeline), - ("flux", FluxImg2ImgPipeline), - ("flux-controlnet", FluxControlNetImg2ImgPipeline), - ("flux-control", FluxControlImg2ImgPipeline), - ] -) - -AUTO_INPAINT_PIPELINES_MAPPING = OrderedDict( - [ - ("stable-diffusion", StableDiffusionInpaintPipeline), - ("stable-diffusion-xl", StableDiffusionXLInpaintPipeline), - ("stable-diffusion-3", StableDiffusion3InpaintPipeline), - ("if", IFInpaintingPipeline), - ("kandinsky", KandinskyInpaintCombinedPipeline), - ("kandinsky22", KandinskyV22InpaintCombinedPipeline), - ("stable-diffusion-controlnet", StableDiffusionControlNetInpaintPipeline), - ("stable-diffusion-controlnet-pag", StableDiffusionControlNetPAGInpaintPipeline), - ("stable-diffusion-xl-controlnet", StableDiffusionXLControlNetInpaintPipeline), - ("stable-diffusion-xl-controlnet-union", StableDiffusionXLControlNetUnionInpaintPipeline), - ("stable-diffusion-3-controlnet", StableDiffusion3ControlNetInpaintingPipeline), - ("stable-diffusion-xl-pag", StableDiffusionXLPAGInpaintPipeline), - ("flux", FluxInpaintPipeline), - ("flux-controlnet", FluxControlNetInpaintPipeline), - ("flux-control", FluxControlInpaintPipeline), - ("stable-diffusion-pag", StableDiffusionPAGInpaintPipeline), - ] -) - -_AUTO_TEXT2IMAGE_DECODER_PIPELINES_MAPPING = OrderedDict( - [ - ("kandinsky", KandinskyPipeline), - ("kandinsky22", KandinskyV22Pipeline), - ("wuerstchen", WuerstchenDecoderPipeline), - ("cascade", StableCascadeDecoderPipeline), - ] -) -_AUTO_IMAGE2IMAGE_DECODER_PIPELINES_MAPPING = OrderedDict( - [ - ("kandinsky", KandinskyImg2ImgPipeline), - ("kandinsky22", KandinskyV22Img2ImgPipeline), - ] -) -_AUTO_INPAINT_DECODER_PIPELINES_MAPPING = OrderedDict( - [ - ("kandinsky", KandinskyInpaintPipeline), - ("kandinsky22", KandinskyV22InpaintPipeline), - ] -) - -if is_sentencepiece_available(): - from .kolors import KolorsImg2ImgPipeline, KolorsPipeline - from .pag import KolorsPAGPipeline - - AUTO_TEXT2IMAGE_PIPELINES_MAPPING["kolors"] = KolorsPipeline - AUTO_TEXT2IMAGE_PIPELINES_MAPPING["kolors-pag"] = KolorsPAGPipeline - AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["kolors"] = KolorsImg2ImgPipeline - -SUPPORTED_TASKS_MAPPINGS = [ - AUTO_TEXT2IMAGE_PIPELINES_MAPPING, - AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, - AUTO_INPAINT_PIPELINES_MAPPING, - _AUTO_TEXT2IMAGE_DECODER_PIPELINES_MAPPING, - _AUTO_IMAGE2IMAGE_DECODER_PIPELINES_MAPPING, - _AUTO_INPAINT_DECODER_PIPELINES_MAPPING, -] - - -def _get_connected_pipeline(pipeline_cls): - # for now connected pipelines can only be loaded from decoder pipelines, such as kandinsky-community/kandinsky-2-2-decoder - if pipeline_cls in _AUTO_TEXT2IMAGE_DECODER_PIPELINES_MAPPING.values(): - return _get_task_class( - AUTO_TEXT2IMAGE_PIPELINES_MAPPING, pipeline_cls.__name__, throw_error_if_not_exist=False - ) - if pipeline_cls in _AUTO_IMAGE2IMAGE_DECODER_PIPELINES_MAPPING.values(): - return _get_task_class( - AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, pipeline_cls.__name__, throw_error_if_not_exist=False - ) - if pipeline_cls in _AUTO_INPAINT_DECODER_PIPELINES_MAPPING.values(): - return _get_task_class(AUTO_INPAINT_PIPELINES_MAPPING, pipeline_cls.__name__, throw_error_if_not_exist=False) - - -def _get_task_class(mapping, pipeline_class_name, throw_error_if_not_exist: bool = True): - def get_model(pipeline_class_name): - for task_mapping in SUPPORTED_TASKS_MAPPINGS: - for model_name, pipeline in task_mapping.items(): - if pipeline.__name__ == pipeline_class_name: - return model_name - - model_name = get_model(pipeline_class_name) - - if model_name is not None: - task_class = mapping.get(model_name, None) - if task_class is not None: - return task_class - - if throw_error_if_not_exist: - raise ValueError(f"AutoPipeline can't find a pipeline linked to {pipeline_class_name} for {model_name}") - - -class AutoPipelineForText2Image(ConfigMixin): - r""" - - [`AutoPipelineForText2Image`] is a generic pipeline class that instantiates a text-to-image pipeline class. The - specific underlying pipeline class is automatically selected from either the - [`~AutoPipelineForText2Image.from_pretrained`] or [`~AutoPipelineForText2Image.from_pipe`] methods. - - This class cannot be instantiated using `__init__()` (throws an error). - - Class attributes: - - - **config_name** (`str`) -- The configuration filename that stores the class and module names of all the - diffusion pipeline's components. - - """ - - config_name = "model_index.json" - - 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"`{self.__class__.__name__}.from_pipe(pipeline)` methods." - ) - - @classmethod - @validate_hf_hub_args - def from_pretrained(cls, pretrained_model_or_path, **kwargs): - r""" - Instantiates a text-to-image Pytorch diffusion pipeline from pretrained pipeline weight. - - The from_pretrained() method takes care of returning the correct pipeline class instance by: - 1. Detect the pipeline class of the pretrained_model_or_path based on the _class_name property of its - config object - 2. Find the text-to-image pipeline linked to the pipeline class using pattern matching on pipeline class - name. - - If a `controlnet` argument is passed, it will instantiate a [`StableDiffusionControlNetPipeline`] object. - - The pipeline is set in evaluation mode (`model.eval()`) by default. - - If you get the error message below, you need to finetune the weights for your downstream task: - - ``` - Some weights of UNet2DConditionModel were not initialized from the model checkpoint at stable-diffusion-v1-5/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: - - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated - You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. - ``` - - Parameters: - pretrained_model_or_path (`str` or `os.PathLike`, *optional*): - Can be either: - - - A string, the *repo id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained pipeline - hosted on the Hub. - - A path to a *directory* (for example `./my_pipeline_directory/`) containing pipeline weights - saved using - [`~DiffusionPipeline.save_pretrained`]. - torch_dtype (`torch.dtype`, *optional*): - Override the default `torch.dtype` and load the model with another dtype. - 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*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - - 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`): - Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - custom_revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, or a commit id similar to - `revision` when loading a custom pipeline from the Hub. It can be a 🤗 Diffusers version when loading a - custom pipeline from GitHub, otherwise it defaults to `"main"` when loading from the Hub. - mirror (`str`, *optional*): - 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*): - 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. - - Set `device_map="auto"` to have 🤗 Accelerate automatically compute the most optimized `device_map`. For - more information about each option see [designing a device - map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). - max_memory (`Dict`, *optional*): - A dictionary device identifier for the maximum memory. Will default to the maximum memory available for - each GPU and the available CPU RAM if unset. - offload_folder (`str` or `os.PathLike`, *optional*): - The path to offload weights if device_map contains the value `"disk"`. - offload_state_dict (`bool`, *optional*): - If `True`, temporarily offloads the CPU state dict to the hard drive to avoid running out of CPU RAM if - the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to `True` - when there is some disk offload. - low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): - Speed up model loading only loading the pretrained weights and not initializing the weights. This also - tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. - Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this - argument to `True` will raise an error. - use_safetensors (`bool`, *optional*, defaults to `None`): - If set to `None`, the safetensors weights are downloaded if they're available **and** if the - safetensors library is installed. If set to `True`, the model is forcibly loaded from safetensors - weights. If set to `False`, safetensors weights are not loaded. - kwargs (remaining dictionary of keyword arguments, *optional*): - Can be used to overwrite load and saveable variables (the pipeline components of the specific pipeline - class). The overwritten components are passed directly to the pipelines `__init__` method. See example - below for more information. - variant (`str`, *optional*): - Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. - - - - To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with - `huggingface-cli login`. - - - - Examples: - - ```py - >>> from diffusers import AutoPipelineForText2Image - - >>> pipeline = AutoPipelineForText2Image.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") - >>> image = pipeline(prompt).images[0] - ``` - """ - cache_dir = kwargs.pop("cache_dir", None) - force_download = kwargs.pop("force_download", False) - proxies = kwargs.pop("proxies", None) - token = kwargs.pop("token", None) - local_files_only = kwargs.pop("local_files_only", False) - revision = kwargs.pop("revision", None) - - load_config_kwargs = { - "cache_dir": cache_dir, - "force_download": force_download, - "proxies": proxies, - "token": token, - "local_files_only": local_files_only, - "revision": revision, - } - - config = cls.load_config(pretrained_model_or_path, **load_config_kwargs) - orig_class_name = config["_class_name"] - if "ControlPipeline" in orig_class_name: - to_replace = "ControlPipeline" - else: - to_replace = "Pipeline" - - if "controlnet" in kwargs: - if isinstance(kwargs["controlnet"], ControlNetUnionModel): - orig_class_name = config["_class_name"].replace(to_replace, "ControlNetUnionPipeline") - else: - orig_class_name = config["_class_name"].replace(to_replace, "ControlNetPipeline") - if "enable_pag" in kwargs: - enable_pag = kwargs.pop("enable_pag") - if enable_pag: - orig_class_name = orig_class_name.replace(to_replace, "PAGPipeline") - - text_2_image_cls = _get_task_class(AUTO_TEXT2IMAGE_PIPELINES_MAPPING, orig_class_name) - - kwargs = {**load_config_kwargs, **kwargs} - return text_2_image_cls.from_pretrained(pretrained_model_or_path, **kwargs) - - @classmethod - def from_pipe(cls, pipeline, **kwargs): - r""" - Instantiates a text-to-image Pytorch diffusion pipeline from another instantiated diffusion pipeline class. - - The from_pipe() method takes care of returning the correct pipeline class instance by finding the text-to-image - pipeline linked to the pipeline class using pattern matching on pipeline class name. - - All the modules the pipeline contains will be used to initialize the new pipeline without reallocating - additional memory. - - The pipeline is set in evaluation mode (`model.eval()`) by default. - - Parameters: - pipeline (`DiffusionPipeline`): - an instantiated `DiffusionPipeline` object - - ```py - >>> from diffusers import AutoPipelineForText2Image, AutoPipelineForImage2Image - - >>> pipe_i2i = AutoPipelineForImage2Image.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", requires_safety_checker=False - ... ) - - >>> pipe_t2i = AutoPipelineForText2Image.from_pipe(pipe_i2i) - >>> image = pipe_t2i(prompt).images[0] - ``` - """ - - original_config = dict(pipeline.config) - original_cls_name = pipeline.__class__.__name__ - - # derive the pipeline class to instantiate - text_2_image_cls = _get_task_class(AUTO_TEXT2IMAGE_PIPELINES_MAPPING, original_cls_name) - - if "controlnet" in kwargs: - if kwargs["controlnet"] is not None: - to_replace = "PAGPipeline" if "PAG" in text_2_image_cls.__name__ else "Pipeline" - text_2_image_cls = _get_task_class( - AUTO_TEXT2IMAGE_PIPELINES_MAPPING, - text_2_image_cls.__name__.replace("ControlNet", "").replace(to_replace, "ControlNet" + to_replace), - ) - else: - text_2_image_cls = _get_task_class( - AUTO_TEXT2IMAGE_PIPELINES_MAPPING, - text_2_image_cls.__name__.replace("ControlNet", ""), - ) - - if "enable_pag" in kwargs: - enable_pag = kwargs.pop("enable_pag") - if enable_pag: - text_2_image_cls = _get_task_class( - AUTO_TEXT2IMAGE_PIPELINES_MAPPING, - text_2_image_cls.__name__.replace("PAG", "").replace("Pipeline", "PAGPipeline"), - ) - else: - text_2_image_cls = _get_task_class( - AUTO_TEXT2IMAGE_PIPELINES_MAPPING, - text_2_image_cls.__name__.replace("PAG", ""), - ) - - # define expected module and optional kwargs given the pipeline signature - expected_modules, optional_kwargs = text_2_image_cls._get_signature_keys(text_2_image_cls) - - pretrained_model_name_or_path = original_config.pop("_name_or_path", None) - - # allow users pass modules in `kwargs` to override the original pipeline's components - passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs} - original_class_obj = { - k: pipeline.components[k] - for k, v in pipeline.components.items() - if k in expected_modules and k not in passed_class_obj - } - - # allow users pass optional kwargs to override the original pipelines config attribute - passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs} - original_pipe_kwargs = { - k: original_config[k] - for k, v in original_config.items() - if k in optional_kwargs and k not in passed_pipe_kwargs - } - - # config that were not expected by original pipeline is stored as private attribute - # we will pass them as optional arguments if they can be accepted by the pipeline - additional_pipe_kwargs = [ - k[1:] - for k in original_config.keys() - if k.startswith("_") and k[1:] in optional_kwargs and k[1:] not in passed_pipe_kwargs - ] - for k in additional_pipe_kwargs: - original_pipe_kwargs[k] = original_config.pop(f"_{k}") - - text_2_image_kwargs = {**passed_class_obj, **original_class_obj, **passed_pipe_kwargs, **original_pipe_kwargs} - - # store unused config as private attribute - unused_original_config = { - f"{'' if k.startswith('_') else '_'}{k}": original_config[k] - for k, v in original_config.items() - if k not in text_2_image_kwargs - } - - missing_modules = ( - set(expected_modules) - set(text_2_image_cls._optional_components) - set(text_2_image_kwargs.keys()) - ) - - if len(missing_modules) > 0: - raise ValueError( - f"Pipeline {text_2_image_cls} expected {expected_modules}, but only {set(list(passed_class_obj.keys()) + list(original_class_obj.keys()))} were passed" - ) - - model = text_2_image_cls(**text_2_image_kwargs) - model.register_to_config(_name_or_path=pretrained_model_name_or_path) - model.register_to_config(**unused_original_config) - - return model - - -class AutoPipelineForImage2Image(ConfigMixin): - r""" - - [`AutoPipelineForImage2Image`] is a generic pipeline class that instantiates an image-to-image pipeline class. The - specific underlying pipeline class is automatically selected from either the - [`~AutoPipelineForImage2Image.from_pretrained`] or [`~AutoPipelineForImage2Image.from_pipe`] methods. - - This class cannot be instantiated using `__init__()` (throws an error). - - Class attributes: - - - **config_name** (`str`) -- The configuration filename that stores the class and module names of all the - diffusion pipeline's components. - - """ - - config_name = "model_index.json" - - 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"`{self.__class__.__name__}.from_pipe(pipeline)` methods." - ) - - @classmethod - @validate_hf_hub_args - def from_pretrained(cls, pretrained_model_or_path, **kwargs): - r""" - Instantiates a image-to-image Pytorch diffusion pipeline from pretrained pipeline weight. - - The from_pretrained() method takes care of returning the correct pipeline class instance by: - 1. Detect the pipeline class of the pretrained_model_or_path based on the _class_name property of its - config object - 2. Find the image-to-image pipeline linked to the pipeline class using pattern matching on pipeline class - name. - - If a `controlnet` argument is passed, it will instantiate a [`StableDiffusionControlNetImg2ImgPipeline`] - object. - - The pipeline is set in evaluation mode (`model.eval()`) by default. - - If you get the error message below, you need to finetune the weights for your downstream task: - - ``` - Some weights of UNet2DConditionModel were not initialized from the model checkpoint at stable-diffusion-v1-5/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: - - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated - You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. - ``` - - Parameters: - pretrained_model_or_path (`str` or `os.PathLike`, *optional*): - Can be either: - - - A string, the *repo id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained pipeline - hosted on the Hub. - - A path to a *directory* (for example `./my_pipeline_directory/`) containing pipeline weights - saved using - [`~DiffusionPipeline.save_pretrained`]. - torch_dtype (`str` or `torch.dtype`, *optional*): - Override the default `torch.dtype` and load the model with another dtype. - 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*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - - 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`): - Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - custom_revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, or a commit id similar to - `revision` when loading a custom pipeline from the Hub. It can be a 🤗 Diffusers version when loading a - custom pipeline from GitHub, otherwise it defaults to `"main"` when loading from the Hub. - mirror (`str`, *optional*): - 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*): - 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. - - Set `device_map="auto"` to have 🤗 Accelerate automatically compute the most optimized `device_map`. For - more information about each option see [designing a device - map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). - max_memory (`Dict`, *optional*): - A dictionary device identifier for the maximum memory. Will default to the maximum memory available for - each GPU and the available CPU RAM if unset. - offload_folder (`str` or `os.PathLike`, *optional*): - The path to offload weights if device_map contains the value `"disk"`. - offload_state_dict (`bool`, *optional*): - If `True`, temporarily offloads the CPU state dict to the hard drive to avoid running out of CPU RAM if - the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to `True` - when there is some disk offload. - low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): - Speed up model loading only loading the pretrained weights and not initializing the weights. This also - tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. - Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this - argument to `True` will raise an error. - use_safetensors (`bool`, *optional*, defaults to `None`): - If set to `None`, the safetensors weights are downloaded if they're available **and** if the - safetensors library is installed. If set to `True`, the model is forcibly loaded from safetensors - weights. If set to `False`, safetensors weights are not loaded. - kwargs (remaining dictionary of keyword arguments, *optional*): - Can be used to overwrite load and saveable variables (the pipeline components of the specific pipeline - class). The overwritten components are passed directly to the pipelines `__init__` method. See example - below for more information. - variant (`str`, *optional*): - Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. - - - - To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with - `huggingface-cli login`. - - - - Examples: - - ```py - >>> from diffusers import AutoPipelineForImage2Image - - >>> pipeline = AutoPipelineForImage2Image.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") - >>> image = pipeline(prompt, image).images[0] - ``` - """ - cache_dir = kwargs.pop("cache_dir", None) - force_download = kwargs.pop("force_download", False) - proxies = kwargs.pop("proxies", None) - token = kwargs.pop("token", None) - local_files_only = kwargs.pop("local_files_only", False) - revision = kwargs.pop("revision", None) - - load_config_kwargs = { - "cache_dir": cache_dir, - "force_download": force_download, - "proxies": proxies, - "token": token, - "local_files_only": local_files_only, - "revision": revision, - } - - config = cls.load_config(pretrained_model_or_path, **load_config_kwargs) - orig_class_name = config["_class_name"] - - # the `orig_class_name` can be: - # `- *Pipeline` (for regular text-to-image checkpoint) - # - `*ControlPipeline` (for Flux tools specific checkpoint) - # `- *Img2ImgPipeline` (for refiner checkpoint) - if "Img2Img" in orig_class_name: - to_replace = "Img2ImgPipeline" - elif "ControlPipeline" in orig_class_name: - to_replace = "ControlPipeline" - else: - to_replace = "Pipeline" - - if "controlnet" in kwargs: - if isinstance(kwargs["controlnet"], ControlNetUnionModel): - orig_class_name = orig_class_name.replace(to_replace, "ControlNetUnion" + to_replace) - else: - orig_class_name = orig_class_name.replace(to_replace, "ControlNet" + to_replace) - if "enable_pag" in kwargs: - enable_pag = kwargs.pop("enable_pag") - if enable_pag: - orig_class_name = orig_class_name.replace(to_replace, "PAG" + to_replace) - - if to_replace == "ControlPipeline": - orig_class_name = orig_class_name.replace(to_replace, "ControlImg2ImgPipeline") - - image_2_image_cls = _get_task_class(AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, orig_class_name) - - kwargs = {**load_config_kwargs, **kwargs} - return image_2_image_cls.from_pretrained(pretrained_model_or_path, **kwargs) - - @classmethod - def from_pipe(cls, pipeline, **kwargs): - r""" - Instantiates a image-to-image Pytorch diffusion pipeline from another instantiated diffusion pipeline class. - - The from_pipe() method takes care of returning the correct pipeline class instance by finding the - image-to-image pipeline linked to the pipeline class using pattern matching on pipeline class name. - - All the modules the pipeline contains will be used to initialize the new pipeline without reallocating - additional memory. - - The pipeline is set in evaluation mode (`model.eval()`) by default. - - Parameters: - pipeline (`DiffusionPipeline`): - an instantiated `DiffusionPipeline` object - - Examples: - - ```py - >>> from diffusers import AutoPipelineForText2Image, AutoPipelineForImage2Image - - >>> pipe_t2i = AutoPipelineForText2Image.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", requires_safety_checker=False - ... ) - - >>> pipe_i2i = AutoPipelineForImage2Image.from_pipe(pipe_t2i) - >>> image = pipe_i2i(prompt, image).images[0] - ``` - """ - - original_config = dict(pipeline.config) - original_cls_name = pipeline.__class__.__name__ - - # derive the pipeline class to instantiate - image_2_image_cls = _get_task_class(AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, original_cls_name) - - if "controlnet" in kwargs: - if kwargs["controlnet"] is not None: - to_replace = "Img2ImgPipeline" - if "PAG" in image_2_image_cls.__name__: - to_replace = "PAG" + to_replace - image_2_image_cls = _get_task_class( - AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, - image_2_image_cls.__name__.replace("ControlNet", "").replace( - to_replace, "ControlNet" + to_replace - ), - ) - else: - image_2_image_cls = _get_task_class( - AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, - image_2_image_cls.__name__.replace("ControlNet", ""), - ) - - if "enable_pag" in kwargs: - enable_pag = kwargs.pop("enable_pag") - if enable_pag: - image_2_image_cls = _get_task_class( - AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, - image_2_image_cls.__name__.replace("PAG", "").replace("Img2ImgPipeline", "PAGImg2ImgPipeline"), - ) - else: - image_2_image_cls = _get_task_class( - AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, - image_2_image_cls.__name__.replace("PAG", ""), - ) - - # define expected module and optional kwargs given the pipeline signature - expected_modules, optional_kwargs = image_2_image_cls._get_signature_keys(image_2_image_cls) - - pretrained_model_name_or_path = original_config.pop("_name_or_path", None) - - # allow users pass modules in `kwargs` to override the original pipeline's components - passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs} - original_class_obj = { - k: pipeline.components[k] - for k, v in pipeline.components.items() - if k in expected_modules and k not in passed_class_obj - } - - # allow users pass optional kwargs to override the original pipelines config attribute - passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs} - original_pipe_kwargs = { - k: original_config[k] - for k, v in original_config.items() - if k in optional_kwargs and k not in passed_pipe_kwargs - } - - # config attribute that were not expected by original pipeline is stored as its private attribute - # we will pass them as optional arguments if they can be accepted by the pipeline - additional_pipe_kwargs = [ - k[1:] - for k in original_config.keys() - if k.startswith("_") and k[1:] in optional_kwargs and k[1:] not in passed_pipe_kwargs - ] - for k in additional_pipe_kwargs: - original_pipe_kwargs[k] = original_config.pop(f"_{k}") - - image_2_image_kwargs = {**passed_class_obj, **original_class_obj, **passed_pipe_kwargs, **original_pipe_kwargs} - - # store unused config as private attribute - unused_original_config = { - f"{'' if k.startswith('_') else '_'}{k}": original_config[k] - for k, v in original_config.items() - if k not in image_2_image_kwargs - } - - missing_modules = ( - set(expected_modules) - set(image_2_image_cls._optional_components) - set(image_2_image_kwargs.keys()) - ) - - if len(missing_modules) > 0: - raise ValueError( - f"Pipeline {image_2_image_cls} expected {expected_modules}, but only {set(list(passed_class_obj.keys()) + list(original_class_obj.keys()))} were passed" - ) - - model = image_2_image_cls(**image_2_image_kwargs) - model.register_to_config(_name_or_path=pretrained_model_name_or_path) - model.register_to_config(**unused_original_config) - - return model - - -class AutoPipelineForInpainting(ConfigMixin): - r""" - - [`AutoPipelineForInpainting`] is a generic pipeline class that instantiates an inpainting pipeline class. The - specific underlying pipeline class is automatically selected from either the - [`~AutoPipelineForInpainting.from_pretrained`] or [`~AutoPipelineForInpainting.from_pipe`] methods. - - This class cannot be instantiated using `__init__()` (throws an error). - - Class attributes: - - - **config_name** (`str`) -- The configuration filename that stores the class and module names of all the - diffusion pipeline's components. - - """ - - config_name = "model_index.json" - - 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"`{self.__class__.__name__}.from_pipe(pipeline)` methods." - ) - - @classmethod - @validate_hf_hub_args - def from_pretrained(cls, pretrained_model_or_path, **kwargs): - r""" - Instantiates a inpainting Pytorch diffusion pipeline from pretrained pipeline weight. - - The from_pretrained() method takes care of returning the correct pipeline class instance by: - 1. Detect the pipeline class of the pretrained_model_or_path based on the _class_name property of its - config object - 2. Find the inpainting pipeline linked to the pipeline class using pattern matching on pipeline class name. - - If a `controlnet` argument is passed, it will instantiate a [`StableDiffusionControlNetInpaintPipeline`] - object. - - The pipeline is set in evaluation mode (`model.eval()`) by default. - - If you get the error message below, you need to finetune the weights for your downstream task: - - ``` - Some weights of UNet2DConditionModel were not initialized from the model checkpoint at stable-diffusion-v1-5/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: - - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated - You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. - ``` - - Parameters: - pretrained_model_or_path (`str` or `os.PathLike`, *optional*): - Can be either: - - - A string, the *repo id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained pipeline - hosted on the Hub. - - A path to a *directory* (for example `./my_pipeline_directory/`) containing pipeline weights - saved using - [`~DiffusionPipeline.save_pretrained`]. - torch_dtype (`str` or `torch.dtype`, *optional*): - Override the default `torch.dtype` and load the model with another dtype. - 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*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - - 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`): - Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - custom_revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, or a commit id similar to - `revision` when loading a custom pipeline from the Hub. It can be a 🤗 Diffusers version when loading a - custom pipeline from GitHub, otherwise it defaults to `"main"` when loading from the Hub. - mirror (`str`, *optional*): - 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*): - 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. - - Set `device_map="auto"` to have 🤗 Accelerate automatically compute the most optimized `device_map`. For - more information about each option see [designing a device - map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). - max_memory (`Dict`, *optional*): - A dictionary device identifier for the maximum memory. Will default to the maximum memory available for - each GPU and the available CPU RAM if unset. - offload_folder (`str` or `os.PathLike`, *optional*): - The path to offload weights if device_map contains the value `"disk"`. - offload_state_dict (`bool`, *optional*): - If `True`, temporarily offloads the CPU state dict to the hard drive to avoid running out of CPU RAM if - the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to `True` - when there is some disk offload. - low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): - Speed up model loading only loading the pretrained weights and not initializing the weights. This also - tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. - Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this - argument to `True` will raise an error. - use_safetensors (`bool`, *optional*, defaults to `None`): - If set to `None`, the safetensors weights are downloaded if they're available **and** if the - safetensors library is installed. If set to `True`, the model is forcibly loaded from safetensors - weights. If set to `False`, safetensors weights are not loaded. - kwargs (remaining dictionary of keyword arguments, *optional*): - Can be used to overwrite load and saveable variables (the pipeline components of the specific pipeline - class). The overwritten components are passed directly to the pipelines `__init__` method. See example - below for more information. - variant (`str`, *optional*): - Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. - - - - To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with - `huggingface-cli login`. - - - - Examples: - - ```py - >>> from diffusers import AutoPipelineForInpainting - - >>> pipeline = AutoPipelineForInpainting.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") - >>> image = pipeline(prompt, image=init_image, mask_image=mask_image).images[0] - ``` - """ - cache_dir = kwargs.pop("cache_dir", None) - force_download = kwargs.pop("force_download", False) - proxies = kwargs.pop("proxies", None) - token = kwargs.pop("token", None) - local_files_only = kwargs.pop("local_files_only", False) - revision = kwargs.pop("revision", None) - - load_config_kwargs = { - "cache_dir": cache_dir, - "force_download": force_download, - "proxies": proxies, - "token": token, - "local_files_only": local_files_only, - "revision": revision, - } - - config = cls.load_config(pretrained_model_or_path, **load_config_kwargs) - orig_class_name = config["_class_name"] - - # The `orig_class_name`` can be: - # `- *InpaintPipeline` (for inpaint-specific checkpoint) - # - `*ControlPipeline` (for Flux tools specific checkpoint) - # - or *Pipeline (for regular text-to-image checkpoint) - if "Inpaint" in orig_class_name: - to_replace = "InpaintPipeline" - elif "ControlPipeline" in orig_class_name: - to_replace = "ControlPipeline" - else: - to_replace = "Pipeline" - - if "controlnet" in kwargs: - if isinstance(kwargs["controlnet"], ControlNetUnionModel): - orig_class_name = orig_class_name.replace(to_replace, "ControlNetUnion" + to_replace) - else: - orig_class_name = orig_class_name.replace(to_replace, "ControlNet" + to_replace) - if "enable_pag" in kwargs: - enable_pag = kwargs.pop("enable_pag") - if enable_pag: - orig_class_name = orig_class_name.replace(to_replace, "PAG" + to_replace) - if to_replace == "ControlPipeline": - orig_class_name = orig_class_name.replace(to_replace, "ControlInpaintPipeline") - inpainting_cls = _get_task_class(AUTO_INPAINT_PIPELINES_MAPPING, orig_class_name) - - kwargs = {**load_config_kwargs, **kwargs} - return inpainting_cls.from_pretrained(pretrained_model_or_path, **kwargs) - - @classmethod - def from_pipe(cls, pipeline, **kwargs): - r""" - Instantiates a inpainting Pytorch diffusion pipeline from another instantiated diffusion pipeline class. - - The from_pipe() method takes care of returning the correct pipeline class instance by finding the inpainting - pipeline linked to the pipeline class using pattern matching on pipeline class name. - - All the modules the pipeline class contain will be used to initialize the new pipeline without reallocating - additional memory. - - The pipeline is set in evaluation mode (`model.eval()`) by default. - - Parameters: - pipeline (`DiffusionPipeline`): - an instantiated `DiffusionPipeline` object - - Examples: - - ```py - >>> from diffusers import AutoPipelineForText2Image, AutoPipelineForInpainting - - >>> pipe_t2i = AutoPipelineForText2Image.from_pretrained( - ... "DeepFloyd/IF-I-XL-v1.0", requires_safety_checker=False - ... ) - - >>> pipe_inpaint = AutoPipelineForInpainting.from_pipe(pipe_t2i) - >>> image = pipe_inpaint(prompt, image=init_image, mask_image=mask_image).images[0] - ``` - """ - original_config = dict(pipeline.config) - original_cls_name = pipeline.__class__.__name__ - - # derive the pipeline class to instantiate - inpainting_cls = _get_task_class(AUTO_INPAINT_PIPELINES_MAPPING, original_cls_name) - - if "controlnet" in kwargs: - if kwargs["controlnet"] is not None: - inpainting_cls = _get_task_class( - AUTO_INPAINT_PIPELINES_MAPPING, - inpainting_cls.__name__.replace("ControlNet", "").replace( - "InpaintPipeline", "ControlNetInpaintPipeline" - ), - ) - else: - inpainting_cls = _get_task_class( - AUTO_INPAINT_PIPELINES_MAPPING, - inpainting_cls.__name__.replace("ControlNetInpaintPipeline", "InpaintPipeline"), - ) - - if "enable_pag" in kwargs: - enable_pag = kwargs.pop("enable_pag") - if enable_pag: - inpainting_cls = _get_task_class( - AUTO_INPAINT_PIPELINES_MAPPING, - inpainting_cls.__name__.replace("PAG", "").replace("InpaintPipeline", "PAGInpaintPipeline"), - ) - else: - inpainting_cls = _get_task_class( - AUTO_INPAINT_PIPELINES_MAPPING, - inpainting_cls.__name__.replace("PAGInpaintPipeline", "InpaintPipeline"), - ) - - # define expected module and optional kwargs given the pipeline signature - expected_modules, optional_kwargs = inpainting_cls._get_signature_keys(inpainting_cls) - - pretrained_model_name_or_path = original_config.pop("_name_or_path", None) - - # allow users pass modules in `kwargs` to override the original pipeline's components - passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs} - original_class_obj = { - k: pipeline.components[k] - for k, v in pipeline.components.items() - if k in expected_modules and k not in passed_class_obj - } - - # allow users pass optional kwargs to override the original pipelines config attribute - passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs} - original_pipe_kwargs = { - k: original_config[k] - for k, v in original_config.items() - if k in optional_kwargs and k not in passed_pipe_kwargs - } - - # config that were not expected by original pipeline is stored as private attribute - # we will pass them as optional arguments if they can be accepted by the pipeline - additional_pipe_kwargs = [ - k[1:] - for k in original_config.keys() - if k.startswith("_") and k[1:] in optional_kwargs and k[1:] not in passed_pipe_kwargs - ] - for k in additional_pipe_kwargs: - original_pipe_kwargs[k] = original_config.pop(f"_{k}") - - inpainting_kwargs = {**passed_class_obj, **original_class_obj, **passed_pipe_kwargs, **original_pipe_kwargs} - - # store unused config as private attribute - unused_original_config = { - f"{'' if k.startswith('_') else '_'}{k}": original_config[k] - for k, v in original_config.items() - if k not in inpainting_kwargs - } - - missing_modules = ( - set(expected_modules) - set(inpainting_cls._optional_components) - set(inpainting_kwargs.keys()) - ) - - if len(missing_modules) > 0: - raise ValueError( - f"Pipeline {inpainting_cls} expected {expected_modules}, but only {set(list(passed_class_obj.keys()) + list(original_class_obj.keys()))} were passed" - ) - - model = inpainting_cls(**inpainting_kwargs) - model.register_to_config(_name_or_path=pretrained_model_name_or_path) - model.register_to_config(**unused_original_config) - - return model diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__init__.py deleted file mode 100644 index af6c879d5ce88aa8edec0691e987444ff1d3dfec..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -from dataclasses import dataclass -from typing import List, Optional, Union - -import numpy as np -import PIL -from PIL import Image - -from ...utils import OptionalDependencyNotAvailable, is_torch_available, is_transformers_available - - -try: - if not (is_transformers_available() and is_torch_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_objects import ShapEPipeline -else: - from .blip_image_processing import BlipImageProcessor - from .modeling_blip2 import Blip2QFormerModel - from .modeling_ctx_clip import ContextCLIPTextModel - from .pipeline_blip_diffusion import BlipDiffusionPipeline diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index ac7da8dc6c8a32802ed482cd0ec7300a81564435..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/blip_image_processing.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/blip_image_processing.cpython-310.pyc deleted file mode 100644 index 027e78474562f88971bdf2cf1819300e82b5c562..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/blip_image_processing.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/modeling_blip2.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/modeling_blip2.cpython-310.pyc deleted file mode 100644 index 5c2a8768564540e9b097a8a38858fc1b9c37f01d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/modeling_blip2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/modeling_ctx_clip.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/modeling_ctx_clip.cpython-310.pyc deleted file mode 100644 index 8515dd9d9bc1c9005ff1cf9d5c650e80a37444b5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/modeling_ctx_clip.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/pipeline_blip_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/pipeline_blip_diffusion.cpython-310.pyc deleted file mode 100644 index b0d6d7aadbf52b764e68a98b3804aa0723d1c088..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/__pycache__/pipeline_blip_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/blip_image_processing.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/blip_image_processing.py deleted file mode 100644 index e45f431d0b9d21bd7566aef0a881434916d2ab46..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/blip_image_processing.py +++ /dev/null @@ -1,318 +0,0 @@ -# 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. -"""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 -from transformers.image_transforms import convert_to_rgb, resize, to_channel_dimension_format -from transformers.image_utils import ( - OPENAI_CLIP_MEAN, - OPENAI_CLIP_STD, - ChannelDimension, - ImageInput, - PILImageResampling, - infer_channel_dimension_format, - is_scaled_image, - make_list_of_images, - to_numpy_array, - valid_images, -) -from transformers.utils import TensorType, is_vision_available, logging - -from diffusers.utils import numpy_to_pil - - -if is_vision_available(): - import PIL.Image - - -logger = logging.get_logger(__name__) - - -# We needed some extra functions on top of the ones in transformers.image_processing_utils.BaseImageProcessor, namely center crop -# Copy-pasted from transformers.models.blip.image_processing_blip.BlipImageProcessor -class BlipImageProcessor(BaseImageProcessor): - r""" - Constructs a BLIP image processor. - - Args: - do_resize (`bool`, *optional*, defaults to `True`): - Whether to resize the image's (height, width) dimensions to the specified `size`. Can be overridden by the - `do_resize` parameter in the `preprocess` method. - size (`dict`, *optional*, defaults to `{"height": 384, "width": 384}`): - Size of the output image after resizing. Can be overridden by the `size` parameter in the `preprocess` - method. - resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`): - Resampling filter to use if resizing the image. Only has an effect if `do_resize` is set to `True`. Can be - overridden by the `resample` parameter in the `preprocess` method. - do_rescale (`bool`, *optional*, defaults to `True`): - Wwhether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the - `do_rescale` parameter in the `preprocess` method. - rescale_factor (`int` or `float`, *optional*, defaults to `1/255`): - Scale factor to use if rescaling the image. Only has an effect if `do_rescale` is set to `True`. Can be - overridden by the `rescale_factor` parameter in the `preprocess` method. - 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`): - 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`): - 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. - do_convert_rgb (`bool`, *optional*, defaults to `True`): - Whether to convert the image to RGB. - """ - - model_input_names = ["pixel_values"] - - def __init__( - self, - do_resize: bool = True, - size: Dict[str, int] = None, - resample: PILImageResampling = PILImageResampling.BICUBIC, - do_rescale: bool = True, - rescale_factor: Union[int, float] = 1 / 255, - do_normalize: bool = True, - image_mean: Optional[Union[float, List[float]]] = None, - image_std: Optional[Union[float, List[float]]] = None, - do_convert_rgb: bool = True, - do_center_crop: bool = True, - **kwargs, - ) -> None: - super().__init__(**kwargs) - size = size if size is not None else {"height": 224, "width": 224} - size = get_size_dict(size, default_to_square=True) - - self.do_resize = do_resize - self.size = size - self.resample = resample - self.do_rescale = do_rescale - self.rescale_factor = rescale_factor - self.do_normalize = do_normalize - self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN - self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD - self.do_convert_rgb = do_convert_rgb - self.do_center_crop = do_center_crop - - # Copy-pasted from transformers.models.vit.image_processing_vit.ViTImageProcessor.resize with PILImageResampling.BILINEAR->PILImageResampling.BICUBIC - def resize( - self, - image: np.ndarray, - size: Dict[str, int], - resample: PILImageResampling = PILImageResampling.BICUBIC, - data_format: Optional[Union[str, ChannelDimension]] = None, - input_data_format: Optional[Union[str, ChannelDimension]] = None, - **kwargs, - ) -> np.ndarray: - """ - Resize an image to `(size["height"], size["width"])`. - - Args: - image (`np.ndarray`): - Image to resize. - 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`. - data_format (`ChannelDimension` or `str`, *optional*): - The channel dimension format for the output image. If unset, the channel dimension format of the input - image is used. Can be one of: - - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. - - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. - - `"none"` or `ChannelDimension.NONE`: image in (height, width) format. - input_data_format (`ChannelDimension` or `str`, *optional*): - The channel dimension format for the input image. If unset, the channel dimension format is inferred - from the input image. Can be one of: - - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. - - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. - - `"none"` or `ChannelDimension.NONE`: image in (height, width) format. - - Returns: - `np.ndarray`: The resized image. - """ - size = get_size_dict(size) - if "height" not in size or "width" not in size: - raise ValueError(f"The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}") - output_size = (size["height"], size["width"]) - return resize( - image, - size=output_size, - resample=resample, - data_format=data_format, - input_data_format=input_data_format, - **kwargs, - ) - - def preprocess( - self, - images: ImageInput, - do_resize: Optional[bool] = None, - size: Optional[Dict[str, int]] = 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_convert_rgb: bool = None, - data_format: ChannelDimension = ChannelDimension.FIRST, - input_data_format: Optional[Union[str, ChannelDimension]] = None, - **kwargs, - ) -> PIL.Image.Image: - """ - Preprocess an image or batch of images. - - Args: - images (`ImageInput`): - Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If - 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`): - 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 - edge equal to `int(size["shortest_edge"] * (1333 / 800))`. - resample (`PILImageResampling`, *optional*, defaults to `self.resample`): - Resampling filter to use if resizing the image. Only has an effect if `do_resize` is set to `True`. - do_rescale (`bool`, *optional*, defaults to `self.do_rescale`): - Whether to rescale the image values between [0 - 1]. - rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`): - 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 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 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. - return_tensors (`str` or `TensorType`, *optional*): - The type of tensors to return. Can be one of: - - Unset: Return a list of `np.ndarray`. - - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`. - - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`. - - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`. - - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`. - data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`): - The channel dimension format for the output image. Can be one of: - - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. - - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. - - Unset: Use the channel dimension format of the input image. - input_data_format (`ChannelDimension` or `str`, *optional*): - The channel dimension format for the input image. If unset, the channel dimension format is inferred - from the input image. Can be one of: - - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. - - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. - - `"none"` or `ChannelDimension.NONE`: image in (height, width) format. - """ - do_resize = do_resize if do_resize is not None else self.do_resize - resample = resample if resample is not None else self.resample - do_rescale = do_rescale if do_rescale is not None else self.do_rescale - rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor - do_normalize = do_normalize if do_normalize is not None else self.do_normalize - image_mean = image_mean if image_mean is not None else self.image_mean - image_std = image_std if image_std is not None else self.image_std - do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb - do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop - - size = size if size is not None else self.size - size = get_size_dict(size, default_to_square=False) - images = make_list_of_images(images) - - if not valid_images(images): - raise ValueError( - "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, " - "torch.Tensor, tf.Tensor or jax.ndarray." - ) - - if do_resize and size is None or resample is None: - raise ValueError("Size and resample must be specified if do_resize is True.") - - if do_rescale and rescale_factor is None: - raise ValueError("Rescale factor must be specified if do_rescale is True.") - - if do_normalize and (image_mean is None or image_std is None): - raise ValueError("Image mean and std must be specified if do_normalize is True.") - - # PIL RGBA images are converted to RGB - if do_convert_rgb: - images = [convert_to_rgb(image) for image in images] - - # All transformations expect numpy arrays. - images = [to_numpy_array(image) for image in images] - - if is_scaled_image(images[0]) and do_rescale: - logger.warning_once( - "It looks like you are trying to rescale already rescaled images. If the input" - " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again." - ) - if input_data_format is None: - # We assume that all images have the same channel dimension format. - input_data_format = infer_channel_dimension_format(images[0]) - - if do_resize: - images = [ - self.resize(image=image, size=size, resample=resample, input_data_format=input_data_format) - for image in images - ] - - if do_rescale: - images = [ - self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format) - for image in images - ] - if do_normalize: - images = [ - self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format) - for image in images - ] - if do_center_crop: - images = [self.center_crop(image, size, input_data_format=input_data_format) for image in images] - - images = [ - to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images - ] - - encoded_outputs = BatchFeature(data={"pixel_values": images}, tensor_type=return_tensors) - return encoded_outputs - - # Follows diffusers.VaeImageProcessor.postprocess - def postprocess(self, sample: torch.Tensor, output_type: str = "pil"): - if output_type not in ["pt", "np", "pil"]: - raise ValueError( - f"output_type={output_type} is not supported. Make sure to choose one of ['pt', 'np', or 'pil']" - ) - - # Equivalent to diffusers.VaeImageProcessor.denormalize - sample = (sample / 2 + 0.5).clamp(0, 1) - if output_type == "pt": - return sample - - # Equivalent to diffusers.VaeImageProcessor.pt_to_numpy - sample = sample.cpu().permute(0, 2, 3, 1).numpy() - if output_type == "np": - return sample - # Output_type must be 'pil' - sample = numpy_to_pil(sample) - return sample diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/modeling_blip2.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/modeling_blip2.py deleted file mode 100644 index 928698e44286f499e77878707bff7f338bf3fb05..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/modeling_blip2.py +++ /dev/null @@ -1,639 +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 - -import torch -import torch.utils.checkpoint -from torch import nn -from transformers import BertTokenizer -from transformers.activations import QuickGELUActivation as QuickGELU -from transformers.modeling_outputs import ( - BaseModelOutputWithPastAndCrossAttentions, - BaseModelOutputWithPooling, - BaseModelOutputWithPoolingAndCrossAttentions, -) -from transformers.models.blip_2.configuration_blip_2 import Blip2Config, Blip2VisionConfig -from transformers.models.blip_2.modeling_blip_2 import ( - Blip2Encoder, - Blip2PreTrainedModel, - Blip2QFormerAttention, - Blip2QFormerIntermediate, - Blip2QFormerOutput, -) -from transformers.pytorch_utils import apply_chunking_to_forward -from transformers.utils import ( - logging, - replace_return_docstrings, -) - - -logger = logging.get_logger(__name__) - - -# There is an implementation of Blip2 in `transformers` : https://github.com/huggingface/transformers/blob/main/src/transformers/models/blip_2/modeling_blip_2.py. -# But it doesn't support getting multimodal embeddings. So, this module can be -# replaced with a future `transformers` version supports that. -class Blip2TextEmbeddings(nn.Module): - """Construct the embeddings from word and position embeddings.""" - - def __init__(self, config): - super().__init__() - self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) - self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) - - # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load - # any TensorFlow checkpoint file - self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - self.dropout = nn.Dropout(config.hidden_dropout_prob) - - # position_ids (1, len position emb) is contiguous in memory and exported when serialized - self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1))) - self.position_embedding_type = getattr(config, "position_embedding_type", "absolute") - - self.config = config - - def forward( - self, - input_ids=None, - position_ids=None, - query_embeds=None, - past_key_values_length=0, - ): - if input_ids is not None: - seq_length = input_ids.size()[1] - else: - seq_length = 0 - - if position_ids is None: - position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length].clone() - - if input_ids is not None: - embeddings = self.word_embeddings(input_ids) - if self.position_embedding_type == "absolute": - position_embeddings = self.position_embeddings(position_ids) - embeddings = embeddings + position_embeddings - - if query_embeds is not None: - batch_size = embeddings.shape[0] - # repeat the query embeddings for batch size - query_embeds = query_embeds.repeat(batch_size, 1, 1) - embeddings = torch.cat((query_embeds, embeddings), dim=1) - else: - embeddings = query_embeds - embeddings = embeddings.to(query_embeds.dtype) - embeddings = self.LayerNorm(embeddings) - embeddings = self.dropout(embeddings) - return embeddings - - -# Copy-pasted from transformers.models.blip.modeling_blip.BlipVisionEmbeddings with Blip->Blip2 -class Blip2VisionEmbeddings(nn.Module): - def __init__(self, config: Blip2VisionConfig): - super().__init__() - self.config = config - self.embed_dim = config.hidden_size - self.image_size = config.image_size - self.patch_size = config.patch_size - - self.class_embedding = nn.Parameter(torch.randn(1, 1, self.embed_dim)) - - self.patch_embedding = nn.Conv2d( - in_channels=3, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size, bias=False - ) - - self.num_patches = (self.image_size // self.patch_size) ** 2 - self.num_positions = self.num_patches + 1 - - self.position_embedding = nn.Parameter(torch.randn(1, self.num_positions, self.embed_dim)) - - def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: - batch_size = pixel_values.shape[0] - target_dtype = self.patch_embedding.weight.dtype - patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid] - patch_embeds = patch_embeds.flatten(2).transpose(1, 2) - - class_embeds = self.class_embedding.expand(batch_size, 1, -1).to(target_dtype) - embeddings = torch.cat([class_embeds, patch_embeds], dim=1) - embeddings = embeddings + self.position_embedding[:, : embeddings.size(1), :].to(target_dtype) - return embeddings - - -# The Qformer encoder, which takes the visual embeddings, and the text input, to get multimodal embeddings -class Blip2QFormerEncoder(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.layer = nn.ModuleList( - [Blip2QFormerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] - ) - self.gradient_checkpointing = False - - def forward( - self, - hidden_states, - attention_mask=None, - head_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - past_key_values=None, - use_cache=None, - output_attentions=False, - output_hidden_states=False, - return_dict=True, - query_length=0, - ): - all_hidden_states = () if output_hidden_states else None - all_self_attentions = () if output_attentions else None - all_cross_attentions = () if output_attentions else None - - next_decoder_cache = () if use_cache else None - - for i in range(self.config.num_hidden_layers): - layer_module = self.layer[i] - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - layer_head_mask = head_mask[i] if head_mask is not None else None - past_key_value = past_key_values[i] if past_key_values is not None else None - - if getattr(self.config, "gradient_checkpointing", False) and torch.is_grad_enabled(): - if use_cache: - logger.warning( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - - layer_outputs = self._gradient_checkpointing_func( - layer_module, - hidden_states, - attention_mask, - layer_head_mask, - encoder_hidden_states, - encoder_attention_mask, - past_key_value, - output_attentions, - query_length, - ) - else: - layer_outputs = layer_module( - hidden_states, - attention_mask, - layer_head_mask, - encoder_hidden_states, - encoder_attention_mask, - past_key_value, - output_attentions, - query_length, - ) - - hidden_states = layer_outputs[0] - if use_cache: - next_decoder_cache += (layer_outputs[-1],) - if output_attentions: - all_self_attentions = all_self_attentions + (layer_outputs[1],) - if layer_module.has_cross_attention: - all_cross_attentions = all_cross_attentions + (layer_outputs[2],) - - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - if not return_dict: - return tuple( - v - for v in [ - hidden_states, - next_decoder_cache, - all_hidden_states, - all_self_attentions, - all_cross_attentions, - ] - if v is not None - ) - return BaseModelOutputWithPastAndCrossAttentions( - last_hidden_state=hidden_states, - past_key_values=next_decoder_cache, - hidden_states=all_hidden_states, - attentions=all_self_attentions, - cross_attentions=all_cross_attentions, - ) - - -# The layers making up the Qformer encoder -class Blip2QFormerLayer(nn.Module): - def __init__(self, config, layer_idx): - super().__init__() - self.chunk_size_feed_forward = config.chunk_size_feed_forward - self.seq_len_dim = 1 - self.attention = Blip2QFormerAttention(config) - - self.layer_idx = layer_idx - - if layer_idx % config.cross_attention_frequency == 0: - self.crossattention = Blip2QFormerAttention(config, is_cross_attention=True) - self.has_cross_attention = True - else: - self.has_cross_attention = False - - self.intermediate = Blip2QFormerIntermediate(config) - self.intermediate_query = Blip2QFormerIntermediate(config) - self.output_query = Blip2QFormerOutput(config) - self.output = Blip2QFormerOutput(config) - - def forward( - self, - hidden_states, - attention_mask=None, - head_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - past_key_value=None, - output_attentions=False, - query_length=0, - ): - # decoder uni-directional self-attention cached key/values tuple is at positions 1,2 - self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None - self_attention_outputs = self.attention( - hidden_states, - attention_mask, - head_mask, - output_attentions=output_attentions, - past_key_value=self_attn_past_key_value, - ) - attention_output = self_attention_outputs[0] - outputs = self_attention_outputs[1:-1] - - present_key_value = self_attention_outputs[-1] - - if query_length > 0: - query_attention_output = attention_output[:, :query_length, :] - - if self.has_cross_attention: - if encoder_hidden_states is None: - raise ValueError("encoder_hidden_states must be given for cross-attention layers") - cross_attention_outputs = self.crossattention( - query_attention_output, - attention_mask, - head_mask, - encoder_hidden_states, - encoder_attention_mask, - output_attentions=output_attentions, - ) - query_attention_output = cross_attention_outputs[0] - # add cross attentions if we output attention weights - outputs = outputs + cross_attention_outputs[1:-1] - - layer_output = apply_chunking_to_forward( - self.feed_forward_chunk_query, - self.chunk_size_feed_forward, - self.seq_len_dim, - query_attention_output, - ) - - if attention_output.shape[1] > query_length: - layer_output_text = apply_chunking_to_forward( - self.feed_forward_chunk, - self.chunk_size_feed_forward, - self.seq_len_dim, - attention_output[:, query_length:, :], - ) - layer_output = torch.cat([layer_output, layer_output_text], dim=1) - else: - layer_output = apply_chunking_to_forward( - self.feed_forward_chunk, - self.chunk_size_feed_forward, - self.seq_len_dim, - attention_output, - ) - outputs = (layer_output,) + outputs - - outputs = outputs + (present_key_value,) - - return outputs - - def feed_forward_chunk(self, attention_output): - intermediate_output = self.intermediate(attention_output) - layer_output = self.output(intermediate_output, attention_output) - return layer_output - - def feed_forward_chunk_query(self, attention_output): - intermediate_output = self.intermediate_query(attention_output) - layer_output = self.output_query(intermediate_output, attention_output) - return layer_output - - -# ProjLayer used to project the multimodal Blip2 embeddings to be used in the text encoder -class ProjLayer(nn.Module): - def __init__(self, in_dim, out_dim, hidden_dim, drop_p=0.1, eps=1e-12): - super().__init__() - - # Dense1 -> Act -> Dense2 -> Drop -> Res -> Norm - self.dense1 = nn.Linear(in_dim, hidden_dim) - self.act_fn = QuickGELU() - self.dense2 = nn.Linear(hidden_dim, out_dim) - self.dropout = nn.Dropout(drop_p) - - self.LayerNorm = nn.LayerNorm(out_dim, eps=eps) - - def forward(self, x): - x_in = x - - x = self.LayerNorm(x) - x = self.dropout(self.dense2(self.act_fn(self.dense1(x)))) + x_in - - return x - - -# Copy-pasted from transformers.models.blip.modeling_blip.BlipVisionModel with Blip->Blip2, BLIP->BLIP_2 -class Blip2VisionModel(Blip2PreTrainedModel): - main_input_name = "pixel_values" - config_class = Blip2VisionConfig - - def __init__(self, config: Blip2VisionConfig): - super().__init__(config) - self.config = config - embed_dim = config.hidden_size - self.embeddings = Blip2VisionEmbeddings(config) - self.pre_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) - self.encoder = Blip2Encoder(config) - self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) - - self.post_init() - - @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]: - r""" - Returns: - - """ - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if pixel_values is None: - raise ValueError("You have to specify pixel_values") - - hidden_states = self.embeddings(pixel_values) - hidden_states = self.pre_layernorm(hidden_states) - encoder_outputs = self.encoder( - inputs_embeds=hidden_states, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - last_hidden_state = encoder_outputs[0] - last_hidden_state = self.post_layernorm(last_hidden_state) - - pooled_output = last_hidden_state[:, 0, :] - pooled_output = self.post_layernorm(pooled_output) - - if not return_dict: - return (last_hidden_state, pooled_output) + encoder_outputs[1:] - - return BaseModelOutputWithPooling( - last_hidden_state=last_hidden_state, - pooler_output=pooled_output, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - ) - - def get_input_embeddings(self): - return self.embeddings - - -# Qformer model, used to get multimodal embeddings from the text and image inputs -class Blip2QFormerModel(Blip2PreTrainedModel): - """ - Querying Transformer (Q-Former), used in BLIP-2. - """ - - def __init__(self, config: Blip2Config): - super().__init__(config) - self.config = config - self.embeddings = Blip2TextEmbeddings(config.qformer_config) - self.visual_encoder = Blip2VisionModel(config.vision_config) - self.query_tokens = nn.Parameter(torch.zeros(1, config.num_query_tokens, config.qformer_config.hidden_size)) - if not hasattr(config, "tokenizer") or config.tokenizer is None: - self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased", truncation_side="right") - else: - self.tokenizer = BertTokenizer.from_pretrained(config.tokenizer, truncation_side="right") - self.tokenizer.add_special_tokens({"bos_token": "[DEC]"}) - self.proj_layer = ProjLayer( - in_dim=config.qformer_config.hidden_size, - out_dim=config.qformer_config.hidden_size, - hidden_dim=config.qformer_config.hidden_size * 4, - drop_p=0.1, - eps=1e-12, - ) - - self.encoder = Blip2QFormerEncoder(config.qformer_config) - - self.post_init() - - def get_input_embeddings(self): - return self.embeddings.word_embeddings - - def set_input_embeddings(self, value): - self.embeddings.word_embeddings = value - - def _prune_heads(self, heads_to_prune): - """ - Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base - class PreTrainedModel - """ - for layer, heads in heads_to_prune.items(): - self.encoder.layer[layer].attention.prune_heads(heads) - - def get_extended_attention_mask( - self, - attention_mask: torch.Tensor, - input_shape: Tuple[int], - device: torch.device, - has_query: bool = False, - ) -> torch.Tensor: - """ - Makes broadcastable attention and causal masks so that future and masked tokens are ignored. - - Arguments: - attention_mask (`torch.Tensor`): - Mask with ones indicating tokens to attend to, zeros for tokens to ignore. - input_shape (`Tuple[int]`): - The shape of the input to the model. - device (`torch.device`): - The device of the input to the model. - - Returns: - `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`. - """ - # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] - # ourselves in which case we just need to make it broadcastable to all heads. - if attention_mask.dim() == 3: - extended_attention_mask = attention_mask[:, None, :, :] - elif attention_mask.dim() == 2: - # Provided a padding mask of dimensions [batch_size, seq_length] - # - the model is an encoder, so make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length] - extended_attention_mask = attention_mask[:, None, None, :] - else: - raise ValueError( - "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format( - input_shape, attention_mask.shape - ) - ) - - # Since attention_mask is 1.0 for positions we want to attend and 0.0 for - # masked positions, this operation will create a tensor which is 0.0 for - # positions we want to attend and -10000.0 for masked positions. - # Since we are adding it to the raw scores before the softmax, this is - # effectively the same as removing these entirely. - extended_attention_mask = extended_attention_mask.to(dtype=self.dtype) # fp16 compatibility - extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0 - return extended_attention_mask - - def forward( - self, - text_input=None, - image_input=None, - head_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - past_key_values=None, - use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - ): - r""" - encoder_hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, `optional`): - Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if - the model is configured as a decoder. - encoder_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, `optional`): - Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in - the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - past_key_values (`tuple(tuple(torch.Tensor))` of length `config.n_layers` with each tuple having 4 tensors of: - shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): Contains precomputed key and - value hidden states of the attention blocks. Can be used to speed up decoding. If `past_key_values` are - used, the user can optionally input only the last `decoder_input_ids` (those that don't have their past key - value states given to this model) of shape `(batch_size, 1)` instead of all `decoder_input_ids` of shape - `(batch_size, sequence_length)`. - use_cache (`bool`, `optional`): - If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see - `past_key_values`). - """ - - text = self.tokenizer(text_input, return_tensors="pt", padding=True) - text = text.to(self.device) - input_ids = text.input_ids - batch_size = input_ids.shape[0] - query_atts = torch.ones((batch_size, self.query_tokens.size()[1]), dtype=torch.long).to(self.device) - attention_mask = torch.cat([query_atts, text.attention_mask], dim=1) - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # past_key_values_length - past_key_values_length = ( - past_key_values[0][0].shape[2] - self.config.query_length if past_key_values is not None else 0 - ) - - query_length = self.query_tokens.shape[1] - - embedding_output = self.embeddings( - input_ids=input_ids, - query_embeds=self.query_tokens, - past_key_values_length=past_key_values_length, - ) - - # embedding_output = self.layernorm(query_embeds) - # embedding_output = self.dropout(embedding_output) - - input_shape = embedding_output.size()[:-1] - batch_size, seq_length = input_shape - device = embedding_output.device - - image_embeds_frozen = self.visual_encoder(image_input).last_hidden_state - # image_embeds_frozen = torch.ones_like(image_embeds_frozen) - encoder_hidden_states = image_embeds_frozen - - if attention_mask is None: - attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device) - - # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] - # ourselves in which case we just need to make it broadcastable to all heads. - extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape, device) - - # If a 2D or 3D attention mask is provided for the cross-attention - # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] - if encoder_hidden_states is not None: - if isinstance(encoder_hidden_states, list): - encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[0].size() - else: - encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() - encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) - - if isinstance(encoder_attention_mask, list): - encoder_extended_attention_mask = [self.invert_attention_mask(mask) for mask in encoder_attention_mask] - elif encoder_attention_mask is None: - encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) - encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) - else: - encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) - else: - encoder_extended_attention_mask = None - - # Prepare head mask if needed - # 1.0 in head_mask indicate we keep the head - # attention_probs has shape bsz x n_heads x N x N - # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] - # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] - head_mask = self.get_head_mask(head_mask, self.config.qformer_config.num_hidden_layers) - - encoder_outputs = self.encoder( - embedding_output, - attention_mask=extended_attention_mask, - head_mask=head_mask, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_extended_attention_mask, - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - query_length=query_length, - ) - sequence_output = encoder_outputs[0] - pooled_output = sequence_output[:, 0, :] - - if not return_dict: - return self.proj_layer(sequence_output[:, :query_length, :]) - - return BaseModelOutputWithPoolingAndCrossAttentions( - last_hidden_state=sequence_output, - pooler_output=pooled_output, - past_key_values=encoder_outputs.past_key_values, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - cross_attentions=encoder_outputs.cross_attentions, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py deleted file mode 100644 index 1b0342ce7a5628210a7607e496a1f71b51ad2ee0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py +++ /dev/null @@ -1,223 +0,0 @@ -# Copyright 2025 Salesforce.com, inc. -# 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 - -import torch -from torch import nn -from transformers import CLIPPreTrainedModel -from transformers.modeling_outputs import BaseModelOutputWithPooling -from transformers.models.clip.configuration_clip import CLIPTextConfig -from transformers.models.clip.modeling_clip import CLIPEncoder - - -def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): - """ - Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. - """ - bsz, src_len = mask.size() - tgt_len = tgt_len if tgt_len is not None else src_len - - expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype) - - inverted_mask = 1.0 - expanded_mask - - return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min) - - -# This is a modified version of the CLIPTextModel from transformers.models.clip.modeling_clip -# Which allows for an extra input of "context embeddings", which are the query embeddings used in Qformer -# They pass through the clip model, along with the text embeddings, and interact with them using self attention -class ContextCLIPTextModel(CLIPPreTrainedModel): - config_class = CLIPTextConfig - - _no_split_modules = ["CLIPEncoderLayer"] - - def __init__(self, config: CLIPTextConfig): - super().__init__(config) - self.text_model = ContextCLIPTextTransformer(config) - # Initialize weights and apply final processing - self.post_init() - - 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]: - return self.text_model( - ctx_embeddings=ctx_embeddings, - ctx_begin_pos=ctx_begin_pos, - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - -class ContextCLIPTextTransformer(nn.Module): - def __init__(self, config: CLIPTextConfig): - super().__init__() - self.config = config - embed_dim = config.hidden_size - self.embeddings = ContextCLIPTextEmbeddings(config) - self.encoder = CLIPEncoder(config) - self.final_layer_norm = nn.LayerNorm(embed_dim) - - 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]: - r""" - Returns: - - """ - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if input_ids is None: - raise ValueError("You have to specify either input_ids") - - input_shape = input_ids.size() - input_ids = input_ids.view(-1, input_shape[-1]) - - hidden_states = self.embeddings( - input_ids=input_ids, - position_ids=position_ids, - ctx_embeddings=ctx_embeddings, - ctx_begin_pos=ctx_begin_pos, - ) - - bsz, seq_len = input_shape - if ctx_embeddings is not None: - seq_len += ctx_embeddings.size(1) - # CLIP's text model uses causal mask, prepare it here. - # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 - causal_attention_mask = self._build_causal_attention_mask(bsz, seq_len, hidden_states.dtype).to( - hidden_states.device - ) - # expand attention_mask - if attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _expand_mask(attention_mask, hidden_states.dtype) - - encoder_outputs = self.encoder( - inputs_embeds=hidden_states, - attention_mask=attention_mask, - causal_attention_mask=causal_attention_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - last_hidden_state = encoder_outputs[0] - last_hidden_state = self.final_layer_norm(last_hidden_state) - - # text_embeds.shape = [batch_size, sequence_length, transformer.width] - # take features from the eot embedding (eot_token is the highest number in each sequence) - # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14 - pooled_output = last_hidden_state[ - torch.arange(last_hidden_state.shape[0], device=input_ids.device), - input_ids.to(torch.int).argmax(dim=-1), - ] - - if not return_dict: - return (last_hidden_state, pooled_output) + encoder_outputs[1:] - - return BaseModelOutputWithPooling( - last_hidden_state=last_hidden_state, - pooler_output=pooled_output, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - ) - - def _build_causal_attention_mask(self, bsz, seq_len, dtype): - # lazily create causal attention mask, with full attention between the vision tokens - # pytorch uses additive attention mask; fill with -inf - mask = torch.empty(bsz, seq_len, seq_len, dtype=dtype) - mask.fill_(torch.tensor(torch.finfo(dtype).min)) - mask.triu_(1) # zero out the lower diagonal - mask = mask.unsqueeze(1) # expand mask - return mask - - -class ContextCLIPTextEmbeddings(nn.Module): - def __init__(self, config: CLIPTextConfig): - super().__init__() - embed_dim = config.hidden_size - - self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) - self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim) - - # position_ids (1, len position emb) is contiguous in memory and exported when serialized - self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1))) - - 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, - ) -> torch.Tensor: - if ctx_embeddings is None: - ctx_len = 0 - else: - ctx_len = ctx_embeddings.shape[1] - - seq_length = (input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2]) + ctx_len - - if position_ids is None: - position_ids = self.position_ids[:, :seq_length] - - if inputs_embeds is None: - inputs_embeds = self.token_embedding(input_ids) - - # for each input embeddings, add the ctx embeddings at the correct position - input_embeds_ctx = [] - bsz = inputs_embeds.shape[0] - - if ctx_embeddings is not None: - for i in range(bsz): - cbp = ctx_begin_pos[i] - - prefix = inputs_embeds[i, :cbp] - # remove the special token embedding - suffix = inputs_embeds[i, cbp:] - - input_embeds_ctx.append(torch.cat([prefix, ctx_embeddings[i], suffix], dim=0)) - - inputs_embeds = torch.stack(input_embeds_ctx, dim=0) - - position_embeddings = self.position_embedding(position_ids) - embeddings = inputs_embeds + position_embeddings - - return embeddings diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py deleted file mode 100644 index 439dc511a0c990d051c1b1af875336d2c1842718..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py +++ /dev/null @@ -1,361 +0,0 @@ -# Copyright 2025 Salesforce.com, inc. -# 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 List, Optional, Union - -import PIL.Image -import torch -from transformers import CLIPTokenizer - -from ...models import AutoencoderKL, UNet2DConditionModel -from ...schedulers import PNDMScheduler -from ...utils import ( - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, ImagePipelineOutput -from .blip_image_processing import BlipImageProcessor -from .modeling_blip2 import Blip2QFormerModel -from .modeling_ctx_clip import ContextCLIPTextModel - - -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 - >>> from diffusers.pipelines import BlipDiffusionPipeline - >>> from diffusers.utils import load_image - >>> import torch - - >>> blip_diffusion_pipe = BlipDiffusionPipeline.from_pretrained( - ... "Salesforce/blipdiffusion", torch_dtype=torch.float16 - ... ).to("cuda") - - - >>> cond_subject = "dog" - >>> tgt_subject = "dog" - >>> text_prompt_input = "swimming underwater" - - >>> cond_image = load_image( - ... "https://huggingface.co/datasets/ayushtues/blipdiffusion_images/resolve/main/dog.jpg" - ... ) - >>> guidance_scale = 7.5 - >>> num_inference_steps = 25 - >>> negative_prompt = "over-exposure, under-exposure, saturated, duplicate, out of frame, lowres, cropped, worst quality, low quality, jpeg artifacts, morbid, mutilated, out of frame, ugly, bad anatomy, bad proportions, deformed, blurry, duplicate" - - - >>> output = blip_diffusion_pipe( - ... text_prompt_input, - ... cond_image, - ... cond_subject, - ... tgt_subject, - ... guidance_scale=guidance_scale, - ... num_inference_steps=num_inference_steps, - ... neg_prompt=negative_prompt, - ... height=512, - ... width=512, - ... ).images - >>> output[0].save("image.png") - ``` -""" - - -class BlipDiffusionPipeline(DeprecatedPipelineMixin, DiffusionPipeline): - """ - Pipeline for Zero-Shot Subject Driven Generation using Blip 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.) - - Args: - tokenizer ([`CLIPTokenizer`]): - Tokenizer for the text encoder - text_encoder ([`ContextCLIPTextModel`]): - Text encoder to encode the text prompt - vae ([`AutoencoderKL`]): - VAE model to map the latents to the image - unet ([`UNet2DConditionModel`]): - Conditional U-Net architecture to denoise the image embedding. - scheduler ([`PNDMScheduler`]): - A scheduler to be used in combination with `unet` to generate image latents. - qformer ([`Blip2QFormerModel`]): - QFormer model to get multi-modal embeddings from the text and image. - image_processor ([`BlipImageProcessor`]): - Image Processor to preprocess and postprocess the image. - ctx_begin_pos (int, `optional`, defaults to 2): - Position of the context token in the text encoder. - """ - - _last_supported_version = "0.33.1" - model_cpu_offload_seq = "qformer->text_encoder->unet->vae" - - def __init__( - self, - tokenizer: CLIPTokenizer, - text_encoder: ContextCLIPTextModel, - vae: AutoencoderKL, - unet: UNet2DConditionModel, - scheduler: PNDMScheduler, - qformer: Blip2QFormerModel, - image_processor: BlipImageProcessor, - ctx_begin_pos: int = 2, - mean: List[float] = None, - std: List[float] = None, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - vae=vae, - unet=unet, - scheduler=scheduler, - qformer=qformer, - image_processor=image_processor, - ) - self.register_to_config(ctx_begin_pos=ctx_begin_pos, mean=mean, std=std) - - def get_query_embeddings(self, input_image, src_subject): - return self.qformer(image_input=input_image, text_input=src_subject, return_dict=False) - - # from the original Blip Diffusion code, specifies the target subject and augments the prompt by repeating it - def _build_prompt(self, prompts, tgt_subjects, prompt_strength=1.0, prompt_reps=20): - rv = [] - for prompt, tgt_subject in zip(prompts, tgt_subjects): - prompt = f"a {tgt_subject} {prompt.strip()}" - # a trick to amplify the prompt - rv.append(", ".join([prompt] * int(prompt_strength * prompt_reps))) - - return rv - - # Copied from diffusers.pipelines.consistency_models.pipeline_consistency_models.ConsistencyModelPipeline.prepare_latents - def prepare_latents(self, batch_size, num_channels, height, width, dtype, device, generator, latents=None): - shape = (batch_size, num_channels, 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." - ) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device=device, dtype=dtype) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def encode_prompt(self, query_embeds, prompt, device=None): - device = device or self._execution_device - - # embeddings for prompt, with query_embeds as context - max_len = self.text_encoder.text_model.config.max_position_embeddings - max_len -= self.qformer.config.num_query_tokens - - tokenized_prompt = self.tokenizer( - prompt, - padding="max_length", - truncation=True, - max_length=max_len, - return_tensors="pt", - ).to(device) - - batch_size = query_embeds.shape[0] - ctx_begin_pos = [self.config.ctx_begin_pos] * batch_size - - text_embeddings = self.text_encoder( - input_ids=tokenized_prompt.input_ids, - ctx_embeddings=query_embeds, - ctx_begin_pos=ctx_begin_pos, - )[0] - - return text_embeddings - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: List[str], - reference_image: PIL.Image.Image, - source_subject_category: List[str], - target_subject_category: List[str], - latents: Optional[torch.Tensor] = 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] = "", - prompt_strength: float = 1.0, - prompt_reps: int = 20, - output_type: Optional[str] = "pil", - return_dict: bool = True, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - 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]`): - The source subject category. - 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 - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor will ge generated by random sampling. - 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. - height (`int`, *optional*, defaults to 512): - The height of the generated image. - width (`int`, *optional*, defaults to 512): - The width 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. - 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 ""): - 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_strength (`float`, *optional*, defaults to 1.0): - The strength of the prompt. Specifies the number of times the prompt is repeated along with prompt_reps - to amplify the prompt. - prompt_reps (`int`, *optional*, defaults to 20): - The number of times the prompt is repeated along with prompt_strength to amplify the prompt. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - device = self._execution_device - - reference_image = self.image_processor.preprocess( - reference_image, image_mean=self.config.mean, image_std=self.config.std, return_tensors="pt" - )["pixel_values"] - reference_image = reference_image.to(device) - - if isinstance(prompt, str): - prompt = [prompt] - if isinstance(source_subject_category, str): - source_subject_category = [source_subject_category] - if isinstance(target_subject_category, str): - target_subject_category = [target_subject_category] - - batch_size = len(prompt) - - prompt = self._build_prompt( - prompts=prompt, - tgt_subjects=target_subject_category, - prompt_strength=prompt_strength, - prompt_reps=prompt_reps, - ) - query_embeds = self.get_query_embeddings(reference_image, source_subject_category) - text_embeddings = self.encode_prompt(query_embeds, prompt, device) - do_classifier_free_guidance = guidance_scale > 1.0 - if do_classifier_free_guidance: - max_length = self.text_encoder.text_model.config.max_position_embeddings - - uncond_input = self.tokenizer( - [neg_prompt] * batch_size, - padding="max_length", - max_length=max_length, - return_tensors="pt", - ) - uncond_embeddings = self.text_encoder( - input_ids=uncond_input.input_ids.to(device), - ctx_embeddings=None, - )[0] - # 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 - text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) - - scale_down_factor = 2 ** (len(self.unet.config.block_out_channels) - 1) - latents = self.prepare_latents( - batch_size=batch_size, - num_channels=self.unet.config.in_channels, - height=height // scale_down_factor, - width=width // scale_down_factor, - generator=generator, - latents=latents, - dtype=self.unet.dtype, - device=device, - ) - # set timesteps - extra_set_kwargs = {} - self.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs) - - for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)): - # expand the latents if we are doing classifier free guidance - do_classifier_free_guidance = guidance_scale > 1.0 - - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - noise_pred = self.unet( - latent_model_input, - timestep=t, - encoder_hidden_states=text_embeddings, - down_block_additional_residuals=None, - mid_block_additional_residual=None, - )["sample"] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - latents = self.scheduler.step( - noise_pred, - t, - latents, - )["prev_sample"] - - if XLA_AVAILABLE: - xm.mark_step() - - image = self.vae.decode(latents / self.vae.config.scaling_factor, 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 ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__init__.py deleted file mode 100644 index d9238b735c4156f25a831671196f418bd565db60..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__init__.py +++ /dev/null @@ -1,49 +0,0 @@ -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 = {} -_additional_imports = {} -_import_structure = {"pipeline_output": ["ChromaPipelineOutput"]} - -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_chroma"] = ["ChromaPipeline"] - _import_structure["pipeline_chroma_img2img"] = ["ChromaImg2ImgPipeline"] -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_chroma import ChromaPipeline - from .pipeline_chroma_img2img import ChromaImg2ImgPipeline -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/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 2e6e06b647160045ee8ff599fbcb465bfdfa9831..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__pycache__/pipeline_chroma.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__pycache__/pipeline_chroma.cpython-310.pyc deleted file mode 100644 index 1872707e8ec17e2f1adbf4d940b05499bd92d050..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__pycache__/pipeline_chroma.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__pycache__/pipeline_chroma_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__pycache__/pipeline_chroma_img2img.cpython-310.pyc deleted file mode 100644 index 5e748b4302a621b0f55af379a21079fc225bd42d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__pycache__/pipeline_chroma_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index e29d17417709bddb41526657d9906e88f0d9e681..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/pipeline_chroma.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/pipeline_chroma.py deleted file mode 100644 index c74834ee82525d607688b331be416b548b970530..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/pipeline_chroma.py +++ /dev/null @@ -1,949 +0,0 @@ -# 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 AutoencoderKL, 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 ..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 ChromaPipeline - - >>> model_id = "lodestones/Chroma" - >>> ckpt_path = "https://huggingface.co/lodestones/Chroma/blob/main/chroma-unlocked-v37.safetensors" - >>> transformer = ChromaTransformer2DModel.from_single_file(ckpt_path, torch_dtype=torch.bfloat16) - >>> pipe = ChromaPipeline.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 ChromaPipeline( - DiffusionPipeline, - FluxLoraLoaderMixin, - FromSingleFileMixin, - TextualInversionLoaderMixin, - FluxIPAdapterMixin, -): - r""" - The Chroma pipeline for text-to-image generation. - - Reference: https://huggingface.co/lodestones/Chroma/ - - Args: - transformer ([`ChromaTransformer2DModel`]): - 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 ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representation - 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`): - Second Tokenizer of class - [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast). - """ - - model_cpu_offload_seq = "text_encoder->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 - # 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 - - 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 - attention_mask = text_inputs.attention_mask.clone() - - # Chroma requires the attention mask to include one padding token - seq_lengths = attention_mask.sum(dim=1) - mask_indices = torch.arange(attention_mask.size(1)).unsqueeze(0).expand(batch_size, -1) - attention_mask = (mask_indices <= seq_lengths.unsqueeze(1)).long() - - prompt_embeds = self.text_encoder( - text_input_ids.to(device), output_hidden_states=False, attention_mask=attention_mask.to(device) - )[0] - - dtype = self.text_encoder.dtype - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - attention_mask = attention_mask.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 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 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) - - @staticmethod - 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 - 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 enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.prepare_latents - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=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_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 // 2, width // 2, 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)], - dim=1, - ) - attention_mask = attention_mask.to(dtype) - - 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 * 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. - guidance_scale (`float`, *optional*, defaults to 3.5): - 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. - 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 ge 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 * 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, - 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 // 4 - latents, latent_image_ids = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 5. Prepare timesteps - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas - image_seq_len = latents.shape[1] - 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 - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/pipeline_chroma_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/pipeline_chroma_img2img.py deleted file mode 100644 index 9936608aaf211d9244c61e3056a255faef583cf6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/pipeline_chroma_img2img.py +++ /dev/null @@ -1,1034 +0,0 @@ -# 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 AutoencoderKL, 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 ..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 ChromaTransformer2DModel, ChromaImg2ImgPipeline - - >>> model_id = "lodestones/Chroma" - >>> ckpt_path = "https://huggingface.co/lodestones/Chroma/blob/main/chroma-unlocked-v37.safetensors" - >>> pipe = ChromaImg2ImgPipeline.from_pretrained( - ... model_id, - ... transformer=transformer, - ... torch_dtype=torch.bfloat16, - ... ) - >>> pipe.enable_model_cpu_offload() - >>> init_image = load_image( - ... "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" - ... ) - >>> prompt = "a scenic fastasy landscape with a river and mountains in the background, vibrant colors, detailed, high resolution" - >>> negative_prompt = "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors" - >>> image = pipe(prompt, image=init_image, negative_prompt=negative_prompt).images[0] - >>> image.save("chroma-img2img.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: Optional[torch.Generator] = 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 ChromaImg2ImgPipeline( - DiffusionPipeline, - FluxLoraLoaderMixin, - FromSingleFileMixin, - TextualInversionLoaderMixin, - FluxIPAdapterMixin, -): - r""" - The Chroma pipeline for image-to-image generation. - - Reference: https://huggingface.co/lodestones/Chroma/ - - Args: - transformer ([`ChromaTransformer2DModel`]): - 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 ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representation - 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`): - Second Tokenizer of class - [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast). - """ - - model_cpu_offload_seq = "text_encoder->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 - - 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 - attention_mask = text_inputs.attention_mask.clone() - - # Chroma requires the attention mask to include one padding token - seq_lengths = attention_mask.sum(dim=1) - mask_indices = torch.arange(attention_mask.size(1)).unsqueeze(0).expand(batch_size, -1) - attention_mask = (mask_indices <= seq_lengths.unsqueeze(1)).long() - - prompt_embeds = self.text_encoder( - text_input_ids.to(device), output_hidden_states=False, attention_mask=attention_mask.to(device) - )[0] - - dtype = self.text_encoder.dtype - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - attention_mask = attention_mask.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 - - # 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 - - 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 - - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt - ): - device = device or self._execution_device - - 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, - strength, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=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 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(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 _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 - 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 enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - # 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_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(height // 2, width // 2, device, dtype) - - if latents is not None: - return latents.to(device=device, dtype=dtype), latent_image_ids - - 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) - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = self.scheduler.scale_noise(image_latents, timestep, noise) - latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) - 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)], - dim=1, - ) - attention_mask = attention_mask.to(dtype) - - 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, - image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 35, - sigmas: Optional[List[float]] = 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", - 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 * 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 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. - 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 ge 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 * 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, - height, - width, - 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, - ) - - self._guidance_scale = guidance_scale - self._joint_attention_kwargs = joint_attention_kwargs - self._current_timestep = None - self._interrupt = False - - # 2. Preprocess image - init_image = self.image_processor.preprocess(image, height=height, width=width) - 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 - 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, - ) - - 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. Prepare image embeddings - 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 - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latents.shape[0]) - - 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] - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/pipeline_output.py deleted file mode 100644 index 951d132dba2eab63a08808f347baa4121ca4dc0b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/chroma/pipeline_output.py +++ /dev/null @@ -1,21 +0,0 @@ -from dataclasses import dataclass -from typing import List, Union - -import numpy as np -import PIL.Image - -from ...utils import BaseOutput - - -@dataclass -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, - num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__init__.py deleted file mode 100644 index e4fa1dda53d3f6500d408a2adc199a27abfeb8c3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__init__.py +++ /dev/null @@ -1,54 +0,0 @@ -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_cogvideox"] = ["CogVideoXPipeline"] - _import_structure["pipeline_cogvideox_fun_control"] = ["CogVideoXFunControlPipeline"] - _import_structure["pipeline_cogvideox_image2video"] = ["CogVideoXImageToVideoPipeline"] - _import_structure["pipeline_cogvideox_video2video"] = ["CogVideoXVideoToVideoPipeline"] - -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_cogvideox import CogVideoXPipeline - from .pipeline_cogvideox_fun_control import CogVideoXFunControlPipeline - from .pipeline_cogvideox_image2video import CogVideoXImageToVideoPipeline - from .pipeline_cogvideox_video2video import CogVideoXVideoToVideoPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index bf0133e739acdae329d90a16fc92f2a2afe0d626..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox.cpython-310.pyc deleted file mode 100644 index b1448a8739cfbf22e8c0dca7442b875ac7a437b7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox_fun_control.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox_fun_control.cpython-310.pyc deleted file mode 100644 index 0a0d19fc03167a709897adaef7c0858ad989bff4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox_fun_control.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox_image2video.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox_image2video.cpython-310.pyc deleted file mode 100644 index 4d364da7c524831e9ffeb8fe2cd4f0cb888b65e5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox_image2video.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox_video2video.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox_video2video.cpython-310.pyc deleted file mode 100644 index db01ef0b9d7ee3654fce27ed1b7f4d222681833f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_cogvideox_video2video.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index d83fb96b6ebfa899b9e5c4618c85056857c64bef..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_cogvideox.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_cogvideox.py deleted file mode 100644 index f08a3c35c2fa792f495c4fd633e84757c52deaf4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_cogvideox.py +++ /dev/null @@ -1,788 +0,0 @@ -# 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 math -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import T5EncoderModel, T5Tokenizer - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...loaders import CogVideoXLoraLoaderMixin -from ...models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel -from ...models.embeddings import get_3d_rotary_pos_embed -from ...pipelines.pipeline_utils import DiffusionPipeline -from ...schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler -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_output import CogVideoXPipelineOutput - - -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 CogVideoXPipeline - >>> from diffusers.utils import export_to_video - - >>> # Models: "THUDM/CogVideoX-2b" or "THUDM/CogVideoX-5b" - >>> pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-2b", torch_dtype=torch.float16).to("cuda") - >>> prompt = ( - ... "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. " - ... "The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other " - ... "pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, " - ... "casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. " - ... "The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical " - ... "atmosphere of this unique musical performance." - ... ) - >>> video = pipe(prompt=prompt, guidance_scale=6, num_inference_steps=50).frames[0] - >>> export_to_video(video, "output.mp4", fps=8) - ``` -""" - - -# Similar to diffusers.pipelines.hunyuandit.pipeline_hunyuandit.get_resize_crop_region_for_grid -def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): - tw = tgt_width - th = tgt_height - h, w = src - r = h / w - if r > (th / tw): - resize_height = th - resize_width = int(round(th / h * w)) - else: - resize_width = tw - resize_height = int(round(tw / w * h)) - - crop_top = int(round((th - resize_height) / 2.0)) - crop_left = int(round((tw - resize_width) / 2.0)) - - return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_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, - **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 CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin): - r""" - Pipeline for text-to-video generation using CogVideoX. - - 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.) - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. - text_encoder ([`T5EncoderModel`]): - Frozen text-encoder. CogVideoX uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the - [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant. - tokenizer (`T5Tokenizer`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - transformer ([`CogVideoXTransformer3DModel`]): - A text conditioned `CogVideoXTransformer3DModel` to denoise the encoded video latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `transformer` to denoise the encoded video latents. - """ - - _optional_components = [] - model_cpu_offload_seq = "text_encoder->transformer->vae" - - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - ] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - vae: AutoencoderKLCogVideoX, - transformer: CogVideoXTransformer3DModel, - scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler], - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler - ) - self.vae_scale_factor_spatial = ( - 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - ) - self.vae_scale_factor_temporal = ( - self.vae.config.temporal_compression_ratio if getattr(self, "vae", None) else 4 - ) - self.vae_scaling_factor_image = self.vae.config.scaling_factor if getattr(self, "vae", None) else 0.7 - - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - - def _get_t5_prompt_embeds( - self, - prompt: Union[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 = 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 - 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) - - return prompt_embeds - - 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: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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 prepare_latents( - self, batch_size, num_channels_latents, num_frames, 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." - ) - - shape = ( - batch_size, - (num_frames - 1) // self.vae_scale_factor_temporal + 1, - num_channels_latents, - height // self.vae_scale_factor_spatial, - width // self.vae_scale_factor_spatial, - ) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def decode_latents(self, latents: torch.Tensor) -> torch.Tensor: - latents = latents.permute(0, 2, 1, 3, 4) # [batch_size, num_channels, num_frames, height, width] - latents = 1 / self.vae_scaling_factor_image * latents - - frames = self.vae.decode(latents).sample - return frames - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.latte.pipeline_latte.LattePipeline.check_inputs - def check_inputs( - self, - prompt, - height, - width, - negative_prompt, - callback_on_step_end_tensor_inputs, - prompt_embeds=None, - negative_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 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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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 fuse_qkv_projections(self) -> None: - r"""Enables fused QKV projections.""" - self.fusing_transformer = True - self.transformer.fuse_qkv_projections() - - def unfuse_qkv_projections(self) -> None: - r"""Disable QKV projection fusion if enabled.""" - if not self.fusing_transformer: - logger.warning("The Transformer was not initially fused for QKV projections. Doing nothing.") - else: - self.transformer.unfuse_qkv_projections() - self.fusing_transformer = False - - def _prepare_rotary_positional_embeddings( - self, - height: int, - width: int, - num_frames: int, - device: torch.device, - ) -> 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) - - p = self.transformer.config.patch_size - p_t = self.transformer.config.patch_size_t - - base_size_width = self.transformer.config.sample_width // p - base_size_height = self.transformer.config.sample_height // p - - if p_t is None: - # CogVideoX 1.0 - grid_crops_coords = get_resize_crop_region_for_grid( - (grid_height, grid_width), base_size_width, base_size_height - ) - freqs_cos, freqs_sin = get_3d_rotary_pos_embed( - embed_dim=self.transformer.config.attention_head_dim, - crops_coords=grid_crops_coords, - grid_size=(grid_height, grid_width), - temporal_size=num_frames, - device=device, - ) - else: - # CogVideoX 1.5 - base_num_frames = (num_frames + p_t - 1) // p_t - - freqs_cos, freqs_sin = get_3d_rotary_pos_embed( - embed_dim=self.transformer.config.attention_head_dim, - crops_coords=None, - grid_size=(grid_height, grid_width), - temporal_size=base_num_frames, - grid_type="slice", - max_size=(base_size_height, base_size_width), - device=device, - ) - - return freqs_cos, freqs_sin - - @property - def guidance_scale(self): - return self._guidance_scale - - @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: 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, - num_inference_steps: int = 50, - timesteps: Optional[List[int]] = 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, - 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"], - max_sequence_length: int = 226, - ) -> Union[CogVideoXPipelineOutput, 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. 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 - less than `1`). - height (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial): - The height in pixels of the generated image. This is set to 480 by default for the best results. - width (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial): - The width in pixels of the generated image. This is set to 720 by default for the best results. - num_frames (`int`, defaults to `48`): - Number of frames to generate. Must be divisible by self.vae_scale_factor_temporal. Generated video will - contain 1 extra frame because CogVideoX is conditioned with (num_seconds * fps + 1) frames where - num_seconds is 6 and fps is 8. However, since videos can be saved at any fps, the only condition that - needs to be satisfied is that of divisibility mentioned above. - 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 7.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 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.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. - 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 `226`): - Maximum sequence length in encoded prompt. Must be consistent with - `self.transformer.config.max_text_seq_length` otherwise may lead to poor results. - - Examples: - - Returns: - [`~pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipelineOutput`] or `tuple`: - [`~pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipelineOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, 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 - - height = height or self.transformer.config.sample_height * self.vae_scale_factor_spatial - width = width or self.transformer.config.sample_width * self.vae_scale_factor_spatial - num_frames = num_frames or self.transformer.config.sample_frames - - num_videos_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - negative_prompt, - callback_on_step_end_tensor_inputs, - prompt_embeds, - negative_prompt_embeds, - ) - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._current_timestep = None - self._interrupt = False - - # 2. Default 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - negative_prompt, - 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, - ) - if do_classifier_free_guidance: - 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) - self._num_timesteps = len(timesteps) - - # 5. Prepare latents - latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - - # For CogVideoX 1.5, the latent frames should be padded to make it divisible by patch_size_t - patch_size_t = self.transformer.config.patch_size_t - additional_frames = 0 - if patch_size_t is not None and latent_frames % patch_size_t != 0: - additional_frames = patch_size_t - latent_frames % patch_size_t - num_frames += additional_frames * self.vae_scale_factor_temporal - - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - latent_channels, - num_frames, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Create rotary embeds if required - image_rotary_emb = ( - self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device) - if self.transformer.config.use_rotary_positional_embeddings - else None - ) - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - # for DPM-solver++ - old_pred_original_sample = None - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - self._current_timestep = t - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - # predict noise model_output - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=timestep, - image_rotary_emb=image_rotary_emb, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - noise_pred = noise_pred.float() - - # perform guidance - if use_dynamic_cfg: - self._guidance_scale = 1 + guidance_scale * ( - (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2 - ) - if 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) - - # compute the previous noisy sample x_t -> x_t-1 - if not isinstance(self.scheduler, CogVideoXDPMScheduler): - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - else: - latents, old_pred_original_sample = self.scheduler.step( - noise_pred, - old_pred_original_sample, - t, - timesteps[i - 1] if i > 0 else None, - latents, - **extra_step_kwargs, - return_dict=False, - ) - latents = latents.to(prompt_embeds.dtype) - - # 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 == 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 not output_type == "latent": - # Discard any padding frames that were added for CogVideoX 1.5 - latents = latents[:, additional_frames:] - video = self.decode_latents(latents) - video = self.video_processor.postprocess_video(video=video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return CogVideoXPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py deleted file mode 100644 index fe3e8ae388c7b98c9661aa328f6196990394edc9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py +++ /dev/null @@ -1,841 +0,0 @@ -# Copyright 2025 The CogVideoX team, Tsinghua University & ZhipuAI, Alibaba-PAI 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 math -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from PIL import Image -from transformers import T5EncoderModel, T5Tokenizer - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...loaders import CogVideoXLoraLoaderMixin -from ...models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel -from ...models.embeddings import get_3d_rotary_pos_embed -from ...pipelines.pipeline_utils import DiffusionPipeline -from ...schedulers import KarrasDiffusionSchedulers -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_output import CogVideoXPipelineOutput - - -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 CogVideoXFunControlPipeline, DDIMScheduler - >>> from diffusers.utils import export_to_video, load_video - - >>> pipe = CogVideoXFunControlPipeline.from_pretrained( - ... "alibaba-pai/CogVideoX-Fun-V1.1-5b-Pose", torch_dtype=torch.bfloat16 - ... ) - >>> pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) - >>> pipe.to("cuda") - - >>> control_video = load_video( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/hiker.mp4" - ... ) - >>> prompt = ( - ... "An astronaut stands triumphantly at the peak of a towering mountain. Panorama of rugged peaks and " - ... "valleys. Very futuristic vibe and animated aesthetic. Highlights of purple and golden colors in " - ... "the scene. The sky is looks like an animated/cartoonish dream of galaxies, nebulae, stars, planets, " - ... "moons, but the remainder of the scene is mostly realistic." - ... ) - - >>> video = pipe(prompt=prompt, control_video=control_video).frames[0] - >>> export_to_video(video, "output.mp4", fps=8) - ``` -""" - - -# Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.get_resize_crop_region_for_grid -def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): - tw = tgt_width - th = tgt_height - h, w = src - r = h / w - if r > (th / tw): - resize_height = th - resize_width = int(round(th / h * w)) - else: - resize_width = tw - resize_height = int(round(tw / w * h)) - - crop_top = int(round((th - resize_height) / 2.0)) - crop_left = int(round((tw - resize_width) / 2.0)) - - return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_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, - **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 CogVideoXFunControlPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin): - r""" - Pipeline for controlled text-to-video generation using CogVideoX Fun. - - 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.) - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. - text_encoder ([`T5EncoderModel`]): - Frozen text-encoder. CogVideoX uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the - [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant. - tokenizer (`T5Tokenizer`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - transformer ([`CogVideoXTransformer3DModel`]): - A text conditioned `CogVideoXTransformer3DModel` to denoise the encoded video latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `transformer` to denoise the encoded video latents. - """ - - _optional_components = [] - model_cpu_offload_seq = "text_encoder->vae->transformer->vae" - - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - ] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - vae: AutoencoderKLCogVideoX, - transformer: CogVideoXTransformer3DModel, - scheduler: KarrasDiffusionSchedulers, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler - ) - self.vae_scale_factor_spatial = ( - 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - ) - self.vae_scale_factor_temporal = ( - self.vae.config.temporal_compression_ratio if getattr(self, "vae", None) else 4 - ) - self.vae_scaling_factor_image = self.vae.config.scaling_factor if getattr(self, "vae", None) else 0.7 - - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds - def _get_t5_prompt_embeds( - self, - prompt: Union[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 = 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 - 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) - - return 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, - 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, - max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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 - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.prepare_latents - def prepare_latents( - self, batch_size, num_channels_latents, num_frames, 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." - ) - - shape = ( - batch_size, - (num_frames - 1) // self.vae_scale_factor_temporal + 1, - num_channels_latents, - height // self.vae_scale_factor_spatial, - width // self.vae_scale_factor_spatial, - ) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return 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]: - if mask is not None: - masks = [] - for i in range(mask.size(0)): - current_mask = mask[i].unsqueeze(0) - current_mask = self.vae.encode(current_mask)[0] - current_mask = current_mask.mode() - masks.append(current_mask) - mask = torch.cat(masks, dim=0) - mask = mask * self.vae.config.scaling_factor - - if masked_image is not None: - mask_pixel_values = [] - for i in range(masked_image.size(0)): - mask_pixel_value = masked_image[i].unsqueeze(0) - mask_pixel_value = self.vae.encode(mask_pixel_value)[0] - mask_pixel_value = mask_pixel_value.mode() - mask_pixel_values.append(mask_pixel_value) - masked_image_latents = torch.cat(mask_pixel_values, dim=0) - masked_image_latents = masked_image_latents * self.vae.config.scaling_factor - else: - masked_image_latents = None - - return mask, masked_image_latents - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.decode_latents - def decode_latents(self, latents: torch.Tensor) -> torch.Tensor: - latents = latents.permute(0, 2, 1, 3, 4) # [batch_size, num_channels, num_frames, height, width] - latents = 1 / self.vae_scaling_factor_image * latents - - frames = self.vae.decode(latents).sample - return frames - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - negative_prompt, - callback_on_step_end_tensor_inputs, - prompt_embeds=None, - negative_prompt_embeds=None, - control_video=None, - control_video_latents=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_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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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 control_video is not None and control_video_latents is not None: - raise ValueError( - "Cannot pass both `control_video` and `control_video_latents`. Please make sure to pass only one of these parameters." - ) - - def fuse_qkv_projections(self) -> None: - r"""Enables fused QKV projections.""" - self.fusing_transformer = True - self.transformer.fuse_qkv_projections() - - def unfuse_qkv_projections(self) -> None: - r"""Disable QKV projection fusion if enabled.""" - if not self.fusing_transformer: - logger.warning("The Transformer was not initially fused for QKV projections. Doing nothing.") - else: - self.transformer.unfuse_qkv_projections() - self.fusing_transformer = False - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._prepare_rotary_positional_embeddings - def _prepare_rotary_positional_embeddings( - self, - height: int, - width: int, - num_frames: int, - device: torch.device, - ) -> 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) - - p = self.transformer.config.patch_size - p_t = self.transformer.config.patch_size_t - - base_size_width = self.transformer.config.sample_width // p - base_size_height = self.transformer.config.sample_height // p - - if p_t is None: - # CogVideoX 1.0 - grid_crops_coords = get_resize_crop_region_for_grid( - (grid_height, grid_width), base_size_width, base_size_height - ) - freqs_cos, freqs_sin = get_3d_rotary_pos_embed( - embed_dim=self.transformer.config.attention_head_dim, - crops_coords=grid_crops_coords, - grid_size=(grid_height, grid_width), - temporal_size=num_frames, - device=device, - ) - else: - # CogVideoX 1.5 - base_num_frames = (num_frames + p_t - 1) // p_t - - freqs_cos, freqs_sin = get_3d_rotary_pos_embed( - embed_dim=self.transformer.config.attention_head_dim, - crops_coords=None, - grid_size=(grid_height, grid_width), - temporal_size=base_num_frames, - grid_type="slice", - max_size=(base_size_height, base_size_width), - device=device, - ) - - return freqs_cos, freqs_sin - - @property - def guidance_scale(self): - return self._guidance_scale - - @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: 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, - num_inference_steps: int = 50, - timesteps: Optional[List[int]] = 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, - 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"], - max_sequence_length: int = 226, - ) -> Union[CogVideoXPipelineOutput, 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. 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 - less than `1`). - 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): - The height in pixels of the generated image. This is set to 480 by default for the best results. - width (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial): - The width in pixels of the generated image. This is set to 720 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 6.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 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 ge generated by sampling using the supplied random `generator`. - control_video_latents (`torch.Tensor`, *optional*): - Pre-generated control latents, sampled from a Gaussian distribution, to be used as inputs for - controlled video generation. If not provided, `control_video` must be provided. - 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_xl.StableDiffusionXLPipelineOutput`] 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 `226`): - Maximum sequence length in encoded prompt. Must be consistent with - `self.transformer.config.max_text_seq_length` otherwise may lead to poor results. - - Examples: - - Returns: - [`~pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipelineOutput`] or `tuple`: - [`~pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipelineOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, 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 - - if control_video is not None and isinstance(control_video[0], Image.Image): - control_video = [control_video] - - height = height or self.transformer.config.sample_height * self.vae_scale_factor_spatial - width = width or self.transformer.config.sample_width * self.vae_scale_factor_spatial - num_frames = len(control_video[0]) if control_video is not None else control_video_latents.size(2) - - num_videos_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - negative_prompt, - callback_on_step_end_tensor_inputs, - prompt_embeds, - negative_prompt_embeds, - control_video, - control_video_latents, - ) - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._current_timestep = None - self._interrupt = False - - # 2. Default 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - negative_prompt, - 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, - ) - if do_classifier_free_guidance: - 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) - self._num_timesteps = len(timesteps) - - # 5. Prepare latents - latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - - # For CogVideoX 1.5, the latent frames should be padded to make it divisible by patch_size_t - patch_size_t = self.transformer.config.patch_size_t - if patch_size_t is not None and latent_frames % patch_size_t != 0: - raise ValueError( - f"The number of latent frames must be divisible by `{patch_size_t=}` but the given video " - f"contains {latent_frames=}, which is not divisible." - ) - - latent_channels = self.transformer.config.in_channels // 2 - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - latent_channels, - num_frames, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - if control_video_latents is None: - control_video = self.video_processor.preprocess_video(control_video, height=height, width=width) - control_video = control_video.to(device=device, dtype=prompt_embeds.dtype) - - _, control_video_latents = self.prepare_control_latents(None, control_video) - control_video_latents = control_video_latents.permute(0, 2, 1, 3, 4) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Create rotary embeds if required - image_rotary_emb = ( - self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device) - if self.transformer.config.use_rotary_positional_embeddings - else None - ) - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - # for DPM-solver++ - old_pred_original_sample = None - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - self._current_timestep = t - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - latent_control_input = ( - torch.cat([control_video_latents] * 2) if do_classifier_free_guidance else control_video_latents - ) - latent_model_input = torch.cat([latent_model_input, latent_control_input], dim=2) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - # predict noise model_output - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=timestep, - image_rotary_emb=image_rotary_emb, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - noise_pred = noise_pred.float() - - # perform guidance - if use_dynamic_cfg: - self._guidance_scale = 1 + guidance_scale * ( - (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2 - ) - if 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) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - latents = latents.to(prompt_embeds.dtype) - - # 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 == 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 not output_type == "latent": - video = self.decode_latents(latents) - video = self.video_processor.postprocess_video(video=video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return CogVideoXPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py deleted file mode 100644 index a982f4b275574a2d9049c9b6baa25101036b34a5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py +++ /dev/null @@ -1,902 +0,0 @@ -# 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 math -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import PIL -import torch -from transformers import T5EncoderModel, T5Tokenizer - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput -from ...loaders import CogVideoXLoraLoaderMixin -from ...models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel -from ...models.embeddings import get_3d_rotary_pos_embed -from ...pipelines.pipeline_utils import DiffusionPipeline -from ...schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler -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_output import CogVideoXPipelineOutput - - -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 CogVideoXImageToVideoPipeline - >>> from diffusers.utils import export_to_video, load_image - - >>> pipe = CogVideoXImageToVideoPipeline.from_pretrained("THUDM/CogVideoX-5b-I2V", torch_dtype=torch.bfloat16) - >>> pipe.to("cuda") - - >>> prompt = "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot." - >>> image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" - ... ) - >>> video = pipe(image, prompt, use_dynamic_cfg=True) - >>> export_to_video(video.frames[0], "output.mp4", fps=8) - ``` -""" - - -# Similar to diffusers.pipelines.hunyuandit.pipeline_hunyuandit.get_resize_crop_region_for_grid -def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): - tw = tgt_width - th = tgt_height - h, w = src - r = h / w - if r > (th / tw): - resize_height = th - resize_width = int(round(th / h * w)) - else: - resize_width = tw - resize_height = int(round(tw / w * h)) - - crop_top = int(round((th - resize_height) / 2.0)) - crop_left = int(round((tw - resize_width) / 2.0)) - - return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_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, - **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: Optional[torch.Generator] = 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 CogVideoXImageToVideoPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin): - r""" - Pipeline for image-to-video generation using CogVideoX. - - 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.) - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. - text_encoder ([`T5EncoderModel`]): - Frozen text-encoder. CogVideoX uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the - [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant. - tokenizer (`T5Tokenizer`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - transformer ([`CogVideoXTransformer3DModel`]): - A text conditioned `CogVideoXTransformer3DModel` to denoise the encoded video latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `transformer` to denoise the encoded video latents. - """ - - _optional_components = [] - model_cpu_offload_seq = "text_encoder->transformer->vae" - - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - ] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - vae: AutoencoderKLCogVideoX, - transformer: CogVideoXTransformer3DModel, - scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler], - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - vae=vae, - transformer=transformer, - scheduler=scheduler, - ) - self.vae_scale_factor_spatial = ( - 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - ) - self.vae_scale_factor_temporal = ( - self.vae.config.temporal_compression_ratio if getattr(self, "vae", None) else 4 - ) - self.vae_scaling_factor_image = self.vae.config.scaling_factor if getattr(self, "vae", None) else 0.7 - - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds - def _get_t5_prompt_embeds( - self, - prompt: Union[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 = 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 - 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) - - return 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, - 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, - max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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 prepare_latents( - self, - image: torch.Tensor, - batch_size: int = 1, - num_channels_latents: int = 16, - 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, - ): - 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." - ) - - num_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - shape = ( - batch_size, - num_frames, - num_channels_latents, - height // self.vae_scale_factor_spatial, - width // self.vae_scale_factor_spatial, - ) - - # For CogVideoX1.5, the latent should add 1 for padding (Not use) - if self.transformer.config.patch_size_t is not None: - shape = shape[:1] + (shape[1] + shape[1] % self.transformer.config.patch_size_t,) + shape[2:] - - image = image.unsqueeze(2) # [B, C, F, H, W] - - if isinstance(generator, list): - image_latents = [ - retrieve_latents(self.vae.encode(image[i].unsqueeze(0)), generator[i]) for i in range(batch_size) - ] - else: - image_latents = [retrieve_latents(self.vae.encode(img.unsqueeze(0)), generator) for img in image] - - image_latents = torch.cat(image_latents, dim=0).to(dtype).permute(0, 2, 1, 3, 4) # [B, F, C, H, W] - - if not self.vae.config.invert_scale_latents: - image_latents = self.vae_scaling_factor_image * image_latents - else: - # This is awkward but required because the CogVideoX team forgot to multiply the - # scaling factor during training :) - image_latents = 1 / self.vae_scaling_factor_image * image_latents - - padding_shape = ( - batch_size, - num_frames - 1, - num_channels_latents, - height // self.vae_scale_factor_spatial, - width // self.vae_scale_factor_spatial, - ) - - latent_padding = torch.zeros(padding_shape, device=device, dtype=dtype) - image_latents = torch.cat([image_latents, latent_padding], dim=1) - - # Select the first frame along the second dimension - if self.transformer.config.patch_size_t is not None: - first_frame = image_latents[:, : image_latents.size(1) % self.transformer.config.patch_size_t, ...] - image_latents = torch.cat([first_frame, image_latents], dim=1) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents, image_latents - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.decode_latents - def decode_latents(self, latents: torch.Tensor) -> torch.Tensor: - latents = latents.permute(0, 2, 1, 3, 4) # [batch_size, num_channels, num_frames, height, width] - latents = 1 / self.vae_scaling_factor_image * latents - - frames = self.vae.decode(latents).sample - return frames - - # Copied from diffusers.pipelines.animatediff.pipeline_animatediff_video2video.AnimateDiffVideoToVideoPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, timesteps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - timesteps = timesteps[t_start * self.scheduler.order :] - - return timesteps, num_inference_steps - t_start - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - image, - prompt, - height, - width, - negative_prompt, - callback_on_step_end_tensor_inputs, - latents=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - if ( - not isinstance(image, torch.Tensor) - and not isinstance(image, PIL.Image.Image) - 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" - f" {type(image)}" - ) - - 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_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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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}." - ) - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.fuse_qkv_projections - def fuse_qkv_projections(self) -> None: - r"""Enables fused QKV projections.""" - self.fusing_transformer = True - self.transformer.fuse_qkv_projections() - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.unfuse_qkv_projections - def unfuse_qkv_projections(self) -> None: - r"""Disable QKV projection fusion if enabled.""" - if not self.fusing_transformer: - logger.warning("The Transformer was not initially fused for QKV projections. Doing nothing.") - else: - self.transformer.unfuse_qkv_projections() - self.fusing_transformer = False - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._prepare_rotary_positional_embeddings - def _prepare_rotary_positional_embeddings( - self, - height: int, - width: int, - num_frames: int, - device: torch.device, - ) -> 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) - - p = self.transformer.config.patch_size - p_t = self.transformer.config.patch_size_t - - base_size_width = self.transformer.config.sample_width // p - base_size_height = self.transformer.config.sample_height // p - - if p_t is None: - # CogVideoX 1.0 - grid_crops_coords = get_resize_crop_region_for_grid( - (grid_height, grid_width), base_size_width, base_size_height - ) - freqs_cos, freqs_sin = get_3d_rotary_pos_embed( - embed_dim=self.transformer.config.attention_head_dim, - crops_coords=grid_crops_coords, - grid_size=(grid_height, grid_width), - temporal_size=num_frames, - device=device, - ) - else: - # CogVideoX 1.5 - base_num_frames = (num_frames + p_t - 1) // p_t - - freqs_cos, freqs_sin = get_3d_rotary_pos_embed( - embed_dim=self.transformer.config.attention_head_dim, - crops_coords=None, - grid_size=(grid_height, grid_width), - temporal_size=base_num_frames, - grid_type="slice", - max_size=(base_size_height, base_size_width), - device=device, - ) - - return freqs_cos, freqs_sin - - @property - def guidance_scale(self): - return self._guidance_scale - - @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, - 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, - num_frames: int = 49, - num_inference_steps: int = 50, - timesteps: Optional[List[int]] = 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, - 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"], - max_sequence_length: int = 226, - ) -> Union[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*): - 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 - less than `1`). - height (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial): - The height in pixels of the generated image. This is set to 480 by default for the best results. - width (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial): - The width in pixels of the generated image. This is set to 720 by default for the best results. - num_frames (`int`, defaults to `48`): - Number of frames to generate. Must be divisible by self.vae_scale_factor_temporal. Generated video will - contain 1 extra frame because CogVideoX is conditioned with (num_seconds * fps + 1) frames where - num_seconds is 6 and fps is 8. However, since videos can be saved at any fps, the only condition that - needs to be satisfied is that of divisibility mentioned above. - 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 7.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 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.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. - 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 `226`): - Maximum sequence length in encoded prompt. Must be consistent with - `self.transformer.config.max_text_seq_length` otherwise may lead to poor results. - - Examples: - - Returns: - [`~pipelines.cogvideo.pipeline_output.CogVideoXPipelineOutput`] or `tuple`: - [`~pipelines.cogvideo.pipeline_output.CogVideoXPipelineOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, 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 - - height = height or self.transformer.config.sample_height * self.vae_scale_factor_spatial - width = width or self.transformer.config.sample_width * self.vae_scale_factor_spatial - num_frames = num_frames or self.transformer.config.sample_frames - - num_videos_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - image=image, - prompt=prompt, - height=height, - width=width, - negative_prompt=negative_prompt, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - latents=latents, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - self._guidance_scale = guidance_scale - self._current_timestep = None - self._attention_kwargs = attention_kwargs - self._interrupt = False - - # 2. Default 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=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, - ) - if do_classifier_free_guidance: - 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) - self._num_timesteps = len(timesteps) - - # 5. Prepare latents - latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - - # For CogVideoX 1.5, the latent frames should be padded to make it divisible by patch_size_t - patch_size_t = self.transformer.config.patch_size_t - additional_frames = 0 - if patch_size_t is not None and latent_frames % patch_size_t != 0: - additional_frames = patch_size_t - latent_frames % patch_size_t - num_frames += additional_frames * self.vae_scale_factor_temporal - - image = self.video_processor.preprocess(image, height=height, width=width).to( - device, dtype=prompt_embeds.dtype - ) - - latent_channels = self.transformer.config.in_channels // 2 - latents, image_latents = self.prepare_latents( - image, - batch_size * num_videos_per_prompt, - latent_channels, - num_frames, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Create rotary embeds if required - image_rotary_emb = ( - self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device) - if self.transformer.config.use_rotary_positional_embeddings - else None - ) - - # 8. Create ofs embeds if required - ofs_emb = None if self.transformer.config.ofs_embed_dim is None else latents.new_full((1,), fill_value=2.0) - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - # for DPM-solver++ - old_pred_original_sample = None - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - self._current_timestep = t - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - latent_image_input = torch.cat([image_latents] * 2) if do_classifier_free_guidance else image_latents - latent_model_input = torch.cat([latent_model_input, latent_image_input], dim=2) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - # predict noise model_output - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=timestep, - ofs=ofs_emb, - image_rotary_emb=image_rotary_emb, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - noise_pred = noise_pred.float() - - # perform guidance - if use_dynamic_cfg: - self._guidance_scale = 1 + guidance_scale * ( - (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2 - ) - if 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) - - # compute the previous noisy sample x_t -> x_t-1 - if not isinstance(self.scheduler, CogVideoXDPMScheduler): - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - else: - latents, old_pred_original_sample = self.scheduler.step( - noise_pred, - old_pred_original_sample, - t, - timesteps[i - 1] if i > 0 else None, - latents, - **extra_step_kwargs, - return_dict=False, - ) - latents = latents.to(prompt_embeds.dtype) - - # 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 == 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 not output_type == "latent": - # Discard any padding frames that were added for CogVideoX 1.5 - latents = latents[:, additional_frames:] - video = self.decode_latents(latents) - video = self.video_processor.postprocess_video(video=video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return CogVideoXPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py deleted file mode 100644 index 7c50bdcb7defafe35556f255b9f441365db160fe..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py +++ /dev/null @@ -1,867 +0,0 @@ -# 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 math -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from PIL import Image -from transformers import T5EncoderModel, T5Tokenizer - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...loaders import CogVideoXLoraLoaderMixin -from ...models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel -from ...models.embeddings import get_3d_rotary_pos_embed -from ...pipelines.pipeline_utils import DiffusionPipeline -from ...schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler -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_output import CogVideoXPipelineOutput - - -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 CogVideoXDPMScheduler, CogVideoXVideoToVideoPipeline - >>> from diffusers.utils import export_to_video, load_video - - >>> # Models: "THUDM/CogVideoX-2b" or "THUDM/CogVideoX-5b" - >>> pipe = CogVideoXVideoToVideoPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16) - >>> pipe.to("cuda") - >>> pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config) - - >>> input_video = load_video( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/hiker.mp4" - ... ) - >>> prompt = ( - ... "An astronaut stands triumphantly at the peak of a towering mountain. Panorama of rugged peaks and " - ... "valleys. Very futuristic vibe and animated aesthetic. Highlights of purple and golden colors in " - ... "the scene. The sky is looks like an animated/cartoonish dream of galaxies, nebulae, stars, planets, " - ... "moons, but the remainder of the scene is mostly realistic." - ... ) - - >>> video = pipe( - ... video=input_video, prompt=prompt, strength=0.8, guidance_scale=6, num_inference_steps=50 - ... ).frames[0] - >>> export_to_video(video, "output.mp4", fps=8) - ``` -""" - - -# Similar to diffusers.pipelines.hunyuandit.pipeline_hunyuandit.get_resize_crop_region_for_grid -def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): - tw = tgt_width - th = tgt_height - h, w = src - r = h / w - if r > (th / tw): - resize_height = th - resize_width = int(round(th / h * w)) - else: - resize_width = tw - resize_height = int(round(tw / w * h)) - - crop_top = int(round((th - resize_height) / 2.0)) - crop_left = int(round((tw - resize_width) / 2.0)) - - return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_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, - **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: Optional[torch.Generator] = 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 CogVideoXVideoToVideoPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin): - r""" - Pipeline for video-to-video generation using CogVideoX. - - 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.) - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. - text_encoder ([`T5EncoderModel`]): - Frozen text-encoder. CogVideoX uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the - [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant. - tokenizer (`T5Tokenizer`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - transformer ([`CogVideoXTransformer3DModel`]): - A text conditioned `CogVideoXTransformer3DModel` to denoise the encoded video latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `transformer` to denoise the encoded video latents. - """ - - _optional_components = [] - model_cpu_offload_seq = "text_encoder->transformer->vae" - - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - ] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - vae: AutoencoderKLCogVideoX, - transformer: CogVideoXTransformer3DModel, - scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler], - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler - ) - - self.vae_scale_factor_spatial = ( - 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - ) - self.vae_scale_factor_temporal = ( - self.vae.config.temporal_compression_ratio if getattr(self, "vae", None) else 4 - ) - self.vae_scaling_factor_image = self.vae.config.scaling_factor if getattr(self, "vae", None) else 0.7 - - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds - def _get_t5_prompt_embeds( - self, - prompt: Union[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 = 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 - 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) - - return 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, - 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, - max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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 prepare_latents( - self, - video: Optional[torch.Tensor] = 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, - ): - 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." - ) - - num_frames = (video.size(2) - 1) // self.vae_scale_factor_temporal + 1 if latents is None else latents.size(1) - - shape = ( - batch_size, - num_frames, - num_channels_latents, - height // self.vae_scale_factor_spatial, - width // self.vae_scale_factor_spatial, - ) - - if latents is None: - if isinstance(generator, list): - 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).permute(0, 2, 1, 3, 4) # [B, F, C, H, W] - init_latents = self.vae_scaling_factor_image * init_latents - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = self.scheduler.add_noise(init_latents, noise, timestep) - else: - latents = latents.to(device) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.decode_latents - def decode_latents(self, latents: torch.Tensor) -> torch.Tensor: - latents = latents.permute(0, 2, 1, 3, 4) # [batch_size, num_channels, num_frames, height, width] - latents = 1 / self.vae_scaling_factor_image * latents - - frames = self.vae.decode(latents).sample - return frames - - # Copied from diffusers.pipelines.animatediff.pipeline_animatediff_video2video.AnimateDiffVideoToVideoPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, timesteps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - timesteps = timesteps[t_start * self.scheduler.order :] - - return timesteps, num_inference_steps - t_start - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - strength, - negative_prompt, - callback_on_step_end_tensor_inputs, - video=None, - latents=None, - prompt_embeds=None, - negative_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 strength < 0 or strength > 1: - raise ValueError(f"The value of strength should 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 prompt is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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 video is not None and latents is not None: - raise ValueError("Only one of `video` or `latents` should be provided") - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.fuse_qkv_projections - def fuse_qkv_projections(self) -> None: - r"""Enables fused QKV projections.""" - self.fusing_transformer = True - self.transformer.fuse_qkv_projections() - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.unfuse_qkv_projections - def unfuse_qkv_projections(self) -> None: - r"""Disable QKV projection fusion if enabled.""" - if not self.fusing_transformer: - logger.warning("The Transformer was not initially fused for QKV projections. Doing nothing.") - else: - self.transformer.unfuse_qkv_projections() - self.fusing_transformer = False - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._prepare_rotary_positional_embeddings - def _prepare_rotary_positional_embeddings( - self, - height: int, - width: int, - num_frames: int, - device: torch.device, - ) -> 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) - - p = self.transformer.config.patch_size - p_t = self.transformer.config.patch_size_t - - base_size_width = self.transformer.config.sample_width // p - base_size_height = self.transformer.config.sample_height // p - - if p_t is None: - # CogVideoX 1.0 - grid_crops_coords = get_resize_crop_region_for_grid( - (grid_height, grid_width), base_size_width, base_size_height - ) - freqs_cos, freqs_sin = get_3d_rotary_pos_embed( - embed_dim=self.transformer.config.attention_head_dim, - crops_coords=grid_crops_coords, - grid_size=(grid_height, grid_width), - temporal_size=num_frames, - device=device, - ) - else: - # CogVideoX 1.5 - base_num_frames = (num_frames + p_t - 1) // p_t - - freqs_cos, freqs_sin = get_3d_rotary_pos_embed( - embed_dim=self.transformer.config.attention_head_dim, - crops_coords=None, - grid_size=(grid_height, grid_width), - temporal_size=base_num_frames, - grid_type="slice", - max_size=(base_size_height, base_size_width), - device=device, - ) - - return freqs_cos, freqs_sin - - @property - def guidance_scale(self): - return self._guidance_scale - - @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, - 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, - num_inference_steps: int = 50, - timesteps: Optional[List[int]] = 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, - 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"], - max_sequence_length: int = 226, - ) -> Union[CogVideoXPipelineOutput, Tuple]: - """ - Function invoked when calling the pipeline for generation. - - Args: - 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*): - 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 - less than `1`). - height (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial): - The height in pixels of the generated image. This is set to 480 by default for the best results. - width (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial): - The width in pixels of the generated image. This is set to 720 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. - strength (`float`, *optional*, defaults to 0.8): - Higher strength leads to more differences between original video and generated video. - guidance_scale (`float`, *optional*, defaults to 7.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 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.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. - 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 `226`): - Maximum sequence length in encoded prompt. Must be consistent with - `self.transformer.config.max_text_seq_length` otherwise may lead to poor results. - - Examples: - - Returns: - [`~pipelines.cogvideo.pipeline_output.CogVideoXPipelineOutput`] or `tuple`: - [`~pipelines.cogvideo.pipeline_output.CogVideoXPipelineOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, 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 - - height = height or self.transformer.config.sample_height * self.vae_scale_factor_spatial - width = width or self.transformer.config.sample_width * self.vae_scale_factor_spatial - num_frames = len(video) if latents is None else latents.size(1) - - num_videos_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt=prompt, - height=height, - width=width, - strength=strength, - negative_prompt=negative_prompt, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - video=video, - latents=latents, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._current_timestep = None - self._interrupt = False - - # 2. Default 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - negative_prompt, - 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, - ) - if do_classifier_free_guidance: - 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) - 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) - - # 5. Prepare latents - latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - - # For CogVideoX 1.5, the latent frames should be padded to make it divisible by patch_size_t - patch_size_t = self.transformer.config.patch_size_t - if patch_size_t is not None and latent_frames % patch_size_t != 0: - raise ValueError( - f"The number of latent frames must be divisible by `{patch_size_t=}` but the given video " - f"contains {latent_frames=}, which is not divisible." - ) - - if latents is None: - video = self.video_processor.preprocess_video(video, height=height, width=width) - video = video.to(device=device, dtype=prompt_embeds.dtype) - - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - video, - batch_size * num_videos_per_prompt, - latent_channels, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - latent_timestep, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Create rotary embeds if required - image_rotary_emb = ( - self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device) - if self.transformer.config.use_rotary_positional_embeddings - else None - ) - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - # for DPM-solver++ - old_pred_original_sample = None - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - self._current_timestep = t - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - # predict noise model_output - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=timestep, - image_rotary_emb=image_rotary_emb, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - noise_pred = noise_pred.float() - - # perform guidance - if use_dynamic_cfg: - self._guidance_scale = 1 + guidance_scale * ( - (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2 - ) - if 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) - - # compute the previous noisy sample x_t -> x_t-1 - if not isinstance(self.scheduler, CogVideoXDPMScheduler): - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - else: - latents, old_pred_original_sample = self.scheduler.step( - noise_pred, - old_pred_original_sample, - t, - timesteps[i - 1] if i > 0 else None, - latents, - **extra_step_kwargs, - return_dict=False, - ) - latents = latents.to(prompt_embeds.dtype) - - # 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 == 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 not output_type == "latent": - video = self.decode_latents(latents) - video = self.video_processor.postprocess_video(video=video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return CogVideoXPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_output.py deleted file mode 100644 index 3de030dd6928db49ab0bc4d11868a93ac98dea50..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogvideo/pipeline_output.py +++ /dev/null @@ -1,20 +0,0 @@ -from dataclasses import dataclass - -import torch - -from diffusers.utils import BaseOutput - - -@dataclass -class CogVideoXPipelineOutput(BaseOutput): - r""" - 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 - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/__init__.py deleted file mode 100644 index 50895251ba0b67dc2848b99654cbdd08312833f5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/__init__.py +++ /dev/null @@ -1,47 +0,0 @@ -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 = {} -_additional_imports = {} -_import_structure = {"pipeline_output": ["CogView3PlusPipelineOutput"]} - -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_cogview3plus"] = ["CogView3PlusPipeline"] -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_cogview3plus import CogView3PlusPipeline -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/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index bce80379ab7372537f9d51a40a8b41b4a89e1aac..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/__pycache__/pipeline_cogview3plus.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/__pycache__/pipeline_cogview3plus.cpython-310.pyc deleted file mode 100644 index 1be053e77252c540f46bd2a1a8f03ce687af6af7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/__pycache__/pipeline_cogview3plus.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index a5147517daf55ed7186ae24a95ad9acc1c14e075..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/pipeline_cogview3plus.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/pipeline_cogview3plus.py deleted file mode 100644 index f2f852c213ad6e3163528c7307d80634f0da0321..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/pipeline_cogview3plus.py +++ /dev/null @@ -1,682 +0,0 @@ -# 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 -from typing import Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import T5EncoderModel, T5Tokenizer - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import VaeImageProcessor -from ...models import AutoencoderKL, CogView3PlusTransformer2DModel -from ...pipelines.pipeline_utils import DiffusionPipeline -from ...schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler -from ...utils import is_torch_xla_available, logging, replace_example_docstring -from ...utils.torch_utils import randn_tensor -from .pipeline_output import CogView3PipelineOutput - - -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 CogView3PlusPipeline - - >>> pipe = CogView3PlusPipeline.from_pretrained("THUDM/CogView3-Plus-3B", 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") - ``` -""" - - -# 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 CogView3PlusPipeline(DiffusionPipeline): - r""" - Pipeline for text-to-image generation using CogView3Plus. - - 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.) - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`T5EncoderModel`]): - Frozen text-encoder. CogView3Plus uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the - [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant. - tokenizer (`T5Tokenizer`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - transformer ([`CogView3PlusTransformer2DModel`]): - A text conditioned `CogView3PlusTransformer2DModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - """ - - _optional_components = [] - model_cpu_offload_seq = "text_encoder->transformer->vae" - - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - ] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - vae: AutoencoderKL, - transformer: CogView3PlusTransformer2DModel, - scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler], - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_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) - - # 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, - num_images_per_prompt: int = 1, - max_sequence_length: int = 226, - 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) - - 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 - 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_images_per_prompt, 1) - prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - - return prompt_embeds - - def encode_prompt( - self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = 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, - max_sequence_length: int = 224, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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_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. - 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. - max_sequence_length (`int`, defaults to `224`): - Maximum sequence length in encoded prompt. Can be set to other values but may lead to poorer results. - 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_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - dtype=dtype, - ) - - if do_classifier_free_guidance and negative_prompt is None: - negative_prompt_embeds = prompt_embeds.new_zeros(prompt_embeds.shape) - - if do_classifier_free_guidance and negative_prompt_embeds is None: - 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_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - dtype=dtype, - ) - - return prompt_embeds, negative_prompt_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.latte.pipeline_latte.LattePipeline.check_inputs - def check_inputs( - self, - prompt, - height, - width, - negative_prompt, - callback_on_step_end_tensor_inputs, - prompt_embeds=None, - negative_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 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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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}." - ) - - @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 do_classifier_free_guidance(self): - return self._guidance_scale > 1 - - @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: Optional[Union[str, List[str]]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: Optional[List[int]] = 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), - 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"], - max_sequence_length: int = 224, - ) -> Union[CogView3PipelineOutput, 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. If not defined, one has to pass `prompt_embeds`. - 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`). - height (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor): - The height in pixels of the generated image. If not provided, it is set to 1024. - width (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor): - The width in pixels of the generated image. If not provided it is set to 1024. - 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. - 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. - 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). - 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. - 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 `224`): - Maximum sequence length in encoded prompt. Can be set to other values but may lead to poorer results. - - Examples: - - Returns: - [`~pipelines.cogview3.pipeline_cogview3plus.CogView3PipelineOutput`] or `tuple`: - [`~pipelines.cogview3.pipeline_cogview3plus.CogView3PipelineOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, 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 - - height = height or self.transformer.config.sample_size * self.vae_scale_factor - width = width or self.transformer.config.sample_size * self.vae_scale_factor - - original_size = original_size or (height, width) - target_size = (height, width) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - negative_prompt, - callback_on_step_end_tensor_inputs, - prompt_embeds, - negative_prompt_embeds, - ) - self._guidance_scale = guidance_scale - self._interrupt = False - - # 2. Default 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - negative_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, - ) - if self.do_classifier_free_guidance: - 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) - self._num_timesteps = len(timesteps) - - # 5. Prepare latents. - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Prepare additional timestep conditions - original_size = torch.tensor([original_size], dtype=prompt_embeds.dtype) - target_size = torch.tensor([target_size], dtype=prompt_embeds.dtype) - crops_coords_top_left = torch.tensor([crops_coords_top_left], dtype=prompt_embeds.dtype) - - if self.do_classifier_free_guidance: - original_size = torch.cat([original_size, original_size]) - target_size = torch.cat([target_size, target_size]) - crops_coords_top_left = torch.cat([crops_coords_top_left, crops_coords_top_left]) - - original_size = original_size.to(device).repeat(batch_size * num_images_per_prompt, 1) - target_size = target_size.to(device).repeat(batch_size * num_images_per_prompt, 1) - crops_coords_top_left = crops_coords_top_left.to(device).repeat(batch_size * num_images_per_prompt, 1) - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - # for DPM-solver++ - old_pred_original_sample = None - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - # predict noise model_output - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=timestep, - original_size=original_size, - target_size=target_size, - crop_coords=crops_coords_top_left, - return_dict=False, - )[0] - noise_pred = noise_pred.float() - - # perform guidance - 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) - - # compute the previous noisy sample x_t -> x_t-1 - if not isinstance(self.scheduler, CogVideoXDPMScheduler): - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - else: - latents, old_pred_original_sample = self.scheduler.step( - noise_pred, - old_pred_original_sample, - t, - timesteps[i - 1] if i > 0 else None, - latents, - **extra_step_kwargs, - return_dict=False, - ) - latents = latents.to(prompt_embeds.dtype) - - # 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 == 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 not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 0 - ] - else: - image = latents - - 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 CogView3PipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/pipeline_output.py deleted file mode 100644 index 3891dd51e6916abbce4d78042ab7f00ce53af6ae..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview3/pipeline_output.py +++ /dev/null @@ -1,21 +0,0 @@ -from dataclasses import dataclass -from typing import List, Union - -import numpy as np -import PIL.Image - -from ...utils import BaseOutput - - -@dataclass -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, - num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__init__.py deleted file mode 100644 index 6a365e17fee7185cbfb67e857ec2d265c35674ab..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__init__.py +++ /dev/null @@ -1,49 +0,0 @@ -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 = {} -_additional_imports = {} -_import_structure = {"pipeline_output": ["CogView4PlusPipelineOutput"]} - -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_cogview4"] = ["CogView4Pipeline"] - _import_structure["pipeline_cogview4_control"] = ["CogView4ControlPipeline"] -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_cogview4 import CogView4Pipeline - from .pipeline_cogview4_control import CogView4ControlPipeline -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/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 5322ff61e263ffa45583793fa1168410ba81abb8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__pycache__/pipeline_cogview4.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__pycache__/pipeline_cogview4.cpython-310.pyc deleted file mode 100644 index 386682d7480985ce0fe22ac9b24713579bbf370d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__pycache__/pipeline_cogview4.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__pycache__/pipeline_cogview4_control.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__pycache__/pipeline_cogview4_control.cpython-310.pyc deleted file mode 100644 index 915bbd502cbe6ecacf537305941ce170653a7950..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__pycache__/pipeline_cogview4_control.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index e2c2145ea624a3d12a5ed09ea7b6b57564507493..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/pipeline_cogview4.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/pipeline_cogview4.py deleted file mode 100644 index 880253459e1e0a839f19aa6b3f0842b6c6355e43..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/pipeline_cogview4.py +++ /dev/null @@ -1,684 +0,0 @@ -# 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 -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import torch -from transformers import AutoTokenizer, GlmModel - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import VaeImageProcessor -from ...loaders import CogView4LoraLoaderMixin -from ...models import AutoencoderKL, CogView4Transformer2DModel -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 CogView4PipelineOutput - - -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 CogView4Pipeline - - >>> pipe = CogView4Pipeline.from_pretrained("THUDM/CogView4-6B", 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 - - -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. - """ - 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 - - -class CogView4Pipeline(DiffusionPipeline, CogView4LoraLoaderMixin): - r""" - Pipeline for text-to-image generation using CogView4. - - 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.) - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`GLMModel`]): - Frozen text-encoder. CogView4 uses [glm-4-9b-hf](https://huggingface.co/THUDM/glm-4-9b-hf). - tokenizer (`PreTrainedTokenizer`): - Tokenizer of class - [PreTrainedTokenizer](https://huggingface.co/docs/transformers/main/en/main_classes/tokenizer#transformers.PreTrainedTokenizer). - transformer ([`CogView4Transformer2DModel`]): - A text conditioned `CogView4Transformer2DModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - """ - - _optional_components = [] - model_cpu_offload_seq = "text_encoder->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - tokenizer: AutoTokenizer, - text_encoder: GlmModel, - vae: AutoencoderKL, - transformer: CogView4Transformer2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_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) - - def _get_glm_embeds( - self, - prompt: Union[str, List[str]] = None, - max_sequence_length: int = 1024, - 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 - - text_inputs = self.tokenizer( - prompt, - padding="longest", # not use max length - max_length=max_sequence_length, - truncation=True, - add_special_tokens=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[:, 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}" - ) - current_length = text_input_ids.shape[1] - pad_length = (16 - (current_length % 16)) % 16 - if pad_length > 0: - pad_ids = torch.full( - (text_input_ids.shape[0], pad_length), - fill_value=self.tokenizer.pad_token_id, - dtype=text_input_ids.dtype, - device=text_input_ids.device, - ) - text_input_ids = torch.cat([pad_ids, text_input_ids], dim=1) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=True).hidden_states[-2] - - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - return prompt_embeds - - def encode_prompt( - self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = 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, - max_sequence_length: int = 1024, - ): - 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_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. - 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 - max_sequence_length (`int`, defaults to `1024`): - 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_glm_embeds(prompt, max_sequence_length, device, dtype) - - seq_len = prompt_embeds.size(1) - 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) - - 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_glm_embeds(negative_prompt, max_sequence_length, device, dtype) - - seq_len = negative_prompt_embeds.size(1) - 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) - - 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, - negative_prompt, - callback_on_step_end_tensor_inputs, - prompt_embeds=None, - negative_prompt_embeds=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 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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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[0] != negative_prompt_embeds.shape[0]: - raise ValueError( - "`prompt_embeds` and `negative_prompt_embeds` must have the same batch size when passed directly, but" - f" got: `prompt_embeds` {prompt_embeds.shape} and `negative_prompt_embeds`" - f" {negative_prompt_embeds.shape}." - ) - if prompt_embeds.shape[-1] != negative_prompt_embeds.shape[-1]: - raise ValueError( - "`prompt_embeds` and `negative_prompt_embeds` must have the same dimension when passed directly, but" - f" got: `prompt_embeds` {prompt_embeds.shape} and `negative_prompt_embeds`" - f" {negative_prompt_embeds.shape}." - ) - - @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 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: Optional[Union[str, List[str]]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = 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), - 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"], - max_sequence_length: int = 1024, - ) -> Union[CogView4PipelineOutput, 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. If not defined, one has to pass `prompt_embeds`. - 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`). - height (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor): - The height in pixels of the generated image. If not provided, it is set to 1024. - width (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor): - The width in pixels of the generated image. If not provided it is set to 1024. - 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. - 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://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.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. - 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). - 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. - 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 `224`): - Maximum sequence length in encoded prompt. Can be set to other values but may lead to poorer results. - - Examples: - - Returns: - [`~pipelines.cogview4.pipeline_CogView4.CogView4PipelineOutput`] or `tuple`: - [`~pipelines.cogview4.pipeline_CogView4.CogView4PipelineOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, 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 - - height = height or self.transformer.config.sample_size * self.vae_scale_factor - width = width or self.transformer.config.sample_size * self.vae_scale_factor - - original_size = original_size or (height, width) - target_size = (height, width) - - # Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - negative_prompt, - callback_on_step_end_tensor_inputs, - prompt_embeds, - negative_prompt_embeds, - ) - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._current_timestep = None - self._interrupt = False - - # Default 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 - - # Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - negative_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, - ) - - # Prepare latents - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - torch.float32, - device, - generator, - latents, - ) - - # Prepare additional timestep conditions - original_size = torch.tensor([original_size], dtype=prompt_embeds.dtype, device=device) - 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) - - original_size = original_size.repeat(batch_size * num_images_per_prompt, 1) - 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) - 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) - - # Denoising loop - transformer_dtype = self.transformer.dtype - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 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 - latent_model_input = latents.to(transformer_dtype) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latents.shape[0]) - - noise_pred_cond = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=timestep, - original_size=original_size, - target_size=target_size, - crop_coords=crops_coords_top_left, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=negative_prompt_embeds, - timestep=timestep, - original_size=original_size, - target_size=target_size, - crop_coords=crops_coords_top_left, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - - 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] - - # 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, self.scheduler.sigmas[i], 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() - - self._current_timestep = None - - if not output_type == "latent": - latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor - image = self.vae.decode(latents, return_dict=False, generator=generator)[0] - else: - image = latents - - 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 CogView4PipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/pipeline_cogview4_control.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/pipeline_cogview4_control.py deleted file mode 100644 index ac8d786f04f7c63d0b884ca8a8dd83b6844018dd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/pipeline_cogview4_control.py +++ /dev/null @@ -1,732 +0,0 @@ -# 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 -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import torch -from transformers import AutoTokenizer, GlmModel - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...models import AutoencoderKL, CogView4Transformer2DModel -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 CogView4PipelineOutput - - -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 CogView4ControlPipeline - - >>> pipe = CogView4ControlPipeline.from_pretrained("THUDM/CogView4-6B-Control", torch_dtype=torch.bfloat16) - >>> control_image = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png" - ... ) - >>> prompt = "A bird in space" - >>> image = pipe(prompt, control_image=control_image, height=1024, width=1024, guidance_scale=3.5).images[0] - >>> image.save("cogview4-control.png") - ``` -""" - - -# Copied from diffusers.pipelines.cogview4.pipeline_cogview4.calculate_shift -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: 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. - """ - 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 - - -class CogView4ControlPipeline(DiffusionPipeline): - r""" - Pipeline for text-to-image generation using CogView4. - - 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.) - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`GLMModel`]): - Frozen text-encoder. CogView4 uses [glm-4-9b-hf](https://huggingface.co/THUDM/glm-4-9b-hf). - tokenizer (`PreTrainedTokenizer`): - Tokenizer of class - [PreTrainedTokenizer](https://huggingface.co/docs/transformers/main/en/main_classes/tokenizer#transformers.PreTrainedTokenizer). - transformer ([`CogView4Transformer2DModel`]): - A text conditioned `CogView4Transformer2DModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - """ - - _optional_components = [] - model_cpu_offload_seq = "text_encoder->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - tokenizer: AutoTokenizer, - text_encoder: GlmModel, - vae: AutoencoderKL, - transformer: CogView4Transformer2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_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) - - # Copied from diffusers.pipelines.cogview4.pipeline_cogview4.CogView4Pipeline._get_glm_embeds - def _get_glm_embeds( - self, - prompt: Union[str, List[str]] = None, - max_sequence_length: int = 1024, - 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 - - text_inputs = self.tokenizer( - prompt, - padding="longest", # not use max length - max_length=max_sequence_length, - truncation=True, - add_special_tokens=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[:, 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}" - ) - current_length = text_input_ids.shape[1] - pad_length = (16 - (current_length % 16)) % 16 - if pad_length > 0: - pad_ids = torch.full( - (text_input_ids.shape[0], pad_length), - fill_value=self.tokenizer.pad_token_id, - dtype=text_input_ids.dtype, - device=text_input_ids.device, - ) - text_input_ids = torch.cat([pad_ids, text_input_ids], dim=1) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=True).hidden_states[-2] - - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - return prompt_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, - 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, - max_sequence_length: int = 1024, - ): - 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_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. - 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 - max_sequence_length (`int`, defaults to `1024`): - 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_glm_embeds(prompt, max_sequence_length, device, dtype) - - seq_len = prompt_embeds.size(1) - 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) - - 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_glm_embeds(negative_prompt, max_sequence_length, device, dtype) - - seq_len = negative_prompt_embeds.size(1) - 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) - - 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 prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - if isinstance(image, torch.Tensor): - pass - else: - image = self.image_processor.preprocess(image, height=height, width=width) - - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0, output_size=image.shape[0] * repeat_by) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - def check_inputs( - self, - prompt, - height, - width, - negative_prompt, - callback_on_step_end_tensor_inputs, - prompt_embeds=None, - negative_prompt_embeds=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 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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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}." - ) - - @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 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: Optional[Union[str, List[str]]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = 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), - 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"], - max_sequence_length: int = 1024, - ) -> Union[CogView4PipelineOutput, 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. If not defined, one has to pass `prompt_embeds`. - 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`). - height (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor): - The height in pixels of the generated image. If not provided, it is set to 1024. - width (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor): - The width in pixels of the generated image. If not provided it is set to 1024. - 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. - 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://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.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. - 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). - 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.pipeline_CogView4.CogView4PipelineOutput`] 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 `224`): - Maximum sequence length in encoded prompt. Can be set to other values but may lead to poorer results. - Examples: - - Returns: - [`~pipelines.cogview4.pipeline_CogView4.CogView4PipelineOutput`] or `tuple`: - [`~pipelines.cogview4.pipeline_CogView4.CogView4PipelineOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, 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 - - height = height or self.transformer.config.sample_size * self.vae_scale_factor - width = width or self.transformer.config.sample_size * self.vae_scale_factor - - original_size = original_size or (height, width) - target_size = (height, width) - - # Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - negative_prompt, - callback_on_step_end_tensor_inputs, - prompt_embeds, - negative_prompt_embeds, - ) - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._current_timestep = None - self._interrupt = False - - # Default 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 - - # Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - negative_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, - ) - - # Prepare latents - latent_channels = self.transformer.config.in_channels // 2 - - control_image = self.prepare_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=self.vae.dtype, - ) - height, width = control_image.shape[-2:] - - vae_shift_factor = 0 - - control_image = self.vae.encode(control_image).latent_dist.sample() - control_image = (control_image - vae_shift_factor) * self.vae.config.scaling_factor - - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - torch.float32, - device, - generator, - latents, - ) - - # Prepare additional timestep conditions - original_size = torch.tensor([original_size], dtype=prompt_embeds.dtype, device=device) - 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) - - original_size = original_size.repeat(batch_size * num_images_per_prompt, 1) - 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) - 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) - # Denoising loop - transformer_dtype = self.transformer.dtype - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 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 - latent_model_input = torch.cat([latents, control_image], dim=1).to(transformer_dtype) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latents.shape[0]) - - noise_pred_cond = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=timestep, - original_size=original_size, - target_size=target_size, - crop_coords=crops_coords_top_left, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=negative_prompt_embeds, - timestep=timestep, - original_size=original_size, - target_size=target_size, - crop_coords=crops_coords_top_left, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - - 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] - - # 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, self.scheduler.sigmas[i], 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() - - self._current_timestep = None - - if not output_type == "latent": - latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor - image = self.vae.decode(latents, return_dict=False, generator=generator)[0] - else: - image = latents - - 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 CogView4PipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/pipeline_output.py deleted file mode 100644 index 4efec131084504b0bbb422a0802ca4b6f0f05f9e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cogview4/pipeline_output.py +++ /dev/null @@ -1,21 +0,0 @@ -from dataclasses import dataclass -from typing import List, Union - -import numpy as np -import PIL.Image - -from ...utils import BaseOutput - - -@dataclass -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, - num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__init__.py deleted file mode 100644 index 7b9ba330fbd129d104180a2c0e5d05672b48fa47..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__init__.py +++ /dev/null @@ -1,49 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - get_objects_from_module, - is_opencv_available, - is_torch_available, - is_transformers_available, -) - - -_dummy_objects = {} -_import_structure = {} - - -try: - if not (is_transformers_available() and is_torch_available() and is_opencv_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_torch_and_transformers_and_opencv_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_and_opencv_objects)) -else: - _import_structure["pipeline_consisid"] = ["ConsisIDPipeline"] - -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_consisid import ConsisIDPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 6f80d68e142622400c053840ab76d916c95b263e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__pycache__/consisid_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__pycache__/consisid_utils.cpython-310.pyc deleted file mode 100644 index f4aa0b6e71c8aecd35f07f0f903d569d992ae283..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__pycache__/consisid_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__pycache__/pipeline_consisid.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__pycache__/pipeline_consisid.cpython-310.pyc deleted file mode 100644 index af1c912a95d9aa6e70398c363c5ab3ed94954d9f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__pycache__/pipeline_consisid.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index bd906a9682884f78e417d163e562b60b59e7ebc3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/consisid_utils.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/consisid_utils.py deleted file mode 100644 index 521d4d787e54c37ae73cae70294f857e2ef0ca82..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/consisid_utils.py +++ /dev/null @@ -1,357 +0,0 @@ -import importlib.util -import os - -import cv2 -import numpy as np -import torch -from PIL import Image, ImageOps -from torchvision.transforms import InterpolationMode -from torchvision.transforms.functional import normalize, resize - -from ...utils import get_logger, load_image - - -logger = get_logger(__name__) - -_insightface_available = importlib.util.find_spec("insightface") is not None -_consisid_eva_clip_available = importlib.util.find_spec("consisid_eva_clip") is not None -_facexlib_available = importlib.util.find_spec("facexlib") is not None - -if _insightface_available: - import insightface - from insightface.app import FaceAnalysis -else: - raise ImportError("insightface is not available. Please install it using 'pip install insightface'.") - -if _consisid_eva_clip_available: - from consisid_eva_clip import create_model_and_transforms - from consisid_eva_clip.constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD -else: - raise ImportError("consisid_eva_clip is not available. Please install it using 'pip install consisid_eva_clip'.") - -if _facexlib_available: - from facexlib.parsing import init_parsing_model - from facexlib.utils.face_restoration_helper import FaceRestoreHelper -else: - raise ImportError("facexlib is not available. Please install it using 'pip install facexlib'.") - - -def resize_numpy_image_long(image, resize_long_edge=768): - """ - Resize the input image to a specified long edge while maintaining aspect ratio. - - Args: - image (numpy.ndarray): Input image (H x W x C or H x W). - resize_long_edge (int): The target size for the long edge of the image. Default is 768. - - Returns: - numpy.ndarray: Resized image with the long edge matching `resize_long_edge`, while maintaining the aspect - ratio. - """ - - h, w = image.shape[:2] - if max(h, w) <= resize_long_edge: - return image - k = resize_long_edge / max(h, w) - h = int(h * k) - w = int(w * k) - image = cv2.resize(image, (w, h), interpolation=cv2.INTER_LANCZOS4) - return image - - -def img2tensor(imgs, bgr2rgb=True, float32=True): - """Numpy array to tensor. - - Args: - imgs (list[ndarray] | ndarray): Input images. - bgr2rgb (bool): Whether to change bgr to rgb. - float32 (bool): Whether to change to float32. - - Returns: - list[tensor] | tensor: Tensor images. If returned results only have - one element, just return tensor. - """ - - def _totensor(img, bgr2rgb, float32): - if img.shape[2] == 3 and bgr2rgb: - if img.dtype == "float64": - img = img.astype("float32") - img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) - img = torch.from_numpy(img.transpose(2, 0, 1)) - if float32: - img = img.float() - return img - - if isinstance(imgs, list): - return [_totensor(img, bgr2rgb, float32) for img in imgs] - return _totensor(imgs, bgr2rgb, float32) - - -def to_gray(img): - """ - Converts an RGB image to grayscale by applying the standard luminosity formula. - - Args: - img (torch.Tensor): The input image tensor with shape (batch_size, channels, height, width). - The image is expected to be in RGB format (3 channels). - - Returns: - torch.Tensor: The grayscale image tensor with shape (batch_size, 3, height, width). - The grayscale values are replicated across all three channels. - """ - x = 0.299 * img[:, 0:1] + 0.587 * img[:, 1:2] + 0.114 * img[:, 2:3] - x = x.repeat(1, 3, 1, 1) - return x - - -def process_face_embeddings( - face_helper_1, - clip_vision_model, - face_helper_2, - eva_transform_mean, - eva_transform_std, - app, - device, - weight_dtype, - image, - original_id_image=None, - is_align_face=True, -): - """ - Process face embeddings from an image, extracting relevant features such as face embeddings, landmarks, and parsed - face features using a series of face detection and alignment tools. - - Args: - face_helper_1: Face helper object (first helper) for alignment and landmark detection. - clip_vision_model: Pre-trained CLIP vision model used for feature extraction. - face_helper_2: Face helper object (second helper) for embedding extraction. - eva_transform_mean: Mean values for image normalization before passing to EVA model. - eva_transform_std: Standard deviation values for image normalization before passing to EVA model. - app: Application instance used for face detection. - device: Device (CPU or GPU) where the computations will be performed. - weight_dtype: Data type of the weights for precision (e.g., `torch.float32`). - image: Input image in RGB format with pixel values in the range [0, 255]. - original_id_image: (Optional) Original image for feature extraction if `is_align_face` is False. - is_align_face: Boolean flag indicating whether face alignment should be performed. - - Returns: - 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. - - face_kps: Keypoints of the face detected in the image. - """ - - face_helper_1.clean_all() - image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) - # get antelopev2 embedding - face_info = app.get(image_bgr) - if len(face_info) > 0: - face_info = sorted(face_info, key=lambda x: (x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]))[ - -1 - ] # only use the maximum face - id_ante_embedding = face_info["embedding"] # (512,) - face_kps = face_info["kps"] - else: - id_ante_embedding = None - face_kps = None - - # using facexlib to detect and align face - face_helper_1.read_image(image_bgr) - face_helper_1.get_face_landmarks_5(only_center_face=True) - if face_kps is None: - face_kps = face_helper_1.all_landmarks_5[0] - face_helper_1.align_warp_face() - if len(face_helper_1.cropped_faces) == 0: - raise RuntimeError("facexlib align face fail") - align_face = face_helper_1.cropped_faces[0] # (512, 512, 3) # RGB - - # in case insightface didn't detect face - if id_ante_embedding is None: - logger.warning("Failed to detect face using insightface. Extracting embedding with align face") - id_ante_embedding = face_helper_2.get_feat(align_face) - - id_ante_embedding = torch.from_numpy(id_ante_embedding).to(device, weight_dtype) # torch.Size([512]) - if id_ante_embedding.ndim == 1: - id_ante_embedding = id_ante_embedding.unsqueeze(0) # torch.Size([1, 512]) - - # parsing - if is_align_face: - input = img2tensor(align_face, bgr2rgb=True).unsqueeze(0) / 255.0 # torch.Size([1, 3, 512, 512]) - input = input.to(device) - parsing_out = face_helper_1.face_parse(normalize(input, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]))[0] - parsing_out = parsing_out.argmax(dim=1, keepdim=True) # torch.Size([1, 1, 512, 512]) - bg_label = [0, 16, 18, 7, 8, 9, 14, 15] - bg = sum(parsing_out == i for i in bg_label).bool() - white_image = torch.ones_like(input) # torch.Size([1, 3, 512, 512]) - # only keep the face features - return_face_features_image = torch.where(bg, white_image, to_gray(input)) # torch.Size([1, 3, 512, 512]) - return_face_features_image_2 = torch.where(bg, white_image, input) # torch.Size([1, 3, 512, 512]) - else: - original_image_bgr = cv2.cvtColor(original_id_image, cv2.COLOR_RGB2BGR) - input = img2tensor(original_image_bgr, bgr2rgb=True).unsqueeze(0) / 255.0 # torch.Size([1, 3, 512, 512]) - input = input.to(device) - return_face_features_image = return_face_features_image_2 = input - - # transform img before sending to eva-clip-vit - face_features_image = resize( - return_face_features_image, clip_vision_model.image_size, InterpolationMode.BICUBIC - ) # torch.Size([1, 3, 336, 336]) - face_features_image = normalize(face_features_image, eva_transform_mean, eva_transform_std) - id_cond_vit, id_vit_hidden = clip_vision_model( - face_features_image.to(weight_dtype), return_all_features=False, return_hidden=True, shuffle=False - ) # torch.Size([1, 768]), list(torch.Size([1, 577, 1024])) - id_cond_vit_norm = torch.norm(id_cond_vit, 2, 1, True) - id_cond_vit = torch.div(id_cond_vit, id_cond_vit_norm) - - id_cond = torch.cat( - [id_ante_embedding, id_cond_vit], dim=-1 - ) # torch.Size([1, 512]), torch.Size([1, 768]) -> torch.Size([1, 1280]) - - return ( - id_cond, - id_vit_hidden, - return_face_features_image_2, - face_kps, - ) # torch.Size([1, 1280]), list(torch.Size([1, 577, 1024])) - - -def process_face_embeddings_infer( - face_helper_1, - clip_vision_model, - face_helper_2, - eva_transform_mean, - eva_transform_std, - app, - device, - weight_dtype, - img_file_path, - is_align_face=True, -): - """ - Process face embeddings from an input image for inference, including alignment, feature extraction, and embedding - concatenation. - - Args: - face_helper_1: Face helper object (first helper) for alignment and landmark detection. - clip_vision_model: Pre-trained CLIP vision model used for feature extraction. - face_helper_2: Face helper object (second helper) for embedding extraction. - eva_transform_mean: Mean values for image normalization before passing to EVA model. - eva_transform_std: Standard deviation values for image normalization before passing to EVA model. - app: Application instance used for face detection. - device: Device (CPU or GPU) where the computations will be performed. - weight_dtype: Data type of the weights for precision (e.g., `torch.float32`). - img_file_path: Path to the input image file (string) or a numpy array representing an image. - is_align_face: Boolean flag indicating whether face alignment should be performed (default: True). - - Returns: - 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. - - face_kps: Keypoints of the face detected in the image. - """ - - # Load and preprocess the input image - if isinstance(img_file_path, str): - image = np.array(load_image(image=img_file_path).convert("RGB")) - else: - image = np.array(ImageOps.exif_transpose(Image.fromarray(img_file_path)).convert("RGB")) - - # Resize image to ensure the longer side is 1024 pixels - image = resize_numpy_image_long(image, 1024) - original_id_image = image - - # Process the image to extract face embeddings and related features - id_cond, id_vit_hidden, align_crop_face_image, face_kps = process_face_embeddings( - face_helper_1, - clip_vision_model, - face_helper_2, - eva_transform_mean, - eva_transform_std, - app, - device, - weight_dtype, - image, - original_id_image, - is_align_face, - ) - - # Convert the aligned cropped face image (torch tensor) to a numpy array - tensor = align_crop_face_image.cpu().detach() - tensor = tensor.squeeze() - tensor = tensor.permute(1, 2, 0) - tensor = tensor.numpy() * 255 - tensor = tensor.astype(np.uint8) - image = ImageOps.exif_transpose(Image.fromarray(tensor)) - - return id_cond, id_vit_hidden, image, face_kps - - -def prepare_face_models(model_path, device, dtype): - """ - Prepare all face models for the facial recognition task. - - Parameters: - - model_path: Path to the directory containing model files. - - device: The device (e.g., 'cuda', 'xpu', 'cpu') where models will be loaded. - - dtype: Data type (e.g., torch.float32) for model inference. - - Returns: - - face_helper_1: First face restoration helper. - - face_helper_2: Second face restoration helper. - - face_clip_model: CLIP model for face extraction. - - eva_transform_mean: Mean value for image normalization. - - eva_transform_std: Standard deviation value for image normalization. - - face_main_model: Main face analysis model. - """ - # get helper model - face_helper_1 = FaceRestoreHelper( - upscale_factor=1, - face_size=512, - crop_ratio=(1, 1), - det_model="retinaface_resnet50", - save_ext="png", - device=device, - model_rootpath=os.path.join(model_path, "face_encoder"), - ) - face_helper_1.face_parse = None - face_helper_1.face_parse = init_parsing_model( - model_name="bisenet", device=device, model_rootpath=os.path.join(model_path, "face_encoder") - ) - face_helper_2 = insightface.model_zoo.get_model( - f"{model_path}/face_encoder/models/antelopev2/glintr100.onnx", providers=["CUDAExecutionProvider"] - ) - face_helper_2.prepare(ctx_id=0) - - # get local facial extractor part 1 - model, _, _ = create_model_and_transforms( - "EVA02-CLIP-L-14-336", - os.path.join(model_path, "face_encoder", "EVA02_CLIP_L_336_psz14_s6B.pt"), - force_custom_clip=True, - ) - face_clip_model = model.visual - eva_transform_mean = getattr(face_clip_model, "image_mean", OPENAI_DATASET_MEAN) - eva_transform_std = getattr(face_clip_model, "image_std", OPENAI_DATASET_STD) - if not isinstance(eva_transform_mean, (list, tuple)): - eva_transform_mean = (eva_transform_mean,) * 3 - if not isinstance(eva_transform_std, (list, tuple)): - eva_transform_std = (eva_transform_std,) * 3 - eva_transform_mean = eva_transform_mean - eva_transform_std = eva_transform_std - - # get local facial extractor part 2 - face_main_model = FaceAnalysis( - name="antelopev2", root=os.path.join(model_path, "face_encoder"), providers=["CUDAExecutionProvider"] - ) - face_main_model.prepare(ctx_id=0, det_size=(640, 640)) - - # move face models to device - face_helper_1.face_det.eval() - face_helper_1.face_parse.eval() - face_clip_model.eval() - face_helper_1.face_det.to(device) - face_helper_1.face_parse.to(device) - face_clip_model.to(device, dtype=dtype) - - return face_helper_1, face_helper_2, face_clip_model, face_main_model, eva_transform_mean, eva_transform_std diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/pipeline_consisid.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/pipeline_consisid.py deleted file mode 100644 index 644bd811f6c72627bc90558e574f4db783ebc8af..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/pipeline_consisid.py +++ /dev/null @@ -1,974 +0,0 @@ -# Copyright 2025 ConsisID Authors 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 math -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL -import torch -from transformers import T5EncoderModel, T5Tokenizer - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput -from ...loaders import CogVideoXLoraLoaderMixin -from ...models import AutoencoderKLCogVideoX, ConsisIDTransformer3DModel -from ...models.embeddings import get_3d_rotary_pos_embed -from ...pipelines.pipeline_utils import DiffusionPipeline -from ...schedulers import CogVideoXDPMScheduler -from ...utils import is_opencv_available, logging, replace_example_docstring -from ...utils.torch_utils import randn_tensor -from ...video_processor import VideoProcessor -from .pipeline_output import ConsisIDPipelineOutput - - -if is_opencv_available(): - import cv2 - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -EXAMPLE_DOC_STRING = """ - Examples: - ```python - >>> import torch - >>> from diffusers import ConsisIDPipeline - >>> from diffusers.pipelines.consisid.consisid_utils import prepare_face_models, process_face_embeddings_infer - >>> from diffusers.utils import export_to_video - >>> from huggingface_hub import snapshot_download - - >>> snapshot_download(repo_id="BestWishYsh/ConsisID-preview", local_dir="BestWishYsh/ConsisID-preview") - >>> ( - ... face_helper_1, - ... face_helper_2, - ... face_clip_model, - ... face_main_model, - ... eva_transform_mean, - ... eva_transform_std, - ... ) = prepare_face_models("BestWishYsh/ConsisID-preview", device="cuda", dtype=torch.bfloat16) - >>> pipe = ConsisIDPipeline.from_pretrained("BestWishYsh/ConsisID-preview", torch_dtype=torch.bfloat16) - >>> pipe.to("cuda") - - >>> # ConsisID works well with long and well-described prompts. Make sure the face in the image is clearly visible (e.g., preferably half-body or full-body). - >>> prompt = "The video captures a boy walking along a city street, filmed in black and white on a classic 35mm camera. His expression is thoughtful, his brow slightly furrowed as if he's lost in contemplation. The film grain adds a textured, timeless quality to the image, evoking a sense of nostalgia. Around him, the cityscape is filled with vintage buildings, cobblestone sidewalks, and softly blurred figures passing by, their outlines faint and indistinct. Streetlights cast a gentle glow, while shadows play across the boy's path, adding depth to the scene. The lighting highlights the boy's subtle smile, hinting at a fleeting moment of curiosity. The overall cinematic atmosphere, complete with classic film still aesthetics and dramatic contrasts, gives the scene an evocative and introspective feel." - >>> image = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/consisid/consisid_input.png?download=true" - - >>> id_cond, id_vit_hidden, image, face_kps = process_face_embeddings_infer( - ... face_helper_1, - ... face_clip_model, - ... face_helper_2, - ... eva_transform_mean, - ... eva_transform_std, - ... face_main_model, - ... "cuda", - ... torch.bfloat16, - ... image, - ... is_align_face=True, - ... ) - - >>> video = pipe( - ... image=image, - ... prompt=prompt, - ... num_inference_steps=50, - ... guidance_scale=6.0, - ... use_dynamic_cfg=False, - ... id_vit_hidden=id_vit_hidden, - ... id_cond=id_cond, - ... kps_cond=face_kps, - ... generator=torch.Generator("cuda").manual_seed(42), - ... ) - >>> export_to_video(video.frames[0], "output.mp4", fps=8) - ``` -""" - - -def draw_kps(image_pil, kps, color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)]): - """ - This function draws keypoints and the limbs connecting them on an image. - - 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 - colors. - - Returns: - - PIL.Image: Image with the keypoints and limbs drawn. - """ - - stickwidth = 4 - limbSeq = np.array([[0, 2], [1, 2], [3, 2], [4, 2]]) - kps = np.array(kps) - - w, h = image_pil.size - out_img = np.zeros([h, w, 3]) - - for i in range(len(limbSeq)): - index = limbSeq[i] - color = color_list[index[0]] - - x = kps[index][:, 0] - y = kps[index][:, 1] - length = ((x[0] - x[1]) ** 2 + (y[0] - y[1]) ** 2) ** 0.5 - angle = math.degrees(math.atan2(y[0] - y[1], x[0] - x[1])) - polygon = cv2.ellipse2Poly( - (int(np.mean(x)), int(np.mean(y))), (int(length / 2), stickwidth), int(angle), 0, 360, 1 - ) - out_img = cv2.fillConvexPoly(out_img.copy(), polygon, color) - out_img = (out_img * 0.6).astype(np.uint8) - - for idx_kp, kp in enumerate(kps): - color = color_list[idx_kp] - x, y = kp - out_img = cv2.circle(out_img.copy(), (int(x), int(y)), 10, color, -1) - - out_img_pil = PIL.Image.fromarray(out_img.astype(np.uint8)) - return out_img_pil - - -# Similar to diffusers.pipelines.hunyuandit.pipeline_hunyuandit.get_resize_crop_region_for_grid -def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): - """ - This function calculates the resize and crop region for an image to fit a target width and height while preserving - the aspect ratio. - - Parameters: - - src (tuple): A tuple containing the source image's height (h) and width (w). - - tgt_width (int): The target width to resize the image. - - tgt_height (int): The target height to resize the image. - - Returns: - - tuple: Two tuples representing the crop region: - 1. The top-left coordinates of the crop region. - 2. The bottom-right coordinates of the crop region. - """ - - tw = tgt_width - th = tgt_height - h, w = src - r = h / w - if r > (th / tw): - resize_height = th - resize_width = int(round(th / h * w)) - else: - resize_width = tw - resize_height = int(round(tw / w * h)) - - crop_top = int(round((th - resize_height) / 2.0)) - crop_left = int(round((tw - resize_width) / 2.0)) - - return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_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, - **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: Optional[torch.Generator] = 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 ConsisIDPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin): - r""" - Pipeline for image-to-video generation using ConsisID. - - 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.) - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. - text_encoder ([`T5EncoderModel`]): - Frozen text-encoder. ConsisID uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the - [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant. - tokenizer (`T5Tokenizer`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - transformer ([`ConsisIDTransformer3DModel`]): - A text conditioned `ConsisIDTransformer3DModel` to denoise the encoded video latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `transformer` to denoise the encoded video latents. - """ - - _optional_components = [] - model_cpu_offload_seq = "text_encoder->transformer->vae" - - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - ] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - vae: AutoencoderKLCogVideoX, - transformer: ConsisIDTransformer3DModel, - scheduler: CogVideoXDPMScheduler, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - vae=vae, - transformer=transformer, - scheduler=scheduler, - ) - self.vae_scale_factor_spatial = ( - 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8 - ) - self.vae_scale_factor_temporal = ( - self.vae.config.temporal_compression_ratio if hasattr(self, "vae") and self.vae is not None else 4 - ) - self.vae_scaling_factor_image = ( - self.vae.config.scaling_factor if hasattr(self, "vae") and self.vae is not None else 0.7 - ) - - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds - def _get_t5_prompt_embeds( - self, - prompt: Union[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 = 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 - 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) - - return 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, - 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, - max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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 prepare_latents( - self, - image: torch.Tensor, - batch_size: int = 1, - num_channels_latents: int = 16, - 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, - ): - 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." - ) - - num_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - shape = ( - batch_size, - num_frames, - num_channels_latents, - height // self.vae_scale_factor_spatial, - width // self.vae_scale_factor_spatial, - ) - - image = image.unsqueeze(2) # [B, C, F, H, W] - - if isinstance(generator, list): - image_latents = [ - retrieve_latents(self.vae.encode(image[i].unsqueeze(0)), generator[i]) for i in range(batch_size) - ] - if kps_cond is not None: - kps_cond = kps_cond.unsqueeze(2) - kps_cond_latents = [ - retrieve_latents(self.vae.encode(kps_cond[i].unsqueeze(0)), generator[i]) - for i in range(batch_size) - ] - else: - image_latents = [retrieve_latents(self.vae.encode(img.unsqueeze(0)), generator) for img in image] - if kps_cond is not None: - kps_cond = kps_cond.unsqueeze(2) - kps_cond_latents = [retrieve_latents(self.vae.encode(img.unsqueeze(0)), generator) for img in kps_cond] - - image_latents = torch.cat(image_latents, dim=0).to(dtype).permute(0, 2, 1, 3, 4) # [B, F, C, H, W] - image_latents = self.vae_scaling_factor_image * image_latents - - if kps_cond is not None: - kps_cond_latents = torch.cat(kps_cond_latents, dim=0).to(dtype).permute(0, 2, 1, 3, 4) # [B, F, C, H, W] - kps_cond_latents = self.vae_scaling_factor_image * kps_cond_latents - - padding_shape = ( - batch_size, - num_frames - 2, - num_channels_latents, - height // self.vae_scale_factor_spatial, - width // self.vae_scale_factor_spatial, - ) - else: - padding_shape = ( - batch_size, - num_frames - 1, - num_channels_latents, - height // self.vae_scale_factor_spatial, - width // self.vae_scale_factor_spatial, - ) - - latent_padding = torch.zeros(padding_shape, device=device, dtype=dtype) - if kps_cond is not None: - image_latents = torch.cat([image_latents, kps_cond_latents, latent_padding], dim=1) - else: - image_latents = torch.cat([image_latents, latent_padding], dim=1) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents, image_latents - - # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.decode_latents - def decode_latents(self, latents: torch.Tensor) -> torch.Tensor: - latents = latents.permute(0, 2, 1, 3, 4) # [batch_size, num_channels, num_frames, height, width] - latents = 1 / self.vae_scaling_factor_image * latents - - frames = self.vae.decode(latents).sample - return frames - - # Copied from diffusers.pipelines.animatediff.pipeline_animatediff_video2video.AnimateDiffVideoToVideoPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, timesteps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - timesteps = timesteps[t_start * self.scheduler.order :] - - return timesteps, num_inference_steps - t_start - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - image, - prompt, - height, - width, - negative_prompt, - callback_on_step_end_tensor_inputs, - latents=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - if ( - not isinstance(image, torch.Tensor) - and not isinstance(image, PIL.Image.Image) - 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" - f" {type(image)}" - ) - - 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_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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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_rotary_positional_embeddings( - self, - height: int, - width: int, - num_frames: int, - device: torch.device, - ) -> 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 - base_size_height = self.transformer.config.sample_height // self.transformer.config.patch_size - - grid_crops_coords = get_resize_crop_region_for_grid( - (grid_height, grid_width), base_size_width, base_size_height - ) - freqs_cos, freqs_sin = get_3d_rotary_pos_embed( - embed_dim=self.transformer.config.attention_head_dim, - crops_coords=grid_crops_coords, - grid_size=(grid_height, grid_width), - temporal_size=num_frames, - device=device, - ) - - return freqs_cos, freqs_sin - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def num_timesteps(self): - return self._num_timesteps - - @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, - prompt: Optional[Union[str, List[str]]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: int = 480, - width: int = 720, - num_frames: int = 49, - num_inference_steps: int = 50, - guidance_scale: float = 6.0, - 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, - 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"], - 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]: - """ - 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. - 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`). - height (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial): - The height in pixels of the generated image. This is set to 480 by default for the best results. - width (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial): - The width in pixels of the generated image. This is set to 720 by default for the best results. - num_frames (`int`, defaults to `49`): - Number of frames to generate. Must be divisible by self.vae_scale_factor_temporal. Generated video will - contain 1 extra frame because ConsisID is conditioned with (num_seconds * fps + 1) frames where - num_seconds is 6 and fps is 4. However, since videos can be saved at any fps, the only condition that - needs to be satisfied is that of divisibility mentioned above. - 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 6): - 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. - use_dynamic_cfg (`bool`, *optional*, defaults to `False`): - If True, dynamically adjusts the guidance scale during inference. This allows the model to use a - progressive guidance scale, improving the balance between text-guided generation and image quality over - the course of the inference steps. Typically, early inference steps use a higher guidance scale for - 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*): - 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. - 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 `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*): - 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*): - 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*): - 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 - information. - - Examples: - - Returns: - [`~pipelines.consisid.pipeline_output.ConsisIDPipelineOutput`] or `tuple`: - [`~pipelines.consisid.pipeline_output.ConsisIDPipelineOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, 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 - - height = height or self.transformer.config.sample_height * self.vae_scale_factor_spatial - width = width or self.transformer.config.sample_width * self.vae_scale_factor_spatial - num_frames = num_frames or self.transformer.config.sample_frames - - num_videos_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - image=image, - prompt=prompt, - height=height, - width=width, - negative_prompt=negative_prompt, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - latents=latents, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._interrupt = False - - # 2. Default 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=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, - ) - if do_classifier_free_guidance: - 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) - self._num_timesteps = len(timesteps) - - # 5. Prepare latents - is_kps = getattr(self.transformer.config, "is_kps", False) - kps_cond = kps_cond if is_kps else None - if kps_cond is not None: - kps_cond = draw_kps(image, kps_cond) - kps_cond = self.video_processor.preprocess(kps_cond, height=height, width=width).to( - device, dtype=prompt_embeds.dtype - ) - - image = self.video_processor.preprocess(image, height=height, width=width).to( - device, dtype=prompt_embeds.dtype - ) - - latent_channels = self.transformer.config.in_channels // 2 - latents, image_latents = self.prepare_latents( - image, - batch_size * num_videos_per_prompt, - latent_channels, - num_frames, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - kps_cond, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Create rotary embeds if required - image_rotary_emb = ( - self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device) - if self.transformer.config.use_rotary_positional_embeddings - else None - ) - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - # for DPM-solver++ - old_pred_original_sample = None - timesteps_cpu = timesteps.cpu() - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - latent_image_input = torch.cat([image_latents] * 2) if do_classifier_free_guidance else image_latents - latent_model_input = torch.cat([latent_model_input, latent_image_input], dim=2) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - # predict noise model_output - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=timestep, - image_rotary_emb=image_rotary_emb, - attention_kwargs=attention_kwargs, - return_dict=False, - id_vit_hidden=id_vit_hidden, - id_cond=id_cond, - )[0] - noise_pred = noise_pred.float() - - # perform guidance - if use_dynamic_cfg: - self._guidance_scale = 1 + guidance_scale * ( - ( - 1 - - math.cos( - math.pi - * ((num_inference_steps - timesteps_cpu[i].item()) / num_inference_steps) ** 5.0 - ) - ) - / 2 - ) - if 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) - - # compute the previous noisy sample x_t -> x_t-1 - if not isinstance(self.scheduler, CogVideoXDPMScheduler): - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - else: - latents, old_pred_original_sample = self.scheduler.step( - noise_pred, - old_pred_original_sample, - t, - timesteps[i - 1] if i > 0 else None, - latents, - **extra_step_kwargs, - return_dict=False, - ) - latents = latents.to(prompt_embeds.dtype) - - # 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 == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): - progress_bar.update() - - if not output_type == "latent": - video = self.decode_latents(latents) - video = self.video_processor.postprocess_video(video=video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return ConsisIDPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/pipeline_output.py deleted file mode 100644 index dd4a63aa50b97432e0dad1a8e27f2c12e84b1b9d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/consisid/pipeline_output.py +++ /dev/null @@ -1,20 +0,0 @@ -from dataclasses import dataclass - -import torch - -from diffusers.utils import BaseOutput - - -@dataclass -class ConsisIDPipelineOutput(BaseOutput): - r""" - 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 - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/consistency_models/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/consistency_models/__init__.py deleted file mode 100644 index 162d91c010acf95aa2daf87c51ab1e0c68361fd5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/consistency_models/__init__.py +++ /dev/null @@ -1,24 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - _LazyModule, -) - - -_import_structure = { - "pipeline_consistency_models": ["ConsistencyModelPipeline"], -} - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - from .pipeline_consistency_models import ConsistencyModelPipeline - -else: - import sys - - sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/consistency_models/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/consistency_models/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 6325821abfefa12b7b065a2cd705d085c25959b2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/consistency_models/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/consistency_models/__pycache__/pipeline_consistency_models.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/consistency_models/__pycache__/pipeline_consistency_models.cpython-310.pyc deleted file mode 100644 index 4a3fe8c2c223efdbe6245fc5c04221a708288579..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/consistency_models/__pycache__/pipeline_consistency_models.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/consistency_models/pipeline_consistency_models.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/consistency_models/pipeline_consistency_models.py deleted file mode 100644 index dec448f3f46d175531d83e698c09ae8fe6b63e34..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/consistency_models/pipeline_consistency_models.py +++ /dev/null @@ -1,286 +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 Callable, List, Optional, Union - -import torch - -from ...models import UNet2DModel -from ...schedulers import CMStochasticIterativeScheduler -from ...utils import ( - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -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 ConsistencyModelPipeline - - >>> device = "cuda" - >>> # Load the cd_imagenet64_l2 checkpoint. - >>> model_id_or_path = "openai/diffusers-cd_imagenet64_l2" - >>> pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16) - >>> pipe.to(device) - - >>> # Onestep Sampling - >>> image = pipe(num_inference_steps=1).images[0] - >>> image.save("cd_imagenet64_l2_onestep_sample.png") - - >>> # Onestep sampling, class-conditional image generation - >>> # ImageNet-64 class label 145 corresponds to king penguins - >>> image = pipe(num_inference_steps=1, class_labels=145).images[0] - >>> image.save("cd_imagenet64_l2_onestep_sample_penguin.png") - - >>> # Multistep sampling, class-conditional image generation - >>> # Timesteps can be explicitly specified; the particular timesteps below are from the original GitHub repo: - >>> # https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L77 - >>> image = pipe(num_inference_steps=None, timesteps=[22, 0], class_labels=145).images[0] - >>> image.save("cd_imagenet64_l2_multistep_sample_penguin.png") - ``` -""" - - -class ConsistencyModelPipeline(DiffusionPipeline): - r""" - Pipeline for unconditional or class-conditional image generation. - - 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: - unet ([`UNet2DModel`]): - A `UNet2DModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. Currently only - compatible with [`CMStochasticIterativeScheduler`]. - """ - - model_cpu_offload_seq = "unet" - - def __init__(self, unet: UNet2DModel, scheduler: CMStochasticIterativeScheduler) -> None: - super().__init__() - - self.register_modules( - unet=unet, - scheduler=scheduler, - ) - - self.safety_checker = None - - def prepare_latents(self, batch_size, num_channels, height, width, dtype, device, generator, latents=None): - shape = (batch_size, num_channels, 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." - ) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device=device, dtype=dtype) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - # Follows diffusers.VaeImageProcessor.postprocess - def postprocess_image(self, sample: torch.Tensor, output_type: str = "pil"): - if output_type not in ["pt", "np", "pil"]: - raise ValueError( - f"output_type={output_type} is not supported. Make sure to choose one of ['pt', 'np', or 'pil']" - ) - - # Equivalent to diffusers.VaeImageProcessor.denormalize - sample = (sample / 2 + 0.5).clamp(0, 1) - if output_type == "pt": - return sample - - # Equivalent to diffusers.VaeImageProcessor.pt_to_numpy - sample = sample.cpu().permute(0, 2, 3, 1).numpy() - if output_type == "np": - return sample - - # Output_type must be 'pil' - sample = self.numpy_to_pil(sample) - return sample - - def prepare_class_labels(self, batch_size, device, class_labels=None): - if self.unet.config.num_class_embeds is not None: - if isinstance(class_labels, list): - class_labels = torch.tensor(class_labels, dtype=torch.int) - elif isinstance(class_labels, int): - assert batch_size == 1, "Batch size must be 1 if classes is an int" - class_labels = torch.tensor([class_labels], dtype=torch.int) - elif class_labels is None: - # Randomly generate batch_size class labels - # TODO: should use generator here? int analogue of randn_tensor is not exposed in ...utils - class_labels = torch.randint(0, self.unet.config.num_class_embeds, size=(batch_size,)) - class_labels = class_labels.to(device) - else: - class_labels = None - return class_labels - - def check_inputs(self, num_inference_steps, timesteps, latents, batch_size, img_size, callback_steps): - if num_inference_steps is None and timesteps is None: - raise ValueError("Exactly one of `num_inference_steps` or `timesteps` must be supplied.") - - if num_inference_steps is not None and timesteps is not None: - logger.warning( - f"Both `num_inference_steps`: {num_inference_steps} and `timesteps`: {timesteps} are supplied;" - " `timesteps` will be used over `num_inference_steps`." - ) - - if latents is not None: - expected_shape = (batch_size, 3, img_size, img_size) - if latents.shape != expected_shape: - raise ValueError(f"The shape of latents is {latents.shape} but is expected to be {expected_shape}.") - - if (callback_steps is None) or ( - 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)}." - ) - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - batch_size: int = 1, - class_labels: Optional[Union[torch.Tensor, List[int], int]] = 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], 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*): - 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*): - 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*): - 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`. - 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 [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images. - """ - # 0. Prepare call parameters - img_size = self.unet.config.sample_size - device = self._execution_device - - # 1. Check inputs - self.check_inputs(num_inference_steps, timesteps, latents, batch_size, img_size, callback_steps) - - # 2. Prepare image latents - # Sample image latents x_0 ~ N(0, sigma_0^2 * I) - sample = self.prepare_latents( - batch_size=batch_size, - num_channels=self.unet.config.in_channels, - height=img_size, - width=img_size, - dtype=self.unet.dtype, - device=device, - generator=generator, - latents=latents, - ) - - # 3. Handle class_labels for class-conditional models - class_labels = self.prepare_class_labels(batch_size, device, class_labels=class_labels) - - # 4. Prepare timesteps - if timesteps is not None: - self.scheduler.set_timesteps(timesteps=timesteps, device=device) - timesteps = self.scheduler.timesteps - num_inference_steps = len(timesteps) - else: - self.scheduler.set_timesteps(num_inference_steps) - timesteps = self.scheduler.timesteps - - # 5. Denoising loop - # Multistep sampling: implements Algorithm 1 in the paper - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - scaled_sample = self.scheduler.scale_model_input(sample, t) - model_output = self.unet(scaled_sample, t, class_labels=class_labels, return_dict=False)[0] - - sample = self.scheduler.step(model_output, t, sample, generator=generator)[0] - - # call the callback, if provided - progress_bar.update() - if callback is not None and i % callback_steps == 0: - callback(i, t, sample) - - if XLA_AVAILABLE: - xm.mark_step() - - # 6. Post-process image sample - image = self.postprocess_image(sample, output_type=output_type) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__init__.py deleted file mode 100644 index a49dccf235a3b84b2695bba6724ba913d3f60155..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__init__.py +++ /dev/null @@ -1,86 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - get_objects_from_module, - is_flax_available, - 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["multicontrolnet"] = ["MultiControlNetModel"] - _import_structure["pipeline_controlnet"] = ["StableDiffusionControlNetPipeline"] - _import_structure["pipeline_controlnet_blip_diffusion"] = ["BlipDiffusionControlNetPipeline"] - _import_structure["pipeline_controlnet_img2img"] = ["StableDiffusionControlNetImg2ImgPipeline"] - _import_structure["pipeline_controlnet_inpaint"] = ["StableDiffusionControlNetInpaintPipeline"] - _import_structure["pipeline_controlnet_inpaint_sd_xl"] = ["StableDiffusionXLControlNetInpaintPipeline"] - _import_structure["pipeline_controlnet_sd_xl"] = ["StableDiffusionXLControlNetPipeline"] - _import_structure["pipeline_controlnet_sd_xl_img2img"] = ["StableDiffusionXLControlNetImg2ImgPipeline"] - _import_structure["pipeline_controlnet_union_inpaint_sd_xl"] = ["StableDiffusionXLControlNetUnionInpaintPipeline"] - _import_structure["pipeline_controlnet_union_sd_xl"] = ["StableDiffusionXLControlNetUnionPipeline"] - _import_structure["pipeline_controlnet_union_sd_xl_img2img"] = ["StableDiffusionXLControlNetUnionImg2ImgPipeline"] -try: - if not (is_transformers_available() and is_flax_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_flax_and_transformers_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_flax_and_transformers_objects)) -else: - _import_structure["pipeline_flax_controlnet"] = ["FlaxStableDiffusionControlNetPipeline"] - - -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 .multicontrolnet import MultiControlNetModel - from .pipeline_controlnet import StableDiffusionControlNetPipeline - from .pipeline_controlnet_blip_diffusion import BlipDiffusionControlNetPipeline - from .pipeline_controlnet_img2img import StableDiffusionControlNetImg2ImgPipeline - from .pipeline_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline - from .pipeline_controlnet_inpaint_sd_xl import StableDiffusionXLControlNetInpaintPipeline - from .pipeline_controlnet_sd_xl import StableDiffusionXLControlNetPipeline - from .pipeline_controlnet_sd_xl_img2img import StableDiffusionXLControlNetImg2ImgPipeline - from .pipeline_controlnet_union_inpaint_sd_xl import StableDiffusionXLControlNetUnionInpaintPipeline - from .pipeline_controlnet_union_sd_xl import StableDiffusionXLControlNetUnionPipeline - from .pipeline_controlnet_union_sd_xl_img2img import StableDiffusionXLControlNetUnionImg2ImgPipeline - - try: - if not (is_transformers_available() and is_flax_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_flax_and_transformers_objects import * # noqa F403 - else: - from .pipeline_flax_controlnet import FlaxStableDiffusionControlNetPipeline - - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 7ce934037cd6c58cebfc7eb8edebe646d744ed38..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/multicontrolnet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/multicontrolnet.cpython-310.pyc deleted file mode 100644 index 1dc9e8b69ce3f0fb6a444d5db3bca791e8a0fb75..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/multicontrolnet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet.cpython-310.pyc deleted file mode 100644 index 8b96e91b786d70c0d34684a778b844364cb34ca5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_blip_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_blip_diffusion.cpython-310.pyc deleted file mode 100644 index d65295e59d1e2123bce6b51b8279986e4da1aba1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_blip_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_img2img.cpython-310.pyc deleted file mode 100644 index 61bbe51692b613627069fd1f6ef9ca2892974e11..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_inpaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_inpaint.cpython-310.pyc deleted file mode 100644 index 349ea2e46ac1681e83e6d7d5527f1c4c2fe045bb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_inpaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_inpaint_sd_xl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_inpaint_sd_xl.cpython-310.pyc deleted file mode 100644 index b941f8f9f0aaa73cf81dcdce42bfa4cca884e53f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_inpaint_sd_xl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_sd_xl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_sd_xl.cpython-310.pyc deleted file mode 100644 index f8b25d598f735e44227ac326c2d314aeb391e452..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_sd_xl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_sd_xl_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_sd_xl_img2img.cpython-310.pyc deleted file mode 100644 index 31c9f8e01408a0b767ac911142116681c5ae8d06..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_sd_xl_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_union_inpaint_sd_xl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_union_inpaint_sd_xl.cpython-310.pyc deleted file mode 100644 index 8f8e8e4514d91d8835e80b04cff3e330fff65e8f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_union_inpaint_sd_xl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_union_sd_xl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_union_sd_xl.cpython-310.pyc deleted file mode 100644 index c41393e23291e7d724739733830339b5936d9d6e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_union_sd_xl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_union_sd_xl_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_union_sd_xl_img2img.cpython-310.pyc deleted file mode 100644 index 7b211215a683ae9f4e99efac8dac79ff8e2a4576..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_controlnet_union_sd_xl_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_flax_controlnet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_flax_controlnet.cpython-310.pyc deleted file mode 100644 index 3a5924cf1d5ff193a1eb83f4a5a0c52cde23750c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/__pycache__/pipeline_flax_controlnet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/multicontrolnet.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/multicontrolnet.py deleted file mode 100644 index 6526dd8c9a576076db072659c81b457b66af7a35..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/multicontrolnet.py +++ /dev/null @@ -1,12 +0,0 @@ -from ...models.controlnets.multicontrolnet import MultiControlNetModel -from ...utils import deprecate, logging - - -logger = logging.get_logger(__name__) - - -class MultiControlNetModel(MultiControlNetModel): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `MultiControlNetModel` from `diffusers.pipelines.controlnet.multicontrolnet` is deprecated and this will be removed in a future version. Please use `from diffusers.models.controlnets.multicontrolnet import MultiControlNetModel`, instead." - deprecate("diffusers.pipelines.controlnet.multicontrolnet.MultiControlNetModel", "0.34", deprecation_message) - super().__init__(*args, **kwargs) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet.py deleted file mode 100644 index fe0e69314ccaae267e502c1da1c636f51a66eb5b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet.py +++ /dev/null @@ -1,1366 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ControlNetModel, ImageProjection, MultiControlNetModel, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import empty_device_cache, is_compiled_module, is_torch_version, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion.pipeline_output import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -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 - >>> # !pip install opencv-python transformers accelerate - >>> from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler - >>> from diffusers.utils import load_image - >>> import numpy as np - >>> import torch - - >>> import cv2 - >>> from PIL import Image - - >>> # download an image - >>> image = load_image( - ... "https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png" - ... ) - >>> image = np.array(image) - - >>> # get canny image - >>> image = cv2.Canny(image, 100, 200) - >>> image = image[:, :, None] - >>> image = np.concatenate([image, image, image], axis=2) - >>> canny_image = Image.fromarray(image) - - >>> # load control net and stable diffusion v1-5 - >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16) - >>> pipe = StableDiffusionControlNetPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16 - ... ) - - >>> # speed up diffusion process with faster scheduler and memory optimization - >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) - >>> # remove following line if xformers is not installed - >>> pipe.enable_xformers_memory_efficient_attention() - - >>> pipe.enable_model_cpu_offload() - - >>> # generate image - >>> generator = torch.manual_seed(0) - >>> image = pipe( - ... "futuristic-looking woman", num_inference_steps=20, generator=generator, image=canny_image - ... ).images[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, - **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 StableDiffusionControlNetPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - IPAdapterMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion with ControlNet guidance. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - 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. - 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "image"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: CLIPVisionModelWithProjection = None, - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - if isinstance(controlnet, (list, tuple)): - controlnet = MultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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, do_convert_rgb=True) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - image, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_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, - ): - 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}." - ) - - # Check `image` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - self.check_image(image, prompt, prompt_embeds) - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if not isinstance(image, list): - raise TypeError("For multiple controlnets: `image` must be type `list`") - - # When `image` is a nested list: - # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]]) - elif any(isinstance(i, list) for i in image): - transposed_image = [list(t) for t in zip(*image)] - if len(transposed_image) != len(self.controlnet.nets): - raise ValueError( - f"For multiple controlnets: if you pass`image` as a list of list, each sublist must have the same length as the number of controlnets, but the sublists in `image` got {len(transposed_image)} images and {len(self.controlnet.nets)} ControlNets." - ) - for image_ in transposed_image: - self.check_image(image_, prompt, prompt_embeds) - 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." - ) - else: - for image_ in image: - self.check_image(image_, prompt, prompt_embeds) - else: - assert False - - # Check `controlnet_conditioning_scale` - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if isinstance(controlnet_conditioning_scale, list): - if any(isinstance(i, list) for i in controlnet_conditioning_scale): - raise ValueError( - "A single batch of varying conditioning scale settings (e.g. [[1.0, 0.5], [0.2, 0.8]]) is not supported at the moment. " - "The conditioning scale must be fixed across the batch." - ) - elif 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" - ) - else: - assert False - - if not isinstance(control_guidance_start, (tuple, list)): - control_guidance_start = [control_guidance_start] - - if not isinstance(control_guidance_end, (tuple, list)): - control_guidance_end = [control_guidance_end] - - 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(self.controlnet, MultiControlNetModel): - 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" - ) - - 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 prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_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, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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"], - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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]]`): - 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. When `prompt` is a list, and if a list of images is passed for a single - ControlNet, each will be paired with each prompt in the `prompt` list. This also applies to multiple - ControlNets, where a list of image lists can be passed to batch for each prompt and each ControlNet. - 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. - 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*): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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): - 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): - The percentage of total steps at which the ControlNet starts applying. - 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 - the output of the pre-final layer will be used for computing the prompt embeddings. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - image, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - 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 - - if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): - controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) - - global_pool_conditions = ( - controlnet.config.global_pool_conditions - if isinstance(controlnet, ControlNetModel) - else controlnet.nets[0].config.global_pool_conditions - ) - guess_mode = guess_mode or global_pool_conditions - - # 3. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - 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, - self.do_classifier_free_guidance, - ) - - # 4. Prepare image - if isinstance(controlnet, ControlNetModel): - image = self.prepare_image( - image=image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - height, width = image.shape[-2:] - elif isinstance(controlnet, MultiControlNetModel): - images = [] - - # Nested lists as ControlNet condition - if isinstance(image[0], list): - # Transpose the nested image list - image = [list(t) for t in zip(*image)] - - for image_ in image: - image_ = self.prepare_image( - image=image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - - images.append(image_) - - image = images - height, width = image[0].shape[-2:] - else: - assert False - - # 5. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - self._num_timesteps = len(timesteps) - - # 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, - ) - - # 6.5 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) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - # 7.2 Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - is_unet_compiled = is_compiled_module(self.unet) - is_controlnet_compiled = is_compiled_module(self.controlnet) - is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1") - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # Relevant thread: - # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428 - if (is_unet_compiled and is_controlnet_compiled) and is_torch_higher_equal_2_1: - torch._inductor.cudagraph_mark_step_begin() - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # controlnet(s) inference - if guess_mode and self.do_classifier_free_guidance: - # Infer ControlNet only for the conditional batch. - control_model_input = latents - control_model_input = self.scheduler.scale_model_input(control_model_input, t) - controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] - else: - control_model_input = latent_model_input - controlnet_prompt_embeds = prompt_embeds - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=image, - conditioning_scale=cond_scale, - guess_mode=guess_mode, - return_dict=False, - ) - - if guess_mode and self.do_classifier_free_guidance: - # Inferred ControlNet only for the conditional batch. - # To apply the output of ControlNet to both the unconditional and conditional batches, - # add 0 to the unconditional batch to keep it unchanged. - down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] - mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - image = callback_outputs.pop("image", image) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - # If we do sequential model offloading, let's offload unet and controlnet - # manually for max memory savings - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.unet.to("cpu") - self.controlnet.to("cpu") - empty_device_cache() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py deleted file mode 100644 index 7d6a29cecadd2c60ba15d0e26666dd1d336392e8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py +++ /dev/null @@ -1,426 +0,0 @@ -# Copyright 2025 Salesforce.com, inc. -# 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 List, Optional, Union - -import PIL.Image -import torch -from transformers import CLIPTokenizer - -from ...models import AutoencoderKL, ControlNetModel, UNet2DConditionModel -from ...schedulers import PNDMScheduler -from ...utils import ( - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..blip_diffusion.blip_image_processing import BlipImageProcessor -from ..blip_diffusion.modeling_blip2 import Blip2QFormerModel -from ..blip_diffusion.modeling_ctx_clip import ContextCLIPTextModel -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -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 - >>> from diffusers.pipelines import BlipDiffusionControlNetPipeline - >>> from diffusers.utils import load_image - >>> from controlnet_aux import CannyDetector - >>> import torch - - >>> blip_diffusion_pipe = BlipDiffusionControlNetPipeline.from_pretrained( - ... "Salesforce/blipdiffusion-controlnet", torch_dtype=torch.float16 - ... ).to("cuda") - - >>> style_subject = "flower" - >>> tgt_subject = "teapot" - >>> text_prompt = "on a marble table" - - >>> cldm_cond_image = load_image( - ... "https://huggingface.co/datasets/ayushtues/blipdiffusion_images/resolve/main/kettle.jpg" - ... ).resize((512, 512)) - >>> canny = CannyDetector() - >>> cldm_cond_image = canny(cldm_cond_image, 30, 70, output_type="pil") - >>> style_image = load_image( - ... "https://huggingface.co/datasets/ayushtues/blipdiffusion_images/resolve/main/flower.jpg" - ... ) - >>> guidance_scale = 7.5 - >>> num_inference_steps = 50 - >>> negative_prompt = "over-exposure, under-exposure, saturated, duplicate, out of frame, lowres, cropped, worst quality, low quality, jpeg artifacts, morbid, mutilated, out of frame, ugly, bad anatomy, bad proportions, deformed, blurry, duplicate" - - - >>> output = blip_diffusion_pipe( - ... text_prompt, - ... style_image, - ... cldm_cond_image, - ... style_subject, - ... tgt_subject, - ... guidance_scale=guidance_scale, - ... num_inference_steps=num_inference_steps, - ... neg_prompt=negative_prompt, - ... height=512, - ... width=512, - ... ).images - >>> output[0].save("image.png") - ``` -""" - - -class BlipDiffusionControlNetPipeline(DiffusionPipeline): - """ - Pipeline for Canny Edge based Controlled subject-driven generation using Blip 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.) - - Args: - tokenizer ([`CLIPTokenizer`]): - Tokenizer for the text encoder - text_encoder ([`ContextCLIPTextModel`]): - Text encoder to encode the text prompt - vae ([`AutoencoderKL`]): - VAE model to map the latents to the image - unet ([`UNet2DConditionModel`]): - Conditional U-Net architecture to denoise the image embedding. - scheduler ([`PNDMScheduler`]): - A scheduler to be used in combination with `unet` to generate image latents. - qformer ([`Blip2QFormerModel`]): - QFormer model to get multi-modal embeddings from the text and image. - controlnet ([`ControlNetModel`]): - ControlNet model to get the conditioning image embedding. - image_processor ([`BlipImageProcessor`]): - Image Processor to preprocess and postprocess the image. - ctx_begin_pos (int, `optional`, defaults to 2): - Position of the context token in the text encoder. - """ - - model_cpu_offload_seq = "qformer->text_encoder->unet->vae" - - def __init__( - self, - tokenizer: CLIPTokenizer, - text_encoder: ContextCLIPTextModel, - vae: AutoencoderKL, - unet: UNet2DConditionModel, - scheduler: PNDMScheduler, - qformer: Blip2QFormerModel, - controlnet: ControlNetModel, - image_processor: BlipImageProcessor, - ctx_begin_pos: int = 2, - mean: List[float] = None, - std: List[float] = None, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - vae=vae, - unet=unet, - scheduler=scheduler, - qformer=qformer, - controlnet=controlnet, - image_processor=image_processor, - ) - self.register_to_config(ctx_begin_pos=ctx_begin_pos, mean=mean, std=std) - - def get_query_embeddings(self, input_image, src_subject): - return self.qformer(image_input=input_image, text_input=src_subject, return_dict=False) - - # from the original Blip Diffusion code, specifies the target subject and augments the prompt by repeating it - def _build_prompt(self, prompts, tgt_subjects, prompt_strength=1.0, prompt_reps=20): - rv = [] - for prompt, tgt_subject in zip(prompts, tgt_subjects): - prompt = f"a {tgt_subject} {prompt.strip()}" - # a trick to amplify the prompt - rv.append(", ".join([prompt] * int(prompt_strength * prompt_reps))) - - return rv - - # Copied from diffusers.pipelines.consistency_models.pipeline_consistency_models.ConsistencyModelPipeline.prepare_latents - def prepare_latents(self, batch_size, num_channels, height, width, dtype, device, generator, latents=None): - shape = (batch_size, num_channels, 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." - ) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device=device, dtype=dtype) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def encode_prompt(self, query_embeds, prompt, device=None): - device = device or self._execution_device - - # embeddings for prompt, with query_embeds as context - max_len = self.text_encoder.text_model.config.max_position_embeddings - max_len -= self.qformer.config.num_query_tokens - - tokenized_prompt = self.tokenizer( - prompt, - padding="max_length", - truncation=True, - max_length=max_len, - return_tensors="pt", - ).to(device) - - batch_size = query_embeds.shape[0] - ctx_begin_pos = [self.config.ctx_begin_pos] * batch_size - - text_embeddings = self.text_encoder( - input_ids=tokenized_prompt.input_ids, - ctx_embeddings=query_embeds, - ctx_begin_pos=ctx_begin_pos, - )[0] - - return text_embeddings - - # Adapted from diffusers.pipelines.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.prepare_image - def prepare_control_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - ): - image = self.image_processor.preprocess( - image, - size={"width": width, "height": height}, - do_rescale=True, - do_center_crop=False, - do_normalize=False, - return_tensors="pt", - )["pixel_values"].to(device) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance: - image = torch.cat([image] * 2) - - return image - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - 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, - 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] = "", - prompt_strength: float = 1.0, - prompt_reps: int = 20, - output_type: Optional[str] = "pil", - return_dict: bool = True, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - 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]`): - The source subject category. - 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 - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor will ge generated by random sampling. - 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. - height (`int`, *optional*, defaults to 512): - The height of the generated image. - width (`int`, *optional*, defaults to 512): - The width of the generated image. - seed (`int`, *optional*, defaults to 42): - The seed to use for random generation. - 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*): - 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 ""): - 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_strength (`float`, *optional*, defaults to 1.0): - The strength of the prompt. Specifies the number of times the prompt is repeated along with prompt_reps - to amplify the prompt. - prompt_reps (`int`, *optional*, defaults to 20): - The number of times the prompt is repeated along with prompt_strength to amplify the prompt. - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - device = self._execution_device - - reference_image = self.image_processor.preprocess( - reference_image, image_mean=self.config.mean, image_std=self.config.std, return_tensors="pt" - )["pixel_values"] - reference_image = reference_image.to(device) - - if isinstance(prompt, str): - prompt = [prompt] - if isinstance(source_subject_category, str): - source_subject_category = [source_subject_category] - if isinstance(target_subject_category, str): - target_subject_category = [target_subject_category] - - batch_size = len(prompt) - - prompt = self._build_prompt( - prompts=prompt, - tgt_subjects=target_subject_category, - prompt_strength=prompt_strength, - prompt_reps=prompt_reps, - ) - query_embeds = self.get_query_embeddings(reference_image, source_subject_category) - text_embeddings = self.encode_prompt(query_embeds, prompt, device) - # 3. unconditional embedding - do_classifier_free_guidance = guidance_scale > 1.0 - if do_classifier_free_guidance: - max_length = self.text_encoder.text_model.config.max_position_embeddings - - uncond_input = self.tokenizer( - [neg_prompt] * batch_size, - padding="max_length", - max_length=max_length, - return_tensors="pt", - ) - uncond_embeddings = self.text_encoder( - input_ids=uncond_input.input_ids.to(device), - ctx_embeddings=None, - )[0] - # 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 - text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) - scale_down_factor = 2 ** (len(self.unet.config.block_out_channels) - 1) - latents = self.prepare_latents( - batch_size=batch_size, - num_channels=self.unet.config.in_channels, - height=height // scale_down_factor, - width=width // scale_down_factor, - generator=generator, - latents=latents, - dtype=self.unet.dtype, - device=device, - ) - # set timesteps - extra_set_kwargs = {} - self.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs) - - cond_image = self.prepare_control_image( - image=condtioning_image, - width=width, - height=height, - batch_size=batch_size, - num_images_per_prompt=1, - device=device, - dtype=self.controlnet.dtype, - do_classifier_free_guidance=do_classifier_free_guidance, - ) - - for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)): - # expand the latents if we are doing classifier free guidance - do_classifier_free_guidance = guidance_scale > 1.0 - - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - down_block_res_samples, mid_block_res_sample = self.controlnet( - latent_model_input, - t, - encoder_hidden_states=text_embeddings, - controlnet_cond=cond_image, - return_dict=False, - ) - - noise_pred = self.unet( - latent_model_input, - timestep=t, - encoder_hidden_states=text_embeddings, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - )["sample"] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - latents = self.scheduler.step( - noise_pred, - t, - latents, - )["prev_sample"] - - if XLA_AVAILABLE: - xm.mark_step() - - image = self.vae.decode(latents / self.vae.config.scaling_factor, 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 ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py deleted file mode 100644 index 12cc6f630d805658b104dcefd2a40a0fb459c893..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py +++ /dev/null @@ -1,1338 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ControlNetModel, ImageProjection, MultiControlNetModel, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import empty_device_cache, is_compiled_module, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -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 - >>> # !pip install opencv-python transformers accelerate - >>> from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, UniPCMultistepScheduler - >>> from diffusers.utils import load_image - >>> import numpy as np - >>> import torch - - >>> import cv2 - >>> from PIL import Image - - >>> # download an image - >>> image = load_image( - ... "https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png" - ... ) - >>> np_image = np.array(image) - - >>> # get canny image - >>> np_image = cv2.Canny(np_image, 100, 200) - >>> np_image = np_image[:, :, None] - >>> np_image = np.concatenate([np_image, np_image, np_image], axis=2) - >>> canny_image = Image.fromarray(np_image) - - >>> # load control net and stable diffusion v1-5 - >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16) - >>> pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16 - ... ) - - >>> # speed up diffusion process with faster scheduler and memory optimization - >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) - >>> pipe.enable_model_cpu_offload() - - >>> # generate image - >>> generator = torch.manual_seed(0) - >>> image = pipe( - ... "futuristic-looking woman", - ... num_inference_steps=20, - ... generator=generator, - ... image=image, - ... control_image=canny_image, - ... ).images[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" -): - 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") - - -def prepare_image(image): - if isinstance(image, torch.Tensor): - # Batch single image - if image.ndim == 3: - image = image.unsqueeze(0) - - image = image.to(dtype=torch.float32) - else: - # preprocess image - if isinstance(image, (PIL.Image.Image, np.ndarray)): - image = [image] - - if isinstance(image, list) and isinstance(image[0], PIL.Image.Image): - image = [np.array(i.convert("RGB"))[None, :] for i in image] - image = np.concatenate(image, axis=0) - elif isinstance(image, list) and isinstance(image[0], np.ndarray): - image = np.concatenate([i[None, :] for i in image], axis=0) - - image = image.transpose(0, 3, 1, 2) - image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 - - return image - - -class StableDiffusionControlNetImg2ImgPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - IPAdapterMixin, - FromSingleFileMixin, -): - r""" - Pipeline for image-to-image generation using Stable Diffusion with ControlNet guidance. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - 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. - 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "control_image"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: CLIPVisionModelWithProjection = None, - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - if isinstance(controlnet, (list, tuple)): - controlnet = MultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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, do_convert_rgb=True) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - image, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_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, - ): - 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}." - ) - - # `prompt` needs more sophisticated handling when there are multiple - # conditionings. - if isinstance(self.controlnet, MultiControlNetModel): - if isinstance(prompt, list): - 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` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - self.check_image(image, prompt, prompt_embeds) - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if not isinstance(image, list): - raise TypeError("For multiple controlnets: `image` must be type `list`") - - # When `image` is a nested list: - # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]]) - elif any(isinstance(i, list) for i in image): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - 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) - else: - assert False - - # Check `controlnet_conditioning_scale` - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if isinstance(controlnet_conditioning_scale, list): - if any(isinstance(i, list) for i in controlnet_conditioning_scale): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - elif 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" - ) - else: - assert False - - 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(self.controlnet, MultiControlNetModel): - 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.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.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}" - ) - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.prepare_image - def prepare_control_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.prepare_latents - def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - 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." - ) - - elif isinstance(generator, list): - if image.shape[0] < batch_size and batch_size % image.shape[0] == 0: - image = torch.cat([image] * (batch_size // image.shape[0]), dim=0) - elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} " - ) - - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - deprecation_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial" - " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" - " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" - " your script to pass as many initial images as text prompts to suppress this warning." - ) - deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - latents = init_latents - - return latents - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_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, - control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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"], - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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]]`): - 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]]`): - 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. - 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. - strength (`float`, *optional*, defaults to 0.8): - 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`. - 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): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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): - 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): - The percentage of total steps at which the ControlNet starts applying. - 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 - the output of the pre-final layer will be used for computing the prompt embeddings. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - control_image, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - 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 - - if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): - controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) - - global_pool_conditions = ( - controlnet.config.global_pool_conditions - if isinstance(controlnet, ControlNetModel) - else controlnet.nets[0].config.global_pool_conditions - ) - guess_mode = guess_mode or global_pool_conditions - - # 3. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - 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, - self.do_classifier_free_guidance, - ) - - # 4. Prepare image - image = self.image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - - # 5. Prepare controlnet_conditioning_image - if isinstance(controlnet, ControlNetModel): - control_image = self.prepare_control_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - elif isinstance(controlnet, MultiControlNetModel): - control_images = [] - - for control_image_ in control_image: - control_image_ = self.prepare_control_image( - image=control_image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - - control_images.append(control_image_) - - control_image = control_images - else: - assert False - - # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - self._num_timesteps = len(timesteps) - - # 6. Prepare latent variables - if latents is None: - latents = self.prepare_latents( - image, - latent_timestep, - batch_size, - num_images_per_prompt, - prompt_embeds.dtype, - device, - generator, - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - # 7.2 Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # controlnet(s) inference - if guess_mode and self.do_classifier_free_guidance: - # Infer ControlNet only for the conditional batch. - control_model_input = latents - control_model_input = self.scheduler.scale_model_input(control_model_input, t) - controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] - else: - control_model_input = latent_model_input - controlnet_prompt_embeds = prompt_embeds - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=control_image, - conditioning_scale=cond_scale, - guess_mode=guess_mode, - return_dict=False, - ) - - if guess_mode and self.do_classifier_free_guidance: - # Inferred ControlNet only for the conditional batch. - # To apply the output of ControlNet to both the unconditional and conditional batches, - # add 0 to the unconditional batch to keep it unchanged. - down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] - mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=self.cross_attention_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - control_image = callback_outputs.pop("control_image", control_image) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # If we do sequential model offloading, let's offload unet and controlnet - # manually for max memory savings - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.unet.to("cpu") - self.controlnet.to("cpu") - empty_device_cache() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py deleted file mode 100644 index 41303d9c5c5a38b1c7c4b2ed32622a5744753934..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py +++ /dev/null @@ -1,1530 +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. - -# 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 - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ControlNetModel, ImageProjection, MultiControlNetModel, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import empty_device_cache, is_compiled_module, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -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 - >>> # !pip install transformers accelerate - >>> from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, DDIMScheduler - >>> from diffusers.utils import load_image - >>> import numpy as np - >>> import torch - - >>> init_image = load_image( - ... "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy.png" - ... ) - >>> init_image = init_image.resize((512, 512)) - - >>> generator = torch.Generator(device="cpu").manual_seed(1) - - >>> mask_image = load_image( - ... "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy_mask.png" - ... ) - >>> mask_image = mask_image.resize((512, 512)) - - - >>> def make_canny_condition(image): - ... image = np.array(image) - ... image = cv2.Canny(image, 100, 200) - ... image = image[:, :, None] - ... image = np.concatenate([image, image, image], axis=2) - ... image = Image.fromarray(image) - ... return image - - - >>> control_image = make_canny_condition(init_image) - - >>> controlnet = ControlNetModel.from_pretrained( - ... "lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16 - ... ) - >>> pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16 - ... ) - - >>> pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) - >>> pipe.enable_model_cpu_offload() - - >>> # generate image - >>> image = pipe( - ... "a handsome man with ray-ban sunglasses", - ... num_inference_steps=20, - ... generator=generator, - ... eta=1.0, - ... image=init_image, - ... mask_image=mask_image, - ... control_image=control_image, - ... ).images[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" -): - 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 StableDiffusionControlNetInpaintPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - IPAdapterMixin, - FromSingleFileMixin, -): - r""" - Pipeline for image inpainting using Stable Diffusion with ControlNet guidance. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters - - - - This pipeline can be used with checkpoints that have been specifically fine-tuned for inpainting - ([stable-diffusion-v1-5/stable-diffusion-inpainting](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-inpainting)) - as well as default text-to-image Stable Diffusion checkpoints - ([stable-diffusion-v1-5/stable-diffusion-v1-5](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5)). - Default text-to-image Stable Diffusion checkpoints might be preferable for ControlNets that have been fine-tuned on - those, such as [lllyasviel/control_v11p_sd15_inpaint](https://huggingface.co/lllyasviel/control_v11p_sd15_inpaint). - - - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - 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. - 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "control_image", - "mask", - "masked_image_latents", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: CLIPVisionModelWithProjection = None, - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - if isinstance(controlnet, (list, tuple)): - controlnet = MultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True - ) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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, - height, - width, - callback_steps, - output_type, - negative_prompt=None, - prompt_embeds=None, - negative_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, - padding_mask_crop=None, - ): - if height is not None and height % 8 != 0 or width is not None and 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}." - ) - - 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}.") - - # `prompt` needs more sophisticated handling when there are multiple - # conditionings. - if isinstance(self.controlnet, MultiControlNetModel): - if isinstance(prompt, list): - 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` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - self.check_image(image, prompt, prompt_embeds) - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if not isinstance(image, list): - raise TypeError("For multiple controlnets: `image` must be type `list`") - - # When `image` is a nested list: - # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]]) - elif any(isinstance(i, list) for i in image): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - 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) - else: - assert False - - # Check `controlnet_conditioning_scale` - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if isinstance(controlnet_conditioning_scale, list): - if any(isinstance(i, list) for i in controlnet_conditioning_scale): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - elif 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" - ) - else: - assert False - - 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(self.controlnet, MultiControlNetModel): - 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.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.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 prepare_control_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - crops_coords, - resize_mode, - do_classifier_free_guidance=False, - guess_mode=False, - ): - image = self.control_image_processor.preprocess( - image, height=height, width=width, crops_coords=crops_coords, resize_mode=resize_mode - ).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.StableDiffusionInpaintPipeline.prepare_latents - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=None, - image=None, - timestep=None, - is_strength_max=True, - return_noise=False, - return_image_latents=False, - ): - 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 (image is None or timestep is None) and not is_strength_max: - raise ValueError( - "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise." - "However, either the image or the noise timestep has not been provided." - ) - - if return_image_latents or (latents is None and not is_strength_max): - image = image.to(device=device, dtype=dtype) - - if image.shape[1] == 4: - image_latents = image - else: - image_latents = self._encode_vae_image(image=image, generator=generator) - image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1) - - if latents is None: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # if strength is 1. then initialise the latents to noise, else initial to image + noise - latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep) - # if pure noise then scale the initial latents by the Scheduler's init sigma - latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents - else: - noise = latents.to(device) - latents = noise * self.scheduler.init_noise_sigma - - outputs = (latents,) - - if return_noise: - outputs += (noise,) - - if return_image_latents: - outputs += (image_latents,) - - return outputs - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.StableDiffusionInpaintPipeline.prepare_mask_latents - def prepare_mask_latents( - self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance - ): - # 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 // self.vae_scale_factor, width // self.vae_scale_factor) - ) - mask = mask.to(device=device, dtype=dtype) - - masked_image = masked_image.to(device=device, dtype=dtype) - - if masked_image.shape[1] == 4: - masked_image_latents = masked_image - else: - masked_image_latents = self._encode_vae_image(masked_image, generator=generator) - - # 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) - - mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask - masked_image_latents = ( - torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents - ) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.StableDiffusionInpaintPipeline._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 = self.vae.config.scaling_factor * image_latents - - return image_latents - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_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, - control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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"], - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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`, 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]`): - `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]]`): - 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. - 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. - padding_mask_crop (`int`, *optional*, defaults to `None`): - The size of margin in the crop to be applied to the image and masking. If `None`, no crop is applied to - image and mask_image. If `padding_mask_crop` is not `None`, it will first find a rectangular region - with the same aspect ration of the image and contains all masked area, and then expand that area based - on `padding_mask_crop`. The image and mask_image will then be cropped based on the expanded area before - resizing to the original image size for inpainting. This is useful when the masked area is small while - the image is large and contain information irrelevant for inpainting, such as background. - 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`. - 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): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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): - 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): - The percentage of total steps at which the ControlNet starts applying. - 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 - the output of the pre-final layer will be used for computing the prompt embeddings. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - control_image, - mask_image, - height, - width, - callback_steps, - output_type, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - callback_on_step_end_tensor_inputs, - padding_mask_crop, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - 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] - - if padding_mask_crop is not None: - height, width = self.image_processor.get_default_height_width(image, height, width) - 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" - - device = self._execution_device - - if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): - controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) - - global_pool_conditions = ( - controlnet.config.global_pool_conditions - if isinstance(controlnet, ControlNetModel) - else controlnet.nets[0].config.global_pool_conditions - ) - guess_mode = guess_mode or global_pool_conditions - - # 3. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - 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, - self.do_classifier_free_guidance, - ) - - # 4. Prepare image - if isinstance(controlnet, ControlNetModel): - control_image = self.prepare_control_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - crops_coords=crops_coords, - resize_mode=resize_mode, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - elif isinstance(controlnet, MultiControlNetModel): - control_images = [] - - for control_image_ in control_image: - control_image_ = self.prepare_control_image( - image=control_image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - crops_coords=crops_coords, - resize_mode=resize_mode, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - - control_images.append(control_image_) - - control_image = control_images - else: - assert False - - # 4.1 Preprocess mask and image - resizes image and mask w.r.t height and width - 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) - - mask = self.mask_processor.preprocess( - mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords - ) - - masked_image = init_image * (mask < 0.5) - _, _, height, width = init_image.shape - - # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps( - num_inference_steps=num_inference_steps, strength=strength, device=device - ) - # at which timestep to set the initial noise (n.b. 50% if strength is 0.5) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise - is_strength_max = strength == 1.0 - self._num_timesteps = len(timesteps) - - # 6. Prepare latent variables - num_channels_latents = self.vae.config.latent_channels - num_channels_unet = self.unet.config.in_channels - return_image_latents = num_channels_unet == 4 - latents_outputs = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - image=init_image, - timestep=latent_timestep, - is_strength_max=is_strength_max, - return_noise=True, - return_image_latents=return_image_latents, - ) - - if return_image_latents: - latents, noise, image_latents = latents_outputs - else: - latents, noise = latents_outputs - - # 7. Prepare mask latent variables - mask, masked_image_latents = self.prepare_mask_latents( - mask, - masked_image, - batch_size * num_images_per_prompt, - height, - width, - prompt_embeds.dtype, - device, - generator, - self.do_classifier_free_guidance, - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - # 7.2 Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # controlnet(s) inference - if guess_mode and self.do_classifier_free_guidance: - # Infer ControlNet only for the conditional batch. - control_model_input = latents - control_model_input = self.scheduler.scale_model_input(control_model_input, t) - controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] - else: - control_model_input = latent_model_input - controlnet_prompt_embeds = prompt_embeds - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=control_image, - conditioning_scale=cond_scale, - guess_mode=guess_mode, - return_dict=False, - ) - - if guess_mode and self.do_classifier_free_guidance: - # Inferred ControlNet only for the conditional batch. - # To apply the output of ControlNet to both the unconditional and conditional batches, - # add 0 to the unconditional batch to keep it unchanged. - down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] - mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) - - # predict the noise residual - if num_channels_unet == 9: - latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) - - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=self.cross_attention_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - if num_channels_unet == 4: - init_latents_proper = image_latents - if self.do_classifier_free_guidance: - init_mask, _ = mask.chunk(2) - else: - init_mask = mask - - if i < len(timesteps) - 1: - noise_timestep = timesteps[i + 1] - init_latents_proper = self.scheduler.add_noise( - init_latents_proper, noise, torch.tensor([noise_timestep]) - ) - - latents = (1 - init_mask) * init_latents_proper + init_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) - control_image = callback_outputs.pop("control_image", control_image) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # If we do sequential model offloading, let's offload unet and controlnet - # manually for max memory savings - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.unet.to("cpu") - self.controlnet.to("cpu") - empty_device_cache() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 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) - - if padding_mask_crop is not None: - image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] - - # 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/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py deleted file mode 100644 index 4aa2a62a53ac4aba3bab17e1a536ec2e20aa0bff..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py +++ /dev/null @@ -1,1883 +0,0 @@ -# Copyright 2025 Harutatsu Akiyama, Jinbin Bai, 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, Tuple, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ControlNetModel, ImageProjection, MultiControlNetModel, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_invisible_watermark_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import empty_device_cache, is_compiled_module, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput - - -if is_invisible_watermark_available(): - from diffusers.pipelines.stable_diffusion_xl.watermark import StableDiffusionXLWatermarker - - -from ...utils import is_torch_xla_available - - -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 - - -# 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" -): - 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") - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> # !pip install transformers accelerate - >>> from diffusers import StableDiffusionXLControlNetInpaintPipeline, ControlNetModel, DDIMScheduler - >>> from diffusers.utils import load_image - >>> from PIL import Image - >>> import numpy as np - >>> import torch - - >>> init_image = load_image( - ... "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy.png" - ... ) - >>> init_image = init_image.resize((1024, 1024)) - - >>> generator = torch.Generator(device="cpu").manual_seed(1) - - >>> mask_image = load_image( - ... "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy_mask.png" - ... ) - >>> mask_image = mask_image.resize((1024, 1024)) - - - >>> def make_canny_condition(image): - ... image = np.array(image) - ... image = cv2.Canny(image, 100, 200) - ... image = image[:, :, None] - ... image = np.concatenate([image, image, image], axis=2) - ... image = Image.fromarray(image) - ... return image - - - >>> control_image = make_canny_condition(init_image) - - >>> controlnet = ControlNetModel.from_pretrained( - ... "diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16 - ... ) - >>> pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", controlnet=controlnet, torch_dtype=torch.float16 - ... ) - - >>> pipe.enable_model_cpu_offload() - - >>> # generate image - >>> image = pipe( - ... "a handsome man with ray-ban sunglasses", - ... num_inference_steps=20, - ... generator=generator, - ... eta=1.0, - ... image=init_image, - ... mask_image=mask_image, - ... control_image=control_image, - ... ).images[0] - ``` -""" - - -# 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 StableDiffusionXLControlNetInpaintPipeline( - DiffusionPipeline, - StableDiffusionMixin, - StableDiffusionXLLoraLoaderMixin, - FromSingleFileMixin, - IPAdapterMixin, - TextualInversionLoaderMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion XL. - - 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.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] for loading LoRA weights - - [`~loaders.StableDiffusionXLLoraLoaderMixin.save_lora_weights`] for saving LoRA weights - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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`]. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae" - - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "image_encoder", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "add_neg_time_ids", - "mask", - "masked_image_latents", - "control_image", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - controlnet: Union[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, - ): - super().__init__() - - if isinstance(controlnet, (list, tuple)): - controlnet = MultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - self.register_to_config(requires_aesthetics_score=requires_aesthetics_score) - 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.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True - ) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - 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, - prompt, - prompt_2, - image, - mask_image, - strength, - num_inference_steps, - callback_steps, - output_type, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - controlnet_conditioning_scale=1.0, - control_guidance_start=0.0, - control_guidance_end=1.0, - callback_on_step_end_tensor_inputs=None, - padding_mask_crop=None, - ): - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - if num_inference_steps is None: - raise ValueError("`num_inference_steps` cannot be None.") - elif not isinstance(num_inference_steps, int) or num_inference_steps <= 0: - raise ValueError( - f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type" - f" {type(num_inference_steps)}." - ) - - 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_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 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 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`." - ) - - # `prompt` needs more sophisticated handling when there are multiple - # conditionings. - if isinstance(self.controlnet, MultiControlNetModel): - if isinstance(prompt, list): - 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` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - self.check_image(image, prompt, prompt_embeds) - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if not isinstance(image, list): - raise TypeError("For multiple controlnets: `image` must be type `list`") - - # When `image` is a nested list: - # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]]) - elif any(isinstance(i, list) for i in image): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - 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) - else: - assert False - - # Check `controlnet_conditioning_scale` - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if isinstance(controlnet_conditioning_scale, list): - if any(isinstance(i, list) for i in controlnet_conditioning_scale): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - elif 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" - ) - else: - assert False - - if not isinstance(control_guidance_start, (tuple, list)): - control_guidance_start = [control_guidance_start] - - if not isinstance(control_guidance_end, (tuple, list)): - control_guidance_end = [control_guidance_end] - - 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(self.controlnet, MultiControlNetModel): - 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" - ) - - def prepare_control_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - crops_coords, - resize_mode, - do_classifier_free_guidance=False, - guess_mode=False, - ): - image = self.control_image_processor.preprocess( - image, height=height, width=width, crops_coords=crops_coords, resize_mode=resize_mode - ).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=None, - image=None, - timestep=None, - is_strength_max=True, - add_noise=True, - return_noise=False, - return_image_latents=False, - ): - 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 (image is None or timestep is None) and not is_strength_max: - raise ValueError( - "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise." - "However, either the image or the noise timestep has not been provided." - ) - - if return_image_latents or (latents is None and not is_strength_max): - image = image.to(device=device, dtype=dtype) - - if image.shape[1] == 4: - image_latents = image - else: - image_latents = self._encode_vae_image(image=image, generator=generator) - image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1) - - if latents is None and add_noise: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # if strength is 1. then initialise the latents to noise, else initial to image + noise - latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep) - # if pure noise then scale the initial latents by the Scheduler's init sigma - latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents - elif add_noise: - noise = latents.to(device) - latents = noise * self.scheduler.init_noise_sigma - else: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = image_latents.to(device) - - outputs = (latents,) - - if return_noise: - outputs += (noise,) - - if return_image_latents: - outputs += (image_latents,) - - return outputs - - def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): - dtype = image.dtype - if self.vae.config.force_upcast: - image = image.float() - self.vae.to(dtype=torch.float32) - - 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) - - if self.vae.config.force_upcast: - self.vae.to(dtype) - - image_latents = image_latents.to(dtype) - image_latents = self.vae.config.scaling_factor * image_latents - - return image_latents - - def prepare_mask_latents( - self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance - ): - # 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 // self.vae_scale_factor, width // self.vae_scale_factor) - ) - mask = mask.to(device=device, dtype=dtype) - - # 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) - - mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask - - masked_image_latents = None - if masked_image is not None: - masked_image = masked_image.to(device=device, dtype=dtype) - masked_image_latents = self._encode_vae_image(masked_image, generator=generator) - 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 - ) - - masked_image_latents = ( - torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents - ) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device, denoising_start=None): - # get the original timestep using init_timestep - if denoising_start is None: - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - t_start = 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 - - else: - # Strength is irrelevant if we directly request a timestep to start at; - # that is, strength is determined by the denoising_start instead. - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (denoising_start * self.scheduler.config.num_train_timesteps) - ) - ) - - num_inference_steps = (self.scheduler.timesteps < discrete_timestep_cutoff).sum().item() - if self.scheduler.order == 2 and num_inference_steps % 2 == 0: - # if the scheduler is a 2nd order scheduler we might have to do +1 - # because `num_inference_steps` might be even given that every timestep - # (except the highest one) is duplicated. If `num_inference_steps` is even it would - # mean that we cut the timesteps in the middle of the denoising step - # (between 1st and 2nd derivative) which leads to incorrect results. By adding 1 - # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler - num_inference_steps = num_inference_steps + 1 - - # because t_n+1 >= t_n, we slice the timesteps starting from the end - t_start = len(self.scheduler.timesteps) - num_inference_steps - timesteps = self.scheduler.timesteps[t_start:] - if hasattr(self.scheduler, "set_begin_index"): - self.scheduler.set_begin_index(t_start) - return timesteps, num_inference_steps - - def _get_add_time_ids( - self, - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - dtype, - text_encoder_projection_dim=None, - ): - if self.config.requires_aesthetics_score: - add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,)) - add_neg_time_ids = list(original_size + crops_coords_top_left + (negative_aesthetic_score,)) - else: - add_time_ids = list(original_size + crops_coords_top_left + target_size) - add_neg_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 - and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model." - ) - elif ( - expected_add_embed_dim < passed_add_embed_dim - and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model." - ) - elif 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) - add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype) - - return add_time_ids, add_neg_time_ids - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_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, - prompt_2: Optional[Union[str, List[str]]] = 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, - strength: float = 0.9999, - num_inference_steps: int = 50, - denoising_start: Optional[float] = None, - denoising_end: Optional[float] = 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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, - 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, - 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"], - **kwargs, - ): - 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 - image (`PIL.Image.Image`): - `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will - be masked out with `mask_image` and repainted according to `prompt`. - mask_image (`PIL.Image.Image`): - `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be - 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)`. - 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. - padding_mask_crop (`int`, *optional*, defaults to `None`): - The size of margin in the crop to be applied to the image and masking. If `None`, no crop is applied to - image and mask_image. If `padding_mask_crop` is not `None`, it will first find a rectangular region - with the same aspect ration of the image and contains all masked area, and then expand that area based - on `padding_mask_crop`. The image and mask_image will then be cropped based on the expanded area before - resizing to the original image size for inpainting. This is useful when the masked area is small while - the image is large and contain information irrelevant for inpainting, such as background. - strength (`float`, *optional*, defaults to 0.9999): - Conceptually, indicates how much to transform the masked portion of 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 the masked - portion of the reference `image`. Note that in the case of `denoising_start` being declared as an - integer, the value of `strength` will be ignored. - 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. - denoising_start (`float`, *optional*): - When specified, indicates the fraction (between 0.0 and 1.0) of the total denoising process to be - bypassed before it is initiated. Consequently, the initial part of the denoising process is skipped and - it is assumed that the passed `image` is a partly denoised image. Note that when this is specified, - strength will be ignored. The `denoising_start` parameter is particularly beneficial when this pipeline - is integrated into a "Mixture of Denoisers" multi-pipeline setup, as detailed in [**Refining the Image - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output). - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - still retain a substantial amount of noise (ca. final 20% of timesteps still needed) and should be - denoised by a successor pipeline that has `denoising_start` set to 0.8 so that it only denoises the - final 20% of 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output). - 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`). - 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. - 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)`. 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. - 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. - 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 ge generated by sampling using the supplied random `generator`. - 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. - cross_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). - 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` 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 `(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). - aesthetic_score (`float`, *optional*, defaults to 6.0): - Used to simulate an aesthetic score of the generated image by influencing the positive text condition. - 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_aesthetic_score (`float`, *optional*, defaults to 2.5): - 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). Can be used to - simulate an aesthetic score of the generated image by influencing the negative text condition. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a - `tuple. `tuple. When returning a tuple, the first element is a list with the generated images. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # # 0.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 - - # 0.1 align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 1. Check inputs - self.check_inputs( - prompt, - prompt_2, - control_image, - mask_image, - strength, - num_inference_steps, - callback_steps, - output_type, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - callback_on_step_end_tensor_inputs, - padding_mask_crop, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - 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 - - if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): - controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) - - # 3. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else 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=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - - # 3.1 Encode ip_adapter_image - 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, - self.do_classifier_free_guidance, - ) - - # 4. set timesteps - def denoising_value_valid(dnv): - return isinstance(dnv, float) and 0 < dnv < 1 - - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps( - num_inference_steps, - strength, - device, - denoising_start=denoising_start if denoising_value_valid(denoising_start) else None, - ) - # check that number of inference steps is not < 1 - as this doesn't make sense - 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." - ) - # at which timestep to set the initial noise (n.b. 50% if strength is 0.5) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise - is_strength_max = strength == 1.0 - self._num_timesteps = len(timesteps) - - # 5. Preprocess mask and image - resizes image and mask w.r.t height and width - # 5.1 Prepare init image - if padding_mask_crop is not None: - height, width = self.image_processor.get_default_height_width(image, height, width) - 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) - - # 5.2 Prepare control images - if isinstance(controlnet, ControlNetModel): - control_image = self.prepare_control_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - crops_coords=crops_coords, - resize_mode=resize_mode, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - elif isinstance(controlnet, MultiControlNetModel): - control_images = [] - - for control_image_ in control_image: - control_image_ = self.prepare_control_image( - image=control_image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - crops_coords=crops_coords, - resize_mode=resize_mode, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - - control_images.append(control_image_) - - control_image = control_images - else: - raise ValueError(f"{controlnet.__class__} is not supported.") - - # 5.3 Prepare mask - mask = self.mask_processor.preprocess( - mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords - ) - - masked_image = init_image * (mask < 0.5) - _, _, height, width = init_image.shape - - # 6. Prepare latent variables - num_channels_latents = self.vae.config.latent_channels - num_channels_unet = self.unet.config.in_channels - return_image_latents = num_channels_unet == 4 - - add_noise = True if denoising_start is None else False - latents_outputs = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - image=init_image, - timestep=latent_timestep, - is_strength_max=is_strength_max, - add_noise=add_noise, - return_noise=True, - return_image_latents=return_image_latents, - ) - - if return_image_latents: - latents, noise, image_latents = latents_outputs - else: - latents, noise = latents_outputs - - # 7. Prepare mask latent variables - mask, masked_image_latents = self.prepare_mask_latents( - mask, - masked_image, - batch_size * num_images_per_prompt, - height, - width, - prompt_embeds.dtype, - device, - generator, - self.do_classifier_free_guidance, - ) - - # 8. Check that sizes of mask, masked image and latents match - if num_channels_unet == 9: - # default case for stable-diffusion-v1-5/stable-diffusion-inpainting - num_channels_mask = mask.shape[1] - num_channels_masked_image = masked_image_latents.shape[1] - if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" - " `pipeline.unet` or your `mask_image` or `image` input." - ) - elif num_channels_unet != 4: - raise ValueError( - f"The unet {self.unet.__class__} should have either 4 or 9 input channels, not {self.unet.config.in_channels}." - ) - # 8.1 Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 8.2 Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps if isinstance(controlnet, MultiControlNetModel) else keeps[0]) - - # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - height, width = latents.shape[-2:] - height = height * self.vae_scale_factor - width = width * self.vae_scale_factor - - original_size = original_size or (height, width) - target_size = target_size or (height, width) - - # 10. 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, add_neg_time_ids = self._get_add_time_ids( - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - dtype=prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1) - - 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_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1) - add_time_ids = torch.cat([add_neg_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) - - # 11. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - if ( - denoising_end is not None - and denoising_start is not None - and denoising_value_valid(denoising_end) - and denoising_value_valid(denoising_start) - and denoising_start >= denoising_end - ): - raise ValueError( - f"`denoising_start`: {denoising_start} cannot be larger than or equal to `denoising_end`: " - + f" {denoising_end} when using type float." - ) - elif denoising_end is not None and denoising_value_valid(denoising_end): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - # concat latents, mask, masked_image_latents in the channel dimension - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - - # controlnet(s) inference - if guess_mode and self.do_classifier_free_guidance: - # Infer ControlNet only for the conditional batch. - control_model_input = latents - control_model_input = self.scheduler.scale_model_input(control_model_input, t) - controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] - controlnet_added_cond_kwargs = { - "text_embeds": add_text_embeds.chunk(2)[1], - "time_ids": add_time_ids.chunk(2)[1], - } - else: - control_model_input = latent_model_input - controlnet_prompt_embeds = prompt_embeds - controlnet_added_cond_kwargs = added_cond_kwargs - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - # # Resize control_image to match the size of the input to the controlnet - # if control_image.shape[-2:] != control_model_input.shape[-2:]: - # control_image = F.interpolate(control_image, size=control_model_input.shape[-2:], mode="bilinear", align_corners=False) - - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=control_image, - conditioning_scale=cond_scale, - guess_mode=guess_mode, - added_cond_kwargs=controlnet_added_cond_kwargs, - return_dict=False, - ) - - if guess_mode and self.do_classifier_free_guidance: - # Inferred ControlNet only for the conditional batch. - # To apply the output of ControlNet to both the unconditional and conditional batches, - # add 0 to the unconditional batch to keep it unchanged. - down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] - mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = image_embeds - - if num_channels_unet == 9: - latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=self.cross_attention_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - if self.do_classifier_free_guidance and guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - if num_channels_unet == 4: - init_latents_proper = image_latents - if self.do_classifier_free_guidance: - init_mask, _ = mask.chunk(2) - else: - init_mask = mask - - if i < len(timesteps) - 1: - noise_timestep = timesteps[i + 1] - init_latents_proper = self.scheduler.add_noise( - init_latents_proper, noise, torch.tensor([noise_timestep]) - ) - - latents = (1 - init_mask) * init_latents_proper + init_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) - control_image = callback_outputs.pop("control_image", control_image) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # make sure the VAE is in float32 mode, as it overflows in float16 - if self.vae.dtype == torch.float16 and self.vae.config.force_upcast: - self.upcast_vae() - latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype) - - # If we do sequential model offloading, let's offload unet and controlnet - # manually for max memory savings - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.unet.to("cpu") - self.controlnet.to("cpu") - empty_device_cache() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - else: - return StableDiffusionXLPipelineOutput(images=latents) - - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - image = self.image_processor.postprocess(image, output_type=output_type) - - if padding_mask_crop is not None: - image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return StableDiffusionXLPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py deleted file mode 100644 index 8b37d38b909924d308ca343eb78a0cb98955262a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py +++ /dev/null @@ -1,1616 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from diffusers.utils.import_utils import is_invisible_watermark_available - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ControlNetModel, ImageProjection, MultiControlNetModel, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import is_compiled_module, is_torch_version, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput - - -if is_invisible_watermark_available(): - from ..stable_diffusion_xl.watermark import StableDiffusionXLWatermarker - - -from ...utils import is_torch_xla_available - - -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 - >>> # !pip install opencv-python transformers accelerate - >>> from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel, AutoencoderKL - >>> from diffusers.utils import load_image - >>> import numpy as np - >>> import torch - - >>> import cv2 - >>> from PIL import Image - - >>> prompt = "aerial view, a futuristic research complex in a bright foggy jungle, hard lighting" - >>> negative_prompt = "low quality, bad quality, sketches" - - >>> # download an image - >>> image = load_image( - ... "https://hf.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/hf-logo.png" - ... ) - - >>> # initialize the models and pipeline - >>> controlnet_conditioning_scale = 0.5 # recommended for good generalization - >>> controlnet = ControlNetModel.from_pretrained( - ... "diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16 - ... ) - >>> vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) - >>> pipe = StableDiffusionXLControlNetPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", controlnet=controlnet, vae=vae, torch_dtype=torch.float16 - ... ) - >>> pipe.enable_model_cpu_offload() - - >>> # get canny image - >>> image = np.array(image) - >>> image = cv2.Canny(image, 100, 200) - >>> image = image[:, :, None] - >>> image = np.concatenate([image, image, image], axis=2) - >>> canny_image = Image.fromarray(image) - - >>> # generate image - >>> image = pipe( - ... prompt, controlnet_conditioning_scale=controlnet_conditioning_scale, image=canny_image - ... ).images[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, - **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 StableDiffusionXLControlNetPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionXLLoraLoaderMixin, - IPAdapterMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion XL with ControlNet guidance. - - 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.). - - The pipeline also inherits the following loading methods: - - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings - - [`~loaders.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] for loading LoRA weights - - [`~loaders.StableDiffusionXLLoraLoaderMixin.save_lora_weights`] for saving LoRA weights - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - text_encoder_2 ([`~transformers.CLIPTextModelWithProjection`]): - Second frozen text-encoder - ([laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - tokenizer_2 ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - 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. - 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 should always be set to 0. Also see the config of - `stabilityai/stable-diffusion-xl-base-1-0`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark](https://github.com/ShieldMnt/invisible-watermark/) library to - watermark output images. If not defined, it defaults to `True` if the package is installed; otherwise no - watermarker is used. - """ - - # leave controlnet out on purpose because it iterates with unet - model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "feature_extractor", - "image_encoder", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "negative_add_time_ids", - "image", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], - scheduler: KarrasDiffusionSchedulers, - force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - feature_extractor: CLIPImageProcessor = None, - image_encoder: CLIPVisionModelWithProjection = None, - ): - super().__init__() - - if isinstance(controlnet, (list, tuple)): - controlnet = MultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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, do_convert_rgb=True) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - prompt_2, - image, - callback_steps, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - negative_pooled_prompt_embeds=None, - controlnet_conditioning_scale=1.0, - control_guidance_start=0.0, - control_guidance_end=1.0, - callback_on_step_end_tensor_inputs=None, - ): - 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_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`." - ) - - # `prompt` needs more sophisticated handling when there are multiple - # conditionings. - if isinstance(self.controlnet, MultiControlNetModel): - if isinstance(prompt, list): - 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` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - self.check_image(image, prompt, prompt_embeds) - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if not isinstance(image, list): - raise TypeError("For multiple controlnets: `image` must be type `list`") - - # When `image` is a nested list: - # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]]) - elif any(isinstance(i, list) for i in image): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - 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) - else: - assert False - - # Check `controlnet_conditioning_scale` - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if isinstance(controlnet_conditioning_scale, list): - if any(isinstance(i, list) for i in controlnet_conditioning_scale): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - elif 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" - ) - else: - assert False - - if not isinstance(control_guidance_start, (tuple, list)): - control_guidance_start = [control_guidance_start] - - if not isinstance(control_guidance_end, (tuple, list)): - control_guidance_end = [control_guidance_end] - - 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(self.controlnet, MultiControlNetModel): - 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.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.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}" - ) - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.prepare_image - def prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - 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.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_end(self): - return self._denoising_end - - @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, - prompt_2: Optional[Union[str, List[str]]] = None, - image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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"], - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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*): - 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]]`): - 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. - height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The height in pixels of the generated image. 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. 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. - 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*): - 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. - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output) - 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*): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - pooled_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated pooled text embeddings. Can be used to easily tweak text inputs (prompt weighting). If - not provided, pooled text embeddings are 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 (prompt - 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*): - 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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): - 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): - The percentage of total steps at which the ControlNet starts applying. - 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)): - 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. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned, - otherwise a `tuple` is returned containing the output images. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - image, - callback_steps, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - negative_pooled_prompt_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - 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 - - if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): - controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) - - global_pool_conditions = ( - controlnet.config.global_pool_conditions - if isinstance(controlnet, ControlNetModel) - else controlnet.nets[0].config.global_pool_conditions - ) - guess_mode = guess_mode or global_pool_conditions - - # 3.1 Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt, - prompt_2, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - 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=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - - # 3.2 Encode ip_adapter_image - 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, - self.do_classifier_free_guidance, - ) - - # 4. Prepare image - if isinstance(controlnet, ControlNetModel): - image = self.prepare_image( - image=image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - height, width = image.shape[-2:] - elif isinstance(controlnet, MultiControlNetModel): - images = [] - - for image_ in image: - image_ = self.prepare_image( - image=image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - - images.append(image_) - - image = images - height, width = image[0].shape[-2:] - else: - assert False - - # 5. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - self._num_timesteps = len(timesteps) - - # 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, - ) - - # 6.5 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) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) - - # 7.2 Prepare added time ids & embeddings - if isinstance(image, list): - original_size = original_size or image[0].shape[-2:] - else: - original_size = original_size or image.shape[-2:] - target_size = target_size or (height, width) - - 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) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - # 8.1 Apply denoising_end - if ( - self.denoising_end is not None - and isinstance(self.denoising_end, float) - and self.denoising_end > 0 - and self.denoising_end < 1 - ): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - is_unet_compiled = is_compiled_module(self.unet) - is_controlnet_compiled = is_compiled_module(self.controlnet) - is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1") - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # Relevant thread: - # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428 - if ( - torch.cuda.is_available() - and (is_unet_compiled and is_controlnet_compiled) - and is_torch_higher_equal_2_1 - ): - torch._inductor.cudagraph_mark_step_begin() - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - - # controlnet(s) inference - if guess_mode and self.do_classifier_free_guidance: - # Infer ControlNet only for the conditional batch. - control_model_input = latents - control_model_input = self.scheduler.scale_model_input(control_model_input, t) - controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] - controlnet_added_cond_kwargs = { - "text_embeds": add_text_embeds.chunk(2)[1], - "time_ids": add_time_ids.chunk(2)[1], - } - else: - control_model_input = latent_model_input - controlnet_prompt_embeds = prompt_embeds - controlnet_added_cond_kwargs = added_cond_kwargs - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=image, - conditioning_scale=cond_scale, - guess_mode=guess_mode, - added_cond_kwargs=controlnet_added_cond_kwargs, - return_dict=False, - ) - - if guess_mode and self.do_classifier_free_guidance: - # Inferred ControlNet only for the conditional batch. - # To apply the output of ControlNet to both the unconditional and conditional batches, - # add 0 to the unconditional batch to keep it unchanged. - down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] - mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = image_embeds - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds - ) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids) - image = callback_outputs.pop("image", image) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - 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) - - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None - has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None - if has_latents_mean and has_latents_std: - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean - else: - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, 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": - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py deleted file mode 100644 index 526e1ffcb2cc30c29d49da7a516ae3b0ea271bcc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py +++ /dev/null @@ -1,1681 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from diffusers.utils.import_utils import is_invisible_watermark_available - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ControlNetModel, ImageProjection, MultiControlNetModel, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import empty_device_cache, is_compiled_module, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput - - -if is_invisible_watermark_available(): - from ..stable_diffusion_xl.watermark import StableDiffusionXLWatermarker - - -from ...utils import is_torch_xla_available - - -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 - >>> # pip install accelerate transformers safetensors diffusers - - >>> import torch - >>> import numpy as np - >>> from PIL import Image - - >>> from transformers import DPTImageProcessor, DPTForDepthEstimation - >>> 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( - ... "diffusers/controlnet-depth-sdxl-1.0-small", - ... variant="fp16", - ... use_safetensors=True, - ... torch_dtype=torch.float16, - ... ) - >>> vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) - >>> pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", - ... controlnet=controlnet, - ... vae=vae, - ... variant="fp16", - ... use_safetensors=True, - ... torch_dtype=torch.float16, - ... ) - >>> pipe.enable_model_cpu_offload() - - - >>> def get_depth_map(image): - ... image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda") - ... with torch.no_grad(), torch.autocast("cuda"): - ... depth_map = depth_estimator(image).predicted_depth - - ... depth_map = torch.nn.functional.interpolate( - ... depth_map.unsqueeze(1), - ... size=(1024, 1024), - ... mode="bicubic", - ... align_corners=False, - ... ) - ... depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True) - ... depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True) - ... depth_map = (depth_map - depth_min) / (depth_max - depth_min) - ... image = torch.cat([depth_map] * 3, dim=1) - ... image = image.permute(0, 2, 3, 1).cpu().numpy()[0] - ... image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8)) - ... return image - - - >>> prompt = "A robot, 4k photo" - >>> image = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/cat.png" - ... ).resize((1024, 1024)) - >>> controlnet_conditioning_scale = 0.5 # recommended for good generalization - >>> depth_image = get_depth_map(image) - - >>> images = pipe( - ... prompt, - ... image=image, - ... control_image=depth_image, - ... strength=0.99, - ... num_inference_steps=50, - ... controlnet_conditioning_scale=controlnet_conditioning_scale, - ... ).images - >>> images[0].save(f"robot_cat.png") - ``` -""" - - -# 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" -): - 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 StableDiffusionXLControlNetImg2ImgPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionXLLoraLoaderMixin, - FromSingleFileMixin, - IPAdapterMixin, -): - r""" - Pipeline for image-to-image generation using Stable Diffusion XL with ControlNet guidance. - - 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.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 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. - 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. - 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`]. - requires_aesthetics_score (`bool`, *optional*, defaults to `"False"`): - Whether the `unet` requires an `aesthetic_score` condition to be passed during inference. Also see the - config of `stabilityai/stable-diffusion-xl-refiner-1-0`. - 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`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to - watermark output images. If not defined, it will default to True if the package is installed, otherwise no - watermarker will be used. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "feature_extractor", - "image_encoder", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "add_neg_time_ids", - "control_image", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - controlnet: Union[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: CLIPImageProcessor = None, - image_encoder: CLIPVisionModelWithProjection = None, - ): - super().__init__() - - if isinstance(controlnet, (list, tuple)): - controlnet = MultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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, do_convert_rgb=True) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - self.register_to_config(requires_aesthetics_score=requires_aesthetics_score) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - prompt_2, - image, - strength, - num_inference_steps, - callback_steps, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - 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, - ): - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - if num_inference_steps is None: - raise ValueError("`num_inference_steps` cannot be None.") - elif not isinstance(num_inference_steps, int) or num_inference_steps <= 0: - raise ValueError( - f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type" - f" {type(num_inference_steps)}." - ) - - 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_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`." - ) - - # `prompt` needs more sophisticated handling when there are multiple - # conditionings. - if isinstance(self.controlnet, MultiControlNetModel): - if isinstance(prompt, list): - 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` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - self.check_image(image, prompt, prompt_embeds) - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if not isinstance(image, list): - raise TypeError("For multiple controlnets: `image` must be type `list`") - - # When `image` is a nested list: - # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]]) - elif any(isinstance(i, list) for i in image): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - 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) - else: - assert False - - # Check `controlnet_conditioning_scale` - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if isinstance(controlnet_conditioning_scale, list): - if any(isinstance(i, list) for i in controlnet_conditioning_scale): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - elif 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" - ) - else: - assert False - - if not isinstance(control_guidance_start, (tuple, list)): - control_guidance_start = [control_guidance_start] - - if not isinstance(control_guidance_end, (tuple, list)): - control_guidance_end = [control_guidance_end] - - 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(self.controlnet, MultiControlNetModel): - 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.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}" - ) - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.prepare_image - def prepare_control_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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 - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.prepare_latents - def prepare_latents( - self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None, add_noise=True - ): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - latents_mean = latents_std = None - if hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None: - latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1) - if hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None: - latents_std = torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1) - - # Offload text encoder if `enable_model_cpu_offload` was enabled - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.text_encoder_2.to("cpu") - empty_device_cache() - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - # make sure the VAE is in float32 mode, as it overflows in float16 - if self.vae.config.force_upcast: - image = image.float() - self.vae.to(dtype=torch.float32) - - 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." - ) - - elif isinstance(generator, list): - if image.shape[0] < batch_size and batch_size % image.shape[0] == 0: - image = torch.cat([image] * (batch_size // image.shape[0]), dim=0) - elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} " - ) - - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - if self.vae.config.force_upcast: - self.vae.to(dtype) - - init_latents = init_latents.to(dtype) - if latents_mean is not None and latents_std is not None: - latents_mean = latents_mean.to(device=device, dtype=dtype) - latents_std = latents_std.to(device=device, dtype=dtype) - init_latents = (init_latents - latents_mean) * self.vae.config.scaling_factor / latents_std - else: - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - if add_noise: - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - - latents = init_latents - - return latents - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline._get_add_time_ids - def _get_add_time_ids( - self, - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype, - text_encoder_projection_dim=None, - ): - if self.config.requires_aesthetics_score: - add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,)) - add_neg_time_ids = list( - negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,) - ) - else: - add_time_ids = list(original_size + crops_coords_top_left + target_size) - add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_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 - and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model." - ) - elif ( - expected_add_embed_dim < passed_add_embed_dim - and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model." - ) - elif 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) - add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype) - - return add_time_ids, add_neg_time_ids - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_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, - prompt_2: Optional[Union[str, List[str]]] = None, - image: PipelineImageInput = None, - control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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, - 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"], - **kwargs, - ): - 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 - 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]]`): - 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 - and/or width are passed, `image` is resized according to them. 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. - height (`int`, *optional*, defaults to the size of control_image): - The height in pixels of the generated image. 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 the size of control_image): - The width in pixels of the generated image. 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. - strength (`float`, *optional*, defaults to 0.8): - 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`. - 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`). - 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. - 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*): - 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 ge 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. - 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)`. 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. - 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. - cross_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_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): - The percentage of total steps at which the controlnet starts applying. - 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)): - 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. - aesthetic_score (`float`, *optional*, defaults to 6.0): - Used to simulate an aesthetic score of the generated image by influencing the positive text condition. - 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_aesthetic_score (`float`, *optional*, defaults to 2.5): - 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). Can be used to - simulate an aesthetic score of the generated image by influencing the negative text condition. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple` - containing the output images. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - control_image, - strength, - num_inference_steps, - callback_steps, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - 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 - - if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): - controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) - - global_pool_conditions = ( - controlnet.config.global_pool_conditions - if isinstance(controlnet, ControlNetModel) - else controlnet.nets[0].config.global_pool_conditions - ) - guess_mode = guess_mode or global_pool_conditions - - # 3.1. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt, - prompt_2, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - 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=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - - # 3.2 Encode ip_adapter_image - 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, - self.do_classifier_free_guidance, - ) - - # 4. Prepare image and controlnet_conditioning_image - image = self.image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - - if isinstance(controlnet, ControlNetModel): - control_image = self.prepare_control_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - height, width = control_image.shape[-2:] - elif isinstance(controlnet, MultiControlNetModel): - control_images = [] - - for control_image_ in control_image: - control_image_ = self.prepare_control_image( - image=control_image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - - control_images.append(control_image_) - - control_image = control_images - height, width = control_image[0].shape[-2:] - else: - assert False - - # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - self._num_timesteps = len(timesteps) - - # 6. Prepare latent variables - if latents is None: - latents = self.prepare_latents( - image, - latent_timestep, - batch_size, - num_images_per_prompt, - prompt_embeds.dtype, - device, - generator, - True, - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) - - # 7.2 Prepare added time ids & embeddings - if isinstance(control_image, list): - original_size = original_size or control_image[0].shape[-2:] - else: - original_size = original_size or control_image.shape[-2:] - target_size = target_size or (height, width) - - if negative_original_size is None: - negative_original_size = original_size - if negative_target_size is None: - negative_target_size = target_size - 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, add_neg_time_ids = self._get_add_time_ids( - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype=prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1) - - 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_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1) - add_time_ids = torch.cat([add_neg_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) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - - # controlnet(s) inference - if guess_mode and self.do_classifier_free_guidance: - # Infer ControlNet only for the conditional batch. - control_model_input = latents - control_model_input = self.scheduler.scale_model_input(control_model_input, t) - controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] - controlnet_added_cond_kwargs = { - "text_embeds": add_text_embeds.chunk(2)[1], - "time_ids": add_time_ids.chunk(2)[1], - } - else: - control_model_input = latent_model_input - controlnet_prompt_embeds = prompt_embeds - controlnet_added_cond_kwargs = added_cond_kwargs - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=control_image, - conditioning_scale=cond_scale, - guess_mode=guess_mode, - added_cond_kwargs=controlnet_added_cond_kwargs, - return_dict=False, - ) - - if guess_mode and self.do_classifier_free_guidance: - # Inferred ControlNet only for the conditional batch. - # To apply the output of ControlNet to both the unconditional and conditional batches, - # add 0 to the unconditional batch to keep it unchanged. - down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] - mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = image_embeds - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=self.cross_attention_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds - ) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - add_neg_time_ids = callback_outputs.pop("add_neg_time_ids", add_neg_time_ids) - control_image = callback_outputs.pop("control_image", control_image) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # If we do sequential model offloading, let's offload unet and controlnet - # manually for max memory savings - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.unet.to("cpu") - self.controlnet.to("cpu") - empty_device_cache() - - 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) - - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None - has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None - if has_latents_mean and has_latents_std: - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean - else: - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, return_dict=False)[0] - - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - else: - image = latents - return StableDiffusionXLPipelineOutput(images=image) - - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py deleted file mode 100644 index fd490e1d5d1edd5e916df6e5f0a03d85643716a2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py +++ /dev/null @@ -1,1791 +0,0 @@ -# Copyright 2025 Harutatsu Akiyama, Jinbin Bai, 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, Tuple, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ControlNetModel, ControlNetUnionModel, ImageProjection, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_invisible_watermark_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import empty_device_cache, is_compiled_module, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput - - -if is_invisible_watermark_available(): - from diffusers.pipelines.stable_diffusion_xl.watermark import StableDiffusionXLWatermarker - - -from ...utils import is_torch_xla_available - - -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 - - -# 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" -): - 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") - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - from diffusers import StableDiffusionXLControlNetUnionInpaintPipeline, ControlNetUnionModel, AutoencoderKL - from diffusers.utils import load_image - import torch - import numpy as np - from PIL import Image - - prompt = "A cat" - # download an image - image = load_image( - "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/in_paint/overture-creations-5sI6fQgYIuo.png" - ).resize((1024, 1024)) - mask = load_image( - "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/in_paint/overture-creations-5sI6fQgYIuo_mask.png" - ).resize((1024, 1024)) - # initialize the models and pipeline - controlnet = ControlNetUnionModel.from_pretrained( - "brad-twinkl/controlnet-union-sdxl-1.0-promax", torch_dtype=torch.float16 - ) - vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) - pipe = StableDiffusionXLControlNetUnionInpaintPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", - controlnet=controlnet, - vae=vae, - torch_dtype=torch.float16, - variant="fp16", - ) - pipe.enable_model_cpu_offload() - controlnet_img = image.copy() - controlnet_img_np = np.array(controlnet_img) - mask_np = np.array(mask) - controlnet_img_np[mask_np > 0] = 0 - controlnet_img = Image.fromarray(controlnet_img_np) - # generate image - image = pipe(prompt, image=image, mask_image=mask, control_image=[controlnet_img], control_mode=[7]).images[0] - image.save("inpaint.png") - ``` -""" - - -# 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 StableDiffusionXLControlNetUnionInpaintPipeline( - DiffusionPipeline, - StableDiffusionMixin, - StableDiffusionXLLoraLoaderMixin, - FromSingleFileMixin, - IPAdapterMixin, - TextualInversionLoaderMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion XL. - - 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.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] for loading LoRA weights - - [`~loaders.StableDiffusionXLLoraLoaderMixin.save_lora_weights`] for saving LoRA weights - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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`]. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae" - - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "image_encoder", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "add_text_embeds", - "add_time_ids", - "mask", - "masked_image_latents", - "control_image", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - controlnet: ControlNetUnionModel, - 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, - ): - super().__init__() - - if not isinstance(controlnet, ControlNetUnionModel): - raise ValueError("Expected `controlnet` to be of type `ControlNetUnionModel`.") - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - self.register_to_config(requires_aesthetics_score=requires_aesthetics_score) - 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.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True - ) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # 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, - prompt, - prompt_2, - image, - mask_image, - strength, - num_inference_steps, - callback_steps, - output_type, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - controlnet_conditioning_scale=1.0, - control_guidance_start=0.0, - control_guidance_end=1.0, - callback_on_step_end_tensor_inputs=None, - padding_mask_crop=None, - ): - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - if num_inference_steps is None: - raise ValueError("`num_inference_steps` cannot be None.") - elif not isinstance(num_inference_steps, int) or num_inference_steps <= 0: - raise ValueError( - f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type" - f" {type(num_inference_steps)}." - ) - - 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_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 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 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`." - ) - - # Check `image` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - self.check_image(image, prompt, prompt_embeds) - elif ( - isinstance(self.controlnet, ControlNetUnionModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetUnionModel) - ): - self.check_image(image, prompt, prompt_embeds) - - else: - assert False - - if not isinstance(control_guidance_start, (tuple, list)): - control_guidance_start = [control_guidance_start] - - if not isinstance(control_guidance_end, (tuple, list)): - control_guidance_end = [control_guidance_end] - - 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." - ) - - 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.controlnet.pipeline_controlnet_inpaint_sd_xl.StableDiffusionXLControlNetInpaintPipeline.prepare_control_image - def prepare_control_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - crops_coords, - resize_mode, - do_classifier_free_guidance=False, - guess_mode=False, - ): - image = self.control_image_processor.preprocess( - image, height=height, width=width, crops_coords=crops_coords, resize_mode=resize_mode - ).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet_inpaint_sd_xl.StableDiffusionXLControlNetInpaintPipeline.prepare_latents - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=None, - image=None, - timestep=None, - is_strength_max=True, - add_noise=True, - return_noise=False, - return_image_latents=False, - ): - 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 (image is None or timestep is None) and not is_strength_max: - raise ValueError( - "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise." - "However, either the image or the noise timestep has not been provided." - ) - - if return_image_latents or (latents is None and not is_strength_max): - image = image.to(device=device, dtype=dtype) - - if image.shape[1] == 4: - image_latents = image - else: - image_latents = self._encode_vae_image(image=image, generator=generator) - image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1) - - if latents is None and add_noise: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # if strength is 1. then initialise the latents to noise, else initial to image + noise - latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep) - # if pure noise then scale the initial latents by the Scheduler's init sigma - latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents - elif add_noise: - noise = latents.to(device) - latents = noise * self.scheduler.init_noise_sigma - else: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = image_latents.to(device) - - outputs = (latents,) - - if return_noise: - outputs += (noise,) - - if return_image_latents: - outputs += (image_latents,) - - return outputs - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet_inpaint_sd_xl.StableDiffusionXLControlNetInpaintPipeline._encode_vae_image - def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): - dtype = image.dtype - if self.vae.config.force_upcast: - image = image.float() - self.vae.to(dtype=torch.float32) - - 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) - - if self.vae.config.force_upcast: - self.vae.to(dtype) - - image_latents = image_latents.to(dtype) - image_latents = self.vae.config.scaling_factor * image_latents - - return image_latents - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet_inpaint_sd_xl.StableDiffusionXLControlNetInpaintPipeline.prepare_mask_latents - def prepare_mask_latents( - self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance - ): - # 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 // self.vae_scale_factor, width // self.vae_scale_factor) - ) - mask = mask.to(device=device, dtype=dtype) - - # 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) - - mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask - - masked_image_latents = None - if masked_image is not None: - masked_image = masked_image.to(device=device, dtype=dtype) - masked_image_latents = self._encode_vae_image(masked_image, generator=generator) - 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 - ) - - masked_image_latents = ( - torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents - ) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device, denoising_start=None): - # get the original timestep using init_timestep - if denoising_start is None: - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - t_start = 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 - - else: - # Strength is irrelevant if we directly request a timestep to start at; - # that is, strength is determined by the denoising_start instead. - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (denoising_start * self.scheduler.config.num_train_timesteps) - ) - ) - - num_inference_steps = (self.scheduler.timesteps < discrete_timestep_cutoff).sum().item() - if self.scheduler.order == 2 and num_inference_steps % 2 == 0: - # if the scheduler is a 2nd order scheduler we might have to do +1 - # because `num_inference_steps` might be even given that every timestep - # (except the highest one) is duplicated. If `num_inference_steps` is even it would - # mean that we cut the timesteps in the middle of the denoising step - # (between 1st and 2nd derivative) which leads to incorrect results. By adding 1 - # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler - num_inference_steps = num_inference_steps + 1 - - # because t_n+1 >= t_n, we slice the timesteps starting from the end - t_start = len(self.scheduler.timesteps) - num_inference_steps - timesteps = self.scheduler.timesteps[t_start:] - if hasattr(self.scheduler, "set_begin_index"): - self.scheduler.set_begin_index(t_start) - return timesteps, num_inference_steps - - def _get_add_time_ids( - self, - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - dtype, - text_encoder_projection_dim=None, - ): - if self.config.requires_aesthetics_score: - add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,)) - add_neg_time_ids = list(original_size + crops_coords_top_left + (negative_aesthetic_score,)) - else: - add_time_ids = list(original_size + crops_coords_top_left + target_size) - add_neg_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 - and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model." - ) - elif ( - expected_add_embed_dim < passed_add_embed_dim - and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model." - ) - elif 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) - add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype) - - return add_time_ids, add_neg_time_ids - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_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, - prompt_2: Optional[Union[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, - strength: float = 0.9999, - num_inference_steps: int = 50, - denoising_start: Optional[float] = None, - denoising_end: Optional[float] = 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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]]] = 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, - 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"], - **kwargs, - ): - 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 - image (`PIL.Image.Image`): - `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will - be masked out with `mask_image` and repainted according to `prompt`. - mask_image (`PIL.Image.Image`): - `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be - 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)`. - 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. - padding_mask_crop (`int`, *optional*, defaults to `None`): - The size of margin in the crop to be applied to the image and masking. If `None`, no crop is applied to - image and mask_image. If `padding_mask_crop` is not `None`, it will first find a rectangular region - with the same aspect ration of the image and contains all masked area, and then expand that area based - on `padding_mask_crop`. The image and mask_image will then be cropped based on the expanded area before - resizing to the original image size for inpainting. This is useful when the masked area is small while - the image is large and contain information irrelevant for inpainting, such as background. - strength (`float`, *optional*, defaults to 0.9999): - Conceptually, indicates how much to transform the masked portion of 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 the masked - portion of the reference `image`. Note that in the case of `denoising_start` being declared as an - integer, the value of `strength` will be ignored. - 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. - denoising_start (`float`, *optional*): - When specified, indicates the fraction (between 0.0 and 1.0) of the total denoising process to be - bypassed before it is initiated. Consequently, the initial part of the denoising process is skipped and - it is assumed that the passed `image` is a partly denoised image. Note that when this is specified, - strength will be ignored. The `denoising_start` parameter is particularly beneficial when this pipeline - is integrated into a "Mixture of Denoisers" multi-pipeline setup, as detailed in [**Refining the Image - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output). - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - still retain a substantial amount of noise (ca. final 20% of timesteps still needed) and should be - denoised by a successor pipeline that has `denoising_start` set to 0.8 so that it only denoises the - final 20% of 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output). - 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`). - 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. - 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)`. 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. - 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. - 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 ge generated by sampling using the supplied random `generator`. - 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. - cross_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). - 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` 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 `(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). - aesthetic_score (`float`, *optional*, defaults to 6.0): - Used to simulate an aesthetic score of the generated image by influencing the positive text condition. - 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_aesthetic_score (`float`, *optional*, defaults to 2.5): - 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). Can be used to - simulate an aesthetic score of the generated image by influencing the negative text condition. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a - `tuple. `tuple. When returning a tuple, the first element is a list with the generated images. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - - # # 0.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 - - # 0.1 align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - - if not isinstance(control_image, list): - control_image = [control_image] - else: - control_image = control_image.copy() - - if not isinstance(control_mode, list): - control_mode = [control_mode] - - if len(control_image) != len(control_mode): - raise ValueError("Expected len(control_image) == len(control_type)") - - num_control_type = controlnet.config.num_control_type - - # 1. Check inputs - control_type = [0 for _ in range(num_control_type)] - for _image, control_idx in zip(control_image, control_mode): - control_type[control_idx] = 1 - self.check_inputs( - prompt, - prompt_2, - _image, - mask_image, - strength, - num_inference_steps, - callback_steps, - output_type, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - callback_on_step_end_tensor_inputs, - padding_mask_crop, - ) - - control_type = torch.Tensor(control_type) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - 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. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else 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=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - - # 3.1 Encode ip_adapter_image - 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, - self.do_classifier_free_guidance, - ) - - # 4. set timesteps - def denoising_value_valid(dnv): - return isinstance(dnv, float) and 0 < dnv < 1 - - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps( - num_inference_steps, - strength, - device, - denoising_start=denoising_start if denoising_value_valid(denoising_start) else None, - ) - # check that number of inference steps is not < 1 - as this doesn't make sense - 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." - ) - # at which timestep to set the initial noise (n.b. 50% if strength is 0.5) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise - is_strength_max = strength == 1.0 - self._num_timesteps = len(timesteps) - - # 5. Preprocess mask and image - resizes image and mask w.r.t height and width - # 5.1 Prepare init image - if padding_mask_crop is not None: - height, width = self.image_processor.get_default_height_width(image, height, width) - 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) - - # 5.2 Prepare control images - for idx, _ in enumerate(control_image): - control_image[idx] = self.prepare_control_image( - image=control_image[idx], - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - crops_coords=crops_coords, - resize_mode=resize_mode, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - height, width = control_image[idx].shape[-2:] - - # 5.3 Prepare mask - mask = self.mask_processor.preprocess( - mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords - ) - - masked_image = init_image * (mask < 0.5) - _, _, height, width = init_image.shape - - # 6. Prepare latent variables - num_channels_latents = self.vae.config.latent_channels - num_channels_unet = self.unet.config.in_channels - return_image_latents = num_channels_unet == 4 - - add_noise = True if denoising_start is None else False - latents_outputs = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - image=init_image, - timestep=latent_timestep, - is_strength_max=is_strength_max, - add_noise=add_noise, - return_noise=True, - return_image_latents=return_image_latents, - ) - - if return_image_latents: - latents, noise, image_latents = latents_outputs - else: - latents, noise = latents_outputs - - # 7. Prepare mask latent variables - mask, _ = self.prepare_mask_latents( - mask, - masked_image, - batch_size * num_images_per_prompt, - height, - width, - prompt_embeds.dtype, - device, - generator, - self.do_classifier_free_guidance, - ) - - # 8. Check that sizes of mask, masked image and latents match - if num_channels_unet != 4: - raise ValueError( - f"The unet {self.unet.__class__} should have either 4 or 9 input channels, not {self.unet.config.in_channels}." - ) - # 8.1 Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 8.2 Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - controlnet_keep.append( - 1.0 - - float(i / len(timesteps) < control_guidance_start or (i + 1) / len(timesteps) > control_guidance_end) - ) - - # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - height, width = latents.shape[-2:] - height = height * self.vae_scale_factor - width = width * self.vae_scale_factor - - original_size = original_size or (height, width) - target_size = target_size or (height, width) - for _image in control_image: - if isinstance(_image, torch.Tensor): - original_size = original_size or _image.shape[-2:] - - # 10. 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, add_neg_time_ids = self._get_add_time_ids( - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - dtype=prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1) - - 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_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1) - add_time_ids = torch.cat([add_neg_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) - - # 11. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - if ( - denoising_end is not None - and denoising_start is not None - and denoising_value_valid(denoising_end) - and denoising_value_valid(denoising_start) - and denoising_start >= denoising_end - ): - raise ValueError( - f"`denoising_start`: {denoising_start} cannot be larger than or equal to `denoising_end`: " - + f" {denoising_end} when using type float." - ) - elif denoising_end is not None and denoising_value_valid(denoising_end): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - control_type = ( - control_type.reshape(1, -1) - .to(device, dtype=prompt_embeds.dtype) - .repeat(batch_size * num_images_per_prompt * 2, 1) - ) - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - added_cond_kwargs = { - "text_embeds": add_text_embeds, - "time_ids": add_time_ids, - } - - # controlnet(s) inference - if guess_mode and self.do_classifier_free_guidance: - # Infer ControlNet only for the conditional batch. - control_model_input = latents - control_model_input = self.scheduler.scale_model_input(control_model_input, t) - controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] - controlnet_added_cond_kwargs = { - "text_embeds": add_text_embeds.chunk(2)[1], - "time_ids": add_time_ids.chunk(2)[1], - } - else: - control_model_input = latent_model_input - controlnet_prompt_embeds = prompt_embeds - controlnet_added_cond_kwargs = added_cond_kwargs - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - # # Resize control_image to match the size of the input to the controlnet - # if control_image.shape[-2:] != control_model_input.shape[-2:]: - # control_image = F.interpolate(control_image, size=control_model_input.shape[-2:], mode="bilinear", align_corners=False) - - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=control_image, - control_type=control_type, - control_type_idx=control_mode, - conditioning_scale=cond_scale, - guess_mode=guess_mode, - added_cond_kwargs=controlnet_added_cond_kwargs, - return_dict=False, - ) - - if guess_mode and self.do_classifier_free_guidance: - # Inferred ControlNet only for the conditional batch. - # To apply the output of ControlNet to both the unconditional and conditional batches, - # add 0 to the unconditional batch to keep it unchanged. - down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] - mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = image_embeds - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=self.cross_attention_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - if self.do_classifier_free_guidance and guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - init_latents_proper = image_latents - if self.do_classifier_free_guidance: - init_mask, _ = mask.chunk(2) - else: - init_mask = mask - - if i < len(timesteps) - 1: - noise_timestep = timesteps[i + 1] - init_latents_proper = self.scheduler.add_noise( - init_latents_proper, noise, torch.tensor([noise_timestep]) - ) - - latents = (1 - init_mask) * init_latents_proper + init_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) - control_image = callback_outputs.pop("control_image", control_image) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # make sure the VAE is in float32 mode, as it overflows in float16 - if self.vae.dtype == torch.float16 and self.vae.config.force_upcast: - self.upcast_vae() - latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype) - - # If we do sequential model offloading, let's offload unet and controlnet - # manually for max memory savings - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.unet.to("cpu") - self.controlnet.to("cpu") - empty_device_cache() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - else: - return StableDiffusionXLPipelineOutput(images=latents) - - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - image = self.image_processor.postprocess(image, output_type=output_type) - - if padding_mask_crop is not None: - image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return StableDiffusionXLPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py deleted file mode 100644 index 60768e43faf9bf64e89542bb50490e6f3987a0fb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py +++ /dev/null @@ -1,1616 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from diffusers.utils.import_utils import is_invisible_watermark_available - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import ( - AutoencoderKL, - ControlNetUnionModel, - ImageProjection, - MultiControlNetUnionModel, - UNet2DConditionModel, -) -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import is_compiled_module, is_torch_version, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput - - -if is_invisible_watermark_available(): - from ..stable_diffusion_xl.watermark import StableDiffusionXLWatermarker - - -from ...utils import is_torch_xla_available - - -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 - >>> # !pip install controlnet_aux - >>> from controlnet_aux import LineartAnimeDetector - >>> from diffusers import StableDiffusionXLControlNetUnionPipeline, ControlNetUnionModel, AutoencoderKL - >>> from diffusers.utils import load_image - >>> import torch - - >>> prompt = "A cat" - >>> # download an image - >>> image = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/kandinsky/cat.png" - ... ).resize((1024, 1024)) - >>> # initialize the models and pipeline - >>> controlnet = ControlNetUnionModel.from_pretrained( - ... "xinsir/controlnet-union-sdxl-1.0", torch_dtype=torch.float16 - ... ) - >>> vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) - >>> pipe = StableDiffusionXLControlNetUnionPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", - ... controlnet=controlnet, - ... vae=vae, - ... torch_dtype=torch.float16, - ... variant="fp16", - ... ) - >>> pipe.enable_model_cpu_offload() - >>> # prepare image - >>> processor = LineartAnimeDetector.from_pretrained("lllyasviel/Annotators") - >>> controlnet_img = processor(image, output_type="pil") - >>> # generate image - >>> image = pipe(prompt, control_image=[controlnet_img], control_mode=[3], height=1024, width=1024).images[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, - **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 StableDiffusionXLControlNetUnionPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionXLLoraLoaderMixin, - IPAdapterMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion XL with ControlNet guidance. - - 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.). - - The pipeline also inherits the following loading methods: - - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings - - [`~loaders.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] for loading LoRA weights - - [`~loaders.StableDiffusionXLLoraLoaderMixin.save_lora_weights`] for saving LoRA weights - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - text_encoder_2 ([`~transformers.CLIPTextModelWithProjection`]): - Second frozen text-encoder - ([laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - tokenizer_2 ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - controlnet ([`ControlNetUnionModel`]`): - Provides additional conditioning to the `unet` during the denoising process. - 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 should always be set to 0. Also see the config of - `stabilityai/stable-diffusion-xl-base-1-0`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark](https://github.com/ShieldMnt/invisible-watermark/) library to - watermark output images. If not defined, it defaults to `True` if the package is installed; otherwise no - watermarker is used. - """ - - # leave controlnet out on purpose because it iterates with unet - model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "feature_extractor", - "image_encoder", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "add_text_embeds", - "add_time_ids", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - controlnet: Union[ - ControlNetUnionModel, List[ControlNetUnionModel], Tuple[ControlNetUnionModel], MultiControlNetUnionModel - ], - scheduler: KarrasDiffusionSchedulers, - force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - feature_extractor: CLIPImageProcessor = None, - image_encoder: CLIPVisionModelWithProjection = None, - ): - super().__init__() - - if isinstance(controlnet, (list, tuple)): - controlnet = MultiControlNetUnionModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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, do_convert_rgb=True) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # 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, - prompt, - prompt_2, - image: PipelineImageInput, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - negative_pooled_prompt_embeds=None, - controlnet_conditioning_scale=1.0, - control_guidance_start=0.0, - control_guidance_end=1.0, - control_mode=None, - callback_on_step_end_tensor_inputs=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_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`." - ) - - # `prompt` needs more sophisticated handling when there are multiple - # conditionings. - if isinstance(self.controlnet, MultiControlNetUnionModel): - if isinstance(prompt, list): - 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, ControlNetUnionModel): - for image_ in image: - self.check_image(image_, prompt, prompt_embeds) - elif isinstance(controlnet, MultiControlNetUnionModel): - if not isinstance(image, list): - raise TypeError("For multiple controlnets: `image` must be type `list`") - elif not all(isinstance(i, list) for i in image): - raise ValueError("For multiple controlnets: elements of `image` must be list of conditionings.") - 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 images_ in image: - for image_ in images_: - self.check_image(image_, prompt, prompt_embeds) - - # Check `controlnet_conditioning_scale` - if isinstance(controlnet, MultiControlNetUnionModel): - if isinstance(controlnet_conditioning_scale, list): - if any(isinstance(i, list) for i in controlnet_conditioning_scale): - raise ValueError("A single batch of multiple conditionings is not supported at the moment.") - elif 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, MultiControlNetUnionModel): - 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.") - - # Check `control_mode` - if isinstance(controlnet, ControlNetUnionModel): - if max(control_mode) >= controlnet.config.num_control_type: - raise ValueError(f"control_mode: must be lower than {controlnet.config.num_control_type}.") - elif isinstance(controlnet, MultiControlNetUnionModel): - for _control_mode, _controlnet in zip(control_mode, self.controlnet.nets): - if max(_control_mode) >= _controlnet.config.num_control_type: - raise ValueError(f"control_mode: must be lower than {_controlnet.config.num_control_type}.") - - # Equal number of `image` and `control_mode` elements - if isinstance(controlnet, ControlNetUnionModel): - if len(image) != len(control_mode): - raise ValueError("Expected len(control_image) == len(control_mode)") - elif isinstance(controlnet, MultiControlNetUnionModel): - if not all(isinstance(i, list) for i in control_mode): - raise ValueError( - "For multiple controlnets: elements of control_mode must be lists representing conditioning mode." - ) - - elif sum(len(x) for x in image) != sum(len(x) for x in control_mode): - raise ValueError("Expected len(control_image) == len(control_mode)") - - 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.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.prepare_image - def prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - 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.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_end(self): - return self._denoising_end - - @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, - prompt_2: Optional[Union[str, List[str]]] = None, - control_image: Union[PipelineImageInput, List[PipelineImageInput]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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"], - ): - r""" - The call function to the pipeline for generation. - - Args: - 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*): - 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*): - 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. - height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The height in pixels of the generated image. 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. 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. - 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*): - 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. - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output) - 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*): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - pooled_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated pooled text embeddings. Can be used to easily tweak text inputs (prompt weighting). If - not provided, pooled text embeddings are 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 (prompt - 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*): - 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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): - 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): - The percentage of total steps at which the ControlNet starts applying. - 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*): - 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)): - 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. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned, - otherwise a `tuple` is returned containing the output images. - """ - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - if not isinstance(control_image, list): - control_image = [control_image] - else: - control_image = control_image.copy() - - if not isinstance(control_mode, list): - control_mode = [control_mode] - - if isinstance(controlnet, MultiControlNetUnionModel): - control_image = [[item] for item in control_image] - control_mode = [[item] for item in control_mode] - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetUnionModel) else len(control_mode) - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - if isinstance(controlnet_conditioning_scale, float): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetUnionModel) else len(control_mode) - controlnet_conditioning_scale = [controlnet_conditioning_scale] * mult - - # 1. Check inputs - self.check_inputs( - prompt, - prompt_2, - control_image, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - negative_pooled_prompt_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - control_mode, - callback_on_step_end_tensor_inputs, - ) - - if isinstance(controlnet, ControlNetUnionModel): - control_type = torch.zeros(controlnet.config.num_control_type).scatter_(0, torch.tensor(control_mode), 1) - elif isinstance(controlnet, MultiControlNetUnionModel): - control_type = [ - torch.zeros(controlnet_.config.num_control_type).scatter_(0, torch.tensor(control_mode_), 1) - for control_mode_, controlnet_ in zip(control_mode, self.controlnet.nets) - ] - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - 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 - - global_pool_conditions = ( - controlnet.config.global_pool_conditions - if isinstance(controlnet, ControlNetUnionModel) - else controlnet.nets[0].config.global_pool_conditions - ) - guess_mode = guess_mode or global_pool_conditions - - # 3.1 Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt, - prompt_2, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - 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=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - - # 3.2 Encode ip_adapter_image - 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, - self.do_classifier_free_guidance, - ) - - # 4. Prepare image - if isinstance(controlnet, ControlNetUnionModel): - control_images = [] - - for image_ in control_image: - image_ = self.prepare_image( - image=image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - - control_images.append(image_) - - control_image = control_images - height, width = control_image[0].shape[-2:] - - elif isinstance(controlnet, MultiControlNetUnionModel): - control_images = [] - - for control_image_ in control_image: - images = [] - - for image_ in control_image_: - image_ = self.prepare_image( - image=image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - - images.append(image_) - control_images.append(images) - - control_image = control_images - height, width = control_image[0][0].shape[-2:] - - # 5. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - self._num_timesteps = len(timesteps) - - # 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, - ) - - # 6.5 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) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps) - - # 7.2 Prepare added time ids & embeddings - original_size = original_size or (height, width) - target_size = target_size or (height, width) - for _image in control_image: - if isinstance(_image, torch.Tensor): - original_size = original_size or _image.shape[-2:] - 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) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - # 8.1 Apply denoising_end - if ( - self.denoising_end is not None - and isinstance(self.denoising_end, float) - and self.denoising_end > 0 - and self.denoising_end < 1 - ): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - is_unet_compiled = is_compiled_module(self.unet) - is_controlnet_compiled = is_compiled_module(self.controlnet) - is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1") - - if isinstance(controlnet, ControlNetUnionModel): - control_type = ( - control_type.reshape(1, -1) - .to(self._execution_device, dtype=prompt_embeds.dtype) - .repeat(batch_size * num_images_per_prompt * 2, 1) - ) - if isinstance(controlnet, MultiControlNetUnionModel): - control_type = [ - _control_type.reshape(1, -1) - .to(self._execution_device, dtype=prompt_embeds.dtype) - .repeat(batch_size * num_images_per_prompt * 2, 1) - for _control_type in control_type - ] - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # Relevant thread: - # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428 - if (is_unet_compiled and is_controlnet_compiled) and is_torch_higher_equal_2_1: - torch._inductor.cudagraph_mark_step_begin() - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - added_cond_kwargs = { - "text_embeds": add_text_embeds, - "time_ids": add_time_ids, - } - - # controlnet(s) inference - if guess_mode and self.do_classifier_free_guidance: - # Infer ControlNet only for the conditional batch. - control_model_input = latents - control_model_input = self.scheduler.scale_model_input(control_model_input, t) - controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] - controlnet_added_cond_kwargs = { - "text_embeds": add_text_embeds.chunk(2)[1], - "time_ids": add_time_ids.chunk(2)[1], - } - else: - control_model_input = latent_model_input - controlnet_prompt_embeds = prompt_embeds - controlnet_added_cond_kwargs = added_cond_kwargs - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=control_image, - control_type=control_type, - control_type_idx=control_mode, - conditioning_scale=cond_scale, - guess_mode=guess_mode, - added_cond_kwargs=controlnet_added_cond_kwargs, - return_dict=False, - ) - - if guess_mode and self.do_classifier_free_guidance: - # Inferred ControlNet only for the conditional batch. - # To apply the output of ControlNet to both the unconditional and conditional batches, - # add 0 to the unconditional batch to keep it unchanged. - down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] - mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = image_embeds - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - - # 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 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) - - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None - has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None - if has_latents_mean and has_latents_std: - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean - else: - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, 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": - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py deleted file mode 100644 index 82ef4b6391eb02ccaa7d44aa1bef78ff046398ed..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py +++ /dev/null @@ -1,1623 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from diffusers.utils.import_utils import is_invisible_watermark_available - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ControlNetModel, ControlNetUnionModel, ImageProjection, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import empty_device_cache, is_compiled_module, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput - - -if is_invisible_watermark_available(): - from ..stable_diffusion_xl.watermark import StableDiffusionXLWatermarker - - -from ...utils import is_torch_xla_available - - -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 - # !pip install controlnet_aux - from diffusers import ( - StableDiffusionXLControlNetUnionImg2ImgPipeline, - ControlNetUnionModel, - AutoencoderKL, - ) - from diffusers.utils import load_image - import torch - from PIL import Image - import numpy as np - - prompt = "A cat" - # download an image - image = load_image( - "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/kandinsky/cat.png" - ) - # initialize the models and pipeline - controlnet = ControlNetUnionModel.from_pretrained( - "brad-twinkl/controlnet-union-sdxl-1.0-promax", torch_dtype=torch.float16 - ) - vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) - pipe = StableDiffusionXLControlNetUnionImg2ImgPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", - controlnet=controlnet, - vae=vae, - torch_dtype=torch.float16, - variant="fp16", - ).to("cuda") - # `enable_model_cpu_offload` is not recommended due to multiple generations - height = image.height - width = image.width - ratio = np.sqrt(1024.0 * 1024.0 / (width * height)) - # 3 * 3 upscale correspond to 16 * 3 multiply, 2 * 2 correspond to 16 * 2 multiply and so on. - scale_image_factor = 3 - base_factor = 16 - factor = scale_image_factor * base_factor - W, H = int(width * ratio) // factor * factor, int(height * ratio) // factor * factor - image = image.resize((W, H)) - target_width = W // scale_image_factor - target_height = H // scale_image_factor - images = [] - crops_coords_list = [ - (0, 0), - (0, width // 2), - (height // 2, 0), - (width // 2, height // 2), - 0, - 0, - 0, - 0, - 0, - ] - for i in range(scale_image_factor): - for j in range(scale_image_factor): - left = j * target_width - top = i * target_height - right = left + target_width - bottom = top + target_height - cropped_image = image.crop((left, top, right, bottom)) - cropped_image = cropped_image.resize((W, H)) - images.append(cropped_image) - # set ControlNetUnion input - result_images = [] - for sub_img, crops_coords in zip(images, crops_coords_list): - new_width, new_height = W, H - out = pipe( - prompt=[prompt] * 1, - image=sub_img, - control_image=[sub_img], - control_mode=[6], - width=new_width, - height=new_height, - num_inference_steps=30, - crops_coords_top_left=(W, H), - target_size=(W, H), - original_size=(W * 2, H * 2), - ) - result_images.append(out.images[0]) - new_im = Image.new("RGB", (new_width * scale_image_factor, new_height * scale_image_factor)) - new_im.paste(result_images[0], (0, 0)) - new_im.paste(result_images[1], (new_width, 0)) - new_im.paste(result_images[2], (new_width * 2, 0)) - new_im.paste(result_images[3], (0, new_height)) - new_im.paste(result_images[4], (new_width, new_height)) - new_im.paste(result_images[5], (new_width * 2, new_height)) - new_im.paste(result_images[6], (0, new_height * 2)) - new_im.paste(result_images[7], (new_width, new_height * 2)) - new_im.paste(result_images[8], (new_width * 2, new_height * 2)) - ``` -""" - - -# 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" -): - 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 StableDiffusionXLControlNetUnionImg2ImgPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionXLLoraLoaderMixin, - FromSingleFileMixin, - IPAdapterMixin, -): - r""" - Pipeline for image-to-image generation using Stable Diffusion XL with ControlNet guidance. - - 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.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 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. - controlnet ([`ControlNetUnionModel`]): - Provides additional conditioning to the unet during the denoising process. - 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`]. - requires_aesthetics_score (`bool`, *optional*, defaults to `"False"`): - Whether the `unet` requires an `aesthetic_score` condition to be passed during inference. Also see the - config of `stabilityai/stable-diffusion-xl-refiner-1-0`. - 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`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to - watermark output images. If not defined, it will default to True if the package is installed, otherwise no - watermarker will be used. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "feature_extractor", - "image_encoder", - ] - _callback_tensor_inputs = ["latents", "prompt_embeds", "add_text_embeds", "add_time_ids", "control_image"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - controlnet: ControlNetUnionModel, - scheduler: KarrasDiffusionSchedulers, - requires_aesthetics_score: bool = False, - force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - feature_extractor: CLIPImageProcessor = None, - image_encoder: CLIPVisionModelWithProjection = None, - ): - super().__init__() - - if not isinstance(controlnet, ControlNetUnionModel): - raise ValueError("Expected `controlnet` to be of type `ControlNetUnionModel`.") - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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, do_convert_rgb=True) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - self.register_to_config(requires_aesthetics_score=requires_aesthetics_score) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - prompt_2, - image, - strength, - num_inference_steps, - callback_steps, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - 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, - ): - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - if num_inference_steps is None: - raise ValueError("`num_inference_steps` cannot be None.") - elif not isinstance(num_inference_steps, int) or num_inference_steps <= 0: - raise ValueError( - f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type" - f" {type(num_inference_steps)}." - ) - - 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_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`." - ) - - # Check `image` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - self.check_image(image, prompt, prompt_embeds) - elif ( - isinstance(self.controlnet, ControlNetUnionModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetUnionModel) - ): - self.check_image(image, prompt, prompt_embeds) - else: - assert False - - if not isinstance(control_guidance_start, (tuple, list)): - control_guidance_start = [control_guidance_start] - - if not isinstance(control_guidance_end, (tuple, list)): - control_guidance_end = [control_guidance_end] - - 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." - ) - - 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.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}" - ) - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.prepare_image - def prepare_control_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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 - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.prepare_latents - def prepare_latents( - self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None, add_noise=True - ): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - latents_mean = latents_std = None - if hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None: - latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1) - if hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None: - latents_std = torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1) - - # Offload text encoder if `enable_model_cpu_offload` was enabled - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.text_encoder_2.to("cpu") - empty_device_cache() - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - # make sure the VAE is in float32 mode, as it overflows in float16 - if self.vae.config.force_upcast: - image = image.float() - self.vae.to(dtype=torch.float32) - - 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." - ) - - elif isinstance(generator, list): - if image.shape[0] < batch_size and batch_size % image.shape[0] == 0: - image = torch.cat([image] * (batch_size // image.shape[0]), dim=0) - elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} " - ) - - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - if self.vae.config.force_upcast: - self.vae.to(dtype) - - init_latents = init_latents.to(dtype) - if latents_mean is not None and latents_std is not None: - latents_mean = latents_mean.to(device=device, dtype=dtype) - latents_std = latents_std.to(device=device, dtype=dtype) - init_latents = (init_latents - latents_mean) * self.vae.config.scaling_factor / latents_std - else: - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - if add_noise: - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - - latents = init_latents - - return latents - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline._get_add_time_ids - def _get_add_time_ids( - self, - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype, - text_encoder_projection_dim=None, - ): - if self.config.requires_aesthetics_score: - add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,)) - add_neg_time_ids = list( - negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,) - ) - else: - add_time_ids = list(original_size + crops_coords_top_left + target_size) - add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_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 - and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model." - ) - elif ( - expected_add_embed_dim < passed_add_embed_dim - and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model." - ) - elif 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) - add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype) - - return add_time_ids, add_neg_time_ids - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_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, - prompt_2: Optional[Union[str, List[str]]] = None, - image: PipelineImageInput = None, - control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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]]] = 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, - 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"], - **kwargs, - ): - 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 - 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`): - 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 - and/or width are passed, `image` is resized according to them. 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. - height (`int`, *optional*, defaults to the size of control_image): - The height in pixels of the generated image. 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 the size of control_image): - The width in pixels of the generated image. 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. - strength (`float`, *optional*, defaults to 0.8): - 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`. - 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`). - 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. - 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*): - 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 ge 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. - 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)`. 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. - 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. - cross_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_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): - The percentage of total steps at which the controlnet starts applying. - 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)): - 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. - aesthetic_score (`float`, *optional*, defaults to 6.0): - Used to simulate an aesthetic score of the generated image by influencing the positive text condition. - 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_aesthetic_score (`float`, *optional*, defaults to 2.5): - 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). Can be used to - simulate an aesthetic score of the generated image by influencing the negative text condition. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple` - containing the output images. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - - if not isinstance(control_image, list): - control_image = [control_image] - else: - control_image = control_image.copy() - - if not isinstance(control_mode, list): - control_mode = [control_mode] - - if len(control_image) != len(control_mode): - raise ValueError("Expected len(control_image) == len(control_type)") - - num_control_type = controlnet.config.num_control_type - - # 1. Check inputs - control_type = [0 for _ in range(num_control_type)] - for _image, control_idx in zip(control_image, control_mode): - control_type[control_idx] = 1 - self.check_inputs( - prompt, - prompt_2, - _image, - strength, - num_inference_steps, - callback_steps, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - callback_on_step_end_tensor_inputs, - ) - - control_type = torch.Tensor(control_type) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - 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 - - global_pool_conditions = controlnet.config.global_pool_conditions - guess_mode = guess_mode or global_pool_conditions - - # 3.1. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt, - prompt_2, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - 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=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - - # 3.2 Encode ip_adapter_image - 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, - self.do_classifier_free_guidance, - ) - - # 4. Prepare image and controlnet_conditioning_image - image = self.image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - - for idx, _ in enumerate(control_image): - control_image[idx] = self.prepare_control_image( - image=control_image[idx], - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - height, width = control_image[idx].shape[-2:] - - # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - self._num_timesteps = len(timesteps) - - # 6. Prepare latent variables - if latents is None: - latents = self.prepare_latents( - image, - latent_timestep, - batch_size, - num_images_per_prompt, - prompt_embeds.dtype, - device, - generator, - True, - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - controlnet_keep.append( - 1.0 - - float(i / len(timesteps) < control_guidance_start or (i + 1) / len(timesteps) > control_guidance_end) - ) - - # 7.2 Prepare added time ids & embeddings - original_size = original_size or (height, width) - target_size = target_size or (height, width) - for _image in control_image: - if isinstance(_image, torch.Tensor): - original_size = original_size or _image.shape[-2:] - - if negative_original_size is None: - negative_original_size = original_size - if negative_target_size is None: - negative_target_size = target_size - 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, add_neg_time_ids = self._get_add_time_ids( - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype=prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1) - - 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_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1) - add_time_ids = torch.cat([add_neg_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) - control_type = ( - control_type.reshape(1, -1) - .to(device, dtype=prompt_embeds.dtype) - .repeat(batch_size * num_images_per_prompt * 2, 1) - ) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - added_cond_kwargs = { - "text_embeds": add_text_embeds, - "time_ids": add_time_ids, - } - - # controlnet(s) inference - if guess_mode and self.do_classifier_free_guidance: - # Infer ControlNet only for the conditional batch. - control_model_input = latents - control_model_input = self.scheduler.scale_model_input(control_model_input, t) - controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] - controlnet_added_cond_kwargs = { - "text_embeds": add_text_embeds.chunk(2)[1], - "time_ids": add_time_ids.chunk(2)[1], - } - else: - control_model_input = latent_model_input - controlnet_prompt_embeds = prompt_embeds - controlnet_added_cond_kwargs = added_cond_kwargs - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=control_image, - control_type=control_type, - control_type_idx=control_mode, - conditioning_scale=cond_scale, - guess_mode=guess_mode, - added_cond_kwargs=controlnet_added_cond_kwargs, - return_dict=False, - ) - - if guess_mode and self.do_classifier_free_guidance: - # Inferred ControlNet only for the conditional batch. - # To apply the output of ControlNet to both the unconditional and conditional batches, - # add 0 to the unconditional batch to keep it unchanged. - down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] - mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = image_embeds - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=self.cross_attention_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - control_image = callback_outputs.pop("control_image", control_image) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # If we do sequential model offloading, let's offload unet and controlnet - # manually for max memory savings - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.unet.to("cpu") - self.controlnet.to("cpu") - empty_device_cache() - - 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) - - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None - has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None - if has_latents_mean and has_latents_std: - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean - else: - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, return_dict=False)[0] - - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - else: - image = latents - return StableDiffusionXLPipelineOutput(images=image) - - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py deleted file mode 100644 index 1de1d4bde7f56e89317405ac96600e47f034659d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py +++ /dev/null @@ -1,535 +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 warnings -from functools import partial -from typing import Dict, List, Optional, Union - -import jax -import jax.numpy as jnp -import numpy as np -from flax.core.frozen_dict import FrozenDict -from flax.jax_utils import unreplicate -from flax.training.common_utils import shard -from PIL import Image -from transformers import CLIPImageProcessor, CLIPTokenizer, FlaxCLIPTextModel - -from ...models import FlaxAutoencoderKL, FlaxControlNetModel, FlaxUNet2DConditionModel -from ...schedulers import ( - FlaxDDIMScheduler, - FlaxDPMSolverMultistepScheduler, - FlaxLMSDiscreteScheduler, - FlaxPNDMScheduler, -) -from ...utils import PIL_INTERPOLATION, logging, replace_example_docstring -from ..pipeline_flax_utils import FlaxDiffusionPipeline -from ..stable_diffusion import FlaxStableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker_flax import FlaxStableDiffusionSafetyChecker - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -# Set to True to use python for loop instead of jax.fori_loop for easier debugging -DEBUG = False - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import jax - >>> import numpy as np - >>> import jax.numpy as jnp - >>> from flax.jax_utils import replicate - >>> from flax.training.common_utils import shard - >>> from diffusers.utils import load_image, make_image_grid - >>> from PIL import Image - >>> from diffusers import FlaxStableDiffusionControlNetPipeline, FlaxControlNetModel - - - >>> def create_key(seed=0): - ... return jax.random.PRNGKey(seed) - - - >>> rng = create_key(0) - - >>> # get canny image - >>> canny_image = load_image( - ... "https://huggingface.co/datasets/YiYiXu/test-doc-assets/resolve/main/blog_post_cell_10_output_0.jpeg" - ... ) - - >>> prompts = "best quality, extremely detailed" - >>> negative_prompts = "monochrome, lowres, bad anatomy, worst quality, low quality" - - >>> # load control net and stable diffusion v1-5 - >>> controlnet, controlnet_params = FlaxControlNetModel.from_pretrained( - ... "lllyasviel/sd-controlnet-canny", from_pt=True, dtype=jnp.float32 - ... ) - >>> pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", - ... controlnet=controlnet, - ... revision="flax", - ... dtype=jnp.float32, - ... ) - >>> params["controlnet"] = controlnet_params - - >>> num_samples = jax.device_count() - >>> rng = jax.random.split(rng, jax.device_count()) - - >>> prompt_ids = pipe.prepare_text_inputs([prompts] * num_samples) - >>> negative_prompt_ids = pipe.prepare_text_inputs([negative_prompts] * num_samples) - >>> processed_image = pipe.prepare_image_inputs([canny_image] * num_samples) - - >>> p_params = replicate(params) - >>> prompt_ids = shard(prompt_ids) - >>> negative_prompt_ids = shard(negative_prompt_ids) - >>> processed_image = shard(processed_image) - - >>> output = pipe( - ... prompt_ids=prompt_ids, - ... image=processed_image, - ... params=p_params, - ... prng_seed=rng, - ... num_inference_steps=50, - ... neg_prompt_ids=negative_prompt_ids, - ... jit=True, - ... ).images - - >>> output_images = pipe.numpy_to_pil(np.asarray(output.reshape((num_samples,) + output.shape[-3:]))) - >>> output_images = make_image_grid(output_images, num_samples // 4, 4) - >>> output_images.save("generated_image.png") - ``` -""" - - -class FlaxStableDiffusionControlNetPipeline(FlaxDiffusionPipeline): - r""" - Flax-based pipeline for text-to-image generation using Stable Diffusion with ControlNet Guidance. - - This model inherits from [`FlaxDiffusionPipeline`]. Check the superclass documentation for the generic methods - implemented for all pipelines (downloading, saving, running on a particular device, etc.). - - Args: - vae ([`FlaxAutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.FlaxCLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`FlaxUNet2DConditionModel`]): - A `FlaxUNet2DConditionModel` to denoise the encoded image latents. - controlnet ([`FlaxControlNetModel`]: - Provides additional conditioning to the `unet` during the denoising process. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of - [`FlaxDDIMScheduler`], [`FlaxLMSDiscreteScheduler`], [`FlaxPNDMScheduler`], or - [`FlaxDPMSolverMultistepScheduler`]. - safety_checker ([`FlaxStableDiffusionSafetyChecker`]): - 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - def __init__( - self, - vae: FlaxAutoencoderKL, - text_encoder: FlaxCLIPTextModel, - tokenizer: CLIPTokenizer, - unet: FlaxUNet2DConditionModel, - controlnet: FlaxControlNetModel, - scheduler: Union[ - FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler - ], - safety_checker: FlaxStableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - dtype: jnp.dtype = jnp.float32, - ): - super().__init__() - self.dtype = dtype - - if safety_checker is None: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) - 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]]): - if not isinstance(prompt, (str, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - text_input = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - - return text_input.input_ids - - def prepare_image_inputs(self, image: Union[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)}") - - if isinstance(image, Image.Image): - image = [image] - - processed_images = jnp.concatenate([preprocess(img, jnp.float32) for img in image]) - - return processed_images - - def _get_has_nsfw_concepts(self, features, params): - has_nsfw_concepts = self.safety_checker(features, params) - return has_nsfw_concepts - - def _run_safety_checker(self, images, safety_model_params, jit=False): - # safety_model_params should already be replicated when jit is True - pil_images = [Image.fromarray(image) for image in images] - features = self.feature_extractor(pil_images, return_tensors="np").pixel_values - - if jit: - features = shard(features) - has_nsfw_concepts = _p_get_has_nsfw_concepts(self, features, safety_model_params) - has_nsfw_concepts = unshard(has_nsfw_concepts) - safety_model_params = unreplicate(safety_model_params) - else: - has_nsfw_concepts = self._get_has_nsfw_concepts(features, safety_model_params) - - images_was_copied = False - for idx, has_nsfw_concept in enumerate(has_nsfw_concepts): - if has_nsfw_concept: - if not images_was_copied: - images_was_copied = True - images = images.copy() - - images[idx] = np.zeros(images[idx].shape, dtype=np.uint8) # black image - - if any(has_nsfw_concepts): - warnings.warn( - "Potential NSFW content was detected in one or more images. A black image will be returned" - " instead. Try again with a different prompt and/or seed." - ) - - return images, has_nsfw_concepts - - def _generate( - self, - prompt_ids: jnp.ndarray, - image: jnp.ndarray, - params: Union[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, - controlnet_conditioning_scale: float = 1.0, - ): - height, width = image.shape[-2:] - if height % 64 != 0 or width % 64 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 64 but are {height} and {width}.") - - # get prompt text embeddings - prompt_embeds = self.text_encoder(prompt_ids, params=params["text_encoder"])[0] - - # TODO: currently it is assumed `do_classifier_free_guidance = guidance_scale > 1.0` - # implement this conditional `do_classifier_free_guidance = guidance_scale > 1.0` - batch_size = prompt_ids.shape[0] - - max_length = prompt_ids.shape[-1] - - if neg_prompt_ids is None: - uncond_input = self.tokenizer( - [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="np" - ).input_ids - else: - uncond_input = neg_prompt_ids - negative_prompt_embeds = self.text_encoder(uncond_input, params=params["text_encoder"])[0] - context = jnp.concatenate([negative_prompt_embeds, prompt_embeds]) - - image = jnp.concatenate([image] * 2) - - latents_shape = ( - batch_size, - self.unet.config.in_channels, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - if latents is None: - latents = jax.random.normal(prng_seed, shape=latents_shape, dtype=jnp.float32) - else: - if latents.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - - def loop_body(step, args): - latents, scheduler_state = args - # 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 - latents_input = jnp.concatenate([latents] * 2) - - t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step] - timestep = jnp.broadcast_to(t, latents_input.shape[0]) - - latents_input = self.scheduler.scale_model_input(scheduler_state, latents_input, t) - - down_block_res_samples, mid_block_res_sample = self.controlnet.apply( - {"params": params["controlnet"]}, - jnp.array(latents_input), - jnp.array(timestep, dtype=jnp.int32), - encoder_hidden_states=context, - controlnet_cond=image, - conditioning_scale=controlnet_conditioning_scale, - return_dict=False, - ) - - # predict the noise residual - noise_pred = self.unet.apply( - {"params": params["unet"]}, - jnp.array(latents_input), - jnp.array(timestep, dtype=jnp.int32), - encoder_hidden_states=context, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - ).sample - - # perform guidance - noise_pred_uncond, noise_prediction_text = jnp.split(noise_pred, 2, axis=0) - noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents, scheduler_state = self.scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() - return latents, scheduler_state - - scheduler_state = self.scheduler.set_timesteps( - params["scheduler"], num_inference_steps=num_inference_steps, shape=latents_shape - ) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * params["scheduler"].init_noise_sigma - - if DEBUG: - # run with python for loop - for i in range(num_inference_steps): - latents, scheduler_state = loop_body(i, (latents, scheduler_state)) - else: - latents, _ = jax.lax.fori_loop(0, num_inference_steps, loop_body, (latents, scheduler_state)) - - # scale and decode the image latents with vae - latents = 1 / self.vae.config.scaling_factor * latents - image = self.vae.apply({"params": params["vae"]}, latents, method=self.vae.decode).sample - - image = (image / 2 + 0.5).clip(0, 1).transpose(0, 2, 3, 1) - return image - - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt_ids: jnp.ndarray, - image: jnp.ndarray, - params: Union[Dict, FrozenDict], - prng_seed: jax.Array, - num_inference_steps: int = 50, - guidance_scale: Union[float, jnp.ndarray] = 7.5, - latents: jnp.ndarray = None, - neg_prompt_ids: jnp.ndarray = None, - controlnet_conditioning_scale: Union[float, jnp.ndarray] = 1.0, - return_dict: bool = True, - jit: bool = False, - ): - r""" - The call function to the pipeline for generation. - - Args: - prompt_ids (`jnp.ndarray`): - The prompt or prompts to guide the image generation. - image (`jnp.ndarray`): - Array representing the ControlNet input condition to provide guidance to the `unet` for generation. - params (`Dict` or `FrozenDict`): - Dictionary containing the model parameters/weights. - prng_seed (`jax.Array`): - Array containing random number generator key. - 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): - 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`. - latents (`jnp.ndarray`, *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 - array is generated by sampling using the supplied random `generator`. - controlnet_conditioning_scale (`float` or `jnp.ndarray`, *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`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] instead of - a plain tuple. - jit (`bool`, defaults to `False`): - Whether to run `pmap` versions of the generation and safety scoring functions. - - - - This argument exists because `__call__` is not yet end-to-end pmap-able. It will be removed in a - future release. - - - - Examples: - - Returns: - [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] 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. - """ - - height, width = image.shape[-2:] - - if isinstance(guidance_scale, float): - # Convert to a tensor so each device gets a copy. Follow the prompt_ids for - # shape information, as they may be sharded (when `jit` is `True`), or not. - guidance_scale = jnp.array([guidance_scale] * prompt_ids.shape[0]) - if len(prompt_ids.shape) > 2: - # Assume sharded - guidance_scale = guidance_scale[:, None] - - if isinstance(controlnet_conditioning_scale, float): - # Convert to a tensor so each device gets a copy. Follow the prompt_ids for - # shape information, as they may be sharded (when `jit` is `True`), or not. - controlnet_conditioning_scale = jnp.array([controlnet_conditioning_scale] * prompt_ids.shape[0]) - if len(prompt_ids.shape) > 2: - # Assume sharded - controlnet_conditioning_scale = controlnet_conditioning_scale[:, None] - - if jit: - images = _p_generate( - self, - prompt_ids, - image, - params, - prng_seed, - num_inference_steps, - guidance_scale, - latents, - neg_prompt_ids, - controlnet_conditioning_scale, - ) - else: - images = self._generate( - prompt_ids, - image, - params, - prng_seed, - num_inference_steps, - guidance_scale, - latents, - neg_prompt_ids, - controlnet_conditioning_scale, - ) - - if self.safety_checker is not None: - safety_params = params["safety_checker"] - images_uint8_casted = (images * 255).round().astype("uint8") - num_devices, batch_size = images.shape[:2] - - images_uint8_casted = np.asarray(images_uint8_casted).reshape(num_devices * batch_size, height, width, 3) - images_uint8_casted, has_nsfw_concept = self._run_safety_checker(images_uint8_casted, safety_params, jit) - images = np.array(images) - - # block images - if any(has_nsfw_concept): - for i, is_nsfw in enumerate(has_nsfw_concept): - if is_nsfw: - images[i] = np.asarray(images_uint8_casted[i]) - - images = images.reshape(num_devices, batch_size, height, width, 3) - else: - images = np.asarray(images) - has_nsfw_concept = False - - if not return_dict: - return (images, has_nsfw_concept) - - return FlaxStableDiffusionPipelineOutput(images=images, nsfw_content_detected=has_nsfw_concept) - - -# Static argnums are pipe, num_inference_steps. A change would trigger recompilation. -# Non-static args are (sharded) input tensors mapped over their first dimension (hence, `0`). -@partial( - jax.pmap, - in_axes=(None, 0, 0, 0, 0, None, 0, 0, 0, 0), - static_broadcasted_argnums=(0, 5), -) -def _p_generate( - pipe, - prompt_ids, - image, - params, - prng_seed, - num_inference_steps, - guidance_scale, - latents, - neg_prompt_ids, - controlnet_conditioning_scale, -): - return pipe._generate( - prompt_ids, - image, - params, - prng_seed, - num_inference_steps, - guidance_scale, - latents, - neg_prompt_ids, - controlnet_conditioning_scale, - ) - - -@partial(jax.pmap, static_broadcasted_argnums=(0,)) -def _p_get_has_nsfw_concepts(pipe, features, params): - return pipe._get_has_nsfw_concepts(features, params) - - -def unshard(x: jnp.ndarray): - # einops.rearrange(x, 'd b ... -> (d b) ...') - num_devices, batch_size = x.shape[:2] - rest = x.shape[2:] - return x.reshape(num_devices * batch_size, *rest) - - -def preprocess(image, dtype): - image = image.convert("RGB") - w, h = image.size - w, h = (x - x % 64 for x in (w, h)) # resize to integer multiple of 64 - image = image.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]) - image = jnp.array(image).astype(dtype) / 255.0 - image = image[None].transpose(0, 3, 1, 2) - return image diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_hunyuandit/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_hunyuandit/__init__.py deleted file mode 100644 index 34c59795de328f2fa2cbe610625d866e4710d4a2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_hunyuandit/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -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_hunyuandit_controlnet"] = ["HunyuanDiTControlNetPipeline"] - -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_hunyuandit_controlnet import HunyuanDiTControlNetPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_hunyuandit/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_hunyuandit/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index fcfc1a8989db1f58548663819f1e5ac30c165ca2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_hunyuandit/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_hunyuandit/__pycache__/pipeline_hunyuandit_controlnet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_hunyuandit/__pycache__/pipeline_hunyuandit_controlnet.cpython-310.pyc deleted file mode 100644 index 2b7ae660330ef22b2ea97bc7ef9f87e6caf3baf6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_hunyuandit/__pycache__/pipeline_hunyuandit_controlnet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py deleted file mode 100644 index 9b9adf49012539a8793e66c4cec8c91b4fb0de11..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py +++ /dev/null @@ -1,1058 +0,0 @@ -# Copyright 2025 HunyuanDiT Authors 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 Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import torch -from transformers import BertModel, BertTokenizer, CLIPImageProcessor, MT5Tokenizer, T5EncoderModel - -from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...models import AutoencoderKL, HunyuanDiT2DControlNetModel, HunyuanDiT2DModel, HunyuanDiT2DMultiControlNetModel -from ...models.embeddings import get_2d_rotary_pos_embed -from ...pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker -from ...schedulers import DDPMScheduler -from ...utils import ( - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -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 - from diffusers import HunyuanDiT2DControlNetModel, HunyuanDiTControlNetPipeline - import torch - - controlnet = HunyuanDiT2DControlNetModel.from_pretrained( - "Tencent-Hunyuan/HunyuanDiT-v1.1-ControlNet-Diffusers-Canny", torch_dtype=torch.float16 - ) - - pipe = HunyuanDiTControlNetPipeline.from_pretrained( - "Tencent-Hunyuan/HunyuanDiT-v1.1-Diffusers", controlnet=controlnet, torch_dtype=torch.float16 - ) - pipe.to("cuda") - - from diffusers.utils import load_image - - cond_image = load_image( - "https://huggingface.co/Tencent-Hunyuan/HunyuanDiT-v1.1-ControlNet-Diffusers-Canny/resolve/main/canny.jpg?download=true" - ) - - ## You may also use English prompt as HunyuanDiT supports both English and Chinese - prompt = "在夜晚的酒店门前，一座古老的中国风格的狮子雕像矗立着，它的眼睛闪烁着光芒，仿佛在守护着这座建筑。背景是夜晚的酒店前，构图方式是特写，平视，居中构图。这张照片呈现了真实摄影风格，蕴含了中国雕塑文化，同时展现了神秘氛围" - # prompt="At night, an ancient Chinese-style lion statue stands in front of the hotel, its eyes gleaming as if guarding the building. The background is the hotel entrance at night, with a close-up, eye-level, and centered composition. This photo presents a realistic photographic style, embodies Chinese sculpture culture, and reveals a mysterious atmosphere." - image = pipe( - prompt, - height=1024, - width=1024, - control_image=cond_image, - num_inference_steps=50, - ).images[0] - ``` -""" - -STANDARD_RATIO = np.array( - [ - 1.0, # 1:1 - 4.0 / 3.0, # 4:3 - 3.0 / 4.0, # 3:4 - 16.0 / 9.0, # 16:9 - 9.0 / 16.0, # 9:16 - ] -) -STANDARD_SHAPE = [ - [(1024, 1024), (1280, 1280)], # 1:1 - [(1024, 768), (1152, 864), (1280, 960)], # 4:3 - [(768, 1024), (864, 1152), (960, 1280)], # 3:4 - [(1280, 768)], # 16:9 - [(768, 1280)], # 9:16 -] -STANDARD_AREA = [np.array([w * h for w, h in shapes]) for shapes in STANDARD_SHAPE] -SUPPORTED_SHAPE = [ - (1024, 1024), - (1280, 1280), # 1:1 - (1024, 768), - (1152, 864), - (1280, 960), # 4:3 - (768, 1024), - (864, 1152), - (960, 1280), # 3:4 - (1280, 768), # 16:9 - (768, 1280), # 9:16 -] - - -def map_to_standard_shapes(target_width, target_height): - target_ratio = target_width / target_height - closest_ratio_idx = np.argmin(np.abs(STANDARD_RATIO - target_ratio)) - closest_area_idx = np.argmin(np.abs(STANDARD_AREA[closest_ratio_idx] - target_width * target_height)) - width, height = STANDARD_SHAPE[closest_ratio_idx][closest_area_idx] - return width, height - - -def get_resize_crop_region_for_grid(src, tgt_size): - th = tw = tgt_size - h, w = src - - r = h / w - - # resize - if r > 1: - resize_height = th - resize_width = int(round(th / h * w)) - else: - resize_width = tw - resize_height = int(round(tw / w * h)) - - crop_top = int(round((th - resize_height) / 2.0)) - crop_left = int(round((tw - resize_width) / 2.0)) - - return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width) - - -# 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 HunyuanDiTControlNetPipeline(DiffusionPipeline): - r""" - Pipeline for English/Chinese-to-image generation using HunyuanDiT. - - 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.) - - HunyuanDiT uses two text encoders: [mT5](https://huggingface.co/google/mt5-base) and [bilingual CLIP](fine-tuned by - ourselves) - - Args: - 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`]): - 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`]): - A `BertTokenizer` or `CLIPTokenizer` to tokenize text. - transformer ([`HunyuanDiT2DModel`]): - The HunyuanDiT model designed by Tencent Hunyuan. - text_encoder_2 (`T5EncoderModel`): - The mT5 embedder. Specifically, it is 't5-v1_1-xxl'. - tokenizer_2 (`MT5Tokenizer`): - 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`]): - 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. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" - _optional_components = [ - "safety_checker", - "feature_extractor", - "text_encoder_2", - "tokenizer_2", - "text_encoder", - "tokenizer", - ] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "prompt_embeds_2", - "negative_prompt_embeds_2", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: BertModel, - tokenizer: BertTokenizer, - transformer: HunyuanDiT2DModel, - 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[MT5Tokenizer] = None, - requires_safety_checker: bool = True, - ): - super().__init__() - if isinstance(controlnet, (list, tuple)): - controlnet = HunyuanDiT2DMultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - transformer=transformer, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - text_encoder_2=text_encoder_2, - controlnet=controlnet, - ) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - 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.register_to_config(requires_safety_checker=requires_safety_checker) - self.default_sample_size = ( - self.transformer.config.sample_size - if hasattr(self, "transformer") and self.transformer is not None - else 128 - ) - - # Copied from diffusers.pipelines.hunyuandit.pipeline_hunyuandit.HunyuanDiTPipeline.encode_prompt - def encode_prompt( - self, - prompt: str, - device: torch.device = None, - 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, - text_encoder_index: int = 0, - ): - 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 - dtype (`torch.dtype`): - torch dtype - 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. - prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the prompt. Required when `prompt_embeds` is passed directly. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the negative prompt. Required when `negative_prompt_embeds` is passed directly. - max_sequence_length (`int`, *optional*): maximum sequence length to use for the prompt. - text_encoder_index (`int`, *optional*): - Index of the text encoder to use. `0` for clip and `1` for T5. - """ - if dtype is None: - if self.text_encoder_2 is not None: - dtype = self.text_encoder_2.dtype - elif self.transformer is not None: - dtype = self.transformer.dtype - else: - dtype = None - - if device is None: - device = self._execution_device - - tokenizers = [self.tokenizer, self.tokenizer_2] - text_encoders = [self.text_encoder, self.text_encoder_2] - - tokenizer = tokenizers[text_encoder_index] - text_encoder = text_encoders[text_encoder_index] - - if max_sequence_length is None: - if text_encoder_index == 0: - max_length = 77 - if text_encoder_index == 1: - max_length = 256 - else: - max_length = max_sequence_length - - 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: - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=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_attention_mask = text_inputs.attention_mask.to(device) - prompt_embeds = text_encoder( - text_input_ids.to(device), - attention_mask=prompt_attention_mask, - ) - prompt_embeds = prompt_embeds[0] - prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1) - - prompt_embeds = prompt_embeds.to(dtype=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 - - max_length = prompt_embeds.shape[1] - uncond_input = tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - - negative_prompt_attention_mask = uncond_input.attention_mask.to(device) - negative_prompt_embeds = text_encoder( - uncond_input.input_ids.to(device), - attention_mask=negative_prompt_attention_mask, - ) - negative_prompt_embeds = negative_prompt_embeds[0] - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 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] - - negative_prompt_embeds = negative_prompt_embeds.to(dtype=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) - - return prompt_embeds, negative_prompt_embeds, prompt_attention_mask, negative_prompt_attention_mask - - # 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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.hunyuandit.pipeline_hunyuandit.HunyuanDiTPipeline.check_inputs - def check_inputs( - self, - prompt, - height, - width, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - prompt_embeds_2=None, - negative_prompt_embeds_2=None, - prompt_attention_mask_2=None, - negative_prompt_attention_mask_2=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_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 None and prompt_embeds_2 is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds_2`. Cannot leave both `prompt` and `prompt_embeds_2` 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 prompt_embeds_2 is not None and prompt_attention_mask_2 is None: - raise ValueError("Must provide `prompt_attention_mask_2` when specifying `prompt_embeds_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." - ) - - 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 negative_prompt_embeds_2 is not None and negative_prompt_attention_mask_2 is None: - raise ValueError( - "Must provide `negative_prompt_attention_mask_2` when specifying `negative_prompt_embeds_2`." - ) - 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_2 is not None and negative_prompt_embeds_2 is not None: - if prompt_embeds_2.shape != negative_prompt_embeds_2.shape: - raise ValueError( - "`prompt_embeds_2` and `negative_prompt_embeds_2` must have the same shape when passed directly, but" - f" got: `prompt_embeds_2` {prompt_embeds_2.shape} != `negative_prompt_embeds_2`" - f" {negative_prompt_embeds_2.shape}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - # Copied from diffusers.pipelines.controlnet_sd3.pipeline_stable_diffusion_3_controlnet.StableDiffusion3ControlNetPipeline.prepare_image - def prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - if isinstance(image, torch.Tensor): - pass - else: - image = self.image_processor.preprocess(image, height=height, width=width) - - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - # 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 - - @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, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 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", - 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"], - 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), - use_resolution_binning: bool = True, - ): - r""" - The call function to the pipeline for generation with HunyuanDiT. - - Args: - 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. - width (`int`): - 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. This parameter is modulated by `strength`. - 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): - The percentage of total steps at which the ControlNet starts applying. - 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]]`): - 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): - 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*): - 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): - The number of images to generate per prompt. - 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*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - 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. - prompt_embeds_2 (`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. - 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. - negative_prompt_embeds_2 (`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. - prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the prompt. Required when `prompt_embeds` is passed directly. - prompt_attention_mask_2 (`torch.Tensor`, *optional*): - Attention mask for the prompt. Required when `prompt_embeds_2` is passed directly. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the negative prompt. Required when `negative_prompt_embeds` is passed directly. - negative_prompt_attention_mask_2 (`torch.Tensor`, *optional*): - Attention mask for the negative prompt. Required when `negative_prompt_embeds_2` is passed directly. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback_on_step_end (`Callable[[int, int, Dict], 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*): - 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)`): - The original size of the image. Used to calculate the time ids. - 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)`): - 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 - standard resolution. Supported resolutions are 1024x1024, 1280x1280, 1024x768, 1152x864, 1280x960, - 768x1024, 864x1152, 960x1280, 1280x768, and 768x1280. It is recommended to set this to `True`. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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 isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - # 0. default height and width - height = height or self.default_sample_size * self.vae_scale_factor - width = width or self.default_sample_size * self.vae_scale_factor - height = int((height // 16) * 16) - width = int((width // 16) * 16) - - if use_resolution_binning and (height, width) not in SUPPORTED_SHAPE: - width, height = map_to_standard_shapes(width, height) - height = int(height) - width = int(width) - logger.warning(f"Reshaped to (height, width)=({height}, {width}), Supported shapes are {SUPPORTED_SHAPE}") - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - prompt_embeds_2, - negative_prompt_embeds_2, - prompt_attention_mask_2, - negative_prompt_attention_mask_2, - callback_on_step_end_tensor_inputs, - ) - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - 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. Encode input prompt - - ( - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt=prompt, - device=device, - dtype=self.transformer.dtype, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - 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, - max_sequence_length=77, - text_encoder_index=0, - ) - ( - prompt_embeds_2, - negative_prompt_embeds_2, - prompt_attention_mask_2, - negative_prompt_attention_mask_2, - ) = self.encode_prompt( - prompt=prompt, - device=device, - dtype=self.transformer.dtype, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds_2, - negative_prompt_embeds=negative_prompt_embeds_2, - prompt_attention_mask=prompt_attention_mask_2, - negative_prompt_attention_mask=negative_prompt_attention_mask_2, - max_sequence_length=256, - text_encoder_index=1, - ) - - # 4. Prepare control image - if isinstance(self.controlnet, HunyuanDiT2DControlNetModel): - control_image = self.prepare_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=self.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=False, - ) - height, width = control_image.shape[-2:] - - control_image = self.vae.encode(control_image).latent_dist.sample() - control_image = control_image * self.vae.config.scaling_factor - - elif isinstance(self.controlnet, HunyuanDiT2DMultiControlNetModel): - control_images = [] - - for control_image_ in control_image: - control_image_ = self.prepare_image( - image=control_image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=self.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=False, - ) - - control_image_ = self.vae.encode(control_image_).latent_dist.sample() - control_image_ = control_image_ * self.vae.config.scaling_factor - - control_images.append(control_image_) - - control_image = control_images - else: - assert False - - # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 6. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 8. create image_rotary_emb, style embedding & time ids - grid_height = height // 8 // self.transformer.config.patch_size - grid_width = width // 8 // self.transformer.config.patch_size - base_size = 512 // 8 // self.transformer.config.patch_size - grid_crops_coords = get_resize_crop_region_for_grid((grid_height, grid_width), base_size) - image_rotary_emb = get_2d_rotary_pos_embed( - self.transformer.inner_dim // self.transformer.num_heads, - grid_crops_coords, - (grid_height, grid_width), - device=device, - output_type="pt", - ) - - style = torch.tensor([0], device=device) - - target_size = target_size or (height, width) - add_time_ids = list(original_size + target_size + crops_coords_top_left) - add_time_ids = torch.tensor([add_time_ids], dtype=prompt_embeds.dtype) - - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask]) - prompt_embeds_2 = torch.cat([negative_prompt_embeds_2, prompt_embeds_2]) - prompt_attention_mask_2 = torch.cat([negative_prompt_attention_mask_2, prompt_attention_mask_2]) - add_time_ids = torch.cat([add_time_ids] * 2, dim=0) - style = torch.cat([style] * 2, dim=0) - - prompt_embeds = prompt_embeds.to(device=device) - prompt_attention_mask = prompt_attention_mask.to(device=device) - prompt_embeds_2 = prompt_embeds_2.to(device=device) - prompt_attention_mask_2 = prompt_attention_mask_2.to(device=device) - add_time_ids = add_time_ids.to(dtype=prompt_embeds.dtype, device=device).repeat( - batch_size * num_images_per_prompt, 1 - ) - style = style.to(device=device).repeat(batch_size * num_images_per_prompt) - - # 9. Denoising loop - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # expand scalar t to 1-D tensor to match the 1st dim of latent_model_input - t_expand = torch.tensor([t] * latent_model_input.shape[0], device=device).to( - dtype=latent_model_input.dtype - ) - - # controlnet(s) inference - control_block_samples = self.controlnet( - latent_model_input, - t_expand, - encoder_hidden_states=prompt_embeds, - text_embedding_mask=prompt_attention_mask, - encoder_hidden_states_t5=prompt_embeds_2, - text_embedding_mask_t5=prompt_attention_mask_2, - image_meta_size=add_time_ids, - style=style, - image_rotary_emb=image_rotary_emb, - return_dict=False, - controlnet_cond=control_image, - conditioning_scale=controlnet_conditioning_scale, - )[0] - - # predict the noise residual - noise_pred = self.transformer( - latent_model_input, - t_expand, - encoder_hidden_states=prompt_embeds, - text_embedding_mask=prompt_attention_mask, - encoder_hidden_states_t5=prompt_embeds_2, - text_embedding_mask_t5=prompt_attention_mask_2, - image_meta_size=add_time_ids, - style=style, - image_rotary_emb=image_rotary_emb, - return_dict=False, - controlnet_block_samples=control_block_samples, - )[0] - - noise_pred, _ = noise_pred.chunk(2, dim=1) - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - if self.do_classifier_free_guidance and guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - prompt_embeds_2 = callback_outputs.pop("prompt_embeds_2", prompt_embeds_2) - negative_prompt_embeds_2 = callback_outputs.pop( - "negative_prompt_embeds_2", negative_prompt_embeds_2 - ) - - 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 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/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/__init__.py deleted file mode 100644 index aeb61dc8e247555e654d2f84c035dff53328ea0d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/__init__.py +++ /dev/null @@ -1,57 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - get_objects_from_module, - is_flax_available, - 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_stable_diffusion_3_controlnet"] = ["StableDiffusion3ControlNetPipeline"] - _import_structure["pipeline_stable_diffusion_3_controlnet_inpainting"] = [ - "StableDiffusion3ControlNetInpaintingPipeline" - ] - -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_stable_diffusion_3_controlnet import StableDiffusion3ControlNetPipeline - from .pipeline_stable_diffusion_3_controlnet_inpainting import StableDiffusion3ControlNetInpaintingPipeline - - try: - if not (is_transformers_available() and is_flax_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_flax_and_transformers_objects import * # noqa F403 - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index b51b4dd11a99b1ca3399d3ad9f75cc56b987b954..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/__pycache__/pipeline_stable_diffusion_3_controlnet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/__pycache__/pipeline_stable_diffusion_3_controlnet.cpython-310.pyc deleted file mode 100644 index 4b84964a4dfe0cf1a78e3fb37805f8fc4a3529fa..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/__pycache__/pipeline_stable_diffusion_3_controlnet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/__pycache__/pipeline_stable_diffusion_3_controlnet_inpainting.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/__pycache__/pipeline_stable_diffusion_3_controlnet_inpainting.cpython-310.pyc deleted file mode 100644 index 8dfa692148b7c6c0391c526ec55adebb4ff7d9a2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/__pycache__/pipeline_stable_diffusion_3_controlnet_inpainting.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py deleted file mode 100644 index e31e3a01787297980ccf4e231499c247671e83be..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py +++ /dev/null @@ -1,1244 +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. - -import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import ( - CLIPTextModelWithProjection, - CLIPTokenizer, - SiglipImageProcessor, - SiglipVisionModel, - T5EncoderModel, - T5TokenizerFast, -) - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, SD3IPAdapterMixin, SD3LoraLoaderMixin -from ...models.autoencoders import AutoencoderKL -from ...models.controlnets.controlnet_sd3 import SD3ControlNetModel, SD3MultiControlNetModel -from ...models.transformers import SD3Transformer2DModel -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 ..stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput - - -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 StableDiffusion3ControlNetPipeline - >>> from diffusers.models import SD3ControlNetModel, SD3MultiControlNetModel - >>> from diffusers.utils import load_image - - >>> controlnet = SD3ControlNetModel.from_pretrained("InstantX/SD3-Controlnet-Canny", torch_dtype=torch.float16) - - >>> pipe = StableDiffusion3ControlNetPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-3-medium-diffusers", controlnet=controlnet, torch_dtype=torch.float16 - ... ) - >>> pipe.to("cuda") - >>> control_image = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png" - ... ) - >>> prompt = "A bird in space" - >>> image = pipe( - ... prompt, control_image=control_image, height=1024, width=768, controlnet_conditioning_scale=0.7 - ... ).images[0] - >>> image.save("sd3.png") - ``` -""" - - -# 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 StableDiffusion3ControlNetPipeline( - DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, SD3IPAdapterMixin -): - r""" - Args: - transformer ([`SD3Transformer2DModel`]): - 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 ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModelWithProjection`]): - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), - specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant, - with an additional added projection layer that is initialized with a diagonal matrix with the `hidden_size` - as its dimension. - text_encoder_2 ([`CLIPTextModelWithProjection`]): - [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. - text_encoder_3 ([`T5EncoderModel`]): - Frozen text-encoder. Stable Diffusion 3 uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the - [t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) 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). - tokenizer_3 (`T5TokenizerFast`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - 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. - image_encoder (`SiglipVisionModel`, *optional*): - Pre-trained Vision Model for IP Adapter. - feature_extractor (`SiglipImageProcessor`, *optional*): - Image processor for IP Adapter. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->image_encoder->transformer->vae" - _optional_components = ["image_encoder", "feature_extractor"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "negative_pooled_prompt_embeds"] - - def __init__( - self, - transformer: SD3Transformer2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - text_encoder_2: CLIPTextModelWithProjection, - 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, - ): - super().__init__() - if isinstance(controlnet, (list, tuple)): - controlnet = SD3MultiControlNetModel(controlnet) - if isinstance(controlnet, SD3MultiControlNetModel): - for controlnet_model in controlnet.nets: - # for SD3.5 8b controlnet, it shares the pos_embed with the transformer - if ( - hasattr(controlnet_model.config, "use_pos_embed") - and controlnet_model.config.use_pos_embed is False - ): - pos_embed = controlnet_model._get_pos_embed_from_transformer(transformer) - controlnet_model.pos_embed = pos_embed.to(controlnet_model.dtype).to(controlnet_model.device) - elif isinstance(controlnet, SD3ControlNetModel): - if hasattr(controlnet.config, "use_pos_embed") and controlnet.config.use_pos_embed is False: - pos_embed = controlnet._get_pos_embed_from_transformer(transformer) - controlnet.pos_embed = pos_embed.to(controlnet.dtype).to(controlnet.device) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - text_encoder_3=text_encoder_3, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - tokenizer_3=tokenizer_3, - transformer=transformer, - scheduler=scheduler, - controlnet=controlnet, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = ( - self.transformer.config.sample_size - if hasattr(self, "transformer") and self.transformer is not None - else 128 - ) - - # 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, - num_images_per_prompt: int = 1, - max_sequence_length: int = 256, - 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 self.text_encoder_3 is None: - return torch.zeros( - ( - batch_size * num_images_per_prompt, - self.tokenizer_max_length, - self.transformer.config.joint_attention_dim, - ), - device=device, - dtype=dtype, - ) - - text_inputs = self.tokenizer_3( - 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 - untruncated_ids = self.tokenizer_3(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_3.batch_decode(untruncated_ids[:, self.tokenizer_max_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_3(text_input_ids.to(device))[0] - - dtype = self.text_encoder_3.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clip_skip: Optional[int] = None, - clip_model_index: int = 0, - ): - device = device or self._execution_device - - clip_tokenizers = [self.tokenizer, self.tokenizer_2] - clip_text_encoders = [self.text_encoder, self.text_encoder_2] - - tokenizer = clip_tokenizers[clip_model_index] - text_encoder = clip_text_encoders[clip_model_index] - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer_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[:, self.tokenizer_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) - pooled_prompt_embeds = prompt_embeds[0] - - if clip_skip is None: - prompt_embeds = prompt_embeds.hidden_states[-2] - else: - prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)] - - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - _, 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(batch_size * num_images_per_prompt, seq_len, -1) - - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) - pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds, pooled_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, - 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, - max_sequence_length: int = 256, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - 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`): - 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 all the text-encoders. - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. - 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. - 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, SD3LoraLoaderMixin): - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - 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] - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - prompt_3 = prompt_3 or prompt - prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3 - - prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=clip_skip, - clip_model_index=0, - ) - prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds( - prompt=prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=clip_skip, - clip_model_index=1, - ) - clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1) - - t5_prompt_embed = self._get_t5_prompt_embeds( - prompt=prompt_3, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - clip_prompt_embeds = torch.nn.functional.pad( - clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1]) - ) - - prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2) - pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1) - - if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt_2 = negative_prompt_2 or negative_prompt - negative_prompt_3 = negative_prompt_3 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 - ) - negative_prompt_3 = ( - batch_size * [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3 - ) - - 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_embed, negative_pooled_prompt_embed = self._get_clip_prompt_embeds( - negative_prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=None, - clip_model_index=0, - ) - negative_prompt_2_embed, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds( - negative_prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=None, - clip_model_index=1, - ) - negative_clip_prompt_embeds = torch.cat([negative_prompt_embed, negative_prompt_2_embed], dim=-1) - - t5_negative_prompt_embed = self._get_t5_prompt_embeds( - prompt=negative_prompt_3, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - negative_clip_prompt_embeds = torch.nn.functional.pad( - negative_clip_prompt_embeds, - (0, t5_negative_prompt_embed.shape[-1] - negative_clip_prompt_embeds.shape[-1]), - ) - - negative_prompt_embeds = torch.cat([negative_clip_prompt_embeds, t5_negative_prompt_embed], dim=-2) - negative_pooled_prompt_embeds = torch.cat( - [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1 - ) - - if self.text_encoder is not None: - if isinstance(self, SD3LoraLoaderMixin) 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, SD3LoraLoaderMixin) 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, - prompt_3, - height, - width, - negative_prompt=None, - negative_prompt_2=None, - negative_prompt_3=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=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_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_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_3 is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt_3`: {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)}") - elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)): - raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}") - - 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." - ) - elif negative_prompt_3 is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} 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`." - ) - - 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}") - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_latents - 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=device, dtype=dtype) - - 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 prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - if isinstance(image, torch.Tensor): - pass - else: - image = self.image_processor.preprocess(image, height=height, width=width) - - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - @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 - - @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 - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_image - def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch.Tensor: - """Encodes the given image into a feature representation using a pre-trained image encoder. - - Args: - image (`PipelineImageInput`): - Input image to be encoded. - device: (`torch.device`): - Torch device. - - Returns: - `torch.Tensor`: The encoded image feature representation. - """ - if not isinstance(image, torch.Tensor): - image = self.feature_extractor(image, return_tensors="pt").pixel_values - - image = image.to(device=device, dtype=self.dtype) - - return self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - - # 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, - num_images_per_prompt: int = 1, - do_classifier_free_guidance: bool = True, - ) -> torch.Tensor: - """Prepares image embeddings for use in the IP-Adapter. - - Either `ip_adapter_image` or `ip_adapter_image_embeds` must be passed. - - Args: - ip_adapter_image (`PipelineImageInput`, *optional*): - The input image to extract features from for IP-Adapter. - ip_adapter_image_embeds (`torch.Tensor`, *optional*): - Precomputed image embeddings. - device: (`torch.device`, *optional*): - Torch device. - num_images_per_prompt (`int`, defaults to 1): - Number of images that should be generated per prompt. - do_classifier_free_guidance (`bool`, defaults to True): - Whether to use classifier free guidance or not. - """ - device = device or self._execution_device - - if ip_adapter_image_embeds is not None: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = ip_adapter_image_embeds.chunk(2) - else: - single_image_embeds = ip_adapter_image_embeds - elif ip_adapter_image is not None: - single_image_embeds = self.encode_image(ip_adapter_image, device) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.zeros_like(single_image_embeds) - else: - raise ValueError("Neither `ip_adapter_image_embeds` or `ip_adapter_image_embeds` were provided.") - - image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - negative_image_embeds = torch.cat([single_negative_image_embeds] * num_images_per_prompt, dim=0) - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - - return image_embeds.to(device=device) - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.enable_sequential_cpu_offload - def enable_sequential_cpu_offload(self, *args, **kwargs): - if self.image_encoder is not None and "image_encoder" not in self._exclude_from_cpu_offload: - logger.warning( - "`pipe.enable_sequential_cpu_offload()` might fail for `image_encoder` if it uses " - "`torch.nn.MultiheadAttention`. You can exclude `image_encoder` from CPU offloading by calling " - "`pipe._exclude_from_cpu_offload.append('image_encoder')` before `pipe.enable_sequential_cpu_offload()`." - ) - - super().enable_sequential_cpu_offload(*args, **kwargs) - - @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, - prompt_3: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 28, - sigmas: Optional[List[float]] = 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_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", - 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"], - max_sequence_length: int = 256, - ): - 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 `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - will be used instead - 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): - 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. - 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. - 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): - 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]]`): - 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): - 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*): - 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 instead - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. - ip_adapter_image (`PipelineImageInput`, *optional*): - Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`torch.Tensor`, *optional*): - Pre-generated image embeddings for IP-Adapter. 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. - 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 256): Maximum sequence length to use with the `prompt`. - - 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. - """ - - height = height or self.default_sample_size * self.vae_scale_factor - width = width or self.default_sample_size * self.vae_scale_factor - - controlnet_config = ( - self.controlnet.config - if isinstance(self.controlnet, SD3ControlNetModel) - else self.controlnet.nets[0].config - ) - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(self.controlnet.nets) if isinstance(self.controlnet, SD3MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - prompt_3, - height, - width, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._joint_attention_kwargs = joint_attention_kwargs - 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 - dtype = self.transformer.dtype - - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_3=prompt_3, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - do_classifier_free_guidance=self.do_classifier_free_guidance, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - device=device, - clip_skip=self.clip_skip, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - ) - - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) - - # 3. Prepare control image - if controlnet_config.force_zeros_for_pooled_projection: - # instantx sd3 controlnet does not apply shift factor - vae_shift_factor = 0 - else: - vae_shift_factor = self.vae.config.shift_factor - if isinstance(self.controlnet, SD3ControlNetModel): - control_image = self.prepare_image( - image=control_image, - 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, - ) - height, width = control_image.shape[-2:] - - control_image = self.vae.encode(control_image).latent_dist.sample() - control_image = (control_image - vae_shift_factor) * self.vae.config.scaling_factor - elif isinstance(self.controlnet, SD3MultiControlNetModel): - control_images = [] - - for control_image_ in control_image: - control_image_ = self.prepare_image( - image=control_image_, - 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_image_ = self.vae.encode(control_image_).latent_dist.sample() - control_image_ = (control_image_ - vae_shift_factor) * self.vae.config.scaling_factor - - control_images.append(control_image_) - - control_image = control_images - else: - assert False - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas) - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - # 5. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(self.controlnet, SD3ControlNetModel) else keeps) - - if controlnet_config.force_zeros_for_pooled_projection: - # instantx sd3 controlnet used zero pooled projection - controlnet_pooled_projections = torch.zeros_like(pooled_prompt_embeds) - else: - controlnet_pooled_projections = controlnet_pooled_projections or pooled_prompt_embeds - - if controlnet_config.joint_attention_dim is not None: - controlnet_encoder_hidden_states = prompt_embeds - else: - # SD35 official 8b controlnet does not use encoder_hidden_states - controlnet_encoder_hidden_states = None - - # 7. Prepare image embeddings - if (ip_adapter_image is not None and self.is_ip_adapter_active) or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - - if self.joint_attention_kwargs is None: - self._joint_attention_kwargs = {"ip_adapter_image_embeds": ip_adapter_image_embeds} - else: - self._joint_attention_kwargs.update(ip_adapter_image_embeds=ip_adapter_image_embeds) - - # 8. Denoising loop - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - # controlnet(s) inference - control_block_samples = self.controlnet( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=controlnet_encoder_hidden_states, - pooled_projections=controlnet_pooled_projections, - joint_attention_kwargs=self.joint_attention_kwargs, - controlnet_cond=control_image, - conditioning_scale=cond_scale, - return_dict=False, - )[0] - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - pooled_projections=pooled_prompt_embeds, - block_controlnet_hidden_states=control_block_samples, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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) - - # 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) - negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_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: - 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 StableDiffusion3PipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py deleted file mode 100644 index 000e080d3aea7fa9dfec280a8d00760245322a2e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py +++ /dev/null @@ -1,1264 +0,0 @@ -# Copyright 2025 Stability AI, The HuggingFace Team and The AlimamaCreative 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, Tuple, Union - -import torch -from transformers import ( - CLIPTextModelWithProjection, - CLIPTokenizer, - SiglipImageProcessor, - SiglipModel, - T5EncoderModel, - T5TokenizerFast, -) - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, SD3IPAdapterMixin, SD3LoraLoaderMixin -from ...models.autoencoders import AutoencoderKL -from ...models.controlnets.controlnet_sd3 import SD3ControlNetModel, SD3MultiControlNetModel -from ...models.transformers import SD3Transformer2DModel -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 ..stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput - - -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.utils import load_image, check_min_version - >>> from diffusers.pipelines import StableDiffusion3ControlNetInpaintingPipeline - >>> from diffusers.models.controlnet_sd3 import SD3ControlNetModel - - >>> controlnet = SD3ControlNetModel.from_pretrained( - ... "alimama-creative/SD3-Controlnet-Inpainting", use_safetensors=True, extra_conditioning_channels=1 - ... ) - >>> pipe = StableDiffusion3ControlNetInpaintingPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-3-medium-diffusers", - ... controlnet=controlnet, - ... torch_dtype=torch.float16, - ... ) - >>> pipe.text_encoder.to(torch.float16) - >>> pipe.controlnet.to(torch.float16) - >>> pipe.to("cuda") - - >>> image = load_image( - ... "https://huggingface.co/alimama-creative/SD3-Controlnet-Inpainting/resolve/main/images/dog.png" - ... ) - >>> mask = load_image( - ... "https://huggingface.co/alimama-creative/SD3-Controlnet-Inpainting/resolve/main/images/dog_mask.png" - ... ) - >>> width = 1024 - >>> height = 1024 - >>> prompt = "A cat is sitting next to a puppy." - >>> generator = torch.Generator(device="cuda").manual_seed(24) - >>> res_image = pipe( - ... negative_prompt="deformed, distorted, disfigured, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, NSFW", - ... prompt=prompt, - ... height=height, - ... width=width, - ... control_image=image, - ... control_mask=mask, - ... num_inference_steps=28, - ... generator=generator, - ... controlnet_conditioning_scale=0.95, - ... guidance_scale=7, - ... ).images[0] - >>> res_image.save(f"sd3.png") - ``` -""" - - -# 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 StableDiffusion3ControlNetInpaintingPipeline( - DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, SD3IPAdapterMixin -): - r""" - Args: - transformer ([`SD3Transformer2DModel`]): - 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 ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModelWithProjection`]): - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), - specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant, - with an additional added projection layer that is initialized with a diagonal matrix with the `hidden_size` - as its dimension. - text_encoder_2 ([`CLIPTextModelWithProjection`]): - [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. - text_encoder_3 ([`T5EncoderModel`]): - Frozen text-encoder. Stable Diffusion 3 uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the - [t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) 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). - tokenizer_3 (`T5TokenizerFast`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - 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. - image_encoder (`PreTrainedModel`, *optional*): - Pre-trained Vision Model for IP Adapter. - feature_extractor (`BaseImageProcessor`, *optional*): - Image processor for IP Adapter. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->image_encoder->transformer->vae" - _optional_components = ["image_encoder", "feature_extractor"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "negative_pooled_prompt_embeds"] - - def __init__( - self, - transformer: SD3Transformer2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer_2: CLIPTokenizer, - text_encoder_3: T5EncoderModel, - tokenizer_3: T5TokenizerFast, - controlnet: Union[ - SD3ControlNetModel, List[SD3ControlNetModel], Tuple[SD3ControlNetModel], SD3MultiControlNetModel - ], - image_encoder: SiglipModel = None, - feature_extractor: Optional[SiglipImageProcessor] = None, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - text_encoder_3=text_encoder_3, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - tokenizer_3=tokenizer_3, - transformer=transformer, - scheduler=scheduler, - controlnet=controlnet, - 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.image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_resize=True, do_convert_rgb=True, do_normalize=True - ) - self.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, - do_resize=True, - do_convert_grayscale=True, - do_normalize=False, - do_binarize=True, - ) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = ( - self.transformer.config.sample_size - if hasattr(self, "transformer") and self.transformer is not None - else 128 - ) - self.patch_size = ( - self.transformer.config.patch_size if hasattr(self, "transformer") and self.transformer is not None else 2 - ) - - # 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, - num_images_per_prompt: int = 1, - max_sequence_length: int = 256, - 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 self.text_encoder_3 is None: - return torch.zeros( - ( - batch_size * num_images_per_prompt, - self.tokenizer_max_length, - self.transformer.config.joint_attention_dim, - ), - device=device, - dtype=dtype, - ) - - text_inputs = self.tokenizer_3( - 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 - untruncated_ids = self.tokenizer_3(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_3.batch_decode(untruncated_ids[:, self.tokenizer_max_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_3(text_input_ids.to(device))[0] - - dtype = self.text_encoder_3.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clip_skip: Optional[int] = None, - clip_model_index: int = 0, - ): - device = device or self._execution_device - - clip_tokenizers = [self.tokenizer, self.tokenizer_2] - clip_text_encoders = [self.text_encoder, self.text_encoder_2] - - tokenizer = clip_tokenizers[clip_model_index] - text_encoder = clip_text_encoders[clip_model_index] - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer_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[:, self.tokenizer_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) - pooled_prompt_embeds = prompt_embeds[0] - - if clip_skip is None: - prompt_embeds = prompt_embeds.hidden_states[-2] - else: - prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)] - - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - _, 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(batch_size * num_images_per_prompt, seq_len, -1) - - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) - pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds, pooled_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, - 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, - max_sequence_length: int = 256, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - 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`): - 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 all the text-encoders. - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. - 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. - 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, SD3LoraLoaderMixin): - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - 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] - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - prompt_3 = prompt_3 or prompt - prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3 - - prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=clip_skip, - clip_model_index=0, - ) - prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds( - prompt=prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=clip_skip, - clip_model_index=1, - ) - clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1) - - t5_prompt_embed = self._get_t5_prompt_embeds( - prompt=prompt_3, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - clip_prompt_embeds = torch.nn.functional.pad( - clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1]) - ) - - prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2) - pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1) - - if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt_2 = negative_prompt_2 or negative_prompt - negative_prompt_3 = negative_prompt_3 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 - ) - negative_prompt_3 = ( - batch_size * [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3 - ) - - 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_embed, negative_pooled_prompt_embed = self._get_clip_prompt_embeds( - negative_prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=None, - clip_model_index=0, - ) - negative_prompt_2_embed, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds( - negative_prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=None, - clip_model_index=1, - ) - negative_clip_prompt_embeds = torch.cat([negative_prompt_embed, negative_prompt_2_embed], dim=-1) - - t5_negative_prompt_embed = self._get_t5_prompt_embeds( - prompt=negative_prompt_3, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - negative_clip_prompt_embeds = torch.nn.functional.pad( - negative_clip_prompt_embeds, - (0, t5_negative_prompt_embed.shape[-1] - negative_clip_prompt_embeds.shape[-1]), - ) - - negative_prompt_embeds = torch.cat([negative_clip_prompt_embeds, t5_negative_prompt_embed], dim=-2) - negative_pooled_prompt_embeds = torch.cat( - [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1 - ) - - if self.text_encoder is not None: - if isinstance(self, SD3LoraLoaderMixin) 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, SD3LoraLoaderMixin) 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 - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.check_inputs - def check_inputs( - self, - prompt, - prompt_2, - prompt_3, - height, - width, - negative_prompt=None, - negative_prompt_2=None, - negative_prompt_3=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - ): - if ( - height % (self.vae_scale_factor * self.patch_size) != 0 - or width % (self.vae_scale_factor * self.patch_size) != 0 - ): - raise ValueError( - f"`height` and `width` have to be divisible by {self.vae_scale_factor * self.patch_size} but are {height} and {width}." - f"You can use height {height - height % (self.vae_scale_factor * self.patch_size)} and width {width - width % (self.vae_scale_factor * self.patch_size)}." - ) - - 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_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_3 is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt_3`: {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)}") - elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)): - raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}") - - 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." - ) - elif negative_prompt_3 is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} 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`." - ) - - 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}") - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_latents - 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=device, dtype=dtype) - - 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 prepare_image_with_mask( - self, - image, - mask, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - if isinstance(image, torch.Tensor): - pass - else: - image = self.image_processor.preprocess(image, height=height, width=width) - - image_batch_size = image.shape[0] - - # Prepare image - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - # Prepare mask - if isinstance(mask, torch.Tensor): - pass - else: - mask = self.mask_processor.preprocess(mask, height=height, width=width) - mask = mask.repeat_interleave(repeat_by, dim=0) - mask = mask.to(device=device, dtype=dtype) - - # Get masked image - masked_image = image.clone() - masked_image[(mask > 0.5).repeat(1, 3, 1, 1)] = -1 - - # Encode to latents - image_latents = self.vae.encode(masked_image).latent_dist.sample() - image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor - image_latents = image_latents.to(dtype) - - mask = torch.nn.functional.interpolate( - mask, size=(height // self.vae_scale_factor, width // self.vae_scale_factor) - ) - mask = 1 - mask - control_image = torch.cat([image_latents, mask], dim=1) - - if do_classifier_free_guidance and not guess_mode: - control_image = torch.cat([control_image] * 2) - - return control_image - - @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 - - @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 - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_image - def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch.Tensor: - """Encodes the given image into a feature representation using a pre-trained image encoder. - - Args: - image (`PipelineImageInput`): - Input image to be encoded. - device: (`torch.device`): - Torch device. - - Returns: - `torch.Tensor`: The encoded image feature representation. - """ - if not isinstance(image, torch.Tensor): - image = self.feature_extractor(image, return_tensors="pt").pixel_values - - image = image.to(device=device, dtype=self.dtype) - - return self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - - # 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, - num_images_per_prompt: int = 1, - do_classifier_free_guidance: bool = True, - ) -> torch.Tensor: - """Prepares image embeddings for use in the IP-Adapter. - - Either `ip_adapter_image` or `ip_adapter_image_embeds` must be passed. - - Args: - ip_adapter_image (`PipelineImageInput`, *optional*): - The input image to extract features from for IP-Adapter. - ip_adapter_image_embeds (`torch.Tensor`, *optional*): - Precomputed image embeddings. - device: (`torch.device`, *optional*): - Torch device. - num_images_per_prompt (`int`, defaults to 1): - Number of images that should be generated per prompt. - do_classifier_free_guidance (`bool`, defaults to True): - Whether to use classifier free guidance or not. - """ - device = device or self._execution_device - - if ip_adapter_image_embeds is not None: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = ip_adapter_image_embeds.chunk(2) - else: - single_image_embeds = ip_adapter_image_embeds - elif ip_adapter_image is not None: - single_image_embeds = self.encode_image(ip_adapter_image, device) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.zeros_like(single_image_embeds) - else: - raise ValueError("Neither `ip_adapter_image_embeds` or `ip_adapter_image_embeds` were provided.") - - image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - negative_image_embeds = torch.cat([single_negative_image_embeds] * num_images_per_prompt, dim=0) - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - - return image_embeds.to(device=device) - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.enable_sequential_cpu_offload - def enable_sequential_cpu_offload(self, *args, **kwargs): - if self.image_encoder is not None and "image_encoder" not in self._exclude_from_cpu_offload: - logger.warning( - "`pipe.enable_sequential_cpu_offload()` might fail for `image_encoder` if it uses " - "`torch.nn.MultiheadAttention`. You can exclude `image_encoder` from CPU offloading by calling " - "`pipe._exclude_from_cpu_offload.append('image_encoder')` before `pipe.enable_sequential_cpu_offload()`." - ) - - super().enable_sequential_cpu_offload(*args, **kwargs) - - @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, - prompt_3: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 28, - sigmas: Optional[List[float]] = 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_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", - 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"], - max_sequence_length: int = 256, - ): - 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 `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - will be used instead - 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): - 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. - 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. - 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): - 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]`): - `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]`): - `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): - 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*): - 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 instead - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. - ip_adapter_image (`PipelineImageInput`, *optional*): - Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`torch.Tensor`, *optional*): - Pre-generated image embeddings for IP-Adapter. 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. - 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 256): Maximum sequence length to use with the `prompt`. - - 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. - """ - - height = height or self.default_sample_size * self.vae_scale_factor - width = width or self.default_sample_size * self.vae_scale_factor - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(self.controlnet.nets) if isinstance(self.controlnet, SD3MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - prompt_3, - height, - width, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._joint_attention_kwargs = joint_attention_kwargs - 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 - dtype = self.transformer.dtype - - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_3=prompt_3, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - do_classifier_free_guidance=self.do_classifier_free_guidance, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - device=device, - clip_skip=self.clip_skip, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - ) - - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) - - # 3. Prepare control image - if isinstance(self.controlnet, SD3ControlNetModel): - 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, - ) - latent_height, latent_width = control_image.shape[-2:] - - height = latent_height * self.vae_scale_factor - width = latent_width * self.vae_scale_factor - - elif isinstance(self.controlnet, SD3MultiControlNetModel): - raise NotImplementedError("MultiControlNetModel is not supported for SD3ControlNetInpaintingPipeline.") - else: - assert False - - if controlnet_pooled_projections is None: - controlnet_pooled_projections = torch.zeros_like(pooled_prompt_embeds) - else: - 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) - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - # 5. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(self.controlnet, SD3ControlNetModel) else keeps) - - # 7. Prepare image embeddings - if (ip_adapter_image is not None and self.is_ip_adapter_active) or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - - if self.joint_attention_kwargs is None: - self._joint_attention_kwargs = {"ip_adapter_image_embeds": ip_adapter_image_embeds} - else: - self._joint_attention_kwargs.update(ip_adapter_image_embeds=ip_adapter_image_embeds) - - # 8. Denoising loop - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - # controlnet(s) inference - control_block_samples = self.controlnet( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - pooled_projections=controlnet_pooled_projections, - joint_attention_kwargs=self.joint_attention_kwargs, - controlnet_cond=control_image, - conditioning_scale=cond_scale, - return_dict=False, - )[0] - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - pooled_projections=pooled_prompt_embeds, - block_controlnet_hidden_states=control_block_samples, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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) - - # 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) - negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_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: - latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor - latents = latents.to(dtype=self.vae.dtype) - - 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 StableDiffusion3PipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/__init__.py deleted file mode 100644 index 978278b184f985a452f9d518a1d0eb4f271c74fd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/__init__.py +++ /dev/null @@ -1,68 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - get_objects_from_module, - is_flax_available, - 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_controlnet_xs"] = ["StableDiffusionControlNetXSPipeline"] - _import_structure["pipeline_controlnet_xs_sd_xl"] = ["StableDiffusionXLControlNetXSPipeline"] -try: - if not (is_transformers_available() and is_flax_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_flax_and_transformers_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_flax_and_transformers_objects)) -else: - pass # _import_structure["pipeline_flax_controlnet"] = ["FlaxStableDiffusionControlNetPipeline"] - - -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_controlnet_xs import StableDiffusionControlNetXSPipeline - from .pipeline_controlnet_xs_sd_xl import StableDiffusionXLControlNetXSPipeline - - try: - if not (is_transformers_available() and is_flax_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_flax_and_transformers_objects import * # noqa F403 - else: - pass # from .pipeline_flax_controlnet import FlaxStableDiffusionControlNetPipeline - - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index ea0256cbaeceb3f73707389c5a721bdf3271600e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/__pycache__/pipeline_controlnet_xs.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/__pycache__/pipeline_controlnet_xs.cpython-310.pyc deleted file mode 100644 index 87696278e1d536663f094457d381caac76a66568..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/__pycache__/pipeline_controlnet_xs.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/__pycache__/pipeline_controlnet_xs_sd_xl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/__pycache__/pipeline_controlnet_xs_sd_xl.cpython-310.pyc deleted file mode 100644 index cfb7f5ebbb61851be5025fe8ccb4e9f134ec8a9a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/__pycache__/pipeline_controlnet_xs_sd_xl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py deleted file mode 100644 index 1545027a280cad61c6d2e5b024ea9fb5da3385c1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py +++ /dev/null @@ -1,929 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ControlNetXSAdapter, UNet2DConditionModel, UNetControlNetXSModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import empty_device_cache, is_compiled_module, is_torch_version, randn_tensor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion.pipeline_output import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -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 - >>> # !pip install opencv-python transformers accelerate - >>> from diffusers import StableDiffusionControlNetXSPipeline, ControlNetXSAdapter - >>> from diffusers.utils import load_image - >>> import numpy as np - >>> import torch - - >>> import cv2 - >>> from PIL import Image - - >>> prompt = "aerial view, a futuristic research complex in a bright foggy jungle, hard lighting" - >>> negative_prompt = "low quality, bad quality, sketches" - - >>> # download an image - >>> image = load_image( - ... "https://hf.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/hf-logo.png" - ... ) - - >>> # initialize the models and pipeline - >>> controlnet_conditioning_scale = 0.5 - - >>> controlnet = ControlNetXSAdapter.from_pretrained( - ... "UmerHA/Testing-ConrolNetXS-SD2.1-canny", torch_dtype=torch.float16 - ... ) - >>> pipe = StableDiffusionControlNetXSPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-2-1-base", controlnet=controlnet, torch_dtype=torch.float16 - ... ) - >>> pipe.enable_model_cpu_offload() - - >>> # get canny image - >>> image = np.array(image) - >>> image = cv2.Canny(image, 100, 200) - >>> image = image[:, :, None] - >>> image = np.concatenate([image, image, image], axis=2) - >>> canny_image = Image.fromarray(image) - >>> # generate image - >>> image = pipe( - ... prompt, controlnet_conditioning_scale=controlnet_conditioning_scale, image=canny_image - ... ).images[0] - ``` -""" - - -class StableDiffusionControlNetXSPipeline( - DeprecatedPipelineMixin, - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion with ControlNet-XS guidance. - - 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.). - - 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 - - [`loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A [`UNet2DConditionModel`] used to create a UNetControlNetXSModel to denoise the encoded image latents. - controlnet ([`ControlNetXSAdapter`]): - A [`ControlNetXSAdapter`] to be used in combination with `unet` 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to 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"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetControlNetXSModel], - controlnet: ControlNetXSAdapter, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - if isinstance(unet, UNet2DConditionModel): - unet = UNetControlNetXSModel.from_unet(unet, controlnet) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - safety_checker=safety_checker, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - image, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - controlnet_conditioning_scale=1.0, - control_guidance_start=0.0, - control_guidance_end=1.0, - callback_on_step_end_tensor_inputs=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 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}." - ) - - # Check `image` and `controlnet_conditioning_scale` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.unet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.unet, UNetControlNetXSModel) - or is_compiled - and isinstance(self.unet._orig_mod, UNetControlNetXSModel) - ): - self.check_image(image, prompt, prompt_embeds) - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - else: - assert False - - start, end = 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.") - - 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 prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - ): - image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance: - image = torch.cat([image] * 2) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @property - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.guidance_scale - def guidance_scale(self): - return self._guidance_scale - - @property - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.clip_skip - def clip_skip(self): - return self._clip_skip - - @property - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.do_classifier_free_guidance - def do_classifier_free_guidance(self): - return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None - - @property - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.cross_attention_kwargs - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.num_timesteps - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - image: PipelineImageInput = 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, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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"], - ): - r""" - The call function to the pipeline for generation. - - Args: - 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]]`): - 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. - 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): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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): - 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): - The percentage of total steps at which the ControlNet starts applying. - 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 - the output of the pre-final layer will be used for computing the prompt embeddings. - 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. - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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 isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - unet = self.unet._orig_mod if is_compiled_module(self.unet) else self.unet - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - image, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - 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 - # 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - 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, - lora_scale=text_encoder_lora_scale, - 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 image - image = self.prepare_image( - image=image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=unet.dtype, - do_classifier_free_guidance=do_classifier_free_guidance, - ) - height, width = image.shape[-2:] - - # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 6. Prepare latent variables - num_channels_latents = self.unet.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - self._num_timesteps = len(timesteps) - is_controlnet_compiled = is_compiled_module(self.unet) - is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1") - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # Relevant thread: - # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428 - if torch.cuda.is_available() and is_controlnet_compiled and is_torch_higher_equal_2_1: - torch._inductor.cudagraph_mark_step_begin() - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - apply_control = ( - i / len(timesteps) >= control_guidance_start and (i + 1) / len(timesteps) <= control_guidance_end - ) - noise_pred = self.unet( - sample=latent_model_input, - timestep=t, - encoder_hidden_states=prompt_embeds, - controlnet_cond=image, - conditioning_scale=controlnet_conditioning_scale, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=True, - apply_control=apply_control, - ).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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 we do sequential model offloading, let's offload unet and controlnet - # manually for max memory savings - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.unet.to("cpu") - self.controlnet.to("cpu") - empty_device_cache() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py deleted file mode 100644 index c10931a0f44a5605f89aa1ee5886e166dbfa5301..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py +++ /dev/null @@ -1,1130 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, -) - -from diffusers.utils.import_utils import is_invisible_watermark_available - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, StableDiffusionXLLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ControlNetXSAdapter, UNet2DConditionModel, UNetControlNetXSModel -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import is_compiled_module, is_torch_version, randn_tensor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline -from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput - - -if is_invisible_watermark_available(): - from ..stable_diffusion_xl.watermark import StableDiffusionXLWatermarker - - -from ...utils import is_torch_xla_available - - -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 - >>> # !pip install opencv-python transformers accelerate - >>> from diffusers import StableDiffusionXLControlNetXSPipeline, ControlNetXSAdapter, AutoencoderKL - >>> from diffusers.utils import load_image - >>> import numpy as np - >>> import torch - - >>> import cv2 - >>> from PIL import Image - - >>> prompt = "aerial view, a futuristic research complex in a bright foggy jungle, hard lighting" - >>> negative_prompt = "low quality, bad quality, sketches" - - >>> # download an image - >>> image = load_image( - ... "https://hf.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/hf-logo.png" - ... ) - - >>> # initialize the models and pipeline - >>> controlnet_conditioning_scale = 0.5 - >>> vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) - >>> controlnet = ControlNetXSAdapter.from_pretrained( - ... "UmerHA/Testing-ConrolNetXS-SDXL-canny", torch_dtype=torch.float16 - ... ) - >>> pipe = StableDiffusionXLControlNetXSPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", controlnet=controlnet, torch_dtype=torch.float16 - ... ) - >>> pipe.enable_model_cpu_offload() - - >>> # get canny image - >>> image = np.array(image) - >>> image = cv2.Canny(image, 100, 200) - >>> image = image[:, :, None] - >>> image = np.concatenate([image, image, image], axis=2) - >>> canny_image = Image.fromarray(image) - - >>> # generate image - >>> image = pipe( - ... prompt, controlnet_conditioning_scale=controlnet_conditioning_scale, image=canny_image - ... ).images[0] - ``` -""" - - -class StableDiffusionXLControlNetXSPipeline( - DeprecatedPipelineMixin, - DiffusionPipeline, - TextualInversionLoaderMixin, - StableDiffusionXLLoraLoaderMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion XL with ControlNet-XS guidance. - - 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.). - - The pipeline also inherits the following loading methods: - - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings - - [`loaders.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] for loading LoRA weights - - [`loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - text_encoder_2 ([`~transformers.CLIPTextModelWithProjection`]): - Second frozen text-encoder - ([laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - tokenizer_2 ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A [`UNet2DConditionModel`] used to create a UNetControlNetXSModel to denoise the encoded image latents. - controlnet ([`ControlNetXSAdapter`]): - A [`ControlNetXSAdapter`] to be used in combination with `unet` 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 should always be set to 0. Also see the config of - `stabilityai/stable-diffusion-xl-base-1-0`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark](https://github.com/ShieldMnt/invisible-watermark/) library to - watermark output images. If not defined, it defaults to `True` if the package is installed; otherwise no - watermarker is used. - """ - - _last_supported_version = "0.33.1" - model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "negative_add_time_ids", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetControlNetXSModel], - controlnet: ControlNetXSAdapter, - scheduler: KarrasDiffusionSchedulers, - force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - feature_extractor: CLIPImageProcessor = None, - ): - super().__init__() - - if isinstance(unet, UNet2DConditionModel): - unet = UNetControlNetXSModel.from_unet(unet, controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - prompt_2, - image, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - controlnet_conditioning_scale=1.0, - control_guidance_start=0.0, - control_guidance_end=1.0, - callback_on_step_end_tensor_inputs=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_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`." - ) - - # Check `image` and ``controlnet_conditioning_scale`` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.unet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.unet, UNetControlNetXSModel) - or is_compiled - and isinstance(self.unet._orig_mod, UNetControlNetXSModel) - ): - self.check_image(image, prompt, prompt_embeds) - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - else: - assert False - - start, end = 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.") - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.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 prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - ): - image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance: - image = torch.cat([image] * 2) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - 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.base_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.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - @property - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.guidance_scale - def guidance_scale(self): - return self._guidance_scale - - @property - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.clip_skip - def clip_skip(self): - return self._clip_skip - - @property - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.do_classifier_free_guidance - def do_classifier_free_guidance(self): - return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None - - @property - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.cross_attention_kwargs - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.num_timesteps - def num_timesteps(self): - return self._num_timesteps - - @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, - image: PipelineImageInput = 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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"], - ): - r""" - The call function to the pipeline for generation. - - Args: - 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*): - 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]]`): - 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. - height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The height in pixels of the generated image. 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. 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): - 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*): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - pooled_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated pooled text embeddings. Can be used to easily tweak text inputs (prompt weighting). If - not provided, pooled text embeddings are 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 (prompt - weighting). If not provided, pooled `negative_prompt_embeds` are 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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): - 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)): - 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` 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 `(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)): - 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. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] is - returned, otherwise a `tuple` is returned containing the output images. - """ - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - unet = self.unet._orig_mod if is_compiled_module(self.unet) else self.unet - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - image, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - 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 - # 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt, - prompt_2, - device, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt, - 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=text_encoder_lora_scale, - clip_skip=clip_skip, - ) - - # 4. Prepare image - if isinstance(unet, UNetControlNetXSModel): - image = self.prepare_image( - image=image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=unet.dtype, - do_classifier_free_guidance=do_classifier_free_guidance, - ) - height, width = image.shape[-2:] - else: - assert False - - # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 6. Prepare latent variables - num_channels_latents = self.unet.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Prepare added time ids & embeddings - if isinstance(image, list): - original_size = original_size or image[0].shape[-2:] - else: - original_size = original_size or image.shape[-2:] - target_size = target_size or (height, width) - - 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 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) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - self._num_timesteps = len(timesteps) - is_controlnet_compiled = is_compiled_module(self.unet) - is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1") - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # Relevant thread: - # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428 - if is_controlnet_compiled and is_torch_higher_equal_2_1: - torch._inductor.cudagraph_mark_step_begin() - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - - # predict the noise residual - apply_control = ( - i / len(timesteps) >= control_guidance_start and (i + 1) / len(timesteps) <= control_guidance_end - ) - noise_pred = self.unet( - sample=latent_model_input, - timestep=t, - encoder_hidden_states=prompt_embeds, - controlnet_cond=image, - conditioning_scale=controlnet_conditioning_scale, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=True, - apply_control=apply_control, - ).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - - # 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() - - # manually for max memory savings - if self.vae.dtype == torch.float16 and self.vae.config.force_upcast: - self.upcast_vae() - latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype) - - 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": - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__init__.py deleted file mode 100644 index 2833c89abd5e9b7f842217d3eb5ef8ee45015641..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__init__.py +++ /dev/null @@ -1,54 +0,0 @@ -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_cosmos2_text2image"] = ["Cosmos2TextToImagePipeline"] - _import_structure["pipeline_cosmos2_video2world"] = ["Cosmos2VideoToWorldPipeline"] - _import_structure["pipeline_cosmos_text2world"] = ["CosmosTextToWorldPipeline"] - _import_structure["pipeline_cosmos_video2world"] = ["CosmosVideoToWorldPipeline"] - -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_cosmos2_text2image import Cosmos2TextToImagePipeline - from .pipeline_cosmos2_video2world import Cosmos2VideoToWorldPipeline - from .pipeline_cosmos_text2world import CosmosTextToWorldPipeline - from .pipeline_cosmos_video2world import CosmosVideoToWorldPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 8091cc46a73fe61d26e2dc8089a9af8393378e91..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos2_text2image.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos2_text2image.cpython-310.pyc deleted file mode 100644 index b825c5f64ecd6a073867eabf0530ebb5f12c3a3c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos2_text2image.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos2_video2world.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos2_video2world.cpython-310.pyc deleted file mode 100644 index 87b8e73626fa4a28858c481a21701dfde4dcac3d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos2_video2world.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos_text2world.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos_text2world.cpython-310.pyc deleted file mode 100644 index 91a0b0d4803f1a8072597fa3d90fe0fc1e39d503..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos_text2world.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos_video2world.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos_video2world.cpython-310.pyc deleted file mode 100644 index 76fd0e0af53b681ce4b4cbc71f9f14e9231f691a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_cosmos_video2world.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index da21c2dc9a8c6b127eba69722874ea45bbdf2359..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py deleted file mode 100644 index 66490c2be159df37659244cef88e2e5458214db5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py +++ /dev/null @@ -1,673 +0,0 @@ -# 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. - -import inspect -from typing import Callable, Dict, List, Optional, Union - -import numpy as np -import torch -from transformers import T5EncoderModel, T5TokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...models import AutoencoderKLWan, CosmosTransformer3DModel -from ...schedulers import FlowMatchEulerDiscreteScheduler -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 CosmosImagePipelineOutput - - -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 - - -EXAMPLE_DOC_STRING = """ - Examples: - ```python - >>> import torch - >>> from diffusers import Cosmos2TextToImagePipeline - - >>> # Available checkpoints: nvidia/Cosmos-Predict2-2B-Text2Image, nvidia/Cosmos-Predict2-14B-Text2Image - >>> model_id = "nvidia/Cosmos-Predict2-2B-Text2Image" - >>> pipe = Cosmos2TextToImagePipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) - >>> pipe.to("cuda") - - >>> prompt = "A close-up shot captures a vibrant yellow scrubber vigorously working on a grimy plate, its bristles moving in circular motions to lift stubborn grease and food residue. The dish, once covered in remnants of a hearty meal, gradually reveals its original glossy surface. Suds form and bubble around the scrubber, creating a satisfying visual of cleanliness in progress. The sound of scrubbing fills the air, accompanied by the gentle clinking of the dish against the sink. As the scrubber continues its task, the dish transforms, gleaming under the bright kitchen lights, symbolizing the triumph of cleanliness over mess." - >>> 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." - - >>> output = pipe( - ... prompt=prompt, negative_prompt=negative_prompt, generator=torch.Generator().manual_seed(1) - ... ).images[0] - >>> output.save("output.png") - ``` -""" - - -# 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 Cosmos2TextToImagePipeline(DiffusionPipeline): - r""" - Pipeline for text-to-image generation using [Cosmos Predict2](https://github.com/nvidia-cosmos/cosmos-predict2). - - 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 ([`T5EncoderModel`]): - Frozen text-encoder. Cosmos uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the - [t5-11b](https://huggingface.co/google-t5/t5-11b) variant. - tokenizer (`T5TokenizerFast`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - transformer ([`CosmosTransformer3DModel`]): - Conditional Transformer to denoise the encoded image latents. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - 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"] - # We mark safety_checker as optional here to get around some test failures, but it is not really optional - _optional_components = ["safety_checker"] - - def __init__( - self, - text_encoder: T5EncoderModel, - tokenizer: T5TokenizerFast, - transformer: CosmosTransformer3DModel, - vae: AutoencoderKLWan, - scheduler: FlowMatchEulerDiscreteScheduler, - safety_checker: 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, - 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) - - self.sigma_max = 80.0 - self.sigma_min = 0.002 - self.sigma_data = 1.0 - self.final_sigmas_type = "sigma_min" - if self.scheduler is not None: - self.scheduler.register_to_config( - sigma_max=self.sigma_max, - sigma_min=self.sigma_min, - sigma_data=self.sigma_data, - final_sigmas_type=self.final_sigmas_type, - ) - - # 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, - 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 - - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - return_length=True, - return_offsets_mapping=False, - ) - text_input_ids = text_inputs.input_ids - prompt_attention_mask = text_inputs.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), attention_mask=prompt_attention_mask - ).last_hidden_state - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - - lengths = prompt_attention_mask.sum(dim=1).cpu() - for i, length in enumerate(lengths): - prompt_embeds[i, length:] = 0 - - return 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, - 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, - max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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_images_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, 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_images_per_prompt, 1) - 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 = 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, 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_images_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - - return prompt_embeds, negative_prompt_embeds - - def prepare_latents( - self, - batch_size: int, - num_channels_latents: 16, - 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, - ) -> torch.Tensor: - if latents is not None: - return latents.to(device=device, dtype=dtype) * self.scheduler.config.sigma_max - - 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 - shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_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) - return latents * self.scheduler.config.sigma_max - - # Copied 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, - ): - 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 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)}") - - @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, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = 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", - 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, - ): - 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, one has to pass `prompt_embeds`. - instead. - height (`int`, defaults to `768`): - The height in pixels of the generated image. - width (`int`, defaults to `1360`): - The width in pixels of the generated image. - num_inference_steps (`int`, defaults to `35`): - 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 `7.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_images_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, *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 [`CosmosImagePipelineOutput`] 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. - - Examples: - - Returns: - [`~CosmosImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`CosmosImagePipelineOutput`] 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 - - num_frames = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs(prompt, height, width, prompt_embeds, callback_on_step_end_tensor_inputs) - - 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." - ) - self.safety_checker.to("cpu") - - # 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_images_per_prompt=num_images_per_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - device=device, - max_sequence_length=max_sequence_length, - ) - - # 4. Prepare timesteps - sigmas_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64 - sigmas = torch.linspace(0, 1, num_inference_steps, dtype=sigmas_dtype) - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, device=device, sigmas=sigmas) - if self.scheduler.config.get("final_sigmas_type", "zero") == "sigma_min": - # Replace the last sigma (which is zero) with the minimum sigma value - self.scheduler.sigmas[-1] = self.scheduler.sigmas[-2] - - # 5. Prepare latent variables - transformer_dtype = self.transformer.dtype - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - num_frames, - torch.float32, - device, - generator, - latents, - ) - - padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype) - - # 6. Denoising loop - 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 - current_sigma = self.scheduler.sigmas[i] - - current_t = current_sigma / (current_sigma + 1) - c_in = 1 - current_t - c_skip = 1 - current_t - c_out = -current_t - timestep = current_t.expand(latents.shape[0]).to(transformer_dtype) # [B, 1, T, 1, 1] - - latent_model_input = latents * c_in - latent_model_input = latent_model_input.to(transformer_dtype) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - padding_mask=padding_mask, - return_dict=False, - )[0] - noise_pred = (c_skip * latents + c_out * noise_pred.float()).to(transformer_dtype) - - if self.do_classifier_free_guidance: - noise_pred_uncond = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, - padding_mask=padding_mask, - return_dict=False, - )[0] - noise_pred_uncond = (c_skip * latents + c_out * noise_pred_uncond.float()).to(transformer_dtype) - noise_pred = noise_pred + self.guidance_scale * (noise_pred - noise_pred_uncond) - - noise_pred = (latents - noise_pred) / current_sigma - 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, 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() - - self._current_timestep = None - - if not output_type == "latent": - 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 = latents / latents_std / self.scheduler.config.sigma_data + latents_mean - video = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0] - - if 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) - 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) - self.safety_checker.to("cpu") - else: - video = self.video_processor.postprocess_video(video, output_type=output_type) - image = [batch[0] for batch in video] - if isinstance(video, torch.Tensor): - image = torch.stack(image) - elif isinstance(video, np.ndarray): - image = np.stack(image) - else: - image = latents[:, :, 0] - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return CosmosImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py deleted file mode 100644 index 23a74ad00f933006efdaccbdf288e25f346c11c6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py +++ /dev/null @@ -1,792 +0,0 @@ -# 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. - -import inspect -from typing import Callable, Dict, List, Optional, Union - -import numpy as np -import torch -from transformers import T5EncoderModel, T5TokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput -from ...models import AutoencoderKLWan, CosmosTransformer3DModel -from ...schedulers import FlowMatchEulerDiscreteScheduler -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 - - -EXAMPLE_DOC_STRING = """ - Examples: - ```python - >>> import torch - >>> from diffusers import Cosmos2VideoToWorldPipeline - >>> from diffusers.utils import export_to_video, load_image - - >>> # Available checkpoints: nvidia/Cosmos-Predict2-2B-Video2World, nvidia/Cosmos-Predict2-14B-Video2World - >>> model_id = "nvidia/Cosmos-Predict2-2B-Video2World" - >>> pipe = Cosmos2VideoToWorldPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) - >>> pipe.to("cuda") - - >>> prompt = "A close-up shot captures a vibrant yellow scrubber vigorously working on a grimy plate, its bristles moving in circular motions to lift stubborn grease and food residue. The dish, once covered in remnants of a hearty meal, gradually reveals its original glossy surface. Suds form and bubble around the scrubber, creating a satisfying visual of cleanliness in progress. The sound of scrubbing fills the air, accompanied by the gentle clinking of the dish against the sink. As the scrubber continues its task, the dish transforms, gleaming under the bright kitchen lights, symbolizing the triumph of cleanliness over mess." - >>> 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." - >>> image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yellow-scrubber.png" - ... ) - - >>> video = pipe( - ... image=image, prompt=prompt, negative_prompt=negative_prompt, generator=torch.Generator().manual_seed(1) - ... ).frames[0] - >>> export_to_video(video, "output.mp4", fps=16) - ``` -""" - - -# 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 - - -# 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" -): - 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 Cosmos2VideoToWorldPipeline(DiffusionPipeline): - r""" - Pipeline for video-to-world generation using [Cosmos Predict2](https://github.com/nvidia-cosmos/cosmos-predict2). - - 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 ([`T5EncoderModel`]): - Frozen text-encoder. Cosmos uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the - [t5-11b](https://huggingface.co/google-t5/t5-11b) variant. - tokenizer (`T5TokenizerFast`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - transformer ([`CosmosTransformer3DModel`]): - Conditional Transformer to denoise the encoded image latents. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - 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"] - # We mark safety_checker as optional here to get around some test failures, but it is not really optional - _optional_components = ["safety_checker"] - - def __init__( - self, - text_encoder: T5EncoderModel, - tokenizer: T5TokenizerFast, - transformer: CosmosTransformer3DModel, - vae: AutoencoderKLWan, - scheduler: FlowMatchEulerDiscreteScheduler, - safety_checker: 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, - 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) - - self.sigma_max = 80.0 - self.sigma_min = 0.002 - self.sigma_data = 1.0 - self.final_sigmas_type = "sigma_min" - if self.scheduler is not None: - self.scheduler.register_to_config( - sigma_max=self.sigma_max, - sigma_min=self.sigma_min, - sigma_data=self.sigma_data, - final_sigmas_type=self.final_sigmas_type, - ) - - # 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, - 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 - - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - return_length=True, - return_offsets_mapping=False, - ) - text_input_ids = text_inputs.input_ids - prompt_attention_mask = text_inputs.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), attention_mask=prompt_attention_mask - ).last_hidden_state - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - - lengths = prompt_attention_mask.sum(dim=1).cpu() - for i, length in enumerate(lengths): - prompt_embeds[i, length:] = 0 - - return 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, - 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, - max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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, 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_t5_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 - - def prepare_latents( - self, - video: torch.Tensor, - batch_size: int, - num_channels_latents: 16, - height: int = 704, - 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, - ) -> 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." - ) - - num_cond_frames = video.size(2) - if num_cond_frames >= num_frames: - # Take the last `num_frames` frames for conditioning - num_cond_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - video = video[:, :, -num_frames:] - else: - num_cond_latent_frames = (num_cond_frames - 1) // self.vae_scale_factor_temporal + 1 - num_padding_frames = num_frames - num_cond_frames - last_frame = video[:, :, -1:] - padding = last_frame.repeat(1, 1, num_padding_frames, 1, 1) - video = torch.cat([video, padding], dim=2) - - if isinstance(generator, list): - init_latents = [ - retrieve_latents(self.vae.encode(video[i].unsqueeze(0)), generator=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) - - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, 1).to(device, dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(device, dtype) - ) - init_latents = (init_latents - latents_mean) / latents_std * self.scheduler.config.sigma_data - - 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 - shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device=device, dtype=dtype) - - latents = latents * self.scheduler.config.sigma_max - - padding_shape = (batch_size, 1, num_latent_frames, latent_height, latent_width) - ones_padding = latents.new_ones(padding_shape) - zeros_padding = latents.new_zeros(padding_shape) - - cond_indicator = latents.new_zeros(1, 1, latents.size(2), 1, 1) - cond_indicator[:, :, :num_cond_latent_frames] = 1.0 - cond_mask = cond_indicator * ones_padding + (1 - cond_indicator) * zeros_padding - - uncond_indicator = uncond_mask = None - if do_classifier_free_guidance: - uncond_indicator = latents.new_zeros(1, 1, latents.size(2), 1, 1) - uncond_indicator[:, :, :num_cond_latent_frames] = 1.0 - uncond_mask = uncond_indicator * ones_padding + (1 - uncond_indicator) * zeros_padding - - return latents, init_latents, cond_indicator, uncond_indicator, cond_mask, uncond_mask - - # Copied 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, - ): - 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 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)}") - - @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, - image: PipelineImageInput = None, - video: List[PipelineImageInput] = None, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = 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", - 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, - sigma_conditioning: float = 0.0001, - ): - r""" - The call function to the pipeline for generation. - - 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*): - The video to be used as a conditioning input 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`, defaults to `704`): - The height in pixels of the generated image. - width (`int`, defaults to `1280`): - The width in pixels of the generated image. - num_frames (`int`, defaults to `93`): - The number of frames in the generated video. - num_inference_steps (`int`, defaults to `35`): - 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 `7.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`. - fps (`int`, defaults to `16`): - 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*): - 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, *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. - sigma_conditioning (`float`, defaults to `0.0001`): - The sigma value used for scaling conditioning latents. Ideally, it should not be changed or should be - set to a small value close to zero. - - 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 - - # 1. Check inputs. Raise error if not correct - self.check_inputs(prompt, height, width, prompt_embeds, callback_on_step_end_tensor_inputs) - - 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." - ) - self.safety_checker.to("cpu") - - # 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, - device=device, - max_sequence_length=max_sequence_length, - ) - - # 4. Prepare timesteps - sigmas_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64 - sigmas = torch.linspace(0, 1, num_inference_steps, dtype=sigmas_dtype) - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, device=device, sigmas=sigmas) - if self.scheduler.config.final_sigmas_type == "sigma_min": - # Replace the last sigma (which is zero) with the minimum sigma value - self.scheduler.sigmas[-1] = self.scheduler.sigmas[-2] - - # 5. Prepare latent variables - vae_dtype = self.vae.dtype - transformer_dtype = self.transformer.dtype - - if image is not None: - video = self.video_processor.preprocess(image, height, width).unsqueeze(2) - else: - video = self.video_processor.preprocess_video(video, height, width) - video = video.to(device=device, dtype=vae_dtype) - - num_channels_latents = self.transformer.config.in_channels - 1 - latents, conditioning_latents, cond_indicator, uncond_indicator, cond_mask, uncond_mask = self.prepare_latents( - video, - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - num_frames, - self.do_classifier_free_guidance, - torch.float32, - device, - generator, - latents, - ) - unconditioning_latents = None - - cond_mask = cond_mask.to(transformer_dtype) - if self.do_classifier_free_guidance: - uncond_mask = uncond_mask.to(transformer_dtype) - unconditioning_latents = conditioning_latents - - padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype) - sigma_conditioning = torch.tensor(sigma_conditioning, dtype=torch.float32, device=device) - t_conditioning = sigma_conditioning / (sigma_conditioning + 1) - - # 6. Denoising loop - 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 - current_sigma = self.scheduler.sigmas[i] - - current_t = current_sigma / (current_sigma + 1) - c_in = 1 - current_t - c_skip = 1 - current_t - c_out = -current_t - timestep = current_t.view(1, 1, 1, 1, 1).expand( - latents.size(0), -1, latents.size(2), -1, -1 - ) # [B, 1, T, 1, 1] - - cond_latent = latents * c_in - cond_latent = cond_indicator * conditioning_latents + (1 - cond_indicator) * cond_latent - cond_latent = cond_latent.to(transformer_dtype) - cond_timestep = cond_indicator * t_conditioning + (1 - cond_indicator) * timestep - cond_timestep = cond_timestep.to(transformer_dtype) - - noise_pred = self.transformer( - hidden_states=cond_latent, - timestep=cond_timestep, - encoder_hidden_states=prompt_embeds, - fps=fps, - condition_mask=cond_mask, - padding_mask=padding_mask, - return_dict=False, - )[0] - noise_pred = (c_skip * latents + c_out * noise_pred.float()).to(transformer_dtype) - noise_pred = cond_indicator * conditioning_latents + (1 - cond_indicator) * noise_pred - - if self.do_classifier_free_guidance: - uncond_latent = latents * c_in - uncond_latent = uncond_indicator * unconditioning_latents + (1 - uncond_indicator) * uncond_latent - uncond_latent = uncond_latent.to(transformer_dtype) - uncond_timestep = uncond_indicator * t_conditioning + (1 - uncond_indicator) * timestep - uncond_timestep = uncond_timestep.to(transformer_dtype) - - noise_pred_uncond = self.transformer( - hidden_states=uncond_latent, - timestep=uncond_timestep, - encoder_hidden_states=negative_prompt_embeds, - fps=fps, - condition_mask=uncond_mask, - padding_mask=padding_mask, - return_dict=False, - )[0] - noise_pred_uncond = (c_skip * latents + c_out * noise_pred_uncond.float()).to(transformer_dtype) - noise_pred_uncond = ( - uncond_indicator * unconditioning_latents + (1 - uncond_indicator) * noise_pred_uncond - ) - noise_pred = noise_pred + self.guidance_scale * (noise_pred - noise_pred_uncond) - - noise_pred = (latents - noise_pred) / current_sigma - 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, 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() - - self._current_timestep = None - - if not output_type == "latent": - 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 = ( - torch.tensor(self.vae.config.latents_std) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.scheduler.config.sigma_data + latents_mean - video = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0] - - if 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) - 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) - self.safety_checker.to("cpu") - else: - video = self.video_processor.postprocess_video(video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return CosmosPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py deleted file mode 100644 index f0aa1ecf0e0fd01f32dc8550c54924cf181c4195..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py +++ /dev/null @@ -1,664 +0,0 @@ -# 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. - -import inspect -from typing import Callable, Dict, List, Optional, Union - -import numpy as np -import torch -from transformers import T5EncoderModel, T5TokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...models import AutoencoderKLCosmos, CosmosTransformer3DModel -from ...schedulers import EDMEulerScheduler -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 - - -EXAMPLE_DOC_STRING = """ - Examples: - ```python - >>> import torch - >>> from diffusers import CosmosTextToWorldPipeline - >>> from diffusers.utils import export_to_video - - >>> model_id = "nvidia/Cosmos-1.0-Diffusion-7B-Text2World" - >>> pipe = CosmosTextToWorldPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) - >>> pipe.to("cuda") - - >>> prompt = "A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves. The robot's metallic body gleams under the bright, even lighting, highlighting its futuristic design and intricate joints. A glowing blue light emanates from its chest, adding a touch of advanced technology. The background is dominated by rows of boxes, suggesting a highly organized storage system. The floor is lined with wooden pallets, enhancing the industrial setting. The camera remains static, capturing the robot's poised stance amidst the orderly environment, with a shallow depth of field that keeps the focus on the robot while subtly blurring the background for a cinematic effect." - - >>> output = pipe(prompt=prompt).frames[0] - >>> export_to_video(output, "output.mp4", fps=30) - ``` -""" - - -# 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 CosmosTextToWorldPipeline(DiffusionPipeline): - r""" - Pipeline for text-to-world generation using [Cosmos Predict1](https://github.com/nvidia-cosmos/cosmos-predict1). - - 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 ([`T5EncoderModel`]): - Frozen text-encoder. Cosmos uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the - [t5-11b](https://huggingface.co/google-t5/t5-11b) variant. - tokenizer (`T5TokenizerFast`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - transformer ([`CosmosTransformer3DModel`]): - Conditional Transformer to denoise the encoded image latents. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - vae ([`AutoencoderKLCosmos`]): - 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"] - # We mark safety_checker as optional here to get around some test failures, but it is not really optional - _optional_components = ["safety_checker"] - - def __init__( - self, - text_encoder: T5EncoderModel, - tokenizer: T5TokenizerFast, - transformer: CosmosTransformer3DModel, - vae: AutoencoderKLCosmos, - scheduler: EDMEulerScheduler, - safety_checker: 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, - scheduler=scheduler, - safety_checker=safety_checker, - ) - - self.vae_scale_factor_temporal = ( - self.vae.config.temporal_compression_ratio if getattr(self, "vae", None) else 8 - ) - self.vae_scale_factor_spatial = self.vae.config.spatial_compression_ratio 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: Union[str, List[str]] = None, - 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 - - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - return_length=True, - return_offsets_mapping=False, - ) - text_input_ids = text_inputs.input_ids - prompt_attention_mask = text_inputs.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), attention_mask=prompt_attention_mask - ).last_hidden_state - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - - lengths = prompt_attention_mask.sum(dim=1).cpu() - for i, length in enumerate(lengths): - prompt_embeds[i, length:] = 0 - - return prompt_embeds - - 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: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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, 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_t5_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 - - def prepare_latents( - self, - batch_size: int, - num_channels_latents: 16, - 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, - ) -> torch.Tensor: - if latents is not None: - return latents.to(device=device, dtype=dtype) * self.scheduler.config.sigma_max - - 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 - shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_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) - return latents * self.scheduler.config.sigma_max - - def check_inputs( - self, - prompt, - height, - width, - prompt_embeds=None, - callback_on_step_end_tensor_inputs=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 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)}") - - @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, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = 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", - 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, - ): - 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, one has to pass `prompt_embeds`. - instead. - height (`int`, defaults to `720`): - The height in pixels of the generated image. - width (`int`, defaults to `1280`): - The width in pixels of the generated image. - num_frames (`int`, defaults to `121`): - The number of frames in the generated video. - 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 `7.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`. - fps (`int`, defaults to `30`): - 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*): - 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, *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. - - 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 - - # 1. Check inputs. Raise error if not correct - self.check_inputs(prompt, height, width, prompt_embeds, callback_on_step_end_tensor_inputs) - - 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." - ) - self.safety_checker.to("cpu") - - # 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, - device=device, - max_sequence_length=max_sequence_length, - ) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device) - - # 5. Prepare latent variables - transformer_dtype = self.transformer.dtype - 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, - torch.float32, - device, - generator, - latents, - ) - - padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype) - - # 6. Denoising loop - 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]).to(transformer_dtype) - - latent_model_input = latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - latent_model_input = latent_model_input.to(transformer_dtype) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - fps=fps, - padding_mask=padding_mask, - return_dict=False, - )[0] - - sample = latents - if self.do_classifier_free_guidance: - noise_pred_uncond = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, - fps=fps, - padding_mask=padding_mask, - return_dict=False, - )[0] - noise_pred = torch.cat([noise_pred_uncond, noise_pred]) - sample = torch.cat([sample, sample]) - - # pred_original_sample (x0) - noise_pred = self.scheduler.step(noise_pred, t, sample, return_dict=False)[1] - self.scheduler._step_index -= 1 - - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2) - noise_pred = noise_pred_cond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond) - - # pred_sample (eps) - latents = self.scheduler.step( - noise_pred, t, latents, return_dict=False, pred_original_sample=noise_pred - )[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) - - # 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 not output_type == "latent": - if self.vae.config.latents_mean is not None: - latents_mean, latents_std = self.vae.config.latents_mean, self.vae.config.latents_std - latents_mean = ( - torch.tensor(latents_mean) - .view(1, self.vae.config.latent_channels, -1, 1, 1)[:, :, : latents.size(2)] - .to(latents) - ) - latents_std = ( - torch.tensor(latents_std) - .view(1, self.vae.config.latent_channels, -1, 1, 1)[:, :, : latents.size(2)] - .to(latents) - ) - latents = latents * latents_std / self.scheduler.config.sigma_data + latents_mean - else: - latents = latents / self.scheduler.config.sigma_data - video = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0] - - if 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) - 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) - self.safety_checker.to("cpu") - else: - video = self.video_processor.postprocess_video(video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return CosmosPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py deleted file mode 100644 index cd5a734cc311ca9ceb335b11dcd0b6f015d0c262..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py +++ /dev/null @@ -1,826 +0,0 @@ -# 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. - -import inspect -from typing import Callable, Dict, List, Optional, Union - -import numpy as np -import torch -from transformers import T5EncoderModel, T5TokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput -from ...models import AutoencoderKLCosmos, CosmosTransformer3DModel -from ...schedulers import EDMEulerScheduler -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 - - -EXAMPLE_DOC_STRING = """ - Examples: - Image conditioning: - - ```python - >>> import torch - >>> from diffusers import CosmosVideoToWorldPipeline - >>> from diffusers.utils import export_to_video, load_image - - >>> model_id = "nvidia/Cosmos-1.0-Diffusion-7B-Video2World" - >>> pipe = CosmosVideoToWorldPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) - >>> pipe.to("cuda") - - >>> 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." - >>> image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cosmos/cosmos-video2world-input.jpg" - ... ) - - >>> video = pipe(image=image, prompt=prompt).frames[0] - >>> export_to_video(video, "output.mp4", fps=30) - ``` - - Video conditioning: - - ```python - >>> import torch - >>> from diffusers import CosmosVideoToWorldPipeline - >>> from diffusers.utils import export_to_video, load_video - - >>> model_id = "nvidia/Cosmos-1.0-Diffusion-7B-Video2World" - >>> pipe = CosmosVideoToWorldPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) - >>> pipe.transformer = torch.compile(pipe.transformer) - >>> pipe.to("cuda") - - >>> prompt = "The video depicts 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." - >>> video = load_video( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cosmos/cosmos-video2world-input-vid.mp4" - ... )[ - ... :21 - ... ] # This example uses only the first 21 frames - - >>> video = pipe(video=video, prompt=prompt).frames[0] - >>> export_to_video(video, "output.mp4", fps=30) - ``` -""" - - -# 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 - - -# 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" -): - 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 CosmosVideoToWorldPipeline(DiffusionPipeline): - r""" - Pipeline for image-to-world and video-to-world generation using [Cosmos - Predict-1](https://github.com/nvidia-cosmos/cosmos-predict1). - - 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 ([`T5EncoderModel`]): - Frozen text-encoder. Cosmos uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the - [t5-11b](https://huggingface.co/google-t5/t5-11b) variant. - tokenizer (`T5TokenizerFast`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - transformer ([`CosmosTransformer3DModel`]): - Conditional Transformer to denoise the encoded image latents. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - vae ([`AutoencoderKLCosmos`]): - 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"] - # We mark safety_checker as optional here to get around some test failures, but it is not really optional - _optional_components = ["safety_checker"] - - def __init__( - self, - text_encoder: T5EncoderModel, - tokenizer: T5TokenizerFast, - transformer: CosmosTransformer3DModel, - vae: AutoencoderKLCosmos, - scheduler: EDMEulerScheduler, - safety_checker: 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, - scheduler=scheduler, - safety_checker=safety_checker, - ) - - self.vae_scale_factor_temporal = ( - self.vae.config.temporal_compression_ratio if getattr(self, "vae", None) else 8 - ) - self.vae_scale_factor_spatial = self.vae.config.spatial_compression_ratio if getattr(self, "vae", None) else 8 - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - - # 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, - 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 - - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - return_length=True, - return_offsets_mapping=False, - ) - text_input_ids = text_inputs.input_ids - prompt_attention_mask = text_inputs.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), attention_mask=prompt_attention_mask - ).last_hidden_state - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - - lengths = prompt_attention_mask.sum(dim=1).cpu() - for i, length in enumerate(lengths): - prompt_embeds[i, length:] = 0 - - return 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, - 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, - max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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, 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_t5_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 - - def prepare_latents( - self, - video: torch.Tensor, - batch_size: int, - num_channels_latents: 16, - height: int = 704, - width: int = 1280, - 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, - ) -> 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." - ) - - num_cond_frames = video.size(2) - if num_cond_frames >= num_frames: - # Take the last `num_frames` frames for conditioning - num_cond_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - video = video[:, :, -num_frames:] - else: - num_cond_latent_frames = (num_cond_frames - 1) // self.vae_scale_factor_temporal + 1 - num_padding_frames = num_frames - num_cond_frames - padding = video.new_zeros(video.size(0), video.size(1), num_padding_frames, video.size(3), video.size(4)) - video = torch.cat([video, padding], dim=2) - - if isinstance(generator, list): - init_latents = [ - retrieve_latents(self.vae.encode(video[i].unsqueeze(0)), generator=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) - - if self.vae.config.latents_mean is not None: - latents_mean, latents_std = self.vae.config.latents_mean, self.vae.config.latents_std - latents_mean = ( - torch.tensor(latents_mean) - .view(1, self.vae.config.latent_channels, -1, 1, 1)[:, :, : init_latents.size(2)] - .to(init_latents) - ) - latents_std = ( - torch.tensor(latents_std) - .view(1, self.vae.config.latent_channels, -1, 1, 1)[:, :, : init_latents.size(2)] - .to(init_latents) - ) - init_latents = (init_latents - latents_mean) * self.scheduler.config.sigma_data / latents_std - else: - init_latents = init_latents * self.scheduler.config.sigma_data - - 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 - shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device=device, dtype=dtype) - - latents = latents * self.scheduler.config.sigma_max - - padding_shape = (batch_size, 1, num_latent_frames, latent_height, latent_width) - ones_padding = latents.new_ones(padding_shape) - zeros_padding = latents.new_zeros(padding_shape) - - cond_indicator = latents.new_zeros(1, 1, latents.size(2), 1, 1) - cond_indicator[:, :, :num_cond_latent_frames] = 1.0 - cond_mask = cond_indicator * ones_padding + (1 - cond_indicator) * zeros_padding - - uncond_indicator = uncond_mask = None - if do_classifier_free_guidance: - uncond_indicator = latents.new_zeros(1, 1, latents.size(2), 1, 1) - uncond_indicator[:, :, :num_cond_latent_frames] = 1.0 - uncond_mask = zeros_padding - if not input_frames_guidance: - uncond_mask = uncond_indicator * ones_padding + (1 - uncond_indicator) * zeros_padding - - return latents, init_latents, cond_indicator, uncond_indicator, cond_mask, uncond_mask - - def check_inputs( - self, - prompt, - height, - width, - 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 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 image is None and video is None: - raise ValueError("Either `image` or `video` has to be provided.") - if image is not None and video is not None: - raise ValueError("Only one of `image` or `video` has to be provided.") - - @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, - image: PipelineImageInput = None, - video: List[PipelineImageInput] = None, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: int = 704, - width: int = 1280, - num_frames: int = 121, - num_inference_steps: int = 36, - guidance_scale: float = 7.0, - 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", - 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, - ): - 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, one has to pass `prompt_embeds`. - instead. - height (`int`, defaults to `720`): - The height in pixels of the generated image. - width (`int`, defaults to `1280`): - The width in pixels of the generated image. - num_frames (`int`, defaults to `121`): - The number of frames in the generated video. - 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 `7.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`. - fps (`int`, defaults to `30`): - 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*): - 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, *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. - - 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 - - # 1. Check inputs. Raise error if not correct - self.check_inputs(prompt, height, width, prompt_embeds, callback_on_step_end_tensor_inputs, image, video) - - 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." - ) - self.safety_checker.to("cpu") - - # 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, - device=device, - max_sequence_length=max_sequence_length, - ) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device) - - # 5. Prepare latent variables - vae_dtype = self.vae.dtype - transformer_dtype = self.transformer.dtype - - if image is not None: - video = self.video_processor.preprocess(image, height, width).unsqueeze(2) - else: - video = self.video_processor.preprocess_video(video, height, width) - video = video.to(device=device, dtype=vae_dtype) - - num_channels_latents = self.transformer.config.in_channels - 1 - latents, conditioning_latents, cond_indicator, uncond_indicator, cond_mask, uncond_mask = self.prepare_latents( - video, - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - num_frames, - self.do_classifier_free_guidance, - input_frames_guidance, - torch.float32, - device, - generator, - latents, - ) - cond_mask = cond_mask.to(transformer_dtype) - if self.do_classifier_free_guidance: - uncond_mask = uncond_mask.to(transformer_dtype) - - augment_sigma = torch.tensor([augment_sigma], device=device, dtype=torch.float32) - padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype) - - # 6. Denoising loop - 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]).to(transformer_dtype) - - current_sigma = self.scheduler.sigmas[i] - is_augment_sigma_greater = augment_sigma >= current_sigma - - c_in_augment = self.scheduler._get_conditioning_c_in(augment_sigma) - c_in_original = self.scheduler._get_conditioning_c_in(current_sigma) - - current_cond_indicator = cond_indicator * 0 if is_augment_sigma_greater else cond_indicator - cond_noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=torch.float32) - cond_latent = conditioning_latents + cond_noise * augment_sigma[:, None, None, None, None] - cond_latent = cond_latent * c_in_augment / c_in_original - cond_latent = current_cond_indicator * cond_latent + (1 - current_cond_indicator) * latents - cond_latent = self.scheduler.scale_model_input(cond_latent, t) - cond_latent = cond_latent.to(transformer_dtype) - - noise_pred = self.transformer( - hidden_states=cond_latent, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - fps=fps, - condition_mask=cond_mask, - padding_mask=padding_mask, - return_dict=False, - )[0] - - sample = latents - if self.do_classifier_free_guidance: - current_uncond_indicator = uncond_indicator * 0 if is_augment_sigma_greater else uncond_indicator - uncond_noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=torch.float32) - uncond_latent = conditioning_latents + uncond_noise * augment_sigma[:, None, None, None, None] - uncond_latent = uncond_latent * c_in_augment / c_in_original - uncond_latent = current_uncond_indicator * uncond_latent + (1 - current_uncond_indicator) * latents - uncond_latent = self.scheduler.scale_model_input(uncond_latent, t) - uncond_latent = uncond_latent.to(transformer_dtype) - - noise_pred_uncond = self.transformer( - hidden_states=uncond_latent, - timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, - fps=fps, - condition_mask=uncond_mask, - padding_mask=padding_mask, - return_dict=False, - )[0] - noise_pred = torch.cat([noise_pred_uncond, noise_pred]) - sample = torch.cat([sample, sample]) - - # pred_original_sample (x0) - noise_pred = self.scheduler.step(noise_pred, t, sample, return_dict=False)[1] - self.scheduler._step_index -= 1 - - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2, dim=0) - noise_pred_uncond = ( - current_uncond_indicator * conditioning_latents - + (1 - current_uncond_indicator) * noise_pred_uncond - ) - noise_pred_cond = ( - current_cond_indicator * conditioning_latents + (1 - current_cond_indicator) * noise_pred_cond - ) - noise_pred = noise_pred_cond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond) - else: - noise_pred = ( - current_cond_indicator * conditioning_latents + (1 - current_cond_indicator) * noise_pred - ) - - # pred_sample (eps) - latents = self.scheduler.step( - noise_pred, t, latents, return_dict=False, pred_original_sample=noise_pred - )[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) - - # 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 not output_type == "latent": - if self.vae.config.latents_mean is not None: - latents_mean, latents_std = self.vae.config.latents_mean, self.vae.config.latents_std - latents_mean = ( - torch.tensor(latents_mean) - .view(1, self.vae.config.latent_channels, -1, 1, 1)[:, :, : latents.size(2)] - .to(latents) - ) - latents_std = ( - torch.tensor(latents_std) - .view(1, self.vae.config.latent_channels, -1, 1, 1)[:, :, : latents.size(2)] - .to(latents) - ) - latents = latents * latents_std / self.scheduler.config.sigma_data + latents_mean - else: - latents = latents / self.scheduler.config.sigma_data - video = self.vae.decode(latents.to(vae_dtype), return_dict=False)[0] - - if 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) - 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) - self.safety_checker.to("cpu") - else: - video = self.video_processor.postprocess_video(video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return CosmosPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_output.py deleted file mode 100644 index ec5f4826f62a5aee258c6ec7380e9311083fce20..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/cosmos/pipeline_output.py +++ /dev/null @@ -1,40 +0,0 @@ -from dataclasses import dataclass -from typing import List, Union - -import numpy as np -import PIL.Image -import torch - -from diffusers.utils import BaseOutput, get_logger - - -logger = get_logger(__name__) - - -@dataclass -class CosmosPipelineOutput(BaseOutput): - r""" - 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 - 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 - - -@dataclass -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, - num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/dance_diffusion/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/dance_diffusion/__init__.py deleted file mode 100644 index 0d3e466dfa65b2e9890451607959ed45d092cae7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/dance_diffusion/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import DIFFUSERS_SLOW_IMPORT, _LazyModule - - -_import_structure = {"pipeline_dance_diffusion": ["DanceDiffusionPipeline"]} - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - from .pipeline_dance_diffusion import DanceDiffusionPipeline -else: - import sys - - sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/dance_diffusion/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/dance_diffusion/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index d5b46e9c9bfcc5d9c2f1209db79a5128effe4272..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/dance_diffusion/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/dance_diffusion/__pycache__/pipeline_dance_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/dance_diffusion/__pycache__/pipeline_dance_diffusion.cpython-310.pyc deleted file mode 100644 index f9c0169838f9b33a1eeb8b7237744cd43fec5260..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/dance_diffusion/__pycache__/pipeline_dance_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py deleted file mode 100644 index 5a70c4f5ff9a00f71aff2c2f60fda8f377493281..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py +++ /dev/null @@ -1,169 +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 List, Optional, Tuple, Union - -import torch - -from ...models import UNet1DModel -from ...schedulers import SchedulerMixin -from ...utils import is_torch_xla_available, logging -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import AudioPipelineOutput, DeprecatedPipelineMixin, 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 - - -class DanceDiffusionPipeline(DeprecatedPipelineMixin, DiffusionPipeline): - r""" - Pipeline for audio generation. - - 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.). - - Parameters: - unet ([`UNet1DModel`]): - A `UNet1DModel` to denoise the encoded audio. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded audio latents. Can be one of - [`IPNDMScheduler`]. - """ - - _last_supported_version = "0.33.1" - model_cpu_offload_seq = "unet" - - def __init__(self, unet: UNet1DModel, scheduler: SchedulerMixin): - super().__init__() - self.register_modules(unet=unet, scheduler=scheduler) - - @torch.no_grad() - 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, - return_dict: bool = True, - ) -> Union[AudioPipelineOutput, Tuple]: - r""" - The call function to the pipeline for generation. - - Args: - batch_size (`int`, *optional*, defaults to 1): - The number of audio samples to generate. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher-quality audio sample at - the expense of slower inference. - generator (`torch.Generator`, *optional*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - audio_length_in_s (`float`, *optional*, defaults to `self.unet.config.sample_size/self.unet.config.sample_rate`): - The length of the generated audio sample in seconds. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.AudioPipelineOutput`] instead of a plain tuple. - - Example: - - ```py - from diffusers import DiffusionPipeline - from scipy.io.wavfile import write - - model_id = "harmonai/maestro-150k" - pipe = DiffusionPipeline.from_pretrained(model_id) - pipe = pipe.to("cuda") - - audios = pipe(audio_length_in_s=4.0).audios - - # To save locally - for i, audio in enumerate(audios): - write(f"maestro_test_{i}.wav", pipe.unet.sample_rate, audio.transpose()) - - # To display in google colab - import IPython.display as ipd - - for audio in audios: - display(ipd.Audio(audio, rate=pipe.unet.sample_rate)) - ``` - - Returns: - [`~pipelines.AudioPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.AudioPipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated audio. - """ - - if audio_length_in_s is None: - audio_length_in_s = self.unet.config.sample_size / self.unet.config.sample_rate - - sample_size = audio_length_in_s * self.unet.config.sample_rate - - down_scale_factor = 2 ** len(self.unet.up_blocks) - if sample_size < 3 * down_scale_factor: - raise ValueError( - f"{audio_length_in_s} is too small. Make sure it's bigger or equal to" - f" {3 * down_scale_factor / self.unet.config.sample_rate}." - ) - - original_sample_size = int(sample_size) - if sample_size % down_scale_factor != 0: - sample_size = ( - (audio_length_in_s * self.unet.config.sample_rate) // down_scale_factor + 1 - ) * down_scale_factor - logger.info( - f"{audio_length_in_s} is increased to {sample_size / self.unet.config.sample_rate} so that it can be handled" - f" by the model. It will be cut to {original_sample_size / self.unet.config.sample_rate} after the denoising" - " process." - ) - sample_size = int(sample_size) - - dtype = next(self.unet.parameters()).dtype - shape = (batch_size, self.unet.config.in_channels, sample_size) - 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." - ) - - audio = randn_tensor(shape, generator=generator, device=self._execution_device, dtype=dtype) - - # set step values - self.scheduler.set_timesteps(num_inference_steps, device=audio.device) - self.scheduler.timesteps = self.scheduler.timesteps.to(dtype) - - for t in self.progress_bar(self.scheduler.timesteps): - # 1. predict noise model_output - model_output = self.unet(audio, t).sample - - # 2. compute previous audio sample: x_t -> t_t-1 - audio = self.scheduler.step(model_output, t, audio).prev_sample - - if XLA_AVAILABLE: - xm.mark_step() - - audio = audio.clamp(-1, 1).float().cpu().numpy() - - audio = audio[:, :, :original_sample_size] - - if not return_dict: - return (audio,) - - return AudioPipelineOutput(audios=audio) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddim/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ddim/__init__.py deleted file mode 100644 index d9eede47c897370a23c47c05291690881c987025..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddim/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import DIFFUSERS_SLOW_IMPORT, _LazyModule - - -_import_structure = {"pipeline_ddim": ["DDIMPipeline"]} - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - from .pipeline_ddim import DDIMPipeline -else: - import sys - - sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddim/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ddim/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index e0df62df4084b2eda80cbe5a7fa28252a9cfc037..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddim/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddim/__pycache__/pipeline_ddim.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ddim/__pycache__/pipeline_ddim.cpython-310.pyc deleted file mode 100644 index 8f098af2ff353944fe6fc8303cd2a6707da2c587..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddim/__pycache__/pipeline_ddim.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddim/pipeline_ddim.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ddim/pipeline_ddim.py deleted file mode 100644 index 39587ca5221df400cf7f90391785629490fea079..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddim/pipeline_ddim.py +++ /dev/null @@ -1,167 +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 List, Optional, Tuple, Union - -import torch - -from ...models import UNet2DModel -from ...schedulers import DDIMScheduler -from ...utils import is_torch_xla_available -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - - XLA_AVAILABLE = True -else: - XLA_AVAILABLE = False - - -class DDIMPipeline(DiffusionPipeline): - r""" - Pipeline for image generation. - - 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.). - - Parameters: - unet ([`UNet2DModel`]): - A `UNet2DModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image. Can be one of - [`DDPMScheduler`], or [`DDIMScheduler`]. - """ - - model_cpu_offload_seq = "unet" - - def __init__(self, unet: UNet2DModel, scheduler: DDIMScheduler): - super().__init__() - - # make sure scheduler can always be converted to DDIM - scheduler = DDIMScheduler.from_config(scheduler.config) - - self.register_modules(unet=unet, scheduler=scheduler) - - @torch.no_grad() - def __call__( - self, - batch_size: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - eta: float = 0.0, - num_inference_steps: int = 50, - use_clipped_model_output: Optional[bool] = None, - output_type: Optional[str] = "pil", - return_dict: bool = True, - ) -> Union[ImagePipelineOutput, Tuple]: - r""" - The call function to the pipeline for generation. - - Args: - batch_size (`int`, *optional*, defaults to 1): - The number of images to generate. - generator (`torch.Generator`, *optional*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - 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. A value of `0` - corresponds to DDIM and `1` corresponds to DDPM. - 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. - use_clipped_model_output (`bool`, *optional*, defaults to `None`): - If `True` or `False`, see documentation for [`DDIMScheduler.step`]. If `None`, nothing is passed - downstream to the scheduler (use `None` for schedulers which don't support this 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 [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Example: - - ```py - >>> from diffusers import DDIMPipeline - >>> import PIL.Image - >>> import numpy as np - - >>> # load model and scheduler - >>> pipe = DDIMPipeline.from_pretrained("fusing/ddim-lsun-bedroom") - - >>> # run pipeline in inference (sample random noise and denoise) - >>> image = pipe(eta=0.0, num_inference_steps=50) - - >>> # process image to PIL - >>> image_processed = image.cpu().permute(0, 2, 3, 1) - >>> image_processed = (image_processed + 1.0) * 127.5 - >>> image_processed = image_processed.numpy().astype(np.uint8) - >>> image_pil = PIL.Image.fromarray(image_processed[0]) - - >>> # save image - >>> image_pil.save("test.png") - ``` - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images - """ - - # Sample gaussian noise to begin loop - if isinstance(self.unet.config.sample_size, int): - image_shape = ( - batch_size, - self.unet.config.in_channels, - self.unet.config.sample_size, - self.unet.config.sample_size, - ) - else: - image_shape = (batch_size, self.unet.config.in_channels, *self.unet.config.sample_size) - - 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." - ) - - image = randn_tensor(image_shape, generator=generator, device=self._execution_device, dtype=self.unet.dtype) - - # set step values - self.scheduler.set_timesteps(num_inference_steps) - - for t in self.progress_bar(self.scheduler.timesteps): - # 1. predict noise model_output - model_output = self.unet(image, t).sample - - # 2. predict previous mean of image x_t-1 and add variance depending on eta - # eta corresponds to η in paper and should be between [0, 1] - # do x_t -> x_t-1 - image = self.scheduler.step( - model_output, t, image, eta=eta, use_clipped_model_output=use_clipped_model_output, generator=generator - ).prev_sample - - if XLA_AVAILABLE: - xm.mark_step() - - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).numpy() - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddpm/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ddpm/__init__.py deleted file mode 100644 index eb41dd1dcf642c791f3d7b0d985efcaf3e4a2c22..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddpm/__init__.py +++ /dev/null @@ -1,22 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - _LazyModule, -) - - -_import_structure = {"pipeline_ddpm": ["DDPMPipeline"]} - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - from .pipeline_ddpm import DDPMPipeline - -else: - import sys - - sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddpm/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ddpm/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 141a9da3e710e666f3e68e4146c6782c858cd9c2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddpm/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddpm/__pycache__/pipeline_ddpm.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ddpm/__pycache__/pipeline_ddpm.cpython-310.pyc deleted file mode 100644 index 94c965bc91187c020174a11d6d7856a157fa6a1d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddpm/__pycache__/pipeline_ddpm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddpm/pipeline_ddpm.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ddpm/pipeline_ddpm.py deleted file mode 100644 index 0d7766a8cfd0343ff67f968e9c823d077904e456..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ddpm/pipeline_ddpm.py +++ /dev/null @@ -1,141 +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 List, Optional, Tuple, Union - -import torch - -from ...models import UNet2DModel -from ...schedulers import DDPMScheduler -from ...utils import is_torch_xla_available -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - - XLA_AVAILABLE = True -else: - XLA_AVAILABLE = False - - -class DDPMPipeline(DiffusionPipeline): - r""" - Pipeline for image generation. - - 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.). - - Parameters: - unet ([`UNet2DModel`]): - A `UNet2DModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image. Can be one of - [`DDPMScheduler`], or [`DDIMScheduler`]. - """ - - model_cpu_offload_seq = "unet" - - def __init__(self, unet: UNet2DModel, scheduler: DDPMScheduler): - super().__init__() - self.register_modules(unet=unet, scheduler=scheduler) - - @torch.no_grad() - def __call__( - self, - batch_size: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - num_inference_steps: int = 1000, - output_type: Optional[str] = "pil", - return_dict: bool = True, - ) -> Union[ImagePipelineOutput, Tuple]: - r""" - The call function to the pipeline for generation. - - Args: - batch_size (`int`, *optional*, defaults to 1): - The number of images to generate. - generator (`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 1000): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - 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 [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Example: - - ```py - >>> from diffusers import DDPMPipeline - - >>> # load model and scheduler - >>> pipe = DDPMPipeline.from_pretrained("google/ddpm-cat-256") - - >>> # run pipeline in inference (sample random noise and denoise) - >>> image = pipe().images[0] - - >>> # save image - >>> image.save("ddpm_generated_image.png") - ``` - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images - """ - # Sample gaussian noise to begin loop - if isinstance(self.unet.config.sample_size, int): - image_shape = ( - batch_size, - self.unet.config.in_channels, - self.unet.config.sample_size, - self.unet.config.sample_size, - ) - else: - image_shape = (batch_size, self.unet.config.in_channels, *self.unet.config.sample_size) - - if self.device.type == "mps": - # randn does not work reproducibly on mps - image = randn_tensor(image_shape, generator=generator, dtype=self.unet.dtype) - image = image.to(self.device) - else: - image = randn_tensor(image_shape, generator=generator, device=self.device, dtype=self.unet.dtype) - - # set step values - self.scheduler.set_timesteps(num_inference_steps) - - for t in self.progress_bar(self.scheduler.timesteps): - # 1. predict noise model_output - model_output = self.unet(image, t).sample - - # 2. compute previous image: x_t -> x_t-1 - image = self.scheduler.step(model_output, t, image, generator=generator).prev_sample - - if XLA_AVAILABLE: - xm.mark_step() - - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).numpy() - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__init__.py deleted file mode 100644 index 79aab1fb186a857dd0a3353c4b5905b4595b5b7b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__init__.py +++ /dev/null @@ -1,85 +0,0 @@ -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 = { - "timesteps": [ - "fast27_timesteps", - "smart100_timesteps", - "smart185_timesteps", - "smart27_timesteps", - "smart50_timesteps", - "super100_timesteps", - "super27_timesteps", - "super40_timesteps", - ] -} - -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_if"] = ["IFPipeline"] - _import_structure["pipeline_if_img2img"] = ["IFImg2ImgPipeline"] - _import_structure["pipeline_if_img2img_superresolution"] = ["IFImg2ImgSuperResolutionPipeline"] - _import_structure["pipeline_if_inpainting"] = ["IFInpaintingPipeline"] - _import_structure["pipeline_if_inpainting_superresolution"] = ["IFInpaintingSuperResolutionPipeline"] - _import_structure["pipeline_if_superresolution"] = ["IFSuperResolutionPipeline"] - _import_structure["pipeline_output"] = ["IFPipelineOutput"] - _import_structure["safety_checker"] = ["IFSafetyChecker"] - _import_structure["watermark"] = ["IFWatermarker"] - - -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_if import IFPipeline - from .pipeline_if_img2img import IFImg2ImgPipeline - from .pipeline_if_img2img_superresolution import IFImg2ImgSuperResolutionPipeline - from .pipeline_if_inpainting import IFInpaintingPipeline - from .pipeline_if_inpainting_superresolution import IFInpaintingSuperResolutionPipeline - from .pipeline_if_superresolution import IFSuperResolutionPipeline - from .pipeline_output import IFPipelineOutput - from .safety_checker import IFSafetyChecker - from .timesteps import ( - fast27_timesteps, - smart27_timesteps, - smart50_timesteps, - smart100_timesteps, - smart185_timesteps, - super27_timesteps, - super40_timesteps, - super100_timesteps, - ) - from .watermark import IFWatermarker - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 33552a413383f863aa66e8d7cf62e492ba8fecb4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if.cpython-310.pyc deleted file mode 100644 index 9b3bcd16527cec210a424e3413ba7645449cfed6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_img2img.cpython-310.pyc deleted file mode 100644 index b4136c0a158149586dbdc3801e390a969f001906..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_img2img_superresolution.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_img2img_superresolution.cpython-310.pyc deleted file mode 100644 index c87b8ccfd6c8f7f145230d17be58ec68d5a12951..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_img2img_superresolution.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_inpainting.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_inpainting.cpython-310.pyc deleted file mode 100644 index 4438a8930df246b553da8fd02c66939ed30a6fc6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_inpainting.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_inpainting_superresolution.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_inpainting_superresolution.cpython-310.pyc deleted file mode 100644 index 740fcd125d39928b73db30588e5721659f0c9061..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_inpainting_superresolution.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_superresolution.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_superresolution.cpython-310.pyc deleted file mode 100644 index 56389699a8331c12f03759456ee7e940fdf8b3ec..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_if_superresolution.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index f3a57e6eba9222a31301e355a48f0dc94d7a003c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/safety_checker.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/safety_checker.cpython-310.pyc deleted file mode 100644 index 2e18182f296eeaac37e4b77a8320ed80ca064144..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/safety_checker.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/timesteps.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/timesteps.cpython-310.pyc deleted file mode 100644 index a2bc0767b65b2360461b43314364ea5d7cb93f2c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/timesteps.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/watermark.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/watermark.cpython-310.pyc deleted file mode 100644 index 55912c4a6efa59cf15d24adf6220ad7ca7bd8426..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/__pycache__/watermark.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if.py deleted file mode 100644 index 8fa31f8504d3e4be5ed269b0bea7c365289b3f3f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if.py +++ /dev/null @@ -1,786 +0,0 @@ -import html -import inspect -import re -import urllib.parse as ul -from typing import Any, Callable, Dict, List, Optional, Union - -import torch -from transformers import CLIPImageProcessor, T5EncoderModel, T5Tokenizer - -from ...loaders import StableDiffusionLoraLoaderMixin -from ...models import UNet2DConditionModel -from ...schedulers import DDPMScheduler -from ...utils import ( - BACKENDS_MAPPING, - is_bs4_available, - is_ftfy_available, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .pipeline_output import IFPipelineOutput -from .safety_checker import IFSafetyChecker -from .watermark import IFWatermarker - - -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_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> from diffusers import IFPipeline, IFSuperResolutionPipeline, DiffusionPipeline - >>> from diffusers.utils import pt_to_pil - >>> import torch - - >>> pipe = IFPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", variant="fp16", torch_dtype=torch.float16) - >>> pipe.enable_model_cpu_offload() - - >>> prompt = 'a photo of a kangaroo wearing an orange hoodie and blue sunglasses standing in front of the eiffel tower holding a sign that says "very deep learning"' - >>> prompt_embeds, negative_embeds = pipe.encode_prompt(prompt) - - >>> image = pipe(prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, output_type="pt").images - - >>> # save intermediate image - >>> pil_image = pt_to_pil(image) - >>> pil_image[0].save("./if_stage_I.png") - - >>> super_res_1_pipe = IFSuperResolutionPipeline.from_pretrained( - ... "DeepFloyd/IF-II-L-v1.0", text_encoder=None, variant="fp16", torch_dtype=torch.float16 - ... ) - >>> super_res_1_pipe.enable_model_cpu_offload() - - >>> image = super_res_1_pipe( - ... image=image, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, output_type="pt" - ... ).images - - >>> # save intermediate image - >>> pil_image = pt_to_pil(image) - >>> pil_image[0].save("./if_stage_I.png") - - >>> safety_modules = { - ... "feature_extractor": pipe.feature_extractor, - ... "safety_checker": pipe.safety_checker, - ... "watermarker": pipe.watermarker, - ... } - >>> super_res_2_pipe = DiffusionPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-x4-upscaler", **safety_modules, torch_dtype=torch.float16 - ... ) - >>> super_res_2_pipe.enable_model_cpu_offload() - - >>> image = super_res_2_pipe( - ... prompt=prompt, - ... image=image, - ... ).images - >>> image[0].save("./if_stage_II.png") - ``` -""" - - -class IFPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin): - tokenizer: T5Tokenizer - text_encoder: T5EncoderModel - - unet: UNet2DConditionModel - scheduler: DDPMScheduler - - feature_extractor: Optional[CLIPImageProcessor] - safety_checker: Optional[IFSafetyChecker] - - watermarker: Optional[IFWatermarker] - - bad_punct_regex = re.compile( - r"[" - + "#®•©™&@·º½¾¿¡§~" - + r"\)" - + r"\(" - + r"\]" - + r"\[" - + r"\}" - + r"\{" - + r"\|" - + "\\" - + r"\/" - + r"\*" - + r"]{1,}" - ) # noqa - - _optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor", "watermarker"] - model_cpu_offload_seq = "text_encoder->unet" - _exclude_from_cpu_offload = ["watermarker"] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - unet: UNet2DConditionModel, - scheduler: DDPMScheduler, - safety_checker: Optional[IFSafetyChecker], - feature_extractor: Optional[CLIPImageProcessor], - watermarker: Optional[IFWatermarker], - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the IF license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - watermarker=watermarker, - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - @torch.no_grad() - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - ): - 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 - device: (`torch.device`, *optional*): - torch device to place the resulting embeddings on - 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`). - 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. - clean_caption (bool, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - """ - if prompt is not None and negative_prompt is not None: - 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)} !=" - f" {type(prompt)}." - ) - - if device is None: - device = self._execution_device - - 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] - - # while T5 can handle much longer input sequences than 77, the text encoder was trained with a max length of 77 for IF - max_length = 77 - - if prompt_embeds is None: - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - add_special_tokens=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[:, max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {max_length} tokens: {removed_text}" - ) - - attention_mask = text_inputs.attention_mask.to(device) - - prompt_embeds = self.text_encoder( - text_input_ids.to(device), - attention_mask=attention_mask, - ) - prompt_embeds = prompt_embeds[0] - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - elif self.unet is not None: - dtype = self.unet.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=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 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 - - uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption) - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - attention_mask = uncond_input.attention_mask.to(device) - - 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=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) - - # 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 - else: - negative_prompt_embeds = None - - return prompt_embeds, negative_prompt_embeds - - def run_safety_checker(self, image, device, dtype): - if self.safety_checker is not None: - safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device) - image, nsfw_detected, watermark_detected = self.safety_checker( - images=image, - clip_input=safety_checker_input.pixel_values.to(dtype=dtype), - ) - else: - nsfw_detected = None - watermark_detected = None - - return image, nsfw_detected, watermark_detected - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - if (callback_steps is None) or ( - 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 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_intermediate_images(self, batch_size, num_channels, height, width, dtype, device, generator): - shape = (batch_size, num_channels, 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." - ) - - intermediate_images = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # scale the initial noise by the standard deviation required by the scheduler - intermediate_images = intermediate_images * self.scheduler.init_noise_sigma - return intermediate_images - - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - num_inference_steps: int = 100, - 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, - 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - clean_caption: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - ): - """ - 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. - 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*): - 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): - 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. - height (`int`, *optional*, defaults to self.unet.config.sample_size): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to self.unet.config.sample_size): - The width in pixels of the generated image. - 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*): - 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*): - 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.IFPipelineOutput`] 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. - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - cross_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). - - Examples: - - Returns: - [`~pipelines.stable_diffusion.IFPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.IFPipelineOutput`] 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) - or watermarked content, according to the `safety_checker`. - """ - # 1. Check inputs. Raise error if not correct - self.check_inputs(prompt, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds) - - # 2. Define call parameters - height = height or self.unet.config.sample_size - width = width or self.unet.config.sample_size - - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - do_classifier_free_guidance, - num_images_per_prompt=num_images_per_prompt, - device=device, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - clean_caption=clean_caption, - ) - - if do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - # 4. Prepare timesteps - if timesteps is not None: - self.scheduler.set_timesteps(timesteps=timesteps, device=device) - timesteps = self.scheduler.timesteps - num_inference_steps = len(timesteps) - else: - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - if hasattr(self.scheduler, "set_begin_index"): - self.scheduler.set_begin_index(0) - - # 5. Prepare intermediate images - intermediate_images = self.prepare_intermediate_images( - batch_size * num_images_per_prompt, - self.unet.config.in_channels, - height, - width, - prompt_embeds.dtype, - device, - generator, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # HACK: see comment in `enable_model_cpu_offload` - if hasattr(self, "text_encoder_offload_hook") and self.text_encoder_offload_hook is not None: - self.text_encoder_offload_hook.offload() - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - model_input = ( - torch.cat([intermediate_images] * 2) if do_classifier_free_guidance else intermediate_images - ) - model_input = self.scheduler.scale_model_input(model_input, t) - - # predict the noise residual - noise_pred = self.unet( - model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred_uncond, _ = noise_pred_uncond.split(model_input.shape[1], dim=1) - noise_pred_text, predicted_variance = noise_pred_text.split(model_input.shape[1], dim=1) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, predicted_variance], dim=1) - - if self.scheduler.config.variance_type not in ["learned", "learned_range"]: - noise_pred, _ = noise_pred.split(model_input.shape[1], dim=1) - - # compute the previous noisy sample x_t -> x_t-1 - intermediate_images = self.scheduler.step( - noise_pred, t, intermediate_images, **extra_step_kwargs, return_dict=False - )[0] - - # 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 callback is not None and i % callback_steps == 0: - callback(i, t, intermediate_images) - - if XLA_AVAILABLE: - xm.mark_step() - - image = intermediate_images - - if output_type == "pil": - # 8. Post-processing - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - # 9. Run safety checker - image, nsfw_detected, watermark_detected = self.run_safety_checker(image, device, prompt_embeds.dtype) - - # 10. Convert to PIL - image = self.numpy_to_pil(image) - - # 11. Apply watermark - if self.watermarker is not None: - image = self.watermarker.apply_watermark(image, self.unet.config.sample_size) - elif output_type == "pt": - nsfw_detected = None - watermark_detected = None - - if hasattr(self, "unet_offload_hook") and self.unet_offload_hook is not None: - self.unet_offload_hook.offload() - else: - # 8. Post-processing - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - # 9. Run safety checker - image, nsfw_detected, watermark_detected = self.run_safety_checker(image, device, prompt_embeds.dtype) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image, nsfw_detected, watermark_detected) - - return IFPipelineOutput(images=image, nsfw_detected=nsfw_detected, watermark_detected=watermark_detected) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py deleted file mode 100644 index 507927faf61b05c6be63b8978c11c34d45b213a0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py +++ /dev/null @@ -1,907 +0,0 @@ -import html -import inspect -import re -import urllib.parse as ul -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import CLIPImageProcessor, T5EncoderModel, T5Tokenizer - -from ...loaders import StableDiffusionLoraLoaderMixin -from ...models import UNet2DConditionModel -from ...schedulers import DDPMScheduler -from ...utils import ( - BACKENDS_MAPPING, - PIL_INTERPOLATION, - is_bs4_available, - is_ftfy_available, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .pipeline_output import IFPipelineOutput -from .safety_checker import IFSafetyChecker -from .watermark import IFWatermarker - - -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_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - - -def resize(images: PIL.Image.Image, img_size: int) -> PIL.Image.Image: - w, h = images.size - - coef = w / h - - w, h = img_size, img_size - - if coef >= 1: - w = int(round(img_size / 8 * coef) * 8) - else: - h = int(round(img_size / 8 / coef) * 8) - - images = images.resize((w, h), resample=PIL_INTERPOLATION["bicubic"], reducing_gap=None) - - return images - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> from diffusers import IFImg2ImgPipeline, IFImg2ImgSuperResolutionPipeline, DiffusionPipeline - >>> from diffusers.utils import pt_to_pil - >>> import torch - >>> from PIL import Image - >>> import requests - >>> from io import BytesIO - - >>> url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" - >>> response = requests.get(url) - >>> original_image = Image.open(BytesIO(response.content)).convert("RGB") - >>> original_image = original_image.resize((768, 512)) - - >>> pipe = IFImg2ImgPipeline.from_pretrained( - ... "DeepFloyd/IF-I-XL-v1.0", - ... variant="fp16", - ... torch_dtype=torch.float16, - ... ) - >>> pipe.enable_model_cpu_offload() - - >>> prompt = "A fantasy landscape in style minecraft" - >>> prompt_embeds, negative_embeds = pipe.encode_prompt(prompt) - - >>> image = pipe( - ... image=original_image, - ... prompt_embeds=prompt_embeds, - ... negative_prompt_embeds=negative_embeds, - ... output_type="pt", - ... ).images - - >>> # save intermediate image - >>> pil_image = pt_to_pil(image) - >>> pil_image[0].save("./if_stage_I.png") - - >>> super_res_1_pipe = IFImg2ImgSuperResolutionPipeline.from_pretrained( - ... "DeepFloyd/IF-II-L-v1.0", - ... text_encoder=None, - ... variant="fp16", - ... torch_dtype=torch.float16, - ... ) - >>> super_res_1_pipe.enable_model_cpu_offload() - - >>> image = super_res_1_pipe( - ... image=image, - ... original_image=original_image, - ... prompt_embeds=prompt_embeds, - ... negative_prompt_embeds=negative_embeds, - ... ).images - >>> image[0].save("./if_stage_II.png") - ``` -""" - - -class IFImg2ImgPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin): - tokenizer: T5Tokenizer - text_encoder: T5EncoderModel - - unet: UNet2DConditionModel - scheduler: DDPMScheduler - - feature_extractor: Optional[CLIPImageProcessor] - safety_checker: Optional[IFSafetyChecker] - - watermarker: Optional[IFWatermarker] - - bad_punct_regex = re.compile( - r"[" - + "#®•©™&@·º½¾¿¡§~" - + r"\)" - + r"\(" - + r"\]" - + r"\[" - + r"\}" - + r"\{" - + r"\|" - + "\\" - + r"\/" - + r"\*" - + r"]{1,}" - ) # noqa - - _optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor", "watermarker"] - model_cpu_offload_seq = "text_encoder->unet" - _exclude_from_cpu_offload = ["watermarker"] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - unet: UNet2DConditionModel, - scheduler: DDPMScheduler, - safety_checker: Optional[IFSafetyChecker], - feature_extractor: Optional[CLIPImageProcessor], - watermarker: Optional[IFWatermarker], - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the IF license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - watermarker=watermarker, - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - @torch.no_grad() - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - ): - 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 - device: (`torch.device`, *optional*): - torch device to place the resulting embeddings on - 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`). - 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. - clean_caption (bool, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - """ - if prompt is not None and negative_prompt is not None: - 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)} !=" - f" {type(prompt)}." - ) - - if device is None: - device = self._execution_device - - 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] - - # while T5 can handle much longer input sequences than 77, the text encoder was trained with a max length of 77 for IF - max_length = 77 - - if prompt_embeds is None: - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - add_special_tokens=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[:, max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {max_length} tokens: {removed_text}" - ) - - attention_mask = text_inputs.attention_mask.to(device) - - prompt_embeds = self.text_encoder( - text_input_ids.to(device), - attention_mask=attention_mask, - ) - prompt_embeds = prompt_embeds[0] - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - elif self.unet is not None: - dtype = self.unet.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=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 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 - - uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption) - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - attention_mask = uncond_input.attention_mask.to(device) - - 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=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) - - # 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 - else: - negative_prompt_embeds = None - - return prompt_embeds, negative_prompt_embeds - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.run_safety_checker - def run_safety_checker(self, image, device, dtype): - if self.safety_checker is not None: - safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device) - image, nsfw_detected, watermark_detected = self.safety_checker( - images=image, - clip_input=safety_checker_input.pixel_values.to(dtype=dtype), - ) - else: - nsfw_detected = None - watermark_detected = None - - return image, nsfw_detected, watermark_detected - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - image, - batch_size, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - if (callback_steps is None) or ( - 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 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 isinstance(image, list): - check_image_type = image[0] - else: - check_image_type = image - - if ( - not isinstance(check_image_type, torch.Tensor) - and not isinstance(check_image_type, PIL.Image.Image) - 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" - f" {type(check_image_type)}" - ) - - if isinstance(image, list): - image_batch_size = len(image) - elif isinstance(image, torch.Tensor): - image_batch_size = image.shape[0] - elif isinstance(image, PIL.Image.Image): - image_batch_size = 1 - elif isinstance(image, np.ndarray): - image_batch_size = image.shape[0] - else: - assert False - - if batch_size != image_batch_size: - raise ValueError(f"image batch size: {image_batch_size} must be same as prompt batch size {batch_size}") - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - def preprocess_image(self, image: PIL.Image.Image) -> torch.Tensor: - if not isinstance(image, list): - image = [image] - - def numpy_to_pt(images): - if images.ndim == 3: - images = images[..., None] - - images = torch.from_numpy(images.transpose(0, 3, 1, 2)) - return images - - if isinstance(image[0], PIL.Image.Image): - new_image = [] - - for image_ in image: - image_ = image_.convert("RGB") - image_ = resize(image_, self.unet.config.sample_size) - image_ = np.array(image_) - image_ = image_.astype(np.float32) - image_ = image_ / 127.5 - 1 - new_image.append(image_) - - image = new_image - - image = np.stack(image, axis=0) # to np - image = numpy_to_pt(image) # to pt - - elif isinstance(image[0], np.ndarray): - image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0) - image = numpy_to_pt(image) - - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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_intermediate_images( - self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None - ): - _, channels, height, width = image.shape - - batch_size = batch_size * num_images_per_prompt - - shape = (batch_size, channels, 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." - ) - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - image = image.repeat_interleave(num_images_per_prompt, dim=0) - image = self.scheduler.add_noise(image, noise, timestep) - - return image - - @torch.no_grad() - @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, - strength: float = 0.7, - num_inference_steps: int = 80, - timesteps: List[int] = None, - guidance_scale: float = 10.0, - 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, - 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, - clean_caption: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - ): - """ - 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 (`torch.Tensor` or `PIL.Image.Image`): - `Image`, or tensor representing an image batch, that will be used as the starting point for the - process. - strength (`float`, *optional*, defaults to 0.7): - 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_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*): - 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): - 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` 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*): - 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.IFPipelineOutput`] 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. - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - cross_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). - - Examples: - - Returns: - [`~pipelines.stable_diffusion.IFPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.IFPipelineOutput`] 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) - or watermarked content, according to the `safety_checker`. - """ - # 1. Check inputs. Raise error if not correct - 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] - - self.check_inputs( - prompt, image, batch_size, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds - ) - - # 2. Define call parameters - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - do_classifier_free_guidance, - num_images_per_prompt=num_images_per_prompt, - device=device, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - clean_caption=clean_caption, - ) - - if do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - dtype = prompt_embeds.dtype - - # 4. Prepare timesteps - if timesteps is not None: - self.scheduler.set_timesteps(timesteps=timesteps, device=device) - timesteps = self.scheduler.timesteps - num_inference_steps = len(timesteps) - else: - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength) - - # 5. Prepare intermediate images - image = self.preprocess_image(image) - image = image.to(device=device, dtype=dtype) - - noise_timestep = timesteps[0:1] - noise_timestep = noise_timestep.repeat(batch_size * num_images_per_prompt) - - intermediate_images = self.prepare_intermediate_images( - image, noise_timestep, batch_size, num_images_per_prompt, dtype, device, generator - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # HACK: see comment in `enable_model_cpu_offload` - if hasattr(self, "text_encoder_offload_hook") and self.text_encoder_offload_hook is not None: - self.text_encoder_offload_hook.offload() - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - model_input = ( - torch.cat([intermediate_images] * 2) if do_classifier_free_guidance else intermediate_images - ) - model_input = self.scheduler.scale_model_input(model_input, t) - - # predict the noise residual - noise_pred = self.unet( - model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred_uncond, _ = noise_pred_uncond.split(model_input.shape[1], dim=1) - noise_pred_text, predicted_variance = noise_pred_text.split(model_input.shape[1], dim=1) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, predicted_variance], dim=1) - - if self.scheduler.config.variance_type not in ["learned", "learned_range"]: - noise_pred, _ = noise_pred.split(model_input.shape[1], dim=1) - - # compute the previous noisy sample x_t -> x_t-1 - intermediate_images = self.scheduler.step( - noise_pred, t, intermediate_images, **extra_step_kwargs, return_dict=False - )[0] - - # 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 callback is not None and i % callback_steps == 0: - callback(i, t, intermediate_images) - - if XLA_AVAILABLE: - xm.mark_step() - - image = intermediate_images - - if output_type == "pil": - # 8. Post-processing - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - # 9. Run safety checker - image, nsfw_detected, watermark_detected = self.run_safety_checker(image, device, prompt_embeds.dtype) - - # 10. Convert to PIL - image = self.numpy_to_pil(image) - - # 11. Apply watermark - if self.watermarker is not None: - self.watermarker.apply_watermark(image, self.unet.config.sample_size) - elif output_type == "pt": - nsfw_detected = None - watermark_detected = None - - if hasattr(self, "unet_offload_hook") and self.unet_offload_hook is not None: - self.unet_offload_hook.offload() - else: - # 8. Post-processing - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - # 9. Run safety checker - image, nsfw_detected, watermark_detected = self.run_safety_checker(image, device, prompt_embeds.dtype) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image, nsfw_detected, watermark_detected) - - return IFPipelineOutput(images=image, nsfw_detected=nsfw_detected, watermark_detected=watermark_detected) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py deleted file mode 100644 index 9bc15c3c6f62ded55c3ffe90a39935575f1e646f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py +++ /dev/null @@ -1,1024 +0,0 @@ -import html -import inspect -import re -import urllib.parse as ul -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import CLIPImageProcessor, T5EncoderModel, T5Tokenizer - -from ...loaders import StableDiffusionLoraLoaderMixin -from ...models import UNet2DConditionModel -from ...schedulers import DDPMScheduler -from ...utils import ( - BACKENDS_MAPPING, - PIL_INTERPOLATION, - is_bs4_available, - is_ftfy_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .pipeline_output import IFPipelineOutput -from .safety_checker import IFSafetyChecker -from .watermark import IFWatermarker - - -if is_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - - -from ...utils import is_torch_xla_available - - -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 - - -# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if_img2img.resize -def resize(images: PIL.Image.Image, img_size: int) -> PIL.Image.Image: - w, h = images.size - - coef = w / h - - w, h = img_size, img_size - - if coef >= 1: - w = int(round(img_size / 8 * coef) * 8) - else: - h = int(round(img_size / 8 / coef) * 8) - - images = images.resize((w, h), resample=PIL_INTERPOLATION["bicubic"], reducing_gap=None) - - return images - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> from diffusers import IFImg2ImgPipeline, IFImg2ImgSuperResolutionPipeline, DiffusionPipeline - >>> from diffusers.utils import pt_to_pil - >>> import torch - >>> from PIL import Image - >>> import requests - >>> from io import BytesIO - - >>> url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" - >>> response = requests.get(url) - >>> original_image = Image.open(BytesIO(response.content)).convert("RGB") - >>> original_image = original_image.resize((768, 512)) - - >>> pipe = IFImg2ImgPipeline.from_pretrained( - ... "DeepFloyd/IF-I-XL-v1.0", - ... variant="fp16", - ... torch_dtype=torch.float16, - ... ) - >>> pipe.enable_model_cpu_offload() - - >>> prompt = "A fantasy landscape in style minecraft" - >>> prompt_embeds, negative_embeds = pipe.encode_prompt(prompt) - - >>> image = pipe( - ... image=original_image, - ... prompt_embeds=prompt_embeds, - ... negative_prompt_embeds=negative_embeds, - ... output_type="pt", - ... ).images - - >>> # save intermediate image - >>> pil_image = pt_to_pil(image) - >>> pil_image[0].save("./if_stage_I.png") - - >>> super_res_1_pipe = IFImg2ImgSuperResolutionPipeline.from_pretrained( - ... "DeepFloyd/IF-II-L-v1.0", - ... text_encoder=None, - ... variant="fp16", - ... torch_dtype=torch.float16, - ... ) - >>> super_res_1_pipe.enable_model_cpu_offload() - - >>> image = super_res_1_pipe( - ... image=image, - ... original_image=original_image, - ... prompt_embeds=prompt_embeds, - ... negative_prompt_embeds=negative_embeds, - ... ).images - >>> image[0].save("./if_stage_II.png") - ``` -""" - - -class IFImg2ImgSuperResolutionPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin): - tokenizer: T5Tokenizer - text_encoder: T5EncoderModel - - unet: UNet2DConditionModel - scheduler: DDPMScheduler - image_noising_scheduler: DDPMScheduler - - feature_extractor: Optional[CLIPImageProcessor] - safety_checker: Optional[IFSafetyChecker] - - watermarker: Optional[IFWatermarker] - - bad_punct_regex = re.compile( - r"[" - + "#®•©™&@·º½¾¿¡§~" - + r"\)" - + r"\(" - + r"\]" - + r"\[" - + r"\}" - + r"\{" - + r"\|" - + "\\" - + r"\/" - + r"\*" - + r"]{1,}" - ) # noqa - - _optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor"] - model_cpu_offload_seq = "text_encoder->unet" - _exclude_from_cpu_offload = ["watermarker"] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - unet: UNet2DConditionModel, - scheduler: DDPMScheduler, - image_noising_scheduler: DDPMScheduler, - safety_checker: Optional[IFSafetyChecker], - feature_extractor: Optional[CLIPImageProcessor], - watermarker: Optional[IFWatermarker], - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the IF license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - if unet is not None and unet.config.in_channels != 6: - logger.warning( - "It seems like you have loaded a checkpoint that shall not be used for super resolution from {unet.config._name_or_path} as it accepts {unet.config.in_channels} input channels instead of 6. Please make sure to pass a super resolution checkpoint as the `'unet'`: IFSuperResolutionPipeline.from_pretrained(unet=super_resolution_unet, ...)`." - ) - - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - unet=unet, - scheduler=scheduler, - image_noising_scheduler=image_noising_scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - watermarker=watermarker, - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - @torch.no_grad() - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - ): - 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 - device: (`torch.device`, *optional*): - torch device to place the resulting embeddings on - 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`). - 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. - clean_caption (bool, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - """ - if prompt is not None and negative_prompt is not None: - 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)} !=" - f" {type(prompt)}." - ) - - if device is None: - device = self._execution_device - - 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] - - # while T5 can handle much longer input sequences than 77, the text encoder was trained with a max length of 77 for IF - max_length = 77 - - if prompt_embeds is None: - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - add_special_tokens=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[:, max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {max_length} tokens: {removed_text}" - ) - - attention_mask = text_inputs.attention_mask.to(device) - - prompt_embeds = self.text_encoder( - text_input_ids.to(device), - attention_mask=attention_mask, - ) - prompt_embeds = prompt_embeds[0] - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - elif self.unet is not None: - dtype = self.unet.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=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 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 - - uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption) - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - attention_mask = uncond_input.attention_mask.to(device) - - 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=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) - - # 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 - else: - negative_prompt_embeds = None - - return prompt_embeds, negative_prompt_embeds - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.run_safety_checker - def run_safety_checker(self, image, device, dtype): - if self.safety_checker is not None: - safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device) - image, nsfw_detected, watermark_detected = self.safety_checker( - images=image, - clip_input=safety_checker_input.pixel_values.to(dtype=dtype), - ) - else: - nsfw_detected = None - watermark_detected = None - - return image, nsfw_detected, watermark_detected - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - image, - original_image, - batch_size, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - if (callback_steps is None) or ( - 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 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}." - ) - - # image - - if isinstance(image, list): - check_image_type = image[0] - else: - check_image_type = image - - if ( - not isinstance(check_image_type, torch.Tensor) - and not isinstance(check_image_type, PIL.Image.Image) - 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" - f" {type(check_image_type)}" - ) - - if isinstance(image, list): - image_batch_size = len(image) - elif isinstance(image, torch.Tensor): - image_batch_size = image.shape[0] - elif isinstance(image, PIL.Image.Image): - image_batch_size = 1 - elif isinstance(image, np.ndarray): - image_batch_size = image.shape[0] - else: - assert False - - if batch_size != image_batch_size: - raise ValueError(f"image batch size: {image_batch_size} must be same as prompt batch size {batch_size}") - - # original_image - - if isinstance(original_image, list): - check_image_type = original_image[0] - else: - check_image_type = original_image - - if ( - not isinstance(check_image_type, torch.Tensor) - and not isinstance(check_image_type, PIL.Image.Image) - 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" - f" {type(check_image_type)}" - ) - - if isinstance(original_image, list): - image_batch_size = len(original_image) - elif isinstance(original_image, torch.Tensor): - image_batch_size = original_image.shape[0] - elif isinstance(original_image, PIL.Image.Image): - image_batch_size = 1 - elif isinstance(original_image, np.ndarray): - image_batch_size = original_image.shape[0] - else: - assert False - - if batch_size != image_batch_size: - raise ValueError( - f"original_image batch size: {image_batch_size} must be same as prompt batch size {batch_size}" - ) - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if_img2img.IFImg2ImgPipeline.preprocess_image with preprocess_image -> preprocess_original_image - def preprocess_original_image(self, image: PIL.Image.Image) -> torch.Tensor: - if not isinstance(image, list): - image = [image] - - def numpy_to_pt(images): - if images.ndim == 3: - images = images[..., None] - - images = torch.from_numpy(images.transpose(0, 3, 1, 2)) - return images - - if isinstance(image[0], PIL.Image.Image): - new_image = [] - - for image_ in image: - image_ = image_.convert("RGB") - image_ = resize(image_, self.unet.config.sample_size) - image_ = np.array(image_) - image_ = image_.astype(np.float32) - image_ = image_ / 127.5 - 1 - new_image.append(image_) - - image = new_image - - image = np.stack(image, axis=0) # to np - image = numpy_to_pt(image) # to pt - - elif isinstance(image[0], np.ndarray): - image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0) - image = numpy_to_pt(image) - - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0) - - return image - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if_superresolution.IFSuperResolutionPipeline.preprocess_image - def preprocess_image(self, image: PIL.Image.Image, num_images_per_prompt, device) -> torch.Tensor: - if not isinstance(image, torch.Tensor) and not isinstance(image, list): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - image = [np.array(i).astype(np.float32) / 127.5 - 1.0 for i in image] - - image = np.stack(image, axis=0) # to np - image = torch.from_numpy(image.transpose(0, 3, 1, 2)) - elif isinstance(image[0], np.ndarray): - image = np.stack(image, axis=0) # to np - if image.ndim == 5: - image = image[0] - - image = torch.from_numpy(image.transpose(0, 3, 1, 2)) - elif isinstance(image, list) and isinstance(image[0], torch.Tensor): - dims = image[0].ndim - - if dims == 3: - image = torch.stack(image, dim=0) - elif dims == 4: - image = torch.concat(image, dim=0) - else: - raise ValueError(f"Image must have 3 or 4 dimensions, instead got {dims}") - - image = image.to(device=device, dtype=self.unet.dtype) - - image = image.repeat_interleave(num_images_per_prompt, dim=0) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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 - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if_img2img.IFImg2ImgPipeline.prepare_intermediate_images - def prepare_intermediate_images( - self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None - ): - _, channels, height, width = image.shape - - batch_size = batch_size * num_images_per_prompt - - shape = (batch_size, channels, 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." - ) - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - image = image.repeat_interleave(num_images_per_prompt, dim=0) - image = self.scheduler.add_noise(image, noise, timestep) - - return image - - @torch.no_grad() - @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, - strength: float = 0.8, - prompt: Union[str, List[str]] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - guidance_scale: float = 4.0, - 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, - 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, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - noise_level: int = 250, - clean_caption: bool = True, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - image (`torch.Tensor` or `PIL.Image.Image`): - `Image`, or tensor representing an image batch, that will be used as the starting point for the - process. - original_image (`torch.Tensor` or `PIL.Image.Image`): - The original image that `image` was varied from. - strength (`float`, *optional*, defaults to 0.8): - 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`. - 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*): - 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): - 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` 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*): - 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.IFPipelineOutput`] 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. - cross_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). - noise_level (`int`, *optional*, defaults to 250): - The amount of noise to add to the upscaled image. Must be in the range `[0, 1000)` - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.IFPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.IFPipelineOutput`] 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) - or watermarked content, according to the `safety_checker`. - """ - # 1. Check inputs. Raise error if not correct - 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] - - self.check_inputs( - prompt, - image, - original_image, - batch_size, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ) - - # 2. Define call parameters - - # 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 = guidance_scale > 1.0 - - device = self._execution_device - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - do_classifier_free_guidance, - num_images_per_prompt=num_images_per_prompt, - device=device, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - clean_caption=clean_caption, - ) - - if do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - dtype = prompt_embeds.dtype - - # 4. Prepare timesteps - if timesteps is not None: - self.scheduler.set_timesteps(timesteps=timesteps, device=device) - timesteps = self.scheduler.timesteps - num_inference_steps = len(timesteps) - else: - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength) - - # 5. prepare original image - original_image = self.preprocess_original_image(original_image) - original_image = original_image.to(device=device, dtype=dtype) - - # 6. Prepare intermediate images - noise_timestep = timesteps[0:1] - noise_timestep = noise_timestep.repeat(batch_size * num_images_per_prompt) - - intermediate_images = self.prepare_intermediate_images( - original_image, - noise_timestep, - batch_size, - num_images_per_prompt, - dtype, - device, - generator, - ) - - # 7. Prepare upscaled image and noise level - _, _, height, width = original_image.shape - - image = self.preprocess_image(image, num_images_per_prompt, device) - - upscaled = F.interpolate(image, (height, width), mode="bilinear", align_corners=True) - - noise_level = torch.tensor([noise_level] * upscaled.shape[0], device=upscaled.device) - noise = randn_tensor(upscaled.shape, generator=generator, device=upscaled.device, dtype=upscaled.dtype) - upscaled = self.image_noising_scheduler.add_noise(upscaled, noise, timesteps=noise_level) - - if do_classifier_free_guidance: - noise_level = torch.cat([noise_level] * 2) - - # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # HACK: see comment in `enable_model_cpu_offload` - if hasattr(self, "text_encoder_offload_hook") and self.text_encoder_offload_hook is not None: - self.text_encoder_offload_hook.offload() - - # 9. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - model_input = torch.cat([intermediate_images, upscaled], dim=1) - - model_input = torch.cat([model_input] * 2) if do_classifier_free_guidance else model_input - model_input = self.scheduler.scale_model_input(model_input, t) - - # predict the noise residual - noise_pred = self.unet( - model_input, - t, - encoder_hidden_states=prompt_embeds, - class_labels=noise_level, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred_uncond, _ = noise_pred_uncond.split(model_input.shape[1] // 2, dim=1) - noise_pred_text, predicted_variance = noise_pred_text.split(model_input.shape[1] // 2, dim=1) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, predicted_variance], dim=1) - - if self.scheduler.config.variance_type not in ["learned", "learned_range"]: - noise_pred, _ = noise_pred.split(intermediate_images.shape[1], dim=1) - - # compute the previous noisy sample x_t -> x_t-1 - intermediate_images = self.scheduler.step( - noise_pred, t, intermediate_images, **extra_step_kwargs, return_dict=False - )[0] - - # 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 callback is not None and i % callback_steps == 0: - callback(i, t, intermediate_images) - - if XLA_AVAILABLE: - xm.mark_step() - - image = intermediate_images - - if output_type == "pil": - # 10. Post-processing - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - # 11. Run safety checker - image, nsfw_detected, watermark_detected = self.run_safety_checker(image, device, prompt_embeds.dtype) - - # 12. Convert to PIL - image = self.numpy_to_pil(image) - - # 13. Apply watermark - if self.watermarker is not None: - self.watermarker.apply_watermark(image, self.unet.config.sample_size) - elif output_type == "pt": - nsfw_detected = None - watermark_detected = None - - else: - # 10. Post-processing - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - # 11. Run safety checker - image, nsfw_detected, watermark_detected = self.run_safety_checker(image, device, prompt_embeds.dtype) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image, nsfw_detected, watermark_detected) - - return IFPipelineOutput(images=image, nsfw_detected=nsfw_detected, watermark_detected=watermark_detected) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py deleted file mode 100644 index 9d6cf62020a97ee267a55bfb31f72ebf12acbb19..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py +++ /dev/null @@ -1,1026 +0,0 @@ -import html -import inspect -import re -import urllib.parse as ul -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import CLIPImageProcessor, T5EncoderModel, T5Tokenizer - -from ...loaders import StableDiffusionLoraLoaderMixin -from ...models import UNet2DConditionModel -from ...schedulers import DDPMScheduler -from ...utils import ( - BACKENDS_MAPPING, - PIL_INTERPOLATION, - is_bs4_available, - is_ftfy_available, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .pipeline_output import IFPipelineOutput -from .safety_checker import IFSafetyChecker -from .watermark import IFWatermarker - - -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_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - - -# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if_img2img.resize -def resize(images: PIL.Image.Image, img_size: int) -> PIL.Image.Image: - w, h = images.size - - coef = w / h - - w, h = img_size, img_size - - if coef >= 1: - w = int(round(img_size / 8 * coef) * 8) - else: - h = int(round(img_size / 8 / coef) * 8) - - images = images.resize((w, h), resample=PIL_INTERPOLATION["bicubic"], reducing_gap=None) - - return images - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> from diffusers import IFInpaintingPipeline, IFInpaintingSuperResolutionPipeline, DiffusionPipeline - >>> from diffusers.utils import pt_to_pil - >>> import torch - >>> from PIL import Image - >>> import requests - >>> from io import BytesIO - - >>> url = "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/if/person.png" - >>> response = requests.get(url) - >>> original_image = Image.open(BytesIO(response.content)).convert("RGB") - >>> original_image = original_image - - >>> url = "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/if/glasses_mask.png" - >>> response = requests.get(url) - >>> mask_image = Image.open(BytesIO(response.content)) - >>> mask_image = mask_image - - >>> pipe = IFInpaintingPipeline.from_pretrained( - ... "DeepFloyd/IF-I-XL-v1.0", variant="fp16", torch_dtype=torch.float16 - ... ) - >>> pipe.enable_model_cpu_offload() - - >>> prompt = "blue sunglasses" - >>> prompt_embeds, negative_embeds = pipe.encode_prompt(prompt) - - >>> image = pipe( - ... image=original_image, - ... mask_image=mask_image, - ... prompt_embeds=prompt_embeds, - ... negative_prompt_embeds=negative_embeds, - ... output_type="pt", - ... ).images - - >>> # save intermediate image - >>> pil_image = pt_to_pil(image) - >>> pil_image[0].save("./if_stage_I.png") - - >>> super_res_1_pipe = IFInpaintingSuperResolutionPipeline.from_pretrained( - ... "DeepFloyd/IF-II-L-v1.0", text_encoder=None, variant="fp16", torch_dtype=torch.float16 - ... ) - >>> super_res_1_pipe.enable_model_cpu_offload() - - >>> image = super_res_1_pipe( - ... image=image, - ... mask_image=mask_image, - ... original_image=original_image, - ... prompt_embeds=prompt_embeds, - ... negative_prompt_embeds=negative_embeds, - ... ).images - >>> image[0].save("./if_stage_II.png") - ``` -""" - - -class IFInpaintingPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin): - tokenizer: T5Tokenizer - text_encoder: T5EncoderModel - - unet: UNet2DConditionModel - scheduler: DDPMScheduler - - feature_extractor: Optional[CLIPImageProcessor] - safety_checker: Optional[IFSafetyChecker] - - watermarker: Optional[IFWatermarker] - - bad_punct_regex = re.compile( - r"[" - + "#®•©™&@·º½¾¿¡§~" - + r"\)" - + r"\(" - + r"\]" - + r"\[" - + r"\}" - + r"\{" - + r"\|" - + "\\" - + r"\/" - + r"\*" - + r"]{1,}" - ) # noqa - - _optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor", "watermarker"] - model_cpu_offload_seq = "text_encoder->unet" - _exclude_from_cpu_offload = ["watermarker"] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - unet: UNet2DConditionModel, - scheduler: DDPMScheduler, - safety_checker: Optional[IFSafetyChecker], - feature_extractor: Optional[CLIPImageProcessor], - watermarker: Optional[IFWatermarker], - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the IF license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - watermarker=watermarker, - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - @torch.no_grad() - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - ): - 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 - device: (`torch.device`, *optional*): - torch device to place the resulting embeddings on - 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`). - 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. - clean_caption (bool, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - """ - if prompt is not None and negative_prompt is not None: - 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)} !=" - f" {type(prompt)}." - ) - - if device is None: - device = self._execution_device - - 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] - - # while T5 can handle much longer input sequences than 77, the text encoder was trained with a max length of 77 for IF - max_length = 77 - - if prompt_embeds is None: - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - add_special_tokens=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[:, max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {max_length} tokens: {removed_text}" - ) - - attention_mask = text_inputs.attention_mask.to(device) - - prompt_embeds = self.text_encoder( - text_input_ids.to(device), - attention_mask=attention_mask, - ) - prompt_embeds = prompt_embeds[0] - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - elif self.unet is not None: - dtype = self.unet.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=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 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 - - uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption) - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - attention_mask = uncond_input.attention_mask.to(device) - - 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=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) - - # 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 - else: - negative_prompt_embeds = None - - return prompt_embeds, negative_prompt_embeds - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.run_safety_checker - def run_safety_checker(self, image, device, dtype): - if self.safety_checker is not None: - safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device) - image, nsfw_detected, watermark_detected = self.safety_checker( - images=image, - clip_input=safety_checker_input.pixel_values.to(dtype=dtype), - ) - else: - nsfw_detected = None - watermark_detected = None - - return image, nsfw_detected, watermark_detected - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - image, - mask_image, - batch_size, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - if (callback_steps is None) or ( - 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 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}." - ) - - # image - - if isinstance(image, list): - check_image_type = image[0] - else: - check_image_type = image - - if ( - not isinstance(check_image_type, torch.Tensor) - and not isinstance(check_image_type, PIL.Image.Image) - 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" - f" {type(check_image_type)}" - ) - - if isinstance(image, list): - image_batch_size = len(image) - elif isinstance(image, torch.Tensor): - image_batch_size = image.shape[0] - elif isinstance(image, PIL.Image.Image): - image_batch_size = 1 - elif isinstance(image, np.ndarray): - image_batch_size = image.shape[0] - else: - assert False - - if batch_size != image_batch_size: - raise ValueError(f"image batch size: {image_batch_size} must be same as prompt batch size {batch_size}") - - # mask_image - - if isinstance(mask_image, list): - check_image_type = mask_image[0] - else: - check_image_type = mask_image - - if ( - not isinstance(check_image_type, torch.Tensor) - and not isinstance(check_image_type, PIL.Image.Image) - 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" - f" {type(check_image_type)}" - ) - - if isinstance(mask_image, list): - image_batch_size = len(mask_image) - elif isinstance(mask_image, torch.Tensor): - image_batch_size = mask_image.shape[0] - elif isinstance(mask_image, PIL.Image.Image): - image_batch_size = 1 - elif isinstance(mask_image, np.ndarray): - image_batch_size = mask_image.shape[0] - else: - assert False - - if image_batch_size != 1 and batch_size != image_batch_size: - raise ValueError( - f"mask_image batch size: {image_batch_size} must be `1` or the same as prompt batch size {batch_size}" - ) - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if_img2img.IFImg2ImgPipeline.preprocess_image - def preprocess_image(self, image: PIL.Image.Image) -> torch.Tensor: - if not isinstance(image, list): - image = [image] - - def numpy_to_pt(images): - if images.ndim == 3: - images = images[..., None] - - images = torch.from_numpy(images.transpose(0, 3, 1, 2)) - return images - - if isinstance(image[0], PIL.Image.Image): - new_image = [] - - for image_ in image: - image_ = image_.convert("RGB") - image_ = resize(image_, self.unet.config.sample_size) - image_ = np.array(image_) - image_ = image_.astype(np.float32) - image_ = image_ / 127.5 - 1 - new_image.append(image_) - - image = new_image - - image = np.stack(image, axis=0) # to np - image = numpy_to_pt(image) # to pt - - elif isinstance(image[0], np.ndarray): - image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0) - image = numpy_to_pt(image) - - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0) - - return image - - def preprocess_mask_image(self, mask_image) -> torch.Tensor: - if not isinstance(mask_image, list): - mask_image = [mask_image] - - if isinstance(mask_image[0], torch.Tensor): - mask_image = torch.cat(mask_image, axis=0) if mask_image[0].ndim == 4 else torch.stack(mask_image, axis=0) - - if mask_image.ndim == 2: - # Batch and add channel dim for single mask - mask_image = mask_image.unsqueeze(0).unsqueeze(0) - elif mask_image.ndim == 3 and mask_image.shape[0] == 1: - # Single mask, the 0'th dimension is considered to be - # the existing batch size of 1 - mask_image = mask_image.unsqueeze(0) - elif mask_image.ndim == 3 and mask_image.shape[0] != 1: - # Batch of mask, the 0'th dimension is considered to be - # the batching dimension - mask_image = mask_image.unsqueeze(1) - - mask_image[mask_image < 0.5] = 0 - mask_image[mask_image >= 0.5] = 1 - - elif isinstance(mask_image[0], PIL.Image.Image): - new_mask_image = [] - - for mask_image_ in mask_image: - mask_image_ = mask_image_.convert("L") - mask_image_ = resize(mask_image_, self.unet.config.sample_size) - mask_image_ = np.array(mask_image_) - mask_image_ = mask_image_[None, None, :] - new_mask_image.append(mask_image_) - - mask_image = new_mask_image - - mask_image = np.concatenate(mask_image, axis=0) - mask_image = mask_image.astype(np.float32) / 255.0 - mask_image[mask_image < 0.5] = 0 - mask_image[mask_image >= 0.5] = 1 - mask_image = torch.from_numpy(mask_image) - - elif isinstance(mask_image[0], np.ndarray): - mask_image = np.concatenate([m[None, None, :] for m in mask_image], axis=0) - - mask_image[mask_image < 0.5] = 0 - mask_image[mask_image >= 0.5] = 1 - mask_image = torch.from_numpy(mask_image) - - return mask_image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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_intermediate_images( - self, image, timestep, batch_size, num_images_per_prompt, dtype, device, mask_image, generator=None - ): - image_batch_size, channels, height, width = image.shape - - batch_size = batch_size * num_images_per_prompt - - shape = (batch_size, channels, 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." - ) - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - image = image.repeat_interleave(num_images_per_prompt, dim=0) - noised_image = self.scheduler.add_noise(image, noise, timestep) - - image = (1 - mask_image) * image + mask_image * noised_image - - return image - - @torch.no_grad() - @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, - strength: float = 1.0, - num_inference_steps: int = 50, - timesteps: List[int] = None, - guidance_scale: float = 7.0, - 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, - 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, - clean_caption: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - ): - """ - 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 (`torch.Tensor` or `PIL.Image.Image`): - `Image`, or tensor representing an image batch, that will be used as the starting point for the - process. - mask_image (`PIL.Image.Image`): - `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be - 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)`. - strength (`float`, *optional*, defaults to 1.0): - 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_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. If not defined, equal spaced `num_inference_steps` - timesteps are used. Must be in descending order. - guidance_scale (`float`, *optional*, defaults to 7.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. - 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*): - 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*): - 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.IFPipelineOutput`] 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. - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - cross_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). - - Examples: - - Returns: - [`~pipelines.stable_diffusion.IFPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.IFPipelineOutput`] 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) - or watermarked content, according to the `safety_checker`. - """ - # 1. Check inputs. Raise error if not correct - 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] - - self.check_inputs( - prompt, - image, - mask_image, - batch_size, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ) - - # 2. Define call parameters - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - do_classifier_free_guidance, - num_images_per_prompt=num_images_per_prompt, - device=device, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - clean_caption=clean_caption, - ) - - if do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - dtype = prompt_embeds.dtype - - # 4. Prepare timesteps - if timesteps is not None: - self.scheduler.set_timesteps(timesteps=timesteps, device=device) - timesteps = self.scheduler.timesteps - num_inference_steps = len(timesteps) - else: - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength) - - # 5. Prepare intermediate images - image = self.preprocess_image(image) - image = image.to(device=device, dtype=dtype) - - mask_image = self.preprocess_mask_image(mask_image) - mask_image = mask_image.to(device=device, dtype=dtype) - - if mask_image.shape[0] == 1: - mask_image = mask_image.repeat_interleave(batch_size * num_images_per_prompt, dim=0) - else: - mask_image = mask_image.repeat_interleave(num_images_per_prompt, dim=0) - - noise_timestep = timesteps[0:1] - noise_timestep = noise_timestep.repeat(batch_size * num_images_per_prompt) - - intermediate_images = self.prepare_intermediate_images( - image, noise_timestep, batch_size, num_images_per_prompt, dtype, device, mask_image, generator - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # HACK: see comment in `enable_model_cpu_offload` - if hasattr(self, "text_encoder_offload_hook") and self.text_encoder_offload_hook is not None: - self.text_encoder_offload_hook.offload() - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - model_input = ( - torch.cat([intermediate_images] * 2) if do_classifier_free_guidance else intermediate_images - ) - model_input = self.scheduler.scale_model_input(model_input, t) - - # predict the noise residual - noise_pred = self.unet( - model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred_uncond, _ = noise_pred_uncond.split(model_input.shape[1], dim=1) - noise_pred_text, predicted_variance = noise_pred_text.split(model_input.shape[1], dim=1) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, predicted_variance], dim=1) - - if self.scheduler.config.variance_type not in ["learned", "learned_range"]: - noise_pred, _ = noise_pred.split(model_input.shape[1], dim=1) - - # compute the previous noisy sample x_t -> x_t-1 - prev_intermediate_images = intermediate_images - - intermediate_images = self.scheduler.step( - noise_pred, t, intermediate_images, **extra_step_kwargs, return_dict=False - )[0] - - intermediate_images = (1 - mask_image) * prev_intermediate_images + mask_image * intermediate_images - - # 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 callback is not None and i % callback_steps == 0: - callback(i, t, intermediate_images) - - if XLA_AVAILABLE: - xm.mark_step() - - image = intermediate_images - - if output_type == "pil": - # 8. Post-processing - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - # 9. Run safety checker - image, nsfw_detected, watermark_detected = self.run_safety_checker(image, device, prompt_embeds.dtype) - - # 10. Convert to PIL - image = self.numpy_to_pil(image) - - # 11. Apply watermark - if self.watermarker is not None: - self.watermarker.apply_watermark(image, self.unet.config.sample_size) - elif output_type == "pt": - nsfw_detected = None - watermark_detected = None - - if hasattr(self, "unet_offload_hook") and self.unet_offload_hook is not None: - self.unet_offload_hook.offload() - else: - # 8. Post-processing - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - # 9. Run safety checker - image, nsfw_detected, watermark_detected = self.run_safety_checker(image, device, prompt_embeds.dtype) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image, nsfw_detected, watermark_detected) - - return IFPipelineOutput(images=image, nsfw_detected=nsfw_detected, watermark_detected=watermark_detected) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py deleted file mode 100644 index 0122c164d8b8029933332736c6550dd09b464a2d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py +++ /dev/null @@ -1,1134 +0,0 @@ -import html -import inspect -import re -import urllib.parse as ul -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import CLIPImageProcessor, T5EncoderModel, T5Tokenizer - -from ...loaders import StableDiffusionLoraLoaderMixin -from ...models import UNet2DConditionModel -from ...schedulers import DDPMScheduler -from ...utils import ( - BACKENDS_MAPPING, - PIL_INTERPOLATION, - is_bs4_available, - is_ftfy_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .pipeline_output import IFPipelineOutput -from .safety_checker import IFSafetyChecker -from .watermark import IFWatermarker - - -if is_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - - -from ...utils import is_torch_xla_available - - -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 - - -# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if_img2img.resize -def resize(images: PIL.Image.Image, img_size: int) -> PIL.Image.Image: - w, h = images.size - - coef = w / h - - w, h = img_size, img_size - - if coef >= 1: - w = int(round(img_size / 8 * coef) * 8) - else: - h = int(round(img_size / 8 / coef) * 8) - - images = images.resize((w, h), resample=PIL_INTERPOLATION["bicubic"], reducing_gap=None) - - return images - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> from diffusers import IFInpaintingPipeline, IFInpaintingSuperResolutionPipeline, DiffusionPipeline - >>> from diffusers.utils import pt_to_pil - >>> import torch - >>> from PIL import Image - >>> import requests - >>> from io import BytesIO - - >>> url = "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/if/person.png" - >>> response = requests.get(url) - >>> original_image = Image.open(BytesIO(response.content)).convert("RGB") - >>> original_image = original_image - - >>> url = "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/if/glasses_mask.png" - >>> response = requests.get(url) - >>> mask_image = Image.open(BytesIO(response.content)) - >>> mask_image = mask_image - - >>> pipe = IFInpaintingPipeline.from_pretrained( - ... "DeepFloyd/IF-I-XL-v1.0", variant="fp16", torch_dtype=torch.float16 - ... ) - >>> pipe.enable_model_cpu_offload() - - >>> prompt = "blue sunglasses" - - >>> prompt_embeds, negative_embeds = pipe.encode_prompt(prompt) - >>> image = pipe( - ... image=original_image, - ... mask_image=mask_image, - ... prompt_embeds=prompt_embeds, - ... negative_prompt_embeds=negative_embeds, - ... output_type="pt", - ... ).images - - >>> # save intermediate image - >>> pil_image = pt_to_pil(image) - >>> pil_image[0].save("./if_stage_I.png") - - >>> super_res_1_pipe = IFInpaintingSuperResolutionPipeline.from_pretrained( - ... "DeepFloyd/IF-II-L-v1.0", text_encoder=None, variant="fp16", torch_dtype=torch.float16 - ... ) - >>> super_res_1_pipe.enable_model_cpu_offload() - - >>> image = super_res_1_pipe( - ... image=image, - ... mask_image=mask_image, - ... original_image=original_image, - ... prompt_embeds=prompt_embeds, - ... negative_prompt_embeds=negative_embeds, - ... ).images - >>> image[0].save("./if_stage_II.png") - ``` - """ - - -class IFInpaintingSuperResolutionPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin): - tokenizer: T5Tokenizer - text_encoder: T5EncoderModel - - unet: UNet2DConditionModel - scheduler: DDPMScheduler - image_noising_scheduler: DDPMScheduler - - feature_extractor: Optional[CLIPImageProcessor] - safety_checker: Optional[IFSafetyChecker] - - watermarker: Optional[IFWatermarker] - - bad_punct_regex = re.compile( - r"[" - + "#®•©™&@·º½¾¿¡§~" - + r"\)" - + r"\(" - + r"\]" - + r"\[" - + r"\}" - + r"\{" - + r"\|" - + "\\" - + r"\/" - + r"\*" - + r"]{1,}" - ) # noqa - - model_cpu_offload_seq = "text_encoder->unet" - _optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor", "watermarker"] - _exclude_from_cpu_offload = ["watermarker"] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - unet: UNet2DConditionModel, - scheduler: DDPMScheduler, - image_noising_scheduler: DDPMScheduler, - safety_checker: Optional[IFSafetyChecker], - feature_extractor: Optional[CLIPImageProcessor], - watermarker: Optional[IFWatermarker], - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the IF license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - if unet is not None and unet.config.in_channels != 6: - logger.warning( - "It seems like you have loaded a checkpoint that shall not be used for super resolution from {unet.config._name_or_path} as it accepts {unet.config.in_channels} input channels instead of 6. Please make sure to pass a super resolution checkpoint as the `'unet'`: IFSuperResolutionPipeline.from_pretrained(unet=super_resolution_unet, ...)`." - ) - - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - unet=unet, - scheduler=scheduler, - image_noising_scheduler=image_noising_scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - watermarker=watermarker, - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - @torch.no_grad() - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - ): - 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 - device: (`torch.device`, *optional*): - torch device to place the resulting embeddings on - 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`). - 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. - clean_caption (bool, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - """ - if prompt is not None and negative_prompt is not None: - 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)} !=" - f" {type(prompt)}." - ) - - if device is None: - device = self._execution_device - - 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] - - # while T5 can handle much longer input sequences than 77, the text encoder was trained with a max length of 77 for IF - max_length = 77 - - if prompt_embeds is None: - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - add_special_tokens=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[:, max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {max_length} tokens: {removed_text}" - ) - - attention_mask = text_inputs.attention_mask.to(device) - - prompt_embeds = self.text_encoder( - text_input_ids.to(device), - attention_mask=attention_mask, - ) - prompt_embeds = prompt_embeds[0] - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - elif self.unet is not None: - dtype = self.unet.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=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 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 - - uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption) - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - attention_mask = uncond_input.attention_mask.to(device) - - 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=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) - - # 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 - else: - negative_prompt_embeds = None - - return prompt_embeds, negative_prompt_embeds - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.run_safety_checker - def run_safety_checker(self, image, device, dtype): - if self.safety_checker is not None: - safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device) - image, nsfw_detected, watermark_detected = self.safety_checker( - images=image, - clip_input=safety_checker_input.pixel_values.to(dtype=dtype), - ) - else: - nsfw_detected = None - watermark_detected = None - - return image, nsfw_detected, watermark_detected - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - image, - original_image, - mask_image, - batch_size, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - if (callback_steps is None) or ( - 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 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}." - ) - - # image - - if isinstance(image, list): - check_image_type = image[0] - else: - check_image_type = image - - if ( - not isinstance(check_image_type, torch.Tensor) - and not isinstance(check_image_type, PIL.Image.Image) - 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" - f" {type(check_image_type)}" - ) - - if isinstance(image, list): - image_batch_size = len(image) - elif isinstance(image, torch.Tensor): - image_batch_size = image.shape[0] - elif isinstance(image, PIL.Image.Image): - image_batch_size = 1 - elif isinstance(image, np.ndarray): - image_batch_size = image.shape[0] - else: - assert False - - if batch_size != image_batch_size: - raise ValueError(f"image batch size: {image_batch_size} must be same as prompt batch size {batch_size}") - - # original_image - - if isinstance(original_image, list): - check_image_type = original_image[0] - else: - check_image_type = original_image - - if ( - not isinstance(check_image_type, torch.Tensor) - and not isinstance(check_image_type, PIL.Image.Image) - 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" - f" {type(check_image_type)}" - ) - - if isinstance(original_image, list): - image_batch_size = len(original_image) - elif isinstance(original_image, torch.Tensor): - image_batch_size = original_image.shape[0] - elif isinstance(original_image, PIL.Image.Image): - image_batch_size = 1 - elif isinstance(original_image, np.ndarray): - image_batch_size = original_image.shape[0] - else: - assert False - - if batch_size != image_batch_size: - raise ValueError( - f"original_image batch size: {image_batch_size} must be same as prompt batch size {batch_size}" - ) - - # mask_image - - if isinstance(mask_image, list): - check_image_type = mask_image[0] - else: - check_image_type = mask_image - - if ( - not isinstance(check_image_type, torch.Tensor) - and not isinstance(check_image_type, PIL.Image.Image) - 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" - f" {type(check_image_type)}" - ) - - if isinstance(mask_image, list): - image_batch_size = len(mask_image) - elif isinstance(mask_image, torch.Tensor): - image_batch_size = mask_image.shape[0] - elif isinstance(mask_image, PIL.Image.Image): - image_batch_size = 1 - elif isinstance(mask_image, np.ndarray): - image_batch_size = mask_image.shape[0] - else: - assert False - - if image_batch_size != 1 and batch_size != image_batch_size: - raise ValueError( - f"mask_image batch size: {image_batch_size} must be `1` or the same as prompt batch size {batch_size}" - ) - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if_img2img.IFImg2ImgPipeline.preprocess_image with preprocess_image -> preprocess_original_image - def preprocess_original_image(self, image: PIL.Image.Image) -> torch.Tensor: - if not isinstance(image, list): - image = [image] - - def numpy_to_pt(images): - if images.ndim == 3: - images = images[..., None] - - images = torch.from_numpy(images.transpose(0, 3, 1, 2)) - return images - - if isinstance(image[0], PIL.Image.Image): - new_image = [] - - for image_ in image: - image_ = image_.convert("RGB") - image_ = resize(image_, self.unet.config.sample_size) - image_ = np.array(image_) - image_ = image_.astype(np.float32) - image_ = image_ / 127.5 - 1 - new_image.append(image_) - - image = new_image - - image = np.stack(image, axis=0) # to np - image = numpy_to_pt(image) # to pt - - elif isinstance(image[0], np.ndarray): - image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0) - image = numpy_to_pt(image) - - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0) - - return image - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if_superresolution.IFSuperResolutionPipeline.preprocess_image - def preprocess_image(self, image: PIL.Image.Image, num_images_per_prompt, device) -> torch.Tensor: - if not isinstance(image, torch.Tensor) and not isinstance(image, list): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - image = [np.array(i).astype(np.float32) / 127.5 - 1.0 for i in image] - - image = np.stack(image, axis=0) # to np - image = torch.from_numpy(image.transpose(0, 3, 1, 2)) - elif isinstance(image[0], np.ndarray): - image = np.stack(image, axis=0) # to np - if image.ndim == 5: - image = image[0] - - image = torch.from_numpy(image.transpose(0, 3, 1, 2)) - elif isinstance(image, list) and isinstance(image[0], torch.Tensor): - dims = image[0].ndim - - if dims == 3: - image = torch.stack(image, dim=0) - elif dims == 4: - image = torch.concat(image, dim=0) - else: - raise ValueError(f"Image must have 3 or 4 dimensions, instead got {dims}") - - image = image.to(device=device, dtype=self.unet.dtype) - - image = image.repeat_interleave(num_images_per_prompt, dim=0) - - return image - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if_inpainting.IFInpaintingPipeline.preprocess_mask_image - def preprocess_mask_image(self, mask_image) -> torch.Tensor: - if not isinstance(mask_image, list): - mask_image = [mask_image] - - if isinstance(mask_image[0], torch.Tensor): - mask_image = torch.cat(mask_image, axis=0) if mask_image[0].ndim == 4 else torch.stack(mask_image, axis=0) - - if mask_image.ndim == 2: - # Batch and add channel dim for single mask - mask_image = mask_image.unsqueeze(0).unsqueeze(0) - elif mask_image.ndim == 3 and mask_image.shape[0] == 1: - # Single mask, the 0'th dimension is considered to be - # the existing batch size of 1 - mask_image = mask_image.unsqueeze(0) - elif mask_image.ndim == 3 and mask_image.shape[0] != 1: - # Batch of mask, the 0'th dimension is considered to be - # the batching dimension - mask_image = mask_image.unsqueeze(1) - - mask_image[mask_image < 0.5] = 0 - mask_image[mask_image >= 0.5] = 1 - - elif isinstance(mask_image[0], PIL.Image.Image): - new_mask_image = [] - - for mask_image_ in mask_image: - mask_image_ = mask_image_.convert("L") - mask_image_ = resize(mask_image_, self.unet.config.sample_size) - mask_image_ = np.array(mask_image_) - mask_image_ = mask_image_[None, None, :] - new_mask_image.append(mask_image_) - - mask_image = new_mask_image - - mask_image = np.concatenate(mask_image, axis=0) - mask_image = mask_image.astype(np.float32) / 255.0 - mask_image[mask_image < 0.5] = 0 - mask_image[mask_image >= 0.5] = 1 - mask_image = torch.from_numpy(mask_image) - - elif isinstance(mask_image[0], np.ndarray): - mask_image = np.concatenate([m[None, None, :] for m in mask_image], axis=0) - - mask_image[mask_image < 0.5] = 0 - mask_image[mask_image >= 0.5] = 1 - mask_image = torch.from_numpy(mask_image) - - return mask_image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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 - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if_inpainting.IFInpaintingPipeline.prepare_intermediate_images - def prepare_intermediate_images( - self, image, timestep, batch_size, num_images_per_prompt, dtype, device, mask_image, generator=None - ): - image_batch_size, channels, height, width = image.shape - - batch_size = batch_size * num_images_per_prompt - - shape = (batch_size, channels, 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." - ) - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - image = image.repeat_interleave(num_images_per_prompt, dim=0) - noised_image = self.scheduler.add_noise(image, noise, timestep) - - image = (1 - mask_image) * image + mask_image * noised_image - - return image - - @torch.no_grad() - @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, - strength: float = 0.8, - prompt: Union[str, List[str]] = None, - num_inference_steps: int = 100, - timesteps: List[int] = None, - guidance_scale: float = 4.0, - 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, - 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, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - noise_level: int = 0, - clean_caption: bool = True, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - image (`torch.Tensor` or `PIL.Image.Image`): - `Image`, or tensor representing an image batch, that will be used as the starting point for the - process. - original_image (`torch.Tensor` or `PIL.Image.Image`): - The original image that `image` was varied from. - mask_image (`PIL.Image.Image`): - `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be - 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)`. - strength (`float`, *optional*, defaults to 0.8): - 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`. - 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*): - 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): - 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` 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*): - 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.IFPipelineOutput`] 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. - cross_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). - noise_level (`int`, *optional*, defaults to 0): - The amount of noise to add to the upscaled image. Must be in the range `[0, 1000)` - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.IFPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.IFPipelineOutput`] 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) - or watermarked content, according to the `safety_checker`. - """ - # 1. Check inputs. Raise error if not correct - 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] - - self.check_inputs( - prompt, - image, - original_image, - mask_image, - batch_size, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ) - - # 2. Define call parameters - - # 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 = guidance_scale > 1.0 - - device = self._execution_device - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - do_classifier_free_guidance, - num_images_per_prompt=num_images_per_prompt, - device=device, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - clean_caption=clean_caption, - ) - - if do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - dtype = prompt_embeds.dtype - - # 4. Prepare timesteps - if timesteps is not None: - self.scheduler.set_timesteps(timesteps=timesteps, device=device) - timesteps = self.scheduler.timesteps - num_inference_steps = len(timesteps) - else: - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength) - - # 5. prepare original image - original_image = self.preprocess_original_image(original_image) - original_image = original_image.to(device=device, dtype=dtype) - - # 6. prepare mask image - mask_image = self.preprocess_mask_image(mask_image) - mask_image = mask_image.to(device=device, dtype=dtype) - - if mask_image.shape[0] == 1: - mask_image = mask_image.repeat_interleave(batch_size * num_images_per_prompt, dim=0) - else: - mask_image = mask_image.repeat_interleave(num_images_per_prompt, dim=0) - - # 6. Prepare intermediate images - noise_timestep = timesteps[0:1] - noise_timestep = noise_timestep.repeat(batch_size * num_images_per_prompt) - - intermediate_images = self.prepare_intermediate_images( - original_image, - noise_timestep, - batch_size, - num_images_per_prompt, - dtype, - device, - mask_image, - generator, - ) - - # 7. Prepare upscaled image and noise level - _, _, height, width = original_image.shape - - image = self.preprocess_image(image, num_images_per_prompt, device) - - upscaled = F.interpolate(image, (height, width), mode="bilinear", align_corners=True) - - noise_level = torch.tensor([noise_level] * upscaled.shape[0], device=upscaled.device) - noise = randn_tensor(upscaled.shape, generator=generator, device=upscaled.device, dtype=upscaled.dtype) - upscaled = self.image_noising_scheduler.add_noise(upscaled, noise, timesteps=noise_level) - - if do_classifier_free_guidance: - noise_level = torch.cat([noise_level] * 2) - - # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # HACK: see comment in `enable_model_cpu_offload` - if hasattr(self, "text_encoder_offload_hook") and self.text_encoder_offload_hook is not None: - self.text_encoder_offload_hook.offload() - - # 9. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - model_input = torch.cat([intermediate_images, upscaled], dim=1) - - model_input = torch.cat([model_input] * 2) if do_classifier_free_guidance else model_input - model_input = self.scheduler.scale_model_input(model_input, t) - - # predict the noise residual - noise_pred = self.unet( - model_input, - t, - encoder_hidden_states=prompt_embeds, - class_labels=noise_level, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred_uncond, _ = noise_pred_uncond.split(model_input.shape[1] // 2, dim=1) - noise_pred_text, predicted_variance = noise_pred_text.split(model_input.shape[1] // 2, dim=1) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, predicted_variance], dim=1) - - if self.scheduler.config.variance_type not in ["learned", "learned_range"]: - noise_pred, _ = noise_pred.split(intermediate_images.shape[1], dim=1) - - # compute the previous noisy sample x_t -> x_t-1 - prev_intermediate_images = intermediate_images - - intermediate_images = self.scheduler.step( - noise_pred, t, intermediate_images, **extra_step_kwargs, return_dict=False - )[0] - - intermediate_images = (1 - mask_image) * prev_intermediate_images + mask_image * intermediate_images - - # 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 callback is not None and i % callback_steps == 0: - callback(i, t, intermediate_images) - - if XLA_AVAILABLE: - xm.mark_step() - - image = intermediate_images - - if output_type == "pil": - # 10. Post-processing - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - # 11. Run safety checker - image, nsfw_detected, watermark_detected = self.run_safety_checker(image, device, prompt_embeds.dtype) - - # 12. Convert to PIL - image = self.numpy_to_pil(image) - - # 13. Apply watermark - if self.watermarker is not None: - self.watermarker.apply_watermark(image, self.unet.config.sample_size) - elif output_type == "pt": - nsfw_detected = None - watermark_detected = None - - else: - # 10. Post-processing - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - # 11. Run safety checker - image, nsfw_detected, watermark_detected = self.run_safety_checker(image, device, prompt_embeds.dtype) - - self.maybe_free_model_hooks() - - if not return_dict: - return (image, nsfw_detected, watermark_detected) - - return IFPipelineOutput(images=image, nsfw_detected=nsfw_detected, watermark_detected=watermark_detected) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py deleted file mode 100644 index ffa60575fe33c9d450ec39684ed654b8f79f6a5d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py +++ /dev/null @@ -1,883 +0,0 @@ -import html -import inspect -import re -import urllib.parse as ul -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import CLIPImageProcessor, T5EncoderModel, T5Tokenizer - -from ...loaders import StableDiffusionLoraLoaderMixin -from ...models import UNet2DConditionModel -from ...schedulers import DDPMScheduler -from ...utils import ( - BACKENDS_MAPPING, - is_bs4_available, - is_ftfy_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .pipeline_output import IFPipelineOutput -from .safety_checker import IFSafetyChecker -from .watermark import IFWatermarker - - -if is_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - - -from ...utils import is_torch_xla_available - - -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 - >>> from diffusers import IFPipeline, IFSuperResolutionPipeline, DiffusionPipeline - >>> from diffusers.utils import pt_to_pil - >>> import torch - - >>> pipe = IFPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", variant="fp16", torch_dtype=torch.float16) - >>> pipe.enable_model_cpu_offload() - - >>> prompt = 'a photo of a kangaroo wearing an orange hoodie and blue sunglasses standing in front of the eiffel tower holding a sign that says "very deep learning"' - >>> prompt_embeds, negative_embeds = pipe.encode_prompt(prompt) - - >>> image = pipe(prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, output_type="pt").images - - >>> # save intermediate image - >>> pil_image = pt_to_pil(image) - >>> pil_image[0].save("./if_stage_I.png") - - >>> super_res_1_pipe = IFSuperResolutionPipeline.from_pretrained( - ... "DeepFloyd/IF-II-L-v1.0", text_encoder=None, variant="fp16", torch_dtype=torch.float16 - ... ) - >>> super_res_1_pipe.enable_model_cpu_offload() - - >>> image = super_res_1_pipe( - ... image=image, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds - ... ).images - >>> image[0].save("./if_stage_II.png") - ``` -""" - - -class IFSuperResolutionPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin): - tokenizer: T5Tokenizer - text_encoder: T5EncoderModel - - unet: UNet2DConditionModel - scheduler: DDPMScheduler - image_noising_scheduler: DDPMScheduler - - feature_extractor: Optional[CLIPImageProcessor] - safety_checker: Optional[IFSafetyChecker] - - watermarker: Optional[IFWatermarker] - - bad_punct_regex = re.compile( - r"[" - + "#®•©™&@·º½¾¿¡§~" - + r"\)" - + r"\(" - + r"\]" - + r"\[" - + r"\}" - + r"\{" - + r"\|" - + "\\" - + r"\/" - + r"\*" - + r"]{1,}" - ) # noqa - - _optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor", "watermarker"] - model_cpu_offload_seq = "text_encoder->unet" - _exclude_from_cpu_offload = ["watermarker"] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - unet: UNet2DConditionModel, - scheduler: DDPMScheduler, - image_noising_scheduler: DDPMScheduler, - safety_checker: Optional[IFSafetyChecker], - feature_extractor: Optional[CLIPImageProcessor], - watermarker: Optional[IFWatermarker], - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the IF license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - if unet is not None and unet.config.in_channels != 6: - logger.warning( - "It seems like you have loaded a checkpoint that shall not be used for super resolution from {unet.config._name_or_path} as it accepts {unet.config.in_channels} input channels instead of 6. Please make sure to pass a super resolution checkpoint as the `'unet'`: IFSuperResolutionPipeline.from_pretrained(unet=super_resolution_unet, ...)`." - ) - - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - unet=unet, - scheduler=scheduler, - image_noising_scheduler=image_noising_scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - watermarker=watermarker, - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - @torch.no_grad() - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - ): - 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 - device: (`torch.device`, *optional*): - torch device to place the resulting embeddings on - 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`). - 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. - clean_caption (bool, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - """ - if prompt is not None and negative_prompt is not None: - 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)} !=" - f" {type(prompt)}." - ) - - if device is None: - device = self._execution_device - - 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] - - # while T5 can handle much longer input sequences than 77, the text encoder was trained with a max length of 77 for IF - max_length = 77 - - if prompt_embeds is None: - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - add_special_tokens=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[:, max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {max_length} tokens: {removed_text}" - ) - - attention_mask = text_inputs.attention_mask.to(device) - - prompt_embeds = self.text_encoder( - text_input_ids.to(device), - attention_mask=attention_mask, - ) - prompt_embeds = prompt_embeds[0] - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - elif self.unet is not None: - dtype = self.unet.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=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 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 - - uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption) - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - attention_mask = uncond_input.attention_mask.to(device) - - 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=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) - - # 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 - else: - negative_prompt_embeds = None - - return prompt_embeds, negative_prompt_embeds - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.run_safety_checker - def run_safety_checker(self, image, device, dtype): - if self.safety_checker is not None: - safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device) - image, nsfw_detected, watermark_detected = self.safety_checker( - images=image, - clip_input=safety_checker_input.pixel_values.to(dtype=dtype), - ) - else: - nsfw_detected = None - watermark_detected = None - - return image, nsfw_detected, watermark_detected - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - image, - batch_size, - noise_level, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - if (callback_steps is None) or ( - 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 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 noise_level < 0 or noise_level >= self.image_noising_scheduler.config.num_train_timesteps: - raise ValueError( - f"`noise_level`: {noise_level} must be a valid timestep in `self.noising_scheduler`, [0, {self.image_noising_scheduler.config.num_train_timesteps})" - ) - - if isinstance(image, list): - check_image_type = image[0] - else: - check_image_type = image - - if ( - not isinstance(check_image_type, torch.Tensor) - and not isinstance(check_image_type, PIL.Image.Image) - 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" - f" {type(check_image_type)}" - ) - - if isinstance(image, list): - image_batch_size = len(image) - elif isinstance(image, torch.Tensor): - image_batch_size = image.shape[0] - elif isinstance(image, PIL.Image.Image): - image_batch_size = 1 - elif isinstance(image, np.ndarray): - image_batch_size = image.shape[0] - else: - assert False - - if batch_size != image_batch_size: - raise ValueError(f"image batch size: {image_batch_size} must be same as prompt batch size {batch_size}") - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.prepare_intermediate_images - def prepare_intermediate_images(self, batch_size, num_channels, height, width, dtype, device, generator): - shape = (batch_size, num_channels, 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." - ) - - intermediate_images = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # scale the initial noise by the standard deviation required by the scheduler - intermediate_images = intermediate_images * self.scheduler.init_noise_sigma - return intermediate_images - - def preprocess_image(self, image, num_images_per_prompt, device): - if not isinstance(image, torch.Tensor) and not isinstance(image, list): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - image = [np.array(i).astype(np.float32) / 127.5 - 1.0 for i in image] - - image = np.stack(image, axis=0) # to np - image = torch.from_numpy(image.transpose(0, 3, 1, 2)) - elif isinstance(image[0], np.ndarray): - image = np.stack(image, axis=0) # to np - if image.ndim == 5: - image = image[0] - - image = torch.from_numpy(image.transpose(0, 3, 1, 2)) - elif isinstance(image, list) and isinstance(image[0], torch.Tensor): - dims = image[0].ndim - - if dims == 3: - image = torch.stack(image, dim=0) - elif dims == 4: - image = torch.concat(image, dim=0) - else: - raise ValueError(f"Image must have 3 or 4 dimensions, instead got {dims}") - - image = image.to(device=device, dtype=self.unet.dtype) - - image = image.repeat_interleave(num_images_per_prompt, dim=0) - - return image - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - height: int = None, - width: int = None, - image: Union[PIL.Image.Image, np.ndarray, torch.Tensor] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - guidance_scale: float = 4.0, - 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, - 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, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - noise_level: int = 250, - clean_caption: bool = True, - ): - """ - 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 None): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to None): - The width in pixels of the generated image. - image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`): - The image 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 - expense of slower inference. - 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): - 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` 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*): - 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.IFPipelineOutput`] 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. - cross_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). - noise_level (`int`, *optional*, defaults to 250): - The amount of noise to add to the upscaled image. Must be in the range `[0, 1000)` - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.IFPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.IFPipelineOutput`] 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) - or watermarked content, according to the `safety_checker`. - """ - # 1. Check inputs. Raise error if not correct - - 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] - - self.check_inputs( - prompt, - image, - batch_size, - noise_level, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ) - - # 2. Define call parameters - - height = height or self.unet.config.sample_size - width = width or self.unet.config.sample_size - - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - do_classifier_free_guidance, - num_images_per_prompt=num_images_per_prompt, - device=device, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - clean_caption=clean_caption, - ) - - if do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - # 4. Prepare timesteps - if timesteps is not None: - self.scheduler.set_timesteps(timesteps=timesteps, device=device) - timesteps = self.scheduler.timesteps - num_inference_steps = len(timesteps) - else: - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - if hasattr(self.scheduler, "set_begin_index"): - self.scheduler.set_begin_index(0) - - # 5. Prepare intermediate images - num_channels = self.unet.config.in_channels // 2 - intermediate_images = self.prepare_intermediate_images( - batch_size * num_images_per_prompt, - num_channels, - height, - width, - prompt_embeds.dtype, - device, - generator, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Prepare upscaled image and noise level - image = self.preprocess_image(image, num_images_per_prompt, device) - upscaled = F.interpolate(image, (height, width), mode="bilinear", align_corners=True) - - noise_level = torch.tensor([noise_level] * upscaled.shape[0], device=upscaled.device) - noise = randn_tensor(upscaled.shape, generator=generator, device=upscaled.device, dtype=upscaled.dtype) - upscaled = self.image_noising_scheduler.add_noise(upscaled, noise, timesteps=noise_level) - - if do_classifier_free_guidance: - noise_level = torch.cat([noise_level] * 2) - - # HACK: see comment in `enable_model_cpu_offload` - if hasattr(self, "text_encoder_offload_hook") and self.text_encoder_offload_hook is not None: - self.text_encoder_offload_hook.offload() - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - model_input = torch.cat([intermediate_images, upscaled], dim=1) - - model_input = torch.cat([model_input] * 2) if do_classifier_free_guidance else model_input - model_input = self.scheduler.scale_model_input(model_input, t) - - # predict the noise residual - noise_pred = self.unet( - model_input, - t, - encoder_hidden_states=prompt_embeds, - class_labels=noise_level, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred_uncond, _ = noise_pred_uncond.split(model_input.shape[1] // 2, dim=1) - noise_pred_text, predicted_variance = noise_pred_text.split(model_input.shape[1] // 2, dim=1) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, predicted_variance], dim=1) - - if self.scheduler.config.variance_type not in ["learned", "learned_range"]: - noise_pred, _ = noise_pred.split(intermediate_images.shape[1], dim=1) - - # compute the previous noisy sample x_t -> x_t-1 - intermediate_images = self.scheduler.step( - noise_pred, t, intermediate_images, **extra_step_kwargs, return_dict=False - )[0] - - # 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 callback is not None and i % callback_steps == 0: - callback(i, t, intermediate_images) - - if XLA_AVAILABLE: - xm.mark_step() - - image = intermediate_images - - if output_type == "pil": - # 9. Post-processing - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - # 10. Run safety checker - image, nsfw_detected, watermark_detected = self.run_safety_checker(image, device, prompt_embeds.dtype) - - # 11. Convert to PIL - image = self.numpy_to_pil(image) - - # 12. Apply watermark - if self.watermarker is not None: - self.watermarker.apply_watermark(image, self.unet.config.sample_size) - elif output_type == "pt": - nsfw_detected = None - watermark_detected = None - - if hasattr(self, "unet_offload_hook") and self.unet_offload_hook is not None: - self.unet_offload_hook.offload() - else: - # 9. Post-processing - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - # 10. Run safety checker - image, nsfw_detected, watermark_detected = self.run_safety_checker(image, device, prompt_embeds.dtype) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image, nsfw_detected, watermark_detected) - - return IFPipelineOutput(images=image, nsfw_detected=nsfw_detected, watermark_detected=watermark_detected) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_output.py deleted file mode 100644 index b8bae89cec03fd638fcf4eaf55481b5ca14ebf3d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/pipeline_output.py +++ /dev/null @@ -1,29 +0,0 @@ -from dataclasses import dataclass -from typing import List, Optional, Union - -import numpy as np -import PIL.Image - -from ...utils import BaseOutput - - -@dataclass -class IFPipelineOutput(BaseOutput): - r""" - 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, - 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) 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 - checking could not be performed. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] - nsfw_detected: Optional[List[bool]] - watermark_detected: Optional[List[bool]] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/safety_checker.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/safety_checker.py deleted file mode 100644 index 8ffeed580bbea1514b11bf7a168a952328d8f424..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/safety_checker.py +++ /dev/null @@ -1,59 +0,0 @@ -import numpy as np -import torch -import torch.nn as nn -from transformers import CLIPConfig, CLIPVisionModelWithProjection, PreTrainedModel - -from ...utils import logging - - -logger = logging.get_logger(__name__) - - -class IFSafetyChecker(PreTrainedModel): - config_class = CLIPConfig - - _no_split_modules = ["CLIPEncoderLayer"] - - def __init__(self, config: CLIPConfig): - super().__init__(config) - - self.vision_model = CLIPVisionModelWithProjection(config.vision_config) - - self.p_head = nn.Linear(config.vision_config.projection_dim, 1) - self.w_head = nn.Linear(config.vision_config.projection_dim, 1) - - @torch.no_grad() - def forward(self, clip_input, images, p_threshold=0.5, w_threshold=0.5): - image_embeds = self.vision_model(clip_input)[0] - - nsfw_detected = self.p_head(image_embeds) - nsfw_detected = nsfw_detected.flatten() - nsfw_detected = nsfw_detected > p_threshold - nsfw_detected = nsfw_detected.tolist() - - if any(nsfw_detected): - logger.warning( - "Potential NSFW content was detected in one or more images. A black image will be returned instead." - " Try again with a different prompt and/or seed." - ) - - for idx, nsfw_detected_ in enumerate(nsfw_detected): - if nsfw_detected_: - images[idx] = np.zeros(images[idx].shape) - - watermark_detected = self.w_head(image_embeds) - watermark_detected = watermark_detected.flatten() - watermark_detected = watermark_detected > w_threshold - watermark_detected = watermark_detected.tolist() - - if any(watermark_detected): - logger.warning( - "Potential watermarked content was detected in one or more images. A black image will be returned instead." - " Try again with a different prompt and/or seed." - ) - - for idx, watermark_detected_ in enumerate(watermark_detected): - if watermark_detected_: - images[idx] = np.zeros(images[idx].shape) - - return images, nsfw_detected, watermark_detected diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/timesteps.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/timesteps.py deleted file mode 100644 index d44285c017bbb2ccffa4ae86dd77792a048625d9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/timesteps.py +++ /dev/null @@ -1,579 +0,0 @@ -fast27_timesteps = [ - 999, - 800, - 799, - 600, - 599, - 500, - 400, - 399, - 377, - 355, - 333, - 311, - 288, - 266, - 244, - 222, - 200, - 199, - 177, - 155, - 133, - 111, - 88, - 66, - 44, - 22, - 0, -] - -smart27_timesteps = [ - 999, - 976, - 952, - 928, - 905, - 882, - 858, - 857, - 810, - 762, - 715, - 714, - 572, - 429, - 428, - 286, - 285, - 238, - 190, - 143, - 142, - 118, - 95, - 71, - 47, - 24, - 0, -] - -smart50_timesteps = [ - 999, - 988, - 977, - 966, - 955, - 944, - 933, - 922, - 911, - 900, - 899, - 879, - 859, - 840, - 820, - 800, - 799, - 766, - 733, - 700, - 699, - 650, - 600, - 599, - 500, - 499, - 400, - 399, - 350, - 300, - 299, - 266, - 233, - 200, - 199, - 179, - 159, - 140, - 120, - 100, - 99, - 88, - 77, - 66, - 55, - 44, - 33, - 22, - 11, - 0, -] - -smart100_timesteps = [ - 999, - 995, - 992, - 989, - 985, - 981, - 978, - 975, - 971, - 967, - 964, - 961, - 957, - 956, - 951, - 947, - 942, - 937, - 933, - 928, - 923, - 919, - 914, - 913, - 908, - 903, - 897, - 892, - 887, - 881, - 876, - 871, - 870, - 864, - 858, - 852, - 846, - 840, - 834, - 828, - 827, - 820, - 813, - 806, - 799, - 792, - 785, - 784, - 777, - 770, - 763, - 756, - 749, - 742, - 741, - 733, - 724, - 716, - 707, - 699, - 698, - 688, - 677, - 666, - 656, - 655, - 645, - 634, - 623, - 613, - 612, - 598, - 584, - 570, - 569, - 555, - 541, - 527, - 526, - 505, - 484, - 483, - 462, - 440, - 439, - 396, - 395, - 352, - 351, - 308, - 307, - 264, - 263, - 220, - 219, - 176, - 132, - 88, - 44, - 0, -] - -smart185_timesteps = [ - 999, - 997, - 995, - 992, - 990, - 988, - 986, - 984, - 981, - 979, - 977, - 975, - 972, - 970, - 968, - 966, - 964, - 961, - 959, - 957, - 956, - 954, - 951, - 949, - 946, - 944, - 941, - 939, - 936, - 934, - 931, - 929, - 926, - 924, - 921, - 919, - 916, - 914, - 913, - 910, - 907, - 905, - 902, - 899, - 896, - 893, - 891, - 888, - 885, - 882, - 879, - 877, - 874, - 871, - 870, - 867, - 864, - 861, - 858, - 855, - 852, - 849, - 846, - 843, - 840, - 837, - 834, - 831, - 828, - 827, - 824, - 821, - 817, - 814, - 811, - 808, - 804, - 801, - 798, - 795, - 791, - 788, - 785, - 784, - 780, - 777, - 774, - 770, - 766, - 763, - 760, - 756, - 752, - 749, - 746, - 742, - 741, - 737, - 733, - 730, - 726, - 722, - 718, - 714, - 710, - 707, - 703, - 699, - 698, - 694, - 690, - 685, - 681, - 677, - 673, - 669, - 664, - 660, - 656, - 655, - 650, - 646, - 641, - 636, - 632, - 627, - 622, - 618, - 613, - 612, - 607, - 602, - 596, - 591, - 586, - 580, - 575, - 570, - 569, - 563, - 557, - 551, - 545, - 539, - 533, - 527, - 526, - 519, - 512, - 505, - 498, - 491, - 484, - 483, - 474, - 466, - 457, - 449, - 440, - 439, - 428, - 418, - 407, - 396, - 395, - 381, - 366, - 352, - 351, - 330, - 308, - 307, - 286, - 264, - 263, - 242, - 220, - 219, - 176, - 175, - 132, - 131, - 88, - 44, - 0, -] - -super27_timesteps = [ - 999, - 991, - 982, - 974, - 966, - 958, - 950, - 941, - 933, - 925, - 916, - 908, - 900, - 899, - 874, - 850, - 825, - 800, - 799, - 700, - 600, - 500, - 400, - 300, - 200, - 100, - 0, -] - -super40_timesteps = [ - 999, - 992, - 985, - 978, - 971, - 964, - 957, - 949, - 942, - 935, - 928, - 921, - 914, - 907, - 900, - 899, - 879, - 859, - 840, - 820, - 800, - 799, - 766, - 733, - 700, - 699, - 650, - 600, - 599, - 500, - 499, - 400, - 399, - 300, - 299, - 200, - 199, - 100, - 99, - 0, -] - -super100_timesteps = [ - 999, - 996, - 992, - 989, - 985, - 982, - 979, - 975, - 972, - 968, - 965, - 961, - 958, - 955, - 951, - 948, - 944, - 941, - 938, - 934, - 931, - 927, - 924, - 920, - 917, - 914, - 910, - 907, - 903, - 900, - 899, - 891, - 884, - 876, - 869, - 861, - 853, - 846, - 838, - 830, - 823, - 815, - 808, - 800, - 799, - 788, - 777, - 766, - 755, - 744, - 733, - 722, - 711, - 700, - 699, - 688, - 677, - 666, - 655, - 644, - 633, - 622, - 611, - 600, - 599, - 585, - 571, - 557, - 542, - 528, - 514, - 500, - 499, - 485, - 471, - 457, - 442, - 428, - 414, - 400, - 399, - 379, - 359, - 340, - 320, - 300, - 299, - 279, - 259, - 240, - 220, - 200, - 199, - 166, - 133, - 100, - 99, - 66, - 33, - 0, -] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/watermark.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/watermark.py deleted file mode 100644 index e03e3fab026a5702eb9c45a4d19fd3d4a0d3da6d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deepfloyd_if/watermark.py +++ /dev/null @@ -1,46 +0,0 @@ -from typing import List - -import PIL.Image -import torch -from PIL import Image - -from ...configuration_utils import ConfigMixin -from ...models.modeling_utils import ModelMixin -from ...utils import PIL_INTERPOLATION - - -class IFWatermarker(ModelMixin, ConfigMixin): - def __init__(self): - super().__init__() - - 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): - # Copied from https://github.com/deep-floyd/IF/blob/b77482e36ca2031cb94dbca1001fc1e6400bf4ab/deepfloyd_if/modules/base.py#L287 - - h = images[0].height - w = images[0].width - - sample_size = sample_size or h - - coef = min(h / sample_size, w / sample_size) - img_h, img_w = (int(h / coef), int(w / coef)) if coef < 1 else (h, w) - - S1, S2 = 1024**2, img_w * img_h - K = (S2 / S1) ** 0.5 - wm_size, wm_x, wm_y = int(K * 62), img_w - int(14 * K), img_h - int(14 * K) - - if self.watermark_image_as_pil is None: - watermark_image = self.watermark_image.to(torch.uint8).cpu().numpy() - watermark_image = Image.fromarray(watermark_image, mode="RGBA") - self.watermark_image_as_pil = watermark_image - - wm_img = self.watermark_image_as_pil.resize( - (wm_size, wm_size), PIL_INTERPOLATION["bicubic"], reducing_gap=None - ) - - for pil_img in images: - pil_img.paste(wm_img, box=(wm_x - wm_size, wm_y - wm_size, wm_x, wm_y), mask=wm_img.split()[-1]) - - return images diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/__init__.py deleted file mode 100644 index 9936323170adbceac2c5c25e3881ea731d8602e1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/__init__.py +++ /dev/null @@ -1,153 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - get_objects_from_module, - is_librosa_available, - is_note_seq_available, - is_torch_available, - is_transformers_available, -) - - -_dummy_objects = {} -_import_structure = {} - -try: - if not is_torch_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_pt_objects - - _dummy_objects.update(get_objects_from_module(dummy_pt_objects)) -else: - _import_structure["latent_diffusion_uncond"] = ["LDMPipeline"] - _import_structure["pndm"] = ["PNDMPipeline"] - _import_structure["repaint"] = ["RePaintPipeline"] - _import_structure["score_sde_ve"] = ["ScoreSdeVePipeline"] - _import_structure["stochastic_karras_ve"] = ["KarrasVePipeline"] - -try: - if not (is_transformers_available() and is_torch_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_torch_and_transformers_objects - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) -else: - _import_structure["alt_diffusion"] = [ - "AltDiffusionImg2ImgPipeline", - "AltDiffusionPipeline", - "AltDiffusionPipelineOutput", - ] - _import_structure["versatile_diffusion"] = [ - "VersatileDiffusionDualGuidedPipeline", - "VersatileDiffusionImageVariationPipeline", - "VersatileDiffusionPipeline", - "VersatileDiffusionTextToImagePipeline", - ] - _import_structure["vq_diffusion"] = ["VQDiffusionPipeline"] - _import_structure["stable_diffusion_variants"] = [ - "CycleDiffusionPipeline", - "StableDiffusionInpaintPipelineLegacy", - "StableDiffusionPix2PixZeroPipeline", - "StableDiffusionParadigmsPipeline", - "StableDiffusionModelEditingPipeline", - ] - -try: - if not (is_torch_available() and is_librosa_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_torch_and_librosa_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_librosa_objects)) - -else: - _import_structure["audio_diffusion"] = ["AudioDiffusionPipeline", "Mel"] - -try: - if not (is_transformers_available() and is_torch_available() and is_note_seq_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_transformers_and_torch_and_note_seq_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_transformers_and_torch_and_note_seq_objects)) - -else: - _import_structure["spectrogram_diffusion"] = ["MidiProcessor", "SpectrogramDiffusionPipeline"] - - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - try: - if not is_torch_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_pt_objects import * - - else: - from .latent_diffusion_uncond import LDMPipeline - from .pndm import PNDMPipeline - from .repaint import RePaintPipeline - from .score_sde_ve import ScoreSdeVePipeline - from .stochastic_karras_ve import KarrasVePipeline - - try: - if not (is_transformers_available() and is_torch_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_objects import * - - else: - from .alt_diffusion import AltDiffusionImg2ImgPipeline, AltDiffusionPipeline, AltDiffusionPipelineOutput - from .audio_diffusion import AudioDiffusionPipeline, Mel - from .spectrogram_diffusion import SpectrogramDiffusionPipeline - from .stable_diffusion_variants import ( - CycleDiffusionPipeline, - StableDiffusionInpaintPipelineLegacy, - StableDiffusionModelEditingPipeline, - StableDiffusionParadigmsPipeline, - StableDiffusionPix2PixZeroPipeline, - ) - from .stochastic_karras_ve import KarrasVePipeline - from .versatile_diffusion import ( - VersatileDiffusionDualGuidedPipeline, - VersatileDiffusionImageVariationPipeline, - VersatileDiffusionPipeline, - VersatileDiffusionTextToImagePipeline, - ) - from .vq_diffusion import VQDiffusionPipeline - - try: - if not (is_torch_available() and is_librosa_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_librosa_objects import * - else: - from .audio_diffusion import AudioDiffusionPipeline, Mel - - try: - if not (is_transformers_available() and is_torch_available() and is_note_seq_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_transformers_and_torch_and_note_seq_objects import * # noqa F403 - else: - from .spectrogram_diffusion import ( - MidiProcessor, - SpectrogramDiffusionPipeline, - ) - - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index a4568c95706016a4499fc05d36ef61ad5a2f1ed5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__init__.py deleted file mode 100644 index 71fa15b3feff08dc4008d1fa02ba61ad1300efed..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__init__.py +++ /dev/null @@ -1,53 +0,0 @@ -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 - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) -else: - _import_structure["modeling_roberta_series"] = ["RobertaSeriesModelWithTransformation"] - _import_structure["pipeline_alt_diffusion"] = ["AltDiffusionPipeline"] - _import_structure["pipeline_alt_diffusion_img2img"] = ["AltDiffusionImg2ImgPipeline"] - - _import_structure["pipeline_output"] = ["AltDiffusionPipelineOutput"] - -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 .modeling_roberta_series import RobertaSeriesModelWithTransformation - from .pipeline_alt_diffusion import AltDiffusionPipeline - from .pipeline_alt_diffusion_img2img import AltDiffusionImg2ImgPipeline - from .pipeline_output import AltDiffusionPipelineOutput - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 71de5e5e9a22eb39c975dec4956413a86cf4350b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/modeling_roberta_series.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/modeling_roberta_series.cpython-310.pyc deleted file mode 100644 index b198215700b36cf5c6244c3a75f8b02952d5e51a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/modeling_roberta_series.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/pipeline_alt_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/pipeline_alt_diffusion.cpython-310.pyc deleted file mode 100644 index 404614647841996ab7cc8521609da54824fee6bd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/pipeline_alt_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/pipeline_alt_diffusion_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/pipeline_alt_diffusion_img2img.cpython-310.pyc deleted file mode 100644 index a012b699aa45c1d1ad9e36544b71698e44a5ee00..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/pipeline_alt_diffusion_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index be2afee77f5a58424d1445701621f57cf255acd9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py deleted file mode 100644 index f69f905b56c528a292340fcf33ee3b016096d1f4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py +++ /dev/null @@ -1,124 +0,0 @@ -from dataclasses import dataclass -from typing import Optional, Tuple - -import torch -from torch import nn -from transformers import RobertaPreTrainedModel, XLMRobertaConfig, XLMRobertaModel -from transformers.utils import ModelOutput - - -@dataclass -class TransformationModelOutput(ModelOutput): - """ - Base class for text model's outputs that also contains a pooling of the last hidden states. - - Args: - text_embeds (`torch.Tensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): - The text embeddings obtained by applying the projection layer to the pooler_output. - 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 - 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, - 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 - last_hidden_state: torch.Tensor = None - hidden_states: Optional[Tuple[torch.Tensor]] = None - attentions: Optional[Tuple[torch.Tensor]] = None - - -class RobertaSeriesConfig(XLMRobertaConfig): - def __init__( - self, - pad_token_id=1, - bos_token_id=0, - eos_token_id=2, - project_dim=512, - pooler_fn="cls", - learn_encoder=False, - use_attention_mask=True, - **kwargs, - ): - super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs) - self.project_dim = project_dim - self.pooler_fn = pooler_fn - self.learn_encoder = learn_encoder - self.use_attention_mask = use_attention_mask - - -class RobertaSeriesModelWithTransformation(RobertaPreTrainedModel): - _keys_to_ignore_on_load_unexpected = [r"pooler", r"logit_scale"] - _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"] - base_model_prefix = "roberta" - config_class = RobertaSeriesConfig - - def __init__(self, config): - super().__init__(config) - self.roberta = XLMRobertaModel(config) - self.transformation = nn.Linear(config.hidden_size, config.project_dim) - self.has_pre_transformation = getattr(config, "has_pre_transformation", False) - if self.has_pre_transformation: - self.transformation_pre = nn.Linear(config.hidden_size, config.project_dim) - self.pre_LN = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - self.post_init() - - 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, - ): - r""" """ - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - outputs = self.base_model( - input_ids=input_ids, - attention_mask=attention_mask, - token_type_ids=token_type_ids, - position_ids=position_ids, - head_mask=head_mask, - inputs_embeds=inputs_embeds, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - output_attentions=output_attentions, - output_hidden_states=True if self.has_pre_transformation else output_hidden_states, - return_dict=return_dict, - ) - - if self.has_pre_transformation: - sequence_output2 = outputs["hidden_states"][-2] - sequence_output2 = self.pre_LN(sequence_output2) - projection_state2 = self.transformation_pre(sequence_output2) - - return TransformationModelOutput( - projection_state=projection_state2, - last_hidden_state=outputs.last_hidden_state, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - else: - projection_state = self.transformation(outputs.last_hidden_state) - return TransformationModelOutput( - projection_state=projection_state, - last_hidden_state=outputs.last_hidden_state, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py deleted file mode 100644 index 003e7482741fb522d89a76bf238d6a9e15206eec..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py +++ /dev/null @@ -1,990 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import torch -from packaging import version -from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection, XLMRobertaTokenizer - -from ....configuration_utils import FrozenDict -from ....image_processor import PipelineImageInput, VaeImageProcessor -from ....loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ....models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ....models.lora import adjust_lora_scale_text_encoder -from ....schedulers import KarrasDiffusionSchedulers -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 DiffusionPipeline, StableDiffusionMixin -from ...stable_diffusion.safety_checker import StableDiffusionSafetyChecker -from .modeling_roberta_series import RobertaSeriesModelWithTransformation -from .pipeline_output import AltDiffusionPipelineOutput - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import AltDiffusionPipeline - - >>> pipe = AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion-m9", torch_dtype=torch.float16) - >>> pipe = pipe.to("cuda") - - >>> # "dark elf princess, highly detailed, d & d, fantasy, highly detailed, digital painting, trending on artstation, concept art, sharp focus, illustration, art by artgerm and greg rutkowski and fuji choko and viktoria gavrilenko and hoang lap" - >>> prompt = "黑暗精灵公主，非常详细，幻想，非常详细，数字绘画，概念艺术，敏锐的焦点，插图" - >>> image = pipe(prompt).images[0] - ``` -""" - - -# 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 - - -# 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 AltDiffusionPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - IPAdapterMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-to-image generation using Alt Diffusion. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.RobertaSeriesModelWithTransformation`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.XLMRobertaTokenizer`]): - A `XLMRobertaTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: RobertaSeriesModelWithTransformation, - tokenizer: XLMRobertaTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: CLIPVisionModelWithProjection = None, - requires_safety_checker: bool = True, - ): - super().__init__() - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Alt Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_config) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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.register_to_config(requires_safety_checker=requires_safety_checker) - - 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 - - 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 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 - - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - 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 - - 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 prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - 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 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=torch.float32): - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - timesteps (`torch.Tensor`): - generate embedding vectors at these timesteps - embedding_dim (`int`, *optional*, defaults to 512): - dimension of the embeddings to generate - dtype: - data type of the generated embeddings - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(timesteps), embedding_dim)` - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - @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: int = 50, - 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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"], - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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. - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - 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 [`~pipelines.stable_diffusion.AltDiffusionPipelineOutput`] instead of a - plain tuple. - 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). - guidance_rescale (`float`, *optional*, defaults to 0.0): - Guidance rescale factor from [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. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.AltDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.AltDiffusionPipelineOutput`] 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. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - - # 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 - # to deal with lora scaling and other possible forward hooks - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - - # 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 = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - if ip_adapter_image is not None: - output_hidden_state = False if isinstance(self.unet.encoder_hid_proj, ImageProjection) else True - image_embeds, negative_image_embeds = self.encode_image( - ip_adapter_image, device, num_images_per_prompt, output_hidden_state - ) - if self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds]) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 5. 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, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 6.1 Add image embeds for IP-Adapter - added_cond_kwargs = {"image_embeds": image_embeds} if ip_adapter_image is not None else None - - # 6.2 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) - - # 7. Denoising loop - 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): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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.do_classifier_free_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 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 AltDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py deleted file mode 100644 index 64bd06d02ea70ec66a761b68960a0f0149364486..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py +++ /dev/null @@ -1,1045 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from packaging import version -from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection, XLMRobertaTokenizer - -from ....configuration_utils import FrozenDict -from ....image_processor import PipelineImageInput, VaeImageProcessor -from ....loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ....models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ....models.lora import adjust_lora_scale_text_encoder -from ....schedulers import KarrasDiffusionSchedulers -from ....utils import ( - PIL_INTERPOLATION, - 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 DiffusionPipeline, StableDiffusionMixin -from ...stable_diffusion.safety_checker import StableDiffusionSafetyChecker -from .modeling_roberta_series import RobertaSeriesModelWithTransformation -from .pipeline_output import AltDiffusionPipelineOutput - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import requests - >>> import torch - >>> from PIL import Image - >>> from io import BytesIO - - >>> from diffusers import AltDiffusionImg2ImgPipeline - - >>> device = "cuda" - >>> model_id_or_path = "BAAI/AltDiffusion-m9" - >>> pipe = AltDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16) - >>> pipe = pipe.to(device) - - >>> url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" - - >>> response = requests.get(url) - >>> init_image = Image.open(BytesIO(response.content)).convert("RGB") - >>> init_image = init_image.resize((768, 512)) - - >>> # "A fantasy landscape, trending on artstation" - >>> prompt = "幻想风景, artstation" - - >>> images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images - >>> images[0].save("幻想风景.png") - ``` -""" - - -# 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" -): - 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_img2img.preprocess -def preprocess(image): - 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): - return image - elif isinstance(image, PIL.Image.Image): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - w, h = image[0].size - w, h = (x - x % 8 for x in (w, h)) # resize to integer multiple of 8 - - image = [np.array(i.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image] - image = np.concatenate(image, axis=0) - image = np.array(image).astype(np.float32) / 255.0 - image = image.transpose(0, 3, 1, 2) - image = 2.0 * image - 1.0 - image = torch.from_numpy(image) - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, dim=0) - return 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, - **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 AltDiffusionImg2ImgPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-guided image-to-image generation using Alt Diffusion. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.RobertaSeriesModelWithTransformation`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.XLMRobertaTokenizer`]): - A `XLMRobertaTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: RobertaSeriesModelWithTransformation, - tokenizer: XLMRobertaTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: CLIPVisionModelWithProjection = None, - requires_safety_checker: bool = True, - ): - super().__init__() - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Alt Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_config) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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.register_to_config(requires_safety_checker=requires_safety_checker) - - 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 - - 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 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 - - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - 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 - - 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 prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - strength, - callback_steps, - 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 in [0.0, 1.0] but is {strength}") - - 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 get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :] - - return timesteps, num_inference_steps - t_start - - def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - 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." - ) - - elif isinstance(generator, list): - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - deprecation_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial" - " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" - " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" - " your script to pass as many initial images as text prompts to suppress this warning." - ) - deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - latents = init_latents - - return latents - - def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=torch.float32): - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - timesteps (`torch.Tensor`): - generate embedding vectors at these timesteps - embedding_dim (`int`, *optional*, defaults to 512): - dimension of the embeddings to generate - dtype: - data type of the generated embeddings - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(timesteps), embedding_dim)` - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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"], - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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`, 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. - strength (`float`, *optional*, defaults to 0.8): - 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`. - 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*): - 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*): - 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): - The number of images to generate per prompt. - 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*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - 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. - 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. - ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - 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 [`~pipelines.stable_diffusion.AltDiffusionPipelineOutput`] instead of a - plain tuple. - 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). - 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. - 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. - Examples: - - Returns: - [`~pipelines.stable_diffusion.AltDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.AltDiffusionPipelineOutput`] 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. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - strength, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - - # 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 - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - if ip_adapter_image is not None: - output_hidden_state = False if isinstance(self.unet.encoder_hid_proj, ImageProjection) else True - image_embeds, negative_image_embeds = self.encode_image( - ip_adapter_image, device, num_images_per_prompt, output_hidden_state - ) - if self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds]) - - # 4. Preprocess image - image = self.image_processor.preprocess(image) - - # 5. set timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, 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) - - # 6. Prepare latent variables - latents = self.prepare_latents( - image, - latent_timestep, - batch_size, - num_images_per_prompt, - prompt_embeds.dtype, - device, - generator, - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Add image embeds for IP-Adapter - added_cond_kwargs = {"image_embeds": image_embeds} if ip_adapter_image is not None else None - - # 7.2 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) - - # 8. Denoising loop - 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): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 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 AltDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/pipeline_output.py deleted file mode 100644 index dd174ae3c21fe4110babd503f0418366472059ff..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/alt_diffusion/pipeline_output.py +++ /dev/null @@ -1,28 +0,0 @@ -from dataclasses import dataclass -from typing import List, Optional, Union - -import numpy as np -import PIL.Image - -from ....utils import ( - BaseOutput, -) - - -@dataclass -# Copied from diffusers.pipelines.stable_diffusion.pipeline_output.StableDiffusionPipelineOutput with Stable->Alt -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, - num_channels)`. - 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]] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/__init__.py deleted file mode 100644 index 3127951863a7db3f9dd8e42ac5ab64fa9ac3ec0c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -from typing import TYPE_CHECKING - -from ....utils import DIFFUSERS_SLOW_IMPORT, _LazyModule - - -_import_structure = { - "mel": ["Mel"], - "pipeline_audio_diffusion": ["AudioDiffusionPipeline"], -} - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - from .mel import Mel - from .pipeline_audio_diffusion import AudioDiffusionPipeline - -else: - import sys - - sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 366a82727171a45f9441af29925513e277c95f7b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/__pycache__/mel.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/__pycache__/mel.cpython-310.pyc deleted file mode 100644 index 831eab697bff89bc70f90c9ca90ad29ee5544c5d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/__pycache__/mel.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/__pycache__/pipeline_audio_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/__pycache__/pipeline_audio_diffusion.cpython-310.pyc deleted file mode 100644 index 240401adcd4311dac4afc9869d79cc24ef620df3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/__pycache__/pipeline_audio_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/mel.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/mel.py deleted file mode 100644 index 0902f993a0600aa08a39a638f4a75bea5d301415..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/mel.py +++ /dev/null @@ -1,179 +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 numpy as np # noqa: E402 - -from ....configuration_utils import ConfigMixin, register_to_config -from ....schedulers.scheduling_utils import SchedulerMixin - - -try: - import librosa # noqa: E402 - - _librosa_can_be_imported = True - _import_error = "" -except Exception as e: - _librosa_can_be_imported = False - _import_error = ( - f"Cannot import librosa because {e}. Make sure to correctly install librosa to be able to install it." - ) - - -from PIL import Image # noqa: E402 - - -class Mel(ConfigMixin, SchedulerMixin): - """ - Parameters: - x_res (`int`): - x resolution of spectrogram (time). - y_res (`int`): - y resolution of spectrogram (frequency bins). - sample_rate (`int`): - Sample rate of audio. - n_fft (`int`): - Number of Fast Fourier Transforms. - hop_length (`int`): - Hop length (a higher number is recommended if `y_res` < 256). - top_db (`int`): - Loudest decibel value. - n_iter (`int`): - Number of iterations for Griffin-Lim Mel inversion. - """ - - config_name = "mel_config.json" - - @register_to_config - def __init__( - self, - x_res: int = 256, - y_res: int = 256, - sample_rate: int = 22050, - n_fft: int = 2048, - hop_length: int = 512, - top_db: int = 80, - n_iter: int = 32, - ): - self.hop_length = hop_length - self.sr = sample_rate - self.n_fft = n_fft - self.top_db = top_db - self.n_iter = n_iter - self.set_resolution(x_res, y_res) - self.audio = None - - if not _librosa_can_be_imported: - raise ValueError(_import_error) - - def set_resolution(self, x_res: int, y_res: int): - """Set resolution. - - Args: - x_res (`int`): - x resolution of spectrogram (time). - y_res (`int`): - y resolution of spectrogram (frequency bins). - """ - self.x_res = x_res - self.y_res = y_res - self.n_mels = self.y_res - self.slice_size = self.x_res * self.hop_length - 1 - - def load_audio(self, audio_file: str = None, raw_audio: np.ndarray = None): - """Load audio. - - Args: - audio_file (`str`): - An audio file that must be on disk due to [Librosa](https://librosa.org/) limitation. - raw_audio (`np.ndarray`): - The raw audio file as a NumPy array. - """ - if audio_file is not None: - self.audio, _ = librosa.load(audio_file, mono=True, sr=self.sr) - else: - self.audio = raw_audio - - # Pad with silence if necessary. - if len(self.audio) < self.x_res * self.hop_length: - self.audio = np.concatenate([self.audio, np.zeros((self.x_res * self.hop_length - len(self.audio),))]) - - def get_number_of_slices(self) -> int: - """Get number of slices in audio. - - Returns: - `int`: - Number of spectograms audio can be sliced into. - """ - return len(self.audio) // self.slice_size - - def get_audio_slice(self, slice: int = 0) -> np.ndarray: - """Get slice of audio. - - Args: - slice (`int`): - Slice number of audio (out of `get_number_of_slices()`). - - Returns: - `np.ndarray`: - The audio slice as a NumPy array. - """ - return self.audio[self.slice_size * slice : self.slice_size * (slice + 1)] - - def get_sample_rate(self) -> int: - """Get sample rate. - - Returns: - `int`: - Sample rate of audio. - """ - return self.sr - - def audio_slice_to_image(self, slice: int) -> Image.Image: - """Convert slice of audio to spectrogram. - - Args: - slice (`int`): - Slice number of audio to convert (out of `get_number_of_slices()`). - - Returns: - `PIL Image`: - A grayscale image of `x_res x y_res`. - """ - S = librosa.feature.melspectrogram( - y=self.get_audio_slice(slice), sr=self.sr, n_fft=self.n_fft, hop_length=self.hop_length, n_mels=self.n_mels - ) - log_S = librosa.power_to_db(S, ref=np.max, top_db=self.top_db) - bytedata = (((log_S + self.top_db) * 255 / self.top_db).clip(0, 255) + 0.5).astype(np.uint8) - image = Image.fromarray(bytedata) - return image - - def image_to_audio(self, image: Image.Image) -> np.ndarray: - """Converts spectrogram to audio. - - Args: - image (`PIL Image`): - An grayscale image of `x_res x y_res`. - - Returns: - audio (`np.ndarray`): - The audio as a NumPy array. - """ - bytedata = np.frombuffer(image.tobytes(), dtype="uint8").reshape((image.height, image.width)) - log_S = bytedata.astype("float") * self.top_db / 255 - self.top_db - S = librosa.db_to_power(log_S) - audio = librosa.feature.inverse.mel_to_audio( - S, sr=self.sr, n_fft=self.n_fft, hop_length=self.hop_length, n_iter=self.n_iter - ) - return audio diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/pipeline_audio_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/pipeline_audio_diffusion.py deleted file mode 100644 index 81fa999eb1fb3849b0130761c66a2446148d0352..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/audio_diffusion/pipeline_audio_diffusion.py +++ /dev/null @@ -1,329 +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 math import acos, sin -from typing import List, Tuple, Union - -import numpy as np -import torch -from PIL import Image - -from ....models import AutoencoderKL, UNet2DConditionModel -from ....schedulers import DDIMScheduler, DDPMScheduler -from ....utils.torch_utils import randn_tensor -from ...pipeline_utils import AudioPipelineOutput, BaseOutput, DiffusionPipeline, ImagePipelineOutput -from .mel import Mel - - -class AudioDiffusionPipeline(DiffusionPipeline): - """ - Pipeline for audio diffusion. - - 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.). - - Parameters: - vqae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - mel ([`Mel`]): - Transform audio into a spectrogram. - scheduler ([`DDIMScheduler`] or [`DDPMScheduler`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of - [`DDIMScheduler`] or [`DDPMScheduler`]. - """ - - _optional_components = ["vqvae"] - - def __init__( - self, - vqvae: AutoencoderKL, - unet: UNet2DConditionModel, - mel: Mel, - scheduler: Union[DDIMScheduler, DDPMScheduler], - ): - super().__init__() - self.register_modules(unet=unet, scheduler=scheduler, mel=mel, vqvae=vqvae) - - def get_default_steps(self) -> int: - """Returns default number of steps recommended for inference. - - Returns: - `int`: - The number of steps. - """ - return 50 if isinstance(self.scheduler, DDIMScheduler) else 1000 - - @torch.no_grad() - def __call__( - self, - batch_size: int = 1, - audio_file: str = None, - raw_audio: np.ndarray = None, - slice: int = 0, - start_step: int = 0, - steps: int = None, - generator: torch.Generator = None, - mask_start_secs: float = 0, - mask_end_secs: float = 0, - step_generator: torch.Generator = None, - eta: float = 0, - noise: torch.Tensor = None, - encoding: torch.Tensor = None, - return_dict=True, - ) -> Union[ - Union[AudioPipelineOutput, ImagePipelineOutput], - Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]], - ]: - """ - The call function to the pipeline for generation. - - Args: - batch_size (`int`): - Number of samples to generate. - audio_file (`str`): - An audio file that must be on disk due to [Librosa](https://librosa.org/) limitation. - raw_audio (`np.ndarray`): - The raw audio file as a NumPy array. - slice (`int`): - Slice number of audio to convert. - start_step (int): - Step to start diffusion from. - steps (`int`): - Number of denoising steps (defaults to `50` for DDIM and `1000` for DDPM). - generator (`torch.Generator`): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - mask_start_secs (`float`): - Number of seconds of audio to mask (not generate) at start. - mask_end_secs (`float`): - Number of seconds of audio to mask (not generate) at end. - step_generator (`torch.Generator`): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) used to denoise. - None - eta (`float`): - 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. - noise (`torch.Tensor`): - A noise tensor of shape `(batch_size, 1, height, width)` or `None`. - encoding (`torch.Tensor`): - A tensor for [`UNet2DConditionModel`] of shape `(batch_size, seq_length, cross_attention_dim)`. - return_dict (`bool`): - Whether or not to return a [`AudioPipelineOutput`], [`ImagePipelineOutput`] or a plain tuple. - - Examples: - - For audio diffusion: - - ```py - import torch - from IPython.display import Audio - from diffusers import DiffusionPipeline - - device = "cuda" if torch.cuda.is_available() else "cpu" - pipe = DiffusionPipeline.from_pretrained("teticio/audio-diffusion-256").to(device) - - output = pipe() - display(output.images[0]) - display(Audio(output.audios[0], rate=mel.get_sample_rate())) - ``` - - For latent audio diffusion: - - ```py - import torch - from IPython.display import Audio - from diffusers import DiffusionPipeline - - device = "cuda" if torch.cuda.is_available() else "cpu" - pipe = DiffusionPipeline.from_pretrained("teticio/latent-audio-diffusion-256").to(device) - - output = pipe() - display(output.images[0]) - display(Audio(output.audios[0], rate=pipe.mel.get_sample_rate())) - ``` - - For other tasks like variation, inpainting, outpainting, etc: - - ```py - output = pipe( - raw_audio=output.audios[0, 0], - start_step=int(pipe.get_default_steps() / 2), - mask_start_secs=1, - mask_end_secs=1, - ) - display(output.images[0]) - display(Audio(output.audios[0], rate=pipe.mel.get_sample_rate())) - ``` - - Returns: - `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() - self.scheduler.set_timesteps(steps) - step_generator = step_generator or generator - # For backwards compatibility - if isinstance(self.unet.config.sample_size, int): - self.unet.config.sample_size = (self.unet.config.sample_size, self.unet.config.sample_size) - if noise is None: - noise = randn_tensor( - ( - batch_size, - self.unet.config.in_channels, - self.unet.config.sample_size[0], - self.unet.config.sample_size[1], - ), - generator=generator, - device=self.device, - ) - images = noise - mask = None - - if audio_file is not None or raw_audio is not None: - self.mel.load_audio(audio_file, raw_audio) - input_image = self.mel.audio_slice_to_image(slice) - input_image = np.frombuffer(input_image.tobytes(), dtype="uint8").reshape( - (input_image.height, input_image.width) - ) - input_image = (input_image / 255) * 2 - 1 - input_images = torch.tensor(input_image[np.newaxis, :, :], dtype=torch.float).to(self.device) - - if self.vqvae is not None: - input_images = self.vqvae.encode(torch.unsqueeze(input_images, 0)).latent_dist.sample( - generator=generator - )[0] - input_images = self.vqvae.config.scaling_factor * input_images - - if start_step > 0: - images[0, 0] = self.scheduler.add_noise(input_images, noise, self.scheduler.timesteps[start_step - 1]) - - pixels_per_second = ( - self.unet.config.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length - ) - mask_start = int(mask_start_secs * pixels_per_second) - mask_end = int(mask_end_secs * pixels_per_second) - mask = self.scheduler.add_noise(input_images, noise, torch.tensor(self.scheduler.timesteps[start_step:])) - - for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:])): - if isinstance(self.unet, UNet2DConditionModel): - model_output = self.unet(images, t, encoding)["sample"] - else: - model_output = self.unet(images, t)["sample"] - - if isinstance(self.scheduler, DDIMScheduler): - images = self.scheduler.step( - model_output=model_output, - timestep=t, - sample=images, - eta=eta, - generator=step_generator, - )["prev_sample"] - else: - images = self.scheduler.step( - model_output=model_output, - timestep=t, - sample=images, - generator=step_generator, - )["prev_sample"] - - if mask is not None: - if mask_start > 0: - images[:, :, :, :mask_start] = mask[:, step, :, :mask_start] - if mask_end > 0: - images[:, :, :, -mask_end:] = mask[:, step, :, -mask_end:] - - if self.vqvae is not None: - # 0.18215 was scaling factor used in training to ensure unit variance - images = 1 / self.vqvae.config.scaling_factor * images - images = self.vqvae.decode(images)["sample"] - - images = (images / 2 + 0.5).clamp(0, 1) - images = images.cpu().permute(0, 2, 3, 1).numpy() - images = (images * 255).round().astype("uint8") - images = list( - (Image.fromarray(_[:, :, 0]) for _ in images) - if images.shape[3] == 1 - else (Image.fromarray(_, mode="RGB").convert("L") for _ in images) - ) - - audios = [self.mel.image_to_audio(_) for _ in images] - if not return_dict: - return images, (self.mel.get_sample_rate(), audios) - - 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: - """ - Reverse the denoising step process to recover a noisy image from the generated image. - - Args: - images (`List[PIL Image]`): - List of images to encode. - steps (`int`): - Number of encoding steps to perform (defaults to `50`). - - Returns: - `np.ndarray`: - A noise tensor of shape `(batch_size, 1, height, width)`. - """ - - # Only works with DDIM as this method is deterministic - assert isinstance(self.scheduler, DDIMScheduler) - self.scheduler.set_timesteps(steps) - sample = np.array( - [np.frombuffer(image.tobytes(), dtype="uint8").reshape((1, image.height, image.width)) for image in images] - ) - sample = (sample / 255) * 2 - 1 - sample = torch.Tensor(sample).to(self.device) - - for t in self.progress_bar(torch.flip(self.scheduler.timesteps, (0,))): - prev_timestep = t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps - alpha_prod_t = self.scheduler.alphas_cumprod[t] - alpha_prod_t_prev = ( - self.scheduler.alphas_cumprod[prev_timestep] - if prev_timestep >= 0 - else self.scheduler.final_alpha_cumprod - ) - beta_prod_t = 1 - alpha_prod_t - model_output = self.unet(sample, t)["sample"] - pred_sample_direction = (1 - alpha_prod_t_prev) ** (0.5) * model_output - sample = (sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5) - sample = sample * alpha_prod_t ** (0.5) + beta_prod_t ** (0.5) * model_output - - return sample - - @staticmethod - def slerp(x0: torch.Tensor, x1: torch.Tensor, alpha: float) -> torch.Tensor: - """Spherical Linear intERPolation. - - Args: - x0 (`torch.Tensor`): - The first tensor to interpolate between. - x1 (`torch.Tensor`): - Second tensor to interpolate between. - alpha (`float`): - Interpolation between 0 and 1 - - Returns: - `torch.Tensor`: - The interpolated tensor. - """ - - theta = acos(torch.dot(torch.flatten(x0), torch.flatten(x1)) / torch.norm(x0) / torch.norm(x1)) - return sin((1 - alpha) * theta) * x0 / sin(theta) + sin(alpha * theta) * x1 / sin(theta) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/latent_diffusion_uncond/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/latent_diffusion_uncond/__init__.py deleted file mode 100644 index 214f5bbca969f9ae0629578c72aaf339f86ded88..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/latent_diffusion_uncond/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -from typing import TYPE_CHECKING - -from ....utils import DIFFUSERS_SLOW_IMPORT, _LazyModule - - -_import_structure = {"pipeline_latent_diffusion_uncond": ["LDMPipeline"]} - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - from .pipeline_latent_diffusion_uncond import LDMPipeline -else: - import sys - - sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/latent_diffusion_uncond/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/latent_diffusion_uncond/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 2f7eea7bfa7b9c05fc698c947a97859c059d984f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/latent_diffusion_uncond/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/latent_diffusion_uncond/__pycache__/pipeline_latent_diffusion_uncond.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/latent_diffusion_uncond/__pycache__/pipeline_latent_diffusion_uncond.cpython-310.pyc deleted file mode 100644 index ec81fdc1c83c2042d36c6c727dd127a9215031fb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/latent_diffusion_uncond/__pycache__/pipeline_latent_diffusion_uncond.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py deleted file mode 100644 index 0bb24ed0b1cec2f15cd6f9c951fe5c2b413b61f5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py +++ /dev/null @@ -1,130 +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 inspect -from typing import List, Optional, Tuple, Union - -import torch - -from ....models import UNet2DModel, VQModel -from ....schedulers import DDIMScheduler -from ....utils.torch_utils import randn_tensor -from ...pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -class LDMPipeline(DiffusionPipeline): - r""" - Pipeline for unconditional image generation using latent diffusion. - - 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.). - - Parameters: - vqvae ([`VQModel`]): - Vector-quantized (VQ) model to encode and decode images to and from latent representations. - unet ([`UNet2DModel`]): - A `UNet2DModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - [`DDIMScheduler`] is used in combination with `unet` to denoise the encoded image latents. - """ - - def __init__(self, vqvae: VQModel, unet: UNet2DModel, scheduler: DDIMScheduler): - super().__init__() - self.register_modules(vqvae=vqvae, unet=unet, scheduler=scheduler) - - @torch.no_grad() - def __call__( - self, - batch_size: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - eta: float = 0.0, - num_inference_steps: int = 50, - output_type: Optional[str] = "pil", - return_dict: bool = True, - **kwargs, - ) -> Union[Tuple, ImagePipelineOutput]: - r""" - The call function to the pipeline for generation. - - Args: - batch_size (`int`, *optional*, defaults to 1): - Number of images to generate. - generator (`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 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - 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 [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Example: - - ```py - >>> from diffusers import LDMPipeline - - >>> # load model and scheduler - >>> pipe = LDMPipeline.from_pretrained("CompVis/ldm-celebahq-256") - - >>> # run pipeline in inference (sample random noise and denoise) - >>> image = pipe().images[0] - ``` - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images - """ - - latents = randn_tensor( - (batch_size, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size), - generator=generator, - ) - latents = latents.to(self.device) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - - self.scheduler.set_timesteps(num_inference_steps) - - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - - extra_kwargs = {} - if accepts_eta: - extra_kwargs["eta"] = eta - - for t in self.progress_bar(self.scheduler.timesteps): - latent_model_input = self.scheduler.scale_model_input(latents, t) - # predict the noise residual - noise_prediction = self.unet(latent_model_input, t).sample - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_prediction, t, latents, **extra_kwargs).prev_sample - - # adjust latents with inverse of vae scale - latents = latents / self.vqvae.config.scaling_factor - # decode the image latents with the VAE - image = self.vqvae.decode(latents).sample - - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).numpy() - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/pndm/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/pndm/__init__.py deleted file mode 100644 index 5e3bdba74079d77576655e22b43014a0438a9c2e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/pndm/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -from typing import TYPE_CHECKING - -from ....utils import DIFFUSERS_SLOW_IMPORT, _LazyModule - - -_import_structure = {"pipeline_pndm": ["PNDMPipeline"]} - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - from .pipeline_pndm import PNDMPipeline -else: - import sys - - sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/pndm/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/pndm/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 424e26d6dace8fd7120754ccc7eaf4a9d4a0afe0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/pndm/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/pndm/__pycache__/pipeline_pndm.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/pndm/__pycache__/pipeline_pndm.cpython-310.pyc deleted file mode 100644 index 12a704bce12dde806b5bb6218568506e51d97a38..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/pndm/__pycache__/pipeline_pndm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/pndm/pipeline_pndm.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/pndm/pipeline_pndm.py deleted file mode 100644 index 71e3e156e0e4a04665f1643eada7923b73be5619..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/pndm/pipeline_pndm.py +++ /dev/null @@ -1,121 +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 List, Optional, Tuple, Union - -import torch - -from ....models import UNet2DModel -from ....schedulers import PNDMScheduler -from ....utils.torch_utils import randn_tensor -from ...pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -class PNDMPipeline(DiffusionPipeline): - r""" - Pipeline for unconditional image generation. - - 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.). - - Parameters: - unet ([`UNet2DModel`]): - A `UNet2DModel` to denoise the encoded image latents. - scheduler ([`PNDMScheduler`]): - A `PNDMScheduler` to be used in combination with `unet` to denoise the encoded image. - """ - - unet: UNet2DModel - scheduler: PNDMScheduler - - def __init__(self, unet: UNet2DModel, scheduler: PNDMScheduler): - super().__init__() - - scheduler = PNDMScheduler.from_config(scheduler.config) - - self.register_modules(unet=unet, scheduler=scheduler) - - @torch.no_grad() - 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", - return_dict: bool = True, - **kwargs, - ) -> Union[ImagePipelineOutput, Tuple]: - r""" - The call function to the pipeline for generation. - - Args: - batch_size (`int`, `optional`, defaults to 1): - The number of images to generate. - 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`, `optional`): - A [`torch.Generator`](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 generated image. Choose between `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`ImagePipelineOutput`] instead of a plain tuple. - - Example: - - ```py - >>> from diffusers import PNDMPipeline - - >>> # load model and scheduler - >>> pndm = PNDMPipeline.from_pretrained("google/ddpm-cifar10-32") - - >>> # run pipeline in inference (sample random noise and denoise) - >>> image = pndm().images[0] - - >>> # save image - >>> image.save("pndm_generated_image.png") - ``` - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images. - """ - # For more information on the sampling method you can take a look at Algorithm 2 of - # the official paper: https://huggingface.co/papers/2202.09778 - - # Sample gaussian noise to begin loop - image = randn_tensor( - (batch_size, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size), - generator=generator, - device=self.device, - ) - - self.scheduler.set_timesteps(num_inference_steps) - for t in self.progress_bar(self.scheduler.timesteps): - model_output = self.unet(image, t).sample - - image = self.scheduler.step(model_output, t, image).prev_sample - - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).numpy() - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/repaint/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/repaint/__init__.py deleted file mode 100644 index 2c6b04af52d40e8a2bfa2aa5812b9fb8b1da06f5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/repaint/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -from typing import TYPE_CHECKING - -from ....utils import DIFFUSERS_SLOW_IMPORT, _LazyModule - - -_import_structure = {"pipeline_repaint": ["RePaintPipeline"]} - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - from .pipeline_repaint import RePaintPipeline - -else: - import sys - - sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/repaint/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/repaint/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 962bd6981fd0f5edceede7f97453a844e7a5dc15..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/repaint/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/repaint/__pycache__/pipeline_repaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/repaint/__pycache__/pipeline_repaint.cpython-310.pyc deleted file mode 100644 index b01422cf50e4bb5f51c09021ec2e296daaef7a74..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/repaint/__pycache__/pipeline_repaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/repaint/pipeline_repaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/repaint/pipeline_repaint.py deleted file mode 100644 index 56c6007ae886f96d0f5142a655f59588a8d91efc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/repaint/pipeline_repaint.py +++ /dev/null @@ -1,231 +0,0 @@ -# Copyright 2025 ETH Zurich Computer Vision Lab 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 List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch - -from ....models import UNet2DModel -from ....schedulers import RePaintScheduler -from ....utils import PIL_INTERPOLATION, deprecate, logging -from ....utils.torch_utils import randn_tensor -from ...pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess -def _preprocess_image(image: Union[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): - return image - elif isinstance(image, PIL.Image.Image): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - w, h = image[0].size - w, h = (x - x % 8 for x in (w, h)) # resize to integer multiple of 8 - - image = [np.array(i.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image] - image = np.concatenate(image, axis=0) - image = np.array(image).astype(np.float32) / 255.0 - image = image.transpose(0, 3, 1, 2) - image = 2.0 * image - 1.0 - image = torch.from_numpy(image) - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, dim=0) - return image - - -def _preprocess_mask(mask: Union[List, PIL.Image.Image, torch.Tensor]): - if isinstance(mask, torch.Tensor): - return mask - elif isinstance(mask, PIL.Image.Image): - mask = [mask] - - if isinstance(mask[0], PIL.Image.Image): - w, h = mask[0].size - w, h = (x - x % 32 for x in (w, h)) # resize to integer multiple of 32 - mask = [np.array(m.convert("L").resize((w, h), resample=PIL_INTERPOLATION["nearest"]))[None, :] for m in mask] - mask = np.concatenate(mask, axis=0) - mask = mask.astype(np.float32) / 255.0 - mask[mask < 0.5] = 0 - mask[mask >= 0.5] = 1 - mask = torch.from_numpy(mask) - elif isinstance(mask[0], torch.Tensor): - mask = torch.cat(mask, dim=0) - return mask - - -class RePaintPipeline(DiffusionPipeline): - r""" - Pipeline for image inpainting using RePaint. - - 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.). - - Parameters: - unet ([`UNet2DModel`]): - A `UNet2DModel` to denoise the encoded image latents. - scheduler ([`RePaintScheduler`]): - A `RePaintScheduler` to be used in combination with `unet` to denoise the encoded image. - """ - - unet: UNet2DModel - scheduler: RePaintScheduler - model_cpu_offload_seq = "unet" - - def __init__(self, unet: UNet2DModel, scheduler: RePaintScheduler): - super().__init__() - self.register_modules(unet=unet, scheduler=scheduler) - - @torch.no_grad() - def __call__( - self, - image: Union[torch.Tensor, PIL.Image.Image], - mask_image: Union[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", - return_dict: bool = True, - ) -> Union[ImagePipelineOutput, Tuple]: - r""" - The call function to the pipeline for generation. - - Args: - image (`torch.Tensor` or `PIL.Image.Image`): - The original image to inpaint on. - mask_image (`torch.Tensor` or `PIL.Image.Image`): - The mask_image where 0.0 define which part of the original image to inpaint. - num_inference_steps (`int`, *optional*, defaults to 1000): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - eta (`float`): - The weight of the added noise in a diffusion step. Its value is between 0.0 and 1.0; 0.0 corresponds to - DDIM and 1.0 is the DDPM scheduler. - jump_length (`int`, *optional*, defaults to 10): - The number of steps taken forward in time before going backward in time for a single jump ("j" in - RePaint paper). Take a look at Figure 9 and 10 in the - [paper](https://huggingface.co/papers/2201.09865). - jump_n_sample (`int`, *optional*, defaults to 10): - The number of times to make a forward time jump for a given chosen time sample. Take a look at Figure 9 - and 10 in the [paper](https://huggingface.co/papers/2201.09865). - generator (`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"`): - 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 [`ImagePipelineOutput`] instead of a plain tuple. - - Example: - - ```py - >>> from io import BytesIO - >>> import torch - >>> import PIL - >>> import requests - >>> from diffusers import RePaintPipeline, RePaintScheduler - - - >>> def download_image(url): - ... response = requests.get(url) - ... return PIL.Image.open(BytesIO(response.content)).convert("RGB") - - - >>> img_url = "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/celeba_hq_256.png" - >>> mask_url = "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/mask_256.png" - - >>> # Load the original image and the mask as PIL images - >>> original_image = download_image(img_url).resize((256, 256)) - >>> mask_image = download_image(mask_url).resize((256, 256)) - - >>> # Load the RePaint scheduler and pipeline based on a pretrained DDPM model - >>> scheduler = RePaintScheduler.from_pretrained("google/ddpm-ema-celebahq-256") - >>> pipe = RePaintPipeline.from_pretrained("google/ddpm-ema-celebahq-256", scheduler=scheduler) - >>> pipe = pipe.to("cuda") - - >>> generator = torch.Generator(device="cuda").manual_seed(0) - >>> output = pipe( - ... image=original_image, - ... mask_image=mask_image, - ... num_inference_steps=250, - ... eta=0.0, - ... jump_length=10, - ... jump_n_sample=10, - ... generator=generator, - ... ) - >>> inpainted_image = output.images[0] - ``` - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images. - """ - - original_image = image - - original_image = _preprocess_image(original_image) - original_image = original_image.to(device=self._execution_device, dtype=self.unet.dtype) - mask_image = _preprocess_mask(mask_image) - mask_image = mask_image.to(device=self._execution_device, dtype=self.unet.dtype) - - batch_size = original_image.shape[0] - - # sample gaussian noise to begin the loop - 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." - ) - - image_shape = original_image.shape - image = randn_tensor(image_shape, generator=generator, device=self._execution_device, dtype=self.unet.dtype) - - # set step values - self.scheduler.set_timesteps(num_inference_steps, jump_length, jump_n_sample, self._execution_device) - self.scheduler.eta = eta - - t_last = self.scheduler.timesteps[0] + 1 - generator = generator[0] if isinstance(generator, list) else generator - for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)): - if t < t_last: - # predict the noise residual - model_output = self.unet(image, t).sample - # compute previous image: x_t -> x_t-1 - image = self.scheduler.step(model_output, t, image, original_image, mask_image, generator).prev_sample - - else: - # compute the reverse: x_t-1 -> x_t - image = self.scheduler.undo_step(image, t_last, generator) - t_last = t - - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).numpy() - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/score_sde_ve/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/score_sde_ve/__init__.py deleted file mode 100644 index 87c167c3dbd26e0408a41ef197a42dc5eb7038d7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/score_sde_ve/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -from typing import TYPE_CHECKING - -from ....utils import DIFFUSERS_SLOW_IMPORT, _LazyModule - - -_import_structure = {"pipeline_score_sde_ve": ["ScoreSdeVePipeline"]} - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - from .pipeline_score_sde_ve import ScoreSdeVePipeline - -else: - import sys - - sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/score_sde_ve/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/score_sde_ve/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 01cbdf40c96050b451ed7cee1f6d80cec40afe28..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/score_sde_ve/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/score_sde_ve/__pycache__/pipeline_score_sde_ve.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/score_sde_ve/__pycache__/pipeline_score_sde_ve.cpython-310.pyc deleted file mode 100644 index 7f93a33e0dcd0008b9fde6a6a46ef29c9b29cce0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/score_sde_ve/__pycache__/pipeline_score_sde_ve.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/score_sde_ve/pipeline_score_sde_ve.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/score_sde_ve/pipeline_score_sde_ve.py deleted file mode 100644 index 3f04db7ad6990863edcc7b19c905475a07171cae..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/score_sde_ve/pipeline_score_sde_ve.py +++ /dev/null @@ -1,109 +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 List, Optional, Tuple, Union - -import torch - -from ....models import UNet2DModel -from ....schedulers import ScoreSdeVeScheduler -from ....utils.torch_utils import randn_tensor -from ...pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -class ScoreSdeVePipeline(DiffusionPipeline): - r""" - Pipeline for unconditional image generation. - - 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.). - - Parameters: - unet ([`UNet2DModel`]): - A `UNet2DModel` to denoise the encoded image. - scheduler ([`ScoreSdeVeScheduler`]): - A `ScoreSdeVeScheduler` to be used in combination with `unet` to denoise the encoded image. - """ - - unet: UNet2DModel - scheduler: ScoreSdeVeScheduler - - def __init__(self, unet: UNet2DModel, scheduler: ScoreSdeVeScheduler): - super().__init__() - self.register_modules(unet=unet, scheduler=scheduler) - - @torch.no_grad() - 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", - return_dict: bool = True, - **kwargs, - ) -> Union[ImagePipelineOutput, Tuple]: - r""" - The call function to the pipeline for generation. - - Args: - batch_size (`int`, *optional*, defaults to 1): - The number of images to generate. - generator (`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"`): - 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 [`ImagePipelineOutput`] instead of a plain tuple. - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images. - """ - - img_size = self.unet.config.sample_size - shape = (batch_size, 3, img_size, img_size) - - model = self.unet - - sample = randn_tensor(shape, generator=generator) * self.scheduler.init_noise_sigma - sample = sample.to(self.device) - - self.scheduler.set_timesteps(num_inference_steps) - self.scheduler.set_sigmas(num_inference_steps) - - for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)): - sigma_t = self.scheduler.sigmas[i] * torch.ones(shape[0], device=self.device) - - # correction step - for _ in range(self.scheduler.config.correct_steps): - model_output = self.unet(sample, sigma_t).sample - sample = self.scheduler.step_correct(model_output, sample, generator=generator).prev_sample - - # prediction step - model_output = model(sample, sigma_t).sample - output = self.scheduler.step_pred(model_output, t, sample, generator=generator) - - sample, sample_mean = output.prev_sample, output.prev_sample_mean - - sample = sample_mean.clamp(0, 1) - sample = sample.cpu().permute(0, 2, 3, 1).numpy() - if output_type == "pil": - sample = self.numpy_to_pil(sample) - - if not return_dict: - return (sample,) - - return ImagePipelineOutput(images=sample) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__init__.py deleted file mode 100644 index 150954baa0eb8f8a7216b4891effc14a71e21b1b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__init__.py +++ /dev/null @@ -1,75 +0,0 @@ -# flake8: noqa -from typing import TYPE_CHECKING -from ....utils import ( - DIFFUSERS_SLOW_IMPORT, - _LazyModule, - is_note_seq_available, - OptionalDependencyNotAvailable, - is_torch_available, - is_transformers_available, - get_objects_from_module, -) - -_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["continous_encoder"] = ["SpectrogramContEncoder"] - _import_structure["notes_encoder"] = ["SpectrogramNotesEncoder"] - _import_structure["pipeline_spectrogram_diffusion"] = [ - "SpectrogramContEncoder", - "SpectrogramDiffusionPipeline", - "T5FilmDecoder", - ] -try: - if not (is_transformers_available() and is_torch_available() and is_note_seq_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ....utils import dummy_transformers_and_torch_and_note_seq_objects - - _dummy_objects.update(get_objects_from_module(dummy_transformers_and_torch_and_note_seq_objects)) -else: - _import_structure["midi_utils"] = ["MidiProcessor"] - - -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_spectrogram_diffusion import SpectrogramDiffusionPipeline - from .pipeline_spectrogram_diffusion import SpectrogramContEncoder - from .pipeline_spectrogram_diffusion import SpectrogramNotesEncoder - from .pipeline_spectrogram_diffusion import T5FilmDecoder - - try: - if not (is_transformers_available() and is_torch_available() and is_note_seq_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ....utils.dummy_transformers_and_torch_and_note_seq_objects import * - - else: - from .midi_utils import MidiProcessor - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 9562c896282ba054a26c10e5874d5514ab0178da..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/continuous_encoder.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/continuous_encoder.cpython-310.pyc deleted file mode 100644 index 5630024df5530c5d73e6b1a7dda250ac9513f83d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/continuous_encoder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/midi_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/midi_utils.cpython-310.pyc deleted file mode 100644 index 46d628fba3594de3b0dfc6fd495adf39584a9819..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/midi_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/notes_encoder.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/notes_encoder.cpython-310.pyc deleted file mode 100644 index 4cc84b34d70801bc50bc8061b10cd5eec0d5437a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/notes_encoder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/pipeline_spectrogram_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/pipeline_spectrogram_diffusion.cpython-310.pyc deleted file mode 100644 index 2f984e97d620665450a5f4affb044a695281e9af..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/__pycache__/pipeline_spectrogram_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/continuous_encoder.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/continuous_encoder.py deleted file mode 100644 index b26e84f728693b29f62cc058c43810fc4983619b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/continuous_encoder.py +++ /dev/null @@ -1,92 +0,0 @@ -# Copyright 2022 The Music Spectrogram Diffusion Authors. -# 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 torch -import torch.nn as nn -from transformers.modeling_utils import ModuleUtilsMixin -from transformers.models.t5.modeling_t5 import ( - T5Block, - T5Config, - T5LayerNorm, -) - -from ....configuration_utils import ConfigMixin, register_to_config -from ....models import ModelMixin - - -class SpectrogramContEncoder(ModelMixin, ConfigMixin, ModuleUtilsMixin): - @register_to_config - def __init__( - self, - input_dims: int, - targets_context_length: int, - d_model: int, - dropout_rate: float, - num_layers: int, - num_heads: int, - d_kv: int, - d_ff: int, - feed_forward_proj: str, - is_decoder: bool = False, - ): - super().__init__() - - self.input_proj = nn.Linear(input_dims, d_model, bias=False) - - self.position_encoding = nn.Embedding(targets_context_length, d_model) - self.position_encoding.weight.requires_grad = False - - self.dropout_pre = nn.Dropout(p=dropout_rate) - - t5config = T5Config( - d_model=d_model, - num_heads=num_heads, - d_kv=d_kv, - d_ff=d_ff, - feed_forward_proj=feed_forward_proj, - dropout_rate=dropout_rate, - is_decoder=is_decoder, - is_encoder_decoder=False, - ) - self.encoders = nn.ModuleList() - for lyr_num in range(num_layers): - lyr = T5Block(t5config) - self.encoders.append(lyr) - - self.layer_norm = T5LayerNorm(d_model) - self.dropout_post = nn.Dropout(p=dropout_rate) - - def forward(self, encoder_inputs, encoder_inputs_mask): - x = self.input_proj(encoder_inputs) - - # terminal relative positional encodings - max_positions = encoder_inputs.shape[1] - input_positions = torch.arange(max_positions, device=encoder_inputs.device) - - seq_lens = encoder_inputs_mask.sum(-1) - input_positions = torch.roll(input_positions.unsqueeze(0), tuple(seq_lens.tolist()), dims=0) - x += self.position_encoding(input_positions) - - x = self.dropout_pre(x) - - # inverted the attention mask - input_shape = encoder_inputs.size() - extended_attention_mask = self.get_extended_attention_mask(encoder_inputs_mask, input_shape) - - for lyr in self.encoders: - x = lyr(x, extended_attention_mask)[0] - x = self.layer_norm(x) - - return self.dropout_post(x), encoder_inputs_mask diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/midi_utils.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/midi_utils.py deleted file mode 100644 index 8985a6f8880076bd11d9254446a98ac8e7d345e4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/midi_utils.py +++ /dev/null @@ -1,667 +0,0 @@ -# Copyright 2022 The Music Spectrogram Diffusion Authors. -# 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 dataclasses -import math -import os -from typing import Any, Callable, List, Mapping, MutableMapping, Optional, Sequence, Tuple, Union - -import numpy as np -import torch -import torch.nn.functional as F - -from ....utils import is_note_seq_available -from .pipeline_spectrogram_diffusion import TARGET_FEATURE_LENGTH - - -if is_note_seq_available(): - import note_seq -else: - raise ImportError("Please install note-seq via `pip install note-seq`") - - -INPUT_FEATURE_LENGTH = 2048 - -SAMPLE_RATE = 16000 -HOP_SIZE = 320 -FRAME_RATE = int(SAMPLE_RATE // HOP_SIZE) - -DEFAULT_STEPS_PER_SECOND = 100 -DEFAULT_MAX_SHIFT_SECONDS = 10 -DEFAULT_NUM_VELOCITY_BINS = 1 - -SLAKH_CLASS_PROGRAMS = { - "Acoustic Piano": 0, - "Electric Piano": 4, - "Chromatic Percussion": 8, - "Organ": 16, - "Acoustic Guitar": 24, - "Clean Electric Guitar": 26, - "Distorted Electric Guitar": 29, - "Acoustic Bass": 32, - "Electric Bass": 33, - "Violin": 40, - "Viola": 41, - "Cello": 42, - "Contrabass": 43, - "Orchestral Harp": 46, - "Timpani": 47, - "String Ensemble": 48, - "Synth Strings": 50, - "Choir and Voice": 52, - "Orchestral Hit": 55, - "Trumpet": 56, - "Trombone": 57, - "Tuba": 58, - "French Horn": 60, - "Brass Section": 61, - "Soprano/Alto Sax": 64, - "Tenor Sax": 66, - "Baritone Sax": 67, - "Oboe": 68, - "English Horn": 69, - "Bassoon": 70, - "Clarinet": 71, - "Pipe": 73, - "Synth Lead": 80, - "Synth Pad": 88, -} - - -@dataclasses.dataclass -class NoteRepresentationConfig: - """Configuration note representations.""" - - onsets_only: bool - include_ties: bool - - -@dataclasses.dataclass -class NoteEventData: - pitch: int - velocity: Optional[int] = None - program: Optional[int] = None - is_drum: Optional[bool] = None - instrument: Optional[int] = None - - -@dataclasses.dataclass -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) - - -@dataclasses.dataclass -class EventRange: - type: str - min_value: int - max_value: int - - -@dataclasses.dataclass -class Event: - type: str - value: int - - -class Tokenizer: - def __init__(self, regular_ids: int): - # The special tokens: 0=PAD, 1=EOS, and 2=UNK - self._num_special_tokens = 3 - self._num_regular_tokens = regular_ids - - def encode(self, token_ids): - encoded = [] - for token_id in token_ids: - if not 0 <= token_id < self._num_regular_tokens: - raise ValueError( - f"token_id {token_id} does not fall within valid range of [0, {self._num_regular_tokens})" - ) - encoded.append(token_id + self._num_special_tokens) - - # Add EOS token - encoded.append(1) - - # Pad to till INPUT_FEATURE_LENGTH - encoded = encoded + [0] * (INPUT_FEATURE_LENGTH - len(encoded)) - - return encoded - - -class Codec: - """Encode and decode events. - - Useful for declaring what certain ranges of a vocabulary should be used for. This is intended to be used from - Python before encoding or after decoding with GenericTokenVocabulary. This class is more lightweight and does not - include things like EOS or UNK token handling. - - To ensure that 'shift' events are always the first block of the vocab and start at 0, that event type is required - and specified separately. - """ - - def __init__(self, max_shift_steps: int, steps_per_second: float, event_ranges: List[EventRange]): - """Define Codec. - - Args: - max_shift_steps: Maximum number of shift steps that can be encoded. - steps_per_second: Shift steps will be interpreted as having a duration of - 1 / steps_per_second. - event_ranges: Other supported event types and their ranges. - """ - self.steps_per_second = steps_per_second - self._shift_range = EventRange(type="shift", min_value=0, max_value=max_shift_steps) - self._event_ranges = [self._shift_range] + event_ranges - # Ensure all event types have unique names. - assert len(self._event_ranges) == len({er.type for er in self._event_ranges}) - - @property - def num_classes(self) -> int: - return sum(er.max_value - er.min_value + 1 for er in self._event_ranges) - - # The next couple methods are simplified special case methods just for shift - # events that are intended to be used from within autograph functions. - - def is_shift_event_index(self, index: int) -> bool: - return (self._shift_range.min_value <= index) and (index <= self._shift_range.max_value) - - @property - def max_shift_steps(self) -> int: - return self._shift_range.max_value - - def encode_event(self, event: Event) -> int: - """Encode an event to an index.""" - offset = 0 - for er in self._event_ranges: - if event.type == er.type: - if not er.min_value <= event.value <= er.max_value: - raise ValueError( - f"Event value {event.value} is not within valid range " - f"[{er.min_value}, {er.max_value}] for type {event.type}" - ) - return offset + event.value - er.min_value - offset += er.max_value - er.min_value + 1 - - raise ValueError(f"Unknown event type: {event.type}") - - 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: - if event_type == er.type: - return offset, offset + (er.max_value - er.min_value) - offset += er.max_value - er.min_value + 1 - - raise ValueError(f"Unknown event type: {event_type}") - - def decode_event_index(self, index: int) -> Event: - """Decode an event index to an Event.""" - offset = 0 - for er in self._event_ranges: - if offset <= index <= offset + er.max_value - er.min_value: - return Event(type=er.type, value=er.min_value + index - offset) - offset += er.max_value - er.min_value + 1 - - raise ValueError(f"Unknown event index: {index}") - - -@dataclasses.dataclass -class ProgramGranularity: - # both tokens_map_fn and program_map_fn should be idempotent - tokens_map_fn: Callable[[Sequence[int], Codec], Sequence[int]] - program_map_fn: Callable[[int], int] - - -def drop_programs(tokens, codec: Codec): - """Drops program change events from a token sequence.""" - min_program_id, max_program_id = codec.event_type_range("program") - return tokens[(tokens < min_program_id) | (tokens > max_program_id)] - - -def programs_to_midi_classes(tokens, codec): - """Modifies program events to be the first program in the MIDI class.""" - min_program_id, max_program_id = codec.event_type_range("program") - is_program = (tokens >= min_program_id) & (tokens <= max_program_id) - return np.where(is_program, min_program_id + 8 * ((tokens - min_program_id) // 8), tokens) - - -PROGRAM_GRANULARITIES = { - # "flat" granularity; drop program change tokens and set NoteSequence - # programs to zero - "flat": ProgramGranularity(tokens_map_fn=drop_programs, program_map_fn=lambda program: 0), - # map each program to the first program in its MIDI class - "midi_class": ProgramGranularity( - tokens_map_fn=programs_to_midi_classes, program_map_fn=lambda program: 8 * (program // 8) - ), - # leave programs as is - "full": ProgramGranularity(tokens_map_fn=lambda tokens, codec: tokens, program_map_fn=lambda program: program), -} - - -def frame(signal, frame_length, frame_step, pad_end=False, pad_value=0, axis=-1): - """ - equivalent of tf.signal.frame - """ - signal_length = signal.shape[axis] - if pad_end: - frames_overlap = frame_length - frame_step - rest_samples = np.abs(signal_length - frames_overlap) % np.abs(frame_length - frames_overlap) - pad_size = int(frame_length - rest_samples) - - if pad_size != 0: - pad_axis = [0] * signal.ndim - pad_axis[axis] = pad_size - signal = F.pad(signal, pad_axis, "constant", pad_value) - frames = signal.unfold(axis, frame_length, frame_step) - return frames - - -def program_to_slakh_program(program): - # this is done very hackily, probably should use a custom mapping - for slakh_program in sorted(SLAKH_CLASS_PROGRAMS.values(), reverse=True): - if program >= slakh_program: - return slakh_program - - -def audio_to_frames( - samples, - hop_size: int, - frame_rate: int, -) -> 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") - - # Split audio into frames. - frames = frame( - torch.Tensor(samples).unsqueeze(0), - frame_length=frame_size, - frame_step=frame_size, - pad_end=False, # TODO check why its off by 1 here when True - ) - - num_frames = len(samples) // frame_size - - times = np.arange(num_frames) / frame_rate - return frames, times - - -def note_sequence_to_onsets_and_offsets_and_programs( - ns: note_seq.NoteSequence, -) -> 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. - - Args: - ns: NoteSequence from which to extract onsets and offsets. - - Returns: - times: A list of note onset and offset times. values: A list of NoteEventData objects where velocity is zero for - note - offsets. - """ - # Sort by program and pitch and put offsets before onsets as a tiebreaker for - # subsequent stable sort. - notes = sorted(ns.notes, key=lambda note: (note.is_drum, note.program, note.pitch)) - times = [note.end_time for note in notes if not note.is_drum] + [note.start_time for note in notes] - values = [ - NoteEventData(pitch=note.pitch, velocity=0, program=note.program, is_drum=False) - for note in notes - if not note.is_drum - ] + [ - NoteEventData(pitch=note.pitch, velocity=note.velocity, program=note.program, is_drum=note.is_drum) - for note in notes - ] - return times, values - - -def num_velocity_bins_from_codec(codec: Codec): - """Get number of velocity bins from event codec.""" - lo, hi = codec.event_type_range("velocity") - return hi - lo - - -# segment an array into segments of length n -def segment(a, n): - return [a[i : i + n] for i in range(0, len(a), n)] - - -def velocity_to_bin(velocity, num_velocity_bins): - if velocity == 0: - return 0 - else: - return math.ceil(num_velocity_bins * velocity / note_seq.MAX_MIDI_VELOCITY) - - -def note_event_data_to_events( - state: Optional[NoteEncodingState], - value: NoteEventData, - codec: Codec, -) -> Sequence[Event]: - """Convert note event data to a sequence of events.""" - if value.velocity is None: - # onsets only, no program or velocity - return [Event("pitch", value.pitch)] - else: - num_velocity_bins = num_velocity_bins_from_codec(codec) - velocity_bin = velocity_to_bin(value.velocity, num_velocity_bins) - if value.program is None: - # onsets + offsets + velocities only, no programs - if state is not None: - state.active_pitches[(value.pitch, 0)] = velocity_bin - return [Event("velocity", velocity_bin), Event("pitch", value.pitch)] - else: - if value.is_drum: - # drum events use a separate vocabulary - return [Event("velocity", velocity_bin), Event("drum", value.pitch)] - else: - # program + velocity + pitch - if state is not None: - state.active_pitches[(value.pitch, value.program)] = velocity_bin - return [ - Event("program", value.program), - Event("velocity", velocity_bin), - Event("pitch", value.pitch), - ] - - -def note_encoding_state_to_events(state: NoteEncodingState) -> Sequence[Event]: - """Output program and pitch events for active notes plus a final tie event.""" - events = [] - for pitch, program in sorted(state.active_pitches.keys(), key=lambda k: k[::-1]): - if state.active_pitches[(pitch, program)]: - events += [Event("program", program), Event("pitch", pitch)] - events.append(Event("tie", 0)) - return events - - -def encode_and_index_events( - state, event_times, event_values, codec, frame_times, encode_event_fn, encoding_state_to_events_fn=None -): - """Encode a sequence of timed events and index to audio frame times. - - Encodes time shifts as repeated single step shifts for later run length encoding. - - Optionally, also encodes a sequence of "state events", keeping track of the current encoding state at each audio - frame. This can be used e.g. to prepend events representing the current state to a targets segment. - - Args: - state: Initial event encoding state. - event_times: Sequence of event times. - event_values: Sequence of event values. - encode_event_fn: Function that transforms event value into a sequence of one - or more Event objects. - codec: An Codec object that maps Event objects to indices. - frame_times: Time for every audio frame. - encoding_state_to_events_fn: Function that transforms encoding state into a - sequence of one or more Event objects. - - Returns: - events: Encoded events and shifts. event_start_indices: Corresponding start event index for every audio frame. - Note: one event can correspond to multiple audio indices due to sampling rate differences. This makes - splitting sequences tricky because the same event can appear at the end of one sequence and the beginning of - another. - event_end_indices: Corresponding end event index for every audio frame. Used - to ensure when slicing that one chunk ends where the next begins. Should always be true that - event_end_indices[i] = event_start_indices[i + 1]. - state_events: Encoded "state" events representing the encoding state before - each event. - state_event_indices: Corresponding state event index for every audio frame. - """ - indices = np.argsort(event_times, kind="stable") - event_steps = [round(event_times[i] * codec.steps_per_second) for i in indices] - event_values = [event_values[i] for i in indices] - - events = [] - state_events = [] - event_start_indices = [] - state_event_indices = [] - - cur_step = 0 - cur_event_idx = 0 - cur_state_event_idx = 0 - - def fill_event_start_indices_to_cur_step(): - while ( - len(event_start_indices) < len(frame_times) - and frame_times[len(event_start_indices)] < cur_step / codec.steps_per_second - ): - event_start_indices.append(cur_event_idx) - state_event_indices.append(cur_state_event_idx) - - for event_step, event_value in zip(event_steps, event_values): - while event_step > cur_step: - events.append(codec.encode_event(Event(type="shift", value=1))) - cur_step += 1 - fill_event_start_indices_to_cur_step() - cur_event_idx = len(events) - cur_state_event_idx = len(state_events) - if encoding_state_to_events_fn: - # Dump state to state events *before* processing the next event, because - # we want to capture the state prior to the occurrence of the event. - for e in encoding_state_to_events_fn(state): - state_events.append(codec.encode_event(e)) - - for e in encode_event_fn(state, event_value, codec): - events.append(codec.encode_event(e)) - - # After the last event, continue filling out the event_start_indices array. - # The inequality is not strict because if our current step lines up exactly - # with (the start of) an audio frame, we need to add an additional shift event - # to "cover" that frame. - while cur_step / codec.steps_per_second <= frame_times[-1]: - events.append(codec.encode_event(Event(type="shift", value=1))) - cur_step += 1 - fill_event_start_indices_to_cur_step() - cur_event_idx = len(events) - - # Now fill in event_end_indices. We need this extra array to make sure that - # when we slice events, each slice ends exactly where the subsequent slice - # begins. - event_end_indices = event_start_indices[1:] + [len(events)] - - events = np.array(events).astype(np.int32) - state_events = np.array(state_events).astype(np.int32) - event_start_indices = segment(np.array(event_start_indices).astype(np.int32), TARGET_FEATURE_LENGTH) - event_end_indices = segment(np.array(event_end_indices).astype(np.int32), TARGET_FEATURE_LENGTH) - state_event_indices = segment(np.array(state_event_indices).astype(np.int32), TARGET_FEATURE_LENGTH) - - outputs = [] - for start_indices, end_indices, event_indices in zip(event_start_indices, event_end_indices, state_event_indices): - outputs.append( - { - "inputs": events, - "event_start_indices": start_indices, - "event_end_indices": end_indices, - "state_events": state_events, - "state_event_indices": event_indices, - } - ) - - return outputs - - -def extract_sequence_with_indices(features, state_events_end_token=None, feature_key="inputs"): - """Extract target sequence corresponding to audio token segment.""" - features = features.copy() - start_idx = features["event_start_indices"][0] - end_idx = features["event_end_indices"][-1] - - features[feature_key] = features[feature_key][start_idx:end_idx] - - if state_events_end_token is not None: - # Extract the state events corresponding to the audio start token, and - # prepend them to the targets array. - state_event_start_idx = features["state_event_indices"][0] - state_event_end_idx = state_event_start_idx + 1 - while features["state_events"][state_event_end_idx - 1] != state_events_end_token: - state_event_end_idx += 1 - features[feature_key] = np.concatenate( - [ - features["state_events"][state_event_start_idx:state_event_end_idx], - features[feature_key], - ], - axis=0, - ) - - return features - - -def map_midi_programs( - feature, codec: Codec, granularity_type: str = "full", feature_key: str = "inputs" -) -> Mapping[str, Any]: - """Apply MIDI program map to token sequences.""" - granularity = PROGRAM_GRANULARITIES[granularity_type] - - feature[feature_key] = granularity.tokens_map_fn(feature[feature_key], codec) - return feature - - -def run_length_encode_shifts_fn( - features, - codec: Codec, - feature_key: str = "inputs", - state_change_event_types: Sequence[str] = (), -) -> Callable[[Mapping[str, Any]], Mapping[str, Any]]: - """Return a function that run-length encodes shifts for a given codec. - - Args: - codec: The Codec to use for shift events. - feature_key: The feature key for which to run-length encode shifts. - state_change_event_types: A list of event types that represent state - changes; tokens corresponding to these event types will be interpreted as state changes and redundant ones - will be removed. - - Returns: - A preprocessing function that run-length encodes single-step shifts. - """ - state_change_event_ranges = [codec.event_type_range(event_type) for event_type in state_change_event_types] - - def run_length_encode_shifts(features: MutableMapping[str, Any]) -> Mapping[str, Any]: - """Combine leading/interior shifts, trim trailing shifts. - - Args: - features: Dict of features to process. - - Returns: - A dict of features. - """ - events = features[feature_key] - - shift_steps = 0 - total_shift_steps = 0 - output = np.array([], dtype=np.int32) - - current_state = np.zeros(len(state_change_event_ranges), dtype=np.int32) - - for event in events: - if codec.is_shift_event_index(event): - shift_steps += 1 - total_shift_steps += 1 - - else: - # If this event is a state change and has the same value as the current - # state, we can skip it entirely. - is_redundant = False - for i, (min_index, max_index) in enumerate(state_change_event_ranges): - if (min_index <= event) and (event <= max_index): - if current_state[i] == event: - is_redundant = True - current_state[i] = event - if is_redundant: - continue - - # Once we've reached a non-shift event, RLE all previous shift events - # before outputting the non-shift event. - if shift_steps > 0: - shift_steps = total_shift_steps - while shift_steps > 0: - output_steps = np.minimum(codec.max_shift_steps, shift_steps) - output = np.concatenate([output, [output_steps]], axis=0) - shift_steps -= output_steps - output = np.concatenate([output, [event]], axis=0) - - features[feature_key] = output - return features - - return run_length_encode_shifts(features) - - -def note_representation_processor_chain(features, codec: Codec, note_representation_config: NoteRepresentationConfig): - tie_token = codec.encode_event(Event("tie", 0)) - state_events_end_token = tie_token if note_representation_config.include_ties else None - - features = extract_sequence_with_indices( - features, state_events_end_token=state_events_end_token, feature_key="inputs" - ) - - features = map_midi_programs(features, codec) - - features = run_length_encode_shifts_fn(features, codec, state_change_event_types=["velocity", "program"]) - - return features - - -class MidiProcessor: - def __init__(self): - self.codec = Codec( - max_shift_steps=DEFAULT_MAX_SHIFT_SECONDS * DEFAULT_STEPS_PER_SECOND, - steps_per_second=DEFAULT_STEPS_PER_SECOND, - event_ranges=[ - EventRange("pitch", note_seq.MIN_MIDI_PITCH, note_seq.MAX_MIDI_PITCH), - EventRange("velocity", 0, DEFAULT_NUM_VELOCITY_BINS), - EventRange("tie", 0, 0), - EventRange("program", note_seq.MIN_MIDI_PROGRAM, note_seq.MAX_MIDI_PROGRAM), - EventRange("drum", note_seq.MIN_MIDI_PITCH, note_seq.MAX_MIDI_PITCH), - ], - ) - 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]): - if not isinstance(midi, bytes): - with open(midi, "rb") as f: - midi = f.read() - - ns = note_seq.midi_to_note_sequence(midi) - ns_sus = note_seq.apply_sustain_control_changes(ns) - - for note in ns_sus.notes: - if not note.is_drum: - note.program = program_to_slakh_program(note.program) - - samples = np.zeros(int(ns_sus.total_time * SAMPLE_RATE)) - - _, frame_times = audio_to_frames(samples, HOP_SIZE, FRAME_RATE) - times, values = note_sequence_to_onsets_and_offsets_and_programs(ns_sus) - - events = encode_and_index_events( - state=NoteEncodingState(), - event_times=times, - event_values=values, - frame_times=frame_times, - codec=self.codec, - encode_event_fn=note_event_data_to_events, - encoding_state_to_events_fn=note_encoding_state_to_events, - ) - - events = [ - note_representation_processor_chain(event, self.codec, self.note_representation_config) for event in events - ] - input_tokens = [self.tokenizer.encode(event["inputs"]) for event in events] - - return input_tokens diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/notes_encoder.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/notes_encoder.py deleted file mode 100644 index 25ad4a4ccfd28b2423ee9771130a24c941c06649..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/notes_encoder.py +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright 2022 The Music Spectrogram Diffusion Authors. -# 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 torch -import torch.nn as nn -from transformers.modeling_utils import ModuleUtilsMixin -from transformers.models.t5.modeling_t5 import T5Block, T5Config, T5LayerNorm - -from ....configuration_utils import ConfigMixin, register_to_config -from ....models import ModelMixin - - -class SpectrogramNotesEncoder(ModelMixin, ConfigMixin, ModuleUtilsMixin): - @register_to_config - def __init__( - self, - max_length: int, - vocab_size: int, - d_model: int, - dropout_rate: float, - num_layers: int, - num_heads: int, - d_kv: int, - d_ff: int, - feed_forward_proj: str, - is_decoder: bool = False, - ): - super().__init__() - - self.token_embedder = nn.Embedding(vocab_size, d_model) - - self.position_encoding = nn.Embedding(max_length, d_model) - self.position_encoding.weight.requires_grad = False - - self.dropout_pre = nn.Dropout(p=dropout_rate) - - t5config = T5Config( - vocab_size=vocab_size, - d_model=d_model, - num_heads=num_heads, - d_kv=d_kv, - d_ff=d_ff, - dropout_rate=dropout_rate, - feed_forward_proj=feed_forward_proj, - is_decoder=is_decoder, - is_encoder_decoder=False, - ) - - self.encoders = nn.ModuleList() - for lyr_num in range(num_layers): - lyr = T5Block(t5config) - self.encoders.append(lyr) - - self.layer_norm = T5LayerNorm(d_model) - self.dropout_post = nn.Dropout(p=dropout_rate) - - def forward(self, encoder_input_tokens, encoder_inputs_mask): - x = self.token_embedder(encoder_input_tokens) - - seq_length = encoder_input_tokens.shape[1] - inputs_positions = torch.arange(seq_length, device=encoder_input_tokens.device) - x += self.position_encoding(inputs_positions) - - x = self.dropout_pre(x) - - # inverted the attention mask - input_shape = encoder_input_tokens.size() - extended_attention_mask = self.get_extended_attention_mask(encoder_inputs_mask, input_shape) - - for lyr in self.encoders: - x = lyr(x, extended_attention_mask)[0] - x = self.layer_norm(x) - - return self.dropout_post(x), encoder_inputs_mask diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py deleted file mode 100644 index be07b1b15ea89b619e85540864a1f885ac68534b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py +++ /dev/null @@ -1,269 +0,0 @@ -# Copyright 2022 The Music Spectrogram Diffusion Authors. -# 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 math -from typing import Any, Callable, List, Optional, Tuple, Union - -import numpy as np -import torch - -from ....models import T5FilmDecoder -from ....schedulers import DDPMScheduler -from ....utils import is_onnx_available, logging -from ....utils.torch_utils import randn_tensor - - -if is_onnx_available(): - from ...onnx_utils import OnnxRuntimeModel - -from ...pipeline_utils import AudioPipelineOutput, DiffusionPipeline -from .continuous_encoder import SpectrogramContEncoder -from .notes_encoder import SpectrogramNotesEncoder - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -TARGET_FEATURE_LENGTH = 256 - - -class SpectrogramDiffusionPipeline(DiffusionPipeline): - r""" - Pipeline for unconditional audio generation. - - 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: - notes_encoder ([`SpectrogramNotesEncoder`]): - continuous_encoder ([`SpectrogramContEncoder`]): - decoder ([`T5FilmDecoder`]): - A [`T5FilmDecoder`] to denoise the encoded audio latents. - scheduler ([`DDPMScheduler`]): - A scheduler to be used in combination with `decoder` to denoise the encoded audio latents. - melgan ([`OnnxRuntimeModel`]): - """ - - _optional_components = ["melgan"] - - def __init__( - self, - notes_encoder: SpectrogramNotesEncoder, - continuous_encoder: SpectrogramContEncoder, - decoder: T5FilmDecoder, - scheduler: DDPMScheduler, - melgan: OnnxRuntimeModel if is_onnx_available() else Any, - ) -> None: - super().__init__() - - # From MELGAN - self.min_value = math.log(1e-5) # Matches MelGAN training. - self.max_value = 4.0 # Largest value for most examples - self.n_dims = 128 - - self.register_modules( - notes_encoder=notes_encoder, - continuous_encoder=continuous_encoder, - decoder=decoder, - scheduler=scheduler, - melgan=melgan, - ) - - def scale_features(self, features, output_range=(-1.0, 1.0), clip=False): - """Linearly scale features to network outputs range.""" - min_out, max_out = output_range - if clip: - features = torch.clip(features, self.min_value, self.max_value) - # Scale to [0, 1]. - zero_one = (features - self.min_value) / (self.max_value - self.min_value) - # Scale to [min_out, max_out]. - return zero_one * (max_out - min_out) + min_out - - def scale_to_features(self, outputs, input_range=(-1.0, 1.0), clip=False): - """Invert by linearly scaling network outputs to features range.""" - min_out, max_out = input_range - outputs = torch.clip(outputs, min_out, max_out) if clip else outputs - # Scale to [0, 1]. - zero_one = (outputs - min_out) / (max_out - min_out) - # Scale to [self.min_value, self.max_value]. - return zero_one * (self.max_value - self.min_value) + self.min_value - - def encode(self, input_tokens, continuous_inputs, continuous_mask): - tokens_mask = input_tokens > 0 - tokens_encoded, tokens_mask = self.notes_encoder( - encoder_input_tokens=input_tokens, encoder_inputs_mask=tokens_mask - ) - - continuous_encoded, continuous_mask = self.continuous_encoder( - encoder_inputs=continuous_inputs, encoder_inputs_mask=continuous_mask - ) - - return [(tokens_encoded, tokens_mask), (continuous_encoded, continuous_mask)] - - def decode(self, encodings_and_masks, input_tokens, noise_time): - timesteps = noise_time - if not torch.is_tensor(timesteps): - timesteps = torch.tensor([timesteps], dtype=torch.long, device=input_tokens.device) - elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0: - timesteps = timesteps[None].to(input_tokens.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps * torch.ones(input_tokens.shape[0], dtype=timesteps.dtype, device=timesteps.device) - - logits = self.decoder( - encodings_and_masks=encodings_and_masks, decoder_input_tokens=input_tokens, decoder_noise_time=timesteps - ) - return logits - - @torch.no_grad() - def __call__( - self, - input_tokens: List[List[int]], - generator: Optional[torch.Generator] = None, - num_inference_steps: int = 100, - return_dict: bool = True, - output_type: str = "np", - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - ) -> Union[AudioPipelineOutput, Tuple]: - if (callback_steps is None) or ( - 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)}." - ) - r""" - The call function to the pipeline for generation. - - Args: - 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): - The number of denoising steps. More denoising steps usually lead to a higher quality audio at the - expense of slower inference. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.AudioPipelineOutput`] instead of a plain tuple. - output_type (`str`, *optional*, defaults to `"np"`): - The output format of the generated audio. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - - Example: - - ```py - >>> from diffusers import SpectrogramDiffusionPipeline, MidiProcessor - - >>> pipe = SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion") - >>> pipe = pipe.to("cuda") - >>> processor = MidiProcessor() - - >>> # Download MIDI from: wget http://www.piano-midi.de/midis/beethoven/beethoven_hammerklavier_2.mid - >>> output = pipe(processor("beethoven_hammerklavier_2.mid")) - - >>> audio = output.audios[0] - ``` - - Returns: - [`pipelines.AudioPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`pipelines.AudioPipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated audio. - """ - - pred_mel = np.zeros([1, TARGET_FEATURE_LENGTH, self.n_dims], dtype=np.float32) - full_pred_mel = np.zeros([1, 0, self.n_dims], np.float32) - ones = torch.ones((1, TARGET_FEATURE_LENGTH), dtype=bool, device=self.device) - - for i, encoder_input_tokens in enumerate(input_tokens): - if i == 0: - encoder_continuous_inputs = torch.from_numpy(pred_mel[:1].copy()).to( - device=self.device, dtype=self.decoder.dtype - ) - # The first chunk has no previous context. - encoder_continuous_mask = torch.zeros((1, TARGET_FEATURE_LENGTH), dtype=bool, device=self.device) - else: - # The full song pipeline does not feed in a context feature, so the mask - # will be all 0s after the feature converter. Because we know we're - # feeding in a full context chunk from the previous prediction, set it - # to all 1s. - encoder_continuous_mask = ones - - encoder_continuous_inputs = self.scale_features( - encoder_continuous_inputs, output_range=[-1.0, 1.0], clip=True - ) - - encodings_and_masks = self.encode( - input_tokens=torch.IntTensor([encoder_input_tokens]).to(device=self.device), - continuous_inputs=encoder_continuous_inputs, - continuous_mask=encoder_continuous_mask, - ) - - # Sample encoder_continuous_inputs shaped gaussian noise to begin loop - x = randn_tensor( - shape=encoder_continuous_inputs.shape, - generator=generator, - device=self.device, - dtype=self.decoder.dtype, - ) - - # set step values - self.scheduler.set_timesteps(num_inference_steps) - - # Denoising diffusion loop - for j, t in enumerate(self.progress_bar(self.scheduler.timesteps)): - output = self.decode( - encodings_and_masks=encodings_and_masks, - input_tokens=x, - noise_time=t / self.scheduler.config.num_train_timesteps, # rescale to [0, 1) - ) - - # Compute previous output: x_t -> x_t-1 - x = self.scheduler.step(output, t, x, generator=generator).prev_sample - - mel = self.scale_to_features(x, input_range=[-1.0, 1.0]) - encoder_continuous_inputs = mel[:1] - pred_mel = mel.cpu().float().numpy() - - full_pred_mel = np.concatenate([full_pred_mel, pred_mel[:1]], axis=1) - - # call the callback, if provided - if callback is not None and i % callback_steps == 0: - callback(i, full_pred_mel) - - logger.info("Generated segment", i) - - if output_type == "np" and not is_onnx_available(): - raise ValueError( - "Cannot return output in 'np' format if ONNX is not available. Make sure to have ONNX installed or set 'output_type' to 'mel'." - ) - elif output_type == "np" and self.melgan is None: - raise ValueError( - "Cannot return output in 'np' format if melgan component is not defined. Make sure to define `self.melgan` or set 'output_type' to 'mel'." - ) - - if output_type == "np": - output = self.melgan(input_features=full_pred_mel.astype(np.float32)) - else: - output = full_pred_mel - - if not return_dict: - return (output,) - - return AudioPipelineOutput(audios=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__init__.py deleted file mode 100644 index 36cf1a33ce6ada8e718aabadb9a706737aee30bd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__init__.py +++ /dev/null @@ -1,55 +0,0 @@ -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 - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) -else: - _import_structure["pipeline_cycle_diffusion"] = ["CycleDiffusionPipeline"] - _import_structure["pipeline_stable_diffusion_inpaint_legacy"] = ["StableDiffusionInpaintPipelineLegacy"] - _import_structure["pipeline_stable_diffusion_model_editing"] = ["StableDiffusionModelEditingPipeline"] - - _import_structure["pipeline_stable_diffusion_paradigms"] = ["StableDiffusionParadigmsPipeline"] - _import_structure["pipeline_stable_diffusion_pix2pix_zero"] = ["StableDiffusionPix2PixZeroPipeline"] - -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_cycle_diffusion import CycleDiffusionPipeline - from .pipeline_stable_diffusion_inpaint_legacy import StableDiffusionInpaintPipelineLegacy - from .pipeline_stable_diffusion_model_editing import StableDiffusionModelEditingPipeline - from .pipeline_stable_diffusion_paradigms import StableDiffusionParadigmsPipeline - from .pipeline_stable_diffusion_pix2pix_zero import StableDiffusionPix2PixZeroPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index ffc0e26e48dcbb5bbddd135b23d15cc116bb5376..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_cycle_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_cycle_diffusion.cpython-310.pyc deleted file mode 100644 index 1b974390dedda5a016b41240799231579468defb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_cycle_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_onnx_stable_diffusion_inpaint_legacy.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_onnx_stable_diffusion_inpaint_legacy.cpython-310.pyc deleted file mode 100644 index 391bc2c7552a23e87e54ede2d2b3397a2a5cc744..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_onnx_stable_diffusion_inpaint_legacy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_inpaint_legacy.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_inpaint_legacy.cpython-310.pyc deleted file mode 100644 index e8c573da68149e4b3b29b2d4846de0fa2bc5d570..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_inpaint_legacy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_model_editing.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_model_editing.cpython-310.pyc deleted file mode 100644 index fb6357b7f68a4a0babad0f6ae322c8f621136b0c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_model_editing.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_paradigms.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_paradigms.cpython-310.pyc deleted file mode 100644 index f55ed46131f3fdd22ff208f5c2ab59dd14e5edd9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_paradigms.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_pix2pix_zero.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_pix2pix_zero.cpython-310.pyc deleted file mode 100644 index 2f28e7d2fdb20595c7f1158f46f7b847aace971b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/__pycache__/pipeline_stable_diffusion_pix2pix_zero.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py deleted file mode 100644 index 525c5e90c51efcfada1abd0113fd69126700e8de..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py +++ /dev/null @@ -1,953 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from packaging import version -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer - -from ....configuration_utils import FrozenDict -from ....image_processor import PipelineImageInput, VaeImageProcessor -from ....loaders import StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ....models import AutoencoderKL, UNet2DConditionModel -from ....models.lora import adjust_lora_scale_text_encoder -from ....schedulers import DDIMScheduler -from ....utils import PIL_INTERPOLATION, USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers -from ....utils.torch_utils import randn_tensor -from ...pipeline_utils import DiffusionPipeline -from ...stable_diffusion.pipeline_output import StableDiffusionPipelineOutput -from ...stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess -def preprocess(image): - 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): - return image - elif isinstance(image, PIL.Image.Image): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - w, h = image[0].size - w, h = (x - x % 8 for x in (w, h)) # resize to integer multiple of 8 - - image = [np.array(i.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image] - image = np.concatenate(image, axis=0) - image = np.array(image).astype(np.float32) / 255.0 - image = image.transpose(0, 3, 1, 2) - image = 2.0 * image - 1.0 - image = torch.from_numpy(image) - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, dim=0) - return 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" -): - 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") - - -def posterior_sample(scheduler, latents, timestep, clean_latents, generator, eta): - # 1. get previous step value (=t-1) - prev_timestep = timestep - scheduler.config.num_train_timesteps // scheduler.num_inference_steps - - if prev_timestep <= 0: - return clean_latents - - # 2. compute alphas, betas - alpha_prod_t = scheduler.alphas_cumprod[timestep] - alpha_prod_t_prev = ( - scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else scheduler.final_alpha_cumprod - ) - - variance = scheduler._get_variance(timestep, prev_timestep) - std_dev_t = eta * variance ** (0.5) - - # direction pointing to x_t - e_t = (latents - alpha_prod_t ** (0.5) * clean_latents) / (1 - alpha_prod_t) ** (0.5) - dir_xt = (1.0 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * e_t - noise = std_dev_t * randn_tensor( - clean_latents.shape, dtype=clean_latents.dtype, device=clean_latents.device, generator=generator - ) - prev_latents = alpha_prod_t_prev ** (0.5) * clean_latents + dir_xt + noise - - return prev_latents - - -def compute_noise(scheduler, prev_latents, latents, timestep, noise_pred, eta): - # 1. get previous step value (=t-1) - prev_timestep = timestep - scheduler.config.num_train_timesteps // scheduler.num_inference_steps - - # 2. compute alphas, betas - alpha_prod_t = scheduler.alphas_cumprod[timestep] - alpha_prod_t_prev = ( - scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else scheduler.final_alpha_cumprod - ) - - beta_prod_t = 1 - alpha_prod_t - - # 3. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (12) from https://huggingface.co/papers/2010.02502 - pred_original_sample = (latents - beta_prod_t ** (0.5) * noise_pred) / alpha_prod_t ** (0.5) - - # 4. Clip "predicted x_0" - if scheduler.config.clip_sample: - pred_original_sample = torch.clamp(pred_original_sample, -1, 1) - - # 5. compute variance: "sigma_t(η)" -> see formula (16) - # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1) - variance = scheduler._get_variance(timestep, prev_timestep) - std_dev_t = eta * variance ** (0.5) - - # 6. compute "direction pointing to x_t" of formula (12) from https://huggingface.co/papers/2010.02502 - pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * noise_pred - - noise = (prev_latents - (alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction)) / ( - variance ** (0.5) * eta - ) - return noise - - -class CycleDiffusionPipeline(DiffusionPipeline, TextualInversionLoaderMixin, StableDiffusionLoraLoaderMixin): - r""" - Pipeline for text-guided image to image generation using Stable Diffusion. - - 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.). - - 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 - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can only be an - instance of [`DDIMScheduler`]. - 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: DDIMScheduler, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_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.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.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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 - - def check_inputs( - self, - prompt, - strength, - callback_steps, - 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 in [0.0, 1.0] but is {strength}") - - 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}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None): - image = image.to(device=device, dtype=dtype) - - batch_size = image.shape[0] - - if image.shape[1] == 4: - init_latents = image - - else: - 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 isinstance(generator, list): - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(image.shape[0]) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - deprecation_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial" - " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" - " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" - " your script to pass as many initial images as text prompts to suppress this warning." - ) - deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt * num_images_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents] * num_images_per_prompt, dim=0) - - # add noise to latents using the timestep - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # get latents - clean_latents = init_latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - latents = init_latents - - return latents, clean_latents - - @torch.no_grad() - def __call__( - self, - prompt: Union[str, List[str]], - source_prompt: Union[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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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` 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): - 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`. - 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`. - 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`. - source_guidance_scale (`float`, *optional*, defaults to 1): - Guidance scale for the source prompt. This is useful to control the amount of influence the source - prompt has for encoding. - 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 (η) 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*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - 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. - 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - 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. - Example: - - ```py - import requests - import torch - from PIL import Image - from io import BytesIO - - from diffusers import CycleDiffusionPipeline, DDIMScheduler - - # load the pipeline - # make sure you're logged in with `huggingface-cli login` - model_id_or_path = "CompVis/stable-diffusion-v1-4" - scheduler = DDIMScheduler.from_pretrained(model_id_or_path, subfolder="scheduler") - pipe = CycleDiffusionPipeline.from_pretrained(model_id_or_path, scheduler=scheduler).to("cuda") - - # let's download an initial image - url = "https://raw.githubusercontent.com/ChenWu98/cycle-diffusion/main/data/dalle2/An%20astronaut%20riding%20a%20horse.png" - response = requests.get(url) - init_image = Image.open(BytesIO(response.content)).convert("RGB") - init_image = init_image.resize((512, 512)) - init_image.save("horse.png") - - # let's specify a prompt - source_prompt = "An astronaut riding a horse" - prompt = "An astronaut riding an elephant" - - # call the pipeline - image = pipe( - prompt=prompt, - source_prompt=source_prompt, - image=init_image, - num_inference_steps=100, - eta=0.1, - strength=0.8, - guidance_scale=2, - source_guidance_scale=1, - ).images[0] - - image.save("horse_to_elephant.png") - - # let's try another example - # See more samples at the original repo: https://github.com/ChenWu98/cycle-diffusion - url = ( - "https://raw.githubusercontent.com/ChenWu98/cycle-diffusion/main/data/dalle2/A%20black%20colored%20car.png" - ) - response = requests.get(url) - init_image = Image.open(BytesIO(response.content)).convert("RGB") - init_image = init_image.resize((512, 512)) - init_image.save("black.png") - - source_prompt = "A black colored car" - prompt = "A blue colored car" - - # call the pipeline - torch.manual_seed(0) - image = pipe( - prompt=prompt, - source_prompt=source_prompt, - image=init_image, - num_inference_steps=100, - eta=0.1, - strength=0.85, - guidance_scale=3, - source_guidance_scale=1, - ).images[0] - - image.save("black_to_blue.png") - ``` - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - # 1. Check inputs - self.check_inputs(prompt, strength, callback_steps) - - # 2. Define call parameters - batch_size = 1 if isinstance(prompt, str) else len(prompt) - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - prompt_embeds_tuple = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - do_classifier_free_guidance, - prompt_embeds=prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=clip_skip, - ) - source_prompt_embeds_tuple = self.encode_prompt( - source_prompt, device, num_images_per_prompt, do_classifier_free_guidance, None, clip_skip=clip_skip - ) - if prompt_embeds_tuple[1] is not None: - prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]]) - else: - prompt_embeds = prompt_embeds_tuple[0] - if source_prompt_embeds_tuple[1] is not None: - source_prompt_embeds = torch.cat([source_prompt_embeds_tuple[1], source_prompt_embeds_tuple[0]]) - else: - source_prompt_embeds = source_prompt_embeds_tuple[0] - - # 4. Preprocess image - image = self.image_processor.preprocess(image) - - # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - - # 6. Prepare latent variables - latents, clean_latents = self.prepare_latents( - image, latent_timestep, batch_size, num_images_per_prompt, prompt_embeds.dtype, device, generator - ) - source_latents = latents - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - generator = extra_step_kwargs.pop("generator", None) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - source_latent_model_input = ( - torch.cat([source_latents] * 2) if do_classifier_free_guidance else source_latents - ) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - source_latent_model_input = self.scheduler.scale_model_input(source_latent_model_input, t) - - # predict the noise residual - if do_classifier_free_guidance: - concat_latent_model_input = torch.stack( - [ - source_latent_model_input[0], - latent_model_input[0], - source_latent_model_input[1], - latent_model_input[1], - ], - dim=0, - ) - concat_prompt_embeds = torch.stack( - [ - source_prompt_embeds[0], - prompt_embeds[0], - source_prompt_embeds[1], - prompt_embeds[1], - ], - dim=0, - ) - else: - concat_latent_model_input = torch.cat( - [ - source_latent_model_input, - latent_model_input, - ], - dim=0, - ) - concat_prompt_embeds = torch.cat( - [ - source_prompt_embeds, - prompt_embeds, - ], - dim=0, - ) - - concat_noise_pred = self.unet( - concat_latent_model_input, - t, - cross_attention_kwargs=cross_attention_kwargs, - encoder_hidden_states=concat_prompt_embeds, - ).sample - - # perform guidance - if do_classifier_free_guidance: - ( - source_noise_pred_uncond, - noise_pred_uncond, - source_noise_pred_text, - noise_pred_text, - ) = concat_noise_pred.chunk(4, dim=0) - - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - source_noise_pred = source_noise_pred_uncond + source_guidance_scale * ( - source_noise_pred_text - source_noise_pred_uncond - ) - - else: - (source_noise_pred, noise_pred) = concat_noise_pred.chunk(2, dim=0) - - # Sample source_latents from the posterior distribution. - prev_source_latents = posterior_sample( - self.scheduler, source_latents, t, clean_latents, generator=generator, **extra_step_kwargs - ) - # Compute noise. - noise = compute_noise( - self.scheduler, prev_source_latents, source_latents, t, source_noise_pred, **extra_step_kwargs - ) - source_latents = prev_source_latents - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step( - noise_pred, t, latents, variance_noise=noise, **extra_step_kwargs - ).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - # 9. Post-processing - 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) - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py deleted file mode 100644 index 2d9eaa493f1ad8b79118f99f8edbdd80eb517dce..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py +++ /dev/null @@ -1,542 +0,0 @@ -import inspect -from typing import Callable, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTokenizer - -from ....configuration_utils import FrozenDict -from ....schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler -from ....utils import deprecate, logging -from ...onnx_utils import ORT_TO_NP_TYPE, OnnxRuntimeModel -from ...pipeline_utils import DiffusionPipeline -from ...stable_diffusion.pipeline_output import StableDiffusionPipelineOutput - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -def preprocess(image): - w, h = image.size - w, h = (x - x % 32 for x in (w, h)) # resize to integer multiple of 32 - image = image.resize((w, h), resample=PIL.Image.LANCZOS) - image = np.array(image).astype(np.float32) / 255.0 - image = image[None].transpose(0, 3, 1, 2) - return 2.0 * image - 1.0 - - -def preprocess_mask(mask, scale_factor=8): - mask = mask.convert("L") - w, h = mask.size - w, h = (x - x % 32 for x in (w, h)) # resize to integer multiple of 32 - mask = mask.resize((w // scale_factor, h // scale_factor), resample=PIL.Image.NEAREST) - mask = np.array(mask).astype(np.float32) / 255.0 - mask = np.tile(mask, (4, 1, 1)) - mask = mask[None].transpose(0, 1, 2, 3) # what does this step do? - mask = 1 - mask # repaint white, keep black - return mask - - -class OnnxStableDiffusionInpaintPipelineLegacy(DiffusionPipeline): - r""" - Pipeline for text-guided image inpainting using Stable Diffusion. This is a *legacy feature* for Onnx pipelines to - provide compatibility with StableDiffusionInpaintPipelineLegacy and may be removed in the future. - - 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.) - - 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/runwayml/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`. - """ - - _optional_components = ["safety_checker", "feature_extractor"] - _is_onnx = True - - vae_encoder: OnnxRuntimeModel - vae_decoder: OnnxRuntimeModel - text_encoder: OnnxRuntimeModel - tokenizer: CLIPTokenizer - unet: OnnxRuntimeModel - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler] - safety_checker: OnnxRuntimeModel - feature_extractor: CLIPImageProcessor - - def __init__( - self, - vae_encoder: OnnxRuntimeModel, - vae_decoder: OnnxRuntimeModel, - text_encoder: OnnxRuntimeModel, - tokenizer: CLIPTokenizer, - unet: OnnxRuntimeModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], - safety_checker: OnnxRuntimeModel, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae_encoder=vae_encoder, - vae_decoder=vae_decoder, - 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) - - # 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], - do_classifier_free_guidance: bool, - negative_prompt: Optional[str], - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - 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]`): - 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*): - 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 (`np.ndarray`, *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. - """ - 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: - # get prompt text embeddings - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - text_input_ids = text_inputs.input_ids - untruncated_ids = self.tokenizer(prompt, padding="max_length", return_tensors="np").input_ids - - if not np.array_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}" - ) - - prompt_embeds = self.text_encoder(input_ids=text_input_ids.astype(np.int32))[0] - - prompt_embeds = np.repeat(prompt_embeds, num_images_per_prompt, axis=0) - - # 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 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] * batch_size - 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 - - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="np", - ) - negative_prompt_embeds = self.text_encoder(input_ids=uncond_input.input_ids.astype(np.int32))[0] - - if do_classifier_free_guidance: - negative_prompt_embeds = np.repeat(negative_prompt_embeds, num_images_per_prompt, axis=0) - - # 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 - prompt_embeds = np.concatenate([negative_prompt_embeds, prompt_embeds]) - - return prompt_embeds - - def check_inputs( - self, - prompt, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - if (callback_steps is None) or ( - 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 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 __call__( - self, - prompt: Union[str, List[str]], - image: Union[np.ndarray, PIL.Image.Image] = None, - mask_image: Union[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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, np.ndarray], None]] = None, - callback_steps: int = 1, - ): - r""" - Function invoked when calling the pipeline for generation. - - Args: - 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 - process. This is the image whose masked region will be inpainted. - mask_image (`nd.ndarray` or `PIL.Image.Image`): - `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be - replaced by noise and therefore 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)`.uu - strength (`float`, *optional*, defaults to 0.8): - 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_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 will be modulated by `strength`. - 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. 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 (`np.random.RandomState`, *optional*): - A np.random.RandomState to make generation deterministic. - prompt_embeds (`np.ndarray`, *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 (`np.ndarray`, *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: np.ndarray)`. - 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. - - 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`. - """ - - # check inputs. Raise error if not correct - self.check_inputs(prompt, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds) - - # 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 strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - - if generator is None: - generator = np.random - - # set timesteps - self.scheduler.set_timesteps(num_inference_steps) - - if isinstance(image, PIL.Image.Image): - image = preprocess(image) - - # 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 = guidance_scale > 1.0 - - prompt_embeds = self._encode_prompt( - prompt, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - - latents_dtype = prompt_embeds.dtype - image = image.astype(latents_dtype) - - # encode the init image into latents and scale the latents - init_latents = self.vae_encoder(sample=image)[0] - init_latents = 0.18215 * init_latents - - # Expand init_latents for batch_size and num_images_per_prompt - init_latents = np.concatenate([init_latents] * num_images_per_prompt, axis=0) - init_latents_orig = init_latents - - # preprocess mask - if not isinstance(mask_image, np.ndarray): - mask_image = preprocess_mask(mask_image, 8) - mask_image = mask_image.astype(latents_dtype) - mask = np.concatenate([mask_image] * num_images_per_prompt, axis=0) - - # check sizes - if not mask.shape == init_latents.shape: - raise ValueError("The mask and image should be the same size!") - - # get the original timestep using init_timestep - offset = self.scheduler.config.get("steps_offset", 0) - init_timestep = int(num_inference_steps * strength) + offset - init_timestep = min(init_timestep, num_inference_steps) - - timesteps = self.scheduler.timesteps.numpy()[-init_timestep] - timesteps = np.array([timesteps] * batch_size * num_images_per_prompt) - - # add noise to latents using the timesteps - noise = generator.randn(*init_latents.shape).astype(latents_dtype) - init_latents = self.scheduler.add_noise( - torch.from_numpy(init_latents), torch.from_numpy(noise), torch.from_numpy(timesteps) - ) - init_latents = init_latents.numpy() - - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (?) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to ? in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - latents = init_latents - - t_start = max(num_inference_steps - init_timestep + offset, 0) - timesteps = self.scheduler.timesteps[t_start:].numpy() - timestep_dtype = next( - (input.type for input in self.unet.model.get_inputs() if input.name == "timestep"), "tensor(float)" - ) - timestep_dtype = ORT_TO_NP_TYPE[timestep_dtype] - - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - timestep = np.array([t], dtype=timestep_dtype) - noise_pred = self.unet(sample=latent_model_input, timestep=timestep, encoder_hidden_states=prompt_embeds)[ - 0 - ] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step( - torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs - ).prev_sample - - latents = latents.numpy() - - init_latents_proper = self.scheduler.add_noise( - torch.from_numpy(init_latents_orig), torch.from_numpy(noise), torch.from_numpy(np.array([t])) - ) - - init_latents_proper = init_latents_proper.numpy() - - latents = (init_latents_proper * mask) + (latents * (1 - mask)) - - # call the callback, if provided - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - latents = 1 / 0.18215 * latents - # image = self.vae_decoder(latent_sample=latents)[0] - # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1 - image = np.concatenate( - [self.vae_decoder(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])] - ) - - image = np.clip(image / 2 + 0.5, 0, 1) - image = image.transpose((0, 2, 3, 1)) - - if self.safety_checker is not None: - safety_checker_input = self.feature_extractor( - self.numpy_to_pil(image), return_tensors="np" - ).pixel_values.astype(image.dtype) - # There will throw an error if use safety_checker batchsize>1 - images, has_nsfw_concept = [], [] - for i in range(image.shape[0]): - image_i, has_nsfw_concept_i = self.safety_checker( - clip_input=safety_checker_input[i : i + 1], images=image[i : i + 1] - ) - images.append(image_i) - has_nsfw_concept.append(has_nsfw_concept_i[0]) - image = np.concatenate(images) - else: - has_nsfw_concept = None - - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py deleted file mode 100644 index 205ace65eeb63832ee1f8a1d05b0264508bd9837..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py +++ /dev/null @@ -1,788 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from packaging import version -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer - -from ....configuration_utils import FrozenDict -from ....image_processor import VaeImageProcessor -from ....loaders import FromSingleFileMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ....models import AutoencoderKL, UNet2DConditionModel -from ....models.lora import adjust_lora_scale_text_encoder -from ....schedulers import KarrasDiffusionSchedulers -from ....utils import PIL_INTERPOLATION, USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers -from ....utils.torch_utils import randn_tensor -from ...pipeline_utils import DiffusionPipeline -from ...stable_diffusion import StableDiffusionPipelineOutput -from ...stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -logger = logging.get_logger(__name__) - - -def preprocess_image(image, batch_size): - w, h = image.size - w, h = (x - x % 8 for x in (w, h)) # resize to integer multiple of 8 - image = image.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]) - image = np.array(image).astype(np.float32) / 255.0 - image = np.vstack([image[None].transpose(0, 3, 1, 2)] * batch_size) - image = torch.from_numpy(image) - return 2.0 * image - 1.0 - - -def preprocess_mask(mask, batch_size, scale_factor=8): - if not isinstance(mask, torch.Tensor): - mask = mask.convert("L") - w, h = mask.size - w, h = (x - x % 8 for x in (w, h)) # resize to integer multiple of 8 - mask = mask.resize((w // scale_factor, h // scale_factor), resample=PIL_INTERPOLATION["nearest"]) - mask = np.array(mask).astype(np.float32) / 255.0 - mask = np.tile(mask, (4, 1, 1)) - mask = np.vstack([mask[None]] * batch_size) - mask = 1 - mask # repaint white, keep black - mask = torch.from_numpy(mask) - return mask - - else: - valid_mask_channel_sizes = [1, 3] - # if mask channel is fourth tensor dimension, permute dimensions to pytorch standard (B, C, H, W) - if mask.shape[3] in valid_mask_channel_sizes: - mask = mask.permute(0, 3, 1, 2) - elif mask.shape[1] not in valid_mask_channel_sizes: - raise ValueError( - f"Mask channel dimension of size in {valid_mask_channel_sizes} should be second or fourth dimension," - f" but received mask of shape {tuple(mask.shape)}" - ) - # (potentially) reduce mask channel dimension from 3 to 1 for broadcasting to latent shape - mask = mask.mean(dim=1, keepdim=True) - h, w = mask.shape[-2:] - h, w = (x - x % 8 for x in (h, w)) # resize to integer multiple of 8 - mask = torch.nn.functional.interpolate(mask, (h // scale_factor, w // scale_factor)) - return mask - - -class StableDiffusionInpaintPipelineLegacy( - DiffusionPipeline, TextualInversionLoaderMixin, StableDiffusionLoraLoaderMixin, FromSingleFileMixin -): - r""" - Pipeline for text-guided image inpainting using Stable Diffusion. *This is an experimental feature*. - - 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.) - - In addition the pipeline inherits the following loading methods: - - *Textual-Inversion*: [`loaders.TextualInversionLoaderMixin.load_textual_inversion`] - - *LoRA*: [`loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] - - *Ckpt*: [`loaders.FromSingleFileMixin.from_single_file`] - - as well as the following saving methods: - - *LoRA*: [`loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] - - 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/runwayml/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`. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["feature_extractor"] - _exclude_from_cpu_offload = ["safety_checker"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - deprecation_message = ( - f"The class {self.__class__} is deprecated and will be removed in v1.0.0. You can achieve exactly the same functionality" - "by loading your model into `StableDiffusionInpaintPipeline` instead. See https://github.com/huggingface/diffusers/pull/3533" - "for more information." - ) - deprecate("legacy is outdated", "1.0.0", deprecation_message, standard_warn=False) - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_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.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.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - strength, - callback_steps, - 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 in [0.0, 1.0] but is {strength}") - - 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 get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :] - - return timesteps, num_inference_steps - t_start - - def prepare_latents(self, image, timestep, num_images_per_prompt, dtype, device, generator): - image = image.to(device=device, dtype=dtype) - init_latent_dist = self.vae.encode(image).latent_dist - init_latents = init_latent_dist.sample(generator=generator) - init_latents = self.vae.config.scaling_factor * init_latents - - # Expand init_latents for batch_size and num_images_per_prompt - init_latents = torch.cat([init_latents] * num_images_per_prompt, dim=0) - init_latents_orig = init_latents - - # add noise to latents using the timesteps - noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=dtype) - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - latents = init_latents - return latents, init_latents_orig, noise - - @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, - 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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. - image (`torch.Tensor` or `PIL.Image.Image`): - `Image`, or tensor representing an image batch, that will be used as the starting point for the - process. This is the image whose masked region will be inpainted. - mask_image (`torch.Tensor` or `PIL.Image.Image`): - `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be - replaced by noise and therefore 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 mask is a tensor, the - expected shape should be either `(B, H, W, C)` or `(B, C, H, W)`, where C is 1 or 3. - strength (`float`, *optional*, defaults to 0.8): - Conceptually, indicates how much to inpaint the masked area. Must be between 0 and 1. When `strength` - is 1, the denoising process will be run on the masked area for the full number of iterations specified - in `num_inference_steps`. `image` will be used as a reference for the masked area, adding more noise to - that region the larger the `strength`. If `strength` is 0, no inpainting will occur. - num_inference_steps (`int`, *optional*, defaults to 50): - The reference number of denoising steps. More denoising steps usually lead to a higher quality image at - the expense of slower inference. This parameter will be modulated by `strength`, as explained above. - 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. - add_predicted_noise (`bool`, *optional*, defaults to True): - Use predicted noise instead of random noise when constructing noisy versions of the original image in - the reverse diffusion process - 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. - 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. - cross_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). - 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`. - """ - # 1. Check inputs - self.check_inputs(prompt, strength, 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - 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, - lora_scale=text_encoder_lora_scale, - 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. Preprocess image and mask - if not isinstance(image, torch.Tensor): - image = preprocess_image(image, batch_size) - - mask_image = preprocess_mask(mask_image, batch_size, self.vae_scale_factor) - - # 5. set timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - - # 6. Prepare latent variables - # encode the init image into latents and scale the latents - latents, init_latents_orig, noise = self.prepare_latents( - image, latent_timestep, num_images_per_prompt, prompt_embeds.dtype, device, generator - ) - - # 7. Prepare mask latent - mask = mask_image.to(device=device, dtype=latents.dtype) - mask = torch.cat([mask] * num_images_per_prompt) - - # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 9. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - # masking - if add_predicted_noise: - init_latents_proper = self.scheduler.add_noise( - init_latents_orig, noise_pred_uncond, torch.tensor([t]) - ) - else: - init_latents_proper = self.scheduler.add_noise(init_latents_orig, noise, torch.tensor([t])) - - latents = (init_latents_proper * mask) + (latents * (1 - mask)) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - # use original latents corresponding to unmasked portions of the image - latents = (init_latents_orig * mask) + (latents * (1 - mask)) - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py deleted file mode 100644 index d81f0d262552778283c66bb3258743a317db03fd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py +++ /dev/null @@ -1,832 +0,0 @@ -# Copyright 2025 TIME Authors 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, Dict, List, Optional, Union - -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer - -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 PNDMScheduler -from ....schedulers.scheduling_utils import SchedulerMixin -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 DiffusionPipeline, StableDiffusionMixin -from ...stable_diffusion.pipeline_output import StableDiffusionPipelineOutput -from ...stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -AUGS_CONST = ["A photo of ", "An image of ", "A picture of "] - - -class StableDiffusionModelEditingPipeline( - DiffusionPipeline, StableDiffusionMixin, TextualInversionLoaderMixin, StableDiffusionLoraLoaderMixin -): - r""" - Pipeline for text-to-image model editing. - - 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.). - - 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 - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - with_to_k ([`bool`]): - Whether to edit the key projection matrices along with the value projection matrices. - with_augs ([`list`]): - Textual augmentations to apply while editing the text-to-image model. Set to `[]` for no augmentations. - """ - - 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: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: SchedulerMixin, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - with_to_k: bool = True, - with_augs: list = AUGS_CONST, - ): - super().__init__() - - if isinstance(scheduler, PNDMScheduler): - logger.error("PNDMScheduler for this pipeline is currently not supported.") - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - self.with_to_k = with_to_k - self.with_augs = with_augs - - # get cross-attention layers - ca_layers = [] - - def append_ca(net_): - if net_.__class__.__name__ == "CrossAttention": - ca_layers.append(net_) - elif hasattr(net_, "children"): - for net__ in net_.children(): - append_ca(net__) - - # recursively find all cross-attention layers in unet - for net in self.unet.named_children(): - if "down" in net[0]: - append_ca(net[1]) - elif "up" in net[0]: - append_ca(net[1]) - elif "mid" in net[0]: - append_ca(net[1]) - - # get projection matrices - self.ca_clip_layers = [l for l in ca_layers if l.to_v.in_features == 768] - self.projection_matrices = [l.to_v for l in self.ca_clip_layers] - self.og_matrices = [copy.deepcopy(l.to_v) for l in self.ca_clip_layers] - if self.with_to_k: - self.projection_matrices = self.projection_matrices + [l.to_k for l in self.ca_clip_layers] - self.og_matrices = self.og_matrices + [copy.deepcopy(l.to_k) for l in self.ca_clip_layers] - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - 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}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @torch.no_grad() - def edit_model( - self, - source_prompt: str, - destination_prompt: str, - lamb: float = 0.1, - restart_params: bool = True, - ): - r""" - Apply model editing via closed-form solution (see Eq. 5 in the TIME - [paper](https://huggingface.co/papers/2303.08084)). - - Args: - source_prompt (`str`): - The source prompt containing the concept to be edited. - destination_prompt (`str`): - The destination prompt. Must contain all words from `source_prompt` with additional ones to specify the - target edit. - lamb (`float`, *optional*, defaults to 0.1): - The lambda parameter specifying the regularization intensity. Smaller values increase the editing - power. - restart_params (`bool`, *optional*, defaults to True): - Restart the model parameters to their pre-trained version before editing. This is done to avoid edit - compounding. When it is `False`, edits accumulate. - """ - - # restart LDM parameters - if restart_params: - num_ca_clip_layers = len(self.ca_clip_layers) - for idx_, l in enumerate(self.ca_clip_layers): - l.to_v = copy.deepcopy(self.og_matrices[idx_]) - self.projection_matrices[idx_] = l.to_v - if self.with_to_k: - l.to_k = copy.deepcopy(self.og_matrices[num_ca_clip_layers + idx_]) - self.projection_matrices[num_ca_clip_layers + idx_] = l.to_k - - # set up sentences - old_texts = [source_prompt] - new_texts = [destination_prompt] - # add augmentations - base = old_texts[0] if old_texts[0][0:1] != "A" else "a" + old_texts[0][1:] - for aug in self.with_augs: - old_texts.append(aug + base) - base = new_texts[0] if new_texts[0][0:1] != "A" else "a" + new_texts[0][1:] - for aug in self.with_augs: - new_texts.append(aug + base) - - # prepare input k* and v* - old_embs, new_embs = [], [] - for old_text, new_text in zip(old_texts, new_texts): - text_input = self.tokenizer( - [old_text, new_text], - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0] - old_emb, new_emb = text_embeddings - old_embs.append(old_emb) - new_embs.append(new_emb) - - # identify corresponding destinations for each token in old_emb - idxs_replaces = [] - for old_text, new_text in zip(old_texts, new_texts): - tokens_a = self.tokenizer(old_text).input_ids - tokens_b = self.tokenizer(new_text).input_ids - tokens_a = [self.tokenizer.encode("a ")[1] if self.tokenizer.decode(t) == "an" else t for t in tokens_a] - tokens_b = [self.tokenizer.encode("a ")[1] if self.tokenizer.decode(t) == "an" else t for t in tokens_b] - num_orig_tokens = len(tokens_a) - idxs_replace = [] - j = 0 - for i in range(num_orig_tokens): - curr_token = tokens_a[i] - while tokens_b[j] != curr_token: - j += 1 - idxs_replace.append(j) - j += 1 - while j < 77: - idxs_replace.append(j) - j += 1 - while len(idxs_replace) < 77: - idxs_replace.append(76) - idxs_replaces.append(idxs_replace) - - # prepare batch: for each pair of sentences, old context and new values - contexts, valuess = [], [] - for old_emb, new_emb, idxs_replace in zip(old_embs, new_embs, idxs_replaces): - context = old_emb.detach() - values = [] - with torch.no_grad(): - for layer in self.projection_matrices: - values.append(layer(new_emb[idxs_replace]).detach()) - contexts.append(context) - valuess.append(values) - - # edit the model - for layer_num in range(len(self.projection_matrices)): - # mat1 = \lambda W + \sum{v k^T} - mat1 = lamb * self.projection_matrices[layer_num].weight - - # mat2 = \lambda I + \sum{k k^T} - mat2 = lamb * torch.eye( - self.projection_matrices[layer_num].weight.shape[1], - device=self.projection_matrices[layer_num].weight.device, - ) - - # aggregate sums for mat1, mat2 - for context, values in zip(contexts, valuess): - context_vector = context.reshape(context.shape[0], context.shape[1], 1) - context_vector_T = context.reshape(context.shape[0], 1, context.shape[1]) - value_vector = values[layer_num].reshape(values[layer_num].shape[0], values[layer_num].shape[1], 1) - for_mat1 = (value_vector @ context_vector_T).sum(dim=0) - for_mat2 = (context_vector @ context_vector_T).sum(dim=0) - mat1 += for_mat1 - mat2 += for_mat2 - - # update projection matrix - self.projection_matrices[layer_num].weight = torch.nn.Parameter(mat1 @ torch.inverse(mat2)) - - @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, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - 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. - - Examples: - - ```py - >>> import torch - >>> from diffusers import StableDiffusionModelEditingPipeline - - >>> model_ckpt = "CompVis/stable-diffusion-v1-4" - >>> pipe = StableDiffusionModelEditingPipeline.from_pretrained(model_ckpt) - - >>> pipe = pipe.to("cuda") - - >>> source_prompt = "A pack of roses" - >>> destination_prompt = "A pack of blue roses" - >>> pipe.edit_model(source_prompt, destination_prompt) - - >>> prompt = "A field of roses" - >>> image = pipe(prompt).images[0] - ``` - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - # 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - 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, - lora_scale=text_encoder_lora_scale, - 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=device) - timesteps = self.scheduler.timesteps - - # 5. 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, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - ).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py deleted file mode 100644 index 2331157ba5e6c579329b23a7d0799c965af0cf0e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py +++ /dev/null @@ -1,796 +0,0 @@ -# Copyright 2025 ParaDiGMS authors 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 CLIPImageProcessor, CLIPTextModel, CLIPTokenizer - -from ....image_processor import VaeImageProcessor -from ....loaders import FromSingleFileMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ....models import AutoencoderKL, UNet2DConditionModel -from ....models.lora import adjust_lora_scale_text_encoder -from ....schedulers import KarrasDiffusionSchedulers -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 DiffusionPipeline, StableDiffusionMixin -from ...stable_diffusion.pipeline_output import StableDiffusionPipelineOutput -from ...stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import DDPMParallelScheduler - >>> from diffusers import StableDiffusionParadigmsPipeline - - >>> scheduler = DDPMParallelScheduler.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="scheduler") - - >>> pipe = StableDiffusionParadigmsPipeline.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", scheduler=scheduler, torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> ngpu, batch_per_device = torch.cuda.device_count(), 5 - >>> pipe.wrapped_unet = torch.nn.DataParallel(pipe.unet, device_ids=[d for d in range(ngpu)]) - - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> image = pipe(prompt, parallel=ngpu * batch_per_device, num_inference_steps=1000).images[0] - ``` -""" - - -class StableDiffusionParadigmsPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-to-image generation using a parallelized version of Stable Diffusion. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - 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: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # attribute to wrap the unet with torch.nn.DataParallel when running multiple denoising steps on multiple GPUs - self.wrapped_unet = self.unet - - # 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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - 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}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def _cumsum(self, input, dim, debug=False): - if debug: - # cumsum_cuda_kernel does not have a deterministic implementation - # so perform cumsum on cpu for debugging purposes - return torch.cumsum(input.cpu().float(), dim=dim).to(input.device) - else: - return torch.cumsum(input, dim=dim) - - @torch.no_grad() - @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: 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, - 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, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - debug: bool = False, - clip_skip: int = None, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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. - parallel (`int`, *optional*, defaults to 10): - The batch size to use when doing parallel sampling. More parallelism may lead to faster inference but - requires higher memory usage and can also require more total FLOPs. - tolerance (`float`, *optional*, defaults to 0.1): - The error tolerance for determining when to slide the batch window forward for parallel sampling. Lower - tolerance usually leads to less or no degradation. Higher tolerance is faster but can risk degradation - of sample quality. The tolerance is specified as a ratio of the scheduler's noise magnitude. - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - debug (`bool`, *optional*, defaults to `False`): - Whether or not to run in debug mode. In debug mode, `torch.cumsum` is evaluated using the CPU. - 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. - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - # 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 = guidance_scale > 1.0 - - # 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=device) - - # 5. 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, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - extra_step_kwargs.pop("generator", None) - - # # 7. Denoising loop - scheduler = self.scheduler - parallel = min(parallel, len(scheduler.timesteps)) - - begin_idx = 0 - end_idx = parallel - latents_time_evolution_buffer = torch.stack([latents] * (len(scheduler.timesteps) + 1)) - - # We must make sure the noise of stochastic schedulers such as DDPM is sampled only once per timestep. - # Sampling inside the parallel denoising loop will mess this up, so we pre-sample the noise vectors outside the denoising loop. - noise_array = torch.zeros_like(latents_time_evolution_buffer) - for j in range(len(scheduler.timesteps)): - base_noise = randn_tensor( - shape=latents.shape, generator=generator, device=latents.device, dtype=prompt_embeds.dtype - ) - noise = (self.scheduler._get_variance(scheduler.timesteps[j]) ** 0.5) * base_noise - noise_array[j] = noise.clone() - - # We specify the error tolerance as a ratio of the scheduler's noise magnitude. We similarly compute the error tolerance - # outside of the denoising loop to avoid recomputing it at every step. - # We will be dividing the norm of the noise, so we store its inverse here to avoid a division at every step. - inverse_variance_norm = 1.0 / torch.tensor( - [scheduler._get_variance(scheduler.timesteps[j]) for j in range(len(scheduler.timesteps))] + [0] - ).to(noise_array.device) - latent_dim = noise_array[0, 0].numel() - inverse_variance_norm = inverse_variance_norm[:, None] / latent_dim - - scaled_tolerance = tolerance**2 - - with self.progress_bar(total=num_inference_steps) as progress_bar: - steps = 0 - while begin_idx < len(scheduler.timesteps): - # these have shape (parallel_dim, 2*batch_size, ...) - # parallel_len is at most parallel, but could be less if we are at the end of the timesteps - # we are processing batch window of timesteps spanning [begin_idx, end_idx) - parallel_len = end_idx - begin_idx - - block_prompt_embeds = torch.stack([prompt_embeds] * parallel_len) - block_latents = latents_time_evolution_buffer[begin_idx:end_idx] - block_t = scheduler.timesteps[begin_idx:end_idx, None].repeat(1, batch_size * num_images_per_prompt) - t_vec = block_t - if do_classifier_free_guidance: - t_vec = t_vec.repeat(1, 2) - - # expand the latents if we are doing classifier free guidance - latent_model_input = ( - torch.cat([block_latents] * 2, dim=1) if do_classifier_free_guidance else block_latents - ) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t_vec) - - # if parallel_len is small, no need to use multiple GPUs - net = self.wrapped_unet if parallel_len > 3 else self.unet - # predict the noise residual, shape is now [parallel_len * 2 * batch_size * num_images_per_prompt, ...] - model_output = net( - latent_model_input.flatten(0, 1), - t_vec.flatten(0, 1), - encoder_hidden_states=block_prompt_embeds.flatten(0, 1), - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - per_latent_shape = model_output.shape[1:] - if do_classifier_free_guidance: - model_output = model_output.reshape( - parallel_len, 2, batch_size * num_images_per_prompt, *per_latent_shape - ) - noise_pred_uncond, noise_pred_text = model_output[:, 0], model_output[:, 1] - model_output = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - model_output = model_output.reshape( - parallel_len * batch_size * num_images_per_prompt, *per_latent_shape - ) - - block_latents_denoise = scheduler.batch_step_no_noise( - model_output=model_output, - timesteps=block_t.flatten(0, 1), - sample=block_latents.flatten(0, 1), - **extra_step_kwargs, - ).reshape(block_latents.shape) - - # back to shape (parallel_dim, batch_size, ...) - # now we want to add the pre-sampled noise - # parallel sampling algorithm requires computing the cumulative drift from the beginning - # of the window, so we need to compute cumulative sum of the deltas and the pre-sampled noises. - delta = block_latents_denoise - block_latents - cumulative_delta = self._cumsum(delta, dim=0, debug=debug) - cumulative_noise = self._cumsum(noise_array[begin_idx:end_idx], dim=0, debug=debug) - - # if we are using an ODE-like scheduler (like DDIM), we don't want to add noise - if scheduler._is_ode_scheduler: - cumulative_noise = 0 - - block_latents_new = ( - latents_time_evolution_buffer[begin_idx][None,] + cumulative_delta + cumulative_noise - ) - cur_error = torch.linalg.norm( - (block_latents_new - latents_time_evolution_buffer[begin_idx + 1 : end_idx + 1]).reshape( - parallel_len, batch_size * num_images_per_prompt, -1 - ), - dim=-1, - ).pow(2) - error_ratio = cur_error * inverse_variance_norm[begin_idx + 1 : end_idx + 1] - - # find the first index of the vector error_ratio that is greater than error tolerance - # we can shift the window for the next iteration up to this index - error_ratio = torch.nn.functional.pad( - error_ratio, (0, 0, 0, 1), value=1e9 - ) # handle the case when everything is below ratio, by padding the end of parallel_len dimension - any_error_at_time = torch.max(error_ratio > scaled_tolerance, dim=1).values.int() - ind = torch.argmax(any_error_at_time).item() - - # compute the new begin and end idxs for the window - new_begin_idx = begin_idx + min(1 + ind, parallel) - new_end_idx = min(new_begin_idx + parallel, len(scheduler.timesteps)) - - # store the computed latents for the current window in the global buffer - latents_time_evolution_buffer[begin_idx + 1 : end_idx + 1] = block_latents_new - # initialize the new sliding window latents with the end of the current window, - # should be better than random initialization - latents_time_evolution_buffer[end_idx : new_end_idx + 1] = latents_time_evolution_buffer[end_idx][ - None, - ] - - steps += 1 - - progress_bar.update(new_begin_idx - begin_idx) - if callback is not None and steps % callback_steps == 0: - callback(begin_idx, block_t[begin_idx], latents_time_evolution_buffer[begin_idx]) - - begin_idx = new_begin_idx - end_idx = new_end_idx - - latents = latents_time_evolution_buffer[-1] - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py deleted file mode 100644 index f9034a58441c21f04ade475bbe75aae6762265fd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py +++ /dev/null @@ -1,1310 +0,0 @@ -# Copyright 2025 Pix2Pix Zero Authors 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, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import ( - BlipForConditionalGeneration, - BlipProcessor, - CLIPImageProcessor, - CLIPTextModel, - CLIPTokenizer, -) - -from ....image_processor import PipelineImageInput, VaeImageProcessor -from ....loaders import StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ....models import AutoencoderKL, UNet2DConditionModel -from ....models.attention_processor import Attention -from ....models.lora import adjust_lora_scale_text_encoder -from ....schedulers import DDIMScheduler, DDPMScheduler, EulerAncestralDiscreteScheduler, LMSDiscreteScheduler -from ....schedulers.scheduling_ddim_inverse import DDIMInverseScheduler -from ....utils import ( - PIL_INTERPOLATION, - USE_PEFT_BACKEND, - BaseOutput, - deprecate, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ....utils.torch_utils import randn_tensor -from ...pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ...stable_diffusion.pipeline_output import StableDiffusionPipelineOutput -from ...stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class Pix2PixInversionPipelineOutput(BaseOutput, TextualInversionLoaderMixin): - """ - Output class for Stable Diffusion pipelines. - - 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, - 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] - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import requests - >>> import torch - - >>> from diffusers import DDIMScheduler, StableDiffusionPix2PixZeroPipeline - - - >>> def download(embedding_url, local_filepath): - ... r = requests.get(embedding_url) - ... with open(local_filepath, "wb") as f: - ... f.write(r.content) - - - >>> model_ckpt = "CompVis/stable-diffusion-v1-4" - >>> pipeline = StableDiffusionPix2PixZeroPipeline.from_pretrained(model_ckpt, torch_dtype=torch.float16) - >>> pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config) - >>> pipeline.to("cuda") - - >>> prompt = "a high resolution painting of a cat in the style of van gough" - >>> source_emb_url = "https://hf.co/datasets/sayakpaul/sample-datasets/resolve/main/cat.pt" - >>> target_emb_url = "https://hf.co/datasets/sayakpaul/sample-datasets/resolve/main/dog.pt" - - >>> for url in [source_emb_url, target_emb_url]: - ... download(url, url.split("/")[-1]) - - >>> src_embeds = torch.load(source_emb_url.split("/")[-1]) - >>> target_embeds = torch.load(target_emb_url.split("/")[-1]) - >>> images = pipeline( - ... prompt, - ... source_embeds=src_embeds, - ... target_embeds=target_embeds, - ... num_inference_steps=50, - ... cross_attention_guidance_amount=0.15, - ... ).images - - >>> images[0].save("edited_image_dog.png") - ``` -""" - -EXAMPLE_INVERT_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from transformers import BlipForConditionalGeneration, BlipProcessor - >>> from diffusers import DDIMScheduler, DDIMInverseScheduler, StableDiffusionPix2PixZeroPipeline - - >>> import requests - >>> from PIL import Image - - >>> captioner_id = "Salesforce/blip-image-captioning-base" - >>> processor = BlipProcessor.from_pretrained(captioner_id) - >>> model = BlipForConditionalGeneration.from_pretrained( - ... captioner_id, torch_dtype=torch.float16, low_cpu_mem_usage=True - ... ) - - >>> sd_model_ckpt = "CompVis/stable-diffusion-v1-4" - >>> pipeline = StableDiffusionPix2PixZeroPipeline.from_pretrained( - ... sd_model_ckpt, - ... caption_generator=model, - ... caption_processor=processor, - ... torch_dtype=torch.float16, - ... safety_checker=None, - ... ) - - >>> pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config) - >>> pipeline.inverse_scheduler = DDIMInverseScheduler.from_config(pipeline.scheduler.config) - >>> pipeline.enable_model_cpu_offload() - - >>> img_url = "https://github.com/pix2pixzero/pix2pix-zero/raw/main/assets/test_images/cats/cat_6.png" - - >>> raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB").resize((512, 512)) - >>> # generate caption - >>> caption = pipeline.generate_caption(raw_image) - - >>> # "a photography of a cat with flowers and dai dai daie - daie - daie kasaii" - >>> inv_latents = pipeline.invert(caption, image=raw_image).latents - >>> # we need to generate source and target embeds - - >>> source_prompts = ["a cat sitting on the street", "a cat playing in the field", "a face of a cat"] - - >>> target_prompts = ["a dog sitting on the street", "a dog playing in the field", "a face of a dog"] - - >>> source_embeds = pipeline.get_embeds(source_prompts) - >>> target_embeds = pipeline.get_embeds(target_prompts) - >>> # the latents can then be used to edit a real image - >>> # when using Stable Diffusion 2 or other models that use v-prediction - >>> # set `cross_attention_guidance_amount` to 0.01 or less to avoid input latent gradient explosion - - >>> image = pipeline( - ... caption, - ... source_embeds=source_embeds, - ... target_embeds=target_embeds, - ... num_inference_steps=50, - ... cross_attention_guidance_amount=0.15, - ... generator=generator, - ... latents=inv_latents, - ... negative_prompt=caption, - ... ).images[0] - >>> image.save("edited_image.png") - ``` -""" - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess -def preprocess(image): - 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): - return image - elif isinstance(image, PIL.Image.Image): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - w, h = image[0].size - w, h = (x - x % 8 for x in (w, h)) # resize to integer multiple of 8 - - image = [np.array(i.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image] - image = np.concatenate(image, axis=0) - image = np.array(image).astype(np.float32) / 255.0 - image = image.transpose(0, 3, 1, 2) - image = 2.0 * image - 1.0 - image = torch.from_numpy(image) - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, dim=0) - return image - - -def prepare_unet(unet: UNet2DConditionModel): - """Modifies the UNet (`unet`) to perform Pix2Pix Zero optimizations.""" - pix2pix_zero_attn_procs = {} - for name in unet.attn_processors.keys(): - module_name = name.replace(".processor", "") - module = unet.get_submodule(module_name) - if "attn2" in name: - pix2pix_zero_attn_procs[name] = Pix2PixZeroAttnProcessor(is_pix2pix_zero=True) - module.requires_grad_(True) - else: - pix2pix_zero_attn_procs[name] = Pix2PixZeroAttnProcessor(is_pix2pix_zero=False) - module.requires_grad_(False) - - unet.set_attn_processor(pix2pix_zero_attn_procs) - return unet - - -class Pix2PixZeroL2Loss: - def __init__(self): - self.loss = 0.0 - - def compute_loss(self, predictions, targets): - self.loss += ((predictions - targets) ** 2).sum((1, 2)).mean(0) - - -class Pix2PixZeroAttnProcessor: - """An attention processor class to store the attention weights. - In Pix2Pix Zero, it happens during computations in the cross-attention blocks.""" - - def __init__(self, is_pix2pix_zero=False): - self.is_pix2pix_zero = is_pix2pix_zero - if self.is_pix2pix_zero: - self.reference_cross_attn_map = {} - - def __call__( - self, - attn: Attention, - hidden_states, - encoder_hidden_states=None, - attention_mask=None, - timestep=None, - loss=None, - ): - batch_size, sequence_length, _ = hidden_states.shape - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query = attn.head_to_batch_dim(query) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - - attention_probs = attn.get_attention_scores(query, key, attention_mask) - if self.is_pix2pix_zero and timestep is not None: - # new bookkeeping to save the attention weights. - if loss is None: - self.reference_cross_attn_map[timestep.item()] = attention_probs.detach().cpu() - # compute loss - elif loss is not None: - prev_attn_probs = self.reference_cross_attn_map.pop(timestep.item()) - loss.compute_loss(attention_probs, prev_attn_probs.to(attention_probs.device)) - - hidden_states = torch.bmm(attention_probs, value) - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - return hidden_states - - -class StableDiffusionPix2PixZeroPipeline(DiffusionPipeline, StableDiffusionMixin): - r""" - Pipeline for pixel-level image editing using Pix2Pix Zero. Based on 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.) - - 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`], [`EulerAncestralDiscreteScheduler`], or [`DDPMScheduler`]. - 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/runwayml/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`. - requires_safety_checker (bool): - Whether the pipeline requires a safety checker. We recommend setting it to True if you're using the - pipeline publicly. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = [ - "safety_checker", - "feature_extractor", - "caption_generator", - "caption_processor", - "inverse_scheduler", - ] - _exclude_from_cpu_offload = ["safety_checker"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: Union[DDPMScheduler, DDIMScheduler, EulerAncestralDiscreteScheduler, LMSDiscreteScheduler], - feature_extractor: CLIPImageProcessor, - safety_checker: StableDiffusionSafetyChecker, - inverse_scheduler: DDIMInverseScheduler, - caption_generator: BlipForConditionalGeneration, - caption_processor: BlipProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - caption_processor=caption_processor, - caption_generator=caption_generator, - inverse_scheduler=inverse_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.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - source_embeds, - target_embeds, - callback_steps, - prompt_embeds=None, - ): - if (callback_steps is None) or ( - 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 source_embeds is None and target_embeds is None: - raise ValueError("`source_embeds` and `target_embeds` cannot be undefined.") - - 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)}") - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @torch.no_grad() - def generate_caption(self, images): - """Generates caption for a given image.""" - text = "a photography of" - - prev_device = self.caption_generator.device - - device = self._execution_device - inputs = self.caption_processor(images, text, return_tensors="pt").to( - device=device, dtype=self.caption_generator.dtype - ) - self.caption_generator.to(device) - outputs = self.caption_generator.generate(**inputs, max_new_tokens=128) - - # offload caption generator - self.caption_generator.to(prev_device) - - caption = self.caption_processor.batch_decode(outputs, skip_special_tokens=True)[0] - return caption - - def construct_direction(self, embs_source: torch.Tensor, embs_target: torch.Tensor): - """Constructs the edit direction to steer the image generation process semantically.""" - 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: - num_prompts = len(prompt) - embeds = [] - for i in range(0, num_prompts, batch_size): - prompt_slice = prompt[i : i + batch_size] - - input_ids = self.tokenizer( - prompt_slice, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ).input_ids - - input_ids = input_ids.to(self.text_encoder.device) - embeds.append(self.text_encoder(input_ids)[0]) - - return torch.cat(embeds, dim=0).mean(0)[None] - - def prepare_image_latents(self, image, batch_size, dtype, device, generator=None): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - if image.shape[1] == 4: - latents = image - - else: - 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 isinstance(generator, list): - latents = [ - self.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size) - ] - latents = torch.cat(latents, dim=0) - else: - latents = self.vae.encode(image).latent_dist.sample(generator) - - latents = self.vae.config.scaling_factor * latents - - if batch_size != latents.shape[0]: - if batch_size % latents.shape[0] == 0: - # expand image_latents for batch_size - deprecation_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {latents.shape[0]} initial" - " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" - " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" - " your script to pass as many initial images as text prompts to suppress this warning." - ) - deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) - additional_latents_per_image = batch_size // latents.shape[0] - latents = torch.cat([latents] * additional_latents_per_image, dim=0) - else: - raise ValueError( - f"Cannot duplicate `image` of batch size {latents.shape[0]} to {batch_size} text prompts." - ) - else: - latents = torch.cat([latents], dim=0) - - return latents - - def get_epsilon(self, model_output: torch.Tensor, sample: torch.Tensor, timestep: int): - pred_type = self.inverse_scheduler.config.prediction_type - alpha_prod_t = self.inverse_scheduler.alphas_cumprod[timestep] - - beta_prod_t = 1 - alpha_prod_t - - if pred_type == "epsilon": - return model_output - elif pred_type == "sample": - return (sample - alpha_prod_t ** (0.5) * model_output) / beta_prod_t ** (0.5) - elif pred_type == "v_prediction": - return (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - else: - raise ValueError( - f"prediction_type given as {pred_type} must be one of `epsilon`, `sample`, or `v_prediction`" - ) - - def auto_corr_loss(self, hidden_states, generator=None): - reg_loss = 0.0 - for i in range(hidden_states.shape[0]): - for j in range(hidden_states.shape[1]): - noise = hidden_states[i : i + 1, j : j + 1, :, :] - while True: - roll_amount = torch.randint(noise.shape[2] // 2, (1,), generator=generator).item() - reg_loss += (noise * torch.roll(noise, shifts=roll_amount, dims=2)).mean() ** 2 - reg_loss += (noise * torch.roll(noise, shifts=roll_amount, dims=3)).mean() ** 2 - - if noise.shape[2] <= 8: - break - noise = F.avg_pool2d(noise, kernel_size=2) - return reg_loss - - def kl_divergence(self, hidden_states): - mean = hidden_states.mean() - var = hidden_states.var() - return var + mean**2 - 1 - torch.log(var + 1e-7) - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Optional[Union[str, List[str]]] = None, - source_embeds: torch.Tensor = None, - target_embeds: torch.Tensor = 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, - cross_attention_guidance_amount: float = 0.1, - output_type: Optional[str] = "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, - ): - 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. - source_embeds (`torch.Tensor`): - Source concept embeddings. Generation of the embeddings as per the [original - paper](https://huggingface.co/papers/2302.03027). Used in discovering the edit direction. - target_embeds (`torch.Tensor`): - Target concept embeddings. Generation of the embeddings as per the [original - paper](https://huggingface.co/papers/2302.03027). Used in discovering the edit direction. - 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` 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 ge 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. - cross_attention_guidance_amount (`float`, defaults to 0.1): - Amount of guidance needed from the reference cross-attention maps. - 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. - 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. - Examples: - - 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. Define the spatial resolutions. - 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, - source_embeds, - target_embeds, - callback_steps, - prompt_embeds, - ) - - # 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] - if cross_attention_kwargs is None: - cross_attention_kwargs = {} - - 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 = guidance_scale > 1.0 - - # 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=device) - timesteps = self.scheduler.timesteps - - # 5. Generate the inverted noise from the input image or any other image - # generated from the input prompt. - 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_init = latents.clone() - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 8. Rejig the UNet so that we can obtain the cross-attenion maps and - # use them for guiding the subsequent image generation. - self.unet = prepare_unet(self.unet) - - # 7. Denoising loop where we obtain the cross-attention maps. - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs={"timestep": t}, - ).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - # 8. Compute the edit directions. - edit_direction = self.construct_direction(source_embeds, target_embeds).to(prompt_embeds.device) - - # 9. Edit the prompt embeddings as per the edit directions discovered. - prompt_embeds_edit = prompt_embeds.clone() - prompt_embeds_edit[1:2] += edit_direction - - # 10. Second denoising loop to generate the edited image. - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - latents = latents_init - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # we want to learn the latent such that it steers the generation - # process towards the edited direction, so make the make initial - # noise learnable - x_in = latent_model_input.detach().clone() - x_in.requires_grad = True - - # optimizer - opt = torch.optim.SGD([x_in], lr=cross_attention_guidance_amount) - - with torch.enable_grad(): - # initialize loss - loss = Pix2PixZeroL2Loss() - - # predict the noise residual - noise_pred = self.unet( - x_in, - t, - encoder_hidden_states=prompt_embeds_edit.detach(), - cross_attention_kwargs={"timestep": t, "loss": loss}, - ).sample - - loss.loss.backward(retain_graph=False) - opt.step() - - # recompute the noise - noise_pred = self.unet( - x_in.detach(), - t, - encoder_hidden_states=prompt_embeds_edit, - cross_attention_kwargs={"timestep": None}, - ).sample - - latents = x_in.detach().chunk(2)[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 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) - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_INVERT_DOC_STRING) - def invert( - self, - prompt: Optional[str] = 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, - cross_attention_guidance_amount: float = 0.1, - output_type: Optional[str] = "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, - lambda_auto_corr: float = 20.0, - lambda_kl: float = 20.0, - num_reg_steps: int = 5, - num_auto_corr_rolls: int = 5, - ): - r""" - Function used to generate inverted latents given a prompt and image. - - 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` `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): - 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 1): - 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. - 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 ge 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. - cross_attention_guidance_amount (`float`, defaults to 0.1): - Amount of guidance needed from the reference cross-attention maps. - 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. - lambda_auto_corr (`float`, *optional*, defaults to 20.0): - Lambda parameter to control auto correction - lambda_kl (`float`, *optional*, defaults to 20.0): - Lambda parameter to control Kullback–Leibler divergence output - num_reg_steps (`int`, *optional*, defaults to 5): - Number of regularization loss steps - num_auto_corr_rolls (`int`, *optional*, defaults to 5): - Number of auto correction roll steps - - Examples: - - Returns: - [`~pipelines.stable_diffusion.pipeline_stable_diffusion_pix2pix_zero.Pix2PixInversionPipelineOutput`] or - `tuple`: - [`~pipelines.stable_diffusion.pipeline_stable_diffusion_pix2pix_zero.Pix2PixInversionPipelineOutput`] if - `return_dict` is True, otherwise a `tuple. When returning a tuple, the first element is the inverted - latents tensor and then second is the corresponding decoded image. - """ - # 1. 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 cross_attention_kwargs is None: - cross_attention_kwargs = {} - - 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 = guidance_scale > 1.0 - - # 3. Preprocess image - image = self.image_processor.preprocess(image) - - # 4. Prepare latent variables - latents = self.prepare_image_latents(image, batch_size, self.vae.dtype, device, generator) - - # 5. Encode input prompt - num_images_per_prompt = 1 - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - do_classifier_free_guidance, - prompt_embeds=prompt_embeds, - ) - # 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.inverse_scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.inverse_scheduler.timesteps - - # 6. Rejig the UNet so that we can obtain the cross-attenion maps and - # use them for guiding the subsequent image generation. - self.unet = prepare_unet(self.unet) - - # 7. Denoising loop where we obtain the cross-attention maps. - num_warmup_steps = len(timesteps) - num_inference_steps * self.inverse_scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.inverse_scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs={"timestep": t}, - ).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # regularization of the noise prediction - with torch.enable_grad(): - for _ in range(num_reg_steps): - if lambda_auto_corr > 0: - for _ in range(num_auto_corr_rolls): - var = torch.autograd.Variable(noise_pred.detach().clone(), requires_grad=True) - - # Derive epsilon from model output before regularizing to IID standard normal - var_epsilon = self.get_epsilon(var, latent_model_input.detach(), t) - - l_ac = self.auto_corr_loss(var_epsilon, generator=generator) - l_ac.backward() - - grad = var.grad.detach() / num_auto_corr_rolls - noise_pred = noise_pred - lambda_auto_corr * grad - - if lambda_kl > 0: - var = torch.autograd.Variable(noise_pred.detach().clone(), requires_grad=True) - - # Derive epsilon from model output before regularizing to IID standard normal - var_epsilon = self.get_epsilon(var, latent_model_input.detach(), t) - - l_kld = self.kl_divergence(var_epsilon) - l_kld.backward() - - grad = var.grad.detach() - noise_pred = noise_pred - lambda_kl * grad - - noise_pred = noise_pred.detach() - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.inverse_scheduler.step(noise_pred, t, latents).prev_sample - - # call the callback, if provided - if i == len(timesteps) - 1 or ( - (i + 1) > num_warmup_steps and (i + 1) % self.inverse_scheduler.order == 0 - ): - progress_bar.update() - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - inverted_latents = latents.detach().clone() - - # 8. Post-processing - image = self.vae.decode(latents / self.vae.config.scaling_factor, 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 (inverted_latents, image) - - return Pix2PixInversionPipelineOutput(latents=inverted_latents, images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stochastic_karras_ve/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stochastic_karras_ve/__init__.py deleted file mode 100644 index 15c9a8c27f98dd7e1913bd57dfd5e8dae71172b4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stochastic_karras_ve/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -from typing import TYPE_CHECKING - -from ....utils import DIFFUSERS_SLOW_IMPORT, _LazyModule - - -_import_structure = {"pipeline_stochastic_karras_ve": ["KarrasVePipeline"]} - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - from .pipeline_stochastic_karras_ve import KarrasVePipeline - -else: - import sys - - sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stochastic_karras_ve/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stochastic_karras_ve/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index d4695d3721934ed9f770b6e765e47a9a64586d66..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stochastic_karras_ve/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stochastic_karras_ve/__pycache__/pipeline_stochastic_karras_ve.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stochastic_karras_ve/__pycache__/pipeline_stochastic_karras_ve.cpython-310.pyc deleted file mode 100644 index 11a6368dff748356288c193d210b9a1ba44bb70e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stochastic_karras_ve/__pycache__/pipeline_stochastic_karras_ve.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stochastic_karras_ve/pipeline_stochastic_karras_ve.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stochastic_karras_ve/pipeline_stochastic_karras_ve.py deleted file mode 100644 index 50b8b0bcbc1df5ed2cc1b702bdb65fd627ef69f8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/stochastic_karras_ve/pipeline_stochastic_karras_ve.py +++ /dev/null @@ -1,128 +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 List, Optional, Tuple, Union - -import torch - -from ....models import UNet2DModel -from ....schedulers import KarrasVeScheduler -from ....utils.torch_utils import randn_tensor -from ...pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -class KarrasVePipeline(DiffusionPipeline): - r""" - Pipeline for unconditional image generation. - - Parameters: - unet ([`UNet2DModel`]): - A `UNet2DModel` to denoise the encoded image. - scheduler ([`KarrasVeScheduler`]): - A scheduler to be used in combination with `unet` to denoise the encoded image. - """ - - # add type hints for linting - unet: UNet2DModel - scheduler: KarrasVeScheduler - - def __init__(self, unet: UNet2DModel, scheduler: KarrasVeScheduler): - super().__init__() - self.register_modules(unet=unet, scheduler=scheduler) - - @torch.no_grad() - 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", - return_dict: bool = True, - **kwargs, - ) -> Union[Tuple, ImagePipelineOutput]: - r""" - The call function to the pipeline for generation. - - Args: - batch_size (`int`, *optional*, defaults to 1): - The number of images to generate. - generator (`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 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - 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 [`ImagePipelineOutput`] instead of a plain tuple. - - Example: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images. - """ - - img_size = self.unet.config.sample_size - shape = (batch_size, 3, img_size, img_size) - - model = self.unet - - # sample x_0 ~ N(0, sigma_0^2 * I) - sample = randn_tensor(shape, generator=generator, device=self.device) * self.scheduler.init_noise_sigma - - self.scheduler.set_timesteps(num_inference_steps) - - for t in self.progress_bar(self.scheduler.timesteps): - # here sigma_t == t_i from the paper - sigma = self.scheduler.schedule[t] - sigma_prev = self.scheduler.schedule[t - 1] if t > 0 else 0 - - # 1. Select temporarily increased noise level sigma_hat - # 2. Add new noise to move from sample_i to sample_hat - sample_hat, sigma_hat = self.scheduler.add_noise_to_input(sample, sigma, generator=generator) - - # 3. Predict the noise residual given the noise magnitude `sigma_hat` - # The model inputs and output are adjusted by following eq. (213) in [1]. - model_output = (sigma_hat / 2) * model((sample_hat + 1) / 2, sigma_hat / 2).sample - - # 4. Evaluate dx/dt at sigma_hat - # 5. Take Euler step from sigma to sigma_prev - step_output = self.scheduler.step(model_output, sigma_hat, sigma_prev, sample_hat) - - if sigma_prev != 0: - # 6. Apply 2nd order correction - # The model inputs and output are adjusted by following eq. (213) in [1]. - model_output = (sigma_prev / 2) * model((step_output.prev_sample + 1) / 2, sigma_prev / 2).sample - step_output = self.scheduler.step_correct( - model_output, - sigma_hat, - sigma_prev, - sample_hat, - step_output.prev_sample, - step_output["derivative"], - ) - sample = step_output.prev_sample - - sample = (sample / 2 + 0.5).clamp(0, 1) - image = sample.cpu().permute(0, 2, 3, 1).numpy() - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__init__.py deleted file mode 100644 index 8ea6ef6e2f65b96aebebdf72cb80135003e4f08d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__init__.py +++ /dev/null @@ -1,71 +0,0 @@ -from typing import TYPE_CHECKING - -from ....utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - is_torch_available, - is_transformers_available, - is_transformers_version, -) - - -_dummy_objects = {} -_import_structure = {} - -try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.25.0")): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ....utils.dummy_torch_and_transformers_objects import ( - VersatileDiffusionDualGuidedPipeline, - VersatileDiffusionImageVariationPipeline, - VersatileDiffusionPipeline, - VersatileDiffusionTextToImagePipeline, - ) - - _dummy_objects.update( - { - "VersatileDiffusionDualGuidedPipeline": VersatileDiffusionDualGuidedPipeline, - "VersatileDiffusionImageVariationPipeline": VersatileDiffusionImageVariationPipeline, - "VersatileDiffusionPipeline": VersatileDiffusionPipeline, - "VersatileDiffusionTextToImagePipeline": VersatileDiffusionTextToImagePipeline, - } - ) -else: - _import_structure["modeling_text_unet"] = ["UNetFlatConditionModel"] - _import_structure["pipeline_versatile_diffusion"] = ["VersatileDiffusionPipeline"] - _import_structure["pipeline_versatile_diffusion_dual_guided"] = ["VersatileDiffusionDualGuidedPipeline"] - _import_structure["pipeline_versatile_diffusion_image_variation"] = ["VersatileDiffusionImageVariationPipeline"] - _import_structure["pipeline_versatile_diffusion_text_to_image"] = ["VersatileDiffusionTextToImagePipeline"] - - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.25.0")): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ....utils.dummy_torch_and_transformers_objects import ( - VersatileDiffusionDualGuidedPipeline, - VersatileDiffusionImageVariationPipeline, - VersatileDiffusionPipeline, - VersatileDiffusionTextToImagePipeline, - ) - else: - from .pipeline_versatile_diffusion import VersatileDiffusionPipeline - from .pipeline_versatile_diffusion_dual_guided import VersatileDiffusionDualGuidedPipeline - from .pipeline_versatile_diffusion_image_variation import VersatileDiffusionImageVariationPipeline - from .pipeline_versatile_diffusion_text_to_image import VersatileDiffusionTextToImagePipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 8796d7ccd650c540174b742b149c9f99a33f8ec7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/modeling_text_unet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/modeling_text_unet.cpython-310.pyc deleted file mode 100644 index ba6bd4e2f38c1d38faaedb8aea5bda20367c91a9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/modeling_text_unet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion.cpython-310.pyc deleted file mode 100644 index e467821820e878924c88a504717eb5636c09eca3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion_dual_guided.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion_dual_guided.cpython-310.pyc deleted file mode 100644 index aba6aa10a348fa34999cb406aeb8b559303e91d9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion_dual_guided.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion_image_variation.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion_image_variation.cpython-310.pyc deleted file mode 100644 index 221d85a13fddaf506f50d0afbe384887f34dadfd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion_image_variation.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion_text_to_image.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion_text_to_image.cpython-310.pyc deleted file mode 100644 index c8a5962653e7ac9157895e773155393a957c68a7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/__pycache__/pipeline_versatile_diffusion_text_to_image.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py deleted file mode 100644 index eda950998d67017d144eff4fcdd0fbe3d1f98b19..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py +++ /dev/null @@ -1,2446 +0,0 @@ -from typing import Any, Dict, List, Optional, Tuple, Union - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F - -from diffusers.utils import deprecate - -from ....configuration_utils import ConfigMixin, register_to_config -from ....models import ModelMixin -from ....models.activations import get_activation -from ....models.attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - Attention, - AttentionProcessor, - AttnAddedKVProcessor, - AttnAddedKVProcessor2_0, - AttnProcessor, -) -from ....models.embeddings import ( - GaussianFourierProjection, - ImageHintTimeEmbedding, - ImageProjection, - ImageTimeEmbedding, - TextImageProjection, - TextImageTimeEmbedding, - TextTimeEmbedding, - TimestepEmbedding, - Timesteps, -) -from ....models.resnet import ResnetBlockCondNorm2D -from ....models.transformers.dual_transformer_2d import DualTransformer2DModel -from ....models.transformers.transformer_2d import Transformer2DModel -from ....models.unets.unet_2d_condition import UNet2DConditionOutput -from ....utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers -from ....utils.torch_utils import apply_freeu - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -def get_down_block( - down_block_type, - num_layers, - in_channels, - out_channels, - temb_channels, - add_downsample, - resnet_eps, - resnet_act_fn, - num_attention_heads, - transformer_layers_per_block, - attention_type, - attention_head_dim, - resnet_groups=None, - cross_attention_dim=None, - downsample_padding=None, - dual_cross_attention=False, - use_linear_projection=False, - only_cross_attention=False, - upcast_attention=False, - resnet_time_scale_shift="default", - resnet_skip_time_act=False, - resnet_out_scale_factor=1.0, - cross_attention_norm=None, - dropout=0.0, -): - down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type - if down_block_type == "DownBlockFlat": - return DownBlockFlat( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - dropout=dropout, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - downsample_padding=downsample_padding, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif down_block_type == "CrossAttnDownBlockFlat": - if cross_attention_dim is None: - raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlockFlat") - return CrossAttnDownBlockFlat( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - dropout=dropout, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - downsample_padding=downsample_padding, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads, - dual_cross_attention=dual_cross_attention, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - raise ValueError(f"{down_block_type} is not supported.") - - -def get_up_block( - up_block_type, - num_layers, - in_channels, - out_channels, - prev_output_channel, - temb_channels, - add_upsample, - resnet_eps, - resnet_act_fn, - num_attention_heads, - transformer_layers_per_block, - resolution_idx, - attention_type, - attention_head_dim, - resnet_groups=None, - cross_attention_dim=None, - dual_cross_attention=False, - use_linear_projection=False, - only_cross_attention=False, - upcast_attention=False, - resnet_time_scale_shift="default", - resnet_skip_time_act=False, - resnet_out_scale_factor=1.0, - cross_attention_norm=None, - dropout=0.0, -): - up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type - if up_block_type == "UpBlockFlat": - return UpBlockFlat( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - elif up_block_type == "CrossAttnUpBlockFlat": - if cross_attention_dim is None: - raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlockFlat") - return CrossAttnUpBlockFlat( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - dropout=dropout, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - resnet_groups=resnet_groups, - cross_attention_dim=cross_attention_dim, - num_attention_heads=num_attention_heads, - dual_cross_attention=dual_cross_attention, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - ) - raise ValueError(f"{up_block_type} is not supported.") - - -class FourierEmbedder(nn.Module): - def __init__(self, num_freqs=64, temperature=100): - super().__init__() - - self.num_freqs = num_freqs - self.temperature = temperature - - freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs) - freq_bands = freq_bands[None, None, None] - self.register_buffer("freq_bands", freq_bands, persistent=False) - - def __call__(self, x): - x = self.freq_bands * x.unsqueeze(-1) - return torch.stack((x.sin(), x.cos()), dim=-1).permute(0, 1, 3, 4, 2).reshape(*x.shape[:2], -1) - - -class GLIGENTextBoundingboxProjection(nn.Module): - def __init__(self, positive_len, out_dim, feature_type, fourier_freqs=8): - super().__init__() - self.positive_len = positive_len - self.out_dim = out_dim - - self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs) - self.position_dim = fourier_freqs * 2 * 4 # 2: sin/cos, 4: xyxy - - if isinstance(out_dim, tuple): - out_dim = out_dim[0] - - if feature_type == "text-only": - self.linears = nn.Sequential( - nn.Linear(self.positive_len + self.position_dim, 512), - nn.SiLU(), - nn.Linear(512, 512), - nn.SiLU(), - nn.Linear(512, out_dim), - ) - self.null_positive_feature = torch.nn.Parameter(torch.zeros([self.positive_len])) - - elif feature_type == "text-image": - self.linears_text = nn.Sequential( - nn.Linear(self.positive_len + self.position_dim, 512), - nn.SiLU(), - nn.Linear(512, 512), - nn.SiLU(), - nn.Linear(512, out_dim), - ) - self.linears_image = nn.Sequential( - nn.Linear(self.positive_len + self.position_dim, 512), - nn.SiLU(), - nn.Linear(512, 512), - nn.SiLU(), - nn.Linear(512, out_dim), - ) - self.null_text_feature = torch.nn.Parameter(torch.zeros([self.positive_len])) - self.null_image_feature = torch.nn.Parameter(torch.zeros([self.positive_len])) - - self.null_position_feature = torch.nn.Parameter(torch.zeros([self.position_dim])) - - def forward( - self, - boxes, - masks, - positive_embeddings=None, - phrases_masks=None, - image_masks=None, - phrases_embeddings=None, - image_embeddings=None, - ): - masks = masks.unsqueeze(-1) - - xyxy_embedding = self.fourier_embedder(boxes) - xyxy_null = self.null_position_feature.view(1, 1, -1) - xyxy_embedding = xyxy_embedding * masks + (1 - masks) * xyxy_null - - if positive_embeddings: - positive_null = self.null_positive_feature.view(1, 1, -1) - positive_embeddings = positive_embeddings * masks + (1 - masks) * positive_null - - objs = self.linears(torch.cat([positive_embeddings, xyxy_embedding], dim=-1)) - else: - phrases_masks = phrases_masks.unsqueeze(-1) - image_masks = image_masks.unsqueeze(-1) - - text_null = self.null_text_feature.view(1, 1, -1) - image_null = self.null_image_feature.view(1, 1, -1) - - phrases_embeddings = phrases_embeddings * phrases_masks + (1 - phrases_masks) * text_null - image_embeddings = image_embeddings * image_masks + (1 - image_masks) * image_null - - objs_text = self.linears_text(torch.cat([phrases_embeddings, xyxy_embedding], dim=-1)) - objs_image = self.linears_image(torch.cat([image_embeddings, xyxy_embedding], dim=-1)) - objs = torch.cat([objs_text, objs_image], dim=1) - - return objs - - -class UNetFlatConditionModel(ModelMixin, ConfigMixin): - r""" - A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample - shaped output. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented - for all models (such as downloading or saving). - - Parameters: - 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. - center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample. - 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")`): - 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")`): - The tuple of upsample blocks to use. - 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)`): - 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. - mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. - 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): - The dimension of the cross attention features. - 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): - 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 - [`~models.unet_2d_blocks.CrossAttnDownBlockFlat`], [`~models.unet_2d_blocks.CrossAttnUpBlockFlat`], - [`~models.unet_2d_blocks.UNetMidBlockFlatCrossAttn`]. - encoder_hid_dim (`int`, *optional*, defaults to None): - If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim` - dimension to `cross_attention_dim`. - 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 (`int`, *optional*, defaults to 8): The dimension of the attention heads. - num_attention_heads (`int`, *optional*): - The number of attention heads. If not defined, defaults to `attention_head_dim` - resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config - for ResNet blocks (see [`~models.resnet.ResnetBlockFlat`]). Choose from `default` or `scale_shift`. - class_embed_type (`str`, *optional*, defaults to `None`): - The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`, - `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`. - addition_embed_type (`str`, *optional*, defaults to `None`): - Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or - "text". "text" will use the `TextTimeEmbedding` layer. - addition_time_embed_dim: (`int`, *optional*, defaults to `None`): - Dimension for the timestep embeddings. - num_class_embeds (`int`, *optional*, defaults to `None`): - Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing - class conditioning with `class_embed_type` equal to `None`. - time_embedding_type (`str`, *optional*, defaults to `positional`): - The type of position embedding to use for timesteps. Choose from `positional` or `fourier`. - time_embedding_dim (`int`, *optional*, defaults to `None`): - An optional override for the dimension of the projected time embedding. - time_embedding_act_fn (`str`, *optional*, defaults to `None`): - Optional activation function to use only once on the time embeddings before they are passed to the rest of - the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`. - timestep_post_act (`str`, *optional*, defaults to `None`): - The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`. - time_cond_proj_dim (`int`, *optional*, defaults to `None`): - The dimension of `cond_proj` layer in the timestep embedding. - conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer. conv_out_kernel (`int`, - *optional*, default to `3`): The kernel size of `conv_out` layer. projection_class_embeddings_input_dim (`int`, - *optional*): The dimension of the `class_labels` input when - `class_embed_type="projection"`. Required when `class_embed_type="projection"`. - class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time - embeddings with the class embeddings. - mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`): - Whether to use cross attention with the mid block when using the `UNetMidBlockFlatSimpleCrossAttn`. If - `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the - `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False` - otherwise. - """ - - _supports_gradient_checkpointing = True - _no_split_modules = ["BasicTransformerBlock", "ResnetBlockFlat", "CrossAttnUpBlockFlat"] - - @register_to_config - def __init__( - self, - sample_size: Optional[int] = 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] = ( - "CrossAttnDownBlockFlat", - "CrossAttnDownBlockFlat", - "CrossAttnDownBlockFlat", - "DownBlockFlat", - ), - mid_block_type: Optional[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, - 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_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, - 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, - 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, - conv_in_kernel: int = 3, - conv_out_kernel: int = 3, - projection_class_embeddings_input_dim: Optional[int] = None, - attention_type: str = "default", - class_embeddings_concat: bool = False, - mid_block_only_cross_attention: Optional[bool] = None, - cross_attention_norm: Optional[str] = None, - addition_embed_type_num_heads=64, - ): - super().__init__() - - self.sample_size = sample_size - - if num_attention_heads is not None: - raise ValueError( - "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19." - ) - - # If `num_attention_heads` is not defined (which is the case for most models) - # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. - # The reason for this behavior is to correct for incorrectly named variables that were introduced - # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 - # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking - # which is why we correct for the naming here. - num_attention_heads = num_attention_heads or attention_head_dim - - # Check inputs - if len(down_block_types) != len(up_block_types): - raise ValueError( - f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}." - ) - - if len(block_out_channels) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}." - ) - - if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}." - ) - - if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types): - raise ValueError( - f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}." - ) - if isinstance(transformer_layers_per_block, list) and reverse_transformer_layers_per_block is None: - for layer_number_per_block in transformer_layers_per_block: - if isinstance(layer_number_per_block, list): - raise ValueError("Must provide 'reverse_transformer_layers_per_block` if using asymmetrical UNet.") - - # input - conv_in_padding = (conv_in_kernel - 1) // 2 - self.conv_in = LinearMultiDim( - in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding - ) - - # time - if time_embedding_type == "fourier": - time_embed_dim = time_embedding_dim or block_out_channels[0] * 2 - if time_embed_dim % 2 != 0: - raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.") - self.time_proj = GaussianFourierProjection( - time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos - ) - timestep_input_dim = time_embed_dim - elif time_embedding_type == "positional": - time_embed_dim = time_embedding_dim or block_out_channels[0] * 4 - - self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) - timestep_input_dim = block_out_channels[0] - else: - raise ValueError( - f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`." - ) - - self.time_embedding = TimestepEmbedding( - timestep_input_dim, - time_embed_dim, - act_fn=act_fn, - post_act_fn=timestep_post_act, - cond_proj_dim=time_cond_proj_dim, - ) - - if encoder_hid_dim_type is None and encoder_hid_dim is not None: - encoder_hid_dim_type = "text_proj" - self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type) - logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.") - - if encoder_hid_dim is None and encoder_hid_dim_type is not None: - raise ValueError( - f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}." - ) - - if encoder_hid_dim_type == "text_proj": - self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim) - elif encoder_hid_dim_type == "text_image_proj": - # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much - # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use - # case when `addition_embed_type == "text_image_proj"` (Kandinsky 2.1)` - self.encoder_hid_proj = TextImageProjection( - text_embed_dim=encoder_hid_dim, - image_embed_dim=cross_attention_dim, - cross_attention_dim=cross_attention_dim, - ) - elif encoder_hid_dim_type == "image_proj": - # Kandinsky 2.2 - self.encoder_hid_proj = ImageProjection( - image_embed_dim=encoder_hid_dim, - cross_attention_dim=cross_attention_dim, - ) - elif encoder_hid_dim_type is not None: - raise ValueError( - f"`encoder_hid_dim_type`: {encoder_hid_dim_type} must be None, 'text_proj', 'text_image_proj' or 'image_proj'." - ) - else: - self.encoder_hid_proj = None - - # class embedding - if class_embed_type is None and num_class_embeds is not None: - self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) - elif class_embed_type == "timestep": - self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn=act_fn) - elif class_embed_type == "identity": - self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) - elif class_embed_type == "projection": - if projection_class_embeddings_input_dim is None: - raise ValueError( - "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set" - ) - # The projection `class_embed_type` is the same as the timestep `class_embed_type` except - # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings - # 2. it projects from an arbitrary input dimension. - # - # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations. - # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings. - # As a result, `TimestepEmbedding` can be passed arbitrary vectors. - self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) - elif class_embed_type == "simple_projection": - if projection_class_embeddings_input_dim is None: - raise ValueError( - "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set" - ) - self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim) - else: - self.class_embedding = None - - if addition_embed_type == "text": - if encoder_hid_dim is not None: - text_time_embedding_from_dim = encoder_hid_dim - else: - text_time_embedding_from_dim = cross_attention_dim - - self.add_embedding = TextTimeEmbedding( - text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads - ) - elif addition_embed_type == "text_image": - # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much - # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use - # case when `addition_embed_type == "text_image"` (Kandinsky 2.1)` - self.add_embedding = TextImageTimeEmbedding( - text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim - ) - elif addition_embed_type == "text_time": - self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift) - self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) - elif addition_embed_type == "image": - # Kandinsky 2.2 - self.add_embedding = ImageTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim) - elif addition_embed_type == "image_hint": - # Kandinsky 2.2 ControlNet - self.add_embedding = ImageHintTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim) - elif addition_embed_type is not None: - raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.") - - if time_embedding_act_fn is None: - self.time_embed_act = None - else: - self.time_embed_act = get_activation(time_embedding_act_fn) - - self.down_blocks = nn.ModuleList([]) - self.up_blocks = nn.ModuleList([]) - - if isinstance(only_cross_attention, bool): - if mid_block_only_cross_attention is None: - mid_block_only_cross_attention = only_cross_attention - - only_cross_attention = [only_cross_attention] * len(down_block_types) - - if mid_block_only_cross_attention is None: - mid_block_only_cross_attention = False - - if isinstance(num_attention_heads, int): - num_attention_heads = (num_attention_heads,) * len(down_block_types) - - if isinstance(attention_head_dim, int): - attention_head_dim = (attention_head_dim,) * len(down_block_types) - - if isinstance(cross_attention_dim, int): - cross_attention_dim = (cross_attention_dim,) * len(down_block_types) - - if isinstance(layers_per_block, int): - layers_per_block = [layers_per_block] * len(down_block_types) - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) - - if class_embeddings_concat: - # The time embeddings are concatenated with the class embeddings. The dimension of the - # time embeddings passed to the down, middle, and up blocks is twice the dimension of the - # regular time embeddings - blocks_time_embed_dim = time_embed_dim * 2 - else: - blocks_time_embed_dim = time_embed_dim - - # down - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = get_down_block( - down_block_type, - num_layers=layers_per_block[i], - transformer_layers_per_block=transformer_layers_per_block[i], - in_channels=input_channel, - out_channels=output_channel, - temb_channels=blocks_time_embed_dim, - add_downsample=not is_final_block, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resnet_groups=norm_num_groups, - cross_attention_dim=cross_attention_dim[i], - num_attention_heads=num_attention_heads[i], - downsample_padding=downsample_padding, - dual_cross_attention=dual_cross_attention, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention[i], - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - attention_type=attention_type, - resnet_skip_time_act=resnet_skip_time_act, - resnet_out_scale_factor=resnet_out_scale_factor, - cross_attention_norm=cross_attention_norm, - attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, - dropout=dropout, - ) - self.down_blocks.append(down_block) - - # mid - if mid_block_type == "UNetMidBlockFlatCrossAttn": - self.mid_block = UNetMidBlockFlatCrossAttn( - transformer_layers_per_block=transformer_layers_per_block[-1], - in_channels=block_out_channels[-1], - temb_channels=blocks_time_embed_dim, - dropout=dropout, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - output_scale_factor=mid_block_scale_factor, - resnet_time_scale_shift=resnet_time_scale_shift, - cross_attention_dim=cross_attention_dim[-1], - num_attention_heads=num_attention_heads[-1], - resnet_groups=norm_num_groups, - dual_cross_attention=dual_cross_attention, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - attention_type=attention_type, - ) - elif mid_block_type == "UNetMidBlockFlatSimpleCrossAttn": - self.mid_block = UNetMidBlockFlatSimpleCrossAttn( - in_channels=block_out_channels[-1], - temb_channels=blocks_time_embed_dim, - dropout=dropout, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - output_scale_factor=mid_block_scale_factor, - cross_attention_dim=cross_attention_dim[-1], - attention_head_dim=attention_head_dim[-1], - resnet_groups=norm_num_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - skip_time_act=resnet_skip_time_act, - only_cross_attention=mid_block_only_cross_attention, - cross_attention_norm=cross_attention_norm, - ) - elif mid_block_type == "UNetMidBlockFlat": - self.mid_block = UNetMidBlockFlat( - in_channels=block_out_channels[-1], - temb_channels=blocks_time_embed_dim, - dropout=dropout, - num_layers=0, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - output_scale_factor=mid_block_scale_factor, - resnet_groups=norm_num_groups, - resnet_time_scale_shift=resnet_time_scale_shift, - add_attention=False, - ) - elif mid_block_type is None: - self.mid_block = None - else: - raise ValueError(f"unknown mid_block_type : {mid_block_type}") - - # count how many layers upsample the images - self.num_upsamplers = 0 - - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - reversed_num_attention_heads = list(reversed(num_attention_heads)) - reversed_layers_per_block = list(reversed(layers_per_block)) - reversed_cross_attention_dim = list(reversed(cross_attention_dim)) - reversed_transformer_layers_per_block = ( - list(reversed(transformer_layers_per_block)) - if reverse_transformer_layers_per_block is None - else reverse_transformer_layers_per_block - ) - only_cross_attention = list(reversed(only_cross_attention)) - - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - is_final_block = i == len(block_out_channels) - 1 - - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] - - # add upsample block for all BUT final layer - if not is_final_block: - add_upsample = True - self.num_upsamplers += 1 - else: - add_upsample = False - - up_block = get_up_block( - up_block_type, - num_layers=reversed_layers_per_block[i] + 1, - transformer_layers_per_block=reversed_transformer_layers_per_block[i], - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - temb_channels=blocks_time_embed_dim, - add_upsample=add_upsample, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - resolution_idx=i, - resnet_groups=norm_num_groups, - cross_attention_dim=reversed_cross_attention_dim[i], - num_attention_heads=reversed_num_attention_heads[i], - dual_cross_attention=dual_cross_attention, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention[i], - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - attention_type=attention_type, - resnet_skip_time_act=resnet_skip_time_act, - resnet_out_scale_factor=resnet_out_scale_factor, - cross_attention_norm=cross_attention_norm, - attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, - dropout=dropout, - ) - self.up_blocks.append(up_block) - prev_output_channel = output_channel - - # out - if norm_num_groups is not None: - self.conv_norm_out = nn.GroupNorm( - num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps - ) - - self.conv_act = get_activation(act_fn) - - else: - self.conv_norm_out = None - self.conv_act = None - - conv_out_padding = (conv_out_kernel - 1) // 2 - self.conv_out = LinearMultiDim( - block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding - ) - - if attention_type in ["gated", "gated-text-image"]: - positive_len = 768 - if isinstance(cross_attention_dim, int): - positive_len = cross_attention_dim - elif isinstance(cross_attention_dim, (list, tuple)): - positive_len = cross_attention_dim[0] - - feature_type = "text-only" if attention_type == "gated" else "text-image" - self.position_net = GLIGENTextBoundingboxProjection( - positive_len=positive_len, out_dim=cross_attention_dim, feature_type=feature_type - ) - - @property - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - self.set_attn_processor(processor) - - def set_attention_slice(self, slice_size): - 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. This is useful for saving some memory in exchange for a small decrease in speed. - - Args: - slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): - When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If - `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is - provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` - must be a multiple of `slice_size`. - """ - sliceable_head_dims = [] - - def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): - if hasattr(module, "set_attention_slice"): - sliceable_head_dims.append(module.sliceable_head_dim) - - for child in module.children(): - fn_recursive_retrieve_sliceable_dims(child) - - # retrieve number of attention layers - for module in self.children(): - fn_recursive_retrieve_sliceable_dims(module) - - num_sliceable_layers = len(sliceable_head_dims) - - if slice_size == "auto": - # half the attention head size is usually a good trade-off between - # speed and memory - slice_size = [dim // 2 for dim in sliceable_head_dims] - elif slice_size == "max": - # make smallest slice possible - slice_size = num_sliceable_layers * [1] - - slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size - - if len(slice_size) != len(sliceable_head_dims): - raise ValueError( - f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" - f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." - ) - - for i in range(len(slice_size)): - size = slice_size[i] - dim = sliceable_head_dims[i] - if size is not None and size > dim: - raise ValueError(f"size {size} has to be smaller or equal to {dim}.") - - # 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]): - if hasattr(module, "set_attention_slice"): - module.set_attention_slice(slice_size.pop()) - - for child in module.children(): - fn_recursive_set_attention_slice(child, slice_size) - - reversed_slice_size = list(reversed(slice_size)) - for module in self.children(): - fn_recursive_set_attention_slice(module, reversed_slice_size) - - def enable_freeu(self, s1, s2, b1, b2): - r"""Enables the FreeU mechanism from https://huggingface.co/papers/2309.11497. - - The suffixes after the scaling factors represent the stage blocks where they are being applied. - - Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that - are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL. - - Args: - s1 (`float`): - Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to - mitigate the "oversmoothing effect" in the enhanced denoising process. - s2 (`float`): - Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to - mitigate the "oversmoothing effect" in the enhanced denoising process. - b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features. - b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features. - """ - for i, upsample_block in enumerate(self.up_blocks): - setattr(upsample_block, "s1", s1) - setattr(upsample_block, "s2", s2) - setattr(upsample_block, "b1", b1) - setattr(upsample_block, "b2", b2) - - def disable_freeu(self): - """Disables the FreeU mechanism.""" - freeu_keys = {"s1", "s2", "b1", "b2"} - for i, upsample_block in enumerate(self.up_blocks): - for k in freeu_keys: - if hasattr(upsample_block, k) or getattr(upsample_block, k, None) is not None: - setattr(upsample_block, k, None) - - def fuse_qkv_projections(self): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - """ - self.original_attn_processors = None - - for _, attn_processor in self.attn_processors.items(): - if "Added" in str(attn_processor.__class__.__name__): - raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") - - self.original_attn_processors = self.attn_processors - - for module in self.modules(): - if isinstance(module, Attention): - module.fuse_projections(fuse=True) - - def unfuse_qkv_projections(self): - """Disables the fused QKV projection if enabled. - - - - This API is 🧪 experimental. - - - - """ - if self.original_attn_processors is not None: - self.set_attn_processor(self.original_attn_processors) - - def unload_lora(self): - """Unloads LoRA weights.""" - deprecate( - "unload_lora", - "0.28.0", - "Calling `unload_lora()` is deprecated and will be removed in a future version. Please install `peft` and then call `disable_adapters().", - ) - for module in self.modules(): - if hasattr(module, "set_lora_layer"): - module.set_lora_layer(None) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[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, - return_dict: bool = True, - ) -> Union[UNet2DConditionOutput, Tuple]: - r""" - The [`UNetFlatConditionModel`] forward method. - - Args: - sample (`torch.Tensor`): - The noisy input tensor with the following shape `(batch, channel, height, width)`. - timestep (`torch.Tensor` or `float` or `int`): The number of timesteps to denoise an input. - encoder_hidden_states (`torch.Tensor`): - The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. - class_labels (`torch.Tensor`, *optional*, defaults to `None`): - Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. - timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`): - Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed - through the `self.time_embedding` layer to obtain the timestep embeddings. - attention_mask (`torch.Tensor`, *optional*, defaults to `None`): - An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask - is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large - negative values to the attention scores corresponding to "discard" tokens. - cross_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). - added_cond_kwargs: (`dict`, *optional*): - A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that - are passed along to the UNet blocks. - down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*): - A tuple of tensors that if specified are added to the residuals of down unet blocks. - mid_block_additional_residual: (`torch.Tensor`, *optional*): - A tensor that if specified is added to the residual of the middle unet block. - encoder_attention_mask (`torch.Tensor`): - A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If - `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias, - which adds large negative values to the attention scores corresponding to "discard" tokens. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain - tuple. - cross_attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the [`AttnProcessor`]. - added_cond_kwargs: (`dict`, *optional*): - A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that - are passed along to the UNet blocks. - down_block_additional_residuals (`tuple` of `torch.Tensor`, *optional*): - additional residuals to be added to UNet long skip connections from down blocks to up blocks for - example from ControlNet side model(s) - mid_block_additional_residual (`torch.Tensor`, *optional*): - additional residual to be added to UNet mid block output, for example from ControlNet side model - down_intrablock_additional_residuals (`tuple` of `torch.Tensor`, *optional*): - additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s) - - Returns: - [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: - If `return_dict` is True, an [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] is returned, - otherwise a `tuple` is returned where the first element is the sample tensor. - """ - # By default samples have to be AT least a multiple of the overall upsampling factor. - # The overall upsampling factor is equal to 2 ** (# num of upsampling layers). - # However, the upsampling interpolation output size can be forced to fit any upsampling size - # on the fly if necessary. - default_overall_up_factor = 2**self.num_upsamplers - - # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` - forward_upsample_size = False - upsample_size = None - - for dim in sample.shape[-2:]: - if dim % default_overall_up_factor != 0: - # Forward upsample size to force interpolation output size. - forward_upsample_size = True - break - - # ensure attention_mask is a bias, and give it a singleton query_tokens dimension - # expects mask of shape: - # [batch, key_tokens] - # adds singleton query_tokens dimension: - # [batch, 1, key_tokens] - # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: - # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) - # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) - if attention_mask is not None: - # assume that mask is expressed as: - # (1 = keep, 0 = discard) - # convert mask into a bias that can be added to attention scores: - # (keep = +0, discard = -10000.0) - attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 - attention_mask = attention_mask.unsqueeze(1) - - # convert encoder_attention_mask to a bias the same way we do for attention_mask - if encoder_attention_mask is not None: - encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0 - encoder_attention_mask = encoder_attention_mask.unsqueeze(1) - - # 0. center input if necessary - if self.config.center_input_sample: - sample = 2 * sample - 1.0 - - # 1. time - 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 - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = sample.device.type == "mps" - is_npu = sample.device.type == "npu" - if isinstance(timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) - elif len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps.expand(sample.shape[0]) - - t_emb = self.time_proj(timesteps) - - # `Timesteps` does not contain any weights and will always return f32 tensors - # but time_embedding might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - t_emb = t_emb.to(dtype=sample.dtype) - - emb = self.time_embedding(t_emb, timestep_cond) - aug_emb = None - - if self.class_embedding is not None: - if class_labels is None: - raise ValueError("class_labels should be provided when num_class_embeds > 0") - - if self.config.class_embed_type == "timestep": - class_labels = self.time_proj(class_labels) - - # `Timesteps` does not contain any weights and will always return f32 tensors - # there might be better ways to encapsulate this. - class_labels = class_labels.to(dtype=sample.dtype) - - class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype) - - if self.config.class_embeddings_concat: - emb = torch.cat([emb, class_emb], dim=-1) - else: - emb = emb + class_emb - - if self.config.addition_embed_type == "text": - aug_emb = self.add_embedding(encoder_hidden_states) - elif self.config.addition_embed_type == "text_image": - # Kandinsky 2.1 - style - if "image_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" - ) - - image_embs = added_cond_kwargs.get("image_embeds") - text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states) - aug_emb = self.add_embedding(text_embs, image_embs) - elif self.config.addition_embed_type == "text_time": - # SDXL - style - if "text_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" - ) - text_embeds = added_cond_kwargs.get("text_embeds") - if "time_ids" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" - ) - time_ids = added_cond_kwargs.get("time_ids") - time_embeds = self.add_time_proj(time_ids.flatten()) - time_embeds = time_embeds.reshape((text_embeds.shape[0], -1)) - add_embeds = torch.concat([text_embeds, time_embeds], dim=-1) - add_embeds = add_embeds.to(emb.dtype) - aug_emb = self.add_embedding(add_embeds) - elif self.config.addition_embed_type == "image": - # Kandinsky 2.2 - style - if "image_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" - ) - image_embs = added_cond_kwargs.get("image_embeds") - aug_emb = self.add_embedding(image_embs) - elif self.config.addition_embed_type == "image_hint": - # Kandinsky 2.2 - style - if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`" - ) - image_embs = added_cond_kwargs.get("image_embeds") - hint = added_cond_kwargs.get("hint") - aug_emb, hint = self.add_embedding(image_embs, hint) - sample = torch.cat([sample, hint], dim=1) - - emb = emb + aug_emb if aug_emb is not None else emb - - if self.time_embed_act is not None: - emb = self.time_embed_act(emb) - - 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) - elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj": - # Kandinsky 2.1 - style - if "image_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" - ) - - image_embeds = added_cond_kwargs.get("image_embeds") - encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds) - elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj": - # Kandinsky 2.2 - style - if "image_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" - ) - image_embeds = added_cond_kwargs.get("image_embeds") - encoder_hidden_states = self.encoder_hid_proj(image_embeds) - elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "ip_image_proj": - if "image_embeds" not in added_cond_kwargs: - raise ValueError( - f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'ip_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" - ) - image_embeds = added_cond_kwargs.get("image_embeds") - image_embeds = self.encoder_hid_proj(image_embeds) - encoder_hidden_states = (encoder_hidden_states, image_embeds) - - # 2. pre-process - sample = self.conv_in(sample) - - # 2.5 GLIGEN position net - if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None: - cross_attention_kwargs = cross_attention_kwargs.copy() - gligen_args = cross_attention_kwargs.pop("gligen") - cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)} - - # 3. down - lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 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 - # maintain backward compatibility for legacy usage, where - # T2I-Adapter and ControlNet both use down_block_additional_residuals arg - # but can only use one or the other - if not is_adapter and mid_block_additional_residual is None and down_block_additional_residuals is not None: - deprecate( - "T2I should not use down_block_additional_residuals", - "1.3.0", - "Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \ - and will be removed in diffusers 1.3.0. `down_block_additional_residuals` should only be used \ - for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ", - standard_warn=False, - ) - down_intrablock_additional_residuals = down_block_additional_residuals - is_adapter = True - - down_block_res_samples = (sample,) - for downsample_block in self.down_blocks: - if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: - # For t2i-adapter CrossAttnDownBlockFlat - additional_residuals = {} - if is_adapter and len(down_intrablock_additional_residuals) > 0: - additional_residuals["additional_residuals"] = down_intrablock_additional_residuals.pop(0) - - sample, res_samples = downsample_block( - hidden_states=sample, - temb=emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - encoder_attention_mask=encoder_attention_mask, - **additional_residuals, - ) - else: - sample, res_samples = downsample_block(hidden_states=sample, temb=emb) - if is_adapter and len(down_intrablock_additional_residuals) > 0: - sample += down_intrablock_additional_residuals.pop(0) - - down_block_res_samples += res_samples - - if is_controlnet: - new_down_block_res_samples = () - - for down_block_res_sample, down_block_additional_residual in zip( - down_block_res_samples, down_block_additional_residuals - ): - down_block_res_sample = down_block_res_sample + down_block_additional_residual - new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,) - - down_block_res_samples = new_down_block_res_samples - - # 4. mid - if self.mid_block is not None: - if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention: - sample = self.mid_block( - sample, - emb, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - cross_attention_kwargs=cross_attention_kwargs, - encoder_attention_mask=encoder_attention_mask, - ) - else: - sample = self.mid_block(sample, emb) - - # To support T2I-Adapter-XL - if ( - is_adapter - and len(down_intrablock_additional_residuals) > 0 - and sample.shape == down_intrablock_additional_residuals[0].shape - ): - sample += down_intrablock_additional_residuals.pop(0) - - if is_controlnet: - sample = sample + mid_block_additional_residual - - # 5. up - for i, upsample_block in enumerate(self.up_blocks): - is_final_block = i == len(self.up_blocks) - 1 - - res_samples = down_block_res_samples[-len(upsample_block.resnets) :] - down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] - - # if we have not reached the final block and need to forward the - # upsample size, we do it here - if not is_final_block and forward_upsample_size: - upsample_size = down_block_res_samples[-1].shape[2:] - - if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - upsample_size=upsample_size, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - ) - else: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - upsample_size=upsample_size, - scale=lora_scale, - ) - - # 6. post-process - if self.conv_norm_out: - sample = self.conv_norm_out(sample) - 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,) - - return UNet2DConditionOutput(sample=sample) - - -class LinearMultiDim(nn.Linear): - def __init__(self, in_features, out_features=None, second_dim=4, *args, **kwargs): - in_features = [in_features, second_dim, 1] if isinstance(in_features, int) else list(in_features) - if out_features is None: - out_features = in_features - out_features = [out_features, second_dim, 1] if isinstance(out_features, int) else list(out_features) - self.in_features_multidim = in_features - self.out_features_multidim = out_features - super().__init__(np.array(in_features).prod(), np.array(out_features).prod()) - - def forward(self, input_tensor, *args, **kwargs): - shape = input_tensor.shape - n_dim = len(self.in_features_multidim) - input_tensor = input_tensor.reshape(*shape[0:-n_dim], self.in_features) - output_tensor = super().forward(input_tensor) - output_tensor = output_tensor.view(*shape[0:-n_dim], *self.out_features_multidim) - return output_tensor - - -class ResnetBlockFlat(nn.Module): - def __init__( - self, - *, - in_channels, - out_channels=None, - dropout=0.0, - temb_channels=512, - groups=32, - groups_out=None, - pre_norm=True, - eps=1e-6, - time_embedding_norm="default", - use_in_shortcut=None, - second_dim=4, - **kwargs, - ): - super().__init__() - self.pre_norm = pre_norm - self.pre_norm = True - - in_channels = [in_channels, second_dim, 1] if isinstance(in_channels, int) else list(in_channels) - self.in_channels_prod = np.array(in_channels).prod() - self.channels_multidim = in_channels - - if out_channels is not None: - out_channels = [out_channels, second_dim, 1] if isinstance(out_channels, int) else list(out_channels) - out_channels_prod = np.array(out_channels).prod() - self.out_channels_multidim = out_channels - else: - out_channels_prod = self.in_channels_prod - self.out_channels_multidim = self.channels_multidim - self.time_embedding_norm = time_embedding_norm - - if groups_out is None: - groups_out = groups - - self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=self.in_channels_prod, eps=eps, affine=True) - self.conv1 = torch.nn.Conv2d(self.in_channels_prod, out_channels_prod, kernel_size=1, padding=0) - - if temb_channels is not None: - self.time_emb_proj = torch.nn.Linear(temb_channels, out_channels_prod) - else: - self.time_emb_proj = None - - self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels_prod, eps=eps, affine=True) - self.dropout = torch.nn.Dropout(dropout) - self.conv2 = torch.nn.Conv2d(out_channels_prod, out_channels_prod, kernel_size=1, padding=0) - - self.nonlinearity = nn.SiLU() - - self.use_in_shortcut = ( - self.in_channels_prod != out_channels_prod if use_in_shortcut is None else use_in_shortcut - ) - - self.conv_shortcut = None - if self.use_in_shortcut: - self.conv_shortcut = torch.nn.Conv2d( - self.in_channels_prod, out_channels_prod, kernel_size=1, stride=1, padding=0 - ) - - def forward(self, input_tensor, temb): - shape = input_tensor.shape - n_dim = len(self.channels_multidim) - input_tensor = input_tensor.reshape(*shape[0:-n_dim], self.in_channels_prod, 1, 1) - input_tensor = input_tensor.view(-1, self.in_channels_prod, 1, 1) - - hidden_states = input_tensor - - hidden_states = self.norm1(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - hidden_states = self.conv1(hidden_states) - - if temb is not None: - temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None] - hidden_states = hidden_states + temb - - hidden_states = self.norm2(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - - hidden_states = self.dropout(hidden_states) - hidden_states = self.conv2(hidden_states) - - if self.conv_shortcut is not None: - input_tensor = self.conv_shortcut(input_tensor) - - output_tensor = input_tensor + hidden_states - - output_tensor = output_tensor.view(*shape[0:-n_dim], -1) - output_tensor = output_tensor.view(*shape[0:-n_dim], *self.out_channels_multidim) - - return output_tensor - - -class DownBlockFlat(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor: float = 1.0, - add_downsample: bool = True, - downsample_padding: int = 1, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlockFlat( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - self.resnets = nn.ModuleList(resnets) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - LinearMultiDim( - out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - def forward( - self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None - ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: - output_states = () - - for resnet in self.resnets: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = resnet(hidden_states, temb) - - output_states = output_states + (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - output_states = output_states + (hidden_states,) - - return hidden_states, output_states - - -class CrossAttnDownBlockFlat(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[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_pre_norm: bool = True, - num_attention_heads: int = 1, - cross_attention_dim: int = 1280, - output_scale_factor: float = 1.0, - downsample_padding: int = 1, - add_downsample: bool = True, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - only_cross_attention: bool = False, - upcast_attention: bool = False, - attention_type: str = "default", - ): - super().__init__() - resnets = [] - attentions = [] - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * num_layers - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlockFlat( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - if not dual_cross_attention: - attentions.append( - Transformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - attention_type=attention_type, - ) - ) - else: - attentions.append( - DualTransformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - ) - ) - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - if add_downsample: - self.downsamplers = nn.ModuleList( - [ - LinearMultiDim( - out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" - ) - ] - ) - else: - self.downsamplers = None - - self.gradient_checkpointing = False - - 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, ...]]: - output_states = () - - blocks = list(zip(self.resnets, self.attentions)) - - for i, (resnet, attn) in enumerate(blocks): - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - else: - hidden_states = resnet(hidden_states, temb) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - - # apply additional residuals to the output of the last pair of resnet and attention blocks - if i == len(blocks) - 1 and additional_residuals is not None: - hidden_states = hidden_states + additional_residuals - - output_states = output_states + (hidden_states,) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - output_states = output_states + (hidden_states,) - - return hidden_states, output_states - - -# Copied from diffusers.models.unets.unet_2d_blocks.UpBlock2D with UpBlock2D->UpBlockFlat, ResnetBlock2D->ResnetBlockFlat, Upsample2D->LinearMultiDim -class UpBlockFlat(nn.Module): - def __init__( - self, - in_channels: int, - prev_output_channel: int, - out_channels: int, - temb_channels: int, - resolution_idx: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor: float = 1.0, - add_upsample: bool = True, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlockFlat( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - self.resnets = nn.ModuleList(resnets) - - if add_upsample: - self.upsamplers = nn.ModuleList([LinearMultiDim(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - def forward( - self, - hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = 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) - - is_freeu_enabled = ( - getattr(self, "s1", None) - and getattr(self, "s2", None) - and getattr(self, "b1", None) - and getattr(self, "b2", None) - ) - - for resnet in self.resnets: - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - - # FreeU: Only operate on the first two stages - if is_freeu_enabled: - hidden_states, res_hidden_states = apply_freeu( - self.resolution_idx, - hidden_states, - res_hidden_states, - s1=self.s1, - s2=self.s2, - b1=self.b1, - b2=self.b2, - ) - - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = resnet(hidden_states, temb) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) - - return hidden_states - - -# Copied from diffusers.models.unets.unet_2d_blocks.CrossAttnUpBlock2D with CrossAttnUpBlock2D->CrossAttnUpBlockFlat, ResnetBlock2D->ResnetBlockFlat, Upsample2D->LinearMultiDim -class CrossAttnUpBlockFlat(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - prev_output_channel: int, - temb_channels: int, - resolution_idx: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[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_pre_norm: bool = True, - num_attention_heads: int = 1, - cross_attention_dim: int = 1280, - output_scale_factor: float = 1.0, - add_upsample: bool = True, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - only_cross_attention: bool = False, - upcast_attention: bool = False, - attention_type: str = "default", - ): - super().__init__() - resnets = [] - attentions = [] - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * num_layers - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlockFlat( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - if not dual_cross_attention: - attentions.append( - Transformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - attention_type=attention_type, - ) - ) - else: - attentions.append( - DualTransformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - ) - ) - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - if add_upsample: - self.upsamplers = nn.ModuleList([LinearMultiDim(out_channels, use_conv=True, out_channels=out_channels)]) - else: - self.upsamplers = None - - self.gradient_checkpointing = False - self.resolution_idx = resolution_idx - - 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, - ) -> 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.") - - is_freeu_enabled = ( - getattr(self, "s1", None) - and getattr(self, "s2", None) - and getattr(self, "b1", None) - and getattr(self, "b2", None) - ) - - for resnet, attn in zip(self.resnets, self.attentions): - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - - # FreeU: Only operate on the first two stages - if is_freeu_enabled: - hidden_states, res_hidden_states = apply_freeu( - self.resolution_idx, - hidden_states, - res_hidden_states, - s1=self.s1, - s2=self.s2, - b1=self.b1, - b2=self.b2, - ) - - hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - else: - hidden_states = resnet(hidden_states, temb) - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) - - return hidden_states - - -# Copied from diffusers.models.unets.unet_2d_blocks.UNetMidBlock2D with UNetMidBlock2D->UNetMidBlockFlat, ResnetBlock2D->ResnetBlockFlat -class UNetMidBlockFlat(nn.Module): - """ - A 2D UNet mid-block [`UNetMidBlockFlat`] with multiple residual blocks and optional attention blocks. - - Args: - in_channels (`int`): The number of input channels. - temb_channels (`int`): The number of temporal embedding channels. - dropout (`float`, *optional*, defaults to 0.0): The dropout rate. - num_layers (`int`, *optional*, defaults to 1): The number of residual blocks. - resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks. - resnet_time_scale_shift (`str`, *optional*, defaults to `default`): - The type of normalization to apply to the time embeddings. This can help to improve the performance of the - model on tasks with long-range temporal dependencies. - 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. - 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. - attention_head_dim (`int`, *optional*, defaults to 1): - Dimension of a single attention head. The number of attention heads is determined based on this value and - the number of input channels. - output_scale_factor (`float`, *optional*, defaults to 1.0): The output scale factor. - - Returns: - `torch.Tensor`: The output of the last residual block, which is a tensor of shape `(batch_size, in_channels, - height, width)`. - - """ - - def __init__( - self, - in_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", # default, spatial - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - attn_groups: Optional[int] = None, - resnet_pre_norm: bool = True, - add_attention: bool = True, - attention_head_dim: int = 1, - output_scale_factor: float = 1.0, - ): - super().__init__() - resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) - self.add_attention = add_attention - - if attn_groups is None: - attn_groups = resnet_groups if resnet_time_scale_shift == "default" else None - - # there is always at least one resnet - if resnet_time_scale_shift == "spatial": - resnets = [ - ResnetBlockCondNorm2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm="spatial", - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - ) - ] - else: - resnets = [ - ResnetBlockFlat( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ] - attentions = [] - - if attention_head_dim is None: - logger.warning( - f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}." - ) - attention_head_dim = in_channels - - for _ in range(num_layers): - if self.add_attention: - attentions.append( - Attention( - in_channels, - heads=in_channels // attention_head_dim, - dim_head=attention_head_dim, - rescale_output_factor=output_scale_factor, - eps=resnet_eps, - norm_num_groups=attn_groups, - spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None, - residual_connection=True, - bias=True, - upcast_softmax=True, - _from_deprecated_attn_block=True, - ) - ) - else: - attentions.append(None) - - if resnet_time_scale_shift == "spatial": - resnets.append( - ResnetBlockCondNorm2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm="spatial", - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - ) - ) - else: - resnets.append( - ResnetBlockFlat( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - self.gradient_checkpointing = False - - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = 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: - if attn is not None: - hidden_states = attn(hidden_states, temb=temb) - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - if attn is not None: - hidden_states = attn(hidden_states, temb=temb) - hidden_states = resnet(hidden_states, temb) - - return hidden_states - - -# Copied from diffusers.models.unets.unet_2d_blocks.UNetMidBlock2DCrossAttn with UNetMidBlock2DCrossAttn->UNetMidBlockFlatCrossAttn, ResnetBlock2D->ResnetBlockFlat -class UNetMidBlockFlatCrossAttn(nn.Module): - def __init__( - self, - in_channels: int, - temb_channels: int, - out_channels: Optional[int] = None, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[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_pre_norm: bool = True, - num_attention_heads: int = 1, - output_scale_factor: float = 1.0, - cross_attention_dim: int = 1280, - dual_cross_attention: bool = False, - use_linear_projection: bool = False, - upcast_attention: bool = False, - attention_type: str = "default", - ): - super().__init__() - - out_channels = out_channels or in_channels - self.in_channels = in_channels - self.out_channels = out_channels - - self.has_cross_attention = True - self.num_attention_heads = num_attention_heads - resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) - - # support for variable transformer layers per block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block] * num_layers - - resnet_groups_out = resnet_groups_out or resnet_groups - - # there is always at least one resnet - resnets = [ - ResnetBlockFlat( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - groups_out=resnet_groups_out, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ] - attentions = [] - - for i in range(num_layers): - if not dual_cross_attention: - attentions.append( - Transformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups_out, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - attention_type=attention_type, - ) - ) - else: - attentions.append( - DualTransformer2DModel( - num_attention_heads, - out_channels // num_attention_heads, - in_channels=out_channels, - num_layers=1, - cross_attention_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - ) - ) - resnets.append( - ResnetBlockFlat( - in_channels=out_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups_out, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - self.gradient_checkpointing = False - - 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, - ) -> 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.") - - 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: - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) - else: - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - cross_attention_kwargs=cross_attention_kwargs, - attention_mask=attention_mask, - encoder_attention_mask=encoder_attention_mask, - return_dict=False, - )[0] - hidden_states = resnet(hidden_states, temb) - - return hidden_states - - -# Copied from diffusers.models.unets.unet_2d_blocks.UNetMidBlock2DSimpleCrossAttn with UNetMidBlock2DSimpleCrossAttn->UNetMidBlockFlatSimpleCrossAttn, ResnetBlock2D->ResnetBlockFlat -class UNetMidBlockFlatSimpleCrossAttn(nn.Module): - def __init__( - self, - in_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attention_head_dim: int = 1, - output_scale_factor: float = 1.0, - cross_attention_dim: int = 1280, - skip_time_act: bool = False, - only_cross_attention: bool = False, - cross_attention_norm: Optional[str] = None, - ): - super().__init__() - - self.has_cross_attention = True - - self.attention_head_dim = attention_head_dim - resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) - - self.num_heads = in_channels // self.attention_head_dim - - # there is always at least one resnet - resnets = [ - ResnetBlockFlat( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - skip_time_act=skip_time_act, - ) - ] - attentions = [] - - for _ in range(num_layers): - processor = ( - AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor() - ) - - attentions.append( - Attention( - query_dim=in_channels, - cross_attention_dim=in_channels, - heads=self.num_heads, - dim_head=self.attention_head_dim, - added_kv_proj_dim=cross_attention_dim, - norm_num_groups=resnet_groups, - bias=True, - upcast_softmax=True, - only_cross_attention=only_cross_attention, - cross_attention_norm=cross_attention_norm, - processor=processor, - ) - ) - resnets.append( - ResnetBlockFlat( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - skip_time_act=skip_time_act, - ) - ) - - self.attentions = nn.ModuleList(attentions) - self.resnets = nn.ModuleList(resnets) - - 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, - ) -> 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.") - - if attention_mask is None: - # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask. - mask = None if encoder_hidden_states is None else encoder_attention_mask - else: - # when attention_mask is defined: we don't even check for encoder_attention_mask. - # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks. - # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask. - # then we can simplify this whole if/else block to: - # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask - mask = attention_mask - - hidden_states = self.resnets[0](hidden_states, temb) - for attn, resnet in zip(self.attentions, self.resnets[1:]): - # attn - hidden_states = attn( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=mask, - **cross_attention_kwargs, - ) - - # resnet - hidden_states = resnet(hidden_states, temb) - - return hidden_states diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py deleted file mode 100644 index 61582853b055a988704e821aab5c072d6c4d1207..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py +++ /dev/null @@ -1,421 +0,0 @@ -import inspect -from typing import Callable, List, Optional, Union - -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModel - -from ....models import AutoencoderKL, UNet2DConditionModel -from ....schedulers import KarrasDiffusionSchedulers -from ....utils import logging -from ...pipeline_utils import DiffusionPipeline -from .pipeline_versatile_diffusion_dual_guided import VersatileDiffusionDualGuidedPipeline -from .pipeline_versatile_diffusion_image_variation import VersatileDiffusionImageVariationPipeline -from .pipeline_versatile_diffusion_text_to_image import VersatileDiffusionTextToImagePipeline - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class VersatileDiffusionPipeline(DiffusionPipeline): - r""" - Pipeline for text-to-image generation using Stable Diffusion. - - 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: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - tokenizer: CLIPTokenizer - image_feature_extractor: CLIPImageProcessor - text_encoder: CLIPTextModel - image_encoder: CLIPVisionModel - image_unet: UNet2DConditionModel - text_unet: UNet2DConditionModel - vae: AutoencoderKL - scheduler: KarrasDiffusionSchedulers - - def __init__( - self, - tokenizer: CLIPTokenizer, - image_feature_extractor: CLIPImageProcessor, - text_encoder: CLIPTextModel, - image_encoder: CLIPVisionModel, - image_unet: UNet2DConditionModel, - text_unet: UNet2DConditionModel, - vae: AutoencoderKL, - scheduler: KarrasDiffusionSchedulers, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, - image_feature_extractor=image_feature_extractor, - text_encoder=text_encoder, - image_encoder=image_encoder, - image_unet=image_unet, - text_unet=text_unet, - vae=vae, - scheduler=scheduler, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - - @torch.no_grad() - def image_variation( - self, - image: Union[torch.Tensor, PIL.Image.Image], - 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, - output_type: Optional[str] = "pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], 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`): - 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. - width (`int`, *optional*, defaults to `self.image_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): - 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*): - 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): - The number of images to generate per prompt. - 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`, *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`. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - - Examples: - - ```py - >>> from diffusers import VersatileDiffusionPipeline - >>> import torch - >>> import requests - >>> from io import BytesIO - >>> from PIL import Image - - >>> # let's download an initial image - >>> url = "https://huggingface.co/datasets/diffusers/images/resolve/main/benz.jpg" - - >>> response = requests.get(url) - >>> image = Image.open(BytesIO(response.content)).convert("RGB") - - >>> pipe = VersatileDiffusionPipeline.from_pretrained( - ... "shi-labs/versatile-diffusion", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> generator = torch.Generator(device="cuda").manual_seed(0) - >>> image = pipe.image_variation(image, generator=generator).images[0] - >>> image.save("./car_variation.png") - ``` - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - expected_components = inspect.signature(VersatileDiffusionImageVariationPipeline.__init__).parameters.keys() - components = {name: component for name, component in self.components.items() if name in expected_components} - return VersatileDiffusionImageVariationPipeline(**components)( - image=image, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - eta=eta, - generator=generator, - latents=latents, - output_type=output_type, - return_dict=return_dict, - callback=callback, - callback_steps=callback_steps, - ) - - @torch.no_grad() - def text_to_image( - self, - prompt: Union[str, List[str]], - 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, - output_type: Optional[str] = "pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - width (`int`, *optional*, defaults to `self.image_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): - 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*): - 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): - The number of images to generate per prompt. - 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`, *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`. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - - Examples: - - ```py - >>> from diffusers import VersatileDiffusionPipeline - >>> import torch - - >>> pipe = VersatileDiffusionPipeline.from_pretrained( - ... "shi-labs/versatile-diffusion", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> generator = torch.Generator(device="cuda").manual_seed(0) - >>> image = pipe.text_to_image("an astronaut riding on a horse on mars", generator=generator).images[0] - >>> image.save("./astronaut.png") - ``` - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - expected_components = inspect.signature(VersatileDiffusionTextToImagePipeline.__init__).parameters.keys() - components = {name: component for name, component in self.components.items() if name in expected_components} - temp_pipeline = VersatileDiffusionTextToImagePipeline(**components) - output = temp_pipeline( - prompt=prompt, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - eta=eta, - generator=generator, - latents=latents, - output_type=output_type, - return_dict=return_dict, - callback=callback, - callback_steps=callback_steps, - ) - # swap the attention blocks back to the original state - temp_pipeline._swap_unet_attention_blocks() - - return output - - @torch.no_grad() - def dual_guided( - self, - prompt: Union[PIL.Image.Image, List[PIL.Image.Image]], - image: Union[str, List[str]], - text_to_image_strength: float = 0.5, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - guidance_scale: float = 7.5, - 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, - output_type: Optional[str] = "pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - width (`int`, *optional*, defaults to `self.image_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): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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`. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - - Examples: - - ```py - >>> from diffusers import VersatileDiffusionPipeline - >>> import torch - >>> import requests - >>> from io import BytesIO - >>> from PIL import Image - - >>> # let's download an initial image - >>> url = "https://huggingface.co/datasets/diffusers/images/resolve/main/benz.jpg" - - >>> response = requests.get(url) - >>> image = Image.open(BytesIO(response.content)).convert("RGB") - >>> text = "a red car in the sun" - - >>> pipe = VersatileDiffusionPipeline.from_pretrained( - ... "shi-labs/versatile-diffusion", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> generator = torch.Generator(device="cuda").manual_seed(0) - >>> text_to_image_strength = 0.75 - - >>> image = pipe.dual_guided( - ... prompt=text, image=image, text_to_image_strength=text_to_image_strength, generator=generator - ... ).images[0] - >>> image.save("./car_variation.png") - ``` - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images. - """ - - expected_components = inspect.signature(VersatileDiffusionDualGuidedPipeline.__init__).parameters.keys() - components = {name: component for name, component in self.components.items() if name in expected_components} - temp_pipeline = VersatileDiffusionDualGuidedPipeline(**components) - output = temp_pipeline( - prompt=prompt, - image=image, - text_to_image_strength=text_to_image_strength, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - num_images_per_prompt=num_images_per_prompt, - eta=eta, - generator=generator, - latents=latents, - output_type=output_type, - return_dict=return_dict, - callback=callback, - callback_steps=callback_steps, - ) - temp_pipeline._revert_dual_attention() - - return output diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py deleted file mode 100644 index 0252f4f6af7f6daedc6b3f6d3912f602c7a27de0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py +++ /dev/null @@ -1,561 +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 inspect -from typing import Callable, List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch -import torch.utils.checkpoint -from transformers import ( - CLIPImageProcessor, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ....image_processor import VaeImageProcessor -from ....models import AutoencoderKL, DualTransformer2DModel, Transformer2DModel, UNet2DConditionModel -from ....schedulers import KarrasDiffusionSchedulers -from ....utils import deprecate, logging -from ....utils.torch_utils import randn_tensor -from ...pipeline_utils import DiffusionPipeline, ImagePipelineOutput -from .modeling_text_unet import UNetFlatConditionModel - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class VersatileDiffusionDualGuidedPipeline(DiffusionPipeline): - r""" - Pipeline for image-text dual-guided generation using Versatile Diffusion. - - 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.). - - Parameters: - vqvae ([`VQModel`]): - Vector-quantized (VQ) model to encode and decode images to and from latent representations. - bert ([`LDMBertModel`]): - Text-encoder model based on [`~transformers.BERT`]. - tokenizer ([`~transformers.BertTokenizer`]): - A `BertTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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`]. - """ - - model_cpu_offload_seq = "bert->unet->vqvae" - - tokenizer: CLIPTokenizer - image_feature_extractor: CLIPImageProcessor - text_encoder: CLIPTextModelWithProjection - image_encoder: CLIPVisionModelWithProjection - image_unet: UNet2DConditionModel - text_unet: UNetFlatConditionModel - vae: AutoencoderKL - scheduler: KarrasDiffusionSchedulers - - _optional_components = ["text_unet"] - - def __init__( - self, - tokenizer: CLIPTokenizer, - image_feature_extractor: CLIPImageProcessor, - text_encoder: CLIPTextModelWithProjection, - image_encoder: CLIPVisionModelWithProjection, - image_unet: UNet2DConditionModel, - text_unet: UNetFlatConditionModel, - vae: AutoencoderKL, - scheduler: KarrasDiffusionSchedulers, - ): - super().__init__() - self.register_modules( - tokenizer=tokenizer, - image_feature_extractor=image_feature_extractor, - text_encoder=text_encoder, - image_encoder=image_encoder, - image_unet=image_unet, - text_unet=text_unet, - vae=vae, - 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) - - if self.text_unet is not None and ( - "dual_cross_attention" not in self.image_unet.config or not self.image_unet.config.dual_cross_attention - ): - # if loading from a universal checkpoint rather than a saved dual-guided pipeline - self._convert_to_dual_attention() - - def remove_unused_weights(self): - self.register_modules(text_unet=None) - - def _convert_to_dual_attention(self): - """ - Replace image_unet's `Transformer2DModel` blocks with `DualTransformer2DModel` that contains transformer blocks - from both `image_unet` and `text_unet` - """ - for name, module in self.image_unet.named_modules(): - if isinstance(module, Transformer2DModel): - parent_name, index = name.rsplit(".", 1) - index = int(index) - - image_transformer = self.image_unet.get_submodule(parent_name)[index] - text_transformer = self.text_unet.get_submodule(parent_name)[index] - - config = image_transformer.config - dual_transformer = DualTransformer2DModel( - num_attention_heads=config.num_attention_heads, - attention_head_dim=config.attention_head_dim, - in_channels=config.in_channels, - num_layers=config.num_layers, - dropout=config.dropout, - norm_num_groups=config.norm_num_groups, - cross_attention_dim=config.cross_attention_dim, - attention_bias=config.attention_bias, - sample_size=config.sample_size, - num_vector_embeds=config.num_vector_embeds, - activation_fn=config.activation_fn, - num_embeds_ada_norm=config.num_embeds_ada_norm, - ) - dual_transformer.transformers[0] = image_transformer - dual_transformer.transformers[1] = text_transformer - - self.image_unet.get_submodule(parent_name)[index] = dual_transformer - self.image_unet.register_to_config(dual_cross_attention=True) - - def _revert_dual_attention(self): - """ - Revert the image_unet `DualTransformer2DModel` blocks back to `Transformer2DModel` with image_unet weights Call - this function if you reuse `image_unet` in another pipeline, e.g. `VersatileDiffusionPipeline` - """ - for name, module in self.image_unet.named_modules(): - if isinstance(module, DualTransformer2DModel): - parent_name, index = name.rsplit(".", 1) - index = int(index) - self.image_unet.get_submodule(parent_name)[index] = module.transformers[0] - - self.image_unet.register_to_config(dual_cross_attention=False) - - def _encode_text_prompt(self, prompt, device, num_images_per_prompt, do_classifier_free_guidance): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`): - 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 - """ - - def normalize_embeddings(encoder_output): - embeds = self.text_encoder.text_projection(encoder_output.last_hidden_state) - embeds_pooled = encoder_output.text_embeds - embeds = embeds / torch.norm(embeds_pooled.unsqueeze(1), dim=-1, keepdim=True) - return embeds - - batch_size = len(prompt) - - 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="max_length", return_tensors="pt").input_ids - - if 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 - - prompt_embeds = self.text_encoder( - text_input_ids.to(device), - attention_mask=attention_mask, - ) - prompt_embeds = normalize_embeddings(prompt_embeds) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - bs_embed, seq_len, _ = prompt_embeds.shape - 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: - uncond_tokens = [""] * batch_size - max_length = text_input_ids.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 = normalize_embeddings(negative_prompt_embeds) - - # 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.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - return prompt_embeds - - def _encode_image_prompt(self, prompt, device, num_images_per_prompt, do_classifier_free_guidance): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`): - 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 - """ - - def normalize_embeddings(encoder_output): - embeds = self.image_encoder.vision_model.post_layernorm(encoder_output.last_hidden_state) - embeds = self.image_encoder.visual_projection(embeds) - embeds_pooled = embeds[:, 0:1] - embeds = embeds / torch.norm(embeds_pooled, dim=-1, keepdim=True) - return embeds - - batch_size = len(prompt) if isinstance(prompt, list) else 1 - - # get prompt text embeddings - image_input = self.image_feature_extractor(images=prompt, return_tensors="pt") - pixel_values = image_input.pixel_values.to(device).to(self.image_encoder.dtype) - image_embeddings = self.image_encoder(pixel_values) - image_embeddings = normalize_embeddings(image_embeddings) - - # duplicate image embeddings for each generation per prompt, using mps friendly method - bs_embed, seq_len, _ = image_embeddings.shape - image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1) - image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) - - # get unconditional embeddings for classifier free guidance - if do_classifier_free_guidance: - uncond_images = [np.zeros((512, 512, 3)) + 0.5] * batch_size - uncond_images = self.image_feature_extractor(images=uncond_images, return_tensors="pt") - pixel_values = uncond_images.pixel_values.to(device).to(self.image_encoder.dtype) - negative_prompt_embeds = self.image_encoder(pixel_values) - negative_prompt_embeds = normalize_embeddings(negative_prompt_embeds) - - # 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.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - - # For classifier free guidance, we need to do two forward passes. - # Here we concatenate the unconditional and conditional embeddings into a single batch - # to avoid doing two forward passes - image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings]) - - return image_embeddings - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs(self, prompt, image, height, width, callback_steps): - if not isinstance(prompt, str) and not isinstance(prompt, PIL.Image.Image) and not isinstance(prompt, list): - raise ValueError(f"`prompt` has to be of type `str` `PIL.Image` or `list` but is {type(prompt)}") - if not isinstance(image, str) and not isinstance(image, PIL.Image.Image) and not isinstance(image, list): - raise ValueError(f"`image` has to be of type `str` `PIL.Image` or `list` but is {type(image)}") - - 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 None) or ( - 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)}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - 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 - - for i, type in enumerate(condition_types): - if type == "text": - module.condition_lengths[i] = self.text_encoder.config.max_position_embeddings - module.transformer_index_for_condition[i] = 1 # use the second (text) transformer - else: - module.condition_lengths[i] = 257 - module.transformer_index_for_condition[i] = 0 # use the first (image) transformer - - @torch.no_grad() - def __call__( - self, - prompt: Union[PIL.Image.Image, List[PIL.Image.Image]], - image: Union[str, List[str]], - text_to_image_strength: float = 0.5, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - guidance_scale: float = 7.5, - 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, - output_type: Optional[str] = "pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - width (`int`, *optional*, defaults to `self.image_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): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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`. - 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 [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - - Examples: - - ```py - >>> from diffusers import VersatileDiffusionDualGuidedPipeline - >>> import torch - >>> import requests - >>> from io import BytesIO - >>> from PIL import Image - - >>> # let's download an initial image - >>> url = "https://huggingface.co/datasets/diffusers/images/resolve/main/benz.jpg" - - >>> response = requests.get(url) - >>> image = Image.open(BytesIO(response.content)).convert("RGB") - >>> text = "a red car in the sun" - - >>> pipe = VersatileDiffusionDualGuidedPipeline.from_pretrained( - ... "shi-labs/versatile-diffusion", torch_dtype=torch.float16 - ... ) - >>> pipe.remove_unused_weights() - >>> pipe = pipe.to("cuda") - - >>> generator = torch.Generator(device="cuda").manual_seed(0) - >>> text_to_image_strength = 0.75 - - >>> image = pipe( - ... prompt=text, image=image, text_to_image_strength=text_to_image_strength, generator=generator - ... ).images[0] - >>> image.save("./car_variation.png") - ``` - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images. - """ - # 0. Default height and width to unet - height = height or self.image_unet.config.sample_size * self.vae_scale_factor - width = width or self.image_unet.config.sample_size * self.vae_scale_factor - - # 1. Check inputs. Raise error if not correct - self.check_inputs(prompt, image, height, width, callback_steps) - - # 2. Define call parameters - prompt = [prompt] if not isinstance(prompt, list) else prompt - image = [image] if not isinstance(image, list) else image - batch_size = len(prompt) - 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 = guidance_scale > 1.0 - - # 3. Encode input prompts - prompt_embeds = self._encode_text_prompt(prompt, device, num_images_per_prompt, do_classifier_free_guidance) - image_embeddings = self._encode_image_prompt(image, device, num_images_per_prompt, do_classifier_free_guidance) - dual_prompt_embeddings = torch.cat([prompt_embeds, image_embeddings], dim=1) - prompt_types = ("text", "image") - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.image_unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - dual_prompt_embeddings.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Combine the attention blocks of the image and text UNets - self.set_transformer_params(text_to_image_strength, prompt_types) - - # 8. Denoising loop - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.image_unet(latent_model_input, t, encoder_hidden_states=dual_prompt_embeddings).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # call the callback, if provided - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - else: - image = latents - - image = self.image_processor.postprocess(image, output_type=output_type) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py deleted file mode 100644 index 2beb0be57b73fcb9b910178f0ffe860d3df5df12..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py +++ /dev/null @@ -1,402 +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 inspect -from typing import Callable, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -import torch.utils.checkpoint -from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection - -from ....image_processor import VaeImageProcessor -from ....models import AutoencoderKL, UNet2DConditionModel -from ....schedulers import KarrasDiffusionSchedulers -from ....utils import deprecate, logging -from ....utils.torch_utils import randn_tensor -from ...pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class VersatileDiffusionImageVariationPipeline(DiffusionPipeline): - r""" - Pipeline for image variation using Versatile Diffusion. - - 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.). - - Parameters: - vqvae ([`VQModel`]): - Vector-quantized (VQ) model to encode and decode images to and from latent representations. - bert ([`LDMBertModel`]): - Text-encoder model based on [`~transformers.BERT`]. - tokenizer ([`~transformers.BertTokenizer`]): - A `BertTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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`]. - """ - - model_cpu_offload_seq = "bert->unet->vqvae" - - image_feature_extractor: CLIPImageProcessor - image_encoder: CLIPVisionModelWithProjection - image_unet: UNet2DConditionModel - vae: AutoencoderKL - scheduler: KarrasDiffusionSchedulers - - def __init__( - self, - image_feature_extractor: CLIPImageProcessor, - image_encoder: CLIPVisionModelWithProjection, - image_unet: UNet2DConditionModel, - vae: AutoencoderKL, - scheduler: KarrasDiffusionSchedulers, - ): - super().__init__() - self.register_modules( - image_feature_extractor=image_feature_extractor, - image_encoder=image_encoder, - image_unet=image_unet, - vae=vae, - 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) - - def _encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`): - 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]`): - 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`). - """ - - def normalize_embeddings(encoder_output): - embeds = self.image_encoder.vision_model.post_layernorm(encoder_output.last_hidden_state) - embeds = self.image_encoder.visual_projection(embeds) - embeds_pooled = embeds[:, 0:1] - embeds = embeds / torch.norm(embeds_pooled, dim=-1, keepdim=True) - return embeds - - if isinstance(prompt, torch.Tensor) and len(prompt.shape) == 4: - prompt = list(prompt) - - batch_size = len(prompt) if isinstance(prompt, list) else 1 - - # get prompt text embeddings - image_input = self.image_feature_extractor(images=prompt, return_tensors="pt") - pixel_values = image_input.pixel_values.to(device).to(self.image_encoder.dtype) - image_embeddings = self.image_encoder(pixel_values) - image_embeddings = normalize_embeddings(image_embeddings) - - # duplicate image embeddings for each generation per prompt, using mps friendly method - bs_embed, seq_len, _ = image_embeddings.shape - image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1) - image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) - - # get unconditional embeddings for classifier free guidance - if do_classifier_free_guidance: - 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): - raise TypeError( - f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" - f" {type(prompt)}." - ) - elif isinstance(negative_prompt, PIL.Image.Image): - uncond_images = [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_images = negative_prompt - - uncond_images = self.image_feature_extractor(images=uncond_images, return_tensors="pt") - pixel_values = uncond_images.pixel_values.to(device).to(self.image_encoder.dtype) - negative_prompt_embeds = self.image_encoder(pixel_values) - negative_prompt_embeds = normalize_embeddings(negative_prompt_embeds) - - # 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.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - - # For classifier free guidance, we need to do two forward passes. - # Here we concatenate the unconditional and conditional embeddings into a single batch - # to avoid doing two forward passes - image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings]) - - return image_embeddings - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_image_variation.StableDiffusionImageVariationPipeline.check_inputs - def check_inputs(self, image, height, width, callback_steps): - if ( - not isinstance(image, torch.Tensor) - and not isinstance(image, PIL.Image.Image) - 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" - f" {type(image)}" - ) - - 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 None) or ( - 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)}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @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, - 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, - output_type: Optional[str] = "pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - width (`int`, *optional*, defaults to `self.image_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): - 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*): - 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): - The number of images to generate per prompt. - 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`, *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`. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - - Examples: - - ```py - >>> from diffusers import VersatileDiffusionImageVariationPipeline - >>> import torch - >>> import requests - >>> from io import BytesIO - >>> from PIL import Image - - >>> # let's download an initial image - >>> url = "https://huggingface.co/datasets/diffusers/images/resolve/main/benz.jpg" - - >>> response = requests.get(url) - >>> image = Image.open(BytesIO(response.content)).convert("RGB") - - >>> pipe = VersatileDiffusionImageVariationPipeline.from_pretrained( - ... "shi-labs/versatile-diffusion", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> generator = torch.Generator(device="cuda").manual_seed(0) - >>> image = pipe(image, generator=generator).images[0] - >>> image.save("./car_variation.png") - ``` - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned, - otherwise a `tuple` is returned where the first element is a list with the generated images. - """ - # 0. Default height and width to unet - height = height or self.image_unet.config.sample_size * self.vae_scale_factor - width = width or self.image_unet.config.sample_size * self.vae_scale_factor - - # 1. Check inputs. Raise error if not correct - self.check_inputs(image, height, width, callback_steps) - - # 2. Define call parameters - batch_size = 1 if isinstance(image, PIL.Image.Image) else len(image) - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - image_embeddings = self._encode_prompt( - image, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt - ) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.image_unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - image_embeddings.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.image_unet(latent_model_input, t, encoder_hidden_states=image_embeddings).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # call the callback, if provided - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - else: - image = latents - - image = self.image_processor.postprocess(image, output_type=output_type) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py deleted file mode 100644 index adfd899e766a0d8644609d091bc91b4dc4f4b0d9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py +++ /dev/null @@ -1,480 +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 inspect -from typing import Callable, List, Optional, Union - -import torch -import torch.utils.checkpoint -from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer - -from ....image_processor import VaeImageProcessor -from ....models import AutoencoderKL, Transformer2DModel, UNet2DConditionModel -from ....schedulers import KarrasDiffusionSchedulers -from ....utils import deprecate, logging -from ....utils.torch_utils import randn_tensor -from ...pipeline_utils import DiffusionPipeline, ImagePipelineOutput -from .modeling_text_unet import UNetFlatConditionModel - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class VersatileDiffusionTextToImagePipeline(DiffusionPipeline): - r""" - Pipeline for text-to-image generation using Versatile Diffusion. - - 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.). - - Parameters: - vqvae ([`VQModel`]): - Vector-quantized (VQ) model to encode and decode images to and from latent representations. - bert ([`LDMBertModel`]): - Text-encoder model based on [`~transformers.BERT`]. - tokenizer ([`~transformers.BertTokenizer`]): - A `BertTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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`]. - """ - - model_cpu_offload_seq = "bert->unet->vqvae" - - tokenizer: CLIPTokenizer - image_feature_extractor: CLIPImageProcessor - text_encoder: CLIPTextModelWithProjection - image_unet: UNet2DConditionModel - text_unet: UNetFlatConditionModel - vae: AutoencoderKL - scheduler: KarrasDiffusionSchedulers - - _optional_components = ["text_unet"] - - def __init__( - self, - tokenizer: CLIPTokenizer, - text_encoder: CLIPTextModelWithProjection, - image_unet: UNet2DConditionModel, - text_unet: UNetFlatConditionModel, - vae: AutoencoderKL, - scheduler: KarrasDiffusionSchedulers, - ): - super().__init__() - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - image_unet=image_unet, - text_unet=text_unet, - vae=vae, - 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) - - if self.text_unet is not None: - self._swap_unet_attention_blocks() - - def _swap_unet_attention_blocks(self): - """ - Swap the `Transformer2DModel` blocks between the image and text UNets - """ - for name, module in self.image_unet.named_modules(): - if isinstance(module, Transformer2DModel): - parent_name, index = name.rsplit(".", 1) - index = int(index) - self.image_unet.get_submodule(parent_name)[index], self.text_unet.get_submodule(parent_name)[index] = ( - self.text_unet.get_submodule(parent_name)[index], - self.image_unet.get_submodule(parent_name)[index], - ) - - def remove_unused_weights(self): - self.register_modules(text_unet=None) - - def _encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`): - 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]`): - 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`). - """ - - def normalize_embeddings(encoder_output): - embeds = self.text_encoder.text_projection(encoder_output.last_hidden_state) - embeds_pooled = encoder_output.text_embeds - embeds = embeds / torch.norm(embeds_pooled.unsqueeze(1), dim=-1, keepdim=True) - return embeds - - batch_size = len(prompt) if isinstance(prompt, list) else 1 - - 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="max_length", return_tensors="pt").input_ids - - if 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 - - prompt_embeds = self.text_encoder( - text_input_ids.to(device), - attention_mask=attention_mask, - ) - prompt_embeds = normalize_embeddings(prompt_embeds) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - bs_embed, seq_len, _ = prompt_embeds.shape - 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: - uncond_tokens: List[str] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif 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 - - max_length = text_input_ids.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 = normalize_embeddings(negative_prompt_embeds) - - # 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.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - return prompt_embeds - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - 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}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @torch.no_grad() - def __call__( - self, - prompt: Union[str, List[str]], - 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, - output_type: Optional[str] = "pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - width (`int`, *optional*, defaults to `self.image_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): - 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*): - 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): - The number of images to generate per prompt. - 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`, *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`. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - - Examples: - - ```py - >>> from diffusers import VersatileDiffusionTextToImagePipeline - >>> import torch - - >>> pipe = VersatileDiffusionTextToImagePipeline.from_pretrained( - ... "shi-labs/versatile-diffusion", torch_dtype=torch.float16 - ... ) - >>> pipe.remove_unused_weights() - >>> pipe = pipe.to("cuda") - - >>> generator = torch.Generator(device="cuda").manual_seed(0) - >>> image = pipe("an astronaut riding on a horse on mars", generator=generator).images[0] - >>> image.save("./astronaut.png") - ``` - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned, - otherwise a `tuple` is returned where the first element is a list with the generated images. - """ - # 0. Default height and width to unet - height = height or self.image_unet.config.sample_size * self.vae_scale_factor - width = width or self.image_unet.config.sample_size * self.vae_scale_factor - - # 1. Check inputs. Raise error if not correct - self.check_inputs(prompt, height, width, callback_steps) - - # 2. Define call parameters - batch_size = 1 if isinstance(prompt, str) else len(prompt) - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds = self._encode_prompt( - prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt - ) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.image_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, - ) - - # 6. Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.image_unet(latent_model_input, t, encoder_hidden_states=prompt_embeds).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # call the callback, if provided - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - else: - image = latents - - image = self.image_processor.postprocess(image, output_type=output_type) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/vq_diffusion/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/vq_diffusion/__init__.py deleted file mode 100644 index 070903377c7188415af0417d4839d74a8a34dc01..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/vq_diffusion/__init__.py +++ /dev/null @@ -1,57 +0,0 @@ -from typing import TYPE_CHECKING - -from ....utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - 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.dummy_torch_and_transformers_objects import ( - LearnedClassifierFreeSamplingEmbeddings, - VQDiffusionPipeline, - ) - - _dummy_objects.update( - { - "LearnedClassifierFreeSamplingEmbeddings": LearnedClassifierFreeSamplingEmbeddings, - "VQDiffusionPipeline": VQDiffusionPipeline, - } - ) -else: - _import_structure["pipeline_vq_diffusion"] = ["LearnedClassifierFreeSamplingEmbeddings", "VQDiffusionPipeline"] - - -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 ( - LearnedClassifierFreeSamplingEmbeddings, - VQDiffusionPipeline, - ) - else: - from .pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings, VQDiffusionPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/vq_diffusion/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/vq_diffusion/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index cbdf8a512e4a6d6714cd2c4d96d4d1110ebab6f5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/vq_diffusion/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/vq_diffusion/__pycache__/pipeline_vq_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/vq_diffusion/__pycache__/pipeline_vq_diffusion.cpython-310.pyc deleted file mode 100644 index 3373adb44740c1067a548fad92c0b555e422347e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/vq_diffusion/__pycache__/pipeline_vq_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py deleted file mode 100644 index e8617a54b69160df2fe4153e3619a7baf2742e0f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py +++ /dev/null @@ -1,325 +0,0 @@ -# Copyright 2025 Microsoft 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, List, Optional, Tuple, Union - -import torch -from transformers import CLIPTextModel, CLIPTokenizer - -from ....configuration_utils import ConfigMixin, register_to_config -from ....models import ModelMixin, Transformer2DModel, VQModel -from ....schedulers import VQDiffusionScheduler -from ....utils import logging -from ...pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class LearnedClassifierFreeSamplingEmbeddings(ModelMixin, ConfigMixin): - """ - Utility class for storing learned text embeddings for classifier free sampling - """ - - @register_to_config - def __init__(self, learnable: bool, hidden_size: Optional[int] = None, length: Optional[int] = None): - super().__init__() - - self.learnable = learnable - - if self.learnable: - assert hidden_size is not None, "learnable=True requires `hidden_size` to be set" - assert length is not None, "learnable=True requires `length` to be set" - - embeddings = torch.zeros(length, hidden_size) - else: - embeddings = None - - self.embeddings = torch.nn.Parameter(embeddings) - - -class VQDiffusionPipeline(DiffusionPipeline): - r""" - Pipeline for text-to-image generation using VQ Diffusion. - - 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: - vqvae ([`VQModel`]): - Vector Quantized Variational Auto-Encoder (VAE) model to encode and decode images to and from latent - representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-base-patch32](https://huggingface.co/openai/clip-vit-base-patch32)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - transformer ([`Transformer2DModel`]): - A conditional `Transformer2DModel` to denoise the encoded image latents. - scheduler ([`VQDiffusionScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - """ - - vqvae: VQModel - text_encoder: CLIPTextModel - tokenizer: CLIPTokenizer - transformer: Transformer2DModel - learned_classifier_free_sampling_embeddings: LearnedClassifierFreeSamplingEmbeddings - scheduler: VQDiffusionScheduler - - def __init__( - self, - vqvae: VQModel, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - transformer: Transformer2DModel, - scheduler: VQDiffusionScheduler, - learned_classifier_free_sampling_embeddings: LearnedClassifierFreeSamplingEmbeddings, - ): - super().__init__() - - self.register_modules( - vqvae=vqvae, - transformer=transformer, - text_encoder=text_encoder, - tokenizer=tokenizer, - scheduler=scheduler, - learned_classifier_free_sampling_embeddings=learned_classifier_free_sampling_embeddings, - ) - - def _encode_prompt(self, prompt, num_images_per_prompt, do_classifier_free_guidance): - batch_size = len(prompt) if isinstance(prompt, list) else 1 - - # get prompt text embeddings - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - - if text_input_ids.shape[-1] > self.tokenizer.model_max_length: - removed_text = self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]) - 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}" - ) - text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length] - prompt_embeds = self.text_encoder(text_input_ids.to(self.device))[0] - - # NOTE: This additional step of normalizing the text embeddings is from VQ-Diffusion. - # While CLIP does normalize the pooled output of the text transformer when combining - # the image and text embeddings, CLIP does not directly normalize the last hidden state. - # - # CLIP normalizing the pooled output. - # https://github.com/huggingface/transformers/blob/d92e22d1f28324f513f3080e5c47c071a3916721/src/transformers/models/clip/modeling_clip.py#L1052-L1053 - prompt_embeds = prompt_embeds / prompt_embeds.norm(dim=-1, keepdim=True) - - # duplicate text embeddings for each generation per prompt - prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - if self.learned_classifier_free_sampling_embeddings.learnable: - negative_prompt_embeds = self.learned_classifier_free_sampling_embeddings.embeddings - negative_prompt_embeds = negative_prompt_embeds.unsqueeze(0).repeat(batch_size, 1, 1) - else: - uncond_tokens = [""] * batch_size - - max_length = text_input_ids.shape[-1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - negative_prompt_embeds = self.text_encoder(uncond_input.input_ids.to(self.device))[0] - # See comment for normalizing text embeddings - negative_prompt_embeds = negative_prompt_embeds / negative_prompt_embeds.norm(dim=-1, keepdim=True) - - # 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.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - return prompt_embeds - - @torch.no_grad() - def __call__( - self, - prompt: Union[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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - ) -> Union[ImagePipelineOutput, Tuple]: - """ - The call function to the pipeline for generation. - - Args: - 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 - expense of slower inference. - 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`. - truncation_rate (`float`, *optional*, defaults to 1.0 (equivalent to no truncation)): - Used to "truncate" the predicted classes for x_0 such that the cumulative probability for a pixel is at - most `truncation_rate`. The lowest probabilities that would increase the cumulative probability above - `truncation_rate` are set to zero. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - generator (`torch.Generator`, *optional*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - latents (`torch.Tensor` of shape (batch), *optional*): - Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image - generation. Must be valid embedding indices.If not provided, a latents tensor will be generated of - completely masked latent pixels. - 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 [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images. - """ - if isinstance(prompt, str): - batch_size = 1 - elif isinstance(prompt, list): - batch_size = len(prompt) - else: - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - batch_size = batch_size * num_images_per_prompt - - do_classifier_free_guidance = guidance_scale > 1.0 - - prompt_embeds = self._encode_prompt(prompt, num_images_per_prompt, do_classifier_free_guidance) - - if (callback_steps is None) or ( - 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)}." - ) - - # get the initial completely masked latents unless the user supplied it - - latents_shape = (batch_size, self.transformer.num_latent_pixels) - if latents is None: - mask_class = self.transformer.num_vector_embeds - 1 - latents = torch.full(latents_shape, mask_class).to(self.device) - else: - if latents.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - if (latents < 0).any() or (latents >= self.transformer.num_vector_embeds).any(): - raise ValueError( - "Unexpected latents value(s). All latents be valid embedding indices i.e. in the range 0," - f" {self.transformer.num_vector_embeds - 1} (inclusive)." - ) - latents = latents.to(self.device) - - # set timesteps - self.scheduler.set_timesteps(num_inference_steps, device=self.device) - - timesteps_tensor = self.scheduler.timesteps.to(self.device) - - sample = latents - - for i, t in enumerate(self.progress_bar(timesteps_tensor)): - # expand the sample if we are doing classifier free guidance - latent_model_input = torch.cat([sample] * 2) if do_classifier_free_guidance else sample - - # predict the un-noised image - # model_output == `log_p_x_0` - model_output = self.transformer(latent_model_input, encoder_hidden_states=prompt_embeds, timestep=t).sample - - if do_classifier_free_guidance: - model_output_uncond, model_output_text = model_output.chunk(2) - model_output = model_output_uncond + guidance_scale * (model_output_text - model_output_uncond) - model_output -= torch.logsumexp(model_output, dim=1, keepdim=True) - - model_output = self.truncate(model_output, truncation_rate) - - # remove `log(0)`'s (`-inf`s) - model_output = model_output.clamp(-70) - - # compute the previous noisy sample x_t -> x_t-1 - sample = self.scheduler.step(model_output, timestep=t, sample=sample, generator=generator).prev_sample - - # call the callback, if provided - if callback is not None and i % callback_steps == 0: - callback(i, t, sample) - - embedding_channels = self.vqvae.config.vq_embed_dim - embeddings_shape = (batch_size, self.transformer.height, self.transformer.width, embedding_channels) - embeddings = self.vqvae.quantize.get_codebook_entry(sample, shape=embeddings_shape) - image = self.vqvae.decode(embeddings, force_not_quantize=True).sample - - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).numpy() - - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) - - def truncate(self, log_p_x_0: torch.Tensor, truncation_rate: float) -> torch.Tensor: - """ - Truncates `log_p_x_0` such that for each column vector, the total cumulative probability is `truncation_rate` - The lowest probabilities that would increase the cumulative probability above `truncation_rate` are set to - zero. - """ - sorted_log_p_x_0, indices = torch.sort(log_p_x_0, 1, descending=True) - sorted_p_x_0 = torch.exp(sorted_log_p_x_0) - keep_mask = sorted_p_x_0.cumsum(dim=1) < truncation_rate - - # Ensure that at least the largest probability is not zeroed out - all_true = torch.full_like(keep_mask[:, 0:1, :], True) - keep_mask = torch.cat((all_true, keep_mask), dim=1) - keep_mask = keep_mask[:, :-1, :] - - keep_mask = keep_mask.gather(1, indices.argsort(1)) - - rv = log_p_x_0.clone() - - rv[~keep_mask] = -torch.inf # -inf = log(0) - - return rv diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/dit/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/dit/__init__.py deleted file mode 100644 index fe2a94f3cba77d867f97111a41895918842adc27..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/dit/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import DIFFUSERS_SLOW_IMPORT, _LazyModule - - -_import_structure = {"pipeline_dit": ["DiTPipeline"]} - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - from .pipeline_dit import DiTPipeline - -else: - import sys - - sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/dit/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/dit/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 99b44baaa108464592ad3040fa7e7aa2c99a15f2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/dit/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/dit/__pycache__/pipeline_dit.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/dit/__pycache__/pipeline_dit.cpython-310.pyc deleted file mode 100644 index 931a110b0ff981b837b34c34ea308605915b9a91..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/dit/__pycache__/pipeline_dit.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/dit/pipeline_dit.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/dit/pipeline_dit.py deleted file mode 100644 index 14f63ea229ac608e1faf4e6de7114ad072282e90..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/dit/pipeline_dit.py +++ /dev/null @@ -1,249 +0,0 @@ -# Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) -# William Peebles and Saining Xie -# -# Copyright (c) 2021 OpenAI -# MIT License -# -# 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 Dict, List, Optional, Tuple, Union - -import torch - -from ...models import AutoencoderKL, DiTTransformer2DModel -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import is_torch_xla_available -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - - XLA_AVAILABLE = True -else: - XLA_AVAILABLE = False - - -class DiTPipeline(DiffusionPipeline): - r""" - Pipeline for image generation based on a Transformer backbone instead of a UNet. - - 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.). - - Parameters: - transformer ([`DiTTransformer2DModel`]): - A class conditioned `DiTTransformer2DModel` to denoise the encoded image latents. - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - scheduler ([`DDIMScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - """ - - model_cpu_offload_seq = "transformer->vae" - - def __init__( - self, - transformer: DiTTransformer2DModel, - vae: AutoencoderKL, - scheduler: KarrasDiffusionSchedulers, - id2label: Optional[Dict[int, str]] = None, - ): - super().__init__() - self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler) - - # create a imagenet -> id dictionary for easier use - self.labels = {} - if id2label is not None: - for key, value in id2label.items(): - for label in value.split(","): - 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]: - r""" - - Map label strings from ImageNet to corresponding class ids. - - Parameters: - label (`str` or `dict` of `str`): - Label strings to be mapped to class ids. - - Returns: - `list` of `int`: - Class ids to be processed by pipeline. - """ - - if not isinstance(label, list): - label = list(label) - - for l in label: - if l not in self.labels: - raise ValueError( - f"{l} does not exist. Please make sure to select one of the following labels: \n {self.labels}." - ) - - return [self.labels[l] for l in label] - - @torch.no_grad() - def __call__( - self, - class_labels: List[int], - guidance_scale: float = 4.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - num_inference_steps: int = 50, - output_type: Optional[str] = "pil", - return_dict: bool = True, - ) -> Union[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. - 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`. - generator (`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 250): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - 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 [`ImagePipelineOutput`] instead of a plain tuple. - - Examples: - - ```py - >>> from diffusers import DiTPipeline, DPMSolverMultistepScheduler - >>> import torch - - >>> pipe = DiTPipeline.from_pretrained("facebook/DiT-XL-2-256", torch_dtype=torch.float16) - >>> pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config) - >>> pipe = pipe.to("cuda") - - >>> # pick words from Imagenet class labels - >>> pipe.labels # to print all available words - - >>> # pick words that exist in ImageNet - >>> words = ["white shark", "umbrella"] - - >>> class_ids = pipe.get_label_ids(words) - - >>> generator = torch.manual_seed(33) - >>> output = pipe(class_labels=class_ids, num_inference_steps=25, generator=generator) - - >>> image = output.images[0] # label 'white shark' - ``` - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images - """ - - batch_size = len(class_labels) - latent_size = self.transformer.config.sample_size - latent_channels = self.transformer.config.in_channels - - latents = randn_tensor( - shape=(batch_size, latent_channels, latent_size, latent_size), - generator=generator, - device=self._execution_device, - dtype=self.transformer.dtype, - ) - latent_model_input = torch.cat([latents] * 2) if guidance_scale > 1 else latents - - class_labels = torch.tensor(class_labels, device=self._execution_device).reshape(-1) - class_null = torch.tensor([1000] * batch_size, device=self._execution_device) - class_labels_input = torch.cat([class_labels, class_null], 0) if guidance_scale > 1 else class_labels - - # set step values - self.scheduler.set_timesteps(num_inference_steps) - for t in self.progress_bar(self.scheduler.timesteps): - if guidance_scale > 1: - half = latent_model_input[: len(latent_model_input) // 2] - latent_model_input = torch.cat([half, half], dim=0) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - timesteps = t - if not torch.is_tensor(timesteps): - # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = latent_model_input.device.type == "mps" - is_npu = latent_model_input.device.type == "npu" - if isinstance(timesteps, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - timesteps = torch.tensor([timesteps], dtype=dtype, device=latent_model_input.device) - elif len(timesteps.shape) == 0: - timesteps = timesteps[None].to(latent_model_input.device) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps.expand(latent_model_input.shape[0]) - # predict noise model_output - noise_pred = self.transformer( - latent_model_input, timestep=timesteps, class_labels=class_labels_input - ).sample - - # perform guidance - if guidance_scale > 1: - eps, rest = noise_pred[:, :latent_channels], noise_pred[:, latent_channels:] - cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0) - - half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps) - eps = torch.cat([half_eps, half_eps], dim=0) - - noise_pred = torch.cat([eps, rest], dim=1) - - # learned sigma - if self.transformer.config.out_channels // 2 == latent_channels: - model_output, _ = torch.split(noise_pred, latent_channels, dim=1) - else: - model_output = noise_pred - - # compute previous image: x_t -> x_t-1 - latent_model_input = self.scheduler.step(model_output, t, latent_model_input).prev_sample - - if XLA_AVAILABLE: - xm.mark_step() - - if guidance_scale > 1: - latents, _ = latent_model_input.chunk(2, dim=0) - else: - latents = latent_model_input - - latents = 1 / self.vae.config.scaling_factor * latents - samples = self.vae.decode(latents).sample - - samples = (samples / 2 + 0.5).clamp(0, 1) - - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - samples = samples.cpu().permute(0, 2, 3, 1).float().numpy() - - if output_type == "pil": - samples = self.numpy_to_pil(samples) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (samples,) - - return ImagePipelineOutput(images=samples) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__init__.py deleted file mode 100644 index 49923423f951c1e6124ddb9b0506f740f28a7a5a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__init__.py +++ /dev/null @@ -1,52 +0,0 @@ -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_easyanimate"] = ["EasyAnimatePipeline"] - _import_structure["pipeline_easyanimate_control"] = ["EasyAnimateControlPipeline"] - _import_structure["pipeline_easyanimate_inpaint"] = ["EasyAnimateInpaintPipeline"] - -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_easyanimate import EasyAnimatePipeline - from .pipeline_easyanimate_control import EasyAnimateControlPipeline - from .pipeline_easyanimate_inpaint import EasyAnimateInpaintPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index ed8a56e57ff866a3a01c5b02abb5718d499c38a4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/pipeline_easyanimate.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/pipeline_easyanimate.cpython-310.pyc deleted file mode 100644 index a457f381d755d1b6ba9624491cfbad66f205970e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/pipeline_easyanimate.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/pipeline_easyanimate_control.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/pipeline_easyanimate_control.cpython-310.pyc deleted file mode 100644 index 7a1b79e6c0960a9cf4506e9c1ffd54c37fb50222..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/pipeline_easyanimate_control.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/pipeline_easyanimate_inpaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/pipeline_easyanimate_inpaint.cpython-310.pyc deleted file mode 100644 index 40baf043f66db00d98a5d3fd2ebb407981626349..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/pipeline_easyanimate_inpaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index 644314876fa290bdef41a61bf449713db1b577d4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/pipeline_easyanimate.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/pipeline_easyanimate.py deleted file mode 100644 index 92239c0d32f054c15eaf2d3cbbaf7075655df1e3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/pipeline_easyanimate.py +++ /dev/null @@ -1,770 +0,0 @@ -# Copyright 2025 The EasyAnimate 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 Callable, Dict, List, Optional, Union - -import torch -from transformers import ( - BertModel, - BertTokenizer, - Qwen2Tokenizer, - Qwen2VLForConditionalGeneration, -) - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...models import AutoencoderKLMagvit, EasyAnimateTransformer3DModel -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 ...video_processor import VideoProcessor -from .pipeline_output import EasyAnimatePipelineOutput - - -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 EasyAnimatePipeline - >>> from diffusers.utils import export_to_video - - >>> # Models: "alibaba-pai/EasyAnimateV5.1-12b-zh" - >>> pipe = EasyAnimatePipeline.from_pretrained( - ... "alibaba-pai/EasyAnimateV5.1-7b-zh-diffusers", torch_dtype=torch.float16 - ... ).to("cuda") - >>> prompt = ( - ... "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. " - ... "The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other " - ... "pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, " - ... "casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. " - ... "The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical " - ... "atmosphere of this unique musical performance." - ... ) - >>> sample_size = (512, 512) - >>> video = pipe( - ... prompt=prompt, - ... guidance_scale=6, - ... negative_prompt="bad detailed", - ... height=sample_size[0], - ... width=sample_size[1], - ... num_inference_steps=50, - ... ).frames[0] - >>> export_to_video(video, "output.mp4", fps=8) - ``` -""" - - -# Similar to diffusers.pipelines.hunyuandit.pipeline_hunyuandit.get_resize_crop_region_for_grid -def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): - tw = tgt_width - th = tgt_height - h, w = src - r = h / w - if r > (th / tw): - resize_height = th - resize_width = int(round(th / h * w)) - else: - resize_width = tw - resize_height = int(round(tw / w * h)) - - crop_top = int(round((th - resize_height) / 2.0)) - crop_left = int(round((tw - resize_width) / 2.0)) - - return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width) - - -# 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 - - -# 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 EasyAnimatePipeline(DiffusionPipeline): - r""" - Pipeline for text-to-video generation using EasyAnimate. - - 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.) - - EasyAnimate uses one text encoder [qwen2 vl](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct) in V5.1. - - 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`]): - EasyAnimate uses [qwen2 vl](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct) in V5.1. - tokenizer (Optional[`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer`]): - A `Qwen2Tokenizer` or `BertTokenizer` to tokenize text. - transformer ([`EasyAnimateTransformer3DModel`]): - The EasyAnimate model designed by EasyAnimate Team. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - A scheduler to be used in combination with EasyAnimate to denoise the encoded image latents. - """ - - model_cpu_offload_seq = "text_encoder->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKLMagvit, - text_encoder: Union[Qwen2VLForConditionalGeneration, BertModel], - tokenizer: Union[Qwen2Tokenizer, BertTokenizer], - transformer: EasyAnimateTransformer3DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - ) - self.enable_text_attention_mask = ( - self.transformer.config.enable_text_attention_mask - if getattr(self, "transformer", None) is not None - else True - ) - self.vae_spatial_compression_ratio = ( - self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 8 - ) - self.vae_temporal_compression_ratio = ( - self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 4 - ) - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) - - def encode_prompt( - self, - prompt: Union[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, - max_sequence_length: int = 256, - ): - 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 - dtype (`torch.dtype`): - torch dtype - 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. - prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the prompt. Required when `prompt_embeds` is passed directly. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the negative prompt. Required when `negative_prompt_embeds` is passed directly. - max_sequence_length (`int`, *optional*): maximum sequence length to use for the prompt. - """ - dtype = dtype or self.text_encoder.dtype - device = device or self.text_encoder.device - - 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: - if isinstance(prompt, str): - messages = [ - { - "role": "user", - "content": [{"type": "text", "text": prompt}], - } - ] - else: - messages = [ - { - "role": "user", - "content": [{"type": "text", "text": _prompt}], - } - for _prompt in prompt - ] - text = [ - self.tokenizer.apply_chat_template([m], tokenize=False, add_generation_prompt=True) for m in messages - ] - - text_inputs = self.tokenizer( - text=text, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_attention_mask=True, - padding_side="right", - return_tensors="pt", - ) - text_inputs = text_inputs.to(self.text_encoder.device) - - text_input_ids = text_inputs.input_ids - prompt_attention_mask = text_inputs.attention_mask - if self.enable_text_attention_mask: - # Inference: Generation of the output - prompt_embeds = self.text_encoder( - input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True - ).hidden_states[-2] - else: - raise ValueError("LLM needs attention_mask") - prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1) - - prompt_embeds = prompt_embeds.to(dtype=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) - prompt_attention_mask = prompt_attention_mask.to(device=device) - - # get unconditional embeddings for classifier free guidance - if do_classifier_free_guidance and negative_prompt_embeds is None: - if negative_prompt is not None and isinstance(negative_prompt, str): - messages = [ - { - "role": "user", - "content": [{"type": "text", "text": negative_prompt}], - } - ] - else: - messages = [ - { - "role": "user", - "content": [{"type": "text", "text": _negative_prompt}], - } - for _negative_prompt in negative_prompt - ] - text = [ - self.tokenizer.apply_chat_template([m], tokenize=False, add_generation_prompt=True) for m in messages - ] - - text_inputs = self.tokenizer( - text=text, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_attention_mask=True, - padding_side="right", - return_tensors="pt", - ) - text_inputs = text_inputs.to(self.text_encoder.device) - - text_input_ids = text_inputs.input_ids - negative_prompt_attention_mask = text_inputs.attention_mask - if self.enable_text_attention_mask: - # Inference: Generation of the output - negative_prompt_embeds = self.text_encoder( - input_ids=text_input_ids, - attention_mask=negative_prompt_attention_mask, - output_hidden_states=True, - ).hidden_states[-2] - else: - raise ValueError("LLM needs attention_mask") - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 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] - - negative_prompt_embeds = negative_prompt_embeds.to(dtype=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) - negative_prompt_attention_mask = negative_prompt_attention_mask.to(device=device) - - return prompt_embeds, negative_prompt_embeds, prompt_attention_mask, negative_prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - callback_on_step_end_tensor_inputs=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 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 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 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}." - ) - - def prepare_latents( - self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None - ): - if latents is not None: - return latents.to(device=device, dtype=dtype) - - shape = ( - batch_size, - num_channels_latents, - (num_frames - 1) // self.vae_temporal_compression_ratio + 1, - height // self.vae_spatial_compression_ratio, - width // self.vae_spatial_compression_ratio, - ) - - 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) - # scale the initial noise by the standard deviation required by the scheduler - if hasattr(self.scheduler, "init_noise_sigma"): - latents = latents * self.scheduler.init_noise_sigma - return latents - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - # 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 - - @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, - 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", - 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"], - 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*): - 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). - height (`int`, *optional*): - Height of the generated image in pixels. - width (`int`, *optional*): - Width of the generated image in pixels. - num_inference_steps (`int`, *optional*, defaults to 50): - Number of denoising steps during generation. More steps generally yield higher quality images but slow - 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*): - 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*): - A generator to ensure reproducibility in image generation. - latents (`torch.Tensor`, *optional*): - Predefined latent tensors to condition generation. - prompt_embeds (`torch.Tensor`, *optional*): - Text embeddings for the prompts. Overrides prompt string inputs for more flexibility. - negative_prompt_embeds (`torch.Tensor`, *optional*): - Embeddings for negative prompts. Overrides string inputs if defined. - prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the primary prompt embeddings. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for negative prompt embeddings. - output_type (`str`, *optional*, defaults to "latent"): - Format of the generated output, either as a PIL image or as a NumPy array. - return_dict (`bool`, *optional*, defaults to `True`): - 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*): - 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 dimensions of the output. - target_size (`Tuple[int, int]`, *optional*): - Desired output dimensions for calculations. - crops_coords_top_left (`Tuple[int, int]`, *optional*, defaults to `(0, 0)`): - Coordinates for cropping. - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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 isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - # 0. default height and width - height = int((height // 16) * 16) - width = int((width // 16) * 16) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - callback_on_step_end_tensor_inputs, - ) - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - 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 - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - else: - dtype = self.transformer.dtype - - # 3. Encode input prompt - ( - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt=prompt, - device=device, - dtype=dtype, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - 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, - ) - - # 4. Prepare timesteps - if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler): - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, mu=1 - ) - else: - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) - - # 5. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - num_frames, - height, - width, - dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask]) - - prompt_embeds = prompt_embeds.to(device=device) - prompt_attention_mask = prompt_attention_mask.to(device=device) - - # 7. Denoising loop - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - if hasattr(self.scheduler, "scale_model_input"): - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # expand scalar t to 1-D tensor to match the 1st dim of latent_model_input - t_expand = torch.tensor([t] * latent_model_input.shape[0], device=device).to( - dtype=latent_model_input.dtype - ) - - # predict the noise residual - noise_pred = self.transformer( - latent_model_input, - t_expand, - encoder_hidden_states=prompt_embeds, - return_dict=False, - )[0] - - if noise_pred.size()[1] != self.vae.config.latent_channels: - noise_pred, _ = noise_pred.chunk(2, dim=1) - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - if self.do_classifier_free_guidance and guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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 not output_type == "latent": - latents = 1 / self.vae.config.scaling_factor * latents - video = self.vae.decode(latents, return_dict=False)[0] - video = self.video_processor.postprocess_video(video=video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return EasyAnimatePipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/pipeline_easyanimate_control.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/pipeline_easyanimate_control.py deleted file mode 100644 index f74a11f87d752acc611c38c635ec83d78469573b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/pipeline_easyanimate_control.py +++ /dev/null @@ -1,994 +0,0 @@ -# Copyright 2025 The EasyAnimate 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 Callable, Dict, List, Optional, Union - -import numpy as np -import torch -import torch.nn.functional as F -from PIL import Image -from transformers import ( - BertModel, - BertTokenizer, - Qwen2Tokenizer, - Qwen2VLForConditionalGeneration, -) - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import VaeImageProcessor -from ...models import AutoencoderKLMagvit, EasyAnimateTransformer3DModel -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 ...video_processor import VideoProcessor -from .pipeline_output import EasyAnimatePipelineOutput - - -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 EasyAnimateControlPipeline - >>> from diffusers.pipelines.easyanimate.pipeline_easyanimate_control import get_video_to_video_latent - >>> from diffusers.utils import export_to_video, load_video - - >>> pipe = EasyAnimateControlPipeline.from_pretrained( - ... "alibaba-pai/EasyAnimateV5.1-12b-zh-Control-diffusers", torch_dtype=torch.bfloat16 - ... ) - >>> pipe.to("cuda") - - >>> control_video = load_video( - ... "https://huggingface.co/alibaba-pai/EasyAnimateV5.1-12b-zh-Control/blob/main/asset/pose.mp4" - ... ) - >>> prompt = ( - ... "In this sunlit outdoor garden, a beautiful woman is dressed in a knee-length, sleeveless white dress. " - ... "The hem of her dress gently sways with her graceful dance, much like a butterfly fluttering in the breeze. " - ... "Sunlight filters through the leaves, casting dappled shadows that highlight her soft features and clear eyes, " - ... "making her appear exceptionally elegant. It seems as if every movement she makes speaks of youth and vitality. " - ... "As she twirls on the grass, her dress flutters, as if the entire garden is rejoicing in her dance. " - ... "The colorful flowers around her sway in the gentle breeze, with roses, chrysanthemums, and lilies each " - ... "releasing their fragrances, creating a relaxed and joyful atmosphere." - ... ) - >>> sample_size = (672, 384) - >>> num_frames = 49 - - >>> input_video, _, _ = get_video_to_video_latent(control_video, num_frames, sample_size) - >>> video = pipe( - ... prompt, - ... num_frames=num_frames, - ... negative_prompt="Twisted body, limb deformities, text subtitles, comics, stillness, ugliness, errors, garbled text.", - ... height=sample_size[0], - ... width=sample_size[1], - ... control_video=input_video, - ... ).frames[0] - >>> export_to_video(video, "output.mp4", fps=8) - ``` -""" - - -def preprocess_image(image, sample_size): - """ - Preprocess a single image (PIL.Image, numpy.ndarray, or torch.Tensor) to a resized tensor. - """ - if isinstance(image, torch.Tensor): - # If input is a tensor, assume it's in CHW format and resize using interpolation - image = torch.nn.functional.interpolate( - image.unsqueeze(0), size=sample_size, mode="bilinear", align_corners=False - ).squeeze(0) - elif isinstance(image, Image.Image): - # If input is a PIL image, resize and convert to numpy array - image = image.resize((sample_size[1], sample_size[0])) - image = np.array(image) - elif isinstance(image, np.ndarray): - # If input is a numpy array, resize using PIL - image = Image.fromarray(image).resize((sample_size[1], sample_size[0])) - image = np.array(image) - else: - raise ValueError("Unsupported input type. Expected PIL.Image, numpy.ndarray, or torch.Tensor.") - - # Convert to tensor if not already - if not isinstance(image, torch.Tensor): - image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 # HWC -> CHW, normalize to [0, 1] - - return image - - -def get_video_to_video_latent(input_video, num_frames, sample_size, validation_video_mask=None, ref_image=None): - if input_video is not None: - # Convert each frame in the list to tensor - input_video = [preprocess_image(frame, sample_size=sample_size) for frame in input_video] - - # Stack all frames into a single tensor (F, C, H, W) - input_video = torch.stack(input_video)[:num_frames] - - # Add batch dimension (B, F, C, H, W) - input_video = input_video.permute(1, 0, 2, 3).unsqueeze(0) - - if validation_video_mask is not None: - # Handle mask input - validation_video_mask = preprocess_image(validation_video_mask, size=sample_size) - input_video_mask = torch.where(validation_video_mask < 240 / 255.0, 0.0, 255) - - # Adjust mask dimensions to match video - input_video_mask = input_video_mask.unsqueeze(0).unsqueeze(-1).permute([3, 0, 1, 2]).unsqueeze(0) - input_video_mask = torch.tile(input_video_mask, [1, 1, input_video.size()[2], 1, 1]) - input_video_mask = input_video_mask.to(input_video.device, input_video.dtype) - else: - input_video_mask = torch.zeros_like(input_video[:, :1]) - input_video_mask[:, :, :] = 255 - else: - input_video, input_video_mask = None, None - - if ref_image is not None: - # Convert reference image to tensor - ref_image = preprocess_image(ref_image, size=sample_size) - ref_image = ref_image.permute(1, 0, 2, 3).unsqueeze(0) # Add batch dimension (B, C, H, W) - else: - ref_image = None - - return input_video, input_video_mask, ref_image - - -# Similar to diffusers.pipelines.hunyuandit.pipeline_hunyuandit.get_resize_crop_region_for_grid -def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): - tw = tgt_width - th = tgt_height - h, w = src - r = h / w - if r > (th / tw): - resize_height = th - resize_width = int(round(th / h * w)) - else: - resize_width = tw - resize_height = int(round(tw / w * h)) - - crop_top = int(round((th - resize_height) / 2.0)) - crop_left = int(round((tw - resize_width) / 2.0)) - - return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width) - - -# 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 - - -# Resize mask information in magvit -def resize_mask(mask, latent, process_first_frame_only=True): - latent_size = latent.size() - - if process_first_frame_only: - target_size = list(latent_size[2:]) - target_size[0] = 1 - first_frame_resized = F.interpolate( - mask[:, :, 0:1, :, :], size=target_size, mode="trilinear", align_corners=False - ) - - target_size = list(latent_size[2:]) - target_size[0] = target_size[0] - 1 - if target_size[0] != 0: - remaining_frames_resized = F.interpolate( - mask[:, :, 1:, :, :], size=target_size, mode="trilinear", align_corners=False - ) - resized_mask = torch.cat([first_frame_resized, remaining_frames_resized], dim=2) - else: - resized_mask = first_frame_resized - else: - target_size = list(latent_size[2:]) - resized_mask = F.interpolate(mask, size=target_size, mode="trilinear", align_corners=False) - return resized_mask - - -# 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 EasyAnimateControlPipeline(DiffusionPipeline): - r""" - Pipeline for text-to-video generation using EasyAnimate. - - 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.) - - EasyAnimate uses one text encoder [qwen2 vl](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct) in V5.1. - - 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`]): - EasyAnimate uses [qwen2 vl](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct) in V5.1. - tokenizer (Optional[`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer`]): - A `Qwen2Tokenizer` or `BertTokenizer` to tokenize text. - transformer ([`EasyAnimateTransformer3DModel`]): - The EasyAnimate model designed by EasyAnimate Team. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - A scheduler to be used in combination with EasyAnimate to denoise the encoded image latents. - """ - - model_cpu_offload_seq = "text_encoder->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKLMagvit, - text_encoder: Union[Qwen2VLForConditionalGeneration, BertModel], - tokenizer: Union[Qwen2Tokenizer, BertTokenizer], - transformer: EasyAnimateTransformer3DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - ) - - self.enable_text_attention_mask = ( - self.transformer.config.enable_text_attention_mask - if getattr(self, "transformer", None) is not None - else True - ) - self.vae_spatial_compression_ratio = ( - self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 8 - ) - self.vae_temporal_compression_ratio = ( - self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 4 - ) - self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) - self.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_spatial_compression_ratio, - do_normalize=False, - do_binarize=True, - do_convert_grayscale=True, - ) - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) - - # Copied from diffusers.pipelines.easyanimate.pipeline_easyanimate.EasyAnimatePipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[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, - max_sequence_length: int = 256, - ): - 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 - dtype (`torch.dtype`): - torch dtype - 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. - prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the prompt. Required when `prompt_embeds` is passed directly. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the negative prompt. Required when `negative_prompt_embeds` is passed directly. - max_sequence_length (`int`, *optional*): maximum sequence length to use for the prompt. - """ - dtype = dtype or self.text_encoder.dtype - device = device or self.text_encoder.device - - 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: - if isinstance(prompt, str): - messages = [ - { - "role": "user", - "content": [{"type": "text", "text": prompt}], - } - ] - else: - messages = [ - { - "role": "user", - "content": [{"type": "text", "text": _prompt}], - } - for _prompt in prompt - ] - text = [ - self.tokenizer.apply_chat_template([m], tokenize=False, add_generation_prompt=True) for m in messages - ] - - text_inputs = self.tokenizer( - text=text, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_attention_mask=True, - padding_side="right", - return_tensors="pt", - ) - text_inputs = text_inputs.to(self.text_encoder.device) - - text_input_ids = text_inputs.input_ids - prompt_attention_mask = text_inputs.attention_mask - if self.enable_text_attention_mask: - # Inference: Generation of the output - prompt_embeds = self.text_encoder( - input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True - ).hidden_states[-2] - else: - raise ValueError("LLM needs attention_mask") - prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1) - - prompt_embeds = prompt_embeds.to(dtype=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) - prompt_attention_mask = prompt_attention_mask.to(device=device) - - # get unconditional embeddings for classifier free guidance - if do_classifier_free_guidance and negative_prompt_embeds is None: - if negative_prompt is not None and isinstance(negative_prompt, str): - messages = [ - { - "role": "user", - "content": [{"type": "text", "text": negative_prompt}], - } - ] - else: - messages = [ - { - "role": "user", - "content": [{"type": "text", "text": _negative_prompt}], - } - for _negative_prompt in negative_prompt - ] - text = [ - self.tokenizer.apply_chat_template([m], tokenize=False, add_generation_prompt=True) for m in messages - ] - - text_inputs = self.tokenizer( - text=text, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_attention_mask=True, - padding_side="right", - return_tensors="pt", - ) - text_inputs = text_inputs.to(self.text_encoder.device) - - text_input_ids = text_inputs.input_ids - negative_prompt_attention_mask = text_inputs.attention_mask - if self.enable_text_attention_mask: - # Inference: Generation of the output - negative_prompt_embeds = self.text_encoder( - input_ids=text_input_ids, - attention_mask=negative_prompt_attention_mask, - output_hidden_states=True, - ).hidden_states[-2] - else: - raise ValueError("LLM needs attention_mask") - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 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] - - negative_prompt_embeds = negative_prompt_embeds.to(dtype=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) - negative_prompt_attention_mask = negative_prompt_attention_mask.to(device=device) - - return prompt_embeds, negative_prompt_embeds, prompt_attention_mask, negative_prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - callback_on_step_end_tensor_inputs=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 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 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 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}." - ) - - def prepare_latents( - self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None - ): - if latents is not None: - return latents.to(device=device, dtype=dtype) - - shape = ( - batch_size, - num_channels_latents, - (num_frames - 1) // self.vae_temporal_compression_ratio + 1, - height // self.vae_spatial_compression_ratio, - width // self.vae_spatial_compression_ratio, - ) - - 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) - # scale the initial noise by the standard deviation required by the scheduler - if hasattr(self.scheduler, "init_noise_sigma"): - latents = latents * self.scheduler.init_noise_sigma - return latents - - def prepare_control_latents( - self, control, control_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance - ): - # resize the control to latents shape as we concatenate the control to the latents - # we do that before converting to dtype to avoid breaking in case we're using cpu_offload - # and half precision - - if control is not None: - control = control.to(device=device, dtype=dtype) - bs = 1 - new_control = [] - for i in range(0, control.shape[0], bs): - control_bs = control[i : i + bs] - control_bs = self.vae.encode(control_bs)[0] - control_bs = control_bs.mode() - new_control.append(control_bs) - control = torch.cat(new_control, dim=0) - control = control * self.vae.config.scaling_factor - - if control_image is not None: - control_image = control_image.to(device=device, dtype=dtype) - bs = 1 - new_control_pixel_values = [] - for i in range(0, control_image.shape[0], bs): - control_pixel_values_bs = control_image[i : i + bs] - control_pixel_values_bs = self.vae.encode(control_pixel_values_bs)[0] - control_pixel_values_bs = control_pixel_values_bs.mode() - new_control_pixel_values.append(control_pixel_values_bs) - control_image_latents = torch.cat(new_control_pixel_values, dim=0) - control_image_latents = control_image_latents * self.vae.config.scaling_factor - else: - control_image_latents = None - - return control, control_image_latents - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - # 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 - - @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, - 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", - 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"], - guidance_rescale: float = 0.0, - timesteps: Optional[List[int]] = None, - ): - r""" - Generates images or video using the EasyAnimate pipeline based on the provided prompts. - - Examples: - 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). - height (`int`, *optional*): - Height of the generated image in pixels. - width (`int`, *optional*): - Width of the generated image in pixels. - num_inference_steps (`int`, *optional*, defaults to 50): - Number of denoising steps during generation. More steps generally yield higher quality images but slow - 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*): - 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*): - A generator to ensure reproducibility in image generation. - latents (`torch.Tensor`, *optional*): - Predefined latent tensors to condition generation. - prompt_embeds (`torch.Tensor`, *optional*): - Text embeddings for the prompts. Overrides prompt string inputs for more flexibility. - negative_prompt_embeds (`torch.Tensor`, *optional*): - Embeddings for negative prompts. Overrides string inputs if defined. - prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the primary prompt embeddings. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for negative prompt embeddings. - output_type (`str`, *optional*, defaults to "latent"): - Format of the generated output, either as a PIL image or as a NumPy array. - return_dict (`bool`, *optional*, defaults to `True`): - 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*): - Tensor names to be included in callback function calls. - guidance_rescale (`float`, *optional*, defaults to 0.0): - Adjusts noise levels based on guidance scale. - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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 isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - # 0. default height and width - height = int((height // 16) * 16) - width = int((width // 16) * 16) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - callback_on_step_end_tensor_inputs, - ) - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - 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 - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - else: - dtype = self.transformer.dtype - - # 3. Encode input prompt - ( - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt=prompt, - device=device, - dtype=dtype, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - 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, - text_encoder_index=0, - ) - - # 4. Prepare timesteps - if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler): - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, mu=1 - ) - else: - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.vae.config.latent_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - num_frames, - height, - width, - dtype, - device, - generator, - latents, - ) - - if control_camera_video is not None: - control_video_latents = resize_mask(control_camera_video, latents, process_first_frame_only=True) - control_video_latents = control_video_latents * 6 - control_latents = ( - torch.cat([control_video_latents] * 2) if self.do_classifier_free_guidance else control_video_latents - ).to(device, dtype) - elif control_video is not None: - batch_size, channels, num_frames, height_video, width_video = control_video.shape - control_video = self.image_processor.preprocess( - control_video.permute(0, 2, 1, 3, 4).reshape( - batch_size * num_frames, channels, height_video, width_video - ), - height=height, - width=width, - ) - control_video = control_video.to(dtype=torch.float32) - control_video = control_video.reshape(batch_size, num_frames, channels, height, width).permute( - 0, 2, 1, 3, 4 - ) - control_video_latents = self.prepare_control_latents( - None, - control_video, - batch_size, - height, - width, - dtype, - device, - generator, - self.do_classifier_free_guidance, - )[1] - control_latents = ( - torch.cat([control_video_latents] * 2) if self.do_classifier_free_guidance else control_video_latents - ).to(device, dtype) - else: - control_video_latents = torch.zeros_like(latents).to(device, dtype) - control_latents = ( - torch.cat([control_video_latents] * 2) if self.do_classifier_free_guidance else control_video_latents - ).to(device, dtype) - - if ref_image is not None: - batch_size, channels, num_frames, height_video, width_video = ref_image.shape - ref_image = self.image_processor.preprocess( - ref_image.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height_video, width_video), - height=height, - width=width, - ) - ref_image = ref_image.to(dtype=torch.float32) - ref_image = ref_image.reshape(batch_size, num_frames, channels, height, width).permute(0, 2, 1, 3, 4) - - ref_image_latents = self.prepare_control_latents( - None, - ref_image, - batch_size, - height, - width, - prompt_embeds.dtype, - device, - generator, - self.do_classifier_free_guidance, - )[1] - - ref_image_latents_conv_in = torch.zeros_like(latents) - if latents.size()[2] != 1: - ref_image_latents_conv_in[:, :, :1] = ref_image_latents - ref_image_latents_conv_in = ( - torch.cat([ref_image_latents_conv_in] * 2) - if self.do_classifier_free_guidance - else ref_image_latents_conv_in - ).to(device, dtype) - control_latents = torch.cat([control_latents, ref_image_latents_conv_in], dim=1) - else: - ref_image_latents_conv_in = torch.zeros_like(latents) - ref_image_latents_conv_in = ( - torch.cat([ref_image_latents_conv_in] * 2) - if self.do_classifier_free_guidance - else ref_image_latents_conv_in - ).to(device, dtype) - control_latents = torch.cat([control_latents, ref_image_latents_conv_in], dim=1) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask]) - - # To latents.device - prompt_embeds = prompt_embeds.to(device=device) - prompt_attention_mask = prompt_attention_mask.to(device=device) - - # 7. Denoising loop - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - if hasattr(self.scheduler, "scale_model_input"): - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # expand scalar t to 1-D tensor to match the 1st dim of latent_model_input - t_expand = torch.tensor([t] * latent_model_input.shape[0], device=device).to( - dtype=latent_model_input.dtype - ) - # predict the noise residual - noise_pred = self.transformer( - latent_model_input, - t_expand, - encoder_hidden_states=prompt_embeds, - control_latents=control_latents, - return_dict=False, - )[0] - if noise_pred.size()[1] != self.vae.config.latent_channels: - noise_pred, _ = noise_pred.chunk(2, dim=1) - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - if self.do_classifier_free_guidance and guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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() - - # Convert to tensor - if not output_type == "latent": - video = self.decode_latents(latents) - video = self.video_processor.postprocess_video(video=video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return EasyAnimatePipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/pipeline_easyanimate_inpaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/pipeline_easyanimate_inpaint.py deleted file mode 100644 index b16ef92d8e6b99e9aabb89ba752e2f0645a647ba..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/pipeline_easyanimate_inpaint.py +++ /dev/null @@ -1,1235 +0,0 @@ -# Copyright 2025 The EasyAnimate 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 Callable, Dict, List, Optional, Union - -import numpy as np -import torch -import torch.nn.functional as F -from PIL import Image -from transformers import ( - BertModel, - BertTokenizer, - Qwen2Tokenizer, - Qwen2VLForConditionalGeneration, -) - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import VaeImageProcessor -from ...models import AutoencoderKLMagvit, EasyAnimateTransformer3DModel -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 ...video_processor import VideoProcessor -from .pipeline_output import EasyAnimatePipelineOutput - - -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 EasyAnimateInpaintPipeline - >>> from diffusers.pipelines.easyanimate.pipeline_easyanimate_inpaint import get_image_to_video_latent - >>> from diffusers.utils import export_to_video, load_image - - >>> pipe = EasyAnimateInpaintPipeline.from_pretrained( - ... "alibaba-pai/EasyAnimateV5.1-12b-zh-InP-diffusers", torch_dtype=torch.bfloat16 - ... ) - >>> pipe.to("cuda") - - >>> prompt = "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot." - >>> validation_image_start = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" - ... ) - - >>> validation_image_end = None - >>> sample_size = (448, 576) - >>> num_frames = 49 - >>> input_video, input_video_mask = get_image_to_video_latent( - ... [validation_image_start], validation_image_end, num_frames, sample_size - ... ) - - >>> video = pipe( - ... prompt, - ... num_frames=num_frames, - ... negative_prompt="Twisted body, limb deformities, text subtitles, comics, stillness, ugliness, errors, garbled text.", - ... height=sample_size[0], - ... width=sample_size[1], - ... video=input_video, - ... mask_video=input_video_mask, - ... ) - >>> export_to_video(video.frames[0], "output.mp4", fps=8) - ``` -""" - - -def preprocess_image(image, sample_size): - """ - Preprocess a single image (PIL.Image, numpy.ndarray, or torch.Tensor) to a resized tensor. - """ - if isinstance(image, torch.Tensor): - # If input is a tensor, assume it's in CHW format and resize using interpolation - image = torch.nn.functional.interpolate( - image.unsqueeze(0), size=sample_size, mode="bilinear", align_corners=False - ).squeeze(0) - elif isinstance(image, Image.Image): - # If input is a PIL image, resize and convert to numpy array - image = image.resize((sample_size[1], sample_size[0])) - image = np.array(image) - elif isinstance(image, np.ndarray): - # If input is a numpy array, resize using PIL - image = Image.fromarray(image).resize((sample_size[1], sample_size[0])) - image = np.array(image) - else: - raise ValueError("Unsupported input type. Expected PIL.Image, numpy.ndarray, or torch.Tensor.") - - # Convert to tensor if not already - if not isinstance(image, torch.Tensor): - image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 # HWC -> CHW, normalize to [0, 1] - - return image - - -def get_image_to_video_latent(validation_image_start, validation_image_end, num_frames, sample_size): - """ - Generate latent representations for video from start and end images. Inputs can be PIL.Image, numpy.ndarray, or - torch.Tensor. - """ - input_video = None - input_video_mask = None - - if validation_image_start is not None: - # Preprocess the starting image(s) - if isinstance(validation_image_start, list): - image_start = [preprocess_image(img, sample_size) for img in validation_image_start] - else: - image_start = preprocess_image(validation_image_start, sample_size) - - # Create video tensor from the starting image(s) - if isinstance(image_start, list): - start_video = torch.cat( - [img.unsqueeze(1).unsqueeze(0) for img in image_start], - dim=2, - ) - input_video = torch.tile(start_video[:, :, :1], [1, 1, num_frames, 1, 1]) - input_video[:, :, : len(image_start)] = start_video - else: - input_video = torch.tile( - image_start.unsqueeze(1).unsqueeze(0), - [1, 1, num_frames, 1, 1], - ) - - # Normalize input video (already normalized in preprocess_image) - - # Create mask for the input video - input_video_mask = torch.zeros_like(input_video[:, :1]) - if isinstance(image_start, list): - input_video_mask[:, :, len(image_start) :] = 255 - else: - input_video_mask[:, :, 1:] = 255 - - # Handle ending image(s) if provided - if validation_image_end is not None: - if isinstance(validation_image_end, list): - image_end = [preprocess_image(img, sample_size) for img in validation_image_end] - end_video = torch.cat( - [img.unsqueeze(1).unsqueeze(0) for img in image_end], - dim=2, - ) - input_video[:, :, -len(end_video) :] = end_video - input_video_mask[:, :, -len(image_end) :] = 0 - else: - image_end = preprocess_image(validation_image_end, sample_size) - input_video[:, :, -1:] = image_end.unsqueeze(1).unsqueeze(0) - input_video_mask[:, :, -1:] = 0 - - elif validation_image_start is None: - # If no starting image is provided, initialize empty tensors - input_video = torch.zeros([1, 3, num_frames, sample_size[0], sample_size[1]]) - input_video_mask = torch.ones([1, 1, num_frames, sample_size[0], sample_size[1]]) * 255 - - return input_video, input_video_mask - - -# Similar to diffusers.pipelines.hunyuandit.pipeline_hunyuandit.get_resize_crop_region_for_grid -def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): - tw = tgt_width - th = tgt_height - h, w = src - r = h / w - if r > (th / tw): - resize_height = th - resize_width = int(round(th / h * w)) - else: - resize_width = tw - resize_height = int(round(tw / w * h)) - - crop_top = int(round((th - resize_height) / 2.0)) - crop_left = int(round((tw - resize_width) / 2.0)) - - return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width) - - -# 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 - - -# Resize mask information in magvit -def resize_mask(mask, latent, process_first_frame_only=True): - latent_size = latent.size() - - if process_first_frame_only: - target_size = list(latent_size[2:]) - target_size[0] = 1 - first_frame_resized = F.interpolate( - mask[:, :, 0:1, :, :], size=target_size, mode="trilinear", align_corners=False - ) - - target_size = list(latent_size[2:]) - target_size[0] = target_size[0] - 1 - if target_size[0] != 0: - remaining_frames_resized = F.interpolate( - mask[:, :, 1:, :, :], size=target_size, mode="trilinear", align_corners=False - ) - resized_mask = torch.cat([first_frame_resized, remaining_frames_resized], dim=2) - else: - resized_mask = first_frame_resized - else: - target_size = list(latent_size[2:]) - resized_mask = F.interpolate(mask, size=target_size, mode="trilinear", align_corners=False) - return resized_mask - - -## Add noise to reference video -def add_noise_to_reference_video(image, ratio=None, generator=None): - if ratio is None: - sigma = torch.normal(mean=-3.0, std=0.5, size=(image.shape[0],)).to(image.device) - sigma = torch.exp(sigma).to(image.dtype) - else: - sigma = torch.ones((image.shape[0],)).to(image.device, image.dtype) * ratio - - if generator is not None: - image_noise = ( - torch.randn(image.size(), generator=generator, dtype=image.dtype, device=image.device) - * sigma[:, None, None, None, None] - ) - else: - image_noise = torch.randn_like(image) * sigma[:, None, None, None, None] - image_noise = torch.where(image == -1, torch.zeros_like(image), image_noise) - image = image + image_noise - return 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, - **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 EasyAnimateInpaintPipeline(DiffusionPipeline): - r""" - Pipeline for text-to-video generation using EasyAnimate. - - 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.) - - EasyAnimate uses one text encoder [qwen2 vl](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct) in V5.1. - - 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`]): - EasyAnimate uses [qwen2 vl](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct) in V5.1. - tokenizer (Optional[`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer`]): - A `Qwen2Tokenizer` or `BertTokenizer` to tokenize text. - transformer ([`EasyAnimateTransformer3DModel`]): - The EasyAnimate model designed by EasyAnimate Team. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - A scheduler to be used in combination with EasyAnimate to denoise the encoded image latents. - """ - - model_cpu_offload_seq = "text_encoder->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKLMagvit, - text_encoder: Union[Qwen2VLForConditionalGeneration, BertModel], - tokenizer: Union[Qwen2Tokenizer, BertTokenizer], - transformer: EasyAnimateTransformer3DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - ) - - self.enable_text_attention_mask = ( - self.transformer.config.enable_text_attention_mask - if getattr(self, "transformer", None) is not None - else True - ) - self.vae_spatial_compression_ratio = ( - self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 8 - ) - self.vae_temporal_compression_ratio = ( - self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 4 - ) - self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) - self.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_spatial_compression_ratio, - do_normalize=False, - do_binarize=True, - do_convert_grayscale=True, - ) - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) - - # Copied from diffusers.pipelines.easyanimate.pipeline_easyanimate.EasyAnimatePipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[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, - max_sequence_length: int = 256, - ): - 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 - dtype (`torch.dtype`): - torch dtype - 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. - prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the prompt. Required when `prompt_embeds` is passed directly. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the negative prompt. Required when `negative_prompt_embeds` is passed directly. - max_sequence_length (`int`, *optional*): maximum sequence length to use for the prompt. - """ - dtype = dtype or self.text_encoder.dtype - device = device or self.text_encoder.device - - 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: - if isinstance(prompt, str): - messages = [ - { - "role": "user", - "content": [{"type": "text", "text": prompt}], - } - ] - else: - messages = [ - { - "role": "user", - "content": [{"type": "text", "text": _prompt}], - } - for _prompt in prompt - ] - text = [ - self.tokenizer.apply_chat_template([m], tokenize=False, add_generation_prompt=True) for m in messages - ] - - text_inputs = self.tokenizer( - text=text, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_attention_mask=True, - padding_side="right", - return_tensors="pt", - ) - text_inputs = text_inputs.to(self.text_encoder.device) - - text_input_ids = text_inputs.input_ids - prompt_attention_mask = text_inputs.attention_mask - if self.enable_text_attention_mask: - # Inference: Generation of the output - prompt_embeds = self.text_encoder( - input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True - ).hidden_states[-2] - else: - raise ValueError("LLM needs attention_mask") - prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1) - - prompt_embeds = prompt_embeds.to(dtype=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) - prompt_attention_mask = prompt_attention_mask.to(device=device) - - # get unconditional embeddings for classifier free guidance - if do_classifier_free_guidance and negative_prompt_embeds is None: - if negative_prompt is not None and isinstance(negative_prompt, str): - messages = [ - { - "role": "user", - "content": [{"type": "text", "text": negative_prompt}], - } - ] - else: - messages = [ - { - "role": "user", - "content": [{"type": "text", "text": _negative_prompt}], - } - for _negative_prompt in negative_prompt - ] - text = [ - self.tokenizer.apply_chat_template([m], tokenize=False, add_generation_prompt=True) for m in messages - ] - - text_inputs = self.tokenizer( - text=text, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_attention_mask=True, - padding_side="right", - return_tensors="pt", - ) - text_inputs = text_inputs.to(self.text_encoder.device) - - text_input_ids = text_inputs.input_ids - negative_prompt_attention_mask = text_inputs.attention_mask - if self.enable_text_attention_mask: - # Inference: Generation of the output - negative_prompt_embeds = self.text_encoder( - input_ids=text_input_ids, - attention_mask=negative_prompt_attention_mask, - output_hidden_states=True, - ).hidden_states[-2] - else: - raise ValueError("LLM needs attention_mask") - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 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] - - negative_prompt_embeds = negative_prompt_embeds.to(dtype=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) - negative_prompt_attention_mask = negative_prompt_attention_mask.to(device=device) - - return prompt_embeds, negative_prompt_embeds, prompt_attention_mask, negative_prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - callback_on_step_end_tensor_inputs=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 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 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 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}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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, - do_classifier_free_guidance, - noise_aug_strength, - ): - # 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 - if mask is not None: - mask = mask.to(device=device, dtype=dtype) - new_mask = [] - bs = 1 - for i in range(0, mask.shape[0], bs): - mask_bs = mask[i : i + bs] - mask_bs = self.vae.encode(mask_bs)[0] - mask_bs = mask_bs.mode() - new_mask.append(mask_bs) - mask = torch.cat(new_mask, dim=0) - mask = mask * self.vae.config.scaling_factor - - if masked_image is not None: - masked_image = masked_image.to(device=device, dtype=dtype) - if self.transformer.config.add_noise_in_inpaint_model: - masked_image = add_noise_to_reference_video( - masked_image, ratio=noise_aug_strength, generator=generator - ) - new_mask_pixel_values = [] - bs = 1 - for i in range(0, masked_image.shape[0], bs): - mask_pixel_values_bs = masked_image[i : i + bs] - mask_pixel_values_bs = self.vae.encode(mask_pixel_values_bs)[0] - mask_pixel_values_bs = mask_pixel_values_bs.mode() - new_mask_pixel_values.append(mask_pixel_values_bs) - masked_image_latents = torch.cat(new_mask_pixel_values, dim=0) - masked_image_latents = masked_image_latents * self.vae.config.scaling_factor - - # 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) - else: - masked_image_latents = None - - return mask, masked_image_latents - - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - num_frames, - dtype, - device, - generator, - latents=None, - video=None, - timestep=None, - is_strength_max=True, - return_noise=False, - return_video_latents=False, - ): - shape = ( - batch_size, - num_channels_latents, - (num_frames - 1) // self.vae_temporal_compression_ratio + 1, - height // self.vae_spatial_compression_ratio, - width // self.vae_spatial_compression_ratio, - ) - - 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 return_video_latents or (latents is None and not is_strength_max): - video = video.to(device=device, dtype=dtype) - bs = 1 - new_video = [] - for i in range(0, video.shape[0], bs): - video_bs = video[i : i + bs] - video_bs = self.vae.encode(video_bs)[0] - video_bs = video_bs.sample() - new_video.append(video_bs) - video = torch.cat(new_video, dim=0) - video = video * self.vae.config.scaling_factor - - video_latents = video.repeat(batch_size // video.shape[0], 1, 1, 1, 1) - video_latents = video_latents.to(device=device, dtype=dtype) - - if latents is None: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # if strength is 1. then initialise the latents to noise, else initial to image + noise - if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler): - latents = noise if is_strength_max else self.scheduler.scale_noise(video_latents, timestep, noise) - else: - latents = noise if is_strength_max else self.scheduler.add_noise(video_latents, noise, timestep) - # if pure noise then scale the initial latents by the Scheduler's init sigma - if hasattr(self.scheduler, "init_noise_sigma"): - latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents - else: - if hasattr(self.scheduler, "init_noise_sigma"): - noise = latents.to(device) - latents = noise * self.scheduler.init_noise_sigma - else: - latents = latents.to(device) - - # scale the initial noise by the standard deviation required by the scheduler - outputs = (latents,) - - if return_noise: - outputs += (noise,) - - if return_video_latents: - outputs += (video_latents,) - - return outputs - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - # 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 - - @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, - 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", - 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"], - guidance_rescale: float = 0.0, - strength: float = 1.0, - noise_aug_strength: float = 0.0563, - timesteps: Optional[List[int]] = None, - ): - r""" - The call function to the pipeline for generation with HunyuanDiT. - - Examples: - 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 - continuity of generated content. - video (`torch.FloatTensor`, *optional*): - A tensor representing an input video, which can be modified depending on the prompts provided. - mask_video (`torch.FloatTensor`, *optional*): - A tensor to specify areas of the video to be masked (omitted from generation). - masked_video_latents (`torch.FloatTensor`, *optional*): - Latents from masked portions of the video, utilized during image generation. - height (`int`, *optional*): - The height in pixels of the generated image or video frames. - width (`int`, *optional*): - The width in pixels of the generated image or video frames. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image but slower - inference time. This parameter is modulated by `strength`. - 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*): - 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): - The number of images to generate per prompt. - eta (`float`, *optional*, defaults to 0.0): - 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*): - 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*): - A pre-computed latent representation which can be used to guide the generation process. - prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not - provided, embeddings are generated from the `prompt` input argument. - negative_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated negative text embeddings, aiding in fine-tuning what should not be represented in the - outputs. If not provided, embeddings are generated from the `negative_prompt` argument. - prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask guiding the focus of the model on specific parts of the prompt text. Required when using - `prompt_embeds`. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the negative prompt, needed when `negative_prompt_embeds` are used. - output_type (`str`, *optional*, defaults to `"latent"`): - The output format of the generated image. Choose between `PIL.Image` and `np.array` to define how you - want the results to be formatted. - 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`, - *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*): - 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): - Rescale parameter for adjusting noise configuration based on guidance rescale. Based on findings from - [Common Diffusion Noise Schedules and Sample Steps are - Flawed](https://huggingface.co/papers/2305.08891). - strength (`float`, *optional*, defaults to 1.0): - Affects the overall styling or quality of the generated output. Values closer to 1 usually provide - direct adherence to prompts. - - Examples: - # Example usage of the function for generating images based on prompts. - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - Returns either a structured output containing generated images and their metadata when `return_dict` is - `True`, or a simpler tuple, where the first element is a list of generated images and the second - element indicates if any of them contain "not-safe-for-work" (NSFW) content. - """ - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - # 0. default height and width - height = int(height // 16 * 16) - width = int(width // 16 * 16) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - callback_on_step_end_tensor_inputs, - ) - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - 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 - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - else: - dtype = self.transformer.dtype - - # 3. Encode input prompt - ( - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt=prompt, - device=device, - dtype=dtype, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - 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, - ) - - # 4. set timesteps - if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler): - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, mu=1 - ) - else: - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) - timesteps, num_inference_steps = self.get_timesteps( - num_inference_steps=num_inference_steps, strength=strength, device=device - ) - - # at which timestep to set the initial noise (n.b. 50% if strength is 0.5) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise - is_strength_max = strength == 1.0 - - if video is not None: - batch_size, channels, num_frames, height_video, width_video = video.shape - init_video = self.image_processor.preprocess( - video.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height_video, width_video), - height=height, - width=width, - ) - init_video = init_video.to(dtype=torch.float32) - init_video = init_video.reshape(batch_size, num_frames, channels, height, width).permute(0, 2, 1, 3, 4) - else: - init_video = None - - # Prepare latent variables - num_channels_latents = self.vae.config.latent_channels - num_channels_transformer = self.transformer.config.in_channels - return_image_latents = num_channels_transformer == num_channels_latents - - # 5. Prepare latents. - latents_outputs = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - num_frames, - dtype, - device, - generator, - latents, - video=init_video, - timestep=latent_timestep, - is_strength_max=is_strength_max, - return_noise=True, - return_video_latents=return_image_latents, - ) - if return_image_latents: - latents, noise, image_latents = latents_outputs - else: - latents, noise = latents_outputs - - # 6. Prepare inpaint latents if it needs. - if mask_video is not None: - if (mask_video == 255).all(): - mask = torch.zeros_like(latents).to(device, dtype) - # Use zero latents if we want to t2v. - if self.transformer.config.resize_inpaint_mask_directly: - mask_latents = torch.zeros_like(latents)[:, :1].to(device, dtype) - else: - mask_latents = torch.zeros_like(latents).to(device, dtype) - masked_video_latents = torch.zeros_like(latents).to(device, dtype) - - mask_input = torch.cat([mask_latents] * 2) if self.do_classifier_free_guidance else mask_latents - masked_video_latents_input = ( - torch.cat([masked_video_latents] * 2) if self.do_classifier_free_guidance else masked_video_latents - ) - inpaint_latents = torch.cat([mask_input, masked_video_latents_input], dim=1).to(dtype) - else: - # Prepare mask latent variables - batch_size, channels, num_frames, height_video, width_video = mask_video.shape - mask_condition = self.mask_processor.preprocess( - mask_video.permute(0, 2, 1, 3, 4).reshape( - batch_size * num_frames, channels, height_video, width_video - ), - height=height, - width=width, - ) - mask_condition = mask_condition.to(dtype=torch.float32) - mask_condition = mask_condition.reshape(batch_size, num_frames, channels, height, width).permute( - 0, 2, 1, 3, 4 - ) - - if num_channels_transformer != num_channels_latents: - mask_condition_tile = torch.tile(mask_condition, [1, 3, 1, 1, 1]) - if masked_video_latents is None: - masked_video = ( - init_video * (mask_condition_tile < 0.5) - + torch.ones_like(init_video) * (mask_condition_tile > 0.5) * -1 - ) - else: - masked_video = masked_video_latents - - if self.transformer.config.resize_inpaint_mask_directly: - _, masked_video_latents = self.prepare_mask_latents( - None, - masked_video, - batch_size, - height, - width, - dtype, - device, - generator, - self.do_classifier_free_guidance, - noise_aug_strength=noise_aug_strength, - ) - mask_latents = resize_mask( - 1 - mask_condition, masked_video_latents, self.vae.config.cache_mag_vae - ) - mask_latents = mask_latents.to(device, dtype) * self.vae.config.scaling_factor - else: - mask_latents, masked_video_latents = self.prepare_mask_latents( - mask_condition_tile, - masked_video, - batch_size, - height, - width, - dtype, - device, - generator, - self.do_classifier_free_guidance, - noise_aug_strength=noise_aug_strength, - ) - - mask_input = torch.cat([mask_latents] * 2) if self.do_classifier_free_guidance else mask_latents - masked_video_latents_input = ( - torch.cat([masked_video_latents] * 2) - if self.do_classifier_free_guidance - else masked_video_latents - ) - inpaint_latents = torch.cat([mask_input, masked_video_latents_input], dim=1).to(dtype) - else: - inpaint_latents = None - - mask = torch.tile(mask_condition, [1, num_channels_latents, 1, 1, 1]) - mask = F.interpolate(mask, size=latents.size()[-3:], mode="trilinear", align_corners=True).to( - device, dtype - ) - else: - if num_channels_transformer != num_channels_latents: - mask = torch.zeros_like(latents).to(device, dtype) - if self.transformer.config.resize_inpaint_mask_directly: - mask_latents = torch.zeros_like(latents)[:, :1].to(device, dtype) - else: - mask_latents = torch.zeros_like(latents).to(device, dtype) - masked_video_latents = torch.zeros_like(latents).to(device, dtype) - - mask_input = torch.cat([mask_latents] * 2) if self.do_classifier_free_guidance else mask_latents - masked_video_latents_input = ( - torch.cat([masked_video_latents] * 2) if self.do_classifier_free_guidance else masked_video_latents - ) - inpaint_latents = torch.cat([mask_input, masked_video_latents_input], dim=1).to(dtype) - else: - mask = torch.zeros_like(init_video[:, :1]) - mask = torch.tile(mask, [1, num_channels_latents, 1, 1, 1]) - mask = F.interpolate(mask, size=latents.size()[-3:], mode="trilinear", align_corners=True).to( - device, dtype - ) - - inpaint_latents = None - - # Check that sizes of mask, masked image and latents match - if num_channels_transformer != num_channels_latents: - num_channels_mask = mask_latents.shape[1] - num_channels_masked_image = masked_video_latents.shape[1] - if ( - num_channels_latents + num_channels_mask + num_channels_masked_image - != self.transformer.config.in_channels - ): - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.transformer`: {self.transformer.config} expects" - f" {self.transformer.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" - " `pipeline.transformer` or your `mask_image` or `image` input." - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask]) - - # To latents.device - prompt_embeds = prompt_embeds.to(device=device) - prompt_attention_mask = prompt_attention_mask.to(device=device) - - # 8. Denoising loop - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - if hasattr(self.scheduler, "scale_model_input"): - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # expand scalar t to 1-D tensor to match the 1st dim of latent_model_input - t_expand = torch.tensor([t] * latent_model_input.shape[0], device=device).to( - dtype=latent_model_input.dtype - ) - - # predict the noise residual - noise_pred = self.transformer( - latent_model_input, - t_expand, - encoder_hidden_states=prompt_embeds, - inpaint_latents=inpaint_latents, - return_dict=False, - )[0] - if noise_pred.size()[1] != self.vae.config.latent_channels: - noise_pred, _ = noise_pred.chunk(2, dim=1) - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - if self.do_classifier_free_guidance and guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - if num_channels_transformer == num_channels_latents: - init_latents_proper = image_latents - init_mask = mask - if i < len(timesteps) - 1: - noise_timestep = timesteps[i + 1] - if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler): - init_latents_proper = self.scheduler.scale_noise( - init_latents_proper, torch.tensor([noise_timestep], noise) - ) - else: - init_latents_proper = self.scheduler.add_noise( - init_latents_proper, noise, torch.tensor([noise_timestep]) - ) - - latents = (1 - init_mask) * init_latents_proper + init_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) - - 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 not output_type == "latent": - latents = 1 / self.vae.config.scaling_factor * latents - video = self.vae.decode(latents, return_dict=False)[0] - video = self.video_processor.postprocess_video(video=video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return EasyAnimatePipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/pipeline_output.py deleted file mode 100644 index c761a3b1079fcf7160aa05ff1f4084035ebae7f0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/easyanimate/pipeline_output.py +++ /dev/null @@ -1,20 +0,0 @@ -from dataclasses import dataclass - -import torch - -from diffusers.utils import BaseOutput - - -@dataclass -class EasyAnimatePipelineOutput(BaseOutput): - r""" - 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 - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__init__.py deleted file mode 100644 index 72e1b578f2cab9d6afcff6f9b64f51de867e4fba..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__init__.py +++ /dev/null @@ -1,69 +0,0 @@ -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 = {} -_additional_imports = {} -_import_structure = {"pipeline_output": ["FluxPipelineOutput", "FluxPriorReduxPipelineOutput"]} - -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["modeling_flux"] = ["ReduxImageEncoder"] - _import_structure["pipeline_flux"] = ["FluxPipeline"] - _import_structure["pipeline_flux_control"] = ["FluxControlPipeline"] - _import_structure["pipeline_flux_control_img2img"] = ["FluxControlImg2ImgPipeline"] - _import_structure["pipeline_flux_control_inpaint"] = ["FluxControlInpaintPipeline"] - _import_structure["pipeline_flux_controlnet"] = ["FluxControlNetPipeline"] - _import_structure["pipeline_flux_controlnet_image_to_image"] = ["FluxControlNetImg2ImgPipeline"] - _import_structure["pipeline_flux_controlnet_inpainting"] = ["FluxControlNetInpaintPipeline"] - _import_structure["pipeline_flux_fill"] = ["FluxFillPipeline"] - _import_structure["pipeline_flux_img2img"] = ["FluxImg2ImgPipeline"] - _import_structure["pipeline_flux_inpaint"] = ["FluxInpaintPipeline"] - _import_structure["pipeline_flux_prior_redux"] = ["FluxPriorReduxPipeline"] -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 .modeling_flux import ReduxImageEncoder - from .pipeline_flux import FluxPipeline - from .pipeline_flux_control import FluxControlPipeline - from .pipeline_flux_control_img2img import FluxControlImg2ImgPipeline - from .pipeline_flux_control_inpaint import FluxControlInpaintPipeline - from .pipeline_flux_controlnet import FluxControlNetPipeline - from .pipeline_flux_controlnet_image_to_image import FluxControlNetImg2ImgPipeline - from .pipeline_flux_controlnet_inpainting import FluxControlNetInpaintPipeline - from .pipeline_flux_fill import FluxFillPipeline - from .pipeline_flux_img2img import FluxImg2ImgPipeline - from .pipeline_flux_inpaint import FluxInpaintPipeline - from .pipeline_flux_prior_redux import FluxPriorReduxPipeline -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/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 367f409a2f9469b1c674803cb4ae31a6afe09c06..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/modeling_flux.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/modeling_flux.cpython-310.pyc deleted file mode 100644 index de54cdd484309805b282c6f89405e30b72166c97..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/modeling_flux.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux.cpython-310.pyc deleted file mode 100644 index 8a9fcb6643d82f1fb5737169f804d3a48109d7ca..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_control.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_control.cpython-310.pyc deleted file mode 100644 index 0d23795ba1af31e5f3547ce72a36a81c3410b20a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_control.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_control_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_control_img2img.cpython-310.pyc deleted file mode 100644 index 1fdb4201258b01af204504939431e1c8a70f45b3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_control_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_control_inpaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_control_inpaint.cpython-310.pyc deleted file mode 100644 index aa5ffe3a2dd4425f681d2f1540848900318e0495..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_control_inpaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_controlnet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_controlnet.cpython-310.pyc deleted file mode 100644 index 4568a988af1d389ee7fa0958a451a62b615ed173..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_controlnet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_controlnet_image_to_image.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_controlnet_image_to_image.cpython-310.pyc deleted file mode 100644 index af363959aabb82fc0d9088248e78a204bb036b9a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_controlnet_image_to_image.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_controlnet_inpainting.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_controlnet_inpainting.cpython-310.pyc deleted file mode 100644 index 1ef42d22e4e12c77e41eb37705c3c9f47d015842..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_controlnet_inpainting.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_fill.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_fill.cpython-310.pyc deleted file mode 100644 index 8bd20b6ec2e5fd8975822f36e7844e0fe2958472..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_fill.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_img2img.cpython-310.pyc deleted file mode 100644 index 8d9830c48d20c23f8b6eb29f194fab5320dfc9de..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_inpaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_inpaint.cpython-310.pyc deleted file mode 100644 index 9dfe261604af2384d82d1dff050ebafc7076d187..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_inpaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_prior_redux.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_prior_redux.cpython-310.pyc deleted file mode 100644 index 89e581509b8600ee810d016eb9f6a85f81d4a668..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_flux_prior_redux.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index 9025a3ee639e0b9cf0488ed1e47d11db95609b37..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/modeling_flux.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/modeling_flux.py deleted file mode 100644 index d7f2f45812b3b9af0fbf4521a8e8507f942a7eac..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/modeling_flux.py +++ /dev/null @@ -1,47 +0,0 @@ -# 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. - - -from dataclasses import dataclass -from typing import Optional - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.modeling_utils import ModelMixin -from ...utils import BaseOutput - - -@dataclass -class ReduxImageEncoderOutput(BaseOutput): - image_embeds: Optional[torch.Tensor] = None - - -class ReduxImageEncoder(ModelMixin, ConfigMixin): - @register_to_config - def __init__( - self, - redux_dim: int = 1152, - txt_in_features: int = 4096, - ) -> None: - super().__init__() - - self.redux_up = nn.Linear(redux_dim, txt_in_features * 3) - self.redux_down = nn.Linear(txt_in_features * 3, txt_in_features) - - def forward(self, x: torch.Tensor) -> ReduxImageEncoderOutput: - projected_x = self.redux_down(nn.functional.silu(self.redux_up(x))) - - return ReduxImageEncoderOutput(image_embeds=projected_x) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux.py deleted file mode 100644 index bdee2ead483c8c7c21b0732c63d3e9fdd02727fe..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux.py +++ /dev/null @@ -1,982 +0,0 @@ -# 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, - CLIPTextModel, - CLIPTokenizer, - CLIPVisionModelWithProjection, - T5EncoderModel, - T5TokenizerFast, -) - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FluxIPAdapterMixin, FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, FluxTransformer2DModel -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 FluxPipelineOutput - - -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 FluxPipeline - - >>> pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", 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=4, guidance_scale=0.0).images[0] - >>> image.save("flux.png") - ``` -""" - - -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 FluxPipeline( - DiffusionPipeline, - FluxLoraLoaderMixin, - FromSingleFileMixin, - TextualInversionLoaderMixin, - FluxIPAdapterMixin, -): - r""" - The Flux pipeline for text-to-image generation. - - Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ - - Args: - transformer ([`FluxTransformer2DModel`]): - 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 ([`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: CLIPTextModel, - tokenizer: CLIPTokenizer, - text_encoder_2: T5EncoderModel, - tokenizer_2: T5TokenizerFast, - transformer: FluxTransformer2DModel, - image_encoder: CLIPVisionModelWithProjection = None, - feature_extractor: CLIPImageProcessor = None, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - 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 - # 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.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = 128 - - 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_2) - - text_inputs = self.tokenizer_2( - 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 - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, self.tokenizer_max_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_2(text_input_ids.to(device), output_hidden_states=False)[0] - - dtype = self.text_encoder_2.dtype - prompt_embeds = prompt_embeds.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) - - return prompt_embeds - - def _get_clip_prompt_embeds( - self, - prompt: Union[str, List[str]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - ): - device = device or self._execution_device - - 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=self.tokenizer_max_length, - truncation=True, - return_overflowing_tokens=False, - return_length=False, - 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_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds.pooler_output - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds - - def encode_prompt( - self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - max_sequence_length: int = 512, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - scale_lora_layers(self.text_encoder_2, lora_scale) - - prompt = [prompt] if isinstance(prompt, str) else prompt - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - # We only use the pooled prompt output from the CLIPTextModel - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - ) - prompt_embeds = self._get_t5_prompt_embeds( - prompt=prompt_2, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - 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) - - if self.text_encoder_2 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_2, lora_scale) - - 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) - - return prompt_embeds, pooled_prompt_embeds, text_ids - - 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 - - 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, - 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, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - ): - 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_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 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`." - ) - - 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) - - @staticmethod - 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 - 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 enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=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_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 // 2, width // 2, device, dtype) - - return latents, latent_image_ids - - @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 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, - prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt: Union[str, List[str]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - true_cfg_scale: float = 1.0, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 28, - sigmas: Optional[List[float]] = 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", - 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. - 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*): - 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*): - 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): - When > 1.0 and a provided `negative_prompt`, enables true classifier-free guidance. - 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. - 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.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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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.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. - negative_pooled_prompt_embeds (`torch.FloatTensor`, *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.flux.FluxPipelineOutput`] 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.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] 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_2, - height, - width, - 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, - 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 - ) - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None - ) - do_true_cfg = true_cfg_scale > 1 and has_neg_prompt - ( - prompt_embeds, - pooled_prompt_embeds, - text_ids, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - device=device, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - if do_true_cfg: - ( - negative_prompt_embeds, - negative_pooled_prompt_embeds, - negative_text_ids, - ) = self.encode_prompt( - prompt=negative_prompt, - prompt_2=negative_prompt_2, - prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=negative_pooled_prompt_embeds, - 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 // 4 - latents, latent_image_ids = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 5. Prepare timesteps - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas - image_seq_len = latents.shape[1] - 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, - ) - 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) - 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 - 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 - 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, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - if do_true_cfg: - 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, - guidance=guidance, - pooled_projections=negative_pooled_prompt_embeds, - encoder_hidden_states=negative_prompt_embeds, - txt_ids=negative_text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - noise_pred = neg_noise_pred + true_cfg_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 - - 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 FluxPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_control.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_control.py deleted file mode 100644 index b4f77cf0194ebeb833d1d69b2b3a5cab81afc0de..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_control.py +++ /dev/null @@ -1,886 +0,0 @@ -# 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 CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin -from ...models.autoencoders import AutoencoderKL -from ...models.transformers import FluxTransformer2DModel -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 FluxPipelineOutput - - -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 controlnet_aux import CannyDetector - >>> from diffusers import FluxControlPipeline - >>> from diffusers.utils import load_image - - >>> pipe = FluxControlPipeline.from_pretrained( - ... "black-forest-labs/FLUX.1-Canny-dev", torch_dtype=torch.bfloat16 - ... ).to("cuda") - - >>> prompt = "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts." - >>> control_image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/robot.png" - ... ) - - >>> processor = CannyDetector() - >>> control_image = processor( - ... control_image, low_threshold=50, high_threshold=200, detect_resolution=1024, image_resolution=1024 - ... ) - - >>> image = pipe( - ... prompt=prompt, - ... control_image=control_image, - ... height=1024, - ... width=1024, - ... num_inference_steps=50, - ... guidance_scale=30.0, - ... ).images[0] - >>> image.save("output.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 FluxControlPipeline( - DiffusionPipeline, - FluxLoraLoaderMixin, - FromSingleFileMixin, - TextualInversionLoaderMixin, -): - r""" - The Flux pipeline for controllable text-to-image generation. - - Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ - - Args: - transformer ([`FluxTransformer2DModel`]): - 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 ([`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->transformer->vae" - _optional_components = [] - _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__( - self, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - text_encoder_2: T5EncoderModel, - tokenizer_2: T5TokenizerFast, - transformer: FluxTransformer2DModel, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - 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.vae_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, vae_latent_channels=self.vae_latent_channels - ) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = 128 - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds - 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_2) - - text_inputs = self.tokenizer_2( - 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 - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, self.tokenizer_max_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_2(text_input_ids.to(device), output_hidden_states=False)[0] - - dtype = self.text_encoder_2.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - ): - device = device or self._execution_device - - 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=self.tokenizer_max_length, - truncation=True, - return_overflowing_tokens=False, - return_length=False, - 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_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds.pooler_output - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - max_sequence_length: int = 512, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - scale_lora_layers(self.text_encoder_2, lora_scale) - - prompt = [prompt] if isinstance(prompt, str) else prompt - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - # We only use the pooled prompt output from the CLIPTextModel - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - ) - prompt_embeds = self._get_t5_prompt_embeds( - prompt=prompt_2, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - 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) - - if self.text_encoder_2 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_2, lora_scale) - - 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) - - return prompt_embeds, pooled_prompt_embeds, text_ids - - def check_inputs( - self, - prompt, - prompt_2, - height, - width, - prompt_embeds=None, - pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - ): - 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_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 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 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 enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.prepare_latents - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=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_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 // 2, width // 2, device, dtype) - - return latents, latent_image_ids - - # Copied from diffusers.pipelines.controlnet_sd3.pipeline_stable_diffusion_3_controlnet.StableDiffusion3ControlNetPipeline.prepare_image - def prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - if isinstance(image, torch.Tensor): - pass - else: - image = self.image_processor.preprocess(image, height=height, width=width) - - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - @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 - - @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, - control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 28, - sigmas: Optional[List[float]] = 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", - 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. - 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]]`): - 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. - 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. - 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.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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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.FluxPipelineOutput`] 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.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] 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_2, - height, - width, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - 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._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 - lora_scale = ( - self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None - ) - ( - prompt_embeds, - pooled_prompt_embeds, - text_ids, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - 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 // 8 - - control_image = self.prepare_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=self.vae.dtype, - ) - - if control_image.ndim == 4: - control_image = self.vae.encode(control_image).latent_dist.sample(generator=generator) - control_image = (control_image - self.vae.config.shift_factor) * self.vae.config.scaling_factor - - height_control_image, width_control_image = control_image.shape[2:] - control_image = self._pack_latents( - control_image, - batch_size * num_images_per_prompt, - num_channels_latents, - height_control_image, - width_control_image, - ) - - latents, latent_image_ids = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 5. Prepare timesteps - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas - image_seq_len = latents.shape[1] - 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, - ) - 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 - - # 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 = torch.cat([latents, control_image], dim=2) - - # 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=latent_model_input, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - # 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() - - 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 FluxPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_control_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_control_img2img.py deleted file mode 100644 index 8d5439daf607b36a5a76b7b3a4143946c268e8c3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_control_img2img.py +++ /dev/null @@ -1,941 +0,0 @@ -# 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 CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin -from ...models.autoencoders import AutoencoderKL -from ...models.transformers import FluxTransformer2DModel -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 FluxPipelineOutput - - -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 controlnet_aux import CannyDetector - >>> from diffusers import FluxControlImg2ImgPipeline - >>> from diffusers.utils import load_image - - >>> pipe = FluxControlImg2ImgPipeline.from_pretrained( - ... "black-forest-labs/FLUX.1-Canny-dev", torch_dtype=torch.bfloat16 - ... ).to("cuda") - - >>> prompt = "A robot made of exotic candies and chocolates of different kinds. Abstract background" - >>> image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/watercolor-painting.jpg" - ... ) - >>> control_image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/robot.png" - ... ) - - >>> processor = CannyDetector() - >>> control_image = processor( - ... control_image, low_threshold=50, high_threshold=200, detect_resolution=1024, image_resolution=1024 - ... ) - - >>> image = pipe( - ... prompt=prompt, - ... image=image, - ... control_image=control_image, - ... strength=0.8, - ... height=1024, - ... width=1024, - ... num_inference_steps=50, - ... guidance_scale=30.0, - ... ).images[0] - >>> image.save("output.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: Optional[torch.Generator] = 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 FluxControlImg2ImgPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin): - r""" - The Flux pipeline for image inpainting. - - Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ - - Args: - transformer ([`FluxTransformer2DModel`]): - 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 ([`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->transformer->vae" - _optional_components = [] - _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__( - self, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - text_encoder_2: T5EncoderModel, - tokenizer_2: T5TokenizerFast, - transformer: FluxTransformer2DModel, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - transformer=transformer, - scheduler=scheduler, - ) - 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 = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = 128 - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds - 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_2) - - text_inputs = self.tokenizer_2( - 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 - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, self.tokenizer_max_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_2(text_input_ids.to(device), output_hidden_states=False)[0] - - dtype = self.text_encoder_2.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - ): - device = device or self._execution_device - - 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=self.tokenizer_max_length, - truncation=True, - return_overflowing_tokens=False, - return_length=False, - 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_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds.pooler_output - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - max_sequence_length: int = 512, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - scale_lora_layers(self.text_encoder_2, lora_scale) - - prompt = [prompt] if isinstance(prompt, str) else prompt - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - # We only use the pooled prompt output from the CLIPTextModel - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - ) - prompt_embeds = self._get_t5_prompt_embeds( - prompt=prompt_2, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - 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) - - if self.text_encoder_2 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_2, lora_scale) - - 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) - - return prompt_embeds, pooled_prompt_embeds, text_ids - - # 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, - prompt_2, - strength, - height, - width, - prompt_embeds=None, - pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=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_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 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 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) - - if latents is not None: - return latents.to(device=device, dtype=dtype), latent_image_ids - - image = image.to(device=device, dtype=dtype) - image_latents = self._encode_vae_image(image=image, generator=generator) - 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) - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = self.scheduler.scale_noise(image_latents, timestep, noise) - latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) - return latents, latent_image_ids - - # Copied from diffusers.pipelines.controlnet_sd3.pipeline_stable_diffusion_3_controlnet.StableDiffusion3ControlNetPipeline.prepare_image - def prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - if isinstance(image, torch.Tensor): - pass - else: - image = self.image_processor.preprocess(image, height=height, width=width) - - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - @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 - - @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, - image: PipelineImageInput = None, - control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - strength: float = 0.6, - num_inference_steps: int = 28, - sigmas: Optional[List[float]] = 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", - 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. - 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`, 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]]`): - 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. - 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. - 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`. - 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 7.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_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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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.FluxPipelineOutput`] 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.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] 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_2, - strength, - height, - width, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - 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._interrupt = False - - # 2. Preprocess image - init_image = self.image_processor.preprocess(image, height=height, width=width) - 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 - - # 3. Prepare text embeddings - lora_scale = ( - self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None - ) - ( - prompt_embeds, - pooled_prompt_embeds, - text_ids, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - 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) - - 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 // 8 - - control_image = self.prepare_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=self.vae.dtype, - ) - - if control_image.ndim == 4: - control_image = self.vae.encode(control_image).latent_dist.sample(generator=generator) - control_image = (control_image - self.vae.config.shift_factor) * self.vae.config.scaling_factor - - height_control_image, width_control_image = control_image.shape[2:] - control_image = self._pack_latents( - control_image, - batch_size * num_images_per_prompt, - num_channels_latents, - height_control_image, - width_control_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, - ) - - 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 - - # 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 = torch.cat([latents, control_image], dim=2) - - # 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=latent_model_input, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - # 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() - - 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 FluxPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py deleted file mode 100644 index 872bcf177c02a2627e0737c7ce6f86a960bc0f77..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py +++ /dev/null @@ -1,1139 +0,0 @@ -# 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 ( - CLIPTextModel, - CLIPTokenizer, - T5EncoderModel, - T5TokenizerFast, -) - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FluxLoraLoaderMixin, - FromSingleFileMixin, - TextualInversionLoaderMixin, -) -from ...models.autoencoders import AutoencoderKL -from ...models.transformers import FluxTransformer2DModel -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 FluxPipelineOutput - - -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 FluxControlInpaintPipeline - from diffusers.models.transformers import FluxTransformer2DModel - from transformers import T5EncoderModel - from diffusers.utils import load_image, make_image_grid - from image_gen_aux import DepthPreprocessor # https://github.com/huggingface/image_gen_aux - from PIL import Image - import numpy as np - - pipe = FluxControlInpaintPipeline.from_pretrained( - "black-forest-labs/FLUX.1-Depth-dev", - torch_dtype=torch.bfloat16, - ) - # use following lines if you have GPU constraints - # --------------------------------------------------------------- - transformer = FluxTransformer2DModel.from_pretrained( - "sayakpaul/FLUX.1-Depth-dev-nf4", subfolder="transformer", torch_dtype=torch.bfloat16 - ) - text_encoder_2 = T5EncoderModel.from_pretrained( - "sayakpaul/FLUX.1-Depth-dev-nf4", subfolder="text_encoder_2", torch_dtype=torch.bfloat16 - ) - pipe.transformer = transformer - pipe.text_encoder_2 = text_encoder_2 - pipe.enable_model_cpu_offload() - # --------------------------------------------------------------- - pipe.to("cuda") - - prompt = "a blue robot singing opera with human-like expressions" - image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/robot.png") - - head_mask = np.zeros_like(image) - head_mask[65:580, 300:642] = 255 - mask_image = Image.fromarray(head_mask) - - processor = DepthPreprocessor.from_pretrained("LiheYoung/depth-anything-large-hf") - control_image = processor(image)[0].convert("RGB") - - output = pipe( - prompt=prompt, - image=image, - control_image=control_image, - mask_image=mask_image, - num_inference_steps=30, - strength=0.9, - guidance_scale=10.0, - generator=torch.Generator().manual_seed(42), - ).images[0] - make_image_grid([image, control_image, mask_image, output.resize(image.size)], rows=1, cols=4).save( - "output.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: Optional[torch.Generator] = 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 FluxControlInpaintPipeline( - DiffusionPipeline, - FluxLoraLoaderMixin, - FromSingleFileMixin, - TextualInversionLoaderMixin, -): - r""" - The Flux pipeline for image inpainting using Flux-dev-Depth/Canny. - - Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ - - Args: - transformer ([`FluxTransformer2DModel`]): - 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 ([`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->transformer->vae" - _optional_components = [] - _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__( - self, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - text_encoder_2: T5EncoderModel, - tokenizer_2: T5TokenizerFast, - transformer: FluxTransformer2DModel, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - transformer=transformer, - scheduler=scheduler, - ) - 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 = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) - latent_channels = self.vae.config.latent_channels if getattr(self, "vae", None) else 16 - self.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor * 2, - vae_latent_channels=latent_channels, - do_normalize=False, - do_binarize=True, - do_convert_grayscale=True, - ) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = 128 - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds - 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_2) - - text_inputs = self.tokenizer_2( - 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 - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, self.tokenizer_max_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_2(text_input_ids.to(device), output_hidden_states=False)[0] - - dtype = self.text_encoder_2.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - ): - device = device or self._execution_device - - 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=self.tokenizer_max_length, - truncation=True, - return_overflowing_tokens=False, - return_length=False, - 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_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds.pooler_output - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - max_sequence_length: int = 512, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - scale_lora_layers(self.text_encoder_2, lora_scale) - - prompt = [prompt] if isinstance(prompt, str) else prompt - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - # We only use the pooled prompt output from the CLIPTextModel - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - ) - prompt_embeds = self._get_t5_prompt_embeds( - prompt=prompt_2, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - 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) - - if self.text_encoder_2 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_2, lora_scale) - - 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) - - return prompt_embeds, pooled_prompt_embeds, text_ids - - # 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, - prompt_2, - strength, - height, - width, - prompt_embeds=None, - pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=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_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 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 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 enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - 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) - - if latents is not None: - return latents.to(device=device, dtype=dtype), latent_image_ids - - image = image.to(device=device, dtype=dtype) - image_latents = self._encode_vae_image(image=image, generator=generator) - 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) - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = self.scheduler.scale_noise(image_latents, timestep, noise) - latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) - return latents, noise, image_latents, latent_image_ids - - # Copied from diffusers.pipelines.controlnet_sd3.pipeline_stable_diffusion_3_controlnet.StableDiffusion3ControlNetPipeline.prepare_image - def prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - if isinstance(image, torch.Tensor): - pass - else: - image = self.image_processor.preprocess(image, height=height, width=width) - - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - def prepare_mask_latents( - self, - image, - mask_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. - image = self.image_processor.preprocess(image, height=height, width=width) - mask_image = self.mask_processor.preprocess(mask_image, height=height, width=width) - - masked_image = image * (1 - mask_image) - masked_image = masked_image.to(device=device, dtype=dtype) - - 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_image = torch.nn.functional.interpolate(mask_image, size=(height, width)) - mask_image = mask_image.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] == num_channels_latents: - 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_image.shape[0] < batch_size: - if not batch_size % mask_image.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_image.shape[0]} mask_image were passed. Make sure the number" - " of masks that you pass is divisible by the total requested batch size." - ) - mask_image = mask_image.repeat(batch_size // mask_image.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_image = self._pack_latents( - mask_image.repeat(1, num_channels_latents, 1, 1), - batch_size, - num_channels_latents, - height, - width, - ) - masked_image_latents = torch.cat((masked_image_latents, mask_image), dim=-1) - - return mask_image, 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 - - @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, - image: PipelineImageInput = None, - control_image: PipelineImageInput = None, - mask_image: PipelineImageInput = None, - masked_image_latents: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - strength: float = 0.6, - num_inference_steps: int = 28, - sigmas: Optional[List[float]] = 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", - 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. - 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`, 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]]`): - 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]`): - `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]`): - `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): - 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. - 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`. - 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 7.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_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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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.FluxPipelineOutput`] 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.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] 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_2, - strength, - height, - width, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - 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._interrupt = False - device = self._execution_device - - # 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 - - # 3. Prepare text embeddings - lora_scale = ( - self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None - ) - ( - prompt_embeds, - pooled_prompt_embeds, - text_ids, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - device=device, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - - # 3. Preprocess mask and image - num_channels_latents = self.vae.config.latent_channels - if masked_image_latents is not None: - # pre computed masked_image_latents and mask_image - masked_image_latents = masked_image_latents.to(latents.device) - mask = mask_image.to(latents.device) - else: - mask, masked_image_latents = self.prepare_mask_latents( - image, - mask_image, - batch_size, - num_channels_latents, - num_images_per_prompt, - height, - width, - prompt_embeds.dtype, - device, - generator, - ) - - init_image = self.image_processor.preprocess(image, height=height, width=width) - init_image = init_image.to(dtype=torch.float32) - - # 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) - - 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 // 8 - - control_image = self.prepare_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=self.vae.dtype, - ) - - if control_image.ndim == 4: - control_image = self.vae.encode(control_image).latent_dist.sample(generator=generator) - control_image = (control_image - self.vae.config.shift_factor) * self.vae.config.scaling_factor - - height_control_image, width_control_image = control_image.shape[2:] - control_image = self._pack_latents( - control_image, - batch_size * num_images_per_prompt, - num_channels_latents, - height_control_image, - width_control_image, - ) - - 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, - ) - - # 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_8 = 2 * (int(height) // (self.vae_scale_factor * 2)) - width_8 = 2 * (int(width) // (self.vae_scale_factor * 2)) - - 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 - - # 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 = torch.cat([latents, control_image], dim=2) - - # 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=latent_model_input, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - # 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_mask = mask - if i < len(timesteps) - 1: - noise_timestep = timesteps[i + 1] - init_latents_proper = self.scheduler.scale_noise( - image_latents, torch.tensor([noise_timestep]), noise - ) - else: - init_latents_proper = image_latents - init_latents_proper = self._pack_latents( - init_latents_proper, batch_size * num_images_per_prompt, num_channels_latents, height_8, width_8 - ) - - 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() - - 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 FluxPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_controlnet.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_controlnet.py deleted file mode 100644 index 1438d4a902fe6d38b850af385cec2ff7a4e59403..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_controlnet.py +++ /dev/null @@ -1,1178 +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. - -import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTokenizer, - CLIPVisionModelWithProjection, - T5EncoderModel, - T5TokenizerFast, -) - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FluxIPAdapterMixin, FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin -from ...models.autoencoders import AutoencoderKL -from ...models.controlnets.controlnet_flux import FluxControlNetModel, FluxMultiControlNetModel -from ...models.transformers import FluxTransformer2DModel -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 FluxPipelineOutput - - -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.utils import load_image - >>> from diffusers import FluxControlNetPipeline - >>> from diffusers import FluxControlNetModel - - >>> base_model = "black-forest-labs/FLUX.1-dev" - >>> controlnet_model = "InstantX/FLUX.1-dev-controlnet-canny" - >>> controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16) - >>> pipe = FluxControlNetPipeline.from_pretrained( - ... base_model, controlnet=controlnet, torch_dtype=torch.bfloat16 - ... ) - >>> pipe.to("cuda") - >>> control_image = load_image("https://huggingface.co/InstantX/SD3-Controlnet-Canny/resolve/main/canny.jpg") - >>> prompt = "A girl in city, 25 years old, cool, futuristic" - >>> image = pipe( - ... prompt, - ... control_image=control_image, - ... control_guidance_start=0.2, - ... control_guidance_end=0.8, - ... controlnet_conditioning_scale=1.0, - ... num_inference_steps=28, - ... guidance_scale=3.5, - ... ).images[0] - >>> image.save("flux.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: Optional[torch.Generator] = 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 FluxControlNetPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin, FluxIPAdapterMixin): - r""" - The Flux pipeline for text-to-image generation. - - Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ - - Args: - transformer ([`FluxTransformer2DModel`]): - 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 ([`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", "control_image"] - - def __init__( - self, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - text_encoder_2: T5EncoderModel, - tokenizer_2: T5TokenizerFast, - transformer: FluxTransformer2DModel, - controlnet: Union[ - FluxControlNetModel, List[FluxControlNetModel], Tuple[FluxControlNetModel], FluxMultiControlNetModel - ], - image_encoder: CLIPVisionModelWithProjection = None, - feature_extractor: CLIPImageProcessor = None, - ): - super().__init__() - if isinstance(controlnet, (list, tuple)): - controlnet = FluxMultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - transformer=transformer, - scheduler=scheduler, - controlnet=controlnet, - 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 - # 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.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = 128 - - 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_2( - 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 - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, self.tokenizer_max_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_2(text_input_ids.to(device), output_hidden_states=False)[0] - - dtype = self.text_encoder_2.dtype - prompt_embeds = prompt_embeds.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) - - return prompt_embeds - - def _get_clip_prompt_embeds( - self, - prompt: Union[str, List[str]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - ): - device = device or self._execution_device - - 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=self.tokenizer_max_length, - truncation=True, - return_overflowing_tokens=False, - return_length=False, - 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_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds.pooler_output - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds - - def encode_prompt( - self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - max_sequence_length: int = 512, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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. - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - scale_lora_layers(self.text_encoder_2, lora_scale) - - prompt = [prompt] if isinstance(prompt, str) else prompt - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - # We only use the pooled prompt output from the CLIPTextModel - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - ) - prompt_embeds = self._get_t5_prompt_embeds( - prompt=prompt_2, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - 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) - - if self.text_encoder_2 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_2, lora_scale) - - 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) - - return prompt_embeds, pooled_prompt_embeds, text_ids - - # 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, - 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, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - ): - 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_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`." - ) - - 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 - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.prepare_latents - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=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_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 // 2, width // 2, device, dtype) - - return latents, latent_image_ids - - # Copied from diffusers.pipelines.controlnet_sd3.pipeline_stable_diffusion_3_controlnet.StableDiffusion3ControlNetPipeline.prepare_image - def prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - if isinstance(image, torch.Tensor): - pass - else: - image = self.image_processor.preprocess(image, height=height, width=width) - - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - @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 - - @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, - negative_prompt: Union[str, List[str]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - true_cfg_scale: float = 1.0, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 28, - sigmas: Optional[List[float]] = 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_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", - 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. - 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): - 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. - guidance_scale (`float`, *optional*, defaults to 7.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. - 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): - 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]]`): - 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): - 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): - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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. - 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.FluxPipelineOutput`] 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.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] 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 - - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(self.controlnet.nets) if isinstance(self.controlnet, FluxMultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - height, - width, - 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, - 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._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 - dtype = self.transformer.dtype - - # 3. Prepare text embeddings - lora_scale = ( - self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None - ) - do_true_cfg = true_cfg_scale > 1 and negative_prompt is not None - ( - prompt_embeds, - pooled_prompt_embeds, - text_ids, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - device=device, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - if do_true_cfg: - ( - negative_prompt_embeds, - negative_pooled_prompt_embeds, - _, - ) = self.encode_prompt( - prompt=negative_prompt, - prompt_2=negative_prompt_2, - prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=negative_pooled_prompt_embeds, - device=device, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - - # 3. Prepare control image - num_channels_latents = self.transformer.config.in_channels // 4 - if isinstance(self.controlnet, FluxControlNetModel): - control_image = self.prepare_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=self.vae.dtype, - ) - height, width = control_image.shape[-2:] - - # xlab controlnet has a input_hint_block and instantx controlnet does not - controlnet_blocks_repeat = False if self.controlnet.input_hint_block is None else True - if self.controlnet.input_hint_block is None: - # vae encode - control_image = retrieve_latents(self.vae.encode(control_image), generator=generator) - control_image = (control_image - self.vae.config.shift_factor) * self.vae.config.scaling_factor - - # pack - height_control_image, width_control_image = control_image.shape[2:] - control_image = self._pack_latents( - control_image, - batch_size * num_images_per_prompt, - num_channels_latents, - height_control_image, - width_control_image, - ) - - # Here we ensure that `control_mode` has the same length as the control_image. - if control_mode is not None: - if not isinstance(control_mode, int): - raise ValueError(" For `FluxControlNet`, `control_mode` should be an `int` or `None`") - control_mode = torch.tensor(control_mode).to(device, dtype=torch.long) - control_mode = control_mode.view(-1, 1).expand(control_image.shape[0], 1) - - elif isinstance(self.controlnet, FluxMultiControlNetModel): - control_images = [] - # xlab controlnet has a input_hint_block and instantx controlnet does not - controlnet_blocks_repeat = False if self.controlnet.nets[0].input_hint_block is None else True - for i, control_image_ in enumerate(control_image): - control_image_ = self.prepare_image( - image=control_image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=self.vae.dtype, - ) - height, width = control_image_.shape[-2:] - - if self.controlnet.nets[0].input_hint_block is None: - # vae encode - control_image_ = retrieve_latents(self.vae.encode(control_image_), generator=generator) - control_image_ = (control_image_ - self.vae.config.shift_factor) * self.vae.config.scaling_factor - - # pack - height_control_image, width_control_image = control_image_.shape[2:] - control_image_ = self._pack_latents( - control_image_, - batch_size * num_images_per_prompt, - num_channels_latents, - height_control_image, - width_control_image, - ) - control_images.append(control_image_) - - control_image = control_images - - # Here we ensure that `control_mode` has the same length as the control_image. - if isinstance(control_mode, list) and len(control_mode) != len(control_image): - raise ValueError( - "For Multi-ControlNet, `control_mode` must be a list of the same " - + " length as the number of controlnets (control images) specified" - ) - if not isinstance(control_mode, list): - control_mode = [control_mode] * len(control_image) - # set control mode - control_modes = [] - for cmode in control_mode: - if cmode is None: - cmode = -1 - control_mode = torch.tensor(cmode).expand(control_images[0].shape[0]).to(device, dtype=torch.long) - control_modes.append(control_mode) - control_mode = control_modes - - # 4. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels // 4 - latents, latent_image_ids = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 5. Prepare timesteps - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas - image_seq_len = latents.shape[1] - 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, - ) - - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - # 6. Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(self.controlnet, FluxControlNetModel) else keeps) - - 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, - ) - - # 7. Denoising loop - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - 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) - - if isinstance(self.controlnet, FluxMultiControlNetModel): - use_guidance = self.controlnet.nets[0].config.guidance_embeds - else: - use_guidance = self.controlnet.config.guidance_embeds - - guidance = torch.tensor([guidance_scale], device=device) if use_guidance else None - guidance = guidance.expand(latents.shape[0]) if guidance is not None else None - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - # controlnet - controlnet_block_samples, controlnet_single_block_samples = self.controlnet( - hidden_states=latents, - controlnet_cond=control_image, - controlnet_mode=control_mode, - conditioning_scale=cond_scale, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - ) - - guidance = ( - torch.tensor([guidance_scale], device=device) if self.transformer.config.guidance_embeds else None - ) - guidance = guidance.expand(latents.shape[0]) if guidance is not None else None - - noise_pred = self.transformer( - hidden_states=latents, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - controlnet_block_samples=controlnet_block_samples, - controlnet_single_block_samples=controlnet_single_block_samples, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - controlnet_blocks_repeat=controlnet_blocks_repeat, - )[0] - - if do_true_cfg: - 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, - guidance=guidance, - pooled_projections=negative_pooled_prompt_embeds, - encoder_hidden_states=negative_prompt_embeds, - controlnet_block_samples=controlnet_block_samples, - controlnet_single_block_samples=controlnet_single_block_samples, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - controlnet_blocks_repeat=controlnet_blocks_repeat, - )[0] - noise_pred = neg_noise_pred + true_cfg_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) - control_image = callback_outputs.pop("control_image", control_image) - - # 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 = 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 FluxPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py deleted file mode 100644 index 52e15de53b0d1cfc233576663508e181e1215473..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py +++ /dev/null @@ -1,1004 +0,0 @@ -import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import torch -from transformers import ( - CLIPTextModel, - CLIPTokenizer, - T5EncoderModel, - T5TokenizerFast, -) - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin -from ...models.autoencoders import AutoencoderKL -from ...models.controlnets.controlnet_flux import FluxControlNetModel, FluxMultiControlNetModel -from ...models.transformers import FluxTransformer2DModel -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 FluxPipelineOutput - - -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 FluxControlNetImg2ImgPipeline, FluxControlNetModel - >>> from diffusers.utils import load_image - - >>> device = "cuda" if torch.cuda.is_available() else "cpu" - - >>> controlnet = FluxControlNetModel.from_pretrained( - ... "InstantX/FLUX.1-dev-Controlnet-Canny-alpha", torch_dtype=torch.bfloat16 - ... ) - - >>> pipe = FluxControlNetImg2ImgPipeline.from_pretrained( - ... "black-forest-labs/FLUX.1-schnell", controlnet=controlnet, torch_dtype=torch.float16 - ... ) - - >>> pipe.text_encoder.to(torch.float16) - >>> pipe.controlnet.to(torch.float16) - >>> pipe.to("cuda") - - >>> control_image = load_image("https://huggingface.co/InstantX/SD3-Controlnet-Canny/resolve/main/canny.jpg") - >>> init_image = load_image( - ... "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" - ... ) - - >>> prompt = "A girl in city, 25 years old, cool, futuristic" - >>> image = pipe( - ... prompt, - ... image=init_image, - ... control_image=control_image, - ... control_guidance_start=0.2, - ... control_guidance_end=0.8, - ... controlnet_conditioning_scale=1.0, - ... strength=0.7, - ... num_inference_steps=2, - ... guidance_scale=3.5, - ... ).images[0] - >>> image.save("flux_controlnet_img2img.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: Optional[torch.Generator] = 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 FluxControlNetImg2ImgPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin): - r""" - The Flux controlnet pipeline for image-to-image generation. - - Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ - - Args: - transformer ([`FluxTransformer2DModel`]): - 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 ([`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->transformer->vae" - _optional_components = [] - _callback_tensor_inputs = ["latents", "prompt_embeds", "control_image"] - - def __init__( - self, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - text_encoder_2: T5EncoderModel, - tokenizer_2: T5TokenizerFast, - transformer: FluxTransformer2DModel, - controlnet: Union[ - FluxControlNetModel, List[FluxControlNetModel], Tuple[FluxControlNetModel], FluxMultiControlNetModel - ], - ): - super().__init__() - if isinstance(controlnet, (list, tuple)): - controlnet = FluxMultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - transformer=transformer, - scheduler=scheduler, - controlnet=controlnet, - ) - 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 = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = 128 - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds - 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_2) - - text_inputs = self.tokenizer_2( - 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 - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, self.tokenizer_max_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_2(text_input_ids.to(device), output_hidden_states=False)[0] - - dtype = self.text_encoder_2.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - ): - device = device or self._execution_device - - 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=self.tokenizer_max_length, - truncation=True, - return_overflowing_tokens=False, - return_length=False, - 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_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds.pooler_output - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - max_sequence_length: int = 512, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - scale_lora_layers(self.text_encoder_2, lora_scale) - - prompt = [prompt] if isinstance(prompt, str) else prompt - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - # We only use the pooled prompt output from the CLIPTextModel - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - ) - prompt_embeds = self._get_t5_prompt_embeds( - prompt=prompt_2, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - 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) - - if self.text_encoder_2 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_2, lora_scale) - - 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) - - return prompt_embeds, pooled_prompt_embeds, text_ids - - # 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, - prompt_2, - strength, - height, - width, - callback_on_step_end_tensor_inputs, - prompt_embeds=None, - pooled_prompt_embeds=None, - max_sequence_length=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_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 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 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) - - if latents is not None: - return latents.to(device=device, dtype=dtype), latent_image_ids - - image = image.to(device=device, dtype=dtype) - image_latents = self._encode_vae_image(image=image, generator=generator) - 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) - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = self.scheduler.scale_noise(image_latents, timestep, noise) - latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) - return latents, latent_image_ids - - # Copied from diffusers.pipelines.controlnet_sd3.pipeline_stable_diffusion_3_controlnet.StableDiffusion3ControlNetPipeline.prepare_image - def prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - if isinstance(image, torch.Tensor): - pass - else: - image = self.image_processor.preprocess(image, height=height, width=width) - - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - @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 - - @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, - image: PipelineImageInput = None, - control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - strength: float = 0.6, - num_inference_steps: int = 28, - sigmas: Optional[List[float]] = 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", - 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, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `List[str]`, *optional*): - The prompt or prompts to guide the image generation. - 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`): - The image(s) to modify with the pipeline. - 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. - width (`int`, *optional*, defaults to self.default_sample_size * self.vae_scale_factor): - The width in pixels of the generated image. - strength (`float`, *optional*, defaults to 0.6): - Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. - 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*): - 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*): - 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): - 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*): - One or more [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. - prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. - pooled_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated pooled text embeddings. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple. - joint_attention_kwargs (`dict`, *optional*): - 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*): - 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. - - Examples: - - Returns: - [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] 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 - - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(self.controlnet.nets) if isinstance(self.controlnet, FluxMultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - self.check_inputs( - prompt, - prompt_2, - strength, - height, - width, - callback_on_step_end_tensor_inputs, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - max_sequence_length=max_sequence_length, - ) - - self._guidance_scale = guidance_scale - self._joint_attention_kwargs = joint_attention_kwargs - 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 - dtype = self.transformer.dtype - - lora_scale = ( - self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None - ) - ( - prompt_embeds, - pooled_prompt_embeds, - text_ids, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - device=device, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - - init_image = self.image_processor.preprocess(image, height=height, width=width) - init_image = init_image.to(dtype=torch.float32) - - num_channels_latents = self.transformer.config.in_channels // 4 - - if isinstance(self.controlnet, FluxControlNetModel): - control_image = self.prepare_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=self.vae.dtype, - ) - height, width = control_image.shape[-2:] - - # xlab controlnet has a input_hint_block and instantx controlnet does not - controlnet_blocks_repeat = False if self.controlnet.input_hint_block is None else True - if self.controlnet.input_hint_block is None: - control_image = retrieve_latents(self.vae.encode(control_image), generator=generator) - control_image = (control_image - self.vae.config.shift_factor) * self.vae.config.scaling_factor - - height_control_image, width_control_image = control_image.shape[2:] - control_image = self._pack_latents( - control_image, - batch_size * num_images_per_prompt, - num_channels_latents, - height_control_image, - width_control_image, - ) - - if control_mode is not None: - control_mode = torch.tensor(control_mode).to(device, dtype=torch.long) - control_mode = control_mode.reshape([-1, 1]) - - elif isinstance(self.controlnet, FluxMultiControlNetModel): - control_images = [] - - # xlab controlnet has a input_hint_block and instantx controlnet does not - controlnet_blocks_repeat = False if self.controlnet.nets[0].input_hint_block is None else True - for i, control_image_ in enumerate(control_image): - control_image_ = self.prepare_image( - image=control_image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=self.vae.dtype, - ) - height, width = control_image_.shape[-2:] - - if self.controlnet.nets[0].input_hint_block is None: - control_image_ = retrieve_latents(self.vae.encode(control_image_), generator=generator) - control_image_ = (control_image_ - self.vae.config.shift_factor) * self.vae.config.scaling_factor - - height_control_image, width_control_image = control_image_.shape[2:] - control_image_ = self._pack_latents( - control_image_, - batch_size * num_images_per_prompt, - num_channels_latents, - height_control_image, - width_control_image, - ) - - control_images.append(control_image_) - - control_image = control_images - - control_mode_ = [] - if isinstance(control_mode, list): - for cmode in control_mode: - if cmode is None: - control_mode_.append(-1) - else: - control_mode_.append(cmode) - control_mode = torch.tensor(control_mode_).to(device, dtype=torch.long) - control_mode = control_mode.reshape([-1, 1]) - - 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) - - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - 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, - ) - - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(self.controlnet, FluxControlNetModel) else keeps) - - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - 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 - - timestep = t.expand(latents.shape[0]).to(latents.dtype) - - if isinstance(self.controlnet, FluxMultiControlNetModel): - use_guidance = self.controlnet.nets[0].config.guidance_embeds - else: - use_guidance = self.controlnet.config.guidance_embeds - - guidance = torch.tensor([guidance_scale], device=device) if use_guidance else None - guidance = guidance.expand(latents.shape[0]) if guidance is not None else None - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - controlnet_block_samples, controlnet_single_block_samples = self.controlnet( - hidden_states=latents, - controlnet_cond=control_image, - controlnet_mode=control_mode, - conditioning_scale=cond_scale, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - ) - - guidance = ( - torch.tensor([guidance_scale], device=device) if self.transformer.config.guidance_embeds else None - ) - guidance = guidance.expand(latents.shape[0]) if guidance is not None else None - - noise_pred = self.transformer( - hidden_states=latents, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - controlnet_block_samples=controlnet_block_samples, - controlnet_single_block_samples=controlnet_single_block_samples, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - controlnet_blocks_repeat=controlnet_blocks_repeat, - )[0] - - 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(): - 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) - control_image = callback_outputs.pop("control_image", control_image) - - 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 = 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) - - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return FluxPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py deleted file mode 100644 index d1e874d0b88f21c0144025dd0a0b9943b95b21a1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py +++ /dev/null @@ -1,1205 +0,0 @@ -import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL -import torch -from transformers import ( - CLIPTextModel, - CLIPTokenizer, - T5EncoderModel, - T5TokenizerFast, -) - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin -from ...models.autoencoders import AutoencoderKL -from ...models.controlnets.controlnet_flux import FluxControlNetModel, FluxMultiControlNetModel -from ...models.transformers import FluxTransformer2DModel -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 FluxPipelineOutput - - -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__) - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import FluxControlNetInpaintPipeline - >>> from diffusers.models import FluxControlNetModel - >>> from diffusers.utils import load_image - - >>> controlnet = FluxControlNetModel.from_pretrained( - ... "InstantX/FLUX.1-dev-controlnet-canny", torch_dtype=torch.float16 - ... ) - >>> pipe = FluxControlNetInpaintPipeline.from_pretrained( - ... "black-forest-labs/FLUX.1-schnell", controlnet=controlnet, torch_dtype=torch.float16 - ... ) - >>> pipe.to("cuda") - - >>> control_image = load_image( - ... "https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Canny-alpha/resolve/main/canny.jpg" - ... ) - >>> init_image = load_image( - ... "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png" - ... ) - >>> mask_image = load_image( - ... "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png" - ... ) - - >>> prompt = "A girl holding a sign that says InstantX" - >>> image = pipe( - ... prompt, - ... image=init_image, - ... mask_image=mask_image, - ... control_image=control_image, - ... control_guidance_start=0.2, - ... control_guidance_end=0.8, - ... controlnet_conditioning_scale=0.7, - ... strength=0.7, - ... num_inference_steps=28, - ... guidance_scale=3.5, - ... ).images[0] - >>> image.save("flux_controlnet_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: Optional[torch.Generator] = 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 FluxControlNetInpaintPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin): - r""" - The Flux controlnet pipeline for inpainting. - - Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ - - Args: - transformer ([`FluxTransformer2DModel`]): - 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 ([`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->transformer->vae" - _optional_components = [] - _callback_tensor_inputs = ["latents", "prompt_embeds", "control_image", "mask", "masked_image_latents"] - - def __init__( - self, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - text_encoder_2: T5EncoderModel, - tokenizer_2: T5TokenizerFast, - transformer: FluxTransformer2DModel, - controlnet: Union[ - FluxControlNetModel, List[FluxControlNetModel], Tuple[FluxControlNetModel], FluxMultiControlNetModel - ], - ): - super().__init__() - if isinstance(controlnet, (list, tuple)): - controlnet = FluxMultiControlNetModel(controlnet) - - self.register_modules( - scheduler=scheduler, - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - text_encoder_2=text_encoder_2, - tokenizer_2=tokenizer_2, - transformer=transformer, - controlnet=controlnet, - ) - - 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 = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) - latent_channels = self.vae.config.latent_channels if getattr(self, "vae", None) else 16 - self.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor * 2, - vae_latent_channels=latent_channels, - do_normalize=False, - do_binarize=True, - do_convert_grayscale=True, - ) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = 128 - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds - 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_2) - - text_inputs = self.tokenizer_2( - 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 - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, self.tokenizer_max_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_2(text_input_ids.to(device), output_hidden_states=False)[0] - - dtype = self.text_encoder_2.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - ): - device = device or self._execution_device - - 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=self.tokenizer_max_length, - truncation=True, - return_overflowing_tokens=False, - return_length=False, - 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_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds.pooler_output - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - max_sequence_length: int = 512, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - scale_lora_layers(self.text_encoder_2, lora_scale) - - prompt = [prompt] if isinstance(prompt, str) else prompt - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - # We only use the pooled prompt output from the CLIPTextModel - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - ) - prompt_embeds = self._get_t5_prompt_embeds( - prompt=prompt_2, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - 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) - - if self.text_encoder_2 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_2, lora_scale) - - 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) - - return prompt_embeds, pooled_prompt_embeds, text_ids - - # 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, - prompt_2, - image, - mask_image, - strength, - height, - width, - output_type, - prompt_embeds=None, - pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - padding_mask_crop=None, - max_sequence_length=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_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 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 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) - image_latents = self._encode_vae_image(image=image, generator=generator) - - 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 - - # Copied from diffusers.pipelines.controlnet_sd3.pipeline_stable_diffusion_3_controlnet.StableDiffusion3ControlNetPipeline.prepare_image - def prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - if isinstance(image, torch.Tensor): - pass - else: - image = self.image_processor.preprocess(image, height=height, width=width) - - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - @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 - - @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, - image: PipelineImageInput = None, - mask_image: PipelineImageInput = None, - masked_image_latents: PipelineImageInput = None, - control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - strength: float = 0.6, - padding_mask_crop: Optional[int] = None, - sigmas: Optional[List[float]] = 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", - 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, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `List[str]`, *optional*): - The prompt or prompts to guide the image generation. - 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`): - The image(s) to inpaint. - 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`): - 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. - width (`int`, *optional*, defaults to self.default_sample_size * self.vae_scale_factor): - The width in pixels of the generated image. - strength (`float`, *optional*, defaults to 0.6): - Conceptually, indicates how much to inpaint the masked area. Must be between 0 and 1. - padding_mask_crop (`int`, *optional*): - The size of the padding to use when cropping the mask. - 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*): - 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): - The percentage of total steps at which the ControlNet starts applying. - 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*): - 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): - 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*): - One or more [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. - prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. - pooled_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated pooled text embeddings. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple. - joint_attention_kwargs (`dict`, *optional*): - 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*): - 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. - - Examples: - - Returns: - [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] 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 - - global_height = height - global_width = width - - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(self.controlnet.nets) if isinstance(self.controlnet, FluxMultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 1. Check inputs - self.check_inputs( - prompt, - prompt_2, - image, - mask_image, - strength, - height, - width, - output_type=output_type, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - padding_mask_crop=padding_mask_crop, - max_sequence_length=max_sequence_length, - ) - - 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 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 - dtype = self.transformer.dtype - - # 3. Encode input prompt - lora_scale = ( - self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None - ) - prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - device=device, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - - # 4. Preprocess mask and image - if padding_mask_crop is not None: - crops_coords = self.mask_processor.get_crop_region( - mask_image, global_width, global_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=global_height, width=global_width, crops_coords=crops_coords, resize_mode=resize_mode - ) - init_image = init_image.to(dtype=torch.float32) - - # 5. Prepare control image - num_channels_latents = self.transformer.config.in_channels // 4 - if isinstance(self.controlnet, FluxControlNetModel): - control_image = self.prepare_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=self.vae.dtype, - ) - height, width = control_image.shape[-2:] - - # xlab controlnet has a input_hint_block and instantx controlnet does not - controlnet_blocks_repeat = False if self.controlnet.input_hint_block is None else True - if self.controlnet.input_hint_block is None: - # vae encode - control_image = retrieve_latents(self.vae.encode(control_image), generator=generator) - control_image = (control_image - self.vae.config.shift_factor) * self.vae.config.scaling_factor - - # pack - height_control_image, width_control_image = control_image.shape[2:] - control_image = self._pack_latents( - control_image, - batch_size * num_images_per_prompt, - num_channels_latents, - height_control_image, - width_control_image, - ) - - # set control mode - if control_mode is not None: - control_mode = torch.tensor(control_mode).to(device, dtype=torch.long) - control_mode = control_mode.reshape([-1, 1]) - - elif isinstance(self.controlnet, FluxMultiControlNetModel): - control_images = [] - - # xlab controlnet has a input_hint_block and instantx controlnet does not - controlnet_blocks_repeat = False if self.controlnet.nets[0].input_hint_block is None else True - for i, control_image_ in enumerate(control_image): - control_image_ = self.prepare_image( - image=control_image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=self.vae.dtype, - ) - height, width = control_image_.shape[-2:] - - if self.controlnet.nets[0].input_hint_block is None: - # vae encode - control_image_ = retrieve_latents(self.vae.encode(control_image_), generator=generator) - control_image_ = (control_image_ - self.vae.config.shift_factor) * self.vae.config.scaling_factor - - # pack - height_control_image, width_control_image = control_image_.shape[2:] - control_image_ = self._pack_latents( - control_image_, - batch_size * num_images_per_prompt, - num_channels_latents, - height_control_image, - width_control_image, - ) - - control_images.append(control_image_) - - control_image = control_images - - # set control mode - control_mode_ = [] - if isinstance(control_mode, list): - for cmode in control_mode: - if cmode is None: - control_mode_.append(-1) - else: - control_mode_.append(cmode) - control_mode = torch.tensor(control_mode_).to(device, dtype=torch.long) - control_mode = control_mode.reshape([-1, 1]) - - # 6. Prepare timesteps - - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas - image_seq_len = (int(global_height) // self.vae_scale_factor // 2) * ( - int(global_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) - - 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) - - # 7. Prepare latent variables - - latents, noise, image_latents, latent_image_ids = self.prepare_latents( - init_image, - latent_timestep, - batch_size * num_images_per_prompt, - num_channels_latents, - global_height, - global_width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 8. Prepare mask latents - mask_condition = self.mask_processor.preprocess( - mask_image, height=global_height, width=global_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, - global_height, - global_width, - prompt_embeds.dtype, - device, - generator, - ) - - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(self.controlnet, FluxControlNetModel) else keeps) - - # 9. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - 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 - - timestep = t.expand(latents.shape[0]).to(latents.dtype) - - # predict the noise residual - if isinstance(self.controlnet, FluxMultiControlNetModel): - use_guidance = self.controlnet.nets[0].config.guidance_embeds - else: - use_guidance = self.controlnet.config.guidance_embeds - if use_guidance: - guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32) - guidance = guidance.expand(latents.shape[0]) - else: - guidance = None - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - controlnet_block_samples, controlnet_single_block_samples = self.controlnet( - hidden_states=latents, - controlnet_cond=control_image, - controlnet_mode=control_mode, - conditioning_scale=cond_scale, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - ) - - 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 - - noise_pred = self.transformer( - hidden_states=latents, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - controlnet_block_samples=controlnet_block_samples, - controlnet_single_block_samples=controlnet_single_block_samples, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - controlnet_blocks_repeat=controlnet_blocks_repeat, - )[0] - - # 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 inpainting, we need to apply the mask and add the masked image latents - 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) - - # 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) - control_image = callback_outputs.pop("control_image", control_image) - mask = callback_outputs.pop("mask", mask) - masked_image_latents = callback_outputs.pop("masked_image_latents", masked_image_latents) - - 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() - - # Post-processing - if output_type == "latent": - image = latents - else: - latents = self._unpack_latents(latents, global_height, global_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 FluxPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_fill.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_fill.py deleted file mode 100644 index ddfb284eafd6a5c180025c20ff8317c8c9a8f64f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_fill.py +++ /dev/null @@ -1,1037 +0,0 @@ -# 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 CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast - -from ...image_processor import VaeImageProcessor -from ...loaders import FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin -from ...models.autoencoders import AutoencoderKL -from ...models.transformers import FluxTransformer2DModel -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 FluxPipelineOutput - - -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 FluxFillPipeline - >>> from diffusers.utils import load_image - - >>> image = load_image("https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/cup.png") - >>> mask = load_image("https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/cup_mask.png") - - >>> pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16) - >>> pipe.enable_model_cpu_offload() # save some VRAM by offloading the model to CPU - - >>> image = pipe( - ... prompt="a white paper cup", - ... image=image, - ... mask_image=mask, - ... height=1632, - ... width=1232, - ... guidance_scale=30, - ... num_inference_steps=50, - ... max_sequence_length=512, - ... generator=torch.Generator("cpu").manual_seed(0), - ... ).images[0] - >>> image.save("flux_fill.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 - - -# 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" -): - 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 FluxFillPipeline( - DiffusionPipeline, - FluxLoraLoaderMixin, - FromSingleFileMixin, - TextualInversionLoaderMixin, -): - r""" - The Flux Fill pipeline for image inpainting/outpainting. - - Reference: https://blackforestlabs.ai/flux-1-tools/ - - Args: - transformer ([`FluxTransformer2DModel`]): - 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 ([`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->transformer->vae" - _optional_components = [] - _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__( - self, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - text_encoder_2: T5EncoderModel, - tokenizer_2: T5TokenizerFast, - transformer: FluxTransformer2DModel, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - transformer=transformer, - scheduler=scheduler, - ) - 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.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, vae_latent_channels=self.latent_channels - ) - 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, - ) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = 128 - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds - 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_2) - - text_inputs = self.tokenizer_2( - 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 - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, self.tokenizer_max_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_2(text_input_ids.to(device), output_hidden_states=False)[0] - - dtype = self.text_encoder_2.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - ): - device = device or self._execution_device - - 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=self.tokenizer_max_length, - truncation=True, - return_overflowing_tokens=False, - return_length=False, - 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_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds.pooler_output - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds - - def prepare_mask_latents( - self, - mask, - masked_image, - batch_size, - num_channels_latents, - num_images_per_prompt, - height, - width, - dtype, - device, - generator, - ): - # 1. calculate the height and width of the latents - # 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)) - - # 2. encode the masked image - if masked_image.shape[1] == num_channels_latents: - 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 - masked_image_latents = masked_image_latents.to(device=device, dtype=dtype) - - # 3. duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method - batch_size = batch_size * num_images_per_prompt - 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) - - # 4. pack the masked_image_latents - # batch_size, num_channels_latents, height, width -> batch_size, height//2 * width//2 , num_channels_latents*4 - masked_image_latents = self._pack_latents( - masked_image_latents, - batch_size, - num_channels_latents, - height, - width, - ) - - # 5.resize mask to latents shape we we concatenate the mask to the latents - mask = mask[:, 0, :, :] # batch_size, 8 * height, 8 * width (mask has not been 8x compressed) - mask = mask.view( - batch_size, height, self.vae_scale_factor, width, self.vae_scale_factor - ) # batch_size, height, 8, width, 8 - mask = mask.permute(0, 2, 4, 1, 3) # batch_size, 8, 8, height, width - mask = mask.reshape( - batch_size, self.vae_scale_factor * self.vae_scale_factor, height, width - ) # batch_size, 8*8, height, width - - # 6. pack the mask: - # batch_size, 64, height, width -> batch_size, height//2 * width//2 , 64*2*2 - mask = self._pack_latents( - mask, - batch_size, - self.vae_scale_factor * self.vae_scale_factor, - height, - width, - ) - mask = mask.to(device=device, dtype=dtype) - - return mask, masked_image_latents - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - max_sequence_length: int = 512, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - scale_lora_layers(self.text_encoder_2, lora_scale) - - prompt = [prompt] if isinstance(prompt, str) else prompt - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - # We only use the pooled prompt output from the CLIPTextModel - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - ) - prompt_embeds = self._get_t5_prompt_embeds( - prompt=prompt_2, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - 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) - - if self.text_encoder_2 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_2, lora_scale) - - 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) - - return prompt_embeds, pooled_prompt_embeds, text_ids - - # 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, - prompt_2, - strength, - height, - width, - prompt_embeds=None, - pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - image=None, - mask_image=None, - masked_image_latents=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_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 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 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}") - - if image is not None and masked_image_latents is not None: - raise ValueError( - "Please provide either `image` or `masked_image_latents`, `masked_image_latents` should not be passed." - ) - - if image is not None and mask_image is None: - raise ValueError("Please provide `mask_image` when passing `image`.") - - @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 enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - # Copied from diffusers.pipelines.flux.pipeline_flux_img2img.FluxImg2ImgPipeline.prepare_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) - - if latents is not None: - return latents.to(device=device, dtype=dtype), latent_image_ids - - 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) - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = self.scheduler.scale_noise(image_latents, timestep, noise) - latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) - return latents, latent_image_ids - - @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 - - @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, - 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, - strength: float = 1.0, - num_inference_steps: int = 50, - sigmas: Optional[List[float]] = 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", - 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. - 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`, 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]`): - `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]`): - `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): - 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. - 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`. - 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 30.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_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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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.FluxPipelineOutput`] 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.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] 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_2, - strength, - height, - width, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - image=image, - mask_image=mask_image, - masked_image_latents=masked_image_latents, - ) - - self._guidance_scale = guidance_scale - self._joint_attention_kwargs = joint_attention_kwargs - self._interrupt = False - - init_image = self.image_processor.preprocess(image, height=height, width=width) - init_image = init_image.to(dtype=torch.float32) - - # 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 prompt embeddings - lora_scale = ( - self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None - ) - ( - prompt_embeds, - pooled_prompt_embeds, - text_ids, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - 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) - - 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.vae.config.latent_channels - 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, - ) - - # 6. Prepare mask and masked image latents - if masked_image_latents is not None: - masked_image_latents = masked_image_latents.to(latents.device) - else: - mask_image = self.mask_processor.preprocess(mask_image, height=height, width=width) - - masked_image = init_image * (1 - mask_image) - masked_image = masked_image.to(device=device, dtype=prompt_embeds.dtype) - - height, width = init_image.shape[-2:] - mask, masked_image_latents = self.prepare_mask_latents( - mask_image, - masked_image, - batch_size, - num_channels_latents, - num_images_per_prompt, - height, - width, - prompt_embeds.dtype, - device, - generator, - ) - masked_image_latents = torch.cat((masked_image_latents, mask), dim=-1) - - 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 - - # 7. 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]).to(latents.dtype) - - noise_pred = self.transformer( - hidden_states=torch.cat((latents, masked_image_latents), dim=2), - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - # 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() - - # 8. Post-process the image - 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 FluxPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_img2img.py deleted file mode 100644 index 1c4cf3b1cd16da81423eb04b61d910b58ee6b0d6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_img2img.py +++ /dev/null @@ -1,1067 +0,0 @@ -# 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, - CLIPTextModel, - CLIPTokenizer, - 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 FluxTransformer2DModel -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 FluxPipelineOutput - - -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 FluxImg2ImgPipeline - >>> from diffusers.utils import load_image - - >>> device = "cuda" - >>> pipe = FluxImg2ImgPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16) - >>> pipe = pipe.to(device) - - >>> 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)) - - >>> prompt = "cat wizard, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney, 8k" - - >>> images = pipe( - ... prompt=prompt, image=init_image, num_inference_steps=4, strength=0.95, guidance_scale=0.0 - ... ).images[0] - ``` -""" - - -# 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: Optional[torch.Generator] = 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 FluxImg2ImgPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin, FluxIPAdapterMixin): - r""" - The Flux pipeline for image inpainting. - - Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ - - Args: - transformer ([`FluxTransformer2DModel`]): - 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 ([`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: CLIPTextModel, - tokenizer: CLIPTokenizer, - text_encoder_2: T5EncoderModel, - tokenizer_2: T5TokenizerFast, - transformer: FluxTransformer2DModel, - image_encoder: CLIPVisionModelWithProjection = None, - feature_extractor: CLIPImageProcessor = None, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - 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 - # 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.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, vae_latent_channels=self.latent_channels - ) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = 128 - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds - 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_2) - - text_inputs = self.tokenizer_2( - 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 - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, self.tokenizer_max_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_2(text_input_ids.to(device), output_hidden_states=False)[0] - - dtype = self.text_encoder_2.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - ): - device = device or self._execution_device - - 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=self.tokenizer_max_length, - truncation=True, - return_overflowing_tokens=False, - return_length=False, - 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_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds.pooler_output - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - max_sequence_length: int = 512, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - scale_lora_layers(self.text_encoder_2, lora_scale) - - prompt = [prompt] if isinstance(prompt, str) else prompt - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - # We only use the pooled prompt output from the CLIPTextModel - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - ) - prompt_embeds = self._get_t5_prompt_embeds( - prompt=prompt_2, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - 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) - - if self.text_encoder_2 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_2, lora_scale) - - 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) - - return prompt_embeds, pooled_prompt_embeds, text_ids - - # 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, - prompt_2, - strength, - 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, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=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_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`." - ) - - 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 - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.enable_vae_slicing - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.disable_vae_slicing - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.enable_vae_tiling - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.disable_vae_tiling - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - 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) - - if latents is not None: - return latents.to(device=device, dtype=dtype), latent_image_ids - - 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) - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = self.scheduler.scale_noise(image_latents, timestep, noise) - latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) - return latents, latent_image_ids - - @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 - - @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, - negative_prompt: Union[str, List[str]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - true_cfg_scale: float = 1.0, - image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - strength: float = 0.6, - num_inference_steps: int = 28, - sigmas: Optional[List[float]] = 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", - 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. - 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`, 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. - 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. - 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`. - 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 7.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_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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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. - 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.FluxPipelineOutput`] 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.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] 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_2, - strength, - height, - width, - 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, - 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._interrupt = False - - # 2. Preprocess image - init_image = self.image_processor.preprocess(image, height=height, width=width) - 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 - ) - do_true_cfg = true_cfg_scale > 1 and negative_prompt is not None - ( - prompt_embeds, - pooled_prompt_embeds, - text_ids, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - device=device, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - if do_true_cfg: - ( - negative_prompt_embeds, - negative_pooled_prompt_embeds, - _, - ) = self.encode_prompt( - prompt=negative_prompt, - prompt_2=negative_prompt_2, - prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=negative_pooled_prompt_embeds, - 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) - - 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 - - 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, - ) - - 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, - ) - - # 6. Denoising loop - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - 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, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - if do_true_cfg: - 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, - guidance=guidance, - pooled_projections=negative_pooled_prompt_embeds, - encoder_hidden_states=negative_prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - noise_pred = neg_noise_pred + true_cfg_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() - - 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 FluxPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_inpaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_inpaint.py deleted file mode 100644 index eeacd9b19b8a8b7a487cf900c5b2021d1f7068f2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_inpaint.py +++ /dev/null @@ -1,1207 +0,0 @@ -# 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, - CLIPTextModel, - CLIPTokenizer, - CLIPVisionModelWithProjection, - T5EncoderModel, - T5TokenizerFast, -) - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FluxIPAdapterMixin, FluxLoraLoaderMixin, TextualInversionLoaderMixin -from ...models.autoencoders import AutoencoderKL -from ...models.transformers import FluxTransformer2DModel -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 FluxPipelineOutput - - -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 FluxInpaintPipeline - >>> from diffusers.utils import load_image - - >>> pipe = FluxInpaintPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", 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("flux_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: Optional[torch.Generator] = 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 FluxInpaintPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FluxIPAdapterMixin): - r""" - The Flux pipeline for image inpainting. - - Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ - - Args: - transformer ([`FluxTransformer2DModel`]): - 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 ([`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: CLIPTextModel, - tokenizer: CLIPTokenizer, - text_encoder_2: T5EncoderModel, - tokenizer_2: T5TokenizerFast, - transformer: FluxTransformer2DModel, - image_encoder: CLIPVisionModelWithProjection = None, - feature_extractor: CLIPImageProcessor = None, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - 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 - # 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.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, vae_latent_channels=self.latent_channels - ) - 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, - ) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = 128 - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds - 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_2) - - text_inputs = self.tokenizer_2( - 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 - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, self.tokenizer_max_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_2(text_input_ids.to(device), output_hidden_states=False)[0] - - dtype = self.text_encoder_2.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - ): - device = device or self._execution_device - - 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=self.tokenizer_max_length, - truncation=True, - return_overflowing_tokens=False, - return_length=False, - 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_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds.pooler_output - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - max_sequence_length: int = 512, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - scale_lora_layers(self.text_encoder_2, lora_scale) - - prompt = [prompt] if isinstance(prompt, str) else prompt - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - # We only use the pooled prompt output from the CLIPTextModel - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - ) - prompt_embeds = self._get_t5_prompt_embeds( - prompt=prompt_2, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - 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) - - if self.text_encoder_2 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_2, lora_scale) - - 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) - - return prompt_embeds, pooled_prompt_embeds, text_ids - - # 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, - prompt_2, - image, - mask_image, - strength, - height, - width, - output_type, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - padding_mask_crop=None, - max_sequence_length=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_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`." - ) - - 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 - - @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, - negative_prompt: Union[str, List[str]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = 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, - strength: float = 0.6, - num_inference_steps: int = 28, - sigmas: Optional[List[float]] = 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", - 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. - 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`, 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]`): - `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]`): - `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): - 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. - padding_mask_crop (`int`, *optional*, defaults to `None`): - The size of margin in the crop to be applied to the image and masking. If `None`, no crop is applied to - image and mask_image. If `padding_mask_crop` is not `None`, it will first find a rectangular region - with the same aspect ration of the image and contains all masked area, and then expand that area based - on `padding_mask_crop`. The image and mask_image will then be cropped based on the expanded area before - resizing to the original image size for inpainting. This is useful when the masked area is small while - the image is large and contain information irrelevant for inpainting, such as background. - 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`. - 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 7.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_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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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. - 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.FluxPipelineOutput`] 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.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] 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_2, - image, - mask_image, - strength, - height, - width, - output_type=output_type, - 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, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - padding_mask_crop=padding_mask_crop, - max_sequence_length=max_sequence_length, - ) - - self._guidance_scale = guidance_scale - self._joint_attention_kwargs = joint_attention_kwargs - 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 - ) - do_true_cfg = true_cfg_scale > 1 and negative_prompt is not None - ( - prompt_embeds, - pooled_prompt_embeds, - text_ids, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - device=device, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - if do_true_cfg: - ( - negative_prompt_embeds, - negative_pooled_prompt_embeds, - _, - ) = self.encode_prompt( - prompt=negative_prompt, - prompt_2=negative_prompt_2, - prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=negative_pooled_prompt_embeds, - 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) - - 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, - ) - - # 6. Denoising loop - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - 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, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - if do_true_cfg: - 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, - guidance=guidance, - pooled_projections=negative_pooled_prompt_embeds, - encoder_hidden_states=negative_prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - noise_pred = neg_noise_pred + true_cfg_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] - - # 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() - - 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) - - if padding_mask_crop is not None: - image = [ - self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image - ] - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return FluxPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_prior_redux.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_prior_redux.py deleted file mode 100644 index b5ccfb31a3521899dc7b7e25afe26c0cfd25942a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_prior_redux.py +++ /dev/null @@ -1,492 +0,0 @@ -# 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. - - -from typing import List, Optional, Union - -import torch -from PIL import Image -from transformers import ( - CLIPTextModel, - CLIPTokenizer, - SiglipImageProcessor, - SiglipVisionModel, - T5EncoderModel, - T5TokenizerFast, -) - -from ...image_processor import PipelineImageInput -from ...loaders import FluxLoraLoaderMixin, TextualInversionLoaderMixin -from ...utils import ( - USE_PEFT_BACKEND, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ..pipeline_utils import DiffusionPipeline -from .modeling_flux import ReduxImageEncoder -from .pipeline_output import FluxPriorReduxPipelineOutput - - -if is_torch_xla_available(): - 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 FluxPriorReduxPipeline, FluxPipeline - >>> from diffusers.utils import load_image - - >>> device = "cuda" - >>> dtype = torch.bfloat16 - - >>> repo_redux = "black-forest-labs/FLUX.1-Redux-dev" - >>> repo_base = "black-forest-labs/FLUX.1-dev" - >>> pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(repo_redux, torch_dtype=dtype).to(device) - >>> pipe = FluxPipeline.from_pretrained( - ... repo_base, text_encoder=None, text_encoder_2=None, torch_dtype=torch.bfloat16 - ... ).to(device) - - >>> image = load_image( - ... "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/style_ziggy/img5.png" - ... ) - >>> pipe_prior_output = pipe_prior_redux(image) - >>> images = pipe( - ... guidance_scale=2.5, - ... num_inference_steps=50, - ... generator=torch.Generator("cpu").manual_seed(0), - ... **pipe_prior_output, - ... ).images - >>> images[0].save("flux-redux.png") - ``` -""" - - -class FluxPriorReduxPipeline(DiffusionPipeline): - r""" - The Flux Redux pipeline for image-to-image generation. - - Reference: https://blackforestlabs.ai/flux-1-tools/ - - Args: - image_encoder ([`SiglipVisionModel`]): - SIGLIP vision model to encode the input image. - feature_extractor ([`SiglipImageProcessor`]): - Image processor for preprocessing images for the SIGLIP model. - image_embedder ([`ReduxImageEncoder`]): - Redux image encoder to process the SIGLIP embeddings. - text_encoder ([`CLIPTextModel`], *optional*): - [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`], *optional*): - [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`, *optional*): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer). - tokenizer_2 (`T5TokenizerFast`, *optional*): - Second Tokenizer of class - [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast). - """ - - model_cpu_offload_seq = "image_encoder->image_embedder" - _optional_components = [ - "text_encoder", - "tokenizer", - "text_encoder_2", - "tokenizer_2", - ] - _callback_tensor_inputs = [] - - def __init__( - self, - image_encoder: SiglipVisionModel, - feature_extractor: SiglipImageProcessor, - image_embedder: ReduxImageEncoder, - text_encoder: CLIPTextModel = None, - tokenizer: CLIPTokenizer = None, - text_encoder_2: T5EncoderModel = None, - tokenizer_2: T5TokenizerFast = None, - ): - super().__init__() - - self.register_modules( - image_encoder=image_encoder, - feature_extractor=feature_extractor, - image_embedder=image_embedder, - text_encoder=text_encoder, - tokenizer=tokenizer, - text_encoder_2=text_encoder_2, - tokenizer_2=tokenizer_2, - ) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - - def check_inputs( - self, - image, - prompt, - prompt_2, - prompt_embeds=None, - pooled_prompt_embeds=None, - prompt_embeds_scale=1.0, - pooled_prompt_embeds_scale=1.0, - ): - 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 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 prompt is not None and (isinstance(prompt, list) and isinstance(image, list) and len(prompt) != len(image)): - raise ValueError( - f"number of prompts must be equal to number of images, but {len(prompt)} prompts were provided and {len(image)} images" - ) - 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 isinstance(prompt_embeds_scale, list) and ( - isinstance(image, list) and len(prompt_embeds_scale) != len(image) - ): - raise ValueError( - f"number of weights must be equal to number of images, but {len(prompt_embeds_scale)} weights were provided and {len(image)} images" - ) - - def encode_image(self, image, device, num_images_per_prompt): - dtype = next(self.image_encoder.parameters()).dtype - image = self.feature_extractor.preprocess( - images=image, do_resize=True, return_tensors="pt", do_convert_rgb=True - ) - image = image.to(device=device, dtype=dtype) - - image_enc_hidden_states = self.image_encoder(**image).last_hidden_state - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - - return image_enc_hidden_states - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds - 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_2) - - text_inputs = self.tokenizer_2( - 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 - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, self.tokenizer_max_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_2(text_input_ids.to(device), output_hidden_states=False)[0] - - dtype = self.text_encoder_2.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - ): - device = device or self._execution_device - - 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=self.tokenizer_max_length, - truncation=True, - return_overflowing_tokens=False, - return_length=False, - 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_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds.pooler_output - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - max_sequence_length: int = 512, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - scale_lora_layers(self.text_encoder_2, lora_scale) - - prompt = [prompt] if isinstance(prompt, str) else prompt - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - # We only use the pooled prompt output from the CLIPTextModel - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - ) - prompt_embeds = self._get_t5_prompt_embeds( - prompt=prompt_2, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - 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) - - if self.text_encoder_2 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_2, lora_scale) - - 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) - - return prompt_embeds, pooled_prompt_embeds, text_ids - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - 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, - 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`, 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*): - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated pooled text embeddings. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.flux.FluxPriorReduxPipelineOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`~pipelines.flux.FluxPriorReduxPipelineOutput`] or `tuple`: - [`~pipelines.flux.FluxPriorReduxPipelineOutput`] 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( - image, - prompt, - prompt_2, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - prompt_embeds_scale=prompt_embeds_scale, - pooled_prompt_embeds_scale=pooled_prompt_embeds_scale, - ) - - # 2. Define call parameters - if image is not None and isinstance(image, Image.Image): - batch_size = 1 - elif image is not None and isinstance(image, list): - batch_size = len(image) - else: - batch_size = image.shape[0] - if prompt is not None and isinstance(prompt, str): - prompt = batch_size * [prompt] - if isinstance(prompt_embeds_scale, float): - prompt_embeds_scale = batch_size * [prompt_embeds_scale] - if isinstance(pooled_prompt_embeds_scale, float): - pooled_prompt_embeds_scale = batch_size * [pooled_prompt_embeds_scale] - - device = self._execution_device - - # 3. Prepare image embeddings - image_latents = self.encode_image(image, device, 1) - - image_embeds = self.image_embedder(image_latents).image_embeds - image_embeds = image_embeds.to(device=device) - - # 3. Prepare (dummy) text embeddings - if hasattr(self, "text_encoder") and self.text_encoder is not None: - ( - prompt_embeds, - pooled_prompt_embeds, - _, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - device=device, - num_images_per_prompt=1, - max_sequence_length=512, - lora_scale=None, - ) - else: - if prompt is not None: - logger.warning( - "prompt input is ignored when text encoders are not loaded to the pipeline. " - "Make sure to explicitly load the text encoders to enable prompt input. " - ) - # max_sequence_length is 512, t5 encoder hidden size is 4096 - prompt_embeds = torch.zeros((batch_size, 512, 4096), device=device, dtype=image_embeds.dtype) - # pooled_prompt_embeds is 768, clip text encoder hidden size - pooled_prompt_embeds = torch.zeros((batch_size, 768), device=device, dtype=image_embeds.dtype) - - # scale & concatenate image and text embeddings - prompt_embeds = torch.cat([prompt_embeds, image_embeds], dim=1) - - prompt_embeds *= torch.tensor(prompt_embeds_scale, device=device, dtype=image_embeds.dtype)[:, None, None] - pooled_prompt_embeds *= torch.tensor(pooled_prompt_embeds_scale, device=device, dtype=image_embeds.dtype)[ - :, None - ] - - # weighted sum - prompt_embeds = torch.sum(prompt_embeds, dim=0, keepdim=True) - pooled_prompt_embeds = torch.sum(pooled_prompt_embeds, dim=0, keepdim=True) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (prompt_embeds, pooled_prompt_embeds) - - return FluxPriorReduxPipelineOutput(prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_output.py deleted file mode 100644 index 388824e89f87c8fd7e264a58795631877ecd4f16..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_output.py +++ /dev/null @@ -1,37 +0,0 @@ -from dataclasses import dataclass -from typing import List, Union - -import numpy as np -import PIL.Image -import torch - -from ...utils import BaseOutput - - -@dataclass -class FluxPipelineOutput(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, - num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] - - -@dataclass -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, - num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - """ - - prompt_embeds: torch.Tensor - pooled_prompt_embeds: torch.Tensor diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/free_init_utils.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/free_init_utils.py deleted file mode 100644 index 4495c5ea2683380fee98ad8ddba4959abe8a801e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/free_init_utils.py +++ /dev/null @@ -1,187 +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 math -from typing import Tuple, Union - -import torch -import torch.fft as fft - -from ..utils.torch_utils import randn_tensor - - -class FreeInitMixin: - r"""Mixin class for FreeInit.""" - - def enable_free_init( - self, - num_iters: int = 3, - use_fast_sampling: bool = False, - method: str = "butterworth", - order: int = 4, - spatial_stop_frequency: float = 0.25, - temporal_stop_frequency: float = 0.25, - ): - """Enables the FreeInit mechanism as in https://huggingface.co/papers/2312.07537. - - This implementation has been adapted from the [official repository](https://github.com/TianxingWu/FreeInit). - - Args: - num_iters (`int`, *optional*, defaults to `3`): - Number of FreeInit noise re-initialization iterations. - use_fast_sampling (`bool`, *optional*, defaults to `False`): - Whether or not to speedup sampling procedure at the cost of probably lower quality results. Enables the - "Coarse-to-Fine Sampling" strategy, as mentioned in the paper, if set to `True`. - method (`str`, *optional*, defaults to `butterworth`): - Must be one of `butterworth`, `ideal` or `gaussian` to use as the filtering method for the FreeInit low - pass filter. - order (`int`, *optional*, defaults to `4`): - Order of the filter used in `butterworth` method. Larger values lead to `ideal` method behaviour - whereas lower values lead to `gaussian` method behaviour. - spatial_stop_frequency (`float`, *optional*, defaults to `0.25`): - Normalized stop frequency for spatial dimensions. Must be between 0 to 1. Referred to as `d_s` in the - original implementation. - temporal_stop_frequency (`float`, *optional*, defaults to `0.25`): - Normalized stop frequency for temporal dimensions. Must be between 0 to 1. Referred to as `d_t` in the - original implementation. - """ - self._free_init_num_iters = num_iters - self._free_init_use_fast_sampling = use_fast_sampling - self._free_init_method = method - self._free_init_order = order - self._free_init_spatial_stop_frequency = spatial_stop_frequency - self._free_init_temporal_stop_frequency = temporal_stop_frequency - - def disable_free_init(self): - """Disables the FreeInit mechanism if enabled.""" - self._free_init_num_iters = None - - @property - def free_init_enabled(self): - return hasattr(self, "_free_init_num_iters") and self._free_init_num_iters is not None - - def _get_free_init_freq_filter( - self, - shape: Tuple[int, ...], - device: Union[str, torch.dtype], - filter_type: str, - order: float, - spatial_stop_frequency: float, - temporal_stop_frequency: float, - ) -> torch.Tensor: - r"""Returns the FreeInit filter based on filter type and other input conditions.""" - - time, height, width = shape[-3], shape[-2], shape[-1] - mask = torch.zeros(shape) - - if spatial_stop_frequency == 0 or temporal_stop_frequency == 0: - return mask - - if filter_type == "butterworth": - - def retrieve_mask(x): - return 1 / (1 + (x / spatial_stop_frequency**2) ** order) - elif filter_type == "gaussian": - - def retrieve_mask(x): - return math.exp(-1 / (2 * spatial_stop_frequency**2) * x) - elif filter_type == "ideal": - - def retrieve_mask(x): - return 1 if x <= spatial_stop_frequency * 2 else 0 - else: - raise NotImplementedError("`filter_type` must be one of gaussian, butterworth or ideal") - - for t in range(time): - for h in range(height): - for w in range(width): - d_square = ( - ((spatial_stop_frequency / temporal_stop_frequency) * (2 * t / time - 1)) ** 2 - + (2 * h / height - 1) ** 2 - + (2 * w / width - 1) ** 2 - ) - mask[..., t, h, w] = retrieve_mask(d_square) - - return mask.to(device) - - def _apply_freq_filter(self, x: torch.Tensor, noise: torch.Tensor, low_pass_filter: torch.Tensor) -> torch.Tensor: - r"""Noise reinitialization.""" - # FFT - x_freq = fft.fftn(x, dim=(-3, -2, -1)) - x_freq = fft.fftshift(x_freq, dim=(-3, -2, -1)) - noise_freq = fft.fftn(noise, dim=(-3, -2, -1)) - noise_freq = fft.fftshift(noise_freq, dim=(-3, -2, -1)) - - # frequency mix - high_pass_filter = 1 - low_pass_filter - x_freq_low = x_freq * low_pass_filter - noise_freq_high = noise_freq * high_pass_filter - x_freq_mixed = x_freq_low + noise_freq_high # mix in freq domain - - # IFFT - x_freq_mixed = fft.ifftshift(x_freq_mixed, dim=(-3, -2, -1)) - x_mixed = fft.ifftn(x_freq_mixed, dim=(-3, -2, -1)).real - - return x_mixed - - def _apply_free_init( - self, - latents: torch.Tensor, - free_init_iteration: int, - num_inference_steps: int, - device: torch.device, - dtype: torch.dtype, - generator: torch.Generator, - ): - if free_init_iteration == 0: - self._free_init_initial_noise = latents.detach().clone() - else: - latent_shape = latents.shape - - free_init_filter_shape = (1, *latent_shape[1:]) - free_init_freq_filter = self._get_free_init_freq_filter( - shape=free_init_filter_shape, - device=device, - filter_type=self._free_init_method, - order=self._free_init_order, - spatial_stop_frequency=self._free_init_spatial_stop_frequency, - temporal_stop_frequency=self._free_init_temporal_stop_frequency, - ) - - current_diffuse_timestep = self.scheduler.config.num_train_timesteps - 1 - diffuse_timesteps = torch.full((latent_shape[0],), current_diffuse_timestep).long() - - z_t = self.scheduler.add_noise( - original_samples=latents, noise=self._free_init_initial_noise, timesteps=diffuse_timesteps.to(device) - ).to(dtype=torch.float32) - - z_rand = randn_tensor( - shape=latent_shape, - generator=generator, - device=device, - dtype=torch.float32, - ) - latents = self._apply_freq_filter(z_t, z_rand, low_pass_filter=free_init_freq_filter) - latents = latents.to(dtype) - - # Coarse-to-Fine Sampling for faster inference (can lead to lower quality) - if self._free_init_use_fast_sampling: - num_inference_steps = max( - 1, int(num_inference_steps / self._free_init_num_iters * (free_init_iteration + 1)) - ) - - if num_inference_steps > 0: - self.scheduler.set_timesteps(num_inference_steps, device=device) - - return latents, self.scheduler.timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/free_noise_utils.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/free_noise_utils.py deleted file mode 100644 index 2910afaf237bada3528628316fa5634946aa38a0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/free_noise_utils.py +++ /dev/null @@ -1,596 +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 Callable, Dict, List, Optional, Tuple, Union - -import torch -import torch.nn as nn - -from ..models.attention import BasicTransformerBlock, FreeNoiseTransformerBlock -from ..models.resnet import Downsample2D, ResnetBlock2D, Upsample2D -from ..models.transformers.transformer_2d import Transformer2DModel -from ..models.unets.unet_motion_model import ( - AnimateDiffTransformer3D, - CrossAttnDownBlockMotion, - DownBlockMotion, - UpBlockMotion, -) -from ..pipelines.pipeline_utils import DiffusionPipeline -from ..utils import logging -from ..utils.torch_utils import randn_tensor - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class SplitInferenceModule(nn.Module): - r""" - A wrapper module class that splits inputs along a specified dimension before performing a forward pass. - - This module is useful when you need to perform inference on large tensors in a memory-efficient way by breaking - them into smaller chunks, processing each chunk separately, and then reassembling the results. - - Args: - module (`nn.Module`): - The underlying PyTorch module that will be applied to each chunk of split inputs. - split_size (`int`, defaults to `1`): - 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"]`): - A list of keyword arguments (strings) that represent the input tensors to be split. - - Workflow: - 1. The keyword arguments specified in `input_kwargs_to_split` are split into smaller chunks using - `torch.split()` along the dimension `split_dim` and with a chunk size of `split_size`. - 2. The `module` is invoked once for each split with both the split inputs and any unchanged arguments - that were passed. - 3. The output tensors from each split are concatenated back together along `split_dim` before returning. - - Example: - ```python - >>> import torch - >>> import torch.nn as nn - - >>> model = nn.Linear(1000, 1000) - >>> split_module = SplitInferenceModule(model, split_size=2, split_dim=0, input_kwargs_to_split=["input"]) - - >>> input_tensor = torch.randn(42, 1000) - >>> # Will split the tensor into 21 slices of shape [2, 1000]. - >>> output = split_module(input=input_tensor) - ``` - - It is also possible to nest `SplitInferenceModule` across different split dimensions for more complex - multi-dimensional splitting. - """ - - def __init__( - self, - module: nn.Module, - split_size: int = 1, - split_dim: int = 0, - input_kwargs_to_split: List[str] = ["hidden_states"], - ) -> None: - super().__init__() - - self.module = module - self.split_size = split_size - 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]]: - r"""Forward method for the `SplitInferenceModule`. - - This method processes the input by splitting specified keyword arguments along a given dimension, running the - underlying module on each split, and then concatenating the results. The splitting is controlled by the - `split_size` and `split_dim` parameters specified during initialization. - - Args: - *args (`Any`): - Positional arguments that are passed directly to the `module` without modification. - **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]]`: - 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 - along the same `split_dim` after processing all splits. - - If the underlying module returns a tuple of tensors, each element of the tuple will be concatenated - along the `split_dim` across all splits, and the final result will be a tuple of concatenated tensors. - """ - split_inputs = {} - - # 1. Split inputs that were specified during initialization and also present in passed kwargs - for key in list(kwargs.keys()): - if key not in self.input_kwargs_to_split or not torch.is_tensor(kwargs[key]): - continue - split_inputs[key] = torch.split(kwargs[key], self.split_size, self.split_dim) - kwargs.pop(key) - - # 2. Invoke forward pass across each split - results = [] - for split_input in zip(*split_inputs.values()): - inputs = dict(zip(split_inputs.keys(), split_input)) - inputs.update(kwargs) - - intermediate_tensor_or_tensor_tuple = self.module(*args, **inputs) - results.append(intermediate_tensor_or_tensor_tuple) - - # 3. Concatenate split restuls to obtain final outputs - if isinstance(results[0], torch.Tensor): - return torch.cat(results, dim=self.split_dim) - elif isinstance(results[0], tuple): - return tuple([torch.cat(x, dim=self.split_dim) for x in zip(*results)]) - else: - raise ValueError( - "In order to use the SplitInferenceModule, it is necessary for the underlying `module` to either return a torch.Tensor or a tuple of torch.Tensor's." - ) - - -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]): - r"""Helper function to enable FreeNoise in transformer blocks.""" - - for motion_module in block.motion_modules: - num_transformer_blocks = len(motion_module.transformer_blocks) - - for i in range(num_transformer_blocks): - if isinstance(motion_module.transformer_blocks[i], FreeNoiseTransformerBlock): - motion_module.transformer_blocks[i].set_free_noise_properties( - self._free_noise_context_length, - self._free_noise_context_stride, - self._free_noise_weighting_scheme, - ) - else: - assert isinstance(motion_module.transformer_blocks[i], BasicTransformerBlock) - basic_transfomer_block = motion_module.transformer_blocks[i] - - motion_module.transformer_blocks[i] = FreeNoiseTransformerBlock( - dim=basic_transfomer_block.dim, - num_attention_heads=basic_transfomer_block.num_attention_heads, - attention_head_dim=basic_transfomer_block.attention_head_dim, - dropout=basic_transfomer_block.dropout, - cross_attention_dim=basic_transfomer_block.cross_attention_dim, - activation_fn=basic_transfomer_block.activation_fn, - attention_bias=basic_transfomer_block.attention_bias, - only_cross_attention=basic_transfomer_block.only_cross_attention, - double_self_attention=basic_transfomer_block.double_self_attention, - positional_embeddings=basic_transfomer_block.positional_embeddings, - num_positional_embeddings=basic_transfomer_block.num_positional_embeddings, - context_length=self._free_noise_context_length, - context_stride=self._free_noise_context_stride, - weighting_scheme=self._free_noise_weighting_scheme, - ).to(device=self.device, dtype=self.dtype) - - motion_module.transformer_blocks[i].load_state_dict( - basic_transfomer_block.state_dict(), strict=True - ) - motion_module.transformer_blocks[i].set_chunk_feed_forward( - basic_transfomer_block._chunk_size, basic_transfomer_block._chunk_dim - ) - - def _disable_free_noise_in_block(self, block: Union[CrossAttnDownBlockMotion, DownBlockMotion, UpBlockMotion]): - r"""Helper function to disable FreeNoise in transformer blocks.""" - - for motion_module in block.motion_modules: - num_transformer_blocks = len(motion_module.transformer_blocks) - - for i in range(num_transformer_blocks): - if isinstance(motion_module.transformer_blocks[i], FreeNoiseTransformerBlock): - free_noise_transfomer_block = motion_module.transformer_blocks[i] - - motion_module.transformer_blocks[i] = BasicTransformerBlock( - dim=free_noise_transfomer_block.dim, - num_attention_heads=free_noise_transfomer_block.num_attention_heads, - attention_head_dim=free_noise_transfomer_block.attention_head_dim, - dropout=free_noise_transfomer_block.dropout, - cross_attention_dim=free_noise_transfomer_block.cross_attention_dim, - activation_fn=free_noise_transfomer_block.activation_fn, - attention_bias=free_noise_transfomer_block.attention_bias, - only_cross_attention=free_noise_transfomer_block.only_cross_attention, - double_self_attention=free_noise_transfomer_block.double_self_attention, - positional_embeddings=free_noise_transfomer_block.positional_embeddings, - num_positional_embeddings=free_noise_transfomer_block.num_positional_embeddings, - ).to(device=self.device, dtype=self.dtype) - - motion_module.transformer_blocks[i].load_state_dict( - free_noise_transfomer_block.state_dict(), strict=True - ) - motion_module.transformer_blocks[i].set_chunk_feed_forward( - free_noise_transfomer_block._chunk_size, free_noise_transfomer_block._chunk_dim - ) - - def _check_inputs_free_noise( - self, - prompt, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - num_frames, - ) -> None: - if not isinstance(prompt, (str, dict)): - raise ValueError(f"Expected `prompt` to have type `str` or `dict` but found {type(prompt)=}") - - if negative_prompt is not None: - if not isinstance(negative_prompt, (str, dict)): - raise ValueError( - f"Expected `negative_prompt` to have type `str` or `dict` but found {type(negative_prompt)=}" - ) - - if prompt_embeds is not None or negative_prompt_embeds is not None: - raise ValueError("`prompt_embeds` and `negative_prompt_embeds` is not supported in FreeNoise yet.") - - frame_indices = [isinstance(x, int) for x in prompt.keys()] - frame_prompts = [isinstance(x, str) for x in prompt.values()] - min_frame = min(list(prompt.keys())) - max_frame = max(list(prompt.keys())) - - if not all(frame_indices): - raise ValueError("Expected integer keys in `prompt` dict for FreeNoise.") - if not all(frame_prompts): - raise ValueError("Expected str values in `prompt` dict for FreeNoise.") - if min_frame != 0: - raise ValueError("The minimum frame index in `prompt` dict must be 0 as a starting prompt is necessary.") - if max_frame >= num_frames: - raise ValueError( - f"The maximum frame index in `prompt` dict must be lesser than {num_frames=} and follow 0-based indexing." - ) - - def _encode_prompt_free_noise( - self, - prompt: Union[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, - ) -> torch.Tensor: - if negative_prompt is None: - negative_prompt = "" - - # Ensure that we have a dictionary of prompts - if isinstance(prompt, str): - prompt = {0: prompt} - if isinstance(negative_prompt, str): - negative_prompt = {0: negative_prompt} - - self._check_inputs_free_noise(prompt, negative_prompt, prompt_embeds, negative_prompt_embeds, num_frames) - - # Sort the prompts based on frame indices - prompt = dict(sorted(prompt.items())) - negative_prompt = dict(sorted(negative_prompt.items())) - - # Ensure that we have a prompt for the last frame index - prompt[num_frames - 1] = prompt[list(prompt.keys())[-1]] - negative_prompt[num_frames - 1] = negative_prompt[list(negative_prompt.keys())[-1]] - - frame_indices = list(prompt.keys()) - frame_prompts = list(prompt.values()) - frame_negative_indices = list(negative_prompt.keys()) - frame_negative_prompts = list(negative_prompt.values()) - - # Generate and interpolate positive prompts - prompt_embeds, _ = self.encode_prompt( - prompt=frame_prompts, - device=device, - num_images_per_prompt=num_videos_per_prompt, - do_classifier_free_guidance=False, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - lora_scale=lora_scale, - clip_skip=clip_skip, - ) - - shape = (num_frames, *prompt_embeds.shape[1:]) - prompt_interpolation_embeds = prompt_embeds.new_zeros(shape) - - for i in range(len(frame_indices) - 1): - start_frame = frame_indices[i] - end_frame = frame_indices[i + 1] - start_tensor = prompt_embeds[i].unsqueeze(0) - end_tensor = prompt_embeds[i + 1].unsqueeze(0) - - prompt_interpolation_embeds[start_frame : end_frame + 1] = self._free_noise_prompt_interpolation_callback( - start_frame, end_frame, start_tensor, end_tensor - ) - - # Generate and interpolate negative prompts - negative_prompt_embeds = None - negative_prompt_interpolation_embeds = None - - if do_classifier_free_guidance: - _, negative_prompt_embeds = self.encode_prompt( - prompt=[""] * len(frame_negative_prompts), - device=device, - num_images_per_prompt=num_videos_per_prompt, - do_classifier_free_guidance=True, - negative_prompt=frame_negative_prompts, - prompt_embeds=None, - negative_prompt_embeds=None, - lora_scale=lora_scale, - clip_skip=clip_skip, - ) - - negative_prompt_interpolation_embeds = negative_prompt_embeds.new_zeros(shape) - - for i in range(len(frame_negative_indices) - 1): - start_frame = frame_negative_indices[i] - end_frame = frame_negative_indices[i + 1] - start_tensor = negative_prompt_embeds[i].unsqueeze(0) - end_tensor = negative_prompt_embeds[i + 1].unsqueeze(0) - - negative_prompt_interpolation_embeds[start_frame : end_frame + 1] = ( - self._free_noise_prompt_interpolation_callback(start_frame, end_frame, start_tensor, end_tensor) - ) - - prompt_embeds = prompt_interpolation_embeds - negative_prompt_embeds = negative_prompt_interpolation_embeds - - if do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - return prompt_embeds, negative_prompt_embeds - - def _prepare_latents_free_noise( - self, - batch_size: int, - num_channels_latents: int, - num_frames: int, - height: int, - width: int, - dtype: torch.dtype, - device: torch.device, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.Tensor] = 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." - ) - - context_num_frames = ( - self._free_noise_context_length if self._free_noise_context_length == "repeat_context" else num_frames - ) - - shape = ( - batch_size, - num_channels_latents, - context_num_frames, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - if self._free_noise_noise_type == "random": - return latents - else: - if latents.size(2) == num_frames: - return latents - elif latents.size(2) != self._free_noise_context_length: - raise ValueError( - f"You have passed `latents` as a parameter to FreeNoise. The expected number of frames is either {num_frames} or {self._free_noise_context_length}, but found {latents.size(2)}" - ) - latents = latents.to(device) - - if self._free_noise_noise_type == "shuffle_context": - for i in range(self._free_noise_context_length, num_frames, self._free_noise_context_stride): - # ensure window is within bounds - window_start = max(0, i - self._free_noise_context_length) - window_end = min(num_frames, window_start + self._free_noise_context_stride) - window_length = window_end - window_start - - if window_length == 0: - break - - indices = torch.LongTensor(list(range(window_start, window_end))) - shuffled_indices = indices[torch.randperm(window_length, generator=generator)] - - current_start = i - current_end = min(num_frames, current_start + window_length) - if current_end == current_start + window_length: - # batch of frames perfectly fits the window - latents[:, :, current_start:current_end] = latents[:, :, shuffled_indices] - else: - # handle the case where the last batch of frames does not fit perfectly with the window - prefix_length = current_end - current_start - shuffled_indices = shuffled_indices[:prefix_length] - latents[:, :, current_start:current_end] = latents[:, :, shuffled_indices] - - elif self._free_noise_noise_type == "repeat_context": - num_repeats = (num_frames + self._free_noise_context_length - 1) // self._free_noise_context_length - latents = torch.cat([latents] * num_repeats, dim=2) - - latents = latents[:, :, :num_frames] - return latents - - def _lerp( - self, start_index: int, end_index: int, start_tensor: torch.Tensor, end_tensor: torch.Tensor - ) -> torch.Tensor: - num_indices = end_index - start_index + 1 - interpolated_tensors = [] - - for i in range(num_indices): - alpha = i / (num_indices - 1) - interpolated_tensor = (1 - alpha) * start_tensor + alpha * end_tensor - interpolated_tensors.append(interpolated_tensor) - - interpolated_tensors = torch.cat(interpolated_tensors) - return interpolated_tensors - - def enable_free_noise( - self, - context_length: Optional[int] = 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, - ) -> None: - r""" - Enable long video generation using FreeNoise. - - Args: - context_length (`int`, defaults to `16`, *optional*): - The number of video frames to process at once. It's recommended to set this to the maximum frames the - Motion Adapter was trained with (usually 16/24/32). If `None`, the default value from the motion - adapter config is used. - context_stride (`int`, *optional*): - Long videos are generated by processing many frames. FreeNoise processes these frames in sliding - windows of size `context_length`. Context stride allows you to specify how many frames to skip between - each window. For example, a context length of 16 and context stride of 4 would process 24 frames as: - [0, 15], [4, 19], [8, 23] (0-based indexing) - weighting_scheme (`str`, defaults to `pyramid`): - Weighting scheme for averaging latents after accumulation in FreeNoise blocks. The following weighting - schemes are supported currently: - - "flat" - Performs weighting averaging with a flat weight pattern: [1, 1, 1, 1, 1]. - - "pyramid" - Performs weighted averaging with a pyramid like weight pattern: [1, 2, 3, 2, 1]. - - "delayed_reverse_sawtooth" - Performs weighted averaging with low weights for earlier frames and high-to-low weights for - later frames: [0.01, 0.01, 3, 2, 1]. - noise_type (`str`, defaults to "shuffle_context"): - Must be one of ["shuffle_context", "repeat_context", "random"]. - - "shuffle_context" - Shuffles a fixed batch of `context_length` latents to create a final latent of size - `num_frames`. This is usually the best setting for most generation scenarios. However, there - might be visible repetition noticeable in the kinds of motion/animation generated. - - "repeated_context" - Repeats a fixed batch of `context_length` latents to create a final latent of size - `num_frames`. - - "random" - The final latents are random without any repetition. - """ - - allowed_weighting_scheme = ["flat", "pyramid", "delayed_reverse_sawtooth"] - allowed_noise_type = ["shuffle_context", "repeat_context", "random"] - - if context_length > self.motion_adapter.config.motion_max_seq_length: - logger.warning( - f"You have set {context_length=} which is greater than {self.motion_adapter.config.motion_max_seq_length=}. This can lead to bad generation results." - ) - if weighting_scheme not in allowed_weighting_scheme: - raise ValueError( - f"The parameter `weighting_scheme` must be one of {allowed_weighting_scheme}, but got {weighting_scheme=}" - ) - if noise_type not in allowed_noise_type: - raise ValueError(f"The parameter `noise_type` must be one of {allowed_noise_type}, but got {noise_type=}") - - self._free_noise_context_length = context_length or self.motion_adapter.config.motion_max_seq_length - self._free_noise_context_stride = context_stride - self._free_noise_weighting_scheme = weighting_scheme - self._free_noise_noise_type = noise_type - self._free_noise_prompt_interpolation_callback = prompt_interpolation_callback or self._lerp - - if hasattr(self.unet.mid_block, "motion_modules"): - blocks = [*self.unet.down_blocks, self.unet.mid_block, *self.unet.up_blocks] - else: - blocks = [*self.unet.down_blocks, *self.unet.up_blocks] - - for block in blocks: - self._enable_free_noise_in_block(block) - - def disable_free_noise(self) -> None: - r"""Disable the FreeNoise sampling mechanism.""" - self._free_noise_context_length = None - - if hasattr(self.unet.mid_block, "motion_modules"): - blocks = [*self.unet.down_blocks, self.unet.mid_block, *self.unet.up_blocks] - else: - blocks = [*self.unet.down_blocks, *self.unet.up_blocks] - - blocks = [*self.unet.down_blocks, self.unet.mid_block, *self.unet.up_blocks] - for block in blocks: - self._disable_free_noise_in_block(block) - - def _enable_split_inference_motion_modules_( - 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"]) - - for i in range(len(motion_module.transformer_blocks)): - motion_module.transformer_blocks[i] = SplitInferenceModule( - motion_module.transformer_blocks[i], - spatial_split_size, - 0, - ["hidden_states", "encoder_hidden_states"], - ) - - 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 - ) -> 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: - 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 - ) -> None: - for i in range(len(samplers)): - samplers[i] = SplitInferenceModule(samplers[i], temporal_split_size, 0, ["hidden_states"]) - - def enable_free_noise_split_inference(self, spatial_split_size: int = 256, temporal_split_size: int = 16) -> None: - r""" - Enable FreeNoise memory optimizations by utilizing - [`~diffusers.pipelines.free_noise_utils.SplitInferenceModule`] across different intermediate modeling blocks. - - Args: - spatial_split_size (`int`, defaults to `256`): - The split size across spatial dimensions for internal blocks. This is used in facilitating split - inference across the effective batch dimension (`[B x H x W, F, C]`) of intermediate tensors in motion - modeling blocks. - temporal_split_size (`int`, defaults to `16`): - The split size across temporal dimensions for internal blocks. This is used in facilitating split - inference across the effective batch dimension (`[B x F, H x W, C]`) of intermediate tensors in spatial - attention, resnets, downsampling and upsampling blocks. - """ - # TODO(aryan): Discuss on what's the best way to provide more control to users - blocks = [*self.unet.down_blocks, self.unet.mid_block, *self.unet.up_blocks] - for block in blocks: - if getattr(block, "motion_modules", None) is not None: - self._enable_split_inference_motion_modules_(block.motion_modules, spatial_split_size) - if getattr(block, "attentions", None) is not None: - self._enable_split_inference_attentions_(block.attentions, temporal_split_size) - if getattr(block, "resnets", None) is not None: - self._enable_split_inference_resnets_(block.resnets, temporal_split_size) - if getattr(block, "downsamplers", None) is not None: - self._enable_split_inference_samplers_(block.downsamplers, temporal_split_size) - if getattr(block, "upsamplers", None) is not None: - self._enable_split_inference_samplers_(block.upsamplers, temporal_split_size) - - @property - def free_noise_enabled(self): - return hasattr(self, "_free_noise_context_length") and self._free_noise_context_length is not None diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/__init__.py deleted file mode 100644 index 498df900e68b70920e5fc49764eb6bdcd2dc3354..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/__init__.py +++ /dev/null @@ -1,47 +0,0 @@ -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 = {} -_additional_imports = {} -_import_structure = {"pipeline_output": ["HiDreamImagePipelineOutput"]} - -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_hidream_image"] = ["HiDreamImagePipeline"] -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_hidream_image import HiDreamImagePipeline -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/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 9570ad4dbee5e115c1540f59c4d3ad57af1a8238..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/__pycache__/pipeline_hidream_image.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/__pycache__/pipeline_hidream_image.cpython-310.pyc deleted file mode 100644 index 7955d7b138b0dc5379cb594fc573d95a405f21d8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/__pycache__/pipeline_hidream_image.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index 241b8143a57e48ca9163ca62212cccd1408e0360..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/pipeline_hidream_image.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/pipeline_hidream_image.py deleted file mode 100644 index 341cdaf1e6efb386bc3b412abadc2ce0cb3b8961..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/pipeline_hidream_image.py +++ /dev/null @@ -1,1026 +0,0 @@ -# Copyright 2025 HiDream-ai 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 -import math -from typing import Any, Callable, Dict, List, Optional, Union - -import torch -from transformers import ( - CLIPTextModelWithProjection, - CLIPTokenizer, - LlamaForCausalLM, - PreTrainedTokenizerFast, - T5EncoderModel, - T5Tokenizer, -) - -from ...image_processor import VaeImageProcessor -from ...loaders import HiDreamImageLoraLoaderMixin -from ...models import AutoencoderKL, HiDreamImageTransformer2DModel -from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .pipeline_output import HiDreamImagePipelineOutput - - -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 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", - ... output_hidden_states=True, - ... output_attentions=True, - ... torch_dtype=torch.bfloat16, - ... ) - - >>> pipe = HiDreamImagePipeline.from_pretrained( - ... "HiDream-ai/HiDream-I1-Full", - ... tokenizer_4=tokenizer_4, - ... text_encoder_4=text_encoder_4, - ... torch_dtype=torch.bfloat16, - ... ) - >>> pipe.enable_model_cpu_offload() - - >>> image = pipe( - ... 'A cat holding a sign that says "Hi-Dreams.ai".', - ... height=1024, - ... width=1024, - ... guidance_scale=5.0, - ... num_inference_steps=50, - ... generator=torch.Generator("cuda").manual_seed(0), - ... ).images[0] - >>> image.save("output.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 HiDreamImagePipeline(DiffusionPipeline, HiDreamImageLoraLoaderMixin): - model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->text_encoder_4->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds_t5", "prompt_embeds_llama3", "pooled_prompt_embeds"] - - def __init__( - self, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer_2: CLIPTokenizer, - text_encoder_3: T5EncoderModel, - tokenizer_3: T5Tokenizer, - text_encoder_4: LlamaForCausalLM, - tokenizer_4: PreTrainedTokenizerFast, - transformer: HiDreamImageTransformer2DModel, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - text_encoder_3=text_encoder_3, - text_encoder_4=text_encoder_4, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - tokenizer_3=tokenizer_3, - tokenizer_4=tokenizer_4, - 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 - ) - # HiDreamImage 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 - if getattr(self, "tokenizer_4", None) is not None: - self.tokenizer_4.pad_token = self.tokenizer_4.eos_token - - def _get_t5_prompt_embeds( - self, - prompt: Union[str, List[str]] = None, - max_sequence_length: int = 128, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - ): - device = device or self._execution_device - dtype = dtype or self.text_encoder_3.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - - text_inputs = self.tokenizer_3( - prompt, - padding="max_length", - max_length=min(max_sequence_length, self.tokenizer_3.model_max_length), - truncation=True, - add_special_tokens=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - attention_mask = text_inputs.attention_mask - untruncated_ids = self.tokenizer_3(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_3.batch_decode( - untruncated_ids[:, min(max_sequence_length, self.tokenizer_3.model_max_length) - 1 : -1] - ) - logger.warning( - "The following part of your input was truncated because `max_sequence_length` is set to " - f" {min(max_sequence_length, self.tokenizer_3.model_max_length)} tokens: {removed_text}" - ) - - prompt_embeds = self.text_encoder_3(text_input_ids.to(device), attention_mask=attention_mask.to(device))[0] - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - return prompt_embeds - - def _get_clip_prompt_embeds( - self, - tokenizer, - text_encoder, - prompt: Union[str, List[str]], - max_sequence_length: int = 128, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - ): - device = device or self._execution_device - dtype = dtype or text_encoder.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=min(max_sequence_length, 218), - 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[:, 218 - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {218} tokens: {removed_text}" - ) - prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds[0] - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - return prompt_embeds - - def _get_llama3_prompt_embeds( - self, - prompt: Union[str, List[str]] = None, - max_sequence_length: int = 128, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - ): - device = device or self._execution_device - dtype = dtype or self.text_encoder_4.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - - text_inputs = self.tokenizer_4( - prompt, - padding="max_length", - max_length=min(max_sequence_length, self.tokenizer_4.model_max_length), - truncation=True, - add_special_tokens=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - attention_mask = text_inputs.attention_mask - untruncated_ids = self.tokenizer_4(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_4.batch_decode( - untruncated_ids[:, min(max_sequence_length, self.tokenizer_4.model_max_length) - 1 : -1] - ) - logger.warning( - "The following part of your input was truncated because `max_sequence_length` is set to " - f" {min(max_sequence_length, self.tokenizer_4.model_max_length)} tokens: {removed_text}" - ) - - outputs = self.text_encoder_4( - text_input_ids.to(device), - attention_mask=attention_mask.to(device), - output_hidden_states=True, - output_attentions=True, - ) - - prompt_embeds = outputs.hidden_states[1:] - prompt_embeds = torch.stack(prompt_embeds, dim=0) - return 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, - 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, - max_sequence_length: int = 128, - lora_scale: Optional[float] = None, - ): - prompt = [prompt] if isinstance(prompt, str) else prompt - if prompt is not None: - batch_size = len(prompt) - else: - batch_size = pooled_prompt_embeds.shape[0] - - device = device or self._execution_device - - if pooled_prompt_embeds is None: - pooled_prompt_embeds_1 = self._get_clip_prompt_embeds( - self.tokenizer, self.text_encoder, prompt, max_sequence_length, device, dtype - ) - - if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt - - if len(negative_prompt) > 1 and len(negative_prompt) != batch_size: - raise ValueError(f"negative_prompt must be of length 1 or {batch_size}") - - negative_pooled_prompt_embeds_1 = self._get_clip_prompt_embeds( - self.tokenizer, self.text_encoder, negative_prompt, max_sequence_length, device, dtype - ) - - if negative_pooled_prompt_embeds_1.shape[0] == 1 and batch_size > 1: - negative_pooled_prompt_embeds_1 = negative_pooled_prompt_embeds_1.repeat(batch_size, 1) - - if pooled_prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - if len(prompt_2) > 1 and len(prompt_2) != batch_size: - raise ValueError(f"prompt_2 must be of length 1 or {batch_size}") - - pooled_prompt_embeds_2 = self._get_clip_prompt_embeds( - self.tokenizer_2, self.text_encoder_2, prompt_2, max_sequence_length, device, dtype - ) - - if pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1: - pooled_prompt_embeds_2 = pooled_prompt_embeds_2.repeat(batch_size, 1) - - if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: - negative_prompt_2 = negative_prompt_2 or negative_prompt - negative_prompt_2 = [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 - - if len(negative_prompt_2) > 1 and len(negative_prompt_2) != batch_size: - raise ValueError(f"negative_prompt_2 must be of length 1 or {batch_size}") - - negative_pooled_prompt_embeds_2 = self._get_clip_prompt_embeds( - self.tokenizer_2, self.text_encoder_2, negative_prompt_2, max_sequence_length, device, dtype - ) - - if negative_pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1: - negative_pooled_prompt_embeds_2 = negative_pooled_prompt_embeds_2.repeat(batch_size, 1) - - if pooled_prompt_embeds is None: - pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_1, pooled_prompt_embeds_2], dim=-1) - - if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: - negative_pooled_prompt_embeds = torch.cat( - [negative_pooled_prompt_embeds_1, negative_pooled_prompt_embeds_2], dim=-1 - ) - - if prompt_embeds_t5 is None: - prompt_3 = prompt_3 or prompt - prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3 - - if len(prompt_3) > 1 and len(prompt_3) != batch_size: - raise ValueError(f"prompt_3 must be of length 1 or {batch_size}") - - prompt_embeds_t5 = self._get_t5_prompt_embeds(prompt_3, max_sequence_length, device, dtype) - - if prompt_embeds_t5.shape[0] == 1 and batch_size > 1: - prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1) - - if do_classifier_free_guidance and negative_prompt_embeds_t5 is None: - negative_prompt_3 = negative_prompt_3 or negative_prompt - negative_prompt_3 = [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3 - - if len(negative_prompt_3) > 1 and len(negative_prompt_3) != batch_size: - raise ValueError(f"negative_prompt_3 must be of length 1 or {batch_size}") - - negative_prompt_embeds_t5 = self._get_t5_prompt_embeds( - negative_prompt_3, max_sequence_length, device, dtype - ) - - if negative_prompt_embeds_t5.shape[0] == 1 and batch_size > 1: - negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1) - - if prompt_embeds_llama3 is None: - prompt_4 = prompt_4 or prompt - prompt_4 = [prompt_4] if isinstance(prompt_4, str) else prompt_4 - - if len(prompt_4) > 1 and len(prompt_4) != batch_size: - raise ValueError(f"prompt_4 must be of length 1 or {batch_size}") - - prompt_embeds_llama3 = self._get_llama3_prompt_embeds(prompt_4, max_sequence_length, device, dtype) - - if prompt_embeds_llama3.shape[0] == 1 and batch_size > 1: - prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1) - - if do_classifier_free_guidance and negative_prompt_embeds_llama3 is None: - negative_prompt_4 = negative_prompt_4 or negative_prompt - negative_prompt_4 = [negative_prompt_4] if isinstance(negative_prompt_4, str) else negative_prompt_4 - - if len(negative_prompt_4) > 1 and len(negative_prompt_4) != batch_size: - raise ValueError(f"negative_prompt_4 must be of length 1 or {batch_size}") - - negative_prompt_embeds_llama3 = self._get_llama3_prompt_embeds( - negative_prompt_4, max_sequence_length, device, dtype - ) - - if negative_prompt_embeds_llama3.shape[0] == 1 and batch_size > 1: - negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1) - - # duplicate pooled_prompt_embeds for each generation per prompt - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt) - pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - # duplicate t5_prompt_embeds for batch_size and num_images_per_prompt - bs_embed, seq_len, _ = prompt_embeds_t5.shape - if bs_embed == 1 and batch_size > 1: - prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate prompt_embeds_t5 of batch size {bs_embed}") - prompt_embeds_t5 = prompt_embeds_t5.repeat(1, num_images_per_prompt, 1) - prompt_embeds_t5 = prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1) - - # duplicate llama3_prompt_embeds for batch_size and num_images_per_prompt - _, bs_embed, seq_len, dim = prompt_embeds_llama3.shape - if bs_embed == 1 and batch_size > 1: - prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate prompt_embeds_llama3 of batch size {bs_embed}") - prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1) - prompt_embeds_llama3 = prompt_embeds_llama3.view(-1, batch_size * num_images_per_prompt, seq_len, dim) - - if do_classifier_free_guidance: - # duplicate negative_pooled_prompt_embeds for batch_size and num_images_per_prompt - bs_embed, seq_len = negative_pooled_prompt_embeds.shape - if bs_embed == 1 and batch_size > 1: - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(batch_size, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate negative_pooled_prompt_embeds of batch size {bs_embed}") - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt) - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - # duplicate negative_t5_prompt_embeds for batch_size and num_images_per_prompt - bs_embed, seq_len, _ = negative_prompt_embeds_t5.shape - if bs_embed == 1 and batch_size > 1: - negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate negative_prompt_embeds_t5 of batch size {bs_embed}") - negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds_t5 = negative_prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1) - - # duplicate negative_prompt_embeds_llama3 for batch_size and num_images_per_prompt - _, bs_embed, seq_len, dim = negative_prompt_embeds_llama3.shape - if bs_embed == 1 and batch_size > 1: - negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate negative_prompt_embeds_llama3 of batch size {bs_embed}") - negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1) - negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.view( - -1, batch_size * num_images_per_prompt, seq_len, dim - ) - - return ( - prompt_embeds_t5, - negative_prompt_embeds_t5, - prompt_embeds_llama3, - negative_prompt_embeds_llama3, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - def check_inputs( - self, - prompt, - prompt_2, - prompt_3, - prompt_4, - negative_prompt=None, - negative_prompt_2=None, - negative_prompt_3=None, - negative_prompt_4=None, - prompt_embeds_t5=None, - prompt_embeds_llama3=None, - negative_prompt_embeds_t5=None, - negative_prompt_embeds_llama3=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=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 pooled_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `pooled_prompt_embeds`: {pooled_prompt_embeds}. Please make sure to" - " only forward one of the two." - ) - elif prompt_2 is not None and pooled_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt_2`: {prompt_2} and `pooled_prompt_embeds`: {pooled_prompt_embeds}. Please make sure to" - " only forward one of the two." - ) - elif prompt_3 is not None and prompt_embeds_t5 is not None: - raise ValueError( - f"Cannot forward both `prompt_3`: {prompt_3} and `prompt_embeds_t5`: {prompt_embeds_t5}. Please make sure to" - " only forward one of the two." - ) - elif prompt_4 is not None and prompt_embeds_llama3 is not None: - raise ValueError( - f"Cannot forward both `prompt_4`: {prompt_4} and `prompt_embeds_llama3`: {prompt_embeds_llama3}. Please make sure to" - " only forward one of the two." - ) - elif prompt is None and pooled_prompt_embeds is None: - raise ValueError( - "Provide either `prompt` or `pooled_prompt_embeds`. Cannot leave both `prompt` and `pooled_prompt_embeds` undefined." - ) - elif prompt is None and prompt_embeds_t5 is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds_t5`. Cannot leave both `prompt` and `prompt_embeds_t5` undefined." - ) - elif prompt is None and prompt_embeds_llama3 is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds_llama3`. Cannot leave both `prompt` and `prompt_embeds_llama3` 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)}") - elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)): - raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}") - elif prompt_4 is not None and (not isinstance(prompt_4, str) and not isinstance(prompt_4, list)): - raise ValueError(f"`prompt_4` has to be of type `str` or `list` but is {type(prompt_4)}") - - if negative_prompt is not None and negative_pooled_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_pooled_prompt_embeds`:" - f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two." - ) - elif negative_prompt_2 is not None and negative_pooled_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_pooled_prompt_embeds`:" - f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two." - ) - elif negative_prompt_3 is not None and negative_prompt_embeds_t5 is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} and `negative_prompt_embeds_t5`:" - f" {negative_prompt_embeds_t5}. Please make sure to only forward one of the two." - ) - elif negative_prompt_4 is not None and negative_prompt_embeds_llama3 is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_4`: {negative_prompt_4} and `negative_prompt_embeds_llama3`:" - f" {negative_prompt_embeds_llama3}. Please make sure to only forward one of the two." - ) - - if pooled_prompt_embeds is not None and negative_pooled_prompt_embeds is not None: - if pooled_prompt_embeds.shape != negative_pooled_prompt_embeds.shape: - raise ValueError( - "`pooled_prompt_embeds` and `negative_pooled_prompt_embeds` must have the same shape when passed directly, but" - f" got: `pooled_prompt_embeds` {pooled_prompt_embeds.shape} != `negative_pooled_prompt_embeds`" - f" {negative_pooled_prompt_embeds.shape}." - ) - if prompt_embeds_t5 is not None and negative_prompt_embeds_t5 is not None: - if prompt_embeds_t5.shape != negative_prompt_embeds_t5.shape: - raise ValueError( - "`prompt_embeds_t5` and `negative_prompt_embeds_t5` must have the same shape when passed directly, but" - f" got: `prompt_embeds_t5` {prompt_embeds_t5.shape} != `negative_prompt_embeds_t5`" - f" {negative_prompt_embeds_t5.shape}." - ) - if prompt_embeds_llama3 is not None and negative_prompt_embeds_llama3 is not None: - if prompt_embeds_llama3.shape != negative_prompt_embeds_llama3.shape: - raise ValueError( - "`prompt_embeds_llama3` and `negative_prompt_embeds_llama3` must have the same shape when passed directly, but" - f" got: `prompt_embeds_llama3` {prompt_embeds_llama3.shape} != `negative_prompt_embeds_llama3`" - f" {negative_prompt_embeds_llama3.shape}." - ) - - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=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_channels_latents, height, width) - - 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) - return latents - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def do_classifier_free_guidance(self): - return self._guidance_scale > 1 - - @property - def attention_kwargs(self): - return self._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, - 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, - num_inference_steps: int = 50, - sigmas: Optional[List[float]] = 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", - 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"], - max_sequence_length: int = 128, - **kwargs, - ): - 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 `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - will be used instead. - 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*): - 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): - 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. - 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. - 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*): - 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*): - 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*): - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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.flux.FluxPipelineOutput`] 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 128): Maximum sequence length to use with the `prompt`. - - Examples: - - Returns: - [`~pipelines.hidream_image.HiDreamImagePipelineOutput`] or `tuple`: - [`~pipelines.hidream_image.HiDreamImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When - returning a tuple, the first element is a list with the generated. images. - """ - - prompt_embeds = kwargs.get("prompt_embeds", None) - negative_prompt_embeds = kwargs.get("negative_prompt_embeds", None) - - if prompt_embeds is not None: - deprecation_message = "The `prompt_embeds` argument is deprecated. Please use `prompt_embeds_t5` and `prompt_embeds_llama3` instead." - deprecate("prompt_embeds", "0.35.0", deprecation_message) - prompt_embeds_t5 = prompt_embeds[0] - prompt_embeds_llama3 = prompt_embeds[1] - - if negative_prompt_embeds is not None: - deprecation_message = "The `negative_prompt_embeds` argument is deprecated. Please use `negative_prompt_embeds_t5` and `negative_prompt_embeds_llama3` instead." - deprecate("negative_prompt_embeds", "0.35.0", deprecation_message) - negative_prompt_embeds_t5 = negative_prompt_embeds[0] - negative_prompt_embeds_llama3 = negative_prompt_embeds[1] - - height = height or self.default_sample_size * self.vae_scale_factor - width = width or self.default_sample_size * self.vae_scale_factor - - division = self.vae_scale_factor * 2 - S_max = (self.default_sample_size * self.vae_scale_factor) ** 2 - scale = S_max / (width * height) - scale = math.sqrt(scale) - width, height = int(width * scale // division * division), int(height * scale // division * division) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - prompt_3, - prompt_4, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - negative_prompt_4=negative_prompt_4, - prompt_embeds_t5=prompt_embeds_t5, - prompt_embeds_llama3=prompt_embeds_llama3, - negative_prompt_embeds_t5=negative_prompt_embeds_t5, - negative_prompt_embeds_llama3=negative_prompt_embeds_llama3, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - 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) - elif pooled_prompt_embeds is not None: - batch_size = pooled_prompt_embeds.shape[0] - - device = self._execution_device - - # 3. Encode prompt - lora_scale = self.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None - ( - prompt_embeds_t5, - negative_prompt_embeds_t5, - prompt_embeds_llama3, - negative_prompt_embeds_llama3, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_3=prompt_3, - prompt_4=prompt_4, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - negative_prompt_4=negative_prompt_4, - do_classifier_free_guidance=self.do_classifier_free_guidance, - prompt_embeds_t5=prompt_embeds_t5, - prompt_embeds_llama3=prompt_embeds_llama3, - negative_prompt_embeds_t5=negative_prompt_embeds_t5, - negative_prompt_embeds_llama3=negative_prompt_embeds_llama3, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - device=device, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - - if self.do_classifier_free_guidance: - prompt_embeds_t5 = torch.cat([negative_prompt_embeds_t5, prompt_embeds_t5], dim=0) - prompt_embeds_llama3 = torch.cat([negative_prompt_embeds_llama3, prompt_embeds_llama3], dim=1) - pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) - - # 4. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - pooled_prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 5. Prepare timesteps - mu = calculate_shift(self.transformer.max_seq) - scheduler_kwargs = {"mu": mu} - if isinstance(self.scheduler, UniPCMultistepScheduler): - self.scheduler.set_timesteps(num_inference_steps, device=device) # , shift=math.exp(mu)) - timesteps = self.scheduler.timesteps - else: - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - device, - sigmas=sigmas, - **scheduler_kwargs, - ) - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - # 6. Denoising loop - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timesteps=timestep, - encoder_hidden_states_t5=prompt_embeds_t5, - encoder_hidden_states_llama3=prompt_embeds_llama3, - pooled_embeds=pooled_prompt_embeds, - return_dict=False, - )[0] - noise_pred = -noise_pred - - # perform guidance - 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) - - # 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_t5 = callback_outputs.pop("prompt_embeds_t5", prompt_embeds_t5) - prompt_embeds_llama3 = callback_outputs.pop("prompt_embeds_llama3", prompt_embeds_llama3) - pooled_prompt_embeds = callback_outputs.pop("pooled_prompt_embeds", pooled_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: - 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 HiDreamImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/pipeline_output.py deleted file mode 100644 index 66f0f1260d1801acb7953cec3447074cb01ee9a9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/hidream_image/pipeline_output.py +++ /dev/null @@ -1,35 +0,0 @@ -# Copyright 2025 HiDream-ai 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 dataclasses import dataclass -from typing import List, Union - -import numpy as np -import PIL.Image - -from ...utils import BaseOutput - - -@dataclass -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, - num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__init__.py deleted file mode 100644 index d42d38fac9793db2c2fb94c4b41d3755be75d898..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__init__.py +++ /dev/null @@ -1,54 +0,0 @@ -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_hunyuan_skyreels_image2video"] = ["HunyuanSkyreelsImageToVideoPipeline"] - _import_structure["pipeline_hunyuan_video"] = ["HunyuanVideoPipeline"] - _import_structure["pipeline_hunyuan_video_framepack"] = ["HunyuanVideoFramepackPipeline"] - _import_structure["pipeline_hunyuan_video_image2video"] = ["HunyuanVideoImageToVideoPipeline"] - -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_hunyuan_skyreels_image2video import HunyuanSkyreelsImageToVideoPipeline - from .pipeline_hunyuan_video import HunyuanVideoPipeline - from .pipeline_hunyuan_video_framepack import HunyuanVideoFramepackPipeline - from .pipeline_hunyuan_video_image2video import HunyuanVideoImageToVideoPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 378764ed90d55b484bb58645e4630790cb427249..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_skyreels_image2video.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_skyreels_image2video.cpython-310.pyc deleted file mode 100644 index df5a5dbf6b789c117210363e0452112446a27e36..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_skyreels_image2video.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_video.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_video.cpython-310.pyc deleted file mode 100644 index 4b43d0a2790d2f15eb4bb880599f2251e3d14a63..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_video.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_video_framepack.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_video_framepack.cpython-310.pyc deleted file mode 100644 index fb6799dcd3ac5cf45f14a95c798799fcca3e51a7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_video_framepack.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_video_image2video.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_video_image2video.cpython-310.pyc deleted file mode 100644 index f6478b3cc7cee6999bb940105790467647c36f09..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_hunyuan_video_image2video.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index 02511b250a15b1cc5fd5cd96860b79eb35c313b5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py deleted file mode 100644 index d8c3548946fe3e018023016801ae902224afcbb4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py +++ /dev/null @@ -1,804 +0,0 @@ -# Copyright 2025 The HunyuanVideo 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, Tuple, Union - -import numpy as np -import torch -from transformers import CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput -from ...loaders import HunyuanVideoLoraLoaderMixin -from ...models import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel -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 .pipeline_output import HunyuanVideoPipelineOutput - - -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 HunyuanSkyreelsImageToVideoPipeline, HunyuanVideoTransformer3DModel - >>> from diffusers.utils import load_image, export_to_video - - >>> model_id = "hunyuanvideo-community/HunyuanVideo" - >>> transformer_model_id = "Skywork/SkyReels-V1-Hunyuan-I2V" - >>> transformer = HunyuanVideoTransformer3DModel.from_pretrained( - ... transformer_model_id, torch_dtype=torch.bfloat16 - ... ) - >>> pipe = HunyuanSkyreelsImageToVideoPipeline.from_pretrained( - ... model_id, transformer=transformer, torch_dtype=torch.float16 - ... ) - >>> pipe.vae.enable_tiling() - >>> pipe.to("cuda") - - >>> prompt = "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot." - >>> negative_prompt = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion" - >>> image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" - ... ) - - >>> output = pipe( - ... image=image, - ... prompt=prompt, - ... negative_prompt=negative_prompt, - ... num_inference_steps=30, - ... true_cfg_scale=6.0, - ... guidance_scale=1.0, - ... ).frames[0] - >>> export_to_video(output, "output.mp4", fps=15) - ``` -""" - - -DEFAULT_PROMPT_TEMPLATE = { - "template": ( - "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: " - "1. The main content and theme of the video." - "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects." - "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects." - "4. background environment, light, style and atmosphere." - "5. camera angles, movements, and transitions used in the video:<|eot_id|>" - "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>" - ), - "crop_start": 95, -} - - -# 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 - - -# 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" -): - 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 HunyuanSkyreelsImageToVideoPipeline(DiffusionPipeline, HunyuanVideoLoraLoaderMixin): - r""" - Pipeline for image-to-video generation using HunyuanVideo. - - 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 ([`LlamaModel`]): - [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). - tokenizer (`LlamaTokenizer`): - Tokenizer from [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). - transformer ([`HunyuanVideoTransformer3DModel`]): - Conditional Transformer to denoise the encoded image latents. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - vae ([`AutoencoderKLHunyuanVideo`]): - Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. - text_encoder_2 ([`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. - tokenizer_2 (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer). - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__( - self, - text_encoder: LlamaModel, - tokenizer: LlamaTokenizerFast, - transformer: HunyuanVideoTransformer3DModel, - vae: AutoencoderKLHunyuanVideo, - scheduler: FlowMatchEulerDiscreteScheduler, - text_encoder_2: CLIPTextModel, - tokenizer_2: CLIPTokenizer, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - text_encoder_2=text_encoder_2, - tokenizer_2=tokenizer_2, - ) - - self.vae_scale_factor_temporal = self.vae.temporal_compression_ratio if getattr(self, "vae", None) else 4 - self.vae_scale_factor_spatial = self.vae.spatial_compression_ratio if getattr(self, "vae", None) else 8 - self.vae_scaling_factor = self.vae.config.scaling_factor if getattr(self, "vae", None) else 0.476986 - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - - # 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], - num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - max_sequence_length: int = 256, - num_hidden_layers_to_skip: int = 2, - ) -> Tuple[torch.Tensor, torch.Tensor]: - 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) - - prompt = [prompt_template["template"].format(p) for p in prompt] - - crop_start = prompt_template.get("crop_start", None) - if crop_start is None: - prompt_template_input = self.tokenizer( - prompt_template["template"], - padding="max_length", - return_tensors="pt", - return_length=False, - return_overflowing_tokens=False, - return_attention_mask=False, - ) - crop_start = prompt_template_input["input_ids"].shape[-1] - # Remove <|eot_id|> token and placeholder {} - crop_start -= 2 - - max_sequence_length += crop_start - text_inputs = self.tokenizer( - prompt, - max_length=max_sequence_length, - padding="max_length", - truncation=True, - return_tensors="pt", - return_length=False, - return_overflowing_tokens=False, - return_attention_mask=True, - ) - text_input_ids = text_inputs.input_ids.to(device=device) - prompt_attention_mask = text_inputs.attention_mask.to(device=device) - - prompt_embeds = self.text_encoder( - input_ids=text_input_ids, - attention_mask=prompt_attention_mask, - output_hidden_states=True, - ).hidden_states[-(num_hidden_layers_to_skip + 1)] - prompt_embeds = prompt_embeds.to(dtype=dtype) - - if crop_start is not None and crop_start > 0: - prompt_embeds = prompt_embeds[:, crop_start:] - prompt_attention_mask = prompt_attention_mask[:, crop_start:] - - # 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.repeat(1, num_videos_per_prompt) - prompt_attention_mask = prompt_attention_mask.view(batch_size * num_videos_per_prompt, seq_len) - - return prompt_embeds, prompt_attention_mask - - # 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]], - num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - max_sequence_length: int = 77, - ) -> torch.Tensor: - device = device or self._execution_device - dtype = dtype or self.text_encoder_2.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - text_inputs = self.tokenizer_2( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - ) - - text_input_ids = text_inputs.input_ids - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {max_sequence_length} tokens: {removed_text}" - ) - - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False).pooler_output - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, -1) - - return 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, - 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, - max_sequence_length: int = 256, - ): - if prompt_embeds is None: - prompt_embeds, prompt_attention_mask = self._get_llama_prompt_embeds( - prompt, - prompt_template, - num_videos_per_prompt, - device=device, - dtype=dtype, - max_sequence_length=max_sequence_length, - ) - - if pooled_prompt_embeds is None: - if prompt_2 is None: - prompt_2 = prompt - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt, - num_videos_per_prompt, - device=device, - dtype=dtype, - max_sequence_length=77, - ) - - return prompt_embeds, pooled_prompt_embeds, prompt_attention_mask - - # Copied from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.HunyuanVideoPipeline.check_inputs - def check_inputs( - self, - prompt, - prompt_2, - height, - width, - prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - prompt_template=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 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 prompt_template is not None: - if not isinstance(prompt_template, dict): - raise ValueError(f"`prompt_template` has to be of type `dict` but is {type(prompt_template)}") - if "template" not in prompt_template: - raise ValueError( - f"`prompt_template` has to contain a key `template` but only found {prompt_template.keys()}" - ) - - def prepare_latents( - self, - image: torch.Tensor, - batch_size: int, - num_channels_latents: int = 32, - 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, - ) -> 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." - ) - - image = image.unsqueeze(2) # [B, C, 1, H, W] - if isinstance(generator, list): - image_latents = [ - retrieve_latents(self.vae.encode(image[i].unsqueeze(0)), generator[i]) for i in range(batch_size) - ] - else: - image_latents = [retrieve_latents(self.vae.encode(img.unsqueeze(0)), generator) for img in image] - - image_latents = torch.cat(image_latents, dim=0).to(dtype) * self.vae_scaling_factor - - num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - latent_height, latent_width = height // self.vae_scale_factor_spatial, width // self.vae_scale_factor_spatial - shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width) - padding_shape = (batch_size, num_channels_latents, num_latent_frames - 1, latent_height, latent_width) - - latents_padding = torch.zeros(padding_shape, dtype=dtype, device=device) - image_latents = torch.cat([image_latents, latents_padding], dim=2) - - if latents is None: - latents = randn_tensor(shape, generator=generator, dtype=dtype, device=device) - else: - latents = latents.to(dtype=dtype, device=device) - - return latents, image_latents - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - @property - def guidance_scale(self): - return self._guidance_scale - - @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, - 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, - height: int = 544, - width: int = 960, - num_frames: int = 97, - num_inference_steps: int = 50, - 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", - 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, - max_sequence_length: int = 256, - ): - 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, one has to pass `prompt_embeds`. - instead. - 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*): - 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*): - 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`): - The height in pixels of the generated image. - width (`int`, defaults to `1280`): - The width in pixels of the generated image. - num_frames (`int`, defaults to `129`): - 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. - 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. - true_cfg_scale (`float`, *optional*, defaults to 1.0): - When > 1.0 and a provided `negative_prompt`, enables true classifier-free guidance. - guidance_scale (`float`, defaults to `6.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. Note that the only available - HunyuanVideo model is CFG-distilled, which means that traditional guidance between unconditional and - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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_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. - negative_pooled_prompt_embeds (`torch.FloatTensor`, *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 generated image. Choose between `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`HunyuanVideoPipelineOutput`] 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). - 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. - 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. - - Examples: - - Returns: - [`~HunyuanVideoPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] 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 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_2, - height, - width, - prompt_embeds, - callback_on_step_end_tensor_inputs, - prompt_template, - ) - - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None - ) - do_true_cfg = true_cfg_scale > 1 and has_neg_prompt - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._current_timestep = None - self._interrupt = False - - device = self._execution_device - - # 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 - transformer_dtype = self.transformer.dtype - prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_template=prompt_template, - num_videos_per_prompt=num_videos_per_prompt, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - device=device, - max_sequence_length=max_sequence_length, - ) - prompt_embeds = prompt_embeds.to(transformer_dtype) - prompt_attention_mask = prompt_attention_mask.to(transformer_dtype) - pooled_prompt_embeds = pooled_prompt_embeds.to(transformer_dtype) - - if do_true_cfg: - negative_prompt_embeds, negative_pooled_prompt_embeds, negative_prompt_attention_mask = self.encode_prompt( - prompt=negative_prompt, - prompt_2=negative_prompt_2, - prompt_template=prompt_template, - num_videos_per_prompt=num_videos_per_prompt, - prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=negative_pooled_prompt_embeds, - prompt_attention_mask=negative_prompt_attention_mask, - device=device, - max_sequence_length=max_sequence_length, - ) - negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype) - negative_prompt_attention_mask = negative_prompt_attention_mask.to(transformer_dtype) - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.to(transformer_dtype) - - # 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) - - # 5. Prepare latent variables - vae_dtype = self.vae.dtype - image = self.video_processor.preprocess(image, height=height, width=width).to(device, vae_dtype) - num_channels_latents = self.transformer.config.in_channels // 2 - latents, image_latents = self.prepare_latents( - image, - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - num_frames, - torch.float32, - device, - generator, - latents, - ) - latent_image_input = image_latents.to(transformer_dtype) - - # 6. Prepare guidance condition - guidance = torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0 - - # 7. Denoising loop - 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 - latent_model_input = latents.to(transformer_dtype) - latent_model_input = torch.cat([latent_model_input, latent_image_input], dim=1) - - # 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=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - encoder_attention_mask=prompt_attention_mask, - pooled_projections=pooled_prompt_embeds, - guidance=guidance, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - - if do_true_cfg: - neg_noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, - encoder_attention_mask=negative_prompt_attention_mask, - pooled_projections=negative_pooled_prompt_embeds, - guidance=guidance, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) - - # compute the previous noisy sample x_t -> x_t-1 - 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, 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 not output_type == "latent": - latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor - video = self.vae.decode(latents, return_dict=False)[0] - video = self.video_processor.postprocess_video(video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return HunyuanVideoPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py deleted file mode 100644 index b617e4f8b26a340f7b6567129783dfc29463dadd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py +++ /dev/null @@ -1,754 +0,0 @@ -# Copyright 2025 The HunyuanVideo 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, Tuple, Union - -import numpy as np -import torch -from transformers import CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...loaders import HunyuanVideoLoraLoaderMixin -from ...models import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel -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 .pipeline_output import HunyuanVideoPipelineOutput - - -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 HunyuanVideoPipeline, HunyuanVideoTransformer3DModel - >>> from diffusers.utils import export_to_video - - >>> 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, torch_dtype=torch.float16) - >>> pipe.vae.enable_tiling() - >>> pipe.to("cuda") - - >>> output = pipe( - ... prompt="A cat walks on the grass, realistic", - ... height=320, - ... width=512, - ... num_frames=61, - ... num_inference_steps=30, - ... ).frames[0] - >>> export_to_video(output, "output.mp4", fps=15) - ``` -""" - - -DEFAULT_PROMPT_TEMPLATE = { - "template": ( - "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: " - "1. The main content and theme of the video." - "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects." - "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects." - "4. background environment, light, style and atmosphere." - "5. camera angles, movements, and transitions used in the video:<|eot_id|>" - "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>" - ), - "crop_start": 95, -} - - -# 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 HunyuanVideoPipeline(DiffusionPipeline, HunyuanVideoLoraLoaderMixin): - r""" - Pipeline for text-to-video generation using HunyuanVideo. - - 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 ([`LlamaModel`]): - [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). - tokenizer (`LlamaTokenizer`): - Tokenizer from [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). - transformer ([`HunyuanVideoTransformer3DModel`]): - Conditional Transformer to denoise the encoded image latents. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - vae ([`AutoencoderKLHunyuanVideo`]): - Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. - text_encoder_2 ([`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. - tokenizer_2 (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer). - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__( - self, - text_encoder: LlamaModel, - tokenizer: LlamaTokenizerFast, - transformer: HunyuanVideoTransformer3DModel, - vae: AutoencoderKLHunyuanVideo, - scheduler: FlowMatchEulerDiscreteScheduler, - text_encoder_2: CLIPTextModel, - tokenizer_2: CLIPTokenizer, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - text_encoder_2=text_encoder_2, - tokenizer_2=tokenizer_2, - ) - - self.vae_scale_factor_temporal = self.vae.temporal_compression_ratio if getattr(self, "vae", None) else 4 - self.vae_scale_factor_spatial = self.vae.spatial_compression_ratio if getattr(self, "vae", None) else 8 - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - - def _get_llama_prompt_embeds( - self, - prompt: Union[str, List[str]], - prompt_template: Dict[str, Any], - num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - max_sequence_length: int = 256, - num_hidden_layers_to_skip: int = 2, - ) -> Tuple[torch.Tensor, torch.Tensor]: - 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) - - prompt = [prompt_template["template"].format(p) for p in prompt] - - crop_start = prompt_template.get("crop_start", None) - if crop_start is None: - prompt_template_input = self.tokenizer( - prompt_template["template"], - padding="max_length", - return_tensors="pt", - return_length=False, - return_overflowing_tokens=False, - return_attention_mask=False, - ) - crop_start = prompt_template_input["input_ids"].shape[-1] - # Remove <|eot_id|> token and placeholder {} - crop_start -= 2 - - max_sequence_length += crop_start - text_inputs = self.tokenizer( - prompt, - max_length=max_sequence_length, - padding="max_length", - truncation=True, - return_tensors="pt", - return_length=False, - return_overflowing_tokens=False, - return_attention_mask=True, - ) - text_input_ids = text_inputs.input_ids.to(device=device) - prompt_attention_mask = text_inputs.attention_mask.to(device=device) - - prompt_embeds = self.text_encoder( - input_ids=text_input_ids, - attention_mask=prompt_attention_mask, - output_hidden_states=True, - ).hidden_states[-(num_hidden_layers_to_skip + 1)] - prompt_embeds = prompt_embeds.to(dtype=dtype) - - if crop_start is not None and crop_start > 0: - prompt_embeds = prompt_embeds[:, crop_start:] - prompt_attention_mask = prompt_attention_mask[:, crop_start:] - - # 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.repeat(1, num_videos_per_prompt) - prompt_attention_mask = prompt_attention_mask.view(batch_size * num_videos_per_prompt, seq_len) - - return prompt_embeds, prompt_attention_mask - - def _get_clip_prompt_embeds( - self, - prompt: Union[str, List[str]], - num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - max_sequence_length: int = 77, - ) -> torch.Tensor: - device = device or self._execution_device - dtype = dtype or self.text_encoder_2.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - text_inputs = self.tokenizer_2( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - ) - - text_input_ids = text_inputs.input_ids - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {max_sequence_length} tokens: {removed_text}" - ) - - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False).pooler_output - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, -1) - - return 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, - 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, - max_sequence_length: int = 256, - ): - if prompt_embeds is None: - prompt_embeds, prompt_attention_mask = self._get_llama_prompt_embeds( - prompt, - prompt_template, - num_videos_per_prompt, - device=device, - dtype=dtype, - max_sequence_length=max_sequence_length, - ) - - if pooled_prompt_embeds is None: - if prompt_2 is None: - prompt_2 = prompt - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt, - num_videos_per_prompt, - device=device, - dtype=dtype, - max_sequence_length=77, - ) - - return prompt_embeds, pooled_prompt_embeds, prompt_attention_mask - - def check_inputs( - self, - prompt, - prompt_2, - height, - width, - prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - prompt_template=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 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 prompt_template is not None: - if not isinstance(prompt_template, dict): - raise ValueError(f"`prompt_template` has to be of type `dict` but is {type(prompt_template)}") - if "template" not in prompt_template: - raise ValueError( - f"`prompt_template` has to contain a key `template` but only found {prompt_template.keys()}" - ) - - def prepare_latents( - self, - batch_size: int, - num_channels_latents: int = 32, - 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, - ) -> torch.Tensor: - if latents is not None: - return latents.to(device=device, dtype=dtype) - - shape = ( - batch_size, - num_channels_latents, - (num_frames - 1) // self.vae_scale_factor_temporal + 1, - 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 enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - @property - def guidance_scale(self): - return self._guidance_scale - - @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: 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, - height: int = 720, - width: int = 1280, - num_frames: int = 129, - num_inference_steps: int = 50, - 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", - 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, - max_sequence_length: int = 256, - ): - 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, one has to pass `prompt_embeds`. - instead. - 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*): - 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*): - 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`): - The height in pixels of the generated image. - width (`int`, defaults to `1280`): - The width in pixels of the generated image. - num_frames (`int`, defaults to `129`): - 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. - 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. - true_cfg_scale (`float`, *optional*, defaults to 1.0): - When > 1.0 and a provided `negative_prompt`, enables true classifier-free guidance. - guidance_scale (`float`, defaults to `6.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. Note that the only available - HunyuanVideo model is CFG-distilled, which means that traditional guidance between unconditional and - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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_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. - negative_pooled_prompt_embeds (`torch.FloatTensor`, *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 generated image. Choose between `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`HunyuanVideoPipelineOutput`] 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). - 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. - 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. - - Examples: - - Returns: - [`~HunyuanVideoPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] 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 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_2, - height, - width, - prompt_embeds, - callback_on_step_end_tensor_inputs, - prompt_template, - ) - - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None - ) - do_true_cfg = true_cfg_scale > 1 and has_neg_prompt - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._current_timestep = None - self._interrupt = False - - device = self._execution_device - - # 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 - transformer_dtype = self.transformer.dtype - prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_template=prompt_template, - num_videos_per_prompt=num_videos_per_prompt, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - device=device, - max_sequence_length=max_sequence_length, - ) - prompt_embeds = prompt_embeds.to(transformer_dtype) - prompt_attention_mask = prompt_attention_mask.to(transformer_dtype) - pooled_prompt_embeds = pooled_prompt_embeds.to(transformer_dtype) - - if do_true_cfg: - negative_prompt_embeds, negative_pooled_prompt_embeds, negative_prompt_attention_mask = self.encode_prompt( - prompt=negative_prompt, - prompt_2=negative_prompt_2, - prompt_template=prompt_template, - num_videos_per_prompt=num_videos_per_prompt, - prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=negative_pooled_prompt_embeds, - prompt_attention_mask=negative_prompt_attention_mask, - device=device, - max_sequence_length=max_sequence_length, - ) - negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype) - negative_prompt_attention_mask = negative_prompt_attention_mask.to(transformer_dtype) - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.to(transformer_dtype) - - # 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) - - # 5. Prepare latent variables - 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, - torch.float32, - device, - generator, - latents, - ) - - # 6. Prepare guidance condition - guidance = torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0 - - # 7. Denoising loop - 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 - latent_model_input = latents.to(transformer_dtype) - # 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=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - encoder_attention_mask=prompt_attention_mask, - pooled_projections=pooled_prompt_embeds, - guidance=guidance, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - - if do_true_cfg: - neg_noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, - encoder_attention_mask=negative_prompt_attention_mask, - pooled_projections=negative_pooled_prompt_embeds, - guidance=guidance, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) - - # compute the previous noisy sample x_t -> x_t-1 - 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, 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 not output_type == "latent": - latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor - video = self.vae.decode(latents, return_dict=False)[0] - video = self.video_processor.postprocess_video(video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return HunyuanVideoPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py deleted file mode 100644 index 40d6534655d604cd607e9406d195ed0834221919..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py +++ /dev/null @@ -1,1114 +0,0 @@ -# Copyright 2025 The Framepack Team, The HunyuanVideo 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 -import math -from enum import Enum -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import torch -from transformers import ( - CLIPTextModel, - CLIPTokenizer, - LlamaModel, - LlamaTokenizerFast, - SiglipImageProcessor, - SiglipVisionModel, -) - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput -from ...loaders import HunyuanVideoLoraLoaderMixin -from ...models import AutoencoderKLHunyuanVideo, HunyuanVideoFramepackTransformer3DModel -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 .pipeline_output import HunyuanVideoFramepackPipelineOutput - - -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 - - -# TODO(yiyi): We can pack the checkpoints nicely with modular loader -EXAMPLE_DOC_STRING = """ - Examples: - ##### Image-to-Video - - ```python - >>> import torch - >>> from diffusers import HunyuanVideoFramepackPipeline, HunyuanVideoFramepackTransformer3DModel - >>> from diffusers.utils import export_to_video, load_image - >>> from transformers import SiglipImageProcessor, SiglipVisionModel - - >>> transformer = HunyuanVideoFramepackTransformer3DModel.from_pretrained( - ... "lllyasviel/FramePackI2V_HY", torch_dtype=torch.bfloat16 - ... ) - >>> feature_extractor = SiglipImageProcessor.from_pretrained( - ... "lllyasviel/flux_redux_bfl", subfolder="feature_extractor" - ... ) - >>> image_encoder = SiglipVisionModel.from_pretrained( - ... "lllyasviel/flux_redux_bfl", subfolder="image_encoder", torch_dtype=torch.float16 - ... ) - >>> pipe = HunyuanVideoFramepackPipeline.from_pretrained( - ... "hunyuanvideo-community/HunyuanVideo", - ... transformer=transformer, - ... feature_extractor=feature_extractor, - ... image_encoder=image_encoder, - ... torch_dtype=torch.float16, - ... ) - >>> pipe.vae.enable_tiling() - >>> pipe.to("cuda") - - >>> image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png" - ... ) - >>> output = pipe( - ... image=image, - ... prompt="A penguin dancing in the snow", - ... height=832, - ... width=480, - ... num_frames=91, - ... num_inference_steps=30, - ... guidance_scale=9.0, - ... generator=torch.Generator().manual_seed(0), - ... sampling_type="inverted_anti_drifting", - ... ).frames[0] - >>> export_to_video(output, "output.mp4", fps=30) - ``` - - ##### First and Last Image-to-Video - - ```python - >>> import torch - >>> from diffusers import HunyuanVideoFramepackPipeline, HunyuanVideoFramepackTransformer3DModel - >>> from diffusers.utils import export_to_video, load_image - >>> from transformers import SiglipImageProcessor, SiglipVisionModel - - >>> transformer = HunyuanVideoFramepackTransformer3DModel.from_pretrained( - ... "lllyasviel/FramePackI2V_HY", torch_dtype=torch.bfloat16 - ... ) - >>> feature_extractor = SiglipImageProcessor.from_pretrained( - ... "lllyasviel/flux_redux_bfl", subfolder="feature_extractor" - ... ) - >>> image_encoder = SiglipVisionModel.from_pretrained( - ... "lllyasviel/flux_redux_bfl", subfolder="image_encoder", torch_dtype=torch.float16 - ... ) - >>> pipe = HunyuanVideoFramepackPipeline.from_pretrained( - ... "hunyuanvideo-community/HunyuanVideo", - ... transformer=transformer, - ... feature_extractor=feature_extractor, - ... image_encoder=image_encoder, - ... torch_dtype=torch.float16, - ... ) - >>> pipe.to("cuda") - - >>> 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" - ... ) - >>> output = pipe( - ... image=first_image, - ... last_image=last_image, - ... prompt=prompt, - ... height=512, - ... width=512, - ... num_frames=91, - ... num_inference_steps=30, - ... guidance_scale=9.0, - ... generator=torch.Generator().manual_seed(0), - ... sampling_type="inverted_anti_drifting", - ... ).frames[0] - >>> export_to_video(output, "output.mp4", fps=30) - ``` -""" - - -DEFAULT_PROMPT_TEMPLATE = { - "template": ( - "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: " - "1. The main content and theme of the video." - "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects." - "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects." - "4. background environment, light, style and atmosphere." - "5. camera angles, movements, and transitions used in the video:<|eot_id|>" - "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>" - ), - "crop_start": 95, -} - - -# 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 FramepackSamplingType(str, Enum): - VANILLA = "vanilla" - INVERTED_ANTI_DRIFTING = "inverted_anti_drifting" - - -class HunyuanVideoFramepackPipeline(DiffusionPipeline, HunyuanVideoLoraLoaderMixin): - r""" - Pipeline for text-to-video generation using HunyuanVideo. - - 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 ([`LlamaModel`]): - [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). - tokenizer (`LlamaTokenizer`): - Tokenizer from [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). - transformer ([`HunyuanVideoTransformer3DModel`]): - Conditional Transformer to denoise the encoded image latents. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - vae ([`AutoencoderKLHunyuanVideo`]): - Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. - text_encoder_2 ([`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. - tokenizer_2 (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer). - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__( - self, - text_encoder: LlamaModel, - tokenizer: LlamaTokenizerFast, - transformer: HunyuanVideoFramepackTransformer3DModel, - vae: AutoencoderKLHunyuanVideo, - scheduler: FlowMatchEulerDiscreteScheduler, - text_encoder_2: CLIPTextModel, - tokenizer_2: CLIPTokenizer, - image_encoder: SiglipVisionModel, - feature_extractor: SiglipImageProcessor, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - text_encoder_2=text_encoder_2, - tokenizer_2=tokenizer_2, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - ) - - self.vae_scale_factor_temporal = self.vae.temporal_compression_ratio if getattr(self, "vae", None) else 4 - self.vae_scale_factor_spatial = self.vae.spatial_compression_ratio if getattr(self, "vae", None) else 8 - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - - # 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], - num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - max_sequence_length: int = 256, - num_hidden_layers_to_skip: int = 2, - ) -> Tuple[torch.Tensor, torch.Tensor]: - 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) - - prompt = [prompt_template["template"].format(p) for p in prompt] - - crop_start = prompt_template.get("crop_start", None) - if crop_start is None: - prompt_template_input = self.tokenizer( - prompt_template["template"], - padding="max_length", - return_tensors="pt", - return_length=False, - return_overflowing_tokens=False, - return_attention_mask=False, - ) - crop_start = prompt_template_input["input_ids"].shape[-1] - # Remove <|eot_id|> token and placeholder {} - crop_start -= 2 - - max_sequence_length += crop_start - text_inputs = self.tokenizer( - prompt, - max_length=max_sequence_length, - padding="max_length", - truncation=True, - return_tensors="pt", - return_length=False, - return_overflowing_tokens=False, - return_attention_mask=True, - ) - text_input_ids = text_inputs.input_ids.to(device=device) - prompt_attention_mask = text_inputs.attention_mask.to(device=device) - - prompt_embeds = self.text_encoder( - input_ids=text_input_ids, - attention_mask=prompt_attention_mask, - output_hidden_states=True, - ).hidden_states[-(num_hidden_layers_to_skip + 1)] - prompt_embeds = prompt_embeds.to(dtype=dtype) - - if crop_start is not None and crop_start > 0: - prompt_embeds = prompt_embeds[:, crop_start:] - prompt_attention_mask = prompt_attention_mask[:, crop_start:] - - # 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.repeat(1, num_videos_per_prompt) - prompt_attention_mask = prompt_attention_mask.view(batch_size * num_videos_per_prompt, seq_len) - - return prompt_embeds, prompt_attention_mask - - # 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]], - num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - max_sequence_length: int = 77, - ) -> torch.Tensor: - device = device or self._execution_device - dtype = dtype or self.text_encoder_2.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - text_inputs = self.tokenizer_2( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - ) - - text_input_ids = text_inputs.input_ids - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {max_sequence_length} tokens: {removed_text}" - ) - - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False).pooler_output - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, -1) - - return 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, - 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, - max_sequence_length: int = 256, - ): - if prompt_embeds is None: - prompt_embeds, prompt_attention_mask = self._get_llama_prompt_embeds( - prompt, - prompt_template, - num_videos_per_prompt, - device=device, - dtype=dtype, - max_sequence_length=max_sequence_length, - ) - - if pooled_prompt_embeds is None: - if prompt_2 is None: - prompt_2 = prompt - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt, - num_videos_per_prompt, - device=device, - dtype=dtype, - max_sequence_length=77, - ) - - 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 - ): - 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( - device=device, dtype=self.image_encoder.dtype - ) - image_embeds = self.image_encoder(**image).last_hidden_state - return image_embeds.to(dtype=dtype) - - def check_inputs( - self, - prompt, - prompt_2, - height, - width, - prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - prompt_template=None, - image=None, - image_latents=None, - last_image=None, - last_image_latents=None, - sampling_type=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 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 prompt_template is not None: - if not isinstance(prompt_template, dict): - raise ValueError(f"`prompt_template` has to be of type `dict` but is {type(prompt_template)}") - if "template" not in prompt_template: - raise ValueError( - f"`prompt_template` has to contain a key `template` but only found {prompt_template.keys()}" - ) - - sampling_types = [x.value for x in FramepackSamplingType.__members__.values()] - if sampling_type not in sampling_types: - raise ValueError(f"`sampling_type` has to be one of '{sampling_types}' but is '{sampling_type}'") - - if image is not None and image_latents is not None: - raise ValueError("Only one of `image` or `image_latents` can be passed.") - if last_image is not None and last_image_latents is not None: - raise ValueError("Only one of `last_image` or `last_image_latents` can be passed.") - if sampling_type != FramepackSamplingType.INVERTED_ANTI_DRIFTING and ( - last_image is not None or last_image_latents is not None - ): - raise ValueError( - 'Only `"inverted_anti_drifting"` inference type supports `last_image` or `last_image_latents`.' - ) - - def prepare_latents( - self, - batch_size: int = 1, - num_channels_latents: int = 16, - 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, - ) -> torch.Tensor: - if latents is not None: - return latents.to(device=device, dtype=dtype) - shape = ( - batch_size, - num_channels_latents, - (num_frames - 1) // self.vae_scale_factor_temporal + 1, - 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, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = 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 * self.vae.config.scaling_factor - return latents.to(device=device, dtype=dtype) - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - @property - def guidance_scale(self): - return self._guidance_scale - - @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, - 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, - height: int = 720, - width: int = 1280, - num_frames: int = 129, - latent_window_size: int = 9, - num_inference_steps: int = 50, - 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", - 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, - max_sequence_length: int = 256, - sampling_type: FramepackSamplingType = FramepackSamplingType.INVERTED_ANTI_DRIFTING, - ): - r""" - The call function to the pipeline for generation. - - Args: - image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`): - The image to be used as the starting point for the video generation. - 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*): - 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 `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - will be used 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 `true_cfg_scale` is - not greater 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 all the text-encoders. - height (`int`, defaults to `720`): - The height in pixels of the generated image. - width (`int`, defaults to `1280`): - The width in pixels of the generated image. - num_frames (`int`, defaults to `129`): - 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. - 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. - true_cfg_scale (`float`, *optional*, defaults to 1.0): - When > 1.0 and a provided `negative_prompt`, enables true classifier-free guidance. - guidance_scale (`float`, defaults to `6.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. Note that the only available - HunyuanVideo model is CFG-distilled, which means that traditional guidance between unconditional and - 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*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - image_latents (`torch.Tensor`, *optional*): - Pre-encoded image latents. If not provided, the image will be encoded using the VAE. - last_image_latents (`torch.Tensor`, *optional*): - Pre-encoded last image latents. If not provided, the last image will be encoded using the VAE. - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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_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. - negative_pooled_prompt_embeds (`torch.FloatTensor`, *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 generated image. Choose between `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`HunyuanVideoFramepackPipelineOutput`] 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). - 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. - 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. - - Examples: - - Returns: - [`~HunyuanVideoFramepackPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`HunyuanVideoFramepackPipelineOutput`] 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 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_2, - height, - width, - prompt_embeds, - callback_on_step_end_tensor_inputs, - prompt_template, - image, - image_latents, - last_image, - last_image_latents, - sampling_type, - ) - - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None - ) - do_true_cfg = true_cfg_scale > 1 and has_neg_prompt - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._current_timestep = None - self._interrupt = False - - device = self._execution_device - transformer_dtype = self.transformer.dtype - vae_dtype = 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 - transformer_dtype = self.transformer.dtype - prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_template=prompt_template, - num_videos_per_prompt=num_videos_per_prompt, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - device=device, - max_sequence_length=max_sequence_length, - ) - prompt_embeds = prompt_embeds.to(transformer_dtype) - prompt_attention_mask = prompt_attention_mask.to(transformer_dtype) - pooled_prompt_embeds = pooled_prompt_embeds.to(transformer_dtype) - - if do_true_cfg: - negative_prompt_embeds, negative_pooled_prompt_embeds, negative_prompt_attention_mask = self.encode_prompt( - prompt=negative_prompt, - prompt_2=negative_prompt_2, - prompt_template=prompt_template, - num_videos_per_prompt=num_videos_per_prompt, - prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=negative_pooled_prompt_embeds, - prompt_attention_mask=negative_prompt_attention_mask, - device=device, - max_sequence_length=max_sequence_length, - ) - negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype) - negative_prompt_attention_mask = negative_prompt_attention_mask.to(transformer_dtype) - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.to(transformer_dtype) - - # 4. Prepare image - image = self.video_processor.preprocess(image, height, width) - image_embeds = self.encode_image(image, device=device).to(transformer_dtype) - if last_image is not None: - # Credits: https://github.com/lllyasviel/FramePack/pull/167 - # Users can modify the weighting strategy applied here - last_image = self.video_processor.preprocess(last_image, height, width) - last_image_embeds = self.encode_image(last_image, device=device).to(transformer_dtype) - last_image_embeds = (image_embeds + last_image_embeds) / 2 - - # 5. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - window_num_frames = (latent_window_size - 1) * self.vae_scale_factor_temporal + 1 - num_latent_sections = max(1, (num_frames + window_num_frames - 1) // window_num_frames) - history_video = None - total_generated_latent_frames = 0 - - image_latents = self.prepare_image_latents( - image, dtype=torch.float32, device=device, generator=generator, latents=image_latents - ) - if last_image is not None: - last_image_latents = self.prepare_image_latents( - last_image, dtype=torch.float32, device=device, generator=generator - ) - - # Specific to the released checkpoints: - # - https://huggingface.co/lllyasviel/FramePackI2V_HY - # - https://huggingface.co/lllyasviel/FramePack_F1_I2V_HY_20250503 - # TODO: find a more generic way in future if there are more checkpoints - if sampling_type == FramepackSamplingType.INVERTED_ANTI_DRIFTING: - history_sizes = [1, 2, 16] - history_latents = torch.zeros( - batch_size, - num_channels_latents, - sum(history_sizes), - height // self.vae_scale_factor_spatial, - width // self.vae_scale_factor_spatial, - device=device, - dtype=torch.float32, - ) - - elif sampling_type == FramepackSamplingType.VANILLA: - history_sizes = [16, 2, 1] - history_latents = torch.zeros( - batch_size, - num_channels_latents, - sum(history_sizes), - height // self.vae_scale_factor_spatial, - width // self.vae_scale_factor_spatial, - device=device, - dtype=torch.float32, - ) - history_latents = torch.cat([history_latents, image_latents], dim=2) - total_generated_latent_frames += 1 - - else: - assert False - - # 6. Prepare guidance condition - guidance = torch.tensor([guidance_scale] * batch_size, dtype=transformer_dtype, device=device) * 1000.0 - - # 7. Denoising loop - for k in range(num_latent_sections): - if sampling_type == FramepackSamplingType.INVERTED_ANTI_DRIFTING: - latent_paddings = list(reversed(range(num_latent_sections))) - if num_latent_sections > 4: - latent_paddings = [3] + [2] * (num_latent_sections - 3) + [1, 0] - - is_first_section = k == 0 - is_last_section = k == num_latent_sections - 1 - latent_padding_size = latent_paddings[k] * latent_window_size - - indices = torch.arange(0, sum([1, latent_padding_size, latent_window_size, *history_sizes])) - ( - indices_prefix, - indices_padding, - indices_latents, - indices_latents_history_1x, - indices_latents_history_2x, - indices_latents_history_4x, - ) = indices.split([1, latent_padding_size, latent_window_size, *history_sizes], dim=0) - # Inverted anti-drifting sampling: Figure 2(c) in the paper - indices_clean_latents = torch.cat([indices_prefix, indices_latents_history_1x], dim=0) - - latents_prefix = image_latents - latents_history_1x, latents_history_2x, latents_history_4x = history_latents[ - :, :, : sum(history_sizes) - ].split(history_sizes, dim=2) - if last_image is not None and is_first_section: - latents_history_1x = last_image_latents - latents_clean = torch.cat([latents_prefix, latents_history_1x], dim=2) - - elif sampling_type == FramepackSamplingType.VANILLA: - indices = torch.arange(0, sum([1, *history_sizes, latent_window_size])) - ( - indices_prefix, - indices_latents_history_4x, - indices_latents_history_2x, - indices_latents_history_1x, - indices_latents, - ) = indices.split([1, *history_sizes, latent_window_size], dim=0) - indices_clean_latents = torch.cat([indices_prefix, indices_latents_history_1x], dim=0) - - latents_prefix = image_latents - latents_history_4x, latents_history_2x, latents_history_1x = history_latents[ - :, :, -sum(history_sizes) : - ].split(history_sizes, dim=2) - latents_clean = torch.cat([latents_prefix, latents_history_1x], dim=2) - - else: - assert False - - latents = self.prepare_latents( - batch_size, - num_channels_latents, - height, - width, - window_num_frames, - dtype=torch.float32, - device=device, - generator=generator, - latents=None, - ) - - sigmas = np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1] if sigmas is None else sigmas - image_seq_len = ( - latents.shape[2] * latents.shape[3] * latents.shape[4] / self.transformer.config.patch_size**2 - ) - exp_max = 7.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), - ) - mu = min(mu, math.log(exp_max)) - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu - ) - 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]) - - noise_pred = self.transformer( - hidden_states=latents.to(transformer_dtype), - timestep=timestep, - encoder_hidden_states=prompt_embeds, - encoder_attention_mask=prompt_attention_mask, - pooled_projections=pooled_prompt_embeds, - image_embeds=image_embeds, - indices_latents=indices_latents, - guidance=guidance, - latents_clean=latents_clean.to(transformer_dtype), - indices_latents_clean=indices_clean_latents, - latents_history_2x=latents_history_2x.to(transformer_dtype), - indices_latents_history_2x=indices_latents_history_2x, - latents_history_4x=latents_history_4x.to(transformer_dtype), - indices_latents_history_4x=indices_latents_history_4x, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - - if do_true_cfg: - neg_noise_pred = self.transformer( - hidden_states=latents.to(transformer_dtype), - timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, - encoder_attention_mask=negative_prompt_attention_mask, - pooled_projections=negative_pooled_prompt_embeds, - image_embeds=image_embeds, - indices_latents=indices_latents, - guidance=guidance, - latents_clean=latents_clean.to(transformer_dtype), - indices_latents_clean=indices_clean_latents, - latents_history_2x=latents_history_2x.to(transformer_dtype), - indices_latents_history_2x=indices_latents_history_2x, - latents_history_4x=latents_history_4x.to(transformer_dtype), - indices_latents_history_4x=indices_latents_history_4x, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred.float(), 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, 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() - - if sampling_type == FramepackSamplingType.INVERTED_ANTI_DRIFTING: - if is_last_section: - latents = torch.cat([image_latents, latents], dim=2) - total_generated_latent_frames += latents.shape[2] - history_latents = torch.cat([latents, history_latents], dim=2) - real_history_latents = history_latents[:, :, :total_generated_latent_frames] - section_latent_frames = ( - (latent_window_size * 2 + 1) if is_last_section else (latent_window_size * 2) - ) - index_slice = (slice(None), slice(None), slice(0, section_latent_frames)) - - elif sampling_type == FramepackSamplingType.VANILLA: - 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:] - section_latent_frames = latent_window_size * 2 - index_slice = (slice(None), slice(None), slice(-section_latent_frames, None)) - - else: - assert False - - if history_video is None: - if not output_type == "latent": - current_latents = real_history_latents.to(vae_dtype) / self.vae.config.scaling_factor - history_video = self.vae.decode(current_latents, return_dict=False)[0] - else: - history_video = [real_history_latents] - else: - if not output_type == "latent": - overlapped_frames = (latent_window_size - 1) * self.vae_scale_factor_temporal + 1 - current_latents = ( - real_history_latents[index_slice].to(vae_dtype) / self.vae.config.scaling_factor - ) - current_video = self.vae.decode(current_latents, return_dict=False)[0] - - if sampling_type == FramepackSamplingType.INVERTED_ANTI_DRIFTING: - history_video = self._soft_append(current_video, history_video, overlapped_frames) - elif sampling_type == FramepackSamplingType.VANILLA: - history_video = self._soft_append(history_video, current_video, overlapped_frames) - else: - assert False - else: - history_video.append(real_history_latents) - - self._current_timestep = None - - if not 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 = history_video - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return HunyuanVideoFramepackPipelineOutput(frames=video) - - def _soft_append(self, history: torch.Tensor, current: torch.Tensor, overlap: int = 0): - if overlap <= 0: - return torch.cat([history, current], dim=2) - - assert history.shape[2] >= overlap, f"Current length ({history.shape[2]}) must be >= overlap ({overlap})" - assert current.shape[2] >= overlap, f"History length ({current.shape[2]}) must be >= overlap ({overlap})" - - weights = torch.linspace(1, 0, overlap, dtype=history.dtype, device=history.device).view(1, 1, -1, 1, 1) - blended = weights * history[:, :, -overlap:] + (1 - weights) * current[:, :, :overlap] - output = torch.cat([history[:, :, :-overlap], blended, current[:, :, overlap:]], dim=2) - - return output.to(history) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py deleted file mode 100644 index b9246e2eb2480a61f6128f6af1ccd3e3d2c517ab..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py +++ /dev/null @@ -1,980 +0,0 @@ -# Copyright 2025 The HunyuanVideo 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, Tuple, Union - -import numpy as np -import PIL.Image -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTokenizer, - LlamaTokenizerFast, - LlavaForConditionalGeneration, -) - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...loaders import HunyuanVideoLoraLoaderMixin -from ...models import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel -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 .pipeline_output import HunyuanVideoPipelineOutput - - -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 HunyuanVideoImageToVideoPipeline, HunyuanVideoTransformer3DModel - >>> from diffusers.utils import load_image, export_to_video - - >>> # Available checkpoints: hunyuanvideo-community/HunyuanVideo-I2V, hunyuanvideo-community/HunyuanVideo-I2V-33ch - >>> model_id = "hunyuanvideo-community/HunyuanVideo-I2V" - >>> transformer = HunyuanVideoTransformer3DModel.from_pretrained( - ... model_id, subfolder="transformer", torch_dtype=torch.bfloat16 - ... ) - >>> pipe = HunyuanVideoImageToVideoPipeline.from_pretrained( - ... model_id, transformer=transformer, torch_dtype=torch.float16 - ... ) - >>> pipe.vae.enable_tiling() - >>> pipe.to("cuda") - - >>> prompt = "A man with short gray hair plays a red electric guitar." - >>> image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/guitar-man.png" - ... ) - - >>> # If using hunyuanvideo-community/HunyuanVideo-I2V - >>> output = pipe(image=image, prompt=prompt, guidance_scale=6.0).frames[0] - - >>> # If using hunyuanvideo-community/HunyuanVideo-I2V-33ch - >>> output = pipe(image=image, prompt=prompt, guidance_scale=1.0, true_cfg_scale=1.0).frames[0] - - >>> export_to_video(output, "output.mp4", fps=15) - ``` -""" - - -DEFAULT_PROMPT_TEMPLATE = { - "template": ( - "<|start_header_id|>system<|end_header_id|>\n\n\nDescribe the video by detailing the following aspects according to the reference image: " - "1. The main content and theme of the video." - "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects." - "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects." - "4. background environment, light, style and atmosphere." - "5. camera angles, movements, and transitions used in the video:<|eot_id|>\n\n" - "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>" - "<|start_header_id|>assistant<|end_header_id|>\n\n" - ), - "crop_start": 103, - "image_emb_start": 5, - "image_emb_end": 581, - "image_emb_len": 576, - "double_return_token_id": 271, -} - - -def _expand_input_ids_with_image_tokens( - text_input_ids, - prompt_attention_mask, - max_sequence_length, - image_token_index, - image_emb_len, - image_emb_start, - image_emb_end, - pad_token_id, -): - special_image_token_mask = text_input_ids == image_token_index - num_special_image_tokens = torch.sum(special_image_token_mask, dim=-1) - batch_indices, non_image_indices = torch.where(text_input_ids != image_token_index) - - max_expanded_length = max_sequence_length + (num_special_image_tokens.max() * (image_emb_len - 1)) - new_token_positions = torch.cumsum((special_image_token_mask * (image_emb_len - 1) + 1), -1) - 1 - text_to_overwrite = new_token_positions[batch_indices, non_image_indices] - - expanded_input_ids = torch.full( - (text_input_ids.shape[0], max_expanded_length), - pad_token_id, - dtype=text_input_ids.dtype, - device=text_input_ids.device, - ) - expanded_input_ids[batch_indices, text_to_overwrite] = text_input_ids[batch_indices, non_image_indices] - expanded_input_ids[batch_indices, image_emb_start:image_emb_end] = image_token_index - - expanded_attention_mask = torch.zeros( - (text_input_ids.shape[0], max_expanded_length), - dtype=prompt_attention_mask.dtype, - device=prompt_attention_mask.device, - ) - attn_batch_indices, attention_indices = torch.where(expanded_input_ids != pad_token_id) - expanded_attention_mask[attn_batch_indices, attention_indices] = 1.0 - expanded_attention_mask = expanded_attention_mask.to(prompt_attention_mask.dtype) - position_ids = (expanded_attention_mask.cumsum(-1) - 1).masked_fill_((expanded_attention_mask == 0), 1) - - return { - "input_ids": expanded_input_ids, - "attention_mask": expanded_attention_mask, - "position_ids": position_ids, - } - - -# 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 - - -# 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" -): - 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 HunyuanVideoImageToVideoPipeline(DiffusionPipeline, HunyuanVideoLoraLoaderMixin): - r""" - Pipeline for image-to-video generation using HunyuanVideo. - - 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 ([`LlavaForConditionalGeneration`]): - [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). - tokenizer (`LlamaTokenizer`): - Tokenizer from [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). - transformer ([`HunyuanVideoTransformer3DModel`]): - Conditional Transformer to denoise the encoded image latents. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - vae ([`AutoencoderKLHunyuanVideo`]): - Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. - text_encoder_2 ([`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. - tokenizer_2 (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer). - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__( - self, - text_encoder: LlavaForConditionalGeneration, - tokenizer: LlamaTokenizerFast, - transformer: HunyuanVideoTransformer3DModel, - vae: AutoencoderKLHunyuanVideo, - scheduler: FlowMatchEulerDiscreteScheduler, - text_encoder_2: CLIPTextModel, - tokenizer_2: CLIPTokenizer, - image_processor: CLIPImageProcessor, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - text_encoder_2=text_encoder_2, - tokenizer_2=tokenizer_2, - image_processor=image_processor, - ) - - self.vae_scaling_factor = self.vae.config.scaling_factor if getattr(self, "vae", None) else 0.476986 - self.vae_scale_factor_temporal = self.vae.temporal_compression_ratio if getattr(self, "vae", None) else 4 - self.vae_scale_factor_spatial = self.vae.spatial_compression_ratio if getattr(self, "vae", None) else 8 - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - - def _get_llama_prompt_embeds( - self, - image: torch.Tensor, - prompt: Union[str, List[str]], - prompt_template: Dict[str, Any], - num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - max_sequence_length: int = 256, - num_hidden_layers_to_skip: int = 2, - image_embed_interleave: int = 2, - ) -> Tuple[torch.Tensor, torch.Tensor]: - device = device or self._execution_device - dtype = dtype or self.text_encoder.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - prompt = [prompt_template["template"].format(p) for p in prompt] - - crop_start = prompt_template.get("crop_start", None) - - image_emb_len = prompt_template.get("image_emb_len", 576) - image_emb_start = prompt_template.get("image_emb_start", 5) - image_emb_end = prompt_template.get("image_emb_end", 581) - double_return_token_id = prompt_template.get("double_return_token_id", 271) - - if crop_start is None: - prompt_template_input = self.tokenizer( - prompt_template["template"], - padding="max_length", - return_tensors="pt", - return_length=False, - return_overflowing_tokens=False, - return_attention_mask=False, - ) - crop_start = prompt_template_input["input_ids"].shape[-1] - # Remove <|start_header_id|>, <|end_header_id|>, assistant, <|eot_id|>, and placeholder {} - crop_start -= 5 - - max_sequence_length += crop_start - text_inputs = self.tokenizer( - prompt, - max_length=max_sequence_length, - padding="max_length", - truncation=True, - return_tensors="pt", - return_length=False, - return_overflowing_tokens=False, - return_attention_mask=True, - ) - text_input_ids = text_inputs.input_ids.to(device=device) - prompt_attention_mask = text_inputs.attention_mask.to(device=device) - - image_embeds = self.image_processor(image, return_tensors="pt").pixel_values.to(device) - - image_token_index = self.text_encoder.config.image_token_index - pad_token_id = self.text_encoder.config.pad_token_id - expanded_inputs = _expand_input_ids_with_image_tokens( - text_input_ids, - prompt_attention_mask, - max_sequence_length, - image_token_index, - image_emb_len, - image_emb_start, - image_emb_end, - pad_token_id, - ) - prompt_embeds = self.text_encoder( - **expanded_inputs, - pixel_values=image_embeds, - output_hidden_states=True, - ).hidden_states[-(num_hidden_layers_to_skip + 1)] - prompt_embeds = prompt_embeds.to(dtype=dtype) - - if crop_start is not None and crop_start > 0: - text_crop_start = crop_start - 1 + image_emb_len - batch_indices, last_double_return_token_indices = torch.where(text_input_ids == double_return_token_id) - - if last_double_return_token_indices.shape[0] == 3: - # in case the prompt is too long - last_double_return_token_indices = torch.cat( - (last_double_return_token_indices, torch.tensor([text_input_ids.shape[-1]])) - ) - batch_indices = torch.cat((batch_indices, torch.tensor([0]))) - - last_double_return_token_indices = last_double_return_token_indices.reshape(text_input_ids.shape[0], -1)[ - :, -1 - ] - batch_indices = batch_indices.reshape(text_input_ids.shape[0], -1)[:, -1] - assistant_crop_start = last_double_return_token_indices - 1 + image_emb_len - 4 - assistant_crop_end = last_double_return_token_indices - 1 + image_emb_len - attention_mask_assistant_crop_start = last_double_return_token_indices - 4 - attention_mask_assistant_crop_end = last_double_return_token_indices - - prompt_embed_list = [] - prompt_attention_mask_list = [] - image_embed_list = [] - image_attention_mask_list = [] - - for i in range(text_input_ids.shape[0]): - prompt_embed_list.append( - torch.cat( - [ - prompt_embeds[i, text_crop_start : assistant_crop_start[i].item()], - prompt_embeds[i, assistant_crop_end[i].item() :], - ] - ) - ) - prompt_attention_mask_list.append( - torch.cat( - [ - prompt_attention_mask[i, crop_start : attention_mask_assistant_crop_start[i].item()], - prompt_attention_mask[i, attention_mask_assistant_crop_end[i].item() :], - ] - ) - ) - image_embed_list.append(prompt_embeds[i, image_emb_start:image_emb_end]) - image_attention_mask_list.append( - torch.ones(image_embed_list[-1].shape[0]).to(prompt_embeds.device).to(prompt_attention_mask.dtype) - ) - - prompt_embed_list = torch.stack(prompt_embed_list) - prompt_attention_mask_list = torch.stack(prompt_attention_mask_list) - image_embed_list = torch.stack(image_embed_list) - image_attention_mask_list = torch.stack(image_attention_mask_list) - - if 0 < image_embed_interleave < 6: - image_embed_list = image_embed_list[:, ::image_embed_interleave, :] - image_attention_mask_list = image_attention_mask_list[:, ::image_embed_interleave] - - assert ( - prompt_embed_list.shape[0] == prompt_attention_mask_list.shape[0] - and image_embed_list.shape[0] == image_attention_mask_list.shape[0] - ) - - prompt_embeds = torch.cat([image_embed_list, prompt_embed_list], dim=1) - prompt_attention_mask = torch.cat([image_attention_mask_list, prompt_attention_mask_list], dim=1) - - return prompt_embeds, prompt_attention_mask - - def _get_clip_prompt_embeds( - self, - prompt: Union[str, List[str]], - num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - max_sequence_length: int = 77, - ) -> torch.Tensor: - device = device or self._execution_device - dtype = dtype or self.text_encoder_2.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - - text_inputs = self.tokenizer_2( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - ) - - text_input_ids = text_inputs.input_ids - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {max_sequence_length} tokens: {removed_text}" - ) - - prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False).pooler_output - return 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, - 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, - max_sequence_length: int = 256, - image_embed_interleave: int = 2, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if prompt_embeds is None: - prompt_embeds, prompt_attention_mask = self._get_llama_prompt_embeds( - image, - prompt, - prompt_template, - num_videos_per_prompt, - device=device, - dtype=dtype, - max_sequence_length=max_sequence_length, - image_embed_interleave=image_embed_interleave, - ) - - if pooled_prompt_embeds is None: - if prompt_2 is None: - prompt_2 = prompt - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt, - num_videos_per_prompt, - device=device, - dtype=dtype, - max_sequence_length=77, - ) - - return prompt_embeds, pooled_prompt_embeds, prompt_attention_mask - - def check_inputs( - self, - prompt, - prompt_2, - height, - width, - prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - prompt_template=None, - true_cfg_scale=1.0, - guidance_scale=1.0, - ): - 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 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 prompt_template is not None: - if not isinstance(prompt_template, dict): - raise ValueError(f"`prompt_template` has to be of type `dict` but is {type(prompt_template)}") - if "template" not in prompt_template: - raise ValueError( - f"`prompt_template` has to contain a key `template` but only found {prompt_template.keys()}" - ) - - if true_cfg_scale > 1.0 and guidance_scale > 1.0: - logger.warning( - "Both `true_cfg_scale` and `guidance_scale` are greater than 1.0. This will result in both " - "classifier-free guidance and embedded-guidance to be applied. This is not recommended " - "as it may lead to higher memory usage, slower inference and potentially worse results." - ) - - def prepare_latents( - self, - image: torch.Tensor, - batch_size: int, - num_channels_latents: int = 32, - 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, - image_condition_type: str = "latent_concat", - ) -> 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." - ) - - num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 - latent_height, latent_width = height // self.vae_scale_factor_spatial, width // self.vae_scale_factor_spatial - shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width) - - image = image.unsqueeze(2) # [B, C, 1, H, W] - if isinstance(generator, list): - image_latents = [ - retrieve_latents(self.vae.encode(image[i].unsqueeze(0)), generator[i], "argmax") - for i in range(batch_size) - ] - else: - image_latents = [retrieve_latents(self.vae.encode(img.unsqueeze(0)), generator, "argmax") for img in image] - - image_latents = torch.cat(image_latents, dim=0).to(dtype) * self.vae_scaling_factor - image_latents = image_latents.repeat(1, 1, num_latent_frames, 1, 1) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device=device, dtype=dtype) - - t = torch.tensor([0.999]).to(device=device) - latents = latents * t + image_latents * (1 - t) - - if image_condition_type == "token_replace": - image_latents = image_latents[:, :, :1] - - return latents, image_latents - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - @property - def guidance_scale(self): - return self._guidance_scale - - @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, - 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, - height: int = 720, - width: int = 1280, - num_frames: int = 129, - num_inference_steps: int = 50, - 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", - 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, - max_sequence_length: int = 256, - image_embed_interleave: Optional[int] = None, - ): - 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, one has to pass `prompt_embeds`. - instead. - 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*): - 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*): - 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`): - The height in pixels of the generated image. - width (`int`, defaults to `1280`): - The width in pixels of the generated image. - num_frames (`int`, defaults to `129`): - 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. - 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. - true_cfg_scale (`float`, *optional*, defaults to 1.0): - When > 1.0 and a provided `negative_prompt`, enables true classifier-free guidance. - guidance_scale (`float`, defaults to `1.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. Note that the only available - HunyuanVideo model is CFG-distilled, which means that traditional guidance between unconditional and - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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_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. - negative_pooled_prompt_embeds (`torch.FloatTensor`, *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 generated image. Choose between `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`HunyuanVideoPipelineOutput`] 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). - 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. - 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. - - Examples: - - Returns: - [`~HunyuanVideoPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] 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 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_2, - height, - width, - prompt_embeds, - callback_on_step_end_tensor_inputs, - prompt_template, - true_cfg_scale, - guidance_scale, - ) - - image_condition_type = self.transformer.config.image_condition_type - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None - ) - do_true_cfg = true_cfg_scale > 1 and has_neg_prompt - image_embed_interleave = ( - image_embed_interleave - if image_embed_interleave is not None - else ( - 2 if image_condition_type == "latent_concat" else 4 if image_condition_type == "token_replace" else 1 - ) - ) - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._current_timestep = None - self._interrupt = False - - device = self._execution_device - - # 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. Prepare latent variables - vae_dtype = self.vae.dtype - image_tensor = self.video_processor.preprocess(image, height, width).to(device, vae_dtype) - - if image_condition_type == "latent_concat": - num_channels_latents = (self.transformer.config.in_channels - 1) // 2 - elif image_condition_type == "token_replace": - num_channels_latents = self.transformer.config.in_channels - - latents, image_latents = self.prepare_latents( - image_tensor, - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - num_frames, - torch.float32, - device, - generator, - latents, - image_condition_type, - ) - if image_condition_type == "latent_concat": - image_latents[:, :, 1:] = 0 - mask = image_latents.new_ones(image_latents.shape[0], 1, *image_latents.shape[2:]) - mask[:, :, 1:] = 0 - - # 4. Encode input prompt - transformer_dtype = self.transformer.dtype - prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self.encode_prompt( - image=image, - prompt=prompt, - prompt_2=prompt_2, - prompt_template=prompt_template, - num_videos_per_prompt=num_videos_per_prompt, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - device=device, - max_sequence_length=max_sequence_length, - image_embed_interleave=image_embed_interleave, - ) - prompt_embeds = prompt_embeds.to(transformer_dtype) - prompt_attention_mask = prompt_attention_mask.to(transformer_dtype) - pooled_prompt_embeds = pooled_prompt_embeds.to(transformer_dtype) - - if do_true_cfg: - black_image = PIL.Image.new("RGB", (width, height), 0) - negative_prompt_embeds, negative_pooled_prompt_embeds, negative_prompt_attention_mask = self.encode_prompt( - image=black_image, - prompt=negative_prompt, - prompt_2=negative_prompt_2, - prompt_template=prompt_template, - num_videos_per_prompt=num_videos_per_prompt, - prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=negative_pooled_prompt_embeds, - prompt_attention_mask=negative_prompt_attention_mask, - device=device, - max_sequence_length=max_sequence_length, - ) - negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype) - negative_prompt_attention_mask = negative_prompt_attention_mask.to(transformer_dtype) - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.to(transformer_dtype) - - # 5. 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) - - # 6. Prepare guidance condition - guidance = None - if self.transformer.config.guidance_embeds: - guidance = ( - torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0 - ) - - # 7. Denoising loop - 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 - # 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_condition_type == "latent_concat": - latent_model_input = torch.cat([latents, image_latents, mask], dim=1).to(transformer_dtype) - elif image_condition_type == "token_replace": - latent_model_input = torch.cat([image_latents, latents[:, :, 1:]], dim=2).to(transformer_dtype) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - encoder_attention_mask=prompt_attention_mask, - pooled_projections=pooled_prompt_embeds, - guidance=guidance, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - - if do_true_cfg: - neg_noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, - encoder_attention_mask=negative_prompt_attention_mask, - pooled_projections=negative_pooled_prompt_embeds, - guidance=guidance, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) - - # compute the previous noisy sample x_t -> x_t-1 - if image_condition_type == "latent_concat": - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] - elif image_condition_type == "token_replace": - latents = latents = self.scheduler.step( - noise_pred[:, :, 1:], t, latents[:, :, 1:], return_dict=False - )[0] - latents = torch.cat([image_latents, latents], dim=2) - - 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 not output_type == "latent": - latents = latents.to(self.vae.dtype) / self.vae_scaling_factor - video = self.vae.decode(latents, return_dict=False)[0] - if image_condition_type == "latent_concat": - video = video[:, :, 4:, :, :] - video = self.video_processor.postprocess_video(video, output_type=output_type) - else: - if image_condition_type == "latent_concat": - video = latents[:, :, 1:, :, :] - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return HunyuanVideoPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_output.py deleted file mode 100644 index fae0370a53b7ae2163972809f1ee40a4deedbcbf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuan_video/pipeline_output.py +++ /dev/null @@ -1,39 +0,0 @@ -from dataclasses import dataclass -from typing import List, Union - -import numpy as np -import PIL.Image -import torch - -from diffusers.utils import BaseOutput - - -@dataclass -class HunyuanVideoPipelineOutput(BaseOutput): - r""" - 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 - 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 - - -@dataclass -class HunyuanVideoFramepackPipelineOutput(BaseOutput): - r""" - 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 - 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]] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuandit/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuandit/__init__.py deleted file mode 100644 index 8337399106f0585d15fa0a35f607baa2c04b203b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuandit/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -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_hunyuandit"] = ["HunyuanDiTPipeline"] - -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_hunyuandit import HunyuanDiTPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuandit/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuandit/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 302c1d8aace44d11a5075944951a35bdedb66780..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuandit/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuandit/__pycache__/pipeline_hunyuandit.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuandit/__pycache__/pipeline_hunyuandit.cpython-310.pyc deleted file mode 100644 index 7954694ddd2071a38883b7e84bf43a416de46827..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuandit/__pycache__/pipeline_hunyuandit.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py deleted file mode 100644 index c7f84866fe43f610a80f46b1d9b5ee54c0900583..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py +++ /dev/null @@ -1,914 +0,0 @@ -# Copyright 2025 HunyuanDiT Authors 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 Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import torch -from transformers import BertModel, BertTokenizer, CLIPImageProcessor, MT5Tokenizer, T5EncoderModel - -from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import VaeImageProcessor -from ...models import AutoencoderKL, HunyuanDiT2DModel -from ...models.embeddings import get_2d_rotary_pos_embed -from ...pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker -from ...schedulers import DDPMScheduler -from ...utils import ( - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -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 HunyuanDiTPipeline - - >>> pipe = HunyuanDiTPipeline.from_pretrained( - ... "Tencent-Hunyuan/HunyuanDiT-Diffusers", torch_dtype=torch.float16 - ... ) - >>> pipe.to("cuda") - - >>> # You may also use English prompt as HunyuanDiT supports both English and Chinese - >>> # prompt = "An astronaut riding a horse" - >>> prompt = "一个宇航员在骑马" - >>> image = pipe(prompt).images[0] - ``` -""" - -STANDARD_RATIO = np.array( - [ - 1.0, # 1:1 - 4.0 / 3.0, # 4:3 - 3.0 / 4.0, # 3:4 - 16.0 / 9.0, # 16:9 - 9.0 / 16.0, # 9:16 - ] -) -STANDARD_SHAPE = [ - [(1024, 1024), (1280, 1280)], # 1:1 - [(1024, 768), (1152, 864), (1280, 960)], # 4:3 - [(768, 1024), (864, 1152), (960, 1280)], # 3:4 - [(1280, 768)], # 16:9 - [(768, 1280)], # 9:16 -] -STANDARD_AREA = [np.array([w * h for w, h in shapes]) for shapes in STANDARD_SHAPE] -SUPPORTED_SHAPE = [ - (1024, 1024), - (1280, 1280), # 1:1 - (1024, 768), - (1152, 864), - (1280, 960), # 4:3 - (768, 1024), - (864, 1152), - (960, 1280), # 3:4 - (1280, 768), # 16:9 - (768, 1280), # 9:16 -] - - -def map_to_standard_shapes(target_width, target_height): - target_ratio = target_width / target_height - closest_ratio_idx = np.argmin(np.abs(STANDARD_RATIO - target_ratio)) - closest_area_idx = np.argmin(np.abs(STANDARD_AREA[closest_ratio_idx] - target_width * target_height)) - width, height = STANDARD_SHAPE[closest_ratio_idx][closest_area_idx] - return width, height - - -def get_resize_crop_region_for_grid(src, tgt_size): - th = tw = tgt_size - h, w = src - - r = h / w - - # resize - if r > 1: - resize_height = th - resize_width = int(round(th / h * w)) - else: - resize_width = tw - resize_height = int(round(tw / w * h)) - - crop_top = int(round((th - resize_height) / 2.0)) - crop_left = int(round((tw - resize_width) / 2.0)) - - return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width) - - -# 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 HunyuanDiTPipeline(DiffusionPipeline): - r""" - Pipeline for English/Chinese-to-image generation using HunyuanDiT. - - 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.) - - HunyuanDiT uses two text encoders: [mT5](https://huggingface.co/google/mt5-base) and [bilingual CLIP](fine-tuned by - ourselves) - - Args: - 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`]): - 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`]): - A `BertTokenizer` or `CLIPTokenizer` to tokenize text. - transformer ([`HunyuanDiT2DModel`]): - The HunyuanDiT model designed by Tencent Hunyuan. - text_encoder_2 (`T5EncoderModel`): - The mT5 embedder. Specifically, it is 't5-v1_1-xxl'. - tokenizer_2 (`MT5Tokenizer`): - The tokenizer for the mT5 embedder. - scheduler ([`DDPMScheduler`]): - A scheduler to be used in combination with HunyuanDiT to denoise the encoded image latents. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" - _optional_components = [ - "safety_checker", - "feature_extractor", - "text_encoder_2", - "tokenizer_2", - "text_encoder", - "tokenizer", - ] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "prompt_embeds_2", - "negative_prompt_embeds_2", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: BertModel, - tokenizer: BertTokenizer, - transformer: HunyuanDiT2DModel, - scheduler: DDPMScheduler, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - text_encoder_2: Optional[T5EncoderModel] = None, - tokenizer_2: Optional[MT5Tokenizer] = None, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - transformer=transformer, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - text_encoder_2=text_encoder_2, - ) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - 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.register_to_config(requires_safety_checker=requires_safety_checker) - self.default_sample_size = ( - self.transformer.config.sample_size - if hasattr(self, "transformer") and self.transformer is not None - else 128 - ) - - def encode_prompt( - self, - prompt: str, - device: torch.device = None, - 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, - text_encoder_index: int = 0, - ): - 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 - dtype (`torch.dtype`): - torch dtype - 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. - prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the prompt. Required when `prompt_embeds` is passed directly. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the negative prompt. Required when `negative_prompt_embeds` is passed directly. - max_sequence_length (`int`, *optional*): maximum sequence length to use for the prompt. - text_encoder_index (`int`, *optional*): - Index of the text encoder to use. `0` for clip and `1` for T5. - """ - if dtype is None: - if self.text_encoder_2 is not None: - dtype = self.text_encoder_2.dtype - elif self.transformer is not None: - dtype = self.transformer.dtype - else: - dtype = None - - if device is None: - device = self._execution_device - - tokenizers = [self.tokenizer, self.tokenizer_2] - text_encoders = [self.text_encoder, self.text_encoder_2] - - tokenizer = tokenizers[text_encoder_index] - text_encoder = text_encoders[text_encoder_index] - - if max_sequence_length is None: - if text_encoder_index == 0: - max_length = 77 - if text_encoder_index == 1: - max_length = 256 - else: - max_length = max_sequence_length - - 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: - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=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_attention_mask = text_inputs.attention_mask.to(device) - prompt_embeds = text_encoder( - text_input_ids.to(device), - attention_mask=prompt_attention_mask, - ) - prompt_embeds = prompt_embeds[0] - prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1) - - prompt_embeds = prompt_embeds.to(dtype=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 - - max_length = prompt_embeds.shape[1] - uncond_input = tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - - negative_prompt_attention_mask = uncond_input.attention_mask.to(device) - negative_prompt_embeds = text_encoder( - uncond_input.input_ids.to(device), - attention_mask=negative_prompt_attention_mask, - ) - negative_prompt_embeds = negative_prompt_embeds[0] - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 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] - - negative_prompt_embeds = negative_prompt_embeds.to(dtype=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) - - return prompt_embeds, negative_prompt_embeds, prompt_attention_mask, negative_prompt_attention_mask - - # 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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - prompt_embeds_2=None, - negative_prompt_embeds_2=None, - prompt_attention_mask_2=None, - negative_prompt_attention_mask_2=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_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 None and prompt_embeds_2 is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds_2`. Cannot leave both `prompt` and `prompt_embeds_2` 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 prompt_embeds_2 is not None and prompt_attention_mask_2 is None: - raise ValueError("Must provide `prompt_attention_mask_2` when specifying `prompt_embeds_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." - ) - - 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 negative_prompt_embeds_2 is not None and negative_prompt_attention_mask_2 is None: - raise ValueError( - "Must provide `negative_prompt_attention_mask_2` when specifying `negative_prompt_embeds_2`." - ) - 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_2 is not None and negative_prompt_embeds_2 is not None: - if prompt_embeds_2.shape != negative_prompt_embeds_2.shape: - raise ValueError( - "`prompt_embeds_2` and `negative_prompt_embeds_2` must have the same shape when passed directly, but" - f" got: `prompt_embeds_2` {prompt_embeds_2.shape} != `negative_prompt_embeds_2`" - f" {negative_prompt_embeds_2.shape}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - # 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 - - @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, - 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", - 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"], - 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), - use_resolution_binning: bool = True, - ): - r""" - The call function to the pipeline for generation with HunyuanDiT. - - Args: - 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. - width (`int`): - 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. This parameter is modulated by `strength`. - 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*): - 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): - The number of images to generate per prompt. - 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*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - 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. - prompt_embeds_2 (`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. - 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. - negative_prompt_embeds_2 (`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. - prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the prompt. Required when `prompt_embeds` is passed directly. - prompt_attention_mask_2 (`torch.Tensor`, *optional*): - Attention mask for the prompt. Required when `prompt_embeds_2` is passed directly. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the negative prompt. Required when `negative_prompt_embeds` is passed directly. - negative_prompt_attention_mask_2 (`torch.Tensor`, *optional*): - Attention mask for the negative prompt. Required when `negative_prompt_embeds_2` is passed directly. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback_on_step_end (`Callable[[int, int, Dict], 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*): - 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)`): - The original size of the image. Used to calculate the time ids. - 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)`): - 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 - standard resolution. Supported resolutions are 1024x1024, 1280x1280, 1024x768, 1152x864, 1280x960, - 768x1024, 864x1152, 960x1280, 1280x768, and 768x1280. It is recommended to set this to `True`. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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 isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - # 0. default height and width - height = height or self.default_sample_size * self.vae_scale_factor - width = width or self.default_sample_size * self.vae_scale_factor - height = int((height // 16) * 16) - width = int((width // 16) * 16) - - if use_resolution_binning and (height, width) not in SUPPORTED_SHAPE: - width, height = map_to_standard_shapes(width, height) - height = int(height) - width = int(width) - logger.warning(f"Reshaped to (height, width)=({height}, {width}), Supported shapes are {SUPPORTED_SHAPE}") - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - prompt_embeds_2, - negative_prompt_embeds_2, - prompt_attention_mask_2, - negative_prompt_attention_mask_2, - callback_on_step_end_tensor_inputs, - ) - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - 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. Encode input prompt - - ( - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt=prompt, - device=device, - dtype=self.transformer.dtype, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - 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, - max_sequence_length=77, - text_encoder_index=0, - ) - ( - prompt_embeds_2, - negative_prompt_embeds_2, - prompt_attention_mask_2, - negative_prompt_attention_mask_2, - ) = self.encode_prompt( - prompt=prompt, - device=device, - dtype=self.transformer.dtype, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds_2, - negative_prompt_embeds=negative_prompt_embeds_2, - prompt_attention_mask=prompt_attention_mask_2, - negative_prompt_attention_mask=negative_prompt_attention_mask_2, - max_sequence_length=256, - text_encoder_index=1, - ) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7 create image_rotary_emb, style embedding & time ids - grid_height = height // 8 // self.transformer.config.patch_size - grid_width = width // 8 // self.transformer.config.patch_size - base_size = 512 // 8 // self.transformer.config.patch_size - grid_crops_coords = get_resize_crop_region_for_grid((grid_height, grid_width), base_size) - image_rotary_emb = get_2d_rotary_pos_embed( - self.transformer.inner_dim // self.transformer.num_heads, - grid_crops_coords, - (grid_height, grid_width), - device=device, - output_type="pt", - ) - - style = torch.tensor([0], device=device) - - target_size = target_size or (height, width) - add_time_ids = list(original_size + target_size + crops_coords_top_left) - add_time_ids = torch.tensor([add_time_ids], dtype=prompt_embeds.dtype) - - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask]) - prompt_embeds_2 = torch.cat([negative_prompt_embeds_2, prompt_embeds_2]) - prompt_attention_mask_2 = torch.cat([negative_prompt_attention_mask_2, prompt_attention_mask_2]) - add_time_ids = torch.cat([add_time_ids] * 2, dim=0) - style = torch.cat([style] * 2, dim=0) - - prompt_embeds = prompt_embeds.to(device=device) - prompt_attention_mask = prompt_attention_mask.to(device=device) - prompt_embeds_2 = prompt_embeds_2.to(device=device) - prompt_attention_mask_2 = prompt_attention_mask_2.to(device=device) - add_time_ids = add_time_ids.to(dtype=prompt_embeds.dtype, device=device).repeat( - batch_size * num_images_per_prompt, 1 - ) - style = style.to(device=device).repeat(batch_size * num_images_per_prompt) - - # 8. Denoising loop - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # expand scalar t to 1-D tensor to match the 1st dim of latent_model_input - t_expand = torch.tensor([t] * latent_model_input.shape[0], device=device).to( - dtype=latent_model_input.dtype - ) - - # predict the noise residual - noise_pred = self.transformer( - latent_model_input, - t_expand, - encoder_hidden_states=prompt_embeds, - text_embedding_mask=prompt_attention_mask, - encoder_hidden_states_t5=prompt_embeds_2, - text_embedding_mask_t5=prompt_attention_mask_2, - image_meta_size=add_time_ids, - style=style, - image_rotary_emb=image_rotary_emb, - return_dict=False, - )[0] - - noise_pred, _ = noise_pred.chunk(2, dim=1) - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - if self.do_classifier_free_guidance and guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - prompt_embeds_2 = callback_outputs.pop("prompt_embeds_2", prompt_embeds_2) - negative_prompt_embeds_2 = callback_outputs.pop( - "negative_prompt_embeds_2", negative_prompt_embeds_2 - ) - - 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 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/venv/lib/python3.10/site-packages/diffusers/pipelines/i2vgen_xl/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/i2vgen_xl/__init__.py deleted file mode 100644 index b24a7e4cee7fb843b9424469a05f511adfa758de..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/i2vgen_xl/__init__.py +++ /dev/null @@ -1,46 +0,0 @@ -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_i2vgen_xl"] = ["I2VGenXLPipeline"] - - -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_i2vgen_xl import I2VGenXLPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/i2vgen_xl/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/i2vgen_xl/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index dac71d2533755e36fb3d7afee1fea2eb9d2c2502..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/i2vgen_xl/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/i2vgen_xl/__pycache__/pipeline_i2vgen_xl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/i2vgen_xl/__pycache__/pipeline_i2vgen_xl.cpython-310.pyc deleted file mode 100644 index 94d3eaeb33c4f2e3e9a3e2ff9e13196d9f3c28da..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/i2vgen_xl/__pycache__/pipeline_i2vgen_xl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py deleted file mode 100644 index c6cc724a71f0a8c65b5df30e05d5456c34006475..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py +++ /dev/null @@ -1,797 +0,0 @@ -# Copyright 2025 Alibaba DAMO-VILAB 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, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...models import AutoencoderKL -from ...models.unets.unet_i2vgen_xl import I2VGenXLUNet -from ...schedulers import DDIMScheduler -from ...utils import ( - BaseOutput, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ...video_processor import VideoProcessor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin - - -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 I2VGenXLPipeline - >>> from diffusers.utils import export_to_gif, load_image - - >>> pipeline = I2VGenXLPipeline.from_pretrained( - ... "ali-vilab/i2vgen-xl", torch_dtype=torch.float16, variant="fp16" - ... ) - >>> pipeline.enable_model_cpu_offload() - - >>> image_url = ( - ... "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/i2vgen_xl_images/img_0009.png" - ... ) - >>> image = load_image(image_url).convert("RGB") - - >>> prompt = "Papers were floating in the air on a table in the library" - >>> negative_prompt = "Distorted, discontinuous, Ugly, blurry, low resolution, motionless, static, disfigured, disconnected limbs, Ugly faces, incomplete arms" - >>> generator = torch.manual_seed(8888) - - >>> frames = pipeline( - ... prompt=prompt, - ... image=image, - ... num_inference_steps=50, - ... negative_prompt=negative_prompt, - ... guidance_scale=9.0, - ... generator=generator, - ... ).frames[0] - >>> video_path = export_to_gif(frames, "i2v.gif") - ``` -""" - - -@dataclass -class I2VGenXLPipelineOutput(BaseOutput): - r""" - 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 - 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]]] - - -class I2VGenXLPipeline( - DeprecatedPipelineMixin, - DiffusionPipeline, - StableDiffusionMixin, -): - _last_supported_version = "0.33.1" - r""" - Pipeline for image-to-video generation as proposed in [I2VGenXL](https://i2vgen-xl.github.io/). - - 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: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer (`CLIPTokenizer`): - A [`~transformers.CLIPTokenizer`] to tokenize text. - unet ([`I2VGenXLUNet`]): - A [`I2VGenXLUNet`] to denoise the encoded video latents. - scheduler ([`DDIMScheduler`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - image_encoder: CLIPVisionModelWithProjection, - feature_extractor: CLIPImageProcessor, - unet: I2VGenXLUNet, - scheduler: DDIMScheduler, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - unet=unet, - scheduler=scheduler, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - # `do_resize=False` as we do custom resizing. - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor, do_resize=False) - - @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 do_classifier_free_guidance(self): - return self._guidance_scale > 1 - - def encode_prompt( - self, - 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, - ): - 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_videos_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. - 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. - """ - 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: - 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_videos_per_prompt, 1) - prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, seq_len, -1) - - # get unconditional embeddings for classifier free guidance - if self.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 - - 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 - - # Apply clip_skip to negative prompt embeds - if clip_skip is None: - negative_prompt_embeds = self.text_encoder( - uncond_input.input_ids.to(device), - attention_mask=attention_mask, - ) - negative_prompt_embeds = negative_prompt_embeds[0] - else: - negative_prompt_embeds = self.text_encoder( - uncond_input.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. - negative_prompt_embeds = negative_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. - negative_prompt_embeds = self.text_encoder.text_model.final_layer_norm(negative_prompt_embeds) - - if self.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_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 - - def _encode_image(self, image, device, num_videos_per_prompt): - dtype = next(self.image_encoder.parameters()).dtype - - if not isinstance(image, torch.Tensor): - image = self.video_processor.pil_to_numpy(image) - image = self.video_processor.numpy_to_pt(image) - - # Normalize the image with CLIP training stats. - image = self.feature_extractor( - images=image, - do_normalize=True, - do_center_crop=False, - do_resize=False, - do_rescale=False, - return_tensors="pt", - ).pixel_values - - image = image.to(device=device, dtype=dtype) - image_embeddings = self.image_encoder(image).image_embeds - image_embeddings = image_embeddings.unsqueeze(1) - - # duplicate image embeddings for each generation per prompt, using mps friendly method - bs_embed, seq_len, _ = image_embeddings.shape - image_embeddings = image_embeddings.repeat(1, num_videos_per_prompt, 1) - image_embeddings = image_embeddings.view(bs_embed * num_videos_per_prompt, seq_len, -1) - - if self.do_classifier_free_guidance: - negative_image_embeddings = torch.zeros_like(image_embeddings) - image_embeddings = torch.cat([negative_image_embeddings, image_embeddings]) - - return image_embeddings - - def decode_latents(self, latents, decode_chunk_size=None): - latents = 1 / self.vae.config.scaling_factor * latents - - batch_size, channels, num_frames, height, width = latents.shape - latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) - - if decode_chunk_size is not None: - frames = [] - for i in range(0, latents.shape[0], decode_chunk_size): - frame = self.vae.decode(latents[i : i + decode_chunk_size]).sample - frames.append(frame) - image = torch.cat(frames, dim=0) - else: - image = self.vae.decode(latents).sample - - decode_shape = (batch_size, num_frames, -1) + image.shape[2:] - video = image[None, :].reshape(decode_shape).permute(0, 2, 1, 3, 4) - - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - video = video.float() - return video - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - image, - height, - width, - negative_prompt=None, - prompt_embeds=None, - negative_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 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 ( - not isinstance(image, torch.Tensor) - and not isinstance(image, PIL.Image.Image) - 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" - f" {type(image)}" - ) - - def prepare_image_latents( - self, - image, - device, - num_frames, - num_videos_per_prompt, - ): - image = image.to(device=device) - image_latents = self.vae.encode(image).latent_dist.sample() - image_latents = image_latents * self.vae.config.scaling_factor - - # Add frames dimension to image latents - image_latents = image_latents.unsqueeze(2) - - # Append a position mask for each subsequent frame - # after the initial image latent frame - frame_position_mask = [] - for frame_idx in range(num_frames - 1): - scale = (frame_idx + 1) / (num_frames - 1) - frame_position_mask.append(torch.ones_like(image_latents[:, :, :1]) * scale) - if frame_position_mask: - frame_position_mask = torch.cat(frame_position_mask, dim=2) - image_latents = torch.cat([image_latents, frame_position_mask], dim=2) - - # duplicate image_latents for each generation per prompt, using mps friendly method - image_latents = image_latents.repeat(num_videos_per_prompt, 1, 1, 1, 1) - - if self.do_classifier_free_guidance: - image_latents = torch.cat([image_latents] * 2) - - return image_latents - - # Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_synth.TextToVideoSDPipeline.prepare_latents - def prepare_latents( - self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None - ): - shape = ( - batch_size, - num_channels_latents, - num_frames, - height // self.vae_scale_factor, - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - image: PipelineImageInput = None, - height: Optional[int] = 704, - width: Optional[int] = 1280, - target_fps: Optional[int] = 16, - num_frames: int = 16, - num_inference_steps: int = 50, - guidance_scale: float = 9.0, - negative_prompt: Optional[Union[str, List[str]]] = 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = 1, - ): - r""" - The call function to the pipeline for image-to-video generation with [`I2VGenXLPipeline`]. - - Args: - 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 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`): - 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. - target_fps (`int`, *optional*): - Frames per second. The rate at which the generated images shall be exported to a video after - generation. This is also used as a "micro-condition" while generation. - num_frames (`int`, *optional*): - The number of video frames to generate. - num_inference_steps (`int`, *optional*): - The number of denoising steps. - 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*): - 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*): - 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. - num_videos_per_prompt (`int`, *optional*): - The number of images to generate per prompt. - decode_chunk_size (`int`, *optional*): - 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*): - 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. - 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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). - 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. - - Examples: - - Returns: - [`pipelines.i2vgen_xl.pipeline_i2vgen_xl.I2VGenXLPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`pipelines.i2vgen_xl.pipeline_i2vgen_xl.I2VGenXLPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is a list with the generated frames. - """ - # 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, image, height, width, 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. - self._guidance_scale = guidance_scale - - # 3.1 Encode input text prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_videos_per_prompt, - 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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - # 3.2 Encode image prompt - # 3.2.1 Image encodings. - # https://github.com/ali-vilab/i2vgen-xl/blob/2539c9262ff8a2a22fa9daecbfd13f0a2dbc32d0/tools/inferences/inference_i2vgen_entrance.py#L114 - cropped_image = _center_crop_wide(image, (width, width)) - cropped_image = _resize_bilinear( - cropped_image, (self.feature_extractor.crop_size["width"], self.feature_extractor.crop_size["height"]) - ) - image_embeddings = self._encode_image(cropped_image, device, num_videos_per_prompt) - - # 3.2.2 Image latents. - resized_image = _center_crop_wide(image, (width, height)) - image = self.video_processor.preprocess(resized_image).to(device=device, dtype=image_embeddings.dtype) - image_latents = self.prepare_image_latents( - image, - device=device, - num_frames=num_frames, - num_videos_per_prompt=num_videos_per_prompt, - ) - - # 3.3 Prepare additional conditions for the UNet. - if self.do_classifier_free_guidance: - fps_tensor = torch.tensor([target_fps, target_fps]).to(device) - else: - fps_tensor = torch.tensor([target_fps]).to(device) - fps_tensor = fps_tensor.repeat(batch_size * num_videos_per_prompt, 1).ravel() - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - num_channels_latents, - num_frames, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - fps=fps_tensor, - image_latents=image_latents, - image_embeddings=image_embeddings, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # reshape latents - batch_size, channel, frames, width, height = latents.shape - latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * frames, channel, width, height) - noise_pred = noise_pred.permute(0, 2, 1, 3, 4).reshape(batch_size * frames, channel, width, height) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # reshape latents back - latents = latents[None, :].reshape(batch_size, frames, channel, width, height).permute(0, 2, 1, 3, 4) - # 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() - - # 8. Post processing - if output_type == "latent": - video = latents - else: - video_tensor = self.decode_latents(latents, decode_chunk_size=decode_chunk_size) - video = self.video_processor.postprocess_video(video=video_tensor, output_type=output_type) - - # 9. Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return I2VGenXLPipelineOutput(frames=video) - - -# The following utilities are taken and adapted from -# https://github.com/ali-vilab/i2vgen-xl/blob/main/utils/transforms.py. - - -def _convert_pt_to_pil(image: Union[torch.Tensor, List[torch.Tensor]]): - if isinstance(image, list) and isinstance(image[0], torch.Tensor): - image = torch.cat(image, 0) - - if isinstance(image, torch.Tensor): - if image.ndim == 3: - image = image.unsqueeze(0) - - image_numpy = VaeImageProcessor.pt_to_numpy(image) - image_pil = VaeImageProcessor.numpy_to_pil(image_numpy) - image = image_pil - - return image - - -def _resize_bilinear( - image: Union[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) - - if isinstance(image, list): - image = [u.resize(resolution, PIL.Image.BILINEAR) for u in image] - else: - image = image.resize(resolution, PIL.Image.BILINEAR) - return image - - -def _center_crop_wide( - image: Union[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) - - if isinstance(image, list): - scale = min(image[0].size[0] / resolution[0], image[0].size[1] / resolution[1]) - image = [u.resize((round(u.width // scale), round(u.height // scale)), resample=PIL.Image.BOX) for u in image] - - # center crop - x1 = (image[0].width - resolution[0]) // 2 - y1 = (image[0].height - resolution[1]) // 2 - image = [u.crop((x1, y1, x1 + resolution[0], y1 + resolution[1])) for u in image] - return image - else: - scale = min(image.size[0] / resolution[0], image.size[1] / resolution[1]) - image = image.resize((round(image.width // scale), round(image.height // scale)), resample=PIL.Image.BOX) - x1 = (image.width - resolution[0]) // 2 - y1 = (image.height - resolution[1]) // 2 - image = image.crop((x1, y1, x1 + resolution[0], y1 + resolution[1])) - return image diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__init__.py deleted file mode 100644 index 606f7b378a79489bbcbaa87db2040bd4196bbd8a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__init__.py +++ /dev/null @@ -1,66 +0,0 @@ -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_kandinsky"] = ["KandinskyPipeline"] - _import_structure["pipeline_kandinsky_combined"] = [ - "KandinskyCombinedPipeline", - "KandinskyImg2ImgCombinedPipeline", - "KandinskyInpaintCombinedPipeline", - ] - _import_structure["pipeline_kandinsky_img2img"] = ["KandinskyImg2ImgPipeline"] - _import_structure["pipeline_kandinsky_inpaint"] = ["KandinskyInpaintPipeline"] - _import_structure["pipeline_kandinsky_prior"] = ["KandinskyPriorPipeline", "KandinskyPriorPipelineOutput"] - _import_structure["text_encoder"] = ["MultilingualCLIP"] - - -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_kandinsky import KandinskyPipeline - from .pipeline_kandinsky_combined import ( - KandinskyCombinedPipeline, - KandinskyImg2ImgCombinedPipeline, - KandinskyInpaintCombinedPipeline, - ) - from .pipeline_kandinsky_img2img import KandinskyImg2ImgPipeline - from .pipeline_kandinsky_inpaint import KandinskyInpaintPipeline - from .pipeline_kandinsky_prior import KandinskyPriorPipeline, KandinskyPriorPipelineOutput - from .text_encoder import MultilingualCLIP - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index d23f502e61c3a963cff2d41e52bc238187292e26..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky.cpython-310.pyc deleted file mode 100644 index e3d2a8e762954038d20eafcfd9a3efb94b95b9da..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_combined.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_combined.cpython-310.pyc deleted file mode 100644 index 05af219b7c05ef8935ab66f42f29701cc4e186a7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_combined.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_img2img.cpython-310.pyc deleted file mode 100644 index 527395a0bbe5c9ebbc3ce3b871a51f25c722c664..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_inpaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_inpaint.cpython-310.pyc deleted file mode 100644 index 7f31020ee4906ea202b430f2ae7efd6b4d16303f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_inpaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_prior.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_prior.cpython-310.pyc deleted file mode 100644 index 718b4f8d37f7df5247cae994b41c32fdcc380b24..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/pipeline_kandinsky_prior.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/text_encoder.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/text_encoder.cpython-310.pyc deleted file mode 100644 index 67c2fa70e36742dc9ebef642fe8ecf42515870ff..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/__pycache__/text_encoder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky.py deleted file mode 100644 index 89fea893375263bcfe269619926bc430e68d6ba4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky.py +++ /dev/null @@ -1,419 +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 Callable, List, Optional, Union - -import torch -from transformers import ( - XLMRobertaTokenizer, -) - -from ...models import UNet2DConditionModel, VQModel -from ...schedulers import DDIMScheduler, DDPMScheduler -from ...utils import ( - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput -from .text_encoder import MultilingualCLIP - - -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 - >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline - >>> import torch - - >>> pipe_prior = KandinskyPriorPipeline.from_pretrained("kandinsky-community/Kandinsky-2-1-prior") - >>> pipe_prior.to("cuda") - - >>> prompt = "red cat, 4k photo" - >>> out = pipe_prior(prompt) - >>> image_emb = out.image_embeds - >>> negative_image_emb = out.negative_image_embeds - - >>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1") - >>> pipe.to("cuda") - - >>> image = pipe( - ... prompt, - ... image_embeds=image_emb, - ... negative_image_embeds=negative_image_emb, - ... height=768, - ... width=768, - ... num_inference_steps=100, - ... ).images - - >>> image[0].save("cat.png") - ``` -""" - - -def get_new_h_w(h, w, scale_factor=8): - new_h = h // scale_factor**2 - if h % scale_factor**2 != 0: - new_h += 1 - new_w = w // scale_factor**2 - if w % scale_factor**2 != 0: - new_w += 1 - return new_h * scale_factor, new_w * scale_factor - - -class KandinskyPipeline(DiffusionPipeline): - """ - Pipeline for text-to-image generation using Kandinsky - - 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.) - - Args: - text_encoder ([`MultilingualCLIP`]): - Frozen text-encoder. - tokenizer ([`XLMRobertaTokenizer`]): - Tokenizer of class - scheduler (Union[`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. - movq ([`VQModel`]): - MoVQ Decoder to generate the image from the latents. - """ - - model_cpu_offload_seq = "text_encoder->unet->movq" - - def __init__( - self, - text_encoder: MultilingualCLIP, - tokenizer: XLMRobertaTokenizer, - unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, DDPMScheduler], - movq: VQModel, - ): - super().__init__() - - self.register_modules( - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - movq=movq, - ) - self.movq_scale_factor = 2 ** (len(self.movq.config.block_out_channels) - 1) - - # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - def _encode_prompt( - self, - prompt, - device, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt=None, - ): - batch_size = len(prompt) if isinstance(prompt, list) else 1 - # get prompt text embeddings - text_inputs = self.tokenizer( - prompt, - padding="max_length", - truncation=True, - max_length=77, - return_attention_mask=True, - add_special_tokens=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}" - ) - - text_input_ids = text_input_ids.to(device) - text_mask = text_inputs.attention_mask.to(device) - - prompt_embeds, text_encoder_hidden_states = self.text_encoder( - input_ids=text_input_ids, attention_mask=text_mask - ) - - prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - text_encoder_hidden_states = text_encoder_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - uncond_tokens: List[str] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif 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 - - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=77, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - uncond_text_input_ids = uncond_input.input_ids.to(device) - uncond_text_mask = uncond_input.attention_mask.to(device) - - negative_prompt_embeds, uncond_text_encoder_hidden_states = self.text_encoder( - input_ids=uncond_text_input_ids, attention_mask=uncond_text_mask - ) - - # 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.repeat(1, num_images_per_prompt) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len) - - seq_len = uncond_text_encoder_hidden_states.shape[1] - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.repeat(1, num_images_per_prompt, 1) - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.view( - batch_size * num_images_per_prompt, seq_len, -1 - ) - uncond_text_mask = uncond_text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - # done duplicates - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - text_encoder_hidden_states = torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]) - - text_mask = torch.cat([uncond_text_mask, text_mask]) - - return prompt_embeds, text_encoder_hidden_states, text_mask - - @torch.no_grad() - @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, - 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, - callback_steps: int = 1, - return_dict: bool = True, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `List[str]`): - The prompt or prompts to guide the image generation. - 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]`): - The clip image embeddings for negative text prompt, will be used to condition the image generation. - 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): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - 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. - 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. - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - - if isinstance(prompt, str): - batch_size = 1 - elif isinstance(prompt, list): - batch_size = len(prompt) - else: - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - device = self._execution_device - - batch_size = batch_size * num_images_per_prompt - do_classifier_free_guidance = guidance_scale > 1.0 - - prompt_embeds, text_encoder_hidden_states, _ = self._encode_prompt( - prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt - ) - - if isinstance(image_embeds, list): - image_embeds = torch.cat(image_embeds, dim=0) - if isinstance(negative_image_embeds, list): - negative_image_embeds = torch.cat(negative_image_embeds, dim=0) - - if do_classifier_free_guidance: - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - negative_image_embeds = negative_image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0).to( - dtype=prompt_embeds.dtype, device=device - ) - - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps_tensor = self.scheduler.timesteps - - num_channels_latents = self.unet.config.in_channels - - height, width = get_new_h_w(height, width, self.movq_scale_factor) - - # create initial latent - latents = self.prepare_latents( - (batch_size, num_channels_latents, height, width), - text_encoder_hidden_states.dtype, - device, - generator, - latents, - self.scheduler, - ) - - for i, t in enumerate(self.progress_bar(timesteps_tensor)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - added_cond_kwargs = {"text_embeds": prompt_embeds, "image_embeds": image_embeds} - noise_pred = self.unet( - sample=latent_model_input, - timestep=t, - encoder_hidden_states=text_encoder_hidden_states, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - if do_classifier_free_guidance: - noise_pred, variance_pred = noise_pred.split(latents.shape[1], dim=1) - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - _, variance_pred_text = variance_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, variance_pred_text], dim=1) - - if not ( - hasattr(self.scheduler.config, "variance_type") - and self.scheduler.config.variance_type in ["learned", "learned_range"] - ): - noise_pred, _ = noise_pred.split(latents.shape[1], dim=1) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step( - noise_pred, - t, - latents, - generator=generator, - ).prev_sample - - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # post-processing - image = self.movq.decode(latents, force_not_quantize=True)["sample"] - - self.maybe_free_model_hooks() - - if output_type not in ["pt", "np", "pil"]: - raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}") - - if output_type in ["np", "pil"]: - image = image * 0.5 + 0.5 - image = image.clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py deleted file mode 100644 index 90d4042ae2a12e2c5b27804cec9418e9f8c58aaf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py +++ /dev/null @@ -1,817 +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 Callable, List, Optional, Union - -import PIL.Image -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, - XLMRobertaTokenizer, -) - -from ...models import PriorTransformer, UNet2DConditionModel, VQModel -from ...schedulers import DDIMScheduler, DDPMScheduler, UnCLIPScheduler -from ...utils import ( - replace_example_docstring, -) -from ..pipeline_utils import DiffusionPipeline -from .pipeline_kandinsky import KandinskyPipeline -from .pipeline_kandinsky_img2img import KandinskyImg2ImgPipeline -from .pipeline_kandinsky_inpaint import KandinskyInpaintPipeline -from .pipeline_kandinsky_prior import KandinskyPriorPipeline -from .text_encoder import MultilingualCLIP - - -TEXT2IMAGE_EXAMPLE_DOC_STRING = """ - Examples: - ```py - from diffusers import AutoPipelineForText2Image - import torch - - pipe = AutoPipelineForText2Image.from_pretrained( - "kandinsky-community/kandinsky-2-1", torch_dtype=torch.float16 - ) - pipe.enable_model_cpu_offload() - - prompt = "A lion in galaxies, spirals, nebulae, stars, smoke, iridescent, intricate detail, octane render, 8k" - - image = pipe(prompt=prompt, num_inference_steps=25).images[0] - ``` -""" - -IMAGE2IMAGE_EXAMPLE_DOC_STRING = """ - Examples: - ```py - from diffusers import AutoPipelineForImage2Image - import torch - import requests - from io import BytesIO - from PIL import Image - import os - - pipe = AutoPipelineForImage2Image.from_pretrained( - "kandinsky-community/kandinsky-2-1", torch_dtype=torch.float16 - ) - pipe.enable_model_cpu_offload() - - prompt = "A fantasy landscape, Cinematic lighting" - negative_prompt = "low quality, bad quality" - - url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" - - response = requests.get(url) - image = Image.open(BytesIO(response.content)).convert("RGB") - image.thumbnail((768, 768)) - - image = pipe(prompt=prompt, image=original_image, num_inference_steps=25).images[0] - ``` -""" - -INPAINT_EXAMPLE_DOC_STRING = """ - Examples: - ```py - from diffusers import AutoPipelineForInpainting - from diffusers.utils import load_image - import torch - import numpy as np - - pipe = AutoPipelineForInpainting.from_pretrained( - "kandinsky-community/kandinsky-2-1-inpaint", torch_dtype=torch.float16 - ) - pipe.enable_model_cpu_offload() - - prompt = "A fantasy landscape, Cinematic lighting" - negative_prompt = "low quality, bad quality" - - original_image = load_image( - "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/kandinsky/cat.png" - ) - - mask = np.zeros((768, 768), dtype=np.float32) - # Let's mask out an area above the cat's head - mask[:250, 250:-250] = 1 - - image = pipe(prompt=prompt, image=original_image, mask_image=mask, num_inference_steps=25).images[0] - ``` -""" - - -class KandinskyCombinedPipeline(DiffusionPipeline): - """ - Combined Pipeline for text-to-image generation using Kandinsky - - 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.) - - Args: - text_encoder ([`MultilingualCLIP`]): - Frozen text-encoder. - tokenizer ([`XLMRobertaTokenizer`]): - Tokenizer of class - scheduler (Union[`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. - movq ([`VQModel`]): - MoVQ Decoder to generate the image from the latents. - prior_prior ([`PriorTransformer`]): - The canonical unCLIP prior to approximate the image embedding from the text embedding. - prior_image_encoder ([`CLIPVisionModelWithProjection`]): - Frozen image-encoder. - prior_text_encoder ([`CLIPTextModelWithProjection`]): - Frozen text-encoder. - prior_tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - prior_scheduler ([`UnCLIPScheduler`]): - A scheduler to be used in combination with `prior` to generate image embedding. - """ - - _load_connected_pipes = True - model_cpu_offload_seq = "text_encoder->unet->movq->prior_prior->prior_image_encoder->prior_text_encoder" - _exclude_from_cpu_offload = ["prior_prior"] - - def __init__( - self, - text_encoder: MultilingualCLIP, - tokenizer: XLMRobertaTokenizer, - unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, DDPMScheduler], - movq: VQModel, - prior_prior: PriorTransformer, - prior_image_encoder: CLIPVisionModelWithProjection, - prior_text_encoder: CLIPTextModelWithProjection, - prior_tokenizer: CLIPTokenizer, - prior_scheduler: UnCLIPScheduler, - prior_image_processor: CLIPImageProcessor, - ): - super().__init__() - - self.register_modules( - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - movq=movq, - prior_prior=prior_prior, - prior_image_encoder=prior_image_encoder, - prior_text_encoder=prior_text_encoder, - prior_tokenizer=prior_tokenizer, - prior_scheduler=prior_scheduler, - prior_image_processor=prior_image_processor, - ) - self.prior_pipe = KandinskyPriorPipeline( - prior=prior_prior, - image_encoder=prior_image_encoder, - text_encoder=prior_text_encoder, - tokenizer=prior_tokenizer, - scheduler=prior_scheduler, - image_processor=prior_image_processor, - ) - self.decoder_pipe = KandinskyPipeline( - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - movq=movq, - ) - - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = 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): - 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 - GPU only when their specific submodule's `forward` method is called. Offloading happens on a submodule basis. - Memory savings are higher than using `enable_model_cpu_offload`, but performance is lower. - """ - self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - - def progress_bar(self, iterable=None, total=None): - self.prior_pipe.progress_bar(iterable=iterable, total=total) - self.decoder_pipe.progress_bar(iterable=iterable, total=total) - self.decoder_pipe.enable_model_cpu_offload() - - def set_progress_bar_config(self, **kwargs): - self.prior_pipe.set_progress_bar_config(**kwargs) - self.decoder_pipe.set_progress_bar_config(**kwargs) - - @torch.no_grad() - @replace_example_docstring(TEXT2IMAGE_EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, - num_inference_steps: int = 100, - guidance_scale: float = 4.0, - num_images_per_prompt: int = 1, - height: int = 512, - 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, - callback_steps: int = 1, - return_dict: bool = True, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `List[str]`): - The prompt or prompts to guide the image generation. - 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): - 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. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - prior_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. - 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. - 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. - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - prior_outputs = self.prior_pipe( - prompt=prompt, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - num_inference_steps=prior_num_inference_steps, - generator=generator, - latents=latents, - guidance_scale=prior_guidance_scale, - output_type="pt", - return_dict=False, - ) - image_embeds = prior_outputs[0] - negative_image_embeds = prior_outputs[1] - - prompt = [prompt] if not isinstance(prompt, (list, tuple)) else prompt - - if len(prompt) < image_embeds.shape[0] and image_embeds.shape[0] % len(prompt) == 0: - prompt = (image_embeds.shape[0] // len(prompt)) * prompt - - outputs = self.decoder_pipe( - prompt=prompt, - image_embeds=image_embeds, - negative_image_embeds=negative_image_embeds, - width=width, - height=height, - num_inference_steps=num_inference_steps, - generator=generator, - guidance_scale=guidance_scale, - output_type=output_type, - callback=callback, - callback_steps=callback_steps, - return_dict=return_dict, - ) - - self.maybe_free_model_hooks() - - return outputs - - -class KandinskyImg2ImgCombinedPipeline(DiffusionPipeline): - """ - Combined Pipeline for image-to-image generation using Kandinsky - - 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.) - - Args: - text_encoder ([`MultilingualCLIP`]): - Frozen text-encoder. - tokenizer ([`XLMRobertaTokenizer`]): - Tokenizer of class - scheduler (Union[`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. - movq ([`VQModel`]): - MoVQ Decoder to generate the image from the latents. - prior_prior ([`PriorTransformer`]): - The canonical unCLIP prior to approximate the image embedding from the text embedding. - prior_image_encoder ([`CLIPVisionModelWithProjection`]): - Frozen image-encoder. - prior_text_encoder ([`CLIPTextModelWithProjection`]): - Frozen text-encoder. - prior_tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - prior_scheduler ([`UnCLIPScheduler`]): - A scheduler to be used in combination with `prior` to generate image embedding. - """ - - _load_connected_pipes = True - model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->prior_prior->text_encoder->unet->movq" - _exclude_from_cpu_offload = ["prior_prior"] - - def __init__( - self, - text_encoder: MultilingualCLIP, - tokenizer: XLMRobertaTokenizer, - unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, DDPMScheduler], - movq: VQModel, - prior_prior: PriorTransformer, - prior_image_encoder: CLIPVisionModelWithProjection, - prior_text_encoder: CLIPTextModelWithProjection, - prior_tokenizer: CLIPTokenizer, - prior_scheduler: UnCLIPScheduler, - prior_image_processor: CLIPImageProcessor, - ): - super().__init__() - - self.register_modules( - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - movq=movq, - prior_prior=prior_prior, - prior_image_encoder=prior_image_encoder, - prior_text_encoder=prior_text_encoder, - prior_tokenizer=prior_tokenizer, - prior_scheduler=prior_scheduler, - prior_image_processor=prior_image_processor, - ) - self.prior_pipe = KandinskyPriorPipeline( - prior=prior_prior, - image_encoder=prior_image_encoder, - text_encoder=prior_text_encoder, - tokenizer=prior_tokenizer, - scheduler=prior_scheduler, - image_processor=prior_image_processor, - ) - self.decoder_pipe = KandinskyImg2ImgPipeline( - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - movq=movq, - ) - - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = 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): - 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 - `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called. - Note that offloading happens on a submodule basis. Memory savings are higher than with - `enable_model_cpu_offload`, but performance is lower. - """ - self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - - def progress_bar(self, iterable=None, total=None): - self.prior_pipe.progress_bar(iterable=iterable, total=total) - self.decoder_pipe.progress_bar(iterable=iterable, total=total) - self.decoder_pipe.enable_model_cpu_offload() - - def set_progress_bar_config(self, **kwargs): - self.prior_pipe.set_progress_bar_config(**kwargs) - self.decoder_pipe.set_progress_bar_config(**kwargs) - - @torch.no_grad() - @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, - num_inference_steps: int = 100, - guidance_scale: float = 4.0, - num_images_per_prompt: int = 1, - strength: float = 0.3, - height: int = 512, - 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, - callback_steps: int = 1, - return_dict: bool = True, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - 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`, 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*): - 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): - 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. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - strength (`float`, *optional*, defaults to 0.3): - 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`. - prior_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. - 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. - 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. - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - prior_outputs = self.prior_pipe( - prompt=prompt, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - num_inference_steps=prior_num_inference_steps, - generator=generator, - latents=latents, - guidance_scale=prior_guidance_scale, - output_type="pt", - return_dict=False, - ) - image_embeds = prior_outputs[0] - negative_image_embeds = prior_outputs[1] - - prompt = [prompt] if not isinstance(prompt, (list, tuple)) else prompt - image = [image] if isinstance(prompt, PIL.Image.Image) else image - - if len(prompt) < image_embeds.shape[0] and image_embeds.shape[0] % len(prompt) == 0: - prompt = (image_embeds.shape[0] // len(prompt)) * prompt - - if ( - isinstance(image, (list, tuple)) - and len(image) < image_embeds.shape[0] - and image_embeds.shape[0] % len(image) == 0 - ): - image = (image_embeds.shape[0] // len(image)) * image - - outputs = self.decoder_pipe( - prompt=prompt, - image=image, - image_embeds=image_embeds, - negative_image_embeds=negative_image_embeds, - strength=strength, - width=width, - height=height, - num_inference_steps=num_inference_steps, - generator=generator, - guidance_scale=guidance_scale, - output_type=output_type, - callback=callback, - callback_steps=callback_steps, - return_dict=return_dict, - ) - - self.maybe_free_model_hooks() - - return outputs - - -class KandinskyInpaintCombinedPipeline(DiffusionPipeline): - """ - Combined Pipeline for generation using Kandinsky - - 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.) - - Args: - text_encoder ([`MultilingualCLIP`]): - Frozen text-encoder. - tokenizer ([`XLMRobertaTokenizer`]): - Tokenizer of class - scheduler (Union[`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. - movq ([`VQModel`]): - MoVQ Decoder to generate the image from the latents. - prior_prior ([`PriorTransformer`]): - The canonical unCLIP prior to approximate the image embedding from the text embedding. - prior_image_encoder ([`CLIPVisionModelWithProjection`]): - Frozen image-encoder. - prior_text_encoder ([`CLIPTextModelWithProjection`]): - Frozen text-encoder. - prior_tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - prior_scheduler ([`UnCLIPScheduler`]): - A scheduler to be used in combination with `prior` to generate image embedding. - """ - - _load_connected_pipes = True - model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->prior_prior->text_encoder->unet->movq" - _exclude_from_cpu_offload = ["prior_prior"] - - def __init__( - self, - text_encoder: MultilingualCLIP, - tokenizer: XLMRobertaTokenizer, - unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, DDPMScheduler], - movq: VQModel, - prior_prior: PriorTransformer, - prior_image_encoder: CLIPVisionModelWithProjection, - prior_text_encoder: CLIPTextModelWithProjection, - prior_tokenizer: CLIPTokenizer, - prior_scheduler: UnCLIPScheduler, - prior_image_processor: CLIPImageProcessor, - ): - super().__init__() - - self.register_modules( - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - movq=movq, - prior_prior=prior_prior, - prior_image_encoder=prior_image_encoder, - prior_text_encoder=prior_text_encoder, - prior_tokenizer=prior_tokenizer, - prior_scheduler=prior_scheduler, - prior_image_processor=prior_image_processor, - ) - self.prior_pipe = KandinskyPriorPipeline( - prior=prior_prior, - image_encoder=prior_image_encoder, - text_encoder=prior_text_encoder, - tokenizer=prior_tokenizer, - scheduler=prior_scheduler, - image_processor=prior_image_processor, - ) - self.decoder_pipe = KandinskyInpaintPipeline( - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - movq=movq, - ) - - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = 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): - 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 - `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called. - Note that offloading happens on a submodule basis. Memory savings are higher than with - `enable_model_cpu_offload`, but performance is lower. - """ - self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - - def progress_bar(self, iterable=None, total=None): - self.prior_pipe.progress_bar(iterable=iterable, total=total) - self.decoder_pipe.progress_bar(iterable=iterable, total=total) - self.decoder_pipe.enable_model_cpu_offload() - - def set_progress_bar_config(self, **kwargs): - self.prior_pipe.set_progress_bar_config(**kwargs) - self.decoder_pipe.set_progress_bar_config(**kwargs) - - @torch.no_grad() - @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, - num_inference_steps: int = 100, - guidance_scale: float = 4.0, - num_images_per_prompt: int = 1, - height: int = 512, - 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, - callback_steps: int = 1, - return_dict: bool = True, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - 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`, 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. - mask_image (`np.array`): - Tensor representing an image batch, to mask `image`. White pixels in the mask will be 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)`. - 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): - 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. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - prior_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. - 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. - 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. - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - prior_outputs = self.prior_pipe( - prompt=prompt, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - num_inference_steps=prior_num_inference_steps, - generator=generator, - latents=latents, - guidance_scale=prior_guidance_scale, - output_type="pt", - return_dict=False, - ) - image_embeds = prior_outputs[0] - negative_image_embeds = prior_outputs[1] - - prompt = [prompt] if not isinstance(prompt, (list, tuple)) else prompt - image = [image] if isinstance(prompt, PIL.Image.Image) else image - mask_image = [mask_image] if isinstance(mask_image, PIL.Image.Image) else mask_image - - if len(prompt) < image_embeds.shape[0] and image_embeds.shape[0] % len(prompt) == 0: - prompt = (image_embeds.shape[0] // len(prompt)) * prompt - - if ( - isinstance(image, (list, tuple)) - and len(image) < image_embeds.shape[0] - and image_embeds.shape[0] % len(image) == 0 - ): - image = (image_embeds.shape[0] // len(image)) * image - - if ( - isinstance(mask_image, (list, tuple)) - and len(mask_image) < image_embeds.shape[0] - and image_embeds.shape[0] % len(mask_image) == 0 - ): - mask_image = (image_embeds.shape[0] // len(mask_image)) * mask_image - - outputs = self.decoder_pipe( - prompt=prompt, - image=image, - mask_image=mask_image, - image_embeds=image_embeds, - negative_image_embeds=negative_image_embeds, - width=width, - height=height, - num_inference_steps=num_inference_steps, - generator=generator, - guidance_scale=guidance_scale, - output_type=output_type, - callback=callback, - callback_steps=callback_steps, - return_dict=return_dict, - ) - - self.maybe_free_model_hooks() - - return outputs diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py deleted file mode 100644 index 998fc777c022d02f006c7355b9cccac57723e9a4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py +++ /dev/null @@ -1,505 +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 Callable, List, Optional, Union - -import PIL.Image -import torch -from transformers import ( - XLMRobertaTokenizer, -) - -from ...image_processor import VaeImageProcessor -from ...models import UNet2DConditionModel, VQModel -from ...schedulers import DDIMScheduler -from ...utils import ( - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput -from .text_encoder import MultilingualCLIP - - -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 - >>> from diffusers import KandinskyImg2ImgPipeline, KandinskyPriorPipeline - >>> from diffusers.utils import load_image - >>> import torch - - >>> pipe_prior = KandinskyPriorPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-1-prior", torch_dtype=torch.float16 - ... ) - >>> pipe_prior.to("cuda") - - >>> prompt = "A red cartoon frog, 4k" - >>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False) - - >>> pipe = KandinskyImg2ImgPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-1", torch_dtype=torch.float16 - ... ) - >>> pipe.to("cuda") - - >>> init_image = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/frog.png" - ... ) - - >>> image = pipe( - ... prompt, - ... image=init_image, - ... image_embeds=image_emb, - ... negative_image_embeds=zero_image_emb, - ... height=768, - ... width=768, - ... num_inference_steps=100, - ... strength=0.2, - ... ).images - - >>> image[0].save("red_frog.png") - ``` -""" - - -def get_new_h_w(h, w, scale_factor=8): - new_h = h // scale_factor**2 - if h % scale_factor**2 != 0: - new_h += 1 - new_w = w // scale_factor**2 - if w % scale_factor**2 != 0: - new_w += 1 - return new_h * scale_factor, new_w * scale_factor - - -class KandinskyImg2ImgPipeline(DiffusionPipeline): - """ - Pipeline for image-to-image generation using Kandinsky - - 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.) - - Args: - text_encoder ([`MultilingualCLIP`]): - Frozen text-encoder. - tokenizer ([`XLMRobertaTokenizer`]): - Tokenizer of class - scheduler ([`DDIMScheduler`]): - A scheduler to be used in combination with `unet` to generate image latents. - unet ([`UNet2DConditionModel`]): - Conditional U-Net architecture to denoise the image embedding. - movq ([`VQModel`]): - MoVQ image encoder and decoder - """ - - model_cpu_offload_seq = "text_encoder->unet->movq" - - def __init__( - self, - text_encoder: MultilingualCLIP, - movq: VQModel, - tokenizer: XLMRobertaTokenizer, - unet: UNet2DConditionModel, - scheduler: DDIMScheduler, - ): - super().__init__() - - self.register_modules( - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - movq=movq, - ) - self.movq_scale_factor = ( - 2 ** (len(self.movq.config.block_out_channels) - 1) if getattr(self, "movq", None) else 8 - ) - movq_latent_channels = self.movq.config.latent_channels if getattr(self, "movq", None) else 4 - self.image_processor = VaeImageProcessor( - vae_scale_factor=self.movq_scale_factor, - vae_latent_channels=movq_latent_channels, - resample="bicubic", - reducing_gap=1, - ) - - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - timesteps = self.scheduler.timesteps[t_start:] - - return timesteps, num_inference_steps - t_start - - def prepare_latents(self, latents, latent_timestep, shape, dtype, device, generator, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - - shape = latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - latents = self.add_noise(latents, noise, latent_timestep) - return latents - - def _encode_prompt( - self, - prompt, - device, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt=None, - ): - batch_size = len(prompt) if isinstance(prompt, list) else 1 - # get prompt text embeddings - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=77, - truncation=True, - return_attention_mask=True, - add_special_tokens=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}" - ) - - text_input_ids = text_input_ids.to(device) - text_mask = text_inputs.attention_mask.to(device) - - prompt_embeds, text_encoder_hidden_states = self.text_encoder( - input_ids=text_input_ids, attention_mask=text_mask - ) - - prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - text_encoder_hidden_states = text_encoder_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - uncond_tokens: List[str] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif 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 - - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=77, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - uncond_text_input_ids = uncond_input.input_ids.to(device) - uncond_text_mask = uncond_input.attention_mask.to(device) - - negative_prompt_embeds, uncond_text_encoder_hidden_states = self.text_encoder( - input_ids=uncond_text_input_ids, attention_mask=uncond_text_mask - ) - - # 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.repeat(1, num_images_per_prompt) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len) - - seq_len = uncond_text_encoder_hidden_states.shape[1] - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.repeat(1, num_images_per_prompt, 1) - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.view( - batch_size * num_images_per_prompt, seq_len, -1 - ) - uncond_text_mask = uncond_text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - # done duplicates - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - text_encoder_hidden_states = torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]) - - text_mask = torch.cat([uncond_text_mask, text_mask]) - - return prompt_embeds, text_encoder_hidden_states, text_mask - - # add_noise method to overwrite the one in schedule because it use a different beta schedule for adding noise vs sampling - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - betas = torch.linspace(0.0001, 0.02, 1000, dtype=torch.float32) - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - alphas_cumprod = alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype) - timesteps = timesteps.to(original_samples.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(original_samples.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise - - return noisy_samples - - @torch.no_grad() - @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]], - image_embeds: torch.Tensor, - negative_image_embeds: torch.Tensor, - negative_prompt: Optional[Union[str, List[str]]] = 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, - callback_steps: int = 1, - return_dict: bool = True, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - 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]`): - 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]`): - The clip image embeddings for negative text prompt, will be used to condition the image generation. - 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): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - 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. - strength (`float`, *optional*, defaults to 0.3): - 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`. - 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. - 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. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - # 1. Define call parameters - if isinstance(prompt, str): - batch_size = 1 - elif isinstance(prompt, list): - batch_size = len(prompt) - else: - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - device = self._execution_device - - batch_size = batch_size * num_images_per_prompt - - do_classifier_free_guidance = guidance_scale > 1.0 - - # 2. get text and image embeddings - prompt_embeds, text_encoder_hidden_states, _ = self._encode_prompt( - prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt - ) - - if isinstance(image_embeds, list): - image_embeds = torch.cat(image_embeds, dim=0) - if isinstance(negative_image_embeds, list): - negative_image_embeds = torch.cat(negative_image_embeds, dim=0) - - if do_classifier_free_guidance: - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - negative_image_embeds = negative_image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0).to( - dtype=prompt_embeds.dtype, device=device - ) - - # 3. pre-processing initial image - if not isinstance(image, list): - image = [image] - if not all(isinstance(i, (PIL.Image.Image, torch.Tensor)) for i in image): - raise ValueError( - f"Input is in incorrect format: {[type(i) for i in image]}. Currently, we only support PIL image and pytorch tensor" - ) - - image = torch.cat([self.image_processor.preprocess(i, width, height) for i in image], dim=0) - image = image.to(dtype=prompt_embeds.dtype, device=device) - - latents = self.movq.encode(image)["latents"] - latents = latents.repeat_interleave(num_images_per_prompt, dim=0) - - # 4. set timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - - timesteps_tensor, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) - - # the formular to calculate timestep for add_noise is taken from the original kandinsky repo - latent_timestep = int(self.scheduler.config.num_train_timesteps * strength) - 2 - - latent_timestep = torch.tensor([latent_timestep] * batch_size, dtype=timesteps_tensor.dtype, device=device) - - num_channels_latents = self.unet.config.in_channels - - height, width = get_new_h_w(height, width, self.movq_scale_factor) - - # 5. Create initial latent - latents = self.prepare_latents( - latents, - latent_timestep, - (batch_size, num_channels_latents, height, width), - text_encoder_hidden_states.dtype, - device, - generator, - self.scheduler, - ) - - # 6. Denoising loop - for i, t in enumerate(self.progress_bar(timesteps_tensor)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - added_cond_kwargs = {"text_embeds": prompt_embeds, "image_embeds": image_embeds} - noise_pred = self.unet( - sample=latent_model_input, - timestep=t, - encoder_hidden_states=text_encoder_hidden_states, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - if do_classifier_free_guidance: - noise_pred, variance_pred = noise_pred.split(latents.shape[1], dim=1) - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - _, variance_pred_text = variance_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, variance_pred_text], dim=1) - - if not ( - hasattr(self.scheduler.config, "variance_type") - and self.scheduler.config.variance_type in ["learned", "learned_range"] - ): - noise_pred, _ = noise_pred.split(latents.shape[1], dim=1) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step( - noise_pred, - t, - latents, - generator=generator, - ).prev_sample - - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # 7. post-processing - image = self.movq.decode(latents, force_not_quantize=True)["sample"] - - self.maybe_free_model_hooks() - - if output_type not in ["pt", "np", "pil"]: - raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}") - - image = self.image_processor.postprocess(image, output_type) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py deleted file mode 100644 index 5645d2a56edd9e311cea78419a8c910ae2185b21..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py +++ /dev/null @@ -1,647 +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 copy import deepcopy -from typing import Callable, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from packaging import version -from PIL import Image -from transformers import ( - XLMRobertaTokenizer, -) - -from ... import __version__ -from ...models import UNet2DConditionModel, VQModel -from ...schedulers import DDIMScheduler -from ...utils import ( - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput -from .text_encoder import MultilingualCLIP - - -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 - >>> from diffusers import KandinskyInpaintPipeline, KandinskyPriorPipeline - >>> from diffusers.utils import load_image - >>> import torch - >>> import numpy as np - - >>> pipe_prior = KandinskyPriorPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-1-prior", torch_dtype=torch.float16 - ... ) - >>> pipe_prior.to("cuda") - - >>> prompt = "a hat" - >>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False) - - >>> pipe = KandinskyInpaintPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-1-inpaint", torch_dtype=torch.float16 - ... ) - >>> pipe.to("cuda") - - >>> init_image = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/cat.png" - ... ) - - >>> mask = np.zeros((768, 768), dtype=np.float32) - >>> mask[:250, 250:-250] = 1 - - >>> out = pipe( - ... prompt, - ... image=init_image, - ... mask_image=mask, - ... image_embeds=image_emb, - ... negative_image_embeds=zero_image_emb, - ... height=768, - ... width=768, - ... num_inference_steps=50, - ... ) - - >>> image = out.images[0] - >>> image.save("cat_with_hat.png") - ``` -""" - - -def get_new_h_w(h, w, scale_factor=8): - new_h = h // scale_factor**2 - if h % scale_factor**2 != 0: - new_h += 1 - new_w = w // scale_factor**2 - if w % scale_factor**2 != 0: - new_w += 1 - return new_h * scale_factor, new_w * scale_factor - - -def prepare_mask(masks): - prepared_masks = [] - for mask in masks: - old_mask = deepcopy(mask) - for i in range(mask.shape[1]): - for j in range(mask.shape[2]): - if old_mask[0][i][j] == 1: - continue - if i != 0: - mask[:, i - 1, j] = 0 - if j != 0: - mask[:, i, j - 1] = 0 - if i != 0 and j != 0: - mask[:, i - 1, j - 1] = 0 - if i != mask.shape[1] - 1: - mask[:, i + 1, j] = 0 - if j != mask.shape[2] - 1: - mask[:, i, j + 1] = 0 - if i != mask.shape[1] - 1 and j != mask.shape[2] - 1: - mask[:, i + 1, j + 1] = 0 - prepared_masks.append(mask) - return torch.stack(prepared_masks, dim=0) - - -def prepare_mask_and_masked_image(image, mask, height, width): - r""" - Prepares a pair (mask, image) to be consumed by the Kandinsky inpaint pipeline. This means that those inputs will - be converted to ``torch.Tensor`` with shapes ``batch x channels x height x width`` where ``channels`` is ``3`` for - the ``image`` and ``1`` for the ``mask``. - - The ``image`` will be converted to ``torch.float32`` and normalized to be in ``[-1, 1]``. The ``mask`` will be - binarized (``mask > 0.5``) and cast to ``torch.float32`` too. - - Args: - image (Union[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. - It can be a ``PIL.Image``, or a ``height x width`` ``np.array`` or a ``1 x height x width`` - ``torch.Tensor`` or a ``batch x 1 x height x width`` ``torch.Tensor``. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - - - Raises: - ValueError: ``torch.Tensor`` images should be in the ``[-1, 1]`` range. ValueError: ``torch.Tensor`` mask - should be in the ``[0, 1]`` range. ValueError: ``mask`` and ``image`` should have the same spatial dimensions. - TypeError: ``mask`` is a ``torch.Tensor`` but ``image`` is not - (ot the other way around). - - Returns: - tuple[torch.Tensor]: The pair (mask, image) as ``torch.Tensor`` with 4 - dimensions: ``batch x channels x height x width``. - """ - - if image is None: - raise ValueError("`image` input cannot be undefined.") - - if mask is None: - raise ValueError("`mask_image` input cannot be undefined.") - - if isinstance(image, torch.Tensor): - if not isinstance(mask, torch.Tensor): - raise TypeError(f"`image` is a torch.Tensor but `mask` (type: {type(mask)} is not") - - # Batch single image - if image.ndim == 3: - assert image.shape[0] == 3, "Image outside a batch should be of shape (3, H, W)" - image = image.unsqueeze(0) - - # Batch and add channel dim for single mask - if mask.ndim == 2: - mask = mask.unsqueeze(0).unsqueeze(0) - - # Batch single mask or add channel dim - if mask.ndim == 3: - # Single batched mask, no channel dim or single mask not batched but channel dim - if mask.shape[0] == 1: - mask = mask.unsqueeze(0) - - # Batched masks no channel dim - else: - mask = mask.unsqueeze(1) - - assert image.ndim == 4 and mask.ndim == 4, "Image and Mask must have 4 dimensions" - assert image.shape[-2:] == mask.shape[-2:], "Image and Mask must have the same spatial dimensions" - assert image.shape[0] == mask.shape[0], "Image and Mask must have the same batch size" - - # Check image is in [-1, 1] - if image.min() < -1 or image.max() > 1: - raise ValueError("Image should be in [-1, 1] range") - - # Check mask is in [0, 1] - if mask.min() < 0 or mask.max() > 1: - raise ValueError("Mask should be in [0, 1] range") - - # Binarize mask - mask[mask < 0.5] = 0 - mask[mask >= 0.5] = 1 - - # Image as float32 - image = image.to(dtype=torch.float32) - elif isinstance(mask, torch.Tensor): - raise TypeError(f"`mask` is a torch.Tensor but `image` (type: {type(image)} is not") - else: - # preprocess image - if isinstance(image, (PIL.Image.Image, np.ndarray)): - image = [image] - - if isinstance(image, list) and isinstance(image[0], PIL.Image.Image): - # resize all images w.r.t passed height an width - image = [i.resize((width, height), resample=Image.BICUBIC, reducing_gap=1) for i in image] - image = [np.array(i.convert("RGB"))[None, :] for i in image] - image = np.concatenate(image, axis=0) - elif isinstance(image, list) and isinstance(image[0], np.ndarray): - image = np.concatenate([i[None, :] for i in image], axis=0) - - image = image.transpose(0, 3, 1, 2) - image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 - - # preprocess mask - if isinstance(mask, (PIL.Image.Image, np.ndarray)): - mask = [mask] - - if isinstance(mask, list) and isinstance(mask[0], PIL.Image.Image): - mask = [i.resize((width, height), resample=PIL.Image.LANCZOS) for i in mask] - mask = np.concatenate([np.array(m.convert("L"))[None, None, :] for m in mask], axis=0) - mask = mask.astype(np.float32) / 255.0 - elif isinstance(mask, list) and isinstance(mask[0], np.ndarray): - mask = np.concatenate([m[None, None, :] for m in mask], axis=0) - - mask[mask < 0.5] = 0 - mask[mask >= 0.5] = 1 - mask = torch.from_numpy(mask) - - mask = 1 - mask - - return mask, image - - -class KandinskyInpaintPipeline(DiffusionPipeline): - """ - Pipeline for text-guided image inpainting using Kandinsky2.1 - - 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.) - - Args: - text_encoder ([`MultilingualCLIP`]): - Frozen text-encoder. - tokenizer ([`XLMRobertaTokenizer`]): - Tokenizer of class - scheduler ([`DDIMScheduler`]): - A scheduler to be used in combination with `unet` to generate image latents. - unet ([`UNet2DConditionModel`]): - Conditional U-Net architecture to denoise the image embedding. - movq ([`VQModel`]): - MoVQ image encoder and decoder - """ - - model_cpu_offload_seq = "text_encoder->unet->movq" - - def __init__( - self, - text_encoder: MultilingualCLIP, - movq: VQModel, - tokenizer: XLMRobertaTokenizer, - unet: UNet2DConditionModel, - scheduler: DDIMScheduler, - ): - super().__init__() - - self.register_modules( - text_encoder=text_encoder, - movq=movq, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - ) - self.movq_scale_factor = 2 ** (len(self.movq.config.block_out_channels) - 1) - self._warn_has_been_called = False - - # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - def _encode_prompt( - self, - prompt, - device, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt=None, - ): - batch_size = len(prompt) if isinstance(prompt, list) else 1 - # get prompt text embeddings - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=77, - truncation=True, - return_attention_mask=True, - add_special_tokens=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}" - ) - - text_input_ids = text_input_ids.to(device) - text_mask = text_inputs.attention_mask.to(device) - - prompt_embeds, text_encoder_hidden_states = self.text_encoder( - input_ids=text_input_ids, attention_mask=text_mask - ) - - prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - text_encoder_hidden_states = text_encoder_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - uncond_tokens: List[str] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif 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 - - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=77, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - uncond_text_input_ids = uncond_input.input_ids.to(device) - uncond_text_mask = uncond_input.attention_mask.to(device) - - negative_prompt_embeds, uncond_text_encoder_hidden_states = self.text_encoder( - input_ids=uncond_text_input_ids, attention_mask=uncond_text_mask - ) - - # 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.repeat(1, num_images_per_prompt) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len) - - seq_len = uncond_text_encoder_hidden_states.shape[1] - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.repeat(1, num_images_per_prompt, 1) - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.view( - batch_size * num_images_per_prompt, seq_len, -1 - ) - uncond_text_mask = uncond_text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - # done duplicates - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - text_encoder_hidden_states = torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]) - - text_mask = torch.cat([uncond_text_mask, text_mask]) - - return prompt_embeds, text_encoder_hidden_states, text_mask - - @torch.no_grad() - @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], - image_embeds: torch.Tensor, - negative_image_embeds: torch.Tensor, - negative_prompt: Optional[Union[str, List[str]]] = 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, - callback_steps: int = 1, - return_dict: bool = True, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - 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 - process. - mask_image (`PIL.Image.Image`,`torch.Tensor` or `np.ndarray`): - `Image`, or a tensor representing an image batch, to mask `image`. White pixels in the mask will be - repainted, while black pixels will be preserved. You can pass a pytorch tensor as mask only if the - image you passed is a pytorch tensor, and it should contain one color channel (L) instead of 3, so the - expected shape would be either `(B, 1, H, W,)`, `(B, H, W)`, `(1, H, W)` or `(H, W)` If image is an PIL - 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]`): - 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]`): - The clip image embeddings for negative text prompt, will be used to condition the image generation. - 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): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - 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. - 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. - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - if not self._warn_has_been_called and version.parse(version.parse(__version__).base_version) < version.parse( - "0.23.0.dev0" - ): - logger.warning( - "Please note that the expected format of `mask_image` has recently been changed. " - "Before diffusers == 0.19.0, Kandinsky Inpainting pipelines repainted black pixels and preserved black pixels. " - "As of diffusers==0.19.0 this behavior has been inverted. Now white pixels are repainted and black pixels are preserved. " - "This way, Kandinsky's masking behavior is aligned with Stable Diffusion. " - "THIS means that you HAVE to invert the input mask to have the same behavior as before as explained in https://github.com/huggingface/diffusers/pull/4207. " - "This warning will be suppressed after the first inference call and will be removed in diffusers>0.23.0" - ) - self._warn_has_been_called = True - - # Define call parameters - if isinstance(prompt, str): - batch_size = 1 - elif isinstance(prompt, list): - batch_size = len(prompt) - else: - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - device = self._execution_device - - batch_size = batch_size * num_images_per_prompt - do_classifier_free_guidance = guidance_scale > 1.0 - - prompt_embeds, text_encoder_hidden_states, _ = self._encode_prompt( - prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt - ) - - if isinstance(image_embeds, list): - image_embeds = torch.cat(image_embeds, dim=0) - if isinstance(negative_image_embeds, list): - negative_image_embeds = torch.cat(negative_image_embeds, dim=0) - - if do_classifier_free_guidance: - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - negative_image_embeds = negative_image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0).to( - dtype=prompt_embeds.dtype, device=device - ) - - # preprocess image and mask - mask_image, image = prepare_mask_and_masked_image(image, mask_image, height, width) - - image = image.to(dtype=prompt_embeds.dtype, device=device) - image = self.movq.encode(image)["latents"] - - mask_image = mask_image.to(dtype=prompt_embeds.dtype, device=device) - - image_shape = tuple(image.shape[-2:]) - mask_image = F.interpolate( - mask_image, - image_shape, - mode="nearest", - ) - mask_image = prepare_mask(mask_image) - masked_image = image * mask_image - - mask_image = mask_image.repeat_interleave(num_images_per_prompt, dim=0) - masked_image = masked_image.repeat_interleave(num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - mask_image = mask_image.repeat(2, 1, 1, 1) - masked_image = masked_image.repeat(2, 1, 1, 1) - - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps_tensor = self.scheduler.timesteps - - num_channels_latents = self.movq.config.latent_channels - - # get h, w for latents - sample_height, sample_width = get_new_h_w(height, width, self.movq_scale_factor) - - # create initial latent - latents = self.prepare_latents( - (batch_size, num_channels_latents, sample_height, sample_width), - text_encoder_hidden_states.dtype, - device, - generator, - latents, - self.scheduler, - ) - - # Check that sizes of mask, masked image and latents match with expected - num_channels_mask = mask_image.shape[1] - num_channels_masked_image = masked_image.shape[1] - if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" - " `pipeline.unet` or your `mask_image` or `image` input." - ) - - for i, t in enumerate(self.progress_bar(timesteps_tensor)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = torch.cat([latent_model_input, masked_image, mask_image], dim=1) - - added_cond_kwargs = {"text_embeds": prompt_embeds, "image_embeds": image_embeds} - noise_pred = self.unet( - sample=latent_model_input, - timestep=t, - encoder_hidden_states=text_encoder_hidden_states, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - if do_classifier_free_guidance: - noise_pred, variance_pred = noise_pred.split(latents.shape[1], dim=1) - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - _, variance_pred_text = variance_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, variance_pred_text], dim=1) - - if not ( - hasattr(self.scheduler.config, "variance_type") - and self.scheduler.config.variance_type in ["learned", "learned_range"] - ): - noise_pred, _ = noise_pred.split(latents.shape[1], dim=1) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step( - noise_pred, - t, - latents, - generator=generator, - ).prev_sample - - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # post-processing - image = self.movq.decode(latents, force_not_quantize=True)["sample"] - - self.maybe_free_model_hooks() - - if output_type not in ["pt", "np", "pil"]: - raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}") - - if output_type in ["np", "pil"]: - image = image * 0.5 + 0.5 - image = image.clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py deleted file mode 100644 index 8781d706edf5635e9c296b785d025fc13e7b220d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py +++ /dev/null @@ -1,559 +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 dataclasses import dataclass -from typing import List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...models import PriorTransformer -from ...schedulers import UnCLIPScheduler -from ...utils import ( - BaseOutput, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -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 - >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline - >>> import torch - - >>> pipe_prior = KandinskyPriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-1-prior") - >>> pipe_prior.to("cuda") - - >>> prompt = "red cat, 4k photo" - >>> out = pipe_prior(prompt) - >>> image_emb = out.image_embeds - >>> negative_image_emb = out.negative_image_embeds - - >>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1") - >>> pipe.to("cuda") - - >>> image = pipe( - ... prompt, - ... image_embeds=image_emb, - ... negative_image_embeds=negative_image_emb, - ... height=768, - ... width=768, - ... num_inference_steps=100, - ... ).images - - >>> image[0].save("cat.png") - ``` -""" - -EXAMPLE_INTERPOLATE_DOC_STRING = """ - Examples: - ```py - >>> from diffusers import KandinskyPriorPipeline, KandinskyPipeline - >>> from diffusers.utils import load_image - >>> import PIL - - >>> import torch - >>> from torchvision import transforms - - >>> pipe_prior = KandinskyPriorPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-1-prior", torch_dtype=torch.float16 - ... ) - >>> pipe_prior.to("cuda") - - >>> img1 = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/cat.png" - ... ) - - >>> img2 = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/starry_night.jpeg" - ... ) - - >>> images_texts = ["a cat", img1, img2] - >>> weights = [0.3, 0.3, 0.4] - >>> image_emb, zero_image_emb = pipe_prior.interpolate(images_texts, weights) - - >>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1", torch_dtype=torch.float16) - >>> pipe.to("cuda") - - >>> image = pipe( - ... "", - ... image_embeds=image_emb, - ... negative_image_embeds=zero_image_emb, - ... height=768, - ... width=768, - ... num_inference_steps=150, - ... ).images[0] - - >>> image.save("starry_cat.png") - ``` -""" - - -@dataclass -class KandinskyPriorPipelineOutput(BaseOutput): - """ - Output class for KandinskyPriorPipeline. - - Args: - image_embeds (`torch.Tensor`) - clip image embeddings for text prompt - 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] - - -class KandinskyPriorPipeline(DiffusionPipeline): - """ - Pipeline for generating image prior for Kandinsky - - 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.) - - Args: - prior ([`PriorTransformer`]): - The canonical unCLIP prior to approximate the image embedding from the text embedding. - image_encoder ([`CLIPVisionModelWithProjection`]): - Frozen image-encoder. - text_encoder ([`CLIPTextModelWithProjection`]): - Frozen text-encoder. - tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - scheduler ([`UnCLIPScheduler`]): - A scheduler to be used in combination with `prior` to generate image embedding. - """ - - _exclude_from_cpu_offload = ["prior"] - model_cpu_offload_seq = "text_encoder->prior" - - def __init__( - self, - prior: PriorTransformer, - image_encoder: CLIPVisionModelWithProjection, - text_encoder: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - scheduler: UnCLIPScheduler, - image_processor: CLIPImageProcessor, - ): - super().__init__() - - self.register_modules( - prior=prior, - text_encoder=text_encoder, - tokenizer=tokenizer, - scheduler=scheduler, - image_encoder=image_encoder, - image_processor=image_processor, - ) - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_INTERPOLATE_DOC_STRING) - def interpolate( - self, - images_and_prompts: List[Union[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, - negative_prompt: str = "", - guidance_scale: float = 4.0, - device=None, - ): - """ - Function invoked when using the prior pipeline for interpolation. - - Args: - images_and_prompts (`List[Union[str, PIL.Image.Image, torch.Tensor]]`): - list of prompts and images to guide the image generation. - 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*): - 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 ge generated by sampling using the supplied random `generator`. - 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*): - 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): - 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. - - Examples: - - Returns: - [`KandinskyPriorPipelineOutput`] or `tuple` - """ - - device = device or self.device - - if len(images_and_prompts) != len(weights): - raise ValueError( - f"`images_and_prompts` contains {len(images_and_prompts)} items and `weights` contains {len(weights)} items - they should be lists of same length" - ) - - image_embeddings = [] - for cond, weight in zip(images_and_prompts, weights): - if isinstance(cond, str): - image_emb = self( - cond, - num_inference_steps=num_inference_steps, - num_images_per_prompt=num_images_per_prompt, - generator=generator, - latents=latents, - negative_prompt=negative_prior_prompt, - guidance_scale=guidance_scale, - ).image_embeds - - elif isinstance(cond, (PIL.Image.Image, torch.Tensor)): - if isinstance(cond, PIL.Image.Image): - cond = ( - self.image_processor(cond, return_tensors="pt") - .pixel_values[0] - .unsqueeze(0) - .to(dtype=self.image_encoder.dtype, device=device) - ) - - image_emb = self.image_encoder(cond)["image_embeds"] - - else: - raise ValueError( - f"`images_and_prompts` can only contains elements to be of type `str`, `PIL.Image.Image` or `torch.Tensor` but is {type(cond)}" - ) - - image_embeddings.append(image_emb * weight) - - image_emb = torch.cat(image_embeddings).sum(dim=0, keepdim=True) - - out_zero = self( - negative_prompt, - num_inference_steps=num_inference_steps, - num_images_per_prompt=num_images_per_prompt, - generator=generator, - latents=latents, - negative_prompt=negative_prior_prompt, - guidance_scale=guidance_scale, - ) - zero_image_emb = out_zero.negative_image_embeds if negative_prompt == "" else out_zero.image_embeds - - return KandinskyPriorPipelineOutput(image_embeds=image_emb, negative_image_embeds=zero_image_emb) - - # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - def get_zero_embed(self, batch_size=1, device=None): - device = device or self.device - zero_img = torch.zeros(1, 3, self.image_encoder.config.image_size, self.image_encoder.config.image_size).to( - device=device, dtype=self.image_encoder.dtype - ) - zero_image_emb = self.image_encoder(zero_img)["image_embeds"] - zero_image_emb = zero_image_emb.repeat(batch_size, 1) - return zero_image_emb - - def _encode_prompt( - self, - prompt, - device, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt=None, - ): - batch_size = len(prompt) if isinstance(prompt, list) else 1 - # get prompt text embeddings - 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 - text_mask = text_inputs.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[:, 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}" - ) - text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length] - - text_encoder_output = self.text_encoder(text_input_ids.to(device)) - - prompt_embeds = text_encoder_output.text_embeds - text_encoder_hidden_states = text_encoder_output.last_hidden_state - - prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - text_encoder_hidden_states = text_encoder_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - uncond_tokens: List[str] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif 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 - - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - uncond_text_mask = uncond_input.attention_mask.bool().to(device) - negative_prompt_embeds_text_encoder_output = self.text_encoder(uncond_input.input_ids.to(device)) - - negative_prompt_embeds = negative_prompt_embeds_text_encoder_output.text_embeds - uncond_text_encoder_hidden_states = negative_prompt_embeds_text_encoder_output.last_hidden_state - - # 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.repeat(1, num_images_per_prompt) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len) - - seq_len = uncond_text_encoder_hidden_states.shape[1] - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.repeat(1, num_images_per_prompt, 1) - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.view( - batch_size * num_images_per_prompt, seq_len, -1 - ) - uncond_text_mask = uncond_text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - # done duplicates - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - text_encoder_hidden_states = torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]) - - text_mask = torch.cat([uncond_text_mask, text_mask]) - - return prompt_embeds, text_encoder_hidden_states, text_mask - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = 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, - guidance_scale: float = 4.0, - output_type: Optional[str] = "pt", - return_dict: bool = True, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `List[str]`): - The prompt or prompts to guide the image generation. - 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): - 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*): - 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 ge generated by sampling using the supplied random `generator`. - 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. - output_type (`str`, *optional*, defaults to `"pt"`): - The output format of the generate image. Choose between: `"np"` (`np.array`) or `"pt"` - (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`KandinskyPriorPipelineOutput`] or `tuple` - """ - - if isinstance(prompt, str): - prompt = [prompt] - elif not isinstance(prompt, list): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - if isinstance(negative_prompt, str): - negative_prompt = [negative_prompt] - elif not isinstance(negative_prompt, list) and negative_prompt is not None: - raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}") - - # if the negative prompt is defined we double the batch size to - # directly retrieve the negative prompt embedding - if negative_prompt is not None: - prompt = prompt + negative_prompt - negative_prompt = 2 * negative_prompt - - device = self._execution_device - - batch_size = len(prompt) - batch_size = batch_size * num_images_per_prompt - - do_classifier_free_guidance = guidance_scale > 1.0 - prompt_embeds, text_encoder_hidden_states, text_mask = self._encode_prompt( - prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt - ) - - # prior - self.scheduler.set_timesteps(num_inference_steps, device=device) - prior_timesteps_tensor = self.scheduler.timesteps - - embedding_dim = self.prior.config.embedding_dim - - latents = self.prepare_latents( - (batch_size, embedding_dim), - prompt_embeds.dtype, - device, - generator, - latents, - self.scheduler, - ) - - for i, t in enumerate(self.progress_bar(prior_timesteps_tensor)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - predicted_image_embedding = self.prior( - latent_model_input, - timestep=t, - proj_embedding=prompt_embeds, - encoder_hidden_states=text_encoder_hidden_states, - attention_mask=text_mask, - ).predicted_image_embedding - - if do_classifier_free_guidance: - predicted_image_embedding_uncond, predicted_image_embedding_text = predicted_image_embedding.chunk(2) - predicted_image_embedding = predicted_image_embedding_uncond + guidance_scale * ( - predicted_image_embedding_text - predicted_image_embedding_uncond - ) - - if i + 1 == prior_timesteps_tensor.shape[0]: - prev_timestep = None - else: - prev_timestep = prior_timesteps_tensor[i + 1] - - latents = self.scheduler.step( - predicted_image_embedding, - timestep=t, - sample=latents, - generator=generator, - prev_timestep=prev_timestep, - ).prev_sample - - if XLA_AVAILABLE: - xm.mark_step() - - latents = self.prior.post_process_latents(latents) - - image_embeddings = latents - - # if negative prompt has been defined, we retrieve split the image embedding into two - if negative_prompt is None: - zero_embeds = self.get_zero_embed(latents.shape[0], device=latents.device) - - self.maybe_free_model_hooks() - else: - image_embeddings, zero_embeds = image_embeddings.chunk(2) - - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.prior_hook.offload() - - if output_type not in ["pt", "np"]: - raise ValueError(f"Only the output types `pt` and `np` are supported not output_type={output_type}") - - if output_type == "np": - image_embeddings = image_embeddings.cpu().numpy() - zero_embeds = zero_embeds.cpu().numpy() - - if not return_dict: - return (image_embeddings, zero_embeds) - - return KandinskyPriorPipelineOutput(image_embeds=image_embeddings, negative_image_embeds=zero_embeds) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/text_encoder.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/text_encoder.py deleted file mode 100644 index caa0029f00ca22818819d5b76b57ec489c6da1d6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/text_encoder.py +++ /dev/null @@ -1,27 +0,0 @@ -import torch -from transformers import PreTrainedModel, XLMRobertaConfig, XLMRobertaModel - - -class MCLIPConfig(XLMRobertaConfig): - model_type = "M-CLIP" - - def __init__(self, transformerDimSize=1024, imageDimSize=768, **kwargs): - self.transformerDimensions = transformerDimSize - self.numDims = imageDimSize - super().__init__(**kwargs) - - -class MultilingualCLIP(PreTrainedModel): - config_class = MCLIPConfig - - def __init__(self, config, *args, **kwargs): - super().__init__(config, *args, **kwargs) - self.transformer = XLMRobertaModel(config) - self.LinearTransformation = torch.nn.Linear( - in_features=config.transformerDimensions, out_features=config.numDims - ) - - def forward(self, input_ids, attention_mask): - embs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)[0] - embs2 = (embs * attention_mask.unsqueeze(2)).sum(dim=1) / attention_mask.sum(dim=1)[:, None] - return self.LinearTransformation(embs2), embs diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__init__.py deleted file mode 100644 index 67e97f161173ac8981dadf757fd8d6438307c973..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__init__.py +++ /dev/null @@ -1,70 +0,0 @@ -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_kandinsky2_2"] = ["KandinskyV22Pipeline"] - _import_structure["pipeline_kandinsky2_2_combined"] = [ - "KandinskyV22CombinedPipeline", - "KandinskyV22Img2ImgCombinedPipeline", - "KandinskyV22InpaintCombinedPipeline", - ] - _import_structure["pipeline_kandinsky2_2_controlnet"] = ["KandinskyV22ControlnetPipeline"] - _import_structure["pipeline_kandinsky2_2_controlnet_img2img"] = ["KandinskyV22ControlnetImg2ImgPipeline"] - _import_structure["pipeline_kandinsky2_2_img2img"] = ["KandinskyV22Img2ImgPipeline"] - _import_structure["pipeline_kandinsky2_2_inpainting"] = ["KandinskyV22InpaintPipeline"] - _import_structure["pipeline_kandinsky2_2_prior"] = ["KandinskyV22PriorPipeline"] - _import_structure["pipeline_kandinsky2_2_prior_emb2emb"] = ["KandinskyV22PriorEmb2EmbPipeline"] - - -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_kandinsky2_2 import KandinskyV22Pipeline - from .pipeline_kandinsky2_2_combined import ( - KandinskyV22CombinedPipeline, - KandinskyV22Img2ImgCombinedPipeline, - KandinskyV22InpaintCombinedPipeline, - ) - from .pipeline_kandinsky2_2_controlnet import KandinskyV22ControlnetPipeline - from .pipeline_kandinsky2_2_controlnet_img2img import KandinskyV22ControlnetImg2ImgPipeline - from .pipeline_kandinsky2_2_img2img import KandinskyV22Img2ImgPipeline - from .pipeline_kandinsky2_2_inpainting import KandinskyV22InpaintPipeline - from .pipeline_kandinsky2_2_prior import KandinskyV22PriorPipeline - from .pipeline_kandinsky2_2_prior_emb2emb import KandinskyV22PriorEmb2EmbPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index dd8a651ff98eb2a88ea9c6086a0fe0120612f271..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2.cpython-310.pyc deleted file mode 100644 index 858bd65133323049732322dabe7ada47214914dc..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_combined.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_combined.cpython-310.pyc deleted file mode 100644 index 0b4d4144e0b81cf5d54d8bcc4ede96df108d5602..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_combined.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_controlnet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_controlnet.cpython-310.pyc deleted file mode 100644 index 35a2a3957cd2090d1420a421a8f97b27e26e74a5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_controlnet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_controlnet_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_controlnet_img2img.cpython-310.pyc deleted file mode 100644 index e03753a982e04dc69ad16e5e7ab3c7c0eb038665..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_controlnet_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_img2img.cpython-310.pyc deleted file mode 100644 index c58d03b53e64a8030c18bc8bc170bd5a72864b10..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_inpainting.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_inpainting.cpython-310.pyc deleted file mode 100644 index 9c227543c32f80c3ae2aed56797e2736929f2b51..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_inpainting.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_prior.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_prior.cpython-310.pyc deleted file mode 100644 index 9044388f1ac2bef6e735f4a922c29616858652bc..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_prior.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_prior_emb2emb.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_prior_emb2emb.cpython-310.pyc deleted file mode 100644 index e89486acca1feb16abb044c1cc71d33e33df1199..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/__pycache__/pipeline_kandinsky2_2_prior_emb2emb.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py deleted file mode 100644 index 3ecc0ebd5b2558fb91176e3e3ddc226b4b00215f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py +++ /dev/null @@ -1,331 +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 Callable, Dict, List, Optional, Union - -import torch - -from ...models import UNet2DConditionModel, VQModel -from ...schedulers import DDPMScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -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 - >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline - >>> import torch - - >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior") - >>> pipe_prior.to("cuda") - >>> prompt = "red cat, 4k photo" - >>> out = pipe_prior(prompt) - >>> image_emb = out.image_embeds - >>> zero_image_emb = out.negative_image_embeds - >>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder") - >>> pipe.to("cuda") - >>> image = pipe( - ... image_embeds=image_emb, - ... negative_image_embeds=zero_image_emb, - ... height=768, - ... width=768, - ... num_inference_steps=50, - ... ).images - >>> image[0].save("cat.png") - ``` -""" - - -def downscale_height_and_width(height, width, scale_factor=8): - new_height = height // scale_factor**2 - if height % scale_factor**2 != 0: - new_height += 1 - new_width = width // scale_factor**2 - if width % scale_factor**2 != 0: - new_width += 1 - return new_height * scale_factor, new_width * scale_factor - - -class KandinskyV22Pipeline(DiffusionPipeline): - """ - Pipeline for text-to-image generation using Kandinsky - - 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.) - - Args: - scheduler (Union[`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. - movq ([`VQModel`]): - MoVQ Decoder to generate the image from the latents. - """ - - model_cpu_offload_seq = "unet->movq" - _callback_tensor_inputs = ["latents", "image_embeds", "negative_image_embeds"] - - def __init__( - self, - unet: UNet2DConditionModel, - scheduler: DDPMScheduler, - movq: VQModel, - ): - super().__init__() - - self.register_modules( - unet=unet, - scheduler=scheduler, - movq=movq, - ) - self.movq_scale_factor = 2 ** (len(self.movq.config.block_out_channels) - 1) - - # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - @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 - - @torch.no_grad() - @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]], - 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", - return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], 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]`): - 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]`): - 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. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - 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. - 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. - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - 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. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - 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]}" - ) - - device = self._execution_device - - self._guidance_scale = guidance_scale - - if isinstance(image_embeds, list): - image_embeds = torch.cat(image_embeds, dim=0) - batch_size = image_embeds.shape[0] * num_images_per_prompt - if isinstance(negative_image_embeds, list): - negative_image_embeds = torch.cat(negative_image_embeds, dim=0) - - if self.do_classifier_free_guidance: - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - negative_image_embeds = negative_image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0).to( - dtype=self.unet.dtype, device=device - ) - - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - num_channels_latents = self.unet.config.in_channels - - height, width = downscale_height_and_width(height, width, self.movq_scale_factor) - - # create initial latent - latents = self.prepare_latents( - (batch_size, num_channels_latents, height, width), - image_embeds.dtype, - device, - generator, - latents, - self.scheduler, - ) - - self._num_timesteps = len(timesteps) - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - added_cond_kwargs = {"image_embeds": image_embeds} - noise_pred = self.unet( - sample=latent_model_input, - timestep=t, - encoder_hidden_states=None, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - if self.do_classifier_free_guidance: - noise_pred, variance_pred = noise_pred.split(latents.shape[1], dim=1) - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - _, variance_pred_text = variance_pred.chunk(2) - noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, variance_pred_text], dim=1) - - if not ( - hasattr(self.scheduler.config, "variance_type") - and self.scheduler.config.variance_type in ["learned", "learned_range"] - ): - noise_pred, _ = noise_pred.split(latents.shape[1], dim=1) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step( - noise_pred, - t, - latents, - generator=generator, - )[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) - image_embeds = callback_outputs.pop("image_embeds", image_embeds) - negative_image_embeds = callback_outputs.pop("negative_image_embeds", negative_image_embeds) - - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if output_type not in ["pt", "np", "pil", "latent"]: - raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}") - - if not output_type == "latent": - # post-processing - image = self.movq.decode(latents, force_not_quantize=True)["sample"] - if output_type in ["np", "pil"]: - image = image * 0.5 + 0.5 - image = image.clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - if output_type == "pil": - image = self.numpy_to_pil(image) - else: - image = latents - - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py deleted file mode 100644 index e0b88b41e8c539371ef07ed140526fa8d18a45aa..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py +++ /dev/null @@ -1,854 +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 Callable, Dict, List, Optional, Union - -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...models import PriorTransformer, UNet2DConditionModel, VQModel -from ...schedulers import DDPMScheduler, UnCLIPScheduler -from ...utils import deprecate, logging, replace_example_docstring -from ..pipeline_utils import DiffusionPipeline -from .pipeline_kandinsky2_2 import KandinskyV22Pipeline -from .pipeline_kandinsky2_2_img2img import KandinskyV22Img2ImgPipeline -from .pipeline_kandinsky2_2_inpainting import KandinskyV22InpaintPipeline -from .pipeline_kandinsky2_2_prior import KandinskyV22PriorPipeline - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -TEXT2IMAGE_EXAMPLE_DOC_STRING = """ - Examples: - ```py - from diffusers import AutoPipelineForText2Image - import torch - - pipe = AutoPipelineForText2Image.from_pretrained( - "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16 - ) - pipe.enable_model_cpu_offload() - - prompt = "A lion in galaxies, spirals, nebulae, stars, smoke, iridescent, intricate detail, octane render, 8k" - - image = pipe(prompt=prompt, num_inference_steps=25).images[0] - ``` -""" - -IMAGE2IMAGE_EXAMPLE_DOC_STRING = """ - Examples: - ```py - from diffusers import AutoPipelineForImage2Image - import torch - import requests - from io import BytesIO - from PIL import Image - import os - - pipe = AutoPipelineForImage2Image.from_pretrained( - "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16 - ) - pipe.enable_model_cpu_offload() - - prompt = "A fantasy landscape, Cinematic lighting" - negative_prompt = "low quality, bad quality" - - url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" - - response = requests.get(url) - image = Image.open(BytesIO(response.content)).convert("RGB") - image.thumbnail((768, 768)) - - image = pipe(prompt=prompt, image=original_image, num_inference_steps=25).images[0] - ``` -""" - -INPAINT_EXAMPLE_DOC_STRING = """ - Examples: - ```py - from diffusers import AutoPipelineForInpainting - from diffusers.utils import load_image - import torch - import numpy as np - - pipe = AutoPipelineForInpainting.from_pretrained( - "kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch.float16 - ) - pipe.enable_model_cpu_offload() - - prompt = "A fantasy landscape, Cinematic lighting" - negative_prompt = "low quality, bad quality" - - original_image = load_image( - "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/kandinsky/cat.png" - ) - - mask = np.zeros((768, 768), dtype=np.float32) - # Let's mask out an area above the cat's head - mask[:250, 250:-250] = 1 - - image = pipe(prompt=prompt, image=original_image, mask_image=mask, num_inference_steps=25).images[0] - ``` -""" - - -class KandinskyV22CombinedPipeline(DiffusionPipeline): - """ - Combined Pipeline for text-to-image generation using Kandinsky - - 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.) - - Args: - scheduler (Union[`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. - movq ([`VQModel`]): - MoVQ Decoder to generate the image from the latents. - prior_prior ([`PriorTransformer`]): - The canonical unCLIP prior to approximate the image embedding from the text embedding. - prior_image_encoder ([`CLIPVisionModelWithProjection`]): - Frozen image-encoder. - prior_text_encoder ([`CLIPTextModelWithProjection`]): - Frozen text-encoder. - prior_tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - prior_scheduler ([`UnCLIPScheduler`]): - A scheduler to be used in combination with `prior` to generate image embedding. - prior_image_processor ([`CLIPImageProcessor`]): - A image_processor to be used to preprocess image from clip. - """ - - model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->unet->movq" - _load_connected_pipes = True - _exclude_from_cpu_offload = ["prior_prior"] - - def __init__( - self, - unet: UNet2DConditionModel, - scheduler: DDPMScheduler, - movq: VQModel, - prior_prior: PriorTransformer, - prior_image_encoder: CLIPVisionModelWithProjection, - prior_text_encoder: CLIPTextModelWithProjection, - prior_tokenizer: CLIPTokenizer, - prior_scheduler: UnCLIPScheduler, - prior_image_processor: CLIPImageProcessor, - ): - super().__init__() - - self.register_modules( - unet=unet, - scheduler=scheduler, - movq=movq, - prior_prior=prior_prior, - prior_image_encoder=prior_image_encoder, - prior_text_encoder=prior_text_encoder, - prior_tokenizer=prior_tokenizer, - prior_scheduler=prior_scheduler, - prior_image_processor=prior_image_processor, - ) - self.prior_pipe = KandinskyV22PriorPipeline( - prior=prior_prior, - image_encoder=prior_image_encoder, - text_encoder=prior_text_encoder, - tokenizer=prior_tokenizer, - scheduler=prior_scheduler, - image_processor=prior_image_processor, - ) - self.decoder_pipe = KandinskyV22Pipeline( - unet=unet, - scheduler=scheduler, - movq=movq, - ) - - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = 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): - 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 - `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called. - Note that offloading happens on a submodule basis. Memory savings are higher than with - `enable_model_cpu_offload`, but performance is lower. - """ - self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - - def progress_bar(self, iterable=None, total=None): - self.prior_pipe.progress_bar(iterable=iterable, total=total) - self.decoder_pipe.progress_bar(iterable=iterable, total=total) - self.decoder_pipe.enable_model_cpu_offload() - - def set_progress_bar_config(self, **kwargs): - self.prior_pipe.set_progress_bar_config(**kwargs) - self.decoder_pipe.set_progress_bar_config(**kwargs) - - @torch.no_grad() - @replace_example_docstring(TEXT2IMAGE_EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, - num_inference_steps: int = 100, - guidance_scale: float = 4.0, - num_images_per_prompt: int = 1, - height: int = 512, - 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, - 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"], - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `List[str]`): - The prompt or prompts to guide the image generation. - 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): - 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. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - prior_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. - 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. - 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. - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - prior_callback_on_step_end (`Callable`, *optional*): - 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*): - 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. - callback_on_step_end (`Callable`, *optional*): - A function that calls at the end of each denoising steps during the inference of the decoder pipeline. - 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. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - prior_outputs = self.prior_pipe( - prompt=prompt, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - num_inference_steps=prior_num_inference_steps, - generator=generator, - latents=latents, - guidance_scale=prior_guidance_scale, - output_type="pt", - return_dict=False, - callback_on_step_end=prior_callback_on_step_end, - callback_on_step_end_tensor_inputs=prior_callback_on_step_end_tensor_inputs, - ) - image_embeds = prior_outputs[0] - negative_image_embeds = prior_outputs[1] - - prompt = [prompt] if not isinstance(prompt, (list, tuple)) else prompt - - if len(prompt) < image_embeds.shape[0] and image_embeds.shape[0] % len(prompt) == 0: - prompt = (image_embeds.shape[0] // len(prompt)) * prompt - - outputs = self.decoder_pipe( - image_embeds=image_embeds, - negative_image_embeds=negative_image_embeds, - width=width, - height=height, - num_inference_steps=num_inference_steps, - generator=generator, - guidance_scale=guidance_scale, - output_type=output_type, - callback=callback, - callback_steps=callback_steps, - return_dict=return_dict, - callback_on_step_end=callback_on_step_end, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - ) - self.maybe_free_model_hooks() - - return outputs - - -class KandinskyV22Img2ImgCombinedPipeline(DiffusionPipeline): - """ - Combined Pipeline for image-to-image generation using Kandinsky - - 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.) - - Args: - scheduler (Union[`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. - movq ([`VQModel`]): - MoVQ Decoder to generate the image from the latents. - prior_prior ([`PriorTransformer`]): - The canonical unCLIP prior to approximate the image embedding from the text embedding. - prior_image_encoder ([`CLIPVisionModelWithProjection`]): - Frozen image-encoder. - prior_text_encoder ([`CLIPTextModelWithProjection`]): - Frozen text-encoder. - prior_tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - prior_scheduler ([`UnCLIPScheduler`]): - A scheduler to be used in combination with `prior` to generate image embedding. - prior_image_processor ([`CLIPImageProcessor`]): - A image_processor to be used to preprocess image from clip. - """ - - model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->unet->movq" - _load_connected_pipes = True - _exclude_from_cpu_offload = ["prior_prior"] - - def __init__( - self, - unet: UNet2DConditionModel, - scheduler: DDPMScheduler, - movq: VQModel, - prior_prior: PriorTransformer, - prior_image_encoder: CLIPVisionModelWithProjection, - prior_text_encoder: CLIPTextModelWithProjection, - prior_tokenizer: CLIPTokenizer, - prior_scheduler: UnCLIPScheduler, - prior_image_processor: CLIPImageProcessor, - ): - super().__init__() - - self.register_modules( - unet=unet, - scheduler=scheduler, - movq=movq, - prior_prior=prior_prior, - prior_image_encoder=prior_image_encoder, - prior_text_encoder=prior_text_encoder, - prior_tokenizer=prior_tokenizer, - prior_scheduler=prior_scheduler, - prior_image_processor=prior_image_processor, - ) - self.prior_pipe = KandinskyV22PriorPipeline( - prior=prior_prior, - image_encoder=prior_image_encoder, - text_encoder=prior_text_encoder, - tokenizer=prior_tokenizer, - scheduler=prior_scheduler, - image_processor=prior_image_processor, - ) - self.decoder_pipe = KandinskyV22Img2ImgPipeline( - unet=unet, - scheduler=scheduler, - movq=movq, - ) - - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = 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): - 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` - method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with - `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`. - """ - 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): - 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 - `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called. - Note that offloading happens on a submodule basis. Memory savings are higher than with - `enable_model_cpu_offload`, but performance is lower. - """ - self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - - def progress_bar(self, iterable=None, total=None): - self.prior_pipe.progress_bar(iterable=iterable, total=total) - self.decoder_pipe.progress_bar(iterable=iterable, total=total) - self.decoder_pipe.enable_model_cpu_offload() - - def set_progress_bar_config(self, **kwargs): - self.prior_pipe.set_progress_bar_config(**kwargs) - self.decoder_pipe.set_progress_bar_config(**kwargs) - - @torch.no_grad() - @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, - num_inference_steps: int = 100, - guidance_scale: float = 4.0, - strength: float = 0.3, - num_images_per_prompt: int = 1, - height: int = 512, - 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, - 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"], - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - 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`, 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*): - 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): - The number of images to generate per prompt. - 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. - strength (`float`, *optional*, defaults to 0.3): - 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_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. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - prior_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. - 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*): - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - prior_outputs = self.prior_pipe( - prompt=prompt, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - num_inference_steps=prior_num_inference_steps, - generator=generator, - latents=latents, - guidance_scale=prior_guidance_scale, - output_type="pt", - return_dict=False, - callback_on_step_end=prior_callback_on_step_end, - callback_on_step_end_tensor_inputs=prior_callback_on_step_end_tensor_inputs, - ) - image_embeds = prior_outputs[0] - negative_image_embeds = prior_outputs[1] - - prompt = [prompt] if not isinstance(prompt, (list, tuple)) else prompt - image = [image] if isinstance(image, PIL.Image.Image) else image - - if len(prompt) < image_embeds.shape[0] and image_embeds.shape[0] % len(prompt) == 0: - prompt = (image_embeds.shape[0] // len(prompt)) * prompt - - if ( - isinstance(image, (list, tuple)) - and len(image) < image_embeds.shape[0] - and image_embeds.shape[0] % len(image) == 0 - ): - image = (image_embeds.shape[0] // len(image)) * image - - outputs = self.decoder_pipe( - image=image, - image_embeds=image_embeds, - negative_image_embeds=negative_image_embeds, - width=width, - height=height, - strength=strength, - num_inference_steps=num_inference_steps, - generator=generator, - guidance_scale=guidance_scale, - output_type=output_type, - callback=callback, - callback_steps=callback_steps, - return_dict=return_dict, - callback_on_step_end=callback_on_step_end, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - ) - - self.maybe_free_model_hooks() - return outputs - - -class KandinskyV22InpaintCombinedPipeline(DiffusionPipeline): - """ - Combined Pipeline for inpainting generation using Kandinsky - - 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.) - - Args: - scheduler (Union[`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. - movq ([`VQModel`]): - MoVQ Decoder to generate the image from the latents. - prior_prior ([`PriorTransformer`]): - The canonical unCLIP prior to approximate the image embedding from the text embedding. - prior_image_encoder ([`CLIPVisionModelWithProjection`]): - Frozen image-encoder. - prior_text_encoder ([`CLIPTextModelWithProjection`]): - Frozen text-encoder. - prior_tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - prior_scheduler ([`UnCLIPScheduler`]): - A scheduler to be used in combination with `prior` to generate image embedding. - prior_image_processor ([`CLIPImageProcessor`]): - A image_processor to be used to preprocess image from clip. - """ - - model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->unet->movq" - _load_connected_pipes = True - _exclude_from_cpu_offload = ["prior_prior"] - - def __init__( - self, - unet: UNet2DConditionModel, - scheduler: DDPMScheduler, - movq: VQModel, - prior_prior: PriorTransformer, - prior_image_encoder: CLIPVisionModelWithProjection, - prior_text_encoder: CLIPTextModelWithProjection, - prior_tokenizer: CLIPTokenizer, - prior_scheduler: UnCLIPScheduler, - prior_image_processor: CLIPImageProcessor, - ): - super().__init__() - - self.register_modules( - unet=unet, - scheduler=scheduler, - movq=movq, - prior_prior=prior_prior, - prior_image_encoder=prior_image_encoder, - prior_text_encoder=prior_text_encoder, - prior_tokenizer=prior_tokenizer, - prior_scheduler=prior_scheduler, - prior_image_processor=prior_image_processor, - ) - self.prior_pipe = KandinskyV22PriorPipeline( - prior=prior_prior, - image_encoder=prior_image_encoder, - text_encoder=prior_text_encoder, - tokenizer=prior_tokenizer, - scheduler=prior_scheduler, - image_processor=prior_image_processor, - ) - self.decoder_pipe = KandinskyV22InpaintPipeline( - unet=unet, - scheduler=scheduler, - movq=movq, - ) - - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = 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): - 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 - `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called. - Note that offloading happens on a submodule basis. Memory savings are higher than with - `enable_model_cpu_offload`, but performance is lower. - """ - self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - - def progress_bar(self, iterable=None, total=None): - self.prior_pipe.progress_bar(iterable=iterable, total=total) - self.decoder_pipe.progress_bar(iterable=iterable, total=total) - self.decoder_pipe.enable_model_cpu_offload() - - def set_progress_bar_config(self, **kwargs): - self.prior_pipe.set_progress_bar_config(**kwargs) - self.decoder_pipe.set_progress_bar_config(**kwargs) - - @torch.no_grad() - @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, - num_inference_steps: int = 100, - guidance_scale: float = 4.0, - num_images_per_prompt: int = 1, - height: int = 512, - 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", - 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"], - **kwargs, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - 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`, 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. - mask_image (`np.array`): - Tensor representing an image batch, to mask `image`. White pixels in the mask will be 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)`. - 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): - The number of images to generate per prompt. - 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. - 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. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - prior_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. - 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*): - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - prior_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: `prior_callback_on_step_end(self: DiffusionPipeline, step: int, timestep: - int, callback_kwargs: Dict)`. - 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. - 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. - - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - prior_kwargs = {} - if kwargs.get("prior_callback", None) is not None: - prior_kwargs["callback"] = kwargs.pop("prior_callback") - deprecate( - "prior_callback", - "1.0.0", - "Passing `prior_callback` as an input argument to `__call__` is deprecated, consider use `prior_callback_on_step_end`", - ) - if kwargs.get("prior_callback_steps", None) is not None: - deprecate( - "prior_callback_steps", - "1.0.0", - "Passing `prior_callback_steps` as an input argument to `__call__` is deprecated, consider use `prior_callback_on_step_end`", - ) - prior_kwargs["callback_steps"] = kwargs.pop("prior_callback_steps") - - prior_outputs = self.prior_pipe( - prompt=prompt, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - num_inference_steps=prior_num_inference_steps, - generator=generator, - latents=latents, - guidance_scale=prior_guidance_scale, - output_type="pt", - return_dict=False, - callback_on_step_end=prior_callback_on_step_end, - callback_on_step_end_tensor_inputs=prior_callback_on_step_end_tensor_inputs, - **prior_kwargs, - ) - image_embeds = prior_outputs[0] - negative_image_embeds = prior_outputs[1] - - prompt = [prompt] if not isinstance(prompt, (list, tuple)) else prompt - image = [image] if isinstance(image, PIL.Image.Image) else image - mask_image = [mask_image] if isinstance(mask_image, PIL.Image.Image) else mask_image - - if len(prompt) < image_embeds.shape[0] and image_embeds.shape[0] % len(prompt) == 0: - prompt = (image_embeds.shape[0] // len(prompt)) * prompt - - if ( - isinstance(image, (list, tuple)) - and len(image) < image_embeds.shape[0] - and image_embeds.shape[0] % len(image) == 0 - ): - image = (image_embeds.shape[0] // len(image)) * image - - if ( - isinstance(mask_image, (list, tuple)) - and len(mask_image) < image_embeds.shape[0] - and image_embeds.shape[0] % len(mask_image) == 0 - ): - mask_image = (image_embeds.shape[0] // len(mask_image)) * mask_image - - outputs = self.decoder_pipe( - image=image, - mask_image=mask_image, - image_embeds=image_embeds, - negative_image_embeds=negative_image_embeds, - width=width, - height=height, - num_inference_steps=num_inference_steps, - generator=generator, - guidance_scale=guidance_scale, - output_type=output_type, - return_dict=return_dict, - callback_on_step_end=callback_on_step_end, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - **kwargs, - ) - self.maybe_free_model_hooks() - - return outputs diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py deleted file mode 100644 index b9f98f5458e2be332eeda7fba52caad7a819a5f4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py +++ /dev/null @@ -1,333 +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 Callable, List, Optional, Union - -import torch - -from ...models import UNet2DConditionModel, VQModel -from ...schedulers import DDPMScheduler -from ...utils import ( - is_torch_xla_available, - logging, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -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 - >>> import numpy as np - - >>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22ControlnetPipeline - >>> from transformers import pipeline - >>> from diffusers.utils import load_image - - - >>> def make_hint(image, depth_estimator): - ... image = depth_estimator(image)["depth"] - ... image = np.array(image) - ... image = image[:, :, None] - ... image = np.concatenate([image, image, image], axis=2) - ... detected_map = torch.from_numpy(image).float() / 255.0 - ... hint = detected_map.permute(2, 0, 1) - ... return hint - - - >>> depth_estimator = pipeline("depth-estimation") - - >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16 - ... ) - >>> pipe_prior = pipe_prior.to("cuda") - - >>> pipe = KandinskyV22ControlnetPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - - >>> img = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/cat.png" - ... ).resize((768, 768)) - - >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to("cuda") - - >>> prompt = "A robot, 4k photo" - >>> negative_prior_prompt = "lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, 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, long neck, username, watermark, signature" - - >>> generator = torch.Generator(device="cuda").manual_seed(43) - - >>> image_emb, zero_image_emb = pipe_prior( - ... prompt=prompt, negative_prompt=negative_prior_prompt, generator=generator - ... ).to_tuple() - - >>> images = pipe( - ... image_embeds=image_emb, - ... negative_image_embeds=zero_image_emb, - ... hint=hint, - ... num_inference_steps=50, - ... generator=generator, - ... height=768, - ... width=768, - ... ).images - - >>> images[0].save("robot_cat.png") - ``` -""" - - -# Copied from diffusers.pipelines.kandinsky2_2.pipeline_kandinsky2_2.downscale_height_and_width -def downscale_height_and_width(height, width, scale_factor=8): - new_height = height // scale_factor**2 - if height % scale_factor**2 != 0: - new_height += 1 - new_width = width // scale_factor**2 - if width % scale_factor**2 != 0: - new_width += 1 - return new_height * scale_factor, new_width * scale_factor - - -class KandinskyV22ControlnetPipeline(DiffusionPipeline): - """ - Pipeline for text-to-image generation using Kandinsky - - 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.) - - Args: - scheduler ([`DDIMScheduler`]): - A scheduler to be used in combination with `unet` to generate image latents. - unet ([`UNet2DConditionModel`]): - Conditional U-Net architecture to denoise the image embedding. - movq ([`VQModel`]): - MoVQ Decoder to generate the image from the latents. - """ - - model_cpu_offload_seq = "unet->movq" - - def __init__( - self, - unet: UNet2DConditionModel, - scheduler: DDPMScheduler, - movq: VQModel, - ): - super().__init__() - - self.register_modules( - unet=unet, - scheduler=scheduler, - movq=movq, - ) - self.movq_scale_factor = 2 ** (len(self.movq.config.block_out_channels) - 1) - - # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - @torch.no_grad() - def __call__( - self, - image_embeds: Union[torch.Tensor, List[torch.Tensor]], - negative_image_embeds: Union[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, - callback_steps: int = 1, - return_dict: bool = True, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - 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]`): - 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]`): - The clip image embeddings for negative text prompt, will be used to condition the image generation. - 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): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - 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. - 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. - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - device = self._execution_device - - do_classifier_free_guidance = guidance_scale > 1.0 - - if isinstance(image_embeds, list): - image_embeds = torch.cat(image_embeds, dim=0) - if isinstance(negative_image_embeds, list): - negative_image_embeds = torch.cat(negative_image_embeds, dim=0) - if isinstance(hint, list): - hint = torch.cat(hint, dim=0) - - batch_size = image_embeds.shape[0] * num_images_per_prompt - - if do_classifier_free_guidance: - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - negative_image_embeds = negative_image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - hint = hint.repeat_interleave(num_images_per_prompt, dim=0) - - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0).to( - dtype=self.unet.dtype, device=device - ) - hint = torch.cat([hint, hint], dim=0).to(dtype=self.unet.dtype, device=device) - - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps_tensor = self.scheduler.timesteps - - num_channels_latents = self.movq.config.latent_channels - - height, width = downscale_height_and_width(height, width, self.movq_scale_factor) - - # create initial latent - latents = self.prepare_latents( - (batch_size, num_channels_latents, height, width), - image_embeds.dtype, - device, - generator, - latents, - self.scheduler, - ) - - for i, t in enumerate(self.progress_bar(timesteps_tensor)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - added_cond_kwargs = {"image_embeds": image_embeds, "hint": hint} - noise_pred = self.unet( - sample=latent_model_input, - timestep=t, - encoder_hidden_states=None, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - if do_classifier_free_guidance: - noise_pred, variance_pred = noise_pred.split(latents.shape[1], dim=1) - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - _, variance_pred_text = variance_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, variance_pred_text], dim=1) - - if not ( - hasattr(self.scheduler.config, "variance_type") - and self.scheduler.config.variance_type in ["learned", "learned_range"] - ): - noise_pred, _ = noise_pred.split(latents.shape[1], dim=1) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step( - noise_pred, - t, - latents, - generator=generator, - )[0] - - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # post-processing - image = self.movq.decode(latents, force_not_quantize=True)["sample"] - - # Offload all models - self.maybe_free_model_hooks() - - if output_type not in ["pt", "np", "pil"]: - raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}") - - if output_type in ["np", "pil"]: - image = image * 0.5 + 0.5 - image = image.clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py deleted file mode 100644 index 54154c6ec1f2c4aef3e876e61fc1c915296e7039..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py +++ /dev/null @@ -1,371 +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 Callable, List, Optional, Union - -import PIL.Image -import torch - -from ...image_processor import VaeImageProcessor -from ...models import UNet2DConditionModel, VQModel -from ...schedulers import DDPMScheduler -from ...utils import ( - is_torch_xla_available, - logging, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -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 - >>> import numpy as np - - >>> from diffusers import KandinskyV22PriorEmb2EmbPipeline, KandinskyV22ControlnetImg2ImgPipeline - >>> from transformers import pipeline - >>> from diffusers.utils import load_image - - - >>> def make_hint(image, depth_estimator): - ... image = depth_estimator(image)["depth"] - ... image = np.array(image) - ... image = image[:, :, None] - ... image = np.concatenate([image, image, image], axis=2) - ... detected_map = torch.from_numpy(image).float() / 255.0 - ... hint = detected_map.permute(2, 0, 1) - ... return hint - - - >>> depth_estimator = pipeline("depth-estimation") - - >>> pipe_prior = KandinskyV22PriorEmb2EmbPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16 - ... ) - >>> pipe_prior = pipe_prior.to("cuda") - - >>> pipe = KandinskyV22ControlnetImg2ImgPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> img = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/cat.png" - ... ).resize((768, 768)) - - - >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to("cuda") - - >>> prompt = "A robot, 4k photo" - >>> negative_prior_prompt = "lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, 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, long neck, username, watermark, signature" - - >>> generator = torch.Generator(device="cuda").manual_seed(43) - - >>> img_emb = pipe_prior(prompt=prompt, image=img, strength=0.85, generator=generator) - >>> negative_emb = pipe_prior(prompt=negative_prior_prompt, image=img, strength=1, generator=generator) - - >>> images = pipe( - ... image=img, - ... strength=0.5, - ... image_embeds=img_emb.image_embeds, - ... negative_image_embeds=negative_emb.image_embeds, - ... hint=hint, - ... num_inference_steps=50, - ... generator=generator, - ... height=768, - ... width=768, - ... ).images - - >>> images[0].save("robot_cat.png") - ``` -""" - - -class KandinskyV22ControlnetImg2ImgPipeline(DiffusionPipeline): - """ - Pipeline for image-to-image generation using Kandinsky - - 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.) - - Args: - scheduler ([`DDIMScheduler`]): - A scheduler to be used in combination with `unet` to generate image latents. - unet ([`UNet2DConditionModel`]): - Conditional U-Net architecture to denoise the image embedding. - movq ([`VQModel`]): - MoVQ Decoder to generate the image from the latents. - """ - - model_cpu_offload_seq = "unet->movq" - - def __init__( - self, - unet: UNet2DConditionModel, - scheduler: DDPMScheduler, - movq: VQModel, - ): - super().__init__() - - self.register_modules( - unet=unet, - scheduler=scheduler, - movq=movq, - ) - movq_scale_factor = 2 ** (len(self.movq.config.block_out_channels) - 1) if getattr(self, "movq", None) else 8 - movq_latent_channels = self.movq.config.latent_channels if getattr(self, "movq", None) else 4 - self.image_processor = VaeImageProcessor( - vae_scale_factor=movq_scale_factor, - vae_latent_channels=movq_latent_channels, - resample="bicubic", - reducing_gap=1, - ) - - # Copied from diffusers.pipelines.kandinsky.pipeline_kandinsky_img2img.KandinskyImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - timesteps = self.scheduler.timesteps[t_start:] - - return timesteps, num_inference_steps - t_start - - # Copied from diffusers.pipelines.kandinsky2_2.pipeline_kandinsky2_2_img2img.KandinskyV22Img2ImgPipeline.prepare_latents - def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - 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." - ) - - elif isinstance(generator, list): - init_latents = [ - self.movq.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = self.movq.encode(image).latent_dist.sample(generator) - - init_latents = self.movq.config.scaling_factor * init_latents - - init_latents = torch.cat([init_latents], dim=0) - - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - - latents = init_latents - - return latents - - @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]], - hint: 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", - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - return_dict: bool = True, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - 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`, 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. - strength (`float`, *optional*, defaults to 0.8): - 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`. - hint (`torch.Tensor`): - The controlnet condition. - 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. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - 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. - 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. - 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. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - device = self._execution_device - - do_classifier_free_guidance = guidance_scale > 1.0 - - if isinstance(image_embeds, list): - image_embeds = torch.cat(image_embeds, dim=0) - if isinstance(negative_image_embeds, list): - negative_image_embeds = torch.cat(negative_image_embeds, dim=0) - if isinstance(hint, list): - hint = torch.cat(hint, dim=0) - - batch_size = image_embeds.shape[0] - - if do_classifier_free_guidance: - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - negative_image_embeds = negative_image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - hint = hint.repeat_interleave(num_images_per_prompt, dim=0) - - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0).to( - dtype=self.unet.dtype, device=device - ) - hint = torch.cat([hint, hint], dim=0).to(dtype=self.unet.dtype, device=device) - - if not isinstance(image, list): - image = [image] - if not all(isinstance(i, (PIL.Image.Image, torch.Tensor)) for i in image): - raise ValueError( - f"Input is in incorrect format: {[type(i) for i in image]}. Currently, we only support PIL image and pytorch tensor" - ) - - image = torch.cat([self.image_processor.preprocess(i, width, height) for i in image], dim=0) - image = image.to(dtype=image_embeds.dtype, device=device) - - latents = self.movq.encode(image)["latents"] - latents = latents.repeat_interleave(num_images_per_prompt, dim=0) - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - latents = self.prepare_latents( - latents, latent_timestep, batch_size, num_images_per_prompt, image_embeds.dtype, device, generator - ) - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - added_cond_kwargs = {"image_embeds": image_embeds, "hint": hint} - noise_pred = self.unet( - sample=latent_model_input, - timestep=t, - encoder_hidden_states=None, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - if do_classifier_free_guidance: - noise_pred, variance_pred = noise_pred.split(latents.shape[1], dim=1) - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - _, variance_pred_text = variance_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, variance_pred_text], dim=1) - - if not ( - hasattr(self.scheduler.config, "variance_type") - and self.scheduler.config.variance_type in ["learned", "learned_range"] - ): - noise_pred, _ = noise_pred.split(latents.shape[1], dim=1) - - # compute the previous noisy sample x_t -> x_t-1 - - latents = self.scheduler.step( - noise_pred, - t, - latents, - generator=generator, - )[0] - - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # post-processing - image = self.movq.decode(latents, force_not_quantize=True)["sample"] - - # Offload all models - self.maybe_free_model_hooks() - - if output_type not in ["pt", "np", "pil"]: - raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}") - - image = self.image_processor.postprocess(image, output_type) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py deleted file mode 100644 index 3b2509098fd15131ffe11ef1d3882718e31ab67e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py +++ /dev/null @@ -1,382 +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 Callable, Dict, List, Optional, Union - -import PIL.Image -import torch - -from ...image_processor import VaeImageProcessor -from ...models import UNet2DConditionModel, VQModel -from ...schedulers import DDPMScheduler -from ...utils import deprecate, is_torch_xla_available, logging -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -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 - >>> from diffusers import KandinskyV22Img2ImgPipeline, KandinskyV22PriorPipeline - >>> from diffusers.utils import load_image - >>> import torch - - >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16 - ... ) - >>> pipe_prior.to("cuda") - - >>> prompt = "A red cartoon frog, 4k" - >>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False) - - >>> pipe = KandinskyV22Img2ImgPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16 - ... ) - >>> pipe.to("cuda") - - >>> init_image = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/frog.png" - ... ) - - >>> image = pipe( - ... image=init_image, - ... image_embeds=image_emb, - ... negative_image_embeds=zero_image_emb, - ... height=768, - ... width=768, - ... num_inference_steps=100, - ... strength=0.2, - ... ).images - - >>> image[0].save("red_frog.png") - ``` -""" - - -class KandinskyV22Img2ImgPipeline(DiffusionPipeline): - """ - Pipeline for image-to-image generation using Kandinsky - - 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.) - - Args: - scheduler ([`DDIMScheduler`]): - A scheduler to be used in combination with `unet` to generate image latents. - unet ([`UNet2DConditionModel`]): - Conditional U-Net architecture to denoise the image embedding. - movq ([`VQModel`]): - MoVQ Decoder to generate the image from the latents. - """ - - model_cpu_offload_seq = "unet->movq" - _callback_tensor_inputs = ["latents", "image_embeds", "negative_image_embeds"] - - def __init__( - self, - unet: UNet2DConditionModel, - scheduler: DDPMScheduler, - movq: VQModel, - ): - super().__init__() - - self.register_modules( - unet=unet, - scheduler=scheduler, - movq=movq, - ) - movq_scale_factor = 2 ** (len(self.movq.config.block_out_channels) - 1) if getattr(self, "movq", None) else 8 - movq_latent_channels = self.movq.config.latent_channels if getattr(self, "movq", None) else 4 - self.image_processor = VaeImageProcessor( - vae_scale_factor=movq_scale_factor, - vae_latent_channels=movq_latent_channels, - resample="bicubic", - reducing_gap=1, - ) - - # Copied from diffusers.pipelines.kandinsky.pipeline_kandinsky_img2img.KandinskyImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - timesteps = self.scheduler.timesteps[t_start:] - - return timesteps, num_inference_steps - t_start - - def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - 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." - ) - - elif isinstance(generator, list): - init_latents = [ - self.movq.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = self.movq.encode(image).latent_dist.sample(generator) - - init_latents = self.movq.config.scaling_factor * init_latents - - init_latents = torch.cat([init_latents], dim=0) - - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - - latents = init_latents - - return latents - - @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 - - @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]], - 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", - return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], 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]`): - 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`, 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. - strength (`float`, *optional*, defaults to 0.8): - 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`. - 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. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - 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. - 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. - 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. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - 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. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - 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]}" - ) - - device = self._execution_device - - self._guidance_scale = guidance_scale - - if isinstance(image_embeds, list): - image_embeds = torch.cat(image_embeds, dim=0) - batch_size = image_embeds.shape[0] - if isinstance(negative_image_embeds, list): - negative_image_embeds = torch.cat(negative_image_embeds, dim=0) - - if self.do_classifier_free_guidance: - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - negative_image_embeds = negative_image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0).to( - dtype=self.unet.dtype, device=device - ) - - if not isinstance(image, list): - image = [image] - if not all(isinstance(i, (PIL.Image.Image, torch.Tensor)) for i in image): - raise ValueError( - f"Input is in incorrect format: {[type(i) for i in image]}. Currently, we only support PIL image and pytorch tensor" - ) - - image = torch.cat([self.image_processor.preprocess(i, width, height) for i in image], dim=0) - image = image.to(dtype=image_embeds.dtype, device=device) - - latents = self.movq.encode(image)["latents"] - latents = latents.repeat_interleave(num_images_per_prompt, dim=0) - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - latents = self.prepare_latents( - latents, latent_timestep, batch_size, num_images_per_prompt, image_embeds.dtype, device, generator - ) - self._num_timesteps = len(timesteps) - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - added_cond_kwargs = {"image_embeds": image_embeds} - noise_pred = self.unet( - sample=latent_model_input, - timestep=t, - encoder_hidden_states=None, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - if self.do_classifier_free_guidance: - noise_pred, variance_pred = noise_pred.split(latents.shape[1], dim=1) - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - _, variance_pred_text = variance_pred.chunk(2) - noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, variance_pred_text], dim=1) - - if not ( - hasattr(self.scheduler.config, "variance_type") - and self.scheduler.config.variance_type in ["learned", "learned_range"] - ): - noise_pred, _ = noise_pred.split(latents.shape[1], dim=1) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step( - noise_pred, - t, - latents, - generator=generator, - )[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) - image_embeds = callback_outputs.pop("image_embeds", image_embeds) - negative_image_embeds = callback_outputs.pop("negative_image_embeds", negative_image_embeds) - - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if not output_type == "latent": - image = self.movq.decode(latents, force_not_quantize=True)["sample"] - image = self.image_processor.postprocess(image, output_type) - else: - image = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py deleted file mode 100644 index 22171849bbf6a119e1c0d511318be8209d96af9e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py +++ /dev/null @@ -1,567 +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 copy import deepcopy -from typing import Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from packaging import version -from PIL import Image - -from ... import __version__ -from ...models import UNet2DConditionModel, VQModel -from ...schedulers import DDPMScheduler -from ...utils import deprecate, is_torch_xla_available, logging -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -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 - >>> from diffusers import KandinskyV22InpaintPipeline, KandinskyV22PriorPipeline - >>> from diffusers.utils import load_image - >>> import torch - >>> import numpy as np - - >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16 - ... ) - >>> pipe_prior.to("cuda") - - >>> prompt = "a hat" - >>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False) - - >>> pipe = KandinskyV22InpaintPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch.float16 - ... ) - >>> pipe.to("cuda") - - >>> init_image = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/cat.png" - ... ) - - >>> mask = np.zeros((768, 768), dtype=np.float32) - >>> mask[:250, 250:-250] = 1 - - >>> out = pipe( - ... image=init_image, - ... mask_image=mask, - ... image_embeds=image_emb, - ... negative_image_embeds=zero_image_emb, - ... height=768, - ... width=768, - ... num_inference_steps=50, - ... ) - - >>> image = out.images[0] - >>> image.save("cat_with_hat.png") - ``` -""" - - -# Copied from diffusers.pipelines.kandinsky2_2.pipeline_kandinsky2_2.downscale_height_and_width -def downscale_height_and_width(height, width, scale_factor=8): - new_height = height // scale_factor**2 - if height % scale_factor**2 != 0: - new_height += 1 - new_width = width // scale_factor**2 - if width % scale_factor**2 != 0: - new_width += 1 - return new_height * scale_factor, new_width * scale_factor - - -# Copied from diffusers.pipelines.kandinsky.pipeline_kandinsky_inpaint.prepare_mask -def prepare_mask(masks): - prepared_masks = [] - for mask in masks: - old_mask = deepcopy(mask) - for i in range(mask.shape[1]): - for j in range(mask.shape[2]): - if old_mask[0][i][j] == 1: - continue - if i != 0: - mask[:, i - 1, j] = 0 - if j != 0: - mask[:, i, j - 1] = 0 - if i != 0 and j != 0: - mask[:, i - 1, j - 1] = 0 - if i != mask.shape[1] - 1: - mask[:, i + 1, j] = 0 - if j != mask.shape[2] - 1: - mask[:, i, j + 1] = 0 - if i != mask.shape[1] - 1 and j != mask.shape[2] - 1: - mask[:, i + 1, j + 1] = 0 - prepared_masks.append(mask) - return torch.stack(prepared_masks, dim=0) - - -# Copied from diffusers.pipelines.kandinsky.pipeline_kandinsky_inpaint.prepare_mask_and_masked_image -def prepare_mask_and_masked_image(image, mask, height, width): - r""" - Prepares a pair (mask, image) to be consumed by the Kandinsky inpaint pipeline. This means that those inputs will - be converted to ``torch.Tensor`` with shapes ``batch x channels x height x width`` where ``channels`` is ``3`` for - the ``image`` and ``1`` for the ``mask``. - - The ``image`` will be converted to ``torch.float32`` and normalized to be in ``[-1, 1]``. The ``mask`` will be - binarized (``mask > 0.5``) and cast to ``torch.float32`` too. - - Args: - image (Union[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. - It can be a ``PIL.Image``, or a ``height x width`` ``np.array`` or a ``1 x height x width`` - ``torch.Tensor`` or a ``batch x 1 x height x width`` ``torch.Tensor``. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - - - Raises: - ValueError: ``torch.Tensor`` images should be in the ``[-1, 1]`` range. ValueError: ``torch.Tensor`` mask - should be in the ``[0, 1]`` range. ValueError: ``mask`` and ``image`` should have the same spatial dimensions. - TypeError: ``mask`` is a ``torch.Tensor`` but ``image`` is not - (ot the other way around). - - Returns: - tuple[torch.Tensor]: The pair (mask, image) as ``torch.Tensor`` with 4 - dimensions: ``batch x channels x height x width``. - """ - - if image is None: - raise ValueError("`image` input cannot be undefined.") - - if mask is None: - raise ValueError("`mask_image` input cannot be undefined.") - - if isinstance(image, torch.Tensor): - if not isinstance(mask, torch.Tensor): - raise TypeError(f"`image` is a torch.Tensor but `mask` (type: {type(mask)} is not") - - # Batch single image - if image.ndim == 3: - assert image.shape[0] == 3, "Image outside a batch should be of shape (3, H, W)" - image = image.unsqueeze(0) - - # Batch and add channel dim for single mask - if mask.ndim == 2: - mask = mask.unsqueeze(0).unsqueeze(0) - - # Batch single mask or add channel dim - if mask.ndim == 3: - # Single batched mask, no channel dim or single mask not batched but channel dim - if mask.shape[0] == 1: - mask = mask.unsqueeze(0) - - # Batched masks no channel dim - else: - mask = mask.unsqueeze(1) - - assert image.ndim == 4 and mask.ndim == 4, "Image and Mask must have 4 dimensions" - assert image.shape[-2:] == mask.shape[-2:], "Image and Mask must have the same spatial dimensions" - assert image.shape[0] == mask.shape[0], "Image and Mask must have the same batch size" - - # Check image is in [-1, 1] - if image.min() < -1 or image.max() > 1: - raise ValueError("Image should be in [-1, 1] range") - - # Check mask is in [0, 1] - if mask.min() < 0 or mask.max() > 1: - raise ValueError("Mask should be in [0, 1] range") - - # Binarize mask - mask[mask < 0.5] = 0 - mask[mask >= 0.5] = 1 - - # Image as float32 - image = image.to(dtype=torch.float32) - elif isinstance(mask, torch.Tensor): - raise TypeError(f"`mask` is a torch.Tensor but `image` (type: {type(image)} is not") - else: - # preprocess image - if isinstance(image, (PIL.Image.Image, np.ndarray)): - image = [image] - - if isinstance(image, list) and isinstance(image[0], PIL.Image.Image): - # resize all images w.r.t passed height an width - image = [i.resize((width, height), resample=Image.BICUBIC, reducing_gap=1) for i in image] - image = [np.array(i.convert("RGB"))[None, :] for i in image] - image = np.concatenate(image, axis=0) - elif isinstance(image, list) and isinstance(image[0], np.ndarray): - image = np.concatenate([i[None, :] for i in image], axis=0) - - image = image.transpose(0, 3, 1, 2) - image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 - - # preprocess mask - if isinstance(mask, (PIL.Image.Image, np.ndarray)): - mask = [mask] - - if isinstance(mask, list) and isinstance(mask[0], PIL.Image.Image): - mask = [i.resize((width, height), resample=PIL.Image.LANCZOS) for i in mask] - mask = np.concatenate([np.array(m.convert("L"))[None, None, :] for m in mask], axis=0) - mask = mask.astype(np.float32) / 255.0 - elif isinstance(mask, list) and isinstance(mask[0], np.ndarray): - mask = np.concatenate([m[None, None, :] for m in mask], axis=0) - - mask[mask < 0.5] = 0 - mask[mask >= 0.5] = 1 - mask = torch.from_numpy(mask) - - mask = 1 - mask - - return mask, image - - -class KandinskyV22InpaintPipeline(DiffusionPipeline): - """ - Pipeline for text-guided image inpainting using Kandinsky2.1 - - 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.) - - Args: - scheduler ([`DDIMScheduler`]): - A scheduler to be used in combination with `unet` to generate image latents. - unet ([`UNet2DConditionModel`]): - Conditional U-Net architecture to denoise the image embedding. - movq ([`VQModel`]): - MoVQ Decoder to generate the image from the latents. - """ - - model_cpu_offload_seq = "unet->movq" - _callback_tensor_inputs = ["latents", "image_embeds", "negative_image_embeds", "masked_image", "mask_image"] - - def __init__( - self, - unet: UNet2DConditionModel, - scheduler: DDPMScheduler, - movq: VQModel, - ): - super().__init__() - - self.register_modules( - unet=unet, - scheduler=scheduler, - movq=movq, - ) - self.movq_scale_factor = 2 ** (len(self.movq.config.block_out_channels) - 1) - self._warn_has_been_called = False - - # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - @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 - - @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]], - 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", - return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], 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]`): - 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 - be masked out with `mask_image` and repainted according to `prompt`. - mask_image (`np.array`): - Tensor representing an image batch, to mask `image`. White pixels in the mask will be 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)`. - 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. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - 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. - 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. - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - 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. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - """ - if not self._warn_has_been_called and version.parse(version.parse(__version__).base_version) < version.parse( - "0.23.0.dev0" - ): - logger.warning( - "Please note that the expected format of `mask_image` has recently been changed. " - "Before diffusers == 0.19.0, Kandinsky Inpainting pipelines repainted black pixels and preserved black pixels. " - "As of diffusers==0.19.0 this behavior has been inverted. Now white pixels are repainted and black pixels are preserved. " - "This way, Kandinsky's masking behavior is aligned with Stable Diffusion. " - "THIS means that you HAVE to invert the input mask to have the same behavior as before as explained in https://github.com/huggingface/diffusers/pull/4207. " - "This warning will be suppressed after the first inference call and will be removed in diffusers>0.23.0" - ) - self._warn_has_been_called = True - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - 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]}" - ) - - self._guidance_scale = guidance_scale - - device = self._execution_device - - if isinstance(image_embeds, list): - image_embeds = torch.cat(image_embeds, dim=0) - batch_size = image_embeds.shape[0] * num_images_per_prompt - if isinstance(negative_image_embeds, list): - negative_image_embeds = torch.cat(negative_image_embeds, dim=0) - - if self.do_classifier_free_guidance: - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - negative_image_embeds = negative_image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0).to( - dtype=self.unet.dtype, device=device - ) - - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # preprocess image and mask - mask_image, image = prepare_mask_and_masked_image(image, mask_image, height, width) - - image = image.to(dtype=image_embeds.dtype, device=device) - image = self.movq.encode(image)["latents"] - - mask_image = mask_image.to(dtype=image_embeds.dtype, device=device) - - image_shape = tuple(image.shape[-2:]) - mask_image = F.interpolate( - mask_image, - image_shape, - mode="nearest", - ) - mask_image = prepare_mask(mask_image) - masked_image = image * mask_image - - mask_image = mask_image.repeat_interleave(num_images_per_prompt, dim=0) - masked_image = masked_image.repeat_interleave(num_images_per_prompt, dim=0) - if self.do_classifier_free_guidance: - mask_image = mask_image.repeat(2, 1, 1, 1) - masked_image = masked_image.repeat(2, 1, 1, 1) - - num_channels_latents = self.movq.config.latent_channels - - height, width = downscale_height_and_width(height, width, self.movq_scale_factor) - - # create initial latent - latents = self.prepare_latents( - (batch_size, num_channels_latents, height, width), - image_embeds.dtype, - device, - generator, - latents, - self.scheduler, - ) - noise = torch.clone(latents) - - self._num_timesteps = len(timesteps) - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = torch.cat([latent_model_input, masked_image, mask_image], dim=1) - - added_cond_kwargs = {"image_embeds": image_embeds} - noise_pred = self.unet( - sample=latent_model_input, - timestep=t, - encoder_hidden_states=None, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - if self.do_classifier_free_guidance: - noise_pred, variance_pred = noise_pred.split(latents.shape[1], dim=1) - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - _, variance_pred_text = variance_pred.chunk(2) - noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, variance_pred_text], dim=1) - - if not ( - hasattr(self.scheduler.config, "variance_type") - and self.scheduler.config.variance_type in ["learned", "learned_range"] - ): - noise_pred, _ = noise_pred.split(latents.shape[1], dim=1) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step( - noise_pred, - t, - latents, - generator=generator, - )[0] - init_latents_proper = image[:1] - init_mask = mask_image[:1] - - if i < len(timesteps) - 1: - noise_timestep = timesteps[i + 1] - init_latents_proper = self.scheduler.add_noise( - init_latents_proper, noise, torch.tensor([noise_timestep]) - ) - - latents = init_mask * init_latents_proper + (1 - init_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) - image_embeds = callback_outputs.pop("image_embeds", image_embeds) - negative_image_embeds = callback_outputs.pop("negative_image_embeds", negative_image_embeds) - masked_image = callback_outputs.pop("masked_image", masked_image) - mask_image = callback_outputs.pop("mask_image", mask_image) - - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # post-processing - latents = mask_image[:1] * image[:1] + (1 - mask_image[:1]) * latents - - if output_type not in ["pt", "np", "pil", "latent"]: - raise ValueError( - f"Only the output types `pt`, `pil`, `np` and `latent` are supported not output_type={output_type}" - ) - - if not output_type == "latent": - image = self.movq.decode(latents, force_not_quantize=True)["sample"] - - if output_type in ["np", "pil"]: - image = image * 0.5 + 0.5 - image = image.clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - if output_type == "pil": - image = self.numpy_to_pil(image) - else: - image = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py deleted file mode 100644 index 68954c2dc886060dea6085331fca7644ed37aee3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py +++ /dev/null @@ -1,561 +0,0 @@ -from typing import Callable, Dict, List, Optional, Union - -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...models import PriorTransformer -from ...schedulers import UnCLIPScheduler -from ...utils import ( - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..kandinsky import KandinskyPriorPipelineOutput -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 - >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline - >>> import torch - - >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior") - >>> pipe_prior.to("cuda") - >>> prompt = "red cat, 4k photo" - >>> image_emb, negative_image_emb = pipe_prior(prompt).to_tuple() - - >>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder") - >>> pipe.to("cuda") - >>> image = pipe( - ... image_embeds=image_emb, - ... negative_image_embeds=negative_image_emb, - ... height=768, - ... width=768, - ... num_inference_steps=50, - ... ).images - >>> image[0].save("cat.png") - ``` -""" - -EXAMPLE_INTERPOLATE_DOC_STRING = """ - Examples: - ```py - >>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22Pipeline - >>> from diffusers.utils import load_image - >>> import PIL - >>> import torch - >>> from torchvision import transforms - - >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16 - ... ) - >>> pipe_prior.to("cuda") - >>> img1 = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/cat.png" - ... ) - >>> img2 = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/starry_night.jpeg" - ... ) - >>> images_texts = ["a cat", img1, img2] - >>> weights = [0.3, 0.3, 0.4] - >>> out = pipe_prior.interpolate(images_texts, weights) - >>> pipe = KandinskyV22Pipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16 - ... ) - >>> pipe.to("cuda") - >>> image = pipe( - ... image_embeds=out.image_embeds, - ... negative_image_embeds=out.negative_image_embeds, - ... height=768, - ... width=768, - ... num_inference_steps=50, - ... ).images[0] - >>> image.save("starry_cat.png") - ``` -""" - - -class KandinskyV22PriorPipeline(DiffusionPipeline): - """ - Pipeline for generating image prior for Kandinsky - - 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.) - - Args: - prior ([`PriorTransformer`]): - The canonical unCLIP prior to approximate the image embedding from the text embedding. - image_encoder ([`CLIPVisionModelWithProjection`]): - Frozen image-encoder. - text_encoder ([`CLIPTextModelWithProjection`]): - Frozen text-encoder. - tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - scheduler ([`UnCLIPScheduler`]): - A scheduler to be used in combination with `prior` to generate image embedding. - image_processor ([`CLIPImageProcessor`]): - A image_processor to be used to preprocess image from clip. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->prior" - _exclude_from_cpu_offload = ["prior"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "text_encoder_hidden_states", "text_mask"] - - def __init__( - self, - prior: PriorTransformer, - image_encoder: CLIPVisionModelWithProjection, - text_encoder: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - scheduler: UnCLIPScheduler, - image_processor: CLIPImageProcessor, - ): - super().__init__() - - self.register_modules( - prior=prior, - text_encoder=text_encoder, - tokenizer=tokenizer, - scheduler=scheduler, - image_encoder=image_encoder, - image_processor=image_processor, - ) - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_INTERPOLATE_DOC_STRING) - def interpolate( - self, - images_and_prompts: List[Union[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, - negative_prompt: str = "", - guidance_scale: float = 4.0, - device=None, - ): - """ - Function invoked when using the prior pipeline for interpolation. - - Args: - images_and_prompts (`List[Union[str, PIL.Image.Image, torch.Tensor]]`): - list of prompts and images to guide the image generation. - 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*): - 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 ge generated by sampling using the supplied random `generator`. - 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*): - 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): - 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. - - Examples: - - Returns: - [`KandinskyPriorPipelineOutput`] or `tuple` - """ - - device = device or self.device - - if len(images_and_prompts) != len(weights): - raise ValueError( - f"`images_and_prompts` contains {len(images_and_prompts)} items and `weights` contains {len(weights)} items - they should be lists of same length" - ) - - image_embeddings = [] - for cond, weight in zip(images_and_prompts, weights): - if isinstance(cond, str): - image_emb = self( - cond, - num_inference_steps=num_inference_steps, - num_images_per_prompt=num_images_per_prompt, - generator=generator, - latents=latents, - negative_prompt=negative_prior_prompt, - guidance_scale=guidance_scale, - ).image_embeds.unsqueeze(0) - - elif isinstance(cond, (PIL.Image.Image, torch.Tensor)): - if isinstance(cond, PIL.Image.Image): - cond = ( - self.image_processor(cond, return_tensors="pt") - .pixel_values[0] - .unsqueeze(0) - .to(dtype=self.image_encoder.dtype, device=device) - ) - - image_emb = self.image_encoder(cond)["image_embeds"].repeat(num_images_per_prompt, 1).unsqueeze(0) - - else: - raise ValueError( - f"`images_and_prompts` can only contains elements to be of type `str`, `PIL.Image.Image` or `torch.Tensor` but is {type(cond)}" - ) - - image_embeddings.append(image_emb * weight) - - image_emb = torch.cat(image_embeddings).sum(dim=0) - - out_zero = self( - negative_prompt, - num_inference_steps=num_inference_steps, - num_images_per_prompt=num_images_per_prompt, - generator=generator, - latents=latents, - negative_prompt=negative_prior_prompt, - guidance_scale=guidance_scale, - ) - zero_image_emb = out_zero.negative_image_embeds if negative_prompt == "" else out_zero.image_embeds - - return KandinskyPriorPipelineOutput(image_embeds=image_emb, negative_image_embeds=zero_image_emb) - - # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - # Copied from diffusers.pipelines.kandinsky.pipeline_kandinsky_prior.KandinskyPriorPipeline.get_zero_embed - def get_zero_embed(self, batch_size=1, device=None): - device = device or self.device - zero_img = torch.zeros(1, 3, self.image_encoder.config.image_size, self.image_encoder.config.image_size).to( - device=device, dtype=self.image_encoder.dtype - ) - zero_image_emb = self.image_encoder(zero_img)["image_embeds"] - zero_image_emb = zero_image_emb.repeat(batch_size, 1) - return zero_image_emb - - # Copied from diffusers.pipelines.kandinsky.pipeline_kandinsky_prior.KandinskyPriorPipeline._encode_prompt - def _encode_prompt( - self, - prompt, - device, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt=None, - ): - batch_size = len(prompt) if isinstance(prompt, list) else 1 - # get prompt text embeddings - 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 - text_mask = text_inputs.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[:, 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}" - ) - text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length] - - text_encoder_output = self.text_encoder(text_input_ids.to(device)) - - prompt_embeds = text_encoder_output.text_embeds - text_encoder_hidden_states = text_encoder_output.last_hidden_state - - prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - text_encoder_hidden_states = text_encoder_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - uncond_tokens: List[str] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif 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 - - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - uncond_text_mask = uncond_input.attention_mask.bool().to(device) - negative_prompt_embeds_text_encoder_output = self.text_encoder(uncond_input.input_ids.to(device)) - - negative_prompt_embeds = negative_prompt_embeds_text_encoder_output.text_embeds - uncond_text_encoder_hidden_states = negative_prompt_embeds_text_encoder_output.last_hidden_state - - # 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.repeat(1, num_images_per_prompt) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len) - - seq_len = uncond_text_encoder_hidden_states.shape[1] - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.repeat(1, num_images_per_prompt, 1) - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.view( - batch_size * num_images_per_prompt, seq_len, -1 - ) - uncond_text_mask = uncond_text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - # done duplicates - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - text_encoder_hidden_states = torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]) - - text_mask = torch.cat([uncond_text_mask, text_mask]) - - return prompt_embeds, text_encoder_hidden_states, text_mask - - @property - def do_classifier_free_guidance(self): - return self._guidance_scale > 1 - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = 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, - guidance_scale: float = 4.0, - output_type: Optional[str] = "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"], - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `List[str]`): - The prompt or prompts to guide the image generation. - 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): - 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*): - 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 ge generated by sampling using the supplied random `generator`. - 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. - output_type (`str`, *optional*, defaults to `"pt"`): - The output format of the generate image. Choose between: `"np"` (`np.array`) or `"pt"` - (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - 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. - - Examples: - - Returns: - [`KandinskyPriorPipelineOutput`] or `tuple` - """ - - 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 isinstance(prompt, str): - prompt = [prompt] - elif not isinstance(prompt, list): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - if isinstance(negative_prompt, str): - negative_prompt = [negative_prompt] - elif not isinstance(negative_prompt, list) and negative_prompt is not None: - raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}") - - # if the negative prompt is defined we double the batch size to - # directly retrieve the negative prompt embedding - if negative_prompt is not None: - prompt = prompt + negative_prompt - negative_prompt = 2 * negative_prompt - - device = self._execution_device - - batch_size = len(prompt) - batch_size = batch_size * num_images_per_prompt - - self._guidance_scale = guidance_scale - - prompt_embeds, text_encoder_hidden_states, text_mask = self._encode_prompt( - prompt, device, num_images_per_prompt, self.do_classifier_free_guidance, negative_prompt - ) - - # prior - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - embedding_dim = self.prior.config.embedding_dim - - latents = self.prepare_latents( - (batch_size, embedding_dim), - prompt_embeds.dtype, - device, - generator, - latents, - self.scheduler, - ) - self._num_timesteps = len(timesteps) - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - predicted_image_embedding = self.prior( - latent_model_input, - timestep=t, - proj_embedding=prompt_embeds, - encoder_hidden_states=text_encoder_hidden_states, - attention_mask=text_mask, - ).predicted_image_embedding - - if self.do_classifier_free_guidance: - predicted_image_embedding_uncond, predicted_image_embedding_text = predicted_image_embedding.chunk(2) - predicted_image_embedding = predicted_image_embedding_uncond + self.guidance_scale * ( - predicted_image_embedding_text - predicted_image_embedding_uncond - ) - - if i + 1 == timesteps.shape[0]: - prev_timestep = None - else: - prev_timestep = timesteps[i + 1] - - latents = self.scheduler.step( - predicted_image_embedding, - timestep=t, - sample=latents, - generator=generator, - prev_timestep=prev_timestep, - ).prev_sample - - 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) - text_encoder_hidden_states = callback_outputs.pop( - "text_encoder_hidden_states", text_encoder_hidden_states - ) - text_mask = callback_outputs.pop("text_mask", text_mask) - - if XLA_AVAILABLE: - xm.mark_step() - - latents = self.prior.post_process_latents(latents) - - image_embeddings = latents - - # if negative prompt has been defined, we retrieve split the image embedding into two - if negative_prompt is None: - zero_embeds = self.get_zero_embed(latents.shape[0], device=latents.device) - else: - image_embeddings, zero_embeds = image_embeddings.chunk(2) - - self.maybe_free_model_hooks() - - if output_type not in ["pt", "np"]: - raise ValueError(f"Only the output types `pt` and `np` are supported not output_type={output_type}") - - if output_type == "np": - image_embeddings = image_embeddings.cpu().numpy() - zero_embeds = zero_embeds.cpu().numpy() - - if not return_dict: - return (image_embeddings, zero_embeds) - - return KandinskyPriorPipelineOutput(image_embeds=image_embeddings, negative_image_embeds=zero_embeds) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py deleted file mode 100644 index 13ea2ad6af63a7e1ba6d705fe9c821f081f0065a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py +++ /dev/null @@ -1,575 +0,0 @@ -from typing import List, Optional, Union - -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...models import PriorTransformer -from ...schedulers import UnCLIPScheduler -from ...utils import ( - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..kandinsky import KandinskyPriorPipelineOutput -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 - >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorEmb2EmbPipeline - >>> import torch - - >>> pipe_prior = KandinskyPriorPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16 - ... ) - >>> pipe_prior.to("cuda") - - >>> prompt = "red cat, 4k photo" - >>> img = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/cat.png" - ... ) - >>> image_emb, nagative_image_emb = pipe_prior(prompt, image=img, strength=0.2).to_tuple() - - >>> pipe = KandinskyPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-2-decoder, torch_dtype=torch.float16" - ... ) - >>> pipe.to("cuda") - - >>> image = pipe( - ... image_embeds=image_emb, - ... negative_image_embeds=negative_image_emb, - ... height=768, - ... width=768, - ... num_inference_steps=100, - ... ).images - - >>> image[0].save("cat.png") - ``` -""" - -EXAMPLE_INTERPOLATE_DOC_STRING = """ - Examples: - ```py - >>> from diffusers import KandinskyV22PriorEmb2EmbPipeline, KandinskyV22Pipeline - >>> from diffusers.utils import load_image - >>> import PIL - - >>> import torch - >>> from torchvision import transforms - - >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16 - ... ) - >>> pipe_prior.to("cuda") - - >>> img1 = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/cat.png" - ... ) - - >>> img2 = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/starry_night.jpeg" - ... ) - - >>> images_texts = ["a cat", img1, img2] - >>> weights = [0.3, 0.3, 0.4] - >>> image_emb, zero_image_emb = pipe_prior.interpolate(images_texts, weights) - - >>> pipe = KandinskyV22Pipeline.from_pretrained( - ... "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16 - ... ) - >>> pipe.to("cuda") - - >>> image = pipe( - ... image_embeds=image_emb, - ... negative_image_embeds=zero_image_emb, - ... height=768, - ... width=768, - ... num_inference_steps=150, - ... ).images[0] - - >>> image.save("starry_cat.png") - ``` -""" - - -class KandinskyV22PriorEmb2EmbPipeline(DiffusionPipeline): - """ - Pipeline for generating image prior for Kandinsky - - 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.) - - Args: - prior ([`PriorTransformer`]): - The canonical unCLIP prior to approximate the image embedding from the text embedding. - image_encoder ([`CLIPVisionModelWithProjection`]): - Frozen image-encoder. - text_encoder ([`CLIPTextModelWithProjection`]): - Frozen text-encoder. - tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - scheduler ([`UnCLIPScheduler`]): - A scheduler to be used in combination with `prior` to generate image embedding. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->prior" - _exclude_from_cpu_offload = ["prior"] - - def __init__( - self, - prior: PriorTransformer, - image_encoder: CLIPVisionModelWithProjection, - text_encoder: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - scheduler: UnCLIPScheduler, - image_processor: CLIPImageProcessor, - ): - super().__init__() - - self.register_modules( - prior=prior, - text_encoder=text_encoder, - tokenizer=tokenizer, - scheduler=scheduler, - image_encoder=image_encoder, - image_processor=image_processor, - ) - - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - timesteps = self.scheduler.timesteps[t_start:] - - return timesteps, num_inference_steps - t_start - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_INTERPOLATE_DOC_STRING) - def interpolate( - self, - images_and_prompts: List[Union[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, - negative_prompt: str = "", - guidance_scale: float = 4.0, - device=None, - ): - """ - Function invoked when using the prior pipeline for interpolation. - - Args: - images_and_prompts (`List[Union[str, PIL.Image.Image, torch.Tensor]]`): - list of prompts and images to guide the image generation. - 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*): - 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 ge generated by sampling using the supplied random `generator`. - 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*): - 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): - 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. - - Examples: - - Returns: - [`KandinskyPriorPipelineOutput`] or `tuple` - """ - - device = device or self.device - - if len(images_and_prompts) != len(weights): - raise ValueError( - f"`images_and_prompts` contains {len(images_and_prompts)} items and `weights` contains {len(weights)} items - they should be lists of same length" - ) - - image_embeddings = [] - for cond, weight in zip(images_and_prompts, weights): - if isinstance(cond, str): - image_emb = self( - cond, - num_inference_steps=num_inference_steps, - num_images_per_prompt=num_images_per_prompt, - generator=generator, - latents=latents, - negative_prompt=negative_prior_prompt, - guidance_scale=guidance_scale, - ).image_embeds.unsqueeze(0) - - elif isinstance(cond, (PIL.Image.Image, torch.Tensor)): - image_emb = self._encode_image( - cond, device=device, num_images_per_prompt=num_images_per_prompt - ).unsqueeze(0) - - else: - raise ValueError( - f"`images_and_prompts` can only contains elements to be of type `str`, `PIL.Image.Image` or `torch.Tensor` but is {type(cond)}" - ) - - image_embeddings.append(image_emb * weight) - - image_emb = torch.cat(image_embeddings).sum(dim=0) - - return KandinskyPriorPipelineOutput(image_embeds=image_emb, negative_image_embeds=torch.randn_like(image_emb)) - - def _encode_image( - self, - image: Union[torch.Tensor, List[PIL.Image.Image]], - device, - num_images_per_prompt, - ): - if not isinstance(image, torch.Tensor): - image = self.image_processor(image, return_tensors="pt").pixel_values.to( - dtype=self.image_encoder.dtype, device=device - ) - - image_emb = self.image_encoder(image)["image_embeds"] # B, D - image_emb = image_emb.repeat_interleave(num_images_per_prompt, dim=0) - image_emb.to(device=device) - - return image_emb - - def prepare_latents(self, emb, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None): - emb = emb.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - init_latents = emb - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - latents = init_latents - - return latents - - # Copied from diffusers.pipelines.kandinsky.pipeline_kandinsky_prior.KandinskyPriorPipeline.get_zero_embed - def get_zero_embed(self, batch_size=1, device=None): - device = device or self.device - zero_img = torch.zeros(1, 3, self.image_encoder.config.image_size, self.image_encoder.config.image_size).to( - device=device, dtype=self.image_encoder.dtype - ) - zero_image_emb = self.image_encoder(zero_img)["image_embeds"] - zero_image_emb = zero_image_emb.repeat(batch_size, 1) - return zero_image_emb - - # Copied from diffusers.pipelines.kandinsky.pipeline_kandinsky_prior.KandinskyPriorPipeline._encode_prompt - def _encode_prompt( - self, - prompt, - device, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt=None, - ): - batch_size = len(prompt) if isinstance(prompt, list) else 1 - # get prompt text embeddings - 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 - text_mask = text_inputs.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[:, 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}" - ) - text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length] - - text_encoder_output = self.text_encoder(text_input_ids.to(device)) - - prompt_embeds = text_encoder_output.text_embeds - text_encoder_hidden_states = text_encoder_output.last_hidden_state - - prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - text_encoder_hidden_states = text_encoder_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - uncond_tokens: List[str] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif 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 - - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - uncond_text_mask = uncond_input.attention_mask.bool().to(device) - negative_prompt_embeds_text_encoder_output = self.text_encoder(uncond_input.input_ids.to(device)) - - negative_prompt_embeds = negative_prompt_embeds_text_encoder_output.text_embeds - uncond_text_encoder_hidden_states = negative_prompt_embeds_text_encoder_output.last_hidden_state - - # 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.repeat(1, num_images_per_prompt) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len) - - seq_len = uncond_text_encoder_hidden_states.shape[1] - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.repeat(1, num_images_per_prompt, 1) - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.view( - batch_size * num_images_per_prompt, seq_len, -1 - ) - uncond_text_mask = uncond_text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - # done duplicates - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - text_encoder_hidden_states = torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]) - - text_mask = torch.cat([uncond_text_mask, text_mask]) - - return prompt_embeds, text_encoder_hidden_states, text_mask - - @torch.no_grad() - @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]], - strength: float = 0.3, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: int = 1, - num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - guidance_scale: float = 4.0, - output_type: Optional[str] = "pt", # pt only - return_dict: bool = True, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - 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` - 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. - emb (`torch.Tensor`): - The image embedding. - 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): - 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*): - 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): - 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. - output_type (`str`, *optional*, defaults to `"pt"`): - The output format of the generate image. Choose between: `"np"` (`np.array`) or `"pt"` - (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`KandinskyPriorPipelineOutput`] or `tuple` - """ - - if isinstance(prompt, str): - prompt = [prompt] - elif not isinstance(prompt, list): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - if isinstance(negative_prompt, str): - negative_prompt = [negative_prompt] - elif not isinstance(negative_prompt, list) and negative_prompt is not None: - raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}") - - # if the negative prompt is defined we double the batch size to - # directly retrieve the negative prompt embedding - if negative_prompt is not None: - prompt = prompt + negative_prompt - negative_prompt = 2 * negative_prompt - - device = self._execution_device - - batch_size = len(prompt) - batch_size = batch_size * num_images_per_prompt - - do_classifier_free_guidance = guidance_scale > 1.0 - prompt_embeds, text_encoder_hidden_states, text_mask = self._encode_prompt( - prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt - ) - - if not isinstance(image, List): - image = [image] - - if isinstance(image[0], torch.Tensor): - image = torch.cat(image, dim=0) - - if isinstance(image, torch.Tensor) and image.ndim == 2: - # allow user to pass image_embeds directly - image_embeds = image.repeat_interleave(num_images_per_prompt, dim=0) - elif isinstance(image, torch.Tensor) and image.ndim != 4: - raise ValueError( - f" if pass `image` as pytorch tensor, or a list of pytorch tensor, please make sure each tensor has shape [batch_size, channels, height, width], currently {image[0].unsqueeze(0).shape}" - ) - else: - image_embeds = self._encode_image(image, device, num_images_per_prompt) - - # prior - self.scheduler.set_timesteps(num_inference_steps, device=device) - - latents = image_embeds - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) - latent_timestep = timesteps[:1].repeat(batch_size) - latents = self.prepare_latents( - latents, - latent_timestep, - batch_size // num_images_per_prompt, - num_images_per_prompt, - prompt_embeds.dtype, - device, - generator, - ) - - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - predicted_image_embedding = self.prior( - latent_model_input, - timestep=t, - proj_embedding=prompt_embeds, - encoder_hidden_states=text_encoder_hidden_states, - attention_mask=text_mask, - ).predicted_image_embedding - - if do_classifier_free_guidance: - predicted_image_embedding_uncond, predicted_image_embedding_text = predicted_image_embedding.chunk(2) - predicted_image_embedding = predicted_image_embedding_uncond + guidance_scale * ( - predicted_image_embedding_text - predicted_image_embedding_uncond - ) - - if i + 1 == timesteps.shape[0]: - prev_timestep = None - else: - prev_timestep = timesteps[i + 1] - - latents = self.scheduler.step( - predicted_image_embedding, - timestep=t, - sample=latents, - generator=generator, - prev_timestep=prev_timestep, - ).prev_sample - - if XLA_AVAILABLE: - xm.mark_step() - - latents = self.prior.post_process_latents(latents) - - image_embeddings = latents - - # if negative prompt has been defined, we retrieve split the image embedding into two - if negative_prompt is None: - zero_embeds = self.get_zero_embed(latents.shape[0], device=latents.device) - else: - image_embeddings, zero_embeds = image_embeddings.chunk(2) - - self.maybe_free_model_hooks() - - if output_type not in ["pt", "np"]: - raise ValueError(f"Only the output types `pt` and `np` are supported not output_type={output_type}") - - if output_type == "np": - image_embeddings = image_embeddings.cpu().numpy() - zero_embeds = zero_embeds.cpu().numpy() - - if not return_dict: - return (image_embeddings, zero_embeds) - - return KandinskyPriorPipelineOutput(image_embeds=image_embeddings, negative_image_embeds=zero_embeds) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__init__.py deleted file mode 100644 index e8a3063141b5e62682ce3ad8e8e0d02473402665..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__init__.py +++ /dev/null @@ -1,49 +0,0 @@ -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_kandinsky3"] = ["Kandinsky3Pipeline"] - _import_structure["pipeline_kandinsky3_img2img"] = ["Kandinsky3Img2ImgPipeline"] - - -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_kandinsky3 import Kandinsky3Pipeline - from .pipeline_kandinsky3_img2img import Kandinsky3Img2ImgPipeline -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/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 99d5c065d6a27d81fa74faf5736fcedebe93f5e6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__pycache__/convert_kandinsky3_unet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__pycache__/convert_kandinsky3_unet.cpython-310.pyc deleted file mode 100644 index ff19ae8f654594eeeab360c16781a8770de32c87..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__pycache__/convert_kandinsky3_unet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__pycache__/pipeline_kandinsky3.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__pycache__/pipeline_kandinsky3.cpython-310.pyc deleted file mode 100644 index c10c1a41d4b02fcbf2f87b37ce336b8b3a066b25..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__pycache__/pipeline_kandinsky3.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__pycache__/pipeline_kandinsky3_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__pycache__/pipeline_kandinsky3_img2img.cpython-310.pyc deleted file mode 100644 index cc16a4cffc914ab06a008dfd9c0d440e27912cf5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/__pycache__/pipeline_kandinsky3_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/convert_kandinsky3_unet.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/convert_kandinsky3_unet.py deleted file mode 100644 index 5360632275b41b035f94d84d6156a6bafe84c613..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/convert_kandinsky3_unet.py +++ /dev/null @@ -1,98 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import fnmatch - -from safetensors.torch import load_file - -from diffusers import Kandinsky3UNet - - -MAPPING = { - "to_time_embed.1": "time_embedding.linear_1", - "to_time_embed.3": "time_embedding.linear_2", - "in_layer": "conv_in", - "out_layer.0": "conv_norm_out", - "out_layer.2": "conv_out", - "down_samples": "down_blocks", - "up_samples": "up_blocks", - "projection_lin": "encoder_hid_proj.projection_linear", - "projection_ln": "encoder_hid_proj.projection_norm", - "feature_pooling": "add_time_condition", - "to_query": "to_q", - "to_key": "to_k", - "to_value": "to_v", - "output_layer": "to_out.0", - "self_attention_block": "attentions.0", -} - -DYNAMIC_MAP = { - "resnet_attn_blocks.*.0": "resnets_in.*", - "resnet_attn_blocks.*.1": ("attentions.*", 1), - "resnet_attn_blocks.*.2": "resnets_out.*", -} -# MAPPING = {} - - -def convert_state_dict(unet_state_dict): - """ - Args: - Convert the state dict of a U-Net model to match the key format expected by Kandinsky3UNet model. - unet_model (torch.nn.Module): The original U-Net model. unet_kandi3_model (torch.nn.Module): The Kandinsky3UNet - model to match keys with. - - Returns: - OrderedDict: The converted state dictionary. - """ - # Example of renaming logic (this will vary based on your model's architecture) - converted_state_dict = {} - for key in unet_state_dict: - new_key = key - for pattern, new_pattern in MAPPING.items(): - new_key = new_key.replace(pattern, new_pattern) - - for dyn_pattern, dyn_new_pattern in DYNAMIC_MAP.items(): - has_matched = False - if fnmatch.fnmatch(new_key, f"*.{dyn_pattern}.*") and not has_matched: - star = int(new_key.split(dyn_pattern.split(".")[0])[-1].split(".")[1]) - - if isinstance(dyn_new_pattern, tuple): - new_star = star + dyn_new_pattern[-1] - dyn_new_pattern = dyn_new_pattern[0] - else: - new_star = star - - pattern = dyn_pattern.replace("*", str(star)) - new_pattern = dyn_new_pattern.replace("*", str(new_star)) - - new_key = new_key.replace(pattern, new_pattern) - has_matched = True - - converted_state_dict[new_key] = unet_state_dict[key] - - return converted_state_dict - - -def main(model_path, output_path): - # Load your original U-Net model - unet_state_dict = load_file(model_path) - - # Initialize your Kandinsky3UNet model - config = {} - - # Convert the state dict - converted_state_dict = convert_state_dict(unet_state_dict) - - unet = Kandinsky3UNet(config) - unet.load_state_dict(converted_state_dict) - - unet.save_pretrained(output_path) - print(f"Converted model saved to {output_path}") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Convert U-Net PyTorch model to Kandinsky3UNet format") - parser.add_argument("--model_path", type=str, required=True, help="Path to the original U-Net PyTorch model") - parser.add_argument("--output_path", type=str, required=True, help="Path to save the converted model") - - args = parser.parse_args() - main(args.model_path, args.output_path) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py deleted file mode 100644 index 57cc0270442d8ebab810f75fcf5739ef89ceb19c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py +++ /dev/null @@ -1,588 +0,0 @@ -from typing import Callable, Dict, List, Optional, Union - -import torch -from transformers import T5EncoderModel, T5Tokenizer - -from ...loaders import StableDiffusionLoraLoaderMixin -from ...models import Kandinsky3UNet, VQModel -from ...schedulers import DDPMScheduler -from ...utils import ( - deprecate, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -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 - >>> from diffusers import AutoPipelineForText2Image - >>> import torch - - >>> pipe = AutoPipelineForText2Image.from_pretrained( - ... "kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16 - ... ) - >>> pipe.enable_model_cpu_offload() - - >>> prompt = "A photograph of the inside of a subway train. There are raccoons sitting on the seats. One of them is reading a newspaper. The window shows the city in the background." - - >>> generator = torch.Generator(device="cpu").manual_seed(0) - >>> image = pipe(prompt, num_inference_steps=25, generator=generator).images[0] - ``` - -""" - - -def downscale_height_and_width(height, width, scale_factor=8): - new_height = height // scale_factor**2 - if height % scale_factor**2 != 0: - new_height += 1 - new_width = width // scale_factor**2 - if width % scale_factor**2 != 0: - new_width += 1 - return new_height * scale_factor, new_width * scale_factor - - -class Kandinsky3Pipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin): - model_cpu_offload_seq = "text_encoder->unet->movq" - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "negative_attention_mask", - "attention_mask", - ] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - unet: Kandinsky3UNet, - scheduler: DDPMScheduler, - movq: VQModel, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_encoder, unet=unet, scheduler=scheduler, movq=movq - ) - - def process_embeds(self, embeddings, attention_mask, cut_context): - if cut_context: - embeddings[attention_mask == 0] = torch.zeros_like(embeddings[attention_mask == 0]) - max_seq_length = attention_mask.sum(-1).max() + 1 - embeddings = embeddings[:, :max_seq_length] - attention_mask = attention_mask[:, :max_seq_length] - return embeddings, attention_mask - - @torch.no_grad() - def encode_prompt( - self, - prompt, - do_classifier_free_guidance=True, - num_images_per_prompt=1, - device=None, - negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - _cut_context=False, - attention_mask: Optional[torch.Tensor] = None, - negative_attention_mask: Optional[torch.Tensor] = None, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`, *optional*): - prompt to be encoded - device: (`torch.device`, *optional*): - torch device to place the resulting embeddings on - num_images_per_prompt (`int`, *optional*, defaults to 1): - 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*): - 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`). - 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. - attention_mask (`torch.Tensor`, *optional*): - Pre-generated attention mask. Must provide if passing `prompt_embeds` directly. - negative_attention_mask (`torch.Tensor`, *optional*): - Pre-generated negative attention mask. Must provide if passing `negative_prompt_embeds` directly. - """ - if prompt is not None and negative_prompt is not None: - 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)} !=" - f" {type(prompt)}." - ) - - if device is None: - device = self._execution_device - - 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] - - max_length = 128 - - if prompt_embeds is None: - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids.to(device) - attention_mask = text_inputs.attention_mask.to(device) - prompt_embeds = self.text_encoder( - text_input_ids, - attention_mask=attention_mask, - ) - prompt_embeds = prompt_embeds[0] - prompt_embeds, attention_mask = self.process_embeds(prompt_embeds, attention_mask, _cut_context) - prompt_embeds = prompt_embeds * attention_mask.unsqueeze(2) - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=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) - 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] - - if negative_prompt is None: - uncond_tokens = [""] * batch_size - 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 - if negative_prompt is not None: - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=128, - truncation=True, - return_attention_mask=True, - return_tensors="pt", - ) - text_input_ids = uncond_input.input_ids.to(device) - negative_attention_mask = uncond_input.attention_mask.to(device) - - negative_prompt_embeds = self.text_encoder( - text_input_ids, - attention_mask=negative_attention_mask, - ) - negative_prompt_embeds = negative_prompt_embeds[0] - negative_prompt_embeds = negative_prompt_embeds[:, : prompt_embeds.shape[1]] - negative_attention_mask = negative_attention_mask[:, : prompt_embeds.shape[1]] - negative_prompt_embeds = negative_prompt_embeds * negative_attention_mask.unsqueeze(2) - - else: - negative_prompt_embeds = torch.zeros_like(prompt_embeds) - negative_attention_mask = torch.zeros_like(attention_mask) - - 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=dtype, device=device) - if negative_prompt_embeds.shape != prompt_embeds.shape: - 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) - negative_attention_mask = negative_attention_mask.repeat(num_images_per_prompt, 1) - - # 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 - else: - negative_prompt_embeds = None - negative_attention_mask = None - return prompt_embeds, negative_prompt_embeds, attention_mask, negative_attention_mask - - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - def check_inputs( - self, - prompt, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - attention_mask=None, - negative_attention_mask=None, - ): - 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}." - ) - if negative_prompt_embeds is not None and negative_attention_mask is None: - raise ValueError("Please provide `negative_attention_mask` along with `negative_prompt_embeds`") - - if negative_prompt_embeds is not None and negative_attention_mask is not None: - if negative_prompt_embeds.shape[:2] != negative_attention_mask.shape: - raise ValueError( - "`negative_prompt_embeds` and `negative_attention_mask` must have the same batch_size and token length when passed directly, but" - f" got: `negative_prompt_embeds` {negative_prompt_embeds.shape[:2]} != `negative_attention_mask`" - f" {negative_attention_mask.shape}." - ) - - if prompt_embeds is not None and attention_mask is None: - raise ValueError("Please provide `attention_mask` along with `prompt_embeds`") - - if prompt_embeds is not None and attention_mask is not None: - if prompt_embeds.shape[:2] != attention_mask.shape: - raise ValueError( - "`prompt_embeds` and `attention_mask` must have the same batch_size and token length when passed directly, but" - f" got: `prompt_embeds` {prompt_embeds.shape[:2]} != `attention_mask`" - f" {attention_mask.shape}." - ) - - @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 - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[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", - 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"], - **kwargs, - ): - """ - 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. - 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*): - 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): - 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. - height (`int`, *optional*, defaults to self.unet.config.sample_size): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to self.unet.config.sample_size): - The width in pixels of the generated image. - 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*): - 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*): - 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. - attention_mask (`torch.Tensor`, *optional*): - Pre-generated attention mask. Must provide if passing `prompt_embeds` directly. - negative_attention_mask (`torch.Tensor`, *optional*): - Pre-generated negative attention mask. Must provide if passing `negative_prompt_embeds` directly. - 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.IFPipelineOutput`] 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. - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - cross_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). - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - 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]}" - ) - - cut_context = True - device = self._execution_device - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - callback_on_step_end_tensor_inputs, - attention_mask, - negative_attention_mask, - ) - - self._guidance_scale = guidance_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] - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds, attention_mask, negative_attention_mask = self.encode_prompt( - prompt, - self.do_classifier_free_guidance, - num_images_per_prompt=num_images_per_prompt, - device=device, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - _cut_context=cut_context, - attention_mask=attention_mask, - negative_attention_mask=negative_attention_mask, - ) - - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - attention_mask = torch.cat([negative_attention_mask, attention_mask]).bool() - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latents - height, width = downscale_height_and_width(height, width, 8) - - latents = self.prepare_latents( - (batch_size * num_images_per_prompt, 4, height, width), - prompt_embeds.dtype, - device, - generator, - latents, - self.scheduler, - ) - - if hasattr(self, "text_encoder_offload_hook") and self.text_encoder_offload_hook is not None: - self.text_encoder_offload_hook.offload() - - # 7. Denoising loop - 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): - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - encoder_attention_mask=attention_mask, - return_dict=False, - )[0] - - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - - noise_pred = (guidance_scale + 1.0) * noise_pred_text - guidance_scale * noise_pred_uncond - # noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step( - noise_pred, - t, - latents, - generator=generator, - ).prev_sample - - 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) - attention_mask = callback_outputs.pop("attention_mask", attention_mask) - negative_attention_mask = callback_outputs.pop("negative_attention_mask", negative_attention_mask) - - if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): - progress_bar.update() - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # post-processing - if output_type not in ["pt", "np", "pil", "latent"]: - raise ValueError( - f"Only the output types `pt`, `pil`, `np` and `latent` are supported not output_type={output_type}" - ) - - if not output_type == "latent": - image = self.movq.decode(latents, force_not_quantize=True)["sample"] - - if output_type in ["np", "pil"]: - image = image * 0.5 + 0.5 - image = image.clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - if output_type == "pil": - image = self.numpy_to_pil(image) - else: - image = latents - - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py deleted file mode 100644 index c7b8022c22b1c5211d6651b9efcf257470b6b3fa..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py +++ /dev/null @@ -1,634 +0,0 @@ -import inspect -from typing import Callable, Dict, List, Optional, Union - -import PIL -import PIL.Image -import torch -from transformers import T5EncoderModel, T5Tokenizer - -from ...image_processor import VaeImageProcessor -from ...loaders import StableDiffusionLoraLoaderMixin -from ...models import Kandinsky3UNet, VQModel -from ...schedulers import DDPMScheduler -from ...utils import ( - deprecate, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -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 - >>> from diffusers import AutoPipelineForImage2Image - >>> from diffusers.utils import load_image - >>> import torch - - >>> pipe = AutoPipelineForImage2Image.from_pretrained( - ... "kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16 - ... ) - >>> pipe.enable_model_cpu_offload() - - >>> prompt = "A painting of the inside of a subway train with tiny raccoons." - >>> image = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/kandinsky3/t2i.png" - ... ) - - >>> generator = torch.Generator(device="cpu").manual_seed(0) - >>> image = pipe(prompt, image=image, strength=0.75, num_inference_steps=25, generator=generator).images[0] - ``` -""" - - -class Kandinsky3Img2ImgPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin): - model_cpu_offload_seq = "text_encoder->movq->unet->movq" - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "negative_attention_mask", - "attention_mask", - ] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - unet: Kandinsky3UNet, - scheduler: DDPMScheduler, - movq: VQModel, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_encoder, unet=unet, scheduler=scheduler, movq=movq - ) - movq_scale_factor = 2 ** (len(self.movq.config.block_out_channels) - 1) if getattr(self, "movq", None) else 8 - movq_latent_channels = self.movq.config.latent_channels if getattr(self, "movq", None) else 4 - self.image_processor = VaeImageProcessor( - vae_scale_factor=movq_scale_factor, - vae_latent_channels=movq_latent_channels, - resample="bicubic", - reducing_gap=1, - ) - - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - timesteps = self.scheduler.timesteps[t_start:] - - return timesteps, num_inference_steps - t_start - - def _process_embeds(self, embeddings, attention_mask, cut_context): - # return embeddings, attention_mask - if cut_context: - embeddings[attention_mask == 0] = torch.zeros_like(embeddings[attention_mask == 0]) - max_seq_length = attention_mask.sum(-1).max() + 1 - embeddings = embeddings[:, :max_seq_length] - attention_mask = attention_mask[:, :max_seq_length] - return embeddings, attention_mask - - @torch.no_grad() - def encode_prompt( - self, - prompt, - do_classifier_free_guidance=True, - num_images_per_prompt=1, - device=None, - negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - _cut_context=False, - attention_mask: Optional[torch.Tensor] = None, - negative_attention_mask: Optional[torch.Tensor] = None, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`, *optional*): - prompt to be encoded - device: (`torch.device`, *optional*): - torch device to place the resulting embeddings on - num_images_per_prompt (`int`, *optional*, defaults to 1): - 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*): - 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`). - 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. - attention_mask (`torch.Tensor`, *optional*): - Pre-generated attention mask. Must provide if passing `prompt_embeds` directly. - negative_attention_mask (`torch.Tensor`, *optional*): - Pre-generated negative attention mask. Must provide if passing `negative_prompt_embeds` directly. - """ - if prompt is not None and negative_prompt is not None: - 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)} !=" - f" {type(prompt)}." - ) - - if device is None: - device = self._execution_device - - 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] - - max_length = 128 - - if prompt_embeds is None: - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids.to(device) - attention_mask = text_inputs.attention_mask.to(device) - prompt_embeds = self.text_encoder( - text_input_ids, - attention_mask=attention_mask, - ) - prompt_embeds = prompt_embeds[0] - prompt_embeds, attention_mask = self._process_embeds(prompt_embeds, attention_mask, _cut_context) - prompt_embeds = prompt_embeds * attention_mask.unsqueeze(2) - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=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) - 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] - - if negative_prompt is None: - uncond_tokens = [""] * batch_size - 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 - if negative_prompt is not None: - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=128, - truncation=True, - return_attention_mask=True, - return_tensors="pt", - ) - text_input_ids = uncond_input.input_ids.to(device) - negative_attention_mask = uncond_input.attention_mask.to(device) - - negative_prompt_embeds = self.text_encoder( - text_input_ids, - attention_mask=negative_attention_mask, - ) - negative_prompt_embeds = negative_prompt_embeds[0] - negative_prompt_embeds = negative_prompt_embeds[:, : prompt_embeds.shape[1]] - negative_attention_mask = negative_attention_mask[:, : prompt_embeds.shape[1]] - negative_prompt_embeds = negative_prompt_embeds * negative_attention_mask.unsqueeze(2) - - else: - negative_prompt_embeds = torch.zeros_like(prompt_embeds) - negative_attention_mask = torch.zeros_like(attention_mask) - - 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=dtype, device=device) - if negative_prompt_embeds.shape != prompt_embeds.shape: - 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) - negative_attention_mask = negative_attention_mask.repeat(num_images_per_prompt, 1) - - # 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 - else: - negative_prompt_embeds = None - negative_attention_mask = None - return prompt_embeds, negative_prompt_embeds, attention_mask, negative_attention_mask - - def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - 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." - ) - - elif isinstance(generator, list): - init_latents = [ - self.movq.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = self.movq.encode(image).latent_dist.sample(generator) - - init_latents = self.movq.config.scaling_factor * init_latents - - init_latents = torch.cat([init_latents], dim=0) - - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - - latents = init_latents - - return latents - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - attention_mask=None, - negative_attention_mask=None, - ): - 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}." - ) - - if negative_prompt_embeds is not None and negative_attention_mask is None: - raise ValueError("Please provide `negative_attention_mask` along with `negative_prompt_embeds`") - - if negative_prompt_embeds is not None and negative_attention_mask is not None: - if negative_prompt_embeds.shape[:2] != negative_attention_mask.shape: - raise ValueError( - "`negative_prompt_embeds` and `negative_attention_mask` must have the same batch_size and token length when passed directly, but" - f" got: `negative_prompt_embeds` {negative_prompt_embeds.shape[:2]} != `negative_attention_mask`" - f" {negative_attention_mask.shape}." - ) - - if prompt_embeds is not None and attention_mask is None: - raise ValueError("Please provide `attention_mask` along with `prompt_embeds`") - - if prompt_embeds is not None and attention_mask is not None: - if prompt_embeds.shape[:2] != attention_mask.shape: - raise ValueError( - "`prompt_embeds` and `attention_mask` must have the same batch_size and token length when passed directly, but" - f" got: `prompt_embeds` {prompt_embeds.shape[:2]} != `attention_mask`" - f" {attention_mask.shape}." - ) - - @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 - - @torch.no_grad() - @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, - 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", - return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], 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*): - 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`, or tensor representing an image batch, that will be used as the starting point for the - process. - strength (`float`, *optional*, defaults to 0.8): - 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`. - 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 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`. 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. - 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. - attention_mask (`torch.Tensor`, *optional*): - Pre-generated attention mask. Must provide if passing `prompt_embeds` directly. - negative_attention_mask (`torch.Tensor`, *optional*): - Pre-generated negative attention mask. Must provide if passing `negative_prompt_embeds` directly. - 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.IFPipelineOutput`] instead of a plain tuple. - 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. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` - - """ - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - 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]}" - ) - - cut_context = True - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - callback_on_step_end_tensor_inputs, - attention_mask, - negative_attention_mask, - ) - - self._guidance_scale = guidance_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] - - device = self._execution_device - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds, attention_mask, negative_attention_mask = self.encode_prompt( - prompt, - self.do_classifier_free_guidance, - num_images_per_prompt=num_images_per_prompt, - device=device, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - _cut_context=cut_context, - attention_mask=attention_mask, - negative_attention_mask=negative_attention_mask, - ) - - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - attention_mask = torch.cat([negative_attention_mask, attention_mask]).bool() - if not isinstance(image, list): - image = [image] - if not all(isinstance(i, (PIL.Image.Image, torch.Tensor)) for i in image): - raise ValueError( - f"Input is in incorrect format: {[type(i) for i in image]}. Currently, we only support PIL image and pytorch tensor" - ) - - image = torch.cat([self.image_processor.preprocess(i) for i in image], dim=0) - image = image.to(dtype=prompt_embeds.dtype, device=device) - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) - # 5. Prepare latents - latents = self.movq.encode(image)["latents"] - latents = latents.repeat_interleave(num_images_per_prompt, dim=0) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - latents = self.prepare_latents( - latents, latent_timestep, batch_size, num_images_per_prompt, prompt_embeds.dtype, device, generator - ) - if hasattr(self, "text_encoder_offload_hook") and self.text_encoder_offload_hook is not None: - self.text_encoder_offload_hook.offload() - - # 7. Denoising loop - 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): - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - encoder_attention_mask=attention_mask, - )[0] - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - - noise_pred = (guidance_scale + 1.0) * noise_pred_text - guidance_scale * noise_pred_uncond - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step( - noise_pred, - t, - latents, - generator=generator, - ).prev_sample - - 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) - attention_mask = callback_outputs.pop("attention_mask", attention_mask) - negative_attention_mask = callback_outputs.pop("negative_attention_mask", negative_attention_mask) - - if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): - progress_bar.update() - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # post-processing - if not output_type == "latent": - image = self.movq.decode(latents, force_not_quantize=True)["sample"] - image = self.image_processor.postprocess(image, output_type) - else: - image = latents - - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__init__.py deleted file mode 100644 index 671d22e9f433802ec82ebe89fae4cef6c001064b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__init__.py +++ /dev/null @@ -1,54 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - get_objects_from_module, - is_sentencepiece_available, - is_torch_available, - is_transformers_available, -) - - -_dummy_objects = {} -_import_structure = {} - -try: - if not (is_transformers_available() and is_torch_available()) and is_sentencepiece_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_torch_and_transformers_and_sentencepiece_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_and_sentencepiece_objects)) -else: - _import_structure["pipeline_kolors"] = ["KolorsPipeline"] - _import_structure["pipeline_kolors_img2img"] = ["KolorsImg2ImgPipeline"] - _import_structure["text_encoder"] = ["ChatGLMModel"] - _import_structure["tokenizer"] = ["ChatGLMTokenizer"] - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - try: - if not (is_transformers_available() and is_torch_available()) and is_sentencepiece_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_and_sentencepiece_objects import * - - else: - from .pipeline_kolors import KolorsPipeline - from .pipeline_kolors_img2img import KolorsImg2ImgPipeline - from .text_encoder import ChatGLMModel - from .tokenizer import ChatGLMTokenizer - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 3a4e69955b27224f62d0ad05960518d5e5bd0bda..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/pipeline_kolors.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/pipeline_kolors.cpython-310.pyc deleted file mode 100644 index 9ec5ceb817d8bec72ddc1aafa785cb5ffbe81bc0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/pipeline_kolors.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/pipeline_kolors_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/pipeline_kolors_img2img.cpython-310.pyc deleted file mode 100644 index fa11e6925d9e7b43f436439b26aa0d18b009878b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/pipeline_kolors_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index 08bda967b7b9358e5977ef292cfeda1cde626eaa..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/text_encoder.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/text_encoder.cpython-310.pyc deleted file mode 100644 index 3e7af404658add64ba9aba829ddcffabbbe2ffb2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/text_encoder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/tokenizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/tokenizer.cpython-310.pyc deleted file mode 100644 index d48adf377fc529c1b0591b89b3bc48ea62fd5d2c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/__pycache__/tokenizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/pipeline_kolors.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/pipeline_kolors.py deleted file mode 100644 index 1fa9f6ce1d43ddb1e630de73b68447c6d339847a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/pipeline_kolors.py +++ /dev/null @@ -1,1072 +0,0 @@ -# Copyright 2025 Stability AI, Kwai-Kolors 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, Tuple, Union - -import torch -from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import IPAdapterMixin, StableDiffusionLoraLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.attention_processor import AttnProcessor2_0, FusedAttnProcessor2_0, XFormersAttnProcessor -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import is_torch_xla_available, logging, replace_example_docstring -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from .pipeline_output import KolorsPipelineOutput -from .text_encoder import ChatGLMModel -from .tokenizer import ChatGLMTokenizer - - -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 KolorsPipeline - - >>> pipe = KolorsPipeline.from_pretrained( - ... "Kwai-Kolors/Kolors-diffusers", variant="fp16", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = ( - ... "A photo of a ladybug, macro, zoom, high quality, film, holding a wooden sign with the text 'KOLORS'" - ... ) - >>> image = pipe(prompt).images[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, - **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 KolorsPipeline(DiffusionPipeline, StableDiffusionMixin, StableDiffusionLoraLoaderMixin, IPAdapterMixin): - r""" - Pipeline for text-to-image generation using Kolors. - - 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.StableDiffusionLoraLoaderMixin.load_lora_weights`] for loading LoRA weights - - [`~loaders.StableDiffusionLoraLoaderMixin.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 ([`ChatGLMModel`]): - Frozen text-encoder. Kolors uses [ChatGLM3-6B](https://huggingface.co/THUDM/chatglm3-6b). - tokenizer (`ChatGLMTokenizer`): - Tokenizer of class - [ChatGLMTokenizer](https://huggingface.co/THUDM/chatglm3-6b/blob/main/tokenization_chatglm.py). - 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 `"False"`): - Whether the negative prompt embeddings shall be forced to always be set to 0. Also see the config of - `Kwai-Kolors/Kolors-diffusers`. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = [ - "image_encoder", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "negative_add_time_ids", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: ChatGLMModel, - tokenizer: ChatGLMTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - image_encoder: CLIPVisionModelWithProjection = None, - feature_extractor: CLIPImageProcessor = None, - force_zeros_for_empty_prompt: bool = False, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - ) - 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 - ) - - def encode_prompt( - self, - prompt, - device: Optional[torch.device] = 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, - max_sequence_length: int = 256, - ): - 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.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. - 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_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. - 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. - max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`. - """ - # from IPython import embed; embed(); exit() - device = device or self._execution_device - - 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] - - # Define tokenizers and text encoders - tokenizers = [self.tokenizer] - text_encoders = [self.text_encoder] - - if prompt_embeds is None: - prompt_embeds_list = [] - for tokenizer, text_encoder in zip(tokenizers, text_encoders): - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - ).to(device) - output = text_encoder( - input_ids=text_inputs["input_ids"], - attention_mask=text_inputs["attention_mask"], - position_ids=text_inputs["position_ids"], - output_hidden_states=True, - ) - - # [max_sequence_length, batch, hidden_size] -> [batch, max_sequence_length, hidden_size] - # clone to have a contiguous tensor - prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone() - # [max_sequence_length, batch, hidden_size] -> [batch, hidden_size] - pooled_prompt_embeds = output.hidden_states[-1][-1, :, :].clone() - bs_embed, seq_len, _ = prompt_embeds.shape - 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) - - prompt_embeds_list.append(prompt_embeds) - - prompt_embeds = prompt_embeds_list[0] - - # 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) - elif 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 - - negative_prompt_embeds_list = [] - - for tokenizer, text_encoder in zip(tokenizers, text_encoders): - uncond_input = tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - ).to(device) - output = text_encoder( - input_ids=uncond_input["input_ids"], - attention_mask=uncond_input["attention_mask"], - position_ids=uncond_input["position_ids"], - output_hidden_states=True, - ) - - # [max_sequence_length, batch, hidden_size] -> [batch, max_sequence_length, hidden_size] - # clone to have a contiguous tensor - negative_prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone() - # [max_sequence_length, batch, hidden_size] -> [batch, hidden_size] - negative_pooled_prompt_embeds = output.hidden_states[-1][-1, :, :].clone() - - 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=text_encoder.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 - ) - - negative_prompt_embeds_list.append(negative_prompt_embeds) - - negative_prompt_embeds = negative_prompt_embeds_list[0] - - bs_embed = pooled_prompt_embeds.shape[0] - 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 - ) - - return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - num_inference_steps, - height, - width, - negative_prompt=None, - prompt_embeds=None, - pooled_prompt_embeds=None, - negative_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - ): - if not isinstance(num_inference_steps, int) or num_inference_steps <= 0: - raise ValueError( - f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type" - f" {type(num_inference_steps)}." - ) - - 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_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 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`." - ) - - 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" - ) - - if max_sequence_length is not None and max_sequence_length > 256: - raise ValueError(f"`max_sequence_length` cannot be greater than 256 but is {max_sequence_length}") - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - 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): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - FusedAttnProcessor2_0, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 do_classifier_free_guidance(self): - return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_end(self): - return self._denoising_end - - @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, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = None, - guidance_scale: float = 5.0, - 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, - 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", - 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"], - max_sequence_length: int = 256, - ): - 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. - Anything below 512 pixels won't work well for - [Kwai-Kolors/Kolors-diffusers](https://huggingface.co/Kwai-Kolors/Kolors-diffusers) 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 - [Kwai-Kolors/Kolors-diffusers](https://huggingface.co/Kwai-Kolors/Kolors-diffusers) 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. - 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*): - 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. - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output) - 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. - 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*): - 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 ge 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. - 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_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. - 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. - 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)`. 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. - 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.kolors.KolorsPipelineOutput`] instead of a plain tuple. - cross_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). - 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. - 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 256): Maximum sequence length to use with the `prompt`. - - Examples: - - Returns: - [`~pipelines.kolors.KolorsPipelineOutput`] or `tuple`: [`~pipelines.kolors.KolorsPipelineOutput`] if - `return_dict` is True, otherwise a `tuple`. When returning a tuple, 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 - - # 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, - num_inference_steps, - height, - width, - negative_prompt, - prompt_embeds, - pooled_prompt_embeds, - negative_prompt_embeds, - negative_pooled_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - ) - - self._guidance_scale = guidance_scale - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - 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. Encode input prompt - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - ) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 5. 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, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Prepare added time ids & embeddings - add_text_embeds = pooled_prompt_embeds - text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1]) - - 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) - - 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, - self.do_classifier_free_guidance, - ) - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - # 8.1 Apply denoising_end - if ( - self.denoising_end is not None - and isinstance(self.denoising_end, float) - and self.denoising_end > 0 - and self.denoising_end < 1 - ): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - # 9. 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) - - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = image_embeds - - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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) - - # compute the previous noisy sample x_t -> x_t-1 - latents_dtype = latents.dtype - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds - ) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids) - - 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 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) - elif latents.dtype != self.vae.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 - self.vae = self.vae.to(latents.dtype) - - # unscale/denormalize the latents - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, 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 KolorsPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/pipeline_kolors_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/pipeline_kolors_img2img.py deleted file mode 100644 index e3cf4f227624bc55804f18f01b82bbc22d9dabb3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/pipeline_kolors_img2img.py +++ /dev/null @@ -1,1252 +0,0 @@ -# Copyright 2025 Stability AI, Kwai-Kolors 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, Tuple, Union - -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import IPAdapterMixin, StableDiffusionXLLoraLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.attention_processor import AttnProcessor2_0, FusedAttnProcessor2_0, XFormersAttnProcessor -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import is_torch_xla_available, logging, replace_example_docstring -from ...utils.torch_utils import empty_device_cache, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from .pipeline_output import KolorsPipelineOutput -from .text_encoder import ChatGLMModel -from .tokenizer import ChatGLMTokenizer - - -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 KolorsImg2ImgPipeline - >>> from diffusers.utils import load_image - - >>> pipe = KolorsImg2ImgPipeline.from_pretrained( - ... "Kwai-Kolors/Kolors-diffusers", variant="fp16", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - >>> url = ( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/kolors/bunny_source.png" - ... ) - - - >>> init_image = load_image(url) - >>> prompt = "high quality image of a capybara wearing sunglasses. In the background of the image there are trees, poles, grass and other objects. At the bottom of the object there is the road., 8k, highly detailed." - >>> image = pipe(prompt, image=init_image).images[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" -): - 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 KolorsImg2ImgPipeline(DiffusionPipeline, StableDiffusionMixin, StableDiffusionXLLoraLoaderMixin, IPAdapterMixin): - r""" - Pipeline for text-to-image generation using Kolors. - - 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.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 ([`ChatGLMModel`]): - Frozen text-encoder. Kolors uses [ChatGLM3-6B](https://huggingface.co/THUDM/chatglm3-6b). - tokenizer (`ChatGLMTokenizer`): - Tokenizer of class - [ChatGLMTokenizer](https://huggingface.co/THUDM/chatglm3-6b/blob/main/tokenization_chatglm.py). - 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 `"False"`): - Whether the negative prompt embeddings shall be forced to always be set to 0. Also see the config of - `Kwai-Kolors/Kolors-diffusers`. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder-unet->vae" - _optional_components = [ - "image_encoder", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "negative_add_time_ids", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: ChatGLMModel, - tokenizer: ChatGLMTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - image_encoder: CLIPVisionModelWithProjection = None, - feature_extractor: CLIPImageProcessor = None, - force_zeros_for_empty_prompt: bool = False, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - ) - 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 - ) - - # Copied from diffusers.pipelines.kolors.pipeline_kolors.KolorsPipeline.encode_prompt - def encode_prompt( - self, - prompt, - device: Optional[torch.device] = 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, - max_sequence_length: int = 256, - ): - 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.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. - 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_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. - 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. - max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`. - """ - # from IPython import embed; embed(); exit() - device = device or self._execution_device - - 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] - - # Define tokenizers and text encoders - tokenizers = [self.tokenizer] - text_encoders = [self.text_encoder] - - if prompt_embeds is None: - prompt_embeds_list = [] - for tokenizer, text_encoder in zip(tokenizers, text_encoders): - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - ).to(device) - output = text_encoder( - input_ids=text_inputs["input_ids"], - attention_mask=text_inputs["attention_mask"], - position_ids=text_inputs["position_ids"], - output_hidden_states=True, - ) - - # [max_sequence_length, batch, hidden_size] -> [batch, max_sequence_length, hidden_size] - # clone to have a contiguous tensor - prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone() - # [max_sequence_length, batch, hidden_size] -> [batch, hidden_size] - pooled_prompt_embeds = output.hidden_states[-1][-1, :, :].clone() - bs_embed, seq_len, _ = prompt_embeds.shape - 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) - - prompt_embeds_list.append(prompt_embeds) - - prompt_embeds = prompt_embeds_list[0] - - # 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) - elif 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 - - negative_prompt_embeds_list = [] - - for tokenizer, text_encoder in zip(tokenizers, text_encoders): - uncond_input = tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - ).to(device) - output = text_encoder( - input_ids=uncond_input["input_ids"], - attention_mask=uncond_input["attention_mask"], - position_ids=uncond_input["position_ids"], - output_hidden_states=True, - ) - - # [max_sequence_length, batch, hidden_size] -> [batch, max_sequence_length, hidden_size] - # clone to have a contiguous tensor - negative_prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone() - # [max_sequence_length, batch, hidden_size] -> [batch, hidden_size] - negative_pooled_prompt_embeds = output.hidden_states[-1][-1, :, :].clone() - - 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=text_encoder.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 - ) - - negative_prompt_embeds_list.append(negative_prompt_embeds) - - negative_prompt_embeds = negative_prompt_embeds_list[0] - - bs_embed = pooled_prompt_embeds.shape[0] - 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 - ) - - return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - strength, - num_inference_steps, - height, - width, - negative_prompt=None, - prompt_embeds=None, - pooled_prompt_embeds=None, - negative_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - ): - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - - if not isinstance(num_inference_steps, int) or num_inference_steps <= 0: - raise ValueError( - f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type" - f" {type(num_inference_steps)}." - ) - - 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_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 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`." - ) - - 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" - ) - - if max_sequence_length is not None and max_sequence_length > 256: - raise ValueError(f"`max_sequence_length` cannot be greater than 256 but is {max_sequence_length}") - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device, denoising_start=None): - # get the original timestep using init_timestep - if denoising_start is None: - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - t_start = 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 - - else: - # Strength is irrelevant if we directly request a timestep to start at; - # that is, strength is determined by the denoising_start instead. - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (denoising_start * self.scheduler.config.num_train_timesteps) - ) - ) - - num_inference_steps = (self.scheduler.timesteps < discrete_timestep_cutoff).sum().item() - if self.scheduler.order == 2 and num_inference_steps % 2 == 0: - # if the scheduler is a 2nd order scheduler we might have to do +1 - # because `num_inference_steps` might be even given that every timestep - # (except the highest one) is duplicated. If `num_inference_steps` is even it would - # mean that we cut the timesteps in the middle of the denoising step - # (between 1st and 2nd derivative) which leads to incorrect results. By adding 1 - # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler - num_inference_steps = num_inference_steps + 1 - - # because t_n+1 >= t_n, we slice the timesteps starting from the end - t_start = len(self.scheduler.timesteps) - num_inference_steps - timesteps = self.scheduler.timesteps[t_start:] - if hasattr(self.scheduler, "set_begin_index"): - self.scheduler.set_begin_index(t_start) - return timesteps, num_inference_steps - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.prepare_latents - def prepare_latents( - self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None, add_noise=True - ): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - latents_mean = latents_std = None - if hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None: - latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1) - if hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None: - latents_std = torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1) - - # Offload text encoder if `enable_model_cpu_offload` was enabled - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.text_encoder_2.to("cpu") - empty_device_cache() - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - # make sure the VAE is in float32 mode, as it overflows in float16 - if self.vae.config.force_upcast: - image = image.float() - self.vae.to(dtype=torch.float32) - - 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." - ) - - elif isinstance(generator, list): - if image.shape[0] < batch_size and batch_size % image.shape[0] == 0: - image = torch.cat([image] * (batch_size // image.shape[0]), dim=0) - elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} " - ) - - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - if self.vae.config.force_upcast: - self.vae.to(dtype) - - init_latents = init_latents.to(dtype) - if latents_mean is not None and latents_std is not None: - latents_mean = latents_mean.to(device=device, dtype=dtype) - latents_std = latents_std.to(device=device, dtype=dtype) - init_latents = (init_latents - latents_mean) * self.vae.config.scaling_factor / latents_std - else: - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - if add_noise: - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - - latents = init_latents - - 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): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - FusedAttnProcessor2_0, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 do_classifier_free_guidance(self): - return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_start(self): - return self._denoising_start - - @property - def denoising_end(self): - return self._denoising_end - - @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, - strength: float = 0.3, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_start: Optional[float] = None, - denoising_end: Optional[float] = None, - guidance_scale: float = 5.0, - 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, - 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", - 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"], - max_sequence_length: int = 256, - ): - 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 (`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` - 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`. Note that in the case of - `denoising_start` being declared as an integer, the value of `strength` will be 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. - Anything below 512 pixels won't work well for - [Kwai-Kolors/Kolors-diffusers](https://huggingface.co/Kwai-Kolors/Kolors-diffusers) 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 - [Kwai-Kolors/Kolors-diffusers](https://huggingface.co/Kwai-Kolors/Kolors-diffusers) 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. - 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*): - 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. - denoising_start (`float`, *optional*): - When specified, indicates the fraction (between 0.0 and 1.0) of the total denoising process to be - bypassed before it is initiated. Consequently, the initial part of the denoising process is skipped and - it is assumed that the passed `image` is a partly denoised image. Note that when this is specified, - strength will be ignored. The `denoising_start` parameter is particularly beneficial when this pipeline - is integrated into a "Mixture of Denoisers" multi-pipeline setup, as detailed in [**Refine Image - Quality**](https://huggingface.co/docs/diffusers/using-diffusers/sdxl#refine-image-quality). - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output) - 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. - 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*): - 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 ge 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. - 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_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. - 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. - 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)`. 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. - 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.kolors.KolorsPipelineOutput`] instead of a plain tuple. - cross_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). - 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. - 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 256): Maximum sequence length to use with the `prompt`. - - Examples: - - Returns: - [`~pipelines.kolors.KolorsPipelineOutput`] or `tuple`: [`~pipelines.kolors.KolorsPipelineOutput`] if - `return_dict` is True, otherwise a `tuple`. When returning a tuple, 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 - - # 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, - strength, - num_inference_steps, - height, - width, - negative_prompt, - prompt_embeds, - pooled_prompt_embeds, - negative_prompt_embeds, - negative_pooled_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - ) - - self._guidance_scale = guidance_scale - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - self._denoising_start = denoising_start - 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. Encode input prompt - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - - # 4. Preprocess image - image = self.image_processor.preprocess(image) - - # 5. Prepare timesteps - def denoising_value_valid(dnv): - return isinstance(dnv, float) and 0 < dnv < 1 - - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - timesteps, num_inference_steps = self.get_timesteps( - num_inference_steps, - strength, - device, - denoising_start=self.denoising_start if denoising_value_valid(self.denoising_start) else None, - ) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - - add_noise = True if self.denoising_start is None else False - - # 6. Prepare latent variables - if latents is None: - latents = self.prepare_latents( - image, - latent_timestep, - batch_size, - num_images_per_prompt, - prompt_embeds.dtype, - device, - generator, - add_noise, - ) - - # 7. Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - height, width = latents.shape[-2:] - height = height * self.vae_scale_factor - width = width * self.vae_scale_factor - - original_size = original_size or (height, width) - target_size = target_size or (height, width) - - # 8. Prepare added time ids & embeddings - add_text_embeds = pooled_prompt_embeds - text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1]) - - 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) - - 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, - self.do_classifier_free_guidance, - ) - - # 9. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - # 9.1 Apply denoising_end - if ( - self.denoising_end is not None - and self.denoising_start is not None - and denoising_value_valid(self.denoising_end) - and denoising_value_valid(self.denoising_start) - and self.denoising_start >= self.denoising_end - ): - raise ValueError( - f"`denoising_start`: {self.denoising_start} cannot be larger than or equal to `denoising_end`: " - + f" {self.denoising_end} when using type float." - ) - elif self.denoising_end is not None and denoising_value_valid(self.denoising_end): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - # 9.2 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) - - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = image_embeds - - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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) - - # compute the previous noisy sample x_t -> x_t-1 - latents_dtype = latents.dtype - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds - ) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids) - - 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 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) - elif latents.dtype != self.vae.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 - self.vae = self.vae.to(latents.dtype) - - # unscale/denormalize the latents - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, 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 KolorsPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/pipeline_output.py deleted file mode 100644 index 310ee7e8a89b253c55dfb47e42964310c93eb56a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/pipeline_output.py +++ /dev/null @@ -1,21 +0,0 @@ -from dataclasses import dataclass -from typing import List, Union - -import numpy as np -import PIL.Image - -from ...utils import BaseOutput - - -@dataclass -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, - num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/text_encoder.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/text_encoder.py deleted file mode 100644 index 6fd17156a1162ce370a0af89819e72e9607af7e7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/text_encoder.py +++ /dev/null @@ -1,862 +0,0 @@ -# Copyright 2025 ChatGLM3-6B Model Team, Kwai-Kolors 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 typing import List, Optional, Tuple - -import torch -import torch.nn.functional as F -from torch import nn -from torch.nn import LayerNorm -from torch.nn.utils import skip_init -from transformers import PretrainedConfig, PreTrainedModel -from transformers.modeling_outputs import BaseModelOutputWithPast - -from ...utils import logging - - -logger = logging.get_logger(__name__) - - -class ChatGLMConfig(PretrainedConfig): - model_type = "chatglm" - - def __init__( - self, - num_layers=28, - padded_vocab_size=65024, - hidden_size=4096, - ffn_hidden_size=13696, - kv_channels=128, - num_attention_heads=32, - seq_length=2048, - hidden_dropout=0.0, - classifier_dropout=None, - attention_dropout=0.0, - layernorm_epsilon=1e-5, - rmsnorm=True, - apply_residual_connection_post_layernorm=False, - post_layer_norm=True, - add_bias_linear=False, - add_qkv_bias=False, - bias_dropout_fusion=True, - multi_query_attention=False, - multi_query_group_num=1, - apply_query_key_layer_scaling=True, - attention_softmax_in_fp32=True, - fp32_residual_connection=False, - quantization_bit=0, - pre_seq_len=None, - prefix_projection=False, - **kwargs, - ): - self.num_layers = num_layers - self.vocab_size = padded_vocab_size - self.padded_vocab_size = padded_vocab_size - self.hidden_size = hidden_size - self.ffn_hidden_size = ffn_hidden_size - self.kv_channels = kv_channels - self.num_attention_heads = num_attention_heads - self.seq_length = seq_length - self.hidden_dropout = hidden_dropout - self.classifier_dropout = classifier_dropout - self.attention_dropout = attention_dropout - self.layernorm_epsilon = layernorm_epsilon - self.rmsnorm = rmsnorm - self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm - self.post_layer_norm = post_layer_norm - self.add_bias_linear = add_bias_linear - self.add_qkv_bias = add_qkv_bias - self.bias_dropout_fusion = bias_dropout_fusion - self.multi_query_attention = multi_query_attention - self.multi_query_group_num = multi_query_group_num - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.attention_softmax_in_fp32 = attention_softmax_in_fp32 - self.fp32_residual_connection = fp32_residual_connection - self.quantization_bit = quantization_bit - self.pre_seq_len = pre_seq_len - self.prefix_projection = prefix_projection - super().__init__(**kwargs) - - -class RMSNorm(torch.nn.Module): - def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs): - super().__init__() - self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype)) - self.eps = eps - - def forward(self, hidden_states: torch.Tensor): - input_dtype = hidden_states.dtype - variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.eps) - - return (self.weight * hidden_states).to(input_dtype) - - -class CoreAttention(torch.nn.Module): - def __init__(self, config: ChatGLMConfig, layer_number): - super(CoreAttention, self).__init__() - - self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling - self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32 - if self.apply_query_key_layer_scaling: - self.attention_softmax_in_fp32 = True - self.layer_number = max(1, layer_number) - - projection_size = config.kv_channels * config.num_attention_heads - - # Per attention head and per partition values. - self.hidden_size_per_partition = projection_size - self.hidden_size_per_attention_head = projection_size // config.num_attention_heads - self.num_attention_heads_per_partition = config.num_attention_heads - - coeff = None - self.norm_factor = math.sqrt(self.hidden_size_per_attention_head) - if self.apply_query_key_layer_scaling: - coeff = self.layer_number - self.norm_factor *= coeff - self.coeff = coeff - - self.attention_dropout = torch.nn.Dropout(config.attention_dropout) - - def forward(self, query_layer, key_layer, value_layer, attention_mask): - pytorch_major_version = int(torch.__version__.split(".")[0]) - if pytorch_major_version >= 2: - query_layer, key_layer, value_layer = [ - k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer] - ] - if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]: - context_layer = torch.nn.functional.scaled_dot_product_attention( - query_layer, key_layer, value_layer, is_causal=True - ) - else: - if attention_mask is not None: - attention_mask = ~attention_mask - context_layer = torch.nn.functional.scaled_dot_product_attention( - query_layer, key_layer, value_layer, attention_mask - ) - context_layer = context_layer.permute(2, 0, 1, 3) - new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,) - context_layer = context_layer.reshape(*new_context_layer_shape) - else: - # Raw attention scores - - # [b, np, sq, sk] - output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0)) - - # [sq, b, np, hn] -> [sq, b * np, hn] - query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1) - # [sk, b, np, hn] -> [sk, b * np, hn] - key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1) - - # preallocting input tensor: [b * np, sq, sk] - matmul_input_buffer = torch.empty( - output_size[0] * output_size[1], - output_size[2], - output_size[3], - dtype=query_layer.dtype, - device=query_layer.device, - ) - - # Raw attention scores. [b * np, sq, sk] - matmul_result = torch.baddbmm( - matmul_input_buffer, - query_layer.transpose(0, 1), # [b * np, sq, hn] - key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk] - beta=0.0, - alpha=(1.0 / self.norm_factor), - ) - - # change view to [b, np, sq, sk] - attention_scores = matmul_result.view(*output_size) - - # =========================== - # Attention probs and dropout - # =========================== - - # attention scores and attention mask [b, np, sq, sk] - if self.attention_softmax_in_fp32: - attention_scores = attention_scores.float() - if self.coeff is not None: - attention_scores = attention_scores * self.coeff - if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]: - attention_mask = torch.ones( - output_size[0], 1, output_size[2], output_size[3], device=attention_scores.device, dtype=torch.bool - ) - attention_mask.tril_() - attention_mask = ~attention_mask - if attention_mask is not None: - attention_scores = attention_scores.masked_fill(attention_mask, float("-inf")) - attention_probs = F.softmax(attention_scores, dim=-1) - attention_probs = attention_probs.type_as(value_layer) - - # This is actually dropping out entire tokens to attend to, which might - # seem a bit unusual, but is taken from the original Transformer paper. - attention_probs = self.attention_dropout(attention_probs) - # ========================= - # Context layer. [sq, b, hp] - # ========================= - - # value_layer -> context layer. - # [sk, b, np, hn] --> [b, np, sq, hn] - - # context layer shape: [b, np, sq, hn] - output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3)) - # change view [sk, b * np, hn] - value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1) - # change view [b * np, sq, sk] - attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1) - # matmul: [b * np, sq, hn] - context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1)) - # change view [b, np, sq, hn] - context_layer = context_layer.view(*output_size) - # [b, np, sq, hn] --> [sq, b, np, hn] - context_layer = context_layer.permute(2, 0, 1, 3).contiguous() - # [sq, b, np, hn] --> [sq, b, hp] - new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,) - context_layer = context_layer.view(*new_context_layer_shape) - - return context_layer - - -def split_tensor_along_last_dim( - tensor: torch.Tensor, - num_partitions: int, - contiguous_split_chunks: bool = False, -) -> List[torch.Tensor]: - """Split a tensor along its last dimension. - - Arguments: - tensor: input tensor. - num_partitions: number of partitions to split the tensor - contiguous_split_chunks: If True, make each chunk contiguous - in memory. - - Returns: - A list of Tensors - """ - # Get the size and dimension. - last_dim = tensor.dim() - 1 - last_dim_size = tensor.size()[last_dim] // num_partitions - # Split. - tensor_list = torch.split(tensor, last_dim_size, dim=last_dim) - # Note: torch.split does not create contiguous tensors by default. - if contiguous_split_chunks: - return tuple(chunk.contiguous() for chunk in tensor_list) - - return tensor_list - - -@torch.jit.script -def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor: - # x: [sq, b, np, hn] - sq, _b, np, _hn = x.size(0), x.size(1), x.size(2), x.size(3) - rot_dim = rope_cache.shape[-2] * 2 - x, x_pass = x[..., :rot_dim], x[..., rot_dim:] - # truncate to support variable sizes - rope_cache = rope_cache[:sq] - xshaped = x.reshape(sq, -1, np, rot_dim // 2, 2) - rope_cache = rope_cache.view(sq, -1, 1, xshaped.size(3), 2) - x_out2 = torch.stack( - [ - xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1], - xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1], - ], - -1, - ) - x_out2 = x_out2.flatten(3) - return torch.cat((x_out2, x_pass), dim=-1) - - -class SelfAttention(torch.nn.Module): - """Parallel self-attention layer abstract class. - - Self-attention layer takes input with size [s, b, h] and returns output of the same size. - """ - - def __init__(self, config: ChatGLMConfig, layer_number, device=None): - super(SelfAttention, self).__init__() - self.layer_number = max(1, layer_number) - - self.projection_size = config.kv_channels * config.num_attention_heads - - # Per attention head and per partition values. - self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads - self.num_attention_heads_per_partition = config.num_attention_heads - - self.multi_query_attention = config.multi_query_attention - self.qkv_hidden_size = 3 * self.projection_size - if self.multi_query_attention: - self.num_multi_query_groups_per_partition = config.multi_query_group_num - self.qkv_hidden_size = ( - self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num - ) - self.query_key_value = nn.Linear( - config.hidden_size, - self.qkv_hidden_size, - bias=config.add_bias_linear or config.add_qkv_bias, - device=device, - ) - - self.core_attention = CoreAttention(config, self.layer_number) - - # Output. - self.dense = nn.Linear( - self.projection_size, - config.hidden_size, - bias=config.add_bias_linear, - device=device, - ) - - def _allocate_memory(self, inference_max_sequence_len, batch_size, device=None, dtype=None): - if self.multi_query_attention: - num_attention_heads = self.num_multi_query_groups_per_partition - else: - num_attention_heads = self.num_attention_heads_per_partition - return torch.empty( - inference_max_sequence_len, - batch_size, - num_attention_heads, - self.hidden_size_per_attention_head, - dtype=dtype, - device=device, - ) - - def forward(self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True): - # hidden_states: [sq, b, h] - - # ================================================= - # Pre-allocate memory for key-values for inference. - # ================================================= - # ===================== - # Query, Key, and Value - # ===================== - - # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)] - mixed_x_layer = self.query_key_value(hidden_states) - - if self.multi_query_attention: - (query_layer, key_layer, value_layer) = mixed_x_layer.split( - [ - self.num_attention_heads_per_partition * self.hidden_size_per_attention_head, - self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head, - self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head, - ], - dim=-1, - ) - query_layer = query_layer.view( - query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) - ) - key_layer = key_layer.view( - key_layer.size()[:-1] - + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head) - ) - value_layer = value_layer.view( - value_layer.size()[:-1] - + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head) - ) - else: - new_tensor_shape = mixed_x_layer.size()[:-1] + ( - self.num_attention_heads_per_partition, - 3 * self.hidden_size_per_attention_head, - ) - mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) - - # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn] - (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3) - - # apply relative positional encoding (rotary embedding) - if rotary_pos_emb is not None: - query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb) - key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb) - - # adjust key and value for inference - if kv_cache is not None: - cache_k, cache_v = kv_cache - key_layer = torch.cat((cache_k, key_layer), dim=0) - value_layer = torch.cat((cache_v, value_layer), dim=0) - if use_cache: - kv_cache = (key_layer, value_layer) - else: - kv_cache = None - - if self.multi_query_attention: - key_layer = key_layer.unsqueeze(-2) - key_layer = key_layer.expand( - -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1 - ) - key_layer = key_layer.contiguous().view( - key_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) - ) - value_layer = value_layer.unsqueeze(-2) - value_layer = value_layer.expand( - -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1 - ) - value_layer = value_layer.contiguous().view( - value_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) - ) - - # ================================== - # core attention computation - # ================================== - - context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask) - - # ================= - # Output. [sq, b, h] - # ================= - - output = self.dense(context_layer) - - return output, kv_cache - - -class MLP(torch.nn.Module): - """MLP. - - MLP will take the input with h hidden state, project it to 4*h hidden dimension, perform nonlinear transformation, - and project the state back into h hidden dimension. - """ - - def __init__(self, config: ChatGLMConfig, device=None): - super(MLP, self).__init__() - - self.add_bias = config.add_bias_linear - - # Project to 4h. If using swiglu double the output width, see https://huggingface.co/papers/2002.05202 - self.dense_h_to_4h = nn.Linear( - config.hidden_size, - config.ffn_hidden_size * 2, - bias=self.add_bias, - device=device, - ) - - def swiglu(x): - x = torch.chunk(x, 2, dim=-1) - return F.silu(x[0]) * x[1] - - self.activation_func = swiglu - - # Project back to h. - self.dense_4h_to_h = nn.Linear(config.ffn_hidden_size, config.hidden_size, bias=self.add_bias, device=device) - - def forward(self, hidden_states): - # [s, b, 4hp] - intermediate_parallel = self.dense_h_to_4h(hidden_states) - intermediate_parallel = self.activation_func(intermediate_parallel) - # [s, b, h] - output = self.dense_4h_to_h(intermediate_parallel) - return output - - -class GLMBlock(torch.nn.Module): - """A single transformer layer. - - Transformer layer takes input with size [s, b, h] and returns an output of the same size. - """ - - def __init__(self, config: ChatGLMConfig, layer_number, device=None): - super(GLMBlock, self).__init__() - self.layer_number = layer_number - - self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm - - self.fp32_residual_connection = config.fp32_residual_connection - - LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm - # Layernorm on the input data. - self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device) - - # Self attention. - self.self_attention = SelfAttention(config, layer_number, device=device) - self.hidden_dropout = config.hidden_dropout - - # Layernorm on the attention output - self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device) - - # MLP - self.mlp = MLP(config, device=device) - - def forward( - self, - hidden_states, - attention_mask, - rotary_pos_emb, - kv_cache=None, - use_cache=True, - ): - # hidden_states: [s, b, h] - - # Layer norm at the beginning of the transformer layer. - layernorm_output = self.input_layernorm(hidden_states) - # Self attention. - attention_output, kv_cache = self.self_attention( - layernorm_output, attention_mask, rotary_pos_emb, kv_cache=kv_cache, use_cache=use_cache - ) - - # Residual connection. - if self.apply_residual_connection_post_layernorm: - residual = layernorm_output - else: - residual = hidden_states - - layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training) - layernorm_input = residual + layernorm_input - - # Layer norm post the self attention. - layernorm_output = self.post_attention_layernorm(layernorm_input) - - # MLP. - mlp_output = self.mlp(layernorm_output) - - # Second residual connection. - if self.apply_residual_connection_post_layernorm: - residual = layernorm_output - else: - residual = layernorm_input - - output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training) - output = residual + output - - return output, kv_cache - - -class GLMTransformer(torch.nn.Module): - """Transformer class.""" - - def __init__(self, config: ChatGLMConfig, device=None): - super(GLMTransformer, self).__init__() - - self.fp32_residual_connection = config.fp32_residual_connection - self.post_layer_norm = config.post_layer_norm - - # Number of layers. - self.num_layers = config.num_layers - - # Transformer layers. - def build_layer(layer_number): - return GLMBlock(config, layer_number, device=device) - - self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)]) - - if self.post_layer_norm: - LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm - # Final layer norm before output. - self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device) - - self.gradient_checkpointing = False - - def _get_layer(self, layer_number): - return self.layers[layer_number] - - def forward( - self, - hidden_states, - attention_mask, - rotary_pos_emb, - kv_caches=None, - use_cache: Optional[bool] = True, - output_hidden_states: Optional[bool] = False, - ): - if not kv_caches: - kv_caches = [None for _ in range(self.num_layers)] - presents = () if use_cache else None - if torch.is_grad_enabled() and self.gradient_checkpointing: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - - all_self_attentions = None - all_hidden_states = () if output_hidden_states else None - for index in range(self.num_layers): - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - layer = self._get_layer(index) - if torch.is_grad_enabled() and self.gradient_checkpointing: - layer_ret = self._gradient_checkpointing_func( - layer, hidden_states, attention_mask, rotary_pos_emb, kv_caches[index], use_cache - ) - else: - layer_ret = layer( - hidden_states, attention_mask, rotary_pos_emb, kv_cache=kv_caches[index], use_cache=use_cache - ) - hidden_states, kv_cache = layer_ret - if use_cache: - presents = presents + (kv_cache,) - - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - # Final layer norm. - if self.post_layer_norm: - hidden_states = self.final_layernorm(hidden_states) - - return hidden_states, presents, all_hidden_states, all_self_attentions - - -class ChatGLMPreTrainedModel(PreTrainedModel): - """ - An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained - models. - """ - - is_parallelizable = False - supports_gradient_checkpointing = True - config_class = ChatGLMConfig - base_model_prefix = "transformer" - _no_split_modules = ["GLMBlock"] - - def _init_weights(self, module: nn.Module): - """Initialize the weights.""" - return - - def get_masks(self, input_ids, past_key_values, padding_mask=None): - batch_size, seq_length = input_ids.shape - full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device) - full_attention_mask.tril_() - past_length = 0 - if past_key_values: - past_length = past_key_values[0][0].shape[0] - if past_length: - full_attention_mask = torch.cat( - (torch.ones(batch_size, seq_length, past_length, device=input_ids.device), full_attention_mask), dim=-1 - ) - if padding_mask is not None: - full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1) - if not past_length and padding_mask is not None: - full_attention_mask -= padding_mask.unsqueeze(-1) - 1 - full_attention_mask = (full_attention_mask < 0.5).bool() - full_attention_mask.unsqueeze_(1) - return full_attention_mask - - def get_position_ids(self, input_ids, device): - batch_size, seq_length = input_ids.shape - position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1) - return position_ids - - -def default_init(cls, *args, **kwargs): - return cls(*args, **kwargs) - - -class Embedding(torch.nn.Module): - """Language model embeddings.""" - - def __init__(self, config: ChatGLMConfig, device=None): - super(Embedding, self).__init__() - - self.hidden_size = config.hidden_size - # Word embeddings (parallel). - self.word_embeddings = nn.Embedding(config.padded_vocab_size, self.hidden_size, device=device) - self.fp32_residual_connection = config.fp32_residual_connection - - def forward(self, input_ids): - # Embeddings. - words_embeddings = self.word_embeddings(input_ids) - embeddings = words_embeddings - # Data format change to avoid explicit transposes : [b s h] --> [s b h]. - embeddings = embeddings.transpose(0, 1).contiguous() - # If the input flag for fp32 residual connection is set, convert for float. - if self.fp32_residual_connection: - embeddings = embeddings.float() - return embeddings - - -class RotaryEmbedding(nn.Module): - def __init__(self, dim, original_impl=False, device=None, dtype=None): - super().__init__() - inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim)) - self.register_buffer("inv_freq", inv_freq) - self.dim = dim - self.original_impl = original_impl - - def forward_impl(self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000): - """Enhanced Transformer with Rotary Position Embedding. - - Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/ - transformers/rope/__init__.py. MIT License: - https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license. - """ - # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$ - theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=torch.float, device=device) / n_elem)) - - # Create position indexes `[0, 1, ..., seq_len - 1]` - seq_idx = torch.arange(seq_len, dtype=torch.float, device=device) - - # Calculate the product of position index and $\theta_i$ - idx_theta = torch.outer(seq_idx, theta).float() - - cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1) - - # this is to mimic the behaviour of complex32, else we will get different results - if dtype in (torch.float16, torch.bfloat16, torch.int8): - cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half() - return cache - - def forward(self, max_seq_len, offset=0): - return self.forward_impl(max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device) - - -class PrefixEncoder(torch.nn.Module): - """ - The torch.nn model to encode the prefix Input shape: (batch-size, prefix-length) Output shape: (batch-size, - prefix-length, 2*layers*hidden) - """ - - def __init__(self, config: ChatGLMConfig): - super().__init__() - self.prefix_projection = config.prefix_projection - if self.prefix_projection: - # Use a two-layer MLP to encode the prefix - kv_size = config.num_layers * config.kv_channels * config.multi_query_group_num * 2 - self.embedding = torch.nn.Embedding(config.pre_seq_len, kv_size) - self.trans = torch.nn.Sequential( - torch.nn.Linear(kv_size, config.hidden_size), - torch.nn.Tanh(), - torch.nn.Linear(config.hidden_size, kv_size), - ) - else: - self.embedding = torch.nn.Embedding( - config.pre_seq_len, config.num_layers * config.kv_channels * config.multi_query_group_num * 2 - ) - - def forward(self, prefix: torch.Tensor): - if self.prefix_projection: - prefix_tokens = self.embedding(prefix) - past_key_values = self.trans(prefix_tokens) - else: - past_key_values = self.embedding(prefix) - return past_key_values - - -class ChatGLMModel(ChatGLMPreTrainedModel): - def __init__(self, config: ChatGLMConfig, device=None, empty_init=True): - super().__init__(config) - if empty_init: - init_method = skip_init - else: - init_method = default_init - init_kwargs = {} - if device is not None: - init_kwargs["device"] = device - self.embedding = init_method(Embedding, config, **init_kwargs) - self.num_layers = config.num_layers - self.multi_query_group_num = config.multi_query_group_num - self.kv_channels = config.kv_channels - - # Rotary positional embeddings - self.seq_length = config.seq_length - rotary_dim = ( - config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels - ) - - self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device) - self.encoder = init_method(GLMTransformer, config, **init_kwargs) - self.output_layer = init_method( - nn.Linear, - config.hidden_size, - config.padded_vocab_size, - bias=False, - **init_kwargs, - ) - self.pre_seq_len = config.pre_seq_len - self.prefix_projection = config.prefix_projection - if self.pre_seq_len is not None: - for param in self.parameters(): - param.requires_grad = False - self.prefix_tokens = torch.arange(self.pre_seq_len).long() - self.prefix_encoder = PrefixEncoder(config) - self.dropout = torch.nn.Dropout(0.1) - - def get_input_embeddings(self): - return self.embedding.word_embeddings - - def get_prompt(self, batch_size, device, dtype=torch.half): - prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device) - past_key_values = self.prefix_encoder(prefix_tokens).type(dtype) - past_key_values = past_key_values.view( - batch_size, self.pre_seq_len, self.num_layers * 2, self.multi_query_group_num, self.kv_channels - ) - # seq_len, b, nh, hidden_size - past_key_values = self.dropout(past_key_values) - past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2) - return past_key_values - - 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, - ): - 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 - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - batch_size, seq_length = input_ids.shape - - if inputs_embeds is None: - inputs_embeds = self.embedding(input_ids) - - if self.pre_seq_len is not None: - if past_key_values is None: - past_key_values = self.get_prompt( - batch_size=batch_size, device=input_ids.device, dtype=inputs_embeds.dtype - ) - if attention_mask is not None: - attention_mask = torch.cat( - [attention_mask.new_ones((batch_size, self.pre_seq_len)), attention_mask], dim=-1 - ) - - if full_attention_mask is None: - if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1): - full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask) - - # Rotary positional embeddings - rotary_pos_emb = self.rotary_pos_emb(self.seq_length) - if position_ids is not None: - rotary_pos_emb = rotary_pos_emb[position_ids] - else: - rotary_pos_emb = rotary_pos_emb[None, :seq_length] - rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous() - - # Run encoder. - hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder( - inputs_embeds, - full_attention_mask, - rotary_pos_emb=rotary_pos_emb, - kv_caches=past_key_values, - use_cache=use_cache, - output_hidden_states=output_hidden_states, - ) - - if not return_dict: - return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None) - - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=presents, - hidden_states=all_hidden_states, - attentions=all_self_attentions, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/tokenizer.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/tokenizer.py deleted file mode 100644 index b824ba12e079c26a5fe5c8ec234e52c89d3bbc97..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/kolors/tokenizer.py +++ /dev/null @@ -1,338 +0,0 @@ -# Copyright 2025 ChatGLM3-6B Model Team, Kwai-Kolors 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 json -import os -import re -from typing import Dict, List, Optional, Union - -from sentencepiece import SentencePieceProcessor -from transformers import PreTrainedTokenizer -from transformers.tokenization_utils_base import BatchEncoding, EncodedInput -from transformers.utils import PaddingStrategy - - -class SPTokenizer: - def __init__(self, model_path: str): - # reload tokenizer - assert os.path.isfile(model_path), model_path - self.sp_model = SentencePieceProcessor(model_file=model_path) - - # BOS / EOS token IDs - self.n_words: int = self.sp_model.vocab_size() - self.bos_id: int = self.sp_model.bos_id() - self.eos_id: int = self.sp_model.eos_id() - self.pad_id: int = self.sp_model.unk_id() - assert self.sp_model.vocab_size() == self.sp_model.get_piece_size() - - role_special_tokens = ["<|system|>", "<|user|>", "<|assistant|>", "<|observation|>"] - special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "sop", "eop"] + role_special_tokens - self.special_tokens = {} - self.index_special_tokens = {} - for token in special_tokens: - self.special_tokens[token] = self.n_words - self.index_special_tokens[self.n_words] = token - self.n_words += 1 - self.role_special_token_expression = "|".join([re.escape(token) for token in role_special_tokens]) - - def tokenize(self, s: str, encode_special_tokens=False): - if encode_special_tokens: - last_index = 0 - t = [] - for match in re.finditer(self.role_special_token_expression, s): - if last_index < match.start(): - t.extend(self.sp_model.EncodeAsPieces(s[last_index : match.start()])) - t.append(s[match.start() : match.end()]) - last_index = match.end() - if last_index < len(s): - t.extend(self.sp_model.EncodeAsPieces(s[last_index:])) - return t - else: - return self.sp_model.EncodeAsPieces(s) - - 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: - t = [self.bos_id] + t - if eos: - t = t + [self.eos_id] - return t - - def decode(self, t: List[int]) -> str: - text, buffer = "", [] - for token in t: - if token in self.index_special_tokens: - if buffer: - text += self.sp_model.decode(buffer) - buffer = [] - text += self.index_special_tokens[token] - else: - buffer.append(token) - if buffer: - text += self.sp_model.decode(buffer) - return text - - def decode_tokens(self, tokens: List[str]) -> str: - text = self.sp_model.DecodePieces(tokens) - return text - - def convert_token_to_id(self, token): - """Converts a token (str) in an id using the vocab.""" - if token in self.special_tokens: - return self.special_tokens[token] - return self.sp_model.PieceToId(token) - - def convert_id_to_token(self, index): - """Converts an index (integer) in a token (str) using the vocab.""" - if index in self.index_special_tokens: - return self.index_special_tokens[index] - if index in [self.eos_id, self.bos_id, self.pad_id] or index < 0: - return "" - return self.sp_model.IdToPiece(index) - - -class ChatGLMTokenizer(PreTrainedTokenizer): - vocab_files_names = {"vocab_file": "tokenizer.model"} - - model_input_names = ["input_ids", "attention_mask", "position_ids"] - - def __init__( - self, - vocab_file, - padding_side="left", - clean_up_tokenization_spaces=False, - encode_special_tokens=False, - **kwargs, - ): - self.name = "GLMTokenizer" - - self.vocab_file = vocab_file - self.tokenizer = SPTokenizer(vocab_file) - self.special_tokens = { - "": self.tokenizer.bos_id, - "": self.tokenizer.eos_id, - "": self.tokenizer.pad_id, - } - self.encode_special_tokens = encode_special_tokens - super().__init__( - padding_side=padding_side, - clean_up_tokenization_spaces=clean_up_tokenization_spaces, - encode_special_tokens=encode_special_tokens, - **kwargs, - ) - - def get_command(self, token): - if token in self.special_tokens: - return self.special_tokens[token] - assert token in self.tokenizer.special_tokens, f"{token} is not a special token for {self.name}" - return self.tokenizer.special_tokens[token] - - @property - def unk_token(self) -> str: - return "" - - @unk_token.setter - def unk_token(self, value: str): - self._unk_token = value - - @property - def pad_token(self) -> str: - return "" - - @pad_token.setter - def pad_token(self, value: str): - self._pad_token = value - - @property - def pad_token_id(self): - return self.get_command("") - - @property - def eos_token(self) -> str: - return "" - - @eos_token.setter - def eos_token(self, value: str): - self._eos_token = value - - @property - def eos_token_id(self): - return self.get_command("") - - @property - def vocab_size(self): - return self.tokenizer.n_words - - def get_vocab(self): - """Returns vocab as a dict""" - vocab = {self._convert_id_to_token(i): i for i in range(self.vocab_size)} - vocab.update(self.added_tokens_encoder) - return vocab - - def _tokenize(self, text, **kwargs): - return self.tokenizer.tokenize(text, encode_special_tokens=self.encode_special_tokens) - - def _convert_token_to_id(self, token): - """Converts a token (str) in an id using the vocab.""" - return self.tokenizer.convert_token_to_id(token) - - 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: - return self.tokenizer.decode_tokens(tokens) - - def save_vocabulary(self, save_directory, filename_prefix=None): - """ - Save the vocabulary and special tokens file to a directory. - - Args: - save_directory (`str`): - The directory in which to save the vocabulary. - filename_prefix (`str`, *optional*): - An optional prefix to add to the named of the saved files. - - Returns: - `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"]) - else: - vocab_file = save_directory - - with open(self.vocab_file, "rb") as fin: - proto_str = fin.read() - - with open(vocab_file, "wb") as writer: - writer.write(proto_str) - - return (vocab_file,) - - def get_prefix_tokens(self): - prefix_tokens = [self.get_command("[gMASK]"), self.get_command("sop")] - return prefix_tokens - - def build_single_message(self, role, metadata, message): - assert role in ["system", "user", "assistant", "observation"], role - role_tokens = [self.get_command(f"<|{role}|>")] + self.tokenizer.encode(f"{metadata}\n") - message_tokens = self.tokenizer.encode(message) - tokens = role_tokens + message_tokens - return tokens - - def build_chat_input(self, query, history=None, role="user"): - if history is None: - history = [] - input_ids = [] - for item in history: - content = item["content"] - if item["role"] == "system" and "tools" in item: - content = content + "\n" + json.dumps(item["tools"], indent=4, ensure_ascii=False) - input_ids.extend(self.build_single_message(item["role"], item.get("metadata", ""), content)) - input_ids.extend(self.build_single_message(role, "", query)) - input_ids.extend([self.get_command("<|assistant|>")]) - 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]: - """ - 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: - - - single sequence: `[CLS] X [SEP]` - - 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*): - Optional second list of IDs for sequence pairs. - - Returns: - `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 - if token_ids_1 is not None: - token_ids_0 = token_ids_0 + token_ids_1 + [self.get_command("")] - return token_ids_0 - - def _pad( - self, - encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding], - max_length: Optional[int] = 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, - ) -> 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]]`). - 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. - - - PaddingStrategy.LONGEST Pad to the longest sequence in the batch - - PaddingStrategy.MAX_LENGTH: Pad to the max length (default) - - PaddingStrategy.DO_NOT_PAD: Do not pad - The tokenizer padding sides are defined in self.padding_side: - - - 'left': pads on the left of the sequences - - 'right': pads on the right of the sequences - pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value. - This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability - `>= 7.5` (Volta). - padding_side (`str`, *optional*): - The side on which the model should have padding applied. Should be selected between ['right', 'left']. - Default value is picked from the class attribute of the same name. - return_attention_mask: - (optional) Set to False to avoid returning attention mask (default: set to model specifics) - """ - # Load from model defaults - assert self.padding_side == "left" - - required_input = encoded_inputs[self.model_input_names[0]] - seq_length = len(required_input) - - if padding_strategy == PaddingStrategy.LONGEST: - max_length = len(required_input) - - if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0): - max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of - - needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length - - # Initialize attention mask if not present. - if "attention_mask" not in encoded_inputs: - encoded_inputs["attention_mask"] = [1] * seq_length - - if "position_ids" not in encoded_inputs: - encoded_inputs["position_ids"] = list(range(seq_length)) - - if needs_to_be_padded: - difference = max_length - len(required_input) - - if "attention_mask" in encoded_inputs: - encoded_inputs["attention_mask"] = [0] * difference + encoded_inputs["attention_mask"] - if "position_ids" in encoded_inputs: - encoded_inputs["position_ids"] = [0] * difference + encoded_inputs["position_ids"] - encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input - - return encoded_inputs diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/__init__.py deleted file mode 100644 index 8f79d3c4773f393ed689a949041d36ad77e20968..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/__init__.py +++ /dev/null @@ -1,50 +0,0 @@ -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_latent_consistency_img2img"] = ["LatentConsistencyModelImg2ImgPipeline"] - _import_structure["pipeline_latent_consistency_text2img"] = ["LatentConsistencyModelPipeline"] - -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_latent_consistency_img2img import LatentConsistencyModelImg2ImgPipeline - from .pipeline_latent_consistency_text2img import LatentConsistencyModelPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 89848921ca320bea8beab29c927216ea928f617a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/__pycache__/pipeline_latent_consistency_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/__pycache__/pipeline_latent_consistency_img2img.cpython-310.pyc deleted file mode 100644 index ca55ce77be61f3732458ada90789f8fc97594a89..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/__pycache__/pipeline_latent_consistency_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/__pycache__/pipeline_latent_consistency_text2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/__pycache__/pipeline_latent_consistency_text2img.cpython-310.pyc deleted file mode 100644 index 1bce8fd7ab4eef0147ff29fdb740f885d9802573..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/__pycache__/pipeline_latent_consistency_text2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py deleted file mode 100644 index 1d7733982ef54ca7f6d8372d4477d89dca4960cb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py +++ /dev/null @@ -1,987 +0,0 @@ -# Copyright 2025 Stanford University 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. - -# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion -# and https://github.com/hojonathanho/diffusion - -import inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import LCMScheduler -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - 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, StableDiffusionMixin -from ..stable_diffusion import StableDiffusionPipelineOutput, StableDiffusionSafetyChecker - - -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 - - -# 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" -): - 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 - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> from diffusers import AutoPipelineForImage2Image - >>> import torch - >>> import PIL - - >>> pipe = AutoPipelineForImage2Image.from_pretrained("SimianLuo/LCM_Dreamshaper_v7") - >>> # To save GPU memory, torch.float16 can be used, but it may compromise image quality. - >>> pipe.to(torch_device="cuda", torch_dtype=torch.float32) - - >>> prompt = "High altitude snowy mountains" - >>> image = PIL.Image.open("./snowy_mountains.png") - - >>> # Can be set to 1~50 steps. LCM support fast inference even <= 4 steps. Recommend: 1~8 steps. - >>> num_inference_steps = 4 - >>> images = pipe( - ... prompt=prompt, image=image, num_inference_steps=num_inference_steps, guidance_scale=8.0 - ... ).images - - >>> images[0].save("image.png") - ``` - -""" - - -class LatentConsistencyModelImg2ImgPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - FromSingleFileMixin, -): - r""" - Pipeline for image-to-image generation using a latent consistency model. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. Currently only - supports [`LCMScheduler`]. - 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - requires_safety_checker (`bool`, *optional*, defaults to `True`): - Whether the pipeline requires a safety checker component. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "denoised", "prompt_embeds", "w_embedding"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: LCMScheduler, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, - requires_safety_checker: bool = True, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - 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) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.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_img2img.StableDiffusionImg2ImgPipeline.prepare_latents - def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - 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." - ) - - elif isinstance(generator, list): - if image.shape[0] < batch_size and batch_size % image.shape[0] == 0: - image = torch.cat([image] * (batch_size // image.shape[0]), dim=0) - elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} " - ) - - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - deprecation_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial" - " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" - " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" - " your script to pass as many initial images as text prompts to suppress this warning." - ) - deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - latents = init_latents - - return latents - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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: Union[str, List[str]], - strength: float, - callback_steps: int, - prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - ): - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - - 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 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" - ) - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def clip_skip(self): - return self._clip_skip - - @property - def do_classifier_free_guidance(self): - return False - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[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, - 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", - 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"], - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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. - original_inference_steps (`int`, *optional*): - The original number of inference steps use to generate a linearly-spaced timestep schedule, from which - 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*): - 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. - 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`. - Note that the original latent consistency models paper uses a different CFG formulation where the - guidance scales are decreased by 1 (so in the paper formulation CFG is enabled when `guidance_scale > - 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*): - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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). - 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. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - strength, - callback_steps, - prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - - # 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 - - 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, - self.do_classifier_free_guidance, - ) - - # 3. Encode input prompt - lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - - # NOTE: when a LCM is distilled from an LDM via latent consistency distillation (Algorithm 1) with guided - # distillation, the forward pass of the LCM learns to approximate sampling from the LDM using CFG with the - # unconditional prompt "" (the empty string). Due to this, LCMs currently do not support negative prompts. - prompt_embeds, _ = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt=None, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=None, - lora_scale=lora_scale, - clip_skip=self.clip_skip, - ) - - # 4. Encode image - image = self.image_processor.preprocess(image) - - # 5. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - device, - timesteps, - original_inference_steps=original_inference_steps, - strength=strength, - ) - - # 6. Prepare latent variables - original_inference_steps = ( - original_inference_steps - if original_inference_steps is not None - else self.scheduler.config.original_inference_steps - ) - latent_timestep = timesteps[:1] - if latents is None: - latents = self.prepare_latents( - image, latent_timestep, batch_size, num_images_per_prompt, prompt_embeds.dtype, device, generator - ) - bs = batch_size * num_images_per_prompt - - # 6. Get Guidance Scale Embedding - # NOTE: We use the Imagen CFG formulation that StableDiffusionPipeline uses rather than the original LCM paper - # CFG formulation, so we need to subtract 1 from the input guidance_scale. - # LCM CFG formulation: cfg_noise = noise_cond + cfg_scale * (noise_cond - noise_uncond), (cfg_scale > 0.0 using CFG) - w = torch.tensor(self.guidance_scale - 1).repeat(bs) - w_embedding = self.get_guidance_scale_embedding(w, embedding_dim=self.unet.config.time_cond_proj_dim).to( - device=device, dtype=latents.dtype - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, None) - - # 7.1 Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - # 8. LCM Multistep Sampling Loop - 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): - latents = latents.to(prompt_embeds.dtype) - - # model prediction (v-prediction, eps, x) - model_pred = self.unet( - latents, - t, - timestep_cond=w_embedding, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # compute the previous noisy sample x_t -> x_t-1 - latents, denoised = self.scheduler.step(model_pred, t, latents, **extra_step_kwargs, return_dict=False) - 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) - w_embedding = callback_outputs.pop("w_embedding", w_embedding) - denoised = callback_outputs.pop("denoised", denoised) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - denoised = denoised.to(prompt_embeds.dtype) - if not output_type == "latent": - image = self.vae.decode(denoised / self.vae.config.scaling_factor, return_dict=False)[0] - image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) - else: - image = denoised - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py deleted file mode 100644 index 3e96b44663e1d091c42a1b9f7932e1a2fe5f2f2d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py +++ /dev/null @@ -1,917 +0,0 @@ -# Copyright 2025 Stanford University 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. - -# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion -# and https://github.com/hojonathanho/diffusion - -import inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import LCMScheduler -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - 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, StableDiffusionMixin -from ..stable_diffusion import StableDiffusionPipelineOutput, StableDiffusionSafetyChecker - - -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 - >>> from diffusers import DiffusionPipeline - >>> import torch - - >>> pipe = DiffusionPipeline.from_pretrained("SimianLuo/LCM_Dreamshaper_v7") - >>> # To save GPU memory, torch.float16 can be used, but it may compromise image quality. - >>> pipe.to(torch_device="cuda", torch_dtype=torch.float32) - - >>> prompt = "Self-portrait oil painting, a beautiful cyborg with golden hair, 8k" - - >>> # Can be set to 1~50 steps. LCM support fast inference even <= 4 steps. Recommend: 1~8 steps. - >>> num_inference_steps = 4 - >>> images = pipe(prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=8.0).images - >>> images[0].save("image.png") - ``` -""" - - -# 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 LatentConsistencyModelPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-to-image generation using a latent consistency model. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. Currently only - supports [`LCMScheduler`]. - 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - requires_safety_checker (`bool`, *optional*, defaults to `True`): - Whether the pipeline requires a safety checker component. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "denoised", "prompt_embeds", "w_embedding"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: LCMScheduler, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.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.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Currently StableDiffusionPipeline.check_inputs with negative prompt stuff removed - def check_inputs( - self, - prompt: Union[str, List[str]], - height: int, - width: int, - callback_steps: int, - prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image=None, - ip_adapter_image_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 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" - ) - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def clip_skip(self): - return self._clip_skip - - @property - def do_classifier_free_guidance(self): - return False - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - @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: int = 4, - original_inference_steps: 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", - 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"], - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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. - original_inference_steps (`int`, *optional*): - The original number of inference steps use to generate a linearly-spaced timestep schedule, from which - 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*): - 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. - 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`. - Note that the original latent consistency models paper uses a different CFG formulation where the - guidance scales are decreased by 1 (so in the paper formulation CFG is enabled when `guidance_scale > - 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*): - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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). - 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. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - # 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, - prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - - # 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 - - 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, - self.do_classifier_free_guidance, - ) - - # 3. Encode input prompt - lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - - # NOTE: when a LCM is distilled from an LDM via latent consistency distillation (Algorithm 1) with guided - # distillation, the forward pass of the LCM learns to approximate sampling from the LDM using CFG with the - # unconditional prompt "" (the empty string). Due to this, LCMs currently do not support negative prompts. - prompt_embeds, _ = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt=None, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=None, - lora_scale=lora_scale, - clip_skip=self.clip_skip, - ) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, original_inference_steps=original_inference_steps - ) - - # 5. Prepare latent variable - 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, - ) - bs = batch_size * num_images_per_prompt - - # 6. Get Guidance Scale Embedding - # NOTE: We use the Imagen CFG formulation that StableDiffusionPipeline uses rather than the original LCM paper - # CFG formulation, so we need to subtract 1 from the input guidance_scale. - # LCM CFG formulation: cfg_noise = noise_cond + cfg_scale * (noise_cond - noise_uncond), (cfg_scale > 0.0 using CFG) - w = torch.tensor(self.guidance_scale - 1).repeat(bs) - w_embedding = self.get_guidance_scale_embedding(w, embedding_dim=self.unet.config.time_cond_proj_dim).to( - device=device, dtype=latents.dtype - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, None) - - # 7.1 Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - # 8. LCM MultiStep Sampling Loop: - 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): - latents = latents.to(prompt_embeds.dtype) - - # model prediction (v-prediction, eps, x) - model_pred = self.unet( - latents, - t, - timestep_cond=w_embedding, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # compute the previous noisy sample x_t -> x_t-1 - latents, denoised = self.scheduler.step(model_pred, t, latents, **extra_step_kwargs, return_dict=False) - 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) - w_embedding = callback_outputs.pop("w_embedding", w_embedding) - denoised = callback_outputs.pop("denoised", denoised) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - denoised = denoised.to(prompt_embeds.dtype) - if not output_type == "latent": - image = self.vae.decode(denoised / self.vae.config.scaling_factor, return_dict=False)[0] - image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) - else: - image = denoised - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/__init__.py deleted file mode 100644 index 561f96fc71dc7b4404e09571e0b7eaa4ee02fde8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/__init__.py +++ /dev/null @@ -1,50 +0,0 @@ -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_latent_diffusion"] = ["LDMBertModel", "LDMTextToImagePipeline"] - _import_structure["pipeline_latent_diffusion_superresolution"] = ["LDMSuperResolutionPipeline"] - - -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_latent_diffusion import LDMBertModel, LDMTextToImagePipeline - from .pipeline_latent_diffusion_superresolution import LDMSuperResolutionPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 06850ec6846c168686d5380e66703bcb87e5f504..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/__pycache__/pipeline_latent_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/__pycache__/pipeline_latent_diffusion.cpython-310.pyc deleted file mode 100644 index fc5c6acae609dd71e62a3bece7ef29119355f5ba..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/__pycache__/pipeline_latent_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/__pycache__/pipeline_latent_diffusion_superresolution.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/__pycache__/pipeline_latent_diffusion_superresolution.cpython-310.pyc deleted file mode 100644 index ee6d4d6e975c8489eec5086f2522112a45226474..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/__pycache__/pipeline_latent_diffusion_superresolution.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py deleted file mode 100644 index bc50835d19d345312968544a19d9e481b7dc2231..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py +++ /dev/null @@ -1,747 +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 inspect -from typing import List, Optional, Tuple, Union - -import torch -import torch.nn as nn -import torch.utils.checkpoint -from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizer -from transformers.activations import ACT2FN -from transformers.modeling_outputs import BaseModelOutput -from transformers.utils import logging - -from ...models import AutoencoderKL, UNet2DConditionModel, UNet2DModel, VQModel -from ...schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler -from ...utils import is_torch_xla_available -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - - XLA_AVAILABLE = True -else: - XLA_AVAILABLE = False - - -class LDMTextToImagePipeline(DiffusionPipeline): - r""" - Pipeline for text-to-image generation using latent diffusion. - - 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.). - - Parameters: - vqvae ([`VQModel`]): - Vector-quantized (VQ) model to encode and decode images to and from latent representations. - bert ([`LDMBertModel`]): - Text-encoder model based on [`~transformers.BERT`]. - tokenizer ([`~transformers.BertTokenizer`]): - A `BertTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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`]. - """ - - model_cpu_offload_seq = "bert->unet->vqvae" - - def __init__( - self, - vqvae: Union[VQModel, AutoencoderKL], - bert: PreTrainedModel, - tokenizer: PreTrainedTokenizer, - unet: Union[UNet2DModel, UNet2DConditionModel], - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], - ): - super().__init__() - self.register_modules(vqvae=vqvae, bert=bert, tokenizer=tokenizer, unet=unet, scheduler=scheduler) - self.vae_scale_factor = 2 ** (len(self.vqvae.config.block_out_channels) - 1) - - @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", - return_dict: bool = True, - **kwargs, - ) -> Union[Tuple, ImagePipelineOutput]: - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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 1.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`. - generator (`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`. - 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 [`ImagePipelineOutput`] instead of a plain tuple. - - Example: - - ```py - >>> from diffusers import DiffusionPipeline - - >>> # load model and scheduler - >>> ldm = DiffusionPipeline.from_pretrained("CompVis/ldm-text2im-large-256") - - >>> # run pipeline in inference (sample random noise and denoise) - >>> prompt = "A painting of a squirrel eating a burger" - >>> images = ldm([prompt], num_inference_steps=50, eta=0.3, guidance_scale=6).images - - >>> # save images - >>> for idx, image in enumerate(images): - ... image.save(f"squirrel-{idx}.png") - ``` - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images. - """ - # 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 - - if isinstance(prompt, str): - batch_size = 1 - elif isinstance(prompt, list): - batch_size = len(prompt) - else: - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - 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}.") - - # get unconditional embeddings for classifier free guidance - if guidance_scale != 1.0: - uncond_input = self.tokenizer( - [""] * batch_size, padding="max_length", max_length=77, truncation=True, return_tensors="pt" - ) - negative_prompt_embeds = self.bert(uncond_input.input_ids.to(self._execution_device))[0] - - # get prompt text embeddings - text_input = self.tokenizer(prompt, padding="max_length", max_length=77, truncation=True, return_tensors="pt") - prompt_embeds = self.bert(text_input.input_ids.to(self._execution_device))[0] - - # get the initial random noise unless the user supplied it - latents_shape = (batch_size, self.unet.config.in_channels, height // 8, width // 8) - 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( - latents_shape, generator=generator, device=self._execution_device, dtype=prompt_embeds.dtype - ) - else: - if latents.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - latents = latents.to(self._execution_device) - - self.scheduler.set_timesteps(num_inference_steps) - - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - - extra_kwargs = {} - if accepts_eta: - extra_kwargs["eta"] = eta - - for t in self.progress_bar(self.scheduler.timesteps): - if guidance_scale == 1.0: - # guidance_scale of 1 means no guidance - latents_input = latents - context = prompt_embeds - else: - # 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 - latents_input = torch.cat([latents] * 2) - context = torch.cat([negative_prompt_embeds, prompt_embeds]) - - # predict the noise residual - noise_pred = self.unet(latents_input, t, encoder_hidden_states=context).sample - # perform guidance - if guidance_scale != 1.0: - noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_kwargs).prev_sample - - if XLA_AVAILABLE: - xm.mark_step() - - # scale and decode the image latents with vae - latents = 1 / self.vqvae.config.scaling_factor * latents - image = self.vqvae.decode(latents).sample - - image = (image / 2 + 0.5).clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).numpy() - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) - - -################################################################################ -# Code for the text transformer model -################################################################################ -""" PyTorch LDMBERT model.""" - - -logger = logging.get_logger(__name__) - -LDMBERT_PRETRAINED_MODEL_ARCHIVE_LIST = [ - "ldm-bert", - # See all LDMBert models at https://huggingface.co/models?filter=ldmbert -] - - -LDMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = { - "ldm-bert": "https://huggingface.co/valhalla/ldm-bert/blob/main/config.json", -} - - -""" LDMBERT model configuration""" - - -class LDMBertConfig(PretrainedConfig): - model_type = "ldmbert" - keys_to_ignore_at_inference = ["past_key_values"] - attribute_map = {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"} - - def __init__( - self, - vocab_size=30522, - max_position_embeddings=77, - encoder_layers=32, - encoder_ffn_dim=5120, - encoder_attention_heads=8, - head_dim=64, - encoder_layerdrop=0.0, - activation_function="gelu", - d_model=1280, - dropout=0.1, - attention_dropout=0.0, - activation_dropout=0.0, - init_std=0.02, - classifier_dropout=0.0, - scale_embedding=False, - use_cache=True, - pad_token_id=0, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.d_model = d_model - self.encoder_ffn_dim = encoder_ffn_dim - self.encoder_layers = encoder_layers - self.encoder_attention_heads = encoder_attention_heads - self.head_dim = head_dim - self.dropout = dropout - self.attention_dropout = attention_dropout - self.activation_dropout = activation_dropout - self.activation_function = activation_function - self.init_std = init_std - self.encoder_layerdrop = encoder_layerdrop - self.classifier_dropout = classifier_dropout - self.use_cache = use_cache - self.num_hidden_layers = encoder_layers - self.scale_embedding = scale_embedding # scale factor will be sqrt(d_model) if True - - super().__init__(pad_token_id=pad_token_id, **kwargs) - - -def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): - """ - Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. - """ - bsz, src_len = mask.size() - tgt_len = tgt_len if tgt_len is not None else src_len - - expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype) - - inverted_mask = 1.0 - expanded_mask - - return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min) - - -# Copied from transformers.models.bart.modeling_bart.BartAttention with Bart->LDMBert -class LDMBertAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" - - def __init__( - self, - embed_dim: int, - num_heads: int, - head_dim: int, - dropout: float = 0.0, - is_decoder: bool = False, - bias: bool = False, - ): - super().__init__() - self.embed_dim = embed_dim - self.num_heads = num_heads - self.dropout = dropout - self.head_dim = head_dim - self.inner_dim = head_dim * num_heads - - self.scaling = self.head_dim**-0.5 - self.is_decoder = is_decoder - - self.k_proj = nn.Linear(embed_dim, self.inner_dim, bias=bias) - self.v_proj = nn.Linear(embed_dim, self.inner_dim, bias=bias) - self.q_proj = nn.Linear(embed_dim, self.inner_dim, bias=bias) - self.out_proj = nn.Linear(self.inner_dim, embed_dim) - - def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): - return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() - - 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, - output_attentions: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: - """Input shape: Batch x Time x Channel""" - - # if key_value_states are provided this layer is used as a cross-attention layer - # for the decoder - is_cross_attention = key_value_states is not None - - bsz, tgt_len, _ = hidden_states.size() - - # get query proj - query_states = self.q_proj(hidden_states) * self.scaling - # get key, value proj - if is_cross_attention and past_key_value is not None: - # reuse k,v, cross_attentions - key_states = past_key_value[0] - value_states = past_key_value[1] - elif is_cross_attention: - # cross_attentions - key_states = self._shape(self.k_proj(key_value_states), -1, bsz) - value_states = self._shape(self.v_proj(key_value_states), -1, bsz) - elif past_key_value is not None: - # reuse k, v, self_attention - key_states = self._shape(self.k_proj(hidden_states), -1, bsz) - value_states = self._shape(self.v_proj(hidden_states), -1, bsz) - key_states = torch.cat([past_key_value[0], key_states], dim=2) - value_states = torch.cat([past_key_value[1], value_states], dim=2) - else: - # self_attention - key_states = self._shape(self.k_proj(hidden_states), -1, bsz) - 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. - # 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 - # 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` - past_key_value = (key_states, value_states) - - proj_shape = (bsz * self.num_heads, -1, self.head_dim) - query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape) - key_states = key_states.view(*proj_shape) - value_states = value_states.view(*proj_shape) - - src_len = key_states.size(1) - attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) - - if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len): - raise ValueError( - f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is" - f" {attn_weights.size()}" - ) - - if attention_mask is not None: - if attention_mask.size() != (bsz, 1, tgt_len, src_len): - raise ValueError( - f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}" - ) - attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask - attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) - - attn_weights = nn.functional.softmax(attn_weights, dim=-1) - - if layer_head_mask is not None: - if layer_head_mask.size() != (self.num_heads,): - raise ValueError( - f"Head mask for a single layer should be of size {(self.num_heads,)}, but is" - f" {layer_head_mask.size()}" - ) - attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len) - attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) - - if output_attentions: - # this operation is a bit awkward, but it's required to - # make sure that attn_weights keeps its gradient. - # In order to do so, attn_weights have to be reshaped - # twice and have to be reused in the following - attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) - attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len) - else: - attn_weights_reshaped = None - - attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) - - attn_output = torch.bmm(attn_probs, value_states) - - if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim): - raise ValueError( - f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is" - f" {attn_output.size()}" - ) - - attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim) - attn_output = attn_output.transpose(1, 2) - - # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be - # partitioned across GPUs when using tensor-parallelism. - attn_output = attn_output.reshape(bsz, tgt_len, self.inner_dim) - - attn_output = self.out_proj(attn_output) - - return attn_output, attn_weights_reshaped, past_key_value - - -class LDMBertEncoderLayer(nn.Module): - def __init__(self, config: LDMBertConfig): - super().__init__() - self.embed_dim = config.d_model - self.self_attn = LDMBertAttention( - embed_dim=self.embed_dim, - num_heads=config.encoder_attention_heads, - head_dim=config.head_dim, - dropout=config.attention_dropout, - ) - self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) - self.dropout = config.dropout - self.activation_fn = ACT2FN[config.activation_function] - self.activation_dropout = config.activation_dropout - self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim) - self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim) - self.final_layer_norm = nn.LayerNorm(self.embed_dim) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor, - layer_head_mask: torch.Tensor, - output_attentions: Optional[bool] = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """ - Args: - hidden_states (`torch.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` - attention_mask (`torch.Tensor`): attention mask of size - `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. - layer_head_mask (`torch.Tensor`): mask for attention heads in a given layer of size - `(encoder_attention_heads,)`. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. - """ - residual = hidden_states - hidden_states = self.self_attn_layer_norm(hidden_states) - hidden_states, attn_weights, _ = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - layer_head_mask=layer_head_mask, - output_attentions=output_attentions, - ) - hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.final_layer_norm(hidden_states) - hidden_states = self.activation_fn(self.fc1(hidden_states)) - hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) - hidden_states = self.fc2(hidden_states) - hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - hidden_states = residual + hidden_states - - if hidden_states.dtype == torch.float16 and ( - torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any() - ): - clamp_value = torch.finfo(hidden_states.dtype).max - 1000 - hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - - outputs = (hidden_states,) - - if output_attentions: - outputs += (attn_weights,) - - return outputs - - -# Copied from transformers.models.bart.modeling_bart.BartPretrainedModel with Bart->LDMBert -class LDMBertPreTrainedModel(PreTrainedModel): - config_class = LDMBertConfig - base_model_prefix = "model" - _supports_gradient_checkpointing = True - _keys_to_ignore_on_load_unexpected = [r"encoder\.version", r"decoder\.version"] - - def _init_weights(self, module): - std = self.config.init_std - if isinstance(module, nn.Linear): - module.weight.data.normal_(mean=0.0, std=std) - if module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, nn.Embedding): - module.weight.data.normal_(mean=0.0, std=std) - if module.padding_idx is not None: - module.weight.data[module.padding_idx].zero_() - - @property - def dummy_inputs(self): - pad_token = self.config.pad_token_id - input_ids = torch.tensor([[0, 6, 10, 4, 2], [0, 8, 12, 2, pad_token]], device=self.device) - dummy_inputs = { - "attention_mask": input_ids.ne(pad_token), - "input_ids": input_ids, - } - return dummy_inputs - - -class LDMBertEncoder(LDMBertPreTrainedModel): - """ - Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a - [`LDMBertEncoderLayer`]. - - Args: - config: LDMBertConfig - embed_tokens (nn.Embedding): output embedding - """ - - def __init__(self, config: LDMBertConfig): - super().__init__(config) - - self.dropout = config.dropout - - embed_dim = config.d_model - self.padding_idx = config.pad_token_id - self.max_source_positions = config.max_position_embeddings - - self.embed_tokens = nn.Embedding(config.vocab_size, embed_dim) - self.embed_positions = nn.Embedding(config.max_position_embeddings, embed_dim) - self.layers = nn.ModuleList([LDMBertEncoderLayer(config) for _ in range(config.encoder_layers)]) - self.layer_norm = nn.LayerNorm(embed_dim) - - self.gradient_checkpointing = False - # Initialize weights and apply final processing - self.post_init() - - def get_input_embeddings(self): - return self.embed_tokens - - def set_input_embeddings(self, value): - self.embed_tokens = 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]: - r""" - Args: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you - provide it. - - Indices can be obtained using [`BartTokenizer`]. See [`PreTrainedTokenizer.encode`] and - [`PreTrainedTokenizer.__call__`] for details. - - [What are input IDs?](../glossary#input-ids) - attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - - [What are attention masks?](../glossary#attention-mask) - head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*): - Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - inputs_embeds (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): - Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. - This is useful if you want more control over how to convert `input_ids` indices into associated vectors - than the model's internal embedding lookup matrix. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. - output_hidden_states (`bool`, *optional*): - Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors - for more detail. - return_dict (`bool`, *optional*): - Whether or not to return a [`~utils.BaseModelOutput`] instead of a plain tuple. - """ - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # retrieve input_ids and inputs_embeds - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") - elif input_ids is not None: - input_shape = input_ids.size() - input_ids = input_ids.view(-1, input_shape[-1]) - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] - else: - raise ValueError("You have to specify either input_ids or inputs_embeds") - - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - - seq_len = input_shape[1] - if position_ids is None: - position_ids = torch.arange(seq_len, dtype=torch.long, device=inputs_embeds.device).expand((1, -1)) - embed_pos = self.embed_positions(position_ids) - - hidden_states = inputs_embeds + embed_pos - hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - - # expand attention_mask - if attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _expand_mask(attention_mask, inputs_embeds.dtype) - - encoder_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - - # check if head_mask has a correct number of layers specified if desired - if head_mask is not None: - if head_mask.size()[0] != (len(self.layers)): - raise ValueError( - f"The head_mask should be specified for {len(self.layers)} layers, but it is for" - f" {head_mask.size()[0]}." - ) - - for idx, encoder_layer in enumerate(self.layers): - if output_hidden_states: - encoder_states = encoder_states + (hidden_states,) - if torch.is_grad_enabled() and self.gradient_checkpointing: - layer_outputs = self._gradient_checkpointing_func( - encoder_layer, - hidden_states, - attention_mask, - (head_mask[idx] if head_mask is not None else None), - ) - else: - layer_outputs = encoder_layer( - hidden_states, - attention_mask, - layer_head_mask=(head_mask[idx] if head_mask is not None else None), - output_attentions=output_attentions, - ) - - hidden_states = layer_outputs[0] - - if output_attentions: - all_attentions = all_attentions + (layer_outputs[1],) - - hidden_states = self.layer_norm(hidden_states) - - if output_hidden_states: - encoder_states = encoder_states + (hidden_states,) - - if not return_dict: - return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None) - return BaseModelOutput( - last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions - ) - - -class LDMBertModel(LDMBertPreTrainedModel): - _no_split_modules = [] - - def __init__(self, config: LDMBertConfig): - super().__init__(config) - self.model = LDMBertEncoder(config) - self.to_logits = nn.Linear(config.hidden_size, config.vocab_size) - - def forward( - self, - input_ids=None, - attention_mask=None, - position_ids=None, - head_mask=None, - inputs_embeds=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - ): - outputs = self.model( - input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - head_mask=head_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - return outputs diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py deleted file mode 100644 index 273e97f1ec608bd2fe2fd5d9ea80a495357d4c52..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py +++ /dev/null @@ -1,200 +0,0 @@ -import inspect -from typing import List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch -import torch.utils.checkpoint - -from ...models import UNet2DModel, VQModel -from ...schedulers import ( - DDIMScheduler, - DPMSolverMultistepScheduler, - EulerAncestralDiscreteScheduler, - EulerDiscreteScheduler, - LMSDiscreteScheduler, - PNDMScheduler, -) -from ...utils import PIL_INTERPOLATION, is_torch_xla_available -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - - XLA_AVAILABLE = True -else: - XLA_AVAILABLE = False - - -def preprocess(image): - w, h = image.size - w, h = (x - x % 32 for x in (w, h)) # resize to integer multiple of 32 - image = image.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]) - image = np.array(image).astype(np.float32) / 255.0 - image = image[None].transpose(0, 3, 1, 2) - image = torch.from_numpy(image) - return 2.0 * image - 1.0 - - -class LDMSuperResolutionPipeline(DiffusionPipeline): - r""" - A pipeline for image super-resolution using latent diffusion. - - 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.). - - Parameters: - vqvae ([`VQModel`]): - Vector-quantized (VQ) model to encode and decode images to and from latent representations. - unet ([`UNet2DModel`]): - A `UNet2DModel` to denoise the encoded image. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latens. Can be one of - [`DDIMScheduler`], [`LMSDiscreteScheduler`], [`EulerDiscreteScheduler`], - [`EulerAncestralDiscreteScheduler`], [`DPMSolverMultistepScheduler`], or [`PNDMScheduler`]. - """ - - def __init__( - self, - vqvae: VQModel, - unet: UNet2DModel, - scheduler: Union[ - DDIMScheduler, - PNDMScheduler, - LMSDiscreteScheduler, - EulerDiscreteScheduler, - EulerAncestralDiscreteScheduler, - DPMSolverMultistepScheduler, - ], - ): - super().__init__() - self.register_modules(vqvae=vqvae, unet=unet, scheduler=scheduler) - - @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", - return_dict: bool = True, - ) -> Union[Tuple, ImagePipelineOutput]: - r""" - The call function to the pipeline for generation. - - Args: - image (`torch.Tensor` or `PIL.Image.Image`): - `Image` or tensor representing an image batch to be used as the starting point for the process. - batch_size (`int`, *optional*, defaults to 1): - Number of images to generate. - 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. - 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*): - A [`torch.Generator`](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 generated image. Choose between `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`ImagePipelineOutput`] instead of a plain tuple. - - Example: - - ```py - >>> import requests - >>> from PIL import Image - >>> from io import BytesIO - >>> from diffusers import LDMSuperResolutionPipeline - >>> import torch - - >>> # load model and scheduler - >>> pipeline = LDMSuperResolutionPipeline.from_pretrained("CompVis/ldm-super-resolution-4x-openimages") - >>> pipeline = pipeline.to("cuda") - - >>> # let's download an image - >>> url = ( - ... "https://user-images.githubusercontent.com/38061659/199705896-b48e17b8-b231-47cd-a270-4ffa5a93fa3e.png" - ... ) - >>> response = requests.get(url) - >>> low_res_img = Image.open(BytesIO(response.content)).convert("RGB") - >>> low_res_img = low_res_img.resize((128, 128)) - - >>> # run pipeline in inference (sample random noise and denoise) - >>> upscaled_image = pipeline(low_res_img, num_inference_steps=100, eta=1).images[0] - >>> # save image - >>> upscaled_image.save("ldm_generated_image.png") - ``` - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images - """ - if isinstance(image, PIL.Image.Image): - batch_size = 1 - elif isinstance(image, torch.Tensor): - batch_size = image.shape[0] - else: - raise ValueError(f"`image` has to be of type `PIL.Image.Image` or `torch.Tensor` but is {type(image)}") - - if isinstance(image, PIL.Image.Image): - image = preprocess(image) - - height, width = image.shape[-2:] - - # in_channels should be 6: 3 for latents, 3 for low resolution image - latents_shape = (batch_size, self.unet.config.in_channels // 2, height, width) - latents_dtype = next(self.unet.parameters()).dtype - - latents = randn_tensor(latents_shape, generator=generator, device=self.device, dtype=latents_dtype) - - image = image.to(device=self.device, dtype=latents_dtype) - - # set timesteps and move to the correct device - self.scheduler.set_timesteps(num_inference_steps, device=self.device) - timesteps_tensor = self.scheduler.timesteps - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature. - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_kwargs = {} - if accepts_eta: - extra_kwargs["eta"] = eta - - for t in self.progress_bar(timesteps_tensor): - # concat latents and low resolution image in the channel dimension. - latents_input = torch.cat([latents, image], dim=1) - latents_input = self.scheduler.scale_model_input(latents_input, t) - # predict the noise residual - noise_pred = self.unet(latents_input, t).sample - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_kwargs).prev_sample - - if XLA_AVAILABLE: - xm.mark_step() - - # decode the image latents with the VQVAE - image = self.vqvae.decode(latents).sample - image = torch.clamp(image, -1.0, 1.0) - image = image / 2 + 0.5 - image = image.cpu().permute(0, 2, 3, 1).numpy() - - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/latte/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/latte/__init__.py deleted file mode 100644 index 4296b42e125303c0e036b9f2deecc36bdf959de3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/latte/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -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_latte"] = ["LattePipeline"] - -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_latte import LattePipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/latte/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/latte/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 532459dba63f8a1ae69eca51f6cb57fee6bf061e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/latte/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/latte/__pycache__/pipeline_latte.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/latte/__pycache__/pipeline_latte.cpython-310.pyc deleted file mode 100644 index 78e602348ffaeeaa5e7fbd8fc264bced55f882c5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/latte/__pycache__/pipeline_latte.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/latte/pipeline_latte.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/latte/pipeline_latte.py deleted file mode 100644 index 0e60d5c7acbec32e6869423c451a595b9eda9034..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/latte/pipeline_latte.py +++ /dev/null @@ -1,910 +0,0 @@ -# Copyright 2025 the Latte 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 inspect -import re -import urllib.parse as ul -from dataclasses import dataclass -from typing import Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import T5EncoderModel, T5Tokenizer - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...models import AutoencoderKL, LatteTransformer3DModel -from ...pipelines.pipeline_utils import DiffusionPipeline -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - BACKENDS_MAPPING, - BaseOutput, - deprecate, - is_bs4_available, - is_ftfy_available, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import is_compiled_module, randn_tensor -from ...video_processor import VideoProcessor - - -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_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import LattePipeline - >>> from diffusers.utils import export_to_gif - - >>> # You can replace the checkpoint id with "maxin-cn/Latte-1" too. - >>> pipe = LattePipeline.from_pretrained("maxin-cn/Latte-1", torch_dtype=torch.float16) - >>> # Enable memory optimizations. - >>> pipe.enable_model_cpu_offload() - - >>> prompt = "A small cactus with a happy face in the Sahara desert." - >>> videos = pipe(prompt).frames[0] - >>> export_to_gif(videos, "latte.gif") - ``` -""" - - -# 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 - - -@dataclass -class LattePipelineOutput(BaseOutput): - frames: torch.Tensor - - -class LattePipeline(DiffusionPipeline): - r""" - Pipeline for text-to-video generation using Latte. - - 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.) - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. - text_encoder ([`T5EncoderModel`]): - Frozen text-encoder. Latte uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the - [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant. - tokenizer (`T5Tokenizer`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - transformer ([`LatteTransformer3DModel`]): - A text conditioned `LatteTransformer3DModel` to denoise the encoded video latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `transformer` to denoise the encoded video latents. - """ - - bad_punct_regex = re.compile(r"[#®•©™&@·º½¾¿¡§~\)\(\]\[\}\{\|\\/\\*]{1,}") - - _optional_components = ["tokenizer", "text_encoder"] - model_cpu_offload_seq = "text_encoder->transformer->vae" - - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - ] - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - vae: AutoencoderKL, - transformer: LatteTransformer3DModel, - scheduler: KarrasDiffusionSchedulers, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_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.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor) - - # Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/utils.py - def mask_text_embeddings(self, emb, mask): - if emb.shape[0] == 1: - keep_index = mask.sum().item() - return emb[:, :, :keep_index, :], keep_index # 1, 120, 4096 -> 1 7 4096 - else: - masked_feature = emb * mask[:, None, :, None] # 1 120 4096 - return masked_feature, emb.shape[2] - - # Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - mask_feature: bool = True, - dtype=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 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 "". - 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 video that should be generated per prompt - device: (`torch.device`, *optional*): - torch device to place the resulting embeddings on - 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. For Latte, it's should be the embeddings of the "" string. - clean_caption (bool, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - mask_feature: (bool, defaults to `True`): - If `True`, the function will mask the text embeddings. - """ - embeds_initially_provided = prompt_embeds is not None and negative_prompt_embeds is not None - - if device is None: - device = self._execution_device - - 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] - - max_length = 120 - if prompt_embeds is None: - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=True, - add_special_tokens=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[:, max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {max_length} tokens: {removed_text}" - ) - - attention_mask = text_inputs.attention_mask.to(device) - prompt_embeds_attention_mask = attention_mask - - prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask) - prompt_embeds = prompt_embeds[0] - else: - prompt_embeds_attention_mask = torch.ones_like(prompt_embeds) - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - elif self.transformer is not None: - dtype = self.transformer.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - - bs_embed, 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(bs_embed * num_images_per_prompt, seq_len, -1) - prompt_embeds_attention_mask = prompt_embeds_attention_mask.view(bs_embed, -1) - prompt_embeds_attention_mask = prompt_embeds_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 = [negative_prompt] * batch_size if isinstance(negative_prompt, str) else negative_prompt - uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption) - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - attention_mask = uncond_input.attention_mask.to(device) - - 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=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) - - # 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 - else: - negative_prompt_embeds = None - - # Perform additional masking. - if mask_feature and not embeds_initially_provided: - prompt_embeds = prompt_embeds.unsqueeze(1) - masked_prompt_embeds, keep_indices = self.mask_text_embeddings(prompt_embeds, prompt_embeds_attention_mask) - masked_prompt_embeds = masked_prompt_embeds.squeeze(1) - masked_negative_prompt_embeds = ( - negative_prompt_embeds[:, :keep_indices, :] if negative_prompt_embeds is not None else None - ) - - return masked_prompt_embeds, masked_negative_prompt_embeds - - return prompt_embeds, negative_prompt_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - negative_prompt, - callback_on_step_end_tensor_inputs, - prompt_embeds=None, - negative_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 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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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}." - ) - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - # Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_synth.TextToVideoSDPipeline.prepare_latents - def prepare_latents( - self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None - ): - shape = ( - batch_size, - num_channels_latents, - num_frames, - height // self.vae_scale_factor, - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @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 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: str = "", - num_inference_steps: int = 50, - timesteps: Optional[List[int]] = 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, - 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"], - clean_caption: bool = True, - mask_feature: bool = True, - enable_temporal_attentions: bool = True, - decode_chunk_size: int = 14, - ) -> Union[LattePipelineOutput, Tuple]: - """ - Function invoked when calling the pipeline for generation. - - Args: - 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*): - 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*): - 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): - 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 videos that are closely linked to - the text `prompt`, usually at the expense of lower video quality. - video_length (`int`, *optional*, defaults to 16): - The number of video frames that are generated. Defaults to 16 frames which at 8 frames per seconds - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of videos to generate per prompt. - height (`int`, *optional*, defaults to self.unet.config.sample_size): - The height in pixels of the generated video. - width (`int`, *optional*, defaults to self.unet.config.sample_size): - The width in pixels of the generated video. - 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*): - 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 video - 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. For Latte 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 generate video. 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.IFPipelineOutput`] instead of a plain tuple. - callback_on_step_end (`Callable[[int, int, Dict], 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*): - 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`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - mask_feature (`bool` defaults to `True`): If set to `True`, the text embeddings will be masked. - enable_temporal_attentions (`bool`, *optional*, defaults to `True`): Whether to enable temporal attentions - decode_chunk_size (`int`, *optional*): - The number of frames to decode at a time. Higher chunk size leads to better temporal consistency at the - expense of more memory usage. By default, the decoder decodes all frames at once for maximal quality. - For lower memory usage, reduce `decode_chunk_size`. - - Examples: - - Returns: - [`~pipelines.latte.pipeline_latte.LattePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.latte.pipeline_latte.LattePipelineOutput`] 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 - - # 0. Default - decode_chunk_size = decode_chunk_size if decode_chunk_size is not None else video_length - - # 1. Check inputs. Raise error if not correct - height = height or self.transformer.config.sample_size * self.vae_scale_factor - width = width or self.transformer.config.sample_size * self.vae_scale_factor - self.check_inputs( - prompt, - height, - width, - negative_prompt, - callback_on_step_end_tensor_inputs, - prompt_embeds, - negative_prompt_embeds, - ) - self._guidance_scale = guidance_scale - self._current_timestep = None - self._interrupt = False - - # 2. Default height and width to transformer - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - do_classifier_free_guidance, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - clean_caption=clean_caption, - mask_feature=mask_feature, - ) - if do_classifier_free_guidance: - 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) - self._num_timesteps = len(timesteps) - - # 5. Prepare latents. - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - video_length, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 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 - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - current_timestep = t - if not torch.is_tensor(current_timestep): - # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = latent_model_input.device.type == "mps" - is_npu = latent_model_input.device.type == "npu" - if isinstance(current_timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device) - elif len(current_timestep.shape) == 0: - current_timestep = current_timestep[None].to(latent_model_input.device) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - current_timestep = current_timestep.expand(latent_model_input.shape[0]) - - # predict noise model_output - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=current_timestep, - enable_temporal_attentions=enable_temporal_attentions, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # use learned sigma? - if not ( - hasattr(self.scheduler.config, "variance_type") - and self.scheduler.config.variance_type in ["learned", "learned_range"] - ): - noise_pred = noise_pred.chunk(2, dim=1)[0] - - # compute previous video: x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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 == 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 == "latents": - deprecation_message = ( - "Passing `output_type='latents'` is deprecated. Please pass `output_type='latent'` instead." - ) - deprecate("output_type_latents", "1.0.0", deprecation_message, standard_warn=False) - output_type = "latent" - - if not output_type == "latent": - video = self.decode_latents(latents, video_length, decode_chunk_size=decode_chunk_size) - video = self.video_processor.postprocess_video(video=video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return LattePipelineOutput(frames=video) - - # Similar to diffusers.pipelines.stable_video_diffusion.pipeline_stable_video_diffusion.decode_latents - def decode_latents(self, latents: torch.Tensor, video_length: int, decode_chunk_size: int = 14): - # [batch, channels, frames, height, width] -> [batch*frames, channels, height, width] - latents = latents.permute(0, 2, 1, 3, 4).flatten(0, 1) - - latents = 1 / self.vae.config.scaling_factor * latents - - forward_vae_fn = self.vae._orig_mod.forward if is_compiled_module(self.vae) else self.vae.forward - accepts_num_frames = "num_frames" in set(inspect.signature(forward_vae_fn).parameters.keys()) - - # decode decode_chunk_size frames at a time to avoid OOM - frames = [] - for i in range(0, latents.shape[0], decode_chunk_size): - num_frames_in = latents[i : i + decode_chunk_size].shape[0] - decode_kwargs = {} - if accepts_num_frames: - # we only pass num_frames_in if it's expected - decode_kwargs["num_frames"] = num_frames_in - - frame = self.vae.decode(latents[i : i + decode_chunk_size], **decode_kwargs).sample - frames.append(frame) - frames = torch.cat(frames, dim=0) - - # [batch*frames, channels, height, width] -> [batch, channels, frames, height, width] - frames = frames.reshape(-1, video_length, *frames.shape[1:]).permute(0, 2, 1, 3, 4) - - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16 - frames = frames.float() - return frames diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__init__.py deleted file mode 100644 index aae3b1cb18ce96c7d79d3c44bc37cf2ae4db9720..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__init__.py +++ /dev/null @@ -1,55 +0,0 @@ -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_leditspp_stable_diffusion"] = ["LEditsPPPipelineStableDiffusion"] - _import_structure["pipeline_leditspp_stable_diffusion_xl"] = ["LEditsPPPipelineStableDiffusionXL"] - - _import_structure["pipeline_output"] = ["LEditsPPDiffusionPipelineOutput", "LEditsPPDiffusionPipelineOutput"] - -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_leditspp_stable_diffusion import ( - LEditsPPDiffusionPipelineOutput, - LEditsPPInversionPipelineOutput, - LEditsPPPipelineStableDiffusion, - ) - from .pipeline_leditspp_stable_diffusion_xl import LEditsPPPipelineStableDiffusionXL - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 80105297402f9beac902f5a5dfcb0ac5fe01c923..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__pycache__/pipeline_leditspp_stable_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__pycache__/pipeline_leditspp_stable_diffusion.cpython-310.pyc deleted file mode 100644 index cd952ca7c97b7ee801ebf08cfe0f94d3071efa6b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__pycache__/pipeline_leditspp_stable_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__pycache__/pipeline_leditspp_stable_diffusion_xl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__pycache__/pipeline_leditspp_stable_diffusion_xl.cpython-310.pyc deleted file mode 100644 index eaa981d3a94fc6337652d35645f961fc55c6df7f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__pycache__/pipeline_leditspp_stable_diffusion_xl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index 8fe6567cec9aebb09d9b52cc222b9fbca8b43e21..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py deleted file mode 100644 index 341ccabaa146280cf503d41fa85844b968e6db6a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py +++ /dev/null @@ -1,1565 +0,0 @@ -import inspect -import math -from itertools import repeat -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -import torch.nn.functional as F -from packaging import version -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer - -from ...configuration_utils import FrozenDict -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, UNet2DConditionModel -from ...models.attention_processor import Attention, AttnProcessor -from ...models.lora import adjust_lora_scale_text_encoder -from ...pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker -from ...schedulers import DDIMScheduler, DPMSolverMultistepScheduler -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - 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 LEditsPPDiffusionPipelineOutput, LEditsPPInversionPipelineOutput - - -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 LEditsPPPipelineStableDiffusion - >>> from diffusers.utils import load_image - - >>> pipe = LEditsPPPipelineStableDiffusion.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", variant="fp16", torch_dtype=torch.float16 - ... ) - >>> pipe.enable_vae_tiling() - >>> pipe = pipe.to("cuda") - - >>> img_url = "https://www.aiml.informatik.tu-darmstadt.de/people/mbrack/cherry_blossom.png" - >>> image = load_image(img_url).resize((512, 512)) - - >>> _ = pipe.invert(image=image, num_inversion_steps=50, skip=0.1) - - >>> edited_image = pipe( - ... editing_prompt=["cherry blossom"], edit_guidance_scale=10.0, edit_threshold=0.75 - ... ).images[0] - ``` -""" - - -# Modified from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionAttendAndExcitePipeline.AttentionStore -class LeditsAttentionStore: - @staticmethod - def get_empty_store(): - return {"down_cross": [], "mid_cross": [], "up_cross": [], "down_self": [], "mid_self": [], "up_self": []} - - def __call__(self, attn, is_cross: bool, place_in_unet: str, editing_prompts, PnP=False): - # attn.shape = batch_size * head_size, seq_len query, seq_len_key - if attn.shape[1] <= self.max_size: - bs = 1 + int(PnP) + editing_prompts - skip = 2 if PnP else 1 # skip PnP & unconditional - attn = torch.stack(attn.split(self.batch_size)).permute(1, 0, 2, 3) - source_batch_size = int(attn.shape[1] // bs) - self.forward(attn[:, skip * source_batch_size :], is_cross, place_in_unet) - - def forward(self, attn, is_cross: bool, place_in_unet: str): - key = f"{place_in_unet}_{'cross' if is_cross else 'self'}" - - self.step_store[key].append(attn) - - def between_steps(self, store_step=True): - if store_step: - if self.average: - if len(self.attention_store) == 0: - self.attention_store = self.step_store - else: - for key in self.attention_store: - for i in range(len(self.attention_store[key])): - self.attention_store[key][i] += self.step_store[key][i] - else: - if len(self.attention_store) == 0: - self.attention_store = [self.step_store] - else: - self.attention_store.append(self.step_store) - - self.cur_step += 1 - self.step_store = self.get_empty_store() - - def get_attention(self, step: int): - if self.average: - attention = { - key: [item / self.cur_step for item in self.attention_store[key]] for key in self.attention_store - } - else: - assert step is not None - attention = self.attention_store[step] - return attention - - def aggregate_attention( - self, attention_maps, prompts, res: Union[int, Tuple[int]], from_where: List[str], is_cross: bool, select: int - ): - out = [[] for x in range(self.batch_size)] - if isinstance(res, int): - num_pixels = res**2 - resolution = (res, res) - else: - num_pixels = res[0] * res[1] - resolution = res[:2] - - for location in from_where: - for bs_item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]: - for batch, item in enumerate(bs_item): - if item.shape[1] == num_pixels: - cross_maps = item.reshape(len(prompts), -1, *resolution, item.shape[-1])[select] - out[batch].append(cross_maps) - - out = torch.stack([torch.cat(x, dim=0) for x in out]) - # average over heads - out = out.sum(1) / out.shape[1] - return out - - def __init__(self, average: bool, batch_size=1, max_resolution=16, max_size: int = None): - self.step_store = self.get_empty_store() - self.attention_store = [] - self.cur_step = 0 - self.average = average - self.batch_size = batch_size - if max_size is None: - self.max_size = max_resolution**2 - elif max_size is not None and max_resolution is None: - self.max_size = max_size - else: - raise ValueError("Only allowed to set one of max_resolution or max_size") - - -# Modified from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionAttendAndExcitePipeline.GaussianSmoothing -class LeditsGaussianSmoothing: - def __init__(self, device): - kernel_size = [3, 3] - sigma = [0.5, 0.5] - - # The gaussian kernel is the product of the gaussian function of each dimension. - kernel = 1 - meshgrids = torch.meshgrid([torch.arange(size, dtype=torch.float32) for size in kernel_size], indexing="ij") - for size, std, mgrid in zip(kernel_size, sigma, meshgrids): - mean = (size - 1) / 2 - kernel *= 1 / (std * math.sqrt(2 * math.pi)) * torch.exp(-(((mgrid - mean) / (2 * std)) ** 2)) - - # Make sure sum of values in gaussian kernel equals 1. - kernel = kernel / torch.sum(kernel) - - # Reshape to depthwise convolutional weight - kernel = kernel.view(1, 1, *kernel.size()) - kernel = kernel.repeat(1, *[1] * (kernel.dim() - 1)) - - self.weight = kernel.to(device) - - def __call__(self, input): - """ - Arguments: - Apply gaussian filter to input. - input (torch.Tensor): Input to apply gaussian filter on. - Returns: - filtered (torch.Tensor): Filtered output. - """ - return F.conv2d(input, weight=self.weight.to(input.dtype)) - - -class LEDITSCrossAttnProcessor: - def __init__(self, attention_store, place_in_unet, pnp, editing_prompts): - self.attnstore = attention_store - self.place_in_unet = place_in_unet - self.editing_prompts = editing_prompts - self.pnp = pnp - - def __call__( - self, - attn: Attention, - hidden_states, - encoder_hidden_states, - attention_mask=None, - temb=None, - ): - batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query = attn.head_to_batch_dim(query) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - - attention_probs = attn.get_attention_scores(query, key, attention_mask) - self.attnstore( - attention_probs, - is_cross=True, - place_in_unet=self.place_in_unet, - editing_prompts=self.editing_prompts, - PnP=self.pnp, - ) - - hidden_states = torch.bmm(attention_probs, value) - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - hidden_states = hidden_states / attn.rescale_output_factor - return hidden_states - - -# 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 LEditsPPPipelineStableDiffusion( - DiffusionPipeline, TextualInversionLoaderMixin, StableDiffusionLoraLoaderMixin, IPAdapterMixin, FromSingleFileMixin -): - """ - Pipeline for textual image editing using LEDits++ with Stable Diffusion. - - This model inherits from [`DiffusionPipeline`] and builds on the [`StableDiffusionPipeline`]. Check the superclass - documentation for the generic methods implemented for all pipelines (downloading, saving, running on a particular - device, etc.). - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.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 ([`~transformers.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 ([`DPMSolverMultistepScheduler`] or [`DDIMScheduler`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latens. Can be one of - [`DPMSolverMultistepScheduler`] or [`DDIMScheduler`]. If any other scheduler is passed it will - automatically be set to [`DPMSolverMultistepScheduler`]. - 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/CompVis/stable-diffusion-v1-4) for details. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - Model that extracts features from generated images to be used as inputs for the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, DPMSolverMultistepScheduler], - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - if not isinstance(scheduler, DDIMScheduler) and not isinstance(scheduler, DPMSolverMultistepScheduler): - scheduler = DPMSolverMultistepScheduler.from_config( - scheduler.config, algorithm_type="sde-dpmsolver++", solver_order=2 - ) - logger.warning( - "This pipeline only supports DDIMScheduler and DPMSolverMultistepScheduler. " - "The scheduler has been changed to DPMSolverMultistepScheduler." - ) - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_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.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.register_to_config(requires_safety_checker=requires_safety_checker) - - self.inversion_steps = None - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, eta, generator=None): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Modified from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.check_inputs - def check_inputs( - self, - negative_prompt=None, - editing_prompt_embeddings=None, - negative_prompt_embeds=None, - callback_on_step_end_tensor_inputs=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 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 editing_prompt_embeddings is not None and negative_prompt_embeds is not None: - if editing_prompt_embeddings.shape != negative_prompt_embeds.shape: - raise ValueError( - "`editing_prompt_embeddings` and `negative_prompt_embeds` must have the same shape when passed directly, but" - f" got: `editing_prompt_embeddings` {editing_prompt_embeddings.shape} != `negative_prompt_embeds`" - f" {negative_prompt_embeds.shape}." - ) - - # Modified from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, latents): - # shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor) - - # 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 - latents = latents * self.scheduler.init_noise_sigma - return latents - - def prepare_unet(self, attention_store, PnP: bool = False): - attn_procs = {} - for name in self.unet.attn_processors.keys(): - if name.startswith("mid_block"): - place_in_unet = "mid" - elif name.startswith("up_blocks"): - place_in_unet = "up" - elif name.startswith("down_blocks"): - place_in_unet = "down" - else: - continue - - if "attn2" in name and place_in_unet != "mid": - attn_procs[name] = LEDITSCrossAttnProcessor( - attention_store=attention_store, - place_in_unet=place_in_unet, - pnp=PnP, - editing_prompts=self.enabled_editing_prompts, - ) - else: - attn_procs[name] = AttnProcessor() - - self.unet.set_attn_processor(attn_procs) - - def encode_prompt( - self, - device, - num_images_per_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, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - 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*): - 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(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 - 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) - - batch_size = self.batch_size - num_edit_tokens = None - - if negative_prompt_embeds is None: - uncond_tokens: List[str] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - 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 exoected" - f"{batch_size} based on the input images. Please make sure that passed `negative_prompt` matches" - " the batch size of `prompt`." - ) - else: - uncond_tokens = negative_prompt - - # textual inversion: procecss multi-vector tokens if necessary - if isinstance(self, TextualInversionLoaderMixin): - uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer) - - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=self.tokenizer.model_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 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 = negative_prompt_embeds.dtype - - negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device) - - if enable_edit_guidance: - if editing_prompt_embeds is None: - # textual inversion: procecss multi-vector tokens if necessary - # if isinstance(self, TextualInversionLoaderMixin): - # prompt = self.maybe_convert_prompt(prompt, self.tokenizer) - if isinstance(editing_prompt, str): - editing_prompt = [editing_prompt] - - max_length = negative_prompt_embeds.shape[1] - text_inputs = self.tokenizer( - [x for item in editing_prompt for x in repeat(item, batch_size)], - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - return_length=True, - ) - - num_edit_tokens = text_inputs.length - 2 # not counting startoftext and endoftext - text_input_ids = text_inputs.input_ids - untruncated_ids = self.tokenizer( - [x for item in editing_prompt for x in repeat(item, batch_size)], - 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: - editing_prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask) - editing_prompt_embeds = editing_prompt_embeds[0] - else: - editing_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. - editing_prompt_embeds = editing_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. - editing_prompt_embeds = self.text_encoder.text_model.final_layer_norm(editing_prompt_embeds) - - editing_prompt_embeds = editing_prompt_embeds.to(dtype=negative_prompt_embeds.dtype, device=device) - - bs_embed_edit, seq_len, _ = editing_prompt_embeds.shape - editing_prompt_embeds = editing_prompt_embeds.to(dtype=negative_prompt_embeds.dtype, device=device) - editing_prompt_embeds = editing_prompt_embeds.repeat(1, num_images_per_prompt, 1) - editing_prompt_embeds = editing_prompt_embeds.view(bs_embed_edit * num_images_per_prompt, seq_len, -1) - - # get unconditional embeddings for 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 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 editing_prompt_embeds, negative_prompt_embeds, num_edit_tokens - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @property - def clip_skip(self): - return self._clip_skip - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - @torch.no_grad() - @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", - 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, - use_cross_attn_mask: bool = False, - use_intersect_mask: bool = True, - attn_store_steps: Optional[List[int]] = [], - store_averaged_over_steps: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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"], - **kwargs, - ): - r""" - The call function to the pipeline for editing. The - [`~pipelines.ledits_pp.LEditsPPPipelineStableDiffusion.invert`] method has to be called beforehand. Edits will - always be performed for the last inverted image(s). - - Args: - 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*): - 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.ledits_pp.LEditsPPDiffusionPipelineOutput`] instead of a plain - tuple. - 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`. - editing_prompt_embeds (`torch.Tensor>`, *optional*): - Pre-computed embeddings to use for guiding the image generation. Guidance direction of embedding should - be specified via `reverse_editing_direction`. - 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`): - Whether the corresponding prompt in `editing_prompt` should be increased or decreased. - 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): - 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`): - 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): - 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 - 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 - is set to true. Cross-attention masks are defined as 'M^1' of equation 12 of [LEDITS++ - paper](https://huggingface.co/papers/2311.16711). - use_intersect_mask (`bool`, defaults to `True`): - 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*): - 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 - False, attention maps for each step are stores separately. Just for visualization purposes. - 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). - guidance_rescale (`float`, *optional*, defaults to 0.0): - Guidance rescale factor from [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. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.ledits_pp.LEditsPPDiffusionPipelineOutput`] or `tuple`: - [`~pipelines.ledits_pp.LEditsPPDiffusionPipelineOutput`] 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`. - """ - - if self.inversion_steps is None: - raise ValueError( - "You need to invert an input image first before calling the pipeline. The `invert` method has to be called beforehand. Edits will always be performed for the last inverted image(s)." - ) - - eta = self.eta - num_images_per_prompt = 1 - latents = self.init_latents - - zs = self.zs - self.scheduler.set_timesteps(len(self.scheduler.timesteps)) - - if use_intersect_mask: - use_cross_attn_mask = True - - if use_cross_attn_mask: - self.smoothing = LeditsGaussianSmoothing(self.device) - - if user_mask is not None: - user_mask = user_mask.to(self.device) - - org_prompt = "" - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - negative_prompt, - editing_prompt_embeds, - negative_prompt_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - - # 2. Define call parameters - batch_size = self.batch_size - - if editing_prompt: - enable_edit_guidance = True - if isinstance(editing_prompt, str): - editing_prompt = [editing_prompt] - self.enabled_editing_prompts = len(editing_prompt) - elif editing_prompt_embeds is not None: - enable_edit_guidance = True - self.enabled_editing_prompts = editing_prompt_embeds.shape[0] - else: - self.enabled_editing_prompts = 0 - enable_edit_guidance = False - - # 3. Encode input prompt - lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - - edit_concepts, uncond_embeddings, num_edit_tokens = self.encode_prompt( - editing_prompt=editing_prompt, - device=self.device, - num_images_per_prompt=num_images_per_prompt, - enable_edit_guidance=enable_edit_guidance, - negative_prompt=negative_prompt, - editing_prompt_embeds=editing_prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=lora_scale, - clip_skip=self.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 enable_edit_guidance: - text_embeddings = torch.cat([uncond_embeddings, edit_concepts]) - self.text_cross_attention_maps = [editing_prompt] if isinstance(editing_prompt, str) else editing_prompt - else: - text_embeddings = torch.cat([uncond_embeddings]) - - # 4. Prepare timesteps - # self.scheduler.set_timesteps(num_inference_steps, device=self.device) - timesteps = self.inversion_steps - t_to_idx = {int(v): k for k, v in enumerate(timesteps[-zs.shape[0] :])} - - if use_cross_attn_mask: - self.attention_store = LeditsAttentionStore( - average=store_averaged_over_steps, - batch_size=batch_size, - max_size=(latents.shape[-2] / 4.0) * (latents.shape[-1] / 4.0), - max_resolution=None, - ) - self.prepare_unet(self.attention_store, PnP=False) - resolution = latents.shape[-2:] - att_res = (int(resolution[0] / 4), int(resolution[1] / 4)) - - # 5. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - None, - None, - text_embeddings.dtype, - self.device, - latents, - ) - - # 6. Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(eta) - - self.sem_guidance = None - self.activation_mask = None - - # 7. Denoising loop - num_warmup_steps = 0 - with self.progress_bar(total=len(timesteps)) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - - if enable_edit_guidance: - latent_model_input = torch.cat([latents] * (1 + self.enabled_editing_prompts)) - else: - latent_model_input = latents - - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - text_embed_input = text_embeddings - - # predict the noise residual - noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embed_input).sample - - noise_pred_out = noise_pred.chunk(1 + self.enabled_editing_prompts) # [b,4, 64, 64] - noise_pred_uncond = noise_pred_out[0] - noise_pred_edit_concepts = noise_pred_out[1:] - - noise_guidance_edit = torch.zeros( - noise_pred_uncond.shape, - device=self.device, - dtype=noise_pred_uncond.dtype, - ) - - if sem_guidance is not None and len(sem_guidance) > i: - noise_guidance_edit += sem_guidance[i].to(self.device) - - elif enable_edit_guidance: - if self.activation_mask is None: - self.activation_mask = torch.zeros( - (len(timesteps), len(noise_pred_edit_concepts), *noise_pred_edit_concepts[0].shape) - ) - - if self.sem_guidance is None: - self.sem_guidance = torch.zeros((len(timesteps), *noise_pred_uncond.shape)) - - for c, noise_pred_edit_concept in enumerate(noise_pred_edit_concepts): - if isinstance(edit_warmup_steps, list): - edit_warmup_steps_c = edit_warmup_steps[c] - else: - edit_warmup_steps_c = edit_warmup_steps - if i < edit_warmup_steps_c: - continue - - if isinstance(edit_guidance_scale, list): - edit_guidance_scale_c = edit_guidance_scale[c] - else: - edit_guidance_scale_c = edit_guidance_scale - - if isinstance(edit_threshold, list): - edit_threshold_c = edit_threshold[c] - else: - edit_threshold_c = edit_threshold - if isinstance(reverse_editing_direction, list): - reverse_editing_direction_c = reverse_editing_direction[c] - else: - reverse_editing_direction_c = reverse_editing_direction - - if isinstance(edit_cooldown_steps, list): - edit_cooldown_steps_c = edit_cooldown_steps[c] - elif edit_cooldown_steps is None: - edit_cooldown_steps_c = i + 1 - else: - edit_cooldown_steps_c = edit_cooldown_steps - - if i >= edit_cooldown_steps_c: - continue - - noise_guidance_edit_tmp = noise_pred_edit_concept - noise_pred_uncond - - if reverse_editing_direction_c: - noise_guidance_edit_tmp = noise_guidance_edit_tmp * -1 - - noise_guidance_edit_tmp = noise_guidance_edit_tmp * edit_guidance_scale_c - - if user_mask is not None: - noise_guidance_edit_tmp = noise_guidance_edit_tmp * user_mask - - if use_cross_attn_mask: - out = self.attention_store.aggregate_attention( - attention_maps=self.attention_store.step_store, - prompts=self.text_cross_attention_maps, - res=att_res, - from_where=["up", "down"], - is_cross=True, - select=self.text_cross_attention_maps.index(editing_prompt[c]), - ) - attn_map = out[:, :, :, 1 : 1 + num_edit_tokens[c]] # 0 -> startoftext - - # average over all tokens - if attn_map.shape[3] != num_edit_tokens[c]: - raise ValueError( - f"Incorrect shape of attention_map. Expected size {num_edit_tokens[c]}, but found {attn_map.shape[3]}!" - ) - - attn_map = torch.sum(attn_map, dim=3) - - # gaussian_smoothing - attn_map = F.pad(attn_map.unsqueeze(1), (1, 1, 1, 1), mode="reflect") - attn_map = self.smoothing(attn_map).squeeze(1) - - # torch.quantile function expects float32 - if attn_map.dtype == torch.float32: - tmp = torch.quantile(attn_map.flatten(start_dim=1), edit_threshold_c, dim=1) - else: - tmp = torch.quantile( - attn_map.flatten(start_dim=1).to(torch.float32), edit_threshold_c, dim=1 - ).to(attn_map.dtype) - attn_mask = torch.where( - attn_map >= tmp.unsqueeze(1).unsqueeze(1).repeat(1, *att_res), 1.0, 0.0 - ) - - # resolution must match latent space dimension - attn_mask = F.interpolate( - attn_mask.unsqueeze(1), - noise_guidance_edit_tmp.shape[-2:], # 64,64 - ).repeat(1, 4, 1, 1) - self.activation_mask[i, c] = attn_mask.detach().cpu() - if not use_intersect_mask: - noise_guidance_edit_tmp = noise_guidance_edit_tmp * attn_mask - - if use_intersect_mask: - if t <= 800: - noise_guidance_edit_tmp_quantile = torch.abs(noise_guidance_edit_tmp) - noise_guidance_edit_tmp_quantile = torch.sum( - noise_guidance_edit_tmp_quantile, dim=1, keepdim=True - ) - noise_guidance_edit_tmp_quantile = noise_guidance_edit_tmp_quantile.repeat( - 1, self.unet.config.in_channels, 1, 1 - ) - - # torch.quantile function expects float32 - if noise_guidance_edit_tmp_quantile.dtype == torch.float32: - tmp = torch.quantile( - noise_guidance_edit_tmp_quantile.flatten(start_dim=2), - edit_threshold_c, - dim=2, - keepdim=False, - ) - else: - tmp = torch.quantile( - noise_guidance_edit_tmp_quantile.flatten(start_dim=2).to(torch.float32), - edit_threshold_c, - dim=2, - keepdim=False, - ).to(noise_guidance_edit_tmp_quantile.dtype) - - intersect_mask = ( - torch.where( - noise_guidance_edit_tmp_quantile >= tmp[:, :, None, None], - torch.ones_like(noise_guidance_edit_tmp), - torch.zeros_like(noise_guidance_edit_tmp), - ) - * attn_mask - ) - - self.activation_mask[i, c] = intersect_mask.detach().cpu() - - noise_guidance_edit_tmp = noise_guidance_edit_tmp * intersect_mask - - else: - # print(f"only attention mask for step {i}") - noise_guidance_edit_tmp = noise_guidance_edit_tmp * attn_mask - - elif not use_cross_attn_mask: - # calculate quantile - noise_guidance_edit_tmp_quantile = torch.abs(noise_guidance_edit_tmp) - noise_guidance_edit_tmp_quantile = torch.sum( - noise_guidance_edit_tmp_quantile, dim=1, keepdim=True - ) - noise_guidance_edit_tmp_quantile = noise_guidance_edit_tmp_quantile.repeat(1, 4, 1, 1) - - # torch.quantile function expects float32 - if noise_guidance_edit_tmp_quantile.dtype == torch.float32: - tmp = torch.quantile( - noise_guidance_edit_tmp_quantile.flatten(start_dim=2), - edit_threshold_c, - dim=2, - keepdim=False, - ) - else: - tmp = torch.quantile( - noise_guidance_edit_tmp_quantile.flatten(start_dim=2).to(torch.float32), - edit_threshold_c, - dim=2, - keepdim=False, - ).to(noise_guidance_edit_tmp_quantile.dtype) - - self.activation_mask[i, c] = ( - torch.where( - noise_guidance_edit_tmp_quantile >= tmp[:, :, None, None], - torch.ones_like(noise_guidance_edit_tmp), - torch.zeros_like(noise_guidance_edit_tmp), - ) - .detach() - .cpu() - ) - - noise_guidance_edit_tmp = torch.where( - noise_guidance_edit_tmp_quantile >= tmp[:, :, None, None], - noise_guidance_edit_tmp, - torch.zeros_like(noise_guidance_edit_tmp), - ) - - noise_guidance_edit += noise_guidance_edit_tmp - - self.sem_guidance[i] = noise_guidance_edit.detach().cpu() - - noise_pred = noise_pred_uncond + noise_guidance_edit - - if enable_edit_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg( - noise_pred, - noise_pred_edit_concepts.mean(dim=0, keepdim=False), - guidance_rescale=self.guidance_rescale, - ) - - idx = t_to_idx[int(t)] - latents = self.scheduler.step( - noise_pred, t, latents, variance_noise=zs[idx], **extra_step_kwargs - ).prev_sample - - # step callback - if use_cross_attn_mask: - store_step = i in attn_store_steps - self.attention_store.between_steps(store_step) - - 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() - - # 8. Post-processing - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 0 - ] - image, has_nsfw_concept = self.run_safety_checker(image, self.device, text_embeddings.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 LEditsPPDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) - - @torch.no_grad() - def invert( - self, - image: PipelineImageInput, - source_prompt: str = "", - 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, - ): - r""" - The function to the pipeline for image inversion as described by the [LEDITS++ - Paper](https://huggingface.co/papers/2301.12247). If the scheduler is set to [`~schedulers.DDIMScheduler`] the - inversion proposed by [edit-friendly DPDM](https://huggingface.co/papers/2304.06140) will be performed instead. - - Args: - image (`PipelineImageInput`): - Input for the image(s) that are to be edited. Multiple input images have to default to the same aspect - ratio. - source_prompt (`str`, defaults to `""`): - Prompt describing the input image that will be used for guidance during inversion. Guidance is disabled - if the `source_prompt` is `""`. - source_guidance_scale (`float`, defaults to `3.5`): - Strength of guidance during inversion. - num_inversion_steps (`int`, defaults to `30`): - Number of total performed inversion steps after discarding the initial `skip` steps. - skip (`float`, defaults to `0.15`): - Portion of initial steps that will be ignored for inversion and subsequent generation. Lower values - will lead to stronger changes to the input image. `skip` has to be between `0` and `1`. - generator (`torch.Generator`, *optional*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make inversion - deterministic. - 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). - 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. - height (`int`, *optional*, defaults to `None`): - The height in preprocessed image. If `None`, will use the `get_default_height_width()` to get default - height. - width (`int`, *optional*`, defaults to `None`): - The width in preprocessed. If `None`, will use get_default_height_width()` to get the default width. - resize_mode (`str`, *optional*, defaults to `default`): - The resize mode, can be one of `default` or `fill`. If `default`, will resize the image to fit within - the specified width and height, and it may not maintaining the original aspect ratio. If `fill`, will - resize the image to fit within the specified width and height, maintaining the aspect ratio, and then - center the image within the dimensions, filling empty with data from image. If `crop`, will resize the - 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`): - The crop coordinates for each image in the batch. If `None`, will not crop the image. - - Returns: - [`~pipelines.ledits_pp.LEditsPPInversionPipelineOutput`]: Output will contain the resized input image(s) - and respective VAE reconstruction(s). - """ - if height is not None and height % 32 != 0 or width is not None and width % 32 != 0: - raise ValueError("height and width must be a factor of 32.") - # Reset attn processor, we do not want to store attn maps during inversion - self.unet.set_attn_processor(AttnProcessor()) - - self.eta = 1.0 - - self.scheduler.config.timestep_spacing = "leading" - self.scheduler.set_timesteps(int(num_inversion_steps * (1 + skip))) - self.inversion_steps = self.scheduler.timesteps[-num_inversion_steps:] - timesteps = self.inversion_steps - - # 1. encode image - x0, resized = self.encode_image( - image, - dtype=self.text_encoder.dtype, - height=height, - width=width, - resize_mode=resize_mode, - crops_coords=crops_coords, - ) - self.batch_size = x0.shape[0] - - # autoencoder reconstruction - image_rec = self.vae.decode(x0 / self.vae.config.scaling_factor, return_dict=False, generator=generator)[0] - image_rec = self.image_processor.postprocess(image_rec, output_type="pil") - - # 2. get embeddings - do_classifier_free_guidance = source_guidance_scale > 1.0 - - lora_scale = cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - - uncond_embedding, text_embeddings, _ = self.encode_prompt( - num_images_per_prompt=1, - device=self.device, - negative_prompt=None, - enable_edit_guidance=do_classifier_free_guidance, - editing_prompt=source_prompt, - lora_scale=lora_scale, - clip_skip=clip_skip, - ) - - # 3. find zs and xts - variance_noise_shape = (num_inversion_steps, *x0.shape) - - # intermediate latents - t_to_idx = {int(v): k for k, v in enumerate(timesteps)} - xts = torch.zeros(size=variance_noise_shape, device=self.device, dtype=uncond_embedding.dtype) - - for t in reversed(timesteps): - idx = num_inversion_steps - t_to_idx[int(t)] - 1 - noise = randn_tensor(shape=x0.shape, generator=generator, device=self.device, dtype=x0.dtype) - xts[idx] = self.scheduler.add_noise(x0, noise, torch.Tensor([t])) - xts = torch.cat([x0.unsqueeze(0), xts], dim=0) - - self.scheduler.set_timesteps(len(self.scheduler.timesteps)) - # noise maps - zs = torch.zeros(size=variance_noise_shape, device=self.device, dtype=uncond_embedding.dtype) - - with self.progress_bar(total=len(timesteps)) as progress_bar: - for t in timesteps: - idx = num_inversion_steps - t_to_idx[int(t)] - 1 - # 1. predict noise residual - xt = xts[idx + 1] - - noise_pred = self.unet(xt, timestep=t, encoder_hidden_states=uncond_embedding).sample - - if not source_prompt == "": - noise_pred_cond = self.unet(xt, timestep=t, encoder_hidden_states=text_embeddings).sample - noise_pred = noise_pred + source_guidance_scale * (noise_pred_cond - noise_pred) - - xtm1 = xts[idx] - z, xtm1_corrected = compute_noise(self.scheduler, xtm1, xt, t, noise_pred, self.eta) - zs[idx] = z - - # correction to avoid error accumulation - xts[idx] = xtm1_corrected - - progress_bar.update() - - if XLA_AVAILABLE: - xm.mark_step() - - self.init_latents = xts[-1].expand(self.batch_size, -1, -1, -1) - zs = zs.flip(0) - self.zs = zs - - return LEditsPPInversionPipelineOutput(images=resized, vae_reconstruction_images=image_rec) - - @torch.no_grad() - def encode_image(self, image, dtype=None, height=None, width=None, resize_mode="default", crops_coords=None): - image = self.image_processor.preprocess( - image=image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords - ) - height, width = image.shape[-2:] - if height % 32 != 0 or width % 32 != 0: - raise ValueError( - "Image height and width must be a factor of 32. " - "Consider down-sampling the input using the `height` and `width` parameters" - ) - resized = self.image_processor.postprocess(image=image, output_type="pil") - - if max(image.shape[-2:]) > self.vae.config["sample_size"] * 1.5: - logger.warning( - "Your input images far exceed the default resolution of the underlying diffusion model. " - "The output images may contain severe artifacts! " - "Consider down-sampling the input using the `height` and `width` parameters" - ) - image = image.to(dtype) - - x0 = self.vae.encode(image.to(self.device)).latent_dist.mode() - x0 = x0.to(dtype) - x0 = self.vae.config.scaling_factor * x0 - return x0, resized - - -def compute_noise_ddim(scheduler, prev_latents, latents, timestep, noise_pred, eta): - # 1. get previous step value (=t-1) - prev_timestep = timestep - scheduler.config.num_train_timesteps // scheduler.num_inference_steps - - # 2. compute alphas, betas - alpha_prod_t = scheduler.alphas_cumprod[timestep] - alpha_prod_t_prev = ( - scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else scheduler.final_alpha_cumprod - ) - - beta_prod_t = 1 - alpha_prod_t - - # 3. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (12) from https://huggingface.co/papers/2010.02502 - pred_original_sample = (latents - beta_prod_t ** (0.5) * noise_pred) / alpha_prod_t ** (0.5) - - # 4. Clip "predicted x_0" - if scheduler.config.clip_sample: - pred_original_sample = torch.clamp(pred_original_sample, -1, 1) - - # 5. compute variance: "sigma_t(η)" -> see formula (16) - # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1) - variance = scheduler._get_variance(timestep, prev_timestep) - std_dev_t = eta * variance ** (0.5) - - # 6. compute "direction pointing to x_t" of formula (12) from https://huggingface.co/papers/2010.02502 - pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * noise_pred - - # modified so that updated xtm1 is returned as well (to avoid error accumulation) - mu_xt = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction - if variance > 0.0: - noise = (prev_latents - mu_xt) / (variance ** (0.5) * eta) - else: - noise = torch.tensor([0.0]).to(latents.device) - - return noise, mu_xt + (eta * variance**0.5) * noise - - -def compute_noise_sde_dpm_pp_2nd(scheduler, prev_latents, latents, timestep, noise_pred, eta): - def first_order_update(model_output, sample): # timestep, prev_timestep, sample): - sigma_t, sigma_s = scheduler.sigmas[scheduler.step_index + 1], scheduler.sigmas[scheduler.step_index] - alpha_t, sigma_t = scheduler._sigma_to_alpha_sigma_t(sigma_t) - alpha_s, sigma_s = scheduler._sigma_to_alpha_sigma_t(sigma_s) - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s = torch.log(alpha_s) - torch.log(sigma_s) - - h = lambda_t - lambda_s - - mu_xt = (sigma_t / sigma_s * torch.exp(-h)) * sample + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output - - mu_xt = scheduler.dpm_solver_first_order_update( - model_output=model_output, sample=sample, noise=torch.zeros_like(sample) - ) - - sigma = sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) - if sigma > 0.0: - noise = (prev_latents - mu_xt) / sigma - else: - noise = torch.tensor([0.0]).to(sample.device) - - prev_sample = mu_xt + sigma * noise - return noise, prev_sample - - def second_order_update(model_output_list, sample): # timestep_list, prev_timestep, sample): - sigma_t, sigma_s0, sigma_s1 = ( - scheduler.sigmas[scheduler.step_index + 1], - scheduler.sigmas[scheduler.step_index], - scheduler.sigmas[scheduler.step_index - 1], - ) - - alpha_t, sigma_t = scheduler._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = scheduler._sigma_to_alpha_sigma_t(sigma_s0) - alpha_s1, sigma_s1 = scheduler._sigma_to_alpha_sigma_t(sigma_s1) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1) - - m0, m1 = model_output_list[-1], model_output_list[-2] - - h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1 - r0 = h_0 / h - D0, D1 = m0, (1.0 / r0) * (m0 - m1) - - mu_xt = ( - (sigma_t / sigma_s0 * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 - + 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1 - ) - - sigma = sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) - if sigma > 0.0: - noise = (prev_latents - mu_xt) / sigma - else: - noise = torch.tensor([0.0]).to(sample.device) - - prev_sample = mu_xt + sigma * noise - - return noise, prev_sample - - if scheduler.step_index is None: - scheduler._init_step_index(timestep) - - model_output = scheduler.convert_model_output(model_output=noise_pred, sample=latents) - for i in range(scheduler.config.solver_order - 1): - scheduler.model_outputs[i] = scheduler.model_outputs[i + 1] - scheduler.model_outputs[-1] = model_output - - if scheduler.lower_order_nums < 1: - noise, prev_sample = first_order_update(model_output, latents) - else: - noise, prev_sample = second_order_update(scheduler.model_outputs, latents) - - if scheduler.lower_order_nums < scheduler.config.solver_order: - scheduler.lower_order_nums += 1 - - # upon completion increase step index by one - scheduler._step_index += 1 - - return noise, prev_sample - - -def compute_noise(scheduler, *args): - if isinstance(scheduler, DDIMScheduler): - return compute_noise_ddim(scheduler, *args) - elif ( - isinstance(scheduler, DPMSolverMultistepScheduler) - and scheduler.config.algorithm_type == "sde-dpmsolver++" - and scheduler.config.solver_order == 2 - ): - return compute_noise_sde_dpm_pp_2nd(scheduler, *args) - else: - raise NotImplementedError diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py deleted file mode 100644 index ac64844f6fee12e2079984831ce7c809b78cb8d7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py +++ /dev/null @@ -1,1852 +0,0 @@ -# Copyright 2023 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 math -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -import torch.nn.functional as F -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, UNet2DConditionModel -from ...models.attention_processor import ( - Attention, - AttnProcessor, - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import DDIMScheduler, DPMSolverMultistepScheduler -from ...utils import ( - USE_PEFT_BACKEND, - is_invisible_watermark_available, - 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 LEditsPPDiffusionPipelineOutput, LEditsPPInversionPipelineOutput - - -if is_invisible_watermark_available(): - from ..stable_diffusion_xl.watermark import StableDiffusionXLWatermarker - -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 LEditsPPPipelineStableDiffusionXL - >>> from diffusers.utils import load_image - - >>> pipe = LEditsPPPipelineStableDiffusionXL.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", variant="fp16", torch_dtype=torch.float16 - ... ) - >>> pipe.enable_vae_tiling() - >>> pipe = pipe.to("cuda") - - >>> img_url = "https://www.aiml.informatik.tu-darmstadt.de/people/mbrack/tennis.jpg" - >>> image = load_image(img_url).resize((1024, 1024)) - - >>> _ = pipe.invert(image=image, num_inversion_steps=50, skip=0.2) - - >>> edited_image = pipe( - ... editing_prompt=["tennis ball", "tomato"], - ... reverse_editing_direction=[True, False], - ... edit_guidance_scale=[5.0, 10.0], - ... edit_threshold=[0.9, 0.85], - ... ).images[0] - ``` -""" - - -# Copied from diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.LeditsAttentionStore -class LeditsAttentionStore: - @staticmethod - def get_empty_store(): - return {"down_cross": [], "mid_cross": [], "up_cross": [], "down_self": [], "mid_self": [], "up_self": []} - - def __call__(self, attn, is_cross: bool, place_in_unet: str, editing_prompts, PnP=False): - # attn.shape = batch_size * head_size, seq_len query, seq_len_key - if attn.shape[1] <= self.max_size: - bs = 1 + int(PnP) + editing_prompts - skip = 2 if PnP else 1 # skip PnP & unconditional - attn = torch.stack(attn.split(self.batch_size)).permute(1, 0, 2, 3) - source_batch_size = int(attn.shape[1] // bs) - self.forward(attn[:, skip * source_batch_size :], is_cross, place_in_unet) - - def forward(self, attn, is_cross: bool, place_in_unet: str): - key = f"{place_in_unet}_{'cross' if is_cross else 'self'}" - - self.step_store[key].append(attn) - - def between_steps(self, store_step=True): - if store_step: - if self.average: - if len(self.attention_store) == 0: - self.attention_store = self.step_store - else: - for key in self.attention_store: - for i in range(len(self.attention_store[key])): - self.attention_store[key][i] += self.step_store[key][i] - else: - if len(self.attention_store) == 0: - self.attention_store = [self.step_store] - else: - self.attention_store.append(self.step_store) - - self.cur_step += 1 - self.step_store = self.get_empty_store() - - def get_attention(self, step: int): - if self.average: - attention = { - key: [item / self.cur_step for item in self.attention_store[key]] for key in self.attention_store - } - else: - assert step is not None - attention = self.attention_store[step] - return attention - - def aggregate_attention( - self, attention_maps, prompts, res: Union[int, Tuple[int]], from_where: List[str], is_cross: bool, select: int - ): - out = [[] for x in range(self.batch_size)] - if isinstance(res, int): - num_pixels = res**2 - resolution = (res, res) - else: - num_pixels = res[0] * res[1] - resolution = res[:2] - - for location in from_where: - for bs_item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]: - for batch, item in enumerate(bs_item): - if item.shape[1] == num_pixels: - cross_maps = item.reshape(len(prompts), -1, *resolution, item.shape[-1])[select] - out[batch].append(cross_maps) - - out = torch.stack([torch.cat(x, dim=0) for x in out]) - # average over heads - out = out.sum(1) / out.shape[1] - return out - - def __init__(self, average: bool, batch_size=1, max_resolution=16, max_size: int = None): - self.step_store = self.get_empty_store() - self.attention_store = [] - self.cur_step = 0 - self.average = average - self.batch_size = batch_size - if max_size is None: - self.max_size = max_resolution**2 - elif max_size is not None and max_resolution is None: - self.max_size = max_size - else: - raise ValueError("Only allowed to set one of max_resolution or max_size") - - -# Copied from diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.LeditsGaussianSmoothing -class LeditsGaussianSmoothing: - def __init__(self, device): - kernel_size = [3, 3] - sigma = [0.5, 0.5] - - # The gaussian kernel is the product of the gaussian function of each dimension. - kernel = 1 - meshgrids = torch.meshgrid([torch.arange(size, dtype=torch.float32) for size in kernel_size], indexing="ij") - for size, std, mgrid in zip(kernel_size, sigma, meshgrids): - mean = (size - 1) / 2 - kernel *= 1 / (std * math.sqrt(2 * math.pi)) * torch.exp(-(((mgrid - mean) / (2 * std)) ** 2)) - - # Make sure sum of values in gaussian kernel equals 1. - kernel = kernel / torch.sum(kernel) - - # Reshape to depthwise convolutional weight - kernel = kernel.view(1, 1, *kernel.size()) - kernel = kernel.repeat(1, *[1] * (kernel.dim() - 1)) - - self.weight = kernel.to(device) - - def __call__(self, input): - """ - Arguments: - Apply gaussian filter to input. - input (torch.Tensor): Input to apply gaussian filter on. - Returns: - filtered (torch.Tensor): Filtered output. - """ - return F.conv2d(input, weight=self.weight.to(input.dtype)) - - -# Copied from diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.LEDITSCrossAttnProcessor -class LEDITSCrossAttnProcessor: - def __init__(self, attention_store, place_in_unet, pnp, editing_prompts): - self.attnstore = attention_store - self.place_in_unet = place_in_unet - self.editing_prompts = editing_prompts - self.pnp = pnp - - def __call__( - self, - attn: Attention, - hidden_states, - encoder_hidden_states, - attention_mask=None, - temb=None, - ): - batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query = attn.head_to_batch_dim(query) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - - attention_probs = attn.get_attention_scores(query, key, attention_mask) - self.attnstore( - attention_probs, - is_cross=True, - place_in_unet=self.place_in_unet, - editing_prompts=self.editing_prompts, - PnP=self.pnp, - ) - - hidden_states = torch.bmm(attention_probs, value) - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - hidden_states = hidden_states / attn.rescale_output_factor - return hidden_states - - -class LEditsPPPipelineStableDiffusionXL( - DiffusionPipeline, - FromSingleFileMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, -): - """ - Pipeline for textual image editing using LEDits++ with Stable Diffusion XL. - - This model inherits from [`DiffusionPipeline`] and builds on the [`StableDiffusionXLPipeline`]. Check the - superclass documentation for the generic methods implemented for all pipelines (downloading, saving, running on a - particular device, etc.). - - In addition the pipeline inherits the following loading methods: - - *LoRA*: [`LEditsPPPipelineStableDiffusionXL.load_lora_weights`] - - *Ckpt*: [`loaders.FromSingleFileMixin.from_single_file`] - - as well as the following saving methods: - - *LoRA*: [`loaders.StableDiffusionXLPipeline.save_lora_weights`] - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.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 ([`~transformers.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 ([`~transformers.CLIPTokenizer`]): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - tokenizer_2 ([`~transformers.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 ([`DPMSolverMultistepScheduler`] or [`DDIMScheduler`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latens. Can be one of - [`DPMSolverMultistepScheduler`] or [`DDIMScheduler`]. If any other scheduler is passed it will - automatically be set to [`DPMSolverMultistepScheduler`]. - 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`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to - watermark output images. If not defined, it will default to True if the package is installed, otherwise no - watermarker will be used. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "image_encoder", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "negative_add_time_ids", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: Union[DPMSolverMultistepScheduler, DDIMScheduler], - image_encoder: CLIPVisionModelWithProjection = None, - feature_extractor: CLIPImageProcessor = None, - force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - ): - super().__init__() - - 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, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - ) - 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) - - if not isinstance(scheduler, DDIMScheduler) and not isinstance(scheduler, DPMSolverMultistepScheduler): - self.scheduler = DPMSolverMultistepScheduler.from_config( - scheduler.config, algorithm_type="sde-dpmsolver++", solver_order=2 - ) - logger.warning( - "This pipeline only supports DDIMScheduler and DPMSolverMultistepScheduler. " - "The scheduler has been changed to DPMSolverMultistepScheduler." - ) - - 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 - ) - - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - self.inversion_steps = None - - def encode_prompt( - self, - device: Optional[torch.device] = 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, - 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, - ) -> object: - r""" - Encodes the prompt into text encoder hidden states. - - Args: - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - 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*): - 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*): - 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. - 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. - enable_edit_guidance (`bool`): - Whether to guide towards an editing prompt or not. - 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*): - Pre-generated edit text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. - If not provided and 'enable_edit_guidance' is True, editing_prompt_embeds will be generated from - `editing_prompt` input argument. - editing_pooled_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated edit pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, pooled editing_pooled_prompt_embeds will be generated from `editing_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, 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) - - batch_size = self.batch_size - - # 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] - ) - num_edit_tokens = 0 - - # get unconditional embeddings for classifier free guidance - zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt - - if 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 batch_size != len(negative_prompt): - raise ValueError( - f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but image inversion " - f" has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" - " the batch size of the input images." - ) - 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) - - uncond_input = tokenizer( - negative_prompt, - padding="max_length", - max_length=tokenizer.model_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 - 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 zero_out_negative_prompt: - negative_prompt_embeds = torch.zeros_like(negative_prompt_embeds) - negative_pooled_prompt_embeds = torch.zeros_like(negative_pooled_prompt_embeds) - - if enable_edit_guidance and editing_prompt_embeds is None: - editing_prompt_2 = editing_prompt - - editing_prompts = [editing_prompt, editing_prompt_2] - edit_prompt_embeds_list = [] - - for editing_prompt, tokenizer, text_encoder in zip(editing_prompts, tokenizers, text_encoders): - if isinstance(self, TextualInversionLoaderMixin): - editing_prompt = self.maybe_convert_prompt(editing_prompt, tokenizer) - - max_length = negative_prompt_embeds.shape[1] - edit_concepts_input = tokenizer( - # [x for item in editing_prompt for x in repeat(item, batch_size)], - editing_prompt, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - return_length=True, - ) - num_edit_tokens = edit_concepts_input.length - 2 - - edit_concepts_embeds = text_encoder( - edit_concepts_input.input_ids.to(device), - output_hidden_states=True, - ) - # We are only ALWAYS interested in the pooled output of the final text encoder - editing_pooled_prompt_embeds = edit_concepts_embeds[0] - if clip_skip is None: - edit_concepts_embeds = edit_concepts_embeds.hidden_states[-2] - else: - # "2" because SDXL always indexes from the penultimate layer. - edit_concepts_embeds = edit_concepts_embeds.hidden_states[-(clip_skip + 2)] - - edit_prompt_embeds_list.append(edit_concepts_embeds) - - edit_concepts_embeds = torch.concat(edit_prompt_embeds_list, dim=-1) - elif not enable_edit_guidance: - edit_concepts_embeds = None - editing_pooled_prompt_embeds = None - - negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) - bs_embed, seq_len, _ = negative_prompt_embeds.shape - # 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=self.text_encoder_2.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 enable_edit_guidance: - bs_embed_edit, seq_len, _ = edit_concepts_embeds.shape - edit_concepts_embeds = edit_concepts_embeds.to(dtype=self.text_encoder_2.dtype, device=device) - edit_concepts_embeds = edit_concepts_embeds.repeat(1, num_images_per_prompt, 1) - edit_concepts_embeds = edit_concepts_embeds.view(bs_embed_edit * num_images_per_prompt, seq_len, -1) - - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view( - bs_embed * num_images_per_prompt, -1 - ) - - if enable_edit_guidance: - editing_pooled_prompt_embeds = editing_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view( - bs_embed_edit * 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 ( - negative_prompt_embeds, - edit_concepts_embeds, - negative_pooled_prompt_embeds, - editing_pooled_prompt_embeds, - num_edit_tokens, - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, eta, generator=None): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - negative_prompt=None, - negative_prompt_2=None, - negative_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - ): - 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 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`." - ) - - # Modified from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - def prepare_latents(self, device, latents): - latents = latents.to(device) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - 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.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_end(self): - return self._denoising_end - - @property - def num_timesteps(self): - return self._num_timesteps - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - # Copied from diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.LEditsPPPipelineStableDiffusion.prepare_unet - def prepare_unet(self, attention_store, PnP: bool = False): - attn_procs = {} - for name in self.unet.attn_processors.keys(): - if name.startswith("mid_block"): - place_in_unet = "mid" - elif name.startswith("up_blocks"): - place_in_unet = "up" - elif name.startswith("down_blocks"): - place_in_unet = "down" - else: - continue - - if "attn2" in name and place_in_unet != "mid": - attn_procs[name] = LEDITSCrossAttnProcessor( - attention_store=attention_store, - place_in_unet=place_in_unet, - pnp=PnP, - editing_prompts=self.enabled_editing_prompts, - ) - else: - attn_procs[name] = AttnProcessor() - - self.unet.set_attn_processor(attn_procs) - - @torch.no_grad() - @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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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, - use_cross_attn_mask: bool = False, - use_intersect_mask: bool = False, - user_mask: Optional[torch.Tensor] = None, - attn_store_steps: Optional[List[int]] = [], - 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"], - **kwargs, - ): - r""" - The call function to the pipeline for editing. The - [`~pipelines.ledits_pp.LEditsPPPipelineStableDiffusionXL.invert`] method has to be called beforehand. Edits - will always be performed for the last inverted image(s). - - Args: - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - 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*): - 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 - 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. - 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. - ip_adapter_image: (`PipelineImageInput`, *optional*): - Optional image input to work with IP Adapters. - 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. - 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. - cross_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). - guidance_rescale (`float`, *optional*, defaults to 0.7): - 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. - 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 `(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*): - 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`. - editing_prompt_embeddings (`torch.Tensor`, *optional*): - Pre-generated 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. - editing_pooled_prompt_embeddings (`torch.Tensor`, *optional*): - 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`): - Whether the corresponding prompt in `editing_prompt` should be increased or decreased. - 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): - Number of diffusion steps (for each prompt) for which guidance is not applied. - 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): - 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 - 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 - is set to true. Cross-attention masks are defined as 'M^1' of equation 12 of [LEDITS++ - paper](https://huggingface.co/papers/2311.16711). - use_intersect_mask: - 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). - user_mask: - User-provided mask for even better control over the editing process. This is helpful when LEDITS++'s - implicit masks do not meet user preferences. - attn_store_steps: - Steps for which the attention maps are stored in the AttentionStore. Just for visualization purposes. - store_averaged_over_steps: - Whether the attention maps for the 'attn_store_steps' are stored averaged over the diffusion steps. If - False, attention maps for each step are stores separately. Just for visualization purposes. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.ledits_pp.LEditsPPDiffusionPipelineOutput`] or `tuple`: - [`~pipelines.ledits_pp.LEditsPPDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple. When - returning a tuple, the first element is a list with the generated images. - """ - if self.inversion_steps is None: - raise ValueError( - "You need to invert an input image first before calling the pipeline. The `invert` method has to be called beforehand. Edits will always be performed for the last inverted image(s)." - ) - - eta = self.eta - num_images_per_prompt = 1 - latents = self.init_latents - - zs = self.zs - self.scheduler.set_timesteps(len(self.scheduler.timesteps)) - - if use_intersect_mask: - use_cross_attn_mask = True - - if use_cross_attn_mask: - self.smoothing = LeditsGaussianSmoothing(self.device) - - if user_mask is not None: - user_mask = user_mask.to(self.device) - - # TODO: Check inputs - # 1. Check inputs. Raise error if not correct - # self.check_inputs( - # callback_steps, - # negative_prompt, - # negative_prompt_2, - # prompt_embeds, - # negative_prompt_embeds, - # pooled_prompt_embeds, - # negative_pooled_prompt_embeds, - # ) - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - - # 2. Define call parameters - batch_size = self.batch_size - - device = self._execution_device - - if editing_prompt: - enable_edit_guidance = True - if isinstance(editing_prompt, str): - editing_prompt = [editing_prompt] - self.enabled_editing_prompts = len(editing_prompt) - elif editing_prompt_embeddings is not None: - enable_edit_guidance = True - self.enabled_editing_prompts = editing_prompt_embeddings.shape[0] - else: - self.enabled_editing_prompts = 0 - enable_edit_guidance = False - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - ( - prompt_embeds, - edit_prompt_embeds, - negative_pooled_prompt_embeds, - pooled_edit_embeds, - num_edit_tokens, - ) = self.encode_prompt( - device=device, - num_images_per_prompt=num_images_per_prompt, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_embeds=negative_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.clip_skip, - enable_edit_guidance=enable_edit_guidance, - editing_prompt=editing_prompt, - editing_prompt_embeds=editing_prompt_embeddings, - editing_pooled_prompt_embeds=editing_pooled_prompt_embeds, - ) - - # 4. Prepare timesteps - # self.scheduler.set_timesteps(num_inference_steps, device=device) - - timesteps = self.inversion_steps - t_to_idx = {int(v): k for k, v in enumerate(timesteps)} - - if use_cross_attn_mask: - self.attention_store = LeditsAttentionStore( - average=store_averaged_over_steps, - batch_size=batch_size, - max_size=(latents.shape[-2] / 4.0) * (latents.shape[-1] / 4.0), - max_resolution=None, - ) - self.prepare_unet(self.attention_store) - resolution = latents.shape[-2:] - att_res = (int(resolution[0] / 4), int(resolution[1] / 4)) - - # 5. Prepare latent variables - latents = self.prepare_latents(device=device, latents=latents) - - # 6. Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(eta) - - if self.text_encoder_2 is None: - text_encoder_projection_dim = int(negative_pooled_prompt_embeds.shape[-1]) - else: - text_encoder_projection_dim = self.text_encoder_2.config.projection_dim - - # 7. Prepare added time ids & embeddings - add_text_embeds = negative_pooled_prompt_embeds - add_time_ids = self._get_add_time_ids( - self.size, - crops_coords_top_left, - self.size, - dtype=negative_pooled_prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - - if enable_edit_guidance: - prompt_embeds = torch.cat([prompt_embeds, edit_prompt_embeds], dim=0) - add_text_embeds = torch.cat([add_text_embeds, pooled_edit_embeds], dim=0) - edit_concepts_time_ids = add_time_ids.repeat(edit_prompt_embeds.shape[0], 1) - add_time_ids = torch.cat([add_time_ids, edit_concepts_time_ids], dim=0) - self.text_cross_attention_maps = [editing_prompt] if isinstance(editing_prompt, str) else editing_prompt - - 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) - - if ip_adapter_image is not None: - # TODO: fix image encoding - image_embeds, negative_image_embeds = self.encode_image(ip_adapter_image, device, num_images_per_prompt) - if self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds]) - image_embeds = image_embeds.to(device) - - # 8. Denoising loop - self.sem_guidance = None - self.activation_mask = None - - if ( - self.denoising_end is not None - and isinstance(self.denoising_end, float) - and self.denoising_end > 0 - and self.denoising_end < 1 - ): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - # 9. 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) - - self._num_timesteps = len(timesteps) - with self.progress_bar(total=self._num_timesteps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * (1 + self.enabled_editing_prompts)) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - # predict the noise residual - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - if ip_adapter_image is not None: - added_cond_kwargs["image_embeds"] = image_embeds - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - noise_pred_out = noise_pred.chunk(1 + self.enabled_editing_prompts) # [b,4, 64, 64] - noise_pred_uncond = noise_pred_out[0] - noise_pred_edit_concepts = noise_pred_out[1:] - - noise_guidance_edit = torch.zeros( - noise_pred_uncond.shape, - device=self.device, - dtype=noise_pred_uncond.dtype, - ) - - if sem_guidance is not None and len(sem_guidance) > i: - noise_guidance_edit += sem_guidance[i].to(self.device) - - elif enable_edit_guidance: - if self.activation_mask is None: - self.activation_mask = torch.zeros( - (len(timesteps), self.enabled_editing_prompts, *noise_pred_edit_concepts[0].shape) - ) - if self.sem_guidance is None: - self.sem_guidance = torch.zeros((len(timesteps), *noise_pred_uncond.shape)) - - # noise_guidance_edit = torch.zeros_like(noise_guidance) - for c, noise_pred_edit_concept in enumerate(noise_pred_edit_concepts): - if isinstance(edit_warmup_steps, list): - edit_warmup_steps_c = edit_warmup_steps[c] - else: - edit_warmup_steps_c = edit_warmup_steps - if i < edit_warmup_steps_c: - continue - - if isinstance(edit_guidance_scale, list): - edit_guidance_scale_c = edit_guidance_scale[c] - else: - edit_guidance_scale_c = edit_guidance_scale - - if isinstance(edit_threshold, list): - edit_threshold_c = edit_threshold[c] - else: - edit_threshold_c = edit_threshold - if isinstance(reverse_editing_direction, list): - reverse_editing_direction_c = reverse_editing_direction[c] - else: - reverse_editing_direction_c = reverse_editing_direction - - if isinstance(edit_cooldown_steps, list): - edit_cooldown_steps_c = edit_cooldown_steps[c] - elif edit_cooldown_steps is None: - edit_cooldown_steps_c = i + 1 - else: - edit_cooldown_steps_c = edit_cooldown_steps - - if i >= edit_cooldown_steps_c: - continue - - noise_guidance_edit_tmp = noise_pred_edit_concept - noise_pred_uncond - - if reverse_editing_direction_c: - noise_guidance_edit_tmp = noise_guidance_edit_tmp * -1 - - noise_guidance_edit_tmp = noise_guidance_edit_tmp * edit_guidance_scale_c - - if user_mask is not None: - noise_guidance_edit_tmp = noise_guidance_edit_tmp * user_mask - - if use_cross_attn_mask: - out = self.attention_store.aggregate_attention( - attention_maps=self.attention_store.step_store, - prompts=self.text_cross_attention_maps, - res=att_res, - from_where=["up", "down"], - is_cross=True, - select=self.text_cross_attention_maps.index(editing_prompt[c]), - ) - attn_map = out[:, :, :, 1 : 1 + num_edit_tokens[c]] # 0 -> startoftext - - # average over all tokens - if attn_map.shape[3] != num_edit_tokens[c]: - raise ValueError( - f"Incorrect shape of attention_map. Expected size {num_edit_tokens[c]}, but found {attn_map.shape[3]}!" - ) - attn_map = torch.sum(attn_map, dim=3) - - # gaussian_smoothing - attn_map = F.pad(attn_map.unsqueeze(1), (1, 1, 1, 1), mode="reflect") - attn_map = self.smoothing(attn_map).squeeze(1) - - # torch.quantile function expects float32 - if attn_map.dtype == torch.float32: - tmp = torch.quantile(attn_map.flatten(start_dim=1), edit_threshold_c, dim=1) - else: - tmp = torch.quantile( - attn_map.flatten(start_dim=1).to(torch.float32), edit_threshold_c, dim=1 - ).to(attn_map.dtype) - attn_mask = torch.where( - attn_map >= tmp.unsqueeze(1).unsqueeze(1).repeat(1, *att_res), 1.0, 0.0 - ) - - # resolution must match latent space dimension - attn_mask = F.interpolate( - attn_mask.unsqueeze(1), - noise_guidance_edit_tmp.shape[-2:], # 64,64 - ).repeat(1, 4, 1, 1) - self.activation_mask[i, c] = attn_mask.detach().cpu() - if not use_intersect_mask: - noise_guidance_edit_tmp = noise_guidance_edit_tmp * attn_mask - - if use_intersect_mask: - noise_guidance_edit_tmp_quantile = torch.abs(noise_guidance_edit_tmp) - noise_guidance_edit_tmp_quantile = torch.sum( - noise_guidance_edit_tmp_quantile, dim=1, keepdim=True - ) - noise_guidance_edit_tmp_quantile = noise_guidance_edit_tmp_quantile.repeat( - 1, self.unet.config.in_channels, 1, 1 - ) - - # torch.quantile function expects float32 - if noise_guidance_edit_tmp_quantile.dtype == torch.float32: - tmp = torch.quantile( - noise_guidance_edit_tmp_quantile.flatten(start_dim=2), - edit_threshold_c, - dim=2, - keepdim=False, - ) - else: - tmp = torch.quantile( - noise_guidance_edit_tmp_quantile.flatten(start_dim=2).to(torch.float32), - edit_threshold_c, - dim=2, - keepdim=False, - ).to(noise_guidance_edit_tmp_quantile.dtype) - - intersect_mask = ( - torch.where( - noise_guidance_edit_tmp_quantile >= tmp[:, :, None, None], - torch.ones_like(noise_guidance_edit_tmp), - torch.zeros_like(noise_guidance_edit_tmp), - ) - * attn_mask - ) - - self.activation_mask[i, c] = intersect_mask.detach().cpu() - - noise_guidance_edit_tmp = noise_guidance_edit_tmp * intersect_mask - - elif not use_cross_attn_mask: - # calculate quantile - noise_guidance_edit_tmp_quantile = torch.abs(noise_guidance_edit_tmp) - noise_guidance_edit_tmp_quantile = torch.sum( - noise_guidance_edit_tmp_quantile, dim=1, keepdim=True - ) - noise_guidance_edit_tmp_quantile = noise_guidance_edit_tmp_quantile.repeat(1, 4, 1, 1) - - # torch.quantile function expects float32 - if noise_guidance_edit_tmp_quantile.dtype == torch.float32: - tmp = torch.quantile( - noise_guidance_edit_tmp_quantile.flatten(start_dim=2), - edit_threshold_c, - dim=2, - keepdim=False, - ) - else: - tmp = torch.quantile( - noise_guidance_edit_tmp_quantile.flatten(start_dim=2).to(torch.float32), - edit_threshold_c, - dim=2, - keepdim=False, - ).to(noise_guidance_edit_tmp_quantile.dtype) - - self.activation_mask[i, c] = ( - torch.where( - noise_guidance_edit_tmp_quantile >= tmp[:, :, None, None], - torch.ones_like(noise_guidance_edit_tmp), - torch.zeros_like(noise_guidance_edit_tmp), - ) - .detach() - .cpu() - ) - - noise_guidance_edit_tmp = torch.where( - noise_guidance_edit_tmp_quantile >= tmp[:, :, None, None], - noise_guidance_edit_tmp, - torch.zeros_like(noise_guidance_edit_tmp), - ) - - noise_guidance_edit += noise_guidance_edit_tmp - - self.sem_guidance[i] = noise_guidance_edit.detach().cpu() - - noise_pred = noise_pred_uncond + noise_guidance_edit - - # compute the previous noisy sample x_t -> x_t-1 - if enable_edit_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg( - noise_pred, - noise_pred_edit_concepts.mean(dim=0, keepdim=False), - guidance_rescale=self.guidance_rescale, - ) - - idx = t_to_idx[int(t)] - latents = self.scheduler.step( - noise_pred, t, latents, variance_noise=zs[idx], **extra_step_kwargs, return_dict=False - )[0] - - # step callback - if use_cross_attn_mask: - store_step = i in attn_store_steps - self.attention_store.between_steps(store_step) - - 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds - ) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - # negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids) - - # call the callback, if provided - if i == len(timesteps) - 1 or ((i + 1) > 0 and (i + 1) % self.scheduler.order == 0): - progress_bar.update() - - if XLA_AVAILABLE: - xm.mark_step() - - 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": - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - 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 LEditsPPDiffusionPipelineOutput(images=image, nsfw_content_detected=None) - - @torch.no_grad() - # Modified from diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.LEditsPPPipelineStableDiffusion.encode_image - def encode_image(self, image, dtype=None, height=None, width=None, resize_mode="default", crops_coords=None): - image = self.image_processor.preprocess( - image=image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords - ) - height, width = image.shape[-2:] - if height % 32 != 0 or width % 32 != 0: - raise ValueError( - "Image height and width must be a factor of 32. " - "Consider down-sampling the input using the `height` and `width` parameters" - ) - resized = self.image_processor.postprocess(image=image, output_type="pil") - - if max(image.shape[-2:]) > self.vae.config["sample_size"] * 1.5: - logger.warning( - "Your input images far exceed the default resolution of the underlying diffusion model. " - "The output images may contain severe artifacts! " - "Consider down-sampling the input using the `height` and `width` parameters" - ) - image = image.to(self.device, dtype=dtype) - needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast - - if needs_upcasting: - image = image.float() - self.upcast_vae() - - x0 = self.vae.encode(image).latent_dist.mode() - x0 = x0.to(dtype) - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - - x0 = self.vae.config.scaling_factor * x0 - return x0, resized - - @torch.no_grad() - def invert( - self, - image: PipelineImageInput, - source_prompt: str = "", - source_guidance_scale=3.5, - negative_prompt: str = None, - 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), - 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, - ): - r""" - The function to the pipeline for image inversion as described by the [LEDITS++ - Paper](https://huggingface.co/papers/2301.12247). If the scheduler is set to [`~schedulers.DDIMScheduler`] the - inversion proposed by [edit-friendly DPDM](https://huggingface.co/papers/2304.06140) will be performed instead. - - Args: - image (`PipelineImageInput`): - Input for the image(s) that are to be edited. Multiple input images have to default to the same aspect - ratio. - source_prompt (`str`, defaults to `""`): - Prompt describing the input image that will be used for guidance during inversion. Guidance is disabled - 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*): - 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_inversion_steps (`int`, defaults to `50`): - Number of total performed inversion steps after discarding the initial `skip` steps. - skip (`float`, defaults to `0.15`): - Portion of initial steps that will be ignored for inversion and subsequent generation. Lower values - will lead to stronger changes to the input image. `skip` has to be between `0` and `1`. - 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` 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). - num_zero_noise_steps (`int`, defaults to `3`): - Number of final diffusion steps that will not renoise the current image. If no steps are set to zero - SD-XL in combination with [`DPMSolverMultistepScheduler`] will produce noise artifacts. - cross_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). - - Returns: - [`~pipelines.ledits_pp.LEditsPPInversionPipelineOutput`]: Output will contain the resized input image(s) - and respective VAE reconstruction(s). - """ - if height is not None and height % 32 != 0 or width is not None and width % 32 != 0: - raise ValueError("height and width must be a factor of 32.") - - # Reset attn processor, we do not want to store attn maps during inversion - self.unet.set_attn_processor(AttnProcessor()) - - self.eta = 1.0 - - self.scheduler.config.timestep_spacing = "leading" - self.scheduler.set_timesteps(int(num_inversion_steps * (1 + skip))) - self.inversion_steps = self.scheduler.timesteps[-num_inversion_steps:] - timesteps = self.inversion_steps - - num_images_per_prompt = 1 - - device = self._execution_device - - # 0. Ensure that only uncond embedding is used if prompt = "" - if source_prompt == "": - # noise pred should only be noise_pred_uncond - source_guidance_scale = 0.0 - do_classifier_free_guidance = False - else: - do_classifier_free_guidance = source_guidance_scale > 1.0 - - # 1. prepare image - x0, resized = self.encode_image( - image, - dtype=self.text_encoder_2.dtype, - height=height, - width=width, - resize_mode=resize_mode, - crops_coords=crops_coords, - ) - width = x0.shape[2] * self.vae_scale_factor - height = x0.shape[3] * self.vae_scale_factor - self.size = (height, width) - - self.batch_size = x0.shape[0] - - # 2. get embeddings - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - - if isinstance(source_prompt, str): - source_prompt = [source_prompt] * self.batch_size - - ( - negative_prompt_embeds, - prompt_embeds, - negative_pooled_prompt_embeds, - edit_pooled_prompt_embeds, - _, - ) = self.encode_prompt( - device=device, - num_images_per_prompt=num_images_per_prompt, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - editing_prompt=source_prompt, - lora_scale=text_encoder_lora_scale, - enable_edit_guidance=do_classifier_free_guidance, - ) - if self.text_encoder_2 is None: - text_encoder_projection_dim = int(negative_pooled_prompt_embeds.shape[-1]) - else: - text_encoder_projection_dim = self.text_encoder_2.config.projection_dim - - # 3. Prepare added time ids & embeddings - add_text_embeds = negative_pooled_prompt_embeds - add_time_ids = self._get_add_time_ids( - self.size, - crops_coords_top_left, - self.size, - dtype=negative_prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - - if do_classifier_free_guidance: - negative_prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - add_text_embeds = torch.cat([add_text_embeds, edit_pooled_prompt_embeds], dim=0) - add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0) - - negative_prompt_embeds = negative_prompt_embeds.to(device) - - add_text_embeds = add_text_embeds.to(device) - add_time_ids = add_time_ids.to(device).repeat(self.batch_size * num_images_per_prompt, 1) - - # autoencoder reconstruction - if self.vae.dtype == torch.float16 and self.vae.config.force_upcast: - self.upcast_vae() - x0_tmp = x0.to(next(iter(self.vae.post_quant_conv.parameters())).dtype) - image_rec = self.vae.decode( - x0_tmp / self.vae.config.scaling_factor, return_dict=False, generator=generator - )[0] - elif self.vae.config.force_upcast: - x0_tmp = x0.to(next(iter(self.vae.post_quant_conv.parameters())).dtype) - image_rec = self.vae.decode( - x0_tmp / self.vae.config.scaling_factor, return_dict=False, generator=generator - )[0] - else: - image_rec = self.vae.decode(x0 / self.vae.config.scaling_factor, return_dict=False, generator=generator)[0] - - image_rec = self.image_processor.postprocess(image_rec, output_type="pil") - - # 5. find zs and xts - variance_noise_shape = (num_inversion_steps, *x0.shape) - - # intermediate latents - t_to_idx = {int(v): k for k, v in enumerate(timesteps)} - xts = torch.zeros(size=variance_noise_shape, device=self.device, dtype=negative_prompt_embeds.dtype) - - for t in reversed(timesteps): - idx = num_inversion_steps - t_to_idx[int(t)] - 1 - noise = randn_tensor(shape=x0.shape, generator=generator, device=self.device, dtype=x0.dtype) - xts[idx] = self.scheduler.add_noise(x0, noise, t.unsqueeze(0)) - xts = torch.cat([x0.unsqueeze(0), xts], dim=0) - - # noise maps - zs = torch.zeros(size=variance_noise_shape, device=self.device, dtype=negative_prompt_embeds.dtype) - - self.scheduler.set_timesteps(len(self.scheduler.timesteps)) - - for t in self.progress_bar(timesteps): - idx = num_inversion_steps - t_to_idx[int(t)] - 1 - # 1. predict noise residual - xt = xts[idx + 1] - - latent_model_input = torch.cat([xt] * 2) if do_classifier_free_guidance else xt - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=negative_prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # 2. perform guidance - if do_classifier_free_guidance: - noise_pred_out = noise_pred.chunk(2) - noise_pred_uncond, noise_pred_text = noise_pred_out[0], noise_pred_out[1] - noise_pred = noise_pred_uncond + source_guidance_scale * (noise_pred_text - noise_pred_uncond) - - xtm1 = xts[idx] - z, xtm1_corrected = compute_noise(self.scheduler, xtm1, xt, t, noise_pred, self.eta) - zs[idx] = z - - # correction to avoid error accumulation - xts[idx] = xtm1_corrected - - self.init_latents = xts[-1] - zs = zs.flip(0) - - if num_zero_noise_steps > 0: - zs[-num_zero_noise_steps:] = torch.zeros_like(zs[-num_zero_noise_steps:]) - self.zs = zs - return LEditsPPInversionPipelineOutput(images=resized, vae_reconstruction_images=image_rec) - - -# Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.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 - - -# Copied from diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.compute_noise_ddim -def compute_noise_ddim(scheduler, prev_latents, latents, timestep, noise_pred, eta): - # 1. get previous step value (=t-1) - prev_timestep = timestep - scheduler.config.num_train_timesteps // scheduler.num_inference_steps - - # 2. compute alphas, betas - alpha_prod_t = scheduler.alphas_cumprod[timestep] - alpha_prod_t_prev = ( - scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else scheduler.final_alpha_cumprod - ) - - beta_prod_t = 1 - alpha_prod_t - - # 3. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (12) from https://huggingface.co/papers/2010.02502 - pred_original_sample = (latents - beta_prod_t ** (0.5) * noise_pred) / alpha_prod_t ** (0.5) - - # 4. Clip "predicted x_0" - if scheduler.config.clip_sample: - pred_original_sample = torch.clamp(pred_original_sample, -1, 1) - - # 5. compute variance: "sigma_t(η)" -> see formula (16) - # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1) - variance = scheduler._get_variance(timestep, prev_timestep) - std_dev_t = eta * variance ** (0.5) - - # 6. compute "direction pointing to x_t" of formula (12) from https://huggingface.co/papers/2010.02502 - pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * noise_pred - - # modified so that updated xtm1 is returned as well (to avoid error accumulation) - mu_xt = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction - if variance > 0.0: - noise = (prev_latents - mu_xt) / (variance ** (0.5) * eta) - else: - noise = torch.tensor([0.0]).to(latents.device) - - return noise, mu_xt + (eta * variance**0.5) * noise - - -# Copied from diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.compute_noise_sde_dpm_pp_2nd -def compute_noise_sde_dpm_pp_2nd(scheduler, prev_latents, latents, timestep, noise_pred, eta): - def first_order_update(model_output, sample): # timestep, prev_timestep, sample): - sigma_t, sigma_s = scheduler.sigmas[scheduler.step_index + 1], scheduler.sigmas[scheduler.step_index] - alpha_t, sigma_t = scheduler._sigma_to_alpha_sigma_t(sigma_t) - alpha_s, sigma_s = scheduler._sigma_to_alpha_sigma_t(sigma_s) - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s = torch.log(alpha_s) - torch.log(sigma_s) - - h = lambda_t - lambda_s - - mu_xt = (sigma_t / sigma_s * torch.exp(-h)) * sample + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output - - mu_xt = scheduler.dpm_solver_first_order_update( - model_output=model_output, sample=sample, noise=torch.zeros_like(sample) - ) - - sigma = sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) - if sigma > 0.0: - noise = (prev_latents - mu_xt) / sigma - else: - noise = torch.tensor([0.0]).to(sample.device) - - prev_sample = mu_xt + sigma * noise - return noise, prev_sample - - def second_order_update(model_output_list, sample): # timestep_list, prev_timestep, sample): - sigma_t, sigma_s0, sigma_s1 = ( - scheduler.sigmas[scheduler.step_index + 1], - scheduler.sigmas[scheduler.step_index], - scheduler.sigmas[scheduler.step_index - 1], - ) - - alpha_t, sigma_t = scheduler._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = scheduler._sigma_to_alpha_sigma_t(sigma_s0) - alpha_s1, sigma_s1 = scheduler._sigma_to_alpha_sigma_t(sigma_s1) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1) - - m0, m1 = model_output_list[-1], model_output_list[-2] - - h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1 - r0 = h_0 / h - D0, D1 = m0, (1.0 / r0) * (m0 - m1) - - mu_xt = ( - (sigma_t / sigma_s0 * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 - + 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1 - ) - - sigma = sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) - if sigma > 0.0: - noise = (prev_latents - mu_xt) / sigma - else: - noise = torch.tensor([0.0]).to(sample.device) - - prev_sample = mu_xt + sigma * noise - - return noise, prev_sample - - if scheduler.step_index is None: - scheduler._init_step_index(timestep) - - model_output = scheduler.convert_model_output(model_output=noise_pred, sample=latents) - for i in range(scheduler.config.solver_order - 1): - scheduler.model_outputs[i] = scheduler.model_outputs[i + 1] - scheduler.model_outputs[-1] = model_output - - if scheduler.lower_order_nums < 1: - noise, prev_sample = first_order_update(model_output, latents) - else: - noise, prev_sample = second_order_update(scheduler.model_outputs, latents) - - if scheduler.lower_order_nums < scheduler.config.solver_order: - scheduler.lower_order_nums += 1 - - # upon completion increase step index by one - scheduler._step_index += 1 - - return noise, prev_sample - - -# Copied from diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.compute_noise -def compute_noise(scheduler, *args): - if isinstance(scheduler, DDIMScheduler): - return compute_noise_ddim(scheduler, *args) - elif ( - isinstance(scheduler, DPMSolverMultistepScheduler) - and scheduler.config.algorithm_type == "sde-dpmsolver++" - and scheduler.config.solver_order == 2 - ): - return compute_noise_sde_dpm_pp_2nd(scheduler, *args) - else: - raise NotImplementedError diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/pipeline_output.py deleted file mode 100644 index 756be82b0069aa986bc0594476e9060e114a7c84..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ledits_pp/pipeline_output.py +++ /dev/null @@ -1,43 +0,0 @@ -from dataclasses import dataclass -from typing import List, Optional, Union - -import numpy as np -import PIL.Image - -from ...utils import BaseOutput - - -@dataclass -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, - num_channels)`. - 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]] - - -@dataclass -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 ` - (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 - ` (batch_size, height, width, num_channels)`. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] - vae_reconstruction_images: Union[List[PIL.Image.Image], np.ndarray] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__init__.py deleted file mode 100644 index 6001867916b348e0ea1634e9bc5090623ed6c5ae..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__init__.py +++ /dev/null @@ -1,56 +0,0 @@ -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["modeling_latent_upsampler"] = ["LTXLatentUpsamplerModel"] - _import_structure["pipeline_ltx"] = ["LTXPipeline"] - _import_structure["pipeline_ltx_condition"] = ["LTXConditionPipeline"] - _import_structure["pipeline_ltx_image2video"] = ["LTXImageToVideoPipeline"] - _import_structure["pipeline_ltx_latent_upsample"] = ["LTXLatentUpsamplePipeline"] - -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 .modeling_latent_upsampler import LTXLatentUpsamplerModel - from .pipeline_ltx import LTXPipeline - from .pipeline_ltx_condition import LTXConditionPipeline - from .pipeline_ltx_image2video import LTXImageToVideoPipeline - from .pipeline_ltx_latent_upsample import LTXLatentUpsamplePipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 7309451d308ab73e66f985bf392f4dc512db89ed..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/modeling_latent_upsampler.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/modeling_latent_upsampler.cpython-310.pyc deleted file mode 100644 index 991b88765500832a874a49de0af2482ae83c91e2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/modeling_latent_upsampler.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_ltx.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_ltx.cpython-310.pyc deleted file mode 100644 index 9aafa804e52f49e6496b46b1ec28d2818363f36f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_ltx.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_ltx_condition.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_ltx_condition.cpython-310.pyc deleted file mode 100644 index de901546dc59580f5f351866f8854abfda8381d2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_ltx_condition.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_ltx_image2video.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_ltx_image2video.cpython-310.pyc deleted file mode 100644 index b2e8aa025687f61219acf4e1de9f895b79739b98..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_ltx_image2video.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_ltx_latent_upsample.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_ltx_latent_upsample.cpython-310.pyc deleted file mode 100644 index 6c630bd3a6ebc1a6050ecdd687ba5d99f8d56700..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_ltx_latent_upsample.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index 831077e13f608a2a856f1f6ac748c6fd3c2e7716..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/modeling_latent_upsampler.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/modeling_latent_upsampler.py deleted file mode 100644 index 6dce792a2b43757015b0aced8a11fd27a8baed86..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/modeling_latent_upsampler.py +++ /dev/null @@ -1,188 +0,0 @@ -# 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 typing import Optional - -import torch - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.modeling_utils import ModelMixin - - -class ResBlock(torch.nn.Module): - def __init__(self, channels: int, mid_channels: Optional[int] = 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 - - -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 LTXLatentUpsamplerModel(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 = 512, - num_blocks_per_stage: int = 4, - dims: int = 3, - spatial_upsample: bool = True, - temporal_upsample: bool = False, - ): - 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: - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_ltx.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_ltx.py deleted file mode 100644 index 3b58b4a45a453c56dc9b819bb8edfc52d4253bad..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_ltx.py +++ /dev/null @@ -1,846 +0,0 @@ -# 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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -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 -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 LTXPipeline - >>> from diffusers.utils import export_to_video - - >>> pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16) - >>> pipe.to("cuda") - - >>> 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" - - >>> video = pipe( - ... prompt=prompt, - ... negative_prompt=negative_prompt, - ... width=704, - ... height=480, - ... num_frames=161, - ... num_inference_steps=50, - ... ).frames[0] - >>> export_to_video(video, "output.mp4", fps=24) - ``` -""" - - -# 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 - - -# 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 LTXPipeline(DiffusionPipeline, FromSingleFileMixin, LTXVideoLoraLoaderMixin): - 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). - """ - - 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, - ) - - 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") 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 - ) - - def _get_t5_prompt_embeds( - self, - prompt: Union[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 = 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.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, - 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, - max_sequence_length: int = 128, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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 - - 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 - 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 - - def prepare_latents( - self, - 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, - ) -> torch.Tensor: - if latents is not None: - 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 - - @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: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: int = 512, - width: int = 704, - num_frames: int = 161, - frame_rate: int = 25, - num_inference_steps: int = 50, - 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", - 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"], - max_sequence_length: int = 128, - ): - 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`, defaults to `512`): - The height in pixels of the generated image. This is set to 480 by default for the best results. - width (`int`, defaults to `704`): - The width in pixels of the generated image. This is set to 848 by default for the best results. - num_frames (`int`, defaults to `161`): - The number of video frames to generate - 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`, defaults to `3 `): - 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://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of - [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). - Guidance rescale factor should fix overexposure when 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*): - 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 ge 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.LTXPipelineOutput`] 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 `128 `): - Maximum sequence length to use with the `prompt`. - - 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 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) - - # 4. Prepare latent variables - 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, - torch.float32, - device, - generator, - latents, - ) - - # 5. Prepare timesteps - 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 - video_sequence_length = latent_num_frames * latent_height * latent_width - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) - mu = calculate_shift( - video_sequence_length, - 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, - 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, - ) - - # 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) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=timestep, - encoder_attention_mask=prompt_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - rope_interpolation_scale=rope_interpolation_scale, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - noise_pred = noise_pred.float() - - 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: - # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf - noise_pred = rescale_noise_cfg( - noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale - ) - - # compute the previous noisy sample x_t -> x_t-1 - 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, 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() - - if output_type == "latent": - video = latents - else: - 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 - ) - 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) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return LTXPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_ltx_condition.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_ltx_condition.py deleted file mode 100644 index fa9ee4fc7b87be0c484db59e633c9788e827ae45..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_ltx_condition.py +++ /dev/null @@ -1,1272 +0,0 @@ -# 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 inspect -from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import PIL.Image -import torch -from transformers import T5EncoderModel, T5TokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput -from ...loaders import FromSingleFileMixin, LTXVideoLoraLoaderMixin -from ...models.autoencoders import AutoencoderKLLTXVideo -from ...models.transformers import LTXVideoTransformer3DModel -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 .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.pipelines.ltx.pipeline_ltx_condition import LTXConditionPipeline, LTXVideoCondition - >>> from diffusers.utils import export_to_video, load_video, load_image - - >>> pipe = LTXConditionPipeline.from_pretrained("Lightricks/LTX-Video-0.9.5", torch_dtype=torch.bfloat16) - >>> pipe.to("cuda") - - >>> # Load input image and video - >>> video = load_video( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cosmos/cosmos-video2world-input-vid.mp4" - ... ) - >>> image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cosmos/cosmos-video2world-input.jpg" - ... ) - - >>> # Create conditioning objects - >>> condition1 = LTXVideoCondition( - ... image=image, - ... frame_index=0, - ... ) - >>> condition2 = LTXVideoCondition( - ... video=video, - ... frame_index=80, - ... ) - - >>> 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 = "worst quality, inconsistent motion, blurry, jittery, distorted" - - >>> # Generate video - >>> generator = torch.Generator("cuda").manual_seed(0) - >>> # Text-only conditioning is also supported without the need to pass `conditions` - >>> video = pipe( - ... conditions=[condition1, condition2], - ... prompt=prompt, - ... negative_prompt=negative_prompt, - ... width=768, - ... height=512, - ... num_frames=161, - ... num_inference_steps=40, - ... generator=generator, - ... ).frames[0] - - >>> export_to_video(video, "output.mp4", fps=24) - ``` -""" - - -@dataclass -class LTXVideoCondition: - """ - Defines a single frame-conditioning item for LTX Video - a single frame or a sequence of frames. - - Attributes: - image (`PIL.Image.Image`): - The image to condition the video on. - 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. - strength (`float`, defaults to `1.0`): - 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 - frame_index: int = 0 - strength: float = 1.0 - - -# from LTX-Video/ltx_video/schedulers/rf.py -def linear_quadratic_schedule(num_steps, threshold_noise=0.025, linear_steps=None): - if linear_steps is None: - linear_steps = num_steps // 2 - if num_steps < 2: - return torch.tensor([1.0]) - linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)] - threshold_noise_step_diff = linear_steps - threshold_noise * num_steps - quadratic_steps = num_steps - linear_steps - quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2) - linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2) - const = quadratic_coef * (linear_steps**2) - quadratic_sigma_schedule = [ - quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, num_steps) - ] - sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0] - sigma_schedule = [1.0 - x for x in sigma_schedule] - return torch.tensor(sigma_schedule[:-1]) - - -# 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 - - -# 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" -): - 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.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 LTXConditionPipeline(DiffusionPipeline, FromSingleFileMixin, LTXVideoLoraLoaderMixin): - r""" - Pipeline for text/image/video-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). - """ - - 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, - ) - - 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") 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 - - def _get_t5_prompt_embeds( - self, - prompt: Union[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 = 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), attention_mask=prompt_attention_mask)[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.mochi.pipeline_mochi.MochiPipeline.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: 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: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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 - - def check_inputs( - self, - prompt, - conditions, - image, - video, - frame_index, - strength, - denoise_strength, - 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}." - ) - - if conditions is not None and (image is not None or video is not None): - raise ValueError("If `conditions` is provided, `image` and `video` must not be provided.") - - if conditions is None: - if isinstance(image, list) and isinstance(frame_index, list) and len(image) != len(frame_index): - raise ValueError( - "If `conditions` is not provided, `image` and `frame_index` must be of the same length." - ) - elif isinstance(image, list) and isinstance(strength, list) and len(image) != len(strength): - raise ValueError("If `conditions` is not provided, `image` and `strength` must be of the same length.") - elif isinstance(video, list) and isinstance(frame_index, list) and len(video) != len(frame_index): - raise ValueError( - "If `conditions` is not provided, `video` and `frame_index` must be of the same length." - ) - elif isinstance(video, list) and isinstance(strength, list) and len(video) != len(strength): - raise ValueError("If `conditions` is not provided, `video` and `strength` must be of the same length.") - - if denoise_strength < 0 or denoise_strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {denoise_strength}") - - @staticmethod - def _prepare_video_ids( - batch_size: int, - num_frames: int, - height: int, - width: int, - patch_size: int = 1, - patch_size_t: int = 1, - device: torch.device = None, - ) -> torch.Tensor: - latent_sample_coords = torch.meshgrid( - torch.arange(0, num_frames, patch_size_t, device=device), - torch.arange(0, height, patch_size, device=device), - torch.arange(0, width, patch_size, 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.reshape(batch_size, -1, num_frames * height * width) - - return latent_coords - - @staticmethod - def _scale_video_ids( - video_ids: torch.Tensor, - scale_factor: int = 32, - scale_factor_t: int = 8, - frame_index: int = 0, - device: torch.device = None, - ) -> torch.Tensor: - scaled_latent_coords = ( - video_ids - * torch.tensor([scale_factor_t, scale_factor, scale_factor], device=video_ids.device)[None, :, None] - ) - scaled_latent_coords[:, 0] = (scaled_latent_coords[:, 0] + 1 - scale_factor_t).clamp(min=0) - scaled_latent_coords[:, 0] += frame_index - - return scaled_latent_coords - - @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 trim_conditioning_sequence(self, start_frame: int, sequence_num_frames: int, target_num_frames: 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 - - @staticmethod - def add_noise_to_image_conditioning_latents( - t: float, - init_latents: torch.Tensor, - latents: torch.Tensor, - noise_scale: float, - conditioning_mask: torch.Tensor, - generator, - eps=1e-6, - ): - """ - Add timestep-dependent noise to the hard-conditioning latents. This helps with motion continuity, especially - when conditioned on a single frame. - """ - noise = randn_tensor( - latents.shape, - generator=generator, - device=latents.device, - dtype=latents.dtype, - ) - # Add noise only to hard-conditioning latents (conditioning_mask = 1.0) - need_to_noise = (conditioning_mask > 1.0 - eps).unsqueeze(-1) - noised_latents = init_latents + noise_scale * noise * (t**2) - latents = torch.where(need_to_noise, noised_latents, latents) - return latents - - def prepare_latents( - self, - conditions: Optional[List[torch.Tensor]] = None, - condition_strength: Optional[List[float]] = None, - condition_frame_index: Optional[List[int]] = 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]: - 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 - - shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width) - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - if latents is not None and sigma is not None: - if latents.shape != shape: - raise ValueError( - f"Latents shape {latents.shape} does not match expected shape {shape}. Please check the input." - ) - latents = latents.to(device=device, dtype=dtype) - sigma = sigma.to(device=device, dtype=dtype) - latents = sigma * noise + (1 - sigma) * latents - else: - latents = noise - - if len(conditions) > 0: - condition_latent_frames_mask = torch.zeros( - (batch_size, num_latent_frames), device=device, dtype=torch.float32 - ) - - extra_conditioning_latents = [] - extra_conditioning_video_ids = [] - extra_conditioning_mask = [] - extra_conditioning_num_latents = 0 - for data, strength, frame_index in zip(conditions, condition_strength, condition_frame_index): - condition_latents = retrieve_latents(self.vae.encode(data), generator=generator) - condition_latents = self._normalize_latents( - condition_latents, self.vae.latents_mean, self.vae.latents_std - ).to(device, dtype=dtype) - - num_data_frames = data.size(2) - num_cond_frames = condition_latents.size(2) - - if frame_index == 0: - latents[:, :, :num_cond_frames] = torch.lerp( - latents[:, :, :num_cond_frames], condition_latents, strength - ) - condition_latent_frames_mask[:, :num_cond_frames] = strength - - else: - if num_data_frames > 1: - if num_cond_frames < num_prefix_latent_frames: - raise ValueError( - f"Number of latent frames must be at least {num_prefix_latent_frames} but got {num_data_frames}." - ) - - if num_cond_frames > num_prefix_latent_frames: - start_frame = frame_index // self.vae_temporal_compression_ratio + num_prefix_latent_frames - end_frame = start_frame + num_cond_frames - num_prefix_latent_frames - latents[:, :, start_frame:end_frame] = torch.lerp( - latents[:, :, start_frame:end_frame], - condition_latents[:, :, num_prefix_latent_frames:], - strength, - ) - condition_latent_frames_mask[:, start_frame:end_frame] = strength - condition_latents = condition_latents[:, :, :num_prefix_latent_frames] - - noise = randn_tensor(condition_latents.shape, generator=generator, device=device, dtype=dtype) - condition_latents = torch.lerp(noise, condition_latents, strength) - - condition_video_ids = self._prepare_video_ids( - batch_size, - condition_latents.size(2), - latent_height, - latent_width, - patch_size=self.transformer_spatial_patch_size, - patch_size_t=self.transformer_temporal_patch_size, - device=device, - ) - condition_video_ids = self._scale_video_ids( - condition_video_ids, - scale_factor=self.vae_spatial_compression_ratio, - scale_factor_t=self.vae_temporal_compression_ratio, - frame_index=frame_index, - device=device, - ) - condition_latents = self._pack_latents( - condition_latents, - self.transformer_spatial_patch_size, - self.transformer_temporal_patch_size, - ) - condition_conditioning_mask = torch.full( - condition_latents.shape[:2], strength, device=device, dtype=dtype - ) - - extra_conditioning_latents.append(condition_latents) - extra_conditioning_video_ids.append(condition_video_ids) - extra_conditioning_mask.append(condition_conditioning_mask) - extra_conditioning_num_latents += condition_latents.size(1) - - video_ids = self._prepare_video_ids( - batch_size, - num_latent_frames, - latent_height, - latent_width, - patch_size_t=self.transformer_temporal_patch_size, - patch_size=self.transformer_spatial_patch_size, - device=device, - ) - if len(conditions) > 0: - conditioning_mask = condition_latent_frames_mask.gather(1, video_ids[:, 0]) - else: - conditioning_mask, extra_conditioning_num_latents = None, 0 - video_ids = self._scale_video_ids( - video_ids, - scale_factor=self.vae_spatial_compression_ratio, - scale_factor_t=self.vae_temporal_compression_ratio, - frame_index=0, - device=device, - ) - latents = self._pack_latents( - latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size - ) - - if len(conditions) > 0 and len(extra_conditioning_latents) > 0: - latents = torch.cat([*extra_conditioning_latents, latents], dim=1) - video_ids = torch.cat([*extra_conditioning_video_ids, video_ids], dim=2) - conditioning_mask = torch.cat([*extra_conditioning_mask, conditioning_mask], dim=1) - - return latents, conditioning_mask, video_ids, extra_conditioning_num_latents - - def get_timesteps(self, sigmas, timesteps, num_inference_steps, strength): - num_steps = min(int(num_inference_steps * strength), num_inference_steps) - start_index = max(num_inference_steps - num_steps, 0) - sigmas = sigmas[start_index:] - timesteps = timesteps[start_index:] - return sigmas, timesteps, num_inference_steps - start_index - - @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: 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, - denoise_strength: float = 1.0, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: int = 512, - width: int = 704, - num_frames: int = 161, - frame_rate: int = 25, - num_inference_steps: int = 50, - 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", - 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"], - max_sequence_length: int = 256, - ): - r""" - Function invoked when calling the pipeline for generation. - - Args: - 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*): - The image or images to condition the video generation. If not provided, one has to pass `video` or - `conditions`. - 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*): - 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*): - 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*): - The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. - instead. - height (`int`, defaults to `512`): - The height in pixels of the generated image. This is set to 480 by default for the best results. - width (`int`, defaults to `704`): - The width in pixels of the generated image. This is set to 848 by default for the best results. - num_frames (`int`, defaults to `161`): - The number of video frames to generate - 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`, defaults to `3 `): - 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://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of - [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). - Guidance rescale factor should fix overexposure when 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*): - 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 ge 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.LTXPipelineOutput`] 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 `128 `): - Maximum sequence length to use with the `prompt`. - - 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 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, - conditions=conditions, - image=image, - video=video, - frame_index=frame_index, - strength=strength, - denoise_strength=denoise_strength, - 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] - - if conditions is not None: - if not isinstance(conditions, list): - conditions = [conditions] - - strength = [condition.strength for condition in conditions] - frame_index = [condition.frame_index for condition in conditions] - image = [condition.image for condition in conditions] - video = [condition.video for condition in conditions] - elif image is not None or video is not None: - if not isinstance(image, list): - image = [image] - num_conditions = 1 - elif isinstance(image, list): - num_conditions = len(image) - if not isinstance(video, list): - video = [video] - num_conditions = 1 - elif isinstance(video, list): - num_conditions = len(video) - - if not isinstance(frame_index, list): - frame_index = [frame_index] * num_conditions - if not isinstance(strength, list): - strength = [strength] * num_conditions - - device = self._execution_device - vae_dtype = self.vae.dtype - - # 3. Prepare text embeddings & conditioning image/video - ( - 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) - - conditioning_tensors = [] - is_conditioning_image_or_video = image is not None or video is not None - if is_conditioning_image_or_video: - for condition_image, condition_video, condition_frame_index, condition_strength in zip( - image, video, frame_index, strength - ): - if condition_image is not None: - condition_tensor = ( - self.video_processor.preprocess(condition_image, height, width) - .unsqueeze(2) - .to(device, dtype=vae_dtype) - ) - elif condition_video is not None: - condition_tensor = self.video_processor.preprocess_video(condition_video, height, width) - num_frames_input = condition_tensor.size(2) - num_frames_output = self.trim_conditioning_sequence( - condition_frame_index, num_frames_input, num_frames - ) - condition_tensor = condition_tensor[:, :, :num_frames_output] - condition_tensor = condition_tensor.to(device, dtype=vae_dtype) - else: - raise ValueError("Either `image` or `video` must be provided for conditioning.") - - if condition_tensor.size(2) % self.vae_temporal_compression_ratio != 1: - raise ValueError( - f"Number of frames in the video must be of the form (k * {self.vae_temporal_compression_ratio} + 1) " - f"but got {condition_tensor.size(2)} frames." - ) - conditioning_tensors.append(condition_tensor) - - # 4. Prepare timesteps - 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) - sigmas = self.scheduler.sigmas - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - latent_sigma = None - if denoise_strength < 1: - sigmas, timesteps, num_inference_steps = self.get_timesteps( - sigmas, timesteps, num_inference_steps, denoise_strength - ) - latent_sigma = sigmas[:1].repeat(batch_size * num_videos_per_prompt) - - self._num_timesteps = len(timesteps) - - # 5. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - latents, conditioning_mask, video_coords, extra_conditioning_num_latents = self.prepare_latents( - conditioning_tensors, - strength, - frame_index, - batch_size=batch_size * num_videos_per_prompt, - num_channels_latents=num_channels_latents, - height=height, - width=width, - num_frames=num_frames, - sigma=latent_sigma, - latents=latents, - generator=generator, - device=device, - dtype=torch.float32, - ) - - video_coords = video_coords.float() - video_coords[:, 0] = video_coords[:, 0] * (1.0 / frame_rate) - - init_latents = latents.clone() if is_conditioning_image_or_video else None - - if self.do_classifier_free_guidance: - video_coords = torch.cat([video_coords, video_coords], dim=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 - - self._current_timestep = t - - if image_cond_noise_scale > 0 and init_latents is not None: - # Add timestep-dependent noise to the hard-conditioning latents - # This helps with motion continuity, especially when conditioned on a single frame - latents = self.add_noise_to_image_conditioning_latents( - t / 1000.0, - init_latents, - latents, - image_cond_noise_scale, - conditioning_mask, - generator, - ) - - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - if is_conditioning_image_or_video: - conditioning_mask_model_input = ( - torch.cat([conditioning_mask, conditioning_mask]) - if self.do_classifier_free_guidance - else conditioning_mask - ) - latent_model_input = 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]).unsqueeze(-1).float() - if is_conditioning_image_or_video: - timestep = torch.min(timestep, (1 - conditioning_mask_model_input) * 1000.0) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=timestep, - encoder_attention_mask=prompt_attention_mask, - video_coords=video_coords, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - - 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) - timestep, _ = timestep.chunk(2) - - if self.guidance_rescale > 0: - # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf - noise_pred = rescale_noise_cfg( - noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale - ) - - denoised_latents = self.scheduler.step( - -noise_pred, t, latents, per_token_timesteps=timestep, return_dict=False - )[0] - if is_conditioning_image_or_video: - tokens_to_denoise_mask = (t / 1000 - 1e-6 < (1.0 - conditioning_mask)).unsqueeze(-1) - latents = torch.where(tokens_to_denoise_mask, denoised_latents, latents) - else: - latents = denoised_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) - - # 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 is_conditioning_image_or_video: - latents = latents[:, extra_conditioning_num_latents:] - - latents = self._unpack_latents( - latents, - latent_num_frames, - latent_height, - latent_width, - self.transformer_spatial_patch_size, - self.transformer_temporal_patch_size, - ) - - if output_type == "latent": - video = latents - else: - latents = self._denormalize_latents( - latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor - ) - 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 - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_ltx_image2video.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_ltx_image2video.py deleted file mode 100644 index 99412b6962a3fd814c1a89580a89a252af34c6dc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_ltx_image2video.py +++ /dev/null @@ -1,943 +0,0 @@ -# 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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import torch -from transformers import T5EncoderModel, T5TokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput -from ...loaders import FromSingleFileMixin, LTXVideoLoraLoaderMixin -from ...models.autoencoders import AutoencoderKLLTXVideo -from ...models.transformers import LTXVideoTransformer3DModel -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 .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 LTXImageToVideoPipeline - >>> from diffusers.utils import export_to_video, load_image - - >>> pipe = LTXImageToVideoPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16) - >>> pipe.to("cuda") - - >>> 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. Flames engulf the structure, with smoke billowing into the air. Firefighters in protective gear rush to the scene, a fire truck labeled '38' visible behind them. The girl's neutral expression contrasts sharply with the chaos of the fire, creating a poignant and emotionally charged scene." - >>> negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted" - - >>> video = pipe( - ... image=image, - ... prompt=prompt, - ... negative_prompt=negative_prompt, - ... width=704, - ... height=480, - ... num_frames=161, - ... num_inference_steps=50, - ... ).frames[0] - >>> export_to_video(video, "output.mp4", fps=24) - ``` -""" - - -# 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 - - -# 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" -): - 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.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 LTXImageToVideoPipeline(DiffusionPipeline, FromSingleFileMixin, LTXVideoLoraLoaderMixin): - r""" - Pipeline for image-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). - """ - - 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, - ) - - 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") 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 - - def _get_t5_prompt_embeds( - self, - prompt: Union[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 = 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.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, - 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, - max_sequence_length: int = 128, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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 - - # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline.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.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, - image: Optional[torch.Tensor] = 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, - ) -> 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: - 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]) - for i in range(batch_size) - ] - else: - init_latents = [ - retrieve_latents(self.vae.encode(img.unsqueeze(0).unsqueeze(2)), generator) 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) - 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) - 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 - - @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: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: int = 512, - width: int = 704, - num_frames: int = 161, - frame_rate: int = 25, - num_inference_steps: int = 50, - 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", - 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"], - max_sequence_length: int = 128, - ): - 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`, defaults to `512`): - The height in pixels of the generated image. This is set to 480 by default for the best results. - width (`int`, defaults to `704`): - The width in pixels of the generated image. This is set to 848 by default for the best results. - num_frames (`int`, defaults to `161`): - The number of video frames to generate - 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`, defaults to `3 `): - 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://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of - [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). - Guidance rescale factor should fix overexposure when 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*): - 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 ge 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.LTXPipelineOutput`] 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 `128 `): - Maximum sequence length to use with the `prompt`. - - 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 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) - - # 4. Prepare latent variables - 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, - torch.float32, - device, - generator, - latents, - ) - - if self.do_classifier_free_guidance: - conditioning_mask = torch.cat([conditioning_mask, conditioning_mask]) - - # 5. Prepare timesteps - 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 - video_sequence_length = latent_num_frames * latent_height * latent_width - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) - mu = calculate_shift( - video_sequence_length, - 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, - 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, - ) - - # 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) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=timestep, - encoder_attention_mask=prompt_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - rope_interpolation_scale=rope_interpolation_scale, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - noise_pred = noise_pred.float() - - 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) - timestep, _ = timestep.chunk(2) - - if self.guidance_rescale > 0: - # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf - noise_pred = rescale_noise_cfg( - noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale - ) - - # compute the previous noisy sample x_t -> x_t-1 - noise_pred = self._unpack_latents( - noise_pred, - 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 = noise_pred[:, :, 1:] - noise_latents = latents[:, :, 1:] - pred_latents = self.scheduler.step(noise_pred, 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 - ) - - 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() - - if output_type == "latent": - video = latents - else: - 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 - ) - latents = latents.to(prompt_embeds.dtype) - - if not self.vae.config.timestep_conditioning: - timestep = None - else: - noise = torch.randn(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/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py deleted file mode 100644 index 284f33b32631f1cd2db3ff9e07bf3e6fa24f8221..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py +++ /dev/null @@ -1,277 +0,0 @@ -# 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 typing import List, Optional, Union - -import torch - -from ...image_processor import PipelineImageInput -from ...models import AutoencoderKLLTXVideo -from ...utils import get_logger -from ...utils.torch_utils import randn_tensor -from ...video_processor import VideoProcessor -from ..pipeline_utils import DiffusionPipeline -from .modeling_latent_upsampler import LTXLatentUpsamplerModel -from .pipeline_output import LTXPipelineOutput - - -logger = get_logger(__name__) # pylint: disable=invalid-name - - -# 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" -): - 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 LTXLatentUpsamplePipeline(DiffusionPipeline): - model_cpu_offload_seq = "" - - def __init__( - self, - vae: AutoencoderKLLTXVideo, - latent_upsampler: LTXLatentUpsamplerModel, - ) -> 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: Optional[torch.Tensor] = 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, - ) -> torch.Tensor: - if latents is not None: - 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) - 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 - - @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 enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - def check_inputs(self, video, height, width, latents): - 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.") - - @torch.no_grad() - def __call__( - self, - video: Optional[List[PipelineImageInput]] = 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, - adain_factor: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - output_type: Optional[str] = "pil", - return_dict: bool = True, - ): - self.check_inputs( - video=video, - height=height, - width=width, - latents=latents, - ) - - 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 = self.prepare_latents( - video=video, - batch_size=batch_size, - dtype=torch.float32, - device=device, - generator=generator, - latents=latents, - ) - - 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 output_type == "latent": - latents = self._normalize_latents( - latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor - ) - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_output.py deleted file mode 100644 index 36ec3ea884a26b4add4b0729867aaa593cce655c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/ltx/pipeline_output.py +++ /dev/null @@ -1,20 +0,0 @@ -from dataclasses import dataclass - -import torch - -from diffusers.utils import BaseOutput - - -@dataclass -class LTXPipelineOutput(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)`. - """ - - frames: torch.Tensor diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina/__init__.py deleted file mode 100644 index a19dc7e94641f59cad8588e867412829cd4fa793..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -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_lumina"] = ["LuminaPipeline", "LuminaText2ImgPipeline"] - -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_lumina import LuminaPipeline, LuminaText2ImgPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index d4a1aaa909484c8dc73b279efb761b3c115d4a4f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina/__pycache__/pipeline_lumina.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina/__pycache__/pipeline_lumina.cpython-310.pyc deleted file mode 100644 index 66624a7086e671e8a1eb003544ab7fb95dc8bf9f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina/__pycache__/pipeline_lumina.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina/pipeline_lumina.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina/pipeline_lumina.py deleted file mode 100644 index 2067444fa0df316672e4ccb505ab1a16eedb1f9f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina/pipeline_lumina.py +++ /dev/null @@ -1,953 +0,0 @@ -# Copyright 2025 Alpha-VLLM 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 inspect -import math -import re -import urllib.parse as ul -from typing import Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import GemmaPreTrainedModel, GemmaTokenizer, GemmaTokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import VaeImageProcessor -from ...models import AutoencoderKL -from ...models.embeddings import get_2d_rotary_pos_embed_lumina -from ...models.transformers.lumina_nextdit2d import LuminaNextDiT2DModel -from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import ( - BACKENDS_MAPPING, - deprecate, - is_bs4_available, - is_ftfy_available, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -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_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import LuminaPipeline - - >>> pipe = LuminaPipeline.from_pretrained("Alpha-VLLM/Lumina-Next-SFT-diffusers", torch_dtype=torch.bfloat16) - >>> # Enable memory optimizations. - >>> pipe.enable_model_cpu_offload() - - >>> prompt = "Upper body of a young woman in a Victorian-era outfit with brass goggles and leather straps. Background shows an industrial revolution cityscape with smoky skies and tall, metal structures" - >>> image = pipe(prompt).images[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, - **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 LuminaPipeline(DiffusionPipeline): - r""" - Pipeline for text-to-image generation using Lumina-T2I. - - 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.) - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`GemmaPreTrainedModel`]): - Frozen Gemma text-encoder. - tokenizer (`GemmaTokenizer` or `GemmaTokenizerFast`): - Gemma tokenizer. - transformer ([`Transformer2DModel`]): - A text conditioned `Transformer2DModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - """ - - bad_punct_regex = re.compile( - r"[" - + "#®•©™&@·º½¾¿¡§~" - + r"\)" - + r"\(" - + r"\]" - + r"\[" - + r"\}" - + r"\{" - + r"\|" - + "\\" - + r"\/" - + r"\*" - + r"]{1,}" - ) # noqa - - _optional_components = [] - model_cpu_offload_seq = "text_encoder->transformer->vae" - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - ] - - def __init__( - self, - transformer: LuminaNextDiT2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: GemmaPreTrainedModel, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - ) - self.vae_scale_factor = 8 - self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) - self.max_sequence_length = 256 - self.default_sample_size = ( - self.transformer.config.sample_size - if hasattr(self, "transformer") and self.transformer is not None - else 128 - ) - self.default_image_size = self.default_sample_size * self.vae_scale_factor - - def _get_gemma_prompt_embeds( - self, - prompt: Union[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 = device or self._execution_device - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - text_inputs = self.tokenizer( - prompt, - pad_to_multiple_of=8, - max_length=self.max_sequence_length, - truncation=True, - padding=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids.to(device) - untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids.to(device) - - 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.max_sequence_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because Gemma can only handle sequences up to" - f" {self.max_sequence_length} tokens: {removed_text}" - ) - - prompt_attention_mask = text_inputs.attention_mask.to(device) - prompt_embeds = self.text_encoder( - text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True - ) - prompt_embeds = prompt_embeds.hidden_states[-2] - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - elif self.transformer is not None: - dtype = self.transformer.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.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) - prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1) - prompt_attention_mask = prompt_attention_mask.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds, prompt_attention_mask - - # Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt - def encode_prompt( - self, - prompt: Union[str, List[str]], - do_classifier_free_guidance: bool = True, - negative_prompt: Union[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, - clean_caption: bool = False, - **kwargs, - ): - 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 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 "". - 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 - device: (`torch.device`, *optional*): - 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. For Lumina-T2I, it's should be the embeddings of the "" string. - clean_caption (`bool`, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - max_sequence_length (`int`, defaults to 256): Maximum sequence length to use for the prompt. - """ - if device is None: - device = 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_images_per_prompt=num_images_per_prompt, - device=device, - clean_caption=clean_caption, - ) - - # Get negative embeddings for classifier free guidance - if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt if negative_prompt is not None else "" - - # Normalize str to list - 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 isinstance(negative_prompt, str): - negative_prompt = [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`." - ) - # Padding negative prompt to the same length with prompt - prompt_max_length = prompt_embeds.shape[1] - negative_text_inputs = self.tokenizer( - negative_prompt, - padding="max_length", - max_length=prompt_max_length, - truncation=True, - return_tensors="pt", - ) - negative_text_input_ids = negative_text_inputs.input_ids.to(device) - negative_prompt_attention_mask = negative_text_inputs.attention_mask.to(device) - # Get the negative prompt embeddings - negative_prompt_embeds = self.text_encoder( - negative_text_input_ids, - attention_mask=negative_prompt_attention_mask, - output_hidden_states=True, - ) - - negative_dtype = self.text_encoder.dtype - negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2] - _, seq_len, _ = negative_prompt_embeds.shape - - negative_prompt_embeds = negative_prompt_embeds.to(dtype=negative_dtype, device=device) - # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method - 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) - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1) - negative_prompt_attention_mask = negative_prompt_attention_mask.view( - batch_size * num_images_per_prompt, -1 - ) - - return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - negative_prompt, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - callback_on_step_end_tensor_inputs=None, - ): - if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0: - raise ValueError( - f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} 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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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("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}." - ) - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - 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 - - @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 do_classifier_free_guidance(self): - return self._guidance_scale > 1 - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - width: Optional[int] = None, - height: Optional[int] = 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", - 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]: - """ - 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 - 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*): - 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 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. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - height (`int`, *optional*, defaults to self.unet.config.sample_size): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to self.unet.config.sample_size): - The width in pixels of the generated image. - 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*): - 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 ge 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.Tensor`, *optional*): - Pre-generated negative text embeddings. For Lumina-T2I this negative prompt should be "". If not - provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Pre-generated attention mask for negative text embeddings. - 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.IFPipelineOutput`] instead of a plain tuple. - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - max_sequence_length (`int` defaults to 120): - Maximum sequence length to use with the `prompt`. - 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. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where 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, - height, - width, - 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, - ) - - self._guidance_scale = guidance_scale - - cross_attention_kwargs = {} - - # 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 proportional_attn: - cross_attention_kwargs["base_sequence_length"] = (self.default_image_size // 16) ** 2 - - scaling_factor = math.sqrt(width * height / self.default_image_size**2) - - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt, - do_classifier_free_guidance, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - ) - if do_classifier_free_guidance: - prompt_embeds = torch.cat([prompt_embeds, negative_prompt_embeds], dim=0) - 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) - - # 5. Prepare latents. - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - self._num_timesteps = len(timesteps) - - # 6. Denoising loop - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - current_timestep = t - if not torch.is_tensor(current_timestep): - # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = latent_model_input.device.type == "mps" - is_npu = latent_model_input.device.type == "npu" - if isinstance(current_timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - current_timestep = torch.tensor( - [current_timestep], - dtype=dtype, - device=latent_model_input.device, - ) - elif len(current_timestep.shape) == 0: - current_timestep = current_timestep[None].to(latent_model_input.device) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - current_timestep = current_timestep.expand(latent_model_input.shape[0]) - - # reverse the timestep since Lumina uses t=0 as the noise and t=1 as the image - current_timestep = 1 - current_timestep / self.scheduler.config.num_train_timesteps - - # prepare image_rotary_emb for positional encoding - # dynamic scaling_factor for different resolution. - # NOTE: For `Time-aware` denosing mechanism from Lumina-Next - # https://huggingface.co/papers/2406.18583, Sec 2.3 - # NOTE: We should compute different image_rotary_emb with different timestep. - if current_timestep[0] < scaling_watershed: - linear_factor = scaling_factor - ntk_factor = 1.0 - else: - linear_factor = 1.0 - ntk_factor = scaling_factor - image_rotary_emb = get_2d_rotary_pos_embed_lumina( - self.transformer.head_dim, - 384, - 384, - linear_factor=linear_factor, - ntk_factor=ntk_factor, - ) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=current_timestep, - encoder_hidden_states=prompt_embeds, - encoder_mask=prompt_attention_mask, - image_rotary_emb=image_rotary_emb, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - noise_pred = noise_pred.chunk(2, dim=1)[0] - - # perform guidance scale - # NOTE: For exact reproducibility reasons, we apply classifier-free guidance on only - # three channels by default. The standard approach to cfg applies it to all channels. - # This can be done by uncommenting the following line and commenting-out the line following that. - # eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:] - if do_classifier_free_guidance: - noise_pred_eps, noise_pred_rest = noise_pred[:, :3], noise_pred[:, 3:] - noise_pred_cond_eps, noise_pred_uncond_eps = torch.split( - noise_pred_eps, len(noise_pred_eps) // 2, dim=0 - ) - noise_pred_half = noise_pred_uncond_eps + guidance_scale * ( - noise_pred_cond_eps - noise_pred_uncond_eps - ) - noise_pred_eps = torch.cat([noise_pred_half, noise_pred_half], dim=0) - - noise_pred = torch.cat([noise_pred_eps, noise_pred_rest], dim=1) - noise_pred, _ = noise_pred.chunk(2, dim=0) - - # compute the previous noisy sample x_t -> x_t-1 - latents_dtype = latents.dtype - noise_pred = -noise_pred - 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) - - progress_bar.update() - - 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) - - if XLA_AVAILABLE: - xm.mark_step() - - if not output_type == "latent": - latents = latents / self.vae.config.scaling_factor - image = self.vae.decode(latents, return_dict=False)[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 ImagePipelineOutput(images=image) - - -class LuminaText2ImgPipeline(LuminaPipeline): - def __init__( - self, - transformer: LuminaNextDiT2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: GemmaPreTrainedModel, - tokenizer: Union[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) - super().__init__( - transformer=transformer, - scheduler=scheduler, - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina2/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina2/__init__.py deleted file mode 100644 index b1d6bfeb0d5838b1ce22104dc737106b23bc59f3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina2/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -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_lumina2"] = ["Lumina2Pipeline", "Lumina2Text2ImgPipeline"] - -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_lumina2 import Lumina2Pipeline, Lumina2Text2ImgPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina2/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina2/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 9e6f9a02d3565f245a5ad40f5f84d6ada6c76f9d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina2/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina2/__pycache__/pipeline_lumina2.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina2/__pycache__/pipeline_lumina2.cpython-310.pyc deleted file mode 100644 index 5019d879f473ad4c534158194a1d073e1c9c679c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina2/__pycache__/pipeline_lumina2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina2/pipeline_lumina2.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina2/pipeline_lumina2.py deleted file mode 100644 index 0fa0fe97734cab8076bf4a0fc59402774b48396f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/lumina2/pipeline_lumina2.py +++ /dev/null @@ -1,790 +0,0 @@ -# Copyright 2025 Alpha-VLLM 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, Tuple, Union - -import numpy as np -import torch -from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast - -from ...image_processor import VaeImageProcessor -from ...loaders import Lumina2LoraLoaderMixin -from ...models import AutoencoderKL -from ...models.transformers.transformer_lumina2 import Lumina2Transformer2DModel -from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import ( - deprecate, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -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 Lumina2Pipeline - - >>> pipe = Lumina2Pipeline.from_pretrained("Alpha-VLLM/Lumina-Image-2.0", torch_dtype=torch.bfloat16) - >>> # Enable memory optimizations. - >>> pipe.enable_model_cpu_offload() - - >>> prompt = "Upper body of a young woman in a Victorian-era outfit with brass goggles and leather straps. Background shows an industrial revolution cityscape with smoky skies and tall, metal structures" - >>> image = pipe(prompt).images[0] - ``` -""" - - -# 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 Lumina2Pipeline(DiffusionPipeline, Lumina2LoraLoaderMixin): - r""" - Pipeline for text-to-image generation using Lumina-T2I. - - 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.) - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`Gemma2PreTrainedModel`]): - Frozen Gemma2 text-encoder. - tokenizer (`GemmaTokenizer` or `GemmaTokenizerFast`): - Gemma tokenizer. - transformer ([`Transformer2DModel`]): - A text conditioned `Transformer2DModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - """ - - _optional_components = [] - _callback_tensor_inputs = ["latents", "prompt_embeds"] - model_cpu_offload_seq = "text_encoder->transformer->vae" - - def __init__( - self, - transformer: Lumina2Transformer2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: Gemma2PreTrainedModel, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - ) - self.vae_scale_factor = 8 - self.default_sample_size = ( - self.transformer.config.sample_size - if hasattr(self, "transformer") and self.transformer is not None - else 128 - ) - self.default_image_size = self.default_sample_size * self.vae_scale_factor - self.system_prompt = "You are an assistant designed to generate superior images with the superior degree of image-text alignment based on textual prompts or user prompts." - - self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) - - if getattr(self, "tokenizer", None) is not None: - self.tokenizer.padding_side = "right" - - def _get_gemma_prompt_embeds( - self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, - max_sequence_length: int = 256, - ) -> Tuple[torch.Tensor, torch.Tensor]: - device = device or self._execution_device - prompt = [prompt] if isinstance(prompt, str) else prompt - 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) - untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids.to(device) - - 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 Gemma can only handle sequences up to" - f" {max_sequence_length} tokens: {removed_text}" - ) - - prompt_attention_mask = text_inputs.attention_mask.to(device) - prompt_embeds = self.text_encoder( - text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True - ) - prompt_embeds = prompt_embeds.hidden_states[-2] - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - elif self.transformer is not None: - dtype = self.transformer.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - - _, seq_len, _ = prompt_embeds.shape - - return prompt_embeds, prompt_attention_mask - - # Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt - def encode_prompt( - self, - prompt: Union[str, List[str]], - do_classifier_free_guidance: bool = True, - negative_prompt: Union[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, - max_sequence_length: int = 256, - ) -> 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 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`). For - Lumina-T2I, this should be "". - 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 - device: (`torch.device`, *optional*): - 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. For Lumina-T2I, it's should be the embeddings of the "" string. - max_sequence_length (`int`, defaults to `256`): - Maximum sequence length to use for the prompt. - """ - if device is None: - device = 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 system_prompt is None: - system_prompt = self.system_prompt - if prompt is not None: - prompt = [system_prompt + " " + p for p in prompt] - - if prompt_embeds is None: - prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( - prompt=prompt, - device=device, - max_sequence_length=max_sequence_length, - ) - - batch_size, 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) - prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1) - prompt_attention_mask = prompt_attention_mask.view(batch_size * num_images_per_prompt, -1) - - # Get negative embeddings for classifier free guidance - if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt if negative_prompt is not None else "" - - # Normalize str to list - 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 isinstance(negative_prompt, str): - negative_prompt = [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`." - ) - negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( - prompt=negative_prompt, - device=device, - max_sequence_length=max_sequence_length, - ) - - batch_size, seq_len, _ = negative_prompt_embeds.shape - # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method - 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) - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1) - negative_prompt_attention_mask = negative_prompt_attention_mask.view( - batch_size * num_images_per_prompt, -1 - ) - - return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - negative_prompt, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - ): - if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0: - raise ValueError( - f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} 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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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("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 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}") - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=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_channels_latents, 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." - ) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device) - - return latents - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def attention_kwargs(self): - return self._attention_kwargs - - # 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 - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - width: Optional[int] = None, - height: Optional[int] = 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", - 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, - cfg_trunc_ratio: float = 1.0, - cfg_normalization: bool = True, - max_sequence_length: int = 256, - ) -> Union[ImagePipelineOutput, 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. 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 - 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*): - 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 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. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - height (`int`, *optional*, defaults to self.unet.config.sample_size): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to self.unet.config.sample_size): - The width in pixels of the generated image. - 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*): - 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 ge 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.Tensor`, *optional*): - Pre-generated negative text embeddings. For Lumina-T2I this negative prompt should be "". If not - provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Pre-generated attention mask for negative text embeddings. - 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.IFPipelineOutput`] instead of a plain tuple. - attention_kwargs: - 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. - system_prompt (`str`, *optional*): - The system prompt to use for the image generation. - cfg_trunc_ratio (`float`, *optional*, defaults to `1.0`): - The ratio of the timestep interval to apply normalization-based guidance scale. - cfg_normalization (`bool`, *optional*, defaults to `True`): - Whether to apply normalization-based guidance scale. - max_sequence_length (`int`, defaults to `256`): - Maximum sequence length to use with the `prompt`. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where 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 - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - 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, - max_sequence_length=max_sequence_length, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - ) - - # 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 - ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt, - self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - 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, - system_prompt=system_prompt, - ) - - # 4. Prepare latents. - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 5. Prepare timesteps - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas - image_seq_len = latents.shape[1] - 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, - ) - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - # 6. Denoising loop - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # compute whether apply classifier-free truncation on this timestep - do_classifier_free_truncation = (i + 1) / num_inference_steps > cfg_trunc_ratio - # reverse the timestep since Lumina uses t=0 as the noise and t=1 as the image - current_timestep = 1 - t / self.scheduler.config.num_train_timesteps - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - current_timestep = current_timestep.expand(latents.shape[0]) - - noise_pred_cond = self.transformer( - hidden_states=latents, - timestep=current_timestep, - encoder_hidden_states=prompt_embeds, - encoder_attention_mask=prompt_attention_mask, - return_dict=False, - attention_kwargs=self.attention_kwargs, - )[0] - - # perform normalization-based guidance scale on a truncated timestep interval - if self.do_classifier_free_guidance and not do_classifier_free_truncation: - noise_pred_uncond = self.transformer( - hidden_states=latents, - timestep=current_timestep, - encoder_hidden_states=negative_prompt_embeds, - encoder_attention_mask=negative_prompt_attention_mask, - return_dict=False, - attention_kwargs=self.attention_kwargs, - )[0] - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) - # apply normalization after classifier-free guidance - if cfg_normalization: - cond_norm = torch.norm(noise_pred_cond, dim=-1, keepdim=True) - noise_norm = torch.norm(noise_pred, dim=-1, keepdim=True) - noise_pred = noise_pred * (cond_norm / noise_norm) - else: - noise_pred = noise_pred_cond - - # compute the previous noisy sample x_t -> x_t-1 - latents_dtype = latents.dtype - noise_pred = -noise_pred - 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() - - if not output_type == "latent": - 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) - else: - image = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) - - -class Lumina2Text2ImgPipeline(Lumina2Pipeline): - def __init__( - self, - transformer: Lumina2Transformer2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: Gemma2PreTrainedModel, - tokenizer: Union[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) - super().__init__( - transformer=transformer, - scheduler=scheduler, - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__init__.py deleted file mode 100644 index 168a8276be4e24d4983cd48f75d0bd16f08abf35..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__init__.py +++ /dev/null @@ -1,52 +0,0 @@ -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["marigold_image_processing"] = ["MarigoldImageProcessor"] - _import_structure["pipeline_marigold_depth"] = ["MarigoldDepthOutput", "MarigoldDepthPipeline"] - _import_structure["pipeline_marigold_intrinsics"] = ["MarigoldIntrinsicsOutput", "MarigoldIntrinsicsPipeline"] - _import_structure["pipeline_marigold_normals"] = ["MarigoldNormalsOutput", "MarigoldNormalsPipeline"] - -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 .marigold_image_processing import MarigoldImageProcessor - from .pipeline_marigold_depth import MarigoldDepthOutput, MarigoldDepthPipeline - from .pipeline_marigold_intrinsics import MarigoldIntrinsicsOutput, MarigoldIntrinsicsPipeline - from .pipeline_marigold_normals import MarigoldNormalsOutput, MarigoldNormalsPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index f60c24cd7e7154bfa29b7f964d9eab63abfbf72a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/marigold_image_processing.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/marigold_image_processing.cpython-310.pyc deleted file mode 100644 index ba4a803fa867cc7519121a73c14ddc3daf284c08..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/marigold_image_processing.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/pipeline_marigold_depth.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/pipeline_marigold_depth.cpython-310.pyc deleted file mode 100644 index 5f76351d8c8538ebbadabd40753465294655529e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/pipeline_marigold_depth.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/pipeline_marigold_intrinsics.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/pipeline_marigold_intrinsics.cpython-310.pyc deleted file mode 100644 index 27e3f142bb5c8c7ef4446ac834b3b03e92ed6e20..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/pipeline_marigold_intrinsics.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/pipeline_marigold_normals.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/pipeline_marigold_normals.cpython-310.pyc deleted file mode 100644 index 5334245cec55e0b9e8f10be546d5bfc617492bac..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/__pycache__/pipeline_marigold_normals.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/marigold_image_processing.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/marigold_image_processing.py deleted file mode 100644 index 5130a876606a168d382e9e8a10e44e86376aa490..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/marigold_image_processing.py +++ /dev/null @@ -1,689 +0,0 @@ -# Copyright 2023-2025 Marigold Team, ETH Zürich. All rights reserved. -# Copyright 2024-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. -# -------------------------------------------------------------------------- -# 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 - -import numpy as np -import PIL -import torch -import torch.nn.functional as F -from PIL import Image - -from ... import ConfigMixin -from ...configuration_utils import register_to_config -from ...image_processor import PipelineImageInput -from ...utils import CONFIG_NAME, logging -from ...utils.import_utils import is_matplotlib_available - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class MarigoldImageProcessor(ConfigMixin): - config_name = CONFIG_NAME - - @register_to_config - def __init__( - self, - vae_scale_factor: int = 8, - do_normalize: bool = True, - do_range_check: bool = True, - ): - super().__init__() - - @staticmethod - def expand_tensor_or_array(images: Union[torch.Tensor, np.ndarray]) -> Union[torch.Tensor, np.ndarray]: - """ - Expand a tensor or array to a specified number of images. - """ - if isinstance(images, np.ndarray): - if images.ndim == 2: # [H,W] -> [1,H,W,1] - images = images[None, ..., None] - if images.ndim == 3: # [H,W,C] -> [1,H,W,C] - images = images[None] - elif isinstance(images, torch.Tensor): - if images.ndim == 2: # [H,W] -> [1,1,H,W] - images = images[None, None] - elif images.ndim == 3: # [1,H,W] -> [1,1,H,W] - images = images[None] - else: - raise ValueError(f"Unexpected input type: {type(images)}") - return images - - @staticmethod - def pt_to_numpy(images: torch.Tensor) -> np.ndarray: - """ - Convert a PyTorch tensor to a NumPy image. - """ - images = images.cpu().permute(0, 2, 3, 1).float().numpy() - return images - - @staticmethod - def numpy_to_pt(images: np.ndarray) -> torch.Tensor: - """ - Convert a NumPy image to a PyTorch tensor. - """ - if np.issubdtype(images.dtype, np.integer) and not np.issubdtype(images.dtype, np.unsignedinteger): - raise ValueError(f"Input image dtype={images.dtype} cannot be a signed integer.") - if np.issubdtype(images.dtype, np.complexfloating): - raise ValueError(f"Input image dtype={images.dtype} cannot be complex.") - if np.issubdtype(images.dtype, bool): - raise ValueError(f"Input image dtype={images.dtype} cannot be boolean.") - - images = torch.from_numpy(images.transpose(0, 3, 1, 2)) - return images - - @staticmethod - def resize_antialias( - image: torch.Tensor, size: Tuple[int, int], mode: str, is_aa: Optional[bool] = None - ) -> torch.Tensor: - if not torch.is_tensor(image): - raise ValueError(f"Invalid input type={type(image)}.") - if not torch.is_floating_point(image): - raise ValueError(f"Invalid input dtype={image.dtype}.") - if image.dim() != 4: - raise ValueError(f"Invalid input dimensions; shape={image.shape}.") - - antialias = is_aa and mode in ("bilinear", "bicubic") - image = F.interpolate(image, size, mode=mode, antialias=antialias) - - return image - - @staticmethod - def resize_to_max_edge(image: torch.Tensor, max_edge_sz: int, mode: str) -> torch.Tensor: - if not torch.is_tensor(image): - raise ValueError(f"Invalid input type={type(image)}.") - if not torch.is_floating_point(image): - raise ValueError(f"Invalid input dtype={image.dtype}.") - if image.dim() != 4: - raise ValueError(f"Invalid input dimensions; shape={image.shape}.") - - h, w = image.shape[-2:] - max_orig = max(h, w) - new_h = h * max_edge_sz // max_orig - new_w = w * max_edge_sz // max_orig - - if new_h == 0 or new_w == 0: - raise ValueError(f"Extreme aspect ratio of the input image: [{w} x {h}]") - - image = MarigoldImageProcessor.resize_antialias(image, (new_h, new_w), mode, is_aa=True) - - return image - - @staticmethod - 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): - raise ValueError(f"Invalid input dtype={image.dtype}.") - if image.dim() != 4: - raise ValueError(f"Invalid input dimensions; shape={image.shape}.") - - h, w = image.shape[-2:] - ph, pw = -h % align, -w % align - - image = F.pad(image, (0, pw, 0, ph), mode="replicate") - - return image, (ph, pw) - - @staticmethod - 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): - raise ValueError(f"Invalid input dtype={image.dtype}.") - if image.dim() != 4: - raise ValueError(f"Invalid input dimensions; shape={image.shape}.") - - ph, pw = padding - uh = None if ph == 0 else -ph - uw = None if pw == 0 else -pw - - image = image[:, :, :uh, :uw] - - return image - - @staticmethod - def load_image_canonical( - image: Union[torch.Tensor, np.ndarray, Image.Image], - device: torch.device = torch.device("cpu"), - dtype: torch.dtype = torch.float32, - ) -> Tuple[torch.Tensor, int]: - if isinstance(image, Image.Image): - image = np.array(image) - - image_dtype_max = None - if isinstance(image, (np.ndarray, torch.Tensor)): - image = MarigoldImageProcessor.expand_tensor_or_array(image) - if image.ndim != 4: - raise ValueError("Input image is not 2-, 3-, or 4-dimensional.") - if isinstance(image, np.ndarray): - if np.issubdtype(image.dtype, np.integer) and not np.issubdtype(image.dtype, np.unsignedinteger): - raise ValueError(f"Input image dtype={image.dtype} cannot be a signed integer.") - if np.issubdtype(image.dtype, np.complexfloating): - raise ValueError(f"Input image dtype={image.dtype} cannot be complex.") - if np.issubdtype(image.dtype, bool): - raise ValueError(f"Input image dtype={image.dtype} cannot be boolean.") - if np.issubdtype(image.dtype, np.unsignedinteger): - image_dtype_max = np.iinfo(image.dtype).max - image = image.astype(np.float32) # because torch does not have unsigned dtypes beyond torch.uint8 - image = MarigoldImageProcessor.numpy_to_pt(image) - - if torch.is_tensor(image) and not torch.is_floating_point(image) and image_dtype_max is None: - if image.dtype != torch.uint8: - raise ValueError(f"Image dtype={image.dtype} is not supported.") - image_dtype_max = 255 - - if not torch.is_tensor(image): - raise ValueError(f"Input type unsupported: {type(image)}.") - - if image.shape[1] == 1: - image = image.repeat(1, 3, 1, 1) # [N,1,H,W] -> [N,3,H,W] - if image.shape[1] != 3: - raise ValueError(f"Input image is not 1- or 3-channel: {image.shape}.") - - image = image.to(device=device, dtype=dtype) - - if image_dtype_max is not None: - image = image / image_dtype_max - - return image - - @staticmethod - def check_image_values_range(image: torch.Tensor) -> None: - if not torch.is_tensor(image): - raise ValueError(f"Invalid input type={type(image)}.") - if not torch.is_floating_point(image): - raise ValueError(f"Invalid input dtype={image.dtype}.") - if image.min().item() < 0.0 or image.max().item() > 1.0: - raise ValueError("Input image data is partially outside of the [0,1] range.") - - def preprocess( - self, - image: PipelineImageInput, - processing_resolution: Optional[int] = None, - resample_method_input: str = "bilinear", - device: torch.device = torch.device("cpu"), - dtype: torch.dtype = torch.float32, - ): - if isinstance(image, list): - images = None - for i, img in enumerate(image): - img = self.load_image_canonical(img, device, dtype) # [N,3,H,W] - if images is None: - images = img - else: - if images.shape[2:] != img.shape[2:]: - raise ValueError( - f"Input image[{i}] has incompatible dimensions {img.shape[2:]} with the previous images " - f"{images.shape[2:]}" - ) - images = torch.cat((images, img), dim=0) - image = images - del images - else: - image = self.load_image_canonical(image, device, dtype) # [N,3,H,W] - - original_resolution = image.shape[2:] - - if self.config.do_range_check: - self.check_image_values_range(image) - - if self.config.do_normalize: - image = image * 2.0 - 1.0 - - if processing_resolution is not None and processing_resolution > 0: - image = self.resize_to_max_edge(image, processing_resolution, resample_method_input) # [N,3,PH,PW] - - image, padding = self.pad_image(image, self.config.vae_scale_factor) # [N,3,PPH,PPW] - - return image, padding, original_resolution - - @staticmethod - def colormap( - image: Union[np.ndarray, torch.Tensor], - cmap: str = "Spectral", - bytes: bool = False, - _force_method: Optional[str] = None, - ) -> Union[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", - "binary") without having to install or import matplotlib. For all other cases, the function will attempt to use - the native implementation. - - Args: - image: 2D tensor of values between 0 and 1, either as np.ndarray or torch.Tensor. - cmap: Colormap name. - bytes: Whether to return the output as uint8 or floating point image. - _force_method: - Can be used to specify whether to use the native implementation (`"matplotlib"`), the efficient custom - implementation of the select color maps (`"custom"`), or rely on autodetection (`None`, default). - - Returns: - An RGB-colorized tensor corresponding to the input image. - """ - if not (torch.is_tensor(image) or isinstance(image, np.ndarray)): - raise ValueError("Argument must be a numpy array or torch tensor.") - if _force_method not in (None, "matplotlib", "custom"): - raise ValueError("_force_method must be either `None`, `'matplotlib'` or `'custom'`.") - - supported_cmaps = { - "binary": [ - (1.0, 1.0, 1.0), - (0.0, 0.0, 0.0), - ], - "Spectral": [ # Taken from matplotlib/_cm.py - (0.61960784313725492, 0.003921568627450980, 0.25882352941176473), # 0.0 -> [0] - (0.83529411764705885, 0.24313725490196078, 0.30980392156862746), - (0.95686274509803926, 0.42745098039215684, 0.2627450980392157), - (0.99215686274509807, 0.68235294117647061, 0.38039215686274508), - (0.99607843137254903, 0.8784313725490196, 0.54509803921568623), - (1.0, 1.0, 0.74901960784313726), - (0.90196078431372551, 0.96078431372549022, 0.59607843137254901), - (0.6705882352941176, 0.8666666666666667, 0.64313725490196083), - (0.4, 0.76078431372549016, 0.6470588235294118), - (0.19607843137254902, 0.53333333333333333, 0.74117647058823533), - (0.36862745098039218, 0.30980392156862746, 0.63529411764705879), # 1.0 -> [K-1] - ], - } - - def method_matplotlib(image, cmap, bytes=False): - if is_matplotlib_available(): - import matplotlib - else: - return None - - arg_is_pt, device = torch.is_tensor(image), None - if arg_is_pt: - image, device = image.cpu().numpy(), image.device - - if cmap not in matplotlib.colormaps: - raise ValueError( - f"Unexpected color map {cmap}; available options are: {', '.join(list(matplotlib.colormaps.keys()))}" - ) - - cmap = matplotlib.colormaps[cmap] - out = cmap(image, bytes=bytes) # [?,4] - out = out[..., :3] # [?,3] - - if arg_is_pt: - out = torch.tensor(out, device=device) - - return out - - def method_custom(image, cmap, bytes=False): - arg_is_np = isinstance(image, np.ndarray) - if arg_is_np: - image = torch.tensor(image) - if image.dtype == torch.uint8: - image = image.float() / 255 - else: - image = image.float() - - is_cmap_reversed = cmap.endswith("_r") - if is_cmap_reversed: - cmap = cmap[:-2] - - if cmap not in supported_cmaps: - raise ValueError( - f"Only {list(supported_cmaps.keys())} color maps are available without installing matplotlib." - ) - - cmap = supported_cmaps[cmap] - if is_cmap_reversed: - cmap = cmap[::-1] - cmap = torch.tensor(cmap, dtype=torch.float, device=image.device) # [K,3] - K = cmap.shape[0] - - pos = image.clamp(min=0, max=1) * (K - 1) - left = pos.long() - right = (left + 1).clamp(max=K - 1) - - d = (pos - left.float()).unsqueeze(-1) - left_colors = cmap[left] - right_colors = cmap[right] - - out = (1 - d) * left_colors + d * right_colors - - if bytes: - out = (out * 255).to(torch.uint8) - - if arg_is_np: - out = out.numpy() - - return out - - if _force_method is None and torch.is_tensor(image) and cmap == "Spectral": - return method_custom(image, cmap, bytes) - - out = None - if _force_method != "custom": - out = method_matplotlib(image, cmap, bytes) - - if _force_method == "matplotlib" and out is None: - raise ImportError("Make sure to install matplotlib if you want to use a color map other than 'Spectral'.") - - if out is None: - out = method_custom(image, cmap, bytes) - - return out - - @staticmethod - def visualize_depth( - depth: Union[ - 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]: - """ - 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. - 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. - """ - if val_max <= val_min: - raise ValueError(f"Invalid values range: [{val_min}, {val_max}].") - - def visualize_depth_one(img, idx=None): - prefix = "Depth" + (f"[{idx}]" if idx else "") - if isinstance(img, PIL.Image.Image): - if img.mode != "I;16": - raise ValueError(f"{prefix}: invalid PIL mode={img.mode}.") - img = np.array(img).astype(np.float32) / (2**16 - 1) - if isinstance(img, np.ndarray) or torch.is_tensor(img): - if img.ndim != 2: - raise ValueError(f"{prefix}: unexpected shape={img.shape}.") - if isinstance(img, np.ndarray): - img = torch.from_numpy(img) - if not torch.is_floating_point(img): - raise ValueError(f"{prefix}: unexpected dtype={img.dtype}.") - else: - raise ValueError(f"{prefix}: unexpected type={type(img)}.") - if val_min != 0.0 or val_max != 1.0: - img = (img - val_min) / (val_max - val_min) - img = MarigoldImageProcessor.colormap(img, cmap=color_map, bytes=True) # [H,W,3] - img = PIL.Image.fromarray(img.cpu().numpy()) - return img - - if depth is None or isinstance(depth, list) and any(o is None for o in depth): - raise ValueError("Input depth is `None`") - if isinstance(depth, (np.ndarray, torch.Tensor)): - depth = MarigoldImageProcessor.expand_tensor_or_array(depth) - if isinstance(depth, np.ndarray): - depth = MarigoldImageProcessor.numpy_to_pt(depth) # [N,H,W,1] -> [N,1,H,W] - if not (depth.ndim == 4 and depth.shape[1] == 1): # [N,1,H,W] - raise ValueError(f"Unexpected input shape={depth.shape}, expecting [N,1,H,W].") - return [visualize_depth_one(img[0], idx) for idx, img in enumerate(depth)] - elif isinstance(depth, list): - return [visualize_depth_one(img, idx) for idx, img in enumerate(depth)] - else: - raise ValueError(f"Unexpected input type: {type(depth)}") - - @staticmethod - def export_depth_to_16bit_png( - depth: Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]], - val_min: float = 0.0, - val_max: float = 1.0, - ) -> 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): - raise ValueError(f"{prefix}: unexpected type={type(img)}.") - if img.ndim != 2: - raise ValueError(f"{prefix}: unexpected shape={img.shape}.") - if torch.is_tensor(img): - img = img.cpu().numpy() - if not np.issubdtype(img.dtype, np.floating): - raise ValueError(f"{prefix}: unexpected dtype={img.dtype}.") - if val_min != 0.0 or val_max != 1.0: - img = (img - val_min) / (val_max - val_min) - img = (img * (2**16 - 1)).astype(np.uint16) - img = PIL.Image.fromarray(img, mode="I;16") - return img - - if depth is None or isinstance(depth, list) and any(o is None for o in depth): - raise ValueError("Input depth is `None`") - if isinstance(depth, (np.ndarray, torch.Tensor)): - depth = MarigoldImageProcessor.expand_tensor_or_array(depth) - if isinstance(depth, np.ndarray): - depth = MarigoldImageProcessor.numpy_to_pt(depth) # [N,H,W,1] -> [N,1,H,W] - if not (depth.ndim == 4 and depth.shape[1] == 1): - raise ValueError(f"Unexpected input shape={depth.shape}, expecting [N,1,H,W].") - return [export_depth_to_16bit_png_one(img[0], idx) for idx, img in enumerate(depth)] - elif isinstance(depth, list): - return [export_depth_to_16bit_png_one(img, idx) for idx, img in enumerate(depth)] - else: - raise ValueError(f"Unexpected input type: {type(depth)}") - - @staticmethod - def visualize_normals( - normals: Union[ - 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]: - """ - Visualizes surface normals, such as predictions of the `MarigoldNormalsPipeline`. - - Args: - normals (`Union[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. - flip_y (`bool`, *optional*, defaults to `False`): Flips the Y axis of the normals frame of reference. - Default direction is top. - 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. - """ - flip_vec = None - if any((flip_x, flip_y, flip_z)): - flip_vec = torch.tensor( - [ - (-1) ** flip_x, - (-1) ** flip_y, - (-1) ** flip_z, - ], - dtype=torch.float32, - ) - - def visualize_normals_one(img, idx=None): - img = img.permute(1, 2, 0) - if flip_vec is not None: - img *= flip_vec.to(img.device) - img = (img + 1.0) * 0.5 - img = (img * 255).to(dtype=torch.uint8, device="cpu").numpy() - img = PIL.Image.fromarray(img) - return img - - if normals is None or isinstance(normals, list) and any(o is None for o in normals): - raise ValueError("Input normals is `None`") - if isinstance(normals, (np.ndarray, torch.Tensor)): - normals = MarigoldImageProcessor.expand_tensor_or_array(normals) - if isinstance(normals, np.ndarray): - normals = MarigoldImageProcessor.numpy_to_pt(normals) # [N,3,H,W] - if not (normals.ndim == 4 and normals.shape[1] == 3): - raise ValueError(f"Unexpected input shape={normals.shape}, expecting [N,3,H,W].") - return [visualize_normals_one(img, idx) for idx, img in enumerate(normals)] - elif isinstance(normals, list): - return [visualize_normals_one(img, idx) for idx, img in enumerate(normals)] - else: - raise ValueError(f"Unexpected input type: {type(normals)}") - - @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]]: - """ - Visualizes intrinsic image decomposition, such as predictions of the `MarigoldIntrinsicsPipeline`. - - Args: - prediction (`Union[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"`): - 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. - """ - if "target_names" not in target_properties: - raise ValueError("Missing `target_names` in target_properties") - if not isinstance(color_map, str) and not ( - isinstance(color_map, dict) - and all(isinstance(k, str) and isinstance(v, str) for k, v in color_map.items()) - ): - raise ValueError("`color_map` must be a string or a dictionary of strings") - n_targets = len(target_properties["target_names"]) - - def visualize_targets_one(images, idx=None): - # img: [T, 3, H, W] - out = {} - for target_name, img in zip(target_properties["target_names"], images): - img = img.permute(1, 2, 0) # [H, W, 3] - prediction_space = target_properties[target_name].get("prediction_space", "srgb") - if prediction_space == "stack": - sub_target_names = target_properties[target_name]["sub_target_names"] - if len(sub_target_names) != 3 or any( - not (isinstance(s, str) or s is None) for s in sub_target_names - ): - raise ValueError(f"Unexpected target sub-names {sub_target_names} in {target_name}") - for i, sub_target_name in enumerate(sub_target_names): - if sub_target_name is None: - continue - sub_img = img[:, :, i] - sub_prediction_space = target_properties[sub_target_name].get("prediction_space", "srgb") - if sub_prediction_space == "linear": - sub_up_to_scale = target_properties[sub_target_name].get("up_to_scale", False) - if sub_up_to_scale: - sub_img = sub_img / max(sub_img.max().item(), 1e-6) - sub_img = sub_img ** (1 / 2.2) - cmap_name = ( - color_map if isinstance(color_map, str) else color_map.get(sub_target_name, "binary") - ) - sub_img = MarigoldImageProcessor.colormap(sub_img, cmap=cmap_name, bytes=True) - sub_img = PIL.Image.fromarray(sub_img.cpu().numpy()) - out[sub_target_name] = sub_img - elif prediction_space == "linear": - up_to_scale = target_properties[target_name].get("up_to_scale", False) - if up_to_scale: - img = img / max(img.max().item(), 1e-6) - img = img ** (1 / 2.2) - elif prediction_space == "srgb": - pass - img = (img * 255).to(dtype=torch.uint8, device="cpu").numpy() - img = PIL.Image.fromarray(img) - out[target_name] = img - return out - - if prediction is None or isinstance(prediction, list) and any(o is None for o in prediction): - raise ValueError("Input prediction is `None`") - if isinstance(prediction, (np.ndarray, torch.Tensor)): - prediction = MarigoldImageProcessor.expand_tensor_or_array(prediction) - if isinstance(prediction, np.ndarray): - prediction = MarigoldImageProcessor.numpy_to_pt(prediction) # [N*T,3,H,W] - if not (prediction.ndim == 4 and prediction.shape[1] == 3 and prediction.shape[0] % n_targets == 0): - raise ValueError(f"Unexpected input shape={prediction.shape}, expecting [N*T,3,H,W].") - N_T, _, H, W = prediction.shape - N = N_T // n_targets - prediction = prediction.reshape(N, n_targets, 3, H, W) - return [visualize_targets_one(img, idx) for idx, img in enumerate(prediction)] - elif isinstance(prediction, list): - return [visualize_targets_one(img, idx) for idx, img in enumerate(prediction)] - else: - raise ValueError(f"Unexpected input type: {type(prediction)}") - - @staticmethod - def visualize_uncertainty( - uncertainty: Union[ - np.ndarray, - torch.Tensor, - List[np.ndarray], - List[torch.Tensor], - ], - saturation_percentile=95, - ) -> 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 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. - """ - - def visualize_uncertainty_one(img, idx=None): - prefix = "Uncertainty" + (f"[{idx}]" if idx else "") - if img.min() < 0: - raise ValueError(f"{prefix}: unexpected data range, min={img.min()}.") - img = img.permute(1, 2, 0) # [H,W,C] - img = img.squeeze(2).cpu().numpy() # [H,W] or [H,W,3] - saturation_value = np.percentile(img, saturation_percentile) - img = np.clip(img * 255 / saturation_value, 0, 255) - img = img.astype(np.uint8) - img = PIL.Image.fromarray(img) - return img - - if uncertainty is None or isinstance(uncertainty, list) and any(o is None for o in uncertainty): - raise ValueError("Input uncertainty is `None`") - if isinstance(uncertainty, (np.ndarray, torch.Tensor)): - uncertainty = MarigoldImageProcessor.expand_tensor_or_array(uncertainty) - if isinstance(uncertainty, np.ndarray): - uncertainty = MarigoldImageProcessor.numpy_to_pt(uncertainty) # [N,C,H,W] - if not (uncertainty.ndim == 4 and uncertainty.shape[1] in (1, 3)): - raise ValueError(f"Unexpected input shape={uncertainty.shape}, expecting [N,C,H,W] with C in (1,3).") - return [visualize_uncertainty_one(img, idx) for idx, img in enumerate(uncertainty)] - elif isinstance(uncertainty, list): - return [visualize_uncertainty_one(img, idx) for idx, img in enumerate(uncertainty)] - else: - raise ValueError(f"Unexpected input type: {type(uncertainty)}") diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/pipeline_marigold_depth.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/pipeline_marigold_depth.py deleted file mode 100644 index da991aefbd4ab425361e0eda13cfcd49a654f847..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/pipeline_marigold_depth.py +++ /dev/null @@ -1,828 +0,0 @@ -# Copyright 2023-2025 Marigold Team, ETH Zürich. All rights reserved. -# Copyright 2024-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. -# -------------------------------------------------------------------------- -# More information and citation instructions are available on the -# Marigold project website: https://marigoldcomputervision.github.io -# -------------------------------------------------------------------------- -from dataclasses import dataclass -from functools import partial -from typing import Any, Dict, List, Optional, Tuple, Union - -import numpy as np -import torch -from PIL import Image -from tqdm.auto import tqdm -from transformers import CLIPTextModel, CLIPTokenizer - -from ...image_processor import PipelineImageInput -from ...models import ( - AutoencoderKL, - UNet2DConditionModel, -) -from ...schedulers import ( - DDIMScheduler, - LCMScheduler, -) -from ...utils import ( - BaseOutput, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.import_utils import is_scipy_available -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .marigold_image_processing import MarigoldImageProcessor - - -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 diffusers ->>> import torch - ->>> pipe = diffusers.MarigoldDepthPipeline.from_pretrained( -... "prs-eth/marigold-depth-v1-1", variant="fp16", torch_dtype=torch.float16 -... ).to("cuda") - ->>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg") ->>> depth = pipe(image) - ->>> vis = pipe.image_processor.visualize_depth(depth.prediction) ->>> vis[0].save("einstein_depth.png") - ->>> depth_16bit = pipe.image_processor.export_depth_to_16bit_png(depth.prediction) ->>> depth_16bit[0].save("einstein_depth_16bit.png") -``` -""" - - -@dataclass -class MarigoldDepthOutput(BaseOutput): - """ - Output class for Marigold monocular depth prediction pipeline. - - Args: - prediction (`np.ndarray`, `torch.Tensor`): - Predicted depth maps with values in the range [0, 1]. The shape is $numimages \times 1 \times height \times - width$ for `torch.Tensor` or $numimages \times height \times width \times 1$ for `np.ndarray`. - uncertainty (`None`, `np.ndarray`, `torch.Tensor`): - Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is $numimages - \times 1 \times height \times width$ for `torch.Tensor` or $numimages \times height \times width \times 1$ - for `np.ndarray`. - latent (`None`, `torch.Tensor`): - Latent features corresponding to the predictions, compatible with the `latents` argument of the pipeline. - The shape is $numimages * numensemble \times 4 \times latentheight \times latentwidth$. - """ - - prediction: Union[np.ndarray, torch.Tensor] - uncertainty: Union[None, np.ndarray, torch.Tensor] - latent: Union[None, torch.Tensor] - - -class MarigoldDepthPipeline(DiffusionPipeline): - """ - Pipeline for monocular depth estimation using the Marigold method: https://marigoldmonodepth.github.io. - - 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.) - - Args: - unet (`UNet2DConditionModel`): - Conditional U-Net to denoise the depth latent, conditioned on image latent. - vae (`AutoencoderKL`): - Variational Auto-Encoder (VAE) Model to encode and decode images and predictions to and from latent - representations. - scheduler (`DDIMScheduler` or `LCMScheduler`): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. - text_encoder (`CLIPTextModel`): - Text-encoder, for empty text embedding. - tokenizer (`CLIPTokenizer`): - CLIP tokenizer. - prediction_type (`str`, *optional*): - Type of predictions made by the model. - scale_invariant (`bool`, *optional*): - A model property specifying whether the predicted depth maps are scale-invariant. This value must be set in - the model config. When used together with the `shift_invariant=True` flag, the model is also called - "affine-invariant". NB: overriding this value is not supported. - shift_invariant (`bool`, *optional*): - A model property specifying whether the predicted depth maps are shift-invariant. This value must be set in - the model config. When used together with the `scale_invariant=True` flag, the model is also called - "affine-invariant". NB: overriding this value is not supported. - default_denoising_steps (`int`, *optional*): - The minimum number of denoising diffusion steps that are required to produce a prediction of reasonable - quality with the given model. This value must be set in the model config. When the pipeline is called - without explicitly setting `num_inference_steps`, the default value is used. This is required to ensure - reasonable results with various model flavors compatible with the pipeline, such as those relying on very - short denoising schedules (`LCMScheduler`) and those with full diffusion schedules (`DDIMScheduler`). - default_processing_resolution (`int`, *optional*): - The recommended value of the `processing_resolution` parameter of the pipeline. This value must be set in - the model config. When the pipeline is called without explicitly setting `processing_resolution`, the - default value is used. This is required to ensure reasonable results with various model flavors trained - with varying optimal processing resolution values. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - supported_prediction_types = ("depth", "disparity") - - def __init__( - self, - unet: UNet2DConditionModel, - vae: AutoencoderKL, - scheduler: Union[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, - ): - super().__init__() - - if prediction_type not in self.supported_prediction_types: - logger.warning( - f"Potentially unsupported `prediction_type='{prediction_type}'`; values supported by the pipeline: " - f"{self.supported_prediction_types}." - ) - - self.register_modules( - unet=unet, - vae=vae, - scheduler=scheduler, - text_encoder=text_encoder, - tokenizer=tokenizer, - ) - self.register_to_config( - prediction_type=prediction_type, - scale_invariant=scale_invariant, - shift_invariant=shift_invariant, - default_denoising_steps=default_denoising_steps, - default_processing_resolution=default_processing_resolution, - ) - - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - - self.scale_invariant = scale_invariant - self.shift_invariant = shift_invariant - self.default_denoising_steps = default_denoising_steps - self.default_processing_resolution = default_processing_resolution - - self.empty_text_embedding = None - - self.image_processor = MarigoldImageProcessor(vae_scale_factor=self.vae_scale_factor) - - def check_inputs( - self, - image: PipelineImageInput, - num_inference_steps: int, - ensemble_size: int, - processing_resolution: int, - 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]]], - output_type: str, - output_uncertainty: bool, - ) -> int: - actual_vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) - if actual_vae_scale_factor != self.vae_scale_factor: - raise ValueError( - f"`vae_scale_factor` computed at initialization ({self.vae_scale_factor}) differs from the actual one ({actual_vae_scale_factor})." - ) - if num_inference_steps is None: - raise ValueError("`num_inference_steps` is not specified and could not be resolved from the model config.") - if num_inference_steps < 1: - raise ValueError("`num_inference_steps` must be positive.") - if ensemble_size < 1: - raise ValueError("`ensemble_size` must be positive.") - if ensemble_size == 2: - logger.warning( - "`ensemble_size` == 2 results are similar to no ensembling (1); " - "consider increasing the value to at least 3." - ) - if ensemble_size > 1 and (self.scale_invariant or self.shift_invariant) and not is_scipy_available(): - raise ImportError("Make sure to install scipy if you want to use ensembling.") - if ensemble_size == 1 and output_uncertainty: - raise ValueError( - "Computing uncertainty by setting `output_uncertainty=True` also requires setting `ensemble_size` " - "greater than 1." - ) - if processing_resolution is None: - raise ValueError( - "`processing_resolution` is not specified and could not be resolved from the model config." - ) - if processing_resolution < 0: - raise ValueError( - "`processing_resolution` must be non-negative: 0 for native resolution, or any positive value for " - "downsampled processing." - ) - if processing_resolution % self.vae_scale_factor != 0: - raise ValueError(f"`processing_resolution` must be a multiple of {self.vae_scale_factor}.") - if resample_method_input not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"): - raise ValueError( - "`resample_method_input` takes string values compatible with PIL library: " - "nearest, nearest-exact, bilinear, bicubic, area." - ) - if resample_method_output not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"): - raise ValueError( - "`resample_method_output` takes string values compatible with PIL library: " - "nearest, nearest-exact, bilinear, bicubic, area." - ) - if batch_size < 1: - raise ValueError("`batch_size` must be positive.") - if output_type not in ["pt", "np"]: - raise ValueError("`output_type` must be one of `pt` or `np`.") - if latents is not None and generator is not None: - raise ValueError("`latents` and `generator` cannot be used together.") - if ensembling_kwargs is not None: - if not isinstance(ensembling_kwargs, dict): - raise ValueError("`ensembling_kwargs` must be a dictionary.") - if "reduction" in ensembling_kwargs and ensembling_kwargs["reduction"] not in ("mean", "median"): - raise ValueError("`ensembling_kwargs['reduction']` can be either `'mean'` or `'median'`.") - - # image checks - num_images = 0 - W, H = None, None - if not isinstance(image, list): - image = [image] - for i, img in enumerate(image): - if isinstance(img, np.ndarray) or torch.is_tensor(img): - if img.ndim not in (2, 3, 4): - raise ValueError(f"`image[{i}]` has unsupported dimensions or shape: {img.shape}.") - H_i, W_i = img.shape[-2:] - N_i = 1 - if img.ndim == 4: - N_i = img.shape[0] - elif isinstance(img, Image.Image): - W_i, H_i = img.size - N_i = 1 - else: - raise ValueError(f"Unsupported `image[{i}]` type: {type(img)}.") - if W is None: - W, H = W_i, H_i - elif (W, H) != (W_i, H_i): - raise ValueError( - f"Input `image[{i}]` has incompatible dimensions {(W_i, H_i)} with the previous images {(W, H)}" - ) - num_images += N_i - - # latents checks - if latents is not None: - if not torch.is_tensor(latents): - raise ValueError("`latents` must be a torch.Tensor.") - if latents.dim() != 4: - raise ValueError(f"`latents` has unsupported dimensions or shape: {latents.shape}.") - - if processing_resolution > 0: - max_orig = max(H, W) - new_H = H * processing_resolution // max_orig - new_W = W * processing_resolution // max_orig - if new_H == 0 or new_W == 0: - raise ValueError(f"Extreme aspect ratio of the input image: [{W} x {H}]") - W, H = new_W, new_H - w = (W + self.vae_scale_factor - 1) // self.vae_scale_factor - h = (H + self.vae_scale_factor - 1) // self.vae_scale_factor - shape_expected = (num_images * ensemble_size, self.vae.config.latent_channels, h, w) - - if latents.shape != shape_expected: - raise ValueError(f"`latents` has unexpected shape={latents.shape} expected={shape_expected}.") - - # generator checks - if generator is not None: - if isinstance(generator, list): - if len(generator) != num_images * ensemble_size: - raise ValueError( - "The number of generators must match the total number of ensemble members for all input images." - ) - if not all(g.device.type == generator[0].device.type for g in generator): - raise ValueError("`generator` device placement is not consistent in the list.") - elif not isinstance(generator, torch.Generator): - raise ValueError(f"Unsupported generator type: {type(generator)}.") - - return num_images - - @torch.compiler.disable - def progress_bar(self, iterable=None, total=None, desc=None, leave=True): - if not hasattr(self, "_progress_bar_config"): - self._progress_bar_config = {} - elif not isinstance(self._progress_bar_config, dict): - raise ValueError( - 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) - progress_bar_config["desc"] = progress_bar_config.get("desc", desc) - progress_bar_config["leave"] = progress_bar_config.get("leave", leave) - if iterable is not None: - return tqdm(iterable, **progress_bar_config) - elif total is not None: - return tqdm(total=total, **progress_bar_config) - else: - raise ValueError("Either `total` or `iterable` has to be defined.") - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - image: PipelineImageInput, - num_inference_steps: Optional[int] = None, - ensemble_size: int = 1, - processing_resolution: Optional[int] = 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, - output_type: str = "np", - output_uncertainty: bool = False, - output_latent: bool = False, - return_dict: bool = True, - ): - """ - 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 - 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 - same width and height. - num_inference_steps (`int`, *optional*, defaults to `None`): - Number of denoising diffusion steps during inference. The default value `None` results in automatic - selection. - ensemble_size (`int`, defaults to `1`): - Number of ensemble predictions. Higher values result in measurable improvements and visual degradation. - processing_resolution (`int`, *optional*, defaults to `None`): - Effective processing resolution. When set to `0`, matches the larger input image dimension. This - produces crisper predictions, but may also lead to the overall loss of global context. The default - value `None` resolves to the optimal value from the model config. - match_input_resolution (`bool`, *optional*, defaults to `True`): - When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer - side of the output will equal to `processing_resolution`. - resample_method_input (`str`, *optional*, defaults to `"bilinear"`): - Resampling method used to resize input images to `processing_resolution`. The accepted values are: - `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`. - resample_method_output (`str`, *optional*, defaults to `"bilinear"`): - Resampling method used to resize output predictions to match the input resolution. The accepted values - are `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`. - batch_size (`int`, *optional*, defaults to `1`): - Batch size; only matters when setting `ensemble_size` or passing a tensor of images. - ensembling_kwargs (`dict`, *optional*, defaults to `None`) - Extra dictionary with arguments for precise ensembling control. The following options are available: - - reduction (`str`, *optional*, defaults to `"median"`): Defines the ensembling function applied in - every pixel location, can be either `"median"` or `"mean"`. - - regularizer_strength (`float`, *optional*, defaults to `0.02`): Strength of the regularizer that - pulls the aligned predictions to the unit range from 0 to 1. - - max_iter (`int`, *optional*, defaults to `2`): Maximum number of the alignment solver steps. Refer to - `scipy.optimize.minimize` function, `options` argument. - - tol (`float`, *optional*, defaults to `1e-3`): Alignment solver tolerance. The solver stops when the - 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`): - 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`): - 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 - values are: `"np"` (numpy array) or `"pt"` (torch tensor). - output_uncertainty (`bool`, *optional*, defaults to `False`): - When enabled, the output's `uncertainty` field contains the predictive uncertainty map, provided that - the `ensemble_size` argument is set to a value above 2. - output_latent (`bool`, *optional*, defaults to `False`): - When enabled, the output's `latent` field contains the latent codes corresponding to the predictions - within the ensemble. These codes can be saved, modified, and used for subsequent calls with the - `latents` argument. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.marigold.MarigoldDepthOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`~pipelines.marigold.MarigoldDepthOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.marigold.MarigoldDepthOutput`] is returned, otherwise a - `tuple` is returned where the first element is the prediction, the second element is the uncertainty - (or `None`), and the third is the latent (or `None`). - """ - - # 0. Resolving variables. - device = self._execution_device - dtype = self.dtype - - # Model-specific optimal default values leading to fast and reasonable results. - if num_inference_steps is None: - num_inference_steps = self.default_denoising_steps - if processing_resolution is None: - processing_resolution = self.default_processing_resolution - - # 1. Check inputs. - num_images = self.check_inputs( - image, - num_inference_steps, - ensemble_size, - processing_resolution, - resample_method_input, - resample_method_output, - batch_size, - ensembling_kwargs, - latents, - generator, - output_type, - output_uncertainty, - ) - - # 2. Prepare empty text conditioning. - # Model invocation: self.tokenizer, self.text_encoder. - if self.empty_text_embedding is None: - prompt = "" - text_inputs = self.tokenizer( - prompt, - padding="do_not_pad", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids.to(device) - self.empty_text_embedding = self.text_encoder(text_input_ids)[0] # [1,2,1024] - - # 3. Preprocess input images. This function loads input image or images of compatible dimensions `(H, W)`, - # optionally downsamples them to the `processing_resolution` `(PH, PW)`, where - # `max(PH, PW) == processing_resolution`, and pads the dimensions to `(PPH, PPW)` such that these values are - # divisible by the latent space downscaling factor (typically 8 in Stable Diffusion). The default value `None` - # of `processing_resolution` resolves to the optimal value from the model config. It is a recommended mode of - # operation and leads to the most reasonable results. Using the native image resolution or any other processing - # resolution can lead to loss of either fine details or global context in the output predictions. - image, padding, original_resolution = self.image_processor.preprocess( - image, processing_resolution, resample_method_input, device, dtype - ) # [N,3,PPH,PPW] - - # 4. Encode input image into latent space. At this step, each of the `N` input images is represented with `E` - # ensemble members. Each ensemble member is an independent diffused prediction, just initialized independently. - # Latents of each such predictions across all input images and all ensemble members are represented in the - # `pred_latent` variable. The variable `image_latent` is of the same shape: it contains each input image encoded - # into latent space and replicated `E` times. The latents can be either generated (see `generator` to ensure - # reproducibility), or passed explicitly via the `latents` argument. The latter can be set outside the pipeline - # code. This behavior can be achieved by setting the `output_latent` argument to `True`. The latent space - # dimensions are `(h, w)`. Encoding into latent space happens in batches of size `batch_size`. - # Model invocation: self.vae.encoder. - image_latent, pred_latent = self.prepare_latents( - image, latents, generator, ensemble_size, batch_size - ) # [N*E,4,h,w], [N*E,4,h,w] - - del image - - batch_empty_text_embedding = self.empty_text_embedding.to(device=device, dtype=dtype).repeat( - batch_size, 1, 1 - ) # [B,1024,2] - - # 5. Process the denoising loop. All `N * E` latents are processed sequentially in batches of size `batch_size`. - # The unet model takes concatenated latent spaces of the input image and the predicted modality as an input, and - # outputs noise for the predicted modality's latent space. The number of denoising diffusion steps is defined by - # `num_inference_steps`. It is either set directly, or resolves to the optimal value specific to the loaded - # model. - # Model invocation: self.unet. - pred_latents = [] - - for i in self.progress_bar( - range(0, num_images * ensemble_size, batch_size), leave=True, desc="Marigold predictions..." - ): - batch_image_latent = image_latent[i : i + batch_size] # [B,4,h,w] - batch_pred_latent = pred_latent[i : i + batch_size] # [B,4,h,w] - effective_batch_size = batch_image_latent.shape[0] - text = batch_empty_text_embedding[:effective_batch_size] # [B,2,1024] - - self.scheduler.set_timesteps(num_inference_steps, device=device) - for t in self.progress_bar(self.scheduler.timesteps, leave=False, desc="Diffusion steps..."): - batch_latent = torch.cat([batch_image_latent, batch_pred_latent], dim=1) # [B,8,h,w] - noise = self.unet(batch_latent, t, encoder_hidden_states=text, return_dict=False)[0] # [B,4,h,w] - batch_pred_latent = self.scheduler.step( - noise, t, batch_pred_latent, generator=generator - ).prev_sample # [B,4,h,w] - - if XLA_AVAILABLE: - xm.mark_step() - - pred_latents.append(batch_pred_latent) - - pred_latent = torch.cat(pred_latents, dim=0) # [N*E,4,h,w] - - del ( - pred_latents, - image_latent, - batch_empty_text_embedding, - batch_image_latent, - batch_pred_latent, - text, - batch_latent, - noise, - ) - - # 6. Decode predictions from latent into pixel space. The resulting `N * E` predictions have shape `(PPH, PPW)`, - # which requires slight postprocessing. Decoding into pixel space happens in batches of size `batch_size`. - # Model invocation: self.vae.decoder. - prediction = torch.cat( - [ - self.decode_prediction(pred_latent[i : i + batch_size]) - for i in range(0, pred_latent.shape[0], batch_size) - ], - dim=0, - ) # [N*E,1,PPH,PPW] - - if not output_latent: - pred_latent = None - - # 7. Remove padding. The output shape is (PH, PW). - prediction = self.image_processor.unpad_image(prediction, padding) # [N*E,1,PH,PW] - - # 8. Ensemble and compute uncertainty (when `output_uncertainty` is set). This code treats each of the `N` - # groups of `E` ensemble predictions independently. For each group it computes an ensembled prediction of shape - # `(PH, PW)` and an optional uncertainty map of the same dimensions. After computing this pair of outputs for - # each group independently, it stacks them respectively into batches of `N` almost final predictions and - # uncertainty maps. - uncertainty = None - if ensemble_size > 1: - prediction = prediction.reshape(num_images, ensemble_size, *prediction.shape[1:]) # [N,E,1,PH,PW] - prediction = [ - self.ensemble_depth( - prediction[i], - self.scale_invariant, - self.shift_invariant, - output_uncertainty, - **(ensembling_kwargs or {}), - ) - for i in range(num_images) - ] # [ [[1,1,PH,PW], [1,1,PH,PW]], ... ] - prediction, uncertainty = zip(*prediction) # [[1,1,PH,PW], ... ], [[1,1,PH,PW], ... ] - prediction = torch.cat(prediction, dim=0) # [N,1,PH,PW] - if output_uncertainty: - uncertainty = torch.cat(uncertainty, dim=0) # [N,1,PH,PW] - else: - uncertainty = None - - # 9. If `match_input_resolution` is set, the output prediction and the uncertainty are upsampled to match the - # input resolution `(H, W)`. This step may introduce upsampling artifacts, and therefore can be disabled. - # Depending on the downstream use-case, upsampling can be also chosen based on the tolerated artifacts by - # setting the `resample_method_output` parameter (e.g., to `"nearest"`). - if match_input_resolution: - prediction = self.image_processor.resize_antialias( - prediction, original_resolution, resample_method_output, is_aa=False - ) # [N,1,H,W] - if uncertainty is not None and output_uncertainty: - uncertainty = self.image_processor.resize_antialias( - uncertainty, original_resolution, resample_method_output, is_aa=False - ) # [N,1,H,W] - - # 10. Prepare the final outputs. - if output_type == "np": - prediction = self.image_processor.pt_to_numpy(prediction) # [N,H,W,1] - if uncertainty is not None and output_uncertainty: - uncertainty = self.image_processor.pt_to_numpy(uncertainty) # [N,H,W,1] - - # 11. Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (prediction, uncertainty, pred_latent) - - return MarigoldDepthOutput( - prediction=prediction, - uncertainty=uncertainty, - latent=pred_latent, - ) - - def prepare_latents( - self, - image: torch.Tensor, - latents: Optional[torch.Tensor], - generator: Optional[torch.Generator], - ensemble_size: int, - batch_size: int, - ) -> Tuple[torch.Tensor, torch.Tensor]: - def retrieve_latents(encoder_output): - if hasattr(encoder_output, "latent_dist"): - 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") - - image_latent = torch.cat( - [ - retrieve_latents(self.vae.encode(image[i : i + batch_size])) - for i in range(0, image.shape[0], batch_size) - ], - dim=0, - ) # [N,4,h,w] - image_latent = image_latent * self.vae.config.scaling_factor - image_latent = image_latent.repeat_interleave(ensemble_size, dim=0) # [N*E,4,h,w] - - pred_latent = latents - if pred_latent is None: - pred_latent = randn_tensor( - image_latent.shape, - generator=generator, - device=image_latent.device, - dtype=image_latent.dtype, - ) # [N*E,4,h,w] - - return image_latent, pred_latent - - def decode_prediction(self, pred_latent: torch.Tensor) -> torch.Tensor: - if pred_latent.dim() != 4 or pred_latent.shape[1] != self.vae.config.latent_channels: - raise ValueError( - f"Expecting 4D tensor of shape [B,{self.vae.config.latent_channels},H,W]; got {pred_latent.shape}." - ) - - prediction = self.vae.decode(pred_latent / self.vae.config.scaling_factor, return_dict=False)[0] # [B,3,H,W] - - prediction = prediction.mean(dim=1, keepdim=True) # [B,1,H,W] - prediction = torch.clip(prediction, -1.0, 1.0) # [B,1,H,W] - prediction = (prediction + 1.0) / 2.0 - - return prediction # [B,1,H,W] - - @staticmethod - def ensemble_depth( - depth: torch.Tensor, - scale_invariant: bool = True, - shift_invariant: bool = True, - output_uncertainty: bool = False, - reduction: str = "median", - regularizer_strength: float = 0.02, - max_iter: int = 2, - tol: float = 1e-3, - max_res: int = 1024, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """ - 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 - depth maps, it can also be used with disparity maps as long as the input tensor values are non-negative. The - alignment happens when the predictions have one or more degrees of freedom, that is when they are either - affine-invariant (`scale_invariant=True` and `shift_invariant=True`), or just scale-invariant (only - `scale_invariant=True`). For absolute predictions (`scale_invariant=False` and `shift_invariant=False`) - alignment is skipped and only ensembling is performed. - - Args: - depth (`torch.Tensor`): - Input ensemble depth maps. - scale_invariant (`bool`, *optional*, defaults to `True`): - Whether to treat predictions as scale-invariant. - shift_invariant (`bool`, *optional*, defaults to `True`): - Whether to treat predictions as shift-invariant. - output_uncertainty (`bool`, *optional*, defaults to `False`): - Whether to output uncertainty map. - reduction (`str`, *optional*, defaults to `"median"`): - Reduction method used to ensemble aligned predictions. The accepted values are: `"mean"` and - `"median"`. - regularizer_strength (`float`, *optional*, defaults to `0.02`): - Strength of the regularizer that pulls the aligned predictions to the unit range from 0 to 1. - max_iter (`int`, *optional*, defaults to `2`): - Maximum number of the alignment solver steps. Refer to `scipy.optimize.minimize` function, `options` - argument. - tol (`float`, *optional*, defaults to `1e-3`): - Alignment solver tolerance. The solver stops when the tolerance is reached. - max_res (`int`, *optional*, defaults to `1024`): - Resolution at which the alignment is performed; `None` matches the `processing_resolution`. - Returns: - A tensor of aligned and ensembled depth maps and optionally a tensor of uncertainties of the same shape: - `(1, 1, H, W)`. - """ - if depth.dim() != 4 or depth.shape[1] != 1: - raise ValueError(f"Expecting 4D tensor of shape [B,1,H,W]; got {depth.shape}.") - if reduction not in ("mean", "median"): - raise ValueError(f"Unrecognized reduction method: {reduction}.") - if not scale_invariant and shift_invariant: - raise ValueError("Pure shift-invariant ensembling is not supported.") - - def init_param(depth: torch.Tensor): - init_min = depth.reshape(ensemble_size, -1).min(dim=1).values - init_max = depth.reshape(ensemble_size, -1).max(dim=1).values - - if scale_invariant and shift_invariant: - init_s = 1.0 / (init_max - init_min).clamp(min=1e-6) - init_t = -init_s * init_min - param = torch.cat((init_s, init_t)).cpu().numpy() - elif scale_invariant: - init_s = 1.0 / init_max.clamp(min=1e-6) - param = init_s.cpu().numpy() - else: - raise ValueError("Unrecognized alignment.") - param = param.astype(np.float64) - - return param - - def align(depth: torch.Tensor, param: np.ndarray) -> torch.Tensor: - if scale_invariant and shift_invariant: - s, t = np.split(param, 2) - s = torch.from_numpy(s).to(depth).view(ensemble_size, 1, 1, 1) - t = torch.from_numpy(t).to(depth).view(ensemble_size, 1, 1, 1) - out = depth * s + t - elif scale_invariant: - s = torch.from_numpy(param).to(depth).view(ensemble_size, 1, 1, 1) - out = depth * s - else: - raise ValueError("Unrecognized alignment.") - return out - - def ensemble( - depth_aligned: torch.Tensor, return_uncertainty: bool = False - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - uncertainty = None - if reduction == "mean": - prediction = torch.mean(depth_aligned, dim=0, keepdim=True) - if return_uncertainty: - uncertainty = torch.std(depth_aligned, dim=0, keepdim=True) - elif reduction == "median": - prediction = torch.median(depth_aligned, dim=0, keepdim=True).values - if return_uncertainty: - uncertainty = torch.median(torch.abs(depth_aligned - prediction), dim=0, keepdim=True).values - else: - raise ValueError(f"Unrecognized reduction method: {reduction}.") - return prediction, uncertainty - - def cost_fn(param: np.ndarray, depth: torch.Tensor) -> float: - cost = 0.0 - depth_aligned = align(depth, param) - - for i, j in torch.combinations(torch.arange(ensemble_size)): - diff = depth_aligned[i] - depth_aligned[j] - cost += (diff**2).mean().sqrt().item() - - if regularizer_strength > 0: - prediction, _ = ensemble(depth_aligned, return_uncertainty=False) - err_near = prediction.min().abs().item() - err_far = (1.0 - prediction.max()).abs().item() - cost += (err_near + err_far) * regularizer_strength - - return cost - - def compute_param(depth: torch.Tensor): - import scipy - - depth_to_align = depth.to(torch.float32) - if max_res is not None and max(depth_to_align.shape[2:]) > max_res: - depth_to_align = MarigoldImageProcessor.resize_to_max_edge(depth_to_align, max_res, "nearest-exact") - - param = init_param(depth_to_align) - - res = scipy.optimize.minimize( - partial(cost_fn, depth=depth_to_align), - param, - method="BFGS", - tol=tol, - options={"maxiter": max_iter, "disp": False}, - ) - - return res.x - - requires_aligning = scale_invariant or shift_invariant - ensemble_size = depth.shape[0] - - if requires_aligning: - param = compute_param(depth) - depth = align(depth, param) - - depth, uncertainty = ensemble(depth, return_uncertainty=output_uncertainty) - - depth_max = depth.max() - if scale_invariant and shift_invariant: - depth_min = depth.min() - elif scale_invariant: - depth_min = 0 - else: - raise ValueError("Unrecognized alignment.") - depth_range = (depth_max - depth_min).clamp(min=1e-6) - depth = (depth - depth_min) / depth_range - if output_uncertainty: - uncertainty /= depth_range - - return depth, uncertainty # [1,1,H,W], [1,1,H,W] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py deleted file mode 100644 index c809de18f469f63202f77caa2a1dae997da55ab8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py +++ /dev/null @@ -1,721 +0,0 @@ -# Copyright 2023-2025 Marigold Team, ETH Zürich. All rights reserved. -# Copyright 2024-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. -# -------------------------------------------------------------------------- -# More information and citation instructions are available on the -# Marigold project website: https://marigoldcomputervision.github.io -# -------------------------------------------------------------------------- -from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union - -import numpy as np -import torch -from PIL import Image -from tqdm.auto import tqdm -from transformers import CLIPTextModel, CLIPTokenizer - -from ...image_processor import PipelineImageInput -from ...models import ( - AutoencoderKL, - UNet2DConditionModel, -) -from ...schedulers import ( - DDIMScheduler, - LCMScheduler, -) -from ...utils import ( - BaseOutput, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .marigold_image_processing import MarigoldImageProcessor - - -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 diffusers ->>> import torch - ->>> pipe = diffusers.MarigoldIntrinsicsPipeline.from_pretrained( -... "prs-eth/marigold-iid-appearance-v1-1", variant="fp16", torch_dtype=torch.float16 -... ).to("cuda") - ->>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg") ->>> intrinsics = pipe(image) - ->>> vis = pipe.image_processor.visualize_intrinsics(intrinsics.prediction, pipe.target_properties) ->>> vis[0]["albedo"].save("einstein_albedo.png") ->>> vis[0]["roughness"].save("einstein_roughness.png") ->>> vis[0]["metallicity"].save("einstein_metallicity.png") -``` -```py ->>> import diffusers ->>> import torch - ->>> pipe = diffusers.MarigoldIntrinsicsPipeline.from_pretrained( -... "prs-eth/marigold-iid-lighting-v1-1", variant="fp16", torch_dtype=torch.float16 -... ).to("cuda") - ->>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg") ->>> intrinsics = pipe(image) - ->>> vis = pipe.image_processor.visualize_intrinsics(intrinsics.prediction, pipe.target_properties) ->>> vis[0]["albedo"].save("einstein_albedo.png") ->>> vis[0]["shading"].save("einstein_shading.png") ->>> vis[0]["residual"].save("einstein_residual.png") -``` -""" - - -@dataclass -class MarigoldIntrinsicsOutput(BaseOutput): - """ - Output class for Marigold Intrinsic Image Decomposition pipeline. - - Args: - prediction (`np.ndarray`, `torch.Tensor`): - Predicted image intrinsics with values in the range [0, 1]. The shape is $(numimages * numtargets) \times 3 - \times height \times width$ for `torch.Tensor` or $(numimages * numtargets) \times height \times width - \times 3$ for `np.ndarray`, where `numtargets` corresponds to the number of predicted target modalities of - the intrinsic image decomposition. - uncertainty (`None`, `np.ndarray`, `torch.Tensor`): - Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is $(numimages * - numtargets) \times 3 \times height \times width$ for `torch.Tensor` or $(numimages * numtargets) \times - height \times width \times 3$ for `np.ndarray`. - latent (`None`, `torch.Tensor`): - Latent features corresponding to the predictions, compatible with the `latents` argument of the pipeline. - The shape is $(numimages * numensemble) \times (numtargets * 4) \times latentheight \times latentwidth$. - """ - - prediction: Union[np.ndarray, torch.Tensor] - uncertainty: Union[None, np.ndarray, torch.Tensor] - latent: Union[None, torch.Tensor] - - -class MarigoldIntrinsicsPipeline(DiffusionPipeline): - """ - Pipeline for Intrinsic Image Decomposition (IID) using the Marigold method: - https://marigoldcomputervision.github.io. - - 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.) - - Args: - unet (`UNet2DConditionModel`): - Conditional U-Net to denoise the targets latent, conditioned on image latent. - vae (`AutoencoderKL`): - Variational Auto-Encoder (VAE) Model to encode and decode images and predictions to and from latent - representations. - scheduler (`DDIMScheduler` or `LCMScheduler`): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. - text_encoder (`CLIPTextModel`): - Text-encoder, for empty text embedding. - tokenizer (`CLIPTokenizer`): - 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, - order and names of the predicted modalities, and any other metadata that may be required to interpret the - predictions. - default_denoising_steps (`int`, *optional*): - The minimum number of denoising diffusion steps that are required to produce a prediction of reasonable - quality with the given model. This value must be set in the model config. When the pipeline is called - without explicitly setting `num_inference_steps`, the default value is used. This is required to ensure - reasonable results with various model flavors compatible with the pipeline, such as those relying on very - short denoising schedules (`LCMScheduler`) and those with full diffusion schedules (`DDIMScheduler`). - default_processing_resolution (`int`, *optional*): - The recommended value of the `processing_resolution` parameter of the pipeline. This value must be set in - the model config. When the pipeline is called without explicitly setting `processing_resolution`, the - default value is used. This is required to ensure reasonable results with various model flavors trained - with varying optimal processing resolution values. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - supported_prediction_types = ("intrinsics",) - - def __init__( - self, - unet: UNet2DConditionModel, - vae: AutoencoderKL, - scheduler: Union[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, - ): - super().__init__() - - if prediction_type not in self.supported_prediction_types: - logger.warning( - f"Potentially unsupported `prediction_type='{prediction_type}'`; values supported by the pipeline: " - f"{self.supported_prediction_types}." - ) - - self.register_modules( - unet=unet, - vae=vae, - scheduler=scheduler, - text_encoder=text_encoder, - tokenizer=tokenizer, - ) - self.register_to_config( - prediction_type=prediction_type, - target_properties=target_properties, - default_denoising_steps=default_denoising_steps, - default_processing_resolution=default_processing_resolution, - ) - - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - - self.target_properties = target_properties - self.default_denoising_steps = default_denoising_steps - self.default_processing_resolution = default_processing_resolution - - self.empty_text_embedding = None - - self.image_processor = MarigoldImageProcessor(vae_scale_factor=self.vae_scale_factor) - - @property - def n_targets(self): - return self.unet.config.out_channels // self.vae.config.latent_channels - - def check_inputs( - self, - image: PipelineImageInput, - num_inference_steps: int, - ensemble_size: int, - processing_resolution: int, - 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]]], - output_type: str, - output_uncertainty: bool, - ) -> int: - actual_vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) - if actual_vae_scale_factor != self.vae_scale_factor: - raise ValueError( - f"`vae_scale_factor` computed at initialization ({self.vae_scale_factor}) differs from the actual one ({actual_vae_scale_factor})." - ) - if num_inference_steps is None: - raise ValueError("`num_inference_steps` is not specified and could not be resolved from the model config.") - if num_inference_steps < 1: - raise ValueError("`num_inference_steps` must be positive.") - if ensemble_size < 1: - raise ValueError("`ensemble_size` must be positive.") - if ensemble_size == 2: - logger.warning( - "`ensemble_size` == 2 results are similar to no ensembling (1); " - "consider increasing the value to at least 3." - ) - if ensemble_size == 1 and output_uncertainty: - raise ValueError( - "Computing uncertainty by setting `output_uncertainty=True` also requires setting `ensemble_size` " - "greater than 1." - ) - if processing_resolution is None: - raise ValueError( - "`processing_resolution` is not specified and could not be resolved from the model config." - ) - if processing_resolution < 0: - raise ValueError( - "`processing_resolution` must be non-negative: 0 for native resolution, or any positive value for " - "downsampled processing." - ) - if processing_resolution % self.vae_scale_factor != 0: - raise ValueError(f"`processing_resolution` must be a multiple of {self.vae_scale_factor}.") - if resample_method_input not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"): - raise ValueError( - "`resample_method_input` takes string values compatible with PIL library: " - "nearest, nearest-exact, bilinear, bicubic, area." - ) - if resample_method_output not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"): - raise ValueError( - "`resample_method_output` takes string values compatible with PIL library: " - "nearest, nearest-exact, bilinear, bicubic, area." - ) - if batch_size < 1: - raise ValueError("`batch_size` must be positive.") - if output_type not in ["pt", "np"]: - raise ValueError("`output_type` must be one of `pt` or `np`.") - if latents is not None and generator is not None: - raise ValueError("`latents` and `generator` cannot be used together.") - if ensembling_kwargs is not None: - if not isinstance(ensembling_kwargs, dict): - raise ValueError("`ensembling_kwargs` must be a dictionary.") - if "reduction" in ensembling_kwargs and ensembling_kwargs["reduction"] not in ("median", "mean"): - raise ValueError("`ensembling_kwargs['reduction']` can be either `'median'` or `'mean'`.") - - # image checks - num_images = 0 - W, H = None, None - if not isinstance(image, list): - image = [image] - for i, img in enumerate(image): - if isinstance(img, np.ndarray) or torch.is_tensor(img): - if img.ndim not in (2, 3, 4): - raise ValueError(f"`image[{i}]` has unsupported dimensions or shape: {img.shape}.") - H_i, W_i = img.shape[-2:] - N_i = 1 - if img.ndim == 4: - N_i = img.shape[0] - elif isinstance(img, Image.Image): - W_i, H_i = img.size - N_i = 1 - else: - raise ValueError(f"Unsupported `image[{i}]` type: {type(img)}.") - if W is None: - W, H = W_i, H_i - elif (W, H) != (W_i, H_i): - raise ValueError( - f"Input `image[{i}]` has incompatible dimensions {(W_i, H_i)} with the previous images {(W, H)}" - ) - num_images += N_i - - # latents checks - if latents is not None: - if not torch.is_tensor(latents): - raise ValueError("`latents` must be a torch.Tensor.") - if latents.dim() != 4: - raise ValueError(f"`latents` has unsupported dimensions or shape: {latents.shape}.") - - if processing_resolution > 0: - max_orig = max(H, W) - new_H = H * processing_resolution // max_orig - new_W = W * processing_resolution // max_orig - if new_H == 0 or new_W == 0: - raise ValueError(f"Extreme aspect ratio of the input image: [{W} x {H}]") - W, H = new_W, new_H - w = (W + self.vae_scale_factor - 1) // self.vae_scale_factor - h = (H + self.vae_scale_factor - 1) // self.vae_scale_factor - shape_expected = (num_images * ensemble_size, self.unet.config.out_channels, h, w) - - if latents.shape != shape_expected: - raise ValueError(f"`latents` has unexpected shape={latents.shape} expected={shape_expected}.") - - # generator checks - if generator is not None: - if isinstance(generator, list): - if len(generator) != num_images * ensemble_size: - raise ValueError( - "The number of generators must match the total number of ensemble members for all input images." - ) - if not all(g.device.type == generator[0].device.type for g in generator): - raise ValueError("`generator` device placement is not consistent in the list.") - elif not isinstance(generator, torch.Generator): - raise ValueError(f"Unsupported generator type: {type(generator)}.") - - return num_images - - @torch.compiler.disable - def progress_bar(self, iterable=None, total=None, desc=None, leave=True): - if not hasattr(self, "_progress_bar_config"): - self._progress_bar_config = {} - elif not isinstance(self._progress_bar_config, dict): - raise ValueError( - 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) - progress_bar_config["desc"] = progress_bar_config.get("desc", desc) - progress_bar_config["leave"] = progress_bar_config.get("leave", leave) - if iterable is not None: - return tqdm(iterable, **progress_bar_config) - elif total is not None: - return tqdm(total=total, **progress_bar_config) - else: - raise ValueError("Either `total` or `iterable` has to be defined.") - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - image: PipelineImageInput, - num_inference_steps: Optional[int] = None, - ensemble_size: int = 1, - processing_resolution: Optional[int] = 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, - output_type: str = "np", - output_uncertainty: bool = False, - output_latent: bool = False, - return_dict: bool = True, - ): - """ - 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. - 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 - same width and height. - num_inference_steps (`int`, *optional*, defaults to `None`): - Number of denoising diffusion steps during inference. The default value `None` results in automatic - selection. - ensemble_size (`int`, defaults to `1`): - Number of ensemble predictions. Higher values result in measurable improvements and visual degradation. - processing_resolution (`int`, *optional*, defaults to `None`): - Effective processing resolution. When set to `0`, matches the larger input image dimension. This - produces crisper predictions, but may also lead to the overall loss of global context. The default - value `None` resolves to the optimal value from the model config. - match_input_resolution (`bool`, *optional*, defaults to `True`): - When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer - side of the output will equal to `processing_resolution`. - resample_method_input (`str`, *optional*, defaults to `"bilinear"`): - Resampling method used to resize input images to `processing_resolution`. The accepted values are: - `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`. - resample_method_output (`str`, *optional*, defaults to `"bilinear"`): - Resampling method used to resize output predictions to match the input resolution. The accepted values - are `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`. - batch_size (`int`, *optional*, defaults to `1`): - Batch size; only matters when setting `ensemble_size` or passing a tensor of images. - ensembling_kwargs (`dict`, *optional*, defaults to `None`) - Extra dictionary with arguments for precise ensembling control. The following options are available: - - reduction (`str`, *optional*, defaults to `"median"`): Defines the ensembling function applied in - every pixel location, can be either `"median"` or `"mean"`. - 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`): - 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 - values are: `"np"` (numpy array) or `"pt"` (torch tensor). - output_uncertainty (`bool`, *optional*, defaults to `False`): - When enabled, the output's `uncertainty` field contains the predictive uncertainty map, provided that - the `ensemble_size` argument is set to a value above 2. - output_latent (`bool`, *optional*, defaults to `False`): - When enabled, the output's `latent` field contains the latent codes corresponding to the predictions - within the ensemble. These codes can be saved, modified, and used for subsequent calls with the - `latents` argument. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.marigold.MarigoldIntrinsicsOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`~pipelines.marigold.MarigoldIntrinsicsOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.marigold.MarigoldIntrinsicsOutput`] is returned, otherwise a - `tuple` is returned where the first element is the prediction, the second element is the uncertainty - (or `None`), and the third is the latent (or `None`). - """ - - # 0. Resolving variables. - device = self._execution_device - dtype = self.dtype - - # Model-specific optimal default values leading to fast and reasonable results. - if num_inference_steps is None: - num_inference_steps = self.default_denoising_steps - if processing_resolution is None: - processing_resolution = self.default_processing_resolution - - # 1. Check inputs. - num_images = self.check_inputs( - image, - num_inference_steps, - ensemble_size, - processing_resolution, - resample_method_input, - resample_method_output, - batch_size, - ensembling_kwargs, - latents, - generator, - output_type, - output_uncertainty, - ) - - # 2. Prepare empty text conditioning. - # Model invocation: self.tokenizer, self.text_encoder. - if self.empty_text_embedding is None: - prompt = "" - text_inputs = self.tokenizer( - prompt, - padding="do_not_pad", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids.to(device) - self.empty_text_embedding = self.text_encoder(text_input_ids)[0] # [1,2,1024] - - # 3. Preprocess input images. This function loads input image or images of compatible dimensions `(H, W)`, - # optionally downsamples them to the `processing_resolution` `(PH, PW)`, where - # `max(PH, PW) == processing_resolution`, and pads the dimensions to `(PPH, PPW)` such that these values are - # divisible by the latent space downscaling factor (typically 8 in Stable Diffusion). The default value `None` - # of `processing_resolution` resolves to the optimal value from the model config. It is a recommended mode of - # operation and leads to the most reasonable results. Using the native image resolution or any other processing - # resolution can lead to loss of either fine details or global context in the output predictions. - image, padding, original_resolution = self.image_processor.preprocess( - image, processing_resolution, resample_method_input, device, dtype - ) # [N,3,PPH,PPW] - - # 4. Encode input image into latent space. At this step, each of the `N` input images is represented with `E` - # ensemble members. Each ensemble member is an independent diffused prediction, just initialized independently. - # Latents of each such predictions across all input images and all ensemble members are represented in the - # `pred_latent` variable. The variable `image_latent` contains each input image encoded into latent space and - # replicated `E` times. The variable `pred_latent` contains latents initialization, where the latent space is - # replicated `T` times relative to the single latent space of `image_latent`, where `T` is the number of the - # predicted targets. The latents can be either generated (see `generator` to ensure reproducibility), or passed - # explicitly via the `latents` argument. The latter can be set outside the pipeline code. This behavior can be - # achieved by setting the `output_latent` argument to `True`. The latent space dimensions are `(h, w)`. Encoding - # into latent space happens in batches of size `batch_size`. - # Model invocation: self.vae.encoder. - image_latent, pred_latent = self.prepare_latents( - image, latents, generator, ensemble_size, batch_size - ) # [N*E,4,h,w], [N*E,T*4,h,w] - - del image - - batch_empty_text_embedding = self.empty_text_embedding.to(device=device, dtype=dtype).repeat( - batch_size, 1, 1 - ) # [B,1024,2] - - # 5. Process the denoising loop. All `N * E` latents are processed sequentially in batches of size `batch_size`. - # The unet model takes concatenated latent spaces of the input image and the predicted modality as an input, and - # outputs noise for the predicted modality's latent space. The number of denoising diffusion steps is defined by - # `num_inference_steps`. It is either set directly, or resolves to the optimal value specific to the loaded - # model. - # Model invocation: self.unet. - pred_latents = [] - - for i in self.progress_bar( - range(0, num_images * ensemble_size, batch_size), leave=True, desc="Marigold predictions..." - ): - batch_image_latent = image_latent[i : i + batch_size] # [B,4,h,w] - batch_pred_latent = pred_latent[i : i + batch_size] # [B,T*4,h,w] - effective_batch_size = batch_image_latent.shape[0] - text = batch_empty_text_embedding[:effective_batch_size] # [B,2,1024] - - self.scheduler.set_timesteps(num_inference_steps, device=device) - for t in self.progress_bar(self.scheduler.timesteps, leave=False, desc="Diffusion steps..."): - batch_latent = torch.cat([batch_image_latent, batch_pred_latent], dim=1) # [B,(1+T)*4,h,w] - noise = self.unet(batch_latent, t, encoder_hidden_states=text, return_dict=False)[0] # [B,T*4,h,w] - batch_pred_latent = self.scheduler.step( - noise, t, batch_pred_latent, generator=generator - ).prev_sample # [B,T*4,h,w] - - if XLA_AVAILABLE: - xm.mark_step() - - pred_latents.append(batch_pred_latent) - - pred_latent = torch.cat(pred_latents, dim=0) # [N*E,T*4,h,w] - - del ( - pred_latents, - image_latent, - batch_empty_text_embedding, - batch_image_latent, - batch_pred_latent, - text, - batch_latent, - noise, - ) - - # 6. Decode predictions from latent into pixel space. The resulting `N * E` predictions have shape `(PPH, PPW)`, - # which requires slight postprocessing. Decoding into pixel space happens in batches of size `batch_size`. - # Model invocation: self.vae.decoder. - pred_latent_for_decoding = pred_latent.reshape( - num_images * ensemble_size * self.n_targets, self.vae.config.latent_channels, *pred_latent.shape[2:] - ) # [N*E*T,4,PPH,PPW] - prediction = torch.cat( - [ - self.decode_prediction(pred_latent_for_decoding[i : i + batch_size]) - for i in range(0, pred_latent_for_decoding.shape[0], batch_size) - ], - dim=0, - ) # [N*E*T,3,PPH,PPW] - - del pred_latent_for_decoding - if not output_latent: - pred_latent = None - - # 7. Remove padding. The output shape is (PH, PW). - prediction = self.image_processor.unpad_image(prediction, padding) # [N*E*T,3,PH,PW] - - # 8. Ensemble and compute uncertainty (when `output_uncertainty` is set). This code treats each of the `N*T` - # groups of `E` ensemble predictions independently. For each group it computes an ensembled prediction of shape - # `(PH, PW)` and an optional uncertainty map of the same dimensions. After computing this pair of outputs for - # each group independently, it stacks them respectively into batches of `N*T` almost final predictions and - # uncertainty maps. - uncertainty = None - if ensemble_size > 1: - prediction = prediction.reshape( - num_images, ensemble_size, self.n_targets, *prediction.shape[1:] - ) # [N,E,T,3,PH,PW] - prediction = [ - self.ensemble_intrinsics(prediction[i], output_uncertainty, **(ensembling_kwargs or {})) - for i in range(num_images) - ] # [ [[T,3,PH,PW], [T,3,PH,PW]], ... ] - prediction, uncertainty = zip(*prediction) # [[T,3,PH,PW], ... ], [[T,3,PH,PW], ... ] - prediction = torch.cat(prediction, dim=0) # [N*T,3,PH,PW] - if output_uncertainty: - uncertainty = torch.cat(uncertainty, dim=0) # [N*T,3,PH,PW] - else: - uncertainty = None - - # 9. If `match_input_resolution` is set, the output prediction and the uncertainty are upsampled to match the - # input resolution `(H, W)`. This step may introduce upsampling artifacts, and therefore can be disabled. - # Depending on the downstream use-case, upsampling can be also chosen based on the tolerated artifacts by - # setting the `resample_method_output` parameter (e.g., to `"nearest"`). - if match_input_resolution: - prediction = self.image_processor.resize_antialias( - prediction, original_resolution, resample_method_output, is_aa=False - ) # [N*T,3,H,W] - if uncertainty is not None and output_uncertainty: - uncertainty = self.image_processor.resize_antialias( - uncertainty, original_resolution, resample_method_output, is_aa=False - ) # [N*T,1,H,W] - - # 10. Prepare the final outputs. - if output_type == "np": - prediction = self.image_processor.pt_to_numpy(prediction) # [N*T,H,W,3] - if uncertainty is not None and output_uncertainty: - uncertainty = self.image_processor.pt_to_numpy(uncertainty) # [N*T,H,W,3] - - # 11. Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (prediction, uncertainty, pred_latent) - - return MarigoldIntrinsicsOutput( - prediction=prediction, - uncertainty=uncertainty, - latent=pred_latent, - ) - - def prepare_latents( - self, - image: torch.Tensor, - latents: Optional[torch.Tensor], - generator: Optional[torch.Generator], - ensemble_size: int, - batch_size: int, - ) -> Tuple[torch.Tensor, torch.Tensor]: - def retrieve_latents(encoder_output): - if hasattr(encoder_output, "latent_dist"): - 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") - - image_latent = torch.cat( - [ - retrieve_latents(self.vae.encode(image[i : i + batch_size])) - for i in range(0, image.shape[0], batch_size) - ], - dim=0, - ) # [N,4,h,w] - image_latent = image_latent * self.vae.config.scaling_factor - image_latent = image_latent.repeat_interleave(ensemble_size, dim=0) # [N*E,4,h,w] - N_E, C, H, W = image_latent.shape - - pred_latent = latents - if pred_latent is None: - pred_latent = randn_tensor( - (N_E, self.n_targets * C, H, W), - generator=generator, - device=image_latent.device, - dtype=image_latent.dtype, - ) # [N*E,T*4,h,w] - - return image_latent, pred_latent - - def decode_prediction(self, pred_latent: torch.Tensor) -> torch.Tensor: - if pred_latent.dim() != 4 or pred_latent.shape[1] != self.vae.config.latent_channels: - raise ValueError( - f"Expecting 4D tensor of shape [B,{self.vae.config.latent_channels},H,W]; got {pred_latent.shape}." - ) - - prediction = self.vae.decode(pred_latent / self.vae.config.scaling_factor, return_dict=False)[0] # [B,3,H,W] - - prediction = torch.clip(prediction, -1.0, 1.0) # [B,3,H,W] - prediction = (prediction + 1.0) / 2.0 - - return prediction # [B,3,H,W] - - @staticmethod - def ensemble_intrinsics( - targets: torch.Tensor, - output_uncertainty: bool = False, - reduction: str = "median", - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """ - 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 - predicted targets. - - Args: - targets (`torch.Tensor`): - Input ensemble of intrinsic image decomposition maps. - output_uncertainty (`bool`, *optional*, defaults to `False`): - Whether to output uncertainty map. - reduction (`str`, *optional*, defaults to `"mean"`): - Reduction method used to ensemble aligned predictions. The accepted values are: `"median"` and - `"mean"`. - - Returns: - A tensor of aligned and ensembled intrinsic decomposition maps with shape `(T, 3, H, W)` and optionally a - tensor of uncertainties of shape `(T, 3, H, W)`. - """ - if targets.dim() != 5 or targets.shape[2] != 3: - raise ValueError(f"Expecting 4D tensor of shape [B,T,3,H,W]; got {targets.shape}.") - if reduction not in ("median", "mean"): - raise ValueError(f"Unrecognized reduction method: {reduction}.") - - B, T, _, H, W = targets.shape - uncertainty = None - if reduction == "mean": - prediction = torch.mean(targets, dim=0) # [T,3,H,W] - if output_uncertainty: - uncertainty = torch.std(targets, dim=0) # [T,3,H,W] - elif reduction == "median": - prediction = torch.median(targets, dim=0, keepdim=True).values # [1,T,3,H,W] - if output_uncertainty: - uncertainty = torch.abs(targets - prediction) # [B,T,3,H,W] - uncertainty = torch.median(uncertainty, dim=0).values # [T,3,H,W] - prediction = prediction.squeeze(0) # [T,3,H,W] - else: - raise ValueError(f"Unrecognized reduction method: {reduction}.") - return prediction, uncertainty diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/pipeline_marigold_normals.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/pipeline_marigold_normals.py deleted file mode 100644 index 192ed590a489ab382db49c26e9b939841cfc08c0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/marigold/pipeline_marigold_normals.py +++ /dev/null @@ -1,705 +0,0 @@ -# Copyright 2023-2025 Marigold Team, ETH Zürich. All rights reserved. -# Copyright 2024-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. -# -------------------------------------------------------------------------- -# More information and citation instructions are available on the -# Marigold project website: https://marigoldcomputervision.github.io -# -------------------------------------------------------------------------- -from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union - -import numpy as np -import torch -from PIL import Image -from tqdm.auto import tqdm -from transformers import CLIPTextModel, CLIPTokenizer - -from ...image_processor import PipelineImageInput -from ...models import ( - AutoencoderKL, - UNet2DConditionModel, -) -from ...schedulers import ( - DDIMScheduler, - LCMScheduler, -) -from ...utils import ( - BaseOutput, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .marigold_image_processing import MarigoldImageProcessor - - -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 diffusers ->>> import torch - ->>> pipe = diffusers.MarigoldNormalsPipeline.from_pretrained( -... "prs-eth/marigold-normals-v1-1", variant="fp16", torch_dtype=torch.float16 -... ).to("cuda") - ->>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg") ->>> normals = pipe(image) - ->>> vis = pipe.image_processor.visualize_normals(normals.prediction) ->>> vis[0].save("einstein_normals.png") -``` -""" - - -@dataclass -class MarigoldNormalsOutput(BaseOutput): - """ - Output class for Marigold monocular normals prediction pipeline. - - Args: - prediction (`np.ndarray`, `torch.Tensor`): - Predicted normals with values in the range [-1, 1]. The shape is $numimages \times 3 \times height \times - width$ for `torch.Tensor` or $numimages \times height \times width \times 3$ for `np.ndarray`. - uncertainty (`None`, `np.ndarray`, `torch.Tensor`): - Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is $numimages - \times 1 \times height \times width$ for `torch.Tensor` or $numimages \times height \times width \times 1$ - for `np.ndarray`. - latent (`None`, `torch.Tensor`): - Latent features corresponding to the predictions, compatible with the `latents` argument of the pipeline. - The shape is $numimages * numensemble \times 4 \times latentheight \times latentwidth$. - """ - - prediction: Union[np.ndarray, torch.Tensor] - uncertainty: Union[None, np.ndarray, torch.Tensor] - latent: Union[None, torch.Tensor] - - -class MarigoldNormalsPipeline(DiffusionPipeline): - """ - Pipeline for monocular normals estimation using the Marigold method: https://marigoldmonodepth.github.io. - - 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.) - - Args: - unet (`UNet2DConditionModel`): - Conditional U-Net to denoise the normals latent, conditioned on image latent. - vae (`AutoencoderKL`): - Variational Auto-Encoder (VAE) Model to encode and decode images and predictions to and from latent - representations. - scheduler (`DDIMScheduler` or `LCMScheduler`): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. - text_encoder (`CLIPTextModel`): - Text-encoder, for empty text embedding. - tokenizer (`CLIPTokenizer`): - CLIP tokenizer. - prediction_type (`str`, *optional*): - Type of predictions made by the model. - use_full_z_range (`bool`, *optional*): - Whether the normals predicted by this model utilize the full range of the Z dimension, or only its positive - half. - default_denoising_steps (`int`, *optional*): - The minimum number of denoising diffusion steps that are required to produce a prediction of reasonable - quality with the given model. This value must be set in the model config. When the pipeline is called - without explicitly setting `num_inference_steps`, the default value is used. This is required to ensure - reasonable results with various model flavors compatible with the pipeline, such as those relying on very - short denoising schedules (`LCMScheduler`) and those with full diffusion schedules (`DDIMScheduler`). - default_processing_resolution (`int`, *optional*): - The recommended value of the `processing_resolution` parameter of the pipeline. This value must be set in - the model config. When the pipeline is called without explicitly setting `processing_resolution`, the - default value is used. This is required to ensure reasonable results with various model flavors trained - with varying optimal processing resolution values. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - supported_prediction_types = ("normals",) - - def __init__( - self, - unet: UNet2DConditionModel, - vae: AutoencoderKL, - scheduler: Union[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, - ): - super().__init__() - - if prediction_type not in self.supported_prediction_types: - logger.warning( - f"Potentially unsupported `prediction_type='{prediction_type}'`; values supported by the pipeline: " - f"{self.supported_prediction_types}." - ) - - self.register_modules( - unet=unet, - vae=vae, - scheduler=scheduler, - text_encoder=text_encoder, - tokenizer=tokenizer, - ) - self.register_to_config( - prediction_type=prediction_type, - use_full_z_range=use_full_z_range, - default_denoising_steps=default_denoising_steps, - default_processing_resolution=default_processing_resolution, - ) - - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - - self.use_full_z_range = use_full_z_range - self.default_denoising_steps = default_denoising_steps - self.default_processing_resolution = default_processing_resolution - - self.empty_text_embedding = None - - self.image_processor = MarigoldImageProcessor(vae_scale_factor=self.vae_scale_factor) - - def check_inputs( - self, - image: PipelineImageInput, - num_inference_steps: int, - ensemble_size: int, - processing_resolution: int, - 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]]], - output_type: str, - output_uncertainty: bool, - ) -> int: - actual_vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) - if actual_vae_scale_factor != self.vae_scale_factor: - raise ValueError( - f"`vae_scale_factor` computed at initialization ({self.vae_scale_factor}) differs from the actual one ({actual_vae_scale_factor})." - ) - if num_inference_steps is None: - raise ValueError("`num_inference_steps` is not specified and could not be resolved from the model config.") - if num_inference_steps < 1: - raise ValueError("`num_inference_steps` must be positive.") - if ensemble_size < 1: - raise ValueError("`ensemble_size` must be positive.") - if ensemble_size == 2: - logger.warning( - "`ensemble_size` == 2 results are similar to no ensembling (1); " - "consider increasing the value to at least 3." - ) - if ensemble_size == 1 and output_uncertainty: - raise ValueError( - "Computing uncertainty by setting `output_uncertainty=True` also requires setting `ensemble_size` " - "greater than 1." - ) - if processing_resolution is None: - raise ValueError( - "`processing_resolution` is not specified and could not be resolved from the model config." - ) - if processing_resolution < 0: - raise ValueError( - "`processing_resolution` must be non-negative: 0 for native resolution, or any positive value for " - "downsampled processing." - ) - if processing_resolution % self.vae_scale_factor != 0: - raise ValueError(f"`processing_resolution` must be a multiple of {self.vae_scale_factor}.") - if resample_method_input not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"): - raise ValueError( - "`resample_method_input` takes string values compatible with PIL library: " - "nearest, nearest-exact, bilinear, bicubic, area." - ) - if resample_method_output not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"): - raise ValueError( - "`resample_method_output` takes string values compatible with PIL library: " - "nearest, nearest-exact, bilinear, bicubic, area." - ) - if batch_size < 1: - raise ValueError("`batch_size` must be positive.") - if output_type not in ["pt", "np"]: - raise ValueError("`output_type` must be one of `pt` or `np`.") - if latents is not None and generator is not None: - raise ValueError("`latents` and `generator` cannot be used together.") - if ensembling_kwargs is not None: - if not isinstance(ensembling_kwargs, dict): - raise ValueError("`ensembling_kwargs` must be a dictionary.") - if "reduction" in ensembling_kwargs and ensembling_kwargs["reduction"] not in ("closest", "mean"): - raise ValueError("`ensembling_kwargs['reduction']` can be either `'closest'` or `'mean'`.") - - # image checks - num_images = 0 - W, H = None, None - if not isinstance(image, list): - image = [image] - for i, img in enumerate(image): - if isinstance(img, np.ndarray) or torch.is_tensor(img): - if img.ndim not in (2, 3, 4): - raise ValueError(f"`image[{i}]` has unsupported dimensions or shape: {img.shape}.") - H_i, W_i = img.shape[-2:] - N_i = 1 - if img.ndim == 4: - N_i = img.shape[0] - elif isinstance(img, Image.Image): - W_i, H_i = img.size - N_i = 1 - else: - raise ValueError(f"Unsupported `image[{i}]` type: {type(img)}.") - if W is None: - W, H = W_i, H_i - elif (W, H) != (W_i, H_i): - raise ValueError( - f"Input `image[{i}]` has incompatible dimensions {(W_i, H_i)} with the previous images {(W, H)}" - ) - num_images += N_i - - # latents checks - if latents is not None: - if not torch.is_tensor(latents): - raise ValueError("`latents` must be a torch.Tensor.") - if latents.dim() != 4: - raise ValueError(f"`latents` has unsupported dimensions or shape: {latents.shape}.") - - if processing_resolution > 0: - max_orig = max(H, W) - new_H = H * processing_resolution // max_orig - new_W = W * processing_resolution // max_orig - if new_H == 0 or new_W == 0: - raise ValueError(f"Extreme aspect ratio of the input image: [{W} x {H}]") - W, H = new_W, new_H - w = (W + self.vae_scale_factor - 1) // self.vae_scale_factor - h = (H + self.vae_scale_factor - 1) // self.vae_scale_factor - shape_expected = (num_images * ensemble_size, self.vae.config.latent_channels, h, w) - - if latents.shape != shape_expected: - raise ValueError(f"`latents` has unexpected shape={latents.shape} expected={shape_expected}.") - - # generator checks - if generator is not None: - if isinstance(generator, list): - if len(generator) != num_images * ensemble_size: - raise ValueError( - "The number of generators must match the total number of ensemble members for all input images." - ) - if not all(g.device.type == generator[0].device.type for g in generator): - raise ValueError("`generator` device placement is not consistent in the list.") - elif not isinstance(generator, torch.Generator): - raise ValueError(f"Unsupported generator type: {type(generator)}.") - - return num_images - - @torch.compiler.disable - def progress_bar(self, iterable=None, total=None, desc=None, leave=True): - if not hasattr(self, "_progress_bar_config"): - self._progress_bar_config = {} - elif not isinstance(self._progress_bar_config, dict): - raise ValueError( - 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) - progress_bar_config["desc"] = progress_bar_config.get("desc", desc) - progress_bar_config["leave"] = progress_bar_config.get("leave", leave) - if iterable is not None: - return tqdm(iterable, **progress_bar_config) - elif total is not None: - return tqdm(total=total, **progress_bar_config) - else: - raise ValueError("Either `total` or `iterable` has to be defined.") - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - image: PipelineImageInput, - num_inference_steps: Optional[int] = None, - ensemble_size: int = 1, - processing_resolution: Optional[int] = 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, - output_type: str = "np", - output_uncertainty: bool = False, - output_latent: bool = False, - return_dict: bool = True, - ): - """ - 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 - 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 - same width and height. - num_inference_steps (`int`, *optional*, defaults to `None`): - Number of denoising diffusion steps during inference. The default value `None` results in automatic - selection. - ensemble_size (`int`, defaults to `1`): - Number of ensemble predictions. Higher values result in measurable improvements and visual degradation. - processing_resolution (`int`, *optional*, defaults to `None`): - Effective processing resolution. When set to `0`, matches the larger input image dimension. This - produces crisper predictions, but may also lead to the overall loss of global context. The default - value `None` resolves to the optimal value from the model config. - match_input_resolution (`bool`, *optional*, defaults to `True`): - When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer - side of the output will equal to `processing_resolution`. - resample_method_input (`str`, *optional*, defaults to `"bilinear"`): - Resampling method used to resize input images to `processing_resolution`. The accepted values are: - `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`. - resample_method_output (`str`, *optional*, defaults to `"bilinear"`): - Resampling method used to resize output predictions to match the input resolution. The accepted values - are `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`. - batch_size (`int`, *optional*, defaults to `1`): - Batch size; only matters when setting `ensemble_size` or passing a tensor of images. - ensembling_kwargs (`dict`, *optional*, defaults to `None`) - Extra dictionary with arguments for precise ensembling control. The following options are available: - - reduction (`str`, *optional*, defaults to `"closest"`): Defines the ensembling function applied in - every pixel location, can be either `"closest"` or `"mean"`. - 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`): - 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 - values are: `"np"` (numpy array) or `"pt"` (torch tensor). - output_uncertainty (`bool`, *optional*, defaults to `False`): - When enabled, the output's `uncertainty` field contains the predictive uncertainty map, provided that - the `ensemble_size` argument is set to a value above 2. - output_latent (`bool`, *optional*, defaults to `False`): - When enabled, the output's `latent` field contains the latent codes corresponding to the predictions - within the ensemble. These codes can be saved, modified, and used for subsequent calls with the - `latents` argument. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.marigold.MarigoldNormalsOutput`] instead of a plain tuple. - - Examples: - - Returns: - [`~pipelines.marigold.MarigoldNormalsOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.marigold.MarigoldNormalsOutput`] is returned, otherwise a - `tuple` is returned where the first element is the prediction, the second element is the uncertainty - (or `None`), and the third is the latent (or `None`). - """ - - # 0. Resolving variables. - device = self._execution_device - dtype = self.dtype - - # Model-specific optimal default values leading to fast and reasonable results. - if num_inference_steps is None: - num_inference_steps = self.default_denoising_steps - if processing_resolution is None: - processing_resolution = self.default_processing_resolution - - # 1. Check inputs. - num_images = self.check_inputs( - image, - num_inference_steps, - ensemble_size, - processing_resolution, - resample_method_input, - resample_method_output, - batch_size, - ensembling_kwargs, - latents, - generator, - output_type, - output_uncertainty, - ) - - # 2. Prepare empty text conditioning. - # Model invocation: self.tokenizer, self.text_encoder. - if self.empty_text_embedding is None: - prompt = "" - text_inputs = self.tokenizer( - prompt, - padding="do_not_pad", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids.to(device) - self.empty_text_embedding = self.text_encoder(text_input_ids)[0] # [1,2,1024] - - # 3. Preprocess input images. This function loads input image or images of compatible dimensions `(H, W)`, - # optionally downsamples them to the `processing_resolution` `(PH, PW)`, where - # `max(PH, PW) == processing_resolution`, and pads the dimensions to `(PPH, PPW)` such that these values are - # divisible by the latent space downscaling factor (typically 8 in Stable Diffusion). The default value `None` - # of `processing_resolution` resolves to the optimal value from the model config. It is a recommended mode of - # operation and leads to the most reasonable results. Using the native image resolution or any other processing - # resolution can lead to loss of either fine details or global context in the output predictions. - image, padding, original_resolution = self.image_processor.preprocess( - image, processing_resolution, resample_method_input, device, dtype - ) # [N,3,PPH,PPW] - - # 4. Encode input image into latent space. At this step, each of the `N` input images is represented with `E` - # ensemble members. Each ensemble member is an independent diffused prediction, just initialized independently. - # Latents of each such predictions across all input images and all ensemble members are represented in the - # `pred_latent` variable. The variable `image_latent` is of the same shape: it contains each input image encoded - # into latent space and replicated `E` times. The latents can be either generated (see `generator` to ensure - # reproducibility), or passed explicitly via the `latents` argument. The latter can be set outside the pipeline - # code. This behavior can be achieved by setting the `output_latent` argument to `True`. The latent space - # dimensions are `(h, w)`. Encoding into latent space happens in batches of size `batch_size`. - # Model invocation: self.vae.encoder. - image_latent, pred_latent = self.prepare_latents( - image, latents, generator, ensemble_size, batch_size - ) # [N*E,4,h,w], [N*E,4,h,w] - - del image - - batch_empty_text_embedding = self.empty_text_embedding.to(device=device, dtype=dtype).repeat( - batch_size, 1, 1 - ) # [B,1024,2] - - # 5. Process the denoising loop. All `N * E` latents are processed sequentially in batches of size `batch_size`. - # The unet model takes concatenated latent spaces of the input image and the predicted modality as an input, and - # outputs noise for the predicted modality's latent space. The number of denoising diffusion steps is defined by - # `num_inference_steps`. It is either set directly, or resolves to the optimal value specific to the loaded - # model. - # Model invocation: self.unet. - pred_latents = [] - - for i in self.progress_bar( - range(0, num_images * ensemble_size, batch_size), leave=True, desc="Marigold predictions..." - ): - batch_image_latent = image_latent[i : i + batch_size] # [B,4,h,w] - batch_pred_latent = pred_latent[i : i + batch_size] # [B,4,h,w] - effective_batch_size = batch_image_latent.shape[0] - text = batch_empty_text_embedding[:effective_batch_size] # [B,2,1024] - - self.scheduler.set_timesteps(num_inference_steps, device=device) - for t in self.progress_bar(self.scheduler.timesteps, leave=False, desc="Diffusion steps..."): - batch_latent = torch.cat([batch_image_latent, batch_pred_latent], dim=1) # [B,8,h,w] - noise = self.unet(batch_latent, t, encoder_hidden_states=text, return_dict=False)[0] # [B,4,h,w] - batch_pred_latent = self.scheduler.step( - noise, t, batch_pred_latent, generator=generator - ).prev_sample # [B,4,h,w] - - if XLA_AVAILABLE: - xm.mark_step() - - pred_latents.append(batch_pred_latent) - - pred_latent = torch.cat(pred_latents, dim=0) # [N*E,4,h,w] - - del ( - pred_latents, - image_latent, - batch_empty_text_embedding, - batch_image_latent, - batch_pred_latent, - text, - batch_latent, - noise, - ) - - # 6. Decode predictions from latent into pixel space. The resulting `N * E` predictions have shape `(PPH, PPW)`, - # which requires slight postprocessing. Decoding into pixel space happens in batches of size `batch_size`. - # Model invocation: self.vae.decoder. - prediction = torch.cat( - [ - self.decode_prediction(pred_latent[i : i + batch_size]) - for i in range(0, pred_latent.shape[0], batch_size) - ], - dim=0, - ) # [N*E,3,PPH,PPW] - - if not output_latent: - pred_latent = None - - # 7. Remove padding. The output shape is (PH, PW). - prediction = self.image_processor.unpad_image(prediction, padding) # [N*E,3,PH,PW] - - # 8. Ensemble and compute uncertainty (when `output_uncertainty` is set). This code treats each of the `N` - # groups of `E` ensemble predictions independently. For each group it computes an ensembled prediction of shape - # `(PH, PW)` and an optional uncertainty map of the same dimensions. After computing this pair of outputs for - # each group independently, it stacks them respectively into batches of `N` almost final predictions and - # uncertainty maps. - uncertainty = None - if ensemble_size > 1: - prediction = prediction.reshape(num_images, ensemble_size, *prediction.shape[1:]) # [N,E,3,PH,PW] - prediction = [ - self.ensemble_normals(prediction[i], output_uncertainty, **(ensembling_kwargs or {})) - for i in range(num_images) - ] # [ [[1,3,PH,PW], [1,1,PH,PW]], ... ] - prediction, uncertainty = zip(*prediction) # [[1,3,PH,PW], ... ], [[1,1,PH,PW], ... ] - prediction = torch.cat(prediction, dim=0) # [N,3,PH,PW] - if output_uncertainty: - uncertainty = torch.cat(uncertainty, dim=0) # [N,1,PH,PW] - else: - uncertainty = None - - # 9. If `match_input_resolution` is set, the output prediction and the uncertainty are upsampled to match the - # input resolution `(H, W)`. This step may introduce upsampling artifacts, and therefore can be disabled. - # After upsampling, the native resolution normal maps are renormalized to unit length to reduce the artifacts. - # Depending on the downstream use-case, upsampling can be also chosen based on the tolerated artifacts by - # setting the `resample_method_output` parameter (e.g., to `"nearest"`). - if match_input_resolution: - prediction = self.image_processor.resize_antialias( - prediction, original_resolution, resample_method_output, is_aa=False - ) # [N,3,H,W] - prediction = self.normalize_normals(prediction) # [N,3,H,W] - if uncertainty is not None and output_uncertainty: - uncertainty = self.image_processor.resize_antialias( - uncertainty, original_resolution, resample_method_output, is_aa=False - ) # [N,1,H,W] - - # 10. Prepare the final outputs. - if output_type == "np": - prediction = self.image_processor.pt_to_numpy(prediction) # [N,H,W,3] - if uncertainty is not None and output_uncertainty: - uncertainty = self.image_processor.pt_to_numpy(uncertainty) # [N,H,W,1] - - # 11. Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (prediction, uncertainty, pred_latent) - - return MarigoldNormalsOutput( - prediction=prediction, - uncertainty=uncertainty, - latent=pred_latent, - ) - - # Copied from diffusers.pipelines.marigold.pipeline_marigold_depth.MarigoldDepthPipeline.prepare_latents - def prepare_latents( - self, - image: torch.Tensor, - latents: Optional[torch.Tensor], - generator: Optional[torch.Generator], - ensemble_size: int, - batch_size: int, - ) -> Tuple[torch.Tensor, torch.Tensor]: - def retrieve_latents(encoder_output): - if hasattr(encoder_output, "latent_dist"): - 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") - - image_latent = torch.cat( - [ - retrieve_latents(self.vae.encode(image[i : i + batch_size])) - for i in range(0, image.shape[0], batch_size) - ], - dim=0, - ) # [N,4,h,w] - image_latent = image_latent * self.vae.config.scaling_factor - image_latent = image_latent.repeat_interleave(ensemble_size, dim=0) # [N*E,4,h,w] - - pred_latent = latents - if pred_latent is None: - pred_latent = randn_tensor( - image_latent.shape, - generator=generator, - device=image_latent.device, - dtype=image_latent.dtype, - ) # [N*E,4,h,w] - - return image_latent, pred_latent - - def decode_prediction(self, pred_latent: torch.Tensor) -> torch.Tensor: - if pred_latent.dim() != 4 or pred_latent.shape[1] != self.vae.config.latent_channels: - raise ValueError( - f"Expecting 4D tensor of shape [B,{self.vae.config.latent_channels},H,W]; got {pred_latent.shape}." - ) - - prediction = self.vae.decode(pred_latent / self.vae.config.scaling_factor, return_dict=False)[0] # [B,3,H,W] - - prediction = torch.clip(prediction, -1.0, 1.0) - - if not self.use_full_z_range: - prediction[:, 2, :, :] *= 0.5 - prediction[:, 2, :, :] += 0.5 - - prediction = self.normalize_normals(prediction) # [B,3,H,W] - - return prediction # [B,3,H,W] - - @staticmethod - def normalize_normals(normals: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: - if normals.dim() != 4 or normals.shape[1] != 3: - raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.") - - norm = torch.norm(normals, dim=1, keepdim=True) - normals /= norm.clamp(min=eps) - - return normals - - @staticmethod - def ensemble_normals( - normals: torch.Tensor, output_uncertainty: bool, reduction: str = "closest" - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """ - 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)`. - - Args: - normals (`torch.Tensor`): - Input ensemble normals maps. - output_uncertainty (`bool`, *optional*, defaults to `False`): - Whether to output uncertainty map. - reduction (`str`, *optional*, defaults to `"closest"`): - Reduction method used to ensemble aligned predictions. The accepted values are: `"closest"` and - `"mean"`. - - Returns: - A tensor of aligned and ensembled normals maps with shape `(1, 3, H, W)` and optionally a tensor of - uncertainties of shape `(1, 1, H, W)`. - """ - if normals.dim() != 4 or normals.shape[1] != 3: - raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.") - if reduction not in ("closest", "mean"): - raise ValueError(f"Unrecognized reduction method: {reduction}.") - - mean_normals = normals.mean(dim=0, keepdim=True) # [1,3,H,W] - mean_normals = MarigoldNormalsPipeline.normalize_normals(mean_normals) # [1,3,H,W] - - sim_cos = (mean_normals * normals).sum(dim=1, keepdim=True) # [E,1,H,W] - sim_cos = sim_cos.clamp(-1, 1) # required to avoid NaN in uncertainty with fp16 - - uncertainty = None - if output_uncertainty: - uncertainty = sim_cos.arccos() # [E,1,H,W] - uncertainty = uncertainty.mean(dim=0, keepdim=True) / np.pi # [1,1,H,W] - - if reduction == "mean": - return mean_normals, uncertainty # [1,3,H,W], [1,1,H,W] - - closest_indices = sim_cos.argmax(dim=0, keepdim=True) # [1,1,H,W] - closest_indices = closest_indices.repeat(1, 3, 1, 1) # [1,3,H,W] - closest_normals = torch.gather(normals, 0, closest_indices) # [1,3,H,W] - - return closest_normals, uncertainty # [1,3,H,W], [1,1,H,W] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/__init__.py deleted file mode 100644 index a8fd4da9fd36e28479df6258f7de108b63d3c2bb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -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_mochi"] = ["MochiPipeline"] - -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_mochi import MochiPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index ddc7a4b9f39bc6005a1f1161fe0ead09533eef22..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/__pycache__/pipeline_mochi.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/__pycache__/pipeline_mochi.cpython-310.pyc deleted file mode 100644 index fe0f5d26348ef612c2ff42f4dbeb4493ec1a2381..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/__pycache__/pipeline_mochi.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index 3885b5c7c8af0ba013b1b9840e577c9a7950ed3b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/pipeline_mochi.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/pipeline_mochi.py deleted file mode 100644 index 7712b415242c71f927c105d76dc828e4358409ce..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/pipeline_mochi.py +++ /dev/null @@ -1,744 +0,0 @@ -# Copyright 2025 Genmo 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 T5EncoderModel, T5TokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...loaders import Mochi1LoraLoaderMixin -from ...models import AutoencoderKLMochi, MochiTransformer3DModel -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 .pipeline_output import MochiPipelineOutput - - -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 MochiPipeline - >>> from diffusers.utils import export_to_video - - >>> pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview", torch_dtype=torch.bfloat16) - >>> pipe.enable_model_cpu_offload() - >>> pipe.enable_vae_tiling() - >>> prompt = "Close-up of a chameleon's eye, with its scaly skin changing color. Ultra high resolution 4k." - >>> frames = pipe(prompt, num_inference_steps=28, guidance_scale=3.5).frames[0] - >>> export_to_video(frames, "mochi.mp4") - ``` -""" - - -# from: https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77 -def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None): - if linear_steps is None: - linear_steps = num_steps // 2 - linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)] - threshold_noise_step_diff = linear_steps - threshold_noise * num_steps - quadratic_steps = num_steps - linear_steps - quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2) - linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2) - const = quadratic_coef * (linear_steps**2) - quadratic_sigma_schedule = [ - quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, num_steps) - ] - sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule - sigma_schedule = [1.0 - x for x in sigma_schedule] - return sigma_schedule - - -# 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 MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin): - r""" - The mochi pipeline for text-to-video generation. - - Reference: https://github.com/genmoai/models - - Args: - transformer ([`MochiTransformer3DModel`]): - 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 ([`AutoencoderKLMochi`]): - Variational Auto-Encoder (VAE) Model to encode and decode videos 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). - """ - - 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: AutoencoderKLMochi, - text_encoder: T5EncoderModel, - tokenizer: T5TokenizerFast, - transformer: MochiTransformer3DModel, - force_zeros_for_empty_prompt: bool = False, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - ) - # TODO: determine these scaling factors from model parameters - self.vae_spatial_scale_factor = 8 - self.vae_temporal_scale_factor = 6 - self.patch_size = 2 - - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_scale_factor) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 256 - ) - self.default_height = 480 - self.default_width = 848 - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - - def _get_t5_prompt_embeds( - self, - prompt: Union[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 = 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) - - # The original Mochi implementation zeros out empty negative prompts - # but this can lead to overflow when placing the entire pipeline under the autocast context - # adding this here so that we can enable zeroing prompts if necessary - if self.config.force_zeros_for_empty_prompt and (prompt == "" or prompt[-1] == ""): - text_input_ids = torch.zeros_like(text_input_ids, device=device) - prompt_attention_mask = torch.zeros_like(prompt_attention_mask, dtype=torch.bool, device=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), attention_mask=prompt_attention_mask)[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 - - # 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, - 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, - max_sequence_length: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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 - - 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 % 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_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}." - ) - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - num_frames, - dtype, - device, - generator, - latents=None, - ): - height = height // self.vae_spatial_scale_factor - width = width // self.vae_spatial_scale_factor - num_frames = (num_frames - 1) // self.vae_temporal_scale_factor + 1 - - shape = (batch_size, num_channels_latents, num_frames, height, width) - - if latents is not None: - return latents.to(device=device, dtype=dtype) - 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=torch.float32) - latents = latents.to(dtype) - return latents - - @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 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: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_frames: int = 19, - num_inference_steps: int = 64, - 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", - 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"], - max_sequence_length: int = 256, - ): - 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.default_height`): - The height in pixels of the generated image. This is set to 480 by default for the best results. - width (`int`, *optional*, defaults to `self.default_width`): - The width in pixels of the generated image. This is set to 848 by default for the best results. - num_frames (`int`, defaults to `19`): - The number of video frames to generate - 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`, defaults to `4.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_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 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.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. - 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.mochi.MochiPipelineOutput`] 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.mochi.MochiPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.mochi.MochiPipelineOutput`] 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 - - height = height or self.default_height - width = width or self.default_width - - # 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._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, - 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, - ) - # 4. Prepare latent variables - 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, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - 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) - - # 5. Prepare timestep - # from https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77 - threshold_noise = 0.025 - sigmas = linear_quadratic_schedule(num_inference_steps, threshold_noise) - sigmas = np.array(sigmas) - - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - device, - timesteps, - sigmas, - ) - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - # 6. Denoising loop - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # Note: Mochi uses reversed timesteps. To ensure compatibility with methods like FasterCache, we need - # to make sure we're using the correct non-reversed timestep values. - self._current_timestep = 1000 - t - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - encoder_hidden_states=prompt_embeds, - timestep=timestep, - encoder_attention_mask=prompt_attention_mask, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - # Mochi CFG + Sampling runs in FP32 - noise_pred = noise_pred.to(torch.float32) - - 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) - - # compute the previous noisy sample x_t -> x_t-1 - latents_dtype = latents.dtype - latents = self.scheduler.step(noise_pred, t, latents.to(torch.float32), return_dict=False)[0] - latents = latents.to(latents_dtype) - - 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": - video = latents - else: - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None - has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None - if has_latents_mean and has_latents_std: - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean - else: - latents = latents / self.vae.config.scaling_factor - - video = self.vae.decode(latents, 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 MochiPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/pipeline_output.py deleted file mode 100644 index d15827bc0084ee07175252aa54bc589e604df67d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/mochi/pipeline_output.py +++ /dev/null @@ -1,20 +0,0 @@ -from dataclasses import dataclass - -import torch - -from diffusers.utils import BaseOutput - - -@dataclass -class MochiPipelineOutput(BaseOutput): - r""" - 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 - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/musicldm/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/musicldm/__init__.py deleted file mode 100644 index ed71eeb1d99b28f20f7cd94776c0303208620653..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/musicldm/__init__.py +++ /dev/null @@ -1,49 +0,0 @@ -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 = {} -_import_structure = {} - -try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")): - 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_musicldm"] = ["MusicLDMPipeline"] - - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")): - raise OptionalDependencyNotAvailable() - - except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_objects import * - else: - from .pipeline_musicldm import MusicLDMPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/musicldm/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/musicldm/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index f9792889ebf1f32df922ea02d599342b3712169e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/musicldm/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/musicldm/__pycache__/pipeline_musicldm.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/musicldm/__pycache__/pipeline_musicldm.cpython-310.pyc deleted file mode 100644 index ce2f2a433a257f14353a647949e97e6caaff5d25..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/musicldm/__pycache__/pipeline_musicldm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/musicldm/pipeline_musicldm.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/musicldm/pipeline_musicldm.py deleted file mode 100644 index c909e5eb0d26d9c5f79ba8646015ee68ccd1d22c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/musicldm/pipeline_musicldm.py +++ /dev/null @@ -1,653 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import torch -from transformers import ( - ClapFeatureExtractor, - ClapModel, - ClapTextModelWithProjection, - RobertaTokenizer, - RobertaTokenizerFast, - SpeechT5HifiGan, -) - -from ...models import AutoencoderKL, UNet2DConditionModel -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - is_accelerate_available, - is_accelerate_version, - is_librosa_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import empty_device_cache, get_device, randn_tensor -from ..pipeline_utils import AudioPipelineOutput, DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin - - -if is_librosa_available(): - import librosa - - -from ...utils import is_torch_xla_available - - -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 - >>> from diffusers import MusicLDMPipeline - >>> import torch - >>> import scipy - - >>> repo_id = "ucsd-reach/musicldm" - >>> pipe = MusicLDMPipeline.from_pretrained(repo_id, torch_dtype=torch.float16) - >>> pipe = pipe.to("cuda") - - >>> prompt = "Techno music with a strong, upbeat tempo and high melodic riffs" - >>> audio = pipe(prompt, num_inference_steps=10, audio_length_in_s=5.0).audios[0] - - >>> # save the audio sample as a .wav file - >>> scipy.io.wavfile.write("techno.wav", rate=16000, data=audio) - ``` -""" - - -class MusicLDMPipeline(DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin): - _last_supported_version = "0.33.1" - r""" - Pipeline for text-to-audio generation using MusicLDM. - - 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: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.ClapModel`]): - Frozen text-audio embedding model (`ClapTextModel`), specifically the - [laion/clap-htsat-unfused](https://huggingface.co/laion/clap-htsat-unfused) variant. - tokenizer ([`PreTrainedTokenizer`]): - A [`~transformers.RobertaTokenizer`] to tokenize text. - feature_extractor ([`~transformers.ClapFeatureExtractor`]): - Feature extractor to compute mel-spectrograms from audio waveforms. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded audio latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded audio latents. Can be one of - [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. - vocoder ([`~transformers.SpeechT5HifiGan`]): - Vocoder of class `SpeechT5HifiGan`. - """ - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: Union[ClapTextModelWithProjection, ClapModel], - tokenizer: Union[RobertaTokenizer, RobertaTokenizerFast], - feature_extractor: Optional[ClapFeatureExtractor], - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - vocoder: SpeechT5HifiGan, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - feature_extractor=feature_extractor, - unet=unet, - scheduler=scheduler, - vocoder=vocoder, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - - def _encode_prompt( - self, - prompt, - device, - num_waveforms_per_prompt, - do_classifier_free_guidance, - negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = 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_waveforms_per_prompt (`int`): - 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*): - 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`). - 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. - """ - 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: - 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 - attention_mask = text_inputs.attention_mask - 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 CLAP can only handle sequences up to" - f" {self.tokenizer.model_max_length} tokens: {removed_text}" - ) - - prompt_embeds = self.text_encoder.get_text_features( - text_input_ids.to(device), - attention_mask=attention_mask.to(device), - ) - - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.text_model.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_waveforms_per_prompt) - prompt_embeds = prompt_embeds.view(bs_embed * num_waveforms_per_prompt, seq_len) - - # 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 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 - - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - - uncond_input_ids = uncond_input.input_ids.to(device) - attention_mask = uncond_input.attention_mask.to(device) - - negative_prompt_embeds = self.text_encoder.get_text_features( - uncond_input_ids, - attention_mask=attention_mask, - ) - - 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=self.text_encoder.text_model.dtype, device=device) - - negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_waveforms_per_prompt) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_waveforms_per_prompt, seq_len) - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - return prompt_embeds - - # Copied from diffusers.pipelines.audioldm.pipeline_audioldm.AudioLDMPipeline.mel_spectrogram_to_waveform - def mel_spectrogram_to_waveform(self, mel_spectrogram): - if mel_spectrogram.dim() == 4: - mel_spectrogram = mel_spectrogram.squeeze(1) - - waveform = self.vocoder(mel_spectrogram) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - waveform = waveform.cpu().float() - return waveform - - # Copied from diffusers.pipelines.audioldm2.pipeline_audioldm2.AudioLDM2Pipeline.score_waveforms - def score_waveforms(self, text, audio, num_waveforms_per_prompt, device, dtype): - if not is_librosa_available(): - logger.info( - "Automatic scoring of the generated audio waveforms against the input prompt text requires the " - "`librosa` package to resample the generated waveforms. Returning the audios in the order they were " - "generated. To enable automatic scoring, install `librosa` with: `pip install librosa`." - ) - return audio - inputs = self.tokenizer(text, return_tensors="pt", padding=True) - resampled_audio = librosa.resample( - audio.numpy(), orig_sr=self.vocoder.config.sampling_rate, target_sr=self.feature_extractor.sampling_rate - ) - inputs["input_features"] = self.feature_extractor( - list(resampled_audio), return_tensors="pt", sampling_rate=self.feature_extractor.sampling_rate - ).input_features.type(dtype) - inputs = inputs.to(device) - - # compute the audio-text similarity score using the CLAP model - logits_per_text = self.text_encoder(**inputs).logits_per_text - # sort by the highest matching generations per prompt - indices = torch.argsort(logits_per_text, dim=1, descending=True)[:, :num_waveforms_per_prompt] - audio = torch.index_select(audio, 0, indices.reshape(-1).cpu()) - return audio - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.audioldm.pipeline_audioldm.AudioLDMPipeline.check_inputs - def check_inputs( - self, - prompt, - audio_length_in_s, - vocoder_upsample_factor, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - min_audio_length_in_s = vocoder_upsample_factor * self.vae_scale_factor - if audio_length_in_s < min_audio_length_in_s: - raise ValueError( - f"`audio_length_in_s` has to be a positive value greater than or equal to {min_audio_length_in_s}, but " - f"is {audio_length_in_s}." - ) - - if self.vocoder.config.model_in_dim % self.vae_scale_factor != 0: - raise ValueError( - f"The number of frequency bins in the vocoder's log-mel spectrogram has to be divisible by the " - f"VAE scale factor, but got {self.vocoder.config.model_in_dim} bins and a scale factor of " - f"{self.vae_scale_factor}." - ) - - if (callback_steps is None) or ( - 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 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}." - ) - - # Copied from diffusers.pipelines.audioldm.pipeline_audioldm.AudioLDMPipeline.prepare_latents - def prepare_latents(self, batch_size, num_channels_latents, height, dtype, device, generator, latents=None): - shape = ( - batch_size, - num_channels_latents, - int(height) // self.vae_scale_factor, - int(self.vocoder.config.model_in_dim) // 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def enable_model_cpu_offload(self, gpu_id=0): - 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 - `forward` method is called, and the model remains in accelerator until the next model runs. Memory savings are - lower than with `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution - of the `unet`. - """ - if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"): - from accelerate import cpu_offload_with_hook - else: - raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.") - - device_type = get_device() - device = torch.device(f"{device_type}:{gpu_id}") - - if self.device.type != "cpu": - self.to("cpu", silence_dtype_warnings=True) - empty_device_cache() # otherwise we don't see the memory savings (but they probably exist) - - model_sequence = [ - self.text_encoder.text_model, - self.text_encoder.text_projection, - self.unet, - self.vae, - self.vocoder, - self.text_encoder, - ] - - hook = None - for cpu_offloaded_model in model_sequence: - _, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook) - - # We'll offload the last model manually. - self.final_offload_hook = hook - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - audio_length_in_s: Optional[float] = 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, - 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, - 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", - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - num_inference_steps (`int`, *optional*, defaults to 200): - The number of denoising steps. More denoising steps usually lead to a higher quality audio at the - expense of slower inference. - 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*): - 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): - The number of waveforms to generate per prompt. If `num_waveforms_per_prompt > 1`, the text encoding - model is a joint text-audio model ([`~transformers.ClapModel`]), and the tokenizer is a - `[~transformers.ClapProcessor]`, then automatic scoring will be performed between the generated outputs - and the input text. This scoring ranks the generated waveforms based on their cosine similarity to text - input in the joint text-audio embedding space. - 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*): - 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. - 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. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.AudioPipelineOutput`] instead of a plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - output_type (`str`, *optional*, defaults to `"np"`): - The output format of the generated audio. Choose between `"np"` to return a NumPy `np.ndarray` or - `"pt"` to return a PyTorch `torch.Tensor` object. Set to `"latent"` to return the latent diffusion - model (LDM) output. - - Examples: - - Returns: - [`~pipelines.AudioPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.AudioPipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated audio. - """ - # 0. Convert audio input length from seconds to spectrogram height - vocoder_upsample_factor = np.prod(self.vocoder.config.upsample_rates) / self.vocoder.config.sampling_rate - - if audio_length_in_s is None: - audio_length_in_s = self.unet.config.sample_size * self.vae_scale_factor * vocoder_upsample_factor - - height = int(audio_length_in_s / vocoder_upsample_factor) - - original_waveform_length = int(audio_length_in_s * self.vocoder.config.sampling_rate) - if height % self.vae_scale_factor != 0: - height = int(np.ceil(height / self.vae_scale_factor)) * self.vae_scale_factor - logger.info( - f"Audio length in seconds {audio_length_in_s} is increased to {height * vocoder_upsample_factor} " - f"so that it can be handled by the model. It will be cut to {audio_length_in_s} after the " - f"denoising process." - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - audio_length_in_s, - vocoder_upsample_factor, - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds = self._encode_prompt( - prompt, - device, - num_waveforms_per_prompt, - do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_waveforms_per_prompt, - num_channels_latents, - height, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=None, - class_labels=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - self.maybe_free_model_hooks() - - # 8. Post-processing - if not output_type == "latent": - latents = 1 / self.vae.config.scaling_factor * latents - mel_spectrogram = self.vae.decode(latents).sample - else: - return AudioPipelineOutput(audios=latents) - - audio = self.mel_spectrogram_to_waveform(mel_spectrogram) - - audio = audio[:, :original_waveform_length] - - # 9. Automatic scoring - if num_waveforms_per_prompt > 1 and prompt is not None: - audio = self.score_waveforms( - text=prompt, - audio=audio, - num_waveforms_per_prompt=num_waveforms_per_prompt, - device=device, - dtype=prompt_embeds.dtype, - ) - - if output_type == "np": - audio = audio.numpy() - - if not return_dict: - return (audio,) - - return AudioPipelineOutput(audios=audio) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/__init__.py deleted file mode 100644 index 557e7c08dc225f14766701b4953e4ff3c9dddb63..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/__init__.py +++ /dev/null @@ -1,50 +0,0 @@ -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_omnigen"] = ["OmniGenPipeline"] - - -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_omnigen import OmniGenPipeline - - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index efc5cc9de40b29fd5b1218d715bd78ba3664f0ef..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/__pycache__/pipeline_omnigen.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/__pycache__/pipeline_omnigen.cpython-310.pyc deleted file mode 100644 index fc213ccab631cf9314d7b365d6c9eeacd685e4ff..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/__pycache__/pipeline_omnigen.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/__pycache__/processor_omnigen.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/__pycache__/processor_omnigen.cpython-310.pyc deleted file mode 100644 index b6fab173398c4faaf5094dafb2a1822afb431853..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/__pycache__/processor_omnigen.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/pipeline_omnigen.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/pipeline_omnigen.py deleted file mode 100644 index 1254b6725fefb1d56321b17b2fe00aaea4849847..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/pipeline_omnigen.py +++ /dev/null @@ -1,514 +0,0 @@ -# Copyright 2025 OmniGen 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 Callable, Dict, List, Optional, Union - -import numpy as np -import torch -from transformers import LlamaTokenizer - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...models.autoencoders import AutoencoderKL -from ...models.transformers import OmniGenTransformer2DModel -from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import is_torch_xla_available, is_torchvision_available, logging, replace_example_docstring -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -if is_torchvision_available(): - from .processor_omnigen import OmniGenMultiModalProcessor - -if is_torch_xla_available(): - 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 OmniGenPipeline - - >>> pipe = OmniGenPipeline.from_pretrained("Shitao/OmniGen-v1-diffusers", 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=2.5).images[0] - >>> image.save("output.png") - ``` -""" - - -# 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 OmniGenPipeline( - DiffusionPipeline, -): - r""" - The OmniGen pipeline for multimodal-to-image generation. - - Reference: https://huggingface.co/papers/2409.11340 - - Args: - transformer ([`OmniGenTransformer2DModel`]): - Autoregressive Transformer architecture for OmniGen. - 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. - tokenizer (`LlamaTokenizer`): - Text tokenizer of class. - [LlamaTokenizer](https://huggingface.co/docs/transformers/main/model_doc/llama#transformers.LlamaTokenizer). - """ - - model_cpu_offload_seq = "transformer->vae" - _optional_components = [] - _callback_tensor_inputs = ["latents"] - - def __init__( - self, - transformer: OmniGenTransformer2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - tokenizer: LlamaTokenizer, - ): - super().__init__() - - self.register_modules( - vae=vae, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - ) - self.vae_scale_factor = ( - 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) is not None else 8 - ) - # OmniGen 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.multimodal_processor = OmniGenMultiModalProcessor(tokenizer, max_image_size=1024) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 120000 - ) - self.default_sample_size = 128 - - def encode_input_images( - self, - input_pixel_values: List[torch.Tensor], - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - ): - """ - get the continue embedding of input images by VAE - - Args: - input_pixel_values: normalized pixel of input images - device: - Returns: torch.Tensor - """ - device = device or self._execution_device - dtype = dtype or self.vae.dtype - - input_img_latents = [] - for img in input_pixel_values: - img = self.vae.encode(img.to(device, dtype)).latent_dist.sample().mul_(self.vae.config.scaling_factor) - input_img_latents.append(img) - return input_img_latents - - def check_inputs( - self, - prompt, - input_images, - height, - width, - use_input_image_size_as_output, - callback_on_step_end_tensor_inputs=None, - ): - if input_images is not None: - if len(input_images) != len(prompt): - raise ValueError( - f"The number of prompts: {len(prompt)} does not match the number of input images: {len(input_images)}." - ) - for i in range(len(input_images)): - if input_images[i] is not None: - if not all(f"<|image_{k + 1}|>" in prompt[i] for k in range(len(input_images[i]))): - raise ValueError( - f"prompt `{prompt[i]}` doesn't have enough placeholders for the input images `{input_images[i]}`" - ) - - 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 use_input_image_size_as_output: - if input_images is None or input_images[0] is None: - raise ValueError( - "`use_input_image_size_as_output` is set to True, but no input image was found. If you are performing a text-to-image task, please set it to False." - ) - - 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]}" - ) - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_latents - 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=device, dtype=dtype) - - 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 - - @property - def guidance_scale(self): - return self._guidance_scale - - @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]], - input_images: Union[PipelineImageInput, List[PipelineImageInput]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - max_input_image_size: int = 1024, - 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", - return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], 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*): - 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*): - 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. - 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. - 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*): - 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 2.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. - img_guidance_scale (`float`, *optional*, defaults to 1.6): - Defined as equation 3 in [Instrucpix2pix](https://huggingface.co/papers/2211.09800). - use_input_image_size_as_output (bool, defaults to False): - whether to use the input image size as the output image size, which can be used for single-image input, - 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*): - 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 ge generated by sampling using the supplied random `generator`. - 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.FluxPipelineOutput`] instead of a plain tuple. - 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. - - Examples: - - Returns: [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is returned - where 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 - num_cfg = 2 if input_images is not None else 1 - use_img_cfg = True if input_images is not None else False - if isinstance(prompt, str): - prompt = [prompt] - input_images = [input_images] - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - input_images, - height, - width, - use_input_image_size_as_output, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._interrupt = False - - # 2. Define call parameters - batch_size = len(prompt) - device = self._execution_device - - # 3. process multi-modal instructions - if max_input_image_size != self.multimodal_processor.max_image_size: - self.multimodal_processor.reset_max_image_size(max_image_size=max_input_image_size) - processed_data = self.multimodal_processor( - prompt, - input_images, - height=height, - width=width, - use_img_cfg=use_img_cfg, - use_input_image_size_as_output=use_input_image_size_as_output, - num_images_per_prompt=num_images_per_prompt, - ) - processed_data["input_ids"] = processed_data["input_ids"].to(device) - processed_data["attention_mask"] = processed_data["attention_mask"].to(device) - processed_data["position_ids"] = processed_data["position_ids"].to(device) - - # 4. Encode input images - input_img_latents = self.encode_input_images(processed_data["input_pixel_values"], device=device) - - # 5. Prepare timesteps - sigmas = np.linspace(1, 0, num_inference_steps + 1)[:num_inference_steps] - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas=sigmas - ) - self._num_timesteps = len(timesteps) - - # 6. Prepare latents - transformer_dtype = self.transformer.dtype - if use_input_image_size_as_output: - height, width = processed_data["input_pixel_values"][0].shape[-2:] - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - torch.float32, - device, - generator, - latents, - ) - - # 8. Denoising loop - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * (num_cfg + 1)) - latent_model_input = latent_model_input.to(transformer_dtype) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - input_ids=processed_data["input_ids"], - input_img_latents=input_img_latents, - input_image_sizes=processed_data["input_image_sizes"], - attention_mask=processed_data["attention_mask"], - position_ids=processed_data["position_ids"], - return_dict=False, - )[0] - - if num_cfg == 2: - cond, uncond, img_cond = torch.split(noise_pred, len(noise_pred) // 3, dim=0) - noise_pred = uncond + img_guidance_scale * (img_cond - uncond) + guidance_scale * (cond - img_cond) - else: - cond, uncond = torch.split(noise_pred, len(noise_pred) // 2, dim=0) - noise_pred = uncond + guidance_scale * (cond - uncond) - - # compute the previous noisy sample x_t -> x_t-1 - 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, t, callback_kwargs) - - latents = callback_outputs.pop("latents", latents) - - progress_bar.update() - - if not output_type == "latent": - latents = latents.to(self.vae.dtype) - latents = latents / self.vae.config.scaling_factor - image = self.vae.decode(latents, return_dict=False)[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 ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/processor_omnigen.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/processor_omnigen.py deleted file mode 100644 index 7ed11871bb2af43d5525f246928b98e98e9ab4fc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/omnigen/processor_omnigen.py +++ /dev/null @@ -1,332 +0,0 @@ -# Copyright 2025 OmniGen 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 re -from typing import Dict, List - -import numpy as np -import torch -from PIL import Image - -from ...utils import is_torchvision_available - - -if is_torchvision_available(): - from torchvision import transforms - - -def crop_image(pil_image, max_image_size): - """ - Crop the image so that its height and width does not exceed `max_image_size`, while ensuring both the height and - width are multiples of 16. - """ - while min(*pil_image.size) >= 2 * max_image_size: - pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), resample=Image.BOX) - - if max(*pil_image.size) > max_image_size: - scale = max_image_size / max(*pil_image.size) - pil_image = pil_image.resize(tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC) - - if min(*pil_image.size) < 16: - scale = 16 / min(*pil_image.size) - pil_image = pil_image.resize(tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC) - - arr = np.array(pil_image) - crop_y1 = (arr.shape[0] % 16) // 2 - crop_y2 = arr.shape[0] % 16 - crop_y1 - - crop_x1 = (arr.shape[1] % 16) // 2 - crop_x2 = arr.shape[1] % 16 - crop_x1 - - arr = arr[crop_y1 : arr.shape[0] - crop_y2, crop_x1 : arr.shape[1] - crop_x2] - return Image.fromarray(arr) - - -class OmniGenMultiModalProcessor: - def __init__(self, text_tokenizer, max_image_size: int = 1024): - self.text_tokenizer = text_tokenizer - self.max_image_size = max_image_size - - self.image_transform = transforms.Compose( - [ - transforms.Lambda(lambda pil_image: crop_image(pil_image, max_image_size)), - transforms.ToTensor(), - transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), - ] - ) - - self.collator = OmniGenCollator() - - def reset_max_image_size(self, max_image_size): - self.max_image_size = max_image_size - self.image_transform = transforms.Compose( - [ - transforms.Lambda(lambda pil_image: crop_image(pil_image, max_image_size)), - transforms.ToTensor(), - transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), - ] - ) - - def process_image(self, image): - if isinstance(image, str): - image = Image.open(image).convert("RGB") - return self.image_transform(image) - - def process_multi_modal_prompt(self, text, input_images): - text = self.add_prefix_instruction(text) - if input_images is None or len(input_images) == 0: - model_inputs = self.text_tokenizer(text) - return {"input_ids": model_inputs.input_ids, "pixel_values": None, "image_sizes": None} - - pattern = r"<\|image_\d+\|>" - prompt_chunks = [self.text_tokenizer(chunk).input_ids for chunk in re.split(pattern, text)] - - for i in range(1, len(prompt_chunks)): - if prompt_chunks[i][0] == 1: - prompt_chunks[i] = prompt_chunks[i][1:] - - image_tags = re.findall(pattern, text) - image_ids = [int(s.split("|")[1].split("_")[-1]) for s in image_tags] - - unique_image_ids = sorted(set(image_ids)) - assert unique_image_ids == list(range(1, len(unique_image_ids) + 1)), ( - f"image_ids must start from 1, and must be continuous int, e.g. [1, 2, 3], cannot be {unique_image_ids}" - ) - # total images must be the same as the number of image tags - assert len(unique_image_ids) == len(input_images), ( - f"total images must be the same as the number of image tags, got {len(unique_image_ids)} image tags and {len(input_images)} images" - ) - - input_images = [input_images[x - 1] for x in image_ids] - - all_input_ids = [] - img_inx = [] - for i in range(len(prompt_chunks)): - all_input_ids.extend(prompt_chunks[i]) - if i != len(prompt_chunks) - 1: - start_inx = len(all_input_ids) - size = input_images[i].size(-2) * input_images[i].size(-1) // 16 // 16 - img_inx.append([start_inx, start_inx + size]) - all_input_ids.extend([0] * size) - - return {"input_ids": all_input_ids, "pixel_values": input_images, "image_sizes": img_inx} - - def add_prefix_instruction(self, prompt): - user_prompt = "<|user|>\n" - generation_prompt = "Generate an image according to the following instructions\n" - assistant_prompt = "<|assistant|>\n<|diffusion|>" - prompt_suffix = "<|end|>\n" - prompt = f"{user_prompt}{generation_prompt}{prompt}{prompt_suffix}{assistant_prompt}" - return prompt - - def __call__( - self, - 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.", - use_img_cfg: bool = True, - separate_cfg_input: bool = False, - use_input_image_size_as_output: bool = False, - num_images_per_prompt: int = 1, - ) -> Dict: - if isinstance(instructions, str): - instructions = [instructions] - input_images = [input_images] - - input_data = [] - for i in range(len(instructions)): - cur_instruction = instructions[i] - cur_input_images = None if input_images is None else input_images[i] - if cur_input_images is not None and len(cur_input_images) > 0: - cur_input_images = [self.process_image(x) for x in cur_input_images] - else: - cur_input_images = None - assert "<|image_1|>" not in cur_instruction - - mllm_input = self.process_multi_modal_prompt(cur_instruction, cur_input_images) - - neg_mllm_input, img_cfg_mllm_input = None, None - neg_mllm_input = self.process_multi_modal_prompt(negative_prompt, None) - if use_img_cfg: - if cur_input_images is not None and len(cur_input_images) >= 1: - img_cfg_prompt = [f"<|image_{i + 1}|>" for i in range(len(cur_input_images))] - img_cfg_mllm_input = self.process_multi_modal_prompt(" ".join(img_cfg_prompt), cur_input_images) - else: - img_cfg_mllm_input = neg_mllm_input - - for _ in range(num_images_per_prompt): - if use_input_image_size_as_output: - input_data.append( - ( - mllm_input, - neg_mllm_input, - img_cfg_mllm_input, - [mllm_input["pixel_values"][0].size(-2), mllm_input["pixel_values"][0].size(-1)], - ) - ) - else: - input_data.append((mllm_input, neg_mllm_input, img_cfg_mllm_input, [height, width])) - - return self.collator(input_data) - - -class OmniGenCollator: - def __init__(self, pad_token_id=2, hidden_size=3072): - self.pad_token_id = pad_token_id - self.hidden_size = hidden_size - - def create_position(self, attention_mask, num_tokens_for_output_images): - position_ids = [] - text_length = attention_mask.size(-1) - img_length = max(num_tokens_for_output_images) - for mask in attention_mask: - temp_l = torch.sum(mask) - temp_position = [0] * (text_length - temp_l) + list( - range(temp_l + img_length + 1) - ) # we add a time embedding into the sequence, so add one more token - position_ids.append(temp_position) - return torch.LongTensor(position_ids) - - def create_mask(self, attention_mask, num_tokens_for_output_images): - """ - OmniGen applies causal attention to each element in the sequence, but applies bidirectional attention within - each image sequence References: [OmniGen](https://huggingface.co/papers/2409.11340) - """ - extended_mask = [] - padding_images = [] - text_length = attention_mask.size(-1) - img_length = max(num_tokens_for_output_images) - seq_len = text_length + img_length + 1 # we add a time embedding into the sequence, so add one more token - inx = 0 - for mask in attention_mask: - temp_l = torch.sum(mask) - pad_l = text_length - temp_l - - temp_mask = torch.tril(torch.ones(size=(temp_l + 1, temp_l + 1))) - - image_mask = torch.zeros(size=(temp_l + 1, img_length)) - temp_mask = torch.cat([temp_mask, image_mask], dim=-1) - - image_mask = torch.ones(size=(img_length, temp_l + img_length + 1)) - temp_mask = torch.cat([temp_mask, image_mask], dim=0) - - if pad_l > 0: - pad_mask = torch.zeros(size=(temp_l + 1 + img_length, pad_l)) - temp_mask = torch.cat([pad_mask, temp_mask], dim=-1) - - pad_mask = torch.ones(size=(pad_l, seq_len)) - temp_mask = torch.cat([pad_mask, temp_mask], dim=0) - - true_img_length = num_tokens_for_output_images[inx] - pad_img_length = img_length - true_img_length - if pad_img_length > 0: - temp_mask[:, -pad_img_length:] = 0 - temp_padding_imgs = torch.zeros(size=(1, pad_img_length, self.hidden_size)) - else: - temp_padding_imgs = None - - extended_mask.append(temp_mask.unsqueeze(0)) - padding_images.append(temp_padding_imgs) - inx += 1 - return torch.cat(extended_mask, dim=0), padding_images - - def adjust_attention_for_input_images(self, attention_mask, image_sizes): - for b_inx in image_sizes.keys(): - for start_inx, end_inx in image_sizes[b_inx]: - attention_mask[b_inx][start_inx:end_inx, start_inx:end_inx] = 1 - - return attention_mask - - def pad_input_ids(self, input_ids, image_sizes): - max_l = max([len(x) for x in input_ids]) - padded_ids = [] - attention_mask = [] - - for i in range(len(input_ids)): - temp_ids = input_ids[i] - temp_l = len(temp_ids) - pad_l = max_l - temp_l - if pad_l == 0: - attention_mask.append([1] * max_l) - padded_ids.append(temp_ids) - else: - attention_mask.append([0] * pad_l + [1] * temp_l) - padded_ids.append([self.pad_token_id] * pad_l + temp_ids) - - if i in image_sizes: - new_inx = [] - for old_inx in image_sizes[i]: - new_inx.append([x + pad_l for x in old_inx]) - image_sizes[i] = new_inx - - return torch.LongTensor(padded_ids), torch.LongTensor(attention_mask), image_sizes - - def process_mllm_input(self, mllm_inputs, target_img_size): - num_tokens_for_output_images = [] - for img_size in target_img_size: - num_tokens_for_output_images.append(img_size[0] * img_size[1] // 16 // 16) - - pixel_values, image_sizes = [], {} - b_inx = 0 - for x in mllm_inputs: - if x["pixel_values"] is not None: - pixel_values.extend(x["pixel_values"]) - for size in x["image_sizes"]: - if b_inx not in image_sizes: - image_sizes[b_inx] = [size] - else: - image_sizes[b_inx].append(size) - b_inx += 1 - pixel_values = [x.unsqueeze(0) for x in pixel_values] - - input_ids = [x["input_ids"] for x in mllm_inputs] - padded_input_ids, attention_mask, image_sizes = self.pad_input_ids(input_ids, image_sizes) - position_ids = self.create_position(attention_mask, num_tokens_for_output_images) - attention_mask, padding_images = self.create_mask(attention_mask, num_tokens_for_output_images) - attention_mask = self.adjust_attention_for_input_images(attention_mask, image_sizes) - - return padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes - - def __call__(self, features): - mllm_inputs = [f[0] for f in features] - cfg_mllm_inputs = [f[1] for f in features] - img_cfg_mllm_input = [f[2] for f in features] - target_img_size = [f[3] for f in features] - - if img_cfg_mllm_input[0] is not None: - mllm_inputs = mllm_inputs + cfg_mllm_inputs + img_cfg_mllm_input - target_img_size = target_img_size + target_img_size + target_img_size - else: - mllm_inputs = mllm_inputs + cfg_mllm_inputs - target_img_size = target_img_size + target_img_size - - ( - all_padded_input_ids, - all_position_ids, - all_attention_mask, - all_padding_images, - all_pixel_values, - all_image_sizes, - ) = self.process_mllm_input(mllm_inputs, target_img_size) - - data = { - "input_ids": all_padded_input_ids, - "attention_mask": all_attention_mask, - "position_ids": all_position_ids, - "input_pixel_values": all_pixel_values, - "input_image_sizes": all_image_sizes, - } - return data diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/onnx_utils.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/onnx_utils.py deleted file mode 100644 index 74e9f0b97800ce798f14941f6b2c0fb4c2e532fa..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/onnx_utils.py +++ /dev/null @@ -1,230 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# Copyright (c) 2022, NVIDIA CORPORATION. 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 os -import shutil -from pathlib import Path -from typing import Optional, Union - -import numpy as np -from huggingface_hub import hf_hub_download -from huggingface_hub.utils import validate_hf_hub_args - -from ..utils import ONNX_EXTERNAL_WEIGHTS_NAME, ONNX_WEIGHTS_NAME, is_onnx_available, logging - - -if is_onnx_available(): - import onnxruntime as ort - - -logger = logging.get_logger(__name__) - -ORT_TO_NP_TYPE = { - "tensor(bool)": np.bool_, - "tensor(int8)": np.int8, - "tensor(uint8)": np.uint8, - "tensor(int16)": np.int16, - "tensor(uint16)": np.uint16, - "tensor(int32)": np.int32, - "tensor(uint32)": np.uint32, - "tensor(int64)": np.int64, - "tensor(uint64)": np.uint64, - "tensor(float16)": np.float16, - "tensor(float)": np.float32, - "tensor(double)": np.float64, -} - - -class OnnxRuntimeModel: - def __init__(self, model=None, **kwargs): - logger.info("`diffusers.OnnxRuntimeModel` is experimental and might change in the future.") - self.model = model - self.model_save_dir = kwargs.get("model_save_dir", None) - self.latest_model_name = kwargs.get("latest_model_name", ONNX_WEIGHTS_NAME) - - def __call__(self, **kwargs): - inputs = {k: np.array(v) for k, v in kwargs.items()} - return self.model.run(None, inputs) - - @staticmethod - def load_model(path: Union[str, Path], provider=None, sess_options=None, provider_options=None): - """ - Loads an ONNX Inference session with an ExecutionProvider. Default provider is `CPUExecutionProvider` - - Arguments: - path (`str` or `Path`): - Directory from which to load - provider(`str`, *optional*): - Onnxruntime execution provider to use for loading the model, defaults to `CPUExecutionProvider` - """ - if provider is None: - logger.info("No onnxruntime provider specified, using CPUExecutionProvider") - provider = "CPUExecutionProvider" - - if provider_options is None: - provider_options = [] - elif not isinstance(provider_options, list): - provider_options = [provider_options] - - return ort.InferenceSession( - 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): - """ - 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 - latest_model_name. - - Arguments: - save_directory (`str` or `Path`): - Directory where to save the model file. - file_name(`str`, *optional*): - Overwrites the default model file name from `"model.onnx"` to `file_name`. This allows you to save the - model with a different name. - """ - model_file_name = file_name if file_name is not None else ONNX_WEIGHTS_NAME - - src_path = self.model_save_dir.joinpath(self.latest_model_name) - dst_path = Path(save_directory).joinpath(model_file_name) - try: - shutil.copyfile(src_path, dst_path) - except shutil.SameFileError: - pass - - # copy external weights (for models >2GB) - src_path = self.model_save_dir.joinpath(ONNX_EXTERNAL_WEIGHTS_NAME) - if src_path.exists(): - dst_path = Path(save_directory).joinpath(ONNX_EXTERNAL_WEIGHTS_NAME) - try: - shutil.copyfile(src_path, dst_path) - except shutil.SameFileError: - pass - - def save_pretrained( - self, - save_directory: Union[str, os.PathLike], - **kwargs, - ): - """ - Save a model to a directory, so that it can be re-loaded using the [`~OnnxModel.from_pretrained`] class - method.: - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to which to save. Will be created if it doesn't exist. - """ - if os.path.isfile(save_directory): - logger.error(f"Provided path ({save_directory}) should be a directory, not a file") - return - - os.makedirs(save_directory, exist_ok=True) - - # saving model weights/files - self._save_pretrained(save_directory, **kwargs) - - @classmethod - @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, - force_download: bool = False, - cache_dir: Optional[str] = None, - file_name: Optional[str] = None, - provider: Optional[str] = None, - sess_options: Optional["ort.SessionOptions"] = None, - **kwargs, - ): - """ - Load a model from a directory or the HF Hub. - - Arguments: - model_id (`str` or `Path`): - Directory from which to load - token (`str` or `bool`): - 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*): - 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. - file_name(`str`): - Overwrites the default model file name from `"model.onnx"` to `file_name`. This allows you to load - different model files from the same repository or directory. - provider(`str`): - The ONNX runtime provider, e.g. `CPUExecutionProvider` or `CUDAExecutionProvider`. - kwargs (`Dict`, *optional*): - kwargs will be passed to the model during initialization - """ - model_file_name = file_name if file_name is not None else ONNX_WEIGHTS_NAME - # load model from local directory - if os.path.isdir(model_id): - model = OnnxRuntimeModel.load_model( - Path(model_id, model_file_name).as_posix(), - provider=provider, - sess_options=sess_options, - provider_options=kwargs.pop("provider_options"), - ) - kwargs["model_save_dir"] = Path(model_id) - # load model from hub - else: - # download model - model_cache_path = hf_hub_download( - repo_id=model_id, - filename=model_file_name, - token=token, - revision=revision, - cache_dir=cache_dir, - force_download=force_download, - ) - kwargs["model_save_dir"] = Path(model_cache_path).parent - kwargs["latest_model_name"] = Path(model_cache_path).name - model = OnnxRuntimeModel.load_model( - model_cache_path, - provider=provider, - sess_options=sess_options, - provider_options=kwargs.pop("provider_options"), - ) - return cls(model=model, **kwargs) - - @classmethod - @validate_hf_hub_args - def from_pretrained( - cls, - model_id: Union[str, Path], - force_download: bool = True, - token: Optional[str] = None, - cache_dir: Optional[str] = None, - **model_kwargs, - ): - revision = None - if len(str(model_id).split("@")) == 2: - model_id, revision = model_id.split("@") - - return cls._from_pretrained( - model_id=model_id, - revision=revision, - cache_dir=cache_dir, - force_download=force_download, - token=token, - **model_kwargs, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__init__.py deleted file mode 100644 index 176efe3adef6a792ba2e7fef194e2a0a9475fbfd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__init__.py +++ /dev/null @@ -1,80 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - get_objects_from_module, - is_flax_available, - 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_pag_controlnet_sd"] = ["StableDiffusionControlNetPAGPipeline"] - _import_structure["pipeline_pag_controlnet_sd_inpaint"] = ["StableDiffusionControlNetPAGInpaintPipeline"] - _import_structure["pipeline_pag_controlnet_sd_xl"] = ["StableDiffusionXLControlNetPAGPipeline"] - _import_structure["pipeline_pag_controlnet_sd_xl_img2img"] = ["StableDiffusionXLControlNetPAGImg2ImgPipeline"] - _import_structure["pipeline_pag_hunyuandit"] = ["HunyuanDiTPAGPipeline"] - _import_structure["pipeline_pag_kolors"] = ["KolorsPAGPipeline"] - _import_structure["pipeline_pag_pixart_sigma"] = ["PixArtSigmaPAGPipeline"] - _import_structure["pipeline_pag_sana"] = ["SanaPAGPipeline"] - _import_structure["pipeline_pag_sd"] = ["StableDiffusionPAGPipeline"] - _import_structure["pipeline_pag_sd_3"] = ["StableDiffusion3PAGPipeline"] - _import_structure["pipeline_pag_sd_3_img2img"] = ["StableDiffusion3PAGImg2ImgPipeline"] - _import_structure["pipeline_pag_sd_animatediff"] = ["AnimateDiffPAGPipeline"] - _import_structure["pipeline_pag_sd_img2img"] = ["StableDiffusionPAGImg2ImgPipeline"] - _import_structure["pipeline_pag_sd_inpaint"] = ["StableDiffusionPAGInpaintPipeline"] - - _import_structure["pipeline_pag_sd_xl"] = ["StableDiffusionXLPAGPipeline"] - _import_structure["pipeline_pag_sd_xl_img2img"] = ["StableDiffusionXLPAGImg2ImgPipeline"] - _import_structure["pipeline_pag_sd_xl_inpaint"] = ["StableDiffusionXLPAGInpaintPipeline"] - -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_pag_controlnet_sd import StableDiffusionControlNetPAGPipeline - from .pipeline_pag_controlnet_sd_inpaint import StableDiffusionControlNetPAGInpaintPipeline - from .pipeline_pag_controlnet_sd_xl import StableDiffusionXLControlNetPAGPipeline - from .pipeline_pag_controlnet_sd_xl_img2img import StableDiffusionXLControlNetPAGImg2ImgPipeline - from .pipeline_pag_hunyuandit import HunyuanDiTPAGPipeline - from .pipeline_pag_kolors import KolorsPAGPipeline - from .pipeline_pag_pixart_sigma import PixArtSigmaPAGPipeline - from .pipeline_pag_sana import SanaPAGPipeline - from .pipeline_pag_sd import StableDiffusionPAGPipeline - from .pipeline_pag_sd_3 import StableDiffusion3PAGPipeline - from .pipeline_pag_sd_3_img2img import StableDiffusion3PAGImg2ImgPipeline - from .pipeline_pag_sd_animatediff import AnimateDiffPAGPipeline - from .pipeline_pag_sd_img2img import StableDiffusionPAGImg2ImgPipeline - from .pipeline_pag_sd_inpaint import StableDiffusionPAGInpaintPipeline - from .pipeline_pag_sd_xl import StableDiffusionXLPAGPipeline - from .pipeline_pag_sd_xl_img2img import StableDiffusionXLPAGImg2ImgPipeline - from .pipeline_pag_sd_xl_inpaint import StableDiffusionXLPAGInpaintPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 610a71eeef95237a72788cb5399ee49154287434..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pag_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pag_utils.cpython-310.pyc deleted file mode 100644 index bfd80ab9855093221f6396ce359d0d05e4c057dc..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pag_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd.cpython-310.pyc deleted file mode 100644 index f1fd0b8dafa06ebd9060872585c2a69504be89ad..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd_inpaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd_inpaint.cpython-310.pyc deleted file mode 100644 index 5b34079dd36b1c1ca375ee37e6fd628b93da1e4b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd_inpaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd_xl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd_xl.cpython-310.pyc deleted file mode 100644 index 3d0874c82c053eb5a4117f647e34ebca519ba1b5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd_xl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd_xl_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd_xl_img2img.cpython-310.pyc deleted file mode 100644 index dfdce51c2f9a454e993d499e48ac138731a7c6e9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_controlnet_sd_xl_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_hunyuandit.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_hunyuandit.cpython-310.pyc deleted file mode 100644 index ebb9b5dd85b50957f77568dd1b1873d0737e92dd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_hunyuandit.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_kolors.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_kolors.cpython-310.pyc deleted file mode 100644 index ad470e9d120653a2fd0483782873c327dd37cc47..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_kolors.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_pixart_sigma.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_pixart_sigma.cpython-310.pyc deleted file mode 100644 index 0c8f4c136b1785e9628d6b7e6818e034b8c36323..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_pixart_sigma.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sana.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sana.cpython-310.pyc deleted file mode 100644 index 9c599ae718bfe5f5107f39120cc3da7e291090dc..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sana.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd.cpython-310.pyc deleted file mode 100644 index ebd619be8e7838f0bcfb90a3f7657cee1fada96b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_3.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_3.cpython-310.pyc deleted file mode 100644 index fb7c953ce25651a00551a7e041e5018bc2e04d58..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_3.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_3_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_3_img2img.cpython-310.pyc deleted file mode 100644 index ae3c4d1aeecccb92dc459b96944942d87f911614..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_3_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_animatediff.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_animatediff.cpython-310.pyc deleted file mode 100644 index 50f8b5a221cbb00f54eddc9fe16666631ac4e9cf..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_animatediff.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_img2img.cpython-310.pyc deleted file mode 100644 index 77b9e98571f20d058e142fbe83e325059fe55aca..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_inpaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_inpaint.cpython-310.pyc deleted file mode 100644 index ad811ce44df0facb081a463d14836b8007b70ddd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_inpaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_xl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_xl.cpython-310.pyc deleted file mode 100644 index d3cdd1ca12fe8c0b753430af69789044154b87dd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_xl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_xl_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_xl_img2img.cpython-310.pyc deleted file mode 100644 index 2e6d1a15d7eac6d64730f5f9cf73360c0640088f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_xl_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_xl_inpaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_xl_inpaint.cpython-310.pyc deleted file mode 100644 index 915df3394b9f843bcb73cd38977adef5eb06632b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/__pycache__/pipeline_pag_sd_xl_inpaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pag_utils.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pag_utils.py deleted file mode 100644 index 8a56961f321c9901b53392bfc83562cd77a52cc9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pag_utils.py +++ /dev/null @@ -1,243 +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 re -from typing import Dict, List, Tuple, Union - -import torch -import torch.nn as nn - -from ...models.attention_processor import ( - Attention, - AttentionProcessor, - PAGCFGIdentitySelfAttnProcessor2_0, - PAGIdentitySelfAttnProcessor2_0, -) -from ...utils import logging - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class PAGMixin: - r"""Mixin class for [Pertubed Attention Guidance](https://huggingface.co/papers/2403.17377v1).""" - - def _set_pag_attn_processor(self, pag_applied_layers, do_classifier_free_guidance): - r""" - Set the attention processor for the PAG layers. - """ - pag_attn_processors = self._pag_attn_processors - if pag_attn_processors is None: - raise ValueError( - "No PAG attention processors have been set. Set the attention processors by calling `set_pag_applied_layers` and passing the relevant parameters." - ) - - pag_attn_proc = pag_attn_processors[0] if do_classifier_free_guidance else pag_attn_processors[1] - - if hasattr(self, "unet"): - model: nn.Module = self.unet - else: - model: nn.Module = self.transformer - - def is_self_attn(module: nn.Module) -> bool: - r""" - Check if the module is self-attention module based on its name. - """ - return isinstance(module, Attention) and not module.is_cross_attention - - def is_fake_integral_match(layer_id, name): - layer_id = layer_id.split(".")[-1] - name = name.split(".")[-1] - return layer_id.isnumeric() and name.isnumeric() and layer_id == name - - for layer_id in pag_applied_layers: - # for each PAG layer input, we find corresponding self-attention layers in the unet model - target_modules = [] - - for name, module in model.named_modules(): - # Identify the following simple cases: - # (1) Self Attention layer existing - # (2) Whether the module name matches pag layer id even partially - # (3) Make sure it's not a fake integral match if the layer_id ends with a number - # For example, blocks.1, blocks.10 should be differentiable if layer_id="blocks.1" - if ( - is_self_attn(module) - and re.search(layer_id, name) is not None - and not is_fake_integral_match(layer_id, name) - ): - logger.debug(f"Applying PAG to layer: {name}") - target_modules.append(module) - - if len(target_modules) == 0: - raise ValueError(f"Cannot find PAG layer to set attention processor for: {layer_id}") - - for module in target_modules: - module.processor = pag_attn_proc - - def _get_pag_scale(self, t): - r""" - Get the scale factor for the perturbed attention guidance at timestep `t`. - """ - - if self.do_pag_adaptive_scaling: - signal_scale = self.pag_scale - self.pag_adaptive_scale * (1000 - t) - if signal_scale < 0: - signal_scale = 0 - return signal_scale - else: - return self.pag_scale - - def _apply_perturbed_attention_guidance( - self, noise_pred, do_classifier_free_guidance, guidance_scale, t, return_pred_text=False - ): - r""" - Apply perturbed attention guidance to the noise prediction. - - Args: - noise_pred (torch.Tensor): The noise prediction tensor. - do_classifier_free_guidance (bool): Whether to apply classifier-free guidance. - guidance_scale (float): The scale factor for the guidance term. - t (int): The current time step. - 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 - perturbed attention guidance and the text noise prediction. - """ - pag_scale = self._get_pag_scale(t) - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text, noise_pred_perturb = noise_pred.chunk(3) - noise_pred = ( - noise_pred_uncond - + guidance_scale * (noise_pred_text - noise_pred_uncond) - + pag_scale * (noise_pred_text - noise_pred_perturb) - ) - else: - noise_pred_text, noise_pred_perturb = noise_pred.chunk(2) - noise_pred = noise_pred_text + pag_scale * (noise_pred_text - noise_pred_perturb) - if return_pred_text: - return noise_pred, noise_pred_text - return noise_pred - - def _prepare_perturbed_attention_guidance(self, cond, uncond, do_classifier_free_guidance): - """ - Prepares the perturbed attention guidance for the PAG model. - - Args: - cond (torch.Tensor): The conditional input tensor. - uncond (torch.Tensor): The unconditional input tensor. - do_classifier_free_guidance (bool): Flag indicating whether to perform classifier-free guidance. - - Returns: - torch.Tensor: The prepared perturbed attention guidance tensor. - """ - - cond = torch.cat([cond] * 2, dim=0) - - if do_classifier_free_guidance: - cond = torch.cat([uncond, cond], dim=0) - return cond - - def set_pag_applied_layers( - self, - pag_applied_layers: Union[str, List[str]], - pag_attn_processors: Tuple[AttentionProcessor, AttentionProcessor] = ( - PAGCFGIdentitySelfAttnProcessor2_0(), - PAGIdentitySelfAttnProcessor2_0(), - ), - ): - r""" - Set the self-attention layers to apply PAG. Raise ValueError if the input is invalid. - - Args: - 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}" - - Multiple layers as a list - ["blocks.{layers_index_1}", "blocks.{layer_index_2}", ...] - - 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(), - 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). - """ - - if not hasattr(self, "_pag_attn_processors"): - self._pag_attn_processors = None - - if not isinstance(pag_applied_layers, list): - pag_applied_layers = [pag_applied_layers] - if pag_attn_processors is not None: - if not isinstance(pag_attn_processors, tuple) or len(pag_attn_processors) != 2: - raise ValueError("Expected a tuple of two attention processors") - - for i in range(len(pag_applied_layers)): - if not isinstance(pag_applied_layers[i], str): - raise ValueError( - f"Expected either a string or a list of string but got type {type(pag_applied_layers[i])}" - ) - - self.pag_applied_layers = pag_applied_layers - self._pag_attn_processors = pag_attn_processors - - @property - def pag_scale(self) -> float: - r"""Get the scale factor for the perturbed attention guidance.""" - return self._pag_scale - - @property - def pag_adaptive_scale(self) -> float: - r"""Get the adaptive scale factor for the perturbed attention guidance.""" - return self._pag_adaptive_scale - - @property - def do_pag_adaptive_scaling(self) -> bool: - r"""Check if the adaptive scaling is enabled for the perturbed attention guidance.""" - return self._pag_adaptive_scale > 0 and self._pag_scale > 0 and len(self.pag_applied_layers) > 0 - - @property - def do_perturbed_attention_guidance(self) -> bool: - r"""Check if the perturbed attention guidance is enabled.""" - return self._pag_scale > 0 and len(self.pag_applied_layers) > 0 - - @property - 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 - with the key as the name of the layer. - """ - - if self._pag_attn_processors is None: - return {} - - valid_attn_processors = {x.__class__ for x in self._pag_attn_processors} - - processors = {} - # We could have iterated through the self.components.items() and checked if a component is - # `ModelMixin` subclassed but that can include a VAE too. - if hasattr(self, "unet"): - denoiser_module = self.unet - elif hasattr(self, "transformer"): - denoiser_module = self.transformer - else: - raise ValueError("No denoiser module found.") - - for name, proc in denoiser_module.attn_processors.items(): - if proc.__class__ in valid_attn_processors: - processors[name] = proc - - return processors diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py deleted file mode 100644 index de66871922e6bd44b8402eb556c3d0f4768609af..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py +++ /dev/null @@ -1,1343 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ControlNetModel, ImageProjection, MultiControlNetModel, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import 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 empty_device_cache, is_compiled_module, is_torch_version, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion.pipeline_output import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker -from .pag_utils import PAGMixin - - -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 - >>> # !pip install opencv-python transformers accelerate - >>> from diffusers import AutoPipelineForText2Image, ControlNetModel, UniPCMultistepScheduler - >>> from diffusers.utils import load_image - >>> import numpy as np - >>> import torch - - >>> import cv2 - >>> from PIL import Image - - >>> # download an image - >>> image = load_image( - ... "https://hf.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/hf-logo.png" - ... ) - >>> image = np.array(image) - - >>> # get canny image - >>> image = cv2.Canny(image, 100, 200) - >>> image = image[:, :, None] - >>> image = np.concatenate([image, image, image], axis=2) - >>> canny_image = Image.fromarray(image) - - >>> # load control net and stable diffusion v1-5 - >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16) - >>> pipe = AutoPipelineForText2Image.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16, enable_pag=True - ... ) - - >>> # speed up diffusion process with faster scheduler and memory optimization - >>> # remove following line if xformers is not installed - >>> pipe.enable_xformers_memory_efficient_attention() - - >>> pipe.enable_model_cpu_offload() - - >>> # generate image - >>> generator = torch.manual_seed(0) - >>> image = pipe( - ... "aerial view, a futuristic research complex in a bright foggy jungle, hard lighting", - ... guidance_scale=7.5, - ... generator=generator, - ... image=canny_image, - ... pag_scale=10, - ... ).images[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, - **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 StableDiffusionControlNetPAGPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - IPAdapterMixin, - FromSingleFileMixin, - PAGMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion with ControlNet guidance. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - 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. - 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - controlnet: Union[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", - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - if isinstance(controlnet, (list, tuple)): - controlnet = MultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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, do_convert_rgb=True) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - self.set_pag_applied_layers(pag_applied_layers) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - image, - negative_prompt=None, - prompt_embeds=None, - negative_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, - ): - 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}." - ) - - # Check `image` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - self.check_image(image, prompt, prompt_embeds) - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if not isinstance(image, list): - raise TypeError("For multiple controlnets: `image` must be type `list`") - - # When `image` is a nested list: - # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]]) - elif any(isinstance(i, list) for i in image): - transposed_image = [list(t) for t in zip(*image)] - if len(transposed_image) != len(self.controlnet.nets): - raise ValueError( - f"For multiple controlnets: if you pass`image` as a list of list, each sublist must have the same length as the number of controlnets, but the sublists in `image` got {len(transposed_image)} images and {len(self.controlnet.nets)} ControlNets." - ) - for image_ in transposed_image: - self.check_image(image_, prompt, prompt_embeds) - 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." - ) - else: - for image_ in image: - self.check_image(image_, prompt, prompt_embeds) - else: - assert False - - # Check `controlnet_conditioning_scale` - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if isinstance(controlnet_conditioning_scale, list): - if any(isinstance(i, list) for i in controlnet_conditioning_scale): - raise ValueError( - "A single batch of varying conditioning scale settings (e.g. [[1.0, 0.5], [0.2, 0.8]]) is not supported at the moment. " - "The conditioning scale must be fixed across the batch." - ) - elif 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" - ) - else: - assert False - - if not isinstance(control_guidance_start, (tuple, list)): - control_guidance_start = [control_guidance_start] - - if not isinstance(control_guidance_end, (tuple, list)): - control_guidance_end = [control_guidance_end] - - 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(self.controlnet, MultiControlNetModel): - 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.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.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}" - ) - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.prepare_image - def prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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"], - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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]]`): - 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. When `prompt` is a list, and if a list of images is passed for a single - ControlNet, each will be paired with each prompt in the `prompt` list. This also applies to multiple - ControlNets, where a list of image lists can be passed to batch for each prompt and each ControlNet. - 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. - 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*): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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): - 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): - The percentage of total steps at which the ControlNet starts applying. - 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 - the output of the pre-final layer will be used for computing the prompt embeddings. - 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. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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 isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - image, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 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 - - if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): - controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) - - global_pool_conditions = ( - controlnet.config.global_pool_conditions - if isinstance(controlnet, ControlNetModel) - else controlnet.nets[0].config.global_pool_conditions - ) - guess_mode = guess_mode or global_pool_conditions - - # 3. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - # 4. Prepare image - if isinstance(controlnet, ControlNetModel): - image = self.prepare_image( - image=image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - height, width = image.shape[-2:] - elif isinstance(controlnet, MultiControlNetModel): - images = [] - - # Nested lists as ControlNet condition - if isinstance(image[0], list): - # Transpose the nested image list - image = [list(t) for t in zip(*image)] - - for image_ in image: - image_ = self.prepare_image( - image=image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=guess_mode, - ) - - images.append(image_) - - image = images - height, width = image[0].shape[-2:] - else: - assert False - - # 5. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - self._num_timesteps = len(timesteps) - - # 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, - ) - - # 6.5 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) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Add image embeds for IP-Adapter - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - for i, image_embeds in enumerate(ip_adapter_image_embeds): - negative_image_embeds = None - if self.do_classifier_free_guidance: - negative_image_embeds, image_embeds = image_embeds.chunk(2) - - if self.do_perturbed_attention_guidance: - image_embeds = self._prepare_perturbed_attention_guidance( - image_embeds, negative_image_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - image_embeds = image_embeds.to(device) - ip_adapter_image_embeds[i] = image_embeds - - added_cond_kwargs = ( - {"image_embeds": ip_adapter_image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - controlnet_prompt_embeds = prompt_embeds - - # 7.2 Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) - - images = image if isinstance(image, list) else [image] - for i, single_image in enumerate(images): - if self.do_classifier_free_guidance: - single_image = single_image.chunk(2)[0] - - if self.do_perturbed_attention_guidance: - single_image = self._prepare_perturbed_attention_guidance( - single_image, single_image, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - single_image = torch.cat([single_image] * 2) - single_image = single_image.to(device) - images[i] = single_image - - image = images if isinstance(image, list) else images[0] - - # 8. Denoising loop - if self.do_perturbed_attention_guidance: - original_attn_proc = self.unet.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - is_unet_compiled = is_compiled_module(self.unet) - is_controlnet_compiled = is_compiled_module(self.controlnet) - is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1") - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # Relevant thread: - # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428 - if ( - torch.cuda.is_available() - and (is_unet_compiled and is_controlnet_compiled) - and is_torch_higher_equal_2_1 - ): - torch._inductor.cudagraph_mark_step_begin() - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # controlnet(s) inference - control_model_input = latent_model_input - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=image, - conditioning_scale=cond_scale, - guess_mode=guess_mode, - return_dict=False, - ) - - if guess_mode and self.do_classifier_free_guidance: - # Inferred ControlNet only for the conditional batch. - # To apply the output of ControlNet to both the unconditional and conditional batches, - # add 0 to the unconditional batch to keep it unchanged. - down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] - mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t - ) - elif 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) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - - # 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 we do sequential model offloading, let's offload unet and controlnet - # manually for max memory savings - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.unet.to("cpu") - self.controlnet.to("cpu") - empty_device_cache() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 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 self.do_perturbed_attention_guidance: - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py deleted file mode 100644 index 4c02b3dd6dc77455926dd04804c5d8c5b4ca79fe..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py +++ /dev/null @@ -1,1554 +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. - -# 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 - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ControlNetModel, ImageProjection, MultiControlNetModel, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import 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 empty_device_cache, is_compiled_module, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker -from .pag_utils import PAGMixin - - -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 - >>> # !pip install transformers accelerate - >>> import cv2 - >>> from diffusers import AutoPipelineForInpainting, ControlNetModel, DDIMScheduler - >>> from diffusers.utils import load_image - >>> import numpy as np - >>> from PIL import Image - >>> import torch - - >>> init_image = load_image( - ... "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy.png" - ... ) - >>> init_image = init_image.resize((512, 512)) - - >>> generator = torch.Generator(device="cpu").manual_seed(1) - - >>> mask_image = load_image( - ... "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy_mask.png" - ... ) - >>> mask_image = mask_image.resize((512, 512)) - - - >>> def make_canny_condition(image): - ... image = np.array(image) - ... image = cv2.Canny(image, 100, 200) - ... image = image[:, :, None] - ... image = np.concatenate([image, image, image], axis=2) - ... image = Image.fromarray(image) - ... return image - - - >>> control_image = make_canny_condition(init_image) - - >>> controlnet = ControlNetModel.from_pretrained( - ... "lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16 - ... ) - >>> pipe = AutoPipelineForInpainting.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16, enable_pag=True - ... ) - - >>> pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) - >>> pipe.enable_model_cpu_offload() - - >>> # generate image - >>> image = pipe( - ... "a handsome man with ray-ban sunglasses", - ... num_inference_steps=20, - ... generator=generator, - ... eta=1.0, - ... image=init_image, - ... mask_image=mask_image, - ... control_image=control_image, - ... pag_scale=0.3, - ... ).images[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" -): - 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 StableDiffusionControlNetPAGInpaintPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - IPAdapterMixin, - FromSingleFileMixin, - PAGMixin, -): - r""" - Pipeline for image inpainting using Stable Diffusion with ControlNet guidance. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters - - - - This pipeline can be used with checkpoints that have been specifically fine-tuned for inpainting - ([runwayml/stable-diffusion-inpainting](https://huggingface.co/runwayml/stable-diffusion-inpainting)) as well as - default text-to-image Stable Diffusion checkpoints - ([runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5)). Default text-to-image - Stable Diffusion checkpoints might be preferable for ControlNets that have been fine-tuned on those, such as - [lllyasviel/control_v11p_sd15_inpaint](https://huggingface.co/lllyasviel/control_v11p_sd15_inpaint). - - - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - 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. - 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - controlnet: Union[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", - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - if isinstance(controlnet, (list, tuple)): - controlnet = MultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True - ) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - self.set_pag_applied_layers(pag_applied_layers) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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, - height, - width, - output_type, - negative_prompt=None, - prompt_embeds=None, - negative_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, - padding_mask_crop=None, - ): - if height is not None and height % 8 != 0 or width is not None and width % 8 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 8 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 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 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}.") - - # `prompt` needs more sophisticated handling when there are multiple - # conditionings. - if isinstance(self.controlnet, MultiControlNetModel): - if isinstance(prompt, list): - 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` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - self.check_image(image, prompt, prompt_embeds) - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if not isinstance(image, list): - raise TypeError("For multiple controlnets: `image` must be type `list`") - - # When `image` is a nested list: - # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]]) - elif any(isinstance(i, list) for i in image): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - 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) - else: - assert False - - # Check `controlnet_conditioning_scale` - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if isinstance(controlnet_conditioning_scale, list): - if any(isinstance(i, list) for i in controlnet_conditioning_scale): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - elif 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" - ) - else: - assert False - - 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(self.controlnet, MultiControlNetModel): - 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.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.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}" - ) - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet_inpaint.StableDiffusionControlNetInpaintPipeline.prepare_control_image - def prepare_control_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - crops_coords, - resize_mode, - do_classifier_free_guidance=False, - guess_mode=False, - ): - image = self.control_image_processor.preprocess( - image, height=height, width=width, crops_coords=crops_coords, resize_mode=resize_mode - ).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.StableDiffusionInpaintPipeline.prepare_latents - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=None, - image=None, - timestep=None, - is_strength_max=True, - return_noise=False, - return_image_latents=False, - ): - 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 (image is None or timestep is None) and not is_strength_max: - raise ValueError( - "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise." - "However, either the image or the noise timestep has not been provided." - ) - - if return_image_latents or (latents is None and not is_strength_max): - image = image.to(device=device, dtype=dtype) - - if image.shape[1] == 4: - image_latents = image - else: - image_latents = self._encode_vae_image(image=image, generator=generator) - image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1) - - if latents is None: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # if strength is 1. then initialise the latents to noise, else initial to image + noise - latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep) - # if pure noise then scale the initial latents by the Scheduler's init sigma - latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents - else: - noise = latents.to(device) - latents = noise * self.scheduler.init_noise_sigma - - outputs = (latents,) - - if return_noise: - outputs += (noise,) - - if return_image_latents: - outputs += (image_latents,) - - return outputs - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.StableDiffusionInpaintPipeline.prepare_mask_latents - def prepare_mask_latents( - self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance - ): - # 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 // self.vae_scale_factor, width // self.vae_scale_factor) - ) - mask = mask.to(device=device, dtype=dtype) - - masked_image = masked_image.to(device=device, dtype=dtype) - - if masked_image.shape[1] == 4: - masked_image_latents = masked_image - else: - masked_image_latents = self._encode_vae_image(masked_image, generator=generator) - - # 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) - - mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask - masked_image_latents = ( - torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents - ) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.StableDiffusionInpaintPipeline._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 = self.vae.config.scaling_factor * image_latents - - return image_latents - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @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, - control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = 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, - 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", - 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"], - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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`, 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]`): - `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]]`): - 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. - 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. - padding_mask_crop (`int`, *optional*, defaults to `None`): - The size of margin in the crop to be applied to the image and masking. If `None`, no crop is applied to - image and mask_image. If `padding_mask_crop` is not `None`, it will first find a rectangular region - with the same aspect ration of the image and contains all masked area, and then expand that area based - on `padding_mask_crop`. The image and mask_image will then be cropped based on the expanded area before - resizing to the original image size for inpainting. This is useful when the masked area is small while - the image is large and contain information irrelevant for inpainting, such as background. - 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`. - 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): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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): - 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): - The percentage of total steps at which the ControlNet starts applying. - 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 - the output of the pre-final layer will be used for computing the prompt embeddings. - 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. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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 isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - control_image, - mask_image, - height, - width, - output_type, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - callback_on_step_end_tensor_inputs, - padding_mask_crop, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 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 padding_mask_crop is not None: - height, width = self.image_processor.get_default_height_width(image, height, width) - 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" - - device = self._execution_device - - if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): - controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) - - # 3. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - - # 4. Prepare control image - if isinstance(controlnet, ControlNetModel): - control_image = self.prepare_control_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - crops_coords=crops_coords, - resize_mode=resize_mode, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=False, - ) - elif isinstance(controlnet, MultiControlNetModel): - control_images = [] - - for control_image_ in control_image: - control_image_ = self.prepare_control_image( - image=control_image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - crops_coords=crops_coords, - resize_mode=resize_mode, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=False, - ) - - control_images.append(control_image_) - - control_image = control_images - else: - assert False - - # 4.1 Preprocess mask and image - resizes image and mask w.r.t height and width - 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) - - mask = self.mask_processor.preprocess( - mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords - ) - - masked_image = init_image * (mask < 0.5) - _, _, height, width = init_image.shape - - # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps( - num_inference_steps=num_inference_steps, strength=strength, device=device - ) - # at which timestep to set the initial noise (n.b. 50% if strength is 0.5) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise - is_strength_max = strength == 1.0 - self._num_timesteps = len(timesteps) - - # 6. Prepare latent variables - num_channels_latents = self.vae.config.latent_channels - num_channels_unet = self.unet.config.in_channels - return_image_latents = num_channels_unet == 4 - latents_outputs = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - image=init_image, - timestep=latent_timestep, - is_strength_max=is_strength_max, - return_noise=True, - return_image_latents=return_image_latents, - ) - - if return_image_latents: - latents, noise, image_latents = latents_outputs - else: - latents, noise = latents_outputs - - # 7. Prepare mask latent variables - mask, masked_image_latents = self.prepare_mask_latents( - mask, - masked_image, - batch_size * num_images_per_prompt, - height, - width, - prompt_embeds.dtype, - device, - generator, - self.do_classifier_free_guidance, - ) - - # 7.1 Check that sizes of mask, masked image and latents match - if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting - num_channels_mask = mask.shape[1] - num_channels_masked_image = masked_image_latents.shape[1] - if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" - " `pipeline.unet` or your `mask_image` or `image` input." - ) - elif num_channels_unet != 4: - raise ValueError( - f"The unet {self.unet.__class__} should have either 4 or 9 input channels, not {self.unet.config.in_channels}." - ) - - # 7.2 Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.3 Prepare embeddings - # ip-adapter - if ip_adapter_image_embeds is not None: - for i, image_embeds in enumerate(ip_adapter_image_embeds): - negative_image_embeds = None - if self.do_classifier_free_guidance: - negative_image_embeds, image_embeds = image_embeds.chunk(2) - - if self.do_perturbed_attention_guidance: - image_embeds = self._prepare_perturbed_attention_guidance( - image_embeds, negative_image_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - image_embeds = image_embeds.to(device) - ip_adapter_image_embeds[i] = image_embeds - - added_cond_kwargs = ( - {"image_embeds": ip_adapter_image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - # control image - control_images = control_image if isinstance(control_image, list) else [control_image] - for i, single_control_image in enumerate(control_images): - if self.do_classifier_free_guidance: - single_control_image = single_control_image.chunk(2)[0] - - if self.do_perturbed_attention_guidance: - single_control_image = self._prepare_perturbed_attention_guidance( - single_control_image, single_control_image, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - single_control_image = torch.cat([single_control_image] * 2) - single_control_image = single_control_image.to(device) - control_images[i] = single_control_image - - control_image = control_images if isinstance(control_image, list) else control_images[0] - controlnet_prompt_embeds = prompt_embeds - - # 7.4 Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) - - # 7.5 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) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - if self.do_perturbed_attention_guidance: - original_attn_proc = self.unet.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # controlnet(s) inference - control_model_input = latent_model_input - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=control_image, - conditioning_scale=cond_scale, - guess_mode=False, - return_dict=False, - ) - - # concat latents, mask, masked_image_latents in the channel dimension - if num_channels_unet == 9: - first_dim_size = latent_model_input.shape[0] - # Ensure mask and masked_image_latents have the right dimensions - if mask.shape[0] < first_dim_size: - repeat_factor = (first_dim_size + mask.shape[0] - 1) // mask.shape[0] - mask = mask.repeat(repeat_factor, 1, 1, 1)[:first_dim_size] - if masked_image_latents.shape[0] < first_dim_size: - repeat_factor = ( - first_dim_size + masked_image_latents.shape[0] - 1 - ) // masked_image_latents.shape[0] - masked_image_latents = masked_image_latents.repeat(repeat_factor, 1, 1, 1)[:first_dim_size] - # Perform the concatenation - latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) - - # Predict noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t - ) - elif self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - if num_channels_unet == 4: - init_latents_proper = image_latents - if self.do_classifier_free_guidance: - init_mask, _ = mask.chunk(2) - else: - init_mask = mask - - if i < len(timesteps) - 1: - noise_timestep = timesteps[i + 1] - init_latents_proper = self.scheduler.add_noise( - init_latents_proper, noise, torch.tensor([noise_timestep]) - ) - - latents = (1 - init_mask) * init_latents_proper + init_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) - - # 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 we do sequential model offloading, let's offload unet and controlnet - # manually for max memory savings - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.unet.to("cpu") - self.controlnet.to("cpu") - empty_device_cache() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 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) - - if padding_mask_crop is not None: - image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] - - # Offload all models - self.maybe_free_model_hooks() - - if self.do_perturbed_attention_guidance: - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py deleted file mode 100644 index 24cbab43c966938cc1961193be323c4ff049a87b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py +++ /dev/null @@ -1,1631 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from diffusers.utils.import_utils import is_invisible_watermark_available - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ControlNetModel, ImageProjection, MultiControlNetModel, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import is_compiled_module, is_torch_version, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput -from .pag_utils import PAGMixin - - -if is_invisible_watermark_available(): - from ..stable_diffusion_xl.watermark import StableDiffusionXLWatermarker - - -from ...utils import is_torch_xla_available - - -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 - >>> # !pip install opencv-python transformers accelerate - >>> from diffusers import AutoPipelineForText2Image, ControlNetModel, AutoencoderKL - >>> from diffusers.utils import load_image - >>> import numpy as np - >>> import torch - - >>> import cv2 - >>> from PIL import Image - - >>> prompt = "aerial view, a futuristic research complex in a bright foggy jungle, hard lighting" - >>> negative_prompt = "low quality, bad quality, sketches" - - >>> # download an image - >>> image = load_image( - ... "https://hf.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/hf-logo.png" - ... ) - - >>> # initialize the models and pipeline - >>> controlnet_conditioning_scale = 0.5 # recommended for good generalization - >>> controlnet = ControlNetModel.from_pretrained( - ... "diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16 - ... ) - >>> vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) - >>> pipe = AutoPipelineForText2Image.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", - ... controlnet=controlnet, - ... vae=vae, - ... torch_dtype=torch.float16, - ... enable_pag=True, - ... ) - >>> pipe.enable_model_cpu_offload() - - >>> # get canny image - >>> image = np.array(image) - >>> image = cv2.Canny(image, 100, 200) - >>> image = image[:, :, None] - >>> image = np.concatenate([image, image, image], axis=2) - >>> canny_image = Image.fromarray(image) - - >>> # generate image - >>> image = pipe( - ... prompt, controlnet_conditioning_scale=controlnet_conditioning_scale, image=canny_image, pag_scale=0.3 - ... ).images[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, - **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 StableDiffusionXLControlNetPAGPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionXLLoraLoaderMixin, - IPAdapterMixin, - FromSingleFileMixin, - PAGMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion XL with ControlNet guidance. - - 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.). - - The pipeline also inherits the following loading methods: - - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings - - [`~loaders.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] for loading LoRA weights - - [`~loaders.StableDiffusionXLLoraLoaderMixin.save_lora_weights`] for saving LoRA weights - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - text_encoder_2 ([`~transformers.CLIPTextModelWithProjection`]): - Second frozen text-encoder - ([laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - tokenizer_2 ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - 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. - 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 should always be set to 0. Also see the config of - `stabilityai/stable-diffusion-xl-base-1-0`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark](https://github.com/ShieldMnt/invisible-watermark/) library to - watermark output images. If not defined, it defaults to `True` if the package is installed; otherwise no - watermarker is used. - """ - - # leave controlnet out on purpose because it iterates with unet - model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "feature_extractor", - "image_encoder", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "negative_add_time_ids", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], - scheduler: KarrasDiffusionSchedulers, - force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = 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" - ): - super().__init__() - - if isinstance(controlnet, (list, tuple)): - controlnet = MultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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, do_convert_rgb=True) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - self.set_pag_applied_layers(pag_applied_layers) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.check_inputs - def check_inputs( - self, - prompt, - prompt_2, - image, - callback_steps, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - negative_pooled_prompt_embeds=None, - controlnet_conditioning_scale=1.0, - control_guidance_start=0.0, - control_guidance_end=1.0, - callback_on_step_end_tensor_inputs=None, - ): - 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_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`." - ) - - # `prompt` needs more sophisticated handling when there are multiple - # conditionings. - if isinstance(self.controlnet, MultiControlNetModel): - if isinstance(prompt, list): - 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` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - self.check_image(image, prompt, prompt_embeds) - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if not isinstance(image, list): - raise TypeError("For multiple controlnets: `image` must be type `list`") - - # When `image` is a nested list: - # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]]) - elif any(isinstance(i, list) for i in image): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - 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) - else: - assert False - - # Check `controlnet_conditioning_scale` - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if isinstance(controlnet_conditioning_scale, list): - if any(isinstance(i, list) for i in controlnet_conditioning_scale): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - elif 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" - ) - else: - assert False - - if not isinstance(control_guidance_start, (tuple, list)): - control_guidance_start = [control_guidance_start] - - if not isinstance(control_guidance_end, (tuple, list)): - control_guidance_end = [control_guidance_end] - - 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(self.controlnet, MultiControlNetModel): - 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.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.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}" - ) - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet.StableDiffusionControlNetPipeline.prepare_image - def prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - 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.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_end(self): - return self._denoising_end - - @property - def num_timesteps(self): - return self._num_timesteps - - @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, - image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = 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, - 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", - 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"], - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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*): - 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]]`): - 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. - height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The height in pixels of the generated image. 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. 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. - 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*): - 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. - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output) - 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*): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - pooled_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated pooled text embeddings. Can be used to easily tweak text inputs (prompt weighting). If - not provided, pooled text embeddings are 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 (prompt - 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*): - 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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): - 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): - The percentage of total steps at which the ControlNet starts applying. - 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)): - 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. - 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. - 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. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned, - otherwise a `tuple` is returned containing the output images. - """ - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - image, - None, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - negative_pooled_prompt_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 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 - - if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): - controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) - - # 3.1 Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt, - prompt_2, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - 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=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - - # 3.2 Encode ip_adapter_image - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - - # 4. Prepare image - if isinstance(controlnet, ControlNetModel): - image = self.prepare_image( - image=image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=False, - ) - height, width = image.shape[-2:] - elif isinstance(controlnet, MultiControlNetModel): - images = [] - - for image_ in image: - image_ = self.prepare_image( - image=image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=False, - ) - - images.append(image_) - - image = images - height, width = image[0].shape[-2:] - else: - assert False - - # 5. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - self._num_timesteps = len(timesteps) - - # 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, - ) - - # 6.1 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) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) - - # 7.2 Prepare added time ids & embeddings - if isinstance(image, list): - original_size = original_size or image[0].shape[-2:] - else: - original_size = original_size or image.shape[-2:] - target_size = target_size or (height, width) - - 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 - - images = image if isinstance(image, list) else [image] - for i, single_image in enumerate(images): - if self.do_classifier_free_guidance: - single_image = single_image.chunk(2)[0] - - if self.do_perturbed_attention_guidance: - single_image = self._prepare_perturbed_attention_guidance( - single_image, single_image, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - single_image = torch.cat([single_image] * 2) - single_image = single_image.to(device) - images[i] = single_image - - image = images if isinstance(image, list) else images[0] - - if ip_adapter_image_embeds is not None: - for i, image_embeds in enumerate(ip_adapter_image_embeds): - negative_image_embeds = None - if self.do_classifier_free_guidance: - negative_image_embeds, image_embeds = image_embeds.chunk(2) - - if self.do_perturbed_attention_guidance: - image_embeds = self._prepare_perturbed_attention_guidance( - image_embeds, negative_image_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - image_embeds = image_embeds.to(device) - ip_adapter_image_embeds[i] = image_embeds - - if self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - add_text_embeds = self._prepare_perturbed_attention_guidance( - add_text_embeds, negative_pooled_prompt_embeds, self.do_classifier_free_guidance - ) - add_time_ids = self._prepare_perturbed_attention_guidance( - add_time_ids, negative_add_time_ids, self.do_classifier_free_guidance - ) - elif 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} - - controlnet_prompt_embeds = prompt_embeds - controlnet_added_cond_kwargs = added_cond_kwargs - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - # 8.1 Apply denoising_end - if ( - self.denoising_end is not None - and isinstance(self.denoising_end, float) - and self.denoising_end > 0 - and self.denoising_end < 1 - ): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - if self.do_perturbed_attention_guidance: - original_attn_proc = self.unet.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - - is_unet_compiled = is_compiled_module(self.unet) - is_controlnet_compiled = is_compiled_module(self.controlnet) - is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1") - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # Relevant thread: - # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428 - if ( - torch.cuda.is_available() - and (is_unet_compiled and is_controlnet_compiled) - and is_torch_higher_equal_2_1 - ): - torch._inductor.cudagraph_mark_step_begin() - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # controlnet(s) inference - control_model_input = latent_model_input - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=image, - conditioning_scale=cond_scale, - guess_mode=False, - added_cond_kwargs=controlnet_added_cond_kwargs, - return_dict=False, - ) - - if ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = ip_adapter_image_embeds - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t - ) - elif self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds - ) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids) - - # 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 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) - - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None - has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None - if has_latents_mean and has_latents_std: - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean - else: - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, 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": - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - image = self.image_processor.postprocess(image, output_type=output_type) - - # Offload all models - self.maybe_free_model_hooks() - - if self.do_perturbed_attention_guidance: - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image,) - - return StableDiffusionXLPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl_img2img.py deleted file mode 100644 index 913a647fae3e4a5da4a04bbe47c189b31a669a25..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl_img2img.py +++ /dev/null @@ -1,1700 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from diffusers.utils.import_utils import is_invisible_watermark_available - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ControlNetModel, ImageProjection, MultiControlNetModel, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import empty_device_cache, is_compiled_module, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput -from .pag_utils import PAGMixin - - -if is_invisible_watermark_available(): - from ..stable_diffusion_xl.watermark import StableDiffusionXLWatermarker - - -from ...utils import is_torch_xla_available - - -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 - >>> # pip install accelerate transformers safetensors diffusers - - >>> import torch - >>> import numpy as np - >>> from PIL import Image - - >>> from transformers import DPTFeatureExtractor, DPTForDepthEstimation - >>> 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( - ... "diffusers/controlnet-depth-sdxl-1.0-small", - ... variant="fp16", - ... use_safetensors="True", - ... torch_dtype=torch.float16, - ... ) - >>> vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) - >>> pipe = StableDiffusionXLControlNetPAGImg2ImgPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", - ... controlnet=controlnet, - ... vae=vae, - ... variant="fp16", - ... use_safetensors=True, - ... torch_dtype=torch.float16, - ... enable_pag=True, - ... ) - >>> pipe.enable_model_cpu_offload() - - - >>> def get_depth_map(image): - ... image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda") - ... with torch.no_grad(), torch.autocast("cuda"): - ... depth_map = depth_estimator(image).predicted_depth - - ... depth_map = torch.nn.functional.interpolate( - ... depth_map.unsqueeze(1), - ... size=(1024, 1024), - ... mode="bicubic", - ... align_corners=False, - ... ) - ... depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True) - ... depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True) - ... depth_map = (depth_map - depth_min) / (depth_max - depth_min) - ... image = torch.cat([depth_map] * 3, dim=1) - ... image = image.permute(0, 2, 3, 1).cpu().numpy()[0] - ... image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8)) - ... return image - - - >>> prompt = "A robot, 4k photo" - >>> image = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" - ... "/kandinsky/cat.png" - ... ).resize((1024, 1024)) - >>> controlnet_conditioning_scale = 0.5 # recommended for good generalization - >>> depth_image = get_depth_map(image) - - >>> images = pipe( - ... prompt, - ... image=image, - ... control_image=depth_image, - ... strength=0.99, - ... num_inference_steps=50, - ... controlnet_conditioning_scale=controlnet_conditioning_scale, - ... ).images - >>> images[0].save(f"robot_cat.png") - ``` -""" - - -# 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" -): - 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 StableDiffusionXLControlNetPAGImg2ImgPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionXLLoraLoaderMixin, - FromSingleFileMixin, - IPAdapterMixin, - PAGMixin, -): - r""" - Pipeline for image-to-image generation using Stable Diffusion XL with ControlNet guidance. - - 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.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 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. - 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. - 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`]. - requires_aesthetics_score (`bool`, *optional*, defaults to `"False"`): - Whether the `unet` requires an `aesthetic_score` condition to be passed during inference. Also see the - config of `stabilityai/stable-diffusion-xl-refiner-1-0`. - 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`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to - watermark output images. If not defined, it will default to True if the package is installed, otherwise no - watermarker will be used. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "feature_extractor", - "image_encoder", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "add_neg_time_ids", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - controlnet: Union[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: CLIPImageProcessor = None, - image_encoder: CLIPVisionModelWithProjection = None, - pag_applied_layers: Union[str, List[str]] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"] - ): - super().__init__() - - if isinstance(controlnet, (list, tuple)): - controlnet = MultiControlNetModel(controlnet) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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, do_convert_rgb=True) - self.control_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False - ) - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - self.register_to_config(requires_aesthetics_score=requires_aesthetics_score) - - self.set_pag_applied_layers(pag_applied_layers) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl_img2img.StableDiffusionXLControlNetImg2ImgPipeline.check_inputs - def check_inputs( - self, - prompt, - prompt_2, - image, - strength, - num_inference_steps, - callback_steps, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - 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, - ): - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - if num_inference_steps is None: - raise ValueError("`num_inference_steps` cannot be None.") - elif not isinstance(num_inference_steps, int) or num_inference_steps <= 0: - raise ValueError( - f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type" - f" {type(num_inference_steps)}." - ) - - 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_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`." - ) - - # `prompt` needs more sophisticated handling when there are multiple - # conditionings. - if isinstance(self.controlnet, MultiControlNetModel): - if isinstance(prompt, list): - 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` - is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( - self.controlnet, torch._dynamo.eval_frame.OptimizedModule - ) - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - self.check_image(image, prompt, prompt_embeds) - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if not isinstance(image, list): - raise TypeError("For multiple controlnets: `image` must be type `list`") - - # When `image` is a nested list: - # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]]) - elif any(isinstance(i, list) for i in image): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - 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) - else: - assert False - - # Check `controlnet_conditioning_scale` - if ( - isinstance(self.controlnet, ControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, ControlNetModel) - ): - if not isinstance(controlnet_conditioning_scale, float): - raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.") - elif ( - isinstance(self.controlnet, MultiControlNetModel) - or is_compiled - and isinstance(self.controlnet._orig_mod, MultiControlNetModel) - ): - if isinstance(controlnet_conditioning_scale, list): - if any(isinstance(i, list) for i in controlnet_conditioning_scale): - raise ValueError("A single batch of multiple conditionings are supported at the moment.") - elif 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" - ) - else: - assert False - - if not isinstance(control_guidance_start, (tuple, list)): - control_guidance_start = [control_guidance_start] - - if not isinstance(control_guidance_end, (tuple, list)): - control_guidance_end = [control_guidance_end] - - 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(self.controlnet, MultiControlNetModel): - 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.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}" - ) - - # Copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.prepare_image - def prepare_control_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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 - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.prepare_latents - def prepare_latents( - self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None, add_noise=True - ): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - latents_mean = latents_std = None - if hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None: - latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1) - if hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None: - latents_std = torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1) - - # Offload text encoder if `enable_model_cpu_offload` was enabled - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.text_encoder_2.to("cpu") - empty_device_cache() - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - # make sure the VAE is in float32 mode, as it overflows in float16 - if self.vae.config.force_upcast: - image = image.float() - self.vae.to(dtype=torch.float32) - - 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." - ) - - elif isinstance(generator, list): - if image.shape[0] < batch_size and batch_size % image.shape[0] == 0: - image = torch.cat([image] * (batch_size // image.shape[0]), dim=0) - elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} " - ) - - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - if self.vae.config.force_upcast: - self.vae.to(dtype) - - init_latents = init_latents.to(dtype) - if latents_mean is not None and latents_std is not None: - latents_mean = latents_mean.to(device=device, dtype=dtype) - latents_std = latents_std.to(device=device, dtype=dtype) - init_latents = (init_latents - latents_mean) * self.vae.config.scaling_factor / latents_std - else: - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - if add_noise: - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - - latents = init_latents - - return latents - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline._get_add_time_ids - def _get_add_time_ids( - self, - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype, - text_encoder_projection_dim=None, - ): - if self.config.requires_aesthetics_score: - add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,)) - add_neg_time_ids = list( - negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,) - ) - else: - add_time_ids = list(original_size + crops_coords_top_left + target_size) - add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_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 - and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model." - ) - elif ( - expected_add_embed_dim < passed_add_embed_dim - and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model." - ) - elif 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) - add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype) - - return add_time_ids, add_neg_time_ids - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @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, - image: PipelineImageInput = None, - control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[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, - 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"], - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ): - 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 - 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]]`): - 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 - and/or width are passed, `image` is resized according to them. 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. - height (`int`, *optional*, defaults to the size of control_image): - The height in pixels of the generated image. 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 the size of control_image): - The width in pixels of the generated image. 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. - strength (`float`, *optional*, defaults to 0.8): - 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`. - 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`). - 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. - 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*): - 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 ge 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. - 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)`. 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. - 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. - cross_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_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): - The percentage of total steps at which the controlnet starts applying. - 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)): - 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. - aesthetic_score (`float`, *optional*, defaults to 6.0): - Used to simulate an aesthetic score of the generated image by influencing the positive text condition. - 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_aesthetic_score (`float`, *optional*, defaults to 2.5): - 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). Can be used to - simulate an aesthetic score of the generated image by influencing the negative text condition. - 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. - 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. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a - `tuple` containing the output images. - """ - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet - - # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): - control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 - control_guidance_start, control_guidance_end = ( - mult * [control_guidance_start], - mult * [control_guidance_end], - ) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - control_image, - strength, - num_inference_steps, - None, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - controlnet_conditioning_scale, - control_guidance_start, - control_guidance_end, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 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 - - if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): - controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) - - # 3.1 Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt, - prompt_2, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - 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=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - - # 3.2 Encode ip_adapter_image - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - - # 4. Prepare image and controlnet_conditioning_image - image = self.image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - - if isinstance(controlnet, ControlNetModel): - control_image = self.prepare_control_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=False, - ) - height, width = control_image.shape[-2:] - elif isinstance(controlnet, MultiControlNetModel): - control_images = [] - - for control_image_ in control_image: - control_image_ = self.prepare_control_image( - image=control_image_, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=controlnet.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=False, - ) - - control_images.append(control_image_) - - control_image = control_images - height, width = control_image[0].shape[-2:] - else: - assert False - - # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - self._num_timesteps = len(timesteps) - - # 6. Prepare latent variables - if latents is None: - latents = self.prepare_latents( - image, - latent_timestep, - batch_size, - num_images_per_prompt, - prompt_embeds.dtype, - device, - generator, - True, - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Create tensor stating which controlnets to keep - controlnet_keep = [] - for i in range(len(timesteps)): - keeps = [ - 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) - for s, e in zip(control_guidance_start, control_guidance_end) - ] - controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) - - # 7.2 Prepare added time ids & embeddings - if isinstance(control_image, list): - original_size = original_size or control_image[0].shape[-2:] - else: - original_size = original_size or control_image.shape[-2:] - target_size = target_size or (height, width) - - if negative_original_size is None: - negative_original_size = original_size - if negative_target_size is None: - negative_target_size = target_size - 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, add_neg_time_ids = self._get_add_time_ids( - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype=prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1) - add_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1) - - control_images = control_image if isinstance(control_image, list) else [control_image] - for i, single_image in enumerate(control_images): - if self.do_classifier_free_guidance: - single_image = single_image.chunk(2)[0] - - if self.do_perturbed_attention_guidance: - single_image = self._prepare_perturbed_attention_guidance( - single_image, single_image, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - single_image = torch.cat([single_image] * 2) - single_image = single_image.to(device) - control_images[i] = single_image - - control_image = control_images if isinstance(control_image, list) else control_images[0] - - if ip_adapter_image_embeds is not None: - for i, image_embeds in enumerate(ip_adapter_image_embeds): - negative_image_embeds = None - if self.do_classifier_free_guidance: - negative_image_embeds, image_embeds = image_embeds.chunk(2) - - if self.do_perturbed_attention_guidance: - image_embeds = self._prepare_perturbed_attention_guidance( - image_embeds, negative_image_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - image_embeds = image_embeds.to(device) - ip_adapter_image_embeds[i] = image_embeds - - if self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - add_text_embeds = self._prepare_perturbed_attention_guidance( - add_text_embeds, negative_pooled_prompt_embeds, self.do_classifier_free_guidance - ) - add_time_ids = self._prepare_perturbed_attention_guidance( - add_time_ids, add_neg_time_ids, self.do_classifier_free_guidance - ) - elif 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([add_neg_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) - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - - controlnet_prompt_embeds = prompt_embeds - controlnet_added_cond_kwargs = added_cond_kwargs - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - if self.do_perturbed_attention_guidance: - original_attn_proc = self.unet.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # controlnet(s) inference - control_model_input = latent_model_input - - if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] - else: - controlnet_cond_scale = controlnet_conditioning_scale - if isinstance(controlnet_cond_scale, list): - controlnet_cond_scale = controlnet_cond_scale[0] - cond_scale = controlnet_cond_scale * controlnet_keep[i] - down_block_res_samples, mid_block_res_sample = self.controlnet( - control_model_input, - t, - encoder_hidden_states=controlnet_prompt_embeds, - controlnet_cond=control_image, - conditioning_scale=cond_scale, - guess_mode=False, - added_cond_kwargs=controlnet_added_cond_kwargs, - return_dict=False, - ) - - if ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = ip_adapter_image_embeds - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=self.cross_attention_kwargs, - down_block_additional_residuals=down_block_res_samples, - mid_block_additional_residual=mid_block_res_sample, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t - ) - elif self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds - ) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - add_neg_time_ids = callback_outputs.pop("add_neg_time_ids", add_neg_time_ids) - - # 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 we do sequential model offloading, let's offload unet and controlnet - # manually for max memory savings - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.unet.to("cpu") - self.controlnet.to("cpu") - empty_device_cache() - - 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) - - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None - has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None - if has_latents_mean and has_latents_std: - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean - else: - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, return_dict=False)[0] - - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - else: - image = latents - return StableDiffusionXLPipelineOutput(images=image) - - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - image = self.image_processor.postprocess(image, output_type=output_type) - - # Offload all models - self.maybe_free_model_hooks() - - if self.do_perturbed_attention_guidance: - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image,) - - return StableDiffusionXLPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_hunyuandit.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_hunyuandit.py deleted file mode 100644 index 3a084086629d45e2cb7a0a0cea4a89c06ee5940e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_hunyuandit.py +++ /dev/null @@ -1,967 +0,0 @@ -# Copyright 2025 HunyuanDiT Authors 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 Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import torch -from transformers import BertModel, BertTokenizer, CLIPImageProcessor, MT5Tokenizer, T5EncoderModel - -from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import VaeImageProcessor -from ...models import AutoencoderKL, HunyuanDiT2DModel -from ...models.attention_processor import PAGCFGHunyuanAttnProcessor2_0, PAGHunyuanAttnProcessor2_0 -from ...models.embeddings import get_2d_rotary_pos_embed -from ...pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker -from ...schedulers import DDPMScheduler -from ...utils import ( - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .pag_utils import PAGMixin - - -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 AutoPipelineForText2Image - - >>> pipe = AutoPipelineForText2Image.from_pretrained( - ... "Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers", - ... torch_dtype=torch.float16, - ... enable_pag=True, - ... pag_applied_layers=[14], - ... ).to("cuda") - - >>> # prompt = "an astronaut riding a horse" - >>> prompt = "一个宇航员在骑马" - >>> image = pipe(prompt, guidance_scale=4, pag_scale=3).images[0] - ``` -""" - -STANDARD_RATIO = np.array( - [ - 1.0, # 1:1 - 4.0 / 3.0, # 4:3 - 3.0 / 4.0, # 3:4 - 16.0 / 9.0, # 16:9 - 9.0 / 16.0, # 9:16 - ] -) -STANDARD_SHAPE = [ - [(1024, 1024), (1280, 1280)], # 1:1 - [(1024, 768), (1152, 864), (1280, 960)], # 4:3 - [(768, 1024), (864, 1152), (960, 1280)], # 3:4 - [(1280, 768)], # 16:9 - [(768, 1280)], # 9:16 -] -STANDARD_AREA = [np.array([w * h for w, h in shapes]) for shapes in STANDARD_SHAPE] -SUPPORTED_SHAPE = [ - (1024, 1024), - (1280, 1280), # 1:1 - (1024, 768), - (1152, 864), - (1280, 960), # 4:3 - (768, 1024), - (864, 1152), - (960, 1280), # 3:4 - (1280, 768), # 16:9 - (768, 1280), # 9:16 -] - - -def map_to_standard_shapes(target_width, target_height): - target_ratio = target_width / target_height - closest_ratio_idx = np.argmin(np.abs(STANDARD_RATIO - target_ratio)) - closest_area_idx = np.argmin(np.abs(STANDARD_AREA[closest_ratio_idx] - target_width * target_height)) - width, height = STANDARD_SHAPE[closest_ratio_idx][closest_area_idx] - return width, height - - -def get_resize_crop_region_for_grid(src, tgt_size): - th = tw = tgt_size - h, w = src - - r = h / w - - # resize - if r > 1: - resize_height = th - resize_width = int(round(th / h * w)) - else: - resize_width = tw - resize_height = int(round(tw / w * h)) - - crop_top = int(round((th - resize_height) / 2.0)) - crop_left = int(round((tw - resize_width) / 2.0)) - - return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width) - - -# 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 HunyuanDiTPAGPipeline(DiffusionPipeline, PAGMixin): - r""" - Pipeline for English/Chinese-to-image generation using HunyuanDiT and [Perturbed Attention - Guidance](https://huggingface.co/docs/diffusers/en/using-diffusers/pag). - - 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.) - - HunyuanDiT uses two text encoders: [mT5](https://huggingface.co/google/mt5-base) and [bilingual CLIP](fine-tuned by - ourselves) - - Args: - 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`]): - 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`]): - A `BertTokenizer` or `CLIPTokenizer` to tokenize text. - transformer ([`HunyuanDiT2DModel`]): - The HunyuanDiT model designed by Tencent Hunyuan. - text_encoder_2 (`T5EncoderModel`): - The mT5 embedder. Specifically, it is 't5-v1_1-xxl'. - tokenizer_2 (`MT5Tokenizer`): - The tokenizer for the mT5 embedder. - scheduler ([`DDPMScheduler`]): - A scheduler to be used in combination with HunyuanDiT to denoise the encoded image latents. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" - _optional_components = [ - "safety_checker", - "feature_extractor", - "text_encoder_2", - "tokenizer_2", - "text_encoder", - "tokenizer", - ] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "prompt_embeds_2", - "negative_prompt_embeds_2", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: BertModel, - tokenizer: BertTokenizer, - transformer: HunyuanDiT2DModel, - scheduler: DDPMScheduler, - safety_checker: Optional[StableDiffusionSafetyChecker] = None, - feature_extractor: Optional[CLIPImageProcessor] = None, - requires_safety_checker: bool = True, - text_encoder_2: Optional[T5EncoderModel] = None, - tokenizer_2: Optional[MT5Tokenizer] = None, - pag_applied_layers: Union[str, List[str]] = "blocks.1", # "blocks.16.attn1", "blocks.16", "16", 16 - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - transformer=transformer, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - text_encoder_2=text_encoder_2, - ) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - 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.register_to_config(requires_safety_checker=requires_safety_checker) - self.default_sample_size = ( - self.transformer.config.sample_size - if hasattr(self, "transformer") and self.transformer is not None - else 128 - ) - - self.set_pag_applied_layers( - pag_applied_layers, pag_attn_processors=(PAGCFGHunyuanAttnProcessor2_0(), PAGHunyuanAttnProcessor2_0()) - ) - - # Copied from diffusers.pipelines.hunyuandit.pipeline_hunyuandit.HunyuanDiTPipeline.encode_prompt - def encode_prompt( - self, - prompt: str, - device: torch.device = None, - 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, - text_encoder_index: int = 0, - ): - 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 - dtype (`torch.dtype`): - torch dtype - 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. - prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the prompt. Required when `prompt_embeds` is passed directly. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the negative prompt. Required when `negative_prompt_embeds` is passed directly. - max_sequence_length (`int`, *optional*): maximum sequence length to use for the prompt. - text_encoder_index (`int`, *optional*): - Index of the text encoder to use. `0` for clip and `1` for T5. - """ - if dtype is None: - if self.text_encoder_2 is not None: - dtype = self.text_encoder_2.dtype - elif self.transformer is not None: - dtype = self.transformer.dtype - else: - dtype = None - - if device is None: - device = self._execution_device - - tokenizers = [self.tokenizer, self.tokenizer_2] - text_encoders = [self.text_encoder, self.text_encoder_2] - - tokenizer = tokenizers[text_encoder_index] - text_encoder = text_encoders[text_encoder_index] - - if max_sequence_length is None: - if text_encoder_index == 0: - max_length = 77 - if text_encoder_index == 1: - max_length = 256 - else: - max_length = max_sequence_length - - 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: - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=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_attention_mask = text_inputs.attention_mask.to(device) - prompt_embeds = text_encoder( - text_input_ids.to(device), - attention_mask=prompt_attention_mask, - ) - prompt_embeds = prompt_embeds[0] - prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1) - - prompt_embeds = prompt_embeds.to(dtype=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 - - max_length = prompt_embeds.shape[1] - uncond_input = tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - - negative_prompt_attention_mask = uncond_input.attention_mask.to(device) - negative_prompt_embeds = text_encoder( - uncond_input.input_ids.to(device), - attention_mask=negative_prompt_attention_mask, - ) - negative_prompt_embeds = negative_prompt_embeds[0] - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 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] - - negative_prompt_embeds = negative_prompt_embeds.to(dtype=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) - - return prompt_embeds, negative_prompt_embeds, prompt_attention_mask, negative_prompt_attention_mask - - # 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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.hunyuandit.pipeline_hunyuandit.HunyuanDiTPipeline.check_inputs - def check_inputs( - self, - prompt, - height, - width, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - prompt_embeds_2=None, - negative_prompt_embeds_2=None, - prompt_attention_mask_2=None, - negative_prompt_attention_mask_2=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_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 None and prompt_embeds_2 is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds_2`. Cannot leave both `prompt` and `prompt_embeds_2` 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 prompt_embeds_2 is not None and prompt_attention_mask_2 is None: - raise ValueError("Must provide `prompt_attention_mask_2` when specifying `prompt_embeds_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." - ) - - 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 negative_prompt_embeds_2 is not None and negative_prompt_attention_mask_2 is None: - raise ValueError( - "Must provide `negative_prompt_attention_mask_2` when specifying `negative_prompt_embeds_2`." - ) - 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_2 is not None and negative_prompt_embeds_2 is not None: - if prompt_embeds_2.shape != negative_prompt_embeds_2.shape: - raise ValueError( - "`prompt_embeds_2` and `negative_prompt_embeds_2` must have the same shape when passed directly, but" - f" got: `prompt_embeds_2` {prompt_embeds_2.shape} != `negative_prompt_embeds_2`" - f" {negative_prompt_embeds_2.shape}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - # 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 - - @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, - 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", - 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"], - 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), - use_resolution_binning: bool = True, - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ): - r""" - The call function to the pipeline for generation with HunyuanDiT. - - Args: - 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. - width (`int`): - 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. This parameter is modulated by `strength`. - 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*): - 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): - The number of images to generate per prompt. - 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*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - 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. - prompt_embeds_2 (`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. - 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. - negative_prompt_embeds_2 (`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. - prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the prompt. Required when `prompt_embeds` is passed directly. - prompt_attention_mask_2 (`torch.Tensor`, *optional*): - Attention mask for the prompt. Required when `prompt_embeds_2` is passed directly. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Attention mask for the negative prompt. Required when `negative_prompt_embeds` is passed directly. - negative_prompt_attention_mask_2 (`torch.Tensor`, *optional*): - Attention mask for the negative prompt. Required when `negative_prompt_embeds_2` is passed directly. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback_on_step_end (`Callable[[int, int, Dict], 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*): - 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)`): - The original size of the image. Used to calculate the time ids. - 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)`): - 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 - standard resolution. Supported resolutions are 1024x1024, 1280x1280, 1024x768, 1152x864, 1280x960, - 768x1024, 864x1152, 960x1280, 1280x768, and 768x1280. It is recommended to set this to `True`. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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 isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - # 0. Default height and width - height = height or self.default_sample_size * self.vae_scale_factor - width = width or self.default_sample_size * self.vae_scale_factor - height = int((height // 16) * 16) - width = int((width // 16) * 16) - - if use_resolution_binning and (height, width) not in SUPPORTED_SHAPE: - width, height = map_to_standard_shapes(width, height) - height = int(height) - width = int(width) - logger.warning(f"Reshaped to (height, width)=({height}, {width}), Supported shapes are {SUPPORTED_SHAPE}") - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - prompt_embeds_2, - negative_prompt_embeds_2, - prompt_attention_mask_2, - negative_prompt_attention_mask_2, - callback_on_step_end_tensor_inputs, - ) - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._interrupt = False - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 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 - ( - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt=prompt, - device=device, - dtype=self.transformer.dtype, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - 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, - max_sequence_length=77, - text_encoder_index=0, - ) - ( - prompt_embeds_2, - negative_prompt_embeds_2, - prompt_attention_mask_2, - negative_prompt_attention_mask_2, - ) = self.encode_prompt( - prompt=prompt, - device=device, - dtype=self.transformer.dtype, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds_2, - negative_prompt_embeds=negative_prompt_embeds_2, - prompt_attention_mask=prompt_attention_mask_2, - negative_prompt_attention_mask=negative_prompt_attention_mask_2, - max_sequence_length=256, - text_encoder_index=1, - ) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Create image_rotary_emb, style embedding & time ids - grid_height = height // 8 // self.transformer.config.patch_size - grid_width = width // 8 // self.transformer.config.patch_size - base_size = 512 // 8 // self.transformer.config.patch_size - grid_crops_coords = get_resize_crop_region_for_grid((grid_height, grid_width), base_size) - image_rotary_emb = get_2d_rotary_pos_embed( - self.transformer.inner_dim // self.transformer.num_heads, - grid_crops_coords, - (grid_height, grid_width), - device=device, - output_type="pt", - ) - - style = torch.tensor([0], device=device) - - target_size = target_size or (height, width) - add_time_ids = list(original_size + target_size + crops_coords_top_left) - add_time_ids = torch.tensor([add_time_ids], dtype=prompt_embeds.dtype) - - # 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 self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - prompt_attention_mask = self._prepare_perturbed_attention_guidance( - prompt_attention_mask, negative_prompt_attention_mask, self.do_classifier_free_guidance - ) - prompt_embeds_2 = self._prepare_perturbed_attention_guidance( - prompt_embeds_2, negative_prompt_embeds_2, self.do_classifier_free_guidance - ) - prompt_attention_mask_2 = self._prepare_perturbed_attention_guidance( - prompt_attention_mask_2, negative_prompt_attention_mask_2, self.do_classifier_free_guidance - ) - add_time_ids = torch.cat([add_time_ids] * 3, dim=0) - style = torch.cat([style] * 3, dim=0) - elif self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask]) - prompt_embeds_2 = torch.cat([negative_prompt_embeds_2, prompt_embeds_2]) - prompt_attention_mask_2 = torch.cat([negative_prompt_attention_mask_2, prompt_attention_mask_2]) - add_time_ids = torch.cat([add_time_ids] * 2, dim=0) - style = torch.cat([style] * 2, dim=0) - - prompt_embeds = prompt_embeds.to(device=device) - prompt_attention_mask = prompt_attention_mask.to(device=device) - prompt_embeds_2 = prompt_embeds_2.to(device=device) - prompt_attention_mask_2 = prompt_attention_mask_2.to(device=device) - add_time_ids = add_time_ids.to(dtype=prompt_embeds.dtype, device=device).repeat( - batch_size * num_images_per_prompt, 1 - ) - style = style.to(device=device).repeat(batch_size * num_images_per_prompt) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - self._num_timesteps = len(timesteps) - - if self.do_perturbed_attention_guidance: - original_attn_proc = self.transformer.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # expand scalar t to 1-D tensor to match the 1st dim of latent_model_input - t_expand = torch.tensor([t] * latent_model_input.shape[0], device=device).to( - dtype=latent_model_input.dtype - ) - - # predict the noise residual - noise_pred = self.transformer( - latent_model_input, - t_expand, - encoder_hidden_states=prompt_embeds, - text_embedding_mask=prompt_attention_mask, - encoder_hidden_states_t5=prompt_embeds_2, - text_embedding_mask_t5=prompt_attention_mask_2, - image_meta_size=add_time_ids, - style=style, - image_rotary_emb=image_rotary_emb, - return_dict=False, - )[0] - - noise_pred, _ = noise_pred.chunk(2, dim=1) - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred, noise_pred_text = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t, True - ) - elif self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - if self.do_classifier_free_guidance and guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - prompt_embeds_2 = callback_outputs.pop("prompt_embeds_2", prompt_embeds_2) - negative_prompt_embeds_2 = callback_outputs.pop( - "negative_prompt_embeds_2", negative_prompt_embeds_2 - ) - - 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 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) - - # 9. Offload all models - self.maybe_free_model_hooks() - - if self.do_perturbed_attention_guidance: - self.transformer.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_kolors.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_kolors.py deleted file mode 100644 index ed8e33e2ba8bd13ae1e178963c87cb1cdc967ba3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_kolors.py +++ /dev/null @@ -1,1138 +0,0 @@ -# Copyright 2025 Stability AI, Kwai-Kolors 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, Tuple, Union - -import torch -from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import IPAdapterMixin, StableDiffusionXLLoraLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.attention_processor import AttnProcessor2_0, FusedAttnProcessor2_0, XFormersAttnProcessor -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import is_torch_xla_available, logging, replace_example_docstring -from ...utils.torch_utils import randn_tensor -from ..kolors.pipeline_output import KolorsPipelineOutput -from ..kolors.text_encoder import ChatGLMModel -from ..kolors.tokenizer import ChatGLMTokenizer -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from .pag_utils import PAGMixin - - -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 AutoPipelineForText2Image - - >>> pipe = AutoPipelineForText2Image.from_pretrained( - ... "Kwai-Kolors/Kolors-diffusers", - ... variant="fp16", - ... torch_dtype=torch.float16, - ... enable_pag=True, - ... pag_applied_layers=["down.block_2.attentions_1", "up.block_0.attentions_1"], - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = ( - ... "A photo of a ladybug, macro, zoom, high quality, film, holding a wooden sign with the text 'KOLORS'" - ... ) - >>> image = pipe(prompt, guidance_scale=5.5, pag_scale=1.5).images[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, - **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 KolorsPAGPipeline( - DiffusionPipeline, StableDiffusionMixin, StableDiffusionXLLoraLoaderMixin, IPAdapterMixin, PAGMixin -): - r""" - Pipeline for text-to-image generation using Kolors. - - 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.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 ([`ChatGLMModel`]): - Frozen text-encoder. Kolors uses [ChatGLM3-6B](https://huggingface.co/THUDM/chatglm3-6b). - tokenizer (`ChatGLMTokenizer`): - Tokenizer of class - [ChatGLMTokenizer](https://huggingface.co/THUDM/chatglm3-6b/blob/main/tokenization_chatglm.py). - 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 `"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"`): - 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.: - ["mid"] ["down", "mid"] ["down", "mid", "up.block_1"] ["down", "mid", "up.block_1.attentions_0", - "up.block_1.attentions_1"] - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = [ - "image_encoder", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "negative_add_time_ids", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: ChatGLMModel, - tokenizer: ChatGLMTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - image_encoder: CLIPVisionModelWithProjection = None, - feature_extractor: CLIPImageProcessor = None, - force_zeros_for_empty_prompt: bool = False, - pag_applied_layers: Union[str, List[str]] = "mid", - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - ) - 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 - ) - - self.set_pag_applied_layers(pag_applied_layers) - - # Copied from diffusers.pipelines.kolors.pipeline_kolors.KolorsPipeline.encode_prompt - def encode_prompt( - self, - prompt, - device: Optional[torch.device] = 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, - max_sequence_length: int = 256, - ): - 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.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. - 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_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. - 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. - max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`. - """ - # from IPython import embed; embed(); exit() - device = device or self._execution_device - - 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] - - # Define tokenizers and text encoders - tokenizers = [self.tokenizer] - text_encoders = [self.text_encoder] - - if prompt_embeds is None: - prompt_embeds_list = [] - for tokenizer, text_encoder in zip(tokenizers, text_encoders): - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - ).to(device) - output = text_encoder( - input_ids=text_inputs["input_ids"], - attention_mask=text_inputs["attention_mask"], - position_ids=text_inputs["position_ids"], - output_hidden_states=True, - ) - - # [max_sequence_length, batch, hidden_size] -> [batch, max_sequence_length, hidden_size] - # clone to have a contiguous tensor - prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone() - # [max_sequence_length, batch, hidden_size] -> [batch, hidden_size] - pooled_prompt_embeds = output.hidden_states[-1][-1, :, :].clone() - bs_embed, seq_len, _ = prompt_embeds.shape - 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) - - prompt_embeds_list.append(prompt_embeds) - - prompt_embeds = prompt_embeds_list[0] - - # 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) - elif 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 - - negative_prompt_embeds_list = [] - - for tokenizer, text_encoder in zip(tokenizers, text_encoders): - uncond_input = tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - ).to(device) - output = text_encoder( - input_ids=uncond_input["input_ids"], - attention_mask=uncond_input["attention_mask"], - position_ids=uncond_input["position_ids"], - output_hidden_states=True, - ) - - # [max_sequence_length, batch, hidden_size] -> [batch, max_sequence_length, hidden_size] - # clone to have a contiguous tensor - negative_prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone() - # [max_sequence_length, batch, hidden_size] -> [batch, hidden_size] - negative_pooled_prompt_embeds = output.hidden_states[-1][-1, :, :].clone() - - 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=text_encoder.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 - ) - - negative_prompt_embeds_list.append(negative_prompt_embeds) - - negative_prompt_embeds = negative_prompt_embeds_list[0] - - bs_embed = pooled_prompt_embeds.shape[0] - 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 - ) - - return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.kolors.pipeline_kolors.KolorsPipeline.check_inputs - def check_inputs( - self, - prompt, - num_inference_steps, - height, - width, - negative_prompt=None, - prompt_embeds=None, - pooled_prompt_embeds=None, - negative_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - ): - if not isinstance(num_inference_steps, int) or num_inference_steps <= 0: - raise ValueError( - f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type" - f" {type(num_inference_steps)}." - ) - - 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_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 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`." - ) - - 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" - ) - - if max_sequence_length is not None and max_sequence_length > 256: - raise ValueError(f"`max_sequence_length` cannot be greater than 256 but is {max_sequence_length}") - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - 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): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - FusedAttnProcessor2_0, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 do_classifier_free_guidance(self): - return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_end(self): - return self._denoising_end - - @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, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = None, - guidance_scale: float = 5.0, - 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, - 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", - 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"], - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - max_sequence_length: int = 256, - ): - 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. - Anything below 512 pixels won't work well for - [Kwai-Kolors/Kolors-diffusers](https://huggingface.co/Kwai-Kolors/Kolors-diffusers) 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 - [Kwai-Kolors/Kolors-diffusers](https://huggingface.co/Kwai-Kolors/Kolors-diffusers) 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. - 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*): - 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. - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output) - 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. - 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*): - 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 ge 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. - 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_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. - 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. - 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)`. 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. - 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.kolors.KolorsPipelineOutput`] instead of a plain tuple. - cross_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). - 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. - 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. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`. - - Examples: - - Returns: - [`~pipelines.kolors.KolorsPipelineOutput`] or `tuple`: [`~pipelines.kolors.KolorsPipelineOutput`] if - `return_dict` is True, otherwise a `tuple`. When returning a tuple, 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 - - # 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, - num_inference_steps, - height, - width, - negative_prompt, - prompt_embeds, - pooled_prompt_embeds, - negative_prompt_embeds, - negative_pooled_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - ) - - self._guidance_scale = guidance_scale - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - self._interrupt = False - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 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 - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - ) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 5. 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, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Prepare added time ids & embeddings - add_text_embeds = pooled_prompt_embeds - text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1]) - - 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_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - add_text_embeds = self._prepare_perturbed_attention_guidance( - add_text_embeds, negative_pooled_prompt_embeds, self.do_classifier_free_guidance - ) - add_time_ids = self._prepare_perturbed_attention_guidance( - add_time_ids, negative_add_time_ids, self.do_classifier_free_guidance - ) - elif 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) - - 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, - self.do_classifier_free_guidance, - ) - - for i, image_embeds in enumerate(ip_adapter_image_embeds): - negative_image_embeds = None - if self.do_classifier_free_guidance: - negative_image_embeds, image_embeds = image_embeds.chunk(2) - - if self.do_perturbed_attention_guidance: - image_embeds = self._prepare_perturbed_attention_guidance( - image_embeds, negative_image_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - image_embeds = image_embeds.to(device) - ip_adapter_image_embeds[i] = image_embeds - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - # 8.1 Apply denoising_end - if ( - self.denoising_end is not None - and isinstance(self.denoising_end, float) - and self.denoising_end > 0 - and self.denoising_end < 1 - ): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - # 9. 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) - - if self.do_perturbed_attention_guidance: - original_attn_proc = self.unet.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = image_embeds - - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t - ) - elif 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) - - # compute the previous noisy sample x_t -> x_t-1 - latents_dtype = latents.dtype - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds - ) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids) - - 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 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) - elif latents.dtype != self.vae.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 - self.vae = self.vae.to(latents.dtype) - - # unscale/denormalize the latents - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, 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 self.do_perturbed_attention_guidance: - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image,) - - return KolorsPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py deleted file mode 100644 index d9d6d14a38d96a9b0dea995cac3a01cf5f11efff..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py +++ /dev/null @@ -1,878 +0,0 @@ -# Copyright 2025 PixArt-Sigma Authors 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 inspect -import re -import urllib.parse as ul -from typing import Callable, List, Optional, Tuple, Union - -import torch -from transformers import T5EncoderModel, T5Tokenizer - -from ...image_processor import PixArtImageProcessor -from ...models import AutoencoderKL, PixArtTransformer2DModel -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - BACKENDS_MAPPING, - deprecate, - is_bs4_available, - is_ftfy_available, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput -from ..pixart_alpha.pipeline_pixart_alpha import ( - ASPECT_RATIO_256_BIN, - ASPECT_RATIO_512_BIN, - ASPECT_RATIO_1024_BIN, -) -from ..pixart_alpha.pipeline_pixart_sigma import ASPECT_RATIO_2048_BIN -from .pag_utils import PAGMixin - - -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_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import AutoPipelineForText2Image - - >>> pipe = AutoPipelineForText2Image.from_pretrained( - ... "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", - ... torch_dtype=torch.float16, - ... pag_applied_layers=["blocks.14"], - ... enable_pag=True, - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = "A small cactus with a happy face in the Sahara desert" - >>> image = pipe(prompt, pag_scale=4.0, guidance_scale=1.0).images[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, - **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 PixArtSigmaPAGPipeline(DiffusionPipeline, PAGMixin): - r""" - [PAG pipeline](https://huggingface.co/docs/diffusers/main/en/using-diffusers/pag) for text-to-image generation - using PixArt-Sigma. - """ - - bad_punct_regex = re.compile( - r"[" - + "#®•©™&@·º½¾¿¡§~" - + r"\)" - + r"\(" - + r"\]" - + r"\[" - + r"\}" - + r"\{" - + r"\|" - + "\\" - + r"\/" - + r"\*" - + r"]{1,}" - ) # noqa - - _optional_components = ["tokenizer", "text_encoder"] - model_cpu_offload_seq = "text_encoder->transformer->vae" - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - vae: AutoencoderKL, - transformer: PixArtTransformer2DModel, - scheduler: KarrasDiffusionSchedulers, - pag_applied_layers: Union[str, List[str]] = "blocks.1", # 1st transformer block - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_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 = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor) - - self.set_pag_applied_layers(pag_applied_layers) - - # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt with 120->300 - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - max_sequence_length: int = 300, - **kwargs, - ): - 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 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 "". - 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 - device: (`torch.device`, *optional*): - 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. For PixArt-Alpha, it's should be the embeddings of the "" - string. - clean_caption (`bool`, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. - """ - - if "mask_feature" in kwargs: - deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version." - deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False) - - if device is None: - device = self._execution_device - - # See Section 3.1. of the paper. - max_length = max_sequence_length - - if prompt_embeds is None: - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - add_special_tokens=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[:, max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because T5 can only handle sequences up to" - f" {max_length} tokens: {removed_text}" - ) - - prompt_attention_mask = text_inputs.attention_mask - prompt_attention_mask = prompt_attention_mask.to(device) - - prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask) - prompt_embeds = prompt_embeds[0] - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - elif self.transformer is not None: - dtype = self.transformer.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - - bs_embed, 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(bs_embed * num_images_per_prompt, seq_len, -1) - prompt_attention_mask = prompt_attention_mask.repeat(1, num_images_per_prompt) - prompt_attention_mask = prompt_attention_mask.view(bs_embed * 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 = [negative_prompt] * bs_embed if isinstance(negative_prompt, str) else negative_prompt - uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption) - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - negative_prompt_attention_mask = uncond_input.attention_mask - negative_prompt_attention_mask = negative_prompt_attention_mask.to(device) - - negative_prompt_embeds = self.text_encoder( - uncond_input.input_ids.to(device), attention_mask=negative_prompt_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=dtype, device=device) - - negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) - - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(1, num_images_per_prompt) - negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed * num_images_per_prompt, -1) - else: - negative_prompt_embeds = None - negative_prompt_attention_mask = None - - return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.check_inputs - def check_inputs( - self, - prompt, - height, - width, - negative_prompt, - callback_steps, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=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 None) or ( - 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 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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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("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}." - ) - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - negative_prompt: str = "", - num_inference_steps: int = 20, - 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, - 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], 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]: - """ - 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 - 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*): - 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*): - 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 4.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. - height (`int`, *optional*, defaults to self.unet.config.sample_size): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to self.unet.config.sample_size): - The width in pixels of the generated image. - 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*): - 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 ge 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.Tensor`, *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.Tensor`, *optional*): - Pre-generated attention mask for negative text embeddings. - 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.IFPipelineOutput`] 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. - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - use_resolution_binning (`bool` defaults to `True`): - If set to `True`, the requested height and width are first mapped to the closest resolutions using - `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to - the requested resolution. Useful for generating non-square images. - max_sequence_length (`int` defaults to 300): Maximum sequence length to use with the `prompt`. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images - """ - # 1. Check inputs. Raise error if not correct - height = height or self.transformer.config.sample_size * self.vae_scale_factor - width = width or self.transformer.config.sample_size * self.vae_scale_factor - if use_resolution_binning: - if self.transformer.config.sample_size == 256: - aspect_ratio_bin = ASPECT_RATIO_2048_BIN - elif self.transformer.config.sample_size == 128: - aspect_ratio_bin = ASPECT_RATIO_1024_BIN - elif self.transformer.config.sample_size == 64: - aspect_ratio_bin = ASPECT_RATIO_512_BIN - elif self.transformer.config.sample_size == 32: - aspect_ratio_bin = ASPECT_RATIO_256_BIN - else: - raise ValueError("Invalid sample size") - orig_height, orig_width = height, width - height, width = self.image_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin) - - self.check_inputs( - prompt, - height, - width, - negative_prompt, - callback_steps, - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - ) - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 2. Default height and width to transformer - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt, - do_classifier_free_guidance, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - ) - if self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, do_classifier_free_guidance - ) - prompt_attention_mask = self._prepare_perturbed_attention_guidance( - prompt_attention_mask, negative_prompt_attention_mask, do_classifier_free_guidance - ) - elif 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) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 5. Prepare latents. - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - if self.do_perturbed_attention_guidance: - original_attn_proc = self.transformer.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=do_classifier_free_guidance, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 6.1 Prepare micro-conditions. - added_cond_kwargs = {"resolution": None, "aspect_ratio": None} - - # 7. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance, perturbed-attention guidance, or both - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - current_timestep = t - if not torch.is_tensor(current_timestep): - # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = latent_model_input.device.type == "mps" - is_npu = latent_model_input.device.type == "npu" - if isinstance(current_timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device) - elif len(current_timestep.shape) == 0: - current_timestep = current_timestep[None].to(latent_model_input.device) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - current_timestep = current_timestep.expand(latent_model_input.shape[0]) - - # predict noise model_output - noise_pred = self.transformer( - latent_model_input, - encoder_hidden_states=prompt_embeds, - encoder_attention_mask=prompt_attention_mask, - timestep=current_timestep, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred = self._apply_perturbed_attention_guidance( - noise_pred, do_classifier_free_guidance, guidance_scale, current_timestep - ) - elif do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # learned sigma - if self.transformer.config.out_channels // 2 == latent_channels: - noise_pred = noise_pred.chunk(2, dim=1)[0] - else: - noise_pred = noise_pred - - # compute previous image: x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - if use_resolution_binning: - image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height) - 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 self.do_perturbed_attention_guidance: - self.transformer.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sana.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sana.py deleted file mode 100644 index 8dbae13a3f16e29c34659107e17307772cc8cac4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sana.py +++ /dev/null @@ -1,949 +0,0 @@ -# Copyright 2025 PixArt-Sigma Authors 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 inspect -import re -import urllib.parse as ul -import warnings -from typing import Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PixArtImageProcessor -from ...models import AutoencoderDC, SanaTransformer2DModel -from ...models.attention_processor import PAGCFGSanaLinearAttnProcessor2_0, PAGIdentitySanaLinearAttnProcessor2_0 -from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import ( - BACKENDS_MAPPING, - is_bs4_available, - is_ftfy_available, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import get_device, is_torch_version, randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput -from ..pixart_alpha.pipeline_pixart_alpha import ( - ASPECT_RATIO_512_BIN, - ASPECT_RATIO_1024_BIN, -) -from ..pixart_alpha.pipeline_pixart_sigma import ASPECT_RATIO_2048_BIN -from ..sana.pipeline_sana import ASPECT_RATIO_4096_BIN -from .pag_utils import PAGMixin - - -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_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import SanaPAGPipeline - - >>> pipe = SanaPAGPipeline.from_pretrained( - ... "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers", - ... pag_applied_layers=["transformer_blocks.8"], - ... torch_dtype=torch.float32, - ... ) - >>> pipe.to("cuda") - >>> pipe.text_encoder.to(torch.bfloat16) - >>> pipe.transformer = pipe.transformer.to(torch.bfloat16) - - >>> image = pipe(prompt='a cyberpunk cat with a neon sign that says "Sana"')[0] - >>> image[0].save("output.png") - ``` -""" - - -# 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 SanaPAGPipeline(DiffusionPipeline, PAGMixin): - r""" - Pipeline for text-to-image generation using [Sana](https://huggingface.co/papers/2410.10629). This pipeline - supports the use of [Perturbed Attention Guidance - (PAG)](https://huggingface.co/docs/diffusers/main/en/using-diffusers/pag). - """ - - # fmt: off - bad_punct_regex = re.compile(r"[" + "#®•©™&@·º½¾¿¡§~" + r"\)" + r"\(" + r"\]" + r"\[" + r"\}" + r"\{" + r"\|" + "\\" + r"\/" + r"\*" + r"]{1,}") - # fmt: on - - model_cpu_offload_seq = "text_encoder->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], - text_encoder: Gemma2PreTrainedModel, - vae: AutoencoderDC, - transformer: SanaTransformer2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - pag_applied_layers: Union[str, List[str]] = "transformer_blocks.0", - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler - ) - - self.vae_scale_factor = ( - 2 ** (len(self.vae.config.encoder_block_out_channels) - 1) - if hasattr(self, "vae") and self.vae is not None - else 8 - ) - self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor) - - self.set_pag_applied_layers( - pag_applied_layers, - pag_attn_processors=(PAGCFGSanaLinearAttnProcessor2_0(), PAGIdentitySanaLinearAttnProcessor2_0()), - ) - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = 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 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 "". - 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 - device: (`torch.device`, *optional*): - 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. For Sana, it's should be the embeddings of the "" string. - clean_caption (`bool`, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. - complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): - If `complex_human_instruction` is not empty, the function will use the complex Human instruction for - the prompt. - """ - - if device is None: - device = self._execution_device - - 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 getattr(self, "tokenizer", None) is not None: - self.tokenizer.padding_side = "right" - - # See Section 3.1. of the paper. - max_length = max_sequence_length - select_index = [0] + list(range(-max_length + 1, 0)) - - if prompt_embeds is None: - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - - # prepare complex human instruction - if not complex_human_instruction: - max_length_all = max_length - else: - chi_prompt = "\n".join(complex_human_instruction) - prompt = [chi_prompt + p for p in prompt] - num_chi_prompt_tokens = len(self.tokenizer.encode(chi_prompt)) - max_length_all = num_chi_prompt_tokens + max_length - 2 - - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length_all, - 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.to(device) - - prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask) - prompt_embeds = prompt_embeds[0][:, select_index] - prompt_attention_mask = prompt_attention_mask[:, select_index] - - if self.transformer is not None: - dtype = self.transformer.dtype - elif self.text_encoder is not None: - dtype = self.text_encoder.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - - bs_embed, 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(bs_embed * num_images_per_prompt, seq_len, -1) - prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1) - prompt_attention_mask = prompt_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 = [negative_prompt] * batch_size if isinstance(negative_prompt, str) else negative_prompt - uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption) - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - negative_prompt_attention_mask = uncond_input.attention_mask - negative_prompt_attention_mask = negative_prompt_attention_mask.to(device) - - negative_prompt_embeds = self.text_encoder( - uncond_input.input_ids.to(device), attention_mask=negative_prompt_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=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) - - negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1) - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1) - else: - negative_prompt_embeds = None - negative_prompt_attention_mask = None - - return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.check_inputs - def check_inputs( - self, - prompt, - height, - width, - callback_on_step_end_tensor_inputs=None, - negative_prompt=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 is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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("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}." - ) - - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.prepare_latents - 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=device, dtype=dtype) - - 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 - - @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 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: str = "", - num_inference_steps: int = 20, - timesteps: List[int] = None, - sigmas: List[float] = None, - guidance_scale: float = 4.5, - num_images_per_prompt: Optional[int] = 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", - 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"], - max_sequence_length: int = 300, - 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.", - "Here are examples of how to transform or refine prompts:", - "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.", - "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.", - "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", - "User Prompt: ", - ], - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ) -> Union[ImagePipelineOutput, 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. 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 - 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*): - 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*): - 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 4.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. - height (`int`, *optional*, defaults to self.unet.config.sample_size): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to self.unet.config.sample_size): - The width in pixels of the generated image. - 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*): - 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 ge 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.Tensor`, *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.Tensor`, *optional*): - Pre-generated attention mask for negative text embeddings. - 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.IFPipelineOutput`] instead of a plain tuple. - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - use_resolution_binning (`bool` defaults to `True`): - If set to `True`, the requested height and width are first mapped to the closest resolutions using - `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to - the requested resolution. Useful for generating non-square images. - 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 300): Maximum sequence length to use with the `prompt`. - 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): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] 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 - - if use_resolution_binning: - if self.transformer.config.sample_size == 128: - aspect_ratio_bin = ASPECT_RATIO_4096_BIN - elif self.transformer.config.sample_size == 64: - aspect_ratio_bin = ASPECT_RATIO_2048_BIN - elif self.transformer.config.sample_size == 32: - aspect_ratio_bin = ASPECT_RATIO_1024_BIN - elif self.transformer.config.sample_size == 16: - aspect_ratio_bin = ASPECT_RATIO_512_BIN - else: - raise ValueError("Invalid sample size") - orig_height, orig_width = height, width - height, width = self.image_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin) - - self.check_inputs( - prompt, - height, - width, - callback_on_step_end_tensor_inputs, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - ) - - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - self._guidance_scale = guidance_scale - self._interrupt = False - - # 2. Default height and width to transformer - 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 - ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt, - self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - complex_human_instruction=complex_human_instruction, - ) - - if self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - prompt_attention_mask = self._prepare_perturbed_attention_guidance( - prompt_attention_mask, negative_prompt_attention_mask, self.do_classifier_free_guidance - ) - elif 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) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 5. Prepare latents. - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - torch.float32, - device, - generator, - latents, - ) - if self.do_perturbed_attention_guidance: - original_attn_proc = self.transformer.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - 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 - - # expand the latents if we are doing classifier free guidance, perturbed-attention guidance, or both - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - latent_model_input = 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]).to(latents.dtype) - - # predict noise model_output - noise_pred = self.transformer( - latent_model_input, - encoder_hidden_states=prompt_embeds, - encoder_attention_mask=prompt_attention_mask, - timestep=timestep, - return_dict=False, - )[0] - noise_pred = noise_pred.float() - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, guidance_scale, t - ) - elif self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute previous image: x_t -> x_t-1 - latents_dtype = latents.dtype - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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 == 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) - torch_accelerator_module = getattr(torch, get_device(), torch.cuda) - oom_error = ( - torch.OutOfMemoryError - if is_torch_version(">=", "2.5.0") - else torch_accelerator_module.OutOfMemoryError - ) - try: - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - except oom_error as e: - warnings.warn( - f"{e}. \n" - f"Try to use VAE tiling for large images. For example: \n" - f"pipe.vae.enable_tiling(tile_sample_min_width=512, tile_sample_min_height=512)" - ) - if use_resolution_binning: - image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height) - - if not output_type == "latent": - image = self.image_processor.postprocess(image, output_type=output_type) - - # Offload all models - self.maybe_free_model_hooks() - - if self.do_perturbed_attention_guidance: - self.transformer.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd.py deleted file mode 100644 index 349d006aada34aeaeb0d17b92fefb820fe501fa1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd.py +++ /dev/null @@ -1,1078 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import torch -from packaging import version -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...configuration_utils import FrozenDict -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - 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, StableDiffusionMixin -from ..stable_diffusion.pipeline_output import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker -from .pag_utils import PAGMixin - - -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 AutoPipelineForText2Image - - >>> pipe = AutoPipelineForText2Image.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, enable_pag=True - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> image = pipe(prompt, pag_scale=0.3).images[0] - ``` -""" - - -# 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 - - -# 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 StableDiffusionPAGPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - IPAdapterMixin, - FromSingleFileMixin, - PAGMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - 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", - ): - super().__init__() - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_config) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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.register_to_config(requires_safety_checker=requires_safety_checker) - - self.set_pag_applied_layers(pag_applied_layers) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.check_inputs - def check_inputs( - self, - prompt, - height, - width, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_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}." - ) - - 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_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, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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"], - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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. - 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*): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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). - guidance_rescale (`float`, *optional*, defaults to 0.0): - Guidance rescale factor from [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. - 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. - 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. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - - # 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 - # to deal with lora scaling and other possible forward hooks - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - None, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._interrupt = False - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 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 = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=lora_scale, - clip_skip=self.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 self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - - for i, image_embeds in enumerate(ip_adapter_image_embeds): - negative_image_embeds = None - if self.do_classifier_free_guidance: - negative_image_embeds, image_embeds = image_embeds.chunk(2) - if self.do_perturbed_attention_guidance: - image_embeds = self._prepare_perturbed_attention_guidance( - image_embeds, negative_image_embeds, self.do_classifier_free_guidance - ) - - elif self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - image_embeds = image_embeds.to(device) - ip_adapter_image_embeds[i] = image_embeds - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 5. 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, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 6.1 Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": ip_adapter_image_embeds} - if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) - else None - ) - - # 6.2 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) - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - if self.do_perturbed_attention_guidance: - original_attn_proc = self.unet.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred, noise_pred_text = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t, True - ) - - elif 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.do_classifier_free_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - - # 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 not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 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 self.do_perturbed_attention_guidance: - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_3.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_3.py deleted file mode 100644 index 96796f53b0bcbfae137909b2f8bb11300084b94b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_3.py +++ /dev/null @@ -1,992 +0,0 @@ -# Copyright 2025 Stability 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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import torch -from transformers import ( - CLIPTextModelWithProjection, - CLIPTokenizer, - T5EncoderModel, - T5TokenizerFast, -) - -from ...image_processor import VaeImageProcessor -from ...loaders import FromSingleFileMixin, SD3LoraLoaderMixin -from ...models.attention_processor import PAGCFGJointAttnProcessor2_0, PAGJointAttnProcessor2_0 -from ...models.autoencoders import AutoencoderKL -from ...models.transformers import SD3Transformer2DModel -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 ..stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput -from .pag_utils import PAGMixin - - -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 AutoPipelineForText2Image - - >>> pipe = AutoPipelineForText2Image.from_pretrained( - ... "stabilityai/stable-diffusion-3-medium-diffusers", - ... torch_dtype=torch.float16, - ... enable_pag=True, - ... pag_applied_layers=["blocks.13"], - ... ) - >>> pipe.to("cuda") - >>> prompt = "A cat holding a sign that says hello world" - >>> image = pipe(prompt, guidance_scale=5.0, pag_scale=0.7).images[0] - >>> image.save("sd3_pag.png") - ``` -""" - - -# 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 StableDiffusion3PAGPipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, PAGMixin): - r""" - [PAG pipeline](https://huggingface.co/docs/diffusers/main/en/using-diffusers/pag) for text-to-image generation - using Stable Diffusion 3. - - Args: - transformer ([`SD3Transformer2DModel`]): - 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 ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModelWithProjection`]): - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), - specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant, - with an additional added projection layer that is initialized with a diagonal matrix with the `hidden_size` - as its dimension. - text_encoder_2 ([`CLIPTextModelWithProjection`]): - [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. - text_encoder_3 ([`T5EncoderModel`]): - Frozen text-encoder. Stable Diffusion 3 uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the - [t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) 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). - tokenizer_3 (`T5TokenizerFast`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->transformer->vae" - _optional_components = [] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "negative_pooled_prompt_embeds"] - - def __init__( - self, - transformer: SD3Transformer2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer_2: CLIPTokenizer, - text_encoder_3: T5EncoderModel, - tokenizer_3: T5TokenizerFast, - pag_applied_layers: Union[str, List[str]] = "blocks.1", # 1st transformer block - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - text_encoder_3=text_encoder_3, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - tokenizer_3=tokenizer_3, - 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.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = ( - self.transformer.config.sample_size - if hasattr(self, "transformer") and self.transformer is not None - else 128 - ) - self.patch_size = ( - self.transformer.config.patch_size if hasattr(self, "transformer") and self.transformer is not None else 2 - ) - - self.set_pag_applied_layers( - pag_applied_layers, pag_attn_processors=(PAGCFGJointAttnProcessor2_0(), PAGJointAttnProcessor2_0()) - ) - - # 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, - num_images_per_prompt: int = 1, - max_sequence_length: int = 256, - 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 self.text_encoder_3 is None: - return torch.zeros( - ( - batch_size * num_images_per_prompt, - self.tokenizer_max_length, - self.transformer.config.joint_attention_dim, - ), - device=device, - dtype=dtype, - ) - - text_inputs = self.tokenizer_3( - 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 - untruncated_ids = self.tokenizer_3(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_3.batch_decode(untruncated_ids[:, self.tokenizer_max_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_3(text_input_ids.to(device))[0] - - dtype = self.text_encoder_3.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clip_skip: Optional[int] = None, - clip_model_index: int = 0, - ): - device = device or self._execution_device - - clip_tokenizers = [self.tokenizer, self.tokenizer_2] - clip_text_encoders = [self.text_encoder, self.text_encoder_2] - - tokenizer = clip_tokenizers[clip_model_index] - text_encoder = clip_text_encoders[clip_model_index] - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer_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[:, self.tokenizer_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) - pooled_prompt_embeds = prompt_embeds[0] - - if clip_skip is None: - prompt_embeds = prompt_embeds.hidden_states[-2] - else: - prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)] - - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - _, 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(batch_size * num_images_per_prompt, seq_len, -1) - - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) - pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds, pooled_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, - 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, - max_sequence_length: int = 256, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - 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`): - 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 all the text-encoders. - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. - 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. - 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, SD3LoraLoaderMixin): - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - 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] - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - prompt_3 = prompt_3 or prompt - prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3 - - prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=clip_skip, - clip_model_index=0, - ) - prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds( - prompt=prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=clip_skip, - clip_model_index=1, - ) - clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1) - - t5_prompt_embed = self._get_t5_prompt_embeds( - prompt=prompt_3, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - clip_prompt_embeds = torch.nn.functional.pad( - clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1]) - ) - - prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2) - pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1) - - if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt_2 = negative_prompt_2 or negative_prompt - negative_prompt_3 = negative_prompt_3 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 - ) - negative_prompt_3 = ( - batch_size * [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3 - ) - - 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_embed, negative_pooled_prompt_embed = self._get_clip_prompt_embeds( - negative_prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=None, - clip_model_index=0, - ) - negative_prompt_2_embed, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds( - negative_prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=None, - clip_model_index=1, - ) - negative_clip_prompt_embeds = torch.cat([negative_prompt_embed, negative_prompt_2_embed], dim=-1) - - t5_negative_prompt_embed = self._get_t5_prompt_embeds( - prompt=negative_prompt_3, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - negative_clip_prompt_embeds = torch.nn.functional.pad( - negative_clip_prompt_embeds, - (0, t5_negative_prompt_embed.shape[-1] - negative_clip_prompt_embeds.shape[-1]), - ) - - negative_prompt_embeds = torch.cat([negative_clip_prompt_embeds, t5_negative_prompt_embed], dim=-2) - negative_pooled_prompt_embeds = torch.cat( - [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1 - ) - - if self.text_encoder is not None: - if isinstance(self, SD3LoraLoaderMixin) 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, SD3LoraLoaderMixin) 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 - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.check_inputs - def check_inputs( - self, - prompt, - prompt_2, - prompt_3, - height, - width, - negative_prompt=None, - negative_prompt_2=None, - negative_prompt_3=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - ): - if ( - height % (self.vae_scale_factor * self.patch_size) != 0 - or width % (self.vae_scale_factor * self.patch_size) != 0 - ): - raise ValueError( - f"`height` and `width` have to be divisible by {self.vae_scale_factor * self.patch_size} but are {height} and {width}." - f"You can use height {height - height % (self.vae_scale_factor * self.patch_size)} and width {width - width % (self.vae_scale_factor * self.patch_size)}." - ) - - 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_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_3 is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt_3`: {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)}") - elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)): - raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}") - - 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." - ) - elif negative_prompt_3 is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} 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`." - ) - - 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}") - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_latents - 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=device, dtype=dtype) - - 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 - - @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 - - @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, - prompt_2: Optional[Union[str, List[str]]] = None, - prompt_3: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 28, - sigmas: Optional[List[float]] = 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", - 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"], - max_sequence_length: int = 256, - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ): - 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 `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - will be used instead - 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): - 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. - guidance_scale (`float`, *optional*, defaults to 7.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 instead - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. - 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 256): Maximum sequence length to use with the `prompt`. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - Examples: - - Returns: - [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] 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_2, - prompt_3, - height, - width, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._joint_attention_kwargs = joint_attention_kwargs - self._interrupt = False - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale # - - # 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, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_3=prompt_3, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - do_classifier_free_guidance=self.do_classifier_free_guidance, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - device=device, - clip_skip=self.clip_skip, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - - if self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - pooled_prompt_embeds = self._prepare_perturbed_attention_guidance( - pooled_prompt_embeds, negative_pooled_prompt_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - 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) - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - # 5. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - if self.do_perturbed_attention_guidance: - original_attn_proc = self.transformer.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - - # 6. Denoising loop - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance, perturbed-attention guidance, or both - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - pooled_projections=pooled_prompt_embeds, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t - ) - - elif 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) - - # 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) - negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_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: - 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 self.do_perturbed_attention_guidance: - self.transformer.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image,) - - return StableDiffusion3PipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py deleted file mode 100644 index 202120dc2c2bbed4f1b41308ce6a125b6d30ad5b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py +++ /dev/null @@ -1,1056 +0,0 @@ -# Copyright 2025 Stability 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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import PIL.Image -import torch -from transformers import ( - CLIPTextModelWithProjection, - CLIPTokenizer, - T5EncoderModel, - T5TokenizerFast, -) - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, SD3LoraLoaderMixin -from ...models.attention_processor import PAGCFGJointAttnProcessor2_0, PAGJointAttnProcessor2_0 -from ...models.autoencoders import AutoencoderKL -from ...models.transformers import SD3Transformer2DModel -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 ..stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput -from .pag_utils import PAGMixin - - -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 StableDiffusion3PAGImg2ImgPipeline - >>> from diffusers.utils import load_image - - >>> pipe = StableDiffusion3PAGImg2ImgPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-3-medium-diffusers", - ... torch_dtype=torch.float16, - ... pag_applied_layers=["blocks.13"], - ... ) - >>> pipe.to("cuda") - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> url = "https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/aa_xl/000000009.png" - >>> init_image = load_image(url).convert("RGB") - >>> image = pipe(prompt, image=init_image, guidance_scale=5.0, pag_scale=0.7).images[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" -): - 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 StableDiffusion3PAGImg2ImgPipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, PAGMixin): - r""" - [PAG pipeline](https://huggingface.co/docs/diffusers/main/en/using-diffusers/pag) for image-to-image generation - using Stable Diffusion 3. - - Args: - transformer ([`SD3Transformer2DModel`]): - 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 ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModelWithProjection`]): - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), - specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant, - with an additional added projection layer that is initialized with a diagonal matrix with the `hidden_size` - as its dimension. - text_encoder_2 ([`CLIPTextModelWithProjection`]): - [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. - text_encoder_3 ([`T5EncoderModel`]): - Frozen text-encoder. Stable Diffusion 3 uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the - [t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) 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). - tokenizer_3 (`T5TokenizerFast`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->transformer->vae" - _optional_components = [] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "negative_pooled_prompt_embeds"] - - def __init__( - self, - transformer: SD3Transformer2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer_2: CLIPTokenizer, - text_encoder_3: T5EncoderModel, - tokenizer_3: T5TokenizerFast, - pag_applied_layers: Union[str, List[str]] = "blocks.1", # 1st transformer block - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - text_encoder_3=text_encoder_3, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - tokenizer_3=tokenizer_3, - 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.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = ( - self.transformer.config.sample_size - if hasattr(self, "transformer") and self.transformer is not None - else 128 - ) - self.patch_size = ( - self.transformer.config.patch_size if hasattr(self, "transformer") and self.transformer is not None else 2 - ) - - self.set_pag_applied_layers( - pag_applied_layers, pag_attn_processors=(PAGCFGJointAttnProcessor2_0(), PAGJointAttnProcessor2_0()) - ) - - # 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, - num_images_per_prompt: int = 1, - max_sequence_length: int = 256, - 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 self.text_encoder_3 is None: - return torch.zeros( - ( - batch_size * num_images_per_prompt, - self.tokenizer_max_length, - self.transformer.config.joint_attention_dim, - ), - device=device, - dtype=dtype, - ) - - text_inputs = self.tokenizer_3( - 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 - untruncated_ids = self.tokenizer_3(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_3.batch_decode(untruncated_ids[:, self.tokenizer_max_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_3(text_input_ids.to(device))[0] - - dtype = self.text_encoder_3.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clip_skip: Optional[int] = None, - clip_model_index: int = 0, - ): - device = device or self._execution_device - - clip_tokenizers = [self.tokenizer, self.tokenizer_2] - clip_text_encoders = [self.text_encoder, self.text_encoder_2] - - tokenizer = clip_tokenizers[clip_model_index] - text_encoder = clip_text_encoders[clip_model_index] - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer_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[:, self.tokenizer_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) - pooled_prompt_embeds = prompt_embeds[0] - - if clip_skip is None: - prompt_embeds = prompt_embeds.hidden_states[-2] - else: - prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)] - - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - _, 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(batch_size * num_images_per_prompt, seq_len, -1) - - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) - pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds, pooled_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, - 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, - max_sequence_length: int = 256, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - 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`): - 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 all the text-encoders. - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. - 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. - 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, SD3LoraLoaderMixin): - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - 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] - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - prompt_3 = prompt_3 or prompt - prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3 - - prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=clip_skip, - clip_model_index=0, - ) - prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds( - prompt=prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=clip_skip, - clip_model_index=1, - ) - clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1) - - t5_prompt_embed = self._get_t5_prompt_embeds( - prompt=prompt_3, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - clip_prompt_embeds = torch.nn.functional.pad( - clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1]) - ) - - prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2) - pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1) - - if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt_2 = negative_prompt_2 or negative_prompt - negative_prompt_3 = negative_prompt_3 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 - ) - negative_prompt_3 = ( - batch_size * [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3 - ) - - 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_embed, negative_pooled_prompt_embed = self._get_clip_prompt_embeds( - negative_prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=None, - clip_model_index=0, - ) - negative_prompt_2_embed, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds( - negative_prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=None, - clip_model_index=1, - ) - negative_clip_prompt_embeds = torch.cat([negative_prompt_embed, negative_prompt_2_embed], dim=-1) - - t5_negative_prompt_embed = self._get_t5_prompt_embeds( - prompt=negative_prompt_3, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - negative_clip_prompt_embeds = torch.nn.functional.pad( - negative_clip_prompt_embeds, - (0, t5_negative_prompt_embed.shape[-1] - negative_clip_prompt_embeds.shape[-1]), - ) - - negative_prompt_embeds = torch.cat([negative_clip_prompt_embeds, t5_negative_prompt_embed], dim=-2) - negative_pooled_prompt_embeds = torch.cat( - [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1 - ) - - if self.text_encoder is not None: - if isinstance(self, SD3LoraLoaderMixin) 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, SD3LoraLoaderMixin) 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 - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_img2img.StableDiffusion3Img2ImgPipeline.check_inputs - def check_inputs( - self, - prompt, - prompt_2, - prompt_3, - height, - width, - strength, - negative_prompt=None, - negative_prompt_2=None, - negative_prompt_3=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - ): - if ( - height % (self.vae_scale_factor * self.patch_size) != 0 - or width % (self.vae_scale_factor * self.patch_size) != 0 - ): - raise ValueError( - f"`height` and `width` have to be divisible by {self.vae_scale_factor * self.patch_size} but are {height} and {width}." - f"You can use height {height - height % (self.vae_scale_factor * self.patch_size)} and width {width - width % (self.vae_scale_factor * self.patch_size)}." - ) - - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should 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_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_3 is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt_3`: {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)}") - elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)): - raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}") - - 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." - ) - elif negative_prompt_3 is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} 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`." - ) - - 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}") - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_img2img.StableDiffusion3Img2ImgPipeline.prepare_latents - def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - if image.shape[1] == self.vae.config.latent_channels: - init_latents = image - - else: - 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." - ) - - elif isinstance(generator, list): - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - init_latents = (init_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # get latents - init_latents = self.scheduler.scale_noise(init_latents, timestep, noise) - latents = init_latents.to(device=device, dtype=dtype) - - return latents - - @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 - - @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, - prompt_2: Optional[Union[str, List[str]]] = None, - prompt_3: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - image: PipelineImageInput = None, - strength: float = 0.6, - num_inference_steps: int = 50, - sigmas: Optional[List[float]] = 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", - 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"], - max_sequence_length: int = 256, - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ): - 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 `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - will be used instead - 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`, 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. - strength (`float`, *optional*, defaults to 0.8): - 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`. - 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 7.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 instead - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. - 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 256): Maximum sequence length to use with the `prompt`. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - Examples: - - Returns: - [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] 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_2, - prompt_3, - height, - width, - strength, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._joint_attention_kwargs = joint_attention_kwargs - self._interrupt = False - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 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, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_3=prompt_3, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - do_classifier_free_guidance=self.do_classifier_free_guidance, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - device=device, - clip_skip=self.clip_skip, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - - if self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - pooled_prompt_embeds = self._prepare_perturbed_attention_guidance( - pooled_prompt_embeds, negative_pooled_prompt_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) - - # 3. Preprocess image - 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) - 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 - num_channels_latents = self.transformer.config.in_channels - if latents is None: - latents = self.prepare_latents( - image, - latent_timestep, - batch_size, - num_images_per_prompt, - prompt_embeds.dtype, - device, - generator, - ) - - if self.do_perturbed_attention_guidance: - original_attn_proc = self.transformer.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - - # 6. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - 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 - - # expand the latents if we are doing classifier free guidance, perturbed-attention guidance, or both - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - pooled_projections=pooled_prompt_embeds, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t - ) - - elif 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) - - # 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) - negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_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: - 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 self.do_perturbed_attention_guidance: - self.transformer.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image,) - - return StableDiffusion3PipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py deleted file mode 100644 index de13be9c4d2246e987897bb985f3375271959f90..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py +++ /dev/null @@ -1,878 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...image_processor import PipelineImageInput -from ...loaders import IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel, UNetMotionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...models.unets.unet_motion_model import MotionAdapter -from ...schedulers import 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 -from ...video_processor import VideoProcessor -from ..animatediff.pipeline_output import AnimateDiffPipelineOutput -from ..free_init_utils import FreeInitMixin -from ..free_noise_utils import AnimateDiffFreeNoiseMixin -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from .pag_utils import PAGMixin - - -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 AnimateDiffPAGPipeline, MotionAdapter, DDIMScheduler - >>> from diffusers.utils import export_to_gif - - >>> model_id = "SG161222/Realistic_Vision_V5.1_noVAE" - >>> motion_adapter_id = "guoyww/animatediff-motion-adapter-v1-5-2" - >>> motion_adapter = MotionAdapter.from_pretrained(motion_adapter_id) - >>> scheduler = DDIMScheduler.from_pretrained( - ... model_id, subfolder="scheduler", beta_schedule="linear", steps_offset=1, clip_sample=False - ... ) - >>> pipe = AnimateDiffPAGPipeline.from_pretrained( - ... model_id, - ... motion_adapter=motion_adapter, - ... scheduler=scheduler, - ... pag_applied_layers=["mid"], - ... torch_dtype=torch.float16, - ... ).to("cuda") - - >>> video = pipe( - ... prompt="car, futuristic cityscape with neon lights, street, no human", - ... negative_prompt="low quality, bad quality", - ... num_inference_steps=25, - ... guidance_scale=6.0, - ... pag_scale=3.0, - ... generator=torch.Generator().manual_seed(42), - ... ).frames[0] - - >>> export_to_gif(video, "animatediff_pag.gif") - ``` -""" - - -class AnimateDiffPAGPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - FreeInitMixin, - AnimateDiffFreeNoiseMixin, - PAGMixin, -): - r""" - Pipeline for text-to-video generation using - [AnimateDiff](https://huggingface.co/docs/diffusers/en/api/pipelines/animatediff) and [Perturbed Attention - Guidance](https://huggingface.co/docs/diffusers/en/using-diffusers/pag). - - 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.). - - 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 - - [`~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 ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer (`CLIPTokenizer`): - A [`~transformers.CLIPTokenizer`] to tokenize text. - unet ([`UNet2DConditionModel`]): - 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. - 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`]. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["feature_extractor", "image_encoder", "motion_adapter"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: Union[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"] - ): - super().__init__() - if isinstance(unet, UNet2DConditionModel): - unet = UNetMotionModel.from_unet2d(unet, motion_adapter) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - motion_adapter=motion_adapter, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - self.video_processor = VideoProcessor(do_resize=False, vae_scale_factor=self.vae_scale_factor) - - self.set_pag_applied_layers(pag_applied_layers) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt with num_images_per_prompt -> num_videos_per_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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.animatediff.pipeline_animatediff.AnimateDiffPipeline.decode_latents - def decode_latents(self, latents, decode_chunk_size: int = 16): - latents = 1 / self.vae.config.scaling_factor * latents - - batch_size, channels, num_frames, height, width = latents.shape - latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) - - video = [] - for i in range(0, latents.shape[0], decode_chunk_size): - batch_latents = latents[i : i + decode_chunk_size] - batch_latents = self.vae.decode(batch_latents).sample - video.append(batch_latents) - - video = torch.cat(video) - video = video[None, :].reshape((batch_size, num_frames, -1) + video.shape[2:]).permute(0, 2, 1, 3, 4) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - video = video.float() - return video - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.pia.pipeline_pia.PIAPipeline.check_inputs - def check_inputs( - self, - prompt, - height, - width, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_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_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 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.animatediff.pipeline_animatediff.AnimateDiffPipeline.prepare_latents - def prepare_latents( - self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None - ): - # If FreeNoise is enabled, generate latents as described in Equation (7) of [FreeNoise](https://huggingface.co/papers/2310.15169) - if self.free_noise_enabled: - latents = self._prepare_latents_free_noise( - batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents - ) - - 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." - ) - - shape = ( - batch_size, - num_channels_latents, - num_frames, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - @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, - 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, - 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", - 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"], - decode_chunk_size: int = 16, - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The width in pixels of the generated video. - num_frames (`int`, *optional*, defaults to 16): - The number of video frames that are generated. Defaults to 16 frames which at 8 frames per seconds - amounts to 2 seconds of video. - 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. - 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*): - 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*): - 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 video - 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`. Latents should be of shape - `(batch_size, num_channel, num_frames, height, width)`. - 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. - 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. - 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)`. 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. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated video. Choose between `torch.Tensor`, `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.text_to_video_synthesis.TextToVideoSDPipelineOutput`] instead - of a plain tuple. - 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). - 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. - 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. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - Examples: - - Returns: - [`~pipelines.animatediff.pipeline_output.AnimateDiffPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.animatediff.pipeline_output.AnimateDiffPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is a list with the generated frames. - """ - - # 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 - - num_videos_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 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 - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_videos_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - prompt_embeds = prompt_embeds.repeat_interleave(repeats=num_frames, dim=0) - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_videos_per_prompt, - self.do_classifier_free_guidance, - ) - - for i, image_embeds in enumerate(ip_adapter_image_embeds): - negative_image_embeds = None - if self.do_classifier_free_guidance: - negative_image_embeds, image_embeds = image_embeds.chunk(2) - if self.do_perturbed_attention_guidance: - image_embeds = self._prepare_perturbed_attention_guidance( - image_embeds, negative_image_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - image_embeds = image_embeds.to(device) - ip_adapter_image_embeds[i] = image_embeds - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - num_channels_latents, - num_frames, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": ip_adapter_image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - if self.do_perturbed_attention_guidance: - original_attn_proc = self.unet.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - - num_free_init_iters = self._free_init_num_iters if self.free_init_enabled else 1 - for free_init_iter in range(num_free_init_iters): - if self.free_init_enabled: - latents, timesteps = self._apply_free_init( - latents, free_init_iter, num_inference_steps, device, latents.dtype, generator - ) - - self._num_timesteps = len(timesteps) - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - # 8. Denoising loop - with self.progress_bar(total=self._num_timesteps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat( - [latents] * (prompt_embeds.shape[0] // num_frames // latents.shape[0]) - ) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - ).sample - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t - ) - elif self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - 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() - - # 9. Post processing - if output_type == "latent": - video = latents - else: - video_tensor = self.decode_latents(latents, decode_chunk_size) - video = self.video_processor.postprocess_video(video=video_tensor, output_type=output_type) - - # 10. Offload all models - self.maybe_free_model_hooks() - - if self.do_perturbed_attention_guidance: - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (video,) - - return AnimateDiffPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py deleted file mode 100644 index e9a846b5e2c48a8f21a32f04705e1fb5230b9f34..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py +++ /dev/null @@ -1,1110 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import PIL.Image -import torch -from packaging import version -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...configuration_utils import FrozenDict -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - 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, StableDiffusionMixin -from ..stable_diffusion.pipeline_output import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker -from .pag_utils import PAGMixin - - -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 AutoPipelineForImage2Image - >>> from diffusers.utils import load_image - - >>> pipe = AutoPipelineForImage2Image.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", - ... torch_dtype=torch.float16, - ... enable_pag=True, - ... ) - >>> pipe = pipe.to("cuda") - >>> url = "https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/aa_xl/000000009.png" - - >>> init_image = load_image(url).convert("RGB") - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> image = pipe(prompt, image=init_image, pag_scale=0.3).images[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" -): - 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 StableDiffusionPAGImg2ImgPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - FromSingleFileMixin, - PAGMixin, -): - r""" - Pipeline for text-guided image-to-image generation using Stable Diffusion. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - 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", # ["mid"], ["down.block_1", "up.block_0.attentions_0"] - ): - super().__init__() - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_config) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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.register_to_config(requires_safety_checker=requires_safety_checker) - - self.set_pag_applied_layers(pag_applied_layers) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - strength, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - ): - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should 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 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 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.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.prepare_latents - def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - 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." - ) - - elif isinstance(generator, list): - if image.shape[0] < batch_size and batch_size % image.shape[0] == 0: - image = torch.cat([image] * (batch_size // image.shape[0]), dim=0) - elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} " - ) - - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - deprecation_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial" - " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" - " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" - " your script to pass as many initial images as text prompts to suppress this warning." - ) - deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - latents = init_latents - - return latents - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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"], - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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`, 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. - strength (`float`, *optional*, defaults to 0.8): - 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`. - 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*): - 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*): - 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*): - 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): - The number of images to generate per prompt. - 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*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - 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. - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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). - 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. - 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. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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 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, - strength, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._interrupt = False - - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 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 - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - - for i, image_embeds in enumerate(ip_adapter_image_embeds): - negative_image_embeds = None - if self.do_classifier_free_guidance: - negative_image_embeds, image_embeds = image_embeds.chunk(2) - if self.do_perturbed_attention_guidance: - image_embeds = self._prepare_perturbed_attention_guidance( - image_embeds, negative_image_embeds, self.do_classifier_free_guidance - ) - - elif self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - image_embeds = image_embeds.to(device) - ip_adapter_image_embeds[i] = image_embeds - - # 4. Preprocess image - image = self.image_processor.preprocess(image) - - # 5. set timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, 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) - - # 6. Prepare latent variables - latents = self.prepare_latents( - image, - latent_timestep, - batch_size, - num_images_per_prompt, - prompt_embeds.dtype, - device, - generator, - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - # 7.2 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) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - if self.do_perturbed_attention_guidance: - original_attn_proc = self.unet.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - self._num_timesteps = len(timesteps) - - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - if ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = ip_adapter_image_embeds - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t - ) - elif 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) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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 not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 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 self.do_perturbed_attention_guidance: - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py deleted file mode 100644 index ee9d20f3682bdc84a791c72397e6f6b8a5db6911..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py +++ /dev/null @@ -1,1372 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import PIL.Image -import torch -from packaging import version -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...configuration_utils import FrozenDict -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AsymmetricAutoencoderKL, AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - 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, StableDiffusionMixin -from ..stable_diffusion.pipeline_output import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker -from .pag_utils import PAGMixin - - -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 AutoPipelineForInpainting - - >>> pipe = AutoPipelineForInpainting.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, enable_pag=True - ... ) - >>> pipe = pipe.to("cuda") - >>> 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" - >>> init_image = load_image(img_url).convert("RGB") - >>> mask_image = load_image(mask_url).convert("RGB") - >>> prompt = "A majestic tiger sitting on a bench" - >>> image = pipe( - ... prompt=prompt, - ... image=init_image, - ... mask_image=mask_image, - ... strength=0.8, - ... num_inference_steps=50, - ... guidance_scale=guidance_scale, - ... generator=generator, - ... pag_scale=pag_scale, - ... ).images[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" -): - 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.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 - - -# 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 StableDiffusionPAGInpaintPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - IPAdapterMixin, - FromSingleFileMixin, - PAGMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - 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", - ): - super().__init__() - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" - " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_config) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - self.set_pag_applied_layers(pag_applied_layers) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.StableDiffusionInpaintPipeline.check_inputs - def check_inputs( - self, - prompt, - image, - mask_image, - height, - width, - strength, - callback_steps, - output_type, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - padding_mask_crop=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 != 0 or width % self.vae_scale_factor != 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}." - ) - 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 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.pipeline_stable_diffusion_inpaint.StableDiffusionInpaintPipeline.prepare_latents - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=None, - image=None, - timestep=None, - is_strength_max=True, - return_noise=False, - return_image_latents=False, - ): - 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 (image is None or timestep is None) and not is_strength_max: - raise ValueError( - "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise." - "However, either the image or the noise timestep has not been provided." - ) - - if return_image_latents or (latents is None and not is_strength_max): - image = image.to(device=device, dtype=dtype) - - if image.shape[1] == 4: - image_latents = image - else: - image_latents = self._encode_vae_image(image=image, generator=generator) - image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1) - - if latents is None: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # if strength is 1. then initialise the latents to noise, else initial to image + noise - latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep) - # if pure noise then scale the initial latents by the Scheduler's init sigma - latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents - else: - noise = latents.to(device) - latents = noise * self.scheduler.init_noise_sigma - - outputs = (latents,) - - if return_noise: - outputs += (noise,) - - if return_image_latents: - outputs += (image_latents,) - - return outputs - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.StableDiffusionInpaintPipeline._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 = self.vae.config.scaling_factor * image_latents - - return image_latents - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.StableDiffusionInpaintPipeline.prepare_mask_latents - def prepare_mask_latents( - self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance - ): - # 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 // self.vae_scale_factor, width // self.vae_scale_factor) - ) - mask = mask.to(device=device, dtype=dtype) - - masked_image = masked_image.to(device=device, dtype=dtype) - - if masked_image.shape[1] == 4: - masked_image_latents = masked_image - else: - masked_image_latents = self._encode_vae_image(masked_image, generator=generator) - - # 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) - - mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask - masked_image_latents = ( - torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents - ) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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 - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_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, - masked_image_latents: torch.Tensor = None, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = None, - strength: float = 0.9999, - num_inference_steps: int = 50, - 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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"], - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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. - 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*): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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). - guidance_rescale (`float`, *optional*, defaults to 0.0): - Guidance rescale factor from [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. - 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. - 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. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - - # 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 - # to deal with lora scaling and other possible forward hooks - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - image, - mask_image, - height, - width, - strength, - None, - None, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - padding_mask_crop, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._interrupt = False - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 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 = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=lora_scale, - clip_skip=self.clip_skip, - ) - - # 4. set timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - timesteps, num_inference_steps = self.get_timesteps( - num_inference_steps=num_inference_steps, strength=strength, device=device - ) - # check that number of inference steps is not < 1 - as this doesn't make sense - 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) - # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise - is_strength_max = strength == 1.0 - - # 5. 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) - - # 6. Prepare latent variables - num_channels_latents = self.vae.config.latent_channels - num_channels_unet = self.unet.config.in_channels - return_image_latents = num_channels_unet == 4 - - latents_outputs = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - image=init_image, - timestep=latent_timestep, - is_strength_max=is_strength_max, - return_noise=True, - return_image_latents=return_image_latents, - ) - - if return_image_latents: - latents, noise, image_latents = latents_outputs - else: - latents, noise = latents_outputs - - # 7. Prepare mask latent variables - 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_images_per_prompt, - height, - width, - prompt_embeds.dtype, - device, - generator, - self.do_classifier_free_guidance, - ) - if self.do_perturbed_attention_guidance: - if self.do_classifier_free_guidance: - mask, _ = mask.chunk(2) - masked_image_latents, _ = masked_image_latents.chunk(2) - mask = self._prepare_perturbed_attention_guidance(mask, mask, self.do_classifier_free_guidance) - masked_image_latents = self._prepare_perturbed_attention_guidance( - masked_image_latents, masked_image_latents, self.do_classifier_free_guidance - ) - - # 8. Check that sizes of mask, masked image and latents match - if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting - num_channels_mask = mask.shape[1] - num_channels_masked_image = masked_image_latents.shape[1] - if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" - " `pipeline.unet` or your `mask_image` or `image` input." - ) - elif num_channels_unet != 4: - raise ValueError( - f"The unet {self.unet.__class__} should have either 4 or 9 input channels, not {self.unet.config.in_channels}." - ) - # 9 Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 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 self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - - for i, image_embeds in enumerate(ip_adapter_image_embeds): - negative_image_embeds = None - if self.do_classifier_free_guidance: - negative_image_embeds, image_embeds = image_embeds.chunk(2) - if self.do_perturbed_attention_guidance: - image_embeds = self._prepare_perturbed_attention_guidance( - image_embeds, negative_image_embeds, self.do_classifier_free_guidance - ) - - elif self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - image_embeds = image_embeds.to(device) - ip_adapter_image_embeds[i] = image_embeds - - # 9.1 Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": ip_adapter_image_embeds} - if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) - else None - ) - - # 9.2 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) - - # 10. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - if self.do_perturbed_attention_guidance: - original_attn_proc = self.unet.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - if num_channels_unet == 9: - latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t - ) - - elif 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.do_classifier_free_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - if num_channels_unet == 4: - init_latents_proper = image_latents - if self.do_perturbed_attention_guidance: - init_mask, *_ = mask.chunk(3) if self.do_classifier_free_guidance else mask.chunk(2) - else: - init_mask, *_ = mask.chunk(2) if self.do_classifier_free_guidance else mask - - if i < len(timesteps) - 1: - noise_timestep = timesteps[i + 1] - init_latents_proper = self.scheduler.add_noise( - init_latents_proper, noise, torch.tensor([noise_timestep]) - ) - - latents = (1 - init_mask) * init_latents_proper + init_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 not output_type == "latent": - condition_kwargs = {} - if isinstance(self.vae, AsymmetricAutoencoderKL): - init_image = init_image.to(device=device, dtype=masked_image_latents.dtype) - init_image_condition = init_image.clone() - init_image = self._encode_vae_image(init_image, generator=generator) - mask_condition = mask_condition.to(device=device, dtype=masked_image_latents.dtype) - condition_kwargs = {"image": init_image_condition, "mask": mask_condition} - image = self.vae.decode( - latents / self.vae.config.scaling_factor, return_dict=False, generator=generator, **condition_kwargs - )[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) - - if padding_mask_crop is not None: - image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] - - # Offload all models - self.maybe_free_model_hooks() - - if self.do_perturbed_attention_guidance: - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_xl.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_xl.py deleted file mode 100644 index 4504684133807859697c4e2e62c59772ab4d5ed7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_xl.py +++ /dev/null @@ -1,1354 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - FusedAttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - is_invisible_watermark_available, - 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, StableDiffusionMixin -from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput -from .pag_utils import PAGMixin - - -if is_invisible_watermark_available(): - from ..stable_diffusion_xl.watermark import StableDiffusionXLWatermarker - -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 AutoPipelineForText2Image - - >>> pipe = AutoPipelineForText2Image.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", - ... torch_dtype=torch.float16, - ... enable_pag=True, - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> image = pipe(prompt, pag_scale=0.3).images[0] - ``` -""" - - -# 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 - - -# 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 StableDiffusionXLPAGPipeline( - DiffusionPipeline, - StableDiffusionMixin, - FromSingleFileMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - PAGMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion XL. - - 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`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to - watermark output images. If not defined, it will default to True if the package is installed, otherwise no - watermarker will be used. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "image_encoder", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "negative_add_time_ids", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - 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"] - ): - super().__init__() - - 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, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - ) - 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 - ) - - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - self.set_pag_applied_layers(pag_applied_layers) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.check_inputs - def check_inputs( - self, - prompt, - prompt_2, - height, - width, - callback_steps, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_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_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`." - ) - - 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - 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): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - FusedAttnProcessor2_0, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_end(self): - return self._denoising_end - - @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, - prompt_2: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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"], - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ): - 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. - 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*): - 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. - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output) - 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. - 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*): - 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 ge 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. - 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)`. 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. - 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. - cross_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). - 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. - 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. - 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. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - 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, - None, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - self._interrupt = False - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 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 = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else 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 - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 5. 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, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 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_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - add_text_embeds = self._prepare_perturbed_attention_guidance( - add_text_embeds, negative_pooled_prompt_embeds, self.do_classifier_free_guidance - ) - add_time_ids = self._prepare_perturbed_attention_guidance( - add_time_ids, negative_add_time_ids, self.do_classifier_free_guidance - ) - - elif 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) - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - - for i, image_embeds in enumerate(ip_adapter_image_embeds): - negative_image_embeds = None - if self.do_classifier_free_guidance: - negative_image_embeds, image_embeds = image_embeds.chunk(2) - - if self.do_perturbed_attention_guidance: - image_embeds = self._prepare_perturbed_attention_guidance( - image_embeds, negative_image_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - image_embeds = image_embeds.to(device) - ip_adapter_image_embeds[i] = image_embeds - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - # 8.1 Apply denoising_end - if ( - self.denoising_end is not None - and isinstance(self.denoising_end, float) - and self.denoising_end > 0 - and self.denoising_end < 1 - ): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - # 9. 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) - - if self.do_perturbed_attention_guidance: - original_attn_proc = self.unet.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - - 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 - - # expand the latents if we are doing classifier free guidance, perturbed-attention guidance, or both - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - if ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = ip_adapter_image_embeds - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred, noise_pred_text = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t, True - ) - - elif 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.do_classifier_free_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents_dtype = latents.dtype - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds - ) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids) - - 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 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) - elif latents.dtype != self.vae.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 - self.vae = self.vae.to(latents.dtype) - - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None - has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None - if has_latents_mean and has_latents_std: - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean - else: - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, 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": - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - image = self.image_processor.postprocess(image, output_type=output_type) - - # Offload all models - self.maybe_free_model_hooks() - - if self.do_perturbed_attention_guidance: - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image,) - - return StableDiffusionXLPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py deleted file mode 100644 index 8c355a5fb129e80530e65622d1e806a874943e26..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py +++ /dev/null @@ -1,1549 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import PIL.Image -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - is_invisible_watermark_available, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import empty_device_cache, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput -from .pag_utils import PAGMixin - - -if is_invisible_watermark_available(): - from ..stable_diffusion_xl.watermark import StableDiffusionXLWatermarker - -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 AutoPipelineForImage2Image - >>> from diffusers.utils import load_image - - >>> pipe = AutoPipelineForImage2Image.from_pretrained( - ... "stabilityai/stable-diffusion-xl-refiner-1.0", - ... torch_dtype=torch.float16, - ... enable_pag=True, - ... ) - >>> pipe = pipe.to("cuda") - >>> url = "https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/aa_xl/000000009.png" - - >>> init_image = load_image(url).convert("RGB") - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> image = pipe(prompt, image=init_image, pag_scale=0.3).images[0] - ``` -""" - - -# 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 - - -# 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" -): - 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 StableDiffusionXLPAGImg2ImgPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - FromSingleFileMixin, - StableDiffusionXLLoraLoaderMixin, - IPAdapterMixin, - PAGMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion XL. - - 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`]. - requires_aesthetics_score (`bool`, *optional*, defaults to `"False"`): - Whether the `unet` requires an `aesthetic_score` condition to be passed during inference. Also see the - config of `stabilityai/stable-diffusion-xl-refiner-1-0`. - 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`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to - watermark output images. If not defined, it will default to True if the package is installed, otherwise no - watermarker will be used. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "image_encoder", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "add_neg_time_ids", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - image_encoder: CLIPVisionModelWithProjection = None, - 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"] - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - scheduler=scheduler, - ) - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - self.register_to_config(requires_aesthetics_score=requires_aesthetics_score) - 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) - - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - self.set_pag_applied_layers(pag_applied_layers) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.check_inputs - def check_inputs( - self, - prompt, - prompt_2, - strength, - num_inference_steps, - callback_steps, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - ): - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - if num_inference_steps is None: - raise ValueError("`num_inference_steps` cannot be None.") - elif not isinstance(num_inference_steps, int) or num_inference_steps <= 0: - raise ValueError( - f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type" - f" {type(num_inference_steps)}." - ) - 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_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 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_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device, denoising_start=None): - # get the original timestep using init_timestep - if denoising_start is None: - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - t_start = 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 - - else: - # Strength is irrelevant if we directly request a timestep to start at; - # that is, strength is determined by the denoising_start instead. - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (denoising_start * self.scheduler.config.num_train_timesteps) - ) - ) - - num_inference_steps = (self.scheduler.timesteps < discrete_timestep_cutoff).sum().item() - if self.scheduler.order == 2 and num_inference_steps % 2 == 0: - # if the scheduler is a 2nd order scheduler we might have to do +1 - # because `num_inference_steps` might be even given that every timestep - # (except the highest one) is duplicated. If `num_inference_steps` is even it would - # mean that we cut the timesteps in the middle of the denoising step - # (between 1st and 2nd derivative) which leads to incorrect results. By adding 1 - # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler - num_inference_steps = num_inference_steps + 1 - - # because t_n+1 >= t_n, we slice the timesteps starting from the end - t_start = len(self.scheduler.timesteps) - num_inference_steps - timesteps = self.scheduler.timesteps[t_start:] - if hasattr(self.scheduler, "set_begin_index"): - self.scheduler.set_begin_index(t_start) - return timesteps, num_inference_steps - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.prepare_latents - def prepare_latents( - self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None, add_noise=True - ): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - latents_mean = latents_std = None - if hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None: - latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1) - if hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None: - latents_std = torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1) - - # Offload text encoder if `enable_model_cpu_offload` was enabled - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.text_encoder_2.to("cpu") - empty_device_cache() - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - # make sure the VAE is in float32 mode, as it overflows in float16 - if self.vae.config.force_upcast: - image = image.float() - self.vae.to(dtype=torch.float32) - - 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." - ) - - elif isinstance(generator, list): - if image.shape[0] < batch_size and batch_size % image.shape[0] == 0: - image = torch.cat([image] * (batch_size // image.shape[0]), dim=0) - elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} " - ) - - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - if self.vae.config.force_upcast: - self.vae.to(dtype) - - init_latents = init_latents.to(dtype) - if latents_mean is not None and latents_std is not None: - latents_mean = latents_mean.to(device=device, dtype=dtype) - latents_std = latents_std.to(device=device, dtype=dtype) - init_latents = (init_latents - latents_mean) * self.vae.config.scaling_factor / latents_std - else: - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - if add_noise: - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - - latents = init_latents - - return latents - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline._get_add_time_ids - def _get_add_time_ids( - self, - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype, - text_encoder_projection_dim=None, - ): - if self.config.requires_aesthetics_score: - add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,)) - add_neg_time_ids = list( - negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,) - ) - else: - add_time_ids = list(original_size + crops_coords_top_left + target_size) - add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_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 - and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model." - ) - elif ( - expected_add_embed_dim < passed_add_embed_dim - and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model." - ) - elif 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) - add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype) - - return add_time_ids, add_neg_time_ids - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_end(self): - return self._denoising_end - - @property - def denoising_start(self): - return self._denoising_start - - @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, - prompt_2: Optional[Union[str, List[str]]] = 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, - 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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, - 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"], - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ): - 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 - 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` - 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`. Note that in the case of - `denoising_start` being declared as an integer, the value of `strength` will be ignored. - 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. - 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. - denoising_start (`float`, *optional*): - When specified, indicates the fraction (between 0.0 and 1.0) of the total denoising process to be - bypassed before it is initiated. Consequently, the initial part of the denoising process is skipped and - it is assumed that the passed `image` is a partly denoised image. Note that when this is specified, - strength will be ignored. The `denoising_start` parameter is particularly beneficial when this pipeline - is integrated into a "Mixture of Denoisers" multi-pipeline setup, as detailed in [**Refine Image - Quality**](https://huggingface.co/docs/diffusers/using-diffusers/sdxl#refine-image-quality). - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - still retain a substantial amount of noise (ca. final 20% of timesteps still needed) and should be - denoised by a successor pipeline that has `denoising_start` set to 0.8 so that it only denoises the - final 20% of 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 [**Refine Image - Quality**](https://huggingface.co/docs/diffusers/using-diffusers/sdxl#refine-image-quality). - 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`). - 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. - 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*): - 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 ge 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. - 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)`. 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. - 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.StableDiffusionXLPipelineOutput`] instead of a - plain tuple. - cross_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). - 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. - 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. - aesthetic_score (`float`, *optional*, defaults to 6.0): - Used to simulate an aesthetic score of the generated image by influencing the positive text condition. - 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_aesthetic_score (`float`, *optional*, defaults to 2.5): - 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). Can be used to - simulate an aesthetic score of the generated image by influencing the negative text condition. - 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. - 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. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a - `tuple. When returning a tuple, 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_2, - strength, - num_inference_steps, - None, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - self._denoising_start = denoising_start - self._interrupt = False - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 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 - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else 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=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - - # 4. Preprocess image - image = self.image_processor.preprocess(image) - - # 5. Prepare timesteps - def denoising_value_valid(dnv): - return isinstance(dnv, float) and 0 < dnv < 1 - - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - timesteps, num_inference_steps = self.get_timesteps( - num_inference_steps, - strength, - device, - denoising_start=self.denoising_start if denoising_value_valid(self.denoising_start) else None, - ) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - - add_noise = True if self.denoising_start is None else False - - # 6. Prepare latent variables - if latents is None: - latents = self.prepare_latents( - image, - latent_timestep, - batch_size, - num_images_per_prompt, - prompt_embeds.dtype, - device, - generator, - add_noise, - ) - # 7. Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - height, width = latents.shape[-2:] - height = height * self.vae_scale_factor - width = width * self.vae_scale_factor - - original_size = original_size or (height, width) - target_size = target_size or (height, width) - - # 8. Prepare added time ids & embeddings - if negative_original_size is None: - negative_original_size = original_size - if negative_target_size is None: - negative_target_size = target_size - - 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, add_neg_time_ids = self._get_add_time_ids( - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype=prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1) - add_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1) - - if self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - add_text_embeds = self._prepare_perturbed_attention_guidance( - add_text_embeds, negative_pooled_prompt_embeds, self.do_classifier_free_guidance - ) - add_time_ids = self._prepare_perturbed_attention_guidance( - add_time_ids, add_neg_time_ids, self.do_classifier_free_guidance - ) - elif 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([add_neg_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) - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - for i, image_embeds in enumerate(ip_adapter_image_embeds): - negative_image_embeds = None - if self.do_classifier_free_guidance: - negative_image_embeds, image_embeds = image_embeds.chunk(2) - - if self.do_perturbed_attention_guidance: - image_embeds = self._prepare_perturbed_attention_guidance( - image_embeds, negative_image_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - image_embeds = image_embeds.to(device) - ip_adapter_image_embeds[i] = image_embeds - - # 9. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - # 9.1 Apply denoising_end - if ( - self.denoising_end is not None - and self.denoising_start is not None - and denoising_value_valid(self.denoising_end) - and denoising_value_valid(self.denoising_start) - and self.denoising_start >= self.denoising_end - ): - raise ValueError( - f"`denoising_start`: {self.denoising_start} cannot be larger than or equal to `denoising_end`: " - + f" {self.denoising_end} when using type float." - ) - elif self.denoising_end is not None and denoising_value_valid(self.denoising_end): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - # 9.2 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) - - if self.do_perturbed_attention_guidance: - original_attn_proc = self.unet.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - if ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = ip_adapter_image_embeds - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred, noise_pred_text = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t, True - ) - elif 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.do_classifier_free_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents_dtype = latents.dtype - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds - ) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - add_neg_time_ids = callback_outputs.pop("add_neg_time_ids", add_neg_time_ids) - - # 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 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) - elif latents.dtype != self.vae.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 - self.vae = self.vae.to(latents.dtype) - - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None - has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None - if has_latents_mean and has_latents_std: - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean - else: - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, return_dict=False)[0] - - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - else: - image = latents - - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - image = self.image_processor.postprocess(image, output_type=output_type) - - # Offload all models - self.maybe_free_model_hooks() - - if self.do_perturbed_attention_guidance: - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image,) - - return StableDiffusionXLPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py deleted file mode 100644 index 7d42d1876a8267e499a42a99d093383b6406120b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py +++ /dev/null @@ -1,1780 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import PIL.Image -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - is_invisible_watermark_available, - 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, StableDiffusionMixin -from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput -from .pag_utils import PAGMixin - - -if is_invisible_watermark_available(): - from ..stable_diffusion_xl.watermark import StableDiffusionXLWatermarker - -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 AutoPipelineForInpainting - >>> from diffusers.utils import load_image - - >>> pipe = AutoPipelineForInpainting.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", - ... torch_dtype=torch.float16, - ... variant="fp16", - ... enable_pag=True, - ... ) - >>> pipe.to("cuda") - - >>> 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" - - >>> init_image = load_image(img_url).convert("RGB") - >>> mask_image = load_image(mask_url).convert("RGB") - - >>> prompt = "A majestic tiger sitting on a bench" - >>> image = pipe( - ... prompt=prompt, - ... image=init_image, - ... mask_image=mask_image, - ... num_inference_steps=50, - ... strength=0.80, - ... pag_scale=0.3, - ... ).images[0] - ``` -""" - - -# 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 - - -# 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" -): - 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 StableDiffusionXLPAGInpaintPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionXLLoraLoaderMixin, - FromSingleFileMixin, - IPAdapterMixin, - PAGMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion XL. - - 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`]. - requires_aesthetics_score (`bool`, *optional*, defaults to `"False"`): - Whether the `unet` requires a aesthetic_score condition to be passed during inference. Also see the config - of `stabilityai/stable-diffusion-xl-refiner-1-0`. - 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`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to - watermark output images. If not defined, it will default to True if the package is installed, otherwise no - watermarker will be used. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae" - - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "image_encoder", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "add_neg_time_ids", - "mask", - "masked_image_latents", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - image_encoder: CLIPVisionModelWithProjection = None, - 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"] - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - scheduler=scheduler, - ) - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - self.register_to_config(requires_aesthetics_score=requires_aesthetics_score) - 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.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True - ) - - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - self.set_pag_applied_layers(pag_applied_layers) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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_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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipeline.check_inputs - def check_inputs( - self, - prompt, - prompt_2, - image, - mask_image, - height, - width, - strength, - callback_steps, - output_type, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - padding_mask_crop=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 % 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_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 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 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_xl.pipeline_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipeline.prepare_latents - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=None, - image=None, - timestep=None, - is_strength_max=True, - add_noise=True, - return_noise=False, - return_image_latents=False, - ): - 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 (image is None or timestep is None) and not is_strength_max: - raise ValueError( - "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise." - "However, either the image or the noise timestep has not been provided." - ) - - if image.shape[1] == 4: - image_latents = image.to(device=device, dtype=dtype) - image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1) - elif return_image_latents or (latents is None and not is_strength_max): - image = image.to(device=device, dtype=dtype) - image_latents = self._encode_vae_image(image=image, generator=generator) - image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1) - - if latents is None and add_noise: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # if strength is 1. then initialise the latents to noise, else initial to image + noise - latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep) - # if pure noise then scale the initial latents by the Scheduler's init sigma - latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents - elif add_noise: - noise = latents.to(device) - latents = noise * self.scheduler.init_noise_sigma - else: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = image_latents.to(device) - - outputs = (latents,) - - if return_noise: - outputs += (noise,) - - if return_image_latents: - outputs += (image_latents,) - - return outputs - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipeline._encode_vae_image - def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): - dtype = image.dtype - if self.vae.config.force_upcast: - image = image.float() - self.vae.to(dtype=torch.float32) - - 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) - - if self.vae.config.force_upcast: - self.vae.to(dtype) - - image_latents = image_latents.to(dtype) - image_latents = self.vae.config.scaling_factor * image_latents - - return image_latents - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipeline.prepare_mask_latents - def prepare_mask_latents( - self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance - ): - # 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 // self.vae_scale_factor, width // self.vae_scale_factor) - ) - mask = mask.to(device=device, dtype=dtype) - - # 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) - - mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask - - if masked_image is not None and masked_image.shape[1] == 4: - masked_image_latents = masked_image - else: - masked_image_latents = None - - if masked_image is not None: - if masked_image_latents is None: - masked_image = masked_image.to(device=device, dtype=dtype) - masked_image_latents = self._encode_vae_image(masked_image, generator=generator) - - 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 - ) - - masked_image_latents = ( - torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents - ) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device, denoising_start=None): - # get the original timestep using init_timestep - if denoising_start is None: - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - t_start = 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 - - else: - # Strength is irrelevant if we directly request a timestep to start at; - # that is, strength is determined by the denoising_start instead. - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (denoising_start * self.scheduler.config.num_train_timesteps) - ) - ) - - num_inference_steps = (self.scheduler.timesteps < discrete_timestep_cutoff).sum().item() - if self.scheduler.order == 2 and num_inference_steps % 2 == 0: - # if the scheduler is a 2nd order scheduler we might have to do +1 - # because `num_inference_steps` might be even given that every timestep - # (except the highest one) is duplicated. If `num_inference_steps` is even it would - # mean that we cut the timesteps in the middle of the denoising step - # (between 1st and 2nd derivative) which leads to incorrect results. By adding 1 - # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler - num_inference_steps = num_inference_steps + 1 - - # because t_n+1 >= t_n, we slice the timesteps starting from the end - t_start = len(self.scheduler.timesteps) - num_inference_steps - timesteps = self.scheduler.timesteps[t_start:] - if hasattr(self.scheduler, "set_begin_index"): - self.scheduler.set_begin_index(t_start) - return timesteps, num_inference_steps - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline._get_add_time_ids - def _get_add_time_ids( - self, - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype, - text_encoder_projection_dim=None, - ): - if self.config.requires_aesthetics_score: - add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,)) - add_neg_time_ids = list( - negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,) - ) - else: - add_time_ids = list(original_size + crops_coords_top_left + target_size) - add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_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 - and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model." - ) - elif ( - expected_add_embed_dim < passed_add_embed_dim - and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model." - ) - elif 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) - add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype) - - return add_time_ids, add_neg_time_ids - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_end(self): - return self._denoising_end - - @property - def denoising_start(self): - return self._denoising_start - - @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, - prompt_2: Optional[Union[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, - 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, - 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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, - 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"], - pag_scale: float = 3.0, - pag_adaptive_scale: float = 0.0, - ): - 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 - image (`PIL.Image.Image`): - `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will - be masked out with `mask_image` and repainted according to `prompt`. - mask_image (`PIL.Image.Image`): - `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be - 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)`. - 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. - padding_mask_crop (`int`, *optional*, defaults to `None`): - The size of margin in the crop to be applied to the image and masking. If `None`, no crop is applied to - image and mask_image. If `padding_mask_crop` is not `None`, it will first find a rectangular region - with the same aspect ration of the image and contains all masked area, and then expand that area based - on `padding_mask_crop`. The image and mask_image will then be cropped based on the expanded area before - resizing to the original image size for inpainting. This is useful when the masked area is small while - the image is large and contain information irrelevant for inpainting, such as background. - strength (`float`, *optional*, defaults to 0.9999): - Conceptually, indicates how much to transform the masked portion of 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 the masked - portion of the reference `image`. Note that in the case of `denoising_start` being declared as an - integer, the value of `strength` will be ignored. - 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. - 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. - denoising_start (`float`, *optional*): - When specified, indicates the fraction (between 0.0 and 1.0) of the total denoising process to be - bypassed before it is initiated. Consequently, the initial part of the denoising process is skipped and - it is assumed that the passed `image` is a partly denoised image. Note that when this is specified, - strength will be ignored. The `denoising_start` parameter is particularly beneficial when this pipeline - is integrated into a "Mixture of Denoisers" multi-pipeline setup, as detailed in [**Refining the Image - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output). - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - still retain a substantial amount of noise (ca. final 20% of timesteps still needed) and should be - denoised by a successor pipeline that has `denoising_start` set to 0.8 so that it only denoises the - final 20% of 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output). - 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`). - 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. - 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)`. 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. - 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 ge generated by sampling using the supplied random `generator`. - 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. - cross_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). - 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. - aesthetic_score (`float`, *optional*, defaults to 6.0): - Used to simulate an aesthetic score of the generated image by influencing the positive text condition. - 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_aesthetic_score (`float`, *optional*, defaults to 2.5): - 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). Can be used to - simulate an aesthetic score of the generated image by influencing the negative text condition. - 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. - 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. - pag_scale (`float`, *optional*, defaults to 3.0): - The scale factor for the perturbed attention guidance. If it is set to 0.0, the perturbed attention - guidance will not be used. - pag_adaptive_scale (`float`, *optional*, defaults to 0.0): - The adaptive scale factor for the perturbed attention guidance. If it is set to 0.0, `pag_scale` is - used. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a - `tuple. `tuple. When returning a tuple, 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 - - # 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 - self.check_inputs( - prompt, - prompt_2, - image, - mask_image, - height, - width, - strength, - None, - output_type, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - padding_mask_crop, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - self._denoising_start = denoising_start - self._interrupt = False - self._pag_scale = pag_scale - self._pag_adaptive_scale = pag_adaptive_scale - - # 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 - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else 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=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - - # 4. set timesteps - def denoising_value_valid(dnv): - return isinstance(dnv, float) and 0 < dnv < 1 - - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - timesteps, num_inference_steps = self.get_timesteps( - num_inference_steps, - strength, - device, - denoising_start=self.denoising_start if denoising_value_valid(self.denoising_start) else None, - ) - # check that number of inference steps is not < 1 - as this doesn't make sense - 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." - ) - # at which timestep to set the initial noise (n.b. 50% if strength is 0.5) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise - is_strength_max = strength == 1.0 - - # 5. 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) - - mask = self.mask_processor.preprocess( - mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords - ) - - if masked_image_latents is not None: - masked_image = masked_image_latents - elif init_image.shape[1] == 4: - # if images are in latent space, we can't mask it - masked_image = None - else: - masked_image = init_image * (mask < 0.5) - - # 6. Prepare latent variables - num_channels_latents = self.vae.config.latent_channels - num_channels_unet = self.unet.config.in_channels - return_image_latents = num_channels_unet == 4 - - add_noise = True if self.denoising_start is None else False - latents_outputs = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - image=init_image, - timestep=latent_timestep, - is_strength_max=is_strength_max, - add_noise=add_noise, - return_noise=True, - return_image_latents=return_image_latents, - ) - - if return_image_latents: - latents, noise, image_latents = latents_outputs - else: - latents, noise = latents_outputs - - # 7. Prepare mask latent variables - mask, masked_image_latents = self.prepare_mask_latents( - mask, - masked_image, - batch_size * num_images_per_prompt, - height, - width, - prompt_embeds.dtype, - device, - generator, - self.do_classifier_free_guidance, - ) - if self.do_perturbed_attention_guidance: - if self.do_classifier_free_guidance: - mask, _ = mask.chunk(2) - masked_image_latents, _ = masked_image_latents.chunk(2) - mask = self._prepare_perturbed_attention_guidance(mask, mask, self.do_classifier_free_guidance) - masked_image_latents = self._prepare_perturbed_attention_guidance( - masked_image_latents, masked_image_latents, self.do_classifier_free_guidance - ) - - # 8. Check that sizes of mask, masked image and latents match - if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting - num_channels_mask = mask.shape[1] - num_channels_masked_image = masked_image_latents.shape[1] - if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" - " `pipeline.unet` or your `mask_image` or `image` input." - ) - elif num_channels_unet != 4: - raise ValueError( - f"The unet {self.unet.__class__} should have either 4 or 9 input channels, not {self.unet.config.in_channels}." - ) - # 8.1 Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - height, width = latents.shape[-2:] - height = height * self.vae_scale_factor - width = width * self.vae_scale_factor - - original_size = original_size or (height, width) - target_size = target_size or (height, width) - - # 10. Prepare added time ids & embeddings - if negative_original_size is None: - negative_original_size = original_size - if negative_target_size is None: - negative_target_size = target_size - - 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, add_neg_time_ids = self._get_add_time_ids( - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype=prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1) - add_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1) - - if self.do_perturbed_attention_guidance: - prompt_embeds = self._prepare_perturbed_attention_guidance( - prompt_embeds, negative_prompt_embeds, self.do_classifier_free_guidance - ) - add_text_embeds = self._prepare_perturbed_attention_guidance( - add_text_embeds, negative_pooled_prompt_embeds, self.do_classifier_free_guidance - ) - add_time_ids = self._prepare_perturbed_attention_guidance( - add_time_ids, add_neg_time_ids, self.do_classifier_free_guidance - ) - - elif 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([add_neg_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) - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - for i, image_embeds in enumerate(ip_adapter_image_embeds): - negative_image_embeds = None - if self.do_classifier_free_guidance: - negative_image_embeds, image_embeds = image_embeds.chunk(2) - - if self.do_perturbed_attention_guidance: - image_embeds = self._prepare_perturbed_attention_guidance( - image_embeds, negative_image_embeds, self.do_classifier_free_guidance - ) - elif self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - image_embeds = image_embeds.to(device) - ip_adapter_image_embeds[i] = image_embeds - - # 11. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - if ( - self.denoising_end is not None - and self.denoising_start is not None - and denoising_value_valid(self.denoising_end) - and denoising_value_valid(self.denoising_start) - and self.denoising_start >= self.denoising_end - ): - raise ValueError( - f"`denoising_start`: {self.denoising_start} cannot be larger than or equal to `denoising_end`: " - + f" {self.denoising_end} when using type float." - ) - elif self.denoising_end is not None and denoising_value_valid(self.denoising_end): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - # 11.1 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) - - if self.do_perturbed_attention_guidance: - original_attn_proc = self.unet.attn_processors - self._set_pag_attn_processor( - pag_applied_layers=self.pag_applied_layers, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - - 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 - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * (prompt_embeds.shape[0] // latents.shape[0])) - - # concat latents, mask, masked_image_latents in the channel dimension - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - if num_channels_unet == 9: - latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) - - # predict the noise residual - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = ip_adapter_image_embeds - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_perturbed_attention_guidance: - noise_pred, noise_pred_text = self._apply_perturbed_attention_guidance( - noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t, True - ) - elif 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.do_classifier_free_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents_dtype = latents.dtype - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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 num_channels_unet == 4: - init_latents_proper = image_latents - if self.do_perturbed_attention_guidance: - init_mask, *_ = mask.chunk(3) if self.do_classifier_free_guidance else mask.chunk(2) - else: - init_mask, *_ = mask.chunk(2) if self.do_classifier_free_guidance else mask - - if i < len(timesteps) - 1: - noise_timestep = timesteps[i + 1] - init_latents_proper = self.scheduler.add_noise( - init_latents_proper, noise, torch.tensor([noise_timestep]) - ) - - latents = (1 - init_mask) * init_latents_proper + init_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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds - ) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - add_neg_time_ids = callback_outputs.pop("add_neg_time_ids", add_neg_time_ids) - 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 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) - elif latents.dtype != self.vae.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 - self.vae = self.vae.to(latents.dtype) - - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None - has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None - if has_latents_mean and has_latents_std: - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean - else: - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, return_dict=False)[0] - - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - else: - return StableDiffusionXLPipelineOutput(images=latents) - - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - image = self.image_processor.postprocess(image, output_type=output_type) - - if padding_mask_crop is not None: - image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] - - # Offload all models - self.maybe_free_model_hooks() - - if self.do_perturbed_attention_guidance: - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image,) - - return StableDiffusionXLPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/__init__.py deleted file mode 100644 index aaa775f690c3d290074662c029f242df3c61e003..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/__init__.py +++ /dev/null @@ -1,55 +0,0 @@ -from dataclasses import dataclass -from typing import TYPE_CHECKING, List, Optional, Union - -import numpy as np -import PIL -from PIL import Image - -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["image_encoder"] = ["PaintByExampleImageEncoder"] - _import_structure["pipeline_paint_by_example"] = ["PaintByExamplePipeline"] - - -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 .image_encoder import PaintByExampleImageEncoder - from .pipeline_paint_by_example import PaintByExamplePipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 042a2342ab1701dd487730125a473f279d32eb08..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/__pycache__/image_encoder.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/__pycache__/image_encoder.cpython-310.pyc deleted file mode 100644 index 46df6836e9661daba7f70c0ff79bea9ed376f283..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/__pycache__/image_encoder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/__pycache__/pipeline_paint_by_example.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/__pycache__/pipeline_paint_by_example.cpython-310.pyc deleted file mode 100644 index fef0df43bdc8826515f010512e40b8ce817d2175..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/__pycache__/pipeline_paint_by_example.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/image_encoder.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/image_encoder.py deleted file mode 100644 index 74c575ed8653d110ca300e9198a5b13070659adf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/image_encoder.py +++ /dev/null @@ -1,67 +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 torch -from torch import nn -from transformers import CLIPPreTrainedModel, CLIPVisionModel - -from ...models.attention import BasicTransformerBlock -from ...utils import logging - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class PaintByExampleImageEncoder(CLIPPreTrainedModel): - def __init__(self, config, proj_size=None): - super().__init__(config) - self.proj_size = proj_size or getattr(config, "projection_dim", 768) - - self.model = CLIPVisionModel(config) - self.mapper = PaintByExampleMapper(config) - self.final_layer_norm = nn.LayerNorm(config.hidden_size) - self.proj_out = nn.Linear(config.hidden_size, self.proj_size) - - # uncondition for scaling - self.uncond_vector = nn.Parameter(torch.randn((1, 1, self.proj_size))) - - def forward(self, pixel_values, return_uncond_vector=False): - clip_output = self.model(pixel_values=pixel_values) - latent_states = clip_output.pooler_output - latent_states = self.mapper(latent_states[:, None]) - latent_states = self.final_layer_norm(latent_states) - latent_states = self.proj_out(latent_states) - if return_uncond_vector: - return latent_states, self.uncond_vector - - return latent_states - - -class PaintByExampleMapper(nn.Module): - def __init__(self, config): - super().__init__() - num_layers = (config.num_hidden_layers + 1) // 5 - hid_size = config.hidden_size - num_heads = 1 - self.blocks = nn.ModuleList( - [ - BasicTransformerBlock(hid_size, num_heads, hid_size, activation_fn="gelu", attention_bias=True) - for _ in range(num_layers) - ] - ) - - def forward(self, hidden_states): - for block in self.blocks: - hidden_states = block(hidden_states) - - return hidden_states diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py deleted file mode 100644 index 3e22c9a845450c4159c880509ed92ff415034c67..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py +++ /dev/null @@ -1,637 +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 inspect -from typing import Callable, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import CLIPImageProcessor - -from ...image_processor import VaeImageProcessor -from ...models import AutoencoderKL, UNet2DConditionModel -from ...schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler -from ...utils import deprecate, is_torch_xla_available, logging -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker -from .image_encoder import PaintByExampleImageEncoder - - -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 - - -# 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" -): - 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") - - -def prepare_mask_and_masked_image(image, mask): - """ - Prepares a pair (image, mask) to be consumed by the Paint by Example pipeline. This means that those inputs will be - converted to ``torch.Tensor`` with shapes ``batch x channels x height x width`` where ``channels`` is ``3`` for the - ``image`` and ``1`` for the ``mask``. - - The ``image`` will be converted to ``torch.float32`` and normalized to be in ``[-1, 1]``. The ``mask`` will be - binarized (``mask > 0.5``) and cast to ``torch.float32`` too. - - Args: - image (Union[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. - It can be a ``PIL.Image``, or a ``height x width`` ``np.array`` or a ``1 x height x width`` - ``torch.Tensor`` or a ``batch x 1 x height x width`` ``torch.Tensor``. - - - Raises: - ValueError: ``torch.Tensor`` images should be in the ``[-1, 1]`` range. ValueError: ``torch.Tensor`` mask - should be in the ``[0, 1]`` range. ValueError: ``mask`` and ``image`` should have the same spatial dimensions. - TypeError: ``mask`` is a ``torch.Tensor`` but ``image`` is not - (ot the other way around). - - Returns: - tuple[torch.Tensor]: The pair (mask, masked_image) as ``torch.Tensor`` with 4 - dimensions: ``batch x channels x height x width``. - """ - if isinstance(image, torch.Tensor): - if not isinstance(mask, torch.Tensor): - raise TypeError(f"`image` is a torch.Tensor but `mask` (type: {type(mask)} is not") - - # Batch single image - if image.ndim == 3: - assert image.shape[0] == 3, "Image outside a batch should be of shape (3, H, W)" - image = image.unsqueeze(0) - - # Batch and add channel dim for single mask - if mask.ndim == 2: - mask = mask.unsqueeze(0).unsqueeze(0) - - # Batch single mask or add channel dim - if mask.ndim == 3: - # Batched mask - if mask.shape[0] == image.shape[0]: - mask = mask.unsqueeze(1) - else: - mask = mask.unsqueeze(0) - - assert image.ndim == 4 and mask.ndim == 4, "Image and Mask must have 4 dimensions" - assert image.shape[-2:] == mask.shape[-2:], "Image and Mask must have the same spatial dimensions" - assert image.shape[0] == mask.shape[0], "Image and Mask must have the same batch size" - assert mask.shape[1] == 1, "Mask image must have a single channel" - - # Check image is in [-1, 1] - if image.min() < -1 or image.max() > 1: - raise ValueError("Image should be in [-1, 1] range") - - # Check mask is in [0, 1] - if mask.min() < 0 or mask.max() > 1: - raise ValueError("Mask should be in [0, 1] range") - - # paint-by-example inverses the mask - mask = 1 - mask - - # Binarize mask - mask[mask < 0.5] = 0 - mask[mask >= 0.5] = 1 - - # Image as float32 - image = image.to(dtype=torch.float32) - elif isinstance(mask, torch.Tensor): - raise TypeError(f"`mask` is a torch.Tensor but `image` (type: {type(image)} is not") - else: - if isinstance(image, PIL.Image.Image): - image = [image] - - image = np.concatenate([np.array(i.convert("RGB"))[None, :] for i in image], axis=0) - image = image.transpose(0, 3, 1, 2) - image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 - - # preprocess mask - if isinstance(mask, PIL.Image.Image): - mask = [mask] - - mask = np.concatenate([np.array(m.convert("L"))[None, None, :] for m in mask], axis=0) - mask = mask.astype(np.float32) / 255.0 - - # paint-by-example inverses the mask - mask = 1 - mask - - mask[mask < 0.5] = 0 - mask[mask >= 0.5] = 1 - mask = torch.from_numpy(mask) - - masked_image = image * mask - - return mask, masked_image - - -class PaintByExamplePipeline(DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin): - _last_supported_version = "0.33.1" - r""" - - - 🧪 This is an experimental feature! - - - - Pipeline for image-guided image inpainting using Stable Diffusion. - - 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: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - image_encoder ([`PaintByExampleImageEncoder`]): - Encodes the example input image. The `unet` is conditioned on the example image instead of a text prompt. - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - - """ - - # TODO: feature_extractor is required to encode initial images (if they are in PIL format), - # we should give a descriptive message if the pipeline doesn't have one. - - model_cpu_offload_seq = "unet->vae" - _exclude_from_cpu_offload = ["image_encoder"] - _optional_components = ["safety_checker"] - - def __init__( - self, - vae: AutoencoderKL, - image_encoder: PaintByExampleImageEncoder, - unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = False, - ): - super().__init__() - - self.register_modules( - vae=vae, - image_encoder=image_encoder, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_image_variation.StableDiffusionImageVariationPipeline.check_inputs - def check_inputs(self, image, height, width, callback_steps): - if ( - not isinstance(image, torch.Tensor) - and not isinstance(image, PIL.Image.Image) - 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" - f" {type(image)}" - ) - - 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 None) or ( - 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)}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.StableDiffusionInpaintPipeline.prepare_mask_latents - def prepare_mask_latents( - self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance - ): - # 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 // self.vae_scale_factor, width // self.vae_scale_factor) - ) - mask = mask.to(device=device, dtype=dtype) - - masked_image = masked_image.to(device=device, dtype=dtype) - - if masked_image.shape[1] == 4: - masked_image_latents = masked_image - else: - masked_image_latents = self._encode_vae_image(masked_image, generator=generator) - - # 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) - - mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask - masked_image_latents = ( - torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents - ) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.StableDiffusionInpaintPipeline._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 = self.vae.config.scaling_factor * image_latents - - return image_latents - - def _encode_image(self, image, device, num_images_per_prompt, do_classifier_free_guidance): - dtype = next(self.image_encoder.parameters()).dtype - - if not isinstance(image, torch.Tensor): - image = self.feature_extractor(images=image, return_tensors="pt").pixel_values - - image = image.to(device=device, dtype=dtype) - image_embeddings, negative_prompt_embeds = self.image_encoder(image, return_uncond_vector=True) - - # duplicate image embeddings for each generation per prompt, using mps friendly method - bs_embed, seq_len, _ = image_embeddings.shape - image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1) - image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) - - if do_classifier_free_guidance: - negative_prompt_embeds = negative_prompt_embeds.repeat(1, image_embeddings.shape[0], 1) - negative_prompt_embeds = negative_prompt_embeds.view(bs_embed * num_images_per_prompt, 1, -1) - - # 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 - image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings]) - - return image_embeddings - - @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, - 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, - eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: 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, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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` 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]`): - `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 - expected shape would be `(B, H, W, 1)`. - 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): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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`. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - - Example: - - ```py - >>> import PIL - >>> import requests - >>> import torch - >>> from io import BytesIO - >>> from diffusers import PaintByExamplePipeline - - - >>> def download_image(url): - ... response = requests.get(url) - ... return PIL.Image.open(BytesIO(response.content)).convert("RGB") - - - >>> img_url = ( - ... "https://raw.githubusercontent.com/Fantasy-Studio/Paint-by-Example/main/examples/image/example_1.png" - ... ) - >>> mask_url = ( - ... "https://raw.githubusercontent.com/Fantasy-Studio/Paint-by-Example/main/examples/mask/example_1.png" - ... ) - >>> example_url = "https://raw.githubusercontent.com/Fantasy-Studio/Paint-by-Example/main/examples/reference/example_1.jpg" - - >>> init_image = download_image(img_url).resize((512, 512)) - >>> mask_image = download_image(mask_url).resize((512, 512)) - >>> example_image = download_image(example_url).resize((512, 512)) - - >>> pipe = PaintByExamplePipeline.from_pretrained( - ... "Fantasy-Studio/Paint-by-Example", - ... torch_dtype=torch.float16, - ... ) - >>> pipe = pipe.to("cuda") - - >>> image = pipe(image=init_image, mask_image=mask_image, example_image=example_image).images[0] - >>> image - ``` - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - # 1. Define call parameters - if isinstance(image, PIL.Image.Image): - batch_size = 1 - elif isinstance(image, list): - batch_size = len(image) - else: - batch_size = image.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 = guidance_scale > 1.0 - - # 2. Preprocess mask and image - mask, masked_image = prepare_mask_and_masked_image(image, mask_image) - height, width = masked_image.shape[-2:] - - # 3. Check inputs - self.check_inputs(example_image, height, width, callback_steps) - - # 4. Encode input image - image_embeddings = self._encode_image( - example_image, device, num_images_per_prompt, do_classifier_free_guidance - ) - - # 5. set timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 6. Prepare latent variables - num_channels_latents = self.vae.config.latent_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - image_embeddings.dtype, - device, - generator, - latents, - ) - - # 7. Prepare mask latent variables - mask, masked_image_latents = self.prepare_mask_latents( - mask, - masked_image, - batch_size * num_images_per_prompt, - height, - width, - image_embeddings.dtype, - device, - generator, - do_classifier_free_guidance, - ) - - # 8. Check that sizes of mask, masked image and latents match - num_channels_mask = mask.shape[1] - num_channels_masked_image = masked_image_latents.shape[1] - if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" - " `pipeline.unet` or your `mask_image` or `image` input." - ) - - # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 10. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - # concat latents, mask, masked_image_latents in the channel dimension - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - latent_model_input = torch.cat([latent_model_input, masked_image_latents, mask], dim=1) - - # predict the noise residual - noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=image_embeddings).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - self.maybe_free_model_hooks() - - 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, image_embeddings.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) - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pia/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pia/__init__.py deleted file mode 100644 index 16e8004966e58387f7dba2f6ff3175575fe0abee..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pia/__init__.py +++ /dev/null @@ -1,46 +0,0 @@ -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 - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) -else: - _import_structure["pipeline_pia"] = ["PIAPipeline", "PIAPipelineOutput"] - -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_pia import PIAPipeline, PIAPipelineOutput - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/pia/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pia/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index ddf97323c869a20bee523442ccbaa938f409d5d2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pia/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pia/__pycache__/pipeline_pia.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pia/__pycache__/pipeline_pia.cpython-310.pyc deleted file mode 100644 index 8893518f9cd4220ce30eb4c85a04d65f5e132a81..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pia/__pycache__/pipeline_pia.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pia/pipeline_pia.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pia/pipeline_pia.py deleted file mode 100644 index dfc6e83fbd7cc1a43ee5a88f75544d39b9427ead..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pia/pipeline_pia.py +++ /dev/null @@ -1,958 +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 inspect -from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...image_processor import PipelineImageInput -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel, UNetMotionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...models.unets.unet_motion_model import MotionAdapter -from ...schedulers import ( - DDIMScheduler, - DPMSolverMultistepScheduler, - EulerAncestralDiscreteScheduler, - EulerDiscreteScheduler, - LMSDiscreteScheduler, - PNDMScheduler, -) -from ...utils import ( - USE_PEFT_BACKEND, - BaseOutput, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import randn_tensor -from ...video_processor import VideoProcessor -from ..free_init_utils import FreeInitMixin -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin - - -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 EulerDiscreteScheduler, MotionAdapter, PIAPipeline - >>> from diffusers.utils import export_to_gif, load_image - - >>> adapter = MotionAdapter.from_pretrained("openmmlab/PIA-condition-adapter") - >>> pipe = PIAPipeline.from_pretrained( - ... "SG161222/Realistic_Vision_V6.0_B1_noVAE", motion_adapter=adapter, torch_dtype=torch.float16 - ... ) - - >>> pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config) - >>> image = load_image( - ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/pix2pix/cat_6.png?download=true" - ... ) - >>> image = image.resize((512, 512)) - >>> prompt = "cat in a hat" - >>> negative_prompt = "wrong white balance, dark, sketches, worst quality, low quality, deformed, distorted" - >>> generator = torch.Generator("cpu").manual_seed(0) - >>> output = pipe(image=image, prompt=prompt, negative_prompt=negative_prompt, generator=generator) - >>> frames = output.frames[0] - >>> export_to_gif(frames, "pia-animation.gif") - ``` -""" - -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 - [1.0, 0.9, 0.85, 0.85, 0.85, 0.8, 0.8, 0.8, 0.8, 0.8, 0.8, 0.8, 0.85, 0.85, 0.9, 1.0], # Loop - [1.0, 0.8, 0.8, 0.8, 0.79, 0.78, 0.75, 0.75, 0.75, 0.75, 0.75, 0.78, 0.79, 0.8, 0.8, 1.0], # Loop - [1.0, 0.8, 0.7, 0.7, 0.7, 0.7, 0.6, 0.5, 0.5, 0.6, 0.7, 0.7, 0.7, 0.7, 0.8, 1.0], # Loop - [0.5, 0.4, 0.4, 0.4, 0.35, 0.3], # Style Transfer Candidate Small Motion - [0.5, 0.4, 0.4, 0.4, 0.35, 0.35, 0.3, 0.25, 0.2], # Style Transfer Moderate Motion - [0.5, 0.2], # Style Transfer Large Motion -] - - -def prepare_mask_coef_by_statistics(num_frames: int, cond_frame: int, motion_scale: int): - assert num_frames > 0, "video_length should be greater than 0" - - assert num_frames > cond_frame, "video_length should be greater than cond_frame" - - range_list = RANGE_LIST - - assert motion_scale < len(range_list), f"motion_scale type{motion_scale} not implemented" - - coef = range_list[motion_scale] - coef = coef + ([coef[-1]] * (num_frames - len(coef))) - - order = [abs(i - cond_frame) for i in range(num_frames)] - coef = [coef[order[i]] for i in range(num_frames)] - - return coef - - -@dataclass -class PIAPipelineOutput(BaseOutput): - r""" - Output class for PIAPipeline. - - Args: - 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]]] - - -class PIAPipeline( - DeprecatedPipelineMixin, - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - FromSingleFileMixin, - FreeInitMixin, -): - _last_supported_version = "0.33.1" - r""" - Pipeline for text-to-video generation. - - 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.). - - 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 - - [`~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 ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer (`CLIPTokenizer`): - A [`~transformers.CLIPTokenizer`] to tokenize text. - unet ([`UNet2DConditionModel`]): - 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. - 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`]. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["feature_extractor", "image_encoder", "motion_adapter"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], - scheduler: Union[ - DDIMScheduler, - PNDMScheduler, - LMSDiscreteScheduler, - EulerDiscreteScheduler, - EulerAncestralDiscreteScheduler, - DPMSolverMultistepScheduler, - ], - motion_adapter: Optional[MotionAdapter] = None, - feature_extractor: CLIPImageProcessor = None, - image_encoder: CLIPVisionModelWithProjection = None, - ): - super().__init__() - if isinstance(unet, UNet2DConditionModel): - unet = UNetMotionModel.from_unet2d(unet, motion_adapter) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - motion_adapter=motion_adapter, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - self.video_processor = VideoProcessor(do_resize=False, vae_scale_factor=self.vae_scale_factor) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt with num_images_per_prompt -> num_videos_per_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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.text_to_video_synthesis/pipeline_text_to_video_synth.TextToVideoSDPipeline.decode_latents - def decode_latents(self, latents): - latents = 1 / self.vae.config.scaling_factor * latents - - batch_size, channels, num_frames, height, width = latents.shape - latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) - - image = self.vae.decode(latents).sample - video = image[None, :].reshape((batch_size, num_frames, -1) + image.shape[2:]).permute(0, 2, 1, 3, 4) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - video = video.float() - return video - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_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_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 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.text_to_video_synthesis.pipeline_text_to_video_synth.TextToVideoSDPipeline.prepare_latents - def prepare_latents( - self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None - ): - shape = ( - batch_size, - num_channels_latents, - num_frames, - height // self.vae_scale_factor, - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def prepare_masked_condition( - self, - image, - batch_size, - num_channels_latents, - num_frames, - height, - width, - dtype, - device, - generator, - motion_scale=0, - ): - shape = ( - batch_size, - num_channels_latents, - num_frames, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - _, _, _, scaled_height, scaled_width = shape - - image = self.video_processor.preprocess(image) - image = image.to(device, dtype) - - if isinstance(generator, list): - image_latent = [ - self.vae.encode(image[k : k + 1]).latent_dist.sample(generator[k]) for k in range(batch_size) - ] - image_latent = torch.cat(image_latent, dim=0) - else: - image_latent = self.vae.encode(image).latent_dist.sample(generator) - - image_latent = image_latent.to(device=device, dtype=dtype) - image_latent = torch.nn.functional.interpolate(image_latent, size=[scaled_height, scaled_width]) - image_latent_padding = image_latent.clone() * self.vae.config.scaling_factor - - mask = torch.zeros((batch_size, 1, num_frames, scaled_height, scaled_width)).to(device=device, dtype=dtype) - mask_coef = prepare_mask_coef_by_statistics(num_frames, 0, motion_scale) - masked_image = torch.zeros(batch_size, 4, num_frames, scaled_height, scaled_width).to( - device=device, dtype=self.unet.dtype - ) - for f in range(num_frames): - mask[:, :, f, :, :] = mask_coef[f] - masked_image[:, :, f, :, :] = image_latent_padding.clone() - - mask = torch.cat([mask] * 2) if self.do_classifier_free_guidance else mask - masked_image = torch.cat([masked_image] * 2) if self.do_classifier_free_guidance else masked_image - - return mask, masked_image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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 - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - image: PipelineImageInput, - prompt: Union[str, List[str]] = None, - strength: float = 1.0, - num_frames: Optional[int] = 16, - 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_videos_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, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - motion_scale: int = 0, - output_type: Optional[str] = "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"], - ): - r""" - The call function to the pipeline for generation. - - Args: - image (`PipelineImageInput`): - The input image to be used for video generation. - 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. - height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The height in pixels of the generated video. - width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The width in pixels of the generated video. - num_frames (`int`, *optional*, defaults to 16): - The number of video frames that are generated. Defaults to 16 frames which at 8 frames per seconds - amounts to 2 seconds of video. - 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. - 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*): - 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*): - 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 video - 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`. Latents should be of shape - `(batch_size, num_channel, num_frames, height, width)`. - 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. - 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. - 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)`. 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. - motion_scale: (`int`, *optional*, defaults to 0): - Parameter that controls the amount and type of motion that is added to the image. Increasing the value - increases the amount of motion, while specific ranges of values control the type of motion that is - added. Must be between 0 and 8. Set between 0-2 to only increase the amount of motion. Set between 3-5 - to create looping motion. Set between 6-8 to perform motion with image style transfer. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated video. Choose between `torch.Tensor`, `PIL.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.text_to_video_synthesis.TextToVideoSDPipelineOutput`] instead - of a plain tuple. - 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). - 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. - 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. - - Examples: - - Returns: - [`~pipelines.pia.pipeline_pia.PIAPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.pia.pipeline_pia.PIAPipelineOutput`] is returned, otherwise a - `tuple` is returned where the first element is a list with the generated frames. - """ - # 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 - - num_videos_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - - # 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 - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_videos_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - prompt_embeds = prompt_embeds.repeat_interleave(repeats=num_frames, dim=0) - - 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_videos_per_prompt, - self.do_classifier_free_guidance, - ) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) - latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt) - self._num_timesteps = len(timesteps) - - # 5. Prepare latent variables - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - 4, - num_frames, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents=latents, - ) - mask, masked_image = self.prepare_masked_condition( - image, - batch_size * num_videos_per_prompt, - 4, - num_frames=num_frames, - height=height, - width=width, - dtype=self.unet.dtype, - device=device, - generator=generator, - motion_scale=motion_scale, - ) - if strength < 1.0: - noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) - latents = self.scheduler.add_noise(masked_image[0], noise, latent_timestep) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - # 8. Denoising loop - num_free_init_iters = self._free_init_num_iters if self.free_init_enabled else 1 - for free_init_iter in range(num_free_init_iters): - if self.free_init_enabled: - latents, timesteps = self._apply_free_init( - latents, free_init_iter, num_inference_steps, device, latents.dtype, generator - ) - - self._num_timesteps = len(timesteps) - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - with self.progress_bar(total=self._num_timesteps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - latent_model_input = torch.cat([latent_model_input, mask, masked_image], dim=1) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - ).sample - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - 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() - - # 9. Post processing - if output_type == "latent": - video = latents - else: - video_tensor = self.decode_latents(latents) - video = self.video_processor.postprocess_video(video=video_tensor, output_type=output_type) - - # 10. Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return PIAPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pipeline_flax_utils.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pipeline_flax_utils.py deleted file mode 100644 index 7c5ac89602da649fc2d03a540ffd28740b1687b6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pipeline_flax_utils.py +++ /dev/null @@ -1,610 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# Copyright (c) 2022, NVIDIA CORPORATION. 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 -import os -from typing import Any, Dict, List, Optional, Union - -import flax -import numpy as np -import PIL.Image -from flax.core.frozen_dict import FrozenDict -from huggingface_hub import create_repo, snapshot_download -from huggingface_hub.utils import validate_hf_hub_args -from PIL import Image -from tqdm.auto import tqdm - -from ..configuration_utils import ConfigMixin -from ..models.modeling_flax_utils import FLAX_WEIGHTS_NAME, FlaxModelMixin -from ..schedulers.scheduling_utils_flax import SCHEDULER_CONFIG_NAME, FlaxSchedulerMixin -from ..utils import ( - CONFIG_NAME, - BaseOutput, - PushToHubMixin, - http_user_agent, - is_transformers_available, - logging, -) - - -if is_transformers_available(): - from transformers import FlaxPreTrainedModel - -INDEX_FILE = "diffusion_flax_model.bin" - - -logger = logging.get_logger(__name__) - - -LOADABLE_CLASSES = { - "diffusers": { - "FlaxModelMixin": ["save_pretrained", "from_pretrained"], - "FlaxSchedulerMixin": ["save_pretrained", "from_pretrained"], - "FlaxDiffusionPipeline": ["save_pretrained", "from_pretrained"], - }, - "transformers": { - "PreTrainedTokenizer": ["save_pretrained", "from_pretrained"], - "PreTrainedTokenizerFast": ["save_pretrained", "from_pretrained"], - "FlaxPreTrainedModel": ["save_pretrained", "from_pretrained"], - "FeatureExtractionMixin": ["save_pretrained", "from_pretrained"], - "ProcessorMixin": ["save_pretrained", "from_pretrained"], - "ImageProcessingMixin": ["save_pretrained", "from_pretrained"], - }, -} - -ALL_IMPORTABLE_CLASSES = {} -for library in LOADABLE_CLASSES: - ALL_IMPORTABLE_CLASSES.update(LOADABLE_CLASSES[library]) - - -def import_flax_or_no_model(module, class_name): - try: - # 1. First make sure that if a Flax object is present, import this one - class_obj = getattr(module, "Flax" + class_name) - except AttributeError: - # 2. If this doesn't work, it's not a model and we don't append "Flax" - class_obj = getattr(module, class_name) - except AttributeError: - raise ValueError(f"Neither Flax{class_name} nor {class_name} exist in {module}") - - return class_obj - - -@flax.struct.dataclass -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, - num_channels)`. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] - - -class FlaxDiffusionPipeline(ConfigMixin, PushToHubMixin): - r""" - Base class for Flax-based pipelines. - - [`FlaxDiffusionPipeline`] stores all components (models, schedulers, and processors) for diffusion pipelines and - provides methods for loading, downloading and saving models. It also includes methods to: - - - enable/disable the progress bar for the denoising iteration - - Class attributes: - - - **config_name** ([`str`]) -- The configuration filename that stores the class and module names of all the - diffusion pipeline's components. - """ - - config_name = "model_index.json" - - def register_modules(self, **kwargs): - # import it here to avoid circular import - from diffusers import pipelines - - for name, module in kwargs.items(): - if module is None: - register_dict = {name: (None, None)} - else: - # retrieve library - library = module.__module__.split(".")[0] - - # check if the module is a pipeline module - pipeline_dir = module.__module__.split(".")[-2] - path = module.__module__.split(".") - is_pipeline_module = pipeline_dir in path and hasattr(pipelines, pipeline_dir) - - # if library is not in LOADABLE_CLASSES, then it is a custom module. - # Or if it's a pipeline module, then the module is inside the pipeline - # folder so we set the library to module name. - if library not in LOADABLE_CLASSES or is_pipeline_module: - library = pipeline_dir - - # retrieve class_name - class_name = module.__class__.__name__ - - register_dict = {name: (library, class_name)} - - # save model index config - self.register_to_config(**register_dict) - - # set models - setattr(self, name, module) - - def save_pretrained( - self, - save_directory: Union[str, os.PathLike], - params: Union[Dict, FrozenDict], - push_to_hub: bool = False, - **kwargs, - ): - # TODO: handle inference_state - """ - Save all saveable variables of the pipeline to a directory. A pipeline variable can be saved and loaded if its - class implements both a save and loading method. The pipeline is easily reloaded using the - [`~FlaxDiffusionPipeline.from_pretrained`] class method. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to which to save. Will be created if it doesn't exist. - push_to_hub (`bool`, *optional*, defaults to `False`): - 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*): - Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. - """ - self.save_config(save_directory) - - model_index_dict = dict(self.config) - model_index_dict.pop("_class_name") - model_index_dict.pop("_diffusers_version") - model_index_dict.pop("_module", None) - - 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]) - repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id - - for pipeline_component_name in model_index_dict.keys(): - sub_model = getattr(self, pipeline_component_name) - if sub_model is None: - # edge case for saving a pipeline with safety_checker=None - continue - - model_cls = sub_model.__class__ - - save_method_name = None - # search for the model's base class in LOADABLE_CLASSES - for library_name, library_classes in LOADABLE_CLASSES.items(): - library = importlib.import_module(library_name) - 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): - # if we found a suitable base class in LOADABLE_CLASSES then grab its save method - save_method_name = save_load_methods[0] - break - if save_method_name is not None: - break - - save_method = getattr(sub_model, save_method_name) - expects_params = "params" in set(inspect.signature(save_method).parameters.keys()) - - if expects_params: - save_method( - os.path.join(save_directory, pipeline_component_name), params=params[pipeline_component_name] - ) - else: - save_method(os.path.join(save_directory, pipeline_component_name)) - - if push_to_hub: - self._upload_folder( - save_directory, - repo_id, - token=token, - commit_message=commit_message, - create_pr=create_pr, - ) - - @classmethod - @validate_hf_hub_args - def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs): - r""" - Instantiate a Flax-based diffusion pipeline from pretrained pipeline weights. - - The pipeline is set in evaluation mode (`model.eval()) by default and dropout modules are deactivated. - - If you get the error message below, you need to finetune the weights for your downstream task: - - ``` - Some weights of FlaxUNet2DConditionModel were not initialized from the model checkpoint at stable-diffusion-v1-5/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: - ``` - - Parameters: - pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): - Can be either: - - - A string, the *repo id* (for example `stable-diffusion-v1-5/stable-diffusion-v1-5`) of a - pretrained pipeline hosted on the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved - using [`~FlaxDiffusionPipeline.save_pretrained`]. - dtype (`jnp.dtype`, *optional*): - Override the default `jnp.dtype` and load the model under this dtype. - 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*): - 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`): - Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - mirror (`str`, *optional*): - 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. - kwargs (remaining dictionary of keyword arguments, *optional*): - Can be used to overwrite load and saveable variables (the pipeline components) of the specific pipeline - class. The overwritten components are passed directly to the pipelines `__init__` method. - - - - To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with - `huggingface-cli login`. - - - - Examples: - - ```py - >>> from diffusers import FlaxDiffusionPipeline - - >>> # Download pipeline from huggingface.co and cache. - >>> # Requires to be logged in to Hugging Face hub, - >>> # see more in [the documentation](https://huggingface.co/docs/hub/security-tokens) - >>> pipeline, params = FlaxDiffusionPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", - ... variant="bf16", - ... dtype=jnp.bfloat16, - ... ) - - >>> # Download pipeline, but use a different scheduler - >>> from diffusers import FlaxDPMSolverMultistepScheduler - - >>> model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" - >>> dpmpp, dpmpp_state = FlaxDPMSolverMultistepScheduler.from_pretrained( - ... model_id, - ... subfolder="scheduler", - ... ) - - >>> dpm_pipe, dpm_params = FlaxStableDiffusionPipeline.from_pretrained( - ... model_id, variant="bf16", dtype=jnp.bfloat16, scheduler=dpmpp - ... ) - >>> dpm_params["scheduler"] = dpmpp_state - ``` - """ - cache_dir = kwargs.pop("cache_dir", None) - proxies = kwargs.pop("proxies", None) - local_files_only = kwargs.pop("local_files_only", False) - token = kwargs.pop("token", None) - revision = kwargs.pop("revision", None) - from_pt = kwargs.pop("from_pt", False) - use_memory_efficient_attention = kwargs.pop("use_memory_efficient_attention", False) - split_head_dim = kwargs.pop("split_head_dim", False) - dtype = kwargs.pop("dtype", None) - - # 1. Download the checkpoints and configs - # use snapshot download here to get it working from from_pretrained - if not os.path.isdir(pretrained_model_name_or_path): - config_dict = cls.load_config( - pretrained_model_name_or_path, - cache_dir=cache_dir, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - ) - # make sure we only download sub-folders and `diffusers` filenames - folder_names = [k for k in config_dict.keys() if not k.startswith("_")] - allow_patterns = [os.path.join(k, "*") for k in folder_names] - allow_patterns += [FLAX_WEIGHTS_NAME, SCHEDULER_CONFIG_NAME, CONFIG_NAME, cls.config_name] - - ignore_patterns = ["*.bin", "*.safetensors"] if not from_pt else [] - ignore_patterns += ["*.onnx", "*.onnx_data", "*.xml", "*.pb"] - - if cls != FlaxDiffusionPipeline: - requested_pipeline_class = cls.__name__ - else: - requested_pipeline_class = config_dict.get("_class_name", cls.__name__) - requested_pipeline_class = ( - requested_pipeline_class - if requested_pipeline_class.startswith("Flax") - else "Flax" + requested_pipeline_class - ) - - user_agent = {"pipeline_class": requested_pipeline_class} - user_agent = http_user_agent(user_agent) - - # download all allow_patterns - cached_folder = snapshot_download( - pretrained_model_name_or_path, - cache_dir=cache_dir, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - allow_patterns=allow_patterns, - ignore_patterns=ignore_patterns, - user_agent=user_agent, - ) - else: - cached_folder = pretrained_model_name_or_path - - config_dict = cls.load_config(cached_folder) - - # 2. Load the pipeline class, if using custom module then load it from the hub - # if we load from explicit class, let's use it - if cls != FlaxDiffusionPipeline: - pipeline_class = cls - else: - diffusers_module = importlib.import_module(cls.__module__.split(".")[0]) - class_name = ( - config_dict["_class_name"] - if config_dict["_class_name"].startswith("Flax") - else "Flax" + config_dict["_class_name"] - ) - pipeline_class = getattr(diffusers_module, class_name) - - # some modules can be passed directly to the init - # in this case they are already instantiated in `kwargs` - # extract them here - expected_modules, optional_kwargs = cls._get_signature_keys(pipeline_class) - passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs} - passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs} - - init_dict, unused_kwargs, _ = pipeline_class.extract_init_dict(config_dict, **kwargs) - - # define init kwargs - init_kwargs = {k: init_dict.pop(k) for k in optional_kwargs if k in init_dict} - init_kwargs = {**init_kwargs, **passed_pipe_kwargs} - - # remove `null` components - def load_module(name, value): - if value[0] is None: - return False - if name in passed_class_obj and passed_class_obj[name] is None: - return False - return True - - init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)} - - # Throw nice warnings / errors for fast accelerate loading - if len(unused_kwargs) > 0: - logger.warning( - f"Keyword arguments {unused_kwargs} are not expected by {pipeline_class.__name__} and will be ignored." - ) - - # inference_params - params = {} - - # import it here to avoid circular import - from diffusers import pipelines - - # 3. Load each module in the pipeline - for name, (library_name, class_name) in init_dict.items(): - if class_name is None: - # edge case for when the pipeline was saved with safety_checker=None - init_kwargs[name] = None - continue - - is_pipeline_module = hasattr(pipelines, library_name) - loaded_sub_model = None - sub_model_should_be_defined = True - - # if the model is in a pipeline module, then we load it from the pipeline - if name in passed_class_obj: - # 1. check that passed_class_obj has correct parent class - if not is_pipeline_module: - library = importlib.import_module(library_name) - class_obj = getattr(library, class_name) - importable_classes = LOADABLE_CLASSES[library_name] - class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()} - - expected_class_obj = None - for class_name, class_candidate in class_candidates.items(): - if class_candidate is not None and issubclass(class_obj, class_candidate): - expected_class_obj = class_candidate - - if not issubclass(passed_class_obj[name].__class__, expected_class_obj): - raise ValueError( - f"{passed_class_obj[name]} is of type: {type(passed_class_obj[name])}, but should be" - f" {expected_class_obj}" - ) - elif passed_class_obj[name] is None: - logger.warning( - f"You have passed `None` for {name} to disable its functionality in {pipeline_class}. Note" - f" that this might lead to problems when using {pipeline_class} and is not recommended." - ) - sub_model_should_be_defined = False - else: - logger.warning( - f"You have passed a non-standard module {passed_class_obj[name]}. We cannot verify whether it" - " has the correct type" - ) - - # set passed class object - loaded_sub_model = passed_class_obj[name] - elif is_pipeline_module: - pipeline_module = getattr(pipelines, library_name) - class_obj = import_flax_or_no_model(pipeline_module, class_name) - - importable_classes = ALL_IMPORTABLE_CLASSES - class_candidates = dict.fromkeys(importable_classes.keys(), class_obj) - else: - # else we just import it from the library. - library = importlib.import_module(library_name) - class_obj = import_flax_or_no_model(library, class_name) - - importable_classes = LOADABLE_CLASSES[library_name] - class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()} - - if loaded_sub_model is None and sub_model_should_be_defined: - load_method_name = None - for class_name, class_candidate in class_candidates.items(): - if class_candidate is not None and issubclass(class_obj, class_candidate): - load_method_name = importable_classes[class_name][1] - - load_method = getattr(class_obj, load_method_name) - - # check if the module is in a subdirectory - if os.path.isdir(os.path.join(cached_folder, name)): - loadable_folder = os.path.join(cached_folder, name) - else: - loaded_sub_model = cached_folder - - if issubclass(class_obj, FlaxModelMixin): - loaded_sub_model, loaded_params = load_method( - loadable_folder, - from_pt=from_pt, - use_memory_efficient_attention=use_memory_efficient_attention, - split_head_dim=split_head_dim, - dtype=dtype, - ) - params[name] = loaded_params - elif is_transformers_available() and issubclass(class_obj, FlaxPreTrainedModel): - if from_pt: - # TODO(Suraj): Fix this in Transformers. We should be able to use `_do_init=False` here - loaded_sub_model = load_method(loadable_folder, from_pt=from_pt) - loaded_params = loaded_sub_model.params - del loaded_sub_model._params - else: - loaded_sub_model, loaded_params = load_method(loadable_folder, _do_init=False) - params[name] = loaded_params - elif issubclass(class_obj, FlaxSchedulerMixin): - loaded_sub_model, scheduler_state = load_method(loadable_folder) - params[name] = scheduler_state - else: - loaded_sub_model = load_method(loadable_folder) - - init_kwargs[name] = loaded_sub_model # UNet(...), # DiffusionSchedule(...) - - # 4. Potentially add passed objects if expected - missing_modules = set(expected_modules) - set(init_kwargs.keys()) - passed_modules = list(passed_class_obj.keys()) - - if len(missing_modules) > 0 and missing_modules <= set(passed_modules): - for module in missing_modules: - init_kwargs[module] = passed_class_obj.get(module, None) - elif len(missing_modules) > 0: - passed_modules = set(list(init_kwargs.keys()) + list(passed_class_obj.keys())) - optional_kwargs - raise ValueError( - f"Pipeline {pipeline_class} expected {expected_modules}, but only {passed_modules} were passed." - ) - - model = pipeline_class(**init_kwargs, dtype=dtype) - return model, params - - @classmethod - def _get_signature_keys(cls, obj): - parameters = inspect.signature(obj.__init__).parameters - required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty} - optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty}) - expected_modules = set(required_parameters.keys()) - {"self"} - - return expected_modules, optional_parameters - - @property - def components(self) -> Dict[str, Any]: - r""" - - The `self.components` property can be useful to run different pipelines with the same weights and - configurations to not have to re-allocate memory. - - Examples: - - ```py - >>> from diffusers import ( - ... FlaxStableDiffusionPipeline, - ... FlaxStableDiffusionImg2ImgPipeline, - ... ) - - >>> text2img = FlaxStableDiffusionPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", variant="bf16", dtype=jnp.bfloat16 - ... ) - >>> img2img = FlaxStableDiffusionImg2ImgPipeline(**text2img.components) - ``` - - Returns: - A dictionary containing all the modules needed to initialize the pipeline. - """ - expected_modules, optional_parameters = self._get_signature_keys(self) - components = { - k: getattr(self, k) for k in self.config.keys() if not k.startswith("_") and k not in optional_parameters - } - - if set(components.keys()) != expected_modules: - raise ValueError( - f"{self} has been incorrectly initialized or {self.__class__} is incorrectly implemented. Expected" - f" {expected_modules} to be defined, but {components} are defined." - ) - - return components - - @staticmethod - def numpy_to_pil(images): - """ - Convert a NumPy image or a batch of images to a PIL image. - """ - if images.ndim == 3: - images = images[None, ...] - images = (images * 255).round().astype("uint8") - if images.shape[-1] == 1: - # special case for grayscale (single channel) images - pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images] - else: - pil_images = [Image.fromarray(image) for image in images] - - return pil_images - - # TODO: make it compatible with jax.lax - def progress_bar(self, iterable): - if not hasattr(self, "_progress_bar_config"): - self._progress_bar_config = {} - elif not isinstance(self._progress_bar_config, dict): - raise ValueError( - f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}." - ) - - return tqdm(iterable, **self._progress_bar_config) - - def set_progress_bar_config(self, **kwargs): - self._progress_bar_config = kwargs diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pipeline_loading_utils.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pipeline_loading_utils.py deleted file mode 100644 index d1c2c2adb4c30d00bdc382c17044ce3431c0ba5b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pipeline_loading_utils.py +++ /dev/null @@ -1,1156 +0,0 @@ -# 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. -import importlib -import os -import re -import warnings -from pathlib import Path -from typing import Any, Callable, Dict, List, Optional, Union - -import requests -import torch -from huggingface_hub import DDUFEntry, ModelCard, model_info, snapshot_download -from huggingface_hub.utils import OfflineModeIsEnabled, validate_hf_hub_args -from packaging import version -from requests.exceptions import HTTPError - -from .. import __version__ -from ..utils import ( - FLAX_WEIGHTS_NAME, - ONNX_EXTERNAL_WEIGHTS_NAME, - ONNX_WEIGHTS_NAME, - SAFETENSORS_WEIGHTS_NAME, - WEIGHTS_NAME, - deprecate, - get_class_from_dynamic_module, - is_accelerate_available, - is_peft_available, - is_transformers_available, - is_transformers_version, - logging, -) -from ..utils.torch_utils import is_compiled_module -from .transformers_loading_utils import _load_tokenizer_from_dduf, _load_transformers_model_from_dduf - - -if is_transformers_available(): - import transformers - from transformers import PreTrainedModel, PreTrainedTokenizerBase - from transformers.utils import FLAX_WEIGHTS_NAME as TRANSFORMERS_FLAX_WEIGHTS_NAME - from transformers.utils import SAFE_WEIGHTS_NAME as TRANSFORMERS_SAFE_WEIGHTS_NAME - from transformers.utils import WEIGHTS_NAME as TRANSFORMERS_WEIGHTS_NAME - -if is_accelerate_available(): - import accelerate - from accelerate import dispatch_model - from accelerate.hooks import remove_hook_from_module - from accelerate.utils import compute_module_sizes, get_max_memory - - -INDEX_FILE = "diffusion_pytorch_model.bin" -CUSTOM_PIPELINE_FILE_NAME = "pipeline.py" -DUMMY_MODULES_FOLDER = "diffusers.utils" -TRANSFORMERS_DUMMY_MODULES_FOLDER = "transformers.utils" -CONNECTED_PIPES_KEYS = ["prior"] - -logger = logging.get_logger(__name__) - -LOADABLE_CLASSES = { - "diffusers": { - "ModelMixin": ["save_pretrained", "from_pretrained"], - "SchedulerMixin": ["save_pretrained", "from_pretrained"], - "DiffusionPipeline": ["save_pretrained", "from_pretrained"], - "OnnxRuntimeModel": ["save_pretrained", "from_pretrained"], - }, - "transformers": { - "PreTrainedTokenizer": ["save_pretrained", "from_pretrained"], - "PreTrainedTokenizerFast": ["save_pretrained", "from_pretrained"], - "PreTrainedModel": ["save_pretrained", "from_pretrained"], - "FeatureExtractionMixin": ["save_pretrained", "from_pretrained"], - "ProcessorMixin": ["save_pretrained", "from_pretrained"], - "ImageProcessingMixin": ["save_pretrained", "from_pretrained"], - }, - "onnxruntime.training": { - "ORTModule": ["save_pretrained", "from_pretrained"], - }, -} - -ALL_IMPORTABLE_CLASSES = {} -for library in LOADABLE_CLASSES: - ALL_IMPORTABLE_CLASSES.update(LOADABLE_CLASSES[library]) - - -def is_safetensors_compatible(filenames, passed_components=None, folder_names=None, variant=None) -> bool: - """ - Checking for safetensors compatibility: - - The model is safetensors compatible only if there is a safetensors file for each model component present in - filenames. - - Converting default pytorch serialized filenames to safetensors serialized filenames: - - For models from the diffusers library, just replace the ".bin" extension with ".safetensors" - - For models from the transformers library, the filename changes from "pytorch_model" to "model", and the ".bin" - extension is replaced with ".safetensors" - """ - weight_names = [ - WEIGHTS_NAME, - SAFETENSORS_WEIGHTS_NAME, - FLAX_WEIGHTS_NAME, - ONNX_WEIGHTS_NAME, - ONNX_EXTERNAL_WEIGHTS_NAME, - ] - - if is_transformers_available(): - weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME, TRANSFORMERS_FLAX_WEIGHTS_NAME] - - # model_pytorch, diffusion_model_pytorch, ... - weight_prefixes = [w.split(".")[0] for w in weight_names] - # .bin, .safetensors, ... - weight_suffixs = [w.split(".")[-1] for w in weight_names] - # -00001-of-00002 - transformers_index_format = r"\d{5}-of-\d{5}" - # `diffusion_pytorch_model.bin` as well as `model-00001-of-00002.safetensors` - variant_file_re = re.compile( - rf"({'|'.join(weight_prefixes)})\.({variant}|{variant}-{transformers_index_format})\.({'|'.join(weight_suffixs)})$" - ) - non_variant_file_re = re.compile( - rf"({'|'.join(weight_prefixes)})(-{transformers_index_format})?\.({'|'.join(weight_suffixs)})$" - ) - - passed_components = passed_components or [] - if folder_names: - filenames = {f for f in filenames if os.path.split(f)[0] in folder_names} - - # extract all components of the pipeline and their associated files - components = {} - for filename in filenames: - if not len(filename.split("/")) == 2: - continue - - component, component_filename = filename.split("/") - if component in passed_components: - continue - - components.setdefault(component, []) - components[component].append(component_filename) - - # If there are no component folders check the main directory for safetensors files - filtered_filenames = set() - if not components: - if variant is not None: - filtered_filenames = filter_with_regex(filenames, variant_file_re) - - # If no variant filenames exist check if non-variant files are available - if not filtered_filenames: - filtered_filenames = filter_with_regex(filenames, non_variant_file_re) - return any(".safetensors" in filename for filename in filtered_filenames) - - # iterate over all files of a component - # check if safetensor files exist for that component - for component, component_filenames in components.items(): - matches = [] - filtered_component_filenames = set() - # if variant is provided check if the variant of the safetensors exists - if variant is not None: - filtered_component_filenames = filter_with_regex(component_filenames, variant_file_re) - - # if variant safetensor files do not exist check for non-variants - if not filtered_component_filenames: - filtered_component_filenames = filter_with_regex(component_filenames, non_variant_file_re) - for component_filename in filtered_component_filenames: - filename, extension = os.path.splitext(component_filename) - - match_exists = extension == ".safetensors" - matches.append(match_exists) - - if not any(matches): - return False - - return True - - -def filter_model_files(filenames): - """Filter model repo files for just files/folders that contain model weights""" - weight_names = [ - WEIGHTS_NAME, - SAFETENSORS_WEIGHTS_NAME, - FLAX_WEIGHTS_NAME, - ONNX_WEIGHTS_NAME, - ONNX_EXTERNAL_WEIGHTS_NAME, - ] - - if is_transformers_available(): - weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME, TRANSFORMERS_FLAX_WEIGHTS_NAME] - - allowed_extensions = [wn.split(".")[-1] for wn in weight_names] - - return [f for f in filenames if any(f.endswith(extension) for extension in allowed_extensions)] - - -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]: - weight_names = [ - WEIGHTS_NAME, - SAFETENSORS_WEIGHTS_NAME, - FLAX_WEIGHTS_NAME, - ONNX_WEIGHTS_NAME, - ONNX_EXTERNAL_WEIGHTS_NAME, - ] - - if is_transformers_available(): - weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME, TRANSFORMERS_FLAX_WEIGHTS_NAME] - - # model_pytorch, diffusion_model_pytorch, ... - weight_prefixes = [w.split(".")[0] for w in weight_names] - # .bin, .safetensors, ... - weight_suffixs = [w.split(".")[-1] for w in weight_names] - # -00001-of-00002 - transformers_index_format = r"\d{5}-of-\d{5}" - - if variant is not None: - # `diffusion_pytorch_model.fp16.bin` as well as `model.fp16-00001-of-00002.safetensors` - variant_file_re = re.compile( - rf"({'|'.join(weight_prefixes)})\.({variant}|{variant}-{transformers_index_format})\.({'|'.join(weight_suffixs)})$" - ) - # `text_encoder/pytorch_model.bin.index.fp16.json` - variant_index_re = re.compile( - rf"({'|'.join(weight_prefixes)})\.({'|'.join(weight_suffixs)})\.index\.{variant}\.json$" - ) - legacy_variant_file_re = re.compile(rf".*-{transformers_index_format}\.{variant}\.[a-z]+$") - legacy_variant_index_re = re.compile( - rf"({'|'.join(weight_prefixes)})\.({'|'.join(weight_suffixs)})\.{variant}\.index\.json$" - ) - - # `diffusion_pytorch_model.bin` as well as `model-00001-of-00002.safetensors` - non_variant_file_re = re.compile( - rf"({'|'.join(weight_prefixes)})(-{transformers_index_format})?\.({'|'.join(weight_suffixs)})$" - ) - # `text_encoder/pytorch_model.bin.index.json` - non_variant_index_re = re.compile(rf"({'|'.join(weight_prefixes)})\.({'|'.join(weight_suffixs)})\.index\.json") - - def filter_for_compatible_extensions(filenames, ignore_patterns=None): - if not ignore_patterns: - return filenames - - # ignore patterns uses glob style patterns e.g *.safetensors but we're only - # interested in the extension name - return {f for f in filenames if not any(f.endswith(pat.lstrip("*.")) for pat in ignore_patterns)} - - # Group files by component - components = {} - for filename in filenames: - if not len(filename.split("/")) == 2: - components.setdefault("", []).append(filename) - continue - - component, _ = filename.split("/") - components.setdefault(component, []).append(filename) - - usable_filenames = set() - variant_filenames = set() - for component, component_filenames in components.items(): - component_filenames = filter_for_compatible_extensions(component_filenames, ignore_patterns=ignore_patterns) - - component_variants = set() - component_legacy_variants = set() - component_non_variants = set() - if variant is not None: - component_variants = filter_with_regex(component_filenames, variant_file_re) - component_variant_index_files = filter_with_regex(component_filenames, variant_index_re) - - component_legacy_variants = filter_with_regex(component_filenames, legacy_variant_file_re) - component_legacy_variant_index_files = filter_with_regex(component_filenames, legacy_variant_index_re) - - if component_variants or component_legacy_variants: - variant_filenames.update( - component_variants | component_variant_index_files - if component_variants - else component_legacy_variants | component_legacy_variant_index_files - ) - - else: - component_non_variants = filter_with_regex(component_filenames, non_variant_file_re) - component_variant_index_files = filter_with_regex(component_filenames, non_variant_index_re) - - usable_filenames.update(component_non_variants | component_variant_index_files) - - usable_filenames.update(variant_filenames) - - if len(variant_filenames) == 0 and variant is not None: - error_message = f"You are trying to load model files of the `variant={variant}`, but no such modeling files are available. " - raise ValueError(error_message) - - if len(variant_filenames) > 0 and usable_filenames != variant_filenames: - logger.warning( - f"\nA mixture of {variant} and non-{variant} filenames will be loaded.\nLoaded {variant} filenames:\n" - f"[{', '.join(variant_filenames)}]\nLoaded non-{variant} filenames:\n" - f"[{', '.join(usable_filenames - variant_filenames)}\nIf this behavior is not " - f"expected, please check your folder structure." - ) - - return usable_filenames, variant_filenames - - -@validate_hf_hub_args -def warn_deprecated_model_variant(pretrained_model_name_or_path, token, variant, revision, model_filenames): - info = model_info( - pretrained_model_name_or_path, - token=token, - revision=None, - ) - filenames = {sibling.rfilename for sibling in info.siblings} - comp_model_filenames, _ = variant_compatible_siblings(filenames, variant=revision) - comp_model_filenames = [".".join(f.split(".")[:1] + f.split(".")[2:]) for f in comp_model_filenames] - - if set(model_filenames).issubset(set(comp_model_filenames)): - warnings.warn( - f"You are loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'` even though you can load it via `variant=`{revision}`. Loading model variants via `revision='{revision}'` is deprecated and will be removed in diffusers v1. Please use `variant='{revision}'` instead.", - FutureWarning, - ) - else: - warnings.warn( - f"You are loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'`. This behavior is deprecated and will be removed in diffusers v1. One should use `variant='{revision}'` instead. However, it appears that {pretrained_model_name_or_path} currently does not have the required variant filenames in the 'main' branch. \n The Diffusers team and community would be very grateful if you could open an issue: https://github.com/huggingface/diffusers/issues/new with the title '{pretrained_model_name_or_path} is missing {revision} files' so that the correct variant file can be added.", - FutureWarning, - ) - - -def _unwrap_model(model): - """Unwraps a model.""" - if is_compiled_module(model): - model = model._orig_mod - - if is_peft_available(): - from peft import PeftModel - - if isinstance(model, PeftModel): - model = model.base_model.model - - return model - - -def maybe_raise_or_warn( - library_name, library, class_name, importable_classes, passed_class_obj, name, is_pipeline_module -): - """Simple helper method to raise or warn in case incorrect module has been passed""" - if not is_pipeline_module: - library = importlib.import_module(library_name) - class_obj = getattr(library, class_name) - class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()} - - expected_class_obj = None - for class_name, class_candidate in class_candidates.items(): - if class_candidate is not None and issubclass(class_obj, class_candidate): - expected_class_obj = class_candidate - - # Dynamo wraps the original model in a private class. - # I didn't find a public API to get the original class. - sub_model = passed_class_obj[name] - unwrapped_sub_model = _unwrap_model(sub_model) - model_cls = unwrapped_sub_model.__class__ - - if not issubclass(model_cls, expected_class_obj): - raise ValueError(f"{passed_class_obj[name]} is of type: {model_cls}, but should be {expected_class_obj}") - else: - logger.warning( - f"You have passed a non-standard module {passed_class_obj[name]}. We cannot verify whether it" - " has the correct type" - ) - - -def get_class_obj_and_candidates( - library_name, class_name, importable_classes, pipelines, is_pipeline_module, component_name=None, cache_dir=None -): - """Simple helper method to retrieve class object of module as well as potential parent class objects""" - component_folder = os.path.join(cache_dir, component_name) if component_name and cache_dir else None - - if is_pipeline_module: - pipeline_module = getattr(pipelines, library_name) - - class_obj = getattr(pipeline_module, class_name) - class_candidates = dict.fromkeys(importable_classes.keys(), class_obj) - elif component_folder and os.path.isfile(os.path.join(component_folder, library_name + ".py")): - # load custom component - class_obj = get_class_from_dynamic_module( - component_folder, module_file=library_name + ".py", class_name=class_name - ) - class_candidates = dict.fromkeys(importable_classes.keys(), class_obj) - else: - # else we just import it from the library. - library = importlib.import_module(library_name) - - class_obj = getattr(library, class_name) - class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()} - - return class_obj, class_candidates - - -def _get_custom_pipeline_class( - custom_pipeline, - repo_id=None, - hub_revision=None, - class_name=None, - cache_dir=None, - revision=None, -): - if custom_pipeline.endswith(".py"): - path = Path(custom_pipeline) - # decompose into folder & file - file_name = path.name - custom_pipeline = path.parent.absolute() - elif repo_id is not None: - file_name = f"{custom_pipeline}.py" - custom_pipeline = repo_id - else: - file_name = CUSTOM_PIPELINE_FILE_NAME - - if repo_id is not None and hub_revision is not None: - # if we load the pipeline code from the Hub - # make sure to overwrite the `revision` - revision = hub_revision - - return get_class_from_dynamic_module( - custom_pipeline, - module_file=file_name, - class_name=class_name, - cache_dir=cache_dir, - revision=revision, - ) - - -def _get_pipeline_class( - class_obj, - config=None, - load_connected_pipeline=False, - custom_pipeline=None, - repo_id=None, - hub_revision=None, - class_name=None, - cache_dir=None, - revision=None, -): - if custom_pipeline is not None: - return _get_custom_pipeline_class( - custom_pipeline, - repo_id=repo_id, - hub_revision=hub_revision, - class_name=class_name, - cache_dir=cache_dir, - revision=revision, - ) - - if class_obj.__name__ != "DiffusionPipeline": - return class_obj - - diffusers_module = importlib.import_module(class_obj.__module__.split(".")[0]) - class_name = class_name or config["_class_name"] - if not class_name: - raise ValueError( - "The class name could not be found in the configuration file. Please make sure to pass the correct `class_name`." - ) - - class_name = class_name[4:] if class_name.startswith("Flax") else class_name - - pipeline_cls = getattr(diffusers_module, class_name) - - if load_connected_pipeline: - from .auto_pipeline import _get_connected_pipeline - - connected_pipeline_cls = _get_connected_pipeline(pipeline_cls) - if connected_pipeline_cls is not None: - logger.info( - f"Loading connected pipeline {connected_pipeline_cls.__name__} instead of {pipeline_cls.__name__} as specified via `load_connected_pipeline=True`" - ) - else: - logger.info(f"{pipeline_cls.__name__} has no connected pipeline class. Loading {pipeline_cls.__name__}.") - - pipeline_cls = connected_pipeline_cls or pipeline_cls - - return pipeline_cls - - -def _load_empty_model( - library_name: str, - class_name: str, - importable_classes: List[Any], - pipelines: Any, - is_pipeline_module: bool, - name: str, - torch_dtype: Union[str, torch.dtype], - cached_folder: Union[str, os.PathLike], - **kwargs, -): - # retrieve class objects. - class_obj, _ = get_class_obj_and_candidates( - library_name, - class_name, - importable_classes, - pipelines, - is_pipeline_module, - component_name=name, - cache_dir=cached_folder, - ) - - if is_transformers_available(): - transformers_version = version.parse(version.parse(transformers.__version__).base_version) - else: - transformers_version = "N/A" - - # Determine library. - is_transformers_model = ( - is_transformers_available() - and issubclass(class_obj, PreTrainedModel) - and transformers_version >= version.parse("4.20.0") - ) - diffusers_module = importlib.import_module(__name__.split(".")[0]) - is_diffusers_model = issubclass(class_obj, diffusers_module.ModelMixin) - - model = None - config_path = cached_folder - user_agent = { - "diffusers": __version__, - "file_type": "model", - "framework": "pytorch", - } - - if is_diffusers_model: - # Load config and then the model on meta. - config, unused_kwargs, commit_hash = class_obj.load_config( - os.path.join(config_path, name), - cache_dir=cached_folder, - return_unused_kwargs=True, - return_commit_hash=True, - force_download=kwargs.pop("force_download", False), - proxies=kwargs.pop("proxies", None), - local_files_only=kwargs.pop("local_files_only", False), - token=kwargs.pop("token", None), - revision=kwargs.pop("revision", None), - subfolder=kwargs.pop("subfolder", None), - user_agent=user_agent, - ) - with accelerate.init_empty_weights(): - model = class_obj.from_config(config, **unused_kwargs) - elif is_transformers_model: - config_class = getattr(class_obj, "config_class", None) - if config_class is None: - raise ValueError("`config_class` cannot be None. Please double-check the model.") - - config = config_class.from_pretrained( - cached_folder, - subfolder=name, - force_download=kwargs.pop("force_download", False), - proxies=kwargs.pop("proxies", None), - local_files_only=kwargs.pop("local_files_only", False), - token=kwargs.pop("token", None), - revision=kwargs.pop("revision", None), - user_agent=user_agent, - ) - with accelerate.init_empty_weights(): - model = class_obj(config) - - if model is not None: - model = model.to(dtype=torch_dtype) - return model - - -def _assign_components_to_devices( - 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] - device_memory = device_memory.copy() - - device_id_component_mapping = {} - current_device_index = 0 - for component in module_sizes: - device_id = device_cycle[current_device_index % len(device_cycle)] - component_memory = module_sizes[component] - curr_device_memory = device_memory[device_id] - - # If the GPU doesn't fit the current component offload to the CPU. - if component_memory > curr_device_memory: - device_id_component_mapping["cpu"] = [component] - else: - if device_id not in device_id_component_mapping: - device_id_component_mapping[device_id] = [component] - else: - device_id_component_mapping[device_id].append(component) - - # Update the device memory. - device_memory[device_id] -= component_memory - current_device_index += 1 - - return device_id_component_mapping - - -def _get_final_device_map(device_map, pipeline_class, passed_class_obj, init_dict, library, max_memory, **kwargs): - # To avoid circular import problem. - from diffusers import pipelines - - torch_dtype = kwargs.get("torch_dtype", torch.float32) - - # Load each module in the pipeline on a meta device so that we can derive the device map. - init_empty_modules = {} - for name, (library_name, class_name) in init_dict.items(): - if class_name.startswith("Flax"): - raise ValueError("Flax pipelines are not supported with `device_map`.") - - # Define all importable classes - is_pipeline_module = hasattr(pipelines, library_name) - importable_classes = ALL_IMPORTABLE_CLASSES - loaded_sub_model = None - - # Use passed sub model or load class_name from library_name - if name in passed_class_obj: - # if the model is in a pipeline module, then we load it from the pipeline - # check that passed_class_obj has correct parent class - maybe_raise_or_warn( - library_name, - library, - class_name, - importable_classes, - passed_class_obj, - name, - is_pipeline_module, - ) - with accelerate.init_empty_weights(): - loaded_sub_model = passed_class_obj[name] - - else: - sub_model_dtype = ( - torch_dtype.get(name, torch_dtype.get("default", torch.float32)) - if isinstance(torch_dtype, dict) - else torch_dtype - ) - loaded_sub_model = _load_empty_model( - library_name=library_name, - class_name=class_name, - importable_classes=importable_classes, - pipelines=pipelines, - is_pipeline_module=is_pipeline_module, - pipeline_class=pipeline_class, - name=name, - torch_dtype=sub_model_dtype, - cached_folder=kwargs.get("cached_folder", None), - force_download=kwargs.get("force_download", None), - proxies=kwargs.get("proxies", None), - local_files_only=kwargs.get("local_files_only", None), - token=kwargs.get("token", None), - revision=kwargs.get("revision", None), - ) - - if loaded_sub_model is not None: - init_empty_modules[name] = loaded_sub_model - - # determine device map - # Obtain a sorted dictionary for mapping the model-level components - # to their sizes. - module_sizes = { - module_name: compute_module_sizes( - module, - dtype=torch_dtype.get(module_name, torch_dtype.get("default", torch.float32)) - if isinstance(torch_dtype, dict) - else torch_dtype, - )[""] - for module_name, module in init_empty_modules.items() - if isinstance(module, torch.nn.Module) - } - module_sizes = dict(sorted(module_sizes.items(), key=lambda item: item[1], reverse=True)) - - # Obtain maximum memory available per device (GPUs only). - max_memory = get_max_memory(max_memory) - max_memory = dict(sorted(max_memory.items(), key=lambda item: item[1], reverse=True)) - max_memory = {k: v for k, v in max_memory.items() if k != "cpu"} - - # Obtain a dictionary mapping the model-level components to the available - # devices based on the maximum memory and the model sizes. - final_device_map = None - if len(max_memory) > 0: - device_id_component_mapping = _assign_components_to_devices( - module_sizes, max_memory, device_mapping_strategy=device_map - ) - - # Obtain the final device map, e.g., `{"unet": 0, "text_encoder": 1, "vae": 1, ...}` - final_device_map = {} - for device_id, components in device_id_component_mapping.items(): - for component in components: - final_device_map[component] = device_id - - return final_device_map - - -def load_sub_model( - library_name: str, - class_name: str, - 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]], - offload_state_dict: bool, - model_variants: Dict[str, str], - name: str, - from_flax: bool, - variant: str, - low_cpu_mem_usage: bool, - cached_folder: Union[str, os.PathLike], - use_safetensors: bool, - dduf_entries: Optional[Dict[str, DDUFEntry]], - provider_options: Any, - quantization_config: Optional[Any] = None, -): - """Helper method to load the module `name` from `library_name` and `class_name`""" - from ..quantizers import PipelineQuantizationConfig - - # retrieve class candidates - - class_obj, class_candidates = get_class_obj_and_candidates( - library_name, - class_name, - importable_classes, - pipelines, - is_pipeline_module, - component_name=name, - cache_dir=cached_folder, - ) - - load_method_name = None - # retrieve load method name - for class_name, class_candidate in class_candidates.items(): - if class_candidate is not None and issubclass(class_obj, class_candidate): - load_method_name = importable_classes[class_name][1] - - # if load method name is None, then we have a dummy module -> raise Error - if load_method_name is None: - none_module = class_obj.__module__ - is_dummy_path = none_module.startswith(DUMMY_MODULES_FOLDER) or none_module.startswith( - TRANSFORMERS_DUMMY_MODULES_FOLDER - ) - if is_dummy_path and "dummy" in none_module: - # call class_obj for nice error message of missing requirements - class_obj() - - raise ValueError( - f"The component {class_obj} of {pipeline_class} cannot be loaded as it does not seem to have" - f" any of the loading methods defined in {ALL_IMPORTABLE_CLASSES}." - ) - - load_method = _get_load_method(class_obj, load_method_name, is_dduf=dduf_entries is not None) - - # 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) - - if is_transformers_available(): - transformers_version = version.parse(version.parse(transformers.__version__).base_version) - else: - transformers_version = "N/A" - - is_transformers_model = ( - is_transformers_available() - and issubclass(class_obj, PreTrainedModel) - and transformers_version >= version.parse("4.20.0") - ) - - # 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. - if is_diffusers_model or is_transformers_model: - loading_kwargs["device_map"] = device_map - loading_kwargs["max_memory"] = max_memory - loading_kwargs["offload_folder"] = offload_folder - loading_kwargs["offload_state_dict"] = offload_state_dict - loading_kwargs["variant"] = model_variants.pop(name, None) - loading_kwargs["use_safetensors"] = use_safetensors - - if from_flax: - loading_kwargs["from_flax"] = True - - # the following can be deleted once the minimum required `transformers` version - # is higher than 4.27 - if ( - is_transformers_model - and loading_kwargs["variant"] is not None - and transformers_version < version.parse("4.27.0") - ): - raise ImportError( - f"When passing `variant='{variant}'`, please make sure to upgrade your `transformers` version to at least 4.27.0.dev0" - ) - elif is_transformers_model and loading_kwargs["variant"] is None: - loading_kwargs.pop("variant") - - # if `from_flax` and model is transformer model, can currently not load with `low_cpu_mem_usage` - if not (from_flax and is_transformers_model): - loading_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage - else: - loading_kwargs["low_cpu_mem_usage"] = False - - if ( - quantization_config is not None - and isinstance(quantization_config, PipelineQuantizationConfig) - and issubclass(class_obj, torch.nn.Module) - ): - model_quant_config = quantization_config._resolve_quant_config( - is_diffusers=is_diffusers_model, module_name=name - ) - if model_quant_config is not None: - loading_kwargs["quantization_config"] = model_quant_config - - # check if the module is in a subdirectory - if dduf_entries: - loading_kwargs["dduf_entries"] = dduf_entries - loaded_sub_model = load_method(name, **loading_kwargs) - elif os.path.isdir(os.path.join(cached_folder, name)): - loaded_sub_model = load_method(os.path.join(cached_folder, name), **loading_kwargs) - else: - # else load from the root directory - loaded_sub_model = load_method(cached_folder, **loading_kwargs) - - if isinstance(loaded_sub_model, torch.nn.Module) and isinstance(device_map, dict): - # remove hooks - remove_hook_from_module(loaded_sub_model, recurse=True) - needs_offloading_to_cpu = device_map[""] == "cpu" - - if needs_offloading_to_cpu: - dispatch_model( - loaded_sub_model, - state_dict=loaded_sub_model.state_dict(), - device_map=device_map, - force_hooks=True, - main_device=0, - ) - else: - dispatch_model(loaded_sub_model, device_map=device_map, force_hooks=True) - - return loaded_sub_model - - -def _get_load_method(class_obj: object, load_method_name: str, is_dduf: bool) -> Callable: - """ - Return the method to load the sub model. - - In practice, this method will return the `"from_pretrained"` (or `load_method_name`) method of the class object - except if loading from a DDUF checkpoint. In that case, transformers models and tokenizers have a specific loading - method that we need to use. - """ - if is_dduf: - if issubclass(class_obj, PreTrainedTokenizerBase): - return lambda *args, **kwargs: _load_tokenizer_from_dduf(class_obj, *args, **kwargs) - if issubclass(class_obj, PreTrainedModel): - return lambda *args, **kwargs: _load_transformers_model_from_dduf(class_obj, *args, **kwargs) - return getattr(class_obj, load_method_name) - - -def _fetch_class_library_tuple(module): - # import it here to avoid circular import - diffusers_module = importlib.import_module(__name__.split(".")[0]) - pipelines = getattr(diffusers_module, "pipelines") - - # register the config from the original module, not the dynamo compiled one - not_compiled_module = _unwrap_model(module) - library = not_compiled_module.__module__.split(".")[0] - - # check if the module is a pipeline module - module_path_items = not_compiled_module.__module__.split(".") - pipeline_dir = module_path_items[-2] if len(module_path_items) > 2 else None - - path = not_compiled_module.__module__.split(".") - is_pipeline_module = pipeline_dir in path and hasattr(pipelines, pipeline_dir) - - # if library is not in LOADABLE_CLASSES, then it is a custom module. - # Or if it's a pipeline module, then the module is inside the pipeline - # folder so we set the library to module name. - if is_pipeline_module: - library = pipeline_dir - elif library not in LOADABLE_CLASSES: - library = not_compiled_module.__module__ - - # retrieve class_name - class_name = not_compiled_module.__class__.__name__ - - return (library, class_name) - - -def _identify_model_variants(folder: str, variant: str, config: dict) -> dict: - model_variants = {} - if variant is not None: - for sub_folder in os.listdir(folder): - folder_path = os.path.join(folder, sub_folder) - is_folder = os.path.isdir(folder_path) and sub_folder in config - variant_exists = is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) - if variant_exists: - model_variants[sub_folder] = variant - return model_variants - - -def _resolve_custom_pipeline_and_cls(folder, config, custom_pipeline): - custom_class_name = None - if os.path.isfile(os.path.join(folder, f"{custom_pipeline}.py")): - custom_pipeline = os.path.join(folder, f"{custom_pipeline}.py") - elif isinstance(config["_class_name"], (list, tuple)) and os.path.isfile( - os.path.join(folder, f"{config['_class_name'][0]}.py") - ): - custom_pipeline = os.path.join(folder, f"{config['_class_name'][0]}.py") - custom_class_name = config["_class_name"][1] - - return custom_pipeline, custom_class_name - - -def _maybe_raise_warning_for_inpainting(pipeline_class, pretrained_model_name_or_path: str, config: dict): - if pipeline_class.__name__ == "StableDiffusionInpaintPipeline" and version.parse( - version.parse(config["_diffusers_version"]).base_version - ) <= version.parse("0.5.1"): - from diffusers import StableDiffusionInpaintPipeline, StableDiffusionInpaintPipelineLegacy - - pipeline_class = StableDiffusionInpaintPipelineLegacy - - deprecation_message = ( - "You are using a legacy checkpoint for inpainting with Stable Diffusion, therefore we are loading the" - f" {StableDiffusionInpaintPipelineLegacy} class instead of {StableDiffusionInpaintPipeline}. For" - " better inpainting results, we strongly suggest using Stable Diffusion's official inpainting" - " checkpoint: https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-inpainting instead or adapting your" - f" checkpoint {pretrained_model_name_or_path} to the format of" - " https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-inpainting. Note that we do not actively maintain" - " the {StableDiffusionInpaintPipelineLegacy} class and will likely remove it in version 1.0.0." - ) - deprecate("StableDiffusionInpaintPipelineLegacy", "1.0.0", deprecation_message, standard_warn=False) - - -def _update_init_kwargs_with_connected_pipeline( - init_kwargs: dict, passed_pipe_kwargs: dict, passed_class_objs: dict, folder: str, **pipeline_loading_kwargs -) -> dict: - from .pipeline_utils import DiffusionPipeline - - modelcard = ModelCard.load(os.path.join(folder, "README.md")) - connected_pipes = {prefix: getattr(modelcard.data, prefix, [None])[0] for prefix in CONNECTED_PIPES_KEYS} - - # We don't scheduler argument to match the existing logic: - # https://github.com/huggingface/diffusers/blob/867e0c919e1aa7ef8b03c8eb1460f4f875a683ae/src/diffusers/pipelines/pipeline_utils.py#L906C13-L925C14 - pipeline_loading_kwargs_cp = pipeline_loading_kwargs.copy() - if pipeline_loading_kwargs_cp is not None and len(pipeline_loading_kwargs_cp) >= 1: - for k in pipeline_loading_kwargs: - if "scheduler" in k: - _ = pipeline_loading_kwargs_cp.pop(k) - - def get_connected_passed_kwargs(prefix): - connected_passed_class_obj = { - k.replace(f"{prefix}_", ""): w for k, w in passed_class_objs.items() if k.split("_")[0] == prefix - } - connected_passed_pipe_kwargs = { - k.replace(f"{prefix}_", ""): w for k, w in passed_pipe_kwargs.items() if k.split("_")[0] == prefix - } - - connected_passed_kwargs = {**connected_passed_class_obj, **connected_passed_pipe_kwargs} - return connected_passed_kwargs - - connected_pipes = { - prefix: DiffusionPipeline.from_pretrained( - repo_id, **pipeline_loading_kwargs_cp, **get_connected_passed_kwargs(prefix) - ) - for prefix, repo_id in connected_pipes.items() - if repo_id is not None - } - - for prefix, connected_pipe in connected_pipes.items(): - # add connected pipes to `init_kwargs` with _, e.g. "prior_text_encoder" - init_kwargs.update( - {"_".join([prefix, name]): component for name, component in connected_pipe.components.items()} - ) - - return init_kwargs - - -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 = config_dict.copy() - - # retrieve all folder_names that contain relevant files - folder_names = [k for k, v in config_dict.items() if isinstance(v, list) and k != "_class_name"] - - diffusers_module = importlib.import_module(__name__.split(".")[0]) - pipelines = getattr(diffusers_module, "pipelines") - - # optionally create a custom component <> custom file mapping - custom_components = {} - for component in folder_names: - module_candidate = config_dict[component][0] - - if module_candidate is None or not isinstance(module_candidate, str): - continue - - # We compute candidate file path on the Hub. Do not use `os.path.join`. - candidate_file = f"{component}/{module_candidate}.py" - - if candidate_file in filenames: - custom_components[component] = module_candidate - elif module_candidate not in LOADABLE_CLASSES and not hasattr(pipelines, module_candidate): - raise ValueError( - f"{candidate_file} as defined in `model_index.json` does not exist in {pretrained_model_name} and is not a module in 'diffusers/pipelines'." - ) - - return custom_components, folder_names - - -def _get_ignore_patterns( - passed_components, - 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]: - if ( - use_safetensors - and not allow_pickle - and not is_safetensors_compatible( - model_filenames, passed_components=passed_components, folder_names=model_folder_names, variant=variant - ) - ): - raise EnvironmentError( - f"Could not find the necessary `safetensors` weights in {model_filenames} (variant={variant})" - ) - - if from_flax: - ignore_patterns = ["*.bin", "*.safetensors", "*.onnx", "*.pb"] - - elif use_safetensors and is_safetensors_compatible( - model_filenames, passed_components=passed_components, folder_names=model_folder_names, variant=variant - ): - ignore_patterns = ["*.bin", "*.msgpack"] - - use_onnx = use_onnx if use_onnx is not None else is_onnx - if not use_onnx: - ignore_patterns += ["*.onnx", "*.pb"] - - else: - ignore_patterns = ["*.safetensors", "*.msgpack"] - - use_onnx = use_onnx if use_onnx is not None else is_onnx - if not use_onnx: - ignore_patterns += ["*.onnx", "*.pb"] - - return ignore_patterns - - -def _download_dduf_file( - pretrained_model_name: str, - dduf_file: str, - pipeline_class_name: str, - cache_dir: str, - proxies: str, - local_files_only: bool, - token: str, - revision: str, -): - model_info_call_error = None - if not local_files_only: - try: - info = model_info(pretrained_model_name, token=token, revision=revision) - except (HTTPError, OfflineModeIsEnabled, requests.ConnectionError) as e: - logger.warning(f"Couldn't connect to the Hub: {e}.\nWill try to load from local cache.") - local_files_only = True - model_info_call_error = e # save error to reraise it if model is not cached locally - - if ( - not local_files_only - and dduf_file is not None - and dduf_file not in (sibling.rfilename for sibling in info.siblings) - ): - raise ValueError(f"Requested {dduf_file} file is not available in {pretrained_model_name}.") - - try: - user_agent = {"pipeline_class": pipeline_class_name, "dduf": True} - cached_folder = snapshot_download( - pretrained_model_name, - cache_dir=cache_dir, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - allow_patterns=[dduf_file], - user_agent=user_agent, - ) - return cached_folder - except FileNotFoundError: - # Means we tried to load pipeline with `local_files_only=True` but the files have not been found in local cache. - # This can happen in two cases: - # 1. If the user passed `local_files_only=True` => we raise the error directly - # 2. If we forced `local_files_only=True` when `model_info` failed => we raise the initial error - if model_info_call_error is None: - # 1. user passed `local_files_only=True` - raise - else: - # 2. we forced `local_files_only=True` when `model_info` failed - raise EnvironmentError( - f"Cannot load model {pretrained_model_name}: model is not cached locally and an error occurred" - " while trying to fetch metadata from the Hub. Please check out the root cause in the stacktrace" - " above." - ) from model_info_call_error - - -def _maybe_raise_error_for_incorrect_transformers(config_dict): - has_transformers_component = False - for k in config_dict: - if isinstance(config_dict[k], list): - has_transformers_component = config_dict[k][0] == "transformers" - if has_transformers_component: - break - if has_transformers_component and not is_transformers_version(">", "4.47.1"): - raise ValueError("Please upgrade your `transformers` installation to the latest version to use DDUF.") - - -def _maybe_warn_for_wrong_component_in_quant_config(pipe_init_dict, quant_config): - if quant_config is None: - return - - actual_pipe_components = set(pipe_init_dict.keys()) - missing = "" - quant_components = None - if getattr(quant_config, "components_to_quantize", None) is not None: - quant_components = set(quant_config.components_to_quantize) - elif getattr(quant_config, "quant_mapping", None) is not None and isinstance(quant_config.quant_mapping, dict): - quant_components = set(quant_config.quant_mapping.keys()) - - if quant_components and not quant_components.issubset(actual_pipe_components): - missing = quant_components - actual_pipe_components - - if missing: - logger.warning( - f"The following components in the quantization config {missing} will be ignored " - "as they do not belong to the underlying pipeline. Acceptable values for the pipeline " - f"components are: {', '.join(actual_pipe_components)}." - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pipeline_utils.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pipeline_utils.py deleted file mode 100644 index 2c03811b51ab3b3db3a9194d0a8bbcaba42745df..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pipeline_utils.py +++ /dev/null @@ -1,2147 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# Copyright (c) 2022, NVIDIA CORPORATION. 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 fnmatch -import importlib -import inspect -import os -import re -import sys -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Callable, Dict, List, Optional, Union, get_args, get_origin - -import numpy as np -import PIL.Image -import requests -import torch -from huggingface_hub import ( - DDUFEntry, - ModelCard, - create_repo, - hf_hub_download, - model_info, - read_dduf_file, - snapshot_download, -) -from huggingface_hub.utils import OfflineModeIsEnabled, validate_hf_hub_args -from packaging import version -from requests.exceptions import HTTPError -from tqdm.auto import tqdm -from typing_extensions import Self - -from .. import __version__ -from ..configuration_utils import ConfigMixin -from ..models import AutoencoderKL -from ..models.attention_processor import FusedAttnProcessor2_0 -from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, ModelMixin -from ..quantizers import PipelineQuantizationConfig -from ..quantizers.bitsandbytes.utils import _check_bnb_status -from ..schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME -from ..utils import ( - CONFIG_NAME, - DEPRECATED_REVISION_ARGS, - BaseOutput, - PushToHubMixin, - _get_detailed_type, - _is_valid_type, - is_accelerate_available, - is_accelerate_version, - is_hpu_available, - is_torch_npu_available, - is_torch_version, - is_transformers_version, - logging, - numpy_to_pil, -) -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 - - -if is_torch_npu_available(): - import torch_npu # noqa: F401 - -from .pipeline_loading_utils import ( - ALL_IMPORTABLE_CLASSES, - CONNECTED_PIPES_KEYS, - CUSTOM_PIPELINE_FILE_NAME, - LOADABLE_CLASSES, - _download_dduf_file, - _fetch_class_library_tuple, - _get_custom_components_and_folders, - _get_custom_pipeline_class, - _get_final_device_map, - _get_ignore_patterns, - _get_pipeline_class, - _identify_model_variants, - _maybe_raise_error_for_incorrect_transformers, - _maybe_raise_warning_for_inpainting, - _maybe_warn_for_wrong_component_in_quant_config, - _resolve_custom_pipeline_and_cls, - _unwrap_model, - _update_init_kwargs_with_connected_pipeline, - filter_model_files, - load_sub_model, - maybe_raise_or_warn, - variant_compatible_siblings, - warn_deprecated_model_variant, -) - - -if is_accelerate_available(): - import accelerate - - -LIBRARIES = [] -for library in LOADABLE_CLASSES: - LIBRARIES.append(library) - -SUPPORTED_DEVICE_MAP = ["balanced"] - -logger = logging.get_logger(__name__) - - -@dataclass -class ImagePipelineOutput(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, - num_channels)`. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] - - -@dataclass -class AudioPipelineOutput(BaseOutput): - """ - Output class for audio pipelines. - - Args: - audios (`np.ndarray`) - List of denoised audio samples of a NumPy array of shape `(batch_size, num_channels, sample_rate)`. - """ - - audios: np.ndarray - - -class DeprecatedPipelineMixin: - """ - A mixin that can be used to mark a pipeline as deprecated. - - Pipelines inheriting from this mixin will raise a warning when instantiated, indicating that they are deprecated - and won't receive updates past the specified version. Tests will be skipped for pipelines that inherit from this - mixin. - - Example usage: - ```python - class MyDeprecatedPipeline(DeprecatedPipelineMixin, DiffusionPipeline): - _last_supported_version = "0.20.0" - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - ``` - """ - - # Override this in the inheriting class to specify the last version that will support this pipeline - _last_supported_version = None - - def __init__(self, *args, **kwargs): - # Get the class name for the warning message - class_name = self.__class__.__name__ - - # Get the last supported version or use the current version if not specified - version_info = getattr(self.__class__, "_last_supported_version", __version__) - - # Raise a warning that this pipeline is deprecated - logger.warning( - f"The {class_name} has been deprecated and will not receive bug fixes or feature updates after Diffusers version {version_info}. " - ) - - # Call the parent class's __init__ method - super().__init__(*args, **kwargs) - - -class DiffusionPipeline(ConfigMixin, PushToHubMixin): - r""" - Base class for all pipelines. - - [`DiffusionPipeline`] stores all components (models, schedulers, and processors) for diffusion pipelines and - provides methods for loading, downloading and saving models. It also includes methods to: - - - move all PyTorch modules to the device of your choice - - enable/disable the progress bar for the denoising iteration - - Class attributes: - - - **config_name** (`str`) -- The configuration filename that stores the class and module names of all the - diffusion pipeline's components. - - **_optional_components** (`List[str]`) -- List of all optional components that don't have to be passed to the - pipeline to function (should be overridden by subclasses). - """ - - config_name = "model_index.json" - model_cpu_offload_seq = None - hf_device_map = None - _optional_components = [] - _exclude_from_cpu_offload = [] - _load_connected_pipes = False - _is_onnx = False - - def register_modules(self, **kwargs): - for name, module in kwargs.items(): - # retrieve library - if module is None or isinstance(module, (tuple, list)) and module[0] is None: - register_dict = {name: (None, None)} - else: - library, class_name = _fetch_class_library_tuple(module) - register_dict = {name: (library, class_name)} - - # save model index config - self.register_to_config(**register_dict) - - # set models - setattr(self, name, module) - - def __setattr__(self, name: str, value: Any): - if name in self.__dict__ and hasattr(self.config, name): - # We need to overwrite the config if name exists in config - if isinstance(getattr(self.config, name), (tuple, list)): - if value is not None and self.config[name][0] is not None: - class_library_tuple = _fetch_class_library_tuple(value) - else: - class_library_tuple = (None, None) - - self.register_to_config(**{name: class_library_tuple}) - else: - self.register_to_config(**{name: value}) - - super().__setattr__(name, value) - - def save_pretrained( - self, - save_directory: Union[str, os.PathLike], - safe_serialization: bool = True, - variant: Optional[str] = None, - max_shard_size: Optional[Union[int, str]] = None, - push_to_hub: bool = False, - **kwargs, - ): - """ - Save all saveable variables of the pipeline to a directory. A pipeline variable can be saved and loaded if its - class implements both a save and loading method. The pipeline is easily reloaded using the - [`~DiffusionPipeline.from_pretrained`] class method. - - Arguments: - save_directory (`str` or `os.PathLike`): - Directory to save a 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. Note that this limit will be decreased after a certain - period of time (starting from Oct 2024) to allow users to upgrade to the latest version of `diffusers`. - This is to establish a common default size for this argument across different libraries in the Hugging - Face ecosystem (`transformers`, and `accelerate`, for example). - push_to_hub (`bool`, *optional*, defaults to `False`): - 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*): - Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. - """ - model_index_dict = dict(self.config) - model_index_dict.pop("_class_name", None) - model_index_dict.pop("_diffusers_version", None) - model_index_dict.pop("_module", None) - model_index_dict.pop("_name_or_path", None) - - 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]) - repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id - - expected_modules, optional_kwargs = self._get_signature_keys(self) - - def is_saveable_module(name, value): - if name not in expected_modules: - return False - if name in self._optional_components and value[0] is None: - return False - return True - - model_index_dict = {k: v for k, v in model_index_dict.items() if is_saveable_module(k, v)} - for pipeline_component_name in model_index_dict.keys(): - sub_model = getattr(self, pipeline_component_name) - model_cls = sub_model.__class__ - - # Dynamo wraps the original model in a private class. - # I didn't find a public API to get the original class. - if is_compiled_module(sub_model): - sub_model = _unwrap_model(sub_model) - model_cls = sub_model.__class__ - - save_method_name = None - # search for the model's base class in LOADABLE_CLASSES - 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 {pipeline_component_name} as {library_classes} from {library_name}" - ) - - 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): - # if we found a suitable base class in LOADABLE_CLASSES then grab its save method - 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.{pipeline_component_name}={sub_model} of type {type(sub_model)} cannot be saved." - ) - # make sure that unsaveable components are not tried to be loaded afterward - self.register_to_config(**{pipeline_component_name: (None, None)}) - continue - - save_method = getattr(sub_model, save_method_name) - - # Call the save method with the argument safe_serialization only if it's supported - 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: - # max_shard_size is expected to not be None in ModelMixin - save_kwargs["max_shard_size"] = max_shard_size - - save_method(os.path.join(save_directory, pipeline_component_name), **save_kwargs) - - # finally save the config - self.save_config(save_directory) - - if push_to_hub: - # 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")) - - self._upload_folder( - save_directory, - repo_id, - token=token, - commit_message=commit_message, - create_pr=create_pr, - ) - - def to(self, *args, **kwargs) -> Self: - r""" - Performs Pipeline dtype and/or device conversion. A torch.dtype and torch.device are inferred from the - arguments of `self.to(*args, **kwargs).` - - - - If the pipeline already has the correct torch.dtype and torch.device, then it is returned as is. Otherwise, - the returned pipeline is a copy of self with the desired torch.dtype and torch.device. - - - - - Here are the ways to call `to`: - - - `to(dtype, silence_dtype_warnings=False) → DiffusionPipeline` to return a pipeline with the specified - [`dtype`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.dtype) - - `to(device, silence_dtype_warnings=False) → DiffusionPipeline` to return a pipeline with the specified - [`device`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device) - - `to(device=None, dtype=None, silence_dtype_warnings=False) → DiffusionPipeline` to return a pipeline with the - specified [`device`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device) and - [`dtype`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.dtype) - - Arguments: - dtype (`torch.dtype`, *optional*): - Returns a pipeline with the specified - [`dtype`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.dtype) - device (`torch.Device`, *optional*): - Returns a pipeline with the specified - [`device`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device) - silence_dtype_warnings (`str`, *optional*, defaults to `False`): - Whether to omit warnings if the target `dtype` is not compatible with the target `device`. - - Returns: - [`DiffusionPipeline`]: The pipeline converted to specified `dtype` and/or `dtype`. - """ - dtype = kwargs.pop("dtype", None) - device = kwargs.pop("device", None) - silence_dtype_warnings = kwargs.pop("silence_dtype_warnings", False) - - dtype_arg = None - device_arg = None - if len(args) == 1: - if isinstance(args[0], torch.dtype): - dtype_arg = args[0] - else: - device_arg = torch.device(args[0]) if args[0] is not None else None - elif len(args) == 2: - if isinstance(args[0], torch.dtype): - raise ValueError( - "When passing two arguments, make sure the first corresponds to `device` and the second to `dtype`." - ) - device_arg = torch.device(args[0]) if args[0] is not None else None - dtype_arg = args[1] - elif len(args) > 2: - raise ValueError("Please make sure to pass at most two arguments (`device` and `dtype`) `.to(...)`") - - if dtype is not None and dtype_arg is not None: - raise ValueError( - "You have passed `dtype` both as an argument and as a keyword argument. Please only pass one of the two." - ) - - dtype = dtype or dtype_arg - - if device is not None and device_arg is not None: - raise ValueError( - "You have passed `device` both as an argument and as a keyword argument. Please only pass one of the two." - ) - - device = device or device_arg - device_type = torch.device(device).type if device is not None else None - pipeline_has_bnb = any(any((_check_bnb_status(module))) for _, module in self.components.items()) - - # throw warning if pipeline is in "offloaded"-mode but user tries to manually set to GPU. - def module_is_sequentially_offloaded(module): - if not is_accelerate_available() or is_accelerate_version("<", "0.14.0"): - return False - - _, _, is_loaded_in_8bit_bnb = _check_bnb_status(module) - - if is_loaded_in_8bit_bnb: - return False - - return hasattr(module, "_hf_hook") and ( - isinstance(module._hf_hook, accelerate.hooks.AlignDevicesHook) - or hasattr(module._hf_hook, "hooks") - and isinstance(module._hf_hook.hooks[0], accelerate.hooks.AlignDevicesHook) - ) - - def module_is_offloaded(module): - if not is_accelerate_available() or is_accelerate_version("<", "0.17.0.dev0"): - return False - - return hasattr(module, "_hf_hook") and isinstance(module._hf_hook, accelerate.hooks.CpuOffload) - - # .to("cuda") would raise an error if the pipeline is sequentially offloaded, so we raise our own to make it clearer - 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 - 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." - ) - - if device_type in ["cuda", "xpu"]: - if pipeline_is_sequentially_offloaded and not pipeline_has_bnb: - raise ValueError( - "It seems like you have activated sequential model offloading by calling `enable_sequential_cpu_offload`, but are now attempting to move the pipeline to GPU. This is not compatible with offloading. Please, move your pipeline `.to('cpu')` or consider removing the move altogether if you use sequential offloading." - ) - # PR: https://github.com/huggingface/accelerate/pull/3223/ - elif pipeline_has_bnb and is_accelerate_version("<", "1.1.0.dev0"): - raise ValueError( - "You are trying to call `.to('cuda')` on a pipeline that has models quantized with `bitsandbytes`. Your current `accelerate` installation does not support it. Please upgrade the installation." - ) - - # Display a warning in this case (the operation succeeds but the benefits are lost) - pipeline_is_offloaded = any(module_is_offloaded(module) for _, module in self.components.items()) - if pipeline_is_offloaded and device_type in ["cuda", "xpu"]: - logger.warning( - f"It seems like you have activated model offloading by calling `enable_model_cpu_offload`, but are now manually moving the pipeline to GPU. It is strongly recommended against doing so as memory gains from offloading are likely to be lost. Offloading automatically takes care of moving the individual components {', '.join(self.components.keys())} to GPU when needed. To make sure offloading works as expected, you should consider moving the pipeline back to CPU: `pipeline.to('cpu')` or removing the move altogether if you use offloading." - ) - - # Enable generic support for Intel Gaudi accelerator using GPU/HPU migration - if device_type == "hpu" and kwargs.pop("hpu_migration", True) and is_hpu_available(): - os.environ["PT_HPU_GPU_MIGRATION"] = "1" - logger.debug("Environment variable set: PT_HPU_GPU_MIGRATION=1") - - import habana_frameworks.torch # noqa: F401 - - # HPU hardware check - if not (hasattr(torch, "hpu") and torch.hpu.is_available()): - raise ValueError("You are trying to call `.to('hpu')` but HPU device is unavailable.") - - os.environ["PT_HPU_MAX_COMPOUND_OP_SIZE"] = "1" - logger.debug("Environment variable set: PT_HPU_MAX_COMPOUND_OP_SIZE=1") - - 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)] - - is_offloaded = pipeline_is_offloaded or pipeline_is_sequentially_offloaded - for module in modules: - _, is_loaded_in_4bit_bnb, is_loaded_in_8bit_bnb = _check_bnb_status(module) - is_group_offloaded = self._maybe_raise_error_if_group_offload_active(module=module) - - if (is_loaded_in_4bit_bnb or is_loaded_in_8bit_bnb) and dtype is not None: - logger.warning( - 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: - 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}." - ) - - # Note: we also handle this at the ModelMixin level. The reason for doing it here too is that modeling - # components can be from outside diffusers too, but still have group offloading enabled. - if ( - self._maybe_raise_error_if_group_offload_active(raise_error=False, module=module) - and device is not None - ): - logger.warning( - f"The module '{module.__class__.__name__}' is group offloaded and moving it to {device} via `.to()` is not supported." - ) - - # This can happen for `transformer` models. CPU placement was added in - # 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) - elif not is_loaded_in_4bit_bnb and not is_loaded_in_8bit_bnb and not is_group_offloaded: - module.to(device, dtype) - - if ( - module.dtype == torch.float16 - and str(device) in ["cpu"] - and not silence_dtype_warnings - and not is_offloaded - ): - logger.warning( - "Pipelines loaded with `dtype=torch.float16` cannot run with `cpu` device. It" - " is not recommended to move them to `cpu` as running them will fail. Please make" - " sure to use an accelerator to run the pipeline in inference, due to the lack of" - " support for`float16` operations on this device in PyTorch. Please, remove the" - " `torch_dtype=torch.float16` argument, or use another device for inference." - ) - return self - - @property - def device(self) -> torch.device: - r""" - Returns: - `torch.device`: The torch device on which the pipeline is located. - """ - 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)] - - for module in modules: - return module.device - - return torch.device("cpu") - - @property - def dtype(self) -> torch.dtype: - r""" - Returns: - `torch.dtype`: The torch dtype on which the pipeline is located. - """ - 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)] - - for module in modules: - return module.dtype - - return torch.float32 - - @classmethod - @validate_hf_hub_args - def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs) -> Self: - r""" - Instantiate a PyTorch diffusion pipeline from pretrained pipeline weights. - - The pipeline is set in evaluation mode (`model.eval()`) by default. - - If you get the error message below, you need to finetune the weights for your downstream task: - - ``` - Some weights of UNet2DConditionModel were not initialized from the model checkpoint at stable-diffusion-v1-5/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: - - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated - You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. - ``` - - Parameters: - pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): - Can be either: - - - A string, the *repo id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained pipeline - hosted on the Hub. - - A path to a *directory* (for example `./my_pipeline_directory/`) containing pipeline weights - saved using - [`~DiffusionPipeline.save_pretrained`]. - - A path to a *directory* (for example `./my_pipeline_directory/`) containing a dduf file - torch_dtype (`torch.dtype` or `dict[str, Union[str, torch.dtype]]`, *optional*): - Override the default `torch.dtype` and load the model with another dtype. To load submodels with - different dtype pass a `dict` (for example `{'transformer': torch.bfloat16, 'vae': torch.float16}`). - Set the default dtype for unspecified components with `default` (for example `{'transformer': - torch.bfloat16, 'default': torch.float16}`). If a component is not specified and no default is set, - `torch.float32` is used. - custom_pipeline (`str`, *optional*): - - - - 🧪 This is an experimental feature and may change in the future. - - - - Can be either: - - - A string, the *repo id* (for example `hf-internal-testing/diffusers-dummy-pipeline`) of a custom - pipeline hosted on the Hub. The repository must contain a file called pipeline.py that defines - the custom pipeline. - - A string, the *file name* of a community pipeline hosted on GitHub under - [Community](https://github.com/huggingface/diffusers/tree/main/examples/community). Valid file - names must match the file name and not the pipeline script (`clip_guided_stable_diffusion` - instead of `clip_guided_stable_diffusion.py`). Community pipelines are always loaded from the - current main branch of GitHub. - - A path to a directory (`./my_pipeline_directory/`) containing a custom pipeline. The directory - must contain a file called `pipeline.py` that defines the custom pipeline. - - For more information on how to load and create custom pipelines, please have a look at [Loading and - Adding Custom - Pipelines](https://huggingface.co/docs/diffusers/using-diffusers/custom_pipeline_overview) - 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*): - Path to a directory where a downloaded pretrained model configuration is cached if the standard cache - is not used. - - 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`): - Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - custom_revision (`str`, *optional*): - The specific model version to use. It can be a branch name, a tag name, or a commit id similar to - `revision` when loading a custom pipeline from the Hub. Defaults to the latest stable 🤗 Diffusers - version. - mirror (`str`, *optional*): - 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`, *optional*): - Strategy that dictates how the different components of a pipeline should be placed on available - devices. Currently, only "balanced" `device_map` is supported. Check out - [this](https://huggingface.co/docs/diffusers/main/en/tutorials/inference_with_big_models#device-placement) - to know more. - max_memory (`Dict`, *optional*): - A dictionary device identifier for the maximum memory. Will default to the maximum memory available for - each GPU and the available CPU RAM if unset. - offload_folder (`str` or `os.PathLike`, *optional*): - The path to offload weights if device_map contains the value `"disk"`. - offload_state_dict (`bool`, *optional*): - If `True`, temporarily offloads the CPU state dict to the hard drive to avoid running out of CPU RAM if - the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to `True` - when there is some disk offload. - low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): - Speed up model loading only loading the pretrained weights and not initializing the weights. This also - tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. - Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this - argument to `True` will raise an error. - use_safetensors (`bool`, *optional*, defaults to `None`): - If set to `None`, the safetensors weights are downloaded if they're available **and** if the - safetensors library is installed. If set to `True`, the model is forcibly loaded from safetensors - weights. If set to `False`, safetensors weights are not loaded. - use_onnx (`bool`, *optional*, defaults to `None`): - If set to `True`, ONNX weights will always be downloaded if present. If set to `False`, ONNX weights - will never be downloaded. By default `use_onnx` defaults to the `_is_onnx` class attribute which is - `False` for non-ONNX pipelines and `True` for ONNX pipelines. ONNX weights include both files ending - with `.onnx` and `.pb`. - kwargs (remaining dictionary of keyword arguments, *optional*): - Can be used to overwrite load and saveable variables (the pipeline components of the specific pipeline - class). The overwritten components are passed directly to the pipelines `__init__` method. See example - below for more information. - variant (`str`, *optional*): - Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. - dduf_file(`str`, *optional*): - Load weights from the specified dduf file. - - - - To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with - `huggingface-cli login`. - - - - Examples: - - ```py - >>> from diffusers import DiffusionPipeline - - >>> # Download pipeline from huggingface.co and cache. - >>> pipeline = DiffusionPipeline.from_pretrained("CompVis/ldm-text2im-large-256") - - >>> # Download pipeline that requires an authorization token - >>> # For more information on access tokens, please refer to this section - >>> # of the documentation](https://huggingface.co/docs/hub/security-tokens) - >>> pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") - - >>> # Use a different scheduler - >>> from diffusers import LMSDiscreteScheduler - - >>> scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config) - >>> pipeline.scheduler = scheduler - ``` - """ - # Copy the kwargs to re-use during loading connected pipeline. - kwargs_copied = kwargs.copy() - - 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) - from_flax = kwargs.pop("from_flax", False) - torch_dtype = kwargs.pop("torch_dtype", None) - custom_pipeline = kwargs.pop("custom_pipeline", None) - custom_revision = kwargs.pop("custom_revision", None) - provider = kwargs.pop("provider", None) - sess_options = kwargs.pop("sess_options", None) - provider_options = kwargs.pop("provider_options", None) - device_map = kwargs.pop("device_map", None) - max_memory = kwargs.pop("max_memory", None) - offload_folder = kwargs.pop("offload_folder", None) - offload_state_dict = kwargs.pop("offload_state_dict", None) - low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) - variant = kwargs.pop("variant", None) - dduf_file = kwargs.pop("dduf_file", None) - use_safetensors = kwargs.pop("use_safetensors", None) - use_onnx = kwargs.pop("use_onnx", None) - load_connected_pipeline = kwargs.pop("load_connected_pipeline", False) - quantization_config = kwargs.pop("quantization_config", None) - - if torch_dtype is not None and not isinstance(torch_dtype, dict) and not isinstance(torch_dtype, torch.dtype): - torch_dtype = torch.float32 - logger.warning( - f"Passed `torch_dtype` {torch_dtype} is not a `torch.dtype`. Defaulting to `torch.float32`." - ) - - if low_cpu_mem_usage and not is_accelerate_available(): - low_cpu_mem_usage = False - logger.warning( - "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the" - " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install" - " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip" - " install accelerate\n```\n." - ) - - if quantization_config is not None and not isinstance(quantization_config, PipelineQuantizationConfig): - raise ValueError("`quantization_config` must be an instance of `PipelineQuantizationConfig`.") - - if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"): - raise NotImplementedError( - "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set" - " `low_cpu_mem_usage=False`." - ) - - if device_map is not None and not is_torch_version(">=", "1.9.0"): - raise NotImplementedError( - "Loading and dispatching requires torch >= 1.9.0. Please either update your PyTorch version or set" - " `device_map=None`." - ) - - if device_map is not None and not is_accelerate_available(): - raise NotImplementedError( - "Using `device_map` requires the `accelerate` library. Please install it using: `pip install accelerate`." - ) - - if device_map is not None and not isinstance(device_map, str): - raise ValueError("`device_map` must be a string.") - - if device_map is not None and device_map not in SUPPORTED_DEVICE_MAP: - raise NotImplementedError( - f"{device_map} not supported. Supported strategies are: {', '.join(SUPPORTED_DEVICE_MAP)}" - ) - - if device_map is not None and device_map in SUPPORTED_DEVICE_MAP: - if is_accelerate_version("<", "0.28.0"): - raise NotImplementedError("Device placement requires `accelerate` version `0.28.0` or later.") - - if low_cpu_mem_usage is False and device_map is not None: - raise ValueError( - f"You cannot set `low_cpu_mem_usage` to False while using device_map={device_map} for loading and" - " dispatching. Please make sure to set `low_cpu_mem_usage=True`." - ) - - if dduf_file: - if custom_pipeline: - raise NotImplementedError("Custom pipelines are not supported with DDUF at the moment.") - if load_connected_pipeline: - raise NotImplementedError("Connected pipelines are not supported with DDUF at the moment.") - - # 1. Download the checkpoints and configs - # use snapshot download here to get it working from from_pretrained - if not os.path.isdir(pretrained_model_name_or_path): - if pretrained_model_name_or_path.count("/") > 1: - raise ValueError( - f'The provided pretrained_model_name_or_path "{pretrained_model_name_or_path}"' - " is neither a valid local path nor a valid repo id. Please check the parameter." - ) - cached_folder = cls.download( - pretrained_model_name_or_path, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - from_flax=from_flax, - use_safetensors=use_safetensors, - use_onnx=use_onnx, - custom_pipeline=custom_pipeline, - custom_revision=custom_revision, - variant=variant, - dduf_file=dduf_file, - load_connected_pipeline=load_connected_pipeline, - **kwargs, - ) - else: - cached_folder = pretrained_model_name_or_path - - # The variant filenames can have the legacy sharding checkpoint format that we check and throw - # a warning if detected. - if variant is not None and _check_legacy_sharding_variant_format(folder=cached_folder, variant=variant): - warn_msg = ( - f"Warning: The repository contains sharded checkpoints for variant '{variant}' maybe in a deprecated format. " - "Please check your files carefully:\n\n" - "- Correct format example: diffusion_pytorch_model.fp16-00003-of-00003.safetensors\n" - "- Deprecated format example: diffusion_pytorch_model-00001-of-00002.fp16.safetensors\n\n" - "If you find any files in the deprecated format:\n" - "1. Remove all existing checkpoint files for this variant.\n" - "2. Re-obtain the correct files by running `save_pretrained()`.\n\n" - "This will ensure you're using the most up-to-date and compatible checkpoint format." - ) - logger.warning(warn_msg) - - dduf_entries = None - if dduf_file: - dduf_file_path = os.path.join(cached_folder, dduf_file) - dduf_entries = read_dduf_file(dduf_file_path) - # The reader contains already all the files needed, no need to check it again - cached_folder = "" - - config_dict = cls.load_config(cached_folder, dduf_entries=dduf_entries) - - if dduf_file: - _maybe_raise_error_for_incorrect_transformers(config_dict) - - # pop out "_ignore_files" as it is only needed for download - config_dict.pop("_ignore_files", None) - - # 2. Define which model components should load variants - # We retrieve the information by matching whether variant model checkpoints exist in the subfolders. - # Example: `diffusion_pytorch_model.safetensors` -> `diffusion_pytorch_model.fp16.safetensors` - # with variant being `"fp16"`. - model_variants = _identify_model_variants(folder=cached_folder, variant=variant, config=config_dict) - if len(model_variants) == 0 and variant is not None: - error_message = f"You are trying to load the model files of the `variant={variant}`, but no such modeling files are available." - raise ValueError(error_message) - - # 3. Load the pipeline class, if using custom module then load it from the hub - # if we load from explicit class, let's use it - custom_pipeline, custom_class_name = _resolve_custom_pipeline_and_cls( - folder=cached_folder, config=config_dict, custom_pipeline=custom_pipeline - ) - pipeline_class = _get_pipeline_class( - cls, - config=config_dict, - load_connected_pipeline=load_connected_pipeline, - custom_pipeline=custom_pipeline, - class_name=custom_class_name, - cache_dir=cache_dir, - revision=custom_revision, - ) - - if device_map is not None and pipeline_class._load_connected_pipes: - raise NotImplementedError("`device_map` is not yet supported for connected pipelines.") - - # DEPRECATED: To be removed in 1.0.0 - # we are deprecating the `StableDiffusionInpaintPipelineLegacy` pipeline which gets loaded - # when a user requests for a `StableDiffusionInpaintPipeline` with `diffusers` version being <= 0.5.1. - _maybe_raise_warning_for_inpainting( - pipeline_class=pipeline_class, - pretrained_model_name_or_path=pretrained_model_name_or_path, - config=config_dict, - ) - - # 4. Define expected modules given pipeline signature - # and define non-None initialized modules (=`init_kwargs`) - - # some modules can be passed directly to the init - # in this case they are already instantiated in `kwargs` - # extract them here - expected_modules, optional_kwargs = cls._get_signature_keys(pipeline_class) - expected_types = pipeline_class._get_signature_types() - passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs} - passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs} - init_dict, unused_kwargs, _ = pipeline_class.extract_init_dict(config_dict, **kwargs) - - # define init kwargs and make sure that optional component modules are filtered out - init_kwargs = { - k: init_dict.pop(k) - for k in optional_kwargs - if k in init_dict and k not in pipeline_class._optional_components - } - init_kwargs = {**init_kwargs, **passed_pipe_kwargs} - - # remove `null` components - def load_module(name, value): - if value[0] is None: - return False - if name in passed_class_obj and passed_class_obj[name] is None: - return False - return True - - init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)} - - # Special case: safety_checker must be loaded separately when using `from_flax` - if from_flax and "safety_checker" in init_dict and "safety_checker" not in passed_class_obj: - raise NotImplementedError( - "The safety checker cannot be automatically loaded when loading weights `from_flax`." - " Please, pass `safety_checker=None` to `from_pretrained`, and load the safety checker" - " separately if you need it." - ) - - # 5. Throw nice warnings / errors for fast accelerate loading - if len(unused_kwargs) > 0: - logger.warning( - f"Keyword arguments {unused_kwargs} are not expected by {pipeline_class.__name__} and will be ignored." - ) - - # import it here to avoid circular import - from diffusers import pipelines - - # 6. device map delegation - final_device_map = None - if device_map is not None: - final_device_map = _get_final_device_map( - device_map=device_map, - pipeline_class=pipeline_class, - passed_class_obj=passed_class_obj, - init_dict=init_dict, - library=library, - max_memory=max_memory, - torch_dtype=torch_dtype, - cached_folder=cached_folder, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - ) - - # 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..."): - # 7.1 device_map shenanigans - if final_device_map is not None and len(final_device_map) > 0: - component_device = final_device_map.get(name, None) - if component_device is not None: - current_device_map = {"": component_device} - else: - current_device_map = None - - # 7.2 - now that JAX/Flax is an official framework of the library, we might load from Flax names - class_name = class_name[4:] if class_name.startswith("Flax") else class_name - - # 7.3 Define all importable classes - is_pipeline_module = hasattr(pipelines, library_name) - importable_classes = ALL_IMPORTABLE_CLASSES - loaded_sub_model = None - - # 7.4 Use passed sub model or load class_name from library_name - if name in passed_class_obj: - # if the model is in a pipeline module, then we load it from the pipeline - # check that passed_class_obj has correct parent class - maybe_raise_or_warn( - library_name, library, class_name, importable_classes, passed_class_obj, name, is_pipeline_module - ) - - loaded_sub_model = passed_class_obj[name] - else: - # load sub model - sub_model_dtype = ( - torch_dtype.get(name, torch_dtype.get("default", torch.float32)) - if isinstance(torch_dtype, dict) - else torch_dtype - ) - loaded_sub_model = load_sub_model( - library_name=library_name, - class_name=class_name, - importable_classes=importable_classes, - pipelines=pipelines, - is_pipeline_module=is_pipeline_module, - pipeline_class=pipeline_class, - torch_dtype=sub_model_dtype, - provider=provider, - sess_options=sess_options, - device_map=current_device_map, - max_memory=max_memory, - offload_folder=offload_folder, - offload_state_dict=offload_state_dict, - model_variants=model_variants, - name=name, - from_flax=from_flax, - variant=variant, - low_cpu_mem_usage=low_cpu_mem_usage, - cached_folder=cached_folder, - use_safetensors=use_safetensors, - dduf_entries=dduf_entries, - provider_options=provider_options, - quantization_config=quantization_config, - ) - logger.info( - f"Loaded {name} as {class_name} from `{name}` subfolder of {pretrained_model_name_or_path}." - ) - - init_kwargs[name] = loaded_sub_model # UNet(...), # DiffusionSchedule(...) - - # 8. Handle connected pipelines. - if pipeline_class._load_connected_pipes and os.path.isfile(os.path.join(cached_folder, "README.md")): - init_kwargs = _update_init_kwargs_with_connected_pipeline( - init_kwargs=init_kwargs, - passed_pipe_kwargs=passed_pipe_kwargs, - passed_class_objs=passed_class_obj, - folder=cached_folder, - **kwargs_copied, - ) - - # 9. Potentially add passed objects if expected - missing_modules = set(expected_modules) - set(init_kwargs.keys()) - passed_modules = list(passed_class_obj.keys()) - optional_modules = pipeline_class._optional_components - if len(missing_modules) > 0 and missing_modules <= set(passed_modules + optional_modules): - for module in missing_modules: - init_kwargs[module] = passed_class_obj.get(module, None) - elif len(missing_modules) > 0: - passed_modules = set(list(init_kwargs.keys()) + list(passed_class_obj.keys())) - set(optional_kwargs) - raise ValueError( - f"Pipeline {pipeline_class} expected {expected_modules}, but only {passed_modules} were passed." - ) - - # 10. Type checking init arguments - for kw, arg in init_kwargs.items(): - # Too complex to validate with type annotation alone - if "scheduler" in kw: - continue - # Many tokenizer annotations don't include its "Fast" variant, so skip this - # e.g T5Tokenizer but not T5TokenizerFast - elif "tokenizer" in kw: - continue - elif ( - arg is not None # Skip if None - and not expected_types[kw] == (inspect.Signature.empty,) # Skip if no type annotations - and not _is_valid_type(arg, expected_types[kw]) # Check type - ): - logger.warning(f"Expected types for {kw}: {expected_types[kw]}, got {_get_detailed_type(arg)}.") - - # 11. Instantiate the pipeline - model = pipeline_class(**init_kwargs) - - # 12. Save where the model was instantiated from - model.register_to_config(_name_or_path=pretrained_model_name_or_path) - if device_map is not None: - setattr(model, "hf_device_map", final_device_map) - return model - - @property - def name_or_path(self) -> str: - return getattr(self.config, "_name_or_path", None) - - @property - def _execution_device(self): - r""" - Returns the device on which the pipeline's models will be executed. After calling - [`~DiffusionPipeline.enable_sequential_cpu_offload`] the execution device can only be inferred from - Accelerate's module hooks. - """ - from ..hooks.group_offloading import _get_group_onload_device - - # When apply group offloading at the leaf_level, we're in the same situation as accelerate's sequential - # offloading. We need to return the onload device of the group offloading hooks so that the intermediates - # required for computation (latents, prompt embeddings, etc.) can be created on the correct device. - for name, model in self.components.items(): - if not isinstance(model, torch.nn.Module): - continue - try: - return _get_group_onload_device(model) - except ValueError: - pass - - for name, model in self.components.items(): - if not isinstance(model, torch.nn.Module) or name in self._exclude_from_cpu_offload: - continue - - if not hasattr(model, "_hf_hook"): - return self.device - for module in model.modules(): - if ( - hasattr(module, "_hf_hook") - and hasattr(module._hf_hook, "execution_device") - and module._hf_hook.execution_device is not None - ): - return torch.device(module._hf_hook.execution_device) - return self.device - - def remove_all_hooks(self): - r""" - Removes all hooks that were added when using `enable_sequential_cpu_offload` or `enable_model_cpu_offload`. - """ - for _, model in self.components.items(): - if isinstance(model, torch.nn.Module) and hasattr(model, "_hf_hook"): - 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): - 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 - `forward` method is called, and the model remains in accelerator until the next model runs. Memory savings are - lower than with `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution - of the `unet`. - - Arguments: - gpu_id (`int`, *optional*): - The ID of the accelerator that shall be used in inference. If not specified, it will default to 0. - device (`torch.Device` or `str`, *optional*, defaults to None): - The PyTorch device type of the accelerator that shall be used in inference. If not specified, it will - automatically detect the available accelerator and use. - """ - 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 - 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()`." - ) - - if self.model_cpu_offload_seq is None: - raise ValueError( - "Model CPU offload cannot be enabled because no `model_cpu_offload_seq` class attribute is set." - ) - - if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"): - from accelerate import cpu_offload_with_hook - else: - raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.") - - self.remove_all_hooks() - - if device is None: - device = get_device() - if device == "cpu": - raise RuntimeError("`enable_model_cpu_offload` requires accelerator, but not found") - - torch_device = torch.device(device) - device_index = torch_device.index - - if gpu_id is not None and device_index is not None: - raise ValueError( - f"You have passed both `gpu_id`={gpu_id} and an index as part of the passed device `device`={device}" - f"Cannot pass both. Please make sure to either not define `gpu_id` or not pass the index as part of the device: `device`={torch_device.type}" - ) - - # _offload_gpu_id should be set to passed gpu_id (or id in passed `device`) or default to previously set id or default to 0 - self._offload_gpu_id = gpu_id or torch_device.index or getattr(self, "_offload_gpu_id", 0) - - device_type = torch_device.type - device = torch.device(f"{device_type}:{self._offload_gpu_id}") - self._offload_device = device - - self.to("cpu", silence_dtype_warnings=True) - empty_device_cache(device.type) - - all_model_components = {k: v for k, v in self.components.items() if isinstance(v, torch.nn.Module)} - - self._all_hooks = [] - hook = None - for model_str in self.model_cpu_offload_seq.split("->"): - model = all_model_components.pop(model_str, None) - - if not isinstance(model, torch.nn.Module): - continue - - # 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: - logger.info( - f"Skipping the hook placement for the {model.__class__.__name__} as it is loaded in `bitsandbytes` 8bit." - ) - continue - - _, hook = cpu_offload_with_hook(model, device, prev_module_hook=hook) - self._all_hooks.append(hook) - - # CPU offload models that are not in the seq chain unless they are explicitly excluded - # these models will stay on CPU until maybe_free_model_hooks is called - # some models cannot be in the seq chain because they are iteratively called, such as controlnet - for name, model in all_model_components.items(): - if not isinstance(model, torch.nn.Module): - continue - - if name in self._exclude_from_cpu_offload: - model.to(device) - else: - _, hook = cpu_offload_with_hook(model, device) - self._all_hooks.append(hook) - - def maybe_free_model_hooks(self): - r""" - Method that performs the following: - - Offloads all components. - - Removes all model hooks that were added when using `enable_model_cpu_offload`, and then applies them again. - In case the model has not been offloaded, this function is a no-op. - - Resets stateful diffusers hooks of denoiser components if they were added with - [`~hooks.HookRegistry.register_hook`]. - - Make sure to add this function to the end of the `__call__` function of your pipeline so that it functions - correctly when applying `enable_model_cpu_offload`. - """ - for component in self.components.values(): - if hasattr(component, "_reset_stateful_cache"): - component._reset_stateful_cache() - - if not hasattr(self, "_all_hooks") or len(self._all_hooks) == 0: - # `enable_model_cpu_offload` has not be called, so silently do nothing - return - - # 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): - 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 - and then moved to `torch.device('meta')` and loaded to accelerator only when their specific submodule has its - `forward` method called. Offloading happens on a submodule basis. Memory savings are higher than with - `enable_model_cpu_offload`, but performance is lower. - - Arguments: - gpu_id (`int`, *optional*): - The ID of the accelerator that shall be used in inference. If not specified, it will default to 0. - device (`torch.Device` or `str`, *optional*, defaults to None): - The PyTorch device type of the accelerator that shall be used in inference. If not specified, it will - automatically detect the available accelerator and use. - """ - self._maybe_raise_error_if_group_offload_active(raise_error=True) - - if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"): - from accelerate import cpu_offload - else: - 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 - 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()`." - ) - - if device is None: - device = get_device() - if device == "cpu": - raise RuntimeError("`enable_sequential_cpu_offload` requires accelerator, but not found") - - torch_device = torch.device(device) - device_index = torch_device.index - - if gpu_id is not None and device_index is not None: - raise ValueError( - f"You have passed both `gpu_id`={gpu_id} and an index as part of the passed device `device`={device}" - f"Cannot pass both. Please make sure to either not define `gpu_id` or not pass the index as part of the device: `device`={torch_device.type}" - ) - - # _offload_gpu_id should be set to passed gpu_id (or id in passed `device`) or default to previously set id or default to 0 - self._offload_gpu_id = gpu_id or torch_device.index or getattr(self, "_offload_gpu_id", 0) - - device_type = torch_device.type - device = torch.device(f"{device_type}:{self._offload_gpu_id}") - self._offload_device = device - - if self.device.type != "cpu": - orig_device_type = self.device.type - self.to("cpu", silence_dtype_warnings=True) - empty_device_cache(orig_device_type) - - for name, model in self.components.items(): - if not isinstance(model, torch.nn.Module): - continue - - if name in self._exclude_from_cpu_offload: - model.to(device) - else: - # make sure to offload buffers if not all high level weights - # are of type nn.Module - offload_buffers = len(model._parameters) > 0 - cpu_offload(model, device, offload_buffers=offload_buffers) - - def reset_device_map(self): - r""" - Resets the device maps (if any) to None. - """ - if self.hf_device_map is None: - return - else: - self.remove_all_hooks() - for name, component in self.components.items(): - if isinstance(component, torch.nn.Module): - component.to("cpu") - self.hf_device_map = None - - @classmethod - @validate_hf_hub_args - def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]: - r""" - Download and cache a PyTorch diffusion pipeline from pretrained pipeline weights. - - Parameters: - pretrained_model_name (`str` or `os.PathLike`, *optional*): - A string, the *repository id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained pipeline - hosted on the Hub. - custom_pipeline (`str`, *optional*): - Can be either: - - - A string, the *repository id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained - pipeline hosted on the Hub. The repository must contain a file called `pipeline.py` that defines - the custom pipeline. - - - A string, the *file name* of a community pipeline hosted on GitHub under - [Community](https://github.com/huggingface/diffusers/tree/main/examples/community). Valid file - names must match the file name and not the pipeline script (`clip_guided_stable_diffusion` - instead of `clip_guided_stable_diffusion.py`). Community pipelines are always loaded from the - current `main` branch of GitHub. - - - A path to a *directory* (`./my_pipeline_directory/`) containing a custom pipeline. The directory - must contain a file called `pipeline.py` that defines the custom pipeline. - - - - 🧪 This is an experimental feature and may change in the future. - - - - For more information on how to load and create custom pipelines, take a look at [How to contribute a - community pipeline](https://huggingface.co/docs/diffusers/main/en/using-diffusers/contribute_pipeline). - - 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*): - 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`): - Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - custom_revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, or a commit id similar to - `revision` when loading a custom pipeline from the Hub. It can be a 🤗 Diffusers version when loading a - custom pipeline from GitHub, otherwise it defaults to `"main"` when loading from the Hub. - mirror (`str`, *optional*): - 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. - variant (`str`, *optional*): - Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when - loading `from_flax`. - dduf_file(`str`, *optional*): - Load weights from the specified DDUF file. - use_safetensors (`bool`, *optional*, defaults to `None`): - If set to `None`, the safetensors weights are downloaded if they're available **and** if the - safetensors library is installed. If set to `True`, the model is forcibly loaded from safetensors - weights. If set to `False`, safetensors weights are not loaded. - use_onnx (`bool`, *optional*, defaults to `False`): - If set to `True`, ONNX weights will always be downloaded if present. If set to `False`, ONNX weights - will never be downloaded. By default `use_onnx` defaults to the `_is_onnx` class attribute which is - `False` for non-ONNX pipelines and `True` for ONNX pipelines. ONNX weights include both files ending - with `.onnx` and `.pb`. - trust_remote_code (`bool`, *optional*, defaults to `False`): - Whether or not to allow for custom pipelines and components defined on the Hub in their own files. This - option should only be set to `True` for repositories you trust and in which you have read the code, as - it will execute code present on the Hub on your local machine. - - Returns: - `os.PathLike`: - A path to the downloaded pipeline. - - - - To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with - `huggingface-cli login`. - - - - """ - 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) - from_flax = kwargs.pop("from_flax", False) - custom_pipeline = kwargs.pop("custom_pipeline", None) - custom_revision = kwargs.pop("custom_revision", None) - variant = kwargs.pop("variant", None) - use_safetensors = kwargs.pop("use_safetensors", None) - 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) - - if dduf_file: - if custom_pipeline: - raise NotImplementedError("Custom pipelines are not supported with DDUF at the moment.") - if load_connected_pipeline: - raise NotImplementedError("Connected pipelines are not supported with DDUF at the moment.") - return _download_dduf_file( - pretrained_model_name=pretrained_model_name, - dduf_file=dduf_file, - pipeline_class_name=cls.__name__, - cache_dir=cache_dir, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - ) - - allow_pickle = True if (use_safetensors is None or use_safetensors is False) else False - use_safetensors = use_safetensors if use_safetensors is not None else True - - allow_patterns = None - ignore_patterns = None - - model_info_call_error: Optional[Exception] = None - if not local_files_only: - try: - info = model_info(pretrained_model_name, token=token, revision=revision) - except (HTTPError, OfflineModeIsEnabled, requests.ConnectionError) as e: - logger.warning(f"Couldn't connect to the Hub: {e}.\nWill try to load from local cache.") - local_files_only = True - model_info_call_error = e # save error to reraise it if model is not cached locally - - if not local_files_only: - config_file = hf_hub_download( - pretrained_model_name, - cls.config_name, - cache_dir=cache_dir, - revision=revision, - proxies=proxies, - force_download=force_download, - token=token, - ) - config_dict = cls._dict_from_json_file(config_file) - ignore_filenames = config_dict.pop("_ignore_files", []) - - filenames = {sibling.rfilename for sibling in info.siblings} - if variant is not None and _check_legacy_sharding_variant_format(filenames=filenames, variant=variant): - warn_msg = ( - f"Warning: The repository contains sharded checkpoints for variant '{variant}' maybe in a deprecated format. " - "Please check your files carefully:\n\n" - "- Correct format example: diffusion_pytorch_model.fp16-00003-of-00003.safetensors\n" - "- Deprecated format example: diffusion_pytorch_model-00001-of-00002.fp16.safetensors\n\n" - "If you find any files in the deprecated format:\n" - "1. Remove all existing checkpoint files for this variant.\n" - "2. Re-obtain the correct files by running `save_pretrained()`.\n\n" - "This will ensure you're using the most up-to-date and compatible checkpoint format." - ) - logger.warning(warn_msg) - - filenames = set(filenames) - set(ignore_filenames) - if revision in DEPRECATED_REVISION_ARGS and version.parse( - version.parse(__version__).base_version - ) >= version.parse("0.22.0"): - warn_deprecated_model_variant(pretrained_model_name, token, variant, revision, filenames) - - custom_components, folder_names = _get_custom_components_and_folders( - pretrained_model_name, config_dict, filenames, variant - ) - custom_class_name = None - if custom_pipeline is None and isinstance(config_dict["_class_name"], (list, tuple)): - custom_pipeline = config_dict["_class_name"][0] - custom_class_name = config_dict["_class_name"][1] - - load_pipe_from_hub = custom_pipeline is not None and f"{custom_pipeline}.py" in filenames - load_components_from_hub = len(custom_components) > 0 - - if load_pipe_from_hub and not trust_remote_code: - raise ValueError( - f"The repository for {pretrained_model_name} contains custom code in {custom_pipeline}.py which must be executed to correctly " - f"load the model. You can inspect the repository content at https://hf.co/{pretrained_model_name}/blob/main/{custom_pipeline}.py.\n" - f"Please pass the argument `trust_remote_code=True` to allow custom code to be run." - ) - - if load_components_from_hub and not trust_remote_code: - raise ValueError( - f"The repository for {pretrained_model_name} contains custom code in {'.py, '.join([os.path.join(k, v) for k, v in custom_components.items()])} which must be executed to correctly " - f"load the model. You can inspect the repository content at {', '.join([f'https://hf.co/{pretrained_model_name}/{k}/{v}.py' for k, v in custom_components.items()])}.\n" - f"Please pass the argument `trust_remote_code=True` to allow custom code to be run." - ) - - # retrieve passed components that should not be downloaded - pipeline_class = _get_pipeline_class( - cls, - config_dict, - load_connected_pipeline=load_connected_pipeline, - custom_pipeline=custom_pipeline, - repo_id=pretrained_model_name if load_pipe_from_hub else None, - hub_revision=revision, - class_name=custom_class_name, - cache_dir=cache_dir, - revision=custom_revision, - ) - expected_components, _ = cls._get_signature_keys(pipeline_class) - passed_components = [k for k in expected_components if k in kwargs] - - # retrieve the names of the folders containing model weights - model_folder_names = { - os.path.split(f)[0] for f in filter_model_files(filenames) if os.path.split(f)[0] in folder_names - } - # retrieve all patterns that should not be downloaded and error out when needed - ignore_patterns = _get_ignore_patterns( - passed_components, - model_folder_names, - filenames, - use_safetensors, - from_flax, - allow_pickle, - use_onnx, - pipeline_class._is_onnx, - variant, - ) - - model_filenames, variant_filenames = variant_compatible_siblings( - filenames, variant=variant, ignore_patterns=ignore_patterns - ) - - # all filenames compatible with variant will be added - allow_patterns = list(model_filenames) - - # allow all patterns from non-model folders - # this enables downloading schedulers, tokenizers, ... - allow_patterns += [f"{k}/*" for k in folder_names if k not in model_folder_names] - # add custom component files - allow_patterns += [f"{k}/{f}.py" for k, f in custom_components.items()] - # add custom pipeline file - allow_patterns += [f"{custom_pipeline}.py"] if f"{custom_pipeline}.py" in filenames else [] - # also allow downloading config.json files with the model - allow_patterns += [os.path.join(k, "config.json") for k in model_folder_names] - allow_patterns += [ - SCHEDULER_CONFIG_NAME, - CONFIG_NAME, - cls.config_name, - CUSTOM_PIPELINE_FILE_NAME, - ] - - # Don't download any objects that are passed - allow_patterns = [ - p for p in allow_patterns if not (len(p.split("/")) == 2 and p.split("/")[0] in passed_components) - ] - - if pipeline_class._load_connected_pipes: - allow_patterns.append("README.md") - - # Don't download index files of forbidden patterns either - ignore_patterns = ignore_patterns + [f"{i}.index.*json" for i in ignore_patterns] - re_ignore_pattern = [re.compile(fnmatch.translate(p)) for p in ignore_patterns] - re_allow_pattern = [re.compile(fnmatch.translate(p)) for p in allow_patterns] - - expected_files = [f for f in filenames if not any(p.match(f) for p in re_ignore_pattern)] - expected_files = [f for f in expected_files if any(p.match(f) for p in re_allow_pattern)] - - snapshot_folder = Path(config_file).parent - pipeline_is_cached = all((snapshot_folder / f).is_file() for f in expected_files) - - if pipeline_is_cached and not force_download: - # if the pipeline is cached, we can directly return it - # else call snapshot_download - return snapshot_folder - - user_agent = {"pipeline_class": cls.__name__} - if custom_pipeline is not None and not custom_pipeline.endswith(".py"): - user_agent["custom_pipeline"] = custom_pipeline - - # download all allow_patterns - ignore_patterns - try: - cached_folder = snapshot_download( - pretrained_model_name, - cache_dir=cache_dir, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - allow_patterns=allow_patterns, - ignore_patterns=ignore_patterns, - user_agent=user_agent, - ) - - cls_name = cls.load_config(os.path.join(cached_folder, "model_index.json")).get("_class_name", None) - cls_name = cls_name[4:] if isinstance(cls_name, str) and cls_name.startswith("Flax") else cls_name - - diffusers_module = importlib.import_module(__name__.split(".")[0]) - pipeline_class = getattr(diffusers_module, cls_name, None) if isinstance(cls_name, str) else None - - if pipeline_class is not None and pipeline_class._load_connected_pipes: - modelcard = ModelCard.load(os.path.join(cached_folder, "README.md")) - connected_pipes = sum([getattr(modelcard.data, k, []) for k in CONNECTED_PIPES_KEYS], []) - for connected_pipe_repo_id in connected_pipes: - download_kwargs = { - "cache_dir": cache_dir, - "force_download": force_download, - "proxies": proxies, - "local_files_only": local_files_only, - "token": token, - "variant": variant, - "use_safetensors": use_safetensors, - } - DiffusionPipeline.download(connected_pipe_repo_id, **download_kwargs) - - return cached_folder - - except FileNotFoundError: - # Means we tried to load pipeline with `local_files_only=True` but the files have not been found in local cache. - # This can happen in two cases: - # 1. If the user passed `local_files_only=True` => we raise the error directly - # 2. If we forced `local_files_only=True` when `model_info` failed => we raise the initial error - if model_info_call_error is None: - # 1. user passed `local_files_only=True` - raise - else: - # 2. we forced `local_files_only=True` when `model_info` failed - raise EnvironmentError( - f"Cannot load model {pretrained_model_name}: model is not cached locally and an error occurred" - " while trying to fetch metadata from the Hub. Please check out the root cause in the stacktrace" - " above." - ) from model_info_call_error - - @classmethod - def _get_signature_keys(cls, obj): - parameters = inspect.signature(obj.__init__).parameters - required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty} - optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty}) - expected_modules = set(required_parameters.keys()) - {"self"} - - optional_names = list(optional_parameters) - for name in optional_names: - if name in cls._optional_components: - expected_modules.add(name) - optional_parameters.remove(name) - - return sorted(expected_modules), sorted(optional_parameters) - - @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,) - else: - logger.warning(f"cannot get type annotation for Parameter {k} of {cls}.") - return signature_types - - @property - 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. - - Returns (`dict`): - A dictionary containing all the modules needed to initialize the pipeline. - - Examples: - - ```py - >>> from diffusers import ( - ... StableDiffusionPipeline, - ... StableDiffusionImg2ImgPipeline, - ... StableDiffusionInpaintPipeline, - ... ) - - >>> text2img = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") - >>> img2img = StableDiffusionImg2ImgPipeline(**text2img.components) - >>> inpaint = StableDiffusionInpaintPipeline(**text2img.components) - ``` - """ - expected_modules, optional_parameters = self._get_signature_keys(self) - components = { - k: getattr(self, k) for k in self.config.keys() if not k.startswith("_") and k not in optional_parameters - } - - actual = sorted(set(components.keys())) - expected = sorted(expected_modules) - if actual != expected: - raise ValueError( - f"{self} has been incorrectly initialized or {self.__class__} is incorrectly implemented. Expected" - f" {expected} to be defined, but {actual} are defined." - ) - - return components - - @staticmethod - def numpy_to_pil(images): - """ - Convert a NumPy image or a batch of images to a PIL image. - """ - return numpy_to_pil(images) - - @torch.compiler.disable - def progress_bar(self, iterable=None, total=None): - if not hasattr(self, "_progress_bar_config"): - self._progress_bar_config = {} - elif not isinstance(self._progress_bar_config, dict): - raise ValueError( - f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}." - ) - - if iterable is not None: - return tqdm(iterable, **self._progress_bar_config) - elif total is not None: - return tqdm(total=total, **self._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): - 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 - up during training is not guaranteed. - - - - ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient attention takes - precedent. - - - - Parameters: - attention_op (`Callable`, *optional*): - Override the default `None` operator for use as `op` argument to the - [`memory_efficient_attention()`](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.memory_efficient_attention) - function of xFormers. - - Examples: - - ```py - >>> import torch - >>> from diffusers import DiffusionPipeline - >>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp - - >>> pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16) - >>> pipe = pipe.to("cuda") - >>> pipe.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp) - >>> # Workaround for not accepting attention shape using VAE for Flash Attention - >>> pipe.vae.enable_xformers_memory_efficient_attention(attention_op=None) - ``` - """ - self.set_use_memory_efficient_attention_xformers(True, attention_op) - - def disable_xformers_memory_efficient_attention(self): - r""" - Disable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). - """ - self.set_use_memory_efficient_attention_xformers(False) - - def set_use_memory_efficient_attention_xformers( - self, valid: bool, attention_op: Optional[Callable] = None - ) -> None: - # Recursively walk through all the children. - # Any children which exposes the set_use_memory_efficient_attention_xformers method - # gets the message - def fn_recursive_set_mem_eff(module: torch.nn.Module): - if hasattr(module, "set_use_memory_efficient_attention_xformers"): - module.set_use_memory_efficient_attention_xformers(valid, attention_op) - - for child in module.children(): - fn_recursive_set_mem_eff(child) - - 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)] - - for module in modules: - fn_recursive_set_mem_eff(module) - - def enable_attention_slicing(self, slice_size: Optional[Union[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 - sequentially over each head. This is useful to save some memory in exchange for a small speed decrease. - - - - ⚠️ Don't enable attention slicing if you're already using `scaled_dot_product_attention` (SDPA) from PyTorch - 2.0 or xFormers. These attention computations are already very memory efficient so you won't need to enable - this function. If you enable attention slicing with SDPA or xFormers, it can lead to serious slow downs! - - - - Args: - slice_size (`str` or `int`, *optional*, defaults to `"auto"`): - When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If - `"max"`, maximum amount of memory will be saved by running only one slice at a time. If a number is - provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` - must be a multiple of `slice_size`. - - Examples: - - ```py - >>> import torch - >>> from diffusers import StableDiffusionPipeline - - >>> pipe = StableDiffusionPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", - ... torch_dtype=torch.float16, - ... use_safetensors=True, - ... ) - - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> pipe.enable_attention_slicing() - >>> image = pipe(prompt).images[0] - ``` - """ - self.set_attention_slice(slice_size) - - def disable_attention_slicing(self): - r""" - Disable sliced attention computation. If `enable_attention_slicing` was previously called, attention is - computed in one step. - """ - # set slice_size = `None` to disable `attention slicing` - self.enable_attention_slicing(None) - - def set_attention_slice(self, slice_size: Optional[int]): - 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")] - - for module in modules: - module.set_attention_slice(slice_size) - - @classmethod - def from_pipe(cls, pipeline, **kwargs): - r""" - Create a new pipeline from a given pipeline. This method is useful to create a new pipeline from the existing - pipeline components without reallocating additional memory. - - Arguments: - pipeline (`DiffusionPipeline`): - The pipeline from which to create a new pipeline. - - Returns: - `DiffusionPipeline`: - A new pipeline with the same weights and configurations as `pipeline`. - - Examples: - - ```py - >>> from diffusers import StableDiffusionPipeline, StableDiffusionSAGPipeline - - >>> pipe = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") - >>> new_pipe = StableDiffusionSAGPipeline.from_pipe(pipe) - ``` - """ - - original_config = dict(pipeline.config) - torch_dtype = kwargs.pop("torch_dtype", torch.float32) - - # derive the pipeline class to instantiate - custom_pipeline = kwargs.pop("custom_pipeline", None) - custom_revision = kwargs.pop("custom_revision", None) - - if custom_pipeline is not None: - pipeline_class = _get_custom_pipeline_class(custom_pipeline, revision=custom_revision) - else: - pipeline_class = cls - - expected_modules, optional_kwargs = cls._get_signature_keys(pipeline_class) - # true_optional_modules are optional components with default value in signature so it is ok not to pass them to `__init__` - # e.g. `image_encoder` for StableDiffusionPipeline - parameters = inspect.signature(cls.__init__).parameters - true_optional_modules = set( - {k for k, v in parameters.items() if v.default != inspect._empty and k in expected_modules} - ) - - # get the class of each component based on its type hint - # e.g. {"unet": UNet2DConditionModel, "text_encoder": CLIPTextMode} - component_types = pipeline_class._get_signature_types() - - pretrained_model_name_or_path = original_config.pop("_name_or_path", None) - # allow users pass modules in `kwargs` to override the original pipeline's components - passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs} - - original_class_obj = {} - for name, component in pipeline.components.items(): - if name in expected_modules and name not in passed_class_obj: - # for model components, we will not switch over if the class does not matches the type hint in the new pipeline's signature - if ( - not isinstance(component, ModelMixin) - or type(component) in component_types[name] - or (component is None and name in cls._optional_components) - ): - original_class_obj[name] = component - else: - logger.warning( - f"component {name} is not switched over to new pipeline because type does not match the expected." - f" {name} is {type(component)} while the new pipeline expect {component_types[name]}." - f" please pass the component of the correct type to the new pipeline. `from_pipe(..., {name}={name})`" - ) - - # allow users pass optional kwargs to override the original pipelines config attribute - passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs} - original_pipe_kwargs = { - k: original_config[k] - for k in original_config.keys() - if k in optional_kwargs and k not in passed_pipe_kwargs - } - - # config attribute that were not expected by pipeline is stored as its private attribute - # (i.e. when the original pipeline was also instantiated with `from_pipe` from another pipeline that has this config) - # in this case, we will pass them as optional arguments if they can be accepted by the new pipeline - additional_pipe_kwargs = [ - k[1:] - for k in original_config.keys() - if k.startswith("_") and k[1:] in optional_kwargs and k[1:] not in passed_pipe_kwargs - ] - for k in additional_pipe_kwargs: - original_pipe_kwargs[k] = original_config.pop(f"_{k}") - - pipeline_kwargs = { - **passed_class_obj, - **original_class_obj, - **passed_pipe_kwargs, - **original_pipe_kwargs, - **kwargs, - } - - # store unused config as private attribute in the new pipeline - unused_original_config = { - f"{'' if k.startswith('_') else '_'}{k}": v for k, v in original_config.items() if k not in pipeline_kwargs - } - - missing_modules = ( - set(expected_modules) - - set(pipeline._optional_components) - - set(pipeline_kwargs.keys()) - - set(true_optional_modules) - ) - - if len(missing_modules) > 0: - raise ValueError( - f"Pipeline {pipeline_class} expected {expected_modules}, but only {set(list(passed_class_obj.keys()) + list(original_class_obj.keys()))} were passed" - ) - - new_pipeline = pipeline_class(**pipeline_kwargs) - if pretrained_model_name_or_path is not None: - new_pipeline.register_to_config(_name_or_path=pretrained_model_name_or_path) - new_pipeline.register_to_config(**unused_original_config) - - if torch_dtype is not None: - new_pipeline.to(dtype=torch_dtype) - - return new_pipeline - - def _maybe_raise_error_if_group_offload_active( - self, raise_error: bool = False, module: Optional[torch.nn.Module] = None - ) -> bool: - from ..hooks.group_offloading import _is_group_offload_enabled - - components = self.components.values() if module is None else [module] - components = [component for component in components if isinstance(component, torch.nn.Module)] - for component in components: - if _is_group_offload_enabled(component): - if raise_error: - raise ValueError( - "You are trying to apply model/sequential CPU offloading to a pipeline that contains components " - "with group offloading enabled. This is not supported. Please disable group offloading for " - "components of the pipeline to use other offloading methods." - ) - return True - return False - - -class StableDiffusionMixin: - r""" - Helper for DiffusionPipeline with vae and unet.(mainly for LDM such as stable diffusion) - """ - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - def enable_freeu(self, s1: float, s2: float, b1: float, b2: float): - r"""Enables the FreeU mechanism as in https://huggingface.co/papers/2309.11497. - - The suffixes after the scaling factors represent the stages where they are being applied. - - Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of the values - that are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL. - - Args: - s1 (`float`): - Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to - mitigate "oversmoothing effect" in the enhanced denoising process. - s2 (`float`): - Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to - mitigate "oversmoothing effect" in the enhanced denoising process. - b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features. - b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features. - """ - if not hasattr(self, "unet"): - raise ValueError("The pipeline must have `unet` for using FreeU.") - self.unet.enable_freeu(s1=s1, s2=s2, b1=b1, b2=b2) - - def disable_freeu(self): - """Disables the FreeU mechanism if enabled.""" - self.unet.disable_freeu() - - def fuse_qkv_projections(self, unet: bool = True, vae: bool = True): - """ - Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) - are fused. For cross-attention modules, key and value projection matrices are fused. - - - - This API is 🧪 experimental. - - - - Args: - unet (`bool`, defaults to `True`): To apply fusion on the UNet. - vae (`bool`, defaults to `True`): To apply fusion on the VAE. - """ - self.fusing_unet = False - self.fusing_vae = False - - if unet: - self.fusing_unet = True - self.unet.fuse_qkv_projections() - self.unet.set_attn_processor(FusedAttnProcessor2_0()) - - if vae: - if not isinstance(self.vae, AutoencoderKL): - raise ValueError("`fuse_qkv_projections()` is only supported for the VAE of type `AutoencoderKL`.") - - self.fusing_vae = True - self.vae.fuse_qkv_projections() - self.vae.set_attn_processor(FusedAttnProcessor2_0()) - - def unfuse_qkv_projections(self, unet: bool = True, vae: bool = True): - """Disable QKV projection fusion if enabled. - - - - This API is 🧪 experimental. - - - - Args: - unet (`bool`, defaults to `True`): To apply fusion on the UNet. - vae (`bool`, defaults to `True`): To apply fusion on the VAE. - - """ - if unet: - if not self.fusing_unet: - logger.warning("The UNet was not initially fused for QKV projections. Doing nothing.") - else: - self.unet.unfuse_qkv_projections() - self.fusing_unet = False - - if vae: - if not self.fusing_vae: - logger.warning("The VAE was not initially fused for QKV projections. Doing nothing.") - else: - self.vae.unfuse_qkv_projections() - self.fusing_vae = False diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/__init__.py deleted file mode 100644 index 7cb870fc8589aeff366429c182736b2ddd6ce215..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/__init__.py +++ /dev/null @@ -1,55 +0,0 @@ -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_pixart_alpha"] = ["PixArtAlphaPipeline"] - _import_structure["pipeline_pixart_sigma"] = ["PixArtSigmaPipeline"] - -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_pixart_alpha import ( - ASPECT_RATIO_256_BIN, - ASPECT_RATIO_512_BIN, - ASPECT_RATIO_1024_BIN, - PixArtAlphaPipeline, - ) - from .pipeline_pixart_sigma import ASPECT_RATIO_2048_BIN, PixArtSigmaPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 260b6d95fd4947ba4bc690ac1e1aec5b75751f00..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/__pycache__/pipeline_pixart_alpha.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/__pycache__/pipeline_pixart_alpha.cpython-310.pyc deleted file mode 100644 index d6c0eaf541678b479c14bc77a8eaca23b48b6fb1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/__pycache__/pipeline_pixart_alpha.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/__pycache__/pipeline_pixart_sigma.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/__pycache__/pipeline_pixart_sigma.cpython-310.pyc deleted file mode 100644 index 407027ba0bb2be47635f0661fe58807efbe9d814..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/__pycache__/pipeline_pixart_sigma.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py deleted file mode 100644 index f7e70c511bc71f1424b548bb1fe5d19e57904038..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py +++ /dev/null @@ -1,974 +0,0 @@ -# Copyright 2025 PixArt-Alpha Authors 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 inspect -import re -import urllib.parse as ul -from typing import Callable, List, Optional, Tuple, Union - -import torch -from transformers import T5EncoderModel, T5Tokenizer - -from ...image_processor import PixArtImageProcessor -from ...models import AutoencoderKL, PixArtTransformer2DModel -from ...schedulers import DPMSolverMultistepScheduler -from ...utils import ( - BACKENDS_MAPPING, - deprecate, - is_bs4_available, - is_ftfy_available, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -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_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import PixArtAlphaPipeline - - >>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-XL-2-512x512" too. - >>> pipe = PixArtAlphaPipeline.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", torch_dtype=torch.float16) - >>> # Enable memory optimizations. - >>> pipe.enable_model_cpu_offload() - - >>> prompt = "A small cactus with a happy face in the Sahara desert." - >>> image = pipe(prompt).images[0] - ``` -""" - -ASPECT_RATIO_1024_BIN = { - "0.25": [512.0, 2048.0], - "0.28": [512.0, 1856.0], - "0.32": [576.0, 1792.0], - "0.33": [576.0, 1728.0], - "0.35": [576.0, 1664.0], - "0.4": [640.0, 1600.0], - "0.42": [640.0, 1536.0], - "0.48": [704.0, 1472.0], - "0.5": [704.0, 1408.0], - "0.52": [704.0, 1344.0], - "0.57": [768.0, 1344.0], - "0.6": [768.0, 1280.0], - "0.68": [832.0, 1216.0], - "0.72": [832.0, 1152.0], - "0.78": [896.0, 1152.0], - "0.82": [896.0, 1088.0], - "0.88": [960.0, 1088.0], - "0.94": [960.0, 1024.0], - "1.0": [1024.0, 1024.0], - "1.07": [1024.0, 960.0], - "1.13": [1088.0, 960.0], - "1.21": [1088.0, 896.0], - "1.29": [1152.0, 896.0], - "1.38": [1152.0, 832.0], - "1.46": [1216.0, 832.0], - "1.67": [1280.0, 768.0], - "1.75": [1344.0, 768.0], - "2.0": [1408.0, 704.0], - "2.09": [1472.0, 704.0], - "2.4": [1536.0, 640.0], - "2.5": [1600.0, 640.0], - "3.0": [1728.0, 576.0], - "4.0": [2048.0, 512.0], -} - -ASPECT_RATIO_512_BIN = { - "0.25": [256.0, 1024.0], - "0.28": [256.0, 928.0], - "0.32": [288.0, 896.0], - "0.33": [288.0, 864.0], - "0.35": [288.0, 832.0], - "0.4": [320.0, 800.0], - "0.42": [320.0, 768.0], - "0.48": [352.0, 736.0], - "0.5": [352.0, 704.0], - "0.52": [352.0, 672.0], - "0.57": [384.0, 672.0], - "0.6": [384.0, 640.0], - "0.68": [416.0, 608.0], - "0.72": [416.0, 576.0], - "0.78": [448.0, 576.0], - "0.82": [448.0, 544.0], - "0.88": [480.0, 544.0], - "0.94": [480.0, 512.0], - "1.0": [512.0, 512.0], - "1.07": [512.0, 480.0], - "1.13": [544.0, 480.0], - "1.21": [544.0, 448.0], - "1.29": [576.0, 448.0], - "1.38": [576.0, 416.0], - "1.46": [608.0, 416.0], - "1.67": [640.0, 384.0], - "1.75": [672.0, 384.0], - "2.0": [704.0, 352.0], - "2.09": [736.0, 352.0], - "2.4": [768.0, 320.0], - "2.5": [800.0, 320.0], - "3.0": [864.0, 288.0], - "4.0": [1024.0, 256.0], -} - -ASPECT_RATIO_256_BIN = { - "0.25": [128.0, 512.0], - "0.28": [128.0, 464.0], - "0.32": [144.0, 448.0], - "0.33": [144.0, 432.0], - "0.35": [144.0, 416.0], - "0.4": [160.0, 400.0], - "0.42": [160.0, 384.0], - "0.48": [176.0, 368.0], - "0.5": [176.0, 352.0], - "0.52": [176.0, 336.0], - "0.57": [192.0, 336.0], - "0.6": [192.0, 320.0], - "0.68": [208.0, 304.0], - "0.72": [208.0, 288.0], - "0.78": [224.0, 288.0], - "0.82": [224.0, 272.0], - "0.88": [240.0, 272.0], - "0.94": [240.0, 256.0], - "1.0": [256.0, 256.0], - "1.07": [256.0, 240.0], - "1.13": [272.0, 240.0], - "1.21": [272.0, 224.0], - "1.29": [288.0, 224.0], - "1.38": [288.0, 208.0], - "1.46": [304.0, 208.0], - "1.67": [320.0, 192.0], - "1.75": [336.0, 192.0], - "2.0": [352.0, 176.0], - "2.09": [368.0, 176.0], - "2.4": [384.0, 160.0], - "2.5": [400.0, 160.0], - "3.0": [432.0, 144.0], - "4.0": [512.0, 128.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, - **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 PixArtAlphaPipeline(DiffusionPipeline): - r""" - Pipeline for text-to-image generation using PixArt-Alpha. - - 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.) - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`T5EncoderModel`]): - Frozen text-encoder. PixArt-Alpha uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the - [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant. - tokenizer (`T5Tokenizer`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - transformer ([`PixArtTransformer2DModel`]): - A text conditioned `PixArtTransformer2DModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - """ - - bad_punct_regex = re.compile( - r"[" - + "#®•©™&@·º½¾¿¡§~" - + r"\)" - + r"\(" - + r"\]" - + r"\[" - + r"\}" - + r"\{" - + r"\|" - + "\\" - + r"\/" - + r"\*" - + r"]{1,}" - ) # noqa - - _optional_components = ["tokenizer", "text_encoder"] - model_cpu_offload_seq = "text_encoder->transformer->vae" - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - vae: AutoencoderKL, - transformer: PixArtTransformer2DModel, - scheduler: DPMSolverMultistepScheduler, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_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 = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor) - - # Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - max_sequence_length: int = 120, - **kwargs, - ): - 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 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 "". - 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 - device: (`torch.device`, *optional*): - 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. For PixArt-Alpha, it's should be the embeddings of the "" - string. - clean_caption (`bool`, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - max_sequence_length (`int`, defaults to 120): Maximum sequence length to use for the prompt. - """ - - if "mask_feature" in kwargs: - deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version." - deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False) - - if device is None: - device = self._execution_device - - # See Section 3.1. of the paper. - max_length = max_sequence_length - - if prompt_embeds is None: - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - add_special_tokens=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[:, max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because T5 can only handle sequences up to" - f" {max_length} tokens: {removed_text}" - ) - - prompt_attention_mask = text_inputs.attention_mask - prompt_attention_mask = prompt_attention_mask.to(device) - - prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask) - prompt_embeds = prompt_embeds[0] - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - elif self.transformer is not None: - dtype = self.transformer.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - - bs_embed, 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(bs_embed * num_images_per_prompt, seq_len, -1) - prompt_attention_mask = prompt_attention_mask.repeat(1, num_images_per_prompt) - prompt_attention_mask = prompt_attention_mask.view(bs_embed * 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 = [negative_prompt] * bs_embed if isinstance(negative_prompt, str) else negative_prompt - uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption) - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - negative_prompt_attention_mask = uncond_input.attention_mask - negative_prompt_attention_mask = negative_prompt_attention_mask.to(device) - - negative_prompt_embeds = self.text_encoder( - uncond_input.input_ids.to(device), attention_mask=negative_prompt_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=dtype, device=device) - - negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) - - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(1, num_images_per_prompt) - negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed * num_images_per_prompt, -1) - else: - negative_prompt_embeds = None - negative_prompt_attention_mask = None - - return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - negative_prompt, - callback_steps, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=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 None) or ( - 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 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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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("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}." - ) - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - negative_prompt: str = "", - num_inference_steps: int = 20, - 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, - 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - clean_caption: bool = True, - use_resolution_binning: bool = True, - max_sequence_length: int = 120, - **kwargs, - ) -> Union[ImagePipelineOutput, 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. 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 - 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*): - 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*): - 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 4.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. - height (`int`, *optional*, defaults to self.unet.config.sample_size): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to self.unet.config.sample_size): - The width in pixels of the generated image. - 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*): - 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 ge 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.Tensor`, *optional*): - Pre-generated negative text embeddings. For PixArt-Alpha this negative prompt should be "". If not - provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): - Pre-generated attention mask for negative text embeddings. - 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.IFPipelineOutput`] 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. - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - use_resolution_binning (`bool` defaults to `True`): - If set to `True`, the requested height and width are first mapped to the closest resolutions using - `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to - the requested resolution. Useful for generating non-square images. - max_sequence_length (`int` defaults to 120): Maximum sequence length to use with the `prompt`. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images - """ - if "mask_feature" in kwargs: - deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version." - deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False) - # 1. Check inputs. Raise error if not correct - height = height or self.transformer.config.sample_size * self.vae_scale_factor - width = width or self.transformer.config.sample_size * self.vae_scale_factor - if use_resolution_binning: - if self.transformer.config.sample_size == 128: - aspect_ratio_bin = ASPECT_RATIO_1024_BIN - elif self.transformer.config.sample_size == 64: - aspect_ratio_bin = ASPECT_RATIO_512_BIN - elif self.transformer.config.sample_size == 32: - aspect_ratio_bin = ASPECT_RATIO_256_BIN - else: - raise ValueError("Invalid sample size") - orig_height, orig_width = height, width - height, width = self.image_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin) - - self.check_inputs( - prompt, - height, - width, - negative_prompt, - callback_steps, - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - ) - - # 2. Default height and width to transformer - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt, - do_classifier_free_guidance, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - ) - if 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) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 5. Prepare latents. - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 6.1 Prepare micro-conditions. - added_cond_kwargs = {"resolution": None, "aspect_ratio": None} - if self.transformer.config.sample_size == 128: - resolution = torch.tensor([height, width]).repeat(batch_size * num_images_per_prompt, 1) - aspect_ratio = torch.tensor([float(height / width)]).repeat(batch_size * num_images_per_prompt, 1) - resolution = resolution.to(dtype=prompt_embeds.dtype, device=device) - aspect_ratio = aspect_ratio.to(dtype=prompt_embeds.dtype, device=device) - - if do_classifier_free_guidance: - resolution = torch.cat([resolution, resolution], dim=0) - aspect_ratio = torch.cat([aspect_ratio, aspect_ratio], dim=0) - - added_cond_kwargs = {"resolution": resolution, "aspect_ratio": aspect_ratio} - - # 7. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - current_timestep = t - if not torch.is_tensor(current_timestep): - # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = latent_model_input.device.type == "mps" - is_npu = latent_model_input.device.type == "npu" - if isinstance(current_timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device) - elif len(current_timestep.shape) == 0: - current_timestep = current_timestep[None].to(latent_model_input.device) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - current_timestep = current_timestep.expand(latent_model_input.shape[0]) - - # predict noise model_output - noise_pred = self.transformer( - latent_model_input, - encoder_hidden_states=prompt_embeds, - encoder_attention_mask=prompt_attention_mask, - timestep=current_timestep, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # learned sigma - if self.transformer.config.out_channels // 2 == latent_channels: - noise_pred = noise_pred.chunk(2, dim=1)[0] - else: - noise_pred = noise_pred - - # compute previous image: x_t -> x_t-1 - if num_inference_steps == 1: - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[1] - else: - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - if use_resolution_binning: - image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height) - 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 ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py deleted file mode 100644 index c3d235d91bcc69fd33758ed94b6997b0a86c067a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py +++ /dev/null @@ -1,886 +0,0 @@ -# Copyright 2025 PixArt-Sigma Authors 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 inspect -import re -import urllib.parse as ul -from typing import Callable, List, Optional, Tuple, Union - -import torch -from transformers import T5EncoderModel, T5Tokenizer - -from ...image_processor import PixArtImageProcessor -from ...models import AutoencoderKL, PixArtTransformer2DModel -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - BACKENDS_MAPPING, - deprecate, - is_bs4_available, - is_ftfy_available, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput -from .pipeline_pixart_alpha import ( - ASPECT_RATIO_256_BIN, - ASPECT_RATIO_512_BIN, - ASPECT_RATIO_1024_BIN, -) - - -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_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - - -ASPECT_RATIO_2048_BIN = { - "0.25": [1024.0, 4096.0], - "0.26": [1024.0, 3968.0], - "0.27": [1024.0, 3840.0], - "0.28": [1024.0, 3712.0], - "0.32": [1152.0, 3584.0], - "0.33": [1152.0, 3456.0], - "0.35": [1152.0, 3328.0], - "0.4": [1280.0, 3200.0], - "0.42": [1280.0, 3072.0], - "0.48": [1408.0, 2944.0], - "0.5": [1408.0, 2816.0], - "0.52": [1408.0, 2688.0], - "0.57": [1536.0, 2688.0], - "0.6": [1536.0, 2560.0], - "0.68": [1664.0, 2432.0], - "0.72": [1664.0, 2304.0], - "0.78": [1792.0, 2304.0], - "0.82": [1792.0, 2176.0], - "0.88": [1920.0, 2176.0], - "0.94": [1920.0, 2048.0], - "1.0": [2048.0, 2048.0], - "1.07": [2048.0, 1920.0], - "1.13": [2176.0, 1920.0], - "1.21": [2176.0, 1792.0], - "1.29": [2304.0, 1792.0], - "1.38": [2304.0, 1664.0], - "1.46": [2432.0, 1664.0], - "1.67": [2560.0, 1536.0], - "1.75": [2688.0, 1536.0], - "2.0": [2816.0, 1408.0], - "2.09": [2944.0, 1408.0], - "2.4": [3072.0, 1280.0], - "2.5": [3200.0, 1280.0], - "2.89": [3328.0, 1152.0], - "3.0": [3456.0, 1152.0], - "3.11": [3584.0, 1152.0], - "3.62": [3712.0, 1024.0], - "3.75": [3840.0, 1024.0], - "3.88": [3968.0, 1024.0], - "4.0": [4096.0, 1024.0], -} - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import PixArtSigmaPipeline - - >>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" too. - >>> pipe = PixArtSigmaPipeline.from_pretrained( - ... "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16 - ... ) - >>> # Enable memory optimizations. - >>> # pipe.enable_model_cpu_offload() - - >>> prompt = "A small cactus with a happy face in the Sahara desert." - >>> image = pipe(prompt).images[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, - **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 PixArtSigmaPipeline(DiffusionPipeline): - r""" - Pipeline for text-to-image generation using PixArt-Sigma. - """ - - bad_punct_regex = re.compile( - r"[" - + "#®•©™&@·º½¾¿¡§~" - + r"\)" - + r"\(" - + r"\]" - + r"\[" - + r"\}" - + r"\{" - + r"\|" - + "\\" - + r"\/" - + r"\*" - + r"]{1,}" - ) # noqa - - _optional_components = ["tokenizer", "text_encoder"] - model_cpu_offload_seq = "text_encoder->transformer->vae" - - def __init__( - self, - tokenizer: T5Tokenizer, - text_encoder: T5EncoderModel, - vae: AutoencoderKL, - transformer: PixArtTransformer2DModel, - scheduler: KarrasDiffusionSchedulers, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_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 = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor) - - # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt with 120->300 - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - max_sequence_length: int = 300, - **kwargs, - ): - 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 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 "". - 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 - device: (`torch.device`, *optional*): - 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. For PixArt-Alpha, it's should be the embeddings of the "" - string. - clean_caption (`bool`, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. - """ - - if "mask_feature" in kwargs: - deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version." - deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False) - - if device is None: - device = self._execution_device - - # See Section 3.1. of the paper. - max_length = max_sequence_length - - if prompt_embeds is None: - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length, - truncation=True, - add_special_tokens=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[:, max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because T5 can only handle sequences up to" - f" {max_length} tokens: {removed_text}" - ) - - prompt_attention_mask = text_inputs.attention_mask - prompt_attention_mask = prompt_attention_mask.to(device) - - prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask) - prompt_embeds = prompt_embeds[0] - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - elif self.transformer is not None: - dtype = self.transformer.dtype - else: - dtype = None - - prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) - - bs_embed, 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(bs_embed * num_images_per_prompt, seq_len, -1) - prompt_attention_mask = prompt_attention_mask.repeat(1, num_images_per_prompt) - prompt_attention_mask = prompt_attention_mask.view(bs_embed * 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 = [negative_prompt] * bs_embed if isinstance(negative_prompt, str) else negative_prompt - uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption) - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - negative_prompt_attention_mask = uncond_input.attention_mask - negative_prompt_attention_mask = negative_prompt_attention_mask.to(device) - - negative_prompt_embeds = self.text_encoder( - uncond_input.input_ids.to(device), attention_mask=negative_prompt_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=dtype, device=device) - - negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) - - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(1, num_images_per_prompt) - negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed * num_images_per_prompt, -1) - else: - negative_prompt_embeds = None - negative_prompt_attention_mask = None - - return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.check_inputs - def check_inputs( - self, - prompt, - height, - width, - negative_prompt, - callback_steps, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=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 None) or ( - 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 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 not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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("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}." - ) - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - negative_prompt: str = "", - num_inference_steps: int = 20, - 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, - 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - clean_caption: bool = True, - use_resolution_binning: bool = True, - max_sequence_length: int = 300, - **kwargs, - ) -> Union[ImagePipelineOutput, 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. 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 - 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*): - 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*): - 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 4.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. - height (`int`, *optional*, defaults to self.unet.config.sample_size): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to self.unet.config.sample_size): - The width in pixels of the generated image. - 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*): - 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 ge 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.Tensor`, *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.Tensor`, *optional*): - Pre-generated attention mask for negative text embeddings. - 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.IFPipelineOutput`] 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. - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - use_resolution_binning (`bool` defaults to `True`): - If set to `True`, the requested height and width are first mapped to the closest resolutions using - `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to - the requested resolution. Useful for generating non-square images. - max_sequence_length (`int` defaults to 300): Maximum sequence length to use with the `prompt`. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images - """ - # 1. Check inputs. Raise error if not correct - height = height or self.transformer.config.sample_size * self.vae_scale_factor - width = width or self.transformer.config.sample_size * self.vae_scale_factor - if use_resolution_binning: - if self.transformer.config.sample_size == 256: - aspect_ratio_bin = ASPECT_RATIO_2048_BIN - elif self.transformer.config.sample_size == 128: - aspect_ratio_bin = ASPECT_RATIO_1024_BIN - elif self.transformer.config.sample_size == 64: - aspect_ratio_bin = ASPECT_RATIO_512_BIN - elif self.transformer.config.sample_size == 32: - aspect_ratio_bin = ASPECT_RATIO_256_BIN - else: - raise ValueError("Invalid sample size") - orig_height, orig_width = height, width - height, width = self.image_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin) - - self.check_inputs( - prompt, - height, - width, - negative_prompt, - callback_steps, - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - ) - - # 2. Default height and width to transformer - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt, - do_classifier_free_guidance, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - ) - if 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) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 5. Prepare latents. - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 6.1 Prepare micro-conditions. - added_cond_kwargs = {"resolution": None, "aspect_ratio": None} - - # 7. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - current_timestep = t - if not torch.is_tensor(current_timestep): - # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can - # This would be a good case for the `match` statement (Python 3.10+) - is_mps = latent_model_input.device.type == "mps" - is_npu = latent_model_input.device.type == "npu" - if isinstance(current_timestep, float): - dtype = torch.float32 if (is_mps or is_npu) else torch.float64 - else: - dtype = torch.int32 if (is_mps or is_npu) else torch.int64 - current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device) - elif len(current_timestep.shape) == 0: - current_timestep = current_timestep[None].to(latent_model_input.device) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - current_timestep = current_timestep.expand(latent_model_input.shape[0]) - - # predict noise model_output - noise_pred = self.transformer( - latent_model_input, - encoder_hidden_states=prompt_embeds, - encoder_attention_mask=prompt_attention_mask, - timestep=current_timestep, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # learned sigma - if self.transformer.config.out_channels // 2 == latent_channels: - noise_pred = noise_pred.chunk(2, dim=1)[0] - else: - noise_pred = noise_pred - - # compute previous image: x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if not output_type == "latent": - image = self.vae.decode(latents.to(self.vae.dtype) / self.vae.config.scaling_factor, return_dict=False)[0] - if use_resolution_binning: - image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height) - 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 ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__init__.py deleted file mode 100644 index 91684f35f153516de83cd65ebed28164970b7548..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__init__.py +++ /dev/null @@ -1,53 +0,0 @@ -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_sana"] = ["SanaPipeline"] - _import_structure["pipeline_sana_controlnet"] = ["SanaControlNetPipeline"] - _import_structure["pipeline_sana_sprint"] = ["SanaSprintPipeline"] - _import_structure["pipeline_sana_sprint_img2img"] = ["SanaSprintImg2ImgPipeline"] - -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_sana import SanaPipeline - from .pipeline_sana_controlnet import SanaControlNetPipeline - from .pipeline_sana_sprint import SanaSprintPipeline - from .pipeline_sana_sprint_img2img import SanaSprintImg2ImgPipeline -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/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 1e049700cb8ae41bda70e99936788aab4ee98f4b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index c5f32d94f526b412eb53f80bdb787e6a3bfcfaf7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_sana.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_sana.cpython-310.pyc deleted file mode 100644 index 5db51869cdfcaa78a81a9705f0f8f0a5366189cf..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_sana.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_sana_controlnet.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_sana_controlnet.cpython-310.pyc deleted file mode 100644 index e91a06bde861cddf8ff5e1b91a8d8c6f0ddf3ded..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_sana_controlnet.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_sana_sprint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_sana_sprint.cpython-310.pyc deleted file mode 100644 index dcbcc8188215d835409d63ba3f394d6778ae68a9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_sana_sprint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_sana_sprint_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_sana_sprint_img2img.cpython-310.pyc deleted file mode 100644 index f326ad83a01c124eba9624369f75154f9cebbb8f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/__pycache__/pipeline_sana_sprint_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_output.py deleted file mode 100644 index f8ac129516446607b2bed771b3ebef72ef2c542d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_output.py +++ /dev/null @@ -1,21 +0,0 @@ -from dataclasses import dataclass -from typing import List, Union - -import numpy as np -import PIL.Image - -from ...utils import BaseOutput - - -@dataclass -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, - num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana.py deleted file mode 100644 index 103f57a236401a37881c6dd1bae6b94106e78b9e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana.py +++ /dev/null @@ -1,1011 +0,0 @@ -# Copyright 2025 PixArt-Sigma Authors 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 inspect -import re -import urllib.parse as ul -import warnings -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PixArtImageProcessor -from ...loaders import SanaLoraLoaderMixin -from ...models import AutoencoderDC, SanaTransformer2DModel -from ...schedulers import DPMSolverMultistepScheduler -from ...utils import ( - BACKENDS_MAPPING, - USE_PEFT_BACKEND, - is_bs4_available, - is_ftfy_available, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import get_device, is_torch_version, randn_tensor -from ..pipeline_utils import DiffusionPipeline -from ..pixart_alpha.pipeline_pixart_alpha import ( - ASPECT_RATIO_512_BIN, - ASPECT_RATIO_1024_BIN, -) -from ..pixart_alpha.pipeline_pixart_sigma import ASPECT_RATIO_2048_BIN -from .pipeline_output import SanaPipelineOutput - - -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_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - - -ASPECT_RATIO_4096_BIN = { - "0.25": [2048.0, 8192.0], - "0.26": [2048.0, 7936.0], - "0.27": [2048.0, 7680.0], - "0.28": [2048.0, 7424.0], - "0.32": [2304.0, 7168.0], - "0.33": [2304.0, 6912.0], - "0.35": [2304.0, 6656.0], - "0.4": [2560.0, 6400.0], - "0.42": [2560.0, 6144.0], - "0.48": [2816.0, 5888.0], - "0.5": [2816.0, 5632.0], - "0.52": [2816.0, 5376.0], - "0.57": [3072.0, 5376.0], - "0.6": [3072.0, 5120.0], - "0.68": [3328.0, 4864.0], - "0.72": [3328.0, 4608.0], - "0.78": [3584.0, 4608.0], - "0.82": [3584.0, 4352.0], - "0.88": [3840.0, 4352.0], - "0.94": [3840.0, 4096.0], - "1.0": [4096.0, 4096.0], - "1.07": [4096.0, 3840.0], - "1.13": [4352.0, 3840.0], - "1.21": [4352.0, 3584.0], - "1.29": [4608.0, 3584.0], - "1.38": [4608.0, 3328.0], - "1.46": [4864.0, 3328.0], - "1.67": [5120.0, 3072.0], - "1.75": [5376.0, 3072.0], - "2.0": [5632.0, 2816.0], - "2.09": [5888.0, 2816.0], - "2.4": [6144.0, 2560.0], - "2.5": [6400.0, 2560.0], - "2.89": [6656.0, 2304.0], - "3.0": [6912.0, 2304.0], - "3.11": [7168.0, 2304.0], - "3.62": [7424.0, 2048.0], - "3.75": [7680.0, 2048.0], - "3.88": [7936.0, 2048.0], - "4.0": [8192.0, 2048.0], -} - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import SanaPipeline - - >>> pipe = SanaPipeline.from_pretrained( - ... "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers", torch_dtype=torch.float32 - ... ) - >>> pipe.to("cuda") - >>> pipe.text_encoder.to(torch.bfloat16) - >>> pipe.transformer = pipe.transformer.to(torch.bfloat16) - - >>> image = pipe(prompt='a cyberpunk cat with a neon sign that says "Sana"')[0] - >>> image[0].save("output.png") - ``` -""" - - -# 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 SanaPipeline(DiffusionPipeline, SanaLoraLoaderMixin): - r""" - Pipeline for text-to-image generation using [Sana](https://huggingface.co/papers/2410.10629). - """ - - # fmt: off - bad_punct_regex = re.compile(r"[" + "#®•©™&@·º½¾¿¡§~" + r"\)" + r"\(" + r"\]" + r"\[" + r"\}" + r"\{" + r"\|" + "\\" + r"\/" + r"\*" + r"]{1,}") - # fmt: on - - model_cpu_offload_seq = "text_encoder->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], - text_encoder: Gemma2PreTrainedModel, - vae: AutoencoderDC, - transformer: SanaTransformer2DModel, - scheduler: DPMSolverMultistepScheduler, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler - ) - - self.vae_scale_factor = ( - 2 ** (len(self.vae.config.encoder_block_out_channels) - 1) - if hasattr(self, "vae") and self.vae is not None - else 32 - ) - self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor) - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - def _get_gemma_prompt_embeds( - self, - prompt: Union[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, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`, *optional*): - prompt to be encoded - device: (`torch.device`, *optional*): - torch device to place the resulting embeddings on - clean_caption (`bool`, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. - complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): - If `complex_human_instruction` is not empty, the function will use the complex Human instruction for - the prompt. - """ - prompt = [prompt] if isinstance(prompt, str) else prompt - - if getattr(self, "tokenizer", None) is not None: - self.tokenizer.padding_side = "right" - - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - - # prepare complex human instruction - if not complex_human_instruction: - max_length_all = max_sequence_length - else: - chi_prompt = "\n".join(complex_human_instruction) - prompt = [chi_prompt + p for p in prompt] - num_chi_prompt_tokens = len(self.tokenizer.encode(chi_prompt)) - max_length_all = num_chi_prompt_tokens + max_sequence_length - 2 - - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length_all, - 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.to(device) - - prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask) - prompt_embeds = prompt_embeds[0].to(dtype=dtype, device=device) - - return prompt_embeds, prompt_attention_mask - - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, - lora_scale: Optional[float] = 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 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 "". - 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 - device: (`torch.device`, *optional*): - 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. For Sana, it's should be the embeddings of the "" string. - clean_caption (`bool`, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. - complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): - If `complex_human_instruction` is not empty, the function will use the complex Human instruction for - the prompt. - """ - - if device is None: - device = self._execution_device - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - else: - dtype = None - - # 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, SanaLoraLoaderMixin): - 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) - - 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 getattr(self, "tokenizer", None) is not None: - self.tokenizer.padding_side = "right" - - # See Section 3.1. of the paper. - max_length = max_sequence_length - select_index = [0] + list(range(-max_length + 1, 0)) - - if prompt_embeds is None: - prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( - prompt=prompt, - device=device, - dtype=dtype, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - complex_human_instruction=complex_human_instruction, - ) - - prompt_embeds = prompt_embeds[:, select_index] - prompt_attention_mask = prompt_attention_mask[:, select_index] - - bs_embed, 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(bs_embed * num_images_per_prompt, seq_len, -1) - prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1) - prompt_attention_mask = prompt_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: - negative_prompt = [negative_prompt] * batch_size if isinstance(negative_prompt, str) else negative_prompt - negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( - prompt=negative_prompt, - device=device, - dtype=dtype, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - complex_human_instruction=False, - ) - - 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=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) - - negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1) - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1) - else: - negative_prompt_embeds = None - negative_prompt_attention_mask = None - - if self.text_encoder is not None: - if isinstance(self, SanaLoraLoaderMixin) 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, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - callback_on_step_end_tensor_inputs=None, - negative_prompt=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 is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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("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}." - ) - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - 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=device, dtype=dtype) - - 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 - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def attention_kwargs(self): - return self._attention_kwargs - - @property - def do_classifier_free_guidance(self): - return self._guidance_scale > 1.0 - - @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, - negative_prompt: str = "", - num_inference_steps: int = 20, - timesteps: List[int] = None, - sigmas: List[float] = None, - guidance_scale: float = 4.5, - num_images_per_prompt: Optional[int] = 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", - 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"], - max_sequence_length: int = 300, - 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.", - "Here are examples of how to transform or refine prompts:", - "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.", - "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.", - "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", - "User Prompt: ", - ], - ) -> Union[SanaPipelineOutput, 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. 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 - 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*): - 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*): - 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 4.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. - height (`int`, *optional*, defaults to self.unet.config.sample_size): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to self.unet.config.sample_size): - The width in pixels of the generated image. - 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*): - 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 ge 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.Tensor`, *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.Tensor`, *optional*): - Pre-generated attention mask for negative text embeddings. - 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.IFPipelineOutput`] instead of a plain tuple. - attention_kwargs: - 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). - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - use_resolution_binning (`bool` defaults to `True`): - If set to `True`, the requested height and width are first mapped to the closest resolutions using - `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to - the requested resolution. Useful for generating non-square images. - 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 `300`): - Maximum sequence length to use with the `prompt`. - 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. - - Examples: - - Returns: - [`~pipelines.sana.pipeline_output.SanaPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.sana.pipeline_output.SanaPipelineOutput`] 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 - if use_resolution_binning: - if self.transformer.config.sample_size == 128: - aspect_ratio_bin = ASPECT_RATIO_4096_BIN - elif self.transformer.config.sample_size == 64: - aspect_ratio_bin = ASPECT_RATIO_2048_BIN - elif self.transformer.config.sample_size == 32: - aspect_ratio_bin = ASPECT_RATIO_1024_BIN - elif self.transformer.config.sample_size == 16: - aspect_ratio_bin = ASPECT_RATIO_512_BIN - else: - raise ValueError("Invalid sample size") - orig_height, orig_width = height, width - height, width = self.image_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin) - - self.check_inputs( - prompt, - height, - width, - callback_on_step_end_tensor_inputs, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - ) - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._interrupt = False - - # 2. Default height and width to transformer - 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.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None - - # 3. Encode input prompt - ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt, - self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - complex_human_instruction=complex_human_instruction, - lora_scale=lora_scale, - ) - 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) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 5. Prepare latents. - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - torch.float32, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - transformer_dtype = self.transformer.dtype - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - timestep = timestep * self.transformer.config.timestep_scale - - # predict noise model_output - noise_pred = self.transformer( - latent_model_input.to(dtype=transformer_dtype), - encoder_hidden_states=prompt_embeds.to(dtype=transformer_dtype), - encoder_attention_mask=prompt_attention_mask, - timestep=timestep, - return_dict=False, - attention_kwargs=self.attention_kwargs, - )[0] - noise_pred = noise_pred.float() - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # learned sigma - if self.transformer.config.out_channels // 2 == latent_channels: - noise_pred = noise_pred.chunk(2, dim=1)[0] - - # compute previous image: x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - - # 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) - torch_accelerator_module = getattr(torch, get_device(), torch.cuda) - oom_error = ( - torch.OutOfMemoryError - if is_torch_version(">=", "2.5.0") - else torch_accelerator_module.OutOfMemoryError - ) - try: - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - except oom_error as e: - warnings.warn( - f"{e}. \n" - f"Try to use VAE tiling for large images. For example: \n" - f"pipe.vae.enable_tiling(tile_sample_min_width=512, tile_sample_min_height=512)" - ) - if use_resolution_binning: - image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height) - - 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 SanaPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana_controlnet.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana_controlnet.py deleted file mode 100644 index cdc602b964cf826ec77e35bfd6c75d50ca358a55..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana_controlnet.py +++ /dev/null @@ -1,1106 +0,0 @@ -# Copyright 2025 PixArt-Sigma Authors 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 inspect -import re -import urllib.parse as ul -import warnings -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, PixArtImageProcessor -from ...loaders import SanaLoraLoaderMixin -from ...models import AutoencoderDC, SanaControlNetModel, SanaTransformer2DModel -from ...schedulers import DPMSolverMultistepScheduler -from ...utils import ( - BACKENDS_MAPPING, - USE_PEFT_BACKEND, - is_bs4_available, - is_ftfy_available, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import get_device, is_torch_version, randn_tensor -from ..pipeline_utils import DiffusionPipeline -from ..pixart_alpha.pipeline_pixart_alpha import ( - ASPECT_RATIO_512_BIN, - ASPECT_RATIO_1024_BIN, -) -from ..pixart_alpha.pipeline_pixart_sigma import ASPECT_RATIO_2048_BIN -from .pipeline_output import SanaPipelineOutput - - -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_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - - -ASPECT_RATIO_4096_BIN = { - "0.25": [2048.0, 8192.0], - "0.26": [2048.0, 7936.0], - "0.27": [2048.0, 7680.0], - "0.28": [2048.0, 7424.0], - "0.32": [2304.0, 7168.0], - "0.33": [2304.0, 6912.0], - "0.35": [2304.0, 6656.0], - "0.4": [2560.0, 6400.0], - "0.42": [2560.0, 6144.0], - "0.48": [2816.0, 5888.0], - "0.5": [2816.0, 5632.0], - "0.52": [2816.0, 5376.0], - "0.57": [3072.0, 5376.0], - "0.6": [3072.0, 5120.0], - "0.68": [3328.0, 4864.0], - "0.72": [3328.0, 4608.0], - "0.78": [3584.0, 4608.0], - "0.82": [3584.0, 4352.0], - "0.88": [3840.0, 4352.0], - "0.94": [3840.0, 4096.0], - "1.0": [4096.0, 4096.0], - "1.07": [4096.0, 3840.0], - "1.13": [4352.0, 3840.0], - "1.21": [4352.0, 3584.0], - "1.29": [4608.0, 3584.0], - "1.38": [4608.0, 3328.0], - "1.46": [4864.0, 3328.0], - "1.67": [5120.0, 3072.0], - "1.75": [5376.0, 3072.0], - "2.0": [5632.0, 2816.0], - "2.09": [5888.0, 2816.0], - "2.4": [6144.0, 2560.0], - "2.5": [6400.0, 2560.0], - "2.89": [6656.0, 2304.0], - "3.0": [6912.0, 2304.0], - "3.11": [7168.0, 2304.0], - "3.62": [7424.0, 2048.0], - "3.75": [7680.0, 2048.0], - "3.88": [7936.0, 2048.0], - "4.0": [8192.0, 2048.0], -} - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import SanaControlNetPipeline - >>> from diffusers.utils import load_image - - >>> pipe = SanaControlNetPipeline.from_pretrained( - ... "ishan24/Sana_600M_1024px_ControlNetPlus_diffusers", - ... variant="fp16", - ... torch_dtype={"default": torch.bfloat16, "controlnet": torch.float16, "transformer": torch.float16}, - ... device_map="balanced", - ... ) - >>> cond_image = load_image( - ... "https://huggingface.co/ishan24/Sana_600M_1024px_ControlNet_diffusers/resolve/main/hed_example.png" - ... ) - >>> prompt = 'a cat with a neon sign that says "Sana"' - >>> image = pipe( - ... prompt, - ... control_image=cond_image, - ... ).images[0] - >>> image.save("output.png") - ``` -""" - - -# 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 SanaControlNetPipeline(DiffusionPipeline, SanaLoraLoaderMixin): - r""" - Pipeline for text-to-image generation using [Sana](https://huggingface.co/papers/2410.10629). - """ - - # fmt: off - bad_punct_regex = re.compile(r"[" + "#®•©™&@·º½¾¿¡§~" + r"\)" + r"\(" + r"\]" + r"\[" + r"\}" + r"\{" + r"\|" + "\\" + r"\/" + r"\*" + r"]{1,}") - # fmt: on - - model_cpu_offload_seq = "text_encoder->controlnet->transformer->vae" - _callback_tensor_inputs = ["latents", "control_image", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], - text_encoder: Gemma2PreTrainedModel, - vae: AutoencoderDC, - transformer: SanaTransformer2DModel, - controlnet: SanaControlNetModel, - scheduler: DPMSolverMultistepScheduler, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - vae=vae, - transformer=transformer, - controlnet=controlnet, - scheduler=scheduler, - ) - - self.vae_scale_factor = ( - 2 ** (len(self.vae.config.encoder_block_out_channels) - 1) - if hasattr(self, "vae") and self.vae is not None - else 32 - ) - self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor) - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds - def _get_gemma_prompt_embeds( - self, - prompt: Union[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, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`, *optional*): - prompt to be encoded - device: (`torch.device`, *optional*): - torch device to place the resulting embeddings on - clean_caption (`bool`, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. - complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): - If `complex_human_instruction` is not empty, the function will use the complex Human instruction for - the prompt. - """ - prompt = [prompt] if isinstance(prompt, str) else prompt - - if getattr(self, "tokenizer", None) is not None: - self.tokenizer.padding_side = "right" - - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - - # prepare complex human instruction - if not complex_human_instruction: - max_length_all = max_sequence_length - else: - chi_prompt = "\n".join(complex_human_instruction) - prompt = [chi_prompt + p for p in prompt] - num_chi_prompt_tokens = len(self.tokenizer.encode(chi_prompt)) - max_length_all = num_chi_prompt_tokens + max_sequence_length - 2 - - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length_all, - 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.to(device) - - prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask) - prompt_embeds = prompt_embeds[0].to(dtype=dtype, device=device) - - return prompt_embeds, prompt_attention_mask - - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, - lora_scale: Optional[float] = 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 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 "". - 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 - device: (`torch.device`, *optional*): - 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. For Sana, it's should be the embeddings of the "" string. - clean_caption (`bool`, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. - complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): - If `complex_human_instruction` is not empty, the function will use the complex Human instruction for - the prompt. - """ - - if device is None: - device = self._execution_device - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - else: - dtype = None - - # 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, SanaLoraLoaderMixin): - 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) - - 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 getattr(self, "tokenizer", None) is not None: - self.tokenizer.padding_side = "right" - - # See Section 3.1. of the paper. - max_length = max_sequence_length - select_index = [0] + list(range(-max_length + 1, 0)) - - if prompt_embeds is None: - prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( - prompt=prompt, - device=device, - dtype=dtype, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - complex_human_instruction=complex_human_instruction, - ) - - prompt_embeds = prompt_embeds[:, select_index] - prompt_attention_mask = prompt_attention_mask[:, select_index] - - bs_embed, 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(bs_embed * num_images_per_prompt, seq_len, -1) - prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1) - prompt_attention_mask = prompt_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: - negative_prompt = [negative_prompt] * batch_size if isinstance(negative_prompt, str) else negative_prompt - negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( - prompt=negative_prompt, - device=device, - dtype=dtype, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - complex_human_instruction=False, - ) - - 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=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) - - negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1) - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1) - else: - negative_prompt_embeds = None - negative_prompt_attention_mask = None - - if self.text_encoder is not None: - if isinstance(self, SanaLoraLoaderMixin) 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, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - callback_on_step_end_tensor_inputs=None, - negative_prompt=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 is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - 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("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}." - ) - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - def prepare_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - do_classifier_free_guidance=False, - guess_mode=False, - ): - if isinstance(image, torch.Tensor): - pass - else: - image = self.image_processor.preprocess(image, height=height, width=width) - - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - if do_classifier_free_guidance and not guess_mode: - image = torch.cat([image] * 2) - - return image - - 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=device, dtype=dtype) - - 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 - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def attention_kwargs(self): - return self._attention_kwargs - - @property - def do_classifier_free_guidance(self): - return self._guidance_scale > 1.0 - - @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, - negative_prompt: str = "", - num_inference_steps: int = 20, - 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, - 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", - 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"], - max_sequence_length: int = 300, - 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.", - "Here are examples of how to transform or refine prompts:", - "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.", - "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.", - "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", - "User Prompt: ", - ], - ) -> Union[SanaPipelineOutput, 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. 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 - 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*): - 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*): - 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 4.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. - 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): - 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. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - height (`int`, *optional*, defaults to self.unet.config.sample_size): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to self.unet.config.sample_size): - The width in pixels of the generated image. - 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*): - 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 ge 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.Tensor`, *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.Tensor`, *optional*): - Pre-generated attention mask for negative text embeddings. - 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.IFPipelineOutput`] instead of a plain tuple. - attention_kwargs: - 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). - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - use_resolution_binning (`bool` defaults to `True`): - If set to `True`, the requested height and width are first mapped to the closest resolutions using - `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to - the requested resolution. Useful for generating non-square images. - 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 `300`): - Maximum sequence length to use with the `prompt`. - 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. - - Examples: - - Returns: - [`~pipelines.sana.pipeline_output.SanaPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.sana.pipeline_output.SanaPipelineOutput`] 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 - if use_resolution_binning: - if self.transformer.config.sample_size == 128: - aspect_ratio_bin = ASPECT_RATIO_4096_BIN - elif self.transformer.config.sample_size == 64: - aspect_ratio_bin = ASPECT_RATIO_2048_BIN - elif self.transformer.config.sample_size == 32: - aspect_ratio_bin = ASPECT_RATIO_1024_BIN - elif self.transformer.config.sample_size == 16: - aspect_ratio_bin = ASPECT_RATIO_512_BIN - else: - raise ValueError("Invalid sample size") - orig_height, orig_width = height, width - height, width = self.image_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin) - - self.check_inputs( - prompt, - height, - width, - callback_on_step_end_tensor_inputs, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - prompt_attention_mask, - negative_prompt_attention_mask, - ) - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._interrupt = False - - # 2. Default height and width to transformer - 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.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None - - # 3. Encode input prompt - ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt, - self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - complex_human_instruction=complex_human_instruction, - lora_scale=lora_scale, - ) - 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) - - # 4. Prepare control image - if isinstance(self.controlnet, SanaControlNetModel): - control_image = self.prepare_image( - image=control_image, - width=width, - height=height, - batch_size=batch_size * num_images_per_prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - dtype=self.vae.dtype, - do_classifier_free_guidance=self.do_classifier_free_guidance, - guess_mode=False, - ) - height, width = control_image.shape[-2:] - - control_image = self.vae.encode(control_image).latent - control_image = control_image * self.vae.config.scaling_factor - else: - raise ValueError("`controlnet` must be of type `SanaControlNetModel`.") - - # 5. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 6. Prepare latents. - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - torch.float32, - device, - generator, - latents, - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - controlnet_dtype = self.controlnet.dtype - transformer_dtype = self.transformer.dtype - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - # controlnet(s) inference - controlnet_block_samples = self.controlnet( - latent_model_input.to(dtype=controlnet_dtype), - encoder_hidden_states=prompt_embeds.to(dtype=controlnet_dtype), - encoder_attention_mask=prompt_attention_mask, - timestep=timestep, - return_dict=False, - attention_kwargs=self.attention_kwargs, - controlnet_cond=control_image, - conditioning_scale=controlnet_conditioning_scale, - )[0] - - # predict noise model_output - noise_pred = self.transformer( - latent_model_input.to(dtype=transformer_dtype), - encoder_hidden_states=prompt_embeds.to(dtype=transformer_dtype), - encoder_attention_mask=prompt_attention_mask, - timestep=timestep, - return_dict=False, - attention_kwargs=self.attention_kwargs, - controlnet_block_samples=tuple(t.to(dtype=transformer_dtype) for t in controlnet_block_samples), - )[0] - noise_pred = noise_pred.float() - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # learned sigma - if self.transformer.config.out_channels // 2 == latent_channels: - noise_pred = noise_pred.chunk(2, dim=1)[0] - - # compute previous image: x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - - # 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) - torch_accelerator_module = getattr(torch, get_device(), torch.cuda) - oom_error = ( - torch.OutOfMemoryError - if is_torch_version(">=", "2.5.0") - else torch_accelerator_module.OutOfMemoryError - ) - try: - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - except oom_error as e: - warnings.warn( - f"{e}. \n" - f"Try to use VAE tiling for large images. For example: \n" - f"pipe.vae.enable_tiling(tile_sample_min_width=512, tile_sample_min_height=512)" - ) - if use_resolution_binning: - image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height) - - 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 SanaPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana_sprint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana_sprint.py deleted file mode 100644 index fcf854a54cad178a084e70684e9eae9a1c3f0601..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana_sprint.py +++ /dev/null @@ -1,893 +0,0 @@ -# Copyright 2025 PixArt-Sigma Authors 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 inspect -import re -import urllib.parse as ul -import warnings -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PixArtImageProcessor -from ...loaders import SanaLoraLoaderMixin -from ...models import AutoencoderDC, SanaTransformer2DModel -from ...schedulers import DPMSolverMultistepScheduler -from ...utils import ( - BACKENDS_MAPPING, - USE_PEFT_BACKEND, - is_bs4_available, - is_ftfy_available, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import get_device, is_torch_version, randn_tensor -from ..pipeline_utils import DiffusionPipeline -from ..pixart_alpha.pipeline_pixart_alpha import ASPECT_RATIO_1024_BIN -from .pipeline_output import SanaPipelineOutput - - -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_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import SanaSprintPipeline - - >>> pipe = SanaSprintPipeline.from_pretrained( - ... "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers", torch_dtype=torch.bfloat16 - ... ) - >>> pipe.to("cuda") - - >>> image = pipe(prompt="a tiny astronaut hatching from an egg on the moon")[0] - >>> image[0].save("output.png") - ``` -""" - - -# 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 SanaSprintPipeline(DiffusionPipeline, SanaLoraLoaderMixin): - r""" - Pipeline for text-to-image generation using [SANA-Sprint](https://huggingface.co/papers/2503.09641). - """ - - # fmt: off - bad_punct_regex = re.compile(r"[" + "#®•©™&@·º½¾¿¡§~" + r"\)" + r"\(" + r"\]" + r"\[" + r"\}" + r"\{" + r"\|" + "\\" + r"\/" + r"\*" + r"]{1,}") - # fmt: on - - model_cpu_offload_seq = "text_encoder->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__( - self, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], - text_encoder: Gemma2PreTrainedModel, - vae: AutoencoderDC, - transformer: SanaTransformer2DModel, - scheduler: DPMSolverMultistepScheduler, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler - ) - - self.vae_scale_factor = ( - 2 ** (len(self.vae.config.encoder_block_out_channels) - 1) - if hasattr(self, "vae") and self.vae is not None - else 32 - ) - self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor) - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds - def _get_gemma_prompt_embeds( - self, - prompt: Union[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, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`, *optional*): - prompt to be encoded - device: (`torch.device`, *optional*): - torch device to place the resulting embeddings on - clean_caption (`bool`, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. - complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): - If `complex_human_instruction` is not empty, the function will use the complex Human instruction for - the prompt. - """ - prompt = [prompt] if isinstance(prompt, str) else prompt - - if getattr(self, "tokenizer", None) is not None: - self.tokenizer.padding_side = "right" - - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - - # prepare complex human instruction - if not complex_human_instruction: - max_length_all = max_sequence_length - else: - chi_prompt = "\n".join(complex_human_instruction) - prompt = [chi_prompt + p for p in prompt] - num_chi_prompt_tokens = len(self.tokenizer.encode(chi_prompt)) - max_length_all = num_chi_prompt_tokens + max_sequence_length - 2 - - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length_all, - 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.to(device) - - prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask) - prompt_embeds = prompt_embeds[0].to(dtype=dtype, device=device) - - return prompt_embeds, prompt_attention_mask - - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, - lora_scale: Optional[float] = None, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`, *optional*): - prompt to be encoded - - num_images_per_prompt (`int`, *optional*, defaults to 1): - number of images that should be generated per prompt - device: (`torch.device`, *optional*): - 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. - clean_caption (`bool`, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. - complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): - If `complex_human_instruction` is not empty, the function will use the complex Human instruction for - the prompt. - """ - - if device is None: - device = self._execution_device - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - else: - dtype = None - - # 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, SanaLoraLoaderMixin): - 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) - - if getattr(self, "tokenizer", None) is not None: - self.tokenizer.padding_side = "right" - - # See Section 3.1. of the paper. - max_length = max_sequence_length - select_index = [0] + list(range(-max_length + 1, 0)) - - if prompt_embeds is None: - prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( - prompt=prompt, - device=device, - dtype=dtype, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - complex_human_instruction=complex_human_instruction, - ) - - prompt_embeds = prompt_embeds[:, select_index] - prompt_attention_mask = prompt_attention_mask[:, select_index] - - bs_embed, 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(bs_embed * num_images_per_prompt, seq_len, -1) - prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1) - prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1) - - if self.text_encoder is not None: - if isinstance(self, SanaLoraLoaderMixin) 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, prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - num_inference_steps, - timesteps, - max_timesteps, - intermediate_timesteps, - callback_on_step_end_tensor_inputs=None, - prompt_embeds=None, - 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 timesteps is not None and len(timesteps) != num_inference_steps + 1: - raise ValueError("If providing custom timesteps, `timesteps` must be of length `num_inference_steps + 1`.") - - if timesteps is not None and max_timesteps is not None: - raise ValueError("If providing custom timesteps, `max_timesteps` should not be provided.") - - if timesteps is None and max_timesteps is None: - raise ValueError("Should provide either `timesteps` or `max_timesteps`.") - - if intermediate_timesteps is not None and num_inference_steps != 2: - raise ValueError("Intermediate timesteps for SCM is not supported when num_inference_steps != 2.") - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.prepare_latents - 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=device, dtype=dtype) - - 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 - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def attention_kwargs(self): - return self._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, - num_inference_steps: int = 2, - 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, - 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", - 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"], - max_sequence_length: int = 300, - 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.", - "Here are examples of how to transform or refine prompts:", - "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.", - "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.", - "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", - "User Prompt: ", - ], - ) -> Union[SanaPipelineOutput, 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. If not defined, one has to pass `prompt_embeds`. - instead. - 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. - max_timesteps (`float`, *optional*, defaults to 1.57080): - 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*): - 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.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. - height (`int`, *optional*, defaults to self.unet.config.sample_size): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to self.unet.config.sample_size): - The width in pixels of the generated image. - 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*): - 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 ge 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. - 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.IFPipelineOutput`] instead of a plain tuple. - attention_kwargs: - 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). - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - use_resolution_binning (`bool` defaults to `True`): - If set to `True`, the requested height and width are first mapped to the closest resolutions using - `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to - the requested resolution. Useful for generating non-square images. - 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 `300`): - Maximum sequence length to use with the `prompt`. - 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. - - Examples: - - Returns: - [`~pipelines.sana.pipeline_output.SanaPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.sana.pipeline_output.SanaPipelineOutput`] 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 - if use_resolution_binning: - if self.transformer.config.sample_size == 32: - aspect_ratio_bin = ASPECT_RATIO_1024_BIN - else: - raise ValueError("Invalid sample size") - orig_height, orig_width = height, width - height, width = self.image_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin) - - self.check_inputs( - prompt=prompt, - height=height, - width=width, - num_inference_steps=num_inference_steps, - timesteps=timesteps, - max_timesteps=max_timesteps, - intermediate_timesteps=intermediate_timesteps, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - prompt_embeds=prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - ) - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._interrupt = False - - # 2. Default height and width to transformer - 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.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None - - # 3. Encode input prompt - ( - prompt_embeds, - prompt_attention_mask, - ) = self.encode_prompt( - prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - prompt_embeds=prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - complex_human_instruction=complex_human_instruction, - lora_scale=lora_scale, - ) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - device, - timesteps, - sigmas=None, - max_timesteps=max_timesteps, - intermediate_timesteps=intermediate_timesteps, - ) - if hasattr(self.scheduler, "set_begin_index"): - self.scheduler.set_begin_index(0) - - # 5. Prepare latents. - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - torch.float32, - device, - generator, - latents, - ) - - latents = latents * self.scheduler.config.sigma_data - - guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32) - guidance = guidance.expand(latents.shape[0]).to(prompt_embeds.dtype) - guidance = guidance * self.transformer.config.guidance_embeds_scale - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - timesteps = timesteps[:-1] - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - transformer_dtype = self.transformer.dtype - 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]) - latents_model_input = latents / self.scheduler.config.sigma_data - - scm_timestep = torch.sin(timestep) / (torch.cos(timestep) + torch.sin(timestep)) - - scm_timestep_expanded = scm_timestep.view(-1, 1, 1, 1) - latent_model_input = latents_model_input * torch.sqrt( - scm_timestep_expanded**2 + (1 - scm_timestep_expanded) ** 2 - ) - - # predict noise model_output - noise_pred = self.transformer( - latent_model_input.to(dtype=transformer_dtype), - encoder_hidden_states=prompt_embeds.to(dtype=transformer_dtype), - encoder_attention_mask=prompt_attention_mask, - guidance=guidance, - timestep=scm_timestep, - return_dict=False, - attention_kwargs=self.attention_kwargs, - )[0] - - noise_pred = ( - (1 - 2 * scm_timestep_expanded) * latent_model_input - + (1 - 2 * scm_timestep_expanded + 2 * scm_timestep_expanded**2) * noise_pred - ) / torch.sqrt(scm_timestep_expanded**2 + (1 - scm_timestep_expanded) ** 2) - noise_pred = noise_pred.float() * self.scheduler.config.sigma_data - - # compute previous image: x_t -> x_t-1 - latents, denoised = self.scheduler.step( - noise_pred, timestep, latents, **extra_step_kwargs, return_dict=False - ) - - 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 = denoised / self.scheduler.config.sigma_data - if output_type == "latent": - image = latents - else: - latents = latents.to(self.vae.dtype) - torch_accelerator_module = getattr(torch, get_device(), torch.cuda) - oom_error = ( - torch.OutOfMemoryError - if is_torch_version(">=", "2.5.0") - else torch_accelerator_module.OutOfMemoryError - ) - try: - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - except oom_error as e: - warnings.warn( - f"{e}. \n" - f"Try to use VAE tiling for large images. For example: \n" - f"pipe.vae.enable_tiling(tile_sample_min_width=512, tile_sample_min_height=512)" - ) - if use_resolution_binning: - image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height) - - 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 SanaPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py deleted file mode 100644 index bf290c3ced56da0013469625578093fa7f186434..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py +++ /dev/null @@ -1,981 +0,0 @@ -# Copyright 2025 PixArt-Sigma Authors 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 inspect -import re -import urllib.parse as ul -import warnings -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -import torch.nn.functional as F -from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, PixArtImageProcessor -from ...loaders import SanaLoraLoaderMixin -from ...models import AutoencoderDC, SanaTransformer2DModel -from ...schedulers import DPMSolverMultistepScheduler -from ...utils import ( - BACKENDS_MAPPING, - USE_PEFT_BACKEND, - is_bs4_available, - is_ftfy_available, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import get_device, is_torch_version, randn_tensor -from ..pipeline_utils import DiffusionPipeline -from ..pixart_alpha.pipeline_pixart_alpha import ASPECT_RATIO_1024_BIN -from .pipeline_output import SanaPipelineOutput - - -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_bs4_available(): - from bs4 import BeautifulSoup - -if is_ftfy_available(): - import ftfy - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import SanaSprintImg2ImgPipeline - >>> from diffusers.utils.loading_utils import load_image - - >>> pipe = SanaSprintImg2ImgPipeline.from_pretrained( - ... "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers", torch_dtype=torch.bfloat16 - ... ) - >>> pipe.to("cuda") - - >>> image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png" - ... ) - - - >>> image = pipe(prompt="a cute pink bear", image=image, strength=0.5, height=832, width=480).images[0] - >>> image[0].save("output.png") - ``` -""" - - -# 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 SanaSprintImg2ImgPipeline(DiffusionPipeline, SanaLoraLoaderMixin): - r""" - Pipeline for text-to-image generation using [SANA-Sprint](https://huggingface.co/papers/2503.09641). - """ - - # fmt: off - bad_punct_regex = re.compile( - r"[" + "#®•©™&@·º½¾¿¡§~" + r"\)" + r"\(" + r"\]" + r"\[" + r"\}" + r"\{" + r"\|" + "\\" + r"\/" + r"\*" + r"]{1,}") - # fmt: on - - model_cpu_offload_seq = "text_encoder->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__( - self, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], - text_encoder: Gemma2PreTrainedModel, - vae: AutoencoderDC, - transformer: SanaTransformer2DModel, - scheduler: DPMSolverMultistepScheduler, - ): - super().__init__() - - self.register_modules( - tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler - ) - - self.vae_scale_factor = ( - 2 ** (len(self.vae.config.encoder_block_out_channels) - 1) - if hasattr(self, "vae") and self.vae is not None - else 32 - ) - self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor) - - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.enable_vae_slicing - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.disable_vae_slicing - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.enable_vae_tiling - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds - def _get_gemma_prompt_embeds( - self, - prompt: Union[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, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`, *optional*): - prompt to be encoded - device: (`torch.device`, *optional*): - torch device to place the resulting embeddings on - clean_caption (`bool`, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. - complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): - If `complex_human_instruction` is not empty, the function will use the complex Human instruction for - the prompt. - """ - prompt = [prompt] if isinstance(prompt, str) else prompt - - if getattr(self, "tokenizer", None) is not None: - self.tokenizer.padding_side = "right" - - prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) - - # prepare complex human instruction - if not complex_human_instruction: - max_length_all = max_sequence_length - else: - chi_prompt = "\n".join(complex_human_instruction) - prompt = [chi_prompt + p for p in prompt] - num_chi_prompt_tokens = len(self.tokenizer.encode(chi_prompt)) - max_length_all = num_chi_prompt_tokens + max_sequence_length - 2 - - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_length_all, - 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.to(device) - - prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask) - prompt_embeds = prompt_embeds[0].to(dtype=dtype, device=device) - - return prompt_embeds, prompt_attention_mask - - # Copied from diffusers.pipelines.sana.pipeline_sana_sprint.SanaSprintPipeline.encode_prompt - def encode_prompt( - self, - prompt: Union[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, - clean_caption: bool = False, - max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, - lora_scale: Optional[float] = None, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`, *optional*): - prompt to be encoded - - num_images_per_prompt (`int`, *optional*, defaults to 1): - number of images that should be generated per prompt - device: (`torch.device`, *optional*): - 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. - clean_caption (`bool`, defaults to `False`): - If `True`, the function will preprocess and clean the provided caption before encoding. - max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. - complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): - If `complex_human_instruction` is not empty, the function will use the complex Human instruction for - the prompt. - """ - - if device is None: - device = self._execution_device - - if self.text_encoder is not None: - dtype = self.text_encoder.dtype - else: - dtype = None - - # 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, SanaLoraLoaderMixin): - 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) - - if getattr(self, "tokenizer", None) is not None: - self.tokenizer.padding_side = "right" - - # See Section 3.1. of the paper. - max_length = max_sequence_length - select_index = [0] + list(range(-max_length + 1, 0)) - - if prompt_embeds is None: - prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( - prompt=prompt, - device=device, - dtype=dtype, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - complex_human_instruction=complex_human_instruction, - ) - - prompt_embeds = prompt_embeds[:, select_index] - prompt_attention_mask = prompt_attention_mask[:, select_index] - - bs_embed, 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(bs_embed * num_images_per_prompt, seq_len, -1) - prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1) - prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1) - - if self.text_encoder is not None: - if isinstance(self, SanaLoraLoaderMixin) 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, prompt_attention_mask - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # 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, - strength, - height, - width, - num_inference_steps, - timesteps, - max_timesteps, - intermediate_timesteps, - callback_on_step_end_tensor_inputs=None, - prompt_embeds=None, - 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 % 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 timesteps is not None and len(timesteps) != num_inference_steps + 1: - raise ValueError("If providing custom timesteps, `timesteps` must be of length `num_inference_steps + 1`.") - - if timesteps is not None and max_timesteps is not None: - raise ValueError("If providing custom timesteps, `max_timesteps` should not be provided.") - - if timesteps is None and max_timesteps is None: - raise ValueError("Should provide either `timesteps` or `max_timesteps`.") - - if intermediate_timesteps is not None and num_inference_steps != 2: - raise ValueError("Intermediate timesteps for SCM is not supported when num_inference_steps != 2.") - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing - def _text_preprocessing(self, text, clean_caption=False): - if clean_caption and not is_bs4_available(): - logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if clean_caption and not is_ftfy_available(): - logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) - logger.warning("Setting `clean_caption` to False...") - clean_caption = False - - if not isinstance(text, (tuple, list)): - text = [text] - - def process(text: str): - if clean_caption: - text = self._clean_caption(text) - text = self._clean_caption(text) - else: - text = text.lower().strip() - return text - - return [process(t) for t in text] - - # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption - def _clean_caption(self, caption): - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - - # \n - caption = re.sub(r"\\n", " ", caption) - - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) - - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - - caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - - caption = ftfy.fix_text(caption) - caption = html.unescape(html.unescape(caption)) - - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... - - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - - caption.strip() - - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - - return caption.strip() - - def prepare_image( - self, - image: PipelineImageInput, - width: int, - height: int, - device: torch.device, - dtype: torch.dtype, - ): - if isinstance(image, torch.Tensor): - if image.ndim == 3: - image = image.unsqueeze(0) - # Resize if current dimensions do not match target dimensions. - if image.shape[2] != height or image.shape[3] != width: - image = F.interpolate(image, size=(height, width), mode="bilinear", align_corners=False) - - image = self.image_processor.preprocess(image, height=height, width=width) - - else: - image = self.image_processor.preprocess(image, height=height, width=width) - - image = image.to(device=device, dtype=dtype) - - return image - - def prepare_latents( - self, image, timestep, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None - ): - if latents is not None: - return latents.to(device=device, dtype=dtype) - - shape = ( - batch_size, - num_channels_latents, - int(height) // self.vae_scale_factor, - int(width) // self.vae_scale_factor, - ) - - if image.shape[1] != num_channels_latents: - image = self.vae.encode(image).latent - image_latents = image * self.vae.config.scaling_factor * self.scheduler.config.sigma_data - 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 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." - ) - - # adapt from https://github.com/huggingface/diffusers/blob/c36f8487df35895421c15f351c7d360bd680[…]/examples/research_projects/sana/train_sana_sprint_diffusers.py - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) * self.scheduler.config.sigma_data - latents = torch.cos(timestep) * image_latents + torch.sin(timestep) * noise - return latents - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def attention_kwargs(self): - return self._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, - num_inference_steps: int = 2, - 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, - 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", - 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"], - max_sequence_length: int = 300, - 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.", - "Here are examples of how to transform or refine prompts:", - "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.", - "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.", - "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", - "User Prompt: ", - ], - ) -> Union[SanaPipelineOutput, 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. If not defined, one has to pass `prompt_embeds`. - instead. - 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. - max_timesteps (`float`, *optional*, defaults to 1.57080): - 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*): - 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.5): - 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. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - height (`int`, *optional*, defaults to self.unet.config.sample_size): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to self.unet.config.sample_size): - The width in pixels of the generated image. - 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*): - 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 ge 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. - 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.IFPipelineOutput`] instead of a plain tuple. - attention_kwargs: - 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). - clean_caption (`bool`, *optional*, defaults to `True`): - Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to - be installed. If the dependencies are not installed, the embeddings will be created from the raw - prompt. - use_resolution_binning (`bool` defaults to `True`): - If set to `True`, the requested height and width are first mapped to the closest resolutions using - `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to - the requested resolution. Useful for generating non-square images. - 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 `300`): - Maximum sequence length to use with the `prompt`. - 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. - - Examples: - - Returns: - [`~pipelines.sana.pipeline_output.SanaPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.sana.pipeline_output.SanaPipelineOutput`] 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 - if use_resolution_binning: - if self.transformer.config.sample_size == 32: - aspect_ratio_bin = ASPECT_RATIO_1024_BIN - else: - raise ValueError("Invalid sample size") - orig_height, orig_width = height, width - height, width = self.image_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin) - - self.check_inputs( - prompt=prompt, - strength=strength, - height=height, - width=width, - num_inference_steps=num_inference_steps, - timesteps=timesteps, - max_timesteps=max_timesteps, - intermediate_timesteps=intermediate_timesteps, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - prompt_embeds=prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - ) - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._interrupt = False - - # 2. Default height and width to transformer - 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.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None - - # 2. Preprocess image - init_image = self.prepare_image(image, width, height, device, self.vae.dtype) - - # 3. Encode input prompt - ( - prompt_embeds, - prompt_attention_mask, - ) = self.encode_prompt( - prompt, - num_images_per_prompt=num_images_per_prompt, - device=device, - prompt_embeds=prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - clean_caption=clean_caption, - max_sequence_length=max_sequence_length, - complex_human_instruction=complex_human_instruction, - lora_scale=lora_scale, - ) - - # 5. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - device, - timesteps, - sigmas=None, - max_timesteps=max_timesteps, - intermediate_timesteps=intermediate_timesteps, - ) - if hasattr(self.scheduler, "set_begin_index"): - self.scheduler.set_begin_index(0) - - 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] - - # 5. Prepare latents. - latent_channels = self.transformer.config.in_channels - latents = self.prepare_latents( - init_image, - latent_timestep, - batch_size * num_images_per_prompt, - latent_channels, - height, - width, - torch.float32, - device, - generator, - latents, - ) - - # I think this is redundant given the scaling in prepare_latents - # latents = latents * self.scheduler.config.sigma_data - - guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32) - guidance = guidance.expand(latents.shape[0]).to(prompt_embeds.dtype) - guidance = guidance * self.transformer.config.guidance_embeds_scale - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - timesteps = timesteps[:-1] - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - transformer_dtype = self.transformer.dtype - 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]) - latents_model_input = latents / self.scheduler.config.sigma_data - - scm_timestep = torch.sin(timestep) / (torch.cos(timestep) + torch.sin(timestep)) - - scm_timestep_expanded = scm_timestep.view(-1, 1, 1, 1) - latent_model_input = latents_model_input * torch.sqrt( - scm_timestep_expanded**2 + (1 - scm_timestep_expanded) ** 2 - ) - - # predict noise model_output - noise_pred = self.transformer( - latent_model_input.to(dtype=transformer_dtype), - encoder_hidden_states=prompt_embeds.to(dtype=transformer_dtype), - encoder_attention_mask=prompt_attention_mask, - guidance=guidance, - timestep=scm_timestep, - return_dict=False, - attention_kwargs=self.attention_kwargs, - )[0] - - noise_pred = ( - (1 - 2 * scm_timestep_expanded) * latent_model_input - + (1 - 2 * scm_timestep_expanded + 2 * scm_timestep_expanded**2) * noise_pred - ) / torch.sqrt(scm_timestep_expanded**2 + (1 - scm_timestep_expanded) ** 2) - noise_pred = noise_pred.float() * self.scheduler.config.sigma_data - - # compute previous image: x_t -> x_t-1 - latents, denoised = self.scheduler.step( - noise_pred, timestep, latents, **extra_step_kwargs, return_dict=False - ) - - 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 = denoised / self.scheduler.config.sigma_data - if output_type == "latent": - image = latents - else: - latents = latents.to(self.vae.dtype) - torch_accelerator_module = getattr(torch, get_device(), torch.cuda) - oom_error = ( - torch.OutOfMemoryError - if is_torch_version(">=", "2.5.0") - else torch_accelerator_module.OutOfMemoryError - ) - try: - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - except oom_error as e: - warnings.warn( - f"{e}. \n" - f"Try to use VAE tiling for large images. For example: \n" - f"pipe.vae.enable_tiling(tile_sample_min_width=512, tile_sample_min_height=512)" - ) - if use_resolution_binning: - image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height) - - 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 SanaPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/__init__.py deleted file mode 100644 index 70f5b1a547c4b90e28109843ae3be2fca2e98c88..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/__init__.py +++ /dev/null @@ -1,49 +0,0 @@ -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_output"] = ["SemanticStableDiffusionPipelineOutput"] - _import_structure["pipeline_semantic_stable_diffusion"] = ["SemanticStableDiffusionPipeline"] - - -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_semantic_stable_diffusion import SemanticStableDiffusionPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 8e840fb86cddbaea146204e55e9ee3a52a02fd8c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index df20fd0a87dd2faaeddbed0d520352dcb28839da..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/__pycache__/pipeline_semantic_stable_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/__pycache__/pipeline_semantic_stable_diffusion.cpython-310.pyc deleted file mode 100644 index 000477e67eaa9d58928a4fc22e5ee1cc32b1483c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/__pycache__/pipeline_semantic_stable_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/pipeline_output.py deleted file mode 100644 index 34991299398115f439537b77e1f1fc8a83e0d431..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/pipeline_output.py +++ /dev/null @@ -1,25 +0,0 @@ -from dataclasses import dataclass -from typing import List, Optional, Union - -import numpy as np -import PIL.Image - -from ...utils import BaseOutput - - -@dataclass -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, - num_channels)`. - 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]] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py deleted file mode 100644 index a5f67bffe6f07b50aa8e94a87cf3da69f26da325..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py +++ /dev/null @@ -1,733 +0,0 @@ -import inspect -from itertools import repeat -from typing import Callable, List, Optional, Union - -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer - -from ...image_processor import VaeImageProcessor -from ...models import AutoencoderKL, UNet2DConditionModel -from ...pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import deprecate, is_torch_xla_available, logging -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from .pipeline_output import SemanticStableDiffusionPipelineOutput - - -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 - - -class SemanticStableDiffusionPipeline(DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin): - _last_supported_version = "0.33.1" - r""" - Pipeline for text-to-image generation using Stable Diffusion with latent editing. - - This model inherits from [`DiffusionPipeline`] and builds on the [`StableDiffusionPipeline`]. Check the superclass - documentation for the generic methods implemented for all pipelines (downloading, saving, running on a particular - device, etc.). - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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 ([`Q16SafetyChecker`]): - Classification module that estimates whether generated images could be considered offensive or harmful. - Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details - about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - 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, - 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}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @torch.no_grad() - def __call__( - self, - prompt: Union[str, List[str]], - 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: 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], 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, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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`. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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`): - Whether the corresponding prompt in `editing_prompt` should be increased or decreased. - 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): - 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`): - 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): - 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, - momentum is disabled. Momentum is already built up during warmup (for diffusion steps smaller than - `sld_warmup_steps`). Momentum is only added to latent guidance once all warmup periods are finished. - edit_mom_beta (`float`, *optional*, defaults to 0.4): - 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`): - 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 - correspond to `num_inference_steps`. - - Examples: - - ```py - >>> import torch - >>> from diffusers import SemanticStableDiffusionPipeline - - >>> pipe = SemanticStableDiffusionPipeline.from_pretrained( - ... "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> out = pipe( - ... prompt="a photo of the face of a woman", - ... num_images_per_prompt=1, - ... guidance_scale=7, - ... editing_prompt=[ - ... "smiling, smile", # Concepts to apply - ... "glasses, wearing glasses", - ... "curls, wavy hair, curly hair", - ... "beard, full beard, mustache", - ... ], - ... reverse_editing_direction=[ - ... False, - ... False, - ... False, - ... False, - ... ], # Direction of guidance i.e. increase all concepts - ... edit_warmup_steps=[10, 10, 10, 10], # Warmup period for each concept - ... edit_guidance_scale=[4, 5, 5, 5.4], # Guidance scale for each concept - ... edit_threshold=[ - ... 0.99, - ... 0.975, - ... 0.925, - ... 0.96, - ... ], # Threshold for each concept. Threshold equals the percentile of the latent space that will be discarded. I.e. threshold=0.99 uses 1% of the latent dimensions - ... edit_momentum_scale=0.3, # Momentum scale that will be added to the latent guidance - ... edit_mom_beta=0.6, # Momentum beta - ... edit_weights=[1, 1, 1, 1, 1], # Weights of the individual concepts against each other - ... ) - >>> image = out.images[0] - ``` - - Returns: - [`~pipelines.semantic_stable_diffusion.SemanticStableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, - [`~pipelines.semantic_stable_diffusion.SemanticStableDiffusionPipelineOutput`] 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. - """ - # 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) - - # 2. Define call parameters - batch_size = 1 if isinstance(prompt, str) else len(prompt) - device = self._execution_device - - if editing_prompt: - enable_edit_guidance = True - if isinstance(editing_prompt, str): - editing_prompt = [editing_prompt] - enabled_editing_prompts = len(editing_prompt) - elif editing_prompt_embeddings is not None: - enable_edit_guidance = True - enabled_editing_prompts = editing_prompt_embeddings.shape[0] - else: - enabled_editing_prompts = 0 - enable_edit_guidance = False - - # get prompt text embeddings - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - - if text_input_ids.shape[-1] > self.tokenizer.model_max_length: - removed_text = self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]) - 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}" - ) - text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length] - text_embeddings = self.text_encoder(text_input_ids.to(device))[0] - - # duplicate text embeddings for each generation per prompt, using mps friendly method - bs_embed, seq_len, _ = text_embeddings.shape - text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1) - text_embeddings = text_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) - - if enable_edit_guidance: - # get safety text embeddings - if editing_prompt_embeddings is None: - edit_concepts_input = self.tokenizer( - [x for item in editing_prompt for x in repeat(item, batch_size)], - padding="max_length", - max_length=self.tokenizer.model_max_length, - return_tensors="pt", - ) - - edit_concepts_input_ids = edit_concepts_input.input_ids - - if edit_concepts_input_ids.shape[-1] > self.tokenizer.model_max_length: - removed_text = self.tokenizer.batch_decode( - edit_concepts_input_ids[:, self.tokenizer.model_max_length :] - ) - 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}" - ) - edit_concepts_input_ids = edit_concepts_input_ids[:, : self.tokenizer.model_max_length] - edit_concepts = self.text_encoder(edit_concepts_input_ids.to(device))[0] - else: - edit_concepts = editing_prompt_embeddings.to(device).repeat(batch_size, 1, 1) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - bs_embed_edit, seq_len_edit, _ = edit_concepts.shape - edit_concepts = edit_concepts.repeat(1, num_images_per_prompt, 1) - edit_concepts = edit_concepts.view(bs_embed_edit * num_images_per_prompt, seq_len_edit, -1) - - # 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 = guidance_scale > 1.0 - # get unconditional embeddings for classifier free guidance - - if do_classifier_free_guidance: - uncond_tokens: List[str] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif 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 - - max_length = text_input_ids.shape[-1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(device))[0] - - # duplicate unconditional embeddings for each generation per prompt, using mps friendly method - seq_len = uncond_embeddings.shape[1] - uncond_embeddings = uncond_embeddings.repeat(1, num_images_per_prompt, 1) - uncond_embeddings = uncond_embeddings.view(batch_size * num_images_per_prompt, seq_len, -1) - - # 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 enable_edit_guidance: - text_embeddings = torch.cat([uncond_embeddings, text_embeddings, edit_concepts]) - else: - text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) - # get the initial random noise unless the user supplied it - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. 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, - text_embeddings.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # Initialize edit_momentum to None - edit_momentum = None - - self.uncond_estimates = None - self.text_estimates = None - self.edit_estimates = None - self.sem_guidance = None - - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = ( - torch.cat([latents] * (2 + enabled_editing_prompts)) if do_classifier_free_guidance else latents - ) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_out = noise_pred.chunk(2 + enabled_editing_prompts) # [b,4, 64, 64] - noise_pred_uncond, noise_pred_text = noise_pred_out[0], noise_pred_out[1] - noise_pred_edit_concepts = noise_pred_out[2:] - - # default text guidance - noise_guidance = guidance_scale * (noise_pred_text - noise_pred_uncond) - # noise_guidance = (noise_pred_text - noise_pred_edit_concepts[0]) - - if self.uncond_estimates is None: - self.uncond_estimates = torch.zeros((num_inference_steps + 1, *noise_pred_uncond.shape)) - self.uncond_estimates[i] = noise_pred_uncond.detach().cpu() - - if self.text_estimates is None: - self.text_estimates = torch.zeros((num_inference_steps + 1, *noise_pred_text.shape)) - self.text_estimates[i] = noise_pred_text.detach().cpu() - - if self.edit_estimates is None and enable_edit_guidance: - self.edit_estimates = torch.zeros( - (num_inference_steps + 1, len(noise_pred_edit_concepts), *noise_pred_edit_concepts[0].shape) - ) - - if self.sem_guidance is None: - self.sem_guidance = torch.zeros((num_inference_steps + 1, *noise_pred_text.shape)) - - if edit_momentum is None: - edit_momentum = torch.zeros_like(noise_guidance) - - if enable_edit_guidance: - concept_weights = torch.zeros( - (len(noise_pred_edit_concepts), noise_guidance.shape[0]), - device=device, - dtype=noise_guidance.dtype, - ) - noise_guidance_edit = torch.zeros( - (len(noise_pred_edit_concepts), *noise_guidance.shape), - device=device, - dtype=noise_guidance.dtype, - ) - # noise_guidance_edit = torch.zeros_like(noise_guidance) - warmup_inds = [] - for c, noise_pred_edit_concept in enumerate(noise_pred_edit_concepts): - self.edit_estimates[i, c] = noise_pred_edit_concept - if isinstance(edit_guidance_scale, list): - edit_guidance_scale_c = edit_guidance_scale[c] - else: - edit_guidance_scale_c = edit_guidance_scale - - if isinstance(edit_threshold, list): - edit_threshold_c = edit_threshold[c] - else: - edit_threshold_c = edit_threshold - if isinstance(reverse_editing_direction, list): - reverse_editing_direction_c = reverse_editing_direction[c] - else: - reverse_editing_direction_c = reverse_editing_direction - if edit_weights: - edit_weight_c = edit_weights[c] - else: - edit_weight_c = 1.0 - if isinstance(edit_warmup_steps, list): - edit_warmup_steps_c = edit_warmup_steps[c] - else: - edit_warmup_steps_c = edit_warmup_steps - - if isinstance(edit_cooldown_steps, list): - edit_cooldown_steps_c = edit_cooldown_steps[c] - elif edit_cooldown_steps is None: - edit_cooldown_steps_c = i + 1 - else: - edit_cooldown_steps_c = edit_cooldown_steps - if i >= edit_warmup_steps_c: - warmup_inds.append(c) - if i >= edit_cooldown_steps_c: - noise_guidance_edit[c, :, :, :, :] = torch.zeros_like(noise_pred_edit_concept) - continue - - noise_guidance_edit_tmp = noise_pred_edit_concept - noise_pred_uncond - # tmp_weights = (noise_pred_text - noise_pred_edit_concept).sum(dim=(1, 2, 3)) - tmp_weights = (noise_guidance - noise_pred_edit_concept).sum(dim=(1, 2, 3)) - - tmp_weights = torch.full_like(tmp_weights, edit_weight_c) # * (1 / enabled_editing_prompts) - if reverse_editing_direction_c: - noise_guidance_edit_tmp = noise_guidance_edit_tmp * -1 - concept_weights[c, :] = tmp_weights - - noise_guidance_edit_tmp = noise_guidance_edit_tmp * edit_guidance_scale_c - - # torch.quantile function expects float32 - if noise_guidance_edit_tmp.dtype == torch.float32: - tmp = torch.quantile( - torch.abs(noise_guidance_edit_tmp).flatten(start_dim=2), - edit_threshold_c, - dim=2, - keepdim=False, - ) - else: - tmp = torch.quantile( - torch.abs(noise_guidance_edit_tmp).flatten(start_dim=2).to(torch.float32), - edit_threshold_c, - dim=2, - keepdim=False, - ).to(noise_guidance_edit_tmp.dtype) - - noise_guidance_edit_tmp = torch.where( - torch.abs(noise_guidance_edit_tmp) >= tmp[:, :, None, None], - noise_guidance_edit_tmp, - torch.zeros_like(noise_guidance_edit_tmp), - ) - noise_guidance_edit[c, :, :, :, :] = noise_guidance_edit_tmp - - # noise_guidance_edit = noise_guidance_edit + noise_guidance_edit_tmp - - warmup_inds = torch.tensor(warmup_inds).to(device) - if len(noise_pred_edit_concepts) > warmup_inds.shape[0] > 0: - concept_weights = concept_weights.to("cpu") # Offload to cpu - noise_guidance_edit = noise_guidance_edit.to("cpu") - - concept_weights_tmp = torch.index_select(concept_weights.to(device), 0, warmup_inds) - concept_weights_tmp = torch.where( - concept_weights_tmp < 0, torch.zeros_like(concept_weights_tmp), concept_weights_tmp - ) - concept_weights_tmp = concept_weights_tmp / concept_weights_tmp.sum(dim=0) - # concept_weights_tmp = torch.nan_to_num(concept_weights_tmp) - - noise_guidance_edit_tmp = torch.index_select(noise_guidance_edit.to(device), 0, warmup_inds) - noise_guidance_edit_tmp = torch.einsum( - "cb,cbijk->bijk", concept_weights_tmp, noise_guidance_edit_tmp - ) - noise_guidance = noise_guidance + noise_guidance_edit_tmp - - self.sem_guidance[i] = noise_guidance_edit_tmp.detach().cpu() - - del noise_guidance_edit_tmp - del concept_weights_tmp - concept_weights = concept_weights.to(device) - noise_guidance_edit = noise_guidance_edit.to(device) - - concept_weights = torch.where( - concept_weights < 0, torch.zeros_like(concept_weights), concept_weights - ) - - concept_weights = torch.nan_to_num(concept_weights) - - noise_guidance_edit = torch.einsum("cb,cbijk->bijk", concept_weights, noise_guidance_edit) - noise_guidance_edit = noise_guidance_edit.to(edit_momentum.device) - - noise_guidance_edit = noise_guidance_edit + edit_momentum_scale * edit_momentum - - edit_momentum = edit_mom_beta * edit_momentum + (1 - edit_mom_beta) * noise_guidance_edit - - if warmup_inds.shape[0] == len(noise_pred_edit_concepts): - noise_guidance = noise_guidance + noise_guidance_edit - self.sem_guidance[i] = noise_guidance_edit.detach().cpu() - - if sem_guidance is not None: - edit_guidance = sem_guidance[i].to(device) - noise_guidance = noise_guidance + edit_guidance - - noise_pred = noise_pred_uncond + noise_guidance - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # call the callback, if provided - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # 8. Post-processing - 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, text_embeddings.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) - - if not return_dict: - return (image, has_nsfw_concept) - - return SemanticStableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__init__.py deleted file mode 100644 index 4ed563c4a51f6e627c06711b60fe3a0709ff22f7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__init__.py +++ /dev/null @@ -1,71 +0,0 @@ -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["camera"] = ["create_pan_cameras"] - _import_structure["pipeline_shap_e"] = ["ShapEPipeline"] - _import_structure["pipeline_shap_e_img2img"] = ["ShapEImg2ImgPipeline"] - _import_structure["renderer"] = [ - "BoundingBoxVolume", - "ImportanceRaySampler", - "MLPNeRFModelOutput", - "MLPNeRSTFModel", - "ShapEParamsProjModel", - "ShapERenderer", - "StratifiedRaySampler", - "VoidNeRFModel", - ] - -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 .camera import create_pan_cameras - from .pipeline_shap_e import ShapEPipeline - from .pipeline_shap_e_img2img import ShapEImg2ImgPipeline - from .renderer import ( - BoundingBoxVolume, - ImportanceRaySampler, - MLPNeRFModelOutput, - MLPNeRSTFModel, - ShapEParamsProjModel, - ShapERenderer, - StratifiedRaySampler, - VoidNeRFModel, - ) - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index e1c90ad164b4f22304ca6e0c27df5422f4a51663..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/camera.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/camera.cpython-310.pyc deleted file mode 100644 index b84be25e00bcd4da098dcec196d4e614a5d98313..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/camera.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/pipeline_shap_e.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/pipeline_shap_e.cpython-310.pyc deleted file mode 100644 index a6ef6d1bac621d36027cf690a4b3ed44f161350e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/pipeline_shap_e.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/pipeline_shap_e_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/pipeline_shap_e_img2img.cpython-310.pyc deleted file mode 100644 index 955621a24ec933944c6add0f37b6b8d6da193934..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/pipeline_shap_e_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/renderer.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/renderer.cpython-310.pyc deleted file mode 100644 index ffa9d6bb4cf29267e8ae4179d0a7bc9122864381..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/__pycache__/renderer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/camera.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/camera.py deleted file mode 100644 index 31e1759d6154df151bbe13f61c67785ab5078a12..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/camera.py +++ /dev/null @@ -1,147 +0,0 @@ -# Copyright 2025 Open 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. - -from dataclasses import dataclass -from typing import Tuple - -import numpy as np -import torch - - -@dataclass -class DifferentiableProjectiveCamera: - """ - Implements a batch, differentiable, standard pinhole camera - """ - - origin: torch.Tensor # [batch_size x 3] - x: torch.Tensor # [batch_size x 3] - y: torch.Tensor # [batch_size x 3] - z: torch.Tensor # [batch_size x 3] - width: int - height: int - x_fov: float - y_fov: float - shape: Tuple[int] - - def __post_init__(self): - assert self.x.shape[0] == self.y.shape[0] == self.z.shape[0] == self.origin.shape[0] - assert self.x.shape[1] == self.y.shape[1] == self.z.shape[1] == self.origin.shape[1] == 3 - assert len(self.x.shape) == len(self.y.shape) == len(self.z.shape) == len(self.origin.shape) == 2 - - def resolution(self): - return torch.from_numpy(np.array([self.width, self.height], dtype=np.float32)) - - def fov(self): - return torch.from_numpy(np.array([self.x_fov, self.y_fov], dtype=np.float32)) - - def get_image_coords(self) -> torch.Tensor: - """ - :return: coords of shape (width * height, 2) - """ - pixel_indices = torch.arange(self.height * self.width) - coords = torch.stack( - [ - pixel_indices % self.width, - torch.div(pixel_indices, self.width, rounding_mode="trunc"), - ], - axis=1, - ) - return coords - - @property - def camera_rays(self): - batch_size, *inner_shape = self.shape - inner_batch_size = int(np.prod(inner_shape)) - - coords = self.get_image_coords() - coords = torch.broadcast_to(coords.unsqueeze(0), [batch_size * inner_batch_size, *coords.shape]) - rays = self.get_camera_rays(coords) - - rays = rays.view(batch_size, inner_batch_size * self.height * self.width, 2, 3) - - return rays - - def get_camera_rays(self, coords: torch.Tensor) -> torch.Tensor: - batch_size, *shape, n_coords = coords.shape - assert n_coords == 2 - assert batch_size == self.origin.shape[0] - - flat = coords.view(batch_size, -1, 2) - - res = self.resolution() - fov = self.fov() - - fracs = (flat.float() / (res - 1)) * 2 - 1 - fracs = fracs * torch.tan(fov / 2) - - fracs = fracs.view(batch_size, -1, 2) - directions = ( - self.z.view(batch_size, 1, 3) - + self.x.view(batch_size, 1, 3) * fracs[:, :, :1] - + self.y.view(batch_size, 1, 3) * fracs[:, :, 1:] - ) - directions = directions / directions.norm(dim=-1, keepdim=True) - rays = torch.stack( - [ - torch.broadcast_to(self.origin.view(batch_size, 1, 3), [batch_size, directions.shape[1], 3]), - directions, - ], - dim=2, - ) - return rays.view(batch_size, *shape, 2, 3) - - def resize_image(self, width: int, height: int) -> "DifferentiableProjectiveCamera": - """ - Creates a new camera for the resized view assuming the aspect ratio does not change. - """ - assert width * self.height == height * self.width, "The aspect ratio should not change." - return DifferentiableProjectiveCamera( - origin=self.origin, - x=self.x, - y=self.y, - z=self.z, - width=width, - height=height, - x_fov=self.x_fov, - y_fov=self.y_fov, - ) - - -def create_pan_cameras(size: int) -> DifferentiableProjectiveCamera: - origins = [] - xs = [] - ys = [] - zs = [] - for theta in np.linspace(0, 2 * np.pi, num=20): - z = np.array([np.sin(theta), np.cos(theta), -0.5]) - z /= np.sqrt(np.sum(z**2)) - origin = -z * 4 - x = np.array([np.cos(theta), -np.sin(theta), 0.0]) - y = np.cross(z, x) - origins.append(origin) - xs.append(x) - ys.append(y) - zs.append(z) - return DifferentiableProjectiveCamera( - origin=torch.from_numpy(np.stack(origins, axis=0)).float(), - x=torch.from_numpy(np.stack(xs, axis=0)).float(), - y=torch.from_numpy(np.stack(ys, axis=0)).float(), - z=torch.from_numpy(np.stack(zs, axis=0)).float(), - width=size, - height=size, - x_fov=0.7, - y_fov=0.7, - shape=(1, len(xs)), - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/pipeline_shap_e.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/pipeline_shap_e.py deleted file mode 100644 index 49ddfd1196bf693b5115c971fa7b60f1342c9d71..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/pipeline_shap_e.py +++ /dev/null @@ -1,346 +0,0 @@ -# Copyright 2025 Open 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 math -from dataclasses import dataclass -from typing import List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import CLIPTextModelWithProjection, CLIPTokenizer - -from ...models import PriorTransformer -from ...schedulers import HeunDiscreteScheduler -from ...utils import ( - BaseOutput, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .renderer import ShapERenderer - - -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 DiffusionPipeline - >>> from diffusers.utils import export_to_gif - - >>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - >>> repo = "openai/shap-e" - >>> pipe = DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.float16) - >>> pipe = pipe.to(device) - - >>> guidance_scale = 15.0 - >>> prompt = "a shark" - - >>> images = pipe( - ... prompt, - ... guidance_scale=guidance_scale, - ... num_inference_steps=64, - ... frame_size=256, - ... ).images - - >>> gif_path = export_to_gif(images[0], "shark_3d.gif") - ``` -""" - - -@dataclass -class ShapEPipelineOutput(BaseOutput): - """ - Output class for [`ShapEPipeline`] and [`ShapEImg2ImgPipeline`]. - - Args: - images (`torch.Tensor`) - A list of images for 3D rendering. - """ - - images: Union[List[List[PIL.Image.Image]], List[List[np.ndarray]]] - - -class ShapEPipeline(DiffusionPipeline): - """ - Pipeline for generating latent representation of a 3D asset and rendering with the NeRF method. - - 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: - prior ([`PriorTransformer`]): - The canonical unCLIP prior to approximate the image embedding from the text embedding. - text_encoder ([`~transformers.CLIPTextModelWithProjection`]): - Frozen text-encoder. - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - scheduler ([`HeunDiscreteScheduler`]): - A scheduler to be used in combination with the `prior` model to generate image embedding. - shap_e_renderer ([`ShapERenderer`]): - Shap-E renderer projects the generated latents into parameters of a MLP to create 3D objects with the NeRF - rendering method. - """ - - model_cpu_offload_seq = "text_encoder->prior" - _exclude_from_cpu_offload = ["shap_e_renderer"] - - def __init__( - self, - prior: PriorTransformer, - text_encoder: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - scheduler: HeunDiscreteScheduler, - shap_e_renderer: ShapERenderer, - ): - super().__init__() - - self.register_modules( - prior=prior, - text_encoder=text_encoder, - tokenizer=tokenizer, - scheduler=scheduler, - shap_e_renderer=shap_e_renderer, - ) - - # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - def _encode_prompt( - self, - prompt, - device, - num_images_per_prompt, - do_classifier_free_guidance, - ): - len(prompt) if isinstance(prompt, list) else 1 - - # YiYi Notes: set pad_token_id to be 0, not sure why I can't set in the config file - self.tokenizer.pad_token_id = 0 - # get prompt text embeddings - 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}" - ) - - text_encoder_output = self.text_encoder(text_input_ids.to(device)) - prompt_embeds = text_encoder_output.text_embeds - - prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - # in Shap-E it normalize the prompt_embeds and then later rescale it - prompt_embeds = prompt_embeds / torch.linalg.norm(prompt_embeds, dim=-1, keepdim=True) - - if do_classifier_free_guidance: - negative_prompt_embeds = torch.zeros_like(prompt_embeds) - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - # Rescale the features to have unit variance - prompt_embeds = math.sqrt(prompt_embeds.shape[1]) * prompt_embeds - - return prompt_embeds - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - 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, - guidance_scale: float = 4.0, - frame_size: int = 64, - output_type: Optional[str] = "pil", # pil, np, latent, mesh - return_dict: bool = True, - ): - """ - The call function to the pipeline for generation. - - Args: - 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*): - 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`. - 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`. - frame_size (`int`, *optional*, default to 64): - The width and height of each image frame of the generated 3D output. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated image. Choose between `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`), `"latent"` (`torch.Tensor`), or mesh ([`MeshDecoderOutput`]). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.shap_e.pipeline_shap_e.ShapEPipelineOutput`] instead of a plain - tuple. - - Examples: - - Returns: - [`~pipelines.shap_e.pipeline_shap_e.ShapEPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.shap_e.pipeline_shap_e.ShapEPipelineOutput`] is returned, - otherwise a `tuple` is returned where the first element is a list with the generated images. - """ - - if isinstance(prompt, str): - batch_size = 1 - elif isinstance(prompt, list): - batch_size = len(prompt) - else: - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - device = self._execution_device - - batch_size = batch_size * num_images_per_prompt - - do_classifier_free_guidance = guidance_scale > 1.0 - prompt_embeds = self._encode_prompt(prompt, device, num_images_per_prompt, do_classifier_free_guidance) - - # prior - - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - num_embeddings = self.prior.config.num_embeddings - embedding_dim = self.prior.config.embedding_dim - - latents = self.prepare_latents( - (batch_size, num_embeddings * embedding_dim), - prompt_embeds.dtype, - device, - generator, - latents, - self.scheduler, - ) - - # YiYi notes: for testing only to match ldm, we can directly create a latents with desired shape: batch_size, num_embeddings, embedding_dim - latents = latents.reshape(latents.shape[0], num_embeddings, embedding_dim) - - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - scaled_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - noise_pred = self.prior( - scaled_model_input, - timestep=t, - proj_embedding=prompt_embeds, - ).predicted_image_embedding - - # remove the variance - noise_pred, _ = noise_pred.split( - scaled_model_input.shape[2], dim=2 - ) # batch_size, num_embeddings, embedding_dim - - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond) - - latents = self.scheduler.step( - noise_pred, - timestep=t, - sample=latents, - ).prev_sample - - if XLA_AVAILABLE: - xm.mark_step() - - # Offload all models - self.maybe_free_model_hooks() - - if output_type not in ["np", "pil", "latent", "mesh"]: - raise ValueError( - f"Only the output types `pil`, `np`, `latent` and `mesh` are supported not output_type={output_type}" - ) - - if output_type == "latent": - return ShapEPipelineOutput(images=latents) - - images = [] - if output_type == "mesh": - for i, latent in enumerate(latents): - mesh = self.shap_e_renderer.decode_to_mesh( - latent[None, :], - device, - ) - images.append(mesh) - - else: - # np, pil - for i, latent in enumerate(latents): - image = self.shap_e_renderer.decode_to_image( - latent[None, :], - device, - size=frame_size, - ) - images.append(image) - - images = torch.stack(images) - - images = images.cpu().numpy() - - if output_type == "pil": - images = [self.numpy_to_pil(image) for image in images] - - if not return_dict: - return (images,) - - return ShapEPipelineOutput(images=images) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py deleted file mode 100644 index 55d8b85822c47561caa918eacc9f7c1af5b83302..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py +++ /dev/null @@ -1,333 +0,0 @@ -# Copyright 2025 Open 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. - -from dataclasses import dataclass -from typing import List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPVisionModel - -from ...models import PriorTransformer -from ...schedulers import HeunDiscreteScheduler -from ...utils import ( - BaseOutput, - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .renderer import ShapERenderer - - -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 - >>> from PIL import Image - >>> import torch - >>> from diffusers import DiffusionPipeline - >>> from diffusers.utils import export_to_gif, load_image - - >>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - >>> repo = "openai/shap-e-img2img" - >>> pipe = DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.float16) - >>> pipe = pipe.to(device) - - >>> guidance_scale = 3.0 - >>> image_url = "https://hf.co/datasets/diffusers/docs-images/resolve/main/shap-e/corgi.png" - >>> image = load_image(image_url).convert("RGB") - - >>> images = pipe( - ... image, - ... guidance_scale=guidance_scale, - ... num_inference_steps=64, - ... frame_size=256, - ... ).images - - >>> gif_path = export_to_gif(images[0], "corgi_3d.gif") - ``` -""" - - -@dataclass -class ShapEPipelineOutput(BaseOutput): - """ - Output class for [`ShapEPipeline`] and [`ShapEImg2ImgPipeline`]. - - Args: - images (`torch.Tensor`) - A list of images for 3D rendering. - """ - - images: Union[PIL.Image.Image, np.ndarray] - - -class ShapEImg2ImgPipeline(DiffusionPipeline): - """ - Pipeline for generating latent representation of a 3D asset and rendering with the NeRF method from an image. - - 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: - prior ([`PriorTransformer`]): - The canonical unCLIP prior to approximate the image embedding from the text embedding. - image_encoder ([`~transformers.CLIPVisionModel`]): - Frozen image-encoder. - image_processor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to process images. - scheduler ([`HeunDiscreteScheduler`]): - A scheduler to be used in combination with the `prior` model to generate image embedding. - shap_e_renderer ([`ShapERenderer`]): - Shap-E renderer projects the generated latents into parameters of a MLP to create 3D objects with the NeRF - rendering method. - """ - - model_cpu_offload_seq = "image_encoder->prior" - _exclude_from_cpu_offload = ["shap_e_renderer"] - - def __init__( - self, - prior: PriorTransformer, - image_encoder: CLIPVisionModel, - image_processor: CLIPImageProcessor, - scheduler: HeunDiscreteScheduler, - shap_e_renderer: ShapERenderer, - ): - super().__init__() - - self.register_modules( - prior=prior, - image_encoder=image_encoder, - image_processor=image_processor, - scheduler=scheduler, - shap_e_renderer=shap_e_renderer, - ) - - # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - def _encode_image( - self, - image, - device, - num_images_per_prompt, - do_classifier_free_guidance, - ): - 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): - image = self.image_processor(image, return_tensors="pt").pixel_values[0].unsqueeze(0) - - image = image.to(dtype=self.image_encoder.dtype, device=device) - - image_embeds = self.image_encoder(image)["last_hidden_state"] - image_embeds = image_embeds[:, 1:, :].contiguous() # batch_size, dim, 256 - - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - negative_image_embeds = torch.zeros_like(image_embeds) - - # 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 - image_embeds = torch.cat([negative_image_embeds, image_embeds]) - - return image_embeds - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - image: Union[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, - guidance_scale: float = 4.0, - frame_size: int = 64, - output_type: Optional[str] = "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` 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): - 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*): - 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`. - 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`. - frame_size (`int`, *optional*, default to 64): - The width and height of each image frame of the generated 3D output. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated image. Choose between `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`), `"latent"` (`torch.Tensor`), or mesh ([`MeshDecoderOutput`]). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.shap_e.pipeline_shap_e.ShapEPipelineOutput`] instead of a plain - tuple. - - Examples: - - Returns: - [`~pipelines.shap_e.pipeline_shap_e.ShapEPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.shap_e.pipeline_shap_e.ShapEPipelineOutput`] is returned, - otherwise a `tuple` is returned where the first element is a list with the generated images. - """ - - if isinstance(image, PIL.Image.Image): - batch_size = 1 - elif isinstance(image, torch.Tensor): - batch_size = image.shape[0] - elif isinstance(image, list) and isinstance(image[0], (torch.Tensor, PIL.Image.Image)): - 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)}" - ) - - device = self._execution_device - - batch_size = batch_size * num_images_per_prompt - - do_classifier_free_guidance = guidance_scale > 1.0 - image_embeds = self._encode_image(image, device, num_images_per_prompt, do_classifier_free_guidance) - - # prior - - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - num_embeddings = self.prior.config.num_embeddings - embedding_dim = self.prior.config.embedding_dim - if latents is None: - latents = self.prepare_latents( - (batch_size, num_embeddings * embedding_dim), - image_embeds.dtype, - device, - generator, - latents, - self.scheduler, - ) - - # YiYi notes: for testing only to match ldm, we can directly create a latents with desired shape: batch_size, num_embeddings, embedding_dim - latents = latents.reshape(latents.shape[0], num_embeddings, embedding_dim) - - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - scaled_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - noise_pred = self.prior( - scaled_model_input, - timestep=t, - proj_embedding=image_embeds, - ).predicted_image_embedding - - # remove the variance - noise_pred, _ = noise_pred.split( - scaled_model_input.shape[2], dim=2 - ) # batch_size, num_embeddings, embedding_dim - - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond) - - latents = self.scheduler.step( - noise_pred, - timestep=t, - sample=latents, - ).prev_sample - - if XLA_AVAILABLE: - xm.mark_step() - - if output_type not in ["np", "pil", "latent", "mesh"]: - raise ValueError( - f"Only the output types `pil`, `np`, `latent` and `mesh` are supported not output_type={output_type}" - ) - - # Offload all models - self.maybe_free_model_hooks() - - if output_type == "latent": - return ShapEPipelineOutput(images=latents) - - images = [] - if output_type == "mesh": - for i, latent in enumerate(latents): - mesh = self.shap_e_renderer.decode_to_mesh( - latent[None, :], - device, - ) - images.append(mesh) - - else: - # np, pil - for i, latent in enumerate(latents): - image = self.shap_e_renderer.decode_to_image( - latent[None, :], - device, - size=frame_size, - ) - images.append(image) - - images = torch.stack(images) - - images = images.cpu().numpy() - - if output_type == "pil": - images = [self.numpy_to_pil(image) for image in images] - - if not return_dict: - return (images,) - - return ShapEPipelineOutput(images=images) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/renderer.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/renderer.py deleted file mode 100644 index b268eae806a690a1104ae51da7f7d6ea2333993e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/shap_e/renderer.py +++ /dev/null @@ -1,1050 +0,0 @@ -# Copyright 2025 Open 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 math -from dataclasses import dataclass -from typing import Dict, Optional, Tuple - -import numpy as np -import torch -import torch.nn.functional as F -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models import ModelMixin -from ...utils import BaseOutput -from .camera import create_pan_cameras - - -def sample_pmf(pmf: torch.Tensor, n_samples: int) -> torch.Tensor: - r""" - Sample from the given discrete probability distribution with replacement. - - The i-th bin is assumed to have mass pmf[i]. - - Args: - pmf: [batch_size, *shape, n_samples, 1] where (pmf.sum(dim=-2) == 1).all() - n_samples: number of samples - - Return: - indices sampled with replacement - """ - - *shape, support_size, last_dim = pmf.shape - assert last_dim == 1 - - cdf = torch.cumsum(pmf.view(-1, support_size), dim=1) - inds = torch.searchsorted(cdf, torch.rand(cdf.shape[0], n_samples, device=cdf.device)) - - return inds.view(*shape, n_samples, 1).clamp(0, support_size - 1) - - -def posenc_nerf(x: torch.Tensor, min_deg: int = 0, max_deg: int = 15) -> torch.Tensor: - """ - Concatenate x and its positional encodings, following NeRF. - - Reference: https://huggingface.co/papers/2210.04628 - """ - if min_deg == max_deg: - return x - - scales = 2.0 ** torch.arange(min_deg, max_deg, dtype=x.dtype, device=x.device) - *shape, dim = x.shape - xb = (x.reshape(-1, 1, dim) * scales.view(1, -1, 1)).reshape(*shape, -1) - assert xb.shape[-1] == dim * (max_deg - min_deg) - emb = torch.cat([xb, xb + math.pi / 2.0], axis=-1).sin() - return torch.cat([x, emb], dim=-1) - - -def encode_position(position): - return posenc_nerf(position, min_deg=0, max_deg=15) - - -def encode_direction(position, direction=None): - if direction is None: - return torch.zeros_like(posenc_nerf(position, min_deg=0, max_deg=8)) - else: - return posenc_nerf(direction, min_deg=0, max_deg=8) - - -def _sanitize_name(x: str) -> str: - return x.replace(".", "__") - - -def integrate_samples(volume_range, ts, density, channels): - r""" - Function integrating the model output. - - Args: - volume_range: Specifies the integral range [t0, t1] - ts: timesteps - density: torch.Tensor [batch_size, *shape, n_samples, 1] - channels: torch.Tensor [batch_size, *shape, n_samples, n_channels] - returns: - channels: integrated rgb output weights: torch.Tensor [batch_size, *shape, n_samples, 1] (density - *transmittance)[i] weight for each rgb output at [..., i, :]. transmittance: transmittance of this volume - ) - """ - - # 1. Calculate the weights - _, _, dt = volume_range.partition(ts) - ddensity = density * dt - - mass = torch.cumsum(ddensity, dim=-2) - transmittance = torch.exp(-mass[..., -1, :]) - - alphas = 1.0 - torch.exp(-ddensity) - Ts = torch.exp(torch.cat([torch.zeros_like(mass[..., :1, :]), -mass[..., :-1, :]], dim=-2)) - # This is the probability of light hitting and reflecting off of - # something at depth [..., i, :]. - weights = alphas * Ts - - # 2. Integrate channels - channels = torch.sum(channels * weights, dim=-2) - - return channels, weights, transmittance - - -def volume_query_points(volume, grid_size): - indices = torch.arange(grid_size**3, device=volume.bbox_min.device) - zs = indices % grid_size - ys = torch.div(indices, grid_size, rounding_mode="trunc") % grid_size - xs = torch.div(indices, grid_size**2, rounding_mode="trunc") % grid_size - combined = torch.stack([xs, ys, zs], dim=1) - return (combined.float() / (grid_size - 1)) * (volume.bbox_max - volume.bbox_min) + volume.bbox_min - - -def _convert_srgb_to_linear(u: torch.Tensor): - return torch.where(u <= 0.04045, u / 12.92, ((u + 0.055) / 1.055) ** 2.4) - - -def _create_flat_edge_indices( - flat_cube_indices: torch.Tensor, - grid_size: Tuple[int, int, int], -): - num_xs = (grid_size[0] - 1) * grid_size[1] * grid_size[2] - y_offset = num_xs - num_ys = grid_size[0] * (grid_size[1] - 1) * grid_size[2] - z_offset = num_xs + num_ys - return torch.stack( - [ - # Edges spanning x-axis. - flat_cube_indices[:, 0] * grid_size[1] * grid_size[2] - + flat_cube_indices[:, 1] * grid_size[2] - + flat_cube_indices[:, 2], - flat_cube_indices[:, 0] * grid_size[1] * grid_size[2] - + (flat_cube_indices[:, 1] + 1) * grid_size[2] - + flat_cube_indices[:, 2], - flat_cube_indices[:, 0] * grid_size[1] * grid_size[2] - + flat_cube_indices[:, 1] * grid_size[2] - + flat_cube_indices[:, 2] - + 1, - flat_cube_indices[:, 0] * grid_size[1] * grid_size[2] - + (flat_cube_indices[:, 1] + 1) * grid_size[2] - + flat_cube_indices[:, 2] - + 1, - # Edges spanning y-axis. - ( - y_offset - + flat_cube_indices[:, 0] * (grid_size[1] - 1) * grid_size[2] - + flat_cube_indices[:, 1] * grid_size[2] - + flat_cube_indices[:, 2] - ), - ( - y_offset - + (flat_cube_indices[:, 0] + 1) * (grid_size[1] - 1) * grid_size[2] - + flat_cube_indices[:, 1] * grid_size[2] - + flat_cube_indices[:, 2] - ), - ( - y_offset - + flat_cube_indices[:, 0] * (grid_size[1] - 1) * grid_size[2] - + flat_cube_indices[:, 1] * grid_size[2] - + flat_cube_indices[:, 2] - + 1 - ), - ( - y_offset - + (flat_cube_indices[:, 0] + 1) * (grid_size[1] - 1) * grid_size[2] - + flat_cube_indices[:, 1] * grid_size[2] - + flat_cube_indices[:, 2] - + 1 - ), - # Edges spanning z-axis. - ( - z_offset - + flat_cube_indices[:, 0] * grid_size[1] * (grid_size[2] - 1) - + flat_cube_indices[:, 1] * (grid_size[2] - 1) - + flat_cube_indices[:, 2] - ), - ( - z_offset - + (flat_cube_indices[:, 0] + 1) * grid_size[1] * (grid_size[2] - 1) - + flat_cube_indices[:, 1] * (grid_size[2] - 1) - + flat_cube_indices[:, 2] - ), - ( - z_offset - + flat_cube_indices[:, 0] * grid_size[1] * (grid_size[2] - 1) - + (flat_cube_indices[:, 1] + 1) * (grid_size[2] - 1) - + flat_cube_indices[:, 2] - ), - ( - z_offset - + (flat_cube_indices[:, 0] + 1) * grid_size[1] * (grid_size[2] - 1) - + (flat_cube_indices[:, 1] + 1) * (grid_size[2] - 1) - + flat_cube_indices[:, 2] - ), - ], - dim=-1, - ) - - -class VoidNeRFModel(nn.Module): - """ - Implements the default empty space model where all queries are rendered as background. - """ - - def __init__(self, background, channel_scale=255.0): - super().__init__() - background = nn.Parameter(torch.from_numpy(np.array(background)).to(dtype=torch.float32) / channel_scale) - - self.register_buffer("background", background) - - def forward(self, position): - background = self.background[None].to(position.device) - - shape = position.shape[:-1] - ones = [1] * (len(shape) - 1) - n_channels = background.shape[-1] - background = torch.broadcast_to(background.view(background.shape[0], *ones, n_channels), [*shape, n_channels]) - - return background - - -@dataclass -class VolumeRange: - t0: torch.Tensor - t1: torch.Tensor - intersected: torch.Tensor - - def __post_init__(self): - assert self.t0.shape == self.t1.shape == self.intersected.shape - - def partition(self, ts): - """ - Partitions t0 and t1 into n_samples intervals. - - Args: - ts: [batch_size, *shape, n_samples, 1] - - Return: - - lower: [batch_size, *shape, n_samples, 1] upper: [batch_size, *shape, n_samples, 1] delta: [batch_size, - *shape, n_samples, 1] - - where - ts \\in [lower, upper] deltas = upper - lower - """ - - mids = (ts[..., 1:, :] + ts[..., :-1, :]) * 0.5 - lower = torch.cat([self.t0[..., None, :], mids], dim=-2) - upper = torch.cat([mids, self.t1[..., None, :]], dim=-2) - delta = upper - lower - assert lower.shape == upper.shape == delta.shape == ts.shape - return lower, upper, delta - - -class BoundingBoxVolume(nn.Module): - """ - Axis-aligned bounding box defined by the two opposite corners. - """ - - def __init__( - self, - *, - bbox_min, - bbox_max, - min_dist: float = 0.0, - min_t_range: float = 1e-3, - ): - """ - Args: - bbox_min: the left/bottommost corner of the bounding box - bbox_max: the other corner of the bounding box - min_dist: all rays should start at least this distance away from the origin. - """ - super().__init__() - - self.min_dist = min_dist - self.min_t_range = min_t_range - - self.bbox_min = torch.tensor(bbox_min) - self.bbox_max = torch.tensor(bbox_max) - self.bbox = torch.stack([self.bbox_min, self.bbox_max]) - assert self.bbox.shape == (2, 3) - assert min_dist >= 0.0 - assert min_t_range > 0.0 - - def intersect( - self, - origin: torch.Tensor, - direction: torch.Tensor, - t0_lower: Optional[torch.Tensor] = None, - epsilon=1e-6, - ): - """ - Args: - origin: [batch_size, *shape, 3] - direction: [batch_size, *shape, 3] - t0_lower: Optional [batch_size, *shape, 1] lower bound of t0 when intersecting this volume. - params: Optional meta parameters in case Volume is parametric - epsilon: to stabilize calculations - - Return: - A tuple of (t0, t1, intersected) where each has a shape [batch_size, *shape, 1]. If a ray intersects with - the volume, `o + td` is in the volume for all t in [t0, t1]. If the volume is bounded, t1 is guaranteed to - be on the boundary of the volume. - """ - - batch_size, *shape, _ = origin.shape - ones = [1] * len(shape) - bbox = self.bbox.view(1, *ones, 2, 3).to(origin.device) - - def _safe_divide(a, b, epsilon=1e-6): - return a / torch.where(b < 0, b - epsilon, b + epsilon) - - ts = _safe_divide(bbox - origin[..., None, :], direction[..., None, :], epsilon=epsilon) - - # Cases to think about: - # - # 1. t1 <= t0: the ray does not pass through the AABB. - # 2. t0 < t1 <= 0: the ray intersects but the BB is behind the origin. - # 3. t0 <= 0 <= t1: the ray starts from inside the BB - # 4. 0 <= t0 < t1: the ray is not inside and intersects with the BB twice. - # - # 1 and 4 are clearly handled from t0 < t1 below. - # Making t0 at least min_dist (>= 0) takes care of 2 and 3. - t0 = ts.min(dim=-2).values.max(dim=-1, keepdim=True).values.clamp(self.min_dist) - t1 = ts.max(dim=-2).values.min(dim=-1, keepdim=True).values - assert t0.shape == t1.shape == (batch_size, *shape, 1) - if t0_lower is not None: - assert t0.shape == t0_lower.shape - t0 = torch.maximum(t0, t0_lower) - - intersected = t0 + self.min_t_range < t1 - t0 = torch.where(intersected, t0, torch.zeros_like(t0)) - t1 = torch.where(intersected, t1, torch.ones_like(t1)) - - return VolumeRange(t0=t0, t1=t1, intersected=intersected) - - -class StratifiedRaySampler(nn.Module): - """ - Instead of fixed intervals, a sample is drawn uniformly at random from each interval. - """ - - def __init__(self, depth_mode: str = "linear"): - """ - :param depth_mode: linear samples ts linearly in depth. harmonic ensures - closer points are sampled more densely. - """ - self.depth_mode = depth_mode - assert self.depth_mode in ("linear", "geometric", "harmonic") - - def sample( - self, - t0: torch.Tensor, - t1: torch.Tensor, - n_samples: int, - epsilon: float = 1e-3, - ) -> torch.Tensor: - """ - Args: - t0: start time has shape [batch_size, *shape, 1] - t1: finish time has shape [batch_size, *shape, 1] - n_samples: number of ts to sample - Return: - sampled ts of shape [batch_size, *shape, n_samples, 1] - """ - ones = [1] * (len(t0.shape) - 1) - ts = torch.linspace(0, 1, n_samples).view(*ones, n_samples).to(t0.dtype).to(t0.device) - - if self.depth_mode == "linear": - ts = t0 * (1.0 - ts) + t1 * ts - elif self.depth_mode == "geometric": - ts = (t0.clamp(epsilon).log() * (1.0 - ts) + t1.clamp(epsilon).log() * ts).exp() - elif self.depth_mode == "harmonic": - # The original NeRF recommends this interpolation scheme for - # spherical scenes, but there could be some weird edge cases when - # the observer crosses from the inner to outer volume. - ts = 1.0 / (1.0 / t0.clamp(epsilon) * (1.0 - ts) + 1.0 / t1.clamp(epsilon) * ts) - - mids = 0.5 * (ts[..., 1:] + ts[..., :-1]) - upper = torch.cat([mids, t1], dim=-1) - lower = torch.cat([t0, mids], dim=-1) - # yiyi notes: add a random seed here for testing, don't forget to remove - torch.manual_seed(0) - t_rand = torch.rand_like(ts) - - ts = lower + (upper - lower) * t_rand - return ts.unsqueeze(-1) - - -class ImportanceRaySampler(nn.Module): - """ - Given the initial estimate of densities, this samples more from regions/bins expected to have objects. - """ - - def __init__( - self, - volume_range: VolumeRange, - ts: torch.Tensor, - weights: torch.Tensor, - blur_pool: bool = False, - alpha: float = 1e-5, - ): - """ - Args: - volume_range: the range in which a ray intersects the given volume. - ts: earlier samples from the coarse rendering step - weights: discretized version of density * transmittance - blur_pool: if true, use 2-tap max + 2-tap blur filter from mip-NeRF. - alpha: small value to add to weights. - """ - self.volume_range = volume_range - self.ts = ts.clone().detach() - self.weights = weights.clone().detach() - self.blur_pool = blur_pool - self.alpha = alpha - - @torch.no_grad() - def sample(self, t0: torch.Tensor, t1: torch.Tensor, n_samples: int) -> torch.Tensor: - """ - Args: - t0: start time has shape [batch_size, *shape, 1] - t1: finish time has shape [batch_size, *shape, 1] - n_samples: number of ts to sample - Return: - sampled ts of shape [batch_size, *shape, n_samples, 1] - """ - lower, upper, _ = self.volume_range.partition(self.ts) - - batch_size, *shape, n_coarse_samples, _ = self.ts.shape - - weights = self.weights - if self.blur_pool: - padded = torch.cat([weights[..., :1, :], weights, weights[..., -1:, :]], dim=-2) - maxes = torch.maximum(padded[..., :-1, :], padded[..., 1:, :]) - weights = 0.5 * (maxes[..., :-1, :] + maxes[..., 1:, :]) - weights = weights + self.alpha - pmf = weights / weights.sum(dim=-2, keepdim=True) - inds = sample_pmf(pmf, n_samples) - assert inds.shape == (batch_size, *shape, n_samples, 1) - assert (inds >= 0).all() and (inds < n_coarse_samples).all() - - t_rand = torch.rand(inds.shape, device=inds.device) - lower_ = torch.gather(lower, -2, inds) - upper_ = torch.gather(upper, -2, inds) - - ts = lower_ + (upper_ - lower_) * t_rand - ts = torch.sort(ts, dim=-2).values - return ts - - -@dataclass -class MeshDecoderOutput(BaseOutput): - """ - A 3D triangle mesh with optional data at the vertices and faces. - - Args: - verts (`torch.Tensor` of shape `(N, 3)`): - array of vertext coordinates - faces (`torch.Tensor` of shape `(N, 3)`): - array of triangles, pointing to indices in verts. - vertext_channels (Dict): - vertext coordinates for each color channel - """ - - verts: torch.Tensor - faces: torch.Tensor - vertex_channels: Dict[str, torch.Tensor] - - -class MeshDecoder(nn.Module): - """ - Construct meshes from Signed distance functions (SDFs) using marching cubes method - """ - - def __init__(self): - super().__init__() - cases = torch.zeros(256, 5, 3, dtype=torch.long) - masks = torch.zeros(256, 5, dtype=torch.bool) - - self.register_buffer("cases", cases) - self.register_buffer("masks", masks) - - def forward(self, field: torch.Tensor, min_point: torch.Tensor, size: torch.Tensor): - """ - For a signed distance field, produce a mesh using marching cubes. - - :param field: a 3D tensor of field values, where negative values correspond - to the outside of the shape. The dimensions correspond to the x, y, and z directions, respectively. - :param min_point: a tensor of shape [3] containing the point corresponding - to (0, 0, 0) in the field. - :param size: a tensor of shape [3] containing the per-axis distance from the - (0, 0, 0) field corner and the (-1, -1, -1) field corner. - """ - assert len(field.shape) == 3, "input must be a 3D scalar field" - dev = field.device - - cases = self.cases.to(dev) - masks = self.masks.to(dev) - - min_point = min_point.to(dev) - size = size.to(dev) - - grid_size = field.shape - grid_size_tensor = torch.tensor(grid_size).to(size) - - # Create bitmasks between 0 and 255 (inclusive) indicating the state - # of the eight corners of each cube. - bitmasks = (field > 0).to(torch.uint8) - bitmasks = bitmasks[:-1, :, :] | (bitmasks[1:, :, :] << 1) - bitmasks = bitmasks[:, :-1, :] | (bitmasks[:, 1:, :] << 2) - bitmasks = bitmasks[:, :, :-1] | (bitmasks[:, :, 1:] << 4) - - # Compute corner coordinates across the entire grid. - corner_coords = torch.empty(*grid_size, 3, device=dev, dtype=field.dtype) - corner_coords[range(grid_size[0]), :, :, 0] = torch.arange(grid_size[0], device=dev, dtype=field.dtype)[ - :, None, None - ] - corner_coords[:, range(grid_size[1]), :, 1] = torch.arange(grid_size[1], device=dev, dtype=field.dtype)[ - :, None - ] - corner_coords[:, :, range(grid_size[2]), 2] = torch.arange(grid_size[2], device=dev, dtype=field.dtype) - - # Compute all vertices across all edges in the grid, even though we will - # throw some out later. We have (X-1)*Y*Z + X*(Y-1)*Z + X*Y*(Z-1) vertices. - # These are all midpoints, and don't account for interpolation (which is - # done later based on the used edge midpoints). - edge_midpoints = torch.cat( - [ - ((corner_coords[:-1] + corner_coords[1:]) / 2).reshape(-1, 3), - ((corner_coords[:, :-1] + corner_coords[:, 1:]) / 2).reshape(-1, 3), - ((corner_coords[:, :, :-1] + corner_coords[:, :, 1:]) / 2).reshape(-1, 3), - ], - dim=0, - ) - - # Create a flat array of [X, Y, Z] indices for each cube. - cube_indices = torch.zeros( - grid_size[0] - 1, grid_size[1] - 1, grid_size[2] - 1, 3, device=dev, dtype=torch.long - ) - cube_indices[range(grid_size[0] - 1), :, :, 0] = torch.arange(grid_size[0] - 1, device=dev)[:, None, None] - cube_indices[:, range(grid_size[1] - 1), :, 1] = torch.arange(grid_size[1] - 1, device=dev)[:, None] - cube_indices[:, :, range(grid_size[2] - 1), 2] = torch.arange(grid_size[2] - 1, device=dev) - flat_cube_indices = cube_indices.reshape(-1, 3) - - # Create a flat array mapping each cube to 12 global edge indices. - edge_indices = _create_flat_edge_indices(flat_cube_indices, grid_size) - - # Apply the LUT to figure out the triangles. - flat_bitmasks = bitmasks.reshape(-1).long() # must cast to long for indexing to believe this not a mask - local_tris = cases[flat_bitmasks] - local_masks = masks[flat_bitmasks] - # Compute the global edge indices for the triangles. - global_tris = torch.gather(edge_indices, 1, local_tris.reshape(local_tris.shape[0], -1)).reshape( - local_tris.shape - ) - # Select the used triangles for each cube. - selected_tris = global_tris.reshape(-1, 3)[local_masks.reshape(-1)] - - # Now we have a bunch of indices into the full list of possible vertices, - # but we want to reduce this list to only the used vertices. - used_vertex_indices = torch.unique(selected_tris.view(-1)) - used_edge_midpoints = edge_midpoints[used_vertex_indices] - old_index_to_new_index = torch.zeros(len(edge_midpoints), device=dev, dtype=torch.long) - old_index_to_new_index[used_vertex_indices] = torch.arange( - len(used_vertex_indices), device=dev, dtype=torch.long - ) - - # Rewrite the triangles to use the new indices - faces = torch.gather(old_index_to_new_index, 0, selected_tris.view(-1)).reshape(selected_tris.shape) - - # Compute the actual interpolated coordinates corresponding to edge midpoints. - v1 = torch.floor(used_edge_midpoints).to(torch.long) - v2 = torch.ceil(used_edge_midpoints).to(torch.long) - s1 = field[v1[:, 0], v1[:, 1], v1[:, 2]] - s2 = field[v2[:, 0], v2[:, 1], v2[:, 2]] - p1 = (v1.float() / (grid_size_tensor - 1)) * size + min_point - p2 = (v2.float() / (grid_size_tensor - 1)) * size + min_point - # The signs of s1 and s2 should be different. We want to find - # t such that t*s2 + (1-t)*s1 = 0. - t = (s1 / (s1 - s2))[:, None] - verts = t * p2 + (1 - t) * p1 - - return MeshDecoderOutput(verts=verts, faces=faces, vertex_channels=None) - - -@dataclass -class MLPNeRFModelOutput(BaseOutput): - density: torch.Tensor - signed_distance: torch.Tensor - channels: torch.Tensor - ts: torch.Tensor - - -class MLPNeRSTFModel(ModelMixin, ConfigMixin): - @register_to_config - def __init__( - self, - d_hidden: int = 256, - n_output: int = 12, - n_hidden_layers: int = 6, - act_fn: str = "swish", - insert_direction_at: int = 4, - ): - super().__init__() - - # Instantiate the MLP - - # Find out the dimension of encoded position and direction - dummy = torch.eye(1, 3) - d_posenc_pos = encode_position(position=dummy).shape[-1] - d_posenc_dir = encode_direction(position=dummy).shape[-1] - - mlp_widths = [d_hidden] * n_hidden_layers - input_widths = [d_posenc_pos] + mlp_widths - output_widths = mlp_widths + [n_output] - - if insert_direction_at is not None: - input_widths[insert_direction_at] += d_posenc_dir - - self.mlp = nn.ModuleList([nn.Linear(d_in, d_out) for d_in, d_out in zip(input_widths, output_widths)]) - - if act_fn == "swish": - # self.activation = swish - # yiyi testing: - self.activation = lambda x: F.silu(x) - else: - raise ValueError(f"Unsupported activation function {act_fn}") - - self.sdf_activation = torch.tanh - self.density_activation = torch.nn.functional.relu - self.channel_activation = torch.sigmoid - - def map_indices_to_keys(self, output): - h_map = { - "sdf": (0, 1), - "density_coarse": (1, 2), - "density_fine": (2, 3), - "stf": (3, 6), - "nerf_coarse": (6, 9), - "nerf_fine": (9, 12), - } - - mapped_output = {k: output[..., start:end] for k, (start, end) in h_map.items()} - - return mapped_output - - def forward(self, *, position, direction, ts, nerf_level="coarse", rendering_mode="nerf"): - h = encode_position(position) - - h_preact = h - h_directionless = None - for i, layer in enumerate(self.mlp): - if i == self.config.insert_direction_at: # 4 in the config - h_directionless = h_preact - h_direction = encode_direction(position, direction=direction) - h = torch.cat([h, h_direction], dim=-1) - - h = layer(h) - - h_preact = h - - if i < len(self.mlp) - 1: - h = self.activation(h) - - h_final = h - if h_directionless is None: - h_directionless = h_preact - - activation = self.map_indices_to_keys(h_final) - - if nerf_level == "coarse": - h_density = activation["density_coarse"] - else: - h_density = activation["density_fine"] - - if rendering_mode == "nerf": - if nerf_level == "coarse": - h_channels = activation["nerf_coarse"] - else: - h_channels = activation["nerf_fine"] - - elif rendering_mode == "stf": - h_channels = activation["stf"] - - density = self.density_activation(h_density) - signed_distance = self.sdf_activation(activation["sdf"]) - channels = self.channel_activation(h_channels) - - # yiyi notes: I think signed_distance is not used - return MLPNeRFModelOutput(density=density, signed_distance=signed_distance, channels=channels, ts=ts) - - -class ChannelsProj(nn.Module): - def __init__( - self, - *, - vectors: int, - channels: int, - d_latent: int, - ): - super().__init__() - self.proj = nn.Linear(d_latent, vectors * channels) - self.norm = nn.LayerNorm(channels) - self.d_latent = d_latent - self.vectors = vectors - self.channels = channels - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x_bvd = x - w_vcd = self.proj.weight.view(self.vectors, self.channels, self.d_latent) - b_vc = self.proj.bias.view(1, self.vectors, self.channels) - h = torch.einsum("bvd,vcd->bvc", x_bvd, w_vcd) - h = self.norm(h) - - h = h + b_vc - return h - - -class ShapEParamsProjModel(ModelMixin, ConfigMixin): - """ - project the latent representation of a 3D asset to obtain weights of a multi-layer perceptron (MLP). - - For more details, see the original paper: - """ - - @register_to_config - def __init__( - self, - *, - 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]] = ( - (256, 93), - (256, 256), - (256, 256), - (256, 256), - ), - d_latent: int = 1024, - ): - super().__init__() - - # check inputs - if len(param_names) != len(param_shapes): - raise ValueError("Must provide same number of `param_names` as `param_shapes`") - self.projections = nn.ModuleDict({}) - for k, (vectors, channels) in zip(param_names, param_shapes): - self.projections[_sanitize_name(k)] = ChannelsProj( - vectors=vectors, - channels=channels, - d_latent=d_latent, - ) - - def forward(self, x: torch.Tensor): - out = {} - start = 0 - for k, shape in zip(self.config.param_names, self.config.param_shapes): - vectors, _ = shape - end = start + vectors - x_bvd = x[:, start:end] - out[k] = self.projections[_sanitize_name(k)](x_bvd).reshape(len(x), *shape) - start = end - return out - - -class ShapERenderer(ModelMixin, ConfigMixin): - @register_to_config - def __init__( - self, - *, - 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]] = ( - (256, 93), - (256, 256), - (256, 256), - (256, 256), - ), - d_latent: int = 1024, - d_hidden: int = 256, - n_output: int = 12, - n_hidden_layers: int = 6, - act_fn: str = "swish", - insert_direction_at: int = 4, - background: Tuple[float] = ( - 255.0, - 255.0, - 255.0, - ), - ): - super().__init__() - - self.params_proj = ShapEParamsProjModel( - param_names=param_names, - param_shapes=param_shapes, - d_latent=d_latent, - ) - self.mlp = MLPNeRSTFModel(d_hidden, n_output, n_hidden_layers, act_fn, insert_direction_at) - self.void = VoidNeRFModel(background=background, channel_scale=255.0) - self.volume = BoundingBoxVolume(bbox_max=[1.0, 1.0, 1.0], bbox_min=[-1.0, -1.0, -1.0]) - self.mesh_decoder = MeshDecoder() - - @torch.no_grad() - def render_rays(self, rays, sampler, n_samples, prev_model_out=None, render_with_direction=False): - """ - Perform volumetric rendering over a partition of possible t's in the union of rendering volumes (written below - with some abuse of notations) - - C(r) := sum( - transmittance(t[i]) * integrate( - lambda t: density(t) * channels(t) * transmittance(t), [t[i], t[i + 1]], - ) for i in range(len(parts)) - ) + transmittance(t[-1]) * void_model(t[-1]).channels - - where - - 1) transmittance(s) := exp(-integrate(density, [t[0], s])) calculates the probability of light passing through - the volume specified by [t[0], s]. (transmittance of 1 means light can pass freely) 2) density and channels are - obtained by evaluating the appropriate part.model at time t. 3) [t[i], t[i + 1]] is defined as the range of t - where the ray intersects (parts[i].volume \\ union(part.volume for part in parts[:i])) at the surface of the - shell (if bounded). If the ray does not intersect, the integral over this segment is evaluated as 0 and - transmittance(t[i + 1]) := transmittance(t[i]). 4) The last term is integration to infinity (e.g. [t[-1], - math.inf]) that is evaluated by the void_model (i.e. we consider this space to be empty). - - Args: - rays: [batch_size x ... x 2 x 3] origin and direction. sampler: disjoint volume integrals. n_samples: - number of ts to sample. prev_model_outputs: model outputs from the previous rendering step, including - - :return: A tuple of - - `channels` - - A importance samplers for additional fine-grained rendering - - raw model output - """ - origin, direction = rays[..., 0, :], rays[..., 1, :] - - # Integrate over [t[i], t[i + 1]] - - # 1 Intersect the rays with the current volume and sample ts to integrate along. - vrange = self.volume.intersect(origin, direction, t0_lower=None) - ts = sampler.sample(vrange.t0, vrange.t1, n_samples) - ts = ts.to(rays.dtype) - - if prev_model_out is not None: - # Append the previous ts now before fprop because previous - # rendering used a different model and we can't reuse the output. - ts = torch.sort(torch.cat([ts, prev_model_out.ts], dim=-2), dim=-2).values - - batch_size, *_shape, _t0_dim = vrange.t0.shape - _, *ts_shape, _ts_dim = ts.shape - - # 2. Get the points along the ray and query the model - directions = torch.broadcast_to(direction.unsqueeze(-2), [batch_size, *ts_shape, 3]) - positions = origin.unsqueeze(-2) + ts * directions - - directions = directions.to(self.mlp.dtype) - positions = positions.to(self.mlp.dtype) - - optional_directions = directions if render_with_direction else None - - model_out = self.mlp( - position=positions, - direction=optional_directions, - ts=ts, - nerf_level="coarse" if prev_model_out is None else "fine", - ) - - # 3. Integrate the model results - channels, weights, transmittance = integrate_samples( - vrange, model_out.ts, model_out.density, model_out.channels - ) - - # 4. Clean up results that do not intersect with the volume. - transmittance = torch.where(vrange.intersected, transmittance, torch.ones_like(transmittance)) - channels = torch.where(vrange.intersected, channels, torch.zeros_like(channels)) - # 5. integration to infinity (e.g. [t[-1], math.inf]) that is evaluated by the void_model (i.e. we consider this space to be empty). - channels = channels + transmittance * self.void(origin) - - weighted_sampler = ImportanceRaySampler(vrange, ts=model_out.ts, weights=weights) - - return channels, weighted_sampler, model_out - - @torch.no_grad() - def decode_to_image( - self, - latents, - device, - size: int = 64, - ray_batch_size: int = 4096, - n_coarse_samples=64, - n_fine_samples=128, - ): - # project the parameters from the generated latents - projected_params = self.params_proj(latents) - - # update the mlp layers of the renderer - for name, param in self.mlp.state_dict().items(): - if f"nerstf.{name}" in projected_params.keys(): - param.copy_(projected_params[f"nerstf.{name}"].squeeze(0)) - - # create cameras object - camera = create_pan_cameras(size) - rays = camera.camera_rays - rays = rays.to(device) - n_batches = rays.shape[1] // ray_batch_size - - coarse_sampler = StratifiedRaySampler() - - images = [] - - for idx in range(n_batches): - rays_batch = rays[:, idx * ray_batch_size : (idx + 1) * ray_batch_size] - - # render rays with coarse, stratified samples. - _, fine_sampler, coarse_model_out = self.render_rays(rays_batch, coarse_sampler, n_coarse_samples) - # Then, render with additional importance-weighted ray samples. - channels, _, _ = self.render_rays( - rays_batch, fine_sampler, n_fine_samples, prev_model_out=coarse_model_out - ) - - images.append(channels) - - images = torch.cat(images, dim=1) - images = images.view(*camera.shape, camera.height, camera.width, -1).squeeze(0) - - return images - - @torch.no_grad() - def decode_to_mesh( - self, - latents, - device, - grid_size: int = 128, - query_batch_size: int = 4096, - texture_channels: Tuple = ("R", "G", "B"), - ): - # 1. project the parameters from the generated latents - projected_params = self.params_proj(latents) - - # 2. update the mlp layers of the renderer - for name, param in self.mlp.state_dict().items(): - if f"nerstf.{name}" in projected_params.keys(): - param.copy_(projected_params[f"nerstf.{name}"].squeeze(0)) - - # 3. decoding with STF rendering - # 3.1 query the SDF values at vertices along a regular 128**3 grid - - query_points = volume_query_points(self.volume, grid_size) - query_positions = query_points[None].repeat(1, 1, 1).to(device=device, dtype=self.mlp.dtype) - - fields = [] - - for idx in range(0, query_positions.shape[1], query_batch_size): - query_batch = query_positions[:, idx : idx + query_batch_size] - - model_out = self.mlp( - position=query_batch, direction=None, ts=None, nerf_level="fine", rendering_mode="stf" - ) - fields.append(model_out.signed_distance) - - # predicted SDF values - fields = torch.cat(fields, dim=1) - fields = fields.float() - - assert len(fields.shape) == 3 and fields.shape[-1] == 1, ( - f"expected [meta_batch x inner_batch] SDF results, but got {fields.shape}" - ) - - fields = fields.reshape(1, *([grid_size] * 3)) - - # create grid 128 x 128 x 128 - # - force a negative border around the SDFs to close off all the models. - full_grid = torch.zeros( - 1, - grid_size + 2, - grid_size + 2, - grid_size + 2, - device=fields.device, - dtype=fields.dtype, - ) - full_grid.fill_(-1.0) - full_grid[:, 1:-1, 1:-1, 1:-1] = fields - fields = full_grid - - # apply a differentiable implementation of Marching Cubes to construct meshs - raw_meshes = [] - mesh_mask = [] - - for field in fields: - raw_mesh = self.mesh_decoder(field, self.volume.bbox_min, self.volume.bbox_max - self.volume.bbox_min) - mesh_mask.append(True) - raw_meshes.append(raw_mesh) - - mesh_mask = torch.tensor(mesh_mask, device=fields.device) - max_vertices = max(len(m.verts) for m in raw_meshes) - - # 3.2. query the texture color head at each vertex of the resulting mesh. - texture_query_positions = torch.stack( - [m.verts[torch.arange(0, max_vertices) % len(m.verts)] for m in raw_meshes], - dim=0, - ) - texture_query_positions = texture_query_positions.to(device=device, dtype=self.mlp.dtype) - - textures = [] - - for idx in range(0, texture_query_positions.shape[1], query_batch_size): - query_batch = texture_query_positions[:, idx : idx + query_batch_size] - - texture_model_out = self.mlp( - position=query_batch, direction=None, ts=None, nerf_level="fine", rendering_mode="stf" - ) - textures.append(texture_model_out.channels) - - # predict texture color - textures = torch.cat(textures, dim=1) - - textures = _convert_srgb_to_linear(textures) - textures = textures.float() - - # 3.3 augment the mesh with texture data - assert len(textures.shape) == 3 and textures.shape[-1] == len(texture_channels), ( - f"expected [meta_batch x inner_batch x texture_channels] field results, but got {textures.shape}" - ) - - for m, texture in zip(raw_meshes, textures): - texture = texture[: len(m.verts)] - m.vertex_channels = dict(zip(texture_channels, texture.unbind(-1))) - - return raw_meshes[0] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/__init__.py deleted file mode 100644 index dfdd419ae9914e64c7fdcf7c152ac308b75d75d2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/__init__.py +++ /dev/null @@ -1,50 +0,0 @@ -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 = {} -_import_structure = {} - -try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_torch_and_transformers_objects - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) -else: - _import_structure["modeling_stable_audio"] = ["StableAudioProjectionModel"] - _import_structure["pipeline_stable_audio"] = ["StableAudioPipeline"] - - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_objects import * - - else: - from .modeling_stable_audio import StableAudioProjectionModel - from .pipeline_stable_audio import StableAudioPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index d8713587792e14a044dafeb77e5792e14252fe4d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/__pycache__/modeling_stable_audio.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/__pycache__/modeling_stable_audio.cpython-310.pyc deleted file mode 100644 index aa14c134bcf626f989934b9da5bb5144802a3ad2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/__pycache__/modeling_stable_audio.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/__pycache__/pipeline_stable_audio.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/__pycache__/pipeline_stable_audio.cpython-310.pyc deleted file mode 100644 index 48a17a5462eb89a5ebdbcc84d294f23bdc051f7f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/__pycache__/pipeline_stable_audio.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/modeling_stable_audio.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/modeling_stable_audio.py deleted file mode 100644 index 89d4d2dca522deefaab10a585bfb8ee53b7cec8d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/modeling_stable_audio.py +++ /dev/null @@ -1,158 +0,0 @@ -# Copyright 2025 Stability 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. - -from dataclasses import dataclass -from math import pi -from typing import Optional - -import torch -import torch.nn as nn -import torch.utils.checkpoint - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.modeling_utils import ModelMixin -from ...utils import BaseOutput, logging - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class StableAudioPositionalEmbedding(nn.Module): - """Used for continuous time""" - - def __init__(self, dim: int): - super().__init__() - assert (dim % 2) == 0 - half_dim = dim // 2 - self.weights = nn.Parameter(torch.randn(half_dim)) - - def forward(self, times: torch.Tensor) -> torch.Tensor: - times = times[..., None] - freqs = times * self.weights[None] * 2 * pi - fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1) - fouriered = torch.cat((times, fouriered), dim=-1) - return fouriered - - -@dataclass -class StableAudioProjectionModelOutput(BaseOutput): - """ - Args: - Class for StableAudio projection layer's outputs. - text_hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): - Sequence of hidden-states obtained by linearly projecting the hidden-states for the text encoder. - seconds_start_hidden_states (`torch.Tensor` of shape `(batch_size, 1, hidden_size)`, *optional*): - Sequence of hidden-states obtained by linearly projecting the audio start hidden states. - seconds_end_hidden_states (`torch.Tensor` of shape `(batch_size, 1, hidden_size)`, *optional*): - 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 - - -class StableAudioNumberConditioner(nn.Module): - """ - A simple linear projection model to map numbers to a latent space. - - Args: - number_embedding_dim (`int`): - Dimensionality of the number embeddings. - min_value (`int`): - The minimum value of the seconds number conditioning modules. - max_value (`int`): - The maximum value of the seconds number conditioning modules - internal_dim (`int`): - Dimensionality of the intermediate number hidden states. - """ - - def __init__( - self, - number_embedding_dim, - min_value, - max_value, - internal_dim: Optional[int] = 256, - ): - super().__init__() - self.time_positional_embedding = nn.Sequential( - StableAudioPositionalEmbedding(internal_dim), - nn.Linear(in_features=internal_dim + 1, out_features=number_embedding_dim), - ) - - self.number_embedding_dim = number_embedding_dim - self.min_value = min_value - self.max_value = max_value - - def forward( - self, - floats: torch.Tensor, - ): - floats = floats.clamp(self.min_value, self.max_value) - - normalized_floats = (floats - self.min_value) / (self.max_value - self.min_value) - - # Cast floats to same type as embedder - embedder_dtype = next(self.time_positional_embedding.parameters()).dtype - normalized_floats = normalized_floats.to(embedder_dtype) - - embedding = self.time_positional_embedding(normalized_floats) - float_embeds = embedding.view(-1, 1, self.number_embedding_dim) - - return float_embeds - - -class StableAudioProjectionModel(ModelMixin, ConfigMixin): - """ - A simple linear projection model to map the conditioning values to a shared latent space. - - Args: - text_encoder_dim (`int`): - Dimensionality of the text embeddings from the text encoder (T5). - conditioning_dim (`int`): - Dimensionality of the output conditioning tensors. - min_value (`int`): - The minimum value of the seconds number conditioning modules. - max_value (`int`): - The maximum value of the seconds number conditioning modules - """ - - @register_to_config - def __init__(self, text_encoder_dim, conditioning_dim, min_value, max_value): - super().__init__() - self.text_projection = ( - nn.Identity() if conditioning_dim == text_encoder_dim else nn.Linear(text_encoder_dim, conditioning_dim) - ) - self.start_number_conditioner = StableAudioNumberConditioner(conditioning_dim, min_value, max_value) - self.end_number_conditioner = StableAudioNumberConditioner(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 = ( - text_hidden_states if text_hidden_states is None else self.text_projection(text_hidden_states) - ) - seconds_start_hidden_states = ( - start_seconds if start_seconds is None else self.start_number_conditioner(start_seconds) - ) - seconds_end_hidden_states = end_seconds if end_seconds is None else self.end_number_conditioner(end_seconds) - - return StableAudioProjectionModelOutput( - text_hidden_states=text_hidden_states, - seconds_start_hidden_states=seconds_start_hidden_states, - seconds_end_hidden_states=seconds_end_hidden_states, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/pipeline_stable_audio.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/pipeline_stable_audio.py deleted file mode 100644 index 8861ecae7de15bb3854a39b656763f108619ff68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_audio/pipeline_stable_audio.py +++ /dev/null @@ -1,756 +0,0 @@ -# Copyright 2025 Stability 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 inspect -from typing import Callable, List, Optional, Union - -import torch -from transformers import ( - T5EncoderModel, - T5Tokenizer, - T5TokenizerFast, -) - -from ...models import AutoencoderOobleck, StableAudioDiTModel -from ...models.embeddings import get_1d_rotary_pos_embed -from ...schedulers import EDMDPMSolverMultistepScheduler -from ...utils import ( - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline -from .modeling_stable_audio import StableAudioProjectionModel - - -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 scipy - >>> import torch - >>> import soundfile as sf - >>> from diffusers import StableAudioPipeline - - >>> repo_id = "stabilityai/stable-audio-open-1.0" - >>> pipe = StableAudioPipeline.from_pretrained(repo_id, torch_dtype=torch.float16) - >>> pipe = pipe.to("cuda") - - >>> # define the prompts - >>> prompt = "The sound of a hammer hitting a wooden surface." - >>> negative_prompt = "Low quality." - - >>> # set the seed for generator - >>> generator = torch.Generator("cuda").manual_seed(0) - - >>> # run the generation - >>> audio = pipe( - ... prompt, - ... negative_prompt=negative_prompt, - ... num_inference_steps=200, - ... audio_end_in_s=10.0, - ... num_waveforms_per_prompt=3, - ... generator=generator, - ... ).audios - - >>> output = audio[0].T.float().cpu().numpy() - >>> sf.write("hammer.wav", output, pipe.vae.sampling_rate) - ``` -""" - - -class StableAudioPipeline(DiffusionPipeline): - r""" - Pipeline for text-to-audio generation using StableAudio. - - 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: - vae ([`AutoencoderOobleck`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.T5EncoderModel`]): - Frozen text-encoder. StableAudio uses the encoder of - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the - [google-t5/t5-base](https://huggingface.co/google-t5/t5-base) variant. - projection_model ([`StableAudioProjectionModel`]): - A trained model used to linearly project the hidden-states from the text encoder model and the start and - end seconds. The projected hidden-states from the encoder and the conditional seconds are concatenated to - give the input to the transformer model. - tokenizer ([`~transformers.T5Tokenizer`]): - Tokenizer to tokenize text for the frozen text-encoder. - transformer ([`StableAudioDiTModel`]): - A `StableAudioDiTModel` to denoise the encoded audio latents. - scheduler ([`EDMDPMSolverMultistepScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded audio latents. - """ - - model_cpu_offload_seq = "text_encoder->projection_model->transformer->vae" - - def __init__( - self, - vae: AutoencoderOobleck, - text_encoder: T5EncoderModel, - projection_model: StableAudioProjectionModel, - tokenizer: Union[T5Tokenizer, T5TokenizerFast], - transformer: StableAudioDiTModel, - scheduler: EDMDPMSolverMultistepScheduler, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - projection_model=projection_model, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - ) - self.rotary_embed_dim = self.transformer.config.attention_head_dim // 2 - - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.enable_vae_slicing - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.disable_vae_slicing - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def encode_prompt( - self, - 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, - ): - 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: - # 1. Tokenize text - 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 - attention_mask = text_inputs.attention_mask - 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( - f"The following part of your input was truncated because {self.text_encoder.config.model_type} can " - f"only handle sequences up to {self.tokenizer.model_max_length} tokens: {removed_text}" - ) - - text_input_ids = text_input_ids.to(device) - attention_mask = attention_mask.to(device) - - # 2. Text encoder forward - self.text_encoder.eval() - prompt_embeds = self.text_encoder( - text_input_ids, - attention_mask=attention_mask, - ) - prompt_embeds = prompt_embeds[0] - - if do_classifier_free_guidance and negative_prompt is not None: - 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)} !=" - 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 - - # 1. Tokenize text - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - - uncond_input_ids = uncond_input.input_ids.to(device) - negative_attention_mask = uncond_input.attention_mask.to(device) - - # 2. Text encoder forward - self.text_encoder.eval() - negative_prompt_embeds = self.text_encoder( - uncond_input_ids, - attention_mask=negative_attention_mask, - ) - negative_prompt_embeds = negative_prompt_embeds[0] - - if negative_attention_mask is not None: - # set the masked tokens to the null embed - negative_prompt_embeds = torch.where( - negative_attention_mask.to(torch.bool).unsqueeze(2), negative_prompt_embeds, 0.0 - ) - - # 3. Project prompt_embeds and negative_prompt_embeds - if do_classifier_free_guidance and negative_prompt_embeds is not None: - # For classifier free guidance, we need to do two forward passes. - # Here we concatenate the negative and text embeddings into a single batch - # to avoid doing two forward passes - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - if attention_mask is not None and negative_attention_mask is None: - negative_attention_mask = torch.ones_like(attention_mask) - elif attention_mask is None and negative_attention_mask is not None: - attention_mask = torch.ones_like(negative_attention_mask) - - if attention_mask is not None: - attention_mask = torch.cat([negative_attention_mask, attention_mask]) - - prompt_embeds = self.projection_model( - text_hidden_states=prompt_embeds, - ).text_hidden_states - if attention_mask is not None: - prompt_embeds = prompt_embeds * attention_mask.unsqueeze(-1).to(prompt_embeds.dtype) - prompt_embeds = prompt_embeds * attention_mask.unsqueeze(-1).to(prompt_embeds.dtype) - - return prompt_embeds - - def encode_duration( - self, - audio_start_in_s, - audio_end_in_s, - device, - do_classifier_free_guidance, - batch_size, - ): - audio_start_in_s = audio_start_in_s if isinstance(audio_start_in_s, list) else [audio_start_in_s] - audio_end_in_s = audio_end_in_s if isinstance(audio_end_in_s, list) else [audio_end_in_s] - - if len(audio_start_in_s) == 1: - audio_start_in_s = audio_start_in_s * batch_size - if len(audio_end_in_s) == 1: - audio_end_in_s = audio_end_in_s * batch_size - - # Cast the inputs to floats - audio_start_in_s = [float(x) for x in audio_start_in_s] - audio_start_in_s = torch.tensor(audio_start_in_s).to(device) - - audio_end_in_s = [float(x) for x in audio_end_in_s] - audio_end_in_s = torch.tensor(audio_end_in_s).to(device) - - projection_output = self.projection_model( - start_seconds=audio_start_in_s, - end_seconds=audio_end_in_s, - ) - seconds_start_hidden_states = projection_output.seconds_start_hidden_states - seconds_end_hidden_states = projection_output.seconds_end_hidden_states - - # For classifier free guidance, we need to do two forward passes. - # Here we repeat the audio hidden states to avoid doing two forward passes - if do_classifier_free_guidance: - seconds_start_hidden_states = torch.cat([seconds_start_hidden_states, seconds_start_hidden_states], dim=0) - seconds_end_hidden_states = torch.cat([seconds_end_hidden_states, seconds_end_hidden_states], dim=0) - - return seconds_start_hidden_states, seconds_end_hidden_states - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - audio_start_in_s, - audio_end_in_s, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - attention_mask=None, - negative_attention_mask=None, - initial_audio_waveforms=None, - initial_audio_sampling_rate=None, - ): - if audio_end_in_s < audio_start_in_s: - raise ValueError( - f"`audio_end_in_s={audio_end_in_s}' must be higher than 'audio_start_in_s={audio_start_in_s}` but " - ) - - if ( - audio_start_in_s < self.projection_model.config.min_value - or audio_start_in_s > self.projection_model.config.max_value - ): - raise ValueError( - f"`audio_start_in_s` must be greater than or equal to {self.projection_model.config.min_value}, and lower than or equal to {self.projection_model.config.max_value} but " - f"is {audio_start_in_s}." - ) - - if ( - audio_end_in_s < self.projection_model.config.min_value - or audio_end_in_s > self.projection_model.config.max_value - ): - raise ValueError( - f"`audio_end_in_s` must be greater than or equal to {self.projection_model.config.min_value}, and lower than or equal to {self.projection_model.config.max_value} but " - f"is {audio_end_in_s}." - ) - - if (callback_steps is None) or ( - 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 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" - "`prompt` undefined without specifying `prompt_embeds`." - ) - 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 attention_mask is not None and attention_mask.shape != prompt_embeds.shape[:2]: - raise ValueError( - "`attention_mask should have the same batch size and sequence length as `prompt_embeds`, but got:" - f"`attention_mask: {attention_mask.shape} != `prompt_embeds` {prompt_embeds.shape}" - ) - - if initial_audio_sampling_rate is None and initial_audio_waveforms is not None: - raise ValueError( - "`initial_audio_waveforms' is provided but the sampling rate is not. Make sure to pass `initial_audio_sampling_rate`." - ) - - if initial_audio_sampling_rate is not None and initial_audio_sampling_rate != self.vae.sampling_rate: - raise ValueError( - f"`initial_audio_sampling_rate` must be {self.vae.hop_length}' but is `{initial_audio_sampling_rate}`." - "Make sure to resample the `initial_audio_waveforms` and to correct the sampling rate. " - ) - - def prepare_latents( - self, - batch_size, - num_channels_vae, - sample_size, - dtype, - device, - generator, - latents=None, - initial_audio_waveforms=None, - num_waveforms_per_prompt=None, - audio_channels=None, - ): - shape = (batch_size, num_channels_vae, sample_size) - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - - # encode the initial audio for use by the model - if initial_audio_waveforms is not None: - # check dimension - if initial_audio_waveforms.ndim == 2: - initial_audio_waveforms = initial_audio_waveforms.unsqueeze(1) - elif initial_audio_waveforms.ndim != 3: - raise ValueError( - f"`initial_audio_waveforms` must be of shape `(batch_size, num_channels, audio_length)` or `(batch_size, audio_length)` but has `{initial_audio_waveforms.ndim}` dimensions" - ) - - audio_vae_length = int(self.transformer.config.sample_size) * self.vae.hop_length - audio_shape = (batch_size // num_waveforms_per_prompt, audio_channels, audio_vae_length) - - # check num_channels - if initial_audio_waveforms.shape[1] == 1 and audio_channels == 2: - initial_audio_waveforms = initial_audio_waveforms.repeat(1, 2, 1) - elif initial_audio_waveforms.shape[1] == 2 and audio_channels == 1: - initial_audio_waveforms = initial_audio_waveforms.mean(1, keepdim=True) - - if initial_audio_waveforms.shape[:2] != audio_shape[:2]: - raise ValueError( - f"`initial_audio_waveforms` must be of shape `(batch_size, num_channels, audio_length)` or `(batch_size, audio_length)` but is of shape `{initial_audio_waveforms.shape}`" - ) - - # crop or pad - audio_length = initial_audio_waveforms.shape[-1] - if audio_length < audio_vae_length: - logger.warning( - f"The provided input waveform is shorter ({audio_length}) than the required audio length ({audio_vae_length}) of the model and will thus be padded." - ) - elif audio_length > audio_vae_length: - logger.warning( - f"The provided input waveform is longer ({audio_length}) than the required audio length ({audio_vae_length}) of the model and will thus be cropped." - ) - - audio = initial_audio_waveforms.new_zeros(audio_shape) - audio[:, :, : min(audio_length, audio_vae_length)] = initial_audio_waveforms[:, :, :audio_vae_length] - - encoded_audio = self.vae.encode(audio).latent_dist.sample(generator) - encoded_audio = encoded_audio.repeat((num_waveforms_per_prompt, 1, 1)) - latents = encoded_audio + latents - return latents - - @torch.no_grad() - @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, - 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, - 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, - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: Optional[int] = 1, - output_type: Optional[str] = "pt", - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - audio_start_in_s (`float`, *optional*, defaults to 0): - Audio start index in seconds. - num_inference_steps (`int`, *optional*, defaults to 100): - The number of denoising steps. More denoising steps usually lead to a higher quality audio at the - expense of slower inference. - 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*): - 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): - The number of waveforms to generate per prompt. - 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*): - 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 audio - 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`. - initial_audio_waveforms (`torch.Tensor`, *optional*): - Optional initial audio waveforms to use as the initial audio waveform for generation. Must be of shape - `(batch_size, num_channels, audio_length)` or `(batch_size, audio_length)`, where `batch_size` - corresponds to the number of prompts passed to the model. - initial_audio_sampling_rate (`int`, *optional*): - Sampling rate of the `initial_audio_waveforms`, if they are provided. Must be the same as the model. - prompt_embeds (`torch.Tensor`, *optional*): - Pre-computed text embeddings from the text encoder model. Can be used to easily tweak text inputs, - *e.g.* prompt weighting. If not provided, text embeddings will be computed from `prompt` input - argument. - negative_prompt_embeds (`torch.Tensor`, *optional*): - Pre-computed negative text embeddings from the text encoder model. Can be used to easily tweak text - inputs, *e.g.* prompt weighting. If not provided, negative_prompt_embeds will be computed from - `negative_prompt` input argument. - attention_mask (`torch.LongTensor`, *optional*): - Pre-computed attention mask to be applied to the `prompt_embeds`. If not provided, attention mask will - be computed from `prompt` input argument. - negative_attention_mask (`torch.LongTensor`, *optional*): - Pre-computed attention mask to be applied to the `negative_text_audio_duration_embeds`. - 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 calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - output_type (`str`, *optional*, defaults to `"pt"`): - The output format of the generated audio. Choose between `"np"` to return a NumPy `np.ndarray` or - `"pt"` to return a PyTorch `torch.Tensor` object. Set to `"latent"` to return the latent diffusion - model (LDM) output. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned, - otherwise a `tuple` is returned where the first element is a list with the generated audio. - """ - # 0. Convert audio input length from seconds to latent length - downsample_ratio = self.vae.hop_length - - max_audio_length_in_s = self.transformer.config.sample_size * downsample_ratio / self.vae.config.sampling_rate - if audio_end_in_s is None: - audio_end_in_s = max_audio_length_in_s - - if audio_end_in_s - audio_start_in_s > max_audio_length_in_s: - raise ValueError( - f"The total audio length requested ({audio_end_in_s - audio_start_in_s}s) is longer than the model maximum possible length ({max_audio_length_in_s}). Make sure that 'audio_end_in_s-audio_start_in_s<={max_audio_length_in_s}'." - ) - - waveform_start = int(audio_start_in_s * self.vae.config.sampling_rate) - waveform_end = int(audio_end_in_s * self.vae.config.sampling_rate) - waveform_length = int(self.transformer.config.sample_size) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - audio_start_in_s, - audio_end_in_s, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - attention_mask, - negative_attention_mask, - initial_audio_waveforms, - initial_audio_sampling_rate, - ) - - # 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - prompt_embeds = self.encode_prompt( - prompt, - device, - do_classifier_free_guidance, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - attention_mask, - negative_attention_mask, - ) - - # Encode duration - seconds_start_hidden_states, seconds_end_hidden_states = self.encode_duration( - audio_start_in_s, - audio_end_in_s, - device, - do_classifier_free_guidance and (negative_prompt is not None or negative_prompt_embeds is not None), - batch_size, - ) - - # Create text_audio_duration_embeds and audio_duration_embeds - text_audio_duration_embeds = torch.cat( - [prompt_embeds, seconds_start_hidden_states, seconds_end_hidden_states], dim=1 - ) - - audio_duration_embeds = torch.cat([seconds_start_hidden_states, seconds_end_hidden_states], dim=2) - - # In case of classifier free guidance without negative prompt, we need to create unconditional embeddings and - # to concatenate it to the embeddings - if do_classifier_free_guidance and negative_prompt_embeds is None and negative_prompt is None: - negative_text_audio_duration_embeds = torch.zeros_like( - text_audio_duration_embeds, device=text_audio_duration_embeds.device - ) - text_audio_duration_embeds = torch.cat( - [negative_text_audio_duration_embeds, text_audio_duration_embeds], dim=0 - ) - audio_duration_embeds = torch.cat([audio_duration_embeds, audio_duration_embeds], dim=0) - - bs_embed, seq_len, hidden_size = text_audio_duration_embeds.shape - # duplicate audio_duration_embeds and text_audio_duration_embeds for each generation per prompt, using mps friendly method - text_audio_duration_embeds = text_audio_duration_embeds.repeat(1, num_waveforms_per_prompt, 1) - text_audio_duration_embeds = text_audio_duration_embeds.view( - bs_embed * num_waveforms_per_prompt, seq_len, hidden_size - ) - - audio_duration_embeds = audio_duration_embeds.repeat(1, num_waveforms_per_prompt, 1) - audio_duration_embeds = audio_duration_embeds.view( - bs_embed * num_waveforms_per_prompt, -1, audio_duration_embeds.shape[-1] - ) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_vae = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_waveforms_per_prompt, - num_channels_vae, - waveform_length, - text_audio_duration_embeds.dtype, - device, - generator, - latents, - initial_audio_waveforms, - num_waveforms_per_prompt, - audio_channels=self.vae.config.audio_channels, - ) - - # 6. Prepare extra step kwargs - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Prepare rotary positional embedding - rotary_embedding = get_1d_rotary_pos_embed( - self.rotary_embed_dim, - latents.shape[2] + audio_duration_embeds.shape[1], - use_real=True, - repeat_interleave_real=False, - ) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.transformer( - latent_model_input, - t.unsqueeze(0), - encoder_hidden_states=text_audio_duration_embeds, - global_hidden_states=audio_duration_embeds, - rotary_embedding=rotary_embedding, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # 9. Post-processing - if not output_type == "latent": - audio = self.vae.decode(latents).sample - else: - return AudioPipelineOutput(audios=latents) - - audio = audio[:, :, waveform_start:waveform_end] - - if output_type == "np": - audio = audio.cpu().float().numpy() - - self.maybe_free_model_hooks() - - if not return_dict: - return (audio,) - - return AudioPipelineOutput(audios=audio) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__init__.py deleted file mode 100644 index 5270cb94af01fd94ed6c8e76c243a86cad8ec348..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__init__.py +++ /dev/null @@ -1,50 +0,0 @@ -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 - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) -else: - _import_structure["pipeline_stable_cascade"] = ["StableCascadeDecoderPipeline"] - _import_structure["pipeline_stable_cascade_combined"] = ["StableCascadeCombinedPipeline"] - _import_structure["pipeline_stable_cascade_prior"] = ["StableCascadePriorPipeline"] - - -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_stable_cascade import StableCascadeDecoderPipeline - from .pipeline_stable_cascade_combined import StableCascadeCombinedPipeline - from .pipeline_stable_cascade_prior import StableCascadePriorPipeline -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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 6a240b1c793669853d9cc22dc568f86695fae229..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__pycache__/pipeline_stable_cascade.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__pycache__/pipeline_stable_cascade.cpython-310.pyc deleted file mode 100644 index 37fa394b787f2864ee1021bf05765f28c4eaf17f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__pycache__/pipeline_stable_cascade.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__pycache__/pipeline_stable_cascade_combined.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__pycache__/pipeline_stable_cascade_combined.cpython-310.pyc deleted file mode 100644 index a484580ed1b38d6b395b97f0c0cc2cf33083988d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__pycache__/pipeline_stable_cascade_combined.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__pycache__/pipeline_stable_cascade_prior.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__pycache__/pipeline_stable_cascade_prior.cpython-310.pyc deleted file mode 100644 index 4ca24d77c2b104ba63f7e45e409faad0aa3dbb34..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/__pycache__/pipeline_stable_cascade_prior.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py deleted file mode 100644 index 6130a9873cb0cd33ec9391508a0aef9b79506da4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py +++ /dev/null @@ -1,539 +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 Callable, Dict, List, Optional, Union - -import torch -from transformers import CLIPTextModelWithProjection, CLIPTokenizer - -from ...models import StableCascadeUNet -from ...schedulers import DDPMWuerstchenScheduler -from ...utils import is_torch_version, is_torch_xla_available, logging, replace_example_docstring -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput -from ..wuerstchen.modeling_paella_vq_model import PaellaVQModel - - -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 StableCascadePriorPipeline, StableCascadeDecoderPipeline - - >>> prior_pipe = StableCascadePriorPipeline.from_pretrained( - ... "stabilityai/stable-cascade-prior", torch_dtype=torch.bfloat16 - ... ).to("cuda") - >>> gen_pipe = StableCascadeDecoderPipeline.from_pretrain( - ... "stabilityai/stable-cascade", torch_dtype=torch.float16 - ... ).to("cuda") - - >>> prompt = "an image of a shiba inu, donning a spacesuit and helmet" - >>> prior_output = pipe(prompt) - >>> images = gen_pipe(prior_output.image_embeddings, prompt=prompt) - ``` -""" - - -class StableCascadeDecoderPipeline(DiffusionPipeline): - """ - Pipeline for generating images from the Stable Cascade model. - - 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.) - - Args: - tokenizer (`CLIPTokenizer`): - The CLIP tokenizer. - text_encoder (`CLIPTextModelWithProjection`): - The CLIP text encoder. - decoder ([`StableCascadeUNet`]): - The Stable Cascade decoder unet. - vqgan ([`PaellaVQModel`]): - The VQGAN model. - scheduler ([`DDPMWuerstchenScheduler`]): - A scheduler to be used in combination with `prior` to generate image embedding. - latent_dim_scale (float, `optional`, defaults to 10.67): - Multiplier to determine the VQ latent space size from the image embeddings. If the image embeddings are - height=24 and width=24, the VQ latent shape needs to be height=int(24*10.67)=256 and - width=int(24*10.67)=256 in order to match the training conditions. - """ - - unet_name = "decoder" - text_encoder_name = "text_encoder" - model_cpu_offload_seq = "text_encoder->decoder->vqgan" - _callback_tensor_inputs = [ - "latents", - "prompt_embeds_pooled", - "negative_prompt_embeds", - "image_embeddings", - ] - - def __init__( - self, - decoder: StableCascadeUNet, - tokenizer: CLIPTokenizer, - text_encoder: CLIPTextModelWithProjection, - scheduler: DDPMWuerstchenScheduler, - vqgan: PaellaVQModel, - latent_dim_scale: float = 10.67, - ) -> None: - super().__init__() - self.register_modules( - decoder=decoder, - tokenizer=tokenizer, - text_encoder=text_encoder, - scheduler=scheduler, - vqgan=vqgan, - ) - self.register_to_config(latent_dim_scale=latent_dim_scale) - - def prepare_latents( - self, batch_size, image_embeddings, num_images_per_prompt, dtype, device, generator, latents, scheduler - ): - _, channels, height, width = image_embeddings.shape - latents_shape = ( - batch_size * num_images_per_prompt, - 4, - int(height * self.config.latent_dim_scale), - int(width * self.config.latent_dim_scale), - ) - - if latents is None: - latents = randn_tensor(latents_shape, generator=generator, device=device, dtype=dtype) - else: - if latents.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - latents = latents.to(device) - - latents = latents * scheduler.init_noise_sigma - return latents - - def encode_prompt( - self, - device, - batch_size, - num_images_per_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, - ): - if prompt_embeds is None: - # get prompt text embeddings - 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 - attention_mask = text_inputs.attention_mask - - 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}" - ) - text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length] - attention_mask = attention_mask[:, : self.tokenizer.model_max_length] - - text_encoder_output = self.text_encoder( - text_input_ids.to(device), attention_mask=attention_mask.to(device), output_hidden_states=True - ) - prompt_embeds = text_encoder_output.hidden_states[-1] - if prompt_embeds_pooled is None: - prompt_embeds_pooled = text_encoder_output.text_embeds.unsqueeze(1) - - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - prompt_embeds_pooled = prompt_embeds_pooled.to(dtype=self.text_encoder.dtype, device=device) - prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - 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] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif 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 - - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - negative_prompt_embeds_text_encoder_output = self.text_encoder( - uncond_input.input_ids.to(device), - attention_mask=uncond_input.attention_mask.to(device), - output_hidden_states=True, - ) - - negative_prompt_embeds = negative_prompt_embeds_text_encoder_output.hidden_states[-1] - negative_prompt_embeds_pooled = negative_prompt_embeds_text_encoder_output.text_embeds.unsqueeze(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] - negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.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) - - seq_len = negative_prompt_embeds_pooled.shape[1] - negative_prompt_embeds_pooled = negative_prompt_embeds_pooled.to( - dtype=self.text_encoder.dtype, device=device - ) - negative_prompt_embeds_pooled = negative_prompt_embeds_pooled.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds_pooled = negative_prompt_embeds_pooled.view( - batch_size * num_images_per_prompt, seq_len, -1 - ) - # done duplicates - - return prompt_embeds, prompt_embeds_pooled, negative_prompt_embeds, negative_prompt_embeds_pooled - - def check_inputs( - self, - prompt, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - callback_on_step_end_tensor_inputs=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 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}." - ) - - @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 - - def get_timestep_ratio_conditioning(self, t, alphas_cumprod): - s = torch.tensor([0.008]) - clamp_range = [0, 1] - min_var = torch.cos(s / (1 + s) * torch.pi * 0.5) ** 2 - var = alphas_cumprod[t] - var = var.clamp(*clamp_range) - s, min_var = s.to(var.device), min_var.to(var.device) - ratio = (((var * min_var) ** 0.5).acos() / (torch.pi * 0.5)) * (1 + s) - s - return ratio - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - image_embeddings: Union[torch.Tensor, List[torch.Tensor]], - prompt: Union[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, - 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", - return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], 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 Embeddings either extracted from an image or generated by a Prior Model. - 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. - guidance_scale (`float`, *optional*, defaults to 0.0): - Guidance scale as defined in [Classifier-Free Diffusion - Guidance](https://huggingface.co/papers/2207.12598). `decoder_guidance_scale` is defined as `w` of - 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*): - 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*): - 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_embeds_pooled (`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_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. - negative_prompt_embeds_pooled (`torch.Tensor`, *optional*): - Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, negative_prompt_embeds_pooled will be generated from `negative_prompt` - 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*): - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - 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. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` [`~pipelines.ImagePipelineOutput`] if `return_dict` is True, - otherwise a `tuple`. When returning a tuple, the first element is a list with the generated image - embeddings. - """ - - # 0. Define commonly used variables - device = self._execution_device - dtype = self.decoder.dtype - self._guidance_scale = guidance_scale - if is_torch_version("<", "2.2.0") and dtype == torch.bfloat16: - raise ValueError("`StableCascadeDecoderPipeline` requires torch>=2.2.0 when using `torch.bfloat16` dtype.") - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - 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, - ) - if isinstance(image_embeddings, list): - image_embeddings = torch.cat(image_embeddings, dim=0) - - 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] - - # Compute the effective number of images per prompt - # We must account for the fact that the image embeddings from the prior can be generated with num_images_per_prompt > 1 - # This results in a case where a single prompt is associated with multiple image embeddings - # Divide the number of image embeddings by the batch size to determine if this is the case. - num_images_per_prompt = num_images_per_prompt * (image_embeddings.shape[0] // batch_size) - - # 2. Encode caption - if prompt_embeds is None and negative_prompt_embeds is None: - _, prompt_embeds_pooled, _, negative_prompt_embeds_pooled = self.encode_prompt( - prompt=prompt, - device=device, - batch_size=batch_size, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - prompt_embeds_pooled=prompt_embeds_pooled, - negative_prompt_embeds=negative_prompt_embeds, - negative_prompt_embeds_pooled=negative_prompt_embeds_pooled, - ) - - # The pooled embeds from the prior are pooled again before being passed to the decoder - prompt_embeds_pooled = ( - torch.cat([prompt_embeds_pooled, negative_prompt_embeds_pooled]) - if self.do_classifier_free_guidance - else prompt_embeds_pooled - ) - effnet = ( - torch.cat([image_embeddings, torch.zeros_like(image_embeddings)]) - if self.do_classifier_free_guidance - else image_embeddings - ) - - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latents - latents = self.prepare_latents( - batch_size, image_embeddings, num_images_per_prompt, dtype, device, generator, latents, self.scheduler - ) - - if isinstance(self.scheduler, DDPMWuerstchenScheduler): - timesteps = timesteps[:-1] - else: - if hasattr(self.scheduler.config, "clip_sample") and self.scheduler.config.clip_sample: - self.scheduler.config.clip_sample = False # disample sample clipping - logger.warning(" set `clip_sample` to be False") - - # 6. Run denoising loop - if hasattr(self.scheduler, "betas"): - alphas = 1.0 - self.scheduler.betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - else: - alphas_cumprod = [] - - self._num_timesteps = len(timesteps) - for i, t in enumerate(self.progress_bar(timesteps)): - if not isinstance(self.scheduler, DDPMWuerstchenScheduler): - if len(alphas_cumprod) > 0: - timestep_ratio = self.get_timestep_ratio_conditioning(t.long().cpu(), alphas_cumprod) - timestep_ratio = timestep_ratio.expand(latents.size(0)).to(dtype).to(device) - else: - timestep_ratio = t.float().div(self.scheduler.timesteps[-1]).expand(latents.size(0)).to(dtype) - else: - timestep_ratio = t.expand(latents.size(0)).to(dtype) - - # 7. Denoise latents - predicted_latents = self.decoder( - sample=torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents, - timestep_ratio=torch.cat([timestep_ratio] * 2) if self.do_classifier_free_guidance else timestep_ratio, - clip_text_pooled=prompt_embeds_pooled, - effnet=effnet, - return_dict=False, - )[0] - - # 8. Check for classifier free guidance and apply it - if self.do_classifier_free_guidance: - predicted_latents_text, predicted_latents_uncond = predicted_latents.chunk(2) - predicted_latents = torch.lerp(predicted_latents_uncond, predicted_latents_text, self.guidance_scale) - - # 9. Renoise latents to next timestep - if not isinstance(self.scheduler, DDPMWuerstchenScheduler): - timestep_ratio = t - latents = self.scheduler.step( - model_output=predicted_latents, - timestep=timestep_ratio, - sample=latents, - generator=generator, - ).prev_sample - - 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 XLA_AVAILABLE: - xm.mark_step() - - if output_type not in ["pt", "np", "pil", "latent"]: - raise ValueError( - f"Only the output types `pt`, `np`, `pil` and `latent` are supported not output_type={output_type}" - ) - - if not output_type == "latent": - # 10. Scale and decode the image latents with vq-vae - latents = self.vqgan.config.scale_factor * latents - images = self.vqgan.decode(latents).sample.clamp(0, 1) - if output_type == "np": - images = images.permute(0, 2, 3, 1).cpu().float().numpy() # float() as bfloat16-> numpy doesn't work - elif output_type == "pil": - images = images.permute(0, 2, 3, 1).cpu().float().numpy() # float() as bfloat16-> numpy doesn't work - images = self.numpy_to_pil(images) - else: - images = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return images - return ImagePipelineOutput(images) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py deleted file mode 100644 index b705c7e6e5f6f447371b5db8204e91706427c586..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py +++ /dev/null @@ -1,319 +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 Callable, Dict, List, Optional, Union - -import PIL -import torch -from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...models import StableCascadeUNet -from ...schedulers import DDPMWuerstchenScheduler -from ...utils import is_torch_version, replace_example_docstring -from ..pipeline_utils import DiffusionPipeline -from ..wuerstchen.modeling_paella_vq_model import PaellaVQModel -from .pipeline_stable_cascade import StableCascadeDecoderPipeline -from .pipeline_stable_cascade_prior import StableCascadePriorPipeline - - -TEXT2IMAGE_EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import StableCascadeCombinedPipeline - - >>> pipe = StableCascadeCombinedPipeline.from_pretrained( - ... "stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.bfloat16 - ... ) - >>> pipe.enable_model_cpu_offload() - >>> prompt = "an image of a shiba inu, donning a spacesuit and helmet" - >>> images = pipe(prompt=prompt) - ``` -""" - - -class StableCascadeCombinedPipeline(DiffusionPipeline): - """ - Combined Pipeline for text-to-image generation using Stable Cascade. - - 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.) - - Args: - tokenizer (`CLIPTokenizer`): - The decoder tokenizer to be used for text inputs. - text_encoder (`CLIPTextModelWithProjection`): - The decoder text encoder to be used for text inputs. - decoder (`StableCascadeUNet`): - The decoder model to be used for decoder image generation pipeline. - scheduler (`DDPMWuerstchenScheduler`): - The scheduler to be used for decoder image generation pipeline. - vqgan (`PaellaVQModel`): - The VQGAN model to be used for decoder image generation pipeline. - prior_prior (`StableCascadeUNet`): - The prior model to be used for prior pipeline. - prior_text_encoder (`CLIPTextModelWithProjection`): - The prior text encoder to be used for text inputs. - prior_tokenizer (`CLIPTokenizer`): - The prior tokenizer to be used for text inputs. - prior_scheduler (`DDPMWuerstchenScheduler`): - The scheduler to be used for prior pipeline. - prior_feature_extractor ([`~transformers.CLIPImageProcessor`]): - Model that extracts features from generated images to be used as inputs for the `image_encoder`. - prior_image_encoder ([`CLIPVisionModelWithProjection`]): - Frozen CLIP image-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - """ - - _load_connected_pipes = True - _optional_components = ["prior_feature_extractor", "prior_image_encoder"] - - def __init__( - self, - tokenizer: CLIPTokenizer, - text_encoder: CLIPTextModelWithProjection, - decoder: StableCascadeUNet, - scheduler: DDPMWuerstchenScheduler, - vqgan: PaellaVQModel, - prior_prior: StableCascadeUNet, - prior_text_encoder: CLIPTextModelWithProjection, - prior_tokenizer: CLIPTokenizer, - prior_scheduler: DDPMWuerstchenScheduler, - prior_feature_extractor: Optional[CLIPImageProcessor] = None, - prior_image_encoder: Optional[CLIPVisionModelWithProjection] = None, - ): - super().__init__() - - self.register_modules( - text_encoder=text_encoder, - tokenizer=tokenizer, - decoder=decoder, - scheduler=scheduler, - vqgan=vqgan, - prior_text_encoder=prior_text_encoder, - prior_tokenizer=prior_tokenizer, - prior_prior=prior_prior, - prior_scheduler=prior_scheduler, - prior_feature_extractor=prior_feature_extractor, - prior_image_encoder=prior_image_encoder, - ) - self.prior_pipe = StableCascadePriorPipeline( - prior=prior_prior, - text_encoder=prior_text_encoder, - tokenizer=prior_tokenizer, - scheduler=prior_scheduler, - image_encoder=prior_image_encoder, - feature_extractor=prior_feature_extractor, - ) - self.decoder_pipe = StableCascadeDecoderPipeline( - text_encoder=text_encoder, - tokenizer=tokenizer, - decoder=decoder, - scheduler=scheduler, - vqgan=vqgan, - ) - - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = 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): - 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` - method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with - `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`. - """ - 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): - 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 - GPU only when their specific submodule's `forward` method is called. Offloading happens on a submodule basis. - Memory savings are higher than using `enable_model_cpu_offload`, but performance is lower. - """ - self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - - def progress_bar(self, iterable=None, total=None): - self.prior_pipe.progress_bar(iterable=iterable, total=total) - self.decoder_pipe.progress_bar(iterable=iterable, total=total) - - def set_progress_bar_config(self, **kwargs): - self.prior_pipe.set_progress_bar_config(**kwargs) - self.decoder_pipe.set_progress_bar_config(**kwargs) - - @torch.no_grad() - @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, - 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, - 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", - 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"], - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - 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*): - The images to guide the image generation for the prior. - 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*): - Pre-generated text embeddings for the prior. 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_embeds_pooled (`torch.Tensor`, *optional*): - Pre-generated text embeddings for the prior. 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 for the prior. 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. - negative_prompt_embeds_pooled (`torch.Tensor`, *optional*): - Pre-generated negative text embeddings for the prior. 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. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - prior_guidance_scale (`float`, *optional*, defaults to 4.0): - Guidance scale as defined in [Classifier-Free Diffusion - Guidance](https://huggingface.co/papers/2207.12598). `prior_guidance_scale` is defined as `w` of - 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): - 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` - num_inference_steps (`int`, *optional*, defaults to 12): - 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` - decoder_guidance_scale (`float`, *optional*, defaults to 0.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. - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - prior_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: `prior_callback_on_step_end(self: DiffusionPipeline, step: int, timestep: - int, callback_kwargs: Dict)`. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` [`~pipelines.ImagePipelineOutput`] if `return_dict` is True, - otherwise a `tuple`. When returning a tuple, the first element is a list with the generated images. - """ - dtype = self.decoder_pipe.decoder.dtype - if is_torch_version("<", "2.2.0") and dtype == torch.bfloat16: - raise ValueError( - "`StableCascadeCombinedPipeline` requires torch>=2.2.0 when using `torch.bfloat16` dtype." - ) - - prior_outputs = self.prior_pipe( - prompt=prompt if prompt_embeds is None else None, - images=images, - height=height, - width=width, - num_inference_steps=prior_num_inference_steps, - guidance_scale=prior_guidance_scale, - negative_prompt=negative_prompt if negative_prompt_embeds is None else None, - prompt_embeds=prompt_embeds, - prompt_embeds_pooled=prompt_embeds_pooled, - negative_prompt_embeds=negative_prompt_embeds, - negative_prompt_embeds_pooled=negative_prompt_embeds_pooled, - num_images_per_prompt=num_images_per_prompt, - generator=generator, - latents=latents, - output_type="pt", - return_dict=True, - callback_on_step_end=prior_callback_on_step_end, - callback_on_step_end_tensor_inputs=prior_callback_on_step_end_tensor_inputs, - ) - image_embeddings = prior_outputs.image_embeddings - prompt_embeds = prior_outputs.get("prompt_embeds", None) - prompt_embeds_pooled = prior_outputs.get("prompt_embeds_pooled", None) - negative_prompt_embeds = prior_outputs.get("negative_prompt_embeds", None) - negative_prompt_embeds_pooled = prior_outputs.get("negative_prompt_embeds_pooled", None) - - outputs = self.decoder_pipe( - image_embeddings=image_embeddings, - prompt=prompt if prompt_embeds is None else None, - num_inference_steps=num_inference_steps, - guidance_scale=decoder_guidance_scale, - negative_prompt=negative_prompt if negative_prompt_embeds is None else None, - prompt_embeds=prompt_embeds, - prompt_embeds_pooled=prompt_embeds_pooled, - negative_prompt_embeds=negative_prompt_embeds, - negative_prompt_embeds_pooled=negative_prompt_embeds_pooled, - generator=generator, - output_type=output_type, - return_dict=return_dict, - callback_on_step_end=callback_on_step_end, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - ) - - return outputs diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py deleted file mode 100644 index b3b46af206edf575011add5291385c3ec1f2ef52..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py +++ /dev/null @@ -1,650 +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 dataclasses import dataclass -from math import ceil -from typing import Callable, Dict, List, Optional, Union - -import numpy as np -import PIL -import torch -from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...models import StableCascadeUNet -from ...schedulers import DDPMWuerstchenScheduler -from ...utils import BaseOutput, is_torch_xla_available, logging, replace_example_docstring -from ...utils.torch_utils import randn_tensor -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 - - -DEFAULT_STAGE_C_TIMESTEPS = list(np.linspace(1.0, 2 / 3, 20)) + list(np.linspace(2 / 3, 0.0, 11))[1:] - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import StableCascadePriorPipeline - - >>> prior_pipe = StableCascadePriorPipeline.from_pretrained( - ... "stabilityai/stable-cascade-prior", torch_dtype=torch.bfloat16 - ... ).to("cuda") - - >>> prompt = "an image of a shiba inu, donning a spacesuit and helmet" - >>> prior_output = pipe(prompt) - ``` -""" - - -@dataclass -class StableCascadePriorPipelineOutput(BaseOutput): - """ - Output class for WuerstchenPriorPipeline. - - Args: - image_embeddings (`torch.Tensor` or `np.ndarray`) - Prior image embeddings for text prompt - prompt_embeds (`torch.Tensor`): - Text embeddings for the prompt. - negative_prompt_embeds (`torch.Tensor`): - 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] - - -class StableCascadePriorPipeline(DiffusionPipeline): - """ - Pipeline for generating image prior for Stable Cascade. - - 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.) - - Args: - prior ([`StableCascadeUNet`]): - The Stable Cascade prior to approximate the image embedding from the text and/or image embedding. - text_encoder ([`CLIPTextModelWithProjection`]): - Frozen text-encoder - ([laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)). - feature_extractor ([`~transformers.CLIPImageProcessor`]): - Model that extracts features from generated images to be used as inputs for the `image_encoder`. - image_encoder ([`CLIPVisionModelWithProjection`]): - Frozen CLIP image-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - scheduler ([`DDPMWuerstchenScheduler`]): - A scheduler to be used in combination with `prior` to generate image embedding. - resolution_multiple ('float', *optional*, defaults to 42.67): - Default resolution for multiple images generated. - """ - - unet_name = "prior" - text_encoder_name = "text_encoder" - model_cpu_offload_seq = "image_encoder->text_encoder->prior" - _optional_components = ["image_encoder", "feature_extractor"] - _callback_tensor_inputs = ["latents", "text_encoder_hidden_states", "negative_prompt_embeds"] - - def __init__( - self, - tokenizer: CLIPTokenizer, - text_encoder: CLIPTextModelWithProjection, - prior: StableCascadeUNet, - scheduler: DDPMWuerstchenScheduler, - resolution_multiple: float = 42.67, - feature_extractor: Optional[CLIPImageProcessor] = None, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, - ) -> None: - super().__init__() - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - prior=prior, - scheduler=scheduler, - ) - self.register_to_config(resolution_multiple=resolution_multiple) - - def prepare_latents( - self, batch_size, height, width, num_images_per_prompt, dtype, device, generator, latents, scheduler - ): - latent_shape = ( - num_images_per_prompt * batch_size, - self.prior.config.in_channels, - ceil(height / self.config.resolution_multiple), - ceil(width / self.config.resolution_multiple), - ) - - if latents is None: - latents = randn_tensor(latent_shape, generator=generator, device=device, dtype=dtype) - else: - if latents.shape != latent_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latent_shape}") - latents = latents.to(device) - - latents = latents * scheduler.init_noise_sigma - return latents - - def encode_prompt( - self, - device, - batch_size, - num_images_per_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, - ): - if prompt_embeds is None: - # get prompt text embeddings - 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 - attention_mask = text_inputs.attention_mask - - 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}" - ) - text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length] - attention_mask = attention_mask[:, : self.tokenizer.model_max_length] - - text_encoder_output = self.text_encoder( - text_input_ids.to(device), attention_mask=attention_mask.to(device), output_hidden_states=True - ) - prompt_embeds = text_encoder_output.hidden_states[-1] - if prompt_embeds_pooled is None: - prompt_embeds_pooled = text_encoder_output.text_embeds.unsqueeze(1) - - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - prompt_embeds_pooled = prompt_embeds_pooled.to(dtype=self.text_encoder.dtype, device=device) - prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - 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] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif 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 - - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - negative_prompt_embeds_text_encoder_output = self.text_encoder( - uncond_input.input_ids.to(device), - attention_mask=uncond_input.attention_mask.to(device), - output_hidden_states=True, - ) - - negative_prompt_embeds = negative_prompt_embeds_text_encoder_output.hidden_states[-1] - negative_prompt_embeds_pooled = negative_prompt_embeds_text_encoder_output.text_embeds.unsqueeze(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] - negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.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) - - seq_len = negative_prompt_embeds_pooled.shape[1] - negative_prompt_embeds_pooled = negative_prompt_embeds_pooled.to( - dtype=self.text_encoder.dtype, device=device - ) - negative_prompt_embeds_pooled = negative_prompt_embeds_pooled.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds_pooled = negative_prompt_embeds_pooled.view( - batch_size * num_images_per_prompt, seq_len, -1 - ) - # done duplicates - - return prompt_embeds, prompt_embeds_pooled, negative_prompt_embeds, negative_prompt_embeds_pooled - - def encode_image(self, images, device, dtype, batch_size, num_images_per_prompt): - image_embeds = [] - for image in images: - image = self.feature_extractor(image, return_tensors="pt").pixel_values - image = image.to(device=device, dtype=dtype) - image_embed = self.image_encoder(image).image_embeds.unsqueeze(1) - image_embeds.append(image_embed) - image_embeds = torch.cat(image_embeds, dim=1) - - image_embeds = image_embeds.repeat(batch_size * num_images_per_prompt, 1, 1) - negative_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, negative_image_embeds - - def check_inputs( - self, - prompt, - images=None, - image_embeds=None, - negative_prompt=None, - prompt_embeds=None, - prompt_embeds_pooled=None, - negative_prompt_embeds=None, - negative_prompt_embeds_pooled=None, - callback_on_step_end_tensor_inputs=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 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_embeds_pooled is None: - raise ValueError( - "If `prompt_embeds` are provided, `prompt_embeds_pooled` must also be provided. Make sure to generate `prompt_embeds_pooled` from the same text encoder that was used to generate `prompt_embeds`" - ) - - if negative_prompt_embeds is not None and negative_prompt_embeds_pooled is None: - raise ValueError( - "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_pooled` must also be provided. Make sure to generate `prompt_embeds_pooled` from the same text encoder that was used to generate `prompt_embeds`" - ) - - if prompt_embeds_pooled is not None and negative_prompt_embeds_pooled is not None: - if prompt_embeds_pooled.shape != negative_prompt_embeds_pooled.shape: - raise ValueError( - "`prompt_embeds_pooled` and `negative_prompt_embeds_pooled` must have the same shape when passed" - f"directly, but got: `prompt_embeds_pooled` {prompt_embeds_pooled.shape} !=" - f"`negative_prompt_embeds_pooled` {negative_prompt_embeds_pooled.shape}." - ) - - if image_embeds is not None and images is not None: - raise ValueError( - f"Cannot forward both `images`: {images} and `image_embeds`: {image_embeds}. Please make sure to" - " only forward one of the two." - ) - - if images: - for i, image in enumerate(images): - if not isinstance(image, torch.Tensor) and not isinstance(image, PIL.Image.Image): - raise TypeError( - f"'images' must contain images of type 'torch.Tensor' or 'PIL.Image.Image, but got" - f"{type(image)} for image number {i}." - ) - - @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 - - def get_timestep_ratio_conditioning(self, t, alphas_cumprod): - s = torch.tensor([0.008]) - clamp_range = [0, 1] - min_var = torch.cos(s / (1 + s) * torch.pi * 0.5) ** 2 - var = alphas_cumprod[t] - var = var.clamp(*clamp_range) - s, min_var = s.to(var.device), min_var.to(var.device) - ratio = (((var * min_var) ** 0.5).acos() / (torch.pi * 0.5)) * (1 + s) - s - return ratio - - @torch.no_grad() - @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, - height: int = 1024, - width: int = 1024, - num_inference_steps: int = 20, - 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", - return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], 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]`): - The prompt or prompts to guide the image generation. - height (`int`, *optional*, defaults to 1024): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 1024): - The width in pixels of the generated image. - 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. - guidance_scale (`float`, *optional*, defaults to 8.0): - Guidance scale as defined in [Classifier-Free Diffusion - Guidance](https://huggingface.co/papers/2207.12598). `decoder_guidance_scale` is defined as `w` of - 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*): - 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*): - 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_embeds_pooled (`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_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. - negative_prompt_embeds_pooled (`torch.Tensor`, *optional*): - Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, negative_prompt_embeds_pooled will be generated from `negative_prompt` - input argument. - image_embeds (`torch.Tensor`, *optional*): - Pre-generated image embeddings. Can be used to easily tweak image inputs, *e.g.* prompt weighting. If - 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*): - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - 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. - - Examples: - - Returns: - [`StableCascadePriorPipelineOutput`] or `tuple` [`StableCascadePriorPipelineOutput`] if `return_dict` is - True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated image - embeddings. - """ - - # 0. Define commonly used variables - device = self._execution_device - dtype = next(self.prior.parameters()).dtype - self._guidance_scale = guidance_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] - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - images=images, - image_embeds=image_embeds, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - prompt_embeds_pooled=prompt_embeds_pooled, - negative_prompt_embeds=negative_prompt_embeds, - negative_prompt_embeds_pooled=negative_prompt_embeds_pooled, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - ) - - # 2. Encode caption + images - ( - prompt_embeds, - prompt_embeds_pooled, - negative_prompt_embeds, - negative_prompt_embeds_pooled, - ) = self.encode_prompt( - prompt=prompt, - device=device, - batch_size=batch_size, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - prompt_embeds_pooled=prompt_embeds_pooled, - negative_prompt_embeds=negative_prompt_embeds, - negative_prompt_embeds_pooled=negative_prompt_embeds_pooled, - ) - - if images is not None: - image_embeds_pooled, uncond_image_embeds_pooled = self.encode_image( - images=images, - device=device, - dtype=dtype, - batch_size=batch_size, - num_images_per_prompt=num_images_per_prompt, - ) - elif image_embeds is not None: - image_embeds_pooled = image_embeds.repeat(batch_size * num_images_per_prompt, 1, 1) - uncond_image_embeds_pooled = torch.zeros_like(image_embeds_pooled) - else: - image_embeds_pooled = torch.zeros( - batch_size * num_images_per_prompt, - 1, - self.prior.config.clip_image_in_channels, - device=device, - dtype=dtype, - ) - uncond_image_embeds_pooled = torch.zeros( - batch_size * num_images_per_prompt, - 1, - self.prior.config.clip_image_in_channels, - device=device, - dtype=dtype, - ) - - if self.do_classifier_free_guidance: - image_embeds = torch.cat([image_embeds_pooled, uncond_image_embeds_pooled], dim=0) - else: - image_embeds = image_embeds_pooled - - # 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 - text_encoder_hidden_states = ( - torch.cat([prompt_embeds, negative_prompt_embeds]) if negative_prompt_embeds is not None else prompt_embeds - ) - text_encoder_pooled = ( - torch.cat([prompt_embeds_pooled, negative_prompt_embeds_pooled]) - if negative_prompt_embeds is not None - else prompt_embeds_pooled - ) - - # 4. Prepare and set timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latents - latents = self.prepare_latents( - batch_size, height, width, num_images_per_prompt, dtype, device, generator, latents, self.scheduler - ) - - if isinstance(self.scheduler, DDPMWuerstchenScheduler): - timesteps = timesteps[:-1] - else: - if hasattr(self.scheduler.config, "clip_sample") and self.scheduler.config.clip_sample: - self.scheduler.config.clip_sample = False # disample sample clipping - logger.warning(" set `clip_sample` to be False") - # 6. Run denoising loop - if hasattr(self.scheduler, "betas"): - alphas = 1.0 - self.scheduler.betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - else: - alphas_cumprod = [] - - self._num_timesteps = len(timesteps) - for i, t in enumerate(self.progress_bar(timesteps)): - if not isinstance(self.scheduler, DDPMWuerstchenScheduler): - if len(alphas_cumprod) > 0: - timestep_ratio = self.get_timestep_ratio_conditioning(t.long().cpu(), alphas_cumprod) - timestep_ratio = timestep_ratio.expand(latents.size(0)).to(dtype).to(device) - else: - timestep_ratio = t.float().div(self.scheduler.timesteps[-1]).expand(latents.size(0)).to(dtype) - else: - timestep_ratio = t.expand(latents.size(0)).to(dtype) - # 7. Denoise image embeddings - predicted_image_embedding = self.prior( - sample=torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents, - timestep_ratio=torch.cat([timestep_ratio] * 2) if self.do_classifier_free_guidance else timestep_ratio, - clip_text_pooled=text_encoder_pooled, - clip_text=text_encoder_hidden_states, - clip_img=image_embeds, - return_dict=False, - )[0] - - # 8. Check for classifier free guidance and apply it - if self.do_classifier_free_guidance: - predicted_image_embedding_text, predicted_image_embedding_uncond = predicted_image_embedding.chunk(2) - predicted_image_embedding = torch.lerp( - predicted_image_embedding_uncond, predicted_image_embedding_text, self.guidance_scale - ) - - # 9. Renoise latents to next timestep - if not isinstance(self.scheduler, DDPMWuerstchenScheduler): - timestep_ratio = t - latents = self.scheduler.step( - model_output=predicted_image_embedding, timestep=timestep_ratio, sample=latents, generator=generator - ).prev_sample - - 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 XLA_AVAILABLE: - xm.mark_step() - - # Offload all models - self.maybe_free_model_hooks() - - if output_type == "np": - latents = latents.cpu().float().numpy() # float() as bfloat16-> numpy doesn't work - prompt_embeds = prompt_embeds.cpu().float().numpy() # float() as bfloat16-> numpy doesn't work - negative_prompt_embeds = ( - negative_prompt_embeds.cpu().float().numpy() if negative_prompt_embeds is not None else None - ) # float() as bfloat16-> numpy doesn't work - - if not return_dict: - return ( - latents, - prompt_embeds, - prompt_embeds_pooled, - negative_prompt_embeds, - negative_prompt_embeds_pooled, - ) - - return StableCascadePriorPipelineOutput( - image_embeddings=latents, - prompt_embeds=prompt_embeds, - prompt_embeds_pooled=prompt_embeds_pooled, - negative_prompt_embeds=negative_prompt_embeds, - negative_prompt_embeds_pooled=negative_prompt_embeds_pooled, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__init__.py deleted file mode 100644 index b05a3ce2a7c9d7b388d422268b77f306419e40a6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__init__.py +++ /dev/null @@ -1,195 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _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, - is_transformers_version, -) - - -_dummy_objects = {} -_additional_imports = {} -_import_structure = {"pipeline_output": ["StableDiffusionPipelineOutput"]} - -if is_transformers_available() and is_flax_available(): - _import_structure["pipeline_output"].extend(["FlaxStableDiffusionPipelineOutput"]) -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["clip_image_project_model"] = ["CLIPImageProjection"] - _import_structure["pipeline_stable_diffusion"] = ["StableDiffusionPipeline"] - _import_structure["pipeline_stable_diffusion_img2img"] = ["StableDiffusionImg2ImgPipeline"] - _import_structure["pipeline_stable_diffusion_inpaint"] = ["StableDiffusionInpaintPipeline"] - _import_structure["pipeline_stable_diffusion_instruct_pix2pix"] = ["StableDiffusionInstructPix2PixPipeline"] - _import_structure["pipeline_stable_diffusion_latent_upscale"] = ["StableDiffusionLatentUpscalePipeline"] - _import_structure["pipeline_stable_diffusion_upscale"] = ["StableDiffusionUpscalePipeline"] - _import_structure["pipeline_stable_unclip"] = ["StableUnCLIPPipeline"] - _import_structure["pipeline_stable_unclip_img2img"] = ["StableUnCLIPImg2ImgPipeline"] - _import_structure["safety_checker"] = ["StableDiffusionSafetyChecker"] - _import_structure["stable_unclip_image_normalizer"] = ["StableUnCLIPImageNormalizer"] -try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.25.0")): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_objects import ( - StableDiffusionImageVariationPipeline, - ) - - _dummy_objects.update({"StableDiffusionImageVariationPipeline": StableDiffusionImageVariationPipeline}) -else: - _import_structure["pipeline_stable_diffusion_image_variation"] = ["StableDiffusionImageVariationPipeline"] -try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.26.0")): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_objects import ( - StableDiffusionDepth2ImgPipeline, - ) - - _dummy_objects.update( - { - "StableDiffusionDepth2ImgPipeline": StableDiffusionDepth2ImgPipeline, - } - ) -else: - _import_structure["pipeline_stable_diffusion_depth2img"] = ["StableDiffusionDepth2ImgPipeline"] - -try: - if not (is_transformers_available() and is_onnx_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_onnx_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_onnx_objects)) -else: - _import_structure["pipeline_onnx_stable_diffusion"] = [ - "OnnxStableDiffusionPipeline", - "StableDiffusionOnnxPipeline", - ] - _import_structure["pipeline_onnx_stable_diffusion_img2img"] = ["OnnxStableDiffusionImg2ImgPipeline"] - _import_structure["pipeline_onnx_stable_diffusion_inpaint"] = ["OnnxStableDiffusionInpaintPipeline"] - _import_structure["pipeline_onnx_stable_diffusion_inpaint_legacy"] = ["OnnxStableDiffusionInpaintPipelineLegacy"] - _import_structure["pipeline_onnx_stable_diffusion_upscale"] = ["OnnxStableDiffusionUpscalePipeline"] - -if is_transformers_available() and is_flax_available(): - from ...schedulers.scheduling_pndm_flax import PNDMSchedulerState - - _additional_imports.update({"PNDMSchedulerState": PNDMSchedulerState}) - _import_structure["pipeline_flax_stable_diffusion"] = ["FlaxStableDiffusionPipeline"] - _import_structure["pipeline_flax_stable_diffusion_img2img"] = ["FlaxStableDiffusionImg2ImgPipeline"] - _import_structure["pipeline_flax_stable_diffusion_inpaint"] = ["FlaxStableDiffusionInpaintPipeline"] - _import_structure["safety_checker_flax"] = ["FlaxStableDiffusionSafetyChecker"] - -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 .clip_image_project_model import CLIPImageProjection - from .pipeline_stable_diffusion import ( - StableDiffusionPipeline, - StableDiffusionPipelineOutput, - ) - from .pipeline_stable_diffusion_img2img import StableDiffusionImg2ImgPipeline - from .pipeline_stable_diffusion_inpaint import StableDiffusionInpaintPipeline - from .pipeline_stable_diffusion_instruct_pix2pix import ( - StableDiffusionInstructPix2PixPipeline, - ) - from .pipeline_stable_diffusion_latent_upscale import ( - StableDiffusionLatentUpscalePipeline, - ) - from .pipeline_stable_diffusion_upscale import StableDiffusionUpscalePipeline - from .pipeline_stable_unclip import StableUnCLIPPipeline - from .pipeline_stable_unclip_img2img import StableUnCLIPImg2ImgPipeline - from .safety_checker import StableDiffusionSafetyChecker - from .stable_unclip_image_normalizer import StableUnCLIPImageNormalizer - - try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.25.0")): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_objects import ( - StableDiffusionImageVariationPipeline, - ) - else: - from .pipeline_stable_diffusion_image_variation import ( - StableDiffusionImageVariationPipeline, - ) - - try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.26.0")): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_objects import StableDiffusionDepth2ImgPipeline - else: - from .pipeline_stable_diffusion_depth2img import ( - StableDiffusionDepth2ImgPipeline, - ) - - try: - if not (is_transformers_available() and is_onnx_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_onnx_objects import * - else: - from .pipeline_onnx_stable_diffusion import ( - OnnxStableDiffusionPipeline, - StableDiffusionOnnxPipeline, - ) - from .pipeline_onnx_stable_diffusion_img2img import ( - OnnxStableDiffusionImg2ImgPipeline, - ) - from .pipeline_onnx_stable_diffusion_inpaint import ( - OnnxStableDiffusionInpaintPipeline, - ) - from .pipeline_onnx_stable_diffusion_upscale import ( - OnnxStableDiffusionUpscalePipeline, - ) - - try: - if not (is_transformers_available() and is_flax_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_flax_objects import * - else: - from .pipeline_flax_stable_diffusion import FlaxStableDiffusionPipeline - from .pipeline_flax_stable_diffusion_img2img import ( - FlaxStableDiffusionImg2ImgPipeline, - ) - from .pipeline_flax_stable_diffusion_inpaint import ( - FlaxStableDiffusionInpaintPipeline, - ) - from .pipeline_output import FlaxStableDiffusionPipelineOutput - from .safety_checker_flax import FlaxStableDiffusionSafetyChecker - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 583d415741f1bc40eba6fd6726c716f56988417a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/clip_image_project_model.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/clip_image_project_model.cpython-310.pyc deleted file mode 100644 index 59fb85b6882a5702ce44911861897ab72198c678..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/clip_image_project_model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/convert_from_ckpt.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/convert_from_ckpt.cpython-310.pyc deleted file mode 100644 index a3d7d266d2dcde326d9e9599bdfbe3f4a69fcb8e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/convert_from_ckpt.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_flax_stable_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_flax_stable_diffusion.cpython-310.pyc deleted file mode 100644 index b8ee8925bcc4a252a37226890546efb6cfd0e186..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_flax_stable_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_flax_stable_diffusion_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_flax_stable_diffusion_img2img.cpython-310.pyc deleted file mode 100644 index 50b94c9701a7a0ed623f3b580b03db2d6e8142af..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_flax_stable_diffusion_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_flax_stable_diffusion_inpaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_flax_stable_diffusion_inpaint.cpython-310.pyc deleted file mode 100644 index 8b83ca6ec4305cfa7fd52aee9fb923d8c251c9e5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_flax_stable_diffusion_inpaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion.cpython-310.pyc deleted file mode 100644 index d535ad6f3566c007b3f02b4242433fe458de073e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion_img2img.cpython-310.pyc deleted file mode 100644 index c6518855471e0ff1419b8ff0fb6516bdac22050c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion_inpaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion_inpaint.cpython-310.pyc deleted file mode 100644 index d6c27340c6433ffdbdddc1f5f96301a1b45a84a9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion_inpaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion_upscale.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion_upscale.cpython-310.pyc deleted file mode 100644 index f89bb96ea1436e1022455e6f3ffe8a671a90ce2b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_onnx_stable_diffusion_upscale.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index 1f6d3f8c4696e2269e9bfd1a1560dd2328e6e8f1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion.cpython-310.pyc deleted file mode 100644 index c2a0d75d5ed65befebe7887353d778655d894c29..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_depth2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_depth2img.cpython-310.pyc deleted file mode 100644 index 9227ba687234ccfa15f57f6f025e54575e1d8d4a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_depth2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_image_variation.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_image_variation.cpython-310.pyc deleted file mode 100644 index 9845b925415c7005946ef03da3f073c7eb55b993..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_image_variation.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_img2img.cpython-310.pyc deleted file mode 100644 index 5350b35b376bf211e7a69fad18a22e59a36ac156..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_inpaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_inpaint.cpython-310.pyc deleted file mode 100644 index f43d8a691e8246acf37d7c055006cc696298b8a6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_inpaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_instruct_pix2pix.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_instruct_pix2pix.cpython-310.pyc deleted file mode 100644 index e15fb9a2f0adc9170fbb763a41cea2cf89eac4a2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_instruct_pix2pix.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_latent_upscale.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_latent_upscale.cpython-310.pyc deleted file mode 100644 index 2c38c32e934eb19d709bb05676eb59ccb57c0576..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_latent_upscale.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_upscale.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_upscale.cpython-310.pyc deleted file mode 100644 index bae34a7c2279b8a3fcc4ffe85ca111b989220830..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_diffusion_upscale.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_unclip.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_unclip.cpython-310.pyc deleted file mode 100644 index 90f238be3e75115b0cf96fc251a407a911e3673d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_unclip.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_unclip_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_unclip_img2img.cpython-310.pyc deleted file mode 100644 index 43ab2d3482a313e4e243f34e216d0a5d01f97d77..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/pipeline_stable_unclip_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/safety_checker.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/safety_checker.cpython-310.pyc deleted file mode 100644 index 353bccce9b68210a838566c36c67811c37cb81b2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/safety_checker.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/safety_checker_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/safety_checker_flax.cpython-310.pyc deleted file mode 100644 index 0be189bc7c7a94df948c18428b6e5f1dedec72d9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/safety_checker_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/stable_unclip_image_normalizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/stable_unclip_image_normalizer.cpython-310.pyc deleted file mode 100644 index a684feb69ed2b5cfce050e772ff274cb3b2081bb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/__pycache__/stable_unclip_image_normalizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/clip_image_project_model.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/clip_image_project_model.py deleted file mode 100644 index 30dd90242d07d249e62e976c148a9d2cde02271e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/clip_image_project_model.py +++ /dev/null @@ -1,29 +0,0 @@ -# Copyright 2025 The GLIGEN Authors and 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 torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.modeling_utils import ModelMixin - - -class CLIPImageProjection(ModelMixin, ConfigMixin): - @register_to_config - def __init__(self, hidden_size: int = 768): - super().__init__() - self.hidden_size = hidden_size - self.project = nn.Linear(self.hidden_size, self.hidden_size, bias=False) - - def forward(self, x): - return self.project(x) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py deleted file mode 100644 index 6c0221d2092a46dfac6129ab5bbe783c753ab97c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py +++ /dev/null @@ -1,1877 +0,0 @@ -# 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. -"""Conversion script for the Stable Diffusion checkpoints.""" - -import re -from contextlib import nullcontext -from io import BytesIO -from typing import Dict, Optional, Union - -import requests -import torch -import yaml -from transformers import ( - AutoFeatureExtractor, - BertTokenizerFast, - CLIPImageProcessor, - CLIPTextConfig, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionConfig, - CLIPVisionModelWithProjection, -) - -from ...models import ( - AutoencoderKL, - ControlNetModel, - PriorTransformer, - UNet2DConditionModel, -) -from ...schedulers import ( - DDIMScheduler, - DDPMScheduler, - DPMSolverMultistepScheduler, - EulerAncestralDiscreteScheduler, - EulerDiscreteScheduler, - HeunDiscreteScheduler, - LMSDiscreteScheduler, - PNDMScheduler, - UnCLIPScheduler, -) -from ...utils import is_accelerate_available, logging -from ...utils.constants import DIFFUSERS_REQUEST_TIMEOUT -from ...utils.torch_utils import get_device -from ..latent_diffusion.pipeline_latent_diffusion import LDMBertConfig, LDMBertModel -from ..paint_by_example import PaintByExampleImageEncoder -from ..pipeline_utils import DiffusionPipeline -from .safety_checker import StableDiffusionSafetyChecker -from .stable_unclip_image_normalizer import StableUnCLIPImageNormalizer - - -if is_accelerate_available(): - from accelerate import init_empty_weights - from accelerate.utils import set_module_tensor_to_device - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -def shave_segments(path, n_shave_prefix_segments=1): - """ - Removes segments. Positive values shave the first segments, negative shave the last segments. - """ - if n_shave_prefix_segments >= 0: - return ".".join(path.split(".")[n_shave_prefix_segments:]) - else: - return ".".join(path.split(".")[:n_shave_prefix_segments]) - - -def renew_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item.replace("in_layers.0", "norm1") - new_item = new_item.replace("in_layers.2", "conv1") - - new_item = new_item.replace("out_layers.0", "norm2") - new_item = new_item.replace("out_layers.3", "conv2") - - new_item = new_item.replace("emb_layers.1", "time_emb_proj") - new_item = new_item.replace("skip_connection", "conv_shortcut") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("nin_shortcut", "conv_shortcut") - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - # new_item = new_item.replace('norm.weight', 'group_norm.weight') - # new_item = new_item.replace('norm.bias', 'group_norm.bias') - - # new_item = new_item.replace('proj_out.weight', 'proj_attn.weight') - # new_item = new_item.replace('proj_out.bias', 'proj_attn.bias') - - # new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("norm.weight", "group_norm.weight") - new_item = new_item.replace("norm.bias", "group_norm.bias") - - new_item = new_item.replace("q.weight", "to_q.weight") - new_item = new_item.replace("q.bias", "to_q.bias") - - new_item = new_item.replace("k.weight", "to_k.weight") - new_item = new_item.replace("k.bias", "to_k.bias") - - new_item = new_item.replace("v.weight", "to_v.weight") - new_item = new_item.replace("v.bias", "to_v.bias") - - new_item = new_item.replace("proj_out.weight", "to_out.0.weight") - new_item = new_item.replace("proj_out.bias", "to_out.0.bias") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def assign_to_checkpoint( - paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None -): - """ - This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits - attention layers, and takes into account additional replacements that may arise. - - Assigns the weights to the new checkpoint. - """ - assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." - - # Splits the attention layers into three variables. - if attention_paths_to_split is not None: - for path, path_map in attention_paths_to_split.items(): - old_tensor = old_checkpoint[path] - channels = old_tensor.shape[0] // 3 - - target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) - - num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 - - old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) - query, key, value = old_tensor.split(channels // num_heads, dim=1) - - checkpoint[path_map["query"]] = query.reshape(target_shape) - checkpoint[path_map["key"]] = key.reshape(target_shape) - checkpoint[path_map["value"]] = value.reshape(target_shape) - - for path in paths: - new_path = path["new"] - - # These have already been assigned - if attention_paths_to_split is not None and new_path in attention_paths_to_split: - continue - - # Global renaming happens here - new_path = new_path.replace("middle_block.0", "mid_block.resnets.0") - new_path = new_path.replace("middle_block.1", "mid_block.attentions.0") - new_path = new_path.replace("middle_block.2", "mid_block.resnets.1") - - if additional_replacements is not None: - for replacement in additional_replacements: - new_path = new_path.replace(replacement["old"], replacement["new"]) - - # proj_attn.weight has to be converted from conv 1D to linear - is_attn_weight = "proj_attn.weight" in new_path or ("attentions" in new_path and "to_" in new_path) - shape = old_checkpoint[path["old"]].shape - if is_attn_weight and len(shape) == 3: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0] - elif is_attn_weight and len(shape) == 4: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0, 0] - else: - checkpoint[new_path] = old_checkpoint[path["old"]] - - -def conv_attn_to_linear(checkpoint): - keys = list(checkpoint.keys()) - attn_keys = ["query.weight", "key.weight", "value.weight"] - for key in keys: - if ".".join(key.split(".")[-2:]) in attn_keys: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0, 0] - elif "proj_attn.weight" in key: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0] - - -def create_unet_diffusers_config(original_config, image_size: int, controlnet=False): - """ - Creates a config for the diffusers based on the config of the LDM model. - """ - if controlnet: - unet_params = original_config["model"]["params"]["control_stage_config"]["params"] - else: - if ( - "unet_config" in original_config["model"]["params"] - and original_config["model"]["params"]["unet_config"] is not None - ): - unet_params = original_config["model"]["params"]["unet_config"]["params"] - else: - unet_params = original_config["model"]["params"]["network_config"]["params"] - - vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] - - block_out_channels = [unet_params["model_channels"] * mult for mult in unet_params["channel_mult"]] - - down_block_types = [] - resolution = 1 - for i in range(len(block_out_channels)): - block_type = "CrossAttnDownBlock2D" if resolution in unet_params["attention_resolutions"] else "DownBlock2D" - down_block_types.append(block_type) - if i != len(block_out_channels) - 1: - resolution *= 2 - - up_block_types = [] - for i in range(len(block_out_channels)): - block_type = "CrossAttnUpBlock2D" if resolution in unet_params["attention_resolutions"] else "UpBlock2D" - up_block_types.append(block_type) - resolution //= 2 - - if unet_params["transformer_depth"] is not None: - transformer_layers_per_block = ( - unet_params["transformer_depth"] - if isinstance(unet_params["transformer_depth"], int) - else list(unet_params["transformer_depth"]) - ) - else: - transformer_layers_per_block = 1 - - vae_scale_factor = 2 ** (len(vae_params["ch_mult"]) - 1) - - head_dim = unet_params["num_heads"] if "num_heads" in unet_params else None - use_linear_projection = ( - unet_params["use_linear_in_transformer"] if "use_linear_in_transformer" in unet_params else False - ) - if use_linear_projection: - # stable diffusion 2-base-512 and 2-768 - if head_dim is None: - head_dim_mult = unet_params["model_channels"] // unet_params["num_head_channels"] - head_dim = [head_dim_mult * c for c in list(unet_params["channel_mult"])] - - class_embed_type = None - addition_embed_type = None - addition_time_embed_dim = None - projection_class_embeddings_input_dim = None - context_dim = None - - if unet_params["context_dim"] is not None: - context_dim = ( - unet_params["context_dim"] - if isinstance(unet_params["context_dim"], int) - else unet_params["context_dim"][0] - ) - - if "num_classes" in unet_params: - if unet_params["num_classes"] == "sequential": - if context_dim in [2048, 1280]: - # SDXL - addition_embed_type = "text_time" - addition_time_embed_dim = 256 - else: - class_embed_type = "projection" - assert "adm_in_channels" in unet_params - projection_class_embeddings_input_dim = unet_params["adm_in_channels"] - - config = { - "sample_size": image_size // vae_scale_factor, - "in_channels": unet_params["in_channels"], - "down_block_types": tuple(down_block_types), - "block_out_channels": tuple(block_out_channels), - "layers_per_block": unet_params["num_res_blocks"], - "cross_attention_dim": context_dim, - "attention_head_dim": head_dim, - "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, - "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim, - "transformer_layers_per_block": transformer_layers_per_block, - } - - if "disable_self_attentions" in unet_params: - config["only_cross_attention"] = unet_params["disable_self_attentions"] - - if "num_classes" in unet_params and isinstance(unet_params["num_classes"], int): - config["num_class_embeds"] = unet_params["num_classes"] - - if controlnet: - config["conditioning_channels"] = unet_params["hint_channels"] - else: - config["out_channels"] = unet_params["out_channels"] - config["up_block_types"] = tuple(up_block_types) - - return config - - -def create_vae_diffusers_config(original_config, image_size: int): - """ - Creates a config for the diffusers based on the config of the LDM model. - """ - vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] - _ = original_config["model"]["params"]["first_stage_config"]["params"]["embed_dim"] - - block_out_channels = [vae_params["ch"] * mult for mult in vae_params["ch_mult"]] - down_block_types = [ - "DownEncoderBlock2D" if image_size // 2**i not in vae_params["attn_resolutions"] else "AttnDownEncoderBlock2D" - for i, _ in enumerate(block_out_channels) - ] - up_block_types = [ - "UpDecoderBlock2D" if image_size // 2**i not in vae_params["attn_resolutions"] else "AttnUpDecoderBlock2D" - for i, _ in enumerate(block_out_channels) - ][::-1] - - config = { - "sample_size": image_size, - "in_channels": vae_params["in_channels"], - "out_channels": vae_params["out_ch"], - "down_block_types": tuple(down_block_types), - "up_block_types": tuple(up_block_types), - "block_out_channels": tuple(block_out_channels), - "latent_channels": vae_params["z_channels"], - "layers_per_block": vae_params["num_res_blocks"], - } - return config - - -def create_diffusers_schedular(original_config): - schedular = DDIMScheduler( - num_train_timesteps=original_config["model"]["params"]["timesteps"], - beta_start=original_config["model"]["params"]["linear_start"], - beta_end=original_config["model"]["params"]["linear_end"], - beta_schedule="scaled_linear", - ) - return schedular - - -def create_ldm_bert_config(original_config): - bert_params = original_config["model"]["params"]["cond_stage_config"]["params"] - config = LDMBertConfig( - d_model=bert_params.n_embed, - encoder_layers=bert_params.n_layer, - encoder_ffn_dim=bert_params.n_embed * 4, - ) - return config - - -def convert_ldm_unet_checkpoint( - checkpoint, config, path=None, extract_ema=False, controlnet=False, skip_extract_state_dict=False -): - """ - Takes a state dict and a config, and returns a converted checkpoint. - """ - - if skip_extract_state_dict: - unet_state_dict = checkpoint - else: - # extract state_dict for UNet - unet_state_dict = {} - keys = list(checkpoint.keys()) - - if controlnet: - unet_key = "control_model." - else: - unet_key = "model.diffusion_model." - - # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA - if sum(k.startswith("model_ema") for k in keys) > 100 and extract_ema: - logger.warning(f"Checkpoint {path} has both EMA and non-EMA weights.") - logger.warning( - "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA" - " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag." - ) - for key in keys: - if key.startswith("model.diffusion_model"): - flat_ema_key = "model_ema." + "".join(key.split(".")[1:]) - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(flat_ema_key) - else: - if sum(k.startswith("model_ema") for k in keys) > 100: - logger.warning( - "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA" - " weights (usually better for inference), please make sure to add the `--extract_ema` flag." - ) - - for key in keys: - if key.startswith(unet_key): - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key) - - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"] - - if config["class_embed_type"] is None: - # No parameters to port - ... - elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection": - new_checkpoint["class_embedding.linear_1.weight"] = unet_state_dict["label_emb.0.0.weight"] - new_checkpoint["class_embedding.linear_1.bias"] = unet_state_dict["label_emb.0.0.bias"] - new_checkpoint["class_embedding.linear_2.weight"] = unet_state_dict["label_emb.0.2.weight"] - new_checkpoint["class_embedding.linear_2.bias"] = unet_state_dict["label_emb.0.2.bias"] - else: - raise NotImplementedError(f"Not implemented `class_embed_type`: {config['class_embed_type']}") - - if config["addition_embed_type"] == "text_time": - new_checkpoint["add_embedding.linear_1.weight"] = unet_state_dict["label_emb.0.0.weight"] - new_checkpoint["add_embedding.linear_1.bias"] = unet_state_dict["label_emb.0.0.bias"] - new_checkpoint["add_embedding.linear_2.weight"] = unet_state_dict["label_emb.0.2.weight"] - new_checkpoint["add_embedding.linear_2.bias"] = unet_state_dict["label_emb.0.2.bias"] - - # Relevant to StableDiffusionUpscalePipeline - if "num_class_embeds" in config: - if (config["num_class_embeds"] is not None) and ("label_emb.weight" in unet_state_dict): - new_checkpoint["class_embedding.weight"] = unet_state_dict["label_emb.weight"] - - new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"] - - if not controlnet: - new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"] - new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"] - new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key] - for layer_id in range(num_output_blocks) - } - - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - - if f"input_blocks.{i}.0.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.bias" - ) - - paths = renew_resnet_paths(resnets) - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(attentions): - paths = renew_attention_paths(attentions) - - meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - resnet_0 = middle_blocks[0] - attentions = middle_blocks[1] - resnet_1 = middle_blocks[2] - - resnet_0_paths = renew_resnet_paths(resnet_0) - assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config) - - resnet_1_paths = renew_resnet_paths(resnet_1) - assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config) - - attentions_paths = renew_attention_paths(attentions) - meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"} - assign_to_checkpoint( - attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - for i in range(num_output_blocks): - block_id = i // (config["layers_per_block"] + 1) - layer_in_block_id = i % (config["layers_per_block"] + 1) - output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]] - output_block_list = {} - - for layer in output_block_layers: - layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1) - if layer_id in output_block_list: - output_block_list[layer_id].append(layer_name) - else: - output_block_list[layer_id] = [layer_name] - - if len(output_block_list) > 1: - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key] - attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key] - - resnet_0_paths = renew_resnet_paths(resnets) - paths = renew_resnet_paths(resnets) - - meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - output_block_list = {k: sorted(v) for k, v in sorted(output_block_list.items())} - if ["conv.bias", "conv.weight"] in output_block_list.values(): - index = list(output_block_list.values()).index(["conv.bias", "conv.weight"]) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.bias" - ] - - # Clear attentions as they have been attributed above. - if len(attentions) == 2: - attentions = [] - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = { - "old": f"output_blocks.{i}.1", - "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - else: - resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1) - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]]) - - new_checkpoint[new_path] = unet_state_dict[old_path] - - if controlnet: - # conditioning embedding - - orig_index = 0 - - new_checkpoint["controlnet_cond_embedding.conv_in.weight"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.weight" - ) - new_checkpoint["controlnet_cond_embedding.conv_in.bias"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.bias" - ) - - orig_index += 2 - - diffusers_index = 0 - - while diffusers_index < 6: - new_checkpoint[f"controlnet_cond_embedding.blocks.{diffusers_index}.weight"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.weight" - ) - new_checkpoint[f"controlnet_cond_embedding.blocks.{diffusers_index}.bias"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.bias" - ) - diffusers_index += 1 - orig_index += 2 - - new_checkpoint["controlnet_cond_embedding.conv_out.weight"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.weight" - ) - new_checkpoint["controlnet_cond_embedding.conv_out.bias"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.bias" - ) - - # down blocks - for i in range(num_input_blocks): - new_checkpoint[f"controlnet_down_blocks.{i}.weight"] = unet_state_dict.pop(f"zero_convs.{i}.0.weight") - new_checkpoint[f"controlnet_down_blocks.{i}.bias"] = unet_state_dict.pop(f"zero_convs.{i}.0.bias") - - # mid block - new_checkpoint["controlnet_mid_block.weight"] = unet_state_dict.pop("middle_block_out.0.weight") - new_checkpoint["controlnet_mid_block.bias"] = unet_state_dict.pop("middle_block_out.0.bias") - - return new_checkpoint - - -def convert_ldm_vae_checkpoint(checkpoint, config): - # extract state dict for VAE - vae_state_dict = {} - keys = list(checkpoint.keys()) - vae_key = "first_stage_model." if any(k.startswith("first_stage_model.") for k in keys) else "" - for key in keys: - if key.startswith(vae_key): - vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - - new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] - new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] - new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] - new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] - new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] - new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] - - new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] - new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] - new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] - new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] - new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] - new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] - - new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] - new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] - new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] - new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] - - # Retrieves the keys for the encoder down blocks only - num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) - down_blocks = { - layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - # Retrieves the keys for the decoder up blocks only - num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) - up_blocks = { - layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) - } - - for i in range(num_down_blocks): - resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] - - if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.weight" - ) - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.bias" - ) - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - resnets = [ - key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key - ] - - if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.bias" - ] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - return new_checkpoint - - -def convert_ldm_bert_checkpoint(checkpoint, config): - def _copy_attn_layer(hf_attn_layer, pt_attn_layer): - hf_attn_layer.q_proj.weight.data = pt_attn_layer.to_q.weight - hf_attn_layer.k_proj.weight.data = pt_attn_layer.to_k.weight - hf_attn_layer.v_proj.weight.data = pt_attn_layer.to_v.weight - - hf_attn_layer.out_proj.weight = pt_attn_layer.to_out.weight - hf_attn_layer.out_proj.bias = pt_attn_layer.to_out.bias - - def _copy_linear(hf_linear, pt_linear): - hf_linear.weight = pt_linear.weight - hf_linear.bias = pt_linear.bias - - def _copy_layer(hf_layer, pt_layer): - # copy layer norms - _copy_linear(hf_layer.self_attn_layer_norm, pt_layer[0][0]) - _copy_linear(hf_layer.final_layer_norm, pt_layer[1][0]) - - # copy attn - _copy_attn_layer(hf_layer.self_attn, pt_layer[0][1]) - - # copy MLP - pt_mlp = pt_layer[1][1] - _copy_linear(hf_layer.fc1, pt_mlp.net[0][0]) - _copy_linear(hf_layer.fc2, pt_mlp.net[2]) - - def _copy_layers(hf_layers, pt_layers): - for i, hf_layer in enumerate(hf_layers): - if i != 0: - i += i - pt_layer = pt_layers[i : i + 2] - _copy_layer(hf_layer, pt_layer) - - hf_model = LDMBertModel(config).eval() - - # copy embeds - hf_model.model.embed_tokens.weight = checkpoint.transformer.token_emb.weight - hf_model.model.embed_positions.weight.data = checkpoint.transformer.pos_emb.emb.weight - - # copy layer norm - _copy_linear(hf_model.model.layer_norm, checkpoint.transformer.norm) - - # copy hidden layers - _copy_layers(hf_model.model.layers, checkpoint.transformer.attn_layers.layers) - - _copy_linear(hf_model.to_logits, checkpoint.transformer.to_logits) - - return hf_model - - -def convert_ldm_clip_checkpoint(checkpoint, local_files_only=False, text_encoder=None): - if text_encoder is None: - config_name = "openai/clip-vit-large-patch14" - try: - config = CLIPTextConfig.from_pretrained(config_name, local_files_only=local_files_only) - except Exception: - raise ValueError( - f"With local_files_only set to {local_files_only}, you must first locally save the configuration in the following path: 'openai/clip-vit-large-patch14'." - ) - - ctx = init_empty_weights if is_accelerate_available() else nullcontext - with ctx(): - text_model = CLIPTextModel(config) - else: - text_model = text_encoder - - keys = list(checkpoint.keys()) - - text_model_dict = {} - - remove_prefixes = ["cond_stage_model.transformer", "conditioner.embedders.0.transformer"] - - for key in keys: - for prefix in remove_prefixes: - if key.startswith(prefix): - text_model_dict[key[len(prefix + ".") :]] = checkpoint[key] - - if is_accelerate_available(): - for param_name, param in text_model_dict.items(): - set_module_tensor_to_device(text_model, param_name, "cpu", value=param) - else: - if not (hasattr(text_model, "embeddings") and hasattr(text_model.embeddings.position_ids)): - text_model_dict.pop("text_model.embeddings.position_ids", None) - - text_model.load_state_dict(text_model_dict) - - return text_model - - -textenc_conversion_lst = [ - ("positional_embedding", "text_model.embeddings.position_embedding.weight"), - ("token_embedding.weight", "text_model.embeddings.token_embedding.weight"), - ("ln_final.weight", "text_model.final_layer_norm.weight"), - ("ln_final.bias", "text_model.final_layer_norm.bias"), - ("text_projection", "text_projection.weight"), -] -textenc_conversion_map = {x[0]: x[1] for x in textenc_conversion_lst} - -textenc_transformer_conversion_lst = [ - # (stable-diffusion, HF Diffusers) - ("resblocks.", "text_model.encoder.layers."), - ("ln_1", "layer_norm1"), - ("ln_2", "layer_norm2"), - (".c_fc.", ".fc1."), - (".c_proj.", ".fc2."), - (".attn", ".self_attn"), - ("ln_final.", "transformer.text_model.final_layer_norm."), - ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"), - ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"), -] -protected = {re.escape(x[0]): x[1] for x in textenc_transformer_conversion_lst} -textenc_pattern = re.compile("|".join(protected.keys())) - - -def convert_paint_by_example_checkpoint(checkpoint, local_files_only=False): - config = CLIPVisionConfig.from_pretrained("openai/clip-vit-large-patch14", local_files_only=local_files_only) - model = PaintByExampleImageEncoder(config) - - keys = list(checkpoint.keys()) - - text_model_dict = {} - - for key in keys: - if key.startswith("cond_stage_model.transformer"): - text_model_dict[key[len("cond_stage_model.transformer.") :]] = checkpoint[key] - - # load clip vision - model.model.load_state_dict(text_model_dict) - - # load mapper - keys_mapper = { - k[len("cond_stage_model.mapper.res") :]: v - for k, v in checkpoint.items() - if k.startswith("cond_stage_model.mapper") - } - - MAPPING = { - "attn.c_qkv": ["attn1.to_q", "attn1.to_k", "attn1.to_v"], - "attn.c_proj": ["attn1.to_out.0"], - "ln_1": ["norm1"], - "ln_2": ["norm3"], - "mlp.c_fc": ["ff.net.0.proj"], - "mlp.c_proj": ["ff.net.2"], - } - - mapped_weights = {} - for key, value in keys_mapper.items(): - prefix = key[: len("blocks.i")] - suffix = key.split(prefix)[-1].split(".")[-1] - name = key.split(prefix)[-1].split(suffix)[0][1:-1] - mapped_names = MAPPING[name] - - num_splits = len(mapped_names) - for i, mapped_name in enumerate(mapped_names): - new_name = ".".join([prefix, mapped_name, suffix]) - shape = value.shape[0] // num_splits - mapped_weights[new_name] = value[i * shape : (i + 1) * shape] - - model.mapper.load_state_dict(mapped_weights) - - # load final layer norm - model.final_layer_norm.load_state_dict( - { - "bias": checkpoint["cond_stage_model.final_ln.bias"], - "weight": checkpoint["cond_stage_model.final_ln.weight"], - } - ) - - # load final proj - model.proj_out.load_state_dict( - { - "bias": checkpoint["proj_out.bias"], - "weight": checkpoint["proj_out.weight"], - } - ) - - # load uncond vector - model.uncond_vector.data = torch.nn.Parameter(checkpoint["learnable_vector"]) - return model - - -def convert_open_clip_checkpoint( - checkpoint, - config_name, - prefix="cond_stage_model.model.", - has_projection=False, - local_files_only=False, - **config_kwargs, -): - # text_model = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-2", subfolder="text_encoder") - # text_model = CLIPTextModelWithProjection.from_pretrained( - # "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", projection_dim=1280 - # ) - try: - config = CLIPTextConfig.from_pretrained(config_name, **config_kwargs, local_files_only=local_files_only) - except Exception: - raise ValueError( - f"With local_files_only set to {local_files_only}, you must first locally save the configuration in the following path: '{config_name}'." - ) - - ctx = init_empty_weights if is_accelerate_available() else nullcontext - with ctx(): - text_model = CLIPTextModelWithProjection(config) if has_projection else CLIPTextModel(config) - - keys = list(checkpoint.keys()) - - keys_to_ignore = [] - if config_name == "stabilityai/stable-diffusion-2" and config.num_hidden_layers == 23: - # make sure to remove all keys > 22 - keys_to_ignore += [k for k in keys if k.startswith("cond_stage_model.model.transformer.resblocks.23")] - keys_to_ignore += ["cond_stage_model.model.text_projection"] - - text_model_dict = {} - - if prefix + "text_projection" in checkpoint: - d_model = int(checkpoint[prefix + "text_projection"].shape[0]) - else: - d_model = 1024 - - text_model_dict["text_model.embeddings.position_ids"] = text_model.text_model.embeddings.get_buffer("position_ids") - - for key in keys: - if key in keys_to_ignore: - continue - if key[len(prefix) :] in textenc_conversion_map: - if key.endswith("text_projection"): - value = checkpoint[key].T.contiguous() - else: - value = checkpoint[key] - - text_model_dict[textenc_conversion_map[key[len(prefix) :]]] = value - - if key.startswith(prefix + "transformer."): - new_key = key[len(prefix + "transformer.") :] - if new_key.endswith(".in_proj_weight"): - new_key = new_key[: -len(".in_proj_weight")] - new_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], new_key) - text_model_dict[new_key + ".q_proj.weight"] = checkpoint[key][:d_model, :] - text_model_dict[new_key + ".k_proj.weight"] = checkpoint[key][d_model : d_model * 2, :] - text_model_dict[new_key + ".v_proj.weight"] = checkpoint[key][d_model * 2 :, :] - elif new_key.endswith(".in_proj_bias"): - new_key = new_key[: -len(".in_proj_bias")] - new_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], new_key) - text_model_dict[new_key + ".q_proj.bias"] = checkpoint[key][:d_model] - text_model_dict[new_key + ".k_proj.bias"] = checkpoint[key][d_model : d_model * 2] - text_model_dict[new_key + ".v_proj.bias"] = checkpoint[key][d_model * 2 :] - else: - new_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], new_key) - - text_model_dict[new_key] = checkpoint[key] - - if is_accelerate_available(): - for param_name, param in text_model_dict.items(): - set_module_tensor_to_device(text_model, param_name, "cpu", value=param) - else: - if not (hasattr(text_model, "embeddings") and hasattr(text_model.embeddings.position_ids)): - text_model_dict.pop("text_model.embeddings.position_ids", None) - - text_model.load_state_dict(text_model_dict) - - return text_model - - -def stable_unclip_image_encoder(original_config, local_files_only=False): - """ - Returns the image processor and clip image encoder for the img2img unclip pipeline. - - We currently know of two types of stable unclip models which separately use the clip and the openclip image - encoders. - """ - - image_embedder_config = original_config["model"]["params"]["embedder_config"] - - sd_clip_image_embedder_class = image_embedder_config["target"] - sd_clip_image_embedder_class = sd_clip_image_embedder_class.split(".")[-1] - - if sd_clip_image_embedder_class == "ClipImageEmbedder": - clip_model_name = image_embedder_config.params.model - - if clip_model_name == "ViT-L/14": - feature_extractor = CLIPImageProcessor() - image_encoder = CLIPVisionModelWithProjection.from_pretrained( - "openai/clip-vit-large-patch14", local_files_only=local_files_only - ) - else: - raise NotImplementedError(f"Unknown CLIP checkpoint name in stable diffusion checkpoint {clip_model_name}") - - elif sd_clip_image_embedder_class == "FrozenOpenCLIPImageEmbedder": - feature_extractor = CLIPImageProcessor() - image_encoder = CLIPVisionModelWithProjection.from_pretrained( - "laion/CLIP-ViT-H-14-laion2B-s32B-b79K", local_files_only=local_files_only - ) - else: - raise NotImplementedError( - f"Unknown CLIP image embedder class in stable diffusion checkpoint {sd_clip_image_embedder_class}" - ) - - return feature_extractor, image_encoder - - -def stable_unclip_image_noising_components( - original_config, clip_stats_path: Optional[str] = None, device: Optional[str] = None -): - """ - Returns the noising components for the img2img and txt2img unclip pipelines. - - Converts the stability noise augmentor into - 1. a `StableUnCLIPImageNormalizer` for holding the CLIP stats - 2. a `DDPMScheduler` for holding the noise schedule - - If the noise augmentor config specifies a clip stats path, the `clip_stats_path` must be provided. - """ - noise_aug_config = original_config["model"]["params"]["noise_aug_config"] - noise_aug_class = noise_aug_config["target"] - noise_aug_class = noise_aug_class.split(".")[-1] - - if noise_aug_class == "CLIPEmbeddingNoiseAugmentation": - noise_aug_config = noise_aug_config.params - embedding_dim = noise_aug_config.timestep_dim - max_noise_level = noise_aug_config.noise_schedule_config.timesteps - beta_schedule = noise_aug_config.noise_schedule_config.beta_schedule - - image_normalizer = StableUnCLIPImageNormalizer(embedding_dim=embedding_dim) - image_noising_scheduler = DDPMScheduler(num_train_timesteps=max_noise_level, beta_schedule=beta_schedule) - - if "clip_stats_path" in noise_aug_config: - if clip_stats_path is None: - raise ValueError("This stable unclip config requires a `clip_stats_path`") - - clip_mean, clip_std = torch.load(clip_stats_path, map_location=device) - clip_mean = clip_mean[None, :] - clip_std = clip_std[None, :] - - clip_stats_state_dict = { - "mean": clip_mean, - "std": clip_std, - } - - image_normalizer.load_state_dict(clip_stats_state_dict) - else: - raise NotImplementedError(f"Unknown noise augmentor class: {noise_aug_class}") - - return image_normalizer, image_noising_scheduler - - -def convert_controlnet_checkpoint( - checkpoint, - original_config, - checkpoint_path, - image_size, - upcast_attention, - extract_ema, - use_linear_projection=None, - cross_attention_dim=None, -): - ctrlnet_config = create_unet_diffusers_config(original_config, image_size=image_size, controlnet=True) - ctrlnet_config["upcast_attention"] = upcast_attention - - ctrlnet_config.pop("sample_size") - - if use_linear_projection is not None: - ctrlnet_config["use_linear_projection"] = use_linear_projection - - if cross_attention_dim is not None: - ctrlnet_config["cross_attention_dim"] = cross_attention_dim - - ctx = init_empty_weights if is_accelerate_available() else nullcontext - with ctx(): - controlnet = ControlNetModel(**ctrlnet_config) - - # Some controlnet ckpt files are distributed independently from the rest of the - # model components i.e. https://huggingface.co/thibaud/controlnet-sd21/ - if "time_embed.0.weight" in checkpoint: - skip_extract_state_dict = True - else: - skip_extract_state_dict = False - - converted_ctrl_checkpoint = convert_ldm_unet_checkpoint( - checkpoint, - ctrlnet_config, - path=checkpoint_path, - extract_ema=extract_ema, - controlnet=True, - skip_extract_state_dict=skip_extract_state_dict, - ) - - if is_accelerate_available(): - for param_name, param in converted_ctrl_checkpoint.items(): - set_module_tensor_to_device(controlnet, param_name, "cpu", value=param) - else: - controlnet.load_state_dict(converted_ctrl_checkpoint) - - return controlnet - - -def download_from_original_stable_diffusion_ckpt( - checkpoint_path_or_dict: Union[str, Dict[str, torch.Tensor]], - original_config_file: str = None, - image_size: Optional[int] = 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, - 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, - load_safety_checker: bool = True, - safety_checker: Optional[StableDiffusionSafetyChecker] = None, - feature_extractor: Optional[AutoFeatureExtractor] = None, - pipeline_class: DiffusionPipeline = None, - local_files_only=False, - vae_path=None, - vae=None, - text_encoder=None, - text_encoder_2=None, - tokenizer=None, - tokenizer_2=None, - config_files=None, -) -> DiffusionPipeline: - """ - Load a Stable Diffusion pipeline object from a CompVis-style `.ckpt`/`.safetensors` file and (ideally) a `.yaml` - config file. - - Although many of the arguments can be automatically inferred, some of these rely on brittle checks against the - global step count, which will likely fail for models that have undergone further fine-tuning. Therefore, it is - recommended that you override the default values and/or supply an `original_config_file` wherever possible. - - Args: - checkpoint_path_or_dict (`str` or `dict`): Path to `.ckpt` file, or the state dict. - original_config_file (`str`): - Path to `.yaml` config file corresponding to the original architecture. If `None`, will be automatically - inferred by looking for a key that only exists in SD2.0 models. - image_size (`int`, *optional*, defaults to 512): - The image size that the model was trained on. Use 512 for Stable Diffusion v1.X and Stable Diffusion v2 - Base. Use 768 for Stable Diffusion v2. - prediction_type (`str`, *optional*): - The prediction type that the model was trained on. Use `'epsilon'` for Stable Diffusion v1.X and Stable - Diffusion v2 Base. Use `'v_prediction'` for Stable Diffusion v2. - num_in_channels (`int`, *optional*, defaults to None): - The number of input channels. If `None`, it will be automatically inferred. - scheduler_type (`str`, *optional*, defaults to 'pndm'): - Type of scheduler to use. Should be one of `["pndm", "lms", "heun", "euler", "euler-ancestral", "dpm", - "ddim"]`. - model_type (`str`, *optional*, defaults to `None`): - The pipeline type. `None` to automatically infer, or one of `["FrozenOpenCLIPEmbedder", - "FrozenCLIPEmbedder", "PaintByExample"]`. - is_img2img (`bool`, *optional*, defaults to `False`): - Whether the model should be loaded as an img2img pipeline. - extract_ema (`bool`, *optional*, defaults to `False`): Only relevant for - checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights or not. Defaults to - `False`. Pass `True` to extract the EMA weights. EMA weights usually yield higher quality images for - inference. Non-EMA weights are usually better to continue fine-tuning. - upcast_attention (`bool`, *optional*, defaults to `None`): - Whether the attention computation should always be upcasted. This is necessary when running stable - diffusion 2.1. - device (`str`, *optional*, defaults to `None`): - The device to use. Pass `None` to determine automatically. - from_safetensors (`str`, *optional*, defaults to `False`): - If `checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch. - load_safety_checker (`bool`, *optional*, defaults to `True`): - Whether to load the safety checker or not. Defaults to `True`. - safety_checker (`StableDiffusionSafetyChecker`, *optional*, defaults to `None`): - Safety checker to use. If this parameter is `None`, the function will load a new instance of - [StableDiffusionSafetyChecker] by itself, if needed. - feature_extractor (`AutoFeatureExtractor`, *optional*, defaults to `None`): - Feature extractor to use. If this parameter is `None`, the function will load a new instance of - [AutoFeatureExtractor] by itself, if needed. - pipeline_class (`str`, *optional*, defaults to `None`): - The pipeline class to use. Pass `None` to determine automatically. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether or not to only look at local files (i.e., do not try to download the model). - vae (`AutoencoderKL`, *optional*, defaults to `None`): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. If - this parameter is `None`, the function will load a new instance of [CLIP] by itself, if needed. - text_encoder (`CLIPTextModel`, *optional*, defaults to `None`): - An instance of [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel) - to use, specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) - variant. If this parameter is `None`, the function will load a new instance of [CLIP] by itself, if needed. - tokenizer (`CLIPTokenizer`, *optional*, defaults to `None`): - An instance of - [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`): - 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 - - `v2`: Config file for Stable Diffusion v2 - - `xl`: Config file for Stable Diffusion XL - - `xl_refiner`: Config file for Stable Diffusion XL Refiner - return: A StableDiffusionPipeline object representing the passed-in `.ckpt`/`.safetensors` file. - """ - - # import pipelines here to avoid circular import error when using from_single_file method - from diffusers import ( - LDMTextToImagePipeline, - PaintByExamplePipeline, - StableDiffusionControlNetPipeline, - StableDiffusionInpaintPipeline, - StableDiffusionPipeline, - StableDiffusionUpscalePipeline, - StableDiffusionXLControlNetInpaintPipeline, - StableDiffusionXLImg2ImgPipeline, - StableDiffusionXLInpaintPipeline, - StableDiffusionXLPipeline, - StableUnCLIPImg2ImgPipeline, - StableUnCLIPPipeline, - ) - - if prediction_type == "v-prediction": - prediction_type = "v_prediction" - - if isinstance(checkpoint_path_or_dict, str): - if from_safetensors: - from safetensors.torch import load_file as safe_load - - checkpoint = safe_load(checkpoint_path_or_dict, device="cpu") - else: - if device is None: - device = get_device() - checkpoint = torch.load(checkpoint_path_or_dict, map_location=device) - else: - checkpoint = torch.load(checkpoint_path_or_dict, map_location=device) - elif isinstance(checkpoint_path_or_dict, dict): - checkpoint = checkpoint_path_or_dict - - # Sometimes models don't have the global_step item - if "global_step" in checkpoint: - global_step = checkpoint["global_step"] - else: - logger.debug("global_step key not found in model") - global_step = None - - # NOTE: this while loop isn't great but this controlnet checkpoint has one additional - # "state_dict" key https://huggingface.co/thibaud/controlnet-canny-sd21 - while "state_dict" in checkpoint: - checkpoint = checkpoint["state_dict"] - - if original_config_file is None: - key_name_v2_1 = "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_k.weight" - key_name_sd_xl_base = "conditioner.embedders.1.model.transformer.resblocks.9.mlp.c_proj.bias" - key_name_sd_xl_refiner = "conditioner.embedders.0.model.transformer.resblocks.9.mlp.c_proj.bias" - is_upscale = pipeline_class == StableDiffusionUpscalePipeline - - config_url = None - - # model_type = "v1" - if config_files is not None and "v1" in config_files: - original_config_file = config_files["v1"] - else: - config_url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml" - - if key_name_v2_1 in checkpoint and checkpoint[key_name_v2_1].shape[-1] == 1024: - # model_type = "v2" - if config_files is not None and "v2" in config_files: - original_config_file = config_files["v2"] - else: - config_url = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/v2-inference-v.yaml" - if global_step == 110000: - # v2.1 needs to upcast attention - upcast_attention = True - elif key_name_sd_xl_base in checkpoint: - # only base xl has two text embedders - if config_files is not None and "xl" in config_files: - original_config_file = config_files["xl"] - else: - config_url = "https://raw.githubusercontent.com/Stability-AI/generative-models/main/configs/inference/sd_xl_base.yaml" - elif key_name_sd_xl_refiner in checkpoint: - # only refiner xl has embedder and one text embedders - if config_files is not None and "xl_refiner" in config_files: - original_config_file = config_files["xl_refiner"] - else: - config_url = "https://raw.githubusercontent.com/Stability-AI/generative-models/main/configs/inference/sd_xl_refiner.yaml" - - if is_upscale: - config_url = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/x4-upscaling.yaml" - - if config_url is not None: - original_config_file = BytesIO(requests.get(config_url, timeout=DIFFUSERS_REQUEST_TIMEOUT).content) - else: - with open(original_config_file, "r") as f: - original_config_file = f.read() - else: - with open(original_config_file, "r") as f: - original_config_file = f.read() - - original_config = yaml.safe_load(original_config_file) - - # Convert the text model. - if ( - model_type is None - and "cond_stage_config" in original_config["model"]["params"] - and original_config["model"]["params"]["cond_stage_config"] is not None - ): - model_type = original_config["model"]["params"]["cond_stage_config"]["target"].split(".")[-1] - logger.debug(f"no `model_type` given, `model_type` inferred as: {model_type}") - elif model_type is None and original_config["model"]["params"]["network_config"] is not None: - if original_config["model"]["params"]["network_config"]["params"]["context_dim"] == 2048: - model_type = "SDXL" - else: - model_type = "SDXL-Refiner" - if image_size is None: - image_size = 1024 - - if pipeline_class is None: - # Check if we have a SDXL or SD model and initialize default pipeline - if model_type not in ["SDXL", "SDXL-Refiner"]: - pipeline_class = StableDiffusionPipeline if not controlnet else StableDiffusionControlNetPipeline - else: - pipeline_class = StableDiffusionXLPipeline if model_type == "SDXL" else StableDiffusionXLImg2ImgPipeline - - if num_in_channels is None and pipeline_class in [ - StableDiffusionInpaintPipeline, - StableDiffusionXLInpaintPipeline, - StableDiffusionXLControlNetInpaintPipeline, - ]: - num_in_channels = 9 - if num_in_channels is None and pipeline_class == StableDiffusionUpscalePipeline: - num_in_channels = 7 - elif num_in_channels is None: - num_in_channels = 4 - - if "unet_config" in original_config["model"]["params"]: - original_config["model"]["params"]["unet_config"]["params"]["in_channels"] = num_in_channels - elif "network_config" in original_config["model"]["params"]: - original_config["model"]["params"]["network_config"]["params"]["in_channels"] = num_in_channels - - if ( - "parameterization" in original_config["model"]["params"] - and original_config["model"]["params"]["parameterization"] == "v" - ): - if prediction_type is None: - # NOTE: For stable diffusion 2 base it is recommended to pass `prediction_type=="epsilon"` - # as it relies on a brittle global step parameter here - prediction_type = "epsilon" if global_step == 875000 else "v_prediction" - if image_size is None: - # NOTE: For stable diffusion 2 base one has to pass `image_size==512` - # as it relies on a brittle global step parameter here - image_size = 512 if global_step == 875000 else 768 - else: - if prediction_type is None: - prediction_type = "epsilon" - if image_size is None: - image_size = 512 - - if controlnet is None and "control_stage_config" in original_config["model"]["params"]: - path = checkpoint_path_or_dict if isinstance(checkpoint_path_or_dict, str) else "" - controlnet = convert_controlnet_checkpoint( - checkpoint, original_config, path, image_size, upcast_attention, extract_ema - ) - - if "timesteps" in original_config["model"]["params"]: - num_train_timesteps = original_config["model"]["params"]["timesteps"] - else: - num_train_timesteps = 1000 - - if model_type in ["SDXL", "SDXL-Refiner"]: - scheduler_dict = { - "beta_schedule": "scaled_linear", - "beta_start": 0.00085, - "beta_end": 0.012, - "interpolation_type": "linear", - "num_train_timesteps": num_train_timesteps, - "prediction_type": "epsilon", - "sample_max_value": 1.0, - "set_alpha_to_one": False, - "skip_prk_steps": True, - "steps_offset": 1, - "timestep_spacing": "leading", - } - scheduler = EulerDiscreteScheduler.from_config(scheduler_dict) - scheduler_type = "euler" - else: - if "linear_start" in original_config["model"]["params"]: - beta_start = original_config["model"]["params"]["linear_start"] - else: - beta_start = 0.02 - - if "linear_end" in original_config["model"]["params"]: - beta_end = original_config["model"]["params"]["linear_end"] - else: - beta_end = 0.085 - scheduler = DDIMScheduler( - beta_end=beta_end, - beta_schedule="scaled_linear", - beta_start=beta_start, - num_train_timesteps=num_train_timesteps, - steps_offset=1, - clip_sample=False, - set_alpha_to_one=False, - prediction_type=prediction_type, - ) - # make sure scheduler works correctly with DDIM - scheduler.register_to_config(clip_sample=False) - - if scheduler_type == "pndm": - config = dict(scheduler.config) - config["skip_prk_steps"] = True - scheduler = PNDMScheduler.from_config(config) - elif scheduler_type == "lms": - scheduler = LMSDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "heun": - scheduler = HeunDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "euler": - scheduler = EulerDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "euler-ancestral": - scheduler = EulerAncestralDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "dpm": - scheduler = DPMSolverMultistepScheduler.from_config(scheduler.config) - elif scheduler_type == "ddim": - scheduler = scheduler - else: - raise ValueError(f"Scheduler of type {scheduler_type} doesn't exist!") - - if pipeline_class == StableDiffusionUpscalePipeline: - image_size = original_config["model"]["params"]["unet_config"]["params"]["image_size"] - - # Convert the UNet2DConditionModel model. - unet_config = create_unet_diffusers_config(original_config, image_size=image_size) - unet_config["upcast_attention"] = upcast_attention - - path = checkpoint_path_or_dict if isinstance(checkpoint_path_or_dict, str) else "" - converted_unet_checkpoint = convert_ldm_unet_checkpoint( - checkpoint, unet_config, path=path, extract_ema=extract_ema - ) - - ctx = init_empty_weights if is_accelerate_available() else nullcontext - with ctx(): - unet = UNet2DConditionModel(**unet_config) - - if is_accelerate_available(): - if model_type not in ["SDXL", "SDXL-Refiner"]: # SBM Delay this. - for param_name, param in converted_unet_checkpoint.items(): - set_module_tensor_to_device(unet, param_name, "cpu", value=param) - else: - unet.load_state_dict(converted_unet_checkpoint) - - # Convert the VAE model. - if vae_path is None and vae is None: - vae_config = create_vae_diffusers_config(original_config, image_size=image_size) - converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config) - - if ( - "model" in original_config - and "params" in original_config["model"] - and "scale_factor" in original_config["model"]["params"] - ): - vae_scaling_factor = original_config["model"]["params"]["scale_factor"] - else: - vae_scaling_factor = 0.18215 # default SD scaling factor - - vae_config["scaling_factor"] = vae_scaling_factor - - ctx = init_empty_weights if is_accelerate_available() else nullcontext - with ctx(): - vae = AutoencoderKL(**vae_config) - - if is_accelerate_available(): - for param_name, param in converted_vae_checkpoint.items(): - set_module_tensor_to_device(vae, param_name, "cpu", value=param) - else: - vae.load_state_dict(converted_vae_checkpoint) - elif vae is None: - vae = AutoencoderKL.from_pretrained(vae_path, local_files_only=local_files_only) - - if model_type == "FrozenOpenCLIPEmbedder": - config_name = "stabilityai/stable-diffusion-2" - config_kwargs = {"subfolder": "text_encoder"} - - if text_encoder is None: - text_model = convert_open_clip_checkpoint( - checkpoint, config_name, local_files_only=local_files_only, **config_kwargs - ) - else: - text_model = text_encoder - - try: - tokenizer = CLIPTokenizer.from_pretrained( - "stabilityai/stable-diffusion-2", subfolder="tokenizer", local_files_only=local_files_only - ) - except Exception: - raise ValueError( - f"With local_files_only set to {local_files_only}, you must first locally save the tokenizer in the following path: 'stabilityai/stable-diffusion-2'." - ) - - if stable_unclip is None: - if controlnet: - pipe = pipeline_class( - vae=vae, - text_encoder=text_model, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - controlnet=controlnet, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) - if hasattr(pipe, "requires_safety_checker"): - pipe.requires_safety_checker = False - - elif pipeline_class == StableDiffusionUpscalePipeline: - scheduler = DDIMScheduler.from_pretrained( - "stabilityai/stable-diffusion-x4-upscaler", subfolder="scheduler" - ) - low_res_scheduler = DDPMScheduler.from_pretrained( - "stabilityai/stable-diffusion-x4-upscaler", subfolder="low_res_scheduler" - ) - - pipe = pipeline_class( - vae=vae, - text_encoder=text_model, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - low_res_scheduler=low_res_scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) - - else: - pipe = pipeline_class( - vae=vae, - text_encoder=text_model, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) - if hasattr(pipe, "requires_safety_checker"): - pipe.requires_safety_checker = False - - else: - image_normalizer, image_noising_scheduler = stable_unclip_image_noising_components( - original_config, clip_stats_path=clip_stats_path, device=device - ) - - if stable_unclip == "img2img": - feature_extractor, image_encoder = stable_unclip_image_encoder(original_config) - - pipe = StableUnCLIPImg2ImgPipeline( - # image encoding components - feature_extractor=feature_extractor, - image_encoder=image_encoder, - # image noising components - image_normalizer=image_normalizer, - image_noising_scheduler=image_noising_scheduler, - # regular denoising components - tokenizer=tokenizer, - text_encoder=text_model, - unet=unet, - scheduler=scheduler, - # vae - vae=vae, - ) - elif stable_unclip == "txt2img": - if stable_unclip_prior is None or stable_unclip_prior == "karlo": - karlo_model = "kakaobrain/karlo-v1-alpha" - prior = PriorTransformer.from_pretrained( - karlo_model, subfolder="prior", local_files_only=local_files_only - ) - - try: - prior_tokenizer = CLIPTokenizer.from_pretrained( - "openai/clip-vit-large-patch14", local_files_only=local_files_only - ) - except Exception: - raise ValueError( - f"With local_files_only set to {local_files_only}, you must first locally save the tokenizer in the following path: 'openai/clip-vit-large-patch14'." - ) - prior_text_model = CLIPTextModelWithProjection.from_pretrained( - "openai/clip-vit-large-patch14", local_files_only=local_files_only - ) - - prior_scheduler = UnCLIPScheduler.from_pretrained( - karlo_model, subfolder="prior_scheduler", local_files_only=local_files_only - ) - prior_scheduler = DDPMScheduler.from_config(prior_scheduler.config) - else: - raise NotImplementedError(f"unknown prior for stable unclip model: {stable_unclip_prior}") - - pipe = StableUnCLIPPipeline( - # prior components - prior_tokenizer=prior_tokenizer, - prior_text_encoder=prior_text_model, - prior=prior, - prior_scheduler=prior_scheduler, - # image noising components - image_normalizer=image_normalizer, - image_noising_scheduler=image_noising_scheduler, - # regular denoising components - tokenizer=tokenizer, - text_encoder=text_model, - unet=unet, - scheduler=scheduler, - # vae - vae=vae, - ) - else: - raise NotImplementedError(f"unknown `stable_unclip` type: {stable_unclip}") - elif model_type == "PaintByExample": - vision_model = convert_paint_by_example_checkpoint(checkpoint) - try: - tokenizer = CLIPTokenizer.from_pretrained( - "openai/clip-vit-large-patch14", local_files_only=local_files_only - ) - except Exception: - raise ValueError( - f"With local_files_only set to {local_files_only}, you must first locally save the tokenizer in the following path: 'openai/clip-vit-large-patch14'." - ) - try: - feature_extractor = AutoFeatureExtractor.from_pretrained( - "CompVis/stable-diffusion-safety-checker", local_files_only=local_files_only - ) - except Exception: - raise ValueError( - f"With local_files_only set to {local_files_only}, you must first locally save the feature_extractor in the following path: 'CompVis/stable-diffusion-safety-checker'." - ) - pipe = PaintByExamplePipeline( - vae=vae, - image_encoder=vision_model, - unet=unet, - scheduler=scheduler, - safety_checker=None, - feature_extractor=feature_extractor, - ) - elif model_type == "FrozenCLIPEmbedder": - text_model = convert_ldm_clip_checkpoint( - checkpoint, local_files_only=local_files_only, text_encoder=text_encoder - ) - try: - tokenizer = ( - CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", local_files_only=local_files_only) - if tokenizer is None - else tokenizer - ) - except Exception: - raise ValueError( - f"With local_files_only set to {local_files_only}, you must first locally save the tokenizer in the following path: 'openai/clip-vit-large-patch14'." - ) - - if load_safety_checker: - safety_checker = StableDiffusionSafetyChecker.from_pretrained( - "CompVis/stable-diffusion-safety-checker", local_files_only=local_files_only - ) - feature_extractor = AutoFeatureExtractor.from_pretrained( - "CompVis/stable-diffusion-safety-checker", local_files_only=local_files_only - ) - - if controlnet: - pipe = pipeline_class( - vae=vae, - text_encoder=text_model, - tokenizer=tokenizer, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) - else: - pipe = pipeline_class( - vae=vae, - text_encoder=text_model, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) - elif model_type in ["SDXL", "SDXL-Refiner"]: - is_refiner = model_type == "SDXL-Refiner" - - if (is_refiner is False) and (tokenizer is None): - try: - tokenizer = CLIPTokenizer.from_pretrained( - "openai/clip-vit-large-patch14", local_files_only=local_files_only - ) - except Exception: - raise ValueError( - f"With local_files_only set to {local_files_only}, you must first locally save the tokenizer in the following path: 'openai/clip-vit-large-patch14'." - ) - - if (is_refiner is False) and (text_encoder is None): - text_encoder = convert_ldm_clip_checkpoint(checkpoint, local_files_only=local_files_only) - - if tokenizer_2 is None: - try: - tokenizer_2 = CLIPTokenizer.from_pretrained( - "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", pad_token="!", local_files_only=local_files_only - ) - except Exception: - raise ValueError( - f"With local_files_only set to {local_files_only}, you must first locally save the tokenizer in the following path: 'laion/CLIP-ViT-bigG-14-laion2B-39B-b160k' with `pad_token` set to '!'." - ) - - if text_encoder_2 is None: - config_name = "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k" - config_kwargs = {"projection_dim": 1280} - prefix = "conditioner.embedders.0.model." if is_refiner else "conditioner.embedders.1.model." - - text_encoder_2 = convert_open_clip_checkpoint( - checkpoint, - config_name, - prefix=prefix, - has_projection=True, - local_files_only=local_files_only, - **config_kwargs, - ) - - if is_accelerate_available(): # SBM Now move model to cpu. - for param_name, param in converted_unet_checkpoint.items(): - set_module_tensor_to_device(unet, param_name, "cpu", value=param) - - if controlnet: - pipe = pipeline_class( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - text_encoder_2=text_encoder_2, - tokenizer_2=tokenizer_2, - unet=unet, - controlnet=controlnet, - scheduler=scheduler, - force_zeros_for_empty_prompt=True, - ) - elif adapter: - pipe = pipeline_class( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - text_encoder_2=text_encoder_2, - tokenizer_2=tokenizer_2, - unet=unet, - adapter=adapter, - scheduler=scheduler, - force_zeros_for_empty_prompt=True, - ) - - else: - pipeline_kwargs = { - "vae": vae, - "text_encoder": text_encoder, - "tokenizer": tokenizer, - "text_encoder_2": text_encoder_2, - "tokenizer_2": tokenizer_2, - "unet": unet, - "scheduler": scheduler, - } - - if (pipeline_class == StableDiffusionXLImg2ImgPipeline) or ( - pipeline_class == StableDiffusionXLInpaintPipeline - ): - pipeline_kwargs.update({"requires_aesthetics_score": is_refiner}) - - if is_refiner: - pipeline_kwargs.update({"force_zeros_for_empty_prompt": False}) - - pipe = pipeline_class(**pipeline_kwargs) - else: - text_config = create_ldm_bert_config(original_config) - text_model = convert_ldm_bert_checkpoint(checkpoint, text_config) - tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased", local_files_only=local_files_only) - pipe = LDMTextToImagePipeline(vqvae=vae, bert=text_model, tokenizer=tokenizer, unet=unet, scheduler=scheduler) - - return pipe - - -def download_controlnet_from_original_ckpt( - checkpoint_path: str, - original_config_file: str, - image_size: int = 512, - extract_ema: bool = False, - num_in_channels: Optional[int] = None, - upcast_attention: Optional[bool] = None, - device: str = None, - from_safetensors: bool = False, - use_linear_projection: Optional[bool] = None, - cross_attention_dim: Optional[bool] = None, -) -> DiffusionPipeline: - if from_safetensors: - from safetensors import safe_open - - checkpoint = {} - with safe_open(checkpoint_path, framework="pt", device="cpu") as f: - for key in f.keys(): - checkpoint[key] = f.get_tensor(key) - else: - if device is None: - device = get_device() - checkpoint = torch.load(checkpoint_path, map_location=device) - else: - checkpoint = torch.load(checkpoint_path, map_location=device) - - # NOTE: this while loop isn't great but this controlnet checkpoint has one additional - # "state_dict" key https://huggingface.co/thibaud/controlnet-canny-sd21 - while "state_dict" in checkpoint: - checkpoint = checkpoint["state_dict"] - - with open(original_config_file, "r") as f: - original_config_file = f.read() - original_config = yaml.safe_load(original_config_file) - - if num_in_channels is not None: - original_config["model"]["params"]["unet_config"]["params"]["in_channels"] = num_in_channels - - if "control_stage_config" not in original_config["model"]["params"]: - raise ValueError("`control_stage_config` not present in original config") - - controlnet = convert_controlnet_checkpoint( - checkpoint, - original_config, - checkpoint_path, - image_size, - upcast_attention, - extract_ema, - use_linear_projection=use_linear_projection, - cross_attention_dim=cross_attention_dim, - ) - - return controlnet diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py deleted file mode 100644 index 1afa7698da7c15e890c2832a9bd10b9675d96c2c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py +++ /dev/null @@ -1,477 +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 warnings -from functools import partial -from typing import Dict, List, Optional, Union - -import jax -import jax.numpy as jnp -import numpy as np -from flax.core.frozen_dict import FrozenDict -from flax.jax_utils import unreplicate -from flax.training.common_utils import shard -from packaging import version -from PIL import Image -from transformers import CLIPImageProcessor, CLIPTokenizer, FlaxCLIPTextModel - -from ...models import FlaxAutoencoderKL, FlaxUNet2DConditionModel -from ...schedulers import ( - FlaxDDIMScheduler, - FlaxDPMSolverMultistepScheduler, - FlaxLMSDiscreteScheduler, - FlaxPNDMScheduler, -) -from ...utils import deprecate, logging, replace_example_docstring -from ..pipeline_flax_utils import FlaxDiffusionPipeline -from .pipeline_output import FlaxStableDiffusionPipelineOutput -from .safety_checker_flax import FlaxStableDiffusionSafetyChecker - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -# Set to True to use python for loop instead of jax.fori_loop for easier debugging -DEBUG = False - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import jax - >>> import numpy as np - >>> from flax.jax_utils import replicate - >>> from flax.training.common_utils import shard - - >>> from diffusers import FlaxStableDiffusionPipeline - - >>> pipeline, params = FlaxStableDiffusionPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", variant="bf16", dtype=jax.numpy.bfloat16 - ... ) - - >>> prompt = "a photo of an astronaut riding a horse on mars" - - >>> prng_seed = jax.random.PRNGKey(0) - >>> num_inference_steps = 50 - - >>> num_samples = jax.device_count() - >>> prompt = num_samples * [prompt] - >>> prompt_ids = pipeline.prepare_inputs(prompt) - # shard inputs and rng - - >>> params = replicate(params) - >>> prng_seed = jax.random.split(prng_seed, jax.device_count()) - >>> prompt_ids = shard(prompt_ids) - - >>> images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images - >>> images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]))) - ``` -""" - - -class FlaxStableDiffusionPipeline(FlaxDiffusionPipeline): - r""" - Flax-based pipeline for text-to-image generation using Stable Diffusion. - - This model inherits from [`FlaxDiffusionPipeline`]. Check the superclass documentation for the generic methods - implemented for all pipelines (downloading, saving, running on a particular device, etc.). - - Args: - vae ([`FlaxAutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.FlaxCLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`FlaxUNet2DConditionModel`]): - A `FlaxUNet2DConditionModel` 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 - [`FlaxDDIMScheduler`], [`FlaxLMSDiscreteScheduler`], [`FlaxPNDMScheduler`], or - [`FlaxDPMSolverMultistepScheduler`]. - safety_checker ([`FlaxStableDiffusionSafetyChecker`]): - 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - def __init__( - self, - vae: FlaxAutoencoderKL, - text_encoder: FlaxCLIPTextModel, - tokenizer: CLIPTokenizer, - unet: FlaxUNet2DConditionModel, - scheduler: Union[ - FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler - ], - safety_checker: FlaxStableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - dtype: jnp.dtype = jnp.float32, - ): - super().__init__() - self.dtype = dtype - - if safety_checker is None: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" - " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_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.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]]): - if not isinstance(prompt, (str, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - text_input = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - return text_input.input_ids - - def _get_has_nsfw_concepts(self, features, params): - has_nsfw_concepts = self.safety_checker(features, params) - return has_nsfw_concepts - - def _run_safety_checker(self, images, safety_model_params, jit=False): - # safety_model_params should already be replicated when jit is True - pil_images = [Image.fromarray(image) for image in images] - features = self.feature_extractor(pil_images, return_tensors="np").pixel_values - - if jit: - features = shard(features) - has_nsfw_concepts = _p_get_has_nsfw_concepts(self, features, safety_model_params) - has_nsfw_concepts = unshard(has_nsfw_concepts) - safety_model_params = unreplicate(safety_model_params) - else: - has_nsfw_concepts = self._get_has_nsfw_concepts(features, safety_model_params) - - images_was_copied = False - for idx, has_nsfw_concept in enumerate(has_nsfw_concepts): - if has_nsfw_concept: - if not images_was_copied: - images_was_copied = True - images = images.copy() - - images[idx] = np.zeros(images[idx].shape, dtype=np.uint8) # black image - - if any(has_nsfw_concepts): - warnings.warn( - "Potential NSFW content was detected in one or more images. A black image will be returned" - " instead. Try again with a different prompt and/or seed." - ) - - return images, has_nsfw_concepts - - def _generate( - self, - prompt_ids: jnp.array, - params: Union[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, - ): - 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}.") - - # get prompt text embeddings - prompt_embeds = self.text_encoder(prompt_ids, params=params["text_encoder"])[0] - - # TODO: currently it is assumed `do_classifier_free_guidance = guidance_scale > 1.0` - # implement this conditional `do_classifier_free_guidance = guidance_scale > 1.0` - batch_size = prompt_ids.shape[0] - - max_length = prompt_ids.shape[-1] - - if neg_prompt_ids is None: - uncond_input = self.tokenizer( - [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="np" - ).input_ids - else: - uncond_input = neg_prompt_ids - negative_prompt_embeds = self.text_encoder(uncond_input, params=params["text_encoder"])[0] - context = jnp.concatenate([negative_prompt_embeds, prompt_embeds]) - - # Ensure model output will be `float32` before going into the scheduler - guidance_scale = jnp.array([guidance_scale], dtype=jnp.float32) - - latents_shape = ( - batch_size, - self.unet.config.in_channels, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - if latents is None: - latents = jax.random.normal(prng_seed, shape=latents_shape, dtype=jnp.float32) - else: - if latents.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - - def loop_body(step, args): - latents, scheduler_state = args - # 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 - latents_input = jnp.concatenate([latents] * 2) - - t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step] - timestep = jnp.broadcast_to(t, latents_input.shape[0]) - - latents_input = self.scheduler.scale_model_input(scheduler_state, latents_input, t) - - # predict the noise residual - noise_pred = self.unet.apply( - {"params": params["unet"]}, - jnp.array(latents_input), - jnp.array(timestep, dtype=jnp.int32), - encoder_hidden_states=context, - ).sample - # perform guidance - noise_pred_uncond, noise_prediction_text = jnp.split(noise_pred, 2, axis=0) - noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents, scheduler_state = self.scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() - return latents, scheduler_state - - scheduler_state = self.scheduler.set_timesteps( - params["scheduler"], num_inference_steps=num_inference_steps, shape=latents.shape - ) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * params["scheduler"].init_noise_sigma - - if DEBUG: - # run with python for loop - for i in range(num_inference_steps): - latents, scheduler_state = loop_body(i, (latents, scheduler_state)) - else: - latents, _ = jax.lax.fori_loop(0, num_inference_steps, loop_body, (latents, scheduler_state)) - - # scale and decode the image latents with vae - latents = 1 / self.vae.config.scaling_factor * latents - image = self.vae.apply({"params": params["vae"]}, latents, method=self.vae.decode).sample - - image = (image / 2 + 0.5).clip(0, 1).transpose(0, 2, 3, 1) - return image - - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt_ids: jnp.array, - params: Union[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, - latents: jnp.ndarray = None, - neg_prompt_ids: jnp.ndarray = None, - return_dict: bool = True, - jit: bool = False, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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): - 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`. - latents (`jnp.ndarray`, *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 - array is generated by sampling using the supplied random `generator`. - jit (`bool`, defaults to `False`): - Whether to run `pmap` versions of the generation and safety scoring functions. - - - - This argument exists because `__call__` is not yet end-to-end pmap-able. It will be removed in a - future release. - - - - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] instead of - a plain tuple. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] 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. - """ - # 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 - - if isinstance(guidance_scale, float): - # Convert to a tensor so each device gets a copy. Follow the prompt_ids for - # shape information, as they may be sharded (when `jit` is `True`), or not. - guidance_scale = jnp.array([guidance_scale] * prompt_ids.shape[0]) - if len(prompt_ids.shape) > 2: - # Assume sharded - guidance_scale = guidance_scale[:, None] - - if jit: - images = _p_generate( - self, - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - ) - else: - images = self._generate( - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - ) - - if self.safety_checker is not None: - safety_params = params["safety_checker"] - images_uint8_casted = (images * 255).round().astype("uint8") - num_devices, batch_size = images.shape[:2] - - images_uint8_casted = np.asarray(images_uint8_casted).reshape(num_devices * batch_size, height, width, 3) - images_uint8_casted, has_nsfw_concept = self._run_safety_checker(images_uint8_casted, safety_params, jit) - images = np.asarray(images).copy() - - # block images - if any(has_nsfw_concept): - for i, is_nsfw in enumerate(has_nsfw_concept): - if is_nsfw: - images[i, 0] = np.asarray(images_uint8_casted[i]) - - images = images.reshape(num_devices, batch_size, height, width, 3) - else: - images = np.asarray(images) - has_nsfw_concept = False - - if not return_dict: - return (images, has_nsfw_concept) - - return FlaxStableDiffusionPipelineOutput(images=images, nsfw_content_detected=has_nsfw_concept) - - -# Static argnums are pipe, num_inference_steps, height, width. A change would trigger recompilation. -# Non-static args are (sharded) input tensors mapped over their first dimension (hence, `0`). -@partial( - jax.pmap, - in_axes=(None, 0, 0, 0, None, None, None, 0, 0, 0), - static_broadcasted_argnums=(0, 4, 5, 6), -) -def _p_generate( - pipe, - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, -): - return pipe._generate( - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - ) - - -@partial(jax.pmap, static_broadcasted_argnums=(0,)) -def _p_get_has_nsfw_concepts(pipe, features, params): - return pipe._get_has_nsfw_concepts(features, params) - - -def unshard(x: jnp.ndarray): - # einops.rearrange(x, 'd b ... -> (d b) ...') - num_devices, batch_size = x.shape[:2] - rest = x.shape[2:] - return x.reshape(num_devices * batch_size, *rest) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py deleted file mode 100644 index 78e3ba239c4e6ff35a194ff64b601a7cbe584149..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py +++ /dev/null @@ -1,532 +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 warnings -from functools import partial -from typing import Dict, List, Optional, Union - -import jax -import jax.numpy as jnp -import numpy as np -from flax.core.frozen_dict import FrozenDict -from flax.jax_utils import unreplicate -from flax.training.common_utils import shard -from PIL import Image -from transformers import CLIPImageProcessor, CLIPTokenizer, FlaxCLIPTextModel - -from ...models import FlaxAutoencoderKL, FlaxUNet2DConditionModel -from ...schedulers import ( - FlaxDDIMScheduler, - FlaxDPMSolverMultistepScheduler, - FlaxLMSDiscreteScheduler, - FlaxPNDMScheduler, -) -from ...utils import PIL_INTERPOLATION, logging, replace_example_docstring -from ..pipeline_flax_utils import FlaxDiffusionPipeline -from .pipeline_output import FlaxStableDiffusionPipelineOutput -from .safety_checker_flax import FlaxStableDiffusionSafetyChecker - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -# Set to True to use python for loop instead of jax.fori_loop for easier debugging -DEBUG = False - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import jax - >>> import numpy as np - >>> import jax.numpy as jnp - >>> from flax.jax_utils import replicate - >>> from flax.training.common_utils import shard - >>> import requests - >>> from io import BytesIO - >>> from PIL import Image - >>> from diffusers import FlaxStableDiffusionImg2ImgPipeline - - - >>> def create_key(seed=0): - ... return jax.random.PRNGKey(seed) - - - >>> rng = create_key(0) - - >>> url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" - >>> response = requests.get(url) - >>> init_img = Image.open(BytesIO(response.content)).convert("RGB") - >>> init_img = init_img.resize((768, 512)) - - >>> prompts = "A fantasy landscape, trending on artstation" - - >>> pipeline, params = FlaxStableDiffusionImg2ImgPipeline.from_pretrained( - ... "CompVis/stable-diffusion-v1-4", - ... revision="flax", - ... dtype=jnp.bfloat16, - ... ) - - >>> num_samples = jax.device_count() - >>> rng = jax.random.split(rng, jax.device_count()) - >>> prompt_ids, processed_image = pipeline.prepare_inputs( - ... prompt=[prompts] * num_samples, image=[init_img] * num_samples - ... ) - >>> p_params = replicate(params) - >>> prompt_ids = shard(prompt_ids) - >>> processed_image = shard(processed_image) - - >>> output = pipeline( - ... prompt_ids=prompt_ids, - ... image=processed_image, - ... params=p_params, - ... prng_seed=rng, - ... strength=0.75, - ... num_inference_steps=50, - ... jit=True, - ... height=512, - ... width=768, - ... ).images - - >>> output_images = pipeline.numpy_to_pil(np.asarray(output.reshape((num_samples,) + output.shape[-3:]))) - ``` -""" - - -class FlaxStableDiffusionImg2ImgPipeline(FlaxDiffusionPipeline): - r""" - Flax-based pipeline for text-guided image-to-image generation using Stable Diffusion. - - This model inherits from [`FlaxDiffusionPipeline`]. Check the superclass documentation for the generic methods - implemented for all pipelines (downloading, saving, running on a particular device, etc.). - - Args: - vae ([`FlaxAutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.FlaxCLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`FlaxUNet2DConditionModel`]): - A `FlaxUNet2DConditionModel` 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 - [`FlaxDDIMScheduler`], [`FlaxLMSDiscreteScheduler`], [`FlaxPNDMScheduler`], or - [`FlaxDPMSolverMultistepScheduler`]. - safety_checker ([`FlaxStableDiffusionSafetyChecker`]): - 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - def __init__( - self, - vae: FlaxAutoencoderKL, - text_encoder: FlaxCLIPTextModel, - tokenizer: CLIPTokenizer, - unet: FlaxUNet2DConditionModel, - scheduler: Union[ - FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler - ], - safety_checker: FlaxStableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - dtype: jnp.dtype = jnp.float32, - ): - super().__init__() - self.dtype = dtype - - if safety_checker is None: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) - 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]]): - if not isinstance(prompt, (str, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - 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)}") - - if isinstance(image, Image.Image): - image = [image] - - processed_images = jnp.concatenate([preprocess(img, jnp.float32) for img in image]) - - text_input = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - return text_input.input_ids, processed_images - - def _get_has_nsfw_concepts(self, features, params): - has_nsfw_concepts = self.safety_checker(features, params) - return has_nsfw_concepts - - def _run_safety_checker(self, images, safety_model_params, jit=False): - # safety_model_params should already be replicated when jit is True - pil_images = [Image.fromarray(image) for image in images] - features = self.feature_extractor(pil_images, return_tensors="np").pixel_values - - if jit: - features = shard(features) - has_nsfw_concepts = _p_get_has_nsfw_concepts(self, features, safety_model_params) - has_nsfw_concepts = unshard(has_nsfw_concepts) - safety_model_params = unreplicate(safety_model_params) - else: - has_nsfw_concepts = self._get_has_nsfw_concepts(features, safety_model_params) - - images_was_copied = False - for idx, has_nsfw_concept in enumerate(has_nsfw_concepts): - if has_nsfw_concept: - if not images_was_copied: - images_was_copied = True - images = images.copy() - - images[idx] = np.zeros(images[idx].shape, dtype=np.uint8) # black image - - if any(has_nsfw_concepts): - warnings.warn( - "Potential NSFW content was detected in one or more images. A black image will be returned" - " instead. Try again with a different prompt and/or seed." - ) - - return images, has_nsfw_concepts - - def get_timestep_start(self, num_inference_steps, strength): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - - return t_start - - def _generate( - self, - prompt_ids: jnp.ndarray, - image: jnp.ndarray, - params: Union[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, - ): - 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}.") - - # get prompt text embeddings - prompt_embeds = self.text_encoder(prompt_ids, params=params["text_encoder"])[0] - - # TODO: currently it is assumed `do_classifier_free_guidance = guidance_scale > 1.0` - # implement this conditional `do_classifier_free_guidance = guidance_scale > 1.0` - batch_size = prompt_ids.shape[0] - - max_length = prompt_ids.shape[-1] - - if neg_prompt_ids is None: - uncond_input = self.tokenizer( - [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="np" - ).input_ids - else: - uncond_input = neg_prompt_ids - negative_prompt_embeds = self.text_encoder(uncond_input, params=params["text_encoder"])[0] - context = jnp.concatenate([negative_prompt_embeds, prompt_embeds]) - - latents_shape = ( - batch_size, - self.unet.config.in_channels, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - if noise is None: - noise = jax.random.normal(prng_seed, shape=latents_shape, dtype=jnp.float32) - else: - if noise.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {noise.shape}, expected {latents_shape}") - - # Create init_latents - init_latent_dist = self.vae.apply({"params": params["vae"]}, image, method=self.vae.encode).latent_dist - init_latents = init_latent_dist.sample(key=prng_seed).transpose((0, 3, 1, 2)) - init_latents = self.vae.config.scaling_factor * init_latents - - def loop_body(step, args): - latents, scheduler_state = args - # 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 - latents_input = jnp.concatenate([latents] * 2) - - t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step] - timestep = jnp.broadcast_to(t, latents_input.shape[0]) - - latents_input = self.scheduler.scale_model_input(scheduler_state, latents_input, t) - - # predict the noise residual - noise_pred = self.unet.apply( - {"params": params["unet"]}, - jnp.array(latents_input), - jnp.array(timestep, dtype=jnp.int32), - encoder_hidden_states=context, - ).sample - # perform guidance - noise_pred_uncond, noise_prediction_text = jnp.split(noise_pred, 2, axis=0) - noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents, scheduler_state = self.scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() - return latents, scheduler_state - - scheduler_state = self.scheduler.set_timesteps( - params["scheduler"], num_inference_steps=num_inference_steps, shape=latents_shape - ) - - latent_timestep = scheduler_state.timesteps[start_timestep : start_timestep + 1].repeat(batch_size) - - latents = self.scheduler.add_noise(params["scheduler"], init_latents, noise, latent_timestep) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * params["scheduler"].init_noise_sigma - - if DEBUG: - # run with python for loop - for i in range(start_timestep, num_inference_steps): - latents, scheduler_state = loop_body(i, (latents, scheduler_state)) - else: - latents, _ = jax.lax.fori_loop(start_timestep, num_inference_steps, loop_body, (latents, scheduler_state)) - - # scale and decode the image latents with vae - latents = 1 / self.vae.config.scaling_factor * latents - image = self.vae.apply({"params": params["vae"]}, latents, method=self.vae.decode).sample - - image = (image / 2 + 0.5).clip(0, 1).transpose(0, 2, 3, 1) - return image - - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt_ids: jnp.ndarray, - image: jnp.ndarray, - params: Union[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, - noise: jnp.ndarray = None, - neg_prompt_ids: jnp.ndarray = None, - return_dict: bool = True, - jit: bool = False, - ): - r""" - The call function to the pipeline for generation. - - Args: - prompt_ids (`jnp.ndarray`): - The prompt or prompts to guide image generation. - image (`jnp.ndarray`): - Array representing an image batch to be used as the starting point. - params (`Dict` or `FrozenDict`): - Dictionary containing the model parameters/weights. - prng_seed (`jax.Array` or `jax.Array`): - Array containing random number generator key. - strength (`float`, *optional*, defaults to 0.8): - 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`. - 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`. - 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. - 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`. - noise (`jnp.ndarray`, *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. The array is generated by - sampling using the supplied random `generator`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] instead of - a plain tuple. - jit (`bool`, defaults to `False`): - Whether to run `pmap` versions of the generation and safety scoring functions. - - - - This argument exists because `__call__` is not yet end-to-end pmap-able. It will be removed in a - future release. - - - - Examples: - - Returns: - [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] 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. - """ - # 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 - - if isinstance(guidance_scale, float): - # Convert to a tensor so each device gets a copy. Follow the prompt_ids for - # shape information, as they may be sharded (when `jit` is `True`), or not. - guidance_scale = jnp.array([guidance_scale] * prompt_ids.shape[0]) - if len(prompt_ids.shape) > 2: - # Assume sharded - guidance_scale = guidance_scale[:, None] - - start_timestep = self.get_timestep_start(num_inference_steps, strength) - - if jit: - images = _p_generate( - self, - prompt_ids, - image, - params, - prng_seed, - start_timestep, - num_inference_steps, - height, - width, - guidance_scale, - noise, - neg_prompt_ids, - ) - else: - images = self._generate( - prompt_ids, - image, - params, - prng_seed, - start_timestep, - num_inference_steps, - height, - width, - guidance_scale, - noise, - neg_prompt_ids, - ) - - if self.safety_checker is not None: - safety_params = params["safety_checker"] - images_uint8_casted = (images * 255).round().astype("uint8") - num_devices, batch_size = images.shape[:2] - - images_uint8_casted = np.asarray(images_uint8_casted).reshape(num_devices * batch_size, height, width, 3) - images_uint8_casted, has_nsfw_concept = self._run_safety_checker(images_uint8_casted, safety_params, jit) - images = np.asarray(images) - - # block images - if any(has_nsfw_concept): - for i, is_nsfw in enumerate(has_nsfw_concept): - if is_nsfw: - images[i] = np.asarray(images_uint8_casted[i]) - - images = images.reshape(num_devices, batch_size, height, width, 3) - else: - images = np.asarray(images) - has_nsfw_concept = False - - if not return_dict: - return (images, has_nsfw_concept) - - return FlaxStableDiffusionPipelineOutput(images=images, nsfw_content_detected=has_nsfw_concept) - - -# Static argnums are pipe, start_timestep, num_inference_steps, height, width. A change would trigger recompilation. -# Non-static args are (sharded) input tensors mapped over their first dimension (hence, `0`). -@partial( - jax.pmap, - in_axes=(None, 0, 0, 0, 0, None, None, None, None, 0, 0, 0), - static_broadcasted_argnums=(0, 5, 6, 7, 8), -) -def _p_generate( - pipe, - prompt_ids, - image, - params, - prng_seed, - start_timestep, - num_inference_steps, - height, - width, - guidance_scale, - noise, - neg_prompt_ids, -): - return pipe._generate( - prompt_ids, - image, - params, - prng_seed, - start_timestep, - num_inference_steps, - height, - width, - guidance_scale, - noise, - neg_prompt_ids, - ) - - -@partial(jax.pmap, static_broadcasted_argnums=(0,)) -def _p_get_has_nsfw_concepts(pipe, features, params): - return pipe._get_has_nsfw_concepts(features, params) - - -def unshard(x: jnp.ndarray): - # einops.rearrange(x, 'd b ... -> (d b) ...') - num_devices, batch_size = x.shape[:2] - rest = x.shape[2:] - return x.reshape(num_devices * batch_size, *rest) - - -def preprocess(image, dtype): - w, h = image.size - w, h = (x - x % 32 for x in (w, h)) # resize to integer multiple of 32 - image = image.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]) - image = jnp.array(image).astype(dtype) / 255.0 - image = image[None].transpose(0, 3, 1, 2) - return 2.0 * image - 1.0 diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py deleted file mode 100644 index b7e17ba681a20b1382859b47b59fe815955adc6c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py +++ /dev/null @@ -1,593 +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 warnings -from functools import partial -from typing import Dict, List, Optional, Union - -import jax -import jax.numpy as jnp -import numpy as np -from flax.core.frozen_dict import FrozenDict -from flax.jax_utils import unreplicate -from flax.training.common_utils import shard -from packaging import version -from PIL import Image -from transformers import CLIPImageProcessor, CLIPTokenizer, FlaxCLIPTextModel - -from ...models import FlaxAutoencoderKL, FlaxUNet2DConditionModel -from ...schedulers import ( - FlaxDDIMScheduler, - FlaxDPMSolverMultistepScheduler, - FlaxLMSDiscreteScheduler, - FlaxPNDMScheduler, -) -from ...utils import PIL_INTERPOLATION, deprecate, logging, replace_example_docstring -from ..pipeline_flax_utils import FlaxDiffusionPipeline -from .pipeline_output import FlaxStableDiffusionPipelineOutput -from .safety_checker_flax import FlaxStableDiffusionSafetyChecker - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -# Set to True to use python for loop instead of jax.fori_loop for easier debugging -DEBUG = False - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import jax - >>> import numpy as np - >>> from flax.jax_utils import replicate - >>> from flax.training.common_utils import shard - >>> import PIL - >>> import requests - >>> from io import BytesIO - >>> from diffusers import FlaxStableDiffusionInpaintPipeline - - - >>> def download_image(url): - ... response = requests.get(url) - ... return PIL.Image.open(BytesIO(response.content)).convert("RGB") - - - >>> 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" - - >>> init_image = download_image(img_url).resize((512, 512)) - >>> mask_image = download_image(mask_url).resize((512, 512)) - - >>> pipeline, params = FlaxStableDiffusionInpaintPipeline.from_pretrained( - ... "xvjiarui/stable-diffusion-2-inpainting" - ... ) - - >>> prompt = "Face of a yellow cat, high resolution, sitting on a park bench" - >>> prng_seed = jax.random.PRNGKey(0) - >>> num_inference_steps = 50 - - >>> num_samples = jax.device_count() - >>> prompt = num_samples * [prompt] - >>> init_image = num_samples * [init_image] - >>> mask_image = num_samples * [mask_image] - >>> prompt_ids, processed_masked_images, processed_masks = pipeline.prepare_inputs( - ... prompt, init_image, mask_image - ... ) - # shard inputs and rng - - >>> params = replicate(params) - >>> prng_seed = jax.random.split(prng_seed, jax.device_count()) - >>> prompt_ids = shard(prompt_ids) - >>> processed_masked_images = shard(processed_masked_images) - >>> processed_masks = shard(processed_masks) - - >>> images = pipeline( - ... prompt_ids, processed_masks, processed_masked_images, params, prng_seed, num_inference_steps, jit=True - ... ).images - >>> images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]))) - ``` -""" - - -class FlaxStableDiffusionInpaintPipeline(FlaxDiffusionPipeline): - r""" - Flax-based pipeline for text-guided image inpainting using Stable Diffusion. - - - - 🧪 This is an experimental feature! - - - - This model inherits from [`FlaxDiffusionPipeline`]. Check the superclass documentation for the generic methods - implemented for all pipelines (downloading, saving, running on a particular device, etc.). - - Args: - vae ([`FlaxAutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.FlaxCLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`FlaxUNet2DConditionModel`]): - A `FlaxUNet2DConditionModel` 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 - [`FlaxDDIMScheduler`], [`FlaxLMSDiscreteScheduler`], [`FlaxPNDMScheduler`], or - [`FlaxDPMSolverMultistepScheduler`]. - safety_checker ([`FlaxStableDiffusionSafetyChecker`]): - 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - def __init__( - self, - vae: FlaxAutoencoderKL, - text_encoder: FlaxCLIPTextModel, - tokenizer: CLIPTokenizer, - unet: FlaxUNet2DConditionModel, - scheduler: Union[ - FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler - ], - safety_checker: FlaxStableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - dtype: jnp.dtype = jnp.float32, - ): - super().__init__() - self.dtype = dtype - - if safety_checker is None: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" - " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_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.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]], - mask: Union[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)}") - - 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)}") - - if isinstance(image, Image.Image): - image = [image] - - if not isinstance(mask, (Image.Image, list)): - raise ValueError(f"image has to be of type `PIL.Image.Image` or list but is {type(image)}") - - if isinstance(mask, Image.Image): - mask = [mask] - - processed_images = jnp.concatenate([preprocess_image(img, jnp.float32) for img in image]) - processed_masks = jnp.concatenate([preprocess_mask(m, jnp.float32) for m in mask]) - # processed_masks[processed_masks < 0.5] = 0 - processed_masks = processed_masks.at[processed_masks < 0.5].set(0) - # processed_masks[processed_masks >= 0.5] = 1 - processed_masks = processed_masks.at[processed_masks >= 0.5].set(1) - - processed_masked_images = processed_images * (processed_masks < 0.5) - - text_input = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - return text_input.input_ids, processed_masked_images, processed_masks - - def _get_has_nsfw_concepts(self, features, params): - has_nsfw_concepts = self.safety_checker(features, params) - return has_nsfw_concepts - - def _run_safety_checker(self, images, safety_model_params, jit=False): - # safety_model_params should already be replicated when jit is True - pil_images = [Image.fromarray(image) for image in images] - features = self.feature_extractor(pil_images, return_tensors="np").pixel_values - - if jit: - features = shard(features) - has_nsfw_concepts = _p_get_has_nsfw_concepts(self, features, safety_model_params) - has_nsfw_concepts = unshard(has_nsfw_concepts) - safety_model_params = unreplicate(safety_model_params) - else: - has_nsfw_concepts = self._get_has_nsfw_concepts(features, safety_model_params) - - images_was_copied = False - for idx, has_nsfw_concept in enumerate(has_nsfw_concepts): - if has_nsfw_concept: - if not images_was_copied: - images_was_copied = True - images = images.copy() - - images[idx] = np.zeros(images[idx].shape, dtype=np.uint8) # black image - - if any(has_nsfw_concepts): - warnings.warn( - "Potential NSFW content was detected in one or more images. A black image will be returned" - " instead. Try again with a different prompt and/or seed." - ) - - return images, has_nsfw_concepts - - def _generate( - self, - prompt_ids: jnp.ndarray, - mask: jnp.ndarray, - masked_image: jnp.ndarray, - params: Union[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, - ): - 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}.") - - # get prompt text embeddings - prompt_embeds = self.text_encoder(prompt_ids, params=params["text_encoder"])[0] - - # TODO: currently it is assumed `do_classifier_free_guidance = guidance_scale > 1.0` - # implement this conditional `do_classifier_free_guidance = guidance_scale > 1.0` - batch_size = prompt_ids.shape[0] - - max_length = prompt_ids.shape[-1] - - if neg_prompt_ids is None: - uncond_input = self.tokenizer( - [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="np" - ).input_ids - else: - uncond_input = neg_prompt_ids - negative_prompt_embeds = self.text_encoder(uncond_input, params=params["text_encoder"])[0] - context = jnp.concatenate([negative_prompt_embeds, prompt_embeds]) - - latents_shape = ( - batch_size, - self.vae.config.latent_channels, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - if latents is None: - latents = jax.random.normal(prng_seed, shape=latents_shape, dtype=self.dtype) - else: - if latents.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - - prng_seed, mask_prng_seed = jax.random.split(prng_seed) - - masked_image_latent_dist = self.vae.apply( - {"params": params["vae"]}, masked_image, method=self.vae.encode - ).latent_dist - masked_image_latents = masked_image_latent_dist.sample(key=mask_prng_seed).transpose((0, 3, 1, 2)) - masked_image_latents = self.vae.config.scaling_factor * masked_image_latents - del mask_prng_seed - - mask = jax.image.resize(mask, (*mask.shape[:-2], *masked_image_latents.shape[-2:]), method="nearest") - - # 8. Check that sizes of mask, masked image and latents match - num_channels_latents = self.vae.config.latent_channels - num_channels_mask = mask.shape[1] - num_channels_masked_image = masked_image_latents.shape[1] - if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" - " `pipeline.unet` or your `mask_image` or `image` input." - ) - - def loop_body(step, args): - latents, mask, masked_image_latents, scheduler_state = args - # 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 - latents_input = jnp.concatenate([latents] * 2) - mask_input = jnp.concatenate([mask] * 2) - masked_image_latents_input = jnp.concatenate([masked_image_latents] * 2) - - t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step] - timestep = jnp.broadcast_to(t, latents_input.shape[0]) - - latents_input = self.scheduler.scale_model_input(scheduler_state, latents_input, t) - # concat latents, mask, masked_image_latents in the channel dimension - latents_input = jnp.concatenate([latents_input, mask_input, masked_image_latents_input], axis=1) - - # predict the noise residual - noise_pred = self.unet.apply( - {"params": params["unet"]}, - jnp.array(latents_input), - jnp.array(timestep, dtype=jnp.int32), - encoder_hidden_states=context, - ).sample - # perform guidance - noise_pred_uncond, noise_prediction_text = jnp.split(noise_pred, 2, axis=0) - noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents, scheduler_state = self.scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() - return latents, mask, masked_image_latents, scheduler_state - - scheduler_state = self.scheduler.set_timesteps( - params["scheduler"], num_inference_steps=num_inference_steps, shape=latents.shape - ) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * params["scheduler"].init_noise_sigma - - if DEBUG: - # run with python for loop - for i in range(num_inference_steps): - latents, mask, masked_image_latents, scheduler_state = loop_body( - i, (latents, mask, masked_image_latents, scheduler_state) - ) - else: - latents, _, _, _ = jax.lax.fori_loop( - 0, num_inference_steps, loop_body, (latents, mask, masked_image_latents, scheduler_state) - ) - - # scale and decode the image latents with vae - latents = 1 / self.vae.config.scaling_factor * latents - image = self.vae.apply({"params": params["vae"]}, latents, method=self.vae.decode).sample - - image = (image / 2 + 0.5).clip(0, 1).transpose(0, 2, 3, 1) - return image - - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt_ids: jnp.ndarray, - mask: jnp.ndarray, - masked_image: jnp.ndarray, - params: Union[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, - latents: jnp.ndarray = None, - neg_prompt_ids: jnp.ndarray = None, - return_dict: bool = True, - jit: bool = False, - ): - r""" - Function invoked when calling the pipeline for generation. - - Args: - 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. - 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. This parameter is modulated by `strength`. - 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`. - latents (`jnp.ndarray`, *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 - array is generated by sampling using the supplied random `generator`. - jit (`bool`, defaults to `False`): - Whether to run `pmap` versions of the generation and safety scoring functions. - - - - This argument exists because `__call__` is not yet end-to-end pmap-able. It will be removed in a - future release. - - - - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] instead of - a plain tuple. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.FlaxStableDiffusionPipelineOutput`] 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. - """ - # 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 - - masked_image = jax.image.resize(masked_image, (*masked_image.shape[:-2], height, width), method="bicubic") - mask = jax.image.resize(mask, (*mask.shape[:-2], height, width), method="nearest") - - if isinstance(guidance_scale, float): - # Convert to a tensor so each device gets a copy. Follow the prompt_ids for - # shape information, as they may be sharded (when `jit` is `True`), or not. - guidance_scale = jnp.array([guidance_scale] * prompt_ids.shape[0]) - if len(prompt_ids.shape) > 2: - # Assume sharded - guidance_scale = guidance_scale[:, None] - - if jit: - images = _p_generate( - self, - prompt_ids, - mask, - masked_image, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - ) - else: - images = self._generate( - prompt_ids, - mask, - masked_image, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - ) - - if self.safety_checker is not None: - safety_params = params["safety_checker"] - images_uint8_casted = (images * 255).round().astype("uint8") - num_devices, batch_size = images.shape[:2] - - images_uint8_casted = np.asarray(images_uint8_casted).reshape(num_devices * batch_size, height, width, 3) - images_uint8_casted, has_nsfw_concept = self._run_safety_checker(images_uint8_casted, safety_params, jit) - images = np.asarray(images) - - # block images - if any(has_nsfw_concept): - for i, is_nsfw in enumerate(has_nsfw_concept): - if is_nsfw: - images[i] = np.asarray(images_uint8_casted[i]) - - images = images.reshape(num_devices, batch_size, height, width, 3) - else: - images = np.asarray(images) - has_nsfw_concept = False - - if not return_dict: - return (images, has_nsfw_concept) - - return FlaxStableDiffusionPipelineOutput(images=images, nsfw_content_detected=has_nsfw_concept) - - -# Static argnums are pipe, num_inference_steps, height, width. A change would trigger recompilation. -# Non-static args are (sharded) input tensors mapped over their first dimension (hence, `0`). -@partial( - jax.pmap, - in_axes=(None, 0, 0, 0, 0, 0, None, None, None, 0, 0, 0), - static_broadcasted_argnums=(0, 6, 7, 8), -) -def _p_generate( - pipe, - prompt_ids, - mask, - masked_image, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, -): - return pipe._generate( - prompt_ids, - mask, - masked_image, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - ) - - -@partial(jax.pmap, static_broadcasted_argnums=(0,)) -def _p_get_has_nsfw_concepts(pipe, features, params): - return pipe._get_has_nsfw_concepts(features, params) - - -def unshard(x: jnp.ndarray): - # einops.rearrange(x, 'd b ... -> (d b) ...') - num_devices, batch_size = x.shape[:2] - rest = x.shape[2:] - return x.reshape(num_devices * batch_size, *rest) - - -def preprocess_image(image, dtype): - w, h = image.size - w, h = (x - x % 32 for x in (w, h)) # resize to integer multiple of 32 - image = image.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]) - image = jnp.array(image).astype(dtype) / 255.0 - image = image[None].transpose(0, 3, 1, 2) - return 2.0 * image - 1.0 - - -def preprocess_mask(mask, dtype): - w, h = mask.size - w, h = (x - x % 32 for x in (w, h)) # resize to integer multiple of 32 - mask = mask.resize((w, h)) - mask = jnp.array(mask.convert("L")).astype(dtype) / 255.0 - mask = jnp.expand_dims(mask, axis=(0, 1)) - - return mask diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py deleted file mode 100644 index bd8609a11a200f25abe65fc2e595974cd802181a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py +++ /dev/null @@ -1,487 +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 inspect -from typing import Callable, List, Optional, Union - -import numpy as np -import torch -from transformers import CLIPImageProcessor, CLIPTokenizer - -from ...configuration_utils import FrozenDict -from ...schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler -from ...utils import deprecate, logging -from ..onnx_utils import ORT_TO_NP_TYPE, OnnxRuntimeModel -from ..pipeline_utils import DiffusionPipeline -from . import StableDiffusionPipelineOutput - - -logger = logging.get_logger(__name__) - - -class OnnxStableDiffusionPipeline(DiffusionPipeline): - vae_encoder: OnnxRuntimeModel - vae_decoder: OnnxRuntimeModel - text_encoder: OnnxRuntimeModel - tokenizer: CLIPTokenizer - unet: OnnxRuntimeModel - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler] - safety_checker: OnnxRuntimeModel - feature_extractor: CLIPImageProcessor - - _optional_components = ["safety_checker", "feature_extractor"] - _is_onnx = True - - def __init__( - self, - vae_encoder: OnnxRuntimeModel, - vae_decoder: OnnxRuntimeModel, - text_encoder: OnnxRuntimeModel, - tokenizer: CLIPTokenizer, - unet: OnnxRuntimeModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], - safety_checker: OnnxRuntimeModel, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae_encoder=vae_encoder, - vae_decoder=vae_decoder, - 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) - - def _encode_prompt( - self, - prompt: Union[str, List[str]], - num_images_per_prompt: Optional[int], - do_classifier_free_guidance: bool, - negative_prompt: Optional[str], - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - 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]`): - 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*): - 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 (`np.ndarray`, *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. - """ - 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: - # get prompt text embeddings - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - text_input_ids = text_inputs.input_ids - untruncated_ids = self.tokenizer(prompt, padding="max_length", return_tensors="np").input_ids - - if not np.array_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}" - ) - - prompt_embeds = self.text_encoder(input_ids=text_input_ids.astype(np.int32))[0] - - prompt_embeds = np.repeat(prompt_embeds, num_images_per_prompt, axis=0) - - # 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 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] * batch_size - 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 - - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="np", - ) - negative_prompt_embeds = self.text_encoder(input_ids=uncond_input.input_ids.astype(np.int32))[0] - - if do_classifier_free_guidance: - negative_prompt_embeds = np.repeat(negative_prompt_embeds, num_images_per_prompt, axis=0) - - # 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 - prompt_embeds = np.concatenate([negative_prompt_embeds, prompt_embeds]) - - return prompt_embeds - - def check_inputs( - self, - prompt: Union[str, List[str]], - height: Optional[int], - width: Optional[int], - callback_steps: int, - negative_prompt: Optional[str] = None, - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = 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 None) or ( - 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 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 __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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, np.ndarray], None]] = None, - callback_steps: int = 1, - ): - 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 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 - 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 (`np.random.RandomState`, *optional*): - One or a list of [numpy generator(s)](TODO) to make generation deterministic. - latents (`np.ndarray`, *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 (`np.ndarray`, *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 (`np.ndarray`, *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. - - 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`. - """ - - # check inputs. Raise error if not correct - self.check_inputs( - prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds - ) - - # 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 generator is None: - generator = np.random - - # 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 = guidance_scale > 1.0 - - prompt_embeds = self._encode_prompt( - prompt, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - - # get the initial random noise unless the user supplied it - latents_dtype = prompt_embeds.dtype - latents_shape = (batch_size * num_images_per_prompt, 4, height // 8, width // 8) - if latents is None: - latents = generator.randn(*latents_shape).astype(latents_dtype) - elif latents.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - - # set timesteps - self.scheduler.set_timesteps(num_inference_steps) - - latents = latents * np.float64(self.scheduler.init_noise_sigma) - - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - timestep_dtype = next( - (input.type for input in self.unet.model.get_inputs() if input.name == "timestep"), "tensor(float)" - ) - timestep_dtype = ORT_TO_NP_TYPE[timestep_dtype] - - for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(torch.from_numpy(latent_model_input), t) - latent_model_input = latent_model_input.cpu().numpy() - - # predict the noise residual - timestep = np.array([t], dtype=timestep_dtype) - noise_pred = self.unet(sample=latent_model_input, timestep=timestep, encoder_hidden_states=prompt_embeds) - noise_pred = noise_pred[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - scheduler_output = self.scheduler.step( - torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs - ) - latents = scheduler_output.prev_sample.numpy() - - # call the callback, if provided - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - latents = 1 / 0.18215 * latents - # image = self.vae_decoder(latent_sample=latents)[0] - # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1 - image = np.concatenate( - [self.vae_decoder(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])] - ) - - image = np.clip(image / 2 + 0.5, 0, 1) - image = image.transpose((0, 2, 3, 1)) - - if self.safety_checker is not None: - safety_checker_input = self.feature_extractor( - self.numpy_to_pil(image), return_tensors="np" - ).pixel_values.astype(image.dtype) - - images, has_nsfw_concept = [], [] - for i in range(image.shape[0]): - image_i, has_nsfw_concept_i = self.safety_checker( - clip_input=safety_checker_input[i : i + 1], images=image[i : i + 1] - ) - images.append(image_i) - has_nsfw_concept.append(has_nsfw_concept_i[0]) - image = np.concatenate(images) - else: - has_nsfw_concept = None - - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) - - -class StableDiffusionOnnxPipeline(OnnxStableDiffusionPipeline): - def __init__( - self, - vae_encoder: OnnxRuntimeModel, - vae_decoder: OnnxRuntimeModel, - text_encoder: OnnxRuntimeModel, - tokenizer: CLIPTokenizer, - unet: OnnxRuntimeModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], - safety_checker: OnnxRuntimeModel, - feature_extractor: CLIPImageProcessor, - ): - deprecation_message = "Please use `OnnxStableDiffusionPipeline` instead of `StableDiffusionOnnxPipeline`." - deprecate("StableDiffusionOnnxPipeline", "1.0.0", deprecation_message) - super().__init__( - vae_encoder=vae_encoder, - vae_decoder=vae_decoder, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py deleted file mode 100644 index d63bf3bf456459317790668ca8a7a39eb8024d69..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py +++ /dev/null @@ -1,550 +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 inspect -from typing import Callable, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTokenizer - -from ...configuration_utils import FrozenDict -from ...schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler -from ...utils import PIL_INTERPOLATION, deprecate, logging -from ..onnx_utils import ORT_TO_NP_TYPE, OnnxRuntimeModel -from ..pipeline_utils import DiffusionPipeline -from . import StableDiffusionPipelineOutput - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess with 8->64 -def preprocess(image): - 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): - return image - elif isinstance(image, PIL.Image.Image): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - w, h = image[0].size - w, h = (x - x % 64 for x in (w, h)) # resize to integer multiple of 64 - - image = [np.array(i.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image] - image = np.concatenate(image, axis=0) - image = np.array(image).astype(np.float32) / 255.0 - image = image.transpose(0, 3, 1, 2) - image = 2.0 * image - 1.0 - image = torch.from_numpy(image) - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, dim=0) - return image - - -class OnnxStableDiffusionImg2ImgPipeline(DiffusionPipeline): - r""" - Pipeline for text-guided image 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.) - - 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`. - """ - - vae_encoder: OnnxRuntimeModel - vae_decoder: OnnxRuntimeModel - text_encoder: OnnxRuntimeModel - tokenizer: CLIPTokenizer - unet: OnnxRuntimeModel - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler] - safety_checker: OnnxRuntimeModel - feature_extractor: CLIPImageProcessor - - _optional_components = ["safety_checker", "feature_extractor"] - _is_onnx = True - - def __init__( - self, - vae_encoder: OnnxRuntimeModel, - vae_decoder: OnnxRuntimeModel, - text_encoder: OnnxRuntimeModel, - tokenizer: CLIPTokenizer, - unet: OnnxRuntimeModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], - safety_checker: OnnxRuntimeModel, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae_encoder=vae_encoder, - vae_decoder=vae_decoder, - 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) - - # 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], - do_classifier_free_guidance: bool, - negative_prompt: Optional[str], - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - 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]`): - 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*): - 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 (`np.ndarray`, *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. - """ - 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: - # get prompt text embeddings - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - text_input_ids = text_inputs.input_ids - untruncated_ids = self.tokenizer(prompt, padding="max_length", return_tensors="np").input_ids - - if not np.array_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}" - ) - - prompt_embeds = self.text_encoder(input_ids=text_input_ids.astype(np.int32))[0] - - prompt_embeds = np.repeat(prompt_embeds, num_images_per_prompt, axis=0) - - # 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 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] * batch_size - 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 - - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="np", - ) - negative_prompt_embeds = self.text_encoder(input_ids=uncond_input.input_ids.astype(np.int32))[0] - - if do_classifier_free_guidance: - negative_prompt_embeds = np.repeat(negative_prompt_embeds, num_images_per_prompt, axis=0) - - # 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 - prompt_embeds = np.concatenate([negative_prompt_embeds, prompt_embeds]) - - return prompt_embeds - - def check_inputs( - self, - prompt: Union[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, - ): - if (callback_steps is None) or ( - 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 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 __call__( - self, - prompt: Union[str, List[str]], - image: Union[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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, np.ndarray], None]] = None, - callback_steps: int = 1, - ): - r""" - Function invoked when calling the pipeline for generation. - - Args: - 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 - process. - strength (`float`, *optional*, defaults to 0.8): - 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_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 will be modulated by `strength`. - 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. 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 (`np.random.RandomState`, *optional*): - A np.random.RandomState to make generation deterministic. - prompt_embeds (`np.ndarray`, *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 (`np.ndarray`, *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: np.ndarray)`. - 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. - - 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`. - """ - - # check inputs. Raise error if not correct - self.check_inputs(prompt, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds) - - # 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 strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - - if generator is None: - generator = np.random - - # set timesteps - self.scheduler.set_timesteps(num_inference_steps) - - image = preprocess(image).cpu().numpy() - - # 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 = guidance_scale > 1.0 - - prompt_embeds = self._encode_prompt( - prompt, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - - latents_dtype = prompt_embeds.dtype - image = image.astype(latents_dtype) - # encode the init image into latents and scale the latents - init_latents = self.vae_encoder(sample=image)[0] - init_latents = 0.18215 * init_latents - - if isinstance(prompt, str): - prompt = [prompt] - if len(prompt) > init_latents.shape[0] and len(prompt) % init_latents.shape[0] == 0: - # expand init_latents for batch_size - deprecation_message = ( - f"You have passed {len(prompt)} text prompts (`prompt`), but only {init_latents.shape[0]} initial" - " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" - " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" - " your script to pass as many initial images as text prompts to suppress this warning." - ) - deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) - additional_image_per_prompt = len(prompt) // init_latents.shape[0] - init_latents = np.concatenate([init_latents] * additional_image_per_prompt * num_images_per_prompt, axis=0) - elif len(prompt) > init_latents.shape[0] and len(prompt) % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {len(prompt)} text prompts." - ) - else: - init_latents = np.concatenate([init_latents] * num_images_per_prompt, axis=0) - - # get the original timestep using init_timestep - offset = self.scheduler.config.get("steps_offset", 0) - init_timestep = int(num_inference_steps * strength) + offset - init_timestep = min(init_timestep, num_inference_steps) - - timesteps = self.scheduler.timesteps.numpy()[-init_timestep] - timesteps = np.array([timesteps] * batch_size * num_images_per_prompt) - - # add noise to latents using the timesteps - noise = generator.randn(*init_latents.shape).astype(latents_dtype) - init_latents = self.scheduler.add_noise( - torch.from_numpy(init_latents), torch.from_numpy(noise), torch.from_numpy(timesteps) - ) - init_latents = init_latents.numpy() - - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - latents = init_latents - - t_start = max(num_inference_steps - init_timestep + offset, 0) - timesteps = self.scheduler.timesteps[t_start:].numpy() - - timestep_dtype = next( - (input.type for input in self.unet.model.get_inputs() if input.name == "timestep"), "tensor(float)" - ) - timestep_dtype = ORT_TO_NP_TYPE[timestep_dtype] - - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(torch.from_numpy(latent_model_input), t) - latent_model_input = latent_model_input.cpu().numpy() - - # predict the noise residual - timestep = np.array([t], dtype=timestep_dtype) - noise_pred = self.unet(sample=latent_model_input, timestep=timestep, encoder_hidden_states=prompt_embeds)[ - 0 - ] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - scheduler_output = self.scheduler.step( - torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs - ) - latents = scheduler_output.prev_sample.numpy() - - # call the callback, if provided - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - latents = 1 / 0.18215 * latents - # image = self.vae_decoder(latent_sample=latents)[0] - # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1 - image = np.concatenate( - [self.vae_decoder(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])] - ) - - image = np.clip(image / 2 + 0.5, 0, 1) - image = image.transpose((0, 2, 3, 1)) - - if self.safety_checker is not None: - safety_checker_input = self.feature_extractor( - self.numpy_to_pil(image), return_tensors="np" - ).pixel_values.astype(image.dtype) - # safety_checker does not support batched inputs yet - images, has_nsfw_concept = [], [] - for i in range(image.shape[0]): - image_i, has_nsfw_concept_i = self.safety_checker( - clip_input=safety_checker_input[i : i + 1], images=image[i : i + 1] - ) - images.append(image_i) - has_nsfw_concept.append(has_nsfw_concept_i[0]) - image = np.concatenate(images) - else: - has_nsfw_concept = None - - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py deleted file mode 100644 index 6a952a7ae650e08b4aab96ccd65b24f20b31663f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py +++ /dev/null @@ -1,564 +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 inspect -from typing import Callable, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTokenizer - -from ...configuration_utils import FrozenDict -from ...schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler -from ...utils import PIL_INTERPOLATION, deprecate, logging -from ..onnx_utils import ORT_TO_NP_TYPE, OnnxRuntimeModel -from ..pipeline_utils import DiffusionPipeline -from . import StableDiffusionPipelineOutput - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -NUM_UNET_INPUT_CHANNELS = 9 -NUM_LATENT_CHANNELS = 4 - - -def prepare_mask_and_masked_image(image, mask, latents_shape): - image = np.array(image.convert("RGB").resize((latents_shape[1] * 8, latents_shape[0] * 8))) - image = image[None].transpose(0, 3, 1, 2) - image = image.astype(np.float32) / 127.5 - 1.0 - - image_mask = np.array(mask.convert("L").resize((latents_shape[1] * 8, latents_shape[0] * 8))) - masked_image = image * (image_mask < 127.5) - - mask = mask.resize((latents_shape[1], latents_shape[0]), PIL_INTERPOLATION["nearest"]) - mask = np.array(mask.convert("L")) - mask = mask.astype(np.float32) / 255.0 - mask = mask[None, None] - mask[mask < 0.5] = 0 - mask[mask >= 0.5] = 1 - - return mask, masked_image - - -class OnnxStableDiffusionInpaintPipeline(DiffusionPipeline): - r""" - Pipeline for text-guided image inpainting using Stable Diffusion. *This is an experimental feature*. - - 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.) - - 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`. - """ - - vae_encoder: OnnxRuntimeModel - vae_decoder: OnnxRuntimeModel - text_encoder: OnnxRuntimeModel - tokenizer: CLIPTokenizer - unet: OnnxRuntimeModel - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler] - safety_checker: OnnxRuntimeModel - feature_extractor: CLIPImageProcessor - - _optional_components = ["safety_checker", "feature_extractor"] - _is_onnx = True - - def __init__( - self, - vae_encoder: OnnxRuntimeModel, - vae_decoder: OnnxRuntimeModel, - text_encoder: OnnxRuntimeModel, - tokenizer: CLIPTokenizer, - unet: OnnxRuntimeModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], - safety_checker: OnnxRuntimeModel, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - logger.info("`OnnxStableDiffusionInpaintPipeline` is experimental and will very likely change in the future.") - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae_encoder=vae_encoder, - vae_decoder=vae_decoder, - 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) - - # 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], - do_classifier_free_guidance: bool, - negative_prompt: Optional[str], - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - 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]`): - 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*): - 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 (`np.ndarray`, *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. - """ - 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: - # get prompt text embeddings - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - text_input_ids = text_inputs.input_ids - untruncated_ids = self.tokenizer(prompt, padding="max_length", return_tensors="np").input_ids - - if not np.array_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}" - ) - - prompt_embeds = self.text_encoder(input_ids=text_input_ids.astype(np.int32))[0] - - prompt_embeds = np.repeat(prompt_embeds, num_images_per_prompt, axis=0) - - # 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 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] * batch_size - 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 - - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="np", - ) - negative_prompt_embeds = self.text_encoder(input_ids=uncond_input.input_ids.astype(np.int32))[0] - - if do_classifier_free_guidance: - negative_prompt_embeds = np.repeat(negative_prompt_embeds, num_images_per_prompt, axis=0) - - # 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 - prompt_embeds = np.concatenate([negative_prompt_embeds, prompt_embeds]) - - return prompt_embeds - - # 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], - callback_steps: int, - negative_prompt: Optional[str] = None, - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = 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 None) or ( - 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 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}." - ) - - @torch.no_grad() - def __call__( - self, - prompt: Union[str, List[str]], - image: PIL.Image.Image, - mask_image: PIL.Image.Image, - height: Optional[int] = 512, - width: Optional[int] = 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, - 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, np.ndarray], None]] = None, - callback_steps: int = 1, - ): - r""" - Function invoked when calling the pipeline for generation. - - Args: - 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 - be masked out with `mask_image` and repainted according to `prompt`. - mask_image (`PIL.Image.Image`): - `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be - 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)`. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - 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. 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 (`np.random.RandomState`, *optional*): - A np.random.RandomState to make generation deterministic. - latents (`np.ndarray`, *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 (`np.ndarray`, *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 (`np.ndarray`, *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: np.ndarray)`. - 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. - - 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`. - """ - - # check inputs. Raise error if not correct - self.check_inputs( - prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds - ) - - # 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 generator is None: - generator = np.random - - # set timesteps - self.scheduler.set_timesteps(num_inference_steps) - - # 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 = guidance_scale > 1.0 - - prompt_embeds = self._encode_prompt( - prompt, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - - num_channels_latents = NUM_LATENT_CHANNELS - latents_shape = (batch_size * num_images_per_prompt, num_channels_latents, height // 8, width // 8) - latents_dtype = prompt_embeds.dtype - if latents is None: - latents = generator.randn(*latents_shape).astype(latents_dtype) - else: - if latents.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - - # prepare mask and masked_image - mask, masked_image = prepare_mask_and_masked_image(image, mask_image, latents_shape[-2:]) - mask = mask.astype(latents.dtype) - masked_image = masked_image.astype(latents.dtype) - - masked_image_latents = self.vae_encoder(sample=masked_image)[0] - masked_image_latents = 0.18215 * masked_image_latents - - # duplicate mask and masked_image_latents for each generation per prompt - mask = mask.repeat(batch_size * num_images_per_prompt, 0) - masked_image_latents = masked_image_latents.repeat(batch_size * num_images_per_prompt, 0) - - mask = np.concatenate([mask] * 2) if do_classifier_free_guidance else mask - masked_image_latents = ( - np.concatenate([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents - ) - - num_channels_mask = mask.shape[1] - num_channels_masked_image = masked_image_latents.shape[1] - - unet_input_channels = NUM_UNET_INPUT_CHANNELS - if num_channels_latents + num_channels_mask + num_channels_masked_image != unet_input_channels: - raise ValueError( - "Incorrect configuration settings! The config of `pipeline.unet` expects" - f" {unet_input_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" - " `pipeline.unet` or your `mask_image` or `image` input." - ) - - # set timesteps - self.scheduler.set_timesteps(num_inference_steps) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * np.float64(self.scheduler.init_noise_sigma) - - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - timestep_dtype = next( - (input.type for input in self.unet.model.get_inputs() if input.name == "timestep"), "tensor(float)" - ) - timestep_dtype = ORT_TO_NP_TYPE[timestep_dtype] - - for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents - # concat latents, mask, masked_image_latnets in the channel dimension - latent_model_input = self.scheduler.scale_model_input(torch.from_numpy(latent_model_input), t) - latent_model_input = latent_model_input.cpu().numpy() - latent_model_input = np.concatenate([latent_model_input, mask, masked_image_latents], axis=1) - - # predict the noise residual - timestep = np.array([t], dtype=timestep_dtype) - noise_pred = self.unet(sample=latent_model_input, timestep=timestep, encoder_hidden_states=prompt_embeds)[ - 0 - ] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - scheduler_output = self.scheduler.step( - torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs - ) - latents = scheduler_output.prev_sample.numpy() - - # call the callback, if provided - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - latents = 1 / 0.18215 * latents - # image = self.vae_decoder(latent_sample=latents)[0] - # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1 - image = np.concatenate( - [self.vae_decoder(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])] - ) - - image = np.clip(image / 2 + 0.5, 0, 1) - image = image.transpose((0, 2, 3, 1)) - - if self.safety_checker is not None: - safety_checker_input = self.feature_extractor( - self.numpy_to_pil(image), return_tensors="np" - ).pixel_values.astype(image.dtype) - # safety_checker does not support batched inputs yet - images, has_nsfw_concept = [], [] - for i in range(image.shape[0]): - image_i, has_nsfw_concept_i = self.safety_checker( - clip_input=safety_checker_input[i : i + 1], images=image[i : i + 1] - ) - images.append(image_i) - has_nsfw_concept.append(has_nsfw_concept_i[0]) - image = np.concatenate(images) - else: - has_nsfw_concept = None - - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py deleted file mode 100644 index 3f10764dc7b436e1aacd200ecd442183d3755d23..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py +++ /dev/null @@ -1,586 +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 inspect -from typing import Any, Callable, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTokenizer - -from ...configuration_utils import FrozenDict -from ...schedulers import DDPMScheduler, KarrasDiffusionSchedulers -from ...utils import deprecate, logging -from ..onnx_utils import ORT_TO_NP_TYPE, OnnxRuntimeModel -from ..pipeline_utils import DiffusionPipeline -from . import StableDiffusionPipelineOutput - - -logger = logging.get_logger(__name__) - - -def preprocess(image): - if isinstance(image, torch.Tensor): - return image - elif isinstance(image, PIL.Image.Image): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - w, h = image[0].size - w, h = (x - x % 64 for x in (w, h)) # resize to integer multiple of 32 - - image = [np.array(i.resize((w, h)))[None, :] for i in image] - image = np.concatenate(image, axis=0) - image = np.array(image).astype(np.float32) / 255.0 - image = image.transpose(0, 3, 1, 2) - image = 2.0 * image - 1.0 - image = torch.from_numpy(image) - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, dim=0) - - return image - - -class OnnxStableDiffusionUpscalePipeline(DiffusionPipeline): - vae: OnnxRuntimeModel - text_encoder: OnnxRuntimeModel - tokenizer: CLIPTokenizer - unet: OnnxRuntimeModel - low_res_scheduler: DDPMScheduler - scheduler: KarrasDiffusionSchedulers - safety_checker: OnnxRuntimeModel - feature_extractor: CLIPImageProcessor - - _optional_components = ["safety_checker", "feature_extractor"] - _is_onnx = True - - def __init__( - self, - vae: OnnxRuntimeModel, - text_encoder: OnnxRuntimeModel, - tokenizer: Any, - unet: OnnxRuntimeModel, - low_res_scheduler: DDPMScheduler, - scheduler: KarrasDiffusionSchedulers, - safety_checker: Optional[OnnxRuntimeModel] = None, - feature_extractor: Optional[CLIPImageProcessor] = None, - max_noise_level: int = 350, - num_latent_channels=4, - num_unet_input_channels=7, - requires_safety_checker: bool = True, - ): - super().__init__() - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - low_res_scheduler=low_res_scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) - self.register_to_config( - max_noise_level=max_noise_level, - num_latent_channels=num_latent_channels, - num_unet_input_channels=num_unet_input_channels, - ) - - def check_inputs( - self, - prompt: Union[str, List[str]], - image, - noise_level, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - if (callback_steps is None) or ( - 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 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 ( - not isinstance(image, torch.Tensor) - and not isinstance(image, PIL.Image.Image) - and not isinstance(image, np.ndarray) - and not isinstance(image, list) - ): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `np.ndarray`, `PIL.Image.Image` or `list` but is {type(image)}" - ) - - # verify batch size of prompt and image are same if image is a list or tensor or numpy array - if isinstance(image, (list, np.ndarray)): - 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 isinstance(image, list): - image_batch_size = len(image) - else: - image_batch_size = image.shape[0] - if batch_size != image_batch_size: - raise ValueError( - f"`prompt` has batch size {batch_size} and `image` has batch size {image_batch_size}." - " Please make sure that passed `prompt` matches the batch size of `image`." - ) - - # check noise level - if noise_level > self.config.max_noise_level: - raise ValueError(f"`noise_level` has to be <= {self.config.max_noise_level} but is {noise_level}") - - if (callback_steps is None) or ( - 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)}." - ) - - def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, generator, latents=None): - shape = (batch_size, num_channels_latents, height, width) - if latents is None: - latents = generator.randn(*shape).astype(dtype) - elif latents.shape != shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") - - return latents - - def decode_latents(self, latents): - latents = 1 / 0.08333 * latents - image = self.vae(latent_sample=latents)[0] - image = np.clip(image / 2 + 0.5, 0, 1) - image = image.transpose((0, 2, 3, 1)) - return image - - def _encode_prompt( - self, - prompt: Union[str, List[str]], - num_images_per_prompt: Optional[int], - do_classifier_free_guidance: bool, - negative_prompt: Optional[str], - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - 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]`): - 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*): - 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 (`np.ndarray`, *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. - """ - 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: - # get prompt text embeddings - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - text_input_ids = text_inputs.input_ids - untruncated_ids = self.tokenizer(prompt, padding="max_length", return_tensors="np").input_ids - - if not np.array_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}" - ) - - prompt_embeds = self.text_encoder(input_ids=text_input_ids.astype(np.int32))[0] - - prompt_embeds = np.repeat(prompt_embeds, num_images_per_prompt, axis=0) - - # 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 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] * batch_size - 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 - - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="np", - ) - negative_prompt_embeds = self.text_encoder(input_ids=uncond_input.input_ids.astype(np.int32))[0] - - if do_classifier_free_guidance: - negative_prompt_embeds = np.repeat(negative_prompt_embeds, num_images_per_prompt, axis=0) - - # 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 - prompt_embeds = np.concatenate([negative_prompt_embeds, prompt_embeds]) - - return prompt_embeds - - def __call__( - self, - prompt: Union[str, List[str]], - image: Union[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, - 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, np.ndarray], None]] = None, - callback_steps: Optional[int] = 1, - ): - r""" - Function invoked when calling the pipeline for generation. - - Args: - 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 - process. - 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 will be modulated by `strength`. - 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. - 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*): - 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): - 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 (`np.random.RandomState`, *optional*): - A np.random.RandomState 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 ge generated by sampling using the supplied random `generator`. - prompt_embeds (`np.ndarray`, *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 (`np.ndarray`, *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: np.ndarray)`. - 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. - - 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`. - """ - - # 1. Check inputs - self.check_inputs( - prompt, - image, - noise_level, - 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] - - if generator is None: - generator = np.random - - # 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 = guidance_scale > 1.0 - - prompt_embeds = self._encode_prompt( - prompt, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - - latents_dtype = prompt_embeds.dtype - image = preprocess(image).cpu().numpy() - height, width = image.shape[2:] - - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - self.config.num_latent_channels, - height, - width, - latents_dtype, - generator, - ) - image = image.astype(latents_dtype) - - self.scheduler.set_timesteps(num_inference_steps) - timesteps = self.scheduler.timesteps - - # Scale the initial noise by the standard deviation required by the scheduler - latents = latents * np.float64(self.scheduler.init_noise_sigma) - - # 5. Add noise to image - noise_level = np.array([noise_level]).astype(np.int64) - noise = generator.randn(*image.shape).astype(latents_dtype) - - image = self.low_res_scheduler.add_noise( - torch.from_numpy(image), torch.from_numpy(noise), torch.from_numpy(noise_level) - ) - image = image.numpy() - - batch_multiplier = 2 if do_classifier_free_guidance else 1 - image = np.concatenate([image] * batch_multiplier * num_images_per_prompt) - noise_level = np.concatenate([noise_level] * image.shape[0]) - - # 7. Check that sizes of image and latents match - num_channels_image = image.shape[1] - if self.config.num_latent_channels + num_channels_image != self.config.num_unet_input_channels: - raise ValueError( - "Incorrect configuration settings! The config of `pipeline.unet` expects" - f" {self.config.num_unet_input_channels} but received `num_channels_latents`: {self.config.num_latent_channels} +" - f" `num_channels_image`: {num_channels_image} " - f" = {self.config.num_latent_channels + num_channels_image}. Please verify the config of" - " `pipeline.unet` or your `image` input." - ) - - # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - timestep_dtype = next( - (input.type for input in self.unet.model.get_inputs() if input.name == "timestep"), "tensor(float)" - ) - timestep_dtype = ORT_TO_NP_TYPE[timestep_dtype] - - # 9. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents - - # concat latents, mask, masked_image_latents in the channel dimension - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - latent_model_input = np.concatenate([latent_model_input, image], axis=1) - - # timestep to tensor - timestep = np.array([t], dtype=timestep_dtype) - - # predict the noise residual - noise_pred = self.unet( - sample=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - class_labels=noise_level, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step( - torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs - ).prev_sample - latents = latents.numpy() - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - # 10. Post-processing - image = self.decode_latents(latents) - - if self.safety_checker is not None: - safety_checker_input = self.feature_extractor( - self.numpy_to_pil(image), return_tensors="np" - ).pixel_values.astype(image.dtype) - - images, has_nsfw_concept = [], [] - for i in range(image.shape[0]): - image_i, has_nsfw_concept_i = self.safety_checker( - clip_input=safety_checker_input[i : i + 1], images=image[i : i + 1] - ) - images.append(image_i) - has_nsfw_concept.append(has_nsfw_concept_i[0]) - image = np.concatenate(images) - else: - has_nsfw_concept = None - - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_output.py deleted file mode 100644 index 5fb9b1a1412d96b69144a4c2e960dcc8b75a615c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_output.py +++ /dev/null @@ -1,45 +0,0 @@ -from dataclasses import dataclass -from typing import List, Optional, Union - -import numpy as np -import PIL.Image - -from ...utils import BaseOutput, is_flax_available - - -@dataclass -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, - num_channels)`. - 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]] - - -if is_flax_available(): - import flax - - @flax.struct.dataclass - class FlaxStableDiffusionPipelineOutput(BaseOutput): - """ - Output class for Flax-based Stable Diffusion pipelines. - - 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 - or `None` if safety checking could not be performed. - """ - - images: np.ndarray - nsfw_content_detected: List[bool] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py deleted file mode 100644 index acf685784eb13e171fea0bce8dc5b9afdf5be71d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ /dev/null @@ -1,1103 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import torch -from packaging import version -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...configuration_utils import FrozenDict -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - 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, StableDiffusionMixin -from .pipeline_output import StableDiffusionPipelineOutput -from .safety_checker import StableDiffusionSafetyChecker - - -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 StableDiffusionPipeline - - >>> pipe = StableDiffusionPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> image = pipe(prompt).images[0] - ``` -""" - - -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 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 StableDiffusionPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - IPAdapterMixin, - FromSingleFileMixin, -): - """ - Pipeline for text-to-image generation using Stable Diffusion. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: CLIPVisionModelWithProjection = None, - requires_safety_checker: bool = True, - ): - super().__init__() - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - self._is_unet_config_sample_size_int = unet is not None and isinstance(unet.config.sample_size, int) - is_unet_sample_size_less_64 = ( - unet is not None - and hasattr(unet.config, "sample_size") - and self._is_unet_config_sample_size_int - and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" - " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_config) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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.register_to_config(requires_safety_checker=requires_safety_checker) - - 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 - - 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 - - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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 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 - - 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 prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_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}." - ) - - 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" - ) - - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_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, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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"], - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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. - 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*): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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). - guidance_rescale (`float`, *optional*, defaults to 0.0): - Guidance rescale factor from [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. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - # 0. Default height and width to unet - if not height or not width: - height = ( - self.unet.config.sample_size - if self._is_unet_config_sample_size_int - else self.unet.config.sample_size[0] - ) - width = ( - self.unet.config.sample_size - if self._is_unet_config_sample_size_int - else self.unet.config.sample_size[1] - ) - height, width = height * self.vae_scale_factor, width * self.vae_scale_factor - # to deal with lora scaling and other possible forward hooks - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - height, - width, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - 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. Encode input prompt - lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - 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, - self.do_classifier_free_guidance, - ) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 5. 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, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 6.1 Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) - else None - ) - - # 6.2 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) - - # 7. Denoising loop - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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.do_classifier_free_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py deleted file mode 100644 index e957c6661f873f35dd81488dd40b1fd1256a3a91..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py +++ /dev/null @@ -1,897 +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 contextlib -import inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from packaging import version -from transformers import CLIPTextModel, CLIPTokenizer, DPTForDepthEstimation, DPTImageProcessor - -from ...configuration_utils import FrozenDict -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - PIL_INTERPOLATION, - USE_PEFT_BACKEND, - deprecate, - is_torch_xla_available, - logging, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput - - -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 - - -# 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" -): - 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_img2img.preprocess -def preprocess(image): - 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): - return image - elif isinstance(image, PIL.Image.Image): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - w, h = image[0].size - w, h = (x - x % 8 for x in (w, h)) # resize to integer multiple of 8 - - image = [np.array(i.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image] - image = np.concatenate(image, axis=0) - image = np.array(image).astype(np.float32) / 255.0 - image = image.transpose(0, 3, 1, 2) - image = 2.0 * image - 1.0 - image = torch.from_numpy(image) - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, dim=0) - return image - - -class StableDiffusionDepth2ImgPipeline(DiffusionPipeline, TextualInversionLoaderMixin, StableDiffusionLoraLoaderMixin): - r""" - Pipeline for text-guided depth-based image-to-image generation using Stable Diffusion. - - 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.). - - 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 - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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`]. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "depth_mask"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - depth_estimator: DPTForDepthEstimation, - feature_extractor: DPTImageProcessor, - ): - super().__init__() - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" - " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_config) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - depth_estimator=depth_estimator, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - strength, - callback_steps, - 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 in [0.0, 1.0] but is {strength}") - - 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}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.prepare_latents - def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - 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." - ) - - elif isinstance(generator, list): - if image.shape[0] < batch_size and batch_size % image.shape[0] == 0: - image = torch.cat([image] * (batch_size // image.shape[0]), dim=0) - elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} " - ) - - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - deprecation_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial" - " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" - " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" - " your script to pass as many initial images as text prompts to suppress this warning." - ) - deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - latents = init_latents - - return latents - - def prepare_depth_map(self, image, depth_map, batch_size, do_classifier_free_guidance, dtype, device): - if isinstance(image, PIL.Image.Image): - image = [image] - else: - image = list(image) - - if isinstance(image[0], PIL.Image.Image): - width, height = image[0].size - elif isinstance(image[0], np.ndarray): - width, height = image[0].shape[:-1] - else: - height, width = image[0].shape[-2:] - - if depth_map is None: - pixel_values = self.feature_extractor(images=image, return_tensors="pt").pixel_values - pixel_values = pixel_values.to(device=device, dtype=dtype) - # The DPT-Hybrid model uses batch-norm layers which are not compatible with fp16. - # So we use `torch.autocast` here for half precision inference. - if torch.backends.mps.is_available(): - autocast_ctx = contextlib.nullcontext() - logger.warning( - "The DPT-Hybrid model uses batch-norm layers which are not compatible with fp16, but autocast is not yet supported on MPS." - ) - else: - autocast_ctx = torch.autocast(device.type, dtype=dtype) - - with autocast_ctx: - depth_map = self.depth_estimator(pixel_values).predicted_depth - else: - depth_map = depth_map.to(device=device, dtype=dtype) - - depth_map = torch.nn.functional.interpolate( - depth_map.unsqueeze(1), - size=(height // self.vae_scale_factor, width // self.vae_scale_factor), - mode="bicubic", - align_corners=False, - ) - - depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True) - depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True) - depth_map = 2.0 * (depth_map - depth_min) / (depth_max - depth_min) - 1.0 - depth_map = depth_map.to(dtype) - - # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method - if depth_map.shape[0] < batch_size: - repeat_by = batch_size // depth_map.shape[0] - depth_map = depth_map.repeat(repeat_by, 1, 1, 1) - - depth_map = torch.cat([depth_map] * 2) if do_classifier_free_guidance else depth_map - return depth_map - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - def __call__( - self, - prompt: Union[str, List[str]] = None, - image: PipelineImageInput = None, - depth_map: Optional[torch.Tensor] = 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", - 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"], - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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` 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*): - Depth prediction to be used as additional conditioning for the image generation process. If not - defined, it automatically predicts the depth with `self.depth_estimator`. - strength (`float`, *optional*, defaults to 0.8): - 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`. - 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`. - 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*): - 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): - The number of images to generate per prompt. - 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*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - 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. - 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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). - 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. - 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. - Examples: - - ```py - >>> import torch - >>> import requests - >>> from PIL import Image - - >>> from diffusers import StableDiffusionDepth2ImgPipeline - - >>> pipe = StableDiffusionDepth2ImgPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-2-depth", - ... torch_dtype=torch.float16, - ... ) - >>> pipe.to("cuda") - - - >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" - >>> init_image = Image.open(requests.get(url, stream=True).raw) - >>> prompt = "two tigers" - >>> n_prompt = "bad, deformed, ugly, bad anotomy" - >>> image = pipe(prompt=prompt, image=init_image, negative_prompt=n_prompt, strength=0.7).images[0] - ``` - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned, - otherwise a `tuple` is returned where the first element is a list with the generated images. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - # 1. Check inputs - self.check_inputs( - prompt, - strength, - callback_steps, - 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, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - - if image is None: - raise ValueError("`image` input cannot be undefined.") - - # 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 - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - # 4. Prepare depth mask - depth_mask = self.prepare_depth_map( - image, - depth_map, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - prompt_embeds.dtype, - device, - ) - - # 5. Preprocess image - image = self.image_processor.preprocess(image) - - # 6. Set timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - - # 7. Prepare latent variables - latents = self.prepare_latents( - image, latent_timestep, batch_size, num_images_per_prompt, prompt_embeds.dtype, device, generator - ) - - # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 9. Denoising loop - 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): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - latent_model_input = torch.cat([latent_model_input, depth_mask], dim=1) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=self.cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - depth_mask = callback_outputs.pop("depth_mask", depth_mask) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - else: - image = latents - - image = self.image_processor.postprocess(image, output_type=output_type) - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py deleted file mode 100644 index d47e2f0593dd1e2f86bcfc4332eaa78a6c60563d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py +++ /dev/null @@ -1,439 +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 inspect -from typing import Callable, List, Optional, Union - -import PIL.Image -import torch -from packaging import version -from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection - -from ...configuration_utils import FrozenDict -from ...image_processor import VaeImageProcessor -from ...models import AutoencoderKL, UNet2DConditionModel -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import deprecate, is_torch_xla_available, logging -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from . import StableDiffusionPipelineOutput -from .safety_checker import StableDiffusionSafetyChecker - - -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 - - -class StableDiffusionImageVariationPipeline(DiffusionPipeline, StableDiffusionMixin): - r""" - Pipeline to generate image variations from an input image using Stable Diffusion. - - 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: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - image_encoder ([`~transformers.CLIPVisionModelWithProjection`]): - Frozen CLIP image-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - # TODO: feature_extractor is required to encode images (if they are in PIL format), - # we should give a descriptive message if the pipeline doesn't have one. - _optional_components = ["safety_checker"] - model_cpu_offload_seq = "image_encoder->unet->vae" - _exclude_from_cpu_offload = ["safety_checker"] - - def __init__( - self, - vae: AutoencoderKL, - image_encoder: CLIPVisionModelWithProjection, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" - " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_config) - - self.register_modules( - vae=vae, - image_encoder=image_encoder, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - def _encode_image(self, image, device, num_images_per_prompt, do_classifier_free_guidance): - dtype = next(self.image_encoder.parameters()).dtype - - if not isinstance(image, torch.Tensor): - image = self.feature_extractor(images=image, return_tensors="pt").pixel_values - - image = image.to(device=device, dtype=dtype) - image_embeddings = self.image_encoder(image).image_embeds - image_embeddings = image_embeddings.unsqueeze(1) - - # duplicate image embeddings for each generation per prompt, using mps friendly method - bs_embed, seq_len, _ = image_embeddings.shape - image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1) - image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) - - if do_classifier_free_guidance: - negative_prompt_embeds = torch.zeros_like(image_embeddings) - - # 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 - image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings]) - - return image_embeddings - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs(self, image, height, width, callback_steps): - if ( - not isinstance(image, torch.Tensor) - and not isinstance(image, PIL.Image.Image) - 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" - f" {type(image)}" - ) - - 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 None) or ( - 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)}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @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, - num_inference_steps: int = 50, - guidance_scale: float = 7.5, - 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, - output_type: Optional[str] = "pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], 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 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`): - 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. This parameter is modulated by `strength`. - 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`. - 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 (η) 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*): - 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`. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - - Examples: - - ```py - from diffusers import StableDiffusionImageVariationPipeline - from PIL import Image - from io import BytesIO - import requests - - pipe = StableDiffusionImageVariationPipeline.from_pretrained( - "lambdalabs/sd-image-variations-diffusers", revision="v2.0" - ) - pipe = pipe.to("cuda") - - url = "https://lh3.googleusercontent.com/y-iFOHfLTwkuQSUegpwDdgKmOjRSTvPxat63dQLB25xkTs4lhIbRUFeNBWZzYf370g=s1200" - - response = requests.get(url) - image = Image.open(BytesIO(response.content)).convert("RGB") - - out = pipe(image, num_images_per_prompt=3, guidance_scale=15) - out["images"][0].save("result.jpg") - ``` - """ - # 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(image, height, width, callback_steps) - - # 2. Define call parameters - if isinstance(image, PIL.Image.Image): - batch_size = 1 - elif isinstance(image, list): - batch_size = len(image) - else: - batch_size = image.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 = guidance_scale > 1.0 - - # 3. Encode input image - image_embeddings = self._encode_image(image, device, num_images_per_prompt, do_classifier_free_guidance) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. 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, - image_embeddings.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=image_embeddings).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - self.maybe_free_model_hooks() - - 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, image_embeddings.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) - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py deleted file mode 100644 index 95d3ab06f02a2f34a5d71f5796a79f578f276a73..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +++ /dev/null @@ -1,1161 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from packaging import version -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...configuration_utils import FrozenDict -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - PIL_INTERPOLATION, - USE_PEFT_BACKEND, - deprecate, - 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, StableDiffusionMixin -from . import StableDiffusionPipelineOutput -from .safety_checker import StableDiffusionSafetyChecker - - -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 requests - >>> import torch - >>> from PIL import Image - >>> from io import BytesIO - - >>> from diffusers import StableDiffusionImg2ImgPipeline - - >>> device = "cuda" - >>> model_id_or_path = "stable-diffusion-v1-5/stable-diffusion-v1-5" - >>> pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16) - >>> pipe = pipe.to(device) - - >>> url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" - - >>> response = requests.get(url) - >>> init_image = Image.open(BytesIO(response.content)).convert("RGB") - >>> init_image = init_image.resize((768, 512)) - - >>> prompt = "A fantasy landscape, trending on artstation" - - >>> images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images - >>> images[0].save("fantasy_landscape.png") - ``` -""" - - -def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = 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") - - -def preprocess(image): - 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): - return image - elif isinstance(image, PIL.Image.Image): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - w, h = image[0].size - w, h = (x - x % 8 for x in (w, h)) # resize to integer multiple of 8 - - image = [np.array(i.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image] - image = np.concatenate(image, axis=0) - image = np.array(image).astype(np.float32) / 255.0 - image = image.transpose(0, 3, 1, 2) - image = 2.0 * image - 1.0 - image = torch.from_numpy(image) - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, dim=0) - return 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, - **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 StableDiffusionImg2ImgPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-guided image-to-image generation using Stable Diffusion. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: CLIPVisionModelWithProjection = None, - requires_safety_checker: bool = True, - ): - super().__init__() - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" - " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_config) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - strength, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - ): - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - - 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}." - ) - - 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" - ) - - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - 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." - ) - - elif isinstance(generator, list): - if image.shape[0] < batch_size and batch_size % image.shape[0] == 0: - image = torch.cat([image] * (batch_size // image.shape[0]), dim=0) - elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} " - ) - - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - deprecation_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial" - " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" - " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" - " your script to pass as many initial images as text prompts to suppress this warning." - ) - deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - latents = init_latents - - return latents - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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"], - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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`, 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. - strength (`float`, *optional*, defaults to 0.8): - 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`. - 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*): - 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*): - 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*): - 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): - The number of images to generate per prompt. - 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*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - 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. - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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). - 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. - 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. - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - 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, - strength, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - 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. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - 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, - self.do_classifier_free_guidance, - ) - - # 4. Preprocess image - image = self.image_processor.preprocess(image) - - # 5. set timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, 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) - - # 6. Prepare latent variables - latents = self.prepare_latents( - image, - latent_timestep, - batch_size, - num_images_per_prompt, - prompt_embeds.dtype, - device, - generator, - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - # 7.2 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) - - # 8. Denoising loop - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py deleted file mode 100644 index 148d7386a732dc5aea824d6dd4a2dda919551c2a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +++ /dev/null @@ -1,1359 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import PIL.Image -import torch -from packaging import version -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...configuration_utils import FrozenDict -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AsymmetricAutoencoderKL, AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_torch_xla_available, - logging, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from . import StableDiffusionPipelineOutput -from .safety_checker import StableDiffusionSafetyChecker - - -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 - - -# 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" -): - 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 StableDiffusionInpaintPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-guided image inpainting using Stable Diffusion. - - 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.). - - 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 - - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - Args: - vae ([`AutoencoderKL`, `AsymmetricAutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "mask", "masked_image_latents"] - - def __init__( - self, - vae: Union[AutoencoderKL, AsymmetricAutoencoderKL], - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: CLIPVisionModelWithProjection = None, - requires_safety_checker: bool = True, - ): - super().__init__() - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "skip_prk_steps", True) is False: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration" - " `skip_prk_steps`. `skip_prk_steps` should be set to True in the configuration file. Please make" - " sure to update the config accordingly as not setting `skip_prk_steps` in the config might lead to" - " incorrect results in future versions. If you have downloaded this checkpoint from the Hugging Face" - " Hub, it would be very nice if you could open a Pull request for the" - " `scheduler/scheduler_config.json` file" - ) - deprecate("skip_prk_steps not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["skip_prk_steps"] = True - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" - " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_config) - - # Check shapes, assume num_channels_latents == 4, num_channels_mask == 1, num_channels_masked == 4 - if unet is not None and unet.config.in_channels != 9: - logger.info(f"You have loaded a UNet with {unet.config.in_channels} input channels which.") - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - image, - mask_image, - height, - width, - strength, - callback_steps, - output_type, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - padding_mask_crop=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 != 0 or width % self.vae_scale_factor != 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}." - ) - 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 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" - ) - - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=None, - image=None, - timestep=None, - is_strength_max=True, - return_noise=False, - return_image_latents=False, - ): - 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 (image is None or timestep is None) and not is_strength_max: - raise ValueError( - "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise." - "However, either the image or the noise timestep has not been provided." - ) - - if return_image_latents or (latents is None and not is_strength_max): - image = image.to(device=device, dtype=dtype) - - if image.shape[1] == 4: - image_latents = image - else: - image_latents = self._encode_vae_image(image=image, generator=generator) - image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1) - - if latents is None: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # if strength is 1. then initialise the latents to noise, else initial to image + noise - latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep) - # if pure noise then scale the initial latents by the Scheduler's init sigma - latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents - else: - noise = latents.to(device) - latents = noise * self.scheduler.init_noise_sigma - - outputs = (latents,) - - if return_noise: - outputs += (noise,) - - if return_image_latents: - outputs += (image_latents,) - - return outputs - - 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 = self.vae.config.scaling_factor * image_latents - - return image_latents - - def prepare_mask_latents( - self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance - ): - # 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 // self.vae_scale_factor, width // self.vae_scale_factor) - ) - mask = mask.to(device=device, dtype=dtype) - - masked_image = masked_image.to(device=device, dtype=dtype) - - if masked_image.shape[1] == 4: - masked_image_latents = masked_image - else: - masked_image_latents = self._encode_vae_image(masked_image, generator=generator) - - # 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) - - mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask - masked_image_latents = ( - torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents - ) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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 - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @property - def interrupt(self): - return self._interrupt - - @torch.no_grad() - def __call__( - self, - prompt: Union[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, - strength: float = 1.0, - num_inference_steps: int = 50, - 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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"], - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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`, 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]`): - `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)`. - 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. - padding_mask_crop (`int`, *optional*, defaults to `None`): - The size of margin in the crop to be applied to the image and masking. If `None`, no crop is applied to - image and mask_image. If `padding_mask_crop` is not `None`, it will first find a rectangular region - with the same aspect ration of the image and contains all masked area, and then expand that area based - on `padding_mask_crop`. The image and mask_image will then be cropped based on the expanded area before - resizing to the original image size for inpainting. This is useful when the masked area is small while - the image is large and contain information irrelevant for inpainting, such as background. - 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`. - 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*): - 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*): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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). - 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. - 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. - Examples: - - ```py - >>> import PIL - >>> import requests - >>> import torch - >>> from io import BytesIO - - >>> from diffusers import StableDiffusionInpaintPipeline - - - >>> def download_image(url): - ... response = requests.get(url) - ... return PIL.Image.open(BytesIO(response.content)).convert("RGB") - - - >>> 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" - - >>> init_image = download_image(img_url).resize((512, 512)) - >>> mask_image = download_image(mask_url).resize((512, 512)) - - >>> pipe = StableDiffusionInpaintPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-inpainting", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = "Face of a yellow cat, high resolution, sitting on a park bench" - >>> image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0] - ``` - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - # 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 - self.check_inputs( - prompt, - image, - mask_image, - height, - width, - strength, - callback_steps, - output_type, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - padding_mask_crop, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - 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. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=text_encoder_lora_scale, - clip_skip=self.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 self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - 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, - self.do_classifier_free_guidance, - ) - - # 4. set timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - timesteps, num_inference_steps = self.get_timesteps( - num_inference_steps=num_inference_steps, strength=strength, device=device - ) - # check that number of inference steps is not < 1 - as this doesn't make sense - 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." - ) - # at which timestep to set the initial noise (n.b. 50% if strength is 0.5) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise - is_strength_max = strength == 1.0 - - # 5. 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) - - # 6. Prepare latent variables - num_channels_latents = self.vae.config.latent_channels - num_channels_unet = self.unet.config.in_channels - return_image_latents = num_channels_unet == 4 - - latents_outputs = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - image=init_image, - timestep=latent_timestep, - is_strength_max=is_strength_max, - return_noise=True, - return_image_latents=return_image_latents, - ) - - if return_image_latents: - latents, noise, image_latents = latents_outputs - else: - latents, noise = latents_outputs - - # 7. Prepare mask latent variables - 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_images_per_prompt, - height, - width, - prompt_embeds.dtype, - device, - generator, - self.do_classifier_free_guidance, - ) - - # 8. Check that sizes of mask, masked image and latents match - if num_channels_unet == 9: - # default case for stable-diffusion-v1-5/stable-diffusion-inpainting - num_channels_mask = mask.shape[1] - num_channels_masked_image = masked_image_latents.shape[1] - if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" - " `pipeline.unet` or your `mask_image` or `image` input." - ) - elif num_channels_unet != 4: - raise ValueError( - f"The unet {self.unet.__class__} should have either 4 or 9 input channels, not {self.unet.config.in_channels}." - ) - - # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 9.1 Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - # 9.2 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) - - # 10. Denoising loop - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - # concat latents, mask, masked_image_latents in the channel dimension - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - if num_channels_unet == 9: - latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - if num_channels_unet == 4: - init_latents_proper = image_latents - if self.do_classifier_free_guidance: - init_mask, _ = mask.chunk(2) - else: - init_mask = mask - - if i < len(timesteps) - 1: - noise_timestep = timesteps[i + 1] - init_latents_proper = self.scheduler.add_noise( - init_latents_proper, noise, torch.tensor([noise_timestep]) - ) - - latents = (1 - init_mask) * init_latents_proper + init_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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if not output_type == "latent": - condition_kwargs = {} - if isinstance(self.vae, AsymmetricAutoencoderKL): - init_image = init_image.to(device=device, dtype=masked_image_latents.dtype) - init_image_condition = init_image.clone() - init_image = self._encode_vae_image(init_image, generator=generator) - mask_condition = mask_condition.to(device=device, dtype=masked_image_latents.dtype) - condition_kwargs = {"image": init_image_condition, "mask": mask_condition} - image = self.vae.decode( - latents / self.vae.config.scaling_factor, return_dict=False, generator=generator, **condition_kwargs - )[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) - - if padding_mask_crop is not None: - image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] - - # 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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py deleted file mode 100644 index 843d25d67c1040d5f3a4906daa400e1afae644c6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py +++ /dev/null @@ -1,917 +0,0 @@ -# Copyright 2025 The InstructPix2Pix Authors 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, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import PIL_INTERPOLATION, deprecate, is_torch_xla_available, logging -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from . import StableDiffusionPipelineOutput -from .safety_checker import StableDiffusionSafetyChecker - - -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 - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess -def preprocess(image): - 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): - return image - elif isinstance(image, PIL.Image.Image): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - w, h = image[0].size - w, h = (x - x % 8 for x in (w, h)) # resize to integer multiple of 8 - - image = [np.array(i.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image] - image = np.concatenate(image, axis=0) - image = np.array(image).astype(np.float32) / 255.0 - image = image.transpose(0, 3, 1, 2) - image = 2.0 * image - 1.0 - image = torch.from_numpy(image) - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, dim=0) - return 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" -): - 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 StableDiffusionInstructPix2PixPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - IPAdapterMixin, - FromSingleFileMixin, -): - r""" - Pipeline for pixel-level image editing by following text instructions (based on Stable Diffusion). - - 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.). - - 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 - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "image_latents"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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.register_to_config(requires_safety_checker=requires_safety_checker) - - @torch.no_grad() - def __call__( - self, - prompt: Union[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, - 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", - 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, - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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` 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): - 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): - 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`. - image_guidance_scale (`float`, *optional*, defaults to 1.5): - Push the generated image towards the initial `image`. Image guidance scale is enabled by setting - `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*): - 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): - The number of images to generate per prompt. - 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`, *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. - 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. - ip_adapter_image: (`PipelineImageInput`, *optional*): - Optional image input to work with IP Adapters. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - 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). - - Examples: - - ```py - >>> import PIL - >>> import requests - >>> import torch - >>> from io import BytesIO - - >>> from diffusers import StableDiffusionInstructPix2PixPipeline - - - >>> def download_image(url): - ... response = requests.get(url) - ... return PIL.Image.open(BytesIO(response.content)).convert("RGB") - - - >>> img_url = "https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/mountain.png" - - >>> image = download_image(img_url).resize((512, 512)) - - >>> pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained( - ... "timbrooks/instruct-pix2pix", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = "make the mountains snowy" - >>> image = pipe(prompt=prompt, image=image).images[0] - ``` - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - # 0. Check inputs - self.check_inputs( - prompt, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - self._guidance_scale = guidance_scale - self._image_guidance_scale = image_guidance_scale - - device = self._execution_device - - if image is None: - raise ValueError("`image` input cannot be undefined.") - - # 1. 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 - - # 2. Encode input prompt - prompt_embeds = self._encode_prompt( - prompt, - device, - num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - - 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, - self.do_classifier_free_guidance, - ) - # 3. Preprocess image - image = self.image_processor.preprocess(image) - - # 4. set timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare Image latents - image_latents = self.prepare_image_latents( - image, - batch_size, - num_images_per_prompt, - prompt_embeds.dtype, - device, - self.do_classifier_free_guidance, - ) - - height, width = image_latents.shape[-2:] - height = height * self.vae_scale_factor - width = width * self.vae_scale_factor - - # 6. Prepare latent variables - num_channels_latents = self.vae.config.latent_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 7. Check that shapes of latents and image match the UNet channels - num_channels_image = image_latents.shape[1] - if num_channels_latents + num_channels_image != self.unet.config.in_channels: - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_image`: {num_channels_image} " - f" = {num_channels_latents + num_channels_image}. Please verify the config of" - " `pipeline.unet` or your `image` input." - ) - - # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 8.1 Add image embeds for IP-Adapter - added_cond_kwargs = {"image_embeds": image_embeds} if ip_adapter_image is not None else None - - # 9. Denoising loop - 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): - # Expand the latents if we are doing classifier free guidance. - # The latents are expanded 3 times because for pix2pix the guidance\ - # is applied for both the text and the input image. - latent_model_input = torch.cat([latents] * 3) if self.do_classifier_free_guidance else latents - - # concat latents, image_latents in the channel dimension - scaled_latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - scaled_latent_model_input = torch.cat([scaled_latent_model_input, image_latents], dim=1) - - # predict the noise residual - noise_pred = self.unet( - scaled_latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - added_cond_kwargs=added_cond_kwargs, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_text, noise_pred_image, noise_pred_uncond = noise_pred.chunk(3) - noise_pred = ( - noise_pred_uncond - + self.guidance_scale * (noise_pred_text - noise_pred_image) - + self.image_guidance_scale * (noise_pred_image - noise_pred_uncond) - ) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - image_latents = callback_outputs.pop("image_latents", 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - 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) - - 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, - ): - 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. - """ - 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 - - prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask) - prompt_embeds = prompt_embeds[0] - - if self.text_encoder is not None: - prompt_embeds_dtype = self.text_encoder.dtype - else: - prompt_embeds_dtype = self.unet.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 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) - - # 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 - # pix2pix has two negative embeddings, and unlike in other pipelines latents are ordered [prompt_embeds, negative_prompt_embeds, negative_prompt_embeds] - prompt_embeds = torch.cat([prompt_embeds, negative_prompt_embeds, negative_prompt_embeds]) - - return prompt_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - 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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - image_embeds = [] - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - single_image_embeds = torch.stack([single_image_embeds] * num_images_per_prompt, dim=0) - single_negative_image_embeds = torch.stack( - [single_negative_image_embeds] * num_images_per_prompt, dim=0 - ) - - if do_classifier_free_guidance: - single_image_embeds = torch.cat( - [single_image_embeds, single_negative_image_embeds, single_negative_image_embeds] - ) - single_image_embeds = single_image_embeds.to(device) - - image_embeds.append(single_image_embeds) - else: - repeat_dims = [1] - image_embeds = [] - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - ( - single_image_embeds, - single_negative_image_embeds, - single_negative_image_embeds, - ) = single_image_embeds.chunk(3) - single_image_embeds = single_image_embeds.repeat( - num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:])) - ) - single_negative_image_embeds = single_negative_image_embeds.repeat( - num_images_per_prompt, *(repeat_dims * len(single_negative_image_embeds.shape[1:])) - ) - single_image_embeds = torch.cat( - [single_image_embeds, single_negative_image_embeds, single_negative_image_embeds] - ) - else: - single_image_embeds = single_image_embeds.repeat( - num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:])) - ) - image_embeds.append(single_image_embeds) - - return image_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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # 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, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - ): - 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}." - ) - - 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def prepare_image_latents( - self, image, batch_size, num_images_per_prompt, dtype, device, do_classifier_free_guidance, generator=None - ): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - image_latents = image - else: - image_latents = retrieve_latents(self.vae.encode(image), sample_mode="argmax") - - if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: - # expand image_latents for batch_size - deprecation_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {image_latents.shape[0]} initial" - " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" - " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" - " your script to pass as many initial images as text prompts to suppress this warning." - ) - deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) - 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 do_classifier_free_guidance: - uncond_image_latents = torch.zeros_like(image_latents) - image_latents = torch.cat([image_latents, image_latents, uncond_image_latents], dim=0) - - return image_latents - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def image_guidance_scale(self): - return self._image_guidance_scale - - @property - def num_timesteps(self): - return self._num_timesteps - - # 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.0 and self.image_guidance_scale >= 1.0 diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py deleted file mode 100644 index 66d5ffa6b8498ac9f9f9392b1c92f84fd52963a6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py +++ /dev/null @@ -1,665 +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 warnings -from typing import Callable, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from transformers import CLIPTextModel, CLIPTokenizer - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin -from ...models import AutoencoderKL, UNet2DConditionModel -from ...schedulers import EulerDiscreteScheduler -from ...utils import deprecate, is_torch_xla_available, logging -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput, StableDiffusionMixin - - -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 - - -# 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" -): - 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_upscale.preprocess -def preprocess(image): - warnings.warn( - "The preprocess method is deprecated and will be removed in a future version. Please" - " use VaeImageProcessor.preprocess instead", - FutureWarning, - ) - if isinstance(image, torch.Tensor): - return image - elif isinstance(image, PIL.Image.Image): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - w, h = image[0].size - w, h = (x - x % 64 for x in (w, h)) # resize to integer multiple of 64 - - image = [np.array(i.resize((w, h)))[None, :] for i in image] - image = np.concatenate(image, axis=0) - image = np.array(image).astype(np.float32) / 255.0 - image = image.transpose(0, 3, 1, 2) - image = 2.0 * image - 1.0 - image = torch.from_numpy(image) - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, dim=0) - return image - - -class StableDiffusionLatentUpscalePipeline(DiffusionPipeline, StableDiffusionMixin, FromSingleFileMixin): - r""" - Pipeline for upscaling Stable Diffusion output image resolution by a factor of 2. - - 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.). - - The pipeline also inherits the following loading methods: - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A [`EulerDiscreteScheduler`] to be used in combination with `unet` to denoise the encoded image latents. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: EulerDiscreteScheduler, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - 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, resample="bicubic") - - def _encode_prompt( - self, - 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, - **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, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - device=device, - do_classifier_free_guidance=do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - **kwargs, - ) - - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds]) - - return prompt_embeds, pooled_prompt_embeds - - def encode_prompt( - self, - 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, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `list(int)`): - prompt to be encoded - device: (`torch.device`): - torch device - do_classifier_free_guidance (`bool`): - whether to use classifier free guidance or not - 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*): - 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. - 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. - """ - 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 or pooled_prompt_embeds is None: - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_length=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}" - ) - - text_encoder_out = self.text_encoder( - text_input_ids.to(device), - output_hidden_states=True, - ) - prompt_embeds = text_encoder_out.hidden_states[-1] - pooled_prompt_embeds = text_encoder_out.pooler_output - - # 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] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif 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 - - max_length = text_input_ids.shape[-1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_length=True, - return_tensors="pt", - ) - - uncond_encoder_out = self.text_encoder( - uncond_input.input_ids.to(device), - output_hidden_states=True, - ) - - negative_prompt_embeds = uncond_encoder_out.hidden_states[-1] - negative_pooled_prompt_embeds = uncond_encoder_out.pooler_output - - return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds - - # 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, - image, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_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 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 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`." - ) - - if ( - not isinstance(image, torch.Tensor) - and not isinstance(image, np.ndarray) - and not isinstance(image, PIL.Image.Image) - and not isinstance(image, list) - ): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or `list` but is {type(image)}" - ) - - # verify batch size of prompt and image are same if image is a list or tensor - if isinstance(image, (list, torch.Tensor)): - if prompt is not None: - if isinstance(prompt, str): - batch_size = 1 - else: - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - if isinstance(image, list): - image_batch_size = len(image) - else: - image_batch_size = image.shape[0] if image.ndim == 4 else 1 - if batch_size != image_batch_size: - raise ValueError( - f"`prompt` has batch size {batch_size} and `image` has batch size {image_batch_size}." - " Please make sure that passed `prompt` matches the batch size of `image`." - ) - - if (callback_steps is None) or ( - 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)}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.prepare_latents - def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): - shape = (batch_size, num_channels_latents, height, width) - 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 - latents = latents * self.scheduler.init_noise_sigma - return latents - - @torch.no_grad() - def __call__( - self, - prompt: Union[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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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` 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 - encoded using this pipeline's `vae` encoder. - 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): - 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*): - 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*): - 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`. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - - Examples: - ```py - >>> from diffusers import StableDiffusionLatentUpscalePipeline, StableDiffusionPipeline - >>> import torch - - - >>> pipeline = StableDiffusionPipeline.from_pretrained( - ... "CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16 - ... ) - >>> pipeline.to("cuda") - - >>> model_id = "stabilityai/sd-x2-latent-upscaler" - >>> upscaler = StableDiffusionLatentUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16) - >>> upscaler.to("cuda") - - >>> prompt = "a photo of an astronaut high resolution, unreal engine, ultra realistic" - >>> generator = torch.manual_seed(33) - - >>> low_res_latents = pipeline(prompt, generator=generator, output_type="latent").images - - >>> with torch.no_grad(): - ... image = pipeline.decode_latents(low_res_latents) - >>> image = pipeline.numpy_to_pil(image)[0] - - >>> image.save("../images/a1.png") - - >>> upscaled_image = upscaler( - ... prompt=prompt, - ... image=low_res_latents, - ... num_inference_steps=20, - ... guidance_scale=0, - ... generator=generator, - ... ).images[0] - - >>> upscaled_image.save("../images/a2.png") - ``` - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned, - otherwise a `tuple` is returned where the first element is a list with the generated images. - """ - - # 1. Check inputs - self.check_inputs( - prompt, - image, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) - - # 2. Define call parameters - if prompt is not None: - batch_size = 1 if isinstance(prompt, str) else 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 = guidance_scale > 1.0 - - if guidance_scale == 0: - prompt = [""] * batch_size - - # 3. Encode input prompt - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt, - device, - do_classifier_free_guidance, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) - - if do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds]) - - # 4. Preprocess image - image = self.image_processor.preprocess(image) - image = image.to(dtype=prompt_embeds.dtype, device=device) - if image.shape[1] == 3: - # encode image if not in latent-space yet - image = retrieve_latents(self.vae.encode(image), generator=generator) * self.vae.config.scaling_factor - - # 5. set timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - batch_multiplier = 2 if do_classifier_free_guidance else 1 - image = image[None, :] if image.ndim == 3 else image - image = torch.cat([image] * batch_multiplier) - - # 5. Add noise to image (set to be 0): - # (see below notes from the author): - # "the This step theoretically can make the model work better on out-of-distribution inputs, but mostly just seems to make it match the input less, so it's turned off by default." - noise_level = torch.tensor([0.0], dtype=torch.float32, device=device) - noise_level = torch.cat([noise_level] * image.shape[0]) - inv_noise_level = (noise_level**2 + 1) ** (-0.5) - - image_cond = F.interpolate(image, scale_factor=2, mode="nearest") * inv_noise_level[:, None, None, None] - image_cond = image_cond.to(prompt_embeds.dtype) - - noise_level_embed = torch.cat( - [ - torch.ones(pooled_prompt_embeds.shape[0], 64, dtype=pooled_prompt_embeds.dtype, device=device), - torch.zeros(pooled_prompt_embeds.shape[0], 64, dtype=pooled_prompt_embeds.dtype, device=device), - ], - dim=1, - ) - - timestep_condition = torch.cat([noise_level_embed, pooled_prompt_embeds], dim=1) - - # 6. Prepare latent variables - height, width = image.shape[2:] - num_channels_latents = self.vae.config.latent_channels - latents = self.prepare_latents( - batch_size, - num_channels_latents, - height * 2, # 2x upscale - width * 2, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 7. Check that sizes of image and latents match - num_channels_image = image.shape[1] - if num_channels_latents + num_channels_image != self.unet.config.in_channels: - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_image`: {num_channels_image} " - f" = {num_channels_latents + num_channels_image}. Please verify the config of" - " `pipeline.unet` or your `image` input." - ) - - # 9. Denoising loop - num_warmup_steps = 0 - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - sigma = self.scheduler.sigmas[i] - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - scaled_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - scaled_model_input = torch.cat([scaled_model_input, image_cond], dim=1) - # preconditioning parameter based on Karras et al. (2022) (table 1) - timestep = torch.log(sigma) * 0.25 - - noise_pred = self.unet( - scaled_model_input, - timestep, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_condition, - ).sample - - # in original repo, the output contains a variance channel that's not used - noise_pred = noise_pred[:, :-1] - - # apply preconditioning, based on table 1 in Karras et al. (2022) - inv_sigma = 1 / (sigma**2 + 1) - noise_pred = inv_sigma * latent_model_input + self.scheduler.scale_model_input(sigma, t) * noise_pred - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - else: - image = latents - - image = self.image_processor.postprocess(image, output_type=output_type) - - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py deleted file mode 100644 index 02f7526b71028b2a1eb7aa69e5cf62c380850c94..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +++ /dev/null @@ -1,826 +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 inspect -import warnings -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import DDPMScheduler, KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_torch_xla_available, - logging, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from . import StableDiffusionPipelineOutput - - -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 preprocess(image): - warnings.warn( - "The preprocess method is deprecated and will be removed in a future version. Please" - " use VaeImageProcessor.preprocess instead", - FutureWarning, - ) - if isinstance(image, torch.Tensor): - return image - elif isinstance(image, PIL.Image.Image): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - w, h = image[0].size - w, h = (x - x % 64 for x in (w, h)) # resize to integer multiple of 64 - - image = [np.array(i.resize((w, h)))[None, :] for i in image] - image = np.concatenate(image, axis=0) - image = np.array(image).astype(np.float32) / 255.0 - image = image.transpose(0, 3, 1, 2) - image = 2.0 * image - 1.0 - image = torch.from_numpy(image) - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, dim=0) - return image - - -class StableDiffusionUpscalePipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-guided image super-resolution using Stable Diffusion 2. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - low_res_scheduler ([`SchedulerMixin`]): - A scheduler used to add initial noise to the low resolution conditioning image. It must be an instance of - [`DDPMScheduler`]. - 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`]. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["watermarker", "safety_checker", "feature_extractor"] - _exclude_from_cpu_offload = ["safety_checker"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - low_res_scheduler: DDPMScheduler, - scheduler: KarrasDiffusionSchedulers, - safety_checker: Optional[Any] = None, - feature_extractor: Optional[CLIPImageProcessor] = None, - watermarker: Optional[Any] = None, - max_noise_level: int = 350, - ): - super().__init__() - - if hasattr( - vae, "config" - ): # check if vae has a config attribute `scaling_factor` and if it is set to 0.08333, else set it to 0.08333 and deprecate - is_vae_scaling_factor_set_to_0_08333 = ( - hasattr(vae.config, "scaling_factor") and vae.config.scaling_factor == 0.08333 - ) - if not is_vae_scaling_factor_set_to_0_08333: - deprecation_message = ( - "The configuration file of the vae does not contain `scaling_factor` or it is set to" - f" {vae.config.scaling_factor}, which seems highly unlikely. If your checkpoint is a fine-tuned" - " version of `stabilityai/stable-diffusion-x4-upscaler` you should change 'scaling_factor' to" - " 0.08333 Please make sure to update the config accordingly, as not doing so might lead to" - " incorrect results in future versions. If you have downloaded this checkpoint from the Hugging" - " Face Hub, it would be very nice if you could open a Pull Request for the `vae/config.json` file" - ) - deprecate("wrong scaling_factor", "1.0.0", deprecation_message, standard_warn=False) - vae.register_to_config(scaling_factor=0.08333) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - low_res_scheduler=low_res_scheduler, - scheduler=scheduler, - safety_checker=safety_checker, - watermarker=watermarker, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, resample="bicubic") - self.register_to_config(max_noise_level=max_noise_level) - - def run_safety_checker(self, image, device, dtype): - if self.safety_checker is not None: - feature_extractor_input = self.image_processor.postprocess(image, output_type="pil") - safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device) - image, nsfw_detected, watermark_detected = self.safety_checker( - images=image, - clip_input=safety_checker_input.pixel_values.to(dtype=dtype), - ) - else: - nsfw_detected = None - watermark_detected = None - - if hasattr(self, "unet_offload_hook") and self.unet_offload_hook is not None: - self.unet_offload_hook.offload() - - return image, nsfw_detected, watermark_detected - - # 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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # 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, - image, - noise_level, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - if (callback_steps is None) or ( - 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 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 ( - not isinstance(image, torch.Tensor) - and not isinstance(image, PIL.Image.Image) - and not isinstance(image, np.ndarray) - and not isinstance(image, list) - ): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `np.ndarray`, `PIL.Image.Image` or `list` but is {type(image)}" - ) - - # verify batch size of prompt and image are same if image is a list or tensor or numpy array - if isinstance(image, (list, np.ndarray, torch.Tensor)): - 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 isinstance(image, list): - image_batch_size = len(image) - else: - image_batch_size = image.shape[0] - if batch_size != image_batch_size: - raise ValueError( - f"`prompt` has batch size {batch_size} and `image` has batch size {image_batch_size}." - " Please make sure that passed `prompt` matches the batch size of `image`." - ) - - # check noise level - if noise_level > self.config.max_noise_level: - raise ValueError(f"`noise_level` has to be <= {self.config.max_noise_level} but is {noise_level}") - - if (callback_steps is None) or ( - 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)}." - ) - - def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): - shape = (batch_size, num_channels_latents, height, width) - 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 - latents = latents * self.scheduler.init_noise_sigma - return latents - - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - @torch.no_grad() - def __call__( - self, - prompt: Union[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, - 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, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: int = None, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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` 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 - expense of slower inference. - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - 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. - Examples: - ```py - >>> import requests - >>> from PIL import Image - >>> from io import BytesIO - >>> from diffusers import StableDiffusionUpscalePipeline - >>> import torch - - >>> # load model and scheduler - >>> model_id = "stabilityai/stable-diffusion-x4-upscaler" - >>> pipeline = StableDiffusionUpscalePipeline.from_pretrained( - ... model_id, variant="fp16", torch_dtype=torch.float16 - ... ) - >>> pipeline = pipeline.to("cuda") - - >>> # let's download an image - >>> url = "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale/low_res_cat.png" - >>> response = requests.get(url) - >>> low_res_img = Image.open(BytesIO(response.content)).convert("RGB") - >>> low_res_img = low_res_img.resize((128, 128)) - >>> prompt = "a white cat" - - >>> upscaled_image = pipeline(prompt=prompt, image=low_res_img).images[0] - >>> upscaled_image.save("upsampled_cat.png") - ``` - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - - # 1. Check inputs - self.check_inputs( - prompt, - image, - noise_level, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ) - - if image is None: - raise ValueError("`image` input cannot be undefined.") - - # 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - 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, - lora_scale=text_encoder_lora_scale, - 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. Preprocess image - image = self.image_processor.preprocess(image) - image = image.to(dtype=prompt_embeds.dtype, device=device) - - # 5. set timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Add noise to image - noise_level = torch.tensor([noise_level], dtype=torch.long, device=device) - noise = randn_tensor(image.shape, generator=generator, device=device, dtype=prompt_embeds.dtype) - image = self.low_res_scheduler.add_noise(image, noise, noise_level) - - batch_multiplier = 2 if do_classifier_free_guidance else 1 - image = torch.cat([image] * batch_multiplier * num_images_per_prompt) - noise_level = torch.cat([noise_level] * image.shape[0]) - - # 6. Prepare latent variables - height, width = image.shape[2:] - num_channels_latents = self.vae.config.latent_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 7. Check that sizes of image and latents match - num_channels_image = image.shape[1] - if num_channels_latents + num_channels_image != self.unet.config.in_channels: - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_image`: {num_channels_image} " - f" = {num_channels_latents + num_channels_image}. Please verify the config of" - " `pipeline.unet` or your `image` input." - ) - - # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 9. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - # concat latents, mask, masked_image_latents in the channel dimension - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - latent_model_input = torch.cat([latent_model_input, image], dim=1) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - class_labels=noise_level, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - 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() - - # Ensure latents are always the same type as the 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) - - 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) - - # 11. Apply watermark - if output_type == "pil" and self.watermarker is not None: - image = self.watermarker.apply_watermark(image) - - # 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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py deleted file mode 100644 index a134244e3ee4ecfbc881fe0c4b2f917a8fb153a4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py +++ /dev/null @@ -1,952 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer -from transformers.models.clip.modeling_clip import CLIPTextModelOutput - -from ...image_processor import VaeImageProcessor -from ...loaders import StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, PriorTransformer, UNet2DConditionModel -from ...models.embeddings import get_timestep_embedding -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - 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, ImagePipelineOutput, StableDiffusionMixin -from .stable_unclip_image_normalizer import StableUnCLIPImageNormalizer - - -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 StableUnCLIPPipeline - - >>> pipe = StableUnCLIPPipeline.from_pretrained( - ... "fusing/stable-unclip-2-1-l", torch_dtype=torch.float16 - ... ) # TODO update model path - >>> pipe = pipe.to("cuda") - - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> images = pipe(prompt).images - >>> images[0].save("astronaut_horse.png") - ``` -""" - - -class StableUnCLIPPipeline( - DiffusionPipeline, StableDiffusionMixin, TextualInversionLoaderMixin, StableDiffusionLoraLoaderMixin -): - """ - Pipeline for text-to-image generation using stable unCLIP. - - 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.). - - 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 - - Args: - prior_tokenizer ([`CLIPTokenizer`]): - A [`CLIPTokenizer`]. - prior_text_encoder ([`CLIPTextModelWithProjection`]): - Frozen [`CLIPTextModelWithProjection`] text-encoder. - prior ([`PriorTransformer`]): - The canonical unCLIP prior to approximate the image embedding from the text embedding. - prior_scheduler ([`KarrasDiffusionSchedulers`]): - Scheduler used in the prior denoising process. - image_normalizer ([`StableUnCLIPImageNormalizer`]): - Used to normalize the predicted image embeddings before the noise is applied and un-normalize the image - embeddings after the noise has been applied. - image_noising_scheduler ([`KarrasDiffusionSchedulers`]): - Noise schedule for adding noise to the predicted image embeddings. The amount of noise to add is determined - by the `noise_level`. - tokenizer ([`CLIPTokenizer`]): - A [`CLIPTokenizer`]. - text_encoder ([`CLIPTextModel`]): - Frozen [`CLIPTextModel`] text-encoder. - unet ([`UNet2DConditionModel`]): - A [`UNet2DConditionModel`] to denoise the encoded image latents. - scheduler ([`KarrasDiffusionSchedulers`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - """ - - _exclude_from_cpu_offload = ["prior", "image_normalizer"] - model_cpu_offload_seq = "text_encoder->prior_text_encoder->unet->vae" - - # prior components - prior_tokenizer: CLIPTokenizer - prior_text_encoder: CLIPTextModelWithProjection - prior: PriorTransformer - prior_scheduler: KarrasDiffusionSchedulers - - # image noising components - image_normalizer: StableUnCLIPImageNormalizer - image_noising_scheduler: KarrasDiffusionSchedulers - - # regular denoising components - tokenizer: CLIPTokenizer - text_encoder: CLIPTextModel - unet: UNet2DConditionModel - scheduler: KarrasDiffusionSchedulers - - vae: AutoencoderKL - - def __init__( - self, - # prior components - prior_tokenizer: CLIPTokenizer, - prior_text_encoder: CLIPTextModelWithProjection, - prior: PriorTransformer, - prior_scheduler: KarrasDiffusionSchedulers, - # image noising components - image_normalizer: StableUnCLIPImageNormalizer, - image_noising_scheduler: KarrasDiffusionSchedulers, - # regular denoising components - tokenizer: CLIPTokenizer, - text_encoder: CLIPTextModel, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - # vae - vae: AutoencoderKL, - ): - super().__init__() - - self.register_modules( - prior_tokenizer=prior_tokenizer, - prior_text_encoder=prior_text_encoder, - prior=prior, - prior_scheduler=prior_scheduler, - image_normalizer=image_normalizer, - image_noising_scheduler=image_noising_scheduler, - tokenizer=tokenizer, - text_encoder=text_encoder, - unet=unet, - scheduler=scheduler, - vae=vae, - ) - - 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) - - # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline._encode_prompt with _encode_prompt->_encode_prior_prompt, tokenizer->prior_tokenizer, text_encoder->prior_text_encoder - def _encode_prior_prompt( - self, - 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, - ): - if text_model_output is None: - batch_size = len(prompt) if isinstance(prompt, list) else 1 - # get prompt text embeddings - text_inputs = self.prior_tokenizer( - prompt, - padding="max_length", - max_length=self.prior_tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - text_mask = text_inputs.attention_mask.bool().to(device) - - untruncated_ids = self.prior_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.prior_tokenizer.batch_decode( - untruncated_ids[:, self.prior_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.prior_tokenizer.model_max_length} tokens: {removed_text}" - ) - text_input_ids = text_input_ids[:, : self.prior_tokenizer.model_max_length] - - prior_text_encoder_output = self.prior_text_encoder(text_input_ids.to(device)) - - prompt_embeds = prior_text_encoder_output.text_embeds - text_enc_hid_states = prior_text_encoder_output.last_hidden_state - - else: - batch_size = text_model_output[0].shape[0] - prompt_embeds, text_enc_hid_states = text_model_output[0], text_model_output[1] - text_mask = text_attention_mask - - prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - text_enc_hid_states = text_enc_hid_states.repeat_interleave(num_images_per_prompt, dim=0) - text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - uncond_tokens = [""] * batch_size - - uncond_input = self.prior_tokenizer( - uncond_tokens, - padding="max_length", - max_length=self.prior_tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - uncond_text_mask = uncond_input.attention_mask.bool().to(device) - negative_prompt_embeds_prior_text_encoder_output = self.prior_text_encoder( - uncond_input.input_ids.to(device) - ) - - negative_prompt_embeds = negative_prompt_embeds_prior_text_encoder_output.text_embeds - uncond_text_enc_hid_states = negative_prompt_embeds_prior_text_encoder_output.last_hidden_state - - # 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.repeat(1, num_images_per_prompt) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len) - - seq_len = uncond_text_enc_hid_states.shape[1] - uncond_text_enc_hid_states = uncond_text_enc_hid_states.repeat(1, num_images_per_prompt, 1) - uncond_text_enc_hid_states = uncond_text_enc_hid_states.view( - batch_size * num_images_per_prompt, seq_len, -1 - ) - uncond_text_mask = uncond_text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - # done duplicates - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - text_enc_hid_states = torch.cat([uncond_text_enc_hid_states, text_enc_hid_states]) - - text_mask = torch.cat([uncond_text_mask, text_mask]) - - return prompt_embeds, text_enc_hid_states, text_mask - - # 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.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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs with prepare_extra_step_kwargs->prepare_prior_extra_step_kwargs, scheduler->prior_scheduler - def prepare_prior_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the prior_scheduler step, since not all prior_schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other prior_schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.prior_scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the prior_scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.prior_scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - callback_steps, - noise_level, - negative_prompt=None, - prompt_embeds=None, - negative_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 (callback_steps is None) or ( - 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 prompt is not None and prompt_embeds is not None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds`. Please make sure to define only one of the two." - ) - - if prompt is None and prompt_embeds is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." - ) - - 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)}") - - if negative_prompt is not None and negative_prompt_embeds is not None: - raise ValueError( - "Provide either `negative_prompt` or `negative_prompt_embeds`. Cannot leave both `negative_prompt` and `negative_prompt_embeds` undefined." - ) - - if prompt is not None and negative_prompt is not None: - 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)} !=" - f" {type(prompt)}." - ) - - 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 noise_level < 0 or noise_level >= self.image_noising_scheduler.config.num_train_timesteps: - raise ValueError( - f"`noise_level` must be between 0 and {self.image_noising_scheduler.config.num_train_timesteps - 1}, inclusive." - ) - - # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - def noise_image_embeddings( - self, - image_embeds: torch.Tensor, - noise_level: int, - noise: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, - ): - """ - Add noise to the image embeddings. The amount of noise is controlled by a `noise_level` input. A higher - `noise_level` increases the variance in the final un-noised images. - - The noise is applied in two ways: - 1. A noise schedule is applied directly to the embeddings. - 2. A vector of sinusoidal time embeddings are appended to the output. - - In both cases, the amount of noise is controlled by the same `noise_level`. - - The embeddings are normalized before the noise is applied and un-normalized after the noise is applied. - """ - if noise is None: - noise = randn_tensor( - image_embeds.shape, generator=generator, device=image_embeds.device, dtype=image_embeds.dtype - ) - - noise_level = torch.tensor([noise_level] * image_embeds.shape[0], device=image_embeds.device) - - self.image_normalizer.to(image_embeds.device) - image_embeds = self.image_normalizer.scale(image_embeds) - - image_embeds = self.image_noising_scheduler.add_noise(image_embeds, timesteps=noise_level, noise=noise) - - image_embeds = self.image_normalizer.unscale(image_embeds) - - noise_level = get_timestep_embedding( - timesteps=noise_level, embedding_dim=image_embeds.shape[-1], flip_sin_to_cos=True, downscale_freq_shift=0 - ) - - # `get_timestep_embeddings` does not contain any weights and will always return f32 tensors, - # but we might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - noise_level = noise_level.to(image_embeds.dtype) - - image_embeds = torch.cat((image_embeds, noise_level), 1) - - return image_embeds - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - # regular denoising process args - prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = 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, - 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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, - ): - """ - The call function to the pipeline for generation. - - Args: - 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. - 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 20): - 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 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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 [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - noise_level (`int`, *optional*, defaults to `0`): - The amount of noise to add to the image embeddings. A higher `noise_level` increases the variance in - the final un-noised images. See [`StableUnCLIPPipeline.noise_image_embeddings`] for more details. - prior_num_inference_steps (`int`, *optional*, defaults to 25): - The number of denoising steps in the prior denoising process. More denoising steps usually lead to a - higher quality image at the expense of slower inference. - prior_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`. - prior_latents (`torch.Tensor`, *optional*): - Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image - embedding generation in the prior denoising process. 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`. - 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. - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - [`~ pipeline_utils.ImagePipelineOutput`] 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.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=prompt, - height=height, - width=width, - callback_steps=callback_steps, - noise_level=noise_level, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_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] - - batch_size = batch_size * num_images_per_prompt - - 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. - prior_do_classifier_free_guidance = prior_guidance_scale > 1.0 - - # 3. Encode input prompt - prior_prompt_embeds, prior_text_encoder_hidden_states, prior_text_mask = self._encode_prior_prompt( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=prior_do_classifier_free_guidance, - ) - - # 4. Prepare prior timesteps - self.prior_scheduler.set_timesteps(prior_num_inference_steps, device=device) - prior_timesteps_tensor = self.prior_scheduler.timesteps - - # 5. Prepare prior latent variables - embedding_dim = self.prior.config.embedding_dim - prior_latents = self.prepare_latents( - (batch_size, embedding_dim), - prior_prompt_embeds.dtype, - device, - generator, - prior_latents, - self.prior_scheduler, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - prior_extra_step_kwargs = self.prepare_prior_extra_step_kwargs(generator, eta) - - # 7. Prior denoising loop - for i, t in enumerate(self.progress_bar(prior_timesteps_tensor)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([prior_latents] * 2) if prior_do_classifier_free_guidance else prior_latents - latent_model_input = self.prior_scheduler.scale_model_input(latent_model_input, t) - - predicted_image_embedding = self.prior( - latent_model_input, - timestep=t, - proj_embedding=prior_prompt_embeds, - encoder_hidden_states=prior_text_encoder_hidden_states, - attention_mask=prior_text_mask, - ).predicted_image_embedding - - if prior_do_classifier_free_guidance: - predicted_image_embedding_uncond, predicted_image_embedding_text = predicted_image_embedding.chunk(2) - predicted_image_embedding = predicted_image_embedding_uncond + prior_guidance_scale * ( - predicted_image_embedding_text - predicted_image_embedding_uncond - ) - - prior_latents = self.prior_scheduler.step( - predicted_image_embedding, - timestep=t, - sample=prior_latents, - **prior_extra_step_kwargs, - return_dict=False, - )[0] - - if callback is not None and i % callback_steps == 0: - callback(i, t, prior_latents) - - prior_latents = self.prior.post_process_latents(prior_latents) - - image_embeds = prior_latents - - # done prior - - # 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 = guidance_scale > 1.0 - - # 8. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = 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=text_encoder_lora_scale, - 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]) - - # 9. Prepare image embeddings - image_embeds = self.noise_image_embeddings( - image_embeds=image_embeds, - noise_level=noise_level, - generator=generator, - ) - - if do_classifier_free_guidance: - negative_prompt_embeds = torch.zeros_like(image_embeds) - - # 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 - image_embeds = torch.cat([negative_prompt_embeds, image_embeds]) - - # 10. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 11. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - shape = ( - batch_size, - num_channels_latents, - int(height) // self.vae_scale_factor, - int(width) // self.vae_scale_factor, - ) - latents = self.prepare_latents( - shape=shape, - dtype=prompt_embeds.dtype, - device=device, - generator=generator, - latents=latents, - scheduler=self.scheduler, - ) - - # 12. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 13. Denoising loop - for i, t in enumerate(self.progress_bar(timesteps)): - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - class_labels=image_embeds, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - else: - image = latents - - 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 ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py deleted file mode 100644 index abb4cc3a05d51ed0164120abb688267ec5e2af58..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py +++ /dev/null @@ -1,858 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...image_processor import VaeImageProcessor -from ...loaders import StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, UNet2DConditionModel -from ...models.embeddings import get_timestep_embedding -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - 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, ImagePipelineOutput, StableDiffusionMixin -from .stable_unclip_image_normalizer import StableUnCLIPImageNormalizer - - -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 requests - >>> import torch - >>> from PIL import Image - >>> from io import BytesIO - - >>> from diffusers import StableUnCLIPImg2ImgPipeline - - >>> pipe = StableUnCLIPImg2ImgPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-2-1-unclip-small", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" - - >>> response = requests.get(url) - >>> init_image = Image.open(BytesIO(response.content)).convert("RGB") - >>> init_image = init_image.resize((768, 512)) - - >>> prompt = "A fantasy landscape, trending on artstation" - - >>> images = pipe(init_image, prompt).images - >>> images[0].save("fantasy_landscape.png") - ``` -""" - - -class StableUnCLIPImg2ImgPipeline( - DiffusionPipeline, StableDiffusionMixin, TextualInversionLoaderMixin, StableDiffusionLoraLoaderMixin -): - """ - Pipeline for text-guided image-to-image generation using stable unCLIP. - - 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.). - - 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 - - Args: - feature_extractor ([`CLIPImageProcessor`]): - Feature extractor for image pre-processing before being encoded. - image_encoder ([`CLIPVisionModelWithProjection`]): - CLIP vision model for encoding images. - image_normalizer ([`StableUnCLIPImageNormalizer`]): - Used to normalize the predicted image embeddings before the noise is applied and un-normalize the image - embeddings after the noise has been applied. - image_noising_scheduler ([`KarrasDiffusionSchedulers`]): - Noise schedule for adding noise to the predicted image embeddings. The amount of noise to add is determined - by the `noise_level`. - tokenizer (`~transformers.CLIPTokenizer`): - A [`~transformers.CLIPTokenizer`)]. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen [`~transformers.CLIPTextModel`] text-encoder. - unet ([`UNet2DConditionModel`]): - A [`UNet2DConditionModel`] to denoise the encoded image latents. - scheduler ([`KarrasDiffusionSchedulers`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" - _exclude_from_cpu_offload = ["image_normalizer"] - - # image encoding components - feature_extractor: CLIPImageProcessor - image_encoder: CLIPVisionModelWithProjection - - # image noising components - image_normalizer: StableUnCLIPImageNormalizer - image_noising_scheduler: KarrasDiffusionSchedulers - - # regular denoising components - tokenizer: CLIPTokenizer - text_encoder: CLIPTextModel - unet: UNet2DConditionModel - scheduler: KarrasDiffusionSchedulers - - vae: AutoencoderKL - - def __init__( - self, - # image encoding components - feature_extractor: CLIPImageProcessor, - image_encoder: CLIPVisionModelWithProjection, - # image noising components - image_normalizer: StableUnCLIPImageNormalizer, - image_noising_scheduler: KarrasDiffusionSchedulers, - # regular denoising components - tokenizer: CLIPTokenizer, - text_encoder: CLIPTextModel, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - # vae - vae: AutoencoderKL, - ): - super().__init__() - - self.register_modules( - feature_extractor=feature_extractor, - image_encoder=image_encoder, - image_normalizer=image_normalizer, - image_noising_scheduler=image_noising_scheduler, - tokenizer=tokenizer, - text_encoder=text_encoder, - unet=unet, - scheduler=scheduler, - vae=vae, - ) - - 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) - - # 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 - - def _encode_image( - self, - image, - device, - batch_size, - num_images_per_prompt, - do_classifier_free_guidance, - noise_level, - generator, - image_embeds, - ): - dtype = next(self.image_encoder.parameters()).dtype - - if isinstance(image, PIL.Image.Image): - # the image embedding should repeated so it matches the total batch size of the prompt - repeat_by = batch_size - else: - # assume the image input is already properly batched and just needs to be repeated so - # it matches the num_images_per_prompt. - # - # NOTE(will) this is probably missing a few number of side cases. I.e. batched/non-batched - # `image_embeds`. If those happen to be common use cases, let's think harder about - # what the expected dimensions of inputs should be and how we handle the encoding. - repeat_by = num_images_per_prompt - - if image_embeds is None: - if not isinstance(image, torch.Tensor): - image = self.feature_extractor(images=image, return_tensors="pt").pixel_values - - image = image.to(device=device, dtype=dtype) - image_embeds = self.image_encoder(image).image_embeds - - image_embeds = self.noise_image_embeddings( - image_embeds=image_embeds, - noise_level=noise_level, - generator=generator, - ) - - # duplicate image embeddings for each generation per prompt, using mps friendly method - image_embeds = image_embeds.unsqueeze(1) - bs_embed, seq_len, _ = image_embeds.shape - image_embeds = image_embeds.repeat(1, repeat_by, 1) - image_embeds = image_embeds.view(bs_embed * repeat_by, seq_len, -1) - image_embeds = image_embeds.squeeze(1) - - if do_classifier_free_guidance: - negative_prompt_embeds = torch.zeros_like(image_embeds) - - # 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 - image_embeds = torch.cat([negative_prompt_embeds, image_embeds]) - - return image_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.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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - image, - height, - width, - callback_steps, - noise_level, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - image_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 (callback_steps is None) or ( - 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 prompt is not None and prompt_embeds is not None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds`. Please make sure to define only one of the two." - ) - - if prompt is None and prompt_embeds is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." - ) - - 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)}") - - if negative_prompt is not None and negative_prompt_embeds is not None: - raise ValueError( - "Provide either `negative_prompt` or `negative_prompt_embeds`. Cannot leave both `negative_prompt` and `negative_prompt_embeds` undefined." - ) - - if prompt is not None and negative_prompt is not None: - 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)} !=" - f" {type(prompt)}." - ) - - 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 noise_level < 0 or noise_level >= self.image_noising_scheduler.config.num_train_timesteps: - raise ValueError( - f"`noise_level` must be between 0 and {self.image_noising_scheduler.config.num_train_timesteps - 1}, inclusive." - ) - - if image is not None and image_embeds is not None: - raise ValueError( - "Provide either `image` or `image_embeds`. Please make sure to define only one of the two." - ) - - if image is None and image_embeds is None: - raise ValueError( - "Provide either `image` or `image_embeds`. Cannot leave both `image` and `image_embeds` undefined." - ) - - if image is not None: - if ( - not isinstance(image, torch.Tensor) - and not isinstance(image, PIL.Image.Image) - 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" - f" {type(image)}" - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_unclip.StableUnCLIPPipeline.noise_image_embeddings - def noise_image_embeddings( - self, - image_embeds: torch.Tensor, - noise_level: int, - noise: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, - ): - """ - Add noise to the image embeddings. The amount of noise is controlled by a `noise_level` input. A higher - `noise_level` increases the variance in the final un-noised images. - - The noise is applied in two ways: - 1. A noise schedule is applied directly to the embeddings. - 2. A vector of sinusoidal time embeddings are appended to the output. - - In both cases, the amount of noise is controlled by the same `noise_level`. - - The embeddings are normalized before the noise is applied and un-normalized after the noise is applied. - """ - if noise is None: - noise = randn_tensor( - image_embeds.shape, generator=generator, device=image_embeds.device, dtype=image_embeds.dtype - ) - - noise_level = torch.tensor([noise_level] * image_embeds.shape[0], device=image_embeds.device) - - self.image_normalizer.to(image_embeds.device) - image_embeds = self.image_normalizer.scale(image_embeds) - - image_embeds = self.image_noising_scheduler.add_noise(image_embeds, timesteps=noise_level, noise=noise) - - image_embeds = self.image_normalizer.unscale(image_embeds) - - noise_level = get_timestep_embedding( - timesteps=noise_level, embedding_dim=image_embeds.shape[-1], flip_sin_to_cos=True, downscale_freq_shift=0 - ) - - # `get_timestep_embeddings` does not contain any weights and will always return f32 tensors, - # but we might actually be running in fp16. so we need to cast here. - # there might be better ways to encapsulate this. - noise_level = noise_level.to(image_embeds.dtype) - - image_embeds = torch.cat((image_embeds, noise_level), 1) - - return image_embeds - - @torch.no_grad() - @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, - num_inference_steps: int = 20, - guidance_scale: float = 10, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - noise_level: int = 0, - image_embeds: Optional[torch.Tensor] = None, - clip_skip: Optional[int] = None, - ): - 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, either `prompt_embeds` will be - used or prompt is initialized to `""`. - image (`torch.Tensor` or `PIL.Image.Image`): - `Image` or tensor representing an image batch. The image is encoded to its CLIP embedding which the - `unet` is conditioned on. The image is _not_ encoded by the `vae` and then used as the latents in the - denoising process like it is in the standard Stable Diffusion text-guided image variation process. - 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 20): - 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 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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 [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - noise_level (`int`, *optional*, defaults to `0`): - The amount of noise to add to the image embeddings. A higher `noise_level` increases the variance in - the final un-noised images. See [`StableUnCLIPPipeline.noise_image_embeddings`] for more details. - image_embeds (`torch.Tensor`, *optional*): - Pre-generated CLIP embeddings to condition the `unet` on. These latents are not used in the denoising - process. If you want to provide pre-generated latents, pass them to `__call__` as `latents`. - 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. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - [`~ pipeline_utils.ImagePipelineOutput`] 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.unet.config.sample_size * self.vae_scale_factor - width = width or self.unet.config.sample_size * self.vae_scale_factor - - if prompt is None and prompt_embeds is None: - prompt = len(image) * [""] if isinstance(image, list) else "" - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt=prompt, - image=image, - height=height, - width=width, - callback_steps=callback_steps, - noise_level=noise_level, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - image_embeds=image_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] - - batch_size = batch_size * num_images_per_prompt - - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - prompt_embeds, negative_prompt_embeds = 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=text_encoder_lora_scale, - ) - # 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. Encoder input image - noise_level = torch.tensor([noise_level], device=device) - image_embeds = self._encode_image( - image=image, - device=device, - batch_size=batch_size, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=do_classifier_free_guidance, - noise_level=noise_level, - generator=generator, - image_embeds=image_embeds, - ) - - # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 6. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - if latents is None: - latents = self.prepare_latents( - batch_size=batch_size, - num_channels_latents=num_channels_latents, - height=height, - width=width, - dtype=prompt_embeds.dtype, - device=device, - generator=generator, - latents=latents, - ) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 8. Denoising loop - for i, t in enumerate(self.progress_bar(timesteps)): - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - class_labels=image_embeds, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # 9. Post-processing - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - else: - image = latents - - 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 ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/safety_checker.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/safety_checker.py deleted file mode 100644 index 16aff102599ca7fe3399dff9de54aafb9e5333df..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/safety_checker.py +++ /dev/null @@ -1,126 +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 numpy as np -import torch -import torch.nn as nn -from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel - -from ...utils import logging - - -logger = logging.get_logger(__name__) - - -def cosine_distance(image_embeds, text_embeds): - normalized_image_embeds = nn.functional.normalize(image_embeds) - normalized_text_embeds = nn.functional.normalize(text_embeds) - return torch.mm(normalized_image_embeds, normalized_text_embeds.t()) - - -class StableDiffusionSafetyChecker(PreTrainedModel): - config_class = CLIPConfig - main_input_name = "clip_input" - - _no_split_modules = ["CLIPEncoderLayer"] - - def __init__(self, config: CLIPConfig): - super().__init__(config) - - self.vision_model = CLIPVisionModel(config.vision_config) - self.visual_projection = nn.Linear(config.vision_config.hidden_size, config.projection_dim, bias=False) - - self.concept_embeds = nn.Parameter(torch.ones(17, config.projection_dim), requires_grad=False) - self.special_care_embeds = nn.Parameter(torch.ones(3, config.projection_dim), requires_grad=False) - - self.concept_embeds_weights = nn.Parameter(torch.ones(17), requires_grad=False) - self.special_care_embeds_weights = nn.Parameter(torch.ones(3), requires_grad=False) - - @torch.no_grad() - def forward(self, clip_input, images): - pooled_output = self.vision_model(clip_input)[1] # pooled_output - image_embeds = self.visual_projection(pooled_output) - - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds).cpu().float().numpy() - cos_dist = cosine_distance(image_embeds, self.concept_embeds).cpu().float().numpy() - - result = [] - batch_size = image_embeds.shape[0] - for i in range(batch_size): - result_img = {"special_scores": {}, "special_care": [], "concept_scores": {}, "bad_concepts": []} - - # increase this value to create a stronger `nfsw` filter - # at the cost of increasing the possibility of filtering benign images - adjustment = 0.0 - - for concept_idx in range(len(special_cos_dist[0])): - concept_cos = special_cos_dist[i][concept_idx] - concept_threshold = self.special_care_embeds_weights[concept_idx].item() - result_img["special_scores"][concept_idx] = round(concept_cos - concept_threshold + adjustment, 3) - if result_img["special_scores"][concept_idx] > 0: - result_img["special_care"].append({concept_idx, result_img["special_scores"][concept_idx]}) - adjustment = 0.01 - - for concept_idx in range(len(cos_dist[0])): - concept_cos = cos_dist[i][concept_idx] - concept_threshold = self.concept_embeds_weights[concept_idx].item() - result_img["concept_scores"][concept_idx] = round(concept_cos - concept_threshold + adjustment, 3) - if result_img["concept_scores"][concept_idx] > 0: - result_img["bad_concepts"].append(concept_idx) - - result.append(result_img) - - has_nsfw_concepts = [len(res["bad_concepts"]) > 0 for res in result] - - for idx, has_nsfw_concept in enumerate(has_nsfw_concepts): - if has_nsfw_concept: - if torch.is_tensor(images) or torch.is_tensor(images[0]): - images[idx] = torch.zeros_like(images[idx]) # black image - else: - images[idx] = np.zeros(images[idx].shape) # black image - - if any(has_nsfw_concepts): - logger.warning( - "Potential NSFW content was detected in one or more images. A black image will be returned instead." - " Try again with a different prompt and/or seed." - ) - - return images, has_nsfw_concepts - - @torch.no_grad() - def forward_onnx(self, clip_input: torch.Tensor, images: torch.Tensor): - pooled_output = self.vision_model(clip_input)[1] # pooled_output - image_embeds = self.visual_projection(pooled_output) - - special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds) - cos_dist = cosine_distance(image_embeds, self.concept_embeds) - - # increase this value to create a stronger `nsfw` filter - # at the cost of increasing the possibility of filtering benign images - adjustment = 0.0 - - special_scores = special_cos_dist - self.special_care_embeds_weights + adjustment - # special_scores = special_scores.round(decimals=3) - special_care = torch.any(special_scores > 0, dim=1) - special_adjustment = special_care * 0.01 - special_adjustment = special_adjustment.unsqueeze(1).expand(-1, cos_dist.shape[1]) - - concept_scores = (cos_dist - self.concept_embeds_weights) + special_adjustment - # concept_scores = concept_scores.round(decimals=3) - has_nsfw_concepts = torch.any(concept_scores > 0, dim=1) - - images[has_nsfw_concepts] = 0.0 # black image - - return images, has_nsfw_concepts diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/safety_checker_flax.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/safety_checker_flax.py deleted file mode 100644 index 55d5023b686994a731041a83bd716862812e3fac..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/safety_checker_flax.py +++ /dev/null @@ -1,112 +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 - -import jax -import jax.numpy as jnp -from flax import linen as nn -from flax.core.frozen_dict import FrozenDict -from transformers import CLIPConfig, FlaxPreTrainedModel -from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule - - -def jax_cosine_distance(emb_1, emb_2, eps=1e-12): - norm_emb_1 = jnp.divide(emb_1.T, jnp.clip(jnp.linalg.norm(emb_1, axis=1), a_min=eps)).T - norm_emb_2 = jnp.divide(emb_2.T, jnp.clip(jnp.linalg.norm(emb_2, axis=1), a_min=eps)).T - return jnp.matmul(norm_emb_1, norm_emb_2.T) - - -class FlaxStableDiffusionSafetyCheckerModule(nn.Module): - config: CLIPConfig - dtype: jnp.dtype = jnp.float32 - - def setup(self): - self.vision_model = FlaxCLIPVisionModule(self.config.vision_config) - self.visual_projection = nn.Dense(self.config.projection_dim, use_bias=False, dtype=self.dtype) - - self.concept_embeds = self.param("concept_embeds", jax.nn.initializers.ones, (17, self.config.projection_dim)) - self.special_care_embeds = self.param( - "special_care_embeds", jax.nn.initializers.ones, (3, self.config.projection_dim) - ) - - self.concept_embeds_weights = self.param("concept_embeds_weights", jax.nn.initializers.ones, (17,)) - self.special_care_embeds_weights = self.param("special_care_embeds_weights", jax.nn.initializers.ones, (3,)) - - def __call__(self, clip_input): - pooled_output = self.vision_model(clip_input)[1] - image_embeds = self.visual_projection(pooled_output) - - special_cos_dist = jax_cosine_distance(image_embeds, self.special_care_embeds) - cos_dist = jax_cosine_distance(image_embeds, self.concept_embeds) - - # increase this value to create a stronger `nfsw` filter - # at the cost of increasing the possibility of filtering benign image inputs - adjustment = 0.0 - - special_scores = special_cos_dist - self.special_care_embeds_weights[None, :] + adjustment - special_scores = jnp.round(special_scores, 3) - is_special_care = jnp.any(special_scores > 0, axis=1, keepdims=True) - # Use a lower threshold if an image has any special care concept - special_adjustment = is_special_care * 0.01 - - concept_scores = cos_dist - self.concept_embeds_weights[None, :] + special_adjustment - concept_scores = jnp.round(concept_scores, 3) - has_nsfw_concepts = jnp.any(concept_scores > 0, axis=1) - - return has_nsfw_concepts - - -class FlaxStableDiffusionSafetyChecker(FlaxPreTrainedModel): - config_class = CLIPConfig - main_input_name = "clip_input" - module_class = FlaxStableDiffusionSafetyCheckerModule - - def __init__( - self, - config: CLIPConfig, - input_shape: Optional[Tuple] = None, - seed: int = 0, - dtype: jnp.dtype = jnp.float32, - _do_init: bool = True, - **kwargs, - ): - if input_shape is None: - input_shape = (1, 224, 224, 3) - 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: - # init input tensor - clip_input = jax.random.normal(rng, input_shape) - - params_rng, dropout_rng = jax.random.split(rng) - rngs = {"params": params_rng, "dropout": dropout_rng} - - random_params = self.module.init(rngs, clip_input)["params"] - - return random_params - - def __call__( - self, - clip_input, - params: dict = None, - ): - clip_input = jnp.transpose(clip_input, (0, 2, 3, 1)) - - return self.module.apply( - {"params": params or self.params}, - jnp.array(clip_input, dtype=jnp.float32), - rngs={}, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py deleted file mode 100644 index ffd66792fe463c775027362b1b63fff250f3b52f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py +++ /dev/null @@ -1,57 +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, Union - -import torch -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.modeling_utils import ModelMixin - - -class StableUnCLIPImageNormalizer(ModelMixin, ConfigMixin): - """ - This class is used to hold the mean and standard deviation of the CLIP embedder used in stable unCLIP. - - It is used to normalize the image embeddings before the noise is applied and un-normalize the noised image - embeddings. - """ - - @register_to_config - def __init__( - self, - embedding_dim: int = 768, - ): - super().__init__() - - self.mean = nn.Parameter(torch.zeros(1, embedding_dim)) - self.std = nn.Parameter(torch.ones(1, embedding_dim)) - - def to( - self, - torch_device: Optional[Union[str, torch.device]] = None, - torch_dtype: Optional[torch.dtype] = 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)) - return self - - def scale(self, embeds): - embeds = (embeds - self.mean) * 1.0 / self.std - return embeds - - def unscale(self, embeds): - embeds = (embeds * self.std) + self.mean - return embeds diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__init__.py deleted file mode 100644 index b0604589a208d0c5c1270629dcad144395b750cb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__init__.py +++ /dev/null @@ -1,54 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - get_objects_from_module, - is_flax_available, - is_torch_available, - is_transformers_available, -) - - -_dummy_objects = {} -_additional_imports = {} -_import_structure = {"pipeline_output": ["StableDiffusion3PipelineOutput"]} - -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_stable_diffusion_3"] = ["StableDiffusion3Pipeline"] - _import_structure["pipeline_stable_diffusion_3_img2img"] = ["StableDiffusion3Img2ImgPipeline"] - _import_structure["pipeline_stable_diffusion_3_inpaint"] = ["StableDiffusion3InpaintPipeline"] - -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_stable_diffusion_3 import StableDiffusion3Pipeline - from .pipeline_stable_diffusion_3_img2img import StableDiffusion3Img2ImgPipeline - from .pipeline_stable_diffusion_3_inpaint import StableDiffusion3InpaintPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index a6e1e1d1aec20eabf1b3936c54c802f20b4e743d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index e6dd46d56117831d2deb81585da22cdbb628a214..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_stable_diffusion_3.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_stable_diffusion_3.cpython-310.pyc deleted file mode 100644 index 3c50861590ef21ab107333e928cebde5d2f7ca5a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_stable_diffusion_3.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_stable_diffusion_3_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_stable_diffusion_3_img2img.cpython-310.pyc deleted file mode 100644 index 635ec0213d9252f90e821ae256cdce6b89a4ad67..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_stable_diffusion_3_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_stable_diffusion_3_inpaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_stable_diffusion_3_inpaint.cpython-310.pyc deleted file mode 100644 index 0392b221d6d23e1b34b0d1d7a62cda8e7232db64..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/__pycache__/pipeline_stable_diffusion_3_inpaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/pipeline_output.py deleted file mode 100644 index 4655f446102a7f2db93743abdba9acc038270862..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/pipeline_output.py +++ /dev/null @@ -1,21 +0,0 @@ -from dataclasses import dataclass -from typing import List, Union - -import numpy as np -import PIL.Image - -from ...utils import BaseOutput - - -@dataclass -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, - num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py deleted file mode 100644 index ad4c4b09172ee78b9f5eec5397b5afaa449a7c5b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py +++ /dev/null @@ -1,1139 +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. - -import inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import torch -from transformers import ( - CLIPTextModelWithProjection, - CLIPTokenizer, - SiglipImageProcessor, - SiglipVisionModel, - T5EncoderModel, - T5TokenizerFast, -) - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, SD3IPAdapterMixin, SD3LoraLoaderMixin -from ...models.autoencoders import AutoencoderKL -from ...models.transformers import SD3Transformer2DModel -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 StableDiffusion3PipelineOutput - - -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 StableDiffusion3Pipeline - - >>> pipe = StableDiffusion3Pipeline.from_pretrained( - ... "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16 - ... ) - >>> pipe.to("cuda") - >>> prompt = "A cat holding a sign that says hello world" - >>> image = pipe(prompt).images[0] - >>> image.save("sd3.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 StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, SD3IPAdapterMixin): - r""" - Args: - transformer ([`SD3Transformer2DModel`]): - 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 ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModelWithProjection`]): - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), - specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant, - with an additional added projection layer that is initialized with a diagonal matrix with the `hidden_size` - as its dimension. - text_encoder_2 ([`CLIPTextModelWithProjection`]): - [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. - text_encoder_3 ([`T5EncoderModel`]): - Frozen text-encoder. Stable Diffusion 3 uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the - [t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) 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). - tokenizer_3 (`T5TokenizerFast`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - image_encoder (`SiglipVisionModel`, *optional*): - Pre-trained Vision Model for IP Adapter. - feature_extractor (`SiglipImageProcessor`, *optional*): - Image processor for IP Adapter. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->image_encoder->transformer->vae" - _optional_components = ["image_encoder", "feature_extractor"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "negative_pooled_prompt_embeds"] - - def __init__( - self, - transformer: SD3Transformer2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer_2: CLIPTokenizer, - text_encoder_3: T5EncoderModel, - tokenizer_3: T5TokenizerFast, - image_encoder: SiglipVisionModel = None, - feature_extractor: SiglipImageProcessor = None, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - text_encoder_3=text_encoder_3, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - tokenizer_3=tokenizer_3, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = ( - self.transformer.config.sample_size - if hasattr(self, "transformer") and self.transformer is not None - else 128 - ) - self.patch_size = ( - self.transformer.config.patch_size if hasattr(self, "transformer") and self.transformer is not None else 2 - ) - - def _get_t5_prompt_embeds( - self, - prompt: Union[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 = 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 self.text_encoder_3 is None: - return torch.zeros( - ( - batch_size * num_images_per_prompt, - self.tokenizer_max_length, - self.transformer.config.joint_attention_dim, - ), - device=device, - dtype=dtype, - ) - - text_inputs = self.tokenizer_3( - 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 - untruncated_ids = self.tokenizer_3(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_3.batch_decode(untruncated_ids[:, self.tokenizer_max_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_3(text_input_ids.to(device))[0] - - dtype = self.text_encoder_3.dtype - prompt_embeds = prompt_embeds.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) - - return prompt_embeds - - def _get_clip_prompt_embeds( - self, - prompt: Union[str, List[str]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clip_skip: Optional[int] = None, - clip_model_index: int = 0, - ): - device = device or self._execution_device - - clip_tokenizers = [self.tokenizer, self.tokenizer_2] - clip_text_encoders = [self.text_encoder, self.text_encoder_2] - - tokenizer = clip_tokenizers[clip_model_index] - text_encoder = clip_text_encoders[clip_model_index] - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer_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[:, self.tokenizer_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) - pooled_prompt_embeds = prompt_embeds[0] - - if clip_skip is None: - prompt_embeds = prompt_embeds.hidden_states[-2] - else: - prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)] - - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - _, 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(batch_size * num_images_per_prompt, seq_len, -1) - - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) - pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds, pooled_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, - 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, - max_sequence_length: int = 256, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - 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`): - 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 all the text-encoders. - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. - 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. - 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, SD3LoraLoaderMixin): - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - 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] - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - prompt_3 = prompt_3 or prompt - prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3 - - prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=clip_skip, - clip_model_index=0, - ) - prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds( - prompt=prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=clip_skip, - clip_model_index=1, - ) - clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1) - - t5_prompt_embed = self._get_t5_prompt_embeds( - prompt=prompt_3, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - clip_prompt_embeds = torch.nn.functional.pad( - clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1]) - ) - - prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2) - pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1) - - if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt_2 = negative_prompt_2 or negative_prompt - negative_prompt_3 = negative_prompt_3 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 - ) - negative_prompt_3 = ( - batch_size * [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3 - ) - - 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_embed, negative_pooled_prompt_embed = self._get_clip_prompt_embeds( - negative_prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=None, - clip_model_index=0, - ) - negative_prompt_2_embed, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds( - negative_prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=None, - clip_model_index=1, - ) - negative_clip_prompt_embeds = torch.cat([negative_prompt_embed, negative_prompt_2_embed], dim=-1) - - t5_negative_prompt_embed = self._get_t5_prompt_embeds( - prompt=negative_prompt_3, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - negative_clip_prompt_embeds = torch.nn.functional.pad( - negative_clip_prompt_embeds, - (0, t5_negative_prompt_embed.shape[-1] - negative_clip_prompt_embeds.shape[-1]), - ) - - negative_prompt_embeds = torch.cat([negative_clip_prompt_embeds, t5_negative_prompt_embed], dim=-2) - negative_pooled_prompt_embeds = torch.cat( - [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1 - ) - - if self.text_encoder is not None: - if isinstance(self, SD3LoraLoaderMixin) 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, SD3LoraLoaderMixin) 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, - prompt_3, - height, - width, - negative_prompt=None, - negative_prompt_2=None, - negative_prompt_3=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - ): - if ( - height % (self.vae_scale_factor * self.patch_size) != 0 - or width % (self.vae_scale_factor * self.patch_size) != 0 - ): - raise ValueError( - f"`height` and `width` have to be divisible by {self.vae_scale_factor * self.patch_size} but are {height} and {width}." - f"You can use height {height - height % (self.vae_scale_factor * self.patch_size)} and width {width - width % (self.vae_scale_factor * self.patch_size)}." - ) - - 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_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_3 is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt_3`: {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)}") - elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)): - raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}") - - 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." - ) - elif negative_prompt_3 is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} 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`." - ) - - 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}") - - 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=device, dtype=dtype) - - 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 - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def skip_guidance_layers(self): - return self._skip_guidance_layers - - @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 - - @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 - - # Adapted from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.encode_image - def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch.Tensor: - """Encodes the given image into a feature representation using a pre-trained image encoder. - - Args: - image (`PipelineImageInput`): - Input image to be encoded. - device: (`torch.device`): - Torch device. - - Returns: - `torch.Tensor`: The encoded image feature representation. - """ - if not isinstance(image, torch.Tensor): - image = self.feature_extractor(image, return_tensors="pt").pixel_values - - image = image.to(device=device, dtype=self.dtype) - - return self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - - # 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, - num_images_per_prompt: int = 1, - do_classifier_free_guidance: bool = True, - ) -> torch.Tensor: - """Prepares image embeddings for use in the IP-Adapter. - - Either `ip_adapter_image` or `ip_adapter_image_embeds` must be passed. - - Args: - ip_adapter_image (`PipelineImageInput`, *optional*): - The input image to extract features from for IP-Adapter. - ip_adapter_image_embeds (`torch.Tensor`, *optional*): - Precomputed image embeddings. - device: (`torch.device`, *optional*): - Torch device. - num_images_per_prompt (`int`, defaults to 1): - Number of images that should be generated per prompt. - do_classifier_free_guidance (`bool`, defaults to True): - Whether to use classifier free guidance or not. - """ - device = device or self._execution_device - - if ip_adapter_image_embeds is not None: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = ip_adapter_image_embeds.chunk(2) - else: - single_image_embeds = ip_adapter_image_embeds - elif ip_adapter_image is not None: - single_image_embeds = self.encode_image(ip_adapter_image, device) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.zeros_like(single_image_embeds) - else: - raise ValueError("Neither `ip_adapter_image_embeds` or `ip_adapter_image_embeds` were provided.") - - image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - negative_image_embeds = torch.cat([single_negative_image_embeds] * num_images_per_prompt, dim=0) - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - - return image_embeds.to(device=device) - - def enable_sequential_cpu_offload(self, *args, **kwargs): - if self.image_encoder is not None and "image_encoder" not in self._exclude_from_cpu_offload: - logger.warning( - "`pipe.enable_sequential_cpu_offload()` might fail for `image_encoder` if it uses " - "`torch.nn.MultiheadAttention`. You can exclude `image_encoder` from CPU offloading by calling " - "`pipe._exclude_from_cpu_offload.append('image_encoder')` before `pipe.enable_sequential_cpu_offload()`." - ) - - super().enable_sequential_cpu_offload(*args, **kwargs) - - @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, - prompt_3: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 28, - sigmas: Optional[List[float]] = 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", - 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"], - max_sequence_length: int = 256, - 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, - ): - 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 `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - will be used instead - 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): - 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. - guidance_scale (`float`, *optional*, defaults to 7.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 instead - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. - ip_adapter_image (`PipelineImageInput`, *optional*): - Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`torch.Tensor`, *optional*): - Pre-generated image embeddings for IP-Adapter. 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. - 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_3.StableDiffusion3PipelineOutput`] 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 256): Maximum sequence length to use with the `prompt`. - 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]. - skip_layer_guidance_scale (`int`, *optional*): The scale of the guidance for the layers specified in - `skip_guidance_layers`. The guidance will be applied to the layers specified in `skip_guidance_layers` - with a scale of `skip_layer_guidance_scale`. The guidance will be applied to the rest of the layers - with a scale of `1`. - skip_layer_guidance_stop (`int`, *optional*): The step at which the guidance for the layers specified in - `skip_guidance_layers` will stop. The guidance will be applied to the layers specified in - `skip_guidance_layers` until the fraction specified in `skip_layer_guidance_stop`. Recommended value by - StabiltyAI for Stable Diffusion 3.5 Medium is 0.2. - skip_layer_guidance_start (`int`, *optional*): The step at which the guidance for the layers specified in - `skip_guidance_layers` will start. The guidance will be applied to the layers specified in - `skip_guidance_layers` from the fraction specified in `skip_layer_guidance_start`. Recommended value by - StabiltyAI for Stable Diffusion 3.5 Medium is 0.01. - mu (`float`, *optional*): `mu` value used for `dynamic_shifting`. - - Examples: - - Returns: - [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] 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_2, - prompt_3, - height, - width, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - ) - - self._guidance_scale = guidance_scale - self._skip_layer_guidance_scale = skip_layer_guidance_scale - self._clip_skip = clip_skip - self._joint_attention_kwargs = joint_attention_kwargs - 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, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_3=prompt_3, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - do_classifier_free_guidance=self.do_classifier_free_guidance, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - device=device, - clip_skip=self.clip_skip, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - - if self.do_classifier_free_guidance: - if skip_guidance_layers is not None: - original_prompt_embeds = prompt_embeds - original_pooled_prompt_embeds = pooled_prompt_embeds - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) - - # 4. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 5. Prepare timesteps - scheduler_kwargs = {} - if self.scheduler.config.get("use_dynamic_shifting", None) and mu is None: - _, _, height, width = latents.shape - image_seq_len = (height // self.transformer.config.patch_size) * ( - width // self.transformer.config.patch_size - ) - 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.16), - ) - scheduler_kwargs["mu"] = mu - elif mu is not None: - scheduler_kwargs["mu"] = mu - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - device, - sigmas=sigmas, - **scheduler_kwargs, - ) - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - # 6. Prepare image embeddings - if (ip_adapter_image is not None and self.is_ip_adapter_active) or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - - if self.joint_attention_kwargs is None: - self._joint_attention_kwargs = {"ip_adapter_image_embeds": ip_adapter_image_embeds} - else: - self._joint_attention_kwargs.update(ip_adapter_image_embeds=ip_adapter_image_embeds) - - # 7. Denoising loop - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - pooled_projections=pooled_prompt_embeds, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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) - should_skip_layers = ( - True - if i > num_inference_steps * skip_layer_guidance_start - and i < num_inference_steps * skip_layer_guidance_stop - else False - ) - if skip_guidance_layers is not None and should_skip_layers: - timestep = t.expand(latents.shape[0]) - latent_model_input = latents - noise_pred_skip_layers = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=original_prompt_embeds, - pooled_projections=original_pooled_prompt_embeds, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - skip_layers=skip_guidance_layers, - )[0] - noise_pred = ( - noise_pred + (noise_pred_text - noise_pred_skip_layers) * self._skip_layer_guidance_scale - ) - - # 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) - negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_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: - 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 StableDiffusion3PipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py deleted file mode 100644 index fa1e0a4f3270c5c5c84d5c546962ba4030f9cf1c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py +++ /dev/null @@ -1,1154 +0,0 @@ -# Copyright 2025 Stability 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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import PIL.Image -import torch -from transformers import ( - CLIPTextModelWithProjection, - CLIPTokenizer, - SiglipImageProcessor, - SiglipVisionModel, - T5EncoderModel, - T5TokenizerFast, -) - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, SD3IPAdapterMixin, SD3LoraLoaderMixin -from ...models.autoencoders import AutoencoderKL -from ...models.transformers import SD3Transformer2DModel -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 StableDiffusion3PipelineOutput - - -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 AutoPipelineForImage2Image - >>> from diffusers.utils import load_image - - >>> device = "cuda" - >>> model_id_or_path = "stabilityai/stable-diffusion-3-medium-diffusers" - >>> pipe = AutoPipelineForImage2Image.from_pretrained(model_id_or_path, torch_dtype=torch.float16) - >>> pipe = pipe.to(device) - - >>> 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)) - - >>> prompt = "cat wizard, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney, 8k" - - >>> images = pipe(prompt=prompt, image=init_image, strength=0.95, guidance_scale=7.5).images[0] - ``` -""" - - -# 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: Optional[torch.Generator] = 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 StableDiffusion3Img2ImgPipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, SD3IPAdapterMixin): - r""" - Args: - transformer ([`SD3Transformer2DModel`]): - 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 ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModelWithProjection`]): - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), - specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant, - with an additional added projection layer that is initialized with a diagonal matrix with the `hidden_size` - as its dimension. - text_encoder_2 ([`CLIPTextModelWithProjection`]): - [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. - text_encoder_3 ([`T5EncoderModel`]): - Frozen text-encoder. Stable Diffusion 3 uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the - [t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) 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). - tokenizer_3 (`T5TokenizerFast`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - image_encoder (`SiglipVisionModel`, *optional*): - Pre-trained Vision Model for IP Adapter. - feature_extractor (`SiglipImageProcessor`, *optional*): - Image processor for IP Adapter. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->image_encoder->transformer->vae" - _optional_components = ["image_encoder", "feature_extractor"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "negative_pooled_prompt_embeds"] - - def __init__( - self, - transformer: SD3Transformer2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer_2: CLIPTokenizer, - text_encoder_3: T5EncoderModel, - tokenizer_3: T5TokenizerFast, - image_encoder: Optional[SiglipVisionModel] = None, - feature_extractor: Optional[SiglipImageProcessor] = None, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - text_encoder_3=text_encoder_3, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - tokenizer_3=tokenizer_3, - 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 - 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, vae_latent_channels=latent_channels - ) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = ( - self.transformer.config.sample_size - if hasattr(self, "transformer") and self.transformer is not None - else 128 - ) - self.patch_size = ( - self.transformer.config.patch_size if hasattr(self, "transformer") and self.transformer is not None else 2 - ) - - # 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, - num_images_per_prompt: int = 1, - max_sequence_length: int = 256, - 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 self.text_encoder_3 is None: - return torch.zeros( - ( - batch_size * num_images_per_prompt, - self.tokenizer_max_length, - self.transformer.config.joint_attention_dim, - ), - device=device, - dtype=dtype, - ) - - text_inputs = self.tokenizer_3( - 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 - untruncated_ids = self.tokenizer_3(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_3.batch_decode(untruncated_ids[:, self.tokenizer_max_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_3(text_input_ids.to(device))[0] - - dtype = self.text_encoder_3.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clip_skip: Optional[int] = None, - clip_model_index: int = 0, - ): - device = device or self._execution_device - - clip_tokenizers = [self.tokenizer, self.tokenizer_2] - clip_text_encoders = [self.text_encoder, self.text_encoder_2] - - tokenizer = clip_tokenizers[clip_model_index] - text_encoder = clip_text_encoders[clip_model_index] - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer_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[:, self.tokenizer_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) - pooled_prompt_embeds = prompt_embeds[0] - - if clip_skip is None: - prompt_embeds = prompt_embeds.hidden_states[-2] - else: - prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)] - - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - _, 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(batch_size * num_images_per_prompt, seq_len, -1) - - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) - pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds, pooled_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, - 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, - max_sequence_length: int = 256, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - 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`): - 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 all the text-encoders. - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. - 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. - 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, SD3LoraLoaderMixin): - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - 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] - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - prompt_3 = prompt_3 or prompt - prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3 - - prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=clip_skip, - clip_model_index=0, - ) - prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds( - prompt=prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=clip_skip, - clip_model_index=1, - ) - clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1) - - t5_prompt_embed = self._get_t5_prompt_embeds( - prompt=prompt_3, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - clip_prompt_embeds = torch.nn.functional.pad( - clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1]) - ) - - prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2) - pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1) - - if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt_2 = negative_prompt_2 or negative_prompt - negative_prompt_3 = negative_prompt_3 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 - ) - negative_prompt_3 = ( - batch_size * [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3 - ) - - 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_embed, negative_pooled_prompt_embed = self._get_clip_prompt_embeds( - negative_prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=None, - clip_model_index=0, - ) - negative_prompt_2_embed, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds( - negative_prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=None, - clip_model_index=1, - ) - negative_clip_prompt_embeds = torch.cat([negative_prompt_embed, negative_prompt_2_embed], dim=-1) - - t5_negative_prompt_embed = self._get_t5_prompt_embeds( - prompt=negative_prompt_3, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - negative_clip_prompt_embeds = torch.nn.functional.pad( - negative_clip_prompt_embeds, - (0, t5_negative_prompt_embed.shape[-1] - negative_clip_prompt_embeds.shape[-1]), - ) - - negative_prompt_embeds = torch.cat([negative_clip_prompt_embeds, t5_negative_prompt_embed], dim=-2) - negative_pooled_prompt_embeds = torch.cat( - [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1 - ) - - if self.text_encoder is not None: - if isinstance(self, SD3LoraLoaderMixin) 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, SD3LoraLoaderMixin) 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, - prompt_3, - height, - width, - strength, - negative_prompt=None, - negative_prompt_2=None, - negative_prompt_3=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - ): - if ( - height % (self.vae_scale_factor * self.patch_size) != 0 - or width % (self.vae_scale_factor * self.patch_size) != 0 - ): - raise ValueError( - f"`height` and `width` have to be divisible by {self.vae_scale_factor * self.patch_size} but are {height} and {width}." - f"You can use height {height - height % (self.vae_scale_factor * self.patch_size)} and width {width - width % (self.vae_scale_factor * self.patch_size)}." - ) - - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should 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_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_3 is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt_3`: {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)}") - elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)): - raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}") - - 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." - ) - elif negative_prompt_3 is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} 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`." - ) - - 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}") - - 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_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - if image.shape[1] == self.vae.config.latent_channels: - init_latents = image - - else: - 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." - ) - - elif isinstance(generator, list): - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - init_latents = (init_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # get latents - init_latents = self.scheduler.scale_noise(init_latents, timestep, noise) - latents = init_latents.to(device=device, dtype=dtype) - - return latents - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def joint_attention_kwargs(self): - return self._joint_attention_kwargs - - @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 - - @property - def num_timesteps(self): - return self._num_timesteps - - @property - def interrupt(self): - return self._interrupt - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_image - def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch.Tensor: - """Encodes the given image into a feature representation using a pre-trained image encoder. - - Args: - image (`PipelineImageInput`): - Input image to be encoded. - device: (`torch.device`): - Torch device. - - Returns: - `torch.Tensor`: The encoded image feature representation. - """ - if not isinstance(image, torch.Tensor): - image = self.feature_extractor(image, return_tensors="pt").pixel_values - - image = image.to(device=device, dtype=self.dtype) - - return self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - - # 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, - num_images_per_prompt: int = 1, - do_classifier_free_guidance: bool = True, - ) -> torch.Tensor: - """Prepares image embeddings for use in the IP-Adapter. - - Either `ip_adapter_image` or `ip_adapter_image_embeds` must be passed. - - Args: - ip_adapter_image (`PipelineImageInput`, *optional*): - The input image to extract features from for IP-Adapter. - ip_adapter_image_embeds (`torch.Tensor`, *optional*): - Precomputed image embeddings. - device: (`torch.device`, *optional*): - Torch device. - num_images_per_prompt (`int`, defaults to 1): - Number of images that should be generated per prompt. - do_classifier_free_guidance (`bool`, defaults to True): - Whether to use classifier free guidance or not. - """ - device = device or self._execution_device - - if ip_adapter_image_embeds is not None: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = ip_adapter_image_embeds.chunk(2) - else: - single_image_embeds = ip_adapter_image_embeds - elif ip_adapter_image is not None: - single_image_embeds = self.encode_image(ip_adapter_image, device) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.zeros_like(single_image_embeds) - else: - raise ValueError("Neither `ip_adapter_image_embeds` or `ip_adapter_image_embeds` were provided.") - - image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - negative_image_embeds = torch.cat([single_negative_image_embeds] * num_images_per_prompt, dim=0) - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - - return image_embeds.to(device=device) - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.enable_sequential_cpu_offload - def enable_sequential_cpu_offload(self, *args, **kwargs): - if self.image_encoder is not None and "image_encoder" not in self._exclude_from_cpu_offload: - logger.warning( - "`pipe.enable_sequential_cpu_offload()` might fail for `image_encoder` if it uses " - "`torch.nn.MultiheadAttention`. You can exclude `image_encoder` from CPU offloading by calling " - "`pipe._exclude_from_cpu_offload.append('image_encoder')` before `pipe.enable_sequential_cpu_offload()`." - ) - - super().enable_sequential_cpu_offload(*args, **kwargs) - - @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, - prompt_3: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - image: PipelineImageInput = None, - strength: float = 0.6, - num_inference_steps: int = 50, - sigmas: Optional[List[float]] = 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, - 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"], - max_sequence_length: int = 256, - mu: Optional[float] = 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 `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - will be used instead - 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): - 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.transformer.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. - guidance_scale (`float`, *optional*, defaults to 7.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 instead - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. - ip_adapter_image (`PipelineImageInput`, *optional*): - Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`torch.Tensor`, *optional*): - Pre-generated image embeddings for IP-Adapter. 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. - 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_3.StableDiffusion3PipelineOutput`] 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 256): Maximum sequence length to use with the `prompt`. - mu (`float`, *optional*): `mu` value used for `dynamic_shifting`. - - Examples: - - Returns: - [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] 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_2, - prompt_3, - height, - width, - strength, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._joint_attention_kwargs = joint_attention_kwargs - 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, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_3=prompt_3, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - do_classifier_free_guidance=self.do_classifier_free_guidance, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - device=device, - clip_skip=self.clip_skip, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) - - # 3. Preprocess image - image = self.image_processor.preprocess(image, height=height, width=width) - - # 4. Prepare timesteps - scheduler_kwargs = {} - if self.scheduler.config.get("use_dynamic_shifting", None) and mu is None: - image_seq_len = (int(height) // self.vae_scale_factor // self.transformer.config.patch_size) * ( - int(width) // self.vae_scale_factor // self.transformer.config.patch_size - ) - 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.16), - ) - scheduler_kwargs["mu"] = mu - elif mu is not None: - scheduler_kwargs["mu"] = mu - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, 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) - - # 5. Prepare latent variables - if latents is None: - latents = self.prepare_latents( - image, - latent_timestep, - batch_size, - num_images_per_prompt, - prompt_embeds.dtype, - device, - generator, - ) - - # 6. Prepare image embeddings - if (ip_adapter_image is not None and self.is_ip_adapter_active) or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - - if self.joint_attention_kwargs is None: - self._joint_attention_kwargs = {"ip_adapter_image_embeds": ip_adapter_image_embeds} - else: - self._joint_attention_kwargs.update(ip_adapter_image_embeds=ip_adapter_image_embeds) - - # 7. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - pooled_projections=pooled_prompt_embeds, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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) - - # 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) - negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_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: - 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 StableDiffusion3PipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py deleted file mode 100644 index 937f7195b21d94b9103f3b130ec6b1951a5bda40..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py +++ /dev/null @@ -1,1379 +0,0 @@ -# Copyright 2025 Stability 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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import torch -from transformers import ( - CLIPTextModelWithProjection, - CLIPTokenizer, - SiglipImageProcessor, - SiglipVisionModel, - T5EncoderModel, - T5TokenizerFast, -) - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, SD3IPAdapterMixin, SD3LoraLoaderMixin -from ...models.autoencoders import AutoencoderKL -from ...models.transformers import SD3Transformer2DModel -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 StableDiffusion3PipelineOutput - - -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 StableDiffusion3InpaintPipeline - >>> from diffusers.utils import load_image - - >>> pipe = StableDiffusion3InpaintPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16 - ... ) - >>> 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("sd3_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: Optional[torch.Generator] = 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 StableDiffusion3InpaintPipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, SD3IPAdapterMixin): - r""" - Args: - transformer ([`SD3Transformer2DModel`]): - 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 ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModelWithProjection`]): - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), - specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant, - with an additional added projection layer that is initialized with a diagonal matrix with the `hidden_size` - as its dimension. - text_encoder_2 ([`CLIPTextModelWithProjection`]): - [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. - text_encoder_3 ([`T5EncoderModel`]): - Frozen text-encoder. Stable Diffusion 3 uses - [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the - [t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) 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). - tokenizer_3 (`T5TokenizerFast`): - Tokenizer of class - [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - image_encoder (`SiglipVisionModel`, *optional*): - Pre-trained Vision Model for IP Adapter. - feature_extractor (`SiglipImageProcessor`, *optional*): - Image processor for IP Adapter. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->image_encoder->transformer->vae" - _optional_components = ["image_encoder", "feature_extractor"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "negative_pooled_prompt_embeds"] - - def __init__( - self, - transformer: SD3Transformer2DModel, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer_2: CLIPTokenizer, - text_encoder_3: T5EncoderModel, - tokenizer_3: T5TokenizerFast, - image_encoder: Optional[SiglipVisionModel] = None, - feature_extractor: Optional[SiglipImageProcessor] = None, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - text_encoder_3=text_encoder_3, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - tokenizer_3=tokenizer_3, - 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 - 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, vae_latent_channels=latent_channels - ) - self.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, - vae_latent_channels=latent_channels, - do_normalize=False, - do_binarize=True, - do_convert_grayscale=True, - ) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = ( - self.transformer.config.sample_size - if hasattr(self, "transformer") and self.transformer is not None - else 128 - ) - self.patch_size = ( - self.transformer.config.patch_size if hasattr(self, "transformer") and self.transformer is not None else 2 - ) - - # 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, - num_images_per_prompt: int = 1, - max_sequence_length: int = 256, - 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 self.text_encoder_3 is None: - return torch.zeros( - ( - batch_size * num_images_per_prompt, - self.tokenizer_max_length, - self.transformer.config.joint_attention_dim, - ), - device=device, - dtype=dtype, - ) - - text_inputs = self.tokenizer_3( - 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 - untruncated_ids = self.tokenizer_3(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_3.batch_decode(untruncated_ids[:, self.tokenizer_max_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_3(text_input_ids.to(device))[0] - - dtype = self.text_encoder_3.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clip_skip: Optional[int] = None, - clip_model_index: int = 0, - ): - device = device or self._execution_device - - clip_tokenizers = [self.tokenizer, self.tokenizer_2] - clip_text_encoders = [self.text_encoder, self.text_encoder_2] - - tokenizer = clip_tokenizers[clip_model_index] - text_encoder = clip_text_encoders[clip_model_index] - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer_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[:, self.tokenizer_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) - pooled_prompt_embeds = prompt_embeds[0] - - if clip_skip is None: - prompt_embeds = prompt_embeds.hidden_states[-2] - else: - prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)] - - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - _, 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(batch_size * num_images_per_prompt, seq_len, -1) - - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1) - pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return prompt_embeds, pooled_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, - 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, - max_sequence_length: int = 256, - lora_scale: Optional[float] = None, - ): - r""" - - 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 all text-encoders - 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`): - 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 all the text-encoders. - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. - 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. - 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, SD3LoraLoaderMixin): - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - 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] - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - prompt_3 = prompt_3 or prompt - prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3 - - prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=clip_skip, - clip_model_index=0, - ) - prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds( - prompt=prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=clip_skip, - clip_model_index=1, - ) - clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1) - - t5_prompt_embed = self._get_t5_prompt_embeds( - prompt=prompt_3, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - clip_prompt_embeds = torch.nn.functional.pad( - clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1]) - ) - - prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2) - pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1) - - if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt_2 = negative_prompt_2 or negative_prompt - negative_prompt_3 = negative_prompt_3 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 - ) - negative_prompt_3 = ( - batch_size * [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3 - ) - - 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_embed, negative_pooled_prompt_embed = self._get_clip_prompt_embeds( - negative_prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=None, - clip_model_index=0, - ) - negative_prompt_2_embed, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds( - negative_prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - clip_skip=None, - clip_model_index=1, - ) - negative_clip_prompt_embeds = torch.cat([negative_prompt_embed, negative_prompt_2_embed], dim=-1) - - t5_negative_prompt_embed = self._get_t5_prompt_embeds( - prompt=negative_prompt_3, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - negative_clip_prompt_embeds = torch.nn.functional.pad( - negative_clip_prompt_embeds, - (0, t5_negative_prompt_embed.shape[-1] - negative_clip_prompt_embeds.shape[-1]), - ) - - negative_prompt_embeds = torch.cat([negative_clip_prompt_embeds, t5_negative_prompt_embed], dim=-2) - negative_pooled_prompt_embeds = torch.cat( - [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1 - ) - - if self.text_encoder is not None: - if isinstance(self, SD3LoraLoaderMixin) 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, SD3LoraLoaderMixin) 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 - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_img2img.StableDiffusion3Img2ImgPipeline.check_inputs - def check_inputs( - self, - prompt, - prompt_2, - prompt_3, - height, - width, - strength, - negative_prompt=None, - negative_prompt_2=None, - negative_prompt_3=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - ): - if ( - height % (self.vae_scale_factor * self.patch_size) != 0 - or width % (self.vae_scale_factor * self.patch_size) != 0 - ): - raise ValueError( - f"`height` and `width` have to be divisible by {self.vae_scale_factor * self.patch_size} but are {height} and {width}." - f"You can use height {height - height % (self.vae_scale_factor * self.patch_size)} and width {width - width % (self.vae_scale_factor * self.patch_size)}." - ) - - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should 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_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_3 is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt_3`: {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)}") - elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)): - raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}") - - 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." - ) - elif negative_prompt_3 is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} 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`." - ) - - 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}") - - # 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_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=None, - image=None, - timestep=None, - is_strength_max=True, - return_noise=False, - return_image_latents=False, - ): - 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 (image is None or timestep is None) and not is_strength_max: - raise ValueError( - "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise." - "However, either the image or the noise timestep has not been provided." - ) - - if return_image_latents or (latents is None and not is_strength_max): - image = image.to(device=device, dtype=dtype) - - if image.shape[1] == 16: - image_latents = image - else: - image_latents = self._encode_vae_image(image=image, generator=generator) - image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1) - - if latents is None: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # if strength is 1. then initialise the latents to noise, else initial to image + noise - latents = noise if is_strength_max else self.scheduler.scale_noise(image_latents, timestep, noise) - else: - noise = latents.to(device) - latents = noise - - outputs = (latents,) - - if return_noise: - outputs += (noise,) - - if return_image_latents: - outputs += (image_latents,) - - return outputs - - 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 - - def prepare_mask_latents( - self, - mask, - masked_image, - batch_size, - num_images_per_prompt, - height, - width, - dtype, - device, - generator, - do_classifier_free_guidance, - ): - # 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 // self.vae_scale_factor, width // self.vae_scale_factor) - ) - 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) - - mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask - masked_image_latents = ( - torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents - ) - - # 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 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 - - @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 - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_image - def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch.Tensor: - """Encodes the given image into a feature representation using a pre-trained image encoder. - - Args: - image (`PipelineImageInput`): - Input image to be encoded. - device: (`torch.device`): - Torch device. - - Returns: - `torch.Tensor`: The encoded image feature representation. - """ - if not isinstance(image, torch.Tensor): - image = self.feature_extractor(image, return_tensors="pt").pixel_values - - image = image.to(device=device, dtype=self.dtype) - - return self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - - # 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, - num_images_per_prompt: int = 1, - do_classifier_free_guidance: bool = True, - ) -> torch.Tensor: - """Prepares image embeddings for use in the IP-Adapter. - - Either `ip_adapter_image` or `ip_adapter_image_embeds` must be passed. - - Args: - ip_adapter_image (`PipelineImageInput`, *optional*): - The input image to extract features from for IP-Adapter. - ip_adapter_image_embeds (`torch.Tensor`, *optional*): - Precomputed image embeddings. - device: (`torch.device`, *optional*): - Torch device. - num_images_per_prompt (`int`, defaults to 1): - Number of images that should be generated per prompt. - do_classifier_free_guidance (`bool`, defaults to True): - Whether to use classifier free guidance or not. - """ - device = device or self._execution_device - - if ip_adapter_image_embeds is not None: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = ip_adapter_image_embeds.chunk(2) - else: - single_image_embeds = ip_adapter_image_embeds - elif ip_adapter_image is not None: - single_image_embeds = self.encode_image(ip_adapter_image, device) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.zeros_like(single_image_embeds) - else: - raise ValueError("Neither `ip_adapter_image_embeds` or `ip_adapter_image_embeds` were provided.") - - image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - negative_image_embeds = torch.cat([single_negative_image_embeds] * num_images_per_prompt, dim=0) - image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0) - - return image_embeds.to(device=device) - - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.enable_sequential_cpu_offload - def enable_sequential_cpu_offload(self, *args, **kwargs): - if self.image_encoder is not None and "image_encoder" not in self._exclude_from_cpu_offload: - logger.warning( - "`pipe.enable_sequential_cpu_offload()` might fail for `image_encoder` if it uses " - "`torch.nn.MultiheadAttention`. You can exclude `image_encoder` from CPU offloading by calling " - "`pipe._exclude_from_cpu_offload.append('image_encoder')` before `pipe.enable_sequential_cpu_offload()`." - ) - - super().enable_sequential_cpu_offload(*args, **kwargs) - - @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, - prompt_3: Optional[Union[str, List[str]]] = 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, - strength: float = 0.6, - num_inference_steps: int = 50, - sigmas: Optional[List[float]] = 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", - 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"], - max_sequence_length: int = 256, - mu: Optional[float] = 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 `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - will be used instead - 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`, 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]`): - `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]`): - `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.transformer.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.transformer.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. - padding_mask_crop (`int`, *optional*, defaults to `None`): - The size of margin in the crop to be applied to the image and masking. If `None`, no crop is applied to - image and mask_image. If `padding_mask_crop` is not `None`, it will first find a rectangular region - with the same aspect ration of the image and contains all masked area, and then expand that area based - on `padding_mask_crop`. The image and mask_image will then be cropped based on the expanded area before - resizing to the original image size for inpainting. This is useful when the masked area is small while - the image is large and contain information irrelevant for inpainting, such as background. - 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`. - 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 7.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 instead - 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*): - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. - ip_adapter_image (`PipelineImageInput`, *optional*): - Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`torch.Tensor`, *optional*): - Pre-generated image embeddings for IP-Adapter. 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. - 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_3.StableDiffusion3PipelineOutput`] 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 256): Maximum sequence length to use with the `prompt`. - mu (`float`, *optional*): `mu` value used for `dynamic_shifting`. - - Examples: - - Returns: - [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, 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 - - height = height or self.transformer.config.sample_size * self.vae_scale_factor - width = width or self.transformer.config.sample_size * self.vae_scale_factor - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - prompt_3, - height, - width, - strength, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - ) - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - self._joint_attention_kwargs = joint_attention_kwargs - 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 - - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_3=prompt_3, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - do_classifier_free_guidance=self.do_classifier_free_guidance, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - device=device, - clip_skip=self.clip_skip, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - ) - - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) - - # 3. Prepare timesteps - scheduler_kwargs = {} - if self.scheduler.config.get("use_dynamic_shifting", None) and mu is None: - image_seq_len = (int(height) // self.vae_scale_factor // self.transformer.config.patch_size) * ( - int(width) // self.vae_scale_factor // self.transformer.config.patch_size - ) - 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.16), - ) - scheduler_kwargs["mu"] = mu - elif mu is not None: - scheduler_kwargs["mu"] = mu - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, 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 - 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) - - # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise - is_strength_max = strength == 1.0 - - # 4. 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) - - # 5. Prepare latent variables - num_channels_latents = self.vae.config.latent_channels - num_channels_transformer = self.transformer.config.in_channels - return_image_latents = num_channels_transformer == 16 - - latents_outputs = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - image=init_image, - timestep=latent_timestep, - is_strength_max=is_strength_max, - return_noise=True, - return_image_latents=return_image_latents, - ) - - if return_image_latents: - latents, noise, image_latents = latents_outputs - else: - latents, noise = latents_outputs - - # 6. Prepare mask latent variables - 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_images_per_prompt, - height, - width, - prompt_embeds.dtype, - device, - generator, - self.do_classifier_free_guidance, - ) - - # match the inpainting pipeline and will be updated with input + mask inpainting model later - if num_channels_transformer == 33: - # default case for runwayml/stable-diffusion-inpainting - num_channels_mask = mask.shape[1] - num_channels_masked_image = masked_image_latents.shape[1] - if ( - num_channels_latents + num_channels_mask + num_channels_masked_image - != self.transformer.config.in_channels - ): - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.transformer`: {self.transformer.config} expects" - f" {self.transformer.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" - " `pipeline.transformer` or your `mask_image` or `image` input." - ) - elif num_channels_transformer != 16: - raise ValueError( - f"The transformer {self.transformer.__class__} should have 16 input channels or 33 input channels, not {self.transformer.config.in_channels}." - ) - - # 7. Prepare image embeddings - if (ip_adapter_image is not None and self.is_ip_adapter_active) or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - self.do_classifier_free_guidance, - ) - - if self.joint_attention_kwargs is None: - self._joint_attention_kwargs = {"ip_adapter_image_embeds": ip_adapter_image_embeds} - else: - self._joint_attention_kwargs.update(ip_adapter_image_embeds=ip_adapter_image_embeds) - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - if num_channels_transformer == 33: - latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - pooled_projections=pooled_prompt_embeds, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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) - - # 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 num_channels_transformer == 16: - init_latents_proper = image_latents - if self.do_classifier_free_guidance: - init_mask, _ = mask.chunk(2) - else: - 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) - negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", negative_pooled_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 not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 0 - ] - else: - image = latents - - do_denormalize = [True] * image.shape[0] - - image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) - - if padding_mask_crop is not None: - image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return StableDiffusion3PipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_attend_and_excite/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_attend_and_excite/__init__.py deleted file mode 100644 index cce556fceb2379be482d383e380a09836d25ce3b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_attend_and_excite/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -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_stable_diffusion_attend_and_excite"] = ["StableDiffusionAttendAndExcitePipeline"] - -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_stable_diffusion_attend_and_excite import StableDiffusionAttendAndExcitePipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_attend_and_excite/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_attend_and_excite/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 1a4f8e8d8780411f0e6f551142032c2ea0b85223..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_attend_and_excite/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_attend_and_excite/__pycache__/pipeline_stable_diffusion_attend_and_excite.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_attend_and_excite/__pycache__/pipeline_stable_diffusion_attend_and_excite.cpython-310.pyc deleted file mode 100644 index 9cc1895cc65e2e90fe31897e208d25fee6fb7974..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_attend_and_excite/__pycache__/pipeline_stable_diffusion_attend_and_excite.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py deleted file mode 100644 index a1ff99b6aa3492ec44dacc75b8fc52705635731a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py +++ /dev/null @@ -1,1113 +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 inspect -import math -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import torch -from torch.nn import functional as F -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer - -from ...image_processor import VaeImageProcessor -from ...loaders import StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, UNet2DConditionModel -from ...models.attention_processor import Attention -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - 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 DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -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__) - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import StableDiffusionAttendAndExcitePipeline - - >>> pipe = StableDiffusionAttendAndExcitePipeline.from_pretrained( - ... "CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16 - ... ).to("cuda") - - - >>> prompt = "a cat and a frog" - - >>> # use get_indices function to find out indices of the tokens you want to alter - >>> pipe.get_indices(prompt) - {0: '<|startoftext|>', 1: 'a', 2: 'cat', 3: 'and', 4: 'a', 5: 'frog', 6: '<|endoftext|>'} - - >>> token_indices = [2, 5] - >>> seed = 6141 - >>> generator = torch.Generator("cuda").manual_seed(seed) - - >>> images = pipe( - ... prompt=prompt, - ... token_indices=token_indices, - ... guidance_scale=7.5, - ... generator=generator, - ... num_inference_steps=50, - ... max_iter_to_alter=25, - ... ).images - - >>> image = images[0] - >>> image.save(f"../images/{prompt}_{seed}.png") - ``` -""" - - -class AttentionStore: - @staticmethod - def get_empty_store(): - return {"down": [], "mid": [], "up": []} - - def __call__(self, attn, is_cross: bool, place_in_unet: str): - if self.cur_att_layer >= 0 and is_cross: - if attn.shape[1] == np.prod(self.attn_res): - self.step_store[place_in_unet].append(attn) - - self.cur_att_layer += 1 - if self.cur_att_layer == self.num_att_layers: - self.cur_att_layer = 0 - self.between_steps() - - def between_steps(self): - self.attention_store = self.step_store - self.step_store = self.get_empty_store() - - def get_average_attention(self): - average_attention = self.attention_store - return average_attention - - 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() - for location in from_where: - for item in attention_maps[location]: - cross_maps = item.reshape(-1, self.attn_res[0], self.attn_res[1], item.shape[-1]) - out.append(cross_maps) - out = torch.cat(out, dim=0) - out = out.sum(0) / out.shape[0] - return out - - def reset(self): - self.cur_att_layer = 0 - self.step_store = self.get_empty_store() - self.attention_store = {} - - def __init__(self, attn_res): - """ - Initialize an empty AttentionStore :param step_index: used to visualize only a specific step in the diffusion - process - """ - self.num_att_layers = -1 - self.cur_att_layer = 0 - self.step_store = self.get_empty_store() - self.attention_store = {} - self.curr_step_index = 0 - self.attn_res = attn_res - - -class AttendExciteAttnProcessor: - def __init__(self, attnstore, place_in_unet): - super().__init__() - self.attnstore = attnstore - self.place_in_unet = place_in_unet - - def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None): - batch_size, sequence_length, _ = hidden_states.shape - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - - query = attn.to_q(hidden_states) - - is_cross = encoder_hidden_states is not None - encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query = attn.head_to_batch_dim(query) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - - attention_probs = attn.get_attention_scores(query, key, attention_mask) - - # only need to store attention maps during the Attend and Excite process - if attention_probs.requires_grad: - self.attnstore(attention_probs, is_cross, self.place_in_unet) - - hidden_states = torch.bmm(attention_probs, value) - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - return hidden_states - - -class StableDiffusionAttendAndExcitePipeline( - DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin, TextualInversionLoaderMixin -): - r""" - Pipeline for text-to-image generation using Stable Diffusion and Attend-and-Excite. - - 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.). - - The pipeline also inherits the following loading methods: - - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to 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: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - indices, - height, - width, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_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 (callback_steps is None) or ( - 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 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}." - ) - - indices_is_list_ints = isinstance(indices, list) and isinstance(indices[0], int) - indices_is_list_list_ints = ( - isinstance(indices, list) and isinstance(indices[0], list) and isinstance(indices[0][0], int) - ) - - if not indices_is_list_ints and not indices_is_list_list_ints: - raise TypeError("`indices` must be a list of ints or a list of a list of ints") - - if indices_is_list_ints: - indices_batch_size = 1 - elif indices_is_list_list_ints: - indices_batch_size = len(indices) - - 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 indices_batch_size != prompt_batch_size: - raise ValueError( - f"indices batch size must be same as prompt batch size. indices batch size: {indices_batch_size}, prompt batch size: {prompt_batch_size}" - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @staticmethod - def _compute_max_attention_per_index( - attention_maps: 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 - attention_for_text = torch.nn.functional.softmax(attention_for_text, dim=-1) - - # Shift indices since we removed the first token - indices = [index - 1 for index in indices] - - # Extract the maximum values - max_indices_list = [] - for i in indices: - image = attention_for_text[:, :, i] - smoothing = GaussianSmoothing().to(attention_maps.device) - input = F.pad(image.unsqueeze(0).unsqueeze(0), (1, 1, 1, 1), mode="reflect") - image = smoothing(input).squeeze(0).squeeze(0) - max_indices_list.append(image.max()) - return max_indices_list - - def _aggregate_and_get_max_attention_per_token( - self, - 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( - from_where=("up", "down", "mid"), - ) - max_attention_per_index = self._compute_max_attention_per_index( - attention_maps=attention_maps, - indices=indices, - ) - return max_attention_per_index - - @staticmethod - 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) - return loss - - @staticmethod - def _update_latent(latents: torch.Tensor, loss: torch.Tensor, step_size: float) -> torch.Tensor: - """Update the latent according to the computed loss.""" - grad_cond = torch.autograd.grad(loss.requires_grad_(True), [latents], retain_graph=True)[0] - latents = latents - step_size * grad_cond - return latents - - def _perform_iterative_refinement_step( - self, - latents: torch.Tensor, - indices: List[int], - loss: torch.Tensor, - threshold: float, - text_embeddings: torch.Tensor, - step_size: float, - t: int, - max_refinement_steps: int = 20, - ): - """ - Performs the iterative latent refinement introduced in the paper. Here, we continuously update the latent code - according to our loss objective until the given threshold is reached for all tokens. - """ - iteration = 0 - target_loss = max(0, 1.0 - threshold) - while loss > target_loss: - iteration += 1 - - latents = latents.clone().detach().requires_grad_(True) - self.unet(latents, t, encoder_hidden_states=text_embeddings).sample - self.unet.zero_grad() - - # Get max activation value for each subject token - max_attention_per_index = self._aggregate_and_get_max_attention_per_token( - indices=indices, - ) - - loss = self._compute_loss(max_attention_per_index) - - if loss != 0: - latents = self._update_latent(latents, loss, step_size) - - logger.info(f"\t Try {iteration}. loss: {loss}") - - if iteration >= max_refinement_steps: - logger.info(f"\t Exceeded max number of iterations ({max_refinement_steps})! ") - break - - # Run one more time but don't compute gradients and update the latents. - # We just need to compute the new loss - the grad update will occur below - latents = latents.clone().detach().requires_grad_(True) - _ = self.unet(latents, t, encoder_hidden_states=text_embeddings).sample - self.unet.zero_grad() - - # Get max activation value for each subject token - max_attention_per_index = self._aggregate_and_get_max_attention_per_token( - indices=indices, - ) - loss = self._compute_loss(max_attention_per_index) - logger.info(f"\t Finished with loss of: {loss}") - return loss, latents, max_attention_per_index - - def register_attention_control(self): - attn_procs = {} - cross_att_count = 0 - for name in self.unet.attn_processors.keys(): - if name.startswith("mid_block"): - place_in_unet = "mid" - elif name.startswith("up_blocks"): - place_in_unet = "up" - elif name.startswith("down_blocks"): - place_in_unet = "down" - else: - continue - - cross_att_count += 1 - attn_procs[name] = AttendExciteAttnProcessor(attnstore=self.attention_store, place_in_unet=place_in_unet) - - 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]: - """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)))} - return indices - - @torch.no_grad() - @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, - num_inference_steps: int = 50, - guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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]`): - 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. - 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): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - max_iter_to_alter (`int`, *optional*, defaults to `25`): - Number of denoising steps to apply attend-and-excite. The `max_iter_to_alter` denoising steps are when - attend-and-excite is applied. For example, if `max_iter_to_alter` is `25` and there are a total of `30` - denoising steps, the first `25` denoising steps applies attend-and-excite and the last `5` will not. - thresholds (`dict`, *optional*, defaults to `{0: 0.05, 10: 0.5, 20: 0.8}`): - Dictionary defining the iterations and desired thresholds to apply iterative latent refinement in. - scale_factor (`int`, *optional*, default to 20): - Scale factor to control the step size of each attend-and-excite update. - attn_res (`tuple`, *optional*, default computed from width and height): - The 2D resolution of the semantic attention map. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - - # 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, - token_indices, - 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 = guidance_scale > 1.0 - - # 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=device) - timesteps = self.scheduler.timesteps - - # 5. 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, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - if attn_res is None: - attn_res = int(np.ceil(width / 32)), int(np.ceil(height / 32)) - self.attention_store = AttentionStore(attn_res) - original_attn_proc = self.unet.attn_processors - self.register_attention_control() - - # default config for step size from original repo - scale_range = np.linspace(1.0, 0.5, len(self.scheduler.timesteps)) - step_size = scale_factor * np.sqrt(scale_range) - - text_embeddings = ( - prompt_embeds[batch_size * num_images_per_prompt :] if do_classifier_free_guidance else prompt_embeds - ) - - if isinstance(token_indices[0], int): - token_indices = [token_indices] - - indices = [] - - for ind in token_indices: - indices = indices + [ind] * num_images_per_prompt - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # Attend and excite process - with torch.enable_grad(): - latents = latents.clone().detach().requires_grad_(True) - updated_latents = [] - for latent, index, text_embedding in zip(latents, indices, text_embeddings): - # Forward pass of denoising with text conditioning - latent = latent.unsqueeze(0) - text_embedding = text_embedding.unsqueeze(0) - - self.unet( - latent, - t, - encoder_hidden_states=text_embedding, - cross_attention_kwargs=cross_attention_kwargs, - ).sample - self.unet.zero_grad() - - # Get max activation value for each subject token - max_attention_per_index = self._aggregate_and_get_max_attention_per_token( - indices=index, - ) - - loss = self._compute_loss(max_attention_per_index=max_attention_per_index) - - # If this is an iterative refinement step, verify we have reached the desired threshold for all - if i in thresholds.keys() and loss > 1.0 - thresholds[i]: - loss, latent, max_attention_per_index = self._perform_iterative_refinement_step( - latents=latent, - indices=index, - loss=loss, - threshold=thresholds[i], - text_embeddings=text_embedding, - step_size=step_size[i], - t=t, - ) - - # Perform gradient update - if i < max_iter_to_alter: - if loss != 0: - latent = self._update_latent( - latents=latent, - loss=loss, - step_size=step_size[i], - ) - logger.info(f"Iteration {i} | Loss: {loss:0.4f}") - - updated_latents.append(latent) - - latents = torch.cat(updated_latents, dim=0) - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - ).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # 8. Post-processing - 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) - self.maybe_free_model_hooks() - # make sure to set the original attention processors back - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) - - -class GaussianSmoothing(torch.nn.Module): - """ - Arguments: - Apply gaussian smoothing on a 1d, 2d or 3d tensor. Filtering is performed separately for each channel in the input - using a depthwise convolution. - channels (int, sequence): Number of channels of the input tensors. Output will - have this number of channels as well. - kernel_size (int, sequence): Size of the gaussian kernel. sigma (float, sequence): Standard deviation of the - gaussian kernel. dim (int, optional): The number of dimensions of the data. - Default value is 2 (spatial). - """ - - # channels=1, kernel_size=kernel_size, sigma=sigma, dim=2 - def __init__( - self, - channels: int = 1, - kernel_size: int = 3, - sigma: float = 0.5, - dim: int = 2, - ): - super().__init__() - - if isinstance(kernel_size, int): - kernel_size = [kernel_size] * dim - if isinstance(sigma, float): - sigma = [sigma] * dim - - # The gaussian kernel is the product of the - # gaussian function of each dimension. - kernel = 1 - meshgrids = torch.meshgrid([torch.arange(size, dtype=torch.float32) for size in kernel_size]) - for size, std, mgrid in zip(kernel_size, sigma, meshgrids): - mean = (size - 1) / 2 - kernel *= 1 / (std * math.sqrt(2 * math.pi)) * torch.exp(-(((mgrid - mean) / (2 * std)) ** 2)) - - # Make sure sum of values in gaussian kernel equals 1. - kernel = kernel / torch.sum(kernel) - - # Reshape to depthwise convolutional weight - kernel = kernel.view(1, 1, *kernel.size()) - kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1)) - - self.register_buffer("weight", kernel) - self.groups = channels - - if dim == 1: - self.conv = F.conv1d - elif dim == 2: - self.conv = F.conv2d - elif dim == 3: - self.conv = F.conv3d - else: - raise RuntimeError("Only 1, 2 and 3 dimensions are supported. Received {}.".format(dim)) - - def forward(self, input): - """ - Arguments: - Apply gaussian filter to input. - input (torch.Tensor): Input to apply gaussian filter on. - Returns: - filtered (torch.Tensor): Filtered output. - """ - return self.conv(input, weight=self.weight.to(input.dtype), groups=self.groups) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_diffedit/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_diffedit/__init__.py deleted file mode 100644 index e2145edb96c6be124abf9e9a21b9a5e8a3f3d641..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_diffedit/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -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_stable_diffusion_diffedit"] = ["StableDiffusionDiffEditPipeline"] - -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_stable_diffusion_diffedit import StableDiffusionDiffEditPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_diffedit/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_diffedit/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 61866386c3a7fa83a87ea47d814aa5f72d322f43..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_diffedit/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_diffedit/__pycache__/pipeline_stable_diffusion_diffedit.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_diffedit/__pycache__/pipeline_stable_diffusion_diffedit.cpython-310.pyc deleted file mode 100644 index a9a64fd539aa6ed9577da209ad1e0f6559c68bac..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_diffedit/__pycache__/pipeline_stable_diffusion_diffedit.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py deleted file mode 100644 index 87bd9f4444ac954da3bdd067cbace2dcb6f2df44..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py +++ /dev/null @@ -1,1552 +0,0 @@ -# Copyright 2025 DiffEdit Authors and Pix2Pix Zero Authors 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, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from packaging import version -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer - -from ...configuration_utils import FrozenDict -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 DDIMInverseScheduler, KarrasDiffusionSchedulers -from ...utils import ( - PIL_INTERPOLATION, - USE_PEFT_BACKEND, - BaseOutput, - deprecate, - 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 DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -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 - - -@dataclass -class DiffEditInversionPipelineOutput(BaseOutput): - """ - Output class for Stable Diffusion pipelines. - - 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, - 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] - - -EXAMPLE_DOC_STRING = """ - - ```py - >>> import PIL - >>> import requests - >>> import torch - >>> from io import BytesIO - - >>> from diffusers import StableDiffusionDiffEditPipeline - - - >>> def download_image(url): - ... response = requests.get(url) - ... return PIL.Image.open(BytesIO(response.content)).convert("RGB") - - - >>> img_url = "https://github.com/Xiang-cd/DiffEdit-stable-diffusion/raw/main/assets/origin.png" - - >>> init_image = download_image(img_url).resize((768, 768)) - - >>> pipeline = StableDiffusionDiffEditPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16 - ... ) - - >>> pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config) - >>> pipeline.inverse_scheduler = DDIMInverseScheduler.from_config(pipeline.scheduler.config) - >>> pipeline.enable_model_cpu_offload() - - >>> mask_prompt = "A bowl of fruits" - >>> prompt = "A bowl of pears" - - >>> mask_image = pipeline.generate_mask(image=init_image, source_prompt=prompt, target_prompt=mask_prompt) - >>> image_latents = pipeline.invert(image=init_image, prompt=mask_prompt).latents - >>> image = pipeline(prompt=prompt, mask_image=mask_image, image_latents=image_latents).images[0] - ``` -""" - -EXAMPLE_INVERT_DOC_STRING = """ - ```py - >>> import PIL - >>> import requests - >>> import torch - >>> from io import BytesIO - - >>> from diffusers import StableDiffusionDiffEditPipeline - - - >>> def download_image(url): - ... response = requests.get(url) - ... return PIL.Image.open(BytesIO(response.content)).convert("RGB") - - - >>> img_url = "https://github.com/Xiang-cd/DiffEdit-stable-diffusion/raw/main/assets/origin.png" - - >>> init_image = download_image(img_url).resize((768, 768)) - - >>> pipeline = StableDiffusionDiffEditPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16 - ... ) - - >>> pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config) - >>> pipeline.inverse_scheduler = DDIMInverseScheduler.from_config(pipeline.scheduler.config) - >>> pipeline.enable_model_cpu_offload() - - >>> prompt = "A bowl of fruits" - - >>> inverted_latents = pipeline.invert(image=init_image, prompt=prompt).latents - ``` -""" - - -def auto_corr_loss(hidden_states, generator=None): - reg_loss = 0.0 - for i in range(hidden_states.shape[0]): - for j in range(hidden_states.shape[1]): - noise = hidden_states[i : i + 1, j : j + 1, :, :] - while True: - roll_amount = torch.randint(noise.shape[2] // 2, (1,), generator=generator).item() - reg_loss += (noise * torch.roll(noise, shifts=roll_amount, dims=2)).mean() ** 2 - reg_loss += (noise * torch.roll(noise, shifts=roll_amount, dims=3)).mean() ** 2 - - if noise.shape[2] <= 8: - break - noise = torch.nn.functional.avg_pool2d(noise, kernel_size=2) - return reg_loss - - -def kl_divergence(hidden_states): - return hidden_states.var() + hidden_states.mean() ** 2 - 1 - torch.log(hidden_states.var() + 1e-7) - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess -def preprocess(image): - 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): - return image - elif isinstance(image, PIL.Image.Image): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - w, h = image[0].size - w, h = (x - x % 8 for x in (w, h)) # resize to integer multiple of 8 - - image = [np.array(i.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image] - image = np.concatenate(image, axis=0) - image = np.array(image).astype(np.float32) / 255.0 - image = image.transpose(0, 3, 1, 2) - image = 2.0 * image - 1.0 - image = torch.from_numpy(image) - elif isinstance(image[0], torch.Tensor): - image = torch.cat(image, dim=0) - return image - - -def preprocess_mask(mask, batch_size: int = 1): - if not isinstance(mask, torch.Tensor): - # preprocess mask - if isinstance(mask, (PIL.Image.Image, np.ndarray)): - mask = [mask] - - if isinstance(mask, list): - if isinstance(mask[0], PIL.Image.Image): - mask = [np.array(m.convert("L")).astype(np.float32) / 255.0 for m in mask] - if isinstance(mask[0], np.ndarray): - mask = np.stack(mask, axis=0) if mask[0].ndim < 3 else np.concatenate(mask, axis=0) - mask = torch.from_numpy(mask) - elif isinstance(mask[0], torch.Tensor): - mask = torch.stack(mask, dim=0) if mask[0].ndim < 3 else torch.cat(mask, dim=0) - - # Batch and add channel dim for single mask - if mask.ndim == 2: - mask = mask.unsqueeze(0).unsqueeze(0) - - # Batch single mask or add channel dim - if mask.ndim == 3: - # Single batched mask, no channel dim or single mask not batched but channel dim - if mask.shape[0] == 1: - mask = mask.unsqueeze(0) - - # Batched masks no channel dim - else: - mask = mask.unsqueeze(1) - - # Check mask shape - if batch_size > 1: - if mask.shape[0] == 1: - mask = torch.cat([mask] * batch_size) - elif mask.shape[0] > 1 and mask.shape[0] != batch_size: - raise ValueError( - f"`mask_image` with batch size {mask.shape[0]} cannot be broadcasted to batch size {batch_size} " - f"inferred by prompt inputs" - ) - - if mask.shape[1] != 1: - raise ValueError(f"`mask_image` must have 1 channel, but has {mask.shape[1]} channels") - - # Check mask is in [0, 1] - if mask.min() < 0 or mask.max() > 1: - raise ValueError("`mask_image` should be in [0, 1] range") - - # Binarize mask - mask[mask < 0.5] = 0 - mask[mask >= 0.5] = 1 - - return mask - - -class StableDiffusionDiffEditPipeline( - DeprecatedPipelineMixin, - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, -): - r""" - - - This is an experimental feature! - - - - Pipeline for text-guided image inpainting using Stable Diffusion and DiffEdit. - - 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.). - - The pipeline also inherits the following loading and saving 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 - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. - inverse_scheduler ([`DDIMInverseScheduler`]): - A scheduler to be used in combination with `unet` to fill in the unmasked part of the input latents. - 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - _last_supported_version = "0.33.1" - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "inverse_scheduler"] - _exclude_from_cpu_offload = ["safety_checker"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - inverse_scheduler: DDIMInverseScheduler, - requires_safety_checker: bool = True, - ): - super().__init__() - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "skip_prk_steps", True) is False: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration" - " `skip_prk_steps`. `skip_prk_steps` should be set to True in the configuration file. Please make" - " sure to update the config accordingly as not setting `skip_prk_steps` in the config might lead to" - " incorrect results in future versions. If you have downloaded this checkpoint from the Hugging Face" - " Hub, it would be very nice if you could open a Pull request for the" - " `scheduler/scheduler_config.json` file" - ) - deprecate("skip_prk_steps not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["skip_prk_steps"] = True - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" - " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_config) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - inverse_scheduler=inverse_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.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # 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, - strength, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ): - if (strength is None) or (strength is not None and (strength < 0 or strength > 1)): - raise ValueError( - f"The value of `strength` should in [0.0, 1.0] but is, but is {strength} of type {type(strength)}." - ) - - if (callback_steps is None) or ( - 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 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 check_source_inputs( - self, - source_prompt=None, - source_negative_prompt=None, - source_prompt_embeds=None, - source_negative_prompt_embeds=None, - ): - if source_prompt is not None and source_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `source_prompt`: {source_prompt} and `source_prompt_embeds`: {source_prompt_embeds}." - " Please make sure to only forward one of the two." - ) - elif source_prompt is None and source_prompt_embeds is None: - raise ValueError( - "Provide either `source_image` or `source_prompt_embeds`. Cannot leave all both of the arguments undefined." - ) - elif source_prompt is not None and ( - not isinstance(source_prompt, str) and not isinstance(source_prompt, list) - ): - raise ValueError(f"`source_prompt` has to be of type `str` or `list` but is {type(source_prompt)}") - - if source_negative_prompt is not None and source_negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `source_negative_prompt`: {source_negative_prompt} and `source_negative_prompt_embeds`:" - f" {source_negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - if source_prompt_embeds is not None and source_negative_prompt_embeds is not None: - if source_prompt_embeds.shape != source_negative_prompt_embeds.shape: - raise ValueError( - "`source_prompt_embeds` and `source_negative_prompt_embeds` must have the same shape when passed" - f" directly, but got: `source_prompt_embeds` {source_prompt_embeds.shape} !=" - f" `source_negative_prompt_embeds` {source_negative_prompt_embeds.shape}." - ) - - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :] - - return timesteps, num_inference_steps - t_start - - def get_inverse_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - - # safety for t_start overflow to prevent empty timsteps slice - if t_start == 0: - return self.inverse_scheduler.timesteps, num_inference_steps - timesteps = self.inverse_scheduler.timesteps[:-t_start] - - return timesteps, num_inference_steps - t_start - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def prepare_image_latents(self, image, batch_size, dtype, device, generator=None): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - if image.shape[1] == 4: - latents = image - - else: - 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 isinstance(generator, list): - latents = [ - self.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size) - ] - latents = torch.cat(latents, dim=0) - else: - latents = self.vae.encode(image).latent_dist.sample(generator) - - latents = self.vae.config.scaling_factor * latents - - if batch_size != latents.shape[0]: - if batch_size % latents.shape[0] == 0: - # expand image_latents for batch_size - deprecation_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {latents.shape[0]} initial" - " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" - " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" - " your script to pass as many initial images as text prompts to suppress this warning." - ) - deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) - additional_latents_per_image = batch_size // latents.shape[0] - latents = torch.cat([latents] * additional_latents_per_image, dim=0) - else: - raise ValueError( - f"Cannot duplicate `image` of batch size {latents.shape[0]} to {batch_size} text prompts." - ) - else: - latents = torch.cat([latents], dim=0) - - return latents - - def get_epsilon(self, model_output: torch.Tensor, sample: torch.Tensor, timestep: int): - pred_type = self.inverse_scheduler.config.prediction_type - alpha_prod_t = self.inverse_scheduler.alphas_cumprod[timestep] - - beta_prod_t = 1 - alpha_prod_t - - if pred_type == "epsilon": - return model_output - elif pred_type == "sample": - return (sample - alpha_prod_t ** (0.5) * model_output) / beta_prod_t ** (0.5) - elif pred_type == "v_prediction": - return (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - else: - raise ValueError( - f"prediction_type given as {pred_type} must be one of `epsilon`, `sample`, or `v_prediction`" - ) - - @torch.no_grad() - @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, - 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, - ): - r""" - Generate a latent mask given a mask prompt, a target prompt, and an image. - - 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*): - 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*): - 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*): - 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. - 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*): - 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*): - 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*): - Pre-generated text embeddings to guide the semantic mask generation. Can be used to easily tweak text - inputs (prompt weighting). If not provided, text embeddings are generated from `source_prompt` input - argument. - source_negative_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated text embeddings to negatively guide the semantic mask generation. Can be used to easily - tweak text inputs (prompt weighting). If not provided, text embeddings are generated from - `source_negative_prompt` input argument. - num_maps_per_mask (`int`, *optional*, defaults to 10): - The number of noise maps sampled to generate the semantic mask using DiffEdit. - mask_encode_strength (`float`, *optional*, defaults to 0.5): - The strength of the noise maps sampled to generate the semantic mask using DiffEdit. Must be between 0 - and 1. - mask_thresholding_ratio (`float`, *optional*, defaults to 3.0): - The maximum multiple of the mean absolute difference used to clamp the semantic guidance map before - mask binarization. - 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): - 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*): - A [`torch.Generator`](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 generated image. Choose between `PIL.Image` or `np.array`. - cross_attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the - [`~models.attention_processor.AttnProcessor`] as defined in - [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - - 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 - 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)`. - """ - - # 1. Check inputs (Provide dummy argument for callback_steps) - self.check_inputs( - target_prompt, - mask_encode_strength, - 1, - target_negative_prompt, - target_prompt_embeds, - target_negative_prompt_embeds, - ) - - self.check_source_inputs( - source_prompt, - source_negative_prompt, - source_prompt_embeds, - source_negative_prompt_embeds, - ) - - if (num_maps_per_mask is None) or ( - num_maps_per_mask is not None and (not isinstance(num_maps_per_mask, int) or num_maps_per_mask <= 0) - ): - raise ValueError( - f"`num_maps_per_mask` has to be a positive integer but is {num_maps_per_mask} of type" - f" {type(num_maps_per_mask)}." - ) - - if mask_thresholding_ratio is None or mask_thresholding_ratio <= 0: - raise ValueError( - f"`mask_thresholding_ratio` has to be positive but is {mask_thresholding_ratio} of type" - f" {type(mask_thresholding_ratio)}." - ) - - # 2. Define call parameters - if target_prompt is not None and isinstance(target_prompt, str): - batch_size = 1 - elif target_prompt is not None and isinstance(target_prompt, list): - batch_size = len(target_prompt) - else: - batch_size = target_prompt_embeds.shape[0] - if cross_attention_kwargs is None: - cross_attention_kwargs = {} - - 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 = guidance_scale > 1.0 - - # 3. Encode input prompts - (cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None) - target_negative_prompt_embeds, target_prompt_embeds = self.encode_prompt( - target_prompt, - device, - num_maps_per_mask, - do_classifier_free_guidance, - target_negative_prompt, - prompt_embeds=target_prompt_embeds, - negative_prompt_embeds=target_negative_prompt_embeds, - ) - # 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: - target_prompt_embeds = torch.cat([target_negative_prompt_embeds, target_prompt_embeds]) - - source_negative_prompt_embeds, source_prompt_embeds = self.encode_prompt( - source_prompt, - device, - num_maps_per_mask, - do_classifier_free_guidance, - source_negative_prompt, - prompt_embeds=source_prompt_embeds, - negative_prompt_embeds=source_negative_prompt_embeds, - ) - if do_classifier_free_guidance: - source_prompt_embeds = torch.cat([source_negative_prompt_embeds, source_prompt_embeds]) - - # 4. Preprocess image - image = self.image_processor.preprocess(image).repeat_interleave(num_maps_per_mask, dim=0) - - # 5. Set timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, _ = self.get_timesteps(num_inference_steps, mask_encode_strength, device) - encode_timestep = timesteps[0] - - # 6. Prepare image latents and add noise with specified strength - image_latents = self.prepare_image_latents( - image, batch_size * num_maps_per_mask, self.vae.dtype, device, generator - ) - noise = randn_tensor(image_latents.shape, generator=generator, device=device, dtype=self.vae.dtype) - image_latents = self.scheduler.add_noise(image_latents, noise, encode_timestep) - - latent_model_input = torch.cat([image_latents] * (4 if do_classifier_free_guidance else 2)) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, encode_timestep) - - # 7. Predict the noise residual - prompt_embeds = torch.cat([source_prompt_embeds, target_prompt_embeds]) - noise_pred = self.unet( - latent_model_input, - encode_timestep, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - ).sample - - if do_classifier_free_guidance: - noise_pred_neg_src, noise_pred_source, noise_pred_uncond, noise_pred_target = noise_pred.chunk(4) - noise_pred_source = noise_pred_neg_src + guidance_scale * (noise_pred_source - noise_pred_neg_src) - noise_pred_target = noise_pred_uncond + guidance_scale * (noise_pred_target - noise_pred_uncond) - else: - noise_pred_source, noise_pred_target = noise_pred.chunk(2) - - # 8. Compute the mask from the absolute difference of predicted noise residuals - # TODO: Consider smoothing mask guidance map - mask_guidance_map = ( - torch.abs(noise_pred_target - noise_pred_source) - .reshape(batch_size, num_maps_per_mask, *noise_pred_target.shape[-3:]) - .mean([1, 2]) - ) - clamp_magnitude = mask_guidance_map.mean() * mask_thresholding_ratio - semantic_mask_image = mask_guidance_map.clamp(0, clamp_magnitude) / clamp_magnitude - semantic_mask_image = torch.where(semantic_mask_image <= 0.5, 0, 1) - mask_image = semantic_mask_image.cpu().numpy() - - # 9. Convert to Numpy array or PIL. - if output_type == "pil": - mask_image = self.image_processor.numpy_to_pil(mask_image) - - # Offload all models - self.maybe_free_model_hooks() - - return mask_image - - @torch.no_grad() - @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, - 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, - decode_latents: bool = False, - output_type: Optional[str] = "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, - lambda_auto_corr: float = 20.0, - lambda_kl: float = 20.0, - num_reg_steps: int = 0, - num_auto_corr_rolls: int = 5, - ): - r""" - Generate inverted latents given a prompt and image. - - Args: - 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`. - inpaint_strength (`float`, *optional*, defaults to 0.8): - Indicates extent of the noising process to run latent inversion. Must be between 0 and 1. When - `inpaint_strength` is 1, the inversion process is run for the full number of iterations specified in - `num_inference_steps`. `image` is used as a reference for the inversion process, and adding more noise - increases `inpaint_strength`. If `inpaint_strength` is 0, no inpainting occurs. - 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): - 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*): - 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*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - 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. - 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. - decode_latents (`bool`, *optional*, defaults to `False`): - Whether or not to decode the inverted latents into a generated image. Setting this argument to `True` - decodes all inverted latents for each timestep into a list of generated images. - 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 [`~pipelines.stable_diffusion.DiffEditInversionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - cross_attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the - [`~models.attention_processor.AttnProcessor`] as defined in - [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - lambda_auto_corr (`float`, *optional*, defaults to 20.0): - Lambda parameter to control auto correction. - lambda_kl (`float`, *optional*, defaults to 20.0): - Lambda parameter to control Kullback-Leibler divergence output. - num_reg_steps (`int`, *optional*, defaults to 0): - Number of regularization loss steps. - num_auto_corr_rolls (`int`, *optional*, defaults to 5): - Number of auto correction roll steps. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.pipeline_stable_diffusion_diffedit.DiffEditInversionPipelineOutput`] or - `tuple`: - If `return_dict` is `True`, - [`~pipelines.stable_diffusion.pipeline_stable_diffusion_diffedit.DiffEditInversionPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is the inverted latents tensors - ordered by increasing noise, and the second is the corresponding decoded images if `decode_latents` is - `True`, otherwise `None`. - """ - - # 1. Check inputs - self.check_inputs( - prompt, - inpaint_strength, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ) - - if image is None: - raise ValueError("`image` input cannot be undefined.") - - # 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 cross_attention_kwargs is None: - cross_attention_kwargs = {} - - 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 = guidance_scale > 1.0 - - # 3. Preprocess image - image = self.image_processor.preprocess(image) - - # 4. Prepare latent variables - num_images_per_prompt = 1 - latents = self.prepare_image_latents( - image, batch_size * num_images_per_prompt, self.vae.dtype, device, generator - ) - - # 5. 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, - ) - # 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]) - - # 6. Prepare timesteps - self.inverse_scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_inverse_timesteps(num_inference_steps, inpaint_strength, device) - - # 7. Noising loop where we obtain the intermediate noised latent image for each timestep. - num_warmup_steps = len(timesteps) - num_inference_steps * self.inverse_scheduler.order - inverted_latents = [] - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.inverse_scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - ).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # regularization of the noise prediction (not in original code or paper but borrowed from Pix2PixZero) - if num_reg_steps > 0: - with torch.enable_grad(): - for _ in range(num_reg_steps): - if lambda_auto_corr > 0: - for _ in range(num_auto_corr_rolls): - var = torch.autograd.Variable(noise_pred.detach().clone(), requires_grad=True) - - # Derive epsilon from model output before regularizing to IID standard normal - var_epsilon = self.get_epsilon(var, latent_model_input.detach(), t) - - l_ac = auto_corr_loss(var_epsilon, generator=generator) - l_ac.backward() - - grad = var.grad.detach() / num_auto_corr_rolls - noise_pred = noise_pred - lambda_auto_corr * grad - - if lambda_kl > 0: - var = torch.autograd.Variable(noise_pred.detach().clone(), requires_grad=True) - - # Derive epsilon from model output before regularizing to IID standard normal - var_epsilon = self.get_epsilon(var, latent_model_input.detach(), t) - - l_kld = kl_divergence(var_epsilon) - l_kld.backward() - - grad = var.grad.detach() - noise_pred = noise_pred - lambda_kl * grad - - noise_pred = noise_pred.detach() - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.inverse_scheduler.step(noise_pred, t, latents).prev_sample - inverted_latents.append(latents.detach().clone()) - - # call the callback, if provided - if i == len(timesteps) - 1 or ( - (i + 1) > num_warmup_steps and (i + 1) % self.inverse_scheduler.order == 0 - ): - progress_bar.update() - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - assert len(inverted_latents) == len(timesteps) - latents = torch.stack(list(reversed(inverted_latents)), 1) - - # 8. Post-processing - image = None - if decode_latents: - image = self.decode_latents(latents.flatten(0, 1)) - - # 9. Convert to PIL. - if decode_latents and output_type == "pil": - image = self.image_processor.numpy_to_pil(image) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (latents, image) - - return DiffEditInversionPipelineOutput(latents=latents, images=image) - - @torch.no_grad() - @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, - 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, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: int = None, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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 - 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 expected shape would be `(B, 1, H, W)`. - image_latents (`PIL.Image.Image` or `torch.Tensor`): - Partially noised image latents from the inversion process to be used as inputs for image generation. - inpaint_strength (`float`, *optional*, defaults to 0.8): - Indicates extent to inpaint the masked area. Must be between 0 and 1. When `inpaint_strength` is 1, the - denoising process is run on the masked area for the full number of iterations specified in - `num_inference_steps`. `image_latents` is used as a reference for the masked area, and adding more - noise to a region increases `inpaint_strength`. If `inpaint_strength` is 0, no inpainting occurs. - 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): - 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*): - 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): - The number of images to generate per prompt. - 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`, *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. - 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - 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. - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - - # 1. Check inputs - self.check_inputs( - prompt, - inpaint_strength, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - ) - - if mask_image is None: - raise ValueError( - "`mask_image` input cannot be undefined. Use `generate_mask()` to compute `mask_image` from text prompts." - ) - if image_latents is None: - raise ValueError( - "`image_latents` input cannot be undefined. Use `invert()` to compute `image_latents` from input images." - ) - - # 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 cross_attention_kwargs is None: - cross_attention_kwargs = {} - - 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - 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, - lora_scale=text_encoder_lora_scale, - 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. Preprocess mask - mask_image = preprocess_mask(mask_image, batch_size) - latent_height, latent_width = mask_image.shape[-2:] - mask_image = torch.cat([mask_image] * num_images_per_prompt) - mask_image = mask_image.to(device=device, dtype=prompt_embeds.dtype) - - # 5. Set timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, inpaint_strength, device) - - # 6. Preprocess image latents - if isinstance(image_latents, list) and any(isinstance(l, torch.Tensor) and l.ndim == 5 for l in image_latents): - image_latents = torch.cat(image_latents).detach() - elif isinstance(image_latents, torch.Tensor) and image_latents.ndim == 5: - image_latents = image_latents.detach() - else: - image_latents = self.image_processor.preprocess(image_latents).detach() - - latent_shape = (self.vae.config.latent_channels, latent_height, latent_width) - if image_latents.shape[-3:] != latent_shape: - raise ValueError( - f"Each latent image in `image_latents` must have shape {latent_shape}, " - f"but has shape {image_latents.shape[-3:]}" - ) - if image_latents.ndim == 4: - image_latents = image_latents.reshape(batch_size, len(timesteps), *latent_shape) - if image_latents.shape[:2] != (batch_size, len(timesteps)): - raise ValueError( - f"`image_latents` must have batch size {batch_size} with latent images from {len(timesteps)}" - f" timesteps, but has batch size {image_latents.shape[0]} with latent images from" - f" {image_latents.shape[1]} timesteps." - ) - image_latents = image_latents.transpose(0, 1).repeat_interleave(num_images_per_prompt, dim=1) - image_latents = image_latents.to(device=device, dtype=prompt_embeds.dtype) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 8. Denoising loop - latents = image_latents[0].clone() - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - ).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # mask with inverted latents from appropriate timestep - use original image latent for last step - latents = latents * mask_image + image_latents[i] * (1 - mask_image) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/__init__.py deleted file mode 100644 index 147980cbf9e5c3418fc1854787ae37b25e4fed56..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/__init__.py +++ /dev/null @@ -1,50 +0,0 @@ -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_stable_diffusion_gligen"] = ["StableDiffusionGLIGENPipeline"] - _import_structure["pipeline_stable_diffusion_gligen_text_image"] = ["StableDiffusionGLIGENTextImagePipeline"] - -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_stable_diffusion_gligen import StableDiffusionGLIGENPipeline - from .pipeline_stable_diffusion_gligen_text_image import StableDiffusionGLIGENTextImagePipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index c753d6f6890599f4eaac42c342260007e671fb78..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/__pycache__/pipeline_stable_diffusion_gligen.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/__pycache__/pipeline_stable_diffusion_gligen.cpython-310.pyc deleted file mode 100644 index 03e2784dcbb4f68f174601dd74bb928c81c37b07..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/__pycache__/pipeline_stable_diffusion_gligen.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/__pycache__/pipeline_stable_diffusion_gligen_text_image.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/__pycache__/pipeline_stable_diffusion_gligen_text_image.cpython-310.pyc deleted file mode 100644 index 4dd97f39547d6a619758c851db7f461492af602e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/__pycache__/pipeline_stable_diffusion_gligen_text_image.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py deleted file mode 100644 index 78b026684cfaa8aef49bf4360bb472ade9fa24e9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py +++ /dev/null @@ -1,865 +0,0 @@ -# Copyright 2025 The GLIGEN Authors and 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 warnings -from typing import Any, Callable, Dict, List, Optional, Union - -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer - -from ...image_processor import VaeImageProcessor -from ...loaders import StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, UNet2DConditionModel -from ...models.attention import GatedSelfAttentionDense -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - 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 DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -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 StableDiffusionGLIGENPipeline - >>> from diffusers.utils import load_image - - >>> # Insert objects described by text at the region defined by bounding boxes - >>> pipe = StableDiffusionGLIGENPipeline.from_pretrained( - ... "masterful/gligen-1-4-inpainting-text-box", variant="fp16", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> input_image = load_image( - ... "https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/gligen/livingroom_modern.png" - ... ) - >>> prompt = "a birthday cake" - >>> boxes = [[0.2676, 0.6088, 0.4773, 0.7183]] - >>> phrases = ["a birthday cake"] - - >>> images = pipe( - ... prompt=prompt, - ... gligen_phrases=phrases, - ... gligen_inpaint_image=input_image, - ... gligen_boxes=boxes, - ... gligen_scheduled_sampling_beta=1, - ... output_type="pil", - ... num_inference_steps=50, - ... ).images - - >>> images[0].save("./gligen-1-4-inpainting-text-box.jpg") - - >>> # Generate an image described by the prompt and - >>> # insert objects described by text at the region defined by bounding boxes - >>> pipe = StableDiffusionGLIGENPipeline.from_pretrained( - ... "masterful/gligen-1-4-generation-text-box", variant="fp16", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = "a waterfall and a modern high speed train running through the tunnel in a beautiful forest with fall foliage" - >>> boxes = [[0.1387, 0.2051, 0.4277, 0.7090], [0.4980, 0.4355, 0.8516, 0.7266]] - >>> phrases = ["a waterfall", "a modern high speed train running through the tunnel"] - - >>> images = pipe( - ... prompt=prompt, - ... gligen_phrases=phrases, - ... gligen_boxes=boxes, - ... gligen_scheduled_sampling_beta=1, - ... output_type="pil", - ... num_inference_steps=50, - ... ).images - - >>> images[0].save("./gligen-1-4-generation-text-box.jpg") - ``` -""" - - -class StableDiffusionGLIGENPipeline(DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin): - r""" - Pipeline for text-to-image generation using Stable Diffusion with Grounded-Language-to-Image Generation (GLIGEN). - - 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.). - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - _last_supported_version = "0.33.1" - - _optional_components = ["safety_checker", "feature_extractor"] - model_cpu_offload_seq = "text_encoder->unet->vae" - _exclude_from_cpu_offload = ["safety_checker"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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 - - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - callback_steps, - gligen_phrases, - gligen_boxes, - negative_prompt=None, - prompt_embeds=None, - negative_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 (callback_steps is None) or ( - 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 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 len(gligen_phrases) != len(gligen_boxes): - raise ValueError( - "length of `gligen_phrases` and `gligen_boxes` has to be same, but" - f" got: `gligen_phrases` {len(gligen_phrases)} != `gligen_boxes` {len(gligen_boxes)}" - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def enable_fuser(self, enabled=True): - for module in self.unet.modules(): - if type(module) is GatedSelfAttentionDense: - module.enabled = enabled - - def draw_inpaint_mask_from_boxes(self, boxes, size): - inpaint_mask = torch.ones(size[0], size[1]) - for box in boxes: - x0, x1 = box[0] * size[0], box[2] * size[0] - y0, y1 = box[1] * size[1], box[3] * size[1] - inpaint_mask[int(y0) : int(y1), int(x0) : int(x1)] = 0 - return inpaint_mask - - def crop(self, im, new_width, new_height): - width, height = im.size - left = (width - new_width) / 2 - top = (height - new_height) / 2 - right = (width + new_width) / 2 - bottom = (height + new_height) / 2 - return im.crop((left, top, right, bottom)) - - def target_size_center_crop(self, im, new_hw): - width, height = im.size - if width != height: - im = self.crop(im, min(height, width), min(height, width)) - return im.resize((new_hw, new_hw), PIL.Image.LANCZOS) - - @torch.no_grad() - @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: 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, - 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, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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): - 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]`): - 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]]`): - 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]. - 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. - gligen_scheduled_sampling_beta (`float`, defaults to 0.3): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - guidance_rescale (`float`, *optional*, defaults to 0.0): - Guidance rescale factor from [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. - 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. - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - # 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, - gligen_phrases, - gligen_boxes, - 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 = guidance_scale > 1.0 - - # 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=device) - timesteps = self.scheduler.timesteps - - # 5. 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, - ) - - # 5.1 Prepare GLIGEN variables - max_objs = 30 - if len(gligen_boxes) > max_objs: - warnings.warn( - f"More that {max_objs} objects found. Only first {max_objs} objects will be processed.", - FutureWarning, - ) - gligen_phrases = gligen_phrases[:max_objs] - gligen_boxes = gligen_boxes[:max_objs] - # prepare batched input to the GLIGENTextBoundingboxProjection (boxes, phrases, mask) - # Get tokens for phrases from pre-trained CLIPTokenizer - tokenizer_inputs = self.tokenizer(gligen_phrases, padding=True, return_tensors="pt").to(device) - # For the token, we use the same pre-trained text encoder - # to obtain its text feature - _text_embeddings = self.text_encoder(**tokenizer_inputs).pooler_output - n_objs = len(gligen_boxes) - # For each entity, described in phrases, is denoted with a bounding box, - # we represent the location information as (xmin,ymin,xmax,ymax) - boxes = torch.zeros(max_objs, 4, device=device, dtype=self.text_encoder.dtype) - boxes[:n_objs] = torch.tensor(gligen_boxes) - text_embeddings = torch.zeros( - max_objs, self.unet.config.cross_attention_dim, device=device, dtype=self.text_encoder.dtype - ) - text_embeddings[:n_objs] = _text_embeddings - # Generate a mask for each object that is entity described by phrases - masks = torch.zeros(max_objs, device=device, dtype=self.text_encoder.dtype) - masks[:n_objs] = 1 - - repeat_batch = batch_size * num_images_per_prompt - boxes = boxes.unsqueeze(0).expand(repeat_batch, -1, -1).clone() - text_embeddings = text_embeddings.unsqueeze(0).expand(repeat_batch, -1, -1).clone() - masks = masks.unsqueeze(0).expand(repeat_batch, -1).clone() - if do_classifier_free_guidance: - repeat_batch = repeat_batch * 2 - boxes = torch.cat([boxes] * 2) - text_embeddings = torch.cat([text_embeddings] * 2) - masks = torch.cat([masks] * 2) - masks[: repeat_batch // 2] = 0 - if cross_attention_kwargs is None: - cross_attention_kwargs = {} - cross_attention_kwargs["gligen"] = {"boxes": boxes, "positive_embeddings": text_embeddings, "masks": masks} - - # Prepare latent variables for GLIGEN inpainting - if gligen_inpaint_image is not None: - # if the given input image is not of the same size as expected by VAE - # center crop and resize the input image to expected shape - if gligen_inpaint_image.size != (self.vae.sample_size, self.vae.sample_size): - gligen_inpaint_image = self.target_size_center_crop(gligen_inpaint_image, self.vae.sample_size) - # Convert a single image into a batch of images with a batch size of 1 - # The resulting shape becomes (1, C, H, W), where C is the number of channels, - # and H and W are the height and width of the image. - # scales the pixel values to a range [-1, 1] - gligen_inpaint_image = self.image_processor.preprocess(gligen_inpaint_image) - gligen_inpaint_image = gligen_inpaint_image.to(dtype=self.vae.dtype, device=self.vae.device) - # Run AutoEncoder to get corresponding latents - gligen_inpaint_latent = self.vae.encode(gligen_inpaint_image).latent_dist.sample() - gligen_inpaint_latent = self.vae.config.scaling_factor * gligen_inpaint_latent - # Generate an inpainting mask - # pixel value = 0, where the object is present (defined by bounding boxes above) - # 1, everywhere else - gligen_inpaint_mask = self.draw_inpaint_mask_from_boxes(gligen_boxes, gligen_inpaint_latent.shape[2:]) - gligen_inpaint_mask = gligen_inpaint_mask.to( - dtype=gligen_inpaint_latent.dtype, device=gligen_inpaint_latent.device - ) - gligen_inpaint_mask = gligen_inpaint_mask[None, None] - gligen_inpaint_mask_addition = torch.cat( - (gligen_inpaint_latent * gligen_inpaint_mask, gligen_inpaint_mask), dim=1 - ) - # Convert a single mask into a batch of masks with a batch size of 1 - gligen_inpaint_mask_addition = gligen_inpaint_mask_addition.expand(repeat_batch, -1, -1, -1).clone() - - num_grounding_steps = int(gligen_scheduled_sampling_beta * len(timesteps)) - self.enable_fuser(True) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # Scheduled sampling - if i == num_grounding_steps: - self.enable_fuser(False) - - if latents.shape[1] != 4: - latents = torch.randn_like(latents[:, :4]) - - if gligen_inpaint_image is not None: - gligen_inpaint_latent_with_noise = ( - self.scheduler.add_noise( - gligen_inpaint_latent, torch.randn_like(gligen_inpaint_latent), torch.tensor([t]) - ) - .expand(latents.shape[0], -1, -1, -1) - .clone() - ) - latents = gligen_inpaint_latent_with_noise * gligen_inpaint_mask + latents * ( - 1 - gligen_inpaint_mask - ) - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - if gligen_inpaint_image is not None: - latent_model_input = torch.cat((latent_model_input, gligen_inpaint_mask_addition), dim=1) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - ).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py deleted file mode 100644 index 05cbad139d923bdcb6c7e5a10ac14fb531bfcaa8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py +++ /dev/null @@ -1,1053 +0,0 @@ -# Copyright 2025 The GLIGEN Authors and 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 warnings -from typing import Any, Callable, Dict, List, Optional, Union - -import PIL.Image -import torch -from transformers import ( - CLIPImageProcessor, - CLIPProcessor, - CLIPTextModel, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ...image_processor import VaeImageProcessor -from ...loaders import StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, UNet2DConditionModel -from ...models.attention import GatedSelfAttentionDense -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import 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 -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion import StableDiffusionPipelineOutput -from ..stable_diffusion.clip_image_project_model import CLIPImageProjection -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -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 StableDiffusionGLIGENTextImagePipeline - >>> from diffusers.utils import load_image - - >>> # Insert objects described by image at the region defined by bounding boxes - >>> pipe = StableDiffusionGLIGENTextImagePipeline.from_pretrained( - ... "anhnct/Gligen_Inpainting_Text_Image", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> input_image = load_image( - ... "https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/gligen/livingroom_modern.png" - ... ) - >>> prompt = "a backpack" - >>> boxes = [[0.2676, 0.4088, 0.4773, 0.7183]] - >>> phrases = None - >>> gligen_image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/gligen/backpack.jpeg" - ... ) - - >>> images = pipe( - ... prompt=prompt, - ... gligen_phrases=phrases, - ... gligen_inpaint_image=input_image, - ... gligen_boxes=boxes, - ... gligen_images=[gligen_image], - ... gligen_scheduled_sampling_beta=1, - ... output_type="pil", - ... num_inference_steps=50, - ... ).images - - >>> images[0].save("./gligen-inpainting-text-image-box.jpg") - - >>> # Generate an image described by the prompt and - >>> # insert objects described by text and image at the region defined by bounding boxes - >>> pipe = StableDiffusionGLIGENTextImagePipeline.from_pretrained( - ... "anhnct/Gligen_Text_Image", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = "a flower sitting on the beach" - >>> boxes = [[0.0, 0.09, 0.53, 0.76]] - >>> phrases = ["flower"] - >>> gligen_image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/gligen/pexels-pixabay-60597.jpg" - ... ) - - >>> images = pipe( - ... prompt=prompt, - ... gligen_phrases=phrases, - ... gligen_images=[gligen_image], - ... gligen_boxes=boxes, - ... gligen_scheduled_sampling_beta=1, - ... output_type="pil", - ... num_inference_steps=50, - ... ).images - - >>> images[0].save("./gligen-generation-text-image-box.jpg") - - >>> # Generate an image described by the prompt and - >>> # transfer style described by image at the region defined by bounding boxes - >>> pipe = StableDiffusionGLIGENTextImagePipeline.from_pretrained( - ... "anhnct/Gligen_Text_Image", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = "a dragon flying on the sky" - >>> boxes = [[0.4, 0.2, 1.0, 0.8], [0.0, 1.0, 0.0, 1.0]] # Set `[0.0, 1.0, 0.0, 1.0]` for the style - - >>> gligen_image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/landscape.png" - ... ) - - >>> gligen_placeholder = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/landscape.png" - ... ) - - >>> images = pipe( - ... prompt=prompt, - ... gligen_phrases=[ - ... "dragon", - ... "placeholder", - ... ], # Can use any text instead of `placeholder` token, because we will use mask here - ... gligen_images=[ - ... gligen_placeholder, - ... gligen_image, - ... ], # Can use any image in gligen_placeholder, because we will use mask here - ... input_phrases_mask=[1, 0], # Set 0 for the placeholder token - ... input_images_mask=[0, 1], # Set 0 for the placeholder image - ... gligen_boxes=boxes, - ... gligen_scheduled_sampling_beta=1, - ... output_type="pil", - ... num_inference_steps=50, - ... ).images - - >>> images[0].save("./gligen-generation-text-image-box-style-transfer.jpg") - ``` -""" - - -class StableDiffusionGLIGENTextImagePipeline(DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin): - r""" - Pipeline for text-to-image generation using Stable Diffusion with Grounded-Language-to-Image Generation (GLIGEN). - - 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.). - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - text_encoder ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - processor ([`~transformers.CLIPProcessor`]): - A `CLIPProcessor` to process reference image. - image_encoder ([`~transformers.CLIPVisionModelWithProjection`]): - Frozen image-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - image_project ([`CLIPImageProjection`]): - A `CLIPImageProjection` to project image embedding into phrases embedding space. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to 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: CLIPTokenizer, - processor: CLIPProcessor, - image_encoder: CLIPVisionModelWithProjection, - image_project: CLIPImageProjection, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - image_encoder=image_encoder, - processor=processor, - image_project=image_project, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - callback_steps, - gligen_images, - gligen_phrases, - 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}." - ) - - if gligen_images is not None and gligen_phrases is not None: - if len(gligen_images) != len(gligen_phrases): - raise ValueError( - "`gligen_images` and `gligen_phrases` must have the same length when both are provided, but" - f" got: `gligen_images` with length {len(gligen_images)} != `gligen_phrases` with length {len(gligen_phrases)}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def enable_fuser(self, enabled=True): - for module in self.unet.modules(): - if type(module) is GatedSelfAttentionDense: - module.enabled = enabled - - def draw_inpaint_mask_from_boxes(self, boxes, size): - """ - Create an inpainting mask based on given boxes. This function generates an inpainting mask using the provided - boxes to mark regions that need to be inpainted. - """ - inpaint_mask = torch.ones(size[0], size[1]) - for box in boxes: - x0, x1 = box[0] * size[0], box[2] * size[0] - y0, y1 = box[1] * size[1], box[3] * size[1] - inpaint_mask[int(y0) : int(y1), int(x0) : int(x1)] = 0 - return inpaint_mask - - def crop(self, im, new_width, new_height): - """ - Crop the input image to the specified dimensions. - """ - width, height = im.size - left = (width - new_width) / 2 - top = (height - new_height) / 2 - right = (width + new_width) / 2 - bottom = (height + new_height) / 2 - return im.crop((left, top, right, bottom)) - - def target_size_center_crop(self, im, new_hw): - """ - Crop and resize the image to the target size while keeping the center. - """ - width, height = im.size - if width != height: - im = self.crop(im, min(height, width), min(height, width)) - return im.resize((new_hw, new_hw), PIL.Image.LANCZOS) - - def complete_mask(self, has_mask, max_objs, device): - """ - Based on the input mask corresponding value `0 or 1` for each phrases and image, mask the features - corresponding to phrases and images. - """ - mask = torch.ones(1, max_objs).type(self.text_encoder.dtype).to(device) - if has_mask is None: - return mask - - if isinstance(has_mask, int): - return mask * has_mask - else: - for idx, value in enumerate(has_mask): - mask[0, idx] = value - return mask - - def get_clip_feature(self, input, normalize_constant, device, is_image=False): - """ - Get image and phrases embedding by using CLIP pretrain model. The image embedding is transformed into the - phrases embedding space through a projection. - """ - if is_image: - if input is None: - return None - inputs = self.processor(images=[input], return_tensors="pt").to(device) - inputs["pixel_values"] = inputs["pixel_values"].to(self.image_encoder.dtype) - - outputs = self.image_encoder(**inputs) - feature = outputs.image_embeds - feature = self.image_project(feature).squeeze(0) - feature = (feature / feature.norm()) * normalize_constant - feature = feature.unsqueeze(0) - else: - if input is None: - return None - inputs = self.tokenizer(input, return_tensors="pt", padding=True).to(device) - outputs = self.text_encoder(**inputs) - feature = outputs.pooler_output - return feature - - def get_cross_attention_kwargs_with_grounded( - self, - hidden_size, - gligen_phrases, - gligen_images, - gligen_boxes, - input_phrases_mask, - input_images_mask, - repeat_batch, - normalize_constant, - max_objs, - device, - ): - """ - Prepare the cross-attention kwargs containing information about the grounded input (boxes, mask, image - embedding, phrases embedding). - """ - phrases, images = gligen_phrases, gligen_images - images = [None] * len(phrases) if images is None else images - phrases = [None] * len(images) if phrases is None else phrases - - boxes = torch.zeros(max_objs, 4, device=device, dtype=self.text_encoder.dtype) - masks = torch.zeros(max_objs, device=device, dtype=self.text_encoder.dtype) - phrases_masks = torch.zeros(max_objs, device=device, dtype=self.text_encoder.dtype) - image_masks = torch.zeros(max_objs, device=device, dtype=self.text_encoder.dtype) - phrases_embeddings = torch.zeros(max_objs, hidden_size, device=device, dtype=self.text_encoder.dtype) - image_embeddings = torch.zeros(max_objs, hidden_size, device=device, dtype=self.text_encoder.dtype) - - text_features = [] - image_features = [] - for phrase, image in zip(phrases, images): - text_features.append(self.get_clip_feature(phrase, normalize_constant, device, is_image=False)) - image_features.append(self.get_clip_feature(image, normalize_constant, device, is_image=True)) - - for idx, (box, text_feature, image_feature) in enumerate(zip(gligen_boxes, text_features, image_features)): - boxes[idx] = torch.tensor(box) - masks[idx] = 1 - if text_feature is not None: - phrases_embeddings[idx] = text_feature - phrases_masks[idx] = 1 - if image_feature is not None: - image_embeddings[idx] = image_feature - image_masks[idx] = 1 - - input_phrases_mask = self.complete_mask(input_phrases_mask, max_objs, device) - phrases_masks = phrases_masks.unsqueeze(0).repeat(repeat_batch, 1) * input_phrases_mask - input_images_mask = self.complete_mask(input_images_mask, max_objs, device) - image_masks = image_masks.unsqueeze(0).repeat(repeat_batch, 1) * input_images_mask - boxes = boxes.unsqueeze(0).repeat(repeat_batch, 1, 1) - masks = masks.unsqueeze(0).repeat(repeat_batch, 1) - phrases_embeddings = phrases_embeddings.unsqueeze(0).repeat(repeat_batch, 1, 1) - image_embeddings = image_embeddings.unsqueeze(0).repeat(repeat_batch, 1, 1) - - out = { - "boxes": boxes, - "masks": masks, - "phrases_masks": phrases_masks, - "image_masks": image_masks, - "phrases_embeddings": phrases_embeddings, - "image_embeddings": image_embeddings, - } - - return out - - def get_cross_attention_kwargs_without_grounded(self, hidden_size, repeat_batch, max_objs, device): - """ - Prepare the cross-attention kwargs without information about the grounded input (boxes, mask, image embedding, - phrases embedding) (All are zero tensor). - """ - boxes = torch.zeros(max_objs, 4, device=device, dtype=self.text_encoder.dtype) - masks = torch.zeros(max_objs, device=device, dtype=self.text_encoder.dtype) - phrases_masks = torch.zeros(max_objs, device=device, dtype=self.text_encoder.dtype) - image_masks = torch.zeros(max_objs, device=device, dtype=self.text_encoder.dtype) - phrases_embeddings = torch.zeros(max_objs, hidden_size, device=device, dtype=self.text_encoder.dtype) - image_embeddings = torch.zeros(max_objs, hidden_size, device=device, dtype=self.text_encoder.dtype) - - out = { - "boxes": boxes.unsqueeze(0).repeat(repeat_batch, 1, 1), - "masks": masks.unsqueeze(0).repeat(repeat_batch, 1), - "phrases_masks": phrases_masks.unsqueeze(0).repeat(repeat_batch, 1), - "image_masks": image_masks.unsqueeze(0).repeat(repeat_batch, 1), - "phrases_embeddings": phrases_embeddings.unsqueeze(0).repeat(repeat_batch, 1, 1), - "image_embeddings": image_embeddings.unsqueeze(0).repeat(repeat_batch, 1, 1), - } - - return out - - @torch.no_grad() - @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: 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, - 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, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - gligen_normalize_constant: float = 28.7, - clip_skip: int = None, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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): - 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]`): - 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]`): - 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]`): - pre phrases mask input defined by the correspongding `input_phrases_mask` - input_images_mask (`int` or `List[int]`): - pre images mask input defined by the correspongding `input_images_mask` - 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]. - 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. - gligen_scheduled_sampling_beta (`float`, defaults to 0.3): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - gligen_normalize_constant (`float`, *optional*, defaults to 28.7): - The normalize value of the image embedding. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - # 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, - gligen_images, - gligen_phrases, - 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 = guidance_scale > 1.0 - - # 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, - ) - - 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=device) - timesteps = self.scheduler.timesteps - - # 5. 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, - ) - - # 5.1 Prepare GLIGEN variables - max_objs = 30 - if len(gligen_boxes) > max_objs: - warnings.warn( - f"More that {max_objs} objects found. Only first {max_objs} objects will be processed.", - FutureWarning, - ) - gligen_phrases = gligen_phrases[:max_objs] - gligen_boxes = gligen_boxes[:max_objs] - gligen_images = gligen_images[:max_objs] - - repeat_batch = batch_size * num_images_per_prompt - - if do_classifier_free_guidance: - repeat_batch = repeat_batch * 2 - - if cross_attention_kwargs is None: - cross_attention_kwargs = {} - - hidden_size = prompt_embeds.shape[2] - - cross_attention_kwargs["gligen"] = self.get_cross_attention_kwargs_with_grounded( - hidden_size=hidden_size, - gligen_phrases=gligen_phrases, - gligen_images=gligen_images, - gligen_boxes=gligen_boxes, - input_phrases_mask=input_phrases_mask, - input_images_mask=input_images_mask, - repeat_batch=repeat_batch, - normalize_constant=gligen_normalize_constant, - max_objs=max_objs, - device=device, - ) - - cross_attention_kwargs_without_grounded = {} - cross_attention_kwargs_without_grounded["gligen"] = self.get_cross_attention_kwargs_without_grounded( - hidden_size=hidden_size, repeat_batch=repeat_batch, max_objs=max_objs, device=device - ) - - # Prepare latent variables for GLIGEN inpainting - if gligen_inpaint_image is not None: - # if the given input image is not of the same size as expected by VAE - # center crop and resize the input image to expected shape - if gligen_inpaint_image.size != (self.vae.sample_size, self.vae.sample_size): - gligen_inpaint_image = self.target_size_center_crop(gligen_inpaint_image, self.vae.sample_size) - # Convert a single image into a batch of images with a batch size of 1 - # The resulting shape becomes (1, C, H, W), where C is the number of channels, - # and H and W are the height and width of the image. - # scales the pixel values to a range [-1, 1] - gligen_inpaint_image = self.image_processor.preprocess(gligen_inpaint_image) - gligen_inpaint_image = gligen_inpaint_image.to(dtype=self.vae.dtype, device=self.vae.device) - # Run AutoEncoder to get corresponding latents - gligen_inpaint_latent = self.vae.encode(gligen_inpaint_image).latent_dist.sample() - gligen_inpaint_latent = self.vae.config.scaling_factor * gligen_inpaint_latent - # Generate an inpainting mask - # pixel value = 0, where the object is present (defined by bounding boxes above) - # 1, everywhere else - gligen_inpaint_mask = self.draw_inpaint_mask_from_boxes(gligen_boxes, gligen_inpaint_latent.shape[2:]) - gligen_inpaint_mask = gligen_inpaint_mask.to( - dtype=gligen_inpaint_latent.dtype, device=gligen_inpaint_latent.device - ) - gligen_inpaint_mask = gligen_inpaint_mask[None, None] - gligen_inpaint_mask_addition = torch.cat( - (gligen_inpaint_latent * gligen_inpaint_mask, gligen_inpaint_mask), dim=1 - ) - # Convert a single mask into a batch of masks with a batch size of 1 - gligen_inpaint_mask_addition = gligen_inpaint_mask_addition.expand(repeat_batch, -1, -1, -1).clone() - - int(gligen_scheduled_sampling_beta * len(timesteps)) - self.enable_fuser(True) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if latents.shape[1] != 4: - latents = torch.randn_like(latents[:, :4]) - - if gligen_inpaint_image is not None: - gligen_inpaint_latent_with_noise = ( - self.scheduler.add_noise( - gligen_inpaint_latent, torch.randn_like(gligen_inpaint_latent), torch.tensor([t]) - ) - .expand(latents.shape[0], -1, -1, -1) - .clone() - ) - latents = gligen_inpaint_latent_with_noise * gligen_inpaint_mask + latents * ( - 1 - gligen_inpaint_mask - ) - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - if gligen_inpaint_image is not None: - latent_model_input = torch.cat((latent_model_input, gligen_inpaint_mask_addition), dim=1) - - # predict the noise residual with grounded information - noise_pred_with_grounding = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - ).sample - - # predict the noise residual without grounded information - noise_pred_without_grounding = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs_without_grounded, - ).sample - - # perform guidance - if do_classifier_free_guidance: - # Using noise_pred_text from noise residual with grounded information and noise_pred_uncond from noise residual without grounded information - _, noise_pred_text = noise_pred_with_grounding.chunk(2) - noise_pred_uncond, _ = noise_pred_without_grounding.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - else: - noise_pred = noise_pred_with_grounding - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/__init__.py deleted file mode 100644 index 7eb5bf8c229b584c934e33806efa73a3727dbb5b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 6833e43130ffe1f894361d829fab564441aba03a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/__pycache__/pipeline_stable_diffusion_k_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/__pycache__/pipeline_stable_diffusion_k_diffusion.cpython-310.pyc deleted file mode 100644 index 7e331b2a032b69c412715ac5166a4672b321032c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/__pycache__/pipeline_stable_diffusion_k_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/__pycache__/pipeline_stable_diffusion_xl_k_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/__pycache__/pipeline_stable_diffusion_xl_k_diffusion.cpython-310.pyc deleted file mode 100644 index 7940b1169c54a1228a800151e9d3cdf398495807..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/__pycache__/pipeline_stable_diffusion_xl_k_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py deleted file mode 100644 index 350a49282693868ca93fa9a8b6a4829f1a919825..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py +++ /dev/null @@ -1,693 +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 - - - - 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 ge 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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py deleted file mode 100644 index 3b57555071f350339568fb1f99e15a2e3fecc449..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py +++ /dev/null @@ -1,903 +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.attention_processor import ( - AttnProcessor2_0, - FusedAttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers, LMSDiscreteScheduler -from ...utils import ( - USE_PEFT_BACKEND, - 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): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - FusedAttnProcessor2_0, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - @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 ge 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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_ldm3d/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_ldm3d/__init__.py deleted file mode 100644 index dae2affddd1fd5952f454ed9cee906277dcceb16..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_ldm3d/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -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_stable_diffusion_ldm3d"] = ["StableDiffusionLDM3DPipeline"] - -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_stable_diffusion_ldm3d import StableDiffusionLDM3DPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_ldm3d/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_ldm3d/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 5d220c5f3348b5b2c295788a0638d44e07925cdd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_ldm3d/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_ldm3d/__pycache__/pipeline_stable_diffusion_ldm3d.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_ldm3d/__pycache__/pipeline_stable_diffusion_ldm3d.cpython-310.pyc deleted file mode 100644 index e8412b4f54660acc9f5d37348e36b1a814bd913a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_ldm3d/__pycache__/pipeline_stable_diffusion_ldm3d.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py deleted file mode 100644 index c32121c88c9bb173f567311518db2ae05d3a1da6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py +++ /dev/null @@ -1,1040 +0,0 @@ -# Copyright 2025 The Intel Labs Team Authors 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, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...image_processor import PipelineImageInput, VaeImageProcessorLDM3D -from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - BaseOutput, - deprecate, - 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 DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -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 - >>> from diffusers import StableDiffusionLDM3DPipeline - - >>> pipe = StableDiffusionLDM3DPipeline.from_pretrained("Intel/ldm3d-4c") - >>> pipe = pipe.to("cuda") - - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> output = pipe(prompt) - >>> rgb_image, depth_image = output.rgb, output.depth - >>> rgb_image[0].save("astronaut_ldm3d_rgb.jpg") - >>> depth_image[0].save("astronaut_ldm3d_depth.png") - ``` -""" - - -# 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 - - -# 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 - - -@dataclass -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, - 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, - num_channels)`. - 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]] - - -class StableDiffusionLDM3DPipeline( - DeprecatedPipelineMixin, - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, - StableDiffusionLoraLoaderMixin, - FromSingleFileMixin, -): - _last_supported_version = "0.33.1" - - r""" - Pipeline for text-to-image and 3D generation using LDM3D. - - 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.). - - 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 - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: Optional[CLIPVisionModelWithProjection], - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - self.image_processor = VaeImageProcessorLDM3D(vae_scale_factor=self.vae_scale_factor) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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 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) - rgb_feature_extractor_input = feature_extractor_input[0] - safety_checker_input = self.feature_extractor(rgb_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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.check_inputs - def check_inputs( - self, - prompt, - height, - width, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_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}." - ) - - 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" - ) - - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_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, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 49, - 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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"], - **kwargs, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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. - 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*): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - 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). - 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. - 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. - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - - # 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, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - 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 - - 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, - self.do_classifier_free_guidance, - ) - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - self.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 self.do_classifier_free_guidance: - 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, sigmas - ) - - # 5. 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, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 6.1 Add image embeds for IP-Adapter - added_cond_kwargs = {"image_embeds": image_embeds} if ip_adapter_image is not None else None - - # 6.2 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) - - # 7. Denoising loop - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - if self.do_classifier_free_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - 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] - - rgb, depth = 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 ((rgb, depth), has_nsfw_concept) - - return LDM3DPipelineOutput(rgb=rgb, depth=depth, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_panorama/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_panorama/__init__.py deleted file mode 100644 index f7572db7236cd6bcfd7dd032abcb29fd5f67cf1c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_panorama/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -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_stable_diffusion_panorama"] = ["StableDiffusionPanoramaPipeline"] - -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_stable_diffusion_panorama import StableDiffusionPanoramaPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_panorama/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_panorama/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 23839fd61058d62f0de8bc5a9179bc44f22b2c95..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_panorama/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_panorama/__pycache__/pipeline_stable_diffusion_panorama.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_panorama/__pycache__/pipeline_stable_diffusion_panorama.cpython-310.pyc deleted file mode 100644 index 45be563f812fc162aa2fe5a9ef63ee918e7dcfae..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_panorama/__pycache__/pipeline_stable_diffusion_panorama.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py deleted file mode 100644 index 295095947a121243d5b415c1fbaa3d4b2d324ea2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py +++ /dev/null @@ -1,1195 +0,0 @@ -# Copyright 2025 MultiDiffusion Authors 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, Dict, List, Optional, Tuple, Union - -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import DDIMScheduler -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - 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 DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -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 StableDiffusionPanoramaPipeline, DDIMScheduler - - >>> model_ckpt = "stabilityai/stable-diffusion-2-base" - >>> scheduler = DDIMScheduler.from_pretrained(model_ckpt, subfolder="scheduler") - >>> pipe = StableDiffusionPanoramaPipeline.from_pretrained( - ... model_ckpt, scheduler=scheduler, torch_dtype=torch.float16 - ... ) - - >>> pipe = pipe.to("cuda") - - >>> prompt = "a photo of the dolomites" - >>> image = pipe(prompt).images[0] - ``` -""" - - -# 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 - - -# 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 StableDiffusionPanoramaPipeline( - DeprecatedPipelineMixin, - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - IPAdapterMixin, -): - _last_supported_version = "0.33.1" - - r""" - Pipeline for text-to-image generation using MultiDiffusion. - - 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.). - - 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 - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: DDIMScheduler, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.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 decode_latents_with_padding(self, latents: torch.Tensor, padding: int = 8) -> torch.Tensor: - """ - Decode the given latents with padding for circular inference. - - Args: - latents (torch.Tensor): The input latents to decode. - padding (int, optional): The number of latents to add on each side for padding. Defaults to 8. - - Returns: - torch.Tensor: The decoded image with padding removed. - - Notes: - - The padding is added to remove boundary artifacts and improve the output quality. - - This would slightly increase the memory usage. - - The padding pixels are then removed from the decoded image. - - """ - latents = 1 / self.vae.config.scaling_factor * latents - latents_left = latents[..., :padding] - latents_right = latents[..., -padding:] - latents = torch.cat((latents_right, latents, latents_left), axis=-1) - image = self.vae.decode(latents, return_dict=False)[0] - padding_pix = self.vae_scale_factor * padding - image = image[..., padding_pix:-padding_pix] - return image - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.check_inputs - def check_inputs( - self, - prompt, - height, - width, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_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}." - ) - - 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - def get_views( - self, - panorama_height: int, - panorama_width: int, - window_size: int = 64, - stride: int = 8, - circular_padding: bool = False, - ) -> 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, - num_blocks of height/width should return 1. - - Args: - panorama_height (int): The height of the panorama. - panorama_width (int): The width of the panorama. - window_size (int, optional): The size of the window. Defaults to 64. - stride (int, optional): The stride value. Defaults to 8. - 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 - representing the start and end coordinates of the window in the panorama. - - """ - panorama_height /= 8 - panorama_width /= 8 - num_blocks_height = (panorama_height - window_size) // stride + 1 if panorama_height > window_size else 1 - if circular_padding: - num_blocks_width = panorama_width // stride if panorama_width > window_size else 1 - else: - num_blocks_width = (panorama_width - window_size) // stride + 1 if panorama_width > window_size else 1 - total_num_blocks = int(num_blocks_height * num_blocks_width) - views = [] - for i in range(total_num_blocks): - h_start = int((i // num_blocks_width) * stride) - h_end = h_start + window_size - w_start = int((i % num_blocks_width) * stride) - w_end = w_start + window_size - views.append((h_start, h_end, w_start, w_end)) - return views - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def clip_skip(self): - return self._clip_skip - - @property - def do_classifier_free_guidance(self): - return False - - @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, - height: Optional[int] = 512, - width: Optional[int] = 2048, - num_inference_steps: int = 50, - 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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"], - **kwargs: Any, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - width (`int`, *optional*, defaults to 2048): - The width in pixels of the generated image. The width is kept high because the pipeline is supposed - generate panorama-like images. - 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*): - 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): - 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`. - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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)`. 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. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - cross_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). - guidance_rescale (`float`, *optional*, defaults to 0.0): - A rescaling factor for the guidance embeddings. A value of 0.0 means no rescaling is applied. - circular_padding (`bool`, *optional*, defaults to `False`): - If set to `True`, circular padding is applied to ensure there are no stitching artifacts. Circular - padding allows the model to seamlessly generate a transition from the rightmost part of the image to - the leftmost part, maintaining consistency in a 360-degree sense. - 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. - 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[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. - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", - ) - - # 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, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - 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 - # 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 = guidance_scale > 1.0 - - 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, - self.do_classifier_free_guidance, - ) - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - 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, - lora_scale=text_encoder_lora_scale, - 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 - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) - - # 5. 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, - ) - - # 6. Define panorama grid and initialize views for synthesis. - # prepare batch grid - views = self.get_views(height, width, circular_padding=circular_padding) - views_batch = [views[i : i + view_batch_size] for i in range(0, len(views), view_batch_size)] - views_scheduler_status = [copy.deepcopy(self.scheduler.__dict__)] * len(views_batch) - count = torch.zeros_like(latents) - value = torch.zeros_like(latents) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7.1 Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - # 7.2 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) - - # 8. Denoising loop - # Each denoising step also includes refinement of the latents with respect to the - # views. - 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 - count.zero_() - value.zero_() - - # generate views - # Here, we iterate through different spatial crops of the latents and denoise them. These - # denoised (latent) crops are then averaged to produce the final latent - # for the current timestep via MultiDiffusion. Please see Sec. 4.1 in the - # MultiDiffusion paper for more details: https://huggingface.co/papers/2302.08113 - # Batch views denoise - for j, batch_view in enumerate(views_batch): - vb_size = len(batch_view) - # get the latents corresponding to the current view coordinates - if circular_padding: - latents_for_view = [] - for h_start, h_end, w_start, w_end in batch_view: - if w_end > latents.shape[3]: - # Add circular horizontal padding - latent_view = torch.cat( - ( - latents[:, :, h_start:h_end, w_start:], - latents[:, :, h_start:h_end, : w_end - latents.shape[3]], - ), - axis=-1, - ) - else: - latent_view = latents[:, :, h_start:h_end, w_start:w_end] - latents_for_view.append(latent_view) - latents_for_view = torch.cat(latents_for_view) - else: - latents_for_view = torch.cat( - [ - latents[:, :, h_start:h_end, w_start:w_end] - for h_start, h_end, w_start, w_end in batch_view - ] - ) - - # rematch block's scheduler status - self.scheduler.__dict__.update(views_scheduler_status[j]) - - # expand the latents if we are doing classifier free guidance - latent_model_input = ( - latents_for_view.repeat_interleave(2, dim=0) - if do_classifier_free_guidance - else latents_for_view - ) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # repeat prompt_embeds for batch - prompt_embeds_input = torch.cat([prompt_embeds] * vb_size) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds_input, - timestep_cond=timestep_cond, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - ).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred[::2], noise_pred[1::2] - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - if self.do_classifier_free_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg( - noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale - ) - - # compute the previous noisy sample x_t -> x_t-1 - latents_denoised_batch = self.scheduler.step( - noise_pred, t, latents_for_view, **extra_step_kwargs - ).prev_sample - - # save views scheduler status after sample - views_scheduler_status[j] = copy.deepcopy(self.scheduler.__dict__) - - # extract value from batch - for latents_view_denoised, (h_start, h_end, w_start, w_end) in zip( - latents_denoised_batch.chunk(vb_size), batch_view - ): - if circular_padding and w_end > latents.shape[3]: - # Case for circular padding - value[:, :, h_start:h_end, w_start:] += latents_view_denoised[ - :, :, h_start:h_end, : latents.shape[3] - w_start - ] - value[:, :, h_start:h_end, : w_end - latents.shape[3]] += latents_view_denoised[ - :, :, h_start:h_end, latents.shape[3] - w_start : - ] - count[:, :, h_start:h_end, w_start:] += 1 - count[:, :, h_start:h_end, : w_end - latents.shape[3]] += 1 - else: - value[:, :, h_start:h_end, w_start:w_end] += latents_view_denoised - count[:, :, h_start:h_end, w_start:w_end] += 1 - - # take the MultiDiffusion step. Eq. 5 in MultiDiffusion paper: https://huggingface.co/papers/2302.08113 - latents = torch.where(count > 0, value / count, value) - - 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if output_type != "latent": - if circular_padding: - image = self.decode_latents_with_padding(latents) - else: - 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) - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__init__.py deleted file mode 100644 index b432b9418c46257913d81c5bf56edc0f1fa74ed1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__init__.py +++ /dev/null @@ -1,99 +0,0 @@ -from dataclasses import dataclass -from enum import Enum -from typing import TYPE_CHECKING, List, Optional, Union - -import numpy as np -import PIL -from PIL import Image - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - BaseOutput, - OptionalDependencyNotAvailable, - _LazyModule, - get_objects_from_module, - is_torch_available, - is_transformers_available, -) - - -@dataclass -class SafetyConfig(object): - WEAK = { - "sld_warmup_steps": 15, - "sld_guidance_scale": 20, - "sld_threshold": 0.0, - "sld_momentum_scale": 0.0, - "sld_mom_beta": 0.0, - } - MEDIUM = { - "sld_warmup_steps": 10, - "sld_guidance_scale": 1000, - "sld_threshold": 0.01, - "sld_momentum_scale": 0.3, - "sld_mom_beta": 0.4, - } - STRONG = { - "sld_warmup_steps": 7, - "sld_guidance_scale": 2000, - "sld_threshold": 0.025, - "sld_momentum_scale": 0.5, - "sld_mom_beta": 0.7, - } - MAX = { - "sld_warmup_steps": 0, - "sld_guidance_scale": 5000, - "sld_threshold": 1.0, - "sld_momentum_scale": 0.5, - "sld_mom_beta": 0.7, - } - - -_dummy_objects = {} -_additional_imports = {} -_import_structure = {} - -_additional_imports.update({"SafetyConfig": SafetyConfig}) - -try: - if not (is_transformers_available() and is_torch_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_torch_and_transformers_objects - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) -else: - _import_structure.update( - { - "pipeline_output": ["StableDiffusionSafePipelineOutput"], - "pipeline_stable_diffusion_safe": ["StableDiffusionPipelineSafe"], - "safety_checker": ["StableDiffusionSafetyChecker"], - } - ) - - -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_output import StableDiffusionSafePipelineOutput - from .pipeline_stable_diffusion_safe import StableDiffusionPipelineSafe - from .safety_checker import SafeStableDiffusionSafetyChecker - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 834fb9a945c4f08236b9f101cf099eeb0db6e8e1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index d1e98b19acf9ba0d781005e60f9b17300b0654d0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__pycache__/pipeline_stable_diffusion_safe.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__pycache__/pipeline_stable_diffusion_safe.cpython-310.pyc deleted file mode 100644 index 9ecf423fba723c69844682357c4a53d2bdeb7213..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__pycache__/pipeline_stable_diffusion_safe.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__pycache__/safety_checker.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__pycache__/safety_checker.cpython-310.pyc deleted file mode 100644 index 1951800804dabcaf476afc4fd9439f3acea29323..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/__pycache__/safety_checker.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py deleted file mode 100644 index 69a064d6638df556d3007f59daf7e767ec7c298b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py +++ /dev/null @@ -1,34 +0,0 @@ -from dataclasses import dataclass -from typing import List, Optional, Union - -import numpy as np -import PIL.Image - -from ...utils import ( - BaseOutput, -) - - -@dataclass -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, - 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, 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" - (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] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py deleted file mode 100644 index d334107b07036aaaed2960cbaca887fa1584d780..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py +++ /dev/null @@ -1,785 +0,0 @@ -import inspect -import warnings -from typing import Callable, List, Optional, Union - -import numpy as np -import torch -from packaging import version -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...configuration_utils import FrozenDict -from ...image_processor import PipelineImageInput -from ...loaders import IPAdapterMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import deprecate, is_torch_xla_available, logging -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from . import StableDiffusionSafePipelineOutput -from .safety_checker import SafeStableDiffusionSafetyChecker - - -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 - - -class StableDiffusionPipelineSafe(DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin, IPAdapterMixin): - _last_supported_version = "0.33.1" - - r""" - Pipeline based on the [`StableDiffusionPipeline`] for text-to-image generation using Safe Latent Diffusion. - - 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.). - - The pipeline also inherits the following loading methods: - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: SafeStableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, - requires_safety_checker: bool = True, - ): - super().__init__() - safety_concept: Optional[str] = ( - "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" - ) - - if scheduler is not None and getattr(scheduler.config, "steps_offset", 1) != 1: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" - f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " - "to update the config accordingly as leaving `steps_offset` might led to incorrect results" - " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," - " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" - " file" - ) - deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["steps_offset"] = 1 - scheduler._internal_dict = FrozenDict(new_config) - - if scheduler is not None and getattr(scheduler.config, "clip_sample", False) is True: - deprecation_message = ( - f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." - " `clip_sample` should be set to False in the configuration file. Please make sure to update the" - " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" - " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" - " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" - ) - deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(scheduler.config) - new_config["clip_sample"] = False - scheduler._internal_dict = FrozenDict(new_config) - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - is_unet_version_less_0_9_0 = ( - unet is not None - and hasattr(unet.config, "_diffusers_version") - and version.parse(version.parse(unet.config._diffusers_version).base_version) < version.parse("0.9.0.dev0") - ) - is_unet_sample_size_less_64 = ( - unet is not None and hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 - ) - if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: - deprecation_message = ( - "The configuration file of the unet has set the default `sample_size` to smaller than" - " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the" - " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" - " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- stable-diffusion-v1-5/stable-diffusion-v1-5" - " \n- stable-diffusion-v1-5/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" - " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" - " in the config might lead to incorrect results in future versions. If you have downloaded this" - " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" - " the `unet/config.json` file" - ) - deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) - new_config = dict(unet.config) - new_config["sample_size"] = 64 - unet._internal_dict = FrozenDict(new_config) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - self._safety_text_concept = safety_concept - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - self.register_to_config(requires_safety_checker=requires_safety_checker) - - @property - def safety_concept(self): - r""" - Getter method for the safety concept used with SLD - - Returns: - `str`: The text describing the safety concept - """ - return self._safety_text_concept - - @safety_concept.setter - def safety_concept(self, concept): - r""" - Setter method for the safety concept used with SLD - - Args: - concept (`str`): - The text of the new safety concept - """ - self._safety_text_concept = concept - - def _encode_prompt( - self, - prompt, - device, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt, - enable_safety_guidance, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`): - 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]`): - 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`). - """ - batch_size = len(prompt) if isinstance(prompt, list) else 1 - - 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="max_length", return_tensors="pt").input_ids - - if 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 - - prompt_embeds = self.text_encoder( - text_input_ids.to(device), - attention_mask=attention_mask, - ) - prompt_embeds = prompt_embeds[0] - - # duplicate text embeddings for each generation per prompt, using mps friendly method - bs_embed, seq_len, _ = prompt_embeds.shape - 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: - uncond_tokens: List[str] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif 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 - - max_length = text_input_ids.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] - - # 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.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - - # Encode the safety concept text - if enable_safety_guidance: - safety_concept_input = self.tokenizer( - [self._safety_text_concept], - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - safety_embeddings = self.text_encoder(safety_concept_input.input_ids.to(self.device))[0] - - # duplicate safety embeddings for each generation per prompt, using mps friendly method - seq_len = safety_embeddings.shape[1] - safety_embeddings = safety_embeddings.repeat(batch_size, num_images_per_prompt, 1) - safety_embeddings = safety_embeddings.view(batch_size * num_images_per_prompt, seq_len, -1) - - # For classifier free guidance + sld, we need to do three forward passes. - # Here we concatenate the unconditional and text embeddings into a single batch - # to avoid doing three forward passes - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds, safety_embeddings]) - - else: - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - return prompt_embeds - - def run_safety_checker(self, image, device, dtype, enable_safety_guidance): - if self.safety_checker is not None: - images = image.copy() - safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device) - image, has_nsfw_concept = self.safety_checker( - images=image, clip_input=safety_checker_input.pixel_values.to(dtype) - ) - flagged_images = np.zeros((2, *image.shape[1:])) - if any(has_nsfw_concept): - logger.warning( - "Potential NSFW content was detected in one or more images. A black image will be returned" - " instead." - f"{'You may look at this images in the `unsafe_images` variable of the output at your own discretion.' if enable_safety_guidance else 'Try again with a different prompt and/or seed.'}" - ) - for idx, has_nsfw_concept in enumerate(has_nsfw_concept): - if has_nsfw_concept: - flagged_images[idx] = images[idx] - image[idx] = np.zeros(image[idx].shape) # black image - else: - has_nsfw_concept = None - flagged_images = None - return image, has_nsfw_concept, flagged_images - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - 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, - 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}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def perform_safety_guidance( - self, - enable_safety_guidance, - safety_momentum, - noise_guidance, - noise_pred_out, - i, - sld_guidance_scale, - sld_warmup_steps, - sld_threshold, - sld_momentum_scale, - sld_mom_beta, - ): - # Perform SLD guidance - if enable_safety_guidance: - if safety_momentum is None: - safety_momentum = torch.zeros_like(noise_guidance) - noise_pred_text, noise_pred_uncond = noise_pred_out[0], noise_pred_out[1] - noise_pred_safety_concept = noise_pred_out[2] - - # Equation 6 - scale = torch.clamp(torch.abs((noise_pred_text - noise_pred_safety_concept)) * sld_guidance_scale, max=1.0) - - # Equation 6 - safety_concept_scale = torch.where( - (noise_pred_text - noise_pred_safety_concept) >= sld_threshold, torch.zeros_like(scale), scale - ) - - # Equation 4 - noise_guidance_safety = torch.mul((noise_pred_safety_concept - noise_pred_uncond), safety_concept_scale) - - # Equation 7 - noise_guidance_safety = noise_guidance_safety + sld_momentum_scale * safety_momentum - - # Equation 8 - safety_momentum = sld_mom_beta * safety_momentum + (1 - sld_mom_beta) * noise_guidance_safety - - if i >= sld_warmup_steps: # Warmup - # Equation 3 - noise_guidance = noise_guidance - noise_guidance_safety - return noise_guidance, safety_momentum - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - @torch.no_grad() - def __call__( - self, - prompt: Union[str, List[str]], - 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, - ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], 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, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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): - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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`. - ip_adapter_image: (`PipelineImageInput`, *optional*): - Optional image input to work with IP Adapters. - 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - sld_guidance_scale (`float`, *optional*, defaults to 1000): - If `sld_guidance_scale < 1`, safety guidance is disabled. - sld_warmup_steps (`int`, *optional*, defaults to 10): - Number of warmup steps for safety guidance. SLD is only be applied for diffusion steps greater than - `sld_warmup_steps`. - sld_threshold (`float`, *optional*, defaults to 0.01): - Threshold that separates the hyperplane between appropriate and inappropriate images. - sld_momentum_scale (`float`, *optional*, defaults to 0.3): - Scale of the SLD momentum to be added to the safety guidance at each diffusion step. If set to 0.0, - momentum is disabled. Momentum is built up during warmup for diffusion steps smaller than - `sld_warmup_steps`. - sld_mom_beta (`float`, *optional*, defaults to 0.4): - Defines how safety guidance momentum builds up. `sld_mom_beta` indicates how much of the previous - momentum is kept. Momentum is built up during warmup for diffusion steps smaller than - `sld_warmup_steps`. - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - - Examples: - - ```py - import torch - from diffusers import StableDiffusionPipelineSafe - from diffusers.pipelines.stable_diffusion_safe import SafetyConfig - - pipeline = StableDiffusionPipelineSafe.from_pretrained( - "AIML-TUDA/stable-diffusion-safe", torch_dtype=torch.float16 - ).to("cuda") - prompt = "the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c. leyendecker" - image = pipeline(prompt=prompt, **SafetyConfig.MEDIUM).images[0] - ``` - """ - # 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) - - # 2. Define call parameters - batch_size = 1 if isinstance(prompt, str) else len(prompt) - 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 = guidance_scale > 1.0 - - enable_safety_guidance = sld_guidance_scale > 1.0 and do_classifier_free_guidance - if not enable_safety_guidance: - warnings.warn("Safety checker disabled!") - - if ip_adapter_image is not None: - output_hidden_state = False if isinstance(self.unet.encoder_hid_proj, ImageProjection) else True - image_embeds, negative_image_embeds = self.encode_image( - ip_adapter_image, device, num_images_per_prompt, output_hidden_state - ) - if do_classifier_free_guidance: - if enable_safety_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds, image_embeds]) - else: - image_embeds = torch.cat([negative_image_embeds, image_embeds]) - - # 3. Encode input prompt - prompt_embeds = self._encode_prompt( - prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt, enable_safety_guidance - ) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. 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, - ) - - # 6. Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 6.1 Add image embeds for IP-Adapter - added_cond_kwargs = {"image_embeds": image_embeds} if ip_adapter_image is not None else None - - safety_momentum = None - - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = ( - torch.cat([latents] * (3 if enable_safety_guidance else 2)) - if do_classifier_free_guidance - else latents - ) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, t, encoder_hidden_states=prompt_embeds, added_cond_kwargs=added_cond_kwargs - ).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_out = noise_pred.chunk((3 if enable_safety_guidance else 2)) - noise_pred_uncond, noise_pred_text = noise_pred_out[0], noise_pred_out[1] - - # default classifier free guidance - noise_guidance = noise_pred_text - noise_pred_uncond - - # Perform SLD guidance - if enable_safety_guidance: - if safety_momentum is None: - safety_momentum = torch.zeros_like(noise_guidance) - noise_pred_safety_concept = noise_pred_out[2] - - # Equation 6 - scale = torch.clamp( - torch.abs((noise_pred_text - noise_pred_safety_concept)) * sld_guidance_scale, max=1.0 - ) - - # Equation 6 - safety_concept_scale = torch.where( - (noise_pred_text - noise_pred_safety_concept) >= sld_threshold, - torch.zeros_like(scale), - scale, - ) - - # Equation 4 - noise_guidance_safety = torch.mul( - (noise_pred_safety_concept - noise_pred_uncond), safety_concept_scale - ) - - # Equation 7 - noise_guidance_safety = noise_guidance_safety + sld_momentum_scale * safety_momentum - - # Equation 8 - safety_momentum = sld_mom_beta * safety_momentum + (1 - sld_mom_beta) * noise_guidance_safety - - if i >= sld_warmup_steps: # Warmup - # Equation 3 - noise_guidance = noise_guidance - noise_guidance_safety - - noise_pred = noise_pred_uncond + guidance_scale * noise_guidance - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # 8. Post-processing - image = self.decode_latents(latents) - - # 9. Run safety checker - image, has_nsfw_concept, flagged_images = self.run_safety_checker( - image, device, prompt_embeds.dtype, enable_safety_guidance - ) - - # 10. Convert to PIL - if output_type == "pil": - image = self.numpy_to_pil(image) - if flagged_images is not None: - flagged_images = self.numpy_to_pil(flagged_images) - - if not return_dict: - return ( - image, - has_nsfw_concept, - self._safety_text_concept if enable_safety_guidance else None, - flagged_images, - ) - - return StableDiffusionSafePipelineOutput( - images=image, - nsfw_content_detected=has_nsfw_concept, - applied_safety_concept=self._safety_text_concept if enable_safety_guidance else None, - unsafe_images=flagged_images, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/safety_checker.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/safety_checker.py deleted file mode 100644 index 1f6ad5f2a348c3a0aed208011e40e507fe0195e5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_safe/safety_checker.py +++ /dev/null @@ -1,109 +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 torch -import torch.nn as nn -from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel - -from ...utils import logging - - -logger = logging.get_logger(__name__) - - -def cosine_distance(image_embeds, text_embeds): - normalized_image_embeds = nn.functional.normalize(image_embeds) - normalized_text_embeds = nn.functional.normalize(text_embeds) - return torch.mm(normalized_image_embeds, normalized_text_embeds.t()) - - -class SafeStableDiffusionSafetyChecker(PreTrainedModel): - config_class = CLIPConfig - - _no_split_modules = ["CLIPEncoderLayer"] - - def __init__(self, config: CLIPConfig): - super().__init__(config) - - self.vision_model = CLIPVisionModel(config.vision_config) - self.visual_projection = nn.Linear(config.vision_config.hidden_size, config.projection_dim, bias=False) - - self.concept_embeds = nn.Parameter(torch.ones(17, config.projection_dim), requires_grad=False) - self.special_care_embeds = nn.Parameter(torch.ones(3, config.projection_dim), requires_grad=False) - - self.concept_embeds_weights = nn.Parameter(torch.ones(17), requires_grad=False) - self.special_care_embeds_weights = nn.Parameter(torch.ones(3), requires_grad=False) - - @torch.no_grad() - def forward(self, clip_input, images): - pooled_output = self.vision_model(clip_input)[1] # pooled_output - image_embeds = self.visual_projection(pooled_output) - - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds).cpu().float().numpy() - cos_dist = cosine_distance(image_embeds, self.concept_embeds).cpu().float().numpy() - - result = [] - batch_size = image_embeds.shape[0] - for i in range(batch_size): - result_img = {"special_scores": {}, "special_care": [], "concept_scores": {}, "bad_concepts": []} - - # increase this value to create a stronger `nfsw` filter - # at the cost of increasing the possibility of filtering benign images - adjustment = 0.0 - - for concept_idx in range(len(special_cos_dist[0])): - concept_cos = special_cos_dist[i][concept_idx] - concept_threshold = self.special_care_embeds_weights[concept_idx].item() - result_img["special_scores"][concept_idx] = round(concept_cos - concept_threshold + adjustment, 3) - if result_img["special_scores"][concept_idx] > 0: - result_img["special_care"].append({concept_idx, result_img["special_scores"][concept_idx]}) - adjustment = 0.01 - - for concept_idx in range(len(cos_dist[0])): - concept_cos = cos_dist[i][concept_idx] - concept_threshold = self.concept_embeds_weights[concept_idx].item() - result_img["concept_scores"][concept_idx] = round(concept_cos - concept_threshold + adjustment, 3) - if result_img["concept_scores"][concept_idx] > 0: - result_img["bad_concepts"].append(concept_idx) - - result.append(result_img) - - has_nsfw_concepts = [len(res["bad_concepts"]) > 0 for res in result] - - return images, has_nsfw_concepts - - @torch.no_grad() - def forward_onnx(self, clip_input: torch.Tensor, images: torch.Tensor): - pooled_output = self.vision_model(clip_input)[1] # pooled_output - image_embeds = self.visual_projection(pooled_output) - - special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds) - cos_dist = cosine_distance(image_embeds, self.concept_embeds) - - # increase this value to create a stronger `nsfw` filter - # at the cost of increasing the possibility of filtering benign images - adjustment = 0.0 - - special_scores = special_cos_dist - self.special_care_embeds_weights + adjustment - # special_scores = special_scores.round(decimals=3) - special_care = torch.any(special_scores > 0, dim=1) - special_adjustment = special_care * 0.01 - special_adjustment = special_adjustment.unsqueeze(1).expand(-1, cos_dist.shape[1]) - - concept_scores = (cos_dist - self.concept_embeds_weights) + special_adjustment - # concept_scores = concept_scores.round(decimals=3) - has_nsfw_concepts = torch.any(concept_scores > 0, dim=1) - - return images, has_nsfw_concepts diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_sag/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_sag/__init__.py deleted file mode 100644 index 378e0e57817f58a0a28afed5d6110f6ee3effb3a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_sag/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -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_stable_diffusion_sag"] = ["StableDiffusionSAGPipeline"] - -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_stable_diffusion_sag import StableDiffusionSAGPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_sag/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_sag/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index c89d88646427dd4b9f80ee5015d1aef2bc0f397a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_sag/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_sag/__pycache__/pipeline_stable_diffusion_sag.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_sag/__pycache__/pipeline_stable_diffusion_sag.cpython-310.pyc deleted file mode 100644 index bbfd1dbce53ff467f8a27cc2ae291bdb3c72b3a6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_sag/__pycache__/pipeline_stable_diffusion_sag.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py deleted file mode 100644 index 48add535a81da083f65b2fb9843ee72bf5577819..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py +++ /dev/null @@ -1,970 +0,0 @@ -# Copyright 2025 Susung Hong 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 -import torch.nn.functional as F -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - 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 DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion import StableDiffusionPipelineOutput -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -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 StableDiffusionSAGPipeline - - >>> pipe = StableDiffusionSAGPipeline.from_pretrained( - ... "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> image = pipe(prompt, sag_scale=0.75).images[0] - ``` -""" - - -# processes and stores attention probabilities -class CrossAttnStoreProcessor: - def __init__(self): - self.attention_probs = None - - def __call__( - self, - attn, - hidden_states, - encoder_hidden_states=None, - attention_mask=None, - ): - batch_size, sequence_length, _ = hidden_states.shape - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query = attn.head_to_batch_dim(query) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - - self.attention_probs = attn.get_attention_scores(query, key, attention_mask) - hidden_states = torch.bmm(self.attention_probs, value) - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - return hidden_states - - -# Modified to get self-attention guidance scale in this paper (https://huggingface.co/papers/2210.00939) as an input -class StableDiffusionSAGPipeline( - DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin, TextualInversionLoaderMixin, IPAdapterMixin -): - _last_supported_version = "0.33.1" - - r""" - Pipeline for text-to-image generation using Stable Diffusion. - - 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.). - - The pipeline also inherits the following loading methods: - - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings - - [`~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 ([`~transformers.CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - unet ([`UNet2DConditionModel`]): - A `UNet2DConditionModel` 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 - more details about a model's potential harms. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] - _exclude_from_cpu_offload = ["safety_checker"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, - requires_safety_checker: bool = True, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - 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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - image_embeds = [] - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - single_image_embeds = torch.stack([single_image_embeds] * num_images_per_prompt, dim=0) - single_negative_image_embeds = torch.stack( - [single_negative_image_embeds] * num_images_per_prompt, dim=0 - ) - - if do_classifier_free_guidance: - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds]) - single_image_embeds = single_image_embeds.to(device) - - image_embeds.append(single_image_embeds) - else: - image_embeds = ip_adapter_image_embeds - return image_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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - 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, - 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}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @torch.no_grad() - @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: 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, - 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", - 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, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - 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): - 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`. - sag_scale (`float`, *optional*, defaults to 0.75): - Chosen between [0, 1.0] for better quality. - 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): - The number of images to generate per prompt. - 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*): - 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. - 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. - 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. If not provided, embeddings are computed from the - `ip_adapter_image` 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 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - 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. - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] 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. - """ - # 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 = guidance_scale > 1.0 - # and `sag_scale` is` `s` of equation (16) - # of the self-attention guidance paper: https://huggingface.co/papers/2210.00939 - # `sag_scale = 0` means no self-attention guidance - do_self_attention_guidance = sag_scale > 0.0 - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds( - ip_adapter_image, - ip_adapter_image_embeds, - device, - batch_size * num_images_per_prompt, - do_classifier_free_guidance, - ) - - if do_classifier_free_guidance: - image_embeds = [] - negative_image_embeds = [] - for tmp_image_embeds in ip_adapter_image_embeds: - single_negative_image_embeds, single_image_embeds = tmp_image_embeds.chunk(2) - image_embeds.append(single_image_embeds) - negative_image_embeds.append(single_negative_image_embeds) - else: - image_embeds = ip_adapter_image_embeds - - # 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=device) - timesteps = self.scheduler.timesteps - - if timesteps.dtype not in [torch.int16, torch.int32, torch.int64]: - raise ValueError( - f"{self.__class__.__name__} does not support using a scheduler of type {self.scheduler.__class__.__name__}. Please make sure to use one of 'DDIMScheduler, PNDMScheduler, DDPMScheduler, DEISMultistepScheduler, UniPCMultistepScheduler, DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler'." - ) - - # 5. 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, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 6.1 Add image embeds for IP-Adapter - added_cond_kwargs = ( - {"image_embeds": image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - if do_classifier_free_guidance: - added_uncond_kwargs = ( - {"image_embeds": negative_image_embeds} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None - else None - ) - - # 7. Denoising loop - original_attn_proc = self.unet.attn_processors - store_processor = CrossAttnStoreProcessor() - self.unet.mid_block.attentions[0].transformer_blocks[0].attn1.processor = store_processor - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - map_size = None - - def get_map_size(module, input, output): - nonlocal map_size - map_size = output[0].shape[-2:] - - with self.unet.mid_block.attentions[0].register_forward_hook(get_map_size): - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - ).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # perform self-attention guidance with the stored self-attention map - if do_self_attention_guidance: - # classifier-free guidance produces two chunks of attention map - # and we only use unconditional one according to equation (25) - # in https://huggingface.co/papers/2210.00939 - if do_classifier_free_guidance: - # DDIM-like prediction of x0 - pred_x0 = self.pred_x0(latents, noise_pred_uncond, t) - # get the stored attention maps - uncond_attn, cond_attn = store_processor.attention_probs.chunk(2) - # self-attention-based degrading of latents - degraded_latents = self.sag_masking( - pred_x0, uncond_attn, map_size, t, self.pred_epsilon(latents, noise_pred_uncond, t) - ) - uncond_emb, _ = prompt_embeds.chunk(2) - # forward and give guidance - degraded_pred = self.unet( - degraded_latents, - t, - encoder_hidden_states=uncond_emb, - added_cond_kwargs=added_uncond_kwargs, - ).sample - noise_pred += sag_scale * (noise_pred_uncond - degraded_pred) - else: - # DDIM-like prediction of x0 - pred_x0 = self.pred_x0(latents, noise_pred, t) - # get the stored attention maps - cond_attn = store_processor.attention_probs - # self-attention-based degrading of latents - degraded_latents = self.sag_masking( - pred_x0, cond_attn, map_size, t, self.pred_epsilon(latents, noise_pred, t) - ) - # forward and give guidance - degraded_pred = self.unet( - degraded_latents, - t, - encoder_hidden_states=prompt_embeds, - added_cond_kwargs=added_cond_kwargs, - ).sample - noise_pred += sag_scale * (noise_pred - degraded_pred) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - 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) - - self.maybe_free_model_hooks() - # make sure to set the original attention processors back - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) - - def sag_masking(self, original_latents, attn_map, map_size, t, eps): - # Same masking process as in SAG paper: https://huggingface.co/papers/2210.00939 - bh, hw1, hw2 = attn_map.shape - b, latent_channel, latent_h, latent_w = original_latents.shape - h = self.unet.config.attention_head_dim - if isinstance(h, list): - h = h[-1] - - # Produce attention mask - attn_map = attn_map.reshape(b, h, hw1, hw2) - attn_mask = attn_map.mean(1, keepdim=False).sum(1, keepdim=False) > 1.0 - attn_mask = ( - attn_mask.reshape(b, map_size[0], map_size[1]) - .unsqueeze(1) - .repeat(1, latent_channel, 1, 1) - .type(attn_map.dtype) - ) - attn_mask = F.interpolate(attn_mask, (latent_h, latent_w)) - - # Blur according to the self-attention mask - degraded_latents = gaussian_blur_2d(original_latents, kernel_size=9, sigma=1.0) - degraded_latents = degraded_latents * attn_mask + original_latents * (1 - attn_mask) - - # Noise it again to match the noise level - degraded_latents = self.scheduler.add_noise(degraded_latents, noise=eps, timesteps=t[None]) - - return degraded_latents - - # Modified from diffusers.schedulers.scheduling_ddim.DDIMScheduler.step - # Note: there are some schedulers that clip or do not return x_0 (PNDMScheduler, DDIMScheduler, etc.) - def pred_x0(self, sample, model_output, timestep): - alpha_prod_t = self.scheduler.alphas_cumprod[timestep].to(sample.device) - - beta_prod_t = 1 - alpha_prod_t - if self.scheduler.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - elif self.scheduler.config.prediction_type == "sample": - pred_original_sample = model_output - elif self.scheduler.config.prediction_type == "v_prediction": - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - # predict V - model_output = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - else: - raise ValueError( - f"prediction_type given as {self.scheduler.config.prediction_type} must be one of `epsilon`, `sample`," - " or `v_prediction`" - ) - - return pred_original_sample - - def pred_epsilon(self, sample, model_output, timestep): - alpha_prod_t = self.scheduler.alphas_cumprod[timestep] - - beta_prod_t = 1 - alpha_prod_t - if self.scheduler.config.prediction_type == "epsilon": - pred_eps = model_output - elif self.scheduler.config.prediction_type == "sample": - pred_eps = (sample - (alpha_prod_t**0.5) * model_output) / (beta_prod_t**0.5) - elif self.scheduler.config.prediction_type == "v_prediction": - pred_eps = (beta_prod_t**0.5) * sample + (alpha_prod_t**0.5) * model_output - else: - raise ValueError( - f"prediction_type given as {self.scheduler.config.prediction_type} must be one of `epsilon`, `sample`," - " or `v_prediction`" - ) - - return pred_eps - - -# Gaussian blur -def gaussian_blur_2d(img, kernel_size, sigma): - ksize_half = (kernel_size - 1) * 0.5 - - x = torch.linspace(-ksize_half, ksize_half, steps=kernel_size) - - pdf = torch.exp(-0.5 * (x / sigma).pow(2)) - - x_kernel = pdf / pdf.sum() - x_kernel = x_kernel.to(device=img.device, dtype=img.dtype) - - kernel2d = torch.mm(x_kernel[:, None], x_kernel[None, :]) - kernel2d = kernel2d.expand(img.shape[-3], 1, kernel2d.shape[0], kernel2d.shape[1]) - - padding = [kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2] - - img = F.pad(img, padding, mode="reflect") - img = F.conv2d(img, kernel2d, groups=img.shape[-3]) - - return img diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__init__.py deleted file mode 100644 index 8088fbcfceba205b9b908613f4ca3fdc579120e8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__init__.py +++ /dev/null @@ -1,76 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - get_objects_from_module, - is_flax_available, - is_torch_available, - is_transformers_available, -) - - -_dummy_objects = {} -_additional_imports = {} -_import_structure = {"pipeline_output": ["StableDiffusionXLPipelineOutput"]} - -if is_transformers_available() and is_flax_available(): - _import_structure["pipeline_output"].extend(["FlaxStableDiffusionXLPipelineOutput"]) -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_stable_diffusion_xl"] = ["StableDiffusionXLPipeline"] - _import_structure["pipeline_stable_diffusion_xl_img2img"] = ["StableDiffusionXLImg2ImgPipeline"] - _import_structure["pipeline_stable_diffusion_xl_inpaint"] = ["StableDiffusionXLInpaintPipeline"] - _import_structure["pipeline_stable_diffusion_xl_instruct_pix2pix"] = ["StableDiffusionXLInstructPix2PixPipeline"] - -if is_transformers_available() and is_flax_available(): - from ...schedulers.scheduling_pndm_flax import PNDMSchedulerState - - _additional_imports.update({"PNDMSchedulerState": PNDMSchedulerState}) - _import_structure["pipeline_flax_stable_diffusion_xl"] = ["FlaxStableDiffusionXLPipeline"] - - -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_stable_diffusion_xl import StableDiffusionXLPipeline - from .pipeline_stable_diffusion_xl_img2img import StableDiffusionXLImg2ImgPipeline - from .pipeline_stable_diffusion_xl_inpaint import StableDiffusionXLInpaintPipeline - from .pipeline_stable_diffusion_xl_instruct_pix2pix import StableDiffusionXLInstructPix2PixPipeline - - try: - if not (is_transformers_available() and is_flax_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_flax_objects import * - else: - from .pipeline_flax_stable_diffusion_xl import ( - FlaxStableDiffusionXLPipeline, - ) - from .pipeline_output import FlaxStableDiffusionXLPipelineOutput - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index b6fcf27f9e718da6f1a5b1bfbaf130cb8357dca3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_flax_stable_diffusion_xl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_flax_stable_diffusion_xl.cpython-310.pyc deleted file mode 100644 index f455f6f4f2d10a11d48c4dc9c2c22bda9825502b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_flax_stable_diffusion_xl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index e5dc8857d8eb0c9ea9884a3e121d52a64a7065e0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl.cpython-310.pyc deleted file mode 100644 index 6d105cc8677fbc82f01dc0999da852aad223e944..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl_img2img.cpython-310.pyc deleted file mode 100644 index a4877c1d526a993ac01b36e6015637577cd9e56d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl_inpaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl_inpaint.cpython-310.pyc deleted file mode 100644 index f9da4ec6e4db2516677f5ddc97ff06f8cf5e0ea2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl_inpaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl_instruct_pix2pix.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl_instruct_pix2pix.cpython-310.pyc deleted file mode 100644 index c9d5074192e2f129802e04718f4f8f0a80d89256..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/pipeline_stable_diffusion_xl_instruct_pix2pix.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/watermark.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/watermark.cpython-310.pyc deleted file mode 100644 index 5dd201b2ed520d291c10bb8226cc59925d442720..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/__pycache__/watermark.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py deleted file mode 100644 index 3227fd9a08a47306eeb0804f301d16cdd1f4ce5f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py +++ /dev/null @@ -1,308 +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 functools import partial -from typing import Dict, List, Optional, Union - -import jax -import jax.numpy as jnp -from flax.core.frozen_dict import FrozenDict -from transformers import CLIPTokenizer, FlaxCLIPTextModel - -from diffusers.utils import logging - -from ...models import FlaxAutoencoderKL, FlaxUNet2DConditionModel -from ...schedulers import ( - FlaxDDIMScheduler, - FlaxDPMSolverMultistepScheduler, - FlaxLMSDiscreteScheduler, - FlaxPNDMScheduler, -) -from ..pipeline_flax_utils import FlaxDiffusionPipeline -from .pipeline_output import FlaxStableDiffusionXLPipelineOutput - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -# Set to True to use python for loop instead of jax.fori_loop for easier debugging -DEBUG = False - - -class FlaxStableDiffusionXLPipeline(FlaxDiffusionPipeline): - def __init__( - self, - text_encoder: FlaxCLIPTextModel, - text_encoder_2: FlaxCLIPTextModel, - vae: FlaxAutoencoderKL, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: FlaxUNet2DConditionModel, - scheduler: Union[ - FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler - ], - dtype: jnp.dtype = jnp.float32, - ): - super().__init__() - self.dtype = dtype - - 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.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]]): - if not isinstance(prompt, (str, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - # Assume we have the two encoders - inputs = [] - for tokenizer in [self.tokenizer, self.tokenizer_2]: - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="np", - ) - inputs.append(text_inputs.input_ids) - inputs = jnp.stack(inputs, axis=1) - return inputs - - def __call__( - self, - prompt_ids: jax.Array, - params: Union[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, - latents: jnp.array = None, - neg_prompt_ids: jnp.array = None, - return_dict: bool = True, - output_type: str = None, - jit: bool = False, - ): - # 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 - - if isinstance(guidance_scale, float) and jit: - # Convert to a tensor so each device gets a copy. - guidance_scale = jnp.array([guidance_scale] * prompt_ids.shape[0]) - guidance_scale = guidance_scale[:, None] - - return_latents = output_type == "latent" - - if jit: - images = _p_generate( - self, - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - return_latents, - ) - else: - images = self._generate( - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - return_latents, - ) - - if not return_dict: - return (images,) - - return FlaxStableDiffusionXLPipelineOutput(images=images) - - def get_embeddings(self, prompt_ids: jnp.array, params): - # We assume we have the two encoders - - # bs, encoder_input, seq_length - te_1_inputs = prompt_ids[:, 0, :] - te_2_inputs = prompt_ids[:, 1, :] - - prompt_embeds = self.text_encoder(te_1_inputs, params=params["text_encoder"], output_hidden_states=True) - prompt_embeds = prompt_embeds["hidden_states"][-2] - prompt_embeds_2_out = self.text_encoder_2( - te_2_inputs, params=params["text_encoder_2"], output_hidden_states=True - ) - prompt_embeds_2 = prompt_embeds_2_out["hidden_states"][-2] - text_embeds = prompt_embeds_2_out["text_embeds"] - prompt_embeds = jnp.concatenate([prompt_embeds, prompt_embeds_2], axis=-1) - return prompt_embeds, text_embeds - - def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, bs, dtype): - add_time_ids = list(original_size + crops_coords_top_left + target_size) - add_time_ids = jnp.array([add_time_ids] * bs, dtype=dtype) - return add_time_ids - - def _generate( - self, - prompt_ids: jnp.array, - params: Union[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, - return_latents=False, - ): - 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}.") - - # Encode input prompt - prompt_embeds, pooled_embeds = self.get_embeddings(prompt_ids, params) - - # Get unconditional embeddings - batch_size = prompt_embeds.shape[0] - if neg_prompt_ids is None: - neg_prompt_embeds = jnp.zeros_like(prompt_embeds) - negative_pooled_embeds = jnp.zeros_like(pooled_embeds) - else: - neg_prompt_embeds, negative_pooled_embeds = self.get_embeddings(neg_prompt_ids, params) - - add_time_ids = self._get_add_time_ids( - (height, width), (0, 0), (height, width), prompt_embeds.shape[0], dtype=prompt_embeds.dtype - ) - - prompt_embeds = jnp.concatenate([neg_prompt_embeds, prompt_embeds], axis=0) # (2, 77, 2048) - add_text_embeds = jnp.concatenate([negative_pooled_embeds, pooled_embeds], axis=0) - add_time_ids = jnp.concatenate([add_time_ids, add_time_ids], axis=0) - - # Ensure model output will be `float32` before going into the scheduler - guidance_scale = jnp.array([guidance_scale], dtype=jnp.float32) - - # Create random latents - latents_shape = ( - batch_size, - self.unet.config.in_channels, - height // self.vae_scale_factor, - width // self.vae_scale_factor, - ) - if latents is None: - latents = jax.random.normal(prng_seed, shape=latents_shape, dtype=jnp.float32) - else: - if latents.shape != latents_shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - - # Prepare scheduler state - scheduler_state = self.scheduler.set_timesteps( - params["scheduler"], num_inference_steps=num_inference_steps, shape=latents.shape - ) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * scheduler_state.init_noise_sigma - - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - - # Denoising loop - def loop_body(step, args): - latents, scheduler_state = args - # 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 - latents_input = jnp.concatenate([latents] * 2) - - t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step] - timestep = jnp.broadcast_to(t, latents_input.shape[0]) - - latents_input = self.scheduler.scale_model_input(scheduler_state, latents_input, t) - - # predict the noise residual - noise_pred = self.unet.apply( - {"params": params["unet"]}, - jnp.array(latents_input), - jnp.array(timestep, dtype=jnp.int32), - encoder_hidden_states=prompt_embeds, - added_cond_kwargs=added_cond_kwargs, - ).sample - # perform guidance - noise_pred_uncond, noise_prediction_text = jnp.split(noise_pred, 2, axis=0) - noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents, scheduler_state = self.scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() - return latents, scheduler_state - - if DEBUG: - # run with python for loop - for i in range(num_inference_steps): - latents, scheduler_state = loop_body(i, (latents, scheduler_state)) - else: - latents, _ = jax.lax.fori_loop(0, num_inference_steps, loop_body, (latents, scheduler_state)) - - if return_latents: - return latents - - # Decode latents - latents = 1 / self.vae.config.scaling_factor * latents - image = self.vae.apply({"params": params["vae"]}, latents, method=self.vae.decode).sample - - image = (image / 2 + 0.5).clip(0, 1).transpose(0, 2, 3, 1) - return image - - -# Static argnums are pipe, num_inference_steps, height, width, return_latents. A change would trigger recompilation. -# Non-static args are (sharded) input tensors mapped over their first dimension (hence, `0`). -@partial( - jax.pmap, - in_axes=(None, 0, 0, 0, None, None, None, 0, 0, 0, None), - static_broadcasted_argnums=(0, 4, 5, 6, 10), -) -def _p_generate( - pipe, - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - return_latents, -): - return pipe._generate( - prompt_ids, - params, - prng_seed, - num_inference_steps, - height, - width, - guidance_scale, - latents, - neg_prompt_ids, - return_latents, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py deleted file mode 100644 index 0783f44486ee1448bd15529f745af381ee7fa69f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py +++ /dev/null @@ -1,37 +0,0 @@ -from dataclasses import dataclass -from typing import List, Union - -import numpy as np -import PIL.Image - -from ...utils import BaseOutput, is_flax_available - - -@dataclass -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, - num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] - - -if is_flax_available(): - import flax - - @flax.struct.dataclass - class FlaxStableDiffusionXLPipelineOutput(BaseOutput): - """ - Output class for Flax Stable Diffusion XL pipelines. - - Args: - images (`np.ndarray`) - Array of shape `(batch_size, height, width, num_channels)` with images from the diffusion pipeline. - """ - - images: np.ndarray diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py deleted file mode 100644 index 9ac64a0d842023a4b9b600039b9d27f6b2fec583..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +++ /dev/null @@ -1,1313 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - FusedAttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_invisible_watermark_available, - 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, StableDiffusionMixin -from .pipeline_output import StableDiffusionXLPipelineOutput - - -if is_invisible_watermark_available(): - from .watermark import StableDiffusionXLWatermarker - -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 StableDiffusionXLPipeline - - >>> pipe = StableDiffusionXLPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> image = pipe(prompt).images[0] - ``` -""" - - -# 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 - - -# 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 StableDiffusionXLPipeline( - DiffusionPipeline, - StableDiffusionMixin, - FromSingleFileMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion XL. - - 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`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to - watermark output images. If not defined, it will default to True if the package is installed, otherwise no - watermarker will be used. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "image_encoder", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "add_text_embeds", - "add_time_ids", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - image_encoder: CLIPVisionModelWithProjection = None, - feature_extractor: CLIPImageProcessor = None, - force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - ): - super().__init__() - - 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, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - ) - 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 - ) - - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - prompt_2, - height, - width, - callback_steps, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_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_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`." - ) - - 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - 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 - - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - FusedAttnProcessor2_0, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_end(self): - return self._denoising_end - - @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, - prompt_2: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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"], - **kwargs, - ): - 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. - 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*): - 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. - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output) - 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. - 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*): - 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 ge 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. - 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)`. 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. - 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. - cross_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). - 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. - 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. - 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. - - 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. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - # 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, - callback_steps, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - 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. Encode input prompt - lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else 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 - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 5. 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, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 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) - - 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, - self.do_classifier_free_guidance, - ) - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - # 8.1 Apply denoising_end - if ( - self.denoising_end is not None - and isinstance(self.denoising_end, float) - and self.denoising_end > 0 - and self.denoising_end < 1 - ): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - # 9. 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) - - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = image_embeds - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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.do_classifier_free_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents_dtype = latents.dtype - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - 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) - elif latents.dtype != self.vae.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 - self.vae = self.vae.to(latents.dtype) - - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None - has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None - if has_latents_mean and has_latents_std: - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean - else: - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, 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": - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py deleted file mode 100644 index e63c7a55ce7b76ca2f5b3cffcc254f3e8feee829..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +++ /dev/null @@ -1,1503 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import PIL.Image -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_invisible_watermark_available, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import empty_device_cache, randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from .pipeline_output import StableDiffusionXLPipelineOutput - - -if is_invisible_watermark_available(): - from .watermark import StableDiffusionXLWatermarker - -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 StableDiffusionXLImg2ImgPipeline - >>> from diffusers.utils import load_image - - >>> pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - >>> url = "https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/aa_xl/000000009.png" - - >>> init_image = load_image(url).convert("RGB") - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> image = pipe(prompt, image=init_image).images[0] - ``` -""" - - -# 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 - - -# 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" -): - 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 StableDiffusionXLImg2ImgPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - FromSingleFileMixin, - StableDiffusionXLLoraLoaderMixin, - IPAdapterMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion XL. - - 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`]. - requires_aesthetics_score (`bool`, *optional*, defaults to `"False"`): - Whether the `unet` requires an `aesthetic_score` condition to be passed during inference. Also see the - config of `stabilityai/stable-diffusion-xl-refiner-1-0`. - 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`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to - watermark output images. If not defined, it will default to True if the package is installed, otherwise no - watermarker will be used. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "image_encoder", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "add_text_embeds", - "add_time_ids", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - image_encoder: CLIPVisionModelWithProjection = None, - feature_extractor: CLIPImageProcessor = None, - requires_aesthetics_score: bool = False, - force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - scheduler=scheduler, - ) - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - self.register_to_config(requires_aesthetics_score=requires_aesthetics_score) - 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) - - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - prompt_2, - strength, - num_inference_steps, - callback_steps, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - ): - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - if num_inference_steps is None: - raise ValueError("`num_inference_steps` cannot be None.") - elif not isinstance(num_inference_steps, int) or num_inference_steps <= 0: - raise ValueError( - f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type" - f" {type(num_inference_steps)}." - ) - 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_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 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" - ) - - def get_timesteps(self, num_inference_steps, strength, device, denoising_start=None): - # get the original timestep using init_timestep - if denoising_start is None: - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - t_start = 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 - - else: - # Strength is irrelevant if we directly request a timestep to start at; - # that is, strength is determined by the denoising_start instead. - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (denoising_start * self.scheduler.config.num_train_timesteps) - ) - ) - - num_inference_steps = (self.scheduler.timesteps < discrete_timestep_cutoff).sum().item() - if self.scheduler.order == 2 and num_inference_steps % 2 == 0: - # if the scheduler is a 2nd order scheduler we might have to do +1 - # because `num_inference_steps` might be even given that every timestep - # (except the highest one) is duplicated. If `num_inference_steps` is even it would - # mean that we cut the timesteps in the middle of the denoising step - # (between 1st and 2nd derivative) which leads to incorrect results. By adding 1 - # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler - num_inference_steps = num_inference_steps + 1 - - # because t_n+1 >= t_n, we slice the timesteps starting from the end - t_start = len(self.scheduler.timesteps) - num_inference_steps - timesteps = self.scheduler.timesteps[t_start:] - if hasattr(self.scheduler, "set_begin_index"): - self.scheduler.set_begin_index(t_start) - return timesteps, num_inference_steps - - def prepare_latents( - self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None, add_noise=True - ): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - latents_mean = latents_std = None - if hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None: - latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1) - if hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None: - latents_std = torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1) - - # Offload text encoder if `enable_model_cpu_offload` was enabled - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.text_encoder_2.to("cpu") - empty_device_cache() - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - # make sure the VAE is in float32 mode, as it overflows in float16 - if self.vae.config.force_upcast: - image = image.float() - self.vae.to(dtype=torch.float32) - - 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." - ) - - elif isinstance(generator, list): - if image.shape[0] < batch_size and batch_size % image.shape[0] == 0: - image = torch.cat([image] * (batch_size // image.shape[0]), dim=0) - elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} " - ) - - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - if self.vae.config.force_upcast: - self.vae.to(dtype) - - init_latents = init_latents.to(dtype) - if latents_mean is not None and latents_std is not None: - latents_mean = latents_mean.to(device=device, dtype=dtype) - latents_std = latents_std.to(device=device, dtype=dtype) - init_latents = (init_latents - latents_mean) * self.vae.config.scaling_factor / latents_std - else: - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - if add_noise: - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - - latents = init_latents - - return latents - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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 _get_add_time_ids( - self, - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype, - text_encoder_projection_dim=None, - ): - if self.config.requires_aesthetics_score: - add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,)) - add_neg_time_ids = list( - negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,) - ) - else: - add_time_ids = list(original_size + crops_coords_top_left + target_size) - add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_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 - and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model." - ) - elif ( - expected_add_embed_dim < passed_add_embed_dim - and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model." - ) - elif 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) - add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype) - - return add_time_ids, add_neg_time_ids - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_end(self): - return self._denoising_end - - @property - def denoising_start(self): - return self._denoising_start - - @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, - prompt_2: Optional[Union[str, List[str]]] = 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, - 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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, - 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"], - **kwargs, - ): - 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 - 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` - 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`. Note that in the case of - `denoising_start` being declared as an integer, the value of `strength` will be ignored. - 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. - 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. - denoising_start (`float`, *optional*): - When specified, indicates the fraction (between 0.0 and 1.0) of the total denoising process to be - bypassed before it is initiated. Consequently, the initial part of the denoising process is skipped and - it is assumed that the passed `image` is a partly denoised image. Note that when this is specified, - strength will be ignored. The `denoising_start` parameter is particularly beneficial when this pipeline - is integrated into a "Mixture of Denoisers" multi-pipeline setup, as detailed in [**Refine Image - Quality**](https://huggingface.co/docs/diffusers/using-diffusers/sdxl#refine-image-quality). - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - still retain a substantial amount of noise (ca. final 20% of timesteps still needed) and should be - denoised by a successor pipeline that has `denoising_start` set to 0.8 so that it only denoises the - final 20% of 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 [**Refine Image - Quality**](https://huggingface.co/docs/diffusers/using-diffusers/sdxl#refine-image-quality). - 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`). - 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. - 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*): - 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 ge 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. - 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)`. 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. - 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.StableDiffusionXLPipelineOutput`] instead of a - plain tuple. - cross_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). - 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. - 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. - aesthetic_score (`float`, *optional*, defaults to 6.0): - Used to simulate an aesthetic score of the generated image by influencing the positive text condition. - 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_aesthetic_score (`float`, *optional*, defaults to 2.5): - 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). Can be used to - simulate an aesthetic score of the generated image by influencing the negative text condition. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a - `tuple. When returning a tuple, the first element is a list with the generated images. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - 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_2, - strength, - num_inference_steps, - callback_steps, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - self._denoising_start = denoising_start - 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. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else 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=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - - # 4. Preprocess image - image = self.image_processor.preprocess(image) - - # 5. Prepare timesteps - def denoising_value_valid(dnv): - return isinstance(dnv, float) and 0 < dnv < 1 - - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - timesteps, num_inference_steps = self.get_timesteps( - num_inference_steps, - strength, - device, - denoising_start=self.denoising_start if denoising_value_valid(self.denoising_start) else None, - ) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - - add_noise = True if self.denoising_start is None else False - - # 6. Prepare latent variables - if latents is None: - latents = self.prepare_latents( - image, - latent_timestep, - batch_size, - num_images_per_prompt, - prompt_embeds.dtype, - device, - generator, - add_noise, - ) - # 7. Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - height, width = latents.shape[-2:] - height = height * self.vae_scale_factor - width = width * self.vae_scale_factor - - original_size = original_size or (height, width) - target_size = target_size or (height, width) - - # 8. Prepare added time ids & embeddings - if negative_original_size is None: - negative_original_size = original_size - if negative_target_size is None: - negative_target_size = target_size - - 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, add_neg_time_ids = self._get_add_time_ids( - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype=prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1) - - 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_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1) - add_time_ids = torch.cat([add_neg_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) - - 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, - self.do_classifier_free_guidance, - ) - - # 9. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - # 9.1 Apply denoising_end - if ( - self.denoising_end is not None - and self.denoising_start is not None - and denoising_value_valid(self.denoising_end) - and denoising_value_valid(self.denoising_start) - and self.denoising_start >= self.denoising_end - ): - raise ValueError( - f"`denoising_start`: {self.denoising_start} cannot be larger than or equal to `denoising_end`: " - + f" {self.denoising_end} when using type float." - ) - elif self.denoising_end is not None and denoising_value_valid(self.denoising_end): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - # 9.2 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) - - 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 - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = image_embeds - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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.do_classifier_free_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents_dtype = latents.dtype - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - 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) - elif latents.dtype != self.vae.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 - self.vae = self.vae.to(latents.dtype) - - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None - has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None - if has_latents_mean and has_latents_std: - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean - else: - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, return_dict=False)[0] - - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - else: - image = latents - - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py deleted file mode 100644 index f0bc9b9bb3e2f623c5d45dc22798c3d685d9e347..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py +++ /dev/null @@ -1,1740 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_invisible_watermark_available, - 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, StableDiffusionMixin -from .pipeline_output import StableDiffusionXLPipelineOutput - - -if is_invisible_watermark_available(): - from .watermark import StableDiffusionXLWatermarker - -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 StableDiffusionXLInpaintPipeline - >>> from diffusers.utils import load_image - - >>> pipe = StableDiffusionXLInpaintPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", - ... torch_dtype=torch.float16, - ... variant="fp16", - ... use_safetensors=True, - ... ) - >>> pipe.to("cuda") - - >>> 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" - - >>> init_image = load_image(img_url).convert("RGB") - >>> mask_image = load_image(mask_url).convert("RGB") - - >>> prompt = "A majestic tiger sitting on a bench" - >>> image = pipe( - ... prompt=prompt, image=init_image, mask_image=mask_image, num_inference_steps=50, strength=0.80 - ... ).images[0] - ``` -""" - - -# 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 - - -def mask_pil_to_torch(mask, height, width): - # preprocess mask - if isinstance(mask, (PIL.Image.Image, np.ndarray)): - mask = [mask] - - if isinstance(mask, list) and isinstance(mask[0], PIL.Image.Image): - mask = [i.resize((width, height), resample=PIL.Image.LANCZOS) for i in mask] - mask = np.concatenate([np.array(m.convert("L"))[None, None, :] for m in mask], axis=0) - mask = mask.astype(np.float32) / 255.0 - elif isinstance(mask, list) and isinstance(mask[0], np.ndarray): - mask = np.concatenate([m[None, None, :] for m in mask], axis=0) - - mask = torch.from_numpy(mask) - return mask - - -# 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" -): - 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 StableDiffusionXLInpaintPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionXLLoraLoaderMixin, - FromSingleFileMixin, - IPAdapterMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion XL. - - 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`]. - requires_aesthetics_score (`bool`, *optional*, defaults to `"False"`): - Whether the `unet` requires a aesthetic_score condition to be passed during inference. Also see the config - of `stabilityai/stable-diffusion-xl-refiner-1-0`. - 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`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to - watermark output images. If not defined, it will default to True if the package is installed, otherwise no - watermarker will be used. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae" - - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "image_encoder", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "add_text_embeds", - "add_time_ids", - "mask", - "masked_image_latents", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - image_encoder: CLIPVisionModelWithProjection = None, - feature_extractor: CLIPImageProcessor = None, - requires_aesthetics_score: bool = False, - force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - scheduler=scheduler, - ) - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - self.register_to_config(requires_aesthetics_score=requires_aesthetics_score) - 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.mask_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True - ) - - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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_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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - prompt_2, - image, - mask_image, - height, - width, - strength, - callback_steps, - output_type, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_embeds=None, - callback_on_step_end_tensor_inputs=None, - padding_mask_crop=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 % 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_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 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 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" - ) - - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=None, - image=None, - timestep=None, - is_strength_max=True, - add_noise=True, - return_noise=False, - return_image_latents=False, - ): - 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 (image is None or timestep is None) and not is_strength_max: - raise ValueError( - "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise." - "However, either the image or the noise timestep has not been provided." - ) - - if image.shape[1] == 4: - image_latents = image.to(device=device, dtype=dtype) - image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1) - elif return_image_latents or (latents is None and not is_strength_max): - image = image.to(device=device, dtype=dtype) - image_latents = self._encode_vae_image(image=image, generator=generator) - image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1) - - if latents is None and add_noise: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # if strength is 1. then initialise the latents to noise, else initial to image + noise - latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep) - # if pure noise then scale the initial latents by the Scheduler's init sigma - latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents - elif add_noise: - noise = latents.to(device) - latents = noise * self.scheduler.init_noise_sigma - else: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = image_latents.to(device) - - outputs = (latents,) - - if return_noise: - outputs += (noise,) - - if return_image_latents: - outputs += (image_latents,) - - return outputs - - def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): - dtype = image.dtype - if self.vae.config.force_upcast: - image = image.float() - self.vae.to(dtype=torch.float32) - - 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) - - if self.vae.config.force_upcast: - self.vae.to(dtype) - - image_latents = image_latents.to(dtype) - image_latents = self.vae.config.scaling_factor * image_latents - - return image_latents - - def prepare_mask_latents( - self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance - ): - # 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 // self.vae_scale_factor, width // self.vae_scale_factor) - ) - mask = mask.to(device=device, dtype=dtype) - - # 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) - - mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask - - if masked_image is not None and masked_image.shape[1] == 4: - masked_image_latents = masked_image - else: - masked_image_latents = None - - if masked_image is not None: - if masked_image_latents is None: - masked_image = masked_image.to(device=device, dtype=dtype) - masked_image_latents = self._encode_vae_image(masked_image, generator=generator) - - 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 - ) - - masked_image_latents = ( - torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents - ) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device, denoising_start=None): - # get the original timestep using init_timestep - if denoising_start is None: - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - t_start = 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 - - else: - # Strength is irrelevant if we directly request a timestep to start at; - # that is, strength is determined by the denoising_start instead. - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (denoising_start * self.scheduler.config.num_train_timesteps) - ) - ) - - num_inference_steps = (self.scheduler.timesteps < discrete_timestep_cutoff).sum().item() - if self.scheduler.order == 2 and num_inference_steps % 2 == 0: - # if the scheduler is a 2nd order scheduler we might have to do +1 - # because `num_inference_steps` might be even given that every timestep - # (except the highest one) is duplicated. If `num_inference_steps` is even it would - # mean that we cut the timesteps in the middle of the denoising step - # (between 1st and 2nd derivative) which leads to incorrect results. By adding 1 - # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler - num_inference_steps = num_inference_steps + 1 - - # because t_n+1 >= t_n, we slice the timesteps starting from the end - t_start = len(self.scheduler.timesteps) - num_inference_steps - timesteps = self.scheduler.timesteps[t_start:] - if hasattr(self.scheduler, "set_begin_index"): - self.scheduler.set_begin_index(t_start) - return timesteps, num_inference_steps - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline._get_add_time_ids - def _get_add_time_ids( - self, - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype, - text_encoder_projection_dim=None, - ): - if self.config.requires_aesthetics_score: - add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,)) - add_neg_time_ids = list( - negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,) - ) - else: - add_time_ids = list(original_size + crops_coords_top_left + target_size) - add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_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 - and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model." - ) - elif ( - expected_add_embed_dim < passed_add_embed_dim - and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_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. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model." - ) - elif 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) - add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype) - - return add_time_ids, add_neg_time_ids - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @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 - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_end(self): - return self._denoising_end - - @property - def denoising_start(self): - return self._denoising_start - - @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, - prompt_2: Optional[Union[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, - 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, - 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, - 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", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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, - 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"], - **kwargs, - ): - 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 - image (`PIL.Image.Image`): - `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will - be masked out with `mask_image` and repainted according to `prompt`. - mask_image (`PIL.Image.Image`): - `Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be - 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)`. - 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. - padding_mask_crop (`int`, *optional*, defaults to `None`): - The size of margin in the crop to be applied to the image and masking. If `None`, no crop is applied to - image and mask_image. If `padding_mask_crop` is not `None`, it will first find a rectangular region - with the same aspect ration of the image and contains all masked area, and then expand that area based - on `padding_mask_crop`. The image and mask_image will then be cropped based on the expanded area before - resizing to the original image size for inpainting. This is useful when the masked area is small while - the image is large and contain information irrelevant for inpainting, such as background. - strength (`float`, *optional*, defaults to 0.9999): - Conceptually, indicates how much to transform the masked portion of 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 the masked - portion of the reference `image`. Note that in the case of `denoising_start` being declared as an - integer, the value of `strength` will be ignored. - 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. - 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. - denoising_start (`float`, *optional*): - When specified, indicates the fraction (between 0.0 and 1.0) of the total denoising process to be - bypassed before it is initiated. Consequently, the initial part of the denoising process is skipped and - it is assumed that the passed `image` is a partly denoised image. Note that when this is specified, - strength will be ignored. The `denoising_start` parameter is particularly beneficial when this pipeline - is integrated into a "Mixture of Denoisers" multi-pipeline setup, as detailed in [**Refining the Image - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output). - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - still retain a substantial amount of noise (ca. final 20% of timesteps still needed) and should be - denoised by a successor pipeline that has `denoising_start` set to 0.8 so that it only denoises the - final 20% of 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output). - 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`). - 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. - 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)`. 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. - 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 ge generated by sampling using the supplied random `generator`. - 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. - cross_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). - 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. - aesthetic_score (`float`, *optional*, defaults to 6.0): - Used to simulate an aesthetic score of the generated image by influencing the positive text condition. - 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_aesthetic_score (`float`, *optional*, defaults to 2.5): - 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). Can be used to - simulate an aesthetic score of the generated image by influencing the negative text condition. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a - `tuple. `tuple. When returning a tuple, the first element is a list with the generated images. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - # 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 - self.check_inputs( - prompt, - prompt_2, - image, - mask_image, - height, - width, - strength, - callback_steps, - output_type, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - callback_on_step_end_tensor_inputs, - padding_mask_crop, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - self._denoising_start = denoising_start - 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. Encode input prompt - text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else 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=text_encoder_lora_scale, - clip_skip=self.clip_skip, - ) - - # 4. set timesteps - def denoising_value_valid(dnv): - return isinstance(dnv, float) and 0 < dnv < 1 - - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - timesteps, num_inference_steps = self.get_timesteps( - num_inference_steps, - strength, - device, - denoising_start=self.denoising_start if denoising_value_valid(self.denoising_start) else None, - ) - # check that number of inference steps is not < 1 - as this doesn't make sense - 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." - ) - # at which timestep to set the initial noise (n.b. 50% if strength is 0.5) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise - is_strength_max = strength == 1.0 - - # 5. 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) - - mask = self.mask_processor.preprocess( - mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords - ) - - if masked_image_latents is not None: - masked_image = masked_image_latents - elif init_image.shape[1] == 4: - # if images are in latent space, we can't mask it - masked_image = None - else: - masked_image = init_image * (mask < 0.5) - - # 6. Prepare latent variables - num_channels_latents = self.vae.config.latent_channels - num_channels_unet = self.unet.config.in_channels - return_image_latents = num_channels_unet == 4 - - add_noise = True if self.denoising_start is None else False - latents_outputs = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - image=init_image, - timestep=latent_timestep, - is_strength_max=is_strength_max, - add_noise=add_noise, - return_noise=True, - return_image_latents=return_image_latents, - ) - - if return_image_latents: - latents, noise, image_latents = latents_outputs - else: - latents, noise = latents_outputs - - # 7. Prepare mask latent variables - mask, masked_image_latents = self.prepare_mask_latents( - mask, - masked_image, - batch_size * num_images_per_prompt, - height, - width, - prompt_embeds.dtype, - device, - generator, - self.do_classifier_free_guidance, - ) - - # 8. Check that sizes of mask, masked image and latents match - if num_channels_unet == 9: - # default case for runwayml/stable-diffusion-inpainting - num_channels_mask = mask.shape[1] - num_channels_masked_image = masked_image_latents.shape[1] - if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" - " `pipeline.unet` or your `mask_image` or `image` input." - ) - elif num_channels_unet != 4: - raise ValueError( - f"The unet {self.unet.__class__} should have either 4 or 9 input channels, not {self.unet.config.in_channels}." - ) - # 8.1 Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - height, width = latents.shape[-2:] - height = height * self.vae_scale_factor - width = width * self.vae_scale_factor - - original_size = original_size or (height, width) - target_size = target_size or (height, width) - - # 10. Prepare added time ids & embeddings - if negative_original_size is None: - negative_original_size = original_size - if negative_target_size is None: - negative_target_size = target_size - - 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, add_neg_time_ids = self._get_add_time_ids( - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype=prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1) - - 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_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1) - add_time_ids = torch.cat([add_neg_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) - - 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, - self.do_classifier_free_guidance, - ) - - # 11. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - if ( - self.denoising_end is not None - and self.denoising_start is not None - and denoising_value_valid(self.denoising_end) - and denoising_value_valid(self.denoising_start) - and self.denoising_start >= self.denoising_end - ): - raise ValueError( - f"`denoising_start`: {self.denoising_start} cannot be larger than or equal to `denoising_end`: " - + f" {self.denoising_end} when using type float." - ) - elif self.denoising_end is not None and denoising_value_valid(self.denoising_end): - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - # 11.1 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) - - 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 - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - # concat latents, mask, masked_image_latents in the channel dimension - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - if num_channels_unet == 9: - latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) - - # predict the noise residual - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = image_embeds - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=self.cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - 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.do_classifier_free_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents_dtype = latents.dtype - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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 num_channels_unet == 4: - init_latents_proper = image_latents - if self.do_classifier_free_guidance: - init_mask, _ = mask.chunk(2) - else: - init_mask = mask - - if i < len(timesteps) - 1: - noise_timestep = timesteps[i + 1] - init_latents_proper = self.scheduler.add_noise( - init_latents_proper, noise, torch.tensor([noise_timestep]) - ) - - latents = (1 - init_mask) * init_latents_proper + init_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) - add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) - add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) - 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - 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) - elif latents.dtype != self.vae.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 - self.vae = self.vae.to(latents.dtype) - - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None - has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None - if has_latents_mean and has_latents_std: - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean - else: - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, return_dict=False)[0] - - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - else: - return StableDiffusionXLPipelineOutput(images=latents) - - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - image = self.image_processor.postprocess(image, output_type=output_type) - - if padding_mask_crop is not None: - image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return StableDiffusionXLPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py deleted file mode 100644 index b1379d1b2955e1f2d67794dd3c3b6c8ec23848ff..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py +++ /dev/null @@ -1,1000 +0,0 @@ -# Copyright 2025 Harutatsu Akiyama 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, Tuple, Union - -import PIL.Image -import torch -from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin, StableDiffusionXLLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - FusedAttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_invisible_watermark_available, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin -from .pipeline_output import StableDiffusionXLPipelineOutput - - -if is_invisible_watermark_available(): - from .watermark import StableDiffusionXLWatermarker - -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 StableDiffusionXLInstructPix2PixPipeline - >>> from diffusers.utils import load_image - - >>> resolution = 768 - >>> image = load_image( - ... "https://hf.co/datasets/diffusers/diffusers-images-docs/resolve/main/mountain.png" - ... ).resize((resolution, resolution)) - >>> edit_instruction = "Turn sky into a cloudy one" - - >>> pipe = StableDiffusionXLInstructPix2PixPipeline.from_pretrained( - ... "diffusers/sdxl-instructpix2pix-768", torch_dtype=torch.float16 - ... ).to("cuda") - - >>> edited_image = pipe( - ... prompt=edit_instruction, - ... image=image, - ... height=resolution, - ... width=resolution, - ... guidance_scale=3.0, - ... image_guidance_scale=1.5, - ... num_inference_steps=30, - ... ).images[0] - >>> edited_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" -): - 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") - - -def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): - """ - Rescale `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 - """ - 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 StableDiffusionXLInstructPix2PixPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - FromSingleFileMixin, - StableDiffusionXLLoraLoaderMixin, -): - r""" - Pipeline for pixel-level image editing by following text instructions. Based on Stable Diffusion XL. - - 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 - - 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`]. - requires_aesthetics_score (`bool`, *optional*, defaults to `"False"`): - Whether the `unet` requires a aesthetic_score condition to be passed during inference. Also see the config - of `stabilityai/stable-diffusion-xl-refiner-1-0`. - 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`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to - watermark output images. If not defined, it will default to True if the package is installed, otherwise no - watermarker will be used. - is_cosxl_edit (`bool`, *optional*): - When set the image latents are scaled. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae" - _optional_components = ["tokenizer", "tokenizer_2", "text_encoder", "text_encoder_2"] - - 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, - add_watermarker: Optional[bool] = None, - is_cosxl_edit: Optional[bool] = False, - ): - super().__init__() - - 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 - ) - self.is_cosxl_edit = is_cosxl_edit - - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - 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, - ): - 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. - """ - 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) - - 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] - - # 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 - # 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] - - prompt_embeds = prompt_embeds.hidden_states[-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 - - 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 isinstance(negative_prompt, str): - uncond_tokens = [negative_prompt, negative_prompt_2] - 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) - - prompt_embeds_dtype = self.text_encoder_2.dtype if self.text_encoder_2 is not None else self.unet.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) - - 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) - - 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 - ) - - return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - ): - 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}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def prepare_image_latents( - self, image, batch_size, num_images_per_prompt, dtype, device, do_classifier_free_guidance, generator=None - ): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - image_latents = image - else: - # 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: - image = image.float() - self.upcast_vae() - - image_latents = retrieve_latents(self.vae.encode(image), sample_mode="argmax") - - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - - if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: - # expand image_latents for batch_size - deprecation_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {image_latents.shape[0]} initial" - " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" - " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" - " your script to pass as many initial images as text prompts to suppress this warning." - ) - deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) - 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 do_classifier_free_guidance: - uncond_image_latents = torch.zeros_like(image_latents) - image_latents = torch.cat([image_latents, image_latents, uncond_image_latents], dim=0) - - if image_latents.dtype != self.vae.dtype: - image_latents = image_latents.to(dtype=self.vae.dtype) - - if self.is_cosxl_edit: - image_latents = image_latents * self.vae.config.scaling_factor - - return image_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): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - FusedAttnProcessor2_0, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - @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, - image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 100, - denoising_end: Optional[float] = 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, - 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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, - ): - 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 - 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. - 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. - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output) - 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. - image_guidance_scale (`float`, *optional*, defaults to 1.5): - Image guidance scale is to push the generated image towards the initial image `image`. Image guidance - 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*): - 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. - 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*): - 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 ge 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.StableDiffusionXLPipelineOutput`] 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. - cross_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). - 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. - 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). - aesthetic_score (`float`, *optional*, defaults to 6.0): - Used to simulate an aesthetic score of the generated image by influencing the positive text condition. - 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_aesthetic_score (`float`, *optional*, defaults to 2.5): - 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). Can be used to - simulate an aesthetic score of the generated image by influencing the negative text condition. - - 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, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds) - - if image is None: - raise ValueError("`image` input cannot be undefined.") - - # 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 = guidance_scale > 1.0 and image_guidance_scale >= 1.0 - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else 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=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=text_encoder_lora_scale, - ) - - # 4. Preprocess image - image = self.image_processor.preprocess(image, height=height, width=width).to(device) - - # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 6. Prepare Image latents - image_latents = self.prepare_image_latents( - image, - batch_size, - num_images_per_prompt, - prompt_embeds.dtype, - device, - do_classifier_free_guidance, - ) - - # 7. Prepare latent variables - num_channels_latents = self.vae.config.latent_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 8. Check that shapes of latents and image match the UNet channels - num_channels_image = image_latents.shape[1] - if num_channels_latents + num_channels_image != self.unet.config.in_channels: - raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_image`: {num_channels_image} " - f" = {num_channels_latents + num_channels_image}. Please verify the config of" - " `pipeline.unet` or your `image` input." - ) - - # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 10. 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 do_classifier_free_guidance: - # The extra concat similar to how it's done in SD InstructPix2Pix. - prompt_embeds = torch.cat([prompt_embeds, negative_prompt_embeds, negative_prompt_embeds], dim=0) - add_text_embeds = torch.cat( - [add_text_embeds, negative_pooled_prompt_embeds, negative_pooled_prompt_embeds], dim=0 - ) - add_time_ids = torch.cat([add_time_ids, 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) - - # 11. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - if denoising_end is not None and isinstance(denoising_end, float) and denoising_end > 0 and denoising_end < 1: - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # Expand the latents if we are doing classifier free guidance. - # The latents are expanded 3 times because for pix2pix the guidance - # is applied for both the text and the input image. - latent_model_input = torch.cat([latents] * 3) if do_classifier_free_guidance else latents - - # concat latents, image_latents in the channel dimension - scaled_latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - scaled_latent_model_input = torch.cat([scaled_latent_model_input, image_latents], dim=1) - - # predict the noise residual - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - noise_pred = self.unet( - scaled_latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_text, noise_pred_image, noise_pred_uncond = noise_pred.chunk(3) - noise_pred = ( - noise_pred_uncond - + guidance_scale * (noise_pred_text - noise_pred_image) - + image_guidance_scale * (noise_pred_image - noise_pred_uncond) - ) - - if do_classifier_free_guidance and guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents_dtype = latents.dtype - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, 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) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - 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) - elif latents.dtype != self.vae.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 - self.vae = self.vae.to(latents.dtype) - - # unscale/denormalize the latents - # denormalize with the mean and std if available and not None - has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None - has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None - if has_latents_mean and has_latents_std: - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) - ) - latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean - else: - latents = latents / self.vae.config.scaling_factor - - image = self.vae.decode(latents, return_dict=False)[0] - - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - else: - return StableDiffusionXLPipelineOutput(images=latents) - - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/watermark.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/watermark.py deleted file mode 100644 index 70d06bb6320d890ea90542cc7551bc42d7c59451..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/watermark.py +++ /dev/null @@ -1,42 +0,0 @@ -import numpy as np -import torch - -from ...utils import is_invisible_watermark_available - - -if is_invisible_watermark_available(): - from imwatermark import WatermarkEncoder - - -# Copied from https://github.com/Stability-AI/generative-models/blob/613af104c6b85184091d42d374fef420eddb356d/scripts/demo/streamlit_helpers.py#L66 -WATERMARK_MESSAGE = 0b101100111110110010010000011110111011000110011110 -# bin(x)[2:] gives bits of x as str, use int to convert them to 0/1 -WATERMARK_BITS = [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]] - - -class StableDiffusionXLWatermarker: - def __init__(self): - self.watermark = WATERMARK_BITS - self.encoder = WatermarkEncoder() - - self.encoder.set_watermark("bits", self.watermark) - - def apply_watermark(self, images: torch.Tensor): - # can't encode images that are smaller than 256 - if images.shape[-1] < 256: - return images - - images = (255 * (images / 2 + 0.5)).cpu().permute(0, 2, 3, 1).float().numpy() - - # Convert RGB to BGR, which is the channel order expected by the watermark encoder. - images = images[:, :, :, ::-1] - - # Add watermark and convert BGR back to RGB - images = [self.encoder.encode(image, "dwtDct")[:, :, ::-1] for image in images] - - images = np.array(images) - - images = torch.from_numpy(images).permute(0, 3, 1, 2) - - images = torch.clamp(2 * (images / 255 - 0.5), min=-1.0, max=1.0) - return images diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_video_diffusion/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_video_diffusion/__init__.py deleted file mode 100644 index 3bd4dc78966e217d85769691b98ed8fb0b6ac05c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_video_diffusion/__init__.py +++ /dev/null @@ -1,58 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - BaseOutput, - 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 - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) -else: - _import_structure.update( - { - "pipeline_stable_video_diffusion": [ - "StableVideoDiffusionPipeline", - "StableVideoDiffusionPipelineOutput", - ], - } - ) - - -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_stable_video_diffusion import ( - StableVideoDiffusionPipeline, - StableVideoDiffusionPipelineOutput, - ) - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_video_diffusion/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_video_diffusion/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 4922a2f3b8789c96cdca9dbe753177dd945e3ba3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_video_diffusion/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_video_diffusion/__pycache__/pipeline_stable_video_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_video_diffusion/__pycache__/pipeline_stable_video_diffusion.cpython-310.pyc deleted file mode 100644 index 6b6e0eaf5c1f660781ef596a61261f2db2fa7370..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_video_diffusion/__pycache__/pipeline_stable_video_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py deleted file mode 100644 index 6d9053faaec8a8af0fce4b941b25667473a5e4a8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py +++ /dev/null @@ -1,737 +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 inspect -from dataclasses import dataclass -from typing import Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection - -from ...image_processor import PipelineImageInput -from ...models import AutoencoderKLTemporalDecoder, UNetSpatioTemporalConditionModel -from ...schedulers import EulerDiscreteScheduler -from ...utils import BaseOutput, is_torch_xla_available, logging, replace_example_docstring -from ...utils.torch_utils import is_compiled_module, randn_tensor -from ...video_processor import VideoProcessor -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 - >>> from diffusers import StableVideoDiffusionPipeline - >>> from diffusers.utils import load_image, export_to_video - - >>> pipe = StableVideoDiffusionPipeline.from_pretrained( - ... "stabilityai/stable-video-diffusion-img2vid-xt", torch_dtype=torch.float16, variant="fp16" - ... ) - >>> pipe.to("cuda") - - >>> image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/svd-docstring-example.jpeg" - ... ) - >>> image = image.resize((1024, 576)) - - >>> frames = pipe(image, num_frames=25, decode_chunk_size=8).frames[0] - >>> export_to_video(frames, "generated.mp4", fps=7) - ``` -""" - - -def _append_dims(x, target_dims): - """Appends dimensions to the end of a tensor until it has target_dims dimensions.""" - dims_to_append = target_dims - x.ndim - if dims_to_append < 0: - raise ValueError(f"input has {x.ndim} dims but target_dims is {target_dims}, which is less") - return x[(...,) + (None,) * dims_to_append] - - -# 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 - - -@dataclass -class StableVideoDiffusionPipelineOutput(BaseOutput): - r""" - 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, - num_frames, height, width, num_channels)`. - """ - - frames: Union[List[List[PIL.Image.Image]], np.ndarray, torch.Tensor] - - -class StableVideoDiffusionPipeline(DiffusionPipeline): - r""" - Pipeline to generate video from an input image using Stable Video Diffusion. - - 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: - vae ([`AutoencoderKLTemporalDecoder`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. - image_encoder ([`~transformers.CLIPVisionModelWithProjection`]): - Frozen CLIP image-encoder - ([laion/CLIP-ViT-H-14-laion2B-s32B-b79K](https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K)). - unet ([`UNetSpatioTemporalConditionModel`]): - A `UNetSpatioTemporalConditionModel` to denoise the encoded image latents. - scheduler ([`EulerDiscreteScheduler`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - A `CLIPImageProcessor` to extract features from generated images. - """ - - model_cpu_offload_seq = "image_encoder->unet->vae" - _callback_tensor_inputs = ["latents"] - - def __init__( - self, - vae: AutoencoderKLTemporalDecoder, - image_encoder: CLIPVisionModelWithProjection, - unet: UNetSpatioTemporalConditionModel, - scheduler: EulerDiscreteScheduler, - feature_extractor: CLIPImageProcessor, - ): - super().__init__() - - self.register_modules( - vae=vae, - image_encoder=image_encoder, - unet=unet, - scheduler=scheduler, - 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.video_processor = VideoProcessor(do_resize=True, vae_scale_factor=self.vae_scale_factor) - - def _encode_image( - self, - image: PipelineImageInput, - device: Union[str, torch.device], - num_videos_per_prompt: int, - do_classifier_free_guidance: bool, - ) -> torch.Tensor: - dtype = next(self.image_encoder.parameters()).dtype - - if not isinstance(image, torch.Tensor): - image = self.video_processor.pil_to_numpy(image) - image = self.video_processor.numpy_to_pt(image) - - # We normalize the image before resizing to match with the original implementation. - # Then we unnormalize it after resizing. - image = image * 2.0 - 1.0 - image = _resize_with_antialiasing(image, (224, 224)) - image = (image + 1.0) / 2.0 - - # Normalize the image with for CLIP input - image = self.feature_extractor( - images=image, - do_normalize=True, - do_center_crop=False, - do_resize=False, - do_rescale=False, - return_tensors="pt", - ).pixel_values - - image = image.to(device=device, dtype=dtype) - image_embeddings = self.image_encoder(image).image_embeds - image_embeddings = image_embeddings.unsqueeze(1) - - # duplicate image embeddings for each generation per prompt, using mps friendly method - bs_embed, seq_len, _ = image_embeddings.shape - image_embeddings = image_embeddings.repeat(1, num_videos_per_prompt, 1) - image_embeddings = image_embeddings.view(bs_embed * num_videos_per_prompt, seq_len, -1) - - if do_classifier_free_guidance: - negative_image_embeddings = torch.zeros_like(image_embeddings) - - # 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 - image_embeddings = torch.cat([negative_image_embeddings, image_embeddings]) - - return image_embeddings - - def _encode_vae_image( - self, - image: torch.Tensor, - device: Union[str, torch.device], - num_videos_per_prompt: int, - do_classifier_free_guidance: bool, - ): - image = image.to(device=device) - image_latents = self.vae.encode(image).latent_dist.mode() - - # duplicate image_latents for each generation per prompt, using mps friendly method - image_latents = image_latents.repeat(num_videos_per_prompt, 1, 1, 1) - - if do_classifier_free_guidance: - negative_image_latents = torch.zeros_like(image_latents) - - # 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 - image_latents = torch.cat([negative_image_latents, image_latents]) - - return image_latents - - def _get_add_time_ids( - self, - fps: int, - motion_bucket_id: int, - noise_aug_strength: float, - dtype: torch.dtype, - batch_size: int, - num_videos_per_prompt: int, - do_classifier_free_guidance: bool, - ): - add_time_ids = [fps, motion_bucket_id, noise_aug_strength] - - passed_add_embed_dim = self.unet.config.addition_time_embed_dim * len(add_time_ids) - 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) - add_time_ids = add_time_ids.repeat(batch_size * num_videos_per_prompt, 1) - - if do_classifier_free_guidance: - add_time_ids = torch.cat([add_time_ids, add_time_ids]) - - return add_time_ids - - def decode_latents(self, latents: torch.Tensor, num_frames: int, decode_chunk_size: int = 14): - # [batch, frames, channels, height, width] -> [batch*frames, channels, height, width] - latents = latents.flatten(0, 1) - - latents = 1 / self.vae.config.scaling_factor * latents - - forward_vae_fn = self.vae._orig_mod.forward if is_compiled_module(self.vae) else self.vae.forward - accepts_num_frames = "num_frames" in set(inspect.signature(forward_vae_fn).parameters.keys()) - - # decode decode_chunk_size frames at a time to avoid OOM - frames = [] - for i in range(0, latents.shape[0], decode_chunk_size): - num_frames_in = latents[i : i + decode_chunk_size].shape[0] - decode_kwargs = {} - if accepts_num_frames: - # we only pass num_frames_in if it's expected - decode_kwargs["num_frames"] = num_frames_in - - frame = self.vae.decode(latents[i : i + decode_chunk_size], **decode_kwargs).sample - frames.append(frame) - frames = torch.cat(frames, dim=0) - - # [batch*frames, channels, height, width] -> [batch, channels, frames, height, width] - frames = frames.reshape(-1, num_frames, *frames.shape[1:]).permute(0, 2, 1, 3, 4) - - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - frames = frames.float() - return frames - - def check_inputs(self, image, height, width): - if ( - not isinstance(image, torch.Tensor) - and not isinstance(image, PIL.Image.Image) - 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" - f" {type(image)}" - ) - - 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}.") - - def prepare_latents( - self, - batch_size: int, - num_frames: int, - num_channels_latents: int, - height: int, - width: int, - dtype: torch.dtype, - device: Union[str, torch.device], - generator: torch.Generator, - latents: Optional[torch.Tensor] = None, - ): - shape = ( - batch_size, - num_frames, - num_channels_latents // 2, - height // self.vae_scale_factor, - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @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 do_classifier_free_guidance(self): - if isinstance(self.guidance_scale, (int, float)): - return self.guidance_scale > 1 - return self.guidance_scale.max() > 1 - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - image: Union[PIL.Image.Image, List[PIL.Image.Image], torch.Tensor], - height: int = 576, - width: int = 1024, - num_frames: Optional[int] = None, - num_inference_steps: int = 25, - sigmas: Optional[List[float]] = 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"], - 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(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`): - 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_frames (`int`, *optional*): - The number of video frames to generate. Defaults to `self.unet.config.num_frames` (14 for - `stable-video-diffusion-img2vid` and to 25 for `stable-video-diffusion-img2vid-xt`). - 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*): - 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. - min_guidance_scale (`float`, *optional*, defaults to 1.0): - The minimum guidance scale. Used for the classifier free guidance with first frame. - max_guidance_scale (`float`, *optional*, defaults to 3.0): - The maximum guidance scale. Used for the classifier free guidance with last frame. - fps (`int`, *optional*, defaults to 7): - Frames per second. The rate at which the generated images shall be exported to a video after - generation. Note that Stable Diffusion Video's UNet was micro-conditioned on fps-1 during training. - motion_bucket_id (`int`, *optional*, defaults to 127): - Used for conditioning the amount of motion for the generation. The higher the number the more motion - will be in the video. - noise_aug_strength (`float`, *optional*, defaults to 0.02): - The amount of noise added to the init image, the higher it is the less the video will look like the - init image. Increase it for more motion. - decode_chunk_size (`int`, *optional*): - The number of frames to decode at a time. Higher chunk size leads to better temporal consistency at the - expense of more memory usage. By default, the decoder decodes all frames at once for maximal quality. - 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*): - 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 video - 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`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated image. Choose between `pil`, `np` or `pt`. - callback_on_step_end (`Callable`, *optional*): - A function that is called at the end of each denoising step during 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. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - - Examples: - - 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. - """ - # 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 - - num_frames = num_frames if num_frames is not None else self.unet.config.num_frames - decode_chunk_size = decode_chunk_size if decode_chunk_size is not None else num_frames - - # 1. Check inputs. Raise error if not correct - self.check_inputs(image, height, width) - - # 2. Define call parameters - if isinstance(image, PIL.Image.Image): - batch_size = 1 - elif isinstance(image, list): - batch_size = len(image) - else: - batch_size = image.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. - self._guidance_scale = max_guidance_scale - - # 3. Encode input image - image_embeddings = self._encode_image(image, device, num_videos_per_prompt, self.do_classifier_free_guidance) - - # NOTE: Stable Video Diffusion was conditioned on fps - 1, which is why it is reduced here. - # See: https://github.com/Stability-AI/generative-models/blob/ed0997173f98eaf8f4edf7ba5fe8f15c6b877fd3/scripts/sampling/simple_video_sample.py#L188 - fps = fps - 1 - - # 4. Encode input image using VAE - image = self.video_processor.preprocess(image, height=height, width=width).to(device) - noise = randn_tensor(image.shape, generator=generator, device=device, dtype=image.dtype) - image = image + noise_aug_strength * noise - - needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast - if needs_upcasting: - self.vae.to(dtype=torch.float32) - - image_latents = self._encode_vae_image( - image, - device=device, - num_videos_per_prompt=num_videos_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - ) - image_latents = image_latents.to(image_embeddings.dtype) - - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - - # Repeat the image latents for each frame so we can concatenate them with the noise - # image_latents [batch, channels, height, width] ->[batch, num_frames, channels, height, width] - image_latents = image_latents.unsqueeze(1).repeat(1, num_frames, 1, 1, 1) - - # 5. Get Added Time IDs - added_time_ids = self._get_add_time_ids( - fps, - motion_bucket_id, - noise_aug_strength, - image_embeddings.dtype, - batch_size, - num_videos_per_prompt, - self.do_classifier_free_guidance, - ) - 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) - - # 7. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - num_frames, - num_channels_latents, - height, - width, - image_embeddings.dtype, - device, - generator, - latents, - ) - - # 8. Prepare guidance scale - guidance_scale = torch.linspace(min_guidance_scale, max_guidance_scale, num_frames).unsqueeze(0) - guidance_scale = guidance_scale.to(device, latents.dtype) - guidance_scale = guidance_scale.repeat(batch_size * num_videos_per_prompt, 1) - guidance_scale = _append_dims(guidance_scale, latents.ndim) - - self._guidance_scale = guidance_scale - - # 9. Denoising loop - 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): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # Concatenate image_latents over channels dimension - latent_model_input = torch.cat([latent_model_input, image_latents], dim=2) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=image_embeddings, - added_time_ids=added_time_ids, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents).prev_sample - - 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) - - 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 not output_type == "latent": - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - frames = self.decode_latents(latents, num_frames, decode_chunk_size) - frames = self.video_processor.postprocess_video(video=frames, output_type=output_type) - else: - frames = latents - - self.maybe_free_model_hooks() - - if not return_dict: - return frames - - return StableVideoDiffusionPipelineOutput(frames=frames) - - -# resizing utils -# TODO: clean up later -def _resize_with_antialiasing(input, size, interpolation="bicubic", align_corners=True): - h, w = input.shape[-2:] - factors = (h / size[0], w / size[1]) - - # First, we have to determine sigma - # Taken from skimage: https://github.com/scikit-image/scikit-image/blob/v0.19.2/skimage/transform/_warps.py#L171 - sigmas = ( - max((factors[0] - 1.0) / 2.0, 0.001), - max((factors[1] - 1.0) / 2.0, 0.001), - ) - - # Now kernel size. Good results are for 3 sigma, but that is kind of slow. Pillow uses 1 sigma - # https://github.com/python-pillow/Pillow/blob/master/src/libImaging/Resample.c#L206 - # But they do it in the 2 passes, which gives better results. Let's try 2 sigmas for now - ks = int(max(2.0 * 2 * sigmas[0], 3)), int(max(2.0 * 2 * sigmas[1], 3)) - - # Make sure it is odd - if (ks[0] % 2) == 0: - ks = ks[0] + 1, ks[1] - - if (ks[1] % 2) == 0: - ks = ks[0], ks[1] + 1 - - input = _gaussian_blur2d(input, ks, sigmas) - - output = torch.nn.functional.interpolate(input, size=size, mode=interpolation, align_corners=align_corners) - return output - - -def _compute_padding(kernel_size): - """Compute padding tuple.""" - # 4 or 6 ints: (padding_left, padding_right,padding_top,padding_bottom) - # https://pytorch.org/docs/stable/nn.html#torch.nn.functional.pad - if len(kernel_size) < 2: - raise AssertionError(kernel_size) - computed = [k - 1 for k in kernel_size] - - # for even kernels we need to do asymmetric padding :( - out_padding = 2 * len(kernel_size) * [0] - - for i in range(len(kernel_size)): - computed_tmp = computed[-(i + 1)] - - pad_front = computed_tmp // 2 - pad_rear = computed_tmp - pad_front - - out_padding[2 * i + 0] = pad_front - out_padding[2 * i + 1] = pad_rear - - return out_padding - - -def _filter2d(input, kernel): - # prepare kernel - b, c, h, w = input.shape - tmp_kernel = kernel[:, None, ...].to(device=input.device, dtype=input.dtype) - - tmp_kernel = tmp_kernel.expand(-1, c, -1, -1) - - height, width = tmp_kernel.shape[-2:] - - 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 - tmp_kernel = tmp_kernel.reshape(-1, 1, height, width) - input = input.view(-1, tmp_kernel.size(0), input.size(-2), input.size(-1)) - - # convolve the tensor with the kernel. - output = torch.nn.functional.conv2d(input, tmp_kernel, groups=tmp_kernel.size(0), padding=0, stride=1) - - out = output.view(b, c, h, w) - return out - - -def _gaussian(window_size: int, sigma): - if isinstance(sigma, float): - sigma = torch.tensor([[sigma]]) - - batch_size = sigma.shape[0] - - x = (torch.arange(window_size, device=sigma.device, dtype=sigma.dtype) - window_size // 2).expand(batch_size, -1) - - if window_size % 2 == 0: - x = x + 0.5 - - gauss = torch.exp(-x.pow(2.0) / (2 * sigma.pow(2.0))) - - return gauss / gauss.sum(-1, keepdim=True) - - -def _gaussian_blur2d(input, kernel_size, sigma): - if isinstance(sigma, tuple): - sigma = torch.tensor([sigma], dtype=input.dtype) - else: - sigma = sigma.to(dtype=input.dtype) - - ky, kx = int(kernel_size[0]), int(kernel_size[1]) - bs = sigma.shape[0] - kernel_x = _gaussian(kx, sigma[:, 1].view(bs, 1)) - kernel_y = _gaussian(ky, sigma[:, 0].view(bs, 1)) - out_x = _filter2d(input, kernel_x[..., None, :]) - out = _filter2d(out_x, kernel_y[..., None]) - - return out diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/__init__.py deleted file mode 100644 index 08c22a2707fe55770a519db481954881c1cad26e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/__init__.py +++ /dev/null @@ -1,47 +0,0 @@ -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_stable_diffusion_adapter"] = ["StableDiffusionAdapterPipeline"] - _import_structure["pipeline_stable_diffusion_xl_adapter"] = ["StableDiffusionXLAdapterPipeline"] - - -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_stable_diffusion_adapter import StableDiffusionAdapterPipeline - from .pipeline_stable_diffusion_xl_adapter import StableDiffusionXLAdapterPipeline -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/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 75c6a0d06b12b69865d3b3e92f775095fb71644f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/__pycache__/pipeline_stable_diffusion_adapter.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/__pycache__/pipeline_stable_diffusion_adapter.cpython-310.pyc deleted file mode 100644 index f97402643f14b5fb324dae0db7f3518c3d0b97ef..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/__pycache__/pipeline_stable_diffusion_adapter.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/__pycache__/pipeline_stable_diffusion_xl_adapter.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/__pycache__/pipeline_stable_diffusion_xl_adapter.cpython-310.pyc deleted file mode 100644 index 12941edcb8c390da8b236f907a2a4d91dd31e5f9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/__pycache__/pipeline_stable_diffusion_xl_adapter.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py deleted file mode 100644 index 5c561721fcc76f860c762cb0347a37fa70c8a121..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py +++ /dev/null @@ -1,956 +0,0 @@ -# Copyright 2025 TencentARC 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, Callable, Dict, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer - -from ...image_processor import VaeImageProcessor -from ...loaders import FromSingleFileMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, MultiAdapter, T2IAdapter, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - PIL_INTERPOLATION, - USE_PEFT_BACKEND, - BaseOutput, - deprecate, - 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, StableDiffusionMixin -from ..stable_diffusion.safety_checker import StableDiffusionSafetyChecker - - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - - XLA_AVAILABLE = True -else: - XLA_AVAILABLE = False - - -@dataclass -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, - 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, or `None` if safety checking could not be performed. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] - nsfw_content_detected: Optional[List[bool]] - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> from PIL import Image - >>> from diffusers.utils import load_image - >>> import torch - >>> from diffusers import StableDiffusionAdapterPipeline, T2IAdapter - - >>> image = load_image( - ... "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/t2i-adapter/color_ref.png" - ... ) - - >>> color_palette = image.resize((8, 8)) - >>> color_palette = color_palette.resize((512, 512), resample=Image.Resampling.NEAREST) - - >>> adapter = T2IAdapter.from_pretrained("TencentARC/t2iadapter_color_sd14v1", torch_dtype=torch.float16) - >>> pipe = StableDiffusionAdapterPipeline.from_pretrained( - ... "CompVis/stable-diffusion-v1-4", - ... adapter=adapter, - ... torch_dtype=torch.float16, - ... ) - - >>> pipe.to("cuda") - - >>> out_image = pipe( - ... "At night, glowing cubes in front of the beach", - ... image=color_palette, - ... ).images[0] - ``` -""" - - -def _preprocess_adapter_image(image, height, width): - if isinstance(image, torch.Tensor): - return image - elif isinstance(image, PIL.Image.Image): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - image = [np.array(i.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])) for i in image] - image = [ - i[None, ..., None] if i.ndim == 2 else i[None, ...] for i in image - ] # expand [h, w] or [h, w, c] to [b, h, w, c] - image = np.concatenate(image, axis=0) - image = np.array(image).astype(np.float32) / 255.0 - image = image.transpose(0, 3, 1, 2) - image = torch.from_numpy(image) - elif isinstance(image[0], torch.Tensor): - if image[0].ndim == 3: - image = torch.stack(image, dim=0) - elif image[0].ndim == 4: - image = torch.cat(image, dim=0) - else: - raise ValueError( - f"Invalid image tensor! Expecting image tensor with 3 or 4 dimension, but receive: {image[0].ndim}" - ) - return 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, - **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 StableDiffusionAdapterPipeline(DiffusionPipeline, StableDiffusionMixin, FromSingleFileMixin): - r""" - Pipeline for text-to-image generation using Stable Diffusion augmented with T2I-Adapter - https://huggingface.co/papers/2302.08453 - - 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.) - - Args: - 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 - together. - 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`. - """ - - model_cpu_offload_seq = "text_encoder->adapter->unet->vae" - _optional_components = ["safety_checker", "feature_extractor"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - adapter: Union[T2IAdapter, MultiAdapter, List[T2IAdapter]], - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - - if safety_checker is not None and feature_extractor is None: - raise ValueError( - "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" - " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." - ) - - if isinstance(adapter, (list, tuple)): - adapter = MultiAdapter(adapter) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - adapter=adapter, - scheduler=scheduler, - safety_checker=safety_checker, - 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.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) - self.register_to_config(requires_safety_checker=requires_safety_checker) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - height, - width, - callback_steps, - image, - negative_prompt=None, - prompt_embeds=None, - negative_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 (callback_steps is None) or ( - 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 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 isinstance(self.adapter, MultiAdapter): - if not isinstance(image, list): - raise ValueError( - "MultiAdapter is enabled, but `image` is not a list. Please pass a list of images to `image`." - ) - - if len(image) != len(self.adapter.adapters): - raise ValueError( - f"MultiAdapter requires passing the same number of images as adapters. Given {len(image)} images and {len(self.adapter.adapters)} adapters." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def _default_height_width(self, height, width, image): - # NOTE: It is possible that a list of images have different - # dimensions for each image, so just checking the first image - # is not _exactly_ correct, but it is simple. - while isinstance(image, list): - image = image[0] - - if height is None: - if isinstance(image, PIL.Image.Image): - height = image.height - elif isinstance(image, torch.Tensor): - height = image.shape[-2] - - # round down to nearest multiple of `self.adapter.downscale_factor` - height = (height // self.adapter.downscale_factor) * self.adapter.downscale_factor - - if width is None: - if isinstance(image, PIL.Image.Image): - width = image.width - elif isinstance(image, torch.Tensor): - width = image.shape[-1] - - # round down to nearest multiple of `self.adapter.downscale_factor` - width = (width // self.adapter.downscale_factor) * self.adapter.downscale_factor - - return height, width - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 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, - image: Union[torch.Tensor, PIL.Image.Image, List[PIL.Image.Image]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - 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, - 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, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - adapter_conditioning_scale: Union[float, List[float]] = 1.0, - 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. - 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. - 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. - 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*): - 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): - 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. 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` 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 ge 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.StableDiffusionAdapterPipelineOutput`] 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. - cross_attention_kwargs (`dict`, *optional*): - 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): - 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. - 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. - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionAdapterPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionAdapterPipelineOutput`] 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, width = self._default_height_width(height, width, image) - device = self._execution_device - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, height, width, callback_steps, image, negative_prompt, prompt_embeds, negative_prompt_embeds - ) - - self._guidance_scale = guidance_scale - - if isinstance(self.adapter, MultiAdapter): - adapter_input = [] - - for one_image in image: - one_image = _preprocess_adapter_image(one_image, height, width) - one_image = one_image.to(device=device, dtype=self.adapter.dtype) - adapter_input.append(one_image) - else: - adapter_input = _preprocess_adapter_image(image, height, width) - adapter_input = adapter_input.to(device=device, dtype=self.adapter.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, - device, - num_images_per_prompt, - self.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 self.do_classifier_free_guidance: - 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, sigmas - ) - - # 5. 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, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 6.5 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) - - # 7. Denoising loop - if isinstance(self.adapter, MultiAdapter): - adapter_state = self.adapter(adapter_input, adapter_conditioning_scale) - for k, v in enumerate(adapter_state): - adapter_state[k] = v - else: - adapter_state = self.adapter(adapter_input) - for k, v in enumerate(adapter_state): - adapter_state[k] = v * adapter_conditioning_scale - if num_images_per_prompt > 1: - for k, v in enumerate(adapter_state): - adapter_state[k] = v.repeat(num_images_per_prompt, 1, 1, 1) - if self.do_classifier_free_guidance: - for k, v in enumerate(adapter_state): - adapter_state[k] = torch.cat([v] * 2, dim=0) - - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=cross_attention_kwargs, - down_intrablock_additional_residuals=[state.clone() for state in adapter_state], - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if output_type == "latent": - image = latents - has_nsfw_concept = None - elif output_type == "pil": - # 8. Post-processing - image = self.decode_latents(latents) - - # 9. Run safety checker - image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) - - # 10. Convert to PIL - image = self.numpy_to_pil(image) - else: - # 8. Post-processing - image = self.decode_latents(latents) - - # 9. Run safety checker - image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionAdapterPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py deleted file mode 100644 index 13183df47d4baef54118de17ef08d82632feb558..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py +++ /dev/null @@ -1,1311 +0,0 @@ -# Copyright 2025 TencentARC 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, Tuple, Union - -import numpy as np -import PIL.Image -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, ImageProjection, MultiAdapter, T2IAdapter, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - PIL_INTERPOLATION, - 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, StableDiffusionMixin -from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput - - -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 T2IAdapter, StableDiffusionXLAdapterPipeline, DDPMScheduler - >>> from diffusers.utils import load_image - - >>> sketch_image = load_image("https://huggingface.co/Adapter/t2iadapter/resolve/main/sketch.png").convert("L") - - >>> model_id = "stabilityai/stable-diffusion-xl-base-1.0" - - >>> adapter = T2IAdapter.from_pretrained( - ... "Adapter/t2iadapter", - ... subfolder="sketch_sdxl_1.0", - ... torch_dtype=torch.float16, - ... adapter_type="full_adapter_xl", - ... ) - >>> scheduler = DDPMScheduler.from_pretrained(model_id, subfolder="scheduler") - - >>> pipe = StableDiffusionXLAdapterPipeline.from_pretrained( - ... model_id, adapter=adapter, torch_dtype=torch.float16, variant="fp16", scheduler=scheduler - ... ).to("cuda") - - >>> generator = torch.manual_seed(42) - >>> sketch_image_out = pipe( - ... prompt="a photo of a dog in real world, high quality", - ... negative_prompt="extra digit, fewer digits, cropped, worst quality, low quality", - ... image=sketch_image, - ... generator=generator, - ... guidance_scale=7.5, - ... ).images[0] - ``` -""" - - -def _preprocess_adapter_image(image, height, width): - if isinstance(image, torch.Tensor): - return image - elif isinstance(image, PIL.Image.Image): - image = [image] - - if isinstance(image[0], PIL.Image.Image): - image = [np.array(i.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])) for i in image] - image = [ - i[None, ..., None] if i.ndim == 2 else i[None, ...] for i in image - ] # expand [h, w] or [h, w, c] to [b, h, w, c] - image = np.concatenate(image, axis=0) - image = np.array(image).astype(np.float32) / 255.0 - image = image.transpose(0, 3, 1, 2) - image = torch.from_numpy(image) - elif isinstance(image[0], torch.Tensor): - if image[0].ndim == 3: - image = torch.stack(image, dim=0) - elif image[0].ndim == 4: - image = torch.cat(image, dim=0) - else: - raise ValueError( - f"Invalid image tensor! Expecting image tensor with 3 or 4 dimension, but receive: {image[0].ndim}" - ) - return image - - -# 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 - - -# 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 StableDiffusionXLAdapterPipeline( - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionXLLoraLoaderMixin, - IPAdapterMixin, - FromSingleFileMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion augmented with T2I-Adapter - https://huggingface.co/papers/2302.08453 - - 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: - 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 - together. - 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`. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "feature_extractor", - "image_encoder", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - adapter: Union[T2IAdapter, MultiAdapter, List[T2IAdapter]], - scheduler: KarrasDiffusionSchedulers, - force_zeros_for_empty_prompt: bool = True, - feature_extractor: CLIPImageProcessor = None, - image_encoder: CLIPVisionModelWithProjection = None, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - adapter=adapter, - scheduler=scheduler, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - ) - 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 - ) - - # 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 - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - 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) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_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) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.check_inputs - def check_inputs( - self, - prompt, - prompt_2, - height, - width, - callback_steps, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - ip_adapter_image=None, - ip_adapter_image_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_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`." - ) - - 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - 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.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.t2i_adapter.pipeline_stable_diffusion_adapter.StableDiffusionAdapterPipeline._default_height_width - def _default_height_width(self, height, width, image): - # NOTE: It is possible that a list of images have different - # dimensions for each image, so just checking the first image - # is not _exactly_ correct, but it is simple. - while isinstance(image, list): - image = image[0] - - if height is None: - if isinstance(image, PIL.Image.Image): - height = image.height - elif isinstance(image, torch.Tensor): - height = image.shape[-2] - - # round down to nearest multiple of `self.adapter.downscale_factor` - height = (height // self.adapter.downscale_factor) * self.adapter.downscale_factor - - if width is None: - if isinstance(image, PIL.Image.Image): - width = image.width - elif isinstance(image, torch.Tensor): - width = image.shape[-1] - - # round down to nearest multiple of `self.adapter.downscale_factor` - width = (width // self.adapter.downscale_factor) * self.adapter.downscale_factor - - return height, width - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @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 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, - image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = 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, - 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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, - adapter_conditioning_factor: float = 1.0, - 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 - 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. - height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The height in pixels of the generated image. 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. 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. - 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*): - 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. - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output) - 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. - 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*): - 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 ge 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. - 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)`. 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. - 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.StableDiffusionAdapterPipelineOutput`] - 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. - cross_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). - 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. - 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). - 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. - 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. - adapter_conditioning_factor (`float`, *optional*, defaults to 1.0): - The fraction of timesteps for which adapter should be applied. If `adapter_conditioning_factor` is - `0.0`, adapter is not applied at all. If `adapter_conditioning_factor` is `1.0`, adapter is applied for - all timesteps. If `adapter_conditioning_factor` is `0.5`, adapter is applied for half of the timesteps. - 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. - - Examples: - - Returns: - [`~pipelines.stable_diffusion.StableDiffusionAdapterPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionAdapterPipelineOutput`] 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, width = self._default_height_width(height, width, image) - device = self._execution_device - - if isinstance(self.adapter, MultiAdapter): - adapter_input = [] - - for one_image in image: - one_image = _preprocess_adapter_image(one_image, height, width) - one_image = one_image.to(device=device, dtype=self.adapter.dtype) - adapter_input.append(one_image) - else: - adapter_input = _preprocess_adapter_image(image, height, width) - adapter_input = adapter_input.to(device=device, dtype=self.adapter.dtype) - 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, - callback_steps, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ip_adapter_image, - ip_adapter_image_embeds, - ) - - self._guidance_scale = guidance_scale - - # 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.1 Encode input prompt - ( - 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, - clip_skip=clip_skip, - ) - - # 3.2 Encode ip_adapter_image - 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, - self.do_classifier_free_guidance, - ) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas - ) - - # 5. 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, - ) - - # 6.1 Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 6.2 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) - - # 7. Prepare added time ids & embeddings & adapter features - if isinstance(self.adapter, MultiAdapter): - adapter_state = self.adapter(adapter_input, adapter_conditioning_scale) - for k, v in enumerate(adapter_state): - adapter_state[k] = v - else: - adapter_state = self.adapter(adapter_input) - for k, v in enumerate(adapter_state): - adapter_state[k] = v * adapter_conditioning_scale - if num_images_per_prompt > 1: - for k, v in enumerate(adapter_state): - adapter_state[k] = v.repeat(num_images_per_prompt, 1, 1, 1) - if self.do_classifier_free_guidance: - for k, v in enumerate(adapter_state): - adapter_state[k] = torch.cat([v] * 2, dim=0) - - 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) - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - # Apply denoising_end - if denoising_end is not None and isinstance(denoising_end, float) and denoising_end > 0 and denoising_end < 1: - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - - if ip_adapter_image is not None or ip_adapter_image_embeds is not None: - added_cond_kwargs["image_embeds"] = image_embeds - - # predict the noise residual - if i < int(num_inference_steps * adapter_conditioning_factor): - down_intrablock_additional_residuals = [state.clone() for state in adapter_state] - else: - down_intrablock_additional_residuals = None - - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - timestep_cond=timestep_cond, - cross_attention_kwargs=cross_attention_kwargs, - down_intrablock_additional_residuals=down_intrablock_additional_residuals, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - if self.do_classifier_free_guidance and guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - 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 - return StableDiffusionXLPipelineOutput(images=image) - - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__init__.py deleted file mode 100644 index 8d8fdb92769bb3dcb2dd7696115b16f197062262..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__init__.py +++ /dev/null @@ -1,54 +0,0 @@ -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_output"] = ["TextToVideoSDPipelineOutput"] - _import_structure["pipeline_text_to_video_synth"] = ["TextToVideoSDPipeline"] - _import_structure["pipeline_text_to_video_synth_img2img"] = ["VideoToVideoSDPipeline"] - _import_structure["pipeline_text_to_video_zero"] = ["TextToVideoZeroPipeline"] - _import_structure["pipeline_text_to_video_zero_sdxl"] = ["TextToVideoZeroSDXLPipeline"] - - -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_output import TextToVideoSDPipelineOutput - from .pipeline_text_to_video_synth import TextToVideoSDPipeline - from .pipeline_text_to_video_synth_img2img import VideoToVideoSDPipeline - from .pipeline_text_to_video_zero import TextToVideoZeroPipeline - from .pipeline_text_to_video_zero_sdxl import TextToVideoZeroSDXLPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 8087ce05d888efa31510f5bb4e5112b40002ba54..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index 9dea9dfb831a3e8c93a4bcb8217b452e6bb6a992..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_synth.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_synth.cpython-310.pyc deleted file mode 100644 index 8a1578c1bad4e59fa9d62f18ae2554958e8b985b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_synth.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_synth_img2img.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_synth_img2img.cpython-310.pyc deleted file mode 100644 index 2bf2493e91c51dceb09f12c6539eaa048dc51f5c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_synth_img2img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_zero.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_zero.cpython-310.pyc deleted file mode 100644 index 4f019e02fceb1af5dcc81f6c6e242b364c603d47..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_zero.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_zero_sdxl.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_zero_sdxl.cpython-310.pyc deleted file mode 100644 index 640486f7064e9e179e7ffff330931b214fb0d574..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/__pycache__/pipeline_text_to_video_zero_sdxl.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_output.py deleted file mode 100644 index 040bf0efba8465d830af5e6686bc3f544585f63f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_output.py +++ /dev/null @@ -1,26 +0,0 @@ -from dataclasses import dataclass -from typing import List, Union - -import numpy as np -import PIL -import torch - -from ...utils import ( - BaseOutput, -) - - -@dataclass -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 - 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]]] diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py deleted file mode 100644 index 3ce7b4d1990fb5fcf9a96f6beb2934d653e52e30..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py +++ /dev/null @@ -1,660 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import torch -from transformers import CLIPTextModel, CLIPTokenizer - -from ...loaders import StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, UNet3DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import randn_tensor -from ...video_processor import VideoProcessor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from . import TextToVideoSDPipelineOutput - - -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 TextToVideoSDPipeline - >>> from diffusers.utils import export_to_video - - >>> pipe = TextToVideoSDPipeline.from_pretrained( - ... "damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16, variant="fp16" - ... ) - >>> pipe.enable_model_cpu_offload() - - >>> prompt = "Spiderman is surfing" - >>> video_frames = pipe(prompt).frames[0] - >>> video_path = export_to_video(video_frames) - >>> video_path - ``` -""" - - -class TextToVideoSDPipeline( - DeprecatedPipelineMixin, - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, -): - _last_supported_version = "0.33.1" - r""" - Pipeline for text-to-video generation. - - 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.). - - 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 - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer (`CLIPTokenizer`): - A [`~transformers.CLIPTokenizer`] to tokenize text. - unet ([`UNet3DConditionModel`]): - A [`UNet3DConditionModel`] to denoise the encoded video 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`]. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet3DConditionModel, - scheduler: KarrasDiffusionSchedulers, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - self.video_processor = VideoProcessor(do_resize=False, vae_scale_factor=self.vae_scale_factor) - - # 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 - - def decode_latents(self, latents): - latents = 1 / self.vae.config.scaling_factor * latents - - batch_size, channels, num_frames, height, width = latents.shape - latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) - - image = self.vae.decode(latents).sample - video = image[None, :].reshape((batch_size, num_frames, -1) + image.shape[2:]).permute(0, 2, 1, 3, 4) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - video = video.float() - return video - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - 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, - 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, num_frames, height, width, dtype, device, generator, latents=None - ): - shape = ( - batch_size, - num_channels_latents, - num_frames, - height // self.vae_scale_factor, - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_frames: int = 16, - num_inference_steps: int = 50, - guidance_scale: float = 9.0, - negative_prompt: Optional[Union[str, List[str]]] = 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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. - width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): - The width in pixels of the generated video. - num_frames (`int`, *optional*, defaults to 16): - The number of video frames that are generated. Defaults to 16 frames which at 8 frames per seconds - amounts to 2 seconds of video. - 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. - 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*): - 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): - The number of images to generate per prompt. - 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*): - 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 video - 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`. Latents should be of shape - `(batch_size, num_channel, num_frames, height, width)`. - 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. - 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. - output_type (`str`, *optional*, defaults to `"np"`): - The output format of the generated video. Choose between `torch.Tensor` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.text_to_video_synthesis.TextToVideoSDPipelineOutput`] instead - of a plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - 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. - Examples: - - Returns: - [`~pipelines.text_to_video_synthesis.TextToVideoSDPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.text_to_video_synthesis.TextToVideoSDPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is a list with the generated frames. - """ - # 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 - - num_images_per_prompt = 1 - - # 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - 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, - lora_scale=text_encoder_lora_scale, - 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=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - num_frames, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # reshape latents - bsz, channel, frames, width, height = latents.shape - latents = latents.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height) - noise_pred = noise_pred.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # reshape latents back - latents = latents[None, :].reshape(bsz, frames, channel, width, height).permute(0, 2, 1, 3, 4) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # 8. Post processing - if output_type == "latent": - video = latents - else: - video_tensor = self.decode_latents(latents) - video = self.video_processor.postprocess_video(video=video_tensor, output_type=output_type) - - # 9. Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return TextToVideoSDPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py deleted file mode 100644 index 9d0b7e3dbc32ac782fcaa683dde74c74d3f87b19..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py +++ /dev/null @@ -1,716 +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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import numpy as np -import torch -from transformers import CLIPTextModel, CLIPTokenizer - -from ...loaders import StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, UNet3DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_torch_xla_available, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import randn_tensor -from ...video_processor import VideoProcessor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from . import TextToVideoSDPipelineOutput - - -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 DiffusionPipeline, DPMSolverMultistepScheduler - >>> from diffusers.utils import export_to_video - - >>> pipe = DiffusionPipeline.from_pretrained("cerspense/zeroscope_v2_576w", torch_dtype=torch.float16) - >>> pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config) - >>> pipe.to("cuda") - - >>> prompt = "spiderman running in the desert" - >>> video_frames = pipe(prompt, num_inference_steps=40, height=320, width=576, num_frames=24).frames[0] - >>> # safe low-res video - >>> video_path = export_to_video(video_frames, output_video_path="./video_576_spiderman.mp4") - - >>> # let's offload the text-to-image model - >>> pipe.to("cpu") - - >>> # and load the image-to-image model - >>> pipe = DiffusionPipeline.from_pretrained( - ... "cerspense/zeroscope_v2_XL", torch_dtype=torch.float16, revision="refs/pr/15" - ... ) - >>> pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config) - >>> pipe.enable_model_cpu_offload() - - >>> # The VAE consumes A LOT of memory, let's make sure we run it in sliced mode - >>> pipe.vae.enable_slicing() - - >>> # now let's upscale it - >>> video = [Image.fromarray(frame).resize((1024, 576)) for frame in video_frames] - - >>> # and denoise it - >>> video_frames = pipe(prompt, video=video, strength=0.6).frames[0] - >>> video_path = export_to_video(video_frames, output_video_path="./video_1024_spiderman.mp4") - >>> video_path - ``` -""" - - -# 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" -): - 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 VideoToVideoSDPipeline( - DeprecatedPipelineMixin, - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, -): - _last_supported_version = "0.33.1" - r""" - Pipeline for text-guided video-to-video generation. - - 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.). - - 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 - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. - text_encoder ([`CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer (`CLIPTokenizer`): - A [`~transformers.CLIPTokenizer`] to tokenize text. - unet ([`UNet3DConditionModel`]): - A [`UNet3DConditionModel`] to denoise the encoded video 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`]. - """ - - model_cpu_offload_seq = "text_encoder->unet->vae" - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet3DConditionModel, - scheduler: KarrasDiffusionSchedulers, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - self.video_processor = VideoProcessor(do_resize=False, vae_scale_factor=self.vae_scale_factor) - - # 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.text_to_video_synthesis.pipeline_text_to_video_synth.TextToVideoSDPipeline.decode_latents - def decode_latents(self, latents): - latents = 1 / self.vae.config.scaling_factor * latents - - batch_size, channels, num_frames, height, width = latents.shape - latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) - - image = self.vae.decode(latents).sample - video = image[None, :].reshape((batch_size, num_frames, -1) + image.shape[2:]).permute(0, 2, 1, 3, 4) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - video = video.float() - return video - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - strength, - callback_steps, - 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 in [0.0, 1.0] but is {strength}") - - 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}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = 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_latents(self, video, timestep, batch_size, dtype, device, generator=None): - video = video.to(device=device, dtype=dtype) - - # change from (b, c, f, h, w) -> (b * f, c, w, h) - bsz, channel, frames, width, height = video.shape - video = video.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height) - - if video.shape[1] == 4: - init_latents = video - else: - 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." - ) - elif isinstance(generator, list): - init_latents = [ - retrieve_latents(self.vae.encode(video[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(video), generator=generator) - - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `video` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - latents = init_latents - - latents = latents[None, :].reshape((bsz, frames, latents.shape[1]) + latents.shape[2:]).permute(0, 2, 1, 3, 4) - - return latents - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - video: Union[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, - 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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` 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): - Indicates extent to transform the reference `video`. Must be between 0 and 1. `video` is 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 is maximum and the - denoising process runs for the full number of iterations specified in `num_inference_steps`. A value of - 1 essentially ignores `video`. - 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. - 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*): - 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*): - 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 video - 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`. Latents should be of shape - `(batch_size, num_channel, num_frames, height, width)`. - 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. - 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. - output_type (`str`, *optional*, defaults to `"np"`): - The output format of the generated video. Choose between `torch.Tensor` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.text_to_video_synthesis.TextToVideoSDPipelineOutput`] instead - of a plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - 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). - 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. - Examples: - - Returns: - [`~pipelines.text_to_video_synthesis.TextToVideoSDPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.text_to_video_synthesis.TextToVideoSDPipelineOutput`] is - returned, otherwise a `tuple` is returned where the first element is a list with the generated frames. - """ - # 0. Default height and width to unet - num_images_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs(prompt, strength, 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - 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, - lora_scale=text_encoder_lora_scale, - 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. Preprocess video - video = self.video_processor.preprocess_video(video) - - # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - - # 6. Prepare latent variables - latents = self.prepare_latents(video, latent_timestep, batch_size, prompt_embeds.dtype, device, generator) - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # reshape latents - bsz, channel, frames, width, height = latents.shape - latents = latents.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height) - noise_pred = noise_pred.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # reshape latents back - latents = latents[None, :].reshape(bsz, frames, channel, width, height).permute(0, 2, 1, 3, 4) - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # manually for max memory savings - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.unet.to("cpu") - - # 9. Post processing - if output_type == "latent": - video = latents - else: - video_tensor = self.decode_latents(latents) - video = self.video_processor.postprocess_video(video=video_tensor, output_type=output_type) - - # 10. Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return TextToVideoSDPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py deleted file mode 100644 index 96316f8e91e59d3a8f33f778f494654fee63dbc4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py +++ /dev/null @@ -1,1005 +0,0 @@ -import copy -import inspect -from dataclasses import dataclass -from typing import Callable, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -import torch.nn.functional as F -from torch.nn.functional import grid_sample -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer - -from ...image_processor import VaeImageProcessor -from ...loaders import FromSingleFileMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - BaseOutput, - is_torch_xla_available, - logging, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import empty_device_cache, randn_tensor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion import StableDiffusionSafetyChecker - - -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 rearrange_0(tensor, f): - F, C, H, W = tensor.size() - tensor = torch.permute(torch.reshape(tensor, (F // f, f, C, H, W)), (0, 2, 1, 3, 4)) - return tensor - - -def rearrange_1(tensor): - B, C, F, H, W = tensor.size() - return torch.reshape(torch.permute(tensor, (0, 2, 1, 3, 4)), (B * F, C, H, W)) - - -def rearrange_3(tensor, f): - F, D, C = tensor.size() - return torch.reshape(tensor, (F // f, f, D, C)) - - -def rearrange_4(tensor): - B, F, D, C = tensor.size() - return torch.reshape(tensor, (B * F, D, C)) - - -class CrossFrameAttnProcessor: - """ - Cross frame attention processor. Each frame attends the first frame. - - Args: - batch_size: The number that represents actual batch size, other than the frames. - For example, calling unet with a single prompt and num_images_per_prompt=1, batch_size should be equal to - 2, due to classifier-free guidance. - """ - - def __init__(self, batch_size=2): - self.batch_size = batch_size - - def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None): - batch_size, sequence_length, _ = hidden_states.shape - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - query = attn.to_q(hidden_states) - - is_cross_attention = encoder_hidden_states is not None - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - # Cross Frame Attention - if not is_cross_attention: - video_length = key.size()[0] // self.batch_size - first_frame_index = [0] * video_length - - # rearrange keys to have batch and frames in the 1st and 2nd dims respectively - key = rearrange_3(key, video_length) - key = key[:, first_frame_index] - # rearrange values to have batch and frames in the 1st and 2nd dims respectively - value = rearrange_3(value, video_length) - value = value[:, first_frame_index] - - # rearrange back to original shape - key = rearrange_4(key) - value = rearrange_4(value) - - query = attn.head_to_batch_dim(query) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - - attention_probs = attn.get_attention_scores(query, key, attention_mask) - hidden_states = torch.bmm(attention_probs, value) - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - return hidden_states - - -class CrossFrameAttnProcessor2_0: - """ - Cross frame attention processor with scaled_dot_product attention of Pytorch 2.0. - - Args: - batch_size: The number that represents actual batch size, other than the frames. - For example, calling unet with a single prompt and num_images_per_prompt=1, batch_size should be equal to - 2, due to classifier-free guidance. - """ - - def __init__(self, batch_size=2): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") - self.batch_size = batch_size - - def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None): - batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - inner_dim = hidden_states.shape[-1] - - if attention_mask is not None: - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - query = attn.to_q(hidden_states) - - is_cross_attention = encoder_hidden_states is not None - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - # Cross Frame Attention - if not is_cross_attention: - video_length = max(1, key.size()[0] // self.batch_size) - first_frame_index = [0] * video_length - - # rearrange keys to have batch and frames in the 1st and 2nd dims respectively - key = rearrange_3(key, video_length) - key = key[:, first_frame_index] - # rearrange values to have batch and frames in the 1st and 2nd dims respectively - value = rearrange_3(value, video_length) - value = value[:, first_frame_index] - - # rearrange back to original shape - key = rearrange_4(key) - value = rearrange_4(value) - - head_dim = inner_dim // attn.heads - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - return hidden_states - - -@dataclass -class TextToVideoPipelineOutput(BaseOutput): - r""" - 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, - num_channels)`. - 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]] - - -def coords_grid(batch, ht, wd, device): - # Adapted from https://github.com/princeton-vl/RAFT/blob/master/core/utils/utils.py - coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device)) - coords = torch.stack(coords[::-1], dim=0).float() - return coords[None].repeat(batch, 1, 1, 1) - - -def warp_single_latent(latent, reference_flow): - """ - Warp latent of a single frame with given flow - - Args: - latent: latent code of a single frame - reference_flow: flow which to warp the latent with - - Returns: - warped: warped latent - """ - _, _, H, W = reference_flow.size() - _, _, h, w = latent.size() - coords0 = coords_grid(1, H, W, device=latent.device).to(latent.dtype) - - coords_t0 = coords0 + reference_flow - coords_t0[:, 0] /= W - coords_t0[:, 1] /= H - - coords_t0 = coords_t0 * 2.0 - 1.0 - coords_t0 = F.interpolate(coords_t0, size=(h, w), mode="bilinear") - coords_t0 = torch.permute(coords_t0, (0, 2, 3, 1)) - - warped = grid_sample(latent, coords_t0, mode="nearest", padding_mode="reflection") - return warped - - -def create_motion_field(motion_field_strength_x, motion_field_strength_y, frame_ids, device, dtype): - """ - Create translation motion field - - Args: - motion_field_strength_x: motion strength along x-axis - motion_field_strength_y: motion strength along y-axis - frame_ids: indexes of the frames the latents of which are being processed. - This is needed when we perform chunk-by-chunk inference - device: device - dtype: dtype - - Returns: - - """ - seq_length = len(frame_ids) - reference_flow = torch.zeros((seq_length, 2, 512, 512), device=device, dtype=dtype) - for fr_idx in range(seq_length): - reference_flow[fr_idx, 0, :, :] = motion_field_strength_x * (frame_ids[fr_idx]) - reference_flow[fr_idx, 1, :, :] = motion_field_strength_y * (frame_ids[fr_idx]) - return reference_flow - - -def create_motion_field_and_warp_latents(motion_field_strength_x, motion_field_strength_y, frame_ids, latents): - """ - Creates translation motion and warps the latents accordingly - - Args: - motion_field_strength_x: motion strength along x-axis - motion_field_strength_y: motion strength along y-axis - frame_ids: indexes of the frames the latents of which are being processed. - This is needed when we perform chunk-by-chunk inference - latents: latent codes of frames - - Returns: - warped_latents: warped latents - """ - motion_field = create_motion_field( - motion_field_strength_x=motion_field_strength_x, - motion_field_strength_y=motion_field_strength_y, - frame_ids=frame_ids, - device=latents.device, - dtype=latents.dtype, - ) - warped_latents = latents.clone().detach() - for i in range(len(warped_latents)): - warped_latents[i] = warp_single_latent(latents[i][None], motion_field[i][None]) - return warped_latents - - -class TextToVideoZeroPipeline( - DeprecatedPipelineMixin, - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, - FromSingleFileMixin, -): - _last_supported_version = "0.33.1" - r""" - Pipeline for zero-shot text-to-video generation using Stable Diffusion. - - 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: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - tokenizer (`CLIPTokenizer`): - A [`~transformers.CLIPTokenizer`] to tokenize text. - unet ([`UNet2DConditionModel`]): - A [`UNet3DConditionModel`] to denoise the encoded video 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 - more details about a model's potential harms. - feature_extractor ([`CLIPImageProcessor`]): - A [`CLIPImageProcessor`] to extract features from generated images; used as inputs to the `safety_checker`. - """ - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) - if safety_checker is None and requires_safety_checker: - logger.warning( - f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" - " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" - " results in services or applications open to the public. Both the diffusers team and Hugging Face" - " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" - " it only for use-cases that involve analyzing network behavior or auditing its results. For more" - " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." - ) - 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) - - def forward_loop(self, x_t0, t0, t1, generator): - """ - Perform DDPM forward process from time t0 to t1. This is the same as adding noise with corresponding variance. - - Args: - x_t0: - Latent code at time t0. - t0: - Timestep at t0. - t1: - Timestamp at t1. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - - Returns: - x_t1: - Forward process applied to x_t0 from time t0 to t1. - """ - eps = randn_tensor(x_t0.size(), generator=generator, dtype=x_t0.dtype, device=x_t0.device) - alpha_vec = torch.prod(self.scheduler.alphas[t0:t1]) - x_t1 = torch.sqrt(alpha_vec) * x_t0 + torch.sqrt(1 - alpha_vec) * eps - return x_t1 - - def backward_loop( - self, - latents, - timesteps, - prompt_embeds, - guidance_scale, - callback, - callback_steps, - num_warmup_steps, - extra_step_kwargs, - cross_attention_kwargs=None, - ): - """ - Perform backward process given list of time steps. - - Args: - latents: - Latents at time timesteps[0]. - timesteps: - Time steps along which to perform backward process. - prompt_embeds: - Pre-generated text embeddings. - guidance_scale: - 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`. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - extra_step_kwargs: - Extra_step_kwargs. - cross_attention_kwargs: - 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). - num_warmup_steps: - number of warmup steps. - - Returns: - latents: - Latents of backward process output at time timesteps[-1]. - """ - do_classifier_free_guidance = guidance_scale > 1.0 - num_steps = (len(timesteps) - num_warmup_steps) // self.scheduler.order - with self.progress_bar(total=num_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - ).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - 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, - 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}." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - @torch.no_grad() - def __call__( - self, - prompt: Union[str, List[str]], - video_length: Optional[int] = 8, - 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_videos_per_prompt: Optional[int] = 1, - eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - motion_field_strength_x: float = 12, - motion_field_strength_y: float = 12, - output_type: Optional[str] = "tensor", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: Optional[int] = 1, - t0: int = 44, - t1: int = 47, - frame_ids: Optional[List[int]] = None, - ): - """ - The call function to the pipeline for generation. - - Args: - 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. - 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): - 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*): - 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): - The number of videos to generate per prompt. - 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*): - 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 video - 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`. - output_type (`str`, *optional*, defaults to `"np"`): - The output format of the generated video. Choose between `"latent"` and `"np"`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a - [`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.TextToVideoPipelineOutput`] instead of - a plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - motion_field_strength_x (`float`, *optional*, defaults to 12): - Strength of motion in generated video along x-axis. See the - [paper](https://huggingface.co/papers/2303.13439), Sect. 3.3.1. - motion_field_strength_y (`float`, *optional*, defaults to 12): - Strength of motion in generated video along y-axis. See the - [paper](https://huggingface.co/papers/2303.13439), Sect. 3.3.1. - t0 (`int`, *optional*, defaults to 44): - Timestep t0. Should be in the range [0, num_inference_steps - 1]. See the - [paper](https://huggingface.co/papers/2303.13439), Sect. 3.3.1. - 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*): - Indexes of the frames that are being generated. This is used when generating longer videos - chunk-by-chunk. - - Returns: - [`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.TextToVideoPipelineOutput`]: - The output contains a `ndarray` of the generated video, when `output_type` != `"latent"`, otherwise a - latent code of generated videos and a list of `bool`s indicating whether the corresponding generated - video contains "not-safe-for-work" (nsfw) content.. - """ - assert video_length > 0 - if frame_ids is None: - frame_ids = list(range(video_length)) - assert len(frame_ids) == video_length - - assert num_videos_per_prompt == 1 - - # set the processor - original_attn_proc = self.unet.attn_processors - processor = ( - CrossFrameAttnProcessor2_0(batch_size=2) - if hasattr(F, "scaled_dot_product_attention") - else CrossFrameAttnProcessor(batch_size=2) - ) - self.unet.set_attn_processor(processor) - - if isinstance(prompt, str): - prompt = [prompt] - if isinstance(negative_prompt, str): - negative_prompt = [negative_prompt] - - # 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 - - # Check inputs. Raise error if not correct - self.check_inputs(prompt, height, width, callback_steps) - - # Define call parameters - batch_size = 1 if isinstance(prompt, str) else len(prompt) - 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 = guidance_scale > 1.0 - - # Encode input prompt - prompt_embeds_tuple = self.encode_prompt( - prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt - ) - prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]]) - - # Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - # Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - # Perform the first backward process up to time T_1 - x_1_t1 = self.backward_loop( - timesteps=timesteps[: -t1 - 1], - prompt_embeds=prompt_embeds, - latents=latents, - guidance_scale=guidance_scale, - callback=callback, - callback_steps=callback_steps, - extra_step_kwargs=extra_step_kwargs, - num_warmup_steps=num_warmup_steps, - ) - scheduler_copy = copy.deepcopy(self.scheduler) - - # Perform the second backward process up to time T_0 - x_1_t0 = self.backward_loop( - timesteps=timesteps[-t1 - 1 : -t0 - 1], - prompt_embeds=prompt_embeds, - latents=x_1_t1, - guidance_scale=guidance_scale, - callback=callback, - callback_steps=callback_steps, - extra_step_kwargs=extra_step_kwargs, - num_warmup_steps=0, - ) - - # Propagate first frame latents at time T_0 to remaining frames - x_2k_t0 = x_1_t0.repeat(video_length - 1, 1, 1, 1) - - # Add motion in latents at time T_0 - x_2k_t0 = create_motion_field_and_warp_latents( - motion_field_strength_x=motion_field_strength_x, - motion_field_strength_y=motion_field_strength_y, - latents=x_2k_t0, - frame_ids=frame_ids[1:], - ) - - # Perform forward process up to time T_1 - x_2k_t1 = self.forward_loop( - x_t0=x_2k_t0, - t0=timesteps[-t0 - 1].item(), - t1=timesteps[-t1 - 1].item(), - generator=generator, - ) - - # Perform backward process from time T_1 to 0 - x_1k_t1 = torch.cat([x_1_t1, x_2k_t1]) - b, l, d = prompt_embeds.size() - prompt_embeds = prompt_embeds[:, None].repeat(1, video_length, 1, 1).reshape(b * video_length, l, d) - - self.scheduler = scheduler_copy - x_1k_0 = self.backward_loop( - timesteps=timesteps[-t1 - 1 :], - prompt_embeds=prompt_embeds, - latents=x_1k_t1, - guidance_scale=guidance_scale, - callback=callback, - callback_steps=callback_steps, - extra_step_kwargs=extra_step_kwargs, - num_warmup_steps=0, - ) - latents = x_1k_0 - - # manually for max memory savings - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.unet.to("cpu") - empty_device_cache() - - if output_type == "latent": - image = latents - has_nsfw_concept = None - else: - image = self.decode_latents(latents) - # Run safety checker - image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) - - # Offload all models - self.maybe_free_model_hooks() - # make sure to set the original attention processors back - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image, has_nsfw_concept) - - return TextToVideoPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) - - # 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.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # 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 - - def decode_latents(self, latents): - 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 diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py deleted file mode 100644 index a9fa43c1f5c528b2393098d1854fb46750637f95..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py +++ /dev/null @@ -1,1354 +0,0 @@ -import copy -import inspect -from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL -import torch -import torch.nn.functional as F -from torch.nn.functional import grid_sample -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ...image_processor import VaeImageProcessor -from ...loaders import StableDiffusionXLLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL, UNet2DConditionModel -from ...models.attention_processor import ( - AttnProcessor2_0, - FusedAttnProcessor2_0, - XFormersAttnProcessor, -) -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - BaseOutput, - is_invisible_watermark_available, - logging, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin - - -if is_invisible_watermark_available(): - from ..stable_diffusion_xl.watermark import StableDiffusionXLWatermarker - - -from ...utils import is_torch_xla_available - - -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 - - -# Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.rearrange_0 -def rearrange_0(tensor, f): - F, C, H, W = tensor.size() - tensor = torch.permute(torch.reshape(tensor, (F // f, f, C, H, W)), (0, 2, 1, 3, 4)) - return tensor - - -# Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.rearrange_1 -def rearrange_1(tensor): - B, C, F, H, W = tensor.size() - return torch.reshape(torch.permute(tensor, (0, 2, 1, 3, 4)), (B * F, C, H, W)) - - -# Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.rearrange_3 -def rearrange_3(tensor, f): - F, D, C = tensor.size() - return torch.reshape(tensor, (F // f, f, D, C)) - - -# Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.rearrange_4 -def rearrange_4(tensor): - B, F, D, C = tensor.size() - return torch.reshape(tensor, (B * F, D, C)) - - -# Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.CrossFrameAttnProcessor -class CrossFrameAttnProcessor: - """ - Cross frame attention processor. Each frame attends the first frame. - - Args: - batch_size: The number that represents actual batch size, other than the frames. - For example, calling unet with a single prompt and num_images_per_prompt=1, batch_size should be equal to - 2, due to classifier-free guidance. - """ - - def __init__(self, batch_size=2): - self.batch_size = batch_size - - def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None): - batch_size, sequence_length, _ = hidden_states.shape - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - query = attn.to_q(hidden_states) - - is_cross_attention = encoder_hidden_states is not None - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - # Cross Frame Attention - if not is_cross_attention: - video_length = key.size()[0] // self.batch_size - first_frame_index = [0] * video_length - - # rearrange keys to have batch and frames in the 1st and 2nd dims respectively - key = rearrange_3(key, video_length) - key = key[:, first_frame_index] - # rearrange values to have batch and frames in the 1st and 2nd dims respectively - value = rearrange_3(value, video_length) - value = value[:, first_frame_index] - - # rearrange back to original shape - key = rearrange_4(key) - value = rearrange_4(value) - - query = attn.head_to_batch_dim(query) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - - attention_probs = attn.get_attention_scores(query, key, attention_mask) - hidden_states = torch.bmm(attention_probs, value) - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - return hidden_states - - -# Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.CrossFrameAttnProcessor2_0 -class CrossFrameAttnProcessor2_0: - """ - Cross frame attention processor with scaled_dot_product attention of Pytorch 2.0. - - Args: - batch_size: The number that represents actual batch size, other than the frames. - For example, calling unet with a single prompt and num_images_per_prompt=1, batch_size should be equal to - 2, due to classifier-free guidance. - """ - - def __init__(self, batch_size=2): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") - self.batch_size = batch_size - - def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None): - batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - inner_dim = hidden_states.shape[-1] - - if attention_mask is not None: - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - query = attn.to_q(hidden_states) - - is_cross_attention = encoder_hidden_states is not None - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - # Cross Frame Attention - if not is_cross_attention: - video_length = max(1, key.size()[0] // self.batch_size) - first_frame_index = [0] * video_length - - # rearrange keys to have batch and frames in the 1st and 2nd dims respectively - key = rearrange_3(key, video_length) - key = key[:, first_frame_index] - # rearrange values to have batch and frames in the 1st and 2nd dims respectively - value = rearrange_3(value, video_length) - value = value[:, first_frame_index] - - # rearrange back to original shape - key = rearrange_4(key) - value = rearrange_4(value) - - head_dim = inner_dim // attn.heads - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - return hidden_states - - -@dataclass -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, - num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - """ - - images: Union[List[PIL.Image.Image], np.ndarray] - - -# Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.coords_grid -def coords_grid(batch, ht, wd, device): - # Adapted from https://github.com/princeton-vl/RAFT/blob/master/core/utils/utils.py - coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device)) - coords = torch.stack(coords[::-1], dim=0).float() - return coords[None].repeat(batch, 1, 1, 1) - - -# Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.warp_single_latent -def warp_single_latent(latent, reference_flow): - """ - Warp latent of a single frame with given flow - - Args: - latent: latent code of a single frame - reference_flow: flow which to warp the latent with - - Returns: - warped: warped latent - """ - _, _, H, W = reference_flow.size() - _, _, h, w = latent.size() - coords0 = coords_grid(1, H, W, device=latent.device).to(latent.dtype) - - coords_t0 = coords0 + reference_flow - coords_t0[:, 0] /= W - coords_t0[:, 1] /= H - - coords_t0 = coords_t0 * 2.0 - 1.0 - coords_t0 = F.interpolate(coords_t0, size=(h, w), mode="bilinear") - coords_t0 = torch.permute(coords_t0, (0, 2, 3, 1)) - - warped = grid_sample(latent, coords_t0, mode="nearest", padding_mode="reflection") - return warped - - -# Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.create_motion_field -def create_motion_field(motion_field_strength_x, motion_field_strength_y, frame_ids, device, dtype): - """ - Create translation motion field - - Args: - motion_field_strength_x: motion strength along x-axis - motion_field_strength_y: motion strength along y-axis - frame_ids: indexes of the frames the latents of which are being processed. - This is needed when we perform chunk-by-chunk inference - device: device - dtype: dtype - - Returns: - - """ - seq_length = len(frame_ids) - reference_flow = torch.zeros((seq_length, 2, 512, 512), device=device, dtype=dtype) - for fr_idx in range(seq_length): - reference_flow[fr_idx, 0, :, :] = motion_field_strength_x * (frame_ids[fr_idx]) - reference_flow[fr_idx, 1, :, :] = motion_field_strength_y * (frame_ids[fr_idx]) - return reference_flow - - -# Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.create_motion_field_and_warp_latents -def create_motion_field_and_warp_latents(motion_field_strength_x, motion_field_strength_y, frame_ids, latents): - """ - Creates translation motion and warps the latents accordingly - - Args: - motion_field_strength_x: motion strength along x-axis - motion_field_strength_y: motion strength along y-axis - frame_ids: indexes of the frames the latents of which are being processed. - This is needed when we perform chunk-by-chunk inference - latents: latent codes of frames - - Returns: - warped_latents: warped latents - """ - motion_field = create_motion_field( - motion_field_strength_x=motion_field_strength_x, - motion_field_strength_y=motion_field_strength_y, - frame_ids=frame_ids, - device=latents.device, - dtype=latents.dtype, - ) - warped_latents = latents.clone().detach() - for i in range(len(warped_latents)): - warped_latents[i] = warp_single_latent(latents[i][None], motion_field[i][None]) - return warped_latents - - -# 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 TextToVideoZeroSDXLPipeline( - DeprecatedPipelineMixin, - DiffusionPipeline, - StableDiffusionMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -): - _last_supported_version = "0.33.1" - r""" - Pipeline for zero-shot text-to-video generation using Stable Diffusion XL. - - 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: - 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`]. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "image_encoder", - "feature_extractor", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - image_encoder: CLIPVisionModelWithProjection = None, - feature_extractor: CLIPImageProcessor = None, - force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - ): - super().__init__() - 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, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - ) - 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 - ) - - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - FusedAttnProcessor2_0, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # 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.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - 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) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def check_inputs( - self, - prompt, - prompt_2, - height, - width, - callback_steps, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_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_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`." - ) - - # 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 - - # Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.TextToVideoZeroPipeline.forward_loop - def forward_loop(self, x_t0, t0, t1, generator): - """ - Perform DDPM forward process from time t0 to t1. This is the same as adding noise with corresponding variance. - - Args: - x_t0: - Latent code at time t0. - t0: - Timestep at t0. - t1: - Timestamp at t1. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - - Returns: - x_t1: - Forward process applied to x_t0 from time t0 to t1. - """ - eps = randn_tensor(x_t0.size(), generator=generator, dtype=x_t0.dtype, device=x_t0.device) - alpha_vec = torch.prod(self.scheduler.alphas[t0:t1]) - x_t1 = torch.sqrt(alpha_vec) * x_t0 + torch.sqrt(1 - alpha_vec) * eps - return x_t1 - - def backward_loop( - self, - latents, - timesteps, - prompt_embeds, - guidance_scale, - callback, - callback_steps, - num_warmup_steps, - extra_step_kwargs, - add_text_embeds, - add_time_ids, - cross_attention_kwargs=None, - guidance_rescale: float = 0.0, - ): - """ - Perform backward process given list of time steps - - Args: - latents: - Latents at time timesteps[0]. - timesteps: - Time steps along which to perform backward process. - prompt_embeds: - Pre-generated text embeddings. - guidance_scale: - 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`. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - extra_step_kwargs: - Extra_step_kwargs. - cross_attention_kwargs: - 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). - num_warmup_steps: - number of warmup steps. - - Returns: - latents: latents of backward process output at time timesteps[-1] - """ - - do_classifier_free_guidance = guidance_scale > 1.0 - num_steps = (len(timesteps) - num_warmup_steps) // self.scheduler.order - - with self.progress_bar(total=num_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - if do_classifier_free_guidance and guidance_rescale > 0.0: - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - callback(i, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - return latents.clone().detach() - - @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, - num_inference_steps: int = 50, - denoising_end: Optional[float] = 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, - 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, - motion_field_strength_x: float = 12, - motion_field_strength_y: float = 12, - output_type: Optional[str] = "tensor", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = 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, - t0: int = 44, - t1: int = 47, - ): - """ - 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 - video_length (`int`, *optional*, defaults to 8): - The number of generated video frames. - 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. - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - 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 - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output) - 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`). - 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): - The number of videos 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` 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*): - Indexes of the frames that are being generated. This is used when generating longer videos - chunk-by-chunk. - 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. - 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 ge generated by sampling using the supplied random `generator`. - motion_field_strength_x (`float`, *optional*, defaults to 12): - Strength of motion in generated video along x-axis. See the - [paper](https://huggingface.co/papers/2303.13439), Sect. 3.3.1. - motion_field_strength_y (`float`, *optional*, defaults to 12): - Strength of motion in generated video along y-axis. See the - [paper](https://huggingface.co/papers/2303.13439), Sect. 3.3.1. - 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. - 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. - cross_attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under - `self.processor` in - [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py). - guidance_rescale (`float`, *optional*, defaults to 0.7): - 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. - 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` 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 `(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). - t0 (`int`, *optional*, defaults to 44): - Timestep t0. Should be in the range [0, num_inference_steps - 1]. See the - [paper](https://huggingface.co/papers/2303.13439), Sect. 3.3.1. - 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. - - Returns: - [`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.TextToVideoSDXLPipelineOutput`] or - `tuple`: [`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.TextToVideoSDXLPipelineOutput`] - if `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the - generated images. - """ - assert video_length > 0 - if frame_ids is None: - frame_ids = list(range(video_length)) - assert len(frame_ids) == video_length - - assert num_videos_per_prompt == 1 - - # set the processor - original_attn_proc = self.unet.attn_processors - processor = ( - CrossFrameAttnProcessor2_0(batch_size=2) - if hasattr(F, "scaled_dot_product_attention") - else CrossFrameAttnProcessor(batch_size=2) - ) - self.unet.set_attn_processor(processor) - - if isinstance(prompt, str): - prompt = [prompt] - if isinstance(negative_prompt, str): - negative_prompt = [negative_prompt] - - # 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, - callback_steps, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) - - # 2. Define call parameters - batch_size = ( - 1 if isinstance(prompt, str) else len(prompt) if isinstance(prompt, list) else 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 = guidance_scale > 1.0 - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else 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_videos_per_prompt, - do_classifier_free_guidance=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=text_encoder_lora_scale, - ) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 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 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([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_videos_per_prompt, 1) - - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - # Perform the first backward process up to time T_1 - x_1_t1 = self.backward_loop( - timesteps=timesteps[: -t1 - 1], - prompt_embeds=prompt_embeds, - latents=latents, - guidance_scale=guidance_scale, - callback=callback, - callback_steps=callback_steps, - extra_step_kwargs=extra_step_kwargs, - num_warmup_steps=num_warmup_steps, - add_text_embeds=add_text_embeds, - add_time_ids=add_time_ids, - ) - - scheduler_copy = copy.deepcopy(self.scheduler) - - # Perform the second backward process up to time T_0 - x_1_t0 = self.backward_loop( - timesteps=timesteps[-t1 - 1 : -t0 - 1], - prompt_embeds=prompt_embeds, - latents=x_1_t1, - guidance_scale=guidance_scale, - callback=callback, - callback_steps=callback_steps, - extra_step_kwargs=extra_step_kwargs, - num_warmup_steps=0, - add_text_embeds=add_text_embeds, - add_time_ids=add_time_ids, - ) - - # Propagate first frame latents at time T_0 to remaining frames - x_2k_t0 = x_1_t0.repeat(video_length - 1, 1, 1, 1) - - # Add motion in latents at time T_0 - x_2k_t0 = create_motion_field_and_warp_latents( - motion_field_strength_x=motion_field_strength_x, - motion_field_strength_y=motion_field_strength_y, - latents=x_2k_t0, - frame_ids=frame_ids[1:], - ) - - # Perform forward process up to time T_1 - x_2k_t1 = self.forward_loop( - x_t0=x_2k_t0, - t0=timesteps[-t0 - 1].to(torch.long), - t1=timesteps[-t1 - 1].to(torch.long), - generator=generator, - ) - - # Perform backward process from time T_1 to 0 - latents = torch.cat([x_1_t1, x_2k_t1]) - - self.scheduler = scheduler_copy - timesteps = timesteps[-t1 - 1 :] - - b, l, d = prompt_embeds.size() - prompt_embeds = prompt_embeds[:, None].repeat(1, video_length, 1, 1).reshape(b * video_length, l, d) - - b, k = add_text_embeds.size() - add_text_embeds = add_text_embeds[:, None].repeat(1, video_length, 1).reshape(b * video_length, k) - - b, k = add_time_ids.size() - add_time_ids = add_time_ids[:, None].repeat(1, video_length, 1).reshape(b * video_length, k) - - # 7.1 Apply denoising_end - if denoising_end is not None and isinstance(denoising_end, float) and denoising_end > 0 and denoising_end < 1: - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - x_1k_0 = self.backward_loop( - timesteps=timesteps, - prompt_embeds=prompt_embeds, - latents=latents, - guidance_scale=guidance_scale, - callback=callback, - callback_steps=callback_steps, - extra_step_kwargs=extra_step_kwargs, - num_warmup_steps=0, - add_text_embeds=add_text_embeds, - add_time_ids=add_time_ids, - ) - - latents = x_1k_0 - - 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 - return TextToVideoSDXLPipelineOutput(images=image) - - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - image = self.image_processor.postprocess(image, output_type=output_type) - - self.maybe_free_model_hooks() - # make sure to set the original attention processors back - self.unet.set_attn_processor(original_attn_proc) - - if not return_dict: - return (image,) - - return TextToVideoSDXLPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/transformers_loading_utils.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/transformers_loading_utils.py deleted file mode 100644 index b52d154d6ba2fb2e91a599731f3e4caf7aef5a77..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/transformers_loading_utils.py +++ /dev/null @@ -1,121 +0,0 @@ -# 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. -import contextlib -import os -import tempfile -from typing import TYPE_CHECKING, Dict - -from huggingface_hub import DDUFEntry -from tqdm import tqdm - -from ..utils import is_safetensors_available, is_transformers_available, is_transformers_version - - -if TYPE_CHECKING: - from transformers import PreTrainedModel, PreTrainedTokenizer - -if is_transformers_available(): - from transformers import PreTrainedModel, PreTrainedTokenizer - -if is_safetensors_available(): - import safetensors.torch - - -def _load_tokenizer_from_dduf( - cls: "PreTrainedTokenizer", name: str, dduf_entries: Dict[str, DDUFEntry], **kwargs -) -> "PreTrainedTokenizer": - """ - Load a tokenizer from a DDUF archive. - - In practice, `transformers` do not provide a way to load a tokenizer from a DDUF archive. This function is a - workaround by extracting the tokenizer files from the DDUF archive and loading the tokenizer from the extracted - files. There is an extra cost of extracting the files, but of limited impact as the tokenizer files are usually - small-ish. - """ - with tempfile.TemporaryDirectory() as tmp_dir: - for entry_name, entry in dduf_entries.items(): - if entry_name.startswith(name + "/"): - tmp_entry_path = os.path.join(tmp_dir, *entry_name.split("/")) - # need to create intermediary directory if they don't exist - os.makedirs(os.path.dirname(tmp_entry_path), exist_ok=True) - with open(tmp_entry_path, "wb") as f: - with entry.as_mmap() as mm: - f.write(mm) - return cls.from_pretrained(os.path.dirname(tmp_entry_path), **kwargs) - - -def _load_transformers_model_from_dduf( - cls: "PreTrainedModel", name: str, dduf_entries: Dict[str, DDUFEntry], **kwargs -) -> "PreTrainedModel": - """ - Load a transformers model from a DDUF archive. - - In practice, `transformers` do not provide a way to load a model from a DDUF archive. This function is a workaround - by instantiating a model from the config file and loading the weights from the DDUF archive directly. - """ - config_file = dduf_entries.get(f"{name}/config.json") - if config_file is None: - raise EnvironmentError( - f"Could not find a config.json file for component {name} in DDUF file (contains {dduf_entries.keys()})." - ) - generation_config = dduf_entries.get(f"{name}/generation_config.json", None) - - weight_files = [ - entry - for entry_name, entry in dduf_entries.items() - if entry_name.startswith(f"{name}/") and entry_name.endswith(".safetensors") - ] - if not weight_files: - raise EnvironmentError( - f"Could not find any weight file for component {name} in DDUF file (contains {dduf_entries.keys()})." - ) - if not is_safetensors_available(): - raise EnvironmentError( - "Safetensors is not available, cannot load model from DDUF. Please `pip install safetensors`." - ) - if is_transformers_version("<", "4.47.0"): - raise ImportError( - "You need to install `transformers>4.47.0` in order to load a transformers model from a DDUF file. " - "You can install it with: `pip install --upgrade transformers`" - ) - - with tempfile.TemporaryDirectory() as tmp_dir: - from transformers import AutoConfig, GenerationConfig - - tmp_config_file = os.path.join(tmp_dir, "config.json") - with open(tmp_config_file, "w") as f: - f.write(config_file.read_text()) - config = AutoConfig.from_pretrained(tmp_config_file) - if generation_config is not None: - tmp_generation_config_file = os.path.join(tmp_dir, "generation_config.json") - with open(tmp_generation_config_file, "w") as f: - f.write(generation_config.read_text()) - generation_config = GenerationConfig.from_pretrained(tmp_generation_config_file) - state_dict = {} - with contextlib.ExitStack() as stack: - for entry in tqdm(weight_files, desc="Loading state_dict"): # Loop over safetensors files - # Memory-map the safetensors file - mmap = stack.enter_context(entry.as_mmap()) - # Load tensors from the memory-mapped file - tensors = safetensors.torch.load(mmap) - # Update the state dictionary with tensors - state_dict.update(tensors) - return cls.from_pretrained( - pretrained_model_name_or_path=None, - config=config, - generation_config=generation_config, - state_dict=state_dict, - **kwargs, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__init__.py deleted file mode 100644 index c89e899463beede59b8ccf02688f6168b8ee3d77..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__init__.py +++ /dev/null @@ -1,52 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - is_torch_available, - is_transformers_available, - is_transformers_version, -) - - -_dummy_objects = {} -_import_structure = {} - -try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.25.0")): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_objects import UnCLIPImageVariationPipeline, UnCLIPPipeline - - _dummy_objects.update( - {"UnCLIPImageVariationPipeline": UnCLIPImageVariationPipeline, "UnCLIPPipeline": UnCLIPPipeline} - ) -else: - _import_structure["pipeline_unclip"] = ["UnCLIPPipeline"] - _import_structure["pipeline_unclip_image_variation"] = ["UnCLIPImageVariationPipeline"] - _import_structure["text_proj"] = ["UnCLIPTextProjModel"] - - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - try: - if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.25.0")): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 - else: - from .pipeline_unclip import UnCLIPPipeline - from .pipeline_unclip_image_variation import UnCLIPImageVariationPipeline - from .text_proj import UnCLIPTextProjModel - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 6279c824a6966c8ca3c0efaa705281e5bc9fd713..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__pycache__/pipeline_unclip.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__pycache__/pipeline_unclip.cpython-310.pyc deleted file mode 100644 index e93639a2a5e68973a9fe0797cb40ada3be9a2724..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__pycache__/pipeline_unclip.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__pycache__/pipeline_unclip_image_variation.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__pycache__/pipeline_unclip_image_variation.cpython-310.pyc deleted file mode 100644 index 0605fc333fc403b8f7066dba5111d35a73ac2bab..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__pycache__/pipeline_unclip_image_variation.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__pycache__/text_proj.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__pycache__/text_proj.cpython-310.pyc deleted file mode 100644 index 4b6e957252b1f4b98675b0a698b18f5878654e0f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/__pycache__/text_proj.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/pipeline_unclip.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/pipeline_unclip.py deleted file mode 100644 index bbb9b0eb3ab27983361511da45b672933eeba472..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/pipeline_unclip.py +++ /dev/null @@ -1,504 +0,0 @@ -# Copyright 2025 Kakao Brain 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 List, Optional, Tuple, Union - -import torch -from torch.nn import functional as F -from transformers import CLIPTextModelWithProjection, CLIPTokenizer -from transformers.models.clip.modeling_clip import CLIPTextModelOutput - -from ...models import PriorTransformer, UNet2DConditionModel, UNet2DModel -from ...schedulers import UnCLIPScheduler -from ...utils import is_torch_xla_available, logging -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, ImagePipelineOutput -from .text_proj import UnCLIPTextProjModel - - -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 - - -class UnCLIPPipeline(DeprecatedPipelineMixin, DiffusionPipeline): - """ - Pipeline for text-to-image generation using unCLIP. - - 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 ([`~transformers.CLIPTextModelWithProjection`]): - Frozen text-encoder. - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - prior ([`PriorTransformer`]): - The canonical unCLIP prior to approximate the image embedding from the text embedding. - text_proj ([`UnCLIPTextProjModel`]): - Utility class to prepare and combine the embeddings before they are passed to the decoder. - decoder ([`UNet2DConditionModel`]): - The decoder to invert the image embedding into an image. - super_res_first ([`UNet2DModel`]): - Super resolution UNet. Used in all but the last step of the super resolution diffusion process. - super_res_last ([`UNet2DModel`]): - Super resolution UNet. Used in the last step of the super resolution diffusion process. - prior_scheduler ([`UnCLIPScheduler`]): - Scheduler used in the prior denoising process (a modified [`DDPMScheduler`]). - decoder_scheduler ([`UnCLIPScheduler`]): - Scheduler used in the decoder denoising process (a modified [`DDPMScheduler`]). - super_res_scheduler ([`UnCLIPScheduler`]): - Scheduler used in the super resolution denoising process (a modified [`DDPMScheduler`]). - - """ - - _last_supported_version = "0.33.1" - _exclude_from_cpu_offload = ["prior"] - - prior: PriorTransformer - decoder: UNet2DConditionModel - text_proj: UnCLIPTextProjModel - text_encoder: CLIPTextModelWithProjection - tokenizer: CLIPTokenizer - super_res_first: UNet2DModel - super_res_last: UNet2DModel - - prior_scheduler: UnCLIPScheduler - decoder_scheduler: UnCLIPScheduler - super_res_scheduler: UnCLIPScheduler - - model_cpu_offload_seq = "text_encoder->text_proj->decoder->super_res_first->super_res_last" - - def __init__( - self, - prior: PriorTransformer, - decoder: UNet2DConditionModel, - text_encoder: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - text_proj: UnCLIPTextProjModel, - super_res_first: UNet2DModel, - super_res_last: UNet2DModel, - prior_scheduler: UnCLIPScheduler, - decoder_scheduler: UnCLIPScheduler, - super_res_scheduler: UnCLIPScheduler, - ): - super().__init__() - - self.register_modules( - prior=prior, - decoder=decoder, - text_encoder=text_encoder, - tokenizer=tokenizer, - text_proj=text_proj, - super_res_first=super_res_first, - super_res_last=super_res_last, - prior_scheduler=prior_scheduler, - decoder_scheduler=decoder_scheduler, - super_res_scheduler=super_res_scheduler, - ) - - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - def _encode_prompt( - self, - 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, - ): - if text_model_output is None: - batch_size = len(prompt) if isinstance(prompt, list) else 1 - # get prompt text embeddings - 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 - text_mask = text_inputs.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[:, 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}" - ) - text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length] - - text_encoder_output = self.text_encoder(text_input_ids.to(device)) - - prompt_embeds = text_encoder_output.text_embeds - text_enc_hid_states = text_encoder_output.last_hidden_state - - else: - batch_size = text_model_output[0].shape[0] - prompt_embeds, text_enc_hid_states = text_model_output[0], text_model_output[1] - text_mask = text_attention_mask - - prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - text_enc_hid_states = text_enc_hid_states.repeat_interleave(num_images_per_prompt, dim=0) - text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - uncond_tokens = [""] * batch_size - - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - uncond_text_mask = uncond_input.attention_mask.bool().to(device) - negative_prompt_embeds_text_encoder_output = self.text_encoder(uncond_input.input_ids.to(device)) - - negative_prompt_embeds = negative_prompt_embeds_text_encoder_output.text_embeds - uncond_text_enc_hid_states = negative_prompt_embeds_text_encoder_output.last_hidden_state - - # 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.repeat(1, num_images_per_prompt) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len) - - seq_len = uncond_text_enc_hid_states.shape[1] - uncond_text_enc_hid_states = uncond_text_enc_hid_states.repeat(1, num_images_per_prompt, 1) - uncond_text_enc_hid_states = uncond_text_enc_hid_states.view( - batch_size * num_images_per_prompt, seq_len, -1 - ) - uncond_text_mask = uncond_text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - # done duplicates - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - text_enc_hid_states = torch.cat([uncond_text_enc_hid_states, text_enc_hid_states]) - - text_mask = torch.cat([uncond_text_mask, text_mask]) - - return prompt_embeds, text_enc_hid_states, text_mask - - @torch.no_grad() - def __call__( - self, - prompt: Optional[Union[str, List[str]]] = 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, - prior_guidance_scale: float = 4.0, - decoder_guidance_scale: float = 8.0, - output_type: Optional[str] = "pil", - return_dict: bool = True, - ): - """ - The call function to the pipeline for generation. - - Args: - 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): - The number of images to generate per prompt. - prior_num_inference_steps (`int`, *optional*, defaults to 25): - The number of denoising steps for the prior. More denoising steps usually lead to a higher quality - image at the expense of slower inference. - decoder_num_inference_steps (`int`, *optional*, defaults to 25): - The number of denoising steps for the decoder. More denoising steps usually lead to a higher quality - image at the expense of slower inference. - 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*): - 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*): - Pre-generated noisy latents to be used as inputs for the prior. - decoder_latents (`torch.Tensor` of shape (batch size, channels, height, width), *optional*): - Pre-generated noisy latents to be used as inputs for the decoder. - super_res_latents (`torch.Tensor` of shape (batch size, channels, super res height, super res width), *optional*): - Pre-generated noisy latents to be used as inputs for the decoder. - prior_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`. - decoder_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`. - text_model_output (`CLIPTextModelOutput`, *optional*): - Pre-defined [`CLIPTextModel`] outputs that can be derived from the text encoder. Pre-defined text - outputs can be passed for tasks like text embedding interpolations. Make sure to also pass - `text_attention_mask` in this case. `prompt` can the be left `None`. - text_attention_mask (`torch.Tensor`, *optional*): - Pre-defined CLIP text attention mask that can be derived from the tokenizer. Pre-defined text attention - masks are necessary when passing `text_model_output`. - 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 [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images. - """ - if prompt is not None: - if isinstance(prompt, str): - batch_size = 1 - elif isinstance(prompt, list): - batch_size = len(prompt) - else: - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - else: - batch_size = text_model_output[0].shape[0] - - device = self._execution_device - - batch_size = batch_size * num_images_per_prompt - - do_classifier_free_guidance = prior_guidance_scale > 1.0 or decoder_guidance_scale > 1.0 - - prompt_embeds, text_enc_hid_states, text_mask = self._encode_prompt( - prompt, device, num_images_per_prompt, do_classifier_free_guidance, text_model_output, text_attention_mask - ) - - # prior - - self.prior_scheduler.set_timesteps(prior_num_inference_steps, device=device) - prior_timesteps_tensor = self.prior_scheduler.timesteps - - embedding_dim = self.prior.config.embedding_dim - - prior_latents = self.prepare_latents( - (batch_size, embedding_dim), - prompt_embeds.dtype, - device, - generator, - prior_latents, - self.prior_scheduler, - ) - - for i, t in enumerate(self.progress_bar(prior_timesteps_tensor)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([prior_latents] * 2) if do_classifier_free_guidance else prior_latents - - predicted_image_embedding = self.prior( - latent_model_input, - timestep=t, - proj_embedding=prompt_embeds, - encoder_hidden_states=text_enc_hid_states, - attention_mask=text_mask, - ).predicted_image_embedding - - if do_classifier_free_guidance: - predicted_image_embedding_uncond, predicted_image_embedding_text = predicted_image_embedding.chunk(2) - predicted_image_embedding = predicted_image_embedding_uncond + prior_guidance_scale * ( - predicted_image_embedding_text - predicted_image_embedding_uncond - ) - - if i + 1 == prior_timesteps_tensor.shape[0]: - prev_timestep = None - else: - prev_timestep = prior_timesteps_tensor[i + 1] - - prior_latents = self.prior_scheduler.step( - predicted_image_embedding, - timestep=t, - sample=prior_latents, - generator=generator, - prev_timestep=prev_timestep, - ).prev_sample - - prior_latents = self.prior.post_process_latents(prior_latents) - - image_embeddings = prior_latents - - # done prior - - # decoder - - text_enc_hid_states, additive_clip_time_embeddings = self.text_proj( - image_embeddings=image_embeddings, - prompt_embeds=prompt_embeds, - text_encoder_hidden_states=text_enc_hid_states, - do_classifier_free_guidance=do_classifier_free_guidance, - ) - - if device.type == "mps": - # HACK: MPS: There is a panic when padding bool tensors, - # so cast to int tensor for the pad and back to bool afterwards - text_mask = text_mask.type(torch.int) - decoder_text_mask = F.pad(text_mask, (self.text_proj.clip_extra_context_tokens, 0), value=1) - decoder_text_mask = decoder_text_mask.type(torch.bool) - else: - decoder_text_mask = F.pad(text_mask, (self.text_proj.clip_extra_context_tokens, 0), value=True) - - self.decoder_scheduler.set_timesteps(decoder_num_inference_steps, device=device) - decoder_timesteps_tensor = self.decoder_scheduler.timesteps - - num_channels_latents = self.decoder.config.in_channels - height = self.decoder.config.sample_size - width = self.decoder.config.sample_size - - decoder_latents = self.prepare_latents( - (batch_size, num_channels_latents, height, width), - text_enc_hid_states.dtype, - device, - generator, - decoder_latents, - self.decoder_scheduler, - ) - - for i, t in enumerate(self.progress_bar(decoder_timesteps_tensor)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([decoder_latents] * 2) if do_classifier_free_guidance else decoder_latents - - noise_pred = self.decoder( - sample=latent_model_input, - timestep=t, - encoder_hidden_states=text_enc_hid_states, - class_labels=additive_clip_time_embeddings, - attention_mask=decoder_text_mask, - ).sample - - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred_uncond, _ = noise_pred_uncond.split(latent_model_input.shape[1], dim=1) - noise_pred_text, predicted_variance = noise_pred_text.split(latent_model_input.shape[1], dim=1) - noise_pred = noise_pred_uncond + decoder_guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, predicted_variance], dim=1) - - if i + 1 == decoder_timesteps_tensor.shape[0]: - prev_timestep = None - else: - prev_timestep = decoder_timesteps_tensor[i + 1] - - # compute the previous noisy sample x_t -> x_t-1 - decoder_latents = self.decoder_scheduler.step( - noise_pred, t, decoder_latents, prev_timestep=prev_timestep, generator=generator - ).prev_sample - - decoder_latents = decoder_latents.clamp(-1, 1) - - image_small = decoder_latents - - # done decoder - - # super res - - self.super_res_scheduler.set_timesteps(super_res_num_inference_steps, device=device) - super_res_timesteps_tensor = self.super_res_scheduler.timesteps - - channels = self.super_res_first.config.in_channels // 2 - height = self.super_res_first.config.sample_size - width = self.super_res_first.config.sample_size - - super_res_latents = self.prepare_latents( - (batch_size, channels, height, width), - image_small.dtype, - device, - generator, - super_res_latents, - self.super_res_scheduler, - ) - - if device.type == "mps": - # MPS does not support many interpolations - image_upscaled = F.interpolate(image_small, size=[height, width]) - else: - interpolate_antialias = {} - if "antialias" in inspect.signature(F.interpolate).parameters: - interpolate_antialias["antialias"] = True - - image_upscaled = F.interpolate( - image_small, size=[height, width], mode="bicubic", align_corners=False, **interpolate_antialias - ) - - for i, t in enumerate(self.progress_bar(super_res_timesteps_tensor)): - # no classifier free guidance - - if i == super_res_timesteps_tensor.shape[0] - 1: - unet = self.super_res_last - else: - unet = self.super_res_first - - latent_model_input = torch.cat([super_res_latents, image_upscaled], dim=1) - - noise_pred = unet( - sample=latent_model_input, - timestep=t, - ).sample - - if i + 1 == super_res_timesteps_tensor.shape[0]: - prev_timestep = None - else: - prev_timestep = super_res_timesteps_tensor[i + 1] - - # compute the previous noisy sample x_t -> x_t-1 - super_res_latents = self.super_res_scheduler.step( - noise_pred, t, super_res_latents, prev_timestep=prev_timestep, generator=generator - ).prev_sample - - if XLA_AVAILABLE: - xm.mark_step() - - image = super_res_latents - # done super res - - self.maybe_free_model_hooks() - - # post processing - image = image * 0.5 + 0.5 - image = image.clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/pipeline_unclip_image_variation.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/pipeline_unclip_image_variation.py deleted file mode 100644 index 31710a000e0a50242eab1623b9b9b8f8d92cd908..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/pipeline_unclip_image_variation.py +++ /dev/null @@ -1,431 +0,0 @@ -# Copyright 2025 Kakao Brain 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 List, Optional, Union - -import PIL.Image -import torch -from torch.nn import functional as F -from transformers import ( - CLIPImageProcessor, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from ...models import UNet2DConditionModel, UNet2DModel -from ...schedulers import UnCLIPScheduler -from ...utils import is_torch_xla_available, logging -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, ImagePipelineOutput -from .text_proj import UnCLIPTextProjModel - - -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 - - -class UnCLIPImageVariationPipeline(DeprecatedPipelineMixin, DiffusionPipeline): - """ - Pipeline to generate image variations from an input image using UnCLIP. - - 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 ([`~transformers.CLIPTextModelWithProjection`]): - Frozen text-encoder. - tokenizer ([`~transformers.CLIPTokenizer`]): - A `CLIPTokenizer` to tokenize text. - feature_extractor ([`~transformers.CLIPImageProcessor`]): - Model that extracts features from generated images to be used as inputs for the `image_encoder`. - image_encoder ([`~transformers.CLIPVisionModelWithProjection`]): - Frozen CLIP image-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - text_proj ([`UnCLIPTextProjModel`]): - Utility class to prepare and combine the embeddings before they are passed to the decoder. - decoder ([`UNet2DConditionModel`]): - The decoder to invert the image embedding into an image. - super_res_first ([`UNet2DModel`]): - Super resolution UNet. Used in all but the last step of the super resolution diffusion process. - super_res_last ([`UNet2DModel`]): - Super resolution UNet. Used in the last step of the super resolution diffusion process. - decoder_scheduler ([`UnCLIPScheduler`]): - Scheduler used in the decoder denoising process (a modified [`DDPMScheduler`]). - super_res_scheduler ([`UnCLIPScheduler`]): - Scheduler used in the super resolution denoising process (a modified [`DDPMScheduler`]). - """ - - _last_supported_version = "0.33.1" - decoder: UNet2DConditionModel - text_proj: UnCLIPTextProjModel - text_encoder: CLIPTextModelWithProjection - tokenizer: CLIPTokenizer - feature_extractor: CLIPImageProcessor - image_encoder: CLIPVisionModelWithProjection - super_res_first: UNet2DModel - super_res_last: UNet2DModel - - decoder_scheduler: UnCLIPScheduler - super_res_scheduler: UnCLIPScheduler - model_cpu_offload_seq = "text_encoder->image_encoder->text_proj->decoder->super_res_first->super_res_last" - - def __init__( - self, - decoder: UNet2DConditionModel, - text_encoder: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - text_proj: UnCLIPTextProjModel, - feature_extractor: CLIPImageProcessor, - image_encoder: CLIPVisionModelWithProjection, - super_res_first: UNet2DModel, - super_res_last: UNet2DModel, - decoder_scheduler: UnCLIPScheduler, - super_res_scheduler: UnCLIPScheduler, - ): - super().__init__() - - self.register_modules( - decoder=decoder, - text_encoder=text_encoder, - tokenizer=tokenizer, - text_proj=text_proj, - feature_extractor=feature_extractor, - image_encoder=image_encoder, - super_res_first=super_res_first, - super_res_last=super_res_last, - decoder_scheduler=decoder_scheduler, - super_res_scheduler=super_res_scheduler, - ) - - # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - def _encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_free_guidance): - batch_size = len(prompt) if isinstance(prompt, list) else 1 - - # get prompt text embeddings - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - text_mask = text_inputs.attention_mask.bool().to(device) - text_encoder_output = self.text_encoder(text_input_ids.to(device)) - - prompt_embeds = text_encoder_output.text_embeds - text_encoder_hidden_states = text_encoder_output.last_hidden_state - - prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - text_encoder_hidden_states = text_encoder_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - if do_classifier_free_guidance: - uncond_tokens = [""] * batch_size - - max_length = text_input_ids.shape[-1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - uncond_text_mask = uncond_input.attention_mask.bool().to(device) - negative_prompt_embeds_text_encoder_output = self.text_encoder(uncond_input.input_ids.to(device)) - - negative_prompt_embeds = negative_prompt_embeds_text_encoder_output.text_embeds - uncond_text_encoder_hidden_states = negative_prompt_embeds_text_encoder_output.last_hidden_state - - # 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.repeat(1, num_images_per_prompt) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len) - - seq_len = uncond_text_encoder_hidden_states.shape[1] - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.repeat(1, num_images_per_prompt, 1) - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.view( - batch_size * num_images_per_prompt, seq_len, -1 - ) - uncond_text_mask = uncond_text_mask.repeat_interleave(num_images_per_prompt, dim=0) - - # done duplicates - - # 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 - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - text_encoder_hidden_states = torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]) - - text_mask = torch.cat([uncond_text_mask, text_mask]) - - 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): - dtype = next(self.image_encoder.parameters()).dtype - - if image_embeddings is None: - if not isinstance(image, torch.Tensor): - image = self.feature_extractor(images=image, return_tensors="pt").pixel_values - - image = image.to(device=device, dtype=dtype) - image_embeddings = self.image_encoder(image).image_embeds - - image_embeddings = image_embeddings.repeat_interleave(num_images_per_prompt, dim=0) - - return image_embeddings - - @torch.no_grad() - def __call__( - self, - image: Optional[Union[PIL.Image.Image, List[PIL.Image.Image], torch.Tensor]] = 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, - decoder_guidance_scale: float = 8.0, - output_type: Optional[str] = "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` 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). - Can be left as `None` only when `image_embeddings` are passed. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - decoder_num_inference_steps (`int`, *optional*, defaults to 25): - The number of denoising steps for the decoder. More denoising steps usually lead to a higher quality - image at the expense of slower inference. - 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`, *optional*): - A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make - generation deterministic. - decoder_latents (`torch.Tensor` of shape (batch size, channels, height, width), *optional*): - Pre-generated noisy latents to be used as inputs for the decoder. - super_res_latents (`torch.Tensor` of shape (batch size, channels, super res height, super res width), *optional*): - Pre-generated noisy latents to be used as inputs for the decoder. - decoder_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`. - image_embeddings (`torch.Tensor`, *optional*): - Pre-defined image embeddings that can be derived from the image encoder. Pre-defined image embeddings - can be passed for tasks like image interpolations. `image` can be left as `None`. - 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 [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images. - """ - if image is not None: - if isinstance(image, PIL.Image.Image): - batch_size = 1 - elif isinstance(image, list): - batch_size = len(image) - else: - batch_size = image.shape[0] - else: - batch_size = image_embeddings.shape[0] - - prompt = [""] * batch_size - - device = self._execution_device - - batch_size = batch_size * num_images_per_prompt - - do_classifier_free_guidance = decoder_guidance_scale > 1.0 - - prompt_embeds, text_encoder_hidden_states, text_mask = self._encode_prompt( - prompt, device, num_images_per_prompt, do_classifier_free_guidance - ) - - image_embeddings = self._encode_image(image, device, num_images_per_prompt, image_embeddings) - - # decoder - text_encoder_hidden_states, additive_clip_time_embeddings = self.text_proj( - image_embeddings=image_embeddings, - prompt_embeds=prompt_embeds, - text_encoder_hidden_states=text_encoder_hidden_states, - do_classifier_free_guidance=do_classifier_free_guidance, - ) - - if device.type == "mps": - # HACK: MPS: There is a panic when padding bool tensors, - # so cast to int tensor for the pad and back to bool afterwards - text_mask = text_mask.type(torch.int) - decoder_text_mask = F.pad(text_mask, (self.text_proj.clip_extra_context_tokens, 0), value=1) - decoder_text_mask = decoder_text_mask.type(torch.bool) - else: - decoder_text_mask = F.pad(text_mask, (self.text_proj.clip_extra_context_tokens, 0), value=True) - - self.decoder_scheduler.set_timesteps(decoder_num_inference_steps, device=device) - decoder_timesteps_tensor = self.decoder_scheduler.timesteps - - num_channels_latents = self.decoder.config.in_channels - height = self.decoder.config.sample_size - width = self.decoder.config.sample_size - - if decoder_latents is None: - decoder_latents = self.prepare_latents( - (batch_size, num_channels_latents, height, width), - text_encoder_hidden_states.dtype, - device, - generator, - decoder_latents, - self.decoder_scheduler, - ) - - for i, t in enumerate(self.progress_bar(decoder_timesteps_tensor)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([decoder_latents] * 2) if do_classifier_free_guidance else decoder_latents - - noise_pred = self.decoder( - sample=latent_model_input, - timestep=t, - encoder_hidden_states=text_encoder_hidden_states, - class_labels=additive_clip_time_embeddings, - attention_mask=decoder_text_mask, - ).sample - - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred_uncond, _ = noise_pred_uncond.split(latent_model_input.shape[1], dim=1) - noise_pred_text, predicted_variance = noise_pred_text.split(latent_model_input.shape[1], dim=1) - noise_pred = noise_pred_uncond + decoder_guidance_scale * (noise_pred_text - noise_pred_uncond) - noise_pred = torch.cat([noise_pred, predicted_variance], dim=1) - - if i + 1 == decoder_timesteps_tensor.shape[0]: - prev_timestep = None - else: - prev_timestep = decoder_timesteps_tensor[i + 1] - - # compute the previous noisy sample x_t -> x_t-1 - decoder_latents = self.decoder_scheduler.step( - noise_pred, t, decoder_latents, prev_timestep=prev_timestep, generator=generator - ).prev_sample - - decoder_latents = decoder_latents.clamp(-1, 1) - - image_small = decoder_latents - - # done decoder - - # super res - - self.super_res_scheduler.set_timesteps(super_res_num_inference_steps, device=device) - super_res_timesteps_tensor = self.super_res_scheduler.timesteps - - channels = self.super_res_first.config.in_channels // 2 - height = self.super_res_first.config.sample_size - width = self.super_res_first.config.sample_size - - if super_res_latents is None: - super_res_latents = self.prepare_latents( - (batch_size, channels, height, width), - image_small.dtype, - device, - generator, - super_res_latents, - self.super_res_scheduler, - ) - - if device.type == "mps": - # MPS does not support many interpolations - image_upscaled = F.interpolate(image_small, size=[height, width]) - else: - interpolate_antialias = {} - if "antialias" in inspect.signature(F.interpolate).parameters: - interpolate_antialias["antialias"] = True - - image_upscaled = F.interpolate( - image_small, size=[height, width], mode="bicubic", align_corners=False, **interpolate_antialias - ) - - for i, t in enumerate(self.progress_bar(super_res_timesteps_tensor)): - # no classifier free guidance - - if i == super_res_timesteps_tensor.shape[0] - 1: - unet = self.super_res_last - else: - unet = self.super_res_first - - latent_model_input = torch.cat([super_res_latents, image_upscaled], dim=1) - - noise_pred = unet( - sample=latent_model_input, - timestep=t, - ).sample - - if i + 1 == super_res_timesteps_tensor.shape[0]: - prev_timestep = None - else: - prev_timestep = super_res_timesteps_tensor[i + 1] - - # compute the previous noisy sample x_t -> x_t-1 - super_res_latents = self.super_res_scheduler.step( - noise_pred, t, super_res_latents, prev_timestep=prev_timestep, generator=generator - ).prev_sample - - if XLA_AVAILABLE: - xm.mark_step() - - image = super_res_latents - - # done super res - self.maybe_free_model_hooks() - - # post processing - - image = image * 0.5 + 0.5 - image = image.clamp(0, 1) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - - if output_type == "pil": - image = self.numpy_to_pil(image) - - if not return_dict: - return (image,) - - return ImagePipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/text_proj.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/text_proj.py deleted file mode 100644 index 5e04e48ba621caa4a617e52c5b83a1573bb11975..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/unclip/text_proj.py +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright 2025 Kakao Brain 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 torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models import ModelMixin - - -class UnCLIPTextProjModel(ModelMixin, ConfigMixin): - """ - Utility class for CLIP embeddings. Used to combine the image and text embeddings into a format usable by the - decoder. - - For more details, see the original paper: https://huggingface.co/papers/2204.06125 section 2.1 - """ - - @register_to_config - def __init__( - self, - *, - clip_extra_context_tokens: int = 4, - clip_embeddings_dim: int = 768, - time_embed_dim: int, - cross_attention_dim, - ): - super().__init__() - - self.learned_classifier_free_guidance_embeddings = nn.Parameter(torch.zeros(clip_embeddings_dim)) - - # parameters for additional clip time embeddings - self.embedding_proj = nn.Linear(clip_embeddings_dim, time_embed_dim) - self.clip_image_embeddings_project_to_time_embeddings = nn.Linear(clip_embeddings_dim, time_embed_dim) - - # parameters for encoder hidden states - self.clip_extra_context_tokens = clip_extra_context_tokens - self.clip_extra_context_tokens_proj = nn.Linear( - clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim - ) - self.encoder_hidden_states_proj = nn.Linear(clip_embeddings_dim, cross_attention_dim) - self.text_encoder_hidden_states_norm = nn.LayerNorm(cross_attention_dim) - - def forward(self, *, image_embeddings, prompt_embeds, text_encoder_hidden_states, do_classifier_free_guidance): - if do_classifier_free_guidance: - # Add the classifier free guidance embeddings to the image embeddings - image_embeddings_batch_size = image_embeddings.shape[0] - classifier_free_guidance_embeddings = self.learned_classifier_free_guidance_embeddings.unsqueeze(0) - classifier_free_guidance_embeddings = classifier_free_guidance_embeddings.expand( - image_embeddings_batch_size, -1 - ) - image_embeddings = torch.cat([classifier_free_guidance_embeddings, image_embeddings], dim=0) - - # The image embeddings batch size and the text embeddings batch size are equal - assert image_embeddings.shape[0] == prompt_embeds.shape[0] - - batch_size = prompt_embeds.shape[0] - - # "Specifically, we modify the architecture described in Nichol et al. (2021) by projecting and - # adding CLIP embeddings to the existing timestep embedding, ... - time_projected_prompt_embeds = self.embedding_proj(prompt_embeds) - time_projected_image_embeddings = self.clip_image_embeddings_project_to_time_embeddings(image_embeddings) - additive_clip_time_embeddings = time_projected_image_embeddings + time_projected_prompt_embeds - - # ... and by projecting CLIP embeddings into four - # extra tokens of context that are concatenated to the sequence of outputs from the GLIDE text encoder" - clip_extra_context_tokens = self.clip_extra_context_tokens_proj(image_embeddings) - clip_extra_context_tokens = clip_extra_context_tokens.reshape(batch_size, -1, self.clip_extra_context_tokens) - clip_extra_context_tokens = clip_extra_context_tokens.permute(0, 2, 1) - - text_encoder_hidden_states = self.encoder_hidden_states_proj(text_encoder_hidden_states) - text_encoder_hidden_states = self.text_encoder_hidden_states_norm(text_encoder_hidden_states) - text_encoder_hidden_states = torch.cat([clip_extra_context_tokens, text_encoder_hidden_states], dim=1) - - return text_encoder_hidden_states, additive_clip_time_embeddings diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__init__.py deleted file mode 100644 index 1ac2b09a6e570087c80bc11bf1a8102dd4970b8f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__init__.py +++ /dev/null @@ -1,58 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - 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.dummy_torch_and_transformers_objects import ( - ImageTextPipelineOutput, - UniDiffuserPipeline, - ) - - _dummy_objects.update( - {"ImageTextPipelineOutput": ImageTextPipelineOutput, "UniDiffuserPipeline": UniDiffuserPipeline} - ) -else: - _import_structure["modeling_text_decoder"] = ["UniDiffuserTextDecoder"] - _import_structure["modeling_uvit"] = ["UniDiffuserModel", "UTransformer2DModel"] - _import_structure["pipeline_unidiffuser"] = ["ImageTextPipelineOutput", "UniDiffuserPipeline"] - - -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 ( - ImageTextPipelineOutput, - UniDiffuserPipeline, - ) - else: - from .modeling_text_decoder import UniDiffuserTextDecoder - from .modeling_uvit import UniDiffuserModel, UTransformer2DModel - from .pipeline_unidiffuser import ImageTextPipelineOutput, UniDiffuserPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index aacf9792405a61862d8c1fb3d751f4ec320341af..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__pycache__/modeling_text_decoder.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__pycache__/modeling_text_decoder.cpython-310.pyc deleted file mode 100644 index 0870b30ce29200ba6f94deb0a02a9ff128b65a9f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__pycache__/modeling_text_decoder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__pycache__/modeling_uvit.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__pycache__/modeling_uvit.cpython-310.pyc deleted file mode 100644 index 7667b7d779748c66f8bf6372e23c2de4727da7aa..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__pycache__/modeling_uvit.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__pycache__/pipeline_unidiffuser.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__pycache__/pipeline_unidiffuser.cpython-310.pyc deleted file mode 100644 index 649c2751b33cb871591a081ea96eaa0d7c3197e0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/__pycache__/pipeline_unidiffuser.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/modeling_text_decoder.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/modeling_text_decoder.py deleted file mode 100644 index 0ddcbf735770aea23c08bbd3af15db9daa8cef9f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/modeling_text_decoder.py +++ /dev/null @@ -1,296 +0,0 @@ -from typing import Optional - -import numpy as np -import torch -from torch import nn -from transformers import GPT2Config, GPT2LMHeadModel -from transformers.modeling_utils import ModuleUtilsMixin - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models import ModelMixin - - -# Modified from ClipCaptionModel in https://github.com/thu-ml/unidiffuser/blob/main/libs/caption_decoder.py -class UniDiffuserTextDecoder(ModelMixin, ConfigMixin, ModuleUtilsMixin): - """ - Text decoder model for a image-text [UniDiffuser](https://huggingface.co/papers/2303.06555) model. This is used to - generate text from the UniDiffuser image-text embedding. - - Parameters: - prefix_length (`int`): - Max number of prefix tokens that will be supplied to the model. - prefix_inner_dim (`int`): - The hidden size of the incoming prefix embeddings. For UniDiffuser, this would be the hidden dim of the - CLIP text encoder. - prefix_hidden_dim (`int`, *optional*): - Hidden dim of the MLP if we encode the prefix. - vocab_size (`int`, *optional*, defaults to 50257): - Vocabulary size of the GPT-2 model. Defines the number of different tokens that can be represented by the - `inputs_ids` passed when calling [`GPT2Model`] or [`TFGPT2Model`]. - n_positions (`int`, *optional*, defaults to 1024): - The maximum sequence length that this model might ever be used with. Typically set this to something large - just in case (e.g., 512 or 1024 or 2048). - n_embd (`int`, *optional*, defaults to 768): - Dimensionality of the embeddings and hidden states. - n_layer (`int`, *optional*, defaults to 12): - Number of hidden layers in the Transformer encoder. - n_head (`int`, *optional*, defaults to 12): - Number of attention heads for each attention layer in the Transformer encoder. - n_inner (`int`, *optional*, defaults to None): - Dimensionality of the inner feed-forward layers. `None` will set it to 4 times n_embd - activation_function (`str`, *optional*, defaults to `"gelu"`): - Activation function, to be selected in the list `["relu", "silu", "gelu", "tanh", "gelu_new"]`. - resid_pdrop (`float`, *optional*, defaults to 0.1): - The dropout probability for all fully connected layers in the embeddings, encoder, and pooler. - embd_pdrop (`float`, *optional*, defaults to 0.1): - The dropout ratio for the embeddings. - attn_pdrop (`float`, *optional*, defaults to 0.1): - The dropout ratio for the attention. - layer_norm_epsilon (`float`, *optional*, defaults to 1e-5): - The epsilon to use in the layer normalization layers. - initializer_range (`float`, *optional*, defaults to 0.02): - The standard deviation of the truncated_normal_initializer for initializing all weight matrices. - scale_attn_weights (`bool`, *optional*, defaults to `True`): - Scale attention weights by dividing by sqrt(hidden_size).. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values attentions (not used by all models). - scale_attn_by_inverse_layer_idx (`bool`, *optional*, defaults to `False`): - Whether to additionally scale attention weights by `1 / layer_idx + 1`. - reorder_and_upcast_attn (`bool`, *optional*, defaults to `False`): - Whether to scale keys (K) prior to computing attention (dot-product) and upcast attention - dot-product/softmax to float() when training with mixed precision. - """ - - _keys_to_ignore_on_load_unexpected = [r"h\.\d+\.attn\.bias", r"h\.\d+\.attn\.masked_bias"] - - @register_to_config - def __init__( - self, - prefix_length: int, - prefix_inner_dim: int, - prefix_hidden_dim: Optional[int] = 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, - activation_function: str = "gelu_new", - resid_pdrop: float = 0.1, - embd_pdrop: float = 0.1, - attn_pdrop: float = 0.1, - layer_norm_epsilon: float = 1e-5, - initializer_range: float = 0.02, - scale_attn_weights: bool = True, - use_cache: bool = True, - scale_attn_by_inverse_layer_idx: bool = False, - reorder_and_upcast_attn: bool = False, - ): - super().__init__() - - self.prefix_length = prefix_length - - if prefix_inner_dim != n_embd and prefix_hidden_dim is None: - raise ValueError( - f"`prefix_hidden_dim` cannot be `None` when `prefix_inner_dim`: {prefix_hidden_dim} and" - f" `n_embd`: {n_embd} are not equal." - ) - - self.prefix_inner_dim = prefix_inner_dim - self.prefix_hidden_dim = prefix_hidden_dim - - self.encode_prefix = ( - nn.Linear(self.prefix_inner_dim, self.prefix_hidden_dim) - if self.prefix_hidden_dim is not None - else nn.Identity() - ) - self.decode_prefix = ( - nn.Linear(self.prefix_hidden_dim, n_embd) if self.prefix_hidden_dim is not None else nn.Identity() - ) - - gpt_config = GPT2Config( - vocab_size=vocab_size, - n_positions=n_positions, - n_embd=n_embd, - n_layer=n_layer, - n_head=n_head, - n_inner=n_inner, - activation_function=activation_function, - resid_pdrop=resid_pdrop, - embd_pdrop=embd_pdrop, - attn_pdrop=attn_pdrop, - layer_norm_epsilon=layer_norm_epsilon, - initializer_range=initializer_range, - scale_attn_weights=scale_attn_weights, - use_cache=use_cache, - scale_attn_by_inverse_layer_idx=scale_attn_by_inverse_layer_idx, - reorder_and_upcast_attn=reorder_and_upcast_attn, - ) - self.transformer = GPT2LMHeadModel(gpt_config) - - def forward( - self, - input_ids: torch.Tensor, - prefix_embeds: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - ): - """ - Args: - input_ids (`torch.Tensor` of shape `(N, max_seq_len)`): - Text tokens to use for inference. - prefix_embeds (`torch.Tensor` of shape `(N, prefix_length, 768)`): - Prefix embedding to prepend to the embedded tokens. - attention_mask (`torch.Tensor` of shape `(N, prefix_length + max_seq_len, 768)`, *optional*): - Attention mask for the prefix embedding. - labels (`torch.Tensor`, *optional*): - Labels to use for language modeling. - """ - embedding_text = self.transformer.transformer.wte(input_ids) - hidden = self.encode_prefix(prefix_embeds) - prefix_embeds = self.decode_prefix(hidden) - embedding_cat = torch.cat((prefix_embeds, embedding_text), dim=1) - - if labels is not None: - dummy_token = self.get_dummy_token(input_ids.shape[0], input_ids.device) - labels = torch.cat((dummy_token, input_ids), dim=1) - out = self.transformer(inputs_embeds=embedding_cat, labels=labels, attention_mask=attention_mask) - if self.prefix_hidden_dim is not None: - return out, hidden - else: - return out - - def get_dummy_token(self, batch_size: int, device: torch.device) -> torch.Tensor: - return torch.zeros(batch_size, self.prefix_length, dtype=torch.int64, device=device) - - def encode(self, prefix): - return self.encode_prefix(prefix) - - @torch.no_grad() - def generate_captions(self, features, eos_token_id, device): - """ - Generate captions given text embedding features. Returns list[L]. - - Args: - features (`torch.Tensor` of shape `(B, L, D)`): - Text embedding features to generate captions from. - eos_token_id (`int`): - The token ID of the EOS token for the text decoder model. - device: - Device to perform text generation on. - - Returns: - `List[str]`: A list of strings generated from the decoder model. - """ - - features = torch.split(features, 1, dim=0) - generated_tokens = [] - generated_seq_lengths = [] - for feature in features: - feature = self.decode_prefix(feature.to(device)) # back to the clip feature - # Only support beam search for now - output_tokens, seq_lengths = self.generate_beam( - input_embeds=feature, device=device, eos_token_id=eos_token_id - ) - generated_tokens.append(output_tokens[0]) - generated_seq_lengths.append(seq_lengths[0]) - generated_tokens = torch.stack(generated_tokens) - generated_seq_lengths = torch.stack(generated_seq_lengths) - return generated_tokens, generated_seq_lengths - - @torch.no_grad() - def generate_beam( - self, - input_ids=None, - input_embeds=None, - device=None, - beam_size: int = 5, - entry_length: int = 67, - temperature: float = 1.0, - eos_token_id: Optional[int] = None, - ): - """ - Generates text using the given tokenizer and text prompt or token embedding via beam search. This - implementation is based on the beam search implementation from the [original UniDiffuser - code](https://github.com/thu-ml/unidiffuser/blob/main/libs/caption_decoder.py#L89). - - Args: - eos_token_id (`int`, *optional*): - The token ID of the EOS token for the text decoder model. - input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`, *optional*): - Tokenizer indices of input sequence tokens in the vocabulary. One of `input_ids` and `input_embeds` - must be supplied. - input_embeds (`torch.Tensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*): - An embedded representation to directly pass to the transformer as a prefix for beam search. One of - `input_ids` and `input_embeds` must be supplied. - device: - The device to perform beam search on. - beam_size (`int`, *optional*, defaults to `5`): - The number of best states to store during beam search. - entry_length (`int`, *optional*, defaults to `67`): - The number of iterations to run beam search. - temperature (`float`, *optional*, defaults to 1.0): - 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 - token sequences sorted by score in descending order, and the second element is the sequence lengths - corresponding to those sequences. - """ - # Generates text until stop_token is reached using beam search with the desired beam size. - stop_token_index = eos_token_id - tokens = None - scores = None - seq_lengths = torch.ones(beam_size, device=device, dtype=torch.int) - is_stopped = torch.zeros(beam_size, device=device, dtype=torch.bool) - - if input_embeds is not None: - generated = input_embeds - else: - generated = self.transformer.transformer.wte(input_ids) - - for i in range(entry_length): - outputs = self.transformer(inputs_embeds=generated) - logits = outputs.logits - logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0) - logits = logits.softmax(-1).log() - - if scores is None: - scores, next_tokens = logits.topk(beam_size, -1) - generated = generated.expand(beam_size, *generated.shape[1:]) - next_tokens, scores = next_tokens.permute(1, 0), scores.squeeze(0) - if tokens is None: - tokens = next_tokens - else: - tokens = tokens.expand(beam_size, *tokens.shape[1:]) - tokens = torch.cat((tokens, next_tokens), dim=1) - else: - logits[is_stopped] = -float(np.inf) - logits[is_stopped, 0] = 0 - scores_sum = scores[:, None] + logits - seq_lengths[~is_stopped] += 1 - scores_sum_average = scores_sum / seq_lengths[:, None] - scores_sum_average, next_tokens = scores_sum_average.view(-1).topk(beam_size, -1) - next_tokens_source = next_tokens // scores_sum.shape[1] - seq_lengths = seq_lengths[next_tokens_source] - next_tokens = next_tokens % scores_sum.shape[1] - next_tokens = next_tokens.unsqueeze(1) - tokens = tokens[next_tokens_source] - tokens = torch.cat((tokens, next_tokens), dim=1) - generated = generated[next_tokens_source] - scores = scores_sum_average * seq_lengths - is_stopped = is_stopped[next_tokens_source] - - next_token_embed = self.transformer.transformer.wte(next_tokens.squeeze()).view(generated.shape[0], 1, -1) - generated = torch.cat((generated, next_token_embed), dim=1) - is_stopped = is_stopped + next_tokens.eq(stop_token_index).squeeze() - if is_stopped.all(): - break - - scores = scores / seq_lengths - order = scores.argsort(descending=True) - # tokens tensors are already padded to max_seq_length - output_texts = [tokens[i] for i in order] - output_texts = torch.stack(output_texts, dim=0) - seq_lengths = torch.tensor([seq_lengths[i] for i in order], dtype=seq_lengths.dtype) - return output_texts, seq_lengths diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/modeling_uvit.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/modeling_uvit.py deleted file mode 100644 index 2a04ec2e40309d41e4dea041c7f3398c351bc5dd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/modeling_uvit.py +++ /dev/null @@ -1,1197 +0,0 @@ -import math -from typing import Optional, Union - -import torch -from torch import nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models import ModelMixin -from ...models.attention import FeedForward -from ...models.attention_processor import Attention -from ...models.embeddings import TimestepEmbedding, Timesteps, get_2d_sincos_pos_embed -from ...models.modeling_outputs import Transformer2DModelOutput -from ...models.normalization import AdaLayerNorm -from ...utils import logging - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -def _no_grad_trunc_normal_(tensor, mean, std, a, b): - # Cut & paste from PyTorch official master until it's in a few official releases - RW - # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf - def norm_cdf(x): - # Computes standard normal cumulative distribution function - return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0 - - if (mean < a - 2 * std) or (mean > b + 2 * std): - logger.warning( - "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " - "The distribution of values may be incorrect." - ) - - with torch.no_grad(): - # Values are generated by using a truncated uniform distribution and - # then using the inverse CDF for the normal distribution. - # Get upper and lower cdf values - l = norm_cdf((a - mean) / std) - u = norm_cdf((b - mean) / std) - - # Uniformly fill tensor with values from [l, u], then translate to - # [2l-1, 2u-1]. - tensor.uniform_(2 * l - 1, 2 * u - 1) - - # Use inverse cdf transform for normal distribution to get truncated - # standard normal - tensor.erfinv_() - - # Transform to proper mean, std - tensor.mul_(std * math.sqrt(2.0)) - tensor.add_(mean) - - # Clamp to ensure it's in the proper range - tensor.clamp_(min=a, max=b) - return tensor - - -def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0): - # type: (torch.Tensor, float, float, float, float) -> torch.Tensor - r"""Fills the input Tensor with values drawn from a truncated - normal distribution. The values are effectively drawn from the normal distribution :math:`\mathcal{N}(\text{mean}, - \text{std}^2)` with values outside :math:`[a, b]` redrawn until they are within the bounds. The method used for - generating the random values works best when :math:`a \leq \text{mean} \leq b`. - - Args: - tensor: an n-dimensional `torch.Tensor` - mean: the mean of the normal distribution - std: the standard deviation of the normal distribution - a: the minimum cutoff value - b: the maximum cutoff value - Examples: - >>> w = torch.empty(3, 5) >>> nn.init.trunc_normal_(w) - """ - return _no_grad_trunc_normal_(tensor, mean, std, a, b) - - -class PatchEmbed(nn.Module): - """2D Image to Patch Embedding""" - - def __init__( - self, - height=224, - width=224, - patch_size=16, - in_channels=3, - embed_dim=768, - layer_norm=False, - flatten=True, - bias=True, - use_pos_embed=True, - ): - super().__init__() - - num_patches = (height // patch_size) * (width // patch_size) - self.flatten = flatten - self.layer_norm = layer_norm - - self.proj = nn.Conv2d( - in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias - ) - if layer_norm: - self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6) - else: - self.norm = None - - self.use_pos_embed = use_pos_embed - if self.use_pos_embed: - pos_embed = get_2d_sincos_pos_embed(embed_dim, int(num_patches**0.5), output_type="pt") - self.register_buffer("pos_embed", pos_embed.float().unsqueeze(0), persistent=False) - - def forward(self, latent): - latent = self.proj(latent) - if self.flatten: - latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC - if self.layer_norm: - latent = self.norm(latent) - if self.use_pos_embed: - return latent + self.pos_embed - else: - return latent - - -class SkipBlock(nn.Module): - def __init__(self, dim: int): - super().__init__() - - self.skip_linear = nn.Linear(2 * dim, dim) - - # Use torch.nn.LayerNorm for now, following the original code - self.norm = nn.LayerNorm(dim) - - def forward(self, x, skip): - x = self.skip_linear(torch.cat([x, skip], dim=-1)) - x = self.norm(x) - - return x - - -# Modified to support both pre-LayerNorm and post-LayerNorm configurations -# Don't support AdaLayerNormZero for now -# Modified from diffusers.models.attention.BasicTransformerBlock -class UTransformerBlock(nn.Module): - r""" - A modification of BasicTransformerBlock which supports pre-LayerNorm and post-LayerNorm configurations. - - Parameters: - 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. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. - activation_fn (`str`, *optional*, defaults to `"geglu"`): - Activation function to be used in feed-forward. - num_embeds_ada_norm (:obj: `int`, *optional*): - The number of diffusion steps used during training. See `Transformer2DModel`. - attention_bias (:obj: `bool`, *optional*, defaults to `False`): - Configure if the attentions should contain a bias parameter. - only_cross_attention (`bool`, *optional*): - Whether to use only cross-attention layers. In this case two cross attention layers are used. - double_self_attention (`bool`, *optional*): - Whether to use two self-attention layers. In this case no cross attention layers are used. - upcast_attention (`bool`, *optional*): - Whether to upcast the query and key to float32 when performing the attention calculation. - norm_elementwise_affine (`bool`, *optional*): - Whether to use learnable per-element affine parameters during layer normalization. - norm_type (`str`, defaults to `"layer_norm"`): - The layer norm implementation to use. - pre_layer_norm (`bool`, *optional*): - Whether to perform layer normalization before the attention and feedforward operations ("pre-LayerNorm"), - as opposed to after ("post-LayerNorm"). Note that `BasicTransformerBlock` uses pre-LayerNorm, e.g. - `pre_layer_norm = True`. - final_dropout (`bool`, *optional*): - Whether to use a final Dropout layer after the feedforward network. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - dropout=0.0, - cross_attention_dim: Optional[int] = None, - activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, - attention_bias: bool = False, - only_cross_attention: bool = False, - double_self_attention: bool = False, - upcast_attention: bool = False, - norm_elementwise_affine: bool = True, - norm_type: str = "layer_norm", - pre_layer_norm: bool = True, - final_dropout: bool = False, - ): - super().__init__() - self.only_cross_attention = only_cross_attention - - self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" - - self.pre_layer_norm = pre_layer_norm - - if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: - raise ValueError( - f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" - f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." - ) - - # 1. Self-Attn - self.attn1 = Attention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - cross_attention_dim=cross_attention_dim if only_cross_attention else None, - upcast_attention=upcast_attention, - ) - - # 2. Cross-Attn - if cross_attention_dim is not None or double_self_attention: - self.attn2 = Attention( - query_dim=dim, - cross_attention_dim=cross_attention_dim if not double_self_attention else None, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - upcast_attention=upcast_attention, - ) # is self-attn if encoder_hidden_states is none - else: - self.attn2 = None - - if self.use_ada_layer_norm: - self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) - else: - self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) - - if cross_attention_dim is not None or double_self_attention: - # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. - # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during - # the second cross attention block. - self.norm2 = ( - AdaLayerNorm(dim, num_embeds_ada_norm) - if self.use_ada_layer_norm - else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) - ) - else: - self.norm2 = None - - # 3. Feed-forward - self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) - self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout) - - def forward( - self, - hidden_states, - attention_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - timestep=None, - cross_attention_kwargs=None, - class_labels=None, - ): - # Pre-LayerNorm - if self.pre_layer_norm: - if self.use_ada_layer_norm: - norm_hidden_states = self.norm1(hidden_states, timestep) - else: - norm_hidden_states = self.norm1(hidden_states) - else: - norm_hidden_states = hidden_states - - # 1. Self-Attention - cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} - attn_output = self.attn1( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, - attention_mask=attention_mask, - **cross_attention_kwargs, - ) - - # Post-LayerNorm - if not self.pre_layer_norm: - if self.use_ada_layer_norm: - attn_output = self.norm1(attn_output, timestep) - else: - attn_output = self.norm1(attn_output) - - hidden_states = attn_output + hidden_states - - if self.attn2 is not None: - # Pre-LayerNorm - if self.pre_layer_norm: - norm_hidden_states = ( - self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states) - ) - else: - norm_hidden_states = hidden_states - # TODO (Birch-San): Here we should prepare the encoder_attention mask correctly - # prepare attention mask here - - # 2. Cross-Attention - attn_output = self.attn2( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - **cross_attention_kwargs, - ) - - # Post-LayerNorm - if not self.pre_layer_norm: - attn_output = self.norm2(attn_output, timestep) if self.use_ada_layer_norm else self.norm2(attn_output) - - hidden_states = attn_output + hidden_states - - # 3. Feed-forward - # Pre-LayerNorm - if self.pre_layer_norm: - norm_hidden_states = self.norm3(hidden_states) - else: - norm_hidden_states = hidden_states - - ff_output = self.ff(norm_hidden_states) - - # Post-LayerNorm - if not self.pre_layer_norm: - ff_output = self.norm3(ff_output) - - hidden_states = ff_output + hidden_states - - return hidden_states - - -# Like UTransformerBlock except with LayerNorms on the residual backbone of the block -# Modified from diffusers.models.attention.BasicTransformerBlock -class UniDiffuserBlock(nn.Module): - r""" - A modification of BasicTransformerBlock which supports pre-LayerNorm and post-LayerNorm configurations and puts the - LayerNorms on the residual backbone of the block. This matches the transformer block in the [original UniDiffuser - implementation](https://github.com/thu-ml/unidiffuser/blob/main/libs/uvit_multi_post_ln_v1.py#L104). - - Parameters: - 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. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. - activation_fn (`str`, *optional*, defaults to `"geglu"`): - Activation function to be used in feed-forward. - num_embeds_ada_norm (:obj: `int`, *optional*): - The number of diffusion steps used during training. See `Transformer2DModel`. - attention_bias (:obj: `bool`, *optional*, defaults to `False`): - Configure if the attentions should contain a bias parameter. - only_cross_attention (`bool`, *optional*): - Whether to use only cross-attention layers. In this case two cross attention layers are used. - double_self_attention (`bool`, *optional*): - Whether to use two self-attention layers. In this case no cross attention layers are used. - upcast_attention (`bool`, *optional*): - Whether to upcast the query and key to float() when performing the attention calculation. - norm_elementwise_affine (`bool`, *optional*): - Whether to use learnable per-element affine parameters during layer normalization. - norm_type (`str`, defaults to `"layer_norm"`): - The layer norm implementation to use. - pre_layer_norm (`bool`, *optional*): - Whether to perform layer normalization before the attention and feedforward operations ("pre-LayerNorm"), - as opposed to after ("post-LayerNorm"). The original UniDiffuser implementation is post-LayerNorm - (`pre_layer_norm = False`). - final_dropout (`bool`, *optional*): - Whether to use a final Dropout layer after the feedforward network. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - dropout=0.0, - cross_attention_dim: Optional[int] = None, - activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, - attention_bias: bool = False, - only_cross_attention: bool = False, - double_self_attention: bool = False, - upcast_attention: bool = False, - norm_elementwise_affine: bool = True, - norm_type: str = "layer_norm", - pre_layer_norm: bool = False, - final_dropout: bool = True, - ): - super().__init__() - self.only_cross_attention = only_cross_attention - - self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" - - self.pre_layer_norm = pre_layer_norm - - if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: - raise ValueError( - f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" - f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." - ) - - # 1. Self-Attn - self.attn1 = Attention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - cross_attention_dim=cross_attention_dim if only_cross_attention else None, - upcast_attention=upcast_attention, - ) - - # 2. Cross-Attn - if cross_attention_dim is not None or double_self_attention: - self.attn2 = Attention( - query_dim=dim, - cross_attention_dim=cross_attention_dim if not double_self_attention else None, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - upcast_attention=upcast_attention, - ) # is self-attn if encoder_hidden_states is none - else: - self.attn2 = None - - if self.use_ada_layer_norm: - self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) - else: - self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) - - if cross_attention_dim is not None or double_self_attention: - # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. - # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during - # the second cross attention block. - self.norm2 = ( - AdaLayerNorm(dim, num_embeds_ada_norm) - if self.use_ada_layer_norm - else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) - ) - else: - self.norm2 = None - - # 3. Feed-forward - self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) - self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout) - - def forward( - self, - hidden_states, - attention_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - timestep=None, - cross_attention_kwargs=None, - class_labels=None, - ): - # Following the diffusers transformer block implementation, put the LayerNorm on the - # residual backbone - # Pre-LayerNorm - if self.pre_layer_norm: - if self.use_ada_layer_norm: - hidden_states = self.norm1(hidden_states, timestep) - else: - hidden_states = self.norm1(hidden_states) - - # 1. Self-Attention - cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} - attn_output = self.attn1( - hidden_states, - encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, - attention_mask=attention_mask, - **cross_attention_kwargs, - ) - - hidden_states = attn_output + hidden_states - - # Following the diffusers transformer block implementation, put the LayerNorm on the - # residual backbone - # Post-LayerNorm - if not self.pre_layer_norm: - if self.use_ada_layer_norm: - hidden_states = self.norm1(hidden_states, timestep) - else: - hidden_states = self.norm1(hidden_states) - - if self.attn2 is not None: - # Pre-LayerNorm - if self.pre_layer_norm: - hidden_states = ( - self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states) - ) - # TODO (Birch-San): Here we should prepare the encoder_attention mask correctly - # prepare attention mask here - - # 2. Cross-Attention - attn_output = self.attn2( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - **cross_attention_kwargs, - ) - - hidden_states = attn_output + hidden_states - - # Post-LayerNorm - if not self.pre_layer_norm: - hidden_states = ( - self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states) - ) - - # 3. Feed-forward - # Pre-LayerNorm - if self.pre_layer_norm: - hidden_states = self.norm3(hidden_states) - - ff_output = self.ff(hidden_states) - - hidden_states = ff_output + hidden_states - - # Post-LayerNorm - if not self.pre_layer_norm: - hidden_states = self.norm3(hidden_states) - - return hidden_states - - -# Modified from diffusers.models.transformer_2d.Transformer2DModel -# Modify the transformer block structure to be U-Net like following U-ViT -# Only supports patch-style input and torch.nn.LayerNorm currently -# https://github.com/baofff/U-ViT -class UTransformer2DModel(ModelMixin, ConfigMixin): - """ - Transformer model based on the [U-ViT](https://github.com/baofff/U-ViT) architecture for image-like data. Compared - to [`Transformer2DModel`], this model has skip connections between transformer blocks in a "U"-shaped fashion, - similar to a U-Net. Supports only continuous (actual embeddings) inputs, which are embedded via a [`PatchEmbed`] - layer and then reshaped to (b, t, d). - - Parameters: - num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. - attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. - in_channels (`int`, *optional*): - Pass if the input is continuous. The number of channels in the input. - out_channels (`int`, *optional*): - The number of output channels; if `None`, defaults to `in_channels`. - num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - norm_num_groups (`int`, *optional*, defaults to `32`): - The number of groups to use when performing Group Normalization. - cross_attention_dim (`int`, *optional*): The number of encoder_hidden_states dimensions to use. - attention_bias (`bool`, *optional*): - Configure if the TransformerBlocks' attention should contain a bias parameter. - sample_size (`int`, *optional*): Pass if the input is discrete. The width of the latent images. - Note that this is fixed at training time as it is used for learning a number of position embeddings. See - `ImagePositionalEmbeddings`. - num_vector_embeds (`int`, *optional*): - Pass if the input is discrete. The number of classes of the vector embeddings of the latent pixels. - Includes the class for the masked latent pixel. - patch_size (`int`, *optional*, defaults to 2): - The patch size to use in the patch embedding. - activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. - num_embeds_ada_norm ( `int`, *optional*): Pass if at least one of the norm_layers is `AdaLayerNorm`. - The number of diffusion steps used during training. Note that this is fixed at training time as it is used - to learn a number of embeddings that are added to the hidden states. During inference, you can denoise for - up to but not more than steps than `num_embeds_ada_norm`. - use_linear_projection (int, *optional*): TODO: Not used - only_cross_attention (`bool`, *optional*): - Whether to use only cross-attention layers. In this case two cross attention layers are used in each - transformer block. - upcast_attention (`bool`, *optional*): - Whether to upcast the query and key to float() when performing the attention calculation. - norm_type (`str`, *optional*, defaults to `"layer_norm"`): - The Layer Normalization implementation to use. Defaults to `torch.nn.LayerNorm`. - block_type (`str`, *optional*, defaults to `"unidiffuser"`): - The transformer block implementation to use. If `"unidiffuser"`, has the LayerNorms on the residual - backbone of each transformer block; otherwise has them in the attention/feedforward branches (the standard - behavior in `diffusers`.) - pre_layer_norm (`bool`, *optional*): - Whether to perform layer normalization before the attention and feedforward operations ("pre-LayerNorm"), - as opposed to after ("post-LayerNorm"). The original UniDiffuser implementation is post-LayerNorm - (`pre_layer_norm = False`). - norm_elementwise_affine (`bool`, *optional*): - Whether to use learnable per-element affine parameters during layer normalization. - use_patch_pos_embed (`bool`, *optional*): - Whether to use position embeddings inside the patch embedding layer (`PatchEmbed`). - final_dropout (`bool`, *optional*): - Whether to use a final Dropout layer after the feedforward network. - """ - - @register_to_config - def __init__( - self, - num_attention_heads: int = 16, - attention_head_dim: int = 88, - in_channels: Optional[int] = None, - out_channels: Optional[int] = None, - num_layers: int = 1, - dropout: float = 0.0, - norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = None, - attention_bias: bool = False, - sample_size: Optional[int] = None, - num_vector_embeds: Optional[int] = None, - patch_size: Optional[int] = 2, - activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, - use_linear_projection: bool = False, - only_cross_attention: bool = False, - upcast_attention: bool = False, - norm_type: str = "layer_norm", - block_type: str = "unidiffuser", - pre_layer_norm: bool = False, - norm_elementwise_affine: bool = True, - use_patch_pos_embed=False, - ff_final_dropout: bool = False, - ): - super().__init__() - self.use_linear_projection = use_linear_projection - self.num_attention_heads = num_attention_heads - self.attention_head_dim = attention_head_dim - inner_dim = num_attention_heads * attention_head_dim - - # 1. Input - # Only support patch input of shape (batch_size, num_channels, height, width) for now - assert in_channels is not None and patch_size is not None, "Patch input requires in_channels and patch_size." - - assert sample_size is not None, "UTransformer2DModel over patched input must provide sample_size" - - # 2. Define input layers - self.height = sample_size - self.width = sample_size - - self.patch_size = patch_size - self.pos_embed = PatchEmbed( - height=sample_size, - width=sample_size, - patch_size=patch_size, - in_channels=in_channels, - embed_dim=inner_dim, - use_pos_embed=use_patch_pos_embed, - ) - - # 3. Define transformers blocks - # Modify this to have in_blocks ("downsample" blocks, even though we don't actually downsample), a mid_block, - # and out_blocks ("upsample" blocks). Like a U-Net, there are skip connections from in_blocks to out_blocks in - # a "U"-shaped fashion (e.g. first in_block to last out_block, etc.). - # Quick hack to make the transformer block type configurable - if block_type == "unidiffuser": - block_cls = UniDiffuserBlock - else: - block_cls = UTransformerBlock - self.transformer_in_blocks = nn.ModuleList( - [ - block_cls( - inner_dim, - num_attention_heads, - attention_head_dim, - dropout=dropout, - cross_attention_dim=cross_attention_dim, - activation_fn=activation_fn, - num_embeds_ada_norm=num_embeds_ada_norm, - attention_bias=attention_bias, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - norm_type=norm_type, - pre_layer_norm=pre_layer_norm, - norm_elementwise_affine=norm_elementwise_affine, - final_dropout=ff_final_dropout, - ) - for d in range(num_layers // 2) - ] - ) - - self.transformer_mid_block = block_cls( - inner_dim, - num_attention_heads, - attention_head_dim, - dropout=dropout, - cross_attention_dim=cross_attention_dim, - activation_fn=activation_fn, - num_embeds_ada_norm=num_embeds_ada_norm, - attention_bias=attention_bias, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - norm_type=norm_type, - pre_layer_norm=pre_layer_norm, - norm_elementwise_affine=norm_elementwise_affine, - final_dropout=ff_final_dropout, - ) - - # For each skip connection, we use a SkipBlock (concatenation + Linear + LayerNorm) to process the inputs - # before each transformer out_block. - self.transformer_out_blocks = nn.ModuleList( - [ - nn.ModuleDict( - { - "skip": SkipBlock( - inner_dim, - ), - "block": block_cls( - inner_dim, - num_attention_heads, - attention_head_dim, - dropout=dropout, - cross_attention_dim=cross_attention_dim, - activation_fn=activation_fn, - num_embeds_ada_norm=num_embeds_ada_norm, - attention_bias=attention_bias, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - norm_type=norm_type, - pre_layer_norm=pre_layer_norm, - norm_elementwise_affine=norm_elementwise_affine, - final_dropout=ff_final_dropout, - ), - } - ) - for d in range(num_layers // 2) - ] - ) - - # 4. Define output layers - self.out_channels = in_channels if out_channels is None else out_channels - - # Following the UniDiffuser U-ViT implementation, we process the transformer output with - # a LayerNorm layer with per-element affine params - self.norm_out = nn.LayerNorm(inner_dim) - - def forward( - self, - hidden_states, - encoder_hidden_states=None, - timestep=None, - class_labels=None, - cross_attention_kwargs=None, - return_dict: bool = True, - hidden_states_is_embedding: bool = False, - unpatchify: bool = True, - ): - """ - Args: - hidden_states ( When discrete, `torch.LongTensor` of shape `(batch size, num latent pixels)`. - When continuous, `torch.Tensor` of shape `(batch size, channel, height, width)`): Input hidden_states - encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*): - Conditional embeddings for cross attention layer. If not given, cross-attention defaults to - self-attention. - timestep ( `torch.long`, *optional*): - Optional timestep to be applied as an embedding in AdaLayerNorm's. Used to indicate denoising step. - class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*): - Optional class labels to be applied as an embedding in AdaLayerZeroNorm. Used to indicate class labels - conditioning. - cross_attention_kwargs (*optional*): - Keyword arguments to supply to the cross attention layers, if used. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain - tuple. - hidden_states_is_embedding (`bool`, *optional*, defaults to `False`): - Whether or not hidden_states is an embedding directly usable by the transformer. In this case we will - ignore input handling (e.g. continuous, vectorized, etc.) and directly feed hidden_states into the - transformer blocks. - unpatchify (`bool`, *optional*, defaults to `True`): - Whether to unpatchify the transformer output. - - Returns: - [`~models.transformer_2d.Transformer2DModelOutput`] or `tuple`: - [`~models.transformer_2d.Transformer2DModelOutput`] if `return_dict` is True, otherwise a `tuple`. When - returning a tuple, the first element is the sample tensor. - """ - # 0. Check inputs - - if not unpatchify and return_dict: - raise ValueError( - f"Cannot both define `unpatchify`: {unpatchify} and `return_dict`: {return_dict} since when" - f" `unpatchify` is {unpatchify} the returned output is of shape (batch_size, seq_len, hidden_dim)" - " rather than (batch_size, num_channels, height, width)." - ) - - # 1. Input - if not hidden_states_is_embedding: - hidden_states = self.pos_embed(hidden_states) - - # 2. Blocks - - # In ("downsample") blocks - skips = [] - for in_block in self.transformer_in_blocks: - hidden_states = in_block( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=timestep, - cross_attention_kwargs=cross_attention_kwargs, - class_labels=class_labels, - ) - skips.append(hidden_states) - - # Mid block - hidden_states = self.transformer_mid_block(hidden_states) - - # Out ("upsample") blocks - for out_block in self.transformer_out_blocks: - hidden_states = out_block["skip"](hidden_states, skips.pop()) - hidden_states = out_block["block"]( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=timestep, - cross_attention_kwargs=cross_attention_kwargs, - class_labels=class_labels, - ) - - # 3. Output - # Don't support AdaLayerNorm for now, so no conditioning/scale/shift logic - hidden_states = self.norm_out(hidden_states) - # hidden_states = self.proj_out(hidden_states) - - if unpatchify: - # unpatchify - height = width = int(hidden_states.shape[1] ** 0.5) - hidden_states = hidden_states.reshape( - shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels) - ) - hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) - output = hidden_states.reshape( - shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size) - ) - else: - output = hidden_states - - if not return_dict: - return (output,) - - return Transformer2DModelOutput(sample=output) - - -class UniDiffuserModel(ModelMixin, ConfigMixin): - """ - Transformer model for a image-text [UniDiffuser](https://huggingface.co/papers/2303.06555) model. This is a - modification of [`UTransformer2DModel`] with input and output heads for the VAE-embedded latent image, the - CLIP-embedded image, and the CLIP-embedded prompt (see paper for more details). - - Parameters: - text_dim (`int`): The hidden dimension of the CLIP text model used to embed images. - clip_img_dim (`int`): The hidden dimension of the CLIP vision model used to embed prompts. - num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. - attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. - in_channels (`int`, *optional*): - Pass if the input is continuous. The number of channels in the input. - out_channels (`int`, *optional*): - The number of output channels; if `None`, defaults to `in_channels`. - num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - norm_num_groups (`int`, *optional*, defaults to `32`): - The number of groups to use when performing Group Normalization. - cross_attention_dim (`int`, *optional*): The number of encoder_hidden_states dimensions to use. - attention_bias (`bool`, *optional*): - Configure if the TransformerBlocks' attention should contain a bias parameter. - sample_size (`int`, *optional*): Pass if the input is discrete. The width of the latent images. - Note that this is fixed at training time as it is used for learning a number of position embeddings. See - `ImagePositionalEmbeddings`. - num_vector_embeds (`int`, *optional*): - Pass if the input is discrete. The number of classes of the vector embeddings of the latent pixels. - Includes the class for the masked latent pixel. - patch_size (`int`, *optional*, defaults to 2): - The patch size to use in the patch embedding. - activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. - num_embeds_ada_norm ( `int`, *optional*): Pass if at least one of the norm_layers is `AdaLayerNorm`. - The number of diffusion steps used during training. Note that this is fixed at training time as it is used - to learn a number of embeddings that are added to the hidden states. During inference, you can denoise for - up to but not more than steps than `num_embeds_ada_norm`. - use_linear_projection (int, *optional*): TODO: Not used - only_cross_attention (`bool`, *optional*): - Whether to use only cross-attention layers. In this case two cross attention layers are used in each - transformer block. - upcast_attention (`bool`, *optional*): - Whether to upcast the query and key to float32 when performing the attention calculation. - norm_type (`str`, *optional*, defaults to `"layer_norm"`): - The Layer Normalization implementation to use. Defaults to `torch.nn.LayerNorm`. - block_type (`str`, *optional*, defaults to `"unidiffuser"`): - The transformer block implementation to use. If `"unidiffuser"`, has the LayerNorms on the residual - backbone of each transformer block; otherwise has them in the attention/feedforward branches (the standard - behavior in `diffusers`.) - pre_layer_norm (`bool`, *optional*): - Whether to perform layer normalization before the attention and feedforward operations ("pre-LayerNorm"), - as opposed to after ("post-LayerNorm"). The original UniDiffuser implementation is post-LayerNorm - (`pre_layer_norm = False`). - norm_elementwise_affine (`bool`, *optional*): - Whether to use learnable per-element affine parameters during layer normalization. - use_patch_pos_embed (`bool`, *optional*): - Whether to use position embeddings inside the patch embedding layer (`PatchEmbed`). - ff_final_dropout (`bool`, *optional*): - Whether to use a final Dropout layer after the feedforward network. - use_data_type_embedding (`bool`, *optional*): - Whether to use a data type embedding. This is only relevant for UniDiffuser-v1 style models; UniDiffuser-v1 - is continue-trained from UniDiffuser-v0 on non-publically-available data and accepts a `data_type` - argument, which can either be `1` to use the weights trained on non-publically-available data or `0` - otherwise. This argument is subsequently embedded by the data type embedding, if used. - """ - - @register_to_config - def __init__( - self, - text_dim: int = 768, - clip_img_dim: int = 512, - 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, - num_layers: int = 1, - dropout: float = 0.0, - norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = None, - attention_bias: bool = False, - sample_size: Optional[int] = None, - num_vector_embeds: Optional[int] = None, - patch_size: Optional[int] = None, - activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, - use_linear_projection: bool = False, - only_cross_attention: bool = False, - upcast_attention: bool = False, - norm_type: str = "layer_norm", - block_type: str = "unidiffuser", - pre_layer_norm: bool = False, - use_timestep_embedding=False, - norm_elementwise_affine: bool = True, - use_patch_pos_embed=False, - ff_final_dropout: bool = True, - use_data_type_embedding: bool = False, - ): - super().__init__() - - # 0. Handle dimensions - self.inner_dim = num_attention_heads * attention_head_dim - - assert sample_size is not None, "UniDiffuserModel over patched input must provide sample_size" - self.sample_size = sample_size - self.in_channels = in_channels - self.out_channels = in_channels if out_channels is None else out_channels - - self.patch_size = patch_size - # Assume image is square... - self.num_patches = (self.sample_size // patch_size) * (self.sample_size // patch_size) - - # 1. Define input layers - # 1.1 Input layers for text and image input - # For now, only support patch input for VAE latent image input - self.vae_img_in = PatchEmbed( - height=sample_size, - width=sample_size, - patch_size=patch_size, - in_channels=in_channels, - embed_dim=self.inner_dim, - use_pos_embed=use_patch_pos_embed, - ) - self.clip_img_in = nn.Linear(clip_img_dim, self.inner_dim) - self.text_in = nn.Linear(text_dim, self.inner_dim) - - # 1.2. Timestep embeddings for t_img, t_text - self.timestep_img_proj = Timesteps( - self.inner_dim, - flip_sin_to_cos=True, - downscale_freq_shift=0, - ) - self.timestep_img_embed = ( - TimestepEmbedding( - self.inner_dim, - 4 * self.inner_dim, - out_dim=self.inner_dim, - ) - if use_timestep_embedding - else nn.Identity() - ) - - self.timestep_text_proj = Timesteps( - self.inner_dim, - flip_sin_to_cos=True, - downscale_freq_shift=0, - ) - self.timestep_text_embed = ( - TimestepEmbedding( - self.inner_dim, - 4 * self.inner_dim, - out_dim=self.inner_dim, - ) - if use_timestep_embedding - else nn.Identity() - ) - - # 1.3. Positional embedding - self.num_text_tokens = num_text_tokens - self.num_tokens = 1 + 1 + num_text_tokens + 1 + self.num_patches - self.pos_embed = nn.Parameter(torch.zeros(1, self.num_tokens, self.inner_dim)) - self.pos_embed_drop = nn.Dropout(p=dropout) - trunc_normal_(self.pos_embed, std=0.02) - - # 1.4. Handle data type token embeddings for UniDiffuser-V1, if necessary - self.use_data_type_embedding = use_data_type_embedding - if self.use_data_type_embedding: - self.data_type_token_embedding = nn.Embedding(2, self.inner_dim) - self.data_type_pos_embed_token = nn.Parameter(torch.zeros(1, 1, self.inner_dim)) - - # 2. Define transformer blocks - self.transformer = UTransformer2DModel( - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - in_channels=in_channels, - out_channels=out_channels, - num_layers=num_layers, - dropout=dropout, - norm_num_groups=norm_num_groups, - cross_attention_dim=cross_attention_dim, - attention_bias=attention_bias, - sample_size=sample_size, - num_vector_embeds=num_vector_embeds, - patch_size=patch_size, - activation_fn=activation_fn, - num_embeds_ada_norm=num_embeds_ada_norm, - use_linear_projection=use_linear_projection, - only_cross_attention=only_cross_attention, - upcast_attention=upcast_attention, - norm_type=norm_type, - block_type=block_type, - pre_layer_norm=pre_layer_norm, - norm_elementwise_affine=norm_elementwise_affine, - use_patch_pos_embed=use_patch_pos_embed, - ff_final_dropout=ff_final_dropout, - ) - - # 3. Define output layers - patch_dim = (patch_size**2) * out_channels - self.vae_img_out = nn.Linear(self.inner_dim, patch_dim) - self.clip_img_out = nn.Linear(self.inner_dim, clip_img_dim) - self.text_out = nn.Linear(self.inner_dim, text_dim) - - @torch.jit.ignore - def no_weight_decay(self): - return {"pos_embed"} - - def forward( - self, - 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, - encoder_hidden_states=None, - cross_attention_kwargs=None, - ): - """ - Args: - latent_image_embeds (`torch.Tensor` of shape `(batch size, latent channels, height, width)`): - Latent image representation from the VAE encoder. - image_embeds (`torch.Tensor` of shape `(batch size, 1, clip_img_dim)`): - CLIP-embedded image representation (unsqueezed in the first dimension). - prompt_embeds (`torch.Tensor` of shape `(batch size, seq_len, text_dim)`): - CLIP-embedded text representation. - timestep_img (`torch.long` or `float` or `int`): - Current denoising step for the image. - timestep_text (`torch.long` or `float` or `int`): - Current denoising step for the text. - data_type: (`torch.int` or `float` or `int`, *optional*, defaults to `1`): - Only used in UniDiffuser-v1-style models. Can be either `1`, to use weights trained on nonpublic data, - or `0` otherwise. - encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*): - Conditional embeddings for cross attention layer. If not given, cross-attention defaults to - self-attention. - cross_attention_kwargs (*optional*): - Keyword arguments to supply to the cross attention layers, if used. - - - Returns: - `tuple`: Returns relevant parts of the model's noise prediction: the first element of the tuple is tbe VAE - image embedding, the second element is the CLIP image embedding, and the third element is the CLIP text - embedding. - """ - batch_size = latent_image_embeds.shape[0] - - # 1. Input - # 1.1. Map inputs to shape (B, N, inner_dim) - vae_hidden_states = self.vae_img_in(latent_image_embeds) - clip_hidden_states = self.clip_img_in(image_embeds) - text_hidden_states = self.text_in(prompt_embeds) - - num_text_tokens, num_img_tokens = text_hidden_states.size(1), vae_hidden_states.size(1) - - # 1.2. Encode image timesteps to single token (B, 1, inner_dim) - if not torch.is_tensor(timestep_img): - timestep_img = torch.tensor([timestep_img], dtype=torch.long, device=vae_hidden_states.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep_img = timestep_img * torch.ones(batch_size, dtype=timestep_img.dtype, device=timestep_img.device) - - timestep_img_token = self.timestep_img_proj(timestep_img) - # t_img_token does not contain any weights and will always return f32 tensors - # but time_embedding might be fp16, so we need to cast here. - timestep_img_token = timestep_img_token.to(dtype=self.dtype) - timestep_img_token = self.timestep_img_embed(timestep_img_token) - timestep_img_token = timestep_img_token.unsqueeze(dim=1) - - # 1.3. Encode text timesteps to single token (B, 1, inner_dim) - if not torch.is_tensor(timestep_text): - timestep_text = torch.tensor([timestep_text], dtype=torch.long, device=vae_hidden_states.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep_text = timestep_text * torch.ones(batch_size, dtype=timestep_text.dtype, device=timestep_text.device) - - timestep_text_token = self.timestep_text_proj(timestep_text) - # t_text_token does not contain any weights and will always return f32 tensors - # but time_embedding might be fp16, so we need to cast here. - timestep_text_token = timestep_text_token.to(dtype=self.dtype) - timestep_text_token = self.timestep_text_embed(timestep_text_token) - timestep_text_token = timestep_text_token.unsqueeze(dim=1) - - # 1.4. Concatenate all of the embeddings together. - if self.use_data_type_embedding: - assert data_type is not None, "data_type must be supplied if the model uses a data type embedding" - if not torch.is_tensor(data_type): - data_type = torch.tensor([data_type], dtype=torch.int, device=vae_hidden_states.device) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - data_type = data_type * torch.ones(batch_size, dtype=data_type.dtype, device=data_type.device) - - data_type_token = self.data_type_token_embedding(data_type).unsqueeze(dim=1) - hidden_states = torch.cat( - [ - timestep_img_token, - timestep_text_token, - data_type_token, - text_hidden_states, - clip_hidden_states, - vae_hidden_states, - ], - dim=1, - ) - else: - hidden_states = torch.cat( - [timestep_img_token, timestep_text_token, text_hidden_states, clip_hidden_states, vae_hidden_states], - dim=1, - ) - - # 1.5. Prepare the positional embeddings and add to hidden states - # Note: I think img_vae should always have the proper shape, so there's no need to interpolate - # the position embeddings. - if self.use_data_type_embedding: - pos_embed = torch.cat( - [self.pos_embed[:, : 1 + 1, :], self.data_type_pos_embed_token, self.pos_embed[:, 1 + 1 :, :]], dim=1 - ) - else: - pos_embed = self.pos_embed - hidden_states = hidden_states + pos_embed - hidden_states = self.pos_embed_drop(hidden_states) - - # 2. Blocks - hidden_states = self.transformer( - hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=None, - class_labels=None, - cross_attention_kwargs=cross_attention_kwargs, - return_dict=False, - hidden_states_is_embedding=True, - unpatchify=False, - )[0] - - # 3. Output - # Split out the predicted noise representation. - if self.use_data_type_embedding: - ( - t_img_token_out, - t_text_token_out, - data_type_token_out, - text_out, - img_clip_out, - img_vae_out, - ) = hidden_states.split((1, 1, 1, num_text_tokens, 1, num_img_tokens), dim=1) - else: - t_img_token_out, t_text_token_out, text_out, img_clip_out, img_vae_out = hidden_states.split( - (1, 1, num_text_tokens, 1, num_img_tokens), dim=1 - ) - - img_vae_out = self.vae_img_out(img_vae_out) - - # unpatchify - height = width = int(img_vae_out.shape[1] ** 0.5) - img_vae_out = img_vae_out.reshape( - shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels) - ) - img_vae_out = torch.einsum("nhwpqc->nchpwq", img_vae_out) - img_vae_out = img_vae_out.reshape( - shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size) - ) - - img_clip_out = self.clip_img_out(img_clip_out) - - text_out = self.text_out(text_out) - - return img_vae_out, img_clip_out, text_out diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py deleted file mode 100644 index 40fd3b337301b3a08a955e34990f0cbc1fbf8ce6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py +++ /dev/null @@ -1,1438 +0,0 @@ -import inspect -from dataclasses import dataclass -from typing import Callable, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTokenizer, - CLIPVisionModelWithProjection, - GPT2Tokenizer, -) - -from ...image_processor import VaeImageProcessor -from ...loaders import StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from ...models import AutoencoderKL -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - is_torch_xla_available, - logging, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.outputs import BaseOutput -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline -from .modeling_text_decoder import UniDiffuserTextDecoder -from .modeling_uvit import UniDiffuserModel - - -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 - - -# New BaseOutput child class for joint image-text output -@dataclass -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, - 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 - length `batch_size`. - """ - - images: Optional[Union[List[PIL.Image.Image], np.ndarray]] - text: Optional[Union[List[str], List[List[str]]]] - - -class UniDiffuserPipeline(DeprecatedPipelineMixin, DiffusionPipeline): - r""" - Pipeline for a bimodal image-text model which supports unconditional text and image generation, text-conditioned - image generation, image-conditioned text generation, and joint image-text generation. - - 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: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. This - is part of the UniDiffuser image representation along with the CLIP vision encoding. - text_encoder ([`CLIPTextModel`]): - Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). - image_encoder ([`CLIPVisionModel`]): - A [`~transformers.CLIPVisionModel`] to encode images as part of its image representation along with the VAE - latent representation. - image_processor ([`CLIPImageProcessor`]): - [`~transformers.CLIPImageProcessor`] to preprocess an image before CLIP encoding it with `image_encoder`. - clip_tokenizer ([`CLIPTokenizer`]): - A [`~transformers.CLIPTokenizer`] to tokenize the prompt before encoding it with `text_encoder`. - text_decoder ([`UniDiffuserTextDecoder`]): - Frozen text decoder. This is a GPT-style model which is used to generate text from the UniDiffuser - embedding. - text_tokenizer ([`GPT2Tokenizer`]): - A [`~transformers.GPT2Tokenizer`] to decode text for text generation; used along with the `text_decoder`. - unet ([`UniDiffuserModel`]): - A [U-ViT](https://github.com/baofff/U-ViT) model with UNNet-style skip connections between transformer - layers to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image and/or text latents. The - original UniDiffuser paper uses the [`DPMSolverMultistepScheduler`] scheduler. - """ - - _last_supported_version = "0.33.1" - # TODO: support for moving submodules for components with enable_model_cpu_offload - model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae->text_decoder" - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - image_encoder: CLIPVisionModelWithProjection, - clip_image_processor: CLIPImageProcessor, - clip_tokenizer: CLIPTokenizer, - text_decoder: UniDiffuserTextDecoder, - text_tokenizer: GPT2Tokenizer, - unet: UniDiffuserModel, - scheduler: KarrasDiffusionSchedulers, - ): - super().__init__() - - if text_encoder.config.hidden_size != text_decoder.prefix_inner_dim: - raise ValueError( - f"The text encoder hidden size and text decoder prefix inner dim must be the same, but" - f" `text_encoder.config.hidden_size`: {text_encoder.config.hidden_size} and `text_decoder.prefix_inner_dim`: {text_decoder.prefix_inner_dim}" - ) - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - image_encoder=image_encoder, - clip_image_processor=clip_image_processor, - clip_tokenizer=clip_tokenizer, - text_decoder=text_decoder, - text_tokenizer=text_tokenizer, - unet=unet, - 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.num_channels_latents = vae.config.latent_channels - self.text_encoder_seq_len = text_encoder.config.max_position_embeddings - self.text_encoder_hidden_size = text_encoder.config.hidden_size - self.image_encoder_projection_dim = image_encoder.config.projection_dim - self.unet_resolution = unet.config.sample_size - - self.text_intermediate_dim = self.text_encoder_hidden_size - if self.text_decoder.prefix_hidden_dim is not None: - self.text_intermediate_dim = self.text_decoder.prefix_hidden_dim - - self.mode = None - - # TODO: handle safety checking? - self.safety_checker = None - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def _infer_mode(self, prompt, prompt_embeds, image, latents, prompt_latents, vae_latents, clip_latents): - r""" - Infer the generation task ('mode') from the inputs to `__call__`. If the mode has been manually set, the set - mode will be used. - """ - prompt_available = (prompt is not None) or (prompt_embeds is not None) - image_available = image is not None - input_available = prompt_available or image_available - - prompt_latents_available = prompt_latents is not None - vae_latents_available = vae_latents is not None - clip_latents_available = clip_latents is not None - full_latents_available = latents is not None - image_latents_available = vae_latents_available and clip_latents_available - all_indv_latents_available = prompt_latents_available and image_latents_available - - if self.mode is not None: - # Preferentially use the mode set by the user - mode = self.mode - elif prompt_available: - mode = "text2img" - elif image_available: - mode = "img2text" - else: - # Neither prompt nor image supplied, infer based on availability of latents - if full_latents_available or all_indv_latents_available: - mode = "joint" - elif prompt_latents_available: - mode = "text" - elif image_latents_available: - mode = "img" - else: - # No inputs or latents available - mode = "joint" - - # Give warnings for ambiguous cases - if self.mode is None and prompt_available and image_available: - logger.warning( - f"You have supplied both a text prompt and image to the pipeline and mode has not been set manually," - f" defaulting to mode '{mode}'." - ) - - if self.mode is None and not input_available: - if vae_latents_available != clip_latents_available: - # Exactly one of vae_latents and clip_latents is supplied - logger.warning( - f"You have supplied exactly one of `vae_latents` and `clip_latents`, whereas either both or none" - f" are expected to be supplied. Defaulting to mode '{mode}'." - ) - elif not prompt_latents_available and not vae_latents_available and not clip_latents_available: - # No inputs or latents supplied - logger.warning( - f"No inputs or latents have been supplied, and mode has not been manually set," - f" defaulting to mode '{mode}'." - ) - - return mode - - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.enable_vae_slicing - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.disable_vae_slicing - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.enable_vae_tiling - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.disable_vae_tiling - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - # Functions to manually set the mode - def set_text_mode(self): - r"""Manually set the generation mode to unconditional ("marginal") text generation.""" - self.mode = "text" - - def set_image_mode(self): - r"""Manually set the generation mode to unconditional ("marginal") image generation.""" - self.mode = "img" - - def set_text_to_image_mode(self): - r"""Manually set the generation mode to text-conditioned image generation.""" - self.mode = "text2img" - - def set_image_to_text_mode(self): - r"""Manually set the generation mode to image-conditioned text generation.""" - self.mode = "img2text" - - def set_joint_mode(self): - r"""Manually set the generation mode to unconditional joint image-text generation.""" - self.mode = "joint" - - def reset_mode(self): - r"""Removes a manually set mode; after calling this, the pipeline will infer the mode from inputs.""" - self.mode = None - - def _infer_batch_size( - self, - mode, - prompt, - prompt_embeds, - image, - num_images_per_prompt, - num_prompts_per_image, - latents, - prompt_latents, - vae_latents, - clip_latents, - ): - r"""Infers the batch size and multiplier depending on mode and supplied arguments to `__call__`.""" - if num_images_per_prompt is None: - num_images_per_prompt = 1 - if num_prompts_per_image is None: - num_prompts_per_image = 1 - - assert num_images_per_prompt > 0, "num_images_per_prompt must be a positive integer" - assert num_prompts_per_image > 0, "num_prompts_per_image must be a positive integer" - - if mode in ["text2img"]: - 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: - # Either prompt or prompt_embeds must be present for text2img. - batch_size = prompt_embeds.shape[0] - multiplier = num_images_per_prompt - elif mode in ["img2text"]: - if isinstance(image, PIL.Image.Image): - batch_size = 1 - else: - # Image must be available and type either PIL.Image.Image or torch.Tensor. - # Not currently supporting something like image_embeds. - batch_size = image.shape[0] - multiplier = num_prompts_per_image - elif mode in ["img"]: - if vae_latents is not None: - batch_size = vae_latents.shape[0] - elif clip_latents is not None: - batch_size = clip_latents.shape[0] - else: - batch_size = 1 - multiplier = num_images_per_prompt - elif mode in ["text"]: - if prompt_latents is not None: - batch_size = prompt_latents.shape[0] - else: - batch_size = 1 - multiplier = num_prompts_per_image - elif mode in ["joint"]: - if latents is not None: - batch_size = latents.shape[0] - elif prompt_latents is not None: - batch_size = prompt_latents.shape[0] - elif vae_latents is not None: - batch_size = vae_latents.shape[0] - elif clip_latents is not None: - batch_size = clip_latents.shape[0] - else: - batch_size = 1 - - if num_images_per_prompt == num_prompts_per_image: - multiplier = num_images_per_prompt - else: - multiplier = min(num_images_per_prompt, num_prompts_per_image) - logger.warning( - f"You are using mode `{mode}` and `num_images_per_prompt`: {num_images_per_prompt} and" - f" num_prompts_per_image: {num_prompts_per_image} are not equal. Using batch size equal to" - f" `min(num_images_per_prompt, num_prompts_per_image) = {batch_size}." - ) - return batch_size, multiplier - - # 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 with self.tokenizer->self.clip_tokenizer - 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.clip_tokenizer) - - text_inputs = self.clip_tokenizer( - prompt, - padding="max_length", - max_length=self.clip_tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - untruncated_ids = self.clip_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.clip_tokenizer.batch_decode( - untruncated_ids[:, self.clip_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.clip_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.clip_tokenizer) - - max_length = prompt_embeds.shape[1] - uncond_input = self.clip_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 - - # Modified from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_instruct_pix2pix.StableDiffusionInstructPix2PixPipeline.prepare_image_latents - # Add num_prompts_per_image argument, sample from autoencoder moment distribution - def encode_image_vae_latents( - self, - image, - batch_size, - num_prompts_per_image, - dtype, - device, - do_classifier_free_guidance, - generator=None, - ): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_prompts_per_image - 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 isinstance(generator, list): - image_latents = [ - self.vae.encode(image[i : i + 1]).latent_dist.sample(generator=generator[i]) - * self.vae.config.scaling_factor - for i in range(batch_size) - ] - image_latents = torch.cat(image_latents, dim=0) - else: - image_latents = self.vae.encode(image).latent_dist.sample(generator=generator) - # Scale image_latents by the VAE's scaling factor - image_latents = image_latents * self.vae.config.scaling_factor - - if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: - # expand image_latents for batch_size - deprecation_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {image_latents.shape[0]} initial" - " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" - " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" - " your script to pass as many initial images as text prompts to suppress this warning." - ) - deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) - 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 do_classifier_free_guidance: - uncond_image_latents = torch.zeros_like(image_latents) - image_latents = torch.cat([image_latents, image_latents, uncond_image_latents], dim=0) - - return image_latents - - def encode_image_clip_latents( - self, - image, - batch_size, - num_prompts_per_image, - dtype, - device, - generator=None, - ): - # Map image to CLIP embedding. - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - preprocessed_image = self.clip_image_processor.preprocess( - image, - return_tensors="pt", - ) - preprocessed_image = preprocessed_image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_prompts_per_image - if isinstance(generator, list): - image_latents = [ - self.image_encoder(**preprocessed_image[i : i + 1]).image_embeds for i in range(batch_size) - ] - image_latents = torch.cat(image_latents, dim=0) - else: - image_latents = self.image_encoder(**preprocessed_image).image_embeds - - if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: - # expand image_latents for batch_size - deprecation_message = ( - f"You have passed {batch_size} text prompts (`prompt`), but only {image_latents.shape[0]} initial" - " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" - " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" - " your script to pass as many initial images as text prompts to suppress this warning." - ) - deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) - 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 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." - ) - - return image_latents - - def prepare_text_latents( - self, batch_size, num_images_per_prompt, seq_len, hidden_size, dtype, device, generator, latents=None - ): - # Prepare latents for the CLIP embedded prompt. - shape = (batch_size * num_images_per_prompt, seq_len, hidden_size) - 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 is assumed to have shace (B, L, D) - latents = latents.repeat(num_images_per_prompt, 1, 1) - latents = latents.to(device=device, dtype=dtype) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - # Modified from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - # Rename prepare_latents -> prepare_image_vae_latents and add num_prompts_per_image argument. - def prepare_image_vae_latents( - self, - batch_size, - num_prompts_per_image, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=None, - ): - shape = ( - batch_size * num_prompts_per_image, - num_channels_latents, - height // self.vae_scale_factor, - 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 is assumed to have shape (B, C, H, W) - latents = latents.repeat(num_prompts_per_image, 1, 1, 1) - latents = latents.to(device=device, dtype=dtype) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def prepare_image_clip_latents( - self, batch_size, num_prompts_per_image, clip_img_dim, dtype, device, generator, latents=None - ): - # Prepare latents for the CLIP embedded image. - shape = (batch_size * num_prompts_per_image, 1, clip_img_dim) - 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 is assumed to have shape (B, L, D) - latents = latents.repeat(num_prompts_per_image, 1, 1) - latents = latents.to(device=device, dtype=dtype) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def decode_text_latents(self, text_latents, device): - output_token_list, seq_lengths = self.text_decoder.generate_captions( - text_latents, self.text_tokenizer.eos_token_id, device=device - ) - output_list = output_token_list.cpu().numpy() - generated_text = [ - self.text_tokenizer.decode(output[: int(length)], skip_special_tokens=True) - for output, length in zip(output_list, seq_lengths) - ] - return generated_text - - def _split(self, x, height, width): - r""" - Splits a flattened embedding x of shape (B, C * H * W + clip_img_dim) into two tensors of shape (B, C, H, W) - and (B, 1, clip_img_dim) - """ - batch_size = x.shape[0] - latent_height = height // self.vae_scale_factor - latent_width = width // self.vae_scale_factor - img_vae_dim = self.num_channels_latents * latent_height * latent_width - - img_vae, img_clip = x.split([img_vae_dim, self.image_encoder_projection_dim], dim=1) - - img_vae = torch.reshape(img_vae, (batch_size, self.num_channels_latents, latent_height, latent_width)) - img_clip = torch.reshape(img_clip, (batch_size, 1, self.image_encoder_projection_dim)) - return img_vae, img_clip - - def _combine(self, img_vae, img_clip): - r""" - Combines a latent image img_vae of shape (B, C, H, W) and a CLIP-embedded image img_clip of shape (B, 1, - clip_img_dim) into a single tensor of shape (B, C * H * W + clip_img_dim). - """ - img_vae = torch.reshape(img_vae, (img_vae.shape[0], -1)) - img_clip = torch.reshape(img_clip, (img_clip.shape[0], -1)) - return torch.concat([img_vae, img_clip], dim=-1) - - def _split_joint(self, x, height, width): - r""" - Splits a flattened embedding x of shape (B, C * H * W + clip_img_dim + text_seq_len * text_dim] into (img_vae, - img_clip, text) where img_vae is of shape (B, C, H, W), img_clip is of shape (B, 1, clip_img_dim), and text is - of shape (B, text_seq_len, text_dim). - """ - batch_size = x.shape[0] - latent_height = height // self.vae_scale_factor - latent_width = width // self.vae_scale_factor - img_vae_dim = self.num_channels_latents * latent_height * latent_width - text_dim = self.text_encoder_seq_len * self.text_intermediate_dim - - img_vae, img_clip, text = x.split([img_vae_dim, self.image_encoder_projection_dim, text_dim], dim=1) - - img_vae = torch.reshape(img_vae, (batch_size, self.num_channels_latents, latent_height, latent_width)) - img_clip = torch.reshape(img_clip, (batch_size, 1, self.image_encoder_projection_dim)) - text = torch.reshape(text, (batch_size, self.text_encoder_seq_len, self.text_intermediate_dim)) - return img_vae, img_clip, text - - def _combine_joint(self, img_vae, img_clip, text): - r""" - Combines a latent image img_vae of shape (B, C, H, W), a CLIP-embedded image img_clip of shape (B, L_img, - clip_img_dim), and a text embedding text of shape (B, L_text, text_dim) into a single embedding x of shape (B, - C * H * W + L_img * clip_img_dim + L_text * text_dim). - """ - img_vae = torch.reshape(img_vae, (img_vae.shape[0], -1)) - img_clip = torch.reshape(img_clip, (img_clip.shape[0], -1)) - text = torch.reshape(text, (text.shape[0], -1)) - return torch.concat([img_vae, img_clip, text], dim=-1) - - def _get_noise_pred( - self, - mode, - latents, - t, - prompt_embeds, - img_vae, - img_clip, - max_timestep, - data_type, - guidance_scale, - generator, - device, - height, - width, - ): - r""" - Gets the noise prediction using the `unet` and performs classifier-free guidance, if necessary. - """ - if mode == "joint": - # Joint text-image generation - img_vae_latents, img_clip_latents, text_latents = self._split_joint(latents, height, width) - - img_vae_out, img_clip_out, text_out = self.unet( - img_vae_latents, img_clip_latents, text_latents, timestep_img=t, timestep_text=t, data_type=data_type - ) - - x_out = self._combine_joint(img_vae_out, img_clip_out, text_out) - - if guidance_scale <= 1.0: - return x_out - - # Classifier-free guidance - img_vae_T = randn_tensor(img_vae.shape, generator=generator, device=device, dtype=img_vae.dtype) - img_clip_T = randn_tensor(img_clip.shape, generator=generator, device=device, dtype=img_clip.dtype) - text_T = randn_tensor(prompt_embeds.shape, generator=generator, device=device, dtype=prompt_embeds.dtype) - - _, _, text_out_uncond = self.unet( - img_vae_T, img_clip_T, text_latents, timestep_img=max_timestep, timestep_text=t, data_type=data_type - ) - - img_vae_out_uncond, img_clip_out_uncond, _ = self.unet( - img_vae_latents, - img_clip_latents, - text_T, - timestep_img=t, - timestep_text=max_timestep, - data_type=data_type, - ) - - x_out_uncond = self._combine_joint(img_vae_out_uncond, img_clip_out_uncond, text_out_uncond) - - return guidance_scale * x_out + (1.0 - guidance_scale) * x_out_uncond - elif mode == "text2img": - # Text-conditioned image generation - img_vae_latents, img_clip_latents = self._split(latents, height, width) - - img_vae_out, img_clip_out, text_out = self.unet( - img_vae_latents, img_clip_latents, prompt_embeds, timestep_img=t, timestep_text=0, data_type=data_type - ) - - img_out = self._combine(img_vae_out, img_clip_out) - - if guidance_scale <= 1.0: - return img_out - - # Classifier-free guidance - text_T = randn_tensor(prompt_embeds.shape, generator=generator, device=device, dtype=prompt_embeds.dtype) - - img_vae_out_uncond, img_clip_out_uncond, text_out_uncond = self.unet( - img_vae_latents, - img_clip_latents, - text_T, - timestep_img=t, - timestep_text=max_timestep, - data_type=data_type, - ) - - img_out_uncond = self._combine(img_vae_out_uncond, img_clip_out_uncond) - - return guidance_scale * img_out + (1.0 - guidance_scale) * img_out_uncond - elif mode == "img2text": - # Image-conditioned text generation - img_vae_out, img_clip_out, text_out = self.unet( - img_vae, img_clip, latents, timestep_img=0, timestep_text=t, data_type=data_type - ) - - if guidance_scale <= 1.0: - return text_out - - # Classifier-free guidance - img_vae_T = randn_tensor(img_vae.shape, generator=generator, device=device, dtype=img_vae.dtype) - img_clip_T = randn_tensor(img_clip.shape, generator=generator, device=device, dtype=img_clip.dtype) - - img_vae_out_uncond, img_clip_out_uncond, text_out_uncond = self.unet( - img_vae_T, img_clip_T, latents, timestep_img=max_timestep, timestep_text=t, data_type=data_type - ) - - return guidance_scale * text_out + (1.0 - guidance_scale) * text_out_uncond - elif mode == "text": - # Unconditional ("marginal") text generation (no CFG) - img_vae_out, img_clip_out, text_out = self.unet( - img_vae, img_clip, latents, timestep_img=max_timestep, timestep_text=t, data_type=data_type - ) - - return text_out - elif mode == "img": - # Unconditional ("marginal") image generation (no CFG) - img_vae_latents, img_clip_latents = self._split(latents, height, width) - - img_vae_out, img_clip_out, text_out = self.unet( - img_vae_latents, - img_clip_latents, - prompt_embeds, - timestep_img=t, - timestep_text=max_timestep, - data_type=data_type, - ) - - img_out = self._combine(img_vae_out, img_clip_out) - return img_out - - def check_latents_shape(self, latents_name, latents, expected_shape): - latents_shape = latents.shape - expected_num_dims = len(expected_shape) + 1 # expected dimensions plus the batch dimension - expected_shape_str = ", ".join(str(dim) for dim in expected_shape) - if len(latents_shape) != expected_num_dims: - raise ValueError( - f"`{latents_name}` should have shape (batch_size, {expected_shape_str}), but the current shape" - f" {latents_shape} has {len(latents_shape)} dimensions." - ) - for i in range(1, expected_num_dims): - if latents_shape[i] != expected_shape[i - 1]: - raise ValueError( - f"`{latents_name}` should have shape (batch_size, {expected_shape_str}), but the current shape" - f" {latents_shape} has {latents_shape[i]} != {expected_shape[i - 1]} at dimension {i}." - ) - - def check_inputs( - self, - mode, - prompt, - image, - height, - width, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - latents=None, - prompt_latents=None, - vae_latents=None, - clip_latents=None, - ): - # Check inputs before running the generative process. - if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0: - raise ValueError( - f"`height` and `width` have to be divisible by {self.vae_scale_factor} but are {height} and {width}." - ) - - if (callback_steps is None) or ( - 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 mode == "text2img": - 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 mode == "img2text": - if image is None: - raise ValueError("`img2text` mode requires an image to be provided.") - - # Check provided latents - latent_height = height // self.vae_scale_factor - latent_width = width // self.vae_scale_factor - full_latents_available = latents is not None - prompt_latents_available = prompt_latents is not None - vae_latents_available = vae_latents is not None - clip_latents_available = clip_latents is not None - - if full_latents_available: - individual_latents_available = ( - prompt_latents is not None or vae_latents is not None or clip_latents is not None - ) - if individual_latents_available: - logger.warning( - "You have supplied both `latents` and at least one of `prompt_latents`, `vae_latents`, and" - " `clip_latents`. The value of `latents` will override the value of any individually supplied latents." - ) - # Check shape of full latents - img_vae_dim = self.num_channels_latents * latent_height * latent_width - text_dim = self.text_encoder_seq_len * self.text_encoder_hidden_size - latents_dim = img_vae_dim + self.image_encoder_projection_dim + text_dim - latents_expected_shape = (latents_dim,) - self.check_latents_shape("latents", latents, latents_expected_shape) - - # Check individual latent shapes, if present - if prompt_latents_available: - prompt_latents_expected_shape = (self.text_encoder_seq_len, self.text_encoder_hidden_size) - self.check_latents_shape("prompt_latents", prompt_latents, prompt_latents_expected_shape) - - if vae_latents_available: - vae_latents_expected_shape = (self.num_channels_latents, latent_height, latent_width) - self.check_latents_shape("vae_latents", vae_latents, vae_latents_expected_shape) - - if clip_latents_available: - clip_latents_expected_shape = (1, self.image_encoder_projection_dim) - self.check_latents_shape("clip_latents", clip_latents, clip_latents_expected_shape) - - if mode in ["text2img", "img"] and vae_latents_available and clip_latents_available: - if vae_latents.shape[0] != clip_latents.shape[0]: - raise ValueError( - f"Both `vae_latents` and `clip_latents` are supplied, but their batch dimensions are not equal:" - f" {vae_latents.shape[0]} != {clip_latents.shape[0]}." - ) - - if mode == "joint" and prompt_latents_available and vae_latents_available and clip_latents_available: - if prompt_latents.shape[0] != vae_latents.shape[0] or prompt_latents.shape[0] != clip_latents.shape[0]: - raise ValueError( - f"All of `prompt_latents`, `vae_latents`, and `clip_latents` are supplied, but their batch" - f" dimensions are not equal: {prompt_latents.shape[0]} != {vae_latents.shape[0]}" - f" != {clip_latents.shape[0]}." - ) - - @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, - 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, - 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", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - ): - r""" - The call function to the pipeline for generation. - - Args: - 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*): - `Image` or tensor representing an image batch. Required for image-conditioned text generation - (`img2text`) mode. - 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. - data_type (`int`, *optional*, defaults to 1): - The data type (either 0 or 1). Only used if you are loading a checkpoint which supports a data type - embedding; this is added for compatibility with the - [UniDiffuser-v1](https://huggingface.co/thu-ml/unidiffuser-v1) checkpoint. - 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 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*): - 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. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. Used in `text2img` (text-conditioned image generation) and - `img` mode. If the mode is joint and both `num_images_per_prompt` and `num_prompts_per_image` are - supplied, `min(num_images_per_prompt, num_prompts_per_image)` samples are generated. - num_prompts_per_image (`int`, *optional*, defaults to 1): - The number of prompts to generate per image. Used in `img2text` (image-conditioned text generation) and - `text` mode. If the mode is joint and both `num_images_per_prompt` and `num_prompts_per_image` are - supplied, `min(num_images_per_prompt, num_prompts_per_image)` samples are generated. - 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*): - 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 joint - image-text 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`. This assumes - a full set of VAE, CLIP, and text latents, if supplied, overrides the value of `prompt_latents`, - `vae_latents`, and `clip_latents`. - prompt_latents (`torch.Tensor`, *optional*): - Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for text - 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`. - vae_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`. - clip_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. Used in text-conditioned - image generation (`text2img`) mode. - 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 be generated from the `negative_prompt` input argument. Used - in text-conditioned image generation (`text2img`) mode. - 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 [`~pipelines.ImageTextPipelineOutput`] instead of a plain tuple. - callback (`Callable`, *optional*): - A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at - every step. - - Returns: - [`~pipelines.unidiffuser.ImageTextPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`~pipelines.unidiffuser.ImageTextPipelineOutput`] 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 generated texts. - """ - - # 0. Default height and width to unet - height = height or self.unet_resolution * self.vae_scale_factor - width = width or self.unet_resolution * self.vae_scale_factor - - # 1. Check inputs - # Recalculate mode for each call to the pipeline. - mode = self._infer_mode(prompt, prompt_embeds, image, latents, prompt_latents, vae_latents, clip_latents) - self.check_inputs( - mode, - prompt, - image, - height, - width, - callback_steps, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - latents, - prompt_latents, - vae_latents, - clip_latents, - ) - - # 2. Define call parameters - batch_size, multiplier = self._infer_batch_size( - mode, - prompt, - prompt_embeds, - image, - num_images_per_prompt, - num_prompts_per_image, - latents, - prompt_latents, - vae_latents, - clip_latents, - ) - device = self._execution_device - reduce_text_emb_dim = self.text_intermediate_dim < self.text_encoder_hidden_size or self.mode != "text2img" - - # 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. - # Note that this differs from the formulation in the unidiffusers paper! - do_classifier_free_guidance = guidance_scale > 1.0 - - # check if scheduler is in sigmas space - # scheduler_is_in_sigma_space = hasattr(self.scheduler, "sigmas") - - # 3. Encode input prompt, if available; otherwise prepare text latents - if latents is not None: - # Overwrite individual latents - vae_latents, clip_latents, prompt_latents = self._split_joint(latents, height, width) - - if mode in ["text2img"]: - # 3.1. Encode input prompt, if available - assert prompt is not None or prompt_embeds is not None - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - device=device, - num_images_per_prompt=multiplier, - do_classifier_free_guidance=do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - - # if do_classifier_free_guidance: - # prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - else: - # 3.2. Prepare text latent variables, if input not available - prompt_embeds = self.prepare_text_latents( - batch_size=batch_size, - num_images_per_prompt=multiplier, - seq_len=self.text_encoder_seq_len, - hidden_size=self.text_encoder_hidden_size, - dtype=self.text_encoder.dtype, # Should work with both full precision and mixed precision - device=device, - generator=generator, - latents=prompt_latents, - ) - - if reduce_text_emb_dim: - prompt_embeds = self.text_decoder.encode(prompt_embeds) - - # 4. Encode image, if available; otherwise prepare image latents - if mode in ["img2text"]: - # 4.1. Encode images, if available - assert image is not None, "`img2text` requires a conditioning image" - # Encode image using VAE - image_vae = self.image_processor.preprocess(image) - height, width = image_vae.shape[-2:] - image_vae_latents = self.encode_image_vae_latents( - image=image_vae, - batch_size=batch_size, - num_prompts_per_image=multiplier, - dtype=prompt_embeds.dtype, - device=device, - do_classifier_free_guidance=False, # Copied from InstructPix2Pix, don't use their version of CFG - generator=generator, - ) - - # Encode image using CLIP - image_clip_latents = self.encode_image_clip_latents( - image=image, - batch_size=batch_size, - num_prompts_per_image=multiplier, - dtype=prompt_embeds.dtype, - device=device, - generator=generator, - ) - # (batch_size, clip_hidden_size) => (batch_size, 1, clip_hidden_size) - image_clip_latents = image_clip_latents.unsqueeze(1) - else: - # 4.2. Prepare image latent variables, if input not available - # Prepare image VAE latents in latent space - image_vae_latents = self.prepare_image_vae_latents( - batch_size=batch_size, - num_prompts_per_image=multiplier, - num_channels_latents=self.num_channels_latents, - height=height, - width=width, - dtype=prompt_embeds.dtype, - device=device, - generator=generator, - latents=vae_latents, - ) - - # Prepare image CLIP latents - image_clip_latents = self.prepare_image_clip_latents( - batch_size=batch_size, - num_prompts_per_image=multiplier, - clip_img_dim=self.image_encoder_projection_dim, - dtype=prompt_embeds.dtype, - device=device, - generator=generator, - latents=clip_latents, - ) - - # 5. Set timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - # max_timestep = timesteps[0] - max_timestep = self.scheduler.config.num_train_timesteps - - # 6. Prepare latent variables - if mode == "joint": - latents = self._combine_joint(image_vae_latents, image_clip_latents, prompt_embeds) - elif mode in ["text2img", "img"]: - latents = self._combine(image_vae_latents, image_clip_latents) - elif mode in ["img2text", "text"]: - latents = prompt_embeds - - # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - logger.debug(f"Scheduler extra step kwargs: {extra_step_kwargs}") - - # 8. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # predict the noise residual - # Also applies classifier-free guidance as described in the UniDiffuser paper - noise_pred = self._get_noise_pred( - mode, - latents, - t, - prompt_embeds, - image_vae_latents, - image_clip_latents, - max_timestep, - data_type, - guidance_scale, - generator, - device, - height, - width, - ) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # 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 callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # 9. Post-processing - image = None - text = None - if mode == "joint": - image_vae_latents, image_clip_latents, text_latents = self._split_joint(latents, height, width) - - if not output_type == "latent": - # Map latent VAE image back to pixel space - image = self.vae.decode(image_vae_latents / self.vae.config.scaling_factor, return_dict=False)[0] - else: - image = image_vae_latents - - text = self.decode_text_latents(text_latents, device) - elif mode in ["text2img", "img"]: - image_vae_latents, image_clip_latents = self._split(latents, height, width) - - if not output_type == "latent": - # Map latent VAE image back to pixel space - image = self.vae.decode(image_vae_latents / self.vae.config.scaling_factor, return_dict=False)[0] - else: - image = image_vae_latents - elif mode in ["img2text", "text"]: - text_latents = latents - text = self.decode_text_latents(text_latents, device) - - self.maybe_free_model_hooks() - - # 10. Postprocess the image, if necessary - if image is not None: - do_denormalize = [True] * image.shape[0] - image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) - - # Offload last model to CPU - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.final_offload_hook.offload() - - if not return_dict: - return (image, text) - - return ImageTextPipelineOutput(images=image, text=text) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__init__.py deleted file mode 100644 index ab765a1bbad9ae0b8c7f2e7dec5019e3a5e4ea01..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__init__.py +++ /dev/null @@ -1,52 +0,0 @@ -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_visualcloze_combined"] = ["VisualClozePipeline"] - _import_structure["pipeline_visualcloze_generation"] = ["VisualClozeGenerationPipeline"] - - -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_visualcloze_combined import VisualClozePipeline - from .pipeline_visualcloze_generation import VisualClozeGenerationPipeline - - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 0dd0bcec854c6e2bcaf8d57a2950281f20de73c9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__pycache__/pipeline_visualcloze_combined.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__pycache__/pipeline_visualcloze_combined.cpython-310.pyc deleted file mode 100644 index bca8fd0038dd3f1384f893f727787270d3f13eea..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__pycache__/pipeline_visualcloze_combined.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__pycache__/pipeline_visualcloze_generation.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__pycache__/pipeline_visualcloze_generation.cpython-310.pyc deleted file mode 100644 index d5f26e480921fe305dc18e916711cd02fd04d0c0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__pycache__/pipeline_visualcloze_generation.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__pycache__/visualcloze_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__pycache__/visualcloze_utils.cpython-310.pyc deleted file mode 100644 index 3ed6d8b20db6ef16481655b614f531ba5c34f52f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/__pycache__/visualcloze_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py deleted file mode 100644 index 68130baad7093b892ad8895bb6b65311c59a3cde..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py +++ /dev/null @@ -1,444 +0,0 @@ -# Copyright 2025 VisualCloze 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 Any, Callable, Dict, List, Optional, Union - -import torch -from PIL import Image -from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast - -from ...loaders import FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin -from ...models.autoencoders import AutoencoderKL -from ...models.transformers import FluxTransformer2DModel -from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import ( - is_torch_xla_available, - logging, - replace_example_docstring, -) -from ..flux.pipeline_flux_fill import FluxFillPipeline as VisualClozeUpsamplingPipeline -from ..flux.pipeline_output import FluxPipelineOutput -from ..pipeline_utils import DiffusionPipeline -from .pipeline_visualcloze_generation import VisualClozeGenerationPipeline - - -if is_torch_xla_available(): - 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 VisualClozePipeline - >>> from diffusers.utils import load_image - - >>> image_paths = [ - ... # in-context examples - ... [ - ... load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/visualcloze/visualcloze_mask2image_incontext-example-1_mask.jpg" - ... ), - ... load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/visualcloze/visualcloze_mask2image_incontext-example-1_image.jpg" - ... ), - ... ], - ... # query with the target image - ... [ - ... load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/visualcloze/visualcloze_mask2image_query_mask.jpg" - ... ), - ... None, # No image needed for the target image - ... ], - ... ] - >>> task_prompt = "In each row, a logical task is demonstrated to achieve [IMAGE2] an aesthetically pleasing photograph based on [IMAGE1] sam 2-generated masks with rich color coding." - >>> content_prompt = "Majestic photo of a golden eagle perched on a rocky outcrop in a mountainous landscape. The eagle is positioned in the right foreground, facing left, with its sharp beak and keen eyes prominently visible. Its plumage is a mix of dark brown and golden hues, with intricate feather details. The background features a soft-focus view of snow-capped mountains under a cloudy sky, creating a serene and grandiose atmosphere. The foreground includes rugged rocks and patches of green moss. Photorealistic, medium depth of field, soft natural lighting, cool color palette, high contrast, sharp focus on the eagle, blurred background, tranquil, majestic, wildlife photography." - >>> pipe = VisualClozePipeline.from_pretrained( - ... "VisualCloze/VisualClozePipeline-384", resolution=384, torch_dtype=torch.bfloat16 - ... ) - >>> pipe.to("cuda") - - >>> image = pipe( - ... task_prompt=task_prompt, - ... content_prompt=content_prompt, - ... image=image_paths, - ... upsampling_width=1344, - ... upsampling_height=768, - ... upsampling_strength=0.4, - ... guidance_scale=30, - ... num_inference_steps=30, - ... max_sequence_length=512, - ... generator=torch.Generator("cpu").manual_seed(0), - ... ).images[0][0] - >>> image.save("visualcloze.png") - ``` -""" - - -class VisualClozePipeline( - DiffusionPipeline, - FluxLoraLoaderMixin, - FromSingleFileMixin, - TextualInversionLoaderMixin, -): - r""" - The VisualCloze pipeline for image generation with visual context. Reference: - https://github.com/lzyhha/VisualCloze/tree/main. This pipeline is designed to generate images based on visual - in-context examples. - - Args: - transformer ([`FluxTransformer2DModel`]): - 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 ([`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). - resolution (`int`, *optional*, defaults to 384): - The resolution of each image when concatenating images from the query and in-context examples. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" - _optional_components = [] - _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__( - self, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - text_encoder_2: T5EncoderModel, - tokenizer_2: T5TokenizerFast, - transformer: FluxTransformer2DModel, - resolution: int = 384, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - transformer=transformer, - scheduler=scheduler, - ) - - self.generation_pipe = VisualClozeGenerationPipeline( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - transformer=transformer, - scheduler=scheduler, - resolution=resolution, - ) - self.upsampling_pipe = VisualClozeUpsamplingPipeline( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - transformer=transformer, - scheduler=scheduler, - ) - - def check_inputs( - self, - image, - task_prompt, - content_prompt, - upsampling_height, - upsampling_width, - strength, - prompt_embeds=None, - pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - max_sequence_length=None, - ): - if strength < 0 or strength > 1: - raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") - - if upsampling_height is not None and upsampling_height % (self.vae_scale_factor * 2) != 0: - logger.warning( - f"`upsampling_height`has to be divisible by {self.vae_scale_factor * 2} but are {upsampling_height}. Dimensions will be resized accordingly" - ) - if upsampling_width is not None and upsampling_width % (self.vae_scale_factor * 2) != 0: - logger.warning( - f"`upsampling_width` have to be divisible by {self.vae_scale_factor * 2} but are {upsampling_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]}" - ) - - # Validate prompt inputs - if (task_prompt is not None or content_prompt is not None) and prompt_embeds is not None: - raise ValueError("Cannot provide both text `task_prompt` + `content_prompt` and `prompt_embeds`. ") - - if task_prompt is None and content_prompt is None and prompt_embeds is None: - raise ValueError("Must provide either `task_prompt` + `content_prompt` or pre-computed `prompt_embeds`. ") - - # Validate prompt types and consistency - if task_prompt is None: - raise ValueError("`task_prompt` is missing.") - - if task_prompt is not None and not isinstance(task_prompt, (str, list)): - raise ValueError(f"`task_prompt` must be str or list, got {type(task_prompt)}") - - if content_prompt is not None and not isinstance(content_prompt, (str, list)): - raise ValueError(f"`content_prompt` must be str or list, got {type(content_prompt)}") - - if isinstance(task_prompt, list) or isinstance(content_prompt, list): - if not isinstance(task_prompt, list) or not isinstance(content_prompt, list): - raise ValueError( - f"`task_prompt` and `content_prompt` must both be lists, or both be of type str or None, " - f"got {type(task_prompt)} and {type(content_prompt)}" - ) - if len(content_prompt) != len(task_prompt): - raise ValueError("`task_prompt` and `content_prompt` must have the same length whe they are lists.") - - for sample in image: - if not isinstance(sample, list) or not isinstance(sample[0], list): - raise ValueError("Each sample in the batch must have a 2D list of images.") - if len({len(row) for row in sample}) != 1: - raise ValueError("Each in-context example and query should contain the same number of images.") - if not any(img is None for img in sample[-1]): - raise ValueError("There are no targets in the query, which should be represented as None.") - for row in sample[:-1]: - if any(img is None for img in row): - raise ValueError("Images are missing in in-context examples.") - - # Validate embeddings - 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`." - ) - - # Validate sequence length - if max_sequence_length is not None and max_sequence_length > 512: - raise ValueError(f"max_sequence_length cannot exceed 512, got {max_sequence_length}") - - @torch.no_grad() - @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, - num_inference_steps: int = 50, - sigmas: Optional[List[float]] = 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", - 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, - upsampling_strength: float = 1.0, - ): - r""" - Function invoked when calling the VisualCloze pipeline for generation. - - Args: - task_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts to define the task intention. - 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`, 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)`. - upsampling_height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The height in pixels of the generated image (i.e., output image) after upsampling via SDEdit. By - default, the image is upsampled by a factor of three, and the base resolution is determined by the - resolution parameter of the pipeline. When only one of `upsampling_height` or `upsampling_width` is - specified, the other will be automatically set based on the aspect ratio. - upsampling_width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The width in pixels of the generated image (i.e., output image) after upsampling via SDEdit. By - default, the image is upsampled by a factor of three, and the base resolution is determined by the - resolution parameter of the pipeline. When only one of `upsampling_height` or `upsampling_width` is - specified, the other will be automatically set based on the aspect ratio. - 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 30.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_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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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.FluxPipelineOutput`] 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`. - upsampling_strength (`float`, *optional*, defaults to 1.0): - Indicates extent to transform the reference `image` when upsampling the results. Must be between 0 and - 1. The generated image is used as a starting point and more noise is added the higher the - `upsampling_strength`. The number of denoising steps depends on the amount of noise initially added. - When `upsampling_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 0 skips the upsampling step and - output the results at the resolution of `self.resolution`. - - Examples: - - Returns: - [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict` - is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated - images. - """ - - generation_output = self.generation_pipe( - task_prompt=task_prompt, - content_prompt=content_prompt, - image=image, - num_inference_steps=num_inference_steps, - sigmas=sigmas, - guidance_scale=guidance_scale, - num_images_per_prompt=num_images_per_prompt, - generator=generator, - latents=latents, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - joint_attention_kwargs=joint_attention_kwargs, - callback_on_step_end=callback_on_step_end, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - output_type=output_type if upsampling_strength == 0 else "pil", - ) - if upsampling_strength == 0: - if not return_dict: - return (generation_output,) - - return FluxPipelineOutput(images=generation_output) - - # Upsampling the generated images - # 1. Prepare the input images and prompts - if not isinstance(content_prompt, (list)): - content_prompt = [content_prompt] - n_target_per_sample = [] - upsampling_image = [] - upsampling_mask = [] - upsampling_prompt = [] - upsampling_generator = generator if isinstance(generator, (torch.Generator,)) else [] - for i in range(len(generation_output.images)): - n_target_per_sample.append(len(generation_output.images[i])) - for image in generation_output.images[i]: - upsampling_image.append(image) - upsampling_mask.append(Image.new("RGB", image.size, (255, 255, 255))) - upsampling_prompt.append( - content_prompt[i % len(content_prompt)] if content_prompt[i % len(content_prompt)] else "" - ) - if not isinstance(generator, (torch.Generator,)): - upsampling_generator.append(generator[i % len(content_prompt)]) - - # 2. Apply the denosing loop - upsampling_output = self.upsampling_pipe( - prompt=upsampling_prompt, - image=upsampling_image, - mask_image=upsampling_mask, - height=upsampling_height, - width=upsampling_width, - strength=upsampling_strength, - num_inference_steps=num_inference_steps, - sigmas=sigmas, - guidance_scale=guidance_scale, - generator=upsampling_generator, - output_type=output_type, - joint_attention_kwargs=joint_attention_kwargs, - callback_on_step_end=callback_on_step_end, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - ) - image = upsampling_output.images - - output = [] - if output_type == "pil": - # Each sample in the batch may have multiple output images. When returning as PIL images, - # these images cannot be concatenated. Therefore, for each sample, - # a list is used to represent all the output images. - output = [] - start = 0 - for n in n_target_per_sample: - output.append(image[start : start + n]) - start += n - else: - output = image - - if not return_dict: - return (output,) - - return FluxPipelineOutput(images=output) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py deleted file mode 100644 index e7a1d4a4b24873c9f6a69c437c3b9ef272fc6b19..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py +++ /dev/null @@ -1,952 +0,0 @@ -# Copyright 2025 VisualCloze 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 Any, Callable, Dict, List, Optional, Union - -import numpy as np -import torch -from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast - -from ...loaders import FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin -from ...models.autoencoders import AutoencoderKL -from ...models.transformers import FluxTransformer2DModel -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 ..flux.pipeline_flux_fill import calculate_shift, retrieve_latents, retrieve_timesteps -from ..flux.pipeline_output import FluxPipelineOutput -from ..pipeline_utils import DiffusionPipeline -from .visualcloze_utils import VisualClozeProcessor - - -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 VisualClozeGenerationPipeline, FluxFillPipeline as VisualClozeUpsamplingPipeline - >>> from diffusers.utils import load_image - >>> from PIL import Image - - >>> image_paths = [ - ... # in-context examples - ... [ - ... load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/visualcloze/visualcloze_mask2image_incontext-example-1_mask.jpg" - ... ), - ... load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/visualcloze/visualcloze_mask2image_incontext-example-1_image.jpg" - ... ), - ... ], - ... # query with the target image - ... [ - ... load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/visualcloze/visualcloze_mask2image_query_mask.jpg" - ... ), - ... None, # No image needed for the target image - ... ], - ... ] - >>> task_prompt = "In each row, a logical task is demonstrated to achieve [IMAGE2] an aesthetically pleasing photograph based on [IMAGE1] sam 2-generated masks with rich color coding." - >>> content_prompt = "Majestic photo of a golden eagle perched on a rocky outcrop in a mountainous landscape. The eagle is positioned in the right foreground, facing left, with its sharp beak and keen eyes prominently visible. Its plumage is a mix of dark brown and golden hues, with intricate feather details. The background features a soft-focus view of snow-capped mountains under a cloudy sky, creating a serene and grandiose atmosphere. The foreground includes rugged rocks and patches of green moss. Photorealistic, medium depth of field, soft natural lighting, cool color palette, high contrast, sharp focus on the eagle, blurred background, tranquil, majestic, wildlife photography." - >>> pipe = VisualClozeGenerationPipeline.from_pretrained( - ... "VisualCloze/VisualClozePipeline-384", resolution=384, torch_dtype=torch.bfloat16 - ... ) - >>> pipe.to("cuda") - - >>> image = pipe( - ... task_prompt=task_prompt, - ... content_prompt=content_prompt, - ... image=image_paths, - ... guidance_scale=30, - ... num_inference_steps=30, - ... max_sequence_length=512, - ... generator=torch.Generator("cpu").manual_seed(0), - ... ).images[0][0] - - >>> # optional, upsampling the generated image - >>> pipe_upsample = VisualClozeUpsamplingPipeline.from_pipe(pipe) - >>> pipe_upsample.to("cuda") - - >>> mask_image = Image.new("RGB", image.size, (255, 255, 255)) - - >>> image = pipe_upsample( - ... image=image, - ... mask_image=mask_image, - ... prompt=content_prompt, - ... width=1344, - ... height=768, - ... strength=0.4, - ... guidance_scale=30, - ... num_inference_steps=30, - ... max_sequence_length=512, - ... generator=torch.Generator("cpu").manual_seed(0), - ... ).images[0] - - >>> image.save("visualcloze.png") - ``` -""" - - -class VisualClozeGenerationPipeline( - DiffusionPipeline, - FluxLoraLoaderMixin, - FromSingleFileMixin, - TextualInversionLoaderMixin, -): - r""" - The VisualCloze pipeline for image generation with visual context. Reference: - https://github.com/lzyhha/VisualCloze/tree/main This pipeline is designed to generate images based on visual - in-context examples. - - Args: - transformer ([`FluxTransformer2DModel`]): - 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 ([`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). - resolution (`int`, *optional*, defaults to 384): - The resolution of each image when concatenating images from the query and in-context examples. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" - _optional_components = [] - _callback_tensor_inputs = ["latents", "prompt_embeds"] - - def __init__( - self, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: CLIPTokenizer, - text_encoder_2: T5EncoderModel, - tokenizer_2: T5TokenizerFast, - transformer: FluxTransformer2DModel, - resolution: int = 384, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - transformer=transformer, - scheduler=scheduler, - ) - self.resolution = resolution - 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.latent_channels = self.vae.config.latent_channels if getattr(self, "vae", None) else 16 - self.image_processor = VisualClozeProcessor( - vae_scale_factor=self.vae_scale_factor * 2, vae_latent_channels=self.latent_channels, resolution=resolution - ) - self.tokenizer_max_length = ( - self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77 - ) - self.default_sample_size = 128 - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds - 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_2) - - text_inputs = self.tokenizer_2( - 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 - untruncated_ids = self.tokenizer_2(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_2.batch_decode(untruncated_ids[:, self.tokenizer_max_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_2(text_input_ids.to(device), output_hidden_states=False)[0] - - dtype = self.text_encoder_2.dtype - prompt_embeds = prompt_embeds.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) - - return 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]], - num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - ): - device = device or self._execution_device - - 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=self.tokenizer_max_length, - truncation=True, - return_overflowing_tokens=False, - return_length=False, - 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_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_max_length} tokens: {removed_text}" - ) - prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds.pooler_output - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt) - prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - return 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, - num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - max_sequence_length: int = 512, - lora_scale: Optional[float] = None, - ): - r""" - - Args: - 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*): - The prompt or prompts to define the task intention. - 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 - num_images_per_prompt (`int`): - number of images that should be generated per prompt - 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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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) - if self.text_encoder_2 is not None and USE_PEFT_BACKEND: - scale_lora_layers(self.text_encoder_2, lora_scale) - - if isinstance(layout_prompt, str): - layout_prompt = [layout_prompt] - task_prompt = [task_prompt] - content_prompt = [content_prompt] - - def _preprocess(prompt, content=False): - if prompt is not None: - return f"The last image of the last row depicts: {prompt}" if content else prompt - else: - return "" - - prompt = [ - f"{_preprocess(layout_prompt[i])} {_preprocess(task_prompt[i])} {_preprocess(content_prompt[i], content=True)}".strip() - for i in range(len(layout_prompt)) - ] - pooled_prompt_embeds = self._get_clip_prompt_embeds( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - ) - prompt_embeds = self._get_t5_prompt_embeds( - prompt=prompt, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - ) - - 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) - - if self.text_encoder_2 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_2, lora_scale) - - 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) - - return prompt_embeds, pooled_prompt_embeds, text_ids - - # 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, - image, - task_prompt, - content_prompt, - prompt_embeds=None, - pooled_prompt_embeds=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]}" - ) - - # Validate prompt inputs - if (task_prompt is not None or content_prompt is not None) and prompt_embeds is not None: - raise ValueError("Cannot provide both text `task_prompt` + `content_prompt` and `prompt_embeds`. ") - - if task_prompt is None and content_prompt is None and prompt_embeds is None: - raise ValueError("Must provide either `task_prompt` + `content_prompt` or pre-computed `prompt_embeds`. ") - - # Validate prompt types and consistency - if task_prompt is None: - raise ValueError("`task_prompt` is missing.") - - if task_prompt is not None and not isinstance(task_prompt, (str, list)): - raise ValueError(f"`task_prompt` must be str or list, got {type(task_prompt)}") - - if content_prompt is not None and not isinstance(content_prompt, (str, list)): - raise ValueError(f"`content_prompt` must be str or list, got {type(content_prompt)}") - - if isinstance(task_prompt, list) or isinstance(content_prompt, list): - if not isinstance(task_prompt, list) or not isinstance(content_prompt, list): - raise ValueError( - f"`task_prompt` and `content_prompt` must both be lists, or both be of type str or None, " - f"got {type(task_prompt)} and {type(content_prompt)}" - ) - if len(content_prompt) != len(task_prompt): - raise ValueError("`task_prompt` and `content_prompt` must have the same length whe they are lists.") - - for sample in image: - if not isinstance(sample, list) or not isinstance(sample[0], list): - raise ValueError("Each sample in the batch must have a 2D list of images.") - if len({len(row) for row in sample}) != 1: - raise ValueError("Each in-context example and query should contain the same number of images.") - if not any(img is None for img in sample[-1]): - raise ValueError("There are no targets in the query, which should be represented as None.") - for row in sample[:-1]: - if any(img is None for img in row): - raise ValueError("Images are missing in in-context examples.") - - # Validate embeddings - 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`." - ) - - # Validate sequence length - if max_sequence_length is not None and max_sequence_length > 512: - raise ValueError(f"max_sequence_length cannot exceed 512, got {max_sequence_length}") - - @staticmethod - def _prepare_latent_image_ids(image, vae_scale_factor, device, dtype): - latent_image_ids = [] - - for idx, img in enumerate(image, start=1): - img = img.squeeze(0) - channels, height, width = img.shape - - num_patches_h = height // vae_scale_factor // 2 - num_patches_w = width // vae_scale_factor // 2 - - patch_ids = torch.zeros(num_patches_h, num_patches_w, 3, device=device, dtype=dtype) - patch_ids[..., 0] = idx - patch_ids[..., 1] = torch.arange(num_patches_h, device=device, dtype=dtype)[:, None] - patch_ids[..., 2] = torch.arange(num_patches_w, device=device, dtype=dtype)[None, :] - - patch_ids = patch_ids.reshape(-1, 3) - latent_image_ids.append(patch_ids) - - return torch.cat(latent_image_ids, dim=0) - - @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 - def _unpack_latents(latents, sizes, vae_scale_factor): - batch_size, num_patches, channels = latents.shape - - start = 0 - unpacked_latents = [] - for i in range(len(sizes)): - cur_size = sizes[i] - height = cur_size[0][0] // vae_scale_factor - width = sum([size[1] for size in cur_size]) // vae_scale_factor - - end = start + (height * width) // 4 - - cur_latents = latents[:, start:end] - cur_latents = cur_latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2) - cur_latents = cur_latents.permute(0, 3, 1, 4, 2, 5) - cur_latents = cur_latents.reshape(batch_size, channels // (2 * 2), height, width) - - unpacked_latents.append(cur_latents) - - start = end - - return unpacked_latents - - def enable_vae_slicing(self): - 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.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - 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. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - def _prepare_latents(self, image, mask, gen, vae_scale_factor, device, dtype): - """Helper function to prepare latents for a single batch.""" - # Concatenate images and masks along width dimension - image = [torch.cat(img, dim=3).to(device=device, dtype=dtype) for img in image] - mask = [torch.cat(m, dim=3).to(device=device, dtype=dtype) for m in mask] - - # Generate latent image IDs - latent_image_ids = self._prepare_latent_image_ids(image, vae_scale_factor, device, dtype) - - # For initial encoding, use actual images - image_latent = [self._encode_vae_image(img, gen) for img in image] - masked_image_latent = [img.clone() for img in image_latent] - - for i in range(len(image_latent)): - # Rearrange latents and masks for patch processing - num_channels_latents, height, width = image_latent[i].shape[1:] - image_latent[i] = self._pack_latents(image_latent[i], 1, num_channels_latents, height, width) - masked_image_latent[i] = self._pack_latents(masked_image_latent[i], 1, num_channels_latents, height, width) - - # Rearrange masks for patch processing - num_channels_latents, height, width = mask[i].shape[1:] - mask[i] = mask[i].view( - 1, - num_channels_latents, - height // vae_scale_factor, - vae_scale_factor, - width // vae_scale_factor, - vae_scale_factor, - ) - mask[i] = mask[i].permute(0, 1, 3, 5, 2, 4) - mask[i] = mask[i].reshape( - 1, - num_channels_latents * (vae_scale_factor**2), - height // vae_scale_factor, - width // vae_scale_factor, - ) - mask[i] = self._pack_latents( - mask[i], - 1, - num_channels_latents * (vae_scale_factor**2), - height // vae_scale_factor, - width // vae_scale_factor, - ) - - # Concatenate along batch dimension - image_latent = torch.cat(image_latent, dim=1) - masked_image_latent = torch.cat(masked_image_latent, dim=1) - mask = torch.cat(mask, dim=1) - - return image_latent, masked_image_latent, mask, latent_image_ids - - def prepare_latents( - self, - input_image, - input_mask, - timestep, - batch_size, - dtype, - device, - generator, - 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." - ) - - # Process each batch - masked_image_latents = [] - image_latents = [] - masks = [] - latent_image_ids = [] - - for i in range(len(input_image)): - _image_latent, _masked_image_latent, _mask, _latent_image_ids = self._prepare_latents( - input_image[i], - input_mask[i], - generator if isinstance(generator, torch.Generator) else generator[i], - vae_scale_factor, - device, - dtype, - ) - masked_image_latents.append(_masked_image_latent) - image_latents.append(_image_latent) - masks.append(_mask) - latent_image_ids.append(_latent_image_ids) - - # Concatenate all batches - masked_image_latents = torch.cat(masked_image_latents, dim=0) - image_latents = torch.cat(image_latents, dim=0) - masks = torch.cat(masks, dim=0) - - # Handle batch size expansion - if batch_size > masked_image_latents.shape[0]: - if batch_size % masked_image_latents.shape[0] == 0: - # Expand batches by repeating - additional_image_per_prompt = batch_size // masked_image_latents.shape[0] - masked_image_latents = torch.cat([masked_image_latents] * additional_image_per_prompt, dim=0) - image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0) - masks = torch.cat([masks] * additional_image_per_prompt, dim=0) - else: - raise ValueError( - f"Cannot expand batch size from {masked_image_latents.shape[0]} to {batch_size}. " - "Batch sizes must be multiples of each other." - ) - - # Add noise to latents - noises = randn_tensor(image_latents.shape, generator=generator, device=device, dtype=dtype) - latents = self.scheduler.scale_noise(image_latents, timestep, noises).to(dtype=dtype) - - # Combine masked latents with masks - masked_image_latents = torch.cat((masked_image_latents, masks), dim=-1).to(dtype=dtype) - - return latents, masked_image_latents, latent_image_ids[0] - - @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 - - @torch.no_grad() - @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, - num_inference_steps: int = 50, - sigmas: Optional[List[float]] = 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", - 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 VisualCloze pipeline for generation. - - Args: - task_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts to define the task intention. - 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`, 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)`. - 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 30.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_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. - pooled_prompt_embeds (`torch.FloatTensor`, *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. - 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.FluxPipelineOutput`] 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.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] 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( - image, - task_prompt, - content_prompt, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - 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._interrupt = False - - processor_output = self.image_processor.preprocess( - task_prompt, content_prompt, image, vae_scale_factor=self.vae_scale_factor - ) - - # 2. Define call parameters - if processor_output["task_prompt"] is not None and isinstance(processor_output["task_prompt"], str): - batch_size = 1 - elif processor_output["task_prompt"] is not None and isinstance(processor_output["task_prompt"], list): - batch_size = len(processor_output["task_prompt"]) - - device = self._execution_device - - # 3. Prepare prompt embeddings - lora_scale = ( - self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None - ) - prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt( - layout_prompt=processor_output["layout_prompt"], - task_prompt=processor_output["task_prompt"], - content_prompt=processor_output["content_prompt"], - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - device=device, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - - # 4. Prepare timesteps - # Calculate sequence length and shift factor - image_seq_len = sum( - (size[0] // self.vae_scale_factor // 2) * (size[1] // self.vae_scale_factor // 2) - for sample in processor_output["image_size"][0] - for size in sample - ) - - # Calculate noise schedule parameters - 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), - ) - - # Get timesteps - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas - 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, 1.0, device) - - # 5. Prepare latent variables - latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) - latents, masked_image_latents, latent_image_ids = self.prepare_latents( - processor_output["init_image"], - processor_output["mask"], - latent_timestep, - batch_size * num_images_per_prompt, - prompt_embeds.dtype, - device, - generator, - vae_scale_factor=self.vae_scale_factor, - ) - - # Calculate warmup steps - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - # Prepare 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 - - # 6. 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]).to(latents.dtype) - latent_model_input = torch.cat((latents, masked_image_latents), dim=2) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - - # 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() - - # XLA optimization - if XLA_AVAILABLE: - xm.mark_step() - - # 7. Post-process the image - # Crop the target image - # Since the generated image is a concatenation of the conditional and target regions, - # we need to extract only the target regions based on their positions - image = [] - if output_type == "latent": - image = latents - else: - for b in range(len(latents)): - cur_image_size = processor_output["image_size"][b % batch_size] - cur_target_position = processor_output["target_position"][b % batch_size] - cur_latent = self._unpack_latents(latents[b].unsqueeze(0), cur_image_size, self.vae_scale_factor)[-1] - cur_latent = (cur_latent / self.vae.config.scaling_factor) + self.vae.config.shift_factor - cur_image = self.vae.decode(cur_latent, return_dict=False)[0] - cur_image = self.image_processor.postprocess(cur_image, output_type=output_type)[0] - - start = 0 - cropped = [] - for i, size in enumerate(cur_image_size[-1]): - if cur_target_position[i]: - if output_type == "pil": - cropped.append(cur_image.crop((start, 0, start + size[1], size[0]))) - else: - cropped.append(cur_image[0 : size[0], start : start + size[1]]) - start += size[1] - image.append(cropped) - if output_type != "pil": - image = np.concatenate([arr[None] for sub_image in image for arr in sub_image], axis=0) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return FluxPipelineOutput(images=image) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/visualcloze_utils.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/visualcloze_utils.py deleted file mode 100644 index 5d221bc1e8b122642285760f137464b59c3f747d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/visualcloze/visualcloze_utils.py +++ /dev/null @@ -1,251 +0,0 @@ -# Copyright 2025 VisualCloze 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 Dict, List, Optional, Tuple, Union - -import torch -from PIL import Image - -from ...image_processor import VaeImageProcessor - - -class VisualClozeProcessor(VaeImageProcessor): - """ - Image processor for the VisualCloze pipeline. - - This processor handles the preprocessing of images for visual cloze tasks, including resizing, normalization, and - mask generation. - - Args: - resolution (int, optional): - Target resolution for processing images. Each image will be resized to this resolution before being - concatenated to avoid the out-of-memory error. Defaults to 384. - *args: Additional arguments passed to [~image_processor.VaeImageProcessor] - **kwargs: Additional keyword arguments passed to [~image_processor.VaeImageProcessor] - """ - - def __init__(self, *args, resolution: int = 384, **kwargs): - super().__init__(*args, **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]]: - """ - Preprocesses input images for the VisualCloze pipeline. - - This function handles the preprocessing of input images by: - 1. Resizing and cropping images to maintain consistent dimensions - 2. Converting images to the Tensor format for the VAE - 3. Normalizing pixel values - 4. Tracking image sizes and positions of target images - - Args: - input_images (List[List[Optional[Image.Image]]]): - 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 - - In the last row, condition images are provided and the target images are placed as None - vae_scale_factor (int): - 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 - """ - 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] = [] - - # Process each sample - for i in range(n_samples): - # Determine size from first non-None image - for j in range(n_task_images): - if input_images[i][j] is not None: - aspect_ratio = input_images[i][j].width / input_images[i][j].height - target_area = self.resolution * self.resolution - new_h = int((target_area / aspect_ratio) ** 0.5) - new_w = int(new_h * aspect_ratio) - - new_w = max(new_w // divisible, 1) * divisible - new_h = max(new_h // divisible, 1) * divisible - resize_size[i] = (new_w, new_h) - break - - # Process all images in the sample - for j in range(n_task_images): - if input_images[i][j] is not None: - target = self._resize_and_crop(input_images[i][j], resize_size[i][0], resize_size[i][1]) - processed_images[i].append(target) - if i == n_samples - 1: - target_position.append(0) - else: - blank = Image.new("RGB", resize_size[i] or (self.resolution, self.resolution), (0, 0, 0)) - processed_images[i].append(blank) - if i == n_samples - 1: - target_position.append(1) - - # Ensure consistent width for multiple target images when there are multiple target images - if len(target_position) > 1 and sum(target_position) > 1: - new_w = resize_size[n_samples - 1][0] or 384 - for i in range(len(processed_images)): - for j in range(len(processed_images[i])): - if processed_images[i][j] is not None: - new_h = int(processed_images[i][j].height * (new_w / processed_images[i][j].width)) - new_w = int(new_w / 16) * 16 - new_h = int(new_h / 16) * 16 - processed_images[i][j] = self.height(processed_images[i][j], new_h, new_w) - - # Convert to tensors and normalize - image_sizes = [] - for i in range(len(processed_images)): - image_sizes.append([[img.height, img.width] for img in processed_images[i]]) - for j, image in enumerate(processed_images[i]): - image = self.pil_to_numpy(image) - image = self.numpy_to_pt(image) - image = self.normalize(image) - processed_images[i][j] = 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]]: - """ - Generate masks for the VisualCloze pipeline. - - Args: - input_images (List[List[Image.Image]]): - Processed images from preprocess_image - target_position (List[int]): - Binary list marking the positions of target images (1 for target, 0 for condition) - - Returns: - List[List[torch.Tensor]]: - A nested list of mask tensors (1 for target positions, 0 for condition images) - """ - mask = [] - for i, row in enumerate(input_images): - if i == len(input_images) - 1: # Query row - row_masks = [ - torch.full((1, 1, row[0].shape[2], row[0].shape[3]), fill_value=m) for m in target_position - ] - else: # In-context examples - row_masks = [ - torch.full((1, 1, row[0].shape[2], row[0].shape[3]), fill_value=0) for _ in target_position - ] - mask.append(row_masks) - return mask - - def preprocess_image_upsampling( - self, - input_images: List[List[Image.Image]], - height: int, - width: int, - ) -> Tuple[List[List[Image.Image]], List[List[List[int]]]]: - """Process images for the upsampling stage in the VisualCloze pipeline. - - Args: - input_images: Input image to process - height: Target height - width: Target width - - Returns: - Tuple of processed image and its size - """ - image = self.resize(input_images[0][0], height, width) - image = self.pil_to_numpy(image) # to np - image = self.numpy_to_pt(image) # to pt - image = self.normalize(image) - - input_images[0][0] = image - image_sizes = [[[height, width]]] - return input_images, image_sizes - - 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: - layout_instruction = ( - f"A grid layout with {size[0]} rows and {size[1]} columns, displaying {size[0] * size[1]} images arranged side by side.", - ) - return layout_instruction - - 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, - upsampling: bool = False, - vae_scale_factor: int = 16, - ) -> 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 - 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 - """ - if isinstance(task_prompt, str): - task_prompt = [task_prompt] - content_prompt = [content_prompt] - input_images = [input_images] - - output = { - "init_image": [], - "mask": [], - "task_prompt": task_prompt if not upsampling else [None for _ in range(len(task_prompt))], - "content_prompt": content_prompt, - "layout_prompt": [], - "target_position": [], - "image_size": [], - } - for i in range(len(task_prompt)): - if upsampling: - layout_prompt = None - else: - layout_prompt = self.get_layout_prompt((len(input_images[i]), len(input_images[i][0]))) - - if upsampling: - cur_processed_images, cur_image_size = self.preprocess_image_upsampling( - input_images[i], height=height, width=width - ) - cur_mask = self.preprocess_mask_upsampling(cur_processed_images) - else: - cur_processed_images, cur_image_size, cur_target_position = self.preprocess_image( - input_images[i], vae_scale_factor=vae_scale_factor - ) - cur_mask = self.preprocess_mask(cur_processed_images, cur_target_position) - - output["target_position"].append(cur_target_position) - - output["image_size"].append(cur_image_size) - output["init_image"].append(cur_processed_images) - output["mask"].append(cur_mask) - output["layout_prompt"].append(layout_prompt) - - return output diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__init__.py deleted file mode 100644 index bb96372b1db2b2c29f86cd7d75428161d4bbc824..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__init__.py +++ /dev/null @@ -1,53 +0,0 @@ -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_wan"] = ["WanPipeline"] - _import_structure["pipeline_wan_i2v"] = ["WanImageToVideoPipeline"] - _import_structure["pipeline_wan_vace"] = ["WanVACEPipeline"] - _import_structure["pipeline_wan_video2video"] = ["WanVideoToVideoPipeline"] -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_wan import WanPipeline - from .pipeline_wan_i2v import WanImageToVideoPipeline - from .pipeline_wan_vace import WanVACEPipeline - from .pipeline_wan_video2video import WanVideoToVideoPipeline - -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/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 148b697e940ca0511560c7a2b3307a870ce07be2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_output.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_output.cpython-310.pyc deleted file mode 100644 index 6fa4bc0e8d84b3ea76d7b25e2a3de4e40465cca9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_output.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_wan.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_wan.cpython-310.pyc deleted file mode 100644 index 303475ab7877949883bec6c0b3ef6677737b6880..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_wan.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_wan_i2v.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_wan_i2v.cpython-310.pyc deleted file mode 100644 index c02de1ac80b7b938ae89e937780b43194037f2f9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_wan_i2v.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_wan_vace.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_wan_vace.cpython-310.pyc deleted file mode 100644 index edb9ccd0d419d1512c62fe2796a4ea867065c69e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_wan_vace.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_wan_video2video.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_wan_video2video.cpython-310.pyc deleted file mode 100644 index 2bcdf54f4d52e25e23f2c590f351ecd3f95c6140..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/__pycache__/pipeline_wan_video2video.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_output.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_output.py deleted file mode 100644 index 88907ad0f0a1dfc828879f30bdd31b9043396144..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_output.py +++ /dev/null @@ -1,20 +0,0 @@ -from dataclasses import dataclass - -import torch - -from diffusers.utils import BaseOutput - - -@dataclass -class WanPipelineOutput(BaseOutput): - r""" - 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 - 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/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_wan.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_wan.py deleted file mode 100644 index 6df66118b068d6f802382b54739c93f436ff3fa4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_wan.py +++ /dev/null @@ -1,598 +0,0 @@ -# Copyright 2025 The Wan 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, Dict, List, Optional, Union - -import regex as re -import torch -from transformers import AutoTokenizer, UMT5EncoderModel - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...loaders import WanLoraLoaderMixin -from ...models import AutoencoderKLWan, WanTransformer3DModel -from ...schedulers import FlowMatchEulerDiscreteScheduler -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 WanPipelineOutput - - -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, WanPipeline - >>> from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler - - >>> # Available models: Wan-AI/Wan2.1-T2V-14B-Diffusers, Wan-AI/Wan2.1-T2V-1.3B-Diffusers - >>> model_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers" - >>> vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32) - >>> pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16) - >>> flow_shift = 5.0 # 5.0 for 720P, 3.0 for 480P - >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift) - >>> 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=720, - ... width=1280, - ... num_frames=81, - ... guidance_scale=5.0, - ... ).frames[0] - >>> export_to_video(output, "output.mp4", fps=16) - ``` -""" - - -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 - - -class WanPipeline(DiffusionPipeline, WanLoraLoaderMixin): - r""" - Pipeline for text-to-video generation using Wan. - - 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 ([`WanTransformer3DModel`]): - Conditional Transformer to denoise the input 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. - """ - - model_cpu_offload_seq = "text_encoder->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - tokenizer: AutoTokenizer, - text_encoder: UMT5EncoderModel, - transformer: WanTransformer3DModel, - vae: AutoencoderKLWan, - scheduler: FlowMatchEulerDiscreteScheduler, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - ) - - 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) - - def _get_t5_prompt_embeds( - self, - prompt: Union[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 = 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 - - 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: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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, - ): - 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)}") - - def prepare_latents( - self, - batch_size: int, - num_channels_latents: int = 16, - 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, - ) -> 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 - - @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: Union[str, List[str]] = None, - negative_prompt: Union[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, - output_type: Optional[str] = "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"], - max_sequence_length: int = 512, - ): - 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 `480`): - The height in pixels of the generated image. - width (`int`, defaults to `832`): - The width in pixels of the generated image. - num_frames (`int`, defaults to `81`): - 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 [`WanPipelineOutput`] 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: - [`~WanPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`WanPipelineOutput`] 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 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, - ) - - if num_frames % self.vae_scale_factor_temporal != 1: - logger.warning( - f"`num_frames - 1` has to be divisible by {self.vae_scale_factor_temporal}. Rounding to the nearest number." - ) - num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 - 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 - - # 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 timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - 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, - torch.float32, - device, - generator, - latents, - ) - - # 6. Denoising loop - 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 - 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( - 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 - 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, 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() - - self._current_timestep = None - - 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.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 = latents / latents_std + latents_mean - video = self.vae.decode(latents, return_dict=False)[0] - video = self.video_processor.postprocess_video(video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return WanPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_wan_i2v.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_wan_i2v.py deleted file mode 100644 index c71138a97dd9ade607fc10a5f8a6ca05447f1653..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_wan_i2v.py +++ /dev/null @@ -1,744 +0,0 @@ -# Copyright 2025 The Wan 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, Dict, List, Optional, Tuple, Union - -import PIL -import regex as re -import torch -from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput -from ...loaders import WanLoraLoaderMixin -from ...models import AutoencoderKLWan, WanTransformer3DModel -from ...schedulers import FlowMatchEulerDiscreteScheduler -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 WanPipelineOutput - - -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 - >>> import numpy as np - >>> from diffusers import AutoencoderKLWan, WanImageToVideoPipeline - >>> from diffusers.utils import export_to_video, load_image - >>> from transformers import CLIPVisionModel - - >>> # Available models: Wan-AI/Wan2.1-I2V-14B-480P-Diffusers, Wan-AI/Wan2.1-I2V-14B-720P-Diffusers - >>> model_id = "Wan-AI/Wan2.1-I2V-14B-480P-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) - >>> pipe = WanImageToVideoPipeline.from_pretrained( - ... model_id, vae=vae, image_encoder=image_encoder, torch_dtype=torch.bfloat16 - ... ) - >>> pipe.to("cuda") - - >>> image = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" - ... ) - >>> max_area = 480 * 832 - >>> 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 - >>> image = image.resize((width, height)) - >>> prompt = ( - ... "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in " - ... "the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot." - ... ) - >>> 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( - ... image=image, - ... prompt=prompt, - ... negative_prompt=negative_prompt, - ... height=height, - ... width=width, - ... num_frames=81, - ... guidance_scale=5.0, - ... ).frames[0] - >>> export_to_video(output, "output.mp4", fps=16) - ``` -""" - - -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.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents -def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = 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 WanImageToVideoPipeline(DiffusionPipeline, WanLoraLoaderMixin): - r""" - Pipeline for image-to-video generation using Wan. - - 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. - image_encoder ([`CLIPVisionModel`]): - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPVisionModel), specifically - the - [clip-vit-huge-patch14](https://github.com/mlfoundations/open_clip/blob/main/docs/PRETRAINED.md#vit-h14-xlm-roberta-large) - variant. - transformer ([`WanTransformer3DModel`]): - Conditional Transformer to denoise the input 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. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - tokenizer: AutoTokenizer, - text_encoder: UMT5EncoderModel, - image_encoder: CLIPVisionModel, - image_processor: CLIPImageProcessor, - transformer: WanTransformer3DModel, - vae: AutoencoderKLWan, - scheduler: FlowMatchEulerDiscreteScheduler, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - image_encoder=image_encoder, - transformer=transformer, - scheduler=scheduler, - image_processor=image_processor, - ) - - 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) - self.image_processor = image_processor - - def _get_t5_prompt_embeds( - self, - prompt: Union[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 = 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 - - def encode_image( - self, - image: PipelineImageInput, - device: Optional[torch.device] = None, - ): - device = device or self._execution_device - image = self.image_processor(images=image, return_tensors="pt").to(device) - image_embeds = self.image_encoder(**image, output_hidden_states=True) - return image_embeds.hidden_states[-2] - - # 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, - 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, - max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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, - image, - height, - width, - prompt_embeds=None, - negative_prompt_embeds=None, - image_embeds=None, - callback_on_step_end_tensor_inputs=None, - ): - if image is not None and image_embeds is not None: - raise ValueError( - f"Cannot forward both `image`: {image} and `image_embeds`: {image_embeds}. Please make sure to" - " only forward one of the two." - ) - if image is None and image_embeds is None: - raise ValueError( - "Provide either `image` or `prompt_embeds`. Cannot leave both `image` and `image_embeds` undefined." - ) - if image is not None and not isinstance(image, torch.Tensor) and not isinstance(image, PIL.Image.Image): - raise ValueError(f"`image` has to be of type `torch.Tensor` or `PIL.Image.Image` but is {type(image)}") - 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)}") - - def prepare_latents( - self, - image: PipelineImageInput, - batch_size: int, - num_channels_latents: int = 16, - 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]: - 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 - - shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_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." - ) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device=device, dtype=dtype) - - image = image.unsqueeze(2) - if last_image is None: - video_condition = torch.cat( - [image, image.new_zeros(image.shape[0], image.shape[1], num_frames - 1, height, width)], dim=2 - ) - else: - last_image = last_image.unsqueeze(2) - video_condition = torch.cat( - [image, image.new_zeros(image.shape[0], image.shape[1], num_frames - 2, height, width), last_image], - dim=2, - ) - video_condition = video_condition.to(device=device, dtype=self.vae.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 - ) - - if isinstance(generator, list): - latent_condition = [ - retrieve_latents(self.vae.encode(video_condition), sample_mode="argmax") for _ in generator - ] - latent_condition = torch.cat(latent_condition) - else: - latent_condition = retrieve_latents(self.vae.encode(video_condition), sample_mode="argmax") - latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1) - - latent_condition = latent_condition.to(dtype) - latent_condition = (latent_condition - latents_mean) * latents_std - - mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, latent_width) - - if last_image is None: - mask_lat_size[:, :, list(range(1, num_frames))] = 0 - else: - mask_lat_size[:, :, list(range(1, num_frames - 1))] = 0 - first_frame_mask = mask_lat_size[:, :, 0:1] - first_frame_mask = torch.repeat_interleave(first_frame_mask, dim=2, repeats=self.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, self.vae_scale_factor_temporal, latent_height, latent_width) - mask_lat_size = mask_lat_size.transpose(1, 2) - mask_lat_size = mask_lat_size.to(latent_condition.device) - - return latents, torch.concat([mask_lat_size, latent_condition], dim=1) - - @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 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, - image: PipelineImageInput, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[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, - last_image: Optional[torch.Tensor] = None, - output_type: Optional[str] = "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"], - max_sequence_length: int = 512, - ): - r""" - The call function to 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. - 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`). - height (`int`, defaults to `480`): - The height of the generated video. - width (`int`, defaults to `832`): - The width of the generated video. - num_frames (`int`, defaults to `81`): - 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. - negative_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 `negative_prompt` input argument. - image_embeds (`torch.Tensor`, *optional*): - Pre-generated image embeddings. Can be used to easily tweak image inputs (weighting). If not provided, - image embeddings are generated from the `image` 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 [`WanPipelineOutput`] 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: - [`~WanPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`WanPipelineOutput`] 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 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, - image, - height, - width, - prompt_embeds, - negative_prompt_embeds, - image_embeds, - callback_on_step_end_tensor_inputs, - ) - - if num_frames % self.vae_scale_factor_temporal != 1: - logger.warning( - f"`num_frames - 1` has to be divisible by {self.vae_scale_factor_temporal}. Rounding to the nearest number." - ) - num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 - 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 - - # 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, - ) - - # Encode image embedding - 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) - - if image_embeds is None: - if last_image is None: - image_embeds = self.encode_image(image, device) - else: - image_embeds = self.encode_image([image, last_image], device) - image_embeds = image_embeds.repeat(batch_size, 1, 1) - image_embeds = image_embeds.to(transformer_dtype) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - num_channels_latents = self.vae.config.z_dim - image = self.video_processor.preprocess(image, height=height, width=width).to(device, dtype=torch.float32) - if last_image is not None: - last_image = self.video_processor.preprocess(last_image, height=height, width=width).to( - device, dtype=torch.float32 - ) - latents, condition = self.prepare_latents( - image, - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - num_frames, - torch.float32, - device, - generator, - latents, - last_image, - ) - - # 6. Denoising loop - 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 - 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( - 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 - 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, 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() - - self._current_timestep = None - - 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.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 = latents / latents_std + latents_mean - video = self.vae.decode(latents, return_dict=False)[0] - video = self.video_processor.postprocess_video(video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return WanPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_wan_vace.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_wan_vace.py deleted file mode 100644 index e5f83dd401ad0406e709f75b6645dc20316b857e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_wan_vace.py +++ /dev/null @@ -1,976 +0,0 @@ -# Copyright 2025 The Wan 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, Dict, List, Optional, Union - -import PIL.Image -import regex as re -import torch -from transformers import AutoTokenizer, UMT5EncoderModel - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...image_processor import PipelineImageInput -from ...loaders import WanLoraLoaderMixin -from ...models import AutoencoderKLWan, WanVACETransformer3DModel -from ...schedulers import FlowMatchEulerDiscreteScheduler -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 WanPipelineOutput - - -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 - >>> import PIL.Image - >>> from diffusers import AutoencoderKLWan, WanVACEPipeline - >>> from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler - >>> from diffusers.utils import export_to_video, load_image - def prepare_video_and_mask(first_img: PIL.Image.Image, last_img: PIL.Image.Image, height: int, width: int, num_frames: int): - first_img = first_img.resize((width, height)) - last_img = last_img.resize((width, height)) - frames = [] - frames.append(first_img) - # Ideally, this should be 127.5 to match original code, but they perform computation on numpy arrays - # whereas we are passing PIL images. If you choose to pass numpy arrays, you can set it to 127.5 to - # match the original code. - frames.extend([PIL.Image.new("RGB", (width, height), (128, 128, 128))] * (num_frames - 2)) - frames.append(last_img) - mask_black = PIL.Image.new("L", (width, height), 0) - mask_white = PIL.Image.new("L", (width, height), 255) - mask = [mask_black, *[mask_white] * (num_frames - 2), mask_black] - return frames, mask - - >>> # Available checkpoints: Wan-AI/Wan2.1-VACE-1.3B-diffusers, Wan-AI/Wan2.1-VACE-14B-diffusers - >>> model_id = "Wan-AI/Wan2.1-VACE-1.3B-diffusers" - >>> vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32) - >>> pipe = WanVACEPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16) - >>> flow_shift = 3.0 # 5.0 for 720P, 3.0 for 480P - >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift) - >>> pipe.to("cuda") - - >>> 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." - >>> 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" - >>> first_frame = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_first_frame.png" - ... ) - >>> last_frame = load_image( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_last_frame.png>>> " - ... ) - - >>> height = 512 - >>> width = 512 - >>> num_frames = 81 - >>> video, mask = prepare_video_and_mask(first_frame, last_frame, height, width, num_frames) - - >>> output = pipe( - ... video=video, - ... mask=mask, - ... prompt=prompt, - ... negative_prompt=negative_prompt, - ... height=height, - ... width=width, - ... num_frames=num_frames, - ... num_inference_steps=30, - ... guidance_scale=5.0, - ... generator=torch.Generator().manual_seed(42), - ... ).frames[0] - >>> export_to_video(output, "output.mp4", fps=16) - ``` -""" - - -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.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents -def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = 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 WanVACEPipeline(DiffusionPipeline, WanLoraLoaderMixin): - r""" - Pipeline for controllable generation using Wan. - - 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 ([`WanTransformer3DModel`]): - Conditional Transformer to denoise the input 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. - """ - - model_cpu_offload_seq = "text_encoder->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - tokenizer: AutoTokenizer, - text_encoder: UMT5EncoderModel, - transformer: WanVACETransformer3DModel, - vae: AutoencoderKLWan, - scheduler: FlowMatchEulerDiscreteScheduler, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - ) - - 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) - - # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds - def _get_t5_prompt_embeds( - self, - prompt: Union[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 = 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 - - # 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, - 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, - max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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, - video=None, - mask=None, - reference_images=None, - ): - base = self.vae_scale_factor_spatial * self.transformer.config.patch_size[1] - if height % base != 0 or width % base != 0: - raise ValueError(f"`height` and `width` have to be divisible by {base} 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 video is not None: - if mask is not None: - if len(video) != len(mask): - raise ValueError( - f"Length of `video` {len(video)} and `mask` {len(mask)} do not match. Please make sure that" - " they have the same length." - ) - if reference_images is not None: - is_pil_image = isinstance(reference_images, PIL.Image.Image) - is_list_of_pil_images = isinstance(reference_images, list) and all( - isinstance(ref_img, PIL.Image.Image) for ref_img in reference_images - ) - is_list_of_list_of_pil_images = isinstance(reference_images, list) and all( - isinstance(ref_img, list) and all(isinstance(ref_img_, PIL.Image.Image) for ref_img_ in ref_img) - for ref_img in reference_images - ) - if not (is_pil_image or is_list_of_pil_images or is_list_of_list_of_pil_images): - raise ValueError( - "`reference_images` has to be of type `PIL.Image.Image` or `list` of `PIL.Image.Image`, or " - "`list` of `list` of `PIL.Image.Image`, but is {type(reference_images)}" - ) - if is_list_of_list_of_pil_images and len(reference_images) != 1: - raise ValueError( - "The pipeline only supports generating one video at a time at the moment. When passing a list " - "of list of reference images, where the outer list corresponds to the batch size and the inner " - "list corresponds to list of conditioning images per video, please make sure to only pass " - "one inner list of reference images (i.e., `[[, , ...]]`" - ) - elif mask is not None: - raise ValueError("`mask` can only be passed if `video` is passed as well.") - - 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, - batch_size: int = 1, - height: int = 480, - width: int = 832, - num_frames: int = 81, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - ): - if video is not None: - base = self.vae_scale_factor_spatial * self.transformer.config.patch_size[1] - video_height, video_width = self.video_processor.get_default_height_width(video[0]) - - if video_height * video_width > height * width: - scale = min(width / video_width, height / video_height) - video_height, video_width = int(video_height * scale), int(video_width * scale) - - if video_height % base != 0 or video_width % base != 0: - logger.warning( - f"Video height and width should be divisible by {base}, but got {video_height} and {video_width}. " - ) - video_height = (video_height // base) * base - video_width = (video_width // base) * base - - assert video_height * video_width <= height * width - - video = self.video_processor.preprocess_video(video, video_height, video_width) - image_size = (video_height, video_width) # Use the height/width of video (with possible rescaling) - else: - video = torch.zeros(batch_size, 3, num_frames, height, width, dtype=dtype, device=device) - image_size = (height, width) # Use the height/width provider by user - - if mask is not None: - mask = self.video_processor.preprocess_video(mask, image_size[0], image_size[1]) - mask = torch.clamp((mask + 1) / 2, min=0, max=1) - else: - mask = torch.ones_like(video) - - video = video.to(dtype=dtype, device=device) - mask = mask.to(dtype=dtype, device=device) - - # Make a list of list of images where the outer list corresponds to video batch size and the inner list - # corresponds to list of conditioning images per video - if reference_images is None or isinstance(reference_images, PIL.Image.Image): - reference_images = [[reference_images] for _ in range(video.shape[0])] - elif isinstance(reference_images, (list, tuple)) and isinstance(next(iter(reference_images)), PIL.Image.Image): - reference_images = [reference_images] - elif ( - isinstance(reference_images, (list, tuple)) - and isinstance(next(iter(reference_images)), list) - and isinstance(next(iter(reference_images[0])), PIL.Image.Image) - ): - reference_images = reference_images - else: - raise ValueError( - "`reference_images` has to be of type `PIL.Image.Image` or `list` of `PIL.Image.Image`, or " - "`list` of `list` of `PIL.Image.Image`, but is {type(reference_images)}" - ) - - if video.shape[0] != len(reference_images): - raise ValueError( - f"Batch size of `video` {video.shape[0]} and length of `reference_images` {len(reference_images)} does not match." - ) - - ref_images_lengths = [len(reference_images_batch) for reference_images_batch in reference_images] - if any(l != ref_images_lengths[0] for l in ref_images_lengths): - raise ValueError( - f"All batches of `reference_images` should have the same length, but got {ref_images_lengths}. Support for this " - "may be added in the future." - ) - - reference_images_preprocessed = [] - for i, reference_images_batch in enumerate(reference_images): - preprocessed_images = [] - for j, image in enumerate(reference_images_batch): - if image is None: - continue - image = self.video_processor.preprocess(image, None, None) - img_height, img_width = image.shape[-2:] - scale = min(image_size[0] / img_height, image_size[1] / img_width) - new_height, new_width = int(img_height * scale), int(img_width * scale) - resized_image = torch.nn.functional.interpolate( - image, size=(new_height, new_width), mode="bilinear", align_corners=False - ).squeeze(0) # [C, H, W] - top = (image_size[0] - new_height) // 2 - left = (image_size[1] - new_width) // 2 - canvas = torch.ones(3, *image_size, device=device, dtype=dtype) - canvas[:, top : top + new_height, left : left + new_width] = resized_image - preprocessed_images.append(canvas) - reference_images_preprocessed.append(preprocessed_images) - - return video, mask, reference_images_preprocessed - - 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, - ) -> torch.Tensor: - device = device or self._execution_device - - if isinstance(generator, list): - # TODO: support this - raise ValueError("Passing a list of generators is not yet supported. This may be supported in the future.") - - if reference_images is None: - # For each batch of video, we set no re - # ference image (as one or more can be passed by user) - reference_images = [[None] for _ in range(video.shape[0])] - else: - if video.shape[0] != len(reference_images): - raise ValueError( - f"Batch size of `video` {video.shape[0]} and length of `reference_images` {len(reference_images)} does not match." - ) - - if video.shape[0] != 1: - # TODO: support this - raise ValueError( - "Generating with more than one video is not yet supported. This may be supported in the future." - ) - - vae_dtype = self.vae.dtype - video = video.to(dtype=vae_dtype) - - latents_mean = torch.tensor(self.vae.config.latents_mean, device=device, dtype=torch.float32).view( - 1, self.vae.config.z_dim, 1, 1, 1 - ) - latents_std = 1.0 / torch.tensor(self.vae.config.latents_std, device=device, dtype=torch.float32).view( - 1, self.vae.config.z_dim, 1, 1, 1 - ) - - if mask is None: - latents = retrieve_latents(self.vae.encode(video), generator, sample_mode="argmax").unbind(0) - latents = ((latents.float() - latents_mean) * latents_std).to(vae_dtype) - else: - mask = mask.to(dtype=vae_dtype) - mask = torch.where(mask > 0.5, 1.0, 0.0) - inactive = video * (1 - mask) - reactive = video * mask - inactive = retrieve_latents(self.vae.encode(inactive), generator, sample_mode="argmax") - reactive = retrieve_latents(self.vae.encode(reactive), generator, sample_mode="argmax") - inactive = ((inactive.float() - latents_mean) * latents_std).to(vae_dtype) - reactive = ((reactive.float() - latents_mean) * latents_std).to(vae_dtype) - latents = torch.cat([inactive, reactive], dim=1) - - latent_list = [] - for latent, reference_images_batch in zip(latents, reference_images): - for reference_image in reference_images_batch: - assert reference_image.ndim == 3 - reference_image = reference_image.to(dtype=vae_dtype) - reference_image = reference_image[None, :, None, :, :] # [1, C, 1, H, W] - reference_latent = retrieve_latents(self.vae.encode(reference_image), generator, sample_mode="argmax") - reference_latent = ((reference_latent.float() - latents_mean) * latents_std).to(vae_dtype) - reference_latent = reference_latent.squeeze(0) # [C, 1, H, W] - reference_latent = torch.cat([reference_latent, torch.zeros_like(reference_latent)], dim=0) - latent = torch.cat([reference_latent.squeeze(0), latent], dim=1) - latent_list.append(latent) - return torch.stack(latent_list) - - def prepare_masks( - self, - mask: torch.Tensor, - reference_images: Optional[List[torch.Tensor]] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - ) -> torch.Tensor: - if isinstance(generator, list): - # TODO: support this - raise ValueError("Passing a list of generators is not yet supported. This may be supported in the future.") - - if reference_images is None: - # For each batch of video, we set no reference image (as one or more can be passed by user) - reference_images = [[None] for _ in range(mask.shape[0])] - else: - if mask.shape[0] != len(reference_images): - raise ValueError( - f"Batch size of `mask` {mask.shape[0]} and length of `reference_images` {len(reference_images)} does not match." - ) - - if mask.shape[0] != 1: - # TODO: support this - raise ValueError( - "Generating with more than one video is not yet supported. This may be supported in the future." - ) - - transformer_patch_size = self.transformer.config.patch_size[1] - - mask_list = [] - for mask_, reference_images_batch in zip(mask, reference_images): - num_channels, num_frames, height, width = mask_.shape - new_num_frames = (num_frames + self.vae_scale_factor_temporal - 1) // self.vae_scale_factor_temporal - new_height = height // (self.vae_scale_factor_spatial * transformer_patch_size) * transformer_patch_size - new_width = width // (self.vae_scale_factor_spatial * transformer_patch_size) * transformer_patch_size - mask_ = mask_[0, :, :, :] - mask_ = mask_.view( - num_frames, new_height, self.vae_scale_factor_spatial, new_width, self.vae_scale_factor_spatial - ) - mask_ = mask_.permute(2, 4, 0, 1, 3).flatten(0, 1) # [8x8, num_frames, new_height, new_width] - mask_ = torch.nn.functional.interpolate( - mask_.unsqueeze(0), size=(new_num_frames, new_height, new_width), mode="nearest-exact" - ).squeeze(0) - num_ref_images = len(reference_images_batch) - if num_ref_images > 0: - mask_padding = torch.zeros_like(mask_[:, :num_ref_images, :, :]) - mask_ = torch.cat([mask_padding, mask_], dim=1) - mask_list.append(mask_) - return torch.stack(mask_list) - - def prepare_latents( - self, - batch_size: int, - num_channels_latents: int = 16, - 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, - ) -> 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 - - @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: 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, - 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, - output_type: Optional[str] = "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"], - max_sequence_length: int = 512, - ): - 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, 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 - less than `1`). - 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*): - 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*): - 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`): - 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 - of control layers in the model (`len(transformer.config.vace_layers)`). - height (`int`, defaults to `480`): - The height in pixels of the generated image. - width (`int`, defaults to `832`): - The width in pixels of the generated image. - num_frames (`int`, defaults to `81`): - 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://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. - 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 [`WanPipelineOutput`] 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: - [`~WanPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`WanPipelineOutput`] 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 isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - # Simplification of implementation for now - if not isinstance(prompt, str): - raise ValueError("Passing a list of prompts is not yet supported. This may be supported in the future.") - if num_videos_per_prompt != 1: - raise ValueError( - "Generating multiple videos per prompt is not yet supported. This may be supported in the future." - ) - - # 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, - video, - mask, - reference_images, - ) - - if num_frames % self.vae_scale_factor_temporal != 1: - logger.warning( - f"`num_frames - 1` has to be divisible by {self.vae_scale_factor_temporal}. Rounding to the nearest number." - ) - num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 - 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 - - # 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] - - vae_dtype = self.vae.dtype - transformer_dtype = self.transformer.dtype - - if isinstance(conditioning_scale, (int, float)): - conditioning_scale = [conditioning_scale] * len(self.transformer.config.vace_layers) - if isinstance(conditioning_scale, list): - if len(conditioning_scale) != len(self.transformer.config.vace_layers): - raise ValueError( - f"Length of `conditioning_scale` {len(conditioning_scale)} does not match number of layers {len(self.transformer.config.vace_layers)}." - ) - conditioning_scale = torch.tensor(conditioning_scale) - if isinstance(conditioning_scale, torch.Tensor): - if conditioning_scale.size(0) != len(self.transformer.config.vace_layers): - raise ValueError( - f"Length of `conditioning_scale` {conditioning_scale.size(0)} does not match number of layers {len(self.transformer.config.vace_layers)}." - ) - conditioning_scale = conditioning_scale.to(device=device, dtype=transformer_dtype) - - # 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, - ) - - 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 timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - video, mask, reference_images = self.preprocess_conditions( - video, - mask, - reference_images, - batch_size, - height, - width, - num_frames, - torch.float32, - device, - ) - num_reference_images = len(reference_images[0]) - - conditioning_latents = self.prepare_video_latents(video, mask, reference_images, generator, device) - mask = self.prepare_masks(mask, reference_images, generator) - conditioning_latents = torch.cat([conditioning_latents, mask], dim=1) - conditioning_latents = conditioning_latents.to(transformer_dtype) - - 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 + num_reference_images * self.vae_scale_factor_temporal, - torch.float32, - device, - generator, - latents, - ) - - if conditioning_latents.shape[2] != latents.shape[2]: - logger.warning( - "The number of frames in the conditioning latents does not match the number of frames to be generated. Generation quality may be affected." - ) - - # 6. Denoising loop - 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 - 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, - control_hidden_states=conditioning_latents, - control_hidden_states_scale=conditioning_scale, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - - if self.do_classifier_free_guidance: - noise_uncond = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, - control_hidden_states=conditioning_latents, - control_hidden_states_scale=conditioning_scale, - 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 - 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, 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() - - self._current_timestep = None - - if not output_type == "latent": - latents = latents[:, :, num_reference_images:] - latents = latents.to(vae_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 = latents / latents_std + latents_mean - video = self.vae.decode(latents, return_dict=False)[0] - video = self.video_processor.postprocess_video(video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return WanPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_wan_video2video.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_wan_video2video.py deleted file mode 100644 index 1a2d2e9c2232396ea55b3255fc5e72c0b5ea5c1c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/wan/pipeline_wan_video2video.py +++ /dev/null @@ -1,725 +0,0 @@ -# Copyright 2025 The Wan 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 inspect -from typing import Any, Callable, Dict, List, Optional, Union - -import regex as re -import torch -from PIL import Image -from transformers import AutoTokenizer, UMT5EncoderModel - -from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...loaders import WanLoraLoaderMixin -from ...models import AutoencoderKLWan, WanTransformer3DModel -from ...schedulers import FlowMatchEulerDiscreteScheduler -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 WanPipelineOutput - - -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, WanVideoToVideoPipeline - >>> from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler - - >>> # Available models: Wan-AI/Wan2.1-T2V-14B-Diffusers, Wan-AI/Wan2.1-T2V-1.3B-Diffusers - >>> model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" - >>> vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32) - >>> pipe = WanVideoToVideoPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16) - >>> flow_shift = 3.0 # 5.0 for 720P, 3.0 for 480P - >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift) - >>> pipe.to("cuda") - - >>> prompt = "A robot standing on a mountain top. The sun is setting in the background" - >>> 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" - >>> video = load_video( - ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/hiker.mp4" - ... ) - >>> output = pipe( - ... video=video, - ... prompt=prompt, - ... negative_prompt=negative_prompt, - ... height=480, - ... width=720, - ... guidance_scale=5.0, - ... strength=0.7, - ... ).frames[0] - >>> export_to_video(output, "output.mp4", fps=16) - ``` -""" - - -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.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 - - -# 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" -): - 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 WanVideoToVideoPipeline(DiffusionPipeline, WanLoraLoaderMixin): - r""" - Pipeline for video-to-video generation using Wan. - - 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 ([`WanTransformer3DModel`]): - Conditional Transformer to denoise the input 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. - """ - - model_cpu_offload_seq = "text_encoder->transformer->vae" - _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - tokenizer: AutoTokenizer, - text_encoder: UMT5EncoderModel, - transformer: WanTransformer3DModel, - vae: AutoencoderKLWan, - scheduler: FlowMatchEulerDiscreteScheduler, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - ) - - 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) - - # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds - def _get_t5_prompt_embeds( - self, - prompt: Union[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 = 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 - - # 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, - 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, - max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = 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, - video=None, - latents=None, - prompt_embeds=None, - negative_prompt_embeds=None, - callback_on_step_end_tensor_inputs=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 video is not None and latents is not None: - raise ValueError("Only one of `video` or `latents` should be provided") - - def prepare_latents( - self, - video: Optional[torch.Tensor] = 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, - ): - 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." - ) - - num_latent_frames = ( - (video.size(2) - 1) // self.vae_scale_factor_temporal + 1 if latents is None else latents.size(1) - ) - shape = ( - batch_size, - num_channels_latents, - num_latent_frames, - height // self.vae_scale_factor_spatial, - width // self.vae_scale_factor_spatial, - ) - - if latents is None: - init_latents = [retrieve_latents(self.vae.encode(vid.unsqueeze(0)), sample_mode="argmax") for vid in video] - - init_latents = torch.cat(init_latents, dim=0).to(dtype) - - 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 - ) - - init_latents = (init_latents - latents_mean) * latents_std - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - if hasattr(self.scheduler, "add_noise"): - latents = self.scheduler.add_noise(init_latents, noise, timestep) - else: - latents = self.scheduler.scale_noise(init_latents, timestep, noise) - else: - latents = latents.to(device) - - return latents - - # Copied from diffusers.pipelines.animatediff.pipeline_animatediff_video2video.AnimateDiffVideoToVideoPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, timesteps, strength, device): - # get the original timestep using init_timestep - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - - t_start = max(num_inference_steps - init_timestep, 0) - timesteps = timesteps[t_start * self.scheduler.order :] - - return timesteps, num_inference_steps - t_start - - @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, - video: List[Image.Image] = None, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, - height: int = 480, - width: int = 832, - num_inference_steps: int = 50, - timesteps: Optional[List[int]] = 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", - 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"], - max_sequence_length: int = 512, - ): - 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, one has to pass `prompt_embeds` - instead. - height (`int`, defaults to `480`): - The height in pixels of the generated image. - width (`int`, defaults to `832`): - The width in pixels of the generated image. - num_frames (`int`, defaults to `81`): - 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. - strength (`float`, defaults to `0.8`): - 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*): - 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 [`WanPipelineOutput`] 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: - [`~WanPipelineOutput`] or `tuple`: - If `return_dict` is `True`, [`WanPipelineOutput`] 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 isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): - callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs - - height = height or self.transformer.config.sample_height * self.vae_scale_factor_spatial - width = width or self.transformer.config.sample_width * self.vae_scale_factor_spatial - num_videos_per_prompt = 1 - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - negative_prompt, - height, - width, - video, - latents, - prompt_embeds, - negative_prompt_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._current_timestep = None - self._interrupt = False - - device = self._execution_device - - # 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 timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, 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) - - if latents is None: - video = self.video_processor.preprocess_video(video, height=height, width=width).to( - device, dtype=torch.float32 - ) - - # 5. Prepare latent variables - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - video, - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - torch.float32, - device, - generator, - latents, - latent_timestep, - ) - - # 6. Denoising loop - 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 - 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( - 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 - 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, 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() - - self._current_timestep = None - - 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.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 = latents / latents_std + latents_mean - video = self.vae.decode(latents, return_dict=False)[0] - video = self.video_processor.postprocess_video(video, output_type=output_type) - else: - video = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return WanPipelineOutput(frames=video) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__init__.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__init__.py deleted file mode 100644 index ddb852d1931558fe0948e81e16cf9a92fc2a114b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__init__.py +++ /dev/null @@ -1,56 +0,0 @@ -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 - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) -else: - _import_structure["modeling_paella_vq_model"] = ["PaellaVQModel"] - _import_structure["modeling_wuerstchen_diffnext"] = ["WuerstchenDiffNeXt"] - _import_structure["modeling_wuerstchen_prior"] = ["WuerstchenPrior"] - _import_structure["pipeline_wuerstchen"] = ["WuerstchenDecoderPipeline"] - _import_structure["pipeline_wuerstchen_combined"] = ["WuerstchenCombinedPipeline"] - _import_structure["pipeline_wuerstchen_prior"] = ["DEFAULT_STAGE_C_TIMESTEPS", "WuerstchenPriorPipeline"] - - -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 .modeling_paella_vq_model import PaellaVQModel - from .modeling_wuerstchen_diffnext import WuerstchenDiffNeXt - from .modeling_wuerstchen_prior import WuerstchenPrior - from .pipeline_wuerstchen import WuerstchenDecoderPipeline - from .pipeline_wuerstchen_combined import WuerstchenCombinedPipeline - from .pipeline_wuerstchen_prior import DEFAULT_STAGE_C_TIMESTEPS, WuerstchenPriorPipeline -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/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 064ef6e0f6ce571f0af2f0ab0be5ade106d84982..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/modeling_paella_vq_model.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/modeling_paella_vq_model.cpython-310.pyc deleted file mode 100644 index 27917068256ae1ce7b039b106d91a1fd5aefad5e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/modeling_paella_vq_model.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/modeling_wuerstchen_common.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/modeling_wuerstchen_common.cpython-310.pyc deleted file mode 100644 index 48b4a8330b017915e6ec0dc0472c9b6f6a5e3b9d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/modeling_wuerstchen_common.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/modeling_wuerstchen_diffnext.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/modeling_wuerstchen_diffnext.cpython-310.pyc deleted file mode 100644 index e8c8fe30898853fc9d1a2f39b7c65c064393c0cb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/modeling_wuerstchen_diffnext.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/modeling_wuerstchen_prior.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/modeling_wuerstchen_prior.cpython-310.pyc deleted file mode 100644 index 13b530928eff125fc2a8f2c76364cd11e838fcc0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/modeling_wuerstchen_prior.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/pipeline_wuerstchen.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/pipeline_wuerstchen.cpython-310.pyc deleted file mode 100644 index fd7f9997f6b710388aa502e693b0080f66746ab4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/pipeline_wuerstchen.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/pipeline_wuerstchen_combined.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/pipeline_wuerstchen_combined.cpython-310.pyc deleted file mode 100644 index 077a75e79f56b548fa71a073e03fa88d7b32b2e9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/pipeline_wuerstchen_combined.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/pipeline_wuerstchen_prior.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/pipeline_wuerstchen_prior.cpython-310.pyc deleted file mode 100644 index 322f19e9dba9be59a318376fd38240cff3930246..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/__pycache__/pipeline_wuerstchen_prior.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py deleted file mode 100644 index 5ab206b15176e5ab9cb636762d38d3c38618124a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py +++ /dev/null @@ -1,172 +0,0 @@ -# Copyright (c) 2022 Dominic Rampas MIT License -# 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 Union - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.autoencoders.vae import DecoderOutput, VectorQuantizer -from ...models.modeling_utils import ModelMixin -from ...models.vq_model import VQEncoderOutput -from ...utils.accelerate_utils import apply_forward_hook - - -class MixingResidualBlock(nn.Module): - """ - Residual block with mixing used by Paella's VQ-VAE. - """ - - def __init__(self, inp_channels, embed_dim): - super().__init__() - # depthwise - self.norm1 = nn.LayerNorm(inp_channels, elementwise_affine=False, eps=1e-6) - self.depthwise = nn.Sequential( - nn.ReplicationPad2d(1), nn.Conv2d(inp_channels, inp_channels, kernel_size=3, groups=inp_channels) - ) - - # channelwise - self.norm2 = nn.LayerNorm(inp_channels, elementwise_affine=False, eps=1e-6) - self.channelwise = nn.Sequential( - nn.Linear(inp_channels, embed_dim), nn.GELU(), nn.Linear(embed_dim, inp_channels) - ) - - self.gammas = nn.Parameter(torch.zeros(6), requires_grad=True) - - def forward(self, x): - mods = self.gammas - x_temp = self.norm1(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) * (1 + mods[0]) + mods[1] - x = x + self.depthwise(x_temp) * mods[2] - x_temp = self.norm2(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) * (1 + mods[3]) + mods[4] - x = x + self.channelwise(x_temp.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) * mods[5] - return x - - -class PaellaVQModel(ModelMixin, ConfigMixin): - r"""VQ-VAE model from Paella model. - - This model inherits from [`ModelMixin`]. Check the superclass documentation for the generic methods the library - implements for all the model (such as downloading or saving, etc.) - - 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. - up_down_scale_factor (int, *optional*, defaults to 2): Up and Downscale factor of the input image. - levels (int, *optional*, defaults to 2): Number of levels in the model. - bottleneck_blocks (int, *optional*, defaults to 12): Number of bottleneck blocks in the model. - embed_dim (int, *optional*, defaults to 384): Number of hidden channels in the model. - latent_channels (int, *optional*, defaults to 4): Number of latent channels in the VQ-VAE model. - num_vq_embeddings (int, *optional*, defaults to 8192): Number of codebook vectors in the VQ-VAE. - scale_factor (float, *optional*, defaults to 0.3764): Scaling factor of the latent space. - """ - - @register_to_config - def __init__( - self, - in_channels: int = 3, - out_channels: int = 3, - up_down_scale_factor: int = 2, - levels: int = 2, - bottleneck_blocks: int = 12, - embed_dim: int = 384, - latent_channels: int = 4, - num_vq_embeddings: int = 8192, - scale_factor: float = 0.3764, - ): - super().__init__() - - c_levels = [embed_dim // (2**i) for i in reversed(range(levels))] - # Encoder blocks - self.in_block = nn.Sequential( - nn.PixelUnshuffle(up_down_scale_factor), - nn.Conv2d(in_channels * up_down_scale_factor**2, c_levels[0], kernel_size=1), - ) - down_blocks = [] - for i in range(levels): - if i > 0: - down_blocks.append(nn.Conv2d(c_levels[i - 1], c_levels[i], kernel_size=4, stride=2, padding=1)) - block = MixingResidualBlock(c_levels[i], c_levels[i] * 4) - down_blocks.append(block) - down_blocks.append( - nn.Sequential( - nn.Conv2d(c_levels[-1], latent_channels, kernel_size=1, bias=False), - nn.BatchNorm2d(latent_channels), # then normalize them to have mean 0 and std 1 - ) - ) - self.down_blocks = nn.Sequential(*down_blocks) - - # Vector Quantizer - self.vquantizer = VectorQuantizer(num_vq_embeddings, vq_embed_dim=latent_channels, legacy=False, beta=0.25) - - # Decoder blocks - up_blocks = [nn.Sequential(nn.Conv2d(latent_channels, c_levels[-1], kernel_size=1))] - for i in range(levels): - for j in range(bottleneck_blocks if i == 0 else 1): - block = MixingResidualBlock(c_levels[levels - 1 - i], c_levels[levels - 1 - i] * 4) - up_blocks.append(block) - if i < levels - 1: - up_blocks.append( - nn.ConvTranspose2d( - c_levels[levels - 1 - i], c_levels[levels - 2 - i], kernel_size=4, stride=2, padding=1 - ) - ) - self.up_blocks = nn.Sequential(*up_blocks) - self.out_block = nn.Sequential( - nn.Conv2d(c_levels[0], out_channels * up_down_scale_factor**2, kernel_size=1), - nn.PixelShuffle(up_down_scale_factor), - ) - - @apply_forward_hook - def encode(self, x: torch.Tensor, return_dict: bool = True) -> VQEncoderOutput: - h = self.in_block(x) - h = self.down_blocks(h) - - if not return_dict: - return (h,) - - return VQEncoderOutput(latents=h) - - @apply_forward_hook - def decode( - self, h: torch.Tensor, force_not_quantize: bool = True, return_dict: bool = True - ) -> Union[DecoderOutput, torch.Tensor]: - if not force_not_quantize: - quant, _, _ = self.vquantizer(h) - else: - quant = h - - x = self.up_blocks(quant) - dec = self.out_block(x) - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) - - def forward(self, sample: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: - r""" - Args: - sample (`torch.Tensor`): Input sample. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`DecoderOutput`] instead of a plain tuple. - """ - x = sample - h = self.encode(x).latents - dec = self.decode(h).sample - - if not return_dict: - return (dec,) - - return DecoderOutput(sample=dec) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py deleted file mode 100644 index 73e71b3076fbca259ae76138bc4ab3d3797e2755..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py +++ /dev/null @@ -1,76 +0,0 @@ -import torch -import torch.nn as nn - -from ...models.attention_processor import Attention - - -class WuerstchenLayerNorm(nn.LayerNorm): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def forward(self, x): - x = x.permute(0, 2, 3, 1) - x = super().forward(x) - return x.permute(0, 3, 1, 2) - - -class TimestepBlock(nn.Module): - def __init__(self, c, c_timestep): - super().__init__() - - self.mapper = nn.Linear(c_timestep, c * 2) - - def forward(self, x, t): - a, b = self.mapper(t)[:, :, None, None].chunk(2, dim=1) - return x * (1 + a) + b - - -class ResBlock(nn.Module): - def __init__(self, c, c_skip=0, kernel_size=3, dropout=0.0): - super().__init__() - - self.depthwise = nn.Conv2d(c + c_skip, c, kernel_size=kernel_size, padding=kernel_size // 2, groups=c) - self.norm = WuerstchenLayerNorm(c, elementwise_affine=False, eps=1e-6) - self.channelwise = nn.Sequential( - nn.Linear(c, c * 4), nn.GELU(), GlobalResponseNorm(c * 4), nn.Dropout(dropout), nn.Linear(c * 4, c) - ) - - def forward(self, x, x_skip=None): - x_res = x - if x_skip is not None: - x = torch.cat([x, x_skip], dim=1) - x = self.norm(self.depthwise(x)).permute(0, 2, 3, 1) - x = self.channelwise(x).permute(0, 3, 1, 2) - return x + x_res - - -# from https://github.com/facebookresearch/ConvNeXt-V2/blob/3608f67cc1dae164790c5d0aead7bf2d73d9719b/models/utils.py#L105 -class GlobalResponseNorm(nn.Module): - def __init__(self, dim): - super().__init__() - self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) - self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) - - def forward(self, x): - agg_norm = torch.norm(x, p=2, dim=(1, 2), keepdim=True) - stand_div_norm = agg_norm / (agg_norm.mean(dim=-1, keepdim=True) + 1e-6) - return self.gamma * (x * stand_div_norm) + self.beta + x - - -class AttnBlock(nn.Module): - def __init__(self, c, c_cond, nhead, self_attn=True, dropout=0.0): - super().__init__() - - self.self_attn = self_attn - self.norm = WuerstchenLayerNorm(c, elementwise_affine=False, eps=1e-6) - self.attention = Attention(query_dim=c, heads=nhead, dim_head=c // nhead, dropout=dropout, bias=True) - self.kv_mapper = nn.Sequential(nn.SiLU(), nn.Linear(c_cond, c)) - - def forward(self, x, kv): - kv = self.kv_mapper(kv) - norm_x = self.norm(x) - if self.self_attn: - batch_size, channel, _, _ = x.shape - kv = torch.cat([norm_x.view(batch_size, channel, -1).transpose(1, 2), kv], dim=1) - x = x + self.attention(norm_x, encoder_hidden_states=kv) - return x diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_diffnext.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_diffnext.py deleted file mode 100644 index 77ae597655d164e25add03a641a6b8eca5396200..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_diffnext.py +++ /dev/null @@ -1,254 +0,0 @@ -# Copyright (c) 2023 Dominic Rampas MIT License -# 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 math - -import numpy as np -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.modeling_utils import ModelMixin -from .modeling_wuerstchen_common import AttnBlock, GlobalResponseNorm, TimestepBlock, WuerstchenLayerNorm - - -class WuerstchenDiffNeXt(ModelMixin, ConfigMixin): - @register_to_config - def __init__( - self, - c_in=4, - c_out=4, - c_r=64, - patch_size=2, - c_cond=1024, - c_hidden=[320, 640, 1280, 1280], - nhead=[-1, 10, 20, 20], - blocks=[4, 4, 14, 4], - level_config=["CT", "CTA", "CTA", "CTA"], - inject_effnet=[False, True, True, True], - effnet_embd=16, - clip_embd=1024, - kernel_size=3, - dropout=0.1, - ): - super().__init__() - self.c_r = c_r - self.c_cond = c_cond - if not isinstance(dropout, list): - dropout = [dropout] * len(c_hidden) - - # CONDITIONING - self.clip_mapper = nn.Linear(clip_embd, c_cond) - self.effnet_mappers = nn.ModuleList( - [ - nn.Conv2d(effnet_embd, c_cond, kernel_size=1) if inject else None - for inject in inject_effnet + list(reversed(inject_effnet)) - ] - ) - self.seq_norm = nn.LayerNorm(c_cond, elementwise_affine=False, eps=1e-6) - - self.embedding = nn.Sequential( - nn.PixelUnshuffle(patch_size), - nn.Conv2d(c_in * (patch_size**2), c_hidden[0], kernel_size=1), - WuerstchenLayerNorm(c_hidden[0], elementwise_affine=False, eps=1e-6), - ) - - def get_block(block_type, c_hidden, nhead, c_skip=0, dropout=0): - if block_type == "C": - return ResBlockStageB(c_hidden, c_skip, kernel_size=kernel_size, dropout=dropout) - elif block_type == "A": - return AttnBlock(c_hidden, c_cond, nhead, self_attn=True, dropout=dropout) - elif block_type == "T": - return TimestepBlock(c_hidden, c_r) - else: - raise ValueError(f"Block type {block_type} not supported") - - # BLOCKS - # -- down blocks - self.down_blocks = nn.ModuleList() - for i in range(len(c_hidden)): - down_block = nn.ModuleList() - if i > 0: - down_block.append( - nn.Sequential( - WuerstchenLayerNorm(c_hidden[i - 1], elementwise_affine=False, eps=1e-6), - nn.Conv2d(c_hidden[i - 1], c_hidden[i], kernel_size=2, stride=2), - ) - ) - for _ in range(blocks[i]): - for block_type in level_config[i]: - c_skip = c_cond if inject_effnet[i] else 0 - down_block.append(get_block(block_type, c_hidden[i], nhead[i], c_skip=c_skip, dropout=dropout[i])) - self.down_blocks.append(down_block) - - # -- up blocks - self.up_blocks = nn.ModuleList() - for i in reversed(range(len(c_hidden))): - up_block = nn.ModuleList() - for j in range(blocks[i]): - for k, block_type in enumerate(level_config[i]): - c_skip = c_hidden[i] if i < len(c_hidden) - 1 and j == k == 0 else 0 - c_skip += c_cond if inject_effnet[i] else 0 - up_block.append(get_block(block_type, c_hidden[i], nhead[i], c_skip=c_skip, dropout=dropout[i])) - if i > 0: - up_block.append( - nn.Sequential( - WuerstchenLayerNorm(c_hidden[i], elementwise_affine=False, eps=1e-6), - nn.ConvTranspose2d(c_hidden[i], c_hidden[i - 1], kernel_size=2, stride=2), - ) - ) - self.up_blocks.append(up_block) - - # OUTPUT - self.clf = nn.Sequential( - WuerstchenLayerNorm(c_hidden[0], elementwise_affine=False, eps=1e-6), - nn.Conv2d(c_hidden[0], 2 * c_out * (patch_size**2), kernel_size=1), - nn.PixelShuffle(patch_size), - ) - - # --- WEIGHT INIT --- - self.apply(self._init_weights) - - def _init_weights(self, m): - # General init - if isinstance(m, (nn.Conv2d, nn.Linear)): - nn.init.xavier_uniform_(m.weight) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - - for mapper in self.effnet_mappers: - if mapper is not None: - nn.init.normal_(mapper.weight, std=0.02) # conditionings - nn.init.normal_(self.clip_mapper.weight, std=0.02) # conditionings - nn.init.xavier_uniform_(self.embedding[1].weight, 0.02) # inputs - nn.init.constant_(self.clf[1].weight, 0) # outputs - - # blocks - for level_block in self.down_blocks + self.up_blocks: - for block in level_block: - if isinstance(block, ResBlockStageB): - block.channelwise[-1].weight.data *= np.sqrt(1 / sum(self.config.blocks)) - elif isinstance(block, TimestepBlock): - nn.init.constant_(block.mapper.weight, 0) - - def gen_r_embedding(self, r, max_positions=10000): - r = r * max_positions - half_dim = self.c_r // 2 - emb = math.log(max_positions) / (half_dim - 1) - emb = torch.arange(half_dim, device=r.device).float().mul(-emb).exp() - emb = r[:, None] * emb[None, :] - emb = torch.cat([emb.sin(), emb.cos()], dim=1) - if self.c_r % 2 == 1: # zero pad - emb = nn.functional.pad(emb, (0, 1), mode="constant") - return emb.to(dtype=r.dtype) - - def gen_c_embeddings(self, clip): - clip = self.clip_mapper(clip) - clip = self.seq_norm(clip) - return clip - - def _down_encode(self, x, r_embed, effnet, clip=None): - level_outputs = [] - for i, down_block in enumerate(self.down_blocks): - effnet_c = None - for block in down_block: - if isinstance(block, ResBlockStageB): - if effnet_c is None and self.effnet_mappers[i] is not None: - dtype = effnet.dtype - effnet_c = self.effnet_mappers[i]( - nn.functional.interpolate( - effnet.float(), size=x.shape[-2:], mode="bicubic", antialias=True, align_corners=True - ).to(dtype) - ) - skip = effnet_c if self.effnet_mappers[i] is not None else None - x = block(x, skip) - elif isinstance(block, AttnBlock): - x = block(x, clip) - elif isinstance(block, TimestepBlock): - x = block(x, r_embed) - else: - x = block(x) - level_outputs.insert(0, x) - return level_outputs - - def _up_decode(self, level_outputs, r_embed, effnet, clip=None): - x = level_outputs[0] - for i, up_block in enumerate(self.up_blocks): - effnet_c = None - for j, block in enumerate(up_block): - if isinstance(block, ResBlockStageB): - if effnet_c is None and self.effnet_mappers[len(self.down_blocks) + i] is not None: - dtype = effnet.dtype - effnet_c = self.effnet_mappers[len(self.down_blocks) + i]( - nn.functional.interpolate( - effnet.float(), size=x.shape[-2:], mode="bicubic", antialias=True, align_corners=True - ).to(dtype) - ) - skip = level_outputs[i] if j == 0 and i > 0 else None - if effnet_c is not None: - if skip is not None: - skip = torch.cat([skip, effnet_c], dim=1) - else: - skip = effnet_c - x = block(x, skip) - elif isinstance(block, AttnBlock): - x = block(x, clip) - elif isinstance(block, TimestepBlock): - x = block(x, r_embed) - else: - x = block(x) - return x - - def forward(self, x, r, effnet, clip=None, x_cat=None, eps=1e-3, return_noise=True): - if x_cat is not None: - x = torch.cat([x, x_cat], dim=1) - # Process the conditioning embeddings - r_embed = self.gen_r_embedding(r) - if clip is not None: - clip = self.gen_c_embeddings(clip) - - # Model Blocks - x_in = x - x = self.embedding(x) - level_outputs = self._down_encode(x, r_embed, effnet, clip) - x = self._up_decode(level_outputs, r_embed, effnet, clip) - a, b = self.clf(x).chunk(2, dim=1) - b = b.sigmoid() * (1 - eps * 2) + eps - if return_noise: - return (x_in - a) / b - else: - return a, b - - -class ResBlockStageB(nn.Module): - def __init__(self, c, c_skip=0, kernel_size=3, dropout=0.0): - super().__init__() - self.depthwise = nn.Conv2d(c, c, kernel_size=kernel_size, padding=kernel_size // 2, groups=c) - self.norm = WuerstchenLayerNorm(c, elementwise_affine=False, eps=1e-6) - self.channelwise = nn.Sequential( - nn.Linear(c + c_skip, c * 4), - nn.GELU(), - GlobalResponseNorm(c * 4), - nn.Dropout(dropout), - nn.Linear(c * 4, c), - ) - - def forward(self, x, x_skip=None): - x_res = x - x = self.norm(self.depthwise(x)) - if x_skip is not None: - x = torch.cat([x, x_skip], dim=1) - x = self.channelwise(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) - return x + x_res diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py deleted file mode 100644 index ae821a55788aeddf49d8a7299c6c82da467f2832..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2023 Dominic Rampas MIT License -# 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 math -from typing import Dict, Union - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin -from ...models.attention_processor import ( - ADDED_KV_ATTENTION_PROCESSORS, - CROSS_ATTENTION_PROCESSORS, - AttentionProcessor, - AttnAddedKVProcessor, - AttnProcessor, -) -from ...models.modeling_utils import ModelMixin -from .modeling_wuerstchen_common import AttnBlock, ResBlock, TimestepBlock, WuerstchenLayerNorm - - -class WuerstchenPrior(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin): - unet_name = "prior" - _supports_gradient_checkpointing = True - - @register_to_config - def __init__(self, c_in=16, c=1280, c_cond=1024, c_r=64, depth=16, nhead=16, dropout=0.1): - super().__init__() - - self.c_r = c_r - self.projection = nn.Conv2d(c_in, c, kernel_size=1) - self.cond_mapper = nn.Sequential( - nn.Linear(c_cond, c), - nn.LeakyReLU(0.2), - nn.Linear(c, c), - ) - - self.blocks = nn.ModuleList() - for _ in range(depth): - self.blocks.append(ResBlock(c, dropout=dropout)) - self.blocks.append(TimestepBlock(c, c_r)) - self.blocks.append(AttnBlock(c, c, nhead, self_attn=True, dropout=dropout)) - self.out = nn.Sequential( - WuerstchenLayerNorm(c, elementwise_affine=False, eps=1e-6), - nn.Conv2d(c, c_in * 2, kernel_size=1), - ) - - self.gradient_checkpointing = False - self.set_default_attn_processor() - - @property - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors - def attn_processors(self) -> Dict[str, AttentionProcessor]: - r""" - Returns: - `dict` of attention processors: A dictionary containing all attention processors used in the model with - indexed by its weight name. - """ - # set recursively - processors = {} - - 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() - - for sub_name, child in module.named_children(): - fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) - - return processors - - for name, module in self.named_children(): - fn_recursive_add_processors(name, module, processors) - - return processors - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): - r""" - Sets the attention processor to use to compute attention. - - Parameters: - processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): - The instantiated processor class or a dictionary of processor classes that will be set as the processor - for **all** `Attention` layers. - - If `processor` is a dict, the key needs to define the path to the corresponding cross attention - processor. This is strongly recommended when setting trainable attention processors. - - """ - count = len(self.attn_processors.keys()) - - if isinstance(processor, dict) and len(processor) != count: - raise ValueError( - f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" - f" number of attention layers: {count}. Please make sure to pass {count} processor classes." - ) - - def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): - if hasattr(module, "set_processor"): - if not isinstance(processor, dict): - module.set_processor(processor) - else: - module.set_processor(processor.pop(f"{name}.processor")) - - for sub_name, child in module.named_children(): - fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) - - for name, module in self.named_children(): - fn_recursive_attn_processor(name, module, processor) - - # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor - def set_default_attn_processor(self): - """ - Disables custom attention processors and sets the default attention implementation. - """ - if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnAddedKVProcessor() - elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): - processor = AttnProcessor() - else: - raise ValueError( - f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" - ) - - self.set_attn_processor(processor) - - def gen_r_embedding(self, r, max_positions=10000): - r = r * max_positions - half_dim = self.c_r // 2 - emb = math.log(max_positions) / (half_dim - 1) - emb = torch.arange(half_dim, device=r.device).float().mul(-emb).exp() - emb = r[:, None] * emb[None, :] - emb = torch.cat([emb.sin(), emb.cos()], dim=1) - if self.c_r % 2 == 1: # zero pad - emb = nn.functional.pad(emb, (0, 1), mode="constant") - return emb.to(dtype=r.dtype) - - def forward(self, x, r, c): - x_in = x - x = self.projection(x) - c_embed = self.cond_mapper(c) - r_embed = self.gen_r_embedding(r) - - if torch.is_grad_enabled() and self.gradient_checkpointing: - for block in self.blocks: - if isinstance(block, AttnBlock): - x = self._gradient_checkpointing_func(block, x, c_embed) - elif isinstance(block, TimestepBlock): - x = self._gradient_checkpointing_func(block, x, r_embed) - else: - x = self._gradient_checkpointing_func(block, x) - else: - for block in self.blocks: - if isinstance(block, AttnBlock): - x = block(x, c_embed) - elif isinstance(block, TimestepBlock): - x = block(x, r_embed) - else: - x = block(x) - a, b = self.out(x).chunk(2, dim=1) - return (x_in - a) / ((1 - b).abs() + 1e-5) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py deleted file mode 100644 index b9b02a6dd38a0e1fc28c1136f0419acb4aad6dd5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py +++ /dev/null @@ -1,449 +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 Callable, Dict, List, Optional, Union - -import numpy as np -import torch -from transformers import CLIPTextModel, CLIPTokenizer - -from ...schedulers import DDPMWuerstchenScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, ImagePipelineOutput -from .modeling_paella_vq_model import PaellaVQModel -from .modeling_wuerstchen_diffnext import WuerstchenDiffNeXt - - -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 WuerstchenPriorPipeline, WuerstchenDecoderPipeline - - >>> prior_pipe = WuerstchenPriorPipeline.from_pretrained( - ... "warp-ai/wuerstchen-prior", torch_dtype=torch.float16 - ... ).to("cuda") - >>> gen_pipe = WuerstchenDecoderPipeline.from_pretrain("warp-ai/wuerstchen", torch_dtype=torch.float16).to( - ... "cuda" - ... ) - - >>> prompt = "an image of a shiba inu, donning a spacesuit and helmet" - >>> prior_output = pipe(prompt) - >>> images = gen_pipe(prior_output.image_embeddings, prompt=prompt) - ``` -""" - - -class WuerstchenDecoderPipeline(DeprecatedPipelineMixin, DiffusionPipeline): - """ - Pipeline for generating images from the Wuerstchen model. - - 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.) - - Args: - tokenizer (`CLIPTokenizer`): - The CLIP tokenizer. - text_encoder (`CLIPTextModel`): - The CLIP text encoder. - decoder ([`WuerstchenDiffNeXt`]): - The WuerstchenDiffNeXt unet decoder. - vqgan ([`PaellaVQModel`]): - The VQGAN model. - scheduler ([`DDPMWuerstchenScheduler`]): - A scheduler to be used in combination with `prior` to generate image embedding. - latent_dim_scale (float, `optional`, defaults to 10.67): - Multiplier to determine the VQ latent space size from the image embeddings. If the image embeddings are - height=24 and width=24, the VQ latent shape needs to be height=int(24*10.67)=256 and - width=int(24*10.67)=256 in order to match the training conditions. - """ - - model_cpu_offload_seq = "text_encoder->decoder->vqgan" - _callback_tensor_inputs = [ - "latents", - "text_encoder_hidden_states", - "negative_prompt_embeds", - "image_embeddings", - ] - - def __init__( - self, - tokenizer: CLIPTokenizer, - text_encoder: CLIPTextModel, - decoder: WuerstchenDiffNeXt, - scheduler: DDPMWuerstchenScheduler, - vqgan: PaellaVQModel, - latent_dim_scale: float = 10.67, - ) -> None: - super().__init__() - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - decoder=decoder, - scheduler=scheduler, - vqgan=vqgan, - ) - self.register_to_config(latent_dim_scale=latent_dim_scale) - - # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - def encode_prompt( - self, - prompt, - device, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt=None, - ): - batch_size = len(prompt) if isinstance(prompt, list) else 1 - # get prompt text embeddings - 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 - attention_mask = text_inputs.attention_mask - - 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}" - ) - text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length] - attention_mask = attention_mask[:, : self.tokenizer.model_max_length] - - text_encoder_output = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask.to(device)) - text_encoder_hidden_states = text_encoder_output.last_hidden_state - text_encoder_hidden_states = text_encoder_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - - uncond_text_encoder_hidden_states = None - if do_classifier_free_guidance: - uncond_tokens: List[str] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif 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 - - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - negative_prompt_embeds_text_encoder_output = self.text_encoder( - uncond_input.input_ids.to(device), attention_mask=uncond_input.attention_mask.to(device) - ) - - uncond_text_encoder_hidden_states = negative_prompt_embeds_text_encoder_output.last_hidden_state - - # duplicate unconditional embeddings for each generation per prompt, using mps friendly method - seq_len = uncond_text_encoder_hidden_states.shape[1] - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.repeat(1, num_images_per_prompt, 1) - uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.view( - batch_size * num_images_per_prompt, seq_len, -1 - ) - # done duplicates - - # 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 - return text_encoder_hidden_states, uncond_text_encoder_hidden_states - - @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 - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - image_embeddings: Union[torch.Tensor, List[torch.Tensor]], - prompt: Union[str, List[str]] = None, - num_inference_steps: int = 12, - timesteps: Optional[List[float]] = None, - guidance_scale: float = 0.0, - negative_prompt: Optional[Union[str, List[str]]] = 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", - return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], 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 Embeddings either extracted from an image or generated by a Prior Model. - 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*): - 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): - Guidance scale as defined in [Classifier-Free Diffusion - Guidance](https://huggingface.co/papers/2207.12598). `decoder_guidance_scale` is defined as `w` of - 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*): - 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*): - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - 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. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` [`~pipelines.ImagePipelineOutput`] if `return_dict` is True, - otherwise a `tuple`. When returning a tuple, the first element is a list with the generated image - embeddings. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - 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]}" - ) - - # 0. Define commonly used variables - device = self._execution_device - dtype = self.decoder.dtype - self._guidance_scale = guidance_scale - - # 1. Check inputs. Raise error if not correct - if not isinstance(prompt, list): - if isinstance(prompt, str): - prompt = [prompt] - else: - raise TypeError(f"'prompt' must be of type 'list' or 'str', but got {type(prompt)}.") - - if self.do_classifier_free_guidance: - if negative_prompt is not None and not isinstance(negative_prompt, list): - if isinstance(negative_prompt, str): - negative_prompt = [negative_prompt] - else: - raise TypeError( - f"'negative_prompt' must be of type 'list' or 'str', but got {type(negative_prompt)}." - ) - - if isinstance(image_embeddings, list): - image_embeddings = torch.cat(image_embeddings, dim=0) - if isinstance(image_embeddings, np.ndarray): - image_embeddings = torch.Tensor(image_embeddings, device=device).to(dtype=dtype) - if not isinstance(image_embeddings, torch.Tensor): - raise TypeError( - f"'image_embeddings' must be of type 'torch.Tensor' or 'np.array', but got {type(image_embeddings)}." - ) - - if not isinstance(num_inference_steps, int): - raise TypeError( - f"'num_inference_steps' must be of type 'int', but got {type(num_inference_steps)}\ - In Case you want to provide explicit timesteps, please use the 'timesteps' argument." - ) - - # 2. Encode caption - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - image_embeddings.size(0) * num_images_per_prompt, - self.do_classifier_free_guidance, - negative_prompt, - ) - text_encoder_hidden_states = ( - torch.cat([prompt_embeds, negative_prompt_embeds]) if negative_prompt_embeds is not None else prompt_embeds - ) - effnet = ( - torch.cat([image_embeddings, torch.zeros_like(image_embeddings)]) - if self.do_classifier_free_guidance - else image_embeddings - ) - - # 3. Determine latent shape of latents - latent_height = int(image_embeddings.size(2) * self.config.latent_dim_scale) - latent_width = int(image_embeddings.size(3) * self.config.latent_dim_scale) - latent_features_shape = (image_embeddings.size(0) * num_images_per_prompt, 4, latent_height, latent_width) - - # 4. Prepare and set timesteps - if timesteps is not None: - self.scheduler.set_timesteps(timesteps=timesteps, device=device) - timesteps = self.scheduler.timesteps - num_inference_steps = len(timesteps) - else: - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latents - latents = self.prepare_latents(latent_features_shape, dtype, device, generator, latents, self.scheduler) - - # 6. Run denoising loop - self._num_timesteps = len(timesteps[:-1]) - for i, t in enumerate(self.progress_bar(timesteps[:-1])): - ratio = t.expand(latents.size(0)).to(dtype) - # 7. Denoise latents - predicted_latents = self.decoder( - torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents, - r=torch.cat([ratio] * 2) if self.do_classifier_free_guidance else ratio, - effnet=effnet, - clip=text_encoder_hidden_states, - ) - - # 8. Check for classifier free guidance and apply it - if self.do_classifier_free_guidance: - predicted_latents_text, predicted_latents_uncond = predicted_latents.chunk(2) - predicted_latents = torch.lerp(predicted_latents_uncond, predicted_latents_text, self.guidance_scale) - - # 9. Renoise latents to next timestep - latents = self.scheduler.step( - model_output=predicted_latents, - timestep=ratio, - sample=latents, - generator=generator, - ).prev_sample - - 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) - image_embeddings = callback_outputs.pop("image_embeddings", image_embeddings) - text_encoder_hidden_states = callback_outputs.pop( - "text_encoder_hidden_states", text_encoder_hidden_states - ) - - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if output_type not in ["pt", "np", "pil", "latent"]: - raise ValueError( - f"Only the output types `pt`, `np`, `pil` and `latent` are supported not output_type={output_type}" - ) - - if not output_type == "latent": - # 10. Scale and decode the image latents with vq-vae - latents = self.vqgan.config.scale_factor * latents - images = self.vqgan.decode(latents).sample.clamp(0, 1) - if output_type == "np": - images = images.permute(0, 2, 3, 1).cpu().float().numpy() - elif output_type == "pil": - images = images.permute(0, 2, 3, 1).cpu().float().numpy() - images = self.numpy_to_pil(images) - else: - images = latents - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return images - return ImagePipelineOutput(images) diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py deleted file mode 100644 index 00a88ce34ed2a5a64591624c54b0df425e3179cd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py +++ /dev/null @@ -1,307 +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 Callable, Dict, List, Optional, Union - -import torch -from transformers import CLIPTextModel, CLIPTokenizer - -from ...schedulers import DDPMWuerstchenScheduler -from ...utils import deprecate, replace_example_docstring -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline -from .modeling_paella_vq_model import PaellaVQModel -from .modeling_wuerstchen_diffnext import WuerstchenDiffNeXt -from .modeling_wuerstchen_prior import WuerstchenPrior -from .pipeline_wuerstchen import WuerstchenDecoderPipeline -from .pipeline_wuerstchen_prior import WuerstchenPriorPipeline - - -TEXT2IMAGE_EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> from diffusions import WuerstchenCombinedPipeline - - >>> pipe = WuerstchenCombinedPipeline.from_pretrained("warp-ai/Wuerstchen", torch_dtype=torch.float16).to( - ... "cuda" - ... ) - >>> prompt = "an image of a shiba inu, donning a spacesuit and helmet" - >>> images = pipe(prompt=prompt) - ``` -""" - - -class WuerstchenCombinedPipeline(DeprecatedPipelineMixin, DiffusionPipeline): - """ - Combined Pipeline for text-to-image generation using Wuerstchen - - 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.) - - Args: - tokenizer (`CLIPTokenizer`): - The decoder tokenizer to be used for text inputs. - text_encoder (`CLIPTextModel`): - The decoder text encoder to be used for text inputs. - decoder (`WuerstchenDiffNeXt`): - The decoder model to be used for decoder image generation pipeline. - scheduler (`DDPMWuerstchenScheduler`): - The scheduler to be used for decoder image generation pipeline. - vqgan (`PaellaVQModel`): - The VQGAN model to be used for decoder image generation pipeline. - prior_tokenizer (`CLIPTokenizer`): - The prior tokenizer to be used for text inputs. - prior_text_encoder (`CLIPTextModel`): - The prior text encoder to be used for text inputs. - prior_prior (`WuerstchenPrior`): - The prior model to be used for prior pipeline. - prior_scheduler (`DDPMWuerstchenScheduler`): - The scheduler to be used for prior pipeline. - """ - - _last_supported_version = "0.33.1" - _load_connected_pipes = True - - def __init__( - self, - tokenizer: CLIPTokenizer, - text_encoder: CLIPTextModel, - decoder: WuerstchenDiffNeXt, - scheduler: DDPMWuerstchenScheduler, - vqgan: PaellaVQModel, - prior_tokenizer: CLIPTokenizer, - prior_text_encoder: CLIPTextModel, - prior_prior: WuerstchenPrior, - prior_scheduler: DDPMWuerstchenScheduler, - ): - super().__init__() - - self.register_modules( - text_encoder=text_encoder, - tokenizer=tokenizer, - decoder=decoder, - scheduler=scheduler, - vqgan=vqgan, - prior_prior=prior_prior, - prior_text_encoder=prior_text_encoder, - prior_tokenizer=prior_tokenizer, - prior_scheduler=prior_scheduler, - ) - self.prior_pipe = WuerstchenPriorPipeline( - prior=prior_prior, - text_encoder=prior_text_encoder, - tokenizer=prior_tokenizer, - scheduler=prior_scheduler, - ) - self.decoder_pipe = WuerstchenDecoderPipeline( - text_encoder=text_encoder, - tokenizer=tokenizer, - decoder=decoder, - scheduler=scheduler, - vqgan=vqgan, - ) - - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = 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): - 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` - method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with - `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`. - """ - 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): - 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 - GPU only when their specific submodule's `forward` method is called. Offloading happens on a submodule basis. - Memory savings are higher than using `enable_model_cpu_offload`, but performance is lower. - """ - self.prior_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - self.decoder_pipe.enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) - - def progress_bar(self, iterable=None, total=None): - self.prior_pipe.progress_bar(iterable=iterable, total=total) - self.decoder_pipe.progress_bar(iterable=iterable, total=total) - - def set_progress_bar_config(self, **kwargs): - self.prior_pipe.set_progress_bar_config(**kwargs) - self.decoder_pipe.set_progress_bar_config(**kwargs) - - @torch.no_grad() - @replace_example_docstring(TEXT2IMAGE_EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Optional[Union[str, List[str]]] = None, - height: int = 512, - width: int = 512, - prior_num_inference_steps: int = 60, - prior_timesteps: Optional[List[float]] = None, - prior_guidance_scale: float = 4.0, - num_inference_steps: int = 12, - decoder_timesteps: Optional[List[float]] = 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, - 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", - 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"], - **kwargs, - ): - """ - Function invoked when calling the pipeline for generation. - - Args: - 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*): - 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*): - Pre-generated text embeddings for the prior. 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 for the prior. 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. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 512): - The width in pixels of the generated image. - prior_guidance_scale (`float`, *optional*, defaults to 4.0): - Guidance scale as defined in [Classifier-Free Diffusion - Guidance](https://huggingface.co/papers/2207.12598). `prior_guidance_scale` is defined as `w` of - 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): - 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` - num_inference_steps (`int`, *optional*, defaults to 12): - 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*): - 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*): - 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): - 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. - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - prior_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: `prior_callback_on_step_end(self: DiffusionPipeline, step: int, timestep: - int, callback_kwargs: Dict)`. - 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. - 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. - - Examples: - - Returns: - [`~pipelines.ImagePipelineOutput`] or `tuple` [`~pipelines.ImagePipelineOutput`] if `return_dict` is True, - otherwise a `tuple`. When returning a tuple, the first element is a list with the generated images. - """ - prior_kwargs = {} - if kwargs.get("prior_callback", None) is not None: - prior_kwargs["callback"] = kwargs.pop("prior_callback") - deprecate( - "prior_callback", - "1.0.0", - "Passing `prior_callback` as an input argument to `__call__` is deprecated, consider use `prior_callback_on_step_end`", - ) - if kwargs.get("prior_callback_steps", None) is not None: - deprecate( - "prior_callback_steps", - "1.0.0", - "Passing `prior_callback_steps` as an input argument to `__call__` is deprecated, consider use `prior_callback_on_step_end`", - ) - prior_kwargs["callback_steps"] = kwargs.pop("prior_callback_steps") - - prior_outputs = self.prior_pipe( - prompt=prompt if prompt_embeds is None else None, - height=height, - width=width, - num_inference_steps=prior_num_inference_steps, - timesteps=prior_timesteps, - guidance_scale=prior_guidance_scale, - negative_prompt=negative_prompt if negative_prompt_embeds is None else None, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - num_images_per_prompt=num_images_per_prompt, - generator=generator, - latents=latents, - output_type="pt", - return_dict=False, - callback_on_step_end=prior_callback_on_step_end, - callback_on_step_end_tensor_inputs=prior_callback_on_step_end_tensor_inputs, - **prior_kwargs, - ) - image_embeddings = prior_outputs[0] - - outputs = self.decoder_pipe( - image_embeddings=image_embeddings, - prompt=prompt if prompt is not None else "", - num_inference_steps=num_inference_steps, - timesteps=decoder_timesteps, - guidance_scale=decoder_guidance_scale, - negative_prompt=negative_prompt, - generator=generator, - output_type=output_type, - return_dict=return_dict, - callback_on_step_end=callback_on_step_end, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - **kwargs, - ) - - return outputs diff --git a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py b/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py deleted file mode 100644 index a32f09204d27a3db16bc2c6c4fee7af701c4536f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py +++ /dev/null @@ -1,528 +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 dataclasses import dataclass -from math import ceil -from typing import Callable, Dict, List, Optional, Union - -import numpy as np -import torch -from transformers import CLIPTextModel, CLIPTokenizer - -from ...loaders import StableDiffusionLoraLoaderMixin -from ...schedulers import DDPMWuerstchenScheduler -from ...utils import BaseOutput, deprecate, is_torch_xla_available, logging, replace_example_docstring -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .modeling_wuerstchen_prior import WuerstchenPrior - - -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 - - -DEFAULT_STAGE_C_TIMESTEPS = list(np.linspace(1.0, 2 / 3, 20)) + list(np.linspace(2 / 3, 0.0, 11))[1:] - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import WuerstchenPriorPipeline - - >>> prior_pipe = WuerstchenPriorPipeline.from_pretrained( - ... "warp-ai/wuerstchen-prior", torch_dtype=torch.float16 - ... ).to("cuda") - - >>> prompt = "an image of a shiba inu, donning a spacesuit and helmet" - >>> prior_output = pipe(prompt) - ``` -""" - - -@dataclass -class WuerstchenPriorPipelineOutput(BaseOutput): - """ - Output class for WuerstchenPriorPipeline. - - Args: - image_embeddings (`torch.Tensor` or `np.ndarray`) - Prior image embeddings for text prompt - - """ - - image_embeddings: Union[torch.Tensor, np.ndarray] - - -class WuerstchenPriorPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin): - """ - Pipeline for generating image prior for Wuerstchen. - - 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.StableDiffusionLoraLoaderMixin.load_lora_weights`] for loading LoRA weights - - [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for saving LoRA weights - - Args: - prior ([`Prior`]): - The canonical unCLIP prior to approximate the image embedding from the text embedding. - text_encoder ([`CLIPTextModelWithProjection`]): - Frozen text-encoder. - tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - scheduler ([`DDPMWuerstchenScheduler`]): - A scheduler to be used in combination with `prior` to generate image embedding. - latent_mean ('float', *optional*, defaults to 42.0): - Mean value for latent diffusers. - latent_std ('float', *optional*, defaults to 1.0): - Standard value for latent diffusers. - resolution_multiple ('float', *optional*, defaults to 42.67): - Default resolution for multiple images generated. - """ - - unet_name = "prior" - text_encoder_name = "text_encoder" - model_cpu_offload_seq = "text_encoder->prior" - _callback_tensor_inputs = ["latents", "text_encoder_hidden_states", "negative_prompt_embeds"] - _lora_loadable_modules = ["prior", "text_encoder"] - - def __init__( - self, - tokenizer: CLIPTokenizer, - text_encoder: CLIPTextModel, - prior: WuerstchenPrior, - scheduler: DDPMWuerstchenScheduler, - latent_mean: float = 42.0, - latent_std: float = 1.0, - resolution_multiple: float = 42.67, - ) -> None: - super().__init__() - self.register_modules( - tokenizer=tokenizer, - text_encoder=text_encoder, - prior=prior, - scheduler=scheduler, - ) - self.register_to_config( - latent_mean=latent_mean, latent_std=latent_std, resolution_multiple=resolution_multiple - ) - - # Copied from diffusers.pipelines.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents - def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): - 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) - - latents = latents * scheduler.init_noise_sigma - return latents - - def encode_prompt( - self, - device, - num_images_per_prompt, - do_classifier_free_guidance, - prompt=None, - negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ): - 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: - # get prompt text embeddings - 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 - attention_mask = text_inputs.attention_mask - - 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}" - ) - text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length] - attention_mask = attention_mask[:, : self.tokenizer.model_max_length] - - text_encoder_output = self.text_encoder( - text_input_ids.to(device), attention_mask=attention_mask.to(device) - ) - prompt_embeds = text_encoder_output.last_hidden_state - - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device) - 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] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif 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 - - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - negative_prompt_embeds_text_encoder_output = self.text_encoder( - uncond_input.input_ids.to(device), attention_mask=uncond_input.attention_mask.to(device) - ) - - negative_prompt_embeds = negative_prompt_embeds_text_encoder_output.last_hidden_state - - 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=self.text_encoder.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) - # done duplicates - - return prompt_embeds, negative_prompt_embeds - - def check_inputs( - self, - prompt, - negative_prompt, - num_inference_steps, - do_classifier_free_guidance, - prompt_embeds=None, - negative_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 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 not isinstance(num_inference_steps, int): - raise TypeError( - f"'num_inference_steps' must be of type 'int', but got {type(num_inference_steps)}\ - In Case you want to provide explicit timesteps, please use the 'timesteps' argument." - ) - - @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 - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Optional[Union[str, List[str]]] = None, - height: int = 1024, - width: int = 1024, - num_inference_steps: int = 60, - 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", - return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], 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]`): - The prompt or prompts to guide the image generation. - height (`int`, *optional*, defaults to 1024): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 1024): - The width in pixels of the generated image. - 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*): - 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): - Guidance scale as defined in [Classifier-Free Diffusion - Guidance](https://huggingface.co/papers/2207.12598). `decoder_guidance_scale` is defined as `w` of - 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*): - 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*): - 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. - 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 ge generated by sampling using the supplied random `generator`. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between: `"pil"` (`PIL.Image.Image`), `"np"` - (`np.array`) or `"pt"` (`torch.Tensor`). - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. - 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. - - Examples: - - Returns: - [`~pipelines.WuerstchenPriorPipelineOutput`] or `tuple` [`~pipelines.WuerstchenPriorPipelineOutput`] if - `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the - generated image embeddings. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - 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]}" - ) - - # 0. Define commonly used variables - device = self._execution_device - self._guidance_scale = guidance_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] - - # 1. Check inputs. Raise error if not correct - if prompt is not None and not isinstance(prompt, list): - if isinstance(prompt, str): - prompt = [prompt] - else: - raise TypeError(f"'prompt' must be of type 'list' or 'str', but got {type(prompt)}.") - - if self.do_classifier_free_guidance: - if negative_prompt is not None and not isinstance(negative_prompt, list): - if isinstance(negative_prompt, str): - negative_prompt = [negative_prompt] - else: - raise TypeError( - f"'negative_prompt' must be of type 'list' or 'str', but got {type(negative_prompt)}." - ) - - self.check_inputs( - prompt, - negative_prompt, - num_inference_steps, - self.do_classifier_free_guidance, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - - # 2. Encode caption - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - ) - - # 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 - text_encoder_hidden_states = ( - torch.cat([prompt_embeds, negative_prompt_embeds]) if negative_prompt_embeds is not None else prompt_embeds - ) - - # 3. Determine latent shape of image embeddings - dtype = text_encoder_hidden_states.dtype - latent_height = ceil(height / self.config.resolution_multiple) - latent_width = ceil(width / self.config.resolution_multiple) - num_channels = self.prior.config.c_in - effnet_features_shape = (num_images_per_prompt * batch_size, num_channels, latent_height, latent_width) - - # 4. Prepare and set timesteps - if timesteps is not None: - self.scheduler.set_timesteps(timesteps=timesteps, device=device) - timesteps = self.scheduler.timesteps - num_inference_steps = len(timesteps) - else: - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 5. Prepare latents - latents = self.prepare_latents(effnet_features_shape, dtype, device, generator, latents, self.scheduler) - - # 6. Run denoising loop - self._num_timesteps = len(timesteps[:-1]) - for i, t in enumerate(self.progress_bar(timesteps[:-1])): - ratio = t.expand(latents.size(0)).to(dtype) - - # 7. Denoise image embeddings - predicted_image_embedding = self.prior( - torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents, - r=torch.cat([ratio] * 2) if self.do_classifier_free_guidance else ratio, - c=text_encoder_hidden_states, - ) - - # 8. Check for classifier free guidance and apply it - if self.do_classifier_free_guidance: - predicted_image_embedding_text, predicted_image_embedding_uncond = predicted_image_embedding.chunk(2) - predicted_image_embedding = torch.lerp( - predicted_image_embedding_uncond, predicted_image_embedding_text, self.guidance_scale - ) - - # 9. Renoise latents to next timestep - latents = self.scheduler.step( - model_output=predicted_image_embedding, - timestep=ratio, - sample=latents, - generator=generator, - ).prev_sample - - 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) - text_encoder_hidden_states = callback_outputs.pop( - "text_encoder_hidden_states", text_encoder_hidden_states - ) - negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) - - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - # 10. Denormalize the latents - latents = latents * self.config.latent_mean - self.config.latent_std - - # Offload all models - self.maybe_free_model_hooks() - - if output_type == "np": - latents = latents.cpu().float().numpy() - - if not return_dict: - return (latents,) - - return WuerstchenPriorPipelineOutput(latents) diff --git a/venv/lib/python3.10/site-packages/diffusers/py.typed b/venv/lib/python3.10/site-packages/diffusers/py.typed deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/__init__.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/__init__.py deleted file mode 100644 index efd24187532179892108466c386154e1a54e3bde..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/__init__.py +++ /dev/null @@ -1,194 +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 inspect -from typing import Dict, List, Optional, Union - -from ..utils import is_transformers_available, logging -from .auto import DiffusersAutoQuantizer -from .base import DiffusersQuantizer -from .quantization_config import QuantizationConfigMixin as DiffQuantConfigMixin - - -try: - from transformers.utils.quantization_config import QuantizationConfigMixin as TransformersQuantConfigMixin -except ImportError: - - class TransformersQuantConfigMixin: - pass - - -logger = logging.get_logger(__name__) - - -class PipelineQuantizationConfig: - """ - Configuration class to be used when applying quantization on-the-fly to [`~DiffusionPipeline.from_pretrained`]. - - Args: - quant_backend (`str`): Quantization backend to be used. When using this option, we assume that the backend - is available to both `diffusers` and `transformers`. - quant_kwargs (`dict`): Params to initialize the quantization backend class. - components_to_quantize (`list`): Components of a pipeline to be quantized. - quant_mapping (`dict`): Mapping defining the quantization specs to be used for the pipeline - components. When using this argument, users are not expected to provide `quant_backend`, `quant_kawargs`, - and `components_to_quantize`. - """ - - def __init__( - self, - quant_backend: str = None, - quant_kwargs: Dict[str, Union[str, float, int, dict]] = None, - components_to_quantize: Optional[List[str]] = None, - quant_mapping: Dict[str, Union[DiffQuantConfigMixin, "TransformersQuantConfigMixin"]] = None, - ): - self.quant_backend = quant_backend - # Initialize kwargs to be {} to set to the defaults. - self.quant_kwargs = quant_kwargs or {} - self.components_to_quantize = components_to_quantize - self.quant_mapping = quant_mapping - - self.post_init() - - def post_init(self): - quant_mapping = self.quant_mapping - self.is_granular = True if quant_mapping is not None else False - - self._validate_init_args() - - def _validate_init_args(self): - if self.quant_backend and self.quant_mapping: - raise ValueError("Both `quant_backend` and `quant_mapping` cannot be specified at the same time.") - - if not self.quant_mapping and not self.quant_backend: - raise ValueError("Must provide a `quant_backend` when not providing a `quant_mapping`.") - - if not self.quant_kwargs and not self.quant_mapping: - raise ValueError("Both `quant_kwargs` and `quant_mapping` cannot be None.") - - if self.quant_backend is not None: - self._validate_init_kwargs_in_backends() - - if self.quant_mapping is not None: - self._validate_quant_mapping_args() - - def _validate_init_kwargs_in_backends(self): - quant_backend = self.quant_backend - - self._check_backend_availability(quant_backend) - - quant_config_mapping_transformers, quant_config_mapping_diffusers = self._get_quant_config_list() - - if quant_config_mapping_transformers is not None: - init_kwargs_transformers = inspect.signature(quant_config_mapping_transformers[quant_backend].__init__) - init_kwargs_transformers = {name for name in init_kwargs_transformers.parameters if name != "self"} - else: - init_kwargs_transformers = None - - init_kwargs_diffusers = inspect.signature(quant_config_mapping_diffusers[quant_backend].__init__) - init_kwargs_diffusers = {name for name in init_kwargs_diffusers.parameters if name != "self"} - - if init_kwargs_transformers != init_kwargs_diffusers: - raise ValueError( - "The signatures of the __init__ methods of the quantization config classes in `diffusers` and `transformers` don't match. " - f"Please provide a `quant_mapping` instead, in the {self.__class__.__name__} class. Refer to [the docs](https://huggingface.co/docs/diffusers/main/en/quantization/overview#pipeline-level-quantization) to learn more about how " - "this mapping would look like." - ) - - def _validate_quant_mapping_args(self): - quant_mapping = self.quant_mapping - transformers_map, diffusers_map = self._get_quant_config_list() - - available_transformers = list(transformers_map.values()) if transformers_map else None - available_diffusers = list(diffusers_map.values()) - - for module_name, config in quant_mapping.items(): - if any(isinstance(config, cfg) for cfg in available_diffusers): - continue - - if available_transformers and any(isinstance(config, cfg) for cfg in available_transformers): - continue - - if available_transformers: - raise ValueError( - f"Provided config for module_name={module_name} could not be found. " - f"Available diffusers configs: {available_diffusers}; " - f"Available transformers configs: {available_transformers}." - ) - else: - raise ValueError( - f"Provided config for module_name={module_name} could not be found. " - f"Available diffusers configs: {available_diffusers}." - ) - - def _check_backend_availability(self, quant_backend: str): - quant_config_mapping_transformers, quant_config_mapping_diffusers = self._get_quant_config_list() - - available_backends_transformers = ( - list(quant_config_mapping_transformers.keys()) if quant_config_mapping_transformers else None - ) - available_backends_diffusers = list(quant_config_mapping_diffusers.keys()) - - if ( - available_backends_transformers and quant_backend not in available_backends_transformers - ) or quant_backend not in quant_config_mapping_diffusers: - error_message = f"Provided quant_backend={quant_backend} was not found." - if available_backends_transformers: - error_message += f"\nAvailable ones (transformers): {available_backends_transformers}." - error_message += f"\nAvailable ones (diffusers): {available_backends_diffusers}." - raise ValueError(error_message) - - def _resolve_quant_config(self, is_diffusers: bool = True, module_name: str = None): - quant_config_mapping_transformers, quant_config_mapping_diffusers = self._get_quant_config_list() - - quant_mapping = self.quant_mapping - components_to_quantize = self.components_to_quantize - - # Granular case - if self.is_granular and module_name in quant_mapping: - logger.debug(f"Initializing quantization config class for {module_name}.") - config = quant_mapping[module_name] - return config - - # Global config case - else: - should_quantize = False - # Only quantize the modules requested for. - if components_to_quantize and module_name in components_to_quantize: - should_quantize = True - # No specification for `components_to_quantize` means all modules should be quantized. - elif not self.is_granular and not components_to_quantize: - should_quantize = True - - if should_quantize: - logger.debug(f"Initializing quantization config class for {module_name}.") - mapping_to_use = quant_config_mapping_diffusers if is_diffusers else quant_config_mapping_transformers - quant_config_cls = mapping_to_use[self.quant_backend] - quant_kwargs = self.quant_kwargs - return quant_config_cls(**quant_kwargs) - - # Fallback: no applicable configuration found. - return None - - def _get_quant_config_list(self): - if is_transformers_available(): - from transformers.quantizers.auto import ( - AUTO_QUANTIZATION_CONFIG_MAPPING as quant_config_mapping_transformers, - ) - else: - quant_config_mapping_transformers = None - - from ..quantizers.auto import AUTO_QUANTIZATION_CONFIG_MAPPING as quant_config_mapping_diffusers - - return quant_config_mapping_transformers, quant_config_mapping_diffusers diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 60ed96c629edfe99f74a4dd382728192520c5969..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/__pycache__/auto.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/__pycache__/auto.cpython-310.pyc deleted file mode 100644 index 29bfe7d238c120588f15506192c2c57d3e3e4262..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/__pycache__/auto.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/__pycache__/base.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/__pycache__/base.cpython-310.pyc deleted file mode 100644 index f5cc9c9bedc125751a28a8882a7ca38aa9af89f8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/__pycache__/base.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/__pycache__/quantization_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/__pycache__/quantization_config.cpython-310.pyc deleted file mode 100644 index f0534a97a115013b7dc9830e688c625e613e1a32..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/__pycache__/quantization_config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/auto.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/auto.py deleted file mode 100644 index ce214ae7bc17da37d542fabe0341b7e541c7e708..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/auto.py +++ /dev/null @@ -1,143 +0,0 @@ -# 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. -""" -Adapted from -https://github.com/huggingface/transformers/blob/c409cd81777fb27aadc043ed3d8339dbc020fb3b/src/transformers/quantizers/auto.py -""" - -import warnings -from typing import Dict, Optional, Union - -from .bitsandbytes import BnB4BitDiffusersQuantizer, BnB8BitDiffusersQuantizer -from .gguf import GGUFQuantizer -from .quantization_config import ( - BitsAndBytesConfig, - GGUFQuantizationConfig, - QuantizationConfigMixin, - QuantizationMethod, - QuantoConfig, - TorchAoConfig, -) -from .quanto import QuantoQuantizer -from .torchao import TorchAoHfQuantizer - - -AUTO_QUANTIZER_MAPPING = { - "bitsandbytes_4bit": BnB4BitDiffusersQuantizer, - "bitsandbytes_8bit": BnB8BitDiffusersQuantizer, - "gguf": GGUFQuantizer, - "quanto": QuantoQuantizer, - "torchao": TorchAoHfQuantizer, -} - -AUTO_QUANTIZATION_CONFIG_MAPPING = { - "bitsandbytes_4bit": BitsAndBytesConfig, - "bitsandbytes_8bit": BitsAndBytesConfig, - "gguf": GGUFQuantizationConfig, - "quanto": QuantoConfig, - "torchao": TorchAoConfig, -} - - -class DiffusersAutoQuantizer: - """ - The auto diffusers quantizer class that takes care of automatically instantiating to the correct - `DiffusersQuantizer` given the `QuantizationConfig`. - """ - - @classmethod - 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): - suffix = "_4bit" if quantization_config_dict.get("load_in_4bit", False) else "_8bit" - quant_method = QuantizationMethod.BITS_AND_BYTES + suffix - elif quant_method is None: - raise ValueError( - "The model's quantization config from the arguments has no `quant_method` attribute. Make sure that the model has been correctly quantized" - ) - - if quant_method not in AUTO_QUANTIZATION_CONFIG_MAPPING.keys(): - raise ValueError( - f"Unknown quantization type, got {quant_method} - supported types are:" - f" {list(AUTO_QUANTIZER_MAPPING.keys())}" - ) - - target_cls = AUTO_QUANTIZATION_CONFIG_MAPPING[quant_method] - return target_cls.from_dict(quantization_config_dict) - - @classmethod - def from_config(cls, quantization_config: Union[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) - - quant_method = quantization_config.quant_method - - # Again, we need a special care for bnb as we have a single quantization config - # class for both 4-bit and 8-bit quantization - if quant_method == QuantizationMethod.BITS_AND_BYTES: - if quantization_config.load_in_8bit: - quant_method += "_8bit" - else: - quant_method += "_4bit" - - if quant_method not in AUTO_QUANTIZER_MAPPING.keys(): - raise ValueError( - f"Unknown quantization type, got {quant_method} - supported types are:" - f" {list(AUTO_QUANTIZER_MAPPING.keys())}" - ) - - target_cls = AUTO_QUANTIZER_MAPPING[quant_method] - return target_cls(quantization_config, **kwargs) - - @classmethod - def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): - model_config = cls.load_config(pretrained_model_name_or_path, **kwargs) - if getattr(model_config, "quantization_config", None) is None: - raise ValueError( - f"Did not found a `quantization_config` in {pretrained_model_name_or_path}. Make sure that the model is correctly quantized." - ) - quantization_config_dict = model_config.quantization_config - quantization_config = cls.from_dict(quantization_config_dict) - # Update with potential kwargs that are passed through from_pretrained. - quantization_config.update(kwargs) - - return cls.from_config(quantization_config) - - @classmethod - def merge_quantization_configs( - cls, - quantization_config: Union[dict, QuantizationConfigMixin], - quantization_config_from_args: Optional[QuantizationConfigMixin], - ): - """ - handles situations where both quantization_config from args and quantization_config from model config are - present. - """ - if quantization_config_from_args is not None: - warning_msg = ( - "You passed `quantization_config` or equivalent parameters to `from_pretrained` but the model you're loading" - " already has a `quantization_config` attribute. The `quantization_config` from the model will be used." - ) - else: - warning_msg = "" - - if isinstance(quantization_config, dict): - quantization_config = cls.from_dict(quantization_config) - - if warning_msg != "": - warnings.warn(warning_msg) - - return quantization_config diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/base.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/base.py deleted file mode 100644 index 357d920d29c45dd28d123ae16d76f83d8a069cc1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/base.py +++ /dev/null @@ -1,234 +0,0 @@ -# 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. - -""" -Adapted from -https://github.com/huggingface/transformers/blob/52cb4034ada381fe1ffe8d428a1076e5411a8026/src/transformers/quantizers/base.py -""" - -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union - -from ..utils import is_torch_available -from .quantization_config import QuantizationConfigMixin - - -if TYPE_CHECKING: - from ..models.modeling_utils import ModelMixin - -if is_torch_available(): - import torch - - -class DiffusersQuantizer(ABC): - """ - Abstract class of the HuggingFace quantizer. Supports for now quantizing HF diffusers models for inference and/or - quantization. This class is used only for diffusers.models.modeling_utils.ModelMixin.from_pretrained and cannot be - easily used outside the scope of that method yet. - - 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*): - The list of module names to not convert when quantizing the model. - 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. - """ - - requires_calibration = False - required_packages = None - - def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): - self.quantization_config = quantization_config - - # -- Handle extra kwargs below -- - self.modules_to_not_convert = kwargs.pop("modules_to_not_convert", []) - self.pre_quantized = kwargs.pop("pre_quantized", True) - - if not self.pre_quantized and self.requires_calibration: - raise ValueError( - f"The quantization method {quantization_config.quant_method} does require the model to be pre-quantized." - f" You explicitly passed `pre_quantized=False` meaning your model weights are not quantized. Make sure to " - f"pass `pre_quantized=True` while knowing what you are doing." - ) - - def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": - """ - Some quantization methods require to explicitly set the dtype of the model to a target dtype. You need to - override this method in case you want to make sure that behavior is preserved - - Args: - torch_dtype (`torch.dtype`): - The input dtype that is passed in `from_pretrained` - """ - return torch_dtype - - def update_device_map(self, device_map: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: - """ - 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*): - The device_map that is passed through the `from_pretrained` method. - """ - return device_map - - def adjust_target_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": - """ - Override this method if you want to adjust the `target_dtype` variable used in `from_pretrained` to compute the - device_map in case the device_map is a `str`. E.g. for bitsandbytes we force-set `target_dtype` to `torch.int8` - and for 4-bit we pass a custom enum `accelerate.CustomDtype.int4`. - - Args: - torch_dtype (`torch.dtype`, *optional*): - The torch_dtype that is used to compute the device_map. - """ - return torch_dtype - - 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*): - 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"]: - """ - 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 - `_process_model_before_weight_loading`. `diffusers` models don't have any `modules_to_not_convert` attributes - yet but this can change soon in the future. - - Args: - model (`~diffusers.models.modeling_utils.ModelMixin`): - The model to quantize - torch_dtype (`torch.dtype`): - The dtype passed in `from_pretrained` method. - """ - - return { - name: torch_dtype - for name, _ in model.named_parameters() - 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]]: - """adjust max_memory argument for infer_auto_device_map() if extra memory is needed for quantization""" - return max_memory - - def check_if_quantized_param( - self, - model: "ModelMixin", - param_value: "torch.Tensor", - param_name: str, - state_dict: Dict[str, Any], - **kwargs, - ) -> bool: - """ - checks if a loaded state_dict component is part of quantized param + some validation; only defined for - quantization methods that require to create a new parameters for quantization. - """ - return False - - def create_quantized_param(self, *args, **kwargs) -> "torch.nn.Parameter": - """ - takes needed components from state_dict and creates quantized param. - """ - return - - def check_quantized_param_shape(self, *args, **kwargs): - """ - checks if the quantized param has expected shape. - """ - return True - - def validate_environment(self, *args, **kwargs): - """ - This method is used to potentially check for potential conflicts with arguments that are passed in - `from_pretrained`. You need to define it for all future quantizers that are integrated with diffusers. If no - explicit check are needed, simply return nothing. - """ - return - - def preprocess_model(self, model: "ModelMixin", **kwargs): - """ - Setting model attributes and/or converting model before weights loading. At this point the model should be - initialized on the meta device so you can freely manipulate the skeleton of the model in order to replace - modules in-place. Make sure to override the abstract method `_process_model_before_weight_loading`. - - Args: - model (`~diffusers.models.modeling_utils.ModelMixin`): - The model to quantize - kwargs (`dict`, *optional*): - The keyword arguments that are passed along `_process_model_before_weight_loading`. - """ - model.is_quantized = True - model.quantization_method = self.quantization_config.quant_method - return self._process_model_before_weight_loading(model, **kwargs) - - def postprocess_model(self, model: "ModelMixin", **kwargs): - """ - Post-process the model post weights loading. Make sure to override the abstract method - `_process_model_after_weight_loading`. - - Args: - model (`~diffusers.models.modeling_utils.ModelMixin`): - The model to quantize - kwargs (`dict`, *optional*): - The keyword arguments that are passed along `_process_model_after_weight_loading`. - """ - return self._process_model_after_weight_loading(model, **kwargs) - - def dequantize(self, model): - """ - Potentially dequantize the model to retrieve the original model, with some loss in accuracy / performance. Note - not all quantization schemes support this. - """ - model = self._dequantize(model) - - # Delete quantizer and quantization config - del model.hf_quantizer - - return model - - def _dequantize(self, model): - raise NotImplementedError( - f"{self.quantization_config.quant_method} has no implementation of `dequantize`, please raise an issue on GitHub." - ) - - @abstractmethod - def _process_model_before_weight_loading(self, model, **kwargs): ... - - @abstractmethod - def _process_model_after_weight_loading(self, model, **kwargs): ... - - @property - @abstractmethod - def is_serializable(self): ... - - @property - @abstractmethod - def is_trainable(self): ... - - @property - def is_compileable(self) -> bool: - """Flag indicating whether the quantized model can be compiled""" - return False diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/__init__.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/__init__.py deleted file mode 100644 index 9e745bc810faad76608b00ab9c493369466d0970..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .bnb_quantizer import BnB4BitDiffusersQuantizer, BnB8BitDiffusersQuantizer -from .utils import dequantize_and_replace, dequantize_bnb_weight, replace_with_bnb_linear diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 38abd3b99d6d8fb778e076ee0daeaa8d51ff8fee..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/__pycache__/bnb_quantizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/__pycache__/bnb_quantizer.cpython-310.pyc deleted file mode 100644 index 490b438cdc8790bda2e5408e5b88f3a3987491bc..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/__pycache__/bnb_quantizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index 18b1a46671cb701c1bc18a73762f956dd6e218d0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/bnb_quantizer.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/bnb_quantizer.py deleted file mode 100644 index 0dfdff019b79ef3bfb9b89df6378ab398199beeb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/bnb_quantizer.py +++ /dev/null @@ -1,577 +0,0 @@ -# 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. -""" -Adapted from -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 ...utils import get_module_from_name -from ..base import DiffusersQuantizer - - -if TYPE_CHECKING: - from ...models.modeling_utils import ModelMixin - -from ...utils import ( - is_accelerate_available, - is_accelerate_version, - is_bitsandbytes_available, - is_bitsandbytes_version, - is_torch_available, - logging, -) - - -if is_torch_available(): - import torch - -logger = logging.get_logger(__name__) - - -class BnB4BitDiffusersQuantizer(DiffusersQuantizer): - """ - 4-bit quantization from bitsandbytes.py quantization method: - before loading: converts transformer layers into Linear4bit during loading: load 16bit weight and pass to the - layer object after: quantizes individual weights in Linear4bit into 4bit at the first .cuda() call saving: - from state dict, as usual; saves weights and `quant_state` components - loading: - need to locate `quant_state` components and pass to Param4bit constructor - """ - - use_keep_in_fp32_modules = True - requires_calibration = False - - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - - if self.quantization_config.llm_int8_skip_modules is not None: - self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules - - def validate_environment(self, *args, **kwargs): - if not (torch.cuda.is_available() or torch.xpu.is_available()): - raise RuntimeError("No GPU found. A GPU is needed for quantization.") - if not is_accelerate_available() or is_accelerate_version("<", "0.26.0"): - raise ImportError( - "Using `bitsandbytes` 4-bit quantization requires Accelerate: `pip install 'accelerate>=0.26.0'`" - ) - if not is_bitsandbytes_available() or is_bitsandbytes_version("<", "0.43.3"): - raise ImportError( - "Using `bitsandbytes` 4-bit quantization requires the latest version of bitsandbytes: `pip install -U bitsandbytes`" - ) - - if kwargs.get("from_flax", False): - raise ValueError( - "Converting into 4-bit weights from flax weights is currently not supported, please make" - " sure the weights are in PyTorch format." - ) - - device_map = kwargs.get("device_map", None) - if ( - device_map is not None - and isinstance(device_map, dict) - and not self.quantization_config.llm_int8_enable_fp32_cpu_offload - ): - device_map_without_no_convert = { - key: device_map[key] for key in device_map.keys() if key not in self.modules_to_not_convert - } - if "cpu" in device_map_without_no_convert.values() or "disk" in device_map_without_no_convert.values(): - raise ValueError( - "Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit the " - "quantized model. If you want to dispatch the model on the CPU or the disk while keeping these modules " - "in 32-bit, you need to set `load_in_8bit_fp32_cpu_offload=True` and pass a custom `device_map` to " - "`from_pretrained`. Check " - "https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu " - "for more details. " - ) - - def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": - if target_dtype != torch.int8: - from accelerate.utils import CustomDtype - - logger.info("target_dtype {target_dtype} is replaced by `CustomDtype.INT4` for 4-bit BnB quantization") - return CustomDtype.INT4 - else: - raise ValueError(f"Wrong `target_dtype` ({target_dtype}) provided.") - - def check_if_quantized_param( - self, - model: "ModelMixin", - param_value: "torch.Tensor", - param_name: str, - state_dict: Dict[str, Any], - **kwargs, - ) -> bool: - import bitsandbytes as bnb - - module, tensor_name = get_module_from_name(model, param_name) - if isinstance(module._parameters.get(tensor_name, None), bnb.nn.Params4bit): - # Add here check for loaded components' dtypes once serialization is implemented - return True - elif isinstance(module, bnb.nn.Linear4bit) and tensor_name == "bias": - # bias could be loaded by regular set_module_tensor_to_device() from accelerate, - # but it would wrongly use uninitialized weight there. - return True - else: - return False - - def create_quantized_param( - self, - model: "ModelMixin", - param_value: "torch.Tensor", - param_name: str, - target_device: "torch.device", - state_dict: Dict[str, Any], - unexpected_keys: Optional[List[str]] = None, - **kwargs, - ): - import bitsandbytes as bnb - - module, tensor_name = get_module_from_name(model, param_name) - - if tensor_name not in module._parameters: - raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.") - - old_value = getattr(module, tensor_name) - - if tensor_name == "bias": - if param_value is None: - new_value = old_value.to(target_device) - else: - new_value = param_value.to(target_device) - - new_value = torch.nn.Parameter(new_value, requires_grad=old_value.requires_grad) - module._parameters[tensor_name] = new_value - return - - if not isinstance(module._parameters[tensor_name], bnb.nn.Params4bit): - raise ValueError("this function only loads `Linear4bit components`") - if ( - old_value.device == torch.device("meta") - and target_device not in ["meta", torch.device("meta")] - and param_value is None - ): - raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {target_device}.") - - # construct `new_value` for the module._parameters[tensor_name]: - if self.pre_quantized: - # 4bit loading. Collecting components for restoring quantized weight - # This can be expanded to make a universal call for any quantized weight loading - - if not self.is_serializable: - raise ValueError( - "Detected int4 weights but the version of bitsandbytes is not compatible with int4 serialization. " - "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`." - ) - - if (param_name + ".quant_state.bitsandbytes__fp4" not in state_dict) and ( - param_name + ".quant_state.bitsandbytes__nf4" not in state_dict - ): - raise ValueError( - f"Supplied state dict for {param_name} does not contain `bitsandbytes__*` and possibly other `quantized_stats` components." - ) - - quantized_stats = {} - for k, v in state_dict.items(): - # `startswith` to counter for edge cases where `param_name` - # substring can be present in multiple places in the `state_dict` - if param_name + "." in k and k.startswith(param_name): - quantized_stats[k] = v - if unexpected_keys is not None and k in unexpected_keys: - unexpected_keys.remove(k) - - new_value = bnb.nn.Params4bit.from_prequantized( - data=param_value, - quantized_stats=quantized_stats, - requires_grad=False, - device=target_device, - ) - else: - new_value = param_value.to("cpu") - kwargs = old_value.__dict__ - new_value = bnb.nn.Params4bit(new_value, requires_grad=False, **kwargs).to(target_device) - - module._parameters[tensor_name] = new_value - - def check_quantized_param_shape(self, param_name, current_param, loaded_param): - current_param_shape = current_param.shape - loaded_param_shape = loaded_param.shape - - n = current_param_shape.numel() - inferred_shape = (n,) if "bias" in param_name else ((n + 1) // 2, 1) - if loaded_param_shape != inferred_shape: - raise ValueError( - f"Expected the flattened shape of the current param ({param_name}) to be {loaded_param_shape} but is {inferred_shape}." - ) - else: - return True - - def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[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 - - def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": - if torch_dtype is None: - # We force the `dtype` to be float16, this is a requirement from `bitsandbytes` - logger.info( - "Overriding torch_dtype=%s with `torch_dtype=torch.float16` due to " - "requirements of `bitsandbytes` to enable model loading in 8-bit or 4-bit. " - "Pass your own torch_dtype to specify the dtype of the remaining non-linear layers or pass" - " torch_dtype=torch.float16 to remove this warning.", - torch_dtype, - ) - torch_dtype = torch.float16 - return torch_dtype - - def update_device_map(self, device_map): - if device_map is None: - if torch.xpu.is_available(): - current_device = f"xpu:{torch.xpu.current_device()}" - else: - current_device = f"cuda:{torch.cuda.current_device()}" - device_map = {"": current_device} - logger.info( - "The device_map was not initialized. " - "Setting device_map to {" - ": {current_device}}. " - "If you want to use the model for inference, please set device_map ='auto' " - ) - return device_map - - def _process_model_before_weight_loading( - self, - model: "ModelMixin", - device_map, - keep_in_fp32_modules: List[str] = [], - **kwargs, - ): - from .utils import replace_with_bnb_linear - - load_in_8bit_fp32_cpu_offload = self.quantization_config.llm_int8_enable_fp32_cpu_offload - - # We may keep some modules such as the `proj_out` in their original dtype for numerical stability reasons - self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules - - if not isinstance(self.modules_to_not_convert, list): - self.modules_to_not_convert = [self.modules_to_not_convert] - - self.modules_to_not_convert.extend(keep_in_fp32_modules) - - # Extend `self.modules_to_not_convert` to keys that are supposed to be offloaded to `cpu` or `disk` - if isinstance(device_map, dict) and len(device_map.keys()) > 1: - keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]] - - if len(keys_on_cpu) > 0 and not load_in_8bit_fp32_cpu_offload: - raise ValueError( - "If you want to offload some keys to `cpu` or `disk`, you need to set " - "`llm_int8_enable_fp32_cpu_offload=True`. Note that these modules will not be " - " converted to 8-bit but kept in 32-bit." - ) - self.modules_to_not_convert.extend(keys_on_cpu) - - # Purge `None`. - # Unlike `transformers`, we don't know if we should always keep certain modules in FP32 - # in case of diffusion transformer models. For language models and others alike, `lm_head` - # and tied modules are usually kept in FP32. - self.modules_to_not_convert = [module for module in self.modules_to_not_convert if module is not None] - - model = replace_with_bnb_linear( - model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config - ) - model.config.quantization_config = self.quantization_config - model.is_loaded_in_4bit = True - - def _process_model_after_weight_loading(self, model: "ModelMixin", **kwargs): - model.is_4bit_serializable = self.is_serializable - return model - - @property - def is_serializable(self): - # Because we're mandating `bitsandbytes` 0.43.3. - return True - - @property - def is_trainable(self) -> bool: - # Because we're mandating `bitsandbytes` 0.43.3. - return True - - def _dequantize(self, model): - from .utils import dequantize_and_replace - - is_model_on_cpu = model.device.type == "cpu" - if is_model_on_cpu: - logger.info( - "Model was found to be on CPU (could happen as a result of `enable_model_cpu_offload()`). So, moving it to GPU. After dequantization, will move the model back to CPU again to preserve the previous device." - ) - if torch.xpu.is_available(): - model.to(torch.xpu.current_device()) - else: - model.to(torch.cuda.current_device()) - - model = dequantize_and_replace( - model, self.modules_to_not_convert, quantization_config=self.quantization_config - ) - if is_model_on_cpu: - model.to("cpu") - return model - - -class BnB8BitDiffusersQuantizer(DiffusersQuantizer): - """ - 8-bit quantization from bitsandbytes quantization method: - before loading: converts transformer layers into Linear8bitLt during loading: load 16bit weight and pass to the - layer object after: quantizes individual weights in Linear8bitLt into 8bit at fitst .cuda() call - saving: - from state dict, as usual; saves weights and 'SCB' component - loading: - need to locate SCB component and pass to the Linear8bitLt object - """ - - use_keep_in_fp32_modules = True - requires_calibration = False - - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - - if self.quantization_config.llm_int8_skip_modules is not None: - self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules - - def validate_environment(self, *args, **kwargs): - if not (torch.cuda.is_available() or torch.xpu.is_available()): - raise RuntimeError("No GPU found. A GPU is needed for quantization.") - if not is_accelerate_available() or is_accelerate_version("<", "0.26.0"): - raise ImportError( - "Using `bitsandbytes` 8-bit quantization requires Accelerate: `pip install 'accelerate>=0.26.0'`" - ) - if not is_bitsandbytes_available() or is_bitsandbytes_version("<", "0.43.3"): - raise ImportError( - "Using `bitsandbytes` 8-bit quantization requires the latest version of bitsandbytes: `pip install -U bitsandbytes`" - ) - - if kwargs.get("from_flax", False): - raise ValueError( - "Converting into 8-bit weights from flax weights is currently not supported, please make" - " sure the weights are in PyTorch format." - ) - - device_map = kwargs.get("device_map", None) - if ( - device_map is not None - and isinstance(device_map, dict) - and not self.quantization_config.llm_int8_enable_fp32_cpu_offload - ): - device_map_without_no_convert = { - key: device_map[key] for key in device_map.keys() if key not in self.modules_to_not_convert - } - if "cpu" in device_map_without_no_convert.values() or "disk" in device_map_without_no_convert.values(): - raise ValueError( - "Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit the " - "quantized model. If you want to dispatch the model on the CPU or the disk while keeping these modules " - "in 32-bit, you need to set `load_in_8bit_fp32_cpu_offload=True` and pass a custom `device_map` to " - "`from_pretrained`. Check " - "https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu " - "for more details. " - ) - - # 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]]: - # 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 - - # Copied from diffusers.quantizers.bitsandbytes.bnb_quantizer.BnB4BitDiffusersQuantizer.update_torch_dtype - def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": - if torch_dtype is None: - # We force the `dtype` to be float16, this is a requirement from `bitsandbytes` - logger.info( - "Overriding torch_dtype=%s with `torch_dtype=torch.float16` due to " - "requirements of `bitsandbytes` to enable model loading in 8-bit or 4-bit. " - "Pass your own torch_dtype to specify the dtype of the remaining non-linear layers or pass" - " torch_dtype=torch.float16 to remove this warning.", - torch_dtype, - ) - torch_dtype = torch.float16 - return torch_dtype - - # Copied from diffusers.quantizers.bitsandbytes.bnb_quantizer.BnB4BitDiffusersQuantizer.update_device_map - def update_device_map(self, device_map): - if device_map is None: - if torch.xpu.is_available(): - current_device = f"xpu:{torch.xpu.current_device()}" - else: - current_device = f"cuda:{torch.cuda.current_device()}" - device_map = {"": current_device} - logger.info( - "The device_map was not initialized. " - "Setting device_map to {" - ": {current_device}}. " - "If you want to use the model for inference, please set device_map ='auto' " - ) - return device_map - - def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": - if target_dtype != torch.int8: - logger.info("target_dtype {target_dtype} is replaced by `torch.int8` for 8-bit BnB quantization") - return torch.int8 - - def check_if_quantized_param( - self, - model: "ModelMixin", - param_value: "torch.Tensor", - param_name: str, - state_dict: Dict[str, Any], - **kwargs, - ): - import bitsandbytes as bnb - - module, tensor_name = get_module_from_name(model, param_name) - if isinstance(module._parameters.get(tensor_name, None), bnb.nn.Int8Params): - if self.pre_quantized: - if param_name.replace("weight", "SCB") not in state_dict.keys(): - raise ValueError("Missing quantization component `SCB`") - if param_value.dtype != torch.int8: - raise ValueError( - f"Incompatible dtype `{param_value.dtype}` when loading 8-bit prequantized weight. Expected `torch.int8`." - ) - return True - return False - - def create_quantized_param( - self, - model: "ModelMixin", - param_value: "torch.Tensor", - param_name: str, - target_device: "torch.device", - state_dict: Dict[str, Any], - unexpected_keys: Optional[List[str]] = None, - **kwargs, - ): - import bitsandbytes as bnb - - fp16_statistics_key = param_name.replace("weight", "SCB") - fp16_weights_format_key = param_name.replace("weight", "weight_format") - - fp16_statistics = state_dict.get(fp16_statistics_key, None) - fp16_weights_format = state_dict.get(fp16_weights_format_key, None) - - module, tensor_name = get_module_from_name(model, param_name) - if tensor_name not in module._parameters: - raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.") - - old_value = getattr(module, tensor_name) - - if not isinstance(module._parameters[tensor_name], bnb.nn.Int8Params): - raise ValueError(f"Parameter `{tensor_name}` should only be a `bnb.nn.Int8Params` instance.") - if ( - old_value.device == torch.device("meta") - and target_device not in ["meta", torch.device("meta")] - and param_value is None - ): - raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {target_device}.") - - new_value = param_value.to("cpu") - if self.pre_quantized and not self.is_serializable: - raise ValueError( - "Detected int8 weights but the version of bitsandbytes is not compatible with int8 serialization. " - "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`." - ) - - kwargs = old_value.__dict__ - new_value = bnb.nn.Int8Params(new_value, requires_grad=False, **kwargs).to(target_device) - - module._parameters[tensor_name] = new_value - if fp16_statistics is not None: - setattr(module.weight, "SCB", fp16_statistics.to(target_device)) - if unexpected_keys is not None: - unexpected_keys.remove(fp16_statistics_key) - - # We just need to pop the `weight_format` keys from the state dict to remove unneeded - # messages. The correct format is correctly retrieved during the first forward pass. - if fp16_weights_format is not None and unexpected_keys is not None: - unexpected_keys.remove(fp16_weights_format_key) - - # Copied from diffusers.quantizers.bitsandbytes.bnb_quantizer.BnB4BitDiffusersQuantizer._process_model_after_weight_loading with 4bit->8bit - def _process_model_after_weight_loading(self, model: "ModelMixin", **kwargs): - model.is_8bit_serializable = self.is_serializable - return model - - # Copied from diffusers.quantizers.bitsandbytes.bnb_quantizer.BnB4BitDiffusersQuantizer._process_model_before_weight_loading with 4bit->8bit - def _process_model_before_weight_loading( - self, - model: "ModelMixin", - device_map, - keep_in_fp32_modules: List[str] = [], - **kwargs, - ): - from .utils import replace_with_bnb_linear - - load_in_8bit_fp32_cpu_offload = self.quantization_config.llm_int8_enable_fp32_cpu_offload - - # We may keep some modules such as the `proj_out` in their original dtype for numerical stability reasons - self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules - - if not isinstance(self.modules_to_not_convert, list): - self.modules_to_not_convert = [self.modules_to_not_convert] - - self.modules_to_not_convert.extend(keep_in_fp32_modules) - - # Extend `self.modules_to_not_convert` to keys that are supposed to be offloaded to `cpu` or `disk` - if isinstance(device_map, dict) and len(device_map.keys()) > 1: - keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]] - - if len(keys_on_cpu) > 0 and not load_in_8bit_fp32_cpu_offload: - raise ValueError( - "If you want to offload some keys to `cpu` or `disk`, you need to set " - "`llm_int8_enable_fp32_cpu_offload=True`. Note that these modules will not be " - " converted to 8-bit but kept in 32-bit." - ) - self.modules_to_not_convert.extend(keys_on_cpu) - - # Purge `None`. - # Unlike `transformers`, we don't know if we should always keep certain modules in FP32 - # in case of diffusion transformer models. For language models and others alike, `lm_head` - # and tied modules are usually kept in FP32. - self.modules_to_not_convert = [module for module in self.modules_to_not_convert if module is not None] - - model = replace_with_bnb_linear( - model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config - ) - model.config.quantization_config = self.quantization_config - model.is_loaded_in_8bit = True - - @property - # Copied from diffusers.quantizers.bitsandbytes.bnb_quantizer.BnB4BitDiffusersQuantizer.is_serializable - def is_serializable(self): - # Because we're mandating `bitsandbytes` 0.43.3. - return True - - @property - # Copied from diffusers.quantizers.bitsandbytes.bnb_quantizer.BnB4BitDiffusersQuantizer.is_serializable - def is_trainable(self) -> bool: - # Because we're mandating `bitsandbytes` 0.43.3. - return True - - @property - def is_compileable(self) -> bool: - return True - - def _dequantize(self, model): - from .utils import dequantize_and_replace - - model = dequantize_and_replace( - model, self.modules_to_not_convert, quantization_config=self.quantization_config - ) - return model diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/utils.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/utils.py deleted file mode 100644 index 429aabb8fae6de0a2620cde68fba61db9a0557a7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/bitsandbytes/utils.py +++ /dev/null @@ -1,319 +0,0 @@ -# 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. -""" -Adapted from -https://github.com/huggingface/transformers/blob/c409cd81777fb27aadc043ed3d8339dbc020fb3b/src/transformers/integrations/bitsandbytes.py -""" - -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 - - -if is_torch_available(): - import torch - import torch.nn as nn - -if is_bitsandbytes_available(): - import bitsandbytes as bnb - -if is_accelerate_available(): - import accelerate - from accelerate import init_empty_weights - from accelerate.hooks import add_hook_to_module, remove_hook_from_module - -logger = logging.get_logger(__name__) - - -def _replace_with_bnb_linear( - model, - modules_to_not_convert=None, - current_key_name=None, - quantization_config=None, - has_been_replaced=False, -): - """ - Private method that wraps the recursion for module replacement. - - Returns the converted model and a boolean that indicates if the conversion has been successful or not. - """ - for name, module in model.named_children(): - if current_key_name is None: - current_key_name = [] - current_key_name.append(name) - - if isinstance(module, nn.Linear) and name not in modules_to_not_convert: - # Check if the current key is not in the `modules_to_not_convert` - current_key_name_str = ".".join(current_key_name) - if not any( - (key + "." in current_key_name_str) or (key == current_key_name_str) for key in modules_to_not_convert - ): - with init_empty_weights(): - in_features = module.in_features - out_features = module.out_features - - if quantization_config.quantization_method() == "llm_int8": - model._modules[name] = bnb.nn.Linear8bitLt( - in_features, - out_features, - module.bias is not None, - has_fp16_weights=quantization_config.llm_int8_has_fp16_weight, - threshold=quantization_config.llm_int8_threshold, - ) - has_been_replaced = True - else: - if ( - quantization_config.llm_int8_skip_modules is not None - and name in quantization_config.llm_int8_skip_modules - ): - pass - else: - extra_kwargs = ( - {"quant_storage": quantization_config.bnb_4bit_quant_storage} - if "quant_storage" in list(signature(bnb.nn.Linear4bit).parameters) - else {} - ) - model._modules[name] = bnb.nn.Linear4bit( - in_features, - out_features, - module.bias is not None, - quantization_config.bnb_4bit_compute_dtype, - compress_statistics=quantization_config.bnb_4bit_use_double_quant, - quant_type=quantization_config.bnb_4bit_quant_type, - **extra_kwargs, - ) - has_been_replaced = True - # Store the module class in case we need to transpose the weight later - model._modules[name].source_cls = type(module) - # Force requires grad to False to avoid unexpected errors - model._modules[name].requires_grad_(False) - if len(list(module.children())) > 0: - _, has_been_replaced = _replace_with_bnb_linear( - module, - modules_to_not_convert, - current_key_name, - quantization_config, - has_been_replaced=has_been_replaced, - ) - # Remove the last key for recursion - current_key_name.pop(-1) - return model, has_been_replaced - - -def replace_with_bnb_linear(model, modules_to_not_convert=None, current_key_name=None, quantization_config=None): - """ - Helper function to replace the `nn.Linear` layers within `model` with either `bnb.nn.Linear8bit` or - `bnb.nn.Linear4bit` using the `bitsandbytes` library. - - References: - * `bnb.nn.Linear8bit`: [LLM.int8(): 8-bit Matrix Multiplication for Transformers at - Scale](https://huggingface.co/papers/2208.07339) - * `bnb.nn.Linear4bit`: [QLoRA: Efficient Finetuning of Quantized - LLMs](https://huggingface.co/papers/2305.14314) - - 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 `[]`): - 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*): - 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`). - quantization_config ('transformers.utils.quantization_config.BitsAndBytesConfig'): - To configure and manage settings related to quantization, a technique used to compress neural network - models by reducing the precision of the weights and activations, thus making models more efficient in terms - of both storage and computation. - """ - model, _ = _replace_with_bnb_linear(model, modules_to_not_convert, current_key_name, quantization_config) - - has_been_replaced = any( - isinstance(replaced_module, (bnb.nn.Linear4bit, bnb.nn.Linear8bitLt)) - for _, replaced_module in model.named_modules() - ) - if not has_been_replaced: - logger.warning( - "You are loading your model in 8bit or 4bit but no linear modules were found in your model." - " Please double check your model architecture, or submit an issue on github if you think this is" - " a bug." - ) - - return model - - -# Adapted from PEFT: https://github.com/huggingface/peft/blob/6d458b300fc2ed82e19f796b53af4c97d03ea604/src/peft/utils/integrations.py#L81 -def dequantize_bnb_weight(weight: "torch.nn.Parameter", state=None, dtype: "torch.dtype" = None): - """ - Helper function to dequantize 4bit or 8bit bnb weights. - - If the weight is not a bnb quantized weight, it will be returned as is. - """ - if not isinstance(weight, torch.nn.Parameter): - raise TypeError(f"Input weight should be of type nn.Parameter, got {type(weight)} instead") - - cls_name = weight.__class__.__name__ - if cls_name not in ("Params4bit", "Int8Params"): - return weight - - if cls_name == "Params4bit": - output_tensor = bnb.functional.dequantize_4bit(weight.data, weight.quant_state) - msg = f"The model is going to be dequantized in {output_tensor.dtype} - if you want to upcast it to another dtype, make sure to pass the desired dtype when quantizing the model through `bnb_4bit_quant_type` argument of `BitsAndBytesConfig`" - if dtype: - msg = f"The model is going to be first dequantized in {output_tensor.dtype} and type-casted to {dtype}" - output_tensor = output_tensor.to(dtype) - logger.warning_once(msg) - return output_tensor - - if state.SCB is None: - state.SCB = weight.SCB - - if hasattr(bnb.functional, "int8_vectorwise_dequant"): - # Use bitsandbytes API if available (requires v0.45.0+) - dequantized = bnb.functional.int8_vectorwise_dequant(weight.data, state.SCB) - else: - # Multiply by (scale/127) to dequantize. - dequantized = weight.data * state.SCB.view(-1, 1) * 7.874015718698502e-3 - - if dtype: - dequantized = dequantized.to(dtype) - return dequantized - - -def _create_accelerate_new_hook(old_hook): - r""" - Creates a new hook based on the old hook. Use it only if you know what you are doing ! This method is a copy of: - https://github.com/huggingface/peft/blob/748f7968f3a31ec06a1c2b0328993319ad9a150a/src/peft/utils/other.py#L245 with - some changes - """ - old_hook_cls = getattr(accelerate.hooks, old_hook.__class__.__name__) - old_hook_attr = old_hook.__dict__ - filtered_old_hook_attr = {} - old_hook_init_signature = inspect.signature(old_hook_cls.__init__) - for k in old_hook_attr.keys(): - if k in old_hook_init_signature.parameters: - filtered_old_hook_attr[k] = old_hook_attr[k] - new_hook = old_hook_cls(**filtered_old_hook_attr) - return new_hook - - -def _dequantize_and_replace( - model, - dtype, - modules_to_not_convert=None, - current_key_name=None, - quantization_config=None, - has_been_replaced=False, -): - """ - Converts a quantized model into its dequantized original version. The newly converted model will have some - performance drop compared to the original model before quantization - use it only for specific usecases such as - QLoRA adapters merging. - - Returns the converted model and a boolean that indicates if the conversion has been successful or not. - """ - quant_method = quantization_config.quantization_method() - - target_cls = bnb.nn.Linear8bitLt if quant_method == "llm_int8" else bnb.nn.Linear4bit - - for name, module in model.named_children(): - if current_key_name is None: - current_key_name = [] - current_key_name.append(name) - - if isinstance(module, target_cls) and name not in modules_to_not_convert: - # Check if the current key is not in the `modules_to_not_convert` - current_key_name_str = ".".join(current_key_name) - - if not any( - (key + "." in current_key_name_str) or (key == current_key_name_str) for key in modules_to_not_convert - ): - bias = getattr(module, "bias", None) - - device = module.weight.device - with init_empty_weights(): - new_module = torch.nn.Linear(module.in_features, module.out_features, bias=bias is not None) - - if quant_method == "llm_int8": - state = module.state - else: - state = None - - new_module.weight = torch.nn.Parameter(dequantize_bnb_weight(module.weight, state, dtype)) - - if bias is not None: - new_module.bias = bias - - # Create a new hook and attach it in case we use accelerate - if hasattr(module, "_hf_hook"): - old_hook = module._hf_hook - new_hook = _create_accelerate_new_hook(old_hook) - - remove_hook_from_module(module) - add_hook_to_module(new_module, new_hook) - - new_module.to(device) - model._modules[name] = new_module - has_been_replaced = True - if len(list(module.children())) > 0: - _, has_been_replaced = _dequantize_and_replace( - module, - dtype=dtype, - modules_to_not_convert=modules_to_not_convert, - current_key_name=current_key_name, - quantization_config=quantization_config, - has_been_replaced=has_been_replaced, - ) - # Remove the last key for recursion - current_key_name.pop(-1) - return model, has_been_replaced - - -def dequantize_and_replace( - model, - modules_to_not_convert=None, - quantization_config=None, -): - model, _ = _dequantize_and_replace( - model, - dtype=model.dtype, - modules_to_not_convert=modules_to_not_convert, - quantization_config=quantization_config, - ) - has_been_replaced = any( - isinstance(replaced_module, torch.nn.Linear) for _, replaced_module in model.named_modules() - ) - if not has_been_replaced: - logger.warning( - "Some linear modules were not dequantized. This could lead to unexpected behaviour. Please check your model." - ) - - return model - - -def _check_bnb_status(module) -> Union[bool, bool]: - is_loaded_in_4bit_bnb = ( - hasattr(module, "is_loaded_in_4bit") - and module.is_loaded_in_4bit - and getattr(module, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES - ) - is_loaded_in_8bit_bnb = ( - hasattr(module, "is_loaded_in_8bit") - and module.is_loaded_in_8bit - and getattr(module, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES - ) - return is_loaded_in_4bit_bnb or is_loaded_in_8bit_bnb, is_loaded_in_4bit_bnb, is_loaded_in_8bit_bnb diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/__init__.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/__init__.py deleted file mode 100644 index b3d9082ac803c8da97ba7bb37021445307b96f7c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .gguf_quantizer import GGUFQuantizer diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 1174f604c19982645ae6c956c9d0ee56cb8fc6b3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/__pycache__/gguf_quantizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/__pycache__/gguf_quantizer.cpython-310.pyc deleted file mode 100644 index ad6622eb1d932bde851668d74bd0ff5ad9819195..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/__pycache__/gguf_quantizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index 4b206f9348b90b1c4e3b3c468f80383cdd50bd6f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/gguf_quantizer.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/gguf_quantizer.py deleted file mode 100644 index aa5ebf5711a3731c4768da9ede0206b2edc64ade..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/gguf_quantizer.py +++ /dev/null @@ -1,169 +0,0 @@ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union - -from ..base import DiffusersQuantizer - - -if TYPE_CHECKING: - from ...models.modeling_utils import ModelMixin - - -from ...utils import ( - get_module_from_name, - is_accelerate_available, - is_accelerate_version, - is_gguf_available, - is_gguf_version, - is_torch_available, - logging, -) - - -if is_torch_available() and is_gguf_available(): - import torch - - from .utils import ( - GGML_QUANT_SIZES, - GGUFParameter, - _dequantize_gguf_and_restore_linear, - _quant_shape_from_byte_shape, - _replace_with_gguf_linear, - ) - - -logger = logging.get_logger(__name__) - - -class GGUFQuantizer(DiffusersQuantizer): - use_keep_in_fp32_modules = True - - def __init__(self, quantization_config, **kwargs): - super().__init__(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 - - if not isinstance(self.modules_to_not_convert, list): - self.modules_to_not_convert = [self.modules_to_not_convert] - - def validate_environment(self, *args, **kwargs): - if not is_accelerate_available() or is_accelerate_version("<", "0.26.0"): - raise ImportError( - "Loading GGUF Parameters requires `accelerate` installed in your environment: `pip install 'accelerate>=0.26.0'`" - ) - if not is_gguf_available() or is_gguf_version("<", "0.10.0"): - raise ImportError( - "To load GGUF format files you must have `gguf` installed in your environment: `pip install gguf>=0.10.0`" - ) - - # 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]]: - # 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 - - def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": - if target_dtype != torch.uint8: - logger.info(f"target_dtype {target_dtype} is replaced by `torch.uint8` for GGUF quantization") - return torch.uint8 - - def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": - if torch_dtype is None: - torch_dtype = self.compute_dtype - return torch_dtype - - def check_quantized_param_shape(self, param_name, current_param, loaded_param): - loaded_param_shape = loaded_param.shape - current_param_shape = current_param.shape - quant_type = loaded_param.quant_type - - block_size, type_size = GGML_QUANT_SIZES[quant_type] - - inferred_shape = _quant_shape_from_byte_shape(loaded_param_shape, type_size, block_size) - if inferred_shape != current_param_shape: - raise ValueError( - f"{param_name} has an expected quantized shape of: {inferred_shape}, but received shape: {loaded_param_shape}" - ) - - return True - - def check_if_quantized_param( - self, - model: "ModelMixin", - param_value: Union["GGUFParameter", "torch.Tensor"], - param_name: str, - state_dict: Dict[str, Any], - **kwargs, - ) -> bool: - if isinstance(param_value, GGUFParameter): - return True - - return False - - def create_quantized_param( - self, - model: "ModelMixin", - param_value: Union["GGUFParameter", "torch.Tensor"], - param_name: str, - target_device: "torch.device", - state_dict: Optional[Dict[str, Any]] = None, - unexpected_keys: Optional[List[str]] = None, - **kwargs, - ): - module, tensor_name = get_module_from_name(model, param_name) - if tensor_name not in module._parameters and tensor_name not in module._buffers: - raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.") - - if tensor_name in module._parameters: - module._parameters[tensor_name] = param_value.to(target_device) - if tensor_name in module._buffers: - module._buffers[tensor_name] = param_value.to(target_device) - - def _process_model_before_weight_loading( - self, - model: "ModelMixin", - device_map, - keep_in_fp32_modules: List[str] = [], - **kwargs, - ): - state_dict = kwargs.get("state_dict", None) - - self.modules_to_not_convert.extend(keep_in_fp32_modules) - self.modules_to_not_convert = [module for module in self.modules_to_not_convert if module is not None] - - _replace_with_gguf_linear( - model, self.compute_dtype, state_dict, modules_to_not_convert=self.modules_to_not_convert - ) - - def _process_model_after_weight_loading(self, model: "ModelMixin", **kwargs): - return model - - @property - def is_serializable(self): - return False - - @property - def is_trainable(self) -> bool: - return False - - @property - def is_compileable(self) -> bool: - return True - - def _dequantize(self, model): - is_model_on_cpu = model.device.type == "cpu" - if is_model_on_cpu: - logger.info( - "Model was found to be on CPU (could happen as a result of `enable_model_cpu_offload()`). So, moving it to accelerator. After dequantization, will move the model back to CPU again to preserve the previous device." - ) - device = ( - torch.accelerator.current_accelerator() - if hasattr(torch, "accelerator") - else torch.cuda.current_device() - ) - model.to(device) - - model = _dequantize_gguf_and_restore_linear(model, self.modules_to_not_convert) - if is_model_on_cpu: - model.to("cpu") - return model diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/utils.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/utils.py deleted file mode 100644 index 41d351712961a1a8ecd1330a97c9341eb896b639..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/gguf/utils.py +++ /dev/null @@ -1,461 +0,0 @@ -# Copyright 2025 The HuggingFace Team and City96. 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 contextlib import nullcontext - -import gguf -import torch -import torch.nn as nn - -from ...utils import is_accelerate_available - - -if is_accelerate_available(): - import accelerate - from accelerate import init_empty_weights - from accelerate.hooks import add_hook_to_module, remove_hook_from_module - - -# Copied from diffusers.quantizers.bitsandbytes.utils._create_accelerate_new_hook -def _create_accelerate_new_hook(old_hook): - r""" - Creates a new hook based on the old hook. Use it only if you know what you are doing ! This method is a copy of: - https://github.com/huggingface/peft/blob/748f7968f3a31ec06a1c2b0328993319ad9a150a/src/peft/utils/other.py#L245 with - some changes - """ - old_hook_cls = getattr(accelerate.hooks, old_hook.__class__.__name__) - old_hook_attr = old_hook.__dict__ - filtered_old_hook_attr = {} - old_hook_init_signature = inspect.signature(old_hook_cls.__init__) - for k in old_hook_attr.keys(): - if k in old_hook_init_signature.parameters: - filtered_old_hook_attr[k] = old_hook_attr[k] - new_hook = old_hook_cls(**filtered_old_hook_attr) - return new_hook - - -def _replace_with_gguf_linear(model, compute_dtype, state_dict, prefix="", modules_to_not_convert=[]): - def _should_convert_to_gguf(state_dict, prefix): - weight_key = prefix + "weight" - return weight_key in state_dict and isinstance(state_dict[weight_key], GGUFParameter) - - has_children = list(model.children()) - if not has_children: - return - - for name, module in model.named_children(): - module_prefix = prefix + name + "." - _replace_with_gguf_linear(module, compute_dtype, state_dict, module_prefix, modules_to_not_convert) - - if ( - isinstance(module, nn.Linear) - and _should_convert_to_gguf(state_dict, module_prefix) - and name not in modules_to_not_convert - ): - ctx = init_empty_weights if is_accelerate_available() else nullcontext - with ctx(): - model._modules[name] = GGUFLinear( - module.in_features, - module.out_features, - module.bias is not None, - compute_dtype=compute_dtype, - ) - model._modules[name].source_cls = type(module) - # Force requires_grad to False to avoid unexpected errors - model._modules[name].requires_grad_(False) - - return model - - -def _dequantize_gguf_and_restore_linear(model, modules_to_not_convert=[]): - for name, module in model.named_children(): - if isinstance(module, GGUFLinear) and name not in modules_to_not_convert: - device = module.weight.device - bias = getattr(module, "bias", None) - - ctx = init_empty_weights if is_accelerate_available() else nullcontext - with ctx(): - new_module = nn.Linear( - module.in_features, - module.out_features, - module.bias is not None, - device=device, - ) - new_module.weight = nn.Parameter(dequantize_gguf_tensor(module.weight)) - if bias is not None: - new_module.bias = bias - - # Create a new hook and attach it in case we use accelerate - if hasattr(module, "_hf_hook"): - old_hook = module._hf_hook - new_hook = _create_accelerate_new_hook(old_hook) - - remove_hook_from_module(module) - add_hook_to_module(new_module, new_hook) - - new_module.to(device) - model._modules[name] = new_module - - has_children = list(module.children()) - if has_children: - _dequantize_gguf_and_restore_linear(module, modules_to_not_convert) - - return model - - -# dequantize operations based on torch ports of GGUF dequantize_functions -# from City96 -# more info: https://github.com/city96/ComfyUI-GGUF/blob/main/dequant.py - - -QK_K = 256 -K_SCALE_SIZE = 12 - - -def to_uint32(x): - x = x.view(torch.uint8).to(torch.int32) - return (x[:, 0] | x[:, 1] << 8 | x[:, 2] << 16 | x[:, 3] << 24).unsqueeze(1) - - -def split_block_dims(blocks, *args): - n_max = blocks.shape[1] - dims = list(args) + [n_max - sum(args)] - return torch.split(blocks, dims, dim=1) - - -def get_scale_min(scales): - n_blocks = scales.shape[0] - scales = scales.view(torch.uint8) - scales = scales.reshape((n_blocks, 3, 4)) - - d, m, m_d = torch.split(scales, scales.shape[-2] // 3, dim=-2) - - sc = torch.cat([d & 0x3F, (m_d & 0x0F) | ((d >> 2) & 0x30)], dim=-1) - min = torch.cat([m & 0x3F, (m_d >> 4) | ((m >> 2) & 0x30)], dim=-1) - - return (sc.reshape((n_blocks, 8)), min.reshape((n_blocks, 8))) - - -def dequantize_blocks_Q8_0(blocks, block_size, type_size, dtype=None): - d, x = split_block_dims(blocks, 2) - d = d.view(torch.float16).to(dtype) - x = x.view(torch.int8) - return d * x - - -def dequantize_blocks_Q5_1(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - d, m, qh, qs = split_block_dims(blocks, 2, 2, 4) - d = d.view(torch.float16).to(dtype) - m = m.view(torch.float16).to(dtype) - qh = to_uint32(qh) - - qh = qh.reshape((n_blocks, 1)) >> torch.arange(32, device=d.device, dtype=torch.int32).reshape(1, 32) - ql = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor( - [0, 4], device=d.device, dtype=torch.uint8 - ).reshape(1, 1, 2, 1) - qh = (qh & 1).to(torch.uint8) - ql = (ql & 0x0F).reshape((n_blocks, -1)) - - qs = ql | (qh << 4) - return (d * qs) + m - - -def dequantize_blocks_Q5_0(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - d, qh, qs = split_block_dims(blocks, 2, 4) - d = d.view(torch.float16).to(dtype) - qh = to_uint32(qh) - - qh = qh.reshape(n_blocks, 1) >> torch.arange(32, device=d.device, dtype=torch.int32).reshape(1, 32) - ql = qs.reshape(n_blocks, -1, 1, block_size // 2) >> torch.tensor( - [0, 4], device=d.device, dtype=torch.uint8 - ).reshape(1, 1, 2, 1) - - qh = (qh & 1).to(torch.uint8) - ql = (ql & 0x0F).reshape(n_blocks, -1) - - qs = (ql | (qh << 4)).to(torch.int8) - 16 - return d * qs - - -def dequantize_blocks_Q4_1(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - d, m, qs = split_block_dims(blocks, 2, 2) - d = d.view(torch.float16).to(dtype) - m = m.view(torch.float16).to(dtype) - - qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor( - [0, 4], device=d.device, dtype=torch.uint8 - ).reshape(1, 1, 2, 1) - qs = (qs & 0x0F).reshape(n_blocks, -1) - - return (d * qs) + m - - -def dequantize_blocks_Q4_0(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - d, qs = split_block_dims(blocks, 2) - d = d.view(torch.float16).to(dtype) - - qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor( - [0, 4], device=d.device, dtype=torch.uint8 - ).reshape((1, 1, 2, 1)) - qs = (qs & 0x0F).reshape((n_blocks, -1)).to(torch.int8) - 8 - return d * qs - - -def dequantize_blocks_Q6_K(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - ( - ql, - qh, - scales, - d, - ) = split_block_dims(blocks, QK_K // 2, QK_K // 4, QK_K // 16) - - scales = scales.view(torch.int8).to(dtype) - d = d.view(torch.float16).to(dtype) - d = (d * scales).reshape((n_blocks, QK_K // 16, 1)) - - ql = ql.reshape((n_blocks, -1, 1, 64)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape( - (1, 1, 2, 1) - ) - ql = (ql & 0x0F).reshape((n_blocks, -1, 32)) - qh = qh.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape( - (1, 1, 4, 1) - ) - qh = (qh & 0x03).reshape((n_blocks, -1, 32)) - q = (ql | (qh << 4)).to(torch.int8) - 32 - q = q.reshape((n_blocks, QK_K // 16, -1)) - - return (d * q).reshape((n_blocks, QK_K)) - - -def dequantize_blocks_Q5_K(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - d, dmin, scales, qh, qs = split_block_dims(blocks, 2, 2, K_SCALE_SIZE, QK_K // 8) - - d = d.view(torch.float16).to(dtype) - dmin = dmin.view(torch.float16).to(dtype) - - sc, m = get_scale_min(scales) - - d = (d * sc).reshape((n_blocks, -1, 1)) - dm = (dmin * m).reshape((n_blocks, -1, 1)) - - ql = qs.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape( - (1, 1, 2, 1) - ) - qh = qh.reshape((n_blocks, -1, 1, 32)) >> torch.arange(0, 8, device=d.device, dtype=torch.uint8).reshape( - (1, 1, 8, 1) - ) - ql = (ql & 0x0F).reshape((n_blocks, -1, 32)) - qh = (qh & 0x01).reshape((n_blocks, -1, 32)) - q = ql | (qh << 4) - - return (d * q - dm).reshape((n_blocks, QK_K)) - - -def dequantize_blocks_Q4_K(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - d, dmin, scales, qs = split_block_dims(blocks, 2, 2, K_SCALE_SIZE) - d = d.view(torch.float16).to(dtype) - dmin = dmin.view(torch.float16).to(dtype) - - sc, m = get_scale_min(scales) - - d = (d * sc).reshape((n_blocks, -1, 1)) - dm = (dmin * m).reshape((n_blocks, -1, 1)) - - qs = qs.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape( - (1, 1, 2, 1) - ) - qs = (qs & 0x0F).reshape((n_blocks, -1, 32)) - - return (d * qs - dm).reshape((n_blocks, QK_K)) - - -def dequantize_blocks_Q3_K(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - hmask, qs, scales, d = split_block_dims(blocks, QK_K // 8, QK_K // 4, 12) - d = d.view(torch.float16).to(dtype) - - lscales, hscales = scales[:, :8], scales[:, 8:] - lscales = lscales.reshape((n_blocks, 1, 8)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape( - (1, 2, 1) - ) - lscales = lscales.reshape((n_blocks, 16)) - hscales = hscales.reshape((n_blocks, 1, 4)) >> torch.tensor( - [0, 2, 4, 6], device=d.device, dtype=torch.uint8 - ).reshape((1, 4, 1)) - hscales = hscales.reshape((n_blocks, 16)) - scales = (lscales & 0x0F) | ((hscales & 0x03) << 4) - scales = scales.to(torch.int8) - 32 - - dl = (d * scales).reshape((n_blocks, 16, 1)) - - ql = qs.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape( - (1, 1, 4, 1) - ) - qh = hmask.reshape(n_blocks, -1, 1, 32) >> torch.arange(0, 8, device=d.device, dtype=torch.uint8).reshape( - (1, 1, 8, 1) - ) - ql = ql.reshape((n_blocks, 16, QK_K // 16)) & 3 - qh = (qh.reshape((n_blocks, 16, QK_K // 16)) & 1) ^ 1 - q = ql.to(torch.int8) - (qh << 2).to(torch.int8) - - return (dl * q).reshape((n_blocks, QK_K)) - - -def dequantize_blocks_Q2_K(blocks, block_size, type_size, dtype=None): - n_blocks = blocks.shape[0] - - scales, qs, d, dmin = split_block_dims(blocks, QK_K // 16, QK_K // 4, 2) - d = d.view(torch.float16).to(dtype) - dmin = dmin.view(torch.float16).to(dtype) - - # (n_blocks, 16, 1) - dl = (d * (scales & 0xF)).reshape((n_blocks, QK_K // 16, 1)) - ml = (dmin * (scales >> 4)).reshape((n_blocks, QK_K // 16, 1)) - - shift = torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape((1, 1, 4, 1)) - - qs = (qs.reshape((n_blocks, -1, 1, 32)) >> shift) & 3 - qs = qs.reshape((n_blocks, QK_K // 16, 16)) - qs = dl * qs - ml - - return qs.reshape((n_blocks, -1)) - - -def dequantize_blocks_BF16(blocks, block_size, type_size, dtype=None): - return (blocks.view(torch.int16).to(torch.int32) << 16).view(torch.float32) - - -GGML_QUANT_SIZES = gguf.GGML_QUANT_SIZES -dequantize_functions = { - gguf.GGMLQuantizationType.BF16: dequantize_blocks_BF16, - gguf.GGMLQuantizationType.Q8_0: dequantize_blocks_Q8_0, - gguf.GGMLQuantizationType.Q5_1: dequantize_blocks_Q5_1, - gguf.GGMLQuantizationType.Q5_0: dequantize_blocks_Q5_0, - gguf.GGMLQuantizationType.Q4_1: dequantize_blocks_Q4_1, - gguf.GGMLQuantizationType.Q4_0: dequantize_blocks_Q4_0, - gguf.GGMLQuantizationType.Q6_K: dequantize_blocks_Q6_K, - gguf.GGMLQuantizationType.Q5_K: dequantize_blocks_Q5_K, - gguf.GGMLQuantizationType.Q4_K: dequantize_blocks_Q4_K, - gguf.GGMLQuantizationType.Q3_K: dequantize_blocks_Q3_K, - gguf.GGMLQuantizationType.Q2_K: dequantize_blocks_Q2_K, -} -SUPPORTED_GGUF_QUANT_TYPES = list(dequantize_functions.keys()) - - -def _quant_shape_from_byte_shape(shape, type_size, block_size): - return (*shape[:-1], shape[-1] // type_size * block_size) - - -def dequantize_gguf_tensor(tensor): - if not hasattr(tensor, "quant_type"): - return tensor - - quant_type = tensor.quant_type - dequant_fn = dequantize_functions[quant_type] - - block_size, type_size = GGML_QUANT_SIZES[quant_type] - - tensor = tensor.view(torch.uint8) - shape = _quant_shape_from_byte_shape(tensor.shape, type_size, block_size) - - n_blocks = tensor.numel() // type_size - blocks = tensor.reshape((n_blocks, type_size)) - - dequant = dequant_fn(blocks, block_size, type_size) - dequant = dequant.reshape(shape) - - return dequant.as_tensor() - - -class GGUFParameter(torch.nn.Parameter): - def __new__(cls, data, requires_grad=False, quant_type=None): - data = data if data is not None else torch.empty(0) - self = torch.Tensor._make_subclass(cls, data, requires_grad) - self.quant_type = quant_type - block_size, type_size = GGML_QUANT_SIZES[quant_type] - self.quant_shape = _quant_shape_from_byte_shape(self.shape, type_size, block_size) - - return self - - def as_tensor(self): - return torch.Tensor._make_subclass(torch.Tensor, self, self.requires_grad) - - @staticmethod - def _extract_quant_type(args): - # When converting from original format checkpoints we often use splits, cats etc on tensors - # this method ensures that the returned tensor type from those operations remains GGUFParameter - # so that we preserve quant_type information - for arg in args: - if isinstance(arg, list) and isinstance(arg[0], GGUFParameter): - return arg[0].quant_type - if isinstance(arg, GGUFParameter): - return arg.quant_type - return None - - @classmethod - def __torch_function__(cls, func, types, args=(), kwargs=None): - if kwargs is None: - kwargs = {} - - result = super().__torch_function__(func, types, args, kwargs) - - if isinstance(result, torch.Tensor): - quant_type = cls._extract_quant_type(args) - return cls(result, quant_type=quant_type) - # Handle tuples and lists - elif type(result) in (list, tuple): - # Preserve the original type (tuple or list) - quant_type = cls._extract_quant_type(args) - wrapped = [cls(x, quant_type=quant_type) if isinstance(x, torch.Tensor) else x for x in result] - return type(result)(wrapped) - else: - return result - - -class GGUFLinear(nn.Linear): - def __init__( - self, - in_features, - out_features, - bias=False, - compute_dtype=None, - device=None, - ) -> None: - super().__init__(in_features, out_features, bias, device) - self.compute_dtype = compute_dtype - - def forward(self, inputs): - weight = dequantize_gguf_tensor(self.weight) - weight = weight.to(self.compute_dtype) - bias = self.bias.to(self.compute_dtype) if self.bias is not None else None - - output = torch.nn.functional.linear(inputs, weight, bias) - return output diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/quantization_config.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/quantization_config.py deleted file mode 100644 index 871faf076e5aa361efd45656c31dbcd96af9a316..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/quantization_config.py +++ /dev/null @@ -1,726 +0,0 @@ -#!/usr/bin/env python -# coding=utf-8 - -# Copyright 2023 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. - -""" -Adapted from -https://github.com/huggingface/transformers/blob/52cb4034ada381fe1ffe8d428a1076e5411a8026/src/transformers/utils/quantization_config.py -""" - -import copy -import importlib.metadata -import inspect -import json -import os -from dataclasses import dataclass -from enum import Enum -from functools import partial -from typing import Any, Dict, List, Optional, Union - -from packaging import version - -from ..utils import is_torch_available, is_torchao_available, logging - - -if is_torch_available(): - import torch - -logger = logging.get_logger(__name__) - - -class QuantizationMethod(str, Enum): - BITS_AND_BYTES = "bitsandbytes" - GGUF = "gguf" - TORCHAO = "torchao" - QUANTO = "quanto" - - -if is_torchao_available(): - from torchao.quantization.quant_primitives import MappingType - - class TorchAoJSONEncoder(json.JSONEncoder): - def default(self, obj): - if isinstance(obj, MappingType): - return obj.name - return super().default(obj) - - -@dataclass -class QuantizationConfigMixin: - """ - Mixin class for quantization config - """ - - quant_method: QuantizationMethod - _exclude_attributes_at_init = [] - - @classmethod - 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]`): - 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]`): - Additional parameters from which to initialize the configuration object. - - Returns: - [`QuantizationConfigMixin`]: The configuration object instantiated from those parameters. - """ - - config = cls(**config_dict) - - to_remove = [] - for key, value in kwargs.items(): - if hasattr(config, key): - setattr(config, key, value) - to_remove.append(key) - for key in to_remove: - kwargs.pop(key, None) - - if return_unused_kwargs: - return config, kwargs - else: - return config - - def to_json_file(self, json_file_path: Union[str, os.PathLike]): - """ - Save this instance to a JSON file. - - Args: - json_file_path (`str` or `os.PathLike`): - Path to the JSON file in which this configuration instance's parameters will be saved. - use_diff (`bool`, *optional*, defaults to `True`): - If set to `True`, only the difference between the config instance and the default - `QuantizationConfig()` is serialized to JSON file. - """ - with open(json_file_path, "w", encoding="utf-8") as writer: - config_dict = self.to_dict() - json_string = json.dumps(config_dict, indent=2, sort_keys=True) + "\n" - - writer.write(json_string) - - 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. - """ - return copy.deepcopy(self.__dict__) - - def __iter__(self): - """allows `dict(obj)` for situations where obj may be a dict or QuantizationConfigMixin""" - for attr, value in copy.deepcopy(self.__dict__).items(): - yield attr, value - - def __repr__(self): - return f"{self.__class__.__name__} {self.to_json_string()}" - - def to_json_string(self, use_diff: bool = True) -> str: - """ - Serializes this instance to a JSON string. - - Args: - use_diff (`bool`, *optional*, defaults to `True`): - If set to `True`, only the difference between the config instance and the default `PretrainedConfig()` - is serialized to JSON string. - - Returns: - `str`: String containing all the attributes that make up this configuration instance in JSON format. - """ - if use_diff is True: - config_dict = self.to_diff_dict() - else: - config_dict = self.to_dict() - return json.dumps(config_dict, indent=2, sort_keys=True) + "\n" - - def update(self, **kwargs): - """ - Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes, - returning all the unused kwargs. - - Args: - 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. - """ - to_remove = [] - for key, value in kwargs.items(): - if hasattr(self, key): - setattr(self, key, value) - to_remove.append(key) - - # Remove all the attributes that were updated, without modifying the input dict - unused_kwargs = {key: value for key, value in kwargs.items() if key not in to_remove} - return unused_kwargs - - -@dataclass -class BitsAndBytesConfig(QuantizationConfigMixin): - """ - This is a wrapper class about all possible attributes and features that you can play with a model that has been - loaded using `bitsandbytes`. - - This replaces `load_in_8bit` or `load_in_4bit` therefore both options are mutually exclusive. - - Currently only supports `LLM.int8()`, `FP4`, and `NF4` quantization. If more methods are added to `bitsandbytes`, - then more arguments will be added to this class. - - Args: - load_in_8bit (`bool`, *optional*, defaults to `False`): - This flag is used to enable 8-bit quantization with LLM.int8(). - load_in_4bit (`bool`, *optional*, defaults to `False`): - This flag is used to enable 4-bit quantization by replacing the Linear layers with FP4/NF4 layers from - `bitsandbytes`. - llm_int8_threshold (`float`, *optional*, defaults to 6.0): - This corresponds to the outlier threshold for outlier detection as described in `LLM.int8() : 8-bit Matrix - Multiplication for Transformers at Scale` paper: https://huggingface.co/papers/2208.07339 Any hidden states - value that is above this threshold will be considered an outlier and the operation on those values will be - done in fp16. Values are usually normally distributed, that is, most values are in the range [-3.5, 3.5], - but there are some exceptional systematic outliers that are very differently distributed for large models. - 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*): - 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`. - llm_int8_enable_fp32_cpu_offload (`bool`, *optional*, defaults to `False`): - This flag is used for advanced use cases and users that are aware of this feature. If you want to split - your model in different parts and run some parts in int8 on GPU and some parts in fp32 on CPU, you can use - this flag. This is useful for offloading large models such as `google/flan-t5-xxl`. Note that the int8 - operations will not be run on CPU. - llm_int8_has_fp16_weight (`bool`, *optional*, defaults to `False`): - This flag runs LLM.int8() with 16-bit main weights. This is useful for fine-tuning as the weights do not - have to be converted back and forth for the backward pass. - bnb_4bit_compute_dtype (`torch.dtype` or str, *optional*, defaults to `torch.float32`): - This sets the computational type which might be different than the input type. For example, inputs might be - fp32, but computation can be set to bf16 for speedups. - bnb_4bit_quant_type (`str`, *optional*, defaults to `"fp4"`): - This sets the quantization data type in the bnb.nn.Linear4Bit layers. Options are FP4 and NF4 data types - which are specified by `fp4` or `nf4`. - bnb_4bit_use_double_quant (`bool`, *optional*, defaults to `False`): - This flag is used for nested quantization where the quantization constants from the first quantization are - 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*): - Additional parameters from which to initialize the configuration object. - """ - - _exclude_attributes_at_init = ["_load_in_4bit", "_load_in_8bit", "quant_method"] - - def __init__( - self, - load_in_8bit=False, - load_in_4bit=False, - llm_int8_threshold=6.0, - llm_int8_skip_modules=None, - llm_int8_enable_fp32_cpu_offload=False, - llm_int8_has_fp16_weight=False, - bnb_4bit_compute_dtype=None, - bnb_4bit_quant_type="fp4", - bnb_4bit_use_double_quant=False, - bnb_4bit_quant_storage=None, - **kwargs, - ): - self.quant_method = QuantizationMethod.BITS_AND_BYTES - - if load_in_4bit and load_in_8bit: - raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time") - - self._load_in_8bit = load_in_8bit - self._load_in_4bit = load_in_4bit - self.llm_int8_threshold = llm_int8_threshold - self.llm_int8_skip_modules = llm_int8_skip_modules - self.llm_int8_enable_fp32_cpu_offload = llm_int8_enable_fp32_cpu_offload - self.llm_int8_has_fp16_weight = llm_int8_has_fp16_weight - self.bnb_4bit_quant_type = bnb_4bit_quant_type - self.bnb_4bit_use_double_quant = bnb_4bit_use_double_quant - - if bnb_4bit_compute_dtype is None: - self.bnb_4bit_compute_dtype = torch.float32 - elif isinstance(bnb_4bit_compute_dtype, str): - self.bnb_4bit_compute_dtype = getattr(torch, bnb_4bit_compute_dtype) - elif isinstance(bnb_4bit_compute_dtype, torch.dtype): - self.bnb_4bit_compute_dtype = bnb_4bit_compute_dtype - else: - raise ValueError("bnb_4bit_compute_dtype must be a string or a torch.dtype") - - if bnb_4bit_quant_storage is None: - self.bnb_4bit_quant_storage = torch.uint8 - elif isinstance(bnb_4bit_quant_storage, str): - if bnb_4bit_quant_storage not in ["float16", "float32", "int8", "uint8", "float64", "bfloat16"]: - raise ValueError( - "`bnb_4bit_quant_storage` must be a valid string (one of 'float16', 'float32', 'int8', 'uint8', 'float64', 'bfloat16') " - ) - self.bnb_4bit_quant_storage = getattr(torch, bnb_4bit_quant_storage) - elif isinstance(bnb_4bit_quant_storage, torch.dtype): - self.bnb_4bit_quant_storage = bnb_4bit_quant_storage - else: - raise ValueError("bnb_4bit_quant_storage must be a string or a torch.dtype") - - if kwargs and not all(k in self._exclude_attributes_at_init for k in kwargs): - logger.warning(f"Unused kwargs: {list(kwargs.keys())}. These kwargs are not used in {self.__class__}.") - - self.post_init() - - @property - def load_in_4bit(self): - return self._load_in_4bit - - @load_in_4bit.setter - def load_in_4bit(self, value: bool): - if not isinstance(value, bool): - raise TypeError("load_in_4bit must be a boolean") - - if self.load_in_8bit and value: - raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time") - self._load_in_4bit = value - - @property - def load_in_8bit(self): - return self._load_in_8bit - - @load_in_8bit.setter - def load_in_8bit(self, value: bool): - if not isinstance(value, bool): - raise TypeError("load_in_8bit must be a boolean") - - if self.load_in_4bit and value: - raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time") - self._load_in_8bit = value - - def post_init(self): - r""" - Safety checker that arguments are correct - also replaces some NoneType arguments with their default values. - """ - if not isinstance(self.load_in_4bit, bool): - raise TypeError("load_in_4bit must be a boolean") - - if not isinstance(self.load_in_8bit, bool): - raise TypeError("load_in_8bit must be a boolean") - - if not isinstance(self.llm_int8_threshold, float): - raise TypeError("llm_int8_threshold must be a float") - - if self.llm_int8_skip_modules is not None and not isinstance(self.llm_int8_skip_modules, list): - raise TypeError("llm_int8_skip_modules must be a list of strings") - if not isinstance(self.llm_int8_enable_fp32_cpu_offload, bool): - raise TypeError("llm_int8_enable_fp32_cpu_offload must be a boolean") - - if not isinstance(self.llm_int8_has_fp16_weight, bool): - raise TypeError("llm_int8_has_fp16_weight must be a boolean") - - if self.bnb_4bit_compute_dtype is not None and not isinstance(self.bnb_4bit_compute_dtype, torch.dtype): - raise TypeError("bnb_4bit_compute_dtype must be torch.dtype") - - if not isinstance(self.bnb_4bit_quant_type, str): - raise TypeError("bnb_4bit_quant_type must be a string") - - if not isinstance(self.bnb_4bit_use_double_quant, bool): - raise TypeError("bnb_4bit_use_double_quant must be a boolean") - - if self.load_in_4bit and not version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse( - "0.39.0" - ): - raise ValueError( - "4 bit quantization requires bitsandbytes>=0.39.0 - please upgrade your bitsandbytes version" - ) - - def is_quantizable(self): - r""" - Returns `True` if the model is quantizable, `False` otherwise. - """ - return self.load_in_8bit or self.load_in_4bit - - def quantization_method(self): - r""" - This method returns the quantization method used for the model. If the model is not quantizable, it returns - `None`. - """ - if self.load_in_8bit: - return "llm_int8" - elif self.load_in_4bit and self.bnb_4bit_quant_type == "fp4": - return "fp4" - elif self.load_in_4bit and self.bnb_4bit_quant_type == "nf4": - return "nf4" - else: - return None - - 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. - """ - output = copy.deepcopy(self.__dict__) - output["bnb_4bit_compute_dtype"] = str(output["bnb_4bit_compute_dtype"]).split(".")[1] - output["bnb_4bit_quant_storage"] = str(output["bnb_4bit_quant_storage"]).split(".")[1] - output["load_in_4bit"] = self.load_in_4bit - output["load_in_8bit"] = self.load_in_8bit - - return output - - 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]: - """ - 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, - """ - config_dict = self.to_dict() - - # get the default config dict - default_config_dict = BitsAndBytesConfig().to_dict() - - serializable_config_dict = {} - - # only serialize values that differ from the default config - for key, value in config_dict.items(): - if value != default_config_dict[key]: - serializable_config_dict[key] = value - - return serializable_config_dict - - -@dataclass -class GGUFQuantizationConfig(QuantizationConfigMixin): - """This is a config class for GGUF Quantization techniques. - - Args: - compute_dtype: (`torch.dtype`, defaults to `torch.float32`): - This sets the computational type which might be different than the input type. For example, inputs might be - fp32, but computation can be set to bf16 for speedups. - - """ - - def __init__(self, compute_dtype: Optional["torch.dtype"] = None): - self.quant_method = QuantizationMethod.GGUF - self.compute_dtype = compute_dtype - self.pre_quantized = True - - # TODO: (Dhruv) Add this as an init argument when we can support loading unquantized checkpoints. - self.modules_to_not_convert = None - - if self.compute_dtype is None: - self.compute_dtype = torch.float32 - - -@dataclass -class TorchAoConfig(QuantizationConfigMixin): - """This is a config class for torchao quantization/sparsity techniques. - - Args: - quant_type (`str`): - The type of quantization we want to use, currently supporting: - - **Integer quantization:** - - Full function names: `int4_weight_only`, `int8_dynamic_activation_int4_weight`, - `int8_weight_only`, `int8_dynamic_activation_int8_weight` - - Shorthands: `int4wo`, `int4dq`, `int8wo`, `int8dq` - - - **Floating point 8-bit quantization:** - - Full function names: `float8_weight_only`, `float8_dynamic_activation_float8_weight`, - `float8_static_activation_float8_weight` - - Shorthands: `float8wo`, `float8wo_e5m2`, `float8wo_e4m3`, `float8dq`, `float8dq_e4m3`, - `float8_e4m3_tensor`, `float8_e4m3_row`, - - - **Floating point X-bit quantization:** - - 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 - be satisfied for a given shorthand notation. - - - **Unsigned Integer quantization:** - - Full function names: `uintx_weight_only` - - Shorthands: `uint1wo`, `uint2wo`, `uint3wo`, `uint4wo`, `uint5wo`, `uint6wo`, `uint7wo` - 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*): - 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 - https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques - - Example: - ```python - from diffusers import FluxTransformer2DModel, TorchAoConfig - - quantization_config = TorchAoConfig("int8wo") - transformer = FluxTransformer2DModel.from_pretrained( - "black-forest-labs/Flux.1-Dev", - subfolder="transformer", - quantization_config=quantization_config, - torch_dtype=torch.bfloat16, - ) - ``` - """ - - def __init__(self, quant_type: str, modules_to_not_convert: Optional[List[str]] = None, **kwargs) -> None: - self.quant_method = QuantizationMethod.TORCHAO - self.quant_type = quant_type - self.modules_to_not_convert = modules_to_not_convert - - # When we load from serialized config, "quant_type_kwargs" will be the key - if "quant_type_kwargs" in kwargs: - self.quant_type_kwargs = kwargs["quant_type_kwargs"] - else: - self.quant_type_kwargs = kwargs - - 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(): - 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()`." - ) - - raise ValueError( - f"Requested quantization type: {self.quant_type} is not supported or is an incorrect `quant_type` name. If you think the " - f"provided quantization type should be supported, please open an issue at https://github.com/huggingface/diffusers/issues." - ) - - method = TORCHAO_QUANT_TYPE_METHODS[self.quant_type] - signature = inspect.signature(method) - all_kwargs = { - param.name - for param in signature.parameters.values() - if param.kind in [inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD] - } - unsupported_kwargs = list(self.quant_type_kwargs.keys() - all_kwargs) - - if len(unsupported_kwargs) > 0: - raise ValueError( - f'The quantization method "{quant_type}" does not support the following keyword arguments: ' - f"{unsupported_kwargs}. The following keywords arguments are supported: {all_kwargs}." - ) - - @classmethod - def _get_torchao_quant_type_to_method(cls): - r""" - Returns supported torchao quantization types with all commonly used notations. - """ - - if is_torchao_available(): - # TODO(aryan): Support autoquant and 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, - int8_weight_only, - uintx_weight_only, - ) - - # TODO(aryan): Add a note on how to use PerAxis and PerGroup observers - from torchao.quantization.observer import PerRow, PerTensor - - def generate_float8dq_types(dtype: torch.dtype): - name = "e5m2" if dtype == torch.float8_e5m2 else "e4m3" - types = {} - - for granularity_cls in [PerTensor, PerRow]: - # Note: Activation and Weights cannot have different granularities - granularity_name = "tensor" if granularity_cls is PerTensor else "row" - types[f"float8dq_{name}_{granularity_name}"] = partial( - float8_dynamic_activation_float8_weight, - activation_dtype=dtype, - weight_dtype=dtype, - granularity=(granularity_cls(), granularity_cls()), - ) - - return types - - def generate_fpx_quantization_types(bits: int): - 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) - - 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 - - INT4_QUANTIZATION_TYPES = { - # int4 weight + bfloat16/float16 activation - "int4wo": int4_weight_only, - "int4_weight_only": int4_weight_only, - # int4 weight + int8 activation - "int4dq": int8_dynamic_activation_int4_weight, - "int8_dynamic_activation_int4_weight": int8_dynamic_activation_int4_weight, - } - - INT8_QUANTIZATION_TYPES = { - # int8 weight + bfloat16/float16 activation - "int8wo": int8_weight_only, - "int8_weight_only": int8_weight_only, - # int8 weight + int8 activation - "int8dq": int8_dynamic_activation_int8_weight, - "int8_dynamic_activation_int8_weight": int8_dynamic_activation_int8_weight, - } - - # TODO(aryan): handle torch 2.2/2.3 - FLOATX_QUANTIZATION_TYPES = { - # float8_e5m2 weight + bfloat16/float16 activation - "float8wo": partial(float8_weight_only, weight_dtype=torch.float8_e5m2), - "float8_weight_only": float8_weight_only, - "float8wo_e5m2": partial(float8_weight_only, weight_dtype=torch.float8_e5m2), - # float8_e4m3 weight + bfloat16/float16 activation - "float8wo_e4m3": partial(float8_weight_only, weight_dtype=torch.float8_e4m3fn), - # float8_e5m2 weight + float8 activation (dynamic) - "float8dq": float8_dynamic_activation_float8_weight, - "float8_dynamic_activation_float8_weight": float8_dynamic_activation_float8_weight, - # ===== Matrix multiplication is not supported in float8_e5m2 so the following errors out. - # However, changing activation_dtype=torch.float8_e4m3 might work here ===== - # "float8dq_e5m2": partial( - # float8_dynamic_activation_float8_weight, - # activation_dtype=torch.float8_e5m2, - # weight_dtype=torch.float8_e5m2, - # ), - # **generate_float8dq_types(torch.float8_e5m2), - # ===== ===== - # float8_e4m3 weight + float8 activation (dynamic) - "float8dq_e4m3": partial( - float8_dynamic_activation_float8_weight, - activation_dtype=torch.float8_e4m3fn, - weight_dtype=torch.float8_e4m3fn, - ), - **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), - } - - UINTX_QUANTIZATION_DTYPES = { - "uintx_weight_only": uintx_weight_only, - "uint1wo": partial(uintx_weight_only, dtype=torch.uint1), - "uint2wo": partial(uintx_weight_only, dtype=torch.uint2), - "uint3wo": partial(uintx_weight_only, dtype=torch.uint3), - "uint4wo": partial(uintx_weight_only, dtype=torch.uint4), - "uint5wo": partial(uintx_weight_only, dtype=torch.uint5), - "uint6wo": partial(uintx_weight_only, dtype=torch.uint6), - "uint7wo": partial(uintx_weight_only, dtype=torch.uint7), - # "uint8wo": partial(uintx_weight_only, dtype=torch.uint8), # uint8 quantization is not supported - } - - QUANTIZATION_TYPES = {} - QUANTIZATION_TYPES.update(INT4_QUANTIZATION_TYPES) - QUANTIZATION_TYPES.update(INT8_QUANTIZATION_TYPES) - QUANTIZATION_TYPES.update(UINTX_QUANTIZATION_DTYPES) - - if cls._is_xpu_or_cuda_capability_atleast_8_9(): - QUANTIZATION_TYPES.update(FLOATX_QUANTIZATION_TYPES) - - return QUANTIZATION_TYPES - else: - raise ValueError( - "TorchAoConfig requires torchao to be installed, please install with `pip install torchao`" - ) - - @staticmethod - def _is_xpu_or_cuda_capability_atleast_8_9() -> bool: - if torch.cuda.is_available(): - major, minor = torch.cuda.get_device_capability() - if major == 8: - return minor >= 9 - return major >= 9 - elif torch.xpu.is_available(): - return True - else: - raise RuntimeError("TorchAO requires a CUDA compatible GPU or Intel XPU and installation of PyTorch.") - - def get_apply_tensor_subclass(self): - TORCHAO_QUANT_TYPE_METHODS = self._get_torchao_quant_type_to_method() - return TORCHAO_QUANT_TYPE_METHODS[self.quant_type](**self.quant_type_kwargs) - - def __repr__(self): - r""" - Example of how this looks for `TorchAoConfig("uint4wo", group_size=32)`: - - ``` - TorchAoConfig { - "modules_to_not_convert": null, - "quant_method": "torchao", - "quant_type": "uint4wo", - "quant_type_kwargs": { - "group_size": 32 - } - } - ``` - """ - config_dict = self.to_dict() - return ( - f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True, cls=TorchAoJSONEncoder)}\n" - ) - - -@dataclass -class QuantoConfig(QuantizationConfigMixin): - """ - This is a wrapper class about all possible attributes and features that you can play with a model that has been - loaded using `quanto`. - - Args: - weights_dtype (`str`, *optional*, defaults to `"int8"`): - The target dtype for the weights after quantization. Supported values are ("float8","int8","int4","int2") - modules_to_not_convert (`list`, *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 (e.g. Whisper encoder, Llava encoder, Mixtral gate layers). - """ - - def __init__( - self, - weights_dtype: str = "int8", - modules_to_not_convert: Optional[List[str]] = None, - **kwargs, - ): - self.quant_method = QuantizationMethod.QUANTO - self.weights_dtype = weights_dtype - self.modules_to_not_convert = modules_to_not_convert - - self.post_init() - - def post_init(self): - r""" - Safety checker that arguments are correct - """ - accepted_weights = ["float8", "int8", "int4", "int2"] - if self.weights_dtype not in accepted_weights: - raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights_dtype}") diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/__init__.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/__init__.py deleted file mode 100644 index a4e8a1f41a1e4bf3bbb53c5684398277eb671a83..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .quanto_quantizer import QuantoQuantizer diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index a78e8241791383429220fa1fa4c88cf6c6f33128..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/__pycache__/quanto_quantizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/__pycache__/quanto_quantizer.cpython-310.pyc deleted file mode 100644 index 5e0acaa6027b9df9c0899f36b5b4ae0436a84e59..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/__pycache__/quanto_quantizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index a919ef2d5b695809dded5bf58862da1ca99abd67..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/quanto_quantizer.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/quanto_quantizer.py deleted file mode 100644 index c5f71f816fc343e39f38cf69f817322312c182d2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/quanto_quantizer.py +++ /dev/null @@ -1,181 +0,0 @@ -from typing import TYPE_CHECKING, Any, Dict, List, Union - -from diffusers.utils.import_utils import is_optimum_quanto_version - -from ...utils import ( - get_module_from_name, - is_accelerate_available, - is_accelerate_version, - is_optimum_quanto_available, - is_torch_available, - logging, -) -from ..base import DiffusersQuantizer - - -if TYPE_CHECKING: - from ...models.modeling_utils import ModelMixin - - -if is_torch_available(): - import torch - -if is_accelerate_available(): - from accelerate.utils import CustomDtype, set_module_tensor_to_device - -if is_optimum_quanto_available(): - from .utils import _replace_with_quanto_layers - -logger = logging.get_logger(__name__) - - -class QuantoQuantizer(DiffusersQuantizer): - r""" - Diffusers Quantizer for Optimum Quanto - """ - - use_keep_in_fp32_modules = True - requires_calibration = False - required_packages = ["quanto", "accelerate"] - - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - - def validate_environment(self, *args, **kwargs): - if not is_optimum_quanto_available(): - raise ImportError( - "Loading an optimum-quanto quantized model requires optimum-quanto library (`pip install optimum-quanto`)" - ) - if not is_optimum_quanto_version(">=", "0.2.6"): - raise ImportError( - "Loading an optimum-quanto quantized model requires `optimum-quanto>=0.2.6`. " - "Please upgrade your installation with `pip install --upgrade optimum-quanto" - ) - - if not is_accelerate_available(): - raise ImportError( - "Loading an optimum-quanto quantized model requires accelerate library (`pip install accelerate`)" - ) - - device_map = kwargs.get("device_map", None) - if isinstance(device_map, dict) and len(device_map.keys()) > 1: - raise ValueError( - "`device_map` for multi-GPU inference or CPU/disk offload is currently not supported with Diffusers and the Quanto backend" - ) - - def check_if_quantized_param( - self, - model: "ModelMixin", - param_value: "torch.Tensor", - param_name: str, - state_dict: Dict[str, Any], - **kwargs, - ): - # Quanto imports diffusers internally. This is here to prevent circular imports - from optimum.quanto import QModuleMixin, QTensor - from optimum.quanto.tensor.packed import PackedTensor - - module, tensor_name = get_module_from_name(model, param_name) - if self.pre_quantized and any(isinstance(module, t) for t in [QTensor, PackedTensor]): - return True - elif isinstance(module, QModuleMixin) and "weight" in tensor_name: - return not module.frozen - - return False - - def create_quantized_param( - self, - model: "ModelMixin", - param_value: "torch.Tensor", - param_name: str, - target_device: "torch.device", - *args, - **kwargs, - ): - """ - Create the quantized parameter by calling .freeze() after setting it to the module. - """ - - dtype = kwargs.get("dtype", torch.float32) - module, tensor_name = get_module_from_name(model, param_name) - if self.pre_quantized: - setattr(module, tensor_name, param_value) - else: - set_module_tensor_to_device(model, param_name, target_device, param_value, dtype) - module.freeze() - module.weight.requires_grad = False - - def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: - max_memory = {key: val * 0.90 for key, val in max_memory.items()} - return max_memory - - def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": - if is_accelerate_version(">=", "0.27.0"): - mapping = { - "int8": torch.int8, - "float8": CustomDtype.FP8, - "int4": CustomDtype.INT4, - "int2": CustomDtype.INT2, - } - target_dtype = mapping[self.quantization_config.weights_dtype] - - return target_dtype - - def update_torch_dtype(self, torch_dtype: "torch.dtype" = None) -> "torch.dtype": - if torch_dtype is None: - logger.info("You did not specify `torch_dtype` in `from_pretrained`. Setting it to `torch.float32`.") - torch_dtype = torch.float32 - return torch_dtype - - 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 - - not_missing_keys = [] - for name, module in model.named_modules(): - if isinstance(module, QModuleMixin): - for missing in missing_keys: - if ( - (name in missing or name in f"{prefix}.{missing}") - and not missing.endswith(".weight") - and not missing.endswith(".bias") - ): - not_missing_keys.append(missing) - return [k for k in missing_keys if k not in not_missing_keys] - - def _process_model_before_weight_loading( - self, - model: "ModelMixin", - device_map, - keep_in_fp32_modules: List[str] = [], - **kwargs, - ): - self.modules_to_not_convert = self.quantization_config.modules_to_not_convert - - if not isinstance(self.modules_to_not_convert, list): - self.modules_to_not_convert = [self.modules_to_not_convert] - - self.modules_to_not_convert.extend(keep_in_fp32_modules) - - model = _replace_with_quanto_layers( - model, - modules_to_not_convert=self.modules_to_not_convert, - quantization_config=self.quantization_config, - pre_quantized=self.pre_quantized, - ) - model.config.quantization_config = self.quantization_config - - def _process_model_after_weight_loading(self, model, **kwargs): - return model - - @property - def is_trainable(self): - return True - - @property - def is_serializable(self): - return True - - @property - def is_compileable(self) -> bool: - return True diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/utils.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/utils.py deleted file mode 100644 index 6f41fd36b43a7c222de17a297523939d27f34537..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/quanto/utils.py +++ /dev/null @@ -1,60 +0,0 @@ -import torch.nn as nn - -from ...utils import is_accelerate_available, logging - - -logger = logging.get_logger(__name__) - -if is_accelerate_available(): - from accelerate import init_empty_weights - - -def _replace_with_quanto_layers(model, quantization_config, modules_to_not_convert: list, pre_quantized=False): - # Quanto imports diffusers internally. These are placed here to avoid circular imports - from optimum.quanto import QLinear, freeze, qfloat8, qint2, qint4, qint8 - - def _get_weight_type(dtype: str): - return {"float8": qfloat8, "int8": qint8, "int4": qint4, "int2": qint2}[dtype] - - def _replace_layers(model, quantization_config, modules_to_not_convert): - has_children = list(model.children()) - if not has_children: - return model - - for name, module in model.named_children(): - _replace_layers(module, quantization_config, modules_to_not_convert) - - if name in modules_to_not_convert: - continue - - if isinstance(module, nn.Linear): - with init_empty_weights(): - qlinear = QLinear( - in_features=module.in_features, - out_features=module.out_features, - bias=module.bias is not None, - dtype=module.weight.dtype, - weights=_get_weight_type(quantization_config.weights_dtype), - ) - model._modules[name] = qlinear - model._modules[name].source_cls = type(module) - model._modules[name].requires_grad_(False) - - return model - - model = _replace_layers(model, quantization_config, modules_to_not_convert) - has_been_replaced = any(isinstance(replaced_module, QLinear) for _, replaced_module in model.named_modules()) - - if not has_been_replaced: - logger.warning( - f"{model.__class__.__name__} does not appear to have any `nn.Linear` modules. Quantization will not be applied." - " Please check your model architecture, or submit an issue on Github if you think this is a bug." - " https://github.com/huggingface/diffusers/issues/new" - ) - - # We need to freeze the pre_quantized model in order for the loaded state_dict and model state dict - # to match when trying to load weights with load_model_dict_into_meta - if pre_quantized: - freeze(model) - - return model diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/torchao/__init__.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/torchao/__init__.py deleted file mode 100644 index c56bf54c25158fb9e7ef9ff3bc1eb651743764fb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/torchao/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -# 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. - -from .torchao_quantizer import TorchAoHfQuantizer diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/torchao/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/torchao/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 76137c06c82c94aafe48447323f50f06cc02a5c5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/torchao/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/torchao/__pycache__/torchao_quantizer.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/quantizers/torchao/__pycache__/torchao_quantizer.cpython-310.pyc deleted file mode 100644 index 1f885277632913c028c21a97af9b7cdbed6e3b13..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/quantizers/torchao/__pycache__/torchao_quantizer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/quantizers/torchao/torchao_quantizer.py b/venv/lib/python3.10/site-packages/diffusers/quantizers/torchao/torchao_quantizer.py deleted file mode 100644 index c12513f061da6764d5b4a1bf31707c7b9573ad8b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/quantizers/torchao/torchao_quantizer.py +++ /dev/null @@ -1,341 +0,0 @@ -# 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. - -""" -Adapted from -https://github.com/huggingface/transformers/blob/3a8eb74668e9c2cc563b2f5c62fac174797063e0/src/transformers/quantizers/quantizer_torchao.py -""" - -import importlib -import types -from typing import TYPE_CHECKING, Any, Dict, List, Union - -from packaging import version - -from ...utils import ( - get_module_from_name, - is_torch_available, - is_torch_version, - is_torchao_available, - is_torchao_version, - logging, -) -from ..base import DiffusersQuantizer - - -if TYPE_CHECKING: - from ...models.modeling_utils import ModelMixin - - -if is_torch_available(): - import torch - import torch.nn as nn - - if is_torch_version(">=", "2.5"): - SUPPORTED_TORCH_DTYPES_FOR_QUANTIZATION = ( - # At the moment, only int8 is supported for integer quantization dtypes. - # In Torch 2.6, int1-int7 will be introduced, so this can be visited in the future - # to support more quantization methods, such as intx_weight_only. - torch.int8, - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.uint1, - torch.uint2, - torch.uint3, - torch.uint4, - torch.uint5, - torch.uint6, - torch.uint7, - ) - else: - SUPPORTED_TORCH_DTYPES_FOR_QUANTIZATION = ( - torch.int8, - torch.float8_e4m3fn, - torch.float8_e5m2, - ) - -if is_torchao_available(): - from torchao.quantization import quantize_ - - -def _update_torch_safe_globals(): - safe_globals = [ - (torch.uint1, "torch.uint1"), - (torch.uint2, "torch.uint2"), - (torch.uint3, "torch.uint3"), - (torch.uint4, "torch.uint4"), - (torch.uint5, "torch.uint5"), - (torch.uint6, "torch.uint6"), - (torch.uint7, "torch.uint7"), - ] - 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]) - - except (ImportError, ModuleNotFoundError) as e: - logger.warning( - "Unable to import `torchao` Tensor objects. This may affect loading checkpoints serialized with `torchao`" - ) - logger.debug(e) - - finally: - torch.serialization.add_safe_globals(safe_globals=safe_globals) - - -if ( - is_torch_available() - and is_torch_version(">=", "2.6.0") - and is_torchao_available() - and is_torchao_version(">=", "0.7.0") -): - _update_torch_safe_globals() - - -logger = logging.get_logger(__name__) - - -def _quantization_type(weight): - from torchao.dtypes import AffineQuantizedTensor - from torchao.quantization.linear_activation_quantized_tensor import LinearActivationQuantizedTensor - - if isinstance(weight, AffineQuantizedTensor): - return f"{weight.__class__.__name__}({weight._quantization_type()})" - - if isinstance(weight, LinearActivationQuantizedTensor): - return f"{weight.__class__.__name__}(activation={weight.input_quant_func}, weight={_quantization_type(weight.original_weight_tensor)})" - - -def _linear_extra_repr(self): - weight = _quantization_type(self.weight) - if weight is None: - return f"in_features={self.weight.shape[1]}, out_features={self.weight.shape[0]}, weight=None" - else: - return f"in_features={self.weight.shape[1]}, out_features={self.weight.shape[0]}, weight={weight}" - - -class TorchAoHfQuantizer(DiffusersQuantizer): - r""" - Diffusers Quantizer for TorchAO: https://github.com/pytorch/ao/. - """ - - requires_calibration = False - required_packages = ["torchao"] - - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - - def validate_environment(self, *args, **kwargs): - if not is_torchao_available(): - raise ImportError( - "Loading a TorchAO quantized model requires the torchao library. Please install with `pip install torchao`" - ) - torchao_version = version.parse(importlib.metadata.version("torch")) - if torchao_version < version.parse("0.7.0"): - raise RuntimeError( - f"The minimum required version of `torchao` is 0.7.0, but the current version is {torchao_version}. Please upgrade with `pip install -U torchao`." - ) - - self.offload = False - - device_map = kwargs.get("device_map", None) - if isinstance(device_map, dict): - if "cpu" in device_map.values() or "disk" in device_map.values(): - if self.pre_quantized: - raise ValueError( - "You are attempting to perform cpu/disk offload with a pre-quantized torchao model " - "This is not supported yet. Please remove the CPU or disk device from the `device_map` argument." - ) - else: - self.offload = True - - if self.pre_quantized: - weights_only = kwargs.get("weights_only", None) - if weights_only: - torch_version = version.parse(importlib.metadata.version("torch")) - if torch_version < version.parse("2.5.0"): - # TODO(aryan): TorchAO is compatible with Pytorch >= 2.2 for certain quantization types. Try to see if we can support it in future - raise RuntimeError( - f"In order to use TorchAO pre-quantized model, you need to have torch>=2.5.0. However, the current version is {torch_version}." - ) - - def update_torch_dtype(self, torch_dtype): - quant_type = self.quantization_config.quant_type - - if quant_type.startswith("int") or quant_type.startswith("uint"): - if torch_dtype is not None and torch_dtype != torch.bfloat16: - logger.warning( - f"You are trying to set torch_dtype to {torch_dtype} for int4/int8/uintx quantization, but " - f"only bfloat16 is supported right now. Please set `torch_dtype=torch.bfloat16`." - ) - - if torch_dtype is None: - # We need to set the torch_dtype, otherwise we have dtype mismatch when performing the quantized linear op - logger.warning( - "Overriding `torch_dtype` with `torch_dtype=torch.bfloat16` due to requirements of `torchao` " - "to enable model loading in different precisions. Pass your own `torch_dtype` to specify the " - "dtype of the remaining non-linear layers, or pass torch_dtype=torch.bfloat16, to remove this warning." - ) - torch_dtype = torch.bfloat16 - - return torch_dtype - - def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": - quant_type = self.quantization_config.quant_type - - if quant_type.startswith("int8") or quant_type.startswith("int4"): - # Note that int4 weights are created by packing into torch.int8, but since there is no torch.int4, we use torch.int8 - return torch.int8 - elif quant_type == "uintx_weight_only": - return self.quantization_config.quant_type_kwargs.get("dtype", torch.uint8) - elif quant_type.startswith("uint"): - return { - 1: torch.uint1, - 2: torch.uint2, - 3: torch.uint3, - 4: torch.uint4, - 5: torch.uint5, - 6: torch.uint6, - 7: torch.uint7, - }[int(quant_type[4])] - elif quant_type.startswith("float") or quant_type.startswith("fp"): - return torch.bfloat16 - - if isinstance(target_dtype, SUPPORTED_TORCH_DTYPES_FOR_QUANTIZATION): - return target_dtype - - # We need one of the supported dtypes to be selected in order for accelerate to determine - # the total size of modules/parameters for auto device placement. - possible_device_maps = ["auto", "balanced", "balanced_low_0", "sequential"] - raise ValueError( - f"You have set `device_map` as one of {possible_device_maps} on a TorchAO quantized model but a suitable target dtype " - f"could not be inferred. The supported target_dtypes are: {SUPPORTED_TORCH_DTYPES_FOR_QUANTIZATION}. If you think the " - 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]]: - max_memory = {key: val * 0.9 for key, val in max_memory.items()} - return max_memory - - def check_if_quantized_param( - self, - model: "ModelMixin", - param_value: "torch.Tensor", - param_name: str, - state_dict: Dict[str, Any], - **kwargs, - ) -> bool: - param_device = kwargs.pop("param_device", None) - # Check if the param_name is not in self.modules_to_not_convert - if any((key + "." in param_name) or (key == param_name) for key in self.modules_to_not_convert): - return False - elif param_device == "cpu" and self.offload: - # We don't quantize weights that we offload - return False - else: - # We only quantize the weight of nn.Linear - module, tensor_name = get_module_from_name(model, param_name) - return isinstance(module, torch.nn.Linear) and (tensor_name == "weight") - - def create_quantized_param( - self, - model: "ModelMixin", - param_value: "torch.Tensor", - param_name: str, - target_device: "torch.device", - state_dict: Dict[str, Any], - unexpected_keys: List[str], - **kwargs, - ): - r""" - Each nn.Linear layer that needs to be quantized is processed here. First, we set the value the weight tensor, - then we move it to the target device. Finally, we quantize the module. - """ - module, tensor_name = get_module_from_name(model, param_name) - - if self.pre_quantized: - # If we're loading pre-quantized weights, replace the repr of linear layers for pretty printing info - # about AffineQuantizedTensor - module._parameters[tensor_name] = torch.nn.Parameter(param_value.to(device=target_device)) - if isinstance(module, nn.Linear): - module.extra_repr = types.MethodType(_linear_extra_repr, module) - else: - # As we perform quantization here, the repr of linear layers is that of AQT, so we don't have to do it ourselves - module._parameters[tensor_name] = torch.nn.Parameter(param_value).to(device=target_device) - quantize_(module, self.quantization_config.get_apply_tensor_subclass()) - - def _process_model_before_weight_loading( - self, - model: "ModelMixin", - device_map, - keep_in_fp32_modules: List[str] = [], - **kwargs, - ): - self.modules_to_not_convert = self.quantization_config.modules_to_not_convert - - if not isinstance(self.modules_to_not_convert, list): - self.modules_to_not_convert = [self.modules_to_not_convert] - - self.modules_to_not_convert.extend(keep_in_fp32_modules) - - # Extend `self.modules_to_not_convert` to keys that are supposed to be offloaded to `cpu` or `disk` - if isinstance(device_map, dict) and len(device_map.keys()) > 1: - keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]] - self.modules_to_not_convert.extend(keys_on_cpu) - - # Purge `None`. - # Unlike `transformers`, we don't know if we should always keep certain modules in FP32 - # in case of diffusion transformer models. For language models and others alike, `lm_head` - # and tied modules are usually kept in FP32. - self.modules_to_not_convert = [module for module in self.modules_to_not_convert if module is not None] - - model.config.quantization_config = self.quantization_config - - def _process_model_after_weight_loading(self, model: "ModelMixin"): - return model - - def is_serializable(self, safe_serialization=None): - # TODO(aryan): needs to be tested - if safe_serialization: - logger.warning( - "torchao quantized model does not support safe serialization, please set `safe_serialization` to False." - ) - return False - - _is_torchao_serializable = version.parse(importlib.metadata.version("huggingface_hub")) >= version.parse( - "0.25.0" - ) - - if not _is_torchao_serializable: - logger.warning("torchao quantized model is only serializable after huggingface_hub >= 0.25.0 ") - - if self.offload and self.quantization_config.modules_to_not_convert is None: - logger.warning( - "The model contains offloaded modules and these modules are not quantized. We don't recommend saving the model as we won't be able to reload them." - "If you want to specify modules to not quantize, please specify modules_to_not_convert in the quantization_config." - ) - return False - - return _is_torchao_serializable - - @property - def is_trainable(self): - return self.quantization_config.quant_type.startswith("int8") - - @property - def is_compileable(self) -> bool: - return True diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__init__.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/__init__.py deleted file mode 100644 index 29052c1ba0cbf404bd446487a39d2eb7ca3abf40..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/__init__.py +++ /dev/null @@ -1,224 +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 TYPE_CHECKING - -from ..utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - get_objects_from_module, - is_flax_available, - is_scipy_available, - is_torch_available, - is_torchsde_available, -) - - -_dummy_modules = {} -_import_structure = {} - -try: - if not is_torch_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import dummy_pt_objects # noqa F403 - - _dummy_modules.update(get_objects_from_module(dummy_pt_objects)) - -else: - _import_structure["deprecated"] = ["KarrasVeScheduler", "ScoreSdeVpScheduler"] - _import_structure["scheduling_amused"] = ["AmusedScheduler"] - _import_structure["scheduling_consistency_decoder"] = ["ConsistencyDecoderScheduler"] - _import_structure["scheduling_consistency_models"] = ["CMStochasticIterativeScheduler"] - _import_structure["scheduling_ddim"] = ["DDIMScheduler"] - _import_structure["scheduling_ddim_cogvideox"] = ["CogVideoXDDIMScheduler"] - _import_structure["scheduling_ddim_inverse"] = ["DDIMInverseScheduler"] - _import_structure["scheduling_ddim_parallel"] = ["DDIMParallelScheduler"] - _import_structure["scheduling_ddpm"] = ["DDPMScheduler"] - _import_structure["scheduling_ddpm_parallel"] = ["DDPMParallelScheduler"] - _import_structure["scheduling_ddpm_wuerstchen"] = ["DDPMWuerstchenScheduler"] - _import_structure["scheduling_deis_multistep"] = ["DEISMultistepScheduler"] - _import_structure["scheduling_dpm_cogvideox"] = ["CogVideoXDPMScheduler"] - _import_structure["scheduling_dpmsolver_multistep"] = ["DPMSolverMultistepScheduler"] - _import_structure["scheduling_dpmsolver_multistep_inverse"] = ["DPMSolverMultistepInverseScheduler"] - _import_structure["scheduling_dpmsolver_singlestep"] = ["DPMSolverSinglestepScheduler"] - _import_structure["scheduling_edm_dpmsolver_multistep"] = ["EDMDPMSolverMultistepScheduler"] - _import_structure["scheduling_edm_euler"] = ["EDMEulerScheduler"] - _import_structure["scheduling_euler_ancestral_discrete"] = ["EulerAncestralDiscreteScheduler"] - _import_structure["scheduling_euler_discrete"] = ["EulerDiscreteScheduler"] - _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_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_pndm"] = ["PNDMScheduler"] - _import_structure["scheduling_repaint"] = ["RePaintScheduler"] - _import_structure["scheduling_sasolver"] = ["SASolverScheduler"] - _import_structure["scheduling_scm"] = ["SCMScheduler"] - _import_structure["scheduling_sde_ve"] = ["ScoreSdeVeScheduler"] - _import_structure["scheduling_tcd"] = ["TCDScheduler"] - _import_structure["scheduling_unclip"] = ["UnCLIPScheduler"] - _import_structure["scheduling_unipc_multistep"] = ["UniPCMultistepScheduler"] - _import_structure["scheduling_utils"] = ["AysSchedules", "KarrasDiffusionSchedulers", "SchedulerMixin"] - _import_structure["scheduling_vq_diffusion"] = ["VQDiffusionScheduler"] - -try: - if not is_flax_available(): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import dummy_flax_objects # noqa F403 - - _dummy_modules.update(get_objects_from_module(dummy_flax_objects)) - -else: - _import_structure["scheduling_ddim_flax"] = ["FlaxDDIMScheduler"] - _import_structure["scheduling_ddpm_flax"] = ["FlaxDDPMScheduler"] - _import_structure["scheduling_dpmsolver_multistep_flax"] = ["FlaxDPMSolverMultistepScheduler"] - _import_structure["scheduling_euler_discrete_flax"] = ["FlaxEulerDiscreteScheduler"] - _import_structure["scheduling_karras_ve_flax"] = ["FlaxKarrasVeScheduler"] - _import_structure["scheduling_lms_discrete_flax"] = ["FlaxLMSDiscreteScheduler"] - _import_structure["scheduling_pndm_flax"] = ["FlaxPNDMScheduler"] - _import_structure["scheduling_sde_ve_flax"] = ["FlaxScoreSdeVeScheduler"] - _import_structure["scheduling_utils_flax"] = [ - "FlaxKarrasDiffusionSchedulers", - "FlaxSchedulerMixin", - "FlaxSchedulerOutput", - "broadcast_to_shape_from_left", - ] - - -try: - if not (is_torch_available() and is_scipy_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import dummy_torch_and_scipy_objects # noqa F403 - - _dummy_modules.update(get_objects_from_module(dummy_torch_and_scipy_objects)) - -else: - _import_structure["scheduling_lms_discrete"] = ["LMSDiscreteScheduler"] - -try: - if not (is_torch_available() and is_torchsde_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import dummy_torch_and_torchsde_objects # noqa F403 - - _dummy_modules.update(get_objects_from_module(dummy_torch_and_torchsde_objects)) - -else: - _import_structure["scheduling_cosine_dpmsolver_multistep"] = ["CosineDPMSolverMultistepScheduler"] - _import_structure["scheduling_dpmsolver_sde"] = ["DPMSolverSDEScheduler"] - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - from ..utils import ( - OptionalDependencyNotAvailable, - is_flax_available, - is_scipy_available, - is_torch_available, - is_torchsde_available, - ) - - try: - if not is_torch_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_pt_objects import * # noqa F403 - else: - from .deprecated import KarrasVeScheduler, ScoreSdeVpScheduler - from .scheduling_amused import AmusedScheduler - from .scheduling_consistency_decoder import ConsistencyDecoderScheduler - from .scheduling_consistency_models import CMStochasticIterativeScheduler - from .scheduling_ddim import DDIMScheduler - from .scheduling_ddim_cogvideox import CogVideoXDDIMScheduler - from .scheduling_ddim_inverse import DDIMInverseScheduler - from .scheduling_ddim_parallel import DDIMParallelScheduler - from .scheduling_ddpm import DDPMScheduler - from .scheduling_ddpm_parallel import DDPMParallelScheduler - from .scheduling_ddpm_wuerstchen import DDPMWuerstchenScheduler - from .scheduling_deis_multistep import DEISMultistepScheduler - from .scheduling_dpm_cogvideox import CogVideoXDPMScheduler - from .scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler - from .scheduling_dpmsolver_multistep_inverse import DPMSolverMultistepInverseScheduler - from .scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler - from .scheduling_edm_dpmsolver_multistep import EDMDPMSolverMultistepScheduler - from .scheduling_edm_euler import EDMEulerScheduler - from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler - from .scheduling_euler_discrete import EulerDiscreteScheduler - 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_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_pndm import PNDMScheduler - from .scheduling_repaint import RePaintScheduler - from .scheduling_sasolver import SASolverScheduler - from .scheduling_scm import SCMScheduler - from .scheduling_sde_ve import ScoreSdeVeScheduler - from .scheduling_tcd import TCDScheduler - from .scheduling_unclip import UnCLIPScheduler - from .scheduling_unipc_multistep import UniPCMultistepScheduler - from .scheduling_utils import AysSchedules, KarrasDiffusionSchedulers, SchedulerMixin - from .scheduling_vq_diffusion import VQDiffusionScheduler - try: - if not is_flax_available(): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_flax_objects import * # noqa F403 - else: - from .scheduling_ddim_flax import FlaxDDIMScheduler - from .scheduling_ddpm_flax import FlaxDDPMScheduler - from .scheduling_dpmsolver_multistep_flax import FlaxDPMSolverMultistepScheduler - from .scheduling_euler_discrete_flax import FlaxEulerDiscreteScheduler - from .scheduling_karras_ve_flax import FlaxKarrasVeScheduler - from .scheduling_lms_discrete_flax import FlaxLMSDiscreteScheduler - from .scheduling_pndm_flax import FlaxPNDMScheduler - from .scheduling_sde_ve_flax import FlaxScoreSdeVeScheduler - from .scheduling_utils_flax import ( - FlaxKarrasDiffusionSchedulers, - FlaxSchedulerMixin, - FlaxSchedulerOutput, - broadcast_to_shape_from_left, - ) - - try: - if not (is_torch_available() and is_scipy_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_torch_and_scipy_objects import * # noqa F403 - else: - from .scheduling_lms_discrete import LMSDiscreteScheduler - - try: - if not (is_torch_available() and is_torchsde_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_torch_and_torchsde_objects import * # noqa F403 - else: - from .scheduling_cosine_dpmsolver_multistep import CosineDPMSolverMultistepScheduler - from .scheduling_dpmsolver_sde import DPMSolverSDEScheduler - -else: - import sys - - sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) - for name, value in _dummy_modules.items(): - setattr(sys.modules[__name__], name, value) diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index b9b78a39f2ede55c1301beccc88f9a9d060797d2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_amused.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_amused.cpython-310.pyc deleted file mode 100644 index 567350aa83374b1deb7006066fa16f8e4dba07de..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_amused.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_consistency_decoder.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_consistency_decoder.cpython-310.pyc deleted file mode 100644 index bfaff3fd8f4f9c7444f1e06c1d5218c76cba9dc0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_consistency_decoder.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_consistency_models.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_consistency_models.cpython-310.pyc deleted file mode 100644 index c3a81a86fc4f3f09b3ec4ebe9741e891f9baeb67..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_consistency_models.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_cosine_dpmsolver_multistep.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_cosine_dpmsolver_multistep.cpython-310.pyc deleted file mode 100644 index 16aa5471e5b42b5bb8f7b9c78b860fa405668a41..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_cosine_dpmsolver_multistep.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim.cpython-310.pyc deleted file mode 100644 index f69b1b7610f8dce8c43a5094d8c7bebe2fd52e53..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim_cogvideox.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim_cogvideox.cpython-310.pyc deleted file mode 100644 index 1ca9c8aa5f61158483f58bd729c90ce6c28d13a8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim_cogvideox.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim_flax.cpython-310.pyc deleted file mode 100644 index 4d228268a19437683060638f44f69ddd5f7f0ec2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim_inverse.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim_inverse.cpython-310.pyc deleted file mode 100644 index ab6fe086697379d69b0caf830473354574b57843..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim_inverse.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim_parallel.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim_parallel.cpython-310.pyc deleted file mode 100644 index 3bbbcbba1699f70d90867d25526edf9a895fb4c5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddim_parallel.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddpm.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddpm.cpython-310.pyc deleted file mode 100644 index 52ffb39d7f83d7ce1368b0487d36b26980f06475..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddpm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddpm_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddpm_flax.cpython-310.pyc deleted file mode 100644 index 15a7f0405963bf3600ea296e754aa5c548bed114..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddpm_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddpm_parallel.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddpm_parallel.cpython-310.pyc deleted file mode 100644 index ad1dc019e9129873fa58886c62ff7ca823432f80..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddpm_parallel.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddpm_wuerstchen.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddpm_wuerstchen.cpython-310.pyc deleted file mode 100644 index ec07a576ea3786e515806b7380b6a9755073695f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ddpm_wuerstchen.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_deis_multistep.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_deis_multistep.cpython-310.pyc deleted file mode 100644 index 50881ed21c2769b3cdccfc63d248df73f6f43a0c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_deis_multistep.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpm_cogvideox.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpm_cogvideox.cpython-310.pyc deleted file mode 100644 index cf57dcbf0b4fa92e9a3d0641e82c4423d9cbe340..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpm_cogvideox.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_multistep.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_multistep.cpython-310.pyc deleted file mode 100644 index 9e59fa9a9b1fc577a0c471dec284426bb4b46e31..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_multistep.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_multistep_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_multistep_flax.cpython-310.pyc deleted file mode 100644 index 8a1a774b92a0acb20336d36438196d6f69bc424d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_multistep_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_multistep_inverse.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_multistep_inverse.cpython-310.pyc deleted file mode 100644 index e0b035a40556a6ede1d20b00f5865e8091d61703..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_multistep_inverse.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_sde.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_sde.cpython-310.pyc deleted file mode 100644 index e5ca3d752b839e57bad7d6b787753c77dfba048b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_sde.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_singlestep.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_singlestep.cpython-310.pyc deleted file mode 100644 index 0b53084acb26195168f05a2b470a18f31a8bca3f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_dpmsolver_singlestep.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_edm_dpmsolver_multistep.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_edm_dpmsolver_multistep.cpython-310.pyc deleted file mode 100644 index 986bf7549984ec412c0d61e49b42b0e549820050..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_edm_dpmsolver_multistep.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_edm_euler.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_edm_euler.cpython-310.pyc deleted file mode 100644 index 802e80927a04075b990b075a19be50feb5d0e28a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_edm_euler.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_euler_ancestral_discrete.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_euler_ancestral_discrete.cpython-310.pyc deleted file mode 100644 index 07a505ab4d15e81c1bfa65fa263235687a17b8a9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_euler_ancestral_discrete.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_euler_discrete.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_euler_discrete.cpython-310.pyc deleted file mode 100644 index 80901e0af3018d5babfab73123bdeda7ee2fcf61..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_euler_discrete.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_euler_discrete_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_euler_discrete_flax.cpython-310.pyc deleted file mode 100644 index 0cdfdd96903c0f61a389ecf4f3c3759411e71479..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_euler_discrete_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_flow_match_euler_discrete.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_flow_match_euler_discrete.cpython-310.pyc deleted file mode 100644 index 71b13ec8ccc49c11a3928bb0dc077643707fd95d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_flow_match_euler_discrete.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_flow_match_heun_discrete.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_flow_match_heun_discrete.cpython-310.pyc deleted file mode 100644 index bea09eb2cfc848e84b2acf105db4d5b20940a2e2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_flow_match_heun_discrete.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_flow_match_lcm.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_flow_match_lcm.cpython-310.pyc deleted file mode 100644 index 41bf0ce8052c617362eb3ef81ddf2a3238045413..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_flow_match_lcm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_heun_discrete.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_heun_discrete.cpython-310.pyc deleted file mode 100644 index 4cacb731d7e5172aa60ca6a3504b521fa5e53621..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_heun_discrete.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ipndm.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ipndm.cpython-310.pyc deleted file mode 100644 index 3047dd755069c2b972c9e3c035820e9b00651f8f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_ipndm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_k_dpm_2_ancestral_discrete.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_k_dpm_2_ancestral_discrete.cpython-310.pyc deleted file mode 100644 index c2745676be1de1bb7f3c5b7a4722768a360f5bc0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_k_dpm_2_ancestral_discrete.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_k_dpm_2_discrete.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_k_dpm_2_discrete.cpython-310.pyc deleted file mode 100644 index 47e52c027458759c25005c66bbb185180c785d46..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_k_dpm_2_discrete.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_karras_ve_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_karras_ve_flax.cpython-310.pyc deleted file mode 100644 index e2dc04f2c8a53adfcd7f19ad79c7f475eccb60bb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_karras_ve_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_lcm.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_lcm.cpython-310.pyc deleted file mode 100644 index 81266fc5d5ce084f74c8aa8b73d207d9ceac829f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_lcm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_lms_discrete.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_lms_discrete.cpython-310.pyc deleted file mode 100644 index 46747e69516458b65ff0f5e24fd5bb4bb9c55fa8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_lms_discrete.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_lms_discrete_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_lms_discrete_flax.cpython-310.pyc deleted file mode 100644 index 4c60232150092456389b554803182ad6b5488dad..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_lms_discrete_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_pndm.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_pndm.cpython-310.pyc deleted file mode 100644 index 26c54b005ebe5d9a556e1c559fbf8dcea6337b5e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_pndm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_pndm_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_pndm_flax.cpython-310.pyc deleted file mode 100644 index d2c3a12e4048cf62e8d083ec51adbaee0c201ca5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_pndm_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_repaint.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_repaint.cpython-310.pyc deleted file mode 100644 index 5a046228d256a25be8f0bc1e901f4cbdda6d34f1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_repaint.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_sasolver.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_sasolver.cpython-310.pyc deleted file mode 100644 index 707a1e509d75ea841f3e6bd379c795cde41ce7e2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_sasolver.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_scm.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_scm.cpython-310.pyc deleted file mode 100644 index fccb6cbbc00138769e41f7343ea4c8db5a78f85d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_scm.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_sde_ve.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_sde_ve.cpython-310.pyc deleted file mode 100644 index dd0d87d197327649ea91300eb55befe37ed26556..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_sde_ve.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_sde_ve_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_sde_ve_flax.cpython-310.pyc deleted file mode 100644 index 6e5cc8ada8b2ac3ccf505765890109aa9cb8af55..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_sde_ve_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_tcd.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_tcd.cpython-310.pyc deleted file mode 100644 index 1e8a6fe522e75c18b23199fc5f6ccb6827015dcf..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_tcd.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_unclip.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_unclip.cpython-310.pyc deleted file mode 100644 index 8a29bb96de4fd08f120cf4636bd537c3cf6c0d1c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_unclip.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_unipc_multistep.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_unipc_multistep.cpython-310.pyc deleted file mode 100644 index 0195c0a00bc21a2fb10b03e6ae08f7355c3a89e6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_unipc_multistep.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_utils.cpython-310.pyc deleted file mode 100644 index 9c52305f1037645e2d061ea71f4043b08ce8bd13..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_utils_flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_utils_flax.cpython-310.pyc deleted file mode 100644 index 9ee12cd7cf5038a2ad345861604ee235f84aad33..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_utils_flax.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_vq_diffusion.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_vq_diffusion.cpython-310.pyc deleted file mode 100644 index ee2f70aa348f984782dcc204587f14db3166df04..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/__pycache__/scheduling_vq_diffusion.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/__init__.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/__init__.py deleted file mode 100644 index 479cf9bd568be320b60e87a831bf80835e9943f2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/__init__.py +++ /dev/null @@ -1,50 +0,0 @@ -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_pt_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_pt_objects)) -else: - _import_structure["scheduling_karras_ve"] = ["KarrasVeScheduler"] - _import_structure["scheduling_sde_vp"] = ["ScoreSdeVpScheduler"] - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - try: - if not is_torch_available(): - raise OptionalDependencyNotAvailable() - - except OptionalDependencyNotAvailable: - from ...utils.dummy_pt_objects import * # noqa F403 - else: - from .scheduling_karras_ve import KarrasVeScheduler - from .scheduling_sde_vp import ScoreSdeVpScheduler - - -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/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index c906eefd486b8f0f6a651343f22d44cd9db3926a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/__pycache__/scheduling_karras_ve.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/__pycache__/scheduling_karras_ve.cpython-310.pyc deleted file mode 100644 index 029273852c36a82b810f3ec4e3d3a038ede03c70..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/__pycache__/scheduling_karras_ve.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/__pycache__/scheduling_sde_vp.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/__pycache__/scheduling_sde_vp.cpython-310.pyc deleted file mode 100644 index 118d0b9e1ff884762ea98545b593ff09fe6a41ba..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/__pycache__/scheduling_sde_vp.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/scheduling_karras_ve.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/scheduling_karras_ve.py deleted file mode 100644 index 6b968e70814556d93ed2d5937a2321342b97f16b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/scheduling_karras_ve.py +++ /dev/null @@ -1,244 +0,0 @@ -# Copyright 2025 NVIDIA 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 dataclasses import dataclass -from typing import Optional, Tuple, Union - -import numpy as np -import torch - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils import BaseOutput -from ...utils.torch_utils import randn_tensor -from ..scheduling_utils import SchedulerMixin - - -@dataclass -class KarrasVeOutput(BaseOutput): - """ - Output class for the scheduler's step function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - derivative (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Derivative of predicted original image sample (x_0). - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample (x_{0}) based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - derivative: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -class KarrasVeScheduler(SchedulerMixin, ConfigMixin): - """ - A stochastic scheduler tailored to variance-expanding models. - - 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. - - - - For more details on the parameters, see [Appendix E](https://huggingface.co/papers/2206.00364). The grid search - values used to find the optimal `{s_noise, s_churn, s_min, s_max}` for a specific model are described in Table 5 of - the paper. - - - - Args: - sigma_min (`float`, defaults to 0.02): - The minimum noise magnitude. - sigma_max (`float`, defaults to 100): - The maximum noise magnitude. - s_noise (`float`, defaults to 1.007): - The amount of additional noise to counteract loss of detail during sampling. A reasonable range is [1.000, - 1.011]. - s_churn (`float`, defaults to 80): - The parameter controlling the overall amount of stochasticity. A reasonable range is [0, 100]. - s_min (`float`, defaults to 0.05): - The start value of the sigma range to add noise (enable stochasticity). A reasonable range is [0, 10]. - s_max (`float`, defaults to 50): - The end value of the sigma range to add noise. A reasonable range is [0.2, 80]. - """ - - order = 2 - - @register_to_config - def __init__( - self, - sigma_min: float = 0.02, - sigma_max: float = 100, - s_noise: float = 1.007, - s_churn: float = 80, - s_min: float = 0.05, - s_max: float = 50, - ): - # standard deviation of the initial noise distribution - self.init_noise_sigma = sigma_max - - # setable values - self.num_inference_steps: int = None - 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: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - def set_timesteps(self, num_inference_steps: int, device: Union[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 (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - """ - self.num_inference_steps = num_inference_steps - timesteps = np.arange(0, self.num_inference_steps)[::-1].copy() - self.timesteps = torch.from_numpy(timesteps).to(device) - schedule = [ - ( - self.config.sigma_max**2 - * (self.config.sigma_min**2 / self.config.sigma_max**2) ** (i / (num_inference_steps - 1)) - ) - for i in self.timesteps - ] - 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]: - """ - 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`. - - Args: - sample (`torch.Tensor`): - The input sample. - sigma (`float`): - generator (`torch.Generator`, *optional*): - A random number generator. - """ - if self.config.s_min <= sigma <= self.config.s_max: - gamma = min(self.config.s_churn / self.num_inference_steps, 2**0.5 - 1) - else: - gamma = 0 - - # sample eps ~ N(0, S_noise^2 * I) - eps = self.config.s_noise * randn_tensor(sample.shape, generator=generator).to(sample.device) - sigma_hat = sigma + gamma * sigma - sample_hat = sample + ((sigma_hat**2 - sigma**2) ** 0.5 * eps) - - return sample_hat, sigma_hat - - def step( - self, - model_output: torch.Tensor, - sigma_hat: float, - sigma_prev: float, - sample_hat: torch.Tensor, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - sigma_hat (`float`): - sigma_prev (`float`): - sample_hat (`torch.Tensor`): - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~schedulers.scheduling_karras_ve.KarrasVESchedulerOutput`] or `tuple`. - - Returns: - [`~schedulers.scheduling_karras_ve.KarrasVESchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_karras_ve.KarrasVESchedulerOutput`] is returned, - otherwise a tuple is returned where the first element is the sample tensor. - - """ - - pred_original_sample = sample_hat + sigma_hat * model_output - derivative = (sample_hat - pred_original_sample) / sigma_hat - sample_prev = sample_hat + (sigma_prev - sigma_hat) * derivative - - if not return_dict: - return (sample_prev, derivative) - - return KarrasVeOutput( - prev_sample=sample_prev, derivative=derivative, pred_original_sample=pred_original_sample - ) - - def step_correct( - self, - model_output: torch.Tensor, - sigma_hat: float, - sigma_prev: float, - sample_hat: torch.Tensor, - sample_prev: torch.Tensor, - derivative: torch.Tensor, - return_dict: bool = True, - ) -> Union[KarrasVeOutput, Tuple]: - """ - Corrects the predicted sample based on the `model_output` of the network. - - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - sigma_hat (`float`): TODO - sigma_prev (`float`): TODO - sample_hat (`torch.Tensor`): TODO - sample_prev (`torch.Tensor`): TODO - derivative (`torch.Tensor`): TODO - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~schedulers.scheduling_ddpm.DDPMSchedulerOutput`] or `tuple`. - - Returns: - prev_sample (TODO): updated sample in the diffusion chain. derivative (TODO): TODO - - """ - pred_original_sample = sample_prev + sigma_prev * model_output - derivative_corr = (sample_prev - pred_original_sample) / sigma_prev - sample_prev = sample_hat + (sigma_prev - sigma_hat) * (0.5 * derivative + 0.5 * derivative_corr) - - if not return_dict: - return (sample_prev, derivative) - - return KarrasVeOutput( - prev_sample=sample_prev, derivative=derivative, pred_original_sample=pred_original_sample - ) - - def add_noise(self, original_samples, noise, timesteps): - raise NotImplementedError() diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/scheduling_sde_vp.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/scheduling_sde_vp.py deleted file mode 100644 index 5088bdb4976182468cbe64a3b378974c23654888..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/deprecated/scheduling_sde_vp.py +++ /dev/null @@ -1,109 +0,0 @@ -# Copyright 2025 Google Brain 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. - -# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch - -import math -from typing import Union - -import torch - -from ...configuration_utils import ConfigMixin, register_to_config -from ...utils.torch_utils import randn_tensor -from ..scheduling_utils import SchedulerMixin - - -class ScoreSdeVpScheduler(SchedulerMixin, ConfigMixin): - """ - `ScoreSdeVpScheduler` is a variance preserving stochastic differential equation (SDE) scheduler. - - 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`, defaults to 2000): - The number of diffusion steps to train the model. - beta_min (`int`, defaults to 0.1): - beta_max (`int`, defaults to 20): - sampling_eps (`int`, defaults to 1e-3): - The end value of sampling where timesteps decrease progressively from 1 to epsilon. - """ - - order = 1 - - @register_to_config - def __init__(self, num_train_timesteps=2000, beta_min=0.1, beta_max=20, sampling_eps=1e-3): - self.sigmas = None - self.discrete_sigmas = None - self.timesteps = None - - def set_timesteps(self, num_inference_steps, device: Union[str, torch.device] = None): - """ - Sets the continuous 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 timesteps are not moved. - """ - self.timesteps = torch.linspace(1, self.config.sampling_eps, num_inference_steps, device=device) - - def step_pred(self, score, x, t, generator=None): - """ - 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: - score (): - x (): - t (): - generator (`torch.Generator`, *optional*): - A random number generator. - """ - if self.timesteps is None: - raise ValueError( - "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler" - ) - - # TODO(Patrick) better comments + non-PyTorch - # postprocess model score - log_mean_coeff = -0.25 * t**2 * (self.config.beta_max - self.config.beta_min) - 0.5 * t * self.config.beta_min - std = torch.sqrt(1.0 - torch.exp(2.0 * log_mean_coeff)) - std = std.flatten() - while len(std.shape) < len(score.shape): - std = std.unsqueeze(-1) - score = -score / std - - # compute - dt = -1.0 / len(self.timesteps) - - beta_t = self.config.beta_min + t * (self.config.beta_max - self.config.beta_min) - beta_t = beta_t.flatten() - while len(beta_t.shape) < len(x.shape): - beta_t = beta_t.unsqueeze(-1) - drift = -0.5 * beta_t * x - - diffusion = torch.sqrt(beta_t) - drift = drift - diffusion**2 * score - x_mean = x + drift * dt - - # add noise - noise = randn_tensor(x.shape, layout=x.layout, generator=generator, device=x.device, dtype=x.dtype) - x = x_mean + diffusion * math.sqrt(-dt) * noise - - return x, x_mean - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_amused.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_amused.py deleted file mode 100644 index 238b8d8691715a83105df815c0ebcad748092ec1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_amused.py +++ /dev/null @@ -1,162 +0,0 @@ -import math -from dataclasses import dataclass -from typing import List, Optional, Tuple, Union - -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput -from .scheduling_utils import SchedulerMixin - - -def gumbel_noise(t, generator=None): - device = generator.device if generator is not None else t.device - noise = torch.zeros_like(t, device=device).uniform_(0, 1, generator=generator).to(t.device) - return -torch.log((-torch.log(noise.clamp(1e-20))).clamp(1e-20)) - - -def mask_by_random_topk(mask_len, probs, temperature=1.0, generator=None): - confidence = torch.log(probs.clamp(1e-20)) + temperature * gumbel_noise(probs, generator=generator) - sorted_confidence = torch.sort(confidence, dim=-1).values - cut_off = torch.gather(sorted_confidence, 1, mask_len.long()) - masking = confidence < cut_off - return masking - - -@dataclass -class AmusedSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: torch.Tensor = None - - -class AmusedScheduler(SchedulerMixin, ConfigMixin): - order = 1 - - temperatures: torch.Tensor - - @register_to_config - def __init__( - self, - mask_token_id: int, - masking_schedule: str = "cosine", - ): - self.temperatures = None - self.timesteps = None - - def set_timesteps( - self, - num_inference_steps: int, - temperature: Union[int, Tuple[int, int], List[int]] = (2, 0), - device: Union[str, torch.device] = None, - ): - self.timesteps = torch.arange(num_inference_steps, device=device).flip(0) - - if isinstance(temperature, (tuple, list)): - self.temperatures = torch.linspace(temperature[0], temperature[1], num_inference_steps, device=device) - else: - self.temperatures = torch.linspace(temperature, 0.01, num_inference_steps, device=device) - - def step( - self, - model_output: torch.Tensor, - timestep: torch.long, - sample: torch.LongTensor, - starting_mask_ratio: int = 1, - generator: Optional[torch.Generator] = None, - return_dict: bool = True, - ) -> Union[AmusedSchedulerOutput, Tuple]: - two_dim_input = sample.ndim == 3 and model_output.ndim == 4 - - if two_dim_input: - batch_size, codebook_size, height, width = model_output.shape - sample = sample.reshape(batch_size, height * width) - model_output = model_output.reshape(batch_size, codebook_size, height * width).permute(0, 2, 1) - - unknown_map = sample == self.config.mask_token_id - - probs = model_output.softmax(dim=-1) - - device = probs.device - probs_ = probs.to(generator.device) if generator is not None else probs # handles when generator is on CPU - if probs_.device.type == "cpu" and probs_.dtype != torch.float32: - probs_ = probs_.float() # multinomial is not implemented for cpu half precision - probs_ = probs_.reshape(-1, probs.size(-1)) - pred_original_sample = torch.multinomial(probs_, 1, generator=generator).to(device=device) - pred_original_sample = pred_original_sample[:, 0].view(*probs.shape[:-1]) - pred_original_sample = torch.where(unknown_map, pred_original_sample, sample) - - if timestep == 0: - prev_sample = pred_original_sample - else: - seq_len = sample.shape[1] - step_idx = (self.timesteps == timestep).nonzero() - ratio = (step_idx + 1) / len(self.timesteps) - - if self.config.masking_schedule == "cosine": - mask_ratio = torch.cos(ratio * math.pi / 2) - elif self.config.masking_schedule == "linear": - mask_ratio = 1 - ratio - else: - raise ValueError(f"unknown masking schedule {self.config.masking_schedule}") - - mask_ratio = starting_mask_ratio * mask_ratio - - mask_len = (seq_len * mask_ratio).floor() - # do not mask more than amount previously masked - mask_len = torch.min(unknown_map.sum(dim=-1, keepdim=True) - 1, mask_len) - # mask at least one - mask_len = torch.max(torch.tensor([1], device=model_output.device), mask_len) - - selected_probs = torch.gather(probs, -1, pred_original_sample[:, :, None])[:, :, 0] - # Ignores the tokens given in the input by overwriting their confidence. - selected_probs = torch.where(unknown_map, selected_probs, torch.finfo(selected_probs.dtype).max) - - masking = mask_by_random_topk(mask_len, selected_probs, self.temperatures[step_idx], generator) - - # Masks tokens with lower confidence. - prev_sample = torch.where(masking, self.config.mask_token_id, pred_original_sample) - - if two_dim_input: - prev_sample = prev_sample.reshape(batch_size, height, width) - pred_original_sample = pred_original_sample.reshape(batch_size, height, width) - - if not return_dict: - return (prev_sample, pred_original_sample) - - return AmusedSchedulerOutput(prev_sample, pred_original_sample) - - def add_noise(self, sample, timesteps, generator=None): - step_idx = (self.timesteps == timesteps).nonzero() - ratio = (step_idx + 1) / len(self.timesteps) - - if self.config.masking_schedule == "cosine": - mask_ratio = torch.cos(ratio * math.pi / 2) - elif self.config.masking_schedule == "linear": - mask_ratio = 1 - ratio - else: - raise ValueError(f"unknown masking schedule {self.config.masking_schedule}") - - mask_indices = ( - torch.rand( - sample.shape, device=generator.device if generator is not None else sample.device, generator=generator - ).to(sample.device) - < mask_ratio - ) - - masked_sample = sample.clone() - - masked_sample[mask_indices] = self.config.mask_token_id - - return masked_sample diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_consistency_decoder.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_consistency_decoder.py deleted file mode 100644 index d7af018b284a59878c3cae86cec57e02d344e227..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_consistency_decoder.py +++ /dev/null @@ -1,180 +0,0 @@ -import math -from dataclasses import dataclass -from typing import Optional, Tuple, Union - -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import SchedulerMixin - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -@dataclass -class ConsistencyDecoderSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - """ - - prev_sample: torch.Tensor - - -class ConsistencyDecoderScheduler(SchedulerMixin, ConfigMixin): - order = 1 - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1024, - sigma_data: float = 0.5, - ): - betas = betas_for_alpha_bar(num_train_timesteps) - - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - - self.sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod) - self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod) - - sigmas = torch.sqrt(1.0 / alphas_cumprod - 1) - - sqrt_recip_alphas_cumprod = torch.sqrt(1.0 / alphas_cumprod) - - self.c_skip = sqrt_recip_alphas_cumprod * sigma_data**2 / (sigmas**2 + sigma_data**2) - self.c_out = sigmas * sigma_data / (sigmas**2 + sigma_data**2) ** 0.5 - self.c_in = sqrt_recip_alphas_cumprod / (sigmas**2 + sigma_data**2) ** 0.5 - - def set_timesteps( - self, - num_inference_steps: Optional[int] = None, - device: Union[str, torch.device] = None, - ): - if num_inference_steps != 2: - raise ValueError("Currently more than 2 inference steps are not supported.") - - self.timesteps = torch.tensor([1008, 512], dtype=torch.long, device=device) - self.sqrt_alphas_cumprod = self.sqrt_alphas_cumprod.to(device) - self.sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod.to(device) - self.c_skip = self.c_skip.to(device) - self.c_out = self.c_out.to(device) - self.c_in = self.c_in.to(device) - - @property - def init_noise_sigma(self): - return self.sqrt_one_minus_alphas_cumprod[self.timesteps[0]] - - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample * self.c_in[timestep] - - def step( - self, - model_output: torch.Tensor, - timestep: Union[float, torch.Tensor], - sample: torch.Tensor, - generator: Optional[torch.Generator] = None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. - timestep (`float`): - 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`): - Whether or not to return a - [`~schedulers.scheduling_consistency_models.ConsistencyDecoderSchedulerOutput`] or `tuple`. - - Returns: - [`~schedulers.scheduling_consistency_models.ConsistencyDecoderSchedulerOutput`] or `tuple`: - If return_dict is `True`, - [`~schedulers.scheduling_consistency_models.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 - - timestep_idx = torch.where(self.timesteps == timestep)[0] - - if timestep_idx == len(self.timesteps) - 1: - prev_sample = x_0 - else: - noise = randn_tensor(x_0.shape, generator=generator, dtype=x_0.dtype, device=x_0.device) - prev_sample = ( - self.sqrt_alphas_cumprod[self.timesteps[timestep_idx + 1]].to(x_0.dtype) * x_0 - + self.sqrt_one_minus_alphas_cumprod[self.timesteps[timestep_idx + 1]].to(x_0.dtype) * noise - ) - - if not return_dict: - return (prev_sample,) - - return ConsistencyDecoderSchedulerOutput(prev_sample=prev_sample) diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_consistency_models.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_consistency_models.py deleted file mode 100644 index 0f506225880091a9d3635d35de0dc9aaca1a245b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_consistency_models.py +++ /dev/null @@ -1,445 +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 dataclasses import dataclass -from typing import List, Optional, Tuple, Union - -import numpy as np -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 CMStochasticIterativeSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - """ - - prev_sample: torch.Tensor - - -class CMStochasticIterativeScheduler(SchedulerMixin, ConfigMixin): - """ - Multistep and onestep sampling for consistency models. - - 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`, defaults to 40): - The number of diffusion steps to train the model. - sigma_min (`float`, defaults to 0.002): - Minimum noise magnitude in the sigma schedule. Defaults to 0.002 from the original implementation. - sigma_max (`float`, defaults to 80.0): - Maximum noise magnitude in the sigma schedule. Defaults to 80.0 from the original implementation. - sigma_data (`float`, defaults to 0.5): - The standard deviation of the data distribution from the EDM - [paper](https://huggingface.co/papers/2206.00364). Defaults to 0.5 from the original implementation. - s_noise (`float`, defaults to 1.0): - The amount of additional noise to counteract loss of detail during sampling. A reasonable range is [1.000, - 1.011]. Defaults to 1.0 from the original implementation. - rho (`float`, defaults to 7.0): - The parameter for calculating the Karras sigma schedule from the EDM - [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*): - An explicit timestep schedule that can be optionally specified. The timesteps are expected to be in - increasing order. - """ - - order = 1 - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 40, - sigma_min: float = 0.002, - sigma_max: float = 80.0, - sigma_data: float = 0.5, - s_noise: float = 1.0, - rho: float = 7.0, - clip_denoised: bool = True, - ): - # standard deviation of the initial noise distribution - self.init_noise_sigma = sigma_max - - ramp = np.linspace(0, 1, num_train_timesteps) - sigmas = self._convert_to_karras(ramp) - timesteps = self.sigma_to_t(sigmas) - - # setable values - self.num_inference_steps = None - self.sigmas = torch.from_numpy(sigmas) - self.timesteps = torch.from_numpy(timesteps) - self.custom_timesteps = False - self.is_scale_input_called = False - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: - """ - Scales the consistency model input by `(sigma**2 + sigma_data**2) ** 0.5`. - - Args: - sample (`torch.Tensor`): - The input sample. - timestep (`float` or `torch.Tensor`): - The current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - # Get sigma corresponding to timestep - if self.step_index is None: - self._init_step_index(timestep) - - sigma = self.sigmas[self.step_index] - - sample = sample / ((sigma**2 + self.config.sigma_data**2) ** 0.5) - - self.is_scale_input_called = True - return sample - - def sigma_to_t(self, sigmas: Union[float, np.ndarray]): - """ - Gets scaled timesteps from the Karras sigmas for input to the consistency model. - - Args: - sigmas (`float` or `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. - """ - if not isinstance(sigmas, np.ndarray): - sigmas = np.array(sigmas, dtype=np.float64) - - timesteps = 1000 * 0.25 * np.log(sigmas + 1e-44) - - return timesteps - - def set_timesteps( - self, - num_inference_steps: Optional[int] = None, - device: Union[str, torch.device] = None, - timesteps: Optional[List[int]] = None, - ): - """ - Sets the 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 timesteps are not moved. - 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`. - """ - if num_inference_steps is None and timesteps is None: - raise ValueError("Exactly one of `num_inference_steps` or `timesteps` must be supplied.") - - if num_inference_steps is not None and timesteps is not None: - raise ValueError("Can only pass one of `num_inference_steps` or `timesteps`.") - - # Follow DDPMScheduler custom timesteps logic - if timesteps is not None: - for i in range(1, len(timesteps)): - if timesteps[i] >= timesteps[i - 1]: - raise ValueError("`timesteps` must be in descending order.") - - if timesteps[0] >= self.config.num_train_timesteps: - raise ValueError( - f"`timesteps` must start before `self.config.train_timesteps`: {self.config.num_train_timesteps}." - ) - - timesteps = np.array(timesteps, dtype=np.int64) - self.custom_timesteps = True - else: - if num_inference_steps > self.config.num_train_timesteps: - raise ValueError( - f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:" - f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle" - f" maximal {self.config.num_train_timesteps} timesteps." - ) - - self.num_inference_steps = num_inference_steps - - step_ratio = self.config.num_train_timesteps // self.num_inference_steps - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64) - self.custom_timesteps = False - - # Map timesteps to Karras sigmas directly for multistep sampling - # See https://github.com/openai/consistency_models/blob/main/cm/karras_diffusion.py#L675 - num_train_timesteps = self.config.num_train_timesteps - ramp = timesteps[::-1].copy() - ramp = ramp / (num_train_timesteps - 1) - sigmas = self._convert_to_karras(ramp) - timesteps = self.sigma_to_t(sigmas) - - sigmas = np.concatenate([sigmas, [self.config.sigma_min]]).astype(np.float32) - self.sigmas = torch.from_numpy(sigmas).to(device=device) - - if str(device).startswith("mps"): - # mps does not support float64 - self.timesteps = torch.from_numpy(timesteps).to(device, dtype=torch.float32) - else: - self.timesteps = torch.from_numpy(timesteps).to(device=device) - - self._step_index = None - self._begin_index = None - 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).""" - - sigma_min: float = self.config.sigma_min - sigma_max: float = self.config.sigma_max - - rho = self.config.rho - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - def get_scalings(self, sigma): - 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): - """ - 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. - - - - `epsilon` in the equations for `c_skip` and `c_out` is set to `sigma_min`. - - - - Args: - sigma (`torch.Tensor`): - The current sigma in the Karras sigma schedule. - - Returns: - `tuple`: - 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. - """ - sigma_min = self.config.sigma_min - sigma_data = self.config.sigma_data - - c_skip = sigma_data**2 / ((sigma - sigma_min) ** 2 + sigma_data**2) - c_out = (sigma - sigma_min) * sigma_data / (sigma**2 + sigma_data**2) ** 0.5 - return c_skip, c_out - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep - 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() - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - 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( - self, - model_output: torch.Tensor, - timestep: Union[float, torch.Tensor], - sample: torch.Tensor, - generator: Optional[torch.Generator] = None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. - timestep (`float`): - 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`): - Whether or not to return a - [`~schedulers.scheduling_consistency_models.CMStochasticIterativeSchedulerOutput`] or `tuple`. - - Returns: - [`~schedulers.scheduling_consistency_models.CMStochasticIterativeSchedulerOutput`] or `tuple`: - If return_dict is `True`, - [`~schedulers.scheduling_consistency_models.CMStochasticIterativeSchedulerOutput`] is returned, - otherwise a tuple is returned where the first element is the sample tensor. - """ - - if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)): - raise ValueError( - ( - "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" - f" `{self.__class__}.step()` is not supported. Make sure to pass" - " one of the `scheduler.timesteps` as a timestep." - ), - ) - - if not self.is_scale_input_called: - logger.warning( - "The `scale_model_input` function should be called before `step` to ensure correct denoising. " - "See `StableDiffusionPipeline` for a usage example." - ) - - sigma_min = self.config.sigma_min - sigma_max = self.config.sigma_max - - if self.step_index is None: - self._init_step_index(timestep) - - # sigma_next corresponds to next_t in original implementation - sigma = self.sigmas[self.step_index] - if self.step_index + 1 < self.config.num_train_timesteps: - sigma_next = self.sigmas[self.step_index + 1] - else: - # Set sigma_next to sigma_min - sigma_next = self.sigmas[-1] - - # Get scalings for boundary conditions - c_skip, c_out = self.get_scalings_for_boundary_condition(sigma) - - # 1. Denoise model output using boundary conditions - denoised = c_out * model_output + c_skip * sample - if self.config.clip_denoised: - denoised = denoised.clamp(-1, 1) - - # 2. Sample z ~ N(0, s_noise^2 * I) - # 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 - ) - else: - noise = torch.zeros_like(model_output) - z = noise * self.config.s_noise - - sigma_hat = sigma_next.clamp(min=sigma_min, max=sigma_max) - - # 3. Return noisy sample - # tau = sigma_hat, eps = sigma_min - prev_sample = denoised + z * (sigma_hat**2 - sigma_min**2) ** 0.5 - - # upon completion increase step index by one - self._step_index += 1 - - if not return_dict: - return (prev_sample,) - - return CMStochasticIterativeSchedulerOutput(prev_sample=prev_sample) - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.Tensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - noisy_samples = original_samples + noise * sigma - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py deleted file mode 100644 index 66ed296da8ea89d5cfd3f240d698cf640d08a808..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py +++ /dev/null @@ -1,577 +0,0 @@ -# Copyright 2025 TSAIL 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. - -# 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 - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from .scheduling_dpmsolver_sde import BrownianTreeNoiseSampler -from .scheduling_utils import SchedulerMixin, SchedulerOutput - - -class CosineDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): - """ - Implements a variant of `DPMSolverMultistepScheduler` with cosine schedule, proposed by Nichol and Dhariwal (2021). - This scheduler was used in Stable Audio Open [1]. - - [1] Evans, Parker, et al. "Stable Audio Open" https://huggingface.co/papers/2407.14358 - - 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: - sigma_min (`float`, *optional*, 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): - 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): - 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): - The number of diffusion steps to train the model. - 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 - Video](https://imagen.research.google/video/paper.pdf) 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. - 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. - euler_at_final (`bool`, defaults to `False`): - Whether to use Euler's method in the final step. It is a trade-off between numerical stability and detail - 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. - """ - - _compatibles = [] - order = 1 - - @register_to_config - def __init__( - self, - sigma_min: float = 0.3, - sigma_max: float = 500, - sigma_data: float = 1.0, - sigma_schedule: str = "exponential", - num_train_timesteps: int = 1000, - solver_order: int = 2, - prediction_type: str = "v_prediction", - rho: float = 7.0, - solver_type: str = "midpoint", - lower_order_final: bool = True, - euler_at_final: bool = False, - final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min" - ): - if solver_type not in ["midpoint", "heun"]: - if solver_type in ["logrho", "bh1", "bh2"]: - self.register_to_config(solver_type="midpoint") - else: - raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}") - - ramp = torch.linspace(0, 1, num_train_timesteps) - if sigma_schedule == "karras": - sigmas = self._compute_karras_sigmas(ramp) - elif sigma_schedule == "exponential": - sigmas = self._compute_exponential_sigmas(ramp) - - self.timesteps = self.precondition_noise(sigmas) - - self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) - - # setable values - self.num_inference_steps = None - self.model_outputs = [None] * solver_order - self.lower_order_nums = 0 - self._step_index = None - self._begin_index = None - 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 - return (self.config.sigma_max**2 + 1) ** 0.5 - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 - - # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_inputs - def precondition_inputs(self, sample, sigma): - c_in = self._get_conditioning_c_in(sigma) - scaled_sample = sample * c_in - return scaled_sample - - def precondition_noise(self, sigma): - 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): - sigma_data = self.config.sigma_data - c_skip = sigma_data**2 / (sigma**2 + sigma_data**2) - - if self.config.prediction_type == "epsilon": - c_out = sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5 - elif self.config.prediction_type == "v_prediction": - c_out = -sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5 - else: - raise ValueError(f"Prediction type {self.config.prediction_type} is not supported.") - - denoised = c_skip * sample + c_out * model_output - - 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: - """ - 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. - - Args: - sample (`torch.Tensor`): - The input sample. - timestep (`int`, *optional*): - The current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - if self.step_index is None: - self._init_step_index(timestep) - - sigma = self.sigmas[self.step_index] - sample = self.precondition_inputs(sample, sigma) - - self.is_scale_input_called = True - return sample - - def set_timesteps(self, num_inference_steps: int = None, device: Union[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 (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - """ - - self.num_inference_steps = num_inference_steps - - ramp = torch.linspace(0, 1, self.num_inference_steps) - if self.config.sigma_schedule == "karras": - sigmas = self._compute_karras_sigmas(ramp) - elif self.config.sigma_schedule == "exponential": - sigmas = self._compute_exponential_sigmas(ramp) - - sigmas = sigmas.to(dtype=torch.float32, device=device) - self.timesteps = self.precondition_noise(sigmas) - - if self.config.final_sigmas_type == "sigma_min": - sigma_last = self.config.sigma_min - elif self.config.final_sigmas_type == "zero": - sigma_last = 0 - else: - raise ValueError( - f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" - ) - - self.sigmas = torch.cat([sigmas, torch.tensor([sigma_last], dtype=torch.float32, device=device)]) - - self.model_outputs = [ - None, - ] * self.config.solver_order - self.lower_order_nums = 0 - - # add an index counter for schedulers that allow duplicated timesteps - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - # if a noise sampler is used, reinitialise it - 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).""" - sigma_min = sigma_min or self.config.sigma_min - sigma_max = sigma_max or self.config.sigma_max - - rho = self.config.rho - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - 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. - - https://github.com/crowsonkb/k-diffusion/blob/6ab5146d4a5ef63901326489f31f1d8e7dd36b48/k_diffusion/sampling.py#L26 - """ - sigma_min = sigma_min or self.config.sigma_min - sigma_max = sigma_max or self.config.sigma_max - sigmas = torch.linspace(math.log(sigma_min), math.log(sigma_max), len(ramp)).exp().flip(0) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma - log_sigma = np.log(np.maximum(sigma, 1e-10)) - - # get distribution - dists = log_sigma - log_sigmas[:, np.newaxis] - - # get sigmas range - low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) - high_idx = low_idx + 1 - - low = log_sigmas[low_idx] - high = log_sigmas[high_idx] - - # interpolate sigmas - w = (low - log_sigma) / (low - high) - w = np.clip(w, 0, 1) - - # transform interpolation to time range - t = (1 - w) * low_idx + w * high_idx - t = t.reshape(sigma.shape) - return t - - 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, - ) -> torch.Tensor: - """ - Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is - designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an - integral of the data prediction model. - - - - The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise - prediction and data prediction models. - - - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. - sample (`torch.Tensor`): - A current instance of a sample created by the diffusion process. - - Returns: - `torch.Tensor`: - The converted model output. - """ - sigma = self.sigmas[self.step_index] - x0_pred = self.precondition_outputs(sample, model_output, sigma) - - return x0_pred - - def dpm_solver_first_order_update( - self, - model_output: torch.Tensor, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """ - One step for the first-order DPMSolver (equivalent to DDIM). - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. - sample (`torch.Tensor`): - A current instance of a sample created by the diffusion process. - - 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] - 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) - lambda_s = torch.log(alpha_s) - torch.log(sigma_s) - - h = lambda_t - lambda_s - assert noise is not None - x_t = ( - (sigma_t / sigma_s * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - - return x_t - - def multistep_dpm_solver_second_order_update( - self, - model_output_list: List[torch.Tensor], - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """ - One step for the second-order multistep DPMSolver. - - Args: - 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. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - sigma_t, sigma_s0, sigma_s1 = ( - self.sigmas[self.step_index + 1], - self.sigmas[self.step_index], - self.sigmas[self.step_index - 1], - ) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1) - - m0, m1 = model_output_list[-1], model_output_list[-2] - - h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1 - r0 = h_0 / h - D0, D1 = m0, (1.0 / r0) * (m0 - m1) - - # sde-dpmsolver++ - assert noise is not None - if self.config.solver_type == "midpoint": - x_t = ( - (sigma_t / sigma_s0 * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 - + 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1 - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - elif self.config.solver_type == "heun": - x_t = ( - (sigma_t / sigma_s0 * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 - + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1 - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - - return x_t - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep - def index_for_timestep(self, timestep, schedule_timesteps=None): - if schedule_timesteps is None: - schedule_timesteps = self.timesteps - - index_candidates = (schedule_timesteps == timestep).nonzero() - - if len(index_candidates) == 0: - step_index = len(self.timesteps) - 1 - # 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) - elif len(index_candidates) > 1: - step_index = index_candidates[1].item() - else: - step_index = index_candidates[0].item() - - return step_index - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index - def _init_step_index(self, timestep): - """ - Initialize the step_index counter for the scheduler. - """ - - 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( - self, - model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], - sample: torch.Tensor, - generator=None, - return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: - """ - Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with - the multistep DPMSolver. - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - generator (`torch.Generator`, *optional*): - A random number generator. - return_dict (`bool`): - 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( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - if self.step_index is None: - self._init_step_index(timestep) - - # Improve numerical stability for small number of steps - lower_order_final = (self.step_index == len(self.timesteps) - 1) and ( - self.config.euler_at_final - or (self.config.lower_order_final and len(self.timesteps) < 15) - or self.config.final_sigmas_type == "zero" - ) - lower_order_second = ( - (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15 - ) - - model_output = self.convert_model_output(model_output, sample=sample) - for i in range(self.config.solver_order - 1): - self.model_outputs[i] = self.model_outputs[i + 1] - self.model_outputs[-1] = model_output - - if self.noise_sampler is None: - seed = None - if generator is not None: - seed = ( - [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 - ) - noise = self.noise_sampler(self.sigmas[self.step_index], self.sigmas[self.step_index + 1]).to( - model_output.device - ) - - if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final: - prev_sample = self.dpm_solver_first_order_update(model_output, sample=sample, noise=noise) - elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second: - prev_sample = self.multistep_dpm_solver_second_order_update(self.model_outputs, sample=sample, noise=noise) - - if self.lower_order_nums < self.config.solver_order: - self.lower_order_nums += 1 - - # upon completion increase step index by one - self._step_index += 1 - - if not return_dict: - return (prev_sample,) - - return SchedulerOutput(prev_sample=prev_sample) - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.Tensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - noisy_samples = original_samples + noise * sigma - return noisy_samples - - # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._get_conditioning_c_in - def _get_conditioning_c_in(self, sigma): - c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5) - return c_in - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim.py deleted file mode 100644 index 5ee0d084f060920820b686056ecca4b48fef543e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim.py +++ /dev/null @@ -1,521 +0,0 @@ -# Copyright 2025 Stanford University 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. - -# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion -# and https://github.com/hojonathanho/diffusion - -import math -from dataclasses import dataclass -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM -class DDIMSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -def rescale_zero_terminal_snr(betas): - """ - 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. - - Returns: - `torch.Tensor`: rescaled betas with zero terminal SNR - """ - # Convert betas to alphas_bar_sqrt - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - alphas_bar_sqrt = alphas_cumprod.sqrt() - - # Store old values. - alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() - alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() - - # Shift so the last timestep is zero. - alphas_bar_sqrt -= alphas_bar_sqrt_T - - # Scale so the first timestep is back to the old value. - alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - - # Convert alphas_bar_sqrt to betas - alphas_bar = alphas_bar_sqrt**2 # Revert sqrt - alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod - alphas = torch.cat([alphas_bar[0:1], alphas]) - betas = 1 - alphas - - return betas - - -class DDIMScheduler(SchedulerMixin, ConfigMixin): - """ - `DDIMScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with - non-Markovian guidance. - - 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`, 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`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - clip_sample (`bool`, defaults to `True`): - Clip the predicted sample for numerical stability. - clip_sample_range (`float`, defaults to 1.0): - The maximum magnitude for sample clipping. Valid only when `clip_sample=True`. - set_alpha_to_one (`bool`, defaults to `True`): - Each diffusion step uses the alphas product value at that step and at the previous one. For the final step - there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, - otherwise it uses the alpha value at step 0. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - 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://imagen.research.google/video/paper.pdf) paper). - thresholding (`bool`, defaults to `False`): - Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such - as Stable Diffusion. - dynamic_thresholding_ratio (`float`, defaults to 0.995): - The ratio for the dynamic thresholding method. Valid only when `thresholding=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 the [Common Diffusion Noise Schedules and - Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. - 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). - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - 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, - clip_sample: bool = True, - set_alpha_to_one: bool = True, - steps_offset: int = 0, - prediction_type: str = "epsilon", - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - clip_sample_range: float = 1.0, - sample_max_value: float = 1.0, - timestep_spacing: str = "leading", - rescale_betas_zero_snr: bool = False, - ): - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - # Rescale for zero SNR - if rescale_betas_zero_snr: - self.betas = rescale_zero_terminal_snr(self.betas) - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - # At every step in ddim, we are looking into the previous alphas_cumprod - # For the final step, there is no previous alphas_cumprod because we are already at 0 - # `set_alpha_to_one` decides whether we set this parameter simply to one or - # whether we use the final alpha of the "non-previous" one. - self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0] - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # setable values - 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: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - def _get_variance(self, timestep, prev_timestep): - 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 - beta_prod_t_prev = 1 - alpha_prod_t_prev - - variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev) - - return variance - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - """ - "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the - prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by - s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing - pixels from saturation at each step. We find that dynamic thresholding results in significantly better - photorealism as well as better image-text alignment, especially when using very large guidance weights." - - https://huggingface.co/papers/2205.11487 - """ - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half - - # Flatten sample for doing quantile calculation along each image - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - - abs_sample = sample.abs() # "a certain percentile absolute pixel value" - - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] - s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 - sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" - - sample = sample.reshape(batch_size, channels, *remaining_dims) - sample = sample.to(dtype) - - return sample - - def set_timesteps(self, num_inference_steps: int, device: Union[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. - """ - - if num_inference_steps > self.config.num_train_timesteps: - raise ValueError( - f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:" - f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle" - f" maximal {self.config.num_train_timesteps} timesteps." - ) - - self.num_inference_steps = num_inference_steps - - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = ( - np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps) - .round()[::-1] - .copy() - .astype(np.int64) - ) - elif self.config.timestep_spacing == "leading": - step_ratio = self.config.num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64) - timesteps -= 1 - else: - raise ValueError( - f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'leading' or 'trailing'." - ) - - self.timesteps = torch.from_numpy(timesteps).to(device) - - def step( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - eta: float = 0.0, - use_clipped_model_output: bool = False, - generator=None, - variance_noise: Optional[torch.Tensor] = None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - timestep (`float`): - 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. - generator (`torch.Generator`, *optional*): - A random number generator. - 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.DDIMSchedulerOutput`] or `tuple`. - - Returns: - [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] is returned, otherwise a - tuple is returned where the first element is the sample tensor. - - """ - 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" - ) - - # See formulas (12) and (16) of DDIM paper https://huggingface.co/papers/2010.02502 - # Ideally, read DDIM paper in-detail understanding - - # Notation ( -> - # - pred_noise_t -> e_theta(x_t, t) - # - pred_original_sample -> f_theta(x_t, t) or x_0 - # - std_dev_t -> sigma_t - # - eta -> η - # - pred_sample_direction -> "direction pointing to x_t" - # - pred_prev_sample -> "x_t-1" - - # 1. get previous step value (=t-1) - prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps - - # 2. compute alphas, betas - 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 - - # 3. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (12) from https://huggingface.co/papers/2010.02502 - if self.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - pred_epsilon = model_output - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) - elif self.config.prediction_type == "v_prediction": - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" - " `v_prediction`" - ) - - # 4. Clip or threshold "predicted x_0" - if self.config.thresholding: - pred_original_sample = self._threshold_sample(pred_original_sample) - elif self.config.clip_sample: - pred_original_sample = pred_original_sample.clamp( - -self.config.clip_sample_range, self.config.clip_sample_range - ) - - # 5. compute variance: "sigma_t(η)" -> see formula (16) - # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1) - variance = self._get_variance(timestep, prev_timestep) - std_dev_t = eta * variance ** (0.5) - - if use_clipped_model_output: - # the pred_epsilon is always re-derived from the clipped x_0 in Glide - pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) - - # 6. compute "direction pointing to x_t" of formula (12) from https://huggingface.co/papers/2010.02502 - pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * pred_epsilon - - # 7. compute x_t without "random noise" of formula (12) from https://huggingface.co/papers/2010.02502 - prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction - - if eta > 0: - if variance_noise is not None and generator is not None: - raise ValueError( - "Cannot pass both generator and variance_noise. Please make sure that either `generator` or" - " `variance_noise` stays `None`." - ) - - if variance_noise is None: - variance_noise = randn_tensor( - model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype - ) - variance = std_dev_t * variance_noise - - prev_sample = prev_sample + variance - - if not return_dict: - return ( - prev_sample, - pred_original_sample, - ) - - return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as original_samples - # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement - # for the subsequent add_noise calls - self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype) - timesteps = timesteps.to(original_samples.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(original_samples.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise - return noisy_samples - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity - def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as sample - self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype) - timesteps = timesteps.to(sample.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(sample.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample - return velocity - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim_cogvideox.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim_cogvideox.py deleted file mode 100644 index c19efdc7834dadd2c55d32f5c29e7e05f6d2b27f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim_cogvideox.py +++ /dev/null @@ -1,452 +0,0 @@ -# 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. - -# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion -# and https://github.com/hojonathanho/diffusion - -import math -from dataclasses import dataclass -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM -class DDIMSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -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. - - Returns: - `torch.Tensor`: rescaled betas with zero terminal SNR - """ - - alphas_bar_sqrt = alphas_cumprod.sqrt() - - # Store old values. - alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() - alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() - - # Shift so the last timestep is zero. - alphas_bar_sqrt -= alphas_bar_sqrt_T - - # Scale so the first timestep is back to the old value. - alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - - # Convert alphas_bar_sqrt to betas - alphas_bar = alphas_bar_sqrt**2 # Revert sqrt - - return alphas_bar - - -class CogVideoXDDIMScheduler(SchedulerMixin, ConfigMixin): - """ - `DDIMScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with - non-Markovian guidance. - - 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`, 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`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - clip_sample (`bool`, defaults to `True`): - Clip the predicted sample for numerical stability. - clip_sample_range (`float`, defaults to 1.0): - The maximum magnitude for sample clipping. Valid only when `clip_sample=True`. - set_alpha_to_one (`bool`, defaults to `True`): - Each diffusion step uses the alphas product value at that step and at the previous one. For the final step - there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, - otherwise it uses the alpha value at step 0. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - 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://imagen.research.google/video/paper.pdf) paper). - thresholding (`bool`, defaults to `False`): - Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such - as Stable Diffusion. - dynamic_thresholding_ratio (`float`, defaults to 0.995): - The ratio for the dynamic thresholding method. Valid only when `thresholding=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 the [Common Diffusion Noise Schedules and - Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. - 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). - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - 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, - clip_sample: bool = True, - set_alpha_to_one: bool = True, - steps_offset: int = 0, - prediction_type: str = "epsilon", - clip_sample_range: float = 1.0, - sample_max_value: float = 1.0, - timestep_spacing: str = "leading", - rescale_betas_zero_snr: bool = False, - snr_shift_scale: float = 3.0, - ): - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - # Modify: SNR shift following SD3 - self.alphas_cumprod = self.alphas_cumprod / (snr_shift_scale + (1 - snr_shift_scale) * self.alphas_cumprod) - - # Rescale for zero SNR - if rescale_betas_zero_snr: - self.alphas_cumprod = rescale_zero_terminal_snr(self.alphas_cumprod) - - # At every step in ddim, we are looking into the previous alphas_cumprod - # For the final step, there is no previous alphas_cumprod because we are already at 0 - # `set_alpha_to_one` decides whether we set this parameter simply to one or - # whether we use the final alpha of the "non-previous" one. - self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0] - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # setable values - 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): - 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 - beta_prod_t_prev = 1 - alpha_prod_t_prev - - variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev) - - return variance - - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - def set_timesteps(self, num_inference_steps: int, device: Union[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. - """ - - if num_inference_steps > self.config.num_train_timesteps: - raise ValueError( - f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:" - f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle" - f" maximal {self.config.num_train_timesteps} timesteps." - ) - - self.num_inference_steps = num_inference_steps - - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = ( - np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps) - .round()[::-1] - .copy() - .astype(np.int64) - ) - elif self.config.timestep_spacing == "leading": - step_ratio = self.config.num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64) - timesteps -= 1 - else: - raise ValueError( - f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'leading' or 'trailing'." - ) - - self.timesteps = torch.from_numpy(timesteps).to(device) - - def step( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - eta: float = 0.0, - use_clipped_model_output: bool = False, - generator=None, - variance_noise: Optional[torch.Tensor] = None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - timestep (`float`): - 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. - generator (`torch.Generator`, *optional*): - A random number generator. - 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.DDIMSchedulerOutput`] or `tuple`. - - Returns: - [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] is returned, otherwise a - tuple is returned where the first element is the sample tensor. - - """ - 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" - ) - - # See formulas (12) and (16) of DDIM paper https://huggingface.co/papers/2010.02502 - # Ideally, read DDIM paper in-detail understanding - - # Notation ( -> - # - pred_noise_t -> e_theta(x_t, t) - # - pred_original_sample -> f_theta(x_t, t) or x_0 - # - std_dev_t -> sigma_t - # - eta -> η - # - pred_sample_direction -> "direction pointing to x_t" - # - pred_prev_sample -> "x_t-1" - - # 1. get previous step value (=t-1) - prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps - - # 2. compute alphas, betas - 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 - - # 3. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (12) from https://huggingface.co/papers/2010.02502 - # To make style tests pass, commented out `pred_epsilon` as it is an unused variable - if self.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - # pred_epsilon = model_output - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - # pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) - elif self.config.prediction_type == "v_prediction": - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - # pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" - " `v_prediction`" - ) - - a_t = ((1 - alpha_prod_t_prev) / (1 - alpha_prod_t)) ** 0.5 - b_t = alpha_prod_t_prev**0.5 - alpha_prod_t**0.5 * a_t - - prev_sample = a_t * sample + b_t * pred_original_sample - - if not return_dict: - return ( - prev_sample, - pred_original_sample, - ) - - return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as original_samples - # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement - # for the subsequent add_noise calls - self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype) - timesteps = timesteps.to(original_samples.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(original_samples.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise - return noisy_samples - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity - def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as sample - self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype) - timesteps = timesteps.to(sample.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(sample.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample - return velocity - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim_flax.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim_flax.py deleted file mode 100644 index 802d8f79779de52c39ffd1210571518c74e37b02..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim_flax.py +++ /dev/null @@ -1,314 +0,0 @@ -# Copyright 2025 Stanford University 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. - -# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion -# 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 .scheduling_utils_flax import ( - CommonSchedulerState, - FlaxKarrasDiffusionSchedulers, - FlaxSchedulerMixin, - FlaxSchedulerOutput, - add_noise_common, - get_velocity_common, -) - - -@flax.struct.dataclass -class DDIMSchedulerState: - common: CommonSchedulerState - final_alpha_cumprod: jnp.ndarray - - # setable values - init_noise_sigma: jnp.ndarray - timesteps: jnp.ndarray - num_inference_steps: Optional[int] = None - - @classmethod - def create( - cls, - common: CommonSchedulerState, - final_alpha_cumprod: jnp.ndarray, - init_noise_sigma: jnp.ndarray, - timesteps: jnp.ndarray, - ): - return cls( - common=common, - final_alpha_cumprod=final_alpha_cumprod, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - ) - - -@dataclass -class FlaxDDIMSchedulerOutput(FlaxSchedulerOutput): - state: DDIMSchedulerState - - -class FlaxDDIMScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - Denoising diffusion implicit models is a scheduler that extends the denoising procedure introduced in denoising - diffusion probabilistic models (DDPMs) with non-Markovian guidance. - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - For more details, see the original paper: https://huggingface.co/papers/2010.02502 - - 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 - `linear`, `scaled_linear`, or `squaredcos_cap_v2`. - trained_betas (`jnp.ndarray`, optional): - option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. - clip_sample (`bool`, default `True`): - option to clip predicted sample between for numerical stability. The clip range is determined by - `clip_sample_range`. - clip_sample_range (`float`, default `1.0`): - the maximum magnitude for sample clipping. Valid only when `clip_sample=True`. - set_alpha_to_one (`bool`, default `True`): - each diffusion step uses the value of alphas product at that step and at the previous one. For the final - step there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, - otherwise it uses the value of alpha at step 0. - steps_offset (`int`, default `0`): - An offset added to the inference steps, as required by some model families. - prediction_type (`str`, default `epsilon`): - indicates whether the model predicts the noise (epsilon), or the samples. One of `epsilon`, `sample`. - `v-prediction` is not supported for this scheduler. - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - the `dtype` used for params and computation. - """ - - _compatibles = [e.name for e in FlaxKarrasDiffusionSchedulers] - - dtype: jnp.dtype - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: Optional[jnp.ndarray] = None, - clip_sample: bool = True, - clip_sample_range: float = 1.0, - set_alpha_to_one: bool = True, - steps_offset: int = 0, - prediction_type: str = "epsilon", - dtype: jnp.dtype = jnp.float32, - ): - self.dtype = dtype - - def create_state(self, common: Optional[CommonSchedulerState] = None) -> DDIMSchedulerState: - if common is None: - common = CommonSchedulerState.create(self) - - # At every step in ddim, we are looking into the previous alphas_cumprod - # For the final step, there is no previous alphas_cumprod because we are already at 0 - # `set_alpha_to_one` decides whether we set this parameter simply to one or - # whether we use the final alpha of the "non-previous" one. - final_alpha_cumprod = ( - jnp.array(1.0, dtype=self.dtype) if self.config.set_alpha_to_one else common.alphas_cumprod[0] - ) - - # standard deviation of the initial noise distribution - init_noise_sigma = jnp.array(1.0, dtype=self.dtype) - - timesteps = jnp.arange(0, self.config.num_train_timesteps).round()[::-1] - - return DDIMSchedulerState.create( - common=common, - final_alpha_cumprod=final_alpha_cumprod, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - ) - - def scale_model_input( - self, state: DDIMSchedulerState, sample: jnp.ndarray, timestep: Optional[int] = None - ) -> jnp.ndarray: - """ - Args: - state (`PNDMSchedulerState`): the `FlaxPNDMScheduler` state data class instance. - sample (`jnp.ndarray`): input sample - timestep (`int`, optional): current timestep - - Returns: - `jnp.ndarray`: scaled input sample - """ - return sample - - def set_timesteps( - 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. - - Args: - state (`DDIMSchedulerState`): - the `FlaxDDIMScheduler` state data class instance. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - """ - step_ratio = self.config.num_train_timesteps // num_inference_steps - # creates integer timesteps by multiplying by ratio - # rounding to avoid issues when num_inference_step is power of 3 - timesteps = (jnp.arange(0, num_inference_steps) * step_ratio).round()[::-1] + self.config.steps_offset - - return state.replace( - num_inference_steps=num_inference_steps, - timesteps=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 - ) - beta_prod_t = 1 - alpha_prod_t - beta_prod_t_prev = 1 - alpha_prod_t_prev - - variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev) - - return variance - - def step( - self, - state: DDIMSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - eta: float = 0.0, - return_dict: bool = True, - ) -> Union[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). - - Args: - state (`DDIMSchedulerState`): the `FlaxDDIMScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - return_dict (`bool`): option for returning tuple rather than FlaxDDIMSchedulerOutput class - - Returns: - [`FlaxDDIMSchedulerOutput`] or `tuple`: [`FlaxDDIMSchedulerOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - if state.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - # See formulas (12) and (16) of DDIM paper https://huggingface.co/papers/2010.02502 - # Ideally, read DDIM paper in-detail understanding - - # Notation ( -> - # - pred_noise_t -> e_theta(x_t, t) - # - pred_original_sample -> f_theta(x_t, t) or x_0 - # - std_dev_t -> sigma_t - # - eta -> η - # - pred_sample_direction -> "direction pointing to x_t" - # - pred_prev_sample -> "x_t-1" - - # 1. get previous step value (=t-1) - prev_timestep = timestep - self.config.num_train_timesteps // state.num_inference_steps - - alphas_cumprod = state.common.alphas_cumprod - final_alpha_cumprod = state.final_alpha_cumprod - - # 2. compute alphas, betas - alpha_prod_t = alphas_cumprod[timestep] - alpha_prod_t_prev = jnp.where(prev_timestep >= 0, alphas_cumprod[prev_timestep], final_alpha_cumprod) - - beta_prod_t = 1 - alpha_prod_t - - # 3. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (12) from https://huggingface.co/papers/2010.02502 - if self.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - pred_epsilon = model_output - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) - elif self.config.prediction_type == "v_prediction": - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" - " `v_prediction`" - ) - - # 4. Clip or threshold "predicted x_0" - if self.config.clip_sample: - pred_original_sample = pred_original_sample.clip( - -self.config.clip_sample_range, self.config.clip_sample_range - ) - - # 4. compute variance: "sigma_t(η)" -> see formula (16) - # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1) - variance = self._get_variance(state, timestep, prev_timestep) - std_dev_t = eta * variance ** (0.5) - - # 5. compute "direction pointing to x_t" of formula (12) from https://huggingface.co/papers/2010.02502 - pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * pred_epsilon - - # 6. compute x_t without "random noise" of formula (12) from https://huggingface.co/papers/2010.02502 - prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction - - if not return_dict: - return (prev_sample, state) - - return FlaxDDIMSchedulerOutput(prev_sample=prev_sample, state=state) - - def add_noise( - self, - state: DDIMSchedulerState, - original_samples: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - return add_noise_common(state.common, original_samples, noise, timesteps) - - def get_velocity( - self, - state: DDIMSchedulerState, - sample: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - return get_velocity_common(state.common, sample, noise, timesteps) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim_inverse.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim_inverse.py deleted file mode 100644 index 49dba840d08982d18225c05585e407433153327a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim_inverse.py +++ /dev/null @@ -1,374 +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. - -# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion -# and https://github.com/hojonathanho/diffusion -import math -from dataclasses import dataclass -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch - -from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.schedulers.scheduling_utils import SchedulerMixin -from diffusers.utils import BaseOutput, deprecate - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM -class DDIMSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr -def rescale_zero_terminal_snr(betas): - """ - 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. - - Returns: - `torch.Tensor`: rescaled betas with zero terminal SNR - """ - # Convert betas to alphas_bar_sqrt - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - alphas_bar_sqrt = alphas_cumprod.sqrt() - - # Store old values. - alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() - alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() - - # Shift so the last timestep is zero. - alphas_bar_sqrt -= alphas_bar_sqrt_T - - # Scale so the first timestep is back to the old value. - alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - - # Convert alphas_bar_sqrt to betas - alphas_bar = alphas_bar_sqrt**2 # Revert sqrt - alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod - alphas = torch.cat([alphas_bar[0:1], alphas]) - betas = 1 - alphas - - return betas - - -class DDIMInverseScheduler(SchedulerMixin, ConfigMixin): - """ - `DDIMInverseScheduler` is the reverse scheduler of [`DDIMScheduler`]. - - 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`, 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`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - clip_sample (`bool`, defaults to `True`): - Clip the predicted sample for numerical stability. - clip_sample_range (`float`, defaults to 1.0): - The maximum magnitude for sample clipping. Valid only when `clip_sample=True`. - set_alpha_to_one (`bool`, defaults to `True`): - Each diffusion step uses the alphas product value at that step and at the previous one. For the final step - there is no previous alpha. When this option is `True` the previous alpha product is fixed to 0, otherwise - it uses the alpha value at step `num_train_timesteps - 1`. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - 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://imagen.research.google/video/paper.pdf) paper). - 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. - 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). - """ - - order = 1 - ignore_for_config = ["kwargs"] - _deprecated_kwargs = ["set_alpha_to_zero"] - - @register_to_config - def __init__( - self, - 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, - clip_sample: bool = True, - set_alpha_to_one: bool = True, - steps_offset: int = 0, - prediction_type: str = "epsilon", - clip_sample_range: float = 1.0, - timestep_spacing: str = "leading", - rescale_betas_zero_snr: bool = False, - **kwargs, - ): - if kwargs.get("set_alpha_to_zero", None) is not None: - deprecation_message = ( - "The `set_alpha_to_zero` argument is deprecated. Please use `set_alpha_to_one` instead." - ) - deprecate("set_alpha_to_zero", "1.0.0", deprecation_message, standard_warn=False) - set_alpha_to_one = kwargs["set_alpha_to_zero"] - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - # Rescale for zero SNR - if rescale_betas_zero_snr: - self.betas = rescale_zero_terminal_snr(self.betas) - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - # At every step in inverted ddim, we are looking into the next alphas_cumprod - # For the initial step, there is no current alphas_cumprod, and the index is out of bounds - # `set_alpha_to_one` decides whether we set this parameter simply to one - # in this case, self.step() just output the predicted noise - # or whether we use the initial alpha used in training the diffusion model. - self.initial_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0] - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # setable values - self.num_inference_steps = None - 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: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - def set_timesteps(self, num_inference_steps: int, device: Union[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. - """ - - if num_inference_steps > self.config.num_train_timesteps: - raise ValueError( - f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:" - f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle" - f" maximal {self.config.num_train_timesteps} timesteps." - ) - - self.num_inference_steps = num_inference_steps - - # "leading" and "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "leading": - step_ratio = self.config.num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round().copy().astype(np.int64) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)[::-1]).astype(np.int64) - timesteps -= 1 - else: - raise ValueError( - f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'leading' or 'trailing'." - ) - - self.timesteps = torch.from_numpy(timesteps).to(device) - - def step( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - timestep (`float`): - 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`. - - Returns: - [`~schedulers.scheduling_ddim_inverse.DDIMInverseSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_ddim_inverse.DDIMInverseSchedulerOutput`] is - returned, otherwise a tuple is returned where the first element is the sample tensor. - - """ - # 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 - ) - - # 2. compute alphas, betas - # change original implementation to exactly match noise levels for analogous forward process - alpha_prod_t = self.alphas_cumprod[timestep] if timestep >= 0 else self.initial_alpha_cumprod - alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] - - beta_prod_t = 1 - alpha_prod_t - - # 3. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (12) from https://huggingface.co/papers/2010.02502 - if self.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - pred_epsilon = model_output - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) - elif self.config.prediction_type == "v_prediction": - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" - " `v_prediction`" - ) - - # 4. Clip or threshold "predicted x_0" - if self.config.clip_sample: - pred_original_sample = pred_original_sample.clamp( - -self.config.clip_sample_range, self.config.clip_sample_range - ) - - # 5. compute "direction pointing to x_t" of formula (12) from https://huggingface.co/papers/2010.02502 - pred_sample_direction = (1 - alpha_prod_t_prev) ** (0.5) * pred_epsilon - - # 6. compute x_t without "random noise" of formula (12) from https://huggingface.co/papers/2010.02502 - prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction - - if not return_dict: - return (prev_sample, pred_original_sample) - return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim_parallel.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim_parallel.py deleted file mode 100644 index 7c3f03a8dbe19e76946769ce8915bcbdffa7e0c3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddim_parallel.py +++ /dev/null @@ -1,646 +0,0 @@ -# Copyright 2025 ParaDiGMS authors 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. - -# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion -# and https://github.com/hojonathanho/diffusion - -import math -from dataclasses import dataclass -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput -class DDIMParallelSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr -def rescale_zero_terminal_snr(betas): - """ - 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. - - Returns: - `torch.Tensor`: rescaled betas with zero terminal SNR - """ - # Convert betas to alphas_bar_sqrt - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - alphas_bar_sqrt = alphas_cumprod.sqrt() - - # Store old values. - alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() - alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() - - # Shift so the last timestep is zero. - alphas_bar_sqrt -= alphas_bar_sqrt_T - - # Scale so the first timestep is back to the old value. - alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - - # Convert alphas_bar_sqrt to betas - alphas_bar = alphas_bar_sqrt**2 # Revert sqrt - alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod - alphas = torch.cat([alphas_bar[0:1], alphas]) - betas = 1 - alphas - - return betas - - -class DDIMParallelScheduler(SchedulerMixin, ConfigMixin): - """ - Denoising diffusion implicit models is a scheduler that extends the denoising procedure introduced in denoising - diffusion probabilistic models (DDPMs) with non-Markovian guidance. - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - For more details, see the original paper: https://huggingface.co/papers/2010.02502 - - 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 - `linear`, `scaled_linear`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, optional): - option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. - 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`. - set_alpha_to_one (`bool`, default `True`): - each diffusion step uses the value of alphas product at that step and at the previous one. For the final - step there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, - otherwise it uses the value of alpha at step 0. - steps_offset (`int`, default `0`): - An offset added to the inference steps, as required by some model families. - prediction_type (`str`, default `epsilon`, optional): - 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://imagen.research.google/video/paper.pdf) - thresholding (`bool`, default `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 - (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"`): - 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. - rescale_betas_zero_snr (`bool`, default `False`): - whether to rescale the betas to have zero terminal SNR (proposed by - https://huggingface.co/papers/2305.08891). This can enable 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). - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - _is_ode_scheduler = True - - @register_to_config - # Copied from diffusers.schedulers.scheduling_ddim.DDIMScheduler.__init__ - def __init__( - self, - 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, - clip_sample: bool = True, - set_alpha_to_one: bool = True, - steps_offset: int = 0, - prediction_type: str = "epsilon", - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - clip_sample_range: float = 1.0, - sample_max_value: float = 1.0, - timestep_spacing: str = "leading", - rescale_betas_zero_snr: bool = False, - ): - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - # Rescale for zero SNR - if rescale_betas_zero_snr: - self.betas = rescale_zero_terminal_snr(self.betas) - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - # At every step in ddim, we are looking into the previous alphas_cumprod - # For the final step, there is no previous alphas_cumprod because we are already at 0 - # `set_alpha_to_one` decides whether we set this parameter simply to one or - # whether we use the final alpha of the "non-previous" one. - self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0] - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # setable values - self.num_inference_steps = None - 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: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - def _get_variance(self, timestep, prev_timestep=None): - if prev_timestep is None: - prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps - - 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 - beta_prod_t_prev = 1 - alpha_prod_t_prev - - variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev) - - return variance - - def _batch_get_variance(self, t, prev_t): - 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) - beta_prod_t = 1 - alpha_prod_t - beta_prod_t_prev = 1 - alpha_prod_t_prev - - variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev) - - return variance - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - """ - "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the - prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by - s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing - pixels from saturation at each step. We find that dynamic thresholding results in significantly better - photorealism as well as better image-text alignment, especially when using very large guidance weights." - - https://huggingface.co/papers/2205.11487 - """ - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half - - # Flatten sample for doing quantile calculation along each image - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - - abs_sample = sample.abs() # "a certain percentile absolute pixel value" - - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] - s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 - sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" - - sample = sample.reshape(batch_size, channels, *remaining_dims) - sample = sample.to(dtype) - - 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): - """ - 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. - """ - - if num_inference_steps > self.config.num_train_timesteps: - raise ValueError( - f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:" - f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle" - f" maximal {self.config.num_train_timesteps} timesteps." - ) - - self.num_inference_steps = num_inference_steps - - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = ( - np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps) - .round()[::-1] - .copy() - .astype(np.int64) - ) - elif self.config.timestep_spacing == "leading": - step_ratio = self.config.num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64) - timesteps -= 1 - else: - raise ValueError( - f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'leading' or 'trailing'." - ) - - self.timesteps = torch.from_numpy(timesteps).to(device) - - def step( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - eta: float = 0.0, - use_clipped_model_output: bool = False, - generator=None, - variance_noise: Optional[torch.Tensor] = None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - 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. - 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) - return_dict (`bool`): option for returning tuple rather than DDIMParallelSchedulerOutput class - - Returns: - [`~schedulers.scheduling_utils.DDIMParallelSchedulerOutput`] or `tuple`: - [`~schedulers.scheduling_utils.DDIMParallelSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. - When returning a tuple, the first element is the sample tensor. - - """ - 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" - ) - - # See formulas (12) and (16) of DDIM paper https://huggingface.co/papers/2010.02502 - # Ideally, read DDIM paper in-detail understanding - - # Notation ( -> - # - pred_noise_t -> e_theta(x_t, t) - # - pred_original_sample -> f_theta(x_t, t) or x_0 - # - std_dev_t -> sigma_t - # - eta -> η - # - pred_sample_direction -> "direction pointing to x_t" - # - pred_prev_sample -> "x_t-1" - - # 1. get previous step value (=t-1) - prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps - - # 2. compute alphas, betas - 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 - - # 3. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (12) from https://huggingface.co/papers/2010.02502 - if self.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - pred_epsilon = model_output - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) - elif self.config.prediction_type == "v_prediction": - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" - " `v_prediction`" - ) - - # 4. Clip or threshold "predicted x_0" - if self.config.thresholding: - pred_original_sample = self._threshold_sample(pred_original_sample) - elif self.config.clip_sample: - pred_original_sample = pred_original_sample.clamp( - -self.config.clip_sample_range, self.config.clip_sample_range - ) - - # 5. compute variance: "sigma_t(η)" -> see formula (16) - # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1) - variance = self._get_variance(timestep, prev_timestep) - std_dev_t = eta * variance ** (0.5) - - if use_clipped_model_output: - # the pred_epsilon is always re-derived from the clipped x_0 in Glide - pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) - - # 6. compute "direction pointing to x_t" of formula (12) from https://huggingface.co/papers/2010.02502 - pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * pred_epsilon - - # 7. compute x_t without "random noise" of formula (12) from https://huggingface.co/papers/2010.02502 - prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction - - if eta > 0: - if variance_noise is not None and generator is not None: - raise ValueError( - "Cannot pass both generator and variance_noise. Please make sure that either `generator` or" - " `variance_noise` stays `None`." - ) - - if variance_noise is None: - variance_noise = randn_tensor( - model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype - ) - variance = std_dev_t * variance_noise - - prev_sample = prev_sample + variance - - if not return_dict: - return ( - prev_sample, - pred_original_sample, - ) - - return DDIMParallelSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) - - def batch_step_no_noise( - self, - model_output: torch.Tensor, - timesteps: List[int], - sample: torch.Tensor, - eta: float = 0.0, - use_clipped_model_output: bool = False, - ) -> torch.Tensor: - """ - Batched version of the `step` function, to be able to reverse the SDE for multiple samples/timesteps at once. - Also, does not add any noise to the predicted sample, which is necessary for parallel sampling where the noise - is pre-sampled by the pipeline. - - 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). - - 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. - 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. - - Returns: - `torch.Tensor`: sample tensor at previous timestep. - - """ - 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" - ) - - assert eta == 0.0 - - # See formulas (12) and (16) of DDIM paper https://huggingface.co/papers/2010.02502 - # Ideally, read DDIM paper in-detail understanding - - # Notation ( -> - # - pred_noise_t -> e_theta(x_t, t) - # - pred_original_sample -> f_theta(x_t, t) or x_0 - # - std_dev_t -> sigma_t - # - eta -> η - # - pred_sample_direction -> "direction pointing to x_t" - # - pred_prev_sample -> "x_t-1" - - # 1. get previous step value (=t-1) - t = timesteps - prev_t = t - self.config.num_train_timesteps // self.num_inference_steps - - t = t.view(-1, *([1] * (model_output.ndim - 1))) - prev_t = prev_t.view(-1, *([1] * (model_output.ndim - 1))) - - # 1. compute alphas, betas - self.alphas_cumprod = self.alphas_cumprod.to(model_output.device) - self.final_alpha_cumprod = self.final_alpha_cumprod.to(model_output.device) - 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) - - beta_prod_t = 1 - alpha_prod_t - - # 3. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (12) from https://huggingface.co/papers/2010.02502 - if self.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - pred_epsilon = model_output - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) - elif self.config.prediction_type == "v_prediction": - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" - " `v_prediction`" - ) - - # 4. Clip or threshold "predicted x_0" - if self.config.thresholding: - pred_original_sample = self._threshold_sample(pred_original_sample) - elif self.config.clip_sample: - pred_original_sample = pred_original_sample.clamp( - -self.config.clip_sample_range, self.config.clip_sample_range - ) - - # 5. compute variance: "sigma_t(η)" -> see formula (16) - # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1) - variance = self._batch_get_variance(t, prev_t).to(model_output.device).view(*alpha_prod_t_prev.shape) - std_dev_t = eta * variance ** (0.5) - - if use_clipped_model_output: - # the pred_epsilon is always re-derived from the clipped x_0 in Glide - pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) - - # 6. compute "direction pointing to x_t" of formula (12) from https://huggingface.co/papers/2010.02502 - pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * pred_epsilon - - # 7. compute x_t without "random noise" of formula (12) from https://huggingface.co/papers/2010.02502 - prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction - - return prev_sample - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as original_samples - # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement - # for the subsequent add_noise calls - self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype) - timesteps = timesteps.to(original_samples.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(original_samples.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise - return noisy_samples - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity - def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as sample - self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype) - timesteps = timesteps.to(sample.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(sample.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample - return velocity - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddpm.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddpm.py deleted file mode 100644 index 0fab6d910a822e688cd8f2c9ac162b2bf1d0d269..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddpm.py +++ /dev/null @@ -1,558 +0,0 @@ -# Copyright 2025 UC Berkeley 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. - -# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim - -import math -from dataclasses import dataclass -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin - - -@dataclass -class DDPMSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr -def rescale_zero_terminal_snr(betas): - """ - 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. - - Returns: - `torch.Tensor`: rescaled betas with zero terminal SNR - """ - # Convert betas to alphas_bar_sqrt - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - alphas_bar_sqrt = alphas_cumprod.sqrt() - - # Store old values. - alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() - alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() - - # Shift so the last timestep is zero. - alphas_bar_sqrt -= alphas_bar_sqrt_T - - # Scale so the first timestep is back to the old value. - alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - - # Convert alphas_bar_sqrt to betas - alphas_bar = alphas_bar_sqrt**2 # Revert sqrt - alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod - alphas = torch.cat([alphas_bar[0:1], alphas]) - betas = 1 - alphas - - return betas - - -class DDPMScheduler(SchedulerMixin, ConfigMixin): - """ - `DDPMScheduler` explores the connections between denoising score matching and Langevin dynamics sampling. - - 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`, 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*): - An array of betas to pass directly to the constructor without using `beta_start` and `beta_end`. - variance_type (`str`, defaults to `"fixed_small"`): - Clip the variance 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`, defaults to `True`): - Clip the predicted sample for numerical stability. - clip_sample_range (`float`, defaults to 1.0): - The maximum magnitude for sample clipping. Valid only when `clip_sample=True`. - 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://imagen.research.google/video/paper.pdf) paper). - thresholding (`bool`, defaults to `False`): - Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such - as Stable Diffusion. - dynamic_thresholding_ratio (`float`, defaults to 0.995): - The ratio for the dynamic thresholding method. Valid only when `thresholding=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 the [Common Diffusion Noise Schedules and - Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. - 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 - 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). - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - 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, - variance_type: str = "fixed_small", - clip_sample: bool = True, - prediction_type: str = "epsilon", - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - clip_sample_range: float = 1.0, - sample_max_value: float = 1.0, - timestep_spacing: str = "leading", - steps_offset: int = 0, - rescale_betas_zero_snr: bool = False, - ): - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - elif beta_schedule == "sigmoid": - # GeoDiff sigmoid schedule - betas = torch.linspace(-6, 6, num_train_timesteps) - self.betas = torch.sigmoid(betas) * (beta_end - beta_start) + beta_start - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - # Rescale for zero SNR - if rescale_betas_zero_snr: - self.betas = rescale_zero_terminal_snr(self.betas) - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - self.one = torch.tensor(1.0) - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # setable values - self.custom_timesteps = False - self.num_inference_steps = None - self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy()) - - self.variance_type = variance_type - - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - def set_timesteps( - self, - num_inference_steps: Optional[int] = None, - device: Union[str, torch.device] = None, - timesteps: Optional[List[int]] = 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. 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 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`. - - """ - if num_inference_steps is not None and timesteps is not None: - raise ValueError("Can only pass one of `num_inference_steps` or `custom_timesteps`.") - - if timesteps is not None: - for i in range(1, len(timesteps)): - if timesteps[i] >= timesteps[i - 1]: - raise ValueError("`custom_timesteps` must be in descending order.") - - if timesteps[0] >= self.config.num_train_timesteps: - raise ValueError( - f"`timesteps` must start before `self.config.train_timesteps`: {self.config.num_train_timesteps}." - ) - - timesteps = np.array(timesteps, dtype=np.int64) - self.custom_timesteps = True - else: - if num_inference_steps > self.config.num_train_timesteps: - raise ValueError( - f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:" - f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle" - f" maximal {self.config.num_train_timesteps} timesteps." - ) - - self.num_inference_steps = num_inference_steps - self.custom_timesteps = False - - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = ( - np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps) - .round()[::-1] - .copy() - .astype(np.int64) - ) - elif self.config.timestep_spacing == "leading": - step_ratio = self.config.num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64) - timesteps -= 1 - else: - raise ValueError( - f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'." - ) - - self.timesteps = torch.from_numpy(timesteps).to(device) - - def _get_variance(self, t, predicted_variance=None, variance_type=None): - prev_t = self.previous_timestep(t) - - alpha_prod_t = self.alphas_cumprod[t] - alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.one - current_beta_t = 1 - alpha_prod_t / alpha_prod_t_prev - - # For t > 0, compute predicted variance βt (see formula (6) and (7) from https://huggingface.co/papers/2006.11239) - # and sample from it to get previous sample - # x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample - variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * current_beta_t - - # we always take the log of variance, so clamp it to ensure it's not 0 - variance = torch.clamp(variance, min=1e-20) - - if variance_type is None: - variance_type = self.config.variance_type - - # hacks - were probably added for training stability - if variance_type == "fixed_small": - variance = variance - # for rl-diffuser https://huggingface.co/papers/2205.09991 - elif variance_type == "fixed_small_log": - variance = torch.log(variance) - variance = torch.exp(0.5 * variance) - elif variance_type == "fixed_large": - variance = current_beta_t - elif variance_type == "fixed_large_log": - # Glide max_log - variance = torch.log(current_beta_t) - elif variance_type == "learned": - return predicted_variance - elif variance_type == "learned_range": - min_log = torch.log(variance) - max_log = torch.log(current_beta_t) - frac = (predicted_variance + 1) / 2 - variance = frac * max_log + (1 - frac) * min_log - - return variance - - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - """ - "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the - prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by - s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing - pixels from saturation at each step. We find that dynamic thresholding results in significantly better - photorealism as well as better image-text alignment, especially when using very large guidance weights." - - https://huggingface.co/papers/2205.11487 - """ - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half - - # Flatten sample for doing quantile calculation along each image - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - - abs_sample = sample.abs() # "a certain percentile absolute pixel value" - - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] - s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 - sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" - - sample = sample.reshape(batch_size, channels, *remaining_dims) - sample = sample.to(dtype) - - return sample - - def step( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - generator=None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - timestep (`float`): - 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`, *optional*, defaults to `True`): - Whether or not to return a [`~schedulers.scheduling_ddpm.DDPMSchedulerOutput`] or `tuple`. - - Returns: - [`~schedulers.scheduling_ddpm.DDPMSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_ddpm.DDPMSchedulerOutput`] is returned, otherwise a - tuple is returned where the first element is the sample tensor. - - """ - t = timestep - - prev_t = self.previous_timestep(t) - - 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 - - # 1. compute alphas, betas - alpha_prod_t = self.alphas_cumprod[t] - alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.one - beta_prod_t = 1 - alpha_prod_t - beta_prod_t_prev = 1 - alpha_prod_t_prev - current_alpha_t = alpha_prod_t / alpha_prod_t_prev - current_beta_t = 1 - current_alpha_t - - # 2. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (15) from https://huggingface.co/papers/2006.11239 - if self.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - elif self.config.prediction_type == "v_prediction": - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or" - " `v_prediction` for the DDPMScheduler." - ) - - # 3. Clip or threshold "predicted x_0" - if self.config.thresholding: - pred_original_sample = self._threshold_sample(pred_original_sample) - elif self.config.clip_sample: - pred_original_sample = pred_original_sample.clamp( - -self.config.clip_sample_range, self.config.clip_sample_range - ) - - # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t - # See formula (7) from https://huggingface.co/papers/2006.11239 - pred_original_sample_coeff = (alpha_prod_t_prev ** (0.5) * current_beta_t) / beta_prod_t - current_sample_coeff = current_alpha_t ** (0.5) * beta_prod_t_prev / beta_prod_t - - # 5. Compute predicted previous sample µ_t - # See formula (7) from https://huggingface.co/papers/2006.11239 - pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample - - # 6. Add noise - variance = 0 - if t > 0: - device = model_output.device - variance_noise = randn_tensor( - 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 - elif self.variance_type == "learned_range": - variance = self._get_variance(t, predicted_variance=predicted_variance) - variance = torch.exp(0.5 * variance) * variance_noise - else: - variance = (self._get_variance(t, predicted_variance=predicted_variance) ** 0.5) * variance_noise - - pred_prev_sample = pred_prev_sample + variance - - if not return_dict: - return ( - pred_prev_sample, - pred_original_sample, - ) - - return DDPMSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample) - - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as original_samples - # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement - # for the subsequent add_noise calls - self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype) - timesteps = timesteps.to(original_samples.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(original_samples.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise - return noisy_samples - - def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as sample - self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype) - timesteps = timesteps.to(sample.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(sample.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample - return velocity - - def __len__(self): - return self.config.num_train_timesteps - - def previous_timestep(self, timestep): - if self.custom_timesteps or self.num_inference_steps: - index = (self.timesteps == timestep).nonzero(as_tuple=True)[0][0] - if index == self.timesteps.shape[0] - 1: - prev_t = torch.tensor(-1) - else: - prev_t = self.timesteps[index + 1] - else: - prev_t = timestep - 1 - return prev_t diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddpm_flax.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddpm_flax.py deleted file mode 100644 index a3264f54f572f60c1d44898bd220bf512684f1cb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddpm_flax.py +++ /dev/null @@ -1,303 +0,0 @@ -# Copyright 2025 UC Berkeley 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. - -# 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 .scheduling_utils_flax import ( - CommonSchedulerState, - FlaxKarrasDiffusionSchedulers, - FlaxSchedulerMixin, - FlaxSchedulerOutput, - add_noise_common, - get_velocity_common, -) - - -@flax.struct.dataclass -class DDPMSchedulerState: - common: CommonSchedulerState - - # setable values - init_noise_sigma: jnp.ndarray - timesteps: jnp.ndarray - num_inference_steps: Optional[int] = None - - @classmethod - 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) - - -@dataclass -class FlaxDDPMSchedulerOutput(FlaxSchedulerOutput): - state: DDPMSchedulerState - - -class FlaxDDPMScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - Denoising diffusion probabilistic models (DDPMs) explores the connections between denoising score matching and - Langevin dynamics sampling. - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - 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 - `linear`, `scaled_linear`, or `squaredcos_cap_v2`. - 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`, - `fixed_small_log`, `fixed_large`, `fixed_large_log`, `learned` or `learned_range`. - clip_sample (`bool`, default `True`): - option to clip predicted sample between -1 and 1 for numerical stability. - prediction_type (`str`, default `epsilon`): - indicates whether the model predicts the noise (epsilon), or the samples. One of `epsilon`, `sample`. - `v-prediction` is not supported for this scheduler. - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - the `dtype` used for params and computation. - """ - - _compatibles = [e.name for e in FlaxKarrasDiffusionSchedulers] - - dtype: jnp.dtype - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: Optional[jnp.ndarray] = None, - variance_type: str = "fixed_small", - clip_sample: bool = True, - prediction_type: str = "epsilon", - dtype: jnp.dtype = jnp.float32, - ): - self.dtype = dtype - - def create_state(self, common: Optional[CommonSchedulerState] = None) -> DDPMSchedulerState: - if common is None: - common = CommonSchedulerState.create(self) - - # standard deviation of the initial noise distribution - init_noise_sigma = jnp.array(1.0, dtype=self.dtype) - - timesteps = jnp.arange(0, self.config.num_train_timesteps).round()[::-1] - - return DDPMSchedulerState.create( - common=common, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - ) - - def scale_model_input( - self, state: DDPMSchedulerState, sample: jnp.ndarray, timestep: Optional[int] = None - ) -> jnp.ndarray: - """ - Args: - state (`PNDMSchedulerState`): the `FlaxPNDMScheduler` state data class instance. - sample (`jnp.ndarray`): input sample - timestep (`int`, optional): current timestep - - Returns: - `jnp.ndarray`: scaled input sample - """ - return sample - - def set_timesteps( - 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. - - Args: - state (`DDIMSchedulerState`): - the `FlaxDDPMScheduler` state data class instance. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - """ - - step_ratio = self.config.num_train_timesteps // num_inference_steps - # creates integer timesteps by multiplying by ratio - # rounding to avoid issues when num_inference_step is power of 3 - timesteps = (jnp.arange(0, num_inference_steps) * step_ratio).round()[::-1] - - return state.replace( - num_inference_steps=num_inference_steps, - timesteps=timesteps, - ) - - def _get_variance(self, state: DDPMSchedulerState, t, predicted_variance=None, variance_type=None): - alpha_prod_t = state.common.alphas_cumprod[t] - alpha_prod_t_prev = jnp.where(t > 0, state.common.alphas_cumprod[t - 1], jnp.array(1.0, dtype=self.dtype)) - - # For t > 0, compute predicted variance βt (see formula (6) and (7) from https://huggingface.co/papers/2006.11239) - # and sample from it to get previous sample - # x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample - variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * state.common.betas[t] - - if variance_type is None: - variance_type = self.config.variance_type - - # hacks - were probably added for training stability - if variance_type == "fixed_small": - variance = jnp.clip(variance, a_min=1e-20) - # for rl-diffuser https://huggingface.co/papers/2205.09991 - elif variance_type == "fixed_small_log": - variance = jnp.log(jnp.clip(variance, a_min=1e-20)) - elif variance_type == "fixed_large": - variance = state.common.betas[t] - elif variance_type == "fixed_large_log": - # Glide max_log - variance = jnp.log(state.common.betas[t]) - elif variance_type == "learned": - return predicted_variance - elif variance_type == "learned_range": - min_log = variance - max_log = state.common.betas[t] - frac = (predicted_variance + 1) / 2 - variance = frac * max_log + (1 - frac) * min_log - - return variance - - def step( - self, - state: DDPMSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - key: Optional[jax.Array] = None, - return_dict: bool = True, - ) -> Union[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). - - Args: - state (`DDPMSchedulerState`): the `FlaxDDPMScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - key (`jax.Array`): a PRNG key. - return_dict (`bool`): option for returning tuple rather than FlaxDDPMSchedulerOutput class - - Returns: - [`FlaxDDPMSchedulerOutput`] or `tuple`: [`FlaxDDPMSchedulerOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - t = timestep - - if key is None: - key = jax.random.key(0) - - if ( - len(model_output.shape) > 1 - and model_output.shape[1] == sample.shape[1] * 2 - and self.config.variance_type in ["learned", "learned_range"] - ): - model_output, predicted_variance = jnp.split(model_output, sample.shape[1], axis=1) - else: - predicted_variance = None - - # 1. compute alphas, betas - alpha_prod_t = state.common.alphas_cumprod[t] - alpha_prod_t_prev = jnp.where(t > 0, state.common.alphas_cumprod[t - 1], jnp.array(1.0, dtype=self.dtype)) - beta_prod_t = 1 - alpha_prod_t - beta_prod_t_prev = 1 - alpha_prod_t_prev - - # 2. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (15) from https://huggingface.co/papers/2006.11239 - if self.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - elif self.config.prediction_type == "v_prediction": - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` " - " for the FlaxDDPMScheduler." - ) - - # 3. Clip "predicted x_0" - if self.config.clip_sample: - pred_original_sample = jnp.clip(pred_original_sample, -1, 1) - - # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t - # See formula (7) from https://huggingface.co/papers/2006.11239 - pred_original_sample_coeff = (alpha_prod_t_prev ** (0.5) * state.common.betas[t]) / beta_prod_t - current_sample_coeff = state.common.alphas[t] ** (0.5) * beta_prod_t_prev / beta_prod_t - - # 5. Compute predicted previous sample µ_t - # See formula (7) from https://huggingface.co/papers/2006.11239 - pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample - - # 6. Add noise - def random_variance(): - split_key = jax.random.split(key, num=1)[0] - noise = jax.random.normal(split_key, shape=model_output.shape, dtype=self.dtype) - return (self._get_variance(state, t, predicted_variance=predicted_variance) ** 0.5) * noise - - variance = jnp.where(t > 0, random_variance(), jnp.zeros(model_output.shape, dtype=self.dtype)) - - pred_prev_sample = pred_prev_sample + variance - - if not return_dict: - return (pred_prev_sample, state) - - return FlaxDDPMSchedulerOutput(prev_sample=pred_prev_sample, state=state) - - def add_noise( - self, - state: DDPMSchedulerState, - original_samples: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - return add_noise_common(state.common, original_samples, noise, timesteps) - - def get_velocity( - self, - state: DDPMSchedulerState, - sample: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - return get_velocity_common(state.common, sample, noise, timesteps) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddpm_parallel.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddpm_parallel.py deleted file mode 100644 index ec741f9ecb7dd1f28fde6b53a9c7f5c9fa43b5c9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddpm_parallel.py +++ /dev/null @@ -1,649 +0,0 @@ -# Copyright 2025 ParaDiGMS authors 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. - -# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim - -import math -from dataclasses import dataclass -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput -class DDPMParallelSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr -def rescale_zero_terminal_snr(betas): - """ - 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. - - Returns: - `torch.Tensor`: rescaled betas with zero terminal SNR - """ - # Convert betas to alphas_bar_sqrt - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - alphas_bar_sqrt = alphas_cumprod.sqrt() - - # Store old values. - alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() - alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() - - # Shift so the last timestep is zero. - alphas_bar_sqrt -= alphas_bar_sqrt_T - - # Scale so the first timestep is back to the old value. - alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - - # Convert alphas_bar_sqrt to betas - alphas_bar = alphas_bar_sqrt**2 # Revert sqrt - alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod - alphas = torch.cat([alphas_bar[0:1], alphas]) - betas = 1 - alphas - - return betas - - -class DDPMParallelScheduler(SchedulerMixin, ConfigMixin): - """ - Denoising diffusion probabilistic models (DDPMs) explores the connections between denoising score matching and - Langevin dynamics sampling. - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - 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 - `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`, - `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 - process), `sample` (directly predicting the noisy sample`) or `v_prediction` (see section 2.4 - https://imagen.research.google/video/paper.pdf) - thresholding (`bool`, default `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 - (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"`): - 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`): - 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 - 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). - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - _is_ode_scheduler = False - - @register_to_config - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.__init__ - def __init__( - self, - 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, - variance_type: str = "fixed_small", - clip_sample: bool = True, - prediction_type: str = "epsilon", - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - clip_sample_range: float = 1.0, - sample_max_value: float = 1.0, - timestep_spacing: str = "leading", - steps_offset: int = 0, - rescale_betas_zero_snr: bool = False, - ): - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - elif beta_schedule == "sigmoid": - # GeoDiff sigmoid schedule - betas = torch.linspace(-6, 6, num_train_timesteps) - self.betas = torch.sigmoid(betas) * (beta_end - beta_start) + beta_start - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - # Rescale for zero SNR - if rescale_betas_zero_snr: - self.betas = rescale_zero_terminal_snr(self.betas) - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - self.one = torch.tensor(1.0) - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # setable values - self.custom_timesteps = False - self.num_inference_steps = None - self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy()) - - 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: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - # 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, - ): - """ - 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. 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 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`. - - """ - if num_inference_steps is not None and timesteps is not None: - raise ValueError("Can only pass one of `num_inference_steps` or `custom_timesteps`.") - - if timesteps is not None: - for i in range(1, len(timesteps)): - if timesteps[i] >= timesteps[i - 1]: - raise ValueError("`custom_timesteps` must be in descending order.") - - if timesteps[0] >= self.config.num_train_timesteps: - raise ValueError( - f"`timesteps` must start before `self.config.train_timesteps`: {self.config.num_train_timesteps}." - ) - - timesteps = np.array(timesteps, dtype=np.int64) - self.custom_timesteps = True - else: - if num_inference_steps > self.config.num_train_timesteps: - raise ValueError( - f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:" - f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle" - f" maximal {self.config.num_train_timesteps} timesteps." - ) - - self.num_inference_steps = num_inference_steps - self.custom_timesteps = False - - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = ( - np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps) - .round()[::-1] - .copy() - .astype(np.int64) - ) - elif self.config.timestep_spacing == "leading": - step_ratio = self.config.num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64) - timesteps -= 1 - else: - raise ValueError( - f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'." - ) - - self.timesteps = torch.from_numpy(timesteps).to(device) - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._get_variance - def _get_variance(self, t, predicted_variance=None, variance_type=None): - prev_t = self.previous_timestep(t) - - alpha_prod_t = self.alphas_cumprod[t] - alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.one - current_beta_t = 1 - alpha_prod_t / alpha_prod_t_prev - - # For t > 0, compute predicted variance βt (see formula (6) and (7) from https://huggingface.co/papers/2006.11239) - # and sample from it to get previous sample - # x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample - variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * current_beta_t - - # we always take the log of variance, so clamp it to ensure it's not 0 - variance = torch.clamp(variance, min=1e-20) - - if variance_type is None: - variance_type = self.config.variance_type - - # hacks - were probably added for training stability - if variance_type == "fixed_small": - variance = variance - # for rl-diffuser https://huggingface.co/papers/2205.09991 - elif variance_type == "fixed_small_log": - variance = torch.log(variance) - variance = torch.exp(0.5 * variance) - elif variance_type == "fixed_large": - variance = current_beta_t - elif variance_type == "fixed_large_log": - # Glide max_log - variance = torch.log(current_beta_t) - elif variance_type == "learned": - return predicted_variance - elif variance_type == "learned_range": - min_log = torch.log(variance) - max_log = torch.log(current_beta_t) - frac = (predicted_variance + 1) / 2 - variance = frac * max_log + (1 - frac) * min_log - - return variance - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - """ - "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the - prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by - s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing - pixels from saturation at each step. We find that dynamic thresholding results in significantly better - photorealism as well as better image-text alignment, especially when using very large guidance weights." - - https://huggingface.co/papers/2205.11487 - """ - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half - - # Flatten sample for doing quantile calculation along each image - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - - abs_sample = sample.abs() # "a certain percentile absolute pixel value" - - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] - s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 - sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" - - sample = sample.reshape(batch_size, channels, *remaining_dims) - sample = sample.to(dtype) - - return sample - - def step( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - generator=None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): direct output from learned diffusion model. - 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. - return_dict (`bool`): option for returning tuple rather than DDPMParallelSchedulerOutput class - - Returns: - [`~schedulers.scheduling_utils.DDPMParallelSchedulerOutput`] or `tuple`: - [`~schedulers.scheduling_utils.DDPMParallelSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. - When returning a tuple, the first element is the sample tensor. - - """ - t = timestep - - prev_t = self.previous_timestep(t) - - 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 - - # 1. compute alphas, betas - alpha_prod_t = self.alphas_cumprod[t] - alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.one - beta_prod_t = 1 - alpha_prod_t - beta_prod_t_prev = 1 - alpha_prod_t_prev - current_alpha_t = alpha_prod_t / alpha_prod_t_prev - current_beta_t = 1 - current_alpha_t - - # 2. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (15) from https://huggingface.co/papers/2006.11239 - if self.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - elif self.config.prediction_type == "v_prediction": - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or" - " `v_prediction` for the DDPMScheduler." - ) - - # 3. Clip or threshold "predicted x_0" - if self.config.thresholding: - pred_original_sample = self._threshold_sample(pred_original_sample) - elif self.config.clip_sample: - pred_original_sample = pred_original_sample.clamp( - -self.config.clip_sample_range, self.config.clip_sample_range - ) - - # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t - # See formula (7) from https://huggingface.co/papers/2006.11239 - pred_original_sample_coeff = (alpha_prod_t_prev ** (0.5) * current_beta_t) / beta_prod_t - current_sample_coeff = current_alpha_t ** (0.5) * beta_prod_t_prev / beta_prod_t - - # 5. Compute predicted previous sample µ_t - # See formula (7) from https://huggingface.co/papers/2006.11239 - pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample - - # 6. Add noise - variance = 0 - if t > 0: - device = model_output.device - variance_noise = randn_tensor( - 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 - elif self.variance_type == "learned_range": - variance = self._get_variance(t, predicted_variance=predicted_variance) - variance = torch.exp(0.5 * variance) * variance_noise - else: - variance = (self._get_variance(t, predicted_variance=predicted_variance) ** 0.5) * variance_noise - - pred_prev_sample = pred_prev_sample + variance - - if not return_dict: - return ( - pred_prev_sample, - pred_original_sample, - ) - - return DDPMParallelSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample) - - def batch_step_no_noise( - self, - model_output: torch.Tensor, - timesteps: List[int], - sample: torch.Tensor, - ) -> torch.Tensor: - """ - Batched version of the `step` function, to be able to reverse the SDE for multiple samples/timesteps at once. - Also, does not add any noise to the predicted sample, which is necessary for parallel sampling where the noise - is pre-sampled by the pipeline. - - 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). - - 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. - sample (`torch.Tensor`): - current instance of sample being created by diffusion process. - - Returns: - `torch.Tensor`: sample tensor at previous timestep. - """ - t = timesteps - num_inference_steps = self.num_inference_steps if self.num_inference_steps else self.config.num_train_timesteps - prev_t = t - self.config.num_train_timesteps // num_inference_steps - - 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"]: - model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1) - else: - pass - - # 1. compute alphas, betas - self.alphas_cumprod = self.alphas_cumprod.to(model_output.device) - 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) - - beta_prod_t = 1 - alpha_prod_t - beta_prod_t_prev = 1 - alpha_prod_t_prev - current_alpha_t = alpha_prod_t / alpha_prod_t_prev - current_beta_t = 1 - current_alpha_t - - # 2. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (15) from https://huggingface.co/papers/2006.11239 - if self.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - elif self.config.prediction_type == "v_prediction": - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or" - " `v_prediction` for the DDPMParallelScheduler." - ) - - # 3. Clip or threshold "predicted x_0" - if self.config.thresholding: - pred_original_sample = self._threshold_sample(pred_original_sample) - elif self.config.clip_sample: - pred_original_sample = pred_original_sample.clamp( - -self.config.clip_sample_range, self.config.clip_sample_range - ) - - # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t - # See formula (7) from https://huggingface.co/papers/2006.11239 - pred_original_sample_coeff = (alpha_prod_t_prev ** (0.5) * current_beta_t) / beta_prod_t - current_sample_coeff = current_alpha_t ** (0.5) * beta_prod_t_prev / beta_prod_t - - # 5. Compute predicted previous sample µ_t - # See formula (7) from https://huggingface.co/papers/2006.11239 - pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample - - return pred_prev_sample - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as original_samples - # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement - # for the subsequent add_noise calls - self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype) - timesteps = timesteps.to(original_samples.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(original_samples.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise - return noisy_samples - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity - def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as sample - self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype) - timesteps = timesteps.to(sample.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(sample.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample - return velocity - - def __len__(self): - return self.config.num_train_timesteps - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.previous_timestep - def previous_timestep(self, timestep): - if self.custom_timesteps or self.num_inference_steps: - index = (self.timesteps == timestep).nonzero(as_tuple=True)[0][0] - if index == self.timesteps.shape[0] - 1: - prev_t = torch.tensor(-1) - else: - prev_t = self.timesteps[index + 1] - else: - prev_t = timestep - 1 - return prev_t diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddpm_wuerstchen.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddpm_wuerstchen.py deleted file mode 100644 index 71f08277ebd71fc363dcddc4c8dd5fb6fa250e2b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ddpm_wuerstchen.py +++ /dev/null @@ -1,230 +0,0 @@ -# Copyright (c) 2022 Pablo Pernías MIT License -# Copyright 2025 UC Berkeley 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. - -# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim - -import math -from dataclasses import dataclass -from typing import List, Optional, Tuple, Union - -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import SchedulerMixin - - -@dataclass -class DDPMWuerstchenSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's step function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - """ - - prev_sample: torch.Tensor - - -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -class DDPMWuerstchenScheduler(SchedulerMixin, ConfigMixin): - """ - Denoising diffusion probabilistic models (DDPMs) explores the connections between denoising score matching and - Langevin dynamics sampling. - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - For more details, see the original paper: https://huggingface.co/papers/2006.11239 - - Args: - scaler (`float`): .... - s (`float`): .... - """ - - @register_to_config - def __init__( - self, - scaler: float = 1.0, - s: float = 0.008, - ): - self.scaler = scaler - self.s = torch.tensor([s]) - self._init_alpha_cumprod = torch.cos(self.s / (1 + self.s) * torch.pi * 0.5) ** 2 - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - def _alpha_cumprod(self, t, device): - if self.scaler > 1: - t = 1 - (1 - t) ** self.scaler - elif self.scaler < 1: - t = t**self.scaler - alpha_cumprod = torch.cos( - (t + self.s.to(device)) / (1 + self.s.to(device)) * torch.pi * 0.5 - ) ** 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: - """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. - - Args: - sample (`torch.Tensor`): input sample - timestep (`int`, optional): current timestep - - Returns: - `torch.Tensor`: scaled input sample - """ - return sample - - def set_timesteps( - self, - num_inference_steps: int = None, - timesteps: Optional[List[int]] = None, - device: Union[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]`): - 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): - the device to which the timesteps are moved to. {2 / 3: 20, 0.0: 10} - """ - if timesteps is None: - timesteps = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device) - if not isinstance(timesteps, torch.Tensor): - timesteps = torch.Tensor(timesteps).to(device) - self.timesteps = timesteps - - def step( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - generator=None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): direct output from learned diffusion model. - 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. - return_dict (`bool`): option for returning tuple rather than DDPMWuerstchenSchedulerOutput class - - Returns: - [`DDPMWuerstchenSchedulerOutput`] or `tuple`: [`DDPMWuerstchenSchedulerOutput`] if `return_dict` is True, - otherwise a `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - dtype = model_output.dtype - device = model_output.device - t = timestep - - prev_t = self.previous_timestep(t) - - alpha_cumprod = self._alpha_cumprod(t, device).view(t.size(0), *[1 for _ in sample.shape[1:]]) - alpha_cumprod_prev = self._alpha_cumprod(prev_t, device).view(prev_t.size(0), *[1 for _ in sample.shape[1:]]) - alpha = alpha_cumprod / alpha_cumprod_prev - - mu = (1.0 / alpha).sqrt() * (sample - (1 - alpha) * model_output / (1 - alpha_cumprod).sqrt()) - - std_noise = randn_tensor(mu.shape, generator=generator, device=model_output.device, dtype=model_output.dtype) - std = ((1 - alpha) * (1.0 - alpha_cumprod_prev) / (1.0 - alpha_cumprod)).sqrt() * std_noise - pred = mu + std * (prev_t != 0).float().view(prev_t.size(0), *[1 for _ in sample.shape[1:]]) - - if not return_dict: - return (pred.to(dtype),) - - return DDPMWuerstchenSchedulerOutput(prev_sample=pred.to(dtype)) - - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.Tensor, - ) -> torch.Tensor: - device = original_samples.device - dtype = original_samples.dtype - alpha_cumprod = self._alpha_cumprod(timesteps, device=device).view( - timesteps.size(0), *[1 for _ in original_samples.shape[1:]] - ) - noisy_samples = alpha_cumprod.sqrt() * original_samples + (1 - alpha_cumprod).sqrt() * noise - return noisy_samples.to(dtype=dtype) - - def __len__(self): - return self.config.num_train_timesteps - - def previous_timestep(self, timestep): - index = (self.timesteps - timestep[0]).abs().argmin().item() - prev_t = self.timesteps[index + 1][None].expand(timestep.shape[0]) - return prev_t diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_deis_multistep.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_deis_multistep.py deleted file mode 100644 index 748a7e39c0b89414dcdce6c37e63fb4816201f98..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_deis_multistep.py +++ /dev/null @@ -1,886 +0,0 @@ -# Copyright 2025 FLAIR Lab 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. - -# DISCLAIMER: check https://huggingface.co/papers/2204.13902 and https://github.com/qsh-zh/deis for more info -# 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, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import deprecate, is_scipy_available -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput - - -if is_scipy_available(): - import scipy.stats - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -class DEISMultistepScheduler(SchedulerMixin, ConfigMixin): - """ - `DEISMultistepScheduler` is a fast high order solver for diffusion ordinary differential equations (ODEs). - - 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`, 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`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, *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 - sampling, and `solver_order=3` for unconditional sampling. - prediction_type (`str`, defaults to `epsilon`): - 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://imagen.research.google/video/paper.pdf) paper). - thresholding (`bool`, defaults to `False`): - Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such - as Stable Diffusion. - dynamic_thresholding_ratio (`float`, defaults to 0.995): - The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. - sample_max_value (`float`, defaults to 1.0): - The threshold value for dynamic thresholding. Valid only when `thresholding=True`. - algorithm_type (`str`, defaults to `deis`): - The algorithm type for the solver. - lower_order_final (`bool`, defaults to `True`): - Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. - use_karras_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, - the sigmas are determined according to a sequence of noise levels {σi}. - use_exponential_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. - 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. - 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. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: Optional[np.ndarray] = None, - solver_order: int = 2, - prediction_type: str = "epsilon", - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - sample_max_value: float = 1.0, - algorithm_type: str = "deis", - solver_type: str = "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, - timestep_spacing: str = "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: - raise ValueError( - "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." - ) - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - # Currently we only support VP-type noise schedule - self.alpha_t = torch.sqrt(self.alphas_cumprod) - self.sigma_t = torch.sqrt(1 - self.alphas_cumprod) - self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t) - self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5 - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # settings for DEIS - if algorithm_type not in ["deis"]: - if algorithm_type in ["dpmsolver", "dpmsolver++"]: - self.register_to_config(algorithm_type="deis") - else: - raise NotImplementedError(f"{algorithm_type} is not implemented for {self.__class__}") - - if solver_type not in ["logrho"]: - if solver_type in ["midpoint", "heun", "bh1", "bh2"]: - self.register_to_config(solver_type="logrho") - else: - raise NotImplementedError(f"solver type {solver_type} is not implemented for {self.__class__}") - - # setable values - self.num_inference_steps = None - timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy() - self.timesteps = torch.from_numpy(timesteps) - self.model_outputs = [None] * solver_order - self.lower_order_nums = 0 - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 set_timesteps(self, num_inference_steps: int, device: Union[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 (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - """ - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = ( - np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps + 1) - .round()[::-1][:-1] - .copy() - .astype(np.int64) - ) - elif self.config.timestep_spacing == "leading": - step_ratio = self.config.num_train_timesteps // (num_inference_steps + 1) - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int64) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = np.arange(self.config.num_train_timesteps, 0, -step_ratio).round().copy().astype(np.int64) - timesteps -= 1 - else: - raise ValueError( - 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) - log_sigmas = np.log(sigmas) - if self.config.use_karras_sigmas: - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round() - sigmas = np.concatenate([sigmas, sigmas[-1:]]).astype(np.float32) - elif self.config.use_exponential_sigmas: - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - sigmas = np.concatenate([sigmas, sigmas[-1:]]).astype(np.float32) - elif self.config.use_beta_sigmas: - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - sigmas = np.concatenate([sigmas, sigmas[-1:]]).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() - timesteps = (sigmas * self.config.num_train_timesteps).copy() - sigmas = np.concatenate([sigmas, sigmas[-1:]]).astype(np.float32) - else: - sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) - sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5 - sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) - - self.sigmas = torch.from_numpy(sigmas) - self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) - - self.num_inference_steps = len(timesteps) - - self.model_outputs = [ - None, - ] * self.config.solver_order - self.lower_order_nums = 0 - - # add an index counter for schedulers that allow duplicated timesteps - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - """ - "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the - prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by - s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing - pixels from saturation at each step. We find that dynamic thresholding results in significantly better - photorealism as well as better image-text alignment, especially when using very large guidance weights." - - https://huggingface.co/papers/2205.11487 - """ - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half - - # Flatten sample for doing quantile calculation along each image - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - - abs_sample = sample.abs() # "a certain percentile absolute pixel value" - - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] - s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 - sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" - - sample = sample.reshape(batch_size, channels, *remaining_dims) - sample = sample.to(dtype) - - return sample - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma - log_sigma = np.log(np.maximum(sigma, 1e-10)) - - # get distribution - dists = log_sigma - log_sigmas[:, np.newaxis] - - # get sigmas range - low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) - high_idx = low_idx + 1 - - low = log_sigmas[low_idx] - high = log_sigmas[high_idx] - - # interpolate sigmas - w = (low - log_sigma) / (low - high) - w = np.clip(w, 0, 1) - - # transform interpolation to time range - t = (1 - w) * low_idx + w * high_idx - t = t.reshape(sigma.shape) - return t - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t - def _sigma_to_alpha_sigma_t(self, sigma): - if self.config.use_flow_sigmas: - alpha_t = 1 - sigma - sigma_t = sigma - else: - alpha_t = 1 / ((sigma**2 + 1) ** 0.5) - sigma_t = sigma * alpha_t - - 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: - """Constructs the noise schedule of Karras et al. (2022).""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - rho = 7.0 # 7.0 is the value used in the paper - ramp = np.linspace(0, 1, num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Constructs an exponential noise schedule.""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: - """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return sigmas - - def convert_model_output( - self, - model_output: torch.Tensor, - *args, - sample: torch.Tensor = None, - **kwargs, - ) -> torch.Tensor: - """ - Convert the model output to the corresponding type the DEIS 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`", - ) - - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - 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": - sigma_t = self.sigmas[self.step_index] - 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 DEISMultistepScheduler." - ) - - if self.config.thresholding: - x0_pred = self._threshold_sample(x0_pred) - - if self.config.algorithm_type == "deis": - return (sample - alpha_t * x0_pred) / sigma_t - else: - raise NotImplementedError("only support log-rho multistep deis now") - - def deis_first_order_update( - self, - model_output: torch.Tensor, - *args, - sample: torch.Tensor = None, - **kwargs, - ) -> torch.Tensor: - """ - One step for the first-order DEIS (equivalent to DDIM). - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. - timestep (`int`): - The current discrete timestep in the diffusion chain. - 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. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) - prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 2: - sample = args[2] - 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`", - ) - - 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`", - ) - - 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) - lambda_s = torch.log(alpha_s) - torch.log(sigma_s) - - h = lambda_t - lambda_s - if self.config.algorithm_type == "deis": - x_t = (alpha_t / alpha_s) * sample - (sigma_t * (torch.exp(h) - 1.0)) * model_output - else: - raise NotImplementedError("only support log-rho multistep deis now") - return x_t - - def multistep_deis_second_order_update( - self, - model_output_list: List[torch.Tensor], - *args, - sample: torch.Tensor = None, - **kwargs, - ) -> torch.Tensor: - """ - One step for the second-order multistep DEIS. - - Args: - 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. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None) - prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 2: - sample = args[2] - else: - raise ValueError("missing `sample` as a required keyword argument") - if timestep_list is not None: - deprecate( - "timestep_list", - "1.0.0", - "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", - ) - - 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`", - ) - - sigma_t, sigma_s0, sigma_s1 = ( - self.sigmas[self.step_index + 1], - self.sigmas[self.step_index], - self.sigmas[self.step_index - 1], - ) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1) - - m0, m1 = model_output_list[-1], model_output_list[-2] - - rho_t, rho_s0, rho_s1 = sigma_t / alpha_t, sigma_s0 / alpha_s0, sigma_s1 / alpha_s1 - - if self.config.algorithm_type == "deis": - - def ind_fn(t, b, c): - # Integrate[(log(t) - log(c)) / (log(b) - log(c)), {t}] - return t * (-np.log(c) + np.log(t) - 1) / (np.log(b) - np.log(c)) - - coef1 = ind_fn(rho_t, rho_s0, rho_s1) - ind_fn(rho_s0, rho_s0, rho_s1) - coef2 = ind_fn(rho_t, rho_s1, rho_s0) - ind_fn(rho_s0, rho_s1, rho_s0) - - x_t = alpha_t * (sample / alpha_s0 + coef1 * m0 + coef2 * m1) - return x_t - else: - raise NotImplementedError("only support log-rho multistep deis now") - - def multistep_deis_third_order_update( - self, - model_output_list: List[torch.Tensor], - *args, - sample: torch.Tensor = None, - **kwargs, - ) -> torch.Tensor: - """ - One step for the third-order multistep DEIS. - - Args: - 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. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - - timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None) - prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 2: - sample = args[2] - else: - raise ValueError("missing `sample` as a required keyword argument") - if timestep_list is not None: - deprecate( - "timestep_list", - "1.0.0", - "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", - ) - - 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`", - ) - - sigma_t, sigma_s0, sigma_s1, sigma_s2 = ( - self.sigmas[self.step_index + 1], - self.sigmas[self.step_index], - self.sigmas[self.step_index - 1], - self.sigmas[self.step_index - 2], - ) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1) - alpha_s2, sigma_s2 = self._sigma_to_alpha_sigma_t(sigma_s2) - - m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3] - - rho_t, rho_s0, rho_s1, rho_s2 = ( - sigma_t / alpha_t, - sigma_s0 / alpha_s0, - sigma_s1 / alpha_s1, - sigma_s2 / alpha_s2, - ) - - if self.config.algorithm_type == "deis": - - def ind_fn(t, b, c, d): - # Integrate[(log(t) - log(c))(log(t) - log(d)) / (log(b) - log(c))(log(b) - log(d)), {t}] - numerator = t * ( - np.log(c) * (np.log(d) - np.log(t) + 1) - - np.log(d) * np.log(t) - + np.log(d) - + np.log(t) ** 2 - - 2 * np.log(t) - + 2 - ) - denominator = (np.log(b) - np.log(c)) * (np.log(b) - np.log(d)) - return numerator / denominator - - coef1 = ind_fn(rho_t, rho_s0, rho_s1, rho_s2) - ind_fn(rho_s0, rho_s0, rho_s1, rho_s2) - coef2 = ind_fn(rho_t, rho_s1, rho_s2, rho_s0) - ind_fn(rho_s0, rho_s1, rho_s2, rho_s0) - coef3 = ind_fn(rho_t, rho_s2, rho_s0, rho_s1) - ind_fn(rho_s0, rho_s2, rho_s0, rho_s1) - - x_t = alpha_t * (sample / alpha_s0 + coef1 * m0 + coef2 * m1 + coef3 * m2) - - return x_t - else: - raise NotImplementedError("only support log-rho multistep deis now") - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep - def index_for_timestep(self, timestep, schedule_timesteps=None): - if schedule_timesteps is None: - schedule_timesteps = self.timesteps - - index_candidates = (schedule_timesteps == timestep).nonzero() - - if len(index_candidates) == 0: - step_index = len(self.timesteps) - 1 - # 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) - elif len(index_candidates) > 1: - step_index = index_candidates[1].item() - else: - step_index = index_candidates[0].item() - - return step_index - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index - def _init_step_index(self, timestep): - """ - Initialize the step_index counter for the scheduler. - """ - - 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( - self, - model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], - sample: torch.Tensor, - return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: - """ - Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with - the multistep DEIS. - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - return_dict (`bool`): - 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( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - if self.step_index is None: - self._init_step_index(timestep) - - lower_order_final = ( - (self.step_index == len(self.timesteps) - 1) and self.config.lower_order_final and len(self.timesteps) < 15 - ) - lower_order_second = ( - (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15 - ) - - model_output = self.convert_model_output(model_output, sample=sample) - for i in range(self.config.solver_order - 1): - self.model_outputs[i] = self.model_outputs[i + 1] - self.model_outputs[-1] = model_output - - if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final: - prev_sample = self.deis_first_order_update(model_output, sample=sample) - elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second: - prev_sample = self.multistep_deis_second_order_update(self.model_outputs, sample=sample) - else: - prev_sample = self.multistep_deis_third_order_update(self.model_outputs, sample=sample) - - if self.lower_order_nums < self.config.solver_order: - self.lower_order_nums += 1 - - # upon completion increase step index by one - self._step_index += 1 - - if not return_dict: - return (prev_sample,) - - return SchedulerOutput(prev_sample=prev_sample) - - def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: - """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. - - Args: - sample (`torch.Tensor`): - The input sample. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - noisy_samples = alpha_t * original_samples + sigma_t * noise - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpm_cogvideox.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpm_cogvideox.py deleted file mode 100644 index f7b63720e1071890fc7bb6bfe54b9ad808276093..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpm_cogvideox.py +++ /dev/null @@ -1,489 +0,0 @@ -# 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. - -# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion -# and https://github.com/hojonathanho/diffusion - -import math -from dataclasses import dataclass -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM -class DDIMSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -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. - - Returns: - `torch.Tensor`: rescaled betas with zero terminal SNR - """ - - alphas_bar_sqrt = alphas_cumprod.sqrt() - - # Store old values. - alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() - alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() - - # Shift so the last timestep is zero. - alphas_bar_sqrt -= alphas_bar_sqrt_T - - # Scale so the first timestep is back to the old value. - alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - - # Convert alphas_bar_sqrt to betas - alphas_bar = alphas_bar_sqrt**2 # Revert sqrt - - return alphas_bar - - -class CogVideoXDPMScheduler(SchedulerMixin, ConfigMixin): - """ - `DDIMScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with - non-Markovian guidance. - - 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`, 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`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - clip_sample (`bool`, defaults to `True`): - Clip the predicted sample for numerical stability. - clip_sample_range (`float`, defaults to 1.0): - The maximum magnitude for sample clipping. Valid only when `clip_sample=True`. - set_alpha_to_one (`bool`, defaults to `True`): - Each diffusion step uses the alphas product value at that step and at the previous one. For the final step - there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, - otherwise it uses the alpha value at step 0. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - 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://imagen.research.google/video/paper.pdf) paper). - thresholding (`bool`, defaults to `False`): - Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such - as Stable Diffusion. - dynamic_thresholding_ratio (`float`, defaults to 0.995): - The ratio for the dynamic thresholding method. Valid only when `thresholding=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 the [Common Diffusion Noise Schedules and - Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. - 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). - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - 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, - clip_sample: bool = True, - set_alpha_to_one: bool = True, - steps_offset: int = 0, - prediction_type: str = "epsilon", - clip_sample_range: float = 1.0, - sample_max_value: float = 1.0, - timestep_spacing: str = "leading", - rescale_betas_zero_snr: bool = False, - snr_shift_scale: float = 3.0, - ): - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - # Modify: SNR shift following SD3 - self.alphas_cumprod = self.alphas_cumprod / (snr_shift_scale + (1 - snr_shift_scale) * self.alphas_cumprod) - - # Rescale for zero SNR - if rescale_betas_zero_snr: - self.alphas_cumprod = rescale_zero_terminal_snr(self.alphas_cumprod) - - # At every step in ddim, we are looking into the previous alphas_cumprod - # For the final step, there is no previous alphas_cumprod because we are already at 0 - # `set_alpha_to_one` decides whether we set this parameter simply to one or - # whether we use the final alpha of the "non-previous" one. - self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0] - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # setable values - 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): - 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 - beta_prod_t_prev = 1 - alpha_prod_t_prev - - variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev) - - return variance - - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - def set_timesteps(self, num_inference_steps: int, device: Union[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. - """ - - if num_inference_steps > self.config.num_train_timesteps: - raise ValueError( - f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:" - f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle" - f" maximal {self.config.num_train_timesteps} timesteps." - ) - - self.num_inference_steps = num_inference_steps - - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = ( - np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps) - .round()[::-1] - .copy() - .astype(np.int64) - ) - elif self.config.timestep_spacing == "leading": - step_ratio = self.config.num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64) - timesteps -= 1 - else: - raise ValueError( - f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'leading' or 'trailing'." - ) - - self.timesteps = torch.from_numpy(timesteps).to(device) - - def get_variables(self, alpha_prod_t, alpha_prod_t_prev, alpha_prod_t_back=None): - 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 - - if alpha_prod_t_back is not None: - lamb_previous = ((alpha_prod_t_back / (1 - alpha_prod_t_back)) ** 0.5).log() - h_last = lamb - lamb_previous - r = h_last / h - return h, r, lamb, lamb_next - else: - return h, None, lamb, lamb_next - - def get_mult(self, h, r, alpha_prod_t, alpha_prod_t_prev, alpha_prod_t_back): - mult1 = ((1 - alpha_prod_t_prev) / (1 - alpha_prod_t)) ** 0.5 * (-h).exp() - mult2 = (-2 * h).expm1() * alpha_prod_t_prev**0.5 - - if alpha_prod_t_back is not None: - mult3 = 1 + 1 / (2 * r) - mult4 = 1 / (2 * r) - return mult1, mult2, mult3, mult4 - else: - return mult1, mult2 - - def step( - self, - model_output: torch.Tensor, - old_pred_original_sample: torch.Tensor, - timestep: int, - timestep_back: int, - sample: torch.Tensor, - eta: float = 0.0, - use_clipped_model_output: bool = False, - generator=None, - variance_noise: Optional[torch.Tensor] = None, - return_dict: bool = False, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - timestep (`float`): - 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. - generator (`torch.Generator`, *optional*): - A random number generator. - 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.DDIMSchedulerOutput`] or `tuple`. - - Returns: - [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] is returned, otherwise a - tuple is returned where the first element is the sample tensor. - - """ - 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" - ) - - # See formulas (12) and (16) of DDIM paper https://huggingface.co/papers/2010.02502 - # Ideally, read DDIM paper in-detail understanding - - # Notation ( -> - # - pred_noise_t -> e_theta(x_t, t) - # - pred_original_sample -> f_theta(x_t, t) or x_0 - # - std_dev_t -> sigma_t - # - eta -> η - # - pred_sample_direction -> "direction pointing to x_t" - # - pred_prev_sample -> "x_t-1" - - # 1. get previous step value (=t-1) - prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps - - # 2. compute alphas, betas - 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 - alpha_prod_t_back = self.alphas_cumprod[timestep_back] if timestep_back is not None else None - - beta_prod_t = 1 - alpha_prod_t - - # 3. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (12) from https://huggingface.co/papers/2010.02502 - # To make style tests pass, commented out `pred_epsilon` as it is an unused variable - if self.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - # pred_epsilon = model_output - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - # pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) - elif self.config.prediction_type == "v_prediction": - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - # pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" - " `v_prediction`" - ) - - h, r, lamb, lamb_next = self.get_variables(alpha_prod_t, alpha_prod_t_prev, alpha_prod_t_back) - mult = list(self.get_mult(h, r, alpha_prod_t, alpha_prod_t_prev, alpha_prod_t_back)) - mult_noise = (1 - alpha_prod_t_prev) ** 0.5 * (1 - (-2 * h).exp()) ** 0.5 - - noise = randn_tensor(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype) - prev_sample = mult[0] * sample - mult[1] * pred_original_sample + mult_noise * noise - - if old_pred_original_sample is None or prev_timestep < 0: - # Save a network evaluation if all noise levels are 0 or on the first 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) - x_advanced = mult[0] * sample - mult[1] * denoised_d + mult_noise * noise - - prev_sample = x_advanced - - if not return_dict: - return (prev_sample, pred_original_sample) - - return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as original_samples - # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement - # for the subsequent add_noise calls - self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype) - timesteps = timesteps.to(original_samples.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(original_samples.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise - return noisy_samples - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity - def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as sample - self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype) - timesteps = timesteps.to(sample.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(sample.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample - return velocity - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_multistep.py deleted file mode 100644 index 1a648af5a008fb01fc3c14a9f8f72bf588f50a69..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ /dev/null @@ -1,1174 +0,0 @@ -# Copyright 2025 TSAIL 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. - -# DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver - -import math -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import deprecate, is_scipy_available -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput - - -if is_scipy_available(): - import scipy.stats - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr -def rescale_zero_terminal_snr(betas): - """ - 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. - - Returns: - `torch.Tensor`: rescaled betas with zero terminal SNR - """ - # Convert betas to alphas_bar_sqrt - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - alphas_bar_sqrt = alphas_cumprod.sqrt() - - # Store old values. - alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() - alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() - - # Shift so the last timestep is zero. - alphas_bar_sqrt -= alphas_bar_sqrt_T - - # Scale so the first timestep is back to the old value. - alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - - # Convert alphas_bar_sqrt to betas - alphas_bar = alphas_bar_sqrt**2 # Revert sqrt - alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod - alphas = torch.cat([alphas_bar[0:1], alphas]) - betas = 1 - alphas - - return betas - - -class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): - """ - `DPMSolverMultistepScheduler` is a fast dedicated high-order solver for diffusion ODEs. - - 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`, 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`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, *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 - 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), `v_prediction` (see section 2.4 of [Imagen - Video](https://imagen.research.google/video/paper.pdf) paper), or `flow_prediction`. - thresholding (`bool`, defaults to `False`): - Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such - as Stable Diffusion. - dynamic_thresholding_ratio (`float`, defaults to 0.995): - The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. - sample_max_value (`float`, defaults to 1.0): - The threshold value for dynamic thresholding. Valid only when `thresholding=True` and - `algorithm_type="dpmsolver++"`. - algorithm_type (`str`, defaults to `dpmsolver++`): - Algorithm type for the solver; can be `dpmsolver`, `dpmsolver++`, `sde-dpmsolver` or `sde-dpmsolver++`. The - `dpmsolver` type implements the algorithms in the [DPMSolver](https://huggingface.co/papers/2206.00927) - paper, and the `dpmsolver++` type implements the algorithms in the - [DPMSolver++](https://huggingface.co/papers/2211.01095) paper. It is recommended to use `dpmsolver++` or - `sde-dpmsolver++` with `solver_order=2` for guided sampling like in Stable Diffusion. - 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. - 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. - euler_at_final (`bool`, defaults to `False`): - Whether to use Euler's method in the final step. It is a trade-off between numerical stability and detail - richness. This can stabilize the sampling of the SDE variant of DPMSolver for small number of inference - steps, but sometimes may result in blurring. - use_karras_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, - the sigmas are determined according to a sequence of noise levels {σi}. - use_exponential_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. - 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_lu_lambdas (`bool`, *optional*, defaults to `False`): - Whether to use the uniform-logSNR for step sizes proposed by Lu's DPM-Solver in the noise schedule during - the sampling process. If `True`, the sigmas and time steps are determined according to a sequence of - `lambda(t)`. - 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 shift value for the timestep schedule for flow matching. - 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. - 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. - variance_type (`str`, *optional*): - Set to "learned" or "learned_range" for diffusion models that predict variance. If set, the model's output - contains the predicted Gaussian variance. - 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. - 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 - 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). - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - 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, - solver_order: int = 2, - prediction_type: str = "epsilon", - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - sample_max_value: float = 1.0, - algorithm_type: str = "dpmsolver++", - solver_type: str = "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[str] = "zero", # "zero", "sigma_min" - lambda_min_clipped: float = -float("inf"), - variance_type: Optional[str] = None, - timestep_spacing: str = "linspace", - steps_offset: int = 0, - rescale_betas_zero_snr: 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: - 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) - - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - if rescale_betas_zero_snr: - self.betas = rescale_zero_terminal_snr(self.betas) - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - if rescale_betas_zero_snr: - # Close to 0 without being 0 so first sigma is not inf - # FP16 smallest positive subnormal works well here - self.alphas_cumprod[-1] = 2**-24 - - # Currently we only support VP-type noise schedule - self.alpha_t = torch.sqrt(self.alphas_cumprod) - self.sigma_t = torch.sqrt(1 - self.alphas_cumprod) - self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t) - self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5 - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # settings for DPM-Solver - if algorithm_type not in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"]: - if algorithm_type == "deis": - self.register_to_config(algorithm_type="dpmsolver++") - else: - raise NotImplementedError(f"{algorithm_type} is not implemented for {self.__class__}") - - if solver_type not in ["midpoint", "heun"]: - if solver_type in ["logrho", "bh1", "bh2"]: - self.register_to_config(solver_type="midpoint") - else: - raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}") - - if algorithm_type not in ["dpmsolver++", "sde-dpmsolver++"] and final_sigmas_type == "zero": - raise ValueError( - f"`final_sigmas_type` {final_sigmas_type} is not supported for `algorithm_type` {algorithm_type}. Please choose `sigma_min` instead." - ) - - # setable values - self.num_inference_steps = None - timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy() - self.timesteps = torch.from_numpy(timesteps) - self.model_outputs = [None] * solver_order - self.lower_order_nums = 0 - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - @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 set_timesteps( - self, - num_inference_steps: int = None, - device: Union[str, torch.device] = None, - timesteps: Optional[List[int]] = 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 timesteps are not moved. - 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. - """ - if num_inference_steps is None and timesteps is None: - raise ValueError("Must pass exactly one of `num_inference_steps` or `timesteps`.") - if num_inference_steps is not None and timesteps is not None: - raise ValueError("Can only pass one of `num_inference_steps` or `custom_timesteps`.") - if timesteps is not None and self.config.use_karras_sigmas: - raise ValueError("Cannot use `timesteps` with `config.use_karras_sigmas = True`") - if timesteps is not None and self.config.use_lu_lambdas: - raise ValueError("Cannot use `timesteps` with `config.use_lu_lambdas = True`") - if timesteps is not None and self.config.use_exponential_sigmas: - raise ValueError("Cannot set `timesteps` with `config.use_exponential_sigmas = True`.") - if timesteps is not None and self.config.use_beta_sigmas: - raise ValueError("Cannot set `timesteps` with `config.use_beta_sigmas = True`.") - - if timesteps is not None: - timesteps = np.array(timesteps).astype(np.int64) - else: - # Clipping the minimum of all lambda(t) for numerical stability. - # This is critical for cosine (squaredcos_cap_v2) noise schedule. - clipped_idx = torch.searchsorted(torch.flip(self.lambda_t, [0]), self.config.lambda_min_clipped) - last_timestep = ((self.config.num_train_timesteps - clipped_idx).numpy()).item() - - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = ( - np.linspace(0, last_timestep - 1, num_inference_steps + 1) - .round()[::-1][:-1] - .copy() - .astype(np.int64) - ) - elif self.config.timestep_spacing == "leading": - step_ratio = last_timestep // (num_inference_steps + 1) - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = ( - (np.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int64) - ) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = np.arange(last_timestep, 0, -step_ratio).round().copy().astype(np.int64) - timesteps -= 1 - else: - raise ValueError( - 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) - log_sigmas = np.log(sigmas) - - if self.config.use_karras_sigmas: - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - if self.config.beta_schedule != "squaredcos_cap_v2": - timesteps = timesteps.round() - elif self.config.use_lu_lambdas: - lambdas = np.flip(log_sigmas.copy()) - lambdas = self._convert_to_lu(in_lambdas=lambdas, num_inference_steps=num_inference_steps) - sigmas = np.exp(lambdas) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - if self.config.beta_schedule != "squaredcos_cap_v2": - timesteps = timesteps.round() - elif self.config.use_exponential_sigmas: - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - elif self.config.use_beta_sigmas: - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - 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() - timesteps = (sigmas * self.config.num_train_timesteps).copy() - else: - 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 - elif self.config.final_sigmas_type == "zero": - sigma_last = 0 - else: - raise ValueError( - f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" - ) - - sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) - - self.sigmas = torch.from_numpy(sigmas) - self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) - - self.num_inference_steps = len(timesteps) - - self.model_outputs = [ - None, - ] * self.config.solver_order - self.lower_order_nums = 0 - - # add an index counter for schedulers that allow duplicated timesteps - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - """ - "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the - prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by - s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing - pixels from saturation at each step. We find that dynamic thresholding results in significantly better - photorealism as well as better image-text alignment, especially when using very large guidance weights." - - https://huggingface.co/papers/2205.11487 - """ - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half - - # Flatten sample for doing quantile calculation along each image - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - - abs_sample = sample.abs() # "a certain percentile absolute pixel value" - - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] - s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 - sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" - - sample = sample.reshape(batch_size, channels, *remaining_dims) - sample = sample.to(dtype) - - return sample - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma - log_sigma = np.log(np.maximum(sigma, 1e-10)) - - # get distribution - dists = log_sigma - log_sigmas[:, np.newaxis] - - # get sigmas range - low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) - high_idx = low_idx + 1 - - low = log_sigmas[low_idx] - high = log_sigmas[high_idx] - - # interpolate sigmas - w = (low - log_sigma) / (low - high) - w = np.clip(w, 0, 1) - - # transform interpolation to time range - t = (1 - w) * low_idx + w * high_idx - t = t.reshape(sigma.shape) - return t - - def _sigma_to_alpha_sigma_t(self, sigma): - if self.config.use_flow_sigmas: - alpha_t = 1 - sigma - sigma_t = sigma - else: - alpha_t = 1 / ((sigma**2 + 1) ** 0.5) - sigma_t = sigma * alpha_t - - 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: - """Constructs the noise schedule of Karras et al. (2022).""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - rho = 7.0 # 7.0 is the value used in the paper - ramp = np.linspace(0, 1, num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - def _convert_to_lu(self, in_lambdas: torch.Tensor, num_inference_steps) -> torch.Tensor: - """Constructs the noise schedule of Lu et al. (2022).""" - - lambda_min: float = in_lambdas[-1].item() - lambda_max: float = in_lambdas[0].item() - - rho = 1.0 # 1.0 is the value used in the paper - ramp = np.linspace(0, 1, num_inference_steps) - min_inv_rho = lambda_min ** (1 / rho) - max_inv_rho = lambda_max ** (1 / rho) - lambdas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return lambdas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Constructs an exponential noise schedule.""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: - """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return sigmas - - def convert_model_output( - self, - model_output: torch.Tensor, - *args, - sample: torch.Tensor = None, - **kwargs, - ) -> torch.Tensor: - """ - Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is - designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an - integral of the data prediction model. - - - - The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise - prediction and data prediction models. - - - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. - 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`", - ) - - # DPM-Solver++ needs to solve an integral of the data prediction model. - if self.config.algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]: - if self.config.prediction_type == "epsilon": - # DPM-Solver and DPM-Solver++ only need the "mean" output. - if self.config.variance_type in ["learned", "learned_range"]: - model_output = model_output[:, :3] - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - 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": - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - x0_pred = alpha_t * sample - sigma_t * model_output - elif self.config.prediction_type == "flow_prediction": - sigma_t = self.sigmas[self.step_index] - 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 DPMSolverMultistepScheduler." - ) - - if self.config.thresholding: - x0_pred = self._threshold_sample(x0_pred) - - return x0_pred - - # DPM-Solver needs to solve an integral of the noise prediction model. - elif self.config.algorithm_type in ["dpmsolver", "sde-dpmsolver"]: - if self.config.prediction_type == "epsilon": - # DPM-Solver and DPM-Solver++ only need the "mean" output. - if self.config.variance_type in ["learned", "learned_range"]: - epsilon = model_output[:, :3] - else: - epsilon = model_output - elif self.config.prediction_type == "sample": - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - epsilon = (sample - alpha_t * model_output) / sigma_t - elif self.config.prediction_type == "v_prediction": - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - epsilon = alpha_t * model_output + sigma_t * sample - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" - " `v_prediction` for the DPMSolverMultistepScheduler." - ) - - if self.config.thresholding: - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - x0_pred = (sample - sigma_t * epsilon) / alpha_t - x0_pred = self._threshold_sample(x0_pred) - epsilon = (sample - alpha_t * x0_pred) / sigma_t - - return epsilon - - def dpm_solver_first_order_update( - self, - model_output: torch.Tensor, - *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, - **kwargs, - ) -> torch.Tensor: - """ - One step for the first-order DPMSolver (equivalent to DDIM). - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. - sample (`torch.Tensor`): - A current instance of a sample created by the diffusion process. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) - prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 2: - sample = args[2] - 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`", - ) - - 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`", - ) - - 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) - lambda_s = torch.log(alpha_s) - torch.log(sigma_s) - - h = lambda_t - lambda_s - if self.config.algorithm_type == "dpmsolver++": - x_t = (sigma_t / sigma_s) * sample - (alpha_t * (torch.exp(-h) - 1.0)) * model_output - elif self.config.algorithm_type == "dpmsolver": - x_t = (alpha_t / alpha_s) * sample - (sigma_t * (torch.exp(h) - 1.0)) * model_output - elif self.config.algorithm_type == "sde-dpmsolver++": - assert noise is not None - x_t = ( - (sigma_t / sigma_s * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - elif self.config.algorithm_type == "sde-dpmsolver": - assert noise is not None - x_t = ( - (alpha_t / alpha_s) * sample - - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * model_output - + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise - ) - return x_t - - def multistep_dpm_solver_second_order_update( - self, - model_output_list: List[torch.Tensor], - *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, - **kwargs, - ) -> torch.Tensor: - """ - One step for the second-order multistep DPMSolver. - - Args: - 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. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None) - prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 2: - sample = args[2] - else: - raise ValueError("missing `sample` as a required keyword argument") - if timestep_list is not None: - deprecate( - "timestep_list", - "1.0.0", - "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", - ) - - 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`", - ) - - sigma_t, sigma_s0, sigma_s1 = ( - self.sigmas[self.step_index + 1], - self.sigmas[self.step_index], - self.sigmas[self.step_index - 1], - ) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1) - - m0, m1 = model_output_list[-1], model_output_list[-2] - - h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1 - r0 = h_0 / h - D0, D1 = m0, (1.0 / r0) * (m0 - m1) - if self.config.algorithm_type == "dpmsolver++": - # See https://huggingface.co/papers/2211.01095 for detailed derivations - if self.config.solver_type == "midpoint": - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - - 0.5 * (alpha_t * (torch.exp(-h) - 1.0)) * D1 - ) - elif self.config.solver_type == "heun": - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1 - ) - elif self.config.algorithm_type == "dpmsolver": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - if self.config.solver_type == "midpoint": - x_t = ( - (alpha_t / alpha_s0) * sample - - (sigma_t * (torch.exp(h) - 1.0)) * D0 - - 0.5 * (sigma_t * (torch.exp(h) - 1.0)) * D1 - ) - elif self.config.solver_type == "heun": - x_t = ( - (alpha_t / alpha_s0) * sample - - (sigma_t * (torch.exp(h) - 1.0)) * D0 - - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1 - ) - elif self.config.algorithm_type == "sde-dpmsolver++": - assert noise is not None - if self.config.solver_type == "midpoint": - x_t = ( - (sigma_t / sigma_s0 * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 - + 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1 - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - elif self.config.solver_type == "heun": - x_t = ( - (sigma_t / sigma_s0 * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 - + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1 - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - elif self.config.algorithm_type == "sde-dpmsolver": - assert noise is not None - if self.config.solver_type == "midpoint": - x_t = ( - (alpha_t / alpha_s0) * sample - - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * D0 - - (sigma_t * (torch.exp(h) - 1.0)) * D1 - + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise - ) - elif self.config.solver_type == "heun": - x_t = ( - (alpha_t / alpha_s0) * sample - - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * D0 - - 2.0 * (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1 - + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise - ) - return x_t - - def multistep_dpm_solver_third_order_update( - self, - model_output_list: List[torch.Tensor], - *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, - **kwargs, - ) -> torch.Tensor: - """ - One step for the third-order multistep DPMSolver. - - Args: - 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. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - - timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None) - prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 2: - sample = args[2] - else: - raise ValueError("missing `sample` as a required keyword argument") - if timestep_list is not None: - deprecate( - "timestep_list", - "1.0.0", - "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", - ) - - 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`", - ) - - sigma_t, sigma_s0, sigma_s1, sigma_s2 = ( - self.sigmas[self.step_index + 1], - self.sigmas[self.step_index], - self.sigmas[self.step_index - 1], - self.sigmas[self.step_index - 2], - ) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1) - alpha_s2, sigma_s2 = self._sigma_to_alpha_sigma_t(sigma_s2) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1) - lambda_s2 = torch.log(alpha_s2) - torch.log(sigma_s2) - - m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3] - - h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2 - r0, r1 = h_0 / h, h_1 / h - D0 = m0 - D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2) - D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1) - D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1) - if self.config.algorithm_type == "dpmsolver++": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1 - - (alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2 - ) - elif self.config.algorithm_type == "dpmsolver": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - x_t = ( - (alpha_t / alpha_s0) * sample - - (sigma_t * (torch.exp(h) - 1.0)) * D0 - - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1 - - (sigma_t * ((torch.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2 - ) - elif self.config.algorithm_type == "sde-dpmsolver++": - assert noise is not None - x_t = ( - (sigma_t / sigma_s0 * torch.exp(-h)) * sample - + (alpha_t * (1.0 - torch.exp(-2.0 * h))) * D0 - + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1 - + (alpha_t * ((1.0 - torch.exp(-2.0 * h) - 2.0 * h) / (2.0 * h) ** 2 - 0.5)) * D2 - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - return x_t - - def index_for_timestep(self, timestep, schedule_timesteps=None): - if schedule_timesteps is None: - schedule_timesteps = self.timesteps - - index_candidates = (schedule_timesteps == timestep).nonzero() - - if len(index_candidates) == 0: - step_index = len(self.timesteps) - 1 - # 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) - elif len(index_candidates) > 1: - step_index = index_candidates[1].item() - else: - step_index = index_candidates[0].item() - - return step_index - - def _init_step_index(self, timestep): - """ - Initialize the step_index counter for the scheduler. - """ - - 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( - self, - model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], - sample: torch.Tensor, - generator=None, - variance_noise: Optional[torch.Tensor] = None, - return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: - """ - Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with - the multistep DPMSolver. - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - generator (`torch.Generator`, *optional*): - A random number generator. - variance_noise (`torch.Tensor`): - Alternative to generating noise with `generator` by directly providing the noise for the variance - itself. Useful for methods such as [`LEdits++`]. - return_dict (`bool`): - 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( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - if self.step_index is None: - self._init_step_index(timestep) - - # Improve numerical stability for small number of steps - lower_order_final = (self.step_index == len(self.timesteps) - 1) and ( - self.config.euler_at_final - or (self.config.lower_order_final and len(self.timesteps) < 15) - or self.config.final_sigmas_type == "zero" - ) - lower_order_second = ( - (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15 - ) - - model_output = self.convert_model_output(model_output, sample=sample) - for i in range(self.config.solver_order - 1): - self.model_outputs[i] = self.model_outputs[i + 1] - self.model_outputs[-1] = model_output - - # Upcast to avoid precision issues when computing prev_sample - 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 - ) - elif self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]: - noise = variance_noise.to(device=model_output.device, dtype=torch.float32) - else: - noise = None - - if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final: - prev_sample = self.dpm_solver_first_order_update(model_output, sample=sample, noise=noise) - elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second: - prev_sample = self.multistep_dpm_solver_second_order_update(self.model_outputs, sample=sample, noise=noise) - else: - prev_sample = self.multistep_dpm_solver_third_order_update(self.model_outputs, sample=sample, noise=noise) - - if self.lower_order_nums < self.config.solver_order: - self.lower_order_nums += 1 - - # Cast sample back to expected 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 SchedulerOutput(prev_sample=prev_sample) - - def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: - """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. - - Args: - sample (`torch.Tensor`): - The input sample. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - noisy_samples = alpha_t * original_samples + sigma_t * noise - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py deleted file mode 100644 index 71b9960bf2ff21d116caffce05c986d8b3a898e4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py +++ /dev/null @@ -1,645 +0,0 @@ -# Copyright 2025 TSAIL 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. - -# 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 .scheduling_utils_flax import ( - CommonSchedulerState, - FlaxKarrasDiffusionSchedulers, - FlaxSchedulerMixin, - FlaxSchedulerOutput, - add_noise_common, -) - - -@flax.struct.dataclass -class DPMSolverMultistepSchedulerState: - common: CommonSchedulerState - alpha_t: jnp.ndarray - sigma_t: jnp.ndarray - lambda_t: jnp.ndarray - - # setable values - init_noise_sigma: jnp.ndarray - timesteps: jnp.ndarray - num_inference_steps: Optional[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 - - @classmethod - def create( - cls, - common: CommonSchedulerState, - alpha_t: jnp.ndarray, - sigma_t: jnp.ndarray, - lambda_t: jnp.ndarray, - init_noise_sigma: jnp.ndarray, - timesteps: jnp.ndarray, - ): - return cls( - common=common, - alpha_t=alpha_t, - sigma_t=sigma_t, - lambda_t=lambda_t, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - ) - - -@dataclass -class FlaxDPMSolverMultistepSchedulerOutput(FlaxSchedulerOutput): - state: DPMSolverMultistepSchedulerState - - -class FlaxDPMSolverMultistepScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - DPM-Solver (and the improved version DPM-Solver++) is a fast dedicated high-order solver for diffusion ODEs with - the convergence order guarantee. Empirically, sampling by DPM-Solver with only 20 steps can generate high-quality - samples, and it can generate quite good samples even in only 10 steps. - - For more details, see the original paper: https://huggingface.co/papers/2206.00927 and - https://huggingface.co/papers/2211.01095 - - Currently, we support the multistep DPM-Solver for both noise prediction models and data prediction models. We - recommend to use `solver_order=2` for guided sampling, and `solver_order=3` for unconditional sampling. - - We also support the "dynamic thresholding" method in Imagen (https://huggingface.co/papers/2205.11487). For - pixel-space diffusion models, you can set both `algorithm_type="dpmsolver++"` and `thresholding=True` to use the - dynamic thresholding. Note that the thresholding method is unsuitable for latent-space diffusion models (such as - stable-diffusion). - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - For more details, see the original paper: https://huggingface.co/papers/2206.00927 and - https://huggingface.co/papers/2211.01095 - - 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 - `linear`, `scaled_linear`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, optional): - option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. - solver_order (`int`, default `2`): - the order of DPM-Solver; can be `1` or `2` or `3`. We recommend to use `solver_order=2` for guided - sampling, and `solver_order=3` for unconditional sampling. - prediction_type (`str`, default `epsilon`): - indicates whether the model predicts the noise (epsilon), or the data / `x0`. One of `epsilon`, `sample`, - or `v-prediction`. - thresholding (`bool`, default `False`): - whether to use the "dynamic thresholding" method (introduced by Imagen, - https://huggingface.co/papers/2205.11487). For pixel-space diffusion models, you can set both - `algorithm_type=dpmsolver++` and `thresholding=True` to use the dynamic thresholding. 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 - (https://huggingface.co/papers/2205.11487). - sample_max_value (`float`, default `1.0`): - the threshold value for dynamic thresholding. Valid only when `thresholding=True` and - `algorithm_type="dpmsolver++`. - algorithm_type (`str`, default `dpmsolver++`): - the algorithm type for the solver. Either `dpmsolver` or `dpmsolver++`. The `dpmsolver` type implements the - algorithms in https://huggingface.co/papers/2206.00927, and the `dpmsolver++` type implements the - algorithms in https://huggingface.co/papers/2211.01095. We recommend to use `dpmsolver++` with - `solver_order=2` for guided sampling (e.g. stable-diffusion). - solver_type (`str`, default `midpoint`): - the solver type for the second-order solver. Either `midpoint` or `heun`. The solver type slightly affects - the sample quality, especially for small number of steps. We empirically find that `midpoint` solvers are - slightly better, so we recommend to use the `midpoint` type. - lower_order_final (`bool`, default `True`): - whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. We empirically - find this trick can stabilize the sampling of DPM-Solver for steps < 15, especially for steps <= 10. - 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. - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - the `dtype` used for params and computation. - """ - - _compatibles = [e.name for e in FlaxKarrasDiffusionSchedulers] - - dtype: jnp.dtype - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: Optional[jnp.ndarray] = None, - solver_order: int = 2, - prediction_type: str = "epsilon", - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - sample_max_value: float = 1.0, - algorithm_type: str = "dpmsolver++", - solver_type: str = "midpoint", - lower_order_final: bool = True, - timestep_spacing: str = "linspace", - dtype: jnp.dtype = jnp.float32, - ): - self.dtype = dtype - - def create_state(self, common: Optional[CommonSchedulerState] = None) -> DPMSolverMultistepSchedulerState: - if common is None: - common = CommonSchedulerState.create(self) - - # Currently we only support VP-type noise schedule - alpha_t = jnp.sqrt(common.alphas_cumprod) - sigma_t = jnp.sqrt(1 - common.alphas_cumprod) - lambda_t = jnp.log(alpha_t) - jnp.log(sigma_t) - - # settings for DPM-Solver - if self.config.algorithm_type not in ["dpmsolver", "dpmsolver++"]: - raise NotImplementedError(f"{self.config.algorithm_type} is not implemented for {self.__class__}") - if self.config.solver_type not in ["midpoint", "heun"]: - raise NotImplementedError(f"{self.config.solver_type} is not implemented for {self.__class__}") - - # standard deviation of the initial noise distribution - init_noise_sigma = jnp.array(1.0, dtype=self.dtype) - - timesteps = jnp.arange(0, self.config.num_train_timesteps).round()[::-1] - - return DPMSolverMultistepSchedulerState.create( - common=common, - alpha_t=alpha_t, - sigma_t=sigma_t, - lambda_t=lambda_t, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - ) - - def set_timesteps( - 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. - - Args: - state (`DPMSolverMultistepSchedulerState`): - 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`): - the shape of the samples to be generated. - """ - last_timestep = self.config.num_train_timesteps - if self.config.timestep_spacing == "linspace": - timesteps = ( - jnp.linspace(0, last_timestep - 1, num_inference_steps + 1).round()[::-1][:-1].astype(jnp.int32) - ) - elif self.config.timestep_spacing == "leading": - step_ratio = last_timestep // (num_inference_steps + 1) - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = ( - (jnp.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(jnp.int32) - ) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = jnp.arange(last_timestep, 0, -step_ratio).round().copy().astype(jnp.int32) - timesteps -= 1 - else: - raise ValueError( - f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'." - ) - - # initial running values - - model_outputs = jnp.zeros((self.config.solver_order,) + shape, dtype=self.dtype) - lower_order_nums = jnp.int32(0) - prev_timestep = jnp.int32(-1) - cur_sample = jnp.zeros(shape, dtype=self.dtype) - - return state.replace( - num_inference_steps=num_inference_steps, - timesteps=timesteps, - model_outputs=model_outputs, - lower_order_nums=lower_order_nums, - prev_timestep=prev_timestep, - cur_sample=cur_sample, - ) - - def convert_model_output( - self, - state: DPMSolverMultistepSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - ) -> jnp.ndarray: - """ - Convert the model output to the corresponding type that the algorithm (DPM-Solver / DPM-Solver++) needs. - - DPM-Solver is designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to - discretize an integral of the data prediction model. So we need to first convert the model output to the - corresponding type to match the algorithm. - - Note that the algorithm type and the model type is decoupled. That is to say, we can use either DPM-Solver or - DPM-Solver++ for both noise prediction model and data prediction model. - - Args: - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - - Returns: - `jnp.ndarray`: the converted model output. - """ - # DPM-Solver++ needs to solve an integral of the data prediction model. - if self.config.algorithm_type == "dpmsolver++": - if self.config.prediction_type == "epsilon": - alpha_t, sigma_t = state.alpha_t[timestep], state.sigma_t[timestep] - 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": - alpha_t, sigma_t = state.alpha_t[timestep], state.sigma_t[timestep] - x0_pred = alpha_t * sample - sigma_t * model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, " - " or `v_prediction` for the FlaxDPMSolverMultistepScheduler." - ) - - 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)) - ) - dynamic_max_val = jnp.maximum( - 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 - # DPM-Solver needs to solve an integral of the noise prediction model. - elif self.config.algorithm_type == "dpmsolver": - if self.config.prediction_type == "epsilon": - return model_output - elif self.config.prediction_type == "sample": - alpha_t, sigma_t = state.alpha_t[timestep], state.sigma_t[timestep] - epsilon = (sample - alpha_t * model_output) / sigma_t - return epsilon - elif self.config.prediction_type == "v_prediction": - alpha_t, sigma_t = state.alpha_t[timestep], state.sigma_t[timestep] - 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 FlaxDPMSolverMultistepScheduler." - ) - - def dpm_solver_first_order_update( - self, - state: DPMSolverMultistepSchedulerState, - model_output: jnp.ndarray, - timestep: int, - prev_timestep: int, - sample: jnp.ndarray, - ) -> jnp.ndarray: - """ - One step for the first-order DPM-Solver (equivalent to DDIM). - - See https://huggingface.co/papers/2206.00927 for the detailed derivation. - - Args: - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - prev_timestep (`int`): previous discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - - Returns: - `jnp.ndarray`: the sample tensor at the previous timestep. - """ - t, s0 = prev_timestep, timestep - m0 = model_output - lambda_t, lambda_s = state.lambda_t[t], state.lambda_t[s0] - alpha_t, alpha_s = state.alpha_t[t], state.alpha_t[s0] - sigma_t, sigma_s = state.sigma_t[t], state.sigma_t[s0] - h = lambda_t - lambda_s - if self.config.algorithm_type == "dpmsolver++": - x_t = (sigma_t / sigma_s) * sample - (alpha_t * (jnp.exp(-h) - 1.0)) * m0 - elif self.config.algorithm_type == "dpmsolver": - x_t = (alpha_t / alpha_s) * sample - (sigma_t * (jnp.exp(h) - 1.0)) * m0 - return x_t - - def multistep_dpm_solver_second_order_update( - self, - state: DPMSolverMultistepSchedulerState, - model_output_list: jnp.ndarray, - timestep_list: List[int], - prev_timestep: int, - sample: jnp.ndarray, - ) -> jnp.ndarray: - """ - One step for the second-order multistep DPM-Solver. - - Args: - 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. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - - Returns: - `jnp.ndarray`: the sample tensor at the previous timestep. - """ - 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] - 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 - r0 = h_0 / h - D0, D1 = m0, (1.0 / r0) * (m0 - m1) - if self.config.algorithm_type == "dpmsolver++": - # See https://huggingface.co/papers/2211.01095 for detailed derivations - if self.config.solver_type == "midpoint": - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (jnp.exp(-h) - 1.0)) * D0 - - 0.5 * (alpha_t * (jnp.exp(-h) - 1.0)) * D1 - ) - elif self.config.solver_type == "heun": - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (jnp.exp(-h) - 1.0)) * D0 - + (alpha_t * ((jnp.exp(-h) - 1.0) / h + 1.0)) * D1 - ) - elif self.config.algorithm_type == "dpmsolver": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - if self.config.solver_type == "midpoint": - x_t = ( - (alpha_t / alpha_s0) * sample - - (sigma_t * (jnp.exp(h) - 1.0)) * D0 - - 0.5 * (sigma_t * (jnp.exp(h) - 1.0)) * D1 - ) - elif self.config.solver_type == "heun": - x_t = ( - (alpha_t / alpha_s0) * sample - - (sigma_t * (jnp.exp(h) - 1.0)) * D0 - - (sigma_t * ((jnp.exp(h) - 1.0) / h - 1.0)) * D1 - ) - return x_t - - def multistep_dpm_solver_third_order_update( - self, - state: DPMSolverMultistepSchedulerState, - model_output_list: jnp.ndarray, - timestep_list: List[int], - prev_timestep: int, - sample: jnp.ndarray, - ) -> jnp.ndarray: - """ - One step for the third-order multistep DPM-Solver. - - Args: - 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. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - - 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] - 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], - state.lambda_t[s0], - state.lambda_t[s1], - state.lambda_t[s2], - ) - 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, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2 - r0, r1 = h_0 / h, h_1 / h - D0 = m0 - D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2) - D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1) - D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1) - if self.config.algorithm_type == "dpmsolver++": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (jnp.exp(-h) - 1.0)) * D0 - + (alpha_t * ((jnp.exp(-h) - 1.0) / h + 1.0)) * D1 - - (alpha_t * ((jnp.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2 - ) - elif self.config.algorithm_type == "dpmsolver": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - x_t = ( - (alpha_t / alpha_s0) * sample - - (sigma_t * (jnp.exp(h) - 1.0)) * D0 - - (sigma_t * ((jnp.exp(h) - 1.0) / h - 1.0)) * D1 - - (sigma_t * ((jnp.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2 - ) - return x_t - - def step( - self, - state: DPMSolverMultistepSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - return_dict: bool = True, - ) -> Union[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). - - Args: - state (`DPMSolverMultistepSchedulerState`): - the `FlaxDPMSolverMultistepScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - return_dict (`bool`): option for returning tuple rather than FlaxDPMSolverMultistepSchedulerOutput class - - Returns: - [`FlaxDPMSolverMultistepSchedulerOutput`] or `tuple`: [`FlaxDPMSolverMultistepSchedulerOutput`] if - `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - if state.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - (step_index,) = jnp.where(state.timesteps == timestep, size=1) - step_index = step_index[0] - - prev_timestep = jax.lax.select(step_index == len(state.timesteps) - 1, 0, state.timesteps[step_index + 1]) - - model_output = self.convert_model_output(state, model_output, timestep, sample) - - model_outputs_new = jnp.roll(state.model_outputs, -1, axis=0) - model_outputs_new = model_outputs_new.at[-1].set(model_output) - state = state.replace( - model_outputs=model_outputs_new, - prev_timestep=prev_timestep, - cur_sample=sample, - ) - - def step_1(state: DPMSolverMultistepSchedulerState) -> jnp.ndarray: - return self.dpm_solver_first_order_update( - state, - state.model_outputs[-1], - state.timesteps[step_index], - state.prev_timestep, - state.cur_sample, - ) - - def step_23(state: DPMSolverMultistepSchedulerState) -> jnp.ndarray: - def step_2(state: DPMSolverMultistepSchedulerState) -> jnp.ndarray: - timestep_list = jnp.array([state.timesteps[step_index - 1], state.timesteps[step_index]]) - return self.multistep_dpm_solver_second_order_update( - state, - state.model_outputs, - timestep_list, - state.prev_timestep, - state.cur_sample, - ) - - def step_3(state: DPMSolverMultistepSchedulerState) -> jnp.ndarray: - timestep_list = jnp.array( - [ - state.timesteps[step_index - 2], - state.timesteps[step_index - 1], - state.timesteps[step_index], - ] - ) - return self.multistep_dpm_solver_third_order_update( - state, - state.model_outputs, - timestep_list, - state.prev_timestep, - state.cur_sample, - ) - - step_2_output = step_2(state) - step_3_output = step_3(state) - - if self.config.solver_order == 2: - return step_2_output - elif self.config.lower_order_final and len(state.timesteps) < 15: - return jax.lax.select( - state.lower_order_nums < 2, - step_2_output, - jax.lax.select( - step_index == len(state.timesteps) - 2, - step_2_output, - step_3_output, - ), - ) - else: - return jax.lax.select( - state.lower_order_nums < 2, - step_2_output, - step_3_output, - ) - - step_1_output = step_1(state) - step_23_output = step_23(state) - - if self.config.solver_order == 1: - prev_sample = step_1_output - - elif self.config.lower_order_final and len(state.timesteps) < 15: - prev_sample = jax.lax.select( - state.lower_order_nums < 1, - step_1_output, - jax.lax.select( - step_index == len(state.timesteps) - 1, - step_1_output, - step_23_output, - ), - ) - - else: - prev_sample = jax.lax.select( - state.lower_order_nums < 1, - step_1_output, - step_23_output, - ) - - state = state.replace( - lower_order_nums=jnp.minimum(state.lower_order_nums + 1, self.config.solver_order), - ) - - if not return_dict: - return (prev_sample, state) - - return FlaxDPMSolverMultistepSchedulerOutput(prev_sample=prev_sample, state=state) - - def scale_model_input( - self, state: DPMSolverMultistepSchedulerState, sample: jnp.ndarray, timestep: Optional[int] = None - ) -> jnp.ndarray: - """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. - - Args: - state (`DPMSolverMultistepSchedulerState`): - the `FlaxDPMSolverMultistepScheduler` state data class instance. - sample (`jnp.ndarray`): input sample - timestep (`int`, optional): current timestep - - Returns: - `jnp.ndarray`: scaled input sample - """ - return sample - - def add_noise( - self, - state: DPMSolverMultistepSchedulerState, - original_samples: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - return add_noise_common(state.common, original_samples, noise, timesteps) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py deleted file mode 100644 index 9ec95885111172344cb953d96a70eee0cd12776c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py +++ /dev/null @@ -1,1027 +0,0 @@ -# Copyright 2025 TSAIL 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. - -# DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver - -import math -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import deprecate, is_scipy_available -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput - - -if is_scipy_available(): - import scipy.stats - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin): - """ - `DPMSolverMultistepInverseScheduler` is the reverse scheduler of [`DPMSolverMultistepScheduler`]. - - 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`, 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`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, *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 - 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://imagen.research.google/video/paper.pdf) paper). - thresholding (`bool`, defaults to `False`): - Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such - as Stable Diffusion. - dynamic_thresholding_ratio (`float`, defaults to 0.995): - The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. - sample_max_value (`float`, defaults to 1.0): - The threshold value for dynamic thresholding. Valid only when `thresholding=True` and - `algorithm_type="dpmsolver++"`. - algorithm_type (`str`, defaults to `dpmsolver++`): - Algorithm type for the solver; can be `dpmsolver`, `dpmsolver++`, `sde-dpmsolver` or `sde-dpmsolver++`. The - `dpmsolver` type implements the algorithms in the [DPMSolver](https://huggingface.co/papers/2206.00927) - paper, and the `dpmsolver++` type implements the algorithms in the - [DPMSolver++](https://huggingface.co/papers/2211.01095) paper. It is recommended to use `dpmsolver++` or - `sde-dpmsolver++` with `solver_order=2` for guided sampling like in Stable Diffusion. - 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. - 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. - euler_at_final (`bool`, defaults to `False`): - Whether to use Euler's method in the final step. It is a trade-off between numerical stability and detail - richness. This can stabilize the sampling of the SDE variant of DPMSolver for small number of inference - steps, but sometimes may result in blurring. - use_karras_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, - the sigmas are determined according to a sequence of noise levels {σi}. - use_exponential_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. - 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. - 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. - variance_type (`str`, *optional*): - Set to "learned" or "learned_range" for diffusion models that predict variance. If set, the model's output - contains the predicted Gaussian variance. - 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. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - 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, - solver_order: int = 2, - prediction_type: str = "epsilon", - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - sample_max_value: float = 1.0, - algorithm_type: str = "dpmsolver++", - solver_type: str = "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, - lambda_min_clipped: float = -float("inf"), - variance_type: Optional[str] = None, - timestep_spacing: str = "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: - 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) - - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - # Currently we only support VP-type noise schedule - self.alpha_t = torch.sqrt(self.alphas_cumprod) - self.sigma_t = torch.sqrt(1 - self.alphas_cumprod) - self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t) - self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5 - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # settings for DPM-Solver - if algorithm_type not in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"]: - if algorithm_type == "deis": - self.register_to_config(algorithm_type="dpmsolver++") - else: - raise NotImplementedError(f"{algorithm_type} is not implemented for {self.__class__}") - - if solver_type not in ["midpoint", "heun"]: - if solver_type in ["logrho", "bh1", "bh2"]: - self.register_to_config(solver_type="midpoint") - else: - raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}") - - # setable values - self.num_inference_steps = None - timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32).copy() - self.timesteps = torch.from_numpy(timesteps) - self.model_outputs = [None] * solver_order - self.lower_order_nums = 0 - self._step_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - self.use_karras_sigmas = use_karras_sigmas - self.use_exponential_sigmas = use_exponential_sigmas - self.use_beta_sigmas = use_beta_sigmas - - @property - def step_index(self): - """ - The index counter for current timestep. It will increase 1 after each scheduler step. - """ - return self._step_index - - def set_timesteps(self, num_inference_steps: int = None, device: Union[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 (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - """ - # Clipping the minimum of all lambda(t) for numerical stability. - # This is critical for cosine (squaredcos_cap_v2) noise schedule. - clipped_idx = torch.searchsorted(torch.flip(self.lambda_t, [0]), self.config.lambda_min_clipped).item() - self.noisiest_timestep = self.config.num_train_timesteps - 1 - clipped_idx - - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = ( - np.linspace(0, self.noisiest_timestep, num_inference_steps + 1).round()[:-1].copy().astype(np.int64) - ) - elif self.config.timestep_spacing == "leading": - step_ratio = (self.noisiest_timestep + 1) // (num_inference_steps + 1) - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps + 1) * step_ratio).round()[:-1].copy().astype(np.int64) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = np.arange(self.noisiest_timestep + 1, 0, -step_ratio).round()[::-1].copy().astype(np.int64) - timesteps -= 1 - else: - raise ValueError( - 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) - log_sigmas = np.log(sigmas) - - if self.config.use_karras_sigmas: - sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round() - timesteps = timesteps.copy().astype(np.int64) - sigmas = np.concatenate([sigmas, sigmas[-1:]]).astype(np.float32) - elif self.config.use_exponential_sigmas: - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - sigmas = np.concatenate([sigmas, sigmas[-1:]]).astype(np.float32) - elif self.config.use_beta_sigmas: - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - sigmas = np.concatenate([sigmas, sigmas[-1:]]).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() - timesteps = (sigmas * self.config.num_train_timesteps).copy() - sigmas = np.concatenate([sigmas, sigmas[-1:]]).astype(np.float32) - else: - sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) - sigma_max = ( - (1 - self.alphas_cumprod[self.noisiest_timestep]) / self.alphas_cumprod[self.noisiest_timestep] - ) ** 0.5 - sigmas = np.concatenate([sigmas, [sigma_max]]).astype(np.float32) - - self.sigmas = torch.from_numpy(sigmas) - - # when num_inference_steps == num_train_timesteps, we can end up with - # duplicates in timesteps. - _, unique_indices = np.unique(timesteps, return_index=True) - timesteps = timesteps[np.sort(unique_indices)] - - self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) - - self.num_inference_steps = len(timesteps) - - self.model_outputs = [ - None, - ] * self.config.solver_order - self.lower_order_nums = 0 - - # add an index counter for schedulers that allow duplicated timesteps - self._step_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - """ - "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the - prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by - s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing - pixels from saturation at each step. We find that dynamic thresholding results in significantly better - photorealism as well as better image-text alignment, especially when using very large guidance weights." - - https://huggingface.co/papers/2205.11487 - """ - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half - - # Flatten sample for doing quantile calculation along each image - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - - abs_sample = sample.abs() # "a certain percentile absolute pixel value" - - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] - s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 - sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" - - sample = sample.reshape(batch_size, channels, *remaining_dims) - sample = sample.to(dtype) - - return sample - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma - log_sigma = np.log(np.maximum(sigma, 1e-10)) - - # get distribution - dists = log_sigma - log_sigmas[:, np.newaxis] - - # get sigmas range - low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) - high_idx = low_idx + 1 - - low = log_sigmas[low_idx] - high = log_sigmas[high_idx] - - # interpolate sigmas - w = (low - log_sigma) / (low - high) - w = np.clip(w, 0, 1) - - # transform interpolation to time range - t = (1 - w) * low_idx + w * high_idx - t = t.reshape(sigma.shape) - return t - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t - def _sigma_to_alpha_sigma_t(self, sigma): - if self.config.use_flow_sigmas: - alpha_t = 1 - sigma - sigma_t = sigma - else: - alpha_t = 1 / ((sigma**2 + 1) ** 0.5) - sigma_t = sigma * alpha_t - - 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: - """Constructs the noise schedule of Karras et al. (2022).""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - rho = 7.0 # 7.0 is the value used in the paper - ramp = np.linspace(0, 1, num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Constructs an exponential noise schedule.""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: - """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return sigmas - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.convert_model_output - def convert_model_output( - self, - model_output: torch.Tensor, - *args, - sample: torch.Tensor = None, - **kwargs, - ) -> torch.Tensor: - """ - Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is - designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an - integral of the data prediction model. - - - - The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise - prediction and data prediction models. - - - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. - 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`", - ) - - # DPM-Solver++ needs to solve an integral of the data prediction model. - if self.config.algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]: - if self.config.prediction_type == "epsilon": - # DPM-Solver and DPM-Solver++ only need the "mean" output. - if self.config.variance_type in ["learned", "learned_range"]: - model_output = model_output[:, :3] - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - 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": - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - x0_pred = alpha_t * sample - sigma_t * model_output - elif self.config.prediction_type == "flow_prediction": - sigma_t = self.sigmas[self.step_index] - 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 DPMSolverMultistepScheduler." - ) - - if self.config.thresholding: - x0_pred = self._threshold_sample(x0_pred) - - return x0_pred - - # DPM-Solver needs to solve an integral of the noise prediction model. - elif self.config.algorithm_type in ["dpmsolver", "sde-dpmsolver"]: - if self.config.prediction_type == "epsilon": - # DPM-Solver and DPM-Solver++ only need the "mean" output. - if self.config.variance_type in ["learned", "learned_range"]: - epsilon = model_output[:, :3] - else: - epsilon = model_output - elif self.config.prediction_type == "sample": - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - epsilon = (sample - alpha_t * model_output) / sigma_t - elif self.config.prediction_type == "v_prediction": - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - epsilon = alpha_t * model_output + sigma_t * sample - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" - " `v_prediction` for the DPMSolverMultistepScheduler." - ) - - if self.config.thresholding: - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - x0_pred = (sample - sigma_t * epsilon) / alpha_t - x0_pred = self._threshold_sample(x0_pred) - epsilon = (sample - alpha_t * x0_pred) / sigma_t - - return epsilon - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.dpm_solver_first_order_update - def dpm_solver_first_order_update( - self, - model_output: torch.Tensor, - *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, - **kwargs, - ) -> torch.Tensor: - """ - One step for the first-order DPMSolver (equivalent to DDIM). - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. - sample (`torch.Tensor`): - A current instance of a sample created by the diffusion process. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) - prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 2: - sample = args[2] - 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`", - ) - - 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`", - ) - - 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) - lambda_s = torch.log(alpha_s) - torch.log(sigma_s) - - h = lambda_t - lambda_s - if self.config.algorithm_type == "dpmsolver++": - x_t = (sigma_t / sigma_s) * sample - (alpha_t * (torch.exp(-h) - 1.0)) * model_output - elif self.config.algorithm_type == "dpmsolver": - x_t = (alpha_t / alpha_s) * sample - (sigma_t * (torch.exp(h) - 1.0)) * model_output - elif self.config.algorithm_type == "sde-dpmsolver++": - assert noise is not None - x_t = ( - (sigma_t / sigma_s * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - elif self.config.algorithm_type == "sde-dpmsolver": - assert noise is not None - x_t = ( - (alpha_t / alpha_s) * sample - - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * model_output - + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise - ) - return x_t - - # 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], - *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, - **kwargs, - ) -> torch.Tensor: - """ - One step for the second-order multistep DPMSolver. - - Args: - 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. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None) - prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 2: - sample = args[2] - else: - raise ValueError("missing `sample` as a required keyword argument") - if timestep_list is not None: - deprecate( - "timestep_list", - "1.0.0", - "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", - ) - - 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`", - ) - - sigma_t, sigma_s0, sigma_s1 = ( - self.sigmas[self.step_index + 1], - self.sigmas[self.step_index], - self.sigmas[self.step_index - 1], - ) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1) - - m0, m1 = model_output_list[-1], model_output_list[-2] - - h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1 - r0 = h_0 / h - D0, D1 = m0, (1.0 / r0) * (m0 - m1) - if self.config.algorithm_type == "dpmsolver++": - # See https://huggingface.co/papers/2211.01095 for detailed derivations - if self.config.solver_type == "midpoint": - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - - 0.5 * (alpha_t * (torch.exp(-h) - 1.0)) * D1 - ) - elif self.config.solver_type == "heun": - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1 - ) - elif self.config.algorithm_type == "dpmsolver": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - if self.config.solver_type == "midpoint": - x_t = ( - (alpha_t / alpha_s0) * sample - - (sigma_t * (torch.exp(h) - 1.0)) * D0 - - 0.5 * (sigma_t * (torch.exp(h) - 1.0)) * D1 - ) - elif self.config.solver_type == "heun": - x_t = ( - (alpha_t / alpha_s0) * sample - - (sigma_t * (torch.exp(h) - 1.0)) * D0 - - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1 - ) - elif self.config.algorithm_type == "sde-dpmsolver++": - assert noise is not None - if self.config.solver_type == "midpoint": - x_t = ( - (sigma_t / sigma_s0 * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 - + 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1 - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - elif self.config.solver_type == "heun": - x_t = ( - (sigma_t / sigma_s0 * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 - + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1 - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - elif self.config.algorithm_type == "sde-dpmsolver": - assert noise is not None - if self.config.solver_type == "midpoint": - x_t = ( - (alpha_t / alpha_s0) * sample - - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * D0 - - (sigma_t * (torch.exp(h) - 1.0)) * D1 - + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise - ) - elif self.config.solver_type == "heun": - x_t = ( - (alpha_t / alpha_s0) * sample - - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * D0 - - 2.0 * (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1 - + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise - ) - return x_t - - # 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], - *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, - **kwargs, - ) -> torch.Tensor: - """ - One step for the third-order multistep DPMSolver. - - Args: - 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. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - - timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None) - prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 2: - sample = args[2] - else: - raise ValueError("missing `sample` as a required keyword argument") - if timestep_list is not None: - deprecate( - "timestep_list", - "1.0.0", - "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", - ) - - 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`", - ) - - sigma_t, sigma_s0, sigma_s1, sigma_s2 = ( - self.sigmas[self.step_index + 1], - self.sigmas[self.step_index], - self.sigmas[self.step_index - 1], - self.sigmas[self.step_index - 2], - ) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1) - alpha_s2, sigma_s2 = self._sigma_to_alpha_sigma_t(sigma_s2) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1) - lambda_s2 = torch.log(alpha_s2) - torch.log(sigma_s2) - - m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3] - - h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2 - r0, r1 = h_0 / h, h_1 / h - D0 = m0 - D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2) - D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1) - D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1) - if self.config.algorithm_type == "dpmsolver++": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1 - - (alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2 - ) - elif self.config.algorithm_type == "dpmsolver": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - x_t = ( - (alpha_t / alpha_s0) * sample - - (sigma_t * (torch.exp(h) - 1.0)) * D0 - - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1 - - (sigma_t * ((torch.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2 - ) - elif self.config.algorithm_type == "sde-dpmsolver++": - assert noise is not None - x_t = ( - (sigma_t / sigma_s0 * torch.exp(-h)) * sample - + (alpha_t * (1.0 - torch.exp(-2.0 * h))) * D0 - + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1 - + (alpha_t * ((1.0 - torch.exp(-2.0 * h) - 2.0 * h) / (2.0 * h) ** 2 - 0.5)) * D2 - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - return x_t - - def _init_step_index(self, timestep): - if isinstance(timestep, torch.Tensor): - timestep = timestep.to(self.timesteps.device) - - index_candidates = (self.timesteps == timestep).nonzero() - - if len(index_candidates) == 0: - step_index = len(self.timesteps) - 1 - # 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) - elif len(index_candidates) > 1: - step_index = index_candidates[1].item() - else: - step_index = index_candidates[0].item() - - self._step_index = step_index - - def step( - self, - model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], - sample: torch.Tensor, - generator=None, - variance_noise: Optional[torch.Tensor] = None, - return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: - """ - Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with - the multistep DPMSolver. - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - generator (`torch.Generator`, *optional*): - A random number generator. - 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`): - 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( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - if self.step_index is None: - self._init_step_index(timestep) - - # Improve numerical stability for small number of steps - lower_order_final = (self.step_index == len(self.timesteps) - 1) and ( - self.config.euler_at_final or (self.config.lower_order_final and len(self.timesteps) < 15) - ) - lower_order_second = ( - (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15 - ) - - model_output = self.convert_model_output(model_output, sample=sample) - for i in range(self.config.solver_order - 1): - self.model_outputs[i] = self.model_outputs[i + 1] - self.model_outputs[-1] = model_output - - 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 - ) - elif self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]: - noise = variance_noise - else: - noise = None - - if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final: - prev_sample = self.dpm_solver_first_order_update(model_output, sample=sample, noise=noise) - elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second: - prev_sample = self.multistep_dpm_solver_second_order_update(self.model_outputs, sample=sample, noise=noise) - else: - prev_sample = self.multistep_dpm_solver_third_order_update(self.model_outputs, sample=sample) - - if self.lower_order_nums < self.config.solver_order: - self.lower_order_nums += 1 - - # upon completion increase step index by one - self._step_index += 1 - - if not return_dict: - return (prev_sample,) - - return SchedulerOutput(prev_sample=prev_sample) - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.scale_model_input - def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: - """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. - - Args: - sample (`torch.Tensor`): - The input sample. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - step_indices = [] - for timestep in timesteps: - index_candidates = (schedule_timesteps == timestep).nonzero() - if len(index_candidates) == 0: - step_index = len(schedule_timesteps) - 1 - elif len(index_candidates) > 1: - step_index = index_candidates[1].item() - else: - step_index = index_candidates[0].item() - step_indices.append(step_index) - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - noisy_samples = alpha_t * original_samples + sigma_t * noise - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_sde.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_sde.py deleted file mode 100644 index eeb06773d97756a6eafe4f0bcee8a669925b692e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_sde.py +++ /dev/null @@ -1,676 +0,0 @@ -# Copyright 2025 Katherine Crowson, The HuggingFace Team and hlky. 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 List, Optional, Tuple, Union - -import numpy as np -import torch -import torchsde - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput, is_scipy_available -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin - - -if is_scipy_available(): - import scipy.stats - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DPMSolverSDE -class DPMSolverSDESchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -class BatchedBrownianTree: - """A wrapper around torchsde.BrownianTree that enables batches of entropy.""" - - def __init__(self, x, t0, t1, seed=None, **kwargs): - t0, t1, self.sign = self.sort(t0, t1) - w0 = kwargs.get("w0", torch.zeros_like(x)) - if seed is None: - seed = torch.randint(0, 2**63 - 1, []).item() - self.batched = True - try: - assert len(seed) == x.shape[0] - w0 = w0[0] - except TypeError: - seed = [seed] - self.batched = False - self.trees = [ - torchsde.BrownianInterval( - t0=t0, - t1=t1, - size=w0.shape, - dtype=w0.dtype, - device=w0.device, - entropy=s, - tol=1e-6, - pool_size=24, - halfway_tree=True, - ) - for s in seed - ] - - @staticmethod - def sort(a, b): - return (a, b, 1) if a < b else (b, a, -1) - - def __call__(self, t0, t1): - 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] - - -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 - 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 - internal timestep. - """ - - def __init__(self, x, sigma_min, sigma_max, seed=None, transform=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): - t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next)) - return self.tree(t0, t1) / (t1 - t0).abs().sqrt() - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -class DPMSolverSDEScheduler(SchedulerMixin, ConfigMixin): - """ - DPMSolverSDEScheduler implements the stochastic sampler from the [Elucidating the Design Space of Diffusion-Based - Generative Models](https://huggingface.co/papers/2206.00364) paper. - - 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`, defaults to 1000): - The number of diffusion steps to train the model. - beta_start (`float`, defaults to 0.00085): - The starting `beta` value of inference. - beta_end (`float`, defaults to 0.012): - 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` or `scaled_linear`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - 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://imagen.research.google/video/paper.pdf) paper). - use_karras_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, - the sigmas are determined according to a sequence of noise levels {σi}. - use_exponential_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. - 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. - noise_sampler_seed (`int`, *optional*, defaults to `None`): - The random seed to use for the noise sampler. If `None`, a random seed is generated. - 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. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 2 - - @register_to_config - def __init__( - self, - 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", - 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: - raise ValueError( - "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." - ) - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - # set all values - self.set_timesteps(num_train_timesteps, None, num_train_timesteps) - self.use_karras_sigmas = use_karras_sigmas - self.noise_sampler = None - self.noise_sampler_seed = noise_sampler_seed - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep - 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() - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - 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 - - @property - def init_noise_sigma(self): - # standard deviation of the initial noise distribution - if self.config.timestep_spacing in ["linspace", "trailing"]: - return self.sigmas.max() - - return (self.sigmas.max() ** 2 + 1) ** 0.5 - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 scale_model_input( - self, - sample: torch.Tensor, - timestep: Union[float, torch.Tensor], - ) -> torch.Tensor: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - if self.step_index is None: - self._init_step_index(timestep) - - sigma = self.sigmas[self.step_index] - sigma_input = sigma if self.state_in_first_order else self.mid_point_sigma - sample = sample / ((sigma_input**2 + 1) ** 0.5) - return sample - - def set_timesteps( - self, - num_inference_steps: int, - device: Union[str, torch.device] = None, - num_train_timesteps: Optional[int] = 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 timesteps are not moved. - """ - self.num_inference_steps = num_inference_steps - - num_train_timesteps = num_train_timesteps or self.config.num_train_timesteps - - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = np.linspace(0, num_train_timesteps - 1, num_inference_steps, dtype=float)[::-1].copy() - elif self.config.timestep_spacing == "leading": - step_ratio = num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(float) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(num_train_timesteps, 0, -step_ratio)).round().copy().astype(float) - timesteps -= 1 - else: - raise ValueError( - 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) - log_sigmas = np.log(sigmas) - sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) - - if self.config.use_karras_sigmas: - sigmas = self._convert_to_karras(in_sigmas=sigmas) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - elif self.config.use_exponential_sigmas: - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - elif self.config.use_beta_sigmas: - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - - second_order_timesteps = self._second_order_timesteps(sigmas, log_sigmas) - - sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) - sigmas = torch.from_numpy(sigmas).to(device=device) - self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]]) - - timesteps = torch.from_numpy(timesteps) - second_order_timesteps = torch.from_numpy(second_order_timesteps) - timesteps = torch.cat([timesteps[:1], timesteps[1:].repeat_interleave(2)]) - timesteps[1::2] = second_order_timesteps - - if str(device).startswith("mps"): - # mps does not support float64 - self.timesteps = timesteps.to(device, dtype=torch.float32) - else: - self.timesteps = timesteps.to(device=device) - - # empty first order variables - self.sample = None - self.mid_point_sigma = None - - self._step_index = None - self._begin_index = None - 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 sigma_fn(_t): - return np.exp(-_t) - - def t_fn(_sigma): - return -np.log(_sigma) - - midpoint_ratio = 0.5 - t = t_fn(sigmas) - delta_time = np.diff(t) - t_proposed = t[:-1] + delta_time * midpoint_ratio - sig_proposed = sigma_fn(t_proposed) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sig_proposed]) - return timesteps - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma - log_sigma = np.log(np.maximum(sigma, 1e-10)) - - # get distribution - dists = log_sigma - log_sigmas[:, np.newaxis] - - # get sigmas range - low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) - high_idx = low_idx + 1 - - low = log_sigmas[low_idx] - high = log_sigmas[high_idx] - - # interpolate sigmas - w = (low - log_sigma) / (low - high) - w = np.clip(w, 0, 1) - - # transform interpolation to time range - t = (1 - w) * low_idx + w * high_idx - t = t.reshape(sigma.shape) - return t - - # copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras - def _convert_to_karras(self, in_sigmas: torch.Tensor) -> torch.Tensor: - """Constructs the noise schedule of Karras et al. (2022).""" - - sigma_min: float = in_sigmas[-1].item() - sigma_max: float = in_sigmas[0].item() - - rho = 7.0 # 7.0 is the value used in the paper - ramp = np.linspace(0, 1, self.num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Constructs an exponential noise schedule.""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: - """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return sigmas - - @property - def state_in_first_order(self): - 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], - return_dict: bool = True, - s_noise: float = 1.0, - ) -> Union[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`): - 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`): - 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 - tuple. - s_noise (`float`, *optional*, defaults to 1.0): - Scaling factor for noise added to the sample. - - Returns: - [`~schedulers.scheduling_dpmsolver_sde.DPMSolverSDESchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_dpmsolver_sde.DPMSolverSDESchedulerOutput`] is - returned, otherwise a tuple is returned where the first element is the sample tensor. - """ - if self.step_index is None: - self._init_step_index(timestep) - - # 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) - - # Define functions to compute sigma and t from each other - def sigma_fn(_t: torch.Tensor) -> torch.Tensor: - return _t.neg().exp() - - def t_fn(_sigma: torch.Tensor) -> torch.Tensor: - return _sigma.log().neg() - - if self.state_in_first_order: - sigma = self.sigmas[self.step_index] - sigma_next = self.sigmas[self.step_index + 1] - else: - # 2nd order - sigma = self.sigmas[self.step_index - 1] - sigma_next = self.sigmas[self.step_index] - - # Set the midpoint and step size for the current step - midpoint_ratio = 0.5 - t, t_next = t_fn(sigma), t_fn(sigma_next) - delta_time = t_next - t - t_proposed = t + delta_time * midpoint_ratio - - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - if self.config.prediction_type == "epsilon": - sigma_input = sigma if self.state_in_first_order else sigma_fn(t_proposed) - pred_original_sample = sample - sigma_input * model_output - elif self.config.prediction_type == "v_prediction": - sigma_input = sigma if self.state_in_first_order else sigma_fn(t_proposed) - pred_original_sample = model_output * (-sigma_input / (sigma_input**2 + 1) ** 0.5) + ( - sample / (sigma_input**2 + 1) - ) - elif self.config.prediction_type == "sample": - raise NotImplementedError("prediction_type not implemented yet: sample") - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`" - ) - - if sigma_next == 0: - derivative = (sample - pred_original_sample) / sigma - dt = sigma_next - sigma - prev_sample = sample + derivative * dt - else: - if self.state_in_first_order: - t_next = t_proposed - else: - sample = self.sample - - 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_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 - ( - t - ancestral_t - ).expm1() * pred_original_sample - prev_sample = prev_sample + self.noise_sampler(sigma_fn(t), sigma_fn(t_next)) * s_noise * sigma_up - - if self.state_in_first_order: - # store for 2nd order step - self.sample = sample - self.mid_point_sigma = sigma_fn(t_next) - else: - # free for "first order mode" - self.sample = None - self.mid_point_sigma = None - - # upon completion increase step index by one - self._step_index += 1 - - if not return_dict: - return ( - prev_sample, - pred_original_sample, - ) - - return DPMSolverSDESchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.Tensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - noisy_samples = original_samples + noise * sigma - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_singlestep.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_singlestep.py deleted file mode 100644 index 9e3e830039bb141bc3917d64f3bafa539403f1af..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_dpmsolver_singlestep.py +++ /dev/null @@ -1,1168 +0,0 @@ -# Copyright 2025 TSAIL 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. - -# DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver - -import math -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import deprecate, is_scipy_available, logging -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput - - -if is_scipy_available(): - import scipy.stats - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -class DPMSolverSinglestepScheduler(SchedulerMixin, ConfigMixin): - """ - `DPMSolverSinglestepScheduler` is a fast dedicated high-order solver for diffusion ODEs. - - 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`, 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`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, *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 - 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://imagen.research.google/video/paper.pdf) paper). - thresholding (`bool`, defaults to `False`): - Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such - as Stable Diffusion. - dynamic_thresholding_ratio (`float`, defaults to 0.995): - The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. - sample_max_value (`float`, defaults to 1.0): - The threshold value for dynamic thresholding. Valid only when `thresholding=True` and - `algorithm_type="dpmsolver++"`. - algorithm_type (`str`, defaults to `dpmsolver++`): - Algorithm type for the solver; can be `dpmsolver` or `dpmsolver++` or `sde-dpmsolver++`. The `dpmsolver` - type implements the algorithms in the [DPMSolver](https://huggingface.co/papers/2206.00927) paper, and the - `dpmsolver++` type implements the algorithms in the [DPMSolver++](https://huggingface.co/papers/2211.01095) - paper. It is recommended to use `dpmsolver++` or `sde-dpmsolver++` with `solver_order=2` for guided - sampling like in Stable Diffusion. - 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. - 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. - use_karras_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, - the sigmas are determined according to a sequence of noise levels {σi}. - use_exponential_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. - 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. - final_sigmas_type (`str`, *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. - 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. - variance_type (`str`, *optional*): - Set to "learned" or "learned_range" for diffusion models that predict variance. If set, the model's output - contains the predicted Gaussian variance. - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: Optional[np.ndarray] = None, - solver_order: int = 2, - prediction_type: str = "epsilon", - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - sample_max_value: float = 1.0, - algorithm_type: str = "dpmsolver++", - solver_type: str = "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[str] = "zero", # "zero", "sigma_min" - lambda_min_clipped: float = -float("inf"), - variance_type: Optional[str] = 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.") - 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 == "dpmsolver": - deprecation_message = "algorithm_type `dpmsolver` is deprecated and will be removed in a future version. Choose from `dpmsolver++` or `sde-dpmsolver++` instead" - deprecate("algorithm_types=dpmsolver", "1.0.0", deprecation_message) - - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - # Currently we only support VP-type noise schedule - self.alpha_t = torch.sqrt(self.alphas_cumprod) - self.sigma_t = torch.sqrt(1 - self.alphas_cumprod) - self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t) - self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5 - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # settings for DPM-Solver - if algorithm_type not in ["dpmsolver", "dpmsolver++", "sde-dpmsolver++"]: - if algorithm_type == "deis": - self.register_to_config(algorithm_type="dpmsolver++") - else: - raise NotImplementedError(f"{algorithm_type} is not implemented for {self.__class__}") - if solver_type not in ["midpoint", "heun"]: - if solver_type in ["logrho", "bh1", "bh2"]: - self.register_to_config(solver_type="midpoint") - else: - raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}") - - if algorithm_type not in ["dpmsolver++", "sde-dpmsolver++"] and final_sigmas_type == "zero": - raise ValueError( - f"`final_sigmas_type` {final_sigmas_type} is not supported for `algorithm_type` {algorithm_type}. Please choose `sigma_min` instead." - ) - - # setable values - self.num_inference_steps = None - timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy() - self.timesteps = torch.from_numpy(timesteps) - self.model_outputs = [None] * solver_order - self.sample = None - self.order_list = self.get_order_list(num_train_timesteps) - self._step_index = None - 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]: - """ - Computes the solver order at each time step. - - Args: - num_inference_steps (`int`): - The number of diffusion steps used when generating samples with a pre-trained model. - """ - steps = num_inference_steps - order = self.config.solver_order - if order > 3: - raise ValueError("Order > 3 is not supported by this scheduler") - if self.config.lower_order_final: - if order == 3: - if steps % 3 == 0: - orders = [1, 2, 3] * (steps // 3 - 1) + [1, 2] + [1] - elif steps % 3 == 1: - orders = [1, 2, 3] * (steps // 3) + [1] - else: - orders = [1, 2, 3] * (steps // 3) + [1, 2] - elif order == 2: - if steps % 2 == 0: - orders = [1, 2] * (steps // 2 - 1) + [1, 1] - else: - orders = [1, 2] * (steps // 2) + [1] - elif order == 1: - orders = [1] * steps - else: - if order == 3: - orders = [1, 2, 3] * (steps // 3) - elif order == 2: - orders = [1, 2] * (steps // 2) - elif order == 1: - orders = [1] * steps - - if self.config.final_sigmas_type == "zero": - orders[-1] = 1 - - return orders - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 set_timesteps( - self, - num_inference_steps: int = None, - device: Union[str, torch.device] = None, - timesteps: Optional[List[int]] = 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 timesteps are not moved. - 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`. - """ - if num_inference_steps is None and timesteps is None: - raise ValueError("Must pass exactly one of `num_inference_steps` or `timesteps`.") - if num_inference_steps is not None and timesteps is not None: - raise ValueError("Must pass exactly one of `num_inference_steps` or `timesteps`.") - if timesteps is not None and self.config.use_karras_sigmas: - raise ValueError("Cannot use `timesteps` when `config.use_karras_sigmas=True`.") - if timesteps is not None and self.config.use_exponential_sigmas: - raise ValueError("Cannot set `timesteps` with `config.use_exponential_sigmas = True`.") - if timesteps is not None and self.config.use_beta_sigmas: - raise ValueError("Cannot set `timesteps` with `config.use_beta_sigmas = True`.") - - num_inference_steps = num_inference_steps or len(timesteps) - self.num_inference_steps = num_inference_steps - - if timesteps is not None: - timesteps = np.array(timesteps).astype(np.int64) - else: - # Clipping the minimum of all lambda(t) for numerical stability. - # This is critical for cosine (squaredcos_cap_v2) noise schedule. - clipped_idx = torch.searchsorted(torch.flip(self.lambda_t, [0]), self.config.lambda_min_clipped) - clipped_idx = clipped_idx.item() - timesteps = ( - np.linspace(0, self.config.num_train_timesteps - 1 - clipped_idx, num_inference_steps + 1) - .round()[::-1][:-1] - .copy() - .astype(np.int64) - ) - - sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) - log_sigmas = np.log(sigmas) - if self.config.use_karras_sigmas: - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round() - elif self.config.use_exponential_sigmas: - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - elif self.config.use_beta_sigmas: - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - 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() - timesteps = (sigmas * self.config.num_train_timesteps).copy() - else: - 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 - elif self.config.final_sigmas_type == "zero": - sigma_last = 0 - else: - raise ValueError( - f" `final_sigmas_type` must be one of `sigma_min` or `zero`, but got {self.config.final_sigmas_type}" - ) - sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) - - self.sigmas = torch.from_numpy(sigmas).to(device=device) - - self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) - self.model_outputs = [None] * self.config.solver_order - self.sample = None - - if not self.config.lower_order_final and num_inference_steps % self.config.solver_order != 0: - logger.warning( - "Changing scheduler {self.config} to have `lower_order_final` set to True to handle uneven amount of inference steps. Please make sure to always use an even number of `num_inference steps when using `lower_order_final=False`." - ) - self.register_to_config(lower_order_final=True) - - if not self.config.lower_order_final and self.config.final_sigmas_type == "zero": - logger.warning( - " `last_sigmas_type='zero'` is not supported for `lower_order_final=False`. Changing scheduler {self.config} to have `lower_order_final` set to True." - ) - self.register_to_config(lower_order_final=True) - - self.order_list = self.get_order_list(num_inference_steps) - - # add an index counter for schedulers that allow duplicated timesteps - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - """ - "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the - prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by - s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing - pixels from saturation at each step. We find that dynamic thresholding results in significantly better - photorealism as well as better image-text alignment, especially when using very large guidance weights." - - https://huggingface.co/papers/2205.11487 - """ - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half - - # Flatten sample for doing quantile calculation along each image - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - - abs_sample = sample.abs() # "a certain percentile absolute pixel value" - - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] - s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 - sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" - - sample = sample.reshape(batch_size, channels, *remaining_dims) - sample = sample.to(dtype) - - return sample - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma - log_sigma = np.log(np.maximum(sigma, 1e-10)) - - # get distribution - dists = log_sigma - log_sigmas[:, np.newaxis] - - # get sigmas range - low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) - high_idx = low_idx + 1 - - low = log_sigmas[low_idx] - high = log_sigmas[high_idx] - - # interpolate sigmas - w = (low - log_sigma) / (low - high) - w = np.clip(w, 0, 1) - - # transform interpolation to time range - t = (1 - w) * low_idx + w * high_idx - t = t.reshape(sigma.shape) - return t - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t - def _sigma_to_alpha_sigma_t(self, sigma): - if self.config.use_flow_sigmas: - alpha_t = 1 - sigma - sigma_t = sigma - else: - alpha_t = 1 / ((sigma**2 + 1) ** 0.5) - sigma_t = sigma * alpha_t - - 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: - """Constructs the noise schedule of Karras et al. (2022).""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - rho = 7.0 # 7.0 is the value used in the paper - ramp = np.linspace(0, 1, num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Constructs an exponential noise schedule.""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: - """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return sigmas - - def convert_model_output( - self, - model_output: torch.Tensor, - *args, - sample: torch.Tensor = None, - **kwargs, - ) -> torch.Tensor: - """ - Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is - designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an - integral of the data prediction model. - - - - The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise - prediction and data prediction models. - - - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. - 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`", - ) - # DPM-Solver++ needs to solve an integral of the data prediction model. - if self.config.algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]: - if self.config.prediction_type == "epsilon": - # DPM-Solver and DPM-Solver++ only need the "mean" output. - if self.config.variance_type in ["learned", "learned_range"]: - model_output = model_output[:, :3] - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - 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": - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - x0_pred = alpha_t * sample - sigma_t * model_output - elif self.config.prediction_type == "flow_prediction": - sigma_t = self.sigmas[self.step_index] - 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 DPMSolverSinglestepScheduler." - ) - - if self.config.thresholding: - x0_pred = self._threshold_sample(x0_pred) - - return x0_pred - - # DPM-Solver needs to solve an integral of the noise prediction model. - elif self.config.algorithm_type == "dpmsolver": - if self.config.prediction_type == "epsilon": - # DPM-Solver and DPM-Solver++ only need the "mean" output. - if self.config.variance_type in ["learned", "learned_range"]: - epsilon = model_output[:, :3] - else: - epsilon = model_output - elif self.config.prediction_type == "sample": - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - epsilon = (sample - alpha_t * model_output) / sigma_t - elif self.config.prediction_type == "v_prediction": - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - epsilon = alpha_t * model_output + sigma_t * sample - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" - " `v_prediction` for the DPMSolverSinglestepScheduler." - ) - - if self.config.thresholding: - alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep] - x0_pred = (sample - sigma_t * epsilon) / alpha_t - x0_pred = self._threshold_sample(x0_pred) - epsilon = (sample - alpha_t * x0_pred) / sigma_t - - return epsilon - - def dpm_solver_first_order_update( - self, - model_output: torch.Tensor, - *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, - **kwargs, - ) -> torch.Tensor: - """ - One step for the first-order DPMSolver (equivalent to DDIM). - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. - timestep (`int`): - The current discrete timestep in the diffusion chain. - 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. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) - prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 2: - sample = args[2] - 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`", - ) - - 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`", - ) - 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) - lambda_s = torch.log(alpha_s) - torch.log(sigma_s) - h = lambda_t - lambda_s - if self.config.algorithm_type == "dpmsolver++": - x_t = (sigma_t / sigma_s) * sample - (alpha_t * (torch.exp(-h) - 1.0)) * model_output - elif self.config.algorithm_type == "dpmsolver": - x_t = (alpha_t / alpha_s) * sample - (sigma_t * (torch.exp(h) - 1.0)) * model_output - elif self.config.algorithm_type == "sde-dpmsolver++": - assert noise is not None - x_t = ( - (sigma_t / sigma_s * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - return x_t - - def singlestep_dpm_solver_second_order_update( - self, - model_output_list: List[torch.Tensor], - *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, - **kwargs, - ) -> torch.Tensor: - """ - One step for the second-order singlestep DPMSolver that computes the solution at time `prev_timestep` from the - time `timestep_list[-2]`. - - Args: - 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. - 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. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None) - prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 2: - sample = args[2] - else: - raise ValueError("missing `sample` as a required keyword argument") - if timestep_list is not None: - deprecate( - "timestep_list", - "1.0.0", - "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", - ) - - 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`", - ) - sigma_t, sigma_s0, sigma_s1 = ( - self.sigmas[self.step_index + 1], - self.sigmas[self.step_index], - self.sigmas[self.step_index - 1], - ) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1) - - m0, m1 = model_output_list[-1], model_output_list[-2] - - h, h_0 = lambda_t - lambda_s1, lambda_s0 - lambda_s1 - r0 = h_0 / h - D0, D1 = m1, (1.0 / r0) * (m0 - m1) - if self.config.algorithm_type == "dpmsolver++": - # See https://huggingface.co/papers/2211.01095 for detailed derivations - if self.config.solver_type == "midpoint": - x_t = ( - (sigma_t / sigma_s1) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - - 0.5 * (alpha_t * (torch.exp(-h) - 1.0)) * D1 - ) - elif self.config.solver_type == "heun": - x_t = ( - (sigma_t / sigma_s1) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1 - ) - elif self.config.algorithm_type == "dpmsolver": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - if self.config.solver_type == "midpoint": - x_t = ( - (alpha_t / alpha_s1) * sample - - (sigma_t * (torch.exp(h) - 1.0)) * D0 - - 0.5 * (sigma_t * (torch.exp(h) - 1.0)) * D1 - ) - elif self.config.solver_type == "heun": - x_t = ( - (alpha_t / alpha_s1) * sample - - (sigma_t * (torch.exp(h) - 1.0)) * D0 - - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1 - ) - elif self.config.algorithm_type == "sde-dpmsolver++": - assert noise is not None - if self.config.solver_type == "midpoint": - x_t = ( - (sigma_t / sigma_s1 * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 - + 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1 - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - elif self.config.solver_type == "heun": - x_t = ( - (sigma_t / sigma_s1 * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 - + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1 - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - return x_t - - def singlestep_dpm_solver_third_order_update( - self, - model_output_list: List[torch.Tensor], - *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, - **kwargs, - ) -> torch.Tensor: - """ - One step for the third-order singlestep DPMSolver that computes the solution at time `prev_timestep` from the - time `timestep_list[-3]`. - - Args: - 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. - prev_timestep (`int`): - The previous discrete timestep in the diffusion chain. - sample (`torch.Tensor`): - A current instance of a sample created by diffusion process. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - - timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None) - prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 2: - sample = args[2] - else: - raise ValueError("missing `sample` as a required keyword argument") - if timestep_list is not None: - deprecate( - "timestep_list", - "1.0.0", - "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", - ) - - 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`", - ) - - sigma_t, sigma_s0, sigma_s1, sigma_s2 = ( - self.sigmas[self.step_index + 1], - self.sigmas[self.step_index], - self.sigmas[self.step_index - 1], - self.sigmas[self.step_index - 2], - ) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1) - alpha_s2, sigma_s2 = self._sigma_to_alpha_sigma_t(sigma_s2) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1) - lambda_s2 = torch.log(alpha_s2) - torch.log(sigma_s2) - - m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3] - - h, h_0, h_1 = lambda_t - lambda_s2, lambda_s0 - lambda_s2, lambda_s1 - lambda_s2 - r0, r1 = h_0 / h, h_1 / h - D0 = m2 - D1_0, D1_1 = (1.0 / r1) * (m1 - m2), (1.0 / r0) * (m0 - m2) - D1 = (r0 * D1_0 - r1 * D1_1) / (r0 - r1) - D2 = 2.0 * (D1_1 - D1_0) / (r0 - r1) - if self.config.algorithm_type == "dpmsolver++": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - if self.config.solver_type == "midpoint": - x_t = ( - (sigma_t / sigma_s2) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1_1 - ) - elif self.config.solver_type == "heun": - x_t = ( - (sigma_t / sigma_s2) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1 - - (alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2 - ) - elif self.config.algorithm_type == "dpmsolver": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - if self.config.solver_type == "midpoint": - x_t = ( - (alpha_t / alpha_s2) * sample - - (sigma_t * (torch.exp(h) - 1.0)) * D0 - - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1_1 - ) - elif self.config.solver_type == "heun": - x_t = ( - (alpha_t / alpha_s2) * sample - - (sigma_t * (torch.exp(h) - 1.0)) * D0 - - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1 - - (sigma_t * ((torch.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2 - ) - elif self.config.algorithm_type == "sde-dpmsolver++": - assert noise is not None - if self.config.solver_type == "midpoint": - x_t = ( - (sigma_t / sigma_s2 * torch.exp(-h)) * sample - + (alpha_t * (1.0 - torch.exp(-2.0 * h))) * D0 - + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1_1 - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - elif self.config.solver_type == "heun": - x_t = ( - (sigma_t / sigma_s2 * torch.exp(-h)) * sample - + (alpha_t * (1.0 - torch.exp(-2.0 * h))) * D0 - + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1 - + (alpha_t * ((1.0 - torch.exp(-2.0 * h) + (-2.0 * h)) / (-2.0 * h) ** 2 - 0.5)) * D2 - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - return x_t - - def singlestep_dpm_solver_update( - self, - model_output_list: List[torch.Tensor], - *args, - sample: torch.Tensor = None, - order: int = None, - noise: Optional[torch.Tensor] = None, - **kwargs, - ) -> torch.Tensor: - """ - One step for the singlestep DPMSolver. - - Args: - 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. - prev_timestep (`int`): - The previous discrete timestep in the diffusion chain. - sample (`torch.Tensor`): - A current instance of a sample created by diffusion process. - order (`int`): - The solver order at this step. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None) - prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 2: - sample = args[2] - else: - raise ValueError("missing `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 timestep_list is not None: - deprecate( - "timestep_list", - "1.0.0", - "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", - ) - - 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`", - ) - - if order == 1: - return self.dpm_solver_first_order_update(model_output_list[-1], sample=sample, noise=noise) - elif order == 2: - return self.singlestep_dpm_solver_second_order_update(model_output_list, sample=sample, noise=noise) - elif order == 3: - return self.singlestep_dpm_solver_third_order_update(model_output_list, sample=sample, noise=noise) - else: - raise ValueError(f"Order must be 1, 2, 3, got {order}") - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep - def index_for_timestep(self, timestep, schedule_timesteps=None): - if schedule_timesteps is None: - schedule_timesteps = self.timesteps - - index_candidates = (schedule_timesteps == timestep).nonzero() - - if len(index_candidates) == 0: - step_index = len(self.timesteps) - 1 - # 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) - elif len(index_candidates) > 1: - step_index = index_candidates[1].item() - else: - step_index = index_candidates[0].item() - - return step_index - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index - def _init_step_index(self, timestep): - """ - Initialize the step_index counter for the scheduler. - """ - - 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( - self, - model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], - sample: torch.Tensor, - generator=None, - return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: - """ - Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with - the singlestep DPMSolver. - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - return_dict (`bool`): - 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( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - if self.step_index is None: - self._init_step_index(timestep) - - model_output = self.convert_model_output(model_output, sample=sample) - for i in range(self.config.solver_order - 1): - self.model_outputs[i] = self.model_outputs[i + 1] - self.model_outputs[-1] = model_output - - if self.config.algorithm_type == "sde-dpmsolver++": - noise = randn_tensor( - model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype - ) - else: - noise = None - - order = self.order_list[self.step_index] - - # For img2img denoising might start with order>1 which is not possible - # In this case make sure that the first two steps are both order=1 - while self.model_outputs[-order] is None: - order -= 1 - - # For single-step solvers, we use the initial value at each time with order = 1. - if order == 1: - self.sample = sample - - prev_sample = self.singlestep_dpm_solver_update( - self.model_outputs, sample=self.sample, order=order, noise=noise - ) - - # upon completion increase step index by one, noise=noise - self._step_index += 1 - - if not return_dict: - return (prev_sample,) - - return SchedulerOutput(prev_sample=prev_sample) - - def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: - """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. - - Args: - sample (`torch.Tensor`): - The input sample. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - noisy_samples = alpha_t * original_samples + sigma_t * noise - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py deleted file mode 100644 index f1b38aaff56c5c212a0312196d411e4e550d9da4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py +++ /dev/null @@ -1,712 +0,0 @@ -# Copyright 2025 TSAIL 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. - -# 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 - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import SchedulerMixin, SchedulerOutput - - -class EDMDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): - """ - Implements DPMSolverMultistepScheduler in EDM formulation as presented in Karras et al. 2022 [1]. - `EDMDPMSolverMultistepScheduler` is a fast dedicated high-order solver for diffusion ODEs. - - [1] Karras, Tero, et al. "Elucidating the Design Space of Diffusion-Based Generative Models." - https://huggingface.co/papers/2206.00364 - - 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: - 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): - 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): - 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): - 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://imagen.research.google/video/paper.pdf) paper). - thresholding (`bool`, defaults to `False`): - Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such - as Stable Diffusion. - dynamic_thresholding_ratio (`float`, defaults to 0.995): - The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. - sample_max_value (`float`, defaults to 1.0): - The threshold value for dynamic thresholding. Valid only when `thresholding=True` and - `algorithm_type="dpmsolver++"`. - algorithm_type (`str`, defaults to `dpmsolver++`): - Algorithm type for the solver; can be `dpmsolver++` or `sde-dpmsolver++`. The `dpmsolver++` type implements - the algorithms in the [DPMSolver++](https://huggingface.co/papers/2211.01095) paper. It is recommended to - use `dpmsolver++` or `sde-dpmsolver++` with `solver_order=2` for guided sampling like in Stable Diffusion. - 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. - 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. - euler_at_final (`bool`, defaults to `False`): - Whether to use Euler's method in the final step. It is a trade-off between numerical stability and detail - 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. - """ - - _compatibles = [] - order = 1 - - @register_to_config - def __init__( - self, - sigma_min: float = 0.002, - sigma_max: float = 80.0, - sigma_data: float = 0.5, - sigma_schedule: str = "karras", - num_train_timesteps: int = 1000, - prediction_type: str = "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", - lower_order_final: bool = True, - euler_at_final: bool = False, - final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min" - ): - # settings for DPM-Solver - if algorithm_type not in ["dpmsolver++", "sde-dpmsolver++"]: - if algorithm_type == "deis": - self.register_to_config(algorithm_type="dpmsolver++") - else: - raise NotImplementedError(f"{algorithm_type} is not implemented for {self.__class__}") - - if solver_type not in ["midpoint", "heun"]: - if solver_type in ["logrho", "bh1", "bh2"]: - self.register_to_config(solver_type="midpoint") - else: - raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}") - - if algorithm_type not in ["dpmsolver++", "sde-dpmsolver++"] and final_sigmas_type == "zero": - raise ValueError( - f"`final_sigmas_type` {final_sigmas_type} is not supported for `algorithm_type` {algorithm_type}. Please choose `sigma_min` instead." - ) - - ramp = torch.linspace(0, 1, num_train_timesteps) - if sigma_schedule == "karras": - sigmas = self._compute_karras_sigmas(ramp) - elif sigma_schedule == "exponential": - sigmas = self._compute_exponential_sigmas(ramp) - - self.timesteps = self.precondition_noise(sigmas) - - self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) - - # setable values - self.num_inference_steps = None - self.model_outputs = [None] * solver_order - self.lower_order_nums = 0 - self._step_index = None - self._begin_index = None - 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 - return (self.config.sigma_max**2 + 1) ** 0.5 - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 - - # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_inputs - def precondition_inputs(self, sample, sigma): - 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): - if not isinstance(sigma, torch.Tensor): - sigma = torch.tensor([sigma]) - - c_noise = 0.25 * torch.log(sigma) - - return c_noise - - # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_outputs - def precondition_outputs(self, sample, model_output, sigma): - sigma_data = self.config.sigma_data - c_skip = sigma_data**2 / (sigma**2 + sigma_data**2) - - if self.config.prediction_type == "epsilon": - c_out = sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5 - elif self.config.prediction_type == "v_prediction": - c_out = -sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5 - else: - raise ValueError(f"Prediction type {self.config.prediction_type} is not supported.") - - denoised = c_skip * sample + c_out * model_output - - 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: - """ - 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. - - Args: - sample (`torch.Tensor`): - The input sample. - timestep (`int`, *optional*): - The current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - if self.step_index is None: - self._init_step_index(timestep) - - sigma = self.sigmas[self.step_index] - sample = self.precondition_inputs(sample, sigma) - - self.is_scale_input_called = True - return sample - - def set_timesteps(self, num_inference_steps: int = None, device: Union[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 (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - """ - - self.num_inference_steps = num_inference_steps - - ramp = torch.linspace(0, 1, self.num_inference_steps) - if self.config.sigma_schedule == "karras": - sigmas = self._compute_karras_sigmas(ramp) - elif self.config.sigma_schedule == "exponential": - sigmas = self._compute_exponential_sigmas(ramp) - - sigmas = sigmas.to(dtype=torch.float32, device=device) - self.timesteps = self.precondition_noise(sigmas) - - if self.config.final_sigmas_type == "sigma_min": - sigma_last = self.config.sigma_min - elif self.config.final_sigmas_type == "zero": - sigma_last = 0 - else: - raise ValueError( - f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" - ) - - self.sigmas = torch.cat([sigmas, torch.tensor([sigma_last], dtype=torch.float32, device=device)]) - - self.model_outputs = [ - None, - ] * self.config.solver_order - self.lower_order_nums = 0 - - # add an index counter for schedulers that allow duplicated timesteps - self._step_index = None - self._begin_index = None - 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).""" - sigma_min = sigma_min or self.config.sigma_min - sigma_max = sigma_max or self.config.sigma_max - - rho = self.config.rho - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - 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. - - https://github.com/crowsonkb/k-diffusion/blob/6ab5146d4a5ef63901326489f31f1d8e7dd36b48/k_diffusion/sampling.py#L26 - """ - sigma_min = sigma_min or self.config.sigma_min - sigma_max = sigma_max or self.config.sigma_max - sigmas = torch.linspace(math.log(sigma_min), math.log(sigma_max), len(ramp)).exp().flip(0) - return sigmas - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - """ - "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the - prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by - s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing - pixels from saturation at each step. We find that dynamic thresholding results in significantly better - photorealism as well as better image-text alignment, especially when using very large guidance weights." - - https://huggingface.co/papers/2205.11487 - """ - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half - - # Flatten sample for doing quantile calculation along each image - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - - abs_sample = sample.abs() # "a certain percentile absolute pixel value" - - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] - s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 - sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" - - sample = sample.reshape(batch_size, channels, *remaining_dims) - sample = sample.to(dtype) - - return sample - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma - log_sigma = np.log(np.maximum(sigma, 1e-10)) - - # get distribution - dists = log_sigma - log_sigmas[:, np.newaxis] - - # get sigmas range - low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) - high_idx = low_idx + 1 - - low = log_sigmas[low_idx] - high = log_sigmas[high_idx] - - # interpolate sigmas - w = (low - log_sigma) / (low - high) - w = np.clip(w, 0, 1) - - # transform interpolation to time range - t = (1 - w) * low_idx + w * high_idx - t = t.reshape(sigma.shape) - return t - - 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, - ) -> torch.Tensor: - """ - Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is - designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an - integral of the data prediction model. - - - - The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise - prediction and data prediction models. - - - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. - sample (`torch.Tensor`): - A current instance of a sample created by the diffusion process. - - Returns: - `torch.Tensor`: - The converted model output. - """ - sigma = self.sigmas[self.step_index] - x0_pred = self.precondition_outputs(sample, model_output, sigma) - - if self.config.thresholding: - x0_pred = self._threshold_sample(x0_pred) - - return x0_pred - - def dpm_solver_first_order_update( - self, - model_output: torch.Tensor, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """ - One step for the first-order DPMSolver (equivalent to DDIM). - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. - sample (`torch.Tensor`): - A current instance of a sample created by the diffusion process. - - 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] - 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) - lambda_s = torch.log(alpha_s) - torch.log(sigma_s) - - h = lambda_t - lambda_s - if self.config.algorithm_type == "dpmsolver++": - x_t = (sigma_t / sigma_s) * sample - (alpha_t * (torch.exp(-h) - 1.0)) * model_output - elif self.config.algorithm_type == "sde-dpmsolver++": - assert noise is not None - x_t = ( - (sigma_t / sigma_s * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - - return x_t - - def multistep_dpm_solver_second_order_update( - self, - model_output_list: List[torch.Tensor], - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """ - One step for the second-order multistep DPMSolver. - - Args: - 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. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - sigma_t, sigma_s0, sigma_s1 = ( - self.sigmas[self.step_index + 1], - self.sigmas[self.step_index], - self.sigmas[self.step_index - 1], - ) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1) - - m0, m1 = model_output_list[-1], model_output_list[-2] - - h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1 - r0 = h_0 / h - D0, D1 = m0, (1.0 / r0) * (m0 - m1) - if self.config.algorithm_type == "dpmsolver++": - # See https://huggingface.co/papers/2211.01095 for detailed derivations - if self.config.solver_type == "midpoint": - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - - 0.5 * (alpha_t * (torch.exp(-h) - 1.0)) * D1 - ) - elif self.config.solver_type == "heun": - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1 - ) - elif self.config.algorithm_type == "sde-dpmsolver++": - assert noise is not None - if self.config.solver_type == "midpoint": - x_t = ( - (sigma_t / sigma_s0 * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 - + 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1 - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - elif self.config.solver_type == "heun": - x_t = ( - (sigma_t / sigma_s0 * torch.exp(-h)) * sample - + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 - + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1 - + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise - ) - - return x_t - - def multistep_dpm_solver_third_order_update( - self, - model_output_list: List[torch.Tensor], - sample: torch.Tensor = None, - ) -> torch.Tensor: - """ - One step for the third-order multistep DPMSolver. - - Args: - 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. - - Returns: - `torch.Tensor`: - The sample tensor at the previous timestep. - """ - sigma_t, sigma_s0, sigma_s1, sigma_s2 = ( - self.sigmas[self.step_index + 1], - self.sigmas[self.step_index], - self.sigmas[self.step_index - 1], - self.sigmas[self.step_index - 2], - ) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1) - alpha_s2, sigma_s2 = self._sigma_to_alpha_sigma_t(sigma_s2) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1) - lambda_s2 = torch.log(alpha_s2) - torch.log(sigma_s2) - - m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3] - - h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2 - r0, r1 = h_0 / h, h_1 / h - D0 = m0 - D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2) - D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1) - D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1) - if self.config.algorithm_type == "dpmsolver++": - # See https://huggingface.co/papers/2206.00927 for detailed derivations - x_t = ( - (sigma_t / sigma_s0) * sample - - (alpha_t * (torch.exp(-h) - 1.0)) * D0 - + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1 - - (alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2 - ) - - return x_t - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep - def index_for_timestep(self, timestep, schedule_timesteps=None): - if schedule_timesteps is None: - schedule_timesteps = self.timesteps - - index_candidates = (schedule_timesteps == timestep).nonzero() - - if len(index_candidates) == 0: - step_index = len(self.timesteps) - 1 - # 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) - elif len(index_candidates) > 1: - step_index = index_candidates[1].item() - else: - step_index = index_candidates[0].item() - - return step_index - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index - def _init_step_index(self, timestep): - """ - Initialize the step_index counter for the scheduler. - """ - - 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( - self, - model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], - sample: torch.Tensor, - generator=None, - return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: - """ - Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with - the multistep DPMSolver. - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - generator (`torch.Generator`, *optional*): - A random number generator. - return_dict (`bool`): - 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( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - if self.step_index is None: - self._init_step_index(timestep) - - # Improve numerical stability for small number of steps - lower_order_final = (self.step_index == len(self.timesteps) - 1) and ( - self.config.euler_at_final - or (self.config.lower_order_final and len(self.timesteps) < 15) - or self.config.final_sigmas_type == "zero" - ) - lower_order_second = ( - (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15 - ) - - model_output = self.convert_model_output(model_output, sample=sample) - for i in range(self.config.solver_order - 1): - self.model_outputs[i] = self.model_outputs[i + 1] - self.model_outputs[-1] = model_output - - if self.config.algorithm_type == "sde-dpmsolver++": - noise = randn_tensor( - model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype - ) - else: - noise = None - - if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final: - prev_sample = self.dpm_solver_first_order_update(model_output, sample=sample, noise=noise) - elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second: - prev_sample = self.multistep_dpm_solver_second_order_update(self.model_outputs, sample=sample, noise=noise) - else: - prev_sample = self.multistep_dpm_solver_third_order_update(self.model_outputs, sample=sample) - - if self.lower_order_nums < self.config.solver_order: - self.lower_order_nums += 1 - - # upon completion increase step index by one - self._step_index += 1 - - if not return_dict: - return (prev_sample,) - - return SchedulerOutput(prev_sample=prev_sample) - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.Tensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - noisy_samples = original_samples + noise * sigma - return noisy_samples - - # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._get_conditioning_c_in - def _get_conditioning_c_in(self, sigma): - c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5) - return c_in - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_edm_euler.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_edm_euler.py deleted file mode 100644 index dbeff3de56527c3591a44aaa523d449d7b60c374..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_edm_euler.py +++ /dev/null @@ -1,448 +0,0 @@ -# Copyright 2025 Katherine Crowson 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 List, Optional, Tuple, Union - -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 -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->EulerDiscrete -class EDMEulerSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -class EDMEulerScheduler(SchedulerMixin, ConfigMixin): - """ - Implements the Euler scheduler in EDM formulation as presented in Karras et al. 2022 [1]. - - [1] Karras, Tero, et al. "Elucidating the Design Space of Diffusion-Based Generative Models." - https://huggingface.co/papers/2206.00364 - - 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: - 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): - 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): - 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): - 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://imagen.research.google/video/paper.pdf) 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"`): - 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. - """ - - _compatibles = [] - order = 1 - - @register_to_config - def __init__( - self, - sigma_min: float = 0.002, - sigma_max: float = 80.0, - sigma_data: float = 0.5, - sigma_schedule: str = "karras", - num_train_timesteps: int = 1000, - prediction_type: str = "epsilon", - rho: float = 7.0, - final_sigmas_type: str = "zero", # can be "zero" or "sigma_min" - ): - if sigma_schedule not in ["karras", "exponential"]: - raise ValueError(f"Wrong value for provided for `{sigma_schedule=}`.`") - - # setable values - self.num_inference_steps = None - - sigmas_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64 - sigmas = torch.arange(num_train_timesteps + 1, dtype=sigmas_dtype) / num_train_timesteps - if sigma_schedule == "karras": - sigmas = self._compute_karras_sigmas(sigmas) - elif sigma_schedule == "exponential": - sigmas = self._compute_exponential_sigmas(sigmas) - sigmas = sigmas.to(torch.float32) - - self.timesteps = self.precondition_noise(sigmas) - - if self.config.final_sigmas_type == "sigma_min": - sigma_last = sigmas[-1] - elif self.config.final_sigmas_type == "zero": - sigma_last = 0 - else: - raise ValueError( - f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" - ) - - self.sigmas = torch.cat([sigmas, torch.full((1,), fill_value=sigma_last, device=sigmas.device)]) - - self.is_scale_input_called = False - - self._step_index = None - self._begin_index = None - 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 - return (self.config.sigma_max**2 + 1) ** 0.5 - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 precondition_inputs(self, sample, sigma): - c_in = self._get_conditioning_c_in(sigma) - scaled_sample = sample * c_in - return scaled_sample - - def precondition_noise(self, sigma): - if not isinstance(sigma, torch.Tensor): - sigma = torch.tensor([sigma]) - - c_noise = 0.25 * torch.log(sigma) - - return c_noise - - def precondition_outputs(self, sample, model_output, sigma): - sigma_data = self.config.sigma_data - c_skip = sigma_data**2 / (sigma**2 + sigma_data**2) - - if self.config.prediction_type == "epsilon": - c_out = sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5 - elif self.config.prediction_type == "v_prediction": - c_out = -sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5 - else: - raise ValueError(f"Prediction type {self.config.prediction_type} is not supported.") - - denoised = c_skip * sample + c_out * model_output - - return denoised - - def scale_model_input(self, sample: torch.Tensor, timestep: Union[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. - - Args: - sample (`torch.Tensor`): - The input sample. - timestep (`int`, *optional*): - The current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - if self.step_index is None: - self._init_step_index(timestep) - - sigma = self.sigmas[self.step_index] - sample = self.precondition_inputs(sample, sigma) - - self.is_scale_input_called = True - return sample - - def set_timesteps( - self, - num_inference_steps: int = None, - device: Union[str, torch.device] = None, - sigmas: Optional[Union[torch.Tensor, List[float]]] = 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 timesteps are not moved. - sigmas (`Union[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. - """ - self.num_inference_steps = num_inference_steps - - sigmas_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64 - if sigmas is None: - sigmas = torch.linspace(0, 1, self.num_inference_steps, dtype=sigmas_dtype) - elif isinstance(sigmas, float): - sigmas = torch.tensor(sigmas, dtype=sigmas_dtype) - else: - sigmas = sigmas.to(sigmas_dtype) - if self.config.sigma_schedule == "karras": - sigmas = self._compute_karras_sigmas(sigmas) - elif self.config.sigma_schedule == "exponential": - sigmas = self._compute_exponential_sigmas(sigmas) - sigmas = sigmas.to(dtype=torch.float32, device=device) - - self.timesteps = self.precondition_noise(sigmas) - - if self.config.final_sigmas_type == "sigma_min": - sigma_last = sigmas[-1] - elif self.config.final_sigmas_type == "zero": - sigma_last = 0 - else: - raise ValueError( - f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" - ) - - self.sigmas = torch.cat([sigmas, torch.full((1,), fill_value=sigma_last, device=sigmas.device)]) - self._step_index = None - self._begin_index = None - 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).""" - sigma_min = sigma_min or self.config.sigma_min - sigma_max = sigma_max or self.config.sigma_max - - rho = self.config.rho - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - 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. - - https://github.com/crowsonkb/k-diffusion/blob/6ab5146d4a5ef63901326489f31f1d8e7dd36b48/k_diffusion/sampling.py#L26 - """ - sigma_min = sigma_min or self.config.sigma_min - sigma_max = sigma_max or self.config.sigma_max - sigmas = torch.linspace(math.log(sigma_min), math.log(sigma_max), len(ramp)).exp().flip(0) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep - 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() - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - 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( - self, - model_output: torch.Tensor, - timestep: Union[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, - return_dict: bool = True, - pred_original_sample: Optional[torch.Tensor] = None, - ) -> Union[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 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): - 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. - - 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. - """ - - if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)): - raise ValueError( - ( - "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" - " `EDMEulerScheduler.step()` is not supported. Make sure to pass" - " one of the `scheduler.timesteps` as a timestep." - ), - ) - - if not self.is_scale_input_called: - logger.warning( - "The `scale_model_input` function should be called before `step` to ensure correct denoising. " - "See `StableDiffusionPipeline` for a usage example." - ) - - if self.step_index is None: - self._init_step_index(timestep) - - # Upcast to avoid precision issues when computing prev_sample - sample = sample.to(torch.float32) - - sigma = self.sigmas[self.step_index] - - gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0 - - sigma_hat = sigma * (gamma + 1) - - if gamma > 0: - noise = randn_tensor( - 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 - - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - if pred_original_sample is None: - pred_original_sample = self.precondition_outputs(sample, model_output, sigma_hat) - - # 2. Convert to an ODE derivative - derivative = (sample - pred_original_sample) / sigma_hat - - dt = self.sigmas[self.step_index + 1] - sigma_hat - - prev_sample = sample + derivative * dt - - # 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, - pred_original_sample, - ) - - return EDMEulerSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.Tensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - noisy_samples = original_samples + noise * sigma - return noisy_samples - - def _get_conditioning_c_in(self, sigma): - c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5) - return c_in - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_euler_ancestral_discrete.py deleted file mode 100644 index 9cdaa2c5e10116b5b000c976e10f06e4cd8f0138..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ /dev/null @@ -1,482 +0,0 @@ -# Copyright 2025 Katherine Crowson 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 List, Optional, Tuple, Union - -import numpy as np -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 KarrasDiffusionSchedulers, SchedulerMixin - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->EulerAncestralDiscrete -class EulerAncestralDiscreteSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr -def rescale_zero_terminal_snr(betas): - """ - 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. - - Returns: - `torch.Tensor`: rescaled betas with zero terminal SNR - """ - # Convert betas to alphas_bar_sqrt - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - alphas_bar_sqrt = alphas_cumprod.sqrt() - - # Store old values. - alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() - alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() - - # Shift so the last timestep is zero. - alphas_bar_sqrt -= alphas_bar_sqrt_T - - # Scale so the first timestep is back to the old value. - alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - - # Convert alphas_bar_sqrt to betas - alphas_bar = alphas_bar_sqrt**2 # Revert sqrt - alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod - alphas = torch.cat([alphas_bar[0:1], alphas]) - betas = 1 - alphas - - return betas - - -class EulerAncestralDiscreteScheduler(SchedulerMixin, ConfigMixin): - """ - Ancestral sampling with Euler method steps. - - 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`, 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` or `scaled_linear`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - 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://imagen.research.google/video/paper.pdf) paper). - 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. - 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 - 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). - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - 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, - prediction_type: str = "epsilon", - timestep_spacing: str = "linspace", - steps_offset: int = 0, - rescale_betas_zero_snr: bool = False, - ): - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - if rescale_betas_zero_snr: - self.betas = rescale_zero_terminal_snr(self.betas) - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - if rescale_betas_zero_snr: - # Close to 0 without being 0 so first sigma is not inf - # FP16 smallest positive subnormal works well here - self.alphas_cumprod[-1] = 2**-24 - - sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) - sigmas = np.concatenate([sigmas[::-1], [0.0]]).astype(np.float32) - self.sigmas = torch.from_numpy(sigmas) - - # setable values - self.num_inference_steps = None - timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=float)[::-1].copy() - self.timesteps = torch.from_numpy(timesteps) - self.is_scale_input_called = False - - self._step_index = None - self._begin_index = None - 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 - if self.config.timestep_spacing in ["linspace", "trailing"]: - return self.sigmas.max() - - return (self.sigmas.max() ** 2 + 1) ** 0.5 - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 scale_model_input(self, sample: torch.Tensor, timestep: Union[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. - - Args: - sample (`torch.Tensor`): - The input sample. - timestep (`int`, *optional*): - The current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - - if self.step_index is None: - self._init_step_index(timestep) - - sigma = self.sigmas[self.step_index] - sample = sample / ((sigma**2 + 1) ** 0.5) - self.is_scale_input_called = True - return sample - - def set_timesteps(self, num_inference_steps: int, device: Union[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 (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - """ - self.num_inference_steps = num_inference_steps - - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps, dtype=np.float32)[ - ::-1 - ].copy() - elif self.config.timestep_spacing == "leading": - step_ratio = self.config.num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.float32) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(self.config.num_train_timesteps, 0, -step_ratio)).round().copy().astype(np.float32) - timesteps -= 1 - else: - raise ValueError( - 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) - sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) - sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) - self.sigmas = torch.from_numpy(sigmas).to(device=device) - - self.timesteps = torch.from_numpy(timesteps).to(device=device) - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep - 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() - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - 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( - self, - model_output: torch.Tensor, - timestep: Union[float, torch.Tensor], - sample: torch.Tensor, - generator: Optional[torch.Generator] = None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - timestep (`float`): - 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`): - Whether or not to return a - [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] or tuple. - - Returns: - [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] or `tuple`: - If return_dict is `True`, - [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] is returned, - otherwise a tuple is returned where the first element is the sample tensor. - - """ - - if isinstance(timestep, (int, torch.IntTensor, 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 not self.is_scale_input_called: - logger.warning( - "The `scale_model_input` function should be called before `step` to ensure correct denoising. " - "See `StableDiffusionPipeline` for a usage example." - ) - - if self.step_index is None: - self._init_step_index(timestep) - - sigma = self.sigmas[self.step_index] - - # Upcast to avoid precision issues when computing prev_sample - sample = sample.to(torch.float32) - - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - if self.config.prediction_type == "epsilon": - pred_original_sample = sample - sigma * model_output - elif self.config.prediction_type == "v_prediction": - # * c_out + input * c_skip - pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1)) - elif self.config.prediction_type == "sample": - raise NotImplementedError("prediction_type not implemented yet: sample") - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`" - ) - - sigma_from = self.sigmas[self.step_index] - sigma_to = self.sigmas[self.step_index + 1] - sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5 - sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5 - - # 2. Convert to an ODE derivative - derivative = (sample - pred_original_sample) / sigma - - dt = sigma_down - sigma - - prev_sample = sample + derivative * dt - - device = model_output.device - noise = randn_tensor(model_output.shape, dtype=model_output.dtype, device=device, generator=generator) - - prev_sample = prev_sample + noise * sigma_up - - # 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, - pred_original_sample, - ) - - return EulerAncestralDiscreteSchedulerOutput( - prev_sample=prev_sample, pred_original_sample=pred_original_sample - ) - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.Tensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - noisy_samples = original_samples + noise * sigma - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_euler_discrete.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_euler_discrete.py deleted file mode 100644 index f58d918dbfbe54babe068eb0b2d60189bf92e61e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_euler_discrete.py +++ /dev/null @@ -1,757 +0,0 @@ -# Copyright 2025 Katherine Crowson 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 List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput, is_scipy_available, logging -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin - - -if is_scipy_available(): - import scipy.stats - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->EulerDiscrete -class EulerDiscreteSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr -def rescale_zero_terminal_snr(betas): - """ - 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. - - Returns: - `torch.Tensor`: rescaled betas with zero terminal SNR - """ - # Convert betas to alphas_bar_sqrt - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - alphas_bar_sqrt = alphas_cumprod.sqrt() - - # Store old values. - alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() - alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() - - # Shift so the last timestep is zero. - alphas_bar_sqrt -= alphas_bar_sqrt_T - - # Scale so the first timestep is back to the old value. - alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - - # Convert alphas_bar_sqrt to betas - alphas_bar = alphas_bar_sqrt**2 # Revert sqrt - alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod - alphas = torch.cat([alphas_bar[0:1], alphas]) - betas = 1 - alphas - - return betas - - -class EulerDiscreteScheduler(SchedulerMixin, ConfigMixin): - """ - Euler scheduler. - - 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`, 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` or `scaled_linear`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - 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://imagen.research.google/video/paper.pdf) paper). - interpolation_type(`str`, defaults to `"linear"`, *optional*): - The interpolation type to compute intermediate sigmas for the scheduler denoising steps. Should be on of - `"linear"` or `"log_linear"`. - use_karras_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, - the sigmas are determined according to a sequence of noise levels {σi}. - use_exponential_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. - 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. - 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. - 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 - 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). - 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. - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - 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, - prediction_type: str = "epsilon", - interpolation_type: str = "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, - timestep_spacing: str = "linspace", - timestep_type: str = "discrete", # can be "discrete" or "continuous" - steps_offset: int = 0, - rescale_betas_zero_snr: bool = False, - final_sigmas_type: str = "zero", # can be "zero" or "sigma_min" - ): - 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: - raise ValueError( - "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." - ) - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - if rescale_betas_zero_snr: - self.betas = rescale_zero_terminal_snr(self.betas) - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - if rescale_betas_zero_snr: - # Close to 0 without being 0 so first sigma is not inf - # FP16 smallest positive subnormal works well here - self.alphas_cumprod[-1] = 2**-24 - - sigmas = (((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5).flip(0) - timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=float)[::-1].copy() - timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) - - # setable values - self.num_inference_steps = None - - # TODO: Support the full EDM scalings for all prediction types and timestep types - if timestep_type == "continuous" and prediction_type == "v_prediction": - self.timesteps = torch.Tensor([0.25 * sigma.log() for sigma in sigmas]) - else: - self.timesteps = timesteps - - self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) - - self.is_scale_input_called = False - self.use_karras_sigmas = use_karras_sigmas - self.use_exponential_sigmas = use_exponential_sigmas - self.use_beta_sigmas = use_beta_sigmas - - self._step_index = None - self._begin_index = None - 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 - max_sigma = max(self.sigmas) if isinstance(self.sigmas, list) else self.sigmas.max() - if self.config.timestep_spacing in ["linspace", "trailing"]: - return max_sigma - - return (max_sigma**2 + 1) ** 0.5 - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 scale_model_input(self, sample: torch.Tensor, timestep: Union[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. - - Args: - sample (`torch.Tensor`): - The input sample. - timestep (`int`, *optional*): - The current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - if self.step_index is None: - self._init_step_index(timestep) - - sigma = self.sigmas[self.step_index] - sample = sample / ((sigma**2 + 1) ** 0.5) - - self.is_scale_input_called = True - return sample - - def set_timesteps( - self, - num_inference_steps: int = None, - device: Union[str, torch.device] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = 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 timesteps are not moved. - 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 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: - raise ValueError("Only one of `timesteps` or `sigmas` should be set.") - if num_inference_steps is None and timesteps is None and sigmas is None: - raise ValueError("Must pass exactly one of `num_inference_steps` or `timesteps` or `sigmas.") - if num_inference_steps is not None and (timesteps is not None or sigmas is not None): - raise ValueError("Can only pass one of `num_inference_steps` or `timesteps` or `sigmas`.") - if timesteps is not None and self.config.use_karras_sigmas: - raise ValueError("Cannot set `timesteps` with `config.use_karras_sigmas = True`.") - if timesteps is not None and self.config.use_exponential_sigmas: - raise ValueError("Cannot set `timesteps` with `config.use_exponential_sigmas = True`.") - if timesteps is not None and self.config.use_beta_sigmas: - raise ValueError("Cannot set `timesteps` with `config.use_beta_sigmas = True`.") - if ( - timesteps is not None - and self.config.timestep_type == "continuous" - and self.config.prediction_type == "v_prediction" - ): - raise ValueError( - "Cannot set `timesteps` with `config.timestep_type = 'continuous'` and `config.prediction_type = 'v_prediction'`." - ) - - if num_inference_steps is None: - num_inference_steps = len(timesteps) if timesteps is not None else len(sigmas) - 1 - self.num_inference_steps = num_inference_steps - - if sigmas is not None: - log_sigmas = np.log(np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)) - sigmas = np.array(sigmas).astype(np.float32) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas[:-1]]) - - else: - if timesteps is not None: - timesteps = np.array(timesteps).astype(np.float32) - else: - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = np.linspace( - 0, self.config.num_train_timesteps - 1, num_inference_steps, dtype=np.float32 - )[::-1].copy() - elif self.config.timestep_spacing == "leading": - step_ratio = self.config.num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = ( - (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.float32) - ) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = ( - (np.arange(self.config.num_train_timesteps, 0, -step_ratio)).round().copy().astype(np.float32) - ) - timesteps -= 1 - else: - raise ValueError( - 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) - log_sigmas = np.log(sigmas) - if self.config.interpolation_type == "linear": - sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) - elif self.config.interpolation_type == "log_linear": - sigmas = torch.linspace(np.log(sigmas[-1]), np.log(sigmas[0]), num_inference_steps + 1).exp().numpy() - else: - raise ValueError( - f"{self.config.interpolation_type} is not implemented. Please specify interpolation_type to either" - " 'linear' or 'log_linear'" - ) - - if self.config.use_karras_sigmas: - sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=self.num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - - elif self.config.use_exponential_sigmas: - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - - elif self.config.use_beta_sigmas: - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - - if self.config.final_sigmas_type == "sigma_min": - sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5 - elif self.config.final_sigmas_type == "zero": - sigma_last = 0 - else: - raise ValueError( - f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" - ) - - sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) - - sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) - - # TODO: Support the full EDM scalings for all prediction types and timestep types - if self.config.timestep_type == "continuous" and self.config.prediction_type == "v_prediction": - self.timesteps = torch.Tensor([0.25 * sigma.log() for sigma in sigmas[:-1]]).to(device=device) - else: - self.timesteps = torch.from_numpy(timesteps.astype(np.float32)).to(device=device) - - self._step_index = None - self._begin_index = None - self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication - - def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma - log_sigma = np.log(np.maximum(sigma, 1e-10)) - - # get distribution - dists = log_sigma - log_sigmas[:, np.newaxis] - - # get sigmas range - low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) - high_idx = low_idx + 1 - - low = log_sigmas[low_idx] - high = log_sigmas[high_idx] - - # interpolate sigmas - w = (low - log_sigma) / (low - high) - w = np.clip(w, 0, 1) - - # transform interpolation to time range - t = (1 - w) * low_idx + w * high_idx - t = t.reshape(sigma.shape) - return t - - # Copied from https://github.com/crowsonkb/k-diffusion/blob/686dbad0f39640ea25c8a8c6a6e56bb40eacefa2/k_diffusion/sampling.py#L17 - def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor: - """Constructs the noise schedule of Karras et al. (2022).""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - rho = 7.0 # 7.0 is the value used in the paper - ramp = np.linspace(0, 1, num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - # Copied from https://github.com/crowsonkb/k-diffusion/blob/686dbad0f39640ea25c8a8c6a6e56bb40eacefa2/k_diffusion/sampling.py#L26 - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Constructs an exponential noise schedule.""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) - return sigmas - - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: - """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return sigmas - - 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( - self, - model_output: torch.Tensor, - timestep: Union[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, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - timestep (`float`): - 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): - 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.EulerDiscreteSchedulerOutput`] or - tuple. - - Returns: - [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is - returned, otherwise a tuple is returned where the first element is the sample tensor. - """ - - if isinstance(timestep, (int, torch.IntTensor, 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 not self.is_scale_input_called: - logger.warning( - "The `scale_model_input` function should be called before `step` to ensure correct denoising. " - "See `StableDiffusionPipeline` for a usage example." - ) - - if self.step_index is None: - self._init_step_index(timestep) - - # Upcast to avoid precision issues when computing prev_sample - sample = sample.to(torch.float32) - - sigma = self.sigmas[self.step_index] - - gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0 - - sigma_hat = sigma * (gamma + 1) - - if gamma > 0: - noise = randn_tensor( - 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 - - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - # NOTE: "original_sample" should not be an expected prediction_type but is left in for - # backwards compatibility - if self.config.prediction_type == "original_sample" or self.config.prediction_type == "sample": - pred_original_sample = model_output - elif self.config.prediction_type == "epsilon": - pred_original_sample = sample - sigma_hat * model_output - elif self.config.prediction_type == "v_prediction": - # denoised = model_output * c_out + input * c_skip - pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1)) - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`" - ) - - # 2. Convert to an ODE derivative - derivative = (sample - pred_original_sample) / sigma_hat - - dt = self.sigmas[self.step_index + 1] - sigma_hat - - prev_sample = sample + derivative * dt - - # 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, - pred_original_sample, - ) - - return EulerDiscreteSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) - - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.Tensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - noisy_samples = original_samples + noise * sigma - return noisy_samples - - def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor: - if ( - isinstance(timesteps, int) - or isinstance(timesteps, torch.IntTensor) - or isinstance(timesteps, torch.LongTensor) - ): - raise ValueError( - ( - "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" - " `EulerDiscreteScheduler.get_velocity()` is not supported. Make sure to pass" - " one of the `scheduler.timesteps` as a timestep." - ), - ) - - if sample.device.type == "mps" and torch.is_floating_point(timesteps): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32) - timesteps = timesteps.to(sample.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(sample.device) - timesteps = timesteps.to(sample.device) - - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - alphas_cumprod = self.alphas_cumprod.to(sample) - sqrt_alpha_prod = alphas_cumprod[step_indices] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(sample.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[step_indices]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample - return velocity - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_euler_discrete_flax.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_euler_discrete_flax.py deleted file mode 100644 index dae01302ac6a440f7903cf8b1e22da250b9a3434..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_euler_discrete_flax.py +++ /dev/null @@ -1,265 +0,0 @@ -# Copyright 2025 Katherine Crowson 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 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 .scheduling_utils_flax import ( - CommonSchedulerState, - FlaxKarrasDiffusionSchedulers, - FlaxSchedulerMixin, - FlaxSchedulerOutput, - broadcast_to_shape_from_left, -) - - -@flax.struct.dataclass -class EulerDiscreteSchedulerState: - common: CommonSchedulerState - - # setable values - init_noise_sigma: jnp.ndarray - timesteps: jnp.ndarray - sigmas: jnp.ndarray - num_inference_steps: Optional[int] = None - - @classmethod - def create( - 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) - - -@dataclass -class FlaxEulerDiscreteSchedulerOutput(FlaxSchedulerOutput): - state: EulerDiscreteSchedulerState - - -class FlaxEulerDiscreteScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - Euler scheduler (Algorithm 2) from Karras et al. (2022) https://huggingface.co/papers/2206.00364. . Based on the - original k-diffusion implementation by Katherine Crowson: - https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L51 - - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - 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 - `linear` or `scaled_linear`. - trained_betas (`jnp.ndarray`, optional): - option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. - prediction_type (`str`, default `epsilon`, optional): - 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://imagen.research.google/video/paper.pdf) - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - the `dtype` used for params and computation. - """ - - _compatibles = [e.name for e in FlaxKarrasDiffusionSchedulers] - - dtype: jnp.dtype - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: Optional[jnp.ndarray] = None, - prediction_type: str = "epsilon", - timestep_spacing: str = "linspace", - dtype: jnp.dtype = jnp.float32, - ): - self.dtype = dtype - - def create_state(self, common: Optional[CommonSchedulerState] = None) -> EulerDiscreteSchedulerState: - if common is None: - common = CommonSchedulerState.create(self) - - timesteps = jnp.arange(0, self.config.num_train_timesteps).round()[::-1] - sigmas = ((1 - common.alphas_cumprod) / common.alphas_cumprod) ** 0.5 - sigmas = jnp.interp(timesteps, jnp.arange(0, len(sigmas)), sigmas) - sigmas = jnp.concatenate([sigmas, jnp.array([0.0], dtype=self.dtype)]) - - # standard deviation of the initial noise distribution - if self.config.timestep_spacing in ["linspace", "trailing"]: - init_noise_sigma = sigmas.max() - else: - init_noise_sigma = (sigmas.max() ** 2 + 1) ** 0.5 - - return EulerDiscreteSchedulerState.create( - common=common, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - sigmas=sigmas, - ) - - def scale_model_input(self, state: EulerDiscreteSchedulerState, sample: jnp.ndarray, timestep: int) -> jnp.ndarray: - """ - Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the Euler algorithm. - - Args: - state (`EulerDiscreteSchedulerState`): - the `FlaxEulerDiscreteScheduler` state data class instance. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - timestep (`int`): - current discrete timestep in the diffusion chain. - - Returns: - `jnp.ndarray`: scaled input sample - """ - (step_index,) = jnp.where(state.timesteps == timestep, size=1) - step_index = step_index[0] - - sigma = state.sigmas[step_index] - sample = sample / ((sigma**2 + 1) ** 0.5) - return sample - - def set_timesteps( - 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. - - Args: - state (`EulerDiscreteSchedulerState`): - the `FlaxEulerDiscreteScheduler` state data class instance. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - """ - - if self.config.timestep_spacing == "linspace": - 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) - timesteps += 1 - else: - raise ValueError( - f"timestep_spacing must be one of ['linspace', 'leading'], got {self.config.timestep_spacing}" - ) - - sigmas = ((1 - state.common.alphas_cumprod) / state.common.alphas_cumprod) ** 0.5 - sigmas = jnp.interp(timesteps, jnp.arange(0, len(sigmas)), sigmas) - sigmas = jnp.concatenate([sigmas, jnp.array([0.0], dtype=self.dtype)]) - - # standard deviation of the initial noise distribution - if self.config.timestep_spacing in ["linspace", "trailing"]: - init_noise_sigma = sigmas.max() - else: - init_noise_sigma = (sigmas.max() ** 2 + 1) ** 0.5 - - return state.replace( - timesteps=timesteps, - sigmas=sigmas, - num_inference_steps=num_inference_steps, - init_noise_sigma=init_noise_sigma, - ) - - def step( - self, - state: EulerDiscreteSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - return_dict: bool = True, - ) -> Union[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). - - Args: - state (`EulerDiscreteSchedulerState`): - the `FlaxEulerDiscreteScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - order: coefficient for multi-step inference. - return_dict (`bool`): option for returning tuple rather than FlaxEulerDiscreteScheduler class - - Returns: - [`FlaxEulerDiscreteScheduler`] or `tuple`: [`FlaxEulerDiscreteScheduler`] if `return_dict` is True, - otherwise a `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - if state.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - (step_index,) = jnp.where(state.timesteps == timestep, size=1) - step_index = step_index[0] - - sigma = state.sigmas[step_index] - - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - if self.config.prediction_type == "epsilon": - pred_original_sample = sample - sigma * model_output - elif self.config.prediction_type == "v_prediction": - # * c_out + input * c_skip - pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1)) - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`" - ) - - # 2. Convert to an ODE derivative - derivative = (sample - pred_original_sample) / sigma - - # dt = sigma_down - sigma - dt = state.sigmas[step_index + 1] - sigma - - prev_sample = sample + derivative * dt - - if not return_dict: - return (prev_sample, state) - - return FlaxEulerDiscreteSchedulerOutput(prev_sample=prev_sample, state=state) - - def add_noise( - self, - state: EulerDiscreteSchedulerState, - original_samples: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - sigma = state.sigmas[timesteps].flatten() - sigma = broadcast_to_shape_from_left(sigma, noise.shape) - - noisy_samples = original_samples + noise * sigma - - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_flow_match_euler_discrete.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_flow_match_euler_discrete.py deleted file mode 100644 index 1a4f12ddfa539bb671e09e1ed4ab1af76892b406..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_flow_match_euler_discrete.py +++ /dev/null @@ -1,561 +0,0 @@ -# Copyright 2025 Stability AI, Katherine Crowson 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 List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput, is_scipy_available, logging -from .scheduling_utils import SchedulerMixin - - -if is_scipy_available(): - import scipy.stats - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - """ - - prev_sample: torch.FloatTensor - - -class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): - """ - Euler scheduler. - - 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`, defaults to 1000): - The number of diffusion steps to train the model. - shift (`float`, defaults to 1.0): - The shift value for the timestep schedule. - use_dynamic_shifting (`bool`, defaults to False): - Whether to apply timestep shifting on-the-fly based on the image resolution. - base_shift (`float`, defaults to 0.5): - Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent - with desired output. - max_shift (`float`, defaults to 1.15): - Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be - more exaggerated or stylized. - base_image_seq_len (`int`, defaults to 256): - The base image sequence length. - max_image_seq_len (`int`, defaults to 4096): - The maximum image sequence length. - invert_sigmas (`bool`, defaults to False): - Whether to invert the sigmas. - shift_terminal (`float`, defaults to None): - The end value of the shifted timestep schedule. - use_karras_sigmas (`bool`, defaults to False): - Whether to use Karras sigmas for step sizes in the noise schedule during sampling. - use_exponential_sigmas (`bool`, defaults to False): - Whether to use exponential sigmas for step sizes in the noise schedule during sampling. - 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". - stochastic_sampling (`bool`, defaults to False): - Whether to use stochastic sampling. - """ - - _compatibles = [] - order = 1 - - @register_to_config - def __init__( - self, - 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, - 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", - 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: - raise ValueError( - "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." - ) - if time_shift_type not in {"exponential", "linear"}: - raise ValueError("`time_shift_type` must either be 'exponential' or 'linear'.") - - timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy() - timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) - - sigmas = timesteps / num_train_timesteps - if not use_dynamic_shifting: - # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution - sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) - - self.timesteps = sigmas * num_train_timesteps - - self._step_index = None - self._begin_index = None - - self._shift = shift - - self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication - self.sigma_min = self.sigmas[-1].item() - self.sigma_max = self.sigmas[0].item() - - @property - def shift(self): - """ - The value used for shifting. - """ - return self._shift - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 set_shift(self, shift: float): - self._shift = shift - - def scale_noise( - self, - sample: torch.FloatTensor, - timestep: Union[float, torch.FloatTensor], - noise: Optional[torch.FloatTensor] = None, - ) -> torch.FloatTensor: - """ - Forward process in flow-matching - - Args: - sample (`torch.FloatTensor`): - The input sample. - timestep (`int`, *optional*): - The current timestep in the diffusion chain. - - Returns: - `torch.FloatTensor`: - A scaled input sample. - """ - # Make sure sigmas and timesteps have the same device and dtype as original_samples - sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype) - - if sample.device.type == "mps" and torch.is_floating_point(timestep): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32) - timestep = timestep.to(sample.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(sample.device) - timestep = timestep.to(sample.device) - - # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timestep.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timestep.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(sample.shape): - sigma = sigma.unsqueeze(-1) - - sample = sigma * noise + (1.0 - sigma) * sample - - return sample - - def _sigma_to_t(self, sigma): - return sigma * self.config.num_train_timesteps - - def time_shift(self, mu: float, sigma: float, t: 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: - 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 - - 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, - ): - """ - Sets the discrete timesteps used for the diffusion chain (to be run before inference). - - Args: - 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 (`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*): - Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed - automatically. - """ - 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 and timesteps is not None: - if len(sigmas) != len(timesteps): - raise ValueError("`sigmas` and `timesteps` should have the same length") - - if num_inference_steps is not None: - if (sigmas is not None and len(sigmas) != num_inference_steps) or ( - timesteps is not None and len(timesteps) != num_inference_steps - ): - raise ValueError( - "`sigmas` and `timesteps` should have the same length as num_inference_steps, if `num_inference_steps` is provided" - ) - else: - num_inference_steps = len(sigmas) if sigmas is not None else len(timesteps) - - self.num_inference_steps = num_inference_steps - - # 1. Prepare default sigmas - is_timesteps_provided = timesteps is not None - - if is_timesteps_provided: - timesteps = np.array(timesteps).astype(np.float32) - - 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 - ) - sigmas = timesteps / self.config.num_train_timesteps - else: - sigmas = np.array(sigmas).astype(np.float32) - 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) - else: - sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) - - # 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) - - # 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) - elif self.config.use_exponential_sigmas: - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - elif self.config.use_beta_sigmas: - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - - # 5. Convert sigmas and timesteps to tensors and move to specified device - sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) - if not is_timesteps_provided: - timesteps = sigmas * self.config.num_train_timesteps - else: - timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device) - - # 6. Append the terminal sigma value. - # If a model requires inverted sigma schedule for denoising but timesteps without inversion, the - # `invert_sigmas` flag can be set to `True`. This case is only required in Mochi - if self.config.invert_sigmas: - sigmas = 1.0 - sigmas - timesteps = sigmas * self.config.num_train_timesteps - sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)]) - else: - sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) - - self.timesteps = timesteps - self.sigmas = sigmas - self._step_index = None - self._begin_index = None - - 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( - self, - model_output: torch.FloatTensor, - timestep: Union[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, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.FloatTensor`): - The direct output from learned diffusion model. - timestep (`float`): - 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`): - s_noise (`float`, defaults to 1.0): - Scaling factor for noise added to the sample. - generator (`torch.Generator`, *optional*): - A random number generator. - per_token_timesteps (`torch.Tensor`, *optional*): - The timesteps for each token in the sample. - return_dict (`bool`): - Whether or not to return a - [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple. - - Returns: - [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`: - If return_dict is `True`, - [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned, - otherwise a tuple is returned where the first element is the sample tensor. - """ - - 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" - " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass" - " one of the `scheduler.timesteps` as a timestep." - ), - ) - - if self.step_index is None: - self._init_step_index(timestep) - - # Upcast to avoid precision issues when computing prev_sample - sample = sample.to(torch.float32) - - if per_token_timesteps is not None: - per_token_sigmas = per_token_timesteps / self.config.num_train_timesteps - - sigmas = self.sigmas[:, None, None] - lower_mask = sigmas < per_token_sigmas[None] - 1e-6 - lower_sigmas = lower_mask * sigmas - lower_sigmas, _ = lower_sigmas.max(dim=0) - - current_sigma = per_token_sigmas[..., None] - next_sigma = lower_sigmas[..., None] - dt = current_sigma - next_sigma - else: - sigma_idx = self.step_index - sigma = self.sigmas[sigma_idx] - sigma_next = self.sigmas[sigma_idx + 1] - - current_sigma = sigma - next_sigma = sigma_next - dt = sigma_next - sigma - - if self.config.stochastic_sampling: - x0 = sample - current_sigma * model_output - noise = torch.randn_like(sample) - prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise - else: - prev_sample = sample + dt * model_output - - # upon completion increase step index by one - self._step_index += 1 - if per_token_timesteps is None: - # Cast sample back to model compatible dtype - prev_sample = prev_sample.to(model_output.dtype) - - if not return_dict: - return (prev_sample,) - - 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: - """Constructs the noise schedule of Karras et al. (2022).""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - rho = 7.0 # 7.0 is the value used in the paper - ramp = np.linspace(0, 1, num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Constructs an exponential noise schedule.""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: - """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return sigmas - - def _time_shift_exponential(self, mu, sigma, t): - return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) - - def _time_shift_linear(self, mu, sigma, t): - return mu / (mu + (1 / t - 1) ** sigma) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_flow_match_heun_discrete.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_flow_match_heun_discrete.py deleted file mode 100644 index 38e5f1ba77a8511544a3f10d9ff7382ccda5c83b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_flow_match_heun_discrete.py +++ /dev/null @@ -1,321 +0,0 @@ -# Copyright 2025 Stability AI, Katherine Crowson 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 dataclasses import dataclass -from typing import Optional, Tuple, Union - -import numpy as np -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 FlowMatchHeunDiscreteSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - """ - - prev_sample: torch.FloatTensor - - -class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin): - """ - Heun scheduler. - - 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`, 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. - """ - - _compatibles = [] - order = 2 - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - shift: float = 1.0, - ): - timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy() - timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) - - sigmas = timesteps / num_train_timesteps - sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) - - self.timesteps = sigmas * num_train_timesteps - - self._step_index = None - self._begin_index = None - - self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication - self.sigma_min = self.sigmas[-1].item() - self.sigma_max = self.sigmas[0].item() - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 scale_noise( - self, - sample: torch.FloatTensor, - timestep: Union[float, torch.FloatTensor], - noise: Optional[torch.FloatTensor] = None, - ) -> torch.FloatTensor: - """ - Forward process in flow-matching - - Args: - sample (`torch.FloatTensor`): - The input sample. - timestep (`int`, *optional*): - The current timestep in the diffusion chain. - - Returns: - `torch.FloatTensor`: - A scaled input sample. - """ - if self.step_index is None: - 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): - return sigma * self.config.num_train_timesteps - - def set_timesteps(self, num_inference_steps: int, device: Union[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 (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - """ - 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 - ) - - sigmas = timesteps / self.config.num_train_timesteps - sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas) - sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) - - timesteps = sigmas * self.config.num_train_timesteps - timesteps = torch.cat([timesteps[:1], timesteps[1:].repeat_interleave(2)]) - self.timesteps = timesteps.to(device=device) - - sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) - self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]]) - - # empty dt and derivative - self.prev_derivative = None - self.dt = None - - self._step_index = None - self._begin_index = None - - 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 - - @property - def state_in_first_order(self): - return self.dt is None - - def step( - self, - model_output: torch.FloatTensor, - timestep: Union[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, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.FloatTensor`): - The direct output from learned diffusion model. - timestep (`float`): - 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`): - s_noise (`float`, 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_flow_match_heun_discrete.FlowMatchHeunDiscreteSchedulerOutput`] tuple. - - Returns: - [`~schedulers.scheduling_flow_match_heun_discrete.FlowMatchHeunDiscreteSchedulerOutput`] or `tuple`: - If return_dict is `True`, - [`~schedulers.scheduling_flow_match_heun_discrete.FlowMatchHeunDiscreteSchedulerOutput`] is returned, - otherwise a tuple is returned where the first element is the sample tensor. - """ - - 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" - " `FlowMatchHeunDiscreteScheduler.step()` is not supported. Make sure to pass" - " one of the `scheduler.timesteps` as a timestep." - ), - ) - - if self.step_index is None: - self._init_step_index(timestep) - - # Upcast to avoid precision issues when computing prev_sample - sample = sample.to(torch.float32) - - if self.state_in_first_order: - sigma = self.sigmas[self.step_index] - sigma_next = self.sigmas[self.step_index + 1] - else: - # 2nd order / Heun's method - sigma = self.sigmas[self.step_index - 1] - sigma_next = self.sigmas[self.step_index] - - gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0 - - sigma_hat = sigma * (gamma + 1) - - if gamma > 0: - noise = randn_tensor( - 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 - - if self.state_in_first_order: - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - denoised = sample - model_output * sigma - # 2. convert to an ODE derivative for 1st order - derivative = (sample - denoised) / sigma_hat - # 3. Delta timestep - dt = sigma_next - sigma_hat - - # store for 2nd order step - self.prev_derivative = derivative - self.dt = dt - self.sample = sample - else: - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - denoised = sample - model_output * sigma_next - # 2. 2nd order / Heun's method - derivative = (sample - denoised) / sigma_next - derivative = 0.5 * (self.prev_derivative + derivative) - - # 3. take prev timestep & sample - dt = self.dt - sample = self.sample - - # free dt and derivative - # Note, this puts the scheduler in "first order mode" - self.prev_derivative = None - self.dt = None - self.sample = None - - prev_sample = sample + derivative * dt - # 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 FlowMatchHeunDiscreteSchedulerOutput(prev_sample=prev_sample) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_flow_match_lcm.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_flow_match_lcm.py deleted file mode 100644 index 933bb1cf8e3d05d58ca6216ada6c40e241d59fa2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_flow_match_lcm.py +++ /dev/null @@ -1,561 +0,0 @@ -# Copyright 2025 Stability AI, Katherine Crowson 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 List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput, is_scipy_available, logging -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import SchedulerMixin - - -if is_scipy_available(): - import scipy.stats - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class FlowMatchLCMSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - """ - - prev_sample: torch.FloatTensor - - -class FlowMatchLCMScheduler(SchedulerMixin, ConfigMixin): - """ - LCM scheduler for Flow Matching. - - 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`, defaults to 1000): - The number of diffusion steps to train the model. - shift (`float`, defaults to 1.0): - The shift value for the timestep schedule. - use_dynamic_shifting (`bool`, defaults to False): - Whether to apply timestep shifting on-the-fly based on the image resolution. - base_shift (`float`, defaults to 0.5): - Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent - with desired output. - max_shift (`float`, defaults to 1.15): - Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be - more exaggerated or stylized. - base_image_seq_len (`int`, defaults to 256): - The base image sequence length. - max_image_seq_len (`int`, defaults to 4096): - The maximum image sequence length. - invert_sigmas (`bool`, defaults to False): - Whether to invert the sigmas. - shift_terminal (`float`, defaults to None): - The end value of the shifted timestep schedule. - use_karras_sigmas (`bool`, defaults to False): - Whether to use Karras sigmas for step sizes in the noise schedule during sampling. - use_exponential_sigmas (`bool`, defaults to False): - Whether to use exponential sigmas for step sizes in the noise schedule during sampling. - 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) - 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 - """ - - _compatibles = [] - order = 1 - - @register_to_config - def __init__( - self, - 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, - 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", - ): - 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: - raise ValueError( - "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." - ) - if time_shift_type not in {"exponential", "linear"}: - raise ValueError("`time_shift_type` must either be 'exponential' or 'linear'.") - - timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy() - timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) - - sigmas = timesteps / num_train_timesteps - if not use_dynamic_shifting: - # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution - sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) - - self.timesteps = sigmas * num_train_timesteps - - self._step_index = None - self._begin_index = None - - self._shift = shift - - self._init_size = None - self._scale_factors = scale_factors - self._upscale_mode = upscale_mode - - self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication - self.sigma_min = self.sigmas[-1].item() - self.sigma_max = self.sigmas[0].item() - - @property - def shift(self): - """ - The value used for shifting. - """ - return self._shift - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 set_shift(self, shift: float): - self._shift = shift - - def set_scale_factors(self, scale_factors: list, upscale_mode): - """ - Sets scale factors for a scale-wise generation regime. - - Args: - scale_factors (`list`): - The scale factors for each step - upscale_mode (`str`): - Upscaling method - """ - self._scale_factors = scale_factors - self._upscale_mode = upscale_mode - - def scale_noise( - self, - sample: torch.FloatTensor, - timestep: Union[float, torch.FloatTensor], - noise: Optional[torch.FloatTensor] = None, - ) -> torch.FloatTensor: - """ - Forward process in flow-matching - - Args: - sample (`torch.FloatTensor`): - The input sample. - timestep (`int`, *optional*): - The current timestep in the diffusion chain. - - Returns: - `torch.FloatTensor`: - A scaled input sample. - """ - # Make sure sigmas and timesteps have the same device and dtype as original_samples - sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype) - - if sample.device.type == "mps" and torch.is_floating_point(timestep): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32) - timestep = timestep.to(sample.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(sample.device) - timestep = timestep.to(sample.device) - - # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timestep.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timestep.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(sample.shape): - sigma = sigma.unsqueeze(-1) - - sample = sigma * noise + (1.0 - sigma) * sample - - return sample - - def _sigma_to_t(self, sigma): - return sigma * self.config.num_train_timesteps - - def time_shift(self, mu: float, sigma: float, t: 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: - 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 - - 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, - ): - """ - Sets the discrete timesteps used for the diffusion chain (to be run before inference). - - Args: - 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 (`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*): - Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed - automatically. - """ - 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 and timesteps is not None: - if len(sigmas) != len(timesteps): - raise ValueError("`sigmas` and `timesteps` should have the same length") - - if num_inference_steps is not None: - if (sigmas is not None and len(sigmas) != num_inference_steps) or ( - timesteps is not None and len(timesteps) != num_inference_steps - ): - raise ValueError( - "`sigmas` and `timesteps` should have the same length as num_inference_steps, if `num_inference_steps` is provided" - ) - else: - num_inference_steps = len(sigmas) if sigmas is not None else len(timesteps) - - self.num_inference_steps = num_inference_steps - - # 1. Prepare default sigmas - is_timesteps_provided = timesteps is not None - - if is_timesteps_provided: - timesteps = np.array(timesteps).astype(np.float32) - - 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 - ) - sigmas = timesteps / self.config.num_train_timesteps - else: - sigmas = np.array(sigmas).astype(np.float32) - 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) - else: - sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) - - # 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) - - # 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) - elif self.config.use_exponential_sigmas: - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - elif self.config.use_beta_sigmas: - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - - # 5. Convert sigmas and timesteps to tensors and move to specified device - sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) - if not is_timesteps_provided: - timesteps = sigmas * self.config.num_train_timesteps - else: - timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device) - - # 6. Append the terminal sigma value. - # If a model requires inverted sigma schedule for denoising but timesteps without inversion, the - # `invert_sigmas` flag can be set to `True`. This case is only required in Mochi - if self.config.invert_sigmas: - sigmas = 1.0 - sigmas - timesteps = sigmas * self.config.num_train_timesteps - sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)]) - else: - sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) - - self.timesteps = timesteps - self.sigmas = sigmas - self._step_index = None - self._begin_index = None - - 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( - self, - model_output: torch.FloatTensor, - timestep: Union[float, torch.FloatTensor], - sample: torch.FloatTensor, - generator: Optional[torch.Generator] = None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.FloatTensor`): - The direct output from learned diffusion model. - timestep (`float`): - The current discrete timestep in the diffusion chain. - sample (`torch.FloatTensor`): - A current instance of a sample created by the diffusion process. - generator (`torch.Generator`, *optional*): - A random number generator. - return_dict (`bool`): - Whether or not to return a [`~schedulers.scheduling_flow_match_lcm.FlowMatchLCMSchedulerOutput`] or - tuple. - - Returns: - [`~schedulers.scheduling_flow_match_lcm.FlowMatchLCMSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_flow_match_lcm.FlowMatchLCMSchedulerOutput`] is - returned, otherwise a tuple is returned where the first element is the sample tensor. - """ - - 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" - " `FlowMatchLCMScheduler.step()` is not supported. Make sure to pass" - " one of the `scheduler.timesteps` as a timestep." - ), - ) - - if self._scale_factors and self._upscale_mode and len(self.timesteps) != len(self._scale_factors) + 1: - raise ValueError( - "`_scale_factors` should have the same length as `timesteps` - 1, if `_scale_factors` are set." - ) - - if self._init_size is None or self.step_index is None: - self._init_size = model_output.size()[2:] - - if self.step_index is None: - self._init_step_index(timestep) - - # Upcast to avoid precision issues when computing prev_sample - sample = sample.to(torch.float32) - - sigma = self.sigmas[self.step_index] - sigma_next = self.sigmas[self.step_index + 1] - x0_pred = sample - sigma * model_output - - if self._scale_factors and self._upscale_mode: - if self._step_index < len(self._scale_factors): - 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) - prev_sample = (1 - sigma_next) * x0_pred + sigma_next * noise - - # upon completion increase step index by one - self._step_index += 1 - # Cast sample back to model compatible dtype - prev_sample = prev_sample.to(model_output.dtype) - - if not return_dict: - return (prev_sample,) - - 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: - """Constructs the noise schedule of Karras et al. (2022).""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - rho = 7.0 # 7.0 is the value used in the paper - ramp = np.linspace(0, 1, num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Constructs an exponential noise schedule.""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: - """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return sigmas - - def _time_shift_exponential(self, mu, sigma, t): - return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) - - def _time_shift_linear(self, mu, sigma, t): - return mu / (mu + (1 / t - 1) ** sigma) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_heun_discrete.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_heun_discrete.py deleted file mode 100644 index bd1239cfaec7fb71e07583f940e32ad700da1b04..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_heun_discrete.py +++ /dev/null @@ -1,610 +0,0 @@ -# Copyright 2025 Katherine Crowson, The HuggingFace Team and hlky. 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 List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput, is_scipy_available -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin - - -if is_scipy_available(): - import scipy.stats - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->HeunDiscrete -class HeunDiscreteSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -class HeunDiscreteScheduler(SchedulerMixin, ConfigMixin): - """ - Scheduler with Heun steps for discrete beta schedules. - - 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`, 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` or `scaled_linear`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - 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://imagen.research.google/video/paper.pdf) paper). - clip_sample (`bool`, defaults to `True`): - Clip the predicted sample for numerical stability. - clip_sample_range (`float`, defaults to 1.0): - The maximum magnitude for sample clipping. Valid only when `clip_sample=True`. - use_karras_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, - the sigmas are determined according to a sequence of noise levels {σi}. - use_exponential_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. - 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. - 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. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 2 - - @register_to_config - def __init__( - self, - 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, - clip_sample: Optional[bool] = False, - clip_sample_range: float = 1.0, - timestep_spacing: str = "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: - raise ValueError( - "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." - ) - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="cosine") - elif beta_schedule == "exp": - self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="exp") - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - # set all values - self.set_timesteps(num_train_timesteps, None, num_train_timesteps) - self.use_karras_sigmas = use_karras_sigmas - - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep - 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() - - @property - def init_noise_sigma(self): - # standard deviation of the initial noise distribution - if self.config.timestep_spacing in ["linspace", "trailing"]: - return self.sigmas.max() - - return (self.sigmas.max() ** 2 + 1) ** 0.5 - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 scale_model_input( - self, - sample: torch.Tensor, - timestep: Union[float, torch.Tensor], - ) -> torch.Tensor: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - if self.step_index is None: - self._init_step_index(timestep) - - sigma = self.sigmas[self.step_index] - sample = sample / ((sigma**2 + 1) ** 0.5) - return sample - - 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, - ): - """ - 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 timesteps are not moved. - num_train_timesteps (`int`, *optional*): - The number of diffusion steps used when training the model. If `None`, the default - `num_train_timesteps` attribute is used. - 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. - """ - if num_inference_steps is None and timesteps is None: - raise ValueError("Must pass exactly one of `num_inference_steps` or `custom_timesteps`.") - if num_inference_steps is not None and timesteps is not None: - raise ValueError("Can only pass one of `num_inference_steps` or `custom_timesteps`.") - if timesteps is not None and self.config.use_karras_sigmas: - raise ValueError("Cannot use `timesteps` with `config.use_karras_sigmas = True`") - if timesteps is not None and self.config.use_exponential_sigmas: - raise ValueError("Cannot set `timesteps` with `config.use_exponential_sigmas = True`.") - if timesteps is not None and self.config.use_beta_sigmas: - raise ValueError("Cannot set `timesteps` with `config.use_beta_sigmas = True`.") - - num_inference_steps = num_inference_steps or len(timesteps) - self.num_inference_steps = num_inference_steps - num_train_timesteps = num_train_timesteps or self.config.num_train_timesteps - - if timesteps is not None: - timesteps = np.array(timesteps, dtype=np.float32) - else: - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = np.linspace(0, num_train_timesteps - 1, num_inference_steps, dtype=np.float32)[::-1].copy() - elif self.config.timestep_spacing == "leading": - step_ratio = num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.float32) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(num_train_timesteps, 0, -step_ratio)).round().copy().astype(np.float32) - timesteps -= 1 - else: - raise ValueError( - 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) - log_sigmas = np.log(sigmas) - sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) - - if self.config.use_karras_sigmas: - sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=self.num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - elif self.config.use_exponential_sigmas: - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - elif self.config.use_beta_sigmas: - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - - sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) - sigmas = torch.from_numpy(sigmas).to(device=device) - self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]]) - - timesteps = torch.from_numpy(timesteps) - timesteps = torch.cat([timesteps[:1], timesteps[1:].repeat_interleave(2)]) - - self.timesteps = timesteps.to(device=device, dtype=torch.float32) - - # empty dt and derivative - self.prev_derivative = None - self.dt = None - - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma - log_sigma = np.log(np.maximum(sigma, 1e-10)) - - # get distribution - dists = log_sigma - log_sigmas[:, np.newaxis] - - # get sigmas range - low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) - high_idx = low_idx + 1 - - low = log_sigmas[low_idx] - high = log_sigmas[high_idx] - - # interpolate sigmas - w = (low - log_sigma) / (low - high) - w = np.clip(w, 0, 1) - - # transform interpolation to time range - t = (1 - w) * low_idx + w * high_idx - t = t.reshape(sigma.shape) - return 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: - """Constructs the noise schedule of Karras et al. (2022).""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - rho = 7.0 # 7.0 is the value used in the paper - ramp = np.linspace(0, 1, num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Constructs an exponential noise schedule.""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: - """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return sigmas - - @property - 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): - 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( - self, - model_output: Union[torch.Tensor, np.ndarray], - timestep: Union[float, torch.Tensor], - sample: Union[torch.Tensor, np.ndarray], - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - timestep (`float`): - The current discrete timestep in the diffusion chain. - 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_heun_discrete.HeunDiscreteSchedulerOutput`] or - tuple. - - Returns: - [`~schedulers.scheduling_heun_discrete.HeunDiscreteSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_heun_discrete.HeunDiscreteSchedulerOutput`] is - returned, otherwise a tuple is returned where the first element is the sample tensor. - """ - if self.step_index is None: - self._init_step_index(timestep) - - if self.state_in_first_order: - sigma = self.sigmas[self.step_index] - sigma_next = self.sigmas[self.step_index + 1] - else: - # 2nd order / Heun's method - sigma = self.sigmas[self.step_index - 1] - sigma_next = self.sigmas[self.step_index] - - # currently only gamma=0 is supported. This usually works best anyways. - # We can support gamma in the future but then need to scale the timestep before - # passing it to the model which requires a change in API - gamma = 0 - sigma_hat = sigma * (gamma + 1) # Note: sigma_hat == sigma for now - - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - if self.config.prediction_type == "epsilon": - sigma_input = sigma_hat if self.state_in_first_order else sigma_next - pred_original_sample = sample - sigma_input * model_output - elif self.config.prediction_type == "v_prediction": - sigma_input = sigma_hat if self.state_in_first_order else sigma_next - pred_original_sample = model_output * (-sigma_input / (sigma_input**2 + 1) ** 0.5) + ( - sample / (sigma_input**2 + 1) - ) - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`" - ) - - if self.config.clip_sample: - pred_original_sample = pred_original_sample.clamp( - -self.config.clip_sample_range, self.config.clip_sample_range - ) - - if self.state_in_first_order: - # 2. Convert to an ODE derivative for 1st order - derivative = (sample - pred_original_sample) / sigma_hat - # 3. delta timestep - dt = sigma_next - sigma_hat - - # store for 2nd order step - self.prev_derivative = derivative - self.dt = dt - self.sample = sample - else: - # 2. 2nd order / Heun's method - derivative = (sample - pred_original_sample) / sigma_next - derivative = (self.prev_derivative + derivative) / 2 - - # 3. take prev timestep & sample - dt = self.dt - sample = self.sample - - # free dt and derivative - # Note, this puts the scheduler in "first order mode" - self.prev_derivative = None - self.dt = None - self.sample = None - - prev_sample = sample + derivative * dt - - # upon completion increase step index by one - self._step_index += 1 - - if not return_dict: - return ( - prev_sample, - pred_original_sample, - ) - - return HeunDiscreteSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.Tensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - noisy_samples = original_samples + noise * sigma - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ipndm.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ipndm.py deleted file mode 100644 index 23bc21f10ca426c8a3bb864d540f3a6b7688da6f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_ipndm.py +++ /dev/null @@ -1,224 +0,0 @@ -# Copyright 2025 Zhejiang University 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 typing import List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from .scheduling_utils import SchedulerMixin, SchedulerOutput - - -class IPNDMScheduler(SchedulerMixin, ConfigMixin): - """ - A fourth-order Improved Pseudo Linear Multistep scheduler. - - 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`, defaults to 1000): - The number of diffusion steps to train the model. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - """ - - order = 1 - - @register_to_config - def __init__( - self, num_train_timesteps: int = 1000, trained_betas: Optional[Union[np.ndarray, List[float]]] = None - ): - # set `betas`, `alphas`, `timesteps` - self.set_timesteps(num_train_timesteps) - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # For now we only support F-PNDM, i.e. the runge-kutta method - # For more information on the algorithm please take a look at the paper: https://huggingface.co/papers/2202.09778 - # mainly at formula (9), (12), (13) and the Algorithm 2. - self.pndm_order = 4 - - # running values - self.ets = [] - self._step_index = None - self._begin_index = None - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 set_timesteps(self, num_inference_steps: int, device: Union[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 (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - """ - self.num_inference_steps = num_inference_steps - steps = torch.linspace(1, 0, num_inference_steps + 1)[:-1] - steps = torch.cat([steps, torch.tensor([0.0])]) - - if self.config.trained_betas is not None: - self.betas = torch.tensor(self.config.trained_betas, dtype=torch.float32) - else: - self.betas = torch.sin(steps * math.pi / 2) ** 2 - - self.alphas = (1.0 - self.betas**2) ** 0.5 - - timesteps = (torch.atan2(self.betas, self.alphas) / math.pi * 2)[:-1] - self.timesteps = timesteps.to(device) - - self.ets = [] - self._step_index = None - self._begin_index = None - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep - 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() - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - 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( - self, - model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], - sample: torch.Tensor, - return_dict: bool = True, - ) -> Union[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. - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - return_dict (`bool`): - 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( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - if self.step_index is None: - self._init_step_index(timestep) - - timestep_index = self.step_index - prev_timestep_index = self.step_index + 1 - - ets = sample * self.betas[timestep_index] + model_output * self.alphas[timestep_index] - self.ets.append(ets) - - if len(self.ets) == 1: - ets = self.ets[-1] - elif len(self.ets) == 2: - ets = (3 * self.ets[-1] - self.ets[-2]) / 2 - elif len(self.ets) == 3: - ets = (23 * self.ets[-1] - 16 * self.ets[-2] + 5 * self.ets[-3]) / 12 - else: - ets = (1 / 24) * (55 * self.ets[-1] - 59 * self.ets[-2] + 37 * self.ets[-3] - 9 * self.ets[-4]) - - prev_sample = self._get_prev_sample(sample, timestep_index, prev_timestep_index, ets) - - # upon completion increase step index by one - self._step_index += 1 - - if not return_dict: - return (prev_sample,) - - return SchedulerOutput(prev_sample=prev_sample) - - def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: - """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. - - Args: - sample (`torch.Tensor`): - The input sample. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - def _get_prev_sample(self, sample, timestep_index, prev_timestep_index, ets): - alpha = self.alphas[timestep_index] - sigma = self.betas[timestep_index] - - next_alpha = self.alphas[prev_timestep_index] - next_sigma = self.betas[prev_timestep_index] - - pred = (sample - sigma * ets) / max(alpha, 1e-8) - prev_sample = next_alpha * pred + ets * next_sigma - - return prev_sample - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py deleted file mode 100644 index 6588464073a12a759ba1d22675ee98dde240378f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py +++ /dev/null @@ -1,617 +0,0 @@ -# Copyright 2025 Katherine Crowson, The HuggingFace Team and hlky. 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 List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput, is_scipy_available -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin - - -if is_scipy_available(): - import scipy.stats - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->KDPM2AncestralDiscrete -class KDPM2AncestralDiscreteSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -class KDPM2AncestralDiscreteScheduler(SchedulerMixin, ConfigMixin): - """ - KDPM2DiscreteScheduler with ancestral sampling is inspired by the DPMSolver2 and Algorithm 2 from the [Elucidating - the Design Space of Diffusion-Based Generative Models](https://huggingface.co/papers/2206.00364) paper. - - 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`, defaults to 1000): - The number of diffusion steps to train the model. - beta_start (`float`, defaults to 0.00085): - The starting `beta` value of inference. - beta_end (`float`, defaults to 0.012): - 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` or `scaled_linear`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - use_karras_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, - the sigmas are determined according to a sequence of noise levels {σi}. - use_exponential_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. - 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. - 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://imagen.research.google/video/paper.pdf) paper). - 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. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 2 - - @register_to_config - def __init__( - self, - 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, - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, - prediction_type: str = "epsilon", - timestep_spacing: str = "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: - raise ValueError( - "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." - ) - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - # set all values - self.set_timesteps(num_train_timesteps, None, num_train_timesteps) - self._step_index = None - self._begin_index = None - 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 - if self.config.timestep_spacing in ["linspace", "trailing"]: - return self.sigmas.max() - - return (self.sigmas.max() ** 2 + 1) ** 0.5 - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 scale_model_input( - self, - sample: torch.Tensor, - timestep: Union[float, torch.Tensor], - ) -> torch.Tensor: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - if self.step_index is None: - self._init_step_index(timestep) - - if self.state_in_first_order: - sigma = self.sigmas[self.step_index] - else: - sigma = self.sigmas_interpol[self.step_index - 1] - - sample = sample / ((sigma**2 + 1) ** 0.5) - return sample - - def set_timesteps( - self, - num_inference_steps: int, - device: Union[str, torch.device] = None, - num_train_timesteps: Optional[int] = 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 timesteps are not moved. - """ - self.num_inference_steps = num_inference_steps - - num_train_timesteps = num_train_timesteps or self.config.num_train_timesteps - - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = np.linspace(0, num_train_timesteps - 1, num_inference_steps, dtype=np.float32)[::-1].copy() - elif self.config.timestep_spacing == "leading": - step_ratio = num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.float32) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(num_train_timesteps, 0, -step_ratio)).round().copy().astype(np.float32) - timesteps -= 1 - else: - raise ValueError( - 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) - log_sigmas = np.log(sigmas) - - sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) - - if self.config.use_karras_sigmas: - sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round() - elif self.config.use_exponential_sigmas: - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - elif self.config.use_beta_sigmas: - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - - self.log_sigmas = torch.from_numpy(log_sigmas).to(device) - sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) - sigmas = torch.from_numpy(sigmas).to(device=device) - - # compute up and down sigmas - sigmas_next = sigmas.roll(-1) - sigmas_next[-1] = 0.0 - sigmas_up = (sigmas_next**2 * (sigmas**2 - sigmas_next**2) / sigmas**2) ** 0.5 - sigmas_down = (sigmas_next**2 - sigmas_up**2) ** 0.5 - sigmas_down[-1] = 0.0 - - # compute interpolated sigmas - sigmas_interpol = sigmas.log().lerp(sigmas_down.log(), 0.5).exp() - sigmas_interpol[-2:] = 0.0 - - # set sigmas - self.sigmas = torch.cat([sigmas[:1], sigmas[1:].repeat_interleave(2), sigmas[-1:]]) - self.sigmas_interpol = torch.cat( - [sigmas_interpol[:1], sigmas_interpol[1:].repeat_interleave(2), sigmas_interpol[-1:]] - ) - self.sigmas_up = torch.cat([sigmas_up[:1], sigmas_up[1:].repeat_interleave(2), sigmas_up[-1:]]) - self.sigmas_down = torch.cat([sigmas_down[:1], sigmas_down[1:].repeat_interleave(2), sigmas_down[-1:]]) - - if str(device).startswith("mps"): - timesteps = torch.from_numpy(timesteps).to(device, dtype=torch.float32) - else: - timesteps = torch.from_numpy(timesteps).to(device) - - sigmas_interpol = sigmas_interpol.cpu() - log_sigmas = self.log_sigmas.cpu() - timesteps_interpol = np.array( - [self._sigma_to_t(sigma_interpol, log_sigmas) for sigma_interpol in sigmas_interpol] - ) - - timesteps_interpol = torch.from_numpy(timesteps_interpol).to(device, dtype=timesteps.dtype) - interleaved_timesteps = torch.stack((timesteps_interpol[:-2, None], timesteps[1:, None]), dim=-1).flatten() - - self.timesteps = torch.cat([timesteps[:1], interleaved_timesteps]) - - self.sample = None - - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma - log_sigma = np.log(np.maximum(sigma, 1e-10)) - - # get distribution - dists = log_sigma - log_sigmas[:, np.newaxis] - - # get sigmas range - low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) - high_idx = low_idx + 1 - - low = log_sigmas[low_idx] - high = log_sigmas[high_idx] - - # interpolate sigmas - w = (low - log_sigma) / (low - high) - w = np.clip(w, 0, 1) - - # transform interpolation to time range - t = (1 - w) * low_idx + w * high_idx - t = t.reshape(sigma.shape) - return 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: - """Constructs the noise schedule of Karras et al. (2022).""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - rho = 7.0 # 7.0 is the value used in the paper - ramp = np.linspace(0, 1, num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Constructs an exponential noise schedule.""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: - """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return sigmas - - @property - def state_in_first_order(self): - return self.sample is None - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep - 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() - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - 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( - self, - model_output: Union[torch.Tensor, np.ndarray], - timestep: Union[float, torch.Tensor], - sample: Union[torch.Tensor, np.ndarray], - generator: Optional[torch.Generator] = None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - timestep (`float`): - 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`): - Whether or not to return a - [`~schedulers.scheduling_k_dpm_2_ancestral_discrete.KDPM2AncestralDiscreteSchedulerOutput`] or tuple. - - Returns: - [`~schedulers.scheduling_k_dpm_2_ancestral_discrete.KDPM2AncestralDiscreteSchedulerOutput`] or `tuple`: - If return_dict is `True`, - [`~schedulers.scheduling_k_dpm_2_ancestral_discrete.KDPM2AncestralDiscreteSchedulerOutput`] is - returned, otherwise a tuple is returned where the first element is the sample tensor. - """ - if self.step_index is None: - self._init_step_index(timestep) - - if self.state_in_first_order: - sigma = self.sigmas[self.step_index] - sigma_interpol = self.sigmas_interpol[self.step_index] - sigma_up = self.sigmas_up[self.step_index] - sigma_down = self.sigmas_down[self.step_index - 1] - else: - # 2nd order / KPDM2's method - sigma = self.sigmas[self.step_index - 1] - sigma_interpol = self.sigmas_interpol[self.step_index - 1] - sigma_up = self.sigmas_up[self.step_index - 1] - sigma_down = self.sigmas_down[self.step_index - 1] - - # currently only gamma=0 is supported. This usually works best anyways. - # We can support gamma in the future but then need to scale the timestep before - # passing it to the model which requires a change in API - gamma = 0 - sigma_hat = sigma * (gamma + 1) # Note: sigma_hat == sigma for now - - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - if self.config.prediction_type == "epsilon": - sigma_input = sigma_hat if self.state_in_first_order else sigma_interpol - pred_original_sample = sample - sigma_input * model_output - elif self.config.prediction_type == "v_prediction": - sigma_input = sigma_hat if self.state_in_first_order else sigma_interpol - pred_original_sample = model_output * (-sigma_input / (sigma_input**2 + 1) ** 0.5) + ( - sample / (sigma_input**2 + 1) - ) - elif self.config.prediction_type == "sample": - raise NotImplementedError("prediction_type not implemented yet: sample") - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`" - ) - - if self.state_in_first_order: - # 2. Convert to an ODE derivative for 1st order - derivative = (sample - pred_original_sample) / sigma_hat - # 3. delta timestep - dt = sigma_interpol - sigma_hat - - # store for 2nd order step - self.sample = sample - self.dt = dt - prev_sample = sample + derivative * dt - else: - # DPM-Solver-2 - # 2. Convert to an ODE derivative for 2nd order - derivative = (sample - pred_original_sample) / sigma_interpol - # 3. delta timestep - dt = sigma_down - sigma_hat - - sample = self.sample - self.sample = None - - prev_sample = sample + derivative * dt - noise = randn_tensor( - model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator - ) - prev_sample = prev_sample + noise * sigma_up - - # upon completion increase step index by one - self._step_index += 1 - - if not return_dict: - return ( - prev_sample, - pred_original_sample, - ) - - return KDPM2AncestralDiscreteSchedulerOutput( - prev_sample=prev_sample, pred_original_sample=pred_original_sample - ) - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.Tensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - noisy_samples = original_samples + noise * sigma - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_k_dpm_2_discrete.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_k_dpm_2_discrete.py deleted file mode 100644 index 9b4cd4e204d6220596a86055156f6c1d36b0f7a8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_k_dpm_2_discrete.py +++ /dev/null @@ -1,589 +0,0 @@ -# Copyright 2025 Katherine Crowson, The HuggingFace Team and hlky. 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 List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput, is_scipy_available -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin - - -if is_scipy_available(): - import scipy.stats - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->KDPM2Discrete -class KDPM2DiscreteSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -class KDPM2DiscreteScheduler(SchedulerMixin, ConfigMixin): - """ - KDPM2DiscreteScheduler is inspired by the DPMSolver2 and Algorithm 2 from the [Elucidating the Design Space of - Diffusion-Based Generative Models](https://huggingface.co/papers/2206.00364) paper. - - 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`, defaults to 1000): - The number of diffusion steps to train the model. - beta_start (`float`, defaults to 0.00085): - The starting `beta` value of inference. - beta_end (`float`, defaults to 0.012): - 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` or `scaled_linear`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - use_karras_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, - the sigmas are determined according to a sequence of noise levels {σi}. - use_exponential_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. - 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. - 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://imagen.research.google/video/paper.pdf) paper). - 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. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 2 - - @register_to_config - def __init__( - self, - 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, - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, - prediction_type: str = "epsilon", - timestep_spacing: str = "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: - raise ValueError( - "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." - ) - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - # set all values - self.set_timesteps(num_train_timesteps, None, num_train_timesteps) - - self._step_index = None - self._begin_index = None - 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 - if self.config.timestep_spacing in ["linspace", "trailing"]: - return self.sigmas.max() - - return (self.sigmas.max() ** 2 + 1) ** 0.5 - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 scale_model_input( - self, - sample: torch.Tensor, - timestep: Union[float, torch.Tensor], - ) -> torch.Tensor: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - if self.step_index is None: - self._init_step_index(timestep) - - if self.state_in_first_order: - sigma = self.sigmas[self.step_index] - else: - sigma = self.sigmas_interpol[self.step_index] - - sample = sample / ((sigma**2 + 1) ** 0.5) - return sample - - def set_timesteps( - self, - num_inference_steps: int, - device: Union[str, torch.device] = None, - num_train_timesteps: Optional[int] = 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 timesteps are not moved. - """ - self.num_inference_steps = num_inference_steps - - num_train_timesteps = num_train_timesteps or self.config.num_train_timesteps - - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = np.linspace(0, num_train_timesteps - 1, num_inference_steps, dtype=np.float32)[::-1].copy() - elif self.config.timestep_spacing == "leading": - step_ratio = num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.float32) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(num_train_timesteps, 0, -step_ratio)).round().copy().astype(np.float32) - timesteps -= 1 - else: - raise ValueError( - 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) - log_sigmas = np.log(sigmas) - sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) - - if self.config.use_karras_sigmas: - sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round() - elif self.config.use_exponential_sigmas: - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - elif self.config.use_beta_sigmas: - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - - self.log_sigmas = torch.from_numpy(log_sigmas).to(device=device) - sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) - sigmas = torch.from_numpy(sigmas).to(device=device) - - # interpolate sigmas - sigmas_interpol = sigmas.log().lerp(sigmas.roll(1).log(), 0.5).exp() - - self.sigmas = torch.cat([sigmas[:1], sigmas[1:].repeat_interleave(2), sigmas[-1:]]) - self.sigmas_interpol = torch.cat( - [sigmas_interpol[:1], sigmas_interpol[1:].repeat_interleave(2), sigmas_interpol[-1:]] - ) - - timesteps = torch.from_numpy(timesteps).to(device) - - # interpolate timesteps - sigmas_interpol = sigmas_interpol.cpu() - log_sigmas = self.log_sigmas.cpu() - timesteps_interpol = np.array( - [self._sigma_to_t(sigma_interpol, log_sigmas) for sigma_interpol in sigmas_interpol] - ) - timesteps_interpol = torch.from_numpy(timesteps_interpol).to(device, dtype=timesteps.dtype) - interleaved_timesteps = torch.stack((timesteps_interpol[1:-1, None], timesteps[1:, None]), dim=-1).flatten() - - self.timesteps = torch.cat([timesteps[:1], interleaved_timesteps]) - - self.sample = None - - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - @property - def state_in_first_order(self): - return self.sample is None - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep - 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() - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - 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 - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma - log_sigma = np.log(np.maximum(sigma, 1e-10)) - - # get distribution - dists = log_sigma - log_sigmas[:, np.newaxis] - - # get sigmas range - low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) - high_idx = low_idx + 1 - - low = log_sigmas[low_idx] - high = log_sigmas[high_idx] - - # interpolate sigmas - w = (low - log_sigma) / (low - high) - w = np.clip(w, 0, 1) - - # transform interpolation to time range - t = (1 - w) * low_idx + w * high_idx - t = t.reshape(sigma.shape) - return 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: - """Constructs the noise schedule of Karras et al. (2022).""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - rho = 7.0 # 7.0 is the value used in the paper - ramp = np.linspace(0, 1, num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Constructs an exponential noise schedule.""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: - """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return sigmas - - def step( - self, - model_output: Union[torch.Tensor, np.ndarray], - timestep: Union[float, torch.Tensor], - sample: Union[torch.Tensor, np.ndarray], - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - timestep (`float`): - The current discrete timestep in the diffusion chain. - 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_k_dpm_2_discrete.KDPM2DiscreteSchedulerOutput`] or - tuple. - - Returns: - [`~schedulers.scheduling_k_dpm_2_discrete.KDPM2DiscreteSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_k_dpm_2_discrete.KDPM2DiscreteSchedulerOutput`] is - returned, otherwise a tuple is returned where the first element is the sample tensor. - """ - if self.step_index is None: - self._init_step_index(timestep) - - if self.state_in_first_order: - sigma = self.sigmas[self.step_index] - sigma_interpol = self.sigmas_interpol[self.step_index + 1] - sigma_next = self.sigmas[self.step_index + 1] - else: - # 2nd order / KDPM2's method - sigma = self.sigmas[self.step_index - 1] - sigma_interpol = self.sigmas_interpol[self.step_index] - sigma_next = self.sigmas[self.step_index] - - # currently only gamma=0 is supported. This usually works best anyways. - # We can support gamma in the future but then need to scale the timestep before - # passing it to the model which requires a change in API - gamma = 0 - sigma_hat = sigma * (gamma + 1) # Note: sigma_hat == sigma for now - - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - if self.config.prediction_type == "epsilon": - sigma_input = sigma_hat if self.state_in_first_order else sigma_interpol - pred_original_sample = sample - sigma_input * model_output - elif self.config.prediction_type == "v_prediction": - sigma_input = sigma_hat if self.state_in_first_order else sigma_interpol - pred_original_sample = model_output * (-sigma_input / (sigma_input**2 + 1) ** 0.5) + ( - sample / (sigma_input**2 + 1) - ) - elif self.config.prediction_type == "sample": - raise NotImplementedError("prediction_type not implemented yet: sample") - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`" - ) - - if self.state_in_first_order: - # 2. Convert to an ODE derivative for 1st order - derivative = (sample - pred_original_sample) / sigma_hat - # 3. delta timestep - dt = sigma_interpol - sigma_hat - - # store for 2nd order step - self.sample = sample - else: - # DPM-Solver-2 - # 2. Convert to an ODE derivative for 2nd order - derivative = (sample - pred_original_sample) / sigma_interpol - - # 3. delta timestep - dt = sigma_next - sigma_hat - - sample = self.sample - self.sample = None - - # upon completion increase step index by one - self._step_index += 1 - - prev_sample = sample + derivative * dt - - if not return_dict: - return ( - prev_sample, - pred_original_sample, - ) - - return KDPM2DiscreteSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.Tensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - noisy_samples = original_samples + noise * sigma - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_karras_ve_flax.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_karras_ve_flax.py deleted file mode 100644 index bacfbd61006d68bb39ce781103cd382ace6d772a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_karras_ve_flax.py +++ /dev/null @@ -1,238 +0,0 @@ -# Copyright 2025 NVIDIA 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 dataclasses import dataclass -from typing import Optional, Tuple, Union - -import flax -import jax -import jax.numpy as jnp -from jax import random - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput -from .scheduling_utils_flax import FlaxSchedulerMixin - - -@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) - - @classmethod - def create(cls): - return cls() - - -@dataclass -class FlaxKarrasVeOutput(BaseOutput): - """ - Output class for the scheduler's step function output. - - Args: - prev_sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - derivative (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)` for images): - Derivative of predicted original image sample (x_0). - state (`KarrasVeSchedulerState`): the `FlaxKarrasVeScheduler` state data class. - """ - - prev_sample: jnp.ndarray - derivative: jnp.ndarray - state: KarrasVeSchedulerState - - -class FlaxKarrasVeScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - Stochastic sampling from Karras et al. [1] tailored to the Variance-Expanding (VE) models [2]. Use Algorithm 2 and - the VE column of Table 1 from [1] for reference. - - [1] Karras, Tero, et al. "Elucidating the Design Space of Diffusion-Based Generative Models." - https://huggingface.co/papers/2206.00364 [2] Song, Yang, et al. "Score-based generative modeling through stochastic - differential equations." https://huggingface.co/papers/2011.13456 - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - For more details on the parameters, see the original paper's Appendix E.: "Elucidating the Design Space of - Diffusion-Based Generative Models." https://huggingface.co/papers/2206.00364. The grid search values used to find - the optimal {s_noise, s_churn, s_min, s_max} for a specific model are described in Table 5 of the paper. - - Args: - sigma_min (`float`): minimum noise magnitude - sigma_max (`float`): maximum noise magnitude - s_noise (`float`): the amount of additional noise to counteract loss of detail during sampling. - A reasonable range is [1.000, 1.011]. - s_churn (`float`): the parameter controlling the overall amount of stochasticity. - A reasonable range is [0, 100]. - s_min (`float`): the start value of the sigma range where we add noise (enable stochasticity). - A reasonable range is [0, 10]. - s_max (`float`): the end value of the sigma range where we add noise. - A reasonable range is [0.2, 80]. - """ - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - sigma_min: float = 0.02, - sigma_max: float = 100, - s_noise: float = 1.007, - s_churn: float = 80, - s_min: float = 0.05, - s_max: float = 50, - ): - pass - - def create_state(self): - return KarrasVeSchedulerState.create() - - def set_timesteps( - 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. - - Args: - state (`KarrasVeSchedulerState`): - the `FlaxKarrasVeScheduler` state data class. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - - """ - timesteps = jnp.arange(0, num_inference_steps)[::-1].copy() - schedule = [ - ( - self.config.sigma_max**2 - * (self.config.sigma_min**2 / self.config.sigma_max**2) ** (i / (num_inference_steps - 1)) - ) - for i in timesteps - ] - - return state.replace( - num_inference_steps=num_inference_steps, - schedule=jnp.array(schedule, dtype=jnp.float32), - timesteps=timesteps, - ) - - def add_noise_to_input( - self, - state: KarrasVeSchedulerState, - sample: jnp.ndarray, - sigma: float, - key: jax.Array, - ) -> 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. - - TODO Args: - """ - if self.config.s_min <= sigma <= self.config.s_max: - gamma = min(self.config.s_churn / state.num_inference_steps, 2**0.5 - 1) - else: - gamma = 0 - - # sample eps ~ N(0, S_noise^2 * I) - key = random.split(key, num=1) - eps = self.config.s_noise * random.normal(key=key, shape=sample.shape) - sigma_hat = sigma + gamma * sigma - sample_hat = sample + ((sigma_hat**2 - sigma**2) ** 0.5 * eps) - - return sample_hat, sigma_hat - - def step( - self, - state: KarrasVeSchedulerState, - model_output: jnp.ndarray, - sigma_hat: float, - sigma_prev: float, - sample_hat: jnp.ndarray, - return_dict: bool = True, - ) -> Union[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). - - Args: - state (`KarrasVeSchedulerState`): the `FlaxKarrasVeScheduler` state data class. - model_output (`torch.Tensor` or `np.ndarray`): direct output from learned diffusion model. - sigma_hat (`float`): TODO - sigma_prev (`float`): TODO - sample_hat (`torch.Tensor` or `np.ndarray`): TODO - return_dict (`bool`): option for returning tuple rather than FlaxKarrasVeOutput class - - Returns: - [`~schedulers.scheduling_karras_ve_flax.FlaxKarrasVeOutput`] or `tuple`: Updated sample in the diffusion - chain and derivative. [`~schedulers.scheduling_karras_ve_flax.FlaxKarrasVeOutput`] if `return_dict` is - True, otherwise a `tuple`. When returning a tuple, the first element is the sample tensor. - """ - - pred_original_sample = sample_hat + sigma_hat * model_output - derivative = (sample_hat - pred_original_sample) / sigma_hat - sample_prev = sample_hat + (sigma_prev - sigma_hat) * derivative - - if not return_dict: - return (sample_prev, derivative, state) - - return FlaxKarrasVeOutput(prev_sample=sample_prev, derivative=derivative, state=state) - - def step_correct( - self, - state: KarrasVeSchedulerState, - model_output: jnp.ndarray, - sigma_hat: float, - sigma_prev: float, - sample_hat: jnp.ndarray, - sample_prev: jnp.ndarray, - derivative: jnp.ndarray, - return_dict: bool = True, - ) -> Union[FlaxKarrasVeOutput, Tuple]: - """ - Correct the predicted sample based on the output model_output of the network. TODO complete description - - Args: - state (`KarrasVeSchedulerState`): the `FlaxKarrasVeScheduler` state data class. - model_output (`torch.Tensor` or `np.ndarray`): direct output from learned diffusion model. - sigma_hat (`float`): TODO - sigma_prev (`float`): TODO - sample_hat (`torch.Tensor` or `np.ndarray`): TODO - sample_prev (`torch.Tensor` or `np.ndarray`): TODO - derivative (`torch.Tensor` or `np.ndarray`): TODO - return_dict (`bool`): option for returning tuple rather than FlaxKarrasVeOutput class - - Returns: - prev_sample (TODO): updated sample in the diffusion chain. derivative (TODO): TODO - - """ - pred_original_sample = sample_prev + sigma_prev * model_output - derivative_corr = (sample_prev - pred_original_sample) / sigma_prev - sample_prev = sample_hat + (sigma_prev - sigma_hat) * (0.5 * derivative + 0.5 * derivative_corr) - - if not return_dict: - return (sample_prev, derivative, state) - - return FlaxKarrasVeOutput(prev_sample=sample_prev, derivative=derivative, state=state) - - def add_noise(self, state: KarrasVeSchedulerState, original_samples, noise, timesteps): - raise NotImplementedError() diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_lcm.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_lcm.py deleted file mode 100644 index cd7a29fe675f98eb7998b60b86da41170db9641f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_lcm.py +++ /dev/null @@ -1,653 +0,0 @@ -# Copyright 2025 Stanford University 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. - -# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion -# and https://github.com/hojonathanho/diffusion - -import math -from dataclasses import dataclass -from typing import List, Optional, Tuple, Union - -import numpy as np -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 LCMSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - denoised: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr -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) - - - Args: - betas (`torch.Tensor`): - the betas that the scheduler is being initialized with. - - Returns: - `torch.Tensor`: rescaled betas with zero terminal SNR - """ - # Convert betas to alphas_bar_sqrt - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - alphas_bar_sqrt = alphas_cumprod.sqrt() - - # Store old values. - alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() - alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() - - # Shift so the last timestep is zero. - alphas_bar_sqrt -= alphas_bar_sqrt_T - - # Scale so the first timestep is back to the old value. - alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - - # Convert alphas_bar_sqrt to betas - alphas_bar = alphas_bar_sqrt**2 # Revert sqrt - alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod - alphas = torch.cat([alphas_bar[0:1], alphas]) - betas = 1 - alphas - - return betas - - -class LCMScheduler(SchedulerMixin, ConfigMixin): - """ - `LCMScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with - non-Markovian guidance. - - This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. [`~ConfigMixin`] takes care of storing all config - attributes that are passed in the scheduler's `__init__` function, such as `num_train_timesteps`. They can be - accessed via `scheduler.config.num_train_timesteps`. [`SchedulerMixin`] provides general loading and saving - functionality via the [`SchedulerMixin.save_pretrained`] and [`~SchedulerMixin.from_pretrained`] functions. - - Args: - 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`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - original_inference_steps (`int`, *optional*, defaults to 50): - The default number of inference steps used to generate a linearly-spaced timestep schedule, from which we - will ultimately take `num_inference_steps` evenly spaced timesteps to form the final timestep schedule. - clip_sample (`bool`, defaults to `True`): - Clip the predicted sample for numerical stability. - clip_sample_range (`float`, defaults to 1.0): - The maximum magnitude for sample clipping. Valid only when `clip_sample=True`. - set_alpha_to_one (`bool`, defaults to `True`): - Each diffusion step uses the alphas product value at that step and at the previous one. For the final step - there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, - otherwise it uses the alpha value at step 0. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - 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://imagen.research.google/video/paper.pdf) paper). - thresholding (`bool`, defaults to `False`): - Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such - as Stable Diffusion. - dynamic_thresholding_ratio (`float`, defaults to 0.995): - The ratio for the dynamic thresholding method. Valid only when `thresholding=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 the [Common Diffusion Noise Schedules and - Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. - timestep_scaling (`float`, defaults to 10.0): - The factor the timesteps will be multiplied by when calculating the consistency model boundary conditions - `c_skip` and `c_out`. Increasing this will decrease the approximation error (although the approximation - error at the default of `10.0` is already pretty small). - 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). - """ - - order = 1 - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.00085, - beta_end: float = 0.012, - beta_schedule: str = "scaled_linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, - original_inference_steps: int = 50, - clip_sample: bool = False, - clip_sample_range: float = 1.0, - set_alpha_to_one: bool = True, - steps_offset: int = 0, - prediction_type: str = "epsilon", - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - sample_max_value: float = 1.0, - timestep_spacing: str = "leading", - timestep_scaling: float = 10.0, - rescale_betas_zero_snr: bool = False, - ): - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - # Rescale for zero SNR - if rescale_betas_zero_snr: - self.betas = rescale_zero_terminal_snr(self.betas) - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - # At every step in ddim, we are looking into the previous alphas_cumprod - # For the final step, there is no previous alphas_cumprod because we are already at 0 - # `set_alpha_to_one` decides whether we set this parameter simply to one or - # whether we use the final alpha of the "non-previous" one. - self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0] - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # setable values - self.num_inference_steps = None - self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)) - self.custom_timesteps = False - - self._step_index = None - self._begin_index = None - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep - 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() - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - 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 - - @property - def step_index(self): - 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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: - """ - 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 current timestep in the diffusion chain. - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - """ - "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the - prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by - s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing - pixels from saturation at each step. We find that dynamic thresholding results in significantly better - photorealism as well as better image-text alignment, especially when using very large guidance weights." - - https://huggingface.co/papers/2205.11487 - """ - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half - - # Flatten sample for doing quantile calculation along each image - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - - abs_sample = sample.abs() # "a certain percentile absolute pixel value" - - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] - s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 - sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" - - sample = sample.reshape(batch_size, channels, *remaining_dims) - sample = sample.to(dtype) - - return sample - - 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, - strength: int = 1.0, - ): - """ - Sets the discrete timesteps used for the diffusion chain (to be run before inference). - - Args: - 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. - original_inference_steps (`int`, *optional*): - The original number of inference steps, which will be used to generate a linearly-spaced timestep - 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*): - 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`. - """ - # 0. Check inputs - if num_inference_steps is None and timesteps is None: - raise ValueError("Must pass exactly one of `num_inference_steps` or `custom_timesteps`.") - - if num_inference_steps is not None and timesteps is not None: - raise ValueError("Can only pass one of `num_inference_steps` or `custom_timesteps`.") - - # 1. Calculate the LCM original training/distillation timestep schedule. - original_steps = ( - original_inference_steps if original_inference_steps is not None else self.config.original_inference_steps - ) - - if original_steps > self.config.num_train_timesteps: - raise ValueError( - f"`original_steps`: {original_steps} cannot be larger than `self.config.train_timesteps`:" - f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle" - f" maximal {self.config.num_train_timesteps} timesteps." - ) - - # LCM Timesteps Setting - # The skipping step parameter k from the paper. - k = self.config.num_train_timesteps // original_steps - # LCM Training/Distillation Steps Schedule - # Currently, only a linearly-spaced schedule is supported (same as in the LCM distillation scripts). - lcm_origin_timesteps = np.asarray(list(range(1, int(original_steps * strength) + 1))) * k - 1 - - # 2. Calculate the LCM inference timestep schedule. - if timesteps is not None: - # 2.1 Handle custom timestep schedules. - train_timesteps = set(lcm_origin_timesteps) - non_train_timesteps = [] - for i in range(1, len(timesteps)): - if timesteps[i] >= timesteps[i - 1]: - raise ValueError("`custom_timesteps` must be in descending order.") - - if timesteps[i] not in train_timesteps: - non_train_timesteps.append(timesteps[i]) - - if timesteps[0] >= self.config.num_train_timesteps: - raise ValueError( - f"`timesteps` must start before `self.config.train_timesteps`: {self.config.num_train_timesteps}." - ) - - # Raise warning if timestep schedule does not start with self.config.num_train_timesteps - 1 - if strength == 1.0 and timesteps[0] != self.config.num_train_timesteps - 1: - logger.warning( - f"The first timestep on the custom timestep schedule is {timesteps[0]}, not" - f" `self.config.num_train_timesteps - 1`: {self.config.num_train_timesteps - 1}. You may get" - f" unexpected results when using this timestep schedule." - ) - - # Raise warning if custom timestep schedule contains timesteps not on original timestep schedule - if non_train_timesteps: - logger.warning( - f"The custom timestep schedule contains the following timesteps which are not on the original" - f" training/distillation timestep schedule: {non_train_timesteps}. You may get unexpected results" - f" when using this timestep schedule." - ) - - # Raise warning if custom timestep schedule is longer than original_steps - if len(timesteps) > original_steps: - logger.warning( - f"The number of timesteps in the custom timestep schedule is {len(timesteps)}, which exceeds the" - f" the length of the timestep schedule used for training: {original_steps}. You may get some" - f" unexpected results when using this timestep schedule." - ) - - timesteps = np.array(timesteps, dtype=np.int64) - self.num_inference_steps = len(timesteps) - self.custom_timesteps = True - - # Apply strength (e.g. for img2img pipelines) (see StableDiffusionImg2ImgPipeline.get_timesteps) - init_timestep = min(int(self.num_inference_steps * strength), self.num_inference_steps) - t_start = max(self.num_inference_steps - init_timestep, 0) - timesteps = timesteps[t_start * self.order :] - # TODO: also reset self.num_inference_steps? - else: - # 2.2 Create the "standard" LCM inference timestep schedule. - if num_inference_steps > self.config.num_train_timesteps: - raise ValueError( - f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:" - f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle" - f" maximal {self.config.num_train_timesteps} timesteps." - ) - - skipping_step = len(lcm_origin_timesteps) // num_inference_steps - - if skipping_step < 1: - raise ValueError( - f"The combination of `original_steps x strength`: {original_steps} x {strength} is smaller than `num_inference_steps`: {num_inference_steps}. Make sure to either reduce `num_inference_steps` to a value smaller than {int(original_steps * strength)} or increase `strength` to a value higher than {float(num_inference_steps / original_steps)}." - ) - - self.num_inference_steps = num_inference_steps - - if num_inference_steps > original_steps: - raise ValueError( - f"`num_inference_steps`: {num_inference_steps} cannot be larger than `original_inference_steps`:" - f" {original_steps} because the final timestep schedule will be a subset of the" - f" `original_inference_steps`-sized initial timestep schedule." - ) - - # LCM Inference Steps Schedule - lcm_origin_timesteps = lcm_origin_timesteps[::-1].copy() - # Select (approximately) evenly spaced indices from lcm_origin_timesteps. - inference_indices = np.linspace(0, len(lcm_origin_timesteps), num=num_inference_steps, endpoint=False) - inference_indices = np.floor(inference_indices).astype(np.int64) - timesteps = lcm_origin_timesteps[inference_indices] - - self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.long) - - self._step_index = None - self._begin_index = None - - def get_scalings_for_boundary_condition_discrete(self, timestep): - self.sigma_data = 0.5 # Default: 0.5 - scaled_timestep = timestep * self.config.timestep_scaling - - c_skip = self.sigma_data**2 / (scaled_timestep**2 + self.sigma_data**2) - c_out = scaled_timestep / (scaled_timestep**2 + self.sigma_data**2) ** 0.5 - return c_skip, c_out - - def step( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - generator: Optional[torch.Generator] = None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - timestep (`float`): - 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`, *optional*, defaults to `True`): - Whether or not to return a [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`. - Returns: - [`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a - tuple is returned where the first element is the sample tensor. - """ - 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 self.step_index is None: - self._init_step_index(timestep) - - # 1. get previous step value - prev_step_index = self.step_index + 1 - if prev_step_index < len(self.timesteps): - prev_timestep = self.timesteps[prev_step_index] - else: - prev_timestep = timestep - - # 2. compute alphas, betas - 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 - beta_prod_t_prev = 1 - alpha_prod_t_prev - - # 3. Get scalings for boundary conditions - c_skip, c_out = self.get_scalings_for_boundary_condition_discrete(timestep) - - # 4. Compute the predicted original sample x_0 based on the model parameterization - if self.config.prediction_type == "epsilon": # noise-prediction - predicted_original_sample = (sample - beta_prod_t.sqrt() * model_output) / alpha_prod_t.sqrt() - elif self.config.prediction_type == "sample": # x-prediction - predicted_original_sample = model_output - elif self.config.prediction_type == "v_prediction": # v-prediction - predicted_original_sample = alpha_prod_t.sqrt() * sample - beta_prod_t.sqrt() * model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or" - " `v_prediction` for `LCMScheduler`." - ) - - # 5. Clip or threshold "predicted x_0" - if self.config.thresholding: - predicted_original_sample = self._threshold_sample(predicted_original_sample) - elif self.config.clip_sample: - predicted_original_sample = predicted_original_sample.clamp( - -self.config.clip_sample_range, self.config.clip_sample_range - ) - - # 6. Denoise model output using boundary conditions - denoised = c_out * predicted_original_sample + c_skip * sample - - # 7. Sample and inject noise z ~ N(0, I) for MultiStep Inference - # Noise is not used on the final timestep of the timestep schedule. - # This also means that noise is not used for one-step sampling. - if self.step_index != self.num_inference_steps - 1: - noise = randn_tensor( - model_output.shape, generator=generator, device=model_output.device, dtype=denoised.dtype - ) - prev_sample = alpha_prod_t_prev.sqrt() * denoised + beta_prod_t_prev.sqrt() * noise - else: - prev_sample = denoised - - # upon completion increase step index by one - self._step_index += 1 - - if not return_dict: - return (prev_sample, denoised) - - return LCMSchedulerOutput(prev_sample=prev_sample, denoised=denoised) - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as original_samples - # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement - # for the subsequent add_noise calls - self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype) - timesteps = timesteps.to(original_samples.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(original_samples.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise - return noisy_samples - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity - def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as sample - self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype) - timesteps = timesteps.to(sample.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(sample.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample - return velocity - - def __len__(self): - return self.config.num_train_timesteps - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.previous_timestep - def previous_timestep(self, timestep): - if self.custom_timesteps or self.num_inference_steps: - index = (self.timesteps == timestep).nonzero(as_tuple=True)[0][0] - if index == self.timesteps.shape[0] - 1: - prev_t = torch.tensor(-1) - else: - prev_t = self.timesteps[index + 1] - else: - prev_t = timestep - 1 - return prev_t diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_lms_discrete.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_lms_discrete.py deleted file mode 100644 index c2450204aa8f94719a8eeb8b43ce9f62f37583cf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_lms_discrete.py +++ /dev/null @@ -1,552 +0,0 @@ -# Copyright 2025 Katherine Crowson 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 warnings -from dataclasses import dataclass -from typing import List, Optional, Tuple, Union - -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 .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->LMSDiscrete -class LMSDiscreteSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -class LMSDiscreteScheduler(SchedulerMixin, ConfigMixin): - """ - A linear multistep scheduler for discrete beta schedules. - - 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`, 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` or `scaled_linear`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - use_karras_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, - the sigmas are determined according to a sequence of noise levels {σi}. - use_exponential_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. - 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. - 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://imagen.research.google/video/paper.pdf) paper). - 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. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - 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, - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, - prediction_type: str = "epsilon", - timestep_spacing: str = "linspace", - steps_offset: int = 0, - ): - 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 trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) - sigmas = np.concatenate([sigmas[::-1], [0.0]]).astype(np.float32) - self.sigmas = torch.from_numpy(sigmas) - - # setable values - self.num_inference_steps = None - self.use_karras_sigmas = use_karras_sigmas - self.set_timesteps(num_train_timesteps, None) - self.derivatives = [] - self.is_scale_input_called = False - - self._step_index = None - self._begin_index = None - 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 - if self.config.timestep_spacing in ["linspace", "trailing"]: - return self.sigmas.max() - - return (self.sigmas.max() ** 2 + 1) ** 0.5 - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: - """ - 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 (`float` or `torch.Tensor`): - The current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - - if self.step_index is None: - self._init_step_index(timestep) - - sigma = self.sigmas[self.step_index] - sample = sample / ((sigma**2 + 1) ** 0.5) - self.is_scale_input_called = True - return sample - - def get_lms_coefficient(self, order, t, current_order): - """ - Compute the linear multistep coefficient. - - Args: - order (): - t (): - current_order (): - """ - - def lms_derivative(tau): - prod = 1.0 - for k in range(order): - if current_order == k: - continue - prod *= (tau - self.sigmas[t - k]) / (self.sigmas[t - current_order] - self.sigmas[t - k]) - return prod - - integrated_coeff = integrate.quad(lms_derivative, self.sigmas[t], self.sigmas[t + 1], epsrel=1e-4)[0] - - return integrated_coeff - - def set_timesteps(self, num_inference_steps: int, device: Union[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 (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - """ - self.num_inference_steps = num_inference_steps - - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps, dtype=np.float32)[ - ::-1 - ].copy() - elif self.config.timestep_spacing == "leading": - step_ratio = self.config.num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.float32) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(self.config.num_train_timesteps, 0, -step_ratio)).round().copy().astype(np.float32) - timesteps -= 1 - else: - raise ValueError( - 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) - log_sigmas = np.log(sigmas) - sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) - - if self.config.use_karras_sigmas: - sigmas = self._convert_to_karras(in_sigmas=sigmas) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - elif self.config.use_exponential_sigmas: - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - elif self.config.use_beta_sigmas: - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - - sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) - - self.sigmas = torch.from_numpy(sigmas).to(device=device) - self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.float32) - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - self.derivatives = [] - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep - 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() - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - 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 - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma - log_sigma = np.log(np.maximum(sigma, 1e-10)) - - # get distribution - dists = log_sigma - log_sigmas[:, np.newaxis] - - # get sigmas range - low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) - high_idx = low_idx + 1 - - low = log_sigmas[low_idx] - high = log_sigmas[high_idx] - - # interpolate sigmas - w = (low - log_sigma) / (low - high) - w = np.clip(w, 0, 1) - - # transform interpolation to time range - t = (1 - w) * low_idx + w * high_idx - t = t.reshape(sigma.shape) - return t - - # copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras - def _convert_to_karras(self, in_sigmas: torch.Tensor) -> torch.Tensor: - """Constructs the noise schedule of Karras et al. (2022).""" - - sigma_min: float = in_sigmas[-1].item() - sigma_max: float = in_sigmas[0].item() - - rho = 7.0 # 7.0 is the value used in the paper - ramp = np.linspace(0, 1, self.num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Constructs an exponential noise schedule.""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: - """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return sigmas - - def step( - self, - model_output: torch.Tensor, - timestep: Union[float, torch.Tensor], - sample: torch.Tensor, - order: int = 4, - return_dict: bool = True, - ) -> Union[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). - - Args: - 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`): - A current instance of a sample created by the diffusion process. - order (`int`, defaults to 4): - The order of the linear multistep method. - return_dict (`bool`, *optional*, 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 not self.is_scale_input_called: - warnings.warn( - "The `scale_model_input` function should be called before `step` to ensure correct denoising. " - "See `StableDiffusionPipeline` for a usage example." - ) - - if self.step_index is None: - self._init_step_index(timestep) - - sigma = self.sigmas[self.step_index] - - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - if self.config.prediction_type == "epsilon": - pred_original_sample = sample - sigma * model_output - elif self.config.prediction_type == "v_prediction": - # * c_out + input * c_skip - pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1)) - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`" - ) - - # 2. Convert to an ODE derivative - derivative = (sample - pred_original_sample) / sigma - self.derivatives.append(derivative) - if len(self.derivatives) > order: - self.derivatives.pop(0) - - # 3. Compute linear multistep coefficients - order = min(self.step_index + 1, order) - lms_coeffs = [self.get_lms_coefficient(order, self.step_index, curr_order) for curr_order in range(order)] - - # 4. Compute previous sample based on the derivatives path - prev_sample = sample + sum( - coeff * derivative for coeff, derivative in zip(lms_coeffs, reversed(self.derivatives)) - ) - - # upon completion increase step index by one - self._step_index += 1 - - if not return_dict: - return ( - prev_sample, - pred_original_sample, - ) - - return LMSDiscreteSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.Tensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - noisy_samples = original_samples + noise * sigma - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_lms_discrete_flax.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_lms_discrete_flax.py deleted file mode 100644 index b8e08ff9e134f8a77b428179057df589d736a187..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_lms_discrete_flax.py +++ /dev/null @@ -1,283 +0,0 @@ -# Copyright 2025 Katherine Crowson 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 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 .scheduling_utils_flax import ( - CommonSchedulerState, - FlaxKarrasDiffusionSchedulers, - FlaxSchedulerMixin, - FlaxSchedulerOutput, - broadcast_to_shape_from_left, -) - - -@flax.struct.dataclass -class LMSDiscreteSchedulerState: - common: CommonSchedulerState - - # setable values - init_noise_sigma: jnp.ndarray - timesteps: jnp.ndarray - sigmas: jnp.ndarray - num_inference_steps: Optional[int] = None - - # running values - derivatives: Optional[jnp.ndarray] = None - - @classmethod - def create( - 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) - - -@dataclass -class FlaxLMSSchedulerOutput(FlaxSchedulerOutput): - state: LMSDiscreteSchedulerState - - -class FlaxLMSDiscreteScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - Linear Multistep Scheduler for discrete beta schedules. Based on the original k-diffusion implementation by - Katherine Crowson: - https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L181 - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - 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 - `linear` or `scaled_linear`. - trained_betas (`jnp.ndarray`, optional): - option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. - prediction_type (`str`, default `epsilon`, optional): - 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://imagen.research.google/video/paper.pdf) - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - the `dtype` used for params and computation. - """ - - _compatibles = [e.name for e in FlaxKarrasDiffusionSchedulers] - - dtype: jnp.dtype - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: Optional[jnp.ndarray] = None, - prediction_type: str = "epsilon", - dtype: jnp.dtype = jnp.float32, - ): - self.dtype = dtype - - def create_state(self, common: Optional[CommonSchedulerState] = None) -> LMSDiscreteSchedulerState: - if common is None: - common = CommonSchedulerState.create(self) - - timesteps = jnp.arange(0, self.config.num_train_timesteps).round()[::-1] - sigmas = ((1 - common.alphas_cumprod) / common.alphas_cumprod) ** 0.5 - - # standard deviation of the initial noise distribution - init_noise_sigma = sigmas.max() - - return LMSDiscreteSchedulerState.create( - common=common, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - sigmas=sigmas, - ) - - def scale_model_input(self, state: LMSDiscreteSchedulerState, sample: jnp.ndarray, timestep: int) -> jnp.ndarray: - """ - Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the K-LMS algorithm. - - Args: - state (`LMSDiscreteSchedulerState`): - the `FlaxLMSDiscreteScheduler` state data class instance. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - timestep (`int`): - current discrete timestep in the diffusion chain. - - Returns: - `jnp.ndarray`: scaled input sample - """ - (step_index,) = jnp.where(state.timesteps == timestep, size=1) - step_index = step_index[0] - - sigma = state.sigmas[step_index] - sample = sample / ((sigma**2 + 1) ** 0.5) - return sample - - def get_lms_coefficient(self, state: LMSDiscreteSchedulerState, order, t, current_order): - """ - Compute a linear multistep coefficient. - - Args: - order (TODO): - t (TODO): - current_order (TODO): - """ - - def lms_derivative(tau): - prod = 1.0 - for k in range(order): - if current_order == k: - continue - prod *= (tau - state.sigmas[t - k]) / (state.sigmas[t - current_order] - state.sigmas[t - k]) - return prod - - integrated_coeff = integrate.quad(lms_derivative, state.sigmas[t], state.sigmas[t + 1], epsrel=1e-4)[0] - - return integrated_coeff - - def set_timesteps( - 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. - - Args: - state (`LMSDiscreteSchedulerState`): - the `FlaxLMSDiscreteScheduler` state data class instance. - num_inference_steps (`int`): - 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) - - low_idx = jnp.floor(timesteps).astype(jnp.int32) - high_idx = jnp.ceil(timesteps).astype(jnp.int32) - - frac = jnp.mod(timesteps, 1.0) - - sigmas = ((1 - state.common.alphas_cumprod) / state.common.alphas_cumprod) ** 0.5 - sigmas = (1 - frac) * sigmas[low_idx] + frac * sigmas[high_idx] - sigmas = jnp.concatenate([sigmas, jnp.array([0.0], dtype=self.dtype)]) - - timesteps = timesteps.astype(jnp.int32) - - # initial running values - derivatives = jnp.zeros((0,) + shape, dtype=self.dtype) - - return state.replace( - timesteps=timesteps, - sigmas=sigmas, - num_inference_steps=num_inference_steps, - derivatives=derivatives, - ) - - def step( - self, - state: LMSDiscreteSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - order: int = 4, - return_dict: bool = True, - ) -> Union[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). - - Args: - state (`LMSDiscreteSchedulerState`): the `FlaxLMSDiscreteScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - order: coefficient for multi-step inference. - return_dict (`bool`): option for returning tuple rather than FlaxLMSSchedulerOutput class - - Returns: - [`FlaxLMSSchedulerOutput`] or `tuple`: [`FlaxLMSSchedulerOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - if state.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - sigma = state.sigmas[timestep] - - # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise - if self.config.prediction_type == "epsilon": - pred_original_sample = sample - sigma * model_output - elif self.config.prediction_type == "v_prediction": - # * c_out + input * c_skip - pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1)) - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`" - ) - - # 2. Convert to an ODE derivative - derivative = (sample - pred_original_sample) / sigma - state = state.replace(derivatives=jnp.append(state.derivatives, derivative)) - if len(state.derivatives) > order: - state = state.replace(derivatives=jnp.delete(state.derivatives, 0)) - - # 3. Compute linear multistep coefficients - order = min(timestep + 1, order) - lms_coeffs = [self.get_lms_coefficient(state, order, timestep, curr_order) for curr_order in range(order)] - - # 4. Compute previous sample based on the derivatives path - prev_sample = sample + sum( - coeff * derivative for coeff, derivative in zip(lms_coeffs, reversed(state.derivatives)) - ) - - if not return_dict: - return (prev_sample, state) - - return FlaxLMSSchedulerOutput(prev_sample=prev_sample, state=state) - - def add_noise( - self, - state: LMSDiscreteSchedulerState, - original_samples: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - sigma = state.sigmas[timesteps].flatten() - sigma = broadcast_to_shape_from_left(sigma, noise.shape) - - noisy_samples = original_samples + noise * sigma - - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_pndm.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_pndm.py deleted file mode 100644 index c07621179e2b7321f346a42def88cde7cad0bbe8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_pndm.py +++ /dev/null @@ -1,476 +0,0 @@ -# Copyright 2025 Zhejiang University 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. - -# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim - -import math -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -class PNDMScheduler(SchedulerMixin, ConfigMixin): - """ - `PNDMScheduler` uses pseudo numerical methods for diffusion models such as the Runge-Kutta and linear multi-step - method. - - 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`, 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`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - skip_prk_steps (`bool`, defaults to `False`): - Allows the scheduler to skip the Runge-Kutta steps defined in the original paper as being required before - PLMS steps. - set_alpha_to_one (`bool`, defaults to `False`): - Each diffusion step uses the alphas product value at that step and at the previous one. For the final step - there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, - otherwise it uses the alpha value at step 0. - prediction_type (`str`, defaults to `epsilon`, *optional*): - Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process) - or `v_prediction` (see section 2.4 of [Imagen Video](https://imagen.research.google/video/paper.pdf) - paper). - 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. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - 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, - skip_prk_steps: bool = False, - set_alpha_to_one: bool = False, - prediction_type: str = "epsilon", - timestep_spacing: str = "leading", - steps_offset: int = 0, - ): - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0] - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # For now we only support F-PNDM, i.e. the runge-kutta method - # For more information on the algorithm please take a look at the paper: https://huggingface.co/papers/2202.09778 - # mainly at formula (9), (12), (13) and the Algorithm 2. - self.pndm_order = 4 - - # running values - self.cur_model_output = 0 - self.counter = 0 - self.cur_sample = None - self.ets = [] - - # setable values - self.num_inference_steps = None - self._timesteps = np.arange(0, num_train_timesteps)[::-1].copy() - self.prk_timesteps = None - self.plms_timesteps = None - self.timesteps = None - - def set_timesteps(self, num_inference_steps: int, device: Union[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 (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - """ - - self.num_inference_steps = num_inference_steps - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - self._timesteps = ( - np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps).round().astype(np.int64) - ) - elif self.config.timestep_spacing == "leading": - step_ratio = self.config.num_train_timesteps // self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - self._timesteps = (np.arange(0, num_inference_steps) * step_ratio).round() - self._timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / self.num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - self._timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio))[::-1].astype( - np.int64 - ) - self._timesteps -= 1 - else: - raise ValueError( - f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'." - ) - - if self.config.skip_prk_steps: - # for some models like stable diffusion the prk steps can/should be skipped to - # produce better results. When using PNDM with `self.config.skip_prk_steps` the implementation - # is based on crowsonkb's PLMS sampler implementation: https://github.com/CompVis/latent-diffusion/pull/51 - self.prk_timesteps = np.array([]) - self.plms_timesteps = np.concatenate([self._timesteps[:-1], self._timesteps[-2:-1], self._timesteps[-1:]])[ - ::-1 - ].copy() - else: - prk_timesteps = np.array(self._timesteps[-self.pndm_order :]).repeat(2) + np.tile( - np.array([0, self.config.num_train_timesteps // num_inference_steps // 2]), self.pndm_order - ) - self.prk_timesteps = (prk_timesteps[:-1].repeat(2)[1:-1])[::-1].copy() - self.plms_timesteps = self._timesteps[:-3][ - ::-1 - ].copy() # we copy to avoid having negative strides which are not supported by torch.from_numpy - - timesteps = np.concatenate([self.prk_timesteps, self.plms_timesteps]).astype(np.int64) - self.timesteps = torch.from_numpy(timesteps).to(device) - - self.ets = [] - self.counter = 0 - self.cur_model_output = 0 - - def step( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - return_dict: bool = True, - ) -> Union[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`] - or [`~PNDMScheduler.step_plms`] depending on the internal variable `counter`. - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - return_dict (`bool`): - 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.counter < len(self.prk_timesteps) and not self.config.skip_prk_steps: - return self.step_prk(model_output=model_output, timestep=timestep, sample=sample, return_dict=return_dict) - else: - return self.step_plms(model_output=model_output, timestep=timestep, sample=sample, return_dict=return_dict) - - def step_prk( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - return_dict: bool = True, - ) -> Union[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 - equation. - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - return_dict (`bool`): - 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( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - diff_to_prev = 0 if self.counter % 2 else self.config.num_train_timesteps // self.num_inference_steps // 2 - prev_timestep = timestep - diff_to_prev - timestep = self.prk_timesteps[self.counter // 4 * 4] - - if self.counter % 4 == 0: - self.cur_model_output += 1 / 6 * model_output - self.ets.append(model_output) - self.cur_sample = sample - elif (self.counter - 1) % 4 == 0: - self.cur_model_output += 1 / 3 * model_output - elif (self.counter - 2) % 4 == 0: - self.cur_model_output += 1 / 3 * model_output - elif (self.counter - 3) % 4 == 0: - model_output = self.cur_model_output + 1 / 6 * model_output - self.cur_model_output = 0 - - # cur_sample should not be `None` - cur_sample = self.cur_sample if self.cur_sample is not None else sample - - prev_sample = self._get_prev_sample(cur_sample, timestep, prev_timestep, model_output) - self.counter += 1 - - if not return_dict: - return (prev_sample,) - - return SchedulerOutput(prev_sample=prev_sample) - - def step_plms( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - return_dict: bool = True, - ) -> Union[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. - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - return_dict (`bool`): - 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( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - if not self.config.skip_prk_steps and len(self.ets) < 3: - raise ValueError( - f"{self.__class__} can only be run AFTER scheduler has been run " - "in 'prk' mode for at least 12 iterations " - "See: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pipeline_pndm.py " - "for more information." - ) - - prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps - - if self.counter != 1: - self.ets = self.ets[-3:] - self.ets.append(model_output) - else: - prev_timestep = timestep - timestep = timestep + self.config.num_train_timesteps // self.num_inference_steps - - if len(self.ets) == 1 and self.counter == 0: - model_output = model_output - self.cur_sample = sample - elif len(self.ets) == 1 and self.counter == 1: - model_output = (model_output + self.ets[-1]) / 2 - sample = self.cur_sample - self.cur_sample = None - elif len(self.ets) == 2: - model_output = (3 * self.ets[-1] - self.ets[-2]) / 2 - elif len(self.ets) == 3: - model_output = (23 * self.ets[-1] - 16 * self.ets[-2] + 5 * self.ets[-3]) / 12 - else: - model_output = (1 / 24) * (55 * self.ets[-1] - 59 * self.ets[-2] + 37 * self.ets[-3] - 9 * self.ets[-4]) - - prev_sample = self._get_prev_sample(sample, timestep, prev_timestep, model_output) - self.counter += 1 - - if not return_dict: - return (prev_sample,) - - return SchedulerOutput(prev_sample=prev_sample) - - def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: - """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. - - Args: - sample (`torch.Tensor`): - The input sample. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - def _get_prev_sample(self, sample, timestep, prev_timestep, model_output): - # See formula (9) of PNDM paper https://huggingface.co/papers/2202.09778 - # this function computes x_(t−δ) using the formula of (9) - # Note that x_t needs to be added to both sides of the equation - - # Notation ( -> - # alpha_prod_t -> α_t - # alpha_prod_t_prev -> α_(t−δ) - # beta_prod_t -> (1 - α_t) - # beta_prod_t_prev -> (1 - α_(t−δ)) - # sample -> x_t - # model_output -> e_θ(x_t, t) - # prev_sample -> x_(t−δ) - 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 - beta_prod_t_prev = 1 - alpha_prod_t_prev - - if self.config.prediction_type == "v_prediction": - model_output = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - elif self.config.prediction_type != "epsilon": - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon` or `v_prediction`" - ) - - # corresponds to (α_(t−δ) - α_t) divided by - # denominator of x_t in formula (9) and plus 1 - # Note: (α_(t−δ) - α_t) / (sqrt(α_t) * (sqrt(α_(t−δ)) + sqr(α_t))) = - # sqrt(α_(t−δ)) / sqrt(α_t)) - sample_coeff = (alpha_prod_t_prev / alpha_prod_t) ** (0.5) - - # corresponds to denominator of e_θ(x_t, t) in formula (9) - model_output_denom_coeff = alpha_prod_t * beta_prod_t_prev ** (0.5) + ( - alpha_prod_t * beta_prod_t * alpha_prod_t_prev - ) ** (0.5) - - # full formula (9) - prev_sample = ( - sample_coeff * sample - (alpha_prod_t_prev - alpha_prod_t) * model_output / model_output_denom_coeff - ) - - return prev_sample - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as original_samples - # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement - # for the subsequent add_noise calls - self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype) - timesteps = timesteps.to(original_samples.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(original_samples.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_pndm_flax.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_pndm_flax.py deleted file mode 100644 index 12e22005afaf700f8d145433a9c35cc029500cae..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_pndm_flax.py +++ /dev/null @@ -1,509 +0,0 @@ -# Copyright 2025 Zhejiang University 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. - -# 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 .scheduling_utils_flax import ( - CommonSchedulerState, - FlaxKarrasDiffusionSchedulers, - FlaxSchedulerMixin, - FlaxSchedulerOutput, - add_noise_common, -) - - -@flax.struct.dataclass -class PNDMSchedulerState: - common: CommonSchedulerState - final_alpha_cumprod: jnp.ndarray - - # 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 - - # 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 - - @classmethod - def create( - cls, - common: CommonSchedulerState, - final_alpha_cumprod: jnp.ndarray, - init_noise_sigma: jnp.ndarray, - timesteps: jnp.ndarray, - ): - return cls( - common=common, - final_alpha_cumprod=final_alpha_cumprod, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - ) - - -@dataclass -class FlaxPNDMSchedulerOutput(FlaxSchedulerOutput): - state: PNDMSchedulerState - - -class FlaxPNDMScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - Pseudo numerical methods for diffusion models (PNDM) proposes using more advanced ODE integration techniques, - namely Runge-Kutta method and a linear multi-step method. - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - For more details, see the original paper: https://huggingface.co/papers/2202.09778 - - 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 - `linear`, `scaled_linear`, or `squaredcos_cap_v2`. - trained_betas (`jnp.ndarray`, optional): - option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. - skip_prk_steps (`bool`): - allows the scheduler to skip the Runge-Kutta steps that are defined in the original paper as being required - before plms steps; defaults to `False`. - set_alpha_to_one (`bool`, default `False`): - each diffusion step uses the value of alphas product at that step and at the previous one. For the final - step there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, - otherwise it uses the value of alpha at step 0. - steps_offset (`int`, default `0`): - An offset added to the inference steps, as required by some model families. - prediction_type (`str`, default `epsilon`, optional): - 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://imagen.research.google/video/paper.pdf) - dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`): - the `dtype` used for params and computation. - """ - - _compatibles = [e.name for e in FlaxKarrasDiffusionSchedulers] - - dtype: jnp.dtype - pndm_order: int - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: Optional[jnp.ndarray] = 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, - ): - self.dtype = dtype - - # For now we only support F-PNDM, i.e. the runge-kutta method - # For more information on the algorithm please take a look at the paper: https://huggingface.co/papers/2202.09778 - # mainly at formula (9), (12), (13) and the Algorithm 2. - self.pndm_order = 4 - - def create_state(self, common: Optional[CommonSchedulerState] = None) -> PNDMSchedulerState: - if common is None: - common = CommonSchedulerState.create(self) - - # At every step in ddim, we are looking into the previous alphas_cumprod - # For the final step, there is no previous alphas_cumprod because we are already at 0 - # `set_alpha_to_one` decides whether we set this parameter simply to one or - # whether we use the final alpha of the "non-previous" one. - final_alpha_cumprod = ( - jnp.array(1.0, dtype=self.dtype) if self.config.set_alpha_to_one else common.alphas_cumprod[0] - ) - - # standard deviation of the initial noise distribution - init_noise_sigma = jnp.array(1.0, dtype=self.dtype) - - timesteps = jnp.arange(0, self.config.num_train_timesteps).round()[::-1] - - return PNDMSchedulerState.create( - common=common, - final_alpha_cumprod=final_alpha_cumprod, - init_noise_sigma=init_noise_sigma, - timesteps=timesteps, - ) - - 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. - - Args: - state (`PNDMSchedulerState`): - 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`): - the shape of the samples to be generated. - """ - - step_ratio = self.config.num_train_timesteps // num_inference_steps - # creates integer timesteps by multiplying by ratio - # rounding to avoid issues when num_inference_step is power of 3 - _timesteps = (jnp.arange(0, num_inference_steps) * step_ratio).round() + self.config.steps_offset - - if self.config.skip_prk_steps: - # for some models like stable diffusion the prk steps can/should be skipped to - # produce better results. When using PNDM with `self.config.skip_prk_steps` the implementation - # is based on crowsonkb's PLMS sampler implementation: https://github.com/CompVis/latent-diffusion/pull/51 - - prk_timesteps = jnp.array([], dtype=jnp.int32) - plms_timesteps = jnp.concatenate([_timesteps[:-1], _timesteps[-2:-1], _timesteps[-1:]])[::-1] - - 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), - self.pndm_order, - ) - - prk_timesteps = (prk_timesteps[:-1].repeat(2)[1:-1])[::-1] - plms_timesteps = _timesteps[:-3][::-1] - - timesteps = jnp.concatenate([prk_timesteps, plms_timesteps]) - - # initial running values - - cur_model_output = jnp.zeros(shape, dtype=self.dtype) - counter = jnp.int32(0) - cur_sample = jnp.zeros(shape, dtype=self.dtype) - ets = jnp.zeros((4,) + shape, dtype=self.dtype) - - return state.replace( - timesteps=timesteps, - num_inference_steps=num_inference_steps, - prk_timesteps=prk_timesteps, - plms_timesteps=plms_timesteps, - cur_model_output=cur_model_output, - counter=counter, - cur_sample=cur_sample, - ets=ets, - ) - - def scale_model_input( - self, state: PNDMSchedulerState, sample: jnp.ndarray, timestep: Optional[int] = None - ) -> jnp.ndarray: - """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. - - Args: - state (`PNDMSchedulerState`): the `FlaxPNDMScheduler` state data class instance. - sample (`jnp.ndarray`): input sample - timestep (`int`, optional): current timestep - - Returns: - `jnp.ndarray`: scaled input sample - """ - return sample - - def step( - self, - state: PNDMSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - return_dict: bool = True, - ) -> Union[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). - - This function calls `step_prk()` or `step_plms()` depending on the internal variable `counter`. - - Args: - state (`PNDMSchedulerState`): the `FlaxPNDMScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - return_dict (`bool`): option for returning tuple rather than FlaxPNDMSchedulerOutput class - - Returns: - [`FlaxPNDMSchedulerOutput`] or `tuple`: [`FlaxPNDMSchedulerOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - - if state.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - if self.config.skip_prk_steps: - prev_sample, state = self.step_plms(state, model_output, timestep, sample) - else: - prk_prev_sample, prk_state = self.step_prk(state, model_output, timestep, sample) - plms_prev_sample, plms_state = self.step_plms(state, model_output, timestep, sample) - - cond = state.counter < len(state.prk_timesteps) - - prev_sample = jax.lax.select(cond, prk_prev_sample, plms_prev_sample) - - state = state.replace( - cur_model_output=jax.lax.select(cond, prk_state.cur_model_output, plms_state.cur_model_output), - ets=jax.lax.select(cond, prk_state.ets, plms_state.ets), - cur_sample=jax.lax.select(cond, prk_state.cur_sample, plms_state.cur_sample), - counter=jax.lax.select(cond, prk_state.counter, plms_state.counter), - ) - - if not return_dict: - return (prev_sample, state) - - return FlaxPNDMSchedulerOutput(prev_sample=prev_sample, state=state) - - def step_prk( - self, - state: PNDMSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - ) -> Union[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. - - Args: - state (`PNDMSchedulerState`): the `FlaxPNDMScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - return_dict (`bool`): option for returning tuple rather than FlaxPNDMSchedulerOutput class - - Returns: - [`FlaxPNDMSchedulerOutput`] or `tuple`: [`FlaxPNDMSchedulerOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - - if state.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - diff_to_prev = jnp.where( - 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] - - model_output = jax.lax.select( - (state.counter % 4) != 3, - model_output, # remainder 0, 1, 2 - state.cur_model_output + 1 / 6 * model_output, # remainder 3 - ) - - state = state.replace( - cur_model_output=jax.lax.select_n( - state.counter % 4, - state.cur_model_output + 1 / 6 * model_output, # remainder 0 - state.cur_model_output + 1 / 3 * model_output, # remainder 1 - state.cur_model_output + 1 / 3 * model_output, # remainder 2 - jnp.zeros_like(state.cur_model_output), # remainder 3 - ), - ets=jax.lax.select( - (state.counter % 4) == 0, - state.ets.at[0:3].set(state.ets[1:4]).at[3].set(model_output), # remainder 0 - state.ets, # remainder 1, 2, 3 - ), - cur_sample=jax.lax.select( - (state.counter % 4) == 0, - sample, # remainder 0 - state.cur_sample, # remainder 1, 2, 3 - ), - ) - - cur_sample = state.cur_sample - prev_sample = self._get_prev_sample(state, cur_sample, timestep, prev_timestep, model_output) - state = state.replace(counter=state.counter + 1) - - return (prev_sample, state) - - def step_plms( - self, - state: PNDMSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - ) -> Union[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. - - Args: - state (`PNDMSchedulerState`): the `FlaxPNDMScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - return_dict (`bool`): option for returning tuple rather than FlaxPNDMSchedulerOutput class - - Returns: - [`FlaxPNDMSchedulerOutput`] or `tuple`: [`FlaxPNDMSchedulerOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is the sample tensor. - - """ - - if state.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - # NOTE: There is no way to check in the jitted runtime if the prk mode was ran before - - prev_timestep = timestep - self.config.num_train_timesteps // state.num_inference_steps - prev_timestep = jnp.where(prev_timestep > 0, prev_timestep, 0) - - # Reference: - # if state.counter != 1: - # state.ets.append(model_output) - # else: - # prev_timestep = timestep - # timestep = timestep + self.config.num_train_timesteps // state.num_inference_steps - - 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 - ) - - # Reference: - # if len(state.ets) == 1 and state.counter == 0: - # model_output = model_output - # state.cur_sample = sample - # elif len(state.ets) == 1 and state.counter == 1: - # model_output = (model_output + state.ets[-1]) / 2 - # sample = state.cur_sample - # state.cur_sample = None - # elif len(state.ets) == 2: - # model_output = (3 * state.ets[-1] - state.ets[-2]) / 2 - # elif len(state.ets) == 3: - # model_output = (23 * state.ets[-1] - 16 * state.ets[-2] + 5 * state.ets[-3]) / 12 - # else: - # model_output = (1 / 24) * (55 * state.ets[-1] - 59 * state.ets[-2] + 37 * state.ets[-3] - 9 * state.ets[-4]) - - state = state.replace( - ets=jax.lax.select( - state.counter != 1, - state.ets.at[0:3].set(state.ets[1:4]).at[3].set(model_output), # counter != 1 - state.ets, # counter 1 - ), - cur_sample=jax.lax.select( - state.counter != 1, - sample, # counter != 1 - state.cur_sample, # counter 1 - ), - ) - - state = state.replace( - cur_model_output=jax.lax.select_n( - jnp.clip(state.counter, 0, 4), - model_output, # counter 0 - (model_output + state.ets[-1]) / 2, # counter 1 - (3 * state.ets[-1] - state.ets[-2]) / 2, # counter 2 - (23 * state.ets[-1] - 16 * state.ets[-2] + 5 * state.ets[-3]) / 12, # counter 3 - (1 / 24) - * (55 * state.ets[-1] - 59 * state.ets[-2] + 37 * state.ets[-3] - 9 * state.ets[-4]), # counter >= 4 - ), - ) - - sample = state.cur_sample - model_output = state.cur_model_output - prev_sample = self._get_prev_sample(state, sample, timestep, prev_timestep, model_output) - state = state.replace(counter=state.counter + 1) - - return (prev_sample, state) - - def _get_prev_sample(self, state: PNDMSchedulerState, sample, timestep, prev_timestep, model_output): - # See formula (9) of PNDM paper https://huggingface.co/papers/2202.09778 - # this function computes x_(t−δ) using the formula of (9) - # Note that x_t needs to be added to both sides of the equation - - # Notation ( -> - # alpha_prod_t -> α_t - # alpha_prod_t_prev -> α_(t−δ) - # beta_prod_t -> (1 - α_t) - # beta_prod_t_prev -> (1 - α_(t−δ)) - # sample -> x_t - # model_output -> e_θ(x_t, t) - # 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 - ) - beta_prod_t = 1 - alpha_prod_t - beta_prod_t_prev = 1 - alpha_prod_t_prev - - if self.config.prediction_type == "v_prediction": - model_output = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - elif self.config.prediction_type != "epsilon": - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon` or `v_prediction`" - ) - - # corresponds to (α_(t−δ) - α_t) divided by - # denominator of x_t in formula (9) and plus 1 - # Note: (α_(t−δ) - α_t) / (sqrt(α_t) * (sqrt(α_(t−δ)) + sqr(α_t))) = - # sqrt(α_(t−δ)) / sqrt(α_t)) - sample_coeff = (alpha_prod_t_prev / alpha_prod_t) ** (0.5) - - # corresponds to denominator of e_θ(x_t, t) in formula (9) - model_output_denom_coeff = alpha_prod_t * beta_prod_t_prev ** (0.5) + ( - alpha_prod_t * beta_prod_t * alpha_prod_t_prev - ) ** (0.5) - - # full formula (9) - prev_sample = ( - sample_coeff * sample - (alpha_prod_t_prev - alpha_prod_t) * model_output / model_output_denom_coeff - ) - - return prev_sample - - def add_noise( - self, - state: PNDMSchedulerState, - original_samples: jnp.ndarray, - noise: jnp.ndarray, - timesteps: jnp.ndarray, - ) -> jnp.ndarray: - return add_noise_common(state.common, original_samples, noise, timesteps) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_repaint.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_repaint.py deleted file mode 100644 index 6530c5af9e5bc48657de83e1cf21620472534fd3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_repaint.py +++ /dev/null @@ -1,365 +0,0 @@ -# Copyright 2025 ETH Zurich Computer Vision Lab 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 Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import SchedulerMixin - - -@dataclass -class RePaintSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's step function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample (x_{0}) based on the model output from - the current timestep. `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: torch.Tensor - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -class RePaintScheduler(SchedulerMixin, ConfigMixin): - """ - `RePaintScheduler` is a scheduler for DDPM inpainting inside a given mask. - - 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`, 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`. - eta (`float`): - The weight of noise for added noise in diffusion step. If its value is between 0.0 and 1.0 it corresponds - to the DDIM scheduler, and if its value is between -0.0 and 1.0 it corresponds to the DDPM scheduler. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - clip_sample (`bool`, defaults to `True`): - Clip the predicted sample between -1 and 1 for numerical stability. - - """ - - order = 1 - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.0001, - beta_end: float = 0.02, - beta_schedule: str = "linear", - eta: float = 0.0, - trained_betas: Optional[np.ndarray] = None, - clip_sample: bool = True, - ): - if trained_betas is not None: - self.betas = torch.from_numpy(trained_betas) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - elif beta_schedule == "sigmoid": - # GeoDiff sigmoid schedule - betas = torch.linspace(-6, 6, num_train_timesteps) - self.betas = torch.sigmoid(betas) * (beta_end - beta_start) + beta_start - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - self.one = torch.tensor(1.0) - - self.final_alpha_cumprod = torch.tensor(1.0) - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # setable values - self.num_inference_steps = None - self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy()) - - self.eta = eta - - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - def set_timesteps( - self, - num_inference_steps: int, - jump_length: int = 10, - jump_n_sample: int = 10, - device: Union[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. If used, - `timesteps` must be `None`. - jump_length (`int`, defaults to 10): - The number of steps taken forward in time before going backward in time for a single jump (“j” in - RePaint paper). Take a look at Figure 9 and 10 in the paper. - jump_n_sample (`int`, defaults to 10): - The number of times to make a forward time jump for a given chosen time sample. Take a look at Figure 9 - and 10 in the paper. - device (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - - """ - num_inference_steps = min(self.config.num_train_timesteps, num_inference_steps) - self.num_inference_steps = num_inference_steps - - timesteps = [] - - jumps = {} - for j in range(0, num_inference_steps - jump_length, jump_length): - jumps[j] = jump_n_sample - 1 - - t = num_inference_steps - while t >= 1: - t = t - 1 - timesteps.append(t) - - if jumps.get(t, 0) > 0: - jumps[t] = jumps[t] - 1 - for _ in range(jump_length): - t = t + 1 - timesteps.append(t) - - timesteps = np.array(timesteps) * (self.config.num_train_timesteps // self.num_inference_steps) - self.timesteps = torch.from_numpy(timesteps).to(device) - - def _get_variance(self, t): - prev_timestep = t - self.config.num_train_timesteps // self.num_inference_steps - - alpha_prod_t = self.alphas_cumprod[t] - 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 - beta_prod_t_prev = 1 - alpha_prod_t_prev - - # For t > 0, compute predicted variance βt (see formula (6) and (7) from - # https://huggingface.co/papers/2006.11239) and sample from it to get - # previous sample x_{t-1} ~ N(pred_prev_sample, variance) == add - # variance to pred_sample - # Is equivalent to formula (16) in https://huggingface.co/papers/2010.02502 - # without eta. - # variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * self.betas[t] - variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev) - - return variance - - def step( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - original_image: torch.Tensor, - mask: torch.Tensor, - generator: Optional[torch.Generator] = None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - original_image (`torch.Tensor`): - The original image to inpaint on. - mask (`torch.Tensor`): - The mask where a value of 0.0 indicates which part of the original image to inpaint. - generator (`torch.Generator`, *optional*): - A random number generator. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] or `tuple`. - - Returns: - [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] is returned, - otherwise a tuple is returned where the first element is the sample tensor. - - """ - t = timestep - prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps - - # 1. compute alphas, betas - alpha_prod_t = self.alphas_cumprod[t] - 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 - - # 2. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (15) from https://huggingface.co/papers/2006.11239 - pred_original_sample = (sample - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.5 - - # 3. Clip "predicted x_0" - if self.config.clip_sample: - pred_original_sample = torch.clamp(pred_original_sample, -1, 1) - - # We choose to follow RePaint Algorithm 1 to get x_{t-1}, however we - # substitute formula (7) in the algorithm coming from DDPM paper - # (formula (4) Algorithm 2 - Sampling) with formula (12) from DDIM paper. - # DDIM schedule gives the same results as DDPM with eta = 1.0 - # Noise is being reused in 7. and 8., but no impact on quality has - # been observed. - - # 5. Add noise - device = model_output.device - noise = randn_tensor(model_output.shape, generator=generator, device=device, dtype=model_output.dtype) - std_dev_t = self.eta * self._get_variance(timestep) ** 0.5 - - variance = 0 - if t > 0 and self.eta > 0: - variance = std_dev_t * noise - - # 6. compute "direction pointing to x_t" of formula (12) - # from https://huggingface.co/papers/2010.02502 - pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** 0.5 * model_output - - # 7. compute x_{t-1} of formula (12) from https://huggingface.co/papers/2010.02502 - prev_unknown_part = alpha_prod_t_prev**0.5 * pred_original_sample + pred_sample_direction + variance - - # 8. Algorithm 1 Line 5 https://huggingface.co/papers/2201.09865 - # The computation reported in Algorithm 1 Line 5 is incorrect. Line 5 refers to formula (8a) of the same paper, - # which tells to sample from a Gaussian distribution with mean "(alpha_prod_t_prev**0.5) * original_image" - # and variance "(1 - alpha_prod_t_prev)". This means that the standard Gaussian distribution "noise" should be - # scaled by the square root of the variance (as it is done here), however Algorithm 1 Line 5 tells to scale by the variance. - prev_known_part = (alpha_prod_t_prev**0.5) * original_image + ((1 - alpha_prod_t_prev) ** 0.5) * noise - - # 9. Algorithm 1 Line 8 https://huggingface.co/papers/2201.09865 - pred_prev_sample = mask * prev_known_part + (1.0 - mask) * prev_unknown_part - - if not return_dict: - return ( - pred_prev_sample, - pred_original_sample, - ) - - return RePaintSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample) - - def undo_step(self, sample, timestep, generator=None): - n = self.config.num_train_timesteps // self.num_inference_steps - - for i in range(n): - beta = self.betas[timestep + i] - if sample.device.type == "mps": - # randn does not work reproducibly on mps - noise = randn_tensor(sample.shape, dtype=sample.dtype, generator=generator) - noise = noise.to(sample.device) - else: - noise = randn_tensor(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype) - - # 10. Algorithm 1 Line 10 https://huggingface.co/papers/2201.09865 - sample = (1 - beta) ** 0.5 * sample + beta**0.5 * noise - - return sample - - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - raise NotImplementedError("Use `DDPMScheduler.add_noise()` to train for sampling with RePaint.") - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_sasolver.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_sasolver.py deleted file mode 100644 index 2df7d560ddfb68d786ad3e96227a56ead053ae4b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_sasolver.py +++ /dev/null @@ -1,1221 +0,0 @@ -# Copyright 2025 Shuchen Xue, etc. in University of Chinese Academy of Sciences 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. - -# DISCLAIMER: check https://huggingface.co/papers/2309.05019 -# 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, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import deprecate, is_scipy_available -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput - - -if is_scipy_available(): - import scipy.stats - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -class SASolverScheduler(SchedulerMixin, ConfigMixin): - """ - `SASolverScheduler` is a fast dedicated high-order solver for diffusion SDEs. - - 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`, 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`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - predictor_order (`int`, defaults to 2): - The predictor order which can be `1` or `2` or `3` or '4'. It is recommended to use `predictor_order=2` for - guided sampling, and `predictor_order=3` for unconditional sampling. - corrector_order (`int`, defaults to 2): - The corrector order which can be `1` or `2` or `3` or '4'. It is recommended to use `corrector_order=2` for - guided sampling, and `corrector_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://imagen.research.google/video/paper.pdf) paper). - tau_func (`Callable`, *optional*): - Stochasticity during the sampling. Default in init is `lambda t: 1 if t >= 200 and t <= 800 else 0`. - SA-Solver will sample from vanilla diffusion ODE if tau_func is set to `lambda t: 0`. SA-Solver will sample - from vanilla diffusion SDE if tau_func is set to `lambda t: 1`. For more details, please check - https://huggingface.co/papers/2309.05019 - thresholding (`bool`, defaults to `False`): - Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such - as Stable Diffusion. - dynamic_thresholding_ratio (`float`, defaults to 0.995): - The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. - sample_max_value (`float`, defaults to 1.0): - The threshold value for dynamic thresholding. Valid only when `thresholding=True` and - `algorithm_type="dpmsolver++"`. - algorithm_type (`str`, defaults to `data_prediction`): - Algorithm type for the solver; can be `data_prediction` or `noise_prediction`. It is recommended to use - `data_prediction` with `solver_order=2` for guided sampling like in Stable Diffusion. - lower_order_final (`bool`, defaults to `True`): - Whether to use lower-order solvers in the final steps. Default = True. - use_karras_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, - the sigmas are determined according to a sequence of noise levels {σi}. - use_exponential_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. - 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. - 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. - variance_type (`str`, *optional*): - Set to "learned" or "learned_range" for diffusion models that predict variance. If set, the model's output - contains the predicted Gaussian variance. - 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. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - 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, - predictor_order: int = 2, - corrector_order: int = 2, - prediction_type: str = "epsilon", - tau_func: Optional[Callable] = 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, - lambda_min_clipped: float = -float("inf"), - variance_type: Optional[str] = None, - timestep_spacing: str = "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: - raise ValueError( - "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." - ) - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - ) - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - # Currently we only support VP-type noise schedule - self.alpha_t = torch.sqrt(self.alphas_cumprod) - self.sigma_t = torch.sqrt(1 - self.alphas_cumprod) - self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t) - self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5 - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - if algorithm_type not in ["data_prediction", "noise_prediction"]: - raise NotImplementedError(f"{algorithm_type} is not implemented for {self.__class__}") - - # setable values - self.num_inference_steps = None - timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy() - self.timesteps = torch.from_numpy(timesteps) - self.timestep_list = [None] * max(predictor_order, corrector_order - 1) - self.model_outputs = [None] * max(predictor_order, corrector_order - 1) - - if tau_func is None: - self.tau_func = lambda t: 1 if t >= 200 and t <= 800 else 0 - else: - self.tau_func = tau_func - self.predict_x0 = algorithm_type == "data_prediction" - self.lower_order_nums = 0 - self.last_sample = None - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 set_timesteps(self, num_inference_steps: int = None, device: Union[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 (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - """ - # Clipping the minimum of all lambda(t) for numerical stability. - # This is critical for cosine (squaredcos_cap_v2) noise schedule. - clipped_idx = torch.searchsorted(torch.flip(self.lambda_t, [0]), self.config.lambda_min_clipped) - last_timestep = ((self.config.num_train_timesteps - clipped_idx).numpy()).item() - - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = ( - np.linspace(0, last_timestep - 1, num_inference_steps + 1).round()[::-1][:-1].copy().astype(np.int64) - ) - - elif self.config.timestep_spacing == "leading": - step_ratio = last_timestep // (num_inference_steps + 1) - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int64) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = np.arange(last_timestep, 0, -step_ratio).round().copy().astype(np.int64) - timesteps -= 1 - else: - raise ValueError( - 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) - log_sigmas = np.log(sigmas) - if self.config.use_karras_sigmas: - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round() - sigmas = np.concatenate([sigmas, sigmas[-1:]]).astype(np.float32) - elif self.config.use_exponential_sigmas: - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - sigmas = np.concatenate([sigmas, sigmas[-1:]]).astype(np.float32) - elif self.config.use_beta_sigmas: - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - sigmas = np.concatenate([sigmas, sigmas[-1:]]).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() - timesteps = (sigmas * self.config.num_train_timesteps).copy() - sigmas = np.concatenate([sigmas, sigmas[-1:]]).astype(np.float32) - else: - sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) - sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5 - sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) - - self.sigmas = torch.from_numpy(sigmas) - self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) - - self.num_inference_steps = len(timesteps) - self.model_outputs = [ - None, - ] * max(self.config.predictor_order, self.config.corrector_order - 1) - self.lower_order_nums = 0 - self.last_sample = None - - # add an index counter for schedulers that allow duplicated timesteps - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - """ - "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the - prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by - s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing - pixels from saturation at each step. We find that dynamic thresholding results in significantly better - photorealism as well as better image-text alignment, especially when using very large guidance weights." - - https://huggingface.co/papers/2205.11487 - """ - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half - - # Flatten sample for doing quantile calculation along each image - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - - abs_sample = sample.abs() # "a certain percentile absolute pixel value" - - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] - s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 - sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" - - sample = sample.reshape(batch_size, channels, *remaining_dims) - sample = sample.to(dtype) - - return sample - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma - log_sigma = np.log(np.maximum(sigma, 1e-10)) - - # get distribution - dists = log_sigma - log_sigmas[:, np.newaxis] - - # get sigmas range - low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) - high_idx = low_idx + 1 - - low = log_sigmas[low_idx] - high = log_sigmas[high_idx] - - # interpolate sigmas - w = (low - log_sigma) / (low - high) - w = np.clip(w, 0, 1) - - # transform interpolation to time range - t = (1 - w) * low_idx + w * high_idx - t = t.reshape(sigma.shape) - return t - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t - def _sigma_to_alpha_sigma_t(self, sigma): - if self.config.use_flow_sigmas: - alpha_t = 1 - sigma - sigma_t = sigma - else: - alpha_t = 1 / ((sigma**2 + 1) ** 0.5) - sigma_t = sigma * alpha_t - - 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: - """Constructs the noise schedule of Karras et al. (2022).""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - rho = 7.0 # 7.0 is the value used in the paper - ramp = np.linspace(0, 1, num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Constructs an exponential noise schedule.""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: - """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return sigmas - - def convert_model_output( - self, - model_output: torch.Tensor, - *args, - sample: torch.Tensor = None, - **kwargs, - ) -> torch.Tensor: - """ - Convert the model output to the corresponding type the data_prediction/noise_prediction algorithm needs. - Noise_prediction is designed to discretize an integral of the noise prediction model, and data_prediction is - designed to discretize an integral of the data prediction model. - - - - The algorithm and model type are decoupled. You can use either data_prediction or noise_prediction for both - noise prediction and data prediction models. - - - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. - 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`", - ) - - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - # SA-Solver_data_prediction needs to solve an integral of the data prediction model. - if self.config.algorithm_type in ["data_prediction"]: - if self.config.prediction_type == "epsilon": - # SA-Solver only needs the "mean" output. - if self.config.variance_type in ["learned", "learned_range"]: - model_output = model_output[:, :3] - 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": - sigma_t = self.sigmas[self.step_index] - 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 SASolverScheduler." - ) - - if self.config.thresholding: - x0_pred = self._threshold_sample(x0_pred) - - return x0_pred - - # SA-Solver_noise_prediction needs to solve an integral of the noise prediction model. - elif self.config.algorithm_type in ["noise_prediction"]: - if self.config.prediction_type == "epsilon": - # SA-Solver only needs the "mean" output. - if self.config.variance_type in ["learned", "learned_range"]: - epsilon = model_output[:, :3] - else: - epsilon = model_output - elif self.config.prediction_type == "sample": - epsilon = (sample - alpha_t * model_output) / sigma_t - elif self.config.prediction_type == "v_prediction": - epsilon = alpha_t * model_output + sigma_t * sample - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" - " `v_prediction` for the SASolverScheduler." - ) - - if self.config.thresholding: - alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep] - x0_pred = (sample - sigma_t * epsilon) / alpha_t - x0_pred = self._threshold_sample(x0_pred) - epsilon = (sample - alpha_t * x0_pred) / sigma_t - - return epsilon - - def get_coefficients_exponential_negative(self, order, interval_start, interval_end): - """ - Calculate the integral of exp(-x) * x^order dx from interval_start to interval_end - """ - assert order in [0, 1, 2, 3], "order is only supported for 0, 1, 2 and 3" - - if order == 0: - return torch.exp(-interval_end) * (torch.exp(interval_end - interval_start) - 1) - elif order == 1: - return torch.exp(-interval_end) * ( - (interval_start + 1) * torch.exp(interval_end - interval_start) - (interval_end + 1) - ) - elif order == 2: - return torch.exp(-interval_end) * ( - (interval_start**2 + 2 * interval_start + 2) * torch.exp(interval_end - interval_start) - - (interval_end**2 + 2 * interval_end + 2) - ) - elif order == 3: - return torch.exp(-interval_end) * ( - (interval_start**3 + 3 * interval_start**2 + 6 * interval_start + 6) - * torch.exp(interval_end - interval_start) - - (interval_end**3 + 3 * interval_end**2 + 6 * interval_end + 6) - ) - - def get_coefficients_exponential_positive(self, order, interval_start, interval_end, tau): - """ - Calculate the integral of exp(x(1+tau^2)) * x^order dx from interval_start to interval_end - """ - assert order in [0, 1, 2, 3], "order is only supported for 0, 1, 2 and 3" - - # after change of variable(cov) - interval_end_cov = (1 + tau**2) * interval_end - interval_start_cov = (1 + tau**2) * interval_start - - if order == 0: - return ( - torch.exp(interval_end_cov) * (1 - torch.exp(-(interval_end_cov - interval_start_cov))) / (1 + tau**2) - ) - elif order == 1: - return ( - torch.exp(interval_end_cov) - * ( - (interval_end_cov - 1) - - (interval_start_cov - 1) * torch.exp(-(interval_end_cov - interval_start_cov)) - ) - / ((1 + tau**2) ** 2) - ) - elif order == 2: - return ( - torch.exp(interval_end_cov) - * ( - (interval_end_cov**2 - 2 * interval_end_cov + 2) - - (interval_start_cov**2 - 2 * interval_start_cov + 2) - * torch.exp(-(interval_end_cov - interval_start_cov)) - ) - / ((1 + tau**2) ** 3) - ) - elif order == 3: - return ( - torch.exp(interval_end_cov) - * ( - (interval_end_cov**3 - 3 * interval_end_cov**2 + 6 * interval_end_cov - 6) - - (interval_start_cov**3 - 3 * interval_start_cov**2 + 6 * interval_start_cov - 6) - * torch.exp(-(interval_end_cov - interval_start_cov)) - ) - / ((1 + tau**2) ** 4) - ) - - def lagrange_polynomial_coefficient(self, order, lambda_list): - """ - Calculate the coefficient of lagrange polynomial - """ - - assert order in [0, 1, 2, 3] - assert order == len(lambda_list) - 1 - if order == 0: - return [[1]] - elif order == 1: - return [ - [ - 1 / (lambda_list[0] - lambda_list[1]), - -lambda_list[1] / (lambda_list[0] - lambda_list[1]), - ], - [ - 1 / (lambda_list[1] - lambda_list[0]), - -lambda_list[0] / (lambda_list[1] - lambda_list[0]), - ], - ] - elif order == 2: - denominator1 = (lambda_list[0] - lambda_list[1]) * (lambda_list[0] - lambda_list[2]) - denominator2 = (lambda_list[1] - lambda_list[0]) * (lambda_list[1] - lambda_list[2]) - denominator3 = (lambda_list[2] - lambda_list[0]) * (lambda_list[2] - lambda_list[1]) - return [ - [ - 1 / denominator1, - (-lambda_list[1] - lambda_list[2]) / denominator1, - lambda_list[1] * lambda_list[2] / denominator1, - ], - [ - 1 / denominator2, - (-lambda_list[0] - lambda_list[2]) / denominator2, - lambda_list[0] * lambda_list[2] / denominator2, - ], - [ - 1 / denominator3, - (-lambda_list[0] - lambda_list[1]) / denominator3, - lambda_list[0] * lambda_list[1] / denominator3, - ], - ] - elif order == 3: - denominator1 = ( - (lambda_list[0] - lambda_list[1]) - * (lambda_list[0] - lambda_list[2]) - * (lambda_list[0] - lambda_list[3]) - ) - denominator2 = ( - (lambda_list[1] - lambda_list[0]) - * (lambda_list[1] - lambda_list[2]) - * (lambda_list[1] - lambda_list[3]) - ) - denominator3 = ( - (lambda_list[2] - lambda_list[0]) - * (lambda_list[2] - lambda_list[1]) - * (lambda_list[2] - lambda_list[3]) - ) - denominator4 = ( - (lambda_list[3] - lambda_list[0]) - * (lambda_list[3] - lambda_list[1]) - * (lambda_list[3] - lambda_list[2]) - ) - return [ - [ - 1 / denominator1, - (-lambda_list[1] - lambda_list[2] - lambda_list[3]) / denominator1, - ( - lambda_list[1] * lambda_list[2] - + lambda_list[1] * lambda_list[3] - + lambda_list[2] * lambda_list[3] - ) - / denominator1, - (-lambda_list[1] * lambda_list[2] * lambda_list[3]) / denominator1, - ], - [ - 1 / denominator2, - (-lambda_list[0] - lambda_list[2] - lambda_list[3]) / denominator2, - ( - lambda_list[0] * lambda_list[2] - + lambda_list[0] * lambda_list[3] - + lambda_list[2] * lambda_list[3] - ) - / denominator2, - (-lambda_list[0] * lambda_list[2] * lambda_list[3]) / denominator2, - ], - [ - 1 / denominator3, - (-lambda_list[0] - lambda_list[1] - lambda_list[3]) / denominator3, - ( - lambda_list[0] * lambda_list[1] - + lambda_list[0] * lambda_list[3] - + lambda_list[1] * lambda_list[3] - ) - / denominator3, - (-lambda_list[0] * lambda_list[1] * lambda_list[3]) / denominator3, - ], - [ - 1 / denominator4, - (-lambda_list[0] - lambda_list[1] - lambda_list[2]) / denominator4, - ( - lambda_list[0] * lambda_list[1] - + lambda_list[0] * lambda_list[2] - + lambda_list[1] * lambda_list[2] - ) - / denominator4, - (-lambda_list[0] * lambda_list[1] * lambda_list[2]) / denominator4, - ], - ] - - def get_coefficients_fn(self, order, interval_start, interval_end, lambda_list, tau): - assert order in [1, 2, 3, 4] - assert order == len(lambda_list), "the length of lambda list must be equal to the order" - coefficients = [] - lagrange_coefficient = self.lagrange_polynomial_coefficient(order - 1, lambda_list) - for i in range(order): - coefficient = 0 - for j in range(order): - if self.predict_x0: - coefficient += lagrange_coefficient[i][j] * self.get_coefficients_exponential_positive( - order - 1 - j, interval_start, interval_end, tau - ) - else: - coefficient += lagrange_coefficient[i][j] * self.get_coefficients_exponential_negative( - order - 1 - j, interval_start, interval_end - ) - coefficients.append(coefficient) - assert len(coefficients) == order, "the length of coefficients does not match the order" - return coefficients - - def stochastic_adams_bashforth_update( - self, - model_output: torch.Tensor, - *args, - sample: torch.Tensor, - noise: torch.Tensor, - order: int, - tau: torch.Tensor, - **kwargs, - ) -> torch.Tensor: - """ - One step for the SA-Predictor. - - 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 SA-Predictor at this timestep. - - 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 noise is None: - if len(args) > 2: - noise = args[2] - else: - raise ValueError("missing `noise` 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 tau is None: - if len(args) > 4: - tau = args[4] - else: - raise ValueError("missing `tau` 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 - sigma_t, sigma_s0 = ( - self.sigmas[self.step_index + 1], - self.sigmas[self.step_index], - ) - 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) - - gradient_part = torch.zeros_like(sample) - h = lambda_t - lambda_s0 - lambda_list = [] - - for i in range(order): - si = self.step_index - i - alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) - lambda_si = torch.log(alpha_si) - torch.log(sigma_si) - lambda_list.append(lambda_si) - - gradient_coefficients = self.get_coefficients_fn(order, lambda_s0, lambda_t, lambda_list, tau) - - x = sample - - if self.predict_x0: - if ( - order == 2 - ): ## if order = 2 we do a modification that does not influence the convergence order similar to unipc. Note: This is used only for few steps sampling. - # The added term is O(h^3). Empirically we find it will slightly improve the image quality. - # ODE case - # gradient_coefficients[0] += 1.0 * torch.exp(lambda_t) * (h ** 2 / 2 - (h - 1 + torch.exp(-h))) / (ns.marginal_lambda(t_prev_list[-1]) - ns.marginal_lambda(t_prev_list[-2])) - # gradient_coefficients[1] -= 1.0 * torch.exp(lambda_t) * (h ** 2 / 2 - (h - 1 + torch.exp(-h))) / (ns.marginal_lambda(t_prev_list[-1]) - ns.marginal_lambda(t_prev_list[-2])) - temp_sigma = self.sigmas[self.step_index - 1] - temp_alpha_s, temp_sigma_s = self._sigma_to_alpha_sigma_t(temp_sigma) - temp_lambda_s = torch.log(temp_alpha_s) - torch.log(temp_sigma_s) - gradient_coefficients[0] += ( - 1.0 - * torch.exp((1 + tau**2) * lambda_t) - * (h**2 / 2 - (h * (1 + tau**2) - 1 + torch.exp((1 + tau**2) * (-h))) / ((1 + tau**2) ** 2)) - / (lambda_s0 - temp_lambda_s) - ) - gradient_coefficients[1] -= ( - 1.0 - * torch.exp((1 + tau**2) * lambda_t) - * (h**2 / 2 - (h * (1 + tau**2) - 1 + torch.exp((1 + tau**2) * (-h))) / ((1 + tau**2) ** 2)) - / (lambda_s0 - temp_lambda_s) - ) - - for i in range(order): - if self.predict_x0: - gradient_part += ( - (1 + tau**2) - * sigma_t - * torch.exp(-(tau**2) * lambda_t) - * gradient_coefficients[i] - * model_output_list[-(i + 1)] - ) - else: - gradient_part += -(1 + tau**2) * alpha_t * gradient_coefficients[i] * model_output_list[-(i + 1)] - - if self.predict_x0: - noise_part = sigma_t * torch.sqrt(1 - torch.exp(-2 * tau**2 * h)) * noise - else: - noise_part = tau * sigma_t * torch.sqrt(torch.exp(2 * h) - 1) * noise - - if self.predict_x0: - x_t = torch.exp(-(tau**2) * h) * (sigma_t / sigma_s0) * x + gradient_part + noise_part - else: - x_t = (alpha_t / alpha_s0) * x + gradient_part + noise_part - - x_t = x_t.to(x.dtype) - return x_t - - def stochastic_adams_moulton_update( - self, - this_model_output: torch.Tensor, - *args, - last_sample: torch.Tensor, - last_noise: torch.Tensor, - this_sample: torch.Tensor, - order: int, - tau: torch.Tensor, - **kwargs, - ) -> torch.Tensor: - """ - One step for the SA-Corrector. - - 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 order of SA-Corrector at this step. - - 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 last_noise is None: - if len(args) > 2: - last_noise = args[2] - else: - raise ValueError("missing `last_noise` as a required keyword argument") - if this_sample is None: - if len(args) > 3: - this_sample = args[3] - else: - raise ValueError("missing `this_sample` as a required keyword argument") - if order is None: - if len(args) > 4: - order = args[4] - else: - raise ValueError("missing `order` as a required keyword argument") - if tau is None: - if len(args) > 5: - tau = args[5] - else: - raise ValueError("missing `tau` 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 - sigma_t, sigma_s0 = ( - self.sigmas[self.step_index], - self.sigmas[self.step_index - 1], - ) - 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) - gradient_part = torch.zeros_like(this_sample) - h = lambda_t - lambda_s0 - lambda_list = [] - for i in range(order): - si = self.step_index - i - alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) - lambda_si = torch.log(alpha_si) - torch.log(sigma_si) - lambda_list.append(lambda_si) - - model_prev_list = model_output_list + [this_model_output] - - gradient_coefficients = self.get_coefficients_fn(order, lambda_s0, lambda_t, lambda_list, tau) - - x = last_sample - - if self.predict_x0: - if ( - order == 2 - ): ## if order = 2 we do a modification that does not influence the convergence order similar to UniPC. Note: This is used only for few steps sampling. - # The added term is O(h^3). Empirically we find it will slightly improve the image quality. - # ODE case - # gradient_coefficients[0] += 1.0 * torch.exp(lambda_t) * (h / 2 - (h - 1 + torch.exp(-h)) / h) - # gradient_coefficients[1] -= 1.0 * torch.exp(lambda_t) * (h / 2 - (h - 1 + torch.exp(-h)) / h) - gradient_coefficients[0] += ( - 1.0 - * torch.exp((1 + tau**2) * lambda_t) - * (h / 2 - (h * (1 + tau**2) - 1 + torch.exp((1 + tau**2) * (-h))) / ((1 + tau**2) ** 2 * h)) - ) - gradient_coefficients[1] -= ( - 1.0 - * torch.exp((1 + tau**2) * lambda_t) - * (h / 2 - (h * (1 + tau**2) - 1 + torch.exp((1 + tau**2) * (-h))) / ((1 + tau**2) ** 2 * h)) - ) - - for i in range(order): - if self.predict_x0: - gradient_part += ( - (1 + tau**2) - * sigma_t - * torch.exp(-(tau**2) * lambda_t) - * gradient_coefficients[i] - * model_prev_list[-(i + 1)] - ) - else: - gradient_part += -(1 + tau**2) * alpha_t * gradient_coefficients[i] * model_prev_list[-(i + 1)] - - if self.predict_x0: - noise_part = sigma_t * torch.sqrt(1 - torch.exp(-2 * tau**2 * h)) * last_noise - else: - noise_part = tau * sigma_t * torch.sqrt(torch.exp(2 * h) - 1) * last_noise - - if self.predict_x0: - x_t = torch.exp(-(tau**2) * h) * (sigma_t / sigma_s0) * x + gradient_part + noise_part - else: - x_t = (alpha_t / alpha_s0) * x + gradient_part + noise_part - - x_t = x_t.to(x.dtype) - return x_t - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep - def index_for_timestep(self, timestep, schedule_timesteps=None): - if schedule_timesteps is None: - schedule_timesteps = self.timesteps - - index_candidates = (schedule_timesteps == timestep).nonzero() - - if len(index_candidates) == 0: - step_index = len(self.timesteps) - 1 - # 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) - elif len(index_candidates) > 1: - step_index = index_candidates[1].item() - else: - step_index = index_candidates[0].item() - - return step_index - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index - def _init_step_index(self, timestep): - """ - Initialize the step_index counter for the scheduler. - """ - - 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( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - generator=None, - return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: - """ - Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with - the SA-Solver. - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - generator (`torch.Generator`, *optional*): - A random number generator. - return_dict (`bool`): - 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( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - if self.step_index is None: - self._init_step_index(timestep) - - use_corrector = self.step_index > 0 and self.last_sample is not None - - model_output_convert = self.convert_model_output(model_output, sample=sample) - - if use_corrector: - current_tau = self.tau_func(self.timestep_list[-1]) - sample = self.stochastic_adams_moulton_update( - this_model_output=model_output_convert, - last_sample=self.last_sample, - last_noise=self.last_noise, - this_sample=sample, - order=self.this_corrector_order, - tau=current_tau, - ) - - for i in range(max(self.config.predictor_order, self.config.corrector_order - 1) - 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 - - noise = randn_tensor( - model_output.shape, - generator=generator, - device=model_output.device, - dtype=model_output.dtype, - ) - - if self.config.lower_order_final: - this_predictor_order = min(self.config.predictor_order, len(self.timesteps) - self.step_index) - this_corrector_order = min(self.config.corrector_order, len(self.timesteps) - self.step_index + 1) - else: - this_predictor_order = self.config.predictor_order - this_corrector_order = self.config.corrector_order - - self.this_predictor_order = min(this_predictor_order, self.lower_order_nums + 1) # warmup for multistep - self.this_corrector_order = min(this_corrector_order, self.lower_order_nums + 2) # warmup for multistep - assert self.this_predictor_order > 0 - assert self.this_corrector_order > 0 - - self.last_sample = sample - self.last_noise = noise - - current_tau = self.tau_func(self.timestep_list[-1]) - prev_sample = self.stochastic_adams_bashforth_update( - model_output=model_output_convert, - sample=sample, - noise=noise, - order=self.this_predictor_order, - tau=current_tau, - ) - - if self.lower_order_nums < max(self.config.predictor_order, self.config.corrector_order - 1): - self.lower_order_nums += 1 - - # upon completion increase step index by one - self._step_index += 1 - - if not return_dict: - return (prev_sample,) - - return SchedulerOutput(prev_sample=prev_sample) - - def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: - """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. - - Args: - sample (`torch.Tensor`): - The input sample. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as original_samples - # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement - # for the subsequent add_noise calls - self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype) - timesteps = timesteps.to(original_samples.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(original_samples.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_scm.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_scm.py deleted file mode 100644 index acff268c9b9231e3f28a6431574c64521f757b0c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_scm.py +++ /dev/null @@ -1,265 +0,0 @@ -# # Copyright 2025 Sana-Sprint Authors 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. - -# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion -# and https://github.com/hojonathanho/diffusion - -from dataclasses import dataclass -from typing import Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..schedulers.scheduling_utils import SchedulerMixin -from ..utils import BaseOutput, logging -from ..utils.torch_utils import randn_tensor - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->SCM -class SCMSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -class SCMScheduler(SchedulerMixin, ConfigMixin): - """ - `SCMScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with - non-Markovian guidance. 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`, defaults to 1000): - The number of diffusion steps to train the model. - prediction_type (`str`, defaults to `trigflow`): - Prediction type of the scheduler function. Currently only supports "trigflow". - sigma_data (`float`, defaults to 0.5): - The standard deviation of the noise added during multi-step inference. - """ - - # _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - prediction_type: str = "trigflow", - sigma_data: float = 0.5, - ): - """ - Initialize the SCM scheduler. - - Args: - num_train_timesteps (`int`, defaults to 1000): - The number of diffusion steps to train the model. - prediction_type (`str`, defaults to `trigflow`): - Prediction type of the scheduler function. Currently only supports "trigflow". - sigma_data (`float`, defaults to 0.5): - The standard deviation of the noise added during multi-step inference. - """ - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # setable values - self.num_inference_steps = None - self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)) - - self._step_index = None - self._begin_index = None - - @property - def step_index(self): - return self._step_index - - @property - 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): - """ - 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 set_timesteps( - self, - num_inference_steps: int, - timesteps: torch.Tensor = None, - device: Union[str, torch.device] = None, - max_timesteps: float = 1.57080, - intermediate_timesteps: float = 1.3, - ): - """ - 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. - timesteps (`torch.Tensor`, *optional*): - Custom timesteps to use for the denoising process. - max_timesteps (`float`, defaults to 1.57080): - 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). - """ - if num_inference_steps > self.config.num_train_timesteps: - raise ValueError( - f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:" - f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle" - f" maximal {self.config.num_train_timesteps} timesteps." - ) - - if timesteps is not None and len(timesteps) != num_inference_steps + 1: - raise ValueError("If providing custom timesteps, `timesteps` must be of length `num_inference_steps + 1`.") - - if timesteps is not None and max_timesteps is not None: - raise ValueError("If providing custom timesteps, `max_timesteps` should not be provided.") - - if timesteps is None and max_timesteps is None: - raise ValueError("Should provide either `timesteps` or `max_timesteps`.") - - if intermediate_timesteps is not None and num_inference_steps != 2: - raise ValueError("Intermediate timesteps for SCM is not supported when num_inference_steps != 2.") - - self.num_inference_steps = num_inference_steps - - if timesteps is not None: - if isinstance(timesteps, list): - self.timesteps = torch.tensor(timesteps, device=device).float() - elif isinstance(timesteps, torch.Tensor): - self.timesteps = timesteps.to(device).float() - else: - raise ValueError(f"Unsupported timesteps type: {type(timesteps)}") - elif intermediate_timesteps is not None: - self.timesteps = torch.tensor([max_timesteps, intermediate_timesteps, 0], device=device).float() - else: - # max_timesteps=arctan(80/0.5)=1.56454 is the default from sCM paper, we choose a different value here - self.timesteps = torch.linspace(max_timesteps, 0, num_inference_steps + 1, device=device).float() - print(f"Set timesteps: {self.timesteps}") - - self._step_index = None - self._begin_index = None - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - 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 - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep - 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 step( - self, - model_output: torch.FloatTensor, - timestep: float, - sample: torch.FloatTensor, - generator: torch.Generator = None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.FloatTensor`): - The direct output from learned diffusion model. - timestep (`float`): - The current discrete timestep in the diffusion chain. - sample (`torch.FloatTensor`): - A current instance of a sample created by the diffusion process. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~schedulers.scheduling_scm.SCMSchedulerOutput`] or `tuple`. - Returns: - [`~schedulers.scheduling_utils.SCMSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_scm.SCMSchedulerOutput`] is returned, otherwise a - tuple is returned where the first element is the sample tensor. - """ - 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 self.step_index is None: - self._init_step_index(timestep) - - # 2. compute alphas, betas - t = self.timesteps[self.step_index + 1] - s = self.timesteps[self.step_index] - - # 4. Different Parameterization: - parameterization = self.config.prediction_type - - if parameterization == "trigflow": - pred_x0 = torch.cos(s) * sample - torch.sin(s) * model_output - else: - raise ValueError(f"Unsupported parameterization: {parameterization}") - - # 5. Sample z ~ N(0, I), For MultiStep Inference - # Noise is not used for one-step sampling. - if len(self.timesteps) > 1: - noise = ( - randn_tensor(model_output.shape, device=model_output.device, generator=generator) - * self.config.sigma_data - ) - prev_sample = torch.cos(t) * pred_x0 + torch.sin(t) * noise - else: - prev_sample = pred_x0 - - self._step_index += 1 - - if not return_dict: - return (prev_sample, pred_x0) - - return SCMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_x0) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_sde_ve.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_sde_ve.py deleted file mode 100644 index 1bfc08cce5e97c3efcacbb4f18744cac441ff6fa..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_sde_ve.py +++ /dev/null @@ -1,301 +0,0 @@ -# Copyright 2025 Google Brain 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. - -# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch - -import math -from dataclasses import dataclass -from typing import Optional, Tuple, Union - -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import SchedulerMixin, SchedulerOutput - - -@dataclass -class SdeVeOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - prev_sample_mean (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Mean averaged `prev_sample` over previous timesteps. - """ - - prev_sample: torch.Tensor - prev_sample_mean: torch.Tensor - - -class ScoreSdeVeScheduler(SchedulerMixin, ConfigMixin): - """ - `ScoreSdeVeScheduler` is a variance exploding stochastic differential equation (SDE) scheduler. - - 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`, defaults to 1000): - The number of diffusion steps to train the model. - snr (`float`, defaults to 0.15): - A coefficient weighting the step from the `model_output` sample (from the network) to the random noise. - sigma_min (`float`, defaults to 0.01): - The initial noise scale for the sigma sequence in the sampling procedure. The minimum sigma should mirror - the distribution of the data. - sigma_max (`float`, defaults to 1348.0): - The maximum value used for the range of continuous timesteps passed into the model. - sampling_eps (`float`, defaults to 1e-5): - The end value of sampling where timesteps decrease progressively from 1 to epsilon. - correct_steps (`int`, defaults to 1): - The number of correction steps performed on a produced sample. - """ - - order = 1 - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 2000, - snr: float = 0.15, - sigma_min: float = 0.01, - sigma_max: float = 1348.0, - sampling_eps: float = 1e-5, - correct_steps: int = 1, - ): - # standard deviation of the initial noise distribution - self.init_noise_sigma = sigma_max - - # setable values - self.timesteps = None - - 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: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - def set_timesteps( - self, num_inference_steps: int, sampling_eps: float = None, device: Union[str, torch.device] = None - ): - """ - Sets the continuous 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. - sampling_eps (`float`, *optional*): - The final timestep value (overrides value given during scheduler instantiation). - device (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - - """ - sampling_eps = sampling_eps if sampling_eps is not None else self.config.sampling_eps - - self.timesteps = torch.linspace(1, sampling_eps, num_inference_steps, device=device) - - def set_sigmas( - self, num_inference_steps: int, sigma_min: float = None, sigma_max: float = None, sampling_eps: float = None - ): - """ - Sets the noise scales used for the diffusion chain (to be run before inference). The sigmas control the weight - of the `drift` and `diffusion` components of the sample update. - - Args: - num_inference_steps (`int`): - The number of diffusion steps used when generating samples with a pre-trained model. - sigma_min (`float`, optional): - The initial noise scale value (overrides value given during scheduler instantiation). - sigma_max (`float`, optional): - The final noise scale value (overrides value given during scheduler instantiation). - sampling_eps (`float`, optional): - The final timestep value (overrides value given during scheduler instantiation). - - """ - sigma_min = sigma_min if sigma_min is not None else self.config.sigma_min - sigma_max = sigma_max if sigma_max is not None else self.config.sigma_max - sampling_eps = sampling_eps if sampling_eps is not None else self.config.sampling_eps - if self.timesteps is None: - self.set_timesteps(num_inference_steps, sampling_eps) - - self.sigmas = sigma_min * (sigma_max / sigma_min) ** (self.timesteps / sampling_eps) - self.discrete_sigmas = torch.exp(torch.linspace(math.log(sigma_min), math.log(sigma_max), num_inference_steps)) - self.sigmas = torch.tensor([sigma_min * (sigma_max / sigma_min) ** t for t in self.timesteps]) - - def get_adjacent_sigma(self, timesteps, t): - return torch.where( - timesteps == 0, - torch.zeros_like(t.to(timesteps.device)), - self.discrete_sigmas[timesteps - 1].to(timesteps.device), - ) - - def step_pred( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - generator: Optional[torch.Generator] = None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - generator (`torch.Generator`, *optional*): - A random number generator. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`. - - Returns: - [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_sde_ve.SdeVeOutput`] is returned, otherwise a tuple - is returned where the first element is the sample tensor. - - """ - if self.timesteps is None: - raise ValueError( - "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler" - ) - - timestep = timestep * torch.ones( - sample.shape[0], device=sample.device - ) # torch.repeat_interleave(timestep, sample.shape[0]) - timesteps = (timestep * (len(self.timesteps) - 1)).long() - - # mps requires indices to be in the same device, so we use cpu as is the default with cuda - timesteps = timesteps.to(self.discrete_sigmas.device) - - sigma = self.discrete_sigmas[timesteps].to(sample.device) - adjacent_sigma = self.get_adjacent_sigma(timesteps, timestep).to(sample.device) - drift = torch.zeros_like(sample) - diffusion = (sigma**2 - adjacent_sigma**2) ** 0.5 - - # equation 6 in the paper: the model_output modeled by the network is grad_x log pt(x) - # also equation 47 shows the analog from SDE models to ancestral sampling methods - diffusion = diffusion.flatten() - while len(diffusion.shape) < len(sample.shape): - diffusion = diffusion.unsqueeze(-1) - drift = drift - diffusion**2 * model_output - - # equation 6: sample noise for the diffusion term of - noise = randn_tensor( - sample.shape, layout=sample.layout, generator=generator, device=sample.device, dtype=sample.dtype - ) - prev_sample_mean = sample - drift # subtract because `dt` is a small negative timestep - # TODO is the variable diffusion the correct scaling term for the noise? - prev_sample = prev_sample_mean + diffusion * noise # add impact of diffusion field g - - if not return_dict: - return (prev_sample, prev_sample_mean) - - return SdeVeOutput(prev_sample=prev_sample, prev_sample_mean=prev_sample_mean) - - def step_correct( - self, - model_output: torch.Tensor, - sample: torch.Tensor, - generator: Optional[torch.Generator] = None, - return_dict: bool = True, - ) -> Union[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. - - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - 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`): - Whether or not to return a [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`. - - Returns: - [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_sde_ve.SdeVeOutput`] is returned, otherwise a tuple - is returned where the first element is the sample tensor. - - """ - if self.timesteps is None: - raise ValueError( - "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler" - ) - - # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z" - # sample noise for correction - noise = randn_tensor(sample.shape, layout=sample.layout, generator=generator).to(sample.device) - - # compute step size from the model_output, the noise, and the snr - grad_norm = torch.norm(model_output.reshape(model_output.shape[0], -1), dim=-1).mean() - noise_norm = torch.norm(noise.reshape(noise.shape[0], -1), dim=-1).mean() - step_size = (self.config.snr * noise_norm / grad_norm) ** 2 * 2 - step_size = step_size * torch.ones(sample.shape[0]).to(sample.device) - # self.repeat_scalar(step_size, sample.shape[0]) - - # compute corrected sample: model_output term and noise term - step_size = step_size.flatten() - while len(step_size.shape) < len(sample.shape): - step_size = step_size.unsqueeze(-1) - prev_sample_mean = sample + step_size * model_output - prev_sample = prev_sample_mean + ((step_size * 2) ** 0.5) * noise - - if not return_dict: - return (prev_sample,) - - return SchedulerOutput(prev_sample=prev_sample) - - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.Tensor, - ) -> torch.Tensor: - # Make sure sigmas and timesteps have the same device and dtype as original_samples - timesteps = timesteps.to(original_samples.device) - sigmas = self.discrete_sigmas.to(original_samples.device)[timesteps] - noise = ( - noise * sigmas[:, None, None, None] - if noise is not None - else torch.randn_like(original_samples) * sigmas[:, None, None, None] - ) - noisy_samples = noise + original_samples - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_sde_ve_flax.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_sde_ve_flax.py deleted file mode 100644 index 09cd081462b3fa5c5d005e3737735e9f34b4a5ac..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_sde_ve_flax.py +++ /dev/null @@ -1,280 +0,0 @@ -# Copyright 2025 Google Brain 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. - -# 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 -import jax.numpy as jnp -from jax import random - -from ..configuration_utils import ConfigMixin, register_to_config -from .scheduling_utils_flax import FlaxSchedulerMixin, FlaxSchedulerOutput, broadcast_to_shape_from_left - - -@flax.struct.dataclass -class ScoreSdeVeSchedulerState: - # setable values - timesteps: Optional[jnp.ndarray] = None - discrete_sigmas: Optional[jnp.ndarray] = None - sigmas: Optional[jnp.ndarray] = None - - @classmethod - def create(cls): - return cls() - - -@dataclass -class FlaxSdeVeOutput(FlaxSchedulerOutput): - """ - Output class for the ScoreSdeVeScheduler's step function output. - - Args: - state (`ScoreSdeVeSchedulerState`): - prev_sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - prev_sample_mean (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)` for images): - Mean averaged `prev_sample`. Same as `prev_sample`, only mean-averaged over previous timesteps. - """ - - state: ScoreSdeVeSchedulerState - prev_sample: jnp.ndarray - prev_sample_mean: Optional[jnp.ndarray] = None - - -class FlaxScoreSdeVeScheduler(FlaxSchedulerMixin, ConfigMixin): - """ - The variance exploding stochastic differential equation (SDE) scheduler. - - For more information, see the original paper: https://huggingface.co/papers/2011.13456 - - [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` - function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. - [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and - [`~SchedulerMixin.from_pretrained`] functions. - - Args: - num_train_timesteps (`int`): number of diffusion steps used to train the model. - snr (`float`): - coefficient weighting the step from the model_output sample (from the network) to the random noise. - sigma_min (`float`): - initial noise scale for sigma sequence in sampling procedure. The minimum sigma should mirror the - distribution of the data. - sigma_max (`float`): maximum value used for the range of continuous timesteps passed into the model. - sampling_eps (`float`): the end value of sampling, where timesteps decrease progressively from 1 to - epsilon. - correct_steps (`int`): number of correction steps performed on a produced sample. - """ - - @property - def has_state(self): - return True - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 2000, - snr: float = 0.15, - sigma_min: float = 0.01, - sigma_max: float = 1348.0, - sampling_eps: float = 1e-5, - correct_steps: int = 1, - ): - pass - - def create_state(self): - state = ScoreSdeVeSchedulerState.create() - return self.set_sigmas( - state, - self.config.num_train_timesteps, - self.config.sigma_min, - self.config.sigma_max, - self.config.sampling_eps, - ) - - def set_timesteps( - 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. - - Args: - state (`ScoreSdeVeSchedulerState`): the `FlaxScoreSdeVeScheduler` state data class instance. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - sampling_eps (`float`, optional): - final timestep value (overrides value given at Scheduler instantiation). - - """ - sampling_eps = sampling_eps if sampling_eps is not None else self.config.sampling_eps - - timesteps = jnp.linspace(1, sampling_eps, num_inference_steps) - return state.replace(timesteps=timesteps) - - def set_sigmas( - self, - state: ScoreSdeVeSchedulerState, - num_inference_steps: int, - sigma_min: float = None, - sigma_max: float = None, - sampling_eps: float = None, - ) -> ScoreSdeVeSchedulerState: - """ - Sets the noise scales used for the diffusion chain. Supporting function to be run before inference. - - The sigmas control the weight of the `drift` and `diffusion` components of sample update. - - Args: - state (`ScoreSdeVeSchedulerState`): the `FlaxScoreSdeVeScheduler` state data class instance. - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - sigma_min (`float`, optional): - initial noise scale value (overrides value given at Scheduler instantiation). - sigma_max (`float`, optional): - final noise scale value (overrides value given at Scheduler instantiation). - sampling_eps (`float`, optional): - final timestep value (overrides value given at Scheduler instantiation). - """ - sigma_min = sigma_min if sigma_min is not None else self.config.sigma_min - sigma_max = sigma_max if sigma_max is not None else self.config.sigma_max - sampling_eps = sampling_eps if sampling_eps is not None else self.config.sampling_eps - if state.timesteps is None: - state = self.set_timesteps(state, num_inference_steps, sampling_eps) - - discrete_sigmas = jnp.exp(jnp.linspace(jnp.log(sigma_min), jnp.log(sigma_max), num_inference_steps)) - sigmas = jnp.array([sigma_min * (sigma_max / sigma_min) ** t for t in state.timesteps]) - - return state.replace(discrete_sigmas=discrete_sigmas, sigmas=sigmas) - - def get_adjacent_sigma(self, state, timesteps, t): - return jnp.where(timesteps == 0, jnp.zeros_like(t), state.discrete_sigmas[timesteps - 1]) - - def step_pred( - self, - state: ScoreSdeVeSchedulerState, - model_output: jnp.ndarray, - timestep: int, - sample: jnp.ndarray, - key: jax.Array, - return_dict: bool = True, - ) -> Union[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). - - Args: - state (`ScoreSdeVeSchedulerState`): the `FlaxScoreSdeVeScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - generator: random number generator. - return_dict (`bool`): option for returning tuple rather than FlaxSdeVeOutput class - - Returns: - [`FlaxSdeVeOutput`] or `tuple`: [`FlaxSdeVeOutput`] if `return_dict` is True, otherwise a `tuple`. When - returning a tuple, the first element is the sample tensor. - - """ - if state.timesteps is None: - raise ValueError( - "`state.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler" - ) - - timestep = timestep * jnp.ones( - sample.shape[0], - ) - timesteps = (timestep * (len(state.timesteps) - 1)).long() - - sigma = state.discrete_sigmas[timesteps] - adjacent_sigma = self.get_adjacent_sigma(state, timesteps, timestep) - drift = jnp.zeros_like(sample) - diffusion = (sigma**2 - adjacent_sigma**2) ** 0.5 - - # equation 6 in the paper: the model_output modeled by the network is grad_x log pt(x) - # also equation 47 shows the analog from SDE models to ancestral sampling methods - diffusion = diffusion.flatten() - diffusion = broadcast_to_shape_from_left(diffusion, sample.shape) - drift = drift - diffusion**2 * model_output - - # equation 6: sample noise for the diffusion term of - key = random.split(key, num=1) - noise = random.normal(key=key, shape=sample.shape) - prev_sample_mean = sample - drift # subtract because `dt` is a small negative timestep - # TODO is the variable diffusion the correct scaling term for the noise? - prev_sample = prev_sample_mean + diffusion * noise # add impact of diffusion field g - - if not return_dict: - return (prev_sample, prev_sample_mean, state) - - return FlaxSdeVeOutput(prev_sample=prev_sample, prev_sample_mean=prev_sample_mean, state=state) - - def step_correct( - self, - state: ScoreSdeVeSchedulerState, - model_output: jnp.ndarray, - sample: jnp.ndarray, - key: jax.Array, - return_dict: bool = True, - ) -> Union[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. - - Args: - state (`ScoreSdeVeSchedulerState`): the `FlaxScoreSdeVeScheduler` state data class instance. - model_output (`jnp.ndarray`): direct output from learned diffusion model. - sample (`jnp.ndarray`): - current instance of sample being created by diffusion process. - generator: random number generator. - return_dict (`bool`): option for returning tuple rather than FlaxSdeVeOutput class - - Returns: - [`FlaxSdeVeOutput`] or `tuple`: [`FlaxSdeVeOutput`] if `return_dict` is True, otherwise a `tuple`. When - returning a tuple, the first element is the sample tensor. - - """ - if state.timesteps is None: - raise ValueError( - "`state.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler" - ) - - # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z" - # sample noise for correction - key = random.split(key, num=1) - noise = random.normal(key=key, shape=sample.shape) - - # compute step size from the model_output, the noise, and the snr - grad_norm = jnp.linalg.norm(model_output) - noise_norm = jnp.linalg.norm(noise) - step_size = (self.config.snr * noise_norm / grad_norm) ** 2 * 2 - step_size = step_size * jnp.ones(sample.shape[0]) - - # compute corrected sample: model_output term and noise term - step_size = step_size.flatten() - step_size = broadcast_to_shape_from_left(step_size, sample.shape) - prev_sample_mean = sample + step_size * model_output - prev_sample = prev_sample_mean + ((step_size * 2) ** 0.5) * noise - - if not return_dict: - return (prev_sample, state) - - return FlaxSdeVeOutput(prev_sample=prev_sample, state=state) - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_tcd.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_tcd.py deleted file mode 100644 index 3fd5c341eca99d72c45e49ea72ac6ba2488dc9c2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_tcd.py +++ /dev/null @@ -1,690 +0,0 @@ -# Copyright 2025 Stanford University 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. - -# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion -# and https://github.com/hojonathanho/diffusion - -import math -from dataclasses import dataclass -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..schedulers.scheduling_utils import SchedulerMixin -from ..utils import BaseOutput, logging -from ..utils.torch_utils import randn_tensor - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class TCDSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_noised_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted noised sample `(x_{s})` based on the model output from the current timestep. - """ - - prev_sample: torch.Tensor - pred_noised_sample: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr -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) - - - Args: - betas (`torch.Tensor`): - the betas that the scheduler is being initialized with. - - Returns: - `torch.Tensor`: rescaled betas with zero terminal SNR - """ - # Convert betas to alphas_bar_sqrt - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - alphas_bar_sqrt = alphas_cumprod.sqrt() - - # Store old values. - alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() - alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() - - # Shift so the last timestep is zero. - alphas_bar_sqrt -= alphas_bar_sqrt_T - - # Scale so the first timestep is back to the old value. - alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - - # Convert alphas_bar_sqrt to betas - alphas_bar = alphas_bar_sqrt**2 # Revert sqrt - alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod - alphas = torch.cat([alphas_bar[0:1], alphas]) - betas = 1 - alphas - - return betas - - -class TCDScheduler(SchedulerMixin, ConfigMixin): - """ - `TCDScheduler` incorporates the `Strategic Stochastic Sampling` introduced by the paper `Trajectory Consistency - Distillation`, extending the original Multistep Consistency Sampling to enable unrestricted trajectory traversal. - - This code is based on the official repo of TCD(https://github.com/jabir-zheng/TCD). - - This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. [`~ConfigMixin`] takes care of storing all config - attributes that are passed in the scheduler's `__init__` function, such as `num_train_timesteps`. They can be - accessed via `scheduler.config.num_train_timesteps`. [`SchedulerMixin`] provides general loading and saving - functionality via the [`SchedulerMixin.save_pretrained`] and [`~SchedulerMixin.from_pretrained`] functions. - - Args: - 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`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - original_inference_steps (`int`, *optional*, defaults to 50): - The default number of inference steps used to generate a linearly-spaced timestep schedule, from which we - will ultimately take `num_inference_steps` evenly spaced timesteps to form the final timestep schedule. - clip_sample (`bool`, defaults to `True`): - Clip the predicted sample for numerical stability. - clip_sample_range (`float`, defaults to 1.0): - The maximum magnitude for sample clipping. Valid only when `clip_sample=True`. - set_alpha_to_one (`bool`, defaults to `True`): - Each diffusion step uses the alphas product value at that step and at the previous one. For the final step - there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, - otherwise it uses the alpha value at step 0. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - 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://imagen.research.google/video/paper.pdf) paper). - thresholding (`bool`, defaults to `False`): - Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such - as Stable Diffusion. - dynamic_thresholding_ratio (`float`, defaults to 0.995): - The ratio for the dynamic thresholding method. Valid only when `thresholding=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 the [Common Diffusion Noise Schedules and - Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. - timestep_scaling (`float`, defaults to 10.0): - The factor the timesteps will be multiplied by when calculating the consistency model boundary conditions - `c_skip` and `c_out`. Increasing this will decrease the approximation error (although the approximation - error at the default of `10.0` is already pretty small). - 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). - """ - - order = 1 - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - beta_start: float = 0.00085, - beta_end: float = 0.012, - beta_schedule: str = "scaled_linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, - original_inference_steps: int = 50, - clip_sample: bool = False, - clip_sample_range: float = 1.0, - set_alpha_to_one: bool = True, - steps_offset: int = 0, - prediction_type: str = "epsilon", - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - sample_max_value: float = 1.0, - timestep_spacing: str = "leading", - timestep_scaling: float = 10.0, - rescale_betas_zero_snr: bool = False, - ): - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - # Rescale for zero SNR - if rescale_betas_zero_snr: - self.betas = rescale_zero_terminal_snr(self.betas) - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - # At every step in ddim, we are looking into the previous alphas_cumprod - # For the final step, there is no previous alphas_cumprod because we are already at 0 - # `set_alpha_to_one` decides whether we set this parameter simply to one or - # whether we use the final alpha of the "non-previous" one. - self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0] - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # setable values - self.num_inference_steps = None - self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)) - self.custom_timesteps = False - - self._step_index = None - self._begin_index = None - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep - 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() - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - 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 - - @property - def step_index(self): - 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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: - """ - 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 current timestep in the diffusion chain. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - # Copied from diffusers.schedulers.scheduling_ddim.DDIMScheduler._get_variance - def _get_variance(self, timestep, prev_timestep): - 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 - beta_prod_t_prev = 1 - alpha_prod_t_prev - - variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev) - - return variance - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - """ - "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the - prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by - s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing - pixels from saturation at each step. We find that dynamic thresholding results in significantly better - photorealism as well as better image-text alignment, especially when using very large guidance weights." - - https://huggingface.co/papers/2205.11487 - """ - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half - - # Flatten sample for doing quantile calculation along each image - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - - abs_sample = sample.abs() # "a certain percentile absolute pixel value" - - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] - s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 - sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" - - sample = sample.reshape(batch_size, channels, *remaining_dims) - sample = sample.to(dtype) - - return sample - - 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, - strength: float = 1.0, - ): - """ - Sets the discrete timesteps used for the diffusion chain (to be run before inference). - - Args: - 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. - original_inference_steps (`int`, *optional*): - The original number of inference steps, which will be used to generate a linearly-spaced timestep - 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*): - 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`. - strength (`float`, *optional*, defaults to 1.0): - Used to determine the number of timesteps used for inference when using img2img, inpaint, etc. - """ - # 0. Check inputs - if num_inference_steps is None and timesteps is None: - raise ValueError("Must pass exactly one of `num_inference_steps` or `custom_timesteps`.") - - if num_inference_steps is not None and timesteps is not None: - raise ValueError("Can only pass one of `num_inference_steps` or `custom_timesteps`.") - - # 1. Calculate the TCD original training/distillation timestep schedule. - original_steps = ( - original_inference_steps if original_inference_steps is not None else self.config.original_inference_steps - ) - - if original_inference_steps is None: - # default option, timesteps align with discrete inference steps - if original_steps > self.config.num_train_timesteps: - raise ValueError( - f"`original_steps`: {original_steps} cannot be larger than `self.config.train_timesteps`:" - f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle" - f" maximal {self.config.num_train_timesteps} timesteps." - ) - # TCD Timesteps Setting - # The skipping step parameter k from the paper. - k = self.config.num_train_timesteps // original_steps - # TCD Training/Distillation Steps Schedule - tcd_origin_timesteps = np.asarray(list(range(1, int(original_steps * strength) + 1))) * k - 1 - else: - # customised option, sampled timesteps can be any arbitrary value - tcd_origin_timesteps = np.asarray(list(range(0, int(self.config.num_train_timesteps * strength)))) - - # 2. Calculate the TCD inference timestep schedule. - if timesteps is not None: - # 2.1 Handle custom timestep schedules. - train_timesteps = set(tcd_origin_timesteps) - non_train_timesteps = [] - for i in range(1, len(timesteps)): - if timesteps[i] >= timesteps[i - 1]: - raise ValueError("`custom_timesteps` must be in descending order.") - - if timesteps[i] not in train_timesteps: - non_train_timesteps.append(timesteps[i]) - - if timesteps[0] >= self.config.num_train_timesteps: - raise ValueError( - f"`timesteps` must start before `self.config.train_timesteps`: {self.config.num_train_timesteps}." - ) - - # Raise warning if timestep schedule does not start with self.config.num_train_timesteps - 1 - if strength == 1.0 and timesteps[0] != self.config.num_train_timesteps - 1: - logger.warning( - f"The first timestep on the custom timestep schedule is {timesteps[0]}, not" - f" `self.config.num_train_timesteps - 1`: {self.config.num_train_timesteps - 1}. You may get" - f" unexpected results when using this timestep schedule." - ) - - # Raise warning if custom timestep schedule contains timesteps not on original timestep schedule - if non_train_timesteps: - logger.warning( - f"The custom timestep schedule contains the following timesteps which are not on the original" - f" training/distillation timestep schedule: {non_train_timesteps}. You may get unexpected results" - f" when using this timestep schedule." - ) - - # Raise warning if custom timestep schedule is longer than original_steps - if original_steps is not None: - if len(timesteps) > original_steps: - logger.warning( - f"The number of timesteps in the custom timestep schedule is {len(timesteps)}, which exceeds the" - f" the length of the timestep schedule used for training: {original_steps}. You may get some" - f" unexpected results when using this timestep schedule." - ) - else: - if len(timesteps) > self.config.num_train_timesteps: - logger.warning( - f"The number of timesteps in the custom timestep schedule is {len(timesteps)}, which exceeds the" - f" the length of the timestep schedule used for training: {self.config.num_train_timesteps}. You may get some" - f" unexpected results when using this timestep schedule." - ) - - timesteps = np.array(timesteps, dtype=np.int64) - self.num_inference_steps = len(timesteps) - self.custom_timesteps = True - - # Apply strength (e.g. for img2img pipelines) (see StableDiffusionImg2ImgPipeline.get_timesteps) - init_timestep = min(int(self.num_inference_steps * strength), self.num_inference_steps) - t_start = max(self.num_inference_steps - init_timestep, 0) - timesteps = timesteps[t_start * self.order :] - # TODO: also reset self.num_inference_steps? - else: - # 2.2 Create the "standard" TCD inference timestep schedule. - if num_inference_steps > self.config.num_train_timesteps: - raise ValueError( - f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:" - f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle" - f" maximal {self.config.num_train_timesteps} timesteps." - ) - - if original_steps is not None: - skipping_step = len(tcd_origin_timesteps) // num_inference_steps - - if skipping_step < 1: - raise ValueError( - f"The combination of `original_steps x strength`: {original_steps} x {strength} is smaller than `num_inference_steps`: {num_inference_steps}. Make sure to either reduce `num_inference_steps` to a value smaller than {int(original_steps * strength)} or increase `strength` to a value higher than {float(num_inference_steps / original_steps)}." - ) - - self.num_inference_steps = num_inference_steps - - if original_steps is not None: - if num_inference_steps > original_steps: - raise ValueError( - f"`num_inference_steps`: {num_inference_steps} cannot be larger than `original_inference_steps`:" - f" {original_steps} because the final timestep schedule will be a subset of the" - f" `original_inference_steps`-sized initial timestep schedule." - ) - else: - if num_inference_steps > self.config.num_train_timesteps: - raise ValueError( - f"`num_inference_steps`: {num_inference_steps} cannot be larger than `num_train_timesteps`:" - f" {self.config.num_train_timesteps} because the final timestep schedule will be a subset of the" - f" `num_train_timesteps`-sized initial timestep schedule." - ) - - # TCD Inference Steps Schedule - tcd_origin_timesteps = tcd_origin_timesteps[::-1].copy() - # Select (approximately) evenly spaced indices from tcd_origin_timesteps. - inference_indices = np.linspace(0, len(tcd_origin_timesteps), num=num_inference_steps, endpoint=False) - inference_indices = np.floor(inference_indices).astype(np.int64) - timesteps = tcd_origin_timesteps[inference_indices] - - self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.long) - - self._step_index = None - self._begin_index = None - - def step( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - eta: float = 0.3, - generator: Optional[torch.Generator] = None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - eta (`float`): - A stochastic parameter (referred to as `gamma` in the paper) used to control the stochasticity in every - step. When eta = 0, it represents deterministic sampling, whereas eta = 1 indicates full stochastic - sampling. - generator (`torch.Generator`, *optional*): - A random number generator. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~schedulers.scheduling_tcd.TCDSchedulerOutput`] or `tuple`. - Returns: - [`~schedulers.scheduling_utils.TCDSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_tcd.TCDSchedulerOutput`] is returned, otherwise a - tuple is returned where the first element is the sample tensor. - """ - 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 self.step_index is None: - self._init_step_index(timestep) - - assert 0 <= eta <= 1.0, "gamma must be less than or equal to 1.0" - - # 1. get previous step value - prev_step_index = self.step_index + 1 - if prev_step_index < len(self.timesteps): - prev_timestep = self.timesteps[prev_step_index] - else: - prev_timestep = torch.tensor(0) - - timestep_s = torch.floor((1 - eta) * prev_timestep).to(dtype=torch.long) - - # 2. compute alphas, betas - alpha_prod_t = self.alphas_cumprod[timestep] - beta_prod_t = 1 - alpha_prod_t - - alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod - - alpha_prod_s = self.alphas_cumprod[timestep_s] - beta_prod_s = 1 - alpha_prod_s - - # 3. Compute the predicted noised sample x_s based on the model parameterization - if self.config.prediction_type == "epsilon": # noise-prediction - pred_original_sample = (sample - beta_prod_t.sqrt() * model_output) / alpha_prod_t.sqrt() - pred_epsilon = model_output - pred_noised_sample = alpha_prod_s.sqrt() * pred_original_sample + beta_prod_s.sqrt() * pred_epsilon - elif self.config.prediction_type == "sample": # x-prediction - pred_original_sample = model_output - pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) - pred_noised_sample = alpha_prod_s.sqrt() * pred_original_sample + beta_prod_s.sqrt() * pred_epsilon - elif self.config.prediction_type == "v_prediction": # v-prediction - pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output - pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample - pred_noised_sample = alpha_prod_s.sqrt() * pred_original_sample + beta_prod_s.sqrt() * pred_epsilon - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or" - " `v_prediction` for `TCDScheduler`." - ) - - # 4. Sample and inject noise z ~ N(0, I) for MultiStep Inference - # Noise is not used on the final timestep of the timestep schedule. - # This also means that noise is not used for one-step sampling. - # Eta (referred to as "gamma" in the paper) was introduced to control the stochasticity in every step. - # When eta = 0, it represents deterministic sampling, whereas eta = 1 indicates full stochastic sampling. - if eta > 0: - if self.step_index != self.num_inference_steps - 1: - noise = randn_tensor( - model_output.shape, generator=generator, device=model_output.device, dtype=pred_noised_sample.dtype - ) - prev_sample = (alpha_prod_t_prev / alpha_prod_s).sqrt() * pred_noised_sample + ( - 1 - alpha_prod_t_prev / alpha_prod_s - ).sqrt() * noise - else: - prev_sample = pred_noised_sample - else: - prev_sample = pred_noised_sample - - # upon completion increase step index by one - self._step_index += 1 - - if not return_dict: - return (prev_sample, pred_noised_sample) - - return TCDSchedulerOutput(prev_sample=prev_sample, pred_noised_sample=pred_noised_sample) - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as original_samples - # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement - # for the subsequent add_noise calls - self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype) - timesteps = timesteps.to(original_samples.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(original_samples.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise - return noisy_samples - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity - def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as sample - self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype) - timesteps = timesteps.to(sample.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(sample.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample - return velocity - - def __len__(self): - return self.config.num_train_timesteps - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.previous_timestep - def previous_timestep(self, timestep): - if self.custom_timesteps or self.num_inference_steps: - index = (self.timesteps == timestep).nonzero(as_tuple=True)[0][0] - if index == self.timesteps.shape[0] - 1: - prev_t = torch.tensor(-1) - else: - prev_t = self.timesteps[index + 1] - else: - prev_t = timestep - 1 - return prev_t diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_unclip.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_unclip.py deleted file mode 100644 index d78efabfbc57ba215b1b4d3d0348e3db85543e08..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_unclip.py +++ /dev/null @@ -1,355 +0,0 @@ -# Copyright 2025 Kakao Brain 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 Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput -from ..utils.torch_utils import randn_tensor -from .scheduling_utils import SchedulerMixin - - -@dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->UnCLIP -class UnCLIPSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. - """ - - prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -class UnCLIPScheduler(SchedulerMixin, ConfigMixin): - """ - NOTE: do not use this scheduler. The DDPM scheduler has been updated to support the changes made here. This - scheduler will be removed and replaced with DDPM. - - This is a modified DDPM Scheduler specifically for the karlo unCLIP model. - - This scheduler has some minor variations in how it calculates the learned range variance and dynamically - re-calculates betas based off the timesteps it is skipping. - - The scheduler also uses a slightly different step ratio when computing timesteps to use for inference. - - See [`~DDPMScheduler`] for more information on DDPM scheduling - - Args: - num_train_timesteps (`int`): number of diffusion steps used to train the model. - variance_type (`str`): - options to clip the variance used when adding noise to the denoised sample. Choose from `fixed_small_log` - or `learned_range`. - clip_sample (`bool`, default `True`): - option to clip predicted sample between `-clip_sample_range` and `clip_sample_range` for numerical - stability. - clip_sample_range (`float`, default `1.0`): - The range to clip the sample between. See `clip_sample`. - prediction_type (`str`, default `epsilon`, optional): - prediction type of the scheduler function, one of `epsilon` (predicting the noise of the diffusion process) - or `sample` (directly predicting the noisy sample`) - """ - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - variance_type: str = "fixed_small_log", - clip_sample: bool = True, - clip_sample_range: Optional[float] = 1.0, - prediction_type: str = "epsilon", - beta_schedule: str = "squaredcos_cap_v2", - ): - if beta_schedule != "squaredcos_cap_v2": - raise ValueError("UnCLIPScheduler only supports `beta_schedule`: 'squaredcos_cap_v2'") - - self.betas = betas_for_alpha_bar(num_train_timesteps) - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - self.one = torch.tensor(1.0) - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - # setable values - self.num_inference_steps = None - self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy()) - - self.variance_type = variance_type - - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: - """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. - - Args: - sample (`torch.Tensor`): input sample - timestep (`int`, optional): current timestep - - Returns: - `torch.Tensor`: scaled input sample - """ - return sample - - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): - """ - Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference. - - Note that this scheduler uses a slightly different step ratio than the other diffusers schedulers. The - different step ratio is to mimic the original karlo implementation and does not affect the quality or accuracy - of the results. - - Args: - num_inference_steps (`int`): - the number of diffusion steps used when generating samples with a pre-trained model. - """ - self.num_inference_steps = num_inference_steps - step_ratio = (self.config.num_train_timesteps - 1) / (self.num_inference_steps - 1) - timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64) - self.timesteps = torch.from_numpy(timesteps).to(device) - - def _get_variance(self, t, prev_timestep=None, predicted_variance=None, variance_type=None): - if prev_timestep is None: - prev_timestep = t - 1 - - alpha_prod_t = self.alphas_cumprod[t] - alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.one - beta_prod_t = 1 - alpha_prod_t - beta_prod_t_prev = 1 - alpha_prod_t_prev - - if prev_timestep == t - 1: - beta = self.betas[t] - else: - beta = 1 - alpha_prod_t / alpha_prod_t_prev - - # For t > 0, compute predicted variance βt (see formula (6) and (7) from https://huggingface.co/papers/2006.11239) - # and sample from it to get previous sample - # x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample - variance = beta_prod_t_prev / beta_prod_t * beta - - if variance_type is None: - variance_type = self.config.variance_type - - # hacks - were probably added for training stability - if variance_type == "fixed_small_log": - variance = torch.log(torch.clamp(variance, min=1e-20)) - variance = torch.exp(0.5 * variance) - elif variance_type == "learned_range": - # NOTE difference with DDPM scheduler - min_log = variance.log() - max_log = beta.log() - - frac = (predicted_variance + 1) / 2 - variance = frac * max_log + (1 - frac) * min_log - - return variance - - def step( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - prev_timestep: Optional[int] = None, - generator=None, - return_dict: bool = True, - ) -> Union[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). - - Args: - model_output (`torch.Tensor`): direct output from learned diffusion model. - timestep (`int`): current discrete timestep in the diffusion chain. - sample (`torch.Tensor`): - current instance of sample being created by diffusion process. - prev_timestep (`int`, *optional*): The previous timestep to predict the previous sample at. - Used to dynamically compute beta. If not given, `t-1` is used and the pre-computed beta is used. - generator: random number generator. - return_dict (`bool`): option for returning tuple rather than UnCLIPSchedulerOutput class - - Returns: - [`~schedulers.scheduling_utils.UnCLIPSchedulerOutput`] or `tuple`: - [`~schedulers.scheduling_utils.UnCLIPSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. When - returning a tuple, the first element is the sample tensor. - - """ - t = timestep - - if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type == "learned_range": - model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1) - else: - predicted_variance = None - - # 1. compute alphas, betas - if prev_timestep is None: - prev_timestep = t - 1 - - alpha_prod_t = self.alphas_cumprod[t] - alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.one - beta_prod_t = 1 - alpha_prod_t - beta_prod_t_prev = 1 - alpha_prod_t_prev - - if prev_timestep == t - 1: - beta = self.betas[t] - alpha = self.alphas[t] - else: - beta = 1 - alpha_prod_t / alpha_prod_t_prev - alpha = 1 - beta - - # 2. compute predicted original sample from predicted noise also called - # "predicted x_0" of formula (15) from https://huggingface.co/papers/2006.11239 - if self.config.prediction_type == "epsilon": - pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - elif self.config.prediction_type == "sample": - pred_original_sample = model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon` or `sample`" - " for the UnCLIPScheduler." - ) - - # 3. Clip "predicted x_0" - if self.config.clip_sample: - pred_original_sample = torch.clamp( - pred_original_sample, -self.config.clip_sample_range, self.config.clip_sample_range - ) - - # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t - # See formula (7) from https://huggingface.co/papers/2006.11239 - pred_original_sample_coeff = (alpha_prod_t_prev ** (0.5) * beta) / beta_prod_t - current_sample_coeff = alpha ** (0.5) * beta_prod_t_prev / beta_prod_t - - # 5. Compute predicted previous sample µ_t - # See formula (7) from https://huggingface.co/papers/2006.11239 - pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample - - # 6. Add noise - variance = 0 - if t > 0: - variance_noise = randn_tensor( - model_output.shape, dtype=model_output.dtype, generator=generator, device=model_output.device - ) - - variance = self._get_variance( - t, - predicted_variance=predicted_variance, - prev_timestep=prev_timestep, - ) - - if self.variance_type == "fixed_small_log": - variance = variance - elif self.variance_type == "learned_range": - variance = (0.5 * variance).exp() - else: - raise ValueError( - f"variance_type given as {self.variance_type} must be one of `fixed_small_log` or `learned_range`" - " for the UnCLIPScheduler." - ) - - variance = variance * variance_noise - - pred_prev_sample = pred_prev_sample + variance - - if not return_dict: - return ( - pred_prev_sample, - pred_original_sample, - ) - - return UnCLIPSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample) - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # Make sure alphas_cumprod and timestep have same device and dtype as original_samples - # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement - # for the subsequent add_noise calls - self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device) - alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype) - timesteps = timesteps.to(original_samples.device) - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - while len(sqrt_alpha_prod.shape) < len(original_samples.shape): - sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) - - noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise - return noisy_samples diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_unipc_multistep.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_unipc_multistep.py deleted file mode 100644 index 8b1f699b101a00c50787f1f0ac04d5f062a3b922..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_unipc_multistep.py +++ /dev/null @@ -1,1076 +0,0 @@ -# Copyright 2025 TSAIL 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. - -# DISCLAIMER: check https://huggingface.co/papers/2302.04867 and https://github.com/wl-zhao/UniPC for more info -# 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, Optional, Tuple, Union - -import numpy as np -import torch - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import deprecate, is_scipy_available -from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput - - -if is_scipy_available(): - import scipy.stats - - -# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar -def betas_for_alpha_bar( - num_diffusion_timesteps, - max_beta=0.999, - alpha_transform_type="cosine", -): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. - Choose from `cosine` or `exp` - - Returns: - betas (`np.ndarray`): the betas used by the scheduler to step the model outputs - """ - if alpha_transform_type == "cosine": - - def alpha_bar_fn(t): - return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 - - elif alpha_transform_type == "exp": - - def alpha_bar_fn(t): - return math.exp(t * -12.0) - - else: - raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) - return torch.tensor(betas, dtype=torch.float32) - - -# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr -def rescale_zero_terminal_snr(betas): - """ - 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. - - Returns: - `torch.Tensor`: rescaled betas with zero terminal SNR - """ - # Convert betas to alphas_bar_sqrt - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - alphas_bar_sqrt = alphas_cumprod.sqrt() - - # Store old values. - alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() - alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() - - # Shift so the last timestep is zero. - alphas_bar_sqrt -= alphas_bar_sqrt_T - - # Scale so the first timestep is back to the old value. - alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - - # Convert alphas_bar_sqrt to betas - alphas_bar = alphas_bar_sqrt**2 # Revert sqrt - alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod - alphas = torch.cat([alphas_bar[0:1], alphas]) - betas = 1 - alphas - - return betas - - -class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin): - """ - `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models. - - 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`, 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`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, *optional*): - Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - solver_order (`int`, default `2`): - The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1` - due to the UniC. 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://imagen.research.google/video/paper.pdf) paper). - thresholding (`bool`, defaults to `False`): - Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such - as Stable Diffusion. - dynamic_thresholding_ratio (`float`, defaults to 0.995): - The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. - sample_max_value (`float`, defaults to 1.0): - The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`. - predict_x0 (`bool`, defaults to `True`): - Whether to use the updating algorithm on the predicted x0. - solver_type (`str`, default `bh2`): - Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2` - otherwise. - lower_order_final (`bool`, default `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. - disable_corrector (`list`, default `[]`): - Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)` - and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is - usually disabled during the first few steps. - solver_p (`SchedulerMixin`, default `None`): - Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`. - use_karras_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, - the sigmas are determined according to a sequence of noise levels {σi}. - use_exponential_sigmas (`bool`, *optional*, defaults to `False`): - Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. - 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. - 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. - steps_offset (`int`, defaults to 0): - An offset added to the inference steps, as required by some model families. - 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. - 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). - """ - - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - 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, - solver_order: int = 2, - prediction_type: str = "epsilon", - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - sample_max_value: float = 1.0, - predict_x0: bool = True, - solver_type: str = "bh2", - lower_order_final: bool = True, - disable_corrector: List[int] = [], - solver_p: 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, - timestep_spacing: str = "linspace", - steps_offset: int = 0, - final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min" - rescale_betas_zero_snr: 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: - raise ValueError( - "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." - ) - if trained_betas is not None: - self.betas = torch.tensor(trained_betas, dtype=torch.float32) - elif beta_schedule == "linear": - 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 - elif beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - self.betas = betas_for_alpha_bar(num_train_timesteps) - else: - raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") - - if rescale_betas_zero_snr: - self.betas = rescale_zero_terminal_snr(self.betas) - - self.alphas = 1.0 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) - - if rescale_betas_zero_snr: - # Close to 0 without being 0 so first sigma is not inf - # FP16 smallest positive subnormal works well here - self.alphas_cumprod[-1] = 2**-24 - - # Currently we only support VP-type noise schedule - self.alpha_t = torch.sqrt(self.alphas_cumprod) - self.sigma_t = torch.sqrt(1 - self.alphas_cumprod) - self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t) - self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5 - - # standard deviation of the initial noise distribution - self.init_noise_sigma = 1.0 - - 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 - # setable values - self.num_inference_steps = None - timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy() - self.timesteps = torch.from_numpy(timesteps) - 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 - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - @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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_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 set_timesteps(self, num_inference_steps: int, device: Union[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 (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - """ - # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if self.config.timestep_spacing == "linspace": - timesteps = ( - np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps + 1) - .round()[::-1][:-1] - .copy() - .astype(np.int64) - ) - elif self.config.timestep_spacing == "leading": - step_ratio = self.config.num_train_timesteps // (num_inference_steps + 1) - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = (np.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int64) - timesteps += self.config.steps_offset - elif self.config.timestep_spacing == "trailing": - step_ratio = self.config.num_train_timesteps / num_inference_steps - # creates integer timesteps by multiplying by ratio - # casting to int to avoid issues when num_inference_step is power of 3 - timesteps = np.arange(self.config.num_train_timesteps, 0, -step_ratio).round().copy().astype(np.int64) - timesteps -= 1 - else: - raise ValueError( - 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: - log_sigmas = np.log(sigmas) - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round() - if self.config.final_sigmas_type == "sigma_min": - sigma_last = sigmas[-1] - elif self.config.final_sigmas_type == "zero": - sigma_last = 0 - else: - raise ValueError( - f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" - ) - sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) - elif self.config.use_exponential_sigmas: - log_sigmas = np.log(sigmas) - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - if self.config.final_sigmas_type == "sigma_min": - sigma_last = sigmas[-1] - elif self.config.final_sigmas_type == "zero": - sigma_last = 0 - else: - raise ValueError( - f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" - ) - sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) - elif self.config.use_beta_sigmas: - log_sigmas = np.log(sigmas) - sigmas = np.flip(sigmas).copy() - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]) - if self.config.final_sigmas_type == "sigma_min": - sigma_last = sigmas[-1] - elif self.config.final_sigmas_type == "zero": - sigma_last = 0 - else: - raise ValueError( - f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" - ) - 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() - timesteps = (sigmas * self.config.num_train_timesteps).copy() - if self.config.final_sigmas_type == "sigma_min": - sigma_last = sigmas[-1] - elif self.config.final_sigmas_type == "zero": - sigma_last = 0 - else: - raise ValueError( - f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" - ) - sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) - else: - 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 - elif self.config.final_sigmas_type == "zero": - sigma_last = 0 - else: - raise ValueError( - f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" - ) - sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) - - self.sigmas = torch.from_numpy(sigmas) - self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) - - self.num_inference_steps = len(timesteps) - - self.model_outputs = [ - None, - ] * self.config.solver_order - self.lower_order_nums = 0 - self.last_sample = None - if self.solver_p: - self.solver_p.set_timesteps(self.num_inference_steps, device=device) - - # add an index counter for schedulers that allow duplicated timesteps - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - - # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - """ - "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the - prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by - s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing - pixels from saturation at each step. We find that dynamic thresholding results in significantly better - photorealism as well as better image-text alignment, especially when using very large guidance weights." - - https://huggingface.co/papers/2205.11487 - """ - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half - - # Flatten sample for doing quantile calculation along each image - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - - abs_sample = sample.abs() # "a certain percentile absolute pixel value" - - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp( - s, min=1, max=self.config.sample_max_value - ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] - s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 - sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" - - sample = sample.reshape(batch_size, channels, *remaining_dims) - sample = sample.to(dtype) - - return sample - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma - log_sigma = np.log(np.maximum(sigma, 1e-10)) - - # get distribution - dists = log_sigma - log_sigmas[:, np.newaxis] - - # get sigmas range - low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) - high_idx = low_idx + 1 - - low = log_sigmas[low_idx] - high = log_sigmas[high_idx] - - # interpolate sigmas - w = (low - log_sigma) / (low - high) - w = np.clip(w, 0, 1) - - # transform interpolation to time range - t = (1 - w) * low_idx + w * high_idx - t = t.reshape(sigma.shape) - return t - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t - def _sigma_to_alpha_sigma_t(self, sigma): - if self.config.use_flow_sigmas: - alpha_t = 1 - sigma - sigma_t = sigma - else: - alpha_t = 1 / ((sigma**2 + 1) ** 0.5) - sigma_t = sigma * alpha_t - - 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: - """Constructs the noise schedule of Karras et al. (2022).""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - rho = 7.0 # 7.0 is the value used in the paper - ramp = np.linspace(0, 1, num_inference_steps) - min_inv_rho = sigma_min ** (1 / rho) - max_inv_rho = sigma_max ** (1 / rho) - sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Constructs an exponential noise schedule.""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) - return sigmas - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 - ) -> torch.Tensor: - """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if hasattr(self.config, "sigma_min"): - sigma_min = self.config.sigma_min - else: - sigma_min = None - - if hasattr(self.config, "sigma_max"): - sigma_max = self.config.sigma_max - else: - sigma_max = None - - sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() - sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() - - sigmas = np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return sigmas - - def convert_model_output( - self, - model_output: torch.Tensor, - *args, - sample: 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`", - ) - - 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": - sigma_t = self.sigmas[self.step_index] - 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, - **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 - - sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index] - 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, - **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 - - sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[self.step_index - 1] - 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 - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep - def index_for_timestep(self, timestep, schedule_timesteps=None): - if schedule_timesteps is None: - schedule_timesteps = self.timesteps - - index_candidates = (schedule_timesteps == timestep).nonzero() - - if len(index_candidates) == 0: - step_index = len(self.timesteps) - 1 - # 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) - elif len(index_candidates) > 1: - step_index = index_candidates[1].item() - else: - step_index = index_candidates[0].item() - - return step_index - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index - def _init_step_index(self, timestep): - """ - Initialize the step_index counter for the scheduler. - """ - - 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( - self, - model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], - sample: torch.Tensor, - return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: - """ - Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with - the multistep UniPC. - - Args: - model_output (`torch.Tensor`): - The direct output from 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. - return_dict (`bool`): - 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( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - if self.step_index is None: - self._init_step_index(timestep) - - use_corrector = ( - self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None - ) - - model_output_convert = self.convert_model_output(model_output, sample=sample) - 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, - ) - - 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, - ) - - if self.lower_order_nums < self.config.solver_order: - self.lower_order_nums += 1 - - # upon completion increase step index by one - self._step_index += 1 - - if not return_dict: - return (prev_sample,) - - return SchedulerOutput(prev_sample=prev_sample) - - def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: - """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. - - Args: - sample (`torch.Tensor`): - The input sample. - - Returns: - `torch.Tensor`: - A scaled input sample. - """ - return sample - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise - def add_noise( - self, - original_samples: torch.Tensor, - noise: torch.Tensor, - timesteps: torch.IntTensor, - ) -> torch.Tensor: - # 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): - # mps does not support float64 - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - # add_noise is called after first denoising step (for inpainting) - step_indices = [self.step_index] * timesteps.shape[0] - else: - # add noise is called before first denoising step to create initial latent(img2img) - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - noisy_samples = alpha_t * original_samples + sigma_t * noise - return noisy_samples - - def __len__(self): - return self.config.num_train_timesteps diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_utils.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_utils.py deleted file mode 100644 index 61d3e5a22fb07f41faae3afad47269d409aa0fc2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_utils.py +++ /dev/null @@ -1,194 +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 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 -from typing_extensions import Self - -from ..utils import BaseOutput, PushToHubMixin - - -SCHEDULER_CONFIG_NAME = "scheduler_config.json" - - -# NOTE: We make this type an enum because it simplifies usage in docs and prevents -# circular imports when used for `_compatibles` within the schedulers module. -# When it's used as a type in pipelines, it really is a Union because the actual -# scheduler instance is passed in. -class KarrasDiffusionSchedulers(Enum): - DDIMScheduler = 1 - DDPMScheduler = 2 - PNDMScheduler = 3 - LMSDiscreteScheduler = 4 - EulerDiscreteScheduler = 5 - HeunDiscreteScheduler = 6 - EulerAncestralDiscreteScheduler = 7 - DPMSolverMultistepScheduler = 8 - DPMSolverSinglestepScheduler = 9 - KDPM2DiscreteScheduler = 10 - KDPM2AncestralDiscreteScheduler = 11 - DEISMultistepScheduler = 12 - UniPCMultistepScheduler = 13 - DPMSolverSDEScheduler = 14 - EDMEulerScheduler = 15 - - -AysSchedules = { - "StableDiffusionTimesteps": [999, 850, 736, 645, 545, 455, 343, 233, 124, 24], - "StableDiffusionSigmas": [14.615, 6.475, 3.861, 2.697, 1.886, 1.396, 0.963, 0.652, 0.399, 0.152, 0.0], - "StableDiffusionXLTimesteps": [999, 845, 730, 587, 443, 310, 193, 116, 53, 13], - "StableDiffusionXLSigmas": [14.615, 6.315, 3.771, 2.181, 1.342, 0.862, 0.555, 0.380, 0.234, 0.113, 0.0], - "StableDiffusionVideoSigmas": [700.00, 54.5, 15.886, 7.977, 4.248, 1.789, 0.981, 0.403, 0.173, 0.034, 0.0], -} - - -@dataclass -class SchedulerOutput(BaseOutput): - """ - Base class for the output of a scheduler's `step` function. - - Args: - prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - """ - - prev_sample: torch.Tensor - - -class SchedulerMixin(PushToHubMixin): - """ - Base class for all schedulers. - - [`SchedulerMixin`] contains common functions shared by all schedulers such as general loading and saving - functionalities. - - [`ConfigMixin`] takes care of storing the configuration attributes (like `num_train_timesteps`) that are passed to - 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 - class. Use [`~ConfigMixin.from_config`] to load a different compatible scheduler class (should be overridden - by parent class). - """ - - config_name = SCHEDULER_CONFIG_NAME - _compatibles = [] - has_compatibles = True - - @classmethod - @validate_hf_hub_args - def from_pretrained( - cls, - pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None, - subfolder: Optional[str] = None, - return_unused_kwargs=False, - **kwargs, - ) -> Self: - r""" - Instantiate a scheduler from a pre-defined JSON configuration file in a local directory or Hub repository. - - Parameters: - pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): - Can be either: - - - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on - the Hub. - - A path to a *directory* (for example `./my_model_directory`) containing the scheduler - configuration saved with [`~SchedulerMixin.save_pretrained`]. - subfolder (`str`, *optional*): - 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*): - 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*): - 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`): - Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - local_files_only(`bool`, *optional*, defaults to `False`): - Whether to only load local model weights and configuration files or not. If set to `True`, the model - won't be downloaded from the Hub. - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from - `diffusers-cli login` (stored in `~/.huggingface`) is used. - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier - allowed by Git. - - - - To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with - `huggingface-cli login`. You can also activate the special - ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a - firewalled environment. - - - - """ - config, kwargs, commit_hash = cls.load_config( - pretrained_model_name_or_path=pretrained_model_name_or_path, - subfolder=subfolder, - return_unused_kwargs=True, - return_commit_hash=True, - **kwargs, - ) - 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): - """ - Save a scheduler configuration object to a directory so that it can be reloaded using the - [`~SchedulerMixin.from_pretrained`] class method. - - Args: - save_directory (`str` or `os.PathLike`): - Directory where the configuration JSON file will be saved (will be created if it does not exist). - push_to_hub (`bool`, *optional*, defaults to `False`): - 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*): - 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) - - @property - def compatibles(self): - """ - Returns all schedulers that are compatible with this scheduler - - Returns: - `List[SchedulerMixin]`: List of compatible schedulers - """ - return self._get_compatibles() - - @classmethod - def _get_compatibles(cls): - compatible_classes_str = list(set([cls.__name__] + cls._compatibles)) - diffusers_library = importlib.import_module(__name__.split(".")[0]) - compatible_classes = [ - getattr(diffusers_library, c) for c in compatible_classes_str if hasattr(diffusers_library, c) - ] - return compatible_classes diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_utils_flax.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_utils_flax.py deleted file mode 100644 index abcde6c386654a764538ac476e30af2d0e472799..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_utils_flax.py +++ /dev/null @@ -1,291 +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 math -import os -from dataclasses import dataclass -from enum import Enum -from typing import Optional, Tuple, Union - -import flax -import jax.numpy as jnp -from huggingface_hub.utils import validate_hf_hub_args - -from ..utils import BaseOutput, PushToHubMixin - - -SCHEDULER_CONFIG_NAME = "scheduler_config.json" - - -# NOTE: We make this type an enum because it simplifies usage in docs and prevents -# circular imports when used for `_compatibles` within the schedulers module. -# When it's used as a type in pipelines, it really is a Union because the actual -# scheduler instance is passed in. -class FlaxKarrasDiffusionSchedulers(Enum): - FlaxDDIMScheduler = 1 - FlaxDDPMScheduler = 2 - FlaxPNDMScheduler = 3 - FlaxLMSDiscreteScheduler = 4 - FlaxDPMSolverMultistepScheduler = 5 - FlaxEulerDiscreteScheduler = 6 - - -@dataclass -class FlaxSchedulerOutput(BaseOutput): - """ - Base class for the scheduler's step function output. - - Args: - prev_sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - """ - - prev_sample: jnp.ndarray - - -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 - `from_config` can be used from a class different than the one used to save the config (should be overridden - by parent class). - """ - - config_name = SCHEDULER_CONFIG_NAME - ignore_for_config = ["dtype"] - _compatibles = [] - has_compatibles = True - - @classmethod - @validate_hf_hub_args - def from_pretrained( - cls, - pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None, - subfolder: Optional[str] = None, - return_unused_kwargs=False, - **kwargs, - ): - r""" - Instantiate a Scheduler class from a pre-defined JSON-file. - - Parameters: - pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): - Can be either: - - - A string, the *model id* of a model repo on huggingface.co. Valid model ids should have an - organization name, like `google/ddpm-celebahq-256`. - - A path to a *directory* containing model weights saved using [`~SchedulerMixin.save_pretrained`], - e.g., `./my_model_directory/`. - subfolder (`str`, *optional*): - In case the relevant files are located inside a subfolder of the model repo (either remote in - huggingface.co or downloaded locally), you can specify the folder name here. - 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*): - 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*): - 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`): - Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - local_files_only(`bool`, *optional*, defaults to `False`): - Whether or not to only look at local files (i.e., do not try to download the model). - token (`str` or *bool*, *optional*): - The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated - when running `transformers-cli login` (stored in `~/.huggingface`). - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a - git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any - identifier allowed by git. - - - - It is required to be logged in (`huggingface-cli login`) when you want to use private or [gated - models](https://huggingface.co/docs/hub/models-gated#gated-models). - - - - - - Activate the special ["offline-mode"](https://huggingface.co/transformers/installation.html#offline-mode) to - use this method in a firewalled environment. - - - - """ - config, kwargs = cls.load_config( - pretrained_model_name_or_path=pretrained_model_name_or_path, - subfolder=subfolder, - return_unused_kwargs=True, - **kwargs, - ) - scheduler, unused_kwargs = cls.from_config(config, return_unused_kwargs=True, **kwargs) - - if hasattr(scheduler, "create_state") and getattr(scheduler, "has_state", False): - state = scheduler.create_state() - - if return_unused_kwargs: - return scheduler, state, unused_kwargs - - return scheduler, state - - def save_pretrained(self, save_directory: Union[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. - - Args: - save_directory (`str` or `os.PathLike`): - Directory where the configuration JSON file will be saved (will be created if it does not exist). - push_to_hub (`bool`, *optional*, defaults to `False`): - 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*): - 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) - - @property - def compatibles(self): - """ - Returns all schedulers that are compatible with this scheduler - - Returns: - `List[SchedulerMixin]`: List of compatible schedulers - """ - return self._get_compatibles() - - @classmethod - def _get_compatibles(cls): - compatible_classes_str = list(set([cls.__name__] + cls._compatibles)) - diffusers_library = importlib.import_module(__name__.split(".")[0]) - compatible_classes = [ - getattr(diffusers_library, c) for c in compatible_classes_str if hasattr(diffusers_library, c) - ] - return compatible_classes - - -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) - - -def betas_for_alpha_bar(num_diffusion_timesteps: int, max_beta=0.999, dtype=jnp.float32) -> jnp.ndarray: - """ - Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of - (1-beta) over time from t = [0,1]. - - Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up - to that part of the diffusion process. - - - Args: - num_diffusion_timesteps (`int`): the number of betas to produce. - max_beta (`float`): the maximum beta to use; use values lower than 1 to - prevent singularities. - - Returns: - betas (`jnp.ndarray`): the betas used by the scheduler to step the model outputs - """ - - def alpha_bar(time_step): - return math.cos((time_step + 0.008) / 1.008 * math.pi / 2) ** 2 - - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) - return jnp.array(betas, dtype=dtype) - - -@flax.struct.dataclass -class CommonSchedulerState: - alphas: jnp.ndarray - betas: jnp.ndarray - alphas_cumprod: jnp.ndarray - - @classmethod - def create(cls, scheduler): - config = scheduler.config - - if config.trained_betas is not None: - betas = jnp.asarray(config.trained_betas, dtype=scheduler.dtype) - elif config.beta_schedule == "linear": - betas = jnp.linspace(config.beta_start, config.beta_end, config.num_train_timesteps, dtype=scheduler.dtype) - elif config.beta_schedule == "scaled_linear": - # this schedule is very specific to the latent diffusion model. - betas = ( - jnp.linspace( - config.beta_start**0.5, config.beta_end**0.5, config.num_train_timesteps, dtype=scheduler.dtype - ) - ** 2 - ) - elif config.beta_schedule == "squaredcos_cap_v2": - # Glide cosine schedule - betas = betas_for_alpha_bar(config.num_train_timesteps, dtype=scheduler.dtype) - else: - raise NotImplementedError( - f"beta_schedule {config.beta_schedule} is not implemented for scheduler {scheduler.__class__.__name__}" - ) - - alphas = 1.0 - betas - - alphas_cumprod = jnp.cumprod(alphas, axis=0) - - return cls( - alphas=alphas, - betas=betas, - alphas_cumprod=alphas_cumprod, - ) - - -def get_sqrt_alpha_prod( - state: CommonSchedulerState, original_samples: jnp.ndarray, noise: jnp.ndarray, timesteps: jnp.ndarray -): - alphas_cumprod = state.alphas_cumprod - - sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 - sqrt_alpha_prod = sqrt_alpha_prod.flatten() - sqrt_alpha_prod = broadcast_to_shape_from_left(sqrt_alpha_prod, original_samples.shape) - - sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 - sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() - sqrt_one_minus_alpha_prod = broadcast_to_shape_from_left(sqrt_one_minus_alpha_prod, original_samples.shape) - - return sqrt_alpha_prod, sqrt_one_minus_alpha_prod - - -def add_noise_common( - state: CommonSchedulerState, original_samples: jnp.ndarray, noise: jnp.ndarray, timesteps: jnp.ndarray -): - sqrt_alpha_prod, sqrt_one_minus_alpha_prod = get_sqrt_alpha_prod(state, original_samples, noise, timesteps) - noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise - return noisy_samples - - -def get_velocity_common(state: CommonSchedulerState, sample: jnp.ndarray, noise: jnp.ndarray, timesteps: jnp.ndarray): - sqrt_alpha_prod, sqrt_one_minus_alpha_prod = get_sqrt_alpha_prod(state, sample, noise, timesteps) - velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample - return velocity diff --git a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_vq_diffusion.py b/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_vq_diffusion.py deleted file mode 100644 index 57306301d023defde8ace84c62bc0f0662ee91f7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/schedulers/scheduling_vq_diffusion.py +++ /dev/null @@ -1,467 +0,0 @@ -# Copyright 2025 Microsoft 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 dataclasses import dataclass -from typing import Optional, Tuple, Union - -import numpy as np -import torch -import torch.nn.functional as F - -from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput -from .scheduling_utils import SchedulerMixin - - -@dataclass -class VQDiffusionSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's step function output. - - Args: - prev_sample (`torch.LongTensor` of shape `(batch size, num latent pixels)`): - Computed sample x_{t-1} of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - """ - - prev_sample: torch.LongTensor - - -def index_to_log_onehot(x: torch.LongTensor, num_classes: int) -> torch.Tensor: - """ - Convert batch of vector of class indices into batch of log onehot vectors - - Args: - x (`torch.LongTensor` of shape `(batch size, vector length)`): - Batch of class indices - - num_classes (`int`): - number of classes to be used for the onehot vectors - - Returns: - `torch.Tensor` of shape `(batch size, num classes, vector length)`: - Log onehot vectors - """ - x_onehot = F.one_hot(x, num_classes) - x_onehot = x_onehot.permute(0, 2, 1) - log_x = torch.log(x_onehot.float().clamp(min=1e-30)) - return log_x - - -def gumbel_noised(logits: torch.Tensor, generator: Optional[torch.Generator]) -> torch.Tensor: - """ - Apply gumbel noise to `logits` - """ - uniform = torch.rand(logits.shape, device=logits.device, generator=generator) - gumbel_noise = -torch.log(-torch.log(uniform + 1e-30) + 1e-30) - noised = gumbel_noise + logits - return noised - - -def alpha_schedules(num_diffusion_timesteps: int, alpha_cum_start=0.99999, alpha_cum_end=0.000009): - """ - Cumulative and non-cumulative alpha schedules. - - See section 4.1. - """ - att = ( - np.arange(0, num_diffusion_timesteps) / (num_diffusion_timesteps - 1) * (alpha_cum_end - alpha_cum_start) - + alpha_cum_start - ) - att = np.concatenate(([1], att)) - at = att[1:] / att[:-1] - att = np.concatenate((att[1:], [1])) - return at, att - - -def gamma_schedules(num_diffusion_timesteps: int, gamma_cum_start=0.000009, gamma_cum_end=0.99999): - """ - Cumulative and non-cumulative gamma schedules. - - See section 4.1. - """ - ctt = ( - np.arange(0, num_diffusion_timesteps) / (num_diffusion_timesteps - 1) * (gamma_cum_end - gamma_cum_start) - + gamma_cum_start - ) - ctt = np.concatenate(([0], ctt)) - one_minus_ctt = 1 - ctt - one_minus_ct = one_minus_ctt[1:] / one_minus_ctt[:-1] - ct = 1 - one_minus_ct - ctt = np.concatenate((ctt[1:], [0])) - return ct, ctt - - -class VQDiffusionScheduler(SchedulerMixin, ConfigMixin): - """ - A scheduler for vector quantized diffusion. - - 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_vec_classes (`int`): - The number of classes of the vector embeddings of the latent pixels. Includes the class for the masked - latent pixel. - num_train_timesteps (`int`, defaults to 100): - The number of diffusion steps to train the model. - alpha_cum_start (`float`, defaults to 0.99999): - The starting cumulative alpha value. - alpha_cum_end (`float`, defaults to 0.00009): - The ending cumulative alpha value. - gamma_cum_start (`float`, defaults to 0.00009): - The starting cumulative gamma value. - gamma_cum_end (`float`, defaults to 0.99999): - The ending cumulative gamma value. - """ - - order = 1 - - @register_to_config - def __init__( - self, - num_vec_classes: int, - num_train_timesteps: int = 100, - alpha_cum_start: float = 0.99999, - alpha_cum_end: float = 0.000009, - gamma_cum_start: float = 0.000009, - gamma_cum_end: float = 0.99999, - ): - self.num_embed = num_vec_classes - - # By convention, the index for the mask class is the last class index - self.mask_class = self.num_embed - 1 - - at, att = alpha_schedules(num_train_timesteps, alpha_cum_start=alpha_cum_start, alpha_cum_end=alpha_cum_end) - ct, ctt = gamma_schedules(num_train_timesteps, gamma_cum_start=gamma_cum_start, gamma_cum_end=gamma_cum_end) - - num_non_mask_classes = self.num_embed - 1 - bt = (1 - at - ct) / num_non_mask_classes - btt = (1 - att - ctt) / num_non_mask_classes - - at = torch.tensor(at.astype("float64")) - bt = torch.tensor(bt.astype("float64")) - ct = torch.tensor(ct.astype("float64")) - log_at = torch.log(at) - log_bt = torch.log(bt) - log_ct = torch.log(ct) - - att = torch.tensor(att.astype("float64")) - btt = torch.tensor(btt.astype("float64")) - ctt = torch.tensor(ctt.astype("float64")) - log_cumprod_at = torch.log(att) - log_cumprod_bt = torch.log(btt) - log_cumprod_ct = torch.log(ctt) - - self.log_at = log_at.float() - self.log_bt = log_bt.float() - self.log_ct = log_ct.float() - self.log_cumprod_at = log_cumprod_at.float() - self.log_cumprod_bt = log_cumprod_bt.float() - self.log_cumprod_ct = log_cumprod_ct.float() - - # setable values - 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): - """ - 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 and diffusion process parameters (alpha, beta, gamma) should be moved - to. - """ - self.num_inference_steps = num_inference_steps - timesteps = np.arange(0, self.num_inference_steps)[::-1].copy() - self.timesteps = torch.from_numpy(timesteps).to(device) - - self.log_at = self.log_at.to(device) - self.log_bt = self.log_bt.to(device) - self.log_ct = self.log_ct.to(device) - self.log_cumprod_at = self.log_cumprod_at.to(device) - self.log_cumprod_bt = self.log_cumprod_bt.to(device) - self.log_cumprod_ct = self.log_cumprod_ct.to(device) - - def step( - self, - model_output: torch.Tensor, - timestep: torch.long, - sample: torch.LongTensor, - generator: Optional[torch.Generator] = None, - return_dict: bool = True, - ) -> Union[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. - - Args: - log_p_x_0: (`torch.Tensor` of shape `(batch size, num classes - 1, num latent pixels)`): - The log probabilities for the predicted classes of the initial latent pixels. Does not include a - prediction for the masked class as the initial unnoised image cannot be masked. - t (`torch.long`): - The timestep that determines which transition matrices are used. - x_t (`torch.LongTensor` of shape `(batch size, num latent pixels)`): - The classes of each latent pixel at time `t`. - generator (`torch.Generator`, or `None`): - A random number generator for the noise applied to `p(x_{t-1} | x_t)` before it is sampled from. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~schedulers.scheduling_vq_diffusion.VQDiffusionSchedulerOutput`] or - `tuple`. - - Returns: - [`~schedulers.scheduling_vq_diffusion.VQDiffusionSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_vq_diffusion.VQDiffusionSchedulerOutput`] is - returned, otherwise a tuple is returned where the first element is the sample tensor. - """ - if timestep == 0: - log_p_x_t_min_1 = model_output - else: - log_p_x_t_min_1 = self.q_posterior(model_output, sample, timestep) - - log_p_x_t_min_1 = gumbel_noised(log_p_x_t_min_1, generator) - - x_t_min_1 = log_p_x_t_min_1.argmax(dim=1) - - if not return_dict: - return (x_t_min_1,) - - return VQDiffusionSchedulerOutput(prev_sample=x_t_min_1) - - def q_posterior(self, log_p_x_0, x_t, t): - """ - Calculates the log probabilities for the predicted classes of the image at timestep `t-1`: - - ``` - p(x_{t-1} | x_t) = sum( q(x_t | x_{t-1}) * q(x_{t-1} | x_0) * p(x_0) / q(x_t | x_0) ) - ``` - - Args: - log_p_x_0 (`torch.Tensor` of shape `(batch size, num classes - 1, num latent pixels)`): - The log probabilities for the predicted classes of the initial latent pixels. Does not include a - prediction for the masked class as the initial unnoised image cannot be masked. - x_t (`torch.LongTensor` of shape `(batch size, num latent pixels)`): - The classes of each latent pixel at time `t`. - t (`torch.Long`): - The timestep that determines which transition matrix is used. - - Returns: - `torch.Tensor` of shape `(batch size, num classes, num latent pixels)`: - The log probabilities for the predicted classes of the image at timestep `t-1`. - """ - log_onehot_x_t = index_to_log_onehot(x_t, self.num_embed) - - log_q_x_t_given_x_0 = self.log_Q_t_transitioning_to_known_class( - t=t, x_t=x_t, log_onehot_x_t=log_onehot_x_t, cumulative=True - ) - - log_q_t_given_x_t_min_1 = self.log_Q_t_transitioning_to_known_class( - t=t, x_t=x_t, log_onehot_x_t=log_onehot_x_t, cumulative=False - ) - - # p_0(x_0=C_0 | x_t) / q(x_t | x_0=C_0) ... p_n(x_0=C_0 | x_t) / q(x_t | x_0=C_0) - # . . . - # . . . - # . . . - # p_0(x_0=C_{k-1} | x_t) / q(x_t | x_0=C_{k-1}) ... p_n(x_0=C_{k-1} | x_t) / q(x_t | x_0=C_{k-1}) - q = log_p_x_0 - log_q_x_t_given_x_0 - - # sum_0 = p_0(x_0=C_0 | x_t) / q(x_t | x_0=C_0) + ... + p_0(x_0=C_{k-1} | x_t) / q(x_t | x_0=C_{k-1}), ... , - # sum_n = p_n(x_0=C_0 | x_t) / q(x_t | x_0=C_0) + ... + p_n(x_0=C_{k-1} | x_t) / q(x_t | x_0=C_{k-1}) - q_log_sum_exp = torch.logsumexp(q, dim=1, keepdim=True) - - # p_0(x_0=C_0 | x_t) / q(x_t | x_0=C_0) / sum_0 ... p_n(x_0=C_0 | x_t) / q(x_t | x_0=C_0) / sum_n - # . . . - # . . . - # . . . - # p_0(x_0=C_{k-1} | x_t) / q(x_t | x_0=C_{k-1}) / sum_0 ... p_n(x_0=C_{k-1} | x_t) / q(x_t | x_0=C_{k-1}) / sum_n - q = q - q_log_sum_exp - - # (p_0(x_0=C_0 | x_t) / q(x_t | x_0=C_0) / sum_0) * a_cumulative_{t-1} + b_cumulative_{t-1} ... (p_n(x_0=C_0 | x_t) / q(x_t | x_0=C_0) / sum_n) * a_cumulative_{t-1} + b_cumulative_{t-1} - # . . . - # . . . - # . . . - # (p_0(x_0=C_{k-1} | x_t) / q(x_t | x_0=C_{k-1}) / sum_0) * a_cumulative_{t-1} + b_cumulative_{t-1} ... (p_n(x_0=C_{k-1} | x_t) / q(x_t | x_0=C_{k-1}) / sum_n) * a_cumulative_{t-1} + b_cumulative_{t-1} - # c_cumulative_{t-1} ... c_cumulative_{t-1} - q = self.apply_cumulative_transitions(q, t - 1) - - # ((p_0(x_0=C_0 | x_t) / q(x_t | x_0=C_0) / sum_0) * a_cumulative_{t-1} + b_cumulative_{t-1}) * q(x_t | x_{t-1}=C_0) * sum_0 ... ((p_n(x_0=C_0 | x_t) / q(x_t | x_0=C_0) / sum_n) * a_cumulative_{t-1} + b_cumulative_{t-1}) * q(x_t | x_{t-1}=C_0) * sum_n - # . . . - # . . . - # . . . - # ((p_0(x_0=C_{k-1} | x_t) / q(x_t | x_0=C_{k-1}) / sum_0) * a_cumulative_{t-1} + b_cumulative_{t-1}) * q(x_t | x_{t-1}=C_{k-1}) * sum_0 ... ((p_n(x_0=C_{k-1} | x_t) / q(x_t | x_0=C_{k-1}) / sum_n) * a_cumulative_{t-1} + b_cumulative_{t-1}) * q(x_t | x_{t-1}=C_{k-1}) * sum_n - # c_cumulative_{t-1} * q(x_t | x_{t-1}=C_k) * sum_0 ... c_cumulative_{t-1} * q(x_t | x_{t-1}=C_k) * sum_0 - log_p_x_t_min_1 = q + log_q_t_given_x_t_min_1 + q_log_sum_exp - - # For each column, there are two possible cases. - # - # Where: - # - sum(p_n(x_0))) is summing over all classes for x_0 - # - C_i is the class transitioning from (not to be confused with c_t and c_cumulative_t being used for gamma's) - # - C_j is the class transitioning to - # - # 1. x_t is masked i.e. x_t = c_k - # - # Simplifying the expression, the column vector is: - # . - # . - # . - # (c_t / c_cumulative_t) * (a_cumulative_{t-1} * p_n(x_0 = C_i | x_t) + b_cumulative_{t-1} * sum(p_n(x_0))) - # . - # . - # . - # (c_cumulative_{t-1} / c_cumulative_t) * sum(p_n(x_0)) - # - # From equation (11) stated in terms of forward probabilities, the last row is trivially verified. - # - # For the other rows, we can state the equation as ... - # - # (c_t / c_cumulative_t) * [b_cumulative_{t-1} * p(x_0=c_0) + ... + (a_cumulative_{t-1} + b_cumulative_{t-1}) * p(x_0=C_i) + ... + b_cumulative_{k-1} * p(x_0=c_{k-1})] - # - # This verifies the other rows. - # - # 2. x_t is not masked - # - # Simplifying the expression, there are two cases for the rows of the column vector, where C_j = C_i and where C_j != C_i: - # . - # . - # . - # C_j != C_i: b_t * ((b_cumulative_{t-1} / b_cumulative_t) * p_n(x_0 = c_0) + ... + ((a_cumulative_{t-1} + b_cumulative_{t-1}) / b_cumulative_t) * p_n(x_0 = C_i) + ... + (b_cumulative_{t-1} / (a_cumulative_t + b_cumulative_t)) * p_n(c_0=C_j) + ... + (b_cumulative_{t-1} / b_cumulative_t) * p_n(x_0 = c_{k-1})) - # . - # . - # . - # C_j = C_i: (a_t + b_t) * ((b_cumulative_{t-1} / b_cumulative_t) * p_n(x_0 = c_0) + ... + ((a_cumulative_{t-1} + b_cumulative_{t-1}) / (a_cumulative_t + b_cumulative_t)) * p_n(x_0 = C_i = C_j) + ... + (b_cumulative_{t-1} / b_cumulative_t) * p_n(x_0 = c_{k-1})) - # . - # . - # . - # 0 - # - # The last row is trivially verified. The other rows can be verified by directly expanding equation (11) stated in terms of forward probabilities. - return log_p_x_t_min_1 - - def log_Q_t_transitioning_to_known_class( - self, *, t: torch.int, x_t: torch.LongTensor, log_onehot_x_t: torch.Tensor, cumulative: bool - ): - """ - Calculates the log probabilities of the rows from the (cumulative or non-cumulative) transition matrix for each - latent pixel in `x_t`. - - Args: - t (`torch.Long`): - The timestep that determines which transition matrix is used. - x_t (`torch.LongTensor` of shape `(batch size, num latent pixels)`): - The classes of each latent pixel at time `t`. - log_onehot_x_t (`torch.Tensor` of shape `(batch size, num classes, num latent pixels)`): - The log one-hot vectors of `x_t`. - cumulative (`bool`): - If cumulative is `False`, the single step transition matrix `t-1`->`t` is used. If cumulative is - `True`, the cumulative transition matrix `0`->`t` is used. - - Returns: - `torch.Tensor` of shape `(batch size, num classes - 1, num latent pixels)`: - Each _column_ of the returned matrix is a _row_ of log probabilities of the complete probability - transition matrix. - - When non cumulative, returns `self.num_classes - 1` rows because the initial latent pixel cannot be - masked. - - Where: - - `q_n` is the probability distribution for the forward process of the `n`th latent pixel. - - C_0 is a class of a latent pixel embedding - - C_k is the class of the masked latent pixel - - non-cumulative result (omitting logarithms): - ``` - q_0(x_t | x_{t-1} = C_0) ... q_n(x_t | x_{t-1} = C_0) - . . . - . . . - . . . - q_0(x_t | x_{t-1} = C_k) ... q_n(x_t | x_{t-1} = C_k) - ``` - - cumulative result (omitting logarithms): - ``` - q_0_cumulative(x_t | x_0 = C_0) ... q_n_cumulative(x_t | x_0 = C_0) - . . . - . . . - . . . - q_0_cumulative(x_t | x_0 = C_{k-1}) ... q_n_cumulative(x_t | x_0 = C_{k-1}) - ``` - """ - if cumulative: - a = self.log_cumprod_at[t] - b = self.log_cumprod_bt[t] - c = self.log_cumprod_ct[t] - else: - a = self.log_at[t] - b = self.log_bt[t] - c = self.log_ct[t] - - if not cumulative: - # The values in the onehot vector can also be used as the logprobs for transitioning - # from masked latent pixels. If we are not calculating the cumulative transitions, - # we need to save these vectors to be re-appended to the final matrix so the values - # aren't overwritten. - # - # `P(x_t!=mask|x_{t-1=mask}) = 0` and 0 will be the value of the last row of the onehot vector - # if x_t is not masked - # - # `P(x_t=mask|x_{t-1=mask}) = 1` and 1 will be the value of the last row of the onehot vector - # if x_t is masked - log_onehot_x_t_transitioning_from_masked = log_onehot_x_t[:, -1, :].unsqueeze(1) - - # `index_to_log_onehot` will add onehot vectors for masked pixels, - # so the default one hot matrix has one too many rows. See the doc string - # for an explanation of the dimensionality of the returned matrix. - log_onehot_x_t = log_onehot_x_t[:, :-1, :] - - # this is a cheeky trick to produce the transition probabilities using log one-hot vectors. - # - # Don't worry about what values this sets in the columns that mark transitions - # to masked latent pixels. They are overwrote later with the `mask_class_mask`. - # - # Looking at the below logspace formula in non-logspace, each value will evaluate to either - # `1 * a + b = a + b` where `log_Q_t` has the one hot value in the column - # or - # `0 * a + b = b` where `log_Q_t` has the 0 values in the column. - # - # See equation 7 for more details. - log_Q_t = (log_onehot_x_t + a).logaddexp(b) - - # The whole column of each masked pixel is `c` - mask_class_mask = x_t == self.mask_class - mask_class_mask = mask_class_mask.unsqueeze(1).expand(-1, self.num_embed - 1, -1) - log_Q_t[mask_class_mask] = c - - if not cumulative: - log_Q_t = torch.cat((log_Q_t, log_onehot_x_t_transitioning_from_masked), dim=1) - - return log_Q_t - - def apply_cumulative_transitions(self, q, t): - bsz = q.shape[0] - a = self.log_cumprod_at[t] - b = self.log_cumprod_bt[t] - c = self.log_cumprod_ct[t] - - num_latent_pixels = q.shape[2] - c = c.expand(bsz, 1, num_latent_pixels) - - q = (q + a).logaddexp(b) - q = torch.cat((q, c), dim=1) - - return q diff --git a/venv/lib/python3.10/site-packages/diffusers/training_utils.py b/venv/lib/python3.10/site-packages/diffusers/training_utils.py deleted file mode 100644 index bc30411d87267901a24795d30d86bed2d53da9a0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/training_utils.py +++ /dev/null @@ -1,652 +0,0 @@ -import contextlib -import copy -import gc -import math -import random -from typing import Any, Dict, Iterable, List, Optional, Tuple, Union - -import numpy as np -import torch - -from .models import UNet2DConditionModel -from .schedulers import SchedulerMixin -from .utils import ( - convert_state_dict_to_diffusers, - convert_state_dict_to_peft, - deprecate, - is_peft_available, - is_torch_npu_available, - is_torchvision_available, - is_transformers_available, -) - - -if is_transformers_available(): - import transformers - - if transformers.integrations.deepspeed.is_deepspeed_zero3_enabled(): - import deepspeed - -if is_peft_available(): - from peft import set_peft_model_state_dict - -if is_torchvision_available(): - from torchvision import transforms - -if is_torch_npu_available(): - import torch_npu # noqa: F401 - - -def set_seed(seed: int): - """ - Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`. - - Args: - seed (`int`): The seed to set. - - Returns: - `None` - """ - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if is_torch_npu_available(): - torch.npu.manual_seed_all(seed) - else: - torch.cuda.manual_seed_all(seed) - # ^^ safe to call this function even if cuda is not available - - -def compute_snr(noise_scheduler, timesteps): - """ - Computes SNR as per - https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849 - for the given timesteps using the provided noise scheduler. - - Args: - noise_scheduler (`NoiseScheduler`): - An object containing the noise schedule parameters, specifically `alphas_cumprod`, which is used to compute - the SNR values. - timesteps (`torch.Tensor`): - A tensor of timesteps for which the SNR is computed. - - Returns: - `torch.Tensor`: A tensor containing the computed SNR values for each timestep. - """ - alphas_cumprod = noise_scheduler.alphas_cumprod - sqrt_alphas_cumprod = alphas_cumprod**0.5 - sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5 - - # Expand the tensors. - # Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026 - sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[timesteps].float() - while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape): - sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None] - alpha = sqrt_alphas_cumprod.expand(timesteps.shape) - - sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(device=timesteps.device)[timesteps].float() - while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape): - sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None] - sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape) - - # Compute SNR. - snr = (alpha / sigma) ** 2 - return snr - - -def resolve_interpolation_mode(interpolation_type: str): - """ - Maps a string describing an interpolation function to the corresponding torchvision `InterpolationMode` enum. The - full list of supported enums is documented at - https://pytorch.org/vision/0.9/transforms.html#torchvision.transforms.functional.InterpolationMode. - - Args: - interpolation_type (`str`): - A string describing an interpolation method. Currently, `bilinear`, `bicubic`, `box`, `nearest`, - `nearest_exact`, `hamming`, and `lanczos` are supported, corresponding to the supported interpolation modes - in torchvision. - - Returns: - `torchvision.transforms.InterpolationMode`: an `InterpolationMode` enum used by torchvision's `resize` - transform. - """ - if not is_torchvision_available(): - raise ImportError( - "Please make sure to install `torchvision` to be able to use the `resolve_interpolation_mode()` function." - ) - - if interpolation_type == "bilinear": - interpolation_mode = transforms.InterpolationMode.BILINEAR - elif interpolation_type == "bicubic": - interpolation_mode = transforms.InterpolationMode.BICUBIC - elif interpolation_type == "box": - interpolation_mode = transforms.InterpolationMode.BOX - elif interpolation_type == "nearest": - interpolation_mode = transforms.InterpolationMode.NEAREST - elif interpolation_type == "nearest_exact": - interpolation_mode = transforms.InterpolationMode.NEAREST_EXACT - elif interpolation_type == "hamming": - interpolation_mode = transforms.InterpolationMode.HAMMING - elif interpolation_type == "lanczos": - interpolation_mode = transforms.InterpolationMode.LANCZOS - else: - raise ValueError( - f"The given interpolation mode {interpolation_type} is not supported. Currently supported interpolation" - f" modes are `bilinear`, `bicubic`, `box`, `nearest`, `nearest_exact`, `hamming`, and `lanczos`." - ) - - return interpolation_mode - - -def compute_dream_and_update_latents( - unet: UNet2DConditionModel, - noise_scheduler: SchedulerMixin, - timesteps: torch.Tensor, - noise: torch.Tensor, - noisy_latents: torch.Tensor, - target: torch.Tensor, - encoder_hidden_states: torch.Tensor, - dream_detail_preservation: float = 1.0, -) -> Tuple[Optional[torch.Tensor], Optional[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 - efficient and accurate at the cost of an extra forward step without gradients. - - Args: - `unet`: The state unet to use to make a prediction. - `noise_scheduler`: The noise scheduler used to add noise for the given timestep. - `timesteps`: The timesteps for the noise_scheduler to user. - `noise`: A tensor of noise in the shape of noisy_latents. - `noisy_latents`: Previously noise latents from the training loop. - `target`: The ground-truth tensor to predict after eps is removed. - `encoder_hidden_states`: Text embeddings from the text model. - `dream_detail_preservation`: A float value that indicates detail preservation level. - See reference. - - Returns: - `tuple[torch.Tensor, torch.Tensor]`: Adjusted noisy_latents and target. - """ - alphas_cumprod = noise_scheduler.alphas_cumprod.to(timesteps.device)[timesteps, None, None, None] - sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5 - - # The paper uses lambda = sqrt(1 - alpha) ** p, with p = 1 in their experiments. - dream_lambda = sqrt_one_minus_alphas_cumprod**dream_detail_preservation - - pred = None - with torch.no_grad(): - pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample - - _noisy_latents, _target = (None, None) - if noise_scheduler.config.prediction_type == "epsilon": - predicted_noise = pred - delta_noise = (noise - predicted_noise).detach() - delta_noise.mul_(dream_lambda) - _noisy_latents = noisy_latents.add(sqrt_one_minus_alphas_cumprod * delta_noise) - _target = target.add(delta_noise) - elif noise_scheduler.config.prediction_type == "v_prediction": - raise NotImplementedError("DREAM has not been implemented for v-prediction") - else: - raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}") - - return _noisy_latents, _target - - -def unet_lora_state_dict(unet: UNet2DConditionModel) -> Dict[str, torch.Tensor]: - r""" - Returns: - A state dict containing just the LoRA parameters. - """ - lora_state_dict = {} - - for name, module in unet.named_modules(): - if hasattr(module, "set_lora_layer"): - lora_layer = getattr(module, "lora_layer") - if lora_layer is not None: - current_lora_layer_sd = lora_layer.state_dict() - for lora_layer_matrix_name, lora_param in current_lora_layer_sd.items(): - # The matrix name can either be "down" or "up". - lora_state_dict[f"{name}.lora.{lora_layer_matrix_name}"] = lora_param - - return lora_state_dict - - -def cast_training_params(model: Union[torch.nn.Module, List[torch.nn.Module]], dtype=torch.float32): - """ - Casts the training parameters of the model to the specified data type. - - Args: - model: The PyTorch model whose parameters will be cast. - dtype: The data type to which the model parameters will be cast. - """ - if not isinstance(model, list): - model = [model] - for m in model: - for param in m.parameters(): - # only upcast trainable parameters into fp32 - if param.requires_grad: - param.data = param.to(dtype) - - -def _set_state_dict_into_text_encoder( - 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`. - - Args: - lora_state_dict: The state dictionary to be set. - prefix: String identifier to retrieve the portion of the state dict that belongs to `text_encoder`. - text_encoder: Where the `lora_state_dict` is to be set. - """ - - text_encoder_state_dict = { - f"{k.replace(prefix, '')}": v for k, v in lora_state_dict.items() if k.startswith(prefix) - } - text_encoder_state_dict = convert_state_dict_to_peft(convert_state_dict_to_diffusers(text_encoder_state_dict)) - 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]: - metadatas = {} - for module_name, module in modules_to_save.items(): - if module is not None: - metadatas[f"{module_name}_lora_adapter_metadata"] = module.peft_config["default"].to_dict() - return metadatas - - -def compute_density_for_timestep_sampling( - weighting_scheme: str, - batch_size: int, - logit_mean: float = None, - logit_std: float = None, - mode_scale: float = None, - device: Union[torch.device, str] = "cpu", - generator: Optional[torch.Generator] = None, -): - """ - Compute the density for sampling the timesteps when doing SD3 training. - - Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528. - - SD3 paper reference: https://huggingface.co/papers/2403.03206v1. - """ - if weighting_scheme == "logit_normal": - u = torch.normal(mean=logit_mean, std=logit_std, size=(batch_size,), device=device, generator=generator) - u = torch.nn.functional.sigmoid(u) - elif weighting_scheme == "mode": - u = torch.rand(size=(batch_size,), device=device, generator=generator) - u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2) ** 2 - 1 + u) - else: - u = torch.rand(size=(batch_size,), device=device, generator=generator) - return u - - -def compute_loss_weighting_for_sd3(weighting_scheme: str, sigmas=None): - """ - Computes loss weighting scheme for SD3 training. - - Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528. - - SD3 paper reference: https://huggingface.co/papers/2403.03206v1. - """ - if weighting_scheme == "sigma_sqrt": - weighting = (sigmas**-2.0).float() - elif weighting_scheme == "cosmap": - bot = 1 - 2 * sigmas + 2 * sigmas**2 - weighting = 2 / (math.pi * bot) - else: - weighting = torch.ones_like(sigmas) - return weighting - - -def free_memory(): - """ - Runs garbage collection. Then clears the cache of the available accelerator. - """ - gc.collect() - - if torch.cuda.is_available(): - torch.cuda.empty_cache() - elif torch.backends.mps.is_available(): - torch.mps.empty_cache() - elif is_torch_npu_available(): - torch_npu.npu.empty_cache() - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - torch.xpu.empty_cache() - - -# Adapted from torch-ema https://github.com/fadel/pytorch_ema/blob/master/torch_ema/ema.py#L14 -class EMAModel: - """ - Exponential Moving Average of models weights - """ - - def __init__( - self, - parameters: Iterable[torch.nn.Parameter], - decay: float = 0.9999, - 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, - foreach: bool = False, - model_cls: Optional[Any] = None, - model_config: Dict[str, Any] = None, - **kwargs, - ): - """ - Args: - parameters (Iterable[torch.nn.Parameter]): The parameters to track. - decay (float): The decay factor for the exponential moving average. - min_decay (float): The minimum decay factor for the exponential moving average. - update_after_step (int): The number of steps to wait before starting to update the EMA weights. - use_ema_warmup (bool): Whether to use EMA warmup. - inv_gamma (float): - 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 - weights will be stored on CPU. - - @crowsonkb's notes on EMA Warmup: - If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are good values for models you plan - to train for a million or more steps (reaches decay factor 0.999 at 31.6K steps, 0.9999 at 1M steps), - gamma=1, power=3/4 for models you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999 - at 215.4k steps). - """ - - if isinstance(parameters, torch.nn.Module): - deprecation_message = ( - "Passing a `torch.nn.Module` to `ExponentialMovingAverage` is deprecated. " - "Please pass the parameters of the module instead." - ) - deprecate( - "passing a `torch.nn.Module` to `ExponentialMovingAverage`", - "1.0.0", - deprecation_message, - standard_warn=False, - ) - parameters = parameters.parameters() - - # set use_ema_warmup to True if a torch.nn.Module is passed for backwards compatibility - use_ema_warmup = True - - if kwargs.get("max_value", None) is not None: - deprecation_message = "The `max_value` argument is deprecated. Please use `decay` instead." - deprecate("max_value", "1.0.0", deprecation_message, standard_warn=False) - decay = kwargs["max_value"] - - if kwargs.get("min_value", None) is not None: - deprecation_message = "The `min_value` argument is deprecated. Please use `min_decay` instead." - deprecate("min_value", "1.0.0", deprecation_message, standard_warn=False) - min_decay = kwargs["min_value"] - - parameters = list(parameters) - self.shadow_params = [p.clone().detach() for p in parameters] - - if kwargs.get("device", None) is not None: - deprecation_message = "The `device` argument is deprecated. Please use `to` instead." - deprecate("device", "1.0.0", deprecation_message, standard_warn=False) - self.to(device=kwargs["device"]) - - self.temp_stored_params = None - - self.decay = decay - self.min_decay = min_decay - self.update_after_step = update_after_step - self.use_ema_warmup = use_ema_warmup - self.inv_gamma = inv_gamma - self.power = power - self.optimization_step = 0 - self.cur_decay_value = None # set in `step()` - self.foreach = foreach - - self.model_cls = model_cls - self.model_config = model_config - - @classmethod - def from_pretrained(cls, path, model_cls, foreach=False) -> "EMAModel": - _, ema_kwargs = model_cls.from_config(path, return_unused_kwargs=True) - model = model_cls.from_pretrained(path) - - ema_model = cls(model.parameters(), model_cls=model_cls, model_config=model.config, foreach=foreach) - - ema_model.load_state_dict(ema_kwargs) - return ema_model - - def save_pretrained(self, path): - if self.model_cls is None: - raise ValueError("`save_pretrained` can only be used if `model_cls` was defined at __init__.") - - if self.model_config is None: - raise ValueError("`save_pretrained` can only be used if `model_config` was defined at __init__.") - - model = self.model_cls.from_config(self.model_config) - state_dict = self.state_dict() - state_dict.pop("shadow_params", None) - - model.register_to_config(**state_dict) - self.copy_to(model.parameters()) - model.save_pretrained(path) - - def get_decay(self, optimization_step: int) -> float: - """ - Compute the decay factor for the exponential moving average. - """ - step = max(0, optimization_step - self.update_after_step - 1) - - if step <= 0: - return 0.0 - - if self.use_ema_warmup: - cur_decay_value = 1 - (1 + step / self.inv_gamma) ** -self.power - else: - cur_decay_value = (1 + step) / (10 + step) - - cur_decay_value = min(cur_decay_value, self.decay) - # make sure decay is not smaller than min_decay - cur_decay_value = max(cur_decay_value, self.min_decay) - return cur_decay_value - - @torch.no_grad() - def step(self, parameters: Iterable[torch.nn.Parameter]): - if isinstance(parameters, torch.nn.Module): - deprecation_message = ( - "Passing a `torch.nn.Module` to `ExponentialMovingAverage.step` is deprecated. " - "Please pass the parameters of the module instead." - ) - deprecate( - "passing a `torch.nn.Module` to `ExponentialMovingAverage.step`", - "1.0.0", - deprecation_message, - standard_warn=False, - ) - parameters = parameters.parameters() - - parameters = list(parameters) - - self.optimization_step += 1 - - # Compute the decay factor for the exponential moving average. - decay = self.get_decay(self.optimization_step) - self.cur_decay_value = decay - one_minus_decay = 1 - decay - - context_manager = contextlib.nullcontext() - - if self.foreach: - if is_transformers_available() and transformers.integrations.deepspeed.is_deepspeed_zero3_enabled(): - context_manager = deepspeed.zero.GatheredParameters(parameters, modifier_rank=None) - - with context_manager: - params_grad = [param for param in parameters if param.requires_grad] - s_params_grad = [ - s_param for s_param, param in zip(self.shadow_params, parameters) if param.requires_grad - ] - - if len(params_grad) < len(parameters): - torch._foreach_copy_( - [s_param for s_param, param in zip(self.shadow_params, parameters) if not param.requires_grad], - [param for param in parameters if not param.requires_grad], - non_blocking=True, - ) - - torch._foreach_sub_( - s_params_grad, torch._foreach_sub(s_params_grad, params_grad), alpha=one_minus_decay - ) - - else: - for s_param, param in zip(self.shadow_params, parameters): - if is_transformers_available() and transformers.integrations.deepspeed.is_deepspeed_zero3_enabled(): - context_manager = deepspeed.zero.GatheredParameters(param, modifier_rank=None) - - with context_manager: - if param.requires_grad: - s_param.sub_(one_minus_decay * (s_param - param)) - else: - s_param.copy_(param) - - def copy_to(self, parameters: Iterable[torch.nn.Parameter]) -> None: - """ - Copy current averaged parameters into given collection of parameters. - - Args: - parameters: Iterable of `torch.nn.Parameter`; the parameters to be - updated with the stored moving averages. If `None`, the parameters with which this - `ExponentialMovingAverage` was initialized will be used. - """ - parameters = list(parameters) - if self.foreach: - torch._foreach_copy_( - [param.data for param in parameters], - [s_param.to(param.device).data for s_param, param in zip(self.shadow_params, parameters)], - ) - else: - for s_param, param in zip(self.shadow_params, parameters): - param.data.copy_(s_param.to(param.device).data) - - def pin_memory(self) -> None: - r""" - Move internal buffers of the ExponentialMovingAverage to pinned memory. Useful for non-blocking transfers for - offloading EMA params to the host. - """ - - self.shadow_params = [p.pin_memory() for p in self.shadow_params] - - def to(self, device=None, dtype=None, non_blocking=False) -> None: - r""" - Move internal buffers of the ExponentialMovingAverage to `device`. - - Args: - device: like `device` argument to `torch.Tensor.to` - """ - # .to() on the tensors handles None correctly - self.shadow_params = [ - p.to(device=device, dtype=dtype, non_blocking=non_blocking) - if p.is_floating_point() - else p.to(device=device, non_blocking=non_blocking) - for p in self.shadow_params - ] - - def state_dict(self) -> dict: - r""" - Returns the state of the ExponentialMovingAverage as a dict. This method is used by accelerate during - checkpointing to save the ema state dict. - """ - # Following PyTorch conventions, references to tensors are returned: - # "returns a reference to the state and not its copy!" - - # https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict - return { - "decay": self.decay, - "min_decay": self.min_decay, - "optimization_step": self.optimization_step, - "update_after_step": self.update_after_step, - "use_ema_warmup": self.use_ema_warmup, - "inv_gamma": self.inv_gamma, - "power": self.power, - "shadow_params": self.shadow_params, - } - - def store(self, parameters: Iterable[torch.nn.Parameter]) -> None: - r""" - Saves the current parameters for restoring later. - - Args: - parameters: Iterable of `torch.nn.Parameter`. The parameters to be temporarily stored. - """ - self.temp_stored_params = [param.detach().cpu().clone() for param in parameters] - - def restore(self, parameters: Iterable[torch.nn.Parameter]) -> None: - r""" - Restore the parameters stored with the `store` method. Useful to validate the model with EMA parameters - without: affecting the original optimization process. Store the parameters before the `copy_to()` method. After - validation (or model saving), use this to restore the former parameters. - - Args: - parameters: Iterable of `torch.nn.Parameter`; the parameters to be - updated with the stored parameters. If `None`, the parameters with which this - `ExponentialMovingAverage` was initialized will be used. - """ - - if self.temp_stored_params is None: - raise RuntimeError("This ExponentialMovingAverage has no `store()`ed weights to `restore()`") - if self.foreach: - torch._foreach_copy_( - [param.data for param in parameters], [c_param.data for c_param in self.temp_stored_params] - ) - else: - for c_param, param in zip(self.temp_stored_params, parameters): - param.data.copy_(c_param.data) - - # Better memory-wise. - self.temp_stored_params = None - - def load_state_dict(self, state_dict: dict) -> None: - r""" - Loads the ExponentialMovingAverage state. This method is used by accelerate during checkpointing to save the - ema state dict. - - Args: - state_dict (dict): EMA state. Should be an object returned - from a call to :meth:`state_dict`. - """ - # deepcopy, to be consistent with module API - state_dict = copy.deepcopy(state_dict) - - self.decay = state_dict.get("decay", self.decay) - if self.decay < 0.0 or self.decay > 1.0: - raise ValueError("Decay must be between 0 and 1") - - self.min_decay = state_dict.get("min_decay", self.min_decay) - if not isinstance(self.min_decay, float): - raise ValueError("Invalid min_decay") - - self.optimization_step = state_dict.get("optimization_step", self.optimization_step) - if not isinstance(self.optimization_step, int): - raise ValueError("Invalid optimization_step") - - self.update_after_step = state_dict.get("update_after_step", self.update_after_step) - if not isinstance(self.update_after_step, int): - raise ValueError("Invalid update_after_step") - - self.use_ema_warmup = state_dict.get("use_ema_warmup", self.use_ema_warmup) - if not isinstance(self.use_ema_warmup, bool): - raise ValueError("Invalid use_ema_warmup") - - self.inv_gamma = state_dict.get("inv_gamma", self.inv_gamma) - if not isinstance(self.inv_gamma, (float, int)): - raise ValueError("Invalid inv_gamma") - - self.power = state_dict.get("power", self.power) - if not isinstance(self.power, (float, int)): - raise ValueError("Invalid power") - - shadow_params = state_dict.get("shadow_params", None) - if shadow_params is not None: - self.shadow_params = shadow_params - if not isinstance(self.shadow_params, list): - raise ValueError("shadow_params must be a list") - if not all(isinstance(p, torch.Tensor) for p in self.shadow_params): - raise ValueError("shadow_params must all be Tensors") diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__init__.py b/venv/lib/python3.10/site-packages/diffusers/utils/__init__.py deleted file mode 100644 index 2df05cb8eb36fecaa22d12435008dd9903fecc3f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/__init__.py +++ /dev/null @@ -1,153 +0,0 @@ -# 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 os - -from packaging import version - -from .. import __version__ -from .constants import ( - CONFIG_NAME, - DEPRECATED_REVISION_ARGS, - DIFFUSERS_DYNAMIC_MODULE_NAME, - FLAX_WEIGHTS_NAME, - GGUF_FILE_EXTENSION, - HF_MODULES_CACHE, - HUGGINGFACE_CO_RESOLVE_ENDPOINT, - MIN_PEFT_VERSION, - ONNX_EXTERNAL_WEIGHTS_NAME, - ONNX_WEIGHTS_NAME, - SAFE_WEIGHTS_INDEX_NAME, - SAFETENSORS_FILE_EXTENSION, - SAFETENSORS_WEIGHTS_NAME, - USE_PEFT_BACKEND, - WEIGHTS_INDEX_NAME, - WEIGHTS_NAME, -) -from .deprecation_utils import deprecate -from .doc_utils import replace_example_docstring -from .dynamic_modules_utils import get_class_from_dynamic_module -from .export_utils import export_to_gif, export_to_obj, export_to_ply, export_to_video -from .hub_utils import ( - PushToHubMixin, - _add_variant, - _get_checkpoint_shard_files, - _get_model_file, - extract_commit_hash, - http_user_agent, -) -from .import_utils import ( - BACKENDS_MAPPING, - DIFFUSERS_SLOW_IMPORT, - ENV_VARS_TRUE_AND_AUTO_VALUES, - ENV_VARS_TRUE_VALUES, - USE_JAX, - USE_TF, - USE_TORCH, - DummyObject, - OptionalDependencyNotAvailable, - _LazyModule, - get_objects_from_module, - is_accelerate_available, - is_accelerate_version, - is_better_profanity_available, - is_bitsandbytes_available, - is_bitsandbytes_version, - is_bs4_available, - is_cosmos_guardrail_available, - is_flax_available, - is_ftfy_available, - is_gguf_available, - is_gguf_version, - is_google_colab, - is_hf_hub_version, - is_hpu_available, - is_inflect_available, - is_invisible_watermark_available, - is_k_diffusion_available, - is_k_diffusion_version, - is_librosa_available, - is_matplotlib_available, - is_nltk_available, - is_note_seq_available, - is_onnx_available, - is_opencv_available, - is_optimum_quanto_available, - is_optimum_quanto_version, - is_peft_available, - is_peft_version, - is_pytorch_retinaface_available, - is_safetensors_available, - is_scipy_available, - is_sentencepiece_available, - is_tensorboard_available, - is_timm_available, - is_torch_available, - is_torch_npu_available, - is_torch_version, - is_torch_xla_available, - is_torch_xla_version, - is_torchao_available, - is_torchao_version, - is_torchsde_available, - is_torchvision_available, - is_transformers_available, - is_transformers_version, - is_unidecode_available, - is_wandb_available, - is_xformers_available, - requires_backends, -) -from .loading_utils import get_module_from_name, get_submodule_by_name, load_image, load_video -from .logging import get_logger -from .outputs import BaseOutput -from .peft_utils import ( - check_peft_version, - delete_adapter_layers, - get_adapter_name, - get_peft_kwargs, - recurse_remove_peft_layers, - scale_lora_layers, - set_adapter_layers, - set_weights_and_activate_adapters, - unscale_lora_layers, -) -from .pil_utils import PIL_INTERPOLATION, make_image_grid, numpy_to_pil, pt_to_pil -from .remote_utils import remote_decode -from .state_dict_utils import ( - convert_all_state_dict_to_peft, - convert_state_dict_to_diffusers, - convert_state_dict_to_kohya, - convert_state_dict_to_peft, - convert_unet_state_dict_to_peft, - state_dict_all_zero, -) -from .typing_utils import _get_detailed_type, _is_valid_type - - -logger = get_logger(__name__) - - -def check_min_version(min_version): - if version.parse(__version__) < version.parse(min_version): - if "dev" in min_version: - error_message = ( - "This example requires a source install from HuggingFace diffusers (see " - "`https://huggingface.co/docs/diffusers/installation#install-from-source`)," - ) - else: - error_message = f"This example requires a minimum version of {min_version}," - error_message += f" but the version found is {__version__}.\n" - raise ImportError(error_message) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 3c947a1b8d10a6b52b0f92bc15247ef3a40feba2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/accelerate_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/accelerate_utils.cpython-310.pyc deleted file mode 100644 index 0c228a9560595acdf8330bab0785a45cecfacef4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/accelerate_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/constants.cpython-310.pyc deleted file mode 100644 index 709f27d7002b42c88d57a01a20322ac3d49565a6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/constants.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/deprecation_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/deprecation_utils.cpython-310.pyc deleted file mode 100644 index 6e5b0e03f9a7bc6780e1ffd86907a4fc4ad62655..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/deprecation_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/doc_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/doc_utils.cpython-310.pyc deleted file mode 100644 index d63b4bf2eda49aec987fde25b9ac481770b66b50..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/doc_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_bitsandbytes_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_bitsandbytes_objects.cpython-310.pyc deleted file mode 100644 index d62ed68f13a8a0ba0a114c9d51aed0cf4d46ac65..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_bitsandbytes_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_flax_and_transformers_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_flax_and_transformers_objects.cpython-310.pyc deleted file mode 100644 index 254e3338b33373e40f177ad8b543b782c0916913..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_flax_and_transformers_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_flax_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_flax_objects.cpython-310.pyc deleted file mode 100644 index 1e1f2ac479b3131d3b170982a3a9f5e7909057ea..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_flax_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_gguf_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_gguf_objects.cpython-310.pyc deleted file mode 100644 index e4ccdbb6a8bc1bbd918559b09e3372e33bcfcf95..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_gguf_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_note_seq_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_note_seq_objects.cpython-310.pyc deleted file mode 100644 index 28d7fd3c23b4f5b7ca952541c1525159302880ae..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_note_seq_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_onnx_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_onnx_objects.cpython-310.pyc deleted file mode 100644 index 0d033e5deac32116d23b3a6f72bea544b31a5e5b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_onnx_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_optimum_quanto_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_optimum_quanto_objects.cpython-310.pyc deleted file mode 100644 index 29fce5d87e4f629c9ab399de7dd343ff9c4f6a95..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_optimum_quanto_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_pt_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_pt_objects.cpython-310.pyc deleted file mode 100644 index d87295777cb804ab5a42c4c680ff8f13434a0c17..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_pt_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_librosa_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_librosa_objects.cpython-310.pyc deleted file mode 100644 index 9f349ed14ae0eefbc22c8be3218baf2c8bb8bb41..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_librosa_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_scipy_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_scipy_objects.cpython-310.pyc deleted file mode 100644 index 26e80b1cb05816f141528b3f249ec8ea1abc332b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_scipy_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_torchsde_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_torchsde_objects.cpython-310.pyc deleted file mode 100644 index 6c10fb31a29ac4a678df6e16fbefc6caf0b682e3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_torchsde_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_and_k_diffusion_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_and_k_diffusion_objects.cpython-310.pyc deleted file mode 100644 index 30bc1ea543362470e87000ec067caaa2f783a41d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_and_k_diffusion_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_and_onnx_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_and_onnx_objects.cpython-310.pyc deleted file mode 100644 index 14f3dd4e0ceee422e94cd2854b7e705bde7299e1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_and_onnx_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_and_opencv_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_and_opencv_objects.cpython-310.pyc deleted file mode 100644 index c7ecf189f1827e8e66c6ef8ef654ab2db9a37455..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_and_opencv_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_and_sentencepiece_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_and_sentencepiece_objects.cpython-310.pyc deleted file mode 100644 index a217d9202f9ac718103e3d10dabd6cabec71dc1d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_and_sentencepiece_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_objects.cpython-310.pyc deleted file mode 100644 index b18ef521d041c7cbe8b51d9355b63ffa5506bc36..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torch_and_transformers_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torchao_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torchao_objects.cpython-310.pyc deleted file mode 100644 index 20a832c305bb7e923bd4bec124154c3cc6826de1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_torchao_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_transformers_and_torch_and_note_seq_objects.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_transformers_and_torch_and_note_seq_objects.cpython-310.pyc deleted file mode 100644 index 35dc4e735a3eff7d34dd40b404a58697cd2d9d05..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dummy_transformers_and_torch_and_note_seq_objects.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dynamic_modules_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dynamic_modules_utils.cpython-310.pyc deleted file mode 100644 index cfffcc2ca5f78f05ab67ebd586143e19a6212ecf..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/dynamic_modules_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/export_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/export_utils.cpython-310.pyc deleted file mode 100644 index 818f1df202da3d13e90270a6c7e7b901ddd786cb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/export_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/hub_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/hub_utils.cpython-310.pyc deleted file mode 100644 index 7fc19531ab4caf46fc8fad53e0a7a4641adeaf44..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/hub_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/import_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/import_utils.cpython-310.pyc deleted file mode 100644 index 628a69fa74c4f2fd7378743a43f11cde9f75386f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/import_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/loading_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/loading_utils.cpython-310.pyc deleted file mode 100644 index 66b6d3c8c912cd82386c3c50551759c0529b7695..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/loading_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/logging.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/logging.cpython-310.pyc deleted file mode 100644 index 674198c5a2c9d7a8f0b6fa216da3c94806245f9c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/logging.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/outputs.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/outputs.cpython-310.pyc deleted file mode 100644 index bb0afc8b05cca2d6605a87462119955a1fe1eec5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/outputs.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/peft_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/peft_utils.cpython-310.pyc deleted file mode 100644 index 1d0fda61f1edaacf9d5ad0b4b1befeaf3089ffe9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/peft_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/pil_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/pil_utils.cpython-310.pyc deleted file mode 100644 index a1d737a4ec97a99b6183dde272f02a89237b1d44..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/pil_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/remote_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/remote_utils.cpython-310.pyc deleted file mode 100644 index 73af0c80556031ba1cb408028bc9c324268db2bf..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/remote_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/source_code_parsing_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/source_code_parsing_utils.cpython-310.pyc deleted file mode 100644 index fbb1dfea715ed3130ac9027630ff2173ee86bd0f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/source_code_parsing_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/state_dict_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/state_dict_utils.cpython-310.pyc deleted file mode 100644 index 45b94483fad9553cb94830eff71d216de1432845..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/state_dict_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/testing_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/testing_utils.cpython-310.pyc deleted file mode 100644 index 78583ddf8184585b07b0b1621603180772651231..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/testing_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/torch_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/torch_utils.cpython-310.pyc deleted file mode 100644 index 94a5d1d830ca4dd8db952a9edf40707672fd259b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/torch_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/typing_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/typing_utils.cpython-310.pyc deleted file mode 100644 index c4028f05881a68eded00b3612615529bb29bc91d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/typing_utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/versions.cpython-310.pyc b/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/versions.cpython-310.pyc deleted file mode 100644 index bfe94a8cb595a20fc21f45fef5a68c964f9620f0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/diffusers/utils/__pycache__/versions.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/accelerate_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/accelerate_utils.py deleted file mode 100644 index af3b712b5a9ea71ace657c5189fb6ad9f10d8148..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/accelerate_utils.py +++ /dev/null @@ -1,48 +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. -""" -Accelerate utilities: Utilities related to accelerate -""" - -from packaging import version - -from .import_utils import is_accelerate_available - - -if is_accelerate_available(): - import accelerate - - -def apply_forward_hook(method): - """ - Decorator that applies a registered CpuOffload hook to an arbitrary function rather than `forward`. This is useful - for cases where a PyTorch module provides functions other than `forward` that should trigger a move to the - appropriate acceleration device. This is the case for `encode` and `decode` in [`AutoencoderKL`]. - - This decorator looks inside the internal `_hf_hook` property to find a registered offload hook. - - :param method: The method to decorate. This method should be a method of a PyTorch module. - """ - if not is_accelerate_available(): - return method - accelerate_version = version.parse(accelerate.__version__).base_version - if version.parse(accelerate_version) < version.parse("0.17.0"): - return method - - def wrapper(self, *args, **kwargs): - if hasattr(self, "_hf_hook") and hasattr(self._hf_hook, "pre_forward"): - self._hf_hook.pre_forward(self) - return method(self, *args, **kwargs) - - return wrapper diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/constants.py b/venv/lib/python3.10/site-packages/diffusers/utils/constants.py deleted file mode 100644 index 7c04287d33ed9e4b2030d8c76aed5ff06cbbb285..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/constants.py +++ /dev/null @@ -1,70 +0,0 @@ -# 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 importlib -import os - -from huggingface_hub.constants import HF_HOME -from packaging import version - -from ..dependency_versions_check import dep_version_check -from .import_utils import ENV_VARS_TRUE_VALUES, is_peft_available, is_transformers_available - - -MIN_PEFT_VERSION = "0.6.0" -MIN_TRANSFORMERS_VERSION = "4.34.0" -_CHECK_PEFT = os.environ.get("_CHECK_PEFT", "1") in ENV_VARS_TRUE_VALUES - - -CONFIG_NAME = "config.json" -WEIGHTS_NAME = "diffusion_pytorch_model.bin" -WEIGHTS_INDEX_NAME = "diffusion_pytorch_model.bin.index.json" -FLAX_WEIGHTS_NAME = "diffusion_flax_model.msgpack" -ONNX_WEIGHTS_NAME = "model.onnx" -SAFETENSORS_WEIGHTS_NAME = "diffusion_pytorch_model.safetensors" -SAFE_WEIGHTS_INDEX_NAME = "diffusion_pytorch_model.safetensors.index.json" -SAFETENSORS_FILE_EXTENSION = "safetensors" -GGUF_FILE_EXTENSION = "gguf" -ONNX_EXTERNAL_WEIGHTS_NAME = "weights.pb" -HUGGINGFACE_CO_RESOLVE_ENDPOINT = os.environ.get("HF_ENDPOINT", "https://huggingface.co") -DIFFUSERS_DYNAMIC_MODULE_NAME = "diffusers_modules" -HF_MODULES_CACHE = os.getenv("HF_MODULES_CACHE", os.path.join(HF_HOME, "modules")) -DEPRECATED_REVISION_ARGS = ["fp16", "non-ema"] -DIFFUSERS_REQUEST_TIMEOUT = 60 - -# Below should be `True` if the current version of `peft` and `transformers` are compatible with -# PEFT backend. Will automatically fall back to PEFT backend if the correct versions of the libraries are -# available. -# For PEFT it is has to be greater than or equal to 0.6.0 and for transformers it has to be greater than or equal to 4.34.0. -_required_peft_version = is_peft_available() and version.parse( - version.parse(importlib.metadata.version("peft")).base_version -) >= version.parse(MIN_PEFT_VERSION) -_required_transformers_version = is_transformers_available() and version.parse( - version.parse(importlib.metadata.version("transformers")).base_version -) >= version.parse(MIN_TRANSFORMERS_VERSION) - -USE_PEFT_BACKEND = _required_peft_version and _required_transformers_version - -if USE_PEFT_BACKEND and _CHECK_PEFT: - dep_version_check("peft") - - -DECODE_ENDPOINT_SD_V1 = "https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/" -DECODE_ENDPOINT_SD_XL = "https://x2dmsqunjd6k9prw.us-east-1.aws.endpoints.huggingface.cloud/" -DECODE_ENDPOINT_FLUX = "https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud/" -DECODE_ENDPOINT_HUNYUAN_VIDEO = "https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud/" - - -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/" diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/deprecation_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/deprecation_utils.py deleted file mode 100644 index 4f001b3047d60c7572d770ad93595af3e62ae45e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/deprecation_utils.py +++ /dev/null @@ -1,49 +0,0 @@ -import inspect -import warnings -from typing import Any, Dict, Optional, Union - -from packaging import version - - -def deprecate(*args, take_from: Optional[Union[Dict, Any]] = None, standard_warn=True, stacklevel=2): - from .. import __version__ - - deprecated_kwargs = take_from - values = () - if not isinstance(args[0], tuple): - args = (args,) - - for attribute, version_name, message in args: - if version.parse(version.parse(__version__).base_version) >= version.parse(version_name): - raise ValueError( - f"The deprecation tuple {(attribute, version_name, message)} should be removed since diffusers'" - f" version {__version__} is >= {version_name}" - ) - - warning = None - if isinstance(deprecated_kwargs, dict) and attribute in deprecated_kwargs: - values += (deprecated_kwargs.pop(attribute),) - warning = f"The `{attribute}` argument is deprecated and will be removed in version {version_name}." - elif hasattr(deprecated_kwargs, attribute): - values += (getattr(deprecated_kwargs, attribute),) - warning = f"The `{attribute}` attribute is deprecated and will be removed in version {version_name}." - elif deprecated_kwargs is None: - warning = f"`{attribute}` is deprecated and will be removed in version {version_name}." - - if warning is not None: - warning = warning + " " if standard_warn else "" - warnings.warn(warning + message, FutureWarning, stacklevel=stacklevel) - - if isinstance(deprecated_kwargs, dict) and len(deprecated_kwargs) > 0: - call_frame = inspect.getouterframes(inspect.currentframe())[1] - filename = call_frame.filename - line_number = call_frame.lineno - function = call_frame.function - key, value = next(iter(deprecated_kwargs.items())) - raise TypeError(f"{function} in {filename} line {line_number - 1} got an unexpected keyword argument `{key}`") - - if len(values) == 0: - return - elif len(values) == 1: - return values[0] - return values diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/doc_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/doc_utils.py deleted file mode 100644 index 815083e1425801a5715da8ecd99af017427ff740..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/doc_utils.py +++ /dev/null @@ -1,39 +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. -""" -Doc utilities: Utilities related to documentation -""" - -import re - - -def replace_example_docstring(example_docstring): - def docstring_decorator(fn): - func_doc = fn.__doc__ - lines = func_doc.split("\n") - i = 0 - while i < len(lines) and re.search(r"^\s*Examples?:\s*$", lines[i]) is None: - i += 1 - if i < len(lines): - lines[i] = example_docstring - func_doc = "\n".join(lines) - else: - raise ValueError( - f"The function {fn} should have an empty 'Examples:' in its docstring as placeholder, " - f"current docstring is:\n{func_doc}" - ) - fn.__doc__ = func_doc - return fn - - return docstring_decorator diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_bitsandbytes_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_bitsandbytes_objects.py deleted file mode 100644 index 2dc589428de960f9c46f5a22874e2ef6475c418c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_bitsandbytes_objects.py +++ /dev/null @@ -1,17 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class BitsAndBytesConfig(metaclass=DummyObject): - _backends = ["bitsandbytes"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["bitsandbytes"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["bitsandbytes"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["bitsandbytes"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_flax_and_transformers_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_flax_and_transformers_objects.py deleted file mode 100644 index 5e65e5349bb0a6a0bac62cddf0ce0fad64237c68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_flax_and_transformers_objects.py +++ /dev/null @@ -1,77 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class FlaxStableDiffusionControlNetPipeline(metaclass=DummyObject): - _backends = ["flax", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - -class FlaxStableDiffusionImg2ImgPipeline(metaclass=DummyObject): - _backends = ["flax", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - -class FlaxStableDiffusionInpaintPipeline(metaclass=DummyObject): - _backends = ["flax", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - -class FlaxStableDiffusionPipeline(metaclass=DummyObject): - _backends = ["flax", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - -class FlaxStableDiffusionXLPipeline(metaclass=DummyObject): - _backends = ["flax", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax", "transformers"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_flax_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_flax_objects.py deleted file mode 100644 index 5fa8dbc819316e96f7483addba43f90b9d8f397b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_flax_objects.py +++ /dev/null @@ -1,212 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class FlaxControlNetModel(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxModelMixin(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxUNet2DConditionModel(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxAutoencoderKL(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxDiffusionPipeline(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxDDIMScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxDDPMScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxDPMSolverMultistepScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxEulerDiscreteScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxKarrasVeScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxLMSDiscreteScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxPNDMScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxSchedulerMixin(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - -class FlaxScoreSdeVeScheduler(metaclass=DummyObject): - _backends = ["flax"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["flax"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["flax"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_gguf_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_gguf_objects.py deleted file mode 100644 index 4a6d9a060a1397c5ef0e402e67170960bb3de516..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_gguf_objects.py +++ /dev/null @@ -1,17 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class GGUFQuantizationConfig(metaclass=DummyObject): - _backends = ["gguf"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["gguf"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["gguf"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["gguf"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_note_seq_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_note_seq_objects.py deleted file mode 100644 index c02d0b015aedc37c01fb3b843bc79547aae5da68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_note_seq_objects.py +++ /dev/null @@ -1,17 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class MidiProcessor(metaclass=DummyObject): - _backends = ["note_seq"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["note_seq"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["note_seq"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["note_seq"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_onnx_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_onnx_objects.py deleted file mode 100644 index bde5f6ad0793e2d81bc638600b46ff81748d09ee..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_onnx_objects.py +++ /dev/null @@ -1,17 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class OnnxRuntimeModel(metaclass=DummyObject): - _backends = ["onnx"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["onnx"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["onnx"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["onnx"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_optimum_quanto_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_optimum_quanto_objects.py deleted file mode 100644 index 44f8eaffc2460360d17720578b8224496f91cf94..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_optimum_quanto_objects.py +++ /dev/null @@ -1,17 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class QuantoConfig(metaclass=DummyObject): - _backends = ["optimum_quanto"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["optimum_quanto"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["optimum_quanto"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["optimum_quanto"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_pt_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_pt_objects.py deleted file mode 100644 index 2981f3a420d6b1f3c3cf057018f12d52e727f3c8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_pt_objects.py +++ /dev/null @@ -1,2108 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class FasterCacheConfig(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 HookRegistry(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"] - - 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 apply_faster_cache(*args, **kwargs): - requires_backends(apply_faster_cache, ["torch"]) - - -def apply_pyramid_attention_broadcast(*args, **kwargs): - requires_backends(apply_pyramid_attention_broadcast, ["torch"]) - - -class AllegroTransformer3DModel(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 AsymmetricAutoencoderKL(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 AuraFlowTransformer2DModel(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 AutoencoderDC(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 AutoencoderKL(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 AutoencoderKLAllegro(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 AutoencoderKLCogVideoX(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 AutoencoderKLCosmos(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 AutoencoderKLHunyuanVideo(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"] - - 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 AutoencoderKLMagvit(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 AutoencoderKLMochi(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 AutoencoderKLTemporalDecoder(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 AutoencoderKLWan(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 AutoencoderOobleck(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"] - - 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 AutoModel(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 CacheMixin(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 ChromaTransformer2DModel(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 CogVideoXTransformer3DModel(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 CogView3PlusTransformer2DModel(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 CogView4Transformer2DModel(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 ConsisIDTransformer3DModel(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 ConsistencyDecoderVAE(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 ControlNetModel(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 ControlNetUnionModel(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 ControlNetXSAdapter(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"] - - 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 DiTTransformer2DModel(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 EasyAnimateTransformer3DModel(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 FluxControlNetModel(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 FluxMultiControlNetModel(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 FluxTransformer2DModel(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"] - - 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 HunyuanDiT2DControlNetModel(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 HunyuanDiT2DModel(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 HunyuanDiT2DMultiControlNetModel(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 HunyuanVideoFramepackTransformer3DModel(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 HunyuanVideoTransformer3DModel(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 I2VGenXLUNet(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 Kandinsky3UNet(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 LatteTransformer3DModel(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"] - - 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 Lumina2Transformer2DModel(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 LuminaNextDiT2DModel(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 MochiTransformer3DModel(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 ModelMixin(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 MotionAdapter(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 MultiAdapter(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 MultiControlNetModel(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 OmniGenTransformer2DModel(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 PixArtTransformer2DModel(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 PriorTransformer(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 SanaControlNetModel(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 SanaTransformer2DModel(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 SD3ControlNetModel(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 SD3MultiControlNetModel(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 SD3Transformer2DModel(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 SparseControlNetModel(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 StableAudioDiTModel(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 T2IAdapter(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 T5FilmDecoder(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 Transformer2DModel(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 TransformerTemporalModel(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 UNet1DModel(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 UNet2DConditionModel(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 UNet2DModel(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 UNet3DConditionModel(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 UNetControlNetXSModel(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 UNetMotionModel(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 UNetSpatioTemporalConditionModel(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 UVit2DModel(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 VQModel(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 WanTransformer3DModel(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 WanVACETransformer3DModel(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"]) - - -def get_constant_schedule_with_warmup(*args, **kwargs): - requires_backends(get_constant_schedule_with_warmup, ["torch"]) - - -def get_cosine_schedule_with_warmup(*args, **kwargs): - requires_backends(get_cosine_schedule_with_warmup, ["torch"]) - - -def get_cosine_with_hard_restarts_schedule_with_warmup(*args, **kwargs): - requires_backends(get_cosine_with_hard_restarts_schedule_with_warmup, ["torch"]) - - -def get_linear_schedule_with_warmup(*args, **kwargs): - requires_backends(get_linear_schedule_with_warmup, ["torch"]) - - -def get_polynomial_decay_schedule_with_warmup(*args, **kwargs): - requires_backends(get_polynomial_decay_schedule_with_warmup, ["torch"]) - - -def get_scheduler(*args, **kwargs): - requires_backends(get_scheduler, ["torch"]) - - -class AudioPipelineOutput(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 AutoPipelineForImage2Image(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 AutoPipelineForInpainting(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 AutoPipelineForText2Image(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 BlipDiffusionControlNetPipeline(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 BlipDiffusionPipeline(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 CLIPImageProjection(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 ConsistencyModelPipeline(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 DanceDiffusionPipeline(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 DDIMPipeline(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 DDPMPipeline(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 DiffusionPipeline(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 DiTPipeline(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 ImagePipelineOutput(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 KarrasVePipeline(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 LDMPipeline(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 LDMSuperResolutionPipeline(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 PNDMPipeline(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 RePaintPipeline(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 ScoreSdeVePipeline(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 StableDiffusionMixin(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 DiffusersQuantizer(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 AmusedScheduler(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 CMStochasticIterativeScheduler(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 CogVideoXDDIMScheduler(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 CogVideoXDPMScheduler(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 DDIMInverseScheduler(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 DDIMParallelScheduler(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 DDIMScheduler(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 DDPMParallelScheduler(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 DDPMScheduler(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 DDPMWuerstchenScheduler(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 DEISMultistepScheduler(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 DPMSolverMultistepInverseScheduler(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 DPMSolverMultistepScheduler(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 DPMSolverSinglestepScheduler(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 EDMDPMSolverMultistepScheduler(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 EDMEulerScheduler(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 EulerAncestralDiscreteScheduler(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 EulerDiscreteScheduler(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 FlowMatchEulerDiscreteScheduler(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 FlowMatchHeunDiscreteScheduler(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 FlowMatchLCMScheduler(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"] - - 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 IPNDMScheduler(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 KarrasVeScheduler(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 KDPM2AncestralDiscreteScheduler(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 KDPM2DiscreteScheduler(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 LCMScheduler(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"] - - 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 RePaintScheduler(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 SASolverScheduler(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 SchedulerMixin(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 SCMScheduler(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 ScoreSdeVeScheduler(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 TCDScheduler(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 UnCLIPScheduler(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 UniPCMultistepScheduler(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 VQDiffusionScheduler(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 EMAModel(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"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_librosa_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_librosa_objects.py deleted file mode 100644 index 2088bc4a744198284f22fe54e6f1055cf3568566..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_librosa_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 AudioDiffusionPipeline(metaclass=DummyObject): - _backends = ["torch", "librosa"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "librosa"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "librosa"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "librosa"]) - - -class Mel(metaclass=DummyObject): - _backends = ["torch", "librosa"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "librosa"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "librosa"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "librosa"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_scipy_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_scipy_objects.py deleted file mode 100644 index a1ff25863822b04971d2c6dfdc17f5b28774cf05..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_scipy_objects.py +++ /dev/null @@ -1,17 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class LMSDiscreteScheduler(metaclass=DummyObject): - _backends = ["torch", "scipy"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "scipy"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "scipy"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "scipy"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_torchsde_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_torchsde_objects.py deleted file mode 100644 index 6ff14231b9cc8c88effc4d11f7b439fa16bb7f20..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_torchsde_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 CosineDPMSolverMultistepScheduler(metaclass=DummyObject): - _backends = ["torch", "torchsde"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "torchsde"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "torchsde"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "torchsde"]) - - -class DPMSolverSDEScheduler(metaclass=DummyObject): - _backends = ["torch", "torchsde"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "torchsde"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "torchsde"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "torchsde"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_transformers_and_k_diffusion_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_transformers_and_k_diffusion_objects.py deleted file mode 100644 index 2ab00c54ce2d7ca78f674085f0fe294d1f05c1e0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/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/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_transformers_and_onnx_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_transformers_and_onnx_objects.py deleted file mode 100644 index b7afad8226b87292100270e3e7daad6885be0e7f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_transformers_and_onnx_objects.py +++ /dev/null @@ -1,92 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class OnnxStableDiffusionImg2ImgPipeline(metaclass=DummyObject): - _backends = ["torch", "transformers", "onnx"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers", "onnx"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "onnx"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "onnx"]) - - -class OnnxStableDiffusionInpaintPipeline(metaclass=DummyObject): - _backends = ["torch", "transformers", "onnx"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers", "onnx"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "onnx"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "onnx"]) - - -class OnnxStableDiffusionInpaintPipelineLegacy(metaclass=DummyObject): - _backends = ["torch", "transformers", "onnx"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers", "onnx"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "onnx"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "onnx"]) - - -class OnnxStableDiffusionPipeline(metaclass=DummyObject): - _backends = ["torch", "transformers", "onnx"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers", "onnx"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "onnx"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "onnx"]) - - -class OnnxStableDiffusionUpscalePipeline(metaclass=DummyObject): - _backends = ["torch", "transformers", "onnx"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers", "onnx"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "onnx"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "onnx"]) - - -class StableDiffusionOnnxPipeline(metaclass=DummyObject): - _backends = ["torch", "transformers", "onnx"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers", "onnx"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "onnx"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "onnx"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_transformers_and_opencv_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_transformers_and_opencv_objects.py deleted file mode 100644 index c0c4d9df5eeeca4cca0a5147058eba70f89bcbc2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_transformers_and_opencv_objects.py +++ /dev/null @@ -1,17 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class ConsisIDPipeline(metaclass=DummyObject): - _backends = ["torch", "transformers", "opencv"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers", "opencv"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "opencv"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "opencv"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_transformers_and_sentencepiece_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_transformers_and_sentencepiece_objects.py deleted file mode 100644 index 349db10d1fa5b8505173e1578d3c83ea17e95fbc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_transformers_and_sentencepiece_objects.py +++ /dev/null @@ -1,47 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class KolorsImg2ImgPipeline(metaclass=DummyObject): - _backends = ["torch", "transformers", "sentencepiece"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers", "sentencepiece"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "sentencepiece"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "sentencepiece"]) - - -class KolorsPAGPipeline(metaclass=DummyObject): - _backends = ["torch", "transformers", "sentencepiece"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers", "sentencepiece"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "sentencepiece"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "sentencepiece"]) - - -class KolorsPipeline(metaclass=DummyObject): - _backends = ["torch", "transformers", "sentencepiece"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers", "sentencepiece"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "sentencepiece"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "sentencepiece"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_transformers_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_transformers_objects.py deleted file mode 100644 index 656a8ac6c6f6581693eaa00c5a02de10a42ed513..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torch_and_transformers_objects.py +++ /dev/null @@ -1,3032 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class AllegroPipeline(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 AltDiffusionImg2ImgPipeline(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 AltDiffusionPipeline(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 AmusedImg2ImgPipeline(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 AmusedInpaintPipeline(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 AmusedPipeline(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 AnimateDiffControlNetPipeline(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 AnimateDiffPAGPipeline(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 AnimateDiffPipeline(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 AnimateDiffSDXLPipeline(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 AnimateDiffSparseControlNetPipeline(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 AnimateDiffVideoToVideoControlNetPipeline(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 AnimateDiffVideoToVideoPipeline(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 AudioLDM2Pipeline(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 AudioLDM2ProjectionModel(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 AudioLDM2UNet2DConditionModel(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 AudioLDMPipeline(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 AuraFlowPipeline(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 ChromaImg2ImgPipeline(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"] - - 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 CLIPImageProjection(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 CogVideoXFunControlPipeline(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 CogVideoXImageToVideoPipeline(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 CogVideoXPipeline(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 CogVideoXVideoToVideoPipeline(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 CogView3PlusPipeline(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 CogView4ControlPipeline(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 CogView4Pipeline(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 ConsisIDPipeline(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"] - - 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 Cosmos2VideoToWorldPipeline(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 CosmosTextToWorldPipeline(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 CosmosVideoToWorldPipeline(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 CycleDiffusionPipeline(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 EasyAnimateControlPipeline(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 EasyAnimateInpaintPipeline(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 EasyAnimatePipeline(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 FluxControlImg2ImgPipeline(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 FluxControlInpaintPipeline(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 FluxControlNetImg2ImgPipeline(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 FluxControlNetInpaintPipeline(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 FluxControlNetPipeline(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 FluxControlPipeline(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 FluxFillPipeline(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 FluxImg2ImgPipeline(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 FluxInpaintPipeline(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 FluxPipeline(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 FluxPriorReduxPipeline(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"] - - 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 HunyuanDiTControlNetPipeline(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 HunyuanDiTPAGPipeline(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 HunyuanDiTPipeline(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 HunyuanSkyreelsImageToVideoPipeline(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 HunyuanVideoFramepackPipeline(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 HunyuanVideoImageToVideoPipeline(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 HunyuanVideoPipeline(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 I2VGenXLPipeline(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 IFImg2ImgPipeline(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 IFImg2ImgSuperResolutionPipeline(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 IFInpaintingPipeline(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 IFInpaintingSuperResolutionPipeline(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 IFPipeline(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 IFSuperResolutionPipeline(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 ImageTextPipelineOutput(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 Kandinsky3Img2ImgPipeline(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 Kandinsky3Pipeline(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 KandinskyCombinedPipeline(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 KandinskyImg2ImgCombinedPipeline(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 KandinskyImg2ImgPipeline(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 KandinskyInpaintCombinedPipeline(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 KandinskyInpaintPipeline(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 KandinskyPipeline(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 KandinskyPriorPipeline(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 KandinskyV22CombinedPipeline(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 KandinskyV22ControlnetImg2ImgPipeline(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 KandinskyV22ControlnetPipeline(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 KandinskyV22Img2ImgCombinedPipeline(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 KandinskyV22Img2ImgPipeline(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 KandinskyV22InpaintCombinedPipeline(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 KandinskyV22InpaintPipeline(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 KandinskyV22Pipeline(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 KandinskyV22PriorEmb2EmbPipeline(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 KandinskyV22PriorPipeline(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 LatentConsistencyModelImg2ImgPipeline(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 LatentConsistencyModelPipeline(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 LattePipeline(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 LDMTextToImagePipeline(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 LEditsPPPipelineStableDiffusion(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 LEditsPPPipelineStableDiffusionXL(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"] - - 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"] - - 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 LTXLatentUpsamplePipeline(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 LTXPipeline(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 Lumina2Pipeline(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 Lumina2Text2ImgPipeline(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 LuminaPipeline(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 LuminaText2ImgPipeline(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 MarigoldDepthPipeline(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 MarigoldIntrinsicsPipeline(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 MarigoldNormalsPipeline(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 MochiPipeline(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 MusicLDMPipeline(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 OmniGenPipeline(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 PaintByExamplePipeline(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 PIAPipeline(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 PixArtAlphaPipeline(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 PixArtSigmaPAGPipeline(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 PixArtSigmaPipeline(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 ReduxImageEncoder(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 SanaControlNetPipeline(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 SanaPAGPipeline(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 SanaPipeline(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 SanaSprintImg2ImgPipeline(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 SanaSprintPipeline(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 SemanticStableDiffusionPipeline(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 ShapEImg2ImgPipeline(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 ShapEPipeline(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 StableAudioPipeline(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 StableAudioProjectionModel(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 StableCascadeCombinedPipeline(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 StableCascadeDecoderPipeline(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 StableCascadePriorPipeline(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 StableDiffusion3ControlNetInpaintingPipeline(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 StableDiffusion3ControlNetPipeline(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 StableDiffusion3Img2ImgPipeline(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 StableDiffusion3InpaintPipeline(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 StableDiffusion3PAGImg2ImgPipeline(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 StableDiffusion3PAGPipeline(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 StableDiffusion3Pipeline(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 StableDiffusionAdapterPipeline(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 StableDiffusionAttendAndExcitePipeline(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 StableDiffusionControlNetImg2ImgPipeline(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 StableDiffusionControlNetInpaintPipeline(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 StableDiffusionControlNetPAGInpaintPipeline(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 StableDiffusionControlNetPAGPipeline(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 StableDiffusionControlNetPipeline(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 StableDiffusionControlNetXSPipeline(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 StableDiffusionDepth2ImgPipeline(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 StableDiffusionDiffEditPipeline(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 StableDiffusionGLIGENPipeline(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 StableDiffusionGLIGENTextImagePipeline(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 StableDiffusionImageVariationPipeline(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 StableDiffusionImg2ImgPipeline(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 StableDiffusionInpaintPipeline(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 StableDiffusionInpaintPipelineLegacy(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 StableDiffusionInstructPix2PixPipeline(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 StableDiffusionLatentUpscalePipeline(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 StableDiffusionLDM3DPipeline(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 StableDiffusionModelEditingPipeline(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 StableDiffusionPAGImg2ImgPipeline(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 StableDiffusionPAGInpaintPipeline(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 StableDiffusionPAGPipeline(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 StableDiffusionPanoramaPipeline(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 StableDiffusionParadigmsPipeline(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 StableDiffusionPipeline(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 StableDiffusionPipelineSafe(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 StableDiffusionPix2PixZeroPipeline(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 StableDiffusionSAGPipeline(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 StableDiffusionUpscalePipeline(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 StableDiffusionXLAdapterPipeline(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 StableDiffusionXLControlNetImg2ImgPipeline(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 StableDiffusionXLControlNetInpaintPipeline(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 StableDiffusionXLControlNetPAGImg2ImgPipeline(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 StableDiffusionXLControlNetPAGPipeline(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 StableDiffusionXLControlNetPipeline(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 StableDiffusionXLControlNetUnionImg2ImgPipeline(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 StableDiffusionXLControlNetUnionInpaintPipeline(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 StableDiffusionXLControlNetUnionPipeline(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 StableDiffusionXLControlNetXSPipeline(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 StableDiffusionXLImg2ImgPipeline(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 StableDiffusionXLInpaintPipeline(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 StableDiffusionXLInstructPix2PixPipeline(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 StableDiffusionXLPAGImg2ImgPipeline(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 StableDiffusionXLPAGInpaintPipeline(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 StableDiffusionXLPAGPipeline(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 StableDiffusionXLPipeline(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 StableUnCLIPImg2ImgPipeline(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 StableUnCLIPPipeline(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 StableVideoDiffusionPipeline(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 TextToVideoSDPipeline(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 TextToVideoZeroPipeline(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 TextToVideoZeroSDXLPipeline(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 UnCLIPImageVariationPipeline(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 UnCLIPPipeline(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 UniDiffuserModel(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 UniDiffuserPipeline(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 UniDiffuserTextDecoder(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 VersatileDiffusionDualGuidedPipeline(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 VersatileDiffusionImageVariationPipeline(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 VersatileDiffusionPipeline(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 VersatileDiffusionTextToImagePipeline(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 VideoToVideoSDPipeline(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 VisualClozeGenerationPipeline(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 VisualClozePipeline(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 VQDiffusionPipeline(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 WanImageToVideoPipeline(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 WanPipeline(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 WanVACEPipeline(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 WanVideoToVideoPipeline(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 WuerstchenCombinedPipeline(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 WuerstchenDecoderPipeline(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 WuerstchenPriorPipeline(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"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torchao_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torchao_objects.py deleted file mode 100644 index 16f0f6a55f64099c73ce4d81af8c87053b833e44..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_torchao_objects.py +++ /dev/null @@ -1,17 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class TorchAoConfig(metaclass=DummyObject): - _backends = ["torchao"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torchao"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torchao"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torchao"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_transformers_and_torch_and_note_seq_objects.py b/venv/lib/python3.10/site-packages/diffusers/utils/dummy_transformers_and_torch_and_note_seq_objects.py deleted file mode 100644 index fbde04e33f0abd86d12f3dee048a4f0585c9f19d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dummy_transformers_and_torch_and_note_seq_objects.py +++ /dev/null @@ -1,17 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class SpectrogramDiffusionPipeline(metaclass=DummyObject): - _backends = ["transformers", "torch", "note_seq"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["transformers", "torch", "note_seq"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["transformers", "torch", "note_seq"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["transformers", "torch", "note_seq"]) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/dynamic_modules_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/dynamic_modules_utils.py deleted file mode 100644 index 4878937ab2026c1140378f342215a38077aba2d0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/dynamic_modules_utils.py +++ /dev/null @@ -1,475 +0,0 @@ -# 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. -"""Utilities to dynamically load objects from the Hub.""" - -import importlib -import inspect -import json -import os -import re -import shutil -import sys -from pathlib import Path -from typing import Dict, Optional, Union -from urllib import request - -from huggingface_hub import hf_hub_download, model_info -from huggingface_hub.utils import RevisionNotFoundError, validate_hf_hub_args -from packaging import version - -from .. import __version__ -from . import DIFFUSERS_DYNAMIC_MODULE_NAME, HF_MODULES_CACHE, logging - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -# See https://huggingface.co/datasets/diffusers/community-pipelines-mirror -COMMUNITY_PIPELINES_MIRROR_ID = "diffusers/community-pipelines-mirror" - - -def get_diffusers_versions(): - url = "https://pypi.org/pypi/diffusers/json" - releases = json.loads(request.urlopen(url).read())["releases"].keys() - return sorted(releases, key=lambda x: version.Version(x)) - - -def init_hf_modules(): - """ - Creates the cache directory for modules with an init, and adds it to the Python path. - """ - # This function has already been executed if HF_MODULES_CACHE already is in the Python path. - if HF_MODULES_CACHE in sys.path: - return - - sys.path.append(HF_MODULES_CACHE) - os.makedirs(HF_MODULES_CACHE, exist_ok=True) - init_path = Path(HF_MODULES_CACHE) / "__init__.py" - if not init_path.exists(): - init_path.touch() - - -def create_dynamic_module(name: Union[str, os.PathLike]): - """ - Creates a dynamic module in the cache directory for modules. - """ - init_hf_modules() - dynamic_module_path = Path(HF_MODULES_CACHE) / name - # If the parent module does not exist yet, recursively create it. - if not dynamic_module_path.parent.exists(): - create_dynamic_module(dynamic_module_path.parent) - os.makedirs(dynamic_module_path, exist_ok=True) - init_path = dynamic_module_path / "__init__.py" - if not init_path.exists(): - init_path.touch() - - -def get_relative_imports(module_file): - """ - Get the list of modules that are relatively imported in a module file. - - Args: - module_file (`str` or `os.PathLike`): The module file to inspect. - """ - with open(module_file, "r", encoding="utf-8") as f: - content = f.read() - - # Imports of the form `import .xxx` - relative_imports = re.findall(r"^\s*import\s+\.(\S+)\s*$", content, flags=re.MULTILINE) - # Imports of the form `from .xxx import yyy` - relative_imports += re.findall(r"^\s*from\s+\.(\S+)\s+import", content, flags=re.MULTILINE) - # Unique-ify - return list(set(relative_imports)) - - -def get_relative_import_files(module_file): - """ - Get the list of all files that are needed for a given module. Note that this function recurses through the relative - imports (if a imports b and b imports c, it will return module files for b and c). - - Args: - module_file (`str` or `os.PathLike`): The module file to inspect. - """ - no_change = False - files_to_check = [module_file] - all_relative_imports = [] - - # Let's recurse through all relative imports - while not no_change: - new_imports = [] - for f in files_to_check: - new_imports.extend(get_relative_imports(f)) - - module_path = Path(module_file).parent - new_import_files = [str(module_path / m) for m in new_imports] - new_import_files = [f for f in new_import_files if f not in all_relative_imports] - files_to_check = [f"{f}.py" for f in new_import_files] - - no_change = len(new_import_files) == 0 - all_relative_imports.extend(files_to_check) - - return all_relative_imports - - -def check_imports(filename): - """ - Check if the current Python environment contains all the libraries that are imported in a file. - """ - with open(filename, "r", encoding="utf-8") as f: - content = f.read() - - # Imports of the form `import xxx` - imports = re.findall(r"^\s*import\s+(\S+)\s*$", content, flags=re.MULTILINE) - # Imports of the form `from xxx import yyy` - imports += re.findall(r"^\s*from\s+(\S+)\s+import", content, flags=re.MULTILINE) - # Only keep the top-level module - imports = [imp.split(".")[0] for imp in imports if not imp.startswith(".")] - - # Unique-ify and test we got them all - imports = list(set(imports)) - missing_packages = [] - for imp in imports: - try: - importlib.import_module(imp) - except ImportError: - missing_packages.append(imp) - - if len(missing_packages) > 0: - raise ImportError( - "This modeling file requires the following packages that were not found in your environment: " - f"{', '.join(missing_packages)}. Run `pip install {' '.join(missing_packages)}`" - ) - - return get_relative_imports(filename) - - -def get_class_in_module(class_name, module_path, pretrained_model_name_or_path=None): - """ - Import a module on the cache directory for modules and extract a class from it. - """ - module_path = module_path.replace(os.path.sep, ".") - try: - module = importlib.import_module(module_path) - except ModuleNotFoundError as e: - # This can happen when the repo id contains ".", which Python's import machinery interprets as a directory - # separator. We do a bit of monkey patching to detect and fix this case. - if not ( - pretrained_model_name_or_path is not None - and "." in pretrained_model_name_or_path - and module_path.startswith("diffusers_modules") - and pretrained_model_name_or_path.replace("/", "--") in module_path - ): - raise e # We can't figure this one out, just reraise the original error - - corrected_path = os.path.join(HF_MODULES_CACHE, module_path.replace(".", "/")) + ".py" - corrected_path = corrected_path.replace( - pretrained_model_name_or_path.replace("/", "--").replace(".", "/"), - pretrained_model_name_or_path.replace("/", "--"), - ) - module = importlib.machinery.SourceFileLoader(module_path, corrected_path).load_module() - - if class_name is None: - return find_pipeline_class(module) - return getattr(module, class_name) - - -def find_pipeline_class(loaded_module): - """ - Retrieve pipeline class that inherits from `DiffusionPipeline`. Note that there has to be exactly one class - inheriting from `DiffusionPipeline`. - """ - from ..pipelines import DiffusionPipeline - - cls_members = dict(inspect.getmembers(loaded_module, inspect.isclass)) - - pipeline_class = None - for cls_name, cls in cls_members.items(): - if ( - cls_name != DiffusionPipeline.__name__ - and issubclass(cls, DiffusionPipeline) - and cls.__module__.split(".")[0] != "diffusers" - ): - if pipeline_class is not None: - raise ValueError( - f"Multiple classes that inherit from {DiffusionPipeline.__name__} have been found:" - f" {pipeline_class.__name__}, and {cls_name}. Please make sure to define only one in" - f" {loaded_module}." - ) - pipeline_class = cls - - return pipeline_class - - -@validate_hf_hub_args -def get_cached_module_file( - pretrained_model_name_or_path: Union[str, os.PathLike], - module_file: str, - cache_dir: Optional[Union[str, os.PathLike]] = None, - force_download: bool = False, - proxies: Optional[Dict[str, str]] = None, - token: Optional[Union[bool, str]] = None, - revision: Optional[str] = None, - local_files_only: bool = False, -): - """ - Prepares Downloads a module from a local folder or a distant repo and returns its path inside the cached - Transformers module. - - Args: - pretrained_model_name_or_path (`str` or `os.PathLike`): - This can be either: - - - a string, the *model id* of a pretrained model configuration hosted inside a model repo on - huggingface.co. Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced - under a user or organization name, like `dbmdz/bert-base-german-cased`. - - a path to a *directory* containing a configuration file saved using the - [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`. - - module_file (`str`): - The name of the module file containing the class to look for. - cache_dir (`str` or `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 to (re-)download the configuration files and override the cached versions if they - exist. - 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*): - The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated - when running `transformers-cli login` (stored in `~/.huggingface`). - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a - git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any - identifier allowed by git. - local_files_only (`bool`, *optional*, defaults to `False`): - If `True`, will only try to load the tokenizer configuration from local files. - - - - You may pass a token in `token` if you are not logged in (`huggingface-cli login`) and want to use private or - [gated models](https://huggingface.co/docs/hub/models-gated#gated-models). - - - - Returns: - `str`: The path to the module inside the cache. - """ - # 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 os.path.isfile(module_file_or_url): - resolved_module_file = module_file_or_url - submodule = "local" - elif pretrained_model_name_or_path.count("/") == 0: - available_versions = get_diffusers_versions() - # cut ".dev0" - latest_version = "v" + ".".join(__version__.split(".")[:3]) - - # retrieve github version that matches - if revision is None: - revision = latest_version if latest_version[1:] in available_versions else "main" - logger.info(f"Defaulting to latest_version: {revision}.") - elif revision in available_versions: - revision = f"v{revision}" - elif revision == "main": - revision = revision - else: - raise ValueError( - f"`custom_revision`: {revision} does not exist. Please make sure to choose one of" - f" {', '.join(available_versions + ['main'])}." - ) - - try: - resolved_module_file = hf_hub_download( - repo_id=COMMUNITY_PIPELINES_MIRROR_ID, - repo_type="dataset", - filename=f"{revision}/{pretrained_model_name_or_path}.py", - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - ) - submodule = "git" - module_file = pretrained_model_name_or_path + ".py" - except RevisionNotFoundError as e: - raise EnvironmentError( - f"Revision '{revision}' not found in the community pipelines mirror. Check available revisions on" - " https://huggingface.co/datasets/diffusers/community-pipelines-mirror/tree/main." - " If you don't find the revision you are looking for, please open an issue on https://github.com/huggingface/diffusers/issues." - ) from e - except EnvironmentError: - logger.error(f"Could not locate the {module_file} inside {pretrained_model_name_or_path}.") - raise - else: - try: - # Load from URL or cache if already cached - resolved_module_file = hf_hub_download( - pretrained_model_name_or_path, - module_file, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - ) - submodule = os.path.join("local", "--".join(pretrained_model_name_or_path.split("/"))) - except EnvironmentError: - logger.error(f"Could not locate the {module_file} inside {pretrained_model_name_or_path}.") - raise - - # Check we have all the requirements in our environment - modules_needed = check_imports(resolved_module_file) - - # Now we move the module inside our cached dynamic modules. - full_submodule = DIFFUSERS_DYNAMIC_MODULE_NAME + os.path.sep + submodule - create_dynamic_module(full_submodule) - submodule_path = Path(HF_MODULES_CACHE) / full_submodule - if submodule == "local" or submodule == "git": - # We always copy local files (we could hash the file to see if there was a change, and give them the name of - # that hash, to only copy when there is a modification but it seems overkill for now). - # The only reason we do the copy is to avoid putting too many folders in sys.path. - shutil.copyfile(resolved_module_file, submodule_path / module_file) - for module_needed in modules_needed: - if len(module_needed.split(".")) == 2: - module_needed = "/".join(module_needed.split(".")) - module_folder = module_needed.split("/")[0] - 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) - else: - # Get the commit hash - # TODO: we will get this info in the etag soon, so retrieve it from there and not here. - commit_hash = model_info(pretrained_model_name_or_path, revision=revision, token=token).sha - - # The module file will end up being placed in a subfolder with the git hash of the repo. This way we get the - # benefit of versioning. - submodule_path = submodule_path / commit_hash - full_submodule = full_submodule + os.path.sep + commit_hash - create_dynamic_module(full_submodule) - - if not (submodule_path / module_file).exists(): - if len(module_file.split("/")) == 2: - module_folder = module_file.split("/")[0] - if not os.path.exists(submodule_path / module_folder): - os.makedirs(submodule_path / module_folder) - shutil.copyfile(resolved_module_file, submodule_path / module_file) - - # Make sure we also have every file with relative - for module_needed in modules_needed: - if len(module_needed.split(".")) == 2: - module_needed = "/".join(module_needed.split(".")) - if not (submodule_path / module_needed).exists(): - get_cached_module_file( - pretrained_model_name_or_path, - f"{module_needed}.py", - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - token=token, - revision=revision, - local_files_only=local_files_only, - ) - return os.path.join(full_submodule, module_file) - - -@validate_hf_hub_args -def get_class_from_dynamic_module( - pretrained_model_name_or_path: Union[str, os.PathLike], - module_file: str, - class_name: Optional[str] = None, - cache_dir: Optional[Union[str, os.PathLike]] = None, - force_download: bool = False, - proxies: Optional[Dict[str, str]] = None, - token: Optional[Union[bool, str]] = None, - revision: Optional[str] = None, - local_files_only: bool = False, - **kwargs, -): - """ - Extracts a class from a module file, present in the local folder or repository of a model. - - - - Calling this function will execute the code in the module file found locally or downloaded from the Hub. It should - therefore only be called on trusted repos. - - - - Args: - pretrained_model_name_or_path (`str` or `os.PathLike`): - This can be either: - - - a string, the *model id* of a pretrained model configuration hosted inside a model repo on - huggingface.co. Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced - under a user or organization name, like `dbmdz/bert-base-german-cased`. - - a path to a *directory* containing a configuration file saved using the - [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`. - - module_file (`str`): - The name of the module file containing the class to look for. - class_name (`str`): - The name of the class to import in the module. - cache_dir (`str` or `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 to (re-)download the configuration files and override the cached versions if they - exist. - 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*): - The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated - when running `transformers-cli login` (stored in `~/.huggingface`). - revision (`str`, *optional*, defaults to `"main"`): - The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a - git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any - identifier allowed by git. - local_files_only (`bool`, *optional*, defaults to `False`): - If `True`, will only try to load the tokenizer configuration from local files. - - - - You may pass a token in `token` if you are not logged in (`huggingface-cli login`) and want to use private or - [gated models](https://huggingface.co/docs/hub/models-gated#gated-models). - - - - Returns: - `type`: The class, dynamically imported from the module. - - Examples: - - ```python - # Download module `modeling.py` from huggingface.co and cache then extract the class `MyBertModel` from this - # module. - cls = get_class_from_dynamic_module("sgugger/my-bert-model", "modeling.py", "MyBertModel") - ```""" - # And lastly we get the class inside our newly created module - final_module = get_cached_module_file( - pretrained_model_name_or_path, - module_file, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - token=token, - revision=revision, - local_files_only=local_files_only, - ) - return get_class_in_module(class_name, final_module.replace(".py", ""), pretrained_model_name_or_path) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/export_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/export_utils.py deleted file mode 100644 index 07cf46928a44df0b27337e5d007c6a9bdb0e7c28..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/export_utils.py +++ /dev/null @@ -1,209 +0,0 @@ -import io -import random -import struct -import tempfile -from contextlib import contextmanager -from typing import List, Optional, Union - -import numpy as np -import PIL.Image -import PIL.ImageOps - -from .import_utils import BACKENDS_MAPPING, is_imageio_available, is_opencv_available -from .logging import get_logger - - -global_rng = random.Random() - -logger = get_logger(__name__) - - -@contextmanager -def buffered_writer(raw_f): - f = io.BufferedWriter(raw_f) - yield f - f.flush() - - -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 - - image[0].save( - output_gif_path, - save_all=True, - append_images=image[1:], - optimize=False, - duration=1000 // fps, - loop=0, - ) - return output_gif_path - - -def export_to_ply(mesh, output_ply_path: str = None): - """ - Write a PLY file for a mesh. - """ - if output_ply_path is None: - output_ply_path = tempfile.NamedTemporaryFile(suffix=".ply").name - - coords = mesh.verts.detach().cpu().numpy() - faces = mesh.faces.cpu().numpy() - rgb = np.stack([mesh.vertex_channels[x].detach().cpu().numpy() for x in "RGB"], axis=1) - - with buffered_writer(open(output_ply_path, "wb")) as f: - f.write(b"ply\n") - f.write(b"format binary_little_endian 1.0\n") - f.write(bytes(f"element vertex {len(coords)}\n", "ascii")) - f.write(b"property float x\n") - f.write(b"property float y\n") - f.write(b"property float z\n") - if rgb is not None: - f.write(b"property uchar red\n") - f.write(b"property uchar green\n") - f.write(b"property uchar blue\n") - if faces is not None: - f.write(bytes(f"element face {len(faces)}\n", "ascii")) - f.write(b"property list uchar int vertex_index\n") - f.write(b"end_header\n") - - if rgb is not None: - rgb = (rgb * 255.499).round().astype(int) - vertices = [ - (*coord, *rgb) - for coord, rgb in zip( - coords.tolist(), - rgb.tolist(), - ) - ] - format = struct.Struct("<3f3B") - for item in vertices: - f.write(format.pack(*item)) - else: - format = struct.Struct("<3f") - for vertex in coords.tolist(): - f.write(format.pack(*vertex)) - - if faces is not None: - format = struct.Struct(" str: - """ - quality: - Video output quality. Default is 5. Uses variable bit rate. Highest quality is 10, lowest is 0. Set to None to - prevent variable bitrate flags to FFMPEG so you can manually specify them using output_params instead. - Specifying a fixed bitrate using `bitrate` disables this parameter. - - bitrate: - Set a constant bitrate for the video encoding. Default is None causing `quality` parameter to be used instead. - Better quality videos with smaller file sizes will result from using the `quality` variable bitrate parameter - rather than specifying a fixed bitrate with this parameter. - - macro_block_size: - Size constraint for video. Width and height, must be divisible by this number. If not divisible by this number - imageio will tell ffmpeg to scale the image up to the next closest size divisible by this number. Most codecs - are compatible with a macroblock size of 16 (default), some can go smaller (4, 8). To disable this automatic - feature set it to None or 1, however be warned many players can't decode videos that are odd in size and some - codecs will produce poor results or fail. See https://en.wikipedia.org/wiki/Macroblock. - """ - # TODO: Dhruv. Remove by Diffusers release 0.33.0 - # Added to prevent breaking existing code - if not is_imageio_available(): - logger.warning( - ( - "It is recommended to use `export_to_video` with `imageio` and `imageio-ffmpeg` as a backend. \n" - "These libraries are not present in your environment. Attempting to use legacy OpenCV backend to export video. \n" - "Support for the OpenCV backend will be deprecated in a future Diffusers version" - ) - ) - return _legacy_export_to_video(video_frames, output_video_path, fps) - - if is_imageio_available(): - import imageio - else: - raise ImportError(BACKENDS_MAPPING["imageio"][1].format("export_to_video")) - - try: - imageio.plugins.ffmpeg.get_exe() - except AttributeError: - raise AttributeError( - ( - "Found an existing imageio backend in your environment. Attempting to export video with imageio. \n" - "Unable to find a compatible ffmpeg installation in your environment to use with imageio. Please install via `pip install imageio-ffmpeg" - ) - ) - - if output_video_path is None: - output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name - - if isinstance(video_frames[0], np.ndarray): - video_frames = [(frame * 255).astype(np.uint8) for frame in video_frames] - - elif isinstance(video_frames[0], PIL.Image.Image): - video_frames = [np.array(frame) for frame in video_frames] - - with imageio.get_writer( - output_video_path, fps=fps, quality=quality, bitrate=bitrate, macro_block_size=macro_block_size - ) as writer: - for frame in video_frames: - writer.append_data(frame) - - return output_video_path diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/hub_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/hub_utils.py deleted file mode 100644 index f80f96a3425d880c54bcf95d65dc6f3d2ba0f312..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/hub_utils.py +++ /dev/null @@ -1,557 +0,0 @@ -# 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. - - -import json -import os -import re -import sys -import tempfile -import warnings -from pathlib import Path -from typing import Dict, List, Optional, Union -from uuid import uuid4 - -from huggingface_hub import ( - DDUFEntry, - ModelCard, - ModelCardData, - create_repo, - hf_hub_download, - model_info, - snapshot_download, - upload_folder, -) -from huggingface_hub.constants import HF_HUB_DISABLE_TELEMETRY, HF_HUB_OFFLINE -from huggingface_hub.file_download import REGEX_COMMIT_HASH -from huggingface_hub.utils import ( - EntryNotFoundError, - RepositoryNotFoundError, - RevisionNotFoundError, - is_jinja_available, - validate_hf_hub_args, -) -from packaging import version -from requests import HTTPError - -from .. import __version__ -from .constants import ( - DEPRECATED_REVISION_ARGS, - HUGGINGFACE_CO_RESOLVE_ENDPOINT, - SAFETENSORS_WEIGHTS_NAME, - WEIGHTS_NAME, -) -from .import_utils import ( - ENV_VARS_TRUE_VALUES, - _flax_version, - _jax_version, - _onnxruntime_version, - _torch_version, - is_flax_available, - is_onnx_available, - is_torch_available, -) -from .logging import get_logger - - -logger = get_logger(__name__) - -MODEL_CARD_TEMPLATE_PATH = Path(__file__).parent / "model_card_template.md" -SESSION_ID = uuid4().hex - - -def http_user_agent(user_agent: Union[Dict, str, None] = None) -> str: - """ - Formats a user-agent string with basic info about a request. - """ - ua = f"diffusers/{__version__}; python/{sys.version.split()[0]}; session_id/{SESSION_ID}" - if HF_HUB_DISABLE_TELEMETRY or HF_HUB_OFFLINE: - return ua + "; telemetry/off" - if is_torch_available(): - ua += f"; torch/{_torch_version}" - if is_flax_available(): - ua += f"; jax/{_jax_version}" - ua += f"; flax/{_flax_version}" - if is_onnx_available(): - ua += f"; onnxruntime/{_onnxruntime_version}" - # CI will set this value to True - if os.environ.get("DIFFUSERS_IS_CI", "").upper() in ENV_VARS_TRUE_VALUES: - ua += "; is_ci/true" - if isinstance(user_agent, dict): - ua += "; " + "; ".join(f"{k}/{v}" for k, v in user_agent.items()) - elif isinstance(user_agent, str): - ua += "; " + user_agent - return ua - - -def load_or_create_model_card( - repo_id_or_path: str = None, - token: Optional[str] = None, - is_pipeline: bool = False, - from_training: bool = False, - model_description: Optional[str] = None, - base_model: str = None, - prompt: Optional[str] = None, - license: Optional[str] = None, - widget: Optional[List[dict]] = None, - inference: Optional[bool] = None, -) -> ModelCard: - """ - Loads or creates a model card. - - Args: - repo_id_or_path (`str`): - The repo id (e.g., "runwayml/stable-diffusion-v1-5") or local path where to look for the model card. - token (`str`, *optional*): - Authentication token. Will default to the stored token. See https://huggingface.co/settings/token for more - details. - is_pipeline (`bool`): - Boolean to indicate if we're adding tag to a [`DiffusionPipeline`]. - from_training: (`bool`): Boolean flag to denote if the model card is being created from a training script. - model_description (`str`, *optional*): Model description to add to the model card. Helpful when using - `load_or_create_model_card` from a training script. - base_model (`str`): Base model identifier (e.g., "stabilityai/stable-diffusion-xl-base-1.0"). Useful - for DreamBooth-like training. - 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. - inference: (`bool`, optional): Whether to turn on inference widget. Helpful when using - `load_or_create_model_card` from a training script. - """ - if not is_jinja_available(): - raise ValueError( - "Modelcard rendering is based on Jinja templates." - " Please make sure to have `jinja` installed before using `load_or_create_model_card`." - " To install it, please run `pip install Jinja2`." - ) - - try: - # Check if the model card is present on the remote repo - model_card = ModelCard.load(repo_id_or_path, token=token) - except (EntryNotFoundError, RepositoryNotFoundError): - # Otherwise create a model card from template - if from_training: - model_card = ModelCard.from_template( - card_data=ModelCardData( # Card metadata object that will be converted to YAML block - license=license, - library_name="diffusers", - inference=inference, - base_model=base_model, - instance_prompt=prompt, - widget=widget, - ), - template_path=MODEL_CARD_TEMPLATE_PATH, - model_description=model_description, - ) - 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) - - return model_card - - -def populate_model_card(model_card: ModelCard, tags: Union[str, List[str]] = 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" - - if tags is not None: - if isinstance(tags, str): - tags = [tags] - if model_card.data.tags is None: - model_card.data.tags = [] - for tag in tags: - model_card.data.tags.append(tag) - - return model_card - - -def extract_commit_hash(resolved_file: Optional[str], commit_hash: Optional[str] = None): - """ - Extracts the commit hash from a resolved filename toward a cache file. - """ - if resolved_file is None or commit_hash is not None: - return commit_hash - resolved_file = str(Path(resolved_file).as_posix()) - search = re.search(r"snapshots/([^/]+)/", resolved_file) - if search is None: - return None - commit_hash = search.groups()[0] - return commit_hash if REGEX_COMMIT_HASH.match(commit_hash) else None - - -def _add_variant(weights_name: str, variant: Optional[str] = None) -> str: - if variant is not None: - splits = weights_name.split(".") - splits = splits[:-1] + [variant] + splits[-1:] - weights_name = ".".join(splits) - - return weights_name - - -@validate_hf_hub_args -def _get_model_file( - pretrained_model_name_or_path: Union[str, Path], - *, - weights_name: str, - subfolder: Optional[str] = None, - cache_dir: Optional[str] = None, - force_download: bool = False, - proxies: Optional[Dict] = 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, -): - pretrained_model_name_or_path = str(pretrained_model_name_or_path) - - if dduf_entries: - if subfolder is not None: - raise ValueError( - "DDUF file only allow for 1 level of directory (e.g transformer/model1/model.safetentors is not allowed). " - "Please check the DDUF structure" - ) - model_file = ( - weights_name - if pretrained_model_name_or_path == "" - else "/".join([pretrained_model_name_or_path, weights_name]) - ) - if model_file in dduf_entries: - return model_file - else: - raise EnvironmentError(f"Error no file named {weights_name} found in archive {dduf_entries.keys()}.") - elif os.path.isfile(pretrained_model_name_or_path): - return pretrained_model_name_or_path - elif os.path.isdir(pretrained_model_name_or_path): - if os.path.isfile(os.path.join(pretrained_model_name_or_path, weights_name)): - # Load from a PyTorch checkpoint - model_file = os.path.join(pretrained_model_name_or_path, weights_name) - return model_file - elif subfolder is not None and os.path.isfile( - os.path.join(pretrained_model_name_or_path, subfolder, weights_name) - ): - model_file = os.path.join(pretrained_model_name_or_path, subfolder, weights_name) - return model_file - else: - raise EnvironmentError( - f"Error no file named {weights_name} found in directory {pretrained_model_name_or_path}." - ) - else: - # 1. First check if deprecated way of loading from branches is used - if ( - revision in DEPRECATED_REVISION_ARGS - and (weights_name == WEIGHTS_NAME or weights_name == SAFETENSORS_WEIGHTS_NAME) - and version.parse(version.parse(__version__).base_version) >= version.parse("0.22.0") - ): - try: - model_file = hf_hub_download( - pretrained_model_name_or_path, - filename=_add_variant(weights_name, revision), - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - user_agent=user_agent, - subfolder=subfolder, - revision=revision or commit_hash, - ) - warnings.warn( - f"Loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'` is deprecated. Loading instead from `revision='main'` with `variant={revision}`. Loading model variants via `revision='{revision}'` will be removed in diffusers v1. Please use `variant='{revision}'` instead.", - FutureWarning, - ) - return model_file - except: # noqa: E722 - warnings.warn( - f"You are loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'`. This behavior is deprecated and will be removed in diffusers v1. One should use `variant='{revision}'` instead. However, it appears that {pretrained_model_name_or_path} currently does not have a {_add_variant(weights_name, revision)} file in the 'main' branch of {pretrained_model_name_or_path}. \n The Diffusers team and community would be very grateful if you could open an issue: https://github.com/huggingface/diffusers/issues/new with the title '{pretrained_model_name_or_path} is missing {_add_variant(weights_name, revision)}' so that the correct variant file can be added.", - FutureWarning, - ) - try: - # 2. Load model file as usual - model_file = hf_hub_download( - pretrained_model_name_or_path, - filename=weights_name, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - user_agent=user_agent, - subfolder=subfolder, - revision=revision or commit_hash, - ) - return model_file - - except RepositoryNotFoundError as e: - raise EnvironmentError( - f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier " - "listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a " - "token having permission to this repo with `token` or log in with `huggingface-cli " - "login`." - ) from e - except RevisionNotFoundError as e: - raise EnvironmentError( - f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for " - "this model name. Check the model page at " - f"'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions." - ) from e - except EntryNotFoundError as e: - raise EnvironmentError( - f"{pretrained_model_name_or_path} does not appear to have a file named {weights_name}." - ) from e - except HTTPError as e: - raise EnvironmentError( - f"There was a specific connection error when trying to load {pretrained_model_name_or_path}:\n{e}" - ) from e - except ValueError as e: - raise EnvironmentError( - f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it" - f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a" - f" directory containing a file named {weights_name} or" - " \nCheckout your internet connection or see how to run the library in" - " offline mode at 'https://huggingface.co/docs/diffusers/installation#offline-mode'." - ) from e - except EnvironmentError as e: - raise EnvironmentError( - f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it from " - "'https://huggingface.co/models', make sure you don't have a local directory with the same name. " - f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory " - f"containing a file named {weights_name}" - ) from e - - -def _get_checkpoint_shard_files( - pretrained_model_name_or_path, - index_filename, - cache_dir=None, - proxies=None, - local_files_only=False, - token=None, - user_agent=None, - revision=None, - subfolder="", - dduf_entries: Optional[Dict[str, DDUFEntry]] = None, -): - """ - For a given model: - - - download and cache all the shards of a sharded checkpoint if `pretrained_model_name_or_path` is a model ID on the - Hub - - returns the list of paths to all the shards, as well as some metadata. - - For the description of each arg, see [`PreTrainedModel.from_pretrained`]. `index_filename` is the full path to the - index (downloaded and cached if `pretrained_model_name_or_path` is a model ID on the Hub). - """ - if dduf_entries: - if index_filename not in dduf_entries: - raise ValueError(f"Can't find a checkpoint index ({index_filename}) in {pretrained_model_name_or_path}.") - else: - if not os.path.isfile(index_filename): - raise ValueError(f"Can't find a checkpoint index ({index_filename}) in {pretrained_model_name_or_path}.") - - if dduf_entries: - index = json.loads(dduf_entries[index_filename].read_text()) - else: - with open(index_filename, "r") as f: - index = json.loads(f.read()) - - original_shard_filenames = sorted(set(index["weight_map"].values())) - sharded_metadata = index["metadata"] - sharded_metadata["all_checkpoint_keys"] = list(index["weight_map"].keys()) - sharded_metadata["weight_map"] = index["weight_map"].copy() - shards_path = os.path.join(pretrained_model_name_or_path, subfolder) - - # First, let's deal with local folder. - if os.path.isdir(pretrained_model_name_or_path) or dduf_entries: - shard_filenames = [os.path.join(shards_path, f) for f in original_shard_filenames] - for shard_file in shard_filenames: - if dduf_entries: - if shard_file not in dduf_entries: - raise FileNotFoundError( - f"{shards_path} does not appear to have a file named {shard_file} which is " - "required according to the checkpoint index." - ) - else: - if not os.path.exists(shard_file): - raise FileNotFoundError( - f"{shards_path} does not appear to have a file named {shard_file} which is " - "required according to the checkpoint index." - ) - return shard_filenames, sharded_metadata - - # At this stage pretrained_model_name_or_path is a model identifier on the Hub - allow_patterns = original_shard_filenames - if subfolder is not None: - allow_patterns = [os.path.join(subfolder, p) for p in allow_patterns] - - ignore_patterns = ["*.json", "*.md"] - # `model_info` call must guarded with the above condition. - model_files_info = model_info(pretrained_model_name_or_path, revision=revision, token=token) - for shard_file in original_shard_filenames: - shard_file_present = any(shard_file in k.rfilename for k in model_files_info.siblings) - if not shard_file_present: - raise EnvironmentError( - f"{shards_path} does not appear to have a file named {shard_file} which is " - "required according to the checkpoint index." - ) - - try: - # Load from URL - cached_folder = snapshot_download( - pretrained_model_name_or_path, - cache_dir=cache_dir, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - allow_patterns=allow_patterns, - ignore_patterns=ignore_patterns, - user_agent=user_agent, - ) - if subfolder is not None: - cached_folder = os.path.join(cached_folder, subfolder) - - # We have already dealt with RepositoryNotFoundError and RevisionNotFoundError when getting the index, so - # we don't have to catch them here. We have also dealt with EntryNotFoundError. - except HTTPError as e: - raise EnvironmentError( - f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load {pretrained_model_name_or_path}. You should try" - " again after checking your internet connection." - ) from e - - cached_filenames = [os.path.join(cached_folder, f) for f in original_shard_filenames] - - return cached_filenames, sharded_metadata - - -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: - filenames = [] - for _, _, files in os.walk(folder): - for file in files: - filenames.append(os.path.basename(file)) - transformers_index_format = r"\d{5}-of-\d{5}" - variant_file_re = re.compile(rf".*-{transformers_index_format}\.{variant}\.[a-z]+$") - return any(variant_file_re.match(f) is not None for f in filenames) - - -class PushToHubMixin: - """ - A Mixin to push a model, scheduler, or pipeline to the Hugging Face Hub. - """ - - def _upload_folder( - self, - working_dir: Union[str, os.PathLike], - repo_id: str, - token: Optional[str] = None, - commit_message: Optional[str] = None, - create_pr: bool = False, - ): - """ - Uploads all files in `working_dir` to `repo_id`. - """ - if commit_message is None: - if "Model" in self.__class__.__name__: - commit_message = "Upload model" - elif "Scheduler" in self.__class__.__name__: - commit_message = "Upload scheduler" - else: - commit_message = f"Upload {self.__class__.__name__}" - - logger.info(f"Uploading the files of {working_dir} to {repo_id}.") - return upload_folder( - repo_id=repo_id, folder_path=working_dir, token=token, commit_message=commit_message, create_pr=create_pr - ) - - def push_to_hub( - self, - repo_id: str, - commit_message: Optional[str] = None, - private: Optional[bool] = None, - token: Optional[str] = None, - create_pr: bool = False, - safe_serialization: bool = True, - variant: Optional[str] = None, - ) -> str: - """ - Upload model, scheduler, or pipeline files to the 🤗 Hugging Face Hub. - - Parameters: - repo_id (`str`): - The name of the repository you want to push your model, scheduler, or pipeline files to. It should - contain your organization name when pushing to an organization. `repo_id` can also be a path to a local - directory. - commit_message (`str`, *optional*): - Message to commit while pushing. Default to `"Upload {object}"`. - private (`bool`, *optional*): - Whether to make the repo private. If `None` (default), the repo will be public unless the - organization's default is private. This value is ignored if the repo already exists. - token (`str`, *optional*): - The token to use as HTTP bearer authorization for remote files. The token generated when running - `huggingface-cli login` (stored in `~/.huggingface`). - create_pr (`bool`, *optional*, defaults to `False`): - Whether or not to create a PR with the uploaded files or directly commit. - safe_serialization (`bool`, *optional*, defaults to `True`): - Whether or not to convert the model weights to the `safetensors` format. - variant (`str`, *optional*): - If specified, weights are saved in the format `pytorch_model..bin`. - - Examples: - - ```python - from diffusers import UNet2DConditionModel - - unet = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-2", subfolder="unet") - - # Push the `unet` to your namespace with the name "my-finetuned-unet". - unet.push_to_hub("my-finetuned-unet") - - # Push the `unet` to an organization with the name "my-finetuned-unet". - unet.push_to_hub("your-org/my-finetuned-unet") - ``` - """ - repo_id = create_repo(repo_id, private=private, token=token, exist_ok=True).repo_id - - # Create a new empty model card and eventually tag it - model_card = load_or_create_model_card(repo_id, token=token) - model_card = populate_model_card(model_card) - - # Save all files. - save_kwargs = {"safe_serialization": safe_serialization} - if "Scheduler" not in self.__class__.__name__: - save_kwargs.update({"variant": variant}) - - with tempfile.TemporaryDirectory() as tmpdir: - self.save_pretrained(tmpdir, **save_kwargs) - - # Update model card if needed: - model_card.save(os.path.join(tmpdir, "README.md")) - - return self._upload_folder( - tmpdir, - repo_id, - token=token, - commit_message=commit_message, - create_pr=create_pr, - ) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/import_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/import_utils.py deleted file mode 100644 index f12e9de33172a27ea7e809b1ff8c30effa14c9bc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/import_utils.py +++ /dev/null @@ -1,891 +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 utilities: Utilities related to imports and our lazy inits. -""" - -import importlib.util -import inspect -import operator as op -import os -import sys -from collections import OrderedDict, defaultdict -from itertools import chain -from types import ModuleType -from typing import Any, Tuple, Union - -from huggingface_hub.utils import is_jinja_available # noqa: F401 -from packaging.version import Version, parse - -from . import logging - - -# The package importlib_metadata is in a different place, depending on the python version. -if sys.version_info < (3, 8): - import importlib_metadata -else: - import importlib.metadata as importlib_metadata -try: - _package_map = importlib_metadata.packages_distributions() # load-once to avoid expensive calls -except Exception: - _package_map = None - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"} -ENV_VARS_TRUE_AND_AUTO_VALUES = ENV_VARS_TRUE_VALUES.union({"AUTO"}) - -USE_TF = os.environ.get("USE_TF", "AUTO").upper() -USE_TORCH = os.environ.get("USE_TORCH", "AUTO").upper() -USE_JAX = os.environ.get("USE_FLAX", "AUTO").upper() -USE_SAFETENSORS = os.environ.get("USE_SAFETENSORS", "AUTO").upper() -DIFFUSERS_SLOW_IMPORT = os.environ.get("DIFFUSERS_SLOW_IMPORT", "FALSE").upper() -DIFFUSERS_SLOW_IMPORT = DIFFUSERS_SLOW_IMPORT in ENV_VARS_TRUE_VALUES - -STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt} - -_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]: - global _package_map - pkg_exists = importlib.util.find_spec(pkg_name) is not None - pkg_version = "N/A" - - if pkg_exists: - if _package_map is None: - _package_map = defaultdict(list) - try: - # Fallback for Python < 3.10 - for dist in importlib_metadata.distributions(): - _top_level_declared = (dist.read_text("top_level.txt") or "").split() - _infered_opt_names = { - f.parts[0] if len(f.parts) > 1 else inspect.getmodulename(f) for f in (dist.files or []) - } - {None} - _top_level_inferred = filter(lambda name: "." not in name, _infered_opt_names) - for pkg in _top_level_declared or _top_level_inferred: - _package_map[pkg].append(dist.metadata["Name"]) - except Exception as _: - pass - try: - if get_dist_name and pkg_name in _package_map and _package_map[pkg_name]: - if len(_package_map[pkg_name]) > 1: - logger.warning( - f"Multiple distributions found for package {pkg_name}. Picked distribution: {_package_map[pkg_name][0]}" - ) - pkg_name = _package_map[pkg_name][0] - pkg_version = importlib_metadata.version(pkg_name) - logger.debug(f"Successfully imported {pkg_name} version {pkg_version}") - except (ImportError, importlib_metadata.PackageNotFoundError): - pkg_exists = False - - return pkg_exists, pkg_version - - -if USE_TORCH in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TF not in ENV_VARS_TRUE_VALUES: - _torch_available, _torch_version = _is_package_available("torch") - -else: - logger.info("Disabling PyTorch because USE_TORCH is set") - _torch_available = False - _torch_version = "N/A" - -_jax_version = "N/A" -_flax_version = "N/A" -if USE_JAX in ENV_VARS_TRUE_AND_AUTO_VALUES: - _flax_available = importlib.util.find_spec("jax") is not None and importlib.util.find_spec("flax") is not None - if _flax_available: - try: - _jax_version = importlib_metadata.version("jax") - _flax_version = importlib_metadata.version("flax") - logger.info(f"JAX version {_jax_version}, Flax version {_flax_version} available.") - except importlib_metadata.PackageNotFoundError: - _flax_available = False -else: - _flax_available = False - -if USE_SAFETENSORS in ENV_VARS_TRUE_AND_AUTO_VALUES: - _safetensors_available, _safetensors_version = _is_package_available("safetensors") - -else: - logger.info("Disabling Safetensors because USE_TF is set") - _safetensors_available = False - -_onnxruntime_version = "N/A" -_onnx_available = importlib.util.find_spec("onnxruntime") is not None -if _onnx_available: - candidates = ( - "onnxruntime", - "onnxruntime-cann", - "onnxruntime-directml", - "ort_nightly_directml", - "onnxruntime-gpu", - "ort_nightly_gpu", - "onnxruntime-migraphx", - "onnxruntime-openvino", - "onnxruntime-qnn", - "onnxruntime-rocm", - "onnxruntime-training", - "onnxruntime-vitisai", - ) - _onnxruntime_version = None - # For the metadata, we have to look for both onnxruntime and onnxruntime-x - for pkg in candidates: - try: - _onnxruntime_version = importlib_metadata.version(pkg) - break - except importlib_metadata.PackageNotFoundError: - pass - _onnx_available = _onnxruntime_version is not None - if _onnx_available: - logger.debug(f"Successfully imported onnxruntime version {_onnxruntime_version}") - -# (sayakpaul): importlib.util.find_spec("opencv-python") returns None even when it's installed. -# _opencv_available = importlib.util.find_spec("opencv-python") is not None -try: - candidates = ( - "opencv-python", - "opencv-contrib-python", - "opencv-python-headless", - "opencv-contrib-python-headless", - ) - _opencv_version = None - for pkg in candidates: - try: - _opencv_version = importlib_metadata.version(pkg) - break - except importlib_metadata.PackageNotFoundError: - pass - _opencv_available = _opencv_version is not None - if _opencv_available: - logger.debug(f"Successfully imported cv2 version {_opencv_version}") -except importlib_metadata.PackageNotFoundError: - _opencv_available = False - -_bs4_available = importlib.util.find_spec("bs4") is not None -try: - # importlib metadata under different name - _bs4_version = importlib_metadata.version("beautifulsoup4") - logger.debug(f"Successfully imported ftfy version {_bs4_version}") -except importlib_metadata.PackageNotFoundError: - _bs4_available = False - -_invisible_watermark_available = importlib.util.find_spec("imwatermark") is not None -try: - _invisible_watermark_version = importlib_metadata.version("invisible-watermark") - logger.debug(f"Successfully imported invisible-watermark version {_invisible_watermark_version}") -except importlib_metadata.PackageNotFoundError: - _invisible_watermark_available = False - -_torch_xla_available, _torch_xla_version = _is_package_available("torch_xla") -_torch_npu_available, _torch_npu_version = _is_package_available("torch_npu") -_transformers_available, _transformers_version = _is_package_available("transformers") -_hf_hub_available, _hf_hub_version = _is_package_available("huggingface_hub") -_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") -_compel_available, _compel_version = _is_package_available("compel") -_sentencepiece_available, _sentencepiece_version = _is_package_available("sentencepiece") -_torchsde_available, _torchsde_version = _is_package_available("torchsde") -_peft_available, _peft_version = _is_package_available("peft") -_torchvision_available, _torchvision_version = _is_package_available("torchvision") -_matplotlib_available, _matplotlib_version = _is_package_available("matplotlib") -_timm_available, _timm_version = _is_package_available("timm") -_bitsandbytes_available, _bitsandbytes_version = _is_package_available("bitsandbytes") -_imageio_available, _imageio_version = _is_package_available("imageio") -_ftfy_available, _ftfy_version = _is_package_available("ftfy") -_scipy_available, _scipy_version = _is_package_available("scipy") -_librosa_available, _librosa_version = _is_package_available("librosa") -_accelerate_available, _accelerate_version = _is_package_available("accelerate") -_xformers_available, _xformers_version = _is_package_available("xformers") -_gguf_available, _gguf_version = _is_package_available("gguf") -_torchao_available, _torchao_version = _is_package_available("torchao") -_bitsandbytes_available, _bitsandbytes_version = _is_package_available("bitsandbytes") -_optimum_quanto_available, _optimum_quanto_version = _is_package_available("optimum", get_dist_name=True) -_pytorch_retinaface_available, _pytorch_retinaface_version = _is_package_available("pytorch_retinaface") -_better_profanity_available, _better_profanity_version = _is_package_available("better_profanity") -_nltk_available, _nltk_version = _is_package_available("nltk") -_cosmos_guardrail_available, _cosmos_guardrail_version = _is_package_available("cosmos_guardrail") - - -def is_torch_available(): - return _torch_available - - -def is_torch_xla_available(): - return _torch_xla_available - - -def is_torch_npu_available(): - return _torch_npu_available - - -def is_flax_available(): - return _flax_available - - -def is_transformers_available(): - return _transformers_available - - -def is_inflect_available(): - return _inflect_available - - -def is_unidecode_available(): - return _unidecode_available - - -def is_onnx_available(): - return _onnx_available - - -def is_opencv_available(): - return _opencv_available - - -def is_scipy_available(): - return _scipy_available - - -def is_librosa_available(): - return _librosa_available - - -def is_xformers_available(): - return _xformers_available - - -def is_accelerate_available(): - return _accelerate_available - - -def is_k_diffusion_available(): - return _k_diffusion_available - - -def is_note_seq_available(): - return _note_seq_available - - -def is_wandb_available(): - return _wandb_available - - -def is_tensorboard_available(): - return _tensorboard_available - - -def is_compel_available(): - return _compel_available - - -def is_ftfy_available(): - return _ftfy_available - - -def is_bs4_available(): - return _bs4_available - - -def is_torchsde_available(): - return _torchsde_available - - -def is_invisible_watermark_available(): - return _invisible_watermark_available - - -def is_peft_available(): - return _peft_available - - -def is_torchvision_available(): - return _torchvision_available - - -def is_matplotlib_available(): - return _matplotlib_available - - -def is_safetensors_available(): - return _safetensors_available - - -def is_bitsandbytes_available(): - return _bitsandbytes_available - - -def is_google_colab(): - return _is_google_colab - - -def is_sentencepiece_available(): - return _sentencepiece_available - - -def is_imageio_available(): - return _imageio_available - - -def is_gguf_available(): - return _gguf_available - - -def is_torchao_available(): - return _torchao_available - - -def is_optimum_quanto_available(): - return _optimum_quanto_available - - -def is_timm_available(): - return _timm_available - - -def is_pytorch_retinaface_available(): - return _pytorch_retinaface_available - - -def is_better_profanity_available(): - return _better_profanity_available - - -def is_nltk_available(): - return _nltk_available - - -def is_cosmos_guardrail_available(): - return _cosmos_guardrail_available - - -def is_hpu_available(): - return all(importlib.util.find_spec(lib) for lib in ("habana_frameworks", "habana_frameworks.torch")) - - -# docstyle-ignore -FLAX_IMPORT_ERROR = """ -{0} requires the FLAX library but it was not found in your environment. Checkout the instructions on the -installation page: https://github.com/google/flax and follow the ones that match your environment. -""" - -# docstyle-ignore -INFLECT_IMPORT_ERROR = """ -{0} requires the inflect library but it was not found in your environment. You can install it with pip: `pip install -inflect` -""" - -# docstyle-ignore -PYTORCH_IMPORT_ERROR = """ -{0} requires the PyTorch library but it was not found in your environment. Checkout the instructions on the -installation page: https://pytorch.org/get-started/locally/ and follow the ones that match your environment. -""" - -# docstyle-ignore -ONNX_IMPORT_ERROR = """ -{0} requires the onnxruntime library but it was not found in your environment. You can install it with pip: `pip -install onnxruntime` -""" - -# docstyle-ignore -OPENCV_IMPORT_ERROR = """ -{0} requires the OpenCV library but it was not found in your environment. You can install it with pip: `pip -install opencv-python` -""" - -# docstyle-ignore -SCIPY_IMPORT_ERROR = """ -{0} requires the scipy library but it was not found in your environment. You can install it with pip: `pip install -scipy` -""" - -# docstyle-ignore -LIBROSA_IMPORT_ERROR = """ -{0} requires the librosa library but it was not found in your environment. Checkout the instructions on the -installation page: https://librosa.org/doc/latest/install.html and follow the ones that match your environment. -""" - -# docstyle-ignore -TRANSFORMERS_IMPORT_ERROR = """ -{0} requires the transformers library but it was not found in your environment. You can install it with pip: `pip -install transformers` -""" - -# docstyle-ignore -UNIDECODE_IMPORT_ERROR = """ -{0} requires the unidecode library but it was not found in your environment. You can install it with pip: `pip install -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 -install note-seq` -""" - -# docstyle-ignore -WANDB_IMPORT_ERROR = """ -{0} requires the wandb library but it was not found in your environment. You can install it with pip: `pip -install wandb` -""" - -# docstyle-ignore -TENSORBOARD_IMPORT_ERROR = """ -{0} requires the tensorboard library but it was not found in your environment. You can install it with pip: `pip -install tensorboard` -""" - - -# docstyle-ignore -COMPEL_IMPORT_ERROR = """ -{0} requires the compel library but it was not found in your environment. You can install it with pip: `pip install compel` -""" - -# docstyle-ignore -BS4_IMPORT_ERROR = """ -{0} requires the Beautiful Soup library but it was not found in your environment. You can install it with pip: -`pip install beautifulsoup4`. Please note that you may need to restart your runtime after installation. -""" - -# docstyle-ignore -FTFY_IMPORT_ERROR = """ -{0} requires the ftfy library but it was not found in your environment. Checkout the instructions on the -installation section: https://github.com/rspeer/python-ftfy/tree/master#installing and follow the ones -that match your environment. Please note that you may need to restart your runtime after installation. -""" - -# docstyle-ignore -TORCHSDE_IMPORT_ERROR = """ -{0} requires the torchsde library but it was not found in your environment. You can install it with pip: `pip install torchsde` -""" - -# docstyle-ignore -INVISIBLE_WATERMARK_IMPORT_ERROR = """ -{0} requires the invisible-watermark library but it was not found in your environment. You can install it with pip: `pip install invisible-watermark>=0.2.0` -""" - -# docstyle-ignore -PEFT_IMPORT_ERROR = """ -{0} requires the peft library but it was not found in your environment. You can install it with pip: `pip install peft` -""" - -# docstyle-ignore -SAFETENSORS_IMPORT_ERROR = """ -{0} requires the safetensors library but it was not found in your environment. You can install it with pip: `pip install safetensors` -""" - -# docstyle-ignore -SENTENCEPIECE_IMPORT_ERROR = """ -{0} requires the sentencepiece library but it was not found in your environment. You can install it with pip: `pip install sentencepiece` -""" - - -# docstyle-ignore -BITSANDBYTES_IMPORT_ERROR = """ -{0} requires the bitsandbytes library but it was not found in your environment. You can install it with pip: `pip install bitsandbytes` -""" - -# docstyle-ignore -IMAGEIO_IMPORT_ERROR = """ -{0} requires the imageio library and ffmpeg but it was not found in your environment. You can install it with pip: `pip install imageio imageio-ffmpeg` -""" - -# docstyle-ignore -GGUF_IMPORT_ERROR = """ -{0} requires the gguf library but it was not found in your environment. You can install it with pip: `pip install gguf` -""" - -TORCHAO_IMPORT_ERROR = """ -{0} requires the torchao library but it was not found in your environment. You can install it with pip: `pip install -torchao` -""" - -QUANTO_IMPORT_ERROR = """ -{0} requires the optimum-quanto library but it was not found in your environment. You can install it with pip: `pip -install optimum-quanto` -""" - -# docstyle-ignore -PYTORCH_RETINAFACE_IMPORT_ERROR = """ -{0} requires the pytorch_retinaface library but it was not found in your environment. You can install it with pip: `pip install pytorch_retinaface` -""" - -# docstyle-ignore -BETTER_PROFANITY_IMPORT_ERROR = """ -{0} requires the better_profanity library but it was not found in your environment. You can install it with pip: `pip install better_profanity` -""" - -# docstyle-ignore -NLTK_IMPORT_ERROR = """ -{0} requires the nltk library but it was not found in your environment. You can install it with pip: `pip install nltk` -""" - - -BACKENDS_MAPPING = OrderedDict( - [ - ("bs4", (is_bs4_available, BS4_IMPORT_ERROR)), - ("flax", (is_flax_available, FLAX_IMPORT_ERROR)), - ("inflect", (is_inflect_available, INFLECT_IMPORT_ERROR)), - ("onnx", (is_onnx_available, ONNX_IMPORT_ERROR)), - ("opencv", (is_opencv_available, OPENCV_IMPORT_ERROR)), - ("scipy", (is_scipy_available, SCIPY_IMPORT_ERROR)), - ("torch", (is_torch_available, PYTORCH_IMPORT_ERROR)), - ("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)), - ("compel", (is_compel_available, COMPEL_IMPORT_ERROR)), - ("ftfy", (is_ftfy_available, FTFY_IMPORT_ERROR)), - ("torchsde", (is_torchsde_available, TORCHSDE_IMPORT_ERROR)), - ("invisible_watermark", (is_invisible_watermark_available, INVISIBLE_WATERMARK_IMPORT_ERROR)), - ("peft", (is_peft_available, PEFT_IMPORT_ERROR)), - ("safetensors", (is_safetensors_available, SAFETENSORS_IMPORT_ERROR)), - ("bitsandbytes", (is_bitsandbytes_available, BITSANDBYTES_IMPORT_ERROR)), - ("sentencepiece", (is_sentencepiece_available, SENTENCEPIECE_IMPORT_ERROR)), - ("imageio", (is_imageio_available, IMAGEIO_IMPORT_ERROR)), - ("gguf", (is_gguf_available, GGUF_IMPORT_ERROR)), - ("torchao", (is_torchao_available, TORCHAO_IMPORT_ERROR)), - ("quanto", (is_optimum_quanto_available, QUANTO_IMPORT_ERROR)), - ("pytorch_retinaface", (is_pytorch_retinaface_available, PYTORCH_RETINAFACE_IMPORT_ERROR)), - ("better_profanity", (is_better_profanity_available, BETTER_PROFANITY_IMPORT_ERROR)), - ("nltk", (is_nltk_available, NLTK_IMPORT_ERROR)), - ] -) - - -def requires_backends(obj, backends): - if not isinstance(backends, (list, tuple)): - backends = [backends] - - name = obj.__name__ if hasattr(obj, "__name__") else obj.__class__.__name__ - checks = (BACKENDS_MAPPING[backend] for backend in backends) - failed = [msg.format(name) for available, msg in checks if not available()] - if failed: - raise ImportError("".join(failed)) - - if name in [ - "VersatileDiffusionTextToImagePipeline", - "VersatileDiffusionPipeline", - "VersatileDiffusionDualGuidedPipeline", - "StableDiffusionImageVariationPipeline", - "UnCLIPPipeline", - ] and is_transformers_version("<", "4.25.0"): - raise ImportError( - f"You need to install `transformers>=4.25` in order to use {name}: \n```\n pip install" - " --upgrade transformers \n```" - ) - - if name in ["StableDiffusionDepth2ImgPipeline", "StableDiffusionPix2PixZeroPipeline"] and is_transformers_version( - "<", "4.26.0" - ): - raise ImportError( - f"You need to install `transformers>=4.26` in order to use {name}: \n```\n pip install" - " --upgrade transformers \n```" - ) - - -class DummyObject(type): - """ - Metaclass for the dummy objects. Any class inheriting from it will return the ImportError generated by - `requires_backend` each time a user tries to access any method of that class. - """ - - def __getattr__(cls, key): - if key.startswith("_") and key not in ["_load_connected_pipes", "_is_onnx"]: - return super().__getattr__(cls, key) - requires_backends(cls, cls._backends) - - -# 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): - """ - Compares a library version to some requirement using a given operation. - - Args: - library_or_version (`str` or `packaging.version.Version`): - A library name or a version to check. - operation (`str`): - A string representation of an operator, such as `">"` or `"<="`. - requirement_version (`str`): - The version to compare the library version against - """ - if operation not in STR_OPERATION_TO_FUNC.keys(): - raise ValueError(f"`operation` must be one of {list(STR_OPERATION_TO_FUNC.keys())}, received {operation}") - operation = STR_OPERATION_TO_FUNC[operation] - if isinstance(library_or_version, str): - library_or_version = parse(importlib_metadata.version(library_or_version)) - return operation(library_or_version, parse(requirement_version)) - - -# This function was copied from: https://github.com/huggingface/accelerate/blob/874c4967d94badd24f893064cc3bef45f57cadf7/src/accelerate/utils/versions.py#L338 -def is_torch_version(operation: str, version: str): - """ - Compares the current PyTorch version to a given reference with an operation. - - Args: - operation (`str`): - A string representation of an operator, such as `">"` or `"<="` - version (`str`): - A string version of PyTorch - """ - return compare_versions(parse(_torch_version), operation, version) - - -def is_torch_xla_version(operation: str, version: str): - """ - Compares the current torch_xla version to a given reference with an operation. - - Args: - operation (`str`): - A string representation of an operator, such as `">"` or `"<="` - version (`str`): - A string version of torch_xla - """ - if not is_torch_xla_available: - return False - return compare_versions(parse(_torch_xla_version), operation, version) - - -def is_transformers_version(operation: str, version: str): - """ - Compares the current Transformers 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 _transformers_available: - return False - return compare_versions(parse(_transformers_version), operation, version) - - -def is_hf_hub_version(operation: str, version: str): - """ - Compares the current Hugging Face Hub 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 _hf_hub_available: - return False - return compare_versions(parse(_hf_hub_version), operation, version) - - -def is_accelerate_version(operation: str, version: str): - """ - Compares the current Accelerate 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 _accelerate_available: - return False - return compare_versions(parse(_accelerate_version), operation, version) - - -def is_peft_version(operation: str, version: str): - """ - Compares the current PEFT 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 _peft_available: - return False - return compare_versions(parse(_peft_version), operation, version) - - -def is_bitsandbytes_version(operation: str, version: str): - """ - Args: - Compares the current bitsandbytes version to a given reference with an operation. - operation (`str`): - A string representation of an operator, such as `">"` or `"<="` - version (`str`): - A version string - """ - if not _bitsandbytes_available: - return False - return compare_versions(parse(_bitsandbytes_version), operation, version) - - -def is_gguf_version(operation: str, version: str): - """ - Compares the current Accelerate 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 _gguf_available: - return False - return compare_versions(parse(_gguf_version), operation, version) - - -def is_torchao_version(operation: str, version: str): - """ - Compares the current torchao 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 _torchao_available: - return False - return compare_versions(parse(_torchao_version), operation, version) - - -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) - - -def is_optimum_quanto_version(operation: str, version: str): - """ - Compares the current Accelerate 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 _optimum_quanto_available: - return False - return compare_versions(parse(_optimum_quanto_version), operation, version) - - -def get_objects_from_module(module): - """ - Returns a dict of object names and values in a module, while skipping private/internal objects - - Args: - module (ModuleType): - Module to extract the objects from. - - Returns: - dict: Dictionary of object names and corresponding values - """ - - objects = {} - for name in dir(module): - if name.startswith("_"): - continue - objects[name] = getattr(module, name) - - return objects - - -class OptionalDependencyNotAvailable(BaseException): - """ - An error indicating that an optional dependency of Diffusers was not found in the environment. - """ - - -class _LazyModule(ModuleType): - """ - Module class that surfaces all objects but only performs associated imports when the objects are requested. - """ - - # Very heavily inspired by optuna.integration._IntegrationModule - # https://github.com/optuna/optuna/blob/master/optuna/integration/__init__.py - def __init__(self, name, module_file, import_structure, module_spec=None, extra_objects=None): - super().__init__(name) - self._modules = set(import_structure.keys()) - self._class_to_module = {} - for key, values in import_structure.items(): - for value in values: - self._class_to_module[value] = key - # Needed for autocompletion in an IDE - self.__all__ = list(import_structure.keys()) + list(chain(*import_structure.values())) - self.__file__ = module_file - self.__spec__ = module_spec - self.__path__ = [os.path.dirname(module_file)] - self._objects = {} if extra_objects is None else extra_objects - self._name = name - self._import_structure = import_structure - - # Needed for autocompletion in an IDE - def __dir__(self): - result = super().__dir__() - # The elements of self.__all__ that are submodules may or may not be in the dir already, depending on whether - # they have been accessed or not. So we only add the elements of self.__all__ that are not already in the dir. - for attr in self.__all__: - if attr not in result: - result.append(attr) - return result - - def __getattr__(self, name: str) -> Any: - if name in self._objects: - return self._objects[name] - if name in self._modules: - value = self._get_module(name) - elif name in self._class_to_module.keys(): - module = self._get_module(self._class_to_module[name]) - value = getattr(module, name) - else: - raise AttributeError(f"module {self.__name__} has no attribute {name}") - - setattr(self, name, value) - return value - - def _get_module(self, module_name: str): - try: - return importlib.import_module("." + module_name, self.__name__) - except Exception as e: - raise RuntimeError( - f"Failed to import {self.__name__}.{module_name} because of the following error (look up to see its" - f" traceback):\n{e}" - ) from e - - def __reduce__(self): - return (self.__class__, (self._name, self.__file__, self._import_structure)) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/loading_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/loading_utils.py deleted file mode 100644 index dd23ae73c86175a5736af5f98d27bcd3c7a63859..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/loading_utils.py +++ /dev/null @@ -1,163 +0,0 @@ -import os -import tempfile -from typing import Any, Callable, List, Optional, Tuple, Union -from urllib.parse import unquote, urlparse - -import PIL.Image -import PIL.ImageOps -import requests - -from .constants import DIFFUSERS_REQUEST_TIMEOUT -from .import_utils import BACKENDS_MAPPING, is_imageio_available - - -def load_image( - image: Union[str, PIL.Image.Image], convert_method: Optional[Callable[[PIL.Image.Image], PIL.Image.Image]] = None -) -> PIL.Image.Image: - """ - Loads `image` to a PIL Image. - - Args: - image (`str` or `PIL.Image.Image`): - The image to convert to the PIL Image format. - convert_method (Callable[[PIL.Image.Image], PIL.Image.Image], *optional*): - A conversion method to apply to the image after loading it. When set to `None` the image will be converted - "RGB". - - Returns: - `PIL.Image.Image`: - A PIL Image. - """ - if isinstance(image, str): - if image.startswith("http://") or image.startswith("https://"): - image = PIL.Image.open(requests.get(image, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw) - elif os.path.isfile(image): - image = PIL.Image.open(image) - else: - raise ValueError( - f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {image} is not a valid path." - ) - elif isinstance(image, PIL.Image.Image): - image = image - else: - raise ValueError( - "Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image." - ) - - image = PIL.ImageOps.exif_transpose(image) - - if convert_method is not None: - image = convert_method(image) - else: - image = image.convert("RGB") - - return image - - -def load_video( - video: str, - convert_method: Optional[Callable[[List[PIL.Image.Image]], List[PIL.Image.Image]]] = 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*): - 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]`: - The video as a list of PIL images. - """ - is_url = video.startswith("http://") or video.startswith("https://") - is_file = os.path.isfile(video) - was_tempfile_created = False - - if not (is_url or is_file): - raise ValueError( - f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {video} is not a valid path." - ) - - if is_url: - response = requests.get(video, stream=True) - if response.status_code != 200: - raise ValueError(f"Failed to download video. Status code: {response.status_code}") - - parsed_url = urlparse(video) - file_name = os.path.basename(unquote(parsed_url.path)) - - suffix = os.path.splitext(file_name)[1] or ".mp4" - video_path = tempfile.NamedTemporaryFile(suffix=suffix, delete=False).name - - was_tempfile_created = True - - video_data = response.iter_content(chunk_size=8192) - with open(video_path, "wb") as f: - for chunk in video_data: - f.write(chunk) - - video = video_path - - pil_images = [] - if video.endswith(".gif"): - gif = PIL.Image.open(video) - try: - while True: - pil_images.append(gif.copy()) - gif.seek(gif.tell() + 1) - except EOFError: - pass - - else: - if is_imageio_available(): - import imageio - else: - raise ImportError(BACKENDS_MAPPING["imageio"][1].format("load_video")) - - try: - imageio.plugins.ffmpeg.get_exe() - except AttributeError: - raise AttributeError( - "`Unable to find an ffmpeg installation on your machine. Please install via `pip install imageio-ffmpeg" - ) - - with imageio.get_reader(video) as reader: - # Read all frames - for frame in reader: - pil_images.append(PIL.Image.fromarray(frame)) - - if was_tempfile_created: - os.remove(video_path) - - if convert_method is not None: - pil_images = convert_method(pil_images) - - return pil_images - - -# Taken from `transformers`. -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]: - new_module = getattr(module, split) - if new_module is None: - raise ValueError(f"{module} has no attribute {split}.") - module = new_module - tensor_name = splits[-1] - return module, tensor_name - - -def get_submodule_by_name(root_module, module_path: str): - current = root_module - parts = module_path.split(".") - for part in parts: - if part.isdigit(): - idx = int(part) - current = current[idx] # e.g., for nn.ModuleList or nn.Sequential - else: - current = getattr(current, part) - return current diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/logging.py b/venv/lib/python3.10/site-packages/diffusers/utils/logging.py deleted file mode 100644 index 2ad6d3a476076dbc7f4d3e9ab1ab9f48de24d76f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/logging.py +++ /dev/null @@ -1,340 +0,0 @@ -# coding=utf-8 -# Copyright 2025 Optuna, Hugging Face -# -# 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. -"""Logging utilities.""" - -import logging -import os -import sys -import threading -from logging import ( - CRITICAL, # NOQA - DEBUG, # NOQA - ERROR, # NOQA - FATAL, # NOQA - INFO, # NOQA - NOTSET, # NOQA - WARN, # NOQA - WARNING, # NOQA -) -from typing import Dict, Optional - -from tqdm import auto as tqdm_lib - - -_lock = threading.Lock() -_default_handler: Optional[logging.Handler] = None - -log_levels = { - "debug": logging.DEBUG, - "info": logging.INFO, - "warning": logging.WARNING, - "error": logging.ERROR, - "critical": logging.CRITICAL, -} - -_default_log_level = logging.WARNING - -_tqdm_active = True - - -def _get_default_logging_level() -> int: - """ - If DIFFUSERS_VERBOSITY env var is set to one of the valid choices return that as the new default level. If it is - not - fall back to `_default_log_level` - """ - env_level_str = os.getenv("DIFFUSERS_VERBOSITY", None) - if env_level_str: - if env_level_str in log_levels: - return log_levels[env_level_str] - else: - logging.getLogger().warning( - f"Unknown option DIFFUSERS_VERBOSITY={env_level_str}, has to be one of: {', '.join(log_levels.keys())}" - ) - return _default_log_level - - -def _get_library_name() -> str: - return __name__.split(".")[0] - - -def _get_library_root_logger() -> logging.Logger: - return logging.getLogger(_get_library_name()) - - -def _configure_library_root_logger() -> None: - global _default_handler - - with _lock: - if _default_handler: - # This library has already configured the library root logger. - return - _default_handler = logging.StreamHandler() # Set sys.stderr as stream. - - if sys.stderr: # only if sys.stderr exists, e.g. when not using pythonw in windows - _default_handler.flush = sys.stderr.flush - - # Apply our default configuration to the library root logger. - library_root_logger = _get_library_root_logger() - library_root_logger.addHandler(_default_handler) - library_root_logger.setLevel(_get_default_logging_level()) - library_root_logger.propagate = False - - -def _reset_library_root_logger() -> None: - global _default_handler - - with _lock: - if not _default_handler: - return - - library_root_logger = _get_library_root_logger() - library_root_logger.removeHandler(_default_handler) - library_root_logger.setLevel(logging.NOTSET) - _default_handler = None - - -def get_log_levels_dict() -> Dict[str, int]: - return log_levels - - -def get_logger(name: Optional[str] = None) -> logging.Logger: - """ - Return a logger with the specified name. - - This function is not supposed to be directly accessed unless you are writing a custom diffusers module. - """ - - if name is None: - name = _get_library_name() - - _configure_library_root_logger() - return logging.getLogger(name) - - -def get_verbosity() -> int: - """ - Return the current level for the 🤗 Diffusers' root logger as an `int`. - - Returns: - `int`: - Logging level integers which can be one of: - - - `50`: `diffusers.logging.CRITICAL` or `diffusers.logging.FATAL` - - `40`: `diffusers.logging.ERROR` - - `30`: `diffusers.logging.WARNING` or `diffusers.logging.WARN` - - `20`: `diffusers.logging.INFO` - - `10`: `diffusers.logging.DEBUG` - - """ - - _configure_library_root_logger() - return _get_library_root_logger().getEffectiveLevel() - - -def set_verbosity(verbosity: int) -> None: - """ - Set the verbosity level for the 🤗 Diffusers' root logger. - - Args: - verbosity (`int`): - Logging level which can be one of: - - - `diffusers.logging.CRITICAL` or `diffusers.logging.FATAL` - - `diffusers.logging.ERROR` - - `diffusers.logging.WARNING` or `diffusers.logging.WARN` - - `diffusers.logging.INFO` - - `diffusers.logging.DEBUG` - """ - - _configure_library_root_logger() - _get_library_root_logger().setLevel(verbosity) - - -def set_verbosity_info() -> None: - """Set the verbosity to the `INFO` level.""" - return set_verbosity(INFO) - - -def set_verbosity_warning() -> None: - """Set the verbosity to the `WARNING` level.""" - return set_verbosity(WARNING) - - -def set_verbosity_debug() -> None: - """Set the verbosity to the `DEBUG` level.""" - return set_verbosity(DEBUG) - - -def set_verbosity_error() -> None: - """Set the verbosity to the `ERROR` level.""" - return set_verbosity(ERROR) - - -def disable_default_handler() -> None: - """Disable the default handler of the 🤗 Diffusers' root logger.""" - - _configure_library_root_logger() - - assert _default_handler is not None - _get_library_root_logger().removeHandler(_default_handler) - - -def enable_default_handler() -> None: - """Enable the default handler of the 🤗 Diffusers' root logger.""" - - _configure_library_root_logger() - - assert _default_handler is not None - _get_library_root_logger().addHandler(_default_handler) - - -def add_handler(handler: logging.Handler) -> None: - """adds a handler to the HuggingFace Diffusers' root logger.""" - - _configure_library_root_logger() - - assert handler is not None - _get_library_root_logger().addHandler(handler) - - -def remove_handler(handler: logging.Handler) -> None: - """removes given handler from the HuggingFace Diffusers' root logger.""" - - _configure_library_root_logger() - - assert handler is not None and handler in _get_library_root_logger().handlers - _get_library_root_logger().removeHandler(handler) - - -def disable_propagation() -> None: - """ - Disable propagation of the library log outputs. Note that log propagation is disabled by default. - """ - - _configure_library_root_logger() - _get_library_root_logger().propagate = False - - -def enable_propagation() -> None: - """ - Enable propagation of the library log outputs. Please disable the HuggingFace Diffusers' default handler to prevent - double logging if the root logger has been configured. - """ - - _configure_library_root_logger() - _get_library_root_logger().propagate = True - - -def enable_explicit_format() -> None: - """ - Enable explicit formatting for every 🤗 Diffusers' logger. The explicit formatter is as follows: - ``` - [LEVELNAME|FILENAME|LINE NUMBER] TIME >> MESSAGE - ``` - All handlers currently bound to the root logger are affected by this method. - """ - handlers = _get_library_root_logger().handlers - - for handler in handlers: - formatter = logging.Formatter("[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s") - handler.setFormatter(formatter) - - -def reset_format() -> None: - """ - Resets the formatting for 🤗 Diffusers' loggers. - - All handlers currently bound to the root logger are affected by this method. - """ - handlers = _get_library_root_logger().handlers - - for handler in handlers: - handler.setFormatter(None) - - -def warning_advice(self, *args, **kwargs) -> None: - """ - This method is identical to `logger.warning()`, but if env var DIFFUSERS_NO_ADVISORY_WARNINGS=1 is set, this - warning will not be printed - """ - no_advisory_warnings = os.getenv("DIFFUSERS_NO_ADVISORY_WARNINGS", False) - if no_advisory_warnings: - return - self.warning(*args, **kwargs) - - -logging.Logger.warning_advice = warning_advice - - -class EmptyTqdm: - """Dummy tqdm which doesn't do anything.""" - - def __init__(self, *args, **kwargs): # pylint: disable=unused-argument - self._iterator = args[0] if args else None - - def __iter__(self): - return iter(self._iterator) - - def __getattr__(self, _): - """Return empty function.""" - - def empty_fn(*args, **kwargs): # pylint: disable=unused-argument - return - - return empty_fn - - def __enter__(self): - return self - - def __exit__(self, type_, value, traceback): - return - - -class _tqdm_cls: - def __call__(self, *args, **kwargs): - if _tqdm_active: - return tqdm_lib.tqdm(*args, **kwargs) - else: - return EmptyTqdm(*args, **kwargs) - - def set_lock(self, *args, **kwargs): - self._lock = None - if _tqdm_active: - return tqdm_lib.tqdm.set_lock(*args, **kwargs) - - def get_lock(self): - if _tqdm_active: - return tqdm_lib.tqdm.get_lock() - - -tqdm = _tqdm_cls() - - -def is_progress_bar_enabled() -> bool: - """Return a boolean indicating whether tqdm progress bars are enabled.""" - global _tqdm_active - return bool(_tqdm_active) - - -def enable_progress_bar() -> None: - """Enable tqdm progress bar.""" - global _tqdm_active - _tqdm_active = True - - -def disable_progress_bar() -> None: - """Disable tqdm progress bar.""" - global _tqdm_active - _tqdm_active = False diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/model_card_template.md b/venv/lib/python3.10/site-packages/diffusers/utils/model_card_template.md deleted file mode 100644 index f41b71e24e2081425d3049ae51b50036d5d28b6a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/model_card_template.md +++ /dev/null @@ -1,24 +0,0 @@ ---- -{{ card_data }} ---- - - - -{{ model_description }} - -## Intended uses & limitations - -#### How to use - -```python -# TODO: add an example code snippet for running this diffusion pipeline -``` - -#### Limitations and bias - -[TODO: provide examples of latent issues and potential remediations] - -## Training details - -[TODO: describe the data used to train the model] diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/outputs.py b/venv/lib/python3.10/site-packages/diffusers/utils/outputs.py deleted file mode 100644 index 35691496a182cacc209f300deb4fbf1b1a889cb6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/outputs.py +++ /dev/null @@ -1,138 +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. -""" -Generic utilities -""" - -from collections import OrderedDict -from dataclasses import fields, is_dataclass -from typing import Any, Tuple - -import numpy as np - -from .import_utils import is_torch_available, is_torch_version - - -def is_tensor(x) -> bool: - """ - Tests if `x` is a `torch.Tensor` or `np.ndarray`. - """ - if is_torch_available(): - import torch - - if isinstance(x, torch.Tensor): - return True - - return isinstance(x, np.ndarray) - - -class BaseOutput(OrderedDict): - """ - Base class for all model outputs as dataclass. Has a `__getitem__` that allows indexing by integer or slice (like a - tuple) or strings (like a dictionary) that will ignore the `None` attributes. Otherwise behaves like a regular - Python dictionary. - - - - You can't unpack a [`BaseOutput`] directly. Use the [`~utils.BaseOutput.to_tuple`] method to convert it to a tuple - first. - - - """ - - def __init_subclass__(cls) -> None: - """Register subclasses as pytree nodes. - - This is necessary to synchronize gradients when using `torch.nn.parallel.DistributedDataParallel` with - `static_graph=True` with modules that output `ModelOutput` subclasses. - """ - if is_torch_available(): - import torch.utils._pytree - - if is_torch_version("<", "2.2"): - torch.utils._pytree._register_pytree_node( - cls, - torch.utils._pytree._dict_flatten, - lambda values, context: cls(**torch.utils._pytree._dict_unflatten(values, context)), - ) - else: - torch.utils._pytree.register_pytree_node( - cls, - torch.utils._pytree._dict_flatten, - lambda values, context: cls(**torch.utils._pytree._dict_unflatten(values, context)), - serialized_type_name=f"{cls.__module__}.{cls.__name__}", - ) - - def __post_init__(self) -> None: - class_fields = fields(self) - - # Safety and consistency checks - if not len(class_fields): - raise ValueError(f"{self.__class__.__name__} has no fields.") - - first_field = getattr(self, class_fields[0].name) - other_fields_are_none = all(getattr(self, field.name) is None for field in class_fields[1:]) - - if other_fields_are_none and isinstance(first_field, dict): - for key, value in first_field.items(): - self[key] = value - else: - for field in class_fields: - v = getattr(self, field.name) - if v is not None: - self[field.name] = v - - def __delitem__(self, *args, **kwargs): - raise Exception(f"You cannot use ``__delitem__`` on a {self.__class__.__name__} instance.") - - def setdefault(self, *args, **kwargs): - raise Exception(f"You cannot use ``setdefault`` on a {self.__class__.__name__} instance.") - - def pop(self, *args, **kwargs): - raise Exception(f"You cannot use ``pop`` on a {self.__class__.__name__} instance.") - - def update(self, *args, **kwargs): - raise Exception(f"You cannot use ``update`` on a {self.__class__.__name__} instance.") - - def __getitem__(self, k: Any) -> Any: - if isinstance(k, str): - inner_dict = dict(self.items()) - return inner_dict[k] - else: - return self.to_tuple()[k] - - def __setattr__(self, name: Any, value: Any) -> None: - if name in self.keys() and value is not None: - # Don't call self.__setitem__ to avoid recursion errors - super().__setitem__(name, value) - super().__setattr__(name, value) - - def __setitem__(self, key, value): - # Will raise a KeyException if needed - super().__setitem__(key, value) - # Don't call self.__setattr__ to avoid recursion errors - super().__setattr__(key, value) - - def __reduce__(self): - if not is_dataclass(self): - return super().__reduce__() - callable, _args, *remaining = super().__reduce__() - args = tuple(getattr(self, field.name) for field in fields(self)) - return callable, args, *remaining - - def to_tuple(self) -> Tuple[Any, ...]: - """ - Convert self to a tuple containing all the attributes/keys that are not `None`. - """ - return tuple(self[k] for k in self.keys()) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/peft_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/peft_utils.py deleted file mode 100644 index 3907bdd5b333e1f554a2ac70cad2827a41c6cf21..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/peft_utils.py +++ /dev/null @@ -1,373 +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. -""" -PEFT utilities: Utilities related to peft library -""" - -import collections -import importlib -from typing import Optional - -from packaging import version - -from . import logging -from .import_utils import is_peft_available, is_peft_version, is_torch_available -from .torch_utils import empty_device_cache - - -logger = logging.get_logger(__name__) - -if is_torch_available(): - import torch - - -def recurse_remove_peft_layers(model): - r""" - Recursively replace all instances of `LoraLayer` with corresponding new layers in `model`. - """ - from peft.tuners.tuners_utils import BaseTunerLayer - - has_base_layer_pattern = False - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - has_base_layer_pattern = hasattr(module, "base_layer") - break - - if has_base_layer_pattern: - from peft.utils import _get_submodules - - key_list = [key for key, _ in model.named_modules() if "lora" not in key] - for key in key_list: - try: - parent, target, target_name = _get_submodules(model, key) - except AttributeError: - continue - if hasattr(target, "base_layer"): - setattr(parent, target_name, target.get_base_layer()) - else: - # This is for backwards compatibility with PEFT <= 0.6.2. - # TODO can be removed once that PEFT version is no longer supported. - from peft.tuners.lora import LoraLayer - - for name, module in model.named_children(): - if len(list(module.children())) > 0: - ## compound module, go inside it - recurse_remove_peft_layers(module) - - module_replaced = False - - if isinstance(module, LoraLayer) and isinstance(module, torch.nn.Linear): - new_module = torch.nn.Linear( - module.in_features, - module.out_features, - bias=module.bias is not None, - ).to(module.weight.device) - new_module.weight = module.weight - if module.bias is not None: - new_module.bias = module.bias - - module_replaced = True - elif isinstance(module, LoraLayer) and isinstance(module, torch.nn.Conv2d): - new_module = torch.nn.Conv2d( - module.in_channels, - module.out_channels, - module.kernel_size, - module.stride, - module.padding, - module.dilation, - module.groups, - ).to(module.weight.device) - - new_module.weight = module.weight - if module.bias is not None: - new_module.bias = module.bias - - module_replaced = True - - if module_replaced: - setattr(model, name, new_module) - del module - - empty_device_cache() - return model - - -def scale_lora_layers(model, weight): - """ - Adjust the weightage given to the LoRA layers of the model. - - Args: - model (`torch.nn.Module`): - The model to scale. - weight (`float`): - The weight to be given to the LoRA layers. - """ - from peft.tuners.tuners_utils import BaseTunerLayer - - if weight == 1.0: - return - - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - module.scale_layer(weight) - - -def unscale_lora_layers(model, weight: Optional[float] = None): - """ - Removes the previously passed weight given to the LoRA layers of the model. - - Args: - model (`torch.nn.Module`): - The model to scale. - weight (`float`, *optional*): - The weight to be given to the LoRA layers. If no scale is passed the scale of the lora layer will be - re-initialized to the correct value. If 0.0 is passed, we will re-initialize the scale with the correct - value. - """ - from peft.tuners.tuners_utils import BaseTunerLayer - - if weight is None or weight == 1.0: - return - - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - if weight != 0: - module.unscale_layer(weight) - else: - for adapter_name in module.active_adapters: - # if weight == 0 unscale should re-set the scale to the original value. - module.set_scale(adapter_name, 1.0) - - -def get_peft_kwargs(rank_dict, network_alpha_dict, peft_state_dict, is_unet=True): - rank_pattern = {} - alpha_pattern = {} - r = lora_alpha = list(rank_dict.values())[0] - - if len(set(rank_dict.values())) > 1: - # get the rank occurring the most number of times - r = collections.Counter(rank_dict.values()).most_common()[0][0] - - # for modules with rank different from the most occurring rank, add it to the `rank_pattern` - rank_pattern = dict(filter(lambda x: x[1] != r, rank_dict.items())) - rank_pattern = {k.split(".lora_B.")[0]: v for k, v in rank_pattern.items()} - - if network_alpha_dict is not None and len(network_alpha_dict) > 0: - if len(set(network_alpha_dict.values())) > 1: - # get the alpha occurring the most number of times - lora_alpha = collections.Counter(network_alpha_dict.values()).most_common()[0][0] - - # for modules with alpha different from the most occurring alpha, add it to the `alpha_pattern` - alpha_pattern = dict(filter(lambda x: x[1] != lora_alpha, network_alpha_dict.items())) - if is_unet: - alpha_pattern = { - ".".join(k.split(".lora_A.")[0].split(".")).replace(".alpha", ""): v - for k, v in alpha_pattern.items() - } - else: - alpha_pattern = {".".join(k.split(".down.")[0].split(".")[:-1]): v for k, v in alpha_pattern.items()} - else: - lora_alpha = set(network_alpha_dict.values()).pop() - - # layer names without the Diffusers specific - target_modules = list({name.split(".lora")[0] for name in peft_state_dict.keys()}) - use_dora = any("lora_magnitude_vector" in k for k in peft_state_dict) - # for now we know that the "bias" keys are only associated with `lora_B`. - lora_bias = any("lora_B" in k and k.endswith(".bias") for k in peft_state_dict) - - lora_config_kwargs = { - "r": r, - "lora_alpha": lora_alpha, - "rank_pattern": rank_pattern, - "alpha_pattern": alpha_pattern, - "target_modules": target_modules, - "use_dora": use_dora, - "lora_bias": lora_bias, - } - return lora_config_kwargs - - -def get_adapter_name(model): - from peft.tuners.tuners_utils import BaseTunerLayer - - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - return f"default_{len(module.r)}" - return "default_0" - - -def set_adapter_layers(model, enabled=True): - from peft.tuners.tuners_utils import BaseTunerLayer - - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - # The recent version of PEFT needs to call `enable_adapters` instead - if hasattr(module, "enable_adapters"): - module.enable_adapters(enabled=enabled) - else: - module.disable_adapters = not enabled - - -def delete_adapter_layers(model, adapter_name): - from peft.tuners.tuners_utils import BaseTunerLayer - - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - if hasattr(module, "delete_adapter"): - module.delete_adapter(adapter_name) - else: - raise ValueError( - "The version of PEFT you are using is not compatible, please use a version that is greater than 0.6.1" - ) - - # For transformers integration - we need to pop the adapter from the config - if getattr(model, "_hf_peft_config_loaded", False) and hasattr(model, "peft_config"): - model.peft_config.pop(adapter_name, None) - # In case all adapters are deleted, we need to delete the config - # and make sure to set the flag to False - if len(model.peft_config) == 0: - del model.peft_config - model._hf_peft_config_loaded = None - - -def set_weights_and_activate_adapters(model, adapter_names, weights): - from peft.tuners.tuners_utils import BaseTunerLayer - - def get_module_weight(weight_for_adapter, module_name): - if not isinstance(weight_for_adapter, dict): - # If weight_for_adapter is a single number, always return it. - return weight_for_adapter - - for layer_name, weight_ in weight_for_adapter.items(): - if layer_name in module_name: - return weight_ - - parts = module_name.split(".") - # e.g. key = "down_blocks.1.attentions.0" - key = f"{parts[0]}.{parts[1]}.attentions.{parts[3]}" - block_weight = weight_for_adapter.get(key, 1.0) - - return block_weight - - for module_name, module in model.named_modules(): - if isinstance(module, BaseTunerLayer): - # For backward compatibility with previous PEFT versions, set multiple active adapters - if hasattr(module, "set_adapter"): - module.set_adapter(adapter_names) - else: - module.active_adapter = adapter_names - - # Set the scaling weight for each adapter for this module - for adapter_name, weight in zip(adapter_names, weights): - module.set_scale(adapter_name, get_module_weight(weight, module_name)) - - -def check_peft_version(min_version: str) -> None: - r""" - Checks if the version of PEFT is compatible. - - Args: - version (`str`): - The version of PEFT to check against. - """ - if not is_peft_available(): - raise ValueError("PEFT is not installed. Please install it with `pip install peft`") - - is_peft_version_compatible = version.parse(importlib.metadata.version("peft")) > version.parse(min_version) - - if not is_peft_version_compatible: - raise ValueError( - f"The version of PEFT you are using is not compatible, please use a version that is greater" - f" than {min_version}" - ) - - -def _create_lora_config( - state_dict, - network_alphas, - metadata, - rank_pattern_dict, - is_unet: bool = True, -): - from peft import LoraConfig - - if metadata is not None: - lora_config_kwargs = metadata - else: - lora_config_kwargs = get_peft_kwargs( - rank_pattern_dict, network_alpha_dict=network_alphas, peft_state_dict=state_dict, is_unet=is_unet - ) - - _maybe_raise_error_for_ambiguous_keys(lora_config_kwargs) - - # Version checks for DoRA and lora_bias - if "use_dora" in lora_config_kwargs and lora_config_kwargs["use_dora"]: - if is_peft_version("<", "0.9.0"): - raise ValueError("DoRA requires PEFT >= 0.9.0. Please upgrade.") - - if "lora_bias" in lora_config_kwargs and lora_config_kwargs["lora_bias"]: - if is_peft_version("<=", "0.13.2"): - raise ValueError("lora_bias requires PEFT >= 0.14.0. Please upgrade.") - - try: - return LoraConfig(**lora_config_kwargs) - except TypeError as e: - raise TypeError("`LoraConfig` class could not be instantiated.") from e - - -def _maybe_raise_error_for_ambiguous_keys(config): - rank_pattern = config["rank_pattern"].copy() - target_modules = config["target_modules"] - - for key in list(rank_pattern.keys()): - # try to detect ambiguity - # `target_modules` can also be a str, in which case this loop would loop - # over the chars of the str. The technically correct way to match LoRA keys - # in PEFT is to use LoraModel._check_target_module_exists (lora_config, key). - # But this cuts it for now. - exact_matches = [mod for mod in target_modules if mod == key] - substring_matches = [mod for mod in target_modules if key in mod and mod != key] - - if exact_matches and substring_matches: - if is_peft_version("<", "0.14.1"): - raise ValueError( - "There are ambiguous keys present in this LoRA. To load it, please update your `peft` installation - `pip install -U peft`." - ) - - -def _maybe_warn_for_unhandled_keys(incompatible_keys, adapter_name): - warn_msg = "" - if incompatible_keys is not None: - # Check only for unexpected keys. - unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) - if unexpected_keys: - lora_unexpected_keys = [k for k in unexpected_keys if "lora_" in k and adapter_name in k] - if lora_unexpected_keys: - warn_msg = ( - f"Loading adapter weights from state_dict led to unexpected keys found in the model:" - f" {', '.join(lora_unexpected_keys)}. " - ) - - # Filter missing keys specific to the current adapter. - missing_keys = getattr(incompatible_keys, "missing_keys", None) - if missing_keys: - lora_missing_keys = [k for k in missing_keys if "lora_" in k and adapter_name in k] - if lora_missing_keys: - warn_msg += ( - f"Loading adapter weights from state_dict led to missing keys in the model:" - f" {', '.join(lora_missing_keys)}." - ) - - if warn_msg: - logger.warning(warn_msg) diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/pil_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/pil_utils.py deleted file mode 100644 index 76678070b697c7d87fc3691d9bc5bb3bea83c5b1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/pil_utils.py +++ /dev/null @@ -1,67 +0,0 @@ -from typing import List - -import PIL.Image -import PIL.ImageOps -from packaging import version -from PIL import Image - - -if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("9.1.0"): - PIL_INTERPOLATION = { - "linear": PIL.Image.Resampling.BILINEAR, - "bilinear": PIL.Image.Resampling.BILINEAR, - "bicubic": PIL.Image.Resampling.BICUBIC, - "lanczos": PIL.Image.Resampling.LANCZOS, - "nearest": PIL.Image.Resampling.NEAREST, - } -else: - PIL_INTERPOLATION = { - "linear": PIL.Image.LINEAR, - "bilinear": PIL.Image.BILINEAR, - "bicubic": PIL.Image.BICUBIC, - "lanczos": PIL.Image.LANCZOS, - "nearest": PIL.Image.NEAREST, - } - - -def pt_to_pil(images): - """ - Convert a torch image to a PIL image. - """ - images = (images / 2 + 0.5).clamp(0, 1) - images = images.cpu().permute(0, 2, 3, 1).float().numpy() - images = numpy_to_pil(images) - return images - - -def numpy_to_pil(images): - """ - Convert a numpy image or a batch of images to a PIL image. - """ - if images.ndim == 3: - images = images[None, ...] - images = (images * 255).round().astype("uint8") - if images.shape[-1] == 1: - # special case for grayscale (single channel) images - pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images] - else: - pil_images = [Image.fromarray(image) for image in images] - - return pil_images - - -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. - """ - assert len(images) == rows * cols - - if resize is not None: - images = [img.resize((resize, resize)) for img in images] - - w, h = images[0].size - grid = Image.new("RGB", size=(cols * w, rows * h)) - - for i, img in enumerate(images): - grid.paste(img, box=(i % cols * w, i // cols * h)) - return grid diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/remote_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/remote_utils.py deleted file mode 100644 index 6494dc14171a935cc7efc247f8b4870e76b575d8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/remote_utils.py +++ /dev/null @@ -1,425 +0,0 @@ -# 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 io -import json -from typing import List, Literal, Optional, Union, cast - -import requests - -from .deprecation_utils import deprecate -from .import_utils import is_safetensors_available, is_torch_available - - -if is_torch_available(): - import torch - - from ..image_processor import VaeImageProcessor - from ..video_processor import VideoProcessor - - if is_safetensors_available(): - import safetensors.torch - - DTYPE_MAP = { - "float16": torch.float16, - "float32": torch.float32, - "bfloat16": torch.bfloat16, - "uint8": torch.uint8, - } - - -from PIL import Image - - -def detect_image_type(data: bytes) -> str: - if data.startswith(b"\xff\xd8"): - return "jpeg" - elif data.startswith(b"\x89PNG\r\n\x1a\n"): - return "png" - elif data.startswith(b"GIF87a") or data.startswith(b"GIF89a"): - return "gif" - elif data.startswith(b"BM"): - return "bmp" - return "unknown" - - -def check_inputs_decode( - endpoint: str, - tensor: "torch.Tensor", - processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None, - do_scaling: bool = True, - scaling_factor: Optional[float] = None, - shift_factor: Optional[float] = 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, -): - if tensor.ndim == 3 and height is None and width is None: - raise ValueError("`height` and `width` required for packed latents.") - if ( - output_type == "pt" - and return_type == "pil" - and not partial_postprocess - and not isinstance(processor, (VaeImageProcessor, VideoProcessor)) - ): - raise ValueError("`processor` is required.") - if do_scaling and scaling_factor is None: - deprecate( - "do_scaling", - "1.0.0", - "`do_scaling` is deprecated, pass `scaling_factor` and `shift_factor` if required.", - standard_warn=False, - ) - - -def postprocess_decode( - response: requests.Response, - processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None, - output_type: Literal["mp4", "pil", "pt"] = "pil", - return_type: Literal["mp4", "pil", "pt"] = "pil", - partial_postprocess: bool = False, -): - if output_type == "pt" or (output_type == "pil" and processor is not None): - output_tensor = response.content - parameters = response.headers - shape = json.loads(parameters["shape"]) - dtype = parameters["dtype"] - torch_dtype = DTYPE_MAP[dtype] - output_tensor = torch.frombuffer(bytearray(output_tensor), dtype=torch_dtype).reshape(shape) - if output_type == "pt": - if partial_postprocess: - if return_type == "pil": - output = [Image.fromarray(image.numpy()) for image in output_tensor] - if len(output) == 1: - output = output[0] - elif return_type == "pt": - output = output_tensor - else: - if processor is None or return_type == "pt": - output = output_tensor - else: - if isinstance(processor, VideoProcessor): - output = cast( - List[Image.Image], - processor.postprocess_video(output_tensor, output_type="pil")[0], - ) - else: - output = cast( - Image.Image, - processor.postprocess(output_tensor, output_type="pil")[0], - ) - elif output_type == "pil" and return_type == "pil" and processor is None: - output = Image.open(io.BytesIO(response.content)).convert("RGB") - detected_format = detect_image_type(response.content) - output.format = detected_format - elif output_type == "pil" and processor is not None: - if return_type == "pil": - output = [ - Image.fromarray(image) - for image in (output_tensor.permute(0, 2, 3, 1).float().numpy() * 255).round().astype("uint8") - ] - elif return_type == "pt": - output = output_tensor - elif output_type == "mp4" and return_type == "mp4": - output = response.content - return output - - -def prepare_decode( - tensor: "torch.Tensor", - processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None, - do_scaling: bool = True, - scaling_factor: Optional[float] = None, - shift_factor: Optional[float] = 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, -): - headers = {} - parameters = { - "image_format": image_format, - "output_type": output_type, - "partial_postprocess": partial_postprocess, - "shape": list(tensor.shape), - "dtype": str(tensor.dtype).split(".")[-1], - } - if do_scaling and scaling_factor is not None: - parameters["scaling_factor"] = scaling_factor - if do_scaling and shift_factor is not None: - parameters["shift_factor"] = shift_factor - if do_scaling and scaling_factor is None: - parameters["do_scaling"] = do_scaling - elif do_scaling and scaling_factor is None and shift_factor is None: - parameters["do_scaling"] = do_scaling - if height is not None and width is not None: - parameters["height"] = height - parameters["width"] = width - headers["Content-Type"] = "tensor/binary" - headers["Accept"] = "tensor/binary" - if output_type == "pil" and image_format == "jpg" and processor is None: - headers["Accept"] = "image/jpeg" - elif output_type == "pil" and image_format == "png" and processor is None: - headers["Accept"] = "image/png" - elif output_type == "mp4": - headers["Accept"] = "text/plain" - tensor_data = safetensors.torch._tobytes(tensor, "tensor") - return {"data": tensor_data, "params": parameters, "headers": headers} - - -def remote_decode( - endpoint: str, - tensor: "torch.Tensor", - processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None, - do_scaling: bool = True, - scaling_factor: Optional[float] = None, - shift_factor: Optional[float] = 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"]: - """ - Hugging Face Hybrid Inference that allow running VAE decode remotely. - - Args: - endpoint (`str`): - Endpoint for Remote Decode. - tensor (`torch.Tensor`): - Tensor to be decoded. - processor (`VaeImageProcessor` or `VideoProcessor`, *optional*): - Used with `return_type="pt"`, and `return_type="pil"` for Video models. - do_scaling (`bool`, default `True`, *optional*): - **DEPRECATED**. **pass `scaling_factor`/`shift_factor` instead.** **still set - do_scaling=None/do_scaling=False for no scaling until option is removed** When `True` scaling e.g. `latents - / self.vae.config.scaling_factor` is applied remotely. If `False`, input must be passed with scaling - applied. - scaling_factor (`float`, *optional*): - Scaling is applied when passed e.g. [`latents / - self.vae.config.scaling_factor`](https://github.com/huggingface/diffusers/blob/7007febae5cff000d4df9059d9cf35133e8b2ca9/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L1083C37-L1083C77). - - SD v1: 0.18215 - - SD XL: 0.13025 - - Flux: 0.3611 - If `None`, input must be passed with scaling applied. - shift_factor (`float`, *optional*): - Shift is applied when passed e.g. `latents + self.vae.config.shift_factor`. - - Flux: 0.1159 - If `None`, input must be passed with scaling applied. - output_type (`"mp4"` or `"pil"` or `"pt", default `"pil"): - **Endpoint** output type. Subject to change. Report feedback on preferred type. - - `"mp4": Supported by video models. Endpoint returns `bytes` of video. `"pil"`: Supported by image and video - models. - Image models: Endpoint returns `bytes` of an image in `image_format`. Video models: Endpoint returns - `torch.Tensor` with partial `postprocessing` applied. - Requires `processor` as a flag (any `None` value will work). - `"pt"`: Support by image and video models. Endpoint returns `torch.Tensor`. - With `partial_postprocess=True` the tensor is postprocessed `uint8` image tensor. - - Recommendations: - `"pt"` with `partial_postprocess=True` is the smallest transfer for full quality. `"pt"` with - `partial_postprocess=False` is the most compatible with third party code. `"pil"` with - `image_format="jpg"` is the smallest transfer overall. - - return_type (`"mp4"` or `"pil"` or `"pt", default `"pil"): - **Function** return type. - - `"mp4": Function returns `bytes` of video. `"pil"`: Function returns `PIL.Image.Image`. - With `output_type="pil" no further processing is applied. With `output_type="pt" a `PIL.Image.Image` is - created. - `partial_postprocess=False` `processor` is required. `partial_postprocess=True` `processor` is - **not** required. - `"pt"`: Function returns `torch.Tensor`. - `processor` is **not** required. `partial_postprocess=False` tensor is `float16` or `bfloat16`, without - denormalization. `partial_postprocess=True` tensor is `uint8`, denormalized. - - image_format (`"png"` or `"jpg"`, default `jpg`): - Used with `output_type="pil"`. Endpoint returns `jpg` or `png`. - - partial_postprocess (`bool`, default `False`): - Used with `output_type="pt"`. `partial_postprocess=False` tensor is `float16` or `bfloat16`, without - denormalization. `partial_postprocess=True` tensor is `uint8`, denormalized. - - input_tensor_type (`"binary"`, default `"binary"`): - Tensor transfer type. - - output_tensor_type (`"binary"`, default `"binary"`): - Tensor transfer type. - - height (`int`, **optional**): - Required for `"packed"` latents. - - width (`int`, **optional**): - Required for `"packed"` latents. - - Returns: - output (`Image.Image` or `List[Image.Image]` or `bytes` or `torch.Tensor`). - """ - if input_tensor_type == "base64": - deprecate( - "input_tensor_type='base64'", - "1.0.0", - "input_tensor_type='base64' is deprecated. Using `binary`.", - standard_warn=False, - ) - input_tensor_type = "binary" - if output_tensor_type == "base64": - deprecate( - "output_tensor_type='base64'", - "1.0.0", - "output_tensor_type='base64' is deprecated. Using `binary`.", - standard_warn=False, - ) - output_tensor_type = "binary" - check_inputs_decode( - endpoint, - tensor, - processor, - do_scaling, - scaling_factor, - shift_factor, - output_type, - return_type, - image_format, - partial_postprocess, - input_tensor_type, - output_tensor_type, - height, - width, - ) - kwargs = prepare_decode( - tensor=tensor, - processor=processor, - do_scaling=do_scaling, - scaling_factor=scaling_factor, - shift_factor=shift_factor, - output_type=output_type, - image_format=image_format, - partial_postprocess=partial_postprocess, - height=height, - width=width, - ) - response = requests.post(endpoint, **kwargs) - if not response.ok: - raise RuntimeError(response.json()) - output = postprocess_decode( - response=response, - processor=processor, - output_type=output_type, - return_type=return_type, - partial_postprocess=partial_postprocess, - ) - return output - - -def check_inputs_encode( - endpoint: str, - image: Union["torch.Tensor", Image.Image], - scaling_factor: Optional[float] = None, - shift_factor: Optional[float] = None, -): - pass - - -def postprocess_encode( - response: requests.Response, -): - output_tensor = response.content - parameters = response.headers - shape = json.loads(parameters["shape"]) - dtype = parameters["dtype"] - torch_dtype = DTYPE_MAP[dtype] - output_tensor = torch.frombuffer(bytearray(output_tensor), dtype=torch_dtype).reshape(shape) - return output_tensor - - -def prepare_encode( - image: Union["torch.Tensor", Image.Image], - scaling_factor: Optional[float] = None, - shift_factor: Optional[float] = None, -): - headers = {} - parameters = {} - if scaling_factor is not None: - parameters["scaling_factor"] = scaling_factor - if shift_factor is not None: - parameters["shift_factor"] = shift_factor - if isinstance(image, torch.Tensor): - data = safetensors.torch._tobytes(image.contiguous(), "tensor") - parameters["shape"] = list(image.shape) - parameters["dtype"] = str(image.dtype).split(".")[-1] - else: - buffer = io.BytesIO() - image.save(buffer, format="PNG") - data = buffer.getvalue() - return {"data": data, "params": parameters, "headers": headers} - - -def remote_encode( - endpoint: str, - image: Union["torch.Tensor", Image.Image], - scaling_factor: Optional[float] = None, - shift_factor: Optional[float] = None, -) -> "torch.Tensor": - """ - Hugging Face Hybrid Inference that allow running VAE encode remotely. - - Args: - endpoint (`str`): - Endpoint for Remote Decode. - image (`torch.Tensor` or `PIL.Image.Image`): - Image to be encoded. - scaling_factor (`float`, *optional*): - Scaling is applied when passed e.g. [`latents * self.vae.config.scaling_factor`]. - - SD v1: 0.18215 - - SD XL: 0.13025 - - Flux: 0.3611 - If `None`, input must be passed with scaling applied. - shift_factor (`float`, *optional*): - Shift is applied when passed e.g. `latents - self.vae.config.shift_factor`. - - Flux: 0.1159 - If `None`, input must be passed with scaling applied. - - Returns: - output (`torch.Tensor`). - """ - check_inputs_encode( - endpoint, - image, - scaling_factor, - shift_factor, - ) - kwargs = prepare_encode( - image=image, - scaling_factor=scaling_factor, - shift_factor=shift_factor, - ) - response = requests.post(endpoint, **kwargs) - if not response.ok: - raise RuntimeError(response.json()) - output = postprocess_encode( - response=response, - ) - return output diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/source_code_parsing_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/source_code_parsing_utils.py deleted file mode 100644 index 5f94711c21d825f4baa0b32bbc37fd94c89aba03..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/source_code_parsing_utils.py +++ /dev/null @@ -1,52 +0,0 @@ -import ast -import importlib -import inspect -import textwrap - - -class ReturnNameVisitor(ast.NodeVisitor): - """Thanks to ChatGPT for pairing.""" - - def __init__(self): - self.return_names = [] - - def visit_Return(self, node): - # Check if the return value is a tuple. - if isinstance(node.value, ast.Tuple): - for elt in node.value.elts: - if isinstance(elt, ast.Name): - self.return_names.append(elt.id) - else: - try: - self.return_names.append(ast.unparse(elt)) - except Exception: - self.return_names.append(str(elt)) - else: - if isinstance(node.value, ast.Name): - self.return_names.append(node.value.id) - else: - try: - self.return_names.append(ast.unparse(node.value)) - except Exception: - self.return_names.append(str(node.value)) - self.generic_visit(node) - - def _determine_parent_module(self, cls): - from diffusers import DiffusionPipeline - from diffusers.models.modeling_utils import ModelMixin - - if issubclass(cls, DiffusionPipeline): - return "pipelines" - elif issubclass(cls, ModelMixin): - return "models" - else: - raise NotImplementedError - - def get_ast_tree(self, cls, attribute_name="encode_prompt"): - parent_module_name = self._determine_parent_module(cls) - main_module = importlib.import_module(f"diffusers.{parent_module_name}") - current_cls_module = getattr(main_module, cls.__name__) - source_code = inspect.getsource(getattr(current_cls_module, attribute_name)) - source_code = textwrap.dedent(source_code) - tree = ast.parse(source_code) - return tree diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/state_dict_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/state_dict_utils.py deleted file mode 100644 index 50bfce8b15eb09fc2289f7eefa1d497d97bdbd5b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/state_dict_utils.py +++ /dev/null @@ -1,366 +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. -""" -State dict utilities: utility methods for converting state dicts easily -""" - -import enum -import json - -from .import_utils import is_torch_available -from .logging import get_logger - - -if is_torch_available(): - import torch - - -logger = get_logger(__name__) - - -class StateDictType(enum.Enum): - """ - The mode to use when converting state dicts. - """ - - DIFFUSERS_OLD = "diffusers_old" - KOHYA_SS = "kohya_ss" - PEFT = "peft" - DIFFUSERS = "diffusers" - - -# We need to define a proper mapping for Unet since it uses different output keys than text encoder -# e.g. to_q_lora -> q_proj / to_q -UNET_TO_DIFFUSERS = { - ".to_out_lora.up": ".to_out.0.lora_B", - ".to_out_lora.down": ".to_out.0.lora_A", - ".to_q_lora.down": ".to_q.lora_A", - ".to_q_lora.up": ".to_q.lora_B", - ".to_k_lora.down": ".to_k.lora_A", - ".to_k_lora.up": ".to_k.lora_B", - ".to_v_lora.down": ".to_v.lora_A", - ".to_v_lora.up": ".to_v.lora_B", - ".lora.up": ".lora_B", - ".lora.down": ".lora_A", - ".to_out.lora_magnitude_vector": ".to_out.0.lora_magnitude_vector", -} - - -DIFFUSERS_TO_PEFT = { - ".q_proj.lora_linear_layer.up": ".q_proj.lora_B", - ".q_proj.lora_linear_layer.down": ".q_proj.lora_A", - ".k_proj.lora_linear_layer.up": ".k_proj.lora_B", - ".k_proj.lora_linear_layer.down": ".k_proj.lora_A", - ".v_proj.lora_linear_layer.up": ".v_proj.lora_B", - ".v_proj.lora_linear_layer.down": ".v_proj.lora_A", - ".out_proj.lora_linear_layer.up": ".out_proj.lora_B", - ".out_proj.lora_linear_layer.down": ".out_proj.lora_A", - ".lora_linear_layer.up": ".lora_B", - ".lora_linear_layer.down": ".lora_A", - "text_projection.lora.down.weight": "text_projection.lora_A.weight", - "text_projection.lora.up.weight": "text_projection.lora_B.weight", -} - -DIFFUSERS_OLD_TO_PEFT = { - ".to_q_lora.up": ".q_proj.lora_B", - ".to_q_lora.down": ".q_proj.lora_A", - ".to_k_lora.up": ".k_proj.lora_B", - ".to_k_lora.down": ".k_proj.lora_A", - ".to_v_lora.up": ".v_proj.lora_B", - ".to_v_lora.down": ".v_proj.lora_A", - ".to_out_lora.up": ".out_proj.lora_B", - ".to_out_lora.down": ".out_proj.lora_A", - ".lora_linear_layer.up": ".lora_B", - ".lora_linear_layer.down": ".lora_A", -} - -PEFT_TO_DIFFUSERS = { - ".q_proj.lora_B": ".q_proj.lora_linear_layer.up", - ".q_proj.lora_A": ".q_proj.lora_linear_layer.down", - ".k_proj.lora_B": ".k_proj.lora_linear_layer.up", - ".k_proj.lora_A": ".k_proj.lora_linear_layer.down", - ".v_proj.lora_B": ".v_proj.lora_linear_layer.up", - ".v_proj.lora_A": ".v_proj.lora_linear_layer.down", - ".out_proj.lora_B": ".out_proj.lora_linear_layer.up", - ".out_proj.lora_A": ".out_proj.lora_linear_layer.down", - "to_k.lora_A": "to_k.lora.down", - "to_k.lora_B": "to_k.lora.up", - "to_q.lora_A": "to_q.lora.down", - "to_q.lora_B": "to_q.lora.up", - "to_v.lora_A": "to_v.lora.down", - "to_v.lora_B": "to_v.lora.up", - "to_out.0.lora_A": "to_out.0.lora.down", - "to_out.0.lora_B": "to_out.0.lora.up", -} - -DIFFUSERS_OLD_TO_DIFFUSERS = { - ".to_q_lora.up": ".q_proj.lora_linear_layer.up", - ".to_q_lora.down": ".q_proj.lora_linear_layer.down", - ".to_k_lora.up": ".k_proj.lora_linear_layer.up", - ".to_k_lora.down": ".k_proj.lora_linear_layer.down", - ".to_v_lora.up": ".v_proj.lora_linear_layer.up", - ".to_v_lora.down": ".v_proj.lora_linear_layer.down", - ".to_out_lora.up": ".out_proj.lora_linear_layer.up", - ".to_out_lora.down": ".out_proj.lora_linear_layer.down", - ".to_k.lora_magnitude_vector": ".k_proj.lora_magnitude_vector", - ".to_v.lora_magnitude_vector": ".v_proj.lora_magnitude_vector", - ".to_q.lora_magnitude_vector": ".q_proj.lora_magnitude_vector", - ".to_out.lora_magnitude_vector": ".out_proj.lora_magnitude_vector", -} - -PEFT_TO_KOHYA_SS = { - "lora_A": "lora_down", - "lora_B": "lora_up", - # This is not a comprehensive dict as kohya format requires replacing `.` with `_` in keys, - # adding prefixes and adding alpha values - # Check `convert_state_dict_to_kohya` for more -} - -PEFT_STATE_DICT_MAPPINGS = { - StateDictType.DIFFUSERS_OLD: DIFFUSERS_OLD_TO_PEFT, - StateDictType.DIFFUSERS: DIFFUSERS_TO_PEFT, -} - -DIFFUSERS_STATE_DICT_MAPPINGS = { - StateDictType.DIFFUSERS_OLD: DIFFUSERS_OLD_TO_DIFFUSERS, - StateDictType.PEFT: PEFT_TO_DIFFUSERS, -} - -KOHYA_STATE_DICT_MAPPINGS = {StateDictType.PEFT: PEFT_TO_KOHYA_SS} - -KEYS_TO_ALWAYS_REPLACE = { - ".processor.": ".", -} - - -def convert_state_dict(state_dict, mapping): - r""" - Simply iterates over the state dict and replaces the patterns in `mapping` with the corresponding values. - - Args: - state_dict (`dict[str, torch.Tensor]`): - The state dict to convert. - mapping (`dict[str, str]`): - The mapping to use for conversion, the mapping should be a dictionary with the following structure: - - key: the pattern to replace - - value: the pattern to replace with - - Returns: - converted_state_dict (`dict`) - The converted state dict. - """ - converted_state_dict = {} - for k, v in state_dict.items(): - # First, filter out the keys that we always want to replace - for pattern in KEYS_TO_ALWAYS_REPLACE.keys(): - if pattern in k: - new_pattern = KEYS_TO_ALWAYS_REPLACE[pattern] - k = k.replace(pattern, new_pattern) - - for pattern in mapping.keys(): - if pattern in k: - new_pattern = mapping[pattern] - k = k.replace(pattern, new_pattern) - break - converted_state_dict[k] = v - return converted_state_dict - - -def convert_state_dict_to_peft(state_dict, original_type=None, **kwargs): - r""" - Converts a state dict to the PEFT format The state dict can be from previous diffusers format (`OLD_DIFFUSERS`), or - new diffusers format (`DIFFUSERS`). The method only supports the conversion from diffusers old/new to PEFT for now. - - Args: - state_dict (`dict[str, torch.Tensor]`): - The state dict to convert. - original_type (`StateDictType`, *optional*): - The original type of the state dict, if not provided, the method will try to infer it automatically. - """ - if original_type is None: - # Old diffusers to PEFT - if any("to_out_lora" in k for k in state_dict.keys()): - original_type = StateDictType.DIFFUSERS_OLD - elif any("lora_linear_layer" in k for k in state_dict.keys()): - original_type = StateDictType.DIFFUSERS - else: - raise ValueError("Could not automatically infer state dict type") - - if original_type not in PEFT_STATE_DICT_MAPPINGS.keys(): - raise ValueError(f"Original type {original_type} is not supported") - - mapping = PEFT_STATE_DICT_MAPPINGS[original_type] - return convert_state_dict(state_dict, mapping) - - -def convert_state_dict_to_diffusers(state_dict, original_type=None, **kwargs): - r""" - Converts a state dict to new diffusers format. The state dict can be from previous diffusers format - (`OLD_DIFFUSERS`), or PEFT format (`PEFT`) or new diffusers format (`DIFFUSERS`). In the last case the method will - return the state dict as is. - - The method only supports the conversion from diffusers old, PEFT to diffusers new for now. - - Args: - state_dict (`dict[str, torch.Tensor]`): - The state dict to convert. - original_type (`StateDictType`, *optional*): - The original type of the state dict, if not provided, the method will try to infer it automatically. - kwargs (`dict`, *args*): - Additional arguments to pass to the method. - - - **adapter_name**: For example, in case of PEFT, some keys will be prepended - with the adapter name, therefore needs a special handling. By default PEFT also takes care of that in - `get_peft_model_state_dict` method: - https://github.com/huggingface/peft/blob/ba0477f2985b1ba311b83459d29895c809404e99/src/peft/utils/save_and_load.py#L92 - but we add it here in case we don't want to rely on that method. - """ - peft_adapter_name = kwargs.pop("adapter_name", None) - if peft_adapter_name is not None: - peft_adapter_name = "." + peft_adapter_name - else: - peft_adapter_name = "" - - if original_type is None: - # Old diffusers to PEFT - if any("to_out_lora" in k for k in state_dict.keys()): - original_type = StateDictType.DIFFUSERS_OLD - elif any(f".lora_A{peft_adapter_name}.weight" in k for k in state_dict.keys()): - original_type = StateDictType.PEFT - elif any("lora_linear_layer" in k for k in state_dict.keys()): - # nothing to do - return state_dict - else: - raise ValueError("Could not automatically infer state dict type") - - if original_type not in DIFFUSERS_STATE_DICT_MAPPINGS.keys(): - raise ValueError(f"Original type {original_type} is not supported") - - mapping = DIFFUSERS_STATE_DICT_MAPPINGS[original_type] - return convert_state_dict(state_dict, mapping) - - -def convert_unet_state_dict_to_peft(state_dict): - r""" - Converts a state dict from UNet format to diffusers format - i.e. by removing some keys - """ - mapping = UNET_TO_DIFFUSERS - return convert_state_dict(state_dict, mapping) - - -def convert_all_state_dict_to_peft(state_dict): - r""" - Attempts to first `convert_state_dict_to_peft`, and if it doesn't detect `lora_linear_layer` for a valid - `DIFFUSERS` LoRA for example, attempts to exclusively convert the Unet `convert_unet_state_dict_to_peft` - """ - try: - peft_dict = convert_state_dict_to_peft(state_dict) - except Exception as e: - if str(e) == "Could not automatically infer state dict type": - peft_dict = convert_unet_state_dict_to_peft(state_dict) - else: - raise - - if not any("lora_A" in key or "lora_B" in key for key in peft_dict.keys()): - raise ValueError("Your LoRA was not converted to PEFT") - - return peft_dict - - -def convert_state_dict_to_kohya(state_dict, original_type=None, **kwargs): - r""" - Converts a `PEFT` state dict to `Kohya` format that can be used in AUTOMATIC1111, ComfyUI, SD.Next, InvokeAI, etc. - The method only supports the conversion from PEFT to Kohya for now. - - Args: - state_dict (`dict[str, torch.Tensor]`): - The state dict to convert. - original_type (`StateDictType`, *optional*): - The original type of the state dict, if not provided, the method will try to infer it automatically. - kwargs (`dict`, *args*): - Additional arguments to pass to the method. - - - **adapter_name**: For example, in case of PEFT, some keys will be prepended - with the adapter name, therefore needs a special handling. By default PEFT also takes care of that in - `get_peft_model_state_dict` method: - https://github.com/huggingface/peft/blob/ba0477f2985b1ba311b83459d29895c809404e99/src/peft/utils/save_and_load.py#L92 - but we add it here in case we don't want to rely on that method. - """ - try: - import torch - except ImportError: - logger.error("Converting PEFT state dicts to Kohya requires torch to be installed.") - raise - - peft_adapter_name = kwargs.pop("adapter_name", None) - if peft_adapter_name is not None: - peft_adapter_name = "." + peft_adapter_name - else: - peft_adapter_name = "" - - if original_type is None: - if any(f".lora_A{peft_adapter_name}.weight" in k for k in state_dict.keys()): - original_type = StateDictType.PEFT - - if original_type not in KOHYA_STATE_DICT_MAPPINGS.keys(): - raise ValueError(f"Original type {original_type} is not supported") - - # Use the convert_state_dict function with the appropriate mapping - kohya_ss_partial_state_dict = convert_state_dict(state_dict, KOHYA_STATE_DICT_MAPPINGS[StateDictType.PEFT]) - kohya_ss_state_dict = {} - - # Additional logic for replacing header, alpha parameters `.` with `_` in all keys - for kohya_key, weight in kohya_ss_partial_state_dict.items(): - if "text_encoder_2." in kohya_key: - kohya_key = kohya_key.replace("text_encoder_2.", "lora_te2.") - elif "text_encoder." in kohya_key: - kohya_key = kohya_key.replace("text_encoder.", "lora_te1.") - elif "unet" in kohya_key: - kohya_key = kohya_key.replace("unet", "lora_unet") - elif "lora_magnitude_vector" in kohya_key: - kohya_key = kohya_key.replace("lora_magnitude_vector", "dora_scale") - - kohya_key = kohya_key.replace(".", "_", kohya_key.count(".") - 2) - kohya_key = kohya_key.replace(peft_adapter_name, "") # Kohya doesn't take names - kohya_ss_state_dict[kohya_key] = weight - if "lora_down" in kohya_key: - alpha_key = f"{kohya_key.split('.')[0]}.alpha" - kohya_ss_state_dict[alpha_key] = torch.tensor(len(weight)) - - return kohya_ss_state_dict - - -def state_dict_all_zero(state_dict, filter_str=None): - if filter_str is not None: - if isinstance(filter_str, str): - filter_str = [filter_str] - state_dict = {k: v for k, v in state_dict.items() if any(f in k for f in filter_str)} - - return all(torch.all(param == 0).item() for param in state_dict.values()) - - -def _load_sft_state_dict_metadata(model_file: str): - import safetensors.torch - - from ..loaders.lora_base import LORA_ADAPTER_METADATA_KEY - - with safetensors.torch.safe_open(model_file, framework="pt", device="cpu") as f: - metadata = f.metadata() or {} - - metadata.pop("format", None) - if metadata: - raw = metadata.get(LORA_ADAPTER_METADATA_KEY) - return json.loads(raw) if raw else None - else: - return None diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/testing_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/testing_utils.py deleted file mode 100644 index e5da39c1d8654aaedc7b81c78369bd21b2053c31..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/testing_utils.py +++ /dev/null @@ -1,1431 +0,0 @@ -import functools -import importlib -import importlib.metadata -import inspect -import io -import logging -import multiprocessing -import os -import random -import re -import struct -import sys -import tempfile -import time -import unittest -import urllib.parse -from collections import UserDict -from contextlib import contextmanager -from io import BytesIO, StringIO -from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union - -import numpy as np -import PIL.Image -import PIL.ImageOps -import requests -from numpy.linalg import norm -from packaging import version - -from .constants import DIFFUSERS_REQUEST_TIMEOUT -from .import_utils import ( - BACKENDS_MAPPING, - is_accelerate_available, - is_bitsandbytes_available, - is_compel_available, - is_flax_available, - is_gguf_available, - is_note_seq_available, - is_onnx_available, - is_opencv_available, - is_optimum_quanto_available, - is_peft_available, - is_timm_available, - is_torch_available, - is_torch_version, - is_torchao_available, - is_torchsde_available, - is_transformers_available, -) -from .logging import get_logger - - -if is_torch_available(): - import torch - - IS_ROCM_SYSTEM = torch.version.hip is not None - IS_CUDA_SYSTEM = torch.version.cuda is not None - IS_XPU_SYSTEM = getattr(torch.version, "xpu", None) is not None -else: - IS_ROCM_SYSTEM = False - IS_CUDA_SYSTEM = False - IS_XPU_SYSTEM = False - -global_rng = random.Random() - -logger = get_logger(__name__) - -_required_peft_version = is_peft_available() and version.parse( - version.parse(importlib.metadata.version("peft")).base_version -) > version.parse("0.5") -_required_transformers_version = is_transformers_available() and version.parse( - version.parse(importlib.metadata.version("transformers")).base_version -) > version.parse("4.33") - -USE_PEFT_BACKEND = _required_peft_version and _required_transformers_version -BIG_GPU_MEMORY = int(os.getenv("BIG_GPU_MEMORY", 40)) - -if is_torch_available(): - import torch - - # Set a backend environment variable for any extra module import required for a custom accelerator - if "DIFFUSERS_TEST_BACKEND" in os.environ: - backend = os.environ["DIFFUSERS_TEST_BACKEND"] - try: - _ = importlib.import_module(backend) - except ModuleNotFoundError as e: - raise ModuleNotFoundError( - f"Failed to import `DIFFUSERS_TEST_BACKEND` '{backend}'! This should be the name of an installed module \ - to enable a specified backend.):\n{e}" - ) from e - - if "DIFFUSERS_TEST_DEVICE" in os.environ: - torch_device = os.environ["DIFFUSERS_TEST_DEVICE"] - try: - # try creating device to see if provided device is valid - _ = torch.device(torch_device) - except RuntimeError as e: - raise RuntimeError( - f"Unknown testing device specified by environment variable `DIFFUSERS_TEST_DEVICE`: {torch_device}" - ) from e - logger.info(f"torch_device overrode to {torch_device}") - else: - if torch.cuda.is_available(): - torch_device = "cuda" - elif torch.xpu.is_available(): - torch_device = "xpu" - else: - torch_device = "cpu" - is_torch_higher_equal_than_1_12 = version.parse( - version.parse(torch.__version__).base_version - ) >= version.parse("1.12") - - if is_torch_higher_equal_than_1_12: - # Some builds of torch 1.12 don't have the mps backend registered. See #892 for more details - mps_backend_registered = hasattr(torch.backends, "mps") - torch_device = "mps" if (mps_backend_registered and torch.backends.mps.is_available()) else torch_device - - from .torch_utils import get_torch_cuda_device_capability - - -def torch_all_close(a, b, *args, **kwargs): - if not is_torch_available(): - raise ValueError("PyTorch needs to be installed to use this function.") - if not torch.allclose(a, b, *args, **kwargs): - assert False, f"Max diff is absolute {(a - b).abs().max()}. Diff tensor is {(a - b).abs()}." - return True - - -def numpy_cosine_similarity_distance(a, b): - similarity = np.dot(a, b) / (norm(a) * norm(b)) - distance = 1.0 - similarity.mean() - - return distance - - -def check_if_dicts_are_equal(dict1, dict2): - dict1, dict2 = dict1.copy(), dict2.copy() - - for key, value in dict1.items(): - if isinstance(value, set): - dict1[key] = sorted(value) - for key, value in dict2.items(): - if isinstance(value, set): - dict2[key] = sorted(value) - - for key in dict1: - if key not in dict2: - return False - if dict1[key] != dict2[key]: - return False - - for key in dict2: - if key not in dict1: - return False - - return True - - -def print_tensor_test( - tensor, - limit_to_slices=None, - max_torch_print=None, - filename="test_corrections.txt", - expected_tensor_name="expected_slice", -): - if max_torch_print: - torch.set_printoptions(threshold=10_000) - - test_name = os.environ.get("PYTEST_CURRENT_TEST") - if not torch.is_tensor(tensor): - tensor = torch.from_numpy(tensor) - if limit_to_slices: - tensor = tensor[0, -3:, -3:, -1] - - tensor_str = str(tensor.detach().cpu().flatten().to(torch.float32)).replace("\n", "") - # format is usually: - # expected_slice = np.array([-0.5713, -0.3018, -0.9814, 0.04663, -0.879, 0.76, -1.734, 0.1044, 1.161]) - output_str = tensor_str.replace("tensor", f"{expected_tensor_name} = np.array") - test_file, test_class, test_fn = test_name.split("::") - test_fn = test_fn.split()[0] - with open(filename, "a") as f: - print("::".join([test_file, test_class, test_fn, output_str]), file=f) - - -def get_tests_dir(append_path=None): - """ - Args: - append_path: optional path to append to the tests dir path - Return: - The full path to the `tests` dir, so that the tests can be invoked from anywhere. Optionally `append_path` is - joined after the `tests` dir the former is provided. - """ - # this function caller's __file__ - caller__file__ = inspect.stack()[1][1] - tests_dir = os.path.abspath(os.path.dirname(caller__file__)) - - while not tests_dir.endswith("tests"): - tests_dir = os.path.dirname(tests_dir) - - if append_path: - return Path(tests_dir, append_path).as_posix() - else: - return tests_dir - - -# Taken from the following PR: -# https://github.com/huggingface/accelerate/pull/1964 -def str_to_bool(value) -> int: - """ - Converts a string representation of truth to `True` (1) or `False` (0). True values are `y`, `yes`, `t`, `true`, - `on`, and `1`; False value are `n`, `no`, `f`, `false`, `off`, and `0`; - """ - value = value.lower() - if value in ("y", "yes", "t", "true", "on", "1"): - return 1 - elif value in ("n", "no", "f", "false", "off", "0"): - return 0 - else: - raise ValueError(f"invalid truth value {value}") - - -def parse_flag_from_env(key, default=False): - try: - value = os.environ[key] - except KeyError: - # KEY isn't set, default to `default`. - _value = default - else: - # KEY is set, convert it to True or False. - try: - _value = str_to_bool(value) - except ValueError: - # More values are supported, but let's keep the message simple. - raise ValueError(f"If set, {key} must be yes or no.") - return _value - - -_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): - """Creates a random float32 tensor""" - if rng is None: - rng = global_rng - - total_dims = 1 - for dim in shape: - total_dims *= dim - - values = [] - for _ in range(total_dims): - values.append(rng.random() * scale) - - return torch.tensor(data=values, dtype=torch.float).view(shape).contiguous() - - -def slow(test_case): - """ - Decorator marking a test as slow. - - Slow tests are skipped by default. Set the RUN_SLOW environment variable to a truthy value to run them. - - """ - return unittest.skipUnless(_run_slow_tests, "test is slow")(test_case) - - -def nightly(test_case): - """ - Decorator marking a test that runs nightly in the diffusers CI. - - Slow tests are skipped by default. Set the RUN_NIGHTLY environment variable to a truthy value to run them. - - """ - return unittest.skipUnless(_run_nightly_tests, "test is nightly")(test_case) - - -def is_torch_compile(test_case): - """ - Decorator marking a test that runs compile tests in the diffusers CI. - - Compile tests are skipped by default. Set the RUN_COMPILE environment variable to a truthy value to run them. - - """ - return unittest.skipUnless(_run_compile_tests, "test is torch compile")(test_case) - - -def require_torch(test_case): - """ - Decorator marking a test that requires PyTorch. These tests are skipped when PyTorch isn't installed. - """ - return unittest.skipUnless(is_torch_available(), "test requires PyTorch")(test_case) - - -def require_torch_2(test_case): - """ - Decorator marking a test that requires PyTorch 2. These tests are skipped when it isn't installed. - """ - return unittest.skipUnless(is_torch_available() and is_torch_version(">=", "2.0.0"), "test requires PyTorch 2")( - test_case - ) - - -def require_torch_version_greater_equal(torch_version): - """Decorator marking a test that requires torch with a specific version or greater.""" - - def decorator(test_case): - correct_torch_version = is_torch_available() and is_torch_version(">=", torch_version) - return unittest.skipUnless( - correct_torch_version, f"test requires torch with the version greater than or equal to {torch_version}" - )(test_case) - - return decorator - - -def require_torch_version_greater(torch_version): - """Decorator marking a test that requires torch with a specific version greater.""" - - def decorator(test_case): - correct_torch_version = is_torch_available() and is_torch_version(">", torch_version) - return unittest.skipUnless( - correct_torch_version, f"test requires torch with the version greater than {torch_version}" - )(test_case) - - return decorator - - -def require_torch_gpu(test_case): - """Decorator marking a test that requires CUDA and PyTorch.""" - return unittest.skipUnless(is_torch_available() and torch_device == "cuda", "test requires PyTorch+CUDA")( - test_case - ) - - -def require_torch_cuda_compatibility(expected_compute_capability): - def decorator(test_case): - if torch.cuda.is_available(): - current_compute_capability = get_torch_cuda_device_capability() - return unittest.skipUnless( - float(current_compute_capability) == float(expected_compute_capability), - "Test not supported for this compute capability.", - ) - - return decorator - - -# These decorators are for accelerator-specific behaviours that are not GPU-specific -def require_torch_accelerator(test_case): - """Decorator marking a test that requires an accelerator backend and PyTorch.""" - return unittest.skipUnless(is_torch_available() and torch_device != "cpu", "test requires accelerator+PyTorch")( - test_case - ) - - -def require_torch_multi_gpu(test_case): - """ - Decorator marking a test that requires a multi-GPU setup (in PyTorch). These tests are skipped on a machine without - multiple GPUs. To run *only* the multi_gpu tests, assuming all test names contain multi_gpu: $ pytest -sv ./tests - -k "multi_gpu" - """ - if not is_torch_available(): - return unittest.skip("test requires PyTorch")(test_case) - - import torch - - return unittest.skipUnless(torch.cuda.device_count() > 1, "test requires multiple GPUs")(test_case) - - -def require_torch_multi_accelerator(test_case): - """ - Decorator marking a test that requires a multi-accelerator setup (in PyTorch). These tests are skipped on a machine - without multiple hardware accelerators. - """ - if not is_torch_available(): - return unittest.skip("test requires PyTorch")(test_case) - - import torch - - return unittest.skipUnless( - torch.cuda.device_count() > 1 or torch.xpu.device_count() > 1, "test requires multiple hardware accelerators" - )(test_case) - - -def require_torch_accelerator_with_fp16(test_case): - """Decorator marking a test that requires an accelerator with support for the FP16 data type.""" - return unittest.skipUnless(_is_torch_fp16_available(torch_device), "test requires accelerator with fp16 support")( - test_case - ) - - -def require_torch_accelerator_with_fp64(test_case): - """Decorator marking a test that requires an accelerator with support for the FP64 data type.""" - return unittest.skipUnless(_is_torch_fp64_available(torch_device), "test requires accelerator with fp64 support")( - test_case - ) - - -def require_big_gpu_with_torch_cuda(test_case): - """ - Decorator marking a test that requires a bigger GPU (24GB) for execution. Some example pipelines: Flux, SD3, Cog, - etc. - """ - if not is_torch_available(): - return unittest.skip("test requires PyTorch")(test_case) - - import torch - - if not torch.cuda.is_available(): - return unittest.skip("test requires PyTorch CUDA")(test_case) - - device_properties = torch.cuda.get_device_properties(0) - total_memory = device_properties.total_memory / (1024**3) - return unittest.skipUnless( - total_memory >= BIG_GPU_MEMORY, f"test requires a GPU with at least {BIG_GPU_MEMORY} GB memory" - )(test_case) - - -def require_big_accelerator(test_case): - """ - Decorator marking a test that requires a bigger hardware accelerator (24GB) for execution. Some example pipelines: - Flux, SD3, Cog, etc. - """ - if not is_torch_available(): - return unittest.skip("test requires PyTorch")(test_case) - - import torch - - if not (torch.cuda.is_available() or torch.xpu.is_available()): - return unittest.skip("test requires PyTorch CUDA")(test_case) - - if torch.xpu.is_available(): - device_properties = torch.xpu.get_device_properties(0) - else: - device_properties = torch.cuda.get_device_properties(0) - - total_memory = device_properties.total_memory / (1024**3) - return unittest.skipUnless( - total_memory >= BIG_GPU_MEMORY, - f"test requires a hardware accelerator with at least {BIG_GPU_MEMORY} GB memory", - )(test_case) - - -def require_torch_accelerator_with_training(test_case): - """Decorator marking a test that requires an accelerator with support for training.""" - return unittest.skipUnless( - is_torch_available() and backend_supports_training(torch_device), - "test requires accelerator with training support", - )(test_case) - - -def skip_mps(test_case): - """Decorator marking a test to skip if torch_device is 'mps'""" - return unittest.skipUnless(torch_device != "mps", "test requires non 'mps' device")(test_case) - - -def require_flax(test_case): - """ - Decorator marking a test that requires JAX & Flax. These tests are skipped when one / both are not installed - """ - return unittest.skipUnless(is_flax_available(), "test requires JAX & Flax")(test_case) - - -def require_compel(test_case): - """ - Decorator marking a test that requires compel: https://github.com/damian0815/compel. These tests are skipped when - the library is not installed. - """ - return unittest.skipUnless(is_compel_available(), "test requires compel")(test_case) - - -def require_onnxruntime(test_case): - """ - Decorator marking a test that requires onnxruntime. These tests are skipped when onnxruntime isn't installed. - """ - return unittest.skipUnless(is_onnx_available(), "test requires onnxruntime")(test_case) - - -def require_note_seq(test_case): - """ - Decorator marking a test that requires note_seq. These tests are skipped when note_seq isn't installed. - """ - return unittest.skipUnless(is_note_seq_available(), "test requires note_seq")(test_case) - - -def require_accelerator(test_case): - """ - Decorator marking a test that requires a hardware accelerator backend. These tests are skipped when there are no - hardware accelerator available. - """ - return unittest.skipUnless(torch_device != "cpu", "test requires a hardware accelerator")(test_case) - - -def require_torchsde(test_case): - """ - Decorator marking a test that requires torchsde. These tests are skipped when torchsde isn't installed. - """ - return unittest.skipUnless(is_torchsde_available(), "test requires torchsde")(test_case) - - -def require_peft_backend(test_case): - """ - Decorator marking a test that requires PEFT backend, this would require some specific versions of PEFT and - transformers. - """ - return unittest.skipUnless(USE_PEFT_BACKEND, "test requires PEFT backend")(test_case) - - -def require_timm(test_case): - """ - Decorator marking a test that requires timm. These tests are skipped when timm isn't installed. - """ - return unittest.skipUnless(is_timm_available(), "test requires timm")(test_case) - - -def require_bitsandbytes(test_case): - """ - Decorator marking a test that requires bitsandbytes. These tests are skipped when bitsandbytes isn't installed. - """ - return unittest.skipUnless(is_bitsandbytes_available(), "test requires bitsandbytes")(test_case) - - -def require_quanto(test_case): - """ - Decorator marking a test that requires quanto. These tests are skipped when quanto isn't installed. - """ - return unittest.skipUnless(is_optimum_quanto_available(), "test requires quanto")(test_case) - - -def require_accelerate(test_case): - """ - Decorator marking a test that requires accelerate. These tests are skipped when accelerate isn't installed. - """ - return unittest.skipUnless(is_accelerate_available(), "test requires accelerate")(test_case) - - -def require_peft_version_greater(peft_version): - """ - Decorator marking a test that requires PEFT backend with a specific version, this would require some specific - versions of PEFT and transformers. - """ - - def decorator(test_case): - correct_peft_version = is_peft_available() and version.parse( - version.parse(importlib.metadata.version("peft")).base_version - ) > version.parse(peft_version) - return unittest.skipUnless( - correct_peft_version, f"test requires PEFT backend with the version greater than {peft_version}" - )(test_case) - - return decorator - - -def require_transformers_version_greater(transformers_version): - """ - Decorator marking a test that requires transformers with a specific version, this would require some specific - versions of PEFT and transformers. - """ - - def decorator(test_case): - correct_transformers_version = is_transformers_available() and version.parse( - version.parse(importlib.metadata.version("transformers")).base_version - ) > version.parse(transformers_version) - return unittest.skipUnless( - correct_transformers_version, - f"test requires transformers with the version greater than {transformers_version}", - )(test_case) - - return decorator - - -def require_accelerate_version_greater(accelerate_version): - def decorator(test_case): - correct_accelerate_version = is_accelerate_available() and version.parse( - version.parse(importlib.metadata.version("accelerate")).base_version - ) > version.parse(accelerate_version) - return unittest.skipUnless( - correct_accelerate_version, f"Test requires accelerate with the version greater than {accelerate_version}." - )(test_case) - - return decorator - - -def require_bitsandbytes_version_greater(bnb_version): - def decorator(test_case): - correct_bnb_version = is_bitsandbytes_available() and version.parse( - version.parse(importlib.metadata.version("bitsandbytes")).base_version - ) > version.parse(bnb_version) - return unittest.skipUnless( - correct_bnb_version, f"Test requires bitsandbytes with the version greater than {bnb_version}." - )(test_case) - - return decorator - - -def require_hf_hub_version_greater(hf_hub_version): - def decorator(test_case): - correct_hf_hub_version = version.parse( - version.parse(importlib.metadata.version("huggingface_hub")).base_version - ) > version.parse(hf_hub_version) - return unittest.skipUnless( - correct_hf_hub_version, f"Test requires huggingface_hub with the version greater than {hf_hub_version}." - )(test_case) - - return decorator - - -def require_gguf_version_greater_or_equal(gguf_version): - def decorator(test_case): - correct_gguf_version = is_gguf_available() and version.parse( - version.parse(importlib.metadata.version("gguf")).base_version - ) >= version.parse(gguf_version) - return unittest.skipUnless( - correct_gguf_version, f"Test requires gguf with the version greater than {gguf_version}." - )(test_case) - - return decorator - - -def require_torchao_version_greater_or_equal(torchao_version): - def decorator(test_case): - correct_torchao_version = is_torchao_available() and version.parse( - version.parse(importlib.metadata.version("torchao")).base_version - ) >= version.parse(torchao_version) - return unittest.skipUnless( - correct_torchao_version, f"Test requires torchao with version greater than {torchao_version}." - )(test_case) - - return decorator - - -def deprecate_after_peft_backend(test_case): - """ - Decorator marking a test that will be skipped after PEFT backend - """ - return unittest.skipUnless(not USE_PEFT_BACKEND, "test skipped in favor of PEFT backend")(test_case) - - -def get_python_version(): - sys_info = sys.version_info - major, minor = sys_info.major, sys_info.minor - return major, minor - - -def load_numpy(arry: Union[str, np.ndarray], local_path: Optional[str] = None) -> np.ndarray: - if isinstance(arry, str): - if local_path is not None: - # local_path can be passed to correct images of tests - return Path(local_path, arry.split("/")[-5], arry.split("/")[-2], arry.split("/")[-1]).as_posix() - elif arry.startswith("http://") or arry.startswith("https://"): - response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT) - response.raise_for_status() - arry = np.load(BytesIO(response.content)) - elif os.path.isfile(arry): - arry = np.load(arry) - else: - raise ValueError( - f"Incorrect path or url, URLs must start with `http://` or `https://`, and {arry} is not a valid path" - ) - elif isinstance(arry, np.ndarray): - pass - else: - raise ValueError( - "Incorrect format used for numpy ndarray. Should be an url linking to an image, a local path, or a" - " ndarray." - ) - - return arry - - -def load_pt(url: str, map_location: Optional[str] = None, weights_only: Optional[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: - """ - Loads `image` to a PIL Image. - - Args: - image (`str` or `PIL.Image.Image`): - The image to convert to the PIL Image format. - Returns: - `PIL.Image.Image`: - A PIL Image. - """ - if isinstance(image, str): - if image.startswith("http://") or image.startswith("https://"): - image = PIL.Image.open(requests.get(image, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw) - elif os.path.isfile(image): - image = PIL.Image.open(image) - else: - raise ValueError( - f"Incorrect path or url, URLs must start with `http://` or `https://`, and {image} is not a valid path" - ) - elif isinstance(image, PIL.Image.Image): - image = image - else: - raise ValueError( - "Incorrect format used for image. Should be an url linking to an image, a local path, or a PIL image." - ) - image = PIL.ImageOps.exif_transpose(image) - image = image.convert("RGB") - return image - - -def preprocess_image(image: PIL.Image, batch_size: int): - w, h = image.size - w, h = (x - x % 8 for x in (w, h)) # resize to integer multiple of 8 - image = image.resize((w, h), resample=PIL.Image.LANCZOS) - image = np.array(image).astype(np.float32) / 255.0 - image = np.vstack([image[None].transpose(0, 3, 1, 2)] * batch_size) - image = torch.from_numpy(image) - return 2.0 * image - 1.0 - - -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 - - image[0].save( - output_gif_path, - save_all=True, - append_images=image[1:], - optimize=False, - duration=100, - loop=0, - ) - return output_gif_path - - -@contextmanager -def buffered_writer(raw_f): - f = io.BufferedWriter(raw_f) - yield f - f.flush() - - -def export_to_ply(mesh, output_ply_path: str = None): - """ - Write a PLY file for a mesh. - """ - if output_ply_path is None: - output_ply_path = tempfile.NamedTemporaryFile(suffix=".ply").name - - coords = mesh.verts.detach().cpu().numpy() - faces = mesh.faces.cpu().numpy() - rgb = np.stack([mesh.vertex_channels[x].detach().cpu().numpy() for x in "RGB"], axis=1) - - with buffered_writer(open(output_ply_path, "wb")) as f: - f.write(b"ply\n") - f.write(b"format binary_little_endian 1.0\n") - f.write(bytes(f"element vertex {len(coords)}\n", "ascii")) - f.write(b"property float x\n") - f.write(b"property float y\n") - f.write(b"property float z\n") - if rgb is not None: - f.write(b"property uchar red\n") - f.write(b"property uchar green\n") - f.write(b"property uchar blue\n") - if faces is not None: - f.write(bytes(f"element face {len(faces)}\n", "ascii")) - f.write(b"property list uchar int vertex_index\n") - f.write(b"end_header\n") - - if rgb is not None: - rgb = (rgb * 255.499).round().astype(int) - vertices = [ - (*coord, *rgb) - for coord, rgb in zip( - coords.tolist(), - rgb.tolist(), - ) - ] - format = struct.Struct("<3f3B") - for item in vertices: - f.write(format.pack(*item)) - else: - format = struct.Struct("<3f") - for vertex in coords.tolist(): - f.write(format.pack(*vertex)) - - if faces is not None: - format = struct.Struct(" str: - if is_opencv_available(): - import cv2 - else: - raise ImportError(BACKENDS_MAPPING["opencv"][1].format("export_to_video")) - if output_video_path is None: - output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name - - fourcc = cv2.VideoWriter_fourcc(*"mp4v") - h, w, c = video_frames[0].shape - video_writer = cv2.VideoWriter(output_video_path, fourcc, fps=8, frameSize=(w, h)) - for i in range(len(video_frames)): - img = cv2.cvtColor(video_frames[i], cv2.COLOR_RGB2BGR) - video_writer.write(img) - return output_video_path - - -def load_hf_numpy(path) -> np.ndarray: - base_url = "https://huggingface.co/datasets/fusing/diffusers-testing/resolve/main" - - if not path.startswith("http://") and not path.startswith("https://"): - path = os.path.join(base_url, urllib.parse.quote(path)) - - return load_numpy(path) - - -# --- pytest conf functions --- # - -# to avoid multiple invocation from tests/conftest.py and examples/conftest.py - make sure it's called only once -pytest_opt_registered = {} - - -def pytest_addoption_shared(parser): - """ - This function is to be called from `conftest.py` via `pytest_addoption` wrapper that has to be defined there. - - It allows loading both `conftest.py` files at once without causing a failure due to adding the same `pytest` - option. - - """ - option = "--make-reports" - if option not in pytest_opt_registered: - parser.addoption( - option, - action="store", - default=False, - help="generate report files. The value of this option is used as a prefix to report names", - ) - pytest_opt_registered[option] = 1 - - -def pytest_terminal_summary_main(tr, id): - """ - Generate multiple reports at the end of test suite run - each report goes into a dedicated file in the current - directory. The report files are prefixed with the test suite name. - - This function emulates --duration and -rA pytest arguments. - - This function is to be called from `conftest.py` via `pytest_terminal_summary` wrapper that has to be defined - there. - - Args: - - tr: `terminalreporter` passed from `conftest.py` - - id: unique id like `tests` or `examples` that will be incorporated into the final reports filenames - this is - needed as some jobs have multiple runs of pytest, so we can't have them overwrite each other. - - NB: this functions taps into a private _pytest API and while unlikely, it could break should - pytest do internal changes - also it calls default internal methods of terminalreporter which - can be hijacked by various `pytest-` plugins and interfere. - - """ - from _pytest.config import create_terminal_writer - - if not len(id): - id = "tests" - - config = tr.config - orig_writer = config.get_terminal_writer() - orig_tbstyle = config.option.tbstyle - orig_reportchars = tr.reportchars - - dir = "reports" - Path(dir).mkdir(parents=True, exist_ok=True) - report_files = { - k: f"{dir}/{id}_{k}.txt" - for k in [ - "durations", - "errors", - "failures_long", - "failures_short", - "failures_line", - "passes", - "stats", - "summary_short", - "warnings", - ] - } - - # custom durations report - # note: there is no need to call pytest --durations=XX to get this separate report - # adapted from https://github.com/pytest-dev/pytest/blob/897f151e/src/_pytest/runner.py#L66 - dlist = [] - for replist in tr.stats.values(): - for rep in replist: - if hasattr(rep, "duration"): - dlist.append(rep) - if dlist: - dlist.sort(key=lambda x: x.duration, reverse=True) - with open(report_files["durations"], "w") as f: - durations_min = 0.05 # sec - f.write("slowest durations\n") - for i, rep in enumerate(dlist): - if rep.duration < durations_min: - f.write(f"{len(dlist) - i} durations < {durations_min} secs were omitted") - break - f.write(f"{rep.duration:02.2f}s {rep.when:<8} {rep.nodeid}\n") - - def summary_failures_short(tr): - # expecting that the reports were --tb=long (default) so we chop them off here to the last frame - reports = tr.getreports("failed") - if not reports: - return - tr.write_sep("=", "FAILURES SHORT STACK") - for rep in reports: - msg = tr._getfailureheadline(rep) - tr.write_sep("_", msg, red=True, bold=True) - # chop off the optional leading extra frames, leaving only the last one - longrepr = re.sub(r".*_ _ _ (_ ){10,}_ _ ", "", rep.longreprtext, 0, re.M | re.S) - tr._tw.line(longrepr) - # note: not printing out any rep.sections to keep the report short - - # use ready-made report funcs, we are just hijacking the filehandle to log to a dedicated file each - # adapted from https://github.com/pytest-dev/pytest/blob/897f151e/src/_pytest/terminal.py#L814 - # note: some pytest plugins may interfere by hijacking the default `terminalreporter` (e.g. - # pytest-instafail does that) - - # report failures with line/short/long styles - config.option.tbstyle = "auto" # full tb - with open(report_files["failures_long"], "w") as f: - tr._tw = create_terminal_writer(config, f) - tr.summary_failures() - - # config.option.tbstyle = "short" # short tb - with open(report_files["failures_short"], "w") as f: - tr._tw = create_terminal_writer(config, f) - summary_failures_short(tr) - - config.option.tbstyle = "line" # one line per error - with open(report_files["failures_line"], "w") as f: - tr._tw = create_terminal_writer(config, f) - tr.summary_failures() - - with open(report_files["errors"], "w") as f: - tr._tw = create_terminal_writer(config, f) - tr.summary_errors() - - with open(report_files["warnings"], "w") as f: - tr._tw = create_terminal_writer(config, f) - tr.summary_warnings() # normal warnings - tr.summary_warnings() # final warnings - - tr.reportchars = "wPpsxXEf" # emulate -rA (used in summary_passes() and short_test_summary()) - with open(report_files["passes"], "w") as f: - tr._tw = create_terminal_writer(config, f) - tr.summary_passes() - - with open(report_files["summary_short"], "w") as f: - tr._tw = create_terminal_writer(config, f) - tr.short_test_summary() - - with open(report_files["stats"], "w") as f: - tr._tw = create_terminal_writer(config, f) - tr.summary_stats() - - # restore: - tr._tw = orig_writer - tr.reportchars = orig_reportchars - config.option.tbstyle = orig_tbstyle - - -# Copied 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): - """ - To decorate flaky tests. They will be retried on failures. - - Args: - max_attempts (`int`, *optional*, defaults to 5): - The maximum number of attempts to retry the flaky test. - wait_before_retry (`float`, *optional*): - If provided, will wait that number of seconds before retrying the test. - description (`str`, *optional*): - A string to describe the situation (what / where / why is flaky, link to GH issue/PR comments, errors, - etc.) - """ - - def decorator(test_func_ref): - @functools.wraps(test_func_ref) - def wrapper(*args, **kwargs): - retry_count = 1 - - while retry_count < max_attempts: - try: - return test_func_ref(*args, **kwargs) - - except Exception as err: - print(f"Test failed with {err} at try {retry_count}/{max_attempts}.", file=sys.stderr) - if wait_before_retry is not None: - time.sleep(wait_before_retry) - retry_count += 1 - - return test_func_ref(*args, **kwargs) - - return wrapper - - return decorator - - -# Taken from: https://github.com/huggingface/transformers/blob/3658488ff77ff8d45101293e749263acf437f4d5/src/transformers/testing_utils.py#L1787 -def run_test_in_subprocess(test_case, target_func, inputs=None, timeout=None): - """ - To run a test in a subprocess. In particular, this can avoid (GPU) memory issue. - - Args: - test_case (`unittest.TestCase`): - The test that will run `target_func`. - target_func (`Callable`): - The function implementing the actual testing logic. - inputs (`dict`, *optional*, defaults to `None`): - The inputs that will be passed to `target_func` through an (input) queue. - timeout (`int`, *optional*, defaults to `None`): - The timeout (in seconds) that will be passed to the input and output queues. If not specified, the env. - variable `PYTEST_TIMEOUT` will be checked. If still `None`, its value will be set to `600`. - """ - if timeout is None: - timeout = int(os.environ.get("PYTEST_TIMEOUT", 600)) - - start_methohd = "spawn" - ctx = multiprocessing.get_context(start_methohd) - - input_queue = ctx.Queue(1) - output_queue = ctx.JoinableQueue(1) - - # We can't send `unittest.TestCase` to the child, otherwise we get issues regarding pickle. - input_queue.put(inputs, timeout=timeout) - - process = ctx.Process(target=target_func, args=(input_queue, output_queue, timeout)) - process.start() - # Kill the child process if we can't get outputs from it in time: otherwise, the hanging subprocess prevents - # the test to exit properly. - try: - results = output_queue.get(timeout=timeout) - output_queue.task_done() - except Exception as e: - process.terminate() - test_case.fail(e) - process.join(timeout=timeout) - - if results["error"] is not None: - test_case.fail(f"{results['error']}") - - -class CaptureLogger: - """ - Args: - Context manager to capture `logging` streams - logger: 'logging` logger object - Returns: - The captured output is available via `self.out` - Example: - ```python - >>> from diffusers import logging - >>> from diffusers.testing_utils import CaptureLogger - - >>> msg = "Testing 1, 2, 3" - >>> logging.set_verbosity_info() - >>> logger = logging.get_logger("diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.py") - >>> with CaptureLogger(logger) as cl: - ... logger.info(msg) - >>> assert cl.out, msg + "\n" - ``` - """ - - def __init__(self, logger): - self.logger = logger - self.io = StringIO() - self.sh = logging.StreamHandler(self.io) - self.out = "" - - def __enter__(self): - self.logger.addHandler(self.sh) - return self - - def __exit__(self, *exc): - self.logger.removeHandler(self.sh) - self.out = self.io.getvalue() - - def __repr__(self): - return f"captured: {self.out}\n" - - -def enable_full_determinism(): - """ - Helper function for reproducible behavior during distributed training. See - - https://pytorch.org/docs/stable/notes/randomness.html for pytorch - """ - # Enable PyTorch deterministic mode. This potentially requires either the environment - # 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" - torch.use_deterministic_algorithms(True) - - # Enable CUDNN deterministic mode - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False - torch.backends.cuda.matmul.allow_tf32 = False - - -def disable_full_determinism(): - os.environ["CUDA_LAUNCH_BLOCKING"] = "0" - os.environ["CUBLAS_WORKSPACE_CONFIG"] = "" - torch.use_deterministic_algorithms(False) - - -# Utils for custom and alternative accelerator devices -def _is_torch_fp16_available(device): - if not is_torch_available(): - return False - - import torch - - device = torch.device(device) - - try: - x = torch.zeros((2, 2), dtype=torch.float16).to(device) - _ = torch.mul(x, x) - return True - - except Exception as e: - if device.type == "cuda": - raise ValueError( - f"You have passed a device of type 'cuda' which should work with 'fp16', but 'cuda' does not seem to be correctly installed on your machine: {e}" - ) - - return False - - -def _is_torch_fp64_available(device): - if not is_torch_available(): - return False - - import torch - - device = torch.device(device) - - try: - x = torch.zeros((2, 2), dtype=torch.float64).to(device) - _ = torch.mul(x, x) - return True - - except Exception as e: - if device.type == "cuda": - raise ValueError( - f"You have passed a device of type 'cuda' which should work with 'fp64', but 'cuda' does not seem to be correctly installed on your machine: {e}" - ) - - return False - - -# Guard these lookups for when Torch is not used - alternative accelerator support is for PyTorch -if is_torch_available(): - # Behaviour flags - BACKEND_SUPPORTS_TRAINING = {"cuda": True, "xpu": True, "cpu": True, "mps": False, "default": True} - - # Function definitions - BACKEND_EMPTY_CACHE = { - "cuda": torch.cuda.empty_cache, - "xpu": torch.xpu.empty_cache, - "cpu": None, - "mps": torch.mps.empty_cache, - "default": None, - } - BACKEND_DEVICE_COUNT = { - "cuda": torch.cuda.device_count, - "xpu": torch.xpu.device_count, - "cpu": lambda: 0, - "mps": lambda: 0, - "default": 0, - } - BACKEND_MANUAL_SEED = { - "cuda": torch.cuda.manual_seed, - "xpu": torch.xpu.manual_seed, - "cpu": torch.manual_seed, - "mps": torch.mps.manual_seed, - "default": torch.manual_seed, - } - BACKEND_RESET_PEAK_MEMORY_STATS = { - "cuda": torch.cuda.reset_peak_memory_stats, - "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None), - "cpu": None, - "mps": None, - "default": None, - } - BACKEND_RESET_MAX_MEMORY_ALLOCATED = { - "cuda": torch.cuda.reset_max_memory_allocated, - "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None), - "cpu": None, - "mps": None, - "default": None, - } - BACKEND_MAX_MEMORY_ALLOCATED = { - "cuda": torch.cuda.max_memory_allocated, - "xpu": getattr(torch.xpu, "max_memory_allocated", None), - "cpu": 0, - "mps": 0, - "default": 0, - } - BACKEND_SYNCHRONIZE = { - "cuda": torch.cuda.synchronize, - "xpu": getattr(torch.xpu, "synchronize", None), - "cpu": None, - "mps": None, - "default": None, - } - - -# 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): - if device not in dispatch_table: - return dispatch_table["default"](*args, **kwargs) - - fn = dispatch_table[device] - - # Some device agnostic functions return values. Need to guard against 'None' instead at - # user level - if not callable(fn): - return fn - - return fn(*args, **kwargs) - - -# These are callables which automatically dispatch the function specific to the accelerator -def backend_manual_seed(device: str, seed: int): - return _device_agnostic_dispatch(device, BACKEND_MANUAL_SEED, seed) - - -def backend_synchronize(device: str): - return _device_agnostic_dispatch(device, BACKEND_SYNCHRONIZE) - - -def backend_empty_cache(device: str): - return _device_agnostic_dispatch(device, BACKEND_EMPTY_CACHE) - - -def backend_device_count(device: str): - return _device_agnostic_dispatch(device, BACKEND_DEVICE_COUNT) - - -def backend_reset_peak_memory_stats(device: str): - return _device_agnostic_dispatch(device, BACKEND_RESET_PEAK_MEMORY_STATS) - - -def backend_reset_max_memory_allocated(device: str): - return _device_agnostic_dispatch(device, BACKEND_RESET_MAX_MEMORY_ALLOCATED) - - -def backend_max_memory_allocated(device: str): - return _device_agnostic_dispatch(device, BACKEND_MAX_MEMORY_ALLOCATED) - - -# These are callables which return boolean behaviour flags and can be used to specify some -# device agnostic alternative where the feature is unsupported. -def backend_supports_training(device: str): - if not is_torch_available(): - return False - - if device not in BACKEND_SUPPORTS_TRAINING: - device = "default" - - return BACKEND_SUPPORTS_TRAINING[device] - - -# 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): - try: - # Try to import the function directly - spec_fn = getattr(device_spec_module, attribute_name) - device_fn_dict[torch_device] = spec_fn - except AttributeError as e: - # If the function doesn't exist, and there is no default, throw an error - if "default" not in device_fn_dict: - raise AttributeError( - f"`{attribute_name}` not found in '{device_spec_path}' and no default fallback function found." - ) from e - - if "DIFFUSERS_TEST_DEVICE_SPEC" in os.environ: - device_spec_path = os.environ["DIFFUSERS_TEST_DEVICE_SPEC"] - if not Path(device_spec_path).is_file(): - raise ValueError(f"Specified path to device specification file is not found. Received {device_spec_path}") - - try: - import_name = device_spec_path[: device_spec_path.index(".py")] - except ValueError as e: - raise ValueError(f"Provided device spec file is not a Python file! Received {device_spec_path}") from e - - device_spec_module = importlib.import_module(import_name) - - try: - device_name = device_spec_module.DEVICE_NAME - except AttributeError: - raise AttributeError("Device spec file did not contain `DEVICE_NAME`") - - if "DIFFUSERS_TEST_DEVICE" in os.environ and torch_device != device_name: - msg = f"Mismatch between environment variable `DIFFUSERS_TEST_DEVICE` '{torch_device}' and device found in spec '{device_name}'\n" - msg += "Either unset `DIFFUSERS_TEST_DEVICE` or ensure it matches device spec name." - raise ValueError(msg) - - torch_device = device_name - - # Add one entry here for each `BACKEND_*` dictionary. - update_mapping_from_spec(BACKEND_MANUAL_SEED, "MANUAL_SEED_FN") - update_mapping_from_spec(BACKEND_EMPTY_CACHE, "EMPTY_CACHE_FN") - update_mapping_from_spec(BACKEND_DEVICE_COUNT, "DEVICE_COUNT_FN") - update_mapping_from_spec(BACKEND_SUPPORTS_TRAINING, "SUPPORTS_TRAINING") - update_mapping_from_spec(BACKEND_RESET_PEAK_MEMORY_STATS, "RESET_PEAK_MEMORY_STATS_FN") - update_mapping_from_spec(BACKEND_RESET_MAX_MEMORY_ALLOCATED, "RESET_MAX_MEMORY_ALLOCATED_FN") - update_mapping_from_spec(BACKEND_MAX_MEMORY_ALLOCATED, "MAX_MEMORY_ALLOCATED_FN") - - -# 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]] - - -@functools.lru_cache -def get_device_properties() -> DeviceProperties: - """ - Get environment device properties. - """ - if IS_CUDA_SYSTEM or IS_ROCM_SYSTEM: - import torch - - major, _ = torch.cuda.get_device_capability() - if IS_ROCM_SYSTEM: - return ("rocm", major) - else: - return ("cuda", major) - elif IS_XPU_SYSTEM: - import torch - - # To get more info of the architecture meaning and bit allocation, refer to https://github.com/intel/llvm/blob/sycl/sycl/include/sycl/ext/oneapi/experimental/device_architecture.def - arch = torch.xpu.get_device_capability()["architecture"] - gen_mask = 0x000000FF00000000 - gen = (arch & gen_mask) >> 32 - return ("xpu", gen) - else: - return (torch_device, None) - - -if TYPE_CHECKING: - DevicePropertiesUserDict = UserDict[DeviceProperties, Any] -else: - DevicePropertiesUserDict = UserDict - - -class Expectations(DevicePropertiesUserDict): - def get_expectation(self) -> Any: - """ - Find best matching expectation based on environment device properties. - """ - return self.find_expectation(get_device_properties()) - - @staticmethod - def is_default(key: DeviceProperties) -> bool: - return all(p is None for p in key) - - @staticmethod - def score(key: DeviceProperties, other: DeviceProperties) -> int: - """ - Returns score indicating how similar two instances of the `Properties` tuple are. Points are calculated using - bits, but documented as int. Rules are as follows: - * Matching `type` gives 8 points. - * Semi-matching `type`, for example cuda and rocm, gives 4 points. - * Matching `major` (compute capability major version) gives 2 points. - * Default expectation (if present) gives 1 points. - """ - (device_type, major) = key - (other_device_type, other_major) = other - - score = 0b0 - if device_type == other_device_type: - score |= 0b1000 - elif device_type in ["cuda", "rocm"] and other_device_type in ["cuda", "rocm"]: - score |= 0b100 - - if major == other_major and other_major is not None: - score |= 0b10 - - if Expectations.is_default(other): - score |= 0b1 - - return int(score) - - def find_expectation(self, key: DeviceProperties = (None, None)) -> Any: - """ - Find best matching expectation based on provided device properties. - """ - (result_key, result) = max(self.data.items(), key=lambda x: Expectations.score(key, x[0])) - - if Expectations.score(key, result_key) == 0: - raise ValueError(f"No matching expectation found for {key}") - - return result - - def __repr__(self): - return f"{self.data}" diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/torch_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/torch_utils.py deleted file mode 100644 index ffc1119727dfaf9b83beb4041427d6d4343649fd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/torch_utils.py +++ /dev/null @@ -1,181 +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. -""" -PyTorch utilities: Utilities related to PyTorch -""" - -from typing import List, Optional, Tuple, Union - -from . import logging -from .import_utils import is_torch_available, is_torch_npu_available, is_torch_version - - -if is_torch_available(): - import torch - from torch.fft import fftn, fftshift, ifftn, ifftshift - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -try: - from torch._dynamo import allow_in_graph as maybe_allow_in_graph -except (ImportError, ModuleNotFoundError): - - def maybe_allow_in_graph(cls): - return cls - - -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, -): - """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 - is always created on the CPU. - """ - # device on which tensor is created defaults to device - if isinstance(device, str): - device = torch.device(device) - rand_device = device - batch_size = shape[0] - - layout = layout or torch.strided - device = device or torch.device("cpu") - - if generator is not None: - gen_device_type = generator.device.type if not isinstance(generator, list) else generator[0].device.type - if gen_device_type != device.type and gen_device_type == "cpu": - rand_device = "cpu" - if device != "mps": - logger.info( - f"The passed generator was created on 'cpu' even though a tensor on {device} was expected." - f" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably" - f" slightly speed up this function by passing a generator that was created on the {device} device." - ) - elif gen_device_type != device.type and gen_device_type == "cuda": - raise ValueError(f"Cannot generate a {device} tensor from a generator of type {gen_device_type}.") - - # make sure generator list of length 1 is treated like a non-list - if isinstance(generator, list) and len(generator) == 1: - generator = generator[0] - - if isinstance(generator, list): - shape = (1,) + shape[1:] - latents = [ - torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype, layout=layout) - for i in range(batch_size) - ] - latents = torch.cat(latents, dim=0).to(device) - else: - latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype, layout=layout).to(device) - - return latents - - -def is_compiled_module(module) -> bool: - """Check whether the module was compiled with torch.compile()""" - if is_torch_version("<", "2.0.0") or not hasattr(torch, "_dynamo"): - return False - return isinstance(module, torch._dynamo.eval_frame.OptimizedModule) - - -def fourier_filter(x_in: "torch.Tensor", threshold: int, scale: int) -> "torch.Tensor": - """Fourier filter as introduced in FreeU (https://huggingface.co/papers/2309.11497). - - This version of the method comes from here: - https://github.com/huggingface/diffusers/pull/5164#issuecomment-1732638706 - """ - x = x_in - B, C, H, W = x.shape - - # Non-power of 2 images must be float32 - if (W & (W - 1)) != 0 or (H & (H - 1)) != 0: - x = x.to(dtype=torch.float32) - # fftn does not support bfloat16 - elif x.dtype == torch.bfloat16: - x = x.to(dtype=torch.float32) - - # FFT - x_freq = fftn(x, dim=(-2, -1)) - x_freq = fftshift(x_freq, dim=(-2, -1)) - - B, C, H, W = x_freq.shape - mask = torch.ones((B, C, H, W), device=x.device) - - crow, ccol = H // 2, W // 2 - mask[..., crow - threshold : crow + threshold, ccol - threshold : ccol + threshold] = scale - x_freq = x_freq * mask - - # IFFT - x_freq = ifftshift(x_freq, dim=(-2, -1)) - x_filtered = ifftn(x_freq, dim=(-2, -1)).real - - return x_filtered.to(dtype=x_in.dtype) - - -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: - //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. - hidden_states (`torch.Tensor`): Inputs to the underlying block. - res_hidden_states (`torch.Tensor`): Features from the skip block corresponding to the underlying block. - s1 (`float`): Scaling factor for stage 1 to attenuate the contributions of the skip features. - s2 (`float`): Scaling factor for stage 2 to attenuate the contributions of the skip features. - b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features. - b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features. - """ - if resolution_idx == 0: - num_half_channels = hidden_states.shape[1] // 2 - hidden_states[:, :num_half_channels] = hidden_states[:, :num_half_channels] * freeu_kwargs["b1"] - res_hidden_states = fourier_filter(res_hidden_states, threshold=1, scale=freeu_kwargs["s1"]) - if resolution_idx == 1: - num_half_channels = hidden_states.shape[1] // 2 - hidden_states[:, :num_half_channels] = hidden_states[:, :num_half_channels] * freeu_kwargs["b2"] - res_hidden_states = fourier_filter(res_hidden_states, threshold=1, scale=freeu_kwargs["s2"]) - - return hidden_states, res_hidden_states - - -def get_torch_cuda_device_capability(): - if torch.cuda.is_available(): - device = torch.device("cuda") - compute_capability = torch.cuda.get_device_capability(device) - compute_capability = f"{compute_capability[0]}.{compute_capability[1]}" - return float(compute_capability) - else: - return None - - -def get_device(): - if torch.cuda.is_available(): - return "cuda" - elif is_torch_npu_available(): - return "npu" - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - return "xpu" - else: - return "cpu" - - -def empty_device_cache(device_type: Optional[str] = None): - if device_type is None: - device_type = get_device() - device_mod = getattr(torch, device_type, torch.cuda) - device_mod.empty_cache() diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/typing_utils.py b/venv/lib/python3.10/site-packages/diffusers/utils/typing_utils.py deleted file mode 100644 index 2b5b1a4f5ab5510856bf453a7436f6deb9c61abe..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/typing_utils.py +++ /dev/null @@ -1,91 +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. -""" -Typing utilities: Utilities related to type checking and validation -""" - -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: - """ - 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. - """ - if not isinstance(class_or_tuple, tuple): - class_or_tuple = (class_or_tuple,) - - # Unpack unions - unpacked_class_or_tuple = [] - for t in class_or_tuple: - if get_origin(t) is Union: - unpacked_class_or_tuple.extend(get_args(t)) - else: - unpacked_class_or_tuple.append(t) - class_or_tuple = tuple(unpacked_class_or_tuple) - - if Any in class_or_tuple: - return True - - obj_type = type(obj) - # Classes with obj's type - class_or_tuple = {t for t in class_or_tuple if isinstance(obj, get_origin(t) or t)} - - # Singular types (e.g. int, ControlNet, ...) - # Untyped collections (e.g. List, but not List[int]) - elem_class_or_tuple = {get_args(t) for t in class_or_tuple} - if () in elem_class_or_tuple: - return True - # Typed lists or sets - elif obj_type in (list, set): - return any(all(_is_valid_type(x, t) for x in obj) for t in elem_class_or_tuple) - # Typed tuples - elif obj_type is tuple: - return any( - # Tuples with any length and single type (e.g. Tuple[int, ...]) - (len(t) == 2 and t[-1] is Ellipsis and all(_is_valid_type(x, t[0]) for x in obj)) - or - # Tuples with fixed length and any types (e.g. Tuple[int, str]) - (len(obj) == len(t) and all(_is_valid_type(x, tt) for x, tt in zip(obj, t))) - for t in elem_class_or_tuple - ) - # Typed dicts - elif obj_type is dict: - return any( - all(_is_valid_type(k, kt) and _is_valid_type(v, vt) for k, v in obj.items()) - for kt, vt in elem_class_or_tuple - ) - - else: - return False - - -def _get_detailed_type(obj: Any) -> Type: - """ - Gets a detailed type for an object, including nested types for collections. - """ - obj_type = type(obj) - - if obj_type in (list, set): - obj_origin_type = List if obj_type is list else Set - elems_type = Union[tuple({_get_detailed_type(x) for x in obj})] - return obj_origin_type[elems_type] - elif obj_type is tuple: - return Tuple[tuple(_get_detailed_type(x) for x in obj)] - elif obj_type is dict: - keys_type = Union[tuple({_get_detailed_type(k) for k in obj.keys()})] - values_type = Union[tuple({_get_detailed_type(k) for k in obj.values()})] - return Dict[keys_type, values_type] - else: - return obj_type diff --git a/venv/lib/python3.10/site-packages/diffusers/utils/versions.py b/venv/lib/python3.10/site-packages/diffusers/utils/versions.py deleted file mode 100644 index 945a3977ce62a9a55307862193e4be6f12c3c17f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/utils/versions.py +++ /dev/null @@ -1,117 +0,0 @@ -# Copyright 2020 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. -""" -Utilities for working with package versions -""" - -import importlib.metadata -import operator -import re -import sys -from typing import Optional - -from packaging import version - - -ops = { - "<": operator.lt, - "<=": operator.le, - "==": operator.eq, - "!=": operator.ne, - ">=": operator.ge, - ">": operator.gt, -} - - -def _compare_versions(op, got_ver, want_ver, requirement, pkg, hint): - if got_ver is None or want_ver is None: - raise ValueError( - f"Unable to compare versions for {requirement}: need={want_ver} found={got_ver}. This is unusual. Consider" - f" reinstalling {pkg}." - ) - if not ops[op](version.parse(got_ver), version.parse(want_ver)): - raise ImportError( - f"{requirement} is required for a normal functioning of this module, but found {pkg}=={got_ver}.{hint}" - ) - - -def require_version(requirement: str, hint: Optional[str] = None) -> None: - """ - Perform a runtime check of the dependency versions, using the exact same syntax used by pip. - - The installed module version comes from the *site-packages* dir via *importlib.metadata*. - - Args: - requirement (`str`): pip style definition, e.g., "tokenizers==0.9.4", "tqdm>=4.27", "numpy" - hint (`str`, *optional*): what suggestion to print in case of requirements not being met - - Example: - - ```python - require_version("pandas>1.1.2") - require_version("numpy>1.18.5", "this is important to have for whatever reason") - ```""" - - hint = f"\n{hint}" if hint is not None else "" - - # non-versioned check - if re.match(r"^[\w_\-\d]+$", requirement): - pkg, op, want_ver = requirement, None, None - else: - match = re.findall(r"^([^!=<>\s]+)([\s!=<>]{1,2}.+)", requirement) - if not match: - raise ValueError( - "requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23, but" - f" got {requirement}" - ) - pkg, want_full = match[0] - want_range = want_full.split(",") # there could be multiple requirements - wanted = {} - for w in want_range: - match = re.findall(r"^([\s!=<>]{1,2})(.+)", w) - if not match: - raise ValueError( - "requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23," - f" but got {requirement}" - ) - op, want_ver = match[0] - wanted[op] = want_ver - if op not in ops: - raise ValueError(f"{requirement}: need one of {list(ops.keys())}, but got {op}") - - # special case - if pkg == "python": - got_ver = ".".join([str(x) for x in sys.version_info[:3]]) - for op, want_ver in wanted.items(): - _compare_versions(op, got_ver, want_ver, requirement, pkg, hint) - return - - # check if any version is installed - try: - got_ver = importlib.metadata.version(pkg) - except importlib.metadata.PackageNotFoundError: - raise importlib.metadata.PackageNotFoundError( - f"The '{requirement}' distribution was not found and is required by this application. {hint}" - ) - - # check that the right version is installed if version number or a range was provided - if want_ver is not None: - for op, want_ver in wanted.items(): - _compare_versions(op, got_ver, want_ver, requirement, pkg, hint) - - -def require_version_core(requirement): - """require_version wrapper which emits a core-specific hint on failure""" - hint = "Try: pip install transformers -U or pip install -e '.[dev]' if you're working with git main" - return require_version(requirement, hint) diff --git a/venv/lib/python3.10/site-packages/diffusers/video_processor.py b/venv/lib/python3.10/site-packages/diffusers/video_processor.py deleted file mode 100644 index 59b59b47d2c7d194874ab56753ca3efdbff90f5e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/diffusers/video_processor.py +++ /dev/null @@ -1,113 +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 warnings -from typing import List, Optional, Union - -import numpy as np -import PIL -import torch - -from .image_processor import VaeImageProcessor, is_valid_image, is_valid_image_imagelist - - -class VideoProcessor(VaeImageProcessor): - r"""Simple video processor.""" - - def preprocess_video(self, video, height: Optional[int] = None, width: Optional[int] = None) -> torch.Tensor: - r""" - Preprocesses input video(s). - - Args: - 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. - * 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, - width)`). - * 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, - width)`. - height (`int`, *optional*, defaults to `None`): - The height in preprocessed frames of the video. If `None`, will use the `get_default_height_width()` to - get default height. - 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. - """ - if isinstance(video, list) and isinstance(video[0], np.ndarray) and video[0].ndim == 5: - warnings.warn( - "Passing `video` as a list of 5d np.ndarray is deprecated." - "Please concatenate the list along the batch dimension and pass it as a single 5d np.ndarray", - FutureWarning, - ) - video = np.concatenate(video, axis=0) - if isinstance(video, list) and isinstance(video[0], torch.Tensor) and video[0].ndim == 5: - warnings.warn( - "Passing `video` as a list of 5d torch.Tensor is deprecated." - "Please concatenate the list along the batch dimension and pass it as a single 5d torch.Tensor", - FutureWarning, - ) - video = torch.cat(video, axis=0) - - # ensure the input is a list of videos: - # - if it is a batch of videos (5d torch.Tensor or np.ndarray), it is converted to a list of videos (a list of 4d torch.Tensor or np.ndarray) - # - if it is a single video, it is converted to a list of one video. - if isinstance(video, (np.ndarray, torch.Tensor)) and video.ndim == 5: - video = list(video) - elif isinstance(video, list) and is_valid_image(video[0]) or is_valid_image_imagelist(video): - video = [video] - elif isinstance(video, list) and is_valid_image_imagelist(video[0]): - video = video - else: - raise ValueError( - "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) - - # move the number of channels before the number of frames. - video = video.permute(0, 2, 1, 3, 4) - - return video - - def postprocess_video( - self, video: torch.Tensor, output_type: str = "np" - ) -> Union[np.ndarray, torch.Tensor, List[PIL.Image.Image]]: - r""" - Converts a video tensor to a list of frames for export. - - Args: - video (`torch.Tensor`): The video as a tensor. - output_type (`str`, defaults to `"np"`): Output type of the postprocessed `video` tensor. - """ - batch_size = video.shape[0] - 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) - outputs.append(batch_output) - - if output_type == "np": - outputs = np.stack(outputs) - elif output_type == "pt": - outputs = torch.stack(outputs) - elif not output_type == "pil": - raise ValueError(f"{output_type} does not exist. Please choose one of ['np', 'pt', 'pil']") - - return outputs diff --git a/venv/lib/python3.10/site-packages/distutils-precedence.pth b/venv/lib/python3.10/site-packages/distutils-precedence.pth deleted file mode 100644 index 10c404f6ad452c148c46a39e11ddd4bc58530d16..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/distutils-precedence.pth +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7ea7ffef3fe2a117ee12c68ed6553617f0d7fd2f0590257c25c484959a3b7373 -size 152 diff --git a/venv/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so b/venv/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so deleted file mode 100644 index 9e77f2b0856cbe38b976836febec85f523cb32c7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2bcd94f0849c7977a20306a5525ac978912c8ce65575037eeb56c6a3b6fbf9a4 -size 1438299200 diff --git a/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/LICENSE b/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/LICENSE deleted file mode 100644 index 76f67efdc6470081b512a8db5bf2b1d4962d9c3c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2023 Georgi Gerganov - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/METADATA b/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/METADATA deleted file mode 100644 index 9e4fb97c351b77c21be3ba04cdfe8c6c6b420ee4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/METADATA +++ /dev/null @@ -1,126 +0,0 @@ -Metadata-Version: 2.3 -Name: gguf -Version: 0.17.1 -Summary: Read and write ML models in GGUF for GGML -Keywords: ggml,gguf,llama.cpp -Author: GGML -Author-email: ggml@ggml.ai -Requires-Python: >=3.8 -Classifier: License :: OSI Approved :: MIT License -Classifier: Operating System :: OS Independent -Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3.8 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Programming Language :: Python :: 3.10 -Classifier: Programming Language :: Python :: 3.11 -Classifier: Programming Language :: Python :: 3.12 -Classifier: Programming Language :: Python :: 3.13 -Provides-Extra: gui -Requires-Dist: PySide6 (>=6.9,<7.0) ; (python_version >= "3.9" and python_version < "3.14") and (extra == "gui") -Requires-Dist: numpy (>=1.17) -Requires-Dist: pyyaml (>=5.1) -Requires-Dist: tqdm (>=4.27) -Project-URL: Homepage, https://ggml.ai -Project-URL: Repository, https://github.com/ggml-org/llama.cpp -Description-Content-Type: text/markdown - -## gguf - -This is a Python package for writing binary files in the [GGUF](https://github.com/ggml-org/ggml/pull/302) -(GGML Universal File) format. - -See [convert_hf_to_gguf.py](https://github.com/ggml-org/llama.cpp/blob/master/convert_hf_to_gguf.py) -as an example for its usage. - -## Installation -```sh -pip install gguf -``` - -Optionally, you can install gguf with the extra 'gui' to enable the visual GGUF editor. -```sh -pip install gguf[gui] -``` - -## API Examples/Simple Tools - -[examples/writer.py](https://github.com/ggml-org/llama.cpp/blob/master/gguf-py/examples/writer.py) — Generates `example.gguf` in the current directory to demonstrate generating a GGUF file. Note that this file cannot be used as a model. - -[examples/reader.py](https://github.com/ggml-org/llama.cpp/blob/master/gguf-py/examples/reader.py) — Extracts and displays key-value pairs and tensor details from a GGUF file in a readable format. - -[gguf/scripts/gguf_dump.py](https://github.com/ggml-org/llama.cpp/blob/master/gguf-py/gguf/scripts/gguf_dump.py) — Dumps a GGUF file's metadata to the console. - -[gguf/scripts/gguf_set_metadata.py](https://github.com/ggml-org/llama.cpp/blob/master/gguf-py/gguf/scripts/gguf_set_metadata.py) — Allows changing simple metadata values in a GGUF file by key. - -[gguf/scripts/gguf_convert_endian.py](https://github.com/ggml-org/llama.cpp/blob/master/gguf-py/gguf/scripts/gguf_convert_endian.py) — Allows converting the endianness of GGUF files. - -[gguf/scripts/gguf_new_metadata.py](https://github.com/ggml-org/llama.cpp/blob/master/gguf-py/gguf/scripts/gguf_new_metadata.py) — Copies a GGUF file with added/modified/removed metadata values. - -[gguf/scripts/gguf_editor_gui.py](https://github.com/ggml-org/llama.cpp/blob/master/gguf-py/gguf/scripts/gguf_editor_gui.py) — Allows for viewing, editing, adding, or removing metadata values within a GGUF file as well as viewing its tensors with a Qt interface. - -## Development -Maintainers who participate in development of this package are advised to install it in editable mode: - -```sh -cd /path/to/llama.cpp/gguf-py - -pip install --editable . -``` - -**Note**: This may require to upgrade your Pip installation, with a message saying that editable installation currently requires `setup.py`. -In this case, upgrade Pip to the latest: - -```sh -pip install --upgrade pip -``` - -## Automatic publishing with CI - -There's a GitHub workflow to make a release automatically upon creation of tags in a specified format. - -1. Bump the version in `pyproject.toml`. -2. Create a tag named `gguf-vx.x.x` where `x.x.x` is the semantic version number. - -```sh -git tag -a gguf-v1.0.0 -m "Version 1.0 release" -``` - -3. Push the tags. - -```sh -git push origin --tags -``` - -## Manual publishing -If you want to publish the package manually for any reason, you need to have `twine` and `build` installed: - -```sh -pip install build twine -``` - -Then, follow these steps to release a new version: - -1. Bump the version in `pyproject.toml`. -2. Build the package: - -```sh -python -m build -``` - -3. Upload the generated distribution archives: - -```sh -python -m twine upload dist/* -``` - -## Run Unit Tests - -From root of this repository you can run this command to run all the unit tests - -```bash -python -m unittest discover ./gguf-py -v -``` - -## TODO -- [ ] Include conversion scripts as command line entry points in this package. - diff --git a/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/RECORD b/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/RECORD deleted file mode 100644 index b83178ee6ac492fe40229230f0d0b4a2054ee83b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/RECORD +++ /dev/null @@ -1,47 +0,0 @@ -../../../bin/gguf-convert-endian,sha256=dnXJg54diixDuUXImRlYLQCtDYdWr0aIMFVKOUEqHXA,249 -../../../bin/gguf-dump,sha256=dvhrFZ73eFsZ9qcrzDpNsCORWGARak19D3OPqHkiuTM,239 -../../../bin/gguf-editor-gui,sha256=39oEXapsbJzBsHp6UHzA7DXK6h5h6g2B08k8GI5i6xs,245 -../../../bin/gguf-new-metadata,sha256=-pNoLbvw09fEAZwAKKbeHecty6xCsfeE_22U8JKTrX8,247 -../../../bin/gguf-set-metadata,sha256=ooevAx7qxsR5j2cq3E4lXGLBjcV2niEdSl4k3UIhDag,247 -gguf-0.17.1.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -gguf-0.17.1.dist-info/LICENSE,sha256=73jH5mWeNMeYGU8NNE6AfHIt5wy8oTWe9UdyZh4Ryjg,1072 -gguf-0.17.1.dist-info/METADATA,sha256=pa8_ce5ufoWKHq1zzh97He53EZjFtRFONFDWtl7mkAI,4348 -gguf-0.17.1.dist-info/RECORD,, -gguf-0.17.1.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -gguf-0.17.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88 -gguf-0.17.1.dist-info/entry_points.txt,sha256=TozYSFmMxpOaKE3brn9nWp-QkgM_sZ3a_uFDOXUCYig,273 -gguf/__init__.py,sha256=PM_AEEzX6ojGAodDt78_LIm19HRCXeA6IXpgcjINfC8,219 -gguf/__pycache__/__init__.cpython-310.pyc,, -gguf/__pycache__/constants.cpython-310.pyc,, -gguf/__pycache__/gguf.cpython-310.pyc,, -gguf/__pycache__/gguf_reader.cpython-310.pyc,, -gguf/__pycache__/gguf_writer.cpython-310.pyc,, -gguf/__pycache__/lazy.cpython-310.pyc,, -gguf/__pycache__/metadata.cpython-310.pyc,, -gguf/__pycache__/quants.cpython-310.pyc,, -gguf/__pycache__/tensor_mapping.cpython-310.pyc,, -gguf/__pycache__/utility.cpython-310.pyc,, -gguf/__pycache__/vocab.cpython-310.pyc,, -gguf/constants.py,sha256=-92nfNa8XjXbHsMV9Hf0iGHwcmoXzH2JJ73rAv73NPE,89508 -gguf/gguf.py,sha256=8MDu7a0JEXhLUv_tjhYqDrWubVNc41cFvBYZbkZZenI,478 -gguf/gguf_reader.py,sha256=6uI4vaLeRC2MJV-uUjgsRoUZ-Rdszi8mE_bEPwnk6QE,14828 -gguf/gguf_writer.py,sha256=UJXqFa3T2EOLfPZ6ej3vuyV7OBfKXXaszvHKeHIWmSI,43935 -gguf/lazy.py,sha256=Axy_plbHoC34nMeK_ORbbxEyEFm-Fz8BxFIeqeipSYA,9214 -gguf/metadata.py,sha256=v0kxoYfCk_yXZSnzlSO9AhQ6z_HOUpzs0c6ZV4BRSbA,33301 -gguf/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -gguf/quants.py,sha256=BLtCCqhHBtabaYaAp7EckyYOR5idmoKWhu_Hy3glotk,60771 -gguf/scripts/__pycache__/gguf_convert_endian.cpython-310.pyc,, -gguf/scripts/__pycache__/gguf_dump.cpython-310.pyc,, -gguf/scripts/__pycache__/gguf_editor_gui.cpython-310.pyc,, -gguf/scripts/__pycache__/gguf_hash.cpython-310.pyc,, -gguf/scripts/__pycache__/gguf_new_metadata.cpython-310.pyc,, -gguf/scripts/__pycache__/gguf_set_metadata.cpython-310.pyc,, -gguf/scripts/gguf_convert_endian.py,sha256=yzl_MAQ3jyn_9MmOWV1CksHqlohd7DmrG7REwas0rlo,7365 -gguf/scripts/gguf_dump.py,sha256=zDgZSSQLyO3S4YJsSUzdebDfwmdqQPN7_VtFZ5BkGAk,21785 -gguf/scripts/gguf_editor_gui.py,sha256=frdErSIB90N-sAvqUpbLfdDsaUGMMOWQ-0iumwzzm_M,64398 -gguf/scripts/gguf_hash.py,sha256=nyd8kzjRKnOFek5UaD19pNXeAVMXUfFEASZ8konkGX8,3725 -gguf/scripts/gguf_new_metadata.py,sha256=U_v5FgbH292x7bsi2dG4rbQcWc14nmAtZEWdLnbkIZs,9767 -gguf/scripts/gguf_set_metadata.py,sha256=yGEqcQlCimd-pVl23V7u1giJNN3vfvASRqW8em5YWzs,4145 -gguf/tensor_mapping.py,sha256=6pUDYgly0-yErLV8HSZFSnqyFKAp7SieditMEHMsYGI,55999 -gguf/utility.py,sha256=80rZ3MdGZ6bX0_yFvLoPTTOlxga8THcihF0q38y5h6M,10808 -gguf/vocab.py,sha256=Gr43idKqc4yBHQh20HfeW9oegYWLpM-L0ZKCVxPj4Lc,20788 diff --git a/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/REQUESTED deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/WHEEL b/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/WHEEL deleted file mode 100644 index 9ed4d8facdffc976b12b8144fcda70b971ffea8f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/WHEEL +++ /dev/null @@ -1,4 +0,0 @@ -Wheel-Version: 1.0 -Generator: poetry-core 2.1.3 -Root-Is-Purelib: true -Tag: py3-none-any diff --git a/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/entry_points.txt b/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/entry_points.txt deleted file mode 100644 index 5502bee07257f9f8e862720f7253205ff217f4d3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf-0.17.1.dist-info/entry_points.txt +++ /dev/null @@ -1,7 +0,0 @@ -[console_scripts] -gguf-convert-endian=gguf.scripts.gguf_convert_endian:main -gguf-dump=gguf.scripts.gguf_dump:main -gguf-editor-gui=gguf.scripts.gguf_editor_gui:main -gguf-new-metadata=gguf.scripts.gguf_new_metadata:main -gguf-set-metadata=gguf.scripts.gguf_set_metadata:main - diff --git a/venv/lib/python3.10/site-packages/gguf/__init__.py b/venv/lib/python3.10/site-packages/gguf/__init__.py deleted file mode 100644 index 243defc4c1ca42d3713017d8902592f54ac849cd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -from .constants import * -from .lazy import * -from .gguf_reader import * -from .gguf_writer import * -from .quants import * -from .tensor_mapping import * -from .vocab import * -from .utility import * -from .metadata import * diff --git a/venv/lib/python3.10/site-packages/gguf/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 9c1eb627fd02faa46d38a69fb01ea9c1684e768c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/__pycache__/constants.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/__pycache__/constants.cpython-310.pyc deleted file mode 100644 index eddedc5206d63f64c37fbe86ed2f28a220b7ce33..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/__pycache__/constants.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/__pycache__/gguf.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/__pycache__/gguf.cpython-310.pyc deleted file mode 100644 index da2fb023e0613dff7d563a05dbbd913154244922..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/__pycache__/gguf.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/__pycache__/gguf_reader.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/__pycache__/gguf_reader.cpython-310.pyc deleted file mode 100644 index e9d3e53f196bec3f88d2a5eb411da12adc4b30c7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/__pycache__/gguf_reader.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/__pycache__/gguf_writer.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/__pycache__/gguf_writer.cpython-310.pyc deleted file mode 100644 index b029077befaad30188833cc4595b54cb3cedf70f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/__pycache__/gguf_writer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/__pycache__/lazy.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/__pycache__/lazy.cpython-310.pyc deleted file mode 100644 index b82d53ab649dac4fb61a5fb865d56c0b6fbc3f8f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/__pycache__/lazy.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/__pycache__/metadata.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/__pycache__/metadata.cpython-310.pyc deleted file mode 100644 index e8742f7e3fc44d44816514f59f778d69a9805329..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/__pycache__/metadata.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/__pycache__/quants.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/__pycache__/quants.cpython-310.pyc deleted file mode 100644 index fcee09ae4d3aa53c7c9382aa718dbe88cc6ff323..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/__pycache__/quants.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/__pycache__/tensor_mapping.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/__pycache__/tensor_mapping.cpython-310.pyc deleted file mode 100644 index ed2be677a6100f02c03fafb25207bf6a58d553bb..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/__pycache__/tensor_mapping.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/__pycache__/utility.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/__pycache__/utility.cpython-310.pyc deleted file mode 100644 index 040620346acbdbb1650b5d61e015be5e5fc94166..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/__pycache__/utility.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/__pycache__/vocab.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/__pycache__/vocab.cpython-310.pyc deleted file mode 100644 index 217268c022806a809a6698c3d4a5af4915924fc7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/__pycache__/vocab.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/constants.py b/venv/lib/python3.10/site-packages/gguf/constants.py deleted file mode 100644 index 834a1d5e1a97ed98f01196e0a0542afa95739b8c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/constants.py +++ /dev/null @@ -1,2438 +0,0 @@ -from __future__ import annotations - -from enum import Enum, IntEnum, auto -from typing import Any - -# -# constants -# - -GGUF_MAGIC = 0x46554747 # "GGUF" -GGUF_VERSION = 3 -GGUF_DEFAULT_ALIGNMENT = 32 -GGML_QUANT_VERSION = 2 # GGML_QNT_VERSION from ggml.h - -# -# metadata keys -# - - -class Keys: - class General: - TYPE = "general.type" - ARCHITECTURE = "general.architecture" - QUANTIZATION_VERSION = "general.quantization_version" - ALIGNMENT = "general.alignment" - FILE_TYPE = "general.file_type" - - # Authorship Metadata - NAME = "general.name" - AUTHOR = "general.author" - VERSION = "general.version" - ORGANIZATION = "general.organization" - - FINETUNE = "general.finetune" - BASENAME = "general.basename" - - DESCRIPTION = "general.description" - QUANTIZED_BY = "general.quantized_by" - - SIZE_LABEL = "general.size_label" - - # Licensing details - LICENSE = "general.license" - LICENSE_NAME = "general.license.name" - LICENSE_LINK = "general.license.link" - - # Typically represents the converted GGUF repo (Unless native) - URL = "general.url" # Model Website/Paper - DOI = "general.doi" - UUID = "general.uuid" - REPO_URL = "general.repo_url" # Model Source Repository (git/svn/etc...) - - # Model Source during conversion - SOURCE_URL = "general.source.url" # Model Website/Paper - SOURCE_DOI = "general.source.doi" - SOURCE_UUID = "general.source.uuid" - SOURCE_REPO_URL = "general.source.repo_url" # Model Source Repository (git/svn/etc...) - - # Base Model Source. There can be more than one source if it's a merged - # model like with 'Mistral-7B-Merge-14-v0.1'. This will assist in - # tracing linage of models as it is finetuned or merged over time. - BASE_MODEL_COUNT = "general.base_model.count" - BASE_MODEL_NAME = "general.base_model.{id}.name" - BASE_MODEL_AUTHOR = "general.base_model.{id}.author" - BASE_MODEL_VERSION = "general.base_model.{id}.version" - BASE_MODEL_ORGANIZATION = "general.base_model.{id}.organization" - BASE_MODEL_DESCRIPTION = "general.base_model.{id}.description" - BASE_MODEL_URL = "general.base_model.{id}.url" # Model Website/Paper - BASE_MODEL_DOI = "general.base_model.{id}.doi" - BASE_MODEL_UUID = "general.base_model.{id}.uuid" - BASE_MODEL_REPO_URL = "general.base_model.{id}.repo_url" # Model Source Repository (git/svn/etc...) - - # Dataset Source - DATASET_COUNT = "general.dataset.count" - DATASET_NAME = "general.dataset.{id}.name" - DATASET_AUTHOR = "general.dataset.{id}.author" - DATASET_VERSION = "general.dataset.{id}.version" - DATASET_ORGANIZATION = "general.dataset.{id}.organization" - DATASET_DESCRIPTION = "general.dataset.{id}.description" - DATASET_URL = "general.dataset.{id}.url" # Model Website/Paper - DATASET_DOI = "general.dataset.{id}.doi" - DATASET_UUID = "general.dataset.{id}.uuid" - DATASET_REPO_URL = "general.dataset.{id}.repo_url" # Model Source Repository (git/svn/etc...) - - # Array based KV stores - TAGS = "general.tags" - LANGUAGES = "general.languages" - - class LLM: - VOCAB_SIZE = "{arch}.vocab_size" - CONTEXT_LENGTH = "{arch}.context_length" - EMBEDDING_LENGTH = "{arch}.embedding_length" - FEATURES_LENGTH = "{arch}.features_length" - BLOCK_COUNT = "{arch}.block_count" - LEADING_DENSE_BLOCK_COUNT = "{arch}.leading_dense_block_count" - FEED_FORWARD_LENGTH = "{arch}.feed_forward_length" - EXPERT_FEED_FORWARD_LENGTH = "{arch}.expert_feed_forward_length" - EXPERT_SHARED_FEED_FORWARD_LENGTH = "{arch}.expert_shared_feed_forward_length" - USE_PARALLEL_RESIDUAL = "{arch}.use_parallel_residual" - TENSOR_DATA_LAYOUT = "{arch}.tensor_data_layout" - EXPERT_COUNT = "{arch}.expert_count" - EXPERT_USED_COUNT = "{arch}.expert_used_count" - EXPERT_SHARED_COUNT = "{arch}.expert_shared_count" - EXPERT_WEIGHTS_SCALE = "{arch}.expert_weights_scale" - EXPERT_WEIGHTS_NORM = "{arch}.expert_weights_norm" - EXPERT_GATING_FUNC = "{arch}.expert_gating_func" - MOE_EVERY_N_LAYERS = "{arch}.moe_every_n_layers" - POOLING_TYPE = "{arch}.pooling_type" - LOGIT_SCALE = "{arch}.logit_scale" - DECODER_START_TOKEN_ID = "{arch}.decoder_start_token_id" - ATTN_LOGIT_SOFTCAPPING = "{arch}.attn_logit_softcapping" - FINAL_LOGIT_SOFTCAPPING = "{arch}.final_logit_softcapping" - SWIN_NORM = "{arch}.swin_norm" - RESCALE_EVERY_N_LAYERS = "{arch}.rescale_every_n_layers" - TIME_MIX_EXTRA_DIM = "{arch}.time_mix_extra_dim" - TIME_DECAY_EXTRA_DIM = "{arch}.time_decay_extra_dim" - RESIDUAL_SCALE = "{arch}.residual_scale" - EMBEDDING_SCALE = "{arch}.embedding_scale" - TOKEN_SHIFT_COUNT = "{arch}.token_shift_count" - INTERLEAVE_MOE_LAYER_STEP = "{arch}.interleave_moe_layer_step" - - class Attention: - HEAD_COUNT = "{arch}.attention.head_count" - HEAD_COUNT_KV = "{arch}.attention.head_count_kv" - MAX_ALIBI_BIAS = "{arch}.attention.max_alibi_bias" - CLAMP_KQV = "{arch}.attention.clamp_kqv" - KEY_LENGTH = "{arch}.attention.key_length" - VALUE_LENGTH = "{arch}.attention.value_length" - LAYERNORM_EPS = "{arch}.attention.layer_norm_epsilon" - LAYERNORM_RMS_EPS = "{arch}.attention.layer_norm_rms_epsilon" - GROUPNORM_EPS = "{arch}.attention.group_norm_epsilon" - GROUPNORM_GROUPS = "{arch}.attention.group_norm_groups" - CAUSAL = "{arch}.attention.causal" - Q_LORA_RANK = "{arch}.attention.q_lora_rank" - KV_LORA_RANK = "{arch}.attention.kv_lora_rank" - DECAY_LORA_RANK = "{arch}.attention.decay_lora_rank" - ICLR_LORA_RANK = "{arch}.attention.iclr_lora_rank" - VALUE_RESIDUAL_MIX_LORA_RANK = "{arch}.attention.value_residual_mix_lora_rank" - GATE_LORA_RANK = "{arch}.attention.gate_lora_rank" - REL_BUCKETS_COUNT = "{arch}.attention.relative_buckets_count" - SLIDING_WINDOW = "{arch}.attention.sliding_window" - SCALE = "{arch}.attention.scale" - KEY_LENGTH_MLA = "{arch}.attention.key_length_mla" - VALUE_LENGTH_MLA = "{arch}.attention.value_length_mla" - - class Rope: - DIMENSION_COUNT = "{arch}.rope.dimension_count" - DIMENSION_SECTIONS = "{arch}.rope.dimension_sections" - FREQ_BASE = "{arch}.rope.freq_base" - SCALING_TYPE = "{arch}.rope.scaling.type" - SCALING_FACTOR = "{arch}.rope.scaling.factor" - SCALING_ATTN_FACTOR = "{arch}.rope.scaling.attn_factor" - SCALING_ORIG_CTX_LEN = "{arch}.rope.scaling.original_context_length" - SCALING_FINETUNED = "{arch}.rope.scaling.finetuned" - SCALING_YARN_LOG_MUL = "{arch}.rope.scaling.yarn_log_multiplier" - - class Split: - LLM_KV_SPLIT_NO = "split.no" - LLM_KV_SPLIT_COUNT = "split.count" - LLM_KV_SPLIT_TENSORS_COUNT = "split.tensors.count" - - class SSM: - CONV_KERNEL = "{arch}.ssm.conv_kernel" - INNER_SIZE = "{arch}.ssm.inner_size" - STATE_SIZE = "{arch}.ssm.state_size" - TIME_STEP_RANK = "{arch}.ssm.time_step_rank" - DT_B_C_RMS = "{arch}.ssm.dt_b_c_rms" - - class WKV: - HEAD_SIZE = "{arch}.wkv.head_size" - - class PosNet: - EMBEDDING_LENGTH = "{arch}.posnet.embedding_length" - BLOCK_COUNT = "{arch}.posnet.block_count" - - class ConvNext: - EMBEDDING_LENGTH = "{arch}.convnext.embedding_length" - BLOCK_COUNT = "{arch}.convnext.block_count" - - class Classifier: - OUTPUT_LABELS = "{arch}.classifier.output_labels" - - class Tokenizer: - MODEL = "tokenizer.ggml.model" - PRE = "tokenizer.ggml.pre" - LIST = "tokenizer.ggml.tokens" - TOKEN_TYPE = "tokenizer.ggml.token_type" - TOKEN_TYPE_COUNT = "tokenizer.ggml.token_type_count" # for BERT-style token types - SCORES = "tokenizer.ggml.scores" - MERGES = "tokenizer.ggml.merges" - BOS_ID = "tokenizer.ggml.bos_token_id" - EOS_ID = "tokenizer.ggml.eos_token_id" - EOT_ID = "tokenizer.ggml.eot_token_id" - EOM_ID = "tokenizer.ggml.eom_token_id" - UNK_ID = "tokenizer.ggml.unknown_token_id" - SEP_ID = "tokenizer.ggml.seperator_token_id" - PAD_ID = "tokenizer.ggml.padding_token_id" - MASK_ID = "tokenizer.ggml.mask_token_id" - ADD_BOS = "tokenizer.ggml.add_bos_token" - ADD_EOS = "tokenizer.ggml.add_eos_token" - ADD_PREFIX = "tokenizer.ggml.add_space_prefix" - REMOVE_EXTRA_WS = "tokenizer.ggml.remove_extra_whitespaces" - PRECOMPILED_CHARSMAP = "tokenizer.ggml.precompiled_charsmap" - HF_JSON = "tokenizer.huggingface.json" - RWKV = "tokenizer.rwkv.world" - CHAT_TEMPLATE = "tokenizer.chat_template" - CHAT_TEMPLATE_N = "tokenizer.chat_template.{name}" - CHAT_TEMPLATES = "tokenizer.chat_templates" - # FIM/Infill special tokens constants - FIM_PRE_ID = "tokenizer.ggml.fim_pre_token_id" - FIM_SUF_ID = "tokenizer.ggml.fim_suf_token_id" - FIM_MID_ID = "tokenizer.ggml.fim_mid_token_id" - FIM_PAD_ID = "tokenizer.ggml.fim_pad_token_id" - FIM_REP_ID = "tokenizer.ggml.fim_rep_token_id" - FIM_SEP_ID = "tokenizer.ggml.fim_sep_token_id" - # deprecated: - PREFIX_ID = "tokenizer.ggml.prefix_token_id" - SUFFIX_ID = "tokenizer.ggml.suffix_token_id" - MIDDLE_ID = "tokenizer.ggml.middle_token_id" - - class Adapter: - TYPE = "adapter.type" - LORA_ALPHA = "adapter.lora.alpha" - - class Clip: - PROJECTOR_TYPE = "clip.projector_type" - HAS_VISION_ENCODER = "clip.has_vision_encoder" - HAS_AUDIO_ENCODER = "clip.has_audio_encoder" - HAS_LLAVA_PROJECTOR = "clip.has_llava_projector" - - class ClipVision: - IMAGE_SIZE = "clip.vision.image_size" - PATCH_SIZE = "clip.vision.patch_size" - EMBEDDING_LENGTH = "clip.vision.embedding_length" - FEED_FORWARD_LENGTH = "clip.vision.feed_forward_length" - PROJECTION_DIM = "clip.vision.projection_dim" - BLOCK_COUNT = "clip.vision.block_count" - IMAGE_MEAN = "clip.vision.image_mean" - IMAGE_STD = "clip.vision.image_std" - SPATIAL_MERGE_SIZE = "clip.vision.spatial_merge_size" - USE_GELU = "clip.use_gelu" - USE_SILU = "clip.use_silu" - N_WA_PATTERN = "clip.vision.n_wa_pattern" # used by qwen2.5vl - - class Attention: - HEAD_COUNT = "clip.vision.attention.head_count" - LAYERNORM_EPS = "clip.vision.attention.layer_norm_epsilon" - - class Projector: - SCALE_FACTOR = "clip.vision.projector.scale_factor" - - class ClipAudio: - NUM_MEL_BINS = "clip.audio.num_mel_bins" - EMBEDDING_LENGTH = "clip.audio.embedding_length" - FEED_FORWARD_LENGTH = "clip.audio.feed_forward_length" - PROJECTION_DIM = "clip.audio.projection_dim" - BLOCK_COUNT = "clip.audio.block_count" - - class Attention: - HEAD_COUNT = "clip.audio.attention.head_count" - LAYERNORM_EPS = "clip.audio.attention.layer_norm_epsilon" - - class Projector: - STACK_FACTOR = "clip.audio.projector.stack_factor" - -# -# recommended mapping of model tensor names for storage in gguf -# - - -class GGUFType: - MODEL = "model" - ADAPTER = "adapter" - MMPROJ = "mmproj" # dummy, unused for now - - -class MODEL_ARCH(IntEnum): - MMPROJ = auto() # dummy arch for clip.cpp - LLAMA = auto() - LLAMA4 = auto() - DECI = auto() - FALCON = auto() - BAICHUAN = auto() - GROK = auto() - GPT2 = auto() - GPTJ = auto() - GPTNEOX = auto() - MPT = auto() - STARCODER = auto() - REFACT = auto() - BERT = auto() - NOMIC_BERT = auto() - NOMIC_BERT_MOE = auto() - NEO_BERT = auto() - JINA_BERT_V2 = auto() - BLOOM = auto() - STABLELM = auto() - QWEN = auto() - QWEN2 = auto() - QWEN2MOE = auto() - QWEN2VL = auto() - QWEN3 = auto() - QWEN3MOE = auto() - PHI2 = auto() - PHI3 = auto() - PHIMOE = auto() - PLAMO = auto() - CODESHELL = auto() - ORION = auto() - INTERNLM2 = auto() - MINICPM = auto() - MINICPM3 = auto() - GEMMA = auto() - GEMMA2 = auto() - GEMMA3 = auto() - STARCODER2 = auto() - RWKV6 = auto() - RWKV6QWEN2 = auto() - RWKV7 = auto() - ARWKV7 = auto() - MAMBA = auto() - XVERSE = auto() - COMMAND_R = auto() - COHERE2 = auto() - DBRX = auto() - OLMO = auto() - OLMO2 = auto() - OLMOE = auto() - OPENELM = auto() - ARCTIC = auto() - DEEPSEEK = auto() - DEEPSEEK2 = auto() - CHATGLM = auto() - GLM4 = auto() - BITNET = auto() - T5 = auto() - T5ENCODER = auto() - JAIS = auto() - NEMOTRON = auto() - EXAONE = auto() - GRANITE = auto() - GRANITE_MOE = auto() - CHAMELEON = auto() - WAVTOKENIZER_DEC = auto() - PLM = auto() - BAILINGMOE = auto() - DOTS1 = auto() - ARCEE = auto() - - -class VISION_PROJECTOR_TYPE(IntEnum): - MLP = auto() - LDP = auto() - LDPV2 = auto() - RESAMPLER = auto() - GLM_EDGE = auto() - MERGER = auto() - GEMMA3 = auto() - - -class MODEL_TENSOR(IntEnum): - TOKEN_EMBD = auto() - TOKEN_EMBD_NORM = auto() - TOKEN_TYPES = auto() - POS_EMBD = auto() - OUTPUT = auto() - OUTPUT_NORM = auto() - ROPE_FREQS = auto() - ROPE_FACTORS_LONG = auto() - ROPE_FACTORS_SHORT = auto() - ATTN_Q = auto() - ATTN_K = auto() - ATTN_V = auto() - ATTN_QKV = auto() - ATTN_OUT = auto() - ATTN_NORM = auto() - ATTN_NORM_2 = auto() - ATTN_OUT_NORM = auto() - ATTN_POST_NORM = auto() - ATTN_ROT_EMBD = auto() - FFN_GATE_INP = auto() - FFN_GATE_INP_SHEXP = auto() - FFN_NORM = auto() - FFN_PRE_NORM = auto() - FFN_POST_NORM = auto() - FFN_GATE = auto() - FFN_DOWN = auto() - FFN_UP = auto() - FFN_ACT = auto() - FFN_NORM_EXP = auto() - FFN_GATE_EXP = auto() - FFN_DOWN_EXP = auto() - FFN_UP_EXP = auto() - FFN_GATE_SHEXP = auto() - FFN_DOWN_SHEXP = auto() - FFN_UP_SHEXP = auto() - FFN_EXP_PROBS_B = auto() - ATTN_Q_NORM = auto() - ATTN_K_NORM = auto() - LAYER_OUT_NORM = auto() - SSM_IN = auto() - SSM_CONV1D = auto() - SSM_X = auto() - SSM_DT = auto() - SSM_A = auto() - SSM_D = auto() - SSM_OUT = auto() - TIME_MIX_W0 = auto() - TIME_MIX_W1 = auto() - TIME_MIX_W2 = auto() - TIME_MIX_A0 = auto() - TIME_MIX_A1 = auto() - TIME_MIX_A2 = auto() - TIME_MIX_V0 = auto() - TIME_MIX_V1 = auto() - TIME_MIX_V2 = auto() - TIME_MIX_G1 = auto() - TIME_MIX_G2 = auto() - TIME_MIX_K_K = auto() - TIME_MIX_K_A = auto() - TIME_MIX_R_K = auto() - TIME_MIX_LERP_X = auto() - TIME_MIX_LERP_K = auto() - TIME_MIX_LERP_V = auto() - TIME_MIX_LERP_R = auto() - TIME_MIX_LERP_G = auto() - TIME_MIX_LERP_FUSED = auto() - TIME_MIX_LERP_W = auto() - TIME_MIX_FIRST = auto() - TIME_MIX_DECAY = auto() - TIME_MIX_DECAY_W1 = auto() - TIME_MIX_DECAY_W2 = auto() - TIME_MIX_KEY = auto() - TIME_MIX_VALUE = auto() - TIME_MIX_RECEPTANCE = auto() - TIME_MIX_GATE = auto() - TIME_MIX_LN = auto() - TIME_MIX_OUTPUT = auto() - CHANNEL_MIX_LERP_K = auto() - CHANNEL_MIX_LERP_R = auto() - CHANNEL_MIX_KEY = auto() - CHANNEL_MIX_RECEPTANCE = auto() - CHANNEL_MIX_VALUE = auto() - ATTN_Q_A = auto() - ATTN_Q_B = auto() - ATTN_KV_A_MQA = auto() - ATTN_KV_B = auto() - ATTN_K_B = auto() - ATTN_V_B = auto() - ATTN_Q_A_NORM = auto() - ATTN_KV_A_NORM = auto() - FFN_SUB_NORM = auto() - ATTN_SUB_NORM = auto() - DEC_ATTN_NORM = auto() - DEC_ATTN_Q = auto() - DEC_ATTN_K = auto() - DEC_ATTN_V = auto() - DEC_ATTN_OUT = auto() - DEC_ATTN_REL_B = auto() - DEC_CROSS_ATTN_NORM = auto() - DEC_CROSS_ATTN_Q = auto() - DEC_CROSS_ATTN_K = auto() - DEC_CROSS_ATTN_V = auto() - DEC_CROSS_ATTN_OUT = auto() - DEC_CROSS_ATTN_REL_B = auto() - DEC_FFN_NORM = auto() - DEC_FFN_GATE = auto() - DEC_FFN_DOWN = auto() - DEC_FFN_UP = auto() - DEC_OUTPUT_NORM = auto() - ENC_ATTN_NORM = auto() - ENC_ATTN_Q = auto() - ENC_ATTN_K = auto() - ENC_ATTN_V = auto() - ENC_ATTN_OUT = auto() - ENC_ATTN_REL_B = auto() - ENC_FFN_NORM = auto() - ENC_FFN_GATE = auto() - ENC_FFN_DOWN = auto() - ENC_FFN_UP = auto() - ENC_OUTPUT_NORM = auto() - CLS = auto() # classifier - CLS_OUT = auto() # classifier output projection - CONV1D = auto() - CONVNEXT_DW = auto() - CONVNEXT_NORM = auto() - CONVNEXT_PW1 = auto() - CONVNEXT_PW2 = auto() - CONVNEXT_GAMMA = auto() - POSNET_CONV1 = auto() - POSNET_CONV2 = auto() - POSNET_NORM = auto() - POSNET_NORM1 = auto() - POSNET_NORM2 = auto() - POSNET_ATTN_NORM = auto() - POSNET_ATTN_Q = auto() - POSNET_ATTN_K = auto() - POSNET_ATTN_V = auto() - POSNET_ATTN_OUT = auto() - # vision - V_MMPROJ = auto() - V_MMPROJ_FC = auto() - V_MMPROJ_MLP = auto() - V_MMPROJ_PEG = auto() - V_ENC_EMBD_CLS = auto() - V_ENC_EMBD_PATCH = auto() - V_ENC_EMBD_POS = auto() - V_ENC_INPUT_NORM = auto() - V_ENC_ATTN_Q = auto() - V_ENC_ATTN_Q_NORM = auto() - V_ENC_ATTN_K = auto() - V_ENC_ATTN_K_NORM = auto() - V_ENC_ATTN_V = auto() - V_ENC_ATTN_O = auto() - V_ENC_ATTN_O_NORM = auto() - V_ENC_POST_ATTN_NORM = auto() - V_ENC_FFN_UP = auto() - V_ENC_FFN_GATE = auto() - V_ENC_FFN_DOWN = auto() - V_LAYER_SCALE_1 = auto() - V_LAYER_SCALE_2 = auto() - V_PRE_NORM = auto() - V_POST_NORM = auto() - V_MM_INP_NORM = auto() - V_MM_INP_PROJ = auto() # gemma3 - V_MM_SOFT_EMB_NORM = auto() # gemma3 - V_RESMPL_POS_EMBD_K = auto() # minicpmv - V_RESMPL_ATTN_Q = auto() # minicpmv - V_RESMPL_ATTN_K = auto() # minicpmv - V_RESMPL_ATTN_V = auto() # minicpmv - V_RESMPL_ATTN_OUT = auto() # minicpmv - V_RESMPL_KV = auto() # minicpmv - V_RESMPL_KV_NORM = auto() # minicpmv - V_RESMPL_POST_NORM = auto() # minicpmv - V_RESMPL_Q_NORM = auto() # minicpmv - V_RESMPL_PROJ = auto() # minicpmv - V_RESMPL_QUERY = auto() # minicpmv - V_TOK_EMBD_IMG_BREAK = auto() # pixtral - V_MM_PATCH_MERGER = auto() # mistral small 3.1 - # audio (mtmd) - A_ENC_EMBD_POS = auto() - A_ENC_CONV1D = auto() - A_PRE_NORM = auto() - A_POST_NORM = auto() - A_ENC_ATTN_Q = auto() - A_ENC_ATTN_K = auto() - A_ENC_ATTN_V = auto() - A_ENC_INPUT_NORM = auto() - A_ENC_OUTPUT = auto() - A_ENC_OUTPUT_NORM = auto() - A_ENC_FFN_UP = auto() - A_ENC_FFN_GATE = auto() - A_ENC_FFN_DOWN = auto() - A_MMPROJ = auto() - A_MMPROJ_FC = auto() - A_MM_NORM_PRE = auto() - A_MM_NORM_MID = auto() - - -MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { - MODEL_ARCH.MMPROJ: "clip", # dummy arch for clip.cpp - MODEL_ARCH.LLAMA: "llama", - MODEL_ARCH.LLAMA4: "llama4", - MODEL_ARCH.DECI: "deci", - MODEL_ARCH.FALCON: "falcon", - MODEL_ARCH.BAICHUAN: "baichuan", - MODEL_ARCH.GROK: "grok", - MODEL_ARCH.GPT2: "gpt2", - MODEL_ARCH.GPTJ: "gptj", - MODEL_ARCH.GPTNEOX: "gptneox", - MODEL_ARCH.MPT: "mpt", - MODEL_ARCH.STARCODER: "starcoder", - MODEL_ARCH.REFACT: "refact", - MODEL_ARCH.BERT: "bert", - MODEL_ARCH.NOMIC_BERT: "nomic-bert", - MODEL_ARCH.NOMIC_BERT_MOE: "nomic-bert-moe", - MODEL_ARCH.NEO_BERT: "neo-bert", - MODEL_ARCH.JINA_BERT_V2: "jina-bert-v2", - MODEL_ARCH.BLOOM: "bloom", - MODEL_ARCH.STABLELM: "stablelm", - MODEL_ARCH.QWEN: "qwen", - MODEL_ARCH.QWEN2: "qwen2", - MODEL_ARCH.QWEN2MOE: "qwen2moe", - MODEL_ARCH.QWEN2VL: "qwen2vl", - MODEL_ARCH.QWEN3: "qwen3", - MODEL_ARCH.QWEN3MOE: "qwen3moe", - MODEL_ARCH.PHI2: "phi2", - MODEL_ARCH.PHI3: "phi3", - MODEL_ARCH.PHIMOE: "phimoe", - MODEL_ARCH.PLAMO: "plamo", - MODEL_ARCH.CODESHELL: "codeshell", - MODEL_ARCH.ORION: "orion", - MODEL_ARCH.INTERNLM2: "internlm2", - MODEL_ARCH.MINICPM: "minicpm", - MODEL_ARCH.MINICPM3: "minicpm3", - MODEL_ARCH.GEMMA: "gemma", - MODEL_ARCH.GEMMA2: "gemma2", - MODEL_ARCH.GEMMA3: "gemma3", - MODEL_ARCH.STARCODER2: "starcoder2", - MODEL_ARCH.RWKV6: "rwkv6", - MODEL_ARCH.RWKV6QWEN2: "rwkv6qwen2", - MODEL_ARCH.RWKV7: "rwkv7", - MODEL_ARCH.ARWKV7: "arwkv7", - MODEL_ARCH.MAMBA: "mamba", - MODEL_ARCH.XVERSE: "xverse", - MODEL_ARCH.COMMAND_R: "command-r", - MODEL_ARCH.COHERE2: "cohere2", - MODEL_ARCH.DBRX: "dbrx", - MODEL_ARCH.OLMO: "olmo", - MODEL_ARCH.OLMO2: "olmo2", - MODEL_ARCH.OLMOE: "olmoe", - MODEL_ARCH.OPENELM: "openelm", - MODEL_ARCH.ARCTIC: "arctic", - MODEL_ARCH.DEEPSEEK: "deepseek", - MODEL_ARCH.DEEPSEEK2: "deepseek2", - MODEL_ARCH.CHATGLM: "chatglm", - MODEL_ARCH.GLM4: "glm4", - MODEL_ARCH.BITNET: "bitnet", - MODEL_ARCH.T5: "t5", - MODEL_ARCH.T5ENCODER: "t5encoder", - MODEL_ARCH.JAIS: "jais", - MODEL_ARCH.NEMOTRON: "nemotron", - MODEL_ARCH.EXAONE: "exaone", - MODEL_ARCH.GRANITE: "granite", - MODEL_ARCH.GRANITE_MOE: "granitemoe", - MODEL_ARCH.CHAMELEON: "chameleon", - MODEL_ARCH.WAVTOKENIZER_DEC: "wavtokenizer-dec", - MODEL_ARCH.PLM: "plm", - MODEL_ARCH.BAILINGMOE: "bailingmoe", - MODEL_ARCH.DOTS1: "dots1", - MODEL_ARCH.ARCEE: "arcee", -} - -VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { - VISION_PROJECTOR_TYPE.MLP: "mlp", - VISION_PROJECTOR_TYPE.LDP: "ldp", - VISION_PROJECTOR_TYPE.LDPV2: "ldpv2", - VISION_PROJECTOR_TYPE.RESAMPLER: "resampler", - VISION_PROJECTOR_TYPE.GLM_EDGE: "adapter", - VISION_PROJECTOR_TYPE.MERGER: "qwen2vl_merger", - VISION_PROJECTOR_TYPE.GEMMA3: "gemma3", -} - -TENSOR_NAMES: dict[MODEL_TENSOR, str] = { - MODEL_TENSOR.TOKEN_EMBD: "token_embd", - MODEL_TENSOR.TOKEN_EMBD_NORM: "token_embd_norm", - MODEL_TENSOR.TOKEN_TYPES: "token_types", - MODEL_TENSOR.POS_EMBD: "position_embd", - MODEL_TENSOR.OUTPUT_NORM: "output_norm", - MODEL_TENSOR.OUTPUT: "output", - MODEL_TENSOR.ROPE_FREQS: "rope_freqs", - MODEL_TENSOR.ROPE_FACTORS_LONG: "rope_factors_long", - MODEL_TENSOR.ROPE_FACTORS_SHORT: "rope_factors_short", - MODEL_TENSOR.ATTN_NORM: "blk.{bid}.attn_norm", - MODEL_TENSOR.ATTN_NORM_2: "blk.{bid}.attn_norm_2", - MODEL_TENSOR.ATTN_QKV: "blk.{bid}.attn_qkv", - MODEL_TENSOR.ATTN_Q: "blk.{bid}.attn_q", - MODEL_TENSOR.ATTN_K: "blk.{bid}.attn_k", - MODEL_TENSOR.ATTN_V: "blk.{bid}.attn_v", - MODEL_TENSOR.ATTN_OUT: "blk.{bid}.attn_output", - MODEL_TENSOR.ATTN_ROT_EMBD: "blk.{bid}.attn_rot_embd", - MODEL_TENSOR.ATTN_Q_NORM: "blk.{bid}.attn_q_norm", - MODEL_TENSOR.ATTN_K_NORM: "blk.{bid}.attn_k_norm", - MODEL_TENSOR.ATTN_OUT_NORM: "blk.{bid}.attn_output_norm", - MODEL_TENSOR.ATTN_POST_NORM: "blk.{bid}.post_attention_norm", - MODEL_TENSOR.FFN_GATE_INP: "blk.{bid}.ffn_gate_inp", - MODEL_TENSOR.FFN_GATE_INP_SHEXP: "blk.{bid}.ffn_gate_inp_shexp", - MODEL_TENSOR.FFN_NORM: "blk.{bid}.ffn_norm", - MODEL_TENSOR.FFN_PRE_NORM: "blk.{bid}.ffn_norm", - MODEL_TENSOR.FFN_POST_NORM: "blk.{bid}.post_ffw_norm", - MODEL_TENSOR.FFN_GATE: "blk.{bid}.ffn_gate", - MODEL_TENSOR.FFN_DOWN: "blk.{bid}.ffn_down", - MODEL_TENSOR.FFN_UP: "blk.{bid}.ffn_up", - MODEL_TENSOR.FFN_GATE_SHEXP: "blk.{bid}.ffn_gate_shexp", - MODEL_TENSOR.FFN_DOWN_SHEXP: "blk.{bid}.ffn_down_shexp", - MODEL_TENSOR.FFN_UP_SHEXP: "blk.{bid}.ffn_up_shexp", - MODEL_TENSOR.FFN_ACT: "blk.{bid}.ffn", - MODEL_TENSOR.FFN_NORM_EXP: "blk.{bid}.ffn_norm_exps", - MODEL_TENSOR.FFN_GATE_EXP: "blk.{bid}.ffn_gate_exps", - MODEL_TENSOR.FFN_DOWN_EXP: "blk.{bid}.ffn_down_exps", - MODEL_TENSOR.FFN_UP_EXP: "blk.{bid}.ffn_up_exps", - MODEL_TENSOR.FFN_EXP_PROBS_B: "blk.{bid}.exp_probs_b", - MODEL_TENSOR.LAYER_OUT_NORM: "blk.{bid}.layer_output_norm", - MODEL_TENSOR.SSM_IN: "blk.{bid}.ssm_in", - MODEL_TENSOR.SSM_CONV1D: "blk.{bid}.ssm_conv1d", - MODEL_TENSOR.SSM_X: "blk.{bid}.ssm_x", - MODEL_TENSOR.SSM_DT: "blk.{bid}.ssm_dt", - MODEL_TENSOR.SSM_A: "blk.{bid}.ssm_a", - MODEL_TENSOR.SSM_D: "blk.{bid}.ssm_d", - MODEL_TENSOR.SSM_OUT: "blk.{bid}.ssm_out", - MODEL_TENSOR.TIME_MIX_W0: "blk.{bid}.time_mix_w0", - MODEL_TENSOR.TIME_MIX_W1: "blk.{bid}.time_mix_w1", - MODEL_TENSOR.TIME_MIX_W2: "blk.{bid}.time_mix_w2", - MODEL_TENSOR.TIME_MIX_A0: "blk.{bid}.time_mix_a0", - MODEL_TENSOR.TIME_MIX_A1: "blk.{bid}.time_mix_a1", - MODEL_TENSOR.TIME_MIX_A2: "blk.{bid}.time_mix_a2", - MODEL_TENSOR.TIME_MIX_V0: "blk.{bid}.time_mix_v0", - MODEL_TENSOR.TIME_MIX_V1: "blk.{bid}.time_mix_v1", - MODEL_TENSOR.TIME_MIX_V2: "blk.{bid}.time_mix_v2", - MODEL_TENSOR.TIME_MIX_G1: "blk.{bid}.time_mix_g1", - MODEL_TENSOR.TIME_MIX_G2: "blk.{bid}.time_mix_g2", - MODEL_TENSOR.TIME_MIX_K_K: "blk.{bid}.time_mix_k_k", - MODEL_TENSOR.TIME_MIX_K_A: "blk.{bid}.time_mix_k_a", - MODEL_TENSOR.TIME_MIX_R_K: "blk.{bid}.time_mix_r_k", - MODEL_TENSOR.TIME_MIX_LERP_X: "blk.{bid}.time_mix_lerp_x", - MODEL_TENSOR.TIME_MIX_LERP_K: "blk.{bid}.time_mix_lerp_k", - MODEL_TENSOR.TIME_MIX_LERP_V: "blk.{bid}.time_mix_lerp_v", - MODEL_TENSOR.TIME_MIX_LERP_R: "blk.{bid}.time_mix_lerp_r", - MODEL_TENSOR.TIME_MIX_LERP_G: "blk.{bid}.time_mix_lerp_g", - MODEL_TENSOR.TIME_MIX_LERP_FUSED: "blk.{bid}.time_mix_lerp_fused", - MODEL_TENSOR.TIME_MIX_LERP_W: "blk.{bid}.time_mix_lerp_w", - MODEL_TENSOR.TIME_MIX_FIRST: "blk.{bid}.time_mix_first", - MODEL_TENSOR.TIME_MIX_DECAY: "blk.{bid}.time_mix_decay", - MODEL_TENSOR.TIME_MIX_DECAY_W1: "blk.{bid}.time_mix_decay_w1", - MODEL_TENSOR.TIME_MIX_DECAY_W2: "blk.{bid}.time_mix_decay_w2", - MODEL_TENSOR.TIME_MIX_KEY: "blk.{bid}.time_mix_key", - MODEL_TENSOR.TIME_MIX_VALUE: "blk.{bid}.time_mix_value", - MODEL_TENSOR.TIME_MIX_RECEPTANCE: "blk.{bid}.time_mix_receptance", - MODEL_TENSOR.TIME_MIX_GATE: "blk.{bid}.time_mix_gate", - MODEL_TENSOR.TIME_MIX_LN: "blk.{bid}.time_mix_ln", - MODEL_TENSOR.TIME_MIX_OUTPUT: "blk.{bid}.time_mix_output", - MODEL_TENSOR.CHANNEL_MIX_LERP_K: "blk.{bid}.channel_mix_lerp_k", - MODEL_TENSOR.CHANNEL_MIX_LERP_R: "blk.{bid}.channel_mix_lerp_r", - MODEL_TENSOR.CHANNEL_MIX_KEY: "blk.{bid}.channel_mix_key", - MODEL_TENSOR.CHANNEL_MIX_RECEPTANCE: "blk.{bid}.channel_mix_receptance", - MODEL_TENSOR.CHANNEL_MIX_VALUE: "blk.{bid}.channel_mix_value", - MODEL_TENSOR.ATTN_Q_A: "blk.{bid}.attn_q_a", - MODEL_TENSOR.ATTN_Q_B: "blk.{bid}.attn_q_b", - MODEL_TENSOR.ATTN_KV_A_MQA: "blk.{bid}.attn_kv_a_mqa", - MODEL_TENSOR.ATTN_KV_B: "blk.{bid}.attn_kv_b", - MODEL_TENSOR.ATTN_K_B: "blk.{bid}.attn_k_b", - MODEL_TENSOR.ATTN_V_B: "blk.{bid}.attn_v_b", - MODEL_TENSOR.ATTN_Q_A_NORM: "blk.{bid}.attn_q_a_norm", - MODEL_TENSOR.ATTN_KV_A_NORM: "blk.{bid}.attn_kv_a_norm", - MODEL_TENSOR.ATTN_SUB_NORM: "blk.{bid}.attn_sub_norm", - MODEL_TENSOR.FFN_SUB_NORM: "blk.{bid}.ffn_sub_norm", - MODEL_TENSOR.DEC_ATTN_NORM: "dec.blk.{bid}.attn_norm", - MODEL_TENSOR.DEC_ATTN_Q: "dec.blk.{bid}.attn_q", - MODEL_TENSOR.DEC_ATTN_K: "dec.blk.{bid}.attn_k", - MODEL_TENSOR.DEC_ATTN_V: "dec.blk.{bid}.attn_v", - MODEL_TENSOR.DEC_ATTN_OUT: "dec.blk.{bid}.attn_o", - MODEL_TENSOR.DEC_ATTN_REL_B: "dec.blk.{bid}.attn_rel_b", - MODEL_TENSOR.DEC_CROSS_ATTN_NORM: "dec.blk.{bid}.cross_attn_norm", - MODEL_TENSOR.DEC_CROSS_ATTN_Q: "dec.blk.{bid}.cross_attn_q", - MODEL_TENSOR.DEC_CROSS_ATTN_K: "dec.blk.{bid}.cross_attn_k", - MODEL_TENSOR.DEC_CROSS_ATTN_V: "dec.blk.{bid}.cross_attn_v", - MODEL_TENSOR.DEC_CROSS_ATTN_OUT: "dec.blk.{bid}.cross_attn_o", - MODEL_TENSOR.DEC_CROSS_ATTN_REL_B: "dec.blk.{bid}.cross_attn_rel_b", - MODEL_TENSOR.DEC_FFN_NORM: "dec.blk.{bid}.ffn_norm", - MODEL_TENSOR.DEC_FFN_GATE: "dec.blk.{bid}.ffn_gate", - MODEL_TENSOR.DEC_FFN_DOWN: "dec.blk.{bid}.ffn_down", - MODEL_TENSOR.DEC_FFN_UP: "dec.blk.{bid}.ffn_up", - MODEL_TENSOR.DEC_OUTPUT_NORM: "dec.output_norm", - MODEL_TENSOR.ENC_ATTN_NORM: "enc.blk.{bid}.attn_norm", - MODEL_TENSOR.ENC_ATTN_Q: "enc.blk.{bid}.attn_q", - MODEL_TENSOR.ENC_ATTN_K: "enc.blk.{bid}.attn_k", - MODEL_TENSOR.ENC_ATTN_V: "enc.blk.{bid}.attn_v", - MODEL_TENSOR.ENC_ATTN_OUT: "enc.blk.{bid}.attn_o", - MODEL_TENSOR.ENC_ATTN_REL_B: "enc.blk.{bid}.attn_rel_b", - MODEL_TENSOR.ENC_FFN_NORM: "enc.blk.{bid}.ffn_norm", - MODEL_TENSOR.ENC_FFN_GATE: "enc.blk.{bid}.ffn_gate", - MODEL_TENSOR.ENC_FFN_DOWN: "enc.blk.{bid}.ffn_down", - MODEL_TENSOR.ENC_FFN_UP: "enc.blk.{bid}.ffn_up", - MODEL_TENSOR.ENC_OUTPUT_NORM: "enc.output_norm", - MODEL_TENSOR.CLS: "cls", - MODEL_TENSOR.CLS_OUT: "cls.output", - MODEL_TENSOR.CONV1D: "conv1d", - MODEL_TENSOR.CONVNEXT_DW: "convnext.{bid}.dw", - MODEL_TENSOR.CONVNEXT_NORM: "convnext.{bid}.norm", - MODEL_TENSOR.CONVNEXT_PW1: "convnext.{bid}.pw1", - MODEL_TENSOR.CONVNEXT_PW2: "convnext.{bid}.pw2", - MODEL_TENSOR.CONVNEXT_GAMMA: "convnext.{bid}.gamma", - MODEL_TENSOR.POSNET_CONV1: "posnet.{bid}.conv1", - MODEL_TENSOR.POSNET_CONV2: "posnet.{bid}.conv2", - MODEL_TENSOR.POSNET_NORM: "posnet.{bid}.norm", - MODEL_TENSOR.POSNET_NORM1: "posnet.{bid}.norm1", - MODEL_TENSOR.POSNET_NORM2: "posnet.{bid}.norm2", - MODEL_TENSOR.POSNET_ATTN_NORM: "posnet.{bid}.attn_norm", - MODEL_TENSOR.POSNET_ATTN_Q: "posnet.{bid}.attn_q", - MODEL_TENSOR.POSNET_ATTN_K: "posnet.{bid}.attn_k", - MODEL_TENSOR.POSNET_ATTN_V: "posnet.{bid}.attn_v", - MODEL_TENSOR.POSNET_ATTN_OUT: "posnet.{bid}.attn_output", - # vision - MODEL_TENSOR.V_MMPROJ: "mm.{bid}", - MODEL_TENSOR.V_MMPROJ_FC: "mm.model.fc", - MODEL_TENSOR.V_MMPROJ_MLP: "mm.model.mlp.{bid}", - MODEL_TENSOR.V_MMPROJ_PEG: "mm.model.peg.{bid}", - MODEL_TENSOR.V_ENC_EMBD_CLS: "v.class_embd", - MODEL_TENSOR.V_ENC_EMBD_PATCH: "v.patch_embd", - MODEL_TENSOR.V_ENC_EMBD_POS: "v.position_embd", - MODEL_TENSOR.V_ENC_ATTN_Q: "v.blk.{bid}.attn_q", - MODEL_TENSOR.V_ENC_ATTN_Q_NORM: "v.blk.{bid}.attn_q_norm", - MODEL_TENSOR.V_ENC_ATTN_K: "v.blk.{bid}.attn_k", - MODEL_TENSOR.V_ENC_ATTN_K_NORM: "v.blk.{bid}.attn_k_norm", - MODEL_TENSOR.V_ENC_ATTN_V: "v.blk.{bid}.attn_v", - MODEL_TENSOR.V_ENC_INPUT_NORM: "v.blk.{bid}.ln1", - MODEL_TENSOR.V_ENC_ATTN_O: "v.blk.{bid}.attn_out", - MODEL_TENSOR.V_ENC_ATTN_O_NORM: "v.blk.{bid}.attn_out_norm", - MODEL_TENSOR.V_ENC_POST_ATTN_NORM: "v.blk.{bid}.ln2", - MODEL_TENSOR.V_ENC_FFN_UP: "v.blk.{bid}.ffn_up", - MODEL_TENSOR.V_ENC_FFN_GATE: "v.blk.{bid}.ffn_gate", - MODEL_TENSOR.V_ENC_FFN_DOWN: "v.blk.{bid}.ffn_down", - MODEL_TENSOR.V_LAYER_SCALE_1: "v.blk.{bid}.ls1", - MODEL_TENSOR.V_LAYER_SCALE_2: "v.blk.{bid}.ls2", - MODEL_TENSOR.V_PRE_NORM: "v.pre_ln", - MODEL_TENSOR.V_POST_NORM: "v.post_ln", - MODEL_TENSOR.V_MM_INP_PROJ: "mm.input_projection", - MODEL_TENSOR.V_MM_INP_NORM: "mm.input_norm", - MODEL_TENSOR.V_MM_SOFT_EMB_NORM: "mm.soft_emb_norm", - MODEL_TENSOR.V_RESMPL_POS_EMBD_K: "resampler.pos_embd_k", - MODEL_TENSOR.V_RESMPL_ATTN_Q: "resampler.attn.q", - MODEL_TENSOR.V_RESMPL_ATTN_K: "resampler.attn.k", - MODEL_TENSOR.V_RESMPL_ATTN_V: "resampler.attn.v", - MODEL_TENSOR.V_RESMPL_ATTN_OUT: "resampler.attn.out", - MODEL_TENSOR.V_RESMPL_KV: "resampler.kv", - MODEL_TENSOR.V_RESMPL_KV_NORM: "resampler.ln_kv", - MODEL_TENSOR.V_RESMPL_POST_NORM: "resampler.ln_post", - MODEL_TENSOR.V_RESMPL_Q_NORM: "resampler.ln_q", - MODEL_TENSOR.V_RESMPL_PROJ: "resampler.proj", - MODEL_TENSOR.V_RESMPL_QUERY: "resampler.query", - MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK: "v.token_embd.img_break", # pixtral - MODEL_TENSOR.V_MM_PATCH_MERGER: "mm.patch_merger", # mistral small 3.1 - # audio (mtmd) - MODEL_TENSOR.A_ENC_EMBD_POS: "a.position_embd", - MODEL_TENSOR.A_ENC_CONV1D: "a.conv1d.{bid}", - MODEL_TENSOR.A_PRE_NORM: "a.pre_ln", - MODEL_TENSOR.A_POST_NORM: "a.post_ln", - MODEL_TENSOR.A_ENC_ATTN_Q: "a.blk.{bid}.attn_q", - MODEL_TENSOR.A_ENC_ATTN_K: "a.blk.{bid}.attn_k", - MODEL_TENSOR.A_ENC_ATTN_V: "a.blk.{bid}.attn_v", - MODEL_TENSOR.A_ENC_INPUT_NORM: "a.blk.{bid}.ln1", - MODEL_TENSOR.A_ENC_OUTPUT: "a.blk.{bid}.attn_out", - MODEL_TENSOR.A_ENC_OUTPUT_NORM: "a.blk.{bid}.ln2", - MODEL_TENSOR.A_ENC_FFN_UP: "a.blk.{bid}.ffn_up", - MODEL_TENSOR.A_ENC_FFN_GATE: "a.blk.{bid}.ffn_gate", - MODEL_TENSOR.A_ENC_FFN_DOWN: "a.blk.{bid}.ffn_down", - MODEL_TENSOR.A_MMPROJ: "mm.a.mlp.{bid}", - MODEL_TENSOR.A_MMPROJ_FC: "mm.a.fc", - MODEL_TENSOR.A_MM_NORM_PRE: "mm.a.norm_pre", - MODEL_TENSOR.A_MM_NORM_MID: "mm.a.norm_mid", -} - -MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { - MODEL_ARCH.MMPROJ: [ - MODEL_TENSOR.V_MMPROJ, - MODEL_TENSOR.V_MMPROJ_FC, - MODEL_TENSOR.V_MMPROJ_MLP, - MODEL_TENSOR.V_MMPROJ_PEG, - MODEL_TENSOR.V_ENC_EMBD_CLS, - MODEL_TENSOR.V_ENC_EMBD_PATCH, - MODEL_TENSOR.V_ENC_EMBD_POS, - MODEL_TENSOR.V_ENC_INPUT_NORM, - MODEL_TENSOR.V_ENC_ATTN_Q, - MODEL_TENSOR.V_ENC_ATTN_Q_NORM, - MODEL_TENSOR.V_ENC_ATTN_K, - MODEL_TENSOR.V_ENC_ATTN_K_NORM, - MODEL_TENSOR.V_ENC_ATTN_V, - MODEL_TENSOR.V_ENC_ATTN_O, - MODEL_TENSOR.V_ENC_ATTN_O_NORM, - MODEL_TENSOR.V_ENC_POST_ATTN_NORM, - MODEL_TENSOR.V_ENC_FFN_UP, - MODEL_TENSOR.V_ENC_FFN_GATE, - MODEL_TENSOR.V_ENC_FFN_DOWN, - MODEL_TENSOR.V_LAYER_SCALE_1, - MODEL_TENSOR.V_LAYER_SCALE_2, - MODEL_TENSOR.V_PRE_NORM, - MODEL_TENSOR.V_POST_NORM, - MODEL_TENSOR.V_MM_INP_PROJ, - MODEL_TENSOR.V_MM_INP_NORM, - MODEL_TENSOR.V_MM_SOFT_EMB_NORM, - MODEL_TENSOR.V_RESMPL_POS_EMBD_K, - MODEL_TENSOR.V_RESMPL_ATTN_Q, - MODEL_TENSOR.V_RESMPL_ATTN_K, - MODEL_TENSOR.V_RESMPL_ATTN_V, - MODEL_TENSOR.V_RESMPL_ATTN_OUT, - MODEL_TENSOR.V_RESMPL_KV, - MODEL_TENSOR.V_RESMPL_KV_NORM, - MODEL_TENSOR.V_RESMPL_POST_NORM, - MODEL_TENSOR.V_RESMPL_Q_NORM, - MODEL_TENSOR.V_RESMPL_PROJ, - MODEL_TENSOR.V_RESMPL_QUERY, - MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK, - MODEL_TENSOR.V_MM_PATCH_MERGER, - # audio - MODEL_TENSOR.A_ENC_EMBD_POS, - MODEL_TENSOR.A_ENC_CONV1D, - MODEL_TENSOR.A_PRE_NORM, - MODEL_TENSOR.A_POST_NORM, - MODEL_TENSOR.A_ENC_ATTN_Q, - MODEL_TENSOR.A_ENC_ATTN_K, - MODEL_TENSOR.A_ENC_ATTN_V, - MODEL_TENSOR.A_ENC_INPUT_NORM, - MODEL_TENSOR.A_ENC_OUTPUT, - MODEL_TENSOR.A_ENC_OUTPUT_NORM, - MODEL_TENSOR.A_ENC_FFN_UP, - MODEL_TENSOR.A_ENC_FFN_GATE, - MODEL_TENSOR.A_ENC_FFN_DOWN, - MODEL_TENSOR.A_MMPROJ, - MODEL_TENSOR.A_MMPROJ_FC, - MODEL_TENSOR.A_MM_NORM_PRE, - MODEL_TENSOR.A_MM_NORM_MID, - ], - MODEL_ARCH.LLAMA: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - ], - MODEL_ARCH.LLAMA4: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - MODEL_TENSOR.FFN_GATE_SHEXP, - MODEL_TENSOR.FFN_DOWN_SHEXP, - MODEL_TENSOR.FFN_UP_SHEXP, - ], - MODEL_ARCH.DECI: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - ], - MODEL_ARCH.GROK: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.ATTN_OUT_NORM, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - MODEL_TENSOR.LAYER_OUT_NORM, - ], - MODEL_ARCH.GPTNEOX: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.FALCON: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_NORM_2, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.BAICHUAN: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.STARCODER: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.POS_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.BERT: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.TOKEN_EMBD_NORM, - MODEL_TENSOR.TOKEN_TYPES, - MODEL_TENSOR.POS_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.ATTN_OUT_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.LAYER_OUT_NORM, - MODEL_TENSOR.CLS, - MODEL_TENSOR.CLS_OUT, - ], - MODEL_ARCH.NOMIC_BERT: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.TOKEN_EMBD_NORM, - MODEL_TENSOR.TOKEN_TYPES, - MODEL_TENSOR.POS_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.ATTN_OUT_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.LAYER_OUT_NORM, - ], - MODEL_ARCH.NOMIC_BERT_MOE: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.TOKEN_EMBD_NORM, - MODEL_TENSOR.TOKEN_TYPES, - MODEL_TENSOR.POS_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.ATTN_OUT_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - MODEL_TENSOR.LAYER_OUT_NORM, - ], - MODEL_ARCH.NEO_BERT: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.ENC_OUTPUT_NORM, - MODEL_TENSOR.CLS, - MODEL_TENSOR.CLS_OUT, - ], - MODEL_ARCH.JINA_BERT_V2: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.TOKEN_EMBD_NORM, - MODEL_TENSOR.TOKEN_TYPES, - MODEL_TENSOR.ATTN_NORM_2, - MODEL_TENSOR.ATTN_OUT_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_Q_NORM, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_K_NORM, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.LAYER_OUT_NORM, - MODEL_TENSOR.CLS, - ], - MODEL_ARCH.MPT: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FFN_ACT, - MODEL_TENSOR.ATTN_Q_NORM, - MODEL_TENSOR.ATTN_K_NORM, - MODEL_TENSOR.POS_EMBD, - ], - MODEL_ARCH.GPTJ: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.REFACT: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.BLOOM: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.TOKEN_EMBD_NORM, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.STABLELM: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.ATTN_Q_NORM, - MODEL_TENSOR.ATTN_K_NORM, - ], - MODEL_ARCH.QWEN: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.QWEN2: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.QWEN2VL: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.QWEN2MOE: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - MODEL_TENSOR.FFN_GATE_INP_SHEXP, - MODEL_TENSOR.FFN_GATE_SHEXP, - MODEL_TENSOR.FFN_DOWN_SHEXP, - MODEL_TENSOR.FFN_UP_SHEXP, - ], - MODEL_ARCH.QWEN3: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_Q_NORM, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_K_NORM, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.QWEN3MOE: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_Q_NORM, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_K_NORM, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - ], - MODEL_ARCH.PLAMO: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.GPT2: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.POS_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.PHI2: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.PHI3: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FACTORS_LONG, - MODEL_TENSOR.ROPE_FACTORS_SHORT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.PHIMOE: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FACTORS_LONG, - MODEL_TENSOR.ROPE_FACTORS_SHORT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - ], - MODEL_ARCH.CODESHELL: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.POS_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.ORION: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.INTERNLM2: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.MINICPM: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ROPE_FACTORS_LONG, - MODEL_TENSOR.ROPE_FACTORS_SHORT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - ], - MODEL_ARCH.MINICPM3: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FACTORS_LONG, - MODEL_TENSOR.ROPE_FACTORS_SHORT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q_A, - MODEL_TENSOR.ATTN_Q_B, - MODEL_TENSOR.ATTN_KV_A_MQA, - MODEL_TENSOR.ATTN_KV_B, - MODEL_TENSOR.ATTN_Q_A_NORM, - MODEL_TENSOR.ATTN_KV_A_NORM, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.GEMMA: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FFN_NORM, - ], - MODEL_ARCH.GEMMA2: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_POST_NORM, - MODEL_TENSOR.FFN_PRE_NORM, - MODEL_TENSOR.FFN_POST_NORM, - ], - MODEL_ARCH.GEMMA3: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_Q_NORM, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_K_NORM, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_POST_NORM, - MODEL_TENSOR.FFN_PRE_NORM, - MODEL_TENSOR.FFN_POST_NORM, - ], - MODEL_ARCH.STARCODER2: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.RWKV6: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.TOKEN_EMBD_NORM, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_NORM_2, - MODEL_TENSOR.TIME_MIX_W1, - MODEL_TENSOR.TIME_MIX_W2, - MODEL_TENSOR.TIME_MIX_LERP_X, - MODEL_TENSOR.TIME_MIX_LERP_K, - MODEL_TENSOR.TIME_MIX_LERP_V, - MODEL_TENSOR.TIME_MIX_LERP_R, - MODEL_TENSOR.TIME_MIX_LERP_G, - MODEL_TENSOR.TIME_MIX_LERP_W, - MODEL_TENSOR.TIME_MIX_LERP_FUSED, - MODEL_TENSOR.TIME_MIX_FIRST, - MODEL_TENSOR.TIME_MIX_DECAY, - MODEL_TENSOR.TIME_MIX_DECAY_W1, - MODEL_TENSOR.TIME_MIX_DECAY_W2, - MODEL_TENSOR.TIME_MIX_KEY, - MODEL_TENSOR.TIME_MIX_VALUE, - MODEL_TENSOR.TIME_MIX_RECEPTANCE, - MODEL_TENSOR.TIME_MIX_GATE, - MODEL_TENSOR.TIME_MIX_LN, - MODEL_TENSOR.TIME_MIX_OUTPUT, - MODEL_TENSOR.CHANNEL_MIX_LERP_K, - MODEL_TENSOR.CHANNEL_MIX_LERP_R, - MODEL_TENSOR.CHANNEL_MIX_KEY, - MODEL_TENSOR.CHANNEL_MIX_RECEPTANCE, - MODEL_TENSOR.CHANNEL_MIX_VALUE, - ], - MODEL_ARCH.RWKV6QWEN2: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.TIME_MIX_W1, - MODEL_TENSOR.TIME_MIX_W2, - MODEL_TENSOR.TIME_MIX_LERP_X, - MODEL_TENSOR.TIME_MIX_LERP_K, - MODEL_TENSOR.TIME_MIX_LERP_V, - MODEL_TENSOR.TIME_MIX_LERP_R, - MODEL_TENSOR.TIME_MIX_LERP_G, - MODEL_TENSOR.TIME_MIX_LERP_W, - MODEL_TENSOR.TIME_MIX_LERP_FUSED, - MODEL_TENSOR.TIME_MIX_FIRST, - MODEL_TENSOR.TIME_MIX_DECAY, - MODEL_TENSOR.TIME_MIX_DECAY_W1, - MODEL_TENSOR.TIME_MIX_DECAY_W2, - MODEL_TENSOR.TIME_MIX_KEY, - MODEL_TENSOR.TIME_MIX_VALUE, - MODEL_TENSOR.TIME_MIX_RECEPTANCE, - MODEL_TENSOR.TIME_MIX_GATE, - MODEL_TENSOR.TIME_MIX_LN, - MODEL_TENSOR.TIME_MIX_OUTPUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.RWKV7: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.TOKEN_EMBD_NORM, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_NORM_2, - MODEL_TENSOR.TIME_MIX_LERP_FUSED, - MODEL_TENSOR.TIME_MIX_W0, - MODEL_TENSOR.TIME_MIX_W1, - MODEL_TENSOR.TIME_MIX_W2, - MODEL_TENSOR.TIME_MIX_A0, - MODEL_TENSOR.TIME_MIX_A1, - MODEL_TENSOR.TIME_MIX_A2, - MODEL_TENSOR.TIME_MIX_V0, - MODEL_TENSOR.TIME_MIX_V1, - MODEL_TENSOR.TIME_MIX_V2, - MODEL_TENSOR.TIME_MIX_G1, - MODEL_TENSOR.TIME_MIX_G2, - MODEL_TENSOR.TIME_MIX_K_K, - MODEL_TENSOR.TIME_MIX_K_A, - MODEL_TENSOR.TIME_MIX_R_K, - MODEL_TENSOR.TIME_MIX_KEY, - MODEL_TENSOR.TIME_MIX_VALUE, - MODEL_TENSOR.TIME_MIX_RECEPTANCE, - MODEL_TENSOR.TIME_MIX_LN, - MODEL_TENSOR.TIME_MIX_OUTPUT, - MODEL_TENSOR.CHANNEL_MIX_LERP_K, - MODEL_TENSOR.CHANNEL_MIX_KEY, - MODEL_TENSOR.CHANNEL_MIX_VALUE, - ], - MODEL_ARCH.ARWKV7: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.TOKEN_EMBD_NORM, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.TIME_MIX_LERP_FUSED, - MODEL_TENSOR.TIME_MIX_W0, - MODEL_TENSOR.TIME_MIX_W1, - MODEL_TENSOR.TIME_MIX_W2, - MODEL_TENSOR.TIME_MIX_A0, - MODEL_TENSOR.TIME_MIX_A1, - MODEL_TENSOR.TIME_MIX_A2, - MODEL_TENSOR.TIME_MIX_V0, - MODEL_TENSOR.TIME_MIX_V1, - MODEL_TENSOR.TIME_MIX_V2, - MODEL_TENSOR.TIME_MIX_G1, - MODEL_TENSOR.TIME_MIX_G2, - MODEL_TENSOR.TIME_MIX_K_K, - MODEL_TENSOR.TIME_MIX_K_A, - MODEL_TENSOR.TIME_MIX_R_K, - MODEL_TENSOR.TIME_MIX_KEY, - MODEL_TENSOR.TIME_MIX_VALUE, - MODEL_TENSOR.TIME_MIX_RECEPTANCE, - MODEL_TENSOR.TIME_MIX_LN, - MODEL_TENSOR.TIME_MIX_OUTPUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.MAMBA: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.SSM_IN, - MODEL_TENSOR.SSM_CONV1D, - MODEL_TENSOR.SSM_X, - MODEL_TENSOR.SSM_DT, - MODEL_TENSOR.SSM_A, - MODEL_TENSOR.SSM_D, - MODEL_TENSOR.SSM_OUT, - ], - MODEL_ARCH.XVERSE: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.COMMAND_R: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.ATTN_K_NORM, - MODEL_TENSOR.ATTN_Q_NORM, - ], - MODEL_ARCH.COHERE2: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.DBRX: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_OUT_NORM, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - ], - MODEL_ARCH.OLMO: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.OLMO2: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_POST_NORM, - MODEL_TENSOR.ATTN_Q_NORM, - MODEL_TENSOR.ATTN_K_NORM, - MODEL_TENSOR.FFN_POST_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.OLMOE: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q_NORM, - MODEL_TENSOR.ATTN_K_NORM, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_UP_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - ], - MODEL_ARCH.OPENELM: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_Q_NORM, - MODEL_TENSOR.ATTN_K_NORM, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.ARCTIC: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FFN_NORM_EXP, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - ], - MODEL_ARCH.DEEPSEEK: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - MODEL_TENSOR.FFN_GATE_SHEXP, - MODEL_TENSOR.FFN_DOWN_SHEXP, - MODEL_TENSOR.FFN_UP_SHEXP, - ], - MODEL_ARCH.DEEPSEEK2: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_Q_A, - MODEL_TENSOR.ATTN_Q_B, - MODEL_TENSOR.ATTN_KV_A_MQA, - MODEL_TENSOR.ATTN_KV_B, - MODEL_TENSOR.ATTN_K_B, - MODEL_TENSOR.ATTN_V_B, - MODEL_TENSOR.ATTN_Q_A_NORM, - MODEL_TENSOR.ATTN_KV_A_NORM, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - MODEL_TENSOR.FFN_GATE_SHEXP, - MODEL_TENSOR.FFN_DOWN_SHEXP, - MODEL_TENSOR.FFN_UP_SHEXP, - MODEL_TENSOR.FFN_EXP_PROBS_B, - ], - MODEL_ARCH.PLM: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_KV_A_MQA, - MODEL_TENSOR.ATTN_KV_A_NORM, - MODEL_TENSOR.ATTN_KV_B, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FFN_DOWN, - ], - MODEL_ARCH.CHATGLM : [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.GLM4 : [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.ATTN_POST_NORM, - MODEL_TENSOR.FFN_POST_NORM, - ], - MODEL_ARCH.BITNET: [ - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.ATTN_SUB_NORM, - MODEL_TENSOR.FFN_SUB_NORM, - ], - MODEL_ARCH.T5: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.DEC_ATTN_NORM, - MODEL_TENSOR.DEC_ATTN_Q, - MODEL_TENSOR.DEC_ATTN_K, - MODEL_TENSOR.DEC_ATTN_V, - MODEL_TENSOR.DEC_ATTN_OUT, - MODEL_TENSOR.DEC_ATTN_REL_B, - MODEL_TENSOR.DEC_CROSS_ATTN_NORM, - MODEL_TENSOR.DEC_CROSS_ATTN_Q, - MODEL_TENSOR.DEC_CROSS_ATTN_K, - MODEL_TENSOR.DEC_CROSS_ATTN_V, - MODEL_TENSOR.DEC_CROSS_ATTN_OUT, - MODEL_TENSOR.DEC_CROSS_ATTN_REL_B, - MODEL_TENSOR.DEC_FFN_NORM, - MODEL_TENSOR.DEC_FFN_GATE, - MODEL_TENSOR.DEC_FFN_DOWN, - MODEL_TENSOR.DEC_FFN_UP, - MODEL_TENSOR.DEC_OUTPUT_NORM, - MODEL_TENSOR.ENC_ATTN_NORM, - MODEL_TENSOR.ENC_ATTN_Q, - MODEL_TENSOR.ENC_ATTN_K, - MODEL_TENSOR.ENC_ATTN_V, - MODEL_TENSOR.ENC_ATTN_OUT, - MODEL_TENSOR.ENC_ATTN_REL_B, - MODEL_TENSOR.ENC_FFN_NORM, - MODEL_TENSOR.ENC_FFN_GATE, - MODEL_TENSOR.ENC_FFN_DOWN, - MODEL_TENSOR.ENC_FFN_UP, - MODEL_TENSOR.ENC_OUTPUT_NORM, - ], - MODEL_ARCH.T5ENCODER: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ENC_ATTN_NORM, - MODEL_TENSOR.ENC_ATTN_Q, - MODEL_TENSOR.ENC_ATTN_K, - MODEL_TENSOR.ENC_ATTN_V, - MODEL_TENSOR.ENC_ATTN_OUT, - MODEL_TENSOR.ENC_ATTN_REL_B, - MODEL_TENSOR.ENC_FFN_NORM, - MODEL_TENSOR.ENC_FFN_GATE, - MODEL_TENSOR.ENC_FFN_DOWN, - MODEL_TENSOR.ENC_FFN_UP, - MODEL_TENSOR.ENC_OUTPUT_NORM, - ], - MODEL_ARCH.JAIS: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_QKV, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.NEMOTRON: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.EXAONE: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.GRANITE: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.GRANITE_MOE: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - MODEL_TENSOR.FFN_GATE_SHEXP, - MODEL_TENSOR.FFN_UP_SHEXP, - MODEL_TENSOR.FFN_DOWN_SHEXP, - ], - MODEL_ARCH.CHAMELEON: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_Q_NORM, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_K_NORM, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - MODEL_ARCH.WAVTOKENIZER_DEC: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.TOKEN_EMBD_NORM, - MODEL_TENSOR.CONV1D, - MODEL_TENSOR.CONVNEXT_DW, - MODEL_TENSOR.CONVNEXT_NORM, - MODEL_TENSOR.CONVNEXT_PW1, - MODEL_TENSOR.CONVNEXT_PW2, - MODEL_TENSOR.CONVNEXT_GAMMA, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.POSNET_CONV1, - MODEL_TENSOR.POSNET_CONV2, - MODEL_TENSOR.POSNET_NORM, - MODEL_TENSOR.POSNET_NORM1, - MODEL_TENSOR.POSNET_NORM2, - MODEL_TENSOR.POSNET_ATTN_NORM, - MODEL_TENSOR.POSNET_ATTN_Q, - MODEL_TENSOR.POSNET_ATTN_K, - MODEL_TENSOR.POSNET_ATTN_V, - MODEL_TENSOR.POSNET_ATTN_OUT, - ], - MODEL_ARCH.BAILINGMOE: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_UP_EXP, - MODEL_TENSOR.FFN_GATE_SHEXP, - MODEL_TENSOR.FFN_DOWN_SHEXP, - MODEL_TENSOR.FFN_UP_SHEXP, - ], - MODEL_ARCH.DOTS1: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_Q_NORM, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_K_NORM, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.FFN_EXP_PROBS_B, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_GATE_EXP, - MODEL_TENSOR.FFN_GATE_INP, - MODEL_TENSOR.FFN_GATE_SHEXP, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_DOWN_EXP, - MODEL_TENSOR.FFN_DOWN_SHEXP, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FFN_UP_EXP, - MODEL_TENSOR.FFN_UP_SHEXP, - ], - MODEL_ARCH.ARCEE: [ - MODEL_TENSOR.TOKEN_EMBD, - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.OUTPUT, - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_ROT_EMBD, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - ], - # TODO -} - -# tensors that will not be serialized -MODEL_TENSOR_SKIP: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { - MODEL_ARCH.LLAMA: [ - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_ROT_EMBD, - ], - MODEL_ARCH.DECI: [ - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_ROT_EMBD, - ], - MODEL_ARCH.BAICHUAN: [ - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_ROT_EMBD, - ], - MODEL_ARCH.QWEN: [ - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_ROT_EMBD, - ], - MODEL_ARCH.CODESHELL: [ - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_ROT_EMBD, - ], - MODEL_ARCH.ORION: [ - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_ROT_EMBD, - ], - MODEL_ARCH.STARCODER2: [ - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_ROT_EMBD, - ], - MODEL_ARCH.XVERSE: [ - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_ROT_EMBD, - ], - MODEL_ARCH.DEEPSEEK: [ - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_ROT_EMBD, - ], - MODEL_ARCH.DEEPSEEK2: [ - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_ROT_EMBD, - ], - MODEL_ARCH.CHATGLM: [ - MODEL_TENSOR.ROPE_FREQS, - ], - MODEL_ARCH.NEMOTRON: [ - MODEL_TENSOR.ROPE_FREQS, - MODEL_TENSOR.ATTN_ROT_EMBD, - ], - MODEL_ARCH.BAILINGMOE: [ - MODEL_TENSOR.ROPE_FREQS, - ], -} - -# -# types -# - - -class TokenType(IntEnum): - NORMAL = 1 - UNKNOWN = 2 - CONTROL = 3 - USER_DEFINED = 4 - UNUSED = 5 - BYTE = 6 - - -class RopeScalingType(Enum): - NONE = 'none' - LINEAR = 'linear' - YARN = 'yarn' - LONGROPE = 'longrope' - - -class PoolingType(IntEnum): - NONE = 0 - MEAN = 1 - CLS = 2 - LAST = 3 - RANK = 4 - - -class GGMLQuantizationType(IntEnum): - F32 = 0 - F16 = 1 - Q4_0 = 2 - Q4_1 = 3 - Q5_0 = 6 - Q5_1 = 7 - Q8_0 = 8 - Q8_1 = 9 - Q2_K = 10 - Q3_K = 11 - Q4_K = 12 - Q5_K = 13 - Q6_K = 14 - Q8_K = 15 - IQ2_XXS = 16 - IQ2_XS = 17 - IQ3_XXS = 18 - IQ1_S = 19 - IQ4_NL = 20 - IQ3_S = 21 - IQ2_S = 22 - IQ4_XS = 23 - I8 = 24 - I16 = 25 - I32 = 26 - I64 = 27 - F64 = 28 - IQ1_M = 29 - BF16 = 30 - TQ1_0 = 34 - TQ2_0 = 35 - - -class ExpertGatingFuncType(IntEnum): - SOFTMAX = 1 - SIGMOID = 2 - - -# TODO: add GGMLFileType from ggml_ftype in ggml.h - - -# from llama_ftype in llama.h -# ALL VALUES SHOULD BE THE SAME HERE AS THEY ARE OVER THERE. -class LlamaFileType(IntEnum): - ALL_F32 = 0 - MOSTLY_F16 = 1 # except 1d tensors - MOSTLY_Q4_0 = 2 # except 1d tensors - MOSTLY_Q4_1 = 3 # except 1d tensors - # MOSTLY_Q4_1_SOME_F16 = 4 # tok_embeddings.weight and output.weight are F16 - # MOSTLY_Q4_2 = 5 # support has been removed - # MOSTLY_Q4_3 = 6 # support has been removed - MOSTLY_Q8_0 = 7 # except 1d tensors - MOSTLY_Q5_0 = 8 # except 1d tensors - MOSTLY_Q5_1 = 9 # except 1d tensors - MOSTLY_Q2_K = 10 # except 1d tensors - MOSTLY_Q3_K_S = 11 # except 1d tensors - MOSTLY_Q3_K_M = 12 # except 1d tensors - MOSTLY_Q3_K_L = 13 # except 1d tensors - MOSTLY_Q4_K_S = 14 # except 1d tensors - MOSTLY_Q4_K_M = 15 # except 1d tensors - MOSTLY_Q5_K_S = 16 # except 1d tensors - MOSTLY_Q5_K_M = 17 # except 1d tensors - MOSTLY_Q6_K = 18 # except 1d tensors - MOSTLY_IQ2_XXS = 19 # except 1d tensors - MOSTLY_IQ2_XS = 20 # except 1d tensors - MOSTLY_Q2_K_S = 21 # except 1d tensors - MOSTLY_IQ3_XS = 22 # except 1d tensors - MOSTLY_IQ3_XXS = 23 # except 1d tensors - MOSTLY_IQ1_S = 24 # except 1d tensors - MOSTLY_IQ4_NL = 25 # except 1d tensors - MOSTLY_IQ3_S = 26 # except 1d tensors - MOSTLY_IQ3_M = 27 # except 1d tensors - MOSTLY_IQ2_S = 28 # except 1d tensors - MOSTLY_IQ2_M = 29 # except 1d tensors - MOSTLY_IQ4_XS = 30 # except 1d tensors - MOSTLY_IQ1_M = 31 # except 1d tensors - MOSTLY_BF16 = 32 # except 1d tensors - # MOSTLY_Q4_0_4_4 = 33 # removed from gguf files, use Q4_0 and runtime repack - # MOSTLY_Q4_0_4_8 = 34 # removed from gguf files, use Q4_0 and runtime repack - # MOSTLY_Q4_0_8_8 = 35 # removed from gguf files, use Q4_0 and runtime repack - MOSTLY_TQ1_0 = 36 # except 1d tensors - MOSTLY_TQ2_0 = 37 # except 1d tensors - - GUESSED = 1024 # not specified in the model file - - -class GGUFEndian(IntEnum): - LITTLE = 0 - BIG = 1 - - -class GGUFValueType(IntEnum): - UINT8 = 0 - INT8 = 1 - UINT16 = 2 - INT16 = 3 - UINT32 = 4 - INT32 = 5 - FLOAT32 = 6 - BOOL = 7 - STRING = 8 - ARRAY = 9 - UINT64 = 10 - INT64 = 11 - FLOAT64 = 12 - - @staticmethod - def get_type(val: Any) -> GGUFValueType: - if isinstance(val, (str, bytes, bytearray)): - return GGUFValueType.STRING - elif isinstance(val, list): - return GGUFValueType.ARRAY - elif isinstance(val, float): - return GGUFValueType.FLOAT32 - elif isinstance(val, bool): - return GGUFValueType.BOOL - elif isinstance(val, int): - return GGUFValueType.INT32 - # TODO: need help with 64-bit types in Python - else: - raise ValueError(f"Unknown type: {type(val)}") - - -class VisionProjectorType: - GEMMA3 = "gemma3" - IDEFICS3 = "idefics3" - PIXTRAL = "pixtral" - LLAMA4 = "llama4" - QWEN2VL = "qwen2vl_merger" - QWEN25VL = "qwen2.5vl_merger" - ULTRAVOX = "ultravox" - INTERNVL = "internvl" - QWEN2A = "qwen2a" # audio - QWEN25O = "qwen2.5o" # omni - - -# Items here are (block size, type size) -QK_K = 256 -GGML_QUANT_SIZES: dict[GGMLQuantizationType, tuple[int, int]] = { - GGMLQuantizationType.F32: (1, 4), - GGMLQuantizationType.F16: (1, 2), - GGMLQuantizationType.Q4_0: (32, 2 + 16), - GGMLQuantizationType.Q4_1: (32, 2 + 2 + 16), - GGMLQuantizationType.Q5_0: (32, 2 + 4 + 16), - GGMLQuantizationType.Q5_1: (32, 2 + 2 + 4 + 16), - GGMLQuantizationType.Q8_0: (32, 2 + 32), - GGMLQuantizationType.Q8_1: (32, 4 + 4 + 32), - GGMLQuantizationType.Q2_K: (256, 2 + 2 + QK_K // 16 + QK_K // 4), - GGMLQuantizationType.Q3_K: (256, 2 + QK_K // 4 + QK_K // 8 + 12), - GGMLQuantizationType.Q4_K: (256, 2 + 2 + QK_K // 2 + 12), - GGMLQuantizationType.Q5_K: (256, 2 + 2 + QK_K // 2 + QK_K // 8 + 12), - GGMLQuantizationType.Q6_K: (256, 2 + QK_K // 2 + QK_K // 4 + QK_K // 16), - GGMLQuantizationType.Q8_K: (256, 4 + QK_K + QK_K // 8), - GGMLQuantizationType.IQ2_XXS: (256, 2 + QK_K // 4), - GGMLQuantizationType.IQ2_XS: (256, 2 + QK_K // 4 + QK_K // 32), - GGMLQuantizationType.IQ3_XXS: (256, 2 + QK_K // 4 + QK_K // 8), - GGMLQuantizationType.IQ1_S: (256, 2 + QK_K // 8 + QK_K // 16), - GGMLQuantizationType.IQ4_NL: (32, 2 + 16), - GGMLQuantizationType.IQ3_S: (256, 2 + QK_K // 4 + QK_K // 8 + QK_K // 32 + 4), - GGMLQuantizationType.IQ2_S: (256, 2 + QK_K // 4 + QK_K // 16), - GGMLQuantizationType.IQ4_XS: (256, 2 + 2 + QK_K // 2 + QK_K // 64), - GGMLQuantizationType.I8: (1, 1), - GGMLQuantizationType.I16: (1, 2), - GGMLQuantizationType.I32: (1, 4), - GGMLQuantizationType.I64: (1, 8), - GGMLQuantizationType.F64: (1, 8), - GGMLQuantizationType.IQ1_M: (256, QK_K // 8 + QK_K // 16 + QK_K // 32), - GGMLQuantizationType.BF16: (1, 2), - GGMLQuantizationType.TQ1_0: (256, 2 + 4 * 13), - GGMLQuantizationType.TQ2_0: (256, 2 + 64), -} - - -# Aliases for backward compatibility. - -# general -KEY_GENERAL_ARCHITECTURE = Keys.General.ARCHITECTURE -KEY_GENERAL_QUANTIZATION_VERSION = Keys.General.QUANTIZATION_VERSION -KEY_GENERAL_ALIGNMENT = Keys.General.ALIGNMENT -KEY_GENERAL_NAME = Keys.General.NAME -KEY_GENERAL_AUTHOR = Keys.General.AUTHOR -KEY_GENERAL_URL = Keys.General.URL -KEY_GENERAL_DESCRIPTION = Keys.General.DESCRIPTION -KEY_GENERAL_LICENSE = Keys.General.LICENSE -KEY_GENERAL_SOURCE_URL = Keys.General.SOURCE_URL -KEY_GENERAL_FILE_TYPE = Keys.General.FILE_TYPE - -# LLM -KEY_VOCAB_SIZE = Keys.LLM.VOCAB_SIZE -KEY_CONTEXT_LENGTH = Keys.LLM.CONTEXT_LENGTH -KEY_EMBEDDING_LENGTH = Keys.LLM.EMBEDDING_LENGTH -KEY_BLOCK_COUNT = Keys.LLM.BLOCK_COUNT -KEY_FEED_FORWARD_LENGTH = Keys.LLM.FEED_FORWARD_LENGTH -KEY_USE_PARALLEL_RESIDUAL = Keys.LLM.USE_PARALLEL_RESIDUAL -KEY_TENSOR_DATA_LAYOUT = Keys.LLM.TENSOR_DATA_LAYOUT - -# attention -KEY_ATTENTION_HEAD_COUNT = Keys.Attention.HEAD_COUNT -KEY_ATTENTION_HEAD_COUNT_KV = Keys.Attention.HEAD_COUNT_KV -KEY_ATTENTION_MAX_ALIBI_BIAS = Keys.Attention.MAX_ALIBI_BIAS -KEY_ATTENTION_CLAMP_KQV = Keys.Attention.CLAMP_KQV -KEY_ATTENTION_LAYERNORM_EPS = Keys.Attention.LAYERNORM_EPS -KEY_ATTENTION_LAYERNORM_RMS_EPS = Keys.Attention.LAYERNORM_RMS_EPS - -# RoPE -KEY_ROPE_DIMENSION_COUNT = Keys.Rope.DIMENSION_COUNT -KEY_ROPE_FREQ_BASE = Keys.Rope.FREQ_BASE -KEY_ROPE_SCALING_TYPE = Keys.Rope.SCALING_TYPE -KEY_ROPE_SCALING_FACTOR = Keys.Rope.SCALING_FACTOR -KEY_ROPE_SCALING_ORIG_CTX_LEN = Keys.Rope.SCALING_ORIG_CTX_LEN -KEY_ROPE_SCALING_FINETUNED = Keys.Rope.SCALING_FINETUNED - -# SSM -KEY_SSM_CONV_KERNEL = Keys.SSM.CONV_KERNEL -KEY_SSM_INNER_SIZE = Keys.SSM.INNER_SIZE -KEY_SSM_STATE_SIZE = Keys.SSM.STATE_SIZE -KEY_SSM_TIME_STEP_RANK = Keys.SSM.TIME_STEP_RANK -KEY_SSM_DT_B_C_RMS = Keys.SSM.DT_B_C_RMS - -# tokenization -KEY_TOKENIZER_MODEL = Keys.Tokenizer.MODEL -KEY_TOKENIZER_PRE = Keys.Tokenizer.PRE -KEY_TOKENIZER_LIST = Keys.Tokenizer.LIST -KEY_TOKENIZER_TOKEN_TYPE = Keys.Tokenizer.TOKEN_TYPE -KEY_TOKENIZER_SCORES = Keys.Tokenizer.SCORES -KEY_TOKENIZER_MERGES = Keys.Tokenizer.MERGES -KEY_TOKENIZER_BOS_ID = Keys.Tokenizer.BOS_ID -KEY_TOKENIZER_EOS_ID = Keys.Tokenizer.EOS_ID -KEY_TOKENIZER_EOT_ID = Keys.Tokenizer.EOT_ID -KEY_TOKENIZER_EOM_ID = Keys.Tokenizer.EOM_ID -KEY_TOKENIZER_UNK_ID = Keys.Tokenizer.UNK_ID -KEY_TOKENIZER_SEP_ID = Keys.Tokenizer.SEP_ID -KEY_TOKENIZER_PAD_ID = Keys.Tokenizer.PAD_ID -KEY_TOKENIZER_MASK_ID = Keys.Tokenizer.MASK_ID -KEY_TOKENIZER_HF_JSON = Keys.Tokenizer.HF_JSON -KEY_TOKENIZER_RWKV = Keys.Tokenizer.RWKV - -KEY_TOKENIZER_FIM_PRE_ID = Keys.Tokenizer.FIM_PRE_ID -KEY_TOKENIZER_FIM_SUF_ID = Keys.Tokenizer.FIM_SUF_ID -KEY_TOKENIZER_FIM_MID_ID = Keys.Tokenizer.FIM_MID_ID -KEY_TOKENIZER_FIM_PAD_ID = Keys.Tokenizer.FIM_PAD_ID -KEY_TOKENIZER_FIM_REP_ID = Keys.Tokenizer.FIM_REP_ID -KEY_TOKENIZER_FIM_SEP_ID = Keys.Tokenizer.FIM_SEP_ID - -# deprecated -KEY_TOKENIZER_PREFIX_ID = Keys.Tokenizer.PREFIX_ID -KEY_TOKENIZER_SUFFIX_ID = Keys.Tokenizer.SUFFIX_ID -KEY_TOKENIZER_MIDDLE_ID = Keys.Tokenizer.MIDDLE_ID diff --git a/venv/lib/python3.10/site-packages/gguf/gguf.py b/venv/lib/python3.10/site-packages/gguf/gguf.py deleted file mode 100644 index 651a81eb828248728f854c85c1a437b52892f275..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/gguf.py +++ /dev/null @@ -1,15 +0,0 @@ -# This file left for compatibility. If you want to use the GGUF API from Python -# then don't import gguf/gguf.py directly. If you're looking for examples, see the -# examples/ directory for gguf-py - -import importlib -import sys -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).parent.parent)) - -# Compatibility for people trying to import gguf/gguf.py directly instead of as a package. -importlib.invalidate_caches() -import gguf # noqa: E402 - -importlib.reload(gguf) diff --git a/venv/lib/python3.10/site-packages/gguf/gguf_reader.py b/venv/lib/python3.10/site-packages/gguf/gguf_reader.py deleted file mode 100644 index d87e8f72321b3f776cf214e117e21cdde9e1e18b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/gguf_reader.py +++ /dev/null @@ -1,367 +0,0 @@ -# -# GGUF file reading/modification support. For API usage information, -# please see the files scripts/ for some fairly simple examples. -# -from __future__ import annotations - -import logging -import os -import sys -from collections import OrderedDict -from typing import Any, Literal, NamedTuple, TypeVar, Union - -import numpy as np -import numpy.typing as npt - -from .quants import quant_shape_to_byte_shape - -if __name__ == "__main__": - from pathlib import Path - - # Allow running file in package as a script. - sys.path.insert(0, str(Path(__file__).parent.parent)) - -from gguf.constants import ( - GGML_QUANT_SIZES, - GGUF_DEFAULT_ALIGNMENT, - GGUF_MAGIC, - GGUF_VERSION, - GGMLQuantizationType, - GGUFValueType, - GGUFEndian, -) - -logger = logging.getLogger(__name__) - -READER_SUPPORTED_VERSIONS = [2, GGUF_VERSION] - - -class ReaderField(NamedTuple): - # Offset to start of this field. - offset: int - - # Name of the field (not necessarily from file data). - name: str - - # Data parts. Some types have multiple components, such as strings - # that consist of a length followed by the string data. - parts: list[npt.NDArray[Any]] = [] - - # Indexes into parts that we can call the actual data. For example - # an array of strings will be populated with indexes to the actual - # string data. - data: list[int] = [-1] - - types: list[GGUFValueType] = [] - - def contents(self, index_or_slice: int | slice = slice(None)) -> Any: - if self.types: - to_string = lambda x: str(x.tobytes(), encoding='utf-8') # noqa: E731 - main_type = self.types[0] - - if main_type == GGUFValueType.ARRAY: - sub_type = self.types[-1] - - if sub_type == GGUFValueType.STRING: - indices = self.data[index_or_slice] - - if isinstance(index_or_slice, int): - return to_string(self.parts[indices]) # type: ignore - else: - return [to_string(self.parts[idx]) for idx in indices] # type: ignore - else: - # FIXME: When/if _get_field_parts() support multi-dimensional arrays, this must do so too - - # Check if it's unsafe to perform slice optimization on data - # if any(True for idx in self.data if len(self.parts[idx]) != 1): - # optim_slice = slice(None) - # else: - # optim_slice = index_or_slice - # index_or_slice = slice(None) - - # if isinstance(optim_slice, int): - # return self.parts[self.data[optim_slice]].tolist()[0] - # else: - # return [pv for idx in self.data[optim_slice] for pv in self.parts[idx].tolist()][index_or_slice] - - if isinstance(index_or_slice, int): - return self.parts[self.data[index_or_slice]].tolist()[0] - else: - return [pv for idx in self.data[index_or_slice] for pv in self.parts[idx].tolist()] - - if main_type == GGUFValueType.STRING: - return to_string(self.parts[-1]) - else: - return self.parts[-1].tolist()[0] - - return None - - -class ReaderTensor(NamedTuple): - name: str - tensor_type: GGMLQuantizationType - shape: npt.NDArray[np.uint32] - n_elements: int - n_bytes: int - data_offset: int - data: npt.NDArray[Any] - field: ReaderField - - -class GGUFReader: - # I - same as host, S - swapped - byte_order: Literal['I', 'S'] = 'I' - alignment: int = GGUF_DEFAULT_ALIGNMENT - data_offset: int - - # Note: Internal helper, API may change. - gguf_scalar_to_np: dict[GGUFValueType, type[np.generic]] = { - GGUFValueType.UINT8: np.uint8, - GGUFValueType.INT8: np.int8, - GGUFValueType.UINT16: np.uint16, - GGUFValueType.INT16: np.int16, - GGUFValueType.UINT32: np.uint32, - GGUFValueType.INT32: np.int32, - GGUFValueType.FLOAT32: np.float32, - GGUFValueType.UINT64: np.uint64, - GGUFValueType.INT64: np.int64, - GGUFValueType.FLOAT64: np.float64, - GGUFValueType.BOOL: np.bool_, - } - - def __init__(self, path: os.PathLike[str] | str, mode: Literal['r', 'r+', 'c'] = 'r'): - self.data = np.memmap(path, mode = mode) - offs = 0 - - # Check for GGUF magic - if self._get(offs, np.uint32, override_order = '<')[0] != GGUF_MAGIC: - raise ValueError('GGUF magic invalid') - offs += 4 - - # Check GGUF version - temp_version = self._get(offs, np.uint32) - if temp_version[0] & 65535 == 0: - # If we get 0 here that means it's (probably) a GGUF file created for - # the opposite byte order of the machine this script is running on. - self.byte_order = 'S' - temp_version = temp_version.view(temp_version.dtype.newbyteorder(self.byte_order)) - version = temp_version[0] - if version not in READER_SUPPORTED_VERSIONS: - raise ValueError(f'Sorry, file appears to be version {version} which we cannot handle') - if sys.byteorder == "little": - # Host is little endian - host_endian = GGUFEndian.LITTLE - swapped_endian = GGUFEndian.BIG - else: - # Sorry PDP or other weird systems that don't use BE or LE. - host_endian = GGUFEndian.BIG - swapped_endian = GGUFEndian.LITTLE - self.endianess = swapped_endian if self.byte_order == "S" else host_endian - self.fields: OrderedDict[str, ReaderField] = OrderedDict() - self.tensors: list[ReaderTensor] = [] - offs += self._push_field(ReaderField(offs, 'GGUF.version', [temp_version], [0], [GGUFValueType.UINT32])) - - # Check tensor count and kv count - temp_counts = self._get(offs, np.uint64, 2) - offs += self._push_field(ReaderField(offs, 'GGUF.tensor_count', [temp_counts[:1]], [0], [GGUFValueType.UINT64])) - offs += self._push_field(ReaderField(offs, 'GGUF.kv_count', [temp_counts[1:]], [0], [GGUFValueType.UINT64])) - tensor_count, kv_count = temp_counts - offs = self._build_fields(offs, kv_count) - - # Build Tensor Info Fields - offs, tensors_fields = self._build_tensor_info(offs, tensor_count) - new_align = self.fields.get('general.alignment') - if new_align is not None: - if new_align.types != [GGUFValueType.UINT32]: - raise ValueError('Bad type for general.alignment field') - self.alignment = new_align.parts[-1][0] - padding = offs % self.alignment - if padding != 0: - offs += self.alignment - padding - self.data_offset = offs - self._build_tensors(offs, tensors_fields) - - _DT = TypeVar('_DT', bound = npt.DTypeLike) - - # Fetch a key/value metadata field by key. - def get_field(self, key: str) -> Union[ReaderField, None]: - return self.fields.get(key, None) - - # Fetch a tensor from the list by index. - def get_tensor(self, idx: int) -> ReaderTensor: - return self.tensors[idx] - - def _get( - self, offset: int, dtype: npt.DTypeLike, count: int = 1, override_order: None | Literal['I', 'S', '<'] = None, - ) -> npt.NDArray[Any]: - count = int(count) - itemsize = int(np.empty([], dtype = dtype).itemsize) - end_offs = offset + itemsize * count - arr = self.data[offset:end_offs].view(dtype=dtype)[:count] - return arr.view(arr.dtype.newbyteorder(self.byte_order if override_order is None else override_order)) - - def _push_field(self, field: ReaderField, skip_sum: bool = False) -> int: - if field.name in self.fields: - # TODO: add option to generate error on duplicate keys - # raise KeyError(f'Duplicate {field.name} already in list at offset {field.offset}') - - logger.warning(f'Duplicate key {field.name} at offset {field.offset}') - self.fields[field.name + '_{}'.format(field.offset)] = field - else: - self.fields[field.name] = field - return 0 if skip_sum else sum(int(part.nbytes) for part in field.parts) - - def _get_str(self, offset: int) -> tuple[npt.NDArray[np.uint64], npt.NDArray[np.uint8]]: - slen = self._get(offset, np.uint64) - return slen, self._get(offset + 8, np.uint8, slen[0]) - - def _get_field_parts( - self, orig_offs: int, raw_type: int, - ) -> tuple[int, list[npt.NDArray[Any]], list[int], list[GGUFValueType]]: - offs = orig_offs - types: list[GGUFValueType] = [] - gtype = GGUFValueType(raw_type) - types.append(gtype) - # Handle strings. - if gtype == GGUFValueType.STRING: - sparts: list[npt.NDArray[Any]] = list(self._get_str(offs)) - size = sum(int(part.nbytes) for part in sparts) - return size, sparts, [1], types - # Check if it's a simple scalar type. - nptype = self.gguf_scalar_to_np.get(gtype) - if nptype is not None: - val = self._get(offs, nptype) - return int(val.nbytes), [val], [0], types - # Handle arrays. - if gtype == GGUFValueType.ARRAY: - raw_itype = self._get(offs, np.uint32) - offs += int(raw_itype.nbytes) - alen = self._get(offs, np.uint64) - offs += int(alen.nbytes) - aparts: list[npt.NDArray[Any]] = [raw_itype, alen] - data_idxs: list[int] = [] - # FIXME: Handle multi-dimensional arrays properly instead of flattening - for idx in range(alen[0]): - curr_size, curr_parts, curr_idxs, curr_types = self._get_field_parts(offs, raw_itype[0]) - if idx == 0: - types += curr_types - idxs_offs = len(aparts) - aparts += curr_parts - data_idxs += (idx + idxs_offs for idx in curr_idxs) - offs += curr_size - return offs - orig_offs, aparts, data_idxs, types - # We can't deal with this one. - raise ValueError(f'Unknown/unhandled field type {gtype}') - - def _get_tensor_info_field(self, orig_offs: int) -> ReaderField: - offs = orig_offs - - # Get Tensor Name - name_len, name_data = self._get_str(offs) - offs += int(name_len.nbytes + name_data.nbytes) - - # Get Tensor Dimensions Count - n_dims = self._get(offs, np.uint32) - offs += int(n_dims.nbytes) - - # Get Tensor Dimension Array - dims = self._get(offs, np.uint64, n_dims[0]) - offs += int(dims.nbytes) - - # Get Tensor Encoding Scheme Type - raw_dtype = self._get(offs, np.uint32) - offs += int(raw_dtype.nbytes) - - # Get Tensor Offset - offset_tensor = self._get(offs, np.uint64) - offs += int(offset_tensor.nbytes) - - return ReaderField( - orig_offs, - str(bytes(name_data), encoding = 'utf-8'), - [name_len, name_data, n_dims, dims, raw_dtype, offset_tensor], - [1, 3, 4, 5], - ) - - def _build_fields(self, offs: int, count: int) -> int: - for _ in range(count): - orig_offs = offs - kv_klen, kv_kdata = self._get_str(offs) - offs += int(kv_klen.nbytes + kv_kdata.nbytes) - raw_kv_type = self._get(offs, np.uint32) - offs += int(raw_kv_type.nbytes) - parts: list[npt.NDArray[Any]] = [kv_klen, kv_kdata, raw_kv_type] - idxs_offs = len(parts) - field_size, field_parts, field_idxs, field_types = self._get_field_parts(offs, raw_kv_type[0]) - parts += field_parts - self._push_field(ReaderField( - orig_offs, - str(bytes(kv_kdata), encoding = 'utf-8'), - parts, - [idx + idxs_offs for idx in field_idxs], - field_types, - ), skip_sum = True) - offs += field_size - return offs - - def _build_tensor_info(self, offs: int, count: int) -> tuple[int, list[ReaderField]]: - tensor_fields = [] - for _ in range(count): - field = self._get_tensor_info_field(offs) - offs += sum(int(part.nbytes) for part in field.parts) - tensor_fields.append(field) - return offs, tensor_fields - - def _build_tensors(self, start_offs: int, fields: list[ReaderField]) -> None: - tensors = [] - tensor_names = set() # keep track of name to prevent duplicated tensors - for field in fields: - _name_len, name_data, _n_dims, dims, raw_dtype, offset_tensor = field.parts - # check if there's any tensor having same name already in the list - tensor_name = str(bytes(name_data), encoding = 'utf-8') - if tensor_name in tensor_names: - raise ValueError(f'Found duplicated tensor with name {tensor_name}') - tensor_names.add(tensor_name) - ggml_type = GGMLQuantizationType(raw_dtype[0]) - n_elems = int(np.prod(dims)) - np_dims = tuple(reversed(dims.tolist())) - block_size, type_size = GGML_QUANT_SIZES[ggml_type] - n_bytes = n_elems * type_size // block_size - data_offs = int(start_offs + offset_tensor[0]) - item_type: npt.DTypeLike - if ggml_type == GGMLQuantizationType.F16: - item_count = n_elems - item_type = np.float16 - elif ggml_type == GGMLQuantizationType.F32: - item_count = n_elems - item_type = np.float32 - elif ggml_type == GGMLQuantizationType.F64: - item_count = n_elems - item_type = np.float64 - elif ggml_type == GGMLQuantizationType.I8: - item_count = n_elems - item_type = np.int8 - elif ggml_type == GGMLQuantizationType.I16: - item_count = n_elems - item_type = np.int16 - elif ggml_type == GGMLQuantizationType.I32: - item_count = n_elems - item_type = np.int32 - elif ggml_type == GGMLQuantizationType.I64: - item_count = n_elems - item_type = np.int64 - else: - item_count = n_bytes - item_type = np.uint8 - np_dims = quant_shape_to_byte_shape(np_dims, ggml_type) - tensors.append(ReaderTensor( - name = tensor_name, - tensor_type = ggml_type, - shape = dims, - n_elements = n_elems, - n_bytes = n_bytes, - data_offset = data_offs, - data = self._get(data_offs, item_type, item_count).reshape(np_dims), - field = field, - )) - self.tensors = tensors diff --git a/venv/lib/python3.10/site-packages/gguf/gguf_writer.py b/venv/lib/python3.10/site-packages/gguf/gguf_writer.py deleted file mode 100644 index 54ca0c33fd3368daa6a19a15023b792da88951d2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/gguf_writer.py +++ /dev/null @@ -1,1076 +0,0 @@ -from __future__ import annotations - -import logging -import os -import shutil -import struct -import tempfile -from dataclasses import dataclass -from enum import Enum, auto -from math import prod -from pathlib import Path -from io import BufferedWriter -from typing import IO, Any, Sequence, Mapping -from string import ascii_letters, digits - -import numpy as np - -from .constants import ( - GGUF_DEFAULT_ALIGNMENT, - GGUF_MAGIC, - GGUF_VERSION, - GGMLQuantizationType, - GGUFEndian, - GGUFValueType, - Keys, - RopeScalingType, - PoolingType, - TokenType, - ExpertGatingFuncType, -) - -from .quants import quant_shape_from_byte_shape - -logger = logging.getLogger(__name__) - - -SHARD_NAME_FORMAT = "{:s}-{:05d}-of-{:05d}.gguf" - - -@dataclass -class TensorInfo: - shape: Sequence[int] - dtype: GGMLQuantizationType - nbytes: int - tensor: np.ndarray[Any, Any] | None = None - - -@dataclass -class GGUFValue: - value: Any - type: GGUFValueType - sub_type: GGUFValueType | None = None - - -class WriterState(Enum): - NO_FILE = auto() - EMPTY = auto() - HEADER = auto() - KV_DATA = auto() - TI_DATA = auto() - WEIGHTS = auto() - - -class GGUFWriter: - fout: list[BufferedWriter] | None - path: Path | None - temp_file: tempfile.SpooledTemporaryFile[bytes] | None - tensors: list[dict[str, TensorInfo]] - kv_data: list[dict[str, GGUFValue]] - state: WriterState - _simple_value_packing = { - GGUFValueType.UINT8: "B", - GGUFValueType.INT8: "b", - GGUFValueType.UINT16: "H", - GGUFValueType.INT16: "h", - GGUFValueType.UINT32: "I", - GGUFValueType.INT32: "i", - GGUFValueType.FLOAT32: "f", - GGUFValueType.UINT64: "Q", - GGUFValueType.INT64: "q", - GGUFValueType.FLOAT64: "d", - GGUFValueType.BOOL: "?", - } - - def __init__( - self, path: os.PathLike[str] | str | None, arch: str, use_temp_file: bool = False, endianess: GGUFEndian = GGUFEndian.LITTLE, - split_max_tensors: int = 0, split_max_size: int = 0, dry_run: bool = False, small_first_shard: bool = False - ): - self.fout = None - self.path = Path(path) if path else None - self.arch = arch - self.endianess = endianess - self.data_alignment = GGUF_DEFAULT_ALIGNMENT - self.use_temp_file = use_temp_file - self.temp_file = None - self.tensors = [{}] - self.kv_data = [{}] - self.split_max_tensors = split_max_tensors - self.split_max_size = split_max_size - self.dry_run = dry_run - self.small_first_shard = small_first_shard - logger.info("gguf: This GGUF file is for {0} Endian only".format( - "Big" if self.endianess == GGUFEndian.BIG else "Little", - )) - self.state = WriterState.NO_FILE - - if self.small_first_shard: - self.tensors.append({}) - - self.add_architecture() - - def get_total_parameter_count(self) -> tuple[int, int, int, int]: - total_params = 0 - shared_params = 0 - expert_params = 0 - - expert_sum = 0 - n_expert_tensors = 0 - - last_lora_a: tuple[str, TensorInfo] | None = None - - for tensors in self.tensors: - for name, info in tensors.items(): - - shape = info.shape - - if name.endswith(".lora_a"): - last_lora_a = (name, info) - continue - elif name.endswith(".lora_b"): - if last_lora_a is None or last_lora_a[0] != name[:-1] + "a": - # Bail when the LoRA pair can't be found trivially - logger.warning("can't measure LoRA size correctly, tensor order is unusual") - return 0, 0, 0, 0 - else: - shape = (*shape[:-1], last_lora_a[1].shape[-1]) - - size = prod(shape) - - if "_exps." in name: - expert_params += (size // shape[-3]) - expert_sum += shape[-3] - n_expert_tensors += 1 - else: - shared_params += size - - total_params += size - - # Hopefully this should work even for variable-expert-count models - expert_count = (expert_sum // n_expert_tensors) if n_expert_tensors > 0 else 0 - - # Negate the total to signal it's likely not exact - if last_lora_a is not None: - total_params = -total_params - - # NOTE: keep the output in the same order as accepted by 'size_label' in gguf-py/gguf/utility.py - return total_params, shared_params, expert_params, expert_count - - def format_shard_names(self, path: Path) -> list[Path]: - if len(self.tensors) == 1: - return [path] - return [path.with_name(SHARD_NAME_FORMAT.format(path.stem, i + 1, len(self.tensors))) for i in range(len(self.tensors))] - - def open_output_file(self, path: Path | None = None) -> None: - if self.state is WriterState.EMPTY and self.fout is not None and (path is None or path == self.path): - # allow calling this multiple times as long as the path is the same - return - - if self.state is not WriterState.NO_FILE: - raise ValueError(f'Expected output file to be not yet opened, got {self.state}') - - if path is not None: - self.path = path - - if self.path is not None: - filenames = self.print_plan() - self.fout = [open(filename, "wb") for filename in filenames] - self.state = WriterState.EMPTY - - def print_plan(self) -> list[Path]: - logger.info("Writing the following files:") - assert self.path is not None - filenames = self.format_shard_names(self.path) - assert len(filenames) == len(self.tensors) - for name, tensors in zip(filenames, self.tensors): - logger.info(f"{name}: n_tensors = {len(tensors)}, total_size = {GGUFWriter.format_n_bytes_to_str(sum(ti.nbytes for ti in tensors.values()))}") - - if self.dry_run: - logger.info("Dry run, not writing files") - for name in filenames: - print(name) # noqa: NP100 - exit() - - return filenames - - def add_shard_kv_data(self) -> None: - if len(self.tensors) == 1: - return - - total_tensors = sum(len(t) for t in self.tensors) - assert self.fout is not None - total_splits = len(self.fout) - self.kv_data.extend({} for _ in range(len(self.kv_data), total_splits)) - for i, kv_data in enumerate(self.kv_data): - kv_data[Keys.Split.LLM_KV_SPLIT_NO] = GGUFValue(i, GGUFValueType.UINT16) - kv_data[Keys.Split.LLM_KV_SPLIT_COUNT] = GGUFValue(total_splits, GGUFValueType.UINT16) - kv_data[Keys.Split.LLM_KV_SPLIT_TENSORS_COUNT] = GGUFValue(total_tensors, GGUFValueType.INT32) - - def write_header_to_file(self, path: Path | None = None) -> None: - if len(self.tensors) == 1 and (self.split_max_tensors != 0 or self.split_max_size != 0): - logger.warning("Model fails split requirements, not splitting") - - self.open_output_file(path) - - if self.state is not WriterState.EMPTY: - raise ValueError(f'Expected output file to be empty, got {self.state}') - - assert self.fout is not None - assert len(self.fout) == len(self.tensors) - assert len(self.kv_data) == 1 - - self.add_shard_kv_data() - - for fout, tensors, kv_data in zip(self.fout, self.tensors, self.kv_data): - fout.write(self._pack(" None: - if self.state is not WriterState.HEADER: - raise ValueError(f'Expected output file to contain the header, got {self.state}') - assert self.fout is not None - - for fout, kv_data in zip(self.fout, self.kv_data): - kv_bytes = bytearray() - - for key, val in kv_data.items(): - kv_bytes += self._pack_val(key, GGUFValueType.STRING, add_vtype=False) - kv_bytes += self._pack_val(val.value, val.type, add_vtype=True, sub_type=val.sub_type) - - fout.write(kv_bytes) - - self.flush() - self.state = WriterState.KV_DATA - - def write_ti_data_to_file(self) -> None: - if self.state is not WriterState.KV_DATA: - raise ValueError(f'Expected output file to contain KV data, got {self.state}') - assert self.fout is not None - - for fout, tensors in zip(self.fout, self.tensors): - ti_data = bytearray() - offset_tensor = 0 - - for name, ti in tensors.items(): - ti_data += self._pack_val(name, GGUFValueType.STRING, add_vtype=False) - n_dims = len(ti.shape) - ti_data += self._pack("I", n_dims) - for j in range(n_dims): - ti_data += self._pack("Q", ti.shape[n_dims - 1 - j]) - ti_data += self._pack("I", ti.dtype) - ti_data += self._pack("Q", offset_tensor) - offset_tensor += GGUFWriter.ggml_pad(ti.nbytes, self.data_alignment) - - fout.write(ti_data) - fout.flush() - self.state = WriterState.TI_DATA - - def add_key_value(self, key: str, val: Any, vtype: GGUFValueType, sub_type: GGUFValueType | None = None) -> None: - if any(key in kv_data for kv_data in self.kv_data): - logger.warning(f'Duplicated key name {key!r}, overwriting it with new value {val!r} of type {vtype.name}') - - self.kv_data[0][key] = GGUFValue(value=val, type=vtype, sub_type=sub_type) - - def add_uint8(self, key: str, val: int) -> None: - self.add_key_value(key,val, GGUFValueType.UINT8) - - def add_int8(self, key: str, val: int) -> None: - self.add_key_value(key, val, GGUFValueType.INT8) - - def add_uint16(self, key: str, val: int) -> None: - self.add_key_value(key, val, GGUFValueType.UINT16) - - def add_int16(self, key: str, val: int) -> None: - self.add_key_value(key, val, GGUFValueType.INT16) - - def add_uint32(self, key: str, val: int) -> None: - self.add_key_value(key, val, GGUFValueType.UINT32) - - def add_int32(self, key: str, val: int) -> None: - self.add_key_value(key, val, GGUFValueType.INT32) - - def add_float32(self, key: str, val: float) -> None: - self.add_key_value(key, val, GGUFValueType.FLOAT32) - - def add_uint64(self, key: str, val: int) -> None: - self.add_key_value(key, val, GGUFValueType.UINT64) - - def add_int64(self, key: str, val: int) -> None: - self.add_key_value(key, val, GGUFValueType.INT64) - - def add_float64(self, key: str, val: float) -> None: - self.add_key_value(key, val, GGUFValueType.FLOAT64) - - def add_bool(self, key: str, val: bool) -> None: - self.add_key_value(key, val, GGUFValueType.BOOL) - - def add_string(self, key: str, val: str) -> None: - if not val: - return - self.add_key_value(key, val, GGUFValueType.STRING) - - def add_array(self, key: str, val: Sequence[Any]) -> None: - if len(val) == 0: - return - self.add_key_value(key, val, GGUFValueType.ARRAY) - - @staticmethod - def ggml_pad(x: int, n: int) -> int: - return ((x + n - 1) // n) * n - - def add_tensor_info( - self, name: str, tensor_shape: Sequence[int], tensor_dtype: np.dtype, - tensor_nbytes: int, raw_dtype: GGMLQuantizationType | None = None, - ) -> None: - if self.state is not WriterState.NO_FILE: - raise ValueError(f'Expected output file to be not yet opened, got {self.state}') - - if any(name in tensors for tensors in self.tensors): - raise ValueError(f'Duplicated tensor name {name!r}') - - if raw_dtype is None: - if tensor_dtype == np.float16: - dtype = GGMLQuantizationType.F16 - elif tensor_dtype == np.float32: - dtype = GGMLQuantizationType.F32 - elif tensor_dtype == np.float64: - dtype = GGMLQuantizationType.F64 - elif tensor_dtype == np.int8: - dtype = GGMLQuantizationType.I8 - elif tensor_dtype == np.int16: - dtype = GGMLQuantizationType.I16 - elif tensor_dtype == np.int32: - dtype = GGMLQuantizationType.I32 - elif tensor_dtype == np.int64: - dtype = GGMLQuantizationType.I64 - else: - raise ValueError("Only F16, F32, F64, I8, I16, I32, I64 tensors are supported for now") - else: - dtype = raw_dtype - if tensor_dtype == np.uint8: - tensor_shape = quant_shape_from_byte_shape(tensor_shape, raw_dtype) - - # make sure there is at least one tensor before splitting - if len(self.tensors[-1]) > 0: - if ( # split when over tensor limit - self.split_max_tensors != 0 - and len(self.tensors[-1]) >= self.split_max_tensors - ) or ( # split when over size limit - self.split_max_size != 0 - and sum(ti.nbytes for ti in self.tensors[-1].values()) + tensor_nbytes > self.split_max_size - ): - self.tensors.append({}) - - self.tensors[-1][name] = TensorInfo(shape=tensor_shape, dtype=dtype, nbytes=tensor_nbytes) - - def add_tensor( - self, name: str, tensor: np.ndarray[Any, Any], raw_shape: Sequence[int] | None = None, - raw_dtype: GGMLQuantizationType | None = None, - ) -> None: - if self.endianess == GGUFEndian.BIG: - tensor.byteswap(inplace=True) - if self.use_temp_file and self.temp_file is None: - fp = tempfile.SpooledTemporaryFile(mode="w+b", max_size=256 * 1024 * 1024) - fp.seek(0) - self.temp_file = fp - - shape: Sequence[int] = raw_shape if raw_shape is not None else tensor.shape - self.add_tensor_info(name, shape, tensor.dtype, tensor.nbytes, raw_dtype=raw_dtype) - - if self.temp_file is None: - self.tensors[-1][name].tensor = tensor - return - - tensor.tofile(self.temp_file) - self.write_padding(self.temp_file, tensor.nbytes) - - def write_padding(self, fp: IO[bytes], n: int, align: int | None = None) -> None: - pad = GGUFWriter.ggml_pad(n, align if align is not None else self.data_alignment) - n - if pad != 0: - fp.write(bytes([0] * pad)) - - def write_tensor_data(self, tensor: np.ndarray[Any, Any]) -> None: - if self.state is not WriterState.TI_DATA and self.state is not WriterState.WEIGHTS: - raise ValueError(f'Expected output file to contain tensor info or weights, got {self.state}') - assert self.fout is not None - - if self.endianess == GGUFEndian.BIG: - tensor.byteswap(inplace=True) - - file_id = -1 - for i, tensors in enumerate(self.tensors): - if len(tensors) > 0: - file_id = i - break - - fout = self.fout[file_id] - - # pop the first tensor info - # TODO: cleaner way to get the first key - first_tensor_name = [name for name, _ in zip(self.tensors[file_id].keys(), range(1))][0] - ti = self.tensors[file_id].pop(first_tensor_name) - assert ti.nbytes == tensor.nbytes - - self.write_padding(fout, fout.tell()) - tensor.tofile(fout) - self.write_padding(fout, tensor.nbytes) - - self.state = WriterState.WEIGHTS - - def write_tensors_to_file(self, *, progress: bool = False) -> None: - self.write_ti_data_to_file() - - assert self.fout is not None - - for fout in self.fout: - self.write_padding(fout, fout.tell()) - - if self.temp_file is None: - shard_bar = None - bar = None - - if progress: - from tqdm import tqdm - - total_bytes = sum(ti.nbytes for t in self.tensors for ti in t.values()) - - if len(self.fout) > 1: - shard_bar = tqdm(desc=f"Shard (0/{len(self.fout)})", total=None, unit="byte", unit_scale=True) - bar = tqdm(desc="Writing", total=total_bytes, unit="byte", unit_scale=True) - - for i, (fout, tensors) in enumerate(zip(self.fout, self.tensors)): - if shard_bar is not None: - shard_bar.set_description(f"Shard ({i + 1}/{len(self.fout)})") - total = sum(ti.nbytes for ti in tensors.values()) - shard_bar.reset(total=(total if total > 0 else None)) - - # relying on the fact that Python dicts preserve insertion order (since 3.7) - for ti in tensors.values(): - assert ti.tensor is not None # can only iterate once over the tensors - assert ti.tensor.nbytes == ti.nbytes - ti.tensor.tofile(fout) - if shard_bar is not None: - shard_bar.update(ti.nbytes) - if bar is not None: - bar.update(ti.nbytes) - self.write_padding(fout, ti.nbytes) - ti.tensor = None - else: - self.temp_file.seek(0) - - shutil.copyfileobj(self.temp_file, self.fout[0 if not self.small_first_shard else 1]) - self.flush() - self.temp_file.close() - - self.state = WriterState.WEIGHTS - - def flush(self) -> None: - assert self.fout is not None - for fout in self.fout: - fout.flush() - - def close(self) -> None: - if self.fout is not None: - for fout in self.fout: - fout.close() - self.fout = None - - def add_type(self, type_name: str) -> None: - self.add_string(Keys.General.TYPE, type_name) - - def add_architecture(self) -> None: - self.add_string(Keys.General.ARCHITECTURE, self.arch) - - def add_quantization_version(self, quantization_version: int) -> None: - self.add_uint32(Keys.General.QUANTIZATION_VERSION, quantization_version) - - def add_custom_alignment(self, alignment: int) -> None: - self.data_alignment = alignment - self.add_uint32(Keys.General.ALIGNMENT, alignment) - - def add_file_type(self, ftype: int) -> None: - self.add_uint32(Keys.General.FILE_TYPE, ftype) - - def add_name(self, name: str) -> None: - self.add_string(Keys.General.NAME, name) - - def add_author(self, author: str) -> None: - self.add_string(Keys.General.AUTHOR, author) - - def add_version(self, version: str) -> None: - self.add_string(Keys.General.VERSION, version) - - def add_organization(self, organization: str) -> None: - self.add_string(Keys.General.ORGANIZATION, organization) - - def add_finetune(self, finetune: str) -> None: - self.add_string(Keys.General.FINETUNE, finetune) - - def add_basename(self, basename: str) -> None: - self.add_string(Keys.General.BASENAME, basename) - - def add_description(self, description: str) -> None: - self.add_string(Keys.General.DESCRIPTION, description) - - def add_quantized_by(self, quantized: str) -> None: - self.add_string(Keys.General.QUANTIZED_BY, quantized) - - def add_size_label(self, size_label: str) -> None: - self.add_string(Keys.General.SIZE_LABEL, size_label) - - def add_license(self, license: str) -> None: - self.add_string(Keys.General.LICENSE, license) - - def add_license_name(self, license: str) -> None: - self.add_string(Keys.General.LICENSE_NAME, license) - - def add_license_link(self, license: str) -> None: - self.add_string(Keys.General.LICENSE_LINK, license) - - def add_url(self, url: str) -> None: - self.add_string(Keys.General.URL, url) - - def add_doi(self, doi: str) -> None: - self.add_string(Keys.General.DOI, doi) - - def add_uuid(self, uuid: str) -> None: - self.add_string(Keys.General.UUID, uuid) - - def add_repo_url(self, repo_url: str) -> None: - self.add_string(Keys.General.REPO_URL, repo_url) - - def add_source_url(self, url: str) -> None: - self.add_string(Keys.General.SOURCE_URL, url) - - def add_source_doi(self, doi: str) -> None: - self.add_string(Keys.General.SOURCE_DOI, doi) - - def add_source_uuid(self, uuid: str) -> None: - self.add_string(Keys.General.SOURCE_UUID, uuid) - - def add_source_repo_url(self, repo_url: str) -> None: - self.add_string(Keys.General.SOURCE_REPO_URL, repo_url) - - def add_base_model_count(self, source_count: int) -> None: - self.add_uint32(Keys.General.BASE_MODEL_COUNT, source_count) - - def add_base_model_name(self, source_id: int, name: str) -> None: - self.add_string(Keys.General.BASE_MODEL_NAME.format(id=source_id), name) - - def add_base_model_author(self, source_id: int, author: str) -> None: - self.add_string(Keys.General.BASE_MODEL_AUTHOR.format(id=source_id), author) - - def add_base_model_version(self, source_id: int, version: str) -> None: - self.add_string(Keys.General.BASE_MODEL_VERSION.format(id=source_id), version) - - def add_base_model_organization(self, source_id: int, organization: str) -> None: - self.add_string(Keys.General.BASE_MODEL_ORGANIZATION.format(id=source_id), organization) - - def add_base_model_description(self, source_id: int, description: str) -> None: - self.add_string(Keys.General.BASE_MODEL_DESCRIPTION.format(id=source_id), description) - - def add_base_model_url(self, source_id: int, url: str) -> None: - self.add_string(Keys.General.BASE_MODEL_URL.format(id=source_id), url) - - def add_base_model_doi(self, source_id: int, doi: str) -> None: - self.add_string(Keys.General.BASE_MODEL_DOI.format(id=source_id), doi) - - def add_base_model_uuid(self, source_id: int, uuid: str) -> None: - self.add_string(Keys.General.BASE_MODEL_UUID.format(id=source_id), uuid) - - def add_base_model_repo_url(self, source_id: int, repo_url: str) -> None: - self.add_string(Keys.General.BASE_MODEL_REPO_URL.format(id=source_id), repo_url) - - def add_dataset_count(self, source_count: int) -> None: - self.add_uint32(Keys.General.DATASET_COUNT, source_count) - - def add_dataset_name(self, source_id: int, name: str) -> None: - self.add_string(Keys.General.DATASET_NAME.format(id=source_id), name) - - def add_dataset_author(self, source_id: int, author: str) -> None: - self.add_string(Keys.General.DATASET_AUTHOR.format(id=source_id), author) - - def add_dataset_version(self, source_id: int, version: str) -> None: - self.add_string(Keys.General.DATASET_VERSION.format(id=source_id), version) - - def add_dataset_organization(self, source_id: int, organization: str) -> None: - self.add_string(Keys.General.DATASET_ORGANIZATION.format(id=source_id), organization) - - def add_dataset_description(self, source_id: int, description: str) -> None: - self.add_string(Keys.General.DATASET_DESCRIPTION.format(id=source_id), description) - - def add_dataset_url(self, source_id: int, url: str) -> None: - self.add_string(Keys.General.DATASET_URL.format(id=source_id), url) - - def add_dataset_doi(self, source_id: int, doi: str) -> None: - self.add_string(Keys.General.DATASET_DOI.format(id=source_id), doi) - - def add_dataset_uuid(self, source_id: int, uuid: str) -> None: - self.add_string(Keys.General.DATASET_UUID.format(id=source_id), uuid) - - def add_dataset_repo_url(self, source_id: int, repo_url: str) -> None: - self.add_string(Keys.General.DATASET_REPO_URL.format(id=source_id), repo_url) - - def add_tags(self, tags: Sequence[str]) -> None: - self.add_array(Keys.General.TAGS, tags) - - def add_languages(self, languages: Sequence[str]) -> None: - self.add_array(Keys.General.LANGUAGES, languages) - - def add_tensor_data_layout(self, layout: str) -> None: - self.add_string(Keys.LLM.TENSOR_DATA_LAYOUT.format(arch=self.arch), layout) - - def add_vocab_size(self, size: int) -> None: - self.add_uint32(Keys.LLM.VOCAB_SIZE.format(arch=self.arch), size) - - def add_context_length(self, length: int) -> None: - self.add_uint32(Keys.LLM.CONTEXT_LENGTH.format(arch=self.arch), length) - - def add_embedding_length(self, length: int) -> None: - self.add_uint32(Keys.LLM.EMBEDDING_LENGTH.format(arch=self.arch), length) - - def add_features_length(self, length: int) -> None: - self.add_uint32(Keys.LLM.FEATURES_LENGTH.format(arch=self.arch), length) - - def add_posnet_embedding_length(self, length: int) -> None: - self.add_uint32(Keys.PosNet.EMBEDDING_LENGTH.format(arch=self.arch), length) - - def add_posnet_block_count(self, length: int) -> None: - self.add_uint32(Keys.PosNet.BLOCK_COUNT.format(arch=self.arch), length) - - def add_convnext_embedding_length(self, length: int) -> None: - self.add_uint32(Keys.ConvNext.EMBEDDING_LENGTH.format(arch=self.arch), length) - - def add_convnext_block_count(self, length: int) -> None: - self.add_uint32(Keys.ConvNext.BLOCK_COUNT.format(arch=self.arch), length) - - def add_block_count(self, length: int) -> None: - self.add_uint32(Keys.LLM.BLOCK_COUNT.format(arch=self.arch), length) - - def add_leading_dense_block_count(self, length: int) -> None: - self.add_uint32(Keys.LLM.LEADING_DENSE_BLOCK_COUNT.format(arch=self.arch), length) - - def add_feed_forward_length(self, length: int | Sequence[int]) -> None: - if isinstance(length, int): - self.add_uint32(Keys.LLM.FEED_FORWARD_LENGTH.format(arch=self.arch), length) - else: - self.add_array(Keys.LLM.FEED_FORWARD_LENGTH.format(arch=self.arch), length) - - def add_expert_feed_forward_length(self, length: int) -> None: - self.add_uint32(Keys.LLM.EXPERT_FEED_FORWARD_LENGTH.format(arch=self.arch), length) - - def add_expert_shared_feed_forward_length(self, length: int) -> None: - self.add_uint32(Keys.LLM.EXPERT_SHARED_FEED_FORWARD_LENGTH.format(arch=self.arch), length) - - def add_parallel_residual(self, use: bool) -> None: - self.add_bool(Keys.LLM.USE_PARALLEL_RESIDUAL.format(arch=self.arch), use) - - def add_decoder_start_token_id(self, id: int) -> None: - self.add_uint32(Keys.LLM.DECODER_START_TOKEN_ID.format(arch=self.arch), id) - - def add_head_count(self, count: int | Sequence[int]) -> None: - if isinstance(count, int): - self.add_uint32(Keys.Attention.HEAD_COUNT.format(arch=self.arch), count) - else: - self.add_array(Keys.Attention.HEAD_COUNT.format(arch=self.arch), count) - - def add_head_count_kv(self, count: int | Sequence[int]) -> None: - if isinstance(count, int): - self.add_uint32(Keys.Attention.HEAD_COUNT_KV.format(arch=self.arch), count) - else: - self.add_array(Keys.Attention.HEAD_COUNT_KV.format(arch=self.arch), count) - - def add_key_length(self, length: int) -> None: - self.add_uint32(Keys.Attention.KEY_LENGTH.format(arch=self.arch), length) - - def add_value_length(self, length: int) -> None: - self.add_uint32(Keys.Attention.VALUE_LENGTH.format(arch=self.arch), length) - - def add_key_length_mla(self, length: int) -> None: - self.add_uint32(Keys.Attention.KEY_LENGTH_MLA.format(arch=self.arch), length) - - def add_value_length_mla(self, length: int) -> None: - self.add_uint32(Keys.Attention.VALUE_LENGTH_MLA.format(arch=self.arch), length) - - def add_max_alibi_bias(self, bias: float) -> None: - self.add_float32(Keys.Attention.MAX_ALIBI_BIAS.format(arch=self.arch), bias) - - def add_clamp_kqv(self, value: float) -> None: - self.add_float32(Keys.Attention.CLAMP_KQV.format(arch=self.arch), value) - - def add_logit_scale(self, value: float) -> None: - self.add_float32(Keys.LLM.LOGIT_SCALE.format(arch=self.arch), value) - - def add_attn_logit_softcapping(self, value: float) -> None: - self.add_float32(Keys.LLM.ATTN_LOGIT_SOFTCAPPING.format(arch=self.arch), value) - - def add_final_logit_softcapping(self, value: float) -> None: - self.add_float32(Keys.LLM.FINAL_LOGIT_SOFTCAPPING.format(arch=self.arch), value) - - def add_expert_count(self, count: int) -> None: - self.add_uint32(Keys.LLM.EXPERT_COUNT.format(arch=self.arch), count) - - def add_expert_used_count(self, count: int) -> None: - self.add_uint32(Keys.LLM.EXPERT_USED_COUNT.format(arch=self.arch), count) - - def add_expert_shared_count(self, count: int) -> None: - self.add_uint32(Keys.LLM.EXPERT_SHARED_COUNT.format(arch=self.arch), count) - - def add_expert_weights_scale(self, value: float) -> None: - self.add_float32(Keys.LLM.EXPERT_WEIGHTS_SCALE.format(arch=self.arch), value) - - def add_expert_weights_norm(self, value: bool) -> None: - self.add_bool(Keys.LLM.EXPERT_WEIGHTS_NORM.format(arch=self.arch), value) - - def add_expert_gating_func(self, value: ExpertGatingFuncType) -> None: - self.add_uint32(Keys.LLM.EXPERT_GATING_FUNC.format(arch=self.arch), value.value) - - def add_moe_every_n_layers(self, value: int) -> None: - self.add_uint32(Keys.LLM.MOE_EVERY_N_LAYERS.format(arch=self.arch), value) - - def add_swin_norm(self, value: bool) -> None: - self.add_bool(Keys.LLM.SWIN_NORM.format(arch=self.arch), value) - - def add_rescale_every_n_layers(self, count: int) -> None: - self.add_uint32(Keys.LLM.RESCALE_EVERY_N_LAYERS.format(arch=self.arch), count) - - def add_time_mix_extra_dim(self, dim: int) -> None: - self.add_uint32(Keys.LLM.TIME_MIX_EXTRA_DIM.format(arch=self.arch), dim) - - def add_time_decay_extra_dim(self, dim: int) -> None: - self.add_uint32(Keys.LLM.TIME_DECAY_EXTRA_DIM.format(arch=self.arch), dim) - - def add_residual_scale(self, value: float) -> None: - self.add_float32(Keys.LLM.RESIDUAL_SCALE.format(arch=self.arch), value) - - def add_embedding_scale(self, value: float) -> None: - self.add_float32(Keys.LLM.EMBEDDING_SCALE.format(arch=self.arch), value) - - def add_wkv_head_size(self, size: int) -> None: - self.add_uint32(Keys.WKV.HEAD_SIZE.format(arch=self.arch), size) - - def add_token_shift_count(self, count: int) -> None: - self.add_uint32(Keys.LLM.TOKEN_SHIFT_COUNT.format(arch=self.arch), count) - - def add_interleave_moe_layer_step(self, value: int) -> None: - self.add_uint32(Keys.LLM.INTERLEAVE_MOE_LAYER_STEP.format(arch=self.arch), value) - - def add_layer_norm_eps(self, value: float) -> None: - self.add_float32(Keys.Attention.LAYERNORM_EPS.format(arch=self.arch), value) - - def add_layer_norm_rms_eps(self, value: float) -> None: - self.add_float32(Keys.Attention.LAYERNORM_RMS_EPS.format(arch=self.arch), value) - - def add_group_norm_eps(self, value: float) -> None: - self.add_float32(Keys.Attention.GROUPNORM_EPS.format(arch=self.arch), value) - - def add_group_norm_groups(self, value: int) -> None: - self.add_uint32(Keys.Attention.GROUPNORM_GROUPS.format(arch=self.arch), value) - - def add_causal_attention(self, value: bool) -> None: - self.add_bool(Keys.Attention.CAUSAL.format(arch=self.arch), value) - - def add_q_lora_rank(self, length: int) -> None: - self.add_uint32(Keys.Attention.Q_LORA_RANK.format(arch=self.arch), length) - - def add_kv_lora_rank(self, length: int) -> None: - self.add_uint32(Keys.Attention.KV_LORA_RANK.format(arch=self.arch), length) - - def add_decay_lora_rank(self, length: int) -> None: - self.add_uint32(Keys.Attention.DECAY_LORA_RANK.format(arch=self.arch), length) - - def add_iclr_lora_rank(self, length: int) -> None: - self.add_uint32(Keys.Attention.ICLR_LORA_RANK.format(arch=self.arch), length) - - def add_value_residual_mix_lora_rank(self, length: int) -> None: - self.add_uint32(Keys.Attention.VALUE_RESIDUAL_MIX_LORA_RANK.format(arch=self.arch), length) - - def add_gate_lora_rank(self, length: int) -> None: - self.add_uint32(Keys.Attention.GATE_LORA_RANK.format(arch=self.arch), length) - - def add_relative_attn_buckets_count(self, value: int) -> None: - self.add_uint32(Keys.Attention.REL_BUCKETS_COUNT.format(arch=self.arch), value) - - def add_sliding_window(self, value: int) -> None: - self.add_uint32(Keys.Attention.SLIDING_WINDOW.format(arch=self.arch), value) - - def add_attention_scale(self, value: float) -> None: - self.add_float32(Keys.Attention.SCALE.format(arch=self.arch), value) - - def add_pooling_type(self, value: PoolingType) -> None: - self.add_uint32(Keys.LLM.POOLING_TYPE.format(arch=self.arch), value.value) - - def add_rope_dimension_count(self, count: int) -> None: - self.add_uint32(Keys.Rope.DIMENSION_COUNT.format(arch=self.arch), count) - - def add_rope_dimension_sections(self, dims: Sequence[int]) -> None: - self.add_array(Keys.Rope.DIMENSION_SECTIONS.format(arch=self.arch), dims) - - def add_rope_freq_base(self, value: float) -> None: - self.add_float32(Keys.Rope.FREQ_BASE.format(arch=self.arch), value) - - def add_rope_scaling_type(self, value: RopeScalingType) -> None: - self.add_string(Keys.Rope.SCALING_TYPE.format(arch=self.arch), value.value) - - def add_rope_scaling_factor(self, value: float) -> None: - self.add_float32(Keys.Rope.SCALING_FACTOR.format(arch=self.arch), value) - - def add_rope_scaling_attn_factors(self, value: float) -> None: - self.add_float32(Keys.Rope.SCALING_ATTN_FACTOR.format(arch=self.arch), value) - - def add_rope_scaling_orig_ctx_len(self, value: int) -> None: - self.add_uint32(Keys.Rope.SCALING_ORIG_CTX_LEN.format(arch=self.arch), value) - - def add_rope_scaling_finetuned(self, value: bool) -> None: - self.add_bool(Keys.Rope.SCALING_FINETUNED.format(arch=self.arch), value) - - def add_rope_scaling_yarn_log_mul(self, value: float) -> None: - self.add_float32(Keys.Rope.SCALING_YARN_LOG_MUL.format(arch=self.arch), value) - - def add_ssm_conv_kernel(self, value: int) -> None: - self.add_uint32(Keys.SSM.CONV_KERNEL.format(arch=self.arch), value) - - def add_ssm_inner_size(self, value: int) -> None: - self.add_uint32(Keys.SSM.INNER_SIZE.format(arch=self.arch), value) - - def add_ssm_state_size(self, value: int) -> None: - self.add_uint32(Keys.SSM.STATE_SIZE.format(arch=self.arch), value) - - def add_ssm_time_step_rank(self, value: int) -> None: - self.add_uint32(Keys.SSM.TIME_STEP_RANK.format(arch=self.arch), value) - - def add_ssm_dt_b_c_rms(self, value: bool) -> None: - self.add_bool(Keys.SSM.DT_B_C_RMS.format(arch=self.arch), value) - - def add_tokenizer_model(self, model: str) -> None: - self.add_string(Keys.Tokenizer.MODEL, model) - - def add_tokenizer_pre(self, pre: str) -> None: - self.add_string(Keys.Tokenizer.PRE, pre) - - def add_token_list(self, tokens: Sequence[str] | Sequence[bytes] | Sequence[bytearray]) -> None: - self.add_array(Keys.Tokenizer.LIST, tokens) - - def add_token_merges(self, merges: Sequence[str] | Sequence[bytes] | Sequence[bytearray]) -> None: - self.add_array(Keys.Tokenizer.MERGES, merges) - - def add_token_types(self, types: Sequence[TokenType] | Sequence[int]) -> None: - self.add_array(Keys.Tokenizer.TOKEN_TYPE, types) - - def add_token_type_count(self, value: int) -> None: - self.add_uint32(Keys.Tokenizer.TOKEN_TYPE_COUNT, value) - - def add_token_scores(self, scores: Sequence[float]) -> None: - self.add_array(Keys.Tokenizer.SCORES, scores) - - def add_bos_token_id(self, id: int) -> None: - self.add_uint32(Keys.Tokenizer.BOS_ID, id) - - def add_eos_token_id(self, id: int) -> None: - self.add_uint32(Keys.Tokenizer.EOS_ID, id) - - def add_unk_token_id(self, id: int) -> None: - self.add_uint32(Keys.Tokenizer.UNK_ID, id) - - def add_sep_token_id(self, id: int) -> None: - self.add_uint32(Keys.Tokenizer.SEP_ID, id) - - def add_pad_token_id(self, id: int) -> None: - self.add_uint32(Keys.Tokenizer.PAD_ID, id) - - def add_mask_token_id(self, id: int) -> None: - self.add_uint32(Keys.Tokenizer.MASK_ID, id) - - def add_add_bos_token(self, value: bool) -> None: - self.add_bool(Keys.Tokenizer.ADD_BOS, value) - - def add_add_eos_token(self, value: bool) -> None: - self.add_bool(Keys.Tokenizer.ADD_EOS, value) - - def add_add_space_prefix(self, value: bool) -> None: - self.add_bool(Keys.Tokenizer.ADD_PREFIX, value) - - def add_remove_extra_whitespaces(self, value: bool) -> None: - self.add_bool(Keys.Tokenizer.REMOVE_EXTRA_WS, value) - - def add_precompiled_charsmap(self, charsmap: bytes) -> None: - self.add_array(Keys.Tokenizer.PRECOMPILED_CHARSMAP, charsmap) - - def add_chat_template(self, value: str | Sequence[Mapping[str, str]]) -> None: - if not isinstance(value, str): - template_default = None - template_names = set() - - for choice in value: - name = choice.get('name', '') - template = choice.get('template') - - # Allowing non-alphanumerical characters in template name is probably not a good idea, so filter it - name = ''.join((c if c in ascii_letters + digits else '_' for c in name)) - - if name and template is not None: - if name == 'default': - template_default = template - else: - template_names.add(name) - self.add_string(Keys.Tokenizer.CHAT_TEMPLATE_N.format(name=name), template) - - if template_names: - self.add_array(Keys.Tokenizer.CHAT_TEMPLATES, list(template_names)) - - if template_default is None: - return - - value = template_default - - self.add_string(Keys.Tokenizer.CHAT_TEMPLATE, value) - - def add_eot_token_id(self, id: int) -> None: - self.add_uint32(Keys.Tokenizer.EOT_ID, id) - - def add_eom_token_id(self, id: int) -> None: - self.add_uint32(Keys.Tokenizer.EOM_ID, id) - - def add_classifier_output_labels(self, labels: Sequence[str]) -> None: - self.add_array(Keys.Classifier.OUTPUT_LABELS.format(arch=self.arch), labels) - - # for vision models - - def add_clip_has_vision_encoder(self, value: bool) -> None: - self.add_bool(Keys.Clip.HAS_VISION_ENCODER, value) - - def add_clip_has_audio_encoder(self, value: bool) -> None: - self.add_bool(Keys.Clip.HAS_AUDIO_ENCODER, value) - - def add_clip_projector_type(self, value: str) -> None: - self.add_string(Keys.Clip.PROJECTOR_TYPE, value) - - def add_vision_projection_dim(self, value: int) -> None: - self.add_uint32(Keys.ClipVision.PROJECTION_DIM, value) - - def add_vision_patch_size(self, value: int) -> None: - self.add_uint32(Keys.ClipVision.PATCH_SIZE, value) - - def add_vision_embedding_length(self, value: int) -> None: - self.add_uint32(Keys.ClipVision.EMBEDDING_LENGTH, value) - - def add_vision_feed_forward_length(self, value: int) -> None: - self.add_uint32(Keys.ClipVision.FEED_FORWARD_LENGTH, value) - - def add_vision_block_count(self, value: int) -> None: - self.add_uint32(Keys.ClipVision.BLOCK_COUNT, value) - - def add_vision_head_count(self, value: int) -> None: - self.add_uint32(Keys.ClipVision.Attention.HEAD_COUNT, value) - - def add_vision_attention_layernorm_eps(self, value: float) -> None: - self.add_float32(Keys.ClipVision.Attention.LAYERNORM_EPS, value) - - def add_vision_image_size(self, value: int) -> None: - self.add_uint32(Keys.ClipVision.IMAGE_SIZE, value) - - def add_vision_image_mean(self, values: Sequence[float]) -> None: - self.add_array(Keys.ClipVision.IMAGE_MEAN, values) - - def add_vision_image_std(self, values: Sequence[float]) -> None: - self.add_array(Keys.ClipVision.IMAGE_STD, values) - - def add_vision_spatial_merge_size(self, value: int) -> None: - self.add_uint32(Keys.ClipVision.SPATIAL_MERGE_SIZE, value) - - def add_vision_use_gelu(self, value: bool) -> None: - self.add_bool(Keys.ClipVision.USE_GELU, value) - - def add_vision_use_silu(self, value: bool) -> None: - self.add_bool(Keys.ClipVision.USE_SILU, value) - - def add_vision_projector_scale_factor(self, value: int) -> None: - self.add_uint32(Keys.ClipVision.Projector.SCALE_FACTOR, value) - - def add_vision_n_wa_pattern(self, value: int) -> None: - self.add_uint32(Keys.ClipVision.N_WA_PATTERN, value) - - # audio models - - def add_audio_projection_dim(self, value: int) -> None: - self.add_uint32(Keys.ClipAudio.PROJECTION_DIM, value) - - def add_audio_embedding_length(self, value: int) -> None: - self.add_uint32(Keys.ClipAudio.EMBEDDING_LENGTH, value) - - def add_audio_feed_forward_length(self, value: int) -> None: - self.add_uint32(Keys.ClipAudio.FEED_FORWARD_LENGTH, value) - - def add_audio_block_count(self, value: int) -> None: - self.add_uint32(Keys.ClipAudio.BLOCK_COUNT, value) - - def add_audio_head_count(self, value: int) -> None: - self.add_uint32(Keys.ClipAudio.Attention.HEAD_COUNT, value) - - def add_audio_attention_layernorm_eps(self, value: float) -> None: - self.add_float32(Keys.ClipAudio.Attention.LAYERNORM_EPS, value) - - def add_audio_num_mel_bins(self, value: int) -> None: - self.add_uint32(Keys.ClipAudio.NUM_MEL_BINS, value) - - def add_audio_stack_factor(self, value: int) -> None: - self.add_uint32(Keys.ClipAudio.Projector.STACK_FACTOR, value) - - def _pack(self, fmt: str, value: Any, skip_pack_prefix: bool = False) -> bytes: - pack_prefix = '' - if not skip_pack_prefix: - pack_prefix = '<' if self.endianess == GGUFEndian.LITTLE else '>' - return struct.pack(f'{pack_prefix}{fmt}', value) - - def _pack_val(self, val: Any, vtype: GGUFValueType, add_vtype: bool, sub_type: GGUFValueType | None = None) -> bytes: - kv_data = bytearray() - - if add_vtype: - kv_data += self._pack("I", vtype) - - pack_fmt = self._simple_value_packing.get(vtype) - if pack_fmt is not None: - kv_data += self._pack(pack_fmt, val, skip_pack_prefix = vtype == GGUFValueType.BOOL) - elif vtype == GGUFValueType.STRING: - encoded_val = val.encode("utf-8") if isinstance(val, str) else val - kv_data += self._pack("Q", len(encoded_val)) - kv_data += encoded_val - elif vtype == GGUFValueType.ARRAY: - - if not isinstance(val, Sequence): - raise ValueError("Invalid GGUF metadata array, expecting sequence") - - if len(val) == 0: - raise ValueError("Invalid GGUF metadata array. Empty array") - - if sub_type is not None: - ltype = sub_type - elif isinstance(val, bytes): - ltype = GGUFValueType.UINT8 - else: - ltype = GGUFValueType.get_type(val[0]) - if not all(GGUFValueType.get_type(i) is ltype for i in val[1:]): - raise ValueError("All items in a GGUF array should be of the same type") - kv_data += self._pack("I", ltype) - kv_data += self._pack("Q", len(val)) - for item in val: - kv_data += self._pack_val(item, ltype, add_vtype=False) - else: - raise ValueError("Invalid GGUF metadata value type or value") - - return kv_data - - @staticmethod - def format_n_bytes_to_str(num: int) -> str: - if num == 0: - return "negligible - metadata only" - fnum = float(num) - for unit in ("", "K", "M", "G"): - if abs(fnum) < 1000.0: - return f"{fnum:3.1f}{unit}" - fnum /= 1000.0 - return f"{fnum:.1f}T - over 1TB, split recommended" diff --git a/venv/lib/python3.10/site-packages/gguf/lazy.py b/venv/lib/python3.10/site-packages/gguf/lazy.py deleted file mode 100644 index f9bcadae0224bac5f922ce6565dd1e3883a923e1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/lazy.py +++ /dev/null @@ -1,223 +0,0 @@ -from __future__ import annotations -from abc import ABC, ABCMeta, abstractmethod - -import logging -from typing import Any, Callable - -import numpy as np -from numpy.typing import DTypeLike - - -logger = logging.getLogger(__name__) - - -class LazyMeta(ABCMeta): - - def __new__(cls, name: str, bases: tuple[type, ...], namespace: dict[str, Any], **kwargs): - def __getattr__(self, name: str) -> Any: - meta_attr = getattr(self._meta, name) - if callable(meta_attr): - return type(self)._wrap_fn( - (lambda s, *args, **kwargs: getattr(s, name)(*args, **kwargs)), - use_self=self, - ) - elif isinstance(meta_attr, self._tensor_type): - # e.g. self.T with torch.Tensor should still be wrapped - return type(self)._wrap_fn(lambda s: getattr(s, name))(self) - else: - # no need to wrap non-tensor properties, - # and they likely don't depend on the actual contents of the tensor - return meta_attr - - namespace["__getattr__"] = __getattr__ - - # need to make a builder for the wrapped wrapper to copy the name, - # or else it fails with very cryptic error messages, - # because somehow the same string would end up in every closures - def mk_wrap(op_name: str, *, meta_noop: bool = False): - # need to wrap the wrapper to get self - def wrapped_special_op(self, *args, **kwargs): - return type(self)._wrap_fn( - getattr(type(self)._tensor_type, op_name), - meta_noop=meta_noop, - )(self, *args, **kwargs) - return wrapped_special_op - - # special methods bypass __getattr__, so they need to be added manually - # ref: https://docs.python.org/3/reference/datamodel.html#special-lookup - # NOTE: doing this from a metaclass is very convenient - # TODO: make this even more comprehensive - for binary_op in ( - "lt", "le", "eq", "ne", "ge", "gt", "not" - "abs", "add", "and", "floordiv", "invert", "lshift", "mod", "mul", "matmul", - "neg", "or", "pos", "pow", "rshift", "sub", "truediv", "xor", - "iadd", "iand", "ifloordiv", "ilshift", "imod", "imul", "ior", "irshift", "isub", "ixor", - "radd", "rand", "rfloordiv", "rmul", "ror", "rpow", "rsub", "rtruediv", "rxor", - ): - attr_name = f"__{binary_op}__" - # the result of these operators usually has the same shape and dtype as the input, - # so evaluation on the meta tensor can be skipped. - namespace[attr_name] = mk_wrap(attr_name, meta_noop=True) - - for special_op in ( - "getitem", "setitem", "len", - ): - attr_name = f"__{special_op}__" - namespace[attr_name] = mk_wrap(attr_name, meta_noop=False) - - return super().__new__(cls, name, bases, namespace, **kwargs) - - -# Tree of lazy tensors -class LazyBase(ABC, metaclass=LazyMeta): - _tensor_type: type - _meta: Any - _data: Any | None - _args: tuple - _kwargs: dict[str, Any] - _func: Callable[[Any], Any] | None - - def __init__(self, *, meta: Any, data: Any | None = None, args: tuple = (), kwargs: dict[str, Any] | None = None, func: Callable[[Any], Any] | None = None): - super().__init__() - self._meta = meta - self._data = data - self._args = args - self._kwargs = kwargs if kwargs is not None else {} - self._func = func - assert self._func is not None or self._data is not None - - def __init_subclass__(cls) -> None: - if "_tensor_type" not in cls.__dict__: - raise TypeError(f"property '_tensor_type' must be defined for {cls!r}") - return super().__init_subclass__() - - @staticmethod - def _recurse_apply(o: Any, fn: Callable[[Any], Any]) -> Any: - # TODO: dict and set - if isinstance(o, (list, tuple)): - L = [] - for item in o: - L.append(LazyBase._recurse_apply(item, fn)) - if isinstance(o, tuple): - L = tuple(L) - return L - elif isinstance(o, LazyBase): - return fn(o) - else: - return o - - @classmethod - def _wrap_fn(cls, fn: Callable, *, use_self: LazyBase | None = None, meta_noop: bool | DTypeLike | tuple[DTypeLike, Callable[[tuple[int, ...]], tuple[int, ...]]] = False) -> Callable[[Any], Any]: - def wrapped_fn(*args, **kwargs): - if kwargs is None: - kwargs = {} - args = ((use_self,) if use_self is not None else ()) + args - - meta_args = LazyBase._recurse_apply(args, lambda t: t._meta) - # TODO: maybe handle tensors in kwargs too - - if isinstance(meta_noop, bool) and not meta_noop: - try: - res = fn(*meta_args, **kwargs) - except NotImplementedError: - # running some operations on PyTorch's Meta tensors can cause this exception - res = None - else: - # some operators don't need to actually run on the meta tensors - assert len(args) > 0 - res = args[0] - assert isinstance(res, cls) - res = res._meta - # allow operations to override the dtype and shape - if meta_noop is not True: - if isinstance(meta_noop, tuple): - dtype, shape = meta_noop - assert callable(shape) - res = cls.meta_with_dtype_and_shape(dtype, shape(res.shape)) - else: - res = cls.meta_with_dtype_and_shape(meta_noop, res.shape) - - if isinstance(res, cls._tensor_type): - return cls(meta=cls.eager_to_meta(res), args=args, kwargs=kwargs, func=fn) - elif isinstance(res, tuple) and all(isinstance(t, cls._tensor_type) for t in res): - # share the evaluation between lazy tuple elements - shared_args: list = [args, None] - - def eager_tuple_element(a: list[Any], i: int = 0, /, **kw) -> LazyBase: - assert len(a) == 2 - if a[1] is None: - a[1] = fn(*a[0], **kw) - return a[1][i] - return tuple(cls(meta=cls.eager_to_meta(res[i]), args=(shared_args, i), kwargs=kwargs, func=eager_tuple_element) for i in range(len(res))) - else: - del res # not needed - # non-tensor return likely relies on the contents of the args - # (e.g. the result of torch.equal) - eager_args = cls.to_eager(args) - return fn(*eager_args, **kwargs) - return wrapped_fn - - @classmethod - def to_eager(cls, t: Any) -> Any: - def simple_to_eager(_t: LazyBase) -> Any: - if _t._data is not None: - return _t._data - - # NOTE: there's a recursion limit in Python (usually 1000) - - assert _t._func is not None - _t._args = cls._recurse_apply(_t._args, simple_to_eager) - _t._data = _t._func(*_t._args, **_t._kwargs) - # sanity check - assert _t._data is not None - assert _t._data.dtype == _t._meta.dtype - assert _t._data.shape == _t._meta.shape - - return _t._data - - # recurse into lists and/or tuples, keeping their structure - return cls._recurse_apply(t, simple_to_eager) - - @classmethod - def eager_to_meta(cls, t: Any) -> Any: - return cls.meta_with_dtype_and_shape(t.dtype, t.shape) - - # must be overridden, meta tensor init is backend-specific - @classmethod - @abstractmethod - def meta_with_dtype_and_shape(cls, dtype: Any, shape: Any) -> Any: pass - - @classmethod - def from_eager(cls, t: Any) -> Any: - if type(t) is cls: - # already lazy - return t - elif isinstance(t, cls._tensor_type): - return cls(meta=cls.eager_to_meta(t), data=t) - else: - return TypeError(f"{type(t)!r} is not compatible with {cls._tensor_type!r}") - - -class LazyNumpyTensor(LazyBase): - _tensor_type = np.ndarray - - shape: tuple[int, ...] # Makes the type checker happy in quants.py - - @classmethod - def meta_with_dtype_and_shape(cls, dtype: DTypeLike, shape: tuple[int, ...]) -> np.ndarray[Any, Any]: - # The initial idea was to use np.nan as the fill value, - # but non-float types like np.int16 can't use that. - # So zero it is. - cheat = np.zeros(1, dtype) - return np.lib.stride_tricks.as_strided(cheat, shape, (0 for _ in shape)) - - def astype(self, dtype, *args, **kwargs): - meta = type(self).meta_with_dtype_and_shape(dtype, self._meta.shape) - full_args = (self, dtype,) + args - return type(self)(meta=meta, args=full_args, kwargs=kwargs, func=(lambda a, *args, **kwargs: a.astype(*args, **kwargs))) - - def tofile(self, *args, **kwargs): - eager = LazyNumpyTensor.to_eager(self) - return eager.tofile(*args, **kwargs) - - # TODO: __array_function__ diff --git a/venv/lib/python3.10/site-packages/gguf/metadata.py b/venv/lib/python3.10/site-packages/gguf/metadata.py deleted file mode 100644 index e807f434689de669816504921f54809cac15fc71..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/metadata.py +++ /dev/null @@ -1,642 +0,0 @@ -from __future__ import annotations - -import re -import json -import yaml -import logging -from pathlib import Path -from typing import Any, Literal, Optional -from dataclasses import dataclass - -from .constants import Keys - -import gguf - -logger = logging.getLogger("metadata") - - -@dataclass -class Metadata: - # Authorship Metadata to be written to GGUF KV Store - name: Optional[str] = None - author: Optional[str] = None - version: Optional[str] = None - organization: Optional[str] = None - finetune: Optional[str] = None - basename: Optional[str] = None - description: Optional[str] = None - quantized_by: Optional[str] = None - size_label: Optional[str] = None - url: Optional[str] = None - doi: Optional[str] = None - uuid: Optional[str] = None - repo_url: Optional[str] = None - source_url: Optional[str] = None - source_doi: Optional[str] = None - source_uuid: Optional[str] = None - source_repo_url: Optional[str] = None - license: Optional[str] = None - license_name: Optional[str] = None - license_link: Optional[str] = None - base_models: Optional[list[dict]] = None - tags: Optional[list[str]] = None - languages: Optional[list[str]] = None - datasets: Optional[list[dict]] = None - - @staticmethod - def load(metadata_override_path: Optional[Path] = None, model_path: Optional[Path] = None, model_name: Optional[str] = None, total_params: int = 0) -> Metadata: - # This grabs as many contextual authorship metadata as possible from the model repository - # making any conversion as required to match the gguf kv store metadata format - # as well as giving users the ability to override any authorship metadata that may be incorrect - - # Create a new Metadata instance - metadata = Metadata() - - model_card = Metadata.load_model_card(model_path) - hf_params = Metadata.load_hf_parameters(model_path) - # TODO: load adapter_config.json when possible, it usually contains the base model of the LoRA adapter - - # heuristics - metadata = Metadata.apply_metadata_heuristic(metadata, model_card, hf_params, model_path, total_params) - - # Metadata Override File Provided - # This is based on LLM_KV_NAMES mapping in llama.cpp - metadata_override = Metadata.load_metadata_override(metadata_override_path) - - metadata.name = metadata_override.get(Keys.General.NAME, metadata.name) - metadata.author = metadata_override.get(Keys.General.AUTHOR, metadata.author) - metadata.version = metadata_override.get(Keys.General.VERSION, metadata.version) - metadata.organization = metadata_override.get(Keys.General.ORGANIZATION, metadata.organization) - - metadata.finetune = metadata_override.get(Keys.General.FINETUNE, metadata.finetune) - metadata.basename = metadata_override.get(Keys.General.BASENAME, metadata.basename) - - metadata.description = metadata_override.get(Keys.General.DESCRIPTION, metadata.description) - metadata.quantized_by = metadata_override.get(Keys.General.QUANTIZED_BY, metadata.quantized_by) - - metadata.size_label = metadata_override.get(Keys.General.SIZE_LABEL, metadata.size_label) - metadata.license_name = metadata_override.get(Keys.General.LICENSE_NAME, metadata.license_name) - metadata.license_link = metadata_override.get(Keys.General.LICENSE_LINK, metadata.license_link) - - metadata.url = metadata_override.get(Keys.General.URL, metadata.url) - metadata.doi = metadata_override.get(Keys.General.DOI, metadata.doi) - metadata.uuid = metadata_override.get(Keys.General.UUID, metadata.uuid) - metadata.repo_url = metadata_override.get(Keys.General.REPO_URL, metadata.repo_url) - - metadata.source_url = metadata_override.get(Keys.General.SOURCE_URL, metadata.source_url) - metadata.source_doi = metadata_override.get(Keys.General.SOURCE_DOI, metadata.source_doi) - metadata.source_uuid = metadata_override.get(Keys.General.SOURCE_UUID, metadata.source_uuid) - metadata.source_repo_url = metadata_override.get(Keys.General.SOURCE_REPO_URL, metadata.source_repo_url) - - # Base Models is received here as an array of models - metadata.base_models = metadata_override.get("general.base_models", metadata.base_models) - - # Datasets is received here as an array of datasets - metadata.datasets = metadata_override.get("general.datasets", metadata.datasets) - - metadata.tags = metadata_override.get(Keys.General.TAGS, metadata.tags) - metadata.languages = metadata_override.get(Keys.General.LANGUAGES, metadata.languages) - - # Direct Metadata Override (via direct cli argument) - if model_name is not None: - metadata.name = model_name - - return metadata - - @staticmethod - def load_metadata_override(metadata_override_path: Optional[Path] = None) -> dict[str, Any]: - if metadata_override_path is None or not metadata_override_path.is_file(): - return {} - - with open(metadata_override_path, "r", encoding="utf-8") as f: - return json.load(f) - - @staticmethod - def load_model_card(model_path: Optional[Path] = None) -> dict[str, Any]: - if model_path is None or not model_path.is_dir(): - return {} - - model_card_path = model_path / "README.md" - - if not model_card_path.is_file(): - return {} - - # The model card metadata is assumed to always be in YAML (frontmatter) - # ref: https://github.com/huggingface/transformers/blob/a5c642fe7a1f25d3bdcd76991443ba6ff7ee34b2/src/transformers/modelcard.py#L468-L473 - yaml_content: str = "" - with open(model_card_path, "r", encoding="utf-8") as f: - content = f.read() - lines = content.splitlines() - lines_yaml = [] - if len(lines) == 0: - # Empty file - return {} - if len(lines) > 0 and lines[0] != "---": - # No frontmatter - return {} - for line in lines[1:]: - if line == "---": - break # End of frontmatter - else: - lines_yaml.append(line) - yaml_content = "\n".join(lines_yaml) + "\n" - - # Quick hack to fix the Norway problem - # https://hitchdev.com/strictyaml/why/implicit-typing-removed/ - yaml_content = yaml_content.replace("- no\n", "- \"no\"\n") - - if yaml_content: - data = yaml.safe_load(yaml_content) - if isinstance(data, dict): - return data - else: - logger.error(f"while reading YAML model card frontmatter, data is {type(data)} instead of dict") - return {} - else: - return {} - - @staticmethod - def load_hf_parameters(model_path: Optional[Path] = None) -> dict[str, Any]: - if model_path is None or not model_path.is_dir(): - return {} - - config_path = model_path / "config.json" - - if not config_path.is_file(): - return {} - - with open(config_path, "r", encoding="utf-8") as f: - return json.load(f) - - @staticmethod - def id_to_title(string): - # Convert capitalization into title form unless acronym or version number - return ' '.join([w.title() if w.islower() and not re.match(r'^(v\d+(?:\.\d+)*|\d.*)$', w) else w for w in string.strip().replace('-', ' ').split()]) - - @staticmethod - def get_model_id_components(model_id: Optional[str] = None, total_params: int = 0) -> tuple[str | None, str | None, str | None, str | None, str | None, str | None]: - # Huggingface often store model id as '/' - # so let's parse it and apply some heuristics if possible for model name components - - if model_id is None: - # model ID missing - return None, None, None, None, None, None - - if ' ' in model_id: - # model ID is actually a normal human sentence - # which means its most likely a normal model name only - # not part of the hugging face naming standard, but whatever - return model_id, None, None, None, None, None - - if '/' in model_id: - # model ID (huggingface style) - org_component, model_full_name_component = model_id.split('/', 1) - else: - # model ID but missing org components - org_component, model_full_name_component = None, model_id - - # Check if we erroneously matched against './' or '../' etc... - if org_component is not None and len(org_component) > 0 and org_component[0] == '.': - org_component = None - - name_parts: list[str] = model_full_name_component.split('-') - - # Remove empty parts - for i in reversed(range(len(name_parts))): - if len(name_parts[i]) == 0: - del name_parts[i] - - name_types: list[ - set[Literal["basename", "size_label", "finetune", "version", "type"]] - ] = [set() for _ in name_parts] - - # Annotate the name - for i, part in enumerate(name_parts): - # Version - if re.fullmatch(r'(v|iter)?\d+([.]\d+)*', part, re.IGNORECASE): - name_types[i].add("version") - # Quant type (should not be there for base models, but still annotated) - elif re.fullmatch(r'i?q\d(_\w)*|b?fp?(16|32)', part, re.IGNORECASE): - name_types[i].add("type") - name_parts[i] = part.upper() - # Model size - elif i > 0 and re.fullmatch(r'(([A]|\d+[x])?\d+([._]\d+)?[KMBT][\d]?|small|mini|medium|large|x?xl)', part, re.IGNORECASE): - part = part.replace("_", ".") - # Handle weird bloom-7b1 notation - if part[-1].isdecimal(): - part = part[:-2] + "." + part[-1] + part[-2] - # Normalize the size suffixes - if len(part) > 1 and part[-2].isdecimal(): - if part[-1] in "kmbt": - part = part[:-1] + part[-1].upper() - if total_params != 0: - try: - label_params = float(part[:-1]) * pow(1000, " KMBT".find(part[-1])) - # Only use it as a size label if it's close or bigger than the model size - # Note that LoRA adapters don't necessarily include all layers, - # so this is why bigger label sizes are accepted. - # Do not use the size label when it's smaller than 1/8 of the model size - if (total_params < 0 and label_params < abs(total_params) // 8) or ( - # Check both directions when the current model isn't a LoRA adapter - total_params > 0 and abs(label_params - total_params) > 7 * total_params // 8 - ): - # Likely a context length - name_types[i].add("finetune") - # Lowercase the size when it's a context length - part = part[:-1] + part[-1].lower() - except ValueError: - # Failed to convert the size label to float, use it anyway - pass - if len(name_types[i]) == 0: - name_types[i].add("size_label") - name_parts[i] = part - # Some easy to recognize finetune names - elif i > 0 and re.fullmatch(r'chat|instruct|vision|lora', part, re.IGNORECASE): - if total_params < 0 and part.lower() == "lora": - # ignore redundant "lora" in the finetune part when the output is a lora adapter - name_types[i].add("type") - else: - name_types[i].add("finetune") - - # Ignore word-based size labels when there is at least a number-based one present - # TODO: should word-based size labels always be removed instead? - if any(c.isdecimal() for n, t in zip(name_parts, name_types) if "size_label" in t for c in n): - for n, t in zip(name_parts, name_types): - if "size_label" in t: - if all(c.isalpha() for c in n): - t.remove("size_label") - - at_start = True - # Find the basename through the annotated name - for part, t in zip(name_parts, name_types): - if at_start and ((len(t) == 0 and part[0].isalpha()) or "version" in t): - t.add("basename") - else: - if at_start: - at_start = False - if len(t) == 0: - t.add("finetune") - - # Remove the basename annotation from trailing version - for part, t in zip(reversed(name_parts), reversed(name_types)): - if "basename" in t and len(t) > 1: - t.remove("basename") - else: - break - - basename = "-".join(n for n, t in zip(name_parts, name_types) if "basename" in t) or None - # Deduplicate size labels using order-preserving 'dict' ('set' seems to sort the keys) - size_label = "-".join(dict.fromkeys(s for s, t in zip(name_parts, name_types) if "size_label" in t).keys()) or None - finetune = "-".join(f for f, t in zip(name_parts, name_types) if "finetune" in t) or None - # TODO: should the basename version always be excluded? - # NOTE: multiple finetune versions are joined together - version = "-".join(v for v, t, in zip(name_parts, name_types) if "version" in t and "basename" not in t) or None - - if size_label is None and finetune is None and version is None: - # Too ambiguous, output nothing - basename = None - - return model_full_name_component, org_component, basename, finetune, version, size_label - - @staticmethod - def apply_metadata_heuristic(metadata: Metadata, model_card: Optional[dict] = None, hf_params: Optional[dict] = None, model_path: Optional[Path] = None, total_params: int = 0) -> Metadata: - # Reference Model Card Metadata: https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1 - - # Model Card Heuristics - ######################## - if model_card is not None: - - def use_model_card_metadata(metadata_key: str, model_card_key: str): - if model_card_key in model_card and getattr(metadata, metadata_key, None) is None: - setattr(metadata, metadata_key, model_card.get(model_card_key)) - - def use_array_model_card_metadata(metadata_key: str, model_card_key: str): - # Note: Will append rather than replace if already exist - tags_value = model_card.get(model_card_key, None) - if tags_value is None: - return - - current_value = getattr(metadata, metadata_key, None) - if current_value is None: - current_value = [] - - if isinstance(tags_value, str): - current_value.append(tags_value) - elif isinstance(tags_value, list): - current_value.extend(tags_value) - - setattr(metadata, metadata_key, current_value) - - # LLAMA.cpp's direct internal convention - # (Definitely not part of hugging face formal/informal standard) - ######################################### - use_model_card_metadata("name", "name") - use_model_card_metadata("author", "author") - use_model_card_metadata("version", "version") - use_model_card_metadata("organization", "organization") - use_model_card_metadata("description", "description") - use_model_card_metadata("finetune", "finetune") - use_model_card_metadata("basename", "basename") - use_model_card_metadata("size_label", "size_label") - use_model_card_metadata("source_url", "url") - use_model_card_metadata("source_doi", "doi") - use_model_card_metadata("source_uuid", "uuid") - use_model_card_metadata("source_repo_url", "repo_url") - - # LLAMA.cpp's huggingface style convention - # (Definitely not part of hugging face formal/informal standard... but with model_ appended to match their style) - ########################################### - use_model_card_metadata("name", "model_name") - use_model_card_metadata("author", "model_author") - use_model_card_metadata("version", "model_version") - use_model_card_metadata("organization", "model_organization") - use_model_card_metadata("description", "model_description") - use_model_card_metadata("finetune", "model_finetune") - use_model_card_metadata("basename", "model_basename") - use_model_card_metadata("size_label", "model_size_label") - use_model_card_metadata("source_url", "model_url") - use_model_card_metadata("source_doi", "model_doi") - use_model_card_metadata("source_uuid", "model_uuid") - use_model_card_metadata("source_repo_url", "model_repo_url") - - # Hugging Face Direct Convention - ################################# - - # Not part of huggingface model card standard but notice some model creator using it - # such as TheBloke in 'TheBloke/Mistral-7B-Instruct-v0.2-GGUF' - use_model_card_metadata("name", "model_name") - use_model_card_metadata("author", "model_creator") - use_model_card_metadata("basename", "model_type") - - if "base_model" in model_card or "base_models" in model_card or "base_model_sources" in model_card: - # This represents the parent models that this is based on - # Example: stabilityai/stable-diffusion-xl-base-1.0. Can also be a list (for merges) - # Example of merges: https://huggingface.co/EmbeddedLLM/Mistral-7B-Merge-14-v0.1/blob/main/README.md - metadata_base_models = [] - base_model_value = model_card.get("base_model", model_card.get("base_models", model_card.get("base_model_sources", None))) - - if base_model_value is not None: - if isinstance(base_model_value, str): - metadata_base_models.append(base_model_value) - elif isinstance(base_model_value, list): - metadata_base_models.extend(base_model_value) - - if metadata.base_models is None: - metadata.base_models = [] - - for model_id in metadata_base_models: - # NOTE: model size of base model is assumed to be similar to the size of the current model - base_model = {} - if isinstance(model_id, str): - if model_id.startswith("http://") or model_id.startswith("https://") or model_id.startswith("ssh://"): - base_model["repo_url"] = model_id - - # Check if Hugging Face ID is present in URL - if "huggingface.co" in model_id: - match = re.match(r"https?://huggingface.co/([^/]+/[^/]+)$", model_id) - if match: - model_id_component = match.group(1) - model_full_name_component, org_component, basename, finetune, version, size_label = Metadata.get_model_id_components(model_id_component, total_params) - - # Populate model dictionary with extracted components - if model_full_name_component is not None: - base_model["name"] = Metadata.id_to_title(model_full_name_component) - if org_component is not None: - base_model["organization"] = Metadata.id_to_title(org_component) - if version is not None: - base_model["version"] = version - - else: - # Likely a Hugging Face ID - model_full_name_component, org_component, basename, finetune, version, size_label = Metadata.get_model_id_components(model_id, total_params) - - # Populate model dictionary with extracted components - if model_full_name_component is not None: - base_model["name"] = Metadata.id_to_title(model_full_name_component) - if org_component is not None: - base_model["organization"] = Metadata.id_to_title(org_component) - if version is not None: - base_model["version"] = version - if org_component is not None and model_full_name_component is not None: - base_model["repo_url"] = f"https://huggingface.co/{org_component}/{model_full_name_component}" - - elif isinstance(model_id, dict): - base_model = model_id - - else: - logger.error(f"base model entry '{str(model_id)}' not in a known format") - - metadata.base_models.append(base_model) - - if "datasets" in model_card or "dataset" in model_card or "dataset_sources" in model_card: - # This represents the datasets that this was trained from - metadata_datasets = [] - dataset_value = model_card.get("datasets", model_card.get("dataset", model_card.get("dataset_sources", None))) - - if dataset_value is not None: - if isinstance(dataset_value, str): - metadata_datasets.append(dataset_value) - elif isinstance(dataset_value, list): - metadata_datasets.extend(dataset_value) - - if metadata.datasets is None: - metadata.datasets = [] - - for dataset_id in metadata_datasets: - # NOTE: model size of base model is assumed to be similar to the size of the current model - dataset = {} - if isinstance(dataset_id, str): - if dataset_id.startswith(("http://", "https://", "ssh://")): - dataset["repo_url"] = dataset_id - - # Check if Hugging Face ID is present in URL - if "huggingface.co" in dataset_id: - match = re.match(r"https?://huggingface.co/([^/]+/[^/]+)$", dataset_id) - if match: - dataset_id_component = match.group(1) - dataset_name_component, org_component, basename, finetune, version, size_label = Metadata.get_model_id_components(dataset_id_component, total_params) - - # Populate dataset dictionary with extracted components - if dataset_name_component is not None: - dataset["name"] = Metadata.id_to_title(dataset_name_component) - if org_component is not None: - dataset["organization"] = Metadata.id_to_title(org_component) - if version is not None: - dataset["version"] = version - - else: - # Likely a Hugging Face ID - dataset_name_component, org_component, basename, finetune, version, size_label = Metadata.get_model_id_components(dataset_id, total_params) - - # Populate dataset dictionary with extracted components - if dataset_name_component is not None: - dataset["name"] = Metadata.id_to_title(dataset_name_component) - if org_component is not None: - dataset["organization"] = Metadata.id_to_title(org_component) - if version is not None: - dataset["version"] = version - if org_component is not None and dataset_name_component is not None: - dataset["repo_url"] = f"https://huggingface.co/{org_component}/{dataset_name_component}" - - elif isinstance(dataset_id, dict): - dataset = dataset_id - - else: - logger.error(f"dataset entry '{str(dataset_id)}' not in a known format") - - metadata.datasets.append(dataset) - - use_model_card_metadata("license", "license") - use_model_card_metadata("license_name", "license_name") - use_model_card_metadata("license_link", "license_link") - - use_array_model_card_metadata("tags", "tags") - use_array_model_card_metadata("tags", "pipeline_tag") - - use_array_model_card_metadata("languages", "languages") - use_array_model_card_metadata("languages", "language") - - # Hugging Face Parameter Heuristics - #################################### - - if hf_params is not None: - - hf_name_or_path = hf_params.get("_name_or_path") - if hf_name_or_path is not None and hf_name_or_path.count('/') <= 1: - # Use _name_or_path only if its actually a model name and not some computer path - # e.g. 'meta-llama/Llama-2-7b-hf' - model_id = hf_name_or_path - model_full_name_component, org_component, basename, finetune, version, size_label = Metadata.get_model_id_components(model_id, total_params) - if metadata.name is None and model_full_name_component is not None: - metadata.name = Metadata.id_to_title(model_full_name_component) - if metadata.organization is None and org_component is not None: - metadata.organization = Metadata.id_to_title(org_component) - if metadata.basename is None and basename is not None: - metadata.basename = basename - if metadata.finetune is None and finetune is not None: - metadata.finetune = finetune - if metadata.version is None and version is not None: - metadata.version = version - if metadata.size_label is None and size_label is not None: - metadata.size_label = size_label - - # Directory Folder Name Fallback Heuristics - ############################################ - if model_path is not None: - model_id = model_path.name - model_full_name_component, org_component, basename, finetune, version, size_label = Metadata.get_model_id_components(model_id, total_params) - if metadata.name is None and model_full_name_component is not None: - metadata.name = Metadata.id_to_title(model_full_name_component) - if metadata.organization is None and org_component is not None: - metadata.organization = Metadata.id_to_title(org_component) - if metadata.basename is None and basename is not None: - metadata.basename = basename - if metadata.finetune is None and finetune is not None: - metadata.finetune = finetune - if metadata.version is None and version is not None: - metadata.version = version - if metadata.size_label is None and size_label is not None: - metadata.size_label = size_label - - return metadata - - def set_gguf_meta_model(self, gguf_writer: gguf.GGUFWriter): - assert self.name is not None - gguf_writer.add_name(self.name) - - if self.author is not None: - gguf_writer.add_author(self.author) - if self.version is not None: - gguf_writer.add_version(self.version) - if self.organization is not None: - gguf_writer.add_organization(self.organization) - - if self.finetune is not None: - gguf_writer.add_finetune(self.finetune) - if self.basename is not None: - gguf_writer.add_basename(self.basename) - - if self.description is not None: - gguf_writer.add_description(self.description) - if self.quantized_by is not None: - gguf_writer.add_quantized_by(self.quantized_by) - - if self.size_label is not None: - gguf_writer.add_size_label(self.size_label) - - if self.license is not None: - if isinstance(self.license, list): - gguf_writer.add_license(",".join(self.license)) - else: - gguf_writer.add_license(self.license) - if self.license_name is not None: - gguf_writer.add_license_name(self.license_name) - if self.license_link is not None: - gguf_writer.add_license_link(self.license_link) - - if self.url is not None: - gguf_writer.add_url(self.url) - if self.doi is not None: - gguf_writer.add_doi(self.doi) - if self.uuid is not None: - gguf_writer.add_uuid(self.uuid) - if self.repo_url is not None: - gguf_writer.add_repo_url(self.repo_url) - - if self.source_url is not None: - gguf_writer.add_source_url(self.source_url) - if self.source_doi is not None: - gguf_writer.add_source_doi(self.source_doi) - if self.source_uuid is not None: - gguf_writer.add_source_uuid(self.source_uuid) - if self.source_repo_url is not None: - gguf_writer.add_source_repo_url(self.source_repo_url) - - if self.base_models is not None: - gguf_writer.add_base_model_count(len(self.base_models)) - for key, base_model_entry in enumerate(self.base_models): - if "name" in base_model_entry: - gguf_writer.add_base_model_name(key, base_model_entry["name"]) - if "author" in base_model_entry: - gguf_writer.add_base_model_author(key, base_model_entry["author"]) - if "version" in base_model_entry: - gguf_writer.add_base_model_version(key, base_model_entry["version"]) - if "organization" in base_model_entry: - gguf_writer.add_base_model_organization(key, base_model_entry["organization"]) - if "description" in base_model_entry: - gguf_writer.add_base_model_description(key, base_model_entry["description"]) - if "url" in base_model_entry: - gguf_writer.add_base_model_url(key, base_model_entry["url"]) - if "doi" in base_model_entry: - gguf_writer.add_base_model_doi(key, base_model_entry["doi"]) - if "uuid" in base_model_entry: - gguf_writer.add_base_model_uuid(key, base_model_entry["uuid"]) - if "repo_url" in base_model_entry: - gguf_writer.add_base_model_repo_url(key, base_model_entry["repo_url"]) - - if self.datasets is not None: - gguf_writer.add_dataset_count(len(self.datasets)) - for key, dataset_entry in enumerate(self.datasets): - if "name" in dataset_entry: - gguf_writer.add_dataset_name(key, dataset_entry["name"]) - if "author" in dataset_entry: - gguf_writer.add_dataset_author(key, dataset_entry["author"]) - if "version" in dataset_entry: - gguf_writer.add_dataset_version(key, dataset_entry["version"]) - if "organization" in dataset_entry: - gguf_writer.add_dataset_organization(key, dataset_entry["organization"]) - if "description" in dataset_entry: - gguf_writer.add_dataset_description(key, dataset_entry["description"]) - if "url" in dataset_entry: - gguf_writer.add_dataset_url(key, dataset_entry["url"]) - if "doi" in dataset_entry: - gguf_writer.add_dataset_doi(key, dataset_entry["doi"]) - if "uuid" in dataset_entry: - gguf_writer.add_dataset_uuid(key, dataset_entry["uuid"]) - if "repo_url" in dataset_entry: - gguf_writer.add_dataset_repo_url(key, dataset_entry["repo_url"]) - - if self.tags is not None: - gguf_writer.add_tags(self.tags) - if self.languages is not None: - gguf_writer.add_languages(self.languages) diff --git a/venv/lib/python3.10/site-packages/gguf/py.typed b/venv/lib/python3.10/site-packages/gguf/py.typed deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/gguf/quants.py b/venv/lib/python3.10/site-packages/gguf/quants.py deleted file mode 100644 index 3c8ba82e19d3d9e984ba39caf5cf865b0ee8e72a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/quants.py +++ /dev/null @@ -1,1269 +0,0 @@ -from __future__ import annotations -from abc import ABC, abstractmethod -from typing import Any, Callable, Sequence -from math import log2, ceil - -from numpy.typing import DTypeLike - -from .constants import GGML_QUANT_SIZES, GGMLQuantizationType, QK_K -from .lazy import LazyNumpyTensor - -import numpy as np - - -def quant_shape_to_byte_shape(shape: Sequence[int], quant_type: GGMLQuantizationType) -> tuple[int, ...]: - block_size, type_size = GGML_QUANT_SIZES[quant_type] - if shape[-1] % block_size != 0: - raise ValueError(f"Quantized tensor row size ({shape[-1]}) is not a multiple of {quant_type.name} block size ({block_size})") - return (*shape[:-1], shape[-1] // block_size * type_size) - - -def quant_shape_from_byte_shape(shape: Sequence[int], quant_type: GGMLQuantizationType) -> tuple[int, ...]: - block_size, type_size = GGML_QUANT_SIZES[quant_type] - if shape[-1] % type_size != 0: - raise ValueError(f"Quantized tensor bytes per row ({shape[-1]}) is not a multiple of {quant_type.name} type size ({type_size})") - return (*shape[:-1], shape[-1] // type_size * block_size) - - -# This is faster than np.vectorize and np.apply_along_axis because it works on more than one row at a time -def _apply_over_grouped_rows(func: Callable[[np.ndarray], np.ndarray], arr: np.ndarray, otype: DTypeLike, oshape: tuple[int, ...]) -> np.ndarray: - rows = arr.reshape((-1, arr.shape[-1])) - osize = 1 - for dim in oshape: - osize *= dim - out = np.empty(shape=osize, dtype=otype) - # compute over groups of 16 rows (arbitrary, but seems good for performance) - n_groups = (rows.shape[0] // 16) or 1 - np.concatenate([func(group).ravel() for group in np.array_split(rows, n_groups)], axis=0, out=out) - return out.reshape(oshape) - - -# round away from zero -# ref: https://stackoverflow.com/a/59143326/22827863 -def np_roundf(n: np.ndarray) -> np.ndarray: - a = abs(n) - floored = np.floor(a) - b = floored + np.floor(2 * (a - floored)) - return np.sign(n) * b - - -class QuantError(Exception): ... - - -_type_traits: dict[GGMLQuantizationType, type[__Quant]] = {} - - -def quantize(data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray: - if qtype == GGMLQuantizationType.F32: - return data.astype(np.float32, copy=False) - elif qtype == GGMLQuantizationType.F16: - return data.astype(np.float16, copy=False) - elif (q := _type_traits.get(qtype)) is not None: - return q.quantize(data) - else: - raise NotImplementedError(f"Quantization for {qtype.name} is not yet implemented") - - -def dequantize(data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray: - if qtype == GGMLQuantizationType.F32: - return data.view(np.float32) - elif qtype == GGMLQuantizationType.F16: - return data.view(np.float16).astype(np.float32) - elif (q := _type_traits.get(qtype)) is not None: - return q.dequantize(data) - else: - raise NotImplementedError(f"Dequantization for {qtype.name} is not yet implemented") - - -class __Quant(ABC): - qtype: GGMLQuantizationType - block_size: int - type_size: int - - grid: np.ndarray[Any, np.dtype[np.float32]] | None = None - grid_shape: tuple[int, int] = (0, 0) - grid_map: tuple[int | float, ...] = () - grid_hex: bytes | None = None - - def __init__(self): - return TypeError("Quant conversion classes can't have instances") - - def __init_subclass__(cls, qtype: GGMLQuantizationType) -> None: - cls.qtype = qtype - cls.block_size, cls.type_size = GGML_QUANT_SIZES[qtype] - cls.__quantize_lazy = LazyNumpyTensor._wrap_fn( - cls.__quantize_array, - meta_noop=(np.uint8, cls.__shape_to_bytes) - ) - cls.__dequantize_lazy = LazyNumpyTensor._wrap_fn( - cls.__dequantize_array, - meta_noop=(np.float32, cls.__shape_from_bytes) - ) - assert qtype not in _type_traits - _type_traits[qtype] = cls - - @classmethod - def init_grid(cls): - if cls.grid is not None or cls.grid_hex is None: - return - - bits_per_elem = ceil(log2(len(cls.grid_map))) - assert bits_per_elem != 0, cls.qtype.name - elems_per_byte = 8 // bits_per_elem - - grid = np.frombuffer(cls.grid_hex, dtype=np.uint8) - # decode hexadecimal chars from grid - grid = grid.reshape((-1, 2)) - grid = (np.where(grid > 0x40, grid + 9, grid) & 0x0F) << np.array([4, 0], dtype=np.uint8).reshape((1, 2)) - grid = grid[..., 0] | grid[..., 1] - # unpack the grid values - grid = grid.reshape((-1, 1)) >> np.array([i for i in range(0, 8, 8 // elems_per_byte)], dtype=np.uint8).reshape((1, elems_per_byte)) - grid = (grid & ((1 << bits_per_elem) - 1)).reshape((-1, 1)) - grid_map = np.array(cls.grid_map, dtype=np.float32).reshape((1, -1)) - grid = np.take_along_axis(grid_map, grid, axis=-1) - cls.grid = grid.reshape((1, 1, *cls.grid_shape)) - - @classmethod - @abstractmethod - def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - raise NotImplementedError - - @classmethod - @abstractmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - raise NotImplementedError - - @classmethod - def quantize_rows(cls, rows: np.ndarray) -> np.ndarray: - rows = rows.astype(np.float32, copy=False) - shape = rows.shape - n_blocks = rows.size // cls.block_size - blocks = rows.reshape((n_blocks, cls.block_size)) - blocks = cls.quantize_blocks(blocks) - assert blocks.dtype == np.uint8 - assert blocks.shape[-1] == cls.type_size - return blocks.reshape(cls.__shape_to_bytes(shape)) - - @classmethod - def dequantize_rows(cls, rows: np.ndarray) -> np.ndarray: - rows = rows.view(np.uint8) - shape = rows.shape - n_blocks = rows.size // cls.type_size - blocks = rows.reshape((n_blocks, cls.type_size)) - blocks = cls.dequantize_blocks(blocks) - assert blocks.dtype == np.float32 - assert blocks.shape[-1] == cls.block_size - return blocks.reshape(cls.__shape_from_bytes(shape)) - - @classmethod - def __shape_to_bytes(cls, shape: Sequence[int]): - return quant_shape_to_byte_shape(shape, cls.qtype) - - @classmethod - def __shape_from_bytes(cls, shape: Sequence[int]): - return quant_shape_from_byte_shape(shape, cls.qtype) - - @classmethod - def __quantize_array(cls, array: np.ndarray) -> np.ndarray: - return _apply_over_grouped_rows(cls.quantize_rows, arr=array, otype=np.uint8, oshape=cls.__shape_to_bytes(array.shape)) - - @classmethod - def __dequantize_array(cls, array: np.ndarray) -> np.ndarray: - cls.init_grid() - return _apply_over_grouped_rows(cls.dequantize_rows, arr=array, otype=np.float32, oshape=cls.__shape_from_bytes(array.shape)) - - @classmethod - def __quantize_lazy(cls, lazy_tensor: LazyNumpyTensor, /) -> Any: - pass - - @classmethod - def __dequantize_lazy(cls, lazy_tensor: LazyNumpyTensor, /) -> Any: - pass - - @classmethod - def can_quantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> bool: - return tensor.shape[-1] % cls.block_size == 0 - - @classmethod - def quantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> np.ndarray: - if not cls.can_quantize(tensor): - raise QuantError(f"Can't quantize tensor with shape {tensor.shape} to {cls.qtype.name}") - if isinstance(tensor, LazyNumpyTensor): - return cls.__quantize_lazy(tensor) - else: - return cls.__quantize_array(tensor) - - @classmethod - def dequantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> np.ndarray: - if isinstance(tensor, LazyNumpyTensor): - return cls.__dequantize_lazy(tensor) - else: - return cls.__dequantize_array(tensor) - - -class BF16(__Quant, qtype=GGMLQuantizationType.BF16): - @classmethod - # same as ggml_compute_fp32_to_bf16 in ggml-impl.h - def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n = blocks.view(np.uint32) - # force nan to quiet - n = np.where((n & 0x7fffffff) > 0x7f800000, (n & np.uint32(0xffff0000)) | np.uint32(64 << 16), n) - # round to nearest even - n = (np.uint64(n) + (0x7fff + ((n >> 16) & 1))) >> 16 - return n.astype(np.uint16).view(np.uint8) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - return (blocks.view(np.int16).astype(np.int32) << 16).view(np.float32) - - -class Q4_0(__Quant, qtype=GGMLQuantizationType.Q4_0): - @classmethod - def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - imax = abs(blocks).argmax(axis=-1, keepdims=True) - max = np.take_along_axis(blocks, imax, axis=-1) - - d = max / -8 - with np.errstate(divide="ignore"): - id = np.where(d == 0, 0, 1 / d) - # FIXME: Q4_0's reference rounding is cursed and depends on FMA - qs = np.trunc((np.float64(blocks) * np.float64(id)) + np.float64(8.5), dtype=np.float32).astype(np.uint8).clip(0, 15) - - qs = qs.reshape((n_blocks, 2, cls.block_size // 2)) - qs = qs[..., 0, :] | (qs[..., 1, :] << np.uint8(4)) - - d = d.astype(np.float16).view(np.uint8) - - return np.concatenate([d, qs], axis=-1) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d, qs = np.hsplit(blocks, [2]) - - d = d.view(np.float16).astype(np.float32) - - qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) - qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1)).astype(np.int8) - np.int8(8) - - return (d * qs.astype(np.float32)) - - -class Q4_1(__Quant, qtype=GGMLQuantizationType.Q4_1): - @classmethod - def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - max = blocks.max(axis=-1, keepdims=True) - min = blocks.min(axis=-1, keepdims=True) - - d = (max - min) / 15 - with np.errstate(divide="ignore"): - id = np.where(d == 0, 0, 1 / d) - qs = np.trunc((blocks - min) * id + np.float32(0.5), dtype=np.float32).astype(np.uint8).clip(0, 15) - - qs = qs.reshape((n_blocks, 2, cls.block_size // 2)) - qs = qs[..., 0, :] | (qs[..., 1, :] << np.uint8(4)) - - d = d.astype(np.float16).view(np.uint8) - m = min.astype(np.float16).view(np.uint8) - - return np.concatenate([d, m, qs], axis=-1) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d, rest = np.hsplit(blocks, [2]) - m, qs = np.hsplit(rest, [2]) - - d = d.view(np.float16).astype(np.float32) - m = m.view(np.float16).astype(np.float32) - - qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) - qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1)).astype(np.float32) - - return (d * qs) + m - - -class Q5_0(__Quant, qtype=GGMLQuantizationType.Q5_0): - @classmethod - def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - imax = abs(blocks).argmax(axis=-1, keepdims=True) - max = np.take_along_axis(blocks, imax, axis=-1) - - d = max / -16 - with np.errstate(divide="ignore"): - id = np.where(d == 0, 0, 1 / d) - # FIXME: Q5_0's reference rounding is cursed and depends on FMA - q = np.trunc((np.float64(blocks) * np.float64(id)) + np.float64(16.5), dtype=np.float32).astype(np.uint8).clip(0, 31) - - qs = q.reshape((n_blocks, 2, cls.block_size // 2)) - qs = (qs[..., 0, :] & np.uint8(0x0F)) | (qs[..., 1, :] << np.uint8(4)) - - qh = np.packbits(q.reshape((n_blocks, 1, 32)) >> np.uint8(4), axis=-1, bitorder="little").reshape(n_blocks, 4) - - d = d.astype(np.float16).view(np.uint8) - - return np.concatenate([d, qh, qs], axis=-1) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d, rest = np.hsplit(blocks, [2]) - qh, qs = np.hsplit(rest, [4]) - - d = d.view(np.float16).astype(np.float32) - qh = qh.view(np.uint32) - - qh = qh.reshape((n_blocks, 1)) >> np.array([i for i in range(32)], dtype=np.uint32).reshape((1, 32)) - ql = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) - qh = (qh & np.uint32(0x01)).astype(np.uint8) - ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1)) - - qs = (ql | (qh << np.uint8(4))).astype(np.int8) - np.int8(16) - - return (d * qs.astype(np.float32)) - - -class Q5_1(__Quant, qtype=GGMLQuantizationType.Q5_1): - @classmethod - def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - max = blocks.max(axis=-1, keepdims=True) - min = blocks.min(axis=-1, keepdims=True) - - d = (max - min) / 31 - with np.errstate(divide="ignore"): - id = np.where(d == 0, 0, 1 / d) - q = np.trunc((blocks - min) * id + np.float32(0.5), dtype=np.float32).astype(np.uint8).clip(0, 31) - - qs = q.reshape((n_blocks, 2, cls.block_size // 2)) - qs = (qs[..., 0, :] & np.uint8(0x0F)) | (qs[..., 1, :] << np.uint8(4)) - - qh = np.packbits(q.reshape((n_blocks, 1, 32)) >> np.uint8(4), axis=-1, bitorder="little").reshape(n_blocks, 4) - - d = d.astype(np.float16).view(np.uint8) - m = min.astype(np.float16).view(np.uint8) - - return np.concatenate([d, m, qh, qs], axis=-1) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d, rest = np.hsplit(blocks, [2]) - m, rest = np.hsplit(rest, [2]) - qh, qs = np.hsplit(rest, [4]) - - d = d.view(np.float16).astype(np.float32) - m = m.view(np.float16).astype(np.float32) - qh = qh.view(np.uint32) - - qh = qh.reshape((n_blocks, 1)) >> np.array([i for i in range(32)], dtype=np.uint32).reshape((1, 32)) - ql = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) - qh = (qh & np.uint32(0x01)).astype(np.uint8) - ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1)) - - qs = (ql | (qh << np.uint8(4))).astype(np.float32) - - return (d * qs) + m - - -class Q8_0(__Quant, qtype=GGMLQuantizationType.Q8_0): - @classmethod - # Implementation of Q8_0 with bit-exact same results as reference implementation in ggml-quants.c - def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - - d = abs(blocks).max(axis=1, keepdims=True) / 127 - with np.errstate(divide="ignore"): - id = np.where(d == 0, 0, 1 / d) - qs = np_roundf(blocks * id) - - # (n_blocks, 2) - d = d.astype(np.float16).view(np.uint8) - # (n_blocks, block_size) - qs = qs.astype(np.int8).view(np.uint8) - - return np.concatenate([d, qs], axis=1) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - d, x = np.split(blocks, [2], axis=1) - d = d.view(np.float16).astype(np.float32) - x = x.view(np.int8).astype(np.float32) - - return (x * d) - - -class Q2_K(__Quant, qtype=GGMLQuantizationType.Q2_K): - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - scales, rest = np.hsplit(blocks, [QK_K // 16]) - qs, rest = np.hsplit(rest, [QK_K // 4]) - d, dmin = np.hsplit(rest, [2]) - - d = d.view(np.float16).astype(np.float32) - dmin = dmin.view(np.float16).astype(np.float32) - - # (n_blocks, 16, 1) - dl = (d * (scales & 0xF).astype(np.float32)).reshape((n_blocks, QK_K // 16, 1)) - ml = (dmin * (scales >> 4).astype(np.float32)).reshape((n_blocks, QK_K // 16, 1)) - - shift = np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1)) - - qs = (qs.reshape((n_blocks, -1, 1, 32)) >> shift) & np.uint8(3) - - qs = qs.reshape((n_blocks, QK_K // 16, 16)).astype(np.float32) - - qs = dl * qs - ml - - return qs.reshape((n_blocks, -1)) - - -class Q3_K(__Quant, qtype=GGMLQuantizationType.Q3_K): - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - hmask, rest = np.hsplit(blocks, [QK_K // 8]) - qs, rest = np.hsplit(rest, [QK_K // 4]) - scales, d = np.hsplit(rest, [12]) - - d = d.view(np.float16).astype(np.float32) - - # The scales are packed at 6-bit each in this pattern: - # 0: IIIIAAAA - # 1: JJJJBBBB - # 2: KKKKCCCC - # 3: LLLLDDDD - # 4: MMMMEEEE - # 5: NNNNFFFF - # 6: OOOOGGGG - # 7: PPPPHHHH - # 8: MMIIEEAA - # 9: NNJJFFBB - # 10: OOKKGGCC - # 11: PPLLHHDD - lscales, hscales = np.hsplit(scales, [8]) - lscales = lscales.reshape((n_blocks, 1, 8)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 2, 1)) - lscales = lscales.reshape((n_blocks, 16)) - hscales = hscales.reshape((n_blocks, 1, 4)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 4, 1)) - hscales = hscales.reshape((n_blocks, 16)) - scales = (lscales & np.uint8(0x0F)) | ((hscales & np.uint8(0x03)) << np.uint8(4)) - scales = (scales.astype(np.int8) - np.int8(32)).astype(np.float32) - - dl = (d * scales).reshape((n_blocks, 16, 1)) - - ql = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1)) - qh = hmask.reshape(n_blocks, -1, 1, 32) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8, 1)) - ql = ql.reshape((n_blocks, 16, QK_K // 16)) & np.uint8(3) - qh = (qh.reshape((n_blocks, 16, QK_K // 16)) & np.uint8(1)) - qh = qh ^ np.uint8(1) # strangely, the offset is zero when the bitmask is 1 - q = (ql.astype(np.int8) - (qh << np.uint8(2)).astype(np.int8)).astype(np.float32) - - return (dl * q).reshape((n_blocks, QK_K)) - - -class Q4_K(__Quant, qtype=GGMLQuantizationType.Q4_K): - K_SCALE_SIZE = 12 - - @staticmethod - def get_scale_min(scales: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - n_blocks = scales.shape[0] - scales = scales.view(np.uint8) - ### Unpacking the following: ### - # 0 EEAAAAAA - # 1 FFBBBBBB - # 2 GGCCCCCC - # 3 HHDDDDDD - # 4 eeaaaaaa - # 5 ffbbbbbb - # 6 ggcccccc - # 7 hhdddddd - # 8 eeeeEEEE - # 9 ffffFFFF - # 10 ggggGGGG - # 11 hhhhHHHH - scales = scales.reshape((n_blocks, 3, 4)) - d, m, m_d = np.split(scales, 3, axis=-2) - - sc = np.concatenate([d & 0x3F, (m_d & 0x0F) | ((d >> 2) & 0x30)], axis=-1) - min = np.concatenate([m & 0x3F, (m_d >> 4) | ((m >> 2) & 0x30)], axis=-1) - - return (sc.reshape((n_blocks, 8)), min.reshape((n_blocks, 8))) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d, rest = np.hsplit(blocks, [2]) - dmin, rest = np.hsplit(rest, [2]) - scales, qs = np.hsplit(rest, [cls.K_SCALE_SIZE]) - - d = d.view(np.float16).astype(np.float32) - dmin = dmin.view(np.float16).astype(np.float32) - - sc, m = Q4_K.get_scale_min(scales) - - d = (d * sc.astype(np.float32)).reshape((n_blocks, -1, 1)) - dm = (dmin * m.astype(np.float32)).reshape((n_blocks, -1, 1)) - - qs = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) - qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1, 32)).astype(np.float32) - - return (d * qs - dm).reshape((n_blocks, QK_K)) - - -class Q5_K(__Quant, qtype=GGMLQuantizationType.Q5_K): - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d, rest = np.hsplit(blocks, [2]) - dmin, rest = np.hsplit(rest, [2]) - scales, rest = np.hsplit(rest, [Q4_K.K_SCALE_SIZE]) - qh, qs = np.hsplit(rest, [QK_K // 8]) - - d = d.view(np.float16).astype(np.float32) - dmin = dmin.view(np.float16).astype(np.float32) - - sc, m = Q4_K.get_scale_min(scales) - - d = (d * sc.astype(np.float32)).reshape((n_blocks, -1, 1)) - dm = (dmin * m.astype(np.float32)).reshape((n_blocks, -1, 1)) - - ql = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) - qh = qh.reshape((n_blocks, -1, 1, 32)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8, 1)) - ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1, 32)) - qh = (qh & np.uint8(0x01)).reshape((n_blocks, -1, 32)) - q = (ql | (qh << np.uint8(4))).astype(np.float32) - - return (d * q - dm).reshape((n_blocks, QK_K)) - - -class Q6_K(__Quant, qtype=GGMLQuantizationType.Q6_K): - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - ql, rest = np.hsplit(blocks, [QK_K // 2]) - qh, rest = np.hsplit(rest, [QK_K // 4]) - scales, d = np.hsplit(rest, [QK_K // 16]) - - scales = scales.view(np.int8).astype(np.float32) - d = d.view(np.float16).astype(np.float32) - d = (d * scales).reshape((n_blocks, QK_K // 16, 1)) - - ql = ql.reshape((n_blocks, -1, 1, 64)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) - ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1, 32)) - qh = qh.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1)) - qh = (qh & np.uint8(0x03)).reshape((n_blocks, -1, 32)) - q = (ql | (qh << np.uint8(4))).astype(np.int8) - np.int8(32) - q = q.reshape((n_blocks, QK_K // 16, -1)).astype(np.float32) - - return (d * q).reshape((n_blocks, QK_K)) - - -class TQ1_0(__Quant, qtype=GGMLQuantizationType.TQ1_0): - @classmethod - def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d = abs(blocks).max(axis=-1, keepdims=True) - with np.errstate(divide="ignore"): - id = np.where(d == 0, 0, 1 / d) - qs = np_roundf(blocks * id) - qs = (qs.astype(np.int8) + np.int8(1)).astype(np.uint8) - - qs0, qs1, qh = qs[..., :(32 * 5)], qs[..., (32 * 5):(48 * 5)], qs[..., (48 * 5):] - qs0 = qs0.reshape((n_blocks, -1, 5, 32)) * np.array([81, 27, 9, 3, 1], dtype=np.uint8).reshape((1, 1, 5, 1)) - qs0 = np.sum(qs0, axis=-2).reshape((n_blocks, -1)) - qs1 = qs1.reshape((n_blocks, -1, 5, 16)) * np.array([81, 27, 9, 3, 1], dtype=np.uint8).reshape((1, 1, 5, 1)) - qs1 = np.sum(qs1, axis=-2).reshape((n_blocks, -1)) - qh = qh.reshape((n_blocks, -1, 4, 4)) * np.array([81, 27, 9, 3], dtype=np.uint8).reshape((1, 1, 4, 1)) - qh = np.sum(qh, axis=-2).reshape((n_blocks, -1)) - qs = np.concatenate([qs0, qs1, qh], axis=-1) - qs = (qs.astype(np.uint16) * 256 + (243 - 1)) // 243 - - qs = qs.astype(np.uint8) - d = d.astype(np.float16).view(np.uint8) - - return np.concatenate([qs, d], axis=-1) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - qs, rest = np.hsplit(blocks, [(QK_K - 4 * QK_K // 64) // 5]) - qh, d = np.hsplit(rest, [QK_K // 64]) - - d = d.view(np.float16).astype(np.float32) - - qs0, qs1 = qs[..., :32], qs[..., 32:] - qs0 = qs0.reshape((n_blocks, -1, 1, 32)) * np.array([1, 3, 9, 27, 81], dtype=np.uint8).reshape((1, 1, 5, 1)) - qs0 = qs0.reshape((n_blocks, -1)) - qs1 = qs1.reshape((n_blocks, -1, 1, 16)) * np.array([1, 3, 9, 27, 81], dtype=np.uint8).reshape((1, 1, 5, 1)) - qs1 = qs1.reshape((n_blocks, -1)) - qh = qh.reshape((n_blocks, -1, 1, 4)) * np.array([1, 3, 9, 27], dtype=np.uint8).reshape((1, 1, 4, 1)) - qh = qh.reshape((n_blocks, -1)) - qs = np.concatenate([qs0, qs1, qh], axis=-1) - qs = ((qs.astype(np.uint16) * 3) >> 8).astype(np.int8) - np.int8(1) - - return (d * qs.astype(np.float32)) - - -class TQ2_0(__Quant, qtype=GGMLQuantizationType.TQ2_0): - @classmethod - def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d = abs(blocks).max(axis=-1, keepdims=True) - with np.errstate(divide="ignore"): - id = np.where(d == 0, 0, 1 / d) - qs = np_roundf(blocks * id) - qs = (qs.astype(np.int8) + np.int8(1)).astype(np.uint8) - - qs = qs.reshape((n_blocks, -1, 4, 32)) << np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1)) - qs = qs[..., 0, :] | qs[..., 1, :] | qs[..., 2, :] | qs[..., 3, :] - qs = qs.reshape((n_blocks, -1)) - - d = d.astype(np.float16).view(np.uint8) - - return np.concatenate([qs, d], axis=-1) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - qs, d = np.hsplit(blocks, [QK_K // 4]) - - d = d.view(np.float16).astype(np.float32) - - qs = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1)) - qs = (qs & 0x03).reshape((n_blocks, -1)).astype(np.int8) - np.int8(1) - - return (d * qs.astype(np.float32)) - - -class IQ2_XXS(__Quant, qtype=GGMLQuantizationType.IQ2_XXS): - ksigns: bytes = ( - b"\x00\x81\x82\x03\x84\x05\x06\x87\x88\x09\x0a\x8b\x0c\x8d\x8e\x0f" - b"\x90\x11\x12\x93\x14\x95\x96\x17\x18\x99\x9a\x1b\x9c\x1d\x1e\x9f" - b"\xa0\x21\x22\xa3\x24\xa5\xa6\x27\x28\xa9\xaa\x2b\xac\x2d\x2e\xaf" - b"\x30\xb1\xb2\x33\xb4\x35\x36\xb7\xb8\x39\x3a\xbb\x3c\xbd\xbe\x3f" - b"\xc0\x41\x42\xc3\x44\xc5\xc6\x47\x48\xc9\xca\x4b\xcc\x4d\x4e\xcf" - b"\x50\xd1\xd2\x53\xd4\x55\x56\xd7\xd8\x59\x5a\xdb\x5c\xdd\xde\x5f" - b"\x60\xe1\xe2\x63\xe4\x65\x66\xe7\xe8\x69\x6a\xeb\x6c\xed\xee\x6f" - b"\xf0\x71\x72\xf3\x74\xf5\xf6\x77\x78\xf9\xfa\x7b\xfc\x7d\x7e\xff" - ) - - # iq2xxs_grid, but with each byte of the original packed in 2 bits, - # by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2. - grid_shape = (256, 8) - grid_map = (0x08, 0x19, 0x2b) - grid_hex = ( - b"00000200050008000a00110014002000220028002a0041004400500058006100" - b"6400800082008a00a20001010401100115014001840198010002020222028202" - b"010404041004210424044004420448046004810484049004a404000502050805" - b"200546056905800591050906100640068406a406000805080808140828084108" - b"440850085208880804094009020a140a01100410101021104010601084109010" - b"951000110811201150115a118011241245120014081420142514491480141815" - b"6215001616160118041810184018811800190519a019511a002002200a204420" - b"6120802082202921482100220222012404241024402456240025412564259026" - b"082820289428442a014004401040184021402440404048405640604081408440" - b"9040004120416141804185410142104248425642684200440844204480449944" - b"124524450046014804481048404845480049584961498249454a904a00500850" - b"1150195020508050885004514251a4519152905492540a550156545600581158" - b"195864584059085a046010604060686000615561186260620064056410651265" - b"84654268008002800a8041808280048118814081118201840484108415844084" - b"608400854685948509864086608602880489118a0490109024904090a1901691" - b"8091459200942294449451958198209902a050a085a009a100a218a450a804a9" - ) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d, qs = np.hsplit(blocks, [2]) - - d = d.view(np.float16).astype(np.float32) - - qs = qs.view(np.uint32).reshape(n_blocks, -1, 2) - - db = d * (np.float32(0.5) + (qs[..., 1] >> 28).astype(np.float32)) * np.float32(0.25) - db = db.reshape((n_blocks, -1, 1, 1)) - - # get the sign indices and unpack the bits - signs = qs[..., 1].reshape((n_blocks, -1, 1)) >> np.array([0, 7, 14, 21], dtype=np.uint32).reshape((1, 1, 4)) - ksigns = np.frombuffer(cls.ksigns, dtype=np.uint8).reshape((1, 1, 1, 128)) - signs = (signs & np.uint32(0x7F)).reshape((n_blocks, -1, 4, 1)) - signs = np.take_along_axis(ksigns, signs, axis=-1) - signs = signs.reshape((n_blocks, -1, 4, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 1, 8)) - signs = signs & np.uint8(0x01) - signs = np.where(signs == 0, np.float32(1), np.float32(-1)) - signs = signs.reshape((n_blocks, -1, 4, 8)) - - assert cls.grid is not None - grid = np.take_along_axis(cls.grid, qs[..., 0].copy().view(np.uint8).reshape((n_blocks, -1, 1, 1)), axis=-2) - grid = grid.reshape((n_blocks, -1, 4, 8)) - - return (db * grid * signs).reshape((n_blocks, -1)) - - -class IQ2_XS(__Quant, qtype=GGMLQuantizationType.IQ2_XS): - # iq2xs_grid, but with each byte of the original packed in 2 bits, - # by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2. - grid_shape = (512, 8) - grid_map = (0x08, 0x19, 0x2b) - grid_hex = ( - b"00000200050008000a0011001400160019002000220025002800410044004600" - b"49005000520055005800610064008000820085008800910094009900a0000101" - b"04010601090110011201150118011a0121012401400142014501480151015401" - b"6001680181018401900100020202050208021102140220024102440250025502" - b"80028a0201040404060409041004120415041804210424044004420445044804" - b"5104540456046004810484049004000502050505080511051405200541054405" - b"500561058005010604061006260640064206840600080208050808080a081108" - b"14082008250841084408500858088008a008aa08010904091009400981098909" - b"000a200a280a960aa00a01100410061009101010121015101810211024104010" - b"4210451048105110541060106a10811084109010001102110511081111111411" - b"2011411144115011801194119611011204120612101240126012001402140514" - b"0814111414142014411444144914501464148014011504151015401500161416" - b"49160118041810181218401854188618001905196619511aa91a002002200520" - b"08200a201120142020204120442050208020a020012104211021402148216521" - b"002222228022a82201240424102429244024002541255225992501261a26a626" - b"002808280a28202855288828a22868299029082a202a822a882a8a2a01400440" - b"0640094010401240154018402140244040404240454048404a40514054406040" - b"6540814084409040004102410541084111411441204141414441504180418541" - b"a241014204421042124229424042004402440544084411441444194420444144" - b"4444504480449444014504451045244540459a4500460a464446504601480448" - b"1048404845485448624800491149444950496949044a00500250055008501150" - b"145020502850415044505050805001510451105115514051425100524452aa52" - b"0154045410542154405460548154a154005508558055885521566856a1560058" - b"14584158505899581a5940594259855a0160046010604060546062608660a960" - b"006124624a62926200641664106540654565a46501686a682569066a546a626a" - b"00800280058008801180148020802a8041804480508080808280a880aa800181" - b"0481068110814081518159810082208280828282a082a8820184048410841284" - b"158440846084898400854485a58518866a860088088825885a8880888288a888" - b"0689228a808a888a968aa88a0190049010904090569084900091229164915692" - b"89920094059444945094589429959095929541965198a6984999159a609a00a0" - b"02a008a00aa020a02aa0a0a051a159a1a6a100a202a208a22aa280a2a0a240a4" - b"95a465a698a60aa820a822a828a8a0a8a8a804a984a986a928aa2aaa91aaaaaa" - ) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d, rest = np.hsplit(blocks, [2]) - qs, scales = np.hsplit(rest, [2 * QK_K // 8]) - - d = d.view(np.float16).astype(np.float32) - qs = qs.view(np.uint16) - - scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2)) - scales = (scales & 0x0F).reshape((n_blocks, -1)) - db = d * (np.float32(0.5) + scales) * np.float32(0.25) - db = db.reshape((n_blocks, -1, 1, 1)) - - # get the sign indices and unpack the bits - signs = np.frombuffer(IQ2_XXS.ksigns, dtype=np.uint8).reshape(1, 1, 128) - signs = np.take_along_axis(signs, (qs >> 9).reshape((n_blocks, -1, 1)), axis=-1) - signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8)) - signs = signs & np.uint8(0x01) - signs = np.where(signs == 0, np.float32(1), np.float32(-1)) - signs = signs.reshape((n_blocks, -1, 2, 8)) - - assert cls.grid is not None - grid = np.take_along_axis(cls.grid, (qs & np.uint16(511)).reshape((n_blocks, -1, 1, 1)), axis=-2) - grid = grid.reshape((n_blocks, -1, 2, 8)) - - return (db * grid * signs).reshape((n_blocks, -1)) - - -class IQ2_S(__Quant, qtype=GGMLQuantizationType.IQ2_S): - # iq2s_grid, but with each byte of the original packed in 2 bits, - # by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2. - grid_shape = (1024, 8) - grid_map = (0x08, 0x19, 0x2b) - grid_hex = ( - b"00000200050008000a0011001400160019002000220025002800410044004600" - b"490050005200550058006100640066006900800082008500880091009400a000" - b"a500aa0001010401060109011001120115011801210124014001420145014801" - b"510154015601590160016501680181018401900192019501a101a40100020202" - b"050208021102140220022a02410244024602490250025502800285028a029402" - b"a202010404040604090410041204150418042104240426042904400442044504" - b"48044a0451045404560459046004620465048104840486048904900495049804" - b"a104a40400050205050508050a05110514051605190520052505280541054405" - b"46054905500552055505580561056405800582058505880591059405a0050106" - b"0406060609061006150640064506480651065406600681068406900600080208" - b"050808081108140816081908200825082a084108440846084908500852085508" - b"580861086408800885089408aa08010904091009120915091809210940094509" - b"480951095409600981099009000a110a140a220a280a2a0a500a990a01100410" - b"0610091010101210151018102110241026104010421045104810511054105610" - b"59106010621065106810811084108610901095109810a110a410001102110511" - b"08110a1111111411161119112011221125112811411144114611491150115211" - b"5511581161116411801182118511881191119411011204120912101215122112" - b"2412401245125112541281128412901200140214051408141114141416141914" - b"2014251428144114441446144914501452145514581461146414801482148514" - b"881491149414a014011504150615091510151215151518152115241540154215" - b"4515481551155415601581158415901500160516081611161416201641164416" - b"50168016aa160118041806180918101815181818211840184218451848185118" - b"541860188118841800190219051908191119141920194119441950196919a219" - b"041a101a401a561a00200220052008201120142016201920202025202a204120" - b"4420502052205520642080208a209420aa200121042110211221152121214021" - b"4221452151215421602181218421902100220a22222228222a22442250228822" - b"8a22a82201240424062409241024152418242124242440244224452448245124" - b"5424602481248424902400250525082511251425202541254425502566258025" - b"0126042610264026592600280528112814284128442850288a28aa2801290429" - b"102995290a2a222a642a882a8a2a014004400640094010401240154018401a40" - b"21402440264040404240454048404a4051405440564059406040624065408140" - b"8440904095409840a140a4400041024105410841114114411641194120412241" - b"2541414144414641494150415241554158416141644180418241854188419141" - b"9441a04101420442104212421542184224424042454248425142544260428142" - b"844200440244054408440a441144144416441944204422442544284441444444" - b"46444944504452445544584461446444804482448544884491449444a0440145" - b"0445064509451045124515451845214524454045424545454845514554456045" - b"6a4581458445904500460246054608461146144620464146444650468046a546" - b"0148044809481048124815481848214824484048424845484848514854486048" - b"84489048004902490549084911491449204941494449504980499649014a044a" - b"104a404a00500250055008501150145016501950205022502550285041504450" - b"4650495050505250555058506150645080508250855088509150945001510451" - b"0651095110511251155118512151245140514251455148515151545160518151" - b"8451905100520552085211521452205241524452505269528052015404540654" - b"0954105412541554185421542454405442544554485451545454605481548454" - b"9054005502550555085511551455205541554455505580550156045610562656" - b"405600580258055808581158145820584158445850585a588058015904591059" - b"4059005a195a855aa85a01600460066010601260156018602160246040604560" - b"4860516054606060846090600061026105610861116114612061416144615061" - b"806199610462106240625662a162006405640864116414642064416444645064" - b"806401650465106540654a656865926500669466016804681068656898680069" - b"2a69426aa16a0080028005800880118014801980208025804180448050805280" - b"5580588061808080858091809480018104810981108112811581188121812481" - b"408142814581488151815481818184819081a981008205820a82118214824182" - b"4482508201840484068409841084128415841884218440844284458448845184" - b"5484608481848484908400850285058508851185148520854185448550858085" - b"8a85018604861086298640860088058811881488418844885088a28801890489" - b"40896589228a588a5a8a828aa28a019004900990109012901590189024904090" - b"4290459048905190549060908190849090900091059111911491419144915091" - b"5a910192049210924092a6920094029405940894119414942094419444945094" - b"8094969401950495109540959895a19500964696649601980498109826984098" - b"a998009949995299909a00a005a00aa014a022a02aa041a044a050a0a2a0aaa0" - b"40a165a102a20aa222a228a22aa282a288a28aa2a8a201a404a410a440a489a4" - b"a4a400a519a551a60aa828a8a2a854a986a908aa0aaa20aa22aa28aa88aaaaaa" - ) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d, rest = np.hsplit(blocks, [2]) - qs, rest = np.hsplit(rest, [QK_K // 8]) - signs, rest = np.hsplit(rest, [QK_K // 8]) - qh, scales = np.hsplit(rest, [QK_K // 32]) - - d = d.view(np.float16).astype(np.float32) - - scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2)) - scales = (scales & 0x0F).reshape((n_blocks, -1)) - db = d * (np.float32(0.5) + scales) * np.float32(0.25) - db = db.reshape((n_blocks, -1, 1, 1)) - - # unpack the sign bits - signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8)) - signs = signs & np.uint8(0x01) - signs = np.where(signs == 0, np.float32(1), np.float32(-1)) - signs = signs.reshape((n_blocks, -1, 2, 8)) - - qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4)) - qs = qs.astype(np.uint16) | ((qh & 0x03).astype(np.uint16) << 8).reshape((n_blocks, -1)) - - assert cls.grid is not None - grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2) - grid = grid.reshape((n_blocks, -1, 2, 8)) - - return (db * grid * signs).reshape((n_blocks, -1)) - - -class IQ3_XXS(__Quant, qtype=GGMLQuantizationType.IQ3_XXS): - grid_shape = (256, 4) - grid_map = (0x04, 0x0c, 0x14, 0x1c, 0x24, 0x2c, 0x34, 0x3e) - grid_hex = ( - b"0000020004001100130017002000220031004200730075000101030110011201" - b"2101250130013201410154017001000202020402110220022202310233023702" - b"5102570275020103070310031203250370031304370444045704730475040105" - b"0705320552053506640610071407160743076107011003101010121021102310" - b"3010321034104710501000110211111120112211011203121012121221123012" - b"7212001302132013311346136613011405145014201524154615711505162217" - b"4017002002201120132020202220262031204220012103210521102112212121" - b"3021632167217021002202221122172220222222372240225522012310231423" - b"7023742335245324032527254125742501270327162745270130103012302130" - b"2330503065307230003102312031313144314631013203321032253252327232" - b"1133333330344734723400350635223555351436363663363337603704401740" - b"3540374053405740744120423742404260426642074345430444514464442545" - b"4345704505471047124730471250415070500051065126515551145232527252" - b"0253535310542354275472540255315550562457425724604460466064602161" - b"6161176264623063366344640565526533660367216703700570077010703270" - b"5270267140711272457252720073157333736073217441740075027524753076" - ) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d, rest = np.hsplit(blocks, [2]) - qs, scales = np.hsplit(rest, [QK_K // 4]) - - d = d.view(np.float16).astype(np.float32) - scales = scales.view(np.uint32) - - db = d * (np.float32(0.5) + (scales >> 28).astype(np.float32)) * np.float32(0.5) - db = db.reshape((n_blocks, -1, 1, 1)) - - # get the sign indices and unpack the bits - signs = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 7, 14, 21], dtype=np.uint32).reshape((1, 1, 4)) - ksigns = np.frombuffer(IQ2_XXS.ksigns, dtype=np.uint8).reshape((1, 1, 1, 128)) - signs = (signs & np.uint32(0x7F)).reshape((n_blocks, -1, 4, 1)) - signs = np.take_along_axis(ksigns, signs, axis=-1) - signs = signs.reshape((n_blocks, -1, 4, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 1, 8)) - signs = signs & np.uint8(0x01) - signs = np.where(signs == 0, np.float32(1), np.float32(-1)) - signs = signs.reshape((n_blocks, -1, 4, 8)) - - assert cls.grid is not None - grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2) - grid = grid.reshape((n_blocks, -1, 4, 8)) - - return (db * grid * signs).reshape((n_blocks, -1)) - - -class IQ3_S(__Quant, qtype=GGMLQuantizationType.IQ3_S): - grid_shape = (512, 4) - grid_map = (0x01, 0x03, 0x05, 0x07, 0x09, 0x0b, 0x0d, 0x0f) - grid_hex = ( - b"0000010002000500070010001100120014001600200021002500330040004200" - b"4500470051005300600062007100740077000001010102010401100111011501" - b"2001230127013101350144016101650172010002010205020702100213021602" - b"2102250230023402420245024702510253027002730203031103150320032203" - b"3103330336034403500352036703710375030004130417042104240432044004" - b"4304510470040205040520052205260533054105450547056605730506061106" - b"1306310652067106000702070407200722072607330750075407001001100210" - b"0410101011101310151017102010221031103410361054105610611072100011" - b"0111031106111011141121113011331141115011521170117611001212121512" - b"1712201224123212401243125512601272120113041307131013131321132713" - b"3013341341136213701303140514121414143114331442144614501454140115" - b"1015131521153015321551152016241627164416461601170317101712172117" - b"3517411762177017002001200320052007201020122014201620212023202720" - b"3020322041204320452050205220672070207320752000210221102113211721" - b"2221252131213421422151210122042207222122232230223722412253225722" - b"7122742200230223052311232223242331233323422350236623012407242024" - b"2324322435244124722475240425112522253725402553257025002602260726" - b"2126552661260527112726273027432750270230113013301530173022303130" - b"3330353042304430473051306330713001310331053114312131233140316031" - b"7231763100321232203232323432503201331033143321332333273330334133" - b"4333473355337333033411341634223431345234603464340135103512352535" - b"3235443556357335163641360137033720372237353700400440124020402440" - b"2740324041405040704002410741114113412241304135414341514155410142" - b"0342104215422142334240425742624270420443114313432043224331433543" - b"0044024424443744404471440545074521456245134634466046104715473047" - b"4347514702501050145022504050445047505250665074500151035105511251" - b"2151325172510052115223523052365253520253075310532753445351536553" - b"7353015404542054325446541255265551555355425602570457225711601360" - b"1560316033606060006120612761646112623462426255626262706200631463" - b"2163406325644364626400650365346560650566406611671367007004700770" - b"2070227036704070547062700271117124714371457101720472107216722172" - b"3072517202733273357353730174057413742074507422754275027631760077" - ) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d, rest = np.hsplit(blocks, [2]) - qs, rest = np.hsplit(rest, [QK_K // 4]) - qh, rest = np.hsplit(rest, [QK_K // 32]) - signs, scales = np.hsplit(rest, [QK_K // 8]) - - d = d.view(np.float16).astype(np.float32) - - scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2)) - scales = (scales & 0x0F).reshape((n_blocks, -1)) - db = d * (1 + 2 * scales) - db = db.reshape((n_blocks, -1, 1, 1)) - - # unpack the sign bits - signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8)) - signs = signs & np.uint8(0x01) - signs = np.where(signs == 0, np.float32(1), np.float32(-1)) - signs = signs.reshape((n_blocks, -1, 4, 8)) - - qh = qh.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8) - qh = (qh & 0x01).astype(np.uint16).reshape((n_blocks, -1)) - qs = qs.astype(np.uint16) | (qh << 8) - - assert cls.grid is not None - grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2) - grid = grid.reshape((n_blocks, -1, 4, 8)) - - return (db * grid * signs).reshape((n_blocks, -1)) - - -class IQ1_S(__Quant, qtype=GGMLQuantizationType.IQ1_S): - # iq1s_grid, with each byte packed into 2 bits - # -1, 0, 1 <=> 0, 1, 2 - grid_shape = (2048, 8) - grid_map = (-1, 0, 1) - grid_hex = ( - b"00000200050008000a00110015002000220028002a0045005100540056006500" - b"8000820088008a009500a000a200a800aa000401050111011401160119011a01" - b"2501410146014901520155015a0161016401660168018501910194019601a501" - b"0002020208020a0215022002220228022a024502510259026402690280028202" - b"88028a02910295029902a002a202a802aa021104140416042504410449045504" - b"5a046404650491049904a5040105040505050605150518051a05290540054505" - b"4a0550055105540555055605590560056205650568056a058105910595059805" - b"9a05a105a405a505a605a9051406190641064406500652065506580660066106" - b"6606690685069106940699060008020808080a0815082008220828082a084508" - b"5108560865088008820888088a089508a008a208a808aa080509110914091909" - b"2409250941095009510955096109640969099109940996099909a509000a020a" - b"080a0a0a150a200a220a280a2a0a450a510a590a610a650a800a820a850a880a" - b"8a0a950aa00aa20aa80aaa0a1010111014101910241025104110441050105510" - b"58106110641065106910911094109610a110a510011104110611091110111211" - b"1511181121112411291145114a11501151115211541155115611591160116511" - b"841192119511a111a41111121412161225124012461249125212551258125a12" - b"641266128512911294129612a512011406140914141415141814191421142614" - b"41144514461448144a1451145414551456145914621465146814841489149014" - b"94149514981499149a14a114a414a514a914021505150a151115141515151615" - b"191520152215251528152a154115441545154615511552155415551556155915" - b"5a1561156415651566156915801582158415851588158a159015911594159515" - b"961599159a15a015a215a51501160416051606161516161618161a1621162616" - b"401642164416451648164a165116551656165816591661166416651668166916" - b"6a1686168a1692169516a416a916111816182518411844184618491850185518" - b"58185a1860186118641866186918851891189418a5181019121915191a192119" - b"25194219441945194819511954195519561959195a19601965196a1989199119" - b"921995199819a119a619a919091a161a241a261a441a461a491a501a521a551a" - b"581a611a661a691a851a911a961a9a1a0020022008200a201520202022202520" - b"28202a20452051205920612065208020822088208a209520a020a220a520a820" - b"aa2005211121142119212521422144214921552158215a216121642165216621" - b"8521902196219921a521012208220a22112215222022222228222a2245225122" - b"562259226522812288228a2291229522a022a222a822aa220524142416241924" - b"252444244524462449245224552458245a2466248524912494249924a124a524" - b"0925152521252925402545254825512554255525592562256525682589259025" - b"9425952598259a25a125a425a625a92505261026122619262526412649265526" - b"6026612669268426862690269a260028022808280a2815282028222828282a28" - b"45285128542865288028822888288a28a028a228a828aa280929112914291929" - b"2529462949295229552961296429662969298529902996299929a429a529002a" - b"022a082a0a2a202a222a282a2a2a452a512a562a592a652a802a822a882a8a2a" - b"952aa02aa22aa82aaa2a054011401640254049405240554058405a4061406440" - b"664094409940a140a6400041014104410641094112411541164118411a412141" - b"26412941454148414a41514154415541564159415a41654168416a4181418441" - b"8641904192419541a041a141a241054211421442164225424142524255425a42" - b"6442694289429442a5420144154419442944454448444a445144544455445644" - b"61446244654468446a44814486448944904492449544a044a144a94401450245" - b"05450a4511451445154516451945204525452a45414544454545464549455045" - b"5145544555455645584559456145644565456645694582458445854588459145" - b"94459545964599459a45a545a845aa450146054609461446154618461a462146" - b"2446294640464246454648465046514652465546564659466246654668468146" - b"85468a4694469546a146a446a6460548114815481a4825484248494850485548" - b"5848614864486648694885489148944896489948a5480149054906490a491049" - b"144915491849214924492649404945494a495149524954495549564959496049" - b"6249654966496a49864989499249954996499849a149a449a649a949164a444a" - b"464a494a554a584a5a4a644a694a944aa54a0150045005500650095012501550" - b"1a50215024502950405045504850515054505550565059506550685086508950" - b"95509850a050a150a650a9500551085109510a51115114511551165118511951" - b"20512551265128512a5141514451455146514951505151515251545155515651" - b"585159515a51615164516551665169518251855191519451955196519951a051" - b"a551aa5101520652125215521a5221522452425245524a525152545255525652" - b"595262526552855290529252955299529a52a452045405541154145415541654" - b"185419542154255428542a54415444544554465449544a545054515454545554" - b"5654585459545a54615462546454655466546954805488548a54915494549554" - b"96549954a154a454a554aa540155025504550555065509551055115512551455" - b"1555165519551a55215524552555265529554055415542554455455546554855" - b"4955505551555255545555555655585559555a55605561556455655566556855" - b"69556a5581558455855589558a559055915594559555965598559955a155a455" - b"a555a655a9550056015602560456065608560956115614561556185619562056" - b"2156225624562556265628562956415645564656485649564a56505651565256" - b"545655565656585659565a566156645665566956825685568656885689568a56" - b"915695569a56a256a556a656a856a95604580558065809581058155818582158" - b"2a58455848584a58515854585558565858585958605862586458655882588958" - b"9058925895589858a158a9580159025905590a59115914591559165919592559" - b"41594459455946594959505951595259545955595659585959595a5961596459" - b"655966596959815985598959915994599559965998599959a559045a085a155a" - b"1a5a205a255a265a295a455a485a495a515a555a565a585a595a625a655a685a" - b"6a5a815a8a5a925a955a965a985a9a5aa15a0560146016601960256044605060" - b"5560566058605a60616064606660696081609660a56001610461066109611261" - b"15612161226126612961456149615161556156615961656166616a6184618a61" - b"92619561a161a661a96111621662196240624162466255625662586260628562" - b"91629662a56211641264156416641a6421642664296440644264456448644a64" - b"516454645564566459645a646064626465648464856489649064926494649564" - b"966498649a64a164a464a964056508650a651165156516651965446545654665" - b"496550655165546555655665596561656465656566656965866589658a659165" - b"9565966599659a65a265a565a665a86502660966156620662666286629664066" - b"456648664a66516654665566566658665a666066656668668066826685668a66" - b"9466966698669966a066a466a666aa661668196825684168526855685a686168" - b"6968856891689868a66801690469106915692169246926692969406941694569" - b"4669486951695469556956695969606965696a69826984698a699569a169a469" - b"a569a969116a166a186a416a446a496a506a556a586a5a6a646a656a696a866a" - b"946a986a9a6aa66a0080028008800a802080228028802a804580508051805480" - b"5680598065808080828088808a809580a080a280a880aa800581118114811681" - b"1981258141814481498150815281558156815881598164816681698185818981" - b"948196819981a5810082028208820a8215822082228228822a82518254825982" - b"65828082828288828a829582a082a282a882aa82148419844184448451845584" - b"5a846184648469849484998401850985128515851a8526852985408541854585" - b"4885518554855585568559855a856585668568856a8581858485868589859085" - b"928595859885a68511861686198625864186448649864a865086558659865a86" - b"618666866a86858691869a86a4860088028808880a8815882088228828882a88" - b"41884588518854885988658869888088828888888a889588a088a288a888aa88" - b"05890689118914891689258941894489468949895089528955895a8961896489" - b"858996899989a589008a028a088a0a8a158a208a228a288a2a8a458a518a548a" - b"568a808a828a888a8a8a958aa08aa28aa88aaa8a059011901690189019902590" - b"419046904990559058905a9069906a9085909190949096909990a59001910491" - b"069109911091159118911a912191249126912991409145915091519154915591" - b"569159916291659184918691929195919891a191a491a691a991059211921492" - b"19922592449246924992509252925592589266926992859294929692a9920194" - b"04940694109415941894269440944a9451945494559456945894599460946194" - b"62946594849486949294949495949894a194a9940095059508950a9510951195" - b"14951595169519952195259529952a9541954495459546954995509551955295" - b"549555955695589559955a956195649565956695699581958595889591959295" - b"94959595969599959a95a095a295a595a895aa95019604961096159619962096" - b"2696299645964896499651965296559656965996659668968296849689968a96" - b"929694969596a496a696a9960598169819982598419846985098529855985698" - b"5a98649865988598919896989998a59804990699099910991299159918991a99" - b"209921992499269940994299459948994a995199549955995699599962996599" - b"66996a99819984999099929995999a99a199a699059a159a259a449a469a499a" - b"509a559a589a619a859a919a949a959a969a00a002a008a00aa015a020a022a0" - b"28a02aa045a051a054a056a059a080a082a088a08aa095a0a0a0a2a0a8a0aaa0" - b"05a109a111a114a116a119a11aa146a149a151a155a158a15aa161a164a185a1" - b"90a192a196a199a102a208a20aa210a219a222a228a22aa245a251a256a259a2" - b"65a280a282a288a28aa295a2a0a2a2a2a8a2aaa219a425a441a444a450a454a4" - b"55a458a45aa461a465a466a468a469a485a406a509a510a512a515a518a526a5" - b"29a542a545a551a554a555a556a559a565a56aa581a584a585a586a589a592a5" - b"95a598a505a611a616a61aa621a625a644a646a64aa652a655a656a658a660a6" - b"62a686a690a695a696a699a6a1a6a4a6a6a600a802a808a80aa820a822a828a8" - b"2aa851a854a856a859a880a882a888a88aa895a8a0a8a2a8a8a8aaa805a914a9" - b"19a921a925a941a950a955a95aa961a966a969a990a996a900aa02aa08aa0aaa" - b"20aa22aa28aa2aaa51aa54aa56aa80aa82aa88aa8aaa95aaa0aaa2aaa8aaaaaa" - ) - - delta = np.float32(0.125) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d, rest = np.hsplit(blocks, [2]) - qs, qh = np.hsplit(rest, [QK_K // 8]) - - d = d.view(np.float16).astype(np.float32) - qh = qh.view(np.uint16) - - dl = d * (2 * ((qh >> 12) & 7) + 1) - dl = dl.reshape((n_blocks, -1, 1, 1)) - delta = np.where((qh & np.uint16(0x8000)) == 0, cls.delta, -cls.delta) - delta = delta.reshape((n_blocks, -1, 1, 1)) - - qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 3, 6, 9], dtype=np.uint16).reshape((1, 1, 4)) - qs = qs.astype(np.uint16) | ((qh & 7) << 8).reshape((n_blocks, -1)) - - assert cls.grid is not None - grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2) - grid = grid.reshape((n_blocks, -1, 4, 8)) - - return (dl * (grid + delta)).reshape((n_blocks, -1)) - - -class IQ1_M(__Quant, qtype=GGMLQuantizationType.IQ1_M): - grid_shape = IQ1_S.grid_shape - grid_map = IQ1_S.grid_map - grid_hex = IQ1_S.grid_hex - - delta = IQ1_S.delta - - # Okay *this* type is weird. It's the only one which stores the f16 scales in multiple parts. - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - qs, rest = np.hsplit(blocks, [QK_K // 8]) - qh, scales = np.hsplit(rest, [QK_K // 16]) - - # The f16 scale is packed across multiple bytes - scales = scales.view(np.uint16) - d = (scales.reshape((n_blocks, 4)) & np.uint16(0xF000)) >> np.array([12, 8, 4, 0], dtype=np.uint16).reshape((1, 4)) - d = d[..., 0] | d[..., 1] | d[..., 2] | d[..., 3] - d = d.view(np.float16).astype(np.float32).reshape((n_blocks, 1)) - - scales = scales.reshape(n_blocks, -1, 1) >> np.array([0, 3, 6, 9], dtype=np.uint16).reshape((1, 1, 4)) - scales = (scales & 0x07).reshape((n_blocks, -1)) - dl = d * (2 * scales + 1) - dl = dl.reshape((n_blocks, -1, 2, 1, 1)) - - qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2)) - qs = qs.astype(np.uint16) | ((qh & 0x07).astype(np.uint16) << 8).reshape((n_blocks, -1)) - - delta = np.where(qh & 0x08 == 0, cls.delta, -cls.delta) - delta = delta.reshape((n_blocks, -1, 2, 2, 1)) - - assert cls.grid is not None - grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2) - grid = grid.reshape((n_blocks, -1, 2, 2, 8)) - - return (dl * (grid + delta)).reshape((n_blocks, -1)) - - -class IQ4_NL(__Quant, qtype=GGMLQuantizationType.IQ4_NL): - kvalues = (-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113) - - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d, qs = np.hsplit(blocks, [2]) - - d = d.view(np.float16).astype(np.float32) - - qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) - - qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1, 1)) - - kvalues = np.array(cls.kvalues, dtype=np.int8).reshape(1, 1, 16) - qs = np.take_along_axis(kvalues, qs, axis=-1).astype(np.float32).reshape((n_blocks, -1)) - - return (d * qs) - - -class IQ4_XS(__Quant, qtype=GGMLQuantizationType.IQ4_XS): - @classmethod - def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: - n_blocks = blocks.shape[0] - - d, rest = np.hsplit(blocks, [2]) - scales_h, rest = np.hsplit(rest, [2]) - scales_l, qs = np.hsplit(rest, [QK_K // 64]) - - d = d.view(np.float16).astype(np.float32) - scales_h = scales_h.view(np.uint16) - - scales_l = scales_l.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2)) - scales_h = scales_h.reshape((n_blocks, 1, -1)) >> np.array([2 * i for i in range(QK_K // 32)], dtype=np.uint16).reshape((1, -1, 1)) - scales_l = scales_l.reshape((n_blocks, -1)) & np.uint8(0x0F) - scales_h = scales_h.reshape((n_blocks, -1)).astype(np.uint8) & np.uint8(0x03) - - scales = (scales_l | (scales_h << np.uint8(4))).astype(np.int8) - np.int8(32) - dl = (d * scales.astype(np.float32)).reshape((n_blocks, -1, 1)) - - qs = qs.reshape((n_blocks, -1, 1, 16)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1)) - qs = qs.reshape((n_blocks, -1, 32, 1)) & np.uint8(0x0F) - - kvalues = np.array(IQ4_NL.kvalues, dtype=np.int8).reshape((1, 1, 1, -1)) - qs = np.take_along_axis(kvalues, qs, axis=-1).astype(np.float32).reshape((n_blocks, -1, 32)) - - return (dl * qs).reshape((n_blocks, -1)) diff --git a/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_convert_endian.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_convert_endian.cpython-310.pyc deleted file mode 100644 index c838a8224124af5a86e71a3882cc28c4cc220650..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_convert_endian.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_dump.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_dump.cpython-310.pyc deleted file mode 100644 index 47835b43e03d1aae30f35079f083814724b154e1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_dump.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_editor_gui.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_editor_gui.cpython-310.pyc deleted file mode 100644 index bf28e52f2ed73dc8696afbbc570eec2392ed9774..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_editor_gui.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_hash.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_hash.cpython-310.pyc deleted file mode 100644 index 4d7856368336b0a06817bd058d8c755c6daa67e5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_hash.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_new_metadata.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_new_metadata.cpython-310.pyc deleted file mode 100644 index c9013d9dc44b7679f60ef4c7769150ce7d2c9ea8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_new_metadata.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_set_metadata.cpython-310.pyc b/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_set_metadata.cpython-310.pyc deleted file mode 100644 index a8d41e04fe5365260d3a95d5e2dd371b0e0685f3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/gguf/scripts/__pycache__/gguf_set_metadata.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/gguf/scripts/gguf_convert_endian.py b/venv/lib/python3.10/site-packages/gguf/scripts/gguf_convert_endian.py deleted file mode 100644 index 0e0febaa791786fb9d5c713478377e5d54642df8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/scripts/gguf_convert_endian.py +++ /dev/null @@ -1,182 +0,0 @@ -#!/usr/bin/env python3 -from __future__ import annotations - -import logging -import argparse -import os -import sys -from tqdm import tqdm -from pathlib import Path - -import numpy as np - -# Necessary to load the local gguf package -if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent.parent / 'gguf-py').exists(): - sys.path.insert(0, str(Path(__file__).parent.parent.parent)) - -import gguf - -logger = logging.getLogger("gguf-convert-endian") - - -def convert_byteorder(reader: gguf.GGUFReader, args: argparse.Namespace) -> None: - file_endian = reader.endianess.name - if reader.byte_order == 'S': - host_endian = 'BIG' if file_endian == 'LITTLE' else 'LITTLE' - else: - host_endian = file_endian - order = host_endian if args.order == "native" else args.order.upper() - logger.info(f"* Host is {host_endian} endian, GGUF file seems to be {file_endian} endian") - if file_endian == order: - logger.info(f"* File is already {order} endian. Nothing to do.") - sys.exit(0) - logger.info("* Checking tensors for conversion compatibility") - for tensor in reader.tensors: - if tensor.tensor_type not in ( - gguf.GGMLQuantizationType.F32, - gguf.GGMLQuantizationType.F16, - gguf.GGMLQuantizationType.Q8_0, - gguf.GGMLQuantizationType.Q4_K, - gguf.GGMLQuantizationType.Q6_K, - ): - raise ValueError(f"Cannot handle type {tensor.tensor_type.name} for tensor {repr(tensor.name)}") - logger.info(f"* Preparing to convert from {file_endian} to {order}") - if args.dry_run: - return - logger.warning("*** Warning *** Warning *** Warning **") - logger.warning("* This conversion process may damage the file. Ensure you have a backup.") - if order != host_endian: - logger.warning("* Requested endian differs from host, you will not be able to load the model on this machine.") - logger.warning("* The file will be modified immediately, so if conversion fails or is interrupted") - logger.warning("* the file will be corrupted. Enter exactly YES if you are positive you want to proceed:") - response = input("YES, I am sure> ") - if response != "YES": - logger.warning("You didn't enter YES. Okay then, see ya!") - sys.exit(0) - logger.info(f"* Converting fields ({len(reader.fields)})") - for idx, field in enumerate(reader.fields.values()): - logger.info(f"- {idx:4}: Converting field {repr(field.name)}, part count: {len(field.parts)}") - for part in field.parts: - part.byteswap(inplace=True) - logger.info(f"* Converting tensors ({len(reader.tensors)})") - - for idx, tensor in enumerate(pbar := tqdm(reader.tensors, desc="Converting tensor")): - log_message = ( - f"Converting tensor {repr(tensor.name)}, " - f"type={tensor.tensor_type.name}, " - f"elements={tensor.n_elements} " - ) - - # Byte-swap each part of the tensor's field - for part in tensor.field.parts: - part.byteswap(inplace=True) - - # Byte-swap tensor data if necessary - if tensor.tensor_type == gguf.GGMLQuantizationType.Q8_0: - # Handle Q8_0 tensor blocks (block_q8_0) - # Specific handling of block_q8_0 is required. - # Each block_q8_0 consists of an f16 delta (scaling factor) followed by 32 int8 quantizations. - - block_size = 34 # 34 bytes = + 32 * - - n_blocks = len(tensor.data) // block_size - for block_num in (inner_pbar := tqdm(range(n_blocks), desc="Byte-swapping Blocks", leave=False)): - block_offs = block_num * block_size - - # Byte-Swap f16 sized delta field - delta = tensor.data[block_offs:block_offs + 2].view(dtype=np.uint16) - delta.byteswap(inplace=True) - - # Byte-Swap Q8 weights - if block_num % 100000 == 0: - inner_pbar.set_description(f"Byte-swapping Blocks [{(n_blocks - block_num) // n_blocks}]") - - elif tensor.tensor_type == gguf.GGMLQuantizationType.Q4_K: - # Handle Q4_K tensor blocks (block_q4_k) - # Specific handling of block_q4_k is required. - # Each block_q4_k consists of 2 f16 values followed by 140 int8 values. - - # first flatten structure - newshape = 1 - for i in tensor.data.shape: - newshape *= i - - tensor.data.resize(newshape) - - block_size = 144 - n_blocks = len(tensor.data) // block_size - for block_num in (inner_pbar := tqdm(range(n_blocks), desc="Byte-swapping Blocks", leave=False)): - block_offs = block_num * block_size - - # Byte-Swap f16 sized fields - delta = tensor.data[block_offs:block_offs + 2].view(dtype=np.uint16) - delta.byteswap(inplace=True) - - delta = tensor.data[block_offs + 2:block_offs + 4].view(dtype=np.uint16) - delta.byteswap(inplace=True) - - # Byte-Swap - if block_num % 100000 == 0: - inner_pbar.set_description(f"Byte-swapping Blocks [{(n_blocks - block_num) // n_blocks}]") - - elif tensor.tensor_type == gguf.GGMLQuantizationType.Q6_K: - # Handle Q6_K tensor blocks (block_q6_k) - # Specific handling of block_q6_k is required. - # Each block_q6_k consists of 208 int8 values followed by 1 f16 value. - - # first flatten structure - newshape = 1 - for i in tensor.data.shape: - newshape *= i - - tensor.data.resize(newshape) - - block_size = 210 - n_blocks = len(tensor.data) // block_size - for block_num in (inner_pbar := tqdm(range(n_blocks), desc="Byte-swapping Blocks", leave=False)): - block_offs = block_num * block_size - - # Byte-Swap f16 sized field - delta = tensor.data[block_offs + 208:block_offs + 210].view(dtype=np.uint16) - delta.byteswap(inplace=True) - - # Byte-Swap - if block_num % 100000 == 0: - inner_pbar.set_description(f"Byte-swapping Blocks [{(n_blocks - block_num) // n_blocks}]") - - else: - # Handle other tensor types - tensor.data.byteswap(inplace=True) - - pbar.set_description(log_message) - - logger.info("* Completion") - - -def main() -> None: - parser = argparse.ArgumentParser(description="Convert GGUF file byte order") - parser.add_argument( - "model", type=str, - help="GGUF format model filename", - ) - parser.add_argument( - "order", type=str, choices=['big', 'little', 'native'], - help="Requested byte order", - ) - parser.add_argument( - "--dry-run", action="store_true", - help="Don't actually change anything", - ) - parser.add_argument("--verbose", action="store_true", help="increase output verbosity") - - args = parser.parse_args(None if len(sys.argv) > 1 else ["--help"]) - - logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) - - logger.info(f'* Loading: {args.model}') - reader = gguf.GGUFReader(args.model, 'r' if args.dry_run else 'r+') - convert_byteorder(reader, args) - - -if __name__ == "__main__": - main() diff --git a/venv/lib/python3.10/site-packages/gguf/scripts/gguf_dump.py b/venv/lib/python3.10/site-packages/gguf/scripts/gguf_dump.py deleted file mode 100644 index e282892d645c7c31a1d5e00f383dcaaa5890a127..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/scripts/gguf_dump.py +++ /dev/null @@ -1,454 +0,0 @@ -#!/usr/bin/env python3 -from __future__ import annotations - -import logging -import argparse -import os -import re -import sys -from pathlib import Path -from typing import Any - -# Necessary to load the local gguf package -if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent.parent / 'gguf-py').exists(): - sys.path.insert(0, str(Path(__file__).parent.parent.parent)) - -from gguf import GGUFReader, GGUFValueType, ReaderTensor # noqa: E402 - -logger = logging.getLogger("gguf-dump") - - -def get_file_host_endian(reader: GGUFReader) -> tuple[str, str]: - file_endian = reader.endianess.name - if reader.byte_order == 'S': - host_endian = 'BIG' if file_endian == 'LITTLE' else 'LITTLE' - else: - host_endian = file_endian - return (host_endian, file_endian) - - -# For more information about what field.parts and field.data represent, -# please see the comments in the modify_gguf.py example. -def dump_metadata(reader: GGUFReader, args: argparse.Namespace) -> None: - host_endian, file_endian = get_file_host_endian(reader) - print(f'* File is {file_endian} endian, script is running on a {host_endian} endian host.') # noqa: NP100 - print(f'* Dumping {len(reader.fields)} key/value pair(s)') # noqa: NP100 - for n, field in enumerate(reader.fields.values(), 1): - if not field.types: - pretty_type = 'N/A' - elif field.types[0] == GGUFValueType.ARRAY: - nest_count = len(field.types) - 1 - pretty_type = '[' * nest_count + str(field.types[-1].name) + ']' * nest_count - else: - pretty_type = str(field.types[-1].name) - - log_message = f' {n:5}: {pretty_type:10} | {len(field.data):8} | {field.name}' - if field.types: - curr_type = field.types[0] - if curr_type == GGUFValueType.STRING: - content = field.contents() - if len(content) > 60: - content = content[:57] + '...' - log_message += ' = {0}'.format(repr(content)) - elif curr_type in reader.gguf_scalar_to_np: - log_message += ' = {0}'.format(field.contents()) - else: - content = repr(field.contents(slice(6))) - if len(field.data) > 6: - content = content[:-1] + ', ...]' - log_message += ' = {0}'.format(content) - print(log_message) # noqa: NP100 - if args.no_tensors: - return - print(f'* Dumping {len(reader.tensors)} tensor(s)') # noqa: NP100 - for n, tensor in enumerate(reader.tensors, 1): - prettydims = ', '.join('{0:5}'.format(d) for d in list(tensor.shape) + [1] * (4 - len(tensor.shape))) - print(f' {n:5}: {tensor.n_elements:10} | {prettydims} | {tensor.tensor_type.name:7} | {tensor.name}') # noqa: NP100 - - -def dump_metadata_json(reader: GGUFReader, args: argparse.Namespace) -> None: - import json - host_endian, file_endian = get_file_host_endian(reader) - metadata: dict[str, Any] = {} - tensors: dict[str, Any] = {} - result = { - "filename": args.model, - "endian": file_endian, - "metadata": metadata, - "tensors": tensors, - } - for idx, field in enumerate(reader.fields.values()): - curr: dict[str, Any] = { - "index": idx, - "type": field.types[0].name if field.types else 'UNKNOWN', - "offset": field.offset, - } - metadata[field.name] = curr - if field.types[:1] == [GGUFValueType.ARRAY]: - curr["array_types"] = [t.name for t in field.types][1:] - if not args.json_array: - continue - curr["value"] = field.contents() - else: - curr["value"] = field.contents() - if not args.no_tensors: - for idx, tensor in enumerate(reader.tensors): - tensors[tensor.name] = { - "index": idx, - "shape": tensor.shape.tolist(), - "type": tensor.tensor_type.name, - "offset": tensor.field.offset, - } - json.dump(result, sys.stdout) - - -def markdown_table_with_alignment_support(header_map: list[dict[str, str]], data: list[dict[str, Any]]): - # JSON to Markdown table formatting: https://stackoverflow.com/a/72983854/2850957 - - # Alignment Utility Function - def strAlign(padding: int, alignMode: str | None, strVal: str): - if alignMode == 'center': - return strVal.center(padding) - elif alignMode == 'right': - return strVal.rjust(padding - 1) + ' ' - elif alignMode == 'left': - return ' ' + strVal.ljust(padding - 1) - else: # default left - return ' ' + strVal.ljust(padding - 1) - - def dashAlign(padding: int, alignMode: str | None): - if alignMode == 'center': - return ':' + '-' * (padding - 2) + ':' - elif alignMode == 'right': - return '-' * (padding - 1) + ':' - elif alignMode == 'left': - return ':' + '-' * (padding - 1) - else: # default left - return '-' * (padding) - - # Calculate Padding For Each Column Based On Header and Data Length - rowsPadding = {} - for index, columnEntry in enumerate(header_map): - padCount = max([len(str(v)) for d in data for k, v in d.items() if k == columnEntry['key_name']], default=0) + 2 - headerPadCount = len(columnEntry['header_name']) + 2 - rowsPadding[index] = headerPadCount if padCount <= headerPadCount else padCount - - # Render Markdown Header - rows = [] - rows.append('|'.join(strAlign(rowsPadding[index], columnEntry.get('align'), str(columnEntry['header_name'])) for index, columnEntry in enumerate(header_map))) - rows.append('|'.join(dashAlign(rowsPadding[index], columnEntry.get('align')) for index, columnEntry in enumerate(header_map))) - - # Render Tabular Data - for item in data: - rows.append('|'.join(strAlign(rowsPadding[index], columnEntry.get('align'), str(item[columnEntry['key_name']])) for index, columnEntry in enumerate(header_map))) - - # Convert Tabular String Rows Into String - tableString = "" - for row in rows: - tableString += f'|{row}|\n' - - return tableString - - -def element_count_rounded_notation(count: int) -> str: - if count > 1e15 : - # Quadrillion - scaled_amount = count * 1e-15 - scale_suffix = "Q" - elif count > 1e12 : - # Trillions - scaled_amount = count * 1e-12 - scale_suffix = "T" - elif count > 1e9 : - # Billions - scaled_amount = count * 1e-9 - scale_suffix = "B" - elif count > 1e6 : - # Millions - scaled_amount = count * 1e-6 - scale_suffix = "M" - elif count > 1e3 : - # Thousands - scaled_amount = count * 1e-3 - scale_suffix = "K" - else: - # Under Thousands - scaled_amount = count - scale_suffix = "" - return f"{'~' if count > 1e3 else ''}{round(scaled_amount)}{scale_suffix}" - - -def translate_tensor_name(name): - words = name.split(".") - - # Source: https://github.com/ggml-org/ggml/blob/master/docs/gguf.md#standardized-tensor-names - abbreviation_dictionary = { - 'token_embd': 'Token embedding', - 'pos_embd': 'Position embedding', - 'output_norm': 'Output normalization', - 'output': 'Output', - 'attn_norm': 'Attention normalization', - 'attn_norm_2': 'Attention normalization', - 'attn_qkv': 'Attention query-key-value', - 'attn_q': 'Attention query', - 'attn_k': 'Attention key', - 'attn_v': 'Attention value', - 'attn_output': 'Attention output', - 'ffn_norm': 'Feed-forward network normalization', - 'ffn_up': 'Feed-forward network "up"', - 'ffn_gate': 'Feed-forward network "gate"', - 'ffn_down': 'Feed-forward network "down"', - 'ffn_gate_inp': 'Expert-routing layer for the Feed-forward network in Mixture of Expert models', - 'ffn_gate_exp': 'Feed-forward network "gate" layer per expert in Mixture of Expert models', - 'ffn_down_exp': 'Feed-forward network "down" layer per expert in Mixture of Expert models', - 'ffn_up_exp': 'Feed-forward network "up" layer per expert in Mixture of Expert models', - 'ssm_in': 'State space model input projections', - 'ssm_conv1d': 'State space model rolling/shift', - 'ssm_x': 'State space model selective parametrization', - 'ssm_a': 'State space model state compression', - 'ssm_d': 'State space model skip connection', - 'ssm_dt': 'State space model time step', - 'ssm_out': 'State space model output projection', - 'blk': 'Block', - 'enc': 'Encoder', - 'dec': 'Decoder', - } - - expanded_words = [] - for word in words: - word_norm = word.strip().lower() - if word_norm in abbreviation_dictionary: - expanded_words.append(abbreviation_dictionary[word_norm].title()) - else: - expanded_words.append(word.title()) - - return ' '.join(expanded_words) - - -def dump_markdown_metadata(reader: GGUFReader, args: argparse.Namespace) -> None: - host_endian, file_endian = get_file_host_endian(reader) - markdown_content = "" - markdown_content += f'# {args.model} - GGUF Internal File Dump\n\n' - markdown_content += f'- Endian: {file_endian} endian\n' - markdown_content += '\n' - markdown_content += '## Key Value Metadata Store\n\n' - markdown_content += f'There are {len(reader.fields)} key-value pairs in this file\n' - markdown_content += '\n' - - kv_dump_table: list[dict[str, str | int]] = [] - for n, field in enumerate(reader.fields.values(), 1): - if not field.types: - pretty_type = 'N/A' - elif field.types[0] == GGUFValueType.ARRAY: - nest_count = len(field.types) - 1 - pretty_type = '[' * nest_count + str(field.types[-1].name) + ']' * nest_count - else: - pretty_type = str(field.types[-1].name) - - def escape_markdown_inline_code(value_string): - # Find the longest contiguous sequence of backticks in the string then - # wrap string with appropriate number of backticks required to escape it - max_backticks = max((len(match.group(0)) for match in re.finditer(r'`+', value_string)), default=0) - inline_code_marker = '`' * (max_backticks + 1) - - # If the string starts or ends with a backtick, add a space at the beginning and end - if value_string.startswith('`') or value_string.endswith('`'): - value_string = f" {value_string} " - - return f"{inline_code_marker}{value_string}{inline_code_marker}" - - total_elements = len(field.data) - value = "" - if len(field.types) == 1: - curr_type = field.types[0] - if curr_type == GGUFValueType.STRING: - truncate_length = 60 - value_string = str(bytes(field.parts[-1]), encoding='utf-8') - if len(value_string) > truncate_length: - head = escape_markdown_inline_code(value_string[:truncate_length // 2]) - tail = escape_markdown_inline_code(value_string[-truncate_length // 2:]) - value = "{head}...{tail}".format(head=head, tail=tail) - else: - value = escape_markdown_inline_code(value_string) - elif curr_type in reader.gguf_scalar_to_np: - value = str(field.parts[-1][0]) - else: - if field.types[0] == GGUFValueType.ARRAY: - curr_type = field.types[1] - array_elements = [] - - if curr_type == GGUFValueType.STRING: - render_element = min(5, total_elements) - for element_pos in range(render_element): - truncate_length = 30 - value_string = str(bytes(field.parts[-1 - (total_elements - element_pos - 1) * 2]), encoding='utf-8') - if len(value_string) > truncate_length: - head = escape_markdown_inline_code(value_string[:truncate_length // 2]) - tail = escape_markdown_inline_code(value_string[-truncate_length // 2:]) - value = "{head}...{tail}".format(head=head, tail=tail) - else: - value = escape_markdown_inline_code(value_string) - array_elements.append(value) - - elif curr_type in reader.gguf_scalar_to_np: - render_element = min(7, total_elements) - for element_pos in range(render_element): - array_elements.append(str(field.parts[-1 - (total_elements - element_pos - 1)][0])) - - value = f'[ {", ".join(array_elements).strip()}{", ..." if total_elements > len(array_elements) else ""} ]' - - kv_dump_table.append({"n":n, "pretty_type":pretty_type, "total_elements":total_elements, "field_name":field.name, "value":value}) - - kv_dump_table_header_map = [ - {'key_name':'n', 'header_name':'POS', 'align':'right'}, - {'key_name':'pretty_type', 'header_name':'TYPE', 'align':'left'}, - {'key_name':'total_elements', 'header_name':'Count', 'align':'right'}, - {'key_name':'field_name', 'header_name':'Key', 'align':'left'}, - {'key_name':'value', 'header_name':'Value', 'align':'left'}, - ] - - markdown_content += markdown_table_with_alignment_support(kv_dump_table_header_map, kv_dump_table) - - markdown_content += "\n" - - if not args.no_tensors: - # Group tensors by their prefix and maintain order - tensor_prefix_order: list[str] = [] - tensor_name_to_key: dict[str, int] = {} - tensor_groups: dict[str, list[ReaderTensor]] = {} - total_elements = sum(tensor.n_elements for tensor in reader.tensors) - - # Parsing Tensors Record - for key, tensor in enumerate(reader.tensors): - tensor_components = tensor.name.split('.') - - # Classify Tensor Group - tensor_group_name = "base" - if tensor_components[0] == 'blk': - tensor_group_name = f"{tensor_components[0]}.{tensor_components[1]}" - elif tensor_components[0] in ['enc', 'dec'] and tensor_components[1] == 'blk': - tensor_group_name = f"{tensor_components[0]}.{tensor_components[1]}.{tensor_components[2]}" - elif tensor_components[0] in ['enc', 'dec']: - tensor_group_name = f"{tensor_components[0]}" - - # Check if new Tensor Group - if tensor_group_name not in tensor_groups: - tensor_groups[tensor_group_name] = [] - tensor_prefix_order.append(tensor_group_name) - - # Record Tensor and Tensor Position - tensor_groups[tensor_group_name].append(tensor) - tensor_name_to_key[tensor.name] = key - - # Tensors Mapping Dump - markdown_content += f'## Tensors Overview {element_count_rounded_notation(total_elements)} Elements\n\n' - markdown_content += f'Total number of elements in all tensors: {total_elements} Elements\n' - markdown_content += '\n' - - for group in tensor_prefix_order: - tensors = tensor_groups[group] - group_elements = sum(tensor.n_elements for tensor in tensors) - markdown_content += f"- [{translate_tensor_name(group)} Tensor Group - {element_count_rounded_notation(group_elements)} Elements](#{group.replace('.', '_')})\n" - - markdown_content += "\n" - - markdown_content += "### Tensor Data Offset\n" - markdown_content += '\n' - markdown_content += 'This table contains the offset and data segment relative to start of file\n' - markdown_content += '\n' - - tensor_mapping_table: list[dict[str, str | int]] = [] - for key, tensor in enumerate(reader.tensors): - data_offset_pretty = '{0:#16x}'.format(tensor.data_offset) - data_size_pretty = '{0:#16x}'.format(tensor.n_bytes) - tensor_mapping_table.append({"t_id":key, "layer_name":tensor.name, "data_offset":data_offset_pretty, "data_size":data_size_pretty}) - - tensors_mapping_table_header_map = [ - {'key_name':'t_id', 'header_name':'T_ID', 'align':'right'}, - {'key_name':'layer_name', 'header_name':'Tensor Layer Name', 'align':'left'}, - {'key_name':'data_offset', 'header_name':'Data Offset (B)', 'align':'right'}, - {'key_name':'data_size', 'header_name':'Data Size (B)', 'align':'right'}, - ] - - markdown_content += markdown_table_with_alignment_support(tensors_mapping_table_header_map, tensor_mapping_table) - markdown_content += "\n" - - for group in tensor_prefix_order: - tensors = tensor_groups[group] - group_elements = sum(tensor.n_elements for tensor in tensors) - group_percentage = group_elements / total_elements * 100 - markdown_content += f"### {translate_tensor_name(group)} Tensor Group : {element_count_rounded_notation(group_elements)} Elements\n\n" - - # Precalculate column sizing for visual consistency - prettify_element_est_count_size: int = 1 - prettify_element_count_size: int = 1 - prettify_dimension_max_widths: dict[int, int] = {} - for tensor in tensors: - prettify_element_est_count_size = max(prettify_element_est_count_size, len(str(element_count_rounded_notation(tensor.n_elements)))) - prettify_element_count_size = max(prettify_element_count_size, len(str(tensor.n_elements))) - for i, dimension_size in enumerate(list(tensor.shape) + [1] * (4 - len(tensor.shape))): - prettify_dimension_max_widths[i] = max(prettify_dimension_max_widths.get(i,1), len(str(dimension_size))) - - # Generate Tensor Layer Table Content - tensor_dump_table: list[dict[str, str | int]] = [] - for tensor in tensors: - human_friendly_name = translate_tensor_name(tensor.name.replace(".weight", ".(W)").replace(".bias", ".(B)")) - pretty_dimension = ' x '.join(f'{str(d):>{prettify_dimension_max_widths[i]}}' for i, d in enumerate(list(tensor.shape) + [1] * (4 - len(tensor.shape)))) - element_count_est = f"({element_count_rounded_notation(tensor.n_elements):>{prettify_element_est_count_size}})" - element_count_string = f"{element_count_est} {tensor.n_elements:>{prettify_element_count_size}}" - type_name_string = f"{tensor.tensor_type.name}" - tensor_dump_table.append({"t_id":tensor_name_to_key[tensor.name], "layer_name":tensor.name, "human_layer_name":human_friendly_name, "element_count":element_count_string, "pretty_dimension":pretty_dimension, "tensor_type":type_name_string}) - - tensor_dump_table_header_map = [ - {'key_name':'t_id', 'header_name':'T_ID', 'align':'right'}, - {'key_name':'layer_name', 'header_name':'Tensor Layer Name', 'align':'left'}, - {'key_name':'human_layer_name', 'header_name':'Human Friendly Tensor Layer Name', 'align':'left'}, - {'key_name':'element_count', 'header_name':'Elements', 'align':'left'}, - {'key_name':'pretty_dimension', 'header_name':'Shape', 'align':'left'}, - {'key_name':'tensor_type', 'header_name':'Type', 'align':'left'}, - ] - - markdown_content += markdown_table_with_alignment_support(tensor_dump_table_header_map, tensor_dump_table) - - markdown_content += "\n" - markdown_content += f"- Total elements in {group}: ({element_count_rounded_notation(group_elements):>4}) {group_elements}\n" - markdown_content += f"- Percentage of total elements: {group_percentage:.2f}%\n" - markdown_content += "\n\n" - - print(markdown_content) # noqa: NP100 - - -def main() -> None: - parser = argparse.ArgumentParser(description="Dump GGUF file metadata") - parser.add_argument("model", type=str, help="GGUF format model filename") - parser.add_argument("--no-tensors", action="store_true", help="Don't dump tensor metadata") - parser.add_argument("--json", action="store_true", help="Produce JSON output") - parser.add_argument("--json-array", action="store_true", help="Include full array values in JSON output (long)") - parser.add_argument("--data-offset", action="store_true", help="Start of data offset") - parser.add_argument("--data-alignment", action="store_true", help="Data alignment applied globally to data field") - parser.add_argument("--markdown", action="store_true", help="Produce markdown output") - parser.add_argument("--verbose", action="store_true", help="increase output verbosity") - - args = parser.parse_args(None if len(sys.argv) > 1 else ["--help"]) - - logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) - - if not args.json and not args.markdown and not args.data_offset and not args.data_alignment: - logger.info(f'* Loading: {args.model}') - - reader = GGUFReader(args.model, 'r') - - if args.json: - dump_metadata_json(reader, args) - elif args.markdown: - dump_markdown_metadata(reader, args) - elif args.data_offset: - print(reader.data_offset) # noqa: NP100 - elif args.data_alignment: - print(reader.alignment) # noqa: NP100 - else: - dump_metadata(reader, args) - - -if __name__ == '__main__': - main() diff --git a/venv/lib/python3.10/site-packages/gguf/scripts/gguf_editor_gui.py b/venv/lib/python3.10/site-packages/gguf/scripts/gguf_editor_gui.py deleted file mode 100644 index 05f4db0f8cdc89049a7e8719fcaa1fe741391413..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/scripts/gguf_editor_gui.py +++ /dev/null @@ -1,1621 +0,0 @@ -#!/usr/bin/env python3 -from __future__ import annotations - -import logging -import argparse -import os -import sys -import numpy -import enum -from pathlib import Path -from typing import Any, Optional, Tuple, Type -import warnings - -import numpy as np -from PySide6.QtWidgets import ( - QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, - QPushButton, QLabel, QLineEdit, QFileDialog, QTableWidget, - QTableWidgetItem, QComboBox, QMessageBox, QTabWidget, - QTextEdit, QFormLayout, - QHeaderView, QDialog, QDialogButtonBox -) -from PySide6.QtCore import Qt - -# Necessary to load the local gguf package -if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent.parent / 'gguf-py').exists(): - sys.path.insert(0, str(Path(__file__).parent.parent.parent)) - -import gguf -from gguf import GGUFReader, GGUFWriter, GGUFValueType, ReaderField -from gguf.constants import TokenType, RopeScalingType, PoolingType, GGMLQuantizationType - -logger = logging.getLogger("gguf-editor-gui") - -# Map of key names to enum types for automatic enum interpretation -KEY_TO_ENUM_TYPE = { - gguf.Keys.Tokenizer.TOKEN_TYPE: TokenType, - gguf.Keys.Rope.SCALING_TYPE: RopeScalingType, - gguf.Keys.LLM.POOLING_TYPE: PoolingType, - gguf.Keys.General.FILE_TYPE: GGMLQuantizationType, -} - -# Define the tokenizer keys that should be edited together -TOKENIZER_LINKED_KEYS = [ - gguf.Keys.Tokenizer.LIST, - gguf.Keys.Tokenizer.TOKEN_TYPE, - gguf.Keys.Tokenizer.SCORES -] - - -class TokenizerEditorDialog(QDialog): - def __init__(self, tokens, token_types, scores, parent=None): - super().__init__(parent) - self.setWindowTitle("Edit Tokenizer Data") - self.resize(900, 600) - - self.tokens = tokens.copy() if tokens else [] - self.token_types = token_types.copy() if token_types else [] - self.scores = scores.copy() if scores else [] - - # Ensure all arrays have the same length - max_len = max(len(self.tokens), len(self.token_types), len(self.scores)) - if len(self.tokens) < max_len: - self.tokens.extend([""] * (max_len - len(self.tokens))) - if len(self.token_types) < max_len: - self.token_types.extend([0] * (max_len - len(self.token_types))) - if len(self.scores) < max_len: - self.scores.extend([0.0] * (max_len - len(self.scores))) - - layout = QVBoxLayout(self) - - # Add filter controls - filter_layout = QHBoxLayout() - filter_layout.addWidget(QLabel("Filter:")) - self.filter_edit = QLineEdit() - self.filter_edit.setPlaceholderText("Type to filter tokens...") - self.filter_edit.textChanged.connect(self.apply_filter) - filter_layout.addWidget(self.filter_edit) - - # Add page controls - self.page_size = 100 # Show 100 items per page - self.current_page = 0 - self.total_pages = max(1, (len(self.tokens) + self.page_size - 1) // self.page_size) - - self.page_label = QLabel(f"Page 1 of {self.total_pages}") - filter_layout.addWidget(self.page_label) - - prev_page = QPushButton("Previous") - prev_page.clicked.connect(self.previous_page) - filter_layout.addWidget(prev_page) - - next_page = QPushButton("Next") - next_page.clicked.connect(self.next_page) - filter_layout.addWidget(next_page) - - layout.addLayout(filter_layout) - - # Tokenizer data table - self.tokens_table = QTableWidget() - self.tokens_table.setColumnCount(4) - self.tokens_table.setHorizontalHeaderLabels(["Index", "Token", "Type", "Score"]) - self.tokens_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.ResizeToContents) - self.tokens_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch) - self.tokens_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.ResizeToContents) - self.tokens_table.horizontalHeader().setSectionResizeMode(3, QHeaderView.ResizeMode.ResizeToContents) - - layout.addWidget(self.tokens_table) - - # Controls - controls_layout = QHBoxLayout() - - add_button = QPushButton("Add Token") - add_button.clicked.connect(self.add_token) - controls_layout.addWidget(add_button) - - remove_button = QPushButton("Remove Selected") - remove_button.clicked.connect(self.remove_selected) - controls_layout.addWidget(remove_button) - - controls_layout.addStretch() - - layout.addLayout(controls_layout) - - # Buttons - buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel) - buttons.accepted.connect(self.accept) - buttons.rejected.connect(self.reject) - layout.addWidget(buttons) - - # Initialize the filtered values - self.filtered_indices = list(range(len(self.tokens))) - - # Load data for the first page - self.load_page() - - def apply_filter(self): - """Filter the tokens based on the search text.""" - filter_text = self.filter_edit.text().lower() - - if not filter_text: - # No filter, show all values - self.filtered_indices = list(range(len(self.tokens))) - else: - # Apply filter - self.filtered_indices = [] - for i, token in enumerate(self.tokens): - if filter_text in str(token).lower(): - self.filtered_indices.append(i) - - # Reset to first page and reload - self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size) - self.current_page = 0 - self.page_label.setText(f"Page 1 of {self.total_pages}") - self.load_page() - - def previous_page(self): - """Go to the previous page of results.""" - if self.current_page > 0: - self.current_page -= 1 - self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}") - self.load_page() - - def next_page(self): - """Go to the next page of results.""" - if self.current_page < self.total_pages - 1: - self.current_page += 1 - self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}") - self.load_page() - - def load_page(self): - """Load the current page of tokenizer data.""" - self.tokens_table.setRowCount(0) # Clear the table - - # Calculate start and end indices for the current page - start_idx = self.current_page * self.page_size - end_idx = min(start_idx + self.page_size, len(self.filtered_indices)) - - # Pre-allocate rows for better performance - self.tokens_table.setRowCount(end_idx - start_idx) - - for row, i in enumerate(range(start_idx, end_idx)): - orig_idx = self.filtered_indices[i] - - # Index - index_item = QTableWidgetItem(str(orig_idx)) - index_item.setData(Qt.ItemDataRole.UserRole, orig_idx) # Store original index - index_item.setFlags(index_item.flags() & ~Qt.ItemFlag.ItemIsEditable) - self.tokens_table.setItem(row, 0, index_item) - - # Token - token_item = QTableWidgetItem(str(self.tokens[orig_idx])) - self.tokens_table.setItem(row, 1, token_item) - - # Token Type - token_type = self.token_types[orig_idx] if orig_idx < len(self.token_types) else 0 - try: - enum_val = TokenType(token_type) - display_text = f"{enum_val.name} ({token_type})" - except (ValueError, KeyError): - display_text = f"Unknown ({token_type})" - - type_item = QTableWidgetItem(display_text) - type_item.setData(Qt.ItemDataRole.UserRole, token_type) - - # Make type cell editable with a double-click handler - type_item.setFlags(type_item.flags() & ~Qt.ItemFlag.ItemIsEditable) - self.tokens_table.setItem(row, 2, type_item) - - # Score - score = self.scores[orig_idx] if orig_idx < len(self.scores) else 0.0 - score_item = QTableWidgetItem(str(score)) - self.tokens_table.setItem(row, 3, score_item) - - # Connect double-click handler for token type cells - self.tokens_table.cellDoubleClicked.connect(self.handle_cell_double_click) - - def handle_cell_double_click(self, row, column): - """Handle double-click on a cell, specifically for token type editing.""" - if column == 2: # Token Type column - orig_item = self.tokens_table.item(row, 0) - if orig_item: - orig_idx = orig_item.data(Qt.ItemDataRole.UserRole) - self.edit_token_type(row, orig_idx) - - def edit_token_type(self, row, orig_idx): - """Edit a token type using a dialog with a dropdown of all enum options.""" - current_value = self.token_types[orig_idx] if orig_idx < len(self.token_types) else 0 - - # Create a dialog with enum options - dialog = QDialog(self) - dialog.setWindowTitle("Select Token Type") - layout = QVBoxLayout(dialog) - - combo = QComboBox() - for enum_val in TokenType: - combo.addItem(f"{enum_val.name} ({enum_val.value})", enum_val.value) - - # Set current value - try: - if isinstance(current_value, int): - enum_val = TokenType(current_value) - combo.setCurrentText(f"{enum_val.name} ({current_value})") - except (ValueError, KeyError): - pass - - layout.addWidget(combo) - - buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel) - buttons.accepted.connect(dialog.accept) - buttons.rejected.connect(dialog.reject) - layout.addWidget(buttons) - - if dialog.exec() == QDialog.DialogCode.Accepted: - # Get the selected value - new_value = combo.currentData() - enum_val = TokenType(new_value) - display_text = f"{enum_val.name} ({new_value})" - - # Update the display - type_item = self.tokens_table.item(row, 2) - if type_item: - type_item.setText(display_text) - type_item.setData(Qt.ItemDataRole.UserRole, new_value) - - # Update the actual value - self.token_types[orig_idx] = new_value - - def add_token(self): - """Add a new token to the end of the list.""" - # Add to the end of the arrays - self.tokens.append("") - self.token_types.append(0) # Default to normal token - self.scores.append(0.0) - - orig_idx = len(self.tokens) - 1 - - # Add to filtered indices if it matches the current filter - filter_text = self.filter_edit.text().lower() - if not filter_text or filter_text in "": - self.filtered_indices.append(orig_idx) - - # Update pagination - self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size) - - # Go to the last page to show the new item - self.current_page = self.total_pages - 1 - self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}") - - # Reload the page - self.load_page() - - def remove_selected(self): - """Remove selected tokens from all arrays.""" - selected_rows = [] - for item in self.tokens_table.selectedItems(): - row = item.row() - if row not in selected_rows: - selected_rows.append(row) - - if not selected_rows: - return - - # Get original indices in descending order to avoid index shifting - orig_indices = [] - for row in selected_rows: - orig_item = self.tokens_table.item(row, 0) - if orig_item: - orig_indices.append(orig_item.data(Qt.ItemDataRole.UserRole)) - orig_indices.sort(reverse=True) - - # Remove from all arrays - for idx in orig_indices: - if idx < len(self.tokens): - del self.tokens[idx] - if idx < len(self.token_types): - del self.token_types[idx] - if idx < len(self.scores): - del self.scores[idx] - - # Rebuild filtered_indices - self.filtered_indices = [] - filter_text = self.filter_edit.text().lower() - - for i, token in enumerate(self.tokens): - if not filter_text or filter_text in str(token).lower(): - self.filtered_indices.append(i) - - # Update pagination - self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size) - self.current_page = min(self.current_page, self.total_pages - 1) - self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}") - - # Reload the page - self.load_page() - - def get_data(self): - """Return the edited tokenizer data.""" - return self.tokens, self.token_types, self.scores - - -class ArrayEditorDialog(QDialog): - def __init__(self, array_values, element_type, key=None, parent=None): - super().__init__(parent) - self.setWindowTitle("Edit Array Values") - self.resize(700, 500) - - self.array_values = array_values - self.element_type = element_type - self.key = key - - # Get enum type for this array if applicable - self.enum_type = None - if key in KEY_TO_ENUM_TYPE and element_type == GGUFValueType.INT32: - self.enum_type = KEY_TO_ENUM_TYPE[key] - - layout = QVBoxLayout(self) - - # Add enum type information if applicable - if self.enum_type is not None: - enum_info_layout = QHBoxLayout() - enum_label = QLabel(f"Editing {self.enum_type.__name__} values:") - enum_info_layout.addWidget(enum_label) - - # Add a legend for the enum values - enum_values = ", ".join([f"{e.name}={e.value}" for e in self.enum_type]) - enum_values_label = QLabel(f"Available values: {enum_values}") - enum_values_label.setWordWrap(True) - enum_info_layout.addWidget(enum_values_label, 1) - - layout.addLayout(enum_info_layout) - - # Add search/filter controls - filter_layout = QHBoxLayout() - filter_layout.addWidget(QLabel("Filter:")) - self.filter_edit = QLineEdit() - self.filter_edit.setPlaceholderText("Type to filter values...") - self.filter_edit.textChanged.connect(self.apply_filter) - filter_layout.addWidget(self.filter_edit) - - # Add page controls for large arrays - self.page_size = 100 # Show 100 items per page - self.current_page = 0 - self.total_pages = max(1, (len(array_values) + self.page_size - 1) // self.page_size) - - self.page_label = QLabel(f"Page 1 of {self.total_pages}") - filter_layout.addWidget(self.page_label) - - prev_page = QPushButton("Previous") - prev_page.clicked.connect(self.previous_page) - filter_layout.addWidget(prev_page) - - next_page = QPushButton("Next") - next_page.clicked.connect(self.next_page) - filter_layout.addWidget(next_page) - - layout.addLayout(filter_layout) - - # Array items table - self.items_table = QTableWidget() - - # Set up columns based on whether we have an enum type - if self.enum_type is not None: - self.items_table.setColumnCount(3) - self.items_table.setHorizontalHeaderLabels(["Index", "Value", "Actions"]) - self.items_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.ResizeToContents) - self.items_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch) - self.items_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.ResizeToContents) - else: - self.items_table.setColumnCount(2) - self.items_table.setHorizontalHeaderLabels(["Index", "Value"]) - self.items_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.ResizeToContents) - self.items_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch) - - layout.addWidget(self.items_table) - - # Controls - controls_layout = QHBoxLayout() - - add_button = QPushButton("Add Item") - add_button.clicked.connect(self.add_item) - controls_layout.addWidget(add_button) - - remove_button = QPushButton("Remove Selected") - remove_button.clicked.connect(self.remove_selected) - controls_layout.addWidget(remove_button) - - # Add bulk edit button for enum arrays - if self.enum_type is not None: - bulk_edit_button = QPushButton("Bulk Edit Selected") - bulk_edit_button.clicked.connect(self.bulk_edit_selected) - controls_layout.addWidget(bulk_edit_button) - - controls_layout.addStretch() - - layout.addLayout(controls_layout) - - # Buttons - buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel) - buttons.accepted.connect(self.accept) - buttons.rejected.connect(self.reject) - layout.addWidget(buttons) - - # Initialize the filtered values - self.filtered_indices = list(range(len(self.array_values))) - - # Load array values for the first page - self.load_page() - - def apply_filter(self): - """Filter the array values based on the search text.""" - filter_text = self.filter_edit.text().lower() - - if not filter_text: - # No filter, show all values - self.filtered_indices = list(range(len(self.array_values))) - else: - # Apply filter - self.filtered_indices = [] - for i, value in enumerate(self.array_values): - # For enum values, search in both name and value - if self.enum_type is not None and isinstance(value, int): - try: - enum_val = self.enum_type(value) - display_text = f"{enum_val.name} ({value})".lower() - if filter_text in display_text: - self.filtered_indices.append(i) - except (ValueError, KeyError): - # If not a valid enum value, just check the raw value - if filter_text in str(value).lower(): - self.filtered_indices.append(i) - else: - # For non-enum values, just check the string representation - if filter_text in str(value).lower(): - self.filtered_indices.append(i) - - # Reset to first page and reload - self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size) - self.current_page = 0 - self.page_label.setText(f"Page 1 of {self.total_pages}") - self.load_page() - - def previous_page(self): - """Go to the previous page of results.""" - if self.current_page > 0: - self.current_page -= 1 - self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}") - self.load_page() - - def next_page(self): - """Go to the next page of results.""" - if self.current_page < self.total_pages - 1: - self.current_page += 1 - self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}") - self.load_page() - - def load_page(self): - """Load the current page of array values.""" - self.items_table.setRowCount(0) # Clear the table - - # Calculate start and end indices for the current page - start_idx = self.current_page * self.page_size - end_idx = min(start_idx + self.page_size, len(self.filtered_indices)) - - # Pre-allocate rows for better performance - self.items_table.setRowCount(end_idx - start_idx) - - for row, i in enumerate(range(start_idx, end_idx)): - orig_idx = self.filtered_indices[i] - value = self.array_values[orig_idx] - - # Index - index_item = QTableWidgetItem(str(orig_idx)) - index_item.setData(Qt.ItemDataRole.UserRole, orig_idx) # Store original index - index_item.setFlags(index_item.flags() & ~Qt.ItemFlag.ItemIsEditable) - self.items_table.setItem(row, 0, index_item) - - # Value - if self.enum_type is not None: - # Display enum value and name - try: - if isinstance(value, (int, numpy.signedinteger)): - enum_val = self.enum_type(value) - display_text = f"{enum_val.name} ({value})" - else: - display_text = str(value) - except (ValueError, KeyError): - display_text = f"Unknown ({value})" - - # Store the enum value in the item - value_item = QTableWidgetItem(display_text) - value_item.setData(Qt.ItemDataRole.UserRole, value) - value_item.setFlags(value_item.flags() & ~Qt.ItemFlag.ItemIsEditable) - self.items_table.setItem(row, 1, value_item) - - # Add an edit button in a separate column - edit_button = QPushButton("Edit") - edit_button.setProperty("row", row) - edit_button.clicked.connect(self.edit_array_enum_value) - - # Create a widget to hold the button - button_widget = QWidget() - button_layout = QHBoxLayout(button_widget) - button_layout.setContentsMargins(2, 2, 2, 2) - button_layout.addWidget(edit_button) - button_layout.addStretch() - - self.items_table.setCellWidget(row, 2, button_widget) - else: - value_item = QTableWidgetItem(str(value)) - self.items_table.setItem(row, 1, value_item) - - def edit_array_enum_value(self): - """Handle editing an enum value in the array editor.""" - button = self.sender() - row = button.property("row") - - # Get the original index from the table item - orig_item = self.items_table.item(row, 0) - new_item = self.items_table.item(row, 1) - if orig_item and new_item and self.enum_type and self.edit_enum_value(row, self.enum_type): - orig_idx = orig_item.data(Qt.ItemDataRole.UserRole) - new_value = new_item.data(Qt.ItemDataRole.UserRole) - # Update the stored value in the array - if isinstance(new_value, (int, float, str, bool)): - self.array_values[orig_idx] = new_value - - def bulk_edit_selected(self): - """Edit multiple enum values at once.""" - if not self.enum_type: - return - - selected_rows = set() - for item in self.items_table.selectedItems(): - selected_rows.add(item.row()) - - if not selected_rows: - QMessageBox.information(self, "No Selection", "Please select at least one row to edit.") - return - - # Create a dialog with enum options - dialog = QDialog(self) - dialog.setWindowTitle(f"Bulk Edit {self.enum_type.__name__} Values") - layout = QVBoxLayout(dialog) - - layout.addWidget(QLabel(f"Set {len(selected_rows)} selected items to:")) - - combo = QComboBox() - for enum_val in self.enum_type: - combo.addItem(f"{enum_val.name} ({enum_val.value})", enum_val.value) - - layout.addWidget(combo) - - buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel) - buttons.accepted.connect(dialog.accept) - buttons.rejected.connect(dialog.reject) - layout.addWidget(buttons) - - if dialog.exec() == QDialog.DialogCode.Accepted: - # Get the selected value - new_value = combo.currentData() - enum_val = self.enum_type(new_value) - display_text = f"{enum_val.name} ({new_value})" - - # Update all selected rows - for row in selected_rows: - orig_item = self.items_table.item(row, 0) - new_item = self.items_table.item(row, 1) - if orig_item and new_item: - orig_idx = orig_item.data(Qt.ItemDataRole.UserRole) - self.array_values[orig_idx] = new_value - - # Update the display - new_item.setText(display_text) - new_item.setData(Qt.ItemDataRole.UserRole, new_value) - - def add_item(self): - # Add to the end of the array - orig_idx = len(self.array_values) - - # Add default value based on type - if self.enum_type is not None: - # Default to first enum value - default_value = list(self.enum_type)[0].value - self.array_values.append(default_value) - else: - if self.element_type == GGUFValueType.STRING: - self.array_values.append("") - else: - self.array_values.append(0) - - # Add to filtered indices if it matches the current filter - self.filtered_indices.append(orig_idx) - - # Update pagination - self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size) - - # Go to the last page to show the new item - self.current_page = self.total_pages - 1 - self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}") - - # Reload the page - self.load_page() - - def remove_selected(self): - selected_rows = [] - for item in self.items_table.selectedItems(): - row = item.row() - if row not in selected_rows: - selected_rows.append(row) - - if not selected_rows: - return - - # Get original indices in descending order to avoid index shifting - orig_indices = list() - for row in selected_rows: - orig_item = self.items_table.item(row, 0) - if orig_item: - orig_indices.append(orig_item.data(Qt.ItemDataRole.UserRole)) - orig_indices.sort(reverse=True) - - # Remove from array_values - for idx in orig_indices: - del self.array_values[idx] - - # Rebuild filtered_indices - self.filtered_indices = [] - filter_text = self.filter_edit.text().lower() - - for i, value in enumerate(self.array_values): - if not filter_text: - self.filtered_indices.append(i) - else: - # Apply filter - if self.enum_type is not None and isinstance(value, int): - try: - enum_val = self.enum_type(value) - display_text = f"{enum_val.name} ({value})".lower() - if filter_text in display_text: - self.filtered_indices.append(i) - except (ValueError, KeyError): - if filter_text in str(value).lower(): - self.filtered_indices.append(i) - else: - if filter_text in str(value).lower(): - self.filtered_indices.append(i) - - # Update pagination - self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size) - self.current_page = min(self.current_page, self.total_pages - 1) - self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}") - - # Reload the page - self.load_page() - - def edit_enum_value(self, row: int, enum_type: Type[enum.Enum]): - """Edit an enum value using a dialog with a dropdown of all enum options.""" - # Get the original index from the table item - orig_item = self.items_table.item(row, 0) - if orig_item: - orig_idx = orig_item.data(Qt.ItemDataRole.UserRole) - else: - return - current_value = self.array_values[orig_idx] - - # Create a dialog with enum options - dialog = QDialog(self) - dialog.setWindowTitle(f"Select {enum_type.__name__} Value") - layout = QVBoxLayout(dialog) - - # Add description - description = QLabel(f"Select a {enum_type.__name__} value:") - layout.addWidget(description) - - # Use a combo box for quick selection - combo = QComboBox() - for enum_val in enum_type: - combo.addItem(f"{enum_val.name} ({enum_val.value})", enum_val.value) - - # Set current value - try: - if isinstance(current_value, int): - enum_val = enum_type(current_value) - combo.setCurrentText(f"{enum_val.name} ({current_value})") - except (ValueError, KeyError): - pass - - layout.addWidget(combo) - - buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel) - buttons.accepted.connect(dialog.accept) - buttons.rejected.connect(dialog.reject) - layout.addWidget(buttons) - - if dialog.exec() == QDialog.DialogCode.Accepted: - # Update the value display and stored data - new_value = combo.currentData() - enum_val = enum_type(new_value) - display_text = f"{enum_val.name} ({new_value})" - - new_item = self.items_table.item(row, 1) - if new_item: - new_item.setText(display_text) - new_item.setData(Qt.ItemDataRole.UserRole, new_value) - - # Update the actual array value - self.array_values[orig_idx] = new_value - return True - return False - - def get_array_values(self): - # The array_values list is kept up-to-date as edits are made - return self.array_values - - -class AddMetadataDialog(QDialog): - def __init__(self, parent=None): - super().__init__(parent) - self.setWindowTitle("Add Metadata") - self.resize(400, 200) - - layout = QVBoxLayout(self) - - form_layout = QFormLayout() - - self.key_edit = QLineEdit() - form_layout.addRow("Key:", self.key_edit) - - self.type_combo = QComboBox() - for value_type in GGUFValueType: - if value_type != GGUFValueType.ARRAY: # Skip array type for simplicity - self.type_combo.addItem(value_type.name, value_type) - form_layout.addRow("Type:", self.type_combo) - - self.value_edit = QTextEdit() - form_layout.addRow("Value:", self.value_edit) - - layout.addLayout(form_layout) - - buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel) - buttons.accepted.connect(self.accept) - buttons.rejected.connect(self.reject) - layout.addWidget(buttons) - - def get_data(self) -> Tuple[str, GGUFValueType, Any]: - key = self.key_edit.text() - value_type = self.type_combo.currentData() - value_text = self.value_edit.toPlainText() - - # Convert value based on type - if value_type == GGUFValueType.UINT8: - value = np.uint8(int(value_text)) - elif value_type == GGUFValueType.INT8: - value = np.int8(int(value_text)) - elif value_type == GGUFValueType.UINT16: - value = np.uint16(int(value_text)) - elif value_type == GGUFValueType.INT16: - value = np.int16(int(value_text)) - elif value_type == GGUFValueType.UINT32: - value = np.uint32(int(value_text)) - elif value_type == GGUFValueType.INT32: - value = np.int32(int(value_text)) - elif value_type == GGUFValueType.FLOAT32: - value = np.float32(float(value_text)) - elif value_type == GGUFValueType.BOOL: - value = value_text.lower() in ('true', 'yes', '1') - elif value_type == GGUFValueType.STRING: - value = value_text - else: - value = value_text - - return key, value_type, value - - -class GGUFEditorWindow(QMainWindow): - def __init__(self): - super().__init__() - - self.setWindowTitle("GGUF Editor") - self.resize(1000, 800) - - self.current_file = None - self.reader = None - self.modified = False - self.metadata_changes = {} # Store changes to apply when saving - self.metadata_to_remove = set() # Store keys to remove when saving - self.on_metadata_changed_is_connected = False - - self.setup_ui() - - def setup_ui(self): - central_widget = QWidget() - self.setCentralWidget(central_widget) - - main_layout = QVBoxLayout(central_widget) - - # File controls - file_layout = QHBoxLayout() - - self.file_path_edit = QLineEdit() - self.file_path_edit.setReadOnly(True) - file_layout.addWidget(self.file_path_edit) - - open_button = QPushButton("Open GGUF") - open_button.clicked.connect(self.open_file) - file_layout.addWidget(open_button) - - save_button = QPushButton("Save As...") - save_button.clicked.connect(self.save_file) - file_layout.addWidget(save_button) - - main_layout.addLayout(file_layout) - - # Tabs for different views - self.tabs = QTabWidget() - - # Metadata tab - self.metadata_tab = QWidget() - metadata_layout = QVBoxLayout(self.metadata_tab) - - # Metadata table - self.metadata_table = QTableWidget() - self.metadata_table.setColumnCount(4) - self.metadata_table.setHorizontalHeaderLabels(["Key", "Type", "Value", "Actions"]) - self.metadata_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.Stretch) - self.metadata_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.ResizeToContents) - self.metadata_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.Stretch) - self.metadata_table.horizontalHeader().setSectionResizeMode(3, QHeaderView.ResizeMode.ResizeToContents) - metadata_layout.addWidget(self.metadata_table) - - # Metadata controls - metadata_controls = QHBoxLayout() - - add_metadata_button = QPushButton("Add Metadata") - add_metadata_button.clicked.connect(self.add_metadata) - metadata_controls.addWidget(add_metadata_button) - - metadata_controls.addStretch() - - metadata_layout.addLayout(metadata_controls) - - # Tensors tab - self.tensors_tab = QWidget() - tensors_layout = QVBoxLayout(self.tensors_tab) - - self.tensors_table = QTableWidget() - self.tensors_table.setColumnCount(5) - self.tensors_table.setHorizontalHeaderLabels(["Name", "Type", "Shape", "Elements", "Size (bytes)"]) - self.tensors_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.Stretch) - self.tensors_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.ResizeToContents) - self.tensors_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.ResizeToContents) - self.tensors_table.horizontalHeader().setSectionResizeMode(3, QHeaderView.ResizeMode.ResizeToContents) - self.tensors_table.horizontalHeader().setSectionResizeMode(4, QHeaderView.ResizeMode.ResizeToContents) - tensors_layout.addWidget(self.tensors_table) - - # Add tabs to tab widget - self.tabs.addTab(self.metadata_tab, "Metadata") - self.tabs.addTab(self.tensors_tab, "Tensors") - - main_layout.addWidget(self.tabs) - - # Status bar - self.statusBar().showMessage("Ready") - - def load_file(self, file_path): - """Load a GGUF file by path""" - try: - self.statusBar().showMessage(f"Loading {file_path}...") - QApplication.processEvents() - - self.reader = GGUFReader(file_path, 'r') - self.current_file = file_path - self.file_path_edit.setText(file_path) - - self.load_metadata() - self.load_tensors() - - self.metadata_changes = {} - self.metadata_to_remove = set() - self.modified = False - - self.statusBar().showMessage(f"Loaded {file_path}") - return True - except Exception as e: - QMessageBox.critical(self, "Error", f"Failed to open file: {str(e)}") - self.statusBar().showMessage("Error loading file") - return False - - def open_file(self): - file_path, _ = QFileDialog.getOpenFileName( - self, "Open GGUF File", "", "GGUF Files (*.gguf);;All Files (*)" - ) - - if not file_path: - return - - self.load_file(file_path) - - def load_metadata(self): - self.metadata_table.setRowCount(0) - - if not self.reader: - return - - # Disconnect to prevent triggering during loading - if self.on_metadata_changed_is_connected: - with warnings.catch_warnings(): - warnings.filterwarnings('ignore') - self.metadata_table.itemChanged.disconnect(self.on_metadata_changed) - self.on_metadata_changed_is_connected = False - - for i, (key, field) in enumerate(self.reader.fields.items()): - self.metadata_table.insertRow(i) - - # Key - key_item = QTableWidgetItem(key) - key_item.setFlags(key_item.flags() & ~Qt.ItemFlag.ItemIsEditable) - self.metadata_table.setItem(i, 0, key_item) - - # Type - if not field.types: - type_str = "N/A" - elif field.types[0] == GGUFValueType.ARRAY: - nest_count = len(field.types) - 1 - element_type = field.types[-1].name - # Check if this is an enum array - enum_type = self.get_enum_for_key(key) - if enum_type is not None and field.types[-1] == GGUFValueType.INT32: - element_type = enum_type.__name__ - type_str = '[' * nest_count + element_type + ']' * nest_count - else: - type_str = str(field.types[0].name) - # Check if this is an enum field - enum_type = self.get_enum_for_key(key) - if enum_type is not None and field.types[0] == GGUFValueType.INT32: - type_str = enum_type.__name__ - - type_item = QTableWidgetItem(type_str) - type_item.setFlags(type_item.flags() & ~Qt.ItemFlag.ItemIsEditable) - self.metadata_table.setItem(i, 1, type_item) - - # Value - value_str = self.format_field_value(field) - value_item = QTableWidgetItem(value_str) - - # Make only simple values editable - if len(field.types) == 1 and field.types[0] != GGUFValueType.ARRAY: - value_item.setFlags(value_item.flags() | Qt.ItemFlag.ItemIsEditable) - else: - value_item.setFlags(value_item.flags() & ~Qt.ItemFlag.ItemIsEditable) - - self.metadata_table.setItem(i, 2, value_item) - - # Actions - actions_widget = QWidget() - actions_layout = QHBoxLayout(actions_widget) - actions_layout.setContentsMargins(2, 2, 2, 2) - - # Add Edit button for arrays and enum fields - if field.types and field.types[0] == GGUFValueType.ARRAY: - edit_button = QPushButton("Edit") - edit_button.setProperty("row", i) - edit_button.setProperty("key", key) - edit_button.clicked.connect(self.edit_array_metadata) - actions_layout.addWidget(edit_button) - - # Add special label for tokenizer linked fields - if key in TOKENIZER_LINKED_KEYS: - edit_button.setText("Edit Tokenizer") - edit_button.setToolTip("Edit all tokenizer data together") - elif len(field.types) == 1 and self.get_enum_for_key(key) is not None: - edit_button = QPushButton("Edit") - edit_button.setProperty("row", i) - edit_button.setProperty("key", key) - edit_button.clicked.connect(self.edit_metadata_enum) - actions_layout.addWidget(edit_button) - - remove_button = QPushButton("Remove") - remove_button.setProperty("row", i) - remove_button.setProperty("key", key) - remove_button.clicked.connect(self.remove_metadata) - actions_layout.addWidget(remove_button) - - self.metadata_table.setCellWidget(i, 3, actions_widget) - - # Reconnect after loading - self.metadata_table.itemChanged.connect(self.on_metadata_changed) - self.on_metadata_changed_is_connected = True - - def extract_array_values(self, field: ReaderField) -> list: - """Extract all values from an array field.""" - if not field.types or field.types[0] != GGUFValueType.ARRAY: - return [] - - curr_type = field.types[1] - array_values = [] - total_elements = len(field.data) - - if curr_type == GGUFValueType.STRING: - for element_pos in range(total_elements): - value_string = str(bytes(field.parts[-1 - (total_elements - element_pos - 1) * 2]), encoding='utf-8') - array_values.append(value_string) - elif self.reader and curr_type in self.reader.gguf_scalar_to_np: - for element_pos in range(total_elements): - array_values.append(field.parts[-1 - (total_elements - element_pos - 1)][0]) - - return array_values - - def get_enum_for_key(self, key: str) -> Optional[Type[enum.Enum]]: - """Get the enum type for a given key if it exists.""" - return KEY_TO_ENUM_TYPE.get(key) - - def format_enum_value(self, value: Any, enum_type: Type[enum.Enum]) -> str: - """Format a value as an enum if possible.""" - try: - if isinstance(value, (int, str)): - enum_value = enum_type(value) - return f"{enum_value.name} ({value})" - except (ValueError, KeyError): - pass - return str(value) - - def format_field_value(self, field: ReaderField) -> str: - if not field.types: - return "N/A" - - if len(field.types) == 1: - curr_type = field.types[0] - if curr_type == GGUFValueType.STRING: - return str(bytes(field.parts[-1]), encoding='utf-8') - elif self.reader and curr_type in self.reader.gguf_scalar_to_np: - value = field.parts[-1][0] - # Check if this field has an enum type - enum_type = self.get_enum_for_key(field.name) - if enum_type is not None: - return self.format_enum_value(value, enum_type) - return str(value) - - if field.types[0] == GGUFValueType.ARRAY: - array_values = self.extract_array_values(field) - render_element = min(5, len(array_values)) - - # Get enum type for this array if applicable - enum_type = self.get_enum_for_key(field.name) - - if enum_type is not None: - array_elements = [] - for i in range(render_element): - array_elements.append(self.format_enum_value(array_values[i], enum_type)) - else: - array_elements = [str(array_values[i]) for i in range(render_element)] - - return f"[ {', '.join(array_elements).strip()}{', ...' if len(array_values) > len(array_elements) else ''} ]" - - return "Complex value" - - def load_tensors(self): - self.tensors_table.setRowCount(0) - - if not self.reader: - return - - for i, tensor in enumerate(self.reader.tensors): - self.tensors_table.insertRow(i) - - # Name - name_item = QTableWidgetItem(tensor.name) - name_item.setFlags(name_item.flags() & ~Qt.ItemFlag.ItemIsEditable) - self.tensors_table.setItem(i, 0, name_item) - - # Type - type_item = QTableWidgetItem(tensor.tensor_type.name) - type_item.setFlags(type_item.flags() & ~Qt.ItemFlag.ItemIsEditable) - self.tensors_table.setItem(i, 1, type_item) - - # Shape - shape_str = " × ".join(str(d) for d in tensor.shape) - shape_item = QTableWidgetItem(shape_str) - shape_item.setFlags(shape_item.flags() & ~Qt.ItemFlag.ItemIsEditable) - self.tensors_table.setItem(i, 2, shape_item) - - # Elements - elements_item = QTableWidgetItem(str(tensor.n_elements)) - elements_item.setFlags(elements_item.flags() & ~Qt.ItemFlag.ItemIsEditable) - self.tensors_table.setItem(i, 3, elements_item) - - # Size - size_item = QTableWidgetItem(f"{tensor.n_bytes:,}") - size_item.setFlags(size_item.flags() & ~Qt.ItemFlag.ItemIsEditable) - self.tensors_table.setItem(i, 4, size_item) - - def on_metadata_changed(self, item): - if item.column() != 2: # Only handle value column changes - return - - row = item.row() - orig_item = self.metadata_table.item(row, 0) - key = None - if orig_item: - key = orig_item.text() - new_value = item.text() - - field = None - if self.reader and key: - field = self.reader.get_field(key) - if not field or not field.types or not key: - return - - value_type = field.types[0] - - # Check if this is an enum field - enum_type = self.get_enum_for_key(key) - if enum_type is not None and value_type == GGUFValueType.INT32: - # Try to parse the enum value from the text - try: - # Check if it's a name - try: - enum_val = enum_type[new_value] - converted_value = enum_val.value - except (KeyError, AttributeError): - # Check if it's a number or "NAME (value)" format - if '(' in new_value and ')' in new_value: - # Extract the value from "NAME (value)" format - value_part = new_value.split('(')[1].split(')')[0].strip() - converted_value = int(value_part) - else: - # Try to convert directly to int - converted_value = int(new_value) - - # Validate that it's a valid enum value - enum_type(converted_value) - - # Store the change - self.metadata_changes[key] = (value_type, converted_value) - self.modified = True - - # Update display with formatted enum value - formatted_value = self.format_enum_value(converted_value, enum_type) - item.setText(formatted_value) - - self.statusBar().showMessage(f"Changed {key} to {formatted_value}") - return - except (ValueError, KeyError) as e: - QMessageBox.warning( - self, - f"Invalid Enum Value ({e})", - f"'{new_value}' is not a valid {enum_type.__name__} value.\n" - f"Valid values are: {', '.join(v.name for v in enum_type)}") - - # Revert to original value - original_value = self.format_field_value(field) - item.setText(original_value) - return - - try: - # Convert the string value to the appropriate type - if value_type == GGUFValueType.UINT8: - converted_value = np.uint8(int(new_value)) - elif value_type == GGUFValueType.INT8: - converted_value = np.int8(int(new_value)) - elif value_type == GGUFValueType.UINT16: - converted_value = np.uint16(int(new_value)) - elif value_type == GGUFValueType.INT16: - converted_value = np.int16(int(new_value)) - elif value_type == GGUFValueType.UINT32: - converted_value = np.uint32(int(new_value)) - elif value_type == GGUFValueType.INT32: - converted_value = np.int32(int(new_value)) - elif value_type == GGUFValueType.FLOAT32: - converted_value = np.float32(float(new_value)) - elif value_type == GGUFValueType.BOOL: - converted_value = new_value.lower() in ('true', 'yes', '1') - elif value_type == GGUFValueType.STRING: - converted_value = new_value - else: - # Unsupported type for editing - return - - # Store the change - self.metadata_changes[key] = (value_type, converted_value) - self.modified = True - - self.statusBar().showMessage(f"Changed {key} to {new_value}") - except ValueError: - QMessageBox.warning(self, "Invalid Value", f"The value '{new_value}' is not valid for type {value_type.name}") - - # Revert to original value - original_value = self.format_field_value(field) - item.setText(original_value) - - def remove_metadata(self): - button = self.sender() - key = button.property("key") - row = button.property("row") - - reply = QMessageBox.question( - self, "Confirm Removal", - f"Are you sure you want to remove the metadata key '{key}'?", - QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, QMessageBox.StandardButton.No - ) - - if reply == QMessageBox.StandardButton.Yes: - self.metadata_table.removeRow(row) - self.metadata_to_remove.add(key) - - # If we previously had changes for this key, remove them - if key in self.metadata_changes: - del self.metadata_changes[key] - - self.modified = True - self.statusBar().showMessage(f"Marked {key} for removal") - - def edit_metadata_enum(self): - """Edit an enum metadata field.""" - button = self.sender() - key = button.property("key") - row = button.property("row") - - field = None - if self.reader: - field = self.reader.get_field(key) - if not field or not field.types: - return - - enum_type = self.get_enum_for_key(key) - if enum_type is None: - return - - # Get current value - current_value = field.contents() - - # Create a dialog with enum options - dialog = QDialog(self) - dialog.setWindowTitle(f"Select {enum_type.__name__} Value") - layout = QVBoxLayout(dialog) - - combo = QComboBox() - for enum_val in enum_type: - combo.addItem(f"{enum_val.name} ({enum_val.value})", enum_val.value) - - # Set current value - try: - if isinstance(current_value, (int, str)): - enum_val = enum_type(current_value) - combo.setCurrentText(f"{enum_val.name} ({current_value})") - except (ValueError, KeyError): - pass - - layout.addWidget(combo) - - buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel) - buttons.accepted.connect(dialog.accept) - buttons.rejected.connect(dialog.reject) - layout.addWidget(buttons) - - if dialog.exec() == QDialog.DialogCode.Accepted: - # Get the selected value - new_value = combo.currentData() - enum_val = enum_type(new_value) - - # Store the change - self.metadata_changes[key] = (field.types[0], new_value) - self.modified = True - - # Update display - display_text = f"{enum_val.name} ({new_value})" - target_item = self.metadata_table.item(row, 2) - if target_item: - target_item.setText(display_text) - - self.statusBar().showMessage(f"Changed {key} to {display_text}") - - def edit_array_metadata(self): - button = self.sender() - key = button.property("key") - row = button.property("row") - - # Check if this is one of the linked tokenizer keys - if key in TOKENIZER_LINKED_KEYS: - self.edit_tokenizer_metadata(key) - return - - field = None - if self.reader: - field = self.reader.get_field(key) - if not field or not field.types or field.types[0] != GGUFValueType.ARRAY: - return - - # Get array element type - element_type = field.types[1] - - # Extract array values - array_values = self.extract_array_values(field) - - # Open array editor dialog - dialog = ArrayEditorDialog(array_values, element_type, key, self) - if dialog.exec() == QDialog.DialogCode.Accepted: - new_values = dialog.get_array_values() - - # Store the change - self.metadata_changes[key] = (GGUFValueType.ARRAY, (element_type, new_values)) - self.modified = True - - # Update display - enum_type = self.get_enum_for_key(key) - if enum_type is not None and element_type == GGUFValueType.INT32: - value_str = f"[ {', '.join(self.format_enum_value(v, enum_type) for v in new_values[:5])}{', ...' if len(new_values) > 5 else ''} ]" - else: - value_str = f"[ {', '.join(str(v) for v in new_values[:5])}{', ...' if len(new_values) > 5 else ''} ]" - target_item = self.metadata_table.item(row, 2) - if target_item: - target_item.setText(value_str) - - self.statusBar().showMessage(f"Updated array values for {key}") - - def edit_tokenizer_metadata(self, trigger_key): - """Edit the linked tokenizer metadata arrays together.""" - if not self.reader: - return - - # Get all three fields - tokens_field = self.reader.get_field(gguf.Keys.Tokenizer.LIST) - token_types_field = self.reader.get_field(gguf.Keys.Tokenizer.TOKEN_TYPE) - scores_field = self.reader.get_field(gguf.Keys.Tokenizer.SCORES) - - # Extract values from each field - tokens = self.extract_array_values(tokens_field) if tokens_field else [] - token_types = self.extract_array_values(token_types_field) if token_types_field else [] - scores = self.extract_array_values(scores_field) if scores_field else [] - - # Apply any pending changes - if gguf.Keys.Tokenizer.LIST in self.metadata_changes: - _, (_, tokens) = self.metadata_changes[gguf.Keys.Tokenizer.LIST] - if gguf.Keys.Tokenizer.TOKEN_TYPE in self.metadata_changes: - _, (_, token_types) = self.metadata_changes[gguf.Keys.Tokenizer.TOKEN_TYPE] - if gguf.Keys.Tokenizer.SCORES in self.metadata_changes: - _, (_, scores) = self.metadata_changes[gguf.Keys.Tokenizer.SCORES] - - # Open the tokenizer editor dialog - dialog = TokenizerEditorDialog(tokens, token_types, scores, self) - if dialog.exec() == QDialog.DialogCode.Accepted: - new_tokens, new_token_types, new_scores = dialog.get_data() - - # Store changes for all three arrays - if tokens_field: - self.metadata_changes[gguf.Keys.Tokenizer.LIST] = ( - GGUFValueType.ARRAY, - (tokens_field.types[1], new_tokens) - ) - - if token_types_field: - self.metadata_changes[gguf.Keys.Tokenizer.TOKEN_TYPE] = ( - GGUFValueType.ARRAY, - (token_types_field.types[1], new_token_types) - ) - - if scores_field: - self.metadata_changes[gguf.Keys.Tokenizer.SCORES] = ( - GGUFValueType.ARRAY, - (scores_field.types[1], new_scores) - ) - - self.modified = True - - # Update display for all three fields - self.update_tokenizer_display(gguf.Keys.Tokenizer.LIST, new_tokens) - self.update_tokenizer_display(gguf.Keys.Tokenizer.TOKEN_TYPE, new_token_types) - self.update_tokenizer_display(gguf.Keys.Tokenizer.SCORES, new_scores) - - self.statusBar().showMessage("Updated tokenizer data") - - def update_tokenizer_display(self, key, values): - """Update the display of a tokenizer field in the metadata table.""" - for row in range(self.metadata_table.rowCount()): - key_item = self.metadata_table.item(row, 0) - if key_item and key_item.text() == key: - value_str = f"[ {', '.join(str(v) for v in values[:5])}{', ...' if len(values) > 5 else ''} ]" - value_item = self.metadata_table.item(row, 2) - if value_item: - value_item.setText(value_str) - break - - def add_metadata(self): - dialog = AddMetadataDialog(self) - if dialog.exec() == QDialog.DialogCode.Accepted: - key, value_type, value = dialog.get_data() - - if not key: - QMessageBox.warning(self, "Invalid Key", "Key cannot be empty") - return - - # Check if key already exists - for row in range(self.metadata_table.rowCount()): - orig_item = self.metadata_table.item(row, 0) - if orig_item and orig_item.text() == key: - QMessageBox.warning(self, "Duplicate Key", f"Key '{key}' already exists") - return - - # Add to table - row = self.metadata_table.rowCount() - self.metadata_table.insertRow(row) - - # Key - key_item = QTableWidgetItem(key) - key_item.setFlags(key_item.flags() & ~Qt.ItemFlag.ItemIsEditable) - self.metadata_table.setItem(row, 0, key_item) - - # Type - type_item = QTableWidgetItem(value_type.name) - type_item.setFlags(type_item.flags() & ~Qt.ItemFlag.ItemIsEditable) - self.metadata_table.setItem(row, 1, type_item) - - # Value - value_item = QTableWidgetItem(str(value)) - value_item.setFlags(value_item.flags() | Qt.ItemFlag.ItemIsEditable) - self.metadata_table.setItem(row, 2, value_item) - - # Actions - actions_widget = QWidget() - actions_layout = QHBoxLayout(actions_widget) - actions_layout.setContentsMargins(2, 2, 2, 2) - - remove_button = QPushButton("Remove") - remove_button.setProperty("row", row) - remove_button.setProperty("key", key) - remove_button.clicked.connect(self.remove_metadata) - actions_layout.addWidget(remove_button) - - self.metadata_table.setCellWidget(row, 3, actions_widget) - - # Store the change - self.metadata_changes[key] = (value_type, value) - self.modified = True - - self.statusBar().showMessage(f"Added new metadata key {key}") - - def save_file(self): - if not self.reader: - QMessageBox.warning(self, "No File Open", "Please open a GGUF file first") - return - - if not self.modified and not self.metadata_changes and not self.metadata_to_remove: - QMessageBox.information(self, "No Changes", "No changes to save") - return - - file_path, _ = QFileDialog.getSaveFileName( - self, "Save GGUF File As", "", "GGUF Files (*.gguf);;All Files (*)" - ) - - if not file_path: - return - - try: - self.statusBar().showMessage(f"Saving to {file_path}...") - QApplication.processEvents() - - # Get architecture and endianness from the original file - arch = 'unknown' - field = self.reader.get_field(gguf.Keys.General.ARCHITECTURE) - if field: - arch = field.contents() - - # Create writer - writer = GGUFWriter(file_path, arch=arch, endianess=self.reader.endianess) - - # Get alignment if present - alignment = None - field = self.reader.get_field(gguf.Keys.General.ALIGNMENT) - if field: - alignment = field.contents() - if alignment is not None: - writer.data_alignment = alignment - - # Copy metadata with changes - for field in self.reader.fields.values(): - # Skip virtual fields and fields written by GGUFWriter - if field.name == gguf.Keys.General.ARCHITECTURE or field.name.startswith('GGUF.'): - continue - - # Skip fields marked for removal - if field.name in self.metadata_to_remove: - continue - - # Apply changes if any - sub_type = None - if field.name in self.metadata_changes: - value_type, value = self.metadata_changes[field.name] - if value_type == GGUFValueType.ARRAY: - # Handle array values - sub_type, value = value - else: - # Copy original value - value = field.contents() - value_type = field.types[0] - if value_type == GGUFValueType.ARRAY: - sub_type = field.types[-1] - - if value is not None: - writer.add_key_value(field.name, value, value_type, sub_type=sub_type) - - # Add new metadata - for key, (value_type, value) in self.metadata_changes.items(): - # Skip if the key already existed (we handled it above) - if self.reader.get_field(key) is not None: - continue - - sub_type = None - if value_type == GGUFValueType.ARRAY: - # Handle array values - sub_type, value = value - - writer.add_key_value(key, value, value_type, sub_type=sub_type) - - # Add tensors (including data) - for tensor in self.reader.tensors: - writer.add_tensor(tensor.name, tensor.data, raw_shape=tensor.data.shape, raw_dtype=tensor.tensor_type) - - # Write header and metadata - writer.open_output_file(Path(file_path)) - writer.write_header_to_file() - writer.write_kv_data_to_file() - - # Write tensor data using the optimized method - writer.write_tensors_to_file(progress=False) - - writer.close() - - self.statusBar().showMessage(f"Saved to {file_path}") - - # Ask if user wants to open the new file - reply = QMessageBox.question( - self, "Open Saved File", - "Would you like to open the newly saved file?", - QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, QMessageBox.StandardButton.Yes - ) - - if reply == QMessageBox.StandardButton.Yes: - self.reader = GGUFReader(file_path, 'r') - self.current_file = file_path - self.file_path_edit.setText(file_path) - - self.load_metadata() - self.load_tensors() - - self.metadata_changes = {} - self.metadata_to_remove = set() - self.modified = False - - except Exception as e: - QMessageBox.critical(self, "Error", f"Failed to save file: {str(e)}") - self.statusBar().showMessage("Error saving file") - - -def main() -> None: - parser = argparse.ArgumentParser(description="GUI GGUF Editor") - parser.add_argument("model_path", nargs="?", help="path to GGUF model file to load at startup") - parser.add_argument("--verbose", action="store_true", help="increase output verbosity") - - args = parser.parse_args() - - logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) - - app = QApplication(sys.argv) - window = GGUFEditorWindow() - window.show() - - # Load model if specified - if args.model_path: - if os.path.isfile(args.model_path) and args.model_path.endswith('.gguf'): - window.load_file(args.model_path) - else: - logger.error(f"Invalid model path: {args.model_path}") - QMessageBox.warning( - window, - "Invalid Model Path", - f"The specified file does not exist or is not a GGUF file: {args.model_path}") - - sys.exit(app.exec()) - - -if __name__ == '__main__': - main() diff --git a/venv/lib/python3.10/site-packages/gguf/scripts/gguf_hash.py b/venv/lib/python3.10/site-packages/gguf/scripts/gguf_hash.py deleted file mode 100644 index 3ef98992197e987d8aefc8fb5f5bf9c787e9d06a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/scripts/gguf_hash.py +++ /dev/null @@ -1,102 +0,0 @@ -#!/usr/bin/env python3 -from __future__ import annotations - -import uuid -import hashlib - -import logging -import argparse -import os -import sys -from pathlib import Path - -from tqdm import tqdm - -# Necessary to load the local gguf package -if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent.parent / 'gguf-py').exists(): - sys.path.insert(0, str(Path(__file__).parent.parent.parent)) - -from gguf import GGUFReader # noqa: E402 - - -logger = logging.getLogger("gguf-hash") - -# UUID_NAMESPACE_LLAMA_CPP = uuid.uuid5(uuid.NAMESPACE_URL, 'en.wikipedia.org/wiki/Llama.cpp') -UUID_NAMESPACE_LLAMA_CPP = uuid.UUID('ef001206-dadc-5f6d-a15f-3359e577d4e5') - - -# For more information about what field.parts and field.data represent, -# please see the comments in the modify_gguf.py example. -def gguf_hash(reader: GGUFReader, filename: str, disable_progress_bar: bool, no_layer: bool) -> None: - sha1 = hashlib.sha1() - sha256 = hashlib.sha256() - uuidv5_sha1 = hashlib.sha1() - uuidv5_sha1.update(UUID_NAMESPACE_LLAMA_CPP.bytes) - - # Total Weight Calculation For Progress Bar - total_weights = 0 - for n, tensor in enumerate(reader.tensors, 1): - - # We don't need these - if tensor.name.endswith((".attention.masked_bias", ".attention.bias", ".rotary_emb.inv_freq")): - continue - - # Calculate Tensor Volume - sum_weights_in_tensor = 1 - for dim in tensor.shape: - sum_weights_in_tensor *= dim - total_weights += sum_weights_in_tensor - - # Hash Progress Bar - bar = tqdm(desc="Hashing", total=total_weights, unit="weights", unit_scale=True, disable=disable_progress_bar) - - # Hashing Process - for tensor in reader.tensors: - - # We don't need these - if tensor.name.endswith((".attention.masked_bias", ".attention.bias", ".rotary_emb.inv_freq")): - continue - - # Progressbar - sum_weights_in_tensor = 1 - for dim in tensor.shape: - sum_weights_in_tensor *= dim - bar.update(sum_weights_in_tensor) - - if not no_layer: - - sha1_layer = hashlib.sha1() - sha1_layer.update(tensor.data.data) - print("sha1 {0} {1}:{2}".format(sha1_layer.hexdigest(), filename, tensor.name)) # noqa: NP100 - - sha256_layer = hashlib.sha256() - sha256_layer.update(tensor.data.data) - print("sha256 {0} {1}:{2}".format(sha256_layer.hexdigest(), filename, tensor.name)) # noqa: NP100 - - sha1.update(tensor.data.data) - sha256.update(tensor.data.data) - uuidv5_sha1.update(tensor.data.data) - - # Flush Hash Progress Bar - bar.close() - - # Display Hash Output - print("sha1 {0} {1}".format(sha1.hexdigest(), filename)) # noqa: NP100 - print("sha256 {0} {1}".format(sha256.hexdigest(), filename)) # noqa: NP100 - print("uuid {0} {1}".format(uuid.UUID(bytes=uuidv5_sha1.digest()[:16], version=5), filename)) # noqa: NP100 - - -def main() -> None: - parser = argparse.ArgumentParser(description="Dump GGUF file metadata") - parser.add_argument("model", type=str, help="GGUF format model filename") - parser.add_argument("--no-layer", action="store_true", help="exclude per layer hash") - parser.add_argument("--verbose", action="store_true", help="increase output verbosity") - parser.add_argument("--progressbar", action="store_true", help="enable progressbar") - args = parser.parse_args(None if len(sys.argv) > 1 else ["--help"]) - logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) - reader = GGUFReader(args.model, 'r') - gguf_hash(reader, args.model, not args.progressbar, args.no_layer) - - -if __name__ == '__main__': - main() diff --git a/venv/lib/python3.10/site-packages/gguf/scripts/gguf_new_metadata.py b/venv/lib/python3.10/site-packages/gguf/scripts/gguf_new_metadata.py deleted file mode 100644 index 63f2300348ed0ffad3bb9a4be6aa7aae6b827350..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/scripts/gguf_new_metadata.py +++ /dev/null @@ -1,210 +0,0 @@ -#!/usr/bin/env python3 -from __future__ import annotations - -import logging -import argparse -import os -import sys -import json -from pathlib import Path - -from tqdm import tqdm -from typing import Any, Sequence, NamedTuple - -# Necessary to load the local gguf package -if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent.parent / 'gguf-py').exists(): - sys.path.insert(0, str(Path(__file__).parent.parent.parent)) - -import gguf - -logger = logging.getLogger("gguf-new-metadata") - - -class MetadataDetails(NamedTuple): - type: gguf.GGUFValueType - value: Any - description: str = '' - sub_type: gguf.GGUFValueType | None = None - - -def get_field_data(reader: gguf.GGUFReader, key: str) -> Any: - field = reader.get_field(key) - - return field.contents() if field else None - - -def find_token(token_list: Sequence[int], token: str) -> Sequence[int]: - token_ids = [index for index, value in enumerate(token_list) if value == token] - - if len(token_ids) == 0: - raise LookupError(f'Unable to find "{token}" in token list!') - - return token_ids - - -def copy_with_new_metadata(reader: gguf.GGUFReader, writer: gguf.GGUFWriter, new_metadata: dict[str, MetadataDetails], remove_metadata: Sequence[str]) -> None: - for field in reader.fields.values(): - # Suppress virtual fields and fields written by GGUFWriter - if field.name == gguf.Keys.General.ARCHITECTURE or field.name.startswith('GGUF.'): - logger.debug(f'Suppressing {field.name}') - continue - - # Skip old chat templates if we have new ones - if field.name.startswith(gguf.Keys.Tokenizer.CHAT_TEMPLATE) and gguf.Keys.Tokenizer.CHAT_TEMPLATE in new_metadata: - logger.debug(f'Skipping {field.name}') - continue - - if field.name in remove_metadata: - logger.debug(f'Removing {field.name}') - continue - - val_type = field.types[0] - sub_type = field.types[-1] if val_type == gguf.GGUFValueType.ARRAY else None - old_val = MetadataDetails(val_type, field.contents(), sub_type=sub_type) - val = new_metadata.get(field.name, old_val) - - if field.name in new_metadata: - logger.debug(f'Modifying {field.name}: "{old_val.value}" -> "{val.value}" {val.description}') - del new_metadata[field.name] - elif val.value is not None: - logger.debug(f'Copying {field.name}') - - if val.value is not None: - writer.add_key_value(field.name, val.value, val.type, sub_type=sub_type if val.sub_type is None else val.sub_type) - - if gguf.Keys.Tokenizer.CHAT_TEMPLATE in new_metadata: - logger.debug('Adding chat template(s)') - writer.add_chat_template(new_metadata[gguf.Keys.Tokenizer.CHAT_TEMPLATE].value) - del new_metadata[gguf.Keys.Tokenizer.CHAT_TEMPLATE] - - for key, val in new_metadata.items(): - logger.debug(f'Adding {key}: "{val.value}" {val.description}') - writer.add_key_value(key, val.value, val.type) - - total_bytes = 0 - - for tensor in reader.tensors: - total_bytes += tensor.n_bytes - writer.add_tensor_info(tensor.name, tensor.data.shape, tensor.data.dtype, tensor.data.nbytes, tensor.tensor_type) - - bar = tqdm(desc="Writing", total=total_bytes, unit="byte", unit_scale=True) - - writer.write_header_to_file() - writer.write_kv_data_to_file() - writer.write_ti_data_to_file() - - for tensor in reader.tensors: - writer.write_tensor_data(tensor.data) - bar.update(tensor.n_bytes) - - writer.close() - - -def main() -> None: - tokenizer_metadata = (getattr(gguf.Keys.Tokenizer, n) for n in gguf.Keys.Tokenizer.__dict__.keys() if not n.startswith('_')) - token_names = dict((n.split('.')[-1][:-len('_token_id')], n) for n in tokenizer_metadata if n.endswith('_token_id')) - - parser = argparse.ArgumentParser(description="Make a copy of a GGUF file with new metadata") - parser.add_argument("input", type=Path, help="GGUF format model input filename") - parser.add_argument("output", type=Path, help="GGUF format model output filename") - parser.add_argument("--general-name", type=str, help="The models general.name", metavar='"name"') - parser.add_argument("--general-description", type=str, help="The models general.description", metavar='"Description ..."') - parser.add_argument("--chat-template", type=str, help="Chat template string (or JSON string containing templates)", metavar='"{% ... %} ..."') - parser.add_argument("--chat-template-config", type=Path, help="Config file containing chat template(s)", metavar='tokenizer_config.json') - parser.add_argument("--pre-tokenizer", type=str, help="The models tokenizer.ggml.pre", metavar='"pre tokenizer"') - parser.add_argument("--remove-metadata", action="append", type=str, help="Remove metadata (by key name) from output model", metavar='general.url') - parser.add_argument("--special-token", action="append", type=str, help="Special token by value", nargs=2, metavar=(' | '.join(token_names.keys()), '""')) - parser.add_argument("--special-token-by-id", action="append", type=str, help="Special token by id", nargs=2, metavar=(' | '.join(token_names.keys()), '0')) - parser.add_argument("--force", action="store_true", help="Bypass warnings without confirmation") - parser.add_argument("--verbose", action="store_true", help="Increase output verbosity") - args = parser.parse_args(None if len(sys.argv) > 2 else ["--help"]) - - logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) - - new_metadata = {} - remove_metadata = args.remove_metadata or [] - - if args.general_name: - new_metadata[gguf.Keys.General.NAME] = MetadataDetails(gguf.GGUFValueType.STRING, args.general_name) - - if args.general_description: - new_metadata[gguf.Keys.General.DESCRIPTION] = MetadataDetails(gguf.GGUFValueType.STRING, args.general_description) - - if args.chat_template: - new_metadata[gguf.Keys.Tokenizer.CHAT_TEMPLATE] = MetadataDetails(gguf.GGUFValueType.STRING, json.loads(args.chat_template) if args.chat_template.startswith('[') else args.chat_template) - - if args.chat_template_config: - with open(args.chat_template_config, 'r') as fp: - config = json.load(fp) - template = config.get('chat_template') - if template: - new_metadata[gguf.Keys.Tokenizer.CHAT_TEMPLATE] = MetadataDetails(gguf.GGUFValueType.STRING, template) - - if args.pre_tokenizer: - new_metadata[gguf.Keys.Tokenizer.PRE] = MetadataDetails(gguf.GGUFValueType.STRING, args.pre_tokenizer) - - if remove_metadata: - logger.warning('*** Warning *** Warning *** Warning **') - logger.warning('* Most metadata is required for a fully functional GGUF file,') - logger.warning('* removing crucial metadata may result in a corrupt output file!') - - if not args.force: - logger.warning('* Enter exactly YES if you are positive you want to proceed:') - response = input('YES, I am sure> ') - if response != 'YES': - logger.info("You didn't enter YES. Okay then, see ya!") - sys.exit(0) - - logger.info(f'* Loading: {args.input}') - reader = gguf.GGUFReader(args.input, 'r') - - arch = get_field_data(reader, gguf.Keys.General.ARCHITECTURE) - - token_list = get_field_data(reader, gguf.Keys.Tokenizer.LIST) or [] - - for name, token in args.special_token or []: - if name not in token_names: - logger.warning(f'Unknown special token "{name}", ignoring...') - else: - ids = find_token(token_list, token) - new_metadata[token_names[name]] = MetadataDetails(gguf.GGUFValueType.UINT32, ids[0], f'= {token}') - - if len(ids) > 1: - logger.warning(f'Multiple "{token}" tokens found, choosing ID {ids[0]}, use --special-token-by-id if you want another:') - logger.warning(', '.join(str(i) for i in ids)) - - for name, id_string in args.special_token_by_id or []: - if name not in token_names: - logger.warning(f'Unknown special token "{name}", ignoring...') - elif not id_string.isdecimal(): - raise LookupError(f'Token ID "{id_string}" is not a valid ID!') - else: - id_int = int(id_string) - - if id_int >= 0 and id_int < len(token_list): - new_metadata[token_names[name]] = MetadataDetails(gguf.GGUFValueType.UINT32, id_int, f'= {token_list[id_int]}') - else: - raise LookupError(f'Token ID {id_int} is not within token list!') - - if os.path.isfile(args.output) and not args.force: - logger.warning('*** Warning *** Warning *** Warning **') - logger.warning(f'* The "{args.output}" GGUF file already exists, it will be overwritten!') - logger.warning('* Enter exactly YES if you are positive you want to proceed:') - response = input('YES, I am sure> ') - if response != 'YES': - logger.info("You didn't enter YES. Okay then, see ya!") - sys.exit(0) - - logger.info(f'* Writing: {args.output}') - writer = gguf.GGUFWriter(args.output, arch=arch, endianess=reader.endianess) - - alignment = get_field_data(reader, gguf.Keys.General.ALIGNMENT) - if alignment is not None: - logger.debug(f'Setting custom alignment: {alignment}') - writer.data_alignment = alignment - - copy_with_new_metadata(reader, writer, new_metadata, remove_metadata) - - -if __name__ == '__main__': - main() diff --git a/venv/lib/python3.10/site-packages/gguf/scripts/gguf_set_metadata.py b/venv/lib/python3.10/site-packages/gguf/scripts/gguf_set_metadata.py deleted file mode 100644 index f5809c35c887086f8055e828719041f8ae3ea749..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/scripts/gguf_set_metadata.py +++ /dev/null @@ -1,95 +0,0 @@ -#!/usr/bin/env python3 -import logging -import argparse -import os -import sys -from pathlib import Path - -# Necessary to load the local gguf package -if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent.parent / 'gguf-py').exists(): - sys.path.insert(0, str(Path(__file__).parent.parent.parent)) - -from gguf import GGUFReader # noqa: E402 - -logger = logging.getLogger("gguf-set-metadata") - - -def minimal_example(filename: str) -> None: - reader = GGUFReader(filename, 'r+') - field = reader.fields['tokenizer.ggml.bos_token_id'] - if field is None: - return - part_index = field.data[0] - field.parts[part_index][0] = 2 # Set tokenizer.ggml.bos_token_id to 2 - # - # So what's this field.data thing? It's helpful because field.parts contains - # _every_ part of the GGUF field. For example, tokenizer.ggml.bos_token_id consists - # of: - # - # Part index 0: Key length (27) - # Part index 1: Key data ("tokenizer.ggml.bos_token_id") - # Part index 2: Field type (4, the id for GGUFValueType.UINT32) - # Part index 3: Field value - # - # Note also that each part is an NDArray slice, so even a part that - # is only a single value like the key length will be a NDArray of - # the key length type (numpy.uint32). - # - # The .data attribute in the Field is a list of relevant part indexes - # and doesn't contain internal GGUF details like the key length part. - # In this case, .data will be [3] - just the part index of the - # field value itself. - - -def set_metadata(reader: GGUFReader, args: argparse.Namespace) -> None: - field = reader.get_field(args.key) - if field is None: - logger.error(f'! Field {repr(args.key)} not found') - sys.exit(1) - # Note that field.types is a list of types. This is because the GGUF - # format supports arrays. For example, an array of UINT32 would - # look like [GGUFValueType.ARRAY, GGUFValueType.UINT32] - handler = reader.gguf_scalar_to_np.get(field.types[0]) if field.types else None - if handler is None: - logger.error(f'! This tool only supports changing simple values, {repr(args.key)} has unsupported type {field.types}') - sys.exit(1) - current_value = field.parts[field.data[0]][0] - new_value = handler(args.value) - logger.info(f'* Preparing to change field {repr(args.key)} from {current_value} to {new_value}') - if current_value == new_value: - logger.info(f'- Key {repr(args.key)} already set to requested value {current_value}') - sys.exit(0) - if args.dry_run: - sys.exit(0) - if not args.force: - logger.warning('*** Warning *** Warning *** Warning **') - logger.warning('* Changing fields in a GGUF file can make it unusable. Proceed at your own risk.') - logger.warning('* Enter exactly YES if you are positive you want to proceed:') - response = input('YES, I am sure> ') - if response != 'YES': - logger.info("You didn't enter YES. Okay then, see ya!") - sys.exit(0) - field.parts[field.data[0]][0] = new_value - logger.info('* Field changed. Successful completion.') - - -def main() -> None: - parser = argparse.ArgumentParser(description="Set a simple value in GGUF file metadata") - parser.add_argument("model", type=str, help="GGUF format model filename") - parser.add_argument("key", type=str, help="Metadata key to set") - parser.add_argument("value", type=str, help="Metadata value to set") - parser.add_argument("--dry-run", action="store_true", help="Don't actually change anything") - parser.add_argument("--force", action="store_true", help="Change the field without confirmation") - parser.add_argument("--verbose", action="store_true", help="increase output verbosity") - - args = parser.parse_args(None if len(sys.argv) > 1 else ["--help"]) - - logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) - - logger.info(f'* Loading: {args.model}') - reader = GGUFReader(args.model, 'r' if args.dry_run else 'r+') - set_metadata(reader, args) - - -if __name__ == '__main__': - main() diff --git a/venv/lib/python3.10/site-packages/gguf/tensor_mapping.py b/venv/lib/python3.10/site-packages/gguf/tensor_mapping.py deleted file mode 100644 index 79f044d2a5945236b613ac3733ccc2b60ebd9ba4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/tensor_mapping.py +++ /dev/null @@ -1,1280 +0,0 @@ -from __future__ import annotations - -from typing import Sequence - -from .constants import MODEL_ARCH, MODEL_TENSOR, MODEL_TENSORS, TENSOR_NAMES - - -class TensorNameMap: - mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = { - # Token embeddings - MODEL_TENSOR.TOKEN_EMBD: ( - "gpt_neox.embed_in", # gptneox - "transformer.wte", # gpt2 gpt-j mpt refact qwen dbrx jais exaone - "transformer.word_embeddings", # falcon - "word_embeddings", # bloom - "model.embed_tokens", # llama-hf nemotron olmoe olmo2 rwkv6qwen2 glm4-0414 - "tok_embeddings", # llama-pth - "embeddings.word_embeddings", # bert nomic-bert - "language_model.embedding.word_embeddings", # persimmon - "wte", # gpt2 - "transformer.embd.wte", # phi2 - "model.tok_embeddings", # internlm2 - "model.embedding", # mamba-qbert - "backbone.embedding", # mamba - "backbone.embeddings", # mamba-hf - "transformer.in_out_embed", # Grok - "embedding.word_embeddings", # chatglm - "transformer.token_embeddings", # openelm - "shared", # t5 - "rwkv.embeddings", # rwkv6 - "model.embeddings", # rwkv7 - "model.word_embeddings", # bailingmoe - "language_model.model.embed_tokens", # llama4 - "encoder", # neobert - ), - - # Token type embeddings - MODEL_TENSOR.TOKEN_TYPES: ( - "embeddings.token_type_embeddings", # bert nomic-bert - ), - - # Normalization of token embeddings - MODEL_TENSOR.TOKEN_EMBD_NORM: ( - "word_embeddings_layernorm", # bloom - "embeddings.LayerNorm", # bert - "emb_ln", # nomic-bert - "transformer.norm", # openelm - "rwkv.blocks.0.pre_ln", # rwkv - "rwkv.blocks.0.pre_ln", # rwkv6 - "model.pre_ln", # rwkv7 - "model.layers.0.pre_norm", # rwkv7 - "backbone.norm", # wavtokenizer - ), - - # Position embeddings - MODEL_TENSOR.POS_EMBD: ( - "transformer.wpe", # gpt2 - "embeddings.position_embeddings", # bert - "wpe", # gpt2 - ), - - # Output - MODEL_TENSOR.OUTPUT: ( - "embed_out", # gptneox - "lm_head", # gpt2 mpt falcon llama-hf baichuan qwen mamba dbrx jais nemotron exaone olmoe olmo2 phimoe - "output", # llama-pth bloom internlm2 - "word_embeddings_for_head", # persimmon - "lm_head.linear", # phi2 - "output_layer", # chatglm - "head", # rwkv - "head.out", # wavtokenizer - "lm_head", # llama4 - ), - - # Output norm - MODEL_TENSOR.OUTPUT_NORM: ( - "gpt_neox.final_layer_norm", # gptneox - "transformer.ln_f", # gpt2 gpt-j falcon jais exaone - "model.norm", # llama-hf baichuan internlm2 olmoe olmo2 phimoe - "norm", # llama-pth - "transformer.norm_f", # mpt dbrx - "ln_f", # refact bloom qwen gpt2 - "language_model.encoder.final_layernorm", # persimmon - "model.final_layernorm", # persimmon - "lm_head.ln", # phi2 - "model.norm_f", # mamba-qbert - "backbone.norm_f", # mamba - "transformer.rms_norm", # Grok - "encoder.final_layernorm", # chatglm - "transformer.norm", # openelm - "model.norm", # nemotron - "rwkv.ln_out", # rwkv6 - "model.ln_out", # rwkv7 - "backbone.final_layer_norm", # wavtokenizer - "model.norm", # llama4 - ), - - # Rope frequencies - MODEL_TENSOR.ROPE_FREQS: ( - "rope.freqs", # llama-pth - "rotary_pos_emb.inv_freq", # chatglm - ), - - MODEL_TENSOR.ROPE_FACTORS_LONG: (), - MODEL_TENSOR.ROPE_FACTORS_SHORT: (), - - MODEL_TENSOR.CONV1D: ( - "backbone.embed", # roberta - ), - } - - block_mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = { - # Attention norm - MODEL_TENSOR.ATTN_NORM: ( - "gpt_neox.layers.{bid}.input_layernorm", # gptneox - "transformer.h.{bid}.ln_1", # gpt2 gpt-j refact qwen jais exaone - "transformer.blocks.{bid}.norm_1", # mpt - "transformer.h.{bid}.input_layernorm", # falcon7b - "h.{bid}.input_layernorm", # bloom - "transformer.h.{bid}.ln_mlp", # falcon40b - "model.layers.{bid}.input_layernorm", # llama-hf nemotron olmoe phimoe - "layers.{bid}.attention_norm", # llama-pth - "language_model.encoder.layers.{bid}.input_layernorm", # persimmon - "model.layers.{bid}.ln1", # yi - "h.{bid}.ln_1", # gpt2 - "transformer.h.{bid}.ln", # phi2 - "model.layers.layers.{bid}.norm", # plamo - "model.layers.{bid}.attention_norm", # internlm2 - "model.layers.{bid}.norm", # mamba-qbert - "backbone.layers.{bid}.norm", # mamba - "transformer.decoder_layer.{bid}.rms_norm", # Grok - "transformer.blocks.{bid}.norm_attn_norm.norm_1", # dbrx - "encoder.layers.{bid}.input_layernorm", # chatglm - "transformer.layers.{bid}.attn_norm", # openelm - "rwkv.blocks.{bid}.ln1", # rwkv6 - "model.layers.{bid}.ln1", # rwkv7 - "model.layers.{bid}.input_layernorm", # llama4 - "transformer_encoder.{bid}.attention_norm", # neobert - ), - - # Attention norm 2 - MODEL_TENSOR.ATTN_NORM_2: ( - "transformer.h.{bid}.ln_attn", # falcon40b - "encoder.layer.{bid}.layer_norm_1", # jina-v2-code - "rwkv.blocks.{bid}.ln2", # rwkv6 - "model.layers.{bid}.ln2", # rwkv7 - ), - - # Attention query-key-value - MODEL_TENSOR.ATTN_QKV: ( - "gpt_neox.layers.{bid}.attention.query_key_value", # gptneox - "transformer.h.{bid}.attn.c_attn", # gpt2 qwen jais - "transformer.blocks.{bid}.attn.Wqkv", # mpt - "transformer.blocks.{bid}.norm_attn_norm.attn.Wqkv", # dbrx - "transformer.h.{bid}.self_attention.query_key_value", # falcon - "h.{bid}.self_attention.query_key_value", # bloom - "language_model.encoder.layers.{bid}.self_attention.query_key_value", # persimmon - "model.layers.{bid}.self_attn.query_key_value", # persimmon - "h.{bid}.attn.c_attn", # gpt2 - "transformer.h.{bid}.mixer.Wqkv", # phi2 - "encoder.layers.{bid}.attn.Wqkv", # nomic-bert - "encoder.layers.{bid}.mixer.Wqkv", # jina - "model.layers.{bid}.self_attn.qkv_proj", # phi3 - "encoder.layers.{bid}.self_attention.query_key_value", # chatglm - "transformer.layers.{bid}.attn.qkv_proj", # openelm - "transformer_encoder.{bid}.qkv", # neobert - ), - - # Attention query - MODEL_TENSOR.ATTN_Q: ( - "model.layers.{bid}.self_attn.q_proj", # llama-hf nemotron olmoe olmo2 phimoe - "model.layers.{bid}.self_attn.q_proj_no_perm", # llama-custom - "layers.{bid}.attention.wq", # llama-pth - "encoder.layer.{bid}.attention.self.query", # bert - "transformer.layer.{bid}.attention.q_lin", # distillbert - "transformer.h.{bid}.attn.q_proj", # gpt-j - "model.layers.layers.{bid}.self_attn.q_proj", # plamo - "model.layers.{bid}.attention.wq", # internlm2 - "transformer.decoder_layer.{bid}.multi_head_attention.query",# Grok - "transformer.h.{bid}.attn.attention.q_proj", # exaone - "model.layers.{bid}.self_attn.q_proj", # llama4 - ), - - # Attention key - MODEL_TENSOR.ATTN_K: ( - "model.layers.{bid}.self_attn.k_proj", # llama-hf nemotron olmoe olmo2 phimoe - "model.layers.{bid}.self_attn.k_proj_no_perm", # llama-custom - "layers.{bid}.attention.wk", # llama-pth - "encoder.layer.{bid}.attention.self.key", # bert - "transformer.layer.{bid}.attention.k_lin", # distillbert - "transformer.h.{bid}.attn.k_proj", # gpt-j - "transformer.h.{bid}.attn.k", # refact - "model.layers.layers.{bid}.self_attn.k_proj", # plamo - "model.layers.{bid}.attention.wk", # internlm2 - "transformer.decoder_layer.{bid}.multi_head_attention.key",# Grok - "transformer.h.{bid}.attn.attention.k_proj", # exaone - "model.layers.{bid}.self_attn.k_proj", # llama4 - ), - - # Attention value - MODEL_TENSOR.ATTN_V: ( - "model.layers.{bid}.self_attn.v_proj", # llama-hf nemotron olmoe olmo2 phimoe - "layers.{bid}.attention.wv", # llama-pth - "encoder.layer.{bid}.attention.self.value", # bert - "transformer.layer.{bid}.attention.v_lin", # distillbert - "transformer.h.{bid}.attn.v_proj", # gpt-j - "transformer.h.{bid}.attn.v", # refact - "model.layers.layers.{bid}.self_attn.v_proj", # plamo - "model.layers.{bid}.attention.wv", # internlm2 - "transformer.decoder_layer.{bid}.multi_head_attention.value",# Grok - "transformer.h.{bid}.attn.attention.v_proj", # exaone - "model.layers.{bid}.self_attn.v_proj", # llama4 - ), - - # Attention output - MODEL_TENSOR.ATTN_OUT: ( - "gpt_neox.layers.{bid}.attention.dense", # gptneox - "transformer.h.{bid}.attn.c_proj", # gpt2 refact qwen jais - "transformer.blocks.{bid}.attn.out_proj", # mpt - "transformer.h.{bid}.self_attention.dense", # falcon - "h.{bid}.self_attention.dense", # bloom - "model.layers.{bid}.self_attn.o_proj", # llama-hf nemotron olmoe olmo2 phimoe - "model.layers.{bid}.self_attn.linear_attn", # deci - "layers.{bid}.attention.wo", # llama-pth - "encoder.layer.{bid}.attention.output.dense", # bert - "transformer.layer.{bid}.attention.out_lin", # distillbert - "transformer.h.{bid}.attn.out_proj", # gpt-j - "language_model.encoder.layers.{bid}.self_attention.dense", # persimmon - "model.layers.{bid}.self_attn.dense", # persimmon - "h.{bid}.attn.c_proj", # gpt2 - "transformer.h.{bid}.mixer.out_proj", # phi2 - "model.layers.layers.{bid}.self_attn.o_proj", # plamo - "model.layers.{bid}.attention.wo", # internlm2 - "encoder.layers.{bid}.attn.out_proj", # nomic-bert - "encoder.layers.{bid}.mixer.out_proj", # jina - "transformer.decoder_layer.{bid}.multi_head_attention.linear", # Grok - "transformer.blocks.{bid}.norm_attn_norm.attn.out_proj", # dbrx - "encoder.layers.{bid}.self_attention.dense", # chatglm - "transformer.layers.{bid}.attn.out_proj", # openelm - "transformer.h.{bid}.attn.attention.out_proj", # exaone - "model.layers.{bid}.self_attn.o_proj", # llama4 - "transformer_encoder.{bid}.wo", # neobert - ), - - # Attention output norm - MODEL_TENSOR.ATTN_OUT_NORM: ( - "encoder.layer.{bid}.attention.output.LayerNorm", # bert - "transformer.layer.{bid}.sa_layer_norm", # distillbert - "encoder.layers.{bid}.norm1", # nomic-bert - "transformer.decoder_layer.{bid}.rms_norm_1", # Grok - "transformer.blocks.{bid}.norm_attn_norm.norm_2", # dbrx - ), - - MODEL_TENSOR.ATTN_POST_NORM: ( - "model.layers.{bid}.post_attention_layernorm", # gemma2 olmo2 # ge - "model.layers.{bid}.post_self_attn_layernorm", # glm-4-0414 - ), - - # Rotary embeddings - MODEL_TENSOR.ATTN_ROT_EMBD: ( - "model.layers.{bid}.self_attn.rotary_emb.inv_freq", # llama-hf - "layers.{bid}.attention.inner_attention.rope.freqs", # llama-pth - "model.layers.layers.{bid}.self_attn.rotary_emb.inv_freq", # plamo - "transformer.h.{bid}.attn.rotary_emb.inv_freq", # codeshell - ), - - # Feed-forward norm - MODEL_TENSOR.FFN_NORM: ( - "gpt_neox.layers.{bid}.post_attention_layernorm", # gptneox - "transformer.h.{bid}.ln_2", # gpt2 refact qwen jais exaone - "h.{bid}.post_attention_layernorm", # bloom - "transformer.blocks.{bid}.norm_2", # mpt - "model.layers.{bid}.post_attention_layernorm", # llama-hf nemotron olmoe phimoe - "layers.{bid}.ffn_norm", # llama-pth - "language_model.encoder.layers.{bid}.post_attention_layernorm", # persimmon - "model.layers.{bid}.ln2", # yi - "h.{bid}.ln_2", # gpt2 - "model.layers.{bid}.ffn_norm", # internlm2 - "transformer.decoder_layer.{bid}.rms_norm_2", # Grok - "encoder.layers.{bid}.post_attention_layernorm", # chatglm - "transformer.layers.{bid}.ffn_norm", # openelm - "model.layers.{bid}.post_attention_layernorm", # llama4 - "transformer_encoder.{bid}.ffn_norm", # neobert - ), - - # Post feed-forward norm - MODEL_TENSOR.FFN_PRE_NORM: ( - "model.layers.{bid}.pre_feedforward_layernorm", # gemma2 - ), - - # Post feed-forward norm - MODEL_TENSOR.FFN_POST_NORM: ( - "model.layers.{bid}.post_feedforward_layernorm", # gemma2 olmo2 - "model.layers.{bid}.post_mlp_layernorm", # glm-4-0414 - ), - - MODEL_TENSOR.FFN_GATE_INP: ( - "layers.{bid}.feed_forward.gate", # mixtral - "model.layers.{bid}.block_sparse_moe.gate", # mixtral phimoe - "model.layers.{bid}.mlp.gate", # qwen2moe olmoe - "transformer.decoder_layer.{bid}.router", # Grok - "transformer.blocks.{bid}.ffn.router.layer", # dbrx - "model.layers.{bid}.block_sparse_moe.router.layer", # granitemoe - "model.layers.{bid}.feed_forward.router", # llama4 - "encoder.layers.{bid}.mlp.router.layer", # nomic-bert-moe - ), - - MODEL_TENSOR.FFN_GATE_INP_SHEXP: ( - "model.layers.{bid}.mlp.shared_expert_gate", # qwen2moe - ), - - MODEL_TENSOR.FFN_EXP_PROBS_B: ( - "model.layers.{bid}.mlp.gate.e_score_correction", # deepseek-v3 dots1 - ), - - # Feed-forward up - MODEL_TENSOR.FFN_UP: ( - "gpt_neox.layers.{bid}.mlp.dense_h_to_4h", # gptneox - "transformer.h.{bid}.mlp.c_fc", # gpt2 jais - "transformer.blocks.{bid}.ffn.up_proj", # mpt - "transformer.h.{bid}.mlp.dense_h_to_4h", # falcon - "h.{bid}.mlp.dense_h_to_4h", # bloom - "model.layers.{bid}.mlp.up_proj", # llama-hf refact nemotron olmo2 - "layers.{bid}.feed_forward.w3", # llama-pth - "encoder.layer.{bid}.intermediate.dense", # bert - "transformer.layer.{bid}.ffn.lin1", # distillbert - "transformer.h.{bid}.mlp.fc_in", # gpt-j - "transformer.h.{bid}.mlp.linear_3", # refact - "language_model.encoder.layers.{bid}.mlp.dense_h_to_4h", # persimmon - "model.layers.{bid}.mlp.dense_h_to_4h", # persimmon - "transformer.h.{bid}.mlp.w1", # qwen - "h.{bid}.mlp.c_fc", # gpt2 - "transformer.h.{bid}.mlp.fc1", # phi2 - "model.layers.{bid}.mlp.fc1", # phi2 - "model.layers.{bid}.mlp.gate_up_proj", # phi3 glm-4-0414 - "model.layers.layers.{bid}.mlp.up_proj", # plamo - "model.layers.{bid}.feed_forward.w3", # internlm2 - "encoder.layers.{bid}.mlp.fc11", # nomic-bert - "encoder.layers.{bid}.mlp.fc1", # nomic-bert-moe - "model.layers.{bid}.mlp.c_fc", # starcoder2 - "encoder.layer.{bid}.mlp.gated_layers_v", # jina-bert-v2 (split up/gate, no longer used) - "encoder.layer.{bid}.mlp.gated_layers", # jina-bert-v2 (GEGLU) - "encoder.layer.{bid}.mlp.up_gated_layer", # jina-v2-code (GEGLU) - "model.layers.{bid}.residual_mlp.w3", # arctic - "encoder.layers.{bid}.mlp.dense_h_to_4h", # chatglm - "transformer.h.{bid}.mlp.c_fc_1", # exaone - "model.layers.{bid}.feed_forward.up_proj", # llama4 - "transformer_encoder.{bid}.ffn.w12", # neobert - ), - - MODEL_TENSOR.FFN_UP_EXP: ( - "layers.{bid}.feed_forward.experts.w3", # mixtral (merged) - "transformer.decoder_layer.{bid}.moe.linear_v", # Grok (merged) - "transformer.blocks.{bid}.ffn.experts.mlp.v1", # dbrx - "model.layers.{bid}.mlp.experts.up_proj", # qwen2moe olmoe (merged) - "model.layers.{bid}.block_sparse_moe.experts.w3", # phimoe (merged) - "model.layers.{bid}.feed_forward.experts.up_proj", # llama4 - "encoder.layers.{bid}.mlp.experts.mlp.w1", # nomic-bert-moe - ), - - MODEL_TENSOR.FFN_UP_SHEXP: ( - "model.layers.{bid}.mlp.shared_expert.up_proj", # qwen2moe - "model.layers.{bid}.mlp.shared_experts.up_proj", # deepseek deepseek2 - "model.layers.{bid}.feed_forward.shared_expert.up_proj", # llama4 - ), - - # AWQ-activation gate - MODEL_TENSOR.FFN_ACT: ( - "transformer.blocks.{bid}.ffn.act", # mpt - ), - - # Feed-forward gate - MODEL_TENSOR.FFN_GATE: ( - "model.layers.{bid}.mlp.gate_proj", # llama-hf refact olmo2 - "layers.{bid}.feed_forward.w1", # llama-pth - "transformer.h.{bid}.mlp.w2", # qwen - "transformer.h.{bid}.mlp.c_fc2", # jais - "model.layers.layers.{bid}.mlp.gate_proj", # plamo - "model.layers.{bid}.feed_forward.w1", # internlm2 - "encoder.layers.{bid}.mlp.fc12", # nomic-bert - "encoder.layer.{bid}.mlp.gated_layers_w", # jina-bert-v2 (split up/gate, no longer used) - "transformer.h.{bid}.mlp.linear_1", # refact - "model.layers.{bid}.residual_mlp.w1", # arctic - "transformer.h.{bid}.mlp.c_fc_0", # exaone - "model.layers.{bid}.feed_forward.gate_proj", # llama4 - ), - - MODEL_TENSOR.FFN_GATE_EXP: ( - "layers.{bid}.feed_forward.experts.w1", # mixtral (merged) - "transformer.decoder_layer.{bid}.moe.linear", # Grok (merged) - "transformer.blocks.{bid}.ffn.experts.mlp.w1", # dbrx - "model.layers.{bid}.mlp.experts.gate_proj", # qwen2moe olmoe (merged) - "model.layers.{bid}.block_sparse_moe.experts.w1", # phimoe (merged) - "model.layers.{bid}.feed_forward.experts.gate_proj", # llama4 - ), - - MODEL_TENSOR.FFN_GATE_SHEXP: ( - "model.layers.{bid}.mlp.shared_expert.gate_proj", # qwen2moe - "model.layers.{bid}.mlp.shared_experts.gate_proj", # deepseek deepseek2 - "model.layers.{bid}.feed_forward.shared_expert.gate_proj", # llama4 - ), - - # Feed-forward down - MODEL_TENSOR.FFN_DOWN: ( - "gpt_neox.layers.{bid}.mlp.dense_4h_to_h", # gptneox - "transformer.h.{bid}.mlp.c_proj", # gpt2 refact qwen jais - "transformer.blocks.{bid}.ffn.down_proj", # mpt - "transformer.h.{bid}.mlp.dense_4h_to_h", # falcon - "h.{bid}.mlp.dense_4h_to_h", # bloom - "model.layers.{bid}.mlp.down_proj", # llama-hf nemotron olmo2 - "layers.{bid}.feed_forward.w2", # llama-pth - "encoder.layer.{bid}.output.dense", # bert - "transformer.layer.{bid}.ffn.lin2", # distillbert - "transformer.h.{bid}.mlp.fc_out", # gpt-j - "language_model.encoder.layers.{bid}.mlp.dense_4h_to_h", # persimmon - "model.layers.{bid}.mlp.dense_4h_to_h", # persimmon - "h.{bid}.mlp.c_proj", # gpt2 - "transformer.h.{bid}.mlp.fc2", # phi2 - "model.layers.{bid}.mlp.fc2", # phi2 - "model.layers.layers.{bid}.mlp.down_proj", # plamo - "model.layers.{bid}.feed_forward.w2", # internlm2 - "encoder.layers.{bid}.mlp.fc2", # nomic-bert - "model.layers.{bid}.mlp.c_proj", # starcoder2 - "encoder.layer.{bid}.mlp.wo", # jina-bert-v2 - "transformer.layers.{bid}.ffn.proj_2", # openelm - "model.layers.{bid}.residual_mlp.w2", # arctic - "encoder.layer.{bid}.mlp.down_layer", # jina-bert-v2 - "encoder.layers.{bid}.mlp.dense_4h_to_h", # chatglm - "model.layers.h.{bid}.mlp.c_proj", # exaone - "model.layers.{bid}.feed_forward.down_proj", # llama4 - "transformer_encoder.{bid}.ffn.w3", # neobert - ), - - MODEL_TENSOR.FFN_DOWN_EXP: ( - "layers.{bid}.feed_forward.experts.w2", # mixtral (merged) - "transformer.decoder_layer.{bid}.moe.linear_1", # Grok (merged) - "transformer.blocks.{bid}.ffn.experts.mlp.w2", # dbrx - "model.layers.{bid}.mlp.experts.down_proj", # qwen2moe olmoe (merged) - "model.layers.{bid}.block_sparse_moe.output_linear", # granitemoe - "model.layers.{bid}.block_sparse_moe.experts.w2", # phimoe (merged) - "model.layers.{bid}.feed_forward.experts.down_proj", # llama4 - "encoder.layers.{bid}.mlp.experts.mlp.w2", # nomic-bert-moe - ), - - MODEL_TENSOR.FFN_DOWN_SHEXP: ( - "model.layers.{bid}.mlp.shared_expert.down_proj", # qwen2moe - "model.layers.{bid}.mlp.shared_experts.down_proj", # deepseek deepseek2 - "model.layers.{bid}.feed_forward.shared_expert.down_proj", # llama4 - "model.layers.{bid}.shared_mlp.output_linear", # granitemoe - ), - - MODEL_TENSOR.ATTN_Q_NORM: ( - "language_model.encoder.layers.{bid}.self_attention.q_layernorm", - "model.layers.{bid}.self_attn.q_layernorm", # persimmon - "model.layers.{bid}.self_attn.q_norm", # cohere olmoe chameleon olmo2 - "transformer.blocks.{bid}.attn.q_ln", # sea-lion - "encoder.layer.{bid}.attention.self.layer_norm_q", # jina-bert-v2 - "transformer.layers.{bid}.attn.q_norm", # openelm - ), - - MODEL_TENSOR.ATTN_K_NORM: ( - "language_model.encoder.layers.{bid}.self_attention.k_layernorm", - "model.layers.{bid}.self_attn.k_layernorm", # persimmon - "model.layers.{bid}.self_attn.k_norm", # cohere olmoe chameleon olmo2 - "transformer.blocks.{bid}.attn.k_ln", # sea-lion - "encoder.layer.{bid}.attention.self.layer_norm_k", # jina-bert-v2 - "transformer.layers.{bid}.attn.k_norm", # openelm - ), - - MODEL_TENSOR.ROPE_FREQS: ( - "language_model.encoder.layers.{bid}.self_attention.rotary_emb.inv_freq", # persimmon - ), - - MODEL_TENSOR.LAYER_OUT_NORM: ( - "encoder.layer.{bid}.output.LayerNorm", # bert - "transformer.layer.{bid}.output_layer_norm", # distillbert - "encoder.layers.{bid}.norm2", # nomic-bert - "transformer.decoder_layer.{bid}.rms_norm_3", # Grok - "encoder.layer.{bid}.mlp.layernorm", # jina-bert-v2 - "encoder.layer.{bid}.layer_norm_2" # jina-v2-code - ), - - MODEL_TENSOR.SSM_IN: ( - "model.layers.{bid}.in_proj", - "backbone.layers.{bid}.mixer.in_proj", - ), - - MODEL_TENSOR.SSM_CONV1D: ( - "model.layers.{bid}.conv1d", - "backbone.layers.{bid}.mixer.conv1d", - ), - - MODEL_TENSOR.SSM_X: ( - "model.layers.{bid}.x_proj", - "backbone.layers.{bid}.mixer.x_proj", - ), - - MODEL_TENSOR.SSM_DT: ( - "model.layers.{bid}.dt_proj", - "backbone.layers.{bid}.mixer.dt_proj", - ), - - MODEL_TENSOR.SSM_A: ( - "model.layers.{bid}.A_log", - "backbone.layers.{bid}.mixer.A_log", - ), - - MODEL_TENSOR.SSM_D: ( - "model.layers.{bid}.D", - "backbone.layers.{bid}.mixer.D", - ), - - MODEL_TENSOR.SSM_OUT: ( - "model.layers.{bid}.out_proj", - "backbone.layers.{bid}.mixer.out_proj", - ), - - MODEL_TENSOR.TIME_MIX_W0: ( - "model.layers.{bid}.attention.w0", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_W1: ( - "rwkv.blocks.{bid}.attention.time_maa_w1", # rwkv6 - "model.layers.{bid}.self_attn.time_maa_w1", # rwkv6qwen2 - "model.layers.{bid}.attention.w1", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_W2: ( - "rwkv.blocks.{bid}.attention.time_maa_w2", # rwkv6 - "model.layers.{bid}.self_attn.time_maa_w2", # rwkv6qwen2 - "model.layers.{bid}.attention.w2", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_A0: ( - "model.layers.{bid}.attention.a0", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_A1: ( - "model.layers.{bid}.attention.a1", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_A2: ( - "model.layers.{bid}.attention.a2", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_V0: ( - "model.layers.{bid}.attention.v0", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_V1: ( - "model.layers.{bid}.attention.v1", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_V2: ( - "model.layers.{bid}.attention.v2", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_G1: ( - "model.layers.{bid}.attention.g1", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_G2: ( - "model.layers.{bid}.attention.g2", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_K_K: ( - "model.layers.{bid}.attention.k_k", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_K_A: ( - "model.layers.{bid}.attention.k_a", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_R_K: ( - "model.layers.{bid}.attention.r_k", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_LERP_X: ( - "rwkv.blocks.{bid}.attention.time_maa_x", # rwkv6 - "model.layers.{bid}.self_attn.time_maa_x", # rwkv6qwen2 - ), - - MODEL_TENSOR.TIME_MIX_LERP_K: ( - "rwkv.blocks.{bid}.attention.time_maa_k", # rwkv6 - "model.layers.{bid}.self_attn.time_maa_k", # rwkv6qwen2 - ), - - MODEL_TENSOR.TIME_MIX_LERP_V: ( - "rwkv.blocks.{bid}.attention.time_maa_v", # rwkv6 - "model.layers.{bid}.self_attn.time_maa_v", # rwkv6qwen2 - ), - - MODEL_TENSOR.TIME_MIX_LERP_R: ( - "rwkv.blocks.{bid}.attention.time_maa_r", # rwkv6 - "model.layers.{bid}.self_attn.time_maa_r", # rwkv6qwen2 - ), - - MODEL_TENSOR.TIME_MIX_LERP_G: ( - "rwkv.blocks.{bid}.attention.time_maa_g", # rwkv6 - "model.layers.{bid}.self_attn.time_maa_g", # rwkv6qwen2 - ), - - MODEL_TENSOR.TIME_MIX_LERP_W: ( - "rwkv.blocks.{bid}.attention.time_maa_w", # rwkv6 - "model.layers.{bid}.self_attn.time_maa_w", # rwkv6qwen2 - ), - - MODEL_TENSOR.TIME_MIX_FIRST: ( - "rwkv.blocks.{bid}.attention.time_faaaa", # rwkv6 - ), - - MODEL_TENSOR.TIME_MIX_DECAY: ( - "rwkv.blocks.{bid}.attention.time_decay", # rwkv6 - "model.layers.{bid}.self_attn.time_decay", # rwkv6qwen2 - ), - - MODEL_TENSOR.TIME_MIX_DECAY_W1: ( - "rwkv.blocks.{bid}.attention.time_decay_w1", # rwkv6 - "model.layers.{bid}.self_attn.time_decay_w1", # rwkv6qwen2 - ), - - MODEL_TENSOR.TIME_MIX_DECAY_W2: ( - "rwkv.blocks.{bid}.attention.time_decay_w2", # rwkv6 - "model.layers.{bid}.self_attn.time_decay_w2", # rwkv6qwen2 - ), - - MODEL_TENSOR.TIME_MIX_KEY: ( - "rwkv.blocks.{bid}.attention.key", # rwkv6 - "model.layers.{bid}.self_attn.k_proj", # rwkv6qwen2 - "model.layers.{bid}.attention.key", # rwkv7 - "model.layers.{bid}.attention.k_proj", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_VALUE: ( - "rwkv.blocks.{bid}.attention.value", # rwkv6 - "model.layers.{bid}.self_attn.v_proj", # rwkv6qwen2 - "model.layers.{bid}.attention.value", # rwkv7 - "model.layers.{bid}.attention.v_proj", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_RECEPTANCE: ( - "rwkv.blocks.{bid}.attention.receptance", # rwkv6 - "model.layers.{bid}.self_attn.q_proj", # rwkv6qwen2 - "model.layers.{bid}.attention.receptance", # rwkv7 - "model.layers.{bid}.attention.r_proj", # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_GATE: ( - "rwkv.blocks.{bid}.attention.gate", # rwkv6 - "model.layers.{bid}.self_attn.gate", # rwkv6qwen2 - ), - - MODEL_TENSOR.TIME_MIX_LN: ( - "rwkv.blocks.{bid}.attention.ln_x", # rwkv6 - "model.layers.{bid}.attention.ln_x" # rwkv7 - ), - - MODEL_TENSOR.TIME_MIX_OUTPUT: ( - "rwkv.blocks.{bid}.attention.output", # rwkv6 - "model.layers.{bid}.self_attn.o_proj", # rwkv6qwen2 - "model.layers.{bid}.attention.output", # rwkv7 - "model.layers.{bid}.attention.o_proj", # rwkv7 - ), - - MODEL_TENSOR.CHANNEL_MIX_LERP_K: ( - "rwkv.blocks.{bid}.feed_forward.time_maa_k", # rwkv6 - "model.layers.{bid}.feed_forward.x_k", # rwkv7 - ), - - MODEL_TENSOR.CHANNEL_MIX_LERP_R: ( - "rwkv.blocks.{bid}.feed_forward.time_maa_r", # rwkv6 - ), - - MODEL_TENSOR.CHANNEL_MIX_KEY: ( - "rwkv.blocks.{bid}.feed_forward.key", # rwkv6 - "model.layers.{bid}.feed_forward.key", # rwkv7 - ), - - MODEL_TENSOR.CHANNEL_MIX_RECEPTANCE: ( - "rwkv.blocks.{bid}.feed_forward.receptance", # rwkv6 - ), - - MODEL_TENSOR.CHANNEL_MIX_VALUE: ( - "rwkv.blocks.{bid}.feed_forward.value", # rwkv6 - "model.layers.{bid}.feed_forward.value", # rwkv7 - ), - - MODEL_TENSOR.ATTN_Q_A: ( - "model.layers.{bid}.self_attn.q_a_proj", # deepseek2 - ), - - MODEL_TENSOR.ATTN_Q_B: ( - "model.layers.{bid}.self_attn.q_b_proj", # deepseek2 - ), - - MODEL_TENSOR.ATTN_KV_A_MQA: ( - "model.layers.{bid}.self_attn.kv_a_proj_with_mqa", # deepseek2 - ), - - MODEL_TENSOR.ATTN_KV_B: ( - "model.layers.{bid}.self_attn.kv_b_proj", # deepseek2 - ), - - MODEL_TENSOR.ATTN_K_B: ( - "model.layers.{bid}.self_attn.k_b_proj", # deepseek2 - ), - - MODEL_TENSOR.ATTN_V_B: ( - "model.layers.{bid}.self_attn.v_b_proj", # deepseek2 - ), - - MODEL_TENSOR.ATTN_Q_A_NORM: ( - "model.layers.{bid}.self_attn.q_a_layernorm", # deepseek2 - ), - - MODEL_TENSOR.ATTN_KV_A_NORM: ( - "model.layers.{bid}.self_attn.kv_a_layernorm", # deepseek2 - ), - - MODEL_TENSOR.ATTN_SUB_NORM: ( - "model.layers.{bid}.self_attn.inner_attn_ln", # bitnet - ), - - MODEL_TENSOR.FFN_SUB_NORM: ( - "model.layers.{bid}.mlp.ffn_layernorm", # bitnet - ), - - MODEL_TENSOR.DEC_ATTN_NORM: ( - "decoder.block.{bid}.layer.0.layer_norm", # t5 - ), - - MODEL_TENSOR.DEC_ATTN_Q: ( - "decoder.block.{bid}.layer.0.SelfAttention.q", # t5 - ), - - MODEL_TENSOR.DEC_ATTN_K: ( - "decoder.block.{bid}.layer.0.SelfAttention.k", # t5 - ), - - MODEL_TENSOR.DEC_ATTN_V: ( - "decoder.block.{bid}.layer.0.SelfAttention.v", # t5 - ), - - MODEL_TENSOR.DEC_ATTN_OUT: ( - "decoder.block.{bid}.layer.0.SelfAttention.o", # t5 - ), - - MODEL_TENSOR.DEC_ATTN_REL_B: ( - "decoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias", # t5 - ), - - MODEL_TENSOR.DEC_CROSS_ATTN_NORM: ( - "decoder.block.{bid}.layer.1.layer_norm", # t5 - ), - - MODEL_TENSOR.DEC_CROSS_ATTN_Q: ( - "decoder.block.{bid}.layer.1.EncDecAttention.q", # t5 - ), - - MODEL_TENSOR.DEC_CROSS_ATTN_K: ( - "decoder.block.{bid}.layer.1.EncDecAttention.k", # t5 - ), - - MODEL_TENSOR.DEC_CROSS_ATTN_V: ( - "decoder.block.{bid}.layer.1.EncDecAttention.v", # t5 - ), - - MODEL_TENSOR.DEC_CROSS_ATTN_OUT: ( - "decoder.block.{bid}.layer.1.EncDecAttention.o", # t5 - ), - - MODEL_TENSOR.DEC_CROSS_ATTN_REL_B: ( - "decoder.block.{bid}.layer.1.EncDecAttention.relative_attention_bias", # t5 - ), - - MODEL_TENSOR.DEC_FFN_NORM: ( - "decoder.block.{bid}.layer.2.layer_norm", # t5 - ), - - MODEL_TENSOR.DEC_FFN_GATE: ( - "decoder.block.{bid}.layer.2.DenseReluDense.wi_0", # flan-t5 - ), - - MODEL_TENSOR.DEC_FFN_UP: ( - "decoder.block.{bid}.layer.2.DenseReluDense.wi", # t5 - "decoder.block.{bid}.layer.2.DenseReluDense.wi_1", # flan-t5 - ), - - MODEL_TENSOR.DEC_FFN_DOWN: ( - "decoder.block.{bid}.layer.2.DenseReluDense.wo", # t5 - ), - - MODEL_TENSOR.DEC_OUTPUT_NORM: ( - "decoder.final_layer_norm", # t5 - ), - - MODEL_TENSOR.ENC_ATTN_NORM: ( - "encoder.block.{bid}.layer.0.layer_norm", # t5 - ), - - MODEL_TENSOR.ENC_ATTN_Q: ( - "encoder.block.{bid}.layer.0.SelfAttention.q", # t5 - ), - - MODEL_TENSOR.ENC_ATTN_K: ( - "encoder.block.{bid}.layer.0.SelfAttention.k", # t5 - ), - - MODEL_TENSOR.ENC_ATTN_V: ( - "encoder.block.{bid}.layer.0.SelfAttention.v", # t5 - ), - - MODEL_TENSOR.ENC_ATTN_OUT: ( - "encoder.block.{bid}.layer.0.SelfAttention.o", # t5 - ), - - MODEL_TENSOR.ENC_ATTN_REL_B: ( - "encoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias", # t5 - ), - - MODEL_TENSOR.ENC_FFN_NORM: ( - "encoder.block.{bid}.layer.1.layer_norm", # t5 - ), - - MODEL_TENSOR.ENC_FFN_GATE: ( - "encoder.block.{bid}.layer.1.DenseReluDense.wi_0", # flan-t5 - ), - - MODEL_TENSOR.ENC_FFN_UP: ( - "encoder.block.{bid}.layer.1.DenseReluDense.wi", # t5 - "encoder.block.{bid}.layer.1.DenseReluDense.wi_1", # flan-t5 - ), - - MODEL_TENSOR.ENC_FFN_DOWN: ( - "encoder.block.{bid}.layer.1.DenseReluDense.wo", # t5 - ), - - ############################################################################ - # TODO: these do not belong to block_mappings_cfg - move them to mappings_cfg - MODEL_TENSOR.ENC_OUTPUT_NORM: ( - "encoder.final_layer_norm", # t5 - "layer_norm", # neobert - ), - - MODEL_TENSOR.CLS: ( - "classifier", # jina - "classifier.dense", # roberta - "pre_classifier", # distillbert - "dense", # neobert - ), - - MODEL_TENSOR.CLS_OUT: ( - "classifier.out_proj", # roberta - ), - ############################################################################# - - MODEL_TENSOR.CONVNEXT_DW: ( - "backbone.convnext.{bid}.dwconv", # wavtokenizer - ), - - MODEL_TENSOR.CONVNEXT_NORM: ( - "backbone.convnext.{bid}.norm", # wavtokenizer - ), - - MODEL_TENSOR.CONVNEXT_PW1: ( - "backbone.convnext.{bid}.pwconv1", # wavtokenizer - ), - - MODEL_TENSOR.CONVNEXT_PW2: ( - "backbone.convnext.{bid}.pwconv2", # wavtokenizer - ), - - MODEL_TENSOR.CONVNEXT_GAMMA: ( - "backbone.convnext.{bid}.gamma", # wavtokenizer - ), - - MODEL_TENSOR.POSNET_CONV1: ( - "backbone.posnet.{bid}.conv1", # wavtokenizer - ), - - MODEL_TENSOR.POSNET_CONV2: ( - "backbone.posnet.{bid}.conv2", # wavtokenizer - ), - - MODEL_TENSOR.POSNET_NORM: ( - "backbone.posnet.{bid}.norm", # wavtokenizer - ), - - MODEL_TENSOR.POSNET_NORM1: ( - "backbone.posnet.{bid}.norm1", # wavtokenizer - ), - - MODEL_TENSOR.POSNET_NORM2: ( - "backbone.posnet.{bid}.norm2", # wavtokenizer - ), - - MODEL_TENSOR.POSNET_ATTN_NORM: ( - "backbone.posnet.{bid}.norm", # wavtokenizer - ), - - MODEL_TENSOR.POSNET_ATTN_Q: ( - "backbone.posnet.{bid}.q", # wavtokenizer - ), - - MODEL_TENSOR.POSNET_ATTN_K: ( - "backbone.posnet.{bid}.k", # wavtokenizer - ), - - MODEL_TENSOR.POSNET_ATTN_V: ( - "backbone.posnet.{bid}.v", # wavtokenizer - ), - - MODEL_TENSOR.POSNET_ATTN_OUT: ( - "backbone.posnet.{bid}.proj_out", # wavtokenizer - ), - - ############################################################################# - ## Vision encoder - - MODEL_TENSOR.V_MMPROJ: ( - "multi_modal_projector.linear_{bid}", - "visual.merger.mlp.{bid}", # qwen2vl - ), - - MODEL_TENSOR.V_MMPROJ_FC: ( - "model.connector.modality_projection.proj", # SmolVLM - ), - - MODEL_TENSOR.V_MMPROJ_MLP: ( - "model.mm_projector.mlp.mlp.{bid}", - "vision_model.vision_adapter.mlp.fc{bid}", # llama 4 - "mlp1.{bid}", # InternVL - ), - - MODEL_TENSOR.V_MMPROJ_PEG: ( - "model.mm_projector.peg.peg.{bid}", - ), - - MODEL_TENSOR.V_ENC_EMBD_CLS: ( - "vision_tower.vision_model.embeddings.class_embedding", - "vision_model.class_embedding", # llama 4 - ), - - MODEL_TENSOR.V_ENC_EMBD_PATCH: ( - "vision_tower.vision_model.embeddings.patch_embedding", - "vpm.embeddings.patch_embedding", - "model.vision_model.embeddings.patch_embedding", # SmolVLM - "vision_tower.patch_conv", # pixtral - "vision_model.patch_embedding.linear", # llama 4 - "visual.patch_embed.proj", # qwen2vl - ), - - MODEL_TENSOR.V_ENC_EMBD_POS: ( - "vision_tower.vision_model.embeddings.position_embedding", - "vpm.embeddings.position_embedding", - "model.vision_model.embeddings.position_embedding", # SmolVLM - "vision_model.positional_embedding_vlm", # llama 4 - ), - - MODEL_TENSOR.V_ENC_ATTN_Q: ( - "vision_tower.vision_model.encoder.layers.{bid}.self_attn.q_proj", - "vpm.encoder.layers.{bid}.self_attn.q_proj", - "model.vision_model.encoder.layers.{bid}.self_attn.q_proj", # SmolVLM - "vision_model.model.layers.{bid}.self_attn.q_proj", # llama4 - "vision_tower.transformer.layers.{bid}.attention.q_proj", # pixtral - "visual.blocks.{bid}.attn.q", # qwen2vl, generated - ), - - MODEL_TENSOR.V_ENC_ATTN_Q_NORM: ( - "vision_tower.vision_model.encoder.layers.{bid}.attn.q_norm", # InternVL - ), - - MODEL_TENSOR.V_ENC_ATTN_K: ( - "vision_tower.vision_model.encoder.layers.{bid}.self_attn.k_proj", - "vpm.encoder.layers.{bid}.self_attn.k_proj", - "model.vision_model.encoder.layers.{bid}.self_attn.k_proj", # SmolVLM - "vision_model.model.layers.{bid}.self_attn.k_proj", # llama4 - "vision_tower.transformer.layers.{bid}.attention.k_proj", # pixtral - "visual.blocks.{bid}.attn.k", # qwen2vl, generated - ), - - MODEL_TENSOR.V_ENC_ATTN_K_NORM: ( - "vision_tower.vision_model.encoder.layers.{bid}.attn.k_norm", # InternVL - ), - - MODEL_TENSOR.V_ENC_ATTN_V: ( - "vision_tower.vision_model.encoder.layers.{bid}.self_attn.v_proj", - "vpm.encoder.layers.{bid}.self_attn.v_proj", - "model.vision_model.encoder.layers.{bid}.self_attn.v_proj", # SmolVLM - "vision_model.model.layers.{bid}.self_attn.v_proj", # llama4 - "vision_tower.transformer.layers.{bid}.attention.v_proj", # pixtral - "visual.blocks.{bid}.attn.v", # qwen2vl, generated - ), - - MODEL_TENSOR.V_ENC_INPUT_NORM: ( - "vision_tower.vision_model.encoder.layers.{bid}.layer_norm1", - "vision_tower.vision_model.encoder.layers.{bid}.norm1", # InternVL - "vpm.encoder.layers.{bid}.layer_norm1", - "model.vision_model.encoder.layers.{bid}.layer_norm1", # SmolVLM - "vision_tower.transformer.layers.{bid}.attention_norm", # pixtral - "vision_model.model.layers.{bid}.input_layernorm", # llama4 - "visual.blocks.{bid}.norm1", # qwen2vl - ), - - MODEL_TENSOR.V_ENC_ATTN_O: ( - "vision_tower.vision_model.encoder.layers.{bid}.self_attn.out_proj", - "vision_tower.vision_model.encoder.layers.{bid}.attn.proj", # InternVL - "vpm.encoder.layers.{bid}.self_attn.out_proj", - "model.vision_model.encoder.layers.{bid}.self_attn.out_proj", # SmolVLM - "vision_model.model.layers.{bid}.self_attn.o_proj", # llama4 - "vision_tower.transformer.layers.{bid}.attention.o_proj", # pixtral - "visual.blocks.{bid}.attn.proj", # qwen2vl - ), - - MODEL_TENSOR.V_ENC_POST_ATTN_NORM: ( - "vision_tower.vision_model.encoder.layers.{bid}.layer_norm2", - "vision_tower.vision_model.encoder.layers.{bid}.norm2", # InternVL - "vpm.encoder.layers.{bid}.layer_norm2", - "model.vision_model.encoder.layers.{bid}.layer_norm2", # SmolVLM - "vision_model.model.layers.{bid}.post_attention_layernorm", # llama4 - "vision_tower.transformer.layers.{bid}.ffn_norm", # pixtral - "visual.blocks.{bid}.norm2", # qwen2vl - ), - - MODEL_TENSOR.V_ENC_FFN_UP: ( - "vision_tower.vision_model.encoder.layers.{bid}.mlp.fc1", - "vpm.encoder.layers.{bid}.mlp.fc1", - "model.vision_model.encoder.layers.{bid}.mlp.fc1", # SmolVLM, gemma3 - "vision_tower.transformer.layers.{bid}.feed_forward.up_proj", # pixtral - "vision_model.model.layers.{bid}.mlp.fc1", # llama4 - "visual.blocks.{bid}.mlp.fc1", # qwen2vl - "visual.blocks.{bid}.mlp.up_proj", # qwen2.5vl - ), - - MODEL_TENSOR.V_ENC_FFN_GATE: ( - "vision_tower.transformer.layers.{bid}.feed_forward.gate_proj", # pixtral - "visual.blocks.{bid}.mlp.gate_proj", # qwen2.5vl - ), - - MODEL_TENSOR.V_ENC_FFN_DOWN: ( - "vision_tower.vision_model.encoder.layers.{bid}.mlp.fc2", - "vpm.encoder.layers.{bid}.mlp.fc2", - "model.vision_model.encoder.layers.{bid}.mlp.fc2", # SmolVLM, gemma3 - "vision_tower.transformer.layers.{bid}.feed_forward.down_proj", # pixtral - "vision_model.model.layers.{bid}.mlp.fc2", # llama4 - "visual.blocks.{bid}.mlp.fc2", # qwen2vl - "visual.blocks.{bid}.mlp.down_proj", # qwen2.5vl - ), - - MODEL_TENSOR.V_LAYER_SCALE_1: ( - "vision_tower.vision_model.encoder.layers.{bid}.ls1", # InternVL - ), - - MODEL_TENSOR.V_LAYER_SCALE_2: ( - "vision_tower.vision_model.encoder.layers.{bid}.ls2", # InternVL - ), - - MODEL_TENSOR.V_PRE_NORM: ( - "vision_tower.vision_model.pre_layrnorm", - "vision_tower.ln_pre", # pixtral - "vision_model.layernorm_pre", # llama4 - ), - - MODEL_TENSOR.V_POST_NORM: ( - "vision_tower.vision_model.post_layernorm", - "model.vision_model.post_layernorm", # SmolVLM - "vision_model.layernorm_post", # llama4 - "visual.merger.ln_q", # qwen2vl - ), - - MODEL_TENSOR.V_MM_INP_PROJ: ( - "multi_modal_projector.mm_input_projection", - ), - - MODEL_TENSOR.V_MM_INP_NORM: ( - "multi_modal_projector.norm", - ), - - MODEL_TENSOR.V_MM_SOFT_EMB_NORM: ( - "multi_modal_projector.mm_soft_emb_norm", - ), - - MODEL_TENSOR.V_RESMPL_POS_EMBD_K: ( - "resampler.pos_embed_k", - ), - - MODEL_TENSOR.V_RESMPL_ATTN_Q: ( - "resampler.attn.in_proj_q", # tensor generated from resampler.attn.in_proj - ), - - MODEL_TENSOR.V_RESMPL_ATTN_K: ( - "resampler.attn.in_proj_k", # tensor generated from resampler.attn.in_proj - ), - - MODEL_TENSOR.V_RESMPL_ATTN_V: ( - "resampler.attn.in_proj_v", # tensor generated from resampler.attn.in_proj - ), - - MODEL_TENSOR.V_RESMPL_ATTN_OUT: ( - "resampler.attn.out_proj", - ), - - MODEL_TENSOR.V_RESMPL_KV: ( - "resampler.kv_proj", - ), - - MODEL_TENSOR.V_RESMPL_POST_NORM: ( - "resampler.ln_post", - ), - - MODEL_TENSOR.V_RESMPL_KV_NORM: ( - "resampler.ln_kv", - ), - - MODEL_TENSOR.V_RESMPL_Q_NORM: ( - "resampler.ln_q", - ), - - MODEL_TENSOR.V_RESMPL_PROJ: ( - "resampler.proj", - ), - - MODEL_TENSOR.V_RESMPL_QUERY: ( - "resampler.query", - ), - - MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK: ( - "v.token_embd.img_break", # for pixtral, this is a generated vector - ), - - MODEL_TENSOR.V_MM_PATCH_MERGER: ( - "multi_modal_projector.patch_merger.merging_layer", # mistral small 3.1 - ), - - # audio (mtmd) - - MODEL_TENSOR.A_ENC_EMBD_POS: ( - "audio_tower.embed_positions", # ultravox - ), - - MODEL_TENSOR.A_ENC_CONV1D: ( - "audio_tower.conv{bid}", # ultravox - ), - - MODEL_TENSOR.A_PRE_NORM: (), - - MODEL_TENSOR.A_POST_NORM: ( - "audio_tower.layer_norm", # ultravox - "audio_tower.ln_post", # qwen2omni - ), - - MODEL_TENSOR.A_ENC_ATTN_Q: ( - "audio_tower.layers.{bid}.self_attn.q_proj", # ultravox - ), - - MODEL_TENSOR.A_ENC_ATTN_K: ( - "audio_tower.layers.{bid}.self_attn.k_proj", # ultravox - ), - - MODEL_TENSOR.A_ENC_ATTN_V: ( - "audio_tower.layers.{bid}.self_attn.v_proj", # ultravox - ), - - MODEL_TENSOR.A_ENC_INPUT_NORM: ( - "audio_tower.layers.{bid}.self_attn_layer_norm", # ultravox - ), - - MODEL_TENSOR.A_ENC_OUTPUT: ( - "audio_tower.layers.{bid}.self_attn.out_proj", # ultravox - ), - - MODEL_TENSOR.A_ENC_OUTPUT_NORM: ( - "audio_tower.layers.{bid}.final_layer_norm", # ultravox - ), - - MODEL_TENSOR.A_ENC_FFN_UP: ( - "audio_tower.layers.{bid}.fc1", # ultravox - ), - - MODEL_TENSOR.A_ENC_FFN_GATE: (), - - MODEL_TENSOR.A_ENC_FFN_DOWN: ( - "audio_tower.layers.{bid}.fc2", # ultravox - ), - - # note: some tensors below has "audio." pseudo-prefix, to prevent conflicts with vision tensors - # this prefix is added in the conversion code in modify_tensors() - - MODEL_TENSOR.A_MMPROJ: ( - "audio.multi_modal_projector.linear_{bid}", # ultravox - ), - - MODEL_TENSOR.A_MMPROJ_FC: ( - "audio.multi_modal_projector.linear", # qwen2audio - "audio_tower.proj", # qwen2omni - ), - - MODEL_TENSOR.A_MM_NORM_PRE: ( - "audio.multi_modal_projector.ln_pre", # ultravox - ), - - MODEL_TENSOR.A_MM_NORM_MID: ( - "audio.multi_modal_projector.ln_mid", # ultravox - ), - } - - # architecture-specific block mappings - arch_block_mappings_cfg: dict[MODEL_ARCH, dict[MODEL_TENSOR, tuple[str, ...]]] = { - MODEL_ARCH.ARCTIC: { - MODEL_TENSOR.FFN_NORM: ( - "model.layers.{bid}.residual_layernorm", - ), - MODEL_TENSOR.FFN_NORM_EXP: ( - "model.layers.{bid}.post_attention_layernorm", - ), - }, - } - - mapping: dict[str, tuple[MODEL_TENSOR, str]] - - def __init__(self, arch: MODEL_ARCH, n_blocks: int): - self.mapping = {} - for tensor, keys in self.mappings_cfg.items(): - if tensor not in MODEL_TENSORS[arch]: - continue - tensor_name = TENSOR_NAMES[tensor] - self.mapping[tensor_name] = (tensor, tensor_name) - for key in keys: - self.mapping[key] = (tensor, tensor_name) - if arch in self.arch_block_mappings_cfg: - self.block_mappings_cfg.update(self.arch_block_mappings_cfg[arch]) - for bid in range(n_blocks): - for tensor, keys in self.block_mappings_cfg.items(): - if tensor not in MODEL_TENSORS[arch]: - continue - - tensor_name = TENSOR_NAMES[tensor].format(bid = bid) - self.mapping[tensor_name] = (tensor, tensor_name) - for key in keys: - key = key.format(bid = bid) - self.mapping[key] = (tensor, tensor_name) - - def get_type_and_name(self, key: str, try_suffixes: Sequence[str] = ()) -> tuple[MODEL_TENSOR, str] | None: - result = self.mapping.get(key) - if result is not None: - return result - for suffix in try_suffixes: - if key.endswith(suffix): - result = self.mapping.get(key[:-len(suffix)]) - if result is not None: - return result[0], result[1] + suffix - return None - - def get_name(self, key: str, try_suffixes: Sequence[str] = ()) -> str | None: - result = self.get_type_and_name(key, try_suffixes = try_suffixes) - if result is None: - return None - return result[1] - - def get_type(self, key: str, try_suffixes: Sequence[str] = ()) -> MODEL_TENSOR | None: - result = self.get_type_and_name(key, try_suffixes = try_suffixes) - if result is None: - return None - return result[0] - - def __getitem__(self, key: str) -> str: - try: - return self.mapping[key][1] - except KeyError: - raise KeyError(key) - - def __contains__(self, key: str) -> bool: - return key in self.mapping - - def __repr__(self) -> str: - return repr(self.mapping) - - -def get_tensor_name_map(arch: MODEL_ARCH, n_blocks: int) -> TensorNameMap: - return TensorNameMap(arch, n_blocks) diff --git a/venv/lib/python3.10/site-packages/gguf/utility.py b/venv/lib/python3.10/site-packages/gguf/utility.py deleted file mode 100644 index 00adcbc937398d1e7a9d4b7159a54052bb79ae64..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/utility.py +++ /dev/null @@ -1,264 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Literal - -import os -import json - - -def fill_templated_filename(filename: str, output_type: str | None) -> str: - # Given a file name fill in any type templates e.g. 'some-model-name.{ftype}.gguf' - ftype_lowercase: str = output_type.lower() if output_type is not None else "" - ftype_uppercase: str = output_type.upper() if output_type is not None else "" - return filename.format(ftype_lowercase, - outtype=ftype_lowercase, ftype=ftype_lowercase, - OUTTYPE=ftype_uppercase, FTYPE=ftype_uppercase) - - -def model_weight_count_rounded_notation(model_params_count: int, min_digits: int = 2) -> str: - if model_params_count > 1e12 : - # Trillions Of Parameters - scaled_model_params = model_params_count * 1e-12 - scale_suffix = "T" - elif model_params_count > 1e9 : - # Billions Of Parameters - scaled_model_params = model_params_count * 1e-9 - scale_suffix = "B" - elif model_params_count > 1e6 : - # Millions Of Parameters - scaled_model_params = model_params_count * 1e-6 - scale_suffix = "M" - else: - # Thousands Of Parameters - scaled_model_params = model_params_count * 1e-3 - scale_suffix = "K" - - fix = max(min_digits - len(str(round(scaled_model_params)).lstrip('0')), 0) - - return f"{scaled_model_params:.{fix}f}{scale_suffix}" - - -def size_label(total_params: int, shared_params: int, expert_params: int, expert_count: int) -> str: - - if expert_count > 0: - pretty_size = model_weight_count_rounded_notation(abs(shared_params) + abs(expert_params), min_digits=2) - size_class = f"{expert_count}x{pretty_size}" - else: - size_class = model_weight_count_rounded_notation(abs(total_params), min_digits=2) - - return size_class - - -def naming_convention(model_name: str | None, base_name: str | None, finetune_string: str | None, version_string: str | None, size_label: str | None, output_type: str | None, model_type: Literal['vocab', 'LoRA'] | None = None) -> str: - # Reference: https://github.com/ggml-org/ggml/blob/master/docs/gguf.md#gguf-naming-convention - - if base_name is not None: - name = base_name.strip().replace(' ', '-').replace('/', '-') - elif model_name is not None: - name = model_name.strip().replace(' ', '-').replace('/', '-') - else: - name = "ggml-model" - - parameters = f"-{size_label}" if size_label is not None else "" - - finetune = f"-{finetune_string.strip().replace(' ', '-')}" if finetune_string is not None else "" - - version = f"-{version_string.strip().replace(' ', '-')}" if version_string is not None else "" - - encoding = f"-{output_type.strip().replace(' ', '-').upper()}" if output_type is not None else "" - - kind = f"-{model_type.strip().replace(' ', '-')}" if model_type is not None else "" - - return f"{name}{parameters}{finetune}{version}{encoding}{kind}" - - -@dataclass -class RemoteTensor: - dtype: str - shape: tuple[int, ...] - offset_start: int - size: int - url: str - - def data(self) -> bytearray: - # TODO: handle request errors (maybe with limited retries?) - # NOTE: using a bytearray, otherwise PyTorch complains the buffer is not writeable - data = bytearray(SafetensorRemote.get_data_by_range(url=self.url, start=self.offset_start, size=self.size)) - return data - - -class SafetensorRemote: - """ - Uility class to handle remote safetensor files. - This class is designed to work with Hugging Face model repositories. - - Example (one model has single safetensor file, the other has multiple): - for model_id in ["ngxson/TEST-Tiny-Llama4", "Qwen/Qwen2.5-7B-Instruct"]: - tensors = SafetensorRemote.get_list_tensors_hf_model(model_id) - print(tensors) - - Example reading tensor data: - tensors = SafetensorRemote.get_list_tensors_hf_model(model_id) - for name, meta in tensors.items(): - dtype, shape, offset_start, size, remote_safetensor_url = meta - # read the tensor data - data = SafetensorRemote.get_data_by_range(remote_safetensor_url, offset_start, size) - print(data) - """ - - BASE_DOMAIN = "https://huggingface.co" - ALIGNMENT = 8 # bytes - - @classmethod - def get_list_tensors_hf_model(cls, model_id: str) -> dict[str, RemoteTensor]: - """ - Get list of tensors from a Hugging Face model repository. - - Returns a dictionary of tensor names and their metadata. - Each tensor is represented as a tuple of (dtype, shape, offset_start, size, remote_safetensor_url) - """ - # case 1: model has only one single model.safetensor file - is_single_file = cls.check_file_exist(f"{cls.BASE_DOMAIN}/{model_id}/resolve/main/model.safetensors") - if is_single_file: - url = f"{cls.BASE_DOMAIN}/{model_id}/resolve/main/model.safetensors" - return cls.get_list_tensors(url) - - # case 2: model has multiple files - index_url = f"{cls.BASE_DOMAIN}/{model_id}/resolve/main/model.safetensors.index.json" - is_multiple_files = cls.check_file_exist(index_url) - if is_multiple_files: - # read the index file - index_data = cls.get_data_by_range(index_url, 0) - index_str = index_data.decode('utf-8') - index_json = json.loads(index_str) - assert index_json.get("weight_map") is not None, "weight_map not found in index file" - weight_map = index_json["weight_map"] - # get the list of files - all_files = list(set(weight_map.values())) - all_files.sort() # make sure we load shard files in order - # get the list of tensors - tensors: dict[str, RemoteTensor] = {} - for file in all_files: - url = f"{cls.BASE_DOMAIN}/{model_id}/resolve/main/{file}" - for key, val in cls.get_list_tensors(url).items(): - tensors[key] = val - return tensors - - raise ValueError(f"Model {model_id} does not have any safetensor files") - - @classmethod - def get_list_tensors(cls, url: str) -> dict[str, RemoteTensor]: - """ - Get list of tensors from a remote safetensor file. - - Returns a dictionary of tensor names and their metadata. - Each tensor is represented as a tuple of (dtype, shape, offset_start, size) - """ - metadata, data_start_offset = cls.get_metadata(url) - res: dict[str, RemoteTensor] = {} - - for name, meta in metadata.items(): - if name == "__metadata__": - continue - if not isinstance(meta, dict): - raise ValueError(f"Invalid metadata for tensor '{name}': {meta}") - try: - dtype = meta["dtype"] - shape = meta["shape"] - offset_start_relative, offset_end_relative = meta["data_offsets"] - size = offset_end_relative - offset_start_relative - offset_start = data_start_offset + offset_start_relative - res[name] = RemoteTensor(dtype=dtype, shape=tuple(shape), offset_start=offset_start, size=size, url=url) - except KeyError as e: - raise ValueError(f"Missing key in metadata for tensor '{name}': {e}, meta = {meta}") - - return res - - @classmethod - def get_metadata(cls, url: str) -> tuple[dict, int]: - """ - Get JSON metadata from a remote safetensor file. - - Returns tuple of (metadata, data_start_offset) - """ - # Request first 5MB of the file (hopefully enough for metadata) - read_size = 5 * 1024 * 1024 - raw_data = cls.get_data_by_range(url, 0, read_size) - - # Parse header - # First 8 bytes contain the metadata length as u64 little-endian - if len(raw_data) < 8: - raise ValueError("Not enough data to read metadata size") - metadata_length = int.from_bytes(raw_data[:8], byteorder='little') - - # Calculate the data start offset - data_start_offset = 8 + metadata_length - alignment = SafetensorRemote.ALIGNMENT - if data_start_offset % alignment != 0: - data_start_offset += alignment - (data_start_offset % alignment) - - # Check if we have enough data to read the metadata - if len(raw_data) < 8 + metadata_length: - raise ValueError(f"Could not read complete metadata. Need {8 + metadata_length} bytes, got {len(raw_data)}") - - # Extract metadata bytes and parse as JSON - metadata_bytes = raw_data[8:8 + metadata_length] - metadata_str = metadata_bytes.decode('utf-8') - try: - metadata = json.loads(metadata_str) - return metadata, data_start_offset - except json.JSONDecodeError as e: - raise ValueError(f"Failed to parse safetensor metadata as JSON: {e}") - - @classmethod - def get_data_by_range(cls, url: str, start: int, size: int = -1) -> bytes: - """ - Get raw byte data from a remote file by range. - If size is not specified, it will read the entire file. - """ - import requests - from urllib.parse import urlparse - - parsed_url = urlparse(url) - if not parsed_url.scheme or not parsed_url.netloc: - raise ValueError(f"Invalid URL: {url}") - - headers = cls._get_request_headers() - if size > -1: - headers["Range"] = f"bytes={start}-{start + size}" - response = requests.get(url, allow_redirects=True, headers=headers) - response.raise_for_status() - - # Get raw byte data - return response.content[slice(size if size > -1 else None)] - - @classmethod - def check_file_exist(cls, url: str) -> bool: - """ - Check if a file exists at the given URL. - Returns True if the file exists, False otherwise. - """ - import requests - from urllib.parse import urlparse - - parsed_url = urlparse(url) - if not parsed_url.scheme or not parsed_url.netloc: - raise ValueError(f"Invalid URL: {url}") - - try: - headers = cls._get_request_headers() - headers["Range"] = "bytes=0-0" - response = requests.head(url, allow_redirects=True, headers=headers) - # Success (2xx) or redirect (3xx) - return 200 <= response.status_code < 400 - except requests.RequestException: - return False - - @classmethod - def _get_request_headers(cls) -> dict[str, str]: - """Prepare common headers for requests.""" - headers = {"User-Agent": "convert_hf_to_gguf"} - if os.environ.get("HF_TOKEN"): - headers["Authorization"] = f"Bearer {os.environ['HF_TOKEN']}" - return headers diff --git a/venv/lib/python3.10/site-packages/gguf/vocab.py b/venv/lib/python3.10/site-packages/gguf/vocab.py deleted file mode 100644 index 44d066ee75a7ece48dabb168c5054fb4053f5d48..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/gguf/vocab.py +++ /dev/null @@ -1,498 +0,0 @@ -from __future__ import annotations - -import re -import logging -import json -import os -from pathlib import Path -from typing import Any, Callable, Sequence, Mapping, Iterable, Protocol, ClassVar, runtime_checkable - -try: - from sentencepiece import SentencePieceProcessor -except ImportError: - SentencePieceProcessor = None - -import gguf - -from .gguf_writer import GGUFWriter - -logger = logging.getLogger(__name__) - - -class SpecialVocab: - merges: list[str] - add_special_token: dict[str, bool] - special_token_ids: dict[str, int] - chat_template: str | Sequence[Mapping[str, str]] | None - - def __init__( - self, path: str | os.PathLike[str], load_merges: bool = False, - special_token_types: Iterable[str] | None = None, - n_vocab: int | None = None, - ): - self.special_token_ids = {} - self.add_special_token = {} - self.n_vocab = n_vocab - self.load_merges = load_merges - self.merges = [] - self.chat_template = None - if special_token_types is not None: - self.special_token_types = special_token_types - else: - self.special_token_types = ('bos', 'eos', 'unk', 'sep', 'pad', 'cls', 'mask') - self._load(Path(path)) - - def __repr__(self) -> str: - return ''.format( - len(self.merges), self.special_token_ids or "unset", self.add_special_token or "unset", - ) - - def add_to_gguf(self, gw: GGUFWriter, quiet: bool = False) -> None: - if self.merges: - if not quiet: - logger.info(f'Adding {len(self.merges)} merge(s).') - gw.add_token_merges(self.merges) - elif self.load_merges: - logger.warning('Adding merges requested but no merges found, output may be non-functional.') - for typ, tokid in self.special_token_ids.items(): - id_handler: Callable[[int], None] | None = getattr(gw, f'add_{typ}_token_id', None) - if id_handler is None: - logger.warning(f'No handler for special token type {typ} with id {tokid} - skipping') - continue - if not quiet: - logger.info(f'Setting special token type {typ} to {tokid}') - id_handler(tokid) - for typ, value in self.add_special_token.items(): - add_handler: Callable[[bool], None] | None = getattr(gw, f'add_add_{typ}_token', None) - if add_handler is None: - logger.warning(f'No handler for add_{typ}_token with value {value} - skipping') - continue - if not quiet: - logger.info(f'Setting add_{typ}_token to {value}') - add_handler(value) - if self.chat_template is not None: - if not quiet: - logger.info(f'Setting chat_template to {self.chat_template}') - gw.add_chat_template(self.chat_template) - - def _load(self, path: Path) -> None: - self._try_load_from_tokenizer_json(path) - self._try_load_from_config_json(path) - if self.load_merges and not self.merges: - self._try_load_merges_txt(path) - - def _try_load_merges_txt(self, path: Path) -> bool: - merges_file = path / 'merges.txt' - if not merges_file.is_file(): - return False - with open(merges_file, 'r', encoding = 'utf-8') as fp: - first_line = next(fp, '').strip() - if not first_line.startswith('#'): - fp.seek(0) - line_num = 0 - else: - line_num = 1 - merges = [] - for line in fp: - line_num += 1 - line = line.strip() - if not line: - continue - parts = line.split(None, 3) - if len(parts) != 2: - logger.warning(f'{merges_file.name}: Line {line_num}: Entry malformed, ignoring') - continue - merges.append(f'{parts[0]} {parts[1]}') - self.merges = merges - return True - - def _set_special_token(self, typ: str, tid: Any) -> None: - if not isinstance(tid, int): - return - if tid < 0: - raise ValueError(f'invalid value for special token type {typ}: {tid}') - if self.n_vocab is None or tid < self.n_vocab: - if typ in self.special_token_ids: - return - self.special_token_ids[typ] = tid - return - logger.warning(f'Special token type {typ}, id {tid} out of range, must be under {self.n_vocab} - skipping') - - def _try_load_from_tokenizer_json(self, path: Path) -> bool: - tokenizer_file = path / 'tokenizer.json' - if tokenizer_file.is_file(): - with open(tokenizer_file, encoding = 'utf-8') as f: - tokenizer = json.load(f) - if self.load_merges: - merges = tokenizer.get('model', {}).get('merges') - if isinstance(merges, list) and merges: - if isinstance(merges[0], str): - self.merges = merges - elif isinstance(merges[0], list) and len(merges[0]) == 2 and isinstance(merges[0][0], str): - # New format since transformers 4.45 to support spaces in merges - # ref: https://github.com/ggml-org/llama.cpp/issues/9692 - # TODO: internally store as the new format instead of converting to old - if any(' ' in s for pair in merges for s in pair): - logger.warning(f'Spaces in merges detected, encoding as {chr(ord(" ") + 256)!r}') - self.merges = [ - ' '.join( - [ - # ensure the spaces are properly encoded - ''.join( - chr(ord(c) + 256) if c == ' ' else c - for c in part - ) - for part in pair - ] - ) - for pair in merges - ] - else: - raise ValueError("Unknown tokenizer merges format") - added_tokens = tokenizer.get('added_tokens', {}) - else: - added_tokens = {} - tokenizer_config_file = path / 'tokenizer_config.json' - if not tokenizer_config_file.is_file(): - return True - with open(tokenizer_config_file, encoding = 'utf-8') as f: - tokenizer_config = json.load(f) - chat_template_alt = None - chat_template_file = path / 'chat_template.json' - if chat_template_file.is_file(): - with open(chat_template_file, encoding = 'utf-8') as f: - chat_template_alt = json.load(f).get('chat_template') - chat_template = tokenizer_config.get('chat_template', chat_template_alt) - if chat_template is None or isinstance(chat_template, (str, list)): - self.chat_template = chat_template - else: - logger.warning(f'Bad type for chat_template field in {tokenizer_config_file!r} - ignoring') - for typ in self.special_token_types: - add_entry = tokenizer_config.get(f'add_{typ}_token') - if isinstance(add_entry, bool): - self.add_special_token[typ] = add_entry - entry = tokenizer_config.get(f'{typ}_token') - if isinstance(entry, str): - tc_content = entry - elif isinstance(entry, dict): - entry_content = entry.get('content') - if not isinstance(entry_content, str): - continue - tc_content = entry_content - else: - continue - # We only need the first match here. - maybe_token_id = next( - (atok.get('id') for atok in added_tokens if atok.get('content') == tc_content), - None, - ) - self._set_special_token(typ, maybe_token_id) - return True - - def _try_load_from_config_json(self, path: Path) -> bool: - config_file = path / 'config.json' - if not config_file.is_file(): - return False - with open(config_file, encoding = 'utf-8') as f: - config = json.load(f) - for typ in self.special_token_types: - self._set_special_token(typ, config.get(f'{typ}_token_id')) - return True - - -@runtime_checkable -class BaseVocab(Protocol): - tokenizer_model: ClassVar[str] - name: ClassVar[str] - - -@runtime_checkable -class Vocab(BaseVocab, Protocol): - vocab_size: int - added_tokens_dict: dict[str, int] - added_tokens_list: list[str] - fname_tokenizer: Path - - def __init__(self, base_path: Path): ... - def all_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: ... - - -class NoVocab(BaseVocab): - tokenizer_model = "no_vocab" - name = "no_vocab" - - def __repr__(self) -> str: - return "" - - -class BpeVocab(Vocab): - tokenizer_model = "gpt2" - name = "bpe" - - def __init__(self, base_path: Path): - added_tokens: dict[str, int] = {} - - if (fname_tokenizer := base_path / 'vocab.json').exists(): - # "slow" tokenizer - with open(fname_tokenizer, encoding="utf-8") as f: - self.vocab = json.load(f) - - try: - # FIXME: Verify that added tokens here _cannot_ overlap with the main vocab. - with open(base_path / 'added_tokens.json', encoding="utf-8") as f: - added_tokens = json.load(f) - except FileNotFoundError: - pass - else: - # "fast" tokenizer - fname_tokenizer = base_path / 'tokenizer.json' - - # if this fails, FileNotFoundError propagates to caller - with open(fname_tokenizer, encoding="utf-8") as f: - tokenizer_json = json.load(f) - - tokenizer_model: dict[str, Any] = tokenizer_json['model'] - if ( - tokenizer_model['type'] != 'BPE' or tokenizer_model.get('byte_fallback', False) - or tokenizer_json['decoder']['type'] != 'ByteLevel' - ): - raise FileNotFoundError('Cannot find GPT-2 BPE tokenizer') - - self.vocab = tokenizer_model["vocab"] - - if (added := tokenizer_json.get('added_tokens')) is not None: - # Added tokens here can be duplicates of the main vocabulary. - added_tokens = {item['content']: item['id'] - for item in added - if item['content'] not in self.vocab} - - vocab_size = len(self.vocab) - expected_ids = list(range(vocab_size, vocab_size + len(added_tokens))) - actual_ids = sorted(added_tokens.values()) - if expected_ids != actual_ids: - expected_end_id = vocab_size + len(actual_ids) - 1 - raise ValueError(f"Expected the {len(actual_ids)} added token ID(s) to be sequential in the range " - f"{vocab_size} - {expected_end_id}; got {actual_ids}") - - items = sorted(added_tokens.items(), key=lambda text_idx: text_idx[1]) - self.added_tokens_dict = added_tokens - self.added_tokens_list = [text for (text, idx) in items] - self.vocab_size_base = vocab_size - self.vocab_size = self.vocab_size_base + len(self.added_tokens_list) - self.fname_tokenizer = fname_tokenizer - - def bpe_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: - reverse_vocab = {id: encoded_tok for encoded_tok, id in self.vocab.items()} - - for i, _ in enumerate(self.vocab): - yield reverse_vocab[i], 0.0, gguf.TokenType.NORMAL - - def added_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: - for text in self.added_tokens_list: - score = -1000.0 - yield text.encode("utf-8"), score, gguf.TokenType.CONTROL - - def all_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: - yield from self.bpe_tokens() - yield from self.added_tokens() - - def __repr__(self) -> str: - return f"" - - -class SentencePieceVocab(Vocab): - tokenizer_model = "llama" - name = "spm" - - def __init__(self, base_path: Path): - if SentencePieceProcessor is None: - raise RuntimeError("sentencepiece is not installed") - - added_tokens: dict[str, int] = {} - if (fname_tokenizer := base_path / 'tokenizer.model').exists(): - # normal location - try: - with open(base_path / 'added_tokens.json', encoding="utf-8") as f: - added_tokens = json.load(f) - except FileNotFoundError: - pass - elif not (fname_tokenizer := base_path.parent / 'tokenizer.model').exists(): - # not found in alternate location either - raise FileNotFoundError('Cannot find tokenizer.model') - - self.sentencepiece_tokenizer = SentencePieceProcessor() - self.sentencepiece_tokenizer.LoadFromFile(str(fname_tokenizer)) - vocab_size = self.sentencepiece_tokenizer.vocab_size() - - new_tokens = {id: piece for piece, id in added_tokens.items() if id >= vocab_size} - expected_new_ids = list(range(vocab_size, vocab_size + len(new_tokens))) - actual_new_ids = sorted(new_tokens.keys()) - - if expected_new_ids != actual_new_ids: - raise ValueError(f"Expected new token IDs {expected_new_ids} to be sequential; got {actual_new_ids}") - - # Token pieces that were added to the base vocabulary. - self.added_tokens_dict = added_tokens - self.added_tokens_list = [new_tokens[id] for id in actual_new_ids] - self.vocab_size_base = vocab_size - self.vocab_size = self.vocab_size_base + len(self.added_tokens_list) - self.fname_tokenizer = fname_tokenizer - - def sentencepiece_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: - tokenizer = self.sentencepiece_tokenizer - for i in range(tokenizer.vocab_size()): - piece = tokenizer.IdToPiece(i) - text = piece.encode("utf-8") - score: float = tokenizer.GetScore(i) - - toktype = gguf.TokenType.NORMAL - if tokenizer.IsUnknown(i): - toktype = gguf.TokenType.UNKNOWN - if tokenizer.IsControl(i): - toktype = gguf.TokenType.CONTROL - - # NOTE: I think added_tokens are user defined. - # ref: https://github.com/google/sentencepiece/blob/master/src/sentencepiece_model.proto - # if tokenizer.is_user_defined(i): toktype = gguf.TokenType.USER_DEFINED - - if tokenizer.IsUnused(i): - toktype = gguf.TokenType.UNUSED - if tokenizer.IsByte(i): - toktype = gguf.TokenType.BYTE - - yield text, score, toktype - - def added_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: - for text in self.added_tokens_list: - score = -1000.0 - yield text.encode("utf-8"), score, gguf.TokenType.USER_DEFINED - - def all_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: - yield from self.sentencepiece_tokens() - yield from self.added_tokens() - - def __repr__(self) -> str: - return f"" - - -class LlamaHfVocab(Vocab): - tokenizer_model = "llama" - name = "hfft" - - def __init__(self, base_path: Path): - fname_tokenizer = base_path / 'tokenizer.json' - # if this fails, FileNotFoundError propagates to caller - with open(fname_tokenizer, encoding='utf-8') as f: - tokenizer_json = json.load(f) - - # pre-check so we know if we need transformers - tokenizer_model: dict[str, Any] = tokenizer_json['model'] - is_llama3 = ( - tokenizer_model['type'] == 'BPE' and tokenizer_model.get('ignore_merges', False) - and not tokenizer_model.get('byte_fallback', True) - ) - if is_llama3: - raise TypeError('Llama 3 must be converted with BpeVocab') - - if not is_llama3 and ( - tokenizer_model['type'] != 'BPE' or not tokenizer_model.get('byte_fallback', False) - or tokenizer_json['decoder']['type'] != 'Sequence' - ): - raise FileNotFoundError('Cannot find Llama BPE tokenizer') - - try: - from transformers import AutoTokenizer - except ImportError as e: - raise ImportError( - "To use LlamaHfVocab, please install the `transformers` package. " - "You can install it with `pip install transformers`." - ) from e - - # Allow the tokenizer to default to slow or fast versions. - # Explicitly set tokenizer to use local paths. - self.tokenizer = AutoTokenizer.from_pretrained( - base_path, - cache_dir=base_path, - local_files_only=True, - ) - assert self.tokenizer.is_fast # assume tokenizer.json is used - - # Initialize lists and dictionaries for added tokens - self.added_tokens_list = [] - self.added_tokens_dict = dict() - self.added_tokens_ids = set() - - # Process added tokens - for tok, tokidx in sorted( - self.tokenizer.get_added_vocab().items(), key=lambda x: x[1] - ): - # Only consider added tokens that are not in the base vocabulary - if tokidx >= self.tokenizer.vocab_size: - self.added_tokens_list.append(tok) - self.added_tokens_dict[tok] = tokidx - self.added_tokens_ids.add(tokidx) - - # Store special tokens and their IDs - self.specials = { - tok: self.tokenizer.get_vocab()[tok] - for tok in self.tokenizer.all_special_tokens - } - self.special_ids = set(self.tokenizer.all_special_ids) - - # Set vocabulary sizes - self.vocab_size_base = self.tokenizer.vocab_size - self.vocab_size = self.vocab_size_base + len(self.added_tokens_list) - - self.fname_tokenizer = fname_tokenizer - - def hf_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: - reverse_vocab = { - id: encoded_tok for encoded_tok, id in self.tokenizer.get_vocab().items() - } - - for token_id in range(self.vocab_size_base): - # Skip processing added tokens here - if token_id in self.added_tokens_ids: - continue - - # Convert token text to bytes - token_text = reverse_vocab[token_id].encode("utf-8") - - # Yield token text, score, and type - yield token_text, self.get_token_score(token_id), self.get_token_type( - token_id, token_text, self.special_ids # Reuse already stored special IDs - ) - - def get_token_type(self, token_id: int, token_text: bytes, special_ids: set[int]) -> gguf.TokenType: - # Special case for byte tokens - if re.fullmatch(br"<0x[0-9A-Fa-f]{2}>", token_text): - return gguf.TokenType.BYTE - - # Determine token type based on whether it's a special token - return gguf.TokenType.CONTROL if token_id in special_ids else gguf.TokenType.NORMAL - - def get_token_score(self, token_id: int) -> float: - # Placeholder for actual logic to determine the token's score - # This needs to be implemented based on specific requirements - return -1000.0 # Default score - - def added_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: - for text in self.added_tokens_list: - if text in self.specials: - toktype = self.get_token_type(self.specials[text], b'', self.special_ids) - score = self.get_token_score(self.specials[text]) - else: - toktype = gguf.TokenType.USER_DEFINED - score = -1000.0 - - yield text.encode("utf-8"), score, toktype - - def has_newline_token(self): - return "<0x0A>" in self.tokenizer.vocab or "\n" in self.tokenizer.vocab - - def all_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: - yield from self.hf_tokens() - yield from self.added_tokens() - - def __repr__(self) -> str: - return f"" diff --git a/venv/lib/python3.10/site-packages/git/__init__.py b/venv/lib/python3.10/site-packages/git/__init__.py deleted file mode 100644 index d94ef87d4b1dcaf7444d8e4cfb89b3f66d885904..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/git/__init__.py +++ /dev/null @@ -1,300 +0,0 @@ -# Copyright (C) 2008, 2009 Michael Trier (mtrier@gmail.com) and contributors -# -# This module is part of GitPython and is released under the -# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ - -# @PydevCodeAnalysisIgnore - -__all__ = [ - "Actor", - "AmbiguousObjectName", - "BadName", - "BadObject", - "BadObjectType", - "BaseIndexEntry", - "Blob", - "BlobFilter", - "BlockingLockFile", - "CacheError", - "CheckoutError", - "CommandError", - "Commit", - "Diff", - "DiffConstants", - "DiffIndex", - "Diffable", - "FetchInfo", - "Git", - "GitCmdObjectDB", - "GitCommandError", - "GitCommandNotFound", - "GitConfigParser", - "GitDB", - "GitError", - "HEAD", - "Head", - "HookExecutionError", - "INDEX", - "IndexEntry", - "IndexFile", - "IndexObject", - "InvalidDBRoot", - "InvalidGitRepositoryError", - "List", # Deprecated - import this from `typing` instead. - "LockFile", - "NULL_TREE", - "NoSuchPathError", - "ODBError", - "Object", - "Optional", # Deprecated - import this from `typing` instead. - "ParseError", - "PathLike", - "PushInfo", - "RefLog", - "RefLogEntry", - "Reference", - "Remote", - "RemoteProgress", - "RemoteReference", - "Repo", - "RepositoryDirtyError", - "RootModule", - "RootUpdateProgress", - "Sequence", # Deprecated - import from `typing`, or `collections.abc` in 3.9+. - "StageType", - "Stats", - "Submodule", - "SymbolicReference", - "TYPE_CHECKING", # Deprecated - import this from `typing` instead. - "Tag", - "TagObject", - "TagReference", - "Tree", - "TreeModifier", - "Tuple", # Deprecated - import this from `typing` instead. - "Union", # Deprecated - import this from `typing` instead. - "UnmergedEntriesError", - "UnsafeOptionError", - "UnsafeProtocolError", - "UnsupportedOperation", - "UpdateProgress", - "WorkTreeRepositoryUnsupported", - "refresh", - "remove_password_if_present", - "rmtree", - "safe_decode", - "to_hex_sha", -] - -__version__ = '3.1.44' - -from typing import Any, List, Optional, Sequence, TYPE_CHECKING, Tuple, Union - -if TYPE_CHECKING: - from types import ModuleType - -import warnings - -from gitdb.util import to_hex_sha - -from git.exc import ( - AmbiguousObjectName, - BadName, - BadObject, - BadObjectType, - CacheError, - CheckoutError, - CommandError, - GitCommandError, - GitCommandNotFound, - GitError, - HookExecutionError, - InvalidDBRoot, - InvalidGitRepositoryError, - NoSuchPathError, - ODBError, - ParseError, - RepositoryDirtyError, - UnmergedEntriesError, - UnsafeOptionError, - UnsafeProtocolError, - UnsupportedOperation, - WorkTreeRepositoryUnsupported, -) -from git.types import PathLike - -try: - from git.compat import safe_decode # @NoMove - from git.config import GitConfigParser # @NoMove - from git.objects import ( # @NoMove - Blob, - Commit, - IndexObject, - Object, - RootModule, - RootUpdateProgress, - Submodule, - TagObject, - Tree, - TreeModifier, - UpdateProgress, - ) - from git.refs import ( # @NoMove - HEAD, - Head, - RefLog, - RefLogEntry, - Reference, - RemoteReference, - SymbolicReference, - Tag, - TagReference, - ) - from git.diff import ( # @NoMove - INDEX, - NULL_TREE, - Diff, - DiffConstants, - DiffIndex, - Diffable, - ) - from git.db import GitCmdObjectDB, GitDB # @NoMove - from git.cmd import Git # @NoMove - from git.repo import Repo # @NoMove - from git.remote import FetchInfo, PushInfo, Remote, RemoteProgress # @NoMove - from git.index import ( # @NoMove - BaseIndexEntry, - BlobFilter, - CheckoutError, - IndexEntry, - IndexFile, - StageType, - # NOTE: This tells type checkers what util resolves to. We delete it, and it is - # really resolved by __getattr__, which warns. See below on what to use instead. - util, - ) - from git.util import ( # @NoMove - Actor, - BlockingLockFile, - LockFile, - Stats, - remove_password_if_present, - rmtree, - ) -except GitError as _exc: - raise ImportError("%s: %s" % (_exc.__class__.__name__, _exc)) from _exc - - -def _warned_import(message: str, fullname: str) -> "ModuleType": - import importlib - - warnings.warn(message, DeprecationWarning, stacklevel=3) - return importlib.import_module(fullname) - - -def _getattr(name: str) -> Any: - # TODO: If __version__ is made dynamic and lazily fetched, put that case right here. - - if name == "util": - return _warned_import( - "The expression `git.util` and the import `from git import util` actually " - "reference git.index.util, and not the git.util module accessed in " - '`from git.util import XYZ` or `sys.modules["git.util"]`. This potentially ' - "confusing behavior is currently preserved for compatibility, but may be " - "changed in the future and should not be relied on.", - fullname="git.index.util", - ) - - for names, prefix in ( - ({"head", "log", "reference", "symbolic", "tag"}, "git.refs"), - ({"base", "fun", "typ"}, "git.index"), - ): - if name not in names: - continue - - fullname = f"{prefix}.{name}" - - return _warned_import( - f"{__name__}.{name} is a private alias of {fullname} and subject to " - f"immediate removal. Use {fullname} instead.", - fullname=fullname, - ) - - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -if not TYPE_CHECKING: - # NOTE: The expression `git.util` gives git.index.util and `from git import util` - # imports git.index.util, NOT git.util. It may not be feasible to change this until - # the next major version, to avoid breaking code inadvertently relying on it. - # - # - If git.index.util *is* what you want, use (or import from) that, to avoid - # confusion. - # - # - To use the "real" git.util module, write `from git.util import ...`, or if - # necessary access it as `sys.modules["git.util"]`. - # - # Note also that `import git.util` technically imports the "real" git.util... but - # the *expression* `git.util` after doing so is still git.index.util! - # - # (This situation differs from that of other indirect-submodule imports that are - # unambiguously non-public and subject to immediate removal. Here, the public - # git.util module, though different, makes less discoverable that the expression - # `git.util` refers to a non-public attribute of the git module.) - # - # This had originally come about by a wildcard import. Now that all intended imports - # are explicit, the intuitive but potentially incompatible binding occurs due to the - # usual rules for Python submodule bindings. So for now we replace that binding with - # git.index.util, delete that, and let __getattr__ handle it and issue a warning. - # - # For the same runtime behavior, it would be enough to forgo importing util, and - # delete util as created naturally; __getattr__ would behave the same. But type - # checkers would not know what util refers to when accessed as an attribute of git. - del util - - # This is "hidden" to preserve static checking for undefined/misspelled attributes. - __getattr__ = _getattr - -# { Initialize git executable path - -GIT_OK = None - - -def refresh(path: Optional[PathLike] = None) -> None: - """Convenience method for setting the git executable path. - - :param path: - Optional path to the Git executable. If not absolute, it is resolved - immediately, relative to the current directory. - - :note: - The `path` parameter is usually omitted and cannot be used to specify a custom - command whose location is looked up in a path search on each call. See - :meth:`Git.refresh ` for details on how to achieve this. - - :note: - This calls :meth:`Git.refresh ` and sets other global - configuration according to the effect of doing so. As such, this function should - usually be used instead of using :meth:`Git.refresh ` or - :meth:`FetchInfo.refresh ` directly. - - :note: - This function is called automatically, with no arguments, at import time. - """ - global GIT_OK - GIT_OK = False - - if not Git.refresh(path=path): - return - if not FetchInfo.refresh(): # noqa: F405 - return # type: ignore[unreachable] - - GIT_OK = True - - -try: - refresh() -except Exception as _exc: - raise ImportError("Failed to initialize: {0}".format(_exc)) from _exc - -# } END initialize git executable path diff --git a/venv/lib/python3.10/site-packages/git/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/git/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 4dafbcf93d158c1eb7e0b70a48bb76356693d07b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/git/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/git/__pycache__/cmd.cpython-310.pyc b/venv/lib/python3.10/site-packages/git/__pycache__/cmd.cpython-310.pyc deleted file mode 100644 index 713463324f825dd8eb7247a96f6bb22776888313..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/git/__pycache__/cmd.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/git/__pycache__/compat.cpython-310.pyc b/venv/lib/python3.10/site-packages/git/__pycache__/compat.cpython-310.pyc deleted file mode 100644 index 6a1a40793af4c29ec9388207b03f82d91e566c17..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/git/__pycache__/compat.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/git/__pycache__/config.cpython-310.pyc b/venv/lib/python3.10/site-packages/git/__pycache__/config.cpython-310.pyc deleted file mode 100644 index cf9e0304ab888a983b1facc080198540ce88faa5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/git/__pycache__/config.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/git/__pycache__/db.cpython-310.pyc b/venv/lib/python3.10/site-packages/git/__pycache__/db.cpython-310.pyc deleted file mode 100644 index 9a0c5706223063632552b0c31b39f822114d9ede..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/git/__pycache__/db.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/git/__pycache__/diff.cpython-310.pyc b/venv/lib/python3.10/site-packages/git/__pycache__/diff.cpython-310.pyc deleted file mode 100644 index 9b4c9c4bbaef238a25c3eb30d363eaef329a325c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/git/__pycache__/diff.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/git/__pycache__/exc.cpython-310.pyc b/venv/lib/python3.10/site-packages/git/__pycache__/exc.cpython-310.pyc deleted file mode 100644 index 7637175fbe6e96f9dd8349b7ab81bc6de4b981e2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/git/__pycache__/exc.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/git/__pycache__/remote.cpython-310.pyc b/venv/lib/python3.10/site-packages/git/__pycache__/remote.cpython-310.pyc deleted file mode 100644 index b13cf303cd0709298f6a5c629c4a54cda33ffcdf..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/git/__pycache__/remote.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/git/__pycache__/types.cpython-310.pyc b/venv/lib/python3.10/site-packages/git/__pycache__/types.cpython-310.pyc deleted file mode 100644 index 5228c2102d3ee82e447c2b9caebc042d0662c929..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/git/__pycache__/types.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/git/__pycache__/util.cpython-310.pyc b/venv/lib/python3.10/site-packages/git/__pycache__/util.cpython-310.pyc deleted file mode 100644 index c2d4c0c222d0964f09a0a75644d94903dbd2dd8d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/git/__pycache__/util.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/git/cmd.py b/venv/lib/python3.10/site-packages/git/cmd.py deleted file mode 100644 index 2048a43fa9b3d7948dc680b1377c21e10f1e22e9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/git/cmd.py +++ /dev/null @@ -1,1724 +0,0 @@ -# Copyright (C) 2008, 2009 Michael Trier (mtrier@gmail.com) and contributors -# -# This module is part of GitPython and is released under the -# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ - -from __future__ import annotations - -__all__ = ["GitMeta", "Git"] - -import contextlib -import io -import itertools -import logging -import os -import re -import signal -import subprocess -from subprocess import DEVNULL, PIPE, Popen -import sys -from textwrap import dedent -import threading -import warnings - -from git.compat import defenc, force_bytes, safe_decode -from git.exc import ( - CommandError, - GitCommandError, - GitCommandNotFound, - UnsafeOptionError, - UnsafeProtocolError, -) -from git.util import ( - cygpath, - expand_path, - is_cygwin_git, - patch_env, - remove_password_if_present, - stream_copy, -) - -# typing --------------------------------------------------------------------------- - -from typing import ( - Any, - AnyStr, - BinaryIO, - Callable, - Dict, - IO, - Iterator, - List, - Mapping, - Optional, - Sequence, - TYPE_CHECKING, - TextIO, - Tuple, - Union, - cast, - overload, -) - -from git.types import Literal, PathLike, TBD - -if TYPE_CHECKING: - from git.diff import DiffIndex - from git.repo.base import Repo - -# --------------------------------------------------------------------------------- - -execute_kwargs = { - "istream", - "with_extended_output", - "with_exceptions", - "as_process", - "output_stream", - "stdout_as_string", - "kill_after_timeout", - "with_stdout", - "universal_newlines", - "shell", - "env", - "max_chunk_size", - "strip_newline_in_stdout", -} - -_logger = logging.getLogger(__name__) - - -# ============================================================================== -## @name Utilities -# ------------------------------------------------------------------------------ -# Documentation -## @{ - - -def handle_process_output( - process: "Git.AutoInterrupt" | Popen, - stdout_handler: Union[ - None, - Callable[[AnyStr], None], - Callable[[List[AnyStr]], None], - Callable[[bytes, "Repo", "DiffIndex"], None], - ], - stderr_handler: Union[None, Callable[[AnyStr], None], Callable[[List[AnyStr]], None]], - finalizer: Union[None, Callable[[Union[Popen, "Git.AutoInterrupt"]], None]] = None, - decode_streams: bool = True, - kill_after_timeout: Union[None, float] = None, -) -> None: - R"""Register for notifications to learn that process output is ready to read, and - dispatch lines to the respective line handlers. - - This function returns once the finalizer returns. - - :param process: - :class:`subprocess.Popen` instance. - - :param stdout_handler: - f(stdout_line_string), or ``None``. - - :param stderr_handler: - f(stderr_line_string), or ``None``. - - :param finalizer: - f(proc) - wait for proc to finish. - - :param decode_streams: - Assume stdout/stderr streams are binary and decode them before pushing their - contents to handlers. - - This defaults to ``True``. Set it to ``False`` if: - - - ``universal_newlines == True``, as then streams are in text mode, or - - decoding must happen later, such as for :class:`~git.diff.Diff`\s. - - :param kill_after_timeout: - :class:`float` or ``None``, Default = ``None`` - - To specify a timeout in seconds for the git command, after which the process - should be killed. - """ - - # Use 2 "pump" threads and wait for both to finish. - def pump_stream( - cmdline: List[str], - name: str, - stream: Union[BinaryIO, TextIO], - is_decode: bool, - handler: Union[None, Callable[[Union[bytes, str]], None]], - ) -> None: - try: - for line in stream: - if handler: - if is_decode: - assert isinstance(line, bytes) - line_str = line.decode(defenc) - handler(line_str) - else: - handler(line) - - except Exception as ex: - _logger.error(f"Pumping {name!r} of cmd({remove_password_if_present(cmdline)}) failed due to: {ex!r}") - if "I/O operation on closed file" not in str(ex): - # Only reraise if the error was not due to the stream closing. - raise CommandError([f"<{name}-pump>"] + remove_password_if_present(cmdline), ex) from ex - finally: - stream.close() - - if hasattr(process, "proc"): - process = cast("Git.AutoInterrupt", process) - cmdline: str | Tuple[str, ...] | List[str] = getattr(process.proc, "args", "") - p_stdout = process.proc.stdout if process.proc else None - p_stderr = process.proc.stderr if process.proc else None - else: - process = cast(Popen, process) # type: ignore[redundant-cast] - cmdline = getattr(process, "args", "") - p_stdout = process.stdout - p_stderr = process.stderr - - if not isinstance(cmdline, (tuple, list)): - cmdline = cmdline.split() - - pumps: List[Tuple[str, IO, Callable[..., None] | None]] = [] - if p_stdout: - pumps.append(("stdout", p_stdout, stdout_handler)) - if p_stderr: - pumps.append(("stderr", p_stderr, stderr_handler)) - - threads: List[threading.Thread] = [] - - for name, stream, handler in pumps: - t = threading.Thread(target=pump_stream, args=(cmdline, name, stream, decode_streams, handler)) - t.daemon = True - t.start() - threads.append(t) - - # FIXME: Why join? Will block if stdin needs feeding... - for t in threads: - t.join(timeout=kill_after_timeout) - if t.is_alive(): - if isinstance(process, Git.AutoInterrupt): - process._terminate() - else: # Don't want to deal with the other case. - raise RuntimeError( - "Thread join() timed out in cmd.handle_process_output()." - f" kill_after_timeout={kill_after_timeout} seconds" - ) - if stderr_handler: - error_str: Union[str, bytes] = ( - "error: process killed because it timed out." f" kill_after_timeout={kill_after_timeout} seconds" - ) - if not decode_streams and isinstance(p_stderr, BinaryIO): - # Assume stderr_handler needs binary input. - error_str = cast(str, error_str) - error_str = error_str.encode() - # We ignore typing on the next line because mypy does not like the way - # we inferred that stderr takes str or bytes. - stderr_handler(error_str) # type: ignore[arg-type] - - if finalizer: - finalizer(process) - - -safer_popen: Callable[..., Popen] - -if sys.platform == "win32": - - def _safer_popen_windows( - command: Union[str, Sequence[Any]], - *, - shell: bool = False, - env: Optional[Mapping[str, str]] = None, - **kwargs: Any, - ) -> Popen: - """Call :class:`subprocess.Popen` on Windows but don't include a CWD in the - search. - - This avoids an untrusted search path condition where a file like ``git.exe`` in - a malicious repository would be run when GitPython operates on the repository. - The process using GitPython may have an untrusted repository's working tree as - its current working directory. Some operations may temporarily change to that - directory before running a subprocess. In addition, while by default GitPython - does not run external commands with a shell, it can be made to do so, in which - case the CWD of the subprocess, which GitPython usually sets to a repository - working tree, can itself be searched automatically by the shell. This wrapper - covers all those cases. - - :note: - This currently works by setting the - :envvar:`NoDefaultCurrentDirectoryInExePath` environment variable during - subprocess creation. It also takes care of passing Windows-specific process - creation flags, but that is unrelated to path search. - - :note: - The current implementation contains a race condition on :attr:`os.environ`. - GitPython isn't thread-safe, but a program using it on one thread should - ideally be able to mutate :attr:`os.environ` on another, without - unpredictable results. See comments in: - https://github.com/gitpython-developers/GitPython/pull/1650 - """ - # CREATE_NEW_PROCESS_GROUP is needed for some ways of killing it afterwards. - # https://docs.python.org/3/library/subprocess.html#subprocess.Popen.send_signal - # https://docs.python.org/3/library/subprocess.html#subprocess.CREATE_NEW_PROCESS_GROUP - creationflags = subprocess.CREATE_NO_WINDOW | subprocess.CREATE_NEW_PROCESS_GROUP - - # When using a shell, the shell is the direct subprocess, so the variable must - # be set in its environment, to affect its search behavior. - if shell: - # The original may be immutable, or the caller may reuse it. Mutate a copy. - env = {} if env is None else dict(env) - env["NoDefaultCurrentDirectoryInExePath"] = "1" # The "1" can be an value. - - # When not using a shell, the current process does the search in a - # CreateProcessW API call, so the variable must be set in our environment. With - # a shell, that's unnecessary if https://github.com/python/cpython/issues/101283 - # is patched. In Python versions where it is unpatched, and in the rare case the - # ComSpec environment variable is unset, the search for the shell itself is - # unsafe. Setting NoDefaultCurrentDirectoryInExePath in all cases, as done here, - # is simpler and protects against that. (As above, the "1" can be any value.) - with patch_env("NoDefaultCurrentDirectoryInExePath", "1"): - return Popen( - command, - shell=shell, - env=env, - creationflags=creationflags, - **kwargs, - ) - - safer_popen = _safer_popen_windows -else: - safer_popen = Popen - - -def dashify(string: str) -> str: - return string.replace("_", "-") - - -def slots_to_dict(self: "Git", exclude: Sequence[str] = ()) -> Dict[str, Any]: - return {s: getattr(self, s) for s in self.__slots__ if s not in exclude} - - -def dict_to_slots_and__excluded_are_none(self: object, d: Mapping[str, Any], excluded: Sequence[str] = ()) -> None: - for k, v in d.items(): - setattr(self, k, v) - for k in excluded: - setattr(self, k, None) - - -## -- End Utilities -- @} - -_USE_SHELL_DEFAULT_MESSAGE = ( - "Git.USE_SHELL is deprecated, because only its default value of False is safe. " - "It will be removed in a future release." -) - -_USE_SHELL_DANGER_MESSAGE = ( - "Setting Git.USE_SHELL to True is unsafe and insecure, as the effect of special " - "shell syntax cannot usually be accounted for. This can result in a command " - "injection vulnerability and arbitrary code execution. Git.USE_SHELL is deprecated " - "and will be removed in a future release." -) - - -def _warn_use_shell(extra_danger: bool) -> None: - warnings.warn( - _USE_SHELL_DANGER_MESSAGE if extra_danger else _USE_SHELL_DEFAULT_MESSAGE, - DeprecationWarning, - stacklevel=3, - ) - - -class _GitMeta(type): - """Metaclass for :class:`Git`. - - This helps issue :class:`DeprecationWarning` if :attr:`Git.USE_SHELL` is used. - """ - - def __getattribute(cls, name: str) -> Any: - if name == "USE_SHELL": - _warn_use_shell(False) - return super().__getattribute__(name) - - def __setattr(cls, name: str, value: Any) -> Any: - if name == "USE_SHELL": - _warn_use_shell(value) - super().__setattr__(name, value) - - if not TYPE_CHECKING: - # To preserve static checking for undefined/misspelled attributes while letting - # the methods' bodies be type-checked, these are defined as non-special methods, - # then bound to special names out of view of static type checkers. (The original - # names invoke name mangling (leading "__") to avoid confusion in other scopes.) - __getattribute__ = __getattribute - __setattr__ = __setattr - - -GitMeta = _GitMeta -"""Alias of :class:`Git`'s metaclass, whether it is :class:`type` or a custom metaclass. - -Whether the :class:`Git` class has the default :class:`type` as its metaclass or uses a -custom metaclass is not documented and may change at any time. This statically checkable -metaclass alias is equivalent at runtime to ``type(Git)``. This should almost never be -used. Code that benefits from it is likely to be remain brittle even if it is used. - -In view of the :class:`Git` class's intended use and :class:`Git` objects' dynamic -callable attributes representing git subcommands, it rarely makes sense to inherit from -:class:`Git` at all. Using :class:`Git` in multiple inheritance can be especially tricky -to do correctly. Attempting uses of :class:`Git` where its metaclass is relevant, such -as when a sibling class has an unrelated metaclass and a shared lower bound metaclass -might have to be introduced to solve a metaclass conflict, is not recommended. - -:note: - The correct static type of the :class:`Git` class itself, and any subclasses, is - ``Type[Git]``. (This can be written as ``type[Git]`` in Python 3.9 later.) - - :class:`GitMeta` should never be used in any annotation where ``Type[Git]`` is - intended or otherwise possible to use. This alias is truly only for very rare and - inherently precarious situations where it is necessary to deal with the metaclass - explicitly. -""" - - -class Git(metaclass=_GitMeta): - """The Git class manages communication with the Git binary. - - It provides a convenient interface to calling the Git binary, such as in:: - - g = Git( git_dir ) - g.init() # calls 'git init' program - rval = g.ls_files() # calls 'git ls-files' program - - Debugging: - - * Set the :envvar:`GIT_PYTHON_TRACE` environment variable to print each invocation - of the command to stdout. - * Set its value to ``full`` to see details about the returned values. - """ - - __slots__ = ( - "_working_dir", - "cat_file_all", - "cat_file_header", - "_version_info", - "_version_info_token", - "_git_options", - "_persistent_git_options", - "_environment", - ) - - _excluded_ = ( - "cat_file_all", - "cat_file_header", - "_version_info", - "_version_info_token", - ) - - re_unsafe_protocol = re.compile(r"(.+)::.+") - - def __getstate__(self) -> Dict[str, Any]: - return slots_to_dict(self, exclude=self._excluded_) - - def __setstate__(self, d: Dict[str, Any]) -> None: - dict_to_slots_and__excluded_are_none(self, d, excluded=self._excluded_) - - # CONFIGURATION - - git_exec_name = "git" - """Default git command that should work on Linux, Windows, and other systems.""" - - GIT_PYTHON_TRACE = os.environ.get("GIT_PYTHON_TRACE", False) - """Enables debugging of GitPython's git commands.""" - - USE_SHELL: bool = False - """Deprecated. If set to ``True``, a shell will be used when executing git commands. - - Code that uses ``USE_SHELL = True`` or that passes ``shell=True`` to any GitPython - functions should be updated to use the default value of ``False`` instead. ``True`` - is unsafe unless the effect of syntax treated specially by the shell is fully - considered and accounted for, which is not possible under most circumstances. As - detailed below, it is also no longer needed, even where it had been in the past. - - It is in many if not most cases a command injection vulnerability for an application - to set :attr:`USE_SHELL` to ``True``. Any attacker who can cause a specially crafted - fragment of text to make its way into any part of any argument to any git command - (including paths, branch names, etc.) can cause the shell to read and write - arbitrary files and execute arbitrary commands. Innocent input may also accidentally - contain special shell syntax, leading to inadvertent malfunctions. - - In addition, how a value of ``True`` interacts with some aspects of GitPython's - operation is not precisely specified and may change without warning, even before - GitPython 4.0.0 when :attr:`USE_SHELL` may be removed. This includes: - - * Whether or how GitPython automatically customizes the shell environment. - - * Whether, outside of Windows (where :class:`subprocess.Popen` supports lists of - separate arguments even when ``shell=True``), this can be used with any GitPython - functionality other than direct calls to the :meth:`execute` method. - - * Whether any GitPython feature that runs git commands ever attempts to partially - sanitize data a shell may treat specially. Currently this is not done. - - Prior to GitPython 2.0.8, this had a narrow purpose in suppressing console windows - in graphical Windows applications. In 2.0.8 and higher, it provides no benefit, as - GitPython solves that problem more robustly and safely by using the - ``CREATE_NO_WINDOW`` process creation flag on Windows. - - Because Windows path search differs subtly based on whether a shell is used, in rare - cases changing this from ``True`` to ``False`` may keep an unusual git "executable", - such as a batch file, from being found. To fix this, set the command name or full - path in the :envvar:`GIT_PYTHON_GIT_EXECUTABLE` environment variable or pass the - full path to :func:`git.refresh` (or invoke the script using a ``.exe`` shim). - - Further reading: - - * :meth:`Git.execute` (on the ``shell`` parameter). - * https://github.com/gitpython-developers/GitPython/commit/0d9390866f9ce42870d3116094cd49e0019a970a - * https://learn.microsoft.com/en-us/windows/win32/procthread/process-creation-flags - * https://github.com/python/cpython/issues/91558#issuecomment-1100942950 - * https://learn.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-createprocessw - """ - - _git_exec_env_var = "GIT_PYTHON_GIT_EXECUTABLE" - _refresh_env_var = "GIT_PYTHON_REFRESH" - - GIT_PYTHON_GIT_EXECUTABLE = None - """Provide the full path to the git executable. Otherwise it assumes git is in the - executable search path. - - :note: - The git executable is actually found during the refresh step in the top level - ``__init__``. It can also be changed by explicitly calling :func:`git.refresh`. - """ - - _refresh_token = object() # Since None would match an initial _version_info_token. - - @classmethod - def refresh(cls, path: Union[None, PathLike] = None) -> bool: - """Update information about the git executable :class:`Git` objects will use. - - Called by the :func:`git.refresh` function in the top level ``__init__``. - - :param path: - Optional path to the git executable. If not absolute, it is resolved - immediately, relative to the current directory. (See note below.) - - :note: - The top-level :func:`git.refresh` should be preferred because it calls this - method and may also update other state accordingly. - - :note: - There are three different ways to specify the command that refreshing causes - to be used for git: - - 1. Pass no `path` argument and do not set the - :envvar:`GIT_PYTHON_GIT_EXECUTABLE` environment variable. The command - name ``git`` is used. It is looked up in a path search by the system, in - each command run (roughly similar to how git is found when running - ``git`` commands manually). This is usually the desired behavior. - - 2. Pass no `path` argument but set the :envvar:`GIT_PYTHON_GIT_EXECUTABLE` - environment variable. The command given as the value of that variable is - used. This may be a simple command or an arbitrary path. It is looked up - in each command run. Setting :envvar:`GIT_PYTHON_GIT_EXECUTABLE` to - ``git`` has the same effect as not setting it. - - 3. Pass a `path` argument. This path, if not absolute, is immediately - resolved, relative to the current directory. This resolution occurs at - the time of the refresh. When git commands are run, they are run using - that previously resolved path. If a `path` argument is passed, the - :envvar:`GIT_PYTHON_GIT_EXECUTABLE` environment variable is not - consulted. - - :note: - Refreshing always sets the :attr:`Git.GIT_PYTHON_GIT_EXECUTABLE` class - attribute, which can be read on the :class:`Git` class or any of its - instances to check what command is used to run git. This attribute should - not be confused with the related :envvar:`GIT_PYTHON_GIT_EXECUTABLE` - environment variable. The class attribute is set no matter how refreshing is - performed. - """ - # Discern which path to refresh with. - if path is not None: - new_git = os.path.expanduser(path) - new_git = os.path.abspath(new_git) - else: - new_git = os.environ.get(cls._git_exec_env_var, cls.git_exec_name) - - # Keep track of the old and new git executable path. - old_git = cls.GIT_PYTHON_GIT_EXECUTABLE - old_refresh_token = cls._refresh_token - cls.GIT_PYTHON_GIT_EXECUTABLE = new_git - cls._refresh_token = object() - - # Test if the new git executable path is valid. A GitCommandNotFound error is - # raised by us. A PermissionError is raised if the git executable cannot be - # executed for whatever reason. - has_git = False - try: - cls().version() - has_git = True - except (GitCommandNotFound, PermissionError): - pass - - # Warn or raise exception if test failed. - if not has_git: - err = ( - dedent( - """\ - Bad git executable. - The git executable must be specified in one of the following ways: - - be included in your $PATH - - be set via $%s - - explicitly set via git.refresh() - """ - ) - % cls._git_exec_env_var - ) - - # Revert to whatever the old_git was. - cls.GIT_PYTHON_GIT_EXECUTABLE = old_git - cls._refresh_token = old_refresh_token - - if old_git is None: - # On the first refresh (when GIT_PYTHON_GIT_EXECUTABLE is None) we only - # are quiet, warn, or error depending on the GIT_PYTHON_REFRESH value. - - # Determine what the user wants to happen during the initial refresh. We - # expect GIT_PYTHON_REFRESH to either be unset or be one of the - # following values: - # - # 0|q|quiet|s|silence|silent|n|none - # 1|w|warn|warning|l|log - # 2|r|raise|e|error|exception - - mode = os.environ.get(cls._refresh_env_var, "raise").lower() - - quiet = ["quiet", "q", "silence", "s", "silent", "none", "n", "0"] - warn = ["warn", "w", "warning", "log", "l", "1"] - error = ["error", "e", "exception", "raise", "r", "2"] - - if mode in quiet: - pass - elif mode in warn or mode in error: - err = dedent( - """\ - %s - All git commands will error until this is rectified. - - This initial message can be silenced or aggravated in the future by setting the - $%s environment variable. Use one of the following values: - - %s: for no message or exception - - %s: for a warning message (logging level CRITICAL, displayed by default) - - %s: for a raised exception - - Example: - export %s=%s - """ - ) % ( - err, - cls._refresh_env_var, - "|".join(quiet), - "|".join(warn), - "|".join(error), - cls._refresh_env_var, - quiet[0], - ) - - if mode in warn: - _logger.critical(err) - else: - raise ImportError(err) - else: - err = dedent( - """\ - %s environment variable has been set but it has been set with an invalid value. - - Use only the following values: - - %s: for no message or exception - - %s: for a warning message (logging level CRITICAL, displayed by default) - - %s: for a raised exception - """ - ) % ( - cls._refresh_env_var, - "|".join(quiet), - "|".join(warn), - "|".join(error), - ) - raise ImportError(err) - - # We get here if this was the initial refresh and the refresh mode was - # not error. Go ahead and set the GIT_PYTHON_GIT_EXECUTABLE such that we - # discern the difference between the first refresh at import time - # and subsequent calls to git.refresh or this refresh method. - cls.GIT_PYTHON_GIT_EXECUTABLE = cls.git_exec_name - else: - # After the first refresh (when GIT_PYTHON_GIT_EXECUTABLE is no longer - # None) we raise an exception. - raise GitCommandNotFound(new_git, err) - - return has_git - - @classmethod - def is_cygwin(cls) -> bool: - return is_cygwin_git(cls.GIT_PYTHON_GIT_EXECUTABLE) - - @overload - @classmethod - def polish_url(cls, url: str, is_cygwin: Literal[False] = ...) -> str: ... - - @overload - @classmethod - def polish_url(cls, url: str, is_cygwin: Union[None, bool] = None) -> str: ... - - @classmethod - def polish_url(cls, url: str, is_cygwin: Union[None, bool] = None) -> PathLike: - """Remove any backslashes from URLs to be written in config files. - - Windows might create config files containing paths with backslashes, but git - stops liking them as it will escape the backslashes. Hence we undo the escaping - just to be sure. - """ - if is_cygwin is None: - is_cygwin = cls.is_cygwin() - - if is_cygwin: - url = cygpath(url) - else: - url = os.path.expandvars(url) - if url.startswith("~"): - url = os.path.expanduser(url) - url = url.replace("\\\\", "\\").replace("\\", "/") - return url - - @classmethod - def check_unsafe_protocols(cls, url: str) -> None: - """Check for unsafe protocols. - - Apart from the usual protocols (http, git, ssh), Git allows "remote helpers" - that have the form ``::
``. One of these helpers (``ext::``) - can be used to invoke any arbitrary command. - - See: - - - https://git-scm.com/docs/gitremote-helpers - - https://git-scm.com/docs/git-remote-ext - """ - match = cls.re_unsafe_protocol.match(url) - if match: - protocol = match.group(1) - raise UnsafeProtocolError( - f"The `{protocol}::` protocol looks suspicious, use `allow_unsafe_protocols=True` to allow it." - ) - - @classmethod - def check_unsafe_options(cls, options: List[str], unsafe_options: List[str]) -> None: - """Check for unsafe options. - - Some options that are passed to ``git `` can be used to execute - arbitrary commands. These are blocked by default. - """ - # Options can be of the form `foo`, `--foo bar`, or `--foo=bar`, so we need to - # check if they start with "--foo" or if they are equal to "foo". - bare_unsafe_options = [option.lstrip("-") for option in unsafe_options] - for option in options: - for unsafe_option, bare_option in zip(unsafe_options, bare_unsafe_options): - if option.startswith(unsafe_option) or option == bare_option: - raise UnsafeOptionError( - f"{unsafe_option} is not allowed, use `allow_unsafe_options=True` to allow it." - ) - - class AutoInterrupt: - """Process wrapper that terminates the wrapped process on finalization. - - This kills/interrupts the stored process instance once this instance goes out of - scope. It is used to prevent processes piling up in case iterators stop reading. - - All attributes are wired through to the contained process object. - - The wait method is overridden to perform automatic status code checking and - possibly raise. - """ - - __slots__ = ("proc", "args", "status") - - # If this is non-zero it will override any status code during _terminate, used - # to prevent race conditions in testing. - _status_code_if_terminate: int = 0 - - def __init__(self, proc: Union[None, subprocess.Popen], args: Any) -> None: - self.proc = proc - self.args = args - self.status: Union[int, None] = None - - def _terminate(self) -> None: - """Terminate the underlying process.""" - if self.proc is None: - return - - proc = self.proc - self.proc = None - if proc.stdin: - proc.stdin.close() - if proc.stdout: - proc.stdout.close() - if proc.stderr: - proc.stderr.close() - # Did the process finish already so we have a return code? - try: - if proc.poll() is not None: - self.status = self._status_code_if_terminate or proc.poll() - return - except OSError as ex: - _logger.info("Ignored error after process had died: %r", ex) - - # It can be that nothing really exists anymore... - if os is None or getattr(os, "kill", None) is None: - return - - # Try to kill it. - try: - proc.terminate() - status = proc.wait() # Ensure the process goes away. - - self.status = self._status_code_if_terminate or status - except OSError as ex: - _logger.info("Ignored error after process had died: %r", ex) - # END exception handling - - def __del__(self) -> None: - self._terminate() - - def __getattr__(self, attr: str) -> Any: - return getattr(self.proc, attr) - - # TODO: Bad choice to mimic `proc.wait()` but with different args. - def wait(self, stderr: Union[None, str, bytes] = b"") -> int: - """Wait for the process and return its status code. - - :param stderr: - Previously read value of stderr, in case stderr is already closed. - - :warn: - May deadlock if output or error pipes are used and not handled - separately. - - :raise git.exc.GitCommandError: - If the return status is not 0. - """ - if stderr is None: - stderr_b = b"" - stderr_b = force_bytes(data=stderr, encoding="utf-8") - status: Union[int, None] - if self.proc is not None: - status = self.proc.wait() - p_stderr = self.proc.stderr - else: # Assume the underlying proc was killed earlier or never existed. - status = self.status - p_stderr = None - - def read_all_from_possibly_closed_stream(stream: Union[IO[bytes], None]) -> bytes: - if stream: - try: - return stderr_b + force_bytes(stream.read()) - except (OSError, ValueError): - return stderr_b or b"" - else: - return stderr_b or b"" - - # END status handling - - if status != 0: - errstr = read_all_from_possibly_closed_stream(p_stderr) - _logger.debug("AutoInterrupt wait stderr: %r" % (errstr,)) - raise GitCommandError(remove_password_if_present(self.args), status, errstr) - return status - - # END auto interrupt - - class CatFileContentStream: - """Object representing a sized read-only stream returning the contents of - an object. - - This behaves like a stream, but counts the data read and simulates an empty - stream once our sized content region is empty. - - If not all data are read to the end of the object's lifetime, we read the - rest to ensure the underlying stream continues to work. - """ - - __slots__ = ("_stream", "_nbr", "_size") - - def __init__(self, size: int, stream: IO[bytes]) -> None: - self._stream = stream - self._size = size - self._nbr = 0 # Number of bytes read. - - # Special case: If the object is empty, has null bytes, get the final - # newline right away. - if size == 0: - stream.read(1) - # END handle empty streams - - def read(self, size: int = -1) -> bytes: - bytes_left = self._size - self._nbr - if bytes_left == 0: - return b"" - if size > -1: - # Ensure we don't try to read past our limit. - size = min(bytes_left, size) - else: - # They try to read all, make sure it's not more than what remains. - size = bytes_left - # END check early depletion - data = self._stream.read(size) - self._nbr += len(data) - - # Check for depletion, read our final byte to make the stream usable by - # others. - if self._size - self._nbr == 0: - self._stream.read(1) # final newline - # END finish reading - return data - - def readline(self, size: int = -1) -> bytes: - if self._nbr == self._size: - return b"" - - # Clamp size to lowest allowed value. - bytes_left = self._size - self._nbr - if size > -1: - size = min(bytes_left, size) - else: - size = bytes_left - # END handle size - - data = self._stream.readline(size) - self._nbr += len(data) - - # Handle final byte. - if self._size - self._nbr == 0: - self._stream.read(1) - # END finish reading - - return data - - def readlines(self, size: int = -1) -> List[bytes]: - if self._nbr == self._size: - return [] - - # Leave all additional logic to our readline method, we just check the size. - out = [] - nbr = 0 - while True: - line = self.readline() - if not line: - break - out.append(line) - if size > -1: - nbr += len(line) - if nbr > size: - break - # END handle size constraint - # END readline loop - return out - - # skipcq: PYL-E0301 - def __iter__(self) -> "Git.CatFileContentStream": - return self - - def __next__(self) -> bytes: - line = self.readline() - if not line: - raise StopIteration - - return line - - next = __next__ - - def __del__(self) -> None: - bytes_left = self._size - self._nbr - if bytes_left: - # Read and discard - seeking is impossible within a stream. - # This includes any terminating newline. - self._stream.read(bytes_left + 1) - # END handle incomplete read - - def __init__(self, working_dir: Union[None, PathLike] = None) -> None: - """Initialize this instance with: - - :param working_dir: - Git directory we should work in. If ``None``, we always work in the current - directory as returned by :func:`os.getcwd`. - This is meant to be the working tree directory if available, or the - ``.git`` directory in case of bare repositories. - """ - super().__init__() - self._working_dir = expand_path(working_dir) - self._git_options: Union[List[str], Tuple[str, ...]] = () - self._persistent_git_options: List[str] = [] - - # Extra environment variables to pass to git commands - self._environment: Dict[str, str] = {} - - # Cached version slots - self._version_info: Union[Tuple[int, ...], None] = None - self._version_info_token: object = None - - # Cached command slots - self.cat_file_header: Union[None, TBD] = None - self.cat_file_all: Union[None, TBD] = None - - def __getattribute__(self, name: str) -> Any: - if name == "USE_SHELL": - _warn_use_shell(False) - return super().__getattribute__(name) - - def __getattr__(self, name: str) -> Any: - """A convenience method as it allows to call the command as if it was an object. - - :return: - Callable object that will execute call :meth:`_call_process` with your - arguments. - """ - if name.startswith("_"): - return super().__getattribute__(name) - return lambda *args, **kwargs: self._call_process(name, *args, **kwargs) - - def set_persistent_git_options(self, **kwargs: Any) -> None: - """Specify command line options to the git executable for subsequent - subcommand calls. - - :param kwargs: - A dict of keyword arguments. - These arguments are passed as in :meth:`_call_process`, but will be passed - to the git command rather than the subcommand. - """ - - self._persistent_git_options = self.transform_kwargs(split_single_char_options=True, **kwargs) - - @property - def working_dir(self) -> Union[None, PathLike]: - """:return: Git directory we are working on""" - return self._working_dir - - @property - def version_info(self) -> Tuple[int, ...]: - """ - :return: Tuple with integers representing the major, minor and additional - version numbers as parsed from :manpage:`git-version(1)`. Up to four fields - are used. - - This value is generated on demand and is cached. - """ - # Refreshing is global, but version_info caching is per-instance. - refresh_token = self._refresh_token # Copy token in case of concurrent refresh. - - # Use the cached version if obtained after the most recent refresh. - if self._version_info_token is refresh_token: - assert self._version_info is not None, "Bug: corrupted token-check state" - return self._version_info - - # Run "git version" and parse it. - process_version = self._call_process("version") - version_string = process_version.split(" ")[2] - version_fields = version_string.split(".")[:4] - leading_numeric_fields = itertools.takewhile(str.isdigit, version_fields) - self._version_info = tuple(map(int, leading_numeric_fields)) - - # This value will be considered valid until the next refresh. - self._version_info_token = refresh_token - return self._version_info - - @overload - def execute( - self, - command: Union[str, Sequence[Any]], - *, - as_process: Literal[True], - ) -> "AutoInterrupt": ... - - @overload - def execute( - self, - command: Union[str, Sequence[Any]], - *, - as_process: Literal[False] = False, - stdout_as_string: Literal[True], - ) -> Union[str, Tuple[int, str, str]]: ... - - @overload - def execute( - self, - command: Union[str, Sequence[Any]], - *, - as_process: Literal[False] = False, - stdout_as_string: Literal[False] = False, - ) -> Union[bytes, Tuple[int, bytes, str]]: ... - - @overload - def execute( - self, - command: Union[str, Sequence[Any]], - *, - with_extended_output: Literal[False], - as_process: Literal[False], - stdout_as_string: Literal[True], - ) -> str: ... - - @overload - def execute( - self, - command: Union[str, Sequence[Any]], - *, - with_extended_output: Literal[False], - as_process: Literal[False], - stdout_as_string: Literal[False], - ) -> bytes: ... - - def execute( - self, - command: Union[str, Sequence[Any]], - istream: Union[None, BinaryIO] = None, - with_extended_output: bool = False, - with_exceptions: bool = True, - as_process: bool = False, - output_stream: Union[None, BinaryIO] = None, - stdout_as_string: bool = True, - kill_after_timeout: Union[None, float] = None, - with_stdout: bool = True, - universal_newlines: bool = False, - shell: Union[None, bool] = None, - env: Union[None, Mapping[str, str]] = None, - max_chunk_size: int = io.DEFAULT_BUFFER_SIZE, - strip_newline_in_stdout: bool = True, - **subprocess_kwargs: Any, - ) -> Union[str, bytes, Tuple[int, Union[str, bytes], str], AutoInterrupt]: - R"""Handle executing the command, and consume and return the returned - information (stdout). - - :param command: - The command argument list to execute. - It should be a sequence of program arguments, or a string. The - program to execute is the first item in the args sequence or string. - - :param istream: - Standard input filehandle passed to :class:`subprocess.Popen`. - - :param with_extended_output: - Whether to return a (status, stdout, stderr) tuple. - - :param with_exceptions: - Whether to raise an exception when git returns a non-zero status. - - :param as_process: - Whether to return the created process instance directly from which - streams can be read on demand. This will render `with_extended_output` - and `with_exceptions` ineffective - the caller will have to deal with - the details. It is important to note that the process will be placed - into an :class:`AutoInterrupt` wrapper that will interrupt the process - once it goes out of scope. If you use the command in iterators, you - should pass the whole process instance instead of a single stream. - - :param output_stream: - If set to a file-like object, data produced by the git command will be - copied to the given stream instead of being returned as a string. - This feature only has any effect if `as_process` is ``False``. - - :param stdout_as_string: - If ``False``, the command's standard output will be bytes. Otherwise, it - will be decoded into a string using the default encoding (usually UTF-8). - The latter can fail, if the output contains binary data. - - :param kill_after_timeout: - Specifies a timeout in seconds for the git command, after which the process - should be killed. This will have no effect if `as_process` is set to - ``True``. It is set to ``None`` by default and will let the process run - until the timeout is explicitly specified. Uses of this feature should be - carefully considered, due to the following limitations: - - 1. This feature is not supported at all on Windows. - 2. Effectiveness may vary by operating system. ``ps --ppid`` is used to - enumerate child processes, which is available on most GNU/Linux systems - but not most others. - 3. Deeper descendants do not receive signals, though they may sometimes - terminate as a consequence of their parent processes being killed. - 4. `kill_after_timeout` uses ``SIGKILL``, which can have negative side - effects on a repository. For example, stale locks in case of - :manpage:`git-gc(1)` could render the repository incapable of accepting - changes until the lock is manually removed. - - :param with_stdout: - If ``True``, default ``True``, we open stdout on the created process. - - :param universal_newlines: - If ``True``, pipes will be opened as text, and lines are split at all known - line endings. - - :param shell: - Whether to invoke commands through a shell - (see :class:`Popen(..., shell=True) `). - If this is not ``None``, it overrides :attr:`USE_SHELL`. - - Passing ``shell=True`` to this or any other GitPython function should be - avoided, as it is unsafe under most circumstances. This is because it is - typically not feasible to fully consider and account for the effect of shell - expansions, especially when passing ``shell=True`` to other methods that - forward it to :meth:`Git.execute`. Passing ``shell=True`` is also no longer - needed (nor useful) to work around any known operating system specific - issues. - - :param env: - A dictionary of environment variables to be passed to - :class:`subprocess.Popen`. - - :param max_chunk_size: - Maximum number of bytes in one chunk of data passed to the `output_stream` - in one invocation of its ``write()`` method. If the given number is not - positive then the default value is used. - - :param strip_newline_in_stdout: - Whether to strip the trailing ``\n`` of the command stdout. - - :param subprocess_kwargs: - Keyword arguments to be passed to :class:`subprocess.Popen`. Please note - that some of the valid kwargs are already set by this method; the ones you - specify may not be the same ones. - - :return: - * str(output), if `extended_output` is ``False`` (Default) - * tuple(int(status), str(stdout), str(stderr)), - if `extended_output` is ``True`` - - If `output_stream` is ``True``, the stdout value will be your output stream: - - * output_stream, if `extended_output` is ``False`` - * tuple(int(status), output_stream, str(stderr)), - if `extended_output` is ``True`` - - Note that git is executed with ``LC_MESSAGES="C"`` to ensure consistent - output regardless of system language. - - :raise git.exc.GitCommandError: - - :note: - If you add additional keyword arguments to the signature of this method, you - must update the ``execute_kwargs`` variable housed in this module. - """ - # Remove password for the command if present. - redacted_command = remove_password_if_present(command) - if self.GIT_PYTHON_TRACE and (self.GIT_PYTHON_TRACE != "full" or as_process): - _logger.info(" ".join(redacted_command)) - - # Allow the user to have the command executed in their working dir. - try: - cwd = self._working_dir or os.getcwd() # type: Union[None, str] - if not os.access(str(cwd), os.X_OK): - cwd = None - except FileNotFoundError: - cwd = None - - # Start the process. - inline_env = env - env = os.environ.copy() - # Attempt to force all output to plain ASCII English, which is what some parsing - # code may expect. - # According to https://askubuntu.com/a/311796, we are setting LANGUAGE as well - # just to be sure. - env["LANGUAGE"] = "C" - env["LC_ALL"] = "C" - env.update(self._environment) - if inline_env is not None: - env.update(inline_env) - - if sys.platform == "win32": - if kill_after_timeout is not None: - raise GitCommandError( - redacted_command, - '"kill_after_timeout" feature is not supported on Windows.', - ) - cmd_not_found_exception = OSError - else: - cmd_not_found_exception = FileNotFoundError - # END handle - - stdout_sink = PIPE if with_stdout else getattr(subprocess, "DEVNULL", None) or open(os.devnull, "wb") - if shell is None: - # Get the value of USE_SHELL with no deprecation warning. Do this without - # warnings.catch_warnings, to avoid a race condition with application code - # configuring warnings. The value could be looked up in type(self).__dict__ - # or Git.__dict__, but those can break under some circumstances. This works - # the same as self.USE_SHELL in more situations; see Git.__getattribute__. - shell = super().__getattribute__("USE_SHELL") - _logger.debug( - "Popen(%s, cwd=%s, stdin=%s, shell=%s, universal_newlines=%s)", - redacted_command, - cwd, - "" if istream else "None", - shell, - universal_newlines, - ) - try: - proc = safer_popen( - command, - env=env, - cwd=cwd, - bufsize=-1, - stdin=(istream or DEVNULL), - stderr=PIPE, - stdout=stdout_sink, - shell=shell, - universal_newlines=universal_newlines, - encoding=defenc if universal_newlines else None, - **subprocess_kwargs, - ) - except cmd_not_found_exception as err: - raise GitCommandNotFound(redacted_command, err) from err - else: - # Replace with a typeguard for Popen[bytes]? - proc.stdout = cast(BinaryIO, proc.stdout) - proc.stderr = cast(BinaryIO, proc.stderr) - - if as_process: - return self.AutoInterrupt(proc, command) - - if sys.platform != "win32" and kill_after_timeout is not None: - # Help mypy figure out this is not None even when used inside communicate(). - timeout = kill_after_timeout - - def kill_process(pid: int) -> None: - """Callback to kill a process. - - This callback implementation would be ineffective and unsafe on Windows. - """ - p = Popen(["ps", "--ppid", str(pid)], stdout=PIPE) - child_pids = [] - if p.stdout is not None: - for line in p.stdout: - if len(line.split()) > 0: - local_pid = (line.split())[0] - if local_pid.isdigit(): - child_pids.append(int(local_pid)) - try: - os.kill(pid, signal.SIGKILL) - for child_pid in child_pids: - try: - os.kill(child_pid, signal.SIGKILL) - except OSError: - pass - # Tell the main routine that the process was killed. - kill_check.set() - except OSError: - # It is possible that the process gets completed in the duration - # after timeout happens and before we try to kill the process. - pass - return - - def communicate() -> Tuple[AnyStr, AnyStr]: - watchdog.start() - out, err = proc.communicate() - watchdog.cancel() - if kill_check.is_set(): - err = 'Timeout: the command "%s" did not complete in %d ' "secs." % ( - " ".join(redacted_command), - timeout, - ) - if not universal_newlines: - err = err.encode(defenc) - return out, err - - # END helpers - - kill_check = threading.Event() - watchdog = threading.Timer(timeout, kill_process, args=(proc.pid,)) - else: - communicate = proc.communicate - - # Wait for the process to return. - status = 0 - stdout_value: Union[str, bytes] = b"" - stderr_value: Union[str, bytes] = b"" - newline = "\n" if universal_newlines else b"\n" - try: - if output_stream is None: - stdout_value, stderr_value = communicate() - # Strip trailing "\n". - if stdout_value.endswith(newline) and strip_newline_in_stdout: # type: ignore[arg-type] - stdout_value = stdout_value[:-1] - if stderr_value.endswith(newline): # type: ignore[arg-type] - stderr_value = stderr_value[:-1] - - status = proc.returncode - else: - max_chunk_size = max_chunk_size if max_chunk_size and max_chunk_size > 0 else io.DEFAULT_BUFFER_SIZE - stream_copy(proc.stdout, output_stream, max_chunk_size) - stdout_value = proc.stdout.read() - stderr_value = proc.stderr.read() - # Strip trailing "\n". - if stderr_value.endswith(newline): # type: ignore[arg-type] - stderr_value = stderr_value[:-1] - status = proc.wait() - # END stdout handling - finally: - proc.stdout.close() - proc.stderr.close() - - if self.GIT_PYTHON_TRACE == "full": - cmdstr = " ".join(redacted_command) - - def as_text(stdout_value: Union[bytes, str]) -> str: - return not output_stream and safe_decode(stdout_value) or "" - - # END as_text - - if stderr_value: - _logger.info( - "%s -> %d; stdout: '%s'; stderr: '%s'", - cmdstr, - status, - as_text(stdout_value), - safe_decode(stderr_value), - ) - elif stdout_value: - _logger.info("%s -> %d; stdout: '%s'", cmdstr, status, as_text(stdout_value)) - else: - _logger.info("%s -> %d", cmdstr, status) - # END handle debug printing - - if with_exceptions and status != 0: - raise GitCommandError(redacted_command, status, stderr_value, stdout_value) - - if isinstance(stdout_value, bytes) and stdout_as_string: # Could also be output_stream. - stdout_value = safe_decode(stdout_value) - - # Allow access to the command's status code. - if with_extended_output: - return (status, stdout_value, safe_decode(stderr_value)) - else: - return stdout_value - - def environment(self) -> Dict[str, str]: - return self._environment - - def update_environment(self, **kwargs: Any) -> Dict[str, Union[str, None]]: - """Set environment variables for future git invocations. Return all changed - values in a format that can be passed back into this function to revert the - changes. - - Examples:: - - old_env = self.update_environment(PWD='/tmp') - self.update_environment(**old_env) - - :param kwargs: - Environment variables to use for git processes. - - :return: - Dict that maps environment variables to their old values - """ - old_env = {} - for key, value in kwargs.items(): - # Set value if it is None. - if value is not None: - old_env[key] = self._environment.get(key) - self._environment[key] = value - # Remove key from environment if its value is None. - elif key in self._environment: - old_env[key] = self._environment[key] - del self._environment[key] - return old_env - - @contextlib.contextmanager - def custom_environment(self, **kwargs: Any) -> Iterator[None]: - """A context manager around the above :meth:`update_environment` method to - restore the environment back to its previous state after operation. - - Examples:: - - with self.custom_environment(GIT_SSH='/bin/ssh_wrapper'): - repo.remotes.origin.fetch() - - :param kwargs: - See :meth:`update_environment`. - """ - old_env = self.update_environment(**kwargs) - try: - yield - finally: - self.update_environment(**old_env) - - def transform_kwarg(self, name: str, value: Any, split_single_char_options: bool) -> List[str]: - if len(name) == 1: - if value is True: - return ["-%s" % name] - elif value not in (False, None): - if split_single_char_options: - return ["-%s" % name, "%s" % value] - else: - return ["-%s%s" % (name, value)] - else: - if value is True: - return ["--%s" % dashify(name)] - elif value is not False and value is not None: - return ["--%s=%s" % (dashify(name), value)] - return [] - - def transform_kwargs(self, split_single_char_options: bool = True, **kwargs: Any) -> List[str]: - """Transform Python-style kwargs into git command line options.""" - args = [] - for k, v in kwargs.items(): - if isinstance(v, (list, tuple)): - for value in v: - args += self.transform_kwarg(k, value, split_single_char_options) - else: - args += self.transform_kwarg(k, v, split_single_char_options) - return args - - @classmethod - def _unpack_args(cls, arg_list: Sequence[str]) -> List[str]: - outlist = [] - if isinstance(arg_list, (list, tuple)): - for arg in arg_list: - outlist.extend(cls._unpack_args(arg)) - else: - outlist.append(str(arg_list)) - - return outlist - - def __call__(self, **kwargs: Any) -> "Git": - """Specify command line options to the git executable for a subcommand call. - - :param kwargs: - A dict of keyword arguments. - These arguments are passed as in :meth:`_call_process`, but will be passed - to the git command rather than the subcommand. - - Examples:: - - git(work_tree='/tmp').difftool() - """ - self._git_options = self.transform_kwargs(split_single_char_options=True, **kwargs) - return self - - @overload - def _call_process( - self, method: str, *args: None, **kwargs: None - ) -> str: ... # If no args were given, execute the call with all defaults. - - @overload - def _call_process( - self, - method: str, - istream: int, - as_process: Literal[True], - *args: Any, - **kwargs: Any, - ) -> "Git.AutoInterrupt": ... - - @overload - def _call_process( - self, method: str, *args: Any, **kwargs: Any - ) -> Union[str, bytes, Tuple[int, Union[str, bytes], str], "Git.AutoInterrupt"]: ... - - def _call_process( - self, method: str, *args: Any, **kwargs: Any - ) -> Union[str, bytes, Tuple[int, Union[str, bytes], str], "Git.AutoInterrupt"]: - """Run the given git command with the specified arguments and return the result - as a string. - - :param method: - The command. Contained ``_`` characters will be converted to hyphens, such - as in ``ls_files`` to call ``ls-files``. - - :param args: - The list of arguments. If ``None`` is included, it will be pruned. - This allows your commands to call git more conveniently, as ``None`` is - realized as non-existent. - - :param kwargs: - Contains key-values for the following: - - - The :meth:`execute()` kwds, as listed in ``execute_kwargs``. - - "Command options" to be converted by :meth:`transform_kwargs`. - - The ``insert_kwargs_after`` key which its value must match one of - ``*args``. - - It also contains any command options, to be appended after the matched arg. - - Examples:: - - git.rev_list('master', max_count=10, header=True) - - turns into:: - - git rev-list max-count 10 --header master - - :return: - Same as :meth:`execute`. If no args are given, used :meth:`execute`'s - default (especially ``as_process = False``, ``stdout_as_string = True``) and - return :class:`str`. - """ - # Handle optional arguments prior to calling transform_kwargs. - # Otherwise these'll end up in args, which is bad. - exec_kwargs = {k: v for k, v in kwargs.items() if k in execute_kwargs} - opts_kwargs = {k: v for k, v in kwargs.items() if k not in execute_kwargs} - - insert_after_this_arg = opts_kwargs.pop("insert_kwargs_after", None) - - # Prepare the argument list. - - opt_args = self.transform_kwargs(**opts_kwargs) - ext_args = self._unpack_args([a for a in args if a is not None]) - - if insert_after_this_arg is None: - args_list = opt_args + ext_args - else: - try: - index = ext_args.index(insert_after_this_arg) - except ValueError as err: - raise ValueError( - "Couldn't find argument '%s' in args %s to insert cmd options after" - % (insert_after_this_arg, str(ext_args)) - ) from err - # END handle error - args_list = ext_args[: index + 1] + opt_args + ext_args[index + 1 :] - # END handle opts_kwargs - - call = [self.GIT_PYTHON_GIT_EXECUTABLE] - - # Add persistent git options. - call.extend(self._persistent_git_options) - - # Add the git options, then reset to empty to avoid side effects. - call.extend(self._git_options) - self._git_options = () - - call.append(dashify(method)) - call.extend(args_list) - - return self.execute(call, **exec_kwargs) - - def _parse_object_header(self, header_line: str) -> Tuple[str, str, int]: - """ - :param header_line: - A line of the form:: - - type_string size_as_int - - :return: - (hex_sha, type_string, size_as_int) - - :raise ValueError: - If the header contains indication for an error due to incorrect input sha. - """ - tokens = header_line.split() - if len(tokens) != 3: - if not tokens: - err_msg = ( - f"SHA is empty, possible dubious ownership in the repository " - f"""at {self._working_dir}.\n If this is unintended run:\n\n """ - f""" "git config --global --add safe.directory {self._working_dir}" """ - ) - raise ValueError(err_msg) - else: - raise ValueError("SHA %s could not be resolved, git returned: %r" % (tokens[0], header_line.strip())) - # END handle actual return value - # END error handling - - if len(tokens[0]) != 40: - raise ValueError("Failed to parse header: %r" % header_line) - return (tokens[0], tokens[1], int(tokens[2])) - - def _prepare_ref(self, ref: AnyStr) -> bytes: - # Required for command to separate refs on stdin, as bytes. - if isinstance(ref, bytes): - # Assume 40 bytes hexsha - bin-to-ascii for some reason returns bytes, not text. - refstr: str = ref.decode("ascii") - elif not isinstance(ref, str): - refstr = str(ref) # Could be ref-object. - else: - refstr = ref - - if not refstr.endswith("\n"): - refstr += "\n" - return refstr.encode(defenc) - - def _get_persistent_cmd(self, attr_name: str, cmd_name: str, *args: Any, **kwargs: Any) -> "Git.AutoInterrupt": - cur_val = getattr(self, attr_name) - if cur_val is not None: - return cur_val - - options = {"istream": PIPE, "as_process": True} - options.update(kwargs) - - cmd = self._call_process(cmd_name, *args, **options) - setattr(self, attr_name, cmd) - cmd = cast("Git.AutoInterrupt", cmd) - return cmd - - def __get_object_header(self, cmd: "Git.AutoInterrupt", ref: AnyStr) -> Tuple[str, str, int]: - if cmd.stdin and cmd.stdout: - cmd.stdin.write(self._prepare_ref(ref)) - cmd.stdin.flush() - return self._parse_object_header(cmd.stdout.readline()) - else: - raise ValueError("cmd stdin was empty") - - def get_object_header(self, ref: str) -> Tuple[str, str, int]: - """Use this method to quickly examine the type and size of the object behind the - given ref. - - :note: - The method will only suffer from the costs of command invocation once and - reuses the command in subsequent calls. - - :return: - (hexsha, type_string, size_as_int) - """ - cmd = self._get_persistent_cmd("cat_file_header", "cat_file", batch_check=True) - return self.__get_object_header(cmd, ref) - - def get_object_data(self, ref: str) -> Tuple[str, str, int, bytes]: - """Similar to :meth:`get_object_header`, but returns object data as well. - - :return: - (hexsha, type_string, size_as_int, data_string) - - :note: - Not threadsafe. - """ - hexsha, typename, size, stream = self.stream_object_data(ref) - data = stream.read(size) - del stream - return (hexsha, typename, size, data) - - def stream_object_data(self, ref: str) -> Tuple[str, str, int, "Git.CatFileContentStream"]: - """Similar to :meth:`get_object_data`, but returns the data as a stream. - - :return: - (hexsha, type_string, size_as_int, stream) - - :note: - This method is not threadsafe. You need one independent :class:`Git` - instance per thread to be safe! - """ - cmd = self._get_persistent_cmd("cat_file_all", "cat_file", batch=True) - hexsha, typename, size = self.__get_object_header(cmd, ref) - cmd_stdout = cmd.stdout if cmd.stdout is not None else io.BytesIO() - return (hexsha, typename, size, self.CatFileContentStream(size, cmd_stdout)) - - def clear_cache(self) -> "Git": - """Clear all kinds of internal caches to release resources. - - Currently persistent commands will be interrupted. - - :return: - self - """ - for cmd in (self.cat_file_all, self.cat_file_header): - if cmd: - cmd.__del__() - - self.cat_file_all = None - self.cat_file_header = None - return self diff --git a/venv/lib/python3.10/site-packages/git/compat.py b/venv/lib/python3.10/site-packages/git/compat.py deleted file mode 100644 index d7d9a55a9f3cd42c587bcf7ae3de63004f927dd5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/git/compat.py +++ /dev/null @@ -1,165 +0,0 @@ -# Copyright (C) 2008, 2009 Michael Trier (mtrier@gmail.com) and contributors -# -# This module is part of GitPython and is released under the -# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ - -"""Utilities to help provide compatibility with Python 3. - -This module exists for historical reasons. Code outside GitPython may make use of public -members of this module, but is unlikely to benefit from doing so. GitPython continues to -use some of these utilities, in some cases for compatibility across different platforms. -""" - -import locale -import os -import sys -import warnings - -from gitdb.utils.encoding import force_bytes, force_text # noqa: F401 - -# typing -------------------------------------------------------------------- - -from typing import ( - Any, # noqa: F401 - AnyStr, - Dict, # noqa: F401 - IO, # noqa: F401 - List, - Optional, - TYPE_CHECKING, - Tuple, # noqa: F401 - Type, # noqa: F401 - Union, - overload, -) - -# --------------------------------------------------------------------------- - - -_deprecated_platform_aliases = { - "is_win": os.name == "nt", - "is_posix": os.name == "posix", - "is_darwin": sys.platform == "darwin", -} - - -def _getattr(name: str) -> Any: - try: - value = _deprecated_platform_aliases[name] - except KeyError: - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None - - warnings.warn( - f"{__name__}.{name} and other is_ aliases are deprecated. " - "Write the desired os.name or sys.platform check explicitly instead.", - DeprecationWarning, - stacklevel=2, - ) - return value - - -if not TYPE_CHECKING: # Preserve static checking for undefined/misspelled attributes. - __getattr__ = _getattr - - -def __dir__() -> List[str]: - return [*globals(), *_deprecated_platform_aliases] - - -is_win: bool -"""Deprecated alias for ``os.name == "nt"`` to check for native Windows. - -This is deprecated because it is clearer to write out :attr:`os.name` or -:attr:`sys.platform` checks explicitly, especially in cases where it matters which is -used. - -:note: - ``is_win`` is ``False`` on Cygwin, but is often wrongly assumed ``True``. To detect - Cygwin, use ``sys.platform == "cygwin"``. -""" - -is_posix: bool -"""Deprecated alias for ``os.name == "posix"`` to check for Unix-like ("POSIX") systems. - -This is deprecated because it clearer to write out :attr:`os.name` or -:attr:`sys.platform` checks explicitly, especially in cases where it matters which is -used. - -:note: - For POSIX systems, more detailed information is available in :attr:`sys.platform`, - while :attr:`os.name` is always ``"posix"`` on such systems, including macOS - (Darwin). -""" - -is_darwin: bool -"""Deprecated alias for ``sys.platform == "darwin"`` to check for macOS (Darwin). - -This is deprecated because it clearer to write out :attr:`os.name` or -:attr:`sys.platform` checks explicitly. - -:note: - For macOS (Darwin), ``os.name == "posix"`` as in other Unix-like systems, while - ``sys.platform == "darwin"``. -""" - -defenc = sys.getfilesystemencoding() -"""The encoding used to convert between Unicode and bytes filenames.""" - - -@overload -def safe_decode(s: None) -> None: ... - - -@overload -def safe_decode(s: AnyStr) -> str: ... - - -def safe_decode(s: Union[AnyStr, None]) -> Optional[str]: - """Safely decode a binary string to Unicode.""" - if isinstance(s, str): - return s - elif isinstance(s, bytes): - return s.decode(defenc, "surrogateescape") - elif s is None: - return None - else: - raise TypeError("Expected bytes or text, but got %r" % (s,)) - - -@overload -def safe_encode(s: None) -> None: ... - - -@overload -def safe_encode(s: AnyStr) -> bytes: ... - - -def safe_encode(s: Optional[AnyStr]) -> Optional[bytes]: - """Safely encode a binary string to Unicode.""" - if isinstance(s, str): - return s.encode(defenc) - elif isinstance(s, bytes): - return s - elif s is None: - return None - else: - raise TypeError("Expected bytes or text, but got %r" % (s,)) - - -@overload -def win_encode(s: None) -> None: ... - - -@overload -def win_encode(s: AnyStr) -> bytes: ... - - -def win_encode(s: Optional[AnyStr]) -> Optional[bytes]: - """Encode Unicode strings for process arguments on Windows.""" - if isinstance(s, str): - return s.encode(locale.getpreferredencoding(False)) - elif isinstance(s, bytes): - return s - elif s is not None: - raise TypeError("Expected bytes or text, but got %r" % (s,)) - return None diff --git a/venv/lib/python3.10/site-packages/git/config.py b/venv/lib/python3.10/site-packages/git/config.py deleted file mode 100644 index de35083606da4ce602c83bf562669969c55f8417..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/git/config.py +++ /dev/null @@ -1,944 +0,0 @@ -# Copyright (C) 2008, 2009 Michael Trier (mtrier@gmail.com) and contributors -# -# This module is part of GitPython and is released under the -# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ - -"""Parser for reading and writing configuration files.""" - -__all__ = ["GitConfigParser", "SectionConstraint"] - -import abc -import configparser as cp -import fnmatch -from functools import wraps -import inspect -from io import BufferedReader, IOBase -import logging -import os -import os.path as osp -import re -import sys - -from git.compat import defenc, force_text -from git.util import LockFile - -# typing------------------------------------------------------- - -from typing import ( - Any, - Callable, - Generic, - IO, - List, - Dict, - Sequence, - TYPE_CHECKING, - Tuple, - TypeVar, - Union, - cast, -) - -from git.types import Lit_config_levels, ConfigLevels_Tup, PathLike, assert_never, _T - -if TYPE_CHECKING: - from io import BytesIO - - from git.repo.base import Repo - -T_ConfigParser = TypeVar("T_ConfigParser", bound="GitConfigParser") -T_OMD_value = TypeVar("T_OMD_value", str, bytes, int, float, bool) - -if sys.version_info[:3] < (3, 7, 2): - # typing.Ordereddict not added until Python 3.7.2. - from collections import OrderedDict - - OrderedDict_OMD = OrderedDict -else: - from typing import OrderedDict - - OrderedDict_OMD = OrderedDict[str, List[T_OMD_value]] # type: ignore[assignment, misc] - -# ------------------------------------------------------------- - -_logger = logging.getLogger(__name__) - -CONFIG_LEVELS: ConfigLevels_Tup = ("system", "user", "global", "repository") -"""The configuration level of a configuration file.""" - -CONDITIONAL_INCLUDE_REGEXP = re.compile(r"(?<=includeIf )\"(gitdir|gitdir/i|onbranch):(.+)\"") -"""Section pattern to detect conditional includes. - -See: https://git-scm.com/docs/git-config#_conditional_includes -""" - - -class MetaParserBuilder(abc.ABCMeta): # noqa: B024 - """Utility class wrapping base-class methods into decorators that assure read-only - properties.""" - - def __new__(cls, name: str, bases: Tuple, clsdict: Dict[str, Any]) -> "MetaParserBuilder": - """Equip all base-class methods with a needs_values decorator, and all non-const - methods with a :func:`set_dirty_and_flush_changes` decorator in addition to - that. - """ - kmm = "_mutating_methods_" - if kmm in clsdict: - mutating_methods = clsdict[kmm] - for base in bases: - methods = (t for t in inspect.getmembers(base, inspect.isroutine) if not t[0].startswith("_")) - for name, method in methods: - if name in clsdict: - continue - method_with_values = needs_values(method) - if name in mutating_methods: - method_with_values = set_dirty_and_flush_changes(method_with_values) - # END mutating methods handling - - clsdict[name] = method_with_values - # END for each name/method pair - # END for each base - # END if mutating methods configuration is set - - new_type = super().__new__(cls, name, bases, clsdict) - return new_type - - -def needs_values(func: Callable[..., _T]) -> Callable[..., _T]: - """Return a method for ensuring we read values (on demand) before we try to access - them.""" - - @wraps(func) - def assure_data_present(self: "GitConfigParser", *args: Any, **kwargs: Any) -> _T: - self.read() - return func(self, *args, **kwargs) - - # END wrapper method - return assure_data_present - - -def set_dirty_and_flush_changes(non_const_func: Callable[..., _T]) -> Callable[..., _T]: - """Return a method that checks whether given non constant function may be called. - - If so, the instance will be set dirty. Additionally, we flush the changes right to - disk. - """ - - def flush_changes(self: "GitConfigParser", *args: Any, **kwargs: Any) -> _T: - rval = non_const_func(self, *args, **kwargs) - self._dirty = True - self.write() - return rval - - # END wrapper method - flush_changes.__name__ = non_const_func.__name__ - return flush_changes - - -class SectionConstraint(Generic[T_ConfigParser]): - """Constrains a ConfigParser to only option commands which are constrained to - always use the section we have been initialized with. - - It supports all ConfigParser methods that operate on an option. - - :note: - If used as a context manager, will release the wrapped ConfigParser. - """ - - __slots__ = ("_config", "_section_name") - - _valid_attrs_ = ( - "get_value", - "set_value", - "get", - "set", - "getint", - "getfloat", - "getboolean", - "has_option", - "remove_section", - "remove_option", - "options", - ) - - def __init__(self, config: T_ConfigParser, section: str) -> None: - self._config = config - self._section_name = section - - def __del__(self) -> None: - # Yes, for some reason, we have to call it explicitly for it to work in PY3 ! - # Apparently __del__ doesn't get call anymore if refcount becomes 0 - # Ridiculous ... . - self._config.release() - - def __getattr__(self, attr: str) -> Any: - if attr in self._valid_attrs_: - return lambda *args, **kwargs: self._call_config(attr, *args, **kwargs) - return super().__getattribute__(attr) - - def _call_config(self, method: str, *args: Any, **kwargs: Any) -> Any: - """Call the configuration at the given method which must take a section name as - first argument.""" - return getattr(self._config, method)(self._section_name, *args, **kwargs) - - @property - def config(self) -> T_ConfigParser: - """return: ConfigParser instance we constrain""" - return self._config - - def release(self) -> None: - """Equivalent to :meth:`GitConfigParser.release`, which is called on our - underlying parser instance.""" - return self._config.release() - - def __enter__(self) -> "SectionConstraint[T_ConfigParser]": - self._config.__enter__() - return self - - def __exit__(self, exception_type: str, exception_value: str, traceback: str) -> None: - self._config.__exit__(exception_type, exception_value, traceback) - - -class _OMD(OrderedDict_OMD): - """Ordered multi-dict.""" - - def __setitem__(self, key: str, value: _T) -> None: - super().__setitem__(key, [value]) - - def add(self, key: str, value: Any) -> None: - if key not in self: - super().__setitem__(key, [value]) - return - - super().__getitem__(key).append(value) - - def setall(self, key: str, values: List[_T]) -> None: - super().__setitem__(key, values) - - def __getitem__(self, key: str) -> Any: - return super().__getitem__(key)[-1] - - def getlast(self, key: str) -> Any: - return super().__getitem__(key)[-1] - - def setlast(self, key: str, value: Any) -> None: - if key not in self: - super().__setitem__(key, [value]) - return - - prior = super().__getitem__(key) - prior[-1] = value - - def get(self, key: str, default: Union[_T, None] = None) -> Union[_T, None]: - return super().get(key, [default])[-1] - - def getall(self, key: str) -> List[_T]: - return super().__getitem__(key) - - def items(self) -> List[Tuple[str, _T]]: # type: ignore[override] - """List of (key, last value for key).""" - return [(k, self[k]) for k in self] - - def items_all(self) -> List[Tuple[str, List[_T]]]: - """List of (key, list of values for key).""" - return [(k, self.getall(k)) for k in self] - - -def get_config_path(config_level: Lit_config_levels) -> str: - # We do not support an absolute path of the gitconfig on Windows. - # Use the global config instead. - if sys.platform == "win32" and config_level == "system": - config_level = "global" - - if config_level == "system": - return "/etc/gitconfig" - elif config_level == "user": - config_home = os.environ.get("XDG_CONFIG_HOME") or osp.join(os.environ.get("HOME", "~"), ".config") - return osp.normpath(osp.expanduser(osp.join(config_home, "git", "config"))) - elif config_level == "global": - return osp.normpath(osp.expanduser("~/.gitconfig")) - elif config_level == "repository": - raise ValueError("No repo to get repository configuration from. Use Repo._get_config_path") - else: - # Should not reach here. Will raise ValueError if does. Static typing will warn - # about missing elifs. - assert_never( # type: ignore[unreachable] - config_level, - ValueError(f"Invalid configuration level: {config_level!r}"), - ) - - -class GitConfigParser(cp.RawConfigParser, metaclass=MetaParserBuilder): - """Implements specifics required to read git style configuration files. - - This variation behaves much like the :manpage:`git-config(1)` command, such that the - configuration will be read on demand based on the filepath given during - initialization. - - The changes will automatically be written once the instance goes out of scope, but - can be triggered manually as well. - - The configuration file will be locked if you intend to change values preventing - other instances to write concurrently. - - :note: - The config is case-sensitive even when queried, hence section and option names - must match perfectly. - - :note: - If used as a context manager, this will release the locked file. - """ - - # { Configuration - t_lock = LockFile - """The lock type determines the type of lock to use in new configuration readers. - - They must be compatible to the :class:`~git.util.LockFile` interface. - A suitable alternative would be the :class:`~git.util.BlockingLockFile`. - """ - - re_comment = re.compile(r"^\s*[#;]") - # } END configuration - - optvalueonly_source = r"\s*(?P

` tags. - """ - env = {} if env is None else env - return self.renderer.render(self.parseInline(src, env), self.options, env) - - # link methods - - def validateLink(self, url: str) -> bool: - """Validate if the URL link is allowed in output. - - This validator can prohibit more than really needed to prevent XSS. - It's a tradeoff to keep code simple and to be secure by default. - - Note: the url should be normalized at this point, and existing entities decoded. - """ - return normalize_url.validateLink(url) - - def normalizeLink(self, url: str) -> str: - """Normalize destination URLs in links - - :: - - [label]: destination 'title' - ^^^^^^^^^^^ - """ - return normalize_url.normalizeLink(url) - - def normalizeLinkText(self, link: str) -> str: - """Normalize autolink content - - :: - - - ~~~~~~~~~~~ - """ - return normalize_url.normalizeLinkText(link) diff --git a/venv/lib/python3.10/site-packages/markdown_it/parser_block.py b/venv/lib/python3.10/site-packages/markdown_it/parser_block.py deleted file mode 100644 index 72360f9b31bebca77250168bccae646e9a67dc6e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/parser_block.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Block-level tokenizer.""" -from __future__ import annotations - -import logging -from typing import TYPE_CHECKING, Callable - -from . import rules_block -from .ruler import Ruler -from .rules_block.state_block import StateBlock -from .token import Token -from .utils import EnvType - -if TYPE_CHECKING: - from markdown_it import MarkdownIt - -LOGGER = logging.getLogger(__name__) - - -RuleFuncBlockType = Callable[[StateBlock, int, int, bool], bool] -"""(state: StateBlock, startLine: int, endLine: int, silent: bool) -> matched: bool) - -`silent` disables token generation, useful for lookahead. -""" - -_rules: list[tuple[str, RuleFuncBlockType, list[str]]] = [ - # First 2 params - rule name & source. Secondary array - list of rules, - # which can be terminated by this one. - ("table", rules_block.table, ["paragraph", "reference"]), - ("code", rules_block.code, []), - ("fence", rules_block.fence, ["paragraph", "reference", "blockquote", "list"]), - ( - "blockquote", - rules_block.blockquote, - ["paragraph", "reference", "blockquote", "list"], - ), - ("hr", rules_block.hr, ["paragraph", "reference", "blockquote", "list"]), - ("list", rules_block.list_block, ["paragraph", "reference", "blockquote"]), - ("reference", rules_block.reference, []), - ("html_block", rules_block.html_block, ["paragraph", "reference", "blockquote"]), - ("heading", rules_block.heading, ["paragraph", "reference", "blockquote"]), - ("lheading", rules_block.lheading, []), - ("paragraph", rules_block.paragraph, []), -] - - -class ParserBlock: - """ - ParserBlock#ruler -> Ruler - - [[Ruler]] instance. Keep configuration of block rules. - """ - - def __init__(self) -> None: - self.ruler = Ruler[RuleFuncBlockType]() - for name, rule, alt in _rules: - self.ruler.push(name, rule, {"alt": alt}) - - def tokenize(self, state: StateBlock, startLine: int, endLine: int) -> None: - """Generate tokens for input range.""" - rules = self.ruler.getRules("") - line = startLine - maxNesting = state.md.options.maxNesting - hasEmptyLines = False - - while line < endLine: - state.line = line = state.skipEmptyLines(line) - if line >= endLine: - break - if state.sCount[line] < state.blkIndent: - # Termination condition for nested calls. - # Nested calls currently used for blockquotes & lists - break - if state.level >= maxNesting: - # If nesting level exceeded - skip tail to the end. - # That's not ordinary situation and we should not care about content. - state.line = endLine - break - - # Try all possible rules. - # On success, rule should: - # - update `state.line` - # - update `state.tokens` - # - return True - for rule in rules: - if rule(state, line, endLine, False): - break - - # set state.tight if we had an empty line before current tag - # i.e. latest empty line should not count - state.tight = not hasEmptyLines - - line = state.line - - # paragraph might "eat" one newline after it in nested lists - if (line - 1) < endLine and state.isEmpty(line - 1): - hasEmptyLines = True - - if line < endLine and state.isEmpty(line): - hasEmptyLines = True - line += 1 - state.line = line - - def parse( - self, src: str, md: MarkdownIt, env: EnvType, outTokens: list[Token] - ) -> list[Token] | None: - """Process input string and push block tokens into `outTokens`.""" - if not src: - return None - state = StateBlock(src, md, env, outTokens) - self.tokenize(state, state.line, state.lineMax) - return state.tokens diff --git a/venv/lib/python3.10/site-packages/markdown_it/parser_core.py b/venv/lib/python3.10/site-packages/markdown_it/parser_core.py deleted file mode 100644 index ca5ab2566ba5bd00e654a7af39e3603717ae7194..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/parser_core.py +++ /dev/null @@ -1,45 +0,0 @@ -""" - * class Core - * - * Top-level rules executor. Glues block/inline parsers and does intermediate - * transformations. -""" -from __future__ import annotations - -from typing import Callable - -from .ruler import Ruler -from .rules_core import ( - block, - inline, - linkify, - normalize, - replace, - smartquotes, - text_join, -) -from .rules_core.state_core import StateCore - -RuleFuncCoreType = Callable[[StateCore], None] - -_rules: list[tuple[str, RuleFuncCoreType]] = [ - ("normalize", normalize), - ("block", block), - ("inline", inline), - ("linkify", linkify), - ("replacements", replace), - ("smartquotes", smartquotes), - ("text_join", text_join), -] - - -class ParserCore: - def __init__(self) -> None: - self.ruler = Ruler[RuleFuncCoreType]() - for name, rule in _rules: - self.ruler.push(name, rule) - - def process(self, state: StateCore) -> None: - """Executes core chain rules.""" - for rule in self.ruler.getRules(""): - rule(state) diff --git a/venv/lib/python3.10/site-packages/markdown_it/parser_inline.py b/venv/lib/python3.10/site-packages/markdown_it/parser_inline.py deleted file mode 100644 index 0026c3839d14fe8b78ff033abfac0c299b8a0c54..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/parser_inline.py +++ /dev/null @@ -1,147 +0,0 @@ -"""Tokenizes paragraph content. -""" -from __future__ import annotations - -from typing import TYPE_CHECKING, Callable - -from . import rules_inline -from .ruler import Ruler -from .rules_inline.state_inline import StateInline -from .token import Token -from .utils import EnvType - -if TYPE_CHECKING: - from markdown_it import MarkdownIt - - -# Parser rules -RuleFuncInlineType = Callable[[StateInline, bool], bool] -"""(state: StateInline, silent: bool) -> matched: bool) - -`silent` disables token generation, useful for lookahead. -""" -_rules: list[tuple[str, RuleFuncInlineType]] = [ - ("text", rules_inline.text), - ("linkify", rules_inline.linkify), - ("newline", rules_inline.newline), - ("escape", rules_inline.escape), - ("backticks", rules_inline.backtick), - ("strikethrough", rules_inline.strikethrough.tokenize), - ("emphasis", rules_inline.emphasis.tokenize), - ("link", rules_inline.link), - ("image", rules_inline.image), - ("autolink", rules_inline.autolink), - ("html_inline", rules_inline.html_inline), - ("entity", rules_inline.entity), -] - -# Note `rule2` ruleset was created specifically for emphasis/strikethrough -# post-processing and may be changed in the future. -# -# Don't use this for anything except pairs (plugins working with `balance_pairs`). -# -RuleFuncInline2Type = Callable[[StateInline], None] -_rules2: list[tuple[str, RuleFuncInline2Type]] = [ - ("balance_pairs", rules_inline.link_pairs), - ("strikethrough", rules_inline.strikethrough.postProcess), - ("emphasis", rules_inline.emphasis.postProcess), - # rules for pairs separate '**' into its own text tokens, which may be left unused, - # rule below merges unused segments back with the rest of the text - ("fragments_join", rules_inline.fragments_join), -] - - -class ParserInline: - def __init__(self) -> None: - self.ruler = Ruler[RuleFuncInlineType]() - for name, rule in _rules: - self.ruler.push(name, rule) - # Second ruler used for post-processing (e.g. in emphasis-like rules) - self.ruler2 = Ruler[RuleFuncInline2Type]() - for name, rule2 in _rules2: - self.ruler2.push(name, rule2) - - def skipToken(self, state: StateInline) -> None: - """Skip single token by running all rules in validation mode; - returns `True` if any rule reported success - """ - ok = False - pos = state.pos - rules = self.ruler.getRules("") - maxNesting = state.md.options["maxNesting"] - cache = state.cache - - if pos in cache: - state.pos = cache[pos] - return - - if state.level < maxNesting: - for rule in rules: - # Increment state.level and decrement it later to limit recursion. - # It's harmless to do here, because no tokens are created. - # But ideally, we'd need a separate private state variable for this purpose. - state.level += 1 - ok = rule(state, True) - state.level -= 1 - if ok: - break - else: - # Too much nesting, just skip until the end of the paragraph. - # - # NOTE: this will cause links to behave incorrectly in the following case, - # when an amount of `[` is exactly equal to `maxNesting + 1`: - # - # [[[[[[[[[[[[[[[[[[[[[foo]() - # - # TODO: remove this workaround when CM standard will allow nested links - # (we can replace it by preventing links from being parsed in - # validation mode) - # - state.pos = state.posMax - - if not ok: - state.pos += 1 - cache[pos] = state.pos - - def tokenize(self, state: StateInline) -> None: - """Generate tokens for input range.""" - ok = False - rules = self.ruler.getRules("") - end = state.posMax - maxNesting = state.md.options["maxNesting"] - - while state.pos < end: - # Try all possible rules. - # On success, rule should: - # - # - update `state.pos` - # - update `state.tokens` - # - return true - - if state.level < maxNesting: - for rule in rules: - ok = rule(state, False) - if ok: - break - - if ok: - if state.pos >= end: - break - continue - - state.pending += state.src[state.pos] - state.pos += 1 - - if state.pending: - state.pushPending() - - def parse( - self, src: str, md: MarkdownIt, env: EnvType, tokens: list[Token] - ) -> list[Token]: - """Process input string and push inline tokens into `tokens`""" - state = StateInline(src, md, env, tokens) - self.tokenize(state) - rules2 = self.ruler2.getRules("") - for rule in rules2: - rule(state) - return state.tokens diff --git a/venv/lib/python3.10/site-packages/markdown_it/port.yaml b/venv/lib/python3.10/site-packages/markdown_it/port.yaml deleted file mode 100644 index 3e289e9e27d53efea476622d1d6dc60799920835..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/port.yaml +++ /dev/null @@ -1,48 +0,0 @@ -- package: markdown-it/markdown-it - version: 13.0.1 - commit: e843acc9edad115cbf8cf85e676443f01658be08 - date: May 3, 2022 - notes: - - Rename variables that use python built-in names, e.g. - - `max` -> `maximum` - - `len` -> `length` - - `str` -> `string` - - | - Convert JS `for` loops to `while` loops - this is generally the main difference between the codes, - because in python you can't do e.g. `for {i=1;i PresetType: - config = commonmark.make() - config["components"]["core"]["rules"].append("linkify") - config["components"]["block"]["rules"].append("table") - config["components"]["inline"]["rules"].extend(["strikethrough", "linkify"]) - config["components"]["inline"]["rules2"].append("strikethrough") - config["options"]["linkify"] = True - config["options"]["html"] = True - return config diff --git a/venv/lib/python3.10/site-packages/markdown_it/presets/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/presets/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index e7df66ffb21251c6a6b0dca3ab43f76f729d6857..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/presets/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/presets/__pycache__/commonmark.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/presets/__pycache__/commonmark.cpython-310.pyc deleted file mode 100644 index 759c32a358f4bdb093cb0b5819cb37dc0baff322..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/presets/__pycache__/commonmark.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/presets/__pycache__/default.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/presets/__pycache__/default.cpython-310.pyc deleted file mode 100644 index 0602ab30483ce8627843e804135abd88422000ea..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/presets/__pycache__/default.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/presets/__pycache__/zero.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/presets/__pycache__/zero.cpython-310.pyc deleted file mode 100644 index b6cbb64154f7f7795834c9aa9a1ca3ced40178c5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/presets/__pycache__/zero.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/presets/commonmark.py b/venv/lib/python3.10/site-packages/markdown_it/presets/commonmark.py deleted file mode 100644 index 3990d4344aeb9e07449acf8aa749cb27b0a0e66c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/presets/commonmark.py +++ /dev/null @@ -1,74 +0,0 @@ -"""Commonmark default options. - -This differs to presets.default, -primarily in that it allows HTML and does not enable components: - -- block: table -- inline: strikethrough -""" -from ..utils import PresetType - - -def make() -> PresetType: - return { - "options": { - "maxNesting": 20, # Internal protection, recursion limit - "html": True, # Enable HTML tags in source, - # this is just a shorthand for .enable(["html_inline", "html_block"]) - # used by the linkify rule: - "linkify": False, # autoconvert URL-like texts to links - # used by the replacements and smartquotes rules - # Enable some language-neutral replacements + quotes beautification - "typographer": False, - # used by the smartquotes rule: - # Double + single quotes replacement pairs, when typographer enabled, - # and smartquotes on. Could be either a String or an Array. - # - # For example, you can use '«»„“' for Russian, '„“‚‘' for German, - # and ['«\xA0', '\xA0»', '‹\xA0', '\xA0›'] for French (including nbsp). - "quotes": "\u201c\u201d\u2018\u2019", # /* “”‘’ */ - # Renderer specific; these options are used directly in the HTML renderer - "xhtmlOut": True, # Use '/' to close single tags (
) - "breaks": False, # Convert '\n' in paragraphs into
- "langPrefix": "language-", # CSS language prefix for fenced blocks - # Highlighter function. Should return escaped HTML, - # or '' if the source string is not changed and should be escaped externally. - # If result starts with PresetType: - return { - "options": { - "maxNesting": 100, # Internal protection, recursion limit - "html": False, # Enable HTML tags in source - # this is just a shorthand for .disable(["html_inline", "html_block"]) - # used by the linkify rule: - "linkify": False, # autoconvert URL-like texts to links - # used by the replacements and smartquotes rules: - # Enable some language-neutral replacements + quotes beautification - "typographer": False, - # used by the smartquotes rule: - # Double + single quotes replacement pairs, when typographer enabled, - # and smartquotes on. Could be either a String or an Array. - # For example, you can use '«»„“' for Russian, '„“‚‘' for German, - # and ['«\xA0', '\xA0»', '‹\xA0', '\xA0›'] for French (including nbsp). - "quotes": "\u201c\u201d\u2018\u2019", # /* “”‘’ */ - # Renderer specific; these options are used directly in the HTML renderer - "xhtmlOut": False, # Use '/' to close single tags (
) - "breaks": False, # Convert '\n' in paragraphs into
- "langPrefix": "language-", # CSS language prefix for fenced blocks - # Highlighter function. Should return escaped HTML, - # or '' if the source string is not changed and should be escaped externally. - # If result starts with PresetType: - return { - "options": { - "maxNesting": 20, # Internal protection, recursion limit - "html": False, # Enable HTML tags in source - # this is just a shorthand for .disable(["html_inline", "html_block"]) - # used by the linkify rule: - "linkify": False, # autoconvert URL-like texts to links - # used by the replacements and smartquotes rules: - # Enable some language-neutral replacements + quotes beautification - "typographer": False, - # used by the smartquotes rule: - # Double + single quotes replacement pairs, when typographer enabled, - # and smartquotes on. Could be either a String or an Array. - # For example, you can use '«»„“' for Russian, '„“‚‘' for German, - # and ['«\xA0', '\xA0»', '‹\xA0', '\xA0›'] for French (including nbsp). - "quotes": "\u201c\u201d\u2018\u2019", # /* “”‘’ */ - # Renderer specific; these options are used directly in the HTML renderer - "xhtmlOut": False, # Use '/' to close single tags (
) - "breaks": False, # Convert '\n' in paragraphs into
- "langPrefix": "language-", # CSS language prefix for fenced blocks - # Highlighter function. Should return escaped HTML, - # or '' if the source string is not changed and should be escaped externally. - # If result starts with Any: - ... - - -class RendererHTML(RendererProtocol): - """Contains render rules for tokens. Can be updated and extended. - - Example: - - Each rule is called as independent static function with fixed signature: - - :: - - class Renderer: - def token_type_name(self, tokens, idx, options, env) { - # ... - return renderedHTML - - :: - - class CustomRenderer(RendererHTML): - def strong_open(self, tokens, idx, options, env): - return '' - def strong_close(self, tokens, idx, options, env): - return '' - - md = MarkdownIt(renderer_cls=CustomRenderer) - - result = md.render(...) - - See https://github.com/markdown-it/markdown-it/blob/master/lib/renderer.js - for more details and examples. - """ - - __output__ = "html" - - def __init__(self, parser: Any = None): - self.rules = { - k: v - for k, v in inspect.getmembers(self, predicate=inspect.ismethod) - if not (k.startswith("render") or k.startswith("_")) - } - - def render( - self, tokens: Sequence[Token], options: OptionsDict, env: EnvType - ) -> str: - """Takes token stream and generates HTML. - - :param tokens: list on block tokens to render - :param options: params of parser instance - :param env: additional data from parsed input - - """ - result = "" - - for i, token in enumerate(tokens): - if token.type == "inline": - if token.children: - result += self.renderInline(token.children, options, env) - elif token.type in self.rules: - result += self.rules[token.type](tokens, i, options, env) - else: - result += self.renderToken(tokens, i, options, env) - - return result - - def renderInline( - self, tokens: Sequence[Token], options: OptionsDict, env: EnvType - ) -> str: - """The same as ``render``, but for single token of `inline` type. - - :param tokens: list on block tokens to render - :param options: params of parser instance - :param env: additional data from parsed input (references, for example) - """ - result = "" - - for i, token in enumerate(tokens): - if token.type in self.rules: - result += self.rules[token.type](tokens, i, options, env) - else: - result += self.renderToken(tokens, i, options, env) - - return result - - def renderToken( - self, - tokens: Sequence[Token], - idx: int, - options: OptionsDict, - env: EnvType, - ) -> str: - """Default token renderer. - - Can be overridden by custom function - - :param idx: token index to render - :param options: params of parser instance - """ - result = "" - needLf = False - token = tokens[idx] - - # Tight list paragraphs - if token.hidden: - return "" - - # Insert a newline between hidden paragraph and subsequent opening - # block-level tag. - # - # For example, here we should insert a newline before blockquote: - # - a - # > - # - if token.block and token.nesting != -1 and idx and tokens[idx - 1].hidden: - result += "\n" - - # Add token name, e.g. ``. - # - needLf = False - - result += ">\n" if needLf else ">" - - return result - - @staticmethod - def renderAttrs(token: Token) -> str: - """Render token attributes to string.""" - result = "" - - for key, value in token.attrItems(): - result += " " + escapeHtml(key) + '="' + escapeHtml(str(value)) + '"' - - return result - - def renderInlineAsText( - self, - tokens: Sequence[Token] | None, - options: OptionsDict, - env: EnvType, - ) -> str: - """Special kludge for image `alt` attributes to conform CommonMark spec. - - Don't try to use it! Spec requires to show `alt` content with stripped markup, - instead of simple escaping. - - :param tokens: list on block tokens to render - :param options: params of parser instance - :param env: additional data from parsed input - """ - result = "" - - for token in tokens or []: - if token.type == "text": - result += token.content - elif token.type == "image": - if token.children: - result += self.renderInlineAsText(token.children, options, env) - elif token.type == "softbreak": - result += "\n" - - return result - - ################################################### - - def code_inline( - self, tokens: Sequence[Token], idx: int, options: OptionsDict, env: EnvType - ) -> str: - token = tokens[idx] - return ( - "" - + escapeHtml(tokens[idx].content) - + "" - ) - - def code_block( - self, - tokens: Sequence[Token], - idx: int, - options: OptionsDict, - env: EnvType, - ) -> str: - token = tokens[idx] - - return ( - "" - + escapeHtml(tokens[idx].content) - + "\n" - ) - - def fence( - self, - tokens: Sequence[Token], - idx: int, - options: OptionsDict, - env: EnvType, - ) -> str: - token = tokens[idx] - info = unescapeAll(token.info).strip() if token.info else "" - langName = "" - langAttrs = "" - - if info: - arr = info.split(maxsplit=1) - langName = arr[0] - if len(arr) == 2: - langAttrs = arr[1] - - if options.highlight: - highlighted = options.highlight( - token.content, langName, langAttrs - ) or escapeHtml(token.content) - else: - highlighted = escapeHtml(token.content) - - if highlighted.startswith("" - + highlighted - + "\n" - ) - - return ( - "

"
-            + highlighted
-            + "
\n" - ) - - def image( - self, - tokens: Sequence[Token], - idx: int, - options: OptionsDict, - env: EnvType, - ) -> str: - token = tokens[idx] - - # "alt" attr MUST be set, even if empty. Because it's mandatory and - # should be placed on proper position for tests. - if token.children: - token.attrSet("alt", self.renderInlineAsText(token.children, options, env)) - else: - token.attrSet("alt", "") - - return self.renderToken(tokens, idx, options, env) - - def hardbreak( - self, tokens: Sequence[Token], idx: int, options: OptionsDict, env: EnvType - ) -> str: - return "
\n" if options.xhtmlOut else "
\n" - - def softbreak( - self, tokens: Sequence[Token], idx: int, options: OptionsDict, env: EnvType - ) -> str: - return ( - ("
\n" if options.xhtmlOut else "
\n") if options.breaks else "\n" - ) - - def text( - self, tokens: Sequence[Token], idx: int, options: OptionsDict, env: EnvType - ) -> str: - return escapeHtml(tokens[idx].content) - - def html_block( - self, tokens: Sequence[Token], idx: int, options: OptionsDict, env: EnvType - ) -> str: - return tokens[idx].content - - def html_inline( - self, tokens: Sequence[Token], idx: int, options: OptionsDict, env: EnvType - ) -> str: - return tokens[idx].content diff --git a/venv/lib/python3.10/site-packages/markdown_it/ruler.py b/venv/lib/python3.10/site-packages/markdown_it/ruler.py deleted file mode 100644 index bd8baba34e6907d5c5086ed284639518cc802f86..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/ruler.py +++ /dev/null @@ -1,276 +0,0 @@ -""" -class Ruler - -Helper class, used by [[MarkdownIt#core]], [[MarkdownIt#block]] and -[[MarkdownIt#inline]] to manage sequences of functions (rules): - -- keep rules in defined order -- assign the name to each rule -- enable/disable rules -- add/replace rules -- allow assign rules to additional named chains (in the same) -- caching lists of active rules - -You will not need use this class directly until write plugins. For simple -rules control use [[MarkdownIt.disable]], [[MarkdownIt.enable]] and -[[MarkdownIt.use]]. -""" -from __future__ import annotations - -from collections.abc import Iterable -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Generic, TypedDict, TypeVar -import warnings - -from markdown_it._compat import DATACLASS_KWARGS - -from .utils import EnvType - -if TYPE_CHECKING: - from markdown_it import MarkdownIt - - -class StateBase: - def __init__(self, src: str, md: MarkdownIt, env: EnvType): - self.src = src - self.env = env - self.md = md - - @property - def src(self) -> str: - return self._src - - @src.setter - def src(self, value: str) -> None: - self._src = value - self._srcCharCode: tuple[int, ...] | None = None - - @property - def srcCharCode(self) -> tuple[int, ...]: - warnings.warn( - "StateBase.srcCharCode is deprecated. Use StateBase.src instead.", - DeprecationWarning, - stacklevel=2, - ) - if self._srcCharCode is None: - self._srcCharCode = tuple(ord(c) for c in self._src) - return self._srcCharCode - - -class RuleOptionsType(TypedDict, total=False): - alt: list[str] - - -RuleFuncTv = TypeVar("RuleFuncTv") -"""A rule function, whose signature is dependent on the state type.""" - - -@dataclass(**DATACLASS_KWARGS) -class Rule(Generic[RuleFuncTv]): - name: str - enabled: bool - fn: RuleFuncTv = field(repr=False) - alt: list[str] - - -class Ruler(Generic[RuleFuncTv]): - def __init__(self) -> None: - # List of added rules. - self.__rules__: list[Rule[RuleFuncTv]] = [] - # Cached rule chains. - # First level - chain name, '' for default. - # Second level - diginal anchor for fast filtering by charcodes. - self.__cache__: dict[str, list[RuleFuncTv]] | None = None - - def __find__(self, name: str) -> int: - """Find rule index by name""" - for i, rule in enumerate(self.__rules__): - if rule.name == name: - return i - return -1 - - def __compile__(self) -> None: - """Build rules lookup cache""" - chains = {""} - # collect unique names - for rule in self.__rules__: - if not rule.enabled: - continue - for name in rule.alt: - chains.add(name) - self.__cache__ = {} - for chain in chains: - self.__cache__[chain] = [] - for rule in self.__rules__: - if not rule.enabled: - continue - if chain and (chain not in rule.alt): - continue - self.__cache__[chain].append(rule.fn) - - def at( - self, ruleName: str, fn: RuleFuncTv, options: RuleOptionsType | None = None - ) -> None: - """Replace rule by name with new function & options. - - :param ruleName: rule name to replace. - :param fn: new rule function. - :param options: new rule options (not mandatory). - :raises: KeyError if name not found - """ - index = self.__find__(ruleName) - options = options or {} - if index == -1: - raise KeyError(f"Parser rule not found: {ruleName}") - self.__rules__[index].fn = fn - self.__rules__[index].alt = options.get("alt", []) - self.__cache__ = None - - def before( - self, - beforeName: str, - ruleName: str, - fn: RuleFuncTv, - options: RuleOptionsType | None = None, - ) -> None: - """Add new rule to chain before one with given name. - - :param beforeName: new rule will be added before this one. - :param ruleName: new rule will be added before this one. - :param fn: new rule function. - :param options: new rule options (not mandatory). - :raises: KeyError if name not found - """ - index = self.__find__(beforeName) - options = options or {} - if index == -1: - raise KeyError(f"Parser rule not found: {beforeName}") - self.__rules__.insert( - index, Rule[RuleFuncTv](ruleName, True, fn, options.get("alt", [])) - ) - self.__cache__ = None - - def after( - self, - afterName: str, - ruleName: str, - fn: RuleFuncTv, - options: RuleOptionsType | None = None, - ) -> None: - """Add new rule to chain after one with given name. - - :param afterName: new rule will be added after this one. - :param ruleName: new rule will be added after this one. - :param fn: new rule function. - :param options: new rule options (not mandatory). - :raises: KeyError if name not found - """ - index = self.__find__(afterName) - options = options or {} - if index == -1: - raise KeyError(f"Parser rule not found: {afterName}") - self.__rules__.insert( - index + 1, Rule[RuleFuncTv](ruleName, True, fn, options.get("alt", [])) - ) - self.__cache__ = None - - def push( - self, ruleName: str, fn: RuleFuncTv, options: RuleOptionsType | None = None - ) -> None: - """Push new rule to the end of chain. - - :param ruleName: new rule will be added to the end of chain. - :param fn: new rule function. - :param options: new rule options (not mandatory). - - """ - self.__rules__.append( - Rule[RuleFuncTv](ruleName, True, fn, (options or {}).get("alt", [])) - ) - self.__cache__ = None - - def enable( - self, names: str | Iterable[str], ignoreInvalid: bool = False - ) -> list[str]: - """Enable rules with given names. - - :param names: name or list of rule names to enable. - :param ignoreInvalid: ignore errors when rule not found - :raises: KeyError if name not found and not ignoreInvalid - :return: list of found rule names - """ - if isinstance(names, str): - names = [names] - result: list[str] = [] - for name in names: - idx = self.__find__(name) - if (idx < 0) and ignoreInvalid: - continue - if (idx < 0) and not ignoreInvalid: - raise KeyError(f"Rules manager: invalid rule name {name}") - self.__rules__[idx].enabled = True - result.append(name) - self.__cache__ = None - return result - - def enableOnly( - self, names: str | Iterable[str], ignoreInvalid: bool = False - ) -> list[str]: - """Enable rules with given names, and disable everything else. - - :param names: name or list of rule names to enable. - :param ignoreInvalid: ignore errors when rule not found - :raises: KeyError if name not found and not ignoreInvalid - :return: list of found rule names - """ - if isinstance(names, str): - names = [names] - for rule in self.__rules__: - rule.enabled = False - return self.enable(names, ignoreInvalid) - - def disable( - self, names: str | Iterable[str], ignoreInvalid: bool = False - ) -> list[str]: - """Disable rules with given names. - - :param names: name or list of rule names to enable. - :param ignoreInvalid: ignore errors when rule not found - :raises: KeyError if name not found and not ignoreInvalid - :return: list of found rule names - """ - if isinstance(names, str): - names = [names] - result = [] - for name in names: - idx = self.__find__(name) - if (idx < 0) and ignoreInvalid: - continue - if (idx < 0) and not ignoreInvalid: - raise KeyError(f"Rules manager: invalid rule name {name}") - self.__rules__[idx].enabled = False - result.append(name) - self.__cache__ = None - return result - - def getRules(self, chainName: str = "") -> list[RuleFuncTv]: - """Return array of active functions (rules) for given chain name. - It analyzes rules configuration, compiles caches if not exists and returns result. - - Default chain name is `''` (empty string). It can't be skipped. - That's done intentionally, to keep signature monomorphic for high speed. - - """ - if self.__cache__ is None: - self.__compile__() - assert self.__cache__ is not None - # Chain can be empty, if rules disabled. But we still have to return Array. - return self.__cache__.get(chainName, []) or [] - - def get_all_rules(self) -> list[str]: - """Return all available rule names.""" - return [r.name for r in self.__rules__] - - def get_active_rules(self) -> list[str]: - """Return the active rule names.""" - return [r.name for r in self.__rules__ if r.enabled] diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__init__.py b/venv/lib/python3.10/site-packages/markdown_it/rules_block/__init__.py deleted file mode 100644 index bcf138df9098dc0bef30c6785e37f266af6b5168..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__init__.py +++ /dev/null @@ -1,27 +0,0 @@ -__all__ = ( - "StateBlock", - "paragraph", - "heading", - "lheading", - "code", - "fence", - "hr", - "list_block", - "reference", - "blockquote", - "html_block", - "table", -) - -from .blockquote import blockquote -from .code import code -from .fence import fence -from .heading import heading -from .hr import hr -from .html_block import html_block -from .lheading import lheading -from .list import list_block -from .paragraph import paragraph -from .reference import reference -from .state_block import StateBlock -from .table import table diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 12da30334d7b149011e55eaf0406c721536fcf25..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/blockquote.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/blockquote.cpython-310.pyc deleted file mode 100644 index ff3a0cf1faf73e632a89a9bfa1b0f925e1b55df2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/blockquote.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/code.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/code.cpython-310.pyc deleted file mode 100644 index a759334c556b0026a3e1ff53ac131ab17088f5e1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/code.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/fence.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/fence.cpython-310.pyc deleted file mode 100644 index 33d68e8231878e2509c36c2377d134141d16195d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/fence.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/heading.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/heading.cpython-310.pyc deleted file mode 100644 index acf8c92268cf576ae131a0e01b7502c746182a83..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/heading.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/hr.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/hr.cpython-310.pyc deleted file mode 100644 index 748c579645096139c04ede2c42338dd0514d163d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/hr.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/html_block.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/html_block.cpython-310.pyc deleted file mode 100644 index 62bfde5bf8e024db8c972d58a662542706f52c98..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/html_block.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/lheading.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/lheading.cpython-310.pyc deleted file mode 100644 index acade0f569ad61bcae06e481e0afd9c922af4752..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/lheading.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/list.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/list.cpython-310.pyc deleted file mode 100644 index 71a0926912f86bb7413a42e01c3f43d373e0a8ea..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/list.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/paragraph.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/paragraph.cpython-310.pyc deleted file mode 100644 index d13c45d5e79a734792c98277470b3118148bcc9f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/paragraph.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/reference.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/reference.cpython-310.pyc deleted file mode 100644 index 022cef13c535275691c6c37ec8937686f51254e2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/reference.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/state_block.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/state_block.cpython-310.pyc deleted file mode 100644 index d6ee774329e91b917b65ef69c4920d07aba0c42c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/state_block.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/table.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/table.cpython-310.pyc deleted file mode 100644 index 9285510f136e2f76337550f451f30852d34ac6ab..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_block/__pycache__/table.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/blockquote.py b/venv/lib/python3.10/site-packages/markdown_it/rules_block/blockquote.py deleted file mode 100644 index 0c9081b9cbd4b49d39d75427fd806e56c485a5fd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_block/blockquote.py +++ /dev/null @@ -1,299 +0,0 @@ -# Block quotes -from __future__ import annotations - -import logging - -from ..common.utils import isStrSpace -from .state_block import StateBlock - -LOGGER = logging.getLogger(__name__) - - -def blockquote(state: StateBlock, startLine: int, endLine: int, silent: bool) -> bool: - LOGGER.debug( - "entering blockquote: %s, %s, %s, %s", state, startLine, endLine, silent - ) - - oldLineMax = state.lineMax - pos = state.bMarks[startLine] + state.tShift[startLine] - max = state.eMarks[startLine] - - if state.is_code_block(startLine): - return False - - # check the block quote marker - try: - if state.src[pos] != ">": - return False - except IndexError: - return False - pos += 1 - - # we know that it's going to be a valid blockquote, - # so no point trying to find the end of it in silent mode - if silent: - return True - - # set offset past spaces and ">" - initial = offset = state.sCount[startLine] + 1 - - try: - second_char: str | None = state.src[pos] - except IndexError: - second_char = None - - # skip one optional space after '>' - if second_char == " ": - # ' > test ' - # ^ -- position start of line here: - pos += 1 - initial += 1 - offset += 1 - adjustTab = False - spaceAfterMarker = True - elif second_char == "\t": - spaceAfterMarker = True - - if (state.bsCount[startLine] + offset) % 4 == 3: - # ' >\t test ' - # ^ -- position start of line here (tab has width==1) - pos += 1 - initial += 1 - offset += 1 - adjustTab = False - else: - # ' >\t test ' - # ^ -- position start of line here + shift bsCount slightly - # to make extra space appear - adjustTab = True - - else: - spaceAfterMarker = False - - oldBMarks = [state.bMarks[startLine]] - state.bMarks[startLine] = pos - - while pos < max: - ch = state.src[pos] - - if isStrSpace(ch): - if ch == "\t": - offset += ( - 4 - - (offset + state.bsCount[startLine] + (1 if adjustTab else 0)) % 4 - ) - else: - offset += 1 - - else: - break - - pos += 1 - - oldBSCount = [state.bsCount[startLine]] - state.bsCount[startLine] = ( - state.sCount[startLine] + 1 + (1 if spaceAfterMarker else 0) - ) - - lastLineEmpty = pos >= max - - oldSCount = [state.sCount[startLine]] - state.sCount[startLine] = offset - initial - - oldTShift = [state.tShift[startLine]] - state.tShift[startLine] = pos - state.bMarks[startLine] - - terminatorRules = state.md.block.ruler.getRules("blockquote") - - oldParentType = state.parentType - state.parentType = "blockquote" - - # Search the end of the block - # - # Block ends with either: - # 1. an empty line outside: - # ``` - # > test - # - # ``` - # 2. an empty line inside: - # ``` - # > - # test - # ``` - # 3. another tag: - # ``` - # > test - # - - - - # ``` - - # for (nextLine = startLine + 1; nextLine < endLine; nextLine++) { - nextLine = startLine + 1 - while nextLine < endLine: - # check if it's outdented, i.e. it's inside list item and indented - # less than said list item: - # - # ``` - # 1. anything - # > current blockquote - # 2. checking this line - # ``` - isOutdented = state.sCount[nextLine] < state.blkIndent - - pos = state.bMarks[nextLine] + state.tShift[nextLine] - max = state.eMarks[nextLine] - - if pos >= max: - # Case 1: line is not inside the blockquote, and this line is empty. - break - - evaluatesTrue = state.src[pos] == ">" and not isOutdented - pos += 1 - if evaluatesTrue: - # This line is inside the blockquote. - - # set offset past spaces and ">" - initial = offset = state.sCount[nextLine] + 1 - - try: - next_char: str | None = state.src[pos] - except IndexError: - next_char = None - - # skip one optional space after '>' - if next_char == " ": - # ' > test ' - # ^ -- position start of line here: - pos += 1 - initial += 1 - offset += 1 - adjustTab = False - spaceAfterMarker = True - elif next_char == "\t": - spaceAfterMarker = True - - if (state.bsCount[nextLine] + offset) % 4 == 3: - # ' >\t test ' - # ^ -- position start of line here (tab has width==1) - pos += 1 - initial += 1 - offset += 1 - adjustTab = False - else: - # ' >\t test ' - # ^ -- position start of line here + shift bsCount slightly - # to make extra space appear - adjustTab = True - - else: - spaceAfterMarker = False - - oldBMarks.append(state.bMarks[nextLine]) - state.bMarks[nextLine] = pos - - while pos < max: - ch = state.src[pos] - - if isStrSpace(ch): - if ch == "\t": - offset += ( - 4 - - ( - offset - + state.bsCount[nextLine] - + (1 if adjustTab else 0) - ) - % 4 - ) - else: - offset += 1 - else: - break - - pos += 1 - - lastLineEmpty = pos >= max - - oldBSCount.append(state.bsCount[nextLine]) - state.bsCount[nextLine] = ( - state.sCount[nextLine] + 1 + (1 if spaceAfterMarker else 0) - ) - - oldSCount.append(state.sCount[nextLine]) - state.sCount[nextLine] = offset - initial - - oldTShift.append(state.tShift[nextLine]) - state.tShift[nextLine] = pos - state.bMarks[nextLine] - - nextLine += 1 - continue - - # Case 2: line is not inside the blockquote, and the last line was empty. - if lastLineEmpty: - break - - # Case 3: another tag found. - terminate = False - - for terminatorRule in terminatorRules: - if terminatorRule(state, nextLine, endLine, True): - terminate = True - break - - if terminate: - # Quirk to enforce "hard termination mode" for paragraphs; - # normally if you call `tokenize(state, startLine, nextLine)`, - # paragraphs will look below nextLine for paragraph continuation, - # but if blockquote is terminated by another tag, they shouldn't - state.lineMax = nextLine - - if state.blkIndent != 0: - # state.blkIndent was non-zero, we now set it to zero, - # so we need to re-calculate all offsets to appear as - # if indent wasn't changed - oldBMarks.append(state.bMarks[nextLine]) - oldBSCount.append(state.bsCount[nextLine]) - oldTShift.append(state.tShift[nextLine]) - oldSCount.append(state.sCount[nextLine]) - state.sCount[nextLine] -= state.blkIndent - - break - - oldBMarks.append(state.bMarks[nextLine]) - oldBSCount.append(state.bsCount[nextLine]) - oldTShift.append(state.tShift[nextLine]) - oldSCount.append(state.sCount[nextLine]) - - # A negative indentation means that this is a paragraph continuation - # - state.sCount[nextLine] = -1 - - nextLine += 1 - - oldIndent = state.blkIndent - state.blkIndent = 0 - - token = state.push("blockquote_open", "blockquote", 1) - token.markup = ">" - token.map = lines = [startLine, 0] - - state.md.block.tokenize(state, startLine, nextLine) - - token = state.push("blockquote_close", "blockquote", -1) - token.markup = ">" - - state.lineMax = oldLineMax - state.parentType = oldParentType - lines[1] = state.line - - # Restore original tShift; this might not be necessary since the parser - # has already been here, but just to make sure we can do that. - for i, item in enumerate(oldTShift): - state.bMarks[i + startLine] = oldBMarks[i] - state.tShift[i + startLine] = item - state.sCount[i + startLine] = oldSCount[i] - state.bsCount[i + startLine] = oldBSCount[i] - - state.blkIndent = oldIndent - - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/code.py b/venv/lib/python3.10/site-packages/markdown_it/rules_block/code.py deleted file mode 100644 index 89db9cec4e0fda2798446940397a965a2083f1a3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_block/code.py +++ /dev/null @@ -1,35 +0,0 @@ -"""Code block (4 spaces padded).""" -import logging - -from .state_block import StateBlock - -LOGGER = logging.getLogger(__name__) - - -def code(state: StateBlock, startLine: int, endLine: int, silent: bool) -> bool: - LOGGER.debug("entering code: %s, %s, %s, %s", state, startLine, endLine, silent) - - if not state.is_code_block(startLine): - return False - - last = nextLine = startLine + 1 - - while nextLine < endLine: - if state.isEmpty(nextLine): - nextLine += 1 - continue - - if state.is_code_block(nextLine): - nextLine += 1 - last = nextLine - continue - - break - - state.line = last - - token = state.push("code_block", "code", 0) - token.content = state.getLines(startLine, last, 4 + state.blkIndent, False) + "\n" - token.map = [startLine, state.line] - - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/fence.py b/venv/lib/python3.10/site-packages/markdown_it/rules_block/fence.py deleted file mode 100644 index 263f1b8de8dcdd0dd736eeafab2d9da34ec2c205..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_block/fence.py +++ /dev/null @@ -1,101 +0,0 @@ -# fences (``` lang, ~~~ lang) -import logging - -from .state_block import StateBlock - -LOGGER = logging.getLogger(__name__) - - -def fence(state: StateBlock, startLine: int, endLine: int, silent: bool) -> bool: - LOGGER.debug("entering fence: %s, %s, %s, %s", state, startLine, endLine, silent) - - haveEndMarker = False - pos = state.bMarks[startLine] + state.tShift[startLine] - maximum = state.eMarks[startLine] - - if state.is_code_block(startLine): - return False - - if pos + 3 > maximum: - return False - - marker = state.src[pos] - - if marker not in ("~", "`"): - return False - - # scan marker length - mem = pos - pos = state.skipCharsStr(pos, marker) - - length = pos - mem - - if length < 3: - return False - - markup = state.src[mem:pos] - params = state.src[pos:maximum] - - if marker == "`" and marker in params: - return False - - # Since start is found, we can report success here in validation mode - if silent: - return True - - # search end of block - nextLine = startLine - - while True: - nextLine += 1 - if nextLine >= endLine: - # unclosed block should be autoclosed by end of document. - # also block seems to be autoclosed by end of parent - break - - pos = mem = state.bMarks[nextLine] + state.tShift[nextLine] - maximum = state.eMarks[nextLine] - - if pos < maximum and state.sCount[nextLine] < state.blkIndent: - # non-empty line with negative indent should stop the list: - # - ``` - # test - break - - try: - if state.src[pos] != marker: - continue - except IndexError: - break - - if state.is_code_block(nextLine): - continue - - pos = state.skipCharsStr(pos, marker) - - # closing code fence must be at least as long as the opening one - if pos - mem < length: - continue - - # make sure tail has spaces only - pos = state.skipSpaces(pos) - - if pos < maximum: - continue - - haveEndMarker = True - # found! - break - - # If a fence has heading spaces, they should be removed from its inner block - length = state.sCount[startLine] - - state.line = nextLine + (1 if haveEndMarker else 0) - - token = state.push("fence", "code", 0) - token.info = params - token.content = state.getLines(startLine + 1, nextLine, length, True) - token.markup = markup - token.map = [startLine, state.line] - - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/heading.py b/venv/lib/python3.10/site-packages/markdown_it/rules_block/heading.py deleted file mode 100644 index 850ffb504d1b1aa0198672573ec61ff3ab3d0ed7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_block/heading.py +++ /dev/null @@ -1,68 +0,0 @@ -""" Atex heading (#, ##, ...) """ -from __future__ import annotations - -import logging - -from ..common.utils import isStrSpace -from .state_block import StateBlock - -LOGGER = logging.getLogger(__name__) - - -def heading(state: StateBlock, startLine: int, endLine: int, silent: bool) -> bool: - LOGGER.debug("entering heading: %s, %s, %s, %s", state, startLine, endLine, silent) - - pos = state.bMarks[startLine] + state.tShift[startLine] - maximum = state.eMarks[startLine] - - if state.is_code_block(startLine): - return False - - ch: str | None = state.src[pos] - - if ch != "#" or pos >= maximum: - return False - - # count heading level - level = 1 - pos += 1 - try: - ch = state.src[pos] - except IndexError: - ch = None - while ch == "#" and pos < maximum and level <= 6: - level += 1 - pos += 1 - try: - ch = state.src[pos] - except IndexError: - ch = None - - if level > 6 or (pos < maximum and not isStrSpace(ch)): - return False - - if silent: - return True - - # Let's cut tails like ' ### ' from the end of string - - maximum = state.skipSpacesBack(maximum, pos) - tmp = state.skipCharsStrBack(maximum, "#", pos) - if tmp > pos and isStrSpace(state.src[tmp - 1]): - maximum = tmp - - state.line = startLine + 1 - - token = state.push("heading_open", "h" + str(level), 1) - token.markup = "########"[:level] - token.map = [startLine, state.line] - - token = state.push("inline", "", 0) - token.content = state.src[pos:maximum].strip() - token.map = [startLine, state.line] - token.children = [] - - token = state.push("heading_close", "h" + str(level), -1) - token.markup = "########"[:level] - - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/hr.py b/venv/lib/python3.10/site-packages/markdown_it/rules_block/hr.py deleted file mode 100644 index 16df05f2ab2e9e171b4f3f5612bba62aab78b2a3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_block/hr.py +++ /dev/null @@ -1,55 +0,0 @@ -"""Horizontal rule - -At least 3 of these characters on a line * - _ -""" -import logging - -from ..common.utils import isStrSpace -from .state_block import StateBlock - -LOGGER = logging.getLogger(__name__) - - -def hr(state: StateBlock, startLine: int, endLine: int, silent: bool) -> bool: - LOGGER.debug("entering hr: %s, %s, %s, %s", state, startLine, endLine, silent) - - pos = state.bMarks[startLine] + state.tShift[startLine] - maximum = state.eMarks[startLine] - - if state.is_code_block(startLine): - return False - - try: - marker = state.src[pos] - except IndexError: - return False - pos += 1 - - # Check hr marker - if marker not in ("*", "-", "_"): - return False - - # markers can be mixed with spaces, but there should be at least 3 of them - - cnt = 1 - while pos < maximum: - ch = state.src[pos] - pos += 1 - if ch != marker and not isStrSpace(ch): - return False - if ch == marker: - cnt += 1 - - if cnt < 3: - return False - - if silent: - return True - - state.line = startLine + 1 - - token = state.push("hr", "hr", 0) - token.map = [startLine, state.line] - token.markup = marker * (cnt + 1) - - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/html_block.py b/venv/lib/python3.10/site-packages/markdown_it/rules_block/html_block.py deleted file mode 100644 index 3d43f6ee1deb527a42f4d99da40bd052d9b02886..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_block/html_block.py +++ /dev/null @@ -1,90 +0,0 @@ -# HTML block -from __future__ import annotations - -import logging -import re - -from ..common.html_blocks import block_names -from ..common.html_re import HTML_OPEN_CLOSE_TAG_STR -from .state_block import StateBlock - -LOGGER = logging.getLogger(__name__) - -# An array of opening and corresponding closing sequences for html tags, -# last argument defines whether it can terminate a paragraph or not -HTML_SEQUENCES: list[tuple[re.Pattern[str], re.Pattern[str], bool]] = [ - ( - re.compile(r"^<(script|pre|style|textarea)(?=(\s|>|$))", re.IGNORECASE), - re.compile(r"<\/(script|pre|style|textarea)>", re.IGNORECASE), - True, - ), - (re.compile(r"^"), True), - (re.compile(r"^<\?"), re.compile(r"\?>"), True), - (re.compile(r"^"), True), - (re.compile(r"^"), True), - ( - re.compile("^|$))", re.IGNORECASE), - re.compile(r"^$"), - True, - ), - (re.compile(HTML_OPEN_CLOSE_TAG_STR + "\\s*$"), re.compile(r"^$"), False), -] - - -def html_block(state: StateBlock, startLine: int, endLine: int, silent: bool) -> bool: - LOGGER.debug( - "entering html_block: %s, %s, %s, %s", state, startLine, endLine, silent - ) - pos = state.bMarks[startLine] + state.tShift[startLine] - maximum = state.eMarks[startLine] - - if state.is_code_block(startLine): - return False - - if not state.md.options.get("html", None): - return False - - if state.src[pos] != "<": - return False - - lineText = state.src[pos:maximum] - - html_seq = None - for HTML_SEQUENCE in HTML_SEQUENCES: - if HTML_SEQUENCE[0].search(lineText): - html_seq = HTML_SEQUENCE - break - - if not html_seq: - return False - - if silent: - # true if this sequence can be a terminator, false otherwise - return html_seq[2] - - nextLine = startLine + 1 - - # If we are here - we detected HTML block. - # Let's roll down till block end. - if not html_seq[1].search(lineText): - while nextLine < endLine: - if state.sCount[nextLine] < state.blkIndent: - break - - pos = state.bMarks[nextLine] + state.tShift[nextLine] - maximum = state.eMarks[nextLine] - lineText = state.src[pos:maximum] - - if html_seq[1].search(lineText): - if len(lineText) != 0: - nextLine += 1 - break - nextLine += 1 - - state.line = nextLine - - token = state.push("html_block", "", 0) - token.map = [startLine, nextLine] - token.content = state.getLines(startLine, nextLine, state.blkIndent, True) - - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/lheading.py b/venv/lib/python3.10/site-packages/markdown_it/rules_block/lheading.py deleted file mode 100644 index 3522207abb680510decdd6c54d0be81401128ad7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_block/lheading.py +++ /dev/null @@ -1,86 +0,0 @@ -# lheading (---, ==) -import logging - -from .state_block import StateBlock - -LOGGER = logging.getLogger(__name__) - - -def lheading(state: StateBlock, startLine: int, endLine: int, silent: bool) -> bool: - LOGGER.debug("entering lheading: %s, %s, %s, %s", state, startLine, endLine, silent) - - level = None - nextLine = startLine + 1 - ruler = state.md.block.ruler - terminatorRules = ruler.getRules("paragraph") - - if state.is_code_block(startLine): - return False - - oldParentType = state.parentType - state.parentType = "paragraph" # use paragraph to match terminatorRules - - # jump line-by-line until empty one or EOF - while nextLine < endLine and not state.isEmpty(nextLine): - # this would be a code block normally, but after paragraph - # it's considered a lazy continuation regardless of what's there - if state.sCount[nextLine] - state.blkIndent > 3: - nextLine += 1 - continue - - # Check for underline in setext header - if state.sCount[nextLine] >= state.blkIndent: - pos = state.bMarks[nextLine] + state.tShift[nextLine] - maximum = state.eMarks[nextLine] - - if pos < maximum: - marker = state.src[pos] - - if marker in ("-", "="): - pos = state.skipCharsStr(pos, marker) - pos = state.skipSpaces(pos) - - # /* = */ - if pos >= maximum: - level = 1 if marker == "=" else 2 - break - - # quirk for blockquotes, this line should already be checked by that rule - if state.sCount[nextLine] < 0: - nextLine += 1 - continue - - # Some tags can terminate paragraph without empty line. - terminate = False - for terminatorRule in terminatorRules: - if terminatorRule(state, nextLine, endLine, True): - terminate = True - break - if terminate: - break - - nextLine += 1 - - if not level: - # Didn't find valid underline - return False - - content = state.getLines(startLine, nextLine, state.blkIndent, False).strip() - - state.line = nextLine + 1 - - token = state.push("heading_open", "h" + str(level), 1) - token.markup = marker - token.map = [startLine, state.line] - - token = state.push("inline", "", 0) - token.content = content - token.map = [startLine, state.line - 1] - token.children = [] - - token = state.push("heading_close", "h" + str(level), -1) - token.markup = marker - - state.parentType = oldParentType - - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/list.py b/venv/lib/python3.10/site-packages/markdown_it/rules_block/list.py deleted file mode 100644 index d8070d747035dd6b43f11c4bd88d05533b22bc5b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_block/list.py +++ /dev/null @@ -1,345 +0,0 @@ -# Lists -import logging - -from ..common.utils import isStrSpace -from .state_block import StateBlock - -LOGGER = logging.getLogger(__name__) - - -# Search `[-+*][\n ]`, returns next pos after marker on success -# or -1 on fail. -def skipBulletListMarker(state: StateBlock, startLine: int) -> int: - pos = state.bMarks[startLine] + state.tShift[startLine] - maximum = state.eMarks[startLine] - - try: - marker = state.src[pos] - except IndexError: - return -1 - pos += 1 - - if marker not in ("*", "-", "+"): - return -1 - - if pos < maximum: - ch = state.src[pos] - - if not isStrSpace(ch): - # " -test " - is not a list item - return -1 - - return pos - - -# Search `\d+[.)][\n ]`, returns next pos after marker on success -# or -1 on fail. -def skipOrderedListMarker(state: StateBlock, startLine: int) -> int: - start = state.bMarks[startLine] + state.tShift[startLine] - pos = start - maximum = state.eMarks[startLine] - - # List marker should have at least 2 chars (digit + dot) - if pos + 1 >= maximum: - return -1 - - ch = state.src[pos] - pos += 1 - - ch_ord = ord(ch) - # /* 0 */ /* 9 */ - if ch_ord < 0x30 or ch_ord > 0x39: - return -1 - - while True: - # EOL -> fail - if pos >= maximum: - return -1 - - ch = state.src[pos] - pos += 1 - - # /* 0 */ /* 9 */ - ch_ord = ord(ch) - if ch_ord >= 0x30 and ch_ord <= 0x39: - # List marker should have no more than 9 digits - # (prevents integer overflow in browsers) - if pos - start >= 10: - return -1 - - continue - - # found valid marker - if ch in (")", "."): - break - - return -1 - - if pos < maximum: - ch = state.src[pos] - - if not isStrSpace(ch): - # " 1.test " - is not a list item - return -1 - - return pos - - -def markTightParagraphs(state: StateBlock, idx: int) -> None: - level = state.level + 2 - - i = idx + 2 - length = len(state.tokens) - 2 - while i < length: - if state.tokens[i].level == level and state.tokens[i].type == "paragraph_open": - state.tokens[i + 2].hidden = True - state.tokens[i].hidden = True - i += 2 - i += 1 - - -def list_block(state: StateBlock, startLine: int, endLine: int, silent: bool) -> bool: - LOGGER.debug("entering list: %s, %s, %s, %s", state, startLine, endLine, silent) - - isTerminatingParagraph = False - tight = True - - if state.is_code_block(startLine): - return False - - # Special case: - # - item 1 - # - item 2 - # - item 3 - # - item 4 - # - this one is a paragraph continuation - if ( - state.listIndent >= 0 - and state.sCount[startLine] - state.listIndent >= 4 - and state.sCount[startLine] < state.blkIndent - ): - return False - - # limit conditions when list can interrupt - # a paragraph (validation mode only) - # Next list item should still terminate previous list item - # - # This code can fail if plugins use blkIndent as well as lists, - # but I hope the spec gets fixed long before that happens. - # - if ( - silent - and state.parentType == "paragraph" - and state.sCount[startLine] >= state.blkIndent - ): - isTerminatingParagraph = True - - # Detect list type and position after marker - posAfterMarker = skipOrderedListMarker(state, startLine) - if posAfterMarker >= 0: - isOrdered = True - start = state.bMarks[startLine] + state.tShift[startLine] - markerValue = int(state.src[start : posAfterMarker - 1]) - - # If we're starting a new ordered list right after - # a paragraph, it should start with 1. - if isTerminatingParagraph and markerValue != 1: - return False - else: - posAfterMarker = skipBulletListMarker(state, startLine) - if posAfterMarker >= 0: - isOrdered = False - else: - return False - - # If we're starting a new unordered list right after - # a paragraph, first line should not be empty. - if ( - isTerminatingParagraph - and state.skipSpaces(posAfterMarker) >= state.eMarks[startLine] - ): - return False - - # We should terminate list on style change. Remember first one to compare. - markerChar = state.src[posAfterMarker - 1] - - # For validation mode we can terminate immediately - if silent: - return True - - # Start list - listTokIdx = len(state.tokens) - - if isOrdered: - token = state.push("ordered_list_open", "ol", 1) - if markerValue != 1: - token.attrs = {"start": markerValue} - - else: - token = state.push("bullet_list_open", "ul", 1) - - token.map = listLines = [startLine, 0] - token.markup = markerChar - - # - # Iterate list items - # - - nextLine = startLine - prevEmptyEnd = False - terminatorRules = state.md.block.ruler.getRules("list") - - oldParentType = state.parentType - state.parentType = "list" - - while nextLine < endLine: - pos = posAfterMarker - maximum = state.eMarks[nextLine] - - initial = offset = ( - state.sCount[nextLine] - + posAfterMarker - - (state.bMarks[startLine] + state.tShift[startLine]) - ) - - while pos < maximum: - ch = state.src[pos] - - if ch == "\t": - offset += 4 - (offset + state.bsCount[nextLine]) % 4 - elif ch == " ": - offset += 1 - else: - break - - pos += 1 - - contentStart = pos - - # trimming space in "- \n 3" case, indent is 1 here - indentAfterMarker = 1 if contentStart >= maximum else offset - initial - - # If we have more than 4 spaces, the indent is 1 - # (the rest is just indented code block) - if indentAfterMarker > 4: - indentAfterMarker = 1 - - # " - test" - # ^^^^^ - calculating total length of this thing - indent = initial + indentAfterMarker - - # Run subparser & write tokens - token = state.push("list_item_open", "li", 1) - token.markup = markerChar - token.map = itemLines = [startLine, 0] - if isOrdered: - token.info = state.src[start : posAfterMarker - 1] - - # change current state, then restore it after parser subcall - oldTight = state.tight - oldTShift = state.tShift[startLine] - oldSCount = state.sCount[startLine] - - # - example list - # ^ listIndent position will be here - # ^ blkIndent position will be here - # - oldListIndent = state.listIndent - state.listIndent = state.blkIndent - state.blkIndent = indent - - state.tight = True - state.tShift[startLine] = contentStart - state.bMarks[startLine] - state.sCount[startLine] = offset - - if contentStart >= maximum and state.isEmpty(startLine + 1): - # workaround for this case - # (list item is empty, list terminates before "foo"): - # ~~~~~~~~ - # - - # - # foo - # ~~~~~~~~ - state.line = min(state.line + 2, endLine) - else: - # NOTE in list.js this was: - # state.md.block.tokenize(state, startLine, endLine, True) - # but tokeniz does not take the final parameter - state.md.block.tokenize(state, startLine, endLine) - - # If any of list item is tight, mark list as tight - if (not state.tight) or prevEmptyEnd: - tight = False - - # Item become loose if finish with empty line, - # but we should filter last element, because it means list finish - prevEmptyEnd = (state.line - startLine) > 1 and state.isEmpty(state.line - 1) - - state.blkIndent = state.listIndent - state.listIndent = oldListIndent - state.tShift[startLine] = oldTShift - state.sCount[startLine] = oldSCount - state.tight = oldTight - - token = state.push("list_item_close", "li", -1) - token.markup = markerChar - - nextLine = startLine = state.line - itemLines[1] = nextLine - - if nextLine >= endLine: - break - - contentStart = state.bMarks[startLine] - - # - # Try to check if list is terminated or continued. - # - if state.sCount[nextLine] < state.blkIndent: - break - - if state.is_code_block(startLine): - break - - # fail if terminating block found - terminate = False - for terminatorRule in terminatorRules: - if terminatorRule(state, nextLine, endLine, True): - terminate = True - break - - if terminate: - break - - # fail if list has another type - if isOrdered: - posAfterMarker = skipOrderedListMarker(state, nextLine) - if posAfterMarker < 0: - break - start = state.bMarks[nextLine] + state.tShift[nextLine] - else: - posAfterMarker = skipBulletListMarker(state, nextLine) - if posAfterMarker < 0: - break - - if markerChar != state.src[posAfterMarker - 1]: - break - - # Finalize list - if isOrdered: - token = state.push("ordered_list_close", "ol", -1) - else: - token = state.push("bullet_list_close", "ul", -1) - - token.markup = markerChar - - listLines[1] = nextLine - state.line = nextLine - - state.parentType = oldParentType - - # mark paragraphs tight if needed - if tight: - markTightParagraphs(state, listTokIdx) - - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/paragraph.py b/venv/lib/python3.10/site-packages/markdown_it/rules_block/paragraph.py deleted file mode 100644 index 5388a4b1468defccb5ed4d7d68f3e5ac1a747178..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_block/paragraph.py +++ /dev/null @@ -1,65 +0,0 @@ -"""Paragraph.""" -import logging - -from .state_block import StateBlock - -LOGGER = logging.getLogger(__name__) - - -def paragraph(state: StateBlock, startLine: int, endLine: int, silent: bool) -> bool: - LOGGER.debug( - "entering paragraph: %s, %s, %s, %s", state, startLine, endLine, silent - ) - - nextLine = startLine + 1 - ruler = state.md.block.ruler - terminatorRules = ruler.getRules("paragraph") - endLine = state.lineMax - - oldParentType = state.parentType - state.parentType = "paragraph" - - # jump line-by-line until empty one or EOF - while nextLine < endLine: - if state.isEmpty(nextLine): - break - # this would be a code block normally, but after paragraph - # it's considered a lazy continuation regardless of what's there - if state.sCount[nextLine] - state.blkIndent > 3: - nextLine += 1 - continue - - # quirk for blockquotes, this line should already be checked by that rule - if state.sCount[nextLine] < 0: - nextLine += 1 - continue - - # Some tags can terminate paragraph without empty line. - terminate = False - for terminatorRule in terminatorRules: - if terminatorRule(state, nextLine, endLine, True): - terminate = True - break - - if terminate: - break - - nextLine += 1 - - content = state.getLines(startLine, nextLine, state.blkIndent, False).strip() - - state.line = nextLine - - token = state.push("paragraph_open", "p", 1) - token.map = [startLine, state.line] - - token = state.push("inline", "", 0) - token.content = content - token.map = [startLine, state.line] - token.children = [] - - token = state.push("paragraph_close", "p", -1) - - state.parentType = oldParentType - - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/reference.py b/venv/lib/python3.10/site-packages/markdown_it/rules_block/reference.py deleted file mode 100644 index b77944b2ee8133727da934c3c7dbe800f2fc28ce..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_block/reference.py +++ /dev/null @@ -1,215 +0,0 @@ -import logging - -from ..common.utils import charCodeAt, isSpace, normalizeReference -from .state_block import StateBlock - -LOGGER = logging.getLogger(__name__) - - -def reference(state: StateBlock, startLine: int, _endLine: int, silent: bool) -> bool: - LOGGER.debug( - "entering reference: %s, %s, %s, %s", state, startLine, _endLine, silent - ) - - lines = 0 - pos = state.bMarks[startLine] + state.tShift[startLine] - maximum = state.eMarks[startLine] - nextLine = startLine + 1 - - if state.is_code_block(startLine): - return False - - if state.src[pos] != "[": - return False - - # Simple check to quickly interrupt scan on [link](url) at the start of line. - # Can be useful on practice: https:#github.com/markdown-it/markdown-it/issues/54 - while pos < maximum: - # /* ] */ /* \ */ /* : */ - if state.src[pos] == "]" and state.src[pos - 1] != "\\": - if pos + 1 == maximum: - return False - if state.src[pos + 1] != ":": - return False - break - pos += 1 - - endLine = state.lineMax - - # jump line-by-line until empty one or EOF - terminatorRules = state.md.block.ruler.getRules("reference") - - oldParentType = state.parentType - state.parentType = "reference" - - while nextLine < endLine and not state.isEmpty(nextLine): - # this would be a code block normally, but after paragraph - # it's considered a lazy continuation regardless of what's there - if state.sCount[nextLine] - state.blkIndent > 3: - nextLine += 1 - continue - - # quirk for blockquotes, this line should already be checked by that rule - if state.sCount[nextLine] < 0: - nextLine += 1 - continue - - # Some tags can terminate paragraph without empty line. - terminate = False - for terminatorRule in terminatorRules: - if terminatorRule(state, nextLine, endLine, True): - terminate = True - break - - if terminate: - break - - nextLine += 1 - - string = state.getLines(startLine, nextLine, state.blkIndent, False).strip() - maximum = len(string) - - labelEnd = None - pos = 1 - while pos < maximum: - ch = charCodeAt(string, pos) - if ch == 0x5B: # /* [ */ - return False - elif ch == 0x5D: # /* ] */ - labelEnd = pos - break - elif ch == 0x0A: # /* \n */ - lines += 1 - elif ch == 0x5C: # /* \ */ - pos += 1 - if pos < maximum and charCodeAt(string, pos) == 0x0A: - lines += 1 - pos += 1 - - if ( - labelEnd is None or labelEnd < 0 or charCodeAt(string, labelEnd + 1) != 0x3A - ): # /* : */ - return False - - # [label]: destination 'title' - # ^^^ skip optional whitespace here - pos = labelEnd + 2 - while pos < maximum: - ch = charCodeAt(string, pos) - if ch == 0x0A: - lines += 1 - elif isSpace(ch): - pass - else: - break - pos += 1 - - # [label]: destination 'title' - # ^^^^^^^^^^^ parse this - res = state.md.helpers.parseLinkDestination(string, pos, maximum) - if not res.ok: - return False - - href = state.md.normalizeLink(res.str) - if not state.md.validateLink(href): - return False - - pos = res.pos - lines += res.lines - - # save cursor state, we could require to rollback later - destEndPos = pos - destEndLineNo = lines - - # [label]: destination 'title' - # ^^^ skipping those spaces - start = pos - while pos < maximum: - ch = charCodeAt(string, pos) - if ch == 0x0A: - lines += 1 - elif isSpace(ch): - pass - else: - break - pos += 1 - - # [label]: destination 'title' - # ^^^^^^^ parse this - res = state.md.helpers.parseLinkTitle(string, pos, maximum) - if pos < maximum and start != pos and res.ok: - title = res.str - pos = res.pos - lines += res.lines - else: - title = "" - pos = destEndPos - lines = destEndLineNo - - # skip trailing spaces until the rest of the line - while pos < maximum: - ch = charCodeAt(string, pos) - if not isSpace(ch): - break - pos += 1 - - if pos < maximum and charCodeAt(string, pos) != 0x0A and title: - # garbage at the end of the line after title, - # but it could still be a valid reference if we roll back - title = "" - pos = destEndPos - lines = destEndLineNo - while pos < maximum: - ch = charCodeAt(string, pos) - if not isSpace(ch): - break - pos += 1 - - if pos < maximum and charCodeAt(string, pos) != 0x0A: - # garbage at the end of the line - return False - - label = normalizeReference(string[1:labelEnd]) - if not label: - # CommonMark 0.20 disallows empty labels - return False - - # Reference can not terminate anything. This check is for safety only. - if silent: - return True - - if "references" not in state.env: - state.env["references"] = {} - - state.line = startLine + lines + 1 - - # note, this is not part of markdown-it JS, but is useful for renderers - if state.md.options.get("inline_definitions", False): - token = state.push("definition", "", 0) - token.meta = { - "id": label, - "title": title, - "url": href, - "label": string[1:labelEnd], - } - token.map = [startLine, state.line] - - if label not in state.env["references"]: - state.env["references"][label] = { - "title": title, - "href": href, - "map": [startLine, state.line], - } - else: - state.env.setdefault("duplicate_refs", []).append( - { - "title": title, - "href": href, - "label": label, - "map": [startLine, state.line], - } - ) - - state.parentType = oldParentType - - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/state_block.py b/venv/lib/python3.10/site-packages/markdown_it/rules_block/state_block.py deleted file mode 100644 index 445ad265a01e3f1dededf9f72848686a2b5ee901..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_block/state_block.py +++ /dev/null @@ -1,261 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Literal - -from ..common.utils import isStrSpace -from ..ruler import StateBase -from ..token import Token -from ..utils import EnvType - -if TYPE_CHECKING: - from markdown_it.main import MarkdownIt - - -class StateBlock(StateBase): - def __init__( - self, src: str, md: MarkdownIt, env: EnvType, tokens: list[Token] - ) -> None: - self.src = src - - # link to parser instance - self.md = md - - self.env = env - - # - # Internal state variables - # - - self.tokens = tokens - - self.bMarks: list[int] = [] # line begin offsets for fast jumps - self.eMarks: list[int] = [] # line end offsets for fast jumps - # offsets of the first non-space characters (tabs not expanded) - self.tShift: list[int] = [] - self.sCount: list[int] = [] # indents for each line (tabs expanded) - - # An amount of virtual spaces (tabs expanded) between beginning - # of each line (bMarks) and real beginning of that line. - # - # It exists only as a hack because blockquotes override bMarks - # losing information in the process. - # - # It's used only when expanding tabs, you can think about it as - # an initial tab length, e.g. bsCount=21 applied to string `\t123` - # means first tab should be expanded to 4-21%4 === 3 spaces. - # - self.bsCount: list[int] = [] - - # block parser variables - self.blkIndent = 0 # required block content indent (for example, if we are - # inside a list, it would be positioned after list marker) - self.line = 0 # line index in src - self.lineMax = 0 # lines count - self.tight = False # loose/tight mode for lists - self.ddIndent = -1 # indent of the current dd block (-1 if there isn't any) - self.listIndent = -1 # indent of the current list block (-1 if there isn't any) - - # can be 'blockquote', 'list', 'root', 'paragraph' or 'reference' - # used in lists to determine if they interrupt a paragraph - self.parentType = "root" - - self.level = 0 - - # renderer - self.result = "" - - # Create caches - # Generate markers. - indent_found = False - - start = pos = indent = offset = 0 - length = len(self.src) - - for pos, character in enumerate(self.src): - if not indent_found: - if isStrSpace(character): - indent += 1 - - if character == "\t": - offset += 4 - offset % 4 - else: - offset += 1 - continue - else: - indent_found = True - - if character == "\n" or pos == length - 1: - if character != "\n": - pos += 1 - self.bMarks.append(start) - self.eMarks.append(pos) - self.tShift.append(indent) - self.sCount.append(offset) - self.bsCount.append(0) - - indent_found = False - indent = 0 - offset = 0 - start = pos + 1 - - # Push fake entry to simplify cache bounds checks - self.bMarks.append(length) - self.eMarks.append(length) - self.tShift.append(0) - self.sCount.append(0) - self.bsCount.append(0) - - self.lineMax = len(self.bMarks) - 1 # don't count last fake line - - # pre-check if code blocks are enabled, to speed up is_code_block method - self._code_enabled = "code" in self.md["block"].ruler.get_active_rules() - - def __repr__(self) -> str: - return ( - f"{self.__class__.__name__}" - f"(line={self.line},level={self.level},tokens={len(self.tokens)})" - ) - - def push(self, ttype: str, tag: str, nesting: Literal[-1, 0, 1]) -> Token: - """Push new token to "stream".""" - token = Token(ttype, tag, nesting) - token.block = True - if nesting < 0: - self.level -= 1 # closing tag - token.level = self.level - if nesting > 0: - self.level += 1 # opening tag - self.tokens.append(token) - return token - - def isEmpty(self, line: int) -> bool: - """.""" - return (self.bMarks[line] + self.tShift[line]) >= self.eMarks[line] - - def skipEmptyLines(self, from_pos: int) -> int: - """.""" - while from_pos < self.lineMax: - try: - if (self.bMarks[from_pos] + self.tShift[from_pos]) < self.eMarks[ - from_pos - ]: - break - except IndexError: - pass - from_pos += 1 - return from_pos - - def skipSpaces(self, pos: int) -> int: - """Skip spaces from given position.""" - while True: - try: - current = self.src[pos] - except IndexError: - break - if not isStrSpace(current): - break - pos += 1 - return pos - - def skipSpacesBack(self, pos: int, minimum: int) -> int: - """Skip spaces from given position in reverse.""" - if pos <= minimum: - return pos - while pos > minimum: - pos -= 1 - if not isStrSpace(self.src[pos]): - return pos + 1 - return pos - - def skipChars(self, pos: int, code: int) -> int: - """Skip character code from given position.""" - while True: - try: - current = self.srcCharCode[pos] - except IndexError: - break - if current != code: - break - pos += 1 - return pos - - def skipCharsStr(self, pos: int, ch: str) -> int: - """Skip character string from given position.""" - while True: - try: - current = self.src[pos] - except IndexError: - break - if current != ch: - break - pos += 1 - return pos - - def skipCharsBack(self, pos: int, code: int, minimum: int) -> int: - """Skip character code reverse from given position - 1.""" - if pos <= minimum: - return pos - while pos > minimum: - pos -= 1 - if code != self.srcCharCode[pos]: - return pos + 1 - return pos - - def skipCharsStrBack(self, pos: int, ch: str, minimum: int) -> int: - """Skip character string reverse from given position - 1.""" - if pos <= minimum: - return pos - while pos > minimum: - pos -= 1 - if ch != self.src[pos]: - return pos + 1 - return pos - - def getLines(self, begin: int, end: int, indent: int, keepLastLF: bool) -> str: - """Cut lines range from source.""" - line = begin - if begin >= end: - return "" - - queue = [""] * (end - begin) - - i = 1 - while line < end: - lineIndent = 0 - lineStart = first = self.bMarks[line] - last = ( - self.eMarks[line] + 1 - if line + 1 < end or keepLastLF - else self.eMarks[line] - ) - - while (first < last) and (lineIndent < indent): - ch = self.src[first] - if isStrSpace(ch): - if ch == "\t": - lineIndent += 4 - (lineIndent + self.bsCount[line]) % 4 - else: - lineIndent += 1 - elif first - lineStart < self.tShift[line]: - lineIndent += 1 - else: - break - first += 1 - - if lineIndent > indent: - # partially expanding tabs in code blocks, e.g '\t\tfoobar' - # with indent=2 becomes ' \tfoobar' - queue[i - 1] = (" " * (lineIndent - indent)) + self.src[first:last] - else: - queue[i - 1] = self.src[first:last] - - line += 1 - i += 1 - - return "".join(queue) - - def is_code_block(self, line: int) -> bool: - """Check if line is a code block, - i.e. the code block rule is enabled and text is indented by more than 3 spaces. - """ - return self._code_enabled and (self.sCount[line] - self.blkIndent) >= 4 diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_block/table.py b/venv/lib/python3.10/site-packages/markdown_it/rules_block/table.py deleted file mode 100644 index 4b666c1d5d9ede544b27c38d7ffb5d850edac70b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_block/table.py +++ /dev/null @@ -1,236 +0,0 @@ -# GFM table, https://github.github.com/gfm/#tables-extension- -from __future__ import annotations - -import re - -from ..common.utils import charStrAt, isStrSpace -from .state_block import StateBlock - -headerLineRe = re.compile(r"^:?-+:?$") -enclosingPipesRe = re.compile(r"^\||\|$") - - -def getLine(state: StateBlock, line: int) -> str: - pos = state.bMarks[line] + state.tShift[line] - maximum = state.eMarks[line] - - # return state.src.substr(pos, max - pos) - return state.src[pos:maximum] - - -def escapedSplit(string: str) -> list[str]: - result: list[str] = [] - pos = 0 - max = len(string) - isEscaped = False - lastPos = 0 - current = "" - ch = charStrAt(string, pos) - - while pos < max: - if ch == "|": - if not isEscaped: - # pipe separating cells, '|' - result.append(current + string[lastPos:pos]) - current = "" - lastPos = pos + 1 - else: - # escaped pipe, '\|' - current += string[lastPos : pos - 1] - lastPos = pos - - isEscaped = ch == "\\" - pos += 1 - - ch = charStrAt(string, pos) - - result.append(current + string[lastPos:]) - - return result - - -def table(state: StateBlock, startLine: int, endLine: int, silent: bool) -> bool: - tbodyLines = None - - # should have at least two lines - if startLine + 2 > endLine: - return False - - nextLine = startLine + 1 - - if state.sCount[nextLine] < state.blkIndent: - return False - - if state.is_code_block(nextLine): - return False - - # first character of the second line should be '|', '-', ':', - # and no other characters are allowed but spaces; - # basically, this is the equivalent of /^[-:|][-:|\s]*$/ regexp - - pos = state.bMarks[nextLine] + state.tShift[nextLine] - if pos >= state.eMarks[nextLine]: - return False - first_ch = state.src[pos] - pos += 1 - if first_ch not in ("|", "-", ":"): - return False - - if pos >= state.eMarks[nextLine]: - return False - second_ch = state.src[pos] - pos += 1 - if second_ch not in ("|", "-", ":") and not isStrSpace(second_ch): - return False - - # if first character is '-', then second character must not be a space - # (due to parsing ambiguity with list) - if first_ch == "-" and isStrSpace(second_ch): - return False - - while pos < state.eMarks[nextLine]: - ch = state.src[pos] - - if ch not in ("|", "-", ":") and not isStrSpace(ch): - return False - - pos += 1 - - lineText = getLine(state, startLine + 1) - - columns = lineText.split("|") - aligns = [] - for i in range(len(columns)): - t = columns[i].strip() - if not t: - # allow empty columns before and after table, but not in between columns; - # e.g. allow ` |---| `, disallow ` ---||--- ` - if i == 0 or i == len(columns) - 1: - continue - else: - return False - - if not headerLineRe.search(t): - return False - if charStrAt(t, len(t) - 1) == ":": - aligns.append("center" if charStrAt(t, 0) == ":" else "right") - elif charStrAt(t, 0) == ":": - aligns.append("left") - else: - aligns.append("") - - lineText = getLine(state, startLine).strip() - if "|" not in lineText: - return False - if state.is_code_block(startLine): - return False - columns = escapedSplit(lineText) - if columns and columns[0] == "": - columns.pop(0) - if columns and columns[-1] == "": - columns.pop() - - # header row will define an amount of columns in the entire table, - # and align row should be exactly the same (the rest of the rows can differ) - columnCount = len(columns) - if columnCount == 0 or columnCount != len(aligns): - return False - - if silent: - return True - - oldParentType = state.parentType - state.parentType = "table" - - # use 'blockquote' lists for termination because it's - # the most similar to tables - terminatorRules = state.md.block.ruler.getRules("blockquote") - - token = state.push("table_open", "table", 1) - token.map = tableLines = [startLine, 0] - - token = state.push("thead_open", "thead", 1) - token.map = [startLine, startLine + 1] - - token = state.push("tr_open", "tr", 1) - token.map = [startLine, startLine + 1] - - for i in range(len(columns)): - token = state.push("th_open", "th", 1) - if aligns[i]: - token.attrs = {"style": "text-align:" + aligns[i]} - - token = state.push("inline", "", 0) - # note in markdown-it this map was removed in v12.0.0 however, we keep it, - # since it is helpful to propagate to children tokens - token.map = [startLine, startLine + 1] - token.content = columns[i].strip() - token.children = [] - - token = state.push("th_close", "th", -1) - - token = state.push("tr_close", "tr", -1) - token = state.push("thead_close", "thead", -1) - - nextLine = startLine + 2 - while nextLine < endLine: - if state.sCount[nextLine] < state.blkIndent: - break - - terminate = False - for i in range(len(terminatorRules)): - if terminatorRules[i](state, nextLine, endLine, True): - terminate = True - break - - if terminate: - break - lineText = getLine(state, nextLine).strip() - if not lineText: - break - if state.is_code_block(nextLine): - break - columns = escapedSplit(lineText) - if columns and columns[0] == "": - columns.pop(0) - if columns and columns[-1] == "": - columns.pop() - - if nextLine == startLine + 2: - token = state.push("tbody_open", "tbody", 1) - token.map = tbodyLines = [startLine + 2, 0] - - token = state.push("tr_open", "tr", 1) - token.map = [nextLine, nextLine + 1] - - for i in range(columnCount): - token = state.push("td_open", "td", 1) - if aligns[i]: - token.attrs = {"style": "text-align:" + aligns[i]} - - token = state.push("inline", "", 0) - # note in markdown-it this map was removed in v12.0.0 however, we keep it, - # since it is helpful to propagate to children tokens - token.map = [nextLine, nextLine + 1] - try: - token.content = columns[i].strip() if columns[i] else "" - except IndexError: - token.content = "" - token.children = [] - - token = state.push("td_close", "td", -1) - - token = state.push("tr_close", "tr", -1) - - nextLine += 1 - - if tbodyLines: - token = state.push("tbody_close", "tbody", -1) - tbodyLines[1] = nextLine - - token = state.push("table_close", "table", -1) - - tableLines[1] = nextLine - state.parentType = oldParentType - state.line = nextLine - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__init__.py b/venv/lib/python3.10/site-packages/markdown_it/rules_core/__init__.py deleted file mode 100644 index c9c5368c2b694231000626a03594ebad75fe8c71..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -__all__ = ( - "StateCore", - "normalize", - "block", - "inline", - "replace", - "smartquotes", - "linkify", - "text_join", -) - -from .block import block -from .inline import inline -from .linkify import linkify -from .normalize import normalize -from .replacements import replace -from .smartquotes import smartquotes -from .state_core import StateCore -from .text_join import text_join diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 58bbbdf77d752cb7099d1336ce08260e9cea7b6b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/block.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/block.cpython-310.pyc deleted file mode 100644 index e9d9551ac9483461b5de6d740b015389b66792e7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/block.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/inline.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/inline.cpython-310.pyc deleted file mode 100644 index 15529c39213d4da75367b7ff0181f0d8c23707a3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/inline.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/linkify.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/linkify.cpython-310.pyc deleted file mode 100644 index 797deaf5e6621e49bf069af25ea8483f3c6098c7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/linkify.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/normalize.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/normalize.cpython-310.pyc deleted file mode 100644 index 275a8e43ce4ee2d6fca177b8217ccbaf0e45ba29..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/normalize.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/replacements.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/replacements.cpython-310.pyc deleted file mode 100644 index 942eaf81d4c343a9d1149161887268227bab4f6a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/replacements.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/smartquotes.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/smartquotes.cpython-310.pyc deleted file mode 100644 index 56a6d38559e2f1040dce2ae3fde74191f596628d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/smartquotes.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/state_core.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/state_core.cpython-310.pyc deleted file mode 100644 index 9c58e94e65b840301f1eba5fbbbbd97a3d44dd58..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/state_core.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/text_join.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/text_join.cpython-310.pyc deleted file mode 100644 index 7cc85c425d2d6756e47da502ad00733a735e7c6e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_core/__pycache__/text_join.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/block.py b/venv/lib/python3.10/site-packages/markdown_it/rules_core/block.py deleted file mode 100644 index a6c3bb8d7ae18880fd638690fb5b09beb78b103c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_core/block.py +++ /dev/null @@ -1,13 +0,0 @@ -from ..token import Token -from .state_core import StateCore - - -def block(state: StateCore) -> None: - if state.inlineMode: - token = Token("inline", "", 0) - token.content = state.src - token.map = [0, 1] - token.children = [] - state.tokens.append(token) - else: - state.md.block.parse(state.src, state.md, state.env, state.tokens) diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/inline.py b/venv/lib/python3.10/site-packages/markdown_it/rules_core/inline.py deleted file mode 100644 index c3fd0b5e25dda5d8a5a644cc9e460d0f92ae2d1d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_core/inline.py +++ /dev/null @@ -1,10 +0,0 @@ -from .state_core import StateCore - - -def inline(state: StateCore) -> None: - """Parse inlines""" - for token in state.tokens: - if token.type == "inline": - if token.children is None: - token.children = [] - state.md.inline.parse(token.content, state.md, state.env, token.children) diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/linkify.py b/venv/lib/python3.10/site-packages/markdown_it/rules_core/linkify.py deleted file mode 100644 index efbc9d4c9b1cbada1c936401b3421d73fbff5b64..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_core/linkify.py +++ /dev/null @@ -1,149 +0,0 @@ -from __future__ import annotations - -import re -from typing import Protocol - -from ..common.utils import arrayReplaceAt, isLinkClose, isLinkOpen -from ..token import Token -from .state_core import StateCore - -HTTP_RE = re.compile(r"^http://") -MAILTO_RE = re.compile(r"^mailto:") -TEST_MAILTO_RE = re.compile(r"^mailto:", flags=re.IGNORECASE) - - -def linkify(state: StateCore) -> None: - """Rule for identifying plain-text links.""" - if not state.md.options.linkify: - return - - if not state.md.linkify: - raise ModuleNotFoundError("Linkify enabled but not installed.") - - for inline_token in state.tokens: - if inline_token.type != "inline" or not state.md.linkify.pretest( - inline_token.content - ): - continue - - tokens = inline_token.children - - htmlLinkLevel = 0 - - # We scan from the end, to keep position when new tags added. - # Use reversed logic in links start/end match - assert tokens is not None - i = len(tokens) - while i >= 1: - i -= 1 - assert isinstance(tokens, list) - currentToken = tokens[i] - - # Skip content of markdown links - if currentToken.type == "link_close": - i -= 1 - while ( - tokens[i].level != currentToken.level - and tokens[i].type != "link_open" - ): - i -= 1 - continue - - # Skip content of html tag links - if currentToken.type == "html_inline": - if isLinkOpen(currentToken.content) and htmlLinkLevel > 0: - htmlLinkLevel -= 1 - if isLinkClose(currentToken.content): - htmlLinkLevel += 1 - if htmlLinkLevel > 0: - continue - - if currentToken.type == "text" and state.md.linkify.test( - currentToken.content - ): - text = currentToken.content - links: list[_LinkType] = state.md.linkify.match(text) or [] - - # Now split string to nodes - nodes = [] - level = currentToken.level - lastPos = 0 - - # forbid escape sequence at the start of the string, - # this avoids http\://example.com/ from being linkified as - # http://example.com/ - if ( - links - and links[0].index == 0 - and i > 0 - and tokens[i - 1].type == "text_special" - ): - links = links[1:] - - for link in links: - url = link.url - fullUrl = state.md.normalizeLink(url) - if not state.md.validateLink(fullUrl): - continue - - urlText = link.text - - # Linkifier might send raw hostnames like "example.com", where url - # starts with domain name. So we prepend http:// in those cases, - # and remove it afterwards. - if not link.schema: - urlText = HTTP_RE.sub( - "", state.md.normalizeLinkText("http://" + urlText) - ) - elif link.schema == "mailto:" and TEST_MAILTO_RE.search(urlText): - urlText = MAILTO_RE.sub( - "", state.md.normalizeLinkText("mailto:" + urlText) - ) - else: - urlText = state.md.normalizeLinkText(urlText) - - pos = link.index - - if pos > lastPos: - token = Token("text", "", 0) - token.content = text[lastPos:pos] - token.level = level - nodes.append(token) - - token = Token("link_open", "a", 1) - token.attrs = {"href": fullUrl} - token.level = level - level += 1 - token.markup = "linkify" - token.info = "auto" - nodes.append(token) - - token = Token("text", "", 0) - token.content = urlText - token.level = level - nodes.append(token) - - token = Token("link_close", "a", -1) - level -= 1 - token.level = level - token.markup = "linkify" - token.info = "auto" - nodes.append(token) - - lastPos = link.last_index - - if lastPos < len(text): - token = Token("text", "", 0) - token.content = text[lastPos:] - token.level = level - nodes.append(token) - - inline_token.children = tokens = arrayReplaceAt(tokens, i, nodes) - - -class _LinkType(Protocol): - url: str - text: str - index: int - last_index: int - schema: str | None diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/normalize.py b/venv/lib/python3.10/site-packages/markdown_it/rules_core/normalize.py deleted file mode 100644 index c9f8d0d5729b2497b5f4b611b0451dfe92872506..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_core/normalize.py +++ /dev/null @@ -1,18 +0,0 @@ -"""Normalize input string.""" -import re - -from .state_core import StateCore - -# https://spec.commonmark.org/0.29/#line-ending -NEWLINES_RE = re.compile(r"\r\n?|\n") -NULL_RE = re.compile(r"\0") - - -def normalize(state: StateCore) -> None: - # Normalize newlines - string = NEWLINES_RE.sub("\n", state.src) - - # Replace NULL characters - string = NULL_RE.sub("\uFFFD", string) - - state.src = string diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/replacements.py b/venv/lib/python3.10/site-packages/markdown_it/rules_core/replacements.py deleted file mode 100644 index 14912e17ac8eed885a2fd07c74141804f3f9fa72..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_core/replacements.py +++ /dev/null @@ -1,126 +0,0 @@ -"""Simple typographic replacements - -* ``(c)``, ``(C)`` → © -* ``(tm)``, ``(TM)`` → ™ -* ``(r)``, ``(R)`` → ® -* ``+-`` → ± -* ``...`` → … -* ``?....`` → ?.. -* ``!....`` → !.. -* ``????????`` → ??? -* ``!!!!!`` → !!! -* ``,,,`` → , -* ``--`` → &ndash -* ``---`` → &mdash -""" -from __future__ import annotations - -import logging -import re - -from ..token import Token -from .state_core import StateCore - -LOGGER = logging.getLogger(__name__) - -# TODO: -# - fractionals 1/2, 1/4, 3/4 -> ½, ¼, ¾ -# - multiplication 2 x 4 -> 2 × 4 - -RARE_RE = re.compile(r"\+-|\.\.|\?\?\?\?|!!!!|,,|--") - -# Workaround for phantomjs - need regex without /g flag, -# or root check will fail every second time -# SCOPED_ABBR_TEST_RE = r"\((c|tm|r)\)" - -SCOPED_ABBR_RE = re.compile(r"\((c|tm|r)\)", flags=re.IGNORECASE) - -PLUS_MINUS_RE = re.compile(r"\+-") - -ELLIPSIS_RE = re.compile(r"\.{2,}") - -ELLIPSIS_QUESTION_EXCLAMATION_RE = re.compile(r"([?!])…") - -QUESTION_EXCLAMATION_RE = re.compile(r"([?!]){4,}") - -COMMA_RE = re.compile(r",{2,}") - -EM_DASH_RE = re.compile(r"(^|[^-])---(?=[^-]|$)", flags=re.MULTILINE) - -EN_DASH_RE = re.compile(r"(^|\s)--(?=\s|$)", flags=re.MULTILINE) - -EN_DASH_INDENT_RE = re.compile(r"(^|[^-\s])--(?=[^-\s]|$)", flags=re.MULTILINE) - - -SCOPED_ABBR = {"c": "©", "r": "®", "tm": "™"} - - -def replaceFn(match: re.Match[str]) -> str: - return SCOPED_ABBR[match.group(1).lower()] - - -def replace_scoped(inlineTokens: list[Token]) -> None: - inside_autolink = 0 - - for token in inlineTokens: - if token.type == "text" and not inside_autolink: - token.content = SCOPED_ABBR_RE.sub(replaceFn, token.content) - - if token.type == "link_open" and token.info == "auto": - inside_autolink -= 1 - - if token.type == "link_close" and token.info == "auto": - inside_autolink += 1 - - -def replace_rare(inlineTokens: list[Token]) -> None: - inside_autolink = 0 - - for token in inlineTokens: - if ( - token.type == "text" - and (not inside_autolink) - and RARE_RE.search(token.content) - ): - # +- -> ± - token.content = PLUS_MINUS_RE.sub("±", token.content) - - # .., ..., ....... -> … - token.content = ELLIPSIS_RE.sub("…", token.content) - - # but ?..... & !..... -> ?.. & !.. - token.content = ELLIPSIS_QUESTION_EXCLAMATION_RE.sub("\\1..", token.content) - token.content = QUESTION_EXCLAMATION_RE.sub("\\1\\1\\1", token.content) - - # ,, ,,, ,,,, -> , - token.content = COMMA_RE.sub(",", token.content) - - # em-dash - token.content = EM_DASH_RE.sub("\\1\u2014", token.content) - - # en-dash - token.content = EN_DASH_RE.sub("\\1\u2013", token.content) - token.content = EN_DASH_INDENT_RE.sub("\\1\u2013", token.content) - - if token.type == "link_open" and token.info == "auto": - inside_autolink -= 1 - - if token.type == "link_close" and token.info == "auto": - inside_autolink += 1 - - -def replace(state: StateCore) -> None: - if not state.md.options.typographer: - return - - for token in state.tokens: - if token.type != "inline": - continue - if token.children is None: - continue - - if SCOPED_ABBR_RE.search(token.content): - replace_scoped(token.children) - - if RARE_RE.search(token.content): - replace_rare(token.children) diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/smartquotes.py b/venv/lib/python3.10/site-packages/markdown_it/rules_core/smartquotes.py deleted file mode 100644 index c98fbd71e7d2e644ca7c6ac95827962342326059..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_core/smartquotes.py +++ /dev/null @@ -1,202 +0,0 @@ -"""Convert straight quotation marks to typographic ones -""" -from __future__ import annotations - -import re -from typing import Any - -from ..common.utils import charCodeAt, isMdAsciiPunct, isPunctChar, isWhiteSpace -from ..token import Token -from .state_core import StateCore - -QUOTE_TEST_RE = re.compile(r"['\"]") -QUOTE_RE = re.compile(r"['\"]") -APOSTROPHE = "\u2019" # ’ - - -def replaceAt(string: str, index: int, ch: str) -> str: - # When the index is negative, the behavior is different from the js version. - # But basically, the index will not be negative. - assert index >= 0 - return string[:index] + ch + string[index + 1 :] - - -def process_inlines(tokens: list[Token], state: StateCore) -> None: - stack: list[dict[str, Any]] = [] - - for i, token in enumerate(tokens): - thisLevel = token.level - - j = 0 - for j in range(len(stack))[::-1]: - if stack[j]["level"] <= thisLevel: - break - else: - # When the loop is terminated without a "break". - # Subtract 1 to get the same index as the js version. - j -= 1 - - stack = stack[: j + 1] - - if token.type != "text": - continue - - text = token.content - pos = 0 - maximum = len(text) - - while pos < maximum: - goto_outer = False - lastIndex = pos - t = QUOTE_RE.search(text[lastIndex:]) - if not t: - break - - canOpen = canClose = True - pos = t.start(0) + lastIndex + 1 - isSingle = t.group(0) == "'" - - # Find previous character, - # default to space if it's the beginning of the line - lastChar: None | int = 0x20 - - if t.start(0) + lastIndex - 1 >= 0: - lastChar = charCodeAt(text, t.start(0) + lastIndex - 1) - else: - for j in range(i)[::-1]: - if tokens[j].type == "softbreak" or tokens[j].type == "hardbreak": - break - # should skip all tokens except 'text', 'html_inline' or 'code_inline' - if not tokens[j].content: - continue - - lastChar = charCodeAt(tokens[j].content, len(tokens[j].content) - 1) - break - - # Find next character, - # default to space if it's the end of the line - nextChar: None | int = 0x20 - - if pos < maximum: - nextChar = charCodeAt(text, pos) - else: - for j in range(i + 1, len(tokens)): - # nextChar defaults to 0x20 - if tokens[j].type == "softbreak" or tokens[j].type == "hardbreak": - break - # should skip all tokens except 'text', 'html_inline' or 'code_inline' - if not tokens[j].content: - continue - - nextChar = charCodeAt(tokens[j].content, 0) - break - - isLastPunctChar = lastChar is not None and ( - isMdAsciiPunct(lastChar) or isPunctChar(chr(lastChar)) - ) - isNextPunctChar = nextChar is not None and ( - isMdAsciiPunct(nextChar) or isPunctChar(chr(nextChar)) - ) - - isLastWhiteSpace = lastChar is not None and isWhiteSpace(lastChar) - isNextWhiteSpace = nextChar is not None and isWhiteSpace(nextChar) - - if isNextWhiteSpace: # noqa: SIM114 - canOpen = False - elif isNextPunctChar and not (isLastWhiteSpace or isLastPunctChar): - canOpen = False - - if isLastWhiteSpace: # noqa: SIM114 - canClose = False - elif isLastPunctChar and not (isNextWhiteSpace or isNextPunctChar): - canClose = False - - if nextChar == 0x22 and t.group(0) == '"': # 0x22: " # noqa: SIM102 - if ( - lastChar is not None and lastChar >= 0x30 and lastChar <= 0x39 - ): # 0x30: 0, 0x39: 9 - # special case: 1"" - count first quote as an inch - canClose = canOpen = False - - if canOpen and canClose: - # Replace quotes in the middle of punctuation sequence, but not - # in the middle of the words, i.e.: - # - # 1. foo " bar " baz - not replaced - # 2. foo-"-bar-"-baz - replaced - # 3. foo"bar"baz - not replaced - canOpen = isLastPunctChar - canClose = isNextPunctChar - - if not canOpen and not canClose: - # middle of word - if isSingle: - token.content = replaceAt( - token.content, t.start(0) + lastIndex, APOSTROPHE - ) - continue - - if canClose: - # this could be a closing quote, rewind the stack to get a match - for j in range(len(stack))[::-1]: - item = stack[j] - if stack[j]["level"] < thisLevel: - break - if item["single"] == isSingle and stack[j]["level"] == thisLevel: - item = stack[j] - - if isSingle: - openQuote = state.md.options.quotes[2] - closeQuote = state.md.options.quotes[3] - else: - openQuote = state.md.options.quotes[0] - closeQuote = state.md.options.quotes[1] - - # replace token.content *before* tokens[item.token].content, - # because, if they are pointing at the same token, replaceAt - # could mess up indices when quote length != 1 - token.content = replaceAt( - token.content, t.start(0) + lastIndex, closeQuote - ) - tokens[item["token"]].content = replaceAt( - tokens[item["token"]].content, item["pos"], openQuote - ) - - pos += len(closeQuote) - 1 - if item["token"] == i: - pos += len(openQuote) - 1 - - text = token.content - maximum = len(text) - - stack = stack[:j] - goto_outer = True - break - if goto_outer: - goto_outer = False - continue - - if canOpen: - stack.append( - { - "token": i, - "pos": t.start(0) + lastIndex, - "single": isSingle, - "level": thisLevel, - } - ) - elif canClose and isSingle: - token.content = replaceAt( - token.content, t.start(0) + lastIndex, APOSTROPHE - ) - - -def smartquotes(state: StateCore) -> None: - if not state.md.options.typographer: - return - - for token in state.tokens: - if token.type != "inline" or not QUOTE_RE.search(token.content): - continue - if token.children is not None: - process_inlines(token.children, state) diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/state_core.py b/venv/lib/python3.10/site-packages/markdown_it/rules_core/state_core.py deleted file mode 100644 index a938041d992fdf7ae3f2843a2e0f9ef298c45790..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_core/state_core.py +++ /dev/null @@ -1,25 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING - -from ..ruler import StateBase -from ..token import Token -from ..utils import EnvType - -if TYPE_CHECKING: - from markdown_it import MarkdownIt - - -class StateCore(StateBase): - def __init__( - self, - src: str, - md: MarkdownIt, - env: EnvType, - tokens: list[Token] | None = None, - ) -> None: - self.src = src - self.md = md # link to parser instance - self.env = env - self.tokens: list[Token] = tokens or [] - self.inlineMode = False diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_core/text_join.py b/venv/lib/python3.10/site-packages/markdown_it/rules_core/text_join.py deleted file mode 100644 index d54ccbbc376e7c50cf95227a36a11000b9d80496..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_core/text_join.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Join raw text tokens with the rest of the text - -This is set as a separate rule to provide an opportunity for plugins -to run text replacements after text join, but before escape join. - -For example, `\\:)` shouldn't be replaced with an emoji. -""" -from __future__ import annotations - -from ..token import Token -from .state_core import StateCore - - -def text_join(state: StateCore) -> None: - """Join raw text for escape sequences (`text_special`) tokens with the rest of the text""" - - for inline_token in state.tokens[:]: - if inline_token.type != "inline": - continue - - # convert text_special to text and join all adjacent text nodes - new_tokens: list[Token] = [] - for child_token in inline_token.children or []: - if child_token.type == "text_special": - child_token.type = "text" - if ( - child_token.type == "text" - and new_tokens - and new_tokens[-1].type == "text" - ): - new_tokens[-1].content += child_token.content - else: - new_tokens.append(child_token) - inline_token.children = new_tokens diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__init__.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__init__.py deleted file mode 100644 index 3a8026ec3a7afd6f35140bb074780da05c6e86af..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__init__.py +++ /dev/null @@ -1,31 +0,0 @@ -__all__ = ( - "StateInline", - "text", - "fragments_join", - "link_pairs", - "linkify", - "escape", - "newline", - "backtick", - "emphasis", - "image", - "link", - "autolink", - "entity", - "html_inline", - "strikethrough", -) -from . import emphasis, strikethrough -from .autolink import autolink -from .backticks import backtick -from .balance_pairs import link_pairs -from .entity import entity -from .escape import escape -from .fragments_join import fragments_join -from .html_inline import html_inline -from .image import image -from .link import link -from .linkify import linkify -from .newline import newline -from .state_inline import StateInline -from .text import text diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 64c3999fb1b6b957cc88fc04bcb77e5328886109..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/autolink.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/autolink.cpython-310.pyc deleted file mode 100644 index 4c0d40cbe8ae37aa46f0add47bb1021254b1d993..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/autolink.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/backticks.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/backticks.cpython-310.pyc deleted file mode 100644 index fa819b579dd99333a40e7f3195226617dfefa613..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/backticks.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/balance_pairs.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/balance_pairs.cpython-310.pyc deleted file mode 100644 index fe28c6ffd73dde485ad078ed153ca8d9ec926b65..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/balance_pairs.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/emphasis.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/emphasis.cpython-310.pyc deleted file mode 100644 index a6ecf8f98961ef6db02efc63e7167718c1a3b0fa..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/emphasis.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/entity.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/entity.cpython-310.pyc deleted file mode 100644 index 2f9999e099368ed38c4837a38a5d38c3667a06ee..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/entity.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/escape.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/escape.cpython-310.pyc deleted file mode 100644 index 5cf1adb7dd01bbf59a690e53247b0c33e46c141d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/escape.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/fragments_join.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/fragments_join.cpython-310.pyc deleted file mode 100644 index 992287c22f5c491c29d818c9a0ef24b8e76dae3d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/fragments_join.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/html_inline.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/html_inline.cpython-310.pyc deleted file mode 100644 index e6d1edbe1af5e03a94487e1ecd0de1e0b0923120..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/html_inline.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/image.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/image.cpython-310.pyc deleted file mode 100644 index 01f63753e8a2b9c71be830476b7d684bb9e04f66..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/image.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/link.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/link.cpython-310.pyc deleted file mode 100644 index 69a01fbcd56af2dc5ba0914517fd3ef94dc0ecd1..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/link.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/linkify.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/linkify.cpython-310.pyc deleted file mode 100644 index ca0e7d1ae24c7b40ea94acb9cbc54c2ba3715017..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/linkify.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/newline.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/newline.cpython-310.pyc deleted file mode 100644 index b8bdb31a884c624cff3f5e0c9f4d4521b7d829e8..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/newline.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/state_inline.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/state_inline.cpython-310.pyc deleted file mode 100644 index 92e577d6fd91691c6936014a6f3c8c8f50dea13f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/state_inline.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/strikethrough.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/strikethrough.cpython-310.pyc deleted file mode 100644 index 07b982273c13519489867f904539526f534b5c07..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/strikethrough.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/text.cpython-310.pyc b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/text.cpython-310.pyc deleted file mode 100644 index d8e1dfd6148279dbc1f8c18727ad5cbfebaec62e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/__pycache__/text.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/autolink.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/autolink.py deleted file mode 100644 index 295d963f39254e6ddfe9dc36af2bfa5e534c0827..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/autolink.py +++ /dev/null @@ -1,77 +0,0 @@ -# Process autolinks '' -import re - -from .state_inline import StateInline - -EMAIL_RE = re.compile( - r"^([a-zA-Z0-9.!#$%&\'*+\/=?^_`{|}~-]+@[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?(?:\.[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?)*)$" # noqa: E501 -) -AUTOLINK_RE = re.compile(r"^([a-zA-Z][a-zA-Z0-9+.\-]{1,31}):([^<>\x00-\x20]*)$") - - -def autolink(state: StateInline, silent: bool) -> bool: - pos = state.pos - - if state.src[pos] != "<": - return False - - start = state.pos - maximum = state.posMax - - while True: - pos += 1 - if pos >= maximum: - return False - - ch = state.src[pos] - - if ch == "<": - return False - if ch == ">": - break - - url = state.src[start + 1 : pos] - - if AUTOLINK_RE.search(url) is not None: - fullUrl = state.md.normalizeLink(url) - if not state.md.validateLink(fullUrl): - return False - - if not silent: - token = state.push("link_open", "a", 1) - token.attrs = {"href": fullUrl} - token.markup = "autolink" - token.info = "auto" - - token = state.push("text", "", 0) - token.content = state.md.normalizeLinkText(url) - - token = state.push("link_close", "a", -1) - token.markup = "autolink" - token.info = "auto" - - state.pos += len(url) + 2 - return True - - if EMAIL_RE.search(url) is not None: - fullUrl = state.md.normalizeLink("mailto:" + url) - if not state.md.validateLink(fullUrl): - return False - - if not silent: - token = state.push("link_open", "a", 1) - token.attrs = {"href": fullUrl} - token.markup = "autolink" - token.info = "auto" - - token = state.push("text", "", 0) - token.content = state.md.normalizeLinkText(url) - - token = state.push("link_close", "a", -1) - token.markup = "autolink" - token.info = "auto" - - state.pos += len(url) + 2 - return True - - return False diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/backticks.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/backticks.py deleted file mode 100644 index fc60d6b15cdfa7012a05bcf1ccbb06f44d870dfd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/backticks.py +++ /dev/null @@ -1,72 +0,0 @@ -# Parse backticks -import re - -from .state_inline import StateInline - -regex = re.compile("^ (.+) $") - - -def backtick(state: StateInline, silent: bool) -> bool: - pos = state.pos - - if state.src[pos] != "`": - return False - - start = pos - pos += 1 - maximum = state.posMax - - # scan marker length - while pos < maximum and (state.src[pos] == "`"): - pos += 1 - - marker = state.src[start:pos] - openerLength = len(marker) - - if state.backticksScanned and state.backticks.get(openerLength, 0) <= start: - if not silent: - state.pending += marker - state.pos += openerLength - return True - - matchStart = matchEnd = pos - - # Nothing found in the cache, scan until the end of the line (or until marker is found) - while True: - try: - matchStart = state.src.index("`", matchEnd) - except ValueError: - break - matchEnd = matchStart + 1 - - # scan marker length - while matchEnd < maximum and (state.src[matchEnd] == "`"): - matchEnd += 1 - - closerLength = matchEnd - matchStart - - if closerLength == openerLength: - # Found matching closer length. - if not silent: - token = state.push("code_inline", "code", 0) - token.markup = marker - token.content = state.src[pos:matchStart].replace("\n", " ") - if ( - token.content.startswith(" ") - and token.content.endswith(" ") - and len(token.content.strip()) > 0 - ): - token.content = token.content[1:-1] - state.pos = matchEnd - return True - - # Some different length found, put it in cache as upper limit of where closer can be found - state.backticks[closerLength] = matchStart - - # Scanned through the end, didn't find anything - state.backticksScanned = True - - if not silent: - state.pending += marker - state.pos += openerLength - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/balance_pairs.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/balance_pairs.py deleted file mode 100644 index bbb2101c7e1614dde2323d3a8a42b388f354789e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/balance_pairs.py +++ /dev/null @@ -1,137 +0,0 @@ -"""Balance paired characters (*, _, etc) in inline tokens.""" -from __future__ import annotations - -from .state_inline import Delimiter, StateInline - - -def processDelimiters(state: StateInline, delimiters: list[Delimiter]) -> None: - """For each opening emphasis-like marker find a matching closing one.""" - if not delimiters: - return - - openersBottom = {} - maximum = len(delimiters) - - # headerIdx is the first delimiter of the current (where closer is) delimiter run - headerIdx = 0 - lastTokenIdx = -2 # needs any value lower than -1 - jumps: list[int] = [] - closerIdx = 0 - while closerIdx < maximum: - closer = delimiters[closerIdx] - - jumps.append(0) - - # markers belong to same delimiter run if: - # - they have adjacent tokens - # - AND markers are the same - # - if ( - delimiters[headerIdx].marker != closer.marker - or lastTokenIdx != closer.token - 1 - ): - headerIdx = closerIdx - lastTokenIdx = closer.token - - # Length is only used for emphasis-specific "rule of 3", - # if it's not defined (in strikethrough or 3rd party plugins), - # we can default it to 0 to disable those checks. - # - closer.length = closer.length or 0 - - if not closer.close: - closerIdx += 1 - continue - - # Previously calculated lower bounds (previous fails) - # for each marker, each delimiter length modulo 3, - # and for whether this closer can be an opener; - # https://github.com/commonmark/cmark/commit/34250e12ccebdc6372b8b49c44fab57c72443460 - if closer.marker not in openersBottom: - openersBottom[closer.marker] = [-1, -1, -1, -1, -1, -1] - - minOpenerIdx = openersBottom[closer.marker][ - (3 if closer.open else 0) + (closer.length % 3) - ] - - openerIdx = headerIdx - jumps[headerIdx] - 1 - - newMinOpenerIdx = openerIdx - - while openerIdx > minOpenerIdx: - opener = delimiters[openerIdx] - - if opener.marker != closer.marker: - openerIdx -= jumps[openerIdx] + 1 - continue - - if opener.open and opener.end < 0: - isOddMatch = False - - # from spec: - # - # If one of the delimiters can both open and close emphasis, then the - # sum of the lengths of the delimiter runs containing the opening and - # closing delimiters must not be a multiple of 3 unless both lengths - # are multiples of 3. - # - if ( - (opener.close or closer.open) - and ((opener.length + closer.length) % 3 == 0) - and (opener.length % 3 != 0 or closer.length % 3 != 0) - ): - isOddMatch = True - - if not isOddMatch: - # If previous delimiter cannot be an opener, we can safely skip - # the entire sequence in future checks. This is required to make - # sure algorithm has linear complexity (see *_*_*_*_*_... case). - # - if openerIdx > 0 and not delimiters[openerIdx - 1].open: - lastJump = jumps[openerIdx - 1] + 1 - else: - lastJump = 0 - - jumps[closerIdx] = closerIdx - openerIdx + lastJump - jumps[openerIdx] = lastJump - - closer.open = False - opener.end = closerIdx - opener.close = False - newMinOpenerIdx = -1 - - # treat next token as start of run, - # it optimizes skips in **<...>**a**<...>** pathological case - lastTokenIdx = -2 - - break - - openerIdx -= jumps[openerIdx] + 1 - - if newMinOpenerIdx != -1: - # If match for this delimiter run failed, we want to set lower bound for - # future lookups. This is required to make sure algorithm has linear - # complexity. - # - # See details here: - # https:#github.com/commonmark/cmark/issues/178#issuecomment-270417442 - # - openersBottom[closer.marker][ - (3 if closer.open else 0) + ((closer.length or 0) % 3) - ] = newMinOpenerIdx - - closerIdx += 1 - - -def link_pairs(state: StateInline) -> None: - tokens_meta = state.tokens_meta - maximum = len(state.tokens_meta) - - processDelimiters(state, state.delimiters) - - curr = 0 - while curr < maximum: - curr_meta = tokens_meta[curr] - if curr_meta and "delimiters" in curr_meta: - processDelimiters(state, curr_meta["delimiters"]) - curr += 1 diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/emphasis.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/emphasis.py deleted file mode 100644 index 9a98f9e216c94db0217e986270aaaa72fcc99f7f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/emphasis.py +++ /dev/null @@ -1,102 +0,0 @@ -# Process *this* and _that_ -# -from __future__ import annotations - -from .state_inline import Delimiter, StateInline - - -def tokenize(state: StateInline, silent: bool) -> bool: - """Insert each marker as a separate text token, and add it to delimiter list""" - start = state.pos - marker = state.src[start] - - if silent: - return False - - if marker not in ("_", "*"): - return False - - scanned = state.scanDelims(state.pos, marker == "*") - - for _ in range(scanned.length): - token = state.push("text", "", 0) - token.content = marker - state.delimiters.append( - Delimiter( - marker=ord(marker), - length=scanned.length, - token=len(state.tokens) - 1, - end=-1, - open=scanned.can_open, - close=scanned.can_close, - ) - ) - - state.pos += scanned.length - - return True - - -def _postProcess(state: StateInline, delimiters: list[Delimiter]) -> None: - i = len(delimiters) - 1 - while i >= 0: - startDelim = delimiters[i] - - # /* _ */ /* * */ - if startDelim.marker != 0x5F and startDelim.marker != 0x2A: - i -= 1 - continue - - # Process only opening markers - if startDelim.end == -1: - i -= 1 - continue - - endDelim = delimiters[startDelim.end] - - # If the previous delimiter has the same marker and is adjacent to this one, - # merge those into one strong delimiter. - # - # `whatever` -> `whatever` - # - isStrong = ( - i > 0 - and delimiters[i - 1].end == startDelim.end + 1 - # check that first two markers match and adjacent - and delimiters[i - 1].marker == startDelim.marker - and delimiters[i - 1].token == startDelim.token - 1 - # check that last two markers are adjacent (we can safely assume they match) - and delimiters[startDelim.end + 1].token == endDelim.token + 1 - ) - - ch = chr(startDelim.marker) - - token = state.tokens[startDelim.token] - token.type = "strong_open" if isStrong else "em_open" - token.tag = "strong" if isStrong else "em" - token.nesting = 1 - token.markup = ch + ch if isStrong else ch - token.content = "" - - token = state.tokens[endDelim.token] - token.type = "strong_close" if isStrong else "em_close" - token.tag = "strong" if isStrong else "em" - token.nesting = -1 - token.markup = ch + ch if isStrong else ch - token.content = "" - - if isStrong: - state.tokens[delimiters[i - 1].token].content = "" - state.tokens[delimiters[startDelim.end + 1].token].content = "" - i -= 1 - - i -= 1 - - -def postProcess(state: StateInline) -> None: - """Walk through delimiter list and replace text tokens with tags.""" - _postProcess(state, state.delimiters) - - for token in state.tokens_meta: - if token and "delimiters" in token: - _postProcess(state, token["delimiters"]) diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/entity.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/entity.py deleted file mode 100644 index ec9d39650e5bc533e694d3d6699677068d22c69f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/entity.py +++ /dev/null @@ -1,53 +0,0 @@ -# Process html entity - {, ¯, ", ... -import re - -from ..common.entities import entities -from ..common.utils import fromCodePoint, isValidEntityCode -from .state_inline import StateInline - -DIGITAL_RE = re.compile(r"^&#((?:x[a-f0-9]{1,6}|[0-9]{1,7}));", re.IGNORECASE) -NAMED_RE = re.compile(r"^&([a-z][a-z0-9]{1,31});", re.IGNORECASE) - - -def entity(state: StateInline, silent: bool) -> bool: - pos = state.pos - maximum = state.posMax - - if state.src[pos] != "&": - return False - - if pos + 1 >= maximum: - return False - - if state.src[pos + 1] == "#": - if match := DIGITAL_RE.search(state.src[pos:]): - if not silent: - match1 = match.group(1) - code = ( - int(match1[1:], 16) if match1[0].lower() == "x" else int(match1, 10) - ) - - token = state.push("text_special", "", 0) - token.content = ( - fromCodePoint(code) - if isValidEntityCode(code) - else fromCodePoint(0xFFFD) - ) - token.markup = match.group(0) - token.info = "entity" - - state.pos += len(match.group(0)) - return True - - else: - if (match := NAMED_RE.search(state.src[pos:])) and match.group(1) in entities: - if not silent: - token = state.push("text_special", "", 0) - token.content = entities[match.group(1)] - token.markup = match.group(0) - token.info = "entity" - - state.pos += len(match.group(0)) - return True - - return False diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/escape.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/escape.py deleted file mode 100644 index 9f68b5dfad17fdcfe164ab4a8ef0761051b46efb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/escape.py +++ /dev/null @@ -1,92 +0,0 @@ -""" -Process escaped chars and hardbreaks -""" -from ..common.utils import isStrSpace -from .state_inline import StateInline - - -def escape(state: StateInline, silent: bool) -> bool: - """Process escaped chars and hardbreaks.""" - pos = state.pos - maximum = state.posMax - - if state.src[pos] != "\\": - return False - - pos += 1 - - # '\' at the end of the inline block - if pos >= maximum: - return False - - ch1 = state.src[pos] - ch1_ord = ord(ch1) - if ch1 == "\n": - if not silent: - state.push("hardbreak", "br", 0) - pos += 1 - # skip leading whitespaces from next line - while pos < maximum: - ch = state.src[pos] - if not isStrSpace(ch): - break - pos += 1 - - state.pos = pos - return True - - escapedStr = state.src[pos] - - if ch1_ord >= 0xD800 and ch1_ord <= 0xDBFF and pos + 1 < maximum: - ch2 = state.src[pos + 1] - ch2_ord = ord(ch2) - if ch2_ord >= 0xDC00 and ch2_ord <= 0xDFFF: - escapedStr += ch2 - pos += 1 - - origStr = "\\" + escapedStr - - if not silent: - token = state.push("text_special", "", 0) - token.content = escapedStr if ch1 in _ESCAPED else origStr - token.markup = origStr - token.info = "escape" - - state.pos = pos + 1 - return True - - -_ESCAPED = { - "!", - '"', - "#", - "$", - "%", - "&", - "'", - "(", - ")", - "*", - "+", - ",", - "-", - ".", - "/", - ":", - ";", - "<", - "=", - ">", - "?", - "@", - "[", - "\\", - "]", - "^", - "_", - "`", - "{", - "|", - "}", - "~", -} diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/fragments_join.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/fragments_join.py deleted file mode 100644 index f795c1364b8ac098b7a17f34cd31d7070280cf36..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/fragments_join.py +++ /dev/null @@ -1,43 +0,0 @@ -from .state_inline import StateInline - - -def fragments_join(state: StateInline) -> None: - """ - Clean up tokens after emphasis and strikethrough postprocessing: - merge adjacent text nodes into one and re-calculate all token levels - - This is necessary because initially emphasis delimiter markers (``*, _, ~``) - are treated as their own separate text tokens. Then emphasis rule either - leaves them as text (needed to merge with adjacent text) or turns them - into opening/closing tags (which messes up levels inside). - """ - level = 0 - maximum = len(state.tokens) - - curr = last = 0 - while curr < maximum: - # re-calculate levels after emphasis/strikethrough turns some text nodes - # into opening/closing tags - if state.tokens[curr].nesting < 0: - level -= 1 # closing tag - state.tokens[curr].level = level - if state.tokens[curr].nesting > 0: - level += 1 # opening tag - - if ( - state.tokens[curr].type == "text" - and curr + 1 < maximum - and state.tokens[curr + 1].type == "text" - ): - # collapse two adjacent text nodes - state.tokens[curr + 1].content = ( - state.tokens[curr].content + state.tokens[curr + 1].content - ) - else: - if curr != last: - state.tokens[last] = state.tokens[curr] - last += 1 - curr += 1 - - if curr != last: - del state.tokens[last:] diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/html_inline.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/html_inline.py deleted file mode 100644 index 9065e1d034da76270f7d3f1ba528132c8d57d341..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/html_inline.py +++ /dev/null @@ -1,43 +0,0 @@ -# Process html tags -from ..common.html_re import HTML_TAG_RE -from ..common.utils import isLinkClose, isLinkOpen -from .state_inline import StateInline - - -def isLetter(ch: int) -> bool: - lc = ch | 0x20 # to lower case - # /* a */ and /* z */ - return (lc >= 0x61) and (lc <= 0x7A) - - -def html_inline(state: StateInline, silent: bool) -> bool: - pos = state.pos - - if not state.md.options.get("html", None): - return False - - # Check start - maximum = state.posMax - if state.src[pos] != "<" or pos + 2 >= maximum: - return False - - # Quick fail on second char - ch = state.src[pos + 1] - if ch not in ("!", "?", "/") and not isLetter(ord(ch)): # /* / */ - return False - - match = HTML_TAG_RE.search(state.src[pos:]) - if not match: - return False - - if not silent: - token = state.push("html_inline", "", 0) - token.content = state.src[pos : pos + len(match.group(0))] - - if isLinkOpen(token.content): - state.linkLevel += 1 - if isLinkClose(token.content): - state.linkLevel -= 1 - - state.pos += len(match.group(0)) - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/image.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/image.py deleted file mode 100644 index b4a32a9f5880f3977c81f6b94d18751f93fa8561..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/image.py +++ /dev/null @@ -1,148 +0,0 @@ -# Process ![image]( "title") -from __future__ import annotations - -from ..common.utils import isStrSpace, normalizeReference -from ..token import Token -from .state_inline import StateInline - - -def image(state: StateInline, silent: bool) -> bool: - label = None - href = "" - oldPos = state.pos - max = state.posMax - - if state.src[state.pos] != "!": - return False - - if state.pos + 1 < state.posMax and state.src[state.pos + 1] != "[": - return False - - labelStart = state.pos + 2 - labelEnd = state.md.helpers.parseLinkLabel(state, state.pos + 1, False) - - # parser failed to find ']', so it's not a valid link - if labelEnd < 0: - return False - - pos = labelEnd + 1 - - if pos < max and state.src[pos] == "(": - # - # Inline link - # - - # [link]( "title" ) - # ^^ skipping these spaces - pos += 1 - while pos < max: - ch = state.src[pos] - if not isStrSpace(ch) and ch != "\n": - break - pos += 1 - - if pos >= max: - return False - - # [link]( "title" ) - # ^^^^^^ parsing link destination - start = pos - res = state.md.helpers.parseLinkDestination(state.src, pos, state.posMax) - if res.ok: - href = state.md.normalizeLink(res.str) - if state.md.validateLink(href): - pos = res.pos - else: - href = "" - - # [link]( "title" ) - # ^^ skipping these spaces - start = pos - while pos < max: - ch = state.src[pos] - if not isStrSpace(ch) and ch != "\n": - break - pos += 1 - - # [link]( "title" ) - # ^^^^^^^ parsing link title - res = state.md.helpers.parseLinkTitle(state.src, pos, state.posMax) - if pos < max and start != pos and res.ok: - title = res.str - pos = res.pos - - # [link]( "title" ) - # ^^ skipping these spaces - while pos < max: - ch = state.src[pos] - if not isStrSpace(ch) and ch != "\n": - break - pos += 1 - else: - title = "" - - if pos >= max or state.src[pos] != ")": - state.pos = oldPos - return False - - pos += 1 - - else: - # - # Link reference - # - if "references" not in state.env: - return False - - # /* [ */ - if pos < max and state.src[pos] == "[": - start = pos + 1 - pos = state.md.helpers.parseLinkLabel(state, pos) - if pos >= 0: - label = state.src[start:pos] - pos += 1 - else: - pos = labelEnd + 1 - else: - pos = labelEnd + 1 - - # covers label == '' and label == undefined - # (collapsed reference link and shortcut reference link respectively) - if not label: - label = state.src[labelStart:labelEnd] - - label = normalizeReference(label) - - ref = state.env["references"].get(label, None) - if not ref: - state.pos = oldPos - return False - - href = ref["href"] - title = ref["title"] - - # - # We found the end of the link, and know for a fact it's a valid link - # so all that's left to do is to call tokenizer. - # - if not silent: - content = state.src[labelStart:labelEnd] - - tokens: list[Token] = [] - state.md.inline.parse(content, state.md, state.env, tokens) - - token = state.push("image", "img", 0) - token.attrs = {"src": href, "alt": ""} - token.children = tokens or None - token.content = content - - if title: - token.attrSet("title", title) - - # note, this is not part of markdown-it JS, but is useful for renderers - if label and state.md.options.get("store_labels", False): - token.meta["label"] = label - - state.pos = pos - state.posMax = max - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/link.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/link.py deleted file mode 100644 index 78cf9122f3b4ad740769cb7edd3e4cc49af46e00..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/link.py +++ /dev/null @@ -1,151 +0,0 @@ -# Process [link]( "stuff") - -from ..common.utils import isStrSpace, normalizeReference -from .state_inline import StateInline - - -def link(state: StateInline, silent: bool) -> bool: - href = "" - title = "" - label = None - oldPos = state.pos - maximum = state.posMax - start = state.pos - parseReference = True - - if state.src[state.pos] != "[": - return False - - labelStart = state.pos + 1 - labelEnd = state.md.helpers.parseLinkLabel(state, state.pos, True) - - # parser failed to find ']', so it's not a valid link - if labelEnd < 0: - return False - - pos = labelEnd + 1 - - if pos < maximum and state.src[pos] == "(": - # - # Inline link - # - - # might have found a valid shortcut link, disable reference parsing - parseReference = False - - # [link]( "title" ) - # ^^ skipping these spaces - pos += 1 - while pos < maximum: - ch = state.src[pos] - if not isStrSpace(ch) and ch != "\n": - break - pos += 1 - - if pos >= maximum: - return False - - # [link]( "title" ) - # ^^^^^^ parsing link destination - start = pos - res = state.md.helpers.parseLinkDestination(state.src, pos, state.posMax) - if res.ok: - href = state.md.normalizeLink(res.str) - if state.md.validateLink(href): - pos = res.pos - else: - href = "" - - # [link]( "title" ) - # ^^ skipping these spaces - start = pos - while pos < maximum: - ch = state.src[pos] - if not isStrSpace(ch) and ch != "\n": - break - pos += 1 - - # [link]( "title" ) - # ^^^^^^^ parsing link title - res = state.md.helpers.parseLinkTitle(state.src, pos, state.posMax) - if pos < maximum and start != pos and res.ok: - title = res.str - pos = res.pos - - # [link]( "title" ) - # ^^ skipping these spaces - while pos < maximum: - ch = state.src[pos] - if not isStrSpace(ch) and ch != "\n": - break - pos += 1 - - if pos >= maximum or state.src[pos] != ")": - # parsing a valid shortcut link failed, fallback to reference - parseReference = True - - pos += 1 - - if parseReference: - # - # Link reference - # - if "references" not in state.env: - return False - - if pos < maximum and state.src[pos] == "[": - start = pos + 1 - pos = state.md.helpers.parseLinkLabel(state, pos) - if pos >= 0: - label = state.src[start:pos] - pos += 1 - else: - pos = labelEnd + 1 - - else: - pos = labelEnd + 1 - - # covers label == '' and label == undefined - # (collapsed reference link and shortcut reference link respectively) - if not label: - label = state.src[labelStart:labelEnd] - - label = normalizeReference(label) - - ref = ( - state.env["references"][label] if label in state.env["references"] else None - ) - if not ref: - state.pos = oldPos - return False - - href = ref["href"] - title = ref["title"] - - # - # We found the end of the link, and know for a fact it's a valid link - # so all that's left to do is to call tokenizer. - # - if not silent: - state.pos = labelStart - state.posMax = labelEnd - - token = state.push("link_open", "a", 1) - token.attrs = {"href": href} - - if title: - token.attrSet("title", title) - - # note, this is not part of markdown-it JS, but is useful for renderers - if label and state.md.options.get("store_labels", False): - token.meta["label"] = label - - state.linkLevel += 1 - state.md.inline.tokenize(state) - state.linkLevel -= 1 - - token = state.push("link_close", "a", -1) - - state.pos = pos - state.posMax = maximum - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/linkify.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/linkify.py deleted file mode 100644 index a8a181537d78867f8dfdf552b70f83be09bb2dd5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/linkify.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Process links like https://example.org/""" -import re - -from .state_inline import StateInline - -# RFC3986: scheme = ALPHA *( ALPHA / DIGIT / "+" / "-" / "." ) -SCHEME_RE = re.compile(r"(?:^|[^a-z0-9.+-])([a-z][a-z0-9.+-]*)$", re.IGNORECASE) - - -def linkify(state: StateInline, silent: bool) -> bool: - """Rule for identifying plain-text links.""" - if not state.md.options.linkify: - return False - if state.linkLevel > 0: - return False - if not state.md.linkify: - raise ModuleNotFoundError("Linkify enabled but not installed.") - - pos = state.pos - maximum = state.posMax - - if ( - (pos + 3) > maximum - or state.src[pos] != ":" - or state.src[pos + 1] != "/" - or state.src[pos + 2] != "/" - ): - return False - - if not (match := SCHEME_RE.match(state.pending)): - return False - - proto = match.group(1) - if not (link := state.md.linkify.match_at_start(state.src[pos - len(proto) :])): - return False - url: str = link.url - - # disallow '*' at the end of the link (conflicts with emphasis) - url = url.rstrip("*") - - full_url = state.md.normalizeLink(url) - if not state.md.validateLink(full_url): - return False - - if not silent: - state.pending = state.pending[: -len(proto)] - - token = state.push("link_open", "a", 1) - token.attrs = {"href": full_url} - token.markup = "linkify" - token.info = "auto" - - token = state.push("text", "", 0) - token.content = state.md.normalizeLinkText(url) - - token = state.push("link_close", "a", -1) - token.markup = "linkify" - token.info = "auto" - - state.pos += len(url) - len(proto) - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/newline.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/newline.py deleted file mode 100644 index ca8f1db02da07b023aa9fdb08ee7af326f773da8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/newline.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Proceess '\n'.""" -from ..common.utils import charStrAt, isStrSpace -from .state_inline import StateInline - - -def newline(state: StateInline, silent: bool) -> bool: - pos = state.pos - - if state.src[pos] != "\n": - return False - - pmax = len(state.pending) - 1 - maximum = state.posMax - - # ' \n' -> hardbreak - # Lookup in pending chars is bad practice! Don't copy to other rules! - # Pending string is stored in concat mode, indexed lookups will cause - # conversion to flat mode. - if not silent: - if pmax >= 0 and charStrAt(state.pending, pmax) == " ": - if pmax >= 1 and charStrAt(state.pending, pmax - 1) == " ": - # Find whitespaces tail of pending chars. - ws = pmax - 1 - while ws >= 1 and charStrAt(state.pending, ws - 1) == " ": - ws -= 1 - state.pending = state.pending[:ws] - - state.push("hardbreak", "br", 0) - else: - state.pending = state.pending[:-1] - state.push("softbreak", "br", 0) - - else: - state.push("softbreak", "br", 0) - - pos += 1 - - # skip heading spaces for next line - while pos < maximum and isStrSpace(state.src[pos]): - pos += 1 - - state.pos = pos - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/state_inline.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/state_inline.py deleted file mode 100644 index c0c491c4b7c9ae4117d60f447fdbf3c742f66f48..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/state_inline.py +++ /dev/null @@ -1,166 +0,0 @@ -from __future__ import annotations - -from collections import namedtuple -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Literal - -from .._compat import DATACLASS_KWARGS -from ..common.utils import isMdAsciiPunct, isPunctChar, isWhiteSpace -from ..ruler import StateBase -from ..token import Token -from ..utils import EnvType - -if TYPE_CHECKING: - from markdown_it import MarkdownIt - - -@dataclass(**DATACLASS_KWARGS) -class Delimiter: - # Char code of the starting marker (number). - marker: int - - # Total length of these series of delimiters. - length: int - - # A position of the token this delimiter corresponds to. - token: int - - # If this delimiter is matched as a valid opener, `end` will be - # equal to its position, otherwise it's `-1`. - end: int - - # Boolean flags that determine if this delimiter could open or close - # an emphasis. - open: bool - close: bool - - level: bool | None = None - - -Scanned = namedtuple("Scanned", ["can_open", "can_close", "length"]) - - -class StateInline(StateBase): - def __init__( - self, src: str, md: MarkdownIt, env: EnvType, outTokens: list[Token] - ) -> None: - self.src = src - self.env = env - self.md = md - self.tokens = outTokens - self.tokens_meta: list[dict[str, Any] | None] = [None] * len(outTokens) - - self.pos = 0 - self.posMax = len(self.src) - self.level = 0 - self.pending = "" - self.pendingLevel = 0 - - # Stores { start: end } pairs. Useful for backtrack - # optimization of pairs parse (emphasis, strikes). - self.cache: dict[int, int] = {} - - # List of emphasis-like delimiters for current tag - self.delimiters: list[Delimiter] = [] - - # Stack of delimiter lists for upper level tags - self._prev_delimiters: list[list[Delimiter]] = [] - - # backticklength => last seen position - self.backticks: dict[int, int] = {} - self.backticksScanned = False - - # Counter used to disable inline linkify-it execution - # inside and markdown links - self.linkLevel = 0 - - def __repr__(self) -> str: - return ( - f"{self.__class__.__name__}" - f"(pos=[{self.pos} of {self.posMax}], token={len(self.tokens)})" - ) - - def pushPending(self) -> Token: - token = Token("text", "", 0) - token.content = self.pending - token.level = self.pendingLevel - self.tokens.append(token) - self.pending = "" - return token - - def push(self, ttype: str, tag: str, nesting: Literal[-1, 0, 1]) -> Token: - """Push new token to "stream". - If pending text exists - flush it as text token - """ - if self.pending: - self.pushPending() - - token = Token(ttype, tag, nesting) - token_meta = None - - if nesting < 0: - # closing tag - self.level -= 1 - self.delimiters = self._prev_delimiters.pop() - - token.level = self.level - - if nesting > 0: - # opening tag - self.level += 1 - self._prev_delimiters.append(self.delimiters) - self.delimiters = [] - token_meta = {"delimiters": self.delimiters} - - self.pendingLevel = self.level - self.tokens.append(token) - self.tokens_meta.append(token_meta) - return token - - def scanDelims(self, start: int, canSplitWord: bool) -> Scanned: - """ - Scan a sequence of emphasis-like markers, and determine whether - it can start an emphasis sequence or end an emphasis sequence. - - - start - position to scan from (it should point at a valid marker); - - canSplitWord - determine if these markers can be found inside a word - - """ - pos = start - maximum = self.posMax - marker = self.src[start] - - # treat beginning of the line as a whitespace - lastChar = self.src[start - 1] if start > 0 else " " - - while pos < maximum and self.src[pos] == marker: - pos += 1 - - count = pos - start - - # treat end of the line as a whitespace - nextChar = self.src[pos] if pos < maximum else " " - - isLastPunctChar = isMdAsciiPunct(ord(lastChar)) or isPunctChar(lastChar) - isNextPunctChar = isMdAsciiPunct(ord(nextChar)) or isPunctChar(nextChar) - - isLastWhiteSpace = isWhiteSpace(ord(lastChar)) - isNextWhiteSpace = isWhiteSpace(ord(nextChar)) - - left_flanking = not ( - isNextWhiteSpace - or (isNextPunctChar and not (isLastWhiteSpace or isLastPunctChar)) - ) - right_flanking = not ( - isLastWhiteSpace - or (isLastPunctChar and not (isNextWhiteSpace or isNextPunctChar)) - ) - - if not canSplitWord: - can_open = left_flanking and ((not right_flanking) or isLastPunctChar) - can_close = right_flanking and ((not left_flanking) or isNextPunctChar) - else: - can_open = left_flanking - can_close = right_flanking - - return Scanned(can_open, can_close, count) diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/strikethrough.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/strikethrough.py deleted file mode 100644 index ec816281d49b23d0774bf91db6600d996aaf8b06..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/strikethrough.py +++ /dev/null @@ -1,127 +0,0 @@ -# ~~strike through~~ -from __future__ import annotations - -from .state_inline import Delimiter, StateInline - - -def tokenize(state: StateInline, silent: bool) -> bool: - """Insert each marker as a separate text token, and add it to delimiter list""" - start = state.pos - ch = state.src[start] - - if silent: - return False - - if ch != "~": - return False - - scanned = state.scanDelims(state.pos, True) - length = scanned.length - - if length < 2: - return False - - if length % 2: - token = state.push("text", "", 0) - token.content = ch - length -= 1 - - i = 0 - while i < length: - token = state.push("text", "", 0) - token.content = ch + ch - state.delimiters.append( - Delimiter( - marker=ord(ch), - length=0, # disable "rule of 3" length checks meant for emphasis - token=len(state.tokens) - 1, - end=-1, - open=scanned.can_open, - close=scanned.can_close, - ) - ) - - i += 2 - - state.pos += scanned.length - - return True - - -def _postProcess(state: StateInline, delimiters: list[Delimiter]) -> None: - loneMarkers = [] - maximum = len(delimiters) - - i = 0 - while i < maximum: - startDelim = delimiters[i] - - if startDelim.marker != 0x7E: # /* ~ */ - i += 1 - continue - - if startDelim.end == -1: - i += 1 - continue - - endDelim = delimiters[startDelim.end] - - token = state.tokens[startDelim.token] - token.type = "s_open" - token.tag = "s" - token.nesting = 1 - token.markup = "~~" - token.content = "" - - token = state.tokens[endDelim.token] - token.type = "s_close" - token.tag = "s" - token.nesting = -1 - token.markup = "~~" - token.content = "" - - if ( - state.tokens[endDelim.token - 1].type == "text" - and state.tokens[endDelim.token - 1].content == "~" - ): - loneMarkers.append(endDelim.token - 1) - - i += 1 - - # If a marker sequence has an odd number of characters, it's split - # like this: `~~~~~` -> `~` + `~~` + `~~`, leaving one marker at the - # start of the sequence. - # - # So, we have to move all those markers after subsequent s_close tags. - # - while loneMarkers: - i = loneMarkers.pop() - j = i + 1 - - while (j < len(state.tokens)) and (state.tokens[j].type == "s_close"): - j += 1 - - j -= 1 - - if i != j: - token = state.tokens[j] - state.tokens[j] = state.tokens[i] - state.tokens[i] = token - - -def postProcess(state: StateInline) -> None: - """Walk through delimiter list and replace text tokens with tags.""" - tokens_meta = state.tokens_meta - maximum = len(state.tokens_meta) - _postProcess(state, state.delimiters) - - curr = 0 - while curr < maximum: - try: - curr_meta = tokens_meta[curr] - except IndexError: - pass - else: - if curr_meta and "delimiters" in curr_meta: - _postProcess(state, curr_meta["delimiters"]) - curr += 1 diff --git a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/text.py b/venv/lib/python3.10/site-packages/markdown_it/rules_inline/text.py deleted file mode 100644 index f306b2e4cecde67aa3d363d3844e8b40a70a48b1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/rules_inline/text.py +++ /dev/null @@ -1,53 +0,0 @@ -# Skip text characters for text token, place those to pending buffer -# and increment current pos -from .state_inline import StateInline - -# Rule to skip pure text -# '{}$%@~+=:' reserved for extensions - -# !!!! Don't confuse with "Markdown ASCII Punctuation" chars -# http://spec.commonmark.org/0.15/#ascii-punctuation-character - - -_TerminatorChars = { - "\n", - "!", - "#", - "$", - "%", - "&", - "*", - "+", - "-", - ":", - "<", - "=", - ">", - "@", - "[", - "\\", - "]", - "^", - "_", - "`", - "{", - "}", - "~", -} - - -def text(state: StateInline, silent: bool) -> bool: - pos = state.pos - posMax = state.posMax - while (pos < posMax) and state.src[pos] not in _TerminatorChars: - pos += 1 - - if pos == state.pos: - return False - - if not silent: - state.pending += state.src[state.pos : pos] - - state.pos = pos - - return True diff --git a/venv/lib/python3.10/site-packages/markdown_it/token.py b/venv/lib/python3.10/site-packages/markdown_it/token.py deleted file mode 100644 index 90008b7229d025a2f2432516d3a7e0da5f6781f9..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/token.py +++ /dev/null @@ -1,180 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable, MutableMapping -import dataclasses as dc -from typing import Any, Literal -import warnings - -from markdown_it._compat import DATACLASS_KWARGS - - -def convert_attrs(value: Any) -> Any: - """Convert Token.attrs set as ``None`` or ``[[key, value], ...]`` to a dict. - - This improves compatibility with upstream markdown-it. - """ - if not value: - return {} - if isinstance(value, list): - return dict(value) - return value - - -@dc.dataclass(**DATACLASS_KWARGS) -class Token: - type: str - """Type of the token (string, e.g. "paragraph_open")""" - - tag: str - """HTML tag name, e.g. 'p'""" - - nesting: Literal[-1, 0, 1] - """Level change (number in {-1, 0, 1} set), where: - - `1` means the tag is opening - - `0` means the tag is self-closing - - `-1` means the tag is closing - """ - - attrs: dict[str, str | int | float] = dc.field(default_factory=dict) - """HTML attributes. - Note this differs from the upstream "list of lists" format, - although than an instance can still be initialised with this format. - """ - - map: list[int] | None = None - """Source map info. Format: `[ line_begin, line_end ]`""" - - level: int = 0 - """Nesting level, the same as `state.level`""" - - children: list[Token] | None = None - """Array of child nodes (inline and img tokens).""" - - content: str = "" - """Inner content, in the case of a self-closing tag (code, html, fence, etc.),""" - - markup: str = "" - """'*' or '_' for emphasis, fence string for fence, etc.""" - - info: str = "" - """Additional information: - - Info string for "fence" tokens - - The value "auto" for autolink "link_open" and "link_close" tokens - - The string value of the item marker for ordered-list "list_item_open" tokens - """ - - meta: dict[Any, Any] = dc.field(default_factory=dict) - """A place for plugins to store any arbitrary data""" - - block: bool = False - """True for block-level tokens, false for inline tokens. - Used in renderer to calculate line breaks - """ - - hidden: bool = False - """If true, ignore this element when rendering. - Used for tight lists to hide paragraphs. - """ - - def __post_init__(self) -> None: - self.attrs = convert_attrs(self.attrs) - - def attrIndex(self, name: str) -> int: - warnings.warn( # noqa: B028 - "Token.attrIndex should not be used, since Token.attrs is a dictionary", - UserWarning, - ) - if name not in self.attrs: - return -1 - return list(self.attrs.keys()).index(name) - - def attrItems(self) -> list[tuple[str, str | int | float]]: - """Get (key, value) list of attrs.""" - return list(self.attrs.items()) - - def attrPush(self, attrData: tuple[str, str | int | float]) -> None: - """Add `[ name, value ]` attribute to list. Init attrs if necessary.""" - name, value = attrData - self.attrSet(name, value) - - def attrSet(self, name: str, value: str | int | float) -> None: - """Set `name` attribute to `value`. Override old value if exists.""" - self.attrs[name] = value - - def attrGet(self, name: str) -> None | str | int | float: - """Get the value of attribute `name`, or null if it does not exist.""" - return self.attrs.get(name, None) - - def attrJoin(self, name: str, value: str) -> None: - """Join value to existing attribute via space. - Or create new attribute if not exists. - Useful to operate with token classes. - """ - if name in self.attrs: - current = self.attrs[name] - if not isinstance(current, str): - raise TypeError( - f"existing attr 'name' is not a str: {self.attrs[name]}" - ) - self.attrs[name] = f"{current} {value}" - else: - self.attrs[name] = value - - def copy(self, **changes: Any) -> Token: - """Return a shallow copy of the instance.""" - return dc.replace(self, **changes) - - def as_dict( - self, - *, - children: bool = True, - as_upstream: bool = True, - meta_serializer: Callable[[dict[Any, Any]], Any] | None = None, - filter: Callable[[str, Any], bool] | None = None, - dict_factory: Callable[..., MutableMapping[str, Any]] = dict, - ) -> MutableMapping[str, Any]: - """Return the token as a dictionary. - - :param children: Also convert children to dicts - :param as_upstream: Ensure the output dictionary is equal to that created by markdown-it - For example, attrs are converted to null or lists - :param meta_serializer: hook for serializing ``Token.meta`` - :param filter: A callable whose return code determines whether an - attribute or element is included (``True``) or dropped (``False``). - Is called with the (key, value) pair. - :param dict_factory: A callable to produce dictionaries from. - For example, to produce ordered dictionaries instead of normal Python - dictionaries, pass in ``collections.OrderedDict``. - - """ - mapping = dict_factory((f.name, getattr(self, f.name)) for f in dc.fields(self)) - if filter: - mapping = dict_factory((k, v) for k, v in mapping.items() if filter(k, v)) - if as_upstream and "attrs" in mapping: - mapping["attrs"] = ( - None - if not mapping["attrs"] - else [[k, v] for k, v in mapping["attrs"].items()] - ) - if meta_serializer and "meta" in mapping: - mapping["meta"] = meta_serializer(mapping["meta"]) - if children and mapping.get("children", None): - mapping["children"] = [ - child.as_dict( - children=children, - filter=filter, - dict_factory=dict_factory, - as_upstream=as_upstream, - meta_serializer=meta_serializer, - ) - for child in mapping["children"] - ] - return mapping - - @classmethod - def from_dict(cls, dct: MutableMapping[str, Any]) -> Token: - """Convert a dict to a Token.""" - token = cls(**dct) - if token.children: - token.children = [cls.from_dict(c) for c in token.children] # type: ignore[arg-type] - return token diff --git a/venv/lib/python3.10/site-packages/markdown_it/tree.py b/venv/lib/python3.10/site-packages/markdown_it/tree.py deleted file mode 100644 index 6641e5a44654c9414cff07b6abbc633de7108ecb..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/tree.py +++ /dev/null @@ -1,345 +0,0 @@ -"""A tree representation of a linear markdown-it token stream. - -This module is not part of upstream JavaScript markdown-it. -""" -from __future__ import annotations - -from collections.abc import Generator, Sequence -import textwrap -from typing import Any, NamedTuple, TypeVar, overload - -from .token import Token - - -class _NesterTokens(NamedTuple): - opening: Token - closing: Token - - -_NodeType = TypeVar("_NodeType", bound="SyntaxTreeNode") - - -class SyntaxTreeNode: - """A Markdown syntax tree node. - - A class that can be used to construct a tree representation of a linear - `markdown-it-py` token stream. - - Each node in the tree represents either: - - root of the Markdown document - - a single unnested `Token` - - a `Token` "_open" and "_close" token pair, and the tokens nested in - between - """ - - def __init__( - self, tokens: Sequence[Token] = (), *, create_root: bool = True - ) -> None: - """Initialize a `SyntaxTreeNode` from a token stream. - - If `create_root` is True, create a root node for the document. - """ - # Only nodes representing an unnested token have self.token - self.token: Token | None = None - - # Only containers have nester tokens - self.nester_tokens: _NesterTokens | None = None - - # Root node does not have self.parent - self._parent: Any = None - - # Empty list unless a non-empty container, or unnested token that has - # children (i.e. inline or img) - self._children: list[Any] = [] - - if create_root: - self._set_children_from_tokens(tokens) - return - - if not tokens: - raise ValueError( - "Can only create root from empty token sequence." - " Set `create_root=True`." - ) - elif len(tokens) == 1: - inline_token = tokens[0] - if inline_token.nesting: - raise ValueError( - "Unequal nesting level at the start and end of token stream." - ) - self.token = inline_token - if inline_token.children: - self._set_children_from_tokens(inline_token.children) - else: - self.nester_tokens = _NesterTokens(tokens[0], tokens[-1]) - self._set_children_from_tokens(tokens[1:-1]) - - def __repr__(self) -> str: - return f"{type(self).__name__}({self.type})" - - @overload - def __getitem__(self: _NodeType, item: int) -> _NodeType: - ... - - @overload - def __getitem__(self: _NodeType, item: slice) -> list[_NodeType]: - ... - - def __getitem__(self: _NodeType, item: int | slice) -> _NodeType | list[_NodeType]: - return self.children[item] - - def to_tokens(self: _NodeType) -> list[Token]: - """Recover the linear token stream.""" - - def recursive_collect_tokens(node: _NodeType, token_list: list[Token]) -> None: - if node.type == "root": - for child in node.children: - recursive_collect_tokens(child, token_list) - elif node.token: - token_list.append(node.token) - else: - assert node.nester_tokens - token_list.append(node.nester_tokens.opening) - for child in node.children: - recursive_collect_tokens(child, token_list) - token_list.append(node.nester_tokens.closing) - - tokens: list[Token] = [] - recursive_collect_tokens(self, tokens) - return tokens - - @property - def children(self: _NodeType) -> list[_NodeType]: - return self._children - - @children.setter - def children(self: _NodeType, value: list[_NodeType]) -> None: - self._children = value - - @property - def parent(self: _NodeType) -> _NodeType | None: - return self._parent # type: ignore - - @parent.setter - def parent(self: _NodeType, value: _NodeType | None) -> None: - self._parent = value - - @property - def is_root(self) -> bool: - """Is the node a special root node?""" - return not (self.token or self.nester_tokens) - - @property - def is_nested(self) -> bool: - """Is this node nested?. - - Returns `True` if the node represents a `Token` pair and tokens in the - sequence between them, where `Token.nesting` of the first `Token` in - the pair is 1 and nesting of the other `Token` is -1. - """ - return bool(self.nester_tokens) - - @property - def siblings(self: _NodeType) -> Sequence[_NodeType]: - """Get siblings of the node. - - Gets the whole group of siblings, including self. - """ - if not self.parent: - return [self] - return self.parent.children - - @property - def type(self) -> str: - """Get a string type of the represented syntax. - - - "root" for root nodes - - `Token.type` if the node represents an unnested token - - `Token.type` of the opening token, with "_open" suffix stripped, if - the node represents a nester token pair - """ - if self.is_root: - return "root" - if self.token: - return self.token.type - assert self.nester_tokens - return _removesuffix(self.nester_tokens.opening.type, "_open") - - @property - def next_sibling(self: _NodeType) -> _NodeType | None: - """Get the next node in the sequence of siblings. - - Returns `None` if this is the last sibling. - """ - self_index = self.siblings.index(self) - if self_index + 1 < len(self.siblings): - return self.siblings[self_index + 1] - return None - - @property - def previous_sibling(self: _NodeType) -> _NodeType | None: - """Get the previous node in the sequence of siblings. - - Returns `None` if this is the first sibling. - """ - self_index = self.siblings.index(self) - if self_index - 1 >= 0: - return self.siblings[self_index - 1] - return None - - def _add_child( - self, - tokens: Sequence[Token], - ) -> None: - """Make a child node for `self`.""" - child = type(self)(tokens, create_root=False) - child.parent = self - self.children.append(child) - - def _set_children_from_tokens(self, tokens: Sequence[Token]) -> None: - """Convert the token stream to a tree structure and set the resulting - nodes as children of `self`.""" - reversed_tokens = list(reversed(tokens)) - while reversed_tokens: - token = reversed_tokens.pop() - - if not token.nesting: - self._add_child([token]) - continue - if token.nesting != 1: - raise ValueError("Invalid token nesting") - - nested_tokens = [token] - nesting = 1 - while reversed_tokens and nesting: - token = reversed_tokens.pop() - nested_tokens.append(token) - nesting += token.nesting - if nesting: - raise ValueError(f"unclosed tokens starting {nested_tokens[0]}") - - self._add_child(nested_tokens) - - def pretty( - self, *, indent: int = 2, show_text: bool = False, _current: int = 0 - ) -> str: - """Create an XML style string of the tree.""" - prefix = " " * _current - text = prefix + f"<{self.type}" - if not self.is_root and self.attrs: - text += " " + " ".join(f"{k}={v!r}" for k, v in self.attrs.items()) - text += ">" - if ( - show_text - and not self.is_root - and self.type in ("text", "text_special") - and self.content - ): - text += "\n" + textwrap.indent(self.content, prefix + " " * indent) - for child in self.children: - text += "\n" + child.pretty( - indent=indent, show_text=show_text, _current=_current + indent - ) - return text - - def walk( - self: _NodeType, *, include_self: bool = True - ) -> Generator[_NodeType, None, None]: - """Recursively yield all descendant nodes in the tree starting at self. - - The order mimics the order of the underlying linear token - stream (i.e. depth first). - """ - if include_self: - yield self - for child in self.children: - yield from child.walk(include_self=True) - - # NOTE: - # The values of the properties defined below directly map to properties - # of the underlying `Token`s. A root node does not translate to a `Token` - # object, so calling these property getters on a root node will raise an - # `AttributeError`. - # - # There is no mapping for `Token.nesting` because the `is_nested` property - # provides that data, and can be called on any node type, including root. - - def _attribute_token(self) -> Token: - """Return the `Token` that is used as the data source for the - properties defined below.""" - if self.token: - return self.token - if self.nester_tokens: - return self.nester_tokens.opening - raise AttributeError("Root node does not have the accessed attribute") - - @property - def tag(self) -> str: - """html tag name, e.g. \"p\" """ - return self._attribute_token().tag - - @property - def attrs(self) -> dict[str, str | int | float]: - """Html attributes.""" - return self._attribute_token().attrs - - def attrGet(self, name: str) -> None | str | int | float: - """Get the value of attribute `name`, or null if it does not exist.""" - return self._attribute_token().attrGet(name) - - @property - def map(self) -> tuple[int, int] | None: - """Source map info. Format: `tuple[ line_begin, line_end ]`""" - map_ = self._attribute_token().map - if map_: - # Type ignore because `Token`s attribute types are not perfect - return tuple(map_) # type: ignore - return None - - @property - def level(self) -> int: - """nesting level, the same as `state.level`""" - return self._attribute_token().level - - @property - def content(self) -> str: - """In a case of self-closing tag (code, html, fence, etc.), it - has contents of this tag.""" - return self._attribute_token().content - - @property - def markup(self) -> str: - """'*' or '_' for emphasis, fence string for fence, etc.""" - return self._attribute_token().markup - - @property - def info(self) -> str: - """fence infostring""" - return self._attribute_token().info - - @property - def meta(self) -> dict[Any, Any]: - """A place for plugins to store an arbitrary data.""" - return self._attribute_token().meta - - @property - def block(self) -> bool: - """True for block-level tokens, false for inline tokens.""" - return self._attribute_token().block - - @property - def hidden(self) -> bool: - """If it's true, ignore this element when rendering. - Used for tight lists to hide paragraphs.""" - return self._attribute_token().hidden - - -def _removesuffix(string: str, suffix: str) -> str: - """Remove a suffix from a string. - - Replace this with str.removesuffix() from stdlib when minimum Python - version is 3.9. - """ - if suffix and string.endswith(suffix): - return string[: -len(suffix)] - return string diff --git a/venv/lib/python3.10/site-packages/markdown_it/utils.py b/venv/lib/python3.10/site-packages/markdown_it/utils.py deleted file mode 100644 index a97937208a73bf06817d523d03ba868ceb13a65a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it/utils.py +++ /dev/null @@ -1,176 +0,0 @@ -from __future__ import annotations - -from collections.abc import MutableMapping as MutableMappingABC -from pathlib import Path -from typing import Any, Callable, Iterable, MutableMapping, TypedDict, cast - -EnvType = MutableMapping[str, Any] # note: could use TypeAlias in python 3.10 -"""Type for the environment sandbox used in parsing and rendering, -which stores mutable variables for use by plugins and rules. -""" - - -class OptionsType(TypedDict): - """Options for parsing.""" - - maxNesting: int - """Internal protection, recursion limit.""" - html: bool - """Enable HTML tags in source.""" - linkify: bool - """Enable autoconversion of URL-like texts to links.""" - typographer: bool - """Enable smartquotes and replacements.""" - quotes: str - """Quote characters.""" - xhtmlOut: bool - """Use '/' to close single tags (
).""" - breaks: bool - """Convert newlines in paragraphs into
.""" - langPrefix: str - """CSS language prefix for fenced blocks.""" - highlight: Callable[[str, str, str], str] | None - """Highlighter function: (content, lang, attrs) -> str.""" - - -class PresetType(TypedDict): - """Preset configuration for markdown-it.""" - - options: OptionsType - """Options for parsing.""" - components: MutableMapping[str, MutableMapping[str, list[str]]] - """Components for parsing and rendering.""" - - -class OptionsDict(MutableMappingABC): # type: ignore - """A dictionary, with attribute access to core markdownit configuration options.""" - - # Note: ideally we would probably just remove attribute access entirely, - # but we keep it for backwards compatibility. - - def __init__(self, options: OptionsType) -> None: - self._options = cast(OptionsType, dict(options)) - - def __getitem__(self, key: str) -> Any: - return self._options[key] # type: ignore[literal-required] - - def __setitem__(self, key: str, value: Any) -> None: - self._options[key] = value # type: ignore[literal-required] - - def __delitem__(self, key: str) -> None: - del self._options[key] # type: ignore - - def __iter__(self) -> Iterable[str]: # type: ignore - return iter(self._options) - - def __len__(self) -> int: - return len(self._options) - - def __repr__(self) -> str: - return repr(self._options) - - def __str__(self) -> str: - return str(self._options) - - @property - def maxNesting(self) -> int: - """Internal protection, recursion limit.""" - return self._options["maxNesting"] - - @maxNesting.setter - def maxNesting(self, value: int) -> None: - self._options["maxNesting"] = value - - @property - def html(self) -> bool: - """Enable HTML tags in source.""" - return self._options["html"] - - @html.setter - def html(self, value: bool) -> None: - self._options["html"] = value - - @property - def linkify(self) -> bool: - """Enable autoconversion of URL-like texts to links.""" - return self._options["linkify"] - - @linkify.setter - def linkify(self, value: bool) -> None: - self._options["linkify"] = value - - @property - def typographer(self) -> bool: - """Enable smartquotes and replacements.""" - return self._options["typographer"] - - @typographer.setter - def typographer(self, value: bool) -> None: - self._options["typographer"] = value - - @property - def quotes(self) -> str: - """Quote characters.""" - return self._options["quotes"] - - @quotes.setter - def quotes(self, value: str) -> None: - self._options["quotes"] = value - - @property - def xhtmlOut(self) -> bool: - """Use '/' to close single tags (
).""" - return self._options["xhtmlOut"] - - @xhtmlOut.setter - def xhtmlOut(self, value: bool) -> None: - self._options["xhtmlOut"] = value - - @property - def breaks(self) -> bool: - """Convert newlines in paragraphs into
.""" - return self._options["breaks"] - - @breaks.setter - def breaks(self, value: bool) -> None: - self._options["breaks"] = value - - @property - def langPrefix(self) -> str: - """CSS language prefix for fenced blocks.""" - return self._options["langPrefix"] - - @langPrefix.setter - def langPrefix(self, value: str) -> None: - self._options["langPrefix"] = value - - @property - def highlight(self) -> Callable[[str, str, str], str] | None: - """Highlighter function: (content, langName, langAttrs) -> escaped HTML.""" - return self._options["highlight"] - - @highlight.setter - def highlight(self, value: Callable[[str, str, str], str] | None) -> None: - self._options["highlight"] = value - - -def read_fixture_file(path: str | Path) -> list[list[Any]]: - text = Path(path).read_text(encoding="utf-8") - tests = [] - section = 0 - last_pos = 0 - lines = text.splitlines(keepends=True) - for i in range(len(lines)): - if lines[i].rstrip() == ".": - if section == 0: - tests.append([i, lines[i - 1].strip()]) - section = 1 - elif section == 1: - tests[-1].append("".join(lines[last_pos + 1 : i])) - section = 2 - elif section == 2: - tests[-1].append("".join(lines[last_pos + 1 : i])) - section = 0 - - last_pos = i - return tests diff --git a/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/LICENSE b/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/LICENSE deleted file mode 100644 index 582ddf59e08277fe6e78cee924d2c84805fe36fe..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2020 ExecutableBookProject - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/LICENSE.markdown-it b/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/LICENSE.markdown-it deleted file mode 100644 index 7ffa058cb78f8fb9beb974d9fd429004d2d2e585..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/LICENSE.markdown-it +++ /dev/null @@ -1,22 +0,0 @@ -Copyright (c) 2014 Vitaly Puzrin, Alex Kocharin. - -Permission is hereby granted, free of charge, to any person -obtaining a copy of this software and associated documentation -files (the "Software"), to deal in the Software without -restriction, including without limitation the rights to use, -copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the -Software is furnished to do so, subject to the following -conditions: - -The above copyright notice and this permission notice shall be -included in all copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, -EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES -OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND -NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT -HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, -WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR -OTHER DEALINGS IN THE SOFTWARE. diff --git a/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/METADATA b/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/METADATA deleted file mode 100644 index 8a2978bc0655edefc94ddbce62d5dd46259737bf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/METADATA +++ /dev/null @@ -1,205 +0,0 @@ -Metadata-Version: 2.1 -Name: markdown-it-py -Version: 3.0.0 -Summary: Python port of markdown-it. Markdown parsing, done right! -Keywords: markdown,lexer,parser,commonmark,markdown-it -Author-email: Chris Sewell -Requires-Python: >=3.8 -Description-Content-Type: text/markdown -Classifier: Development Status :: 5 - Production/Stable -Classifier: Intended Audience :: Developers -Classifier: License :: OSI Approved :: MIT License -Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3.8 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Programming Language :: Python :: 3.10 -Classifier: Programming Language :: Python :: 3.11 -Classifier: Programming Language :: Python :: Implementation :: CPython -Classifier: Programming Language :: Python :: Implementation :: PyPy -Classifier: Topic :: Software Development :: Libraries :: Python Modules -Classifier: Topic :: Text Processing :: Markup -Requires-Dist: mdurl~=0.1 -Requires-Dist: psutil ; extra == "benchmarking" -Requires-Dist: pytest ; extra == "benchmarking" -Requires-Dist: pytest-benchmark ; extra == "benchmarking" -Requires-Dist: pre-commit~=3.0 ; extra == "code_style" -Requires-Dist: commonmark~=0.9 ; extra == "compare" -Requires-Dist: markdown~=3.4 ; extra == "compare" -Requires-Dist: mistletoe~=1.0 ; extra == "compare" -Requires-Dist: mistune~=2.0 ; extra == "compare" -Requires-Dist: panflute~=2.3 ; extra == "compare" -Requires-Dist: linkify-it-py>=1,<3 ; extra == "linkify" -Requires-Dist: mdit-py-plugins ; extra == "plugins" -Requires-Dist: gprof2dot ; extra == "profiling" -Requires-Dist: mdit-py-plugins ; extra == "rtd" -Requires-Dist: myst-parser ; extra == "rtd" -Requires-Dist: pyyaml ; extra == "rtd" -Requires-Dist: sphinx ; extra == "rtd" -Requires-Dist: sphinx-copybutton ; extra == "rtd" -Requires-Dist: sphinx-design ; extra == "rtd" -Requires-Dist: sphinx_book_theme ; extra == "rtd" -Requires-Dist: jupyter_sphinx ; extra == "rtd" -Requires-Dist: coverage ; extra == "testing" -Requires-Dist: pytest ; extra == "testing" -Requires-Dist: pytest-cov ; extra == "testing" -Requires-Dist: pytest-regressions ; extra == "testing" -Project-URL: Documentation, https://markdown-it-py.readthedocs.io -Project-URL: Homepage, https://github.com/executablebooks/markdown-it-py -Provides-Extra: benchmarking -Provides-Extra: code_style -Provides-Extra: compare -Provides-Extra: linkify -Provides-Extra: plugins -Provides-Extra: profiling -Provides-Extra: rtd -Provides-Extra: testing - -# markdown-it-py - -[![Github-CI][github-ci]][github-link] -[![Coverage Status][codecov-badge]][codecov-link] -[![PyPI][pypi-badge]][pypi-link] -[![Conda][conda-badge]][conda-link] -[![Code style: black][black-badge]][black-link] -[![PyPI - Downloads][install-badge]][install-link] - -> Markdown parser done right. - -- Follows the __[CommonMark spec](http://spec.commonmark.org/)__ for baseline parsing -- Configurable syntax: you can add new rules and even replace existing ones. -- Pluggable: Adds syntax extensions to extend the parser (see the [plugin list][md-plugins]). -- High speed (see our [benchmarking tests][md-performance]) -- [Safe by default][md-security] -- Member of [Google's Assured Open Source Software](https://cloud.google.com/assured-open-source-software/docs/supported-packages) - -This is a Python port of [markdown-it], and some of its associated plugins. -For more details see: . - -For details on [markdown-it] itself, see: - -- The __[Live demo](https://markdown-it.github.io)__ -- [The markdown-it README][markdown-it-readme] - -## Installation - -```bash -conda install -c conda-forge markdown-it-py -``` - -or - -```bash -pip install markdown-it-py[plugins] -``` - -or with extras - -```bash -conda install -c conda-forge markdown-it-py linkify-it-py mdit-py-plugins -pip install markdown-it-py[linkify,plugins] -``` - -## Usage - -### Python API Usage - -Render markdown to HTML with markdown-it-py and a custom configuration -with and without plugins and features: - -```python -from markdown_it import MarkdownIt -from mdit_py_plugins.front_matter import front_matter_plugin -from mdit_py_plugins.footnote import footnote_plugin - -md = ( - MarkdownIt('commonmark' ,{'breaks':True,'html':True}) - .use(front_matter_plugin) - .use(footnote_plugin) - .enable('table') -) -text = (""" ---- -a: 1 ---- - -a | b -- | - -1 | 2 - -A footnote [^1] - -[^1]: some details -""") -tokens = md.parse(text) -html_text = md.render(text) - -## To export the html to a file, uncomment the lines below: -# from pathlib import Path -# Path("output.html").write_text(html_text) -``` - -### Command-line Usage - -Render markdown to HTML with markdown-it-py from the -command-line: - -```console -usage: markdown-it [-h] [-v] [filenames [filenames ...]] - -Parse one or more markdown files, convert each to HTML, and print to stdout - -positional arguments: - filenames specify an optional list of files to convert - -optional arguments: - -h, --help show this help message and exit - -v, --version show program's version number and exit - -Interactive: - - $ markdown-it - markdown-it-py [version 0.0.0] (interactive) - Type Ctrl-D to complete input, or Ctrl-C to exit. - >>> # Example - ... > markdown *input* - ... -

Example

-
-

markdown input

-
- -Batch: - - $ markdown-it README.md README.footer.md > index.html - -``` - -## References / Thanks - -Big thanks to the authors of [markdown-it]: - -- Alex Kocharin [github/rlidwka](https://github.com/rlidwka) -- Vitaly Puzrin [github/puzrin](https://github.com/puzrin) - -Also [John MacFarlane](https://github.com/jgm) for his work on the CommonMark spec and reference implementations. - -[github-ci]: https://github.com/executablebooks/markdown-it-py/workflows/Python%20package/badge.svg?branch=master -[github-link]: https://github.com/executablebooks/markdown-it-py -[pypi-badge]: https://img.shields.io/pypi/v/markdown-it-py.svg -[pypi-link]: https://pypi.org/project/markdown-it-py -[conda-badge]: https://anaconda.org/conda-forge/markdown-it-py/badges/version.svg -[conda-link]: https://anaconda.org/conda-forge/markdown-it-py -[codecov-badge]: https://codecov.io/gh/executablebooks/markdown-it-py/branch/master/graph/badge.svg -[codecov-link]: https://codecov.io/gh/executablebooks/markdown-it-py -[black-badge]: https://img.shields.io/badge/code%20style-black-000000.svg -[black-link]: https://github.com/ambv/black -[install-badge]: https://img.shields.io/pypi/dw/markdown-it-py?label=pypi%20installs -[install-link]: https://pypistats.org/packages/markdown-it-py - -[CommonMark spec]: http://spec.commonmark.org/ -[markdown-it]: https://github.com/markdown-it/markdown-it -[markdown-it-readme]: https://github.com/markdown-it/markdown-it/blob/master/README.md -[md-security]: https://markdown-it-py.readthedocs.io/en/latest/other.html -[md-performance]: https://markdown-it-py.readthedocs.io/en/latest/other.html -[md-plugins]: https://markdown-it-py.readthedocs.io/en/latest/plugins.html - diff --git a/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/RECORD b/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/RECORD deleted file mode 100644 index d71bbbff1e7464d1f85c9c2d488f081b24fd68c3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/RECORD +++ /dev/null @@ -1,142 +0,0 @@ -../../../bin/markdown-it,sha256=dpBVa_JM9lb-htsAmWj_QmVC-fxMsaBfUx4iCJQEsTE,238 -markdown_it/__init__.py,sha256=9v3vCD7XQJujcZLU2F14T8O88JJO93rZaUoKu7cocH8,113 -markdown_it/__pycache__/__init__.cpython-310.pyc,, -markdown_it/__pycache__/_compat.cpython-310.pyc,, -markdown_it/__pycache__/_punycode.cpython-310.pyc,, -markdown_it/__pycache__/main.cpython-310.pyc,, -markdown_it/__pycache__/parser_block.cpython-310.pyc,, -markdown_it/__pycache__/parser_core.cpython-310.pyc,, -markdown_it/__pycache__/parser_inline.cpython-310.pyc,, -markdown_it/__pycache__/renderer.cpython-310.pyc,, -markdown_it/__pycache__/ruler.cpython-310.pyc,, -markdown_it/__pycache__/token.cpython-310.pyc,, -markdown_it/__pycache__/tree.cpython-310.pyc,, -markdown_it/__pycache__/utils.cpython-310.pyc,, -markdown_it/_compat.py,sha256=mfhalPobHpl8uYt2V6SCOZq3HqaGlWP8MjICwDvS_xE,246 -markdown_it/_punycode.py,sha256=Y_m-fzc5Ey_Kw09MPNN5TUMnPXm2cACUZE_qwUkfFrM,2364 -markdown_it/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -markdown_it/cli/__pycache__/__init__.cpython-310.pyc,, -markdown_it/cli/__pycache__/parse.cpython-310.pyc,, -markdown_it/cli/parse.py,sha256=ZiTSx6t7nLk7rGAtIi0a02EB9sDGJn7YLjKKtufdwNA,2901 -markdown_it/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -markdown_it/common/__pycache__/__init__.cpython-310.pyc,, -markdown_it/common/__pycache__/entities.cpython-310.pyc,, -markdown_it/common/__pycache__/html_blocks.cpython-310.pyc,, -markdown_it/common/__pycache__/html_re.cpython-310.pyc,, -markdown_it/common/__pycache__/normalize_url.cpython-310.pyc,, -markdown_it/common/__pycache__/utils.cpython-310.pyc,, -markdown_it/common/entities.py,sha256=6ulEjBAWYH5lVobgpn5lChPYhPKkdPyVHpNT7O1_x90,156 -markdown_it/common/html_blocks.py,sha256=1cMBp6jIdXqCHvEs2mpJqVGqTuFi6ExL4VO754yksgU,932 -markdown_it/common/html_re.py,sha256=0q5QFkSnX_l7Ob37MDSj2UYT0onCPz_07hUod2K-a6A,929 -markdown_it/common/normalize_url.py,sha256=avOXnLd9xw5jU1q5PLftjAM9pvGx8l9QDEkmZSyrMgg,2568 -markdown_it/common/utils.py,sha256=l2ypUup7jVBwGZJb8M2lxllmTWq8wN5TMV40Mndqu1A,10728 -markdown_it/helpers/__init__.py,sha256=9W7GycpZcq2up1CdVcUpdN77i9Vl4N0CT3y3qMkTjY4,253 -markdown_it/helpers/__pycache__/__init__.cpython-310.pyc,, -markdown_it/helpers/__pycache__/parse_link_destination.cpython-310.pyc,, -markdown_it/helpers/__pycache__/parse_link_label.cpython-310.pyc,, -markdown_it/helpers/__pycache__/parse_link_title.cpython-310.pyc,, -markdown_it/helpers/parse_link_destination.py,sha256=whJoEo42RmgiVpyc2TmTb73nPK3L-DZuZqfE2f8B20Q,1977 -markdown_it/helpers/parse_link_label.py,sha256=HXAnGlAL-2Op6I-lfeuzBIJBjESCRZws6xKKS3lKFSg,1036 -markdown_it/helpers/parse_link_title.py,sha256=5h5YctAUPmaeBqAlCDzzB220-i64HxYyJ27ui-xMBn0,1425 -markdown_it/main.py,sha256=7BYAkBbDmgbdVGNrpAaWZ_-u2-DoYsOCtXWg5uv3Gvg,12772 -markdown_it/parser_block.py,sha256=1bvZMDIdfBYrRNHpZPzmX6W4GXXaFUDqq2iaxDWP-BM,3911 -markdown_it/parser_core.py,sha256=asPHnvL0sk2oqWC69VM8OAqKs-Dk--GRHhGwZonLm6s,1010 -markdown_it/parser_inline.py,sha256=0ZAiRx2GkQ35va1QRxDeDZvpa44wZBNnVFsGzCyj_Po,4997 -markdown_it/port.yaml,sha256=F6WvtDFxjuZ5o0NtmJO0P8SkG6OfONCj_ggKgx4iYWU,2446 -markdown_it/presets/__init__.py,sha256=tLc9od5iXBEkKR6wbXOVPC5S5vkGYK-67tQslhWu7LY,970 -markdown_it/presets/__pycache__/__init__.cpython-310.pyc,, -markdown_it/presets/__pycache__/commonmark.cpython-310.pyc,, -markdown_it/presets/__pycache__/default.cpython-310.pyc,, -markdown_it/presets/__pycache__/zero.cpython-310.pyc,, -markdown_it/presets/commonmark.py,sha256=pqVnOnMmCmiZWHvNnXx4E1FS8VV07jcqUW1EsHuv9rE,2868 -markdown_it/presets/default.py,sha256=TgqnjjHX6SWcElk5yjW1hoP8t2-ESZ0QyrmHprmIL18,1810 -markdown_it/presets/zero.py,sha256=2vETQRRW1v9Ug3TsdRhM9r6wgZRuTcvojwcqDcsZfkI,2112 -markdown_it/py.typed,sha256=8PjyZ1aVoQpRVvt71muvuq5qE-jTFZkK-GLHkhdebmc,26 -markdown_it/renderer.py,sha256=jZ62oK-y-qaz8o8TDxCUMvS5OavraJ2-1uaWVqrUCv0,9970 -markdown_it/ruler.py,sha256=J18Pru7u77kqJfnMC5oiFWgnw_58vkbfk5hWQBgED7s,9199 -markdown_it/rules_block/__init__.py,sha256=8su1tOxDw_IR9JSdgqfkGZSNdYpZZCBC9MMmFODdbmE,553 -markdown_it/rules_block/__pycache__/__init__.cpython-310.pyc,, -markdown_it/rules_block/__pycache__/blockquote.cpython-310.pyc,, -markdown_it/rules_block/__pycache__/code.cpython-310.pyc,, -markdown_it/rules_block/__pycache__/fence.cpython-310.pyc,, -markdown_it/rules_block/__pycache__/heading.cpython-310.pyc,, -markdown_it/rules_block/__pycache__/hr.cpython-310.pyc,, -markdown_it/rules_block/__pycache__/html_block.cpython-310.pyc,, -markdown_it/rules_block/__pycache__/lheading.cpython-310.pyc,, -markdown_it/rules_block/__pycache__/list.cpython-310.pyc,, -markdown_it/rules_block/__pycache__/paragraph.cpython-310.pyc,, -markdown_it/rules_block/__pycache__/reference.cpython-310.pyc,, -markdown_it/rules_block/__pycache__/state_block.cpython-310.pyc,, -markdown_it/rules_block/__pycache__/table.cpython-310.pyc,, -markdown_it/rules_block/blockquote.py,sha256=7uymS36dcrned3DsIaRcqcbFU1NlymhvsZpEXTD3_n8,8887 -markdown_it/rules_block/code.py,sha256=ASAnisg4hS2RhnP_7_1_pjx4NbFSYmrs6lHDgtHPXIo,859 -markdown_it/rules_block/fence.py,sha256=BJgU-PqZ4vAlCqGcrc8UtdLpJJyMeRWN-G-Op-zxrMc,2537 -markdown_it/rules_block/heading.py,sha256=e9NnvXLbY1bvowq_Pd4S-g6LbVg3iCx26qzwv3jLUyE,1746 -markdown_it/rules_block/hr.py,sha256=fPJ-tubFKjxJxhKPiTAxVP-_LHYbAq32iZ52J5sFxOU,1226 -markdown_it/rules_block/html_block.py,sha256=wA8pb34LtZr1BkIATgGKQBIGX5jQNOkwZl9UGEqvb5M,2721 -markdown_it/rules_block/lheading.py,sha256=fWoEuUo7S2svr5UMKmyQMkh0hheYAHg2gMM266Mogs4,2625 -markdown_it/rules_block/list.py,sha256=gIodkAJFyOIyKCZCj5lAlL7jIj5kAzrDb-K-2MFNplY,9668 -markdown_it/rules_block/paragraph.py,sha256=pQqTn8yYDI6_mWX-_m6PXY4wvDQB1nZ4dVUp3gKu1GA,1818 -markdown_it/rules_block/reference.py,sha256=qzR-KJ_60W8ZzuwYGLlO1bgHHVQP4qlYG4yFpOpNlsA,6168 -markdown_it/rules_block/state_block.py,sha256=HowsQyy5hGUibH4HRZWKfLIlXeDUnuWL7kpF0-rSwoM,8422 -markdown_it/rules_block/table.py,sha256=Zjkc0378QtfQzrhrNWC2kVYsGnOLaZyD3dXG1ugfX-s,6987 -markdown_it/rules_core/__init__.py,sha256=JZNOpLZ4i1vR56StidUa-A_As1XtbDwQR0iEErOXyOI,394 -markdown_it/rules_core/__pycache__/__init__.cpython-310.pyc,, -markdown_it/rules_core/__pycache__/block.cpython-310.pyc,, -markdown_it/rules_core/__pycache__/inline.cpython-310.pyc,, -markdown_it/rules_core/__pycache__/linkify.cpython-310.pyc,, -markdown_it/rules_core/__pycache__/normalize.cpython-310.pyc,, -markdown_it/rules_core/__pycache__/replacements.cpython-310.pyc,, -markdown_it/rules_core/__pycache__/smartquotes.cpython-310.pyc,, -markdown_it/rules_core/__pycache__/state_core.cpython-310.pyc,, -markdown_it/rules_core/__pycache__/text_join.cpython-310.pyc,, -markdown_it/rules_core/block.py,sha256=0_JY1CUy-H2OooFtIEZAACtuoGUMohgxo4Z6A_UinSg,372 -markdown_it/rules_core/inline.py,sha256=9oWmeBhJHE7x47oJcN9yp6UsAZtrEY_A-VmfoMvKld4,325 -markdown_it/rules_core/linkify.py,sha256=mjQqpk_lHLh2Nxw4UFaLxa47Fgi-OHnmDamlgXnhmv0,5141 -markdown_it/rules_core/normalize.py,sha256=qVkBO4elitPzyP_sQENho-ycUl8s4eNZ1zZrsR2AAgk,402 -markdown_it/rules_core/replacements.py,sha256=NHL9MOuEnPuMFPLDtTYDK9yj7F2FSleMr6bPro-ciaQ,3470 -markdown_it/rules_core/smartquotes.py,sha256=CtawEcTHYgzIWZwxIGs8e8oSKhm0B7th2305I3FNEc0,7443 -markdown_it/rules_core/state_core.py,sha256=HqWZCUr5fW7xG6jeQZDdO0hE9hxxyl3_-bawgOy57HY,570 -markdown_it/rules_core/text_join.py,sha256=JVuq_27LoI0IjJDmCXOuRiTs1rmSFhFUUjh6MdF_YCk,1172 -markdown_it/rules_inline/__init__.py,sha256=Zvl8P8V830vDhcQKEleLKZ_paC-ypTn7eWpmFa9yySQ,696 -markdown_it/rules_inline/__pycache__/__init__.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/autolink.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/backticks.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/balance_pairs.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/emphasis.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/entity.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/escape.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/fragments_join.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/html_inline.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/image.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/link.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/linkify.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/newline.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/state_inline.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/strikethrough.cpython-310.pyc,, -markdown_it/rules_inline/__pycache__/text.cpython-310.pyc,, -markdown_it/rules_inline/autolink.py,sha256=l4EY7OLzuda350cT6Du_dggEwcb96TvD7YPXf2H6P1M,2079 -markdown_it/rules_inline/backticks.py,sha256=J7bezjjNxiXlKqvHc0fJkHZwH7-2nBsXVjcKydk8E4M,2037 -markdown_it/rules_inline/balance_pairs.py,sha256=vifasmne02sNaBBwuZsA4yI02vmv1gvVN4qR-b9m62E,4851 -markdown_it/rules_inline/emphasis.py,sha256=7aDLZx0Jlekuvbu3uEUTDhJp00Z0Pj6g4C3-VLhI8Co,3123 -markdown_it/rules_inline/entity.py,sha256=CE8AIGMi5isEa24RNseo0wRmTTaj5YLbgTFdDmBesAU,1651 -markdown_it/rules_inline/escape.py,sha256=5DEa7O6ByUfXdXZudcF7JZwLxXG1njAuXIOUsNDPPqU,1658 -markdown_it/rules_inline/fragments_join.py,sha256=_3JbwWYJz74gRHeZk6T8edVJT2IVSsi7FfmJJlieQlA,1493 -markdown_it/rules_inline/html_inline.py,sha256=SBg6HR0HRqCdrkkec0dfOYuQdAqyfeLRFLeQggtgjvg,1130 -markdown_it/rules_inline/image.py,sha256=AMO7kls5c-C4A_S2rBiIlKD8BQ8vIgyWUX7vCpXy_Qs,4135 -markdown_it/rules_inline/link.py,sha256=wRdTMxjNYFiV6uouZMML3fIPvyFrPtMHWzJNBQcQLV0,4318 -markdown_it/rules_inline/linkify.py,sha256=gmHw59SsMQLImv6v1RkDY9lcQAmTN-UZ4GwRFiR8jlg,1704 -markdown_it/rules_inline/newline.py,sha256=LEIhBB_3PPLeCAgaC2naheMjW5N6b9UaAB6sh47Ckz8,1296 -markdown_it/rules_inline/state_inline.py,sha256=rXmMX0P0pCf-v-013YB24MBQxMn2dJhoSZQMNYAZ8HQ,5101 -markdown_it/rules_inline/strikethrough.py,sha256=pwcPlyhkh5pqFVxRCSrdW5dNCIOtU4eDit7TVDTPIVA,3214 -markdown_it/rules_inline/text.py,sha256=GwmMVZziAmhj48l9VaXAXwzzUKDkhaA14thv-TCaS2M,901 -markdown_it/token.py,sha256=NEvuoYAeDh8_6zT6fukzdoncusVOjyUKw2zjsNgZmp4,6439 -markdown_it/tree.py,sha256=YxSqq3qSuhHHm1nQpPUhyDA4VIWHu_G_92bKdUcXXGM,11421 -markdown_it/utils.py,sha256=zPoQ8lhvxtJfg6iNSim0LcnAL0Y4XnV3G4DIIKmL8OU,5365 -markdown_it_py-3.0.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -markdown_it_py-3.0.0.dist-info/LICENSE,sha256=SiJg1uLND1oVGh6G2_59PtVSseK-q_mUHBulxJy85IQ,1078 -markdown_it_py-3.0.0.dist-info/LICENSE.markdown-it,sha256=eSxIxahJoV_fnjfovPnm0d0TsytGxkKnSKCkapkZ1HM,1073 -markdown_it_py-3.0.0.dist-info/METADATA,sha256=0-kME4KQNSCGPzfSEFgQc8MhUd5cmG-LO007BFk3_fw,6940 -markdown_it_py-3.0.0.dist-info/RECORD,, -markdown_it_py-3.0.0.dist-info/WHEEL,sha256=4TfKIB_xu-04bc2iKz6_zFt-gEFEEDU_31HGhqzOCE8,81 -markdown_it_py-3.0.0.dist-info/entry_points.txt,sha256=T81l7fHQ3pllpQ4wUtQK6a8g_p6wxQbnjKVHCk2WMG4,58 diff --git a/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/WHEEL b/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/WHEEL deleted file mode 100644 index 668ba4d0151c5c76ed6e758061daa8c1b0bf5d21..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/WHEEL +++ /dev/null @@ -1,4 +0,0 @@ -Wheel-Version: 1.0 -Generator: flit 3.7.1 -Root-Is-Purelib: true -Tag: py3-none-any diff --git a/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/entry_points.txt b/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/entry_points.txt deleted file mode 100644 index 7d829cd792a5d754844f433c6a8dd499564fdcbf..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/markdown_it_py-3.0.0.dist-info/entry_points.txt +++ /dev/null @@ -1,3 +0,0 @@ -[console_scripts] -markdown-it=markdown_it.cli.parse:main - diff --git a/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/LICENSE b/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/LICENSE deleted file mode 100644 index e06d2081865a766a8668acc12878f98b27fc9ea0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/LICENSE +++ /dev/null @@ -1,202 +0,0 @@ -Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - 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 - 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 - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "{}" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright {yyyy} {name of copyright owner} - - 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. - diff --git a/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/METADATA b/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/METADATA deleted file mode 100644 index 89b6c7956b394ff46b874e8afd0425a20721f04c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/METADATA +++ /dev/null @@ -1,157 +0,0 @@ -Metadata-Version: 2.1 -Name: matrix-client -Version: 0.4.0 -Summary: Client-Server SDK for Matrix -Home-page: https://github.com/matrix-org/matrix-python-sdk -Author: The Matrix.org Team -Author-email: team@matrix.org -License: Apache License, Version 2.0 -Keywords: chat sdk matrix matrix.org -Platform: UNKNOWN -Classifier: Development Status :: 3 - Alpha -Classifier: Intended Audience :: Developers -Classifier: License :: OSI Approved :: Apache Software License -Classifier: Programming Language :: Python :: 2 -Classifier: Programming Language :: Python :: 3 -Classifier: Topic :: Communications :: Chat -Classifier: Topic :: Communications :: Conferencing -Description-Content-Type: text/x-rst -License-File: LICENSE -Requires-Dist: requests (~=2.22) -Requires-Dist: urllib3 (~=1.21) -Provides-Extra: doc -Requires-Dist: Sphinx (==1.*,>=1.7.6) ; extra == 'doc' -Requires-Dist: sphinx-rtd-theme (==0.1.*,>=0.1.9) ; extra == 'doc' -Requires-Dist: sphinxcontrib-napoleon (==0.5.*,>=0.5.3) ; extra == 'doc' -Provides-Extra: e2e -Requires-Dist: python-olm (~=3.1) ; extra == 'e2e' -Requires-Dist: canonicaljson (~=1.1) ; extra == 'e2e' -Provides-Extra: test -Requires-Dist: pytest (<6.0.0,>=4.6) ; extra == 'test' -Requires-Dist: responses (==0.10.*,>=0.10.6) ; extra == 'test' - -Matrix Client SDK for Python -============================ - -.. image:: https://img.shields.io/pypi/v/matrix-client.svg?maxAge=600 - :target: https://pypi.python.org/pypi/matrix-client - :alt: Latest Version -.. image:: https://travis-ci.org/matrix-org/matrix-python-sdk.svg?branch=master - :target: https://travis-ci.org/matrix-org/matrix-python-sdk - :alt: Travis-CI Results -.. image:: https://coveralls.io/repos/github/matrix-org/matrix-python-sdk/badge.svg?branch=master - :target: https://coveralls.io/github/matrix-org/matrix-python-sdk?branch=master - :alt: coveralls.io Results -.. image:: https://img.shields.io/matrix/matrix-python-sdk:matrix.org - :target: https://matrix.to/#/%23matrix-python-sdk:matrix.org - :alt: Matrix chatroom -.. image:: https://img.shields.io/badge/docs-stable-blue - :target: https://matrix-org.github.io/matrix-python-sdk/ - :alt: Documentation - - -Matrix client-server SDK for Python 2.7 and 3.4+ - -Project Status --------------- - -We strongly recommend using the `matrix-nio`_ library rather than this -sdk. It is both more featureful and more actively maintained. - -This sdk is currently lightly maintained without any person ultimately -responsible for the project. Pull-requests **may** be reviewed, but no -new-features or bug-fixes are being actively developed. For more info -or to volunteer to help, please see -https://github.com/matrix-org/matrix-python-sdk/issues/279 or come -chat in `#matrix-python-sdk:matrix.org`_. - -.. _`matrix-nio`: https://github.com/poljar/matrix-nio -.. _`#matrix-python-sdk:matrix.org`: https://matrix.to/#/%23matrix-python-sdk:matrix.org - -Installation -============ -Stable release --------------- -Install with pip from pypi. This will install all necessary dependencies as well. - -.. code:: shell - - pip install matrix_client - -Development version -------------------- -Install using ``setup.py`` in root project directory. This will also install all -needed dependencies. - -.. code:: shell - - git clone https://github.com/matrix-org/matrix-python-sdk.git - cd matrix-python-sdk - python setup.py install - -Usage -===== -The SDK provides 2 layers of interaction. The low-level layer just wraps the -raw HTTP API calls. The high-level layer wraps the low-level layer and provides -an object model to perform actions on. - -Client: - -.. code:: python - - from matrix_client.client import MatrixClient - - client = MatrixClient("http://localhost:8008") - - # New user - token = client.register_with_password(username="foobar", password="monkey") - - # Existing user - token = client.login(username="foobar", password="monkey") - - room = client.create_room("my_room_alias") - room.send_text("Hello!") - - -API: - -.. code:: python - - from matrix_client.api import MatrixHttpApi - - matrix = MatrixHttpApi("https://matrix.org", token="some_token") - response = matrix.send_message("!roomid:matrix.org", "Hello!") - - -Structure -========= -The SDK is split into two modules: ``api`` and ``client``. - -API ---- -This contains the raw HTTP API calls and has minimal business logic. You can -set the access token (``token``) to use for requests as well as set a custom -transaction ID (``txn_id``) which will be incremented for each request. - -Client ------- -This encapsulates the API module and provides object models such as ``Room``. - -Samples -======= -A collection of samples are included, written in Python 3. - -You can either install the SDK, or run the sample like this: - -.. code:: shell - - PYTHONPATH=. python samples/samplename.py - -Building the Documentation -========================== - -The documentation can be built by installing ``sphinx`` and ``sphinx_rtd_theme``. - -Simple run ``make`` inside ``docs`` which will list the avaliable output formats. - - diff --git a/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/RECORD b/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/RECORD deleted file mode 100644 index ac3e90efb56901ab8e3616f821729d58f89d3f58..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/RECORD +++ /dev/null @@ -1,37 +0,0 @@ -matrix_client-0.4.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -matrix_client-0.4.0.dist-info/LICENSE,sha256=y16Ofl9KOYjhBjwULGDcLfdWBfTEZRXnduOspt-XbhQ,11325 -matrix_client-0.4.0.dist-info/METADATA,sha256=u50tqp4kU07AzJgHE12jW-BrFqEnWS049T7dOKxuIdk,4956 -matrix_client-0.4.0.dist-info/RECORD,, -matrix_client-0.4.0.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -matrix_client-0.4.0.dist-info/WHEEL,sha256=Z-nyYpwrcSqxfdux5Mbn_DQ525iP7J2DG3JgGvOYyTQ,110 -matrix_client-0.4.0.dist-info/top_level.txt,sha256=dmjf3uMSj0bJanvGvcANIgJLTPFC3k4E4Ifg5k-0cyU,19 -matrix_client/__init__.py,sha256=sdVPVNtj0pA7JoOfZEx2FdWzUoH-qVaxNq6YLt1EChc,656 -matrix_client/__pycache__/__init__.cpython-310.pyc,, -matrix_client/__pycache__/api.cpython-310.pyc,, -matrix_client/__pycache__/checks.cpython-310.pyc,, -matrix_client/__pycache__/client.cpython-310.pyc,, -matrix_client/__pycache__/errors.cpython-310.pyc,, -matrix_client/__pycache__/room.cpython-310.pyc,, -matrix_client/__pycache__/user.cpython-310.pyc,, -matrix_client/api.py,sha256=F3hi-A3yPxj-hCo9topL8F2geYxOtw5UadRxajseN1A,39415 -matrix_client/checks.py,sha256=M_BsgQNTg3hb3yX4obEsxvTjcrFJ1DhnEHtfPIau-gM,1054 -matrix_client/client.py,sha256=jlQiDTpjc-PWfq5IxPmURL7a-L4NWol9jbiOl5PvRQM,24926 -matrix_client/crypto/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -matrix_client/crypto/__pycache__/__init__.cpython-310.pyc,, -matrix_client/crypto/__pycache__/olm_device.cpython-310.pyc,, -matrix_client/crypto/__pycache__/one_time_keys.cpython-310.pyc,, -matrix_client/crypto/olm_device.py,sha256=bjo7R0iWiRYHv9WvHUdrTjUWr2GAa1aMxPyvPET-lQ8,8538 -matrix_client/crypto/one_time_keys.py,sha256=0cDvhMX395C8A8xROnorKf6CMX_MCFLaPXEfCLfi5Eo,1531 -matrix_client/errors.py,sha256=ruFMLa0_zV1xBrLKqyK-z7BJv1KwuLaMFJ1IPZ1cN1Y,1742 -matrix_client/room.py,sha256=Gxu85KC-LCGUJvGuV9a00aPp3L_HyDMvgRCzjlPEYSM,24910 -matrix_client/user.py,sha256=kZ0UycJY0l7dCCZkwM5FRPnbwBgcQEupJGIpL_db77o,2527 -test/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 -test/__pycache__/__init__.cpython-310.pyc,, -test/__pycache__/api_test.cpython-310.pyc,, -test/__pycache__/client_test.cpython-310.pyc,, -test/__pycache__/response_examples.cpython-310.pyc,, -test/__pycache__/user_test.cpython-310.pyc,, -test/api_test.py,sha256=d0LWb7-YNcF2O_q-CFYutJbsvC-PLS1DT8Ii85tNQ4s,18737 -test/client_test.py,sha256=rJ1CneoAzCvfmhT1GlCQjFB4Sd9kF3CayGCOnc3KDR8,19047 -test/response_examples.py,sha256=pUxyy5gEa7qXwf2KpzW3JKmzLkdWaNgr5coMfYKQf_U,5976 -test/user_test.py,sha256=OPeDnDSAaoGfa-42W1DdHewO21AMBY5yv6TcQ1ocppU,1636 diff --git a/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/REQUESTED deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/WHEEL b/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/WHEEL deleted file mode 100644 index 01b8fc7d4a10cb8b4f1d21f11d3398d07d6b3478..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/WHEEL +++ /dev/null @@ -1,6 +0,0 @@ -Wheel-Version: 1.0 -Generator: bdist_wheel (0.36.2) -Root-Is-Purelib: true -Tag: py2-none-any -Tag: py3-none-any - diff --git a/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/top_level.txt deleted file mode 100644 index ba4c019187a9d1fa0080bd217769b2a8598ebe58..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client-0.4.0.dist-info/top_level.txt +++ /dev/null @@ -1,2 +0,0 @@ -matrix_client -test diff --git a/venv/lib/python3.10/site-packages/matrix_client/__init__.py b/venv/lib/python3.10/site-packages/matrix_client/__init__.py deleted file mode 100644 index 7b18c407159a875ba6a2451c2d34d602789d918d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2015 OpenMarket Ltd -# Copyright 2018 Adam Beckmeyer -# -# 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. -__version__ = "0.4.0" diff --git a/venv/lib/python3.10/site-packages/matrix_client/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/matrix_client/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index a98cfd49a8e5322443f6c43ef6a3a8baf3a4a3a7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/matrix_client/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/matrix_client/__pycache__/api.cpython-310.pyc b/venv/lib/python3.10/site-packages/matrix_client/__pycache__/api.cpython-310.pyc deleted file mode 100644 index 18e0a1148df5f3e1f4db8bc2603b86e7529aaa32..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/matrix_client/__pycache__/api.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/matrix_client/__pycache__/checks.cpython-310.pyc b/venv/lib/python3.10/site-packages/matrix_client/__pycache__/checks.cpython-310.pyc deleted file mode 100644 index 92b2db4102f87f0835165b3cdb3b6f92bfd976d2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/matrix_client/__pycache__/checks.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/matrix_client/__pycache__/client.cpython-310.pyc b/venv/lib/python3.10/site-packages/matrix_client/__pycache__/client.cpython-310.pyc deleted file mode 100644 index 5e4f97f16e8d57942e669f79e2690e934915d40a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/matrix_client/__pycache__/client.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/matrix_client/__pycache__/errors.cpython-310.pyc b/venv/lib/python3.10/site-packages/matrix_client/__pycache__/errors.cpython-310.pyc deleted file mode 100644 index 5521f9dbfcfaba789e4905e94bb2d916c744be23..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/matrix_client/__pycache__/errors.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/matrix_client/__pycache__/room.cpython-310.pyc b/venv/lib/python3.10/site-packages/matrix_client/__pycache__/room.cpython-310.pyc deleted file mode 100644 index cdbcbb609e410a55904c38d4ad51d869e12fb317..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/matrix_client/__pycache__/room.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/matrix_client/__pycache__/user.cpython-310.pyc b/venv/lib/python3.10/site-packages/matrix_client/__pycache__/user.cpython-310.pyc deleted file mode 100644 index 4c4f5ec6a486d3be0f2f748fd34095ea1a87378a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/matrix_client/__pycache__/user.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/matrix_client/api.py b/venv/lib/python3.10/site-packages/matrix_client/api.py deleted file mode 100644 index aae52f999bef2ed32984d05cf27367671531b457..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client/api.py +++ /dev/null @@ -1,1098 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2015 OpenMarket Ltd -# Copyright 2017, 2018 Adam Beckmeyer -# -# 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 warnings -from requests import Session, RequestException -from time import time, sleep -from .__init__ import __version__ -from .errors import MatrixError, MatrixRequestError, MatrixHttpLibError -from urllib3.util import parse_url -from urllib3.exceptions import LocationParseError - -try: - from urllib import quote -except ImportError: - from urllib.parse import quote - -MATRIX_V2_API_PATH = "/_matrix/client/r0" - - -class MatrixHttpApi(object): - """Contains all raw Matrix HTTP Client-Server API calls. - - For room and sync handling, consider using MatrixClient. - - Args: - base_url (str): The home server URL e.g. 'http://localhost:8008' - token (str): Optional. The client's access token. - identity (str): Optional. The mxid to act as (For application services only). - default_429_wait_ms (int): Optional. Time in millseconds to wait before retrying - a request when server returns a HTTP 429 - response without a 'retry_after_ms' key. - use_authorization_header (bool): Optional. Use Authorization header instead - ` of access_token query parameter. - - Examples: - Create a client and send a message:: - - matrix = MatrixHttpApi("https://matrix.org", token="foobar") - response = matrix.sync() - response = matrix.send_message("!roomid:matrix.org", "Hello!") - """ - - def __init__( - self, base_url, token=None, identity=None, - default_429_wait_ms=5000, - use_authorization_header=True - ): - try: - scheme, auth, host, port, path, query, fragment = parse_url(base_url) - except LocationParseError: - raise MatrixError("Invalid homeserver url %s" % base_url) - if not scheme: - raise MatrixError("No scheme in homeserver url %s" % base_url) - self._base_url = base_url - - self.token = token - self.identity = identity - self.txn_id = 0 - self.validate_cert = True - self.session = Session() - self.default_429_wait_ms = default_429_wait_ms - self.use_authorization_header = use_authorization_header - - def initial_sync(self, limit=1): - """ - .. warning:: - - Deprecated. Use sync instead. - - Perform /initialSync. - - Args: - limit (int): The limit= param to provide. - """ - warnings.warn("initial_sync is deprecated. Use sync instead.", DeprecationWarning) - return self._send("GET", "/initialSync", query_params={"limit": limit}) - - def sync(self, since=None, timeout_ms=30000, filter=None, - full_state=None, set_presence=None): - """ Perform a sync request. - - Args: - since (str): Optional. A token which specifies where to continue a sync from. - timeout_ms (int): Optional. The time in milliseconds to wait. - filter (int|str): Either a Filter ID or a JSON string. - full_state (bool): Return the full state for every room the user has joined - Defaults to false. - set_presence (str): Should the client be marked as "online" or" offline" - """ - - request = { - # non-integer timeouts appear to cause issues - "timeout": int(timeout_ms) - } - - if since: - request["since"] = since - - if filter: - request["filter"] = filter - - if full_state: - request["full_state"] = json.dumps(full_state) - - if set_presence: - request["set_presence"] = set_presence - - return self._send("GET", "/sync", query_params=request, - api_path=MATRIX_V2_API_PATH) - - def validate_certificate(self, valid): - self.validate_cert = valid - - def register(self, auth_body=None, kind="user", bind_email=None, - username=None, password=None, device_id=None, - initial_device_display_name=None, inhibit_login=None): - """Performs /register. - - Args: - auth_body (dict): Authentication Params. - kind (str): Specify kind of account to register. Can be 'guest' or 'user'. - bind_email (bool): Whether to use email in registration and authentication. - username (str): The localpart of a Matrix ID. - password (str): The desired password of the account. - device_id (str): ID of the client device. - initial_device_display_name (str): Display name to be assigned. - inhibit_login (bool): Whether to login after registration. Defaults to false. - """ - content = {} - content["kind"] = kind - if auth_body: - content["auth"] = auth_body - if username: - content["username"] = username - if password: - content["password"] = password - if device_id: - content["device_id"] = device_id - if initial_device_display_name: - content["initial_device_display_name"] = \ - initial_device_display_name - if bind_email: - content["bind_email"] = bind_email - if inhibit_login: - content["inhibit_login"] = inhibit_login - return self._send( - "POST", - "/register", - content=content, - query_params={'kind': kind} - ) - - def login(self, login_type, **kwargs): - """Perform /login. - - Args: - login_type (str): The value for the 'type' key. - **kwargs: Additional key/values to add to the JSON submitted. - """ - content = { - "type": login_type - } - for key in kwargs: - if kwargs[key]: - content[key] = kwargs[key] - - return self._send("POST", "/login", content) - - def logout(self): - """Perform /logout. - """ - return self._send("POST", "/logout") - - def logout_all(self): - """Perform /logout/all.""" - return self._send("POST", "/logout/all") - - def create_room( - self, - alias=None, - name=None, - is_public=False, - invitees=None, - federate=None - ): - """Perform /createRoom. - - Args: - alias (str): Optional. The room alias name to set for this room. - name (str): Optional. Name for new room. - is_public (bool): Optional. The public/private visibility. - invitees (list): Optional. The list of user IDs to invite. - federate (bool): Optional. Сan a room be federated. - Default to True. - """ - content = { - "visibility": "public" if is_public else "private" - } - if alias: - content["room_alias_name"] = alias - if invitees: - content["invite"] = invitees - if name: - content["name"] = name - if federate is not None: - content["creation_content"] = {'m.federate': federate} - return self._send("POST", "/createRoom", content) - - def join_room(self, room_id_or_alias): - """Performs /join/$room_id - - Args: - room_id_or_alias (str): The room ID or room alias to join. - """ - if not room_id_or_alias: - raise MatrixError("No alias or room ID to join.") - - path = "/join/%s" % quote(room_id_or_alias) - - return self._send("POST", path) - - def event_stream(self, from_token, timeout=30000): - """ Deprecated. Use sync instead. - Performs /events - - Args: - from_token (str): The 'from' query parameter. - timeout (int): Optional. The 'timeout' query parameter. - """ - warnings.warn("event_stream is deprecated. Use sync instead.", - DeprecationWarning) - path = "/events" - return self._send( - "GET", path, query_params={ - "timeout": timeout, - "from": from_token - } - ) - - def send_state_event(self, room_id, event_type, content, state_key="", - timestamp=None): - """Perform PUT /rooms/$room_id/state/$event_type - - Args: - room_id(str): The room ID to send the state event in. - event_type(str): The state event type to send. - content(dict): The JSON content to send. - state_key(str): Optional. The state key for the event. - timestamp (int): Set origin_server_ts (For application services only) - """ - path = "/rooms/%s/state/%s" % ( - quote(room_id), quote(event_type), - ) - if state_key: - path += "/%s" % (quote(state_key)) - params = {} - if timestamp: - params["ts"] = timestamp - return self._send("PUT", path, content, query_params=params) - - def get_state_event(self, room_id, event_type): - """Perform GET /rooms/$room_id/state/$event_type - - Args: - room_id(str): The room ID. - event_type (str): The type of the event. - - Raises: - MatrixRequestError(code=404) if the state event is not found. - """ - return self._send("GET", "/rooms/{}/state/{}".format(quote(room_id), event_type)) - - def send_message_event(self, room_id, event_type, content, txn_id=None, - timestamp=None): - """Perform PUT /rooms/$room_id/send/$event_type - - Args: - room_id (str): The room ID to send the message event in. - event_type (str): The event type to send. - content (dict): The JSON content to send. - txn_id (int): Optional. The transaction ID to use. - timestamp (int): Set origin_server_ts (For application services only) - """ - if not txn_id: - txn_id = self._make_txn_id() - - path = "/rooms/%s/send/%s/%s" % ( - quote(room_id), quote(event_type), quote(str(txn_id)), - ) - params = {} - if timestamp: - params["ts"] = timestamp - return self._send("PUT", path, content, query_params=params) - - def redact_event(self, room_id, event_id, reason=None, txn_id=None, timestamp=None): - """Perform PUT /rooms/$room_id/redact/$event_id/$txn_id/ - - Args: - room_id(str): The room ID to redact the message event in. - event_id(str): The event id to redact. - reason (str): Optional. The reason the message was redacted. - txn_id(int): Optional. The transaction ID to use. - timestamp(int): Optional. Set origin_server_ts (For application services only) - """ - if not txn_id: - txn_id = self._make_txn_id() - - path = '/rooms/%s/redact/%s/%s' % ( - room_id, event_id, txn_id - ) - content = {} - if reason: - content['reason'] = reason - params = {} - if timestamp: - params["ts"] = timestamp - return self._send("PUT", path, content, query_params=params) - - # content_type can be a image,audio or video - # extra information should be supplied, see - # https://matrix.org/docs/spec/r0.0.1/client_server.html - def send_content(self, room_id, item_url, item_name, msg_type, - extra_information=None, timestamp=None): - if extra_information is None: - extra_information = {} - - content_pack = { - "url": item_url, - "msgtype": msg_type, - "body": item_name, - "info": extra_information - } - return self.send_message_event(room_id, "m.room.message", content_pack, - timestamp=timestamp) - - # http://matrix.org/docs/spec/client_server/r0.2.0.html#m-location - def send_location(self, room_id, geo_uri, name, thumb_url=None, thumb_info=None, - timestamp=None): - """Send m.location message event - - Args: - room_id (str): The room ID to send the event in. - geo_uri (str): The geo uri representing the location. - name (str): Description for the location. - thumb_url (str): URL to the thumbnail of the location. - thumb_info (dict): Metadata about the thumbnail, type ImageInfo. - timestamp (int): Set origin_server_ts (For application services only) - """ - content_pack = { - "geo_uri": geo_uri, - "msgtype": "m.location", - "body": name, - } - if thumb_url: - content_pack["thumbnail_url"] = thumb_url - if thumb_info: - content_pack["thumbnail_info"] = thumb_info - - return self.send_message_event(room_id, "m.room.message", content_pack, - timestamp=timestamp) - - def send_message(self, room_id, text_content, msgtype="m.text", timestamp=None): - """Perform PUT /rooms/$room_id/send/m.room.message - - Args: - room_id (str): The room ID to send the event in. - text_content (str): The m.text body to send. - timestamp (int): Set origin_server_ts (For application services only) - """ - return self.send_message_event( - room_id, "m.room.message", - self.get_text_body(text_content, msgtype), - timestamp=timestamp - ) - - def send_emote(self, room_id, text_content, timestamp=None): - """Perform PUT /rooms/$room_id/send/m.room.message with m.emote msgtype - - Args: - room_id (str): The room ID to send the event in. - text_content (str): The m.emote body to send. - timestamp (int): Set origin_server_ts (For application services only) - """ - return self.send_message_event( - room_id, "m.room.message", - self.get_emote_body(text_content), - timestamp=timestamp - ) - - def send_notice(self, room_id, text_content, timestamp=None): - """Perform PUT /rooms/$room_id/send/m.room.message with m.notice msgtype - - Args: - room_id (str): The room ID to send the event in. - text_content (str): The m.notice body to send. - timestamp (int): Set origin_server_ts (For application services only) - """ - body = { - "msgtype": "m.notice", - "body": text_content - } - return self.send_message_event(room_id, "m.room.message", body, - timestamp=timestamp) - - def get_room_messages(self, room_id, token, direction, limit=10, to=None): - """Perform GET /rooms/{roomId}/messages. - - Args: - room_id (str): The room's id. - token (str): The token to start returning events from. - direction (str): The direction to return events from. One of: ["b", "f"]. - limit (int): The maximum number of events to return. - to (str): The token to stop returning events at. - """ - query = { - "roomId": room_id, - "from": token, - "dir": direction, - "limit": limit, - } - - if to: - query["to"] = to - - return self._send("GET", "/rooms/{}/messages".format(quote(room_id)), - query_params=query, api_path="/_matrix/client/r0") - - def get_room_name(self, room_id): - """Perform GET /rooms/$room_id/state/m.room.name - Args: - room_id(str): The room ID - """ - return self.get_state_event(room_id, "m.room.name") - - def set_room_name(self, room_id, name, timestamp=None): - """Perform PUT /rooms/$room_id/state/m.room.name - Args: - room_id (str): The room ID - name (str): The new room name - timestamp (int): Set origin_server_ts (For application services only) - """ - body = { - "name": name - } - return self.send_state_event(room_id, "m.room.name", body, timestamp=timestamp) - - def get_room_topic(self, room_id): - """Perform GET /rooms/$room_id/state/m.room.topic - Args: - room_id (str): The room ID - """ - return self.get_state_event(room_id, "m.room.topic") - - def set_room_topic(self, room_id, topic, timestamp=None): - """Perform PUT /rooms/$room_id/state/m.room.topic - Args: - room_id (str): The room ID - topic (str): The new room topic - timestamp (int): Set origin_server_ts (For application services only) - """ - body = { - "topic": topic - } - return self.send_state_event(room_id, "m.room.topic", body, timestamp=timestamp) - - def get_power_levels(self, room_id): - """Perform GET /rooms/$room_id/state/m.room.power_levels - - Args: - room_id(str): The room ID - """ - return self.get_state_event(room_id, "m.room.power_levels") - - def set_power_levels(self, room_id, content): - """Perform PUT /rooms/$room_id/state/m.room.power_levels - - Note that any power levels which are not explicitly specified - in the content arg are reset to default values. - - Args: - room_id (str): The room ID - content (dict): The JSON content to send. See example content below. - - Example:: - - api = MatrixHttpApi("http://example.com", token="foobar") - api.set_power_levels("!exampleroom:example.com", - { - "ban": 50, # defaults to 50 if unspecified - "events": { - "m.room.name": 100, # must have PL 100 to change room name - "m.room.power_levels": 100 # must have PL 100 to change PLs - }, - "events_default": 0, # defaults to 0 - "invite": 50, # defaults to 50 - "kick": 50, # defaults to 50 - "redact": 50, # defaults to 50 - "state_default": 50, # defaults to 50 if m.room.power_levels exists - "users": { - "@someguy:example.com": 100 # defaults to 0 - }, - "users_default": 0 # defaults to 0 - } - ) - """ - # Synapse returns M_UNKNOWN if body['events'] is omitted, - # as of 2016-10-31 - if "events" not in content: - content["events"] = {} - - return self.send_state_event(room_id, "m.room.power_levels", content) - - def leave_room(self, room_id): - """Perform POST /rooms/$room_id/leave - - Args: - room_id (str): The room ID - """ - return self._send("POST", "/rooms/" + room_id + "/leave", {}) - - def forget_room(self, room_id): - """Perform POST /rooms/$room_id/forget - - Args: - room_id(str): The room ID - """ - return self._send("POST", "/rooms/" + room_id + "/forget", content={}) - - def invite_user(self, room_id, user_id): - """Perform POST /rooms/$room_id/invite - - Args: - room_id (str): The room ID - user_id (str): The user ID of the invitee - """ - body = { - "user_id": user_id - } - return self._send("POST", "/rooms/" + room_id + "/invite", body) - - def kick_user(self, room_id, user_id, reason=""): - """Calls set_membership with membership="leave" for the user_id provided - """ - self.set_membership(room_id, user_id, "leave", reason) - - def get_membership(self, room_id, user_id): - """Perform GET /rooms/$room_id/state/m.room.member/$user_id - - Args: - room_id (str): The room ID - user_id (str): The user ID - """ - return self._send( - "GET", - "/rooms/%s/state/m.room.member/%s" % (room_id, user_id) - ) - - def set_membership(self, room_id, user_id, membership, reason="", profile=None, - timestamp=None): - """Perform PUT /rooms/$room_id/state/m.room.member/$user_id - - Args: - room_id (str): The room ID - user_id (str): The user ID - membership (str): New membership value - reason (str): The reason - timestamp (int): Set origin_server_ts (For application services only) - """ - if profile is None: - profile = {} - body = { - "membership": membership, - "reason": reason - } - if 'displayname' in profile: - body["displayname"] = profile["displayname"] - if 'avatar_url' in profile: - body["avatar_url"] = profile["avatar_url"] - - return self.send_state_event(room_id, "m.room.member", body, state_key=user_id, - timestamp=timestamp) - - def ban_user(self, room_id, user_id, reason=""): - """Perform POST /rooms/$room_id/ban - - Args: - room_id (str): The room ID - user_id (str): The user ID of the banee(sic) - reason (str): The reason for this ban - """ - body = { - "user_id": user_id, - "reason": reason - } - return self._send("POST", "/rooms/" + room_id + "/ban", body) - - def unban_user(self, room_id, user_id): - """Perform POST /rooms/$room_id/unban - - Args: - room_id (str): The room ID - user_id (str): The user ID of the banee(sic) - """ - body = { - "user_id": user_id - } - return self._send("POST", "/rooms/" + room_id + "/unban", body) - - def get_user_tags(self, user_id, room_id): - return self._send( - "GET", - "/user/%s/rooms/%s/tags" % (user_id, room_id), - ) - - def remove_user_tag(self, user_id, room_id, tag): - return self._send( - "DELETE", - "/user/%s/rooms/%s/tags/%s" % (user_id, room_id, tag), - ) - - def add_user_tag(self, user_id, room_id, tag, order=None, body=None): - if body: - pass - elif order: - body = {"order": order} - else: - body = {} - return self._send( - "PUT", - "/user/%s/rooms/%s/tags/%s" % (user_id, room_id, tag), - body, - ) - - def set_account_data(self, user_id, type, account_data): - return self._send( - "PUT", - "/user/%s/account_data/%s" % (user_id, type), - account_data, - ) - - def set_room_account_data(self, user_id, room_id, type, account_data): - return self._send( - "PUT", - "/user/%s/rooms/%s/account_data/%s" % (user_id, room_id, type), - account_data - ) - - def get_room_state(self, room_id): - """Perform GET /rooms/$room_id/state - - Args: - room_id (str): The room ID - """ - return self._send("GET", "/rooms/" + room_id + "/state") - - def get_text_body(self, text, msgtype="m.text"): - return { - "msgtype": msgtype, - "body": text - } - - def get_emote_body(self, text): - return { - "msgtype": "m.emote", - "body": text - } - - def get_filter(self, user_id, filter_id): - return self._send("GET", "/user/{userId}/filter/{filterId}" - .format(userId=user_id, filterId=filter_id)) - - def create_filter(self, user_id, filter_params): - return self._send("POST", - "/user/{userId}/filter".format(userId=user_id), - filter_params) - - def _send(self, method, path, content=None, query_params=None, headers=None, - api_path=MATRIX_V2_API_PATH, return_json=True): - if query_params is None: - query_params = {} - if headers is None: - headers = {} - - if "User-Agent" not in headers: - headers["User-Agent"] = "matrix-python-sdk/%s" % __version__ - - method = method.upper() - if method not in ["GET", "PUT", "DELETE", "POST"]: - raise MatrixError("Unsupported HTTP method: %s" % method) - - if "Content-Type" not in headers: - headers["Content-Type"] = "application/json" - - if self.use_authorization_header: - headers["Authorization"] = 'Bearer %s' % self.token - else: - query_params["access_token"] = self.token - - if self.identity: - query_params["user_id"] = self.identity - - endpoint = self._base_url + api_path + path - - if headers["Content-Type"] == "application/json" and content is not None: - content = json.dumps(content) - - while True: - try: - response = self.session.request( - method, endpoint, - params=query_params, - data=content, - headers=headers, - verify=self.validate_cert - ) - except RequestException as e: - raise MatrixHttpLibError(e, method, endpoint) - - if response.status_code == 429: - waittime = self.default_429_wait_ms / 1000 - try: - waittime = response.json()['retry_after_ms'] / 1000 - except KeyError: - try: - errordata = json.loads(response.json()['error']) - waittime = errordata['retry_after_ms'] / 1000 - except KeyError: - pass - sleep(waittime) - else: - break - - if response.status_code < 200 or response.status_code >= 300: - raise MatrixRequestError( - code=response.status_code, content=response.text - ) - if return_json: - return response.json() - else: - return response - - def media_upload(self, content, content_type, filename=None): - query_params = {} - if filename is not None: - query_params['filename'] = filename - - return self._send( - "POST", "", - content=content, - headers={"Content-Type": content_type}, - api_path="/_matrix/media/r0/upload", - query_params=query_params - ) - - def get_display_name(self, user_id): - content = self._send("GET", "/profile/%s/displayname" % user_id) - return content.get('displayname', None) - - def set_display_name(self, user_id, display_name): - content = {"displayname": display_name} - return self._send("PUT", "/profile/%s/displayname" % user_id, content) - - def get_avatar_url(self, user_id): - content = self._send("GET", "/profile/%s/avatar_url" % user_id) - return content.get('avatar_url', None) - - def set_avatar_url(self, user_id, avatar_url): - content = {"avatar_url": avatar_url} - return self._send("PUT", "/profile/%s/avatar_url" % user_id, content) - - def get_download_url(self, mxcurl): - if mxcurl.startswith('mxc://'): - return self._base_url + "/_matrix/media/r0/download/" + mxcurl[6:] - else: - raise ValueError("MXC URL did not begin with 'mxc://'") - - def media_download(self, mxcurl, allow_remote=True): - """Download raw media from provided mxc URL. - - Args: - mxcurl (str): mxc media URL. - allow_remote (bool): indicates to the server that it should not - attempt to fetch the media if it is deemed remote. Defaults - to true if not provided. - """ - query_params = {} - if not allow_remote: - query_params["allow_remote"] = False - if mxcurl.startswith('mxc://'): - return self._send( - "GET", mxcurl[6:], - api_path="/_matrix/media/r0/download/", - query_params=query_params, - return_json=False - ) - else: - raise ValueError( - "MXC URL '%s' did not begin with 'mxc://'" % mxcurl - ) - - def get_thumbnail(self, mxcurl, width, height, method='scale', allow_remote=True): - """Download raw media thumbnail from provided mxc URL. - - Args: - mxcurl (str): mxc media URL - width (int): desired thumbnail width - height (int): desired thumbnail height - method (str): thumb creation method. Must be - in ['scale', 'crop']. Default 'scale'. - allow_remote (bool): indicates to the server that it should not - attempt to fetch the media if it is deemed remote. Defaults - to true if not provided. - """ - if method not in ['scale', 'crop']: - raise ValueError( - "Unsupported thumb method '%s'" % method - ) - query_params = { - "width": width, - "height": height, - "method": method - } - if not allow_remote: - query_params["allow_remote"] = False - if mxcurl.startswith('mxc://'): - return self._send( - "GET", mxcurl[6:], - query_params=query_params, - api_path="/_matrix/media/r0/thumbnail/", - return_json=False - ) - else: - raise ValueError( - "MXC URL '%s' did not begin with 'mxc://'" % mxcurl - ) - - def get_url_preview(self, url, ts=None): - """Get preview for URL. - - Args: - url (str): URL to get a preview - ts (double): The preferred point in time to return - a preview for. The server may return a newer - version if it does not have the requested - version available. - """ - params = {'url': url} - if ts: - params['ts'] = ts - return self._send( - "GET", "", - query_params=params, - api_path="/_matrix/media/r0/preview_url" - ) - - def get_room_id(self, room_alias): - """Get room id from its alias. - - Args: - room_alias (str): The room alias name. - - Returns: - Wanted room's id. - """ - content = self._send("GET", "/directory/room/{}".format(quote(room_alias))) - return content.get("room_id", None) - - def set_room_alias(self, room_id, room_alias): - """Set alias to room id - - Args: - room_id (str): The room id. - room_alias (str): The room wanted alias name. - """ - data = { - "room_id": room_id - } - - return self._send("PUT", "/directory/room/{}".format(quote(room_alias)), - content=data) - - def remove_room_alias(self, room_alias): - """Remove mapping of an alias - - Args: - room_alias(str): The alias to be removed. - - Raises: - MatrixRequestError - """ - return self._send("DELETE", "/directory/room/{}".format(quote(room_alias))) - - def get_room_members(self, room_id): - """Get the list of members for this room. - - Args: - room_id (str): The room to get the member events for. - """ - return self._send("GET", "/rooms/{}/members".format(quote(room_id))) - - def set_join_rule(self, room_id, join_rule): - """Set the rule for users wishing to join the room. - - Args: - room_id(str): The room to set the rules for. - join_rule(str): The chosen rule. One of: ["public", "knock", - "invite", "private"] - """ - content = { - "join_rule": join_rule - } - return self.send_state_event(room_id, "m.room.join_rules", content) - - def set_guest_access(self, room_id, guest_access): - """Set the guest access policy of the room. - - Args: - room_id(str): The room to set the rules for. - guest_access(str): Wether guests can join. One of: ["can_join", - "forbidden"] - """ - content = { - "guest_access": guest_access - } - return self.send_state_event(room_id, "m.room.guest_access", content) - - def get_devices(self): - """Gets information about all devices for the current user.""" - return self._send("GET", "/devices") - - def get_device(self, device_id): - """Gets information on a single device, by device id.""" - return self._send("GET", "/devices/%s" % device_id) - - def update_device_info(self, device_id, display_name): - """Update the display name of a device. - - Args: - device_id (str): The device ID of the device to update. - display_name (str): New display name for the device. - """ - content = { - "display_name": display_name - } - return self._send("PUT", "/devices/%s" % device_id, content=content) - - def delete_device(self, auth_body, device_id): - """Deletes the given device, and invalidates any access token associated with it. - - NOTE: This endpoint uses the User-Interactive Authentication API. - - Args: - auth_body (dict): Authentication params. - device_id (str): The device ID of the device to delete. - """ - content = { - "auth": auth_body - } - return self._send("DELETE", "/devices/%s" % device_id, content=content) - - def delete_devices(self, auth_body, devices): - """Bulk deletion of devices. - - NOTE: This endpoint uses the User-Interactive Authentication API. - - Args: - auth_body (dict): Authentication params. - devices (list): List of device ID"s to delete. - """ - content = { - "auth": auth_body, - "devices": devices - } - return self._send("POST", "/delete_devices", content=content) - - def upload_keys(self, device_keys=None, one_time_keys=None): - """Publishes end-to-end encryption keys for the device. - - Said device must be the one used when logging in. - - Args: - device_keys (dict): Optional. Identity keys for the device. The required - keys are: - - | user_id (str): The ID of the user the device belongs to. Must match - the user ID used when logging in. - | device_id (str): The ID of the device these keys belong to. Must match - the device ID used when logging in. - | algorithms (list): The encryption algorithms supported by this - device. - | keys (dict): Public identity keys. Should be formatted as - : . - | signatures (dict): Signatures for the device key object. Should be - formatted as : {: } - - one_time_keys (dict): Optional. One-time public keys. Should be - formatted as : , the key format being - determined by the algorithm. - """ - content = {} - if device_keys: - content["device_keys"] = device_keys - if one_time_keys: - content["one_time_keys"] = one_time_keys - return self._send("POST", "/keys/upload", content=content) - - def query_keys(self, user_devices, timeout=None, token=None): - """Query HS for public keys by user and optionally device. - - Args: - user_devices (dict): The devices whose keys to download. Should be - formatted as : []. No device_ids indicates - all devices for the corresponding user. - timeout (int): Optional. The time (in milliseconds) to wait when - downloading keys from remote servers. - token (str): Optional. If the client is fetching keys as a result of - a device update received in a sync request, this should be the - 'since' token of that sync request, or any later sync token. - """ - content = {"device_keys": user_devices} - if timeout: - content["timeout"] = timeout - if token: - content["token"] = token - return self._send("POST", "/keys/query", content=content) - - def claim_keys(self, key_request, timeout=None): - """Claims one-time keys for use in pre-key messages. - - Args: - key_request (dict): The keys to be claimed. Format should be - : { : }. - timeout (int): Optional. The time (in milliseconds) to wait when - downloading keys from remote servers. - """ - content = {"one_time_keys": key_request} - if timeout: - content["timeout"] = timeout - return self._send("POST", "/keys/claim", content=content) - - def key_changes(self, from_token, to_token): - """Gets a list of users who have updated their device identity keys. - - Args: - from_token (str): The desired start point of the list. Should be the - next_batch field from a response to an earlier call to /sync. - to_token (str): The desired end point of the list. Should be the next_batch - field from a recent call to /sync - typically the most recent such call. - """ - params = {"from": from_token, "to": to_token} - return self._send("GET", "/keys/changes", query_params=params) - - def send_to_device(self, event_type, messages, txn_id=None): - """Sends send-to-device events to a set of client devices. - - Args: - event_type (str): The type of event to send. - messages (dict): The messages to send. Format should be - : {: }. - The device ID may also be '*', meaning all known devices for the user. - txn_id (str): Optional. The transaction ID for this event, will be generated - automatically otherwise. - """ - txn_id = txn_id if txn_id else self._make_txn_id() - return self._send( - "PUT", - "/sendToDevice/{}/{}".format(event_type, txn_id), - content={"messages": messages} - ) - - def _make_txn_id(self): - txn_id = str(self.txn_id) + str(int(time() * 1000)) - self.txn_id += 1 - return txn_id - - def whoami(self): - """Determine user_id for authenticated user. - """ - if not self.token: - raise MatrixError("Authentication required.") - return self._send( - "GET", - "/account/whoami" - ) diff --git a/venv/lib/python3.10/site-packages/matrix_client/checks.py b/venv/lib/python3.10/site-packages/matrix_client/checks.py deleted file mode 100644 index 7e51e7c52c1ac4289148fa4917f7c248151b2fb6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client/checks.py +++ /dev/null @@ -1,30 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2015 OpenMarket Ltd -# -# 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. - - -def check_room_id(room_id): - if not room_id.startswith("!"): - raise ValueError("RoomIDs start with !") - - if ":" not in room_id: - raise ValueError("RoomIDs must have a domain component, seperated by a :") - - -def check_user_id(user_id): - if not user_id.startswith("@"): - raise ValueError("UserIDs start with @") - - if ":" not in user_id: - raise ValueError("UserIDs must have a domain component, seperated by a :") diff --git a/venv/lib/python3.10/site-packages/matrix_client/client.py b/venv/lib/python3.10/site-packages/matrix_client/client.py deleted file mode 100644 index 466cc427c35f54b5f7de6e705ee1b967436f88d4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client/client.py +++ /dev/null @@ -1,677 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2015 OpenMarket Ltd -# -# 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 .api import MatrixHttpApi -from .errors import MatrixRequestError, MatrixUnexpectedResponse -from .room import Room -from .user import User -try: - from .crypto.olm_device import OlmDevice - ENCRYPTION_SUPPORT = True -except ImportError: - ENCRYPTION_SUPPORT = False -from threading import Thread -from time import sleep -from uuid import uuid4 -from warnings import warn -import logging -import sys - -logger = logging.getLogger(__name__) - - -# Cache constants used when instantiating Matrix Client to specify level of caching -class CACHE(int): - pass - - -CACHE.NONE = CACHE(-1) -CACHE.SOME = CACHE(0) -CACHE.ALL = CACHE(1) -# TODO: rather than having CACHE.NONE as kwarg to MatrixClient, there should be a separate -# LightweightMatrixClient that only implements global listeners and doesn't hook into -# User, Room, etc. classes at all. - - -class MatrixClient(object): - """ - The client API for Matrix. For the raw HTTP calls, see MatrixHttpApi. - - Args: - base_url (str): The url of the HS preceding /_matrix. - e.g. (ex: https://localhost:8008 ) - token (Optional[str]): If you have an access token - supply it here. - user_id (Optional[str]): Optional. Obsolete. For backward compatibility. - valid_cert_check (bool): Check the homeservers - certificate on connections? - cache_level (CACHE): One of CACHE.NONE, CACHE.SOME, or - CACHE.ALL (defined in module namespace). - encryption (bool): Optional. Whether or not to enable end-to-end encryption - support. - encryption_conf (dict): Optional. Configuration parameters for encryption. - Refer to :func:`~matrix_client.crypto.olm_device.OlmDevice` for supported - options, since it will be passed to this class. - - Returns: - `MatrixClient` - - Raises: - `MatrixRequestError`, `ValueError` - - Examples: - - Create a new user and send a message:: - - client = MatrixClient("https://matrix.org") - token = client.register_with_password(username="foobar", - password="monkey") - room = client.create_room("myroom") - room.send_image(file_like_object) - - Send a message with an already logged in user:: - - client = MatrixClient("https://matrix.org", token="foobar", - user_id="@foobar:matrix.org") - client.add_listener(func) # NB: event stream callback - client.rooms[0].add_listener(func) # NB: callbacks just for this room. - room = client.join_room("#matrix:matrix.org") - response = room.send_text("Hello!") - response = room.kick("@bob:matrix.org") - - Incoming event callbacks (scopes):: - - def user_callback(user, incoming_event): - pass - - def room_callback(room, incoming_event): - pass - - def global_callback(incoming_event): - pass - - Attributes: - users (dict): A map from user ID to :class:`.User` object. - It is populated automatically while tracking the membership in rooms, and - shouldn't be modified directly. - A :class:`.User` object in this dict is shared between all :class:`.Room` - objects where the corresponding user is joined. - """ - - def __init__(self, base_url, token=None, user_id=None, - valid_cert_check=True, sync_filter_limit=20, - cache_level=CACHE.ALL, encryption=False, encryption_conf=None): - if user_id: - warn( - "user_id is deprecated. " - "Now it is requested from the server.", DeprecationWarning - ) - - if encryption and not ENCRYPTION_SUPPORT: - raise ValueError("Failed to enable encryption. Please make sure the olm " - "library is available.") - - self.api = MatrixHttpApi(base_url, token) - self.api.validate_certificate(valid_cert_check) - self.listeners = [] - self.presence_listeners = {} - self.invite_listeners = [] - self.left_listeners = [] - self.ephemeral_listeners = [] - self.device_id = None - self._encryption = encryption - self.encryption_conf = encryption_conf or {} - self.olm_device = None - if isinstance(cache_level, CACHE): - self._cache_level = cache_level - else: - self._cache_level = CACHE.ALL - raise ValueError( - "cache_level must be one of CACHE.NONE, CACHE.SOME, CACHE.ALL" - ) - - self.sync_token = None - self.sync_filter = '{ "room": { "timeline" : { "limit" : %i } } }' \ - % sync_filter_limit - self.sync_thread = None - self.should_listen = False - - """ Time to wait before attempting a /sync request after failing.""" - self.bad_sync_timeout_limit = 60 * 60 - self.rooms = { - # room_id: Room - } - self.users = { - # user_id: User - } - if token: - response = self.api.whoami() - self.user_id = response["user_id"] - self._sync() - - def get_sync_token(self): - warn("get_sync_token is deprecated. Directly access MatrixClient.sync_token.", - DeprecationWarning) - return self.sync_token - - def set_sync_token(self, token): - warn("set_sync_token is deprecated. Directly access MatrixClient.sync_token.", - DeprecationWarning) - self.sync_token = token - - def set_user_id(self, user_id): - warn("set_user_id is deprecated. Directly access MatrixClient.user_id.", - DeprecationWarning) - self.user_id = user_id - - # TODO: combine register methods into single register method controlled by kwargs - def register_as_guest(self): - """ Register a guest account on this HS. - Note: HS must have guest registration enabled. - Returns: - str: Access Token - Raises: - MatrixRequestError - """ - response = self.api.register(auth_body=None, kind='guest') - return self._post_registration(response) - - def register_with_password(self, username, password): - """ Register for a new account on this HS. - - Args: - username (str): Account username - password (str): Account password - - Returns: - str: Access Token - - Raises: - MatrixRequestError - """ - response = self.api.register( - auth_body={"type": "m.login.dummy"}, - kind='user', - username=username, - password=password, - ) - return self._post_registration(response) - - def _post_registration(self, response): - self.user_id = response["user_id"] - self.token = response["access_token"] - self.hs = response["home_server"] - self.api.token = self.token - self._sync() - return self.token - - def login_with_password_no_sync(self, username, password): - """Deprecated. Use ``login`` with ``sync=False``. - - Login to the homeserver. - - Args: - username (str): Account username - password (str): Account password - - Returns: - str: Access token - - Raises: - MatrixRequestError - """ - warn("login_with_password_no_sync is deprecated. Use login with sync=False.", - DeprecationWarning) - return self.login(username, password, sync=False) - - def login_with_password(self, username, password, limit=10): - """Deprecated. Use ``login`` with ``sync=True``. - - Login to the homeserver. - - Args: - username (str): Account username - password (str): Account password - limit (int): Deprecated. How many messages to return when syncing. - This will be replaced by a filter API in a later release. - - Returns: - str: Access token - - Raises: - MatrixRequestError - """ - warn("login_with_password is deprecated. Use login with sync=True.", - DeprecationWarning) - return self.login(username, password, limit, sync=True) - - def login(self, username, password, limit=10, sync=True, device_id=None): - """Login to the homeserver. - - Args: - username (str): Account username - password (str): Account password - limit (int): Deprecated. How many messages to return when syncing. - This will be replaced by a filter API in a later release. - sync (bool): Optional. Whether to initiate a /sync request after logging in. - device_id (str): Optional. ID of the client device. The server will - auto-generate a device_id if this is not specified. - - Returns: - str: Access token - - Raises: - MatrixRequestError - """ - response = self.api.login( - "m.login.password", user=username, password=password, device_id=device_id - ) - self.user_id = response["user_id"] - self.token = response["access_token"] - self.hs = response["home_server"] - self.api.token = self.token - self.device_id = response["device_id"] - - if self._encryption: - self.olm_device = OlmDevice( - self.api, self.user_id, self.device_id, **self.encryption_conf) - self.olm_device.upload_identity_keys() - self.olm_device.upload_one_time_keys() - - if sync: - """ Limit Filter """ - self.sync_filter = '{ "room": { "timeline" : { "limit" : %i } } }' % limit - self._sync() - return self.token - - def logout(self): - """ Logout from the homeserver. - """ - self.stop_listener_thread() - self.api.logout() - - # TODO: move room creation/joining to User class for future application service usage - # NOTE: we may want to leave thin wrappers here for convenience - def create_room(self, alias=None, is_public=False, invitees=None): - """ Create a new room on the homeserver. - - Args: - alias (str): The canonical_alias of the room. - is_public (bool): The public/private visibility of the room. - invitees (str[]): A set of user ids to invite into the room. - - Returns: - Room - - Raises: - MatrixRequestError - """ - response = self.api.create_room(alias=alias, - is_public=is_public, - invitees=invitees) - return self._mkroom(response["room_id"]) - - def join_room(self, room_id_or_alias): - """ Join a room. - - Args: - room_id_or_alias (str): Room ID or an alias. - - Returns: - Room - - Raises: - MatrixRequestError - """ - response = self.api.join_room(room_id_or_alias) - room_id = ( - response["room_id"] if "room_id" in response else room_id_or_alias - ) - return self._mkroom(room_id) - - def get_rooms(self): - """ Deprecated. Return a dict of {room_id: Room objects} that the user has joined. - - Returns: - Room{}: Rooms the user has joined. - """ - warn("get_rooms is deprecated. Directly access MatrixClient.rooms.", - DeprecationWarning) - return self.rooms - - # TODO: create Listener class and push as much of this logic there as possible - # NOTE: listeners related to things in rooms should be attached to Room objects - def add_listener(self, callback, event_type=None): - """ Add a listener that will send a callback when the client recieves - an event. - - Args: - callback (func(roomchunk)): Callback called when an event arrives. - event_type (str): The event_type to filter for. - - Returns: - uuid.UUID: Unique id of the listener, can be used to identify the listener. - """ - listener_uid = uuid4() - # TODO: listeners should be stored in dict and accessed/deleted directly. Add - # convenience method such that MatrixClient.listeners.new(Listener(...)) performs - # MatrixClient.listeners[uuid4()] = Listener(...) - self.listeners.append( - { - 'uid': listener_uid, - 'callback': callback, - 'event_type': event_type - } - ) - return listener_uid - - def remove_listener(self, uid): - """ Remove listener with given uid. - - Args: - uuid.UUID: Unique id of the listener to remove. - """ - self.listeners[:] = (listener for listener in self.listeners - if listener['uid'] != uid) - - def add_presence_listener(self, callback): - """ Add a presence listener that will send a callback when the client receives - a presence update. - - Args: - callback (func(roomchunk)): Callback called when a presence update arrives. - - Returns: - uuid.UUID: Unique id of the listener, can be used to identify the listener. - """ - listener_uid = uuid4() - self.presence_listeners[listener_uid] = callback - return listener_uid - - def remove_presence_listener(self, uid): - """ Remove presence listener with given uid - - Args: - uuid.UUID: Unique id of the listener to remove - """ - self.presence_listeners.pop(uid) - - def add_ephemeral_listener(self, callback, event_type=None): - """ Add an ephemeral listener that will send a callback when the client recieves - an ephemeral event. - - Args: - callback (func(roomchunk)): Callback called when an ephemeral event arrives. - event_type (str): The event_type to filter for. - - Returns: - uuid.UUID: Unique id of the listener, can be used to identify the listener. - """ - listener_id = uuid4() - self.ephemeral_listeners.append( - { - 'uid': listener_id, - 'callback': callback, - 'event_type': event_type - } - ) - return listener_id - - def remove_ephemeral_listener(self, uid): - """ Remove ephemeral listener with given uid. - - Args: - uuid.UUID: Unique id of the listener to remove. - """ - self.ephemeral_listeners[:] = (listener for listener in self.ephemeral_listeners - if listener['uid'] != uid) - - def add_invite_listener(self, callback): - """ Add a listener that will send a callback when the client receives - an invite. - - Args: - callback (func(room_id, state)): Callback called when an invite arrives. - """ - self.invite_listeners.append(callback) - - def add_leave_listener(self, callback): - """ Add a listener that will send a callback when the client has left a room. - - Args: - callback (func(room_id, room)): Callback called when the client - has left a room. - """ - self.left_listeners.append(callback) - - def listen_for_events(self, timeout_ms=30000): - """ - This function just calls _sync() - - In a future version of this sdk, this function will be deprecated and - _sync method will be renamed sync with the intention of it being called - by downstream code. - - Args: - timeout_ms (int): How long to poll the Home Server for before - retrying. - """ - # TODO: see docstring - self._sync(timeout_ms) - - def listen_forever(self, timeout_ms=30000, exception_handler=None, - bad_sync_timeout=5): - """ Keep listening for events forever. - - Args: - timeout_ms (int): How long to poll the Home Server for before - retrying. - exception_handler (func(exception)): Optional exception handler - function which can be used to handle exceptions in the caller - thread. - bad_sync_timeout (int): Base time to wait after an error before - retrying. Will be increased according to exponential backoff. - """ - _bad_sync_timeout = bad_sync_timeout - self.should_listen = True - while (self.should_listen): - try: - self._sync(timeout_ms) - _bad_sync_timeout = bad_sync_timeout - # TODO: we should also handle MatrixHttpLibError for retry in case no response - except MatrixRequestError as e: - logger.warning("A MatrixRequestError occured during sync.") - if e.code >= 500: - logger.warning("Problem occured serverside. Waiting %i seconds", - bad_sync_timeout) - sleep(bad_sync_timeout) - _bad_sync_timeout = min(_bad_sync_timeout * 2, - self.bad_sync_timeout_limit) - elif exception_handler is not None: - exception_handler(e) - else: - raise - except Exception as e: - logger.exception("Exception thrown during sync") - if exception_handler is not None: - exception_handler(e) - else: - raise - - def start_listener_thread(self, timeout_ms=30000, exception_handler=None): - """ Start a listener thread to listen for events in the background. - - Args: - timeout (int): How long to poll the Home Server for before - retrying. - exception_handler (func(exception)): Optional exception handler - function which can be used to handle exceptions in the caller - thread. - """ - try: - thread = Thread(target=self.listen_forever, - args=(timeout_ms, exception_handler)) - thread.daemon = True - self.sync_thread = thread - self.should_listen = True - thread.start() - except RuntimeError: - e = sys.exc_info()[0] - logger.error("Error: unable to start thread. %s", str(e)) - - def stop_listener_thread(self): - """ Stop listener thread running in the background - """ - if self.sync_thread: - self.should_listen = False - self.sync_thread.join() - self.sync_thread = None - - # TODO: move to User class. Consider creating lightweight Media class. - def upload(self, content, content_type, filename=None): - """ Upload content to the home server and recieve a MXC url. - - Args: - content (bytes): The data of the content. - content_type (str): The mimetype of the content. - filename (str): Optional. Filename of the content. - - Raises: - MatrixUnexpectedResponse: If the homeserver gave a strange response - MatrixRequestError: If the upload failed for some reason. - """ - try: - response = self.api.media_upload(content, content_type, filename) - if "content_uri" in response: - return response["content_uri"] - else: - raise MatrixUnexpectedResponse( - "The upload was successful, but content_uri wasn't found." - ) - except MatrixRequestError as e: - raise MatrixRequestError( - code=e.code, - content="Upload failed: %s" % e - ) - - def _mkroom(self, room_id): - room = Room(self, room_id) - if self._encryption: - try: - event = self.api.get_state_event(room_id, "m.room.encryption") - if event["algorithm"] == "m.megolm.v1.aes-sha2": - room.encrypted = True - except MatrixRequestError as e: - if e.code != 404: - raise - self.rooms[room_id] = room - return self.rooms[room_id] - - # TODO better handling of the blocking I/O caused by update_one_time_key_counts - def _sync(self, timeout_ms=30000): - response = self.api.sync(self.sync_token, timeout_ms, filter=self.sync_filter) - self.sync_token = response["next_batch"] - - if 'presence' in response and 'events' in response['presence']: - for presence_update in response['presence']['events']: - for callback in self.presence_listeners.values(): - callback(presence_update) - - if self._encryption and 'device_one_time_keys_count' in response: - self.olm_device.update_one_time_key_counts( - response['device_one_time_keys_count']) - - rooms = response.get("rooms", {}) - if 'invite' in rooms: - for room_id, invite_room in rooms['invite'].items(): - for listener in self.invite_listeners: - listener(room_id, invite_room['invite_state']) - - if 'leave' in rooms: - for room_id, left_room in rooms['leave'].items(): - for listener in self.left_listeners: - listener(room_id, left_room) - if room_id in self.rooms: - del self.rooms[room_id] - - if 'join' in rooms: - for room_id, sync_room in rooms['join'].items(): - if room_id not in self.rooms: - self._mkroom(room_id) - room = self.rooms[room_id] - # TODO: the rest of this for loop should be in room object method - room.prev_batch = sync_room["timeline"]["prev_batch"] - - if "state" in sync_room and "events" in sync_room["state"]: - for event in sync_room["state"]["events"]: - event['room_id'] = room_id - room._process_state_event(event) - - if "timeline" in sync_room and "events" in sync_room["timeline"]: - for event in sync_room["timeline"]["events"]: - event['room_id'] = room_id - room._put_event(event) - - # TODO: global listeners can still exist but work by each - # room.listeners[uuid] having reference to global listener - - # Dispatch for client (global) listeners - for listener in self.listeners: - if ( - listener['event_type'] is None or - listener['event_type'] == event['type'] - ): - listener['callback'](event) - - if "ephemeral" in sync_room and "events" in sync_room["ephemeral"]: - for event in sync_room['ephemeral']['events']: - event['room_id'] = room_id - room._put_ephemeral_event(event) - - for listener in self.ephemeral_listeners: - if ( - listener['event_type'] is None or - listener['event_type'] == event['type'] - ): - listener['callback'](event) - - def get_user(self, user_id): - """Deprecated. Return a User by their id. - - This method only instantiate a User, which should be done directly. - You can also use :attr:`users` in order to access a User object which - was created automatically. - - Args: - user_id (str): The matrix user id of a user. - """ - warn("get_user is deprecated. Directly instantiate a User instead.", - DeprecationWarning) - return User(self.api, user_id) - - # TODO: move to Room class - def remove_room_alias(self, room_alias): - """Remove mapping of an alias - - Args: - room_alias(str): The alias to be removed. - - Returns: - bool: True if the alias is removed, False otherwise. - """ - try: - self.api.remove_room_alias(room_alias) - return True - except MatrixRequestError: - return False diff --git a/venv/lib/python3.10/site-packages/matrix_client/crypto/__init__.py b/venv/lib/python3.10/site-packages/matrix_client/crypto/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/matrix_client/crypto/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/matrix_client/crypto/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 96053e90ac81b83ffc93a31646d1763839206a57..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/matrix_client/crypto/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/matrix_client/crypto/__pycache__/olm_device.cpython-310.pyc b/venv/lib/python3.10/site-packages/matrix_client/crypto/__pycache__/olm_device.cpython-310.pyc deleted file mode 100644 index e1c7360d8592f911bf997bc65e2bf9139c5819c0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/matrix_client/crypto/__pycache__/olm_device.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/matrix_client/crypto/__pycache__/one_time_keys.cpython-310.pyc b/venv/lib/python3.10/site-packages/matrix_client/crypto/__pycache__/one_time_keys.cpython-310.pyc deleted file mode 100644 index f6859d35ce7a8ddcbc7765377b4f2747e1096434..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/matrix_client/crypto/__pycache__/one_time_keys.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/matrix_client/crypto/olm_device.py b/venv/lib/python3.10/site-packages/matrix_client/crypto/olm_device.py deleted file mode 100644 index 514965db880eb339887dcd7e564e5611c44c2140..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client/crypto/olm_device.py +++ /dev/null @@ -1,207 +0,0 @@ -import logging - -import olm -from canonicaljson import encode_canonical_json - -from matrix_client.checks import check_user_id -from matrix_client.crypto.one_time_keys import OneTimeKeysManager - -logger = logging.getLogger(__name__) - - -class OlmDevice(object): - """Manages the Olm cryptographic functions. - - Has a unique Olm account which holds identity keys. - - Args: - api (MatrixHttpApi): The api object used to make requests. - user_id (str): Matrix user ID. Must match the one used when logging in. - device_id (str): Must match the one used when logging in. - signed_keys_proportion (float): Optional. The proportion of signed one-time keys - we should maintain on the HS compared to unsigned keys. The maximum value of - ``1`` means only signed keys will be uploaded, while the minimum value of - ``0`` means only unsigned keys. The actual amount of keys is determined at - runtime from the given proportion and the maximum number of one-time keys - we can physically hold. - keys_threshold (float): Optional. Threshold below which a one-time key - replenishment is triggered. Must be between ``0`` and ``1``. For example, - ``0.1`` means that new one-time keys will be uploaded when there is less than - 10% of the maximum number of one-time keys on the server. - """ - - _olm_algorithm = 'm.olm.v1.curve25519-aes-sha2' - _megolm_algorithm = 'm.megolm.v1.aes-sha2' - _algorithms = [_olm_algorithm, _megolm_algorithm] - - def __init__(self, - api, - user_id, - device_id, - signed_keys_proportion=1, - keys_threshold=0.1): - if not 0 <= signed_keys_proportion <= 1: - raise ValueError('signed_keys_proportion must be between 0 and 1.') - if not 0 <= keys_threshold <= 1: - raise ValueError('keys_threshold must be between 0 and 1.') - self.api = api - check_user_id(user_id) - self.user_id = user_id - self.device_id = device_id - self.olm_account = olm.Account() - logger.info('Initialised Olm Device.') - self.identity_keys = self.olm_account.identity_keys - # Try to maintain half the number of one-time keys libolm can hold uploaded - # on the HS. This is because some keys will be claimed by peers but not - # used instantly, and we want them to stay in libolm, until the limit is reached - # and it starts discarding keys, starting by the oldest. - target_keys_number = self.olm_account.max_one_time_keys // 2 - self.one_time_keys_manager = OneTimeKeysManager(target_keys_number, - signed_keys_proportion, - keys_threshold) - - def upload_identity_keys(self): - """Uploads this device's identity keys to HS. - - This device must be the one used when logging in. - """ - device_keys = { - 'user_id': self.user_id, - 'device_id': self.device_id, - 'algorithms': self._algorithms, - 'keys': {'{}:{}'.format(alg, self.device_id): key - for alg, key in self.identity_keys.items()} - } - self.sign_json(device_keys) - ret = self.api.upload_keys(device_keys=device_keys) - self.one_time_keys_manager.server_counts = ret['one_time_key_counts'] - logger.info('Uploaded identity keys.') - - def upload_one_time_keys(self, force_update=False): - """Uploads new one-time keys to the HS, if needed. - - Args: - force_update (bool): Fetch the number of one-time keys currently on the HS - before uploading, even if we already know one. In most cases this should - not be necessary, as we get this value from sync responses. - - Returns: - A dict containg the number of new keys that were uploaded for each key type - (signed_curve25519 or curve25519). The format is - ``: ``. If no keys of a given type have been - uploaded, the corresponding key will not be present. Consequently, an - empty dict indicates that no keys were uploaded. - """ - if force_update or not self.one_time_keys_manager.server_counts: - counts = self.api.upload_keys()['one_time_key_counts'] - self.one_time_keys_manager.server_counts = counts - - signed_keys_to_upload = self.one_time_keys_manager.signed_curve25519_to_upload - unsigned_keys_to_upload = self.one_time_keys_manager.curve25519_to_upload - - self.olm_account.generate_one_time_keys(signed_keys_to_upload + - unsigned_keys_to_upload) - - one_time_keys = {} - keys = self.olm_account.one_time_keys['curve25519'] - for i, key_id in enumerate(keys): - if i < signed_keys_to_upload: - key = self.sign_json({'key': keys[key_id]}) - key_type = 'signed_curve25519' - else: - key = keys[key_id] - key_type = 'curve25519' - one_time_keys['{}:{}'.format(key_type, key_id)] = key - - ret = self.api.upload_keys(one_time_keys=one_time_keys) - self.one_time_keys_manager.server_counts = ret['one_time_key_counts'] - self.olm_account.mark_keys_as_published() - - keys_uploaded = {} - if unsigned_keys_to_upload: - keys_uploaded['curve25519'] = unsigned_keys_to_upload - if signed_keys_to_upload: - keys_uploaded['signed_curve25519'] = signed_keys_to_upload - logger.info('Uploaded new one-time keys: %s.', keys_uploaded) - return keys_uploaded - - def update_one_time_key_counts(self, counts): - """Update data on one-time keys count and upload new ones if necessary. - - Args: - counts (dict): Counts of keys currently on the HS for each key type. - """ - self.one_time_keys_manager.server_counts = counts - if self.one_time_keys_manager.should_upload(): - logger.info('Uploading new one-time keys.') - self.upload_one_time_keys() - - def sign_json(self, json): - """Signs a JSON object. - - NOTE: The object is modified in-place and the return value can be ignored. - - As specified, this is done by encoding the JSON object without ``signatures`` or - keys grouped as ``unsigned``, using canonical encoding. - - Args: - json (dict): The JSON object to sign. - - Returns: - The same JSON object, with a ``signatures`` key added. It is formatted as - ``"signatures": ed25519:: ``. - """ - signatures = json.pop('signatures', {}) - unsigned = json.pop('unsigned', None) - - signature_base64 = self.olm_account.sign(encode_canonical_json(json)) - - key_id = 'ed25519:{}'.format(self.device_id) - signatures.setdefault(self.user_id, {})[key_id] = signature_base64 - - json['signatures'] = signatures - if unsigned: - json['unsigned'] = unsigned - - return json - - def verify_json(self, json, user_key, user_id, device_id): - """Verifies a signed key object's signature. - - The object must have a 'signatures' key associated with an object of the form - `user_id: {key_id: signature}`. - - Args: - json (dict): The JSON object to verify. - user_key (str): The public ed25519 key which was used to sign the object. - user_id (str): The user who owns the device. - device_id (str): The device who owns the key. - - Returns: - True if the verification was successful, False if not. - """ - try: - signatures = json.pop('signatures') - except KeyError: - return False - - key_id = 'ed25519:{}'.format(device_id) - try: - signature_base64 = signatures[user_id][key_id] - except KeyError: - json['signatures'] = signatures - return False - - unsigned = json.pop('unsigned', None) - - try: - olm.ed25519_verify(user_key, encode_canonical_json(json), signature_base64) - success = True - except olm.utility.OlmVerifyError: - success = False - - json['signatures'] = signatures - if unsigned: - json['unsigned'] = unsigned - - return success diff --git a/venv/lib/python3.10/site-packages/matrix_client/crypto/one_time_keys.py b/venv/lib/python3.10/site-packages/matrix_client/crypto/one_time_keys.py deleted file mode 100644 index 131dc023f4d46377c5d2bf29f3af8329f68c4a19..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client/crypto/one_time_keys.py +++ /dev/null @@ -1,42 +0,0 @@ -class OneTimeKeysManager(object): - """Handles one-time keys accounting for an OlmDevice.""" - - def __init__(self, target_keys_number, signed_keys_proportion, keys_threshold): - self.target_counts = { - 'signed_curve25519': int(round(signed_keys_proportion * target_keys_number)), - 'curve25519': int(round((1 - signed_keys_proportion) * target_keys_number)), - } - self._server_counts = {} - self.to_upload = {} - self.keys_threshold = keys_threshold - - @property - def server_counts(self): - return self._server_counts - - @server_counts.setter - def server_counts(self, server_counts): - self._server_counts = server_counts - self.update_keys_to_upload() - - def update_keys_to_upload(self): - for key_type, target_number in self.target_counts.items(): - num_keys = self._server_counts.get(key_type, 0) - num_to_create = max(target_number - num_keys, 0) - self.to_upload[key_type] = num_to_create - - def should_upload(self): - if not self._server_counts: - return True - for key_type, target_number in self.target_counts.items(): - if self._server_counts.get(key_type, 0) < target_number * self.keys_threshold: - return True - return False - - @property - def curve25519_to_upload(self): - return self.to_upload.get('curve25519', 0) - - @property - def signed_curve25519_to_upload(self): - return self.to_upload.get('signed_curve25519', 0) diff --git a/venv/lib/python3.10/site-packages/matrix_client/errors.py b/venv/lib/python3.10/site-packages/matrix_client/errors.py deleted file mode 100644 index be1356407cbf007ad201827a7ea59993728e0333..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client/errors.py +++ /dev/null @@ -1,48 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2015 OpenMarket Ltd -# -# 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. - - -class MatrixError(Exception): - """A generic Matrix error. Specific errors will subclass this.""" - pass - - -class MatrixUnexpectedResponse(MatrixError): - """The home server gave an unexpected response. """ - - def __init__(self, content=""): - super(MatrixUnexpectedResponse, self).__init__(content) - self.content = content - - -class MatrixRequestError(MatrixError): - """ The home server returned an error response. """ - - def __init__(self, code=0, content=""): - super(MatrixRequestError, self).__init__("%d: %s" % (code, content)) - self.code = code - self.content = content - - -class MatrixHttpLibError(MatrixError): - """The library used for http requests raised an exception.""" - - def __init__(self, original_exception, method, endpoint): - super(MatrixHttpLibError, self).__init__( - "Something went wrong in {} requesting {}: {}".format(method, - endpoint, - original_exception) - ) - self.original_exception = original_exception diff --git a/venv/lib/python3.10/site-packages/matrix_client/room.py b/venv/lib/python3.10/site-packages/matrix_client/room.py deleted file mode 100644 index 9083ed4a1d24d6f85f4ea53964affd30a0d00af6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client/room.py +++ /dev/null @@ -1,689 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2015 OpenMarket Ltd -# Copyright 2018 Adam Beckmeyer -# -# 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 logging -import re -from uuid import uuid4 - -from .checks import check_room_id -from .user import User -from .errors import MatrixRequestError - -logger = logging.getLogger(__name__) - - -class Room(object): - """Call room-specific functions after joining a room from the client. - - NOTE: This should ideally be called from within the Client. - NOTE: This does not verify the room with the Home Server. - """ - - def __init__(self, client, room_id): - check_room_id(room_id) - - self.room_id = room_id - self.client = client - self.listeners = [] - self.state_listeners = [] - self.ephemeral_listeners = [] - self.events = [] - self.event_history_limit = 20 - self.name = None - self.canonical_alias = None - self.aliases = [] - self.topic = None - self.invite_only = None - self.guest_access = None - self._prev_batch = None - self._members = {} - self.members_displaynames = { - # user_id: displayname - } - self.encrypted = False - - def set_user_profile(self, - displayname=None, - avatar_url=None, - reason="Changing room profile information"): - """Set user profile within a room. - - This sets displayname and avatar_url for the logged in user only in a - specific room. It does not change the user's global user profile. - """ - member = self.client.api.get_membership(self.room_id, self.client.user_id) - if member["membership"] != "join": - raise Exception("Can't set profile if you have not joined the room.") - if displayname is None: - displayname = member["displayname"] - if avatar_url is None: - avatar_url = member["avatar_url"] - self.client.api.set_membership( - self.room_id, - self.client.user_id, - 'join', - reason, { - "displayname": displayname, - "avatar_url": avatar_url - } - ) - - @property - def display_name(self): - """Calculates the display name for a room.""" - if self.name: - return self.name - elif self.canonical_alias: - return self.canonical_alias - - # Member display names without me - members = [u.get_display_name(self) for u in self.get_joined_members() if - self.client.user_id != u.user_id] - members.sort() - - if len(members) == 1: - return members[0] - elif len(members) == 2: - return "{0} and {1}".format(members[0], members[1]) - elif len(members) > 2: - return "{0} and {1} others".format(members[0], len(members) - 1) - else: # len(members) <= 0 or not an integer - # TODO i18n - return "Empty room" - - def send_text(self, text): - """Send a plain text message to the room.""" - return self.client.api.send_message(self.room_id, text) - - def get_html_content(self, html, body=None, msgtype="m.text"): - return { - "body": body if body else re.sub('<[^<]+?>', '', html), - "msgtype": msgtype, - "format": "org.matrix.custom.html", - "formatted_body": html - } - - def send_html(self, html, body=None, msgtype="m.text"): - """Send an html formatted message. - - Args: - html (str): The html formatted message to be sent. - body (str): The unformatted body of the message to be sent. - """ - return self.client.api.send_message_event( - self.room_id, "m.room.message", self.get_html_content(html, body, msgtype)) - - def set_account_data(self, type, account_data): - return self.client.api.set_room_account_data( - self.client.user_id, self.room_id, type, account_data) - - def get_tags(self): - return self.client.api.get_user_tags(self.client.user_id, self.room_id) - - def remove_tag(self, tag): - return self.client.api.remove_user_tag( - self.client.user_id, self.room_id, tag - ) - - def add_tag(self, tag, order=None, content=None): - return self.client.api.add_user_tag( - self.client.user_id, self.room_id, - tag, order, content - ) - - def send_emote(self, text): - """Send an emote (/me style) message to the room.""" - return self.client.api.send_emote(self.room_id, text) - - def send_file(self, url, name, **fileinfo): - """Send a pre-uploaded file to the room. - - See http://matrix.org/docs/spec/r0.2.0/client_server.html#m-file for - fileinfo. - - Args: - url (str): The mxc url of the file. - name (str): The filename of the image. - fileinfo (): Extra information about the file - """ - - return self.client.api.send_content( - self.room_id, url, name, "m.file", - extra_information=fileinfo - ) - - def send_notice(self, text): - """Send a notice (from bot) message to the room.""" - return self.client.api.send_notice(self.room_id, text) - - # See http://matrix.org/docs/spec/r0.0.1/client_server.html#m-image for the - # imageinfo args. - def send_image(self, url, name, **imageinfo): - """Send a pre-uploaded image to the room. - - See http://matrix.org/docs/spec/r0.0.1/client_server.html#m-image - for imageinfo - - Args: - url (str): The mxc url of the image. - name (str): The filename of the image. - imageinfo (): Extra information about the image. - """ - return self.client.api.send_content( - self.room_id, url, name, "m.image", - extra_information=imageinfo - ) - - def send_location(self, geo_uri, name, thumb_url=None, **thumb_info): - """Send a location to the room. - - See http://matrix.org/docs/spec/client_server/r0.2.0.html#m-location - for thumb_info - - Args: - geo_uri (str): The geo uri representing the location. - name (str): Description for the location. - thumb_url (str): URL to the thumbnail of the location. - thumb_info (): Metadata about the thumbnail, type ImageInfo. - """ - return self.client.api.send_location(self.room_id, geo_uri, name, - thumb_url, thumb_info) - - def send_video(self, url, name, **videoinfo): - """Send a pre-uploaded video to the room. - - See http://matrix.org/docs/spec/client_server/r0.2.0.html#m-video - for videoinfo - - Args: - url (str): The mxc url of the video. - name (str): The filename of the video. - videoinfo (): Extra information about the video. - """ - return self.client.api.send_content(self.room_id, url, name, "m.video", - extra_information=videoinfo) - - def send_audio(self, url, name, **audioinfo): - """Send a pre-uploaded audio to the room. - - See http://matrix.org/docs/spec/client_server/r0.2.0.html#m-audio - for audioinfo - - Args: - url (str): The mxc url of the audio. - name (str): The filename of the audio. - audioinfo (): Extra information about the audio. - """ - return self.client.api.send_content(self.room_id, url, name, "m.audio", - extra_information=audioinfo) - - def redact_message(self, event_id, reason=None): - """Redacts the message with specified event_id for the given reason. - - See https://matrix.org/docs/spec/r0.0.1/client_server.html#id112 - """ - return self.client.api.redact_event(self.room_id, event_id, reason) - - def add_listener(self, callback, event_type=None): - """Add a callback handler for events going to this room. - - Args: - callback (func(room, event)): Callback called when an event arrives. - event_type (str): The event_type to filter for. - Returns: - uuid.UUID: Unique id of the listener, can be used to identify the listener. - """ - listener_id = uuid4() - self.listeners.append( - { - 'uid': listener_id, - 'callback': callback, - 'event_type': event_type - } - ) - return listener_id - - def remove_listener(self, uid): - """Remove listener with given uid.""" - self.listeners[:] = (listener for listener in self.listeners - if listener['uid'] != uid) - - def add_ephemeral_listener(self, callback, event_type=None): - """Add a callback handler for ephemeral events going to this room. - - Args: - callback (func(room, event)): Callback called when an ephemeral event arrives. - event_type (str): The event_type to filter for. - Returns: - uuid.UUID: Unique id of the listener, can be used to identify the listener. - """ - listener_id = uuid4() - self.ephemeral_listeners.append( - { - 'uid': listener_id, - 'callback': callback, - 'event_type': event_type - } - ) - return listener_id - - def remove_ephemeral_listener(self, uid): - """Remove ephemeral listener with given uid.""" - self.ephemeral_listeners[:] = (listener for listener in self.ephemeral_listeners - if listener['uid'] != uid) - - def add_state_listener(self, callback, event_type=None): - """Add a callback handler for state events going to this room. - - Args: - callback (func(roomchunk)): Callback called when an event arrives. - event_type (str): The event_type to filter for. - """ - self.state_listeners.append( - { - 'callback': callback, - 'event_type': event_type - } - ) - - def _put_event(self, event): - self.events.append(event) - if len(self.events) > self.event_history_limit: - self.events.pop(0) - if 'state_key' in event: - self._process_state_event(event) - - # Dispatch for room-specific listeners - for listener in self.listeners: - if listener['event_type'] is None or listener['event_type'] == event['type']: - listener['callback'](self, event) - - def _put_ephemeral_event(self, event): - # Dispatch for room-specific listeners - for listener in self.ephemeral_listeners: - if listener['event_type'] is None or listener['event_type'] == event['type']: - listener['callback'](self, event) - - def get_events(self): - """Get the most recent events for this room.""" - return self.events - - def invite_user(self, user_id): - """Invite a user to this room. - - Returns: - boolean: Whether invitation was sent. - """ - try: - self.client.api.invite_user(self.room_id, user_id) - return True - except MatrixRequestError: - return False - - def kick_user(self, user_id, reason=""): - """Kick a user from this room. - - - Args: - user_id (str): The matrix user id of a user. - reason (str): A reason for kicking the user. - - Returns: - boolean: Whether user was kicked. - """ - try: - self.client.api.kick_user(self.room_id, user_id) - return True - except MatrixRequestError: - return False - - def ban_user(self, user_id, reason): - """Ban a user from this room - - Args: - user_id (str): The matrix user id of a user. - reason (str): A reason for banning the user. - - Returns: - boolean: The user was banned. - """ - try: - self.client.api.ban_user(self.room_id, user_id, reason) - return True - except MatrixRequestError: - return False - - def unban_user(self, user_id): - """Unban a user from this room - - Returns: - boolean: The user was unbanned. - """ - try: - self.client.api.unban_user(self.room_id, user_id) - return True - except MatrixRequestError: - return False - - def leave(self): - """Leave the room. - - Returns: - boolean: Leaving the room was successful. - """ - try: - self.client.api.leave_room(self.room_id) - del self.client.rooms[self.room_id] - return True - except MatrixRequestError: - return False - - def update_room_name(self): - """Updates self.name and returns True if room name has changed.""" - try: - response = self.client.api.get_room_name(self.room_id) - if "name" in response and response["name"] != self.name: - self.name = response["name"] - return True - else: - return False - except MatrixRequestError: - return False - - def set_room_name(self, name): - """Return True if room name successfully changed.""" - try: - self.client.api.set_room_name(self.room_id, name) - self.name = name - return True - except MatrixRequestError: - return False - - def send_state_event(self, event_type, content, state_key=""): - """Send a state event to the room. - - Args: - event_type (str): The type of event that you are sending. - content (): An object with the content of the message. - state_key (str, optional): A unique key to identify the state. - """ - return self.client.api.send_state_event( - self.room_id, - event_type, - content, - state_key - ) - - def update_room_topic(self): - """Updates self.topic and returns True if room topic has changed.""" - try: - response = self.client.api.get_room_topic(self.room_id) - if "topic" in response and response["topic"] != self.topic: - self.topic = response["topic"] - return True - else: - return False - except MatrixRequestError: - return False - - def set_room_topic(self, topic): - """Set room topic. - - Returns: - boolean: True if the topic changed, False if not - """ - try: - self.client.api.set_room_topic(self.room_id, topic) - self.topic = topic - return True - except MatrixRequestError: - return False - - def update_aliases(self): - """Get aliases information from room state. - - Returns: - boolean: True if the aliases changed, False if not - """ - try: - response = self.client.api.get_room_state(self.room_id) - for chunk in response: - if "content" in chunk and "aliases" in chunk["content"]: - if chunk["content"]["aliases"] != self.aliases: - self.aliases = chunk["content"]["aliases"] - return True - else: - return False - except MatrixRequestError: - return False - - def add_room_alias(self, room_alias): - """Add an alias to the room and return True if successful.""" - try: - self.client.api.set_room_alias(self.room_id, room_alias) - return True - except MatrixRequestError: - return False - - def get_joined_members(self): - """Returns list of joined members (User objects).""" - if self._members: - return list(self._members.values()) - response = self.client.api.get_room_members(self.room_id) - for event in response["chunk"]: - if event["content"]["membership"] == "join": - user_id = event["state_key"] - self._add_member(user_id, event["content"].get("displayname")) - return list(self._members.values()) - - def _add_member(self, user_id, displayname=None): - if displayname: - self.members_displaynames[user_id] = displayname - if user_id in self._members: - return - if user_id in self.client.users: - self._members[user_id] = self.client.users[user_id] - return - self._members[user_id] = User(self.client.api, user_id, displayname) - self.client.users[user_id] = self._members[user_id] - - def backfill_previous_messages(self, reverse=False, limit=10): - """Backfill handling of previous messages. - - Args: - reverse (bool): When false messages will be backfilled in their original - order (old to new), otherwise the order will be reversed (new to old). - limit (int): Number of messages to go back. - """ - res = self.client.api.get_room_messages(self.room_id, self.prev_batch, - direction="b", limit=limit) - events = res["chunk"] - if not reverse: - events = reversed(events) - for event in events: - self._put_event(event) - - def modify_user_power_levels(self, users=None, users_default=None): - """Modify the power level for a subset of users - - Args: - users(dict): Power levels to assign to specific users, in the form - {"@name0:host0": 10, "@name1:host1": 100, "@name3:host3", None} - A level of None causes the user to revert to the default level - as specified by users_default. - users_default(int): Default power level for users in the room - - Returns: - True if successful, False if not - """ - try: - content = self.client.api.get_power_levels(self.room_id) - if users_default: - content["users_default"] = users_default - - if users: - if "users" in content: - content["users"].update(users) - else: - content["users"] = users - - # Remove any keys with value None - for user, power_level in list(content["users"].items()): - if power_level is None: - del content["users"][user] - self.client.api.set_power_levels(self.room_id, content) - return True - except MatrixRequestError: - return False - - def modify_required_power_levels(self, events=None, **kwargs): - """Modifies room power level requirements. - - Args: - events(dict): Power levels required for sending specific event types, - in the form {"m.room.whatever0": 60, "m.room.whatever2": None}. - Overrides events_default and state_default for the specified - events. A level of None causes the target event to revert to the - default level as specified by events_default or state_default. - **kwargs: Key/value pairs specifying the power levels required for - various actions: - - - events_default(int): Default level for sending message events - - state_default(int): Default level for sending state events - - invite(int): Inviting a user - - redact(int): Redacting an event - - ban(int): Banning a user - - kick(int): Kicking a user - - Returns: - True if successful, False if not - """ - try: - content = self.client.api.get_power_levels(self.room_id) - content.update(kwargs) - for key, value in list(content.items()): - if value is None: - del content[key] - - if events: - if "events" in content: - content["events"].update(events) - else: - content["events"] = events - - # Remove any keys with value None - for event, power_level in list(content["events"].items()): - if power_level is None: - del content["events"][event] - - self.client.api.set_power_levels(self.room_id, content) - return True - except MatrixRequestError: - return False - - def set_invite_only(self, invite_only): - """Set how the room can be joined. - - Args: - invite_only(bool): If True, users will have to be invited to join - the room. If False, anyone who knows the room link can join. - - Returns: - True if successful, False if not - """ - join_rule = "invite" if invite_only else "public" - try: - self.client.api.set_join_rule(self.room_id, join_rule) - self.invite_only = invite_only - return True - except MatrixRequestError: - return False - - def set_guest_access(self, allow_guests): - """Set whether guests can join the room and return True if successful.""" - guest_access = "can_join" if allow_guests else "forbidden" - try: - self.client.api.set_guest_access(self.room_id, guest_access) - self.guest_access = allow_guests - return True - except MatrixRequestError: - return False - - def enable_encryption(self): - """Enables encryption in the room. - - NOTE: Once enabled, encryption cannot be disabled. - - Returns: - True if successful, False if not - """ - try: - self.send_state_event("m.room.encryption", - {"algorithm": "m.megolm.v1.aes-sha2"}) - self.encrypted = True - return True - except MatrixRequestError: - return False - - def _process_state_event(self, state_event): - if "type" not in state_event: - return # Ignore event - etype = state_event["type"] - econtent = state_event["content"] - clevel = self.client._cache_level - - # Don't keep track of room state if caching turned off - if clevel >= 0: - try: - if etype == "m.room.name": - self.name = econtent.get("name") - elif etype == "m.room.canonical_alias": - self.canonical_alias = econtent.get("alias") - elif etype == "m.room.topic": - self.topic = econtent.get("topic") - elif etype == "m.room.aliases": - self.aliases = econtent.get("aliases") - elif etype == "m.room.join_rules": - self.invite_only = econtent["join_rule"] == "invite" - elif etype == "m.room.guest_access": - self.guest_access = econtent["guest_access"] == "can_join" - elif etype == "m.room.encryption": - if econtent.get("algorithm") == "m.megolm.v1.aes-sha2": - self.encrypted = True - elif etype == "m.room.member" and clevel == clevel.ALL: - # tracking room members can be large e.g. #matrix:matrix.org - if econtent["membership"] == "join": - user_id = state_event["state_key"] - self._add_member(user_id, econtent.get("displayname")) - elif econtent["membership"] in ("leave", "kick", "invite"): - self._members.pop(state_event["state_key"], None) - except KeyError: - logger.exception("Unable to parse state event %s, passing over.", - state_event['event_id']) - - for listener in self.state_listeners: - if ( - listener['event_type'] is None or - listener['event_type'] == state_event['type'] - ): - listener['callback'](state_event) - - @property - def prev_batch(self): - return self._prev_batch - - @prev_batch.setter - def prev_batch(self, prev_batch): - self._prev_batch = prev_batch diff --git a/venv/lib/python3.10/site-packages/matrix_client/user.py b/venv/lib/python3.10/site-packages/matrix_client/user.py deleted file mode 100644 index e56a89ef32569ee093c09f4311a7c1188fb87112..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/matrix_client/user.py +++ /dev/null @@ -1,77 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2015 OpenMarket Ltd -# -# 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 warnings import warn - -from .checks import check_user_id - - -class User(object): - """ The User class can be used to call user specific functions. - """ - def __init__(self, api, user_id, displayname=None): - check_user_id(user_id) - - self.user_id = user_id - self.displayname = displayname - self.api = api - - def get_display_name(self, room=None): - """Get this user's display name. - - Args: - room (Room): Optional. When specified, return the display name of the user - in this room. - - Returns: - The display name. Defaults to the user ID if not set. - """ - if room: - try: - return room.members_displaynames[self.user_id] - except KeyError: - return self.user_id - if not self.displayname: - self.displayname = self.api.get_display_name(self.user_id) - return self.displayname or self.user_id - - def get_friendly_name(self): - """Deprecated. Use :meth:`get_display_name` instead.""" - warn("get_friendly_name is deprecated. Use get_display_name instead.", - DeprecationWarning) - return self.get_display_name() - - def set_display_name(self, display_name): - """ Set this users display name. - - Args: - display_name (str): Display Name - """ - self.displayname = display_name - return self.api.set_display_name(self.user_id, display_name) - - def get_avatar_url(self): - mxcurl = self.api.get_avatar_url(self.user_id) - url = None - if mxcurl is not None: - url = self.api.get_download_url(mxcurl) - return url - - def set_avatar_url(self, avatar_url): - """ Set this users avatar. - - Args: - avatar_url (str): mxc url from previously uploaded - """ - return self.api.set_avatar_url(self.user_id, avatar_url) diff --git a/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/LICENSE b/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/LICENSE deleted file mode 100644 index 2a920c59d8abdd485a774087915986448495fd7c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/LICENSE +++ /dev/null @@ -1,46 +0,0 @@ -Copyright (c) 2015 Vitaly Puzrin, Alex Kocharin. -Copyright (c) 2021 Taneli Hukkinen - -Permission is hereby granted, free of charge, to any person -obtaining a copy of this software and associated documentation -files (the "Software"), to deal in the Software without -restriction, including without limitation the rights to use, -copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the -Software is furnished to do so, subject to the following -conditions: - -The above copyright notice and this permission notice shall be -included in all copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, -EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES -OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND -NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT -HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, -WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR -OTHER DEALINGS IN THE SOFTWARE. - --------------------------------------------------------------------------------- - -.parse() is based on Joyent's node.js `url` code: - -Copyright Joyent, Inc. and other Node contributors. All rights reserved. -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to -deal in the Software without restriction, including without limitation the -rights to use, copy, modify, merge, publish, distribute, sublicense, and/or -sell copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in -all copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS -IN THE SOFTWARE. diff --git a/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/METADATA b/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/METADATA deleted file mode 100644 index b4670e86b6dc207c944c55c5d3b84911fb41157a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/METADATA +++ /dev/null @@ -1,32 +0,0 @@ -Metadata-Version: 2.1 -Name: mdurl -Version: 0.1.2 -Summary: Markdown URL utilities -Keywords: markdown,commonmark -Author-email: Taneli Hukkinen -Requires-Python: >=3.7 -Description-Content-Type: text/markdown -Classifier: License :: OSI Approved :: MIT License -Classifier: Operating System :: MacOS -Classifier: Operating System :: Microsoft :: Windows -Classifier: Operating System :: POSIX :: Linux -Classifier: Programming Language :: Python :: 3 :: Only -Classifier: Programming Language :: Python :: 3.7 -Classifier: Programming Language :: Python :: 3.8 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Programming Language :: Python :: 3.10 -Classifier: Programming Language :: Python :: Implementation :: CPython -Classifier: Programming Language :: Python :: Implementation :: PyPy -Classifier: Topic :: Software Development :: Libraries :: Python Modules -Classifier: Typing :: Typed -Project-URL: Homepage, https://github.com/executablebooks/mdurl - -# mdurl - -[![Build Status](https://github.com/executablebooks/mdurl/workflows/Tests/badge.svg?branch=master)](https://github.com/executablebooks/mdurl/actions?query=workflow%3ATests+branch%3Amaster+event%3Apush) -[![codecov.io](https://codecov.io/gh/executablebooks/mdurl/branch/master/graph/badge.svg)](https://codecov.io/gh/executablebooks/mdurl) -[![PyPI version](https://img.shields.io/pypi/v/mdurl)](https://pypi.org/project/mdurl) - -This is a Python port of the JavaScript [mdurl](https://www.npmjs.com/package/mdurl) package. -See the [upstream README.md file](https://github.com/markdown-it/mdurl/blob/master/README.md) for API documentation. - diff --git a/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/RECORD b/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/RECORD deleted file mode 100644 index 594bce3c8502cf7a0f3d2cf18c394d389a2b3e82..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/RECORD +++ /dev/null @@ -1,18 +0,0 @@ -mdurl-0.1.2.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -mdurl-0.1.2.dist-info/LICENSE,sha256=fGBd9uKGZ6lgMRjpgnT2SknOPu0NJvzM6VNKNF4O-VU,2338 -mdurl-0.1.2.dist-info/METADATA,sha256=tTsp1I9Jk2cFP9o8gefOJ9JVg4Drv4PmYCOwLrfd0l0,1638 -mdurl-0.1.2.dist-info/RECORD,, -mdurl-0.1.2.dist-info/WHEEL,sha256=4TfKIB_xu-04bc2iKz6_zFt-gEFEEDU_31HGhqzOCE8,81 -mdurl/__init__.py,sha256=1vpE89NyXniIRZNC_4f6BPm3Ub4bPntjfyyhLRR7opU,547 -mdurl/__pycache__/__init__.cpython-310.pyc,, -mdurl/__pycache__/_decode.cpython-310.pyc,, -mdurl/__pycache__/_encode.cpython-310.pyc,, -mdurl/__pycache__/_format.cpython-310.pyc,, -mdurl/__pycache__/_parse.cpython-310.pyc,, -mdurl/__pycache__/_url.cpython-310.pyc,, -mdurl/_decode.py,sha256=3Q_gDQqU__TvDbu7x-b9LjbVl4QWy5g_qFwljcuvN_Y,3004 -mdurl/_encode.py,sha256=goJLUFt1h4rVZNqqm9t15Nw2W-bFXYQEy3aR01ImWvs,2602 -mdurl/_format.py,sha256=xZct0mdePXA0H3kAqxjGtlB5O86G35DAYMGkA44CmB4,626 -mdurl/_parse.py,sha256=ezZSkM2_4NQ2Zx047sEdcJG7NYQRFHiZK7Y8INHFzwY,11374 -mdurl/_url.py,sha256=5kQnRQN2A_G4svLnRzZcG0bfoD9AbBrYDXousDHZ3z0,284 -mdurl/py.typed,sha256=8PjyZ1aVoQpRVvt71muvuq5qE-jTFZkK-GLHkhdebmc,26 diff --git a/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/WHEEL b/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/WHEEL deleted file mode 100644 index 668ba4d0151c5c76ed6e758061daa8c1b0bf5d21..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/mdurl-0.1.2.dist-info/WHEEL +++ /dev/null @@ -1,4 +0,0 @@ -Wheel-Version: 1.0 -Generator: flit 3.7.1 -Root-Is-Purelib: true -Tag: py3-none-any diff --git a/venv/lib/python3.10/site-packages/mdurl/__init__.py b/venv/lib/python3.10/site-packages/mdurl/__init__.py deleted file mode 100644 index cdbb640e004cef0e950a656a53d92d89d82c7472..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/mdurl/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -__all__ = ( - "decode", - "DECODE_DEFAULT_CHARS", - "DECODE_COMPONENT_CHARS", - "encode", - "ENCODE_DEFAULT_CHARS", - "ENCODE_COMPONENT_CHARS", - "format", - "parse", - "URL", -) -__version__ = "0.1.2" # DO NOT EDIT THIS LINE MANUALLY. LET bump2version UTILITY DO IT - -from mdurl._decode import DECODE_COMPONENT_CHARS, DECODE_DEFAULT_CHARS, decode -from mdurl._encode import ENCODE_COMPONENT_CHARS, ENCODE_DEFAULT_CHARS, encode -from mdurl._format import format -from mdurl._parse import url_parse as parse -from mdurl._url import URL diff --git a/venv/lib/python3.10/site-packages/mdurl/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/mdurl/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 5217afc8f1f281d67c1009fafb8ae3f19598c674..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/mdurl/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/mdurl/__pycache__/_decode.cpython-310.pyc b/venv/lib/python3.10/site-packages/mdurl/__pycache__/_decode.cpython-310.pyc deleted file mode 100644 index 3a49d217278f61cde099f9e4123c9638314876f9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/mdurl/__pycache__/_decode.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/mdurl/__pycache__/_encode.cpython-310.pyc b/venv/lib/python3.10/site-packages/mdurl/__pycache__/_encode.cpython-310.pyc deleted file mode 100644 index 2deb55ffb995ce4af043820232494d01bdd3cb83..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/mdurl/__pycache__/_encode.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/mdurl/__pycache__/_format.cpython-310.pyc b/venv/lib/python3.10/site-packages/mdurl/__pycache__/_format.cpython-310.pyc deleted file mode 100644 index a859d112df65a92138970b5374821ebb411e03d9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/mdurl/__pycache__/_format.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/mdurl/__pycache__/_parse.cpython-310.pyc b/venv/lib/python3.10/site-packages/mdurl/__pycache__/_parse.cpython-310.pyc deleted file mode 100644 index d0c3322435fbe3ff86b79028112fef0b62f91eb5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/mdurl/__pycache__/_parse.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/mdurl/__pycache__/_url.cpython-310.pyc b/venv/lib/python3.10/site-packages/mdurl/__pycache__/_url.cpython-310.pyc deleted file mode 100644 index 61e96f1f84bfd713d8b0f79cebf6d95269a7f39c..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/mdurl/__pycache__/_url.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/mdurl/_decode.py b/venv/lib/python3.10/site-packages/mdurl/_decode.py deleted file mode 100644 index 9b50a2dde976a6d43491ec6f20d12e60f6f6597f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/mdurl/_decode.py +++ /dev/null @@ -1,104 +0,0 @@ -from __future__ import annotations - -from collections.abc import Sequence -import functools -import re - -DECODE_DEFAULT_CHARS = ";/?:@&=+$,#" -DECODE_COMPONENT_CHARS = "" - -decode_cache: dict[str, list[str]] = {} - - -def get_decode_cache(exclude: str) -> Sequence[str]: - if exclude in decode_cache: - return decode_cache[exclude] - - cache: list[str] = [] - decode_cache[exclude] = cache - - for i in range(128): - ch = chr(i) - cache.append(ch) - - for i in range(len(exclude)): - ch_code = ord(exclude[i]) - cache[ch_code] = "%" + ("0" + hex(ch_code)[2:].upper())[-2:] - - return cache - - -# Decode percent-encoded string. -# -def decode(string: str, exclude: str = DECODE_DEFAULT_CHARS) -> str: - cache = get_decode_cache(exclude) - repl_func = functools.partial(repl_func_with_cache, cache=cache) - return re.sub(r"(%[a-f0-9]{2})+", repl_func, string, flags=re.IGNORECASE) - - -def repl_func_with_cache(match: re.Match, cache: Sequence[str]) -> str: - seq = match.group() - result = "" - - i = 0 - l = len(seq) # noqa: E741 - while i < l: - b1 = int(seq[i + 1 : i + 3], 16) - - if b1 < 0x80: - result += cache[b1] - i += 3 # emulate JS for loop statement3 - continue - - if (b1 & 0xE0) == 0xC0 and (i + 3 < l): - # 110xxxxx 10xxxxxx - b2 = int(seq[i + 4 : i + 6], 16) - - if (b2 & 0xC0) == 0x80: - all_bytes = bytes((b1, b2)) - try: - result += all_bytes.decode() - except UnicodeDecodeError: - result += "\ufffd" * 2 - - i += 3 - i += 3 # emulate JS for loop statement3 - continue - - if (b1 & 0xF0) == 0xE0 and (i + 6 < l): - # 1110xxxx 10xxxxxx 10xxxxxx - b2 = int(seq[i + 4 : i + 6], 16) - b3 = int(seq[i + 7 : i + 9], 16) - - if (b2 & 0xC0) == 0x80 and (b3 & 0xC0) == 0x80: - all_bytes = bytes((b1, b2, b3)) - try: - result += all_bytes.decode() - except UnicodeDecodeError: - result += "\ufffd" * 3 - - i += 6 - i += 3 # emulate JS for loop statement3 - continue - - if (b1 & 0xF8) == 0xF0 and (i + 9 < l): - # 111110xx 10xxxxxx 10xxxxxx 10xxxxxx - b2 = int(seq[i + 4 : i + 6], 16) - b3 = int(seq[i + 7 : i + 9], 16) - b4 = int(seq[i + 10 : i + 12], 16) - - if (b2 & 0xC0) == 0x80 and (b3 & 0xC0) == 0x80 and (b4 & 0xC0) == 0x80: - all_bytes = bytes((b1, b2, b3, b4)) - try: - result += all_bytes.decode() - except UnicodeDecodeError: - result += "\ufffd" * 4 - - i += 9 - i += 3 # emulate JS for loop statement3 - continue - - result += "\ufffd" - i += 3 # emulate JS for loop statement3 - - return result diff --git a/venv/lib/python3.10/site-packages/mdurl/_encode.py b/venv/lib/python3.10/site-packages/mdurl/_encode.py deleted file mode 100644 index bc2e5b917afe9e9ecaa6f11af7a9ac82704d3914..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/mdurl/_encode.py +++ /dev/null @@ -1,85 +0,0 @@ -from __future__ import annotations - -from collections.abc import Sequence -from string import ascii_letters, digits, hexdigits -from urllib.parse import quote as encode_uri_component - -ASCII_LETTERS_AND_DIGITS = ascii_letters + digits - -ENCODE_DEFAULT_CHARS = ";/?:@&=+$,-_.!~*'()#" -ENCODE_COMPONENT_CHARS = "-_.!~*'()" - -encode_cache: dict[str, list[str]] = {} - - -# Create a lookup array where anything but characters in `chars` string -# and alphanumeric chars is percent-encoded. -def get_encode_cache(exclude: str) -> Sequence[str]: - if exclude in encode_cache: - return encode_cache[exclude] - - cache: list[str] = [] - encode_cache[exclude] = cache - - for i in range(128): - ch = chr(i) - - if ch in ASCII_LETTERS_AND_DIGITS: - # always allow unencoded alphanumeric characters - cache.append(ch) - else: - cache.append("%" + ("0" + hex(i)[2:].upper())[-2:]) - - for i in range(len(exclude)): - cache[ord(exclude[i])] = exclude[i] - - return cache - - -# Encode unsafe characters with percent-encoding, skipping already -# encoded sequences. -# -# - string - string to encode -# - exclude - list of characters to ignore (in addition to a-zA-Z0-9) -# - keepEscaped - don't encode '%' in a correct escape sequence (default: true) -def encode( - string: str, exclude: str = ENCODE_DEFAULT_CHARS, *, keep_escaped: bool = True -) -> str: - result = "" - - cache = get_encode_cache(exclude) - - l = len(string) # noqa: E741 - i = 0 - while i < l: - code = ord(string[i]) - - # % - if keep_escaped and code == 0x25 and i + 2 < l: - if all(c in hexdigits for c in string[i + 1 : i + 3]): - result += string[i : i + 3] - i += 2 - i += 1 # JS for loop statement3 - continue - - if code < 128: - result += cache[code] - i += 1 # JS for loop statement3 - continue - - if code >= 0xD800 and code <= 0xDFFF: - if code >= 0xD800 and code <= 0xDBFF and i + 1 < l: - next_code = ord(string[i + 1]) - if next_code >= 0xDC00 and next_code <= 0xDFFF: - result += encode_uri_component(string[i] + string[i + 1]) - i += 1 - i += 1 # JS for loop statement3 - continue - result += "%EF%BF%BD" - i += 1 # JS for loop statement3 - continue - - result += encode_uri_component(string[i]) - i += 1 # JS for loop statement3 - - return result diff --git a/venv/lib/python3.10/site-packages/mdurl/_format.py b/venv/lib/python3.10/site-packages/mdurl/_format.py deleted file mode 100644 index 12524ca626065183ec9974f3d7d08dadd4a7d3e8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/mdurl/_format.py +++ /dev/null @@ -1,27 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from mdurl._url import URL - - -def format(url: URL) -> str: # noqa: A001 - result = "" - - result += url.protocol or "" - result += "//" if url.slashes else "" - result += url.auth + "@" if url.auth else "" - - if url.hostname and ":" in url.hostname: - # ipv6 address - result += "[" + url.hostname + "]" - else: - result += url.hostname or "" - - result += ":" + url.port if url.port else "" - result += url.pathname or "" - result += url.search or "" - result += url.hash or "" - - return result diff --git a/venv/lib/python3.10/site-packages/mdurl/_parse.py b/venv/lib/python3.10/site-packages/mdurl/_parse.py deleted file mode 100644 index ffeeac768dca3bff60c55c9b1f0bc0fbb4cec7b1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/mdurl/_parse.py +++ /dev/null @@ -1,304 +0,0 @@ -# Copyright Joyent, Inc. and other Node contributors. -# -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the -# "Software"), to deal in the Software without restriction, including -# without limitation the rights to use, copy, modify, merge, publish, -# distribute, sublicense, and/or sell copies of the Software, and to permit -# persons to whom the Software is furnished to do so, subject to the -# following conditions: -# -# The above copyright notice and this permission notice shall be included -# in all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS -# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF -# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN -# NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, -# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR -# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE -# USE OR OTHER DEALINGS IN THE SOFTWARE. - - -# Changes from joyent/node: -# -# 1. No leading slash in paths, -# e.g. in `url.parse('http://foo?bar')` pathname is ``, not `/` -# -# 2. Backslashes are not replaced with slashes, -# so `http:\\example.org\` is treated like a relative path -# -# 3. Trailing colon is treated like a part of the path, -# i.e. in `http://example.org:foo` pathname is `:foo` -# -# 4. Nothing is URL-encoded in the resulting object, -# (in joyent/node some chars in auth and paths are encoded) -# -# 5. `url.parse()` does not have `parseQueryString` argument -# -# 6. Removed extraneous result properties: `host`, `path`, `query`, etc., -# which can be constructed using other parts of the url. - -from __future__ import annotations - -from collections import defaultdict -import re - -from mdurl._url import URL - -# Reference: RFC 3986, RFC 1808, RFC 2396 - -# define these here so at least they only have to be -# compiled once on the first module load. -PROTOCOL_PATTERN = re.compile(r"^([a-z0-9.+-]+:)", flags=re.IGNORECASE) -PORT_PATTERN = re.compile(r":[0-9]*$") - -# Special case for a simple path URL -SIMPLE_PATH_PATTERN = re.compile(r"^(//?(?!/)[^?\s]*)(\?[^\s]*)?$") - -# RFC 2396: characters reserved for delimiting URLs. -# We actually just auto-escape these. -DELIMS = ("<", ">", '"', "`", " ", "\r", "\n", "\t") - -# RFC 2396: characters not allowed for various reasons. -UNWISE = ("{", "}", "|", "\\", "^", "`") + DELIMS - -# Allowed by RFCs, but cause of XSS attacks. Always escape these. -AUTO_ESCAPE = ("'",) + UNWISE -# Characters that are never ever allowed in a hostname. -# Note that any invalid chars are also handled, but these -# are the ones that are *expected* to be seen, so we fast-path -# them. -NON_HOST_CHARS = ("%", "/", "?", ";", "#") + AUTO_ESCAPE -HOST_ENDING_CHARS = ("/", "?", "#") -HOSTNAME_MAX_LEN = 255 -HOSTNAME_PART_PATTERN = re.compile(r"^[+a-z0-9A-Z_-]{0,63}$") -HOSTNAME_PART_START = re.compile(r"^([+a-z0-9A-Z_-]{0,63})(.*)$") -# protocols that can allow "unsafe" and "unwise" chars. - -# protocols that never have a hostname. -HOSTLESS_PROTOCOL = defaultdict( - bool, - { - "javascript": True, - "javascript:": True, - }, -) -# protocols that always contain a // bit. -SLASHED_PROTOCOL = defaultdict( - bool, - { - "http": True, - "https": True, - "ftp": True, - "gopher": True, - "file": True, - "http:": True, - "https:": True, - "ftp:": True, - "gopher:": True, - "file:": True, - }, -) - - -class MutableURL: - def __init__(self) -> None: - self.protocol: str | None = None - self.slashes: bool = False - self.auth: str | None = None - self.port: str | None = None - self.hostname: str | None = None - self.hash: str | None = None - self.search: str | None = None - self.pathname: str | None = None - - def parse(self, url: str, slashes_denote_host: bool) -> "MutableURL": - lower_proto = "" - slashes = False - rest = url - - # trim before proceeding. - # This is to support parse stuff like " http://foo.com \n" - rest = rest.strip() - - if not slashes_denote_host and len(url.split("#")) == 1: - # Try fast path regexp - simple_path = SIMPLE_PATH_PATTERN.match(rest) - if simple_path: - self.pathname = simple_path.group(1) - if simple_path.group(2): - self.search = simple_path.group(2) - return self - - proto = "" - proto_match = PROTOCOL_PATTERN.match(rest) - if proto_match: - proto = proto_match.group() - lower_proto = proto.lower() - self.protocol = proto - rest = rest[len(proto) :] - - # figure out if it's got a host - # user@server is *always* interpreted as a hostname, and url - # resolution will treat //foo/bar as host=foo,path=bar because that's - # how the browser resolves relative URLs. - if slashes_denote_host or proto or re.search(r"^//[^@/]+@[^@/]+", rest): - slashes = rest.startswith("//") - if slashes and not (proto and HOSTLESS_PROTOCOL[proto]): - rest = rest[2:] - self.slashes = True - - if not HOSTLESS_PROTOCOL[proto] and ( - slashes or (proto and not SLASHED_PROTOCOL[proto]) - ): - - # there's a hostname. - # the first instance of /, ?, ;, or # ends the host. - # - # If there is an @ in the hostname, then non-host chars *are* allowed - # to the left of the last @ sign, unless some host-ending character - # comes *before* the @-sign. - # URLs are obnoxious. - # - # ex: - # http://a@b@c/ => user:a@b host:c - # http://a@b?@c => user:a host:c path:/?@c - - # v0.12 TODO(isaacs): This is not quite how Chrome does things. - # Review our test case against browsers more comprehensively. - - # find the first instance of any hostEndingChars - host_end = -1 - for i in range(len(HOST_ENDING_CHARS)): - hec = rest.find(HOST_ENDING_CHARS[i]) - if hec != -1 and (host_end == -1 or hec < host_end): - host_end = hec - - # at this point, either we have an explicit point where the - # auth portion cannot go past, or the last @ char is the decider. - if host_end == -1: - # atSign can be anywhere. - at_sign = rest.rfind("@") - else: - # atSign must be in auth portion. - # http://a@b/c@d => host:b auth:a path:/c@d - at_sign = rest.rfind("@", 0, host_end + 1) - - # Now we have a portion which is definitely the auth. - # Pull that off. - if at_sign != -1: - auth = rest[:at_sign] - rest = rest[at_sign + 1 :] - self.auth = auth - - # the host is the remaining to the left of the first non-host char - host_end = -1 - for i in range(len(NON_HOST_CHARS)): - hec = rest.find(NON_HOST_CHARS[i]) - if hec != -1 and (host_end == -1 or hec < host_end): - host_end = hec - # if we still have not hit it, then the entire thing is a host. - if host_end == -1: - host_end = len(rest) - - if host_end > 0 and rest[host_end - 1] == ":": - host_end -= 1 - host = rest[:host_end] - rest = rest[host_end:] - - # pull out port. - self.parse_host(host) - - # we've indicated that there is a hostname, - # so even if it's empty, it has to be present. - self.hostname = self.hostname or "" - - # if hostname begins with [ and ends with ] - # assume that it's an IPv6 address. - ipv6_hostname = self.hostname.startswith("[") and self.hostname.endswith( - "]" - ) - - # validate a little. - if not ipv6_hostname: - hostparts = self.hostname.split(".") - l = len(hostparts) # noqa: E741 - i = 0 - while i < l: - part = hostparts[i] - if not part: - i += 1 # emulate statement3 in JS for loop - continue - if not HOSTNAME_PART_PATTERN.search(part): - newpart = "" - k = len(part) - j = 0 - while j < k: - if ord(part[j]) > 127: - # we replace non-ASCII char with a temporary placeholder - # we need this to make sure size of hostname is not - # broken by replacing non-ASCII by nothing - newpart += "x" - else: - newpart += part[j] - j += 1 # emulate statement3 in JS for loop - - # we test again with ASCII char only - if not HOSTNAME_PART_PATTERN.search(newpart): - valid_parts = hostparts[:i] - not_host = hostparts[i + 1 :] - bit = HOSTNAME_PART_START.search(part) - if bit: - valid_parts.append(bit.group(1)) - not_host.insert(0, bit.group(2)) - if not_host: - rest = ".".join(not_host) + rest - self.hostname = ".".join(valid_parts) - break - i += 1 # emulate statement3 in JS for loop - - if len(self.hostname) > HOSTNAME_MAX_LEN: - self.hostname = "" - - # strip [ and ] from the hostname - # the host field still retains them, though - if ipv6_hostname: - self.hostname = self.hostname[1:-1] - - # chop off from the tail first. - hash = rest.find("#") # noqa: A001 - if hash != -1: - # got a fragment string. - self.hash = rest[hash:] - rest = rest[:hash] - qm = rest.find("?") - if qm != -1: - self.search = rest[qm:] - rest = rest[:qm] - if rest: - self.pathname = rest - if SLASHED_PROTOCOL[lower_proto] and self.hostname and not self.pathname: - self.pathname = "" - - return self - - def parse_host(self, host: str) -> None: - port_match = PORT_PATTERN.search(host) - if port_match: - port = port_match.group() - if port != ":": - self.port = port[1:] - host = host[: -len(port)] - if host: - self.hostname = host - - -def url_parse(url: URL | str, *, slashes_denote_host: bool = False) -> URL: - if isinstance(url, URL): - return url - u = MutableURL() - u.parse(url, slashes_denote_host) - return URL( - u.protocol, u.slashes, u.auth, u.port, u.hostname, u.hash, u.search, u.pathname - ) diff --git a/venv/lib/python3.10/site-packages/mdurl/_url.py b/venv/lib/python3.10/site-packages/mdurl/_url.py deleted file mode 100644 index f866e7a179c8854e37c9bba6294f48681e5d99d7..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/mdurl/_url.py +++ /dev/null @@ -1,14 +0,0 @@ -from __future__ import annotations - -from typing import NamedTuple - - -class URL(NamedTuple): - protocol: str | None - slashes: bool - auth: str | None - port: str | None - hostname: str | None - hash: str | None # noqa: A003 - search: str | None - pathname: str | None diff --git a/venv/lib/python3.10/site-packages/mdurl/py.typed b/venv/lib/python3.10/site-packages/mdurl/py.typed deleted file mode 100644 index 7632ecf77545c5e5501cb3fc5719df0761104ca2..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/mdurl/py.typed +++ /dev/null @@ -1 +0,0 @@ -# Marker file for PEP 561 diff --git a/venv/lib/python3.10/site-packages/nacl/__init__.py b/venv/lib/python3.10/site-packages/nacl/__init__.py deleted file mode 100644 index cc8b7ea7b35be88f03c2871ca31065b15c113e2a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/__init__.py +++ /dev/null @@ -1,39 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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. - - -__all__ = [ - "__title__", - "__summary__", - "__uri__", - "__version__", - "__author__", - "__email__", - "__license__", - "__copyright__", -] - -__title__ = "PyNaCl" -__summary__ = ( - "Python binding to the Networking and Cryptography (NaCl) library" -) -__uri__ = "https://github.com/pyca/pynacl/" - -__version__ = "1.5.0" - -__author__ = "The PyNaCl developers" -__email__ = "cryptography-dev@python.org" - -__license__ = "Apache License 2.0" -__copyright__ = "Copyright 2013-2018 {}".format(__author__) diff --git a/venv/lib/python3.10/site-packages/nacl/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index f37c95bc860a151e3419ec1473a60196c6ea8f92..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/__pycache__/encoding.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/__pycache__/encoding.cpython-310.pyc deleted file mode 100644 index 031ed7ce030cd63788a16c56cebc57e430eea97b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/__pycache__/encoding.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/__pycache__/exceptions.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/__pycache__/exceptions.cpython-310.pyc deleted file mode 100644 index 67c143cfbb09bdc54720e6623d6d2ab91dfb14a5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/__pycache__/exceptions.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/__pycache__/hash.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/__pycache__/hash.cpython-310.pyc deleted file mode 100644 index 18e4aec0bc68edf8e7b4cc19dff69c138bb03931..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/__pycache__/hash.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/__pycache__/hashlib.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/__pycache__/hashlib.cpython-310.pyc deleted file mode 100644 index 0b5eb48f0927fa524d270ff5bad9b4faff86f59f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/__pycache__/hashlib.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/__pycache__/public.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/__pycache__/public.cpython-310.pyc deleted file mode 100644 index 752d3769722e8a81db77fd27010ecaa15e93b53d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/__pycache__/public.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/__pycache__/secret.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/__pycache__/secret.cpython-310.pyc deleted file mode 100644 index a955fe89a36e43817a6e1ebb5efb93be702f6a8a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/__pycache__/secret.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/__pycache__/signing.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/__pycache__/signing.cpython-310.pyc deleted file mode 100644 index 98b0e93e32a690ad1423ca98c27e3225f8ee56d6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/__pycache__/signing.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index 9d359e97919463281584d34b974f0fc2bb9b8d02..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/_sodium.abi3.so b/venv/lib/python3.10/site-packages/nacl/_sodium.abi3.so deleted file mode 100644 index 3a80ac894f93abb68f99ab4e4cdad1d1b63a67e1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/_sodium.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b89e91c129db6fdc0ee1eb11d1b554aab7c51ed04e1a6df821021d99a1357c66 -size 2740136 diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__init__.py b/venv/lib/python3.10/site-packages/nacl/bindings/__init__.py deleted file mode 100644 index 1e3b10e82fd6e2da413e1f4999a6d25d7fd919c8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/__init__.py +++ /dev/null @@ -1,451 +0,0 @@ -# Copyright 2013-2019 Donald Stufft and individual contributors -# -# 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 nacl.bindings.crypto_aead import ( - crypto_aead_chacha20poly1305_ABYTES, - crypto_aead_chacha20poly1305_KEYBYTES, - crypto_aead_chacha20poly1305_MESSAGEBYTES_MAX, - crypto_aead_chacha20poly1305_NPUBBYTES, - crypto_aead_chacha20poly1305_NSECBYTES, - crypto_aead_chacha20poly1305_decrypt, - crypto_aead_chacha20poly1305_encrypt, - crypto_aead_chacha20poly1305_ietf_ABYTES, - crypto_aead_chacha20poly1305_ietf_KEYBYTES, - crypto_aead_chacha20poly1305_ietf_MESSAGEBYTES_MAX, - crypto_aead_chacha20poly1305_ietf_NPUBBYTES, - crypto_aead_chacha20poly1305_ietf_NSECBYTES, - crypto_aead_chacha20poly1305_ietf_decrypt, - crypto_aead_chacha20poly1305_ietf_encrypt, - crypto_aead_xchacha20poly1305_ietf_ABYTES, - crypto_aead_xchacha20poly1305_ietf_KEYBYTES, - crypto_aead_xchacha20poly1305_ietf_MESSAGEBYTES_MAX, - crypto_aead_xchacha20poly1305_ietf_NPUBBYTES, - crypto_aead_xchacha20poly1305_ietf_NSECBYTES, - crypto_aead_xchacha20poly1305_ietf_decrypt, - crypto_aead_xchacha20poly1305_ietf_encrypt, -) -from nacl.bindings.crypto_box import ( - crypto_box, - crypto_box_BEFORENMBYTES, - crypto_box_BOXZEROBYTES, - crypto_box_NONCEBYTES, - crypto_box_PUBLICKEYBYTES, - crypto_box_SEALBYTES, - crypto_box_SECRETKEYBYTES, - crypto_box_SEEDBYTES, - crypto_box_ZEROBYTES, - crypto_box_afternm, - crypto_box_beforenm, - crypto_box_keypair, - crypto_box_open, - crypto_box_open_afternm, - crypto_box_seal, - crypto_box_seal_open, - crypto_box_seed_keypair, -) -from nacl.bindings.crypto_core import ( - crypto_core_ed25519_BYTES, - crypto_core_ed25519_NONREDUCEDSCALARBYTES, - crypto_core_ed25519_SCALARBYTES, - crypto_core_ed25519_add, - crypto_core_ed25519_is_valid_point, - crypto_core_ed25519_scalar_add, - crypto_core_ed25519_scalar_complement, - crypto_core_ed25519_scalar_invert, - crypto_core_ed25519_scalar_mul, - crypto_core_ed25519_scalar_negate, - crypto_core_ed25519_scalar_reduce, - crypto_core_ed25519_scalar_sub, - crypto_core_ed25519_sub, - has_crypto_core_ed25519, -) -from nacl.bindings.crypto_generichash import ( - crypto_generichash_BYTES, - crypto_generichash_BYTES_MAX, - crypto_generichash_BYTES_MIN, - crypto_generichash_KEYBYTES, - crypto_generichash_KEYBYTES_MAX, - crypto_generichash_KEYBYTES_MIN, - crypto_generichash_PERSONALBYTES, - crypto_generichash_SALTBYTES, - crypto_generichash_STATEBYTES, - generichash_blake2b_final as crypto_generichash_blake2b_final, - generichash_blake2b_init as crypto_generichash_blake2b_init, - generichash_blake2b_salt_personal as crypto_generichash_blake2b_salt_personal, - generichash_blake2b_update as crypto_generichash_blake2b_update, -) -from nacl.bindings.crypto_hash import ( - crypto_hash, - crypto_hash_BYTES, - crypto_hash_sha256, - crypto_hash_sha256_BYTES, - crypto_hash_sha512, - crypto_hash_sha512_BYTES, -) -from nacl.bindings.crypto_kx import ( - crypto_kx_PUBLIC_KEY_BYTES, - crypto_kx_SECRET_KEY_BYTES, - crypto_kx_SEED_BYTES, - crypto_kx_SESSION_KEY_BYTES, - crypto_kx_client_session_keys, - crypto_kx_keypair, - crypto_kx_seed_keypair, - crypto_kx_server_session_keys, -) -from nacl.bindings.crypto_pwhash import ( - crypto_pwhash_ALG_ARGON2I13, - crypto_pwhash_ALG_ARGON2ID13, - crypto_pwhash_ALG_DEFAULT, - crypto_pwhash_BYTES_MAX, - crypto_pwhash_BYTES_MIN, - crypto_pwhash_PASSWD_MAX, - crypto_pwhash_PASSWD_MIN, - crypto_pwhash_SALTBYTES, - crypto_pwhash_STRBYTES, - crypto_pwhash_alg, - crypto_pwhash_argon2i_MEMLIMIT_INTERACTIVE, - crypto_pwhash_argon2i_MEMLIMIT_MAX, - crypto_pwhash_argon2i_MEMLIMIT_MIN, - crypto_pwhash_argon2i_MEMLIMIT_MODERATE, - crypto_pwhash_argon2i_MEMLIMIT_SENSITIVE, - crypto_pwhash_argon2i_OPSLIMIT_INTERACTIVE, - crypto_pwhash_argon2i_OPSLIMIT_MAX, - crypto_pwhash_argon2i_OPSLIMIT_MIN, - crypto_pwhash_argon2i_OPSLIMIT_MODERATE, - crypto_pwhash_argon2i_OPSLIMIT_SENSITIVE, - crypto_pwhash_argon2i_STRPREFIX, - crypto_pwhash_argon2id_MEMLIMIT_INTERACTIVE, - crypto_pwhash_argon2id_MEMLIMIT_MAX, - crypto_pwhash_argon2id_MEMLIMIT_MIN, - crypto_pwhash_argon2id_MEMLIMIT_MODERATE, - crypto_pwhash_argon2id_MEMLIMIT_SENSITIVE, - crypto_pwhash_argon2id_OPSLIMIT_INTERACTIVE, - crypto_pwhash_argon2id_OPSLIMIT_MAX, - crypto_pwhash_argon2id_OPSLIMIT_MIN, - crypto_pwhash_argon2id_OPSLIMIT_MODERATE, - crypto_pwhash_argon2id_OPSLIMIT_SENSITIVE, - crypto_pwhash_argon2id_STRPREFIX, - crypto_pwhash_scryptsalsa208sha256_BYTES_MAX, - crypto_pwhash_scryptsalsa208sha256_BYTES_MIN, - crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_INTERACTIVE, - crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_MAX, - crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_MIN, - crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_SENSITIVE, - crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_INTERACTIVE, - crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_MAX, - crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_MIN, - crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_SENSITIVE, - crypto_pwhash_scryptsalsa208sha256_PASSWD_MAX, - crypto_pwhash_scryptsalsa208sha256_PASSWD_MIN, - crypto_pwhash_scryptsalsa208sha256_SALTBYTES, - crypto_pwhash_scryptsalsa208sha256_STRBYTES, - crypto_pwhash_scryptsalsa208sha256_STRPREFIX, - crypto_pwhash_scryptsalsa208sha256_ll, - crypto_pwhash_scryptsalsa208sha256_str, - crypto_pwhash_scryptsalsa208sha256_str_verify, - crypto_pwhash_str_alg, - crypto_pwhash_str_verify, - has_crypto_pwhash_scryptsalsa208sha256, - nacl_bindings_pick_scrypt_params, -) -from nacl.bindings.crypto_scalarmult import ( - crypto_scalarmult, - crypto_scalarmult_BYTES, - crypto_scalarmult_SCALARBYTES, - crypto_scalarmult_base, - crypto_scalarmult_ed25519, - crypto_scalarmult_ed25519_BYTES, - crypto_scalarmult_ed25519_SCALARBYTES, - crypto_scalarmult_ed25519_base, - crypto_scalarmult_ed25519_base_noclamp, - crypto_scalarmult_ed25519_noclamp, - has_crypto_scalarmult_ed25519, -) -from nacl.bindings.crypto_secretbox import ( - crypto_secretbox, - crypto_secretbox_BOXZEROBYTES, - crypto_secretbox_KEYBYTES, - crypto_secretbox_MACBYTES, - crypto_secretbox_MESSAGEBYTES_MAX, - crypto_secretbox_NONCEBYTES, - crypto_secretbox_ZEROBYTES, - crypto_secretbox_open, -) -from nacl.bindings.crypto_secretstream import ( - crypto_secretstream_xchacha20poly1305_ABYTES, - crypto_secretstream_xchacha20poly1305_HEADERBYTES, - crypto_secretstream_xchacha20poly1305_KEYBYTES, - crypto_secretstream_xchacha20poly1305_STATEBYTES, - crypto_secretstream_xchacha20poly1305_TAG_FINAL, - crypto_secretstream_xchacha20poly1305_TAG_MESSAGE, - crypto_secretstream_xchacha20poly1305_TAG_PUSH, - crypto_secretstream_xchacha20poly1305_TAG_REKEY, - crypto_secretstream_xchacha20poly1305_init_pull, - crypto_secretstream_xchacha20poly1305_init_push, - crypto_secretstream_xchacha20poly1305_keygen, - crypto_secretstream_xchacha20poly1305_pull, - crypto_secretstream_xchacha20poly1305_push, - crypto_secretstream_xchacha20poly1305_rekey, - crypto_secretstream_xchacha20poly1305_state, -) -from nacl.bindings.crypto_shorthash import ( - BYTES as crypto_shorthash_siphash24_BYTES, - KEYBYTES as crypto_shorthash_siphash24_KEYBYTES, - XBYTES as crypto_shorthash_siphashx24_BYTES, - XKEYBYTES as crypto_shorthash_siphashx24_KEYBYTES, - crypto_shorthash_siphash24, - crypto_shorthash_siphashx24, - has_crypto_shorthash_siphashx24, -) -from nacl.bindings.crypto_sign import ( - crypto_sign, - crypto_sign_BYTES, - crypto_sign_PUBLICKEYBYTES, - crypto_sign_SECRETKEYBYTES, - crypto_sign_SEEDBYTES, - crypto_sign_ed25519_pk_to_curve25519, - crypto_sign_ed25519_sk_to_curve25519, - crypto_sign_ed25519_sk_to_pk, - crypto_sign_ed25519_sk_to_seed, - crypto_sign_ed25519ph_STATEBYTES, - crypto_sign_ed25519ph_final_create, - crypto_sign_ed25519ph_final_verify, - crypto_sign_ed25519ph_state, - crypto_sign_ed25519ph_update, - crypto_sign_keypair, - crypto_sign_open, - crypto_sign_seed_keypair, -) -from nacl.bindings.randombytes import ( - randombytes, - randombytes_buf_deterministic, -) -from nacl.bindings.sodium_core import sodium_init -from nacl.bindings.utils import ( - sodium_add, - sodium_increment, - sodium_memcmp, - sodium_pad, - sodium_unpad, -) - - -__all__ = [ - "crypto_aead_chacha20poly1305_ABYTES", - "crypto_aead_chacha20poly1305_KEYBYTES", - "crypto_aead_chacha20poly1305_MESSAGEBYTES_MAX", - "crypto_aead_chacha20poly1305_NPUBBYTES", - "crypto_aead_chacha20poly1305_NSECBYTES", - "crypto_aead_chacha20poly1305_decrypt", - "crypto_aead_chacha20poly1305_encrypt", - "crypto_aead_chacha20poly1305_ietf_ABYTES", - "crypto_aead_chacha20poly1305_ietf_KEYBYTES", - "crypto_aead_chacha20poly1305_ietf_MESSAGEBYTES_MAX", - "crypto_aead_chacha20poly1305_ietf_NPUBBYTES", - "crypto_aead_chacha20poly1305_ietf_NSECBYTES", - "crypto_aead_chacha20poly1305_ietf_decrypt", - "crypto_aead_chacha20poly1305_ietf_encrypt", - "crypto_aead_xchacha20poly1305_ietf_ABYTES", - "crypto_aead_xchacha20poly1305_ietf_KEYBYTES", - "crypto_aead_xchacha20poly1305_ietf_MESSAGEBYTES_MAX", - "crypto_aead_xchacha20poly1305_ietf_NPUBBYTES", - "crypto_aead_xchacha20poly1305_ietf_NSECBYTES", - "crypto_aead_xchacha20poly1305_ietf_decrypt", - "crypto_aead_xchacha20poly1305_ietf_encrypt", - "crypto_box_SECRETKEYBYTES", - "crypto_box_PUBLICKEYBYTES", - "crypto_box_SEEDBYTES", - "crypto_box_NONCEBYTES", - "crypto_box_ZEROBYTES", - "crypto_box_BOXZEROBYTES", - "crypto_box_BEFORENMBYTES", - "crypto_box_SEALBYTES", - "crypto_box_keypair", - "crypto_box", - "crypto_box_open", - "crypto_box_beforenm", - "crypto_box_afternm", - "crypto_box_open_afternm", - "crypto_box_seal", - "crypto_box_seal_open", - "crypto_box_seed_keypair", - "has_crypto_core_ed25519", - "crypto_core_ed25519_BYTES", - "crypto_core_ed25519_UNIFORMBYTES", - "crypto_core_ed25519_SCALARBYTES", - "crypto_core_ed25519_NONREDUCEDSCALARBYTES", - "crypto_core_ed25519_add", - "crypto_core_ed25519_from_uniform", - "crypto_core_ed25519_is_valid_point", - "crypto_core_ed25519_sub", - "crypto_core_ed25519_scalar_invert", - "crypto_core_ed25519_scalar_negate", - "crypto_core_ed25519_scalar_complement", - "crypto_core_ed25519_scalar_add", - "crypto_core_ed25519_scalar_sub", - "crypto_core_ed25519_scalar_mul", - "crypto_core_ed25519_scalar_reduce", - "crypto_hash_BYTES", - "crypto_hash_sha256_BYTES", - "crypto_hash_sha512_BYTES", - "crypto_hash", - "crypto_hash_sha256", - "crypto_hash_sha512", - "crypto_generichash_BYTES", - "crypto_generichash_BYTES_MIN", - "crypto_generichash_BYTES_MAX", - "crypto_generichash_KEYBYTES", - "crypto_generichash_KEYBYTES_MIN", - "crypto_generichash_KEYBYTES_MAX", - "crypto_generichash_SALTBYTES", - "crypto_generichash_PERSONALBYTES", - "crypto_generichash_STATEBYTES", - "crypto_generichash_blake2b_salt_personal", - "crypto_generichash_blake2b_init", - "crypto_generichash_blake2b_update", - "crypto_generichash_blake2b_final", - "crypto_kx_keypair", - "crypto_kx_seed_keypair", - "crypto_kx_client_session_keys", - "crypto_kx_server_session_keys", - "crypto_kx_PUBLIC_KEY_BYTES", - "crypto_kx_SECRET_KEY_BYTES", - "crypto_kx_SEED_BYTES", - "crypto_kx_SESSION_KEY_BYTES", - "has_crypto_scalarmult_ed25519", - "crypto_scalarmult_BYTES", - "crypto_scalarmult_SCALARBYTES", - "crypto_scalarmult", - "crypto_scalarmult_base", - "crypto_scalarmult_ed25519_BYTES", - "crypto_scalarmult_ed25519_SCALARBYTES", - "crypto_scalarmult_ed25519", - "crypto_scalarmult_ed25519_base", - "crypto_scalarmult_ed25519_noclamp", - "crypto_scalarmult_ed25519_base_noclamp", - "crypto_secretbox_KEYBYTES", - "crypto_secretbox_NONCEBYTES", - "crypto_secretbox_ZEROBYTES", - "crypto_secretbox_BOXZEROBYTES", - "crypto_secretbox_MACBYTES", - "crypto_secretbox_MESSAGEBYTES_MAX", - "crypto_secretbox", - "crypto_secretbox_open", - "crypto_secretstream_xchacha20poly1305_ABYTES", - "crypto_secretstream_xchacha20poly1305_HEADERBYTES", - "crypto_secretstream_xchacha20poly1305_KEYBYTES", - "crypto_secretstream_xchacha20poly1305_STATEBYTES", - "crypto_secretstream_xchacha20poly1305_TAG_FINAL", - "crypto_secretstream_xchacha20poly1305_TAG_MESSAGE", - "crypto_secretstream_xchacha20poly1305_TAG_PUSH", - "crypto_secretstream_xchacha20poly1305_TAG_REKEY", - "crypto_secretstream_xchacha20poly1305_init_pull", - "crypto_secretstream_xchacha20poly1305_init_push", - "crypto_secretstream_xchacha20poly1305_keygen", - "crypto_secretstream_xchacha20poly1305_pull", - "crypto_secretstream_xchacha20poly1305_push", - "crypto_secretstream_xchacha20poly1305_rekey", - "crypto_secretstream_xchacha20poly1305_state", - "has_crypto_shorthash_siphashx24", - "crypto_shorthash_siphash24_BYTES", - "crypto_shorthash_siphash24_KEYBYTES", - "crypto_shorthash_siphash24", - "crypto_shorthash_siphashx24_BYTES", - "crypto_shorthash_siphashx24_KEYBYTES", - "crypto_shorthash_siphashx24", - "crypto_sign_BYTES", - "crypto_sign_SEEDBYTES", - "crypto_sign_PUBLICKEYBYTES", - "crypto_sign_SECRETKEYBYTES", - "crypto_sign_keypair", - "crypto_sign_seed_keypair", - "crypto_sign", - "crypto_sign_open", - "crypto_sign_ed25519_pk_to_curve25519", - "crypto_sign_ed25519_sk_to_curve25519", - "crypto_sign_ed25519_sk_to_pk", - "crypto_sign_ed25519_sk_to_seed", - "crypto_sign_ed25519ph_STATEBYTES", - "crypto_sign_ed25519ph_final_create", - "crypto_sign_ed25519ph_final_verify", - "crypto_sign_ed25519ph_state", - "crypto_sign_ed25519ph_update", - "crypto_pwhash_ALG_ARGON2I13", - "crypto_pwhash_ALG_ARGON2ID13", - "crypto_pwhash_ALG_DEFAULT", - "crypto_pwhash_BYTES_MAX", - "crypto_pwhash_BYTES_MIN", - "crypto_pwhash_PASSWD_MAX", - "crypto_pwhash_PASSWD_MIN", - "crypto_pwhash_SALTBYTES", - "crypto_pwhash_STRBYTES", - "crypto_pwhash_alg", - "crypto_pwhash_argon2i_MEMLIMIT_MIN", - "crypto_pwhash_argon2i_MEMLIMIT_MAX", - "crypto_pwhash_argon2i_MEMLIMIT_INTERACTIVE", - "crypto_pwhash_argon2i_MEMLIMIT_MODERATE", - "crypto_pwhash_argon2i_MEMLIMIT_SENSITIVE", - "crypto_pwhash_argon2i_OPSLIMIT_MIN", - "crypto_pwhash_argon2i_OPSLIMIT_MAX", - "crypto_pwhash_argon2i_OPSLIMIT_INTERACTIVE", - "crypto_pwhash_argon2i_OPSLIMIT_MODERATE", - "crypto_pwhash_argon2i_OPSLIMIT_SENSITIVE", - "crypto_pwhash_argon2i_STRPREFIX", - "crypto_pwhash_argon2id_MEMLIMIT_MIN", - "crypto_pwhash_argon2id_MEMLIMIT_MAX", - "crypto_pwhash_argon2id_MEMLIMIT_INTERACTIVE", - "crypto_pwhash_argon2id_MEMLIMIT_MODERATE", - "crypto_pwhash_argon2id_OPSLIMIT_MIN", - "crypto_pwhash_argon2id_OPSLIMIT_MAX", - "crypto_pwhash_argon2id_MEMLIMIT_SENSITIVE", - "crypto_pwhash_argon2id_OPSLIMIT_INTERACTIVE", - "crypto_pwhash_argon2id_OPSLIMIT_MODERATE", - "crypto_pwhash_argon2id_OPSLIMIT_SENSITIVE", - "crypto_pwhash_argon2id_STRPREFIX", - "crypto_pwhash_str_alg", - "crypto_pwhash_str_verify", - "has_crypto_pwhash_scryptsalsa208sha256", - "crypto_pwhash_scryptsalsa208sha256_BYTES_MAX", - "crypto_pwhash_scryptsalsa208sha256_BYTES_MIN", - "crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_INTERACTIVE", - "crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_MAX", - "crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_MIN", - "crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_SENSITIVE", - "crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_INTERACTIVE", - "crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_MAX", - "crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_MIN", - "crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_SENSITIVE", - "crypto_pwhash_scryptsalsa208sha256_PASSWD_MAX", - "crypto_pwhash_scryptsalsa208sha256_PASSWD_MIN", - "crypto_pwhash_scryptsalsa208sha256_SALTBYTES", - "crypto_pwhash_scryptsalsa208sha256_STRBYTES", - "crypto_pwhash_scryptsalsa208sha256_STRPREFIX", - "crypto_pwhash_scryptsalsa208sha256_ll", - "crypto_pwhash_scryptsalsa208sha256_str", - "crypto_pwhash_scryptsalsa208sha256_str_verify", - "nacl_bindings_pick_scrypt_params", - "randombytes", - "randombytes_buf_deterministic", - "sodium_init", - "sodium_add", - "sodium_increment", - "sodium_memcmp", - "sodium_pad", - "sodium_unpad", -] - - -# Initialize Sodium -sodium_init() diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 4c83bb1062d61c142a1db43591417ce5e7a58a5e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_aead.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_aead.cpython-310.pyc deleted file mode 100644 index eca11a71edb42d93f14ff0555cbcf8940277acc2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_aead.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_box.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_box.cpython-310.pyc deleted file mode 100644 index 917a574c22a853eedaa78d38ca333381b06eaef3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_box.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_core.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_core.cpython-310.pyc deleted file mode 100644 index 3ea1291d3d8e363b64b715592ad421013d117fad..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_core.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_generichash.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_generichash.cpython-310.pyc deleted file mode 100644 index e7c58e5dc821837e9cb08205b213d0e7ad057bf3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_generichash.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_hash.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_hash.cpython-310.pyc deleted file mode 100644 index c047f00b1a1f65b37853cc18f11496ce9295b77b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_hash.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_kx.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_kx.cpython-310.pyc deleted file mode 100644 index ced21341d93e8fbdeb168211d49434676573eadf..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_kx.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_pwhash.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_pwhash.cpython-310.pyc deleted file mode 100644 index e721ed25d29aed49bfc6d925d4f9dd60c104a147..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_pwhash.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_scalarmult.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_scalarmult.cpython-310.pyc deleted file mode 100644 index 909945ded6bc4c2029b113b8f8eb4ed37bd42a74..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_scalarmult.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_secretbox.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_secretbox.cpython-310.pyc deleted file mode 100644 index a040882a70778254c12e09c928fe0639335af7fd..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_secretbox.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_secretstream.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_secretstream.cpython-310.pyc deleted file mode 100644 index 990da1ffcf67692bad838c7d503cfb5141d81e94..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_secretstream.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_shorthash.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_shorthash.cpython-310.pyc deleted file mode 100644 index 4c83051112fd7ec42488583e8f885a546a2b0001..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_shorthash.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_sign.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_sign.cpython-310.pyc deleted file mode 100644 index 7312bb0a4beb938e1acb3153b701281061e919fa..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/crypto_sign.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/randombytes.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/randombytes.cpython-310.pyc deleted file mode 100644 index ad3c25e73c8b372e1b39502063ced60fa067cb14..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/randombytes.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/sodium_core.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/sodium_core.cpython-310.pyc deleted file mode 100644 index 8dd49e11c64ece4cfa88914a86c917e23b4319b2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/sodium_core.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/utils.cpython-310.pyc deleted file mode 100644 index ffff64de5b4f4dfa341a6f5e44bc1497664f2b06..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/bindings/__pycache__/utils.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_aead.py b/venv/lib/python3.10/site-packages/nacl/bindings/crypto_aead.py deleted file mode 100644 index 2e7168db74bdae5fa72f8080291c1362cd145021..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_aead.py +++ /dev/null @@ -1,559 +0,0 @@ -# Copyright 2017 Donald Stufft and individual contributors -# -# 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 - -from nacl import exceptions as exc -from nacl._sodium import ffi, lib -from nacl.exceptions import ensure - -""" -Implementations of authenticated encription with associated data (*AEAD*) -constructions building on the chacha20 stream cipher and the poly1305 -authenticator -""" - -crypto_aead_chacha20poly1305_ietf_KEYBYTES: int = ( - lib.crypto_aead_chacha20poly1305_ietf_keybytes() -) -crypto_aead_chacha20poly1305_ietf_NSECBYTES: int = ( - lib.crypto_aead_chacha20poly1305_ietf_nsecbytes() -) -crypto_aead_chacha20poly1305_ietf_NPUBBYTES: int = ( - lib.crypto_aead_chacha20poly1305_ietf_npubbytes() -) -crypto_aead_chacha20poly1305_ietf_ABYTES: int = ( - lib.crypto_aead_chacha20poly1305_ietf_abytes() -) -crypto_aead_chacha20poly1305_ietf_MESSAGEBYTES_MAX: int = ( - lib.crypto_aead_chacha20poly1305_ietf_messagebytes_max() -) -_aead_chacha20poly1305_ietf_CRYPTBYTES_MAX = ( - crypto_aead_chacha20poly1305_ietf_MESSAGEBYTES_MAX - + crypto_aead_chacha20poly1305_ietf_ABYTES -) - -crypto_aead_chacha20poly1305_KEYBYTES: int = ( - lib.crypto_aead_chacha20poly1305_keybytes() -) -crypto_aead_chacha20poly1305_NSECBYTES: int = ( - lib.crypto_aead_chacha20poly1305_nsecbytes() -) -crypto_aead_chacha20poly1305_NPUBBYTES: int = ( - lib.crypto_aead_chacha20poly1305_npubbytes() -) -crypto_aead_chacha20poly1305_ABYTES: int = ( - lib.crypto_aead_chacha20poly1305_abytes() -) -crypto_aead_chacha20poly1305_MESSAGEBYTES_MAX: int = ( - lib.crypto_aead_chacha20poly1305_messagebytes_max() -) -_aead_chacha20poly1305_CRYPTBYTES_MAX = ( - crypto_aead_chacha20poly1305_MESSAGEBYTES_MAX - + crypto_aead_chacha20poly1305_ABYTES -) - -crypto_aead_xchacha20poly1305_ietf_KEYBYTES: int = ( - lib.crypto_aead_xchacha20poly1305_ietf_keybytes() -) -crypto_aead_xchacha20poly1305_ietf_NSECBYTES: int = ( - lib.crypto_aead_xchacha20poly1305_ietf_nsecbytes() -) -crypto_aead_xchacha20poly1305_ietf_NPUBBYTES: int = ( - lib.crypto_aead_xchacha20poly1305_ietf_npubbytes() -) -crypto_aead_xchacha20poly1305_ietf_ABYTES: int = ( - lib.crypto_aead_xchacha20poly1305_ietf_abytes() -) -crypto_aead_xchacha20poly1305_ietf_MESSAGEBYTES_MAX: int = ( - lib.crypto_aead_xchacha20poly1305_ietf_messagebytes_max() -) -_aead_xchacha20poly1305_ietf_CRYPTBYTES_MAX = ( - crypto_aead_xchacha20poly1305_ietf_MESSAGEBYTES_MAX - + crypto_aead_xchacha20poly1305_ietf_ABYTES -) - - -def crypto_aead_chacha20poly1305_ietf_encrypt( - message: bytes, aad: Optional[bytes], nonce: bytes, key: bytes -) -> bytes: - """ - Encrypt the given ``message`` using the IETF ratified chacha20poly1305 - construction described in RFC7539. - - :param message: - :type message: bytes - :param aad: - :type aad: Optional[bytes] - :param nonce: - :type nonce: bytes - :param key: - :type key: bytes - :return: authenticated ciphertext - :rtype: bytes - """ - ensure( - isinstance(message, bytes), - "Input message type must be bytes", - raising=exc.TypeError, - ) - - mlen = len(message) - - ensure( - mlen <= crypto_aead_chacha20poly1305_ietf_MESSAGEBYTES_MAX, - "Message must be at most {} bytes long".format( - crypto_aead_chacha20poly1305_ietf_MESSAGEBYTES_MAX - ), - raising=exc.ValueError, - ) - - ensure( - isinstance(aad, bytes) or (aad is None), - "Additional data must be bytes or None", - raising=exc.TypeError, - ) - - ensure( - isinstance(nonce, bytes) - and len(nonce) == crypto_aead_chacha20poly1305_ietf_NPUBBYTES, - "Nonce must be a {} bytes long bytes sequence".format( - crypto_aead_chacha20poly1305_ietf_NPUBBYTES - ), - raising=exc.TypeError, - ) - - ensure( - isinstance(key, bytes) - and len(key) == crypto_aead_chacha20poly1305_ietf_KEYBYTES, - "Key must be a {} bytes long bytes sequence".format( - crypto_aead_chacha20poly1305_ietf_KEYBYTES - ), - raising=exc.TypeError, - ) - - if aad: - _aad = aad - aalen = len(aad) - else: - _aad = ffi.NULL - aalen = 0 - - mxout = mlen + crypto_aead_chacha20poly1305_ietf_ABYTES - - clen = ffi.new("unsigned long long *") - - ciphertext = ffi.new("unsigned char[]", mxout) - - res = lib.crypto_aead_chacha20poly1305_ietf_encrypt( - ciphertext, clen, message, mlen, _aad, aalen, ffi.NULL, nonce, key - ) - - ensure(res == 0, "Encryption failed.", raising=exc.CryptoError) - return ffi.buffer(ciphertext, clen[0])[:] - - -def crypto_aead_chacha20poly1305_ietf_decrypt( - ciphertext: bytes, aad: Optional[bytes], nonce: bytes, key: bytes -) -> bytes: - """ - Decrypt the given ``ciphertext`` using the IETF ratified chacha20poly1305 - construction described in RFC7539. - - :param ciphertext: - :type ciphertext: bytes - :param aad: - :type aad: Optional[bytes] - :param nonce: - :type nonce: bytes - :param key: - :type key: bytes - :return: message - :rtype: bytes - """ - ensure( - isinstance(ciphertext, bytes), - "Input ciphertext type must be bytes", - raising=exc.TypeError, - ) - - clen = len(ciphertext) - - ensure( - clen <= _aead_chacha20poly1305_ietf_CRYPTBYTES_MAX, - "Ciphertext must be at most {} bytes long".format( - _aead_chacha20poly1305_ietf_CRYPTBYTES_MAX - ), - raising=exc.ValueError, - ) - - ensure( - isinstance(aad, bytes) or (aad is None), - "Additional data must be bytes or None", - raising=exc.TypeError, - ) - - ensure( - isinstance(nonce, bytes) - and len(nonce) == crypto_aead_chacha20poly1305_ietf_NPUBBYTES, - "Nonce must be a {} bytes long bytes sequence".format( - crypto_aead_chacha20poly1305_ietf_NPUBBYTES - ), - raising=exc.TypeError, - ) - - ensure( - isinstance(key, bytes) - and len(key) == crypto_aead_chacha20poly1305_ietf_KEYBYTES, - "Key must be a {} bytes long bytes sequence".format( - crypto_aead_chacha20poly1305_ietf_KEYBYTES - ), - raising=exc.TypeError, - ) - - mxout = clen - crypto_aead_chacha20poly1305_ietf_ABYTES - - mlen = ffi.new("unsigned long long *") - message = ffi.new("unsigned char[]", mxout) - - if aad: - _aad = aad - aalen = len(aad) - else: - _aad = ffi.NULL - aalen = 0 - - res = lib.crypto_aead_chacha20poly1305_ietf_decrypt( - message, mlen, ffi.NULL, ciphertext, clen, _aad, aalen, nonce, key - ) - - ensure(res == 0, "Decryption failed.", raising=exc.CryptoError) - - return ffi.buffer(message, mlen[0])[:] - - -def crypto_aead_chacha20poly1305_encrypt( - message: bytes, aad: Optional[bytes], nonce: bytes, key: bytes -) -> bytes: - """ - Encrypt the given ``message`` using the "legacy" construction - described in draft-agl-tls-chacha20poly1305. - - :param message: - :type message: bytes - :param aad: - :type aad: Optional[bytes] - :param nonce: - :type nonce: bytes - :param key: - :type key: bytes - :return: authenticated ciphertext - :rtype: bytes - """ - ensure( - isinstance(message, bytes), - "Input message type must be bytes", - raising=exc.TypeError, - ) - - mlen = len(message) - - ensure( - mlen <= crypto_aead_chacha20poly1305_MESSAGEBYTES_MAX, - "Message must be at most {} bytes long".format( - crypto_aead_chacha20poly1305_MESSAGEBYTES_MAX - ), - raising=exc.ValueError, - ) - - ensure( - isinstance(aad, bytes) or (aad is None), - "Additional data must be bytes or None", - raising=exc.TypeError, - ) - - ensure( - isinstance(nonce, bytes) - and len(nonce) == crypto_aead_chacha20poly1305_NPUBBYTES, - "Nonce must be a {} bytes long bytes sequence".format( - crypto_aead_chacha20poly1305_NPUBBYTES - ), - raising=exc.TypeError, - ) - - ensure( - isinstance(key, bytes) - and len(key) == crypto_aead_chacha20poly1305_KEYBYTES, - "Key must be a {} bytes long bytes sequence".format( - crypto_aead_chacha20poly1305_KEYBYTES - ), - raising=exc.TypeError, - ) - - if aad: - _aad = aad - aalen = len(aad) - else: - _aad = ffi.NULL - aalen = 0 - - mlen = len(message) - mxout = mlen + crypto_aead_chacha20poly1305_ietf_ABYTES - - clen = ffi.new("unsigned long long *") - - ciphertext = ffi.new("unsigned char[]", mxout) - - res = lib.crypto_aead_chacha20poly1305_encrypt( - ciphertext, clen, message, mlen, _aad, aalen, ffi.NULL, nonce, key - ) - - ensure(res == 0, "Encryption failed.", raising=exc.CryptoError) - return ffi.buffer(ciphertext, clen[0])[:] - - -def crypto_aead_chacha20poly1305_decrypt( - ciphertext: bytes, aad: Optional[bytes], nonce: bytes, key: bytes -) -> bytes: - """ - Decrypt the given ``ciphertext`` using the "legacy" construction - described in draft-agl-tls-chacha20poly1305. - - :param ciphertext: authenticated ciphertext - :type ciphertext: bytes - :param aad: - :type aad: Optional[bytes] - :param nonce: - :type nonce: bytes - :param key: - :type key: bytes - :return: message - :rtype: bytes - """ - ensure( - isinstance(ciphertext, bytes), - "Input ciphertext type must be bytes", - raising=exc.TypeError, - ) - - clen = len(ciphertext) - - ensure( - clen <= _aead_chacha20poly1305_CRYPTBYTES_MAX, - "Ciphertext must be at most {} bytes long".format( - _aead_chacha20poly1305_CRYPTBYTES_MAX - ), - raising=exc.ValueError, - ) - - ensure( - isinstance(aad, bytes) or (aad is None), - "Additional data must be bytes or None", - raising=exc.TypeError, - ) - - ensure( - isinstance(nonce, bytes) - and len(nonce) == crypto_aead_chacha20poly1305_NPUBBYTES, - "Nonce must be a {} bytes long bytes sequence".format( - crypto_aead_chacha20poly1305_NPUBBYTES - ), - raising=exc.TypeError, - ) - - ensure( - isinstance(key, bytes) - and len(key) == crypto_aead_chacha20poly1305_KEYBYTES, - "Key must be a {} bytes long bytes sequence".format( - crypto_aead_chacha20poly1305_KEYBYTES - ), - raising=exc.TypeError, - ) - - mxout = clen - crypto_aead_chacha20poly1305_ABYTES - - mlen = ffi.new("unsigned long long *") - message = ffi.new("unsigned char[]", mxout) - - if aad: - _aad = aad - aalen = len(aad) - else: - _aad = ffi.NULL - aalen = 0 - - res = lib.crypto_aead_chacha20poly1305_decrypt( - message, mlen, ffi.NULL, ciphertext, clen, _aad, aalen, nonce, key - ) - - ensure(res == 0, "Decryption failed.", raising=exc.CryptoError) - - return ffi.buffer(message, mlen[0])[:] - - -def crypto_aead_xchacha20poly1305_ietf_encrypt( - message: bytes, aad: Optional[bytes], nonce: bytes, key: bytes -) -> bytes: - """ - Encrypt the given ``message`` using the long-nonces xchacha20poly1305 - construction. - - :param message: - :type message: bytes - :param aad: - :type aad: Optional[bytes] - :param nonce: - :type nonce: bytes - :param key: - :type key: bytes - :return: authenticated ciphertext - :rtype: bytes - """ - ensure( - isinstance(message, bytes), - "Input message type must be bytes", - raising=exc.TypeError, - ) - - mlen = len(message) - - ensure( - mlen <= crypto_aead_xchacha20poly1305_ietf_MESSAGEBYTES_MAX, - "Message must be at most {} bytes long".format( - crypto_aead_xchacha20poly1305_ietf_MESSAGEBYTES_MAX - ), - raising=exc.ValueError, - ) - - ensure( - isinstance(aad, bytes) or (aad is None), - "Additional data must be bytes or None", - raising=exc.TypeError, - ) - - ensure( - isinstance(nonce, bytes) - and len(nonce) == crypto_aead_xchacha20poly1305_ietf_NPUBBYTES, - "Nonce must be a {} bytes long bytes sequence".format( - crypto_aead_xchacha20poly1305_ietf_NPUBBYTES - ), - raising=exc.TypeError, - ) - - ensure( - isinstance(key, bytes) - and len(key) == crypto_aead_xchacha20poly1305_ietf_KEYBYTES, - "Key must be a {} bytes long bytes sequence".format( - crypto_aead_xchacha20poly1305_ietf_KEYBYTES - ), - raising=exc.TypeError, - ) - - if aad: - _aad = aad - aalen = len(aad) - else: - _aad = ffi.NULL - aalen = 0 - - mlen = len(message) - mxout = mlen + crypto_aead_xchacha20poly1305_ietf_ABYTES - - clen = ffi.new("unsigned long long *") - - ciphertext = ffi.new("unsigned char[]", mxout) - - res = lib.crypto_aead_xchacha20poly1305_ietf_encrypt( - ciphertext, clen, message, mlen, _aad, aalen, ffi.NULL, nonce, key - ) - - ensure(res == 0, "Encryption failed.", raising=exc.CryptoError) - return ffi.buffer(ciphertext, clen[0])[:] - - -def crypto_aead_xchacha20poly1305_ietf_decrypt( - ciphertext: bytes, aad: Optional[bytes], nonce: bytes, key: bytes -) -> bytes: - """ - Decrypt the given ``ciphertext`` using the long-nonces xchacha20poly1305 - construction. - - :param ciphertext: authenticated ciphertext - :type ciphertext: bytes - :param aad: - :type aad: Optional[bytes] - :param nonce: - :type nonce: bytes - :param key: - :type key: bytes - :return: message - :rtype: bytes - """ - ensure( - isinstance(ciphertext, bytes), - "Input ciphertext type must be bytes", - raising=exc.TypeError, - ) - - clen = len(ciphertext) - - ensure( - clen <= _aead_xchacha20poly1305_ietf_CRYPTBYTES_MAX, - "Ciphertext must be at most {} bytes long".format( - _aead_xchacha20poly1305_ietf_CRYPTBYTES_MAX - ), - raising=exc.ValueError, - ) - - ensure( - isinstance(aad, bytes) or (aad is None), - "Additional data must be bytes or None", - raising=exc.TypeError, - ) - - ensure( - isinstance(nonce, bytes) - and len(nonce) == crypto_aead_xchacha20poly1305_ietf_NPUBBYTES, - "Nonce must be a {} bytes long bytes sequence".format( - crypto_aead_xchacha20poly1305_ietf_NPUBBYTES - ), - raising=exc.TypeError, - ) - - ensure( - isinstance(key, bytes) - and len(key) == crypto_aead_xchacha20poly1305_ietf_KEYBYTES, - "Key must be a {} bytes long bytes sequence".format( - crypto_aead_xchacha20poly1305_ietf_KEYBYTES - ), - raising=exc.TypeError, - ) - - mxout = clen - crypto_aead_xchacha20poly1305_ietf_ABYTES - mlen = ffi.new("unsigned long long *") - message = ffi.new("unsigned char[]", mxout) - - if aad: - _aad = aad - aalen = len(aad) - else: - _aad = ffi.NULL - aalen = 0 - - res = lib.crypto_aead_xchacha20poly1305_ietf_decrypt( - message, mlen, ffi.NULL, ciphertext, clen, _aad, aalen, nonce, key - ) - - ensure(res == 0, "Decryption failed.", raising=exc.CryptoError) - - return ffi.buffer(message, mlen[0])[:] diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_box.py b/venv/lib/python3.10/site-packages/nacl/bindings/crypto_box.py deleted file mode 100644 index 74f7f0a639cd7c6e03b1a52bdb3f15101b1ce050..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_box.py +++ /dev/null @@ -1,324 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 Tuple - -from nacl import exceptions as exc -from nacl._sodium import ffi, lib -from nacl.exceptions import ensure - - -__all__ = ["crypto_box_keypair", "crypto_box"] - - -crypto_box_SECRETKEYBYTES: int = lib.crypto_box_secretkeybytes() -crypto_box_PUBLICKEYBYTES: int = lib.crypto_box_publickeybytes() -crypto_box_SEEDBYTES: int = lib.crypto_box_seedbytes() -crypto_box_NONCEBYTES: int = lib.crypto_box_noncebytes() -crypto_box_ZEROBYTES: int = lib.crypto_box_zerobytes() -crypto_box_BOXZEROBYTES: int = lib.crypto_box_boxzerobytes() -crypto_box_BEFORENMBYTES: int = lib.crypto_box_beforenmbytes() -crypto_box_SEALBYTES: int = lib.crypto_box_sealbytes() - - -def crypto_box_keypair() -> Tuple[bytes, bytes]: - """ - Returns a randomly generated public and secret key. - - :rtype: (bytes(public_key), bytes(secret_key)) - """ - pk = ffi.new("unsigned char[]", crypto_box_PUBLICKEYBYTES) - sk = ffi.new("unsigned char[]", crypto_box_SECRETKEYBYTES) - - rc = lib.crypto_box_keypair(pk, sk) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ( - ffi.buffer(pk, crypto_box_PUBLICKEYBYTES)[:], - ffi.buffer(sk, crypto_box_SECRETKEYBYTES)[:], - ) - - -def crypto_box_seed_keypair(seed: bytes) -> Tuple[bytes, bytes]: - """ - Returns a (public, secret) keypair deterministically generated - from an input ``seed``. - - .. warning:: The seed **must** be high-entropy; therefore, - its generator **must** be a cryptographic quality - random function like, for example, :func:`~nacl.utils.random`. - - .. warning:: The seed **must** be protected and remain secret. - Anyone who knows the seed is really in possession of - the corresponding PrivateKey. - - - :param seed: bytes - :rtype: (bytes(public_key), bytes(secret_key)) - """ - ensure(isinstance(seed, bytes), "seed must be bytes", raising=TypeError) - - if len(seed) != crypto_box_SEEDBYTES: - raise exc.ValueError("Invalid seed") - - pk = ffi.new("unsigned char[]", crypto_box_PUBLICKEYBYTES) - sk = ffi.new("unsigned char[]", crypto_box_SECRETKEYBYTES) - - rc = lib.crypto_box_seed_keypair(pk, sk, seed) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ( - ffi.buffer(pk, crypto_box_PUBLICKEYBYTES)[:], - ffi.buffer(sk, crypto_box_SECRETKEYBYTES)[:], - ) - - -def crypto_box(message: bytes, nonce: bytes, pk: bytes, sk: bytes) -> bytes: - """ - Encrypts and returns a message ``message`` using the secret key ``sk``, - public key ``pk``, and the nonce ``nonce``. - - :param message: bytes - :param nonce: bytes - :param pk: bytes - :param sk: bytes - :rtype: bytes - """ - if len(nonce) != crypto_box_NONCEBYTES: - raise exc.ValueError("Invalid nonce size") - - if len(pk) != crypto_box_PUBLICKEYBYTES: - raise exc.ValueError("Invalid public key") - - if len(sk) != crypto_box_SECRETKEYBYTES: - raise exc.ValueError("Invalid secret key") - - padded = (b"\x00" * crypto_box_ZEROBYTES) + message - ciphertext = ffi.new("unsigned char[]", len(padded)) - - rc = lib.crypto_box(ciphertext, padded, len(padded), nonce, pk, sk) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(ciphertext, len(padded))[crypto_box_BOXZEROBYTES:] - - -def crypto_box_open( - ciphertext: bytes, nonce: bytes, pk: bytes, sk: bytes -) -> bytes: - """ - Decrypts and returns an encrypted message ``ciphertext``, using the secret - key ``sk``, public key ``pk``, and the nonce ``nonce``. - - :param ciphertext: bytes - :param nonce: bytes - :param pk: bytes - :param sk: bytes - :rtype: bytes - """ - if len(nonce) != crypto_box_NONCEBYTES: - raise exc.ValueError("Invalid nonce size") - - if len(pk) != crypto_box_PUBLICKEYBYTES: - raise exc.ValueError("Invalid public key") - - if len(sk) != crypto_box_SECRETKEYBYTES: - raise exc.ValueError("Invalid secret key") - - padded = (b"\x00" * crypto_box_BOXZEROBYTES) + ciphertext - plaintext = ffi.new("unsigned char[]", len(padded)) - - res = lib.crypto_box_open(plaintext, padded, len(padded), nonce, pk, sk) - ensure( - res == 0, - "An error occurred trying to decrypt the message", - raising=exc.CryptoError, - ) - - return ffi.buffer(plaintext, len(padded))[crypto_box_ZEROBYTES:] - - -def crypto_box_beforenm(pk: bytes, sk: bytes) -> bytes: - """ - Computes and returns the shared key for the public key ``pk`` and the - secret key ``sk``. This can be used to speed up operations where the same - set of keys is going to be used multiple times. - - :param pk: bytes - :param sk: bytes - :rtype: bytes - """ - if len(pk) != crypto_box_PUBLICKEYBYTES: - raise exc.ValueError("Invalid public key") - - if len(sk) != crypto_box_SECRETKEYBYTES: - raise exc.ValueError("Invalid secret key") - - k = ffi.new("unsigned char[]", crypto_box_BEFORENMBYTES) - - rc = lib.crypto_box_beforenm(k, pk, sk) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(k, crypto_box_BEFORENMBYTES)[:] - - -def crypto_box_afternm(message: bytes, nonce: bytes, k: bytes) -> bytes: - """ - Encrypts and returns the message ``message`` using the shared key ``k`` and - the nonce ``nonce``. - - :param message: bytes - :param nonce: bytes - :param k: bytes - :rtype: bytes - """ - if len(nonce) != crypto_box_NONCEBYTES: - raise exc.ValueError("Invalid nonce") - - if len(k) != crypto_box_BEFORENMBYTES: - raise exc.ValueError("Invalid shared key") - - padded = b"\x00" * crypto_box_ZEROBYTES + message - ciphertext = ffi.new("unsigned char[]", len(padded)) - - rc = lib.crypto_box_afternm(ciphertext, padded, len(padded), nonce, k) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(ciphertext, len(padded))[crypto_box_BOXZEROBYTES:] - - -def crypto_box_open_afternm( - ciphertext: bytes, nonce: bytes, k: bytes -) -> bytes: - """ - Decrypts and returns the encrypted message ``ciphertext``, using the shared - key ``k`` and the nonce ``nonce``. - - :param ciphertext: bytes - :param nonce: bytes - :param k: bytes - :rtype: bytes - """ - if len(nonce) != crypto_box_NONCEBYTES: - raise exc.ValueError("Invalid nonce") - - if len(k) != crypto_box_BEFORENMBYTES: - raise exc.ValueError("Invalid shared key") - - padded = (b"\x00" * crypto_box_BOXZEROBYTES) + ciphertext - plaintext = ffi.new("unsigned char[]", len(padded)) - - res = lib.crypto_box_open_afternm(plaintext, padded, len(padded), nonce, k) - ensure( - res == 0, - "An error occurred trying to decrypt the message", - raising=exc.CryptoError, - ) - - return ffi.buffer(plaintext, len(padded))[crypto_box_ZEROBYTES:] - - -def crypto_box_seal(message: bytes, pk: bytes) -> bytes: - """ - Encrypts and returns a message ``message`` using an ephemeral secret key - and the public key ``pk``. - The ephemeral public key, which is embedded in the sealed box, is also - used, in combination with ``pk``, to derive the nonce needed for the - underlying box construct. - - :param message: bytes - :param pk: bytes - :rtype: bytes - - .. versionadded:: 1.2 - """ - ensure( - isinstance(message, bytes), - "input message must be bytes", - raising=TypeError, - ) - - ensure( - isinstance(pk, bytes), "public key must be bytes", raising=TypeError - ) - - if len(pk) != crypto_box_PUBLICKEYBYTES: - raise exc.ValueError("Invalid public key") - - _mlen = len(message) - _clen = crypto_box_SEALBYTES + _mlen - - ciphertext = ffi.new("unsigned char[]", _clen) - - rc = lib.crypto_box_seal(ciphertext, message, _mlen, pk) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(ciphertext, _clen)[:] - - -def crypto_box_seal_open(ciphertext: bytes, pk: bytes, sk: bytes) -> bytes: - """ - Decrypts and returns an encrypted message ``ciphertext``, using the - recipent's secret key ``sk`` and the sender's ephemeral public key - embedded in the sealed box. The box contruct nonce is derived from - the recipient's public key ``pk`` and the sender's public key. - - :param ciphertext: bytes - :param pk: bytes - :param sk: bytes - :rtype: bytes - - .. versionadded:: 1.2 - """ - ensure( - isinstance(ciphertext, bytes), - "input ciphertext must be bytes", - raising=TypeError, - ) - - ensure( - isinstance(pk, bytes), "public key must be bytes", raising=TypeError - ) - - ensure( - isinstance(sk, bytes), "secret key must be bytes", raising=TypeError - ) - - if len(pk) != crypto_box_PUBLICKEYBYTES: - raise exc.ValueError("Invalid public key") - - if len(sk) != crypto_box_SECRETKEYBYTES: - raise exc.ValueError("Invalid secret key") - - _clen = len(ciphertext) - - ensure( - _clen >= crypto_box_SEALBYTES, - ("Input cyphertext must be at least {} long").format( - crypto_box_SEALBYTES - ), - raising=exc.TypeError, - ) - - _mlen = _clen - crypto_box_SEALBYTES - - # zero-length malloc results are implementation.dependent - plaintext = ffi.new("unsigned char[]", max(1, _mlen)) - - res = lib.crypto_box_seal_open(plaintext, ciphertext, _clen, pk, sk) - ensure( - res == 0, - "An error occurred trying to decrypt the message", - raising=exc.CryptoError, - ) - - return ffi.buffer(plaintext, _mlen)[:] diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_core.py b/venv/lib/python3.10/site-packages/nacl/bindings/crypto_core.py deleted file mode 100644 index d29da9b3fc25ad78b4a33752a957b313a2b042bd..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_core.py +++ /dev/null @@ -1,412 +0,0 @@ -# Copyright 2018 Donald Stufft and individual contributors -# -# 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 nacl import exceptions as exc -from nacl._sodium import ffi, lib -from nacl.exceptions import ensure - - -has_crypto_core_ed25519 = bool(lib.PYNACL_HAS_CRYPTO_CORE_ED25519) - -crypto_core_ed25519_BYTES = 0 -crypto_core_ed25519_SCALARBYTES = 0 -crypto_core_ed25519_NONREDUCEDSCALARBYTES = 0 - -if has_crypto_core_ed25519: - crypto_core_ed25519_BYTES = lib.crypto_core_ed25519_bytes() - crypto_core_ed25519_SCALARBYTES = lib.crypto_core_ed25519_scalarbytes() - crypto_core_ed25519_NONREDUCEDSCALARBYTES = ( - lib.crypto_core_ed25519_nonreducedscalarbytes() - ) - - -def crypto_core_ed25519_is_valid_point(p: bytes) -> bool: - """ - Check if ``p`` represents a point on the edwards25519 curve, in canonical - form, on the main subgroup, and that the point doesn't have a small order. - - :param p: a :py:data:`.crypto_core_ed25519_BYTES` long bytes sequence - representing a point on the edwards25519 curve - :type p: bytes - :return: point validity - :rtype: bool - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_core_ed25519, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - isinstance(p, bytes) and len(p) == crypto_core_ed25519_BYTES, - "Point must be a crypto_core_ed25519_BYTES long bytes sequence", - raising=exc.TypeError, - ) - - rc = lib.crypto_core_ed25519_is_valid_point(p) - return rc == 1 - - -def crypto_core_ed25519_add(p: bytes, q: bytes) -> bytes: - """ - Add two points on the edwards25519 curve. - - :param p: a :py:data:`.crypto_core_ed25519_BYTES` long bytes sequence - representing a point on the edwards25519 curve - :type p: bytes - :param q: a :py:data:`.crypto_core_ed25519_BYTES` long bytes sequence - representing a point on the edwards25519 curve - :type q: bytes - :return: a point on the edwards25519 curve represented as - a :py:data:`.crypto_core_ed25519_BYTES` long bytes sequence - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_core_ed25519, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - isinstance(p, bytes) - and isinstance(q, bytes) - and len(p) == crypto_core_ed25519_BYTES - and len(q) == crypto_core_ed25519_BYTES, - "Each point must be a {} long bytes sequence".format( - "crypto_core_ed25519_BYTES" - ), - raising=exc.TypeError, - ) - - r = ffi.new("unsigned char[]", crypto_core_ed25519_BYTES) - - rc = lib.crypto_core_ed25519_add(r, p, q) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(r, crypto_core_ed25519_BYTES)[:] - - -def crypto_core_ed25519_sub(p: bytes, q: bytes) -> bytes: - """ - Subtract a point from another on the edwards25519 curve. - - :param p: a :py:data:`.crypto_core_ed25519_BYTES` long bytes sequence - representing a point on the edwards25519 curve - :type p: bytes - :param q: a :py:data:`.crypto_core_ed25519_BYTES` long bytes sequence - representing a point on the edwards25519 curve - :type q: bytes - :return: a point on the edwards25519 curve represented as - a :py:data:`.crypto_core_ed25519_BYTES` long bytes sequence - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_core_ed25519, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - isinstance(p, bytes) - and isinstance(q, bytes) - and len(p) == crypto_core_ed25519_BYTES - and len(q) == crypto_core_ed25519_BYTES, - "Each point must be a {} long bytes sequence".format( - "crypto_core_ed25519_BYTES" - ), - raising=exc.TypeError, - ) - - r = ffi.new("unsigned char[]", crypto_core_ed25519_BYTES) - - rc = lib.crypto_core_ed25519_sub(r, p, q) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(r, crypto_core_ed25519_BYTES)[:] - - -def crypto_core_ed25519_scalar_invert(s: bytes) -> bytes: - """ - Return the multiplicative inverse of integer ``s`` modulo ``L``, - i.e an integer ``i`` such that ``s * i = 1 (mod L)``, where ``L`` - is the order of the main subgroup. - - Raises a ``exc.RuntimeError`` if ``s`` is the integer zero. - - :param s: a :py:data:`.crypto_core_ed25519_SCALARBYTES` - long bytes sequence representing an integer - :type s: bytes - :return: an integer represented as a - :py:data:`.crypto_core_ed25519_SCALARBYTES` long bytes sequence - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_core_ed25519, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - isinstance(s, bytes) and len(s) == crypto_core_ed25519_SCALARBYTES, - "Integer s must be a {} long bytes sequence".format( - "crypto_core_ed25519_SCALARBYTES" - ), - raising=exc.TypeError, - ) - - r = ffi.new("unsigned char[]", crypto_core_ed25519_SCALARBYTES) - - rc = lib.crypto_core_ed25519_scalar_invert(r, s) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(r, crypto_core_ed25519_SCALARBYTES)[:] - - -def crypto_core_ed25519_scalar_negate(s: bytes) -> bytes: - """ - Return the integer ``n`` such that ``s + n = 0 (mod L)``, where ``L`` - is the order of the main subgroup. - - :param s: a :py:data:`.crypto_core_ed25519_SCALARBYTES` - long bytes sequence representing an integer - :type s: bytes - :return: an integer represented as a - :py:data:`.crypto_core_ed25519_SCALARBYTES` long bytes sequence - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_core_ed25519, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - isinstance(s, bytes) and len(s) == crypto_core_ed25519_SCALARBYTES, - "Integer s must be a {} long bytes sequence".format( - "crypto_core_ed25519_SCALARBYTES" - ), - raising=exc.TypeError, - ) - - r = ffi.new("unsigned char[]", crypto_core_ed25519_SCALARBYTES) - - lib.crypto_core_ed25519_scalar_negate(r, s) - - return ffi.buffer(r, crypto_core_ed25519_SCALARBYTES)[:] - - -def crypto_core_ed25519_scalar_complement(s: bytes) -> bytes: - """ - Return the complement of integer ``s`` modulo ``L``, i.e. an integer - ``c`` such that ``s + c = 1 (mod L)``, where ``L`` is the order of - the main subgroup. - - :param s: a :py:data:`.crypto_core_ed25519_SCALARBYTES` - long bytes sequence representing an integer - :type s: bytes - :return: an integer represented as a - :py:data:`.crypto_core_ed25519_SCALARBYTES` long bytes sequence - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_core_ed25519, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - isinstance(s, bytes) and len(s) == crypto_core_ed25519_SCALARBYTES, - "Integer s must be a {} long bytes sequence".format( - "crypto_core_ed25519_SCALARBYTES" - ), - raising=exc.TypeError, - ) - - r = ffi.new("unsigned char[]", crypto_core_ed25519_SCALARBYTES) - - lib.crypto_core_ed25519_scalar_complement(r, s) - - return ffi.buffer(r, crypto_core_ed25519_SCALARBYTES)[:] - - -def crypto_core_ed25519_scalar_add(p: bytes, q: bytes) -> bytes: - """ - Add integers ``p`` and ``q`` modulo ``L``, where ``L`` is the order of - the main subgroup. - - :param p: a :py:data:`.crypto_core_ed25519_SCALARBYTES` - long bytes sequence representing an integer - :type p: bytes - :param q: a :py:data:`.crypto_core_ed25519_SCALARBYTES` - long bytes sequence representing an integer - :type q: bytes - :return: an integer represented as a - :py:data:`.crypto_core_ed25519_SCALARBYTES` long bytes sequence - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_core_ed25519, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - isinstance(p, bytes) - and isinstance(q, bytes) - and len(p) == crypto_core_ed25519_SCALARBYTES - and len(q) == crypto_core_ed25519_SCALARBYTES, - "Each integer must be a {} long bytes sequence".format( - "crypto_core_ed25519_SCALARBYTES" - ), - raising=exc.TypeError, - ) - - r = ffi.new("unsigned char[]", crypto_core_ed25519_SCALARBYTES) - - lib.crypto_core_ed25519_scalar_add(r, p, q) - - return ffi.buffer(r, crypto_core_ed25519_SCALARBYTES)[:] - - -def crypto_core_ed25519_scalar_sub(p: bytes, q: bytes) -> bytes: - """ - Subtract integers ``p`` and ``q`` modulo ``L``, where ``L`` is the - order of the main subgroup. - - :param p: a :py:data:`.crypto_core_ed25519_SCALARBYTES` - long bytes sequence representing an integer - :type p: bytes - :param q: a :py:data:`.crypto_core_ed25519_SCALARBYTES` - long bytes sequence representing an integer - :type q: bytes - :return: an integer represented as a - :py:data:`.crypto_core_ed25519_SCALARBYTES` long bytes sequence - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_core_ed25519, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - isinstance(p, bytes) - and isinstance(q, bytes) - and len(p) == crypto_core_ed25519_SCALARBYTES - and len(q) == crypto_core_ed25519_SCALARBYTES, - "Each integer must be a {} long bytes sequence".format( - "crypto_core_ed25519_SCALARBYTES" - ), - raising=exc.TypeError, - ) - - r = ffi.new("unsigned char[]", crypto_core_ed25519_SCALARBYTES) - - lib.crypto_core_ed25519_scalar_sub(r, p, q) - - return ffi.buffer(r, crypto_core_ed25519_SCALARBYTES)[:] - - -def crypto_core_ed25519_scalar_mul(p: bytes, q: bytes) -> bytes: - """ - Multiply integers ``p`` and ``q`` modulo ``L``, where ``L`` is the - order of the main subgroup. - - :param p: a :py:data:`.crypto_core_ed25519_SCALARBYTES` - long bytes sequence representing an integer - :type p: bytes - :param q: a :py:data:`.crypto_core_ed25519_SCALARBYTES` - long bytes sequence representing an integer - :type q: bytes - :return: an integer represented as a - :py:data:`.crypto_core_ed25519_SCALARBYTES` long bytes sequence - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_core_ed25519, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - isinstance(p, bytes) - and isinstance(q, bytes) - and len(p) == crypto_core_ed25519_SCALARBYTES - and len(q) == crypto_core_ed25519_SCALARBYTES, - "Each integer must be a {} long bytes sequence".format( - "crypto_core_ed25519_SCALARBYTES" - ), - raising=exc.TypeError, - ) - - r = ffi.new("unsigned char[]", crypto_core_ed25519_SCALARBYTES) - - lib.crypto_core_ed25519_scalar_mul(r, p, q) - - return ffi.buffer(r, crypto_core_ed25519_SCALARBYTES)[:] - - -def crypto_core_ed25519_scalar_reduce(s: bytes) -> bytes: - """ - Reduce integer ``s`` to ``s`` modulo ``L``, where ``L`` is the order - of the main subgroup. - - :param s: a :py:data:`.crypto_core_ed25519_NONREDUCEDSCALARBYTES` - long bytes sequence representing an integer - :type s: bytes - :return: an integer represented as a - :py:data:`.crypto_core_ed25519_SCALARBYTES` long bytes sequence - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_core_ed25519, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - isinstance(s, bytes) - and len(s) == crypto_core_ed25519_NONREDUCEDSCALARBYTES, - "Integer s must be a {} long bytes sequence".format( - "crypto_core_ed25519_NONREDUCEDSCALARBYTES" - ), - raising=exc.TypeError, - ) - - r = ffi.new("unsigned char[]", crypto_core_ed25519_SCALARBYTES) - - lib.crypto_core_ed25519_scalar_reduce(r, s) - - return ffi.buffer(r, crypto_core_ed25519_SCALARBYTES)[:] diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_generichash.py b/venv/lib/python3.10/site-packages/nacl/bindings/crypto_generichash.py deleted file mode 100644 index 6ab385a59718221f54738f1727c757447427b04d..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_generichash.py +++ /dev/null @@ -1,281 +0,0 @@ -# Copyright 2013-2019 Donald Stufft and individual contributors -# -# 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 NoReturn, TypeVar - -from nacl import exceptions as exc -from nacl._sodium import ffi, lib -from nacl.exceptions import ensure - - -crypto_generichash_BYTES: int = lib.crypto_generichash_blake2b_bytes() -crypto_generichash_BYTES_MIN: int = lib.crypto_generichash_blake2b_bytes_min() -crypto_generichash_BYTES_MAX: int = lib.crypto_generichash_blake2b_bytes_max() -crypto_generichash_KEYBYTES: int = lib.crypto_generichash_blake2b_keybytes() -crypto_generichash_KEYBYTES_MIN: int = ( - lib.crypto_generichash_blake2b_keybytes_min() -) -crypto_generichash_KEYBYTES_MAX: int = ( - lib.crypto_generichash_blake2b_keybytes_max() -) -crypto_generichash_SALTBYTES: int = lib.crypto_generichash_blake2b_saltbytes() -crypto_generichash_PERSONALBYTES: int = ( - lib.crypto_generichash_blake2b_personalbytes() -) -crypto_generichash_STATEBYTES: int = lib.crypto_generichash_statebytes() - -_OVERLONG = "{0} length greater than {1} bytes" -_TOOBIG = "{0} greater than {1}" - - -def _checkparams( - digest_size: int, key: bytes, salt: bytes, person: bytes -) -> None: - """Check hash parameters""" - ensure( - isinstance(key, bytes), - "Key must be a bytes sequence", - raising=exc.TypeError, - ) - - ensure( - isinstance(salt, bytes), - "Salt must be a bytes sequence", - raising=exc.TypeError, - ) - - ensure( - isinstance(person, bytes), - "Person must be a bytes sequence", - raising=exc.TypeError, - ) - - ensure( - isinstance(digest_size, int), - "Digest size must be an integer number", - raising=exc.TypeError, - ) - - ensure( - digest_size <= crypto_generichash_BYTES_MAX, - _TOOBIG.format("Digest_size", crypto_generichash_BYTES_MAX), - raising=exc.ValueError, - ) - - ensure( - len(key) <= crypto_generichash_KEYBYTES_MAX, - _OVERLONG.format("Key", crypto_generichash_KEYBYTES_MAX), - raising=exc.ValueError, - ) - - ensure( - len(salt) <= crypto_generichash_SALTBYTES, - _OVERLONG.format("Salt", crypto_generichash_SALTBYTES), - raising=exc.ValueError, - ) - - ensure( - len(person) <= crypto_generichash_PERSONALBYTES, - _OVERLONG.format("Person", crypto_generichash_PERSONALBYTES), - raising=exc.ValueError, - ) - - -def generichash_blake2b_salt_personal( - data: bytes, - digest_size: int = crypto_generichash_BYTES, - key: bytes = b"", - salt: bytes = b"", - person: bytes = b"", -) -> bytes: - """One shot hash interface - - :param data: the input data to the hash function - :type data: bytes - :param digest_size: must be at most - :py:data:`.crypto_generichash_BYTES_MAX`; - the default digest size is - :py:data:`.crypto_generichash_BYTES` - :type digest_size: int - :param key: must be at most - :py:data:`.crypto_generichash_KEYBYTES_MAX` long - :type key: bytes - :param salt: must be at most - :py:data:`.crypto_generichash_SALTBYTES` long; - will be zero-padded if needed - :type salt: bytes - :param person: must be at most - :py:data:`.crypto_generichash_PERSONALBYTES` long: - will be zero-padded if needed - :type person: bytes - :return: digest_size long digest - :rtype: bytes - """ - - _checkparams(digest_size, key, salt, person) - - ensure( - isinstance(data, bytes), - "Input data must be a bytes sequence", - raising=exc.TypeError, - ) - - digest = ffi.new("unsigned char[]", digest_size) - - # both _salt and _personal must be zero-padded to the correct length - _salt = ffi.new("unsigned char []", crypto_generichash_SALTBYTES) - _person = ffi.new("unsigned char []", crypto_generichash_PERSONALBYTES) - - ffi.memmove(_salt, salt, len(salt)) - ffi.memmove(_person, person, len(person)) - - rc = lib.crypto_generichash_blake2b_salt_personal( - digest, digest_size, data, len(data), key, len(key), _salt, _person - ) - ensure(rc == 0, "Unexpected failure", raising=exc.RuntimeError) - - return ffi.buffer(digest, digest_size)[:] - - -_Blake2State = TypeVar("_Blake2State", bound="Blake2State") - - -class Blake2State: - """ - Python-level wrapper for the crypto_generichash_blake2b state buffer - """ - - __slots__ = ["_statebuf", "digest_size"] - - def __init__(self, digest_size: int): - self._statebuf = ffi.new( - "unsigned char[]", crypto_generichash_STATEBYTES - ) - self.digest_size = digest_size - - def __reduce__(self) -> NoReturn: - """ - Raise the same exception as hashlib's blake implementation - on copy.copy() - """ - raise TypeError( - "can't pickle {} objects".format(self.__class__.__name__) - ) - - def copy(self: _Blake2State) -> _Blake2State: - _st = self.__class__(self.digest_size) - ffi.memmove( - _st._statebuf, self._statebuf, crypto_generichash_STATEBYTES - ) - return _st - - -def generichash_blake2b_init( - key: bytes = b"", - salt: bytes = b"", - person: bytes = b"", - digest_size: int = crypto_generichash_BYTES, -) -> Blake2State: - """ - Create a new initialized blake2b hash state - - :param key: must be at most - :py:data:`.crypto_generichash_KEYBYTES_MAX` long - :type key: bytes - :param salt: must be at most - :py:data:`.crypto_generichash_SALTBYTES` long; - will be zero-padded if needed - :type salt: bytes - :param person: must be at most - :py:data:`.crypto_generichash_PERSONALBYTES` long: - will be zero-padded if needed - :type person: bytes - :param digest_size: must be at most - :py:data:`.crypto_generichash_BYTES_MAX`; - the default digest size is - :py:data:`.crypto_generichash_BYTES` - :type digest_size: int - :return: a initialized :py:class:`.Blake2State` - :rtype: object - """ - - _checkparams(digest_size, key, salt, person) - - state = Blake2State(digest_size) - - # both _salt and _personal must be zero-padded to the correct length - _salt = ffi.new("unsigned char []", crypto_generichash_SALTBYTES) - _person = ffi.new("unsigned char []", crypto_generichash_PERSONALBYTES) - - ffi.memmove(_salt, salt, len(salt)) - ffi.memmove(_person, person, len(person)) - - rc = lib.crypto_generichash_blake2b_init_salt_personal( - state._statebuf, key, len(key), digest_size, _salt, _person - ) - ensure(rc == 0, "Unexpected failure", raising=exc.RuntimeError) - - return state - - -def generichash_blake2b_update(state: Blake2State, data: bytes) -> None: - """Update the blake2b hash state - - :param state: a initialized Blake2bState object as returned from - :py:func:`.crypto_generichash_blake2b_init` - :type state: :py:class:`.Blake2State` - :param data: - :type data: bytes - """ - - ensure( - isinstance(state, Blake2State), - "State must be a Blake2State object", - raising=exc.TypeError, - ) - - ensure( - isinstance(data, bytes), - "Input data must be a bytes sequence", - raising=exc.TypeError, - ) - - rc = lib.crypto_generichash_blake2b_update( - state._statebuf, data, len(data) - ) - ensure(rc == 0, "Unexpected failure", raising=exc.RuntimeError) - - -def generichash_blake2b_final(state: Blake2State) -> bytes: - """Finalize the blake2b hash state and return the digest. - - :param state: a initialized Blake2bState object as returned from - :py:func:`.crypto_generichash_blake2b_init` - :type state: :py:class:`.Blake2State` - :return: the blake2 digest of the passed-in data stream - :rtype: bytes - """ - - ensure( - isinstance(state, Blake2State), - "State must be a Blake2State object", - raising=exc.TypeError, - ) - - _digest = ffi.new("unsigned char[]", crypto_generichash_BYTES_MAX) - rc = lib.crypto_generichash_blake2b_final( - state._statebuf, _digest, state.digest_size - ) - - ensure(rc == 0, "Unexpected failure", raising=exc.RuntimeError) - return ffi.buffer(_digest, state.digest_size)[:] diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_hash.py b/venv/lib/python3.10/site-packages/nacl/bindings/crypto_hash.py deleted file mode 100644 index 2bab3991dff39423e0ef215f67a0b39877471828..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_hash.py +++ /dev/null @@ -1,63 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 nacl import exceptions as exc -from nacl._sodium import ffi, lib -from nacl.exceptions import ensure - - -# crypto_hash_BYTES = lib.crypto_hash_bytes() -crypto_hash_BYTES: int = lib.crypto_hash_sha512_bytes() -crypto_hash_sha256_BYTES: int = lib.crypto_hash_sha256_bytes() -crypto_hash_sha512_BYTES: int = lib.crypto_hash_sha512_bytes() - - -def crypto_hash(message: bytes) -> bytes: - """ - Hashes and returns the message ``message``. - - :param message: bytes - :rtype: bytes - """ - digest = ffi.new("unsigned char[]", crypto_hash_BYTES) - rc = lib.crypto_hash(digest, message, len(message)) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - return ffi.buffer(digest, crypto_hash_BYTES)[:] - - -def crypto_hash_sha256(message: bytes) -> bytes: - """ - Hashes and returns the message ``message``. - - :param message: bytes - :rtype: bytes - """ - digest = ffi.new("unsigned char[]", crypto_hash_sha256_BYTES) - rc = lib.crypto_hash_sha256(digest, message, len(message)) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - return ffi.buffer(digest, crypto_hash_sha256_BYTES)[:] - - -def crypto_hash_sha512(message: bytes) -> bytes: - """ - Hashes and returns the message ``message``. - - :param message: bytes - :rtype: bytes - """ - digest = ffi.new("unsigned char[]", crypto_hash_sha512_BYTES) - rc = lib.crypto_hash_sha512(digest, message, len(message)) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - return ffi.buffer(digest, crypto_hash_sha512_BYTES)[:] diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_kx.py b/venv/lib/python3.10/site-packages/nacl/bindings/crypto_kx.py deleted file mode 100644 index 172a19f804d9958d1ff090f261969a1607b93b7b..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_kx.py +++ /dev/null @@ -1,200 +0,0 @@ -# Copyright 2018 Donald Stufft and individual contributors -# -# 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 Tuple - -from nacl import exceptions as exc -from nacl._sodium import ffi, lib -from nacl.exceptions import ensure - -__all__ = [ - "crypto_kx_keypair", - "crypto_kx_client_session_keys", - "crypto_kx_server_session_keys", - "crypto_kx_PUBLIC_KEY_BYTES", - "crypto_kx_SECRET_KEY_BYTES", - "crypto_kx_SEED_BYTES", - "crypto_kx_SESSION_KEY_BYTES", -] - -""" -Implementations of client, server key exchange -""" -crypto_kx_PUBLIC_KEY_BYTES: int = lib.crypto_kx_publickeybytes() -crypto_kx_SECRET_KEY_BYTES: int = lib.crypto_kx_secretkeybytes() -crypto_kx_SEED_BYTES: int = lib.crypto_kx_seedbytes() -crypto_kx_SESSION_KEY_BYTES: int = lib.crypto_kx_sessionkeybytes() - - -def crypto_kx_keypair() -> Tuple[bytes, bytes]: - """ - Generate a keypair. - This is a duplicate crypto_box_keypair, but - is included for api consistency. - :return: (public_key, secret_key) - :rtype: (bytes, bytes) - """ - public_key = ffi.new("unsigned char[]", crypto_kx_PUBLIC_KEY_BYTES) - secret_key = ffi.new("unsigned char[]", crypto_kx_SECRET_KEY_BYTES) - res = lib.crypto_kx_keypair(public_key, secret_key) - ensure(res == 0, "Key generation failed.", raising=exc.CryptoError) - - return ( - ffi.buffer(public_key, crypto_kx_PUBLIC_KEY_BYTES)[:], - ffi.buffer(secret_key, crypto_kx_SECRET_KEY_BYTES)[:], - ) - - -def crypto_kx_seed_keypair(seed: bytes) -> Tuple[bytes, bytes]: - """ - Generate a keypair with a given seed. - This is functionally the same as crypto_box_seed_keypair, however - it uses the blake2b hash primitive instead of sha512. - It is included mainly for api consistency when using crypto_kx. - :param seed: random seed - :type seed: bytes - :return: (public_key, secret_key) - :rtype: (bytes, bytes) - """ - public_key = ffi.new("unsigned char[]", crypto_kx_PUBLIC_KEY_BYTES) - secret_key = ffi.new("unsigned char[]", crypto_kx_SECRET_KEY_BYTES) - ensure( - isinstance(seed, bytes) and len(seed) == crypto_kx_SEED_BYTES, - "Seed must be a {} byte long bytes sequence".format( - crypto_kx_SEED_BYTES - ), - raising=exc.TypeError, - ) - res = lib.crypto_kx_seed_keypair(public_key, secret_key, seed) - ensure(res == 0, "Key generation failed.", raising=exc.CryptoError) - - return ( - ffi.buffer(public_key, crypto_kx_PUBLIC_KEY_BYTES)[:], - ffi.buffer(secret_key, crypto_kx_SECRET_KEY_BYTES)[:], - ) - - -def crypto_kx_client_session_keys( - client_public_key: bytes, - client_secret_key: bytes, - server_public_key: bytes, -) -> Tuple[bytes, bytes]: - """ - Generate session keys for the client. - :param client_public_key: - :type client_public_key: bytes - :param client_secret_key: - :type client_secret_key: bytes - :param server_public_key: - :type server_public_key: bytes - :return: (rx_key, tx_key) - :rtype: (bytes, bytes) - """ - ensure( - isinstance(client_public_key, bytes) - and len(client_public_key) == crypto_kx_PUBLIC_KEY_BYTES, - "Client public key must be a {} bytes long bytes sequence".format( - crypto_kx_PUBLIC_KEY_BYTES - ), - raising=exc.TypeError, - ) - ensure( - isinstance(client_secret_key, bytes) - and len(client_secret_key) == crypto_kx_SECRET_KEY_BYTES, - "Client secret key must be a {} bytes long bytes sequence".format( - crypto_kx_PUBLIC_KEY_BYTES - ), - raising=exc.TypeError, - ) - ensure( - isinstance(server_public_key, bytes) - and len(server_public_key) == crypto_kx_PUBLIC_KEY_BYTES, - "Server public key must be a {} bytes long bytes sequence".format( - crypto_kx_PUBLIC_KEY_BYTES - ), - raising=exc.TypeError, - ) - - rx_key = ffi.new("unsigned char[]", crypto_kx_SESSION_KEY_BYTES) - tx_key = ffi.new("unsigned char[]", crypto_kx_SESSION_KEY_BYTES) - res = lib.crypto_kx_client_session_keys( - rx_key, tx_key, client_public_key, client_secret_key, server_public_key - ) - ensure( - res == 0, - "Client session key generation failed.", - raising=exc.CryptoError, - ) - - return ( - ffi.buffer(rx_key, crypto_kx_SESSION_KEY_BYTES)[:], - ffi.buffer(tx_key, crypto_kx_SESSION_KEY_BYTES)[:], - ) - - -def crypto_kx_server_session_keys( - server_public_key: bytes, - server_secret_key: bytes, - client_public_key: bytes, -) -> Tuple[bytes, bytes]: - """ - Generate session keys for the server. - :param server_public_key: - :type server_public_key: bytes - :param server_secret_key: - :type server_secret_key: bytes - :param client_public_key: - :type client_public_key: bytes - :return: (rx_key, tx_key) - :rtype: (bytes, bytes) - """ - ensure( - isinstance(server_public_key, bytes) - and len(server_public_key) == crypto_kx_PUBLIC_KEY_BYTES, - "Server public key must be a {} bytes long bytes sequence".format( - crypto_kx_PUBLIC_KEY_BYTES - ), - raising=exc.TypeError, - ) - ensure( - isinstance(server_secret_key, bytes) - and len(server_secret_key) == crypto_kx_SECRET_KEY_BYTES, - "Server secret key must be a {} bytes long bytes sequence".format( - crypto_kx_PUBLIC_KEY_BYTES - ), - raising=exc.TypeError, - ) - ensure( - isinstance(client_public_key, bytes) - and len(client_public_key) == crypto_kx_PUBLIC_KEY_BYTES, - "Client public key must be a {} bytes long bytes sequence".format( - crypto_kx_PUBLIC_KEY_BYTES - ), - raising=exc.TypeError, - ) - - rx_key = ffi.new("unsigned char[]", crypto_kx_SESSION_KEY_BYTES) - tx_key = ffi.new("unsigned char[]", crypto_kx_SESSION_KEY_BYTES) - res = lib.crypto_kx_server_session_keys( - rx_key, tx_key, server_public_key, server_secret_key, client_public_key - ) - ensure( - res == 0, - "Server session key generation failed.", - raising=exc.CryptoError, - ) - - return ( - ffi.buffer(rx_key, crypto_kx_SESSION_KEY_BYTES)[:], - ffi.buffer(tx_key, crypto_kx_SESSION_KEY_BYTES)[:], - ) diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_pwhash.py b/venv/lib/python3.10/site-packages/nacl/bindings/crypto_pwhash.py deleted file mode 100644 index 0c4cc1a24402a1ca63a499a22c1a38c97c3c5ab4..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_pwhash.py +++ /dev/null @@ -1,600 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 -from typing import Tuple - -import nacl.exceptions as exc -from nacl._sodium import ffi, lib -from nacl.exceptions import ensure - - -has_crypto_pwhash_scryptsalsa208sha256 = bool( - lib.PYNACL_HAS_CRYPTO_PWHASH_SCRYPTSALSA208SHA256 -) - -crypto_pwhash_scryptsalsa208sha256_STRPREFIX = b"" -crypto_pwhash_scryptsalsa208sha256_SALTBYTES = 0 -crypto_pwhash_scryptsalsa208sha256_STRBYTES = 0 -crypto_pwhash_scryptsalsa208sha256_PASSWD_MIN = 0 -crypto_pwhash_scryptsalsa208sha256_PASSWD_MAX = 0 -crypto_pwhash_scryptsalsa208sha256_BYTES_MIN = 0 -crypto_pwhash_scryptsalsa208sha256_BYTES_MAX = 0 -crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_MIN = 0 -crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_MAX = 0 -crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_MIN = 0 -crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_MAX = 0 -crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_INTERACTIVE = 0 -crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_INTERACTIVE = 0 -crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_SENSITIVE = 0 -crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_SENSITIVE = 0 - -if has_crypto_pwhash_scryptsalsa208sha256: - crypto_pwhash_scryptsalsa208sha256_STRPREFIX = ffi.string( - ffi.cast("char *", lib.crypto_pwhash_scryptsalsa208sha256_strprefix()) - )[:] - crypto_pwhash_scryptsalsa208sha256_SALTBYTES = ( - lib.crypto_pwhash_scryptsalsa208sha256_saltbytes() - ) - crypto_pwhash_scryptsalsa208sha256_STRBYTES = ( - lib.crypto_pwhash_scryptsalsa208sha256_strbytes() - ) - crypto_pwhash_scryptsalsa208sha256_PASSWD_MIN = ( - lib.crypto_pwhash_scryptsalsa208sha256_passwd_min() - ) - crypto_pwhash_scryptsalsa208sha256_PASSWD_MAX = ( - lib.crypto_pwhash_scryptsalsa208sha256_passwd_max() - ) - crypto_pwhash_scryptsalsa208sha256_BYTES_MIN = ( - lib.crypto_pwhash_scryptsalsa208sha256_bytes_min() - ) - crypto_pwhash_scryptsalsa208sha256_BYTES_MAX = ( - lib.crypto_pwhash_scryptsalsa208sha256_bytes_max() - ) - crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_MIN = ( - lib.crypto_pwhash_scryptsalsa208sha256_memlimit_min() - ) - crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_MAX = ( - lib.crypto_pwhash_scryptsalsa208sha256_memlimit_max() - ) - crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_MIN = ( - lib.crypto_pwhash_scryptsalsa208sha256_opslimit_min() - ) - crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_MAX = ( - lib.crypto_pwhash_scryptsalsa208sha256_opslimit_max() - ) - crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_INTERACTIVE = ( - lib.crypto_pwhash_scryptsalsa208sha256_opslimit_interactive() - ) - crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_INTERACTIVE = ( - lib.crypto_pwhash_scryptsalsa208sha256_memlimit_interactive() - ) - crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_SENSITIVE = ( - lib.crypto_pwhash_scryptsalsa208sha256_opslimit_sensitive() - ) - crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_SENSITIVE = ( - lib.crypto_pwhash_scryptsalsa208sha256_memlimit_sensitive() - ) - -crypto_pwhash_ALG_ARGON2I13: int = lib.crypto_pwhash_alg_argon2i13() -crypto_pwhash_ALG_ARGON2ID13: int = lib.crypto_pwhash_alg_argon2id13() -crypto_pwhash_ALG_DEFAULT: int = lib.crypto_pwhash_alg_default() - -crypto_pwhash_SALTBYTES: int = lib.crypto_pwhash_saltbytes() -crypto_pwhash_STRBYTES: int = lib.crypto_pwhash_strbytes() - -crypto_pwhash_PASSWD_MIN: int = lib.crypto_pwhash_passwd_min() -crypto_pwhash_PASSWD_MAX: int = lib.crypto_pwhash_passwd_max() -crypto_pwhash_BYTES_MIN: int = lib.crypto_pwhash_bytes_min() -crypto_pwhash_BYTES_MAX: int = lib.crypto_pwhash_bytes_max() - -crypto_pwhash_argon2i_STRPREFIX: bytes = ffi.string( - ffi.cast("char *", lib.crypto_pwhash_argon2i_strprefix()) -)[:] -crypto_pwhash_argon2i_MEMLIMIT_MIN: int = ( - lib.crypto_pwhash_argon2i_memlimit_min() -) -crypto_pwhash_argon2i_MEMLIMIT_MAX: int = ( - lib.crypto_pwhash_argon2i_memlimit_max() -) -crypto_pwhash_argon2i_OPSLIMIT_MIN: int = ( - lib.crypto_pwhash_argon2i_opslimit_min() -) -crypto_pwhash_argon2i_OPSLIMIT_MAX: int = ( - lib.crypto_pwhash_argon2i_opslimit_max() -) -crypto_pwhash_argon2i_OPSLIMIT_INTERACTIVE: int = ( - lib.crypto_pwhash_argon2i_opslimit_interactive() -) -crypto_pwhash_argon2i_MEMLIMIT_INTERACTIVE: int = ( - lib.crypto_pwhash_argon2i_memlimit_interactive() -) -crypto_pwhash_argon2i_OPSLIMIT_MODERATE: int = ( - lib.crypto_pwhash_argon2i_opslimit_moderate() -) -crypto_pwhash_argon2i_MEMLIMIT_MODERATE: int = ( - lib.crypto_pwhash_argon2i_memlimit_moderate() -) -crypto_pwhash_argon2i_OPSLIMIT_SENSITIVE: int = ( - lib.crypto_pwhash_argon2i_opslimit_sensitive() -) -crypto_pwhash_argon2i_MEMLIMIT_SENSITIVE: int = ( - lib.crypto_pwhash_argon2i_memlimit_sensitive() -) - -crypto_pwhash_argon2id_STRPREFIX: bytes = ffi.string( - ffi.cast("char *", lib.crypto_pwhash_argon2id_strprefix()) -)[:] -crypto_pwhash_argon2id_MEMLIMIT_MIN: int = ( - lib.crypto_pwhash_argon2id_memlimit_min() -) -crypto_pwhash_argon2id_MEMLIMIT_MAX: int = ( - lib.crypto_pwhash_argon2id_memlimit_max() -) -crypto_pwhash_argon2id_OPSLIMIT_MIN: int = ( - lib.crypto_pwhash_argon2id_opslimit_min() -) -crypto_pwhash_argon2id_OPSLIMIT_MAX: int = ( - lib.crypto_pwhash_argon2id_opslimit_max() -) -crypto_pwhash_argon2id_OPSLIMIT_INTERACTIVE: int = ( - lib.crypto_pwhash_argon2id_opslimit_interactive() -) -crypto_pwhash_argon2id_MEMLIMIT_INTERACTIVE: int = ( - lib.crypto_pwhash_argon2id_memlimit_interactive() -) -crypto_pwhash_argon2id_OPSLIMIT_MODERATE: int = ( - lib.crypto_pwhash_argon2id_opslimit_moderate() -) -crypto_pwhash_argon2id_MEMLIMIT_MODERATE: int = ( - lib.crypto_pwhash_argon2id_memlimit_moderate() -) -crypto_pwhash_argon2id_OPSLIMIT_SENSITIVE: int = ( - lib.crypto_pwhash_argon2id_opslimit_sensitive() -) -crypto_pwhash_argon2id_MEMLIMIT_SENSITIVE: int = ( - lib.crypto_pwhash_argon2id_memlimit_sensitive() -) - -SCRYPT_OPSLIMIT_INTERACTIVE = ( - crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_INTERACTIVE -) -SCRYPT_MEMLIMIT_INTERACTIVE = ( - crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_INTERACTIVE -) -SCRYPT_OPSLIMIT_SENSITIVE = ( - crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_SENSITIVE -) -SCRYPT_MEMLIMIT_SENSITIVE = ( - crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_SENSITIVE -) -SCRYPT_SALTBYTES = crypto_pwhash_scryptsalsa208sha256_SALTBYTES -SCRYPT_STRBYTES = crypto_pwhash_scryptsalsa208sha256_STRBYTES - -SCRYPT_PR_MAX = (1 << 30) - 1 -LOG2_UINT64_MAX = 63 -UINT64_MAX = (1 << 64) - 1 -SCRYPT_MAX_MEM = 32 * (1024 * 1024) - - -def _check_memory_occupation( - n: int, r: int, p: int, maxmem: int = SCRYPT_MAX_MEM -) -> None: - ensure(r != 0, "Invalid block size", raising=exc.ValueError) - - ensure(p != 0, "Invalid parallelization factor", raising=exc.ValueError) - - ensure( - (n & (n - 1)) == 0, - "Cost factor must be a power of 2", - raising=exc.ValueError, - ) - - ensure(n > 1, "Cost factor must be at least 2", raising=exc.ValueError) - - ensure( - p <= SCRYPT_PR_MAX / r, - "p*r is greater than {}".format(SCRYPT_PR_MAX), - raising=exc.ValueError, - ) - - ensure(n < (1 << (16 * r)), raising=exc.ValueError) - - Blen = p * 128 * r - - i = UINT64_MAX / 128 - - ensure(n + 2 <= i / r, raising=exc.ValueError) - - Vlen = 32 * r * (n + 2) * 4 - - ensure(Blen <= UINT64_MAX - Vlen, raising=exc.ValueError) - - ensure(Blen <= sys.maxsize - Vlen, raising=exc.ValueError) - - ensure( - Blen + Vlen <= maxmem, - "Memory limit would be exceeded with the choosen n, r, p", - raising=exc.ValueError, - ) - - -def nacl_bindings_pick_scrypt_params( - opslimit: int, memlimit: int -) -> Tuple[int, int, int]: - """Python implementation of libsodium's pickparams""" - - if opslimit < 32768: - opslimit = 32768 - - r = 8 - - if opslimit < (memlimit // 32): - p = 1 - maxn = opslimit // (4 * r) - for n_log2 in range(1, 63): # pragma: no branch - if (2 ** n_log2) > (maxn // 2): - break - else: - maxn = memlimit // (r * 128) - for n_log2 in range(1, 63): # pragma: no branch - if (2 ** n_log2) > maxn // 2: - break - - maxrp = (opslimit // 4) // (2 ** n_log2) - - if maxrp > 0x3FFFFFFF: # pragma: no cover - maxrp = 0x3FFFFFFF - - p = maxrp // r - - return n_log2, r, p - - -def crypto_pwhash_scryptsalsa208sha256_ll( - passwd: bytes, - salt: bytes, - n: int, - r: int, - p: int, - dklen: int = 64, - maxmem: int = SCRYPT_MAX_MEM, -) -> bytes: - """ - Derive a cryptographic key using the ``passwd`` and ``salt`` - given as input. - - The work factor can be tuned by by picking different - values for the parameters - - :param bytes passwd: - :param bytes salt: - :param bytes salt: *must* be *exactly* :py:const:`.SALTBYTES` long - :param int dklen: - :param int opslimit: - :param int n: - :param int r: block size, - :param int p: the parallelism factor - :param int maxmem: the maximum available memory available for scrypt's - operations - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_pwhash_scryptsalsa208sha256, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure(isinstance(n, int), raising=TypeError) - ensure(isinstance(r, int), raising=TypeError) - ensure(isinstance(p, int), raising=TypeError) - - ensure(isinstance(passwd, bytes), raising=TypeError) - ensure(isinstance(salt, bytes), raising=TypeError) - - _check_memory_occupation(n, r, p, maxmem) - - buf = ffi.new("uint8_t[]", dklen) - - ret = lib.crypto_pwhash_scryptsalsa208sha256_ll( - passwd, len(passwd), salt, len(salt), n, r, p, buf, dklen - ) - - ensure( - ret == 0, - "Unexpected failure in key derivation", - raising=exc.RuntimeError, - ) - - return ffi.buffer(ffi.cast("char *", buf), dklen)[:] - - -def crypto_pwhash_scryptsalsa208sha256_str( - passwd: bytes, - opslimit: int = SCRYPT_OPSLIMIT_INTERACTIVE, - memlimit: int = SCRYPT_MEMLIMIT_INTERACTIVE, -) -> bytes: - """ - Derive a cryptographic key using the ``passwd`` and ``salt`` - given as input, returning a string representation which includes - the salt and the tuning parameters. - - The returned string can be directly stored as a password hash. - - See :py:func:`.crypto_pwhash_scryptsalsa208sha256` for a short - discussion about ``opslimit`` and ``memlimit`` values. - - :param bytes passwd: - :param int opslimit: - :param int memlimit: - :return: serialized key hash, including salt and tuning parameters - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_pwhash_scryptsalsa208sha256, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - buf = ffi.new("char[]", SCRYPT_STRBYTES) - - ret = lib.crypto_pwhash_scryptsalsa208sha256_str( - buf, passwd, len(passwd), opslimit, memlimit - ) - - ensure( - ret == 0, - "Unexpected failure in password hashing", - raising=exc.RuntimeError, - ) - - return ffi.string(buf) - - -def crypto_pwhash_scryptsalsa208sha256_str_verify( - passwd_hash: bytes, passwd: bytes -) -> bool: - """ - Verifies the ``passwd`` against the ``passwd_hash`` that was generated. - Returns True or False depending on the success - - :param passwd_hash: bytes - :param passwd: bytes - :rtype: boolean - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_pwhash_scryptsalsa208sha256, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - len(passwd_hash) == SCRYPT_STRBYTES - 1, - "Invalid password hash", - raising=exc.ValueError, - ) - - ret = lib.crypto_pwhash_scryptsalsa208sha256_str_verify( - passwd_hash, passwd, len(passwd) - ) - ensure(ret == 0, "Wrong password", raising=exc.InvalidkeyError) - # all went well, therefore: - return True - - -def _check_argon2_limits_alg(opslimit: int, memlimit: int, alg: int) -> None: - - if alg == crypto_pwhash_ALG_ARGON2I13: - if memlimit < crypto_pwhash_argon2i_MEMLIMIT_MIN: - raise exc.ValueError( - "memlimit must be at least {} bytes".format( - crypto_pwhash_argon2i_MEMLIMIT_MIN - ) - ) - elif memlimit > crypto_pwhash_argon2i_MEMLIMIT_MAX: - raise exc.ValueError( - "memlimit must be at most {} bytes".format( - crypto_pwhash_argon2i_MEMLIMIT_MAX - ) - ) - if opslimit < crypto_pwhash_argon2i_OPSLIMIT_MIN: - raise exc.ValueError( - "opslimit must be at least {}".format( - crypto_pwhash_argon2i_OPSLIMIT_MIN - ) - ) - elif opslimit > crypto_pwhash_argon2i_OPSLIMIT_MAX: - raise exc.ValueError( - "opslimit must be at most {}".format( - crypto_pwhash_argon2i_OPSLIMIT_MAX - ) - ) - - elif alg == crypto_pwhash_ALG_ARGON2ID13: - if memlimit < crypto_pwhash_argon2id_MEMLIMIT_MIN: - raise exc.ValueError( - "memlimit must be at least {} bytes".format( - crypto_pwhash_argon2id_MEMLIMIT_MIN - ) - ) - elif memlimit > crypto_pwhash_argon2id_MEMLIMIT_MAX: - raise exc.ValueError( - "memlimit must be at most {} bytes".format( - crypto_pwhash_argon2id_MEMLIMIT_MAX - ) - ) - if opslimit < crypto_pwhash_argon2id_OPSLIMIT_MIN: - raise exc.ValueError( - "opslimit must be at least {}".format( - crypto_pwhash_argon2id_OPSLIMIT_MIN - ) - ) - elif opslimit > crypto_pwhash_argon2id_OPSLIMIT_MAX: - raise exc.ValueError( - "opslimit must be at most {}".format( - crypto_pwhash_argon2id_OPSLIMIT_MAX - ) - ) - else: - raise exc.TypeError("Unsupported algorithm") - - -def crypto_pwhash_alg( - outlen: int, - passwd: bytes, - salt: bytes, - opslimit: int, - memlimit: int, - alg: int, -) -> bytes: - """ - Derive a raw cryptographic key using the ``passwd`` and the ``salt`` - given as input to the ``alg`` algorithm. - - :param outlen: the length of the derived key - :type outlen: int - :param passwd: The input password - :type passwd: bytes - :param salt: - :type salt: bytes - :param opslimit: computational cost - :type opslimit: int - :param memlimit: memory cost - :type memlimit: int - :param alg: algorithm identifier - :type alg: int - :return: derived key - :rtype: bytes - """ - ensure(isinstance(outlen, int), raising=exc.TypeError) - ensure(isinstance(opslimit, int), raising=exc.TypeError) - ensure(isinstance(memlimit, int), raising=exc.TypeError) - ensure(isinstance(alg, int), raising=exc.TypeError) - ensure(isinstance(passwd, bytes), raising=exc.TypeError) - - if len(salt) != crypto_pwhash_SALTBYTES: - raise exc.ValueError( - "salt must be exactly {} bytes long".format( - crypto_pwhash_SALTBYTES - ) - ) - - if outlen < crypto_pwhash_BYTES_MIN: - raise exc.ValueError( - "derived key must be at least {} bytes long".format( - crypto_pwhash_BYTES_MIN - ) - ) - - elif outlen > crypto_pwhash_BYTES_MAX: - raise exc.ValueError( - "derived key must be at most {} bytes long".format( - crypto_pwhash_BYTES_MAX - ) - ) - - _check_argon2_limits_alg(opslimit, memlimit, alg) - - outbuf = ffi.new("unsigned char[]", outlen) - - ret = lib.crypto_pwhash( - outbuf, outlen, passwd, len(passwd), salt, opslimit, memlimit, alg - ) - - ensure( - ret == 0, - "Unexpected failure in key derivation", - raising=exc.RuntimeError, - ) - - return ffi.buffer(outbuf, outlen)[:] - - -def crypto_pwhash_str_alg( - passwd: bytes, - opslimit: int, - memlimit: int, - alg: int, -) -> bytes: - """ - Derive a cryptographic key using the ``passwd`` given as input - and a random salt, returning a string representation which - includes the salt, the tuning parameters and the used algorithm. - - :param passwd: The input password - :type passwd: bytes - :param opslimit: computational cost - :type opslimit: int - :param memlimit: memory cost - :type memlimit: int - :param alg: The algorithm to use - :type alg: int - :return: serialized derived key and parameters - :rtype: bytes - """ - ensure(isinstance(opslimit, int), raising=TypeError) - ensure(isinstance(memlimit, int), raising=TypeError) - ensure(isinstance(passwd, bytes), raising=TypeError) - - _check_argon2_limits_alg(opslimit, memlimit, alg) - - outbuf = ffi.new("char[]", 128) - - ret = lib.crypto_pwhash_str_alg( - outbuf, passwd, len(passwd), opslimit, memlimit, alg - ) - - ensure( - ret == 0, - "Unexpected failure in key derivation", - raising=exc.RuntimeError, - ) - - return ffi.string(outbuf) - - -def crypto_pwhash_str_verify(passwd_hash: bytes, passwd: bytes) -> bool: - """ - Verifies the ``passwd`` against a given password hash. - - Returns True on success, raises InvalidkeyError on failure - :param passwd_hash: saved password hash - :type passwd_hash: bytes - :param passwd: password to be checked - :type passwd: bytes - :return: success - :rtype: boolean - """ - ensure(isinstance(passwd_hash, bytes), raising=TypeError) - ensure(isinstance(passwd, bytes), raising=TypeError) - ensure( - len(passwd_hash) <= 127, - "Hash must be at most 127 bytes long", - raising=exc.ValueError, - ) - - ret = lib.crypto_pwhash_str_verify(passwd_hash, passwd, len(passwd)) - - ensure(ret == 0, "Wrong password", raising=exc.InvalidkeyError) - # all went well, therefore: - return True - - -crypto_pwhash_argon2i_str_verify = crypto_pwhash_str_verify diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_scalarmult.py b/venv/lib/python3.10/site-packages/nacl/bindings/crypto_scalarmult.py deleted file mode 100644 index ca4a2819d9a01f05da2df5dea3d13a3f23a64d22..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_scalarmult.py +++ /dev/null @@ -1,240 +0,0 @@ -# Copyright 2013-2018 Donald Stufft and individual contributors -# -# 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 nacl import exceptions as exc -from nacl._sodium import ffi, lib -from nacl.exceptions import ensure - - -has_crypto_scalarmult_ed25519 = bool(lib.PYNACL_HAS_CRYPTO_SCALARMULT_ED25519) - -crypto_scalarmult_BYTES: int = lib.crypto_scalarmult_bytes() -crypto_scalarmult_SCALARBYTES: int = lib.crypto_scalarmult_scalarbytes() - -crypto_scalarmult_ed25519_BYTES = 0 -crypto_scalarmult_ed25519_SCALARBYTES = 0 - -if has_crypto_scalarmult_ed25519: - crypto_scalarmult_ed25519_BYTES = lib.crypto_scalarmult_ed25519_bytes() - crypto_scalarmult_ed25519_SCALARBYTES = ( - lib.crypto_scalarmult_ed25519_scalarbytes() - ) - - -def crypto_scalarmult_base(n: bytes) -> bytes: - """ - Computes and returns the scalar product of a standard group element and an - integer ``n``. - - :param n: bytes - :rtype: bytes - """ - q = ffi.new("unsigned char[]", crypto_scalarmult_BYTES) - - rc = lib.crypto_scalarmult_base(q, n) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(q, crypto_scalarmult_SCALARBYTES)[:] - - -def crypto_scalarmult(n: bytes, p: bytes) -> bytes: - """ - Computes and returns the scalar product of the given group element and an - integer ``n``. - - :param p: bytes - :param n: bytes - :rtype: bytes - """ - q = ffi.new("unsigned char[]", crypto_scalarmult_BYTES) - - rc = lib.crypto_scalarmult(q, n, p) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(q, crypto_scalarmult_SCALARBYTES)[:] - - -def crypto_scalarmult_ed25519_base(n: bytes) -> bytes: - """ - Computes and returns the scalar product of a standard group element and an - integer ``n`` on the edwards25519 curve. - - :param n: a :py:data:`.crypto_scalarmult_ed25519_SCALARBYTES` long bytes - sequence representing a scalar - :type n: bytes - :return: a point on the edwards25519 curve, represented as a - :py:data:`.crypto_scalarmult_ed25519_BYTES` long bytes sequence - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_scalarmult_ed25519, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - isinstance(n, bytes) - and len(n) == crypto_scalarmult_ed25519_SCALARBYTES, - "Input must be a {} long bytes sequence".format( - "crypto_scalarmult_ed25519_SCALARBYTES" - ), - raising=exc.TypeError, - ) - - q = ffi.new("unsigned char[]", crypto_scalarmult_ed25519_BYTES) - - rc = lib.crypto_scalarmult_ed25519_base(q, n) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(q, crypto_scalarmult_ed25519_BYTES)[:] - - -def crypto_scalarmult_ed25519_base_noclamp(n: bytes) -> bytes: - """ - Computes and returns the scalar product of a standard group element and an - integer ``n`` on the edwards25519 curve. The integer ``n`` is not clamped. - - :param n: a :py:data:`.crypto_scalarmult_ed25519_SCALARBYTES` long bytes - sequence representing a scalar - :type n: bytes - :return: a point on the edwards25519 curve, represented as a - :py:data:`.crypto_scalarmult_ed25519_BYTES` long bytes sequence - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_scalarmult_ed25519, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - isinstance(n, bytes) - and len(n) == crypto_scalarmult_ed25519_SCALARBYTES, - "Input must be a {} long bytes sequence".format( - "crypto_scalarmult_ed25519_SCALARBYTES" - ), - raising=exc.TypeError, - ) - - q = ffi.new("unsigned char[]", crypto_scalarmult_ed25519_BYTES) - - rc = lib.crypto_scalarmult_ed25519_base_noclamp(q, n) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(q, crypto_scalarmult_ed25519_BYTES)[:] - - -def crypto_scalarmult_ed25519(n: bytes, p: bytes) -> bytes: - """ - Computes and returns the scalar product of a *clamped* integer ``n`` - and the given group element on the edwards25519 curve. - The scalar is clamped, as done in the public key generation case, - by setting to zero the bits in position [0, 1, 2, 255] and setting - to one the bit in position 254. - - :param n: a :py:data:`.crypto_scalarmult_ed25519_SCALARBYTES` long bytes - sequence representing a scalar - :type n: bytes - :param p: a :py:data:`.crypto_scalarmult_ed25519_BYTES` long bytes sequence - representing a point on the edwards25519 curve - :type p: bytes - :return: a point on the edwards25519 curve, represented as a - :py:data:`.crypto_scalarmult_ed25519_BYTES` long bytes sequence - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_scalarmult_ed25519, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - isinstance(n, bytes) - and len(n) == crypto_scalarmult_ed25519_SCALARBYTES, - "Input must be a {} long bytes sequence".format( - "crypto_scalarmult_ed25519_SCALARBYTES" - ), - raising=exc.TypeError, - ) - - ensure( - isinstance(p, bytes) and len(p) == crypto_scalarmult_ed25519_BYTES, - "Input must be a {} long bytes sequence".format( - "crypto_scalarmult_ed25519_BYTES" - ), - raising=exc.TypeError, - ) - - q = ffi.new("unsigned char[]", crypto_scalarmult_ed25519_BYTES) - - rc = lib.crypto_scalarmult_ed25519(q, n, p) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(q, crypto_scalarmult_ed25519_BYTES)[:] - - -def crypto_scalarmult_ed25519_noclamp(n: bytes, p: bytes) -> bytes: - """ - Computes and returns the scalar product of an integer ``n`` - and the given group element on the edwards25519 curve. The integer - ``n`` is not clamped. - - :param n: a :py:data:`.crypto_scalarmult_ed25519_SCALARBYTES` long bytes - sequence representing a scalar - :type n: bytes - :param p: a :py:data:`.crypto_scalarmult_ed25519_BYTES` long bytes sequence - representing a point on the edwards25519 curve - :type p: bytes - :return: a point on the edwards25519 curve, represented as a - :py:data:`.crypto_scalarmult_ed25519_BYTES` long bytes sequence - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_scalarmult_ed25519, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - isinstance(n, bytes) - and len(n) == crypto_scalarmult_ed25519_SCALARBYTES, - "Input must be a {} long bytes sequence".format( - "crypto_scalarmult_ed25519_SCALARBYTES" - ), - raising=exc.TypeError, - ) - - ensure( - isinstance(p, bytes) and len(p) == crypto_scalarmult_ed25519_BYTES, - "Input must be a {} long bytes sequence".format( - "crypto_scalarmult_ed25519_BYTES" - ), - raising=exc.TypeError, - ) - - q = ffi.new("unsigned char[]", crypto_scalarmult_ed25519_BYTES) - - rc = lib.crypto_scalarmult_ed25519_noclamp(q, n, p) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(q, crypto_scalarmult_ed25519_BYTES)[:] diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_secretbox.py b/venv/lib/python3.10/site-packages/nacl/bindings/crypto_secretbox.py deleted file mode 100644 index 2a632a2bc4d547a12a6402e5fbe01a1744e915a0..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_secretbox.py +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 nacl import exceptions as exc -from nacl._sodium import ffi, lib -from nacl.exceptions import ensure - - -crypto_secretbox_KEYBYTES: int = lib.crypto_secretbox_keybytes() -crypto_secretbox_NONCEBYTES: int = lib.crypto_secretbox_noncebytes() -crypto_secretbox_ZEROBYTES: int = lib.crypto_secretbox_zerobytes() -crypto_secretbox_BOXZEROBYTES: int = lib.crypto_secretbox_boxzerobytes() -crypto_secretbox_MACBYTES: int = lib.crypto_secretbox_macbytes() -crypto_secretbox_MESSAGEBYTES_MAX: int = ( - lib.crypto_secretbox_messagebytes_max() -) - - -def crypto_secretbox(message: bytes, nonce: bytes, key: bytes) -> bytes: - """ - Encrypts and returns the message ``message`` with the secret ``key`` and - the nonce ``nonce``. - - :param message: bytes - :param nonce: bytes - :param key: bytes - :rtype: bytes - """ - if len(key) != crypto_secretbox_KEYBYTES: - raise exc.ValueError("Invalid key") - - if len(nonce) != crypto_secretbox_NONCEBYTES: - raise exc.ValueError("Invalid nonce") - - padded = b"\x00" * crypto_secretbox_ZEROBYTES + message - ciphertext = ffi.new("unsigned char[]", len(padded)) - - res = lib.crypto_secretbox(ciphertext, padded, len(padded), nonce, key) - ensure(res == 0, "Encryption failed", raising=exc.CryptoError) - - ciphertext = ffi.buffer(ciphertext, len(padded)) - return ciphertext[crypto_secretbox_BOXZEROBYTES:] - - -def crypto_secretbox_open( - ciphertext: bytes, nonce: bytes, key: bytes -) -> bytes: - """ - Decrypt and returns the encrypted message ``ciphertext`` with the secret - ``key`` and the nonce ``nonce``. - - :param ciphertext: bytes - :param nonce: bytes - :param key: bytes - :rtype: bytes - """ - if len(key) != crypto_secretbox_KEYBYTES: - raise exc.ValueError("Invalid key") - - if len(nonce) != crypto_secretbox_NONCEBYTES: - raise exc.ValueError("Invalid nonce") - - padded = b"\x00" * crypto_secretbox_BOXZEROBYTES + ciphertext - plaintext = ffi.new("unsigned char[]", len(padded)) - - res = lib.crypto_secretbox_open(plaintext, padded, len(padded), nonce, key) - ensure( - res == 0, - "Decryption failed. Ciphertext failed verification", - raising=exc.CryptoError, - ) - - plaintext = ffi.buffer(plaintext, len(padded)) - return plaintext[crypto_secretbox_ZEROBYTES:] diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_secretstream.py b/venv/lib/python3.10/site-packages/nacl/bindings/crypto_secretstream.py deleted file mode 100644 index d7c6725e54f70a023cffd34f4f9ebd080cd73987..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_secretstream.py +++ /dev/null @@ -1,357 +0,0 @@ -# Copyright 2013-2018 Donald Stufft and individual contributors -# -# 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 ByteString, Optional, Tuple, cast - -from nacl import exceptions as exc -from nacl._sodium import ffi, lib -from nacl.exceptions import ensure - - -crypto_secretstream_xchacha20poly1305_ABYTES: int = ( - lib.crypto_secretstream_xchacha20poly1305_abytes() -) -crypto_secretstream_xchacha20poly1305_HEADERBYTES: int = ( - lib.crypto_secretstream_xchacha20poly1305_headerbytes() -) -crypto_secretstream_xchacha20poly1305_KEYBYTES: int = ( - lib.crypto_secretstream_xchacha20poly1305_keybytes() -) -crypto_secretstream_xchacha20poly1305_MESSAGEBYTES_MAX: int = ( - lib.crypto_secretstream_xchacha20poly1305_messagebytes_max() -) -crypto_secretstream_xchacha20poly1305_STATEBYTES: int = ( - lib.crypto_secretstream_xchacha20poly1305_statebytes() -) - - -crypto_secretstream_xchacha20poly1305_TAG_MESSAGE: int = ( - lib.crypto_secretstream_xchacha20poly1305_tag_message() -) -crypto_secretstream_xchacha20poly1305_TAG_PUSH: int = ( - lib.crypto_secretstream_xchacha20poly1305_tag_push() -) -crypto_secretstream_xchacha20poly1305_TAG_REKEY: int = ( - lib.crypto_secretstream_xchacha20poly1305_tag_rekey() -) -crypto_secretstream_xchacha20poly1305_TAG_FINAL: int = ( - lib.crypto_secretstream_xchacha20poly1305_tag_final() -) - - -def crypto_secretstream_xchacha20poly1305_keygen() -> bytes: - """ - Generate a key for use with - :func:`.crypto_secretstream_xchacha20poly1305_init_push`. - - """ - keybuf = ffi.new( - "unsigned char[]", - crypto_secretstream_xchacha20poly1305_KEYBYTES, - ) - lib.crypto_secretstream_xchacha20poly1305_keygen(keybuf) - return ffi.buffer(keybuf)[:] - - -class crypto_secretstream_xchacha20poly1305_state: - """ - An object wrapping the crypto_secretstream_xchacha20poly1305 state. - - """ - - __slots__ = ["statebuf", "rawbuf", "tagbuf"] - - def __init__(self) -> None: - """Initialize a clean state object.""" - self.statebuf: ByteString = ffi.new( - "unsigned char[]", - crypto_secretstream_xchacha20poly1305_STATEBYTES, - ) - - self.rawbuf: Optional[ByteString] = None - self.tagbuf: Optional[ByteString] = None - - -def crypto_secretstream_xchacha20poly1305_init_push( - state: crypto_secretstream_xchacha20poly1305_state, key: bytes -) -> bytes: - """ - Initialize a crypto_secretstream_xchacha20poly1305 encryption buffer. - - :param state: a secretstream state object - :type state: crypto_secretstream_xchacha20poly1305_state - :param key: must be - :data:`.crypto_secretstream_xchacha20poly1305_KEYBYTES` long - :type key: bytes - :return: header - :rtype: bytes - - """ - ensure( - isinstance(state, crypto_secretstream_xchacha20poly1305_state), - "State must be a crypto_secretstream_xchacha20poly1305_state object", - raising=exc.TypeError, - ) - ensure( - isinstance(key, bytes), - "Key must be a bytes sequence", - raising=exc.TypeError, - ) - ensure( - len(key) == crypto_secretstream_xchacha20poly1305_KEYBYTES, - "Invalid key length", - raising=exc.ValueError, - ) - - headerbuf = ffi.new( - "unsigned char []", - crypto_secretstream_xchacha20poly1305_HEADERBYTES, - ) - - rc = lib.crypto_secretstream_xchacha20poly1305_init_push( - state.statebuf, headerbuf, key - ) - ensure(rc == 0, "Unexpected failure", raising=exc.RuntimeError) - - return ffi.buffer(headerbuf)[:] - - -def crypto_secretstream_xchacha20poly1305_push( - state: crypto_secretstream_xchacha20poly1305_state, - m: bytes, - ad: Optional[bytes] = None, - tag: int = crypto_secretstream_xchacha20poly1305_TAG_MESSAGE, -) -> bytes: - """ - Add an encrypted message to the secret stream. - - :param state: a secretstream state object - :type state: crypto_secretstream_xchacha20poly1305_state - :param m: the message to encrypt, the maximum length of an individual - message is - :data:`.crypto_secretstream_xchacha20poly1305_MESSAGEBYTES_MAX`. - :type m: bytes - :param ad: additional data to include in the authentication tag - :type ad: bytes or None - :param tag: the message tag, usually - :data:`.crypto_secretstream_xchacha20poly1305_TAG_MESSAGE` or - :data:`.crypto_secretstream_xchacha20poly1305_TAG_FINAL`. - :type tag: int - :return: ciphertext - :rtype: bytes - - """ - ensure( - isinstance(state, crypto_secretstream_xchacha20poly1305_state), - "State must be a crypto_secretstream_xchacha20poly1305_state object", - raising=exc.TypeError, - ) - ensure(isinstance(m, bytes), "Message is not bytes", raising=exc.TypeError) - ensure( - len(m) <= crypto_secretstream_xchacha20poly1305_MESSAGEBYTES_MAX, - "Message is too long", - raising=exc.ValueError, - ) - ensure( - ad is None or isinstance(ad, bytes), - "Additional data must be bytes or None", - raising=exc.TypeError, - ) - - clen = len(m) + crypto_secretstream_xchacha20poly1305_ABYTES - if state.rawbuf is None or len(state.rawbuf) < clen: - state.rawbuf = ffi.new("unsigned char[]", clen) - - if ad is None: - ad = ffi.NULL - adlen = 0 - else: - adlen = len(ad) - - rc = lib.crypto_secretstream_xchacha20poly1305_push( - state.statebuf, - state.rawbuf, - ffi.NULL, - m, - len(m), - ad, - adlen, - tag, - ) - ensure(rc == 0, "Unexpected failure", raising=exc.RuntimeError) - - return ffi.buffer(state.rawbuf, clen)[:] - - -def crypto_secretstream_xchacha20poly1305_init_pull( - state: crypto_secretstream_xchacha20poly1305_state, - header: bytes, - key: bytes, -) -> None: - """ - Initialize a crypto_secretstream_xchacha20poly1305 decryption buffer. - - :param state: a secretstream state object - :type state: crypto_secretstream_xchacha20poly1305_state - :param header: must be - :data:`.crypto_secretstream_xchacha20poly1305_HEADERBYTES` long - :type header: bytes - :param key: must be - :data:`.crypto_secretstream_xchacha20poly1305_KEYBYTES` long - :type key: bytes - - """ - ensure( - isinstance(state, crypto_secretstream_xchacha20poly1305_state), - "State must be a crypto_secretstream_xchacha20poly1305_state object", - raising=exc.TypeError, - ) - ensure( - isinstance(header, bytes), - "Header must be a bytes sequence", - raising=exc.TypeError, - ) - ensure( - len(header) == crypto_secretstream_xchacha20poly1305_HEADERBYTES, - "Invalid header length", - raising=exc.ValueError, - ) - ensure( - isinstance(key, bytes), - "Key must be a bytes sequence", - raising=exc.TypeError, - ) - ensure( - len(key) == crypto_secretstream_xchacha20poly1305_KEYBYTES, - "Invalid key length", - raising=exc.ValueError, - ) - - if state.tagbuf is None: - state.tagbuf = ffi.new("unsigned char *") - - rc = lib.crypto_secretstream_xchacha20poly1305_init_pull( - state.statebuf, header, key - ) - ensure(rc == 0, "Unexpected failure", raising=exc.RuntimeError) - - -def crypto_secretstream_xchacha20poly1305_pull( - state: crypto_secretstream_xchacha20poly1305_state, - c: bytes, - ad: Optional[bytes] = None, -) -> Tuple[bytes, int]: - """ - Read a decrypted message from the secret stream. - - :param state: a secretstream state object - :type state: crypto_secretstream_xchacha20poly1305_state - :param c: the ciphertext to decrypt, the maximum length of an individual - ciphertext is - :data:`.crypto_secretstream_xchacha20poly1305_MESSAGEBYTES_MAX` + - :data:`.crypto_secretstream_xchacha20poly1305_ABYTES`. - :type c: bytes - :param ad: additional data to include in the authentication tag - :type ad: bytes or None - :return: (message, tag) - :rtype: (bytes, int) - - """ - ensure( - isinstance(state, crypto_secretstream_xchacha20poly1305_state), - "State must be a crypto_secretstream_xchacha20poly1305_state object", - raising=exc.TypeError, - ) - ensure( - state.tagbuf is not None, - ( - "State must be initialized using " - "crypto_secretstream_xchacha20poly1305_init_pull" - ), - raising=exc.ValueError, - ) - ensure( - isinstance(c, bytes), - "Ciphertext is not bytes", - raising=exc.TypeError, - ) - ensure( - len(c) >= crypto_secretstream_xchacha20poly1305_ABYTES, - "Ciphertext is too short", - raising=exc.ValueError, - ) - ensure( - len(c) - <= ( - crypto_secretstream_xchacha20poly1305_MESSAGEBYTES_MAX - + crypto_secretstream_xchacha20poly1305_ABYTES - ), - "Ciphertext is too long", - raising=exc.ValueError, - ) - ensure( - ad is None or isinstance(ad, bytes), - "Additional data must be bytes or None", - raising=exc.TypeError, - ) - - mlen = len(c) - crypto_secretstream_xchacha20poly1305_ABYTES - if state.rawbuf is None or len(state.rawbuf) < mlen: - state.rawbuf = ffi.new("unsigned char[]", mlen) - - if ad is None: - ad = ffi.NULL - adlen = 0 - else: - adlen = len(ad) - - rc = lib.crypto_secretstream_xchacha20poly1305_pull( - state.statebuf, - state.rawbuf, - ffi.NULL, - state.tagbuf, - c, - len(c), - ad, - adlen, - ) - ensure(rc == 0, "Unexpected failure", raising=exc.RuntimeError) - - # Cast safety: we `ensure` above that `state.tagbuf is not None`. - return ( - ffi.buffer(state.rawbuf, mlen)[:], - int(cast(bytes, state.tagbuf)[0]), - ) - - -def crypto_secretstream_xchacha20poly1305_rekey( - state: crypto_secretstream_xchacha20poly1305_state, -) -> None: - """ - Explicitly change the encryption key in the stream. - - Normally the stream is re-keyed as needed or an explicit ``tag`` of - :data:`.crypto_secretstream_xchacha20poly1305_TAG_REKEY` is added to a - message to ensure forward secrecy, but this method can be used instead - if the re-keying is controlled without adding the tag. - - :param state: a secretstream state object - :type state: crypto_secretstream_xchacha20poly1305_state - - """ - ensure( - isinstance(state, crypto_secretstream_xchacha20poly1305_state), - "State must be a crypto_secretstream_xchacha20poly1305_state object", - raising=exc.TypeError, - ) - lib.crypto_secretstream_xchacha20poly1305_rekey(state.statebuf) diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_shorthash.py b/venv/lib/python3.10/site-packages/nacl/bindings/crypto_shorthash.py deleted file mode 100644 index 8f7d209e44dc277323f4df4d8fc4e1b726603fb3..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_shorthash.py +++ /dev/null @@ -1,81 +0,0 @@ -# Copyright 2016 Donald Stufft and individual contributors -# -# 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 nacl.exceptions as exc -from nacl._sodium import ffi, lib -from nacl.exceptions import ensure - - -has_crypto_shorthash_siphashx24 = bool( - lib.PYNACL_HAS_CRYPTO_SHORTHASH_SIPHASHX24 -) - -BYTES: int = lib.crypto_shorthash_siphash24_bytes() -KEYBYTES: int = lib.crypto_shorthash_siphash24_keybytes() - -XBYTES = 0 -XKEYBYTES = 0 - -if has_crypto_shorthash_siphashx24: - XBYTES = lib.crypto_shorthash_siphashx24_bytes() - XKEYBYTES = lib.crypto_shorthash_siphashx24_keybytes() - - -def crypto_shorthash_siphash24(data: bytes, key: bytes) -> bytes: - """Compute a fast, cryptographic quality, keyed hash of the input data - - :param data: - :type data: bytes - :param key: len(key) must be equal to - :py:data:`.KEYBYTES` (16) - :type key: bytes - """ - if len(key) != KEYBYTES: - raise exc.ValueError( - "Key length must be exactly {} bytes".format(KEYBYTES) - ) - digest = ffi.new("unsigned char[]", BYTES) - rc = lib.crypto_shorthash_siphash24(digest, data, len(data), key) - - ensure(rc == 0, raising=exc.RuntimeError) - return ffi.buffer(digest, BYTES)[:] - - -def crypto_shorthash_siphashx24(data: bytes, key: bytes) -> bytes: - """Compute a fast, cryptographic quality, keyed hash of the input data - - :param data: - :type data: bytes - :param key: len(key) must be equal to - :py:data:`.XKEYBYTES` (16) - :type key: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - """ - ensure( - has_crypto_shorthash_siphashx24, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - if len(key) != XKEYBYTES: - raise exc.ValueError( - "Key length must be exactly {} bytes".format(XKEYBYTES) - ) - digest = ffi.new("unsigned char[]", XBYTES) - rc = lib.crypto_shorthash_siphashx24(digest, data, len(data), key) - - ensure(rc == 0, raising=exc.RuntimeError) - return ffi.buffer(digest, XBYTES)[:] diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_sign.py b/venv/lib/python3.10/site-packages/nacl/bindings/crypto_sign.py deleted file mode 100644 index de3be47a1b24d9953ec8cb30c08d9b855172bae5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/crypto_sign.py +++ /dev/null @@ -1,327 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 Tuple - -from nacl import exceptions as exc -from nacl._sodium import ffi, lib -from nacl.exceptions import ensure - - -crypto_sign_BYTES: int = lib.crypto_sign_bytes() -# crypto_sign_SEEDBYTES = lib.crypto_sign_seedbytes() -crypto_sign_SEEDBYTES: int = lib.crypto_sign_secretkeybytes() // 2 -crypto_sign_PUBLICKEYBYTES: int = lib.crypto_sign_publickeybytes() -crypto_sign_SECRETKEYBYTES: int = lib.crypto_sign_secretkeybytes() - -crypto_sign_curve25519_BYTES: int = lib.crypto_box_secretkeybytes() - -crypto_sign_ed25519ph_STATEBYTES: int = lib.crypto_sign_ed25519ph_statebytes() - - -def crypto_sign_keypair() -> Tuple[bytes, bytes]: - """ - Returns a randomly generated public key and secret key. - - :rtype: (bytes(public_key), bytes(secret_key)) - """ - pk = ffi.new("unsigned char[]", crypto_sign_PUBLICKEYBYTES) - sk = ffi.new("unsigned char[]", crypto_sign_SECRETKEYBYTES) - - rc = lib.crypto_sign_keypair(pk, sk) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ( - ffi.buffer(pk, crypto_sign_PUBLICKEYBYTES)[:], - ffi.buffer(sk, crypto_sign_SECRETKEYBYTES)[:], - ) - - -def crypto_sign_seed_keypair(seed: bytes) -> Tuple[bytes, bytes]: - """ - Computes and returns the public key and secret key using the seed ``seed``. - - :param seed: bytes - :rtype: (bytes(public_key), bytes(secret_key)) - """ - if len(seed) != crypto_sign_SEEDBYTES: - raise exc.ValueError("Invalid seed") - - pk = ffi.new("unsigned char[]", crypto_sign_PUBLICKEYBYTES) - sk = ffi.new("unsigned char[]", crypto_sign_SECRETKEYBYTES) - - rc = lib.crypto_sign_seed_keypair(pk, sk, seed) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ( - ffi.buffer(pk, crypto_sign_PUBLICKEYBYTES)[:], - ffi.buffer(sk, crypto_sign_SECRETKEYBYTES)[:], - ) - - -def crypto_sign(message: bytes, sk: bytes) -> bytes: - """ - Signs the message ``message`` using the secret key ``sk`` and returns the - signed message. - - :param message: bytes - :param sk: bytes - :rtype: bytes - """ - signed = ffi.new("unsigned char[]", len(message) + crypto_sign_BYTES) - signed_len = ffi.new("unsigned long long *") - - rc = lib.crypto_sign(signed, signed_len, message, len(message), sk) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(signed, signed_len[0])[:] - - -def crypto_sign_open(signed: bytes, pk: bytes) -> bytes: - """ - Verifies the signature of the signed message ``signed`` using the public - key ``pk`` and returns the unsigned message. - - :param signed: bytes - :param pk: bytes - :rtype: bytes - """ - message = ffi.new("unsigned char[]", len(signed)) - message_len = ffi.new("unsigned long long *") - - if ( - lib.crypto_sign_open(message, message_len, signed, len(signed), pk) - != 0 - ): - raise exc.BadSignatureError("Signature was forged or corrupt") - - return ffi.buffer(message, message_len[0])[:] - - -def crypto_sign_ed25519_pk_to_curve25519(public_key_bytes: bytes) -> bytes: - """ - Converts a public Ed25519 key (encoded as bytes ``public_key_bytes``) to - a public Curve25519 key as bytes. - - Raises a ValueError if ``public_key_bytes`` is not of length - ``crypto_sign_PUBLICKEYBYTES`` - - :param public_key_bytes: bytes - :rtype: bytes - """ - if len(public_key_bytes) != crypto_sign_PUBLICKEYBYTES: - raise exc.ValueError("Invalid curve public key") - - curve_public_key_len = crypto_sign_curve25519_BYTES - curve_public_key = ffi.new("unsigned char[]", curve_public_key_len) - - rc = lib.crypto_sign_ed25519_pk_to_curve25519( - curve_public_key, public_key_bytes - ) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(curve_public_key, curve_public_key_len)[:] - - -def crypto_sign_ed25519_sk_to_curve25519(secret_key_bytes: bytes) -> bytes: - """ - Converts a secret Ed25519 key (encoded as bytes ``secret_key_bytes``) to - a secret Curve25519 key as bytes. - - Raises a ValueError if ``secret_key_bytes``is not of length - ``crypto_sign_SECRETKEYBYTES`` - - :param secret_key_bytes: bytes - :rtype: bytes - """ - if len(secret_key_bytes) != crypto_sign_SECRETKEYBYTES: - raise exc.ValueError("Invalid curve secret key") - - curve_secret_key_len = crypto_sign_curve25519_BYTES - curve_secret_key = ffi.new("unsigned char[]", curve_secret_key_len) - - rc = lib.crypto_sign_ed25519_sk_to_curve25519( - curve_secret_key, secret_key_bytes - ) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(curve_secret_key, curve_secret_key_len)[:] - - -def crypto_sign_ed25519_sk_to_pk(secret_key_bytes: bytes) -> bytes: - """ - Extract the public Ed25519 key from a secret Ed25519 key (encoded - as bytes ``secret_key_bytes``). - - Raises a ValueError if ``secret_key_bytes``is not of length - ``crypto_sign_SECRETKEYBYTES`` - - :param secret_key_bytes: bytes - :rtype: bytes - """ - if len(secret_key_bytes) != crypto_sign_SECRETKEYBYTES: - raise exc.ValueError("Invalid secret key") - - return secret_key_bytes[crypto_sign_SEEDBYTES:] - - -def crypto_sign_ed25519_sk_to_seed(secret_key_bytes: bytes) -> bytes: - """ - Extract the seed from a secret Ed25519 key (encoded - as bytes ``secret_key_bytes``). - - Raises a ValueError if ``secret_key_bytes``is not of length - ``crypto_sign_SECRETKEYBYTES`` - - :param secret_key_bytes: bytes - :rtype: bytes - """ - if len(secret_key_bytes) != crypto_sign_SECRETKEYBYTES: - raise exc.ValueError("Invalid secret key") - - return secret_key_bytes[:crypto_sign_SEEDBYTES] - - -class crypto_sign_ed25519ph_state: - """ - State object wrapping the sha-512 state used in ed25519ph computation - """ - - __slots__ = ["state"] - - def __init__(self) -> None: - self.state: bytes = ffi.new( - "unsigned char[]", crypto_sign_ed25519ph_STATEBYTES - ) - - rc = lib.crypto_sign_ed25519ph_init(self.state) - - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - -def crypto_sign_ed25519ph_update( - edph: crypto_sign_ed25519ph_state, pmsg: bytes -) -> None: - """ - Update the hash state wrapped in edph - - :param edph: the ed25519ph state being updated - :type edph: crypto_sign_ed25519ph_state - :param pmsg: the partial message - :type pmsg: bytes - :rtype: None - """ - ensure( - isinstance(edph, crypto_sign_ed25519ph_state), - "edph parameter must be a ed25519ph_state object", - raising=exc.TypeError, - ) - ensure( - isinstance(pmsg, bytes), - "pmsg parameter must be a bytes object", - raising=exc.TypeError, - ) - rc = lib.crypto_sign_ed25519ph_update(edph.state, pmsg, len(pmsg)) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - -def crypto_sign_ed25519ph_final_create( - edph: crypto_sign_ed25519ph_state, sk: bytes -) -> bytes: - """ - Create a signature for the data hashed in edph - using the secret key sk - - :param edph: the ed25519ph state for the data - being signed - :type edph: crypto_sign_ed25519ph_state - :param sk: the ed25519 secret part of the signing key - :type sk: bytes - :return: ed25519ph signature - :rtype: bytes - """ - ensure( - isinstance(edph, crypto_sign_ed25519ph_state), - "edph parameter must be a ed25519ph_state object", - raising=exc.TypeError, - ) - ensure( - isinstance(sk, bytes), - "secret key parameter must be a bytes object", - raising=exc.TypeError, - ) - ensure( - len(sk) == crypto_sign_SECRETKEYBYTES, - ("secret key must be {} bytes long").format( - crypto_sign_SECRETKEYBYTES - ), - raising=exc.TypeError, - ) - signature = ffi.new("unsigned char[]", crypto_sign_BYTES) - rc = lib.crypto_sign_ed25519ph_final_create( - edph.state, signature, ffi.NULL, sk - ) - ensure(rc == 0, "Unexpected library error", raising=exc.RuntimeError) - - return ffi.buffer(signature, crypto_sign_BYTES)[:] - - -def crypto_sign_ed25519ph_final_verify( - edph: crypto_sign_ed25519ph_state, signature: bytes, pk: bytes -) -> bool: - """ - Verify a prehashed signature using the public key pk - - :param edph: the ed25519ph state for the data - being verified - :type edph: crypto_sign_ed25519ph_state - :param signature: the signature being verified - :type signature: bytes - :param pk: the ed25519 public part of the signing key - :type pk: bytes - :return: True if the signature is valid - :rtype: boolean - :raises exc.BadSignatureError: if the signature is not valid - """ - ensure( - isinstance(edph, crypto_sign_ed25519ph_state), - "edph parameter must be a ed25519ph_state object", - raising=exc.TypeError, - ) - ensure( - isinstance(signature, bytes), - "signature parameter must be a bytes object", - raising=exc.TypeError, - ) - ensure( - len(signature) == crypto_sign_BYTES, - ("signature must be {} bytes long").format(crypto_sign_BYTES), - raising=exc.TypeError, - ) - ensure( - isinstance(pk, bytes), - "public key parameter must be a bytes object", - raising=exc.TypeError, - ) - ensure( - len(pk) == crypto_sign_PUBLICKEYBYTES, - ("public key must be {} bytes long").format( - crypto_sign_PUBLICKEYBYTES - ), - raising=exc.TypeError, - ) - rc = lib.crypto_sign_ed25519ph_final_verify(edph.state, signature, pk) - if rc != 0: - raise exc.BadSignatureError("Signature was forged or corrupt") - - return True diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/randombytes.py b/venv/lib/python3.10/site-packages/nacl/bindings/randombytes.py deleted file mode 100644 index ed76deb5f978b549e312957aa84588b25e32176f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/randombytes.py +++ /dev/null @@ -1,51 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 nacl import exceptions as exc -from nacl._sodium import ffi, lib - -randombytes_SEEDBYTES: int = lib.randombytes_seedbytes() - - -def randombytes(size: int) -> bytes: - """ - Returns ``size`` number of random bytes from a cryptographically secure - random source. - - :param size: int - :rtype: bytes - """ - buf = ffi.new("unsigned char[]", size) - lib.randombytes(buf, size) - return ffi.buffer(buf, size)[:] - - -def randombytes_buf_deterministic(size: int, seed: bytes) -> bytes: - """ - Returns ``size`` number of deterministically generated pseudorandom bytes - from a seed - - :param size: int - :param seed: bytes - :rtype: bytes - """ - if len(seed) != randombytes_SEEDBYTES: - raise exc.TypeError( - "Deterministic random bytes must be generated from 32 bytes" - ) - - buf = ffi.new("unsigned char[]", size) - lib.randombytes_buf_deterministic(buf, size, seed) - return ffi.buffer(buf, size)[:] diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/sodium_core.py b/venv/lib/python3.10/site-packages/nacl/bindings/sodium_core.py deleted file mode 100644 index 7ebb84c82b2cd91f7b607de82082e93c7364b0c6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/sodium_core.py +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 nacl import exceptions as exc -from nacl._sodium import ffi, lib -from nacl.exceptions import ensure - - -def _sodium_init() -> None: - ensure( - lib.sodium_init() != -1, - "Could not initialize sodium", - raising=exc.RuntimeError, - ) - - -def sodium_init() -> None: - """ - Initializes sodium, picking the best implementations available for this - machine. - """ - ffi.init_once(_sodium_init, "libsodium") diff --git a/venv/lib/python3.10/site-packages/nacl/bindings/utils.py b/venv/lib/python3.10/site-packages/nacl/bindings/utils.py deleted file mode 100644 index 0ff22e34fc05b37cf4b1b667154949b519d72056..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/bindings/utils.py +++ /dev/null @@ -1,141 +0,0 @@ -# Copyright 2013-2017 Donald Stufft and individual contributors -# -# 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 nacl.exceptions as exc -from nacl._sodium import ffi, lib -from nacl.exceptions import ensure - - -def sodium_memcmp(inp1: bytes, inp2: bytes) -> bool: - """ - Compare contents of two memory regions in constant time - """ - ensure(isinstance(inp1, bytes), raising=exc.TypeError) - ensure(isinstance(inp2, bytes), raising=exc.TypeError) - - ln = max(len(inp1), len(inp2)) - - buf1 = ffi.new("char []", ln) - buf2 = ffi.new("char []", ln) - - ffi.memmove(buf1, inp1, len(inp1)) - ffi.memmove(buf2, inp2, len(inp2)) - - eqL = len(inp1) == len(inp2) - eqC = lib.sodium_memcmp(buf1, buf2, ln) == 0 - - return eqL and eqC - - -def sodium_pad(s: bytes, blocksize: int) -> bytes: - """ - Pad the input bytearray ``s`` to a multiple of ``blocksize`` - using the ISO/IEC 7816-4 algorithm - - :param s: input bytes string - :type s: bytes - :param blocksize: - :type blocksize: int - :return: padded string - :rtype: bytes - """ - ensure(isinstance(s, bytes), raising=exc.TypeError) - ensure(isinstance(blocksize, int), raising=exc.TypeError) - if blocksize <= 0: - raise exc.ValueError - s_len = len(s) - m_len = s_len + blocksize - buf = ffi.new("unsigned char []", m_len) - p_len = ffi.new("size_t []", 1) - ffi.memmove(buf, s, s_len) - rc = lib.sodium_pad(p_len, buf, s_len, blocksize, m_len) - ensure(rc == 0, "Padding failure", raising=exc.CryptoError) - return ffi.buffer(buf, p_len[0])[:] - - -def sodium_unpad(s: bytes, blocksize: int) -> bytes: - """ - Remove ISO/IEC 7816-4 padding from the input byte array ``s`` - - :param s: input bytes string - :type s: bytes - :param blocksize: - :type blocksize: int - :return: unpadded string - :rtype: bytes - """ - ensure(isinstance(s, bytes), raising=exc.TypeError) - ensure(isinstance(blocksize, int), raising=exc.TypeError) - s_len = len(s) - u_len = ffi.new("size_t []", 1) - rc = lib.sodium_unpad(u_len, s, s_len, blocksize) - if rc != 0: - raise exc.CryptoError("Unpadding failure") - return s[: u_len[0]] - - -def sodium_increment(inp: bytes) -> bytes: - """ - Increment the value of a byte-sequence interpreted - as the little-endian representation of a unsigned big integer. - - :param inp: input bytes buffer - :type inp: bytes - :return: a byte-sequence representing, as a little-endian - unsigned big integer, the value ``to_int(inp)`` - incremented by one. - :rtype: bytes - - """ - ensure(isinstance(inp, bytes), raising=exc.TypeError) - - ln = len(inp) - buf = ffi.new("unsigned char []", ln) - - ffi.memmove(buf, inp, ln) - - lib.sodium_increment(buf, ln) - - return ffi.buffer(buf, ln)[:] - - -def sodium_add(a: bytes, b: bytes) -> bytes: - """ - Given a couple of *same-sized* byte sequences, interpreted as the - little-endian representation of two unsigned integers, compute - the modular addition of the represented values, in constant time for - a given common length of the byte sequences. - - :param a: input bytes buffer - :type a: bytes - :param b: input bytes buffer - :type b: bytes - :return: a byte-sequence representing, as a little-endian big integer, - the integer value of ``(to_int(a) + to_int(b)) mod 2^(8*len(a))`` - :rtype: bytes - """ - ensure(isinstance(a, bytes), raising=exc.TypeError) - ensure(isinstance(b, bytes), raising=exc.TypeError) - ln = len(a) - ensure(len(b) == ln, raising=exc.TypeError) - - buf_a = ffi.new("unsigned char []", ln) - buf_b = ffi.new("unsigned char []", ln) - - ffi.memmove(buf_a, a, ln) - ffi.memmove(buf_b, b, ln) - - lib.sodium_add(buf_a, buf_b, ln) - - return ffi.buffer(buf_a, ln)[:] diff --git a/venv/lib/python3.10/site-packages/nacl/encoding.py b/venv/lib/python3.10/site-packages/nacl/encoding.py deleted file mode 100644 index 6740cfb31c2d5e9dfd7ea50cc29560d15b45301f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/encoding.py +++ /dev/null @@ -1,105 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 base64 -import binascii -from abc import ABCMeta, abstractmethod -from typing import SupportsBytes, Type - - -# TODO: when the minimum supported version of Python is 3.8, we can import -# Protocol from typing, and replace Encoder with a Protocol instead. -class _Encoder(metaclass=ABCMeta): - @staticmethod - @abstractmethod - def encode(data: bytes) -> bytes: - """Transform raw data to encoded data.""" - - @staticmethod - @abstractmethod - def decode(data: bytes) -> bytes: - """Transform encoded data back to raw data. - - Decoding after encoding should be a no-op, i.e. `decode(encode(x)) == x`. - """ - - -# Functions that use encoders are passed a subclass of _Encoder, not an instance -# (because the methods are all static). Let's gloss over that detail by defining -# an alias for Type[_Encoder]. -Encoder = Type[_Encoder] - - -class RawEncoder(_Encoder): - @staticmethod - def encode(data: bytes) -> bytes: - return data - - @staticmethod - def decode(data: bytes) -> bytes: - return data - - -class HexEncoder(_Encoder): - @staticmethod - def encode(data: bytes) -> bytes: - return binascii.hexlify(data) - - @staticmethod - def decode(data: bytes) -> bytes: - return binascii.unhexlify(data) - - -class Base16Encoder(_Encoder): - @staticmethod - def encode(data: bytes) -> bytes: - return base64.b16encode(data) - - @staticmethod - def decode(data: bytes) -> bytes: - return base64.b16decode(data) - - -class Base32Encoder(_Encoder): - @staticmethod - def encode(data: bytes) -> bytes: - return base64.b32encode(data) - - @staticmethod - def decode(data: bytes) -> bytes: - return base64.b32decode(data) - - -class Base64Encoder(_Encoder): - @staticmethod - def encode(data: bytes) -> bytes: - return base64.b64encode(data) - - @staticmethod - def decode(data: bytes) -> bytes: - return base64.b64decode(data) - - -class URLSafeBase64Encoder(_Encoder): - @staticmethod - def encode(data: bytes) -> bytes: - return base64.urlsafe_b64encode(data) - - @staticmethod - def decode(data: bytes) -> bytes: - return base64.urlsafe_b64decode(data) - - -class Encodable: - def encode(self: SupportsBytes, encoder: Encoder = RawEncoder) -> bytes: - return encoder.encode(bytes(self)) diff --git a/venv/lib/python3.10/site-packages/nacl/exceptions.py b/venv/lib/python3.10/site-packages/nacl/exceptions.py deleted file mode 100644 index e321df7bde364dee597f68b89ed99afc86369422..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/exceptions.py +++ /dev/null @@ -1,88 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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. - - -# We create a clone of various builtin Exception types which additionally -# inherit from CryptoError. Below, we refer to the parent types via the -# `builtins` namespace, so mypy can distinguish between (e.g.) -# `nacl.exceptions.RuntimeError` and `builtins.RuntimeError`. -import builtins -from typing import Type - - -class CryptoError(Exception): - """ - Base exception for all nacl related errors - """ - - -class BadSignatureError(CryptoError): - """ - Raised when the signature was forged or otherwise corrupt. - """ - - -class RuntimeError(builtins.RuntimeError, CryptoError): - pass - - -class AssertionError(builtins.AssertionError, CryptoError): - pass - - -class TypeError(builtins.TypeError, CryptoError): - pass - - -class ValueError(builtins.ValueError, CryptoError): - pass - - -class InvalidkeyError(CryptoError): - pass - - -class CryptPrefixError(InvalidkeyError): - pass - - -class UnavailableError(RuntimeError): - """ - is a subclass of :class:`~nacl.exceptions.RuntimeError`, raised when - trying to call functions not available in a minimal build of - libsodium. - """ - - pass - - -def ensure(cond: bool, *args: object, **kwds: Type[Exception]) -> None: - """ - Return if a condition is true, otherwise raise a caller-configurable - :py:class:`Exception` - :param bool cond: the condition to be checked - :param sequence args: the arguments to be passed to the exception's - constructor - The only accepted named parameter is `raising` used to configure the - exception to be raised if `cond` is not `True` - """ - _CHK_UNEXP = "check_condition() got an unexpected keyword argument {0}" - - raising = kwds.pop("raising", AssertionError) - if kwds: - raise TypeError(_CHK_UNEXP.format(repr(kwds.popitem()[0]))) - - if cond is True: - return - raise raising(*args) diff --git a/venv/lib/python3.10/site-packages/nacl/hash.py b/venv/lib/python3.10/site-packages/nacl/hash.py deleted file mode 100644 index eb8cff6b156d60bb89757eb8bd8bb6bca2805680..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/hash.py +++ /dev/null @@ -1,182 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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. -""" -The :mod:`nacl.hash` module exposes one-shot interfaces -for libsodium selected hash primitives and the constants needed -for their usage. -""" - - -import nacl.bindings -import nacl.encoding - - -BLAKE2B_BYTES = nacl.bindings.crypto_generichash_BYTES -"""Default digest size for :func:`blake2b` hash""" -BLAKE2B_BYTES_MIN = nacl.bindings.crypto_generichash_BYTES_MIN -"""Minimum allowed digest size for :func:`blake2b` hash""" -BLAKE2B_BYTES_MAX = nacl.bindings.crypto_generichash_BYTES_MAX -"""Maximum allowed digest size for :func:`blake2b` hash""" -BLAKE2B_KEYBYTES = nacl.bindings.crypto_generichash_KEYBYTES -"""Default size of the ``key`` byte array for :func:`blake2b` hash""" -BLAKE2B_KEYBYTES_MIN = nacl.bindings.crypto_generichash_KEYBYTES_MIN -"""Minimum allowed size of the ``key`` byte array for :func:`blake2b` hash""" -BLAKE2B_KEYBYTES_MAX = nacl.bindings.crypto_generichash_KEYBYTES_MAX -"""Maximum allowed size of the ``key`` byte array for :func:`blake2b` hash""" -BLAKE2B_SALTBYTES = nacl.bindings.crypto_generichash_SALTBYTES -"""Maximum allowed length of the ``salt`` byte array for -:func:`blake2b` hash""" -BLAKE2B_PERSONALBYTES = nacl.bindings.crypto_generichash_PERSONALBYTES -"""Maximum allowed length of the ``personalization`` -byte array for :func:`blake2b` hash""" - -SIPHASH_BYTES = nacl.bindings.crypto_shorthash_siphash24_BYTES -"""Size of the :func:`siphash24` digest""" -SIPHASH_KEYBYTES = nacl.bindings.crypto_shorthash_siphash24_KEYBYTES -"""Size of the secret ``key`` used by the :func:`siphash24` MAC""" - -SIPHASHX_AVAILABLE = nacl.bindings.has_crypto_shorthash_siphashx24 -"""``True`` if :func:`siphashx24` is available to be called""" - -SIPHASHX_BYTES = nacl.bindings.crypto_shorthash_siphashx24_BYTES -"""Size of the :func:`siphashx24` digest""" -SIPHASHX_KEYBYTES = nacl.bindings.crypto_shorthash_siphashx24_KEYBYTES -"""Size of the secret ``key`` used by the :func:`siphashx24` MAC""" - -_b2b_hash = nacl.bindings.crypto_generichash_blake2b_salt_personal -_sip_hash = nacl.bindings.crypto_shorthash_siphash24 -_sip_hashx = nacl.bindings.crypto_shorthash_siphashx24 - - -def sha256( - message: bytes, encoder: nacl.encoding.Encoder = nacl.encoding.HexEncoder -) -> bytes: - """ - Hashes ``message`` with SHA256. - - :param message: The message to hash. - :type message: bytes - :param encoder: A class that is able to encode the hashed message. - :returns: The hashed message. - :rtype: bytes - """ - return encoder.encode(nacl.bindings.crypto_hash_sha256(message)) - - -def sha512( - message: bytes, encoder: nacl.encoding.Encoder = nacl.encoding.HexEncoder -) -> bytes: - """ - Hashes ``message`` with SHA512. - - :param message: The message to hash. - :type message: bytes - :param encoder: A class that is able to encode the hashed message. - :returns: The hashed message. - :rtype: bytes - """ - return encoder.encode(nacl.bindings.crypto_hash_sha512(message)) - - -def blake2b( - data: bytes, - digest_size: int = BLAKE2B_BYTES, - key: bytes = b"", - salt: bytes = b"", - person: bytes = b"", - encoder: nacl.encoding.Encoder = nacl.encoding.HexEncoder, -) -> bytes: - """ - Hashes ``data`` with blake2b. - - :param data: the digest input byte sequence - :type data: bytes - :param digest_size: the requested digest size; must be at most - :const:`BLAKE2B_BYTES_MAX`; - the default digest size is - :const:`BLAKE2B_BYTES` - :type digest_size: int - :param key: the key to be set for keyed MAC/PRF usage; if set, the key - must be at most :data:`~nacl.hash.BLAKE2B_KEYBYTES_MAX` long - :type key: bytes - :param salt: an initialization salt at most - :const:`BLAKE2B_SALTBYTES` long; - it will be zero-padded if needed - :type salt: bytes - :param person: a personalization string at most - :const:`BLAKE2B_PERSONALBYTES` long; - it will be zero-padded if needed - :type person: bytes - :param encoder: the encoder to use on returned digest - :type encoder: class - :returns: The hashed message. - :rtype: bytes - """ - - digest = _b2b_hash( - data, digest_size=digest_size, key=key, salt=salt, person=person - ) - return encoder.encode(digest) - - -generichash = blake2b - - -def siphash24( - message: bytes, - key: bytes = b"", - encoder: nacl.encoding.Encoder = nacl.encoding.HexEncoder, -) -> bytes: - """ - Computes a keyed MAC of ``message`` using the short-input-optimized - siphash-2-4 construction. - - :param message: The message to hash. - :type message: bytes - :param key: the message authentication key for the siphash MAC construct - :type key: bytes(:const:`SIPHASH_KEYBYTES`) - :param encoder: A class that is able to encode the hashed message. - :returns: The hashed message. - :rtype: bytes(:const:`SIPHASH_BYTES`) - """ - digest = _sip_hash(message, key) - return encoder.encode(digest) - - -shorthash = siphash24 - - -def siphashx24( - message: bytes, - key: bytes = b"", - encoder: nacl.encoding.Encoder = nacl.encoding.HexEncoder, -) -> bytes: - """ - Computes a keyed MAC of ``message`` using the 128 bit variant of the - siphash-2-4 construction. - - :param message: The message to hash. - :type message: bytes - :param key: the message authentication key for the siphash MAC construct - :type key: bytes(:const:`SIPHASHX_KEYBYTES`) - :param encoder: A class that is able to encode the hashed message. - :returns: The hashed message. - :rtype: bytes(:const:`SIPHASHX_BYTES`) - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - - .. versionadded:: 1.2 - """ - digest = _sip_hashx(message, key) - return encoder.encode(digest) diff --git a/venv/lib/python3.10/site-packages/nacl/hashlib.py b/venv/lib/python3.10/site-packages/nacl/hashlib.py deleted file mode 100644 index 4fd4a976b4d285f6d92e6133e3d69c96c282fcea..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/hashlib.py +++ /dev/null @@ -1,143 +0,0 @@ -# Copyright 2016-2019 Donald Stufft and individual contributors -# -# 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 binascii -from typing import NoReturn - -import nacl.bindings -from nacl.utils import bytes_as_string - -BYTES = nacl.bindings.crypto_generichash_BYTES -BYTES_MIN = nacl.bindings.crypto_generichash_BYTES_MIN -BYTES_MAX = nacl.bindings.crypto_generichash_BYTES_MAX -KEYBYTES = nacl.bindings.crypto_generichash_KEYBYTES -KEYBYTES_MIN = nacl.bindings.crypto_generichash_KEYBYTES_MIN -KEYBYTES_MAX = nacl.bindings.crypto_generichash_KEYBYTES_MAX -SALTBYTES = nacl.bindings.crypto_generichash_SALTBYTES -PERSONALBYTES = nacl.bindings.crypto_generichash_PERSONALBYTES - -SCRYPT_AVAILABLE = nacl.bindings.has_crypto_pwhash_scryptsalsa208sha256 - -_b2b_init = nacl.bindings.crypto_generichash_blake2b_init -_b2b_final = nacl.bindings.crypto_generichash_blake2b_final -_b2b_update = nacl.bindings.crypto_generichash_blake2b_update - - -class blake2b: - """ - :py:mod:`hashlib` API compatible blake2b algorithm implementation - """ - - MAX_DIGEST_SIZE = BYTES - MAX_KEY_SIZE = KEYBYTES_MAX - PERSON_SIZE = PERSONALBYTES - SALT_SIZE = SALTBYTES - - def __init__( - self, - data: bytes = b"", - digest_size: int = BYTES, - key: bytes = b"", - salt: bytes = b"", - person: bytes = b"", - ): - """ - :py:class:`.blake2b` algorithm initializer - - :param data: - :type data: bytes - :param int digest_size: the requested digest size; must be - at most :py:attr:`.MAX_DIGEST_SIZE`; - the default digest size is :py:data:`.BYTES` - :param key: the key to be set for keyed MAC/PRF usage; if set, - the key must be at most :py:data:`.KEYBYTES_MAX` long - :type key: bytes - :param salt: a initialization salt at most - :py:attr:`.SALT_SIZE` long; it will be zero-padded - if needed - :type salt: bytes - :param person: a personalization string at most - :py:attr:`.PERSONAL_SIZE` long; it will be zero-padded - if needed - :type person: bytes - """ - - self._state = _b2b_init( - key=key, salt=salt, person=person, digest_size=digest_size - ) - self._digest_size = digest_size - - if data: - self.update(data) - - @property - def digest_size(self) -> int: - return self._digest_size - - @property - def block_size(self) -> int: - return 128 - - @property - def name(self) -> str: - return "blake2b" - - def update(self, data: bytes) -> None: - _b2b_update(self._state, data) - - def digest(self) -> bytes: - _st = self._state.copy() - return _b2b_final(_st) - - def hexdigest(self) -> str: - return bytes_as_string(binascii.hexlify(self.digest())) - - def copy(self) -> "blake2b": - _cp = type(self)(digest_size=self.digest_size) - _st = self._state.copy() - _cp._state = _st - return _cp - - def __reduce__(self) -> NoReturn: - """ - Raise the same exception as hashlib's blake implementation - on copy.copy() - """ - raise TypeError( - "can't pickle {} objects".format(self.__class__.__name__) - ) - - -def scrypt( - password: bytes, - salt: bytes = b"", - n: int = 2 ** 20, - r: int = 8, - p: int = 1, - maxmem: int = 2 ** 25, - dklen: int = 64, -) -> bytes: - """ - Derive a cryptographic key using the scrypt KDF. - - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - - Implements the same signature as the ``hashlib.scrypt`` implemented - in cpython version 3.6 - """ - return nacl.bindings.crypto_pwhash_scryptsalsa208sha256_ll( - password, salt, n, r, p, maxmem=maxmem, dklen=dklen - ) diff --git a/venv/lib/python3.10/site-packages/nacl/public.py b/venv/lib/python3.10/site-packages/nacl/public.py deleted file mode 100644 index be9410fce0a93c7060a7b5c8e459d9faefcfb502..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/public.py +++ /dev/null @@ -1,423 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 ClassVar, Generic, Optional, Type, TypeVar - -import nacl.bindings -from nacl import encoding -from nacl import exceptions as exc -from nacl.encoding import Encoder -from nacl.utils import EncryptedMessage, StringFixer, random - - -class PublicKey(encoding.Encodable, StringFixer): - """ - The public key counterpart to an Curve25519 :class:`nacl.public.PrivateKey` - for encrypting messages. - - :param public_key: [:class:`bytes`] Encoded Curve25519 public key - :param encoder: A class that is able to decode the `public_key` - - :cvar SIZE: The size that the public key is required to be - """ - - SIZE: ClassVar[int] = nacl.bindings.crypto_box_PUBLICKEYBYTES - - def __init__( - self, - public_key: bytes, - encoder: encoding.Encoder = encoding.RawEncoder, - ): - self._public_key = encoder.decode(public_key) - if not isinstance(self._public_key, bytes): - raise exc.TypeError("PublicKey must be created from 32 bytes") - - if len(self._public_key) != self.SIZE: - raise exc.ValueError( - "The public key must be exactly {} bytes long".format( - self.SIZE - ) - ) - - def __bytes__(self) -> bytes: - return self._public_key - - def __hash__(self) -> int: - return hash(bytes(self)) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, self.__class__): - return False - return nacl.bindings.sodium_memcmp(bytes(self), bytes(other)) - - def __ne__(self, other: object) -> bool: - return not (self == other) - - -class PrivateKey(encoding.Encodable, StringFixer): - """ - Private key for decrypting messages using the Curve25519 algorithm. - - .. warning:: This **must** be protected and remain secret. Anyone who - knows the value of your :class:`~nacl.public.PrivateKey` can decrypt - any message encrypted by the corresponding - :class:`~nacl.public.PublicKey` - - :param private_key: The private key used to decrypt messages - :param encoder: The encoder class used to decode the given keys - - :cvar SIZE: The size that the private key is required to be - :cvar SEED_SIZE: The size that the seed used to generate the - private key is required to be - """ - - SIZE: ClassVar[int] = nacl.bindings.crypto_box_SECRETKEYBYTES - SEED_SIZE: ClassVar[int] = nacl.bindings.crypto_box_SEEDBYTES - - def __init__( - self, - private_key: bytes, - encoder: encoding.Encoder = encoding.RawEncoder, - ): - # Decode the secret_key - private_key = encoder.decode(private_key) - # verify the given secret key type and size are correct - if not ( - isinstance(private_key, bytes) and len(private_key) == self.SIZE - ): - raise exc.TypeError( - ( - "PrivateKey must be created from a {} " - "bytes long raw secret key" - ).format(self.SIZE) - ) - - raw_public_key = nacl.bindings.crypto_scalarmult_base(private_key) - - self._private_key = private_key - self.public_key = PublicKey(raw_public_key) - - @classmethod - def from_seed( - cls, - seed: bytes, - encoder: encoding.Encoder = encoding.RawEncoder, - ) -> "PrivateKey": - """ - Generate a PrivateKey using a deterministic construction - starting from a caller-provided seed - - .. warning:: The seed **must** be high-entropy; therefore, - its generator **must** be a cryptographic quality - random function like, for example, :func:`~nacl.utils.random`. - - .. warning:: The seed **must** be protected and remain secret. - Anyone who knows the seed is really in possession of - the corresponding PrivateKey. - - :param seed: The seed used to generate the private key - :rtype: :class:`~nacl.public.PrivateKey` - """ - # decode the seed - seed = encoder.decode(seed) - # Verify the given seed type and size are correct - if not (isinstance(seed, bytes) and len(seed) == cls.SEED_SIZE): - raise exc.TypeError( - ( - "PrivateKey seed must be a {} bytes long " - "binary sequence" - ).format(cls.SEED_SIZE) - ) - # generate a raw keypair from the given seed - raw_pk, raw_sk = nacl.bindings.crypto_box_seed_keypair(seed) - # construct a instance from the raw secret key - return cls(raw_sk) - - def __bytes__(self) -> bytes: - return self._private_key - - def __hash__(self) -> int: - return hash((type(self), bytes(self.public_key))) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, self.__class__): - return False - return self.public_key == other.public_key - - def __ne__(self, other: object) -> bool: - return not (self == other) - - @classmethod - def generate(cls) -> "PrivateKey": - """ - Generates a random :class:`~nacl.public.PrivateKey` object - - :rtype: :class:`~nacl.public.PrivateKey` - """ - return cls(random(PrivateKey.SIZE), encoder=encoding.RawEncoder) - - -_Box = TypeVar("_Box", bound="Box") - - -class Box(encoding.Encodable, StringFixer): - """ - The Box class boxes and unboxes messages between a pair of keys - - The ciphertexts generated by :class:`~nacl.public.Box` include a 16 - byte authenticator which is checked as part of the decryption. An invalid - authenticator will cause the decrypt function to raise an exception. The - authenticator is not a signature. Once you've decrypted the message you've - demonstrated the ability to create arbitrary valid message, so messages you - send are repudiable. For non-repudiable messages, sign them after - encryption. - - :param private_key: :class:`~nacl.public.PrivateKey` used to encrypt and - decrypt messages - :param public_key: :class:`~nacl.public.PublicKey` used to encrypt and - decrypt messages - - :cvar NONCE_SIZE: The size that the nonce is required to be. - """ - - NONCE_SIZE: ClassVar[int] = nacl.bindings.crypto_box_NONCEBYTES - _shared_key: bytes - - def __init__(self, private_key: PrivateKey, public_key: PublicKey): - if not isinstance(private_key, PrivateKey) or not isinstance( - public_key, PublicKey - ): - raise exc.TypeError( - "Box must be created from a PrivateKey and a PublicKey" - ) - self._shared_key = nacl.bindings.crypto_box_beforenm( - public_key.encode(encoder=encoding.RawEncoder), - private_key.encode(encoder=encoding.RawEncoder), - ) - - def __bytes__(self) -> bytes: - return self._shared_key - - @classmethod - def decode( - cls: Type[_Box], encoded: bytes, encoder: Encoder = encoding.RawEncoder - ) -> _Box: - """ - Alternative constructor. Creates a Box from an existing Box's shared key. - """ - # Create an empty box - box: _Box = cls.__new__(cls) - - # Assign our decoded value to the shared key of the box - box._shared_key = encoder.decode(encoded) - - return box - - def encrypt( - self, - plaintext: bytes, - nonce: Optional[bytes] = None, - encoder: encoding.Encoder = encoding.RawEncoder, - ) -> EncryptedMessage: - """ - Encrypts the plaintext message using the given `nonce` (or generates - one randomly if omitted) and returns the ciphertext encoded with the - encoder. - - .. warning:: It is **VITALLY** important that the nonce is a nonce, - i.e. it is a number used only once for any given key. If you fail - to do this, you compromise the privacy of the messages encrypted. - - :param plaintext: [:class:`bytes`] The plaintext message to encrypt - :param nonce: [:class:`bytes`] The nonce to use in the encryption - :param encoder: The encoder to use to encode the ciphertext - :rtype: [:class:`nacl.utils.EncryptedMessage`] - """ - if nonce is None: - nonce = random(self.NONCE_SIZE) - - if len(nonce) != self.NONCE_SIZE: - raise exc.ValueError( - "The nonce must be exactly %s bytes long" % self.NONCE_SIZE - ) - - ciphertext = nacl.bindings.crypto_box_afternm( - plaintext, - nonce, - self._shared_key, - ) - - encoded_nonce = encoder.encode(nonce) - encoded_ciphertext = encoder.encode(ciphertext) - - return EncryptedMessage._from_parts( - encoded_nonce, - encoded_ciphertext, - encoder.encode(nonce + ciphertext), - ) - - def decrypt( - self, - ciphertext: bytes, - nonce: Optional[bytes] = None, - encoder: encoding.Encoder = encoding.RawEncoder, - ) -> bytes: - """ - Decrypts the ciphertext using the `nonce` (explicitly, when passed as a - parameter or implicitly, when omitted, as part of the ciphertext) and - returns the plaintext message. - - :param ciphertext: [:class:`bytes`] The encrypted message to decrypt - :param nonce: [:class:`bytes`] The nonce used when encrypting the - ciphertext - :param encoder: The encoder used to decode the ciphertext. - :rtype: [:class:`bytes`] - """ - # Decode our ciphertext - ciphertext = encoder.decode(ciphertext) - - if nonce is None: - # If we were given the nonce and ciphertext combined, split them. - nonce = ciphertext[: self.NONCE_SIZE] - ciphertext = ciphertext[self.NONCE_SIZE :] - - if len(nonce) != self.NONCE_SIZE: - raise exc.ValueError( - "The nonce must be exactly %s bytes long" % self.NONCE_SIZE - ) - - plaintext = nacl.bindings.crypto_box_open_afternm( - ciphertext, - nonce, - self._shared_key, - ) - - return plaintext - - def shared_key(self) -> bytes: - """ - Returns the Curve25519 shared secret, that can then be used as a key in - other symmetric ciphers. - - .. warning:: It is **VITALLY** important that you use a nonce with your - symmetric cipher. If you fail to do this, you compromise the - privacy of the messages encrypted. Ensure that the key length of - your cipher is 32 bytes. - :rtype: [:class:`bytes`] - """ - - return self._shared_key - - -_Key = TypeVar("_Key", PublicKey, PrivateKey) - - -class SealedBox(Generic[_Key], encoding.Encodable, StringFixer): - """ - The SealedBox class boxes and unboxes messages addressed to - a specified key-pair by using ephemeral sender's keypairs, - whose private part will be discarded just after encrypting - a single plaintext message. - - The ciphertexts generated by :class:`~nacl.public.SecretBox` include - the public part of the ephemeral key before the :class:`~nacl.public.Box` - ciphertext. - - :param recipient_key: a :class:`~nacl.public.PublicKey` used to encrypt - messages and derive nonces, or a :class:`~nacl.public.PrivateKey` used - to decrypt messages. - - .. versionadded:: 1.2 - """ - - _public_key: bytes - _private_key: Optional[bytes] - - def __init__(self, recipient_key: _Key): - if isinstance(recipient_key, PublicKey): - self._public_key = recipient_key.encode( - encoder=encoding.RawEncoder - ) - self._private_key = None - elif isinstance(recipient_key, PrivateKey): - self._private_key = recipient_key.encode( - encoder=encoding.RawEncoder - ) - self._public_key = recipient_key.public_key.encode( - encoder=encoding.RawEncoder - ) - else: - raise exc.TypeError( - "SealedBox must be created from a PublicKey or a PrivateKey" - ) - - def __bytes__(self) -> bytes: - return self._public_key - - def encrypt( - self, - plaintext: bytes, - encoder: encoding.Encoder = encoding.RawEncoder, - ) -> bytes: - """ - Encrypts the plaintext message using a random-generated ephemeral - keypair and returns a "composed ciphertext", containing both - the public part of the keypair and the ciphertext proper, - encoded with the encoder. - - The private part of the ephemeral key-pair will be scrubbed before - returning the ciphertext, therefore, the sender will not be able to - decrypt the generated ciphertext. - - :param plaintext: [:class:`bytes`] The plaintext message to encrypt - :param encoder: The encoder to use to encode the ciphertext - :return bytes: encoded ciphertext - """ - - ciphertext = nacl.bindings.crypto_box_seal(plaintext, self._public_key) - - encoded_ciphertext = encoder.encode(ciphertext) - - return encoded_ciphertext - - def decrypt( - self: "SealedBox[PrivateKey]", - ciphertext: bytes, - encoder: encoding.Encoder = encoding.RawEncoder, - ) -> bytes: - """ - Decrypts the ciphertext using the ephemeral public key enclosed - in the ciphertext and the SealedBox private key, returning - the plaintext message. - - :param ciphertext: [:class:`bytes`] The encrypted message to decrypt - :param encoder: The encoder used to decode the ciphertext. - :return bytes: The original plaintext - :raises TypeError: if this SealedBox was created with a - :class:`~nacl.public.PublicKey` rather than a - :class:`~nacl.public.PrivateKey`. - """ - # Decode our ciphertext - ciphertext = encoder.decode(ciphertext) - - if self._private_key is None: - raise TypeError( - "SealedBoxes created with a public key cannot decrypt" - ) - plaintext = nacl.bindings.crypto_box_seal_open( - ciphertext, - self._public_key, - self._private_key, - ) - - return plaintext diff --git a/venv/lib/python3.10/site-packages/nacl/pwhash/__init__.py b/venv/lib/python3.10/site-packages/nacl/pwhash/__init__.py deleted file mode 100644 index ffd76a64c5ab9b7649f276c8d8a182ae76db9389..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/pwhash/__init__.py +++ /dev/null @@ -1,75 +0,0 @@ -# Copyright 2017 Donald Stufft and individual contributors -# -# 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 nacl.exceptions import CryptPrefixError - -from . import _argon2, argon2i, argon2id, scrypt - -STRPREFIX = argon2id.STRPREFIX - -PWHASH_SIZE = argon2id.PWHASH_SIZE - -assert _argon2.ALG_ARGON2_DEFAULT == _argon2.ALG_ARGON2ID13 -# since version 1.0.15 of libsodium - -PASSWD_MIN = argon2id.PASSWD_MIN -PASSWD_MAX = argon2id.PASSWD_MAX -MEMLIMIT_MAX = argon2id.MEMLIMIT_MAX -MEMLIMIT_MIN = argon2id.MEMLIMIT_MIN -OPSLIMIT_MAX = argon2id.OPSLIMIT_MAX -OPSLIMIT_MIN = argon2id.OPSLIMIT_MIN -OPSLIMIT_INTERACTIVE = argon2id.OPSLIMIT_INTERACTIVE -MEMLIMIT_INTERACTIVE = argon2id.MEMLIMIT_INTERACTIVE -OPSLIMIT_MODERATE = argon2id.OPSLIMIT_MODERATE -MEMLIMIT_MODERATE = argon2id.MEMLIMIT_MODERATE -OPSLIMIT_SENSITIVE = argon2id.OPSLIMIT_SENSITIVE -MEMLIMIT_SENSITIVE = argon2id.MEMLIMIT_SENSITIVE - -str = argon2id.str - -assert argon2i.ALG != argon2id.ALG - -SCRYPT_SALTBYTES = scrypt.SALTBYTES -SCRYPT_PWHASH_SIZE = scrypt.PWHASH_SIZE -SCRYPT_OPSLIMIT_INTERACTIVE = scrypt.OPSLIMIT_INTERACTIVE -SCRYPT_MEMLIMIT_INTERACTIVE = scrypt.MEMLIMIT_INTERACTIVE -SCRYPT_OPSLIMIT_SENSITIVE = scrypt.OPSLIMIT_SENSITIVE -SCRYPT_MEMLIMIT_SENSITIVE = scrypt.MEMLIMIT_SENSITIVE - - -kdf_scryptsalsa208sha256 = scrypt.kdf -scryptsalsa208sha256_str = scrypt.str -verify_scryptsalsa208sha256 = scrypt.verify - - -def verify(password_hash: bytes, password: bytes) -> bool: - """ - Takes a modular crypt encoded stored password hash derived using one - of the algorithms supported by `libsodium` and checks if the user provided - password will hash to the same string when using the parameters saved - in the stored hash - """ - if password_hash.startswith(argon2id.STRPREFIX): - return argon2id.verify(password_hash, password) - elif password_hash.startswith(argon2i.STRPREFIX): - return argon2id.verify(password_hash, password) - elif scrypt.AVAILABLE and password_hash.startswith(scrypt.STRPREFIX): - return scrypt.verify(password_hash, password) - else: - raise ( - CryptPrefixError( - "given password_hash is not in a supported format" - ) - ) diff --git a/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 841d8a046de8d82d8cad5c035bc310eb3c250645..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/_argon2.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/_argon2.cpython-310.pyc deleted file mode 100644 index 0b3735129714465a5c43f6e1ac5e126519ed4530..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/_argon2.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/argon2i.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/argon2i.cpython-310.pyc deleted file mode 100644 index 1831859e50c4b9313cc0fe78e62abee2e0c55b45..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/argon2i.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/argon2id.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/argon2id.cpython-310.pyc deleted file mode 100644 index 9de786045761eb03afc7810977d801f9f9a00878..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/argon2id.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/scrypt.cpython-310.pyc b/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/scrypt.cpython-310.pyc deleted file mode 100644 index 8be708b42f19622d14f8f56571f8c0552f58a5f7..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/nacl/pwhash/__pycache__/scrypt.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/nacl/pwhash/_argon2.py b/venv/lib/python3.10/site-packages/nacl/pwhash/_argon2.py deleted file mode 100644 index 856eda04ea6afe5fa41ae44519f1de1bf662960e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/pwhash/_argon2.py +++ /dev/null @@ -1,49 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 nacl.bindings - -_argon2_strbytes_plus_one = nacl.bindings.crypto_pwhash_STRBYTES - -PWHASH_SIZE = _argon2_strbytes_plus_one - 1 -SALTBYTES = nacl.bindings.crypto_pwhash_SALTBYTES - -PASSWD_MIN = nacl.bindings.crypto_pwhash_PASSWD_MIN -PASSWD_MAX = nacl.bindings.crypto_pwhash_PASSWD_MAX - -PWHASH_SIZE = _argon2_strbytes_plus_one - 1 - -BYTES_MAX = nacl.bindings.crypto_pwhash_BYTES_MAX -BYTES_MIN = nacl.bindings.crypto_pwhash_BYTES_MIN - -ALG_ARGON2I13 = nacl.bindings.crypto_pwhash_ALG_ARGON2I13 -ALG_ARGON2ID13 = nacl.bindings.crypto_pwhash_ALG_ARGON2ID13 -ALG_ARGON2_DEFAULT = nacl.bindings.crypto_pwhash_ALG_DEFAULT - - -def verify(password_hash: bytes, password: bytes) -> bool: - """ - Takes a modular crypt encoded argon2i or argon2id stored password hash - and checks if the user provided password will hash to the same string - when using the stored parameters - - :param password_hash: password hash serialized in modular crypt() format - :type password_hash: bytes - :param password: user provided password - :type password: bytes - :rtype: boolean - - .. versionadded:: 1.2 - """ - return nacl.bindings.crypto_pwhash_str_verify(password_hash, password) diff --git a/venv/lib/python3.10/site-packages/nacl/pwhash/argon2i.py b/venv/lib/python3.10/site-packages/nacl/pwhash/argon2i.py deleted file mode 100644 index f9b3af7f0a6499099cf50b223d78129bfb371142..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/pwhash/argon2i.py +++ /dev/null @@ -1,132 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 nacl.bindings -import nacl.encoding - -from . import _argon2 - -ALG = _argon2.ALG_ARGON2I13 -STRPREFIX = nacl.bindings.crypto_pwhash_argon2i_STRPREFIX - -SALTBYTES = _argon2.SALTBYTES - -PASSWD_MIN = _argon2.PASSWD_MIN -PASSWD_MAX = _argon2.PASSWD_MAX - -PWHASH_SIZE = _argon2.PWHASH_SIZE - -BYTES_MIN = _argon2.BYTES_MIN -BYTES_MAX = _argon2.BYTES_MAX - -verify = _argon2.verify - -MEMLIMIT_MAX = nacl.bindings.crypto_pwhash_argon2i_MEMLIMIT_MAX -MEMLIMIT_MIN = nacl.bindings.crypto_pwhash_argon2i_MEMLIMIT_MIN -OPSLIMIT_MAX = nacl.bindings.crypto_pwhash_argon2i_OPSLIMIT_MAX -OPSLIMIT_MIN = nacl.bindings.crypto_pwhash_argon2i_OPSLIMIT_MIN - -OPSLIMIT_INTERACTIVE = nacl.bindings.crypto_pwhash_argon2i_OPSLIMIT_INTERACTIVE -MEMLIMIT_INTERACTIVE = nacl.bindings.crypto_pwhash_argon2i_MEMLIMIT_INTERACTIVE -OPSLIMIT_SENSITIVE = nacl.bindings.crypto_pwhash_argon2i_OPSLIMIT_SENSITIVE -MEMLIMIT_SENSITIVE = nacl.bindings.crypto_pwhash_argon2i_MEMLIMIT_SENSITIVE - -OPSLIMIT_MODERATE = nacl.bindings.crypto_pwhash_argon2i_OPSLIMIT_MODERATE -MEMLIMIT_MODERATE = nacl.bindings.crypto_pwhash_argon2i_MEMLIMIT_MODERATE - - -def kdf( - size: int, - password: bytes, - salt: bytes, - opslimit: int = OPSLIMIT_SENSITIVE, - memlimit: int = MEMLIMIT_SENSITIVE, - encoder: nacl.encoding.Encoder = nacl.encoding.RawEncoder, -) -> bytes: - """ - Derive a ``size`` bytes long key from a caller-supplied - ``password`` and ``salt`` pair using the argon2i - memory-hard construct. - - the enclosing module provides the constants - - - :py:const:`.OPSLIMIT_INTERACTIVE` - - :py:const:`.MEMLIMIT_INTERACTIVE` - - :py:const:`.OPSLIMIT_MODERATE` - - :py:const:`.MEMLIMIT_MODERATE` - - :py:const:`.OPSLIMIT_SENSITIVE` - - :py:const:`.MEMLIMIT_SENSITIVE` - - as a guidance for correct settings. - - :param size: derived key size, must be between - :py:const:`.BYTES_MIN` and - :py:const:`.BYTES_MAX` - :type size: int - :param password: password used to seed the key derivation procedure; - it length must be between - :py:const:`.PASSWD_MIN` and - :py:const:`.PASSWD_MAX` - :type password: bytes - :param salt: **RANDOM** salt used in the key derivation procedure; - its length must be exactly :py:const:`.SALTBYTES` - :type salt: bytes - :param opslimit: the time component (operation count) - of the key derivation procedure's computational cost; - it must be between - :py:const:`.OPSLIMIT_MIN` and - :py:const:`.OPSLIMIT_MAX` - :type opslimit: int - :param memlimit: the memory occupation component - of the key derivation procedure's computational cost; - it must be between - :py:const:`.MEMLIMIT_MIN` and - :py:const:`.MEMLIMIT_MAX` - :type memlimit: int - :rtype: bytes - - .. versionadded:: 1.2 - """ - - return encoder.encode( - nacl.bindings.crypto_pwhash_alg( - size, password, salt, opslimit, memlimit, ALG - ) - ) - - -def str( - password: bytes, - opslimit: int = OPSLIMIT_INTERACTIVE, - memlimit: int = MEMLIMIT_INTERACTIVE, -) -> bytes: - """ - Hashes a password with a random salt, using the memory-hard - argon2i construct and returning an ascii string that has all - the needed info to check against a future password - - - The default settings for opslimit and memlimit are those deemed - correct for the interactive user login case. - - :param bytes password: - :param int opslimit: - :param int memlimit: - :rtype: bytes - - .. versionadded:: 1.2 - """ - return nacl.bindings.crypto_pwhash_str_alg( - password, opslimit, memlimit, ALG - ) diff --git a/venv/lib/python3.10/site-packages/nacl/pwhash/argon2id.py b/venv/lib/python3.10/site-packages/nacl/pwhash/argon2id.py deleted file mode 100644 index 1b86d69e141b453a1bafd50197103923ff400d69..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/pwhash/argon2id.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 nacl.bindings -import nacl.encoding - -from . import _argon2 - -ALG = _argon2.ALG_ARGON2ID13 -STRPREFIX = nacl.bindings.crypto_pwhash_argon2id_STRPREFIX - -SALTBYTES = _argon2.SALTBYTES - -PASSWD_MIN = _argon2.PASSWD_MIN -PASSWD_MAX = _argon2.PASSWD_MAX - -PWHASH_SIZE = _argon2.PWHASH_SIZE - -BYTES_MIN = _argon2.BYTES_MIN -BYTES_MAX = _argon2.BYTES_MAX - -verify = _argon2.verify - -MEMLIMIT_MIN = nacl.bindings.crypto_pwhash_argon2id_MEMLIMIT_MIN -MEMLIMIT_MAX = nacl.bindings.crypto_pwhash_argon2id_MEMLIMIT_MAX -OPSLIMIT_MIN = nacl.bindings.crypto_pwhash_argon2id_OPSLIMIT_MIN -OPSLIMIT_MAX = nacl.bindings.crypto_pwhash_argon2id_OPSLIMIT_MAX - -OPSLIMIT_INTERACTIVE = ( - nacl.bindings.crypto_pwhash_argon2id_OPSLIMIT_INTERACTIVE -) -MEMLIMIT_INTERACTIVE = ( - nacl.bindings.crypto_pwhash_argon2id_MEMLIMIT_INTERACTIVE -) -OPSLIMIT_SENSITIVE = nacl.bindings.crypto_pwhash_argon2id_OPSLIMIT_SENSITIVE -MEMLIMIT_SENSITIVE = nacl.bindings.crypto_pwhash_argon2id_MEMLIMIT_SENSITIVE - -OPSLIMIT_MODERATE = nacl.bindings.crypto_pwhash_argon2id_OPSLIMIT_MODERATE -MEMLIMIT_MODERATE = nacl.bindings.crypto_pwhash_argon2id_MEMLIMIT_MODERATE - - -def kdf( - size: int, - password: bytes, - salt: bytes, - opslimit: int = OPSLIMIT_SENSITIVE, - memlimit: int = MEMLIMIT_SENSITIVE, - encoder: nacl.encoding.Encoder = nacl.encoding.RawEncoder, -) -> bytes: - """ - Derive a ``size`` bytes long key from a caller-supplied - ``password`` and ``salt`` pair using the argon2i - memory-hard construct. - - the enclosing module provides the constants - - - :py:const:`.OPSLIMIT_INTERACTIVE` - - :py:const:`.MEMLIMIT_INTERACTIVE` - - :py:const:`.OPSLIMIT_MODERATE` - - :py:const:`.MEMLIMIT_MODERATE` - - :py:const:`.OPSLIMIT_SENSITIVE` - - :py:const:`.MEMLIMIT_SENSITIVE` - - as a guidance for correct settings. - - :param size: derived key size, must be between - :py:const:`.BYTES_MIN` and - :py:const:`.BYTES_MAX` - :type size: int - :param password: password used to seed the key derivation procedure; - it length must be between - :py:const:`.PASSWD_MIN` and - :py:const:`.PASSWD_MAX` - :type password: bytes - :param salt: **RANDOM** salt used in the key derivation procedure; - its length must be exactly :py:const:`.SALTBYTES` - :type salt: bytes - :param opslimit: the time component (operation count) - of the key derivation procedure's computational cost; - it must be between - :py:const:`.OPSLIMIT_MIN` and - :py:const:`.OPSLIMIT_MAX` - :type opslimit: int - :param memlimit: the memory occupation component - of the key derivation procedure's computational cost; - it must be between - :py:const:`.MEMLIMIT_MIN` and - :py:const:`.MEMLIMIT_MAX` - :type memlimit: int - :rtype: bytes - - .. versionadded:: 1.2 - """ - - return encoder.encode( - nacl.bindings.crypto_pwhash_alg( - size, password, salt, opslimit, memlimit, ALG - ) - ) - - -def str( - password: bytes, - opslimit: int = OPSLIMIT_INTERACTIVE, - memlimit: int = MEMLIMIT_INTERACTIVE, -) -> bytes: - """ - Hashes a password with a random salt, using the memory-hard - argon2id construct and returning an ascii string that has all - the needed info to check against a future password - - The default settings for opslimit and memlimit are those deemed - correct for the interactive user login case. - - :param bytes password: - :param int opslimit: - :param int memlimit: - :rtype: bytes - - .. versionadded:: 1.2 - """ - return nacl.bindings.crypto_pwhash_str_alg( - password, opslimit, memlimit, ALG - ) diff --git a/venv/lib/python3.10/site-packages/nacl/pwhash/scrypt.py b/venv/lib/python3.10/site-packages/nacl/pwhash/scrypt.py deleted file mode 100644 index 55bdf498a1d14bc63089e428375e9e7a6a06130e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/pwhash/scrypt.py +++ /dev/null @@ -1,211 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 cast - -import nacl.bindings -import nacl.encoding -from nacl import exceptions as exc -from nacl.exceptions import ensure - -_strbytes_plus_one = nacl.bindings.crypto_pwhash_scryptsalsa208sha256_STRBYTES - -AVAILABLE = nacl.bindings.has_crypto_pwhash_scryptsalsa208sha256 - -STRPREFIX = nacl.bindings.crypto_pwhash_scryptsalsa208sha256_STRPREFIX - -SALTBYTES = nacl.bindings.crypto_pwhash_scryptsalsa208sha256_SALTBYTES - -PASSWD_MIN = nacl.bindings.crypto_pwhash_scryptsalsa208sha256_PASSWD_MIN -PASSWD_MAX = nacl.bindings.crypto_pwhash_scryptsalsa208sha256_PASSWD_MAX - -PWHASH_SIZE = _strbytes_plus_one - 1 - -BYTES_MIN = nacl.bindings.crypto_pwhash_scryptsalsa208sha256_BYTES_MIN -BYTES_MAX = nacl.bindings.crypto_pwhash_scryptsalsa208sha256_BYTES_MAX - -MEMLIMIT_MIN = nacl.bindings.crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_MIN -MEMLIMIT_MAX = nacl.bindings.crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_MAX -OPSLIMIT_MIN = nacl.bindings.crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_MIN -OPSLIMIT_MAX = nacl.bindings.crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_MAX - -OPSLIMIT_INTERACTIVE = ( - nacl.bindings.crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_INTERACTIVE -) -MEMLIMIT_INTERACTIVE = ( - nacl.bindings.crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_INTERACTIVE -) -OPSLIMIT_SENSITIVE = ( - nacl.bindings.crypto_pwhash_scryptsalsa208sha256_OPSLIMIT_SENSITIVE -) -MEMLIMIT_SENSITIVE = ( - nacl.bindings.crypto_pwhash_scryptsalsa208sha256_MEMLIMIT_SENSITIVE -) - -OPSLIMIT_MODERATE = 8 * OPSLIMIT_INTERACTIVE -MEMLIMIT_MODERATE = 8 * MEMLIMIT_INTERACTIVE - - -def kdf( - size: int, - password: bytes, - salt: bytes, - opslimit: int = OPSLIMIT_SENSITIVE, - memlimit: int = MEMLIMIT_SENSITIVE, - encoder: nacl.encoding.Encoder = nacl.encoding.RawEncoder, -) -> bytes: - """ - Derive a ``size`` bytes long key from a caller-supplied - ``password`` and ``salt`` pair using the scryptsalsa208sha256 - memory-hard construct. - - - the enclosing module provides the constants - - - :py:const:`.OPSLIMIT_INTERACTIVE` - - :py:const:`.MEMLIMIT_INTERACTIVE` - - :py:const:`.OPSLIMIT_SENSITIVE` - - :py:const:`.MEMLIMIT_SENSITIVE` - - :py:const:`.OPSLIMIT_MODERATE` - - :py:const:`.MEMLIMIT_MODERATE` - - as a guidance for correct settings respectively for the - interactive login and the long term key protecting sensitive data - use cases. - - :param size: derived key size, must be between - :py:const:`.BYTES_MIN` and - :py:const:`.BYTES_MAX` - :type size: int - :param password: password used to seed the key derivation procedure; - it length must be between - :py:const:`.PASSWD_MIN` and - :py:const:`.PASSWD_MAX` - :type password: bytes - :param salt: **RANDOM** salt used in the key derivation procedure; - its length must be exactly :py:const:`.SALTBYTES` - :type salt: bytes - :param opslimit: the time component (operation count) - of the key derivation procedure's computational cost; - it must be between - :py:const:`.OPSLIMIT_MIN` and - :py:const:`.OPSLIMIT_MAX` - :type opslimit: int - :param memlimit: the memory occupation component - of the key derivation procedure's computational cost; - it must be between - :py:const:`.MEMLIMIT_MIN` and - :py:const:`.MEMLIMIT_MAX` - :type memlimit: int - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - - .. versionadded:: 1.2 - """ - ensure( - AVAILABLE, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - len(salt) == SALTBYTES, - "The salt must be exactly %s, not %s bytes long" - % (SALTBYTES, len(salt)), - raising=exc.ValueError, - ) - - n_log2, r, p = nacl.bindings.nacl_bindings_pick_scrypt_params( - opslimit, memlimit - ) - maxmem = memlimit + (2 ** 16) - - return encoder.encode( - nacl.bindings.crypto_pwhash_scryptsalsa208sha256_ll( - password, - salt, - # Cast safety: n_log2 is a positive integer, and so 2 ** n_log2 is also - # a positive integer. Mypy+typeshed can't deduce this, because there's no - # way to for them to know that n_log2: int is positive. - cast(int, 2 ** n_log2), - r, - p, - maxmem=maxmem, - dklen=size, - ) - ) - - -def str( - password: bytes, - opslimit: int = OPSLIMIT_INTERACTIVE, - memlimit: int = MEMLIMIT_INTERACTIVE, -) -> bytes: - """ - Hashes a password with a random salt, using the memory-hard - scryptsalsa208sha256 construct and returning an ascii string - that has all the needed info to check against a future password - - The default settings for opslimit and memlimit are those deemed - correct for the interactive user login case. - - :param bytes password: - :param int opslimit: - :param int memlimit: - :rtype: bytes - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - - .. versionadded:: 1.2 - """ - ensure( - AVAILABLE, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - return nacl.bindings.crypto_pwhash_scryptsalsa208sha256_str( - password, opslimit, memlimit - ) - - -def verify(password_hash: bytes, password: bytes) -> bool: - """ - Takes the output of scryptsalsa208sha256 and compares it against - a user provided password to see if they are the same - - :param password_hash: bytes - :param password: bytes - :rtype: boolean - :raises nacl.exceptions.UnavailableError: If called when using a - minimal build of libsodium. - - .. versionadded:: 1.2 - """ - ensure( - AVAILABLE, - "Not available in minimal build", - raising=exc.UnavailableError, - ) - - ensure( - len(password_hash) == PWHASH_SIZE, - "The password hash must be exactly %s bytes long" - % nacl.bindings.crypto_pwhash_scryptsalsa208sha256_STRBYTES, - raising=exc.ValueError, - ) - - return nacl.bindings.crypto_pwhash_scryptsalsa208sha256_str_verify( - password_hash, password - ) diff --git a/venv/lib/python3.10/site-packages/nacl/py.typed b/venv/lib/python3.10/site-packages/nacl/py.typed deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/venv/lib/python3.10/site-packages/nacl/secret.py b/venv/lib/python3.10/site-packages/nacl/secret.py deleted file mode 100644 index ba536a27d440b4f453f0fac62d860280d321769a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/secret.py +++ /dev/null @@ -1,305 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 ClassVar, Optional - -import nacl.bindings -from nacl import encoding -from nacl import exceptions as exc -from nacl.utils import EncryptedMessage, StringFixer, random - - -class SecretBox(encoding.Encodable, StringFixer): - """ - The SecretBox class encrypts and decrypts messages using the given secret - key. - - The ciphertexts generated by :class:`~nacl.secret.Secretbox` include a 16 - byte authenticator which is checked as part of the decryption. An invalid - authenticator will cause the decrypt function to raise an exception. The - authenticator is not a signature. Once you've decrypted the message you've - demonstrated the ability to create arbitrary valid message, so messages you - send are repudiable. For non-repudiable messages, sign them after - encryption. - - Encryption is done using `XSalsa20-Poly1305`_, and there are no practical - limits on the number or size of messages (up to 2⁶⁴ messages, each up to 2⁶⁴ - bytes). - - .. _XSalsa20-Poly1305: https://doc.libsodium.org/secret-key_cryptography/secretbox#algorithm-details - - :param key: The secret key used to encrypt and decrypt messages - :param encoder: The encoder class used to decode the given key - - :cvar KEY_SIZE: The size that the key is required to be. - :cvar NONCE_SIZE: The size that the nonce is required to be. - :cvar MACBYTES: The size of the authentication MAC tag in bytes. - :cvar MESSAGEBYTES_MAX: The maximum size of a message which can be - safely encrypted with a single key/nonce - pair. - """ - - KEY_SIZE: ClassVar[int] = nacl.bindings.crypto_secretbox_KEYBYTES - NONCE_SIZE: ClassVar[int] = nacl.bindings.crypto_secretbox_NONCEBYTES - MACBYTES: ClassVar[int] = nacl.bindings.crypto_secretbox_MACBYTES - MESSAGEBYTES_MAX: ClassVar[ - int - ] = nacl.bindings.crypto_secretbox_MESSAGEBYTES_MAX - - def __init__( - self, key: bytes, encoder: encoding.Encoder = encoding.RawEncoder - ): - key = encoder.decode(key) - if not isinstance(key, bytes): - raise exc.TypeError("SecretBox must be created from 32 bytes") - - if len(key) != self.KEY_SIZE: - raise exc.ValueError( - "The key must be exactly %s bytes long" % self.KEY_SIZE, - ) - - self._key = key - - def __bytes__(self) -> bytes: - return self._key - - def encrypt( - self, - plaintext: bytes, - nonce: Optional[bytes] = None, - encoder: encoding.Encoder = encoding.RawEncoder, - ) -> EncryptedMessage: - """ - Encrypts the plaintext message using the given `nonce` (or generates - one randomly if omitted) and returns the ciphertext encoded with the - encoder. - - .. warning:: It is **VITALLY** important that the nonce is a nonce, - i.e. it is a number used only once for any given key. If you fail - to do this, you compromise the privacy of the messages encrypted. - Give your nonces a different prefix, or have one side use an odd - counter and one an even counter. Just make sure they are different. - - :param plaintext: [:class:`bytes`] The plaintext message to encrypt - :param nonce: [:class:`bytes`] The nonce to use in the encryption - :param encoder: The encoder to use to encode the ciphertext - :rtype: [:class:`nacl.utils.EncryptedMessage`] - """ - if nonce is None: - nonce = random(self.NONCE_SIZE) - - if len(nonce) != self.NONCE_SIZE: - raise exc.ValueError( - "The nonce must be exactly %s bytes long" % self.NONCE_SIZE, - ) - - ciphertext = nacl.bindings.crypto_secretbox( - plaintext, nonce, self._key - ) - - encoded_nonce = encoder.encode(nonce) - encoded_ciphertext = encoder.encode(ciphertext) - - return EncryptedMessage._from_parts( - encoded_nonce, - encoded_ciphertext, - encoder.encode(nonce + ciphertext), - ) - - def decrypt( - self, - ciphertext: bytes, - nonce: Optional[bytes] = None, - encoder: encoding.Encoder = encoding.RawEncoder, - ) -> bytes: - """ - Decrypts the ciphertext using the `nonce` (explicitly, when passed as a - parameter or implicitly, when omitted, as part of the ciphertext) and - returns the plaintext message. - - :param ciphertext: [:class:`bytes`] The encrypted message to decrypt - :param nonce: [:class:`bytes`] The nonce used when encrypting the - ciphertext - :param encoder: The encoder used to decode the ciphertext. - :rtype: [:class:`bytes`] - """ - # Decode our ciphertext - ciphertext = encoder.decode(ciphertext) - - if nonce is None: - # If we were given the nonce and ciphertext combined, split them. - nonce = ciphertext[: self.NONCE_SIZE] - ciphertext = ciphertext[self.NONCE_SIZE :] - - if len(nonce) != self.NONCE_SIZE: - raise exc.ValueError( - "The nonce must be exactly %s bytes long" % self.NONCE_SIZE, - ) - - plaintext = nacl.bindings.crypto_secretbox_open( - ciphertext, nonce, self._key - ) - - return plaintext - - -class Aead(encoding.Encodable, StringFixer): - """ - The AEAD class encrypts and decrypts messages using the given secret key. - - Unlike :class:`~nacl.secret.SecretBox`, AEAD supports authenticating - non-confidential data received alongside the message, such as a length - or type tag. - - Like :class:`~nacl.secret.Secretbox`, this class provides authenticated - encryption. An inauthentic message will cause the decrypt function to raise - an exception. - - Likewise, the authenticator should not be mistaken for a (public-key) - signature: recipients (with the ability to decrypt messages) are capable of - creating arbitrary valid message; in particular, this means AEAD messages - are repudiable. For non-repudiable messages, sign them after encryption. - - The cryptosystem used is `XChacha20-Poly1305`_ as specified for - `standardization`_. There are `no practical limits`_ to how much can safely - be encrypted under a given key (up to 2⁶⁴ messages each containing up - to 2⁶⁴ bytes). - - .. _standardization: https://datatracker.ietf.org/doc/html/draft-irtf-cfrg-xchacha - .. _XChacha20-Poly1305: https://doc.libsodium.org/secret-key_cryptography/aead#xchacha-20-poly1305 - .. _no practical limits: https://doc.libsodium.org/secret-key_cryptography/aead#limitations - - :param key: The secret key used to encrypt and decrypt messages - :param encoder: The encoder class used to decode the given key - - :cvar KEY_SIZE: The size that the key is required to be. - :cvar NONCE_SIZE: The size that the nonce is required to be. - :cvar MACBYTES: The size of the authentication MAC tag in bytes. - :cvar MESSAGEBYTES_MAX: The maximum size of a message which can be - safely encrypted with a single key/nonce - pair. - """ - - KEY_SIZE = nacl.bindings.crypto_aead_xchacha20poly1305_ietf_KEYBYTES - NONCE_SIZE = nacl.bindings.crypto_aead_xchacha20poly1305_ietf_NPUBBYTES - MACBYTES = nacl.bindings.crypto_aead_xchacha20poly1305_ietf_ABYTES - MESSAGEBYTES_MAX = ( - nacl.bindings.crypto_aead_xchacha20poly1305_ietf_MESSAGEBYTES_MAX - ) - - def __init__( - self, - key: bytes, - encoder: encoding.Encoder = encoding.RawEncoder, - ): - key = encoder.decode(key) - if not isinstance(key, bytes): - raise exc.TypeError("AEAD must be created from 32 bytes") - - if len(key) != self.KEY_SIZE: - raise exc.ValueError( - "The key must be exactly %s bytes long" % self.KEY_SIZE, - ) - - self._key = key - - def __bytes__(self) -> bytes: - return self._key - - def encrypt( - self, - plaintext: bytes, - aad: bytes = b"", - nonce: Optional[bytes] = None, - encoder: encoding.Encoder = encoding.RawEncoder, - ) -> EncryptedMessage: - """ - Encrypts the plaintext message using the given `nonce` (or generates - one randomly if omitted) and returns the ciphertext encoded with the - encoder. - - .. warning:: It is vitally important for :param nonce: to be unique. - By default, it is generated randomly; [:class:`Aead`] uses XChacha20 - for extended (192b) nonce size, so the risk of reusing random nonces - is negligible. It is *strongly recommended* to keep this behaviour, - as nonce reuse will compromise the privacy of encrypted messages. - Should implicit nonces be inadequate for your application, the - second best option is using split counters; e.g. if sending messages - encrypted under a shared key between 2 users, each user can use the - number of messages it sent so far, prefixed or suffixed with a 1bit - user id. Note that the counter must **never** be rolled back (due - to overflow, on-disk state being rolled back to an earlier backup, - ...) - - :param plaintext: [:class:`bytes`] The plaintext message to encrypt - :param nonce: [:class:`bytes`] The nonce to use in the encryption - :param encoder: The encoder to use to encode the ciphertext - :rtype: [:class:`nacl.utils.EncryptedMessage`] - """ - if nonce is None: - nonce = random(self.NONCE_SIZE) - - if len(nonce) != self.NONCE_SIZE: - raise exc.ValueError( - "The nonce must be exactly %s bytes long" % self.NONCE_SIZE, - ) - - ciphertext = nacl.bindings.crypto_aead_xchacha20poly1305_ietf_encrypt( - plaintext, aad, nonce, self._key - ) - - encoded_nonce = encoder.encode(nonce) - encoded_ciphertext = encoder.encode(ciphertext) - - return EncryptedMessage._from_parts( - encoded_nonce, - encoded_ciphertext, - encoder.encode(nonce + ciphertext), - ) - - def decrypt( - self, - ciphertext: bytes, - aad: bytes = b"", - nonce: Optional[bytes] = None, - encoder: encoding.Encoder = encoding.RawEncoder, - ) -> bytes: - """ - Decrypts the ciphertext using the `nonce` (explicitly, when passed as a - parameter or implicitly, when omitted, as part of the ciphertext) and - returns the plaintext message. - - :param ciphertext: [:class:`bytes`] The encrypted message to decrypt - :param nonce: [:class:`bytes`] The nonce used when encrypting the - ciphertext - :param encoder: The encoder used to decode the ciphertext. - :rtype: [:class:`bytes`] - """ - # Decode our ciphertext - ciphertext = encoder.decode(ciphertext) - - if nonce is None: - # If we were given the nonce and ciphertext combined, split them. - nonce = ciphertext[: self.NONCE_SIZE] - ciphertext = ciphertext[self.NONCE_SIZE :] - - if len(nonce) != self.NONCE_SIZE: - raise exc.ValueError( - "The nonce must be exactly %s bytes long" % self.NONCE_SIZE, - ) - - plaintext = nacl.bindings.crypto_aead_xchacha20poly1305_ietf_decrypt( - ciphertext, aad, nonce, self._key - ) - - return plaintext diff --git a/venv/lib/python3.10/site-packages/nacl/signing.py b/venv/lib/python3.10/site-packages/nacl/signing.py deleted file mode 100644 index 12ec8ec477deed570897e8d41a36f672b5b8688f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/signing.py +++ /dev/null @@ -1,250 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 - -import nacl.bindings -from nacl import encoding -from nacl import exceptions as exc -from nacl.public import ( - PrivateKey as _Curve25519_PrivateKey, - PublicKey as _Curve25519_PublicKey, -) -from nacl.utils import StringFixer, random - - -class SignedMessage(bytes): - """ - A bytes subclass that holds a messaged that has been signed by a - :class:`SigningKey`. - """ - - _signature: bytes - _message: bytes - - @classmethod - def _from_parts( - cls, signature: bytes, message: bytes, combined: bytes - ) -> "SignedMessage": - obj = cls(combined) - obj._signature = signature - obj._message = message - return obj - - @property - def signature(self) -> bytes: - """ - The signature contained within the :class:`SignedMessage`. - """ - return self._signature - - @property - def message(self) -> bytes: - """ - The message contained within the :class:`SignedMessage`. - """ - return self._message - - -class VerifyKey(encoding.Encodable, StringFixer): - """ - The public key counterpart to an Ed25519 SigningKey for producing digital - signatures. - - :param key: [:class:`bytes`] Serialized Ed25519 public key - :param encoder: A class that is able to decode the `key` - """ - - def __init__( - self, key: bytes, encoder: encoding.Encoder = encoding.RawEncoder - ): - # Decode the key - key = encoder.decode(key) - if not isinstance(key, bytes): - raise exc.TypeError("VerifyKey must be created from 32 bytes") - - if len(key) != nacl.bindings.crypto_sign_PUBLICKEYBYTES: - raise exc.ValueError( - "The key must be exactly %s bytes long" - % nacl.bindings.crypto_sign_PUBLICKEYBYTES, - ) - - self._key = key - - def __bytes__(self) -> bytes: - return self._key - - def __hash__(self) -> int: - return hash(bytes(self)) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, self.__class__): - return False - return nacl.bindings.sodium_memcmp(bytes(self), bytes(other)) - - def __ne__(self, other: object) -> bool: - return not (self == other) - - def verify( - self, - smessage: bytes, - signature: Optional[bytes] = None, - encoder: encoding.Encoder = encoding.RawEncoder, - ) -> bytes: - """ - Verifies the signature of a signed message, returning the message - if it has not been tampered with else raising - :class:`~nacl.signing.BadSignatureError`. - - :param smessage: [:class:`bytes`] Either the original messaged or a - signature and message concated together. - :param signature: [:class:`bytes`] If an unsigned message is given for - smessage then the detached signature must be provided. - :param encoder: A class that is able to decode the secret message and - signature. - :rtype: :class:`bytes` - """ - if signature is not None: - # If we were given the message and signature separately, validate - # signature size and combine them. - if not isinstance(signature, bytes): - raise exc.TypeError( - "Verification signature must be created from %d bytes" - % nacl.bindings.crypto_sign_BYTES, - ) - - if len(signature) != nacl.bindings.crypto_sign_BYTES: - raise exc.ValueError( - "The signature must be exactly %d bytes long" - % nacl.bindings.crypto_sign_BYTES, - ) - - smessage = signature + encoder.decode(smessage) - else: - # Decode the signed message - smessage = encoder.decode(smessage) - - return nacl.bindings.crypto_sign_open(smessage, self._key) - - def to_curve25519_public_key(self) -> _Curve25519_PublicKey: - """ - Converts a :class:`~nacl.signing.VerifyKey` to a - :class:`~nacl.public.PublicKey` - - :rtype: :class:`~nacl.public.PublicKey` - """ - raw_pk = nacl.bindings.crypto_sign_ed25519_pk_to_curve25519(self._key) - return _Curve25519_PublicKey(raw_pk) - - -class SigningKey(encoding.Encodable, StringFixer): - """ - Private key for producing digital signatures using the Ed25519 algorithm. - - Signing keys are produced from a 32-byte (256-bit) random seed value. This - value can be passed into the :class:`~nacl.signing.SigningKey` as a - :func:`bytes` whose length is 32. - - .. warning:: This **must** be protected and remain secret. Anyone who knows - the value of your :class:`~nacl.signing.SigningKey` or it's seed can - masquerade as you. - - :param seed: [:class:`bytes`] Random 32-byte value (i.e. private key) - :param encoder: A class that is able to decode the seed - - :ivar: verify_key: [:class:`~nacl.signing.VerifyKey`] The verify - (i.e. public) key that corresponds with this signing key. - """ - - def __init__( - self, - seed: bytes, - encoder: encoding.Encoder = encoding.RawEncoder, - ): - # Decode the seed - seed = encoder.decode(seed) - if not isinstance(seed, bytes): - raise exc.TypeError( - "SigningKey must be created from a 32 byte seed" - ) - - # Verify that our seed is the proper size - if len(seed) != nacl.bindings.crypto_sign_SEEDBYTES: - raise exc.ValueError( - "The seed must be exactly %d bytes long" - % nacl.bindings.crypto_sign_SEEDBYTES - ) - - public_key, secret_key = nacl.bindings.crypto_sign_seed_keypair(seed) - - self._seed = seed - self._signing_key = secret_key - self.verify_key = VerifyKey(public_key) - - def __bytes__(self) -> bytes: - return self._seed - - def __hash__(self) -> int: - return hash(bytes(self)) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, self.__class__): - return False - return nacl.bindings.sodium_memcmp(bytes(self), bytes(other)) - - def __ne__(self, other: object) -> bool: - return not (self == other) - - @classmethod - def generate(cls) -> "SigningKey": - """ - Generates a random :class:`~nacl.signing.SigningKey` object. - - :rtype: :class:`~nacl.signing.SigningKey` - """ - return cls( - random(nacl.bindings.crypto_sign_SEEDBYTES), - encoder=encoding.RawEncoder, - ) - - def sign( - self, - message: bytes, - encoder: encoding.Encoder = encoding.RawEncoder, - ) -> SignedMessage: - """ - Sign a message using this key. - - :param message: [:class:`bytes`] The data to be signed. - :param encoder: A class that is used to encode the signed message. - :rtype: :class:`~nacl.signing.SignedMessage` - """ - raw_signed = nacl.bindings.crypto_sign(message, self._signing_key) - - crypto_sign_BYTES = nacl.bindings.crypto_sign_BYTES - signature = encoder.encode(raw_signed[:crypto_sign_BYTES]) - message = encoder.encode(raw_signed[crypto_sign_BYTES:]) - signed = encoder.encode(raw_signed) - - return SignedMessage._from_parts(signature, message, signed) - - def to_curve25519_private_key(self) -> _Curve25519_PrivateKey: - """ - Converts a :class:`~nacl.signing.SigningKey` to a - :class:`~nacl.public.PrivateKey` - - :rtype: :class:`~nacl.public.PrivateKey` - """ - sk = self._signing_key - raw_private = nacl.bindings.crypto_sign_ed25519_sk_to_curve25519(sk) - return _Curve25519_PrivateKey(raw_private) diff --git a/venv/lib/python3.10/site-packages/nacl/utils.py b/venv/lib/python3.10/site-packages/nacl/utils.py deleted file mode 100644 index d19d236a47c47eca5de1de2a64eed9079317364e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/nacl/utils.py +++ /dev/null @@ -1,88 +0,0 @@ -# Copyright 2013 Donald Stufft and individual contributors -# -# 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 -from typing import SupportsBytes, Type, TypeVar - -import nacl.bindings -from nacl import encoding - -_EncryptedMessage = TypeVar("_EncryptedMessage", bound="EncryptedMessage") - - -class EncryptedMessage(bytes): - """ - A bytes subclass that holds a messaged that has been encrypted by a - :class:`SecretBox`. - """ - - _nonce: bytes - _ciphertext: bytes - - @classmethod - def _from_parts( - cls: Type[_EncryptedMessage], - nonce: bytes, - ciphertext: bytes, - combined: bytes, - ) -> _EncryptedMessage: - obj = cls(combined) - obj._nonce = nonce - obj._ciphertext = ciphertext - return obj - - @property - def nonce(self) -> bytes: - """ - The nonce used during the encryption of the :class:`EncryptedMessage`. - """ - return self._nonce - - @property - def ciphertext(self) -> bytes: - """ - The ciphertext contained within the :class:`EncryptedMessage`. - """ - return self._ciphertext - - -class StringFixer: - def __str__(self: SupportsBytes) -> str: - return str(self.__bytes__()) - - -def bytes_as_string(bytes_in: bytes) -> str: - return bytes_in.decode("ascii") - - -def random(size: int = 32) -> bytes: - return os.urandom(size) - - -def randombytes_deterministic( - size: int, seed: bytes, encoder: encoding.Encoder = encoding.RawEncoder -) -> bytes: - """ - Returns ``size`` number of deterministically generated pseudorandom bytes - from a seed - - :param size: int - :param seed: bytes - :param encoder: The encoder class used to encode the produced bytes - :rtype: bytes - """ - raw_data = nacl.bindings.randombytes_buf_deterministic(size, seed) - - return encoder.encode(raw_data) diff --git a/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/INSTALLER deleted file mode 100644 index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/INSTALLER +++ /dev/null @@ -1 +0,0 @@ -pip diff --git a/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/METADATA b/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/METADATA deleted file mode 100644 index 2eff6a0c3865f050611859a41df6bf81aaf4f88e..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/METADATA +++ /dev/null @@ -1,58 +0,0 @@ -Metadata-Version: 2.4 -Name: Pygments -Version: 2.19.2 -Summary: Pygments is a syntax highlighting package written in Python. -Project-URL: Homepage, https://pygments.org -Project-URL: Documentation, https://pygments.org/docs -Project-URL: Source, https://github.com/pygments/pygments -Project-URL: Bug Tracker, https://github.com/pygments/pygments/issues -Project-URL: Changelog, https://github.com/pygments/pygments/blob/master/CHANGES -Author-email: Georg Brandl -Maintainer: Matthäus G. Chajdas -Maintainer-email: Georg Brandl , Jean Abou Samra -License: BSD-2-Clause -License-File: AUTHORS -License-File: LICENSE -Keywords: syntax highlighting -Classifier: Development Status :: 6 - Mature -Classifier: Intended Audience :: Developers -Classifier: Intended Audience :: End Users/Desktop -Classifier: Intended Audience :: System Administrators -Classifier: License :: OSI Approved :: BSD License -Classifier: Operating System :: OS Independent -Classifier: Programming Language :: Python -Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3.8 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Programming Language :: Python :: 3.10 -Classifier: Programming Language :: Python :: 3.11 -Classifier: Programming Language :: Python :: 3.12 -Classifier: Programming Language :: Python :: 3.13 -Classifier: Programming Language :: Python :: Implementation :: CPython -Classifier: Programming Language :: Python :: Implementation :: PyPy -Classifier: Topic :: Text Processing :: Filters -Classifier: Topic :: Utilities -Requires-Python: >=3.8 -Provides-Extra: plugins -Provides-Extra: windows-terminal -Requires-Dist: colorama>=0.4.6; extra == 'windows-terminal' -Description-Content-Type: text/x-rst - -Pygments -~~~~~~~~ - -Pygments is a syntax highlighting package written in Python. - -It is a generic syntax highlighter suitable for use in code hosting, forums, -wikis or other applications that need to prettify source code. Highlights -are: - -* a wide range of over 500 languages and other text formats is supported -* special attention is paid to details, increasing quality by a fair amount -* support for new languages and formats are added easily -* a number of output formats, presently HTML, LaTeX, RTF, SVG, all image - formats that PIL supports and ANSI sequences -* it is usable as a command-line tool and as a library - -Copyright 2006-2025 by the Pygments team, see ``AUTHORS``. -Licensed under the BSD, see ``LICENSE`` for details. diff --git a/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/RECORD b/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/RECORD deleted file mode 100644 index 18ddc0a3dd1982c49a967e0edebf6cbcbdafe88f..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/RECORD +++ /dev/null @@ -1,684 +0,0 @@ -../../../bin/pygmentize,sha256=tKppXig1xsUEqWk2TpJv80x3X3qmzbg9aWh2-xomRK8,233 -pygments-2.19.2.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 -pygments-2.19.2.dist-info/METADATA,sha256=euEA1n1nAGxkeYA92DX89HqbWfrHlEQeqOZqp_WYTYI,2512 -pygments-2.19.2.dist-info/RECORD,, -pygments-2.19.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87 -pygments-2.19.2.dist-info/entry_points.txt,sha256=uUXw-XhMKBEX4pWcCtpuTTnPhL3h7OEE2jWi51VQsa8,53 -pygments-2.19.2.dist-info/licenses/AUTHORS,sha256=BmDjGKbyFYAq3Icxq4XQxl_yfPzKP10oWX8wZHYZW9k,10824 -pygments-2.19.2.dist-info/licenses/LICENSE,sha256=qdZvHVJt8C4p3Oc0NtNOVuhjL0bCdbvf_HBWnogvnxc,1331 -pygments/__init__.py,sha256=_3UT86TGpHuW8FekdZ8uLidEZH1NhmcLiOy2KKNPCt4,2959 -pygments/__main__.py,sha256=p8AJyoyCOMYGvzWHdnq0_A9qaaVqaj02nIu3xhJp1_4,348 -pygments/__pycache__/__init__.cpython-310.pyc,, -pygments/__pycache__/__main__.cpython-310.pyc,, -pygments/__pycache__/cmdline.cpython-310.pyc,, -pygments/__pycache__/console.cpython-310.pyc,, -pygments/__pycache__/filter.cpython-310.pyc,, -pygments/__pycache__/formatter.cpython-310.pyc,, -pygments/__pycache__/lexer.cpython-310.pyc,, -pygments/__pycache__/modeline.cpython-310.pyc,, -pygments/__pycache__/plugin.cpython-310.pyc,, -pygments/__pycache__/regexopt.cpython-310.pyc,, -pygments/__pycache__/scanner.cpython-310.pyc,, -pygments/__pycache__/sphinxext.cpython-310.pyc,, -pygments/__pycache__/style.cpython-310.pyc,, -pygments/__pycache__/token.cpython-310.pyc,, -pygments/__pycache__/unistring.cpython-310.pyc,, -pygments/__pycache__/util.cpython-310.pyc,, -pygments/cmdline.py,sha256=4pL9Kpn2PUEKPobgrsQgg-vCx2NjsrapKzQ6LxQR7Q0,23536 -pygments/console.py,sha256=AagDWqwea2yBWf10KC9ptBgMpMjxKp8yABAmh-NQOVk,1718 -pygments/filter.py,sha256=YLtpTnZiu07nY3oK9nfR6E9Y1FBHhP5PX8gvkJWcfag,1910 -pygments/filters/__init__.py,sha256=B00KqPCQh5E0XhzaDK74Qa1E4fDSTlD6b0Pvr1v-vEQ,40344 -pygments/filters/__pycache__/__init__.cpython-310.pyc,, -pygments/formatter.py,sha256=H_4J-moKkKfRWUOW9J0u7hhw6n1LiO-2Xu1q2B0sE5w,4366 -pygments/formatters/__init__.py,sha256=7OuvmoYLyoPzoOQV_brHG8GSKYB_wjFSkAQng6x2y9g,5349 -pygments/formatters/__pycache__/__init__.cpython-310.pyc,, -pygments/formatters/__pycache__/_mapping.cpython-310.pyc,, -pygments/formatters/__pycache__/bbcode.cpython-310.pyc,, -pygments/formatters/__pycache__/groff.cpython-310.pyc,, -pygments/formatters/__pycache__/html.cpython-310.pyc,, -pygments/formatters/__pycache__/img.cpython-310.pyc,, -pygments/formatters/__pycache__/irc.cpython-310.pyc,, -pygments/formatters/__pycache__/latex.cpython-310.pyc,, -pygments/formatters/__pycache__/other.cpython-310.pyc,, -pygments/formatters/__pycache__/pangomarkup.cpython-310.pyc,, -pygments/formatters/__pycache__/rtf.cpython-310.pyc,, -pygments/formatters/__pycache__/svg.cpython-310.pyc,, -pygments/formatters/__pycache__/terminal.cpython-310.pyc,, -pygments/formatters/__pycache__/terminal256.cpython-310.pyc,, -pygments/formatters/_mapping.py,sha256=1Cw37FuQlNacnxRKmtlPX4nyLoX9_ttko5ZwscNUZZ4,4176 -pygments/formatters/bbcode.py,sha256=s0Ka35OKuIchoSgEAGf6rj0rl2a9ym9L31JVNSRbZFQ,3296 -pygments/formatters/groff.py,sha256=pLcIHj4jJS_lRAVFnyJODKDu1Xlyl9_AEIdOtbl3DT0,5082 -pygments/formatters/html.py,sha256=FrHJ69FUliEyPY0zTfab0C1gPf7LXsKgeRlhwkniqIs,35953 -pygments/formatters/img.py,sha256=aRpFo8mBmWTL3sBUjRCWkeS3rc6FZrSFC4EksDrl53g,23301 -pygments/formatters/irc.py,sha256=R0Js0TYWySlI2yE9sW6tN4d4X-x3k9ZmudsijGPnLmU,4945 -pygments/formatters/latex.py,sha256=BRYtbLeW_YD1kwhhnFInhJIKylurnri8CF1lP069KWE,19258 -pygments/formatters/other.py,sha256=8pYW27sU_7XicLUqOEt2yWSO0h1IEUM3TIv34KODLwo,4986 -pygments/formatters/pangomarkup.py,sha256=pcFvEC7K1Me0EjGeOZth4oCnEY85bfqc77XzZASEPpY,2206 -pygments/formatters/rtf.py,sha256=kcKMCxTXu-2-hpgEftlGJRm7Ss-yA_Sy8OsHH_qzykA,11921 -pygments/formatters/svg.py,sha256=R6A2ME6JsMQWFiyn8wcKwFUOD6vsu-HLwiIztLu-77E,7138 -pygments/formatters/terminal.py,sha256=J_F_dFXwR9LHWvatIDnwqRYJyjVmSo1Zx8K_XDh6SyM,4626 -pygments/formatters/terminal256.py,sha256=7GQFLE5cfmeu53CAzANO74-kBk2BFkXfn5phmZjYkhM,11717 -pygments/lexer.py,sha256=ib-F_0GxHkwGpb6vWP0DeLMLc7EYgjo3hWFKN5IgOq0,35109 -pygments/lexers/__init__.py,sha256=6YhzxGKlWk38P6JpIJUQ1rVvV0DEZjEmdYsdMQ58hSk,12067 -pygments/lexers/__pycache__/__init__.cpython-310.pyc,, -pygments/lexers/__pycache__/_ada_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_asy_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_cl_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_cocoa_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_csound_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_css_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_googlesql_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_julia_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_lasso_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_lilypond_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_lua_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_luau_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_mapping.cpython-310.pyc,, -pygments/lexers/__pycache__/_mql_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_mysql_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_openedge_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_php_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_postgres_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_qlik_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_scheme_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_scilab_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_sourcemod_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_sql_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_stan_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_stata_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_tsql_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_usd_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_vbscript_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/_vim_builtins.cpython-310.pyc,, -pygments/lexers/__pycache__/actionscript.cpython-310.pyc,, -pygments/lexers/__pycache__/ada.cpython-310.pyc,, -pygments/lexers/__pycache__/agile.cpython-310.pyc,, -pygments/lexers/__pycache__/algebra.cpython-310.pyc,, -pygments/lexers/__pycache__/ambient.cpython-310.pyc,, -pygments/lexers/__pycache__/amdgpu.cpython-310.pyc,, -pygments/lexers/__pycache__/ampl.cpython-310.pyc,, -pygments/lexers/__pycache__/apdlexer.cpython-310.pyc,, -pygments/lexers/__pycache__/apl.cpython-310.pyc,, -pygments/lexers/__pycache__/archetype.cpython-310.pyc,, -pygments/lexers/__pycache__/arrow.cpython-310.pyc,, -pygments/lexers/__pycache__/arturo.cpython-310.pyc,, -pygments/lexers/__pycache__/asc.cpython-310.pyc,, -pygments/lexers/__pycache__/asm.cpython-310.pyc,, -pygments/lexers/__pycache__/asn1.cpython-310.pyc,, -pygments/lexers/__pycache__/automation.cpython-310.pyc,, -pygments/lexers/__pycache__/bare.cpython-310.pyc,, -pygments/lexers/__pycache__/basic.cpython-310.pyc,, -pygments/lexers/__pycache__/bdd.cpython-310.pyc,, -pygments/lexers/__pycache__/berry.cpython-310.pyc,, -pygments/lexers/__pycache__/bibtex.cpython-310.pyc,, -pygments/lexers/__pycache__/blueprint.cpython-310.pyc,, -pygments/lexers/__pycache__/boa.cpython-310.pyc,, -pygments/lexers/__pycache__/bqn.cpython-310.pyc,, -pygments/lexers/__pycache__/business.cpython-310.pyc,, -pygments/lexers/__pycache__/c_cpp.cpython-310.pyc,, -pygments/lexers/__pycache__/c_like.cpython-310.pyc,, -pygments/lexers/__pycache__/capnproto.cpython-310.pyc,, -pygments/lexers/__pycache__/carbon.cpython-310.pyc,, -pygments/lexers/__pycache__/cddl.cpython-310.pyc,, -pygments/lexers/__pycache__/chapel.cpython-310.pyc,, -pygments/lexers/__pycache__/clean.cpython-310.pyc,, -pygments/lexers/__pycache__/codeql.cpython-310.pyc,, -pygments/lexers/__pycache__/comal.cpython-310.pyc,, -pygments/lexers/__pycache__/compiled.cpython-310.pyc,, -pygments/lexers/__pycache__/configs.cpython-310.pyc,, -pygments/lexers/__pycache__/console.cpython-310.pyc,, -pygments/lexers/__pycache__/cplint.cpython-310.pyc,, -pygments/lexers/__pycache__/crystal.cpython-310.pyc,, -pygments/lexers/__pycache__/csound.cpython-310.pyc,, -pygments/lexers/__pycache__/css.cpython-310.pyc,, -pygments/lexers/__pycache__/d.cpython-310.pyc,, -pygments/lexers/__pycache__/dalvik.cpython-310.pyc,, -pygments/lexers/__pycache__/data.cpython-310.pyc,, -pygments/lexers/__pycache__/dax.cpython-310.pyc,, -pygments/lexers/__pycache__/devicetree.cpython-310.pyc,, -pygments/lexers/__pycache__/diff.cpython-310.pyc,, -pygments/lexers/__pycache__/dns.cpython-310.pyc,, -pygments/lexers/__pycache__/dotnet.cpython-310.pyc,, -pygments/lexers/__pycache__/dsls.cpython-310.pyc,, -pygments/lexers/__pycache__/dylan.cpython-310.pyc,, -pygments/lexers/__pycache__/ecl.cpython-310.pyc,, -pygments/lexers/__pycache__/eiffel.cpython-310.pyc,, -pygments/lexers/__pycache__/elm.cpython-310.pyc,, -pygments/lexers/__pycache__/elpi.cpython-310.pyc,, -pygments/lexers/__pycache__/email.cpython-310.pyc,, -pygments/lexers/__pycache__/erlang.cpython-310.pyc,, -pygments/lexers/__pycache__/esoteric.cpython-310.pyc,, -pygments/lexers/__pycache__/ezhil.cpython-310.pyc,, -pygments/lexers/__pycache__/factor.cpython-310.pyc,, -pygments/lexers/__pycache__/fantom.cpython-310.pyc,, -pygments/lexers/__pycache__/felix.cpython-310.pyc,, -pygments/lexers/__pycache__/fift.cpython-310.pyc,, -pygments/lexers/__pycache__/floscript.cpython-310.pyc,, -pygments/lexers/__pycache__/forth.cpython-310.pyc,, -pygments/lexers/__pycache__/fortran.cpython-310.pyc,, -pygments/lexers/__pycache__/foxpro.cpython-310.pyc,, -pygments/lexers/__pycache__/freefem.cpython-310.pyc,, -pygments/lexers/__pycache__/func.cpython-310.pyc,, -pygments/lexers/__pycache__/functional.cpython-310.pyc,, -pygments/lexers/__pycache__/futhark.cpython-310.pyc,, -pygments/lexers/__pycache__/gcodelexer.cpython-310.pyc,, -pygments/lexers/__pycache__/gdscript.cpython-310.pyc,, -pygments/lexers/__pycache__/gleam.cpython-310.pyc,, -pygments/lexers/__pycache__/go.cpython-310.pyc,, -pygments/lexers/__pycache__/grammar_notation.cpython-310.pyc,, -pygments/lexers/__pycache__/graph.cpython-310.pyc,, -pygments/lexers/__pycache__/graphics.cpython-310.pyc,, -pygments/lexers/__pycache__/graphql.cpython-310.pyc,, -pygments/lexers/__pycache__/graphviz.cpython-310.pyc,, -pygments/lexers/__pycache__/gsql.cpython-310.pyc,, -pygments/lexers/__pycache__/hare.cpython-310.pyc,, -pygments/lexers/__pycache__/haskell.cpython-310.pyc,, -pygments/lexers/__pycache__/haxe.cpython-310.pyc,, -pygments/lexers/__pycache__/hdl.cpython-310.pyc,, -pygments/lexers/__pycache__/hexdump.cpython-310.pyc,, -pygments/lexers/__pycache__/html.cpython-310.pyc,, -pygments/lexers/__pycache__/idl.cpython-310.pyc,, -pygments/lexers/__pycache__/igor.cpython-310.pyc,, -pygments/lexers/__pycache__/inferno.cpython-310.pyc,, -pygments/lexers/__pycache__/installers.cpython-310.pyc,, -pygments/lexers/__pycache__/int_fiction.cpython-310.pyc,, -pygments/lexers/__pycache__/iolang.cpython-310.pyc,, -pygments/lexers/__pycache__/j.cpython-310.pyc,, -pygments/lexers/__pycache__/javascript.cpython-310.pyc,, -pygments/lexers/__pycache__/jmespath.cpython-310.pyc,, -pygments/lexers/__pycache__/jslt.cpython-310.pyc,, -pygments/lexers/__pycache__/json5.cpython-310.pyc,, -pygments/lexers/__pycache__/jsonnet.cpython-310.pyc,, -pygments/lexers/__pycache__/jsx.cpython-310.pyc,, -pygments/lexers/__pycache__/julia.cpython-310.pyc,, -pygments/lexers/__pycache__/jvm.cpython-310.pyc,, -pygments/lexers/__pycache__/kuin.cpython-310.pyc,, -pygments/lexers/__pycache__/kusto.cpython-310.pyc,, -pygments/lexers/__pycache__/ldap.cpython-310.pyc,, -pygments/lexers/__pycache__/lean.cpython-310.pyc,, -pygments/lexers/__pycache__/lilypond.cpython-310.pyc,, -pygments/lexers/__pycache__/lisp.cpython-310.pyc,, -pygments/lexers/__pycache__/macaulay2.cpython-310.pyc,, -pygments/lexers/__pycache__/make.cpython-310.pyc,, -pygments/lexers/__pycache__/maple.cpython-310.pyc,, -pygments/lexers/__pycache__/markup.cpython-310.pyc,, -pygments/lexers/__pycache__/math.cpython-310.pyc,, -pygments/lexers/__pycache__/matlab.cpython-310.pyc,, -pygments/lexers/__pycache__/maxima.cpython-310.pyc,, -pygments/lexers/__pycache__/meson.cpython-310.pyc,, -pygments/lexers/__pycache__/mime.cpython-310.pyc,, -pygments/lexers/__pycache__/minecraft.cpython-310.pyc,, -pygments/lexers/__pycache__/mips.cpython-310.pyc,, -pygments/lexers/__pycache__/ml.cpython-310.pyc,, -pygments/lexers/__pycache__/modeling.cpython-310.pyc,, -pygments/lexers/__pycache__/modula2.cpython-310.pyc,, -pygments/lexers/__pycache__/mojo.cpython-310.pyc,, -pygments/lexers/__pycache__/monte.cpython-310.pyc,, -pygments/lexers/__pycache__/mosel.cpython-310.pyc,, -pygments/lexers/__pycache__/ncl.cpython-310.pyc,, -pygments/lexers/__pycache__/nimrod.cpython-310.pyc,, -pygments/lexers/__pycache__/nit.cpython-310.pyc,, -pygments/lexers/__pycache__/nix.cpython-310.pyc,, -pygments/lexers/__pycache__/numbair.cpython-310.pyc,, -pygments/lexers/__pycache__/oberon.cpython-310.pyc,, -pygments/lexers/__pycache__/objective.cpython-310.pyc,, -pygments/lexers/__pycache__/ooc.cpython-310.pyc,, -pygments/lexers/__pycache__/openscad.cpython-310.pyc,, -pygments/lexers/__pycache__/other.cpython-310.pyc,, -pygments/lexers/__pycache__/parasail.cpython-310.pyc,, -pygments/lexers/__pycache__/parsers.cpython-310.pyc,, -pygments/lexers/__pycache__/pascal.cpython-310.pyc,, -pygments/lexers/__pycache__/pawn.cpython-310.pyc,, -pygments/lexers/__pycache__/pddl.cpython-310.pyc,, -pygments/lexers/__pycache__/perl.cpython-310.pyc,, -pygments/lexers/__pycache__/phix.cpython-310.pyc,, -pygments/lexers/__pycache__/php.cpython-310.pyc,, -pygments/lexers/__pycache__/pointless.cpython-310.pyc,, -pygments/lexers/__pycache__/pony.cpython-310.pyc,, -pygments/lexers/__pycache__/praat.cpython-310.pyc,, -pygments/lexers/__pycache__/procfile.cpython-310.pyc,, -pygments/lexers/__pycache__/prolog.cpython-310.pyc,, -pygments/lexers/__pycache__/promql.cpython-310.pyc,, -pygments/lexers/__pycache__/prql.cpython-310.pyc,, -pygments/lexers/__pycache__/ptx.cpython-310.pyc,, -pygments/lexers/__pycache__/python.cpython-310.pyc,, -pygments/lexers/__pycache__/q.cpython-310.pyc,, -pygments/lexers/__pycache__/qlik.cpython-310.pyc,, -pygments/lexers/__pycache__/qvt.cpython-310.pyc,, -pygments/lexers/__pycache__/r.cpython-310.pyc,, -pygments/lexers/__pycache__/rdf.cpython-310.pyc,, -pygments/lexers/__pycache__/rebol.cpython-310.pyc,, -pygments/lexers/__pycache__/rego.cpython-310.pyc,, -pygments/lexers/__pycache__/resource.cpython-310.pyc,, -pygments/lexers/__pycache__/ride.cpython-310.pyc,, -pygments/lexers/__pycache__/rita.cpython-310.pyc,, -pygments/lexers/__pycache__/rnc.cpython-310.pyc,, -pygments/lexers/__pycache__/roboconf.cpython-310.pyc,, -pygments/lexers/__pycache__/robotframework.cpython-310.pyc,, -pygments/lexers/__pycache__/ruby.cpython-310.pyc,, -pygments/lexers/__pycache__/rust.cpython-310.pyc,, -pygments/lexers/__pycache__/sas.cpython-310.pyc,, -pygments/lexers/__pycache__/savi.cpython-310.pyc,, -pygments/lexers/__pycache__/scdoc.cpython-310.pyc,, -pygments/lexers/__pycache__/scripting.cpython-310.pyc,, -pygments/lexers/__pycache__/sgf.cpython-310.pyc,, -pygments/lexers/__pycache__/shell.cpython-310.pyc,, -pygments/lexers/__pycache__/sieve.cpython-310.pyc,, -pygments/lexers/__pycache__/slash.cpython-310.pyc,, -pygments/lexers/__pycache__/smalltalk.cpython-310.pyc,, -pygments/lexers/__pycache__/smithy.cpython-310.pyc,, -pygments/lexers/__pycache__/smv.cpython-310.pyc,, -pygments/lexers/__pycache__/snobol.cpython-310.pyc,, -pygments/lexers/__pycache__/solidity.cpython-310.pyc,, -pygments/lexers/__pycache__/soong.cpython-310.pyc,, -pygments/lexers/__pycache__/sophia.cpython-310.pyc,, -pygments/lexers/__pycache__/special.cpython-310.pyc,, -pygments/lexers/__pycache__/spice.cpython-310.pyc,, -pygments/lexers/__pycache__/sql.cpython-310.pyc,, -pygments/lexers/__pycache__/srcinfo.cpython-310.pyc,, -pygments/lexers/__pycache__/stata.cpython-310.pyc,, -pygments/lexers/__pycache__/supercollider.cpython-310.pyc,, -pygments/lexers/__pycache__/tablegen.cpython-310.pyc,, -pygments/lexers/__pycache__/tact.cpython-310.pyc,, -pygments/lexers/__pycache__/tal.cpython-310.pyc,, -pygments/lexers/__pycache__/tcl.cpython-310.pyc,, -pygments/lexers/__pycache__/teal.cpython-310.pyc,, -pygments/lexers/__pycache__/templates.cpython-310.pyc,, -pygments/lexers/__pycache__/teraterm.cpython-310.pyc,, -pygments/lexers/__pycache__/testing.cpython-310.pyc,, -pygments/lexers/__pycache__/text.cpython-310.pyc,, -pygments/lexers/__pycache__/textedit.cpython-310.pyc,, -pygments/lexers/__pycache__/textfmts.cpython-310.pyc,, -pygments/lexers/__pycache__/theorem.cpython-310.pyc,, -pygments/lexers/__pycache__/thingsdb.cpython-310.pyc,, -pygments/lexers/__pycache__/tlb.cpython-310.pyc,, -pygments/lexers/__pycache__/tls.cpython-310.pyc,, -pygments/lexers/__pycache__/tnt.cpython-310.pyc,, -pygments/lexers/__pycache__/trafficscript.cpython-310.pyc,, -pygments/lexers/__pycache__/typoscript.cpython-310.pyc,, -pygments/lexers/__pycache__/typst.cpython-310.pyc,, -pygments/lexers/__pycache__/ul4.cpython-310.pyc,, -pygments/lexers/__pycache__/unicon.cpython-310.pyc,, -pygments/lexers/__pycache__/urbi.cpython-310.pyc,, -pygments/lexers/__pycache__/usd.cpython-310.pyc,, -pygments/lexers/__pycache__/varnish.cpython-310.pyc,, -pygments/lexers/__pycache__/verification.cpython-310.pyc,, -pygments/lexers/__pycache__/verifpal.cpython-310.pyc,, -pygments/lexers/__pycache__/vip.cpython-310.pyc,, -pygments/lexers/__pycache__/vyper.cpython-310.pyc,, -pygments/lexers/__pycache__/web.cpython-310.pyc,, -pygments/lexers/__pycache__/webassembly.cpython-310.pyc,, -pygments/lexers/__pycache__/webidl.cpython-310.pyc,, -pygments/lexers/__pycache__/webmisc.cpython-310.pyc,, -pygments/lexers/__pycache__/wgsl.cpython-310.pyc,, -pygments/lexers/__pycache__/whiley.cpython-310.pyc,, -pygments/lexers/__pycache__/wowtoc.cpython-310.pyc,, -pygments/lexers/__pycache__/wren.cpython-310.pyc,, -pygments/lexers/__pycache__/x10.cpython-310.pyc,, -pygments/lexers/__pycache__/xorg.cpython-310.pyc,, -pygments/lexers/__pycache__/yang.cpython-310.pyc,, -pygments/lexers/__pycache__/yara.cpython-310.pyc,, -pygments/lexers/__pycache__/zig.cpython-310.pyc,, -pygments/lexers/_ada_builtins.py,sha256=CA_OnShtdc7wWh9oYcRlcrkDAQwYUKl6w7tdSbALQd4,1543 -pygments/lexers/_asy_builtins.py,sha256=cd9M00YH19w5ZL7aqucmC3nwpJGTS04U-01NLy5E2_4,27287 -pygments/lexers/_cl_builtins.py,sha256=kQeUIyZjP4kX0frkICDcKxBYQCLqzIDXa5WV5cevhDo,13994 -pygments/lexers/_cocoa_builtins.py,sha256=Ka1lLJe7JfWtdho4IFIB82X9yBvrbfHCCmEG-peXXhQ,105173 -pygments/lexers/_csound_builtins.py,sha256=qnQYKeI26ZHim316uqy_hDiRiCoHo2RHjD3sYBALyXs,18414 -pygments/lexers/_css_builtins.py,sha256=aD-dhLFXVd1Atn_bZd7gEdQn7Mhe60_VHpvZ340WzDI,12446 -pygments/lexers/_googlesql_builtins.py,sha256=IkrOk-T2v1yzbGzUEEQh5_Cf4uC_cmL_uuhwDpZlTug,16132 -pygments/lexers/_julia_builtins.py,sha256=N2WdSw5zgI2fhDat_i4YeVqurRTC_P8x71ez00SCN6U,11883 -pygments/lexers/_lasso_builtins.py,sha256=8q1gbsrMJeaeUhxIYKhaOxC9j_B-NBpq_XFj2Ze41X0,134510 -pygments/lexers/_lilypond_builtins.py,sha256=XTbGL1z1oKMoqWLEktG33jx5GdGTI9CpeO5NheEi4Y0,108094 -pygments/lexers/_lua_builtins.py,sha256=PhFdZV5-Tzz2j_q4lvG9lr84ELGfL41BhnrSDNNTaG4,8108 -pygments/lexers/_luau_builtins.py,sha256=-IDrU04kUVfjXwSQzMMpXmMYhNsQxZVVZk8cuAA0Lo0,955 -pygments/lexers/_mapping.py,sha256=9fv7xYOUAOr6LzfdFS4MDbPu78o4OQQH-2nsI1bNZf4,70438 -pygments/lexers/_mql_builtins.py,sha256=ybRQjlb7Cul0sDstnzxJl3h0qS6Ieqsr811fqrxyumU,24713 -pygments/lexers/_mysql_builtins.py,sha256=y0kAWZVAs0z2dTFJJV42OZpILgRnd8T3zSlBFv-g_oA,25838 -pygments/lexers/_openedge_builtins.py,sha256=Sz4j9-CPWIaxMa-2fZgY66j7igcu1ob1GR2UtI8zAkg,49398 -pygments/lexers/_php_builtins.py,sha256=Jd4BZpjMDELPi4EVoSxK1-8BFTc63HUwYfm1rLrGj0M,107922 -pygments/lexers/_postgres_builtins.py,sha256=Pqh4z0RBRbnW6rCQtWUdzWCJxNyqpJ7_0HOktxHDxk4,13343 -pygments/lexers/_qlik_builtins.py,sha256=xuJy9c9uZDXv6h8z582P5PrxqkxTZ_nS8gPl9OD9VN8,12595 -pygments/lexers/_scheme_builtins.py,sha256=2hNtJOJmP21lUsikpqMJ2gAmLT3Rwn_KEeqhXwCjgfk,32564 -pygments/lexers/_scilab_builtins.py,sha256=oZYPB1XPdIEz3pII11pFDe6extRRyWGA7pY06X8KZ8w,52411 -pygments/lexers/_sourcemod_builtins.py,sha256=H8AFLsNDdEpymIWOpDwbDJGCP1w-x-1gSlzPDioMF4o,26777 -pygments/lexers/_sql_builtins.py,sha256=oe8F9wWuO2iS6nEsZAdJtCUChBTjgM1Sq_aipu74jXM,6767 -pygments/lexers/_stan_builtins.py,sha256=dwi1hllM_NsaCv-aXJy7lEi57X5Hh5gSD97aCQyT9KM,13445 -pygments/lexers/_stata_builtins.py,sha256=Hqrr6j77zWU3cGGpBPohwexZci43YA4_sVYE4E1sNow,27227 -pygments/lexers/_tsql_builtins.py,sha256=Pi2RhTXcLE3glI9oxNhyVsOMn-fK_1TRxJ-EsYP5LcI,15460 -pygments/lexers/_usd_builtins.py,sha256=c9hbU1cwqBUCFIhNfu_Dob8ywv1rlPhi9w2OTj3kR8s,1658 -pygments/lexers/_vbscript_builtins.py,sha256=MqJ2ABywD21aSRtWYZRG64CCbGstC1kfsiHGJmZzxiw,4225 -pygments/lexers/_vim_builtins.py,sha256=bA4mH8t1mPPQfEiUCKEqRO1O0rL2DUG0Ux1Bt8ZSu0E,57066 -pygments/lexers/actionscript.py,sha256=JBngCe5UhYT_0dLD2j7PnPO0xRRJhmypEuQ-C5in8pY,11727 -pygments/lexers/ada.py,sha256=58k5ra1vGS4iLpW3h1ItY9ftzF3WevaeAAXzAYTiYkQ,5353 -pygments/lexers/agile.py,sha256=DN-7AVIqtG1MshA94rtSGYI_884hVHgzq405wD0_dl8,896 -pygments/lexers/algebra.py,sha256=yGTu9Tt-cQzAISQYIC5MS5a3z4QmL-tGcXnd_pkWGbk,9952 -pygments/lexers/ambient.py,sha256=UnzKpIlfSm3iitHvMd7XTMSY8TjZYYhKOC3AiARS_cE,2605 -pygments/lexers/amdgpu.py,sha256=S8qjn2UMLhBFm3Yn_c06XAGf8cl5x_ZeluelWG_-JAw,1723 -pygments/lexers/ampl.py,sha256=ZBRfDXm760gR1a1gqItnsHuoO3JdUcTBjJ5tFY9UtPA,4176 -pygments/lexers/apdlexer.py,sha256=Zr5-jgjxC8PKzRlEeclakZXPHci7FHBZghQ6wwiuT7A,30800 -pygments/lexers/apl.py,sha256=PTQMp-bxT5P-DbrEvFha10HBTcsDJ5srL3I1s9ljz58,3404 -pygments/lexers/archetype.py,sha256=pQVlP1Fb5OA8nn7QwmFaaaOSvvpoIsQVw43FVCQCve4,11538 -pygments/lexers/arrow.py,sha256=2PKdbWq3xQLF1KoDbWvSxpjwKRrznnDiArTflRGZzBo,3564 -pygments/lexers/arturo.py,sha256=U5MtRNHJtnBn4ZOeWmW6MKlVRG7SX6KhTRamDqzn9tA,11414 -pygments/lexers/asc.py,sha256=-DgZl9jccBDHPlDmjCsrEqx0-Q7ap7XVdNKtxLNWG1w,1693 -pygments/lexers/asm.py,sha256=xm2Y5mcT-sF3oQvair4SWs9EWTyndoaUoSsDy5v6shI,41967 -pygments/lexers/asn1.py,sha256=BlcloIX2bu6Q7BxGcksuhYFHGsXLVKyB4B9mFd4Pj6E,4262 -pygments/lexers/automation.py,sha256=Q61qon8EwpfakMh_2MS2E2zUUT16rG3UNIKPYjITeTs,19831 -pygments/lexers/bare.py,sha256=tWoei86JJX1k-ADhaXd5TgX6ItDTici9yFWpkTPhnfM,3020 -pygments/lexers/basic.py,sha256=qpVe5h8Fa7NJo1EihN-4R_UZpHO6my2Ssgkb-BktkKs,27989 -pygments/lexers/bdd.py,sha256=yysefcOFAEyk9kJ2y4EXmzJTecgLYUHlWixt_3YzPMU,1641 -pygments/lexers/berry.py,sha256=zxGowFb8HMIyN15-m8nmWnW6bPRR4esKtSEVugc9uXM,3209 -pygments/lexers/bibtex.py,sha256=yuNoPxwrJf9DCGUT17hxfDzbq_HtCLkQkRbBtiTVmeQ,4811 -pygments/lexers/blueprint.py,sha256=NzvWHMxCLDWt8hc6gB5jokltxVJgNa7Jwh4c61ng388,6188 -pygments/lexers/boa.py,sha256=dOot1XWNZThPIio2UyAX67K6EpISjSRCFjotD7dcnwE,3921 -pygments/lexers/bqn.py,sha256=nJiwrPKKbRF-qdai5tfqipwBkkko2P3weiZAjHUMimY,3671 -pygments/lexers/business.py,sha256=lRtekOJfsDkb12AGbuz10-G67OJrVJgCBtihTQ8_aoY,28345 -pygments/lexers/c_cpp.py,sha256=D7ZIswaHASlGBgoTlwnSqTQHf8_JyvvSt2L2q1W-F6g,18059 -pygments/lexers/c_like.py,sha256=FTGp17ds6X2rDZOHup2hH6BEn3gKK4nLm9pydNEhm0E,32021 -pygments/lexers/capnproto.py,sha256=XQJAh1WS-0ulqbTn9TdzR6gEgWLcuBqb4sj3jNsrhsY,2174 -pygments/lexers/carbon.py,sha256=av12YuTGZGpOa1Cmxp3lppx3LfSJUWbvOu0ixmUVll0,3211 -pygments/lexers/cddl.py,sha256=MKa70IwABgjBjYu15_Q9v8rsu2sr1a-i2jkiaPTI6sM,5076 -pygments/lexers/chapel.py,sha256=0n_fL3ehLC4pw4YKnmq9jxIXOJcxGPka1Wr1t1zsXPc,5156 -pygments/lexers/clean.py,sha256=dkDPAwF5BTALPeuKFoRKOSD3RfsKcGWbaRo6_G8LHng,6418 -pygments/lexers/codeql.py,sha256=ebvghn2zbrnETV4buVozMDmRCVKSdGiIN8ycLlHpGsE,2576 -pygments/lexers/comal.py,sha256=TC3NzcJ58ew5jw7qwK0kJ-okTA47psZje0yAIS39HR4,3179 -pygments/lexers/compiled.py,sha256=Slfo1sjWqcPawUwf0dIIZLBCL5pkOIoAX2S8Lxs02Mc,1426 -pygments/lexers/configs.py,sha256=wW8pY0Sa5a10pnAeTLGf48HhixQTVageIyHEf1aYMCc,50913 -pygments/lexers/console.py,sha256=-jAG120dupvV3kG3zC70brLJvSLwTFqMubBQuj_GVnU,4180 -pygments/lexers/cplint.py,sha256=DkbyE5EKydLgf6BRr1FhQrK-IeQPL7Zmjk0DVdlRFnQ,1389 -pygments/lexers/crystal.py,sha256=xU-RnpIkpjrquoxtOuOcP8fcesSJl4xhU7kO9m42LZY,15754 -pygments/lexers/csound.py,sha256=ioSw4Q04wdwjUAbnTZ1qLhUq1vxdWFxhh3QtEl5RAJc,16998 -pygments/lexers/css.py,sha256=JN1RBYsee-jrpHWrSmhN3TKc4TkOBn-_BEGpgTCzcqE,25376 -pygments/lexers/d.py,sha256=piOy0EJeiAwPHugiM3gVv0z7HNh3u2gZQoCUSASRbY4,9920 -pygments/lexers/dalvik.py,sha256=deFg2JPBktJ9mEGb9EgxNkmd6vaMjJFQVzUHo8NKIa8,4606 -pygments/lexers/data.py,sha256=o0x0SmB5ms_CPUPljEEEenOON4IQWn86DkwFjkJYCOg,27026 -pygments/lexers/dax.py,sha256=ASi73qmr7OA7cVZXF2GTYGt01Ly1vY8CgD_Pnpm8k-4,8098 -pygments/lexers/devicetree.py,sha256=RecSQCidt8DRE1QFCPUbwwR0hiRlNtsFihdGldeUn3k,4019 -pygments/lexers/diff.py,sha256=F6vxZ64wm5Nag_97de1H_3F700ZwCVnYjKvtT5jilww,5382 -pygments/lexers/dns.py,sha256=Hh5hJ7MXfrq36KgfyIRwK3X8o1LdR98IKERcV4eZ7HY,3891 -pygments/lexers/dotnet.py,sha256=NDE0kOmpe96GLO-zwNLazmj77E9ORGmKpa4ZMCXDXxQ,39441 -pygments/lexers/dsls.py,sha256=GnHKhGL5GxsRFnqC7-65NTPZLOZdmnllNrGP86x_fQE,36746 -pygments/lexers/dylan.py,sha256=7zZ1EbHWXeVHqTD36AqykKqo3fhuIh4sM-whcxUaH_Y,10409 -pygments/lexers/ecl.py,sha256=vhmpa2LBrHxsPkYcf3kPZ1ItVaLRDTebi186wY0xGZA,6371 -pygments/lexers/eiffel.py,sha256=5ydYIEFcgcMoEj4BlK31hZ0aJb8OX0RdAvuCNdlxwqw,2690 -pygments/lexers/elm.py,sha256=uRCddU8jK5vVkH6Y66y8KOsDJprIfrOgeYq3hv1PxAM,3152 -pygments/lexers/elpi.py,sha256=O9j_WKBPyvNFjCRuPciVpW4etVSnILm_T79BhCPZYmo,6877 -pygments/lexers/email.py,sha256=ZZL6yvwCRl1CEQyysuOu0lbabp5tjMutS7f3efFKGR4,4804 -pygments/lexers/erlang.py,sha256=bU11eVHvooLwmVknzN6Xkb2DMk7HbenqdNlYSzhThDM,19147 -pygments/lexers/esoteric.py,sha256=Jfp8UUKyKYsqLaqXRZT3GSM9dzkF65zduwfnH1GoGhU,10500 -pygments/lexers/ezhil.py,sha256=22r-xjvvBVpExTqCI-HycAwunDb1p5gY4tIfDmM0vDw,3272 -pygments/lexers/factor.py,sha256=urZ4En4uKFCLXdEkXLWg9EYUFGHQTTDCwNXtyq-ngok,19530 -pygments/lexers/fantom.py,sha256=JJ13-NwykD-iIESnuzCefCYeQDO95cHMJA8TasF4gHA,10231 -pygments/lexers/felix.py,sha256=F-v0si4zPtRelqzDQWXI1-tarCE-BvawziODxRU7378,9655 -pygments/lexers/fift.py,sha256=rOCwp3v5ocK5YOWvt7Td3Md--97_8e-7Sonx52uS8mA,1644 -pygments/lexers/floscript.py,sha256=aHh82k52jMuDuzl9LatrcSANJiXTCyjGU3SO53bwbb0,2667 -pygments/lexers/forth.py,sha256=ZMtsHdNbnS_0IdSYlfAlfTSPEr0MEsRo-YZriQNueTQ,7193 -pygments/lexers/fortran.py,sha256=1PE5dTxf4Df6LUeXFcmNtyeXWsC8tSiK5dYwPHIJeeQ,10382 -pygments/lexers/foxpro.py,sha256=CBkW62Fuibz3yfyelZCaEO8GGdFJWsuRhqwtsSeBwLM,26295 -pygments/lexers/freefem.py,sha256=LFBQk-m1-nNCgrl-VDH3QwnVWurvb7W29i06LoT207A,26913 -pygments/lexers/func.py,sha256=OR2rkM7gf9fKvad5WcFQln-_U_pb-RUCM9eQatToF4A,3700 -pygments/lexers/functional.py,sha256=fYT2AGZ642cRkIAId0rnXFBsx1c8LLEDRN_VuCEkUyM,693 -pygments/lexers/futhark.py,sha256=Vf1i4t-tR3zqaktVjhTzFNg_ts_9CcyA4ZDfDizbCmk,3743 -pygments/lexers/gcodelexer.py,sha256=4Xs9ax4-JZGupW_qSnHon39wQGpb-tNA3xorMKg841E,874 -pygments/lexers/gdscript.py,sha256=Ws7JKxy0M0IyZ_1iMfRvJPrizEwmeCNLDoeMIFaM-CU,7566 -pygments/lexers/gleam.py,sha256=XIlTcq6cB743pCqbNYo8PocSkjZyDPR6hHgdaJNJ1Vc,2392 -pygments/lexers/go.py,sha256=4LezefgyuqZWHzLZHieUkKTi-ssY6aHJxx7Z-LFaLK0,3783 -pygments/lexers/grammar_notation.py,sha256=LvzhRQHgwZzq9oceukZS_hwnKK58ee7Z5d0cwXOR734,8043 -pygments/lexers/graph.py,sha256=WFqoPA1c_hHYrV0i_F7-eUw3Co4_HmZY3GJ-TyDr670,4108 -pygments/lexers/graphics.py,sha256=tmF9NNALnvPnax8ywYC3pLOla45YXtp9UA0H-5EiTQY,39145 -pygments/lexers/graphql.py,sha256=O_zcrGrBaDaKTlUoJGRruxqk7CJi-NR92Y0Cs-KkCvw,5601 -pygments/lexers/graphviz.py,sha256=mzdXOMpwz9_V-be1eTAMyhkKCBl6UxCIXuq6C2yrtsw,1934 -pygments/lexers/gsql.py,sha256=VPZk9sb26-DumRkWfEaSTeoc0lx5xt5n-6eDDLezMtc,3990 -pygments/lexers/hare.py,sha256=PGCOuILktJsmtTpCZZKkMFtObfJuBpei8HM8HHuq1Tw,2649 -pygments/lexers/haskell.py,sha256=MYr74-PAC8kGJRX-dZmvZsHTc7a2u6yFS2B19LfDD7g,33262 -pygments/lexers/haxe.py,sha256=WHCy_nrXHnfLITfbdp3Ji3lqQU4HAsTUpXsLCp2_4sk,30974 -pygments/lexers/hdl.py,sha256=MOWxhmAuE4Ei0CKDqqaON7T8tl43geancrNYM136Z0U,22738 -pygments/lexers/hexdump.py,sha256=1lj9oJ-KiZXSVYvTMfGmEAQzNEW08WlMcC2I5aYvHK4,3653 -pygments/lexers/html.py,sha256=MxYTI4EeT7QxoGleCAyQq-8n_Sgly6tD95H5zanCNmk,21977 -pygments/lexers/idl.py,sha256=rcihUAGhfuGEaSW6pgFq6NzplT_pv0DagUoefg4zAmk,15449 -pygments/lexers/igor.py,sha256=wVefbUjb3ftaW3LCKGtX1JgLgiY4EmRor5gVOn8vQA8,31633 -pygments/lexers/inferno.py,sha256=ChE_5y5SLH_75Uv7D2dKWQMk2dlN6z1gY1IDjlJZ8rU,3135 -pygments/lexers/installers.py,sha256=ZHliit4Pxz1tYKOIjKkDXI5djTkpzYUMVIPR1xvUrL8,14435 -pygments/lexers/int_fiction.py,sha256=0ZzIa1sZDUQsltd1oHuS-BoNiOF8zKQfcVuDyK1Ttv8,56544 -pygments/lexers/iolang.py,sha256=L6dNDCLH0kxkIUi00fI4Z14QnRu79UcNDrgv02c5Zw8,1905 -pygments/lexers/j.py,sha256=DqNdwQGFLiZW3mCNLRg81gpmsy4Hgcai_9NP3LbWhNU,4853 -pygments/lexers/javascript.py,sha256=TGKQLSrCprCKfhLLGAq_0EOdvqvJKX9pOdKo7tCRurQ,63243 -pygments/lexers/jmespath.py,sha256=R5yA5LJ2nTIaDwnFIpSNGAThd0sAYFccwawA9xBptlg,2082 -pygments/lexers/jslt.py,sha256=OeYQf8O2_9FCaf9W6Q3a7rPdAFLthePCtVSgCrOTcl8,3700 -pygments/lexers/json5.py,sha256=8JZbc8EiTEZdKaIdQg3hXEh0mHWSzPlwd473a0nUuT0,2502 -pygments/lexers/jsonnet.py,sha256=bx2G6J4tJqGrJV1PyZrIWzWHXcoefCX-4lIxxtbn2gw,5636 -pygments/lexers/jsx.py,sha256=wGsoGSB40qAJrVfXwRPtan7OcK0O87RVsHHk0m6gogk,2693 -pygments/lexers/julia.py,sha256=0ZDJ9X83V5GqJzA6T6p0TTN8WHy2JAjvu-FSBXvfXdc,11710 -pygments/lexers/jvm.py,sha256=Yt1iQ3QodXRY-x_HUOGedhyuBBHn5jYH-I8NzOzHTlE,72667 -pygments/lexers/kuin.py,sha256=3dKKJVJlskgrvMKv2tY9NOsFfDjyo-3MLcJ1lFKdXSg,11405 -pygments/lexers/kusto.py,sha256=kaxkoPpEBDsBTCvCOkZZx7oGfv0jk_UNIRIRbfVAsBE,3477 -pygments/lexers/ldap.py,sha256=77vF4t_19x9V522cxRCM5d3HW8Ne3giYsFsMPVYYBw4,6551 -pygments/lexers/lean.py,sha256=7HWRgxFsxS1N9XKqw0vfKwaxl27s5YiVYtZeRUoTHFo,8570 -pygments/lexers/lilypond.py,sha256=yd2Tuv67um6EyCIr-VwBnlPhTHxMaQsBJ4nGgO5fjIk,9752 -pygments/lexers/lisp.py,sha256=EHUy1g4pzEsYPE-zGj2rAXm3YATE1j9dCQOr5-JPSkU,157668 -pygments/lexers/macaulay2.py,sha256=zkV-vxjQYa0Jj9TGfFP1iMgpTZ4ApQuAAIdJVGWb2is,33366 -pygments/lexers/make.py,sha256=YMI5DBCrxWca-pz9cVXcyfuHLcikPx9R_3pW_98Myqo,7831 -pygments/lexers/maple.py,sha256=Rs0dEmOMD3C1YQPd0mntN-vzReq4XfHegH6xV4lvJWo,7960 -pygments/lexers/markup.py,sha256=zWtxsyIx_1OxQzS6wLe8bEqglePv4RqvJjbia8AvV5c,65088 -pygments/lexers/math.py,sha256=P3ZK1ePd8ZnLdlmHezo2irCA8T2-nlHBoSaBoT5mEVI,695 -pygments/lexers/matlab.py,sha256=F9KO4qowIhfP8oVhCRRzE_1sqg4zmQbsB2NZH193PiM,133027 -pygments/lexers/maxima.py,sha256=a0h9Ggs9JEovTrzbJT-BLVbOqI29yPnaMZlkU5f_FeY,2715 -pygments/lexers/meson.py,sha256=BMrsDo6BH2lzTFw7JDwQ9SDNMTrRkXCNRDVf4aFHdsI,4336 -pygments/lexers/mime.py,sha256=yGrf3h37LK4b6ERBpFiL_qzn3JgOfGR5KLagnbWFl6c,7582 -pygments/lexers/minecraft.py,sha256=Nu88snDDPzM0D-742fFdUriczL-EE911pAd4_I4-pAw,13696 -pygments/lexers/mips.py,sha256=STKiZT67b3QERXXn7XKVxlPBu7vwbPC5EyCpuf3Jfbw,4656 -pygments/lexers/ml.py,sha256=t8sCv4BjvuBq6AihKKUwStEONIgdXCC2RMtO0RopNbM,35390 -pygments/lexers/modeling.py,sha256=M7B58bGB-Zwd1EmPxKqtRvg7TgNCyem3MVUHv0_H2SQ,13683 -pygments/lexers/modula2.py,sha256=NtpXBRoUCeHfflgB39LknSkCwhBHBKv2Er_pinjVsNE,53072 -pygments/lexers/mojo.py,sha256=8JRVoftN1E-W2woG0K-4n8PQXTUM9iY6Sl5sWb2uGNg,24233 -pygments/lexers/monte.py,sha256=baWU6zlXloenw9MO1MtEVGE9i3CfiXAYhqU621MIjRk,6289 -pygments/lexers/mosel.py,sha256=gjRdedhA1jTjoYoM1Gpaoog_I9o7TRbYMHk97N1TXwg,9297 -pygments/lexers/ncl.py,sha256=zJ6ahlitit4S0pBXc7Wu96PB7xOn59MwfR2HdY5_C60,63999 -pygments/lexers/nimrod.py,sha256=Q1NSqEkLC5wWt7xJyKC-vzWw_Iw2SfDNP_pyMFBuIfA,6413 -pygments/lexers/nit.py,sha256=p_hVD8GzMRl3CABVKHtYgnXFUQk0i5F2FbWFA6WXm6s,2725 -pygments/lexers/nix.py,sha256=NOrv20gdq-2A7eZ6c2gElPHv1Xx2pvv20-qOymL9GMg,4421 -pygments/lexers/numbair.py,sha256=fxkp2CXeXWKBMewfi1H4JSYkmm4kU58wZ2Sh9BDYAWQ,1758 -pygments/lexers/oberon.py,sha256=jw403qUUs7zpTHAs5CbLjb8qiuwtxLk0spDIYqGZwAw,4210 -pygments/lexers/objective.py,sha256=Fo1WB3JMj8sNeYnvB84H4_qwhOt4WNJtJWjVEOwrJGk,23297 -pygments/lexers/ooc.py,sha256=kD1XaJZaihDF_s-Vyu1Bx68S_9zFt2rhox7NF8LpOZM,3002 -pygments/lexers/openscad.py,sha256=h9I1k8kiuQmhX5vZm6VDSr2fa5Finy0sN8ZDIE-jx1c,3700 -pygments/lexers/other.py,sha256=WLVyqPsvm9oSXIbZwbfyJloS6HGgoFW5nVTaU1uQpTw,1763 -pygments/lexers/parasail.py,sha256=DWMGhtyQgGTXbIgQl_mID6CKqi-Dhbvs_dTkmvrZXfE,2719 -pygments/lexers/parsers.py,sha256=feNgxroPoWRf0NEsON2mtmKDUfslIQppukw6ndEsQ3M,26596 -pygments/lexers/pascal.py,sha256=N2tRAjlXnTxggAzzk2tOOAVzeC2MBzrXy97_HQl5n44,30989 -pygments/lexers/pawn.py,sha256=LWUYQYsebMMt2d5oxX1HYWvBqbakR1h7Av_z8Vw94Wg,8253 -pygments/lexers/pddl.py,sha256=Mk4_BzlROJCd0xR4KKRRSrbj0F7LLQcBRjmsmtWmrCg,2989 -pygments/lexers/perl.py,sha256=9BXn3tyHMA49NvzbM9E2czSCHjeU7bvaPLUcoZrhz-4,39192 -pygments/lexers/phix.py,sha256=hZqychqo5sFMBDESzDPXg1DYHQe_9sn294UfbjihaFk,23249 -pygments/lexers/php.py,sha256=l4hzQrlm0525i5dSw9Vmjcai3TzbPT6DkjzxPg9l6Zc,13061 -pygments/lexers/pointless.py,sha256=WSDjqQyGrNIGmTCdaMxl4zk7OZTlJAMzeUZ02kfgcTI,1974 -pygments/lexers/pony.py,sha256=EXrMkacqMZblI7v4AvBRQe-3Py8__bx5FOgjCLdfXxQ,3279 -pygments/lexers/praat.py,sha256=4UFK-nbC6WkZBhJgcQqEGqq9CocJkW7AmT_OJQbjWzk,12676 -pygments/lexers/procfile.py,sha256=05W2fyofLTP-FbEdSXD1eles-PPqVNfF6RWXjQdW2us,1155 -pygments/lexers/prolog.py,sha256=9Kc5YNUFqkfWu2sYoyzC3RX65abf1bm7oHr86z1s4kQ,12866 -pygments/lexers/promql.py,sha256=n-0vo-o8-ZasqP3Va4ujs562UfZSLfZF-RzT71yL0Tk,4738 -pygments/lexers/prql.py,sha256=PFReuvhbv4K5aeu6lvDfw4m-3hULkB3r43bKAy948os,8747 -pygments/lexers/ptx.py,sha256=KSHAvbiNVUntKilQ6EPYoLFocmJpRsBy_7fW6_Nrs1Y,4501 -pygments/lexers/python.py,sha256=WZe7fBAHKZ_BxPg8qIU26UGhk8qwUYyENJ3IyPW64mc,53805 -pygments/lexers/q.py,sha256=WQFUh3JrpK2j-VGW_Ytn3uJ5frUNmQIFnLtMVGRA9DI,6936 -pygments/lexers/qlik.py,sha256=2wqwdfIjrAz6RNBsP4MyeLX8Z7QpIGzxtf1CvaOlr_g,3693 -pygments/lexers/qvt.py,sha256=XMBnsWRrvCDf989OuDeb-KpszAkeETiACyaghZeL1ns,6103 -pygments/lexers/r.py,sha256=B6WgrD9SY1UTCV1fQBSlZbezPfpYsARn3FQIHcFYOiM,6474 -pygments/lexers/rdf.py,sha256=qUzxLna9v071bHhZAjdsBi8dKaJNk_h9g1ZRUAYCfoo,16056 -pygments/lexers/rebol.py,sha256=4u3N4kzui55HapopXDu3Kt0jczxDZ4buzwR7Mt4tQiM,18259 -pygments/lexers/rego.py,sha256=Rx5Gphbktr9ojg5DbqlyxHeQqqtF7g8W-oF0rmloDNY,1748 -pygments/lexers/resource.py,sha256=ioEzgWksB5HCjoz85XNkQPSd7n5kL0SZiuPkJP1hunQ,2927 -pygments/lexers/ride.py,sha256=kCWdxuR3PclVi4wiA0uUx4CYEFwuTqoMsKjhSW4X3yg,5035 -pygments/lexers/rita.py,sha256=Mj1QNxx1sWAZYC02kw8piVckaiw9B0MqQtiIiDFH0pA,1127 -pygments/lexers/rnc.py,sha256=g7ZD334PMGUqy_Ij64laSN1vJerwHqVkegfMCa3E-y8,1972 -pygments/lexers/roboconf.py,sha256=HbYuK5CqmQdd63SRY2nle01r7-p7mil0SnoauYDmEOY,2074 -pygments/lexers/robotframework.py,sha256=c4U1B9Q9ITBCTohqJTZOvkfyeVbenN4xhzSWIoZh5eU,18448 -pygments/lexers/ruby.py,sha256=uG617E5abBZcECRCqkhIfc-IbZcRb5cGuUZq_xpax90,22753 -pygments/lexers/rust.py,sha256=ZY-9vtsreBP0NfDd0WCouLSp_9MChAL8U8Abe-m9PB8,8260 -pygments/lexers/sas.py,sha256=C1Uz2s9DU6_s2kL-cB_PAGPtpyK5THlmhNmCumC1l48,9456 -pygments/lexers/savi.py,sha256=jrmruK0GnXktgBTWXW3oN3TXtofn3HBbkMlHnR84cko,4878 -pygments/lexers/scdoc.py,sha256=DXRmFDmYuc7h3gPAAVhfcL1OEbNBK5RdPpJqQzF3ZTk,2524 -pygments/lexers/scripting.py,sha256=eaYlkDK-_cAwTcCBHP6QXBCz8n6OzbhzdkRe0uV0xWY,81814 -pygments/lexers/sgf.py,sha256=w6C513ENaO2YCnqrduK7k03NaMDf-pgygvfzq2NaSRk,1985 -pygments/lexers/shell.py,sha256=dCS1zwkf5KwTog4__MnMC7h3Xmwv4_d3fnEV29tSwXI,36381 -pygments/lexers/sieve.py,sha256=eob-L84yf2jmhdNyYZUlbUJozdcd6GXcHW68lmAe8WE,2514 -pygments/lexers/slash.py,sha256=I-cRepmaxhL1SgYvD1hHX3gNBFI8NPszdU7hn1o5JlA,8484 -pygments/lexers/smalltalk.py,sha256=ue2PmqDK2sw0j75WdseiiENJBdZ1OwysH2Op1QN1r24,7204 -pygments/lexers/smithy.py,sha256=VREWoeuz7ANap_Uiopn7rs0Tnsfc-xBisDJKRGQY_y8,2659 -pygments/lexers/smv.py,sha256=He_VBSMbWONMWZmkrB5RYR0cfHVnMyKIXz68IFYl-a8,2805 -pygments/lexers/snobol.py,sha256=qDzb41xQQWMNmjB2MtZs23pFoFgZ2gbRZhK_Ir03r7I,2778 -pygments/lexers/solidity.py,sha256=Tixfnwku4Yezj6nNm8xVaw7EdV1qgAgdwahdTFP0St8,3163 -pygments/lexers/soong.py,sha256=Vm18vV4g6T8UPgjjY2yTRlSXGDpZowmuqQUBFfm4A9A,2339 -pygments/lexers/sophia.py,sha256=2YtYIT8iwAoW0B7TZuuoG_ZILhJV-2A7oBGat-98naE,3376 -pygments/lexers/special.py,sha256=8JuR2Vex8X-RWnC36S0HXTHWp2qmZclc90-TrLUWyaY,3585 -pygments/lexers/spice.py,sha256=m4nK0q4Sq_OFQez7kGWfki0No4ZV24YrONfHVj1Piqs,2790 -pygments/lexers/sql.py,sha256=WSG6vOsR87EEEwSQefP_Z7TauUG_BjqMHUFmPaSOVj4,41476 -pygments/lexers/srcinfo.py,sha256=B8vDs-sJogG3mWa5Hp_7JfHHUMyYRwGvKv6cKbFQXLM,1746 -pygments/lexers/stata.py,sha256=Zr9BC52D5O_3BbdW0N-tzoUmy0NTguL2sC-saXRVM-c,6415 -pygments/lexers/supercollider.py,sha256=_H5wDrn0DiGnlhB_cz6Rt_lo2TvqjSm0o6NPTd9R4Ko,3697 -pygments/lexers/tablegen.py,sha256=1JjedXYY18BNiY9JtNGLOtGfiwduNDZpQLBGTeQ6jAw,3987 -pygments/lexers/tact.py,sha256=X_lsxjFUMaC1TmYysXJq9tmAGifRnil83Bt1zA86Xdo,10809 -pygments/lexers/tal.py,sha256=xS9PlaWQOPj8MVr56fUNq31vUQKRWoLTlyWj9ZHm8AM,2904 -pygments/lexers/tcl.py,sha256=lK97ju4nikkt-oGOzIeyFEM98yq4dZSI8uEmYsq0R6c,5512 -pygments/lexers/teal.py,sha256=t3dqy_Arwv8_yExbX_xiFxv1TqJLPv4vh1MVKjKwS4Y,3522 -pygments/lexers/templates.py,sha256=BVdjYeoacIUuFyHTG39j4PxeNCe5E1oUURjH1rITrI4,75731 -pygments/lexers/teraterm.py,sha256=ciwztagW5Drg2gr17Qykrh6GwMsKy7e4xdQshX95GyQ,9718 -pygments/lexers/testing.py,sha256=YZgDgUEaLEYKSKEqpDsUi3Bn-Db_D42IlyiSsr1oX8U,10810 -pygments/lexers/text.py,sha256=nOCQPssIlKdVWU3PKxZiBPkf_KFM2V48IOssSyqhFY8,1068 -pygments/lexers/textedit.py,sha256=ttT4Ph-hIdgFLG6maRy_GskkziTFK0Wcg28yU0s6lek,7760 -pygments/lexers/textfmts.py,sha256=mi9KLEq4mrzDJbEc8G3VM-mSki_Tylkzodu47yH6z84,15524 -pygments/lexers/theorem.py,sha256=51ppBAEdhJmwU_lC916zMyjEoKLXqf89VAE_Lr0PNCc,17855 -pygments/lexers/thingsdb.py,sha256=x_fHNkLA-hIJyeIs6rg_X8n5OLYvFqaSu1FhI3apI5Y,6017 -pygments/lexers/tlb.py,sha256=ue2gqm45BI512lM13O8skAky9zAb7pLMrxZ8pbt5zRU,1450 -pygments/lexers/tls.py,sha256=_uQUVuMRDOhN-XUyGR5DIlVCk1CUZ1fIOSN4_WQYPKk,1540 -pygments/lexers/tnt.py,sha256=pK4LgoKON7u1xF66JYFncAPSbD8DZaeI_WTZ9HqEFlY,10456 -pygments/lexers/trafficscript.py,sha256=X3B8kgxS54ecuok9ic6Hkp-UMn5DvOmCK0p70Tz27Cw,1506 -pygments/lexers/typoscript.py,sha256=mBuePiVZUoAORPKsHwrx6fBWiy3fAIqG-2O67QmMiFI,8332 -pygments/lexers/typst.py,sha256=zIJBEhUXtWp5OiyAmvFA5m8d1EQG-ocwrJ677dvTUAk,7167 -pygments/lexers/ul4.py,sha256=rCaw0J9j3cdql9lX_HTilg65k9-9S118zOA6TAYfxaM,10499 -pygments/lexers/unicon.py,sha256=RAqoCnAAJBYOAGdR8ng0g6FtB39bGemLRlIqv5mcg9E,18625 -pygments/lexers/urbi.py,sha256=ajNP70NJg32jNnFDZsLvr_-4TToSGqRGkFyAPIJLfCU,6082 -pygments/lexers/usd.py,sha256=2eEGouolodYS402P_gtBrn4lLzpg1z8uHwPCKqjUb_k,3304 -pygments/lexers/varnish.py,sha256=dSh0Ku9SrjmlB29Fi_mWdWavN7M0cMKeepR4a34sOyI,7473 -pygments/lexers/verification.py,sha256=Qu433Q_h3EK3uS4bJoLRFZK0kIVwzX5AFKsa4Z-qnxA,3934 -pygments/lexers/verifpal.py,sha256=buyOOzCo_dGnoC40h0tthylHVVpgDt8qXu4olLvYy_4,2661 -pygments/lexers/vip.py,sha256=2lEV4cLV9p4E37wctBL7zkZ4ZU4p3HVsiLJFzB1bie0,5711 -pygments/lexers/vyper.py,sha256=Zq6sQIUBk6mBdpgOVgu3A6swGoBne0kDlRyjZznm2BY,5615 -pygments/lexers/web.py,sha256=4W9a7vcskrGJnxt4KmoE3SZydWB1qLq7lP2XS85J_m8,913 -pygments/lexers/webassembly.py,sha256=zgcMouzLawcbeFr6w_SOvGoUR68ZtqnnsbOcWEVleLk,5698 -pygments/lexers/webidl.py,sha256=ODtVmw4gVzI8HQWxuEckP6KMwm8WP2G2lSZEjagDXts,10516 -pygments/lexers/webmisc.py,sha256=-_-INDVdk47e2jlj-9bFcuLtntqVorBqIjlnwPfZFdI,40564 -pygments/lexers/wgsl.py,sha256=9igd9dzixGIgNewruv9mPnFms-c9BahkZcCCrZygv84,11880 -pygments/lexers/whiley.py,sha256=lMr750lA4MZsB4xqzVsIRtVMJIC3_dArhFYTHvOPwvA,4017 -pygments/lexers/wowtoc.py,sha256=8xxvf0xGeYtf4PE7KtkHZ_ly9xY_XXHrpCitdKE42Ro,4076 -pygments/lexers/wren.py,sha256=goGXnAMKKa13LLL40ybT3aMGPrk3gCRwZQFYAkKB_w0,3229 -pygments/lexers/x10.py,sha256=Q-AmgdF2E-N7mtOPpZ07CsxrTVnikyqC4uRRv6H75sk,1943 -pygments/lexers/xorg.py,sha256=9ttrBd3_Y2nXANsqtMposSgblYmMYqWXQ-Iz5RH9RsU,925 -pygments/lexers/yang.py,sha256=13CWbSaNr9giOHz4o0SXSklh0bfWt0ah14jJGpTvcn0,4499 -pygments/lexers/yara.py,sha256=jUSv78KTDfguCoAoAZKbYzQERkkyxBBWv5dInVrkDxo,2427 -pygments/lexers/zig.py,sha256=f-80MVOSp1KnczAMokQLVM-_wAEOD16EcGFnaCNlsN0,3976 -pygments/modeline.py,sha256=K5eSkR8GS1r5OkXXTHOcV0aM_6xpk9eWNEIAW-OOJ2g,1005 -pygments/plugin.py,sha256=tPx0rJCTIZ9ioRgLNYG4pifCbAwTRUZddvLw-NfAk2w,1891 -pygments/regexopt.py,sha256=wXaP9Gjp_hKAdnICqoDkRxAOQJSc4v3X6mcxx3z-TNs,3072 -pygments/scanner.py,sha256=nNcETRR1tRuiTaHmHSTTECVYFPcLf6mDZu1e4u91A9E,3092 -pygments/sphinxext.py,sha256=VEe_oHNgLoEGMHc2ROfbee2mF2PPREFyE6_m_JN5FvQ,7898 -pygments/style.py,sha256=Cpw9dCAyW3_JAwFRXOJXmtKb5ZwO2_5KSmlq6q4fZw4,6408 -pygments/styles/__init__.py,sha256=f9KCQXN4uKbe8aI8-L3qTC-_XPfT563FwTg6VTGVfwI,2006 -pygments/styles/__pycache__/__init__.cpython-310.pyc,, -pygments/styles/__pycache__/_mapping.cpython-310.pyc,, -pygments/styles/__pycache__/abap.cpython-310.pyc,, -pygments/styles/__pycache__/algol.cpython-310.pyc,, -pygments/styles/__pycache__/algol_nu.cpython-310.pyc,, -pygments/styles/__pycache__/arduino.cpython-310.pyc,, -pygments/styles/__pycache__/autumn.cpython-310.pyc,, -pygments/styles/__pycache__/borland.cpython-310.pyc,, -pygments/styles/__pycache__/bw.cpython-310.pyc,, -pygments/styles/__pycache__/coffee.cpython-310.pyc,, -pygments/styles/__pycache__/colorful.cpython-310.pyc,, -pygments/styles/__pycache__/default.cpython-310.pyc,, -pygments/styles/__pycache__/dracula.cpython-310.pyc,, -pygments/styles/__pycache__/emacs.cpython-310.pyc,, -pygments/styles/__pycache__/friendly.cpython-310.pyc,, -pygments/styles/__pycache__/friendly_grayscale.cpython-310.pyc,, -pygments/styles/__pycache__/fruity.cpython-310.pyc,, -pygments/styles/__pycache__/gh_dark.cpython-310.pyc,, -pygments/styles/__pycache__/gruvbox.cpython-310.pyc,, -pygments/styles/__pycache__/igor.cpython-310.pyc,, -pygments/styles/__pycache__/inkpot.cpython-310.pyc,, -pygments/styles/__pycache__/lightbulb.cpython-310.pyc,, -pygments/styles/__pycache__/lilypond.cpython-310.pyc,, -pygments/styles/__pycache__/lovelace.cpython-310.pyc,, -pygments/styles/__pycache__/manni.cpython-310.pyc,, -pygments/styles/__pycache__/material.cpython-310.pyc,, -pygments/styles/__pycache__/monokai.cpython-310.pyc,, -pygments/styles/__pycache__/murphy.cpython-310.pyc,, -pygments/styles/__pycache__/native.cpython-310.pyc,, -pygments/styles/__pycache__/nord.cpython-310.pyc,, -pygments/styles/__pycache__/onedark.cpython-310.pyc,, -pygments/styles/__pycache__/paraiso_dark.cpython-310.pyc,, -pygments/styles/__pycache__/paraiso_light.cpython-310.pyc,, -pygments/styles/__pycache__/pastie.cpython-310.pyc,, -pygments/styles/__pycache__/perldoc.cpython-310.pyc,, -pygments/styles/__pycache__/rainbow_dash.cpython-310.pyc,, -pygments/styles/__pycache__/rrt.cpython-310.pyc,, -pygments/styles/__pycache__/sas.cpython-310.pyc,, -pygments/styles/__pycache__/solarized.cpython-310.pyc,, -pygments/styles/__pycache__/staroffice.cpython-310.pyc,, -pygments/styles/__pycache__/stata_dark.cpython-310.pyc,, -pygments/styles/__pycache__/stata_light.cpython-310.pyc,, -pygments/styles/__pycache__/tango.cpython-310.pyc,, -pygments/styles/__pycache__/trac.cpython-310.pyc,, -pygments/styles/__pycache__/vim.cpython-310.pyc,, -pygments/styles/__pycache__/vs.cpython-310.pyc,, -pygments/styles/__pycache__/xcode.cpython-310.pyc,, -pygments/styles/__pycache__/zenburn.cpython-310.pyc,, -pygments/styles/_mapping.py,sha256=6lovFUE29tz6EsV3XYY4hgozJ7q1JL7cfO3UOlgnS8w,3312 -pygments/styles/abap.py,sha256=64Uwr8uPdEdcT-tE-Y2VveTXfH3SkqH9qdMgY49YHQI,749 -pygments/styles/algol.py,sha256=fCuk8ITTehvbJSufiaKlgnFsKbl-xFxxR82xhltc-cQ,2262 -pygments/styles/algol_nu.py,sha256=Gv9WfHJvYegGcUk1zcufQgsdXPNjCUNk8sAHyrSGGh4,2283 -pygments/styles/arduino.py,sha256=NoUB8xk7M1HGPoLfuySOLU0sVwoTuLcZqllXl2EO_iE,4557 -pygments/styles/autumn.py,sha256=fLLfjHXjxCl6crBAxEsBLH372ALMkFacA2bG6KFbJi4,2195 -pygments/styles/borland.py,sha256=_0ySKp4KGCSgtYjPe8uzD6gQhlmAIR4T43i-FoRYNOM,1611 -pygments/styles/bw.py,sha256=vhk8Xoj64fLPdA9IQU6mUVsYMel255jR-FDU7BjIHtI,1406 -pygments/styles/coffee.py,sha256=NqLt-fc7LONma1BGggbceVRY9uDE70WBuZXqK4zwaco,2308 -pygments/styles/colorful.py,sha256=mYcSbehtH7itH_QV9NqJp4Wna1X4lrwl2wkVXS2u-5A,2832 -pygments/styles/default.py,sha256=RTgG2zKWWUxPTDCFxhTnyZI_WZBIVgu5XsUpNvFisCA,2588 -pygments/styles/dracula.py,sha256=vRJmixBoSKV9o8NVQhXGViQqchhIYugfikLmvX0DoBw,2182 -pygments/styles/emacs.py,sha256=TiOG9oc83qToMCRMnJrXtWYqnzAqYycRz_50OoCKtxc,2535 -pygments/styles/friendly.py,sha256=oAi-l9anQTs9STDmUzXGDlOegatEOH4hpD0j6o6dZGM,2604 -pygments/styles/friendly_grayscale.py,sha256=a7Cqkzt6-uTiXvj6GoYBXzRvX5_zviCjjRB04Kf_-Q0,2828 -pygments/styles/fruity.py,sha256=GfSUTG0stlJr5Ow_saCaxbI2IB4-34Dp2TuRTpfUJBs,1324 -pygments/styles/gh_dark.py,sha256=ruNX3d4rf22rx-8HnwvGbNbXRQpXCNcHU1HNq6N4uNg,3590 -pygments/styles/gruvbox.py,sha256=KrFoHEoVnZW6XM9udyXncPomeGyZgIDsNWOH3kCrxFQ,3387 -pygments/styles/igor.py,sha256=fYYPhM0dRCvcDTMVrMVO5oFKnYm-8YVlsuVBoczFLtY,737 -pygments/styles/inkpot.py,sha256=jggSeX9NV15eOL2oJaVmZ6vmV7LWRzXJQRUqcWEqGRs,2404 -pygments/styles/lightbulb.py,sha256=Y8u1qdvlHfBqI2jJex55SkvVatVo_FjEUzE6h-X7m-0,3172 -pygments/styles/lilypond.py,sha256=Y6fp_sEL-zESmxAaMxzjtrKk90cuDC_DalNdC8wj0nw,2066 -pygments/styles/lovelace.py,sha256=cA9uhmbnzY04MccsiYSgMY7fvb4WMRbegWBUrGvXh1M,3178 -pygments/styles/manni.py,sha256=g9FyO7plTwfMm2cU4iiKgdlkMlvQLG6l2Lwkgz5ITS4,2443 -pygments/styles/material.py,sha256=LDmgomAbgtJDZhbv446_zIwgYh50UAqEEtgYNUns1rQ,4201 -pygments/styles/monokai.py,sha256=lrxTJpkBarV9gTLkBQryZ6oNSjekAVheJueKJP5iEYA,5184 -pygments/styles/murphy.py,sha256=-AKZiLkpiWej-otjHMsYCE-I-_IzCOLJY-_GBdKRZRw,2805 -pygments/styles/native.py,sha256=l6tezGSQTB8p_SyOXJ0PWI7KzCeEdtsPmVc4Yn4_CwU,2043 -pygments/styles/nord.py,sha256=GDt3WAaqaWsiCeqpIBPxd8TEUX708fGfwaA7S0w0oy0,5391 -pygments/styles/onedark.py,sha256=k80cZEppCEF-HLoxy_FEA0QmQDZze68nHVMNGyUVa28,1719 -pygments/styles/paraiso_dark.py,sha256=Jkrg4nUKIVNF8U4fPNV_Smq_g9NFbb9eiUrjYpVgQZg,5662 -pygments/styles/paraiso_light.py,sha256=MxN964ZEpze3wF0ss-igaa2I7E684MHe-Zq0rWPH3wo,5668 -pygments/styles/pastie.py,sha256=ZvAs9UpBNYFC-5PFrCRGYnm3FoPKb-eKR-ozbWZP-4g,2525 -pygments/styles/perldoc.py,sha256=HSxB93e4UpQkZspReQ34FeJbZ-59ksGvdaH-hToehi8,2230 -pygments/styles/rainbow_dash.py,sha256=4ugL18Or7aNtaLfPfCLFRiFy0Gu2RA4a9G2LQUE9SrM,2390 -pygments/styles/rrt.py,sha256=fgzfpC0PC_SCcLOMCNEIQTjPUMOncRe7SR10GfSRbXY,1006 -pygments/styles/sas.py,sha256=yzoXmbfQ2ND1WWq93b4vVGYkQSZHPqb4ymes9YYRT3w,1440 -pygments/styles/solarized.py,sha256=qupILFZn02WspnAF5SPYb-W8guo9xnUtjb1HeLw3XgE,4247 -pygments/styles/staroffice.py,sha256=CLbBeMoxay21Xyu3Af2p4xUXyG1_6ydCbvs5RJKYe5w,831 -pygments/styles/stata_dark.py,sha256=vX8SwHV__sG92F4CKribG08MJfSVq98dgs7gEA_n9yc,1257 -pygments/styles/stata_light.py,sha256=uV3GE-ylvffQ0yN3py1YAVqBB5wflIKZbceyK1Lqvrc,1289 -pygments/styles/tango.py,sha256=O2wcM4hHuU1Yt071M9CK7JPtiiSCqyxtT9tbiQICV28,7137 -pygments/styles/trac.py,sha256=9kMv1ZZyMKACWlx2fQVjRP0I2pgcRYCNrd7iGGZg9qk,1981 -pygments/styles/vim.py,sha256=J7_TqvrGkTX_XuTHW0In5wqPLAUPRWyr1122XueZWmM,2019 -pygments/styles/vs.py,sha256=s7YnzbIPuFU3LIke27mc4lAQSn2R3vbbHc1baMGSU_U,1130 -pygments/styles/xcode.py,sha256=PbQdzgGaA4a9LAU1i58alY9kM4IFlQX5jHQwOYmf_Rk,1504 -pygments/styles/zenburn.py,sha256=suZEKzBTCYdhf2cxNwcY7UATJK1tq5eYhGdBcXdf6MU,2203 -pygments/token.py,sha256=WbdWGhYm_Vosb0DDxW9lHNPgITXfWTsQmHt6cy9RbcM,6226 -pygments/unistring.py,sha256=al-_rBemRuGvinsrM6atNsHTmJ6DUbw24q2O2Ru1cBc,63208 -pygments/util.py,sha256=oRtSpiAo5jM9ulntkvVbgXUdiAW57jnuYGB7t9fYuhc,10031 diff --git a/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/WHEEL b/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/WHEEL deleted file mode 100644 index 12228d414b6cfed7c39d3781c85c63256a1d7fb5..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/WHEEL +++ /dev/null @@ -1,4 +0,0 @@ -Wheel-Version: 1.0 -Generator: hatchling 1.27.0 -Root-Is-Purelib: true -Tag: py3-none-any diff --git a/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/entry_points.txt b/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/entry_points.txt deleted file mode 100644 index 15498e35f53320bd5e1de176928daabf26be0109..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/entry_points.txt +++ /dev/null @@ -1,2 +0,0 @@ -[console_scripts] -pygmentize = pygments.cmdline:main diff --git a/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/licenses/AUTHORS b/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/licenses/AUTHORS deleted file mode 100644 index 811c66ae178b354c0a0eb9bd5e7a9ac63b4b056c..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/licenses/AUTHORS +++ /dev/null @@ -1,291 +0,0 @@ -Pygments is written and maintained by Georg Brandl . - -Major developers are Tim Hatch and Armin Ronacher -. - -Other contributors, listed alphabetically, are: - -* Sam Aaron -- Ioke lexer -* Jean Abou Samra -- LilyPond lexer -* João Abecasis -- JSLT lexer -* Ali Afshar -- image formatter -* Thomas Aglassinger -- Easytrieve, JCL, Rexx, Transact-SQL and VBScript - lexers -* Maxence Ahlouche -- PostgreSQL Explain lexer -* Muthiah Annamalai -- Ezhil lexer -* Nikolay Antipov -- OpenSCAD lexer -* Kumar Appaiah -- Debian control lexer -* Andreas Amann -- AppleScript lexer -* Timothy Armstrong -- Dart lexer fixes -* Jeffrey Arnold -- R/S, Rd, BUGS, Jags, and Stan lexers -* Eiríkr Åsheim -- Uxntal lexer -* Jeremy Ashkenas -- CoffeeScript lexer -* José Joaquín Atria -- Praat lexer -* Stefan Matthias Aust -- Smalltalk lexer -* Lucas Bajolet -- Nit lexer -* Ben Bangert -- Mako lexers -* Max Battcher -- Darcs patch lexer -* Thomas Baruchel -- APL lexer -* Tim Baumann -- (Literate) Agda lexer -* Paul Baumgart, 280 North, Inc. -- Objective-J lexer -* Michael Bayer -- Myghty lexers -* Thomas Beale -- Archetype lexers -* John Benediktsson -- Factor lexer -* David Benjamin, Google LLC -- TLS lexer -* Trevor Bergeron -- mIRC formatter -* Vincent Bernat -- LessCSS lexer -* Christopher Bertels -- Fancy lexer -* Sébastien Bigaret -- QVT Operational lexer -* Jarrett Billingsley -- MiniD lexer -* Adam Blinkinsop -- Haskell, Redcode lexers -* Stéphane Blondon -- Procfile, SGF and Sieve lexers -* Frits van Bommel -- assembler lexers -* Pierre Bourdon -- bugfixes -* Martijn Braam -- Kernel log lexer, BARE lexer -* JD Browne, Google LLC -- GoogleSQL lexer -* Matthias Bussonnier -- ANSI style handling for terminal-256 formatter -* chebee7i -- Python traceback lexer improvements -* Hiram Chirino -- Scaml and Jade lexers -* Mauricio Caceres -- SAS and Stata lexers. -* Michael Camilleri, John Gabriele, sogaiu -- Janet lexer -* Daren Chandisingh -- Gleam lexer -* Ian Cooper -- VGL lexer -* David Corbett -- Inform, Jasmin, JSGF, Snowball, and TADS 3 lexers -* Leaf Corcoran -- MoonScript lexer -* Fraser Cormack -- TableGen lexer -* Gabriel Corona -- ASN.1 lexer -* Christopher Creutzig -- MuPAD lexer -* Daniël W. Crompton -- Pike lexer -* Pete Curry -- bugfixes -* Bryan Davis -- EBNF lexer -* Bruno Deferrari -- Shen lexer -* Walter Dörwald -- UL4 lexer -* Luke Drummond -- Meson lexer -* Giedrius Dubinskas -- HTML formatter improvements -* Owen Durni -- Haxe lexer -* Alexander Dutton, Oxford University Computing Services -- SPARQL lexer -* James Edwards -- Terraform lexer -* Nick Efford -- Python 3 lexer -* Sven Efftinge -- Xtend lexer -* Artem Egorkine -- terminal256 formatter -* Matthew Fernandez -- CAmkES lexer -* Paweł Fertyk -- GDScript lexer, HTML formatter improvements -* Michael Ficarra -- CPSA lexer -* James H. Fisher -- PostScript lexer -* Amanda Fitch, Google LLC -- GoogleSQL lexer -* William S. Fulton -- SWIG lexer -* Carlos Galdino -- Elixir and Elixir Console lexers -* Michael Galloy -- IDL lexer -* Naveen Garg -- Autohotkey lexer -* Simon Garnotel -- FreeFem++ lexer -* Laurent Gautier -- R/S lexer -* Alex Gaynor -- PyPy log lexer -* Richard Gerkin -- Igor Pro lexer -* Alain Gilbert -- TypeScript lexer -* Alex Gilding -- BlitzBasic lexer -* GitHub, Inc -- DASM16, Augeas, TOML, and Slash lexers -* Bertrand Goetzmann -- Groovy lexer -* Krzysiek Goj -- Scala lexer -* Rostyslav Golda -- FloScript lexer -* Andrey Golovizin -- BibTeX lexers -* Matt Good -- Genshi, Cheetah lexers -* Michał Górny -- vim modeline support -* Alex Gosse -- TrafficScript lexer -* Patrick Gotthardt -- PHP namespaces support -* Hubert Gruniaux -- C and C++ lexer improvements -* Olivier Guibe -- Asymptote lexer -* Phil Hagelberg -- Fennel lexer -* Florian Hahn -- Boogie lexer -* Martin Harriman -- SNOBOL lexer -* Matthew Harrison -- SVG formatter -* Steven Hazel -- Tcl lexer -* Dan Michael Heggø -- Turtle lexer -* Aslak Hellesøy -- Gherkin lexer -* Greg Hendershott -- Racket lexer -* Justin Hendrick -- ParaSail lexer -* Jordi Gutiérrez Hermoso -- Octave lexer -* David Hess, Fish Software, Inc. -- Objective-J lexer -* Ken Hilton -- Typographic Number Theory and Arrow lexers -* Varun Hiremath -- Debian control lexer -* Rob Hoelz -- Perl 6 lexer -* Doug Hogan -- Mscgen lexer -* Ben Hollis -- Mason lexer -* Max Horn -- GAP lexer -* Fred Hornsey -- OMG IDL Lexer -* Alastair Houghton -- Lexer inheritance facility -* Tim Howard -- BlitzMax lexer -* Dustin Howett -- Logos lexer -* Ivan Inozemtsev -- Fantom lexer -* Hiroaki Itoh -- Shell console rewrite, Lexers for PowerShell session, - MSDOS session, BC, WDiff -* Brian R. Jackson -- Tea lexer -* Christian Jann -- ShellSession lexer -* Jonas Camillus Jeppesen -- Line numbers and line highlighting for - RTF-formatter -* Dennis Kaarsemaker -- sources.list lexer -* Dmitri Kabak -- Inferno Limbo lexer -* Igor Kalnitsky -- vhdl lexer -* Colin Kennedy - USD lexer -* Alexander Kit -- MaskJS lexer -* Pekka Klärck -- Robot Framework lexer -* Gerwin Klein -- Isabelle lexer -* Eric Knibbe -- Lasso lexer -* Stepan Koltsov -- Clay lexer -* Oliver Kopp - Friendly grayscale style -* Adam Koprowski -- Opa lexer -* Benjamin Kowarsch -- Modula-2 lexer -* Domen Kožar -- Nix lexer -* Oleh Krekel -- Emacs Lisp lexer -* Alexander Kriegisch -- Kconfig and AspectJ lexers -* Marek Kubica -- Scheme lexer -* Jochen Kupperschmidt -- Markdown processor -* Gerd Kurzbach -- Modelica lexer -* Jon Larimer, Google Inc. -- Smali lexer -* Olov Lassus -- Dart lexer -* Matt Layman -- TAP lexer -* Dan Lazin, Google LLC -- GoogleSQL lexer -* Kristian Lyngstøl -- Varnish lexers -* Sylvestre Ledru -- Scilab lexer -* Chee Sing Lee -- Flatline lexer -* Mark Lee -- Vala lexer -* Thomas Linder Puls -- Visual Prolog lexer -* Pete Lomax -- Phix lexer -* Valentin Lorentz -- C++ lexer improvements -* Ben Mabey -- Gherkin lexer -* Angus MacArthur -- QML lexer -* Louis Mandel -- X10 lexer -* Louis Marchand -- Eiffel lexer -* Simone Margaritelli -- Hybris lexer -* Tim Martin - World of Warcraft TOC lexer -* Kirk McDonald -- D lexer -* Gordon McGregor -- SystemVerilog lexer -* Stephen McKamey -- Duel/JBST lexer -* Brian McKenna -- F# lexer -* Charles McLaughlin -- Puppet lexer -* Kurt McKee -- Tera Term macro lexer, PostgreSQL updates, MySQL overhaul, JSON lexer -* Joe Eli McIlvain -- Savi lexer -* Lukas Meuser -- BBCode formatter, Lua lexer -* Cat Miller -- Pig lexer -* Paul Miller -- LiveScript lexer -* Hong Minhee -- HTTP lexer -* Michael Mior -- Awk lexer -* Bruce Mitchener -- Dylan lexer rewrite -* Reuben Morais -- SourcePawn lexer -* Jon Morton -- Rust lexer -* Paulo Moura -- Logtalk lexer -* Mher Movsisyan -- DTD lexer -* Dejan Muhamedagic -- Crmsh lexer -* Adrien Nayrat -- PostgreSQL Explain lexer -* Ana Nelson -- Ragel, ANTLR, R console lexers -* David Neto, Google LLC -- WebGPU Shading Language lexer -* Kurt Neufeld -- Markdown lexer -* Nam T. Nguyen -- Monokai style -* Jesper Noehr -- HTML formatter "anchorlinenos" -* Mike Nolta -- Julia lexer -* Avery Nortonsmith -- Pointless lexer -* Jonas Obrist -- BBCode lexer -* Edward O'Callaghan -- Cryptol lexer -* David Oliva -- Rebol lexer -* Pat Pannuto -- nesC lexer -* Jon Parise -- Protocol buffers and Thrift lexers -* Benjamin Peterson -- Test suite refactoring -* Ronny Pfannschmidt -- BBCode lexer -* Dominik Picheta -- Nimrod lexer -* Andrew Pinkham -- RTF Formatter Refactoring -* Clément Prévost -- UrbiScript lexer -* Tanner Prynn -- cmdline -x option and loading lexers from files -* Oleh Prypin -- Crystal lexer (based on Ruby lexer) -* Nick Psaris -- K and Q lexers -* Xidorn Quan -- Web IDL lexer -* Elias Rabel -- Fortran fixed form lexer -* raichoo -- Idris lexer -* Daniel Ramirez -- GDScript lexer -* Kashif Rasul -- CUDA lexer -* Nathan Reed -- HLSL lexer -* Justin Reidy -- MXML lexer -* Jonathon Reinhart, Google LLC -- Soong lexer -* Norman Richards -- JSON lexer -* Corey Richardson -- Rust lexer updates -* Fabrizio Riguzzi -- cplint leder -* Lubomir Rintel -- GoodData MAQL and CL lexers -* Andre Roberge -- Tango style -* Georg Rollinger -- HSAIL lexer -* Michiel Roos -- TypoScript lexer -* Konrad Rudolph -- LaTeX formatter enhancements -* Mario Ruggier -- Evoque lexers -* Miikka Salminen -- Lovelace style, Hexdump lexer, lexer enhancements -* Stou Sandalski -- NumPy, FORTRAN, tcsh and XSLT lexers -* Matteo Sasso -- Common Lisp lexer -* Joe Schafer -- Ada lexer -* Max Schillinger -- TiddlyWiki5 lexer -* Andrew Schmidt -- X++ lexer -* Ken Schutte -- Matlab lexers -* René Schwaiger -- Rainbow Dash style -* Sebastian Schweizer -- Whiley lexer -* Tassilo Schweyer -- Io, MOOCode lexers -* Pablo Seminario -- PromQL lexer -* Ted Shaw -- AutoIt lexer -* Joerg Sieker -- ABAP lexer -* Robert Simmons -- Standard ML lexer -* Kirill Simonov -- YAML lexer -* Corbin Simpson -- Monte lexer -* Ville Skyttä -- ASCII armored lexer -* Alexander Smishlajev -- Visual FoxPro lexer -* Steve Spigarelli -- XQuery lexer -* Jerome St-Louis -- eC lexer -* Camil Staps -- Clean and NuSMV lexers; Solarized style -* James Strachan -- Kotlin lexer -* Tom Stuart -- Treetop lexer -* Colin Sullivan -- SuperCollider lexer -* Ben Swift -- Extempore lexer -* tatt61880 -- Kuin lexer -* Edoardo Tenani -- Arduino lexer -* Tiberius Teng -- default style overhaul -* Jeremy Thurgood -- Erlang, Squid config lexers -* Brian Tiffin -- OpenCOBOL lexer -* Bob Tolbert -- Hy lexer -* Doug Torrance -- Macaulay2 lexer -* Matthias Trute -- Forth lexer -* Tuoa Spi T4 -- Bdd lexer -* Erick Tryzelaar -- Felix lexer -* Alexander Udalov -- Kotlin lexer improvements -* Thomas Van Doren -- Chapel lexer -* Dave Van Ee -- Uxntal lexer updates -* Daniele Varrazzo -- PostgreSQL lexers -* Abe Voelker -- OpenEdge ABL lexer -* Pepijn de Vos -- HTML formatter CTags support -* Matthias Vallentin -- Bro lexer -* Benoît Vinot -- AMPL lexer -* Linh Vu Hong -- RSL lexer -* Taavi Väänänen -- Debian control lexer -* Immanuel Washington -- Smithy lexer -* Nathan Weizenbaum -- Haml and Sass lexers -* Nathan Whetsell -- Csound lexers -* Dietmar Winkler -- Modelica lexer -* Nils Winter -- Smalltalk lexer -* Davy Wybiral -- Clojure lexer -* Whitney Young -- ObjectiveC lexer -* Diego Zamboni -- CFengine3 lexer -* Enrique Zamudio -- Ceylon lexer -* Alex Zimin -- Nemerle lexer -* Rob Zimmerman -- Kal lexer -* Evgenii Zheltonozhskii -- Maple lexer -* Vincent Zurczak -- Roboconf lexer -* Hubert Gruniaux -- C and C++ lexer improvements -* Thomas Symalla -- AMDGPU Lexer -* 15b3 -- Image Formatter improvements -* Fabian Neumann -- CDDL lexer -* Thomas Duboucher -- CDDL lexer -* Philipp Imhof -- Pango Markup formatter -* Thomas Voss -- Sed lexer -* Martin Fischer -- WCAG contrast testing -* Marc Auberer -- Spice lexer -* Amr Hesham -- Carbon lexer -* diskdance -- Wikitext lexer -* vanillajonathan -- PRQL lexer -* Nikolay Antipov -- OpenSCAD lexer -* Markus Meyer, Nextron Systems -- YARA lexer -* Hannes Römer -- Mojo lexer -* Jan Frederik Schaefer -- PDDL lexer - -Many thanks for all contributions! diff --git a/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/licenses/LICENSE b/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/licenses/LICENSE deleted file mode 100644 index 446a1a805c8a949579fb8a9799f2bec7777349dc..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments-2.19.2.dist-info/licenses/LICENSE +++ /dev/null @@ -1,25 +0,0 @@ -Copyright (c) 2006-2022 by the respective authors (see AUTHORS file). -All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are -met: - -* Redistributions of source code must retain the above copyright - notice, this list of conditions and the following disclaimer. - -* Redistributions in binary form must reproduce the above copyright - notice, this list of conditions and the following disclaimer in the - documentation and/or other materials provided with the distribution. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS -"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT -LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR -A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT -OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, -SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT -LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, -DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY -THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT -(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/venv/lib/python3.10/site-packages/pygments/__init__.py b/venv/lib/python3.10/site-packages/pygments/__init__.py deleted file mode 100644 index 2a391c3e47997637afeea6a1ec819e5e035514ff..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments/__init__.py +++ /dev/null @@ -1,82 +0,0 @@ -""" - Pygments - ~~~~~~~~ - - Pygments is a syntax highlighting package written in Python. - - It is a generic syntax highlighter for general use in all kinds of software - such as forum systems, wikis or other applications that need to prettify - source code. Highlights are: - - * a wide range of common languages and markup formats is supported - * special attention is paid to details, increasing quality by a fair amount - * support for new languages and formats are added easily - * a number of output formats, presently HTML, LaTeX, RTF, SVG, all image - formats that PIL supports, and ANSI sequences - * it is usable as a command-line tool and as a library - * ... and it highlights even Brainfuck! - - The `Pygments master branch`_ is installable with ``easy_install Pygments==dev``. - - .. _Pygments master branch: - https://github.com/pygments/pygments/archive/master.zip#egg=Pygments-dev - - :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS. - :license: BSD, see LICENSE for details. -""" -from io import StringIO, BytesIO - -__version__ = '2.19.2' -__docformat__ = 'restructuredtext' - -__all__ = ['lex', 'format', 'highlight'] - - -def lex(code, lexer): - """ - Lex `code` with the `lexer` (must be a `Lexer` instance) - and return an iterable of tokens. Currently, this only calls - `lexer.get_tokens()`. - """ - try: - return lexer.get_tokens(code) - except TypeError: - # Heuristic to catch a common mistake. - from pygments.lexer import RegexLexer - if isinstance(lexer, type) and issubclass(lexer, RegexLexer): - raise TypeError('lex() argument must be a lexer instance, ' - 'not a class') - raise - - -def format(tokens, formatter, outfile=None): # pylint: disable=redefined-builtin - """ - Format ``tokens`` (an iterable of tokens) with the formatter ``formatter`` - (a `Formatter` instance). - - If ``outfile`` is given and a valid file object (an object with a - ``write`` method), the result will be written to it, otherwise it - is returned as a string. - """ - try: - if not outfile: - realoutfile = getattr(formatter, 'encoding', None) and BytesIO() or StringIO() - formatter.format(tokens, realoutfile) - return realoutfile.getvalue() - else: - formatter.format(tokens, outfile) - except TypeError: - # Heuristic to catch a common mistake. - from pygments.formatter import Formatter - if isinstance(formatter, type) and issubclass(formatter, Formatter): - raise TypeError('format() argument must be a formatter instance, ' - 'not a class') - raise - - -def highlight(code, lexer, formatter, outfile=None): - """ - This is the most high-level highlighting function. It combines `lex` and - `format` in one function. - """ - return format(lex(code, lexer), formatter, outfile) diff --git a/venv/lib/python3.10/site-packages/pygments/__main__.py b/venv/lib/python3.10/site-packages/pygments/__main__.py deleted file mode 100644 index 4890a6c768509d85d614e6953388e6cda6162bb6..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments/__main__.py +++ /dev/null @@ -1,17 +0,0 @@ -""" - pygments.__main__ - ~~~~~~~~~~~~~~~~~ - - Main entry point for ``python -m pygments``. - - :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS. - :license: BSD, see LICENSE for details. -""" - -import sys -import pygments.cmdline - -try: - sys.exit(pygments.cmdline.main(sys.argv)) -except KeyboardInterrupt: - sys.exit(1) diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 43e3e7cb09afc76e97ecc4dfaa4d09604fc239ad..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/__main__.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/__main__.cpython-310.pyc deleted file mode 100644 index 90b0832369e85ffb4f538895b893736dd8093750..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/__main__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/cmdline.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/cmdline.cpython-310.pyc deleted file mode 100644 index e2805413bda3e0bf1f948fd9de15a6ac31d494c5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/cmdline.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/console.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/console.cpython-310.pyc deleted file mode 100644 index 4b944d900fd364d5830b66fc0b5c88a4734e7c6d..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/console.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/filter.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/filter.cpython-310.pyc deleted file mode 100644 index cb9a13954ab5eb5cd226448806d491435dbf839e..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/filter.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/formatter.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/formatter.cpython-310.pyc deleted file mode 100644 index 1a6a17a78fe6d07c59f60ea7b8dbdd7b8959c3f3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/formatter.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/lexer.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/lexer.cpython-310.pyc deleted file mode 100644 index 53133754148811d400f5320f58fe45754b23f56b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/lexer.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/modeline.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/modeline.cpython-310.pyc deleted file mode 100644 index cc21470c740d54a46c87384eb2382a7b46bdc7e4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/modeline.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/plugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/plugin.cpython-310.pyc deleted file mode 100644 index c468f830b3fbc36a9ba0fabc7f45632095b22345..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/plugin.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/regexopt.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/regexopt.cpython-310.pyc deleted file mode 100644 index e34bfe035ab22f90f0720cd885dd7c0d74e15e30..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/regexopt.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/scanner.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/scanner.cpython-310.pyc deleted file mode 100644 index ffe3352bfb6af0e692710248c98f3fc409885cc0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/scanner.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/sphinxext.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/sphinxext.cpython-310.pyc deleted file mode 100644 index 699bc822ec431fff21bf2398db84eb94911e3d5f..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/sphinxext.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/style.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/style.cpython-310.pyc deleted file mode 100644 index defe6e405327e747f3cbc4cb08d25e1c32f6af01..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/style.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/token.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/token.cpython-310.pyc deleted file mode 100644 index af93e3f8fa71199a9c98e65489cef1f4e4fdb9ba..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/token.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/unistring.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/unistring.cpython-310.pyc deleted file mode 100644 index a773b688444c0262de2d12c9d32b427046698ff3..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/unistring.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/__pycache__/util.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/__pycache__/util.cpython-310.pyc deleted file mode 100644 index ec3e1d57c7f96dfe0ded6b244b9a9f42cd270bf2..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/__pycache__/util.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/cmdline.py b/venv/lib/python3.10/site-packages/pygments/cmdline.py deleted file mode 100644 index 2878fd55a7924d38ed979e26f565539d6bae621a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments/cmdline.py +++ /dev/null @@ -1,668 +0,0 @@ -""" - pygments.cmdline - ~~~~~~~~~~~~~~~~ - - Command line interface. - - :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS. - :license: BSD, see LICENSE for details. -""" - -import os -import sys -import shutil -import argparse -from textwrap import dedent - -from pygments import __version__, highlight -from pygments.util import ClassNotFound, OptionError, docstring_headline, \ - guess_decode, guess_decode_from_terminal, terminal_encoding, \ - UnclosingTextIOWrapper -from pygments.lexers import get_all_lexers, get_lexer_by_name, guess_lexer, \ - load_lexer_from_file, get_lexer_for_filename, find_lexer_class_for_filename -from pygments.lexers.special import TextLexer -from pygments.formatters.latex import LatexEmbeddedLexer, LatexFormatter -from pygments.formatters import get_all_formatters, get_formatter_by_name, \ - load_formatter_from_file, get_formatter_for_filename, find_formatter_class -from pygments.formatters.terminal import TerminalFormatter -from pygments.formatters.terminal256 import Terminal256Formatter, TerminalTrueColorFormatter -from pygments.filters import get_all_filters, find_filter_class -from pygments.styles import get_all_styles, get_style_by_name - - -def _parse_options(o_strs): - opts = {} - if not o_strs: - return opts - for o_str in o_strs: - if not o_str.strip(): - continue - o_args = o_str.split(',') - for o_arg in o_args: - o_arg = o_arg.strip() - try: - o_key, o_val = o_arg.split('=', 1) - o_key = o_key.strip() - o_val = o_val.strip() - except ValueError: - opts[o_arg] = True - else: - opts[o_key] = o_val - return opts - - -def _parse_filters(f_strs): - filters = [] - if not f_strs: - return filters - for f_str in f_strs: - if ':' in f_str: - fname, fopts = f_str.split(':', 1) - filters.append((fname, _parse_options([fopts]))) - else: - filters.append((f_str, {})) - return filters - - -def _print_help(what, name): - try: - if what == 'lexer': - cls = get_lexer_by_name(name) - print(f"Help on the {cls.name} lexer:") - print(dedent(cls.__doc__)) - elif what == 'formatter': - cls = find_formatter_class(name) - print(f"Help on the {cls.name} formatter:") - print(dedent(cls.__doc__)) - elif what == 'filter': - cls = find_filter_class(name) - print(f"Help on the {name} filter:") - print(dedent(cls.__doc__)) - return 0 - except (AttributeError, ValueError): - print(f"{what} not found!", file=sys.stderr) - return 1 - - -def _print_list(what): - if what == 'lexer': - print() - print("Lexers:") - print("~~~~~~~") - - info = [] - for fullname, names, exts, _ in get_all_lexers(): - tup = (', '.join(names)+':', fullname, - exts and '(filenames ' + ', '.join(exts) + ')' or '') - info.append(tup) - info.sort() - for i in info: - print(('* {}\n {} {}').format(*i)) - - elif what == 'formatter': - print() - print("Formatters:") - print("~~~~~~~~~~~") - - info = [] - for cls in get_all_formatters(): - doc = docstring_headline(cls) - tup = (', '.join(cls.aliases) + ':', doc, cls.filenames and - '(filenames ' + ', '.join(cls.filenames) + ')' or '') - info.append(tup) - info.sort() - for i in info: - print(('* {}\n {} {}').format(*i)) - - elif what == 'filter': - print() - print("Filters:") - print("~~~~~~~~") - - for name in get_all_filters(): - cls = find_filter_class(name) - print("* " + name + ':') - print(f" {docstring_headline(cls)}") - - elif what == 'style': - print() - print("Styles:") - print("~~~~~~~") - - for name in get_all_styles(): - cls = get_style_by_name(name) - print("* " + name + ':') - print(f" {docstring_headline(cls)}") - - -def _print_list_as_json(requested_items): - import json - result = {} - if 'lexer' in requested_items: - info = {} - for fullname, names, filenames, mimetypes in get_all_lexers(): - info[fullname] = { - 'aliases': names, - 'filenames': filenames, - 'mimetypes': mimetypes - } - result['lexers'] = info - - if 'formatter' in requested_items: - info = {} - for cls in get_all_formatters(): - doc = docstring_headline(cls) - info[cls.name] = { - 'aliases': cls.aliases, - 'filenames': cls.filenames, - 'doc': doc - } - result['formatters'] = info - - if 'filter' in requested_items: - info = {} - for name in get_all_filters(): - cls = find_filter_class(name) - info[name] = { - 'doc': docstring_headline(cls) - } - result['filters'] = info - - if 'style' in requested_items: - info = {} - for name in get_all_styles(): - cls = get_style_by_name(name) - info[name] = { - 'doc': docstring_headline(cls) - } - result['styles'] = info - - json.dump(result, sys.stdout) - -def main_inner(parser, argns): - if argns.help: - parser.print_help() - return 0 - - if argns.V: - print(f'Pygments version {__version__}, (c) 2006-2024 by Georg Brandl, Matthäus ' - 'Chajdas and contributors.') - return 0 - - def is_only_option(opt): - return not any(v for (k, v) in vars(argns).items() if k != opt) - - # handle ``pygmentize -L`` - if argns.L is not None: - arg_set = set() - for k, v in vars(argns).items(): - if v: - arg_set.add(k) - - arg_set.discard('L') - arg_set.discard('json') - - if arg_set: - parser.print_help(sys.stderr) - return 2 - - # print version - if not argns.json: - main(['', '-V']) - allowed_types = {'lexer', 'formatter', 'filter', 'style'} - largs = [arg.rstrip('s') for arg in argns.L] - if any(arg not in allowed_types for arg in largs): - parser.print_help(sys.stderr) - return 0 - if not largs: - largs = allowed_types - if not argns.json: - for arg in largs: - _print_list(arg) - else: - _print_list_as_json(largs) - return 0 - - # handle ``pygmentize -H`` - if argns.H: - if not is_only_option('H'): - parser.print_help(sys.stderr) - return 2 - what, name = argns.H - if what not in ('lexer', 'formatter', 'filter'): - parser.print_help(sys.stderr) - return 2 - return _print_help(what, name) - - # parse -O options - parsed_opts = _parse_options(argns.O or []) - - # parse -P options - for p_opt in argns.P or []: - try: - name, value = p_opt.split('=', 1) - except ValueError: - parsed_opts[p_opt] = True - else: - parsed_opts[name] = value - - # encodings - inencoding = parsed_opts.get('inencoding', parsed_opts.get('encoding')) - outencoding = parsed_opts.get('outencoding', parsed_opts.get('encoding')) - - # handle ``pygmentize -N`` - if argns.N: - lexer = find_lexer_class_for_filename(argns.N) - if lexer is None: - lexer = TextLexer - - print(lexer.aliases[0]) - return 0 - - # handle ``pygmentize -C`` - if argns.C: - inp = sys.stdin.buffer.read() - try: - lexer = guess_lexer(inp, inencoding=inencoding) - except ClassNotFound: - lexer = TextLexer - - print(lexer.aliases[0]) - return 0 - - # handle ``pygmentize -S`` - S_opt = argns.S - a_opt = argns.a - if S_opt is not None: - f_opt = argns.f - if not f_opt: - parser.print_help(sys.stderr) - return 2 - if argns.l or argns.INPUTFILE: - parser.print_help(sys.stderr) - return 2 - - try: - parsed_opts['style'] = S_opt - fmter = get_formatter_by_name(f_opt, **parsed_opts) - except ClassNotFound as err: - print(err, file=sys.stderr) - return 1 - - print(fmter.get_style_defs(a_opt or '')) - return 0 - - # if no -S is given, -a is not allowed - if argns.a is not None: - parser.print_help(sys.stderr) - return 2 - - # parse -F options - F_opts = _parse_filters(argns.F or []) - - # -x: allow custom (eXternal) lexers and formatters - allow_custom_lexer_formatter = bool(argns.x) - - # select lexer - lexer = None - - # given by name? - lexername = argns.l - if lexername: - # custom lexer, located relative to user's cwd - if allow_custom_lexer_formatter and '.py' in lexername: - try: - filename = None - name = None - if ':' in lexername: - filename, name = lexername.rsplit(':', 1) - - if '.py' in name: - # This can happen on Windows: If the lexername is - # C:\lexer.py -- return to normal load path in that case - name = None - - if filename and name: - lexer = load_lexer_from_file(filename, name, - **parsed_opts) - else: - lexer = load_lexer_from_file(lexername, **parsed_opts) - except ClassNotFound as err: - print('Error:', err, file=sys.stderr) - return 1 - else: - try: - lexer = get_lexer_by_name(lexername, **parsed_opts) - except (OptionError, ClassNotFound) as err: - print('Error:', err, file=sys.stderr) - return 1 - - # read input code - code = None - - if argns.INPUTFILE: - if argns.s: - print('Error: -s option not usable when input file specified', - file=sys.stderr) - return 2 - - infn = argns.INPUTFILE - try: - with open(infn, 'rb') as infp: - code = infp.read() - except Exception as err: - print('Error: cannot read infile:', err, file=sys.stderr) - return 1 - if not inencoding: - code, inencoding = guess_decode(code) - - # do we have to guess the lexer? - if not lexer: - try: - lexer = get_lexer_for_filename(infn, code, **parsed_opts) - except ClassNotFound as err: - if argns.g: - try: - lexer = guess_lexer(code, **parsed_opts) - except ClassNotFound: - lexer = TextLexer(**parsed_opts) - else: - print('Error:', err, file=sys.stderr) - return 1 - except OptionError as err: - print('Error:', err, file=sys.stderr) - return 1 - - elif not argns.s: # treat stdin as full file (-s support is later) - # read code from terminal, always in binary mode since we want to - # decode ourselves and be tolerant with it - code = sys.stdin.buffer.read() # use .buffer to get a binary stream - if not inencoding: - code, inencoding = guess_decode_from_terminal(code, sys.stdin) - # else the lexer will do the decoding - if not lexer: - try: - lexer = guess_lexer(code, **parsed_opts) - except ClassNotFound: - lexer = TextLexer(**parsed_opts) - - else: # -s option needs a lexer with -l - if not lexer: - print('Error: when using -s a lexer has to be selected with -l', - file=sys.stderr) - return 2 - - # process filters - for fname, fopts in F_opts: - try: - lexer.add_filter(fname, **fopts) - except ClassNotFound as err: - print('Error:', err, file=sys.stderr) - return 1 - - # select formatter - outfn = argns.o - fmter = argns.f - if fmter: - # custom formatter, located relative to user's cwd - if allow_custom_lexer_formatter and '.py' in fmter: - try: - filename = None - name = None - if ':' in fmter: - # Same logic as above for custom lexer - filename, name = fmter.rsplit(':', 1) - - if '.py' in name: - name = None - - if filename and name: - fmter = load_formatter_from_file(filename, name, - **parsed_opts) - else: - fmter = load_formatter_from_file(fmter, **parsed_opts) - except ClassNotFound as err: - print('Error:', err, file=sys.stderr) - return 1 - else: - try: - fmter = get_formatter_by_name(fmter, **parsed_opts) - except (OptionError, ClassNotFound) as err: - print('Error:', err, file=sys.stderr) - return 1 - - if outfn: - if not fmter: - try: - fmter = get_formatter_for_filename(outfn, **parsed_opts) - except (OptionError, ClassNotFound) as err: - print('Error:', err, file=sys.stderr) - return 1 - try: - outfile = open(outfn, 'wb') - except Exception as err: - print('Error: cannot open outfile:', err, file=sys.stderr) - return 1 - else: - if not fmter: - if os.environ.get('COLORTERM','') in ('truecolor', '24bit'): - fmter = TerminalTrueColorFormatter(**parsed_opts) - elif '256' in os.environ.get('TERM', ''): - fmter = Terminal256Formatter(**parsed_opts) - else: - fmter = TerminalFormatter(**parsed_opts) - outfile = sys.stdout.buffer - - # determine output encoding if not explicitly selected - if not outencoding: - if outfn: - # output file? use lexer encoding for now (can still be None) - fmter.encoding = inencoding - else: - # else use terminal encoding - fmter.encoding = terminal_encoding(sys.stdout) - - # provide coloring under Windows, if possible - if not outfn and sys.platform in ('win32', 'cygwin') and \ - fmter.name in ('Terminal', 'Terminal256'): # pragma: no cover - # unfortunately colorama doesn't support binary streams on Py3 - outfile = UnclosingTextIOWrapper(outfile, encoding=fmter.encoding) - fmter.encoding = None - try: - import colorama.initialise - except ImportError: - pass - else: - outfile = colorama.initialise.wrap_stream( - outfile, convert=None, strip=None, autoreset=False, wrap=True) - - # When using the LaTeX formatter and the option `escapeinside` is - # specified, we need a special lexer which collects escaped text - # before running the chosen language lexer. - escapeinside = parsed_opts.get('escapeinside', '') - if len(escapeinside) == 2 and isinstance(fmter, LatexFormatter): - left = escapeinside[0] - right = escapeinside[1] - lexer = LatexEmbeddedLexer(left, right, lexer) - - # ... and do it! - if not argns.s: - # process whole input as per normal... - try: - highlight(code, lexer, fmter, outfile) - finally: - if outfn: - outfile.close() - return 0 - else: - # line by line processing of stdin (eg: for 'tail -f')... - try: - while 1: - line = sys.stdin.buffer.readline() - if not line: - break - if not inencoding: - line = guess_decode_from_terminal(line, sys.stdin)[0] - highlight(line, lexer, fmter, outfile) - if hasattr(outfile, 'flush'): - outfile.flush() - return 0 - except KeyboardInterrupt: # pragma: no cover - return 0 - finally: - if outfn: - outfile.close() - - -class HelpFormatter(argparse.HelpFormatter): - def __init__(self, prog, indent_increment=2, max_help_position=16, width=None): - if width is None: - try: - width = shutil.get_terminal_size().columns - 2 - except Exception: - pass - argparse.HelpFormatter.__init__(self, prog, indent_increment, - max_help_position, width) - - -def main(args=sys.argv): - """ - Main command line entry point. - """ - desc = "Highlight an input file and write the result to an output file." - parser = argparse.ArgumentParser(description=desc, add_help=False, - formatter_class=HelpFormatter) - - operation = parser.add_argument_group('Main operation') - lexersel = operation.add_mutually_exclusive_group() - lexersel.add_argument( - '-l', metavar='LEXER', - help='Specify the lexer to use. (Query names with -L.) If not ' - 'given and -g is not present, the lexer is guessed from the filename.') - lexersel.add_argument( - '-g', action='store_true', - help='Guess the lexer from the file contents, or pass through ' - 'as plain text if nothing can be guessed.') - operation.add_argument( - '-F', metavar='FILTER[:options]', action='append', - help='Add a filter to the token stream. (Query names with -L.) ' - 'Filter options are given after a colon if necessary.') - operation.add_argument( - '-f', metavar='FORMATTER', - help='Specify the formatter to use. (Query names with -L.) ' - 'If not given, the formatter is guessed from the output filename, ' - 'and defaults to the terminal formatter if the output is to the ' - 'terminal or an unknown file extension.') - operation.add_argument( - '-O', metavar='OPTION=value[,OPTION=value,...]', action='append', - help='Give options to the lexer and formatter as a comma-separated ' - 'list of key-value pairs. ' - 'Example: `-O bg=light,python=cool`.') - operation.add_argument( - '-P', metavar='OPTION=value', action='append', - help='Give a single option to the lexer and formatter - with this ' - 'you can pass options whose value contains commas and equal signs. ' - 'Example: `-P "heading=Pygments, the Python highlighter"`.') - operation.add_argument( - '-o', metavar='OUTPUTFILE', - help='Where to write the output. Defaults to standard output.') - - operation.add_argument( - 'INPUTFILE', nargs='?', - help='Where to read the input. Defaults to standard input.') - - flags = parser.add_argument_group('Operation flags') - flags.add_argument( - '-v', action='store_true', - help='Print a detailed traceback on unhandled exceptions, which ' - 'is useful for debugging and bug reports.') - flags.add_argument( - '-s', action='store_true', - help='Process lines one at a time until EOF, rather than waiting to ' - 'process the entire file. This only works for stdin, only for lexers ' - 'with no line-spanning constructs, and is intended for streaming ' - 'input such as you get from `tail -f`. ' - 'Example usage: `tail -f sql.log | pygmentize -s -l sql`.') - flags.add_argument( - '-x', action='store_true', - help='Allow custom lexers and formatters to be loaded from a .py file ' - 'relative to the current working directory. For example, ' - '`-l ./customlexer.py -x`. By default, this option expects a file ' - 'with a class named CustomLexer or CustomFormatter; you can also ' - 'specify your own class name with a colon (`-l ./lexer.py:MyLexer`). ' - 'Users should be very careful not to use this option with untrusted ' - 'files, because it will import and run them.') - flags.add_argument('--json', help='Output as JSON. This can ' - 'be only used in conjunction with -L.', - default=False, - action='store_true') - - special_modes_group = parser.add_argument_group( - 'Special modes - do not do any highlighting') - special_modes = special_modes_group.add_mutually_exclusive_group() - special_modes.add_argument( - '-S', metavar='STYLE -f formatter', - help='Print style definitions for STYLE for a formatter ' - 'given with -f. The argument given by -a is formatter ' - 'dependent.') - special_modes.add_argument( - '-L', nargs='*', metavar='WHAT', - help='List lexers, formatters, styles or filters -- ' - 'give additional arguments for the thing(s) you want to list ' - '(e.g. "styles"), or omit them to list everything.') - special_modes.add_argument( - '-N', metavar='FILENAME', - help='Guess and print out a lexer name based solely on the given ' - 'filename. Does not take input or highlight anything. If no specific ' - 'lexer can be determined, "text" is printed.') - special_modes.add_argument( - '-C', action='store_true', - help='Like -N, but print out a lexer name based solely on ' - 'a given content from standard input.') - special_modes.add_argument( - '-H', action='store', nargs=2, metavar=('NAME', 'TYPE'), - help='Print detailed help for the object of type , ' - 'where is one of "lexer", "formatter" or "filter".') - special_modes.add_argument( - '-V', action='store_true', - help='Print the package version.') - special_modes.add_argument( - '-h', '--help', action='store_true', - help='Print this help.') - special_modes_group.add_argument( - '-a', metavar='ARG', - help='Formatter-specific additional argument for the -S (print ' - 'style sheet) mode.') - - argns = parser.parse_args(args[1:]) - - try: - return main_inner(parser, argns) - except BrokenPipeError: - # someone closed our stdout, e.g. by quitting a pager. - return 0 - except Exception: - if argns.v: - print(file=sys.stderr) - print('*' * 65, file=sys.stderr) - print('An unhandled exception occurred while highlighting.', - file=sys.stderr) - print('Please report the whole traceback to the issue tracker at', - file=sys.stderr) - print('.', - file=sys.stderr) - print('*' * 65, file=sys.stderr) - print(file=sys.stderr) - raise - import traceback - info = traceback.format_exception(*sys.exc_info()) - msg = info[-1].strip() - if len(info) >= 3: - # extract relevant file and position info - msg += '\n (f{})'.format(info[-2].split('\n')[0].strip()[1:]) - print(file=sys.stderr) - print('*** Error while highlighting:', file=sys.stderr) - print(msg, file=sys.stderr) - print('*** If this is a bug you want to report, please rerun with -v.', - file=sys.stderr) - return 1 diff --git a/venv/lib/python3.10/site-packages/pygments/console.py b/venv/lib/python3.10/site-packages/pygments/console.py deleted file mode 100644 index ee1ac27a2ffdc119a1db9128a034300b31d464a8..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments/console.py +++ /dev/null @@ -1,70 +0,0 @@ -""" - pygments.console - ~~~~~~~~~~~~~~~~ - - Format colored console output. - - :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS. - :license: BSD, see LICENSE for details. -""" - -esc = "\x1b[" - -codes = {} -codes[""] = "" -codes["reset"] = esc + "39;49;00m" - -codes["bold"] = esc + "01m" -codes["faint"] = esc + "02m" -codes["standout"] = esc + "03m" -codes["underline"] = esc + "04m" -codes["blink"] = esc + "05m" -codes["overline"] = esc + "06m" - -dark_colors = ["black", "red", "green", "yellow", "blue", - "magenta", "cyan", "gray"] -light_colors = ["brightblack", "brightred", "brightgreen", "brightyellow", "brightblue", - "brightmagenta", "brightcyan", "white"] - -x = 30 -for dark, light in zip(dark_colors, light_colors): - codes[dark] = esc + "%im" % x - codes[light] = esc + "%im" % (60 + x) - x += 1 - -del dark, light, x - -codes["white"] = codes["bold"] - - -def reset_color(): - return codes["reset"] - - -def colorize(color_key, text): - return codes[color_key] + text + codes["reset"] - - -def ansiformat(attr, text): - """ - Format ``text`` with a color and/or some attributes:: - - color normal color - *color* bold color - _color_ underlined color - +color+ blinking color - """ - result = [] - if attr[:1] == attr[-1:] == '+': - result.append(codes['blink']) - attr = attr[1:-1] - if attr[:1] == attr[-1:] == '*': - result.append(codes['bold']) - attr = attr[1:-1] - if attr[:1] == attr[-1:] == '_': - result.append(codes['underline']) - attr = attr[1:-1] - result.append(codes[attr]) - result.append(text) - result.append(codes['reset']) - return ''.join(result) diff --git a/venv/lib/python3.10/site-packages/pygments/filter.py b/venv/lib/python3.10/site-packages/pygments/filter.py deleted file mode 100644 index 5efff438d2b3a187c681e445e4607b529f3ce42a..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments/filter.py +++ /dev/null @@ -1,70 +0,0 @@ -""" - pygments.filter - ~~~~~~~~~~~~~~~ - - Module that implements the default filter. - - :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS. - :license: BSD, see LICENSE for details. -""" - - -def apply_filters(stream, filters, lexer=None): - """ - Use this method to apply an iterable of filters to - a stream. If lexer is given it's forwarded to the - filter, otherwise the filter receives `None`. - """ - def _apply(filter_, stream): - yield from filter_.filter(lexer, stream) - for filter_ in filters: - stream = _apply(filter_, stream) - return stream - - -def simplefilter(f): - """ - Decorator that converts a function into a filter:: - - @simplefilter - def lowercase(self, lexer, stream, options): - for ttype, value in stream: - yield ttype, value.lower() - """ - return type(f.__name__, (FunctionFilter,), { - '__module__': getattr(f, '__module__'), - '__doc__': f.__doc__, - 'function': f, - }) - - -class Filter: - """ - Default filter. Subclass this class or use the `simplefilter` - decorator to create own filters. - """ - - def __init__(self, **options): - self.options = options - - def filter(self, lexer, stream): - raise NotImplementedError() - - -class FunctionFilter(Filter): - """ - Abstract class used by `simplefilter` to create simple - function filters on the fly. The `simplefilter` decorator - automatically creates subclasses of this class for - functions passed to it. - """ - function = None - - def __init__(self, **options): - if not hasattr(self, 'function'): - raise TypeError(f'{self.__class__.__name__!r} used without bound function') - Filter.__init__(self, **options) - - def filter(self, lexer, stream): - # pylint: disable=not-callable - yield from self.function(lexer, stream, self.options) diff --git a/venv/lib/python3.10/site-packages/pygments/filters/__init__.py b/venv/lib/python3.10/site-packages/pygments/filters/__init__.py deleted file mode 100644 index 2fed761a01ee5e9e25918b9a0619fcb47c8811b1..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments/filters/__init__.py +++ /dev/null @@ -1,940 +0,0 @@ -""" - pygments.filters - ~~~~~~~~~~~~~~~~ - - Module containing filter lookup functions and default - filters. - - :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS. - :license: BSD, see LICENSE for details. -""" - -import re - -from pygments.token import String, Comment, Keyword, Name, Error, Whitespace, \ - string_to_tokentype -from pygments.filter import Filter -from pygments.util import get_list_opt, get_int_opt, get_bool_opt, \ - get_choice_opt, ClassNotFound, OptionError -from pygments.plugin import find_plugin_filters - - -def find_filter_class(filtername): - """Lookup a filter by name. Return None if not found.""" - if filtername in FILTERS: - return FILTERS[filtername] - for name, cls in find_plugin_filters(): - if name == filtername: - return cls - return None - - -def get_filter_by_name(filtername, **options): - """Return an instantiated filter. - - Options are passed to the filter initializer if wanted. - Raise a ClassNotFound if not found. - """ - cls = find_filter_class(filtername) - if cls: - return cls(**options) - else: - raise ClassNotFound(f'filter {filtername!r} not found') - - -def get_all_filters(): - """Return a generator of all filter names.""" - yield from FILTERS - for name, _ in find_plugin_filters(): - yield name - - -def _replace_special(ttype, value, regex, specialttype, - replacefunc=lambda x: x): - last = 0 - for match in regex.finditer(value): - start, end = match.start(), match.end() - if start != last: - yield ttype, value[last:start] - yield specialttype, replacefunc(value[start:end]) - last = end - if last != len(value): - yield ttype, value[last:] - - -class CodeTagFilter(Filter): - """Highlight special code tags in comments and docstrings. - - Options accepted: - - `codetags` : list of strings - A list of strings that are flagged as code tags. The default is to - highlight ``XXX``, ``TODO``, ``FIXME``, ``BUG`` and ``NOTE``. - - .. versionchanged:: 2.13 - Now recognizes ``FIXME`` by default. - """ - - def __init__(self, **options): - Filter.__init__(self, **options) - tags = get_list_opt(options, 'codetags', - ['XXX', 'TODO', 'FIXME', 'BUG', 'NOTE']) - self.tag_re = re.compile(r'\b({})\b'.format('|'.join([ - re.escape(tag) for tag in tags if tag - ]))) - - def filter(self, lexer, stream): - regex = self.tag_re - for ttype, value in stream: - if ttype in String.Doc or \ - ttype in Comment and \ - ttype not in Comment.Preproc: - yield from _replace_special(ttype, value, regex, Comment.Special) - else: - yield ttype, value - - -class SymbolFilter(Filter): - """Convert mathematical symbols such as \\ in Isabelle - or \\longrightarrow in LaTeX into Unicode characters. - - This is mostly useful for HTML or console output when you want to - approximate the source rendering you'd see in an IDE. - - Options accepted: - - `lang` : string - The symbol language. Must be one of ``'isabelle'`` or - ``'latex'``. The default is ``'isabelle'``. - """ - - latex_symbols = { - '\\alpha' : '\U000003b1', - '\\beta' : '\U000003b2', - '\\gamma' : '\U000003b3', - '\\delta' : '\U000003b4', - '\\varepsilon' : '\U000003b5', - '\\zeta' : '\U000003b6', - '\\eta' : '\U000003b7', - '\\vartheta' : '\U000003b8', - '\\iota' : '\U000003b9', - '\\kappa' : '\U000003ba', - '\\lambda' : '\U000003bb', - '\\mu' : '\U000003bc', - '\\nu' : '\U000003bd', - '\\xi' : '\U000003be', - '\\pi' : '\U000003c0', - '\\varrho' : '\U000003c1', - '\\sigma' : '\U000003c3', - '\\tau' : '\U000003c4', - '\\upsilon' : '\U000003c5', - '\\varphi' : '\U000003c6', - '\\chi' : '\U000003c7', - '\\psi' : '\U000003c8', - '\\omega' : '\U000003c9', - '\\Gamma' : '\U00000393', - '\\Delta' : '\U00000394', - '\\Theta' : '\U00000398', - '\\Lambda' : '\U0000039b', - '\\Xi' : '\U0000039e', - '\\Pi' : '\U000003a0', - '\\Sigma' : '\U000003a3', - '\\Upsilon' : '\U000003a5', - '\\Phi' : '\U000003a6', - '\\Psi' : '\U000003a8', - '\\Omega' : '\U000003a9', - '\\leftarrow' : '\U00002190', - '\\longleftarrow' : '\U000027f5', - '\\rightarrow' : '\U00002192', - '\\longrightarrow' : '\U000027f6', - '\\Leftarrow' : '\U000021d0', - '\\Longleftarrow' : '\U000027f8', - '\\Rightarrow' : '\U000021d2', - '\\Longrightarrow' : '\U000027f9', - '\\leftrightarrow' : '\U00002194', - '\\longleftrightarrow' : '\U000027f7', - '\\Leftrightarrow' : '\U000021d4', - '\\Longleftrightarrow' : '\U000027fa', - '\\mapsto' : '\U000021a6', - '\\longmapsto' : '\U000027fc', - '\\relbar' : '\U00002500', - '\\Relbar' : '\U00002550', - '\\hookleftarrow' : '\U000021a9', - '\\hookrightarrow' : '\U000021aa', - '\\leftharpoondown' : '\U000021bd', - '\\rightharpoondown' : '\U000021c1', - '\\leftharpoonup' : '\U000021bc', - '\\rightharpoonup' : '\U000021c0', - '\\rightleftharpoons' : '\U000021cc', - '\\leadsto' : '\U0000219d', - '\\downharpoonleft' : '\U000021c3', - '\\downharpoonright' : '\U000021c2', - '\\upharpoonleft' : '\U000021bf', - '\\upharpoonright' : '\U000021be', - '\\restriction' : '\U000021be', - '\\uparrow' : '\U00002191', - '\\Uparrow' : '\U000021d1', - '\\downarrow' : '\U00002193', - '\\Downarrow' : '\U000021d3', - '\\updownarrow' : '\U00002195', - '\\Updownarrow' : '\U000021d5', - '\\langle' : '\U000027e8', - '\\rangle' : '\U000027e9', - '\\lceil' : '\U00002308', - '\\rceil' : '\U00002309', - '\\lfloor' : '\U0000230a', - '\\rfloor' : '\U0000230b', - '\\flqq' : '\U000000ab', - '\\frqq' : '\U000000bb', - '\\bot' : '\U000022a5', - '\\top' : '\U000022a4', - '\\wedge' : '\U00002227', - '\\bigwedge' : '\U000022c0', - '\\vee' : '\U00002228', - '\\bigvee' : '\U000022c1', - '\\forall' : '\U00002200', - '\\exists' : '\U00002203', - '\\nexists' : '\U00002204', - '\\neg' : '\U000000ac', - '\\Box' : '\U000025a1', - '\\Diamond' : '\U000025c7', - '\\vdash' : '\U000022a2', - '\\models' : '\U000022a8', - '\\dashv' : '\U000022a3', - '\\surd' : '\U0000221a', - '\\le' : '\U00002264', - '\\ge' : '\U00002265', - '\\ll' : '\U0000226a', - '\\gg' : '\U0000226b', - '\\lesssim' : '\U00002272', - '\\gtrsim' : '\U00002273', - '\\lessapprox' : '\U00002a85', - '\\gtrapprox' : '\U00002a86', - '\\in' : '\U00002208', - '\\notin' : '\U00002209', - '\\subset' : '\U00002282', - '\\supset' : '\U00002283', - '\\subseteq' : '\U00002286', - '\\supseteq' : '\U00002287', - '\\sqsubset' : '\U0000228f', - '\\sqsupset' : '\U00002290', - '\\sqsubseteq' : '\U00002291', - '\\sqsupseteq' : '\U00002292', - '\\cap' : '\U00002229', - '\\bigcap' : '\U000022c2', - '\\cup' : '\U0000222a', - '\\bigcup' : '\U000022c3', - '\\sqcup' : '\U00002294', - '\\bigsqcup' : '\U00002a06', - '\\sqcap' : '\U00002293', - '\\Bigsqcap' : '\U00002a05', - '\\setminus' : '\U00002216', - '\\propto' : '\U0000221d', - '\\uplus' : '\U0000228e', - '\\bigplus' : '\U00002a04', - '\\sim' : '\U0000223c', - '\\doteq' : '\U00002250', - '\\simeq' : '\U00002243', - '\\approx' : '\U00002248', - '\\asymp' : '\U0000224d', - '\\cong' : '\U00002245', - '\\equiv' : '\U00002261', - '\\Join' : '\U000022c8', - '\\bowtie' : '\U00002a1d', - '\\prec' : '\U0000227a', - '\\succ' : '\U0000227b', - '\\preceq' : '\U0000227c', - '\\succeq' : '\U0000227d', - '\\parallel' : '\U00002225', - '\\mid' : '\U000000a6', - '\\pm' : '\U000000b1', - '\\mp' : '\U00002213', - '\\times' : '\U000000d7', - '\\div' : '\U000000f7', - '\\cdot' : '\U000022c5', - '\\star' : '\U000022c6', - '\\circ' : '\U00002218', - '\\dagger' : '\U00002020', - '\\ddagger' : '\U00002021', - '\\lhd' : '\U000022b2', - '\\rhd' : '\U000022b3', - '\\unlhd' : '\U000022b4', - '\\unrhd' : '\U000022b5', - '\\triangleleft' : '\U000025c3', - '\\triangleright' : '\U000025b9', - '\\triangle' : '\U000025b3', - '\\triangleq' : '\U0000225c', - '\\oplus' : '\U00002295', - '\\bigoplus' : '\U00002a01', - '\\otimes' : '\U00002297', - '\\bigotimes' : '\U00002a02', - '\\odot' : '\U00002299', - '\\bigodot' : '\U00002a00', - '\\ominus' : '\U00002296', - '\\oslash' : '\U00002298', - '\\dots' : '\U00002026', - '\\cdots' : '\U000022ef', - '\\sum' : '\U00002211', - '\\prod' : '\U0000220f', - '\\coprod' : '\U00002210', - '\\infty' : '\U0000221e', - '\\int' : '\U0000222b', - '\\oint' : '\U0000222e', - '\\clubsuit' : '\U00002663', - '\\diamondsuit' : '\U00002662', - '\\heartsuit' : '\U00002661', - '\\spadesuit' : '\U00002660', - '\\aleph' : '\U00002135', - '\\emptyset' : '\U00002205', - '\\nabla' : '\U00002207', - '\\partial' : '\U00002202', - '\\flat' : '\U0000266d', - '\\natural' : '\U0000266e', - '\\sharp' : '\U0000266f', - '\\angle' : '\U00002220', - '\\copyright' : '\U000000a9', - '\\textregistered' : '\U000000ae', - '\\textonequarter' : '\U000000bc', - '\\textonehalf' : '\U000000bd', - '\\textthreequarters' : '\U000000be', - '\\textordfeminine' : '\U000000aa', - '\\textordmasculine' : '\U000000ba', - '\\euro' : '\U000020ac', - '\\pounds' : '\U000000a3', - '\\yen' : '\U000000a5', - '\\textcent' : '\U000000a2', - '\\textcurrency' : '\U000000a4', - '\\textdegree' : '\U000000b0', - } - - isabelle_symbols = { - '\\' : '\U0001d7ec', - '\\' : '\U0001d7ed', - '\\' : '\U0001d7ee', - '\\' : '\U0001d7ef', - '\\' : '\U0001d7f0', - '\\' : '\U0001d7f1', - '\\' : '\U0001d7f2', - '\\' : '\U0001d7f3', - '\\' : '\U0001d7f4', - '\\' : '\U0001d7f5', - '\\' : '\U0001d49c', - '\\' : '\U0000212c', - '\\' : '\U0001d49e', - '\\' : '\U0001d49f', - '\\' : '\U00002130', - '\\' : '\U00002131', - '\\' : '\U0001d4a2', - '\\' : '\U0000210b', - '\\' : '\U00002110', - '\\' : '\U0001d4a5', - '\\' : '\U0001d4a6', - '\\' : '\U00002112', - '\\' : '\U00002133', - '\\' : '\U0001d4a9', - '\\' : '\U0001d4aa', - '\\

' : '\U0001d5c9', - '\\' : '\U0001d5ca', - '\\' : '\U0001d5cb', - '\\' : '\U0001d5cc', - '\\' : '\U0001d5cd', - '\\' : '\U0001d5ce', - '\\' : '\U0001d5cf', - '\\' : '\U0001d5d0', - '\\' : '\U0001d5d1', - '\\' : '\U0001d5d2', - '\\' : '\U0001d5d3', - '\\' : '\U0001d504', - '\\' : '\U0001d505', - '\\' : '\U0000212d', - '\\

' : '\U0001d507', - '\\' : '\U0001d508', - '\\' : '\U0001d509', - '\\' : '\U0001d50a', - '\\' : '\U0000210c', - '\\' : '\U00002111', - '\\' : '\U0001d50d', - '\\' : '\U0001d50e', - '\\' : '\U0001d50f', - '\\' : '\U0001d510', - '\\' : '\U0001d511', - '\\' : '\U0001d512', - '\\' : '\U0001d513', - '\\' : '\U0001d514', - '\\' : '\U0000211c', - '\\' : '\U0001d516', - '\\' : '\U0001d517', - '\\' : '\U0001d518', - '\\' : '\U0001d519', - '\\' : '\U0001d51a', - '\\' : '\U0001d51b', - '\\' : '\U0001d51c', - '\\' : '\U00002128', - '\\' : '\U0001d51e', - '\\' : '\U0001d51f', - '\\' : '\U0001d520', - '\\
' : '\U0001d521', - '\\' : '\U0001d522', - '\\' : '\U0001d523', - '\\' : '\U0001d524', - '\\' : '\U0001d525', - '\\' : '\U0001d526', - '\\' : '\U0001d527', - '\\' : '\U0001d528', - '\\' : '\U0001d529', - '\\' : '\U0001d52a', - '\\' : '\U0001d52b', - '\\' : '\U0001d52c', - '\\' : '\U0001d52d', - '\\' : '\U0001d52e', - '\\' : '\U0001d52f', - '\\' : '\U0001d530', - '\\' : '\U0001d531', - '\\' : '\U0001d532', - '\\' : '\U0001d533', - '\\' : '\U0001d534', - '\\' : '\U0001d535', - '\\' : '\U0001d536', - '\\' : '\U0001d537', - '\\' : '\U000003b1', - '\\' : '\U000003b2', - '\\' : '\U000003b3', - '\\' : '\U000003b4', - '\\' : '\U000003b5', - '\\' : '\U000003b6', - '\\' : '\U000003b7', - '\\' : '\U000003b8', - '\\' : '\U000003b9', - '\\' : '\U000003ba', - '\\' : '\U000003bb', - '\\' : '\U000003bc', - '\\' : '\U000003bd', - '\\' : '\U000003be', - '\\' : '\U000003c0', - '\\' : '\U000003c1', - '\\' : '\U000003c3', - '\\' : '\U000003c4', - '\\' : '\U000003c5', - '\\' : '\U000003c6', - '\\' : '\U000003c7', - '\\' : '\U000003c8', - '\\' : '\U000003c9', - '\\' : '\U00000393', - '\\' : '\U00000394', - '\\' : '\U00000398', - '\\' : '\U0000039b', - '\\' : '\U0000039e', - '\\' : '\U000003a0', - '\\' : '\U000003a3', - '\\' : '\U000003a5', - '\\' : '\U000003a6', - '\\' : '\U000003a8', - '\\' : '\U000003a9', - '\\' : '\U0001d539', - '\\' : '\U00002102', - '\\' : '\U00002115', - '\\' : '\U0000211a', - '\\' : '\U0000211d', - '\\' : '\U00002124', - '\\' : '\U00002190', - '\\' : '\U000027f5', - '\\' : '\U00002192', - '\\' : '\U000027f6', - '\\' : '\U000021d0', - '\\' : '\U000027f8', - '\\' : '\U000021d2', - '\\' : '\U000027f9', - '\\' : '\U00002194', - '\\' : '\U000027f7', - '\\' : '\U000021d4', - '\\' : '\U000027fa', - '\\' : '\U000021a6', - '\\' : '\U000027fc', - '\\' : '\U00002500', - '\\' : '\U00002550', - '\\' : '\U000021a9', - '\\' : '\U000021aa', - '\\' : '\U000021bd', - '\\' : '\U000021c1', - '\\' : '\U000021bc', - '\\' : '\U000021c0', - '\\' : '\U000021cc', - '\\' : '\U0000219d', - '\\' : '\U000021c3', - '\\' : '\U000021c2', - '\\' : '\U000021bf', - '\\' : '\U000021be', - '\\' : '\U000021be', - '\\' : '\U00002237', - '\\' : '\U00002191', - '\\' : '\U000021d1', - '\\' : '\U00002193', - '\\' : '\U000021d3', - '\\' : '\U00002195', - '\\' : '\U000021d5', - '\\' : '\U000027e8', - '\\' : '\U000027e9', - '\\' : '\U00002308', - '\\' : '\U00002309', - '\\' : '\U0000230a', - '\\' : '\U0000230b', - '\\' : '\U00002987', - '\\' : '\U00002988', - '\\' : '\U000027e6', - '\\' : '\U000027e7', - '\\' : '\U00002983', - '\\' : '\U00002984', - '\\' : '\U000000ab', - '\\' : '\U000000bb', - '\\' : '\U000022a5', - '\\' : '\U000022a4', - '\\' : '\U00002227', - '\\' : '\U000022c0', - '\\' : '\U00002228', - '\\' : '\U000022c1', - '\\' : '\U00002200', - '\\' : '\U00002203', - '\\' : '\U00002204', - '\\' : '\U000000ac', - '\\' : '\U000025a1', - '\\' : '\U000025c7', - '\\' : '\U000022a2', - '\\' : '\U000022a8', - '\\' : '\U000022a9', - '\\' : '\U000022ab', - '\\' : '\U000022a3', - '\\' : '\U0000221a', - '\\' : '\U00002264', - '\\' : '\U00002265', - '\\' : '\U0000226a', - '\\' : '\U0000226b', - '\\' : '\U00002272', - '\\' : '\U00002273', - '\\' : '\U00002a85', - '\\' : '\U00002a86', - '\\' : '\U00002208', - '\\' : '\U00002209', - '\\' : '\U00002282', - '\\' : '\U00002283', - '\\' : '\U00002286', - '\\' : '\U00002287', - '\\' : '\U0000228f', - '\\' : '\U00002290', - '\\' : '\U00002291', - '\\' : '\U00002292', - '\\' : '\U00002229', - '\\' : '\U000022c2', - '\\' : '\U0000222a', - '\\' : '\U000022c3', - '\\' : '\U00002294', - '\\' : '\U00002a06', - '\\' : '\U00002293', - '\\' : '\U00002a05', - '\\' : '\U00002216', - '\\' : '\U0000221d', - '\\' : '\U0000228e', - '\\' : '\U00002a04', - '\\' : '\U00002260', - '\\' : '\U0000223c', - '\\' : '\U00002250', - '\\' : '\U00002243', - '\\' : '\U00002248', - '\\' : '\U0000224d', - '\\' : '\U00002245', - '\\' : '\U00002323', - '\\' : '\U00002261', - '\\' : '\U00002322', - '\\' : '\U000022c8', - '\\' : '\U00002a1d', - '\\' : '\U0000227a', - '\\' : '\U0000227b', - '\\' : '\U0000227c', - '\\' : '\U0000227d', - '\\' : '\U00002225', - '\\' : '\U000000a6', - '\\' : '\U000000b1', - '\\' : '\U00002213', - '\\' : '\U000000d7', - '\\
' : '\U000000f7', - '\\' : '\U000022c5', - '\\' : '\U000022c6', - '\\' : '\U00002219', - '\\' : '\U00002218', - '\\' : '\U00002020', - '\\' : '\U00002021', - '\\' : '\U000022b2', - '\\' : '\U000022b3', - '\\' : '\U000022b4', - '\\' : '\U000022b5', - '\\' : '\U000025c3', - '\\' : '\U000025b9', - '\\' : '\U000025b3', - '\\' : '\U0000225c', - '\\' : '\U00002295', - '\\' : '\U00002a01', - '\\' : '\U00002297', - '\\' : '\U00002a02', - '\\' : '\U00002299', - '\\' : '\U00002a00', - '\\' : '\U00002296', - '\\' : '\U00002298', - '\\' : '\U00002026', - '\\' : '\U000022ef', - '\\' : '\U00002211', - '\\' : '\U0000220f', - '\\' : '\U00002210', - '\\' : '\U0000221e', - '\\' : '\U0000222b', - '\\' : '\U0000222e', - '\\' : '\U00002663', - '\\' : '\U00002662', - '\\' : '\U00002661', - '\\' : '\U00002660', - '\\' : '\U00002135', - '\\' : '\U00002205', - '\\' : '\U00002207', - '\\' : '\U00002202', - '\\' : '\U0000266d', - '\\' : '\U0000266e', - '\\' : '\U0000266f', - '\\' : '\U00002220', - '\\' : '\U000000a9', - '\\' : '\U000000ae', - '\\' : '\U000000ad', - '\\' : '\U000000af', - '\\' : '\U000000bc', - '\\' : '\U000000bd', - '\\' : '\U000000be', - '\\' : '\U000000aa', - '\\' : '\U000000ba', - '\\
' : '\U000000a7', - '\\' : '\U000000b6', - '\\' : '\U000000a1', - '\\' : '\U000000bf', - '\\' : '\U000020ac', - '\\' : '\U000000a3', - '\\' : '\U000000a5', - '\\' : '\U000000a2', - '\\' : '\U000000a4', - '\\' : '\U000000b0', - '\\' : '\U00002a3f', - '\\' : '\U00002127', - '\\' : '\U000025ca', - '\\' : '\U00002118', - '\\' : '\U00002240', - '\\' : '\U000022c4', - '\\' : '\U000000b4', - '\\' : '\U00000131', - '\\' : '\U000000a8', - '\\' : '\U000000b8', - '\\' : '\U000002dd', - '\\' : '\U000003f5', - '\\' : '\U000023ce', - '\\' : '\U00002039', - '\\' : '\U0000203a', - '\\' : '\U00002302', - '\\<^sub>' : '\U000021e9', - '\\<^sup>' : '\U000021e7', - '\\<^bold>' : '\U00002759', - '\\<^bsub>' : '\U000021d8', - '\\<^esub>' : '\U000021d9', - '\\<^bsup>' : '\U000021d7', - '\\<^esup>' : '\U000021d6', - } - - lang_map = {'isabelle' : isabelle_symbols, 'latex' : latex_symbols} - - def __init__(self, **options): - Filter.__init__(self, **options) - lang = get_choice_opt(options, 'lang', - ['isabelle', 'latex'], 'isabelle') - self.symbols = self.lang_map[lang] - - def filter(self, lexer, stream): - for ttype, value in stream: - if value in self.symbols: - yield ttype, self.symbols[value] - else: - yield ttype, value - - -class KeywordCaseFilter(Filter): - """Convert keywords to lowercase or uppercase or capitalize them, which - means first letter uppercase, rest lowercase. - - This can be useful e.g. if you highlight Pascal code and want to adapt the - code to your styleguide. - - Options accepted: - - `case` : string - The casing to convert keywords to. Must be one of ``'lower'``, - ``'upper'`` or ``'capitalize'``. The default is ``'lower'``. - """ - - def __init__(self, **options): - Filter.__init__(self, **options) - case = get_choice_opt(options, 'case', - ['lower', 'upper', 'capitalize'], 'lower') - self.convert = getattr(str, case) - - def filter(self, lexer, stream): - for ttype, value in stream: - if ttype in Keyword: - yield ttype, self.convert(value) - else: - yield ttype, value - - -class NameHighlightFilter(Filter): - """Highlight a normal Name (and Name.*) token with a different token type. - - Example:: - - filter = NameHighlightFilter( - names=['foo', 'bar', 'baz'], - tokentype=Name.Function, - ) - - This would highlight the names "foo", "bar" and "baz" - as functions. `Name.Function` is the default token type. - - Options accepted: - - `names` : list of strings - A list of names that should be given the different token type. - There is no default. - `tokentype` : TokenType or string - A token type or a string containing a token type name that is - used for highlighting the strings in `names`. The default is - `Name.Function`. - """ - - def __init__(self, **options): - Filter.__init__(self, **options) - self.names = set(get_list_opt(options, 'names', [])) - tokentype = options.get('tokentype') - if tokentype: - self.tokentype = string_to_tokentype(tokentype) - else: - self.tokentype = Name.Function - - def filter(self, lexer, stream): - for ttype, value in stream: - if ttype in Name and value in self.names: - yield self.tokentype, value - else: - yield ttype, value - - -class ErrorToken(Exception): - pass - - -class RaiseOnErrorTokenFilter(Filter): - """Raise an exception when the lexer generates an error token. - - Options accepted: - - `excclass` : Exception class - The exception class to raise. - The default is `pygments.filters.ErrorToken`. - - .. versionadded:: 0.8 - """ - - def __init__(self, **options): - Filter.__init__(self, **options) - self.exception = options.get('excclass', ErrorToken) - try: - # issubclass() will raise TypeError if first argument is not a class - if not issubclass(self.exception, Exception): - raise TypeError - except TypeError: - raise OptionError('excclass option is not an exception class') - - def filter(self, lexer, stream): - for ttype, value in stream: - if ttype is Error: - raise self.exception(value) - yield ttype, value - - -class VisibleWhitespaceFilter(Filter): - """Convert tabs, newlines and/or spaces to visible characters. - - Options accepted: - - `spaces` : string or bool - If this is a one-character string, spaces will be replaces by this string. - If it is another true value, spaces will be replaced by ``·`` (unicode - MIDDLE DOT). If it is a false value, spaces will not be replaced. The - default is ``False``. - `tabs` : string or bool - The same as for `spaces`, but the default replacement character is ``»`` - (unicode RIGHT-POINTING DOUBLE ANGLE QUOTATION MARK). The default value - is ``False``. Note: this will not work if the `tabsize` option for the - lexer is nonzero, as tabs will already have been expanded then. - `tabsize` : int - If tabs are to be replaced by this filter (see the `tabs` option), this - is the total number of characters that a tab should be expanded to. - The default is ``8``. - `newlines` : string or bool - The same as for `spaces`, but the default replacement character is ``¶`` - (unicode PILCROW SIGN). The default value is ``False``. - `wstokentype` : bool - If true, give whitespace the special `Whitespace` token type. This allows - styling the visible whitespace differently (e.g. greyed out), but it can - disrupt background colors. The default is ``True``. - - .. versionadded:: 0.8 - """ - - def __init__(self, **options): - Filter.__init__(self, **options) - for name, default in [('spaces', '·'), - ('tabs', '»'), - ('newlines', '¶')]: - opt = options.get(name, False) - if isinstance(opt, str) and len(opt) == 1: - setattr(self, name, opt) - else: - setattr(self, name, (opt and default or '')) - tabsize = get_int_opt(options, 'tabsize', 8) - if self.tabs: - self.tabs += ' ' * (tabsize - 1) - if self.newlines: - self.newlines += '\n' - self.wstt = get_bool_opt(options, 'wstokentype', True) - - def filter(self, lexer, stream): - if self.wstt: - spaces = self.spaces or ' ' - tabs = self.tabs or '\t' - newlines = self.newlines or '\n' - regex = re.compile(r'\s') - - def replacefunc(wschar): - if wschar == ' ': - return spaces - elif wschar == '\t': - return tabs - elif wschar == '\n': - return newlines - return wschar - - for ttype, value in stream: - yield from _replace_special(ttype, value, regex, Whitespace, - replacefunc) - else: - spaces, tabs, newlines = self.spaces, self.tabs, self.newlines - # simpler processing - for ttype, value in stream: - if spaces: - value = value.replace(' ', spaces) - if tabs: - value = value.replace('\t', tabs) - if newlines: - value = value.replace('\n', newlines) - yield ttype, value - - -class GobbleFilter(Filter): - """Gobbles source code lines (eats initial characters). - - This filter drops the first ``n`` characters off every line of code. This - may be useful when the source code fed to the lexer is indented by a fixed - amount of space that isn't desired in the output. - - Options accepted: - - `n` : int - The number of characters to gobble. - - .. versionadded:: 1.2 - """ - def __init__(self, **options): - Filter.__init__(self, **options) - self.n = get_int_opt(options, 'n', 0) - - def gobble(self, value, left): - if left < len(value): - return value[left:], 0 - else: - return '', left - len(value) - - def filter(self, lexer, stream): - n = self.n - left = n # How many characters left to gobble. - for ttype, value in stream: - # Remove ``left`` tokens from first line, ``n`` from all others. - parts = value.split('\n') - (parts[0], left) = self.gobble(parts[0], left) - for i in range(1, len(parts)): - (parts[i], left) = self.gobble(parts[i], n) - value = '\n'.join(parts) - - if value != '': - yield ttype, value - - -class TokenMergeFilter(Filter): - """Merges consecutive tokens with the same token type in the output - stream of a lexer. - - .. versionadded:: 1.2 - """ - def __init__(self, **options): - Filter.__init__(self, **options) - - def filter(self, lexer, stream): - current_type = None - current_value = None - for ttype, value in stream: - if ttype is current_type: - current_value += value - else: - if current_type is not None: - yield current_type, current_value - current_type = ttype - current_value = value - if current_type is not None: - yield current_type, current_value - - -FILTERS = { - 'codetagify': CodeTagFilter, - 'keywordcase': KeywordCaseFilter, - 'highlight': NameHighlightFilter, - 'raiseonerror': RaiseOnErrorTokenFilter, - 'whitespace': VisibleWhitespaceFilter, - 'gobble': GobbleFilter, - 'tokenmerge': TokenMergeFilter, - 'symbols': SymbolFilter, -} diff --git a/venv/lib/python3.10/site-packages/pygments/filters/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/filters/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index 04be0ab53d64f3d726f221d4406614ce97643bca..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/filters/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatter.py b/venv/lib/python3.10/site-packages/pygments/formatter.py deleted file mode 100644 index a20d3039d3dd3c40881f392136640f9d1a7e5e82..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments/formatter.py +++ /dev/null @@ -1,129 +0,0 @@ -""" - pygments.formatter - ~~~~~~~~~~~~~~~~~~ - - Base formatter class. - - :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS. - :license: BSD, see LICENSE for details. -""" - -import codecs - -from pygments.util import get_bool_opt -from pygments.styles import get_style_by_name - -__all__ = ['Formatter'] - - -def _lookup_style(style): - if isinstance(style, str): - return get_style_by_name(style) - return style - - -class Formatter: - """ - Converts a token stream to text. - - Formatters should have attributes to help selecting them. These - are similar to the corresponding :class:`~pygments.lexer.Lexer` - attributes. - - .. autoattribute:: name - :no-value: - - .. autoattribute:: aliases - :no-value: - - .. autoattribute:: filenames - :no-value: - - You can pass options as keyword arguments to the constructor. - All formatters accept these basic options: - - ``style`` - The style to use, can be a string or a Style subclass - (default: "default"). Not used by e.g. the - TerminalFormatter. - ``full`` - Tells the formatter to output a "full" document, i.e. - a complete self-contained document. This doesn't have - any effect for some formatters (default: false). - ``title`` - If ``full`` is true, the title that should be used to - caption the document (default: ''). - ``encoding`` - If given, must be an encoding name. This will be used to - convert the Unicode token strings to byte strings in the - output. If it is "" or None, Unicode strings will be written - to the output file, which most file-like objects do not - support (default: None). - ``outencoding`` - Overrides ``encoding`` if given. - - """ - - #: Full name for the formatter, in human-readable form. - name = None - - #: A list of short, unique identifiers that can be used to lookup - #: the formatter from a list, e.g. using :func:`.get_formatter_by_name()`. - aliases = [] - - #: A list of fnmatch patterns that match filenames for which this - #: formatter can produce output. The patterns in this list should be unique - #: among all formatters. - filenames = [] - - #: If True, this formatter outputs Unicode strings when no encoding - #: option is given. - unicodeoutput = True - - def __init__(self, **options): - """ - As with lexers, this constructor takes arbitrary optional arguments, - and if you override it, you should first process your own options, then - call the base class implementation. - """ - self.style = _lookup_style(options.get('style', 'default')) - self.full = get_bool_opt(options, 'full', False) - self.title = options.get('title', '') - self.encoding = options.get('encoding', None) or None - if self.encoding in ('guess', 'chardet'): - # can happen for e.g. pygmentize -O encoding=guess - self.encoding = 'utf-8' - self.encoding = options.get('outencoding') or self.encoding - self.options = options - - def get_style_defs(self, arg=''): - """ - This method must return statements or declarations suitable to define - the current style for subsequent highlighted text (e.g. CSS classes - in the `HTMLFormatter`). - - The optional argument `arg` can be used to modify the generation and - is formatter dependent (it is standardized because it can be given on - the command line). - - This method is called by the ``-S`` :doc:`command-line option `, - the `arg` is then given by the ``-a`` option. - """ - return '' - - def format(self, tokensource, outfile): - """ - This method must format the tokens from the `tokensource` iterable and - write the formatted version to the file object `outfile`. - - Formatter options can control how exactly the tokens are converted. - """ - if self.encoding: - # wrap the outfile in a StreamWriter - outfile = codecs.lookup(self.encoding)[3](outfile) - return self.format_unencoded(tokensource, outfile) - - # Allow writing Formatter[str] or Formatter[bytes]. That's equivalent to - # Formatter. This helps when using third-party type stubs from typeshed. - def __class_getitem__(cls, name): - return cls diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__init__.py b/venv/lib/python3.10/site-packages/pygments/formatters/__init__.py deleted file mode 100644 index b24931cd191efe2a83010b46422face50f050f37..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments/formatters/__init__.py +++ /dev/null @@ -1,157 +0,0 @@ -""" - pygments.formatters - ~~~~~~~~~~~~~~~~~~~ - - Pygments formatters. - - :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS. - :license: BSD, see LICENSE for details. -""" - -import re -import sys -import types -import fnmatch -from os.path import basename - -from pygments.formatters._mapping import FORMATTERS -from pygments.plugin import find_plugin_formatters -from pygments.util import ClassNotFound - -__all__ = ['get_formatter_by_name', 'get_formatter_for_filename', - 'get_all_formatters', 'load_formatter_from_file'] + list(FORMATTERS) - -_formatter_cache = {} # classes by name -_pattern_cache = {} - - -def _fn_matches(fn, glob): - """Return whether the supplied file name fn matches pattern filename.""" - if glob not in _pattern_cache: - pattern = _pattern_cache[glob] = re.compile(fnmatch.translate(glob)) - return pattern.match(fn) - return _pattern_cache[glob].match(fn) - - -def _load_formatters(module_name): - """Load a formatter (and all others in the module too).""" - mod = __import__(module_name, None, None, ['__all__']) - for formatter_name in mod.__all__: - cls = getattr(mod, formatter_name) - _formatter_cache[cls.name] = cls - - -def get_all_formatters(): - """Return a generator for all formatter classes.""" - # NB: this returns formatter classes, not info like get_all_lexers(). - for info in FORMATTERS.values(): - if info[1] not in _formatter_cache: - _load_formatters(info[0]) - yield _formatter_cache[info[1]] - for _, formatter in find_plugin_formatters(): - yield formatter - - -def find_formatter_class(alias): - """Lookup a formatter by alias. - - Returns None if not found. - """ - for module_name, name, aliases, _, _ in FORMATTERS.values(): - if alias in aliases: - if name not in _formatter_cache: - _load_formatters(module_name) - return _formatter_cache[name] - for _, cls in find_plugin_formatters(): - if alias in cls.aliases: - return cls - - -def get_formatter_by_name(_alias, **options): - """ - Return an instance of a :class:`.Formatter` subclass that has `alias` in its - aliases list. The formatter is given the `options` at its instantiation. - - Will raise :exc:`pygments.util.ClassNotFound` if no formatter with that - alias is found. - """ - cls = find_formatter_class(_alias) - if cls is None: - raise ClassNotFound(f"no formatter found for name {_alias!r}") - return cls(**options) - - -def load_formatter_from_file(filename, formattername="CustomFormatter", **options): - """ - Return a `Formatter` subclass instance loaded from the provided file, relative - to the current directory. - - The file is expected to contain a Formatter class named ``formattername`` - (by default, CustomFormatter). Users should be very careful with the input, because - this method is equivalent to running ``eval()`` on the input file. The formatter is - given the `options` at its instantiation. - - :exc:`pygments.util.ClassNotFound` is raised if there are any errors loading - the formatter. - - .. versionadded:: 2.2 - """ - try: - # This empty dict will contain the namespace for the exec'd file - custom_namespace = {} - with open(filename, 'rb') as f: - exec(f.read(), custom_namespace) - # Retrieve the class `formattername` from that namespace - if formattername not in custom_namespace: - raise ClassNotFound(f'no valid {formattername} class found in {filename}') - formatter_class = custom_namespace[formattername] - # And finally instantiate it with the options - return formatter_class(**options) - except OSError as err: - raise ClassNotFound(f'cannot read {filename}: {err}') - except ClassNotFound: - raise - except Exception as err: - raise ClassNotFound(f'error when loading custom formatter: {err}') - - -def get_formatter_for_filename(fn, **options): - """ - Return a :class:`.Formatter` subclass instance that has a filename pattern - matching `fn`. The formatter is given the `options` at its instantiation. - - Will raise :exc:`pygments.util.ClassNotFound` if no formatter for that filename - is found. - """ - fn = basename(fn) - for modname, name, _, filenames, _ in FORMATTERS.values(): - for filename in filenames: - if _fn_matches(fn, filename): - if name not in _formatter_cache: - _load_formatters(modname) - return _formatter_cache[name](**options) - for _name, cls in find_plugin_formatters(): - for filename in cls.filenames: - if _fn_matches(fn, filename): - return cls(**options) - raise ClassNotFound(f"no formatter found for file name {fn!r}") - - -class _automodule(types.ModuleType): - """Automatically import formatters.""" - - def __getattr__(self, name): - info = FORMATTERS.get(name) - if info: - _load_formatters(info[0]) - cls = _formatter_cache[info[1]] - setattr(self, name, cls) - return cls - raise AttributeError(name) - - -oldmod = sys.modules[__name__] -newmod = _automodule(__name__) -newmod.__dict__.update(oldmod.__dict__) -sys.modules[__name__] = newmod -del newmod.newmod, newmod.oldmod, newmod.sys, newmod.types diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/__init__.cpython-310.pyc deleted file mode 100644 index c8f370da95efe50efaec8e92193097503730c2dc..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/__init__.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/_mapping.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/_mapping.cpython-310.pyc deleted file mode 100644 index 02ee2cb1bf77b75a63c790184dc67b94d4b521f6..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/_mapping.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/bbcode.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/bbcode.cpython-310.pyc deleted file mode 100644 index 9e5a0b5f71fd87a89d059e30b61280dc16251f01..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/bbcode.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/groff.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/groff.cpython-310.pyc deleted file mode 100644 index 8147011192035e09d1d1c1b3ed694adf196d3b19..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/groff.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/html.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/html.cpython-310.pyc deleted file mode 100644 index 267a9c848cef35c8a0d64913cc6e37fdf4bd568b..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/html.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/img.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/img.cpython-310.pyc deleted file mode 100644 index ef3a33a5ab128027c7e091c3fe8d7945cc6b2857..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/img.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/irc.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/irc.cpython-310.pyc deleted file mode 100644 index 450732a5a2eceafc18efac6246a74c1b1a5955c5..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/irc.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/latex.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/latex.cpython-310.pyc deleted file mode 100644 index dd33e1fb950b2b5e66895afdf3b0948d672970d0..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/latex.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/other.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/other.cpython-310.pyc deleted file mode 100644 index 5cb2a032cb56523f20b58fd9e2488ed166afb2b4..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/other.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/pangomarkup.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/pangomarkup.cpython-310.pyc deleted file mode 100644 index 1d670a067dd2559dd78c7b5e62c68e380eaa87e9..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/pangomarkup.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/rtf.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/rtf.cpython-310.pyc deleted file mode 100644 index 79a55dec052b01d86445b1fececd38fc974d7351..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/rtf.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/svg.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/svg.cpython-310.pyc deleted file mode 100644 index c324cf10cd627793b27b454f757f944e71b233af..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/svg.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/terminal.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/terminal.cpython-310.pyc deleted file mode 100644 index 5625316193e1fb69b18fe4b546a3e17453eaba22..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/terminal.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/terminal256.cpython-310.pyc b/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/terminal256.cpython-310.pyc deleted file mode 100644 index 9c7f908745cf701a4eaf3b45c11cf054c81d745a..0000000000000000000000000000000000000000 Binary files a/venv/lib/python3.10/site-packages/pygments/formatters/__pycache__/terminal256.cpython-310.pyc and /dev/null differ diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/_mapping.py b/venv/lib/python3.10/site-packages/pygments/formatters/_mapping.py deleted file mode 100644 index 72ca84040b626183e3328679db600c13472021be..0000000000000000000000000000000000000000 --- a/venv/lib/python3.10/site-packages/pygments/formatters/_mapping.py +++ /dev/null @@ -1,23 +0,0 @@ -# Automatically generated by scripts/gen_mapfiles.py. -# DO NOT EDIT BY HAND; run `tox -e mapfiles` instead. - -FORMATTERS = { - 'BBCodeFormatter': ('pygments.formatters.bbcode', 'BBCode', ('bbcode', 'bb'), (), 'Format tokens with BBcodes. These formatting codes are used by many bulletin boards, so you can highlight your sourcecode with pygments before posting it there.'), - 'BmpImageFormatter': ('pygments.formatters.img', 'img_bmp', ('bmp', 'bitmap'), ('*.bmp',), 'Create a bitmap image from source code. This uses the Python Imaging Library to generate a pixmap from the source code.'), - 'GifImageFormatter': ('pygments.formatters.img', 'img_gif', ('gif',), ('*.gif',), 'Create a GIF image from source code. This uses the Python Imaging Library to generate a pixmap from the source code.'), - 'GroffFormatter': ('pygments.formatters.groff', 'groff', ('groff', 'troff', 'roff'), (), 'Format tokens with groff escapes to change their color and font style.'), - 'HtmlFormatter': ('pygments.formatters.html', 'HTML', ('html',), ('*.html', '*.htm'), "Format tokens as HTML 4 ```` tags. By default, the content is enclosed in a ``
`` tag, itself wrapped in a ``
`` tag (but see the `nowrap` option). The ``
``'s CSS class can be set by the `cssclass` option."),
-    'IRCFormatter': ('pygments.formatters.irc', 'IRC', ('irc', 'IRC'), (), 'Format tokens with IRC color sequences'),
-    'ImageFormatter': ('pygments.formatters.img', 'img', ('img', 'IMG', 'png'), ('*.png',), 'Create a PNG image from source code. This uses the Python Imaging Library to generate a pixmap from the source code.'),
-    'JpgImageFormatter': ('pygments.formatters.img', 'img_jpg', ('jpg', 'jpeg'), ('*.jpg',), 'Create a JPEG image from source code. This uses the Python Imaging Library to generate a pixmap from the source code.'),
-    'LatexFormatter': ('pygments.formatters.latex', 'LaTeX', ('latex', 'tex'), ('*.tex',), 'Format tokens as LaTeX code. This needs the `fancyvrb` and `color` standard packages.'),
-    'NullFormatter': ('pygments.formatters.other', 'Text only', ('text', 'null'), ('*.txt',), 'Output the text unchanged without any formatting.'),
-    'PangoMarkupFormatter': ('pygments.formatters.pangomarkup', 'Pango Markup', ('pango', 'pangomarkup'), (), 'Format tokens as Pango Markup code. It can then be rendered to an SVG.'),
-    'RawTokenFormatter': ('pygments.formatters.other', 'Raw tokens', ('raw', 'tokens'), ('*.raw',), 'Format tokens as a raw representation for storing token streams.'),
-    'RtfFormatter': ('pygments.formatters.rtf', 'RTF', ('rtf',), ('*.rtf',), 'Format tokens as RTF markup. This formatter automatically outputs full RTF documents with color information and other useful stuff. Perfect for Copy and Paste into Microsoft(R) Word(R) documents.'),
-    'SvgFormatter': ('pygments.formatters.svg', 'SVG', ('svg',), ('*.svg',), 'Format tokens as an SVG graphics file.  This formatter is still experimental. Each line of code is a ```` element with explicit ``x`` and ``y`` coordinates containing ```` elements with the individual token styles.'),
-    'Terminal256Formatter': ('pygments.formatters.terminal256', 'Terminal256', ('terminal256', 'console256', '256'), (), 'Format tokens with ANSI color sequences, for output in a 256-color terminal or console.  Like in `TerminalFormatter` color sequences are terminated at newlines, so that paging the output works correctly.'),
-    'TerminalFormatter': ('pygments.formatters.terminal', 'Terminal', ('terminal', 'console'), (), 'Format tokens with ANSI color sequences, for output in a text console. Color sequences are terminated at newlines, so that paging the output works correctly.'),
-    'TerminalTrueColorFormatter': ('pygments.formatters.terminal256', 'TerminalTrueColor', ('terminal16m', 'console16m', '16m'), (), 'Format tokens with ANSI color sequences, for output in a true-color terminal or console.  Like in `TerminalFormatter` color sequences are terminated at newlines, so that paging the output works correctly.'),
-    'TestcaseFormatter': ('pygments.formatters.other', 'Testcase', ('testcase',), (), 'Format tokens as appropriate for a new testcase.'),
-}
diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/bbcode.py b/venv/lib/python3.10/site-packages/pygments/formatters/bbcode.py
deleted file mode 100644
index 339edf9dc9a87829457de5cf27f2e7081569cc75..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/formatters/bbcode.py
+++ /dev/null
@@ -1,108 +0,0 @@
-"""
-    pygments.formatters.bbcode
-    ~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-    BBcode formatter.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-
-from pygments.formatter import Formatter
-from pygments.util import get_bool_opt
-
-__all__ = ['BBCodeFormatter']
-
-
-class BBCodeFormatter(Formatter):
-    """
-    Format tokens with BBcodes. These formatting codes are used by many
-    bulletin boards, so you can highlight your sourcecode with pygments before
-    posting it there.
-
-    This formatter has no support for background colors and borders, as there
-    are no common BBcode tags for that.
-
-    Some board systems (e.g. phpBB) don't support colors in their [code] tag,
-    so you can't use the highlighting together with that tag.
-    Text in a [code] tag usually is shown with a monospace font (which this
-    formatter can do with the ``monofont`` option) and no spaces (which you
-    need for indentation) are removed.
-
-    Additional options accepted:
-
-    `style`
-        The style to use, can be a string or a Style subclass (default:
-        ``'default'``).
-
-    `codetag`
-        If set to true, put the output into ``[code]`` tags (default:
-        ``false``)
-
-    `monofont`
-        If set to true, add a tag to show the code with a monospace font
-        (default: ``false``).
-    """
-    name = 'BBCode'
-    aliases = ['bbcode', 'bb']
-    filenames = []
-
-    def __init__(self, **options):
-        Formatter.__init__(self, **options)
-        self._code = get_bool_opt(options, 'codetag', False)
-        self._mono = get_bool_opt(options, 'monofont', False)
-
-        self.styles = {}
-        self._make_styles()
-
-    def _make_styles(self):
-        for ttype, ndef in self.style:
-            start = end = ''
-            if ndef['color']:
-                start += '[color=#{}]'.format(ndef['color'])
-                end = '[/color]' + end
-            if ndef['bold']:
-                start += '[b]'
-                end = '[/b]' + end
-            if ndef['italic']:
-                start += '[i]'
-                end = '[/i]' + end
-            if ndef['underline']:
-                start += '[u]'
-                end = '[/u]' + end
-            # there are no common BBcodes for background-color and border
-
-            self.styles[ttype] = start, end
-
-    def format_unencoded(self, tokensource, outfile):
-        if self._code:
-            outfile.write('[code]')
-        if self._mono:
-            outfile.write('[font=monospace]')
-
-        lastval = ''
-        lasttype = None
-
-        for ttype, value in tokensource:
-            while ttype not in self.styles:
-                ttype = ttype.parent
-            if ttype == lasttype:
-                lastval += value
-            else:
-                if lastval:
-                    start, end = self.styles[lasttype]
-                    outfile.write(''.join((start, lastval, end)))
-                lastval = value
-                lasttype = ttype
-
-        if lastval:
-            start, end = self.styles[lasttype]
-            outfile.write(''.join((start, lastval, end)))
-
-        if self._mono:
-            outfile.write('[/font]')
-        if self._code:
-            outfile.write('[/code]')
-        if self._code or self._mono:
-            outfile.write('\n')
diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/groff.py b/venv/lib/python3.10/site-packages/pygments/formatters/groff.py
deleted file mode 100644
index 028fec4ea49c5ba2eb78fb032fb98c886f4de90c..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/formatters/groff.py
+++ /dev/null
@@ -1,170 +0,0 @@
-"""
-    pygments.formatters.groff
-    ~~~~~~~~~~~~~~~~~~~~~~~~~
-
-    Formatter for groff output.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import math
-from pygments.formatter import Formatter
-from pygments.util import get_bool_opt, get_int_opt
-
-__all__ = ['GroffFormatter']
-
-
-class GroffFormatter(Formatter):
-    """
-    Format tokens with groff escapes to change their color and font style.
-
-    .. versionadded:: 2.11
-
-    Additional options accepted:
-
-    `style`
-        The style to use, can be a string or a Style subclass (default:
-        ``'default'``).
-
-    `monospaced`
-        If set to true, monospace font will be used (default: ``true``).
-
-    `linenos`
-        If set to true, print the line numbers (default: ``false``).
-
-    `wrap`
-        Wrap lines to the specified number of characters. Disabled if set to 0
-        (default: ``0``).
-    """
-
-    name = 'groff'
-    aliases = ['groff','troff','roff']
-    filenames = []
-
-    def __init__(self, **options):
-        Formatter.__init__(self, **options)
-
-        self.monospaced = get_bool_opt(options, 'monospaced', True)
-        self.linenos = get_bool_opt(options, 'linenos', False)
-        self._lineno = 0
-        self.wrap = get_int_opt(options, 'wrap', 0)
-        self._linelen = 0
-
-        self.styles = {}
-        self._make_styles()
-
-
-    def _make_styles(self):
-        regular = '\\f[CR]' if self.monospaced else '\\f[R]'
-        bold = '\\f[CB]' if self.monospaced else '\\f[B]'
-        italic = '\\f[CI]' if self.monospaced else '\\f[I]'
-
-        for ttype, ndef in self.style:
-            start = end = ''
-            if ndef['color']:
-                start += '\\m[{}]'.format(ndef['color'])
-                end = '\\m[]' + end
-            if ndef['bold']:
-                start += bold
-                end = regular + end
-            if ndef['italic']:
-                start += italic
-                end = regular + end
-            if ndef['bgcolor']:
-                start += '\\M[{}]'.format(ndef['bgcolor'])
-                end = '\\M[]' + end
-
-            self.styles[ttype] = start, end
-
-
-    def _define_colors(self, outfile):
-        colors = set()
-        for _, ndef in self.style:
-            if ndef['color'] is not None:
-                colors.add(ndef['color'])
-
-        for color in sorted(colors):
-            outfile.write('.defcolor ' + color + ' rgb #' + color + '\n')
-
-
-    def _write_lineno(self, outfile):
-        self._lineno += 1
-        outfile.write("%s% 4d " % (self._lineno != 1 and '\n' or '', self._lineno))
-
-
-    def _wrap_line(self, line):
-        length = len(line.rstrip('\n'))
-        space = '     ' if self.linenos else ''
-        newline = ''
-
-        if length > self.wrap:
-            for i in range(0, math.floor(length / self.wrap)):
-                chunk = line[i*self.wrap:i*self.wrap+self.wrap]
-                newline += (chunk + '\n' + space)
-            remainder = length % self.wrap
-            if remainder > 0:
-                newline += line[-remainder-1:]
-                self._linelen = remainder
-        elif self._linelen + length > self.wrap:
-            newline = ('\n' + space) + line
-            self._linelen = length
-        else:
-            newline = line
-            self._linelen += length
-
-        return newline
-
-
-    def _escape_chars(self, text):
-        text = text.replace('\\', '\\[u005C]'). \
-                    replace('.', '\\[char46]'). \
-                    replace('\'', '\\[u0027]'). \
-                    replace('`', '\\[u0060]'). \
-                    replace('~', '\\[u007E]')
-        copy = text
-
-        for char in copy:
-            if len(char) != len(char.encode()):
-                uni = char.encode('unicode_escape') \
-                    .decode()[1:] \
-                    .replace('x', 'u00') \
-                    .upper()
-                text = text.replace(char, '\\[u' + uni[1:] + ']')
-
-        return text
-
-
-    def format_unencoded(self, tokensource, outfile):
-        self._define_colors(outfile)
-
-        outfile.write('.nf\n\\f[CR]\n')
-
-        if self.linenos:
-            self._write_lineno(outfile)
-
-        for ttype, value in tokensource:
-            while ttype not in self.styles:
-                ttype = ttype.parent
-            start, end = self.styles[ttype]
-
-            for line in value.splitlines(True):
-                if self.wrap > 0:
-                    line = self._wrap_line(line)
-
-                if start and end:
-                    text = self._escape_chars(line.rstrip('\n'))
-                    if text != '':
-                        outfile.write(''.join((start, text, end)))
-                else:
-                    outfile.write(self._escape_chars(line.rstrip('\n')))
-
-                if line.endswith('\n'):
-                    if self.linenos:
-                        self._write_lineno(outfile)
-                        self._linelen = 0
-                    else:
-                        outfile.write('\n')
-                        self._linelen = 0
-
-        outfile.write('\n.fi')
diff --git a/venv/lib/python3.10/site-packages/pygments/formatters/html.py b/venv/lib/python3.10/site-packages/pygments/formatters/html.py
deleted file mode 100644
index 4ef183681fd1e325831caecffc531dcb305516f3..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/formatters/html.py
+++ /dev/null
@@ -1,995 +0,0 @@
-"""
-    pygments.formatters.html
-    ~~~~~~~~~~~~~~~~~~~~~~~~
-
-    Formatter for HTML output.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import functools
-import os
-import sys
-import os.path
-from io import StringIO
-
-from pygments.formatter import Formatter
-from pygments.token import Token, Text, STANDARD_TYPES
-from pygments.util import get_bool_opt, get_int_opt, get_list_opt
-
-try:
-    import ctags
-except ImportError:
-    ctags = None
-
-__all__ = ['HtmlFormatter']
-
-
-_escape_html_table = {
-    ord('&'): '&',
-    ord('<'): '<',
-    ord('>'): '>',
-    ord('"'): '"',
-    ord("'"): ''',
-}
-
-
-def escape_html(text, table=_escape_html_table):
-    """Escape &, <, > as well as single and double quotes for HTML."""
-    return text.translate(table)
-
-
-def webify(color):
-    if color.startswith('calc') or color.startswith('var'):
-        return color
-    else:
-        # Check if the color can be shortened from 6 to 3 characters
-        color = color.upper()
-        if (len(color) == 6 and
-            (    color[0] == color[1]
-             and color[2] == color[3]
-             and color[4] == color[5])):
-            return f'#{color[0]}{color[2]}{color[4]}'
-        else:
-            return f'#{color}'
-
-
-def _get_ttype_class(ttype):
-    fname = STANDARD_TYPES.get(ttype)
-    if fname:
-        return fname
-    aname = ''
-    while fname is None:
-        aname = '-' + ttype[-1] + aname
-        ttype = ttype.parent
-        fname = STANDARD_TYPES.get(ttype)
-    return fname + aname
-
-
-CSSFILE_TEMPLATE = '''\
-/*
-generated by Pygments 
-Copyright 2006-2025 by the Pygments team.
-Licensed under the BSD license, see LICENSE for details.
-*/
-%(styledefs)s
-'''
-
-DOC_HEADER = '''\
-
-
-
-
-  %(title)s
-  
-  
-
-
-
%(title)s

-
-'''
-
-DOC_HEADER_EXTERNALCSS = '''\
-
-
-
-
-  %(title)s
-  
-  
-
-
-
%(title)s

-
-'''
-
-DOC_FOOTER = '''\
-
-
-'''
-
-
-class HtmlFormatter(Formatter):
-    r"""
-    Format tokens as HTML 4 ```` tags. By default, the content is enclosed
-    in a ``
`` tag, itself wrapped in a ``
`` tag (but see the `nowrap` option).
-    The ``
``'s CSS class can be set by the `cssclass` option.
-
-    If the `linenos` option is set to ``"table"``, the ``
`` is
-    additionally wrapped inside a ```` which has one row and two
-    cells: one containing the line numbers and one containing the code.
-    Example:
-
-    .. sourcecode:: html
-
-        

-        

-          
-          
-        

-            
1
-            2

-          		
-            
def foo(bar):
-              pass
-            

-          

-
-    (whitespace added to improve clarity).
-
-    A list of lines can be specified using the `hl_lines` option to make these
-    lines highlighted (as of Pygments 0.11).
-
-    With the `full` option, a complete HTML 4 document is output, including
-    the style definitions inside a ``$)', _handle_cssblock),
-
-            include('keywords'),
-            include('inline'),
-        ],
-        'keywords': [
-            (words((
-                '\\define', '\\end', 'caption', 'created', 'modified', 'tags',
-                'title', 'type'), prefix=r'^', suffix=r'\b'),
-             Keyword),
-        ],
-        'inline': [
-            # escape
-            (r'\\.', Text),
-            # created or modified date
-            (r'\d{17}', Number.Integer),
-            # italics
-            (r'(\s)(//[^/]+//)((?=\W|\n))',
-             bygroups(Text, Generic.Emph, Text)),
-            # superscript
-            (r'(\s)(\^\^[^\^]+\^\^)', bygroups(Text, Generic.Emph)),
-            # subscript
-            (r'(\s)(,,[^,]+,,)', bygroups(Text, Generic.Emph)),
-            # underscore
-            (r'(\s)(__[^_]+__)', bygroups(Text, Generic.Strong)),
-            # bold
-            (r"(\s)(''[^']+'')((?=\W|\n))",
-             bygroups(Text, Generic.Strong, Text)),
-            # strikethrough
-            (r'(\s)(~~[^~]+~~)((?=\W|\n))',
-             bygroups(Text, Generic.Deleted, Text)),
-            # TiddlyWiki variables
-            (r'<<[^>]+>>', Name.Tag),
-            (r'\$\$[^$]+\$\$', Name.Tag),
-            (r'\$\([^)]+\)\$', Name.Tag),
-            # TiddlyWiki style or class
-            (r'^@@.*$', Name.Tag),
-            # HTML tags
-            (r']+>', Name.Tag),
-            # inline code
-            (r'`[^`]+`', String.Backtick),
-            # HTML escaped symbols
-            (r'&\S*?;', String.Regex),
-            # Wiki links
-            (r'(\[{2})([^]\|]+)(\]{2})', bygroups(Text, Name.Tag, Text)),
-            # External links
-            (r'(\[{2})([^]\|]+)(\|)([^]\|]+)(\]{2})',
-            bygroups(Text, Name.Tag, Text, Name.Attribute, Text)),
-            # Transclusion
-            (r'(\{{2})([^}]+)(\}{2})', bygroups(Text, Name.Tag, Text)),
-            # URLs
-            (r'(\b.?.?tps?://[^\s"]+)', bygroups(Name.Attribute)),
-
-            # general text, must come last!
-            (r'[\w]+', Text),
-            (r'.', Text)
-        ],
-    }
-
-    def __init__(self, **options):
-        self.handlecodeblocks = get_bool_opt(options, 'handlecodeblocks', True)
-        RegexLexer.__init__(self, **options)
-
-
-class WikitextLexer(RegexLexer):
-    """
-    For MediaWiki Wikitext.
-
-    Parsing Wikitext is tricky, and results vary between different MediaWiki
-    installations, so we only highlight common syntaxes (built-in or from
-    popular extensions), and also assume templates produce no unbalanced
-    syntaxes.
-    """
-    name = 'Wikitext'
-    url = 'https://www.mediawiki.org/wiki/Wikitext'
-    aliases = ['wikitext', 'mediawiki']
-    filenames = []
-    mimetypes = ['text/x-wiki']
-    version_added = '2.15'
-    flags = re.MULTILINE
-
-    def nowiki_tag_rules(tag_name):
-        return [
-            (rf'(?i)()', bygroups(Punctuation,
-             Name.Tag, Whitespace, Punctuation), '#pop'),
-            include('entity'),
-            include('text'),
-        ]
-
-    def plaintext_tag_rules(tag_name):
-        return [
-            (rf'(?si)(.*?)()', bygroups(Text,
-             Punctuation, Name.Tag, Whitespace, Punctuation), '#pop'),
-        ]
-
-    def delegate_tag_rules(tag_name, lexer, **lexer_kwargs):
-        return [
-            (rf'(?i)()', bygroups(Punctuation,
-             Name.Tag, Whitespace, Punctuation), '#pop'),
-            (rf'(?si).+?(?=)', using(lexer, **lexer_kwargs)),
-        ]
-
-    def text_rules(token):
-        return [
-            (r'\w+', token),
-            (r'[^\S\n]+', token),
-            (r'(?s).', token),
-        ]
-
-    def handle_syntaxhighlight(self, match, ctx):
-        from pygments.lexers import get_lexer_by_name
-
-        attr_content = match.group()
-        start = 0
-        index = 0
-        while True:
-            index = attr_content.find('>', start)
-            # Exclude comment end (-->)
-            if attr_content[index-2:index] != '--':
-                break
-            start = index + 1
-
-        if index == -1:
-            # No tag end
-            yield from self.get_tokens_unprocessed(attr_content, stack=['root', 'attr'])
-            return
-        attr = attr_content[:index]
-        yield from self.get_tokens_unprocessed(attr, stack=['root', 'attr'])
-        yield match.start(3) + index, Punctuation, '>'
-
-        lexer = None
-        content = attr_content[index+1:]
-        lang_match = re.findall(r'\blang=("|\'|)(\w+)(\1)', attr)
-
-        if len(lang_match) >= 1:
-            # Pick the last match in case of multiple matches
-            lang = lang_match[-1][1]
-            try:
-                lexer = get_lexer_by_name(lang)
-            except ClassNotFound:
-                pass
-
-        if lexer is None:
-            yield match.start() + index + 1, Text, content
-        else:
-            yield from lexer.get_tokens_unprocessed(content)
-
-    def handle_score(self, match, ctx):
-        attr_content = match.group()
-        start = 0
-        index = 0
-        while True:
-            index = attr_content.find('>', start)
-            # Exclude comment end (-->)
-            if attr_content[index-2:index] != '--':
-                break
-            start = index + 1
-
-        if index == -1:
-            # No tag end
-            yield from self.get_tokens_unprocessed(attr_content, stack=['root', 'attr'])
-            return
-        attr = attr_content[:index]
-        content = attr_content[index+1:]
-        yield from self.get_tokens_unprocessed(attr, stack=['root', 'attr'])
-        yield match.start(3) + index, Punctuation, '>'
-
-        lang_match = re.findall(r'\blang=("|\'|)(\w+)(\1)', attr)
-        # Pick the last match in case of multiple matches
-        lang = lang_match[-1][1] if len(lang_match) >= 1 else 'lilypond'
-
-        if lang == 'lilypond':  # Case sensitive
-            yield from LilyPondLexer().get_tokens_unprocessed(content)
-        else:  # ABC
-            # FIXME: Use ABC lexer in the future
-            yield match.start() + index + 1, Text, content
-
-    # a-z removed to prevent linter from complaining, REMEMBER to use (?i)
-    title_char = r' %!"$&\'()*,\-./0-9:;=?@A-Z\\\^_`~+\u0080-\uFFFF'
-    nbsp_char = r'(?:\t| |&\#0*160;|&\#[Xx]0*[Aa]0;|[ \xA0\u1680\u2000-\u200A\u202F\u205F\u3000])'
-    link_address = r'(?:[0-9.]+|\[[0-9a-f:.]+\]|[^\x00-\x20"<>\[\]\x7F\xA0\u1680\u2000-\u200A\u202F\u205F\u3000\uFFFD])'
-    link_char_class = r'[^\x00-\x20"<>\[\]\x7F\xA0\u1680\u2000-\u200A\u202F\u205F\u3000\uFFFD]'
-    double_slashes_i = {
-        '__FORCETOC__', '__NOCONTENTCONVERT__', '__NOCC__', '__NOEDITSECTION__', '__NOGALLERY__',
-        '__NOTITLECONVERT__', '__NOTC__', '__NOTOC__', '__TOC__',
-    }
-    double_slashes = {
-        '__EXPECTUNUSEDCATEGORY__',  '__HIDDENCAT__', '__INDEX__',  '__NEWSECTIONLINK__',
-        '__NOINDEX__',  '__NONEWSECTIONLINK__',  '__STATICREDIRECT__', '__NOGLOBAL__',
-        '__DISAMBIG__', '__EXPECTED_UNCONNECTED_PAGE__',
-    }
-    protocols = {
-        'bitcoin:', 'ftp://', 'ftps://', 'geo:', 'git://', 'gopher://', 'http://', 'https://',
-        'irc://', 'ircs://', 'magnet:', 'mailto:', 'mms://', 'news:', 'nntp://', 'redis://',
-        'sftp://', 'sip:', 'sips:', 'sms:', 'ssh://', 'svn://', 'tel:', 'telnet://', 'urn:',
-        'worldwind://', 'xmpp:', '//',
-    }
-    non_relative_protocols = protocols - {'//'}
-    html_tags = {
-        'abbr', 'b', 'bdi', 'bdo', 'big', 'blockquote', 'br', 'caption', 'center', 'cite', 'code',
-        'data', 'dd', 'del', 'dfn', 'div', 'dl', 'dt', 'em', 'font', 'h1', 'h2', 'h3', 'h4', 'h5',
-        'h6', 'hr', 'i', 'ins', 'kbd', 'li', 'link', 'mark', 'meta', 'ol', 'p', 'q', 'rb', 'rp',
-        'rt', 'rtc', 'ruby', 's', 'samp', 'small', 'span', 'strike', 'strong', 'sub', 'sup',
-        'table', 'td', 'th', 'time', 'tr', 'tt', 'u', 'ul', 'var', 'wbr',
-    }
-    parser_tags = {
-        'graph', 'charinsert', 'rss', 'chem', 'categorytree', 'nowiki', 'inputbox', 'math',
-        'hiero', 'score', 'pre', 'ref', 'translate', 'imagemap', 'templatestyles', 'languages',
-        'noinclude', 'mapframe', 'section', 'poem', 'syntaxhighlight', 'includeonly', 'tvar',
-        'onlyinclude', 'templatedata', 'langconvert', 'timeline', 'dynamicpagelist', 'gallery',
-        'maplink', 'ce', 'references',
-    }
-    variant_langs = {
-        # ZhConverter.php
-        'zh', 'zh-hans', 'zh-hant', 'zh-cn', 'zh-hk', 'zh-mo', 'zh-my', 'zh-sg', 'zh-tw',
-        # WuuConverter.php
-        'wuu', 'wuu-hans', 'wuu-hant',
-        # UzConverter.php
-        'uz', 'uz-latn', 'uz-cyrl',
-        # TlyConverter.php
-        'tly', 'tly-cyrl',
-        # TgConverter.php
-        'tg', 'tg-latn',
-        # SrConverter.php
-        'sr', 'sr-ec', 'sr-el',
-        # ShiConverter.php
-        'shi', 'shi-tfng', 'shi-latn',
-        # ShConverter.php
-        'sh-latn', 'sh-cyrl',
-        # KuConverter.php
-        'ku', 'ku-arab', 'ku-latn',
-        # IuConverter.php
-        'iu', 'ike-cans', 'ike-latn',
-        # GanConverter.php
-        'gan', 'gan-hans', 'gan-hant',
-        # EnConverter.php
-        'en', 'en-x-piglatin',
-        # CrhConverter.php
-        'crh', 'crh-cyrl', 'crh-latn',
-        # BanConverter.php
-        'ban', 'ban-bali', 'ban-x-dharma', 'ban-x-palmleaf', 'ban-x-pku',
-    }
-    magic_vars_i = {
-        'ARTICLEPATH', 'INT', 'PAGEID', 'SCRIPTPATH', 'SERVER', 'SERVERNAME', 'STYLEPATH',
-    }
-    magic_vars = {
-        '!', '=', 'BASEPAGENAME', 'BASEPAGENAMEE', 'CASCADINGSOURCES', 'CONTENTLANGUAGE',
-        'CONTENTLANG', 'CURRENTDAY', 'CURRENTDAY2', 'CURRENTDAYNAME', 'CURRENTDOW', 'CURRENTHOUR',
-        'CURRENTMONTH', 'CURRENTMONTH2', 'CURRENTMONTH1', 'CURRENTMONTHABBREV', 'CURRENTMONTHNAME',
-        'CURRENTMONTHNAMEGEN', 'CURRENTTIME', 'CURRENTTIMESTAMP', 'CURRENTVERSION', 'CURRENTWEEK',
-        'CURRENTYEAR', 'DIRECTIONMARK', 'DIRMARK', 'FULLPAGENAME', 'FULLPAGENAMEE', 'LOCALDAY',
-        'LOCALDAY2', 'LOCALDAYNAME', 'LOCALDOW', 'LOCALHOUR', 'LOCALMONTH', 'LOCALMONTH2',
-        'LOCALMONTH1', 'LOCALMONTHABBREV', 'LOCALMONTHNAME', 'LOCALMONTHNAMEGEN', 'LOCALTIME',
-        'LOCALTIMESTAMP', 'LOCALWEEK', 'LOCALYEAR', 'NAMESPACE', 'NAMESPACEE', 'NAMESPACENUMBER',
-        'NUMBEROFACTIVEUSERS', 'NUMBEROFADMINS', 'NUMBEROFARTICLES', 'NUMBEROFEDITS',
-        'NUMBEROFFILES', 'NUMBEROFPAGES', 'NUMBEROFUSERS', 'PAGELANGUAGE', 'PAGENAME', 'PAGENAMEE',
-        'REVISIONDAY', 'REVISIONDAY2', 'REVISIONID', 'REVISIONMONTH', 'REVISIONMONTH1',
-        'REVISIONSIZE', 'REVISIONTIMESTAMP', 'REVISIONUSER', 'REVISIONYEAR', 'ROOTPAGENAME',
-        'ROOTPAGENAMEE', 'SITENAME', 'SUBJECTPAGENAME', 'ARTICLEPAGENAME', 'SUBJECTPAGENAMEE',
-        'ARTICLEPAGENAMEE', 'SUBJECTSPACE', 'ARTICLESPACE', 'SUBJECTSPACEE', 'ARTICLESPACEE',
-        'SUBPAGENAME', 'SUBPAGENAMEE', 'TALKPAGENAME', 'TALKPAGENAMEE', 'TALKSPACE', 'TALKSPACEE',
-    }
-    parser_functions_i = {
-        'ANCHORENCODE', 'BIDI', 'CANONICALURL', 'CANONICALURLE', 'FILEPATH', 'FORMATNUM',
-        'FULLURL', 'FULLURLE', 'GENDER', 'GRAMMAR', 'INT', r'\#LANGUAGE', 'LC', 'LCFIRST', 'LOCALURL',
-        'LOCALURLE', 'NS', 'NSE', 'PADLEFT', 'PADRIGHT', 'PAGEID', 'PLURAL', 'UC', 'UCFIRST',
-        'URLENCODE',
-    }
-    parser_functions = {
-        'BASEPAGENAME', 'BASEPAGENAMEE', 'CASCADINGSOURCES', 'DEFAULTSORT', 'DEFAULTSORTKEY',
-        'DEFAULTCATEGORYSORT', 'FULLPAGENAME', 'FULLPAGENAMEE', 'NAMESPACE', 'NAMESPACEE',
-        'NAMESPACENUMBER', 'NUMBERINGROUP', 'NUMINGROUP', 'NUMBEROFACTIVEUSERS', 'NUMBEROFADMINS',
-        'NUMBEROFARTICLES', 'NUMBEROFEDITS', 'NUMBEROFFILES', 'NUMBEROFPAGES', 'NUMBEROFUSERS',
-        'PAGENAME', 'PAGENAMEE', 'PAGESINCATEGORY', 'PAGESINCAT', 'PAGESIZE', 'PROTECTIONEXPIRY',
-        'PROTECTIONLEVEL', 'REVISIONDAY', 'REVISIONDAY2', 'REVISIONID', 'REVISIONMONTH',
-        'REVISIONMONTH1', 'REVISIONTIMESTAMP', 'REVISIONUSER', 'REVISIONYEAR', 'ROOTPAGENAME',
-        'ROOTPAGENAMEE', 'SUBJECTPAGENAME', 'ARTICLEPAGENAME', 'SUBJECTPAGENAMEE',
-        'ARTICLEPAGENAMEE', 'SUBJECTSPACE', 'ARTICLESPACE', 'SUBJECTSPACEE', 'ARTICLESPACEE',
-        'SUBPAGENAME', 'SUBPAGENAMEE', 'TALKPAGENAME', 'TALKPAGENAMEE', 'TALKSPACE', 'TALKSPACEE',
-        'INT', 'DISPLAYTITLE', 'PAGESINNAMESPACE', 'PAGESINNS',
-    }
-
-    tokens = {
-        'root': [
-            # Redirects
-            (r"""(?xi)
-                (\A\s*?)(\#REDIRECT:?) # may contain a colon
-                (\s+)(\[\[) (?=[^\]\n]* \]\]$)
-            """,
-             bygroups(Whitespace, Keyword, Whitespace, Punctuation), 'redirect-inner'),
-            # Subheadings
-            (r'^(={2,6})(.+?)(\1)(\s*$\n)',
-             bygroups(Generic.Subheading, Generic.Subheading, Generic.Subheading, Whitespace)),
-            # Headings
-            (r'^(=.+?=)(\s*$\n)',
-             bygroups(Generic.Heading, Whitespace)),
-            # Double-slashed magic words
-            (words(double_slashes_i, prefix=r'(?i)'), Name.Function.Magic),
-            (words(double_slashes), Name.Function.Magic),
-            # Raw URLs
-            (r'(?i)\b(?:{}){}{}*'.format('|'.join(protocols),
-             link_address, link_char_class), Name.Label),
-            # Magic links
-            (rf'\b(?:RFC|PMID){nbsp_char}+[0-9]+\b',
-             Name.Function.Magic),
-            (r"""(?x)
-                \bISBN {nbsp_char}
-                (?: 97[89] {nbsp_dash}? )?
-                (?: [0-9] {nbsp_dash}? ){{9}} # escape format()
-                [0-9Xx]\b
-            """.format(nbsp_char=nbsp_char, nbsp_dash=f'(?:-|{nbsp_char})'), Name.Function.Magic),
-            include('list'),
-            include('inline'),
-            include('text'),
-        ],
-        'redirect-inner': [
-            (r'(\]\])(\s*?\n)', bygroups(Punctuation, Whitespace), '#pop'),
-            (r'(\#)([^#]*?)', bygroups(Punctuation, Name.Label)),
-            (rf'(?i)[{title_char}]+', Name.Tag),
-        ],
-        'list': [
-            # Description lists
-            (r'^;', Keyword, 'dt'),
-            # Ordered lists, unordered lists and indents
-            (r'^[#:*]+', Keyword),
-            # Horizontal rules
-            (r'^-{4,}', Keyword),
-        ],
-        'inline': [
-            # Signatures
-            (r'~{3,5}', Keyword),
-            # Entities
-            include('entity'),
-            # Bold & italic
-            (r"('')(''')(?!')", bygroups(Generic.Emph,
-             Generic.EmphStrong), 'inline-italic-bold'),
-            (r"'''(?!')", Generic.Strong, 'inline-bold'),
-            (r"''(?!')", Generic.Emph, 'inline-italic'),
-            # Comments & parameters & templates
-            include('replaceable'),
-            # Media links
-            (
-                r"""(?xi)
-                (\[\[)
-                    (File|Image) (:)
-                    ((?: [{}] | \{{{{2,3}}[^{{}}]*?\}}{{2,3}} |  )*)
-                    (?: (\#) ([{}]*?) )?
-                """.format(title_char, f'{title_char}#'),
-                bygroups(Punctuation, Name.Namespace,  Punctuation,
-                         using(this, state=['wikilink-name']), Punctuation, Name.Label),
-                'medialink-inner'
-            ),
-            # Wikilinks
-            (
-                r"""(?xi)
-                (\[\[)(?!{}) # Should not contain URLs
-                    (?: ([{}]*) (:))?
-                    ((?: [{}] | \{{{{2,3}}[^{{}}]*?\}}{{2,3}} |  )*?)
-                    (?: (\#) ([{}]*?) )?
-                (\]\])
-                """.format('|'.join(protocols), title_char.replace('/', ''),
-                       title_char, f'{title_char}#'),
-                bygroups(Punctuation, Name.Namespace,  Punctuation,
-                         using(this, state=['wikilink-name']), Punctuation, Name.Label, Punctuation)
-            ),
-            (
-                r"""(?xi)
-                (\[\[)(?!{})
-                    (?: ([{}]*) (:))?
-                    ((?: [{}] | \{{{{2,3}}[^{{}}]*?\}}{{2,3}} |  )*?)
-                    (?: (\#) ([{}]*?) )?
-                    (\|)
-                """.format('|'.join(protocols), title_char.replace('/', ''),
-                       title_char, f'{title_char}#'),
-                bygroups(Punctuation, Name.Namespace,  Punctuation,
-                         using(this, state=['wikilink-name']), Punctuation, Name.Label, Punctuation),
-                'wikilink-inner'
-            ),
-            # External links
-            (
-                r"""(?xi)
-                (\[)
-                    ((?:{}) {} {}*)
-                    (\s*)
-                """.format('|'.join(protocols), link_address, link_char_class),
-                bygroups(Punctuation, Name.Label, Whitespace),
-                'extlink-inner'
-            ),
-            # Tables
-            (r'^(:*)(\s*?)(\{\|)([^\n]*)$', bygroups(Keyword,
-             Whitespace, Punctuation, using(this, state=['root', 'attr'])), 'table'),
-            # HTML tags
-            (r'(?i)(<)({})\b'.format('|'.join(html_tags)),
-             bygroups(Punctuation, Name.Tag), 'tag-inner-ordinary'),
-            (r'(?i)()'.format('|'.join(html_tags)),
-             bygroups(Punctuation, Name.Tag, Whitespace, Punctuation)),
-            # 
-            (r'(?i)(<)(nowiki)\b', bygroups(Punctuation,
-             Name.Tag), ('tag-nowiki', 'tag-inner')),
-            # 
-            (r'(?i)(<)(pre)\b', bygroups(Punctuation,
-             Name.Tag), ('tag-pre', 'tag-inner')),
-            # 
-            (r'(?i)(<)(categorytree)\b', bygroups(
-                Punctuation, Name.Tag), ('tag-categorytree', 'tag-inner')),
-            # 
-            (r'(?i)(<)(hiero)\b', bygroups(Punctuation,
-             Name.Tag), ('tag-hiero', 'tag-inner')),
-            # 
-            (r'(?i)(<)(math)\b', bygroups(Punctuation,
-             Name.Tag), ('tag-math', 'tag-inner')),
-            # 
-            (r'(?i)(<)(chem)\b', bygroups(Punctuation,
-             Name.Tag), ('tag-chem', 'tag-inner')),
-            # 
-            (r'(?i)(<)(ce)\b', bygroups(Punctuation,
-             Name.Tag), ('tag-ce', 'tag-inner')),
-            # 
-            (r'(?i)(<)(charinsert)\b', bygroups(
-                Punctuation, Name.Tag), ('tag-charinsert', 'tag-inner')),
-            # 
-            (r'(?i)(<)(templatedata)\b', bygroups(
-                Punctuation, Name.Tag), ('tag-templatedata', 'tag-inner')),
-            # 
-            (r'(?i)(<)(gallery)\b', bygroups(
-                Punctuation, Name.Tag), ('tag-gallery', 'tag-inner')),
-            # 
-            (r'(?i)(<)(gallery)\b', bygroups(
-                Punctuation, Name.Tag), ('tag-graph', 'tag-inner')),
-            # 
-            (r'(?i)(<)(dynamicpagelist)\b', bygroups(
-                Punctuation, Name.Tag), ('tag-dynamicpagelist', 'tag-inner')),
-            # 
-            (r'(?i)(<)(inputbox)\b', bygroups(
-                Punctuation, Name.Tag), ('tag-inputbox', 'tag-inner')),
-            # 
-            (r'(?i)(<)(rss)\b', bygroups(
-                Punctuation, Name.Tag), ('tag-rss', 'tag-inner')),
-            # 
-            (r'(?i)(<)(imagemap)\b', bygroups(
-                Punctuation, Name.Tag), ('tag-imagemap', 'tag-inner')),
-            # 
-            (r'(?i)()',
-             bygroups(Punctuation, Name.Tag, Whitespace, Punctuation)),
-            (r'(?si)(<)(syntaxhighlight)\b([^>]*?(?.*?)(?=)',
-             bygroups(Punctuation, Name.Tag, handle_syntaxhighlight)),
-            # : Fallback case for self-closing tags
-            (r'(?i)(<)(syntaxhighlight)\b(\s*?)((?:[^>]|-->)*?)(/\s*?(?)*?)(/\s*?(?)*?)(/\s*?(?|\Z)', Comment.Multiline),
-            # Parameters
-            (
-                r"""(?x)
-                (\{{3})
-                    ([^|]*?)
-                    (?=\}{3}|\|)
-                """,
-                bygroups(Punctuation, Name.Variable),
-                'parameter-inner',
-            ),
-            # Magic variables
-            (r'(?i)(\{{\{{)(\s*)({})(\s*)(\}}\}})'.format('|'.join(magic_vars_i)),
-             bygroups(Punctuation, Whitespace, Name.Function, Whitespace, Punctuation)),
-            (r'(\{{\{{)(\s*)({})(\s*)(\}}\}})'.format('|'.join(magic_vars)),
-                bygroups(Punctuation, Whitespace, Name.Function, Whitespace, Punctuation)),
-            # Parser functions & templates
-            (r'\{\{', Punctuation, 'template-begin-space'),
-            #  legacy syntax
-            (r'(?i)(<)(tvar)\b(\|)([^>]*?)(>)', bygroups(Punctuation,
-             Name.Tag, Punctuation, String, Punctuation)),
-            (r'', Punctuation, '#pop'),
-            # 
-            (r'(?i)(<)(tvar)\b', bygroups(Punctuation, Name.Tag), 'tag-inner-ordinary'),
-            (r'(?i)()',
-             bygroups(Punctuation, Name.Tag, Whitespace, Punctuation)),
-        ],
-        'parameter-inner': [
-            (r'\}{3}', Punctuation, '#pop'),
-            (r'\|', Punctuation),
-            include('inline'),
-            include('text'),
-        ],
-        'template-begin-space': [
-            # Templates allow line breaks at the beginning, and due to how MediaWiki handles
-            # comments, an extra state is required to handle things like {{\n\n name}}
-            (r'|\Z)', Comment.Multiline),
-            (r'\s+', Whitespace),
-            # Parser functions
-            (
-                r'(?i)(\#[{}]*?|{})(:)'.format(title_char,
-                                           '|'.join(parser_functions_i)),
-                bygroups(Name.Function, Punctuation), ('#pop', 'template-inner')
-            ),
-            (
-                r'({})(:)'.format('|'.join(parser_functions)),
-                bygroups(Name.Function, Punctuation), ('#pop', 'template-inner')
-            ),
-            # Templates
-            (
-                rf'(?i)([{title_char}]*?)(:)',
-                bygroups(Name.Namespace, Punctuation), ('#pop', 'template-name')
-            ),
-            default(('#pop', 'template-name'),),
-        ],
-        'template-name': [
-            (r'(\s*?)(\|)', bygroups(Text, Punctuation), ('#pop', 'template-inner')),
-            (r'\}\}', Punctuation, '#pop'),
-            (r'\n', Text, '#pop'),
-            include('replaceable'),
-            *text_rules(Name.Tag),
-        ],
-        'template-inner': [
-            (r'\}\}', Punctuation, '#pop'),
-            (r'\|', Punctuation),
-            (
-                r"""(?x)
-                    (?<=\|)
-                    ( (?: (?! \{\{ | \}\} )[^=\|<])*? ) # Exclude templates and tags
-                    (=)
-                """,
-                bygroups(Name.Label, Operator)
-            ),
-            include('inline'),
-            include('text'),
-        ],
-        'table': [
-            # Use [ \t\n\r\0\x0B] instead of \s to follow PHP trim() behavior
-            # Endings
-            (r'^([ \t\n\r\0\x0B]*?)(\|\})',
-             bygroups(Whitespace, Punctuation), '#pop'),
-            # Table rows
-            (r'^([ \t\n\r\0\x0B]*?)(\|-+)(.*)$', bygroups(Whitespace, Punctuation,
-             using(this, state=['root', 'attr']))),
-            # Captions
-            (
-                r"""(?x)
-                ^([ \t\n\r\0\x0B]*?)(\|\+)
-                # Exclude links, template and tags
-                (?: ( (?: (?! \[\[ | \{\{ )[^|\n<] )*? )(\|) )?
-                (.*?)$
-                """,
-                bygroups(Whitespace, Punctuation, using(this, state=[
-                         'root', 'attr']), Punctuation, Generic.Heading),
-            ),
-            # Table data
-            (
-                r"""(?x)
-                ( ^(?:[ \t\n\r\0\x0B]*?)\| | \|\| )
-                (?: ( (?: (?! \[\[ | \{\{ )[^|\n<] )*? )(\|)(?!\|) )?
-                """,
-                bygroups(Punctuation, using(this, state=[
-                         'root', 'attr']), Punctuation),
-            ),
-            # Table headers
-            (
-                r"""(?x)
-                ( ^(?:[ \t\n\r\0\x0B]*?)!  )
-                (?: ( (?: (?! \[\[ | \{\{ )[^|\n<] )*? )(\|)(?!\|) )?
-                """,
-                bygroups(Punctuation, using(this, state=[
-                         'root', 'attr']), Punctuation),
-                'table-header',
-            ),
-            include('list'),
-            include('inline'),
-            include('text'),
-        ],
-        'table-header': [
-            # Requires another state for || handling inside headers
-            (r'\n', Text, '#pop'),
-            (
-                r"""(?x)
-                (!!|\|\|)
-                (?:
-                    ( (?: (?! \[\[ | \{\{ )[^|\n<] )*? )
-                    (\|)(?!\|)
-                )?
-                """,
-                bygroups(Punctuation, using(this, state=[
-                         'root', 'attr']), Punctuation)
-            ),
-            *text_rules(Generic.Subheading),
-        ],
-        'entity': [
-            (r'&\S*?;', Name.Entity),
-        ],
-        'dt': [
-            (r'\n', Text, '#pop'),
-            include('inline'),
-            (r':', Keyword, '#pop'),
-            include('text'),
-        ],
-        'extlink-inner': [
-            (r'\]', Punctuation, '#pop'),
-            include('inline'),
-            include('text'),
-        ],
-        'nowiki-ish': [
-            include('entity'),
-            include('text'),
-        ],
-        'attr': [
-            include('replaceable'),
-            (r'\s+', Whitespace),
-            (r'(=)(\s*)(")', bygroups(Operator, Whitespace, String.Double), 'attr-val-2'),
-            (r"(=)(\s*)(')", bygroups(Operator, Whitespace, String.Single), 'attr-val-1'),
-            (r'(=)(\s*)', bygroups(Operator, Whitespace), 'attr-val-0'),
-            (r'[\w:-]+', Name.Attribute),
-
-        ],
-        'attr-val-0': [
-            (r'\s', Whitespace, '#pop'),
-            include('replaceable'),
-            *text_rules(String),
-        ],
-        'attr-val-1': [
-            (r"'", String.Single, '#pop'),
-            include('replaceable'),
-            *text_rules(String.Single),
-        ],
-        'attr-val-2': [
-            (r'"', String.Double, '#pop'),
-            include('replaceable'),
-            *text_rules(String.Double),
-        ],
-        'tag-inner-ordinary': [
-            (r'/?\s*>', Punctuation, '#pop'),
-            include('tag-attr'),
-        ],
-        'tag-inner': [
-            # Return to root state for self-closing tags
-            (r'/\s*>', Punctuation, '#pop:2'),
-            (r'\s*>', Punctuation, '#pop'),
-            include('tag-attr'),
-        ],
-        # There states below are just like their non-tag variants, the key difference is
-        # they forcibly quit when encountering tag closing markup
-        'tag-attr': [
-            include('replaceable'),
-            (r'\s+', Whitespace),
-            (r'(=)(\s*)(")', bygroups(Operator,
-             Whitespace, String.Double), 'tag-attr-val-2'),
-            (r"(=)(\s*)(')", bygroups(Operator,
-             Whitespace, String.Single), 'tag-attr-val-1'),
-            (r'(=)(\s*)', bygroups(Operator, Whitespace), 'tag-attr-val-0'),
-            (r'[\w:-]+', Name.Attribute),
-
-        ],
-        'tag-attr-val-0': [
-            (r'\s', Whitespace, '#pop'),
-            (r'/?>', Punctuation, '#pop:2'),
-            include('replaceable'),
-            *text_rules(String),
-        ],
-        'tag-attr-val-1': [
-            (r"'", String.Single, '#pop'),
-            (r'/?>', Punctuation, '#pop:2'),
-            include('replaceable'),
-            *text_rules(String.Single),
-        ],
-        'tag-attr-val-2': [
-            (r'"', String.Double, '#pop'),
-            (r'/?>', Punctuation, '#pop:2'),
-            include('replaceable'),
-            *text_rules(String.Double),
-        ],
-        'tag-nowiki': nowiki_tag_rules('nowiki'),
-        'tag-pre': nowiki_tag_rules('pre'),
-        'tag-categorytree': plaintext_tag_rules('categorytree'),
-        'tag-dynamicpagelist': plaintext_tag_rules('dynamicpagelist'),
-        'tag-hiero': plaintext_tag_rules('hiero'),
-        'tag-inputbox': plaintext_tag_rules('inputbox'),
-        'tag-imagemap': plaintext_tag_rules('imagemap'),
-        'tag-charinsert': plaintext_tag_rules('charinsert'),
-        'tag-timeline': plaintext_tag_rules('timeline'),
-        'tag-gallery': plaintext_tag_rules('gallery'),
-        'tag-graph': plaintext_tag_rules('graph'),
-        'tag-rss': plaintext_tag_rules('rss'),
-        'tag-math': delegate_tag_rules('math', TexLexer, state='math'),
-        'tag-chem': delegate_tag_rules('chem', TexLexer, state='math'),
-        'tag-ce': delegate_tag_rules('ce', TexLexer, state='math'),
-        'tag-templatedata': delegate_tag_rules('templatedata', JsonLexer),
-        'text-italic': text_rules(Generic.Emph),
-        'text-bold': text_rules(Generic.Strong),
-        'text-bold-italic': text_rules(Generic.EmphStrong),
-        'text': text_rules(Text),
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/math.py b/venv/lib/python3.10/site-packages/pygments/lexers/math.py
deleted file mode 100644
index b225ffcf935a31fca3a909f84b949fd7acb2cd90..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/math.py
+++ /dev/null
@@ -1,21 +0,0 @@
-"""
-    pygments.lexers.math
-    ~~~~~~~~~~~~~~~~~~~~
-
-    Just export lexers that were contained in this module.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-# ruff: noqa: F401
-from pygments.lexers.python import NumPyLexer
-from pygments.lexers.matlab import MatlabLexer, MatlabSessionLexer, \
-    OctaveLexer, ScilabLexer
-from pygments.lexers.julia import JuliaLexer, JuliaConsoleLexer
-from pygments.lexers.r import RConsoleLexer, SLexer, RdLexer
-from pygments.lexers.modeling import BugsLexer, JagsLexer, StanLexer
-from pygments.lexers.idl import IDLLexer
-from pygments.lexers.algebra import MuPADLexer
-
-__all__ = []
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/matlab.py b/venv/lib/python3.10/site-packages/pygments/lexers/matlab.py
deleted file mode 100644
index 8eeffc9d23a0ad2b5ec186a8e3da29a02340b95e..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/matlab.py
+++ /dev/null
@@ -1,3307 +0,0 @@
-"""
-    pygments.lexers.matlab
-    ~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Matlab and related languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import Lexer, RegexLexer, bygroups, default, words, \
-    do_insertions, include
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Generic, Whitespace
-
-from pygments.lexers import _scilab_builtins
-
-__all__ = ['MatlabLexer', 'MatlabSessionLexer', 'OctaveLexer', 'ScilabLexer']
-
-
-class MatlabLexer(RegexLexer):
-    """
-    For Matlab source code.
-    """
-    name = 'Matlab'
-    aliases = ['matlab']
-    filenames = ['*.m']
-    mimetypes = ['text/matlab']
-    url = 'https://www.mathworks.com/products/matlab.html'
-    version_added = '0.10'
-
-    _operators = r'-|==|~=|<=|>=|<|>|&&|&|~|\|\|?|\.\*|\*|\+|\.\^|\^|\.\\|\./|/|\\'
-
-    tokens = {
-        'expressions': [
-            # operators:
-            (_operators, Operator),
-
-            # numbers (must come before punctuation to handle `.5`; cannot use
-            # `\b` due to e.g. `5. + .5`).  The negative lookahead on operators
-            # avoids including the dot in `1./x` (the dot is part of `./`).
-            (rf'(? and then
-            # (equal | open-parenthesis |  | ).
-            (rf'(?:^|(?<=;))(\s*)(\w+)(\s+)(?!=|\(|{_operators}\s|\s)',
-             bygroups(Whitespace, Name, Whitespace), 'commandargs'),
-
-            include('expressions')
-        ],
-        'blockcomment': [
-            (r'^\s*%\}', Comment.Multiline, '#pop'),
-            (r'^.*\n', Comment.Multiline),
-            (r'.', Comment.Multiline),
-        ],
-        'deffunc': [
-            (r'(\s*)(?:(\S+)(\s*)(=)(\s*))?(.+)(\()(.*)(\))(\s*)',
-             bygroups(Whitespace, Text, Whitespace, Punctuation,
-                      Whitespace, Name.Function, Punctuation, Text,
-                      Punctuation, Whitespace), '#pop'),
-            # function with no args
-            (r'(\s*)([a-zA-Z_]\w*)',
-             bygroups(Whitespace, Name.Function), '#pop'),
-        ],
-        'propattrs': [
-            (r'(\w+)(\s*)(=)(\s*)(\d+)',
-             bygroups(Name.Builtin, Whitespace, Punctuation, Whitespace,
-                      Number)),
-            (r'(\w+)(\s*)(=)(\s*)([a-zA-Z]\w*)',
-             bygroups(Name.Builtin, Whitespace, Punctuation, Whitespace,
-                      Keyword)),
-            (r',', Punctuation),
-            (r'\)', Punctuation, '#pop'),
-            (r'\s+', Whitespace),
-            (r'.', Text),
-        ],
-        'defprops': [
-            (r'%\{\s*\n', Comment.Multiline, 'blockcomment'),
-            (r'%.*$', Comment),
-            (r'(?.
-    """
-    name = 'Matlab session'
-    aliases = ['matlabsession']
-    url = 'https://www.mathworks.com/products/matlab.html'
-    version_added = '0.10'
-    _example = "matlabsession/matlabsession_sample.txt"
-
-    def get_tokens_unprocessed(self, text):
-        mlexer = MatlabLexer(**self.options)
-
-        curcode = ''
-        insertions = []
-        continuation = False
-
-        for match in line_re.finditer(text):
-            line = match.group()
-
-            if line.startswith('>> '):
-                insertions.append((len(curcode),
-                                   [(0, Generic.Prompt, line[:3])]))
-                curcode += line[3:]
-
-            elif line.startswith('>>'):
-                insertions.append((len(curcode),
-                                   [(0, Generic.Prompt, line[:2])]))
-                curcode += line[2:]
-
-            elif line.startswith('???'):
-
-                idx = len(curcode)
-
-                # without is showing error on same line as before...?
-                # line = "\n" + line
-                token = (0, Generic.Traceback, line)
-                insertions.append((idx, [token]))
-            elif continuation and insertions:
-                # line_start is the length of the most recent prompt symbol
-                line_start = len(insertions[-1][-1][-1])
-                # Set leading spaces with the length of the prompt to be a generic prompt
-                # This keeps code aligned when prompts are removed, say with some Javascript
-                if line.startswith(' '*line_start):
-                    insertions.append(
-                        (len(curcode), [(0, Generic.Prompt, line[:line_start])]))
-                    curcode += line[line_start:]
-                else:
-                    curcode += line
-            else:
-                if curcode:
-                    yield from do_insertions(
-                        insertions, mlexer.get_tokens_unprocessed(curcode))
-                    curcode = ''
-                    insertions = []
-
-                yield match.start(), Generic.Output, line
-
-            # Does not allow continuation if a comment is included after the ellipses.
-            # Continues any line that ends with ..., even comments (lines that start with %)
-            if line.strip().endswith('...'):
-                continuation = True
-            else:
-                continuation = False
-
-        if curcode:  # or item:
-            yield from do_insertions(
-                insertions, mlexer.get_tokens_unprocessed(curcode))
-
-
-class OctaveLexer(RegexLexer):
-    """
-    For GNU Octave source code.
-    """
-    name = 'Octave'
-    url = 'https://www.gnu.org/software/octave/index'
-    aliases = ['octave']
-    filenames = ['*.m']
-    mimetypes = ['text/octave']
-    version_added = '1.5'
-
-    # These lists are generated automatically.
-    # Run the following in bash shell:
-    #
-    # First dump all of the Octave manual into a plain text file:
-    #
-    #   $ info octave --subnodes -o octave-manual
-    #
-    # Now grep through it:
-
-    # for i in \
-    #     "Built-in Function" "Command" "Function File" \
-    #     "Loadable Function" "Mapping Function";
-    # do
-    #     perl -e '@name = qw('"$i"');
-    #              print lc($name[0]),"_kw = [\n"';
-    #
-    #     perl -n -e 'print "\"$1\",\n" if /-- '"$i"': .* (\w*) \(/;' \
-    #         octave-manual | sort | uniq ;
-    #     echo "]" ;
-    #     echo;
-    # done
-
-    # taken from Octave Mercurial changeset 8cc154f45e37 (30-jan-2011)
-
-    builtin_kw = (
-        "addlistener", "addpath", "addproperty", "all",
-        "and", "any", "argnames", "argv", "assignin",
-        "atexit", "autoload",
-        "available_graphics_toolkits", "beep_on_error",
-        "bitand", "bitmax", "bitor", "bitshift", "bitxor",
-        "cat", "cell", "cellstr", "char", "class", "clc",
-        "columns", "command_line_path",
-        "completion_append_char", "completion_matches",
-        "complex", "confirm_recursive_rmdir", "cputime",
-        "crash_dumps_octave_core", "ctranspose", "cumprod",
-        "cumsum", "debug_on_error", "debug_on_interrupt",
-        "debug_on_warning", "default_save_options",
-        "dellistener", "diag", "diff", "disp",
-        "doc_cache_file", "do_string_escapes", "double",
-        "drawnow", "e", "echo_executing_commands", "eps",
-        "eq", "errno", "errno_list", "error", "eval",
-        "evalin", "exec", "exist", "exit", "eye", "false",
-        "fclear", "fclose", "fcntl", "fdisp", "feof",
-        "ferror", "feval", "fflush", "fgetl", "fgets",
-        "fieldnames", "file_in_loadpath", "file_in_path",
-        "filemarker", "filesep", "find_dir_in_path",
-        "fixed_point_format", "fnmatch", "fopen", "fork",
-        "formula", "fprintf", "fputs", "fread", "freport",
-        "frewind", "fscanf", "fseek", "fskipl", "ftell",
-        "functions", "fwrite", "ge", "genpath", "get",
-        "getegid", "getenv", "geteuid", "getgid",
-        "getpgrp", "getpid", "getppid", "getuid", "glob",
-        "gt", "gui_mode", "history_control",
-        "history_file", "history_size",
-        "history_timestamp_format_string", "home",
-        "horzcat", "hypot", "ifelse",
-        "ignore_function_time_stamp", "inferiorto",
-        "info_file", "info_program", "inline", "input",
-        "intmax", "intmin", "ipermute",
-        "is_absolute_filename", "isargout", "isbool",
-        "iscell", "iscellstr", "ischar", "iscomplex",
-        "isempty", "isfield", "isfloat", "isglobal",
-        "ishandle", "isieee", "isindex", "isinteger",
-        "islogical", "ismatrix", "ismethod", "isnull",
-        "isnumeric", "isobject", "isreal",
-        "is_rooted_relative_filename", "issorted",
-        "isstruct", "isvarname", "kbhit", "keyboard",
-        "kill", "lasterr", "lasterror", "lastwarn",
-        "ldivide", "le", "length", "link", "linspace",
-        "logical", "lstat", "lt", "make_absolute_filename",
-        "makeinfo_program", "max_recursion_depth", "merge",
-        "methods", "mfilename", "minus", "mislocked",
-        "mkdir", "mkfifo", "mkstemp", "mldivide", "mlock",
-        "mouse_wheel_zoom", "mpower", "mrdivide", "mtimes",
-        "munlock", "nargin", "nargout",
-        "native_float_format", "ndims", "ne", "nfields",
-        "nnz", "norm", "not", "numel", "nzmax",
-        "octave_config_info", "octave_core_file_limit",
-        "octave_core_file_name",
-        "octave_core_file_options", "ones", "or",
-        "output_max_field_width", "output_precision",
-        "page_output_immediately", "page_screen_output",
-        "path", "pathsep", "pause", "pclose", "permute",
-        "pi", "pipe", "plus", "popen", "power",
-        "print_empty_dimensions", "printf",
-        "print_struct_array_contents", "prod",
-        "program_invocation_name", "program_name",
-        "putenv", "puts", "pwd", "quit", "rats", "rdivide",
-        "readdir", "readlink", "read_readline_init_file",
-        "realmax", "realmin", "rehash", "rename",
-        "repelems", "re_read_readline_init_file", "reset",
-        "reshape", "resize", "restoredefaultpath",
-        "rethrow", "rmdir", "rmfield", "rmpath", "rows",
-        "save_header_format_string", "save_precision",
-        "saving_history", "scanf", "set", "setenv",
-        "shell_cmd", "sighup_dumps_octave_core",
-        "sigterm_dumps_octave_core", "silent_functions",
-        "single", "size", "size_equal", "sizemax",
-        "sizeof", "sleep", "source", "sparse_auto_mutate",
-        "split_long_rows", "sprintf", "squeeze", "sscanf",
-        "stat", "stderr", "stdin", "stdout", "strcmp",
-        "strcmpi", "string_fill_char", "strncmp",
-        "strncmpi", "struct", "struct_levels_to_print",
-        "strvcat", "subsasgn", "subsref", "sum", "sumsq",
-        "superiorto", "suppress_verbose_help_message",
-        "symlink", "system", "tic", "tilde_expand",
-        "times", "tmpfile", "tmpnam", "toc", "toupper",
-        "transpose", "true", "typeinfo", "umask", "uminus",
-        "uname", "undo_string_escapes", "unlink", "uplus",
-        "upper", "usage", "usleep", "vec", "vectorize",
-        "vertcat", "waitpid", "warning", "warranty",
-        "whos_line_format", "yes_or_no", "zeros",
-        "inf", "Inf", "nan", "NaN")
-
-    command_kw = ("close", "load", "who", "whos")
-
-    function_kw = (
-        "accumarray", "accumdim", "acosd", "acotd",
-        "acscd", "addtodate", "allchild", "ancestor",
-        "anova", "arch_fit", "arch_rnd", "arch_test",
-        "area", "arma_rnd", "arrayfun", "ascii", "asctime",
-        "asecd", "asind", "assert", "atand",
-        "autoreg_matrix", "autumn", "axes", "axis", "bar",
-        "barh", "bartlett", "bartlett_test", "beep",
-        "betacdf", "betainv", "betapdf", "betarnd",
-        "bicgstab", "bicubic", "binary", "binocdf",
-        "binoinv", "binopdf", "binornd", "bitcmp",
-        "bitget", "bitset", "blackman", "blanks",
-        "blkdiag", "bone", "box", "brighten", "calendar",
-        "cast", "cauchy_cdf", "cauchy_inv", "cauchy_pdf",
-        "cauchy_rnd", "caxis", "celldisp", "center", "cgs",
-        "chisquare_test_homogeneity",
-        "chisquare_test_independence", "circshift", "cla",
-        "clabel", "clf", "clock", "cloglog", "closereq",
-        "colon", "colorbar", "colormap", "colperm",
-        "comet", "common_size", "commutation_matrix",
-        "compan", "compare_versions", "compass",
-        "computer", "cond", "condest", "contour",
-        "contourc", "contourf", "contrast", "conv",
-        "convhull", "cool", "copper", "copyfile", "cor",
-        "corrcoef", "cor_test", "cosd", "cotd", "cov",
-        "cplxpair", "cross", "cscd", "cstrcat", "csvread",
-        "csvwrite", "ctime", "cumtrapz", "curl", "cut",
-        "cylinder", "date", "datenum", "datestr",
-        "datetick", "datevec", "dblquad", "deal",
-        "deblank", "deconv", "delaunay", "delaunayn",
-        "delete", "demo", "detrend", "diffpara", "diffuse",
-        "dir", "discrete_cdf", "discrete_inv",
-        "discrete_pdf", "discrete_rnd", "display",
-        "divergence", "dlmwrite", "dos", "dsearch",
-        "dsearchn", "duplication_matrix", "durbinlevinson",
-        "ellipsoid", "empirical_cdf", "empirical_inv",
-        "empirical_pdf", "empirical_rnd", "eomday",
-        "errorbar", "etime", "etreeplot", "example",
-        "expcdf", "expinv", "expm", "exppdf", "exprnd",
-        "ezcontour", "ezcontourf", "ezmesh", "ezmeshc",
-        "ezplot", "ezpolar", "ezsurf", "ezsurfc", "factor",
-        "factorial", "fail", "fcdf", "feather", "fftconv",
-        "fftfilt", "fftshift", "figure", "fileattrib",
-        "fileparts", "fill", "findall", "findobj",
-        "findstr", "finv", "flag", "flipdim", "fliplr",
-        "flipud", "fpdf", "fplot", "fractdiff", "freqz",
-        "freqz_plot", "frnd", "fsolve",
-        "f_test_regression", "ftp", "fullfile", "fzero",
-        "gamcdf", "gaminv", "gampdf", "gamrnd", "gca",
-        "gcbf", "gcbo", "gcf", "genvarname", "geocdf",
-        "geoinv", "geopdf", "geornd", "getfield", "ginput",
-        "glpk", "gls", "gplot", "gradient",
-        "graphics_toolkit", "gray", "grid", "griddata",
-        "griddatan", "gtext", "gunzip", "gzip", "hadamard",
-        "hamming", "hankel", "hanning", "hggroup",
-        "hidden", "hilb", "hist", "histc", "hold", "hot",
-        "hotelling_test", "housh", "hsv", "hurst",
-        "hygecdf", "hygeinv", "hygepdf", "hygernd",
-        "idivide", "ifftshift", "image", "imagesc",
-        "imfinfo", "imread", "imshow", "imwrite", "index",
-        "info", "inpolygon", "inputname", "interpft",
-        "interpn", "intersect", "invhilb", "iqr", "isa",
-        "isdefinite", "isdir", "is_duplicate_entry",
-        "isequal", "isequalwithequalnans", "isfigure",
-        "ishermitian", "ishghandle", "is_leap_year",
-        "isletter", "ismac", "ismember", "ispc", "isprime",
-        "isprop", "isscalar", "issquare", "isstrprop",
-        "issymmetric", "isunix", "is_valid_file_id",
-        "isvector", "jet", "kendall",
-        "kolmogorov_smirnov_cdf",
-        "kolmogorov_smirnov_test", "kruskal_wallis_test",
-        "krylov", "kurtosis", "laplace_cdf", "laplace_inv",
-        "laplace_pdf", "laplace_rnd", "legend", "legendre",
-        "license", "line", "linkprop", "list_primes",
-        "loadaudio", "loadobj", "logistic_cdf",
-        "logistic_inv", "logistic_pdf", "logistic_rnd",
-        "logit", "loglog", "loglogerr", "logm", "logncdf",
-        "logninv", "lognpdf", "lognrnd", "logspace",
-        "lookfor", "ls_command", "lsqnonneg", "magic",
-        "mahalanobis", "manova", "matlabroot",
-        "mcnemar_test", "mean", "meansq", "median", "menu",
-        "mesh", "meshc", "meshgrid", "meshz", "mexext",
-        "mget", "mkpp", "mode", "moment", "movefile",
-        "mpoles", "mput", "namelengthmax", "nargchk",
-        "nargoutchk", "nbincdf", "nbininv", "nbinpdf",
-        "nbinrnd", "nchoosek", "ndgrid", "newplot", "news",
-        "nonzeros", "normcdf", "normest", "norminv",
-        "normpdf", "normrnd", "now", "nthroot", "null",
-        "ocean", "ols", "onenormest", "optimget",
-        "optimset", "orderfields", "orient", "orth",
-        "pack", "pareto", "parseparams", "pascal", "patch",
-        "pathdef", "pcg", "pchip", "pcolor", "pcr",
-        "peaks", "periodogram", "perl", "perms", "pie",
-        "pink", "planerot", "playaudio", "plot",
-        "plotmatrix", "plotyy", "poisscdf", "poissinv",
-        "poisspdf", "poissrnd", "polar", "poly",
-        "polyaffine", "polyarea", "polyderiv", "polyfit",
-        "polygcd", "polyint", "polyout", "polyreduce",
-        "polyval", "polyvalm", "postpad", "powerset",
-        "ppder", "ppint", "ppjumps", "ppplot", "ppval",
-        "pqpnonneg", "prepad", "primes", "print",
-        "print_usage", "prism", "probit", "qp", "qqplot",
-        "quadcc", "quadgk", "quadl", "quadv", "quiver",
-        "qzhess", "rainbow", "randi", "range", "rank",
-        "ranks", "rat", "reallog", "realpow", "realsqrt",
-        "record", "rectangle_lw", "rectangle_sw",
-        "rectint", "refresh", "refreshdata",
-        "regexptranslate", "repmat", "residue", "ribbon",
-        "rindex", "roots", "rose", "rosser", "rotdim",
-        "rref", "run", "run_count", "rundemos", "run_test",
-        "runtests", "saveas", "saveaudio", "saveobj",
-        "savepath", "scatter", "secd", "semilogx",
-        "semilogxerr", "semilogy", "semilogyerr",
-        "setaudio", "setdiff", "setfield", "setxor",
-        "shading", "shift", "shiftdim", "sign_test",
-        "sinc", "sind", "sinetone", "sinewave", "skewness",
-        "slice", "sombrero", "sortrows", "spaugment",
-        "spconvert", "spdiags", "spearman", "spectral_adf",
-        "spectral_xdf", "specular", "speed", "spencer",
-        "speye", "spfun", "sphere", "spinmap", "spline",
-        "spones", "sprand", "sprandn", "sprandsym",
-        "spring", "spstats", "spy", "sqp", "stairs",
-        "statistics", "std", "stdnormal_cdf",
-        "stdnormal_inv", "stdnormal_pdf", "stdnormal_rnd",
-        "stem", "stft", "strcat", "strchr", "strjust",
-        "strmatch", "strread", "strsplit", "strtok",
-        "strtrim", "strtrunc", "structfun", "studentize",
-        "subplot", "subsindex", "subspace", "substr",
-        "substruct", "summer", "surf", "surface", "surfc",
-        "surfl", "surfnorm", "svds", "swapbytes",
-        "sylvester_matrix", "symvar", "synthesis", "table",
-        "tand", "tar", "tcdf", "tempdir", "tempname",
-        "test", "text", "textread", "textscan", "tinv",
-        "title", "toeplitz", "tpdf", "trace", "trapz",
-        "treelayout", "treeplot", "triangle_lw",
-        "triangle_sw", "tril", "trimesh", "triplequad",
-        "triplot", "trisurf", "triu", "trnd", "tsearchn",
-        "t_test", "t_test_regression", "type", "unidcdf",
-        "unidinv", "unidpdf", "unidrnd", "unifcdf",
-        "unifinv", "unifpdf", "unifrnd", "union", "unique",
-        "unix", "unmkpp", "unpack", "untabify", "untar",
-        "unwrap", "unzip", "u_test", "validatestring",
-        "vander", "var", "var_test", "vech", "ver",
-        "version", "view", "voronoi", "voronoin",
-        "waitforbuttonpress", "wavread", "wavwrite",
-        "wblcdf", "wblinv", "wblpdf", "wblrnd", "weekday",
-        "welch_test", "what", "white", "whitebg",
-        "wienrnd", "wilcoxon_test", "wilkinson", "winter",
-        "xlabel", "xlim", "ylabel", "yulewalker", "zip",
-        "zlabel", "z_test")
-
-    loadable_kw = (
-        "airy", "amd", "balance", "besselh", "besseli",
-        "besselj", "besselk", "bessely", "bitpack",
-        "bsxfun", "builtin", "ccolamd", "cellfun",
-        "cellslices", "chol", "choldelete", "cholinsert",
-        "cholinv", "cholshift", "cholupdate", "colamd",
-        "colloc", "convhulln", "convn", "csymamd",
-        "cummax", "cummin", "daspk", "daspk_options",
-        "dasrt", "dasrt_options", "dassl", "dassl_options",
-        "dbclear", "dbdown", "dbstack", "dbstatus",
-        "dbstop", "dbtype", "dbup", "dbwhere", "det",
-        "dlmread", "dmperm", "dot", "eig", "eigs",
-        "endgrent", "endpwent", "etree", "fft", "fftn",
-        "fftw", "filter", "find", "full", "gcd",
-        "getgrent", "getgrgid", "getgrnam", "getpwent",
-        "getpwnam", "getpwuid", "getrusage", "givens",
-        "gmtime", "gnuplot_binary", "hess", "ifft",
-        "ifftn", "inv", "isdebugmode", "issparse", "kron",
-        "localtime", "lookup", "lsode", "lsode_options",
-        "lu", "luinc", "luupdate", "matrix_type", "max",
-        "min", "mktime", "pinv", "qr", "qrdelete",
-        "qrinsert", "qrshift", "qrupdate", "quad",
-        "quad_options", "qz", "rand", "rande", "randg",
-        "randn", "randp", "randperm", "rcond", "regexp",
-        "regexpi", "regexprep", "schur", "setgrent",
-        "setpwent", "sort", "spalloc", "sparse", "spparms",
-        "sprank", "sqrtm", "strfind", "strftime",
-        "strptime", "strrep", "svd", "svd_driver", "syl",
-        "symamd", "symbfact", "symrcm", "time", "tsearch",
-        "typecast", "urlread", "urlwrite")
-
-    mapping_kw = (
-        "abs", "acos", "acosh", "acot", "acoth", "acsc",
-        "acsch", "angle", "arg", "asec", "asech", "asin",
-        "asinh", "atan", "atanh", "beta", "betainc",
-        "betaln", "bincoeff", "cbrt", "ceil", "conj", "cos",
-        "cosh", "cot", "coth", "csc", "csch", "erf", "erfc",
-        "erfcx", "erfinv", "exp", "finite", "fix", "floor",
-        "fmod", "gamma", "gammainc", "gammaln", "imag",
-        "isalnum", "isalpha", "isascii", "iscntrl",
-        "isdigit", "isfinite", "isgraph", "isinf",
-        "islower", "isna", "isnan", "isprint", "ispunct",
-        "isspace", "isupper", "isxdigit", "lcm", "lgamma",
-        "log", "lower", "mod", "real", "rem", "round",
-        "roundb", "sec", "sech", "sign", "sin", "sinh",
-        "sqrt", "tan", "tanh", "toascii", "tolower", "xor")
-
-    builtin_consts = (
-        "EDITOR", "EXEC_PATH", "I", "IMAGE_PATH", "NA",
-        "OCTAVE_HOME", "OCTAVE_VERSION", "PAGER",
-        "PAGER_FLAGS", "SEEK_CUR", "SEEK_END", "SEEK_SET",
-        "SIG", "S_ISBLK", "S_ISCHR", "S_ISDIR", "S_ISFIFO",
-        "S_ISLNK", "S_ISREG", "S_ISSOCK", "WCONTINUE",
-        "WCOREDUMP", "WEXITSTATUS", "WIFCONTINUED",
-        "WIFEXITED", "WIFSIGNALED", "WIFSTOPPED", "WNOHANG",
-        "WSTOPSIG", "WTERMSIG", "WUNTRACED")
-
-    tokens = {
-        'root': [
-            (r'%\{\s*\n', Comment.Multiline, 'percentblockcomment'),
-            (r'#\{\s*\n', Comment.Multiline, 'hashblockcomment'),
-            (r'[%#].*$', Comment),
-            (r'^\s*function\b', Keyword, 'deffunc'),
-
-            # from 'iskeyword' on hg changeset 8cc154f45e37
-            (words((
-                '__FILE__', '__LINE__', 'break', 'case', 'catch', 'classdef',
-                'continue', 'do', 'else', 'elseif', 'end', 'end_try_catch',
-                'end_unwind_protect', 'endclassdef', 'endevents', 'endfor',
-                'endfunction', 'endif', 'endmethods', 'endproperties', 'endswitch',
-                'endwhile', 'events', 'for', 'function', 'get', 'global', 'if',
-                'methods', 'otherwise', 'persistent', 'properties', 'return',
-                'set', 'static', 'switch', 'try', 'until', 'unwind_protect',
-                'unwind_protect_cleanup', 'while'), suffix=r'\b'),
-             Keyword),
-
-            (words(builtin_kw + command_kw + function_kw + loadable_kw + mapping_kw,
-                   suffix=r'\b'),  Name.Builtin),
-
-            (words(builtin_consts, suffix=r'\b'), Name.Constant),
-
-            # operators in Octave but not Matlab:
-            (r'-=|!=|!|/=|--', Operator),
-            # operators:
-            (r'-|==|~=|<|>|<=|>=|&&|&|~|\|\|?', Operator),
-            # operators in Octave but not Matlab requiring escape for re:
-            (r'\*=|\+=|\^=|\/=|\\=|\*\*|\+\+|\.\*\*', Operator),
-            # operators requiring escape for re:
-            (r'\.\*|\*|\+|\.\^|\^|\.\\|\.\/|\/|\\', Operator),
-
-
-            # punctuation:
-            (r'[\[\](){}:@.,]', Punctuation),
-            (r'=|:|;', Punctuation),
-
-            (r'"[^"]*"', String),
-
-            (r'(\d+\.\d*|\d*\.\d+)([eEf][+-]?[0-9]+)?', Number.Float),
-            (r'\d+[eEf][+-]?[0-9]+', Number.Float),
-            (r'\d+', Number.Integer),
-
-            # quote can be transpose, instead of string:
-            # (not great, but handles common cases...)
-            (r'(?<=[\w)\].])\'+', Operator),
-            (r'(?|<=|>=|&&|&|~|\|\|?', Operator),
-            # operators requiring escape for re:
-            (r'\.\*|\*|\+|\.\^|\^|\.\\|\.\/|\/|\\', Operator),
-
-            # punctuation:
-            (r'[\[\](){}@.,=:;]+', Punctuation),
-
-            (r'"[^"]*"', String),
-
-            # quote can be transpose, instead of string:
-            # (not great, but handles common cases...)
-            (r'(?<=[\w)\].])\'+', Operator),
-            (r'(?', r'<', r'|', r'!', r"'")
-
-    operator_words = ('and', 'or', 'not')
-
-    tokens = {
-        'root': [
-            (r'/\*', Comment.Multiline, 'comment'),
-            (r'"(?:[^"\\]|\\.)*"', String),
-            (r'\(|\)|\[|\]|\{|\}', Punctuation),
-            (r'[,;$]', Punctuation),
-            (words (constants), Name.Constant),
-            (words (keywords), Keyword),
-            (words (operators), Operator),
-            (words (operator_words), Operator.Word),
-            (r'''(?x)
-              ((?:[a-zA-Z_#][\w#]*|`[^`]*`)
-              (?:::[a-zA-Z_#][\w#]*|`[^`]*`)*)(\s*)([(])''',
-             bygroups(Name.Function, Text.Whitespace, Punctuation)),
-            (r'''(?x)
-              (?:[a-zA-Z_#%][\w#%]*|`[^`]*`)
-              (?:::[a-zA-Z_#%][\w#%]*|`[^`]*`)*''', Name.Variable),
-            (r'[-+]?(\d*\.\d+([bdefls][-+]?\d+)?|\d+(\.\d*)?[bdefls][-+]?\d+)', Number.Float),
-            (r'[-+]?\d+', Number.Integer),
-            (r'\s+', Text.Whitespace),
-            (r'.', Text)
-        ],
-        'comment': [
-            (r'[^*/]+', Comment.Multiline),
-            (r'/\*', Comment.Multiline, '#push'),
-            (r'\*/', Comment.Multiline, '#pop'),
-            (r'[*/]', Comment.Multiline)
-        ]
-    }
-
-    def analyse_text (text):
-        strength = 0.0
-        # Input expression terminator.
-        if re.search (r'\$\s*$', text, re.MULTILINE):
-            strength += 0.05
-        # Function definition operator.
-        if ':=' in text:
-            strength += 0.02
-        return strength
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/meson.py b/venv/lib/python3.10/site-packages/pygments/lexers/meson.py
deleted file mode 100644
index 6f2c6da374b0e83fc3adc2c577c77884b7d09923..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/meson.py
+++ /dev/null
@@ -1,139 +0,0 @@
-"""
-    pygments.lexers.meson
-    ~~~~~~~~~~~~~~~~~~~~~
-
-    Pygments lexer for the Meson build system
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, words, include
-from pygments.token import Comment, Name, Number, Punctuation, Operator, \
-    Keyword, String, Whitespace
-
-__all__ = ['MesonLexer']
-
-
-class MesonLexer(RegexLexer):
-    """Meson language lexer.
-
-    The grammar definition use to transcribe the syntax was retrieved from
-    https://mesonbuild.com/Syntax.html#grammar for version 0.58.
-    Some of those definitions are improperly transcribed, so the Meson++
-    implementation was also checked: https://github.com/dcbaker/meson-plus-plus.
-    """
-
-    # TODO String interpolation @VARNAME@ inner matches
-    # TODO keyword_arg: value inner matches
-
-    name = 'Meson'
-    url = 'https://mesonbuild.com/'
-    aliases = ['meson', 'meson.build']
-    filenames = ['meson.build', 'meson_options.txt']
-    mimetypes = ['text/x-meson']
-    version_added = '2.10'
-
-    tokens = {
-        'root': [
-            (r'#.*?$', Comment),
-            (r"'''.*'''", String.Single),
-            (r'[1-9][0-9]*', Number.Integer),
-            (r'0o[0-7]+', Number.Oct),
-            (r'0x[a-fA-F0-9]+', Number.Hex),
-            include('string'),
-            include('keywords'),
-            include('expr'),
-            (r'[a-zA-Z_][a-zA-Z_0-9]*', Name),
-            (r'\s+', Whitespace),
-        ],
-        'string': [
-            (r"[']{3}([']{0,2}([^\\']|\\(.|\n)))*[']{3}", String),
-            (r"'.*?(?`_.
-    """
-
-    name = "MCFunction"
-    url = "https://minecraft.wiki/w/Commands"
-    aliases = ["mcfunction", "mcf"]
-    filenames = ["*.mcfunction"]
-    mimetypes = ["text/mcfunction"]
-    version_added = '2.12'
-
-    # Used to denotate the start of a block comment, borrowed from Github's mcfunction
-    _block_comment_prefix = "[>!]"
-
-    tokens = {
-        "root": [
-            include("names"),
-            include("comments"),
-            include("literals"),
-            include("whitespace"),
-            include("property"),
-            include("operators"),
-            include("selectors"),
-        ],
-
-        "names": [
-            # The start of a command (either beginning of line OR after the run keyword)
-            #  We don't encode a list of keywords since mods, plugins, or even pre-processors
-            #  may add new commands, so we have a 'close-enough' regex which catches them.
-            (r"^(\s*)([a-z_]+)", bygroups(Whitespace, Name.Builtin)),
-            (r"(?<=run)\s+[a-z_]+", Name.Builtin),
-
-            # UUID
-            (r"\b[0-9a-fA-F]+(?:-[0-9a-fA-F]+){4}\b", Name.Variable),
-            include("resource-name"),
-            # normal command names and scoreboards
-            #  there's no way to know the differences unfortuntely
-            (r"[A-Za-z_][\w.#%$]+", Keyword.Constant),
-            (r"[#%$][\w.#%$]+", Name.Variable.Magic),
-        ],
-
-        "resource-name": [
-            # resource names have to be lowercase
-            (r"#?[a-z_][a-z_.-]*:[a-z0-9_./-]+", Name.Function),
-            # similar to above except optional `:``
-            #  a `/` must be present "somewhere"
-            (r"#?[a-z0-9_\.\-]+\/[a-z0-9_\.\-\/]+", Name.Function),
-        ],
-
-        "whitespace": [
-            (r"\s+", Whitespace),
-        ],
-
-        "comments": [
-            (rf"^\s*(#{_block_comment_prefix})", Comment.Multiline,
-             ("comments.block", "comments.block.emphasized")),
-            (r"#.*$", Comment.Single),
-        ],
-        "comments.block": [
-            (rf"^\s*#{_block_comment_prefix}", Comment.Multiline,
-             "comments.block.emphasized"),
-            (r"^\s*#", Comment.Multiline, "comments.block.normal"),
-            default("#pop"),
-        ],
-        "comments.block.normal": [
-            include("comments.block.special"),
-            (r"\S+", Comment.Multiline),
-            (r"\n", Text, "#pop"),
-            include("whitespace"),
-        ],
-        "comments.block.emphasized": [
-            include("comments.block.special"),
-            (r"\S+", String.Doc),
-            (r"\n", Text, "#pop"),
-            include("whitespace"),
-        ],
-        "comments.block.special": [
-            # Params
-            (r"@\S+", Name.Decorator),
-
-            include("resource-name"),
-
-            # Scoreboard player names
-            (r"[#%$][\w.#%$]+", Name.Variable.Magic),
-        ],
-
-        "operators": [
-            (r"[\-~%^?!+*<>\\/|&=.]", Operator),
-        ],
-
-        "literals": [
-            (r"\.\.", Literal),
-            (r"(true|false)", Keyword.Pseudo),
-
-            # these are like unquoted strings and appear in many places
-            (r"[A-Za-z_]+", Name.Variable.Class),
-
-            (r"[0-7]b", Number.Byte),
-            (r"[+-]?\d*\.?\d+([eE]?[+-]?\d+)?[df]?\b", Number.Float),
-            (r"[+-]?\d+\b", Number.Integer),
-            (r'"', String.Double, "literals.string-double"),
-            (r"'", String.Single, "literals.string-single"),
-        ],
-        "literals.string-double": [
-            (r"\\.", String.Escape),
-            (r'[^\\"\n]+', String.Double),
-            (r'"', String.Double, "#pop"),
-        ],
-        "literals.string-single": [
-            (r"\\.", String.Escape),
-            (r"[^\\'\n]+", String.Single),
-            (r"'", String.Single, "#pop"),
-        ],
-
-        "selectors": [
-            (r"@[a-z]", Name.Variable),
-        ],
-
-
-        ## Generic Property Container
-        # There are several, differing instances where the language accepts
-        #  specific contained keys or contained key, value pairings.
-        #
-        # Property Maps:
-        # - Starts with either `[` or `{`
-        # - Key separated by `:` or `=`
-        # - Deliminated by `,`
-        #
-        # Property Lists:
-        # - Starts with `[`
-        # - Deliminated by `,`
-        #
-        # For simplicity, these patterns match a generic, nestable structure
-        #  which follow a key, value pattern. For normal lists, there's only keys.
-        # This allow some "illegal" structures, but we'll accept those for
-        #  sake of simplicity
-        #
-        # Examples:
-        # - `[facing=up, powered=true]` (blockstate)
-        # - `[name="hello world", nbt={key: 1b}]` (selector + nbt)
-        # - `[{"text": "value"}, "literal"]` (json)
-        ##
-        "property": [
-            # This state gets included in root and also several substates
-            # We do this to shortcut the starting of new properties
-            #  within other properties. Lists can have sublists and compounds
-            #  and values can start a new property (see the `difficult_1.txt`
-            #  snippet).
-            (r"\{", Punctuation, ("property.curly", "property.key")),
-            (r"\[", Punctuation, ("property.square", "property.key")),
-        ],
-        "property.curly": [
-            include("whitespace"),
-            include("property"),
-            (r"\}", Punctuation, "#pop"),
-        ],
-        "property.square": [
-            include("whitespace"),
-            include("property"),
-            (r"\]", Punctuation, "#pop"),
-
-            # lists can have sequences of items
-            (r",", Punctuation),
-        ],
-        "property.key": [
-            include("whitespace"),
-
-            # resource names (for advancements)
-            #  can omit `:` to default `minecraft:`
-            # must check if there is a future equals sign if `:` is in the name
-            (r"#?[a-z_][a-z_\.\-]*\:[a-z0-9_\.\-/]+(?=\s*\=)", Name.Attribute, "property.delimiter"),
-            (r"#?[a-z_][a-z0-9_\.\-/]+", Name.Attribute, "property.delimiter"),
-
-            # unquoted NBT key
-            (r"[A-Za-z_\-\+]+", Name.Attribute, "property.delimiter"),
-
-            # quoted JSON or NBT key
-            (r'"', Name.Attribute, "property.delimiter", "literals.string-double"),
-            (r"'", Name.Attribute, "property.delimiter", "literals.string-single"),
-
-            # index for a list
-            (r"-?\d+", Number.Integer, "property.delimiter"),
-
-            default("#pop"),
-        ],
-        "property.key.string-double": [
-            (r"\\.", String.Escape),
-            (r'[^\\"\n]+', Name.Attribute),
-            (r'"', Name.Attribute, "#pop"),
-        ],
-        "property.key.string-single": [
-            (r"\\.", String.Escape),
-            (r"[^\\'\n]+", Name.Attribute),
-            (r"'", Name.Attribute, "#pop"),
-        ],
-        "property.delimiter": [
-            include("whitespace"),
-
-            (r"[:=]!?", Punctuation, "property.value"),
-            (r",", Punctuation),
-
-            default("#pop"),
-        ],
-        "property.value": [
-            include("whitespace"),
-
-            # unquoted resource names are valid literals here
-            (r"#?[a-z_][a-z_\.\-]*\:[a-z0-9_\.\-/]+", Name.Tag),
-            (r"#?[a-z_][a-z0-9_\.\-/]+", Name.Tag),
-
-            include("literals"),
-            include("property"),
-
-            default("#pop"),
-        ],
-    }
-
-
-class MCSchemaLexer(RegexLexer):
-    """Lexer for Minecraft Add-ons data Schemas, an interface structure standard used in Minecraft
-    """
-
-    name = 'MCSchema'
-    url = 'https://learn.microsoft.com/en-us/minecraft/creator/reference/content/schemasreference/'
-    aliases = ['mcschema']
-    filenames = ['*.mcschema']
-    mimetypes = ['text/mcschema']
-    version_added = '2.14'
-
-    tokens = {
-        'commentsandwhitespace': [
-            (r'\s+', Whitespace),
-            (r'//.*?$', Comment.Single),
-            (r'/\*.*?\*/', Comment.Multiline)
-        ],
-        'slashstartsregex': [
-            include('commentsandwhitespace'),
-            (r'/(\\.|[^[/\\\n]|\[(\\.|[^\]\\\n])*])+/'
-             r'([gimuysd]+\b|\B)', String.Regex, '#pop'),
-            (r'(?=/)', Text, ('#pop', 'badregex')),
-            default('#pop')
-        ],
-        'badregex': [
-            (r'\n', Whitespace, '#pop')
-        ],
-        'singlestring': [
-            (r'\\.', String.Escape),
-            (r"'", String.Single, '#pop'),
-            (r"[^\\']+", String.Single),
-        ],
-        'doublestring': [
-            (r'\\.', String.Escape),
-            (r'"', String.Double, '#pop'),
-            (r'[^\\"]+', String.Double),
-        ],
-        'root': [
-            (r'^(?=\s|/|', Comment, '#pop'),
-            (r'[^\-]+|-', Comment),
-        ],
-    }
-
-
-class ReasonLexer(RegexLexer):
-    """
-    For the ReasonML language.
-    """
-
-    name = 'ReasonML'
-    url = 'https://reasonml.github.io/'
-    aliases = ['reasonml', 'reason']
-    filenames = ['*.re', '*.rei']
-    mimetypes = ['text/x-reasonml']
-    version_added = '2.6'
-
-    keywords = (
-        'as', 'assert', 'begin', 'class', 'constraint', 'do', 'done', 'downto',
-        'else', 'end', 'exception', 'external', 'false', 'for', 'fun', 'esfun',
-        'function', 'functor', 'if', 'in', 'include', 'inherit', 'initializer', 'lazy',
-        'let', 'switch', 'module', 'pub', 'mutable', 'new', 'nonrec', 'object', 'of',
-        'open', 'pri', 'rec', 'sig', 'struct', 'then', 'to', 'true', 'try',
-        'type', 'val', 'virtual', 'when', 'while', 'with',
-    )
-    keyopts = (
-        '!=', '#', '&', '&&', r'\(', r'\)', r'\*', r'\+', ',', '-',
-        r'-\.', '=>', r'\.', r'\.\.', r'\.\.\.', ':', '::', ':=', ':>', ';', ';;', '<',
-        '<-', '=', '>', '>]', r'>\}', r'\?', r'\?\?', r'\[', r'\[<', r'\[>',
-        r'\[\|', ']', '_', '`', r'\{', r'\{<', r'\|', r'\|\|', r'\|]', r'\}', '~'
-    )
-
-    operators = r'[!$%&*+\./:<=>?@^|~-]'
-    word_operators = ('and', 'asr', 'land', 'lor', 'lsl', 'lsr', 'lxor', 'mod', 'or')
-    prefix_syms = r'[!?~]'
-    infix_syms = r'[=<>@^|&+\*/$%-]'
-    primitives = ('unit', 'int', 'float', 'bool', 'string', 'char', 'list', 'array')
-
-    tokens = {
-        'escape-sequence': [
-            (r'\\[\\"\'ntbr]', String.Escape),
-            (r'\\[0-9]{3}', String.Escape),
-            (r'\\x[0-9a-fA-F]{2}', String.Escape),
-        ],
-        'root': [
-            (r'\s+', Text),
-            (r'false|true|\(\)|\[\]', Name.Builtin.Pseudo),
-            (r'\b([A-Z][\w\']*)(?=\s*\.)', Name.Namespace, 'dotted'),
-            (r'\b([A-Z][\w\']*)', Name.Class),
-            (r'//.*?\n', Comment.Single),
-            (r'\/\*(?!/)', Comment.Multiline, 'comment'),
-            (r'\b({})\b'.format('|'.join(keywords)), Keyword),
-            (r'({})'.format('|'.join(keyopts[::-1])), Operator.Word),
-            (rf'({infix_syms}|{prefix_syms})?{operators}', Operator),
-            (r'\b({})\b'.format('|'.join(word_operators)), Operator.Word),
-            (r'\b({})\b'.format('|'.join(primitives)), Keyword.Type),
-
-            (r"[^\W\d][\w']*", Name),
-
-            (r'-?\d[\d_]*(.[\d_]*)?([eE][+\-]?\d[\d_]*)', Number.Float),
-            (r'0[xX][\da-fA-F][\da-fA-F_]*', Number.Hex),
-            (r'0[oO][0-7][0-7_]*', Number.Oct),
-            (r'0[bB][01][01_]*', Number.Bin),
-            (r'\d[\d_]*', Number.Integer),
-
-            (r"'(?:(\\[\\\"'ntbr ])|(\\[0-9]{3})|(\\x[0-9a-fA-F]{2}))'",
-             String.Char),
-            (r"'.'", String.Char),
-            (r"'", Keyword),
-
-            (r'"', String.Double, 'string'),
-
-            (r'[~?][a-z][\w\']*:', Name.Variable),
-        ],
-        'comment': [
-            (r'[^/*]+', Comment.Multiline),
-            (r'\/\*', Comment.Multiline, '#push'),
-            (r'\*\/', Comment.Multiline, '#pop'),
-            (r'\*', Comment.Multiline),
-        ],
-        'string': [
-            (r'[^\\"]+', String.Double),
-            include('escape-sequence'),
-            (r'\\\n', String.Double),
-            (r'"', String.Double, '#pop'),
-        ],
-        'dotted': [
-            (r'\s+', Text),
-            (r'\.', Punctuation),
-            (r'[A-Z][\w\']*(?=\s*\.)', Name.Namespace),
-            (r'[A-Z][\w\']*', Name.Class, '#pop'),
-            (r'[a-z_][\w\']*', Name, '#pop'),
-            default('#pop'),
-        ],
-    }
-
-
-class FStarLexer(RegexLexer):
-    """
-    For the F* language.
-    """
-
-    name = 'FStar'
-    url = 'https://www.fstar-lang.org/'
-    aliases = ['fstar']
-    filenames = ['*.fst', '*.fsti']
-    mimetypes = ['text/x-fstar']
-    version_added = '2.7'
-
-    keywords = (
-        'abstract', 'attributes', 'noeq', 'unopteq', 'and'
-        'begin', 'by', 'default', 'effect', 'else', 'end', 'ensures',
-        'exception', 'exists', 'false', 'forall', 'fun', 'function', 'if',
-        'in', 'include', 'inline', 'inline_for_extraction', 'irreducible',
-        'logic', 'match', 'module', 'mutable', 'new', 'new_effect', 'noextract',
-        'of', 'open', 'opaque', 'private', 'range_of', 'reifiable',
-        'reify', 'reflectable', 'requires', 'set_range_of', 'sub_effect',
-        'synth', 'then', 'total', 'true', 'try', 'type', 'unfold', 'unfoldable',
-        'val', 'when', 'with', 'not'
-    )
-    decl_keywords = ('let', 'rec')
-    assume_keywords = ('assume', 'admit', 'assert', 'calc')
-    keyopts = (
-        r'~', r'-', r'/\\', r'\\/', r'<:', r'<@', r'\(\|', r'\|\)', r'#', r'u#',
-        r'&', r'\(', r'\)', r'\(\)', r',', r'~>', r'->', r'<-', r'<--', r'<==>',
-        r'==>', r'\.', r'\?', r'\?\.', r'\.\[', r'\.\(', r'\.\(\|', r'\.\[\|',
-        r'\{:pattern', r':', r'::', r':=', r';', r';;', r'=', r'%\[', r'!\{',
-        r'\[', r'\[@', r'\[\|', r'\|>', r'\]', r'\|\]', r'\{', r'\|', r'\}', r'\$'
-    )
-
-    operators = r'[!$%&*+\./:<=>?@^|~-]'
-    prefix_syms = r'[!?~]'
-    infix_syms = r'[=<>@^|&+\*/$%-]'
-    primitives = ('unit', 'int', 'float', 'bool', 'string', 'char', 'list', 'array')
-
-    tokens = {
-        'escape-sequence': [
-            (r'\\[\\"\'ntbr]', String.Escape),
-            (r'\\[0-9]{3}', String.Escape),
-            (r'\\x[0-9a-fA-F]{2}', String.Escape),
-        ],
-        'root': [
-            (r'\s+', Text),
-            (r'false|true|False|True|\(\)|\[\]', Name.Builtin.Pseudo),
-            (r'\b([A-Z][\w\']*)(?=\s*\.)', Name.Namespace, 'dotted'),
-            (r'\b([A-Z][\w\']*)', Name.Class),
-            (r'\(\*(?![)])', Comment, 'comment'),
-            (r'\/\/.+$', Comment),
-            (r'\b({})\b'.format('|'.join(keywords)), Keyword),
-            (r'\b({})\b'.format('|'.join(assume_keywords)), Name.Exception),
-            (r'\b({})\b'.format('|'.join(decl_keywords)), Keyword.Declaration),
-            (r'({})'.format('|'.join(keyopts[::-1])), Operator),
-            (rf'({infix_syms}|{prefix_syms})?{operators}', Operator),
-            (r'\b({})\b'.format('|'.join(primitives)), Keyword.Type),
-
-            (r"[^\W\d][\w']*", Name),
-
-            (r'-?\d[\d_]*(.[\d_]*)?([eE][+\-]?\d[\d_]*)', Number.Float),
-            (r'0[xX][\da-fA-F][\da-fA-F_]*', Number.Hex),
-            (r'0[oO][0-7][0-7_]*', Number.Oct),
-            (r'0[bB][01][01_]*', Number.Bin),
-            (r'\d[\d_]*', Number.Integer),
-
-            (r"'(?:(\\[\\\"'ntbr ])|(\\[0-9]{3})|(\\x[0-9a-fA-F]{2}))'",
-             String.Char),
-            (r"'.'", String.Char),
-            (r"'", Keyword),  # a stray quote is another syntax element
-            (r"\`([\w\'.]+)\`", Operator.Word),  # for infix applications
-            (r"\`", Keyword),  # for quoting
-            (r'"', String.Double, 'string'),
-
-            (r'[~?][a-z][\w\']*:', Name.Variable),
-        ],
-        'comment': [
-            (r'[^(*)]+', Comment),
-            (r'\(\*', Comment, '#push'),
-            (r'\*\)', Comment, '#pop'),
-            (r'[(*)]', Comment),
-        ],
-        'string': [
-            (r'[^\\"]+', String.Double),
-            include('escape-sequence'),
-            (r'\\\n', String.Double),
-            (r'"', String.Double, '#pop'),
-        ],
-        'dotted': [
-            (r'\s+', Text),
-            (r'\.', Punctuation),
-            (r'[A-Z][\w\']*(?=\s*\.)', Name.Namespace),
-            (r'[A-Z][\w\']*', Name.Class, '#pop'),
-            (r'[a-z_][\w\']*', Name, '#pop'),
-            default('#pop'),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/modeling.py b/venv/lib/python3.10/site-packages/pygments/lexers/modeling.py
deleted file mode 100644
index d681e7f3bb1479dc1ada0c56204d51d35df4ec8f..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/modeling.py
+++ /dev/null
@@ -1,366 +0,0 @@
-"""
-    pygments.lexers.modeling
-    ~~~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for modeling languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, include, bygroups, using, default
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Whitespace
-
-from pygments.lexers.html import HtmlLexer
-from pygments.lexers import _stan_builtins
-
-__all__ = ['ModelicaLexer', 'BugsLexer', 'JagsLexer', 'StanLexer']
-
-
-class ModelicaLexer(RegexLexer):
-    """
-    For Modelica source code.
-    """
-    name = 'Modelica'
-    url = 'http://www.modelica.org/'
-    aliases = ['modelica']
-    filenames = ['*.mo']
-    mimetypes = ['text/x-modelica']
-    version_added = '1.1'
-
-    flags = re.DOTALL | re.MULTILINE
-
-    _name = r"(?:'(?:[^\\']|\\.)+'|[a-zA-Z_]\w*)"
-
-    tokens = {
-        'whitespace': [
-            (r'[\s\ufeff]+', Text),
-            (r'//[^\n]*\n?', Comment.Single),
-            (r'/\*.*?\*/', Comment.Multiline)
-        ],
-        'root': [
-            include('whitespace'),
-            (r'"', String.Double, 'string'),
-            (r'[()\[\]{},;]+', Punctuation),
-            (r'\.?[*^/+-]|\.|<>|[<>:=]=?', Operator),
-            (r'\d+(\.?\d*[eE][-+]?\d+|\.\d*)', Number.Float),
-            (r'\d+', Number.Integer),
-            (r'(abs|acos|actualStream|array|asin|assert|AssertionLevel|atan|'
-             r'atan2|backSample|Boolean|cardinality|cat|ceil|change|Clock|'
-             r'Connections|cos|cosh|cross|delay|diagonal|div|edge|exp|'
-             r'ExternalObject|fill|floor|getInstanceName|hold|homotopy|'
-             r'identity|inStream|integer|Integer|interval|inverse|isPresent|'
-             r'linspace|log|log10|matrix|max|min|mod|ndims|noClock|noEvent|'
-             r'ones|outerProduct|pre|previous|product|Real|reinit|rem|rooted|'
-             r'sample|scalar|semiLinear|shiftSample|sign|sin|sinh|size|skew|'
-             r'smooth|spatialDistribution|sqrt|StateSelect|String|subSample|'
-             r'sum|superSample|symmetric|tan|tanh|terminal|terminate|time|'
-             r'transpose|vector|zeros)\b', Name.Builtin),
-            (r'(algorithm|annotation|break|connect|constant|constrainedby|der|'
-             r'discrete|each|else|elseif|elsewhen|encapsulated|enumeration|'
-             r'equation|exit|expandable|extends|external|firstTick|final|flow|for|if|'
-             r'import|impure|in|initial|inner|input|interval|loop|nondiscrete|outer|'
-             r'output|parameter|partial|protected|public|pure|redeclare|'
-             r'replaceable|return|stream|then|when|while)\b',
-             Keyword.Reserved),
-            (r'(and|not|or)\b', Operator.Word),
-            (r'(block|class|connector|end|function|model|operator|package|'
-             r'record|type)\b', Keyword.Reserved, 'class'),
-            (r'(false|true)\b', Keyword.Constant),
-            (r'within\b', Keyword.Reserved, 'package-prefix'),
-            (_name, Name)
-        ],
-        'class': [
-            include('whitespace'),
-            (r'(function|record)\b', Keyword.Reserved),
-            (r'(if|for|when|while)\b', Keyword.Reserved, '#pop'),
-            (_name, Name.Class, '#pop'),
-            default('#pop')
-        ],
-        'package-prefix': [
-            include('whitespace'),
-            (_name, Name.Namespace, '#pop'),
-            default('#pop')
-        ],
-        'string': [
-            (r'"', String.Double, '#pop'),
-            (r'\\[\'"?\\abfnrtv]', String.Escape),
-            (r'(?i)<\s*html\s*>([^\\"]|\\.)+?(<\s*/\s*html\s*>|(?="))',
-             using(HtmlLexer)),
-            (r'<|\\?[^"\\<]+', String.Double)
-        ]
-    }
-
-
-class BugsLexer(RegexLexer):
-    """
-    Pygments Lexer for OpenBugs and WinBugs
-    models.
-    """
-
-    name = 'BUGS'
-    aliases = ['bugs', 'winbugs', 'openbugs']
-    filenames = ['*.bug']
-    url = 'https://www.mrc-bsu.cam.ac.uk/software/bugs/openbugs'
-    version_added = '1.6'
-
-    _FUNCTIONS = (
-        # Scalar functions
-        'abs', 'arccos', 'arccosh', 'arcsin', 'arcsinh', 'arctan', 'arctanh',
-        'cloglog', 'cos', 'cosh', 'cumulative', 'cut', 'density', 'deviance',
-        'equals', 'expr', 'gammap', 'ilogit', 'icloglog', 'integral', 'log',
-        'logfact', 'loggam', 'logit', 'max', 'min', 'phi', 'post.p.value',
-        'pow', 'prior.p.value', 'probit', 'replicate.post', 'replicate.prior',
-        'round', 'sin', 'sinh', 'solution', 'sqrt', 'step', 'tan', 'tanh',
-        'trunc',
-        # Vector functions
-        'inprod', 'interp.lin', 'inverse', 'logdet', 'mean', 'eigen.vals',
-        'ode', 'prod', 'p.valueM', 'rank', 'ranked', 'replicate.postM',
-        'sd', 'sort', 'sum',
-        # Special
-        'D', 'I', 'F', 'T', 'C')
-    """ OpenBUGS built-in functions
-
-    From http://www.openbugs.info/Manuals/ModelSpecification.html#ContentsAII
-
-    This also includes
-
-    - T, C, I : Truncation and censoring.
-      ``T`` and ``C`` are in OpenBUGS. ``I`` in WinBUGS.
-    - D : ODE
-    - F : Functional http://www.openbugs.info/Examples/Functionals.html
-
-    """
-
-    _DISTRIBUTIONS = ('dbern', 'dbin', 'dcat', 'dnegbin', 'dpois',
-                      'dhyper', 'dbeta', 'dchisqr', 'ddexp', 'dexp',
-                      'dflat', 'dgamma', 'dgev', 'df', 'dggamma', 'dgpar',
-                      'dloglik', 'dlnorm', 'dlogis', 'dnorm', 'dpar',
-                      'dt', 'dunif', 'dweib', 'dmulti', 'ddirch', 'dmnorm',
-                      'dmt', 'dwish')
-    """ OpenBUGS built-in distributions
-
-    Functions from
-    http://www.openbugs.info/Manuals/ModelSpecification.html#ContentsAI
-    """
-
-    tokens = {
-        'whitespace': [
-            (r"\s+", Text),
-        ],
-        'comments': [
-            # Comments
-            (r'#.*$', Comment.Single),
-        ],
-        'root': [
-            # Comments
-            include('comments'),
-            include('whitespace'),
-            # Block start
-            (r'(model)(\s+)(\{)',
-             bygroups(Keyword.Namespace, Text, Punctuation)),
-            # Reserved Words
-            (r'(for|in)(?![\w.])', Keyword.Reserved),
-            # Built-in Functions
-            (r'({})(?=\s*\()'.format(r'|'.join(_FUNCTIONS + _DISTRIBUTIONS)),
-             Name.Builtin),
-            # Regular variable names
-            (r'[A-Za-z][\w.]*', Name),
-            # Number Literals
-            (r'[-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?', Number),
-            # Punctuation
-            (r'\[|\]|\(|\)|:|,|;', Punctuation),
-            # Assignment operators
-            # SLexer makes these tokens Operators.
-            (r'<-|~', Operator),
-            # Infix and prefix operators
-            (r'\+|-|\*|/', Operator),
-            # Block
-            (r'[{}]', Punctuation),
-        ]
-    }
-
-    def analyse_text(text):
-        if re.search(r"^\s*model\s*{", text, re.M):
-            return 0.7
-        else:
-            return 0.0
-
-
-class JagsLexer(RegexLexer):
-    """
-    Pygments Lexer for JAGS.
-    """
-
-    name = 'JAGS'
-    aliases = ['jags']
-    filenames = ['*.jag', '*.bug']
-    url = 'https://mcmc-jags.sourceforge.io'
-    version_added = '1.6'
-
-    # JAGS
-    _FUNCTIONS = (
-        'abs', 'arccos', 'arccosh', 'arcsin', 'arcsinh', 'arctan', 'arctanh',
-        'cos', 'cosh', 'cloglog',
-        'equals', 'exp', 'icloglog', 'ifelse', 'ilogit', 'log', 'logfact',
-        'loggam', 'logit', 'phi', 'pow', 'probit', 'round', 'sin', 'sinh',
-        'sqrt', 'step', 'tan', 'tanh', 'trunc', 'inprod', 'interp.lin',
-        'logdet', 'max', 'mean', 'min', 'prod', 'sum', 'sd', 'inverse',
-        'rank', 'sort', 't', 'acos', 'acosh', 'asin', 'asinh', 'atan',
-        # Truncation/Censoring (should I include)
-        'T', 'I')
-    # Distributions with density, probability and quartile functions
-    _DISTRIBUTIONS = tuple(f'[dpq]{x}' for x in
-                           ('bern', 'beta', 'dchiqsqr', 'ddexp', 'dexp',
-                            'df', 'gamma', 'gen.gamma', 'logis', 'lnorm',
-                            'negbin', 'nchisqr', 'norm', 'par', 'pois', 'weib'))
-    # Other distributions without density and probability
-    _OTHER_DISTRIBUTIONS = (
-        'dt', 'dunif', 'dbetabin', 'dbern', 'dbin', 'dcat', 'dhyper',
-        'ddirch', 'dmnorm', 'dwish', 'dmt', 'dmulti', 'dbinom', 'dchisq',
-        'dnbinom', 'dweibull', 'ddirich')
-
-    tokens = {
-        'whitespace': [
-            (r"\s+", Text),
-        ],
-        'names': [
-            # Regular variable names
-            (r'[a-zA-Z][\w.]*\b', Name),
-        ],
-        'comments': [
-            # do not use stateful comments
-            (r'(?s)/\*.*?\*/', Comment.Multiline),
-            # Comments
-            (r'#.*$', Comment.Single),
-        ],
-        'root': [
-            # Comments
-            include('comments'),
-            include('whitespace'),
-            # Block start
-            (r'(model|data)(\s+)(\{)',
-             bygroups(Keyword.Namespace, Text, Punctuation)),
-            (r'var(?![\w.])', Keyword.Declaration),
-            # Reserved Words
-            (r'(for|in)(?![\w.])', Keyword.Reserved),
-            # Builtins
-            # Need to use lookahead because . is a valid char
-            (r'({})(?=\s*\()'.format(r'|'.join(_FUNCTIONS
-                                          + _DISTRIBUTIONS
-                                          + _OTHER_DISTRIBUTIONS)),
-             Name.Builtin),
-            # Names
-            include('names'),
-            # Number Literals
-            (r'[-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?', Number),
-            (r'\[|\]|\(|\)|:|,|;', Punctuation),
-            # Assignment operators
-            (r'<-|~', Operator),
-            # # JAGS includes many more than OpenBUGS
-            (r'\+|-|\*|\/|\|\|[&]{2}|[<>=]=?|\^|%.*?%', Operator),
-            (r'[{}]', Punctuation),
-        ]
-    }
-
-    def analyse_text(text):
-        if re.search(r'^\s*model\s*\{', text, re.M):
-            if re.search(r'^\s*data\s*\{', text, re.M):
-                return 0.9
-            elif re.search(r'^\s*var', text, re.M):
-                return 0.9
-            else:
-                return 0.3
-        else:
-            return 0
-
-
-class StanLexer(RegexLexer):
-    """Pygments Lexer for Stan models.
-
-    The Stan modeling language is specified in the *Stan Modeling Language
-    User's Guide and Reference Manual, v2.17.0*,
-    `pdf `__.
-    """
-
-    name = 'Stan'
-    aliases = ['stan']
-    filenames = ['*.stan']
-    url = 'https://mc-stan.org'
-    version_added = '1.6'
-
-    tokens = {
-        'whitespace': [
-            (r"\s+", Text),
-        ],
-        'comments': [
-            (r'(?s)/\*.*?\*/', Comment.Multiline),
-            # Comments
-            (r'(//|#).*$', Comment.Single),
-        ],
-        'root': [
-            (r'"[^"]*"', String),
-            # Comments
-            include('comments'),
-            # block start
-            include('whitespace'),
-            # Block start
-            (r'({})(\s*)(\{{)'.format(r'|'.join(('functions', 'data', r'transformed\s+?data',
-                        'parameters', r'transformed\s+parameters',
-                        'model', r'generated\s+quantities'))),
-             bygroups(Keyword.Namespace, Text, Punctuation)),
-            # target keyword
-            (r'target\s*\+=', Keyword),
-            # Reserved Words
-            (r'({})\b'.format(r'|'.join(_stan_builtins.KEYWORDS)), Keyword),
-            # Truncation
-            (r'T(?=\s*\[)', Keyword),
-            # Data types
-            (r'({})\b'.format(r'|'.join(_stan_builtins.TYPES)), Keyword.Type),
-             # < should be punctuation, but elsewhere I can't tell if it is in
-             # a range constraint
-            (r'(<)(\s*)(upper|lower|offset|multiplier)(\s*)(=)',
-             bygroups(Operator, Whitespace, Keyword, Whitespace, Punctuation)),
-            (r'(,)(\s*)(upper)(\s*)(=)',
-             bygroups(Punctuation, Whitespace, Keyword, Whitespace, Punctuation)),
-            # Punctuation
-            (r"[;,\[\]()]", Punctuation),
-            # Builtin
-            (r'({})(?=\s*\()'.format('|'.join(_stan_builtins.FUNCTIONS)), Name.Builtin),
-            (r'(~)(\s*)({})(?=\s*\()'.format('|'.join(_stan_builtins.DISTRIBUTIONS)),
-                bygroups(Operator, Whitespace, Name.Builtin)),
-            # Special names ending in __, like lp__
-            (r'[A-Za-z]\w*__\b', Name.Builtin.Pseudo),
-            (r'({})\b'.format(r'|'.join(_stan_builtins.RESERVED)), Keyword.Reserved),
-            # user-defined functions
-            (r'[A-Za-z]\w*(?=\s*\()]', Name.Function),
-            # Imaginary Literals
-            (r'[0-9]+(\.[0-9]*)?([eE][+-]?[0-9]+)?i', Number.Float),
-            (r'\.[0-9]+([eE][+-]?[0-9]+)?i', Number.Float),
-            (r'[0-9]+i', Number.Float),
-            # Real Literals
-            (r'[0-9]+(\.[0-9]*)?([eE][+-]?[0-9]+)?', Number.Float),
-            (r'\.[0-9]+([eE][+-]?[0-9]+)?', Number.Float),
-            # Integer Literals
-            (r'[0-9]+', Number.Integer),
-            # Regular variable names
-            (r'[A-Za-z]\w*\b', Name),
-            # Assignment operators
-            (r'<-|(?:\+|-|\.?/|\.?\*|=)?=|~', Operator),
-            # Infix, prefix and postfix operators (and = )
-            (r"\+|-|\.?\*|\.?/|\\|'|\.?\^|!=?|<=?|>=?|\|\||&&|%|\?|:|%/%|!", Operator),
-            # Block delimiters
-            (r'[{}]', Punctuation),
-            # Distribution |
-            (r'\|', Punctuation)
-        ]
-    }
-
-    def analyse_text(text):
-        if re.search(r'^\s*parameters\s*\{', text, re.M):
-            return 1.0
-        else:
-            return 0.0
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/modula2.py b/venv/lib/python3.10/site-packages/pygments/lexers/modula2.py
deleted file mode 100644
index 713f4722f7d8d4837450a54bde5b1b7160af9911..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/modula2.py
+++ /dev/null
@@ -1,1579 +0,0 @@
-"""
-    pygments.lexers.modula2
-    ~~~~~~~~~~~~~~~~~~~~~~~
-
-    Multi-Dialect Lexer for Modula-2.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, include
-from pygments.util import get_bool_opt, get_list_opt
-from pygments.token import Text, Comment, Operator, Keyword, Name, \
-    String, Number, Punctuation, Error
-
-__all__ = ['Modula2Lexer']
-
-
-# Multi-Dialect Modula-2 Lexer
-class Modula2Lexer(RegexLexer):
-    """
-    For Modula-2 source code.
-
-    The Modula-2 lexer supports several dialects.  By default, it operates in
-    fallback mode, recognising the *combined* literals, punctuation symbols
-    and operators of all supported dialects, and the *combined* reserved words
-    and builtins of PIM Modula-2, ISO Modula-2 and Modula-2 R10, while not
-    differentiating between library defined identifiers.
-
-    To select a specific dialect, a dialect option may be passed
-    or a dialect tag may be embedded into a source file.
-
-    Dialect Options:
-
-    `m2pim`
-        Select PIM Modula-2 dialect.
-    `m2iso`
-        Select ISO Modula-2 dialect.
-    `m2r10`
-        Select Modula-2 R10 dialect.
-    `objm2`
-        Select Objective Modula-2 dialect.
-
-    The PIM and ISO dialect options may be qualified with a language extension.
-
-    Language Extensions:
-
-    `+aglet`
-        Select Aglet Modula-2 extensions, available with m2iso.
-    `+gm2`
-        Select GNU Modula-2 extensions, available with m2pim.
-    `+p1`
-        Select p1 Modula-2 extensions, available with m2iso.
-    `+xds`
-        Select XDS Modula-2 extensions, available with m2iso.
-
-
-    Passing a Dialect Option via Unix Commandline Interface
-
-    Dialect options may be passed to the lexer using the `dialect` key.
-    Only one such option should be passed. If multiple dialect options are
-    passed, the first valid option is used, any subsequent options are ignored.
-
-    Examples:
-
-    `$ pygmentize -O full,dialect=m2iso -f html -o /path/to/output /path/to/input`
-        Use ISO dialect to render input to HTML output
-    `$ pygmentize -O full,dialect=m2iso+p1 -f rtf -o /path/to/output /path/to/input`
-        Use ISO dialect with p1 extensions to render input to RTF output
-
-
-    Embedding a Dialect Option within a source file
-
-    A dialect option may be embedded in a source file in form of a dialect
-    tag, a specially formatted comment that specifies a dialect option.
-
-    Dialect Tag EBNF::
-
-       dialectTag :
-           OpeningCommentDelim Prefix dialectOption ClosingCommentDelim ;
-
-       dialectOption :
-           'm2pim' | 'm2iso' | 'm2r10' | 'objm2' |
-           'm2iso+aglet' | 'm2pim+gm2' | 'm2iso+p1' | 'm2iso+xds' ;
-
-       Prefix : '!' ;
-
-       OpeningCommentDelim : '(*' ;
-
-       ClosingCommentDelim : '*)' ;
-
-    No whitespace is permitted between the tokens of a dialect tag.
-
-    In the event that a source file contains multiple dialect tags, the first
-    tag that contains a valid dialect option will be used and any subsequent
-    dialect tags will be ignored.  Ideally, a dialect tag should be placed
-    at the beginning of a source file.
-
-    An embedded dialect tag overrides a dialect option set via command line.
-
-    Examples:
-
-    ``(*!m2r10*) DEFINITION MODULE Foobar; ...``
-        Use Modula2 R10 dialect to render this source file.
-    ``(*!m2pim+gm2*) DEFINITION MODULE Bazbam; ...``
-        Use PIM dialect with GNU extensions to render this source file.
-
-
-    Algol Publication Mode:
-
-    In Algol publication mode, source text is rendered for publication of
-    algorithms in scientific papers and academic texts, following the format
-    of the Revised Algol-60 Language Report.  It is activated by passing
-    one of two corresponding styles as an option:
-
-    `algol`
-        render reserved words lowercase underline boldface
-        and builtins lowercase boldface italic
-    `algol_nu`
-        render reserved words lowercase boldface (no underlining)
-        and builtins lowercase boldface italic
-
-    The lexer automatically performs the required lowercase conversion when
-    this mode is activated.
-
-    Example:
-
-    ``$ pygmentize -O full,style=algol -f latex -o /path/to/output /path/to/input``
-        Render input file in Algol publication mode to LaTeX output.
-
-
-    Rendering Mode of First Class ADT Identifiers:
-
-    The rendering of standard library first class ADT identifiers is controlled
-    by option flag "treat_stdlib_adts_as_builtins".
-
-    When this option is turned on, standard library ADT identifiers are rendered
-    as builtins.  When it is turned off, they are rendered as ordinary library
-    identifiers.
-
-    `treat_stdlib_adts_as_builtins` (default: On)
-
-    The option is useful for dialects that support ADTs as first class objects
-    and provide ADTs in the standard library that would otherwise be built-in.
-
-    At present, only Modula-2 R10 supports library ADTs as first class objects
-    and therefore, no ADT identifiers are defined for any other dialects.
-
-    Example:
-
-    ``$ pygmentize -O full,dialect=m2r10,treat_stdlib_adts_as_builtins=Off ...``
-        Render standard library ADTs as ordinary library types.
-
-    .. versionchanged:: 2.1
-       Added multi-dialect support.
-    """
-    name = 'Modula-2'
-    url = 'http://www.modula2.org/'
-    aliases = ['modula2', 'm2']
-    filenames = ['*.def', '*.mod']
-    mimetypes = ['text/x-modula2']
-    version_added = '1.3'
-
-    flags = re.MULTILINE | re.DOTALL
-
-    tokens = {
-        'whitespace': [
-            (r'\n+', Text),  # blank lines
-            (r'\s+', Text),  # whitespace
-        ],
-        'dialecttags': [
-            # PIM Dialect Tag
-            (r'\(\*!m2pim\*\)', Comment.Special),
-            # ISO Dialect Tag
-            (r'\(\*!m2iso\*\)', Comment.Special),
-            # M2R10 Dialect Tag
-            (r'\(\*!m2r10\*\)', Comment.Special),
-            # ObjM2 Dialect Tag
-            (r'\(\*!objm2\*\)', Comment.Special),
-            # Aglet Extensions Dialect Tag
-            (r'\(\*!m2iso\+aglet\*\)', Comment.Special),
-            # GNU Extensions Dialect Tag
-            (r'\(\*!m2pim\+gm2\*\)', Comment.Special),
-            # p1 Extensions Dialect Tag
-            (r'\(\*!m2iso\+p1\*\)', Comment.Special),
-            # XDS Extensions Dialect Tag
-            (r'\(\*!m2iso\+xds\*\)', Comment.Special),
-        ],
-        'identifiers': [
-            (r'([a-zA-Z_$][\w$]*)', Name),
-        ],
-        'prefixed_number_literals': [
-            #
-            # Base-2, whole number
-            (r'0b[01]+(\'[01]+)*', Number.Bin),
-            #
-            # Base-16, whole number
-            (r'0[ux][0-9A-F]+(\'[0-9A-F]+)*', Number.Hex),
-        ],
-        'plain_number_literals': [
-            #
-            # Base-10, real number with exponent
-            (r'[0-9]+(\'[0-9]+)*'  # integral part
-             r'\.[0-9]+(\'[0-9]+)*'  # fractional part
-             r'[eE][+-]?[0-9]+(\'[0-9]+)*',  # exponent
-             Number.Float),
-            #
-            # Base-10, real number without exponent
-            (r'[0-9]+(\'[0-9]+)*'  # integral part
-             r'\.[0-9]+(\'[0-9]+)*',  # fractional part
-             Number.Float),
-            #
-            # Base-10, whole number
-            (r'[0-9]+(\'[0-9]+)*', Number.Integer),
-        ],
-        'suffixed_number_literals': [
-            #
-            # Base-8, whole number
-            (r'[0-7]+B', Number.Oct),
-            #
-            # Base-8, character code
-            (r'[0-7]+C', Number.Oct),
-            #
-            # Base-16, number
-            (r'[0-9A-F]+H', Number.Hex),
-        ],
-        'string_literals': [
-            (r'"(\\\\|\\[^\\]|[^"\\])*"', String.Double),
-            (r"'(\\\\|\\[^\\]|[^'\\])*'", String.Single),
-        ],
-        'digraph_operators': [
-            # Dot Product Operator
-            (r'\*\.', Operator),
-            # Array Concatenation Operator
-            (r'\+>', Operator),  # M2R10 + ObjM2
-            # Inequality Operator
-            (r'<>', Operator),  # ISO + PIM
-            # Less-Or-Equal, Subset
-            (r'<=', Operator),
-            # Greater-Or-Equal, Superset
-            (r'>=', Operator),
-            # Identity Operator
-            (r'==', Operator),  # M2R10 + ObjM2
-            # Type Conversion Operator
-            (r'::', Operator),  # M2R10 + ObjM2
-            # Assignment Symbol
-            (r':=', Operator),
-            # Postfix Increment Mutator
-            (r'\+\+', Operator),  # M2R10 + ObjM2
-            # Postfix Decrement Mutator
-            (r'--', Operator),  # M2R10 + ObjM2
-        ],
-        'unigraph_operators': [
-            # Arithmetic Operators
-            (r'[+-]', Operator),
-            (r'[*/]', Operator),
-            # ISO 80000-2 compliant Set Difference Operator
-            (r'\\', Operator),  # M2R10 + ObjM2
-            # Relational Operators
-            (r'[=#<>]', Operator),
-            # Dereferencing Operator
-            (r'\^', Operator),
-            # Dereferencing Operator Synonym
-            (r'@', Operator),  # ISO
-            # Logical AND Operator Synonym
-            (r'&', Operator),  # PIM + ISO
-            # Logical NOT Operator Synonym
-            (r'~', Operator),  # PIM + ISO
-            # Smalltalk Message Prefix
-            (r'`', Operator),  # ObjM2
-        ],
-        'digraph_punctuation': [
-            # Range Constructor
-            (r'\.\.', Punctuation),
-            # Opening Chevron Bracket
-            (r'<<', Punctuation),  # M2R10 + ISO
-            # Closing Chevron Bracket
-            (r'>>', Punctuation),  # M2R10 + ISO
-            # Blueprint Punctuation
-            (r'->', Punctuation),  # M2R10 + ISO
-            # Distinguish |# and # in M2 R10
-            (r'\|#', Punctuation),
-            # Distinguish ## and # in M2 R10
-            (r'##', Punctuation),
-            # Distinguish |* and * in M2 R10
-            (r'\|\*', Punctuation),
-        ],
-        'unigraph_punctuation': [
-            # Common Punctuation
-            (r'[()\[\]{},.:;|]', Punctuation),
-            # Case Label Separator Synonym
-            (r'!', Punctuation),  # ISO
-            # Blueprint Punctuation
-            (r'\?', Punctuation),  # M2R10 + ObjM2
-        ],
-        'comments': [
-            # Single Line Comment
-            (r'^//.*?\n', Comment.Single),  # M2R10 + ObjM2
-            # Block Comment
-            (r'\(\*([^$].*?)\*\)', Comment.Multiline),
-            # Template Block Comment
-            (r'/\*(.*?)\*/', Comment.Multiline),  # M2R10 + ObjM2
-        ],
-        'pragmas': [
-            # ISO Style Pragmas
-            (r'<\*.*?\*>', Comment.Preproc),  # ISO, M2R10 + ObjM2
-            # Pascal Style Pragmas
-            (r'\(\*\$.*?\*\)', Comment.Preproc),  # PIM
-        ],
-        'root': [
-            include('whitespace'),
-            include('dialecttags'),
-            include('pragmas'),
-            include('comments'),
-            include('identifiers'),
-            include('suffixed_number_literals'),  # PIM + ISO
-            include('prefixed_number_literals'),  # M2R10 + ObjM2
-            include('plain_number_literals'),
-            include('string_literals'),
-            include('digraph_punctuation'),
-            include('digraph_operators'),
-            include('unigraph_punctuation'),
-            include('unigraph_operators'),
-        ]
-    }
-
-#  C o m m o n   D a t a s e t s
-
-    # Common Reserved Words Dataset
-    common_reserved_words = (
-        # 37 common reserved words
-        'AND', 'ARRAY', 'BEGIN', 'BY', 'CASE', 'CONST', 'DEFINITION', 'DIV',
-        'DO', 'ELSE', 'ELSIF', 'END', 'EXIT', 'FOR', 'FROM', 'IF',
-        'IMPLEMENTATION', 'IMPORT', 'IN', 'LOOP', 'MOD', 'MODULE', 'NOT',
-        'OF', 'OR', 'POINTER', 'PROCEDURE', 'RECORD', 'REPEAT', 'RETURN',
-        'SET', 'THEN', 'TO', 'TYPE', 'UNTIL', 'VAR', 'WHILE',
-    )
-
-    # Common Builtins Dataset
-    common_builtins = (
-        # 16 common builtins
-        'ABS', 'BOOLEAN', 'CARDINAL', 'CHAR', 'CHR', 'FALSE', 'INTEGER',
-        'LONGINT', 'LONGREAL', 'MAX', 'MIN', 'NIL', 'ODD', 'ORD', 'REAL',
-        'TRUE',
-    )
-
-    # Common Pseudo-Module Builtins Dataset
-    common_pseudo_builtins = (
-        # 4 common pseudo builtins
-        'ADDRESS', 'BYTE', 'WORD', 'ADR'
-    )
-
-#  P I M   M o d u l a - 2   D a t a s e t s
-
-    # Lexemes to Mark as Error Tokens for PIM Modula-2
-    pim_lexemes_to_reject = (
-        '!', '`', '@', '$', '%', '?', '\\', '==', '++', '--', '::', '*.',
-        '+>', '->', '<<', '>>', '|#', '##',
-    )
-
-    # PIM Modula-2 Additional Reserved Words Dataset
-    pim_additional_reserved_words = (
-        # 3 additional reserved words
-        'EXPORT', 'QUALIFIED', 'WITH',
-    )
-
-    # PIM Modula-2 Additional Builtins Dataset
-    pim_additional_builtins = (
-        # 16 additional builtins
-        'BITSET', 'CAP', 'DEC', 'DISPOSE', 'EXCL', 'FLOAT', 'HALT', 'HIGH',
-        'INC', 'INCL', 'NEW', 'NIL', 'PROC', 'SIZE', 'TRUNC', 'VAL',
-    )
-
-    # PIM Modula-2 Additional Pseudo-Module Builtins Dataset
-    pim_additional_pseudo_builtins = (
-        # 5 additional pseudo builtins
-        'SYSTEM', 'PROCESS', 'TSIZE', 'NEWPROCESS', 'TRANSFER',
-    )
-
-#  I S O   M o d u l a - 2   D a t a s e t s
-
-    # Lexemes to Mark as Error Tokens for ISO Modula-2
-    iso_lexemes_to_reject = (
-        '`', '$', '%', '?', '\\', '==', '++', '--', '::', '*.', '+>', '->',
-        '<<', '>>', '|#', '##',
-    )
-
-    # ISO Modula-2 Additional Reserved Words Dataset
-    iso_additional_reserved_words = (
-        # 9 additional reserved words (ISO 10514-1)
-        'EXCEPT', 'EXPORT', 'FINALLY', 'FORWARD', 'PACKEDSET', 'QUALIFIED',
-        'REM', 'RETRY', 'WITH',
-        # 10 additional reserved words (ISO 10514-2 & ISO 10514-3)
-        'ABSTRACT', 'AS', 'CLASS', 'GUARD', 'INHERIT', 'OVERRIDE', 'READONLY',
-        'REVEAL', 'TRACED', 'UNSAFEGUARDED',
-    )
-
-    # ISO Modula-2 Additional Builtins Dataset
-    iso_additional_builtins = (
-        # 26 additional builtins (ISO 10514-1)
-        'BITSET', 'CAP', 'CMPLX', 'COMPLEX', 'DEC', 'DISPOSE', 'EXCL', 'FLOAT',
-        'HALT', 'HIGH', 'IM', 'INC', 'INCL', 'INT', 'INTERRUPTIBLE',  'LENGTH',
-        'LFLOAT', 'LONGCOMPLEX', 'NEW', 'PROC', 'PROTECTION', 'RE', 'SIZE',
-        'TRUNC', 'UNINTERRUBTIBLE', 'VAL',
-        # 5 additional builtins (ISO 10514-2 & ISO 10514-3)
-        'CREATE', 'DESTROY', 'EMPTY', 'ISMEMBER', 'SELF',
-    )
-
-    # ISO Modula-2 Additional Pseudo-Module Builtins Dataset
-    iso_additional_pseudo_builtins = (
-        # 14 additional builtins (SYSTEM)
-        'SYSTEM', 'BITSPERLOC', 'LOCSPERBYTE', 'LOCSPERWORD', 'LOC',
-        'ADDADR', 'SUBADR', 'DIFADR', 'MAKEADR', 'ADR',
-        'ROTATE', 'SHIFT', 'CAST', 'TSIZE',
-        # 13 additional builtins (COROUTINES)
-        'COROUTINES', 'ATTACH', 'COROUTINE', 'CURRENT', 'DETACH', 'HANDLER',
-        'INTERRUPTSOURCE', 'IOTRANSFER', 'IsATTACHED', 'LISTEN',
-        'NEWCOROUTINE', 'PROT', 'TRANSFER',
-        # 9 additional builtins (EXCEPTIONS)
-        'EXCEPTIONS', 'AllocateSource', 'CurrentNumber', 'ExceptionNumber',
-        'ExceptionSource', 'GetMessage', 'IsCurrentSource',
-        'IsExceptionalExecution', 'RAISE',
-        # 3 additional builtins (TERMINATION)
-        'TERMINATION', 'IsTerminating', 'HasHalted',
-        # 4 additional builtins (M2EXCEPTION)
-        'M2EXCEPTION', 'M2Exceptions', 'M2Exception', 'IsM2Exception',
-        'indexException', 'rangeException', 'caseSelectException',
-        'invalidLocation', 'functionException', 'wholeValueException',
-        'wholeDivException', 'realValueException', 'realDivException',
-        'complexValueException', 'complexDivException', 'protException',
-        'sysException', 'coException', 'exException',
-    )
-
-#  M o d u l a - 2   R 1 0   D a t a s e t s
-
-    # Lexemes to Mark as Error Tokens for Modula-2 R10
-    m2r10_lexemes_to_reject = (
-        '!', '`', '@', '$', '%', '&', '<>',
-    )
-
-    # Modula-2 R10 reserved words in addition to the common set
-    m2r10_additional_reserved_words = (
-        # 12 additional reserved words
-        'ALIAS', 'ARGLIST', 'BLUEPRINT', 'COPY', 'GENLIB', 'INDETERMINATE',
-        'NEW', 'NONE', 'OPAQUE', 'REFERENTIAL', 'RELEASE', 'RETAIN',
-        # 2 additional reserved words with symbolic assembly option
-        'ASM', 'REG',
-    )
-
-    # Modula-2 R10 builtins in addition to the common set
-    m2r10_additional_builtins = (
-        # 26 additional builtins
-        'CARDINAL', 'COUNT', 'EMPTY', 'EXISTS', 'INSERT', 'LENGTH', 'LONGCARD',
-        'OCTET', 'PTR', 'PRED', 'READ', 'READNEW', 'REMOVE', 'RETRIEVE', 'SORT',
-        'STORE', 'SUBSET', 'SUCC', 'TLIMIT', 'TMAX', 'TMIN', 'TRUE', 'TSIZE',
-        'UNICHAR', 'WRITE', 'WRITEF',
-    )
-
-    # Modula-2 R10 Additional Pseudo-Module Builtins Dataset
-    m2r10_additional_pseudo_builtins = (
-        # 13 additional builtins (TPROPERTIES)
-        'TPROPERTIES', 'PROPERTY', 'LITERAL', 'TPROPERTY', 'TLITERAL',
-        'TBUILTIN', 'TDYN', 'TREFC', 'TNIL', 'TBASE', 'TPRECISION',
-        'TMAXEXP', 'TMINEXP',
-        # 4 additional builtins (CONVERSION)
-        'CONVERSION', 'TSXFSIZE', 'SXF', 'VAL',
-        # 35 additional builtins (UNSAFE)
-        'UNSAFE', 'CAST', 'INTRINSIC', 'AVAIL', 'ADD', 'SUB', 'ADDC', 'SUBC',
-        'FETCHADD', 'FETCHSUB', 'SHL', 'SHR', 'ASHR', 'ROTL', 'ROTR', 'ROTLC',
-        'ROTRC', 'BWNOT', 'BWAND', 'BWOR', 'BWXOR', 'BWNAND', 'BWNOR',
-        'SETBIT', 'TESTBIT', 'LSBIT', 'MSBIT', 'CSBITS', 'BAIL', 'HALT',
-        'TODO', 'FFI', 'ADDR', 'VARGLIST', 'VARGC',
-        # 11 additional builtins (ATOMIC)
-        'ATOMIC', 'INTRINSIC', 'AVAIL', 'SWAP', 'CAS', 'INC', 'DEC', 'BWAND',
-        'BWNAND', 'BWOR', 'BWXOR',
-        # 7 additional builtins (COMPILER)
-        'COMPILER', 'DEBUG', 'MODNAME', 'PROCNAME', 'LINENUM', 'DEFAULT',
-        'HASH',
-        # 5 additional builtins (ASSEMBLER)
-        'ASSEMBLER', 'REGISTER', 'SETREG', 'GETREG', 'CODE',
-    )
-
-#  O b j e c t i v e   M o d u l a - 2   D a t a s e t s
-
-    # Lexemes to Mark as Error Tokens for Objective Modula-2
-    objm2_lexemes_to_reject = (
-        '!', '$', '%', '&', '<>',
-    )
-
-    # Objective Modula-2 Extensions
-    # reserved words in addition to Modula-2 R10
-    objm2_additional_reserved_words = (
-        # 16 additional reserved words
-        'BYCOPY', 'BYREF', 'CLASS', 'CONTINUE', 'CRITICAL', 'INOUT', 'METHOD',
-        'ON', 'OPTIONAL', 'OUT', 'PRIVATE', 'PROTECTED', 'PROTOCOL', 'PUBLIC',
-        'SUPER', 'TRY',
-    )
-
-    # Objective Modula-2 Extensions
-    # builtins in addition to Modula-2 R10
-    objm2_additional_builtins = (
-        # 3 additional builtins
-        'OBJECT', 'NO', 'YES',
-    )
-
-    # Objective Modula-2 Extensions
-    # pseudo-module builtins in addition to Modula-2 R10
-    objm2_additional_pseudo_builtins = (
-        # None
-    )
-
-#  A g l e t   M o d u l a - 2   D a t a s e t s
-
-    # Aglet Extensions
-    # reserved words in addition to ISO Modula-2
-    aglet_additional_reserved_words = (
-        # None
-    )
-
-    # Aglet Extensions
-    # builtins in addition to ISO Modula-2
-    aglet_additional_builtins = (
-        # 9 additional builtins
-        'BITSET8', 'BITSET16', 'BITSET32', 'CARDINAL8', 'CARDINAL16',
-        'CARDINAL32', 'INTEGER8', 'INTEGER16', 'INTEGER32',
-    )
-
-    # Aglet Modula-2 Extensions
-    # pseudo-module builtins in addition to ISO Modula-2
-    aglet_additional_pseudo_builtins = (
-        # None
-    )
-
-#  G N U   M o d u l a - 2   D a t a s e t s
-
-    # GNU Extensions
-    # reserved words in addition to PIM Modula-2
-    gm2_additional_reserved_words = (
-        # 10 additional reserved words
-        'ASM', '__ATTRIBUTE__', '__BUILTIN__', '__COLUMN__', '__DATE__',
-        '__FILE__', '__FUNCTION__', '__LINE__', '__MODULE__', 'VOLATILE',
-    )
-
-    # GNU Extensions
-    # builtins in addition to PIM Modula-2
-    gm2_additional_builtins = (
-        # 21 additional builtins
-        'BITSET8', 'BITSET16', 'BITSET32', 'CARDINAL8', 'CARDINAL16',
-        'CARDINAL32', 'CARDINAL64', 'COMPLEX32', 'COMPLEX64', 'COMPLEX96',
-        'COMPLEX128', 'INTEGER8', 'INTEGER16', 'INTEGER32', 'INTEGER64',
-        'REAL8', 'REAL16', 'REAL32', 'REAL96', 'REAL128', 'THROW',
-    )
-
-    # GNU Extensions
-    # pseudo-module builtins in addition to PIM Modula-2
-    gm2_additional_pseudo_builtins = (
-        # None
-    )
-
-#  p 1   M o d u l a - 2   D a t a s e t s
-
-    # p1 Extensions
-    # reserved words in addition to ISO Modula-2
-    p1_additional_reserved_words = (
-        # None
-    )
-
-    # p1 Extensions
-    # builtins in addition to ISO Modula-2
-    p1_additional_builtins = (
-        # None
-    )
-
-    # p1 Modula-2 Extensions
-    # pseudo-module builtins in addition to ISO Modula-2
-    p1_additional_pseudo_builtins = (
-        # 1 additional builtin
-        'BCD',
-    )
-
-#  X D S   M o d u l a - 2   D a t a s e t s
-
-    # XDS Extensions
-    # reserved words in addition to ISO Modula-2
-    xds_additional_reserved_words = (
-        # 1 additional reserved word
-        'SEQ',
-    )
-
-    # XDS Extensions
-    # builtins in addition to ISO Modula-2
-    xds_additional_builtins = (
-        # 9 additional builtins
-        'ASH', 'ASSERT', 'DIFFADR_TYPE', 'ENTIER', 'INDEX', 'LEN',
-        'LONGCARD', 'SHORTCARD', 'SHORTINT',
-    )
-
-    # XDS Modula-2 Extensions
-    # pseudo-module builtins in addition to ISO Modula-2
-    xds_additional_pseudo_builtins = (
-        # 22 additional builtins (SYSTEM)
-        'PROCESS', 'NEWPROCESS', 'BOOL8', 'BOOL16', 'BOOL32', 'CARD8',
-        'CARD16', 'CARD32', 'INT8', 'INT16', 'INT32', 'REF', 'MOVE',
-        'FILL', 'GET', 'PUT', 'CC', 'int', 'unsigned', 'size_t', 'void'
-        # 3 additional builtins (COMPILER)
-        'COMPILER', 'OPTION', 'EQUATION'
-    )
-
-#  P I M   S t a n d a r d   L i b r a r y   D a t a s e t s
-
-    # PIM Modula-2 Standard Library Modules Dataset
-    pim_stdlib_module_identifiers = (
-        'Terminal', 'FileSystem', 'InOut', 'RealInOut', 'MathLib0', 'Storage',
-    )
-
-    # PIM Modula-2 Standard Library Types Dataset
-    pim_stdlib_type_identifiers = (
-        'Flag', 'FlagSet', 'Response', 'Command', 'Lock', 'Permission',
-        'MediumType', 'File', 'FileProc', 'DirectoryProc', 'FileCommand',
-        'DirectoryCommand',
-    )
-
-    # PIM Modula-2 Standard Library Procedures Dataset
-    pim_stdlib_proc_identifiers = (
-        'Read', 'BusyRead', 'ReadAgain', 'Write', 'WriteString', 'WriteLn',
-        'Create', 'Lookup', 'Close', 'Delete', 'Rename', 'SetRead', 'SetWrite',
-        'SetModify', 'SetOpen', 'Doio', 'SetPos', 'GetPos', 'Length', 'Reset',
-        'Again', 'ReadWord', 'WriteWord', 'ReadChar', 'WriteChar',
-        'CreateMedium', 'DeleteMedium', 'AssignName', 'DeassignName',
-        'ReadMedium', 'LookupMedium', 'OpenInput', 'OpenOutput', 'CloseInput',
-        'CloseOutput', 'ReadString', 'ReadInt', 'ReadCard', 'ReadWrd',
-        'WriteInt', 'WriteCard', 'WriteOct', 'WriteHex', 'WriteWrd',
-        'ReadReal', 'WriteReal', 'WriteFixPt', 'WriteRealOct', 'sqrt', 'exp',
-        'ln', 'sin', 'cos', 'arctan', 'entier', 'ALLOCATE', 'DEALLOCATE',
-    )
-
-    # PIM Modula-2 Standard Library Variables Dataset
-    pim_stdlib_var_identifiers = (
-        'Done', 'termCH', 'in', 'out'
-    )
-
-    # PIM Modula-2 Standard Library Constants Dataset
-    pim_stdlib_const_identifiers = (
-        'EOL',
-    )
-
-#  I S O   S t a n d a r d   L i b r a r y   D a t a s e t s
-
-    # ISO Modula-2 Standard Library Modules Dataset
-    iso_stdlib_module_identifiers = (
-        # TO DO
-    )
-
-    # ISO Modula-2 Standard Library Types Dataset
-    iso_stdlib_type_identifiers = (
-        # TO DO
-    )
-
-    # ISO Modula-2 Standard Library Procedures Dataset
-    iso_stdlib_proc_identifiers = (
-        # TO DO
-    )
-
-    # ISO Modula-2 Standard Library Variables Dataset
-    iso_stdlib_var_identifiers = (
-        # TO DO
-    )
-
-    # ISO Modula-2 Standard Library Constants Dataset
-    iso_stdlib_const_identifiers = (
-        # TO DO
-    )
-
-#  M 2   R 1 0   S t a n d a r d   L i b r a r y   D a t a s e t s
-
-    # Modula-2 R10 Standard Library ADTs Dataset
-    m2r10_stdlib_adt_identifiers = (
-        'BCD', 'LONGBCD', 'BITSET', 'SHORTBITSET', 'LONGBITSET',
-        'LONGLONGBITSET', 'COMPLEX', 'LONGCOMPLEX', 'SHORTCARD', 'LONGLONGCARD',
-        'SHORTINT', 'LONGLONGINT', 'POSINT', 'SHORTPOSINT', 'LONGPOSINT',
-        'LONGLONGPOSINT', 'BITSET8', 'BITSET16', 'BITSET32', 'BITSET64',
-        'BITSET128', 'BS8', 'BS16', 'BS32', 'BS64', 'BS128', 'CARDINAL8',
-        'CARDINAL16', 'CARDINAL32', 'CARDINAL64', 'CARDINAL128', 'CARD8',
-        'CARD16', 'CARD32', 'CARD64', 'CARD128', 'INTEGER8', 'INTEGER16',
-        'INTEGER32', 'INTEGER64', 'INTEGER128', 'INT8', 'INT16', 'INT32',
-        'INT64', 'INT128', 'STRING', 'UNISTRING',
-    )
-
-    # Modula-2 R10 Standard Library Blueprints Dataset
-    m2r10_stdlib_blueprint_identifiers = (
-        'ProtoRoot', 'ProtoComputational', 'ProtoNumeric', 'ProtoScalar',
-        'ProtoNonScalar', 'ProtoCardinal', 'ProtoInteger', 'ProtoReal',
-        'ProtoComplex', 'ProtoVector', 'ProtoTuple', 'ProtoCompArray',
-        'ProtoCollection', 'ProtoStaticArray', 'ProtoStaticSet',
-        'ProtoStaticString', 'ProtoArray', 'ProtoString', 'ProtoSet',
-        'ProtoMultiSet', 'ProtoDictionary', 'ProtoMultiDict', 'ProtoExtension',
-        'ProtoIO', 'ProtoCardMath', 'ProtoIntMath', 'ProtoRealMath',
-    )
-
-    # Modula-2 R10 Standard Library Modules Dataset
-    m2r10_stdlib_module_identifiers = (
-        'ASCII', 'BooleanIO', 'CharIO', 'UnicharIO', 'OctetIO',
-        'CardinalIO', 'LongCardIO', 'IntegerIO', 'LongIntIO', 'RealIO',
-        'LongRealIO', 'BCDIO', 'LongBCDIO', 'CardMath', 'LongCardMath',
-        'IntMath', 'LongIntMath', 'RealMath', 'LongRealMath', 'BCDMath',
-        'LongBCDMath', 'FileIO', 'FileSystem', 'Storage', 'IOSupport',
-    )
-
-    # Modula-2 R10 Standard Library Types Dataset
-    m2r10_stdlib_type_identifiers = (
-        'File', 'Status',
-        # TO BE COMPLETED
-    )
-
-    # Modula-2 R10 Standard Library Procedures Dataset
-    m2r10_stdlib_proc_identifiers = (
-        'ALLOCATE', 'DEALLOCATE', 'SIZE',
-        # TO BE COMPLETED
-    )
-
-    # Modula-2 R10 Standard Library Variables Dataset
-    m2r10_stdlib_var_identifiers = (
-        'stdIn', 'stdOut', 'stdErr',
-    )
-
-    # Modula-2 R10 Standard Library Constants Dataset
-    m2r10_stdlib_const_identifiers = (
-        'pi', 'tau',
-    )
-
-#  D i a l e c t s
-
-    # Dialect modes
-    dialects = (
-        'unknown',
-        'm2pim', 'm2iso', 'm2r10', 'objm2',
-        'm2iso+aglet', 'm2pim+gm2', 'm2iso+p1', 'm2iso+xds',
-    )
-
-#   D a t a b a s e s
-
-    # Lexemes to Mark as Errors Database
-    lexemes_to_reject_db = {
-        # Lexemes to reject for unknown dialect
-        'unknown': (
-            # LEAVE THIS EMPTY
-        ),
-        # Lexemes to reject for PIM Modula-2
-        'm2pim': (
-            pim_lexemes_to_reject,
-        ),
-        # Lexemes to reject for ISO Modula-2
-        'm2iso': (
-            iso_lexemes_to_reject,
-        ),
-        # Lexemes to reject for Modula-2 R10
-        'm2r10': (
-            m2r10_lexemes_to_reject,
-        ),
-        # Lexemes to reject for Objective Modula-2
-        'objm2': (
-            objm2_lexemes_to_reject,
-        ),
-        # Lexemes to reject for Aglet Modula-2
-        'm2iso+aglet': (
-            iso_lexemes_to_reject,
-        ),
-        # Lexemes to reject for GNU Modula-2
-        'm2pim+gm2': (
-            pim_lexemes_to_reject,
-        ),
-        # Lexemes to reject for p1 Modula-2
-        'm2iso+p1': (
-            iso_lexemes_to_reject,
-        ),
-        # Lexemes to reject for XDS Modula-2
-        'm2iso+xds': (
-            iso_lexemes_to_reject,
-        ),
-    }
-
-    # Reserved Words Database
-    reserved_words_db = {
-        # Reserved words for unknown dialect
-        'unknown': (
-            common_reserved_words,
-            pim_additional_reserved_words,
-            iso_additional_reserved_words,
-            m2r10_additional_reserved_words,
-        ),
-
-        # Reserved words for PIM Modula-2
-        'm2pim': (
-            common_reserved_words,
-            pim_additional_reserved_words,
-        ),
-
-        # Reserved words for Modula-2 R10
-        'm2iso': (
-            common_reserved_words,
-            iso_additional_reserved_words,
-        ),
-
-        # Reserved words for ISO Modula-2
-        'm2r10': (
-            common_reserved_words,
-            m2r10_additional_reserved_words,
-        ),
-
-        # Reserved words for Objective Modula-2
-        'objm2': (
-            common_reserved_words,
-            m2r10_additional_reserved_words,
-            objm2_additional_reserved_words,
-        ),
-
-        # Reserved words for Aglet Modula-2 Extensions
-        'm2iso+aglet': (
-            common_reserved_words,
-            iso_additional_reserved_words,
-            aglet_additional_reserved_words,
-        ),
-
-        # Reserved words for GNU Modula-2 Extensions
-        'm2pim+gm2': (
-            common_reserved_words,
-            pim_additional_reserved_words,
-            gm2_additional_reserved_words,
-        ),
-
-        # Reserved words for p1 Modula-2 Extensions
-        'm2iso+p1': (
-            common_reserved_words,
-            iso_additional_reserved_words,
-            p1_additional_reserved_words,
-        ),
-
-        # Reserved words for XDS Modula-2 Extensions
-        'm2iso+xds': (
-            common_reserved_words,
-            iso_additional_reserved_words,
-            xds_additional_reserved_words,
-        ),
-    }
-
-    # Builtins Database
-    builtins_db = {
-        # Builtins for unknown dialect
-        'unknown': (
-            common_builtins,
-            pim_additional_builtins,
-            iso_additional_builtins,
-            m2r10_additional_builtins,
-        ),
-
-        # Builtins for PIM Modula-2
-        'm2pim': (
-            common_builtins,
-            pim_additional_builtins,
-        ),
-
-        # Builtins for ISO Modula-2
-        'm2iso': (
-            common_builtins,
-            iso_additional_builtins,
-        ),
-
-        # Builtins for ISO Modula-2
-        'm2r10': (
-            common_builtins,
-            m2r10_additional_builtins,
-        ),
-
-        # Builtins for Objective Modula-2
-        'objm2': (
-            common_builtins,
-            m2r10_additional_builtins,
-            objm2_additional_builtins,
-        ),
-
-        # Builtins for Aglet Modula-2 Extensions
-        'm2iso+aglet': (
-            common_builtins,
-            iso_additional_builtins,
-            aglet_additional_builtins,
-        ),
-
-        # Builtins for GNU Modula-2 Extensions
-        'm2pim+gm2': (
-            common_builtins,
-            pim_additional_builtins,
-            gm2_additional_builtins,
-        ),
-
-        # Builtins for p1 Modula-2 Extensions
-        'm2iso+p1': (
-            common_builtins,
-            iso_additional_builtins,
-            p1_additional_builtins,
-        ),
-
-        # Builtins for XDS Modula-2 Extensions
-        'm2iso+xds': (
-            common_builtins,
-            iso_additional_builtins,
-            xds_additional_builtins,
-        ),
-    }
-
-    # Pseudo-Module Builtins Database
-    pseudo_builtins_db = {
-        # Builtins for unknown dialect
-        'unknown': (
-            common_pseudo_builtins,
-            pim_additional_pseudo_builtins,
-            iso_additional_pseudo_builtins,
-            m2r10_additional_pseudo_builtins,
-        ),
-
-        # Builtins for PIM Modula-2
-        'm2pim': (
-            common_pseudo_builtins,
-            pim_additional_pseudo_builtins,
-        ),
-
-        # Builtins for ISO Modula-2
-        'm2iso': (
-            common_pseudo_builtins,
-            iso_additional_pseudo_builtins,
-        ),
-
-        # Builtins for ISO Modula-2
-        'm2r10': (
-            common_pseudo_builtins,
-            m2r10_additional_pseudo_builtins,
-        ),
-
-        # Builtins for Objective Modula-2
-        'objm2': (
-            common_pseudo_builtins,
-            m2r10_additional_pseudo_builtins,
-            objm2_additional_pseudo_builtins,
-        ),
-
-        # Builtins for Aglet Modula-2 Extensions
-        'm2iso+aglet': (
-            common_pseudo_builtins,
-            iso_additional_pseudo_builtins,
-            aglet_additional_pseudo_builtins,
-        ),
-
-        # Builtins for GNU Modula-2 Extensions
-        'm2pim+gm2': (
-            common_pseudo_builtins,
-            pim_additional_pseudo_builtins,
-            gm2_additional_pseudo_builtins,
-        ),
-
-        # Builtins for p1 Modula-2 Extensions
-        'm2iso+p1': (
-            common_pseudo_builtins,
-            iso_additional_pseudo_builtins,
-            p1_additional_pseudo_builtins,
-        ),
-
-        # Builtins for XDS Modula-2 Extensions
-        'm2iso+xds': (
-            common_pseudo_builtins,
-            iso_additional_pseudo_builtins,
-            xds_additional_pseudo_builtins,
-        ),
-    }
-
-    # Standard Library ADTs Database
-    stdlib_adts_db = {
-        # Empty entry for unknown dialect
-        'unknown': (
-            # LEAVE THIS EMPTY
-        ),
-        # Standard Library ADTs for PIM Modula-2
-        'm2pim': (
-            # No first class library types
-        ),
-
-        # Standard Library ADTs for ISO Modula-2
-        'm2iso': (
-            # No first class library types
-        ),
-
-        # Standard Library ADTs for Modula-2 R10
-        'm2r10': (
-            m2r10_stdlib_adt_identifiers,
-        ),
-
-        # Standard Library ADTs for Objective Modula-2
-        'objm2': (
-            m2r10_stdlib_adt_identifiers,
-        ),
-
-        # Standard Library ADTs for Aglet Modula-2
-        'm2iso+aglet': (
-            # No first class library types
-        ),
-
-        # Standard Library ADTs for GNU Modula-2
-        'm2pim+gm2': (
-            # No first class library types
-        ),
-
-        # Standard Library ADTs for p1 Modula-2
-        'm2iso+p1': (
-            # No first class library types
-        ),
-
-        # Standard Library ADTs for XDS Modula-2
-        'm2iso+xds': (
-            # No first class library types
-        ),
-    }
-
-    # Standard Library Modules Database
-    stdlib_modules_db = {
-        # Empty entry for unknown dialect
-        'unknown': (
-            # LEAVE THIS EMPTY
-        ),
-        # Standard Library Modules for PIM Modula-2
-        'm2pim': (
-            pim_stdlib_module_identifiers,
-        ),
-
-        # Standard Library Modules for ISO Modula-2
-        'm2iso': (
-            iso_stdlib_module_identifiers,
-        ),
-
-        # Standard Library Modules for Modula-2 R10
-        'm2r10': (
-            m2r10_stdlib_blueprint_identifiers,
-            m2r10_stdlib_module_identifiers,
-            m2r10_stdlib_adt_identifiers,
-        ),
-
-        # Standard Library Modules for Objective Modula-2
-        'objm2': (
-            m2r10_stdlib_blueprint_identifiers,
-            m2r10_stdlib_module_identifiers,
-        ),
-
-        # Standard Library Modules for Aglet Modula-2
-        'm2iso+aglet': (
-            iso_stdlib_module_identifiers,
-        ),
-
-        # Standard Library Modules for GNU Modula-2
-        'm2pim+gm2': (
-            pim_stdlib_module_identifiers,
-        ),
-
-        # Standard Library Modules for p1 Modula-2
-        'm2iso+p1': (
-            iso_stdlib_module_identifiers,
-        ),
-
-        # Standard Library Modules for XDS Modula-2
-        'm2iso+xds': (
-            iso_stdlib_module_identifiers,
-        ),
-    }
-
-    # Standard Library Types Database
-    stdlib_types_db = {
-        # Empty entry for unknown dialect
-        'unknown': (
-            # LEAVE THIS EMPTY
-        ),
-        # Standard Library Types for PIM Modula-2
-        'm2pim': (
-            pim_stdlib_type_identifiers,
-        ),
-
-        # Standard Library Types for ISO Modula-2
-        'm2iso': (
-            iso_stdlib_type_identifiers,
-        ),
-
-        # Standard Library Types for Modula-2 R10
-        'm2r10': (
-            m2r10_stdlib_type_identifiers,
-        ),
-
-        # Standard Library Types for Objective Modula-2
-        'objm2': (
-            m2r10_stdlib_type_identifiers,
-        ),
-
-        # Standard Library Types for Aglet Modula-2
-        'm2iso+aglet': (
-            iso_stdlib_type_identifiers,
-        ),
-
-        # Standard Library Types for GNU Modula-2
-        'm2pim+gm2': (
-            pim_stdlib_type_identifiers,
-        ),
-
-        # Standard Library Types for p1 Modula-2
-        'm2iso+p1': (
-            iso_stdlib_type_identifiers,
-        ),
-
-        # Standard Library Types for XDS Modula-2
-        'm2iso+xds': (
-            iso_stdlib_type_identifiers,
-        ),
-    }
-
-    # Standard Library Procedures Database
-    stdlib_procedures_db = {
-        # Empty entry for unknown dialect
-        'unknown': (
-            # LEAVE THIS EMPTY
-        ),
-        # Standard Library Procedures for PIM Modula-2
-        'm2pim': (
-            pim_stdlib_proc_identifiers,
-        ),
-
-        # Standard Library Procedures for ISO Modula-2
-        'm2iso': (
-            iso_stdlib_proc_identifiers,
-        ),
-
-        # Standard Library Procedures for Modula-2 R10
-        'm2r10': (
-            m2r10_stdlib_proc_identifiers,
-        ),
-
-        # Standard Library Procedures for Objective Modula-2
-        'objm2': (
-            m2r10_stdlib_proc_identifiers,
-        ),
-
-        # Standard Library Procedures for Aglet Modula-2
-        'm2iso+aglet': (
-            iso_stdlib_proc_identifiers,
-        ),
-
-        # Standard Library Procedures for GNU Modula-2
-        'm2pim+gm2': (
-            pim_stdlib_proc_identifiers,
-        ),
-
-        # Standard Library Procedures for p1 Modula-2
-        'm2iso+p1': (
-            iso_stdlib_proc_identifiers,
-        ),
-
-        # Standard Library Procedures for XDS Modula-2
-        'm2iso+xds': (
-            iso_stdlib_proc_identifiers,
-        ),
-    }
-
-    # Standard Library Variables Database
-    stdlib_variables_db = {
-        # Empty entry for unknown dialect
-        'unknown': (
-            # LEAVE THIS EMPTY
-        ),
-        # Standard Library Variables for PIM Modula-2
-        'm2pim': (
-            pim_stdlib_var_identifiers,
-        ),
-
-        # Standard Library Variables for ISO Modula-2
-        'm2iso': (
-            iso_stdlib_var_identifiers,
-        ),
-
-        # Standard Library Variables for Modula-2 R10
-        'm2r10': (
-            m2r10_stdlib_var_identifiers,
-        ),
-
-        # Standard Library Variables for Objective Modula-2
-        'objm2': (
-            m2r10_stdlib_var_identifiers,
-        ),
-
-        # Standard Library Variables for Aglet Modula-2
-        'm2iso+aglet': (
-            iso_stdlib_var_identifiers,
-        ),
-
-        # Standard Library Variables for GNU Modula-2
-        'm2pim+gm2': (
-            pim_stdlib_var_identifiers,
-        ),
-
-        # Standard Library Variables for p1 Modula-2
-        'm2iso+p1': (
-            iso_stdlib_var_identifiers,
-        ),
-
-        # Standard Library Variables for XDS Modula-2
-        'm2iso+xds': (
-            iso_stdlib_var_identifiers,
-        ),
-    }
-
-    # Standard Library Constants Database
-    stdlib_constants_db = {
-        # Empty entry for unknown dialect
-        'unknown': (
-            # LEAVE THIS EMPTY
-        ),
-        # Standard Library Constants for PIM Modula-2
-        'm2pim': (
-            pim_stdlib_const_identifiers,
-        ),
-
-        # Standard Library Constants for ISO Modula-2
-        'm2iso': (
-            iso_stdlib_const_identifiers,
-        ),
-
-        # Standard Library Constants for Modula-2 R10
-        'm2r10': (
-            m2r10_stdlib_const_identifiers,
-        ),
-
-        # Standard Library Constants for Objective Modula-2
-        'objm2': (
-            m2r10_stdlib_const_identifiers,
-        ),
-
-        # Standard Library Constants for Aglet Modula-2
-        'm2iso+aglet': (
-            iso_stdlib_const_identifiers,
-        ),
-
-        # Standard Library Constants for GNU Modula-2
-        'm2pim+gm2': (
-            pim_stdlib_const_identifiers,
-        ),
-
-        # Standard Library Constants for p1 Modula-2
-        'm2iso+p1': (
-            iso_stdlib_const_identifiers,
-        ),
-
-        # Standard Library Constants for XDS Modula-2
-        'm2iso+xds': (
-            iso_stdlib_const_identifiers,
-        ),
-    }
-
-#   M e t h o d s
-
-    # initialise a lexer instance
-    def __init__(self, **options):
-        #
-        # check dialect options
-        #
-        dialects = get_list_opt(options, 'dialect', [])
-        #
-        for dialect_option in dialects:
-            if dialect_option in self.dialects[1:-1]:
-                # valid dialect option found
-                self.set_dialect(dialect_option)
-                break
-        #
-        # Fallback Mode (DEFAULT)
-        else:
-            # no valid dialect option
-            self.set_dialect('unknown')
-        #
-        self.dialect_set_by_tag = False
-        #
-        # check style options
-        #
-        styles = get_list_opt(options, 'style', [])
-        #
-        # use lowercase mode for Algol style
-        if 'algol' in styles or 'algol_nu' in styles:
-            self.algol_publication_mode = True
-        else:
-            self.algol_publication_mode = False
-        #
-        # Check option flags
-        #
-        self.treat_stdlib_adts_as_builtins = get_bool_opt(
-            options, 'treat_stdlib_adts_as_builtins', True)
-        #
-        # call superclass initialiser
-        RegexLexer.__init__(self, **options)
-
-    # Set lexer to a specified dialect
-    def set_dialect(self, dialect_id):
-        #
-        # if __debug__:
-        #    print 'entered set_dialect with arg: ', dialect_id
-        #
-        # check dialect name against known dialects
-        if dialect_id not in self.dialects:
-            dialect = 'unknown'  # default
-        else:
-            dialect = dialect_id
-        #
-        # compose lexemes to reject set
-        lexemes_to_reject_set = set()
-        # add each list of reject lexemes for this dialect
-        for list in self.lexemes_to_reject_db[dialect]:
-            lexemes_to_reject_set.update(set(list))
-        #
-        # compose reserved words set
-        reswords_set = set()
-        # add each list of reserved words for this dialect
-        for list in self.reserved_words_db[dialect]:
-            reswords_set.update(set(list))
-        #
-        # compose builtins set
-        builtins_set = set()
-        # add each list of builtins for this dialect excluding reserved words
-        for list in self.builtins_db[dialect]:
-            builtins_set.update(set(list).difference(reswords_set))
-        #
-        # compose pseudo-builtins set
-        pseudo_builtins_set = set()
-        # add each list of builtins for this dialect excluding reserved words
-        for list in self.pseudo_builtins_db[dialect]:
-            pseudo_builtins_set.update(set(list).difference(reswords_set))
-        #
-        # compose ADTs set
-        adts_set = set()
-        # add each list of ADTs for this dialect excluding reserved words
-        for list in self.stdlib_adts_db[dialect]:
-            adts_set.update(set(list).difference(reswords_set))
-        #
-        # compose modules set
-        modules_set = set()
-        # add each list of builtins for this dialect excluding builtins
-        for list in self.stdlib_modules_db[dialect]:
-            modules_set.update(set(list).difference(builtins_set))
-        #
-        # compose types set
-        types_set = set()
-        # add each list of types for this dialect excluding builtins
-        for list in self.stdlib_types_db[dialect]:
-            types_set.update(set(list).difference(builtins_set))
-        #
-        # compose procedures set
-        procedures_set = set()
-        # add each list of procedures for this dialect excluding builtins
-        for list in self.stdlib_procedures_db[dialect]:
-            procedures_set.update(set(list).difference(builtins_set))
-        #
-        # compose variables set
-        variables_set = set()
-        # add each list of variables for this dialect excluding builtins
-        for list in self.stdlib_variables_db[dialect]:
-            variables_set.update(set(list).difference(builtins_set))
-        #
-        # compose constants set
-        constants_set = set()
-        # add each list of constants for this dialect excluding builtins
-        for list in self.stdlib_constants_db[dialect]:
-            constants_set.update(set(list).difference(builtins_set))
-        #
-        # update lexer state
-        self.dialect = dialect
-        self.lexemes_to_reject = lexemes_to_reject_set
-        self.reserved_words = reswords_set
-        self.builtins = builtins_set
-        self.pseudo_builtins = pseudo_builtins_set
-        self.adts = adts_set
-        self.modules = modules_set
-        self.types = types_set
-        self.procedures = procedures_set
-        self.variables = variables_set
-        self.constants = constants_set
-        #
-        # if __debug__:
-        #    print 'exiting set_dialect'
-        #    print ' self.dialect: ', self.dialect
-        #    print ' self.lexemes_to_reject: ', self.lexemes_to_reject
-        #    print ' self.reserved_words: ', self.reserved_words
-        #    print ' self.builtins: ', self.builtins
-        #    print ' self.pseudo_builtins: ', self.pseudo_builtins
-        #    print ' self.adts: ', self.adts
-        #    print ' self.modules: ', self.modules
-        #    print ' self.types: ', self.types
-        #    print ' self.procedures: ', self.procedures
-        #    print ' self.variables: ', self.variables
-        #    print ' self.types: ', self.types
-        #    print ' self.constants: ', self.constants
-
-    # Extracts a dialect name from a dialect tag comment string  and checks
-    # the extracted name against known dialects.  If a match is found,  the
-    # matching name is returned, otherwise dialect id 'unknown' is returned
-    def get_dialect_from_dialect_tag(self, dialect_tag):
-        #
-        # if __debug__:
-        #    print 'entered get_dialect_from_dialect_tag with arg: ', dialect_tag
-        #
-        # constants
-        left_tag_delim = '(*!'
-        right_tag_delim = '*)'
-        left_tag_delim_len = len(left_tag_delim)
-        right_tag_delim_len = len(right_tag_delim)
-        indicator_start = left_tag_delim_len
-        indicator_end = -(right_tag_delim_len)
-        #
-        # check comment string for dialect indicator
-        if len(dialect_tag) > (left_tag_delim_len + right_tag_delim_len) \
-           and dialect_tag.startswith(left_tag_delim) \
-           and dialect_tag.endswith(right_tag_delim):
-            #
-            # if __debug__:
-            #    print 'dialect tag found'
-            #
-            # extract dialect indicator
-            indicator = dialect_tag[indicator_start:indicator_end]
-            #
-            # if __debug__:
-            #    print 'extracted: ', indicator
-            #
-            # check against known dialects
-            for index in range(1, len(self.dialects)):
-                #
-                # if __debug__:
-                #    print 'dialects[', index, ']: ', self.dialects[index]
-                #
-                if indicator == self.dialects[index]:
-                    #
-                    # if __debug__:
-                    #    print 'matching dialect found'
-                    #
-                    # indicator matches known dialect
-                    return indicator
-            else:
-                # indicator does not match any dialect
-                return 'unknown'  # default
-        else:
-            # invalid indicator string
-            return 'unknown'  # default
-
-    # intercept the token stream, modify token attributes and return them
-    def get_tokens_unprocessed(self, text):
-        for index, token, value in RegexLexer.get_tokens_unprocessed(self, text):
-            #
-            # check for dialect tag if dialect has not been set by tag
-            if not self.dialect_set_by_tag and token == Comment.Special:
-                indicated_dialect = self.get_dialect_from_dialect_tag(value)
-                if indicated_dialect != 'unknown':
-                    # token is a dialect indicator
-                    # reset reserved words and builtins
-                    self.set_dialect(indicated_dialect)
-                    self.dialect_set_by_tag = True
-            #
-            # check for reserved words, predefined and stdlib identifiers
-            if token is Name:
-                if value in self.reserved_words:
-                    token = Keyword.Reserved
-                    if self.algol_publication_mode:
-                        value = value.lower()
-                #
-                elif value in self.builtins:
-                    token = Name.Builtin
-                    if self.algol_publication_mode:
-                        value = value.lower()
-                #
-                elif value in self.pseudo_builtins:
-                    token = Name.Builtin.Pseudo
-                    if self.algol_publication_mode:
-                        value = value.lower()
-                #
-                elif value in self.adts:
-                    if not self.treat_stdlib_adts_as_builtins:
-                        token = Name.Namespace
-                    else:
-                        token = Name.Builtin.Pseudo
-                        if self.algol_publication_mode:
-                            value = value.lower()
-                #
-                elif value in self.modules:
-                    token = Name.Namespace
-                #
-                elif value in self.types:
-                    token = Name.Class
-                #
-                elif value in self.procedures:
-                    token = Name.Function
-                #
-                elif value in self.variables:
-                    token = Name.Variable
-                #
-                elif value in self.constants:
-                    token = Name.Constant
-            #
-            elif token in Number:
-                #
-                # mark prefix number literals as error for PIM and ISO dialects
-                if self.dialect not in ('unknown', 'm2r10', 'objm2'):
-                    if "'" in value or value[0:2] in ('0b', '0x', '0u'):
-                        token = Error
-                #
-                elif self.dialect in ('m2r10', 'objm2'):
-                    # mark base-8 number literals as errors for M2 R10 and ObjM2
-                    if token is Number.Oct:
-                        token = Error
-                    # mark suffix base-16 literals as errors for M2 R10 and ObjM2
-                    elif token is Number.Hex and 'H' in value:
-                        token = Error
-                    # mark real numbers with E as errors for M2 R10 and ObjM2
-                    elif token is Number.Float and 'E' in value:
-                        token = Error
-            #
-            elif token in Comment:
-                #
-                # mark single line comment as error for PIM and ISO dialects
-                if token is Comment.Single:
-                    if self.dialect not in ('unknown', 'm2r10', 'objm2'):
-                        token = Error
-                #
-                if token is Comment.Preproc:
-                    # mark ISO pragma as error for PIM dialects
-                    if value.startswith('<*') and \
-                       self.dialect.startswith('m2pim'):
-                        token = Error
-                    # mark PIM pragma as comment for other dialects
-                    elif value.startswith('(*$') and \
-                            self.dialect != 'unknown' and \
-                            not self.dialect.startswith('m2pim'):
-                        token = Comment.Multiline
-            #
-            else:  # token is neither Name nor Comment
-                #
-                # mark lexemes matching the dialect's error token set as errors
-                if value in self.lexemes_to_reject:
-                    token = Error
-                #
-                # substitute lexemes when in Algol mode
-                if self.algol_publication_mode:
-                    if value == '#':
-                        value = '≠'
-                    elif value == '<=':
-                        value = '≤'
-                    elif value == '>=':
-                        value = '≥'
-                    elif value == '==':
-                        value = '≡'
-                    elif value == '*.':
-                        value = '•'
-
-            # return result
-            yield index, token, value
-
-    def analyse_text(text):
-        """It's Pascal-like, but does not use FUNCTION -- uses PROCEDURE
-        instead."""
-
-        # Check if this looks like Pascal, if not, bail out early
-        if not ('(*' in text and '*)' in text and ':=' in text):
-            return
-
-        result = 0
-        # Procedure is in Modula2
-        if re.search(r'\bPROCEDURE\b', text):
-            result += 0.6
-
-        # FUNCTION is only valid in Pascal, but not in Modula2
-        if re.search(r'\bFUNCTION\b', text):
-            result = 0.0
-
-        return result
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/mojo.py b/venv/lib/python3.10/site-packages/pygments/lexers/mojo.py
deleted file mode 100644
index 4df18c4f9c94520ab852ba2f66bb50518eb64be7..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/mojo.py
+++ /dev/null
@@ -1,707 +0,0 @@
-"""
-    pygments.lexers.mojo
-    ~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Mojo and related languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import keyword
-
-from pygments import unistring as uni
-from pygments.lexer import (
-    RegexLexer,
-    bygroups,
-    combined,
-    default,
-    include,
-    this,
-    using,
-    words,
-)
-from pygments.token import (
-    Comment,
-    # Error,
-    Keyword,
-    Name,
-    Number,
-    Operator,
-    Punctuation,
-    String,
-    Text,
-    Whitespace,
-)
-from pygments.util import shebang_matches
-
-__all__ = ["MojoLexer"]
-
-
-class MojoLexer(RegexLexer):
-    """
-    For Mojo source code (version 24.2.1).
-    """
-
-    name = "Mojo"
-    url = "https://docs.modular.com/mojo/"
-    aliases = ["mojo", "🔥"]
-    filenames = [
-        "*.mojo",
-        "*.🔥",
-    ]
-    mimetypes = [
-        "text/x-mojo",
-        "application/x-mojo",
-    ]
-    version_added = "2.18"
-
-    uni_name = f"[{uni.xid_start}][{uni.xid_continue}]*"
-
-    def innerstring_rules(ttype):
-        return [
-            # the old style '%s' % (...) string formatting (still valid in Py3)
-            (
-                r"%(\(\w+\))?[-#0 +]*([0-9]+|[*])?(\.([0-9]+|[*]))?"
-                "[hlL]?[E-GXc-giorsaux%]",
-                String.Interpol,
-            ),
-            # the new style '{}'.format(...) string formatting
-            (
-                r"\{"
-                r"((\w+)((\.\w+)|(\[[^\]]+\]))*)?"  # field name
-                r"(\![sra])?"  # conversion
-                r"(\:(.?[<>=\^])?[-+ ]?#?0?(\d+)?,?(\.\d+)?[E-GXb-gnosx%]?)?"
-                r"\}",
-                String.Interpol,
-            ),
-            # backslashes, quotes and formatting signs must be parsed one at a time
-            (r'[^\\\'"%{\n]+', ttype),
-            (r'[\'"\\]', ttype),
-            # unhandled string formatting sign
-            (r"%|(\{{1,2})", ttype),
-            # newlines are an error (use "nl" state)
-        ]
-
-    def fstring_rules(ttype):
-        return [
-            # Assuming that a '}' is the closing brace after format specifier.
-            # Sadly, this means that we won't detect syntax error. But it's
-            # more important to parse correct syntax correctly, than to
-            # highlight invalid syntax.
-            (r"\}", String.Interpol),
-            (r"\{", String.Interpol, "expr-inside-fstring"),
-            # backslashes, quotes and formatting signs must be parsed one at a time
-            (r'[^\\\'"{}\n]+', ttype),
-            (r'[\'"\\]', ttype),
-            # newlines are an error (use "nl" state)
-        ]
-
-    tokens = {
-        "root": [
-            (r"\s+", Whitespace),
-            (
-                r'^(\s*)([rRuUbB]{,2})("""(?:.|\n)*?""")',
-                bygroups(Whitespace, String.Affix, String.Doc),
-            ),
-            (
-                r"^(\s*)([rRuUbB]{,2})('''(?:.|\n)*?''')",
-                bygroups(Whitespace, String.Affix, String.Doc),
-            ),
-            (r"\A#!.+$", Comment.Hashbang),
-            (r"#.*$", Comment.Single),
-            (r"\\\n", Whitespace),
-            (r"\\", Whitespace),
-            include("keywords"),
-            include("soft-keywords"),
-            # In the original PR, all the below here used ((?:\s|\\\s)+) to
-            # designate whitespace, but I can't find any example of this being
-            # needed in the example file, so we're replacing it with `\s+`.
-            (
-                r"(alias)(\s+)",
-                bygroups(Keyword, Whitespace),
-                "varname",  # TODO varname the right fit?
-            ),
-            (r"(var)(\s+)", bygroups(Keyword, Whitespace), "varname"),
-            (r"(def)(\s+)", bygroups(Keyword, Whitespace), "funcname"),
-            (r"(fn)(\s+)", bygroups(Keyword, Whitespace), "funcname"),
-            (
-                r"(class)(\s+)",
-                bygroups(Keyword, Whitespace),
-                "classname",
-            ),  # not implemented yet
-            (r"(struct)(\s+)", bygroups(Keyword, Whitespace), "structname"),
-            (r"(trait)(\s+)", bygroups(Keyword, Whitespace), "structname"),
-            (r"(from)(\s+)", bygroups(Keyword.Namespace, Whitespace), "fromimport"),
-            (r"(import)(\s+)", bygroups(Keyword.Namespace, Whitespace), "import"),
-            include("expr"),
-        ],
-        "expr": [
-            # raw f-strings
-            (
-                '(?i)(rf|fr)(""")',
-                bygroups(String.Affix, String.Double),
-                combined("rfstringescape", "tdqf"),
-            ),
-            (
-                "(?i)(rf|fr)(''')",
-                bygroups(String.Affix, String.Single),
-                combined("rfstringescape", "tsqf"),
-            ),
-            (
-                '(?i)(rf|fr)(")',
-                bygroups(String.Affix, String.Double),
-                combined("rfstringescape", "dqf"),
-            ),
-            (
-                "(?i)(rf|fr)(')",
-                bygroups(String.Affix, String.Single),
-                combined("rfstringescape", "sqf"),
-            ),
-            # non-raw f-strings
-            (
-                '([fF])(""")',
-                bygroups(String.Affix, String.Double),
-                combined("fstringescape", "tdqf"),
-            ),
-            (
-                "([fF])(''')",
-                bygroups(String.Affix, String.Single),
-                combined("fstringescape", "tsqf"),
-            ),
-            (
-                '([fF])(")',
-                bygroups(String.Affix, String.Double),
-                combined("fstringescape", "dqf"),
-            ),
-            (
-                "([fF])(')",
-                bygroups(String.Affix, String.Single),
-                combined("fstringescape", "sqf"),
-            ),
-            # raw bytes and strings
-            ('(?i)(rb|br|r)(""")', bygroups(String.Affix, String.Double), "tdqs"),
-            ("(?i)(rb|br|r)(''')", bygroups(String.Affix, String.Single), "tsqs"),
-            ('(?i)(rb|br|r)(")', bygroups(String.Affix, String.Double), "dqs"),
-            ("(?i)(rb|br|r)(')", bygroups(String.Affix, String.Single), "sqs"),
-            # non-raw strings
-            (
-                '([uU]?)(""")',
-                bygroups(String.Affix, String.Double),
-                combined("stringescape", "tdqs"),
-            ),
-            (
-                "([uU]?)(''')",
-                bygroups(String.Affix, String.Single),
-                combined("stringescape", "tsqs"),
-            ),
-            (
-                '([uU]?)(")',
-                bygroups(String.Affix, String.Double),
-                combined("stringescape", "dqs"),
-            ),
-            (
-                "([uU]?)(')",
-                bygroups(String.Affix, String.Single),
-                combined("stringescape", "sqs"),
-            ),
-            # non-raw bytes
-            (
-                '([bB])(""")',
-                bygroups(String.Affix, String.Double),
-                combined("bytesescape", "tdqs"),
-            ),
-            (
-                "([bB])(''')",
-                bygroups(String.Affix, String.Single),
-                combined("bytesescape", "tsqs"),
-            ),
-            (
-                '([bB])(")',
-                bygroups(String.Affix, String.Double),
-                combined("bytesescape", "dqs"),
-            ),
-            (
-                "([bB])(')",
-                bygroups(String.Affix, String.Single),
-                combined("bytesescape", "sqs"),
-            ),
-            (r"[^\S\n]+", Text),
-            include("numbers"),
-            (r"!=|==|<<|>>|:=|[-~+/*%=<>&^|.]", Operator),
-            (r"([]{}:\(\),;[])+", Punctuation),
-            (r"(in|is|and|or|not)\b", Operator.Word),
-            include("expr-keywords"),
-            include("builtins"),
-            include("magicfuncs"),
-            include("magicvars"),
-            include("name"),
-        ],
-        "expr-inside-fstring": [
-            (r"[{([]", Punctuation, "expr-inside-fstring-inner"),
-            # without format specifier
-            (
-                r"(=\s*)?"  # debug (https://bugs.python.org/issue36817)
-                r"(\![sraf])?"  # conversion
-                r"\}",
-                String.Interpol,
-                "#pop",
-            ),
-            # with format specifier
-            # we'll catch the remaining '}' in the outer scope
-            (
-                r"(=\s*)?"  # debug (https://bugs.python.org/issue36817)
-                r"(\![sraf])?"  # conversion
-                r":",
-                String.Interpol,
-                "#pop",
-            ),
-            (r"\s+", Whitespace),  # allow new lines
-            include("expr"),
-        ],
-        "expr-inside-fstring-inner": [
-            (r"[{([]", Punctuation, "expr-inside-fstring-inner"),
-            (r"[])}]", Punctuation, "#pop"),
-            (r"\s+", Whitespace),  # allow new lines
-            include("expr"),
-        ],
-        "expr-keywords": [
-            # Based on https://docs.python.org/3/reference/expressions.html
-            (
-                words(
-                    (
-                        "async for",  # TODO https://docs.modular.com/mojo/roadmap#no-async-for-or-async-with
-                        "async with",  # TODO https://docs.modular.com/mojo/roadmap#no-async-for-or-async-with
-                        "await",
-                        "else",
-                        "for",
-                        "if",
-                        "lambda",
-                        "yield",
-                        "yield from",
-                    ),
-                    suffix=r"\b",
-                ),
-                Keyword,
-            ),
-            (words(("True", "False", "None"), suffix=r"\b"), Keyword.Constant),
-        ],
-        "keywords": [
-            (
-                words(
-                    (
-                        "assert",
-                        "async",
-                        "await",
-                        "borrowed",
-                        "break",
-                        "continue",
-                        "del",
-                        "elif",
-                        "else",
-                        "except",
-                        "finally",
-                        "for",
-                        "global",
-                        "if",
-                        "lambda",
-                        "pass",
-                        "raise",
-                        "nonlocal",
-                        "return",
-                        "try",
-                        "while",
-                        "yield",
-                        "yield from",
-                        "as",
-                        "with",
-                    ),
-                    suffix=r"\b",
-                ),
-                Keyword,
-            ),
-            (words(("True", "False", "None"), suffix=r"\b"), Keyword.Constant),
-        ],
-        "soft-keywords": [
-            # `match`, `case` and `_` soft keywords
-            (
-                r"(^[ \t]*)"  # at beginning of line + possible indentation
-                r"(match|case)\b"  # a possible keyword
-                r"(?![ \t]*(?:"  # not followed by...
-                r"[:,;=^&|@~)\]}]|(?:" +  # characters and keywords that mean this isn't
-                # pattern matching (but None/True/False is ok)
-                r"|".join(k for k in keyword.kwlist if k[0].islower())
-                + r")\b))",
-                bygroups(Whitespace, Keyword),
-                "soft-keywords-inner",
-            ),
-        ],
-        "soft-keywords-inner": [
-            # optional `_` keyword
-            (r"(\s+)([^\n_]*)(_\b)", bygroups(Whitespace, using(this), Keyword)),
-            default("#pop"),
-        ],
-        "builtins": [
-            (
-                words(
-                    (
-                        "__import__",
-                        "abs",
-                        "aiter",
-                        "all",
-                        "any",
-                        "bin",
-                        "bool",
-                        "bytearray",
-                        "breakpoint",
-                        "bytes",
-                        "callable",
-                        "chr",
-                        "classmethod",
-                        "compile",
-                        "complex",
-                        "delattr",
-                        "dict",
-                        "dir",
-                        "divmod",
-                        "enumerate",
-                        "eval",
-                        "filter",
-                        "float",
-                        "format",
-                        "frozenset",
-                        "getattr",
-                        "globals",
-                        "hasattr",
-                        "hash",
-                        "hex",
-                        "id",
-                        "input",
-                        "int",
-                        "isinstance",
-                        "issubclass",
-                        "iter",
-                        "len",
-                        "list",
-                        "locals",
-                        "map",
-                        "max",
-                        "memoryview",
-                        "min",
-                        "next",
-                        "object",
-                        "oct",
-                        "open",
-                        "ord",
-                        "pow",
-                        "print",
-                        "property",
-                        "range",
-                        "repr",
-                        "reversed",
-                        "round",
-                        "set",
-                        "setattr",
-                        "slice",
-                        "sorted",
-                        "staticmethod",
-                        "str",
-                        "sum",
-                        "super",
-                        "tuple",
-                        "type",
-                        "vars",
-                        "zip",
-                        # Mojo builtin types: https://docs.modular.com/mojo/stdlib/builtin/
-                        "AnyType",
-                        "Coroutine",
-                        "DType",
-                        "Error",
-                        "Int",
-                        "List",
-                        "ListLiteral",
-                        "Scalar",
-                        "Int8",
-                        "UInt8",
-                        "Int16",
-                        "UInt16",
-                        "Int32",
-                        "UInt32",
-                        "Int64",
-                        "UInt64",
-                        "BFloat16",
-                        "Float16",
-                        "Float32",
-                        "Float64",
-                        "SIMD",
-                        "String",
-                        "Tensor",
-                        "Tuple",
-                        "Movable",
-                        "Copyable",
-                        "CollectionElement",
-                    ),
-                    prefix=r"(?>',
-    # Binary augmented
-    '+=', '-=', '*=', '/=', '%=', '**=', '&=', '|=', '^=', '<<=', '>>=',
-    # Comparison
-    '==', '!=', '<', '<=', '>', '>=', '<=>',
-    # Patterns and assignment
-    ':=', '?', '=~', '!~', '=>',
-    # Calls and sends
-    '.', '<-', '->',
-]
-_escape_pattern = (
-    r'(?:\\x[0-9a-fA-F]{2}|\\u[0-9a-fA-F]{4}|\\U[0-9a-fA-F]{8}|'
-    r'\\["\'\\bftnr])')
-# _char = _escape_chars + [('.', String.Char)]
-_identifier = r'[_a-zA-Z]\w*'
-
-_constants = [
-    # Void constants
-    'null',
-    # Bool constants
-    'false', 'true',
-    # Double constants
-    'Infinity', 'NaN',
-    # Special objects
-    'M', 'Ref', 'throw', 'traceln',
-]
-
-_guards = [
-    'Any', 'Binding', 'Bool', 'Bytes', 'Char', 'DeepFrozen', 'Double',
-    'Empty', 'Int', 'List', 'Map', 'Near', 'NullOk', 'Same', 'Selfless',
-    'Set', 'Str', 'SubrangeGuard', 'Transparent', 'Void',
-]
-
-_safeScope = [
-    '_accumulateList', '_accumulateMap', '_auditedBy', '_bind',
-    '_booleanFlow', '_comparer', '_equalizer', '_iterForever', '_loop',
-    '_makeBytes', '_makeDouble', '_makeFinalSlot', '_makeInt', '_makeList',
-    '_makeMap', '_makeMessageDesc', '_makeOrderedSpace', '_makeParamDesc',
-    '_makeProtocolDesc', '_makeSourceSpan', '_makeString', '_makeVarSlot',
-    '_makeVerbFacet', '_mapExtract', '_matchSame', '_quasiMatcher',
-    '_slotToBinding', '_splitList', '_suchThat', '_switchFailed',
-    '_validateFor', 'b__quasiParser', 'eval', 'import', 'm__quasiParser',
-    'makeBrandPair', 'makeLazySlot', 'safeScope', 'simple__quasiParser',
-]
-
-
-class MonteLexer(RegexLexer):
-    """
-    Lexer for the Monte programming language.
-    """
-    name = 'Monte'
-    url = 'https://monte.readthedocs.io/'
-    aliases = ['monte']
-    filenames = ['*.mt']
-    version_added = '2.2'
-
-    tokens = {
-        'root': [
-            # Comments
-            (r'#[^\n]*\n', Comment),
-
-            # Docstrings
-            # Apologies for the non-greedy matcher here.
-            (r'/\*\*.*?\*/', String.Doc),
-
-            # `var` declarations
-            (r'\bvar\b', Keyword.Declaration, 'var'),
-
-            # `interface` declarations
-            (r'\binterface\b', Keyword.Declaration, 'interface'),
-
-            # method declarations
-            (words(_methods, prefix='\\b', suffix='\\b'),
-             Keyword, 'method'),
-
-            # All other declarations
-            (words(_declarations, prefix='\\b', suffix='\\b'),
-             Keyword.Declaration),
-
-            # Keywords
-            (words(_keywords, prefix='\\b', suffix='\\b'), Keyword),
-
-            # Literals
-            ('[+-]?0x[_0-9a-fA-F]+', Number.Hex),
-            (r'[+-]?[_0-9]+\.[_0-9]*([eE][+-]?[_0-9]+)?', Number.Float),
-            ('[+-]?[_0-9]+', Number.Integer),
-            ("'", String.Double, 'char'),
-            ('"', String.Double, 'string'),
-
-            # Quasiliterals
-            ('`', String.Backtick, 'ql'),
-
-            # Operators
-            (words(_operators), Operator),
-
-            # Verb operators
-            (_identifier + '=', Operator.Word),
-
-            # Safe scope constants
-            (words(_constants, prefix='\\b', suffix='\\b'),
-             Keyword.Pseudo),
-
-            # Safe scope guards
-            (words(_guards, prefix='\\b', suffix='\\b'), Keyword.Type),
-
-            # All other safe scope names
-            (words(_safeScope, prefix='\\b', suffix='\\b'),
-             Name.Builtin),
-
-            # Identifiers
-            (_identifier, Name),
-
-            # Punctuation
-            (r'\(|\)|\{|\}|\[|\]|:|,', Punctuation),
-
-            # Whitespace
-            (' +', Whitespace),
-
-            # Definite lexer errors
-            ('=', Error),
-        ],
-        'char': [
-            # It is definitely an error to have a char of width == 0.
-            ("'", Error, 'root'),
-            (_escape_pattern, String.Escape, 'charEnd'),
-            ('.', String.Char, 'charEnd'),
-        ],
-        'charEnd': [
-            ("'", String.Char, '#pop:2'),
-            # It is definitely an error to have a char of width > 1.
-            ('.', Error),
-        ],
-        # The state of things coming into an interface.
-        'interface': [
-            (' +', Whitespace),
-            (_identifier, Name.Class, '#pop'),
-            include('root'),
-        ],
-        # The state of things coming into a method.
-        'method': [
-            (' +', Whitespace),
-            (_identifier, Name.Function, '#pop'),
-            include('root'),
-        ],
-        'string': [
-            ('"', String.Double, 'root'),
-            (_escape_pattern, String.Escape),
-            (r'\n', String.Double),
-            ('.', String.Double),
-        ],
-        'ql': [
-            ('`', String.Backtick, 'root'),
-            (r'\$' + _escape_pattern, String.Escape),
-            (r'\$\$', String.Escape),
-            (r'@@', String.Escape),
-            (r'\$\{', String.Interpol, 'qlNest'),
-            (r'@\{', String.Interpol, 'qlNest'),
-            (r'\$' + _identifier, Name),
-            ('@' + _identifier, Name),
-            ('.', String.Backtick),
-        ],
-        'qlNest': [
-            (r'\}', String.Interpol, '#pop'),
-            include('root'),
-        ],
-        # The state of things immediately following `var`.
-        'var': [
-            (' +', Whitespace),
-            (_identifier, Name.Variable, '#pop'),
-            include('root'),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/mosel.py b/venv/lib/python3.10/site-packages/pygments/lexers/mosel.py
deleted file mode 100644
index 426c9a145261af10f25c2f6e5f16fdccf665b065..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/mosel.py
+++ /dev/null
@@ -1,447 +0,0 @@
-"""
-    pygments.lexers.mosel
-    ~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for the mosel language.
-    http://www.fico.com/en/products/fico-xpress-optimization
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, words
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation
-
-__all__ = ['MoselLexer']
-
-FUNCTIONS = (
-    # core functions
-    '_',
-    'abs',
-    'arctan',
-    'asproc',
-    'assert',
-    'bitflip',
-    'bitneg',
-    'bitset',
-    'bitshift',
-    'bittest',
-    'bitval',
-    'ceil',
-    'cos',
-    'create',
-    'currentdate',
-    'currenttime',
-    'cutelt',
-    'cutfirst',
-    'cuthead',
-    'cutlast',
-    'cuttail',
-    'datablock',
-    'delcell',
-    'exists',
-    'exit',
-    'exp',
-    'exportprob',
-    'fclose',
-    'fflush',
-    'finalize',
-    'findfirst',
-    'findlast',
-    'floor',
-    'fopen',
-    'fselect',
-    'fskipline',
-    'fwrite',
-    'fwrite_',
-    'fwriteln',
-    'fwriteln_',
-    'getact',
-    'getcoeff',
-    'getcoeffs',
-    'getdual',
-    'getelt',
-    'getfid',
-    'getfirst',
-    'getfname',
-    'gethead',
-    'getlast',
-    'getobjval',
-    'getparam',
-    'getrcost',
-    'getreadcnt',
-    'getreverse',
-    'getsize',
-    'getslack',
-    'getsol',
-    'gettail',
-    'gettype',
-    'getvars',
-    'isdynamic',
-    'iseof',
-    'isfinite',
-    'ishidden',
-    'isinf',
-    'isnan',
-    'isodd',
-    'ln',
-    'localsetparam',
-    'log',
-    'makesos1',
-    'makesos2',
-    'maxlist',
-    'memoryuse',
-    'minlist',
-    'newmuid',
-    'publish',
-    'random',
-    'read',
-    'readln',
-    'reset',
-    'restoreparam',
-    'reverse',
-    'round',
-    'setcoeff',
-    'sethidden',
-    'setioerr',
-    'setmatherr',
-    'setname',
-    'setparam',
-    'setrandseed',
-    'setrange',
-    'settype',
-    'sin',
-    'splithead',
-    'splittail',
-    'sqrt',
-    'strfmt',
-    'substr',
-    'timestamp',
-    'unpublish',
-    'versionnum',
-    'versionstr',
-    'write',
-    'write_',
-    'writeln',
-    'writeln_',
-
-    # mosel exam mmxprs | sed -n -e "s/ [pf][a-z]* \([a-zA-Z0-9_]*\).*/'\1',/p" | sort -u
-    'addcut',
-    'addcuts',
-    'addmipsol',
-    'basisstability',
-    'calcsolinfo',
-    'clearmipdir',
-    'clearmodcut',
-    'command',
-    'copysoltoinit',
-    'crossoverlpsol',
-    'defdelayedrows',
-    'defsecurevecs',
-    'delcuts',
-    'dropcuts',
-    'estimatemarginals',
-    'fixglobal',
-    'flushmsgq',
-    'getbstat',
-    'getcnlist',
-    'getcplist',
-    'getdualray',
-    'getiis',
-    'getiissense',
-    'getiistype',
-    'getinfcause',
-    'getinfeas',
-    'getlb',
-    'getlct',
-    'getleft',
-    'getloadedlinctrs',
-    'getloadedmpvars',
-    'getname',
-    'getprimalray',
-    'getprobstat',
-    'getrange',
-    'getright',
-    'getsensrng',
-    'getsize',
-    'getsol',
-    'gettype',
-    'getub',
-    'getvars',
-    'gety',
-    'hasfeature',
-    'implies',
-    'indicator',
-    'initglobal',
-    'ishidden',
-    'isiisvalid',
-    'isintegral',
-    'loadbasis',
-    'loadcuts',
-    'loadlpsol',
-    'loadmipsol',
-    'loadprob',
-    'maximise',
-    'maximize',
-    'minimise',
-    'minimize',
-    'postsolve',
-    'readbasis',
-    'readdirs',
-    'readsol',
-    'refinemipsol',
-    'rejectintsol',
-    'repairinfeas',
-    'repairinfeas_deprec',
-    'resetbasis',
-    'resetiis',
-    'resetsol',
-    'savebasis',
-    'savemipsol',
-    'savesol',
-    'savestate',
-    'selectsol',
-    'setarchconsistency',
-    'setbstat',
-    'setcallback',
-    'setcbcutoff',
-    'setgndata',
-    'sethidden',
-    'setlb',
-    'setmipdir',
-    'setmodcut',
-    'setsol',
-    'setub',
-    'setucbdata',
-    'stopoptimise',
-    'stopoptimize',
-    'storecut',
-    'storecuts',
-    'unloadprob',
-    'uselastbarsol',
-    'writebasis',
-    'writedirs',
-    'writeprob',
-    'writesol',
-    'xor',
-    'xprs_addctr',
-    'xprs_addindic',
-
-    # mosel exam mmsystem | sed -n -e "s/ [pf][a-z]* \([a-zA-Z0-9_]*\).*/'\1',/p" | sort -u
-    'addmonths',
-    'copytext',
-    'cuttext',
-    'deltext',
-    'endswith',
-    'erase',
-    'expandpath',
-    'fcopy',
-    'fdelete',
-    'findfiles',
-    'findtext',
-    'fmove',
-    'formattext',
-    'getasnumber',
-    'getchar',
-    'getcwd',
-    'getdate',
-    'getday',
-    'getdaynum',
-    'getdays',
-    'getdirsep',
-    'getdsoparam',
-    'getendparse',
-    'getenv',
-    'getfsize',
-    'getfstat',
-    'getftime',
-    'gethour',
-    'getminute',
-    'getmonth',
-    'getmsec',
-    'getoserrmsg',
-    'getoserror',
-    'getpathsep',
-    'getqtype',
-    'getsecond',
-    'getsepchar',
-    'getsize',
-    'getstart',
-    'getsucc',
-    'getsysinfo',
-    'getsysstat',
-    'gettime',
-    'gettmpdir',
-    'gettrim',
-    'getweekday',
-    'getyear',
-    'inserttext',
-    'isvalid',
-    'jointext',
-    'makedir',
-    'makepath',
-    'newtar',
-    'newzip',
-    'nextfield',
-    'openpipe',
-    'parseextn',
-    'parseint',
-    'parsereal',
-    'parsetext',
-    'pastetext',
-    'pathmatch',
-    'pathsplit',
-    'qsort',
-    'quote',
-    'readtextline',
-    'regmatch',
-    'regreplace',
-    'removedir',
-    'removefiles',
-    'setchar',
-    'setdate',
-    'setday',
-    'setdsoparam',
-    'setendparse',
-    'setenv',
-    'sethour',
-    'setminute',
-    'setmonth',
-    'setmsec',
-    'setoserror',
-    'setqtype',
-    'setsecond',
-    'setsepchar',
-    'setstart',
-    'setsucc',
-    'settime',
-    'settrim',
-    'setyear',
-    'sleep',
-    'splittext',
-    'startswith',
-    'system',
-    'tarlist',
-    'textfmt',
-    'tolower',
-    'toupper',
-    'trim',
-    'untar',
-    'unzip',
-    'ziplist',
-
-    # mosel exam mmjobs | sed -n -e "s/ [pf][a-z]* \([a-zA-Z0-9_]*\).*/'\1',/p" | sort -u
-    'canceltimer',
-    'clearaliases',
-    'compile',
-    'connect',
-    'detach',
-    'disconnect',
-    'dropnextevent',
-    'findxsrvs',
-    'getaliases',
-    'getannidents',
-    'getannotations',
-    'getbanner',
-    'getclass',
-    'getdsoprop',
-    'getdsopropnum',
-    'getexitcode',
-    'getfromgid',
-    'getfromid',
-    'getfromuid',
-    'getgid',
-    'gethostalias',
-    'getid',
-    'getmodprop',
-    'getmodpropnum',
-    'getnextevent',
-    'getnode',
-    'getrmtid',
-    'getstatus',
-    'getsysinfo',
-    'gettimer',
-    'getuid',
-    'getvalue',
-    'isqueueempty',
-    'load',
-    'nullevent',
-    'peeknextevent',
-    'resetmodpar',
-    'run',
-    'send',
-    'setcontrol',
-    'setdefstream',
-    'setgid',
-    'sethostalias',
-    'setmodpar',
-    'settimer',
-    'setuid',
-    'setworkdir',
-    'stop',
-    'unload',
-    'wait',
-    'waitexpired',
-    'waitfor',
-    'waitforend',
-)
-
-
-class MoselLexer(RegexLexer):
-    """
-    For the Mosel optimization language.
-    """
-    name = 'Mosel'
-    aliases = ['mosel']
-    filenames = ['*.mos']
-    url = 'https://www.fico.com/fico-xpress-optimization/docs/latest/mosel/mosel_lang/dhtml/moselreflang.html'
-    version_added = '2.6'
-
-    tokens = {
-        'root': [
-            (r'\n', Text),
-            (r'\s+', Text.Whitespace),
-            (r'!.*?\n', Comment.Single),
-            (r'\(!(.|\n)*?!\)', Comment.Multiline),
-            (words((
-                'and', 'as', 'break', 'case', 'count', 'declarations', 'do',
-                'dynamic', 'elif', 'else', 'end-', 'end', 'evaluation', 'false',
-                'forall', 'forward', 'from', 'function', 'hashmap', 'if',
-                'imports', 'include', 'initialisations', 'initializations', 'inter',
-                'max', 'min', 'model', 'namespace', 'next', 'not', 'nsgroup',
-                'nssearch', 'of', 'options', 'or', 'package', 'parameters',
-                'procedure', 'public', 'prod', 'record', 'repeat', 'requirements',
-                'return', 'sum', 'then', 'to', 'true', 'union', 'until', 'uses',
-                'version', 'while', 'with'), prefix=r'\b', suffix=r'\b'),
-             Keyword.Builtin),
-            (words((
-                'range', 'array', 'set', 'list', 'mpvar', 'mpproblem', 'linctr',
-                'nlctr', 'integer', 'string', 'real', 'boolean', 'text', 'time',
-                'date', 'datetime', 'returned', 'Model', 'Mosel', 'counter',
-                'xmldoc', 'is_sos1', 'is_sos2', 'is_integer', 'is_binary',
-                'is_continuous', 'is_free', 'is_semcont', 'is_semint',
-                'is_partint'), prefix=r'\b', suffix=r'\b'),
-             Keyword.Type),
-            (r'(\+|\-|\*|/|=|<=|>=|\||\^|<|>|<>|\.\.|\.|:=|::|:|in|mod|div)',
-             Operator),
-            (r'[()\[\]{},;]+', Punctuation),
-            (words(FUNCTIONS,  prefix=r'\b', suffix=r'\b'), Name.Function),
-            (r'(\d+\.(?!\.)\d*|\.(?!.)\d+)([eE][+-]?\d+)?', Number.Float),
-            (r'\d+([eE][+-]?\d+)?', Number.Integer),
-            (r'[+-]?Infinity', Number.Integer),
-            (r'0[xX][0-9a-fA-F]+', Number),
-            (r'"', String.Double, 'double_quote'),
-            (r'\'', String.Single, 'single_quote'),
-            (r'(\w+|(\.(?!\.)))', Text),
-        ],
-        'single_quote': [
-            (r'\'', String.Single, '#pop'),
-            (r'[^\']+', String.Single),
-        ],
-        'double_quote': [
-            (r'(\\"|\\[0-7]{1,3}\D|\\[abfnrtv]|\\\\)', String.Escape),
-            (r'\"', String.Double, '#pop'),
-            (r'[^"\\]+', String.Double),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/ncl.py b/venv/lib/python3.10/site-packages/pygments/lexers/ncl.py
deleted file mode 100644
index d2f476087bf6e6e92328542c52d0c609d1972845..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/ncl.py
+++ /dev/null
@@ -1,894 +0,0 @@
-"""
-    pygments.lexers.ncl
-    ~~~~~~~~~~~~~~~~~~~
-
-    Lexers for NCAR Command Language.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, include, words
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation
-
-__all__ = ['NCLLexer']
-
-
-class NCLLexer(RegexLexer):
-    """
-    Lexer for NCL code.
-    """
-    name = 'NCL'
-    aliases = ['ncl']
-    filenames = ['*.ncl']
-    mimetypes = ['text/ncl']
-    url = 'https://www.ncl.ucar.edu'
-    version_added = '2.2'
-
-    flags = re.MULTILINE
-
-    tokens = {
-        'root': [
-            (r';.*\n', Comment),
-            include('strings'),
-            include('core'),
-            (r'[a-zA-Z_]\w*', Name),
-            include('nums'),
-            (r'[\s]+', Text),
-        ],
-        'core': [
-            # Statements
-            (words((
-                'begin', 'break', 'continue', 'create', 'defaultapp', 'do',
-                'else', 'end', 'external', 'exit', 'True', 'False', 'file', 'function',
-                'getvalues', 'graphic', 'group', 'if', 'list', 'load', 'local',
-                'new', '_Missing', 'Missing', 'noparent', 'procedure',
-                'quit', 'QUIT', 'Quit', 'record', 'return', 'setvalues', 'stop',
-                'then', 'while'), prefix=r'\b', suffix=r'\s*\b'),
-             Keyword),
-
-            # Data Types
-            (words((
-                'ubyte', 'uint', 'uint64', 'ulong', 'string', 'byte',
-                'character', 'double', 'float', 'integer', 'int64', 'logical',
-                'long', 'short', 'ushort', 'enumeric', 'numeric', 'snumeric'),
-                prefix=r'\b', suffix=r'\s*\b'),
-             Keyword.Type),
-
-            # Operators
-            (r'[\%^*+\-/<>]', Operator),
-
-            # punctuation:
-            (r'[\[\]():@$!&|.,\\{}]', Punctuation),
-            (r'[=:]', Punctuation),
-
-            # Intrinsics
-            (words((
-                'abs', 'acos', 'addfile', 'addfiles', 'all', 'angmom_atm', 'any',
-                'area_conserve_remap', 'area_hi2lores', 'area_poly_sphere',
-                'asciiread', 'asciiwrite', 'asin', 'atan', 'atan2', 'attsetvalues',
-                'avg', 'betainc', 'bin_avg', 'bin_sum', 'bw_bandpass_filter',
-                'cancor', 'cbinread', 'cbinwrite', 'cd_calendar', 'cd_inv_calendar',
-                'cdfbin_p', 'cdfbin_pr', 'cdfbin_s', 'cdfbin_xn', 'cdfchi_p',
-                'cdfchi_x', 'cdfgam_p', 'cdfgam_x', 'cdfnor_p', 'cdfnor_x',
-                'cdft_p', 'cdft_t', 'ceil', 'center_finite_diff',
-                'center_finite_diff_n', 'cfftb', 'cfftf', 'cfftf_frq_reorder',
-                'charactertodouble', 'charactertofloat', 'charactertointeger',
-                'charactertolong', 'charactertoshort', 'charactertostring',
-                'chartodouble', 'chartofloat', 'chartoint', 'chartointeger',
-                'chartolong', 'chartoshort', 'chartostring', 'chiinv', 'clear',
-                'color_index_to_rgba', 'conform', 'conform_dims', 'cos', 'cosh',
-                'count_unique_values', 'covcorm', 'covcorm_xy', 'craybinnumrec',
-                'craybinrecread', 'create_graphic', 'csa1', 'csa1d', 'csa1s',
-                'csa1x', 'csa1xd', 'csa1xs', 'csa2', 'csa2d', 'csa2l', 'csa2ld',
-                'csa2ls', 'csa2lx', 'csa2lxd', 'csa2lxs', 'csa2s', 'csa2x',
-                'csa2xd', 'csa2xs', 'csa3', 'csa3d', 'csa3l', 'csa3ld', 'csa3ls',
-                'csa3lx', 'csa3lxd', 'csa3lxs', 'csa3s', 'csa3x', 'csa3xd',
-                'csa3xs', 'csc2s', 'csgetp', 'css2c', 'cssetp', 'cssgrid', 'csstri',
-                'csvoro', 'cumsum', 'cz2ccm', 'datatondc', 'day_of_week',
-                'day_of_year', 'days_in_month', 'default_fillvalue', 'delete',
-                'depth_to_pres', 'destroy', 'determinant', 'dewtemp_trh',
-                'dgeevx_lapack', 'dim_acumrun_n', 'dim_avg', 'dim_avg_n',
-                'dim_avg_wgt', 'dim_avg_wgt_n', 'dim_cumsum', 'dim_cumsum_n',
-                'dim_gamfit_n', 'dim_gbits', 'dim_max', 'dim_max_n', 'dim_median',
-                'dim_median_n', 'dim_min', 'dim_min_n', 'dim_num', 'dim_num_n',
-                'dim_numrun_n', 'dim_pqsort', 'dim_pqsort_n', 'dim_product',
-                'dim_product_n', 'dim_rmsd', 'dim_rmsd_n', 'dim_rmvmean',
-                'dim_rmvmean_n', 'dim_rmvmed', 'dim_rmvmed_n', 'dim_spi_n',
-                'dim_standardize', 'dim_standardize_n', 'dim_stat4', 'dim_stat4_n',
-                'dim_stddev', 'dim_stddev_n', 'dim_sum', 'dim_sum_n', 'dim_sum_wgt',
-                'dim_sum_wgt_n', 'dim_variance', 'dim_variance_n', 'dimsizes',
-                'doubletobyte', 'doubletochar', 'doubletocharacter',
-                'doubletofloat', 'doubletoint', 'doubletointeger', 'doubletolong',
-                'doubletoshort', 'dpres_hybrid_ccm', 'dpres_plevel', 'draw',
-                'draw_color_palette', 'dsgetp', 'dsgrid2', 'dsgrid2d', 'dsgrid2s',
-                'dsgrid3', 'dsgrid3d', 'dsgrid3s', 'dspnt2', 'dspnt2d', 'dspnt2s',
-                'dspnt3', 'dspnt3d', 'dspnt3s', 'dssetp', 'dtrend', 'dtrend_msg',
-                'dtrend_msg_n', 'dtrend_n', 'dtrend_quadratic',
-                'dtrend_quadratic_msg_n', 'dv2uvf', 'dv2uvg', 'dz_height',
-                'echo_off', 'echo_on', 'eof2data', 'eof_varimax', 'eofcor',
-                'eofcor_pcmsg', 'eofcor_ts', 'eofcov', 'eofcov_pcmsg', 'eofcov_ts',
-                'eofunc', 'eofunc_ts', 'eofunc_varimax', 'equiv_sample_size', 'erf',
-                'erfc', 'esacr', 'esacv', 'esccr', 'esccv', 'escorc', 'escorc_n',
-                'escovc', 'exit', 'exp', 'exp_tapersh', 'exp_tapersh_wgts',
-                'exp_tapershC', 'ezfftb', 'ezfftb_n', 'ezfftf', 'ezfftf_n',
-                'f2fosh', 'f2foshv', 'f2fsh', 'f2fshv', 'f2gsh', 'f2gshv', 'fabs',
-                'fbindirread', 'fbindirwrite', 'fbinnumrec', 'fbinread',
-                'fbinrecread', 'fbinrecwrite', 'fbinwrite', 'fft2db', 'fft2df',
-                'fftshift', 'fileattdef', 'filechunkdimdef', 'filedimdef',
-                'fileexists', 'filegrpdef', 'filevarattdef', 'filevarchunkdef',
-                'filevarcompressleveldef', 'filevardef', 'filevardimsizes',
-                'filwgts_lancos', 'filwgts_lanczos', 'filwgts_normal',
-                'floattobyte', 'floattochar', 'floattocharacter', 'floattoint',
-                'floattointeger', 'floattolong', 'floattoshort', 'floor',
-                'fluxEddy', 'fo2fsh', 'fo2fshv', 'fourier_info', 'frame', 'fspan',
-                'ftcurv', 'ftcurvd', 'ftcurvi', 'ftcurvp', 'ftcurvpi', 'ftcurvps',
-                'ftcurvs', 'ftest', 'ftgetp', 'ftkurv', 'ftkurvd', 'ftkurvp',
-                'ftkurvpd', 'ftsetp', 'ftsurf', 'g2fsh', 'g2fshv', 'g2gsh',
-                'g2gshv', 'gamma', 'gammainc', 'gaus', 'gaus_lobat',
-                'gaus_lobat_wgt', 'gc_aangle', 'gc_clkwise', 'gc_dangle',
-                'gc_inout', 'gc_latlon', 'gc_onarc', 'gc_pnt2gc', 'gc_qarea',
-                'gc_tarea', 'generate_2d_array', 'get_color_index',
-                'get_color_rgba', 'get_cpu_time', 'get_isolines', 'get_ncl_version',
-                'get_script_name', 'get_script_prefix_name', 'get_sphere_radius',
-                'get_unique_values', 'getbitsone', 'getenv', 'getfiledimsizes',
-                'getfilegrpnames', 'getfilepath', 'getfilevaratts',
-                'getfilevarchunkdimsizes', 'getfilevardims', 'getfilevardimsizes',
-                'getfilevarnames', 'getfilevartypes', 'getvaratts', 'getvardims',
-                'gradsf', 'gradsg', 'greg2jul', 'grid2triple', 'hlsrgb', 'hsvrgb',
-                'hydro', 'hyi2hyo', 'idsfft', 'igradsf', 'igradsg', 'ilapsf',
-                'ilapsg', 'ilapvf', 'ilapvg', 'ind', 'ind_resolve', 'int2p',
-                'int2p_n', 'integertobyte', 'integertochar', 'integertocharacter',
-                'integertoshort', 'inttobyte', 'inttochar', 'inttoshort',
-                'inverse_matrix', 'isatt', 'isbigendian', 'isbyte', 'ischar',
-                'iscoord', 'isdefined', 'isdim', 'isdimnamed', 'isdouble',
-                'isenumeric', 'isfile', 'isfilepresent', 'isfilevar',
-                'isfilevaratt', 'isfilevarcoord', 'isfilevardim', 'isfloat',
-                'isfunc', 'isgraphic', 'isint', 'isint64', 'isinteger',
-                'isleapyear', 'islogical', 'islong', 'ismissing', 'isnan_ieee',
-                'isnumeric', 'ispan', 'isproc', 'isshort', 'issnumeric', 'isstring',
-                'isubyte', 'isuint', 'isuint64', 'isulong', 'isunlimited',
-                'isunsigned', 'isushort', 'isvar', 'jul2greg', 'kmeans_as136',
-                'kolsm2_n', 'kron_product', 'lapsf', 'lapsg', 'lapvf', 'lapvg',
-                'latlon2utm', 'lclvl', 'lderuvf', 'lderuvg', 'linint1', 'linint1_n',
-                'linint2', 'linint2_points', 'linmsg', 'linmsg_n', 'linrood_latwgt',
-                'linrood_wgt', 'list_files', 'list_filevars', 'list_hlus',
-                'list_procfuncs', 'list_vars', 'ListAppend', 'ListCount',
-                'ListGetType', 'ListIndex', 'ListIndexFromName', 'ListPop',
-                'ListPush', 'ListSetType', 'loadscript', 'local_max', 'local_min',
-                'log', 'log10', 'longtobyte', 'longtochar', 'longtocharacter',
-                'longtoint', 'longtointeger', 'longtoshort', 'lspoly', 'lspoly_n',
-                'mask', 'max', 'maxind', 'min', 'minind', 'mixed_layer_depth',
-                'mixhum_ptd', 'mixhum_ptrh', 'mjo_cross_coh2pha',
-                'mjo_cross_segment', 'moc_globe_atl', 'monthday', 'natgrid',
-                'natgridd', 'natgrids', 'ncargpath', 'ncargversion', 'ndctodata',
-                'ndtooned', 'new', 'NewList', 'ngezlogo', 'nggcog', 'nggetp',
-                'nglogo', 'ngsetp', 'NhlAddAnnotation', 'NhlAddData',
-                'NhlAddOverlay', 'NhlAddPrimitive', 'NhlAppGetDefaultParentId',
-                'NhlChangeWorkstation', 'NhlClassName', 'NhlClearWorkstation',
-                'NhlDataPolygon', 'NhlDataPolyline', 'NhlDataPolymarker',
-                'NhlDataToNDC', 'NhlDestroy', 'NhlDraw', 'NhlFrame', 'NhlFreeColor',
-                'NhlGetBB', 'NhlGetClassResources', 'NhlGetErrorObjectId',
-                'NhlGetNamedColorIndex', 'NhlGetParentId',
-                'NhlGetParentWorkstation', 'NhlGetWorkspaceObjectId',
-                'NhlIsAllocatedColor', 'NhlIsApp', 'NhlIsDataComm', 'NhlIsDataItem',
-                'NhlIsDataSpec', 'NhlIsTransform', 'NhlIsView', 'NhlIsWorkstation',
-                'NhlName', 'NhlNDCPolygon', 'NhlNDCPolyline', 'NhlNDCPolymarker',
-                'NhlNDCToData', 'NhlNewColor', 'NhlNewDashPattern', 'NhlNewMarker',
-                'NhlPalGetDefined', 'NhlRemoveAnnotation', 'NhlRemoveData',
-                'NhlRemoveOverlay', 'NhlRemovePrimitive', 'NhlSetColor',
-                'NhlSetDashPattern', 'NhlSetMarker', 'NhlUpdateData',
-                'NhlUpdateWorkstation', 'nice_mnmxintvl', 'nngetaspectd',
-                'nngetaspects', 'nngetp', 'nngetsloped', 'nngetslopes', 'nngetwts',
-                'nngetwtsd', 'nnpnt', 'nnpntd', 'nnpntend', 'nnpntendd',
-                'nnpntinit', 'nnpntinitd', 'nnpntinits', 'nnpnts', 'nnsetp', 'num',
-                'obj_anal_ic', 'omega_ccm', 'onedtond', 'overlay', 'paleo_outline',
-                'pdfxy_bin', 'poisson_grid_fill', 'pop_remap', 'potmp_insitu_ocn',
-                'prcwater_dp', 'pres2hybrid', 'pres_hybrid_ccm', 'pres_sigma',
-                'print', 'print_table', 'printFileVarSummary', 'printVarSummary',
-                'product', 'pslec', 'pslhor', 'pslhyp', 'qsort', 'rand',
-                'random_chi', 'random_gamma', 'random_normal', 'random_setallseed',
-                'random_uniform', 'rcm2points', 'rcm2rgrid', 'rdsstoi',
-                'read_colormap_file', 'reg_multlin', 'regcoef', 'regCoef_n',
-                'regline', 'relhum', 'replace_ieeenan', 'reshape', 'reshape_ind',
-                'rgba_to_color_index', 'rgbhls', 'rgbhsv', 'rgbyiq', 'rgrid2rcm',
-                'rhomb_trunc', 'rip_cape_2d', 'rip_cape_3d', 'round', 'rtest',
-                'runave', 'runave_n', 'set_default_fillvalue', 'set_sphere_radius',
-                'setfileoption', 'sfvp2uvf', 'sfvp2uvg', 'shaec', 'shagc',
-                'shgetnp', 'shgetp', 'shgrid', 'shorttobyte', 'shorttochar',
-                'shorttocharacter', 'show_ascii', 'shsec', 'shsetp', 'shsgc',
-                'shsgc_R42', 'sigma2hybrid', 'simpeq', 'simpne', 'sin',
-                'sindex_yrmo', 'sinh', 'sizeof', 'sleep', 'smth9', 'snindex_yrmo',
-                'solve_linsys', 'span_color_indexes', 'span_color_rgba',
-                'sparse_matrix_mult', 'spcorr', 'spcorr_n', 'specx_anal',
-                'specxy_anal', 'spei', 'sprintf', 'sprinti', 'sqrt', 'sqsort',
-                'srand', 'stat2', 'stat4', 'stat_medrng', 'stat_trim',
-                'status_exit', 'stdatmus_p2tdz', 'stdatmus_z2tdp', 'stddev',
-                'str_capital', 'str_concat', 'str_fields_count', 'str_get_cols',
-                'str_get_dq', 'str_get_field', 'str_get_nl', 'str_get_sq',
-                'str_get_tab', 'str_index_of_substr', 'str_insert', 'str_is_blank',
-                'str_join', 'str_left_strip', 'str_lower', 'str_match',
-                'str_match_ic', 'str_match_ic_regex', 'str_match_ind',
-                'str_match_ind_ic', 'str_match_ind_ic_regex', 'str_match_ind_regex',
-                'str_match_regex', 'str_right_strip', 'str_split',
-                'str_split_by_length', 'str_split_csv', 'str_squeeze', 'str_strip',
-                'str_sub_str', 'str_switch', 'str_upper', 'stringtochar',
-                'stringtocharacter', 'stringtodouble', 'stringtofloat',
-                'stringtoint', 'stringtointeger', 'stringtolong', 'stringtoshort',
-                'strlen', 'student_t', 'sum', 'svd_lapack', 'svdcov', 'svdcov_sv',
-                'svdstd', 'svdstd_sv', 'system', 'systemfunc', 'tan', 'tanh',
-                'taper', 'taper_n', 'tdclrs', 'tdctri', 'tdcudp', 'tdcurv',
-                'tddtri', 'tdez2d', 'tdez3d', 'tdgetp', 'tdgrds', 'tdgrid',
-                'tdgtrs', 'tdinit', 'tditri', 'tdlbla', 'tdlblp', 'tdlbls',
-                'tdline', 'tdlndp', 'tdlnpa', 'tdlpdp', 'tdmtri', 'tdotri',
-                'tdpara', 'tdplch', 'tdprpa', 'tdprpi', 'tdprpt', 'tdsetp',
-                'tdsort', 'tdstri', 'tdstrs', 'tdttri', 'thornthwaite', 'tobyte',
-                'tochar', 'todouble', 'tofloat', 'toint', 'toint64', 'tointeger',
-                'tolong', 'toshort', 'tosigned', 'tostring', 'tostring_with_format',
-                'totype', 'toubyte', 'touint', 'touint64', 'toulong', 'tounsigned',
-                'toushort', 'trend_manken', 'tri_trunc', 'triple2grid',
-                'triple2grid2d', 'trop_wmo', 'ttest', 'typeof', 'undef',
-                'unique_string', 'update', 'ushorttoint', 'ut_calendar',
-                'ut_inv_calendar', 'utm2latlon', 'uv2dv_cfd', 'uv2dvf', 'uv2dvg',
-                'uv2sfvpf', 'uv2sfvpg', 'uv2vr_cfd', 'uv2vrdvf', 'uv2vrdvg',
-                'uv2vrf', 'uv2vrg', 'v5d_close', 'v5d_create', 'v5d_setLowLev',
-                'v5d_setUnits', 'v5d_write', 'v5d_write_var', 'variance', 'vhaec',
-                'vhagc', 'vhsec', 'vhsgc', 'vibeta', 'vinth2p', 'vinth2p_ecmwf',
-                'vinth2p_ecmwf_nodes', 'vinth2p_nodes', 'vintp2p_ecmwf', 'vr2uvf',
-                'vr2uvg', 'vrdv2uvf', 'vrdv2uvg', 'wavelet', 'wavelet_default',
-                'weibull', 'wgt_area_smooth', 'wgt_areaave', 'wgt_areaave2',
-                'wgt_arearmse', 'wgt_arearmse2', 'wgt_areasum2', 'wgt_runave',
-                'wgt_runave_n', 'wgt_vert_avg_beta', 'wgt_volave', 'wgt_volave_ccm',
-                'wgt_volrmse', 'wgt_volrmse_ccm', 'where', 'wk_smooth121', 'wmbarb',
-                'wmbarbmap', 'wmdrft', 'wmgetp', 'wmlabs', 'wmsetp', 'wmstnm',
-                'wmvect', 'wmvectmap', 'wmvlbl', 'wrf_avo', 'wrf_cape_2d',
-                'wrf_cape_3d', 'wrf_dbz', 'wrf_eth', 'wrf_helicity', 'wrf_ij_to_ll',
-                'wrf_interp_1d', 'wrf_interp_2d_xy', 'wrf_interp_3d_z',
-                'wrf_latlon_to_ij', 'wrf_ll_to_ij', 'wrf_omega', 'wrf_pvo',
-                'wrf_rh', 'wrf_slp', 'wrf_smooth_2d', 'wrf_td', 'wrf_tk',
-                'wrf_updraft_helicity', 'wrf_uvmet', 'wrf_virtual_temp',
-                'wrf_wetbulb', 'wrf_wps_close_int', 'wrf_wps_open_int',
-                'wrf_wps_rddata_int', 'wrf_wps_rdhead_int', 'wrf_wps_read_int',
-                'wrf_wps_write_int', 'write_matrix', 'write_table', 'yiqrgb',
-                'z2geouv', 'zonal_mpsi', 'addfiles_GetVar', 'advect_variable',
-                'area_conserve_remap_Wrap', 'area_hi2lores_Wrap',
-                'array_append_record', 'assignFillValue', 'byte2flt',
-                'byte2flt_hdf', 'calcDayAnomTLL', 'calcMonAnomLLLT',
-                'calcMonAnomLLT', 'calcMonAnomTLL', 'calcMonAnomTLLL',
-                'calculate_monthly_values', 'cd_convert', 'changeCase',
-                'changeCaseChar', 'clmDayTLL', 'clmDayTLLL', 'clmMon2clmDay',
-                'clmMonLLLT', 'clmMonLLT', 'clmMonTLL', 'clmMonTLLL', 'closest_val',
-                'copy_VarAtts', 'copy_VarCoords', 'copy_VarCoords_1',
-                'copy_VarCoords_2', 'copy_VarMeta', 'copyatt', 'crossp3',
-                'cshstringtolist', 'cssgrid_Wrap', 'dble2flt', 'decimalPlaces',
-                'delete_VarAtts', 'dim_avg_n_Wrap', 'dim_avg_wgt_n_Wrap',
-                'dim_avg_wgt_Wrap', 'dim_avg_Wrap', 'dim_cumsum_n_Wrap',
-                'dim_cumsum_Wrap', 'dim_max_n_Wrap', 'dim_min_n_Wrap',
-                'dim_rmsd_n_Wrap', 'dim_rmsd_Wrap', 'dim_rmvmean_n_Wrap',
-                'dim_rmvmean_Wrap', 'dim_rmvmed_n_Wrap', 'dim_rmvmed_Wrap',
-                'dim_standardize_n_Wrap', 'dim_standardize_Wrap',
-                'dim_stddev_n_Wrap', 'dim_stddev_Wrap', 'dim_sum_n_Wrap',
-                'dim_sum_wgt_n_Wrap', 'dim_sum_wgt_Wrap', 'dim_sum_Wrap',
-                'dim_variance_n_Wrap', 'dim_variance_Wrap', 'dpres_plevel_Wrap',
-                'dtrend_leftdim', 'dv2uvF_Wrap', 'dv2uvG_Wrap', 'eof_north',
-                'eofcor_Wrap', 'eofcov_Wrap', 'eofunc_north', 'eofunc_ts_Wrap',
-                'eofunc_varimax_reorder', 'eofunc_varimax_Wrap', 'eofunc_Wrap',
-                'epsZero', 'f2fosh_Wrap', 'f2foshv_Wrap', 'f2fsh_Wrap',
-                'f2fshv_Wrap', 'f2gsh_Wrap', 'f2gshv_Wrap', 'fbindirSwap',
-                'fbinseqSwap1', 'fbinseqSwap2', 'flt2dble', 'flt2string',
-                'fo2fsh_Wrap', 'fo2fshv_Wrap', 'g2fsh_Wrap', 'g2fshv_Wrap',
-                'g2gsh_Wrap', 'g2gshv_Wrap', 'generate_resample_indices',
-                'generate_sample_indices', 'generate_unique_indices',
-                'genNormalDist', 'get1Dindex', 'get1Dindex_Collapse',
-                'get1Dindex_Exclude', 'get_file_suffix', 'GetFillColor',
-                'GetFillColorIndex', 'getFillValue', 'getind_latlon2d',
-                'getVarDimNames', 'getVarFillValue', 'grib_stime2itime',
-                'hyi2hyo_Wrap', 'ilapsF_Wrap', 'ilapsG_Wrap', 'ind_nearest_coord',
-                'indStrSubset', 'int2dble', 'int2flt', 'int2p_n_Wrap', 'int2p_Wrap',
-                'isMonotonic', 'isStrSubset', 'latGau', 'latGauWgt', 'latGlobeF',
-                'latGlobeFo', 'latRegWgt', 'linint1_n_Wrap', 'linint1_Wrap',
-                'linint2_points_Wrap', 'linint2_Wrap', 'local_max_1d',
-                'local_min_1d', 'lonFlip', 'lonGlobeF', 'lonGlobeFo', 'lonPivot',
-                'merge_levels_sfc', 'mod', 'month_to_annual',
-                'month_to_annual_weighted', 'month_to_season', 'month_to_season12',
-                'month_to_seasonN', 'monthly_total_to_daily_mean', 'nameDim',
-                'natgrid_Wrap', 'NewCosWeight', 'niceLatLon2D', 'NormCosWgtGlobe',
-                'numAsciiCol', 'numAsciiRow', 'numeric2int',
-                'obj_anal_ic_deprecated', 'obj_anal_ic_Wrap', 'omega_ccm_driver',
-                'omega_to_w', 'oneDtostring', 'pack_values', 'pattern_cor', 'pdfx',
-                'pdfxy', 'pdfxy_conform', 'pot_temp', 'pot_vort_hybrid',
-                'pot_vort_isobaric', 'pres2hybrid_Wrap', 'print_clock',
-                'printMinMax', 'quadroots', 'rcm2points_Wrap', 'rcm2rgrid_Wrap',
-                'readAsciiHead', 'readAsciiTable', 'reg_multlin_stats',
-                'region_ind', 'regline_stats', 'relhum_ttd', 'replaceSingleChar',
-                'RGBtoCmap', 'rgrid2rcm_Wrap', 'rho_mwjf', 'rm_single_dims',
-                'rmAnnCycle1D', 'rmInsufData', 'rmMonAnnCycLLLT', 'rmMonAnnCycLLT',
-                'rmMonAnnCycTLL', 'runave_n_Wrap', 'runave_Wrap', 'short2flt',
-                'short2flt_hdf', 'shsgc_R42_Wrap', 'sign_f90', 'sign_matlab',
-                'smth9_Wrap', 'smthClmDayTLL', 'smthClmDayTLLL', 'SqrtCosWeight',
-                'stat_dispersion', 'static_stability', 'stdMonLLLT', 'stdMonLLT',
-                'stdMonTLL', 'stdMonTLLL', 'symMinMaxPlt', 'table_attach_columns',
-                'table_attach_rows', 'time_to_newtime', 'transpose',
-                'triple2grid_Wrap', 'ut_convert', 'uv2dvF_Wrap', 'uv2dvG_Wrap',
-                'uv2vrF_Wrap', 'uv2vrG_Wrap', 'vr2uvF_Wrap', 'vr2uvG_Wrap',
-                'w_to_omega', 'wallClockElapseTime', 'wave_number_spc',
-                'wgt_areaave_Wrap', 'wgt_runave_leftdim', 'wgt_runave_n_Wrap',
-                'wgt_runave_Wrap', 'wgt_vertical_n', 'wind_component',
-                'wind_direction', 'yyyyddd_to_yyyymmdd', 'yyyymm_time',
-                'yyyymm_to_yyyyfrac', 'yyyymmdd_time', 'yyyymmdd_to_yyyyddd',
-                'yyyymmdd_to_yyyyfrac', 'yyyymmddhh_time', 'yyyymmddhh_to_yyyyfrac',
-                'zonal_mpsi_Wrap', 'zonalAve', 'calendar_decode2', 'cd_string',
-                'kf_filter', 'run_cor', 'time_axis_labels', 'ut_string',
-                'wrf_contour', 'wrf_map', 'wrf_map_overlay', 'wrf_map_overlays',
-                'wrf_map_resources', 'wrf_map_zoom', 'wrf_overlay', 'wrf_overlays',
-                'wrf_user_getvar', 'wrf_user_ij_to_ll', 'wrf_user_intrp2d',
-                'wrf_user_intrp3d', 'wrf_user_latlon_to_ij', 'wrf_user_list_times',
-                'wrf_user_ll_to_ij', 'wrf_user_unstagger', 'wrf_user_vert_interp',
-                'wrf_vector', 'gsn_add_annotation', 'gsn_add_polygon',
-                'gsn_add_polyline', 'gsn_add_polymarker',
-                'gsn_add_shapefile_polygons', 'gsn_add_shapefile_polylines',
-                'gsn_add_shapefile_polymarkers', 'gsn_add_text', 'gsn_attach_plots',
-                'gsn_blank_plot', 'gsn_contour', 'gsn_contour_map',
-                'gsn_contour_shade', 'gsn_coordinates', 'gsn_create_labelbar',
-                'gsn_create_legend', 'gsn_create_text',
-                'gsn_csm_attach_zonal_means', 'gsn_csm_blank_plot',
-                'gsn_csm_contour', 'gsn_csm_contour_map', 'gsn_csm_contour_map_ce',
-                'gsn_csm_contour_map_overlay', 'gsn_csm_contour_map_polar',
-                'gsn_csm_hov', 'gsn_csm_lat_time', 'gsn_csm_map', 'gsn_csm_map_ce',
-                'gsn_csm_map_polar', 'gsn_csm_pres_hgt',
-                'gsn_csm_pres_hgt_streamline', 'gsn_csm_pres_hgt_vector',
-                'gsn_csm_streamline', 'gsn_csm_streamline_contour_map',
-                'gsn_csm_streamline_contour_map_ce',
-                'gsn_csm_streamline_contour_map_polar', 'gsn_csm_streamline_map',
-                'gsn_csm_streamline_map_ce', 'gsn_csm_streamline_map_polar',
-                'gsn_csm_streamline_scalar', 'gsn_csm_streamline_scalar_map',
-                'gsn_csm_streamline_scalar_map_ce',
-                'gsn_csm_streamline_scalar_map_polar', 'gsn_csm_time_lat',
-                'gsn_csm_vector', 'gsn_csm_vector_map', 'gsn_csm_vector_map_ce',
-                'gsn_csm_vector_map_polar', 'gsn_csm_vector_scalar',
-                'gsn_csm_vector_scalar_map', 'gsn_csm_vector_scalar_map_ce',
-                'gsn_csm_vector_scalar_map_polar', 'gsn_csm_x2y', 'gsn_csm_x2y2',
-                'gsn_csm_xy', 'gsn_csm_xy2', 'gsn_csm_xy3', 'gsn_csm_y',
-                'gsn_define_colormap', 'gsn_draw_colormap', 'gsn_draw_named_colors',
-                'gsn_histogram', 'gsn_labelbar_ndc', 'gsn_legend_ndc', 'gsn_map',
-                'gsn_merge_colormaps', 'gsn_open_wks', 'gsn_panel', 'gsn_polygon',
-                'gsn_polygon_ndc', 'gsn_polyline', 'gsn_polyline_ndc',
-                'gsn_polymarker', 'gsn_polymarker_ndc', 'gsn_retrieve_colormap',
-                'gsn_reverse_colormap', 'gsn_streamline', 'gsn_streamline_map',
-                'gsn_streamline_scalar', 'gsn_streamline_scalar_map', 'gsn_table',
-                'gsn_text', 'gsn_text_ndc', 'gsn_vector', 'gsn_vector_map',
-                'gsn_vector_scalar', 'gsn_vector_scalar_map', 'gsn_xy', 'gsn_y',
-                'hsv2rgb', 'maximize_output', 'namedcolor2rgb', 'namedcolor2rgba',
-                'reset_device_coordinates', 'span_named_colors'), prefix=r'\b'),
-             Name.Builtin),
-
-            # Resources
-            (words((
-                'amDataXF', 'amDataYF', 'amJust', 'amOn', 'amOrthogonalPosF',
-                'amParallelPosF', 'amResizeNotify', 'amSide', 'amTrackData',
-                'amViewId', 'amZone', 'appDefaultParent', 'appFileSuffix',
-                'appResources', 'appSysDir', 'appUsrDir', 'caCopyArrays',
-                'caXArray', 'caXCast', 'caXMaxV', 'caXMinV', 'caXMissingV',
-                'caYArray', 'caYCast', 'caYMaxV', 'caYMinV', 'caYMissingV',
-                'cnCellFillEdgeColor', 'cnCellFillMissingValEdgeColor',
-                'cnConpackParams', 'cnConstFEnableFill', 'cnConstFLabelAngleF',
-                'cnConstFLabelBackgroundColor', 'cnConstFLabelConstantSpacingF',
-                'cnConstFLabelFont', 'cnConstFLabelFontAspectF',
-                'cnConstFLabelFontColor', 'cnConstFLabelFontHeightF',
-                'cnConstFLabelFontQuality', 'cnConstFLabelFontThicknessF',
-                'cnConstFLabelFormat', 'cnConstFLabelFuncCode', 'cnConstFLabelJust',
-                'cnConstFLabelOn', 'cnConstFLabelOrthogonalPosF',
-                'cnConstFLabelParallelPosF', 'cnConstFLabelPerimColor',
-                'cnConstFLabelPerimOn', 'cnConstFLabelPerimSpaceF',
-                'cnConstFLabelPerimThicknessF', 'cnConstFLabelSide',
-                'cnConstFLabelString', 'cnConstFLabelTextDirection',
-                'cnConstFLabelZone', 'cnConstFUseInfoLabelRes',
-                'cnExplicitLabelBarLabelsOn', 'cnExplicitLegendLabelsOn',
-                'cnExplicitLineLabelsOn', 'cnFillBackgroundColor', 'cnFillColor',
-                'cnFillColors', 'cnFillDotSizeF', 'cnFillDrawOrder', 'cnFillMode',
-                'cnFillOn', 'cnFillOpacityF', 'cnFillPalette', 'cnFillPattern',
-                'cnFillPatterns', 'cnFillScaleF', 'cnFillScales', 'cnFixFillBleed',
-                'cnGridBoundFillColor', 'cnGridBoundFillPattern',
-                'cnGridBoundFillScaleF', 'cnGridBoundPerimColor',
-                'cnGridBoundPerimDashPattern', 'cnGridBoundPerimOn',
-                'cnGridBoundPerimThicknessF', 'cnHighLabelAngleF',
-                'cnHighLabelBackgroundColor', 'cnHighLabelConstantSpacingF',
-                'cnHighLabelCount', 'cnHighLabelFont', 'cnHighLabelFontAspectF',
-                'cnHighLabelFontColor', 'cnHighLabelFontHeightF',
-                'cnHighLabelFontQuality', 'cnHighLabelFontThicknessF',
-                'cnHighLabelFormat', 'cnHighLabelFuncCode', 'cnHighLabelPerimColor',
-                'cnHighLabelPerimOn', 'cnHighLabelPerimSpaceF',
-                'cnHighLabelPerimThicknessF', 'cnHighLabelString', 'cnHighLabelsOn',
-                'cnHighLowLabelOverlapMode', 'cnHighUseLineLabelRes',
-                'cnInfoLabelAngleF', 'cnInfoLabelBackgroundColor',
-                'cnInfoLabelConstantSpacingF', 'cnInfoLabelFont',
-                'cnInfoLabelFontAspectF', 'cnInfoLabelFontColor',
-                'cnInfoLabelFontHeightF', 'cnInfoLabelFontQuality',
-                'cnInfoLabelFontThicknessF', 'cnInfoLabelFormat',
-                'cnInfoLabelFuncCode', 'cnInfoLabelJust', 'cnInfoLabelOn',
-                'cnInfoLabelOrthogonalPosF', 'cnInfoLabelParallelPosF',
-                'cnInfoLabelPerimColor', 'cnInfoLabelPerimOn',
-                'cnInfoLabelPerimSpaceF', 'cnInfoLabelPerimThicknessF',
-                'cnInfoLabelSide', 'cnInfoLabelString', 'cnInfoLabelTextDirection',
-                'cnInfoLabelZone', 'cnLabelBarEndLabelsOn', 'cnLabelBarEndStyle',
-                'cnLabelDrawOrder', 'cnLabelMasking', 'cnLabelScaleFactorF',
-                'cnLabelScaleValueF', 'cnLabelScalingMode', 'cnLegendLevelFlags',
-                'cnLevelCount', 'cnLevelFlag', 'cnLevelFlags', 'cnLevelSelectionMode',
-                'cnLevelSpacingF', 'cnLevels', 'cnLineColor', 'cnLineColors',
-                'cnLineDashPattern', 'cnLineDashPatterns', 'cnLineDashSegLenF',
-                'cnLineDrawOrder', 'cnLineLabelAngleF', 'cnLineLabelBackgroundColor',
-                'cnLineLabelConstantSpacingF', 'cnLineLabelCount',
-                'cnLineLabelDensityF', 'cnLineLabelFont', 'cnLineLabelFontAspectF',
-                'cnLineLabelFontColor', 'cnLineLabelFontColors',
-                'cnLineLabelFontHeightF', 'cnLineLabelFontQuality',
-                'cnLineLabelFontThicknessF', 'cnLineLabelFormat',
-                'cnLineLabelFuncCode', 'cnLineLabelInterval', 'cnLineLabelPerimColor',
-                'cnLineLabelPerimOn', 'cnLineLabelPerimSpaceF',
-                'cnLineLabelPerimThicknessF', 'cnLineLabelPlacementMode',
-                'cnLineLabelStrings', 'cnLineLabelsOn', 'cnLinePalette',
-                'cnLineThicknessF', 'cnLineThicknesses', 'cnLinesOn',
-                'cnLowLabelAngleF', 'cnLowLabelBackgroundColor',
-                'cnLowLabelConstantSpacingF', 'cnLowLabelCount', 'cnLowLabelFont',
-                'cnLowLabelFontAspectF', 'cnLowLabelFontColor',
-                'cnLowLabelFontHeightF', 'cnLowLabelFontQuality',
-                'cnLowLabelFontThicknessF', 'cnLowLabelFormat', 'cnLowLabelFuncCode',
-                'cnLowLabelPerimColor', 'cnLowLabelPerimOn', 'cnLowLabelPerimSpaceF',
-                'cnLowLabelPerimThicknessF', 'cnLowLabelString', 'cnLowLabelsOn',
-                'cnLowUseHighLabelRes', 'cnMaxDataValueFormat', 'cnMaxLevelCount',
-                'cnMaxLevelValF', 'cnMaxPointDistanceF', 'cnMinLevelValF',
-                'cnMissingValFillColor', 'cnMissingValFillPattern',
-                'cnMissingValFillScaleF', 'cnMissingValPerimColor',
-                'cnMissingValPerimDashPattern', 'cnMissingValPerimGridBoundOn',
-                'cnMissingValPerimOn', 'cnMissingValPerimThicknessF',
-                'cnMonoFillColor', 'cnMonoFillPattern', 'cnMonoFillScale',
-                'cnMonoLevelFlag', 'cnMonoLineColor', 'cnMonoLineDashPattern',
-                'cnMonoLineLabelFontColor', 'cnMonoLineThickness', 'cnNoDataLabelOn',
-                'cnNoDataLabelString', 'cnOutOfRangeFillColor',
-                'cnOutOfRangeFillPattern', 'cnOutOfRangeFillScaleF',
-                'cnOutOfRangePerimColor', 'cnOutOfRangePerimDashPattern',
-                'cnOutOfRangePerimOn', 'cnOutOfRangePerimThicknessF',
-                'cnRasterCellSizeF', 'cnRasterMinCellSizeF', 'cnRasterModeOn',
-                'cnRasterSampleFactorF', 'cnRasterSmoothingOn', 'cnScalarFieldData',
-                'cnSmoothingDistanceF', 'cnSmoothingOn', 'cnSmoothingTensionF',
-                'cnSpanFillPalette', 'cnSpanLinePalette', 'ctCopyTables',
-                'ctXElementSize', 'ctXMaxV', 'ctXMinV', 'ctXMissingV', 'ctXTable',
-                'ctXTableLengths', 'ctXTableType', 'ctYElementSize', 'ctYMaxV',
-                'ctYMinV', 'ctYMissingV', 'ctYTable', 'ctYTableLengths',
-                'ctYTableType', 'dcDelayCompute', 'errBuffer',
-                'errFileName', 'errFilePtr', 'errLevel', 'errPrint', 'errUnitNumber',
-                'gsClipOn', 'gsColors', 'gsEdgeColor', 'gsEdgeDashPattern',
-                'gsEdgeDashSegLenF', 'gsEdgeThicknessF', 'gsEdgesOn',
-                'gsFillBackgroundColor', 'gsFillColor', 'gsFillDotSizeF',
-                'gsFillIndex', 'gsFillLineThicknessF', 'gsFillOpacityF',
-                'gsFillScaleF', 'gsFont', 'gsFontAspectF', 'gsFontColor',
-                'gsFontHeightF', 'gsFontOpacityF', 'gsFontQuality',
-                'gsFontThicknessF', 'gsLineColor', 'gsLineDashPattern',
-                'gsLineDashSegLenF', 'gsLineLabelConstantSpacingF', 'gsLineLabelFont',
-                'gsLineLabelFontAspectF', 'gsLineLabelFontColor',
-                'gsLineLabelFontHeightF', 'gsLineLabelFontQuality',
-                'gsLineLabelFontThicknessF', 'gsLineLabelFuncCode',
-                'gsLineLabelString', 'gsLineOpacityF', 'gsLineThicknessF',
-                'gsMarkerColor', 'gsMarkerIndex', 'gsMarkerOpacityF', 'gsMarkerSizeF',
-                'gsMarkerThicknessF', 'gsSegments', 'gsTextAngleF',
-                'gsTextConstantSpacingF', 'gsTextDirection', 'gsTextFuncCode',
-                'gsTextJustification', 'gsnAboveYRefLineBarColors',
-                'gsnAboveYRefLineBarFillScales', 'gsnAboveYRefLineBarPatterns',
-                'gsnAboveYRefLineColor', 'gsnAddCyclic', 'gsnAttachBorderOn',
-                'gsnAttachPlotsXAxis', 'gsnBelowYRefLineBarColors',
-                'gsnBelowYRefLineBarFillScales', 'gsnBelowYRefLineBarPatterns',
-                'gsnBelowYRefLineColor', 'gsnBoxMargin', 'gsnCenterString',
-                'gsnCenterStringFontColor', 'gsnCenterStringFontHeightF',
-                'gsnCenterStringFuncCode', 'gsnCenterStringOrthogonalPosF',
-                'gsnCenterStringParallelPosF', 'gsnContourLineThicknessesScale',
-                'gsnContourNegLineDashPattern', 'gsnContourPosLineDashPattern',
-                'gsnContourZeroLineThicknessF', 'gsnDebugWriteFileName', 'gsnDraw',
-                'gsnFrame', 'gsnHistogramBarWidthPercent', 'gsnHistogramBinIntervals',
-                'gsnHistogramBinMissing', 'gsnHistogramBinWidth',
-                'gsnHistogramClassIntervals', 'gsnHistogramCompare',
-                'gsnHistogramComputePercentages',
-                'gsnHistogramComputePercentagesNoMissing',
-                'gsnHistogramDiscreteBinValues', 'gsnHistogramDiscreteClassValues',
-                'gsnHistogramHorizontal', 'gsnHistogramMinMaxBinsOn',
-                'gsnHistogramNumberOfBins', 'gsnHistogramPercentSign',
-                'gsnHistogramSelectNiceIntervals', 'gsnLeftString',
-                'gsnLeftStringFontColor', 'gsnLeftStringFontHeightF',
-                'gsnLeftStringFuncCode', 'gsnLeftStringOrthogonalPosF',
-                'gsnLeftStringParallelPosF', 'gsnMajorLatSpacing',
-                'gsnMajorLonSpacing', 'gsnMaskLambertConformal',
-                'gsnMaskLambertConformalOutlineOn', 'gsnMaximize',
-                'gsnMinorLatSpacing', 'gsnMinorLonSpacing', 'gsnPanelBottom',
-                'gsnPanelCenter', 'gsnPanelDebug', 'gsnPanelFigureStrings',
-                'gsnPanelFigureStringsBackgroundFillColor',
-                'gsnPanelFigureStringsFontHeightF', 'gsnPanelFigureStringsJust',
-                'gsnPanelFigureStringsPerimOn', 'gsnPanelLabelBar', 'gsnPanelLeft',
-                'gsnPanelMainFont', 'gsnPanelMainFontColor',
-                'gsnPanelMainFontHeightF', 'gsnPanelMainString', 'gsnPanelRight',
-                'gsnPanelRowSpec', 'gsnPanelScalePlotIndex', 'gsnPanelTop',
-                'gsnPanelXF', 'gsnPanelXWhiteSpacePercent', 'gsnPanelYF',
-                'gsnPanelYWhiteSpacePercent', 'gsnPaperHeight', 'gsnPaperMargin',
-                'gsnPaperOrientation', 'gsnPaperWidth', 'gsnPolar',
-                'gsnPolarLabelDistance', 'gsnPolarLabelFont',
-                'gsnPolarLabelFontHeightF', 'gsnPolarLabelSpacing', 'gsnPolarTime',
-                'gsnPolarUT', 'gsnRightString', 'gsnRightStringFontColor',
-                'gsnRightStringFontHeightF', 'gsnRightStringFuncCode',
-                'gsnRightStringOrthogonalPosF', 'gsnRightStringParallelPosF',
-                'gsnScalarContour', 'gsnScale', 'gsnShape', 'gsnSpreadColorEnd',
-                'gsnSpreadColorStart', 'gsnSpreadColors', 'gsnStringFont',
-                'gsnStringFontColor', 'gsnStringFontHeightF', 'gsnStringFuncCode',
-                'gsnTickMarksOn', 'gsnXAxisIrregular2Linear', 'gsnXAxisIrregular2Log',
-                'gsnXRefLine', 'gsnXRefLineColor', 'gsnXRefLineDashPattern',
-                'gsnXRefLineThicknessF', 'gsnXYAboveFillColors', 'gsnXYBarChart',
-                'gsnXYBarChartBarWidth', 'gsnXYBarChartColors',
-                'gsnXYBarChartColors2', 'gsnXYBarChartFillDotSizeF',
-                'gsnXYBarChartFillLineThicknessF', 'gsnXYBarChartFillOpacityF',
-                'gsnXYBarChartFillScaleF', 'gsnXYBarChartOutlineOnly',
-                'gsnXYBarChartOutlineThicknessF', 'gsnXYBarChartPatterns',
-                'gsnXYBarChartPatterns2', 'gsnXYBelowFillColors', 'gsnXYFillColors',
-                'gsnXYFillOpacities', 'gsnXYLeftFillColors', 'gsnXYRightFillColors',
-                'gsnYAxisIrregular2Linear', 'gsnYAxisIrregular2Log', 'gsnYRefLine',
-                'gsnYRefLineColor', 'gsnYRefLineColors', 'gsnYRefLineDashPattern',
-                'gsnYRefLineDashPatterns', 'gsnYRefLineThicknessF',
-                'gsnYRefLineThicknesses', 'gsnZonalMean', 'gsnZonalMeanXMaxF',
-                'gsnZonalMeanXMinF', 'gsnZonalMeanYRefLine', 'lbAutoManage',
-                'lbBottomMarginF', 'lbBoxCount', 'lbBoxEndCapStyle', 'lbBoxFractions',
-                'lbBoxLineColor', 'lbBoxLineDashPattern', 'lbBoxLineDashSegLenF',
-                'lbBoxLineThicknessF', 'lbBoxLinesOn', 'lbBoxMajorExtentF',
-                'lbBoxMinorExtentF', 'lbBoxSeparatorLinesOn', 'lbBoxSizing',
-                'lbFillBackground', 'lbFillColor', 'lbFillColors', 'lbFillDotSizeF',
-                'lbFillLineThicknessF', 'lbFillPattern', 'lbFillPatterns',
-                'lbFillScaleF', 'lbFillScales', 'lbJustification', 'lbLabelAlignment',
-                'lbLabelAngleF', 'lbLabelAutoStride', 'lbLabelBarOn',
-                'lbLabelConstantSpacingF', 'lbLabelDirection', 'lbLabelFont',
-                'lbLabelFontAspectF', 'lbLabelFontColor', 'lbLabelFontHeightF',
-                'lbLabelFontQuality', 'lbLabelFontThicknessF', 'lbLabelFuncCode',
-                'lbLabelJust', 'lbLabelOffsetF', 'lbLabelPosition', 'lbLabelStride',
-                'lbLabelStrings', 'lbLabelsOn', 'lbLeftMarginF', 'lbMaxLabelLenF',
-                'lbMinLabelSpacingF', 'lbMonoFillColor', 'lbMonoFillPattern',
-                'lbMonoFillScale', 'lbOrientation', 'lbPerimColor',
-                'lbPerimDashPattern', 'lbPerimDashSegLenF', 'lbPerimFill',
-                'lbPerimFillColor', 'lbPerimOn', 'lbPerimThicknessF',
-                'lbRasterFillOn', 'lbRightMarginF', 'lbTitleAngleF',
-                'lbTitleConstantSpacingF', 'lbTitleDirection', 'lbTitleExtentF',
-                'lbTitleFont', 'lbTitleFontAspectF', 'lbTitleFontColor',
-                'lbTitleFontHeightF', 'lbTitleFontQuality', 'lbTitleFontThicknessF',
-                'lbTitleFuncCode', 'lbTitleJust', 'lbTitleOffsetF', 'lbTitleOn',
-                'lbTitlePosition', 'lbTitleString', 'lbTopMarginF', 'lgAutoManage',
-                'lgBottomMarginF', 'lgBoxBackground', 'lgBoxLineColor',
-                'lgBoxLineDashPattern', 'lgBoxLineDashSegLenF', 'lgBoxLineThicknessF',
-                'lgBoxLinesOn', 'lgBoxMajorExtentF', 'lgBoxMinorExtentF',
-                'lgDashIndex', 'lgDashIndexes', 'lgItemCount', 'lgItemOrder',
-                'lgItemPlacement', 'lgItemPositions', 'lgItemType', 'lgItemTypes',
-                'lgJustification', 'lgLabelAlignment', 'lgLabelAngleF',
-                'lgLabelAutoStride', 'lgLabelConstantSpacingF', 'lgLabelDirection',
-                'lgLabelFont', 'lgLabelFontAspectF', 'lgLabelFontColor',
-                'lgLabelFontHeightF', 'lgLabelFontQuality', 'lgLabelFontThicknessF',
-                'lgLabelFuncCode', 'lgLabelJust', 'lgLabelOffsetF', 'lgLabelPosition',
-                'lgLabelStride', 'lgLabelStrings', 'lgLabelsOn', 'lgLeftMarginF',
-                'lgLegendOn', 'lgLineColor', 'lgLineColors', 'lgLineDashSegLenF',
-                'lgLineDashSegLens', 'lgLineLabelConstantSpacingF', 'lgLineLabelFont',
-                'lgLineLabelFontAspectF', 'lgLineLabelFontColor',
-                'lgLineLabelFontColors', 'lgLineLabelFontHeightF',
-                'lgLineLabelFontHeights', 'lgLineLabelFontQuality',
-                'lgLineLabelFontThicknessF', 'lgLineLabelFuncCode',
-                'lgLineLabelStrings', 'lgLineLabelsOn', 'lgLineThicknessF',
-                'lgLineThicknesses', 'lgMarkerColor', 'lgMarkerColors',
-                'lgMarkerIndex', 'lgMarkerIndexes', 'lgMarkerSizeF', 'lgMarkerSizes',
-                'lgMarkerThicknessF', 'lgMarkerThicknesses', 'lgMonoDashIndex',
-                'lgMonoItemType', 'lgMonoLineColor', 'lgMonoLineDashSegLen',
-                'lgMonoLineLabelFontColor', 'lgMonoLineLabelFontHeight',
-                'lgMonoLineThickness', 'lgMonoMarkerColor', 'lgMonoMarkerIndex',
-                'lgMonoMarkerSize', 'lgMonoMarkerThickness', 'lgOrientation',
-                'lgPerimColor', 'lgPerimDashPattern', 'lgPerimDashSegLenF',
-                'lgPerimFill', 'lgPerimFillColor', 'lgPerimOn', 'lgPerimThicknessF',
-                'lgRightMarginF', 'lgTitleAngleF', 'lgTitleConstantSpacingF',
-                'lgTitleDirection', 'lgTitleExtentF', 'lgTitleFont',
-                'lgTitleFontAspectF', 'lgTitleFontColor', 'lgTitleFontHeightF',
-                'lgTitleFontQuality', 'lgTitleFontThicknessF', 'lgTitleFuncCode',
-                'lgTitleJust', 'lgTitleOffsetF', 'lgTitleOn', 'lgTitlePosition',
-                'lgTitleString', 'lgTopMarginF', 'mpAreaGroupCount',
-                'mpAreaMaskingOn', 'mpAreaNames', 'mpAreaTypes', 'mpBottomAngleF',
-                'mpBottomMapPosF', 'mpBottomNDCF', 'mpBottomNPCF',
-                'mpBottomPointLatF', 'mpBottomPointLonF', 'mpBottomWindowF',
-                'mpCenterLatF', 'mpCenterLonF', 'mpCenterRotF', 'mpCountyLineColor',
-                'mpCountyLineDashPattern', 'mpCountyLineDashSegLenF',
-                'mpCountyLineThicknessF', 'mpDataBaseVersion', 'mpDataResolution',
-                'mpDataSetName', 'mpDefaultFillColor', 'mpDefaultFillPattern',
-                'mpDefaultFillScaleF', 'mpDynamicAreaGroups', 'mpEllipticalBoundary',
-                'mpFillAreaSpecifiers', 'mpFillBoundarySets', 'mpFillColor',
-                'mpFillColors', 'mpFillColors-default', 'mpFillDotSizeF',
-                'mpFillDrawOrder', 'mpFillOn', 'mpFillPatternBackground',
-                'mpFillPattern', 'mpFillPatterns', 'mpFillPatterns-default',
-                'mpFillScaleF', 'mpFillScales', 'mpFillScales-default',
-                'mpFixedAreaGroups', 'mpGeophysicalLineColor',
-                'mpGeophysicalLineDashPattern', 'mpGeophysicalLineDashSegLenF',
-                'mpGeophysicalLineThicknessF', 'mpGreatCircleLinesOn',
-                'mpGridAndLimbDrawOrder', 'mpGridAndLimbOn', 'mpGridLatSpacingF',
-                'mpGridLineColor', 'mpGridLineDashPattern', 'mpGridLineDashSegLenF',
-                'mpGridLineThicknessF', 'mpGridLonSpacingF', 'mpGridMaskMode',
-                'mpGridMaxLatF', 'mpGridPolarLonSpacingF', 'mpGridSpacingF',
-                'mpInlandWaterFillColor', 'mpInlandWaterFillPattern',
-                'mpInlandWaterFillScaleF', 'mpLabelDrawOrder', 'mpLabelFontColor',
-                'mpLabelFontHeightF', 'mpLabelsOn', 'mpLambertMeridianF',
-                'mpLambertParallel1F', 'mpLambertParallel2F', 'mpLandFillColor',
-                'mpLandFillPattern', 'mpLandFillScaleF', 'mpLeftAngleF',
-                'mpLeftCornerLatF', 'mpLeftCornerLonF', 'mpLeftMapPosF',
-                'mpLeftNDCF', 'mpLeftNPCF', 'mpLeftPointLatF',
-                'mpLeftPointLonF', 'mpLeftWindowF', 'mpLimbLineColor',
-                'mpLimbLineDashPattern', 'mpLimbLineDashSegLenF',
-                'mpLimbLineThicknessF', 'mpLimitMode', 'mpMaskAreaSpecifiers',
-                'mpMaskOutlineSpecifiers', 'mpMaxLatF', 'mpMaxLonF',
-                'mpMinLatF', 'mpMinLonF', 'mpMonoFillColor', 'mpMonoFillPattern',
-                'mpMonoFillScale', 'mpNationalLineColor', 'mpNationalLineDashPattern',
-                'mpNationalLineThicknessF', 'mpOceanFillColor', 'mpOceanFillPattern',
-                'mpOceanFillScaleF', 'mpOutlineBoundarySets', 'mpOutlineDrawOrder',
-                'mpOutlineMaskingOn', 'mpOutlineOn', 'mpOutlineSpecifiers',
-                'mpPerimDrawOrder', 'mpPerimLineColor', 'mpPerimLineDashPattern',
-                'mpPerimLineDashSegLenF', 'mpPerimLineThicknessF', 'mpPerimOn',
-                'mpPolyMode', 'mpProjection', 'mpProvincialLineColor',
-                'mpProvincialLineDashPattern', 'mpProvincialLineDashSegLenF',
-                'mpProvincialLineThicknessF', 'mpRelativeCenterLat',
-                'mpRelativeCenterLon', 'mpRightAngleF', 'mpRightCornerLatF',
-                'mpRightCornerLonF', 'mpRightMapPosF', 'mpRightNDCF',
-                'mpRightNPCF', 'mpRightPointLatF', 'mpRightPointLonF',
-                'mpRightWindowF', 'mpSatelliteAngle1F', 'mpSatelliteAngle2F',
-                'mpSatelliteDistF', 'mpShapeMode', 'mpSpecifiedFillColors',
-                'mpSpecifiedFillDirectIndexing', 'mpSpecifiedFillPatterns',
-                'mpSpecifiedFillPriority', 'mpSpecifiedFillScales',
-                'mpTopAngleF', 'mpTopMapPosF', 'mpTopNDCF', 'mpTopNPCF',
-                'mpTopPointLatF', 'mpTopPointLonF', 'mpTopWindowF',
-                'mpUSStateLineColor', 'mpUSStateLineDashPattern',
-                'mpUSStateLineDashSegLenF', 'mpUSStateLineThicknessF',
-                'pmAnnoManagers', 'pmAnnoViews', 'pmLabelBarDisplayMode',
-                'pmLabelBarHeightF', 'pmLabelBarKeepAspect', 'pmLabelBarOrthogonalPosF',
-                'pmLabelBarParallelPosF', 'pmLabelBarSide', 'pmLabelBarWidthF',
-                'pmLabelBarZone', 'pmLegendDisplayMode', 'pmLegendHeightF',
-                'pmLegendKeepAspect', 'pmLegendOrthogonalPosF',
-                'pmLegendParallelPosF', 'pmLegendSide', 'pmLegendWidthF',
-                'pmLegendZone', 'pmOverlaySequenceIds', 'pmTickMarkDisplayMode',
-                'pmTickMarkZone', 'pmTitleDisplayMode', 'pmTitleZone',
-                'prGraphicStyle', 'prPolyType', 'prXArray', 'prYArray',
-                'sfCopyData', 'sfDataArray', 'sfDataMaxV', 'sfDataMinV',
-                'sfElementNodes', 'sfExchangeDimensions', 'sfFirstNodeIndex',
-                'sfMissingValueV', 'sfXArray', 'sfXCActualEndF', 'sfXCActualStartF',
-                'sfXCEndIndex', 'sfXCEndSubsetV', 'sfXCEndV', 'sfXCStartIndex',
-                'sfXCStartSubsetV', 'sfXCStartV', 'sfXCStride', 'sfXCellBounds',
-                'sfYArray', 'sfYCActualEndF', 'sfYCActualStartF', 'sfYCEndIndex',
-                'sfYCEndSubsetV', 'sfYCEndV', 'sfYCStartIndex', 'sfYCStartSubsetV',
-                'sfYCStartV', 'sfYCStride', 'sfYCellBounds', 'stArrowLengthF',
-                'stArrowStride', 'stCrossoverCheckCount',
-                'stExplicitLabelBarLabelsOn', 'stLabelBarEndLabelsOn',
-                'stLabelFormat', 'stLengthCheckCount', 'stLevelColors',
-                'stLevelCount', 'stLevelPalette', 'stLevelSelectionMode',
-                'stLevelSpacingF', 'stLevels', 'stLineColor', 'stLineOpacityF',
-                'stLineStartStride', 'stLineThicknessF', 'stMapDirection',
-                'stMaxLevelCount', 'stMaxLevelValF', 'stMinArrowSpacingF',
-                'stMinDistanceF', 'stMinLevelValF', 'stMinLineSpacingF',
-                'stMinStepFactorF', 'stMonoLineColor', 'stNoDataLabelOn',
-                'stNoDataLabelString', 'stScalarFieldData', 'stScalarMissingValColor',
-                'stSpanLevelPalette', 'stStepSizeF', 'stStreamlineDrawOrder',
-                'stUseScalarArray', 'stVectorFieldData', 'stZeroFLabelAngleF',
-                'stZeroFLabelBackgroundColor', 'stZeroFLabelConstantSpacingF',
-                'stZeroFLabelFont', 'stZeroFLabelFontAspectF',
-                'stZeroFLabelFontColor', 'stZeroFLabelFontHeightF',
-                'stZeroFLabelFontQuality', 'stZeroFLabelFontThicknessF',
-                'stZeroFLabelFuncCode', 'stZeroFLabelJust', 'stZeroFLabelOn',
-                'stZeroFLabelOrthogonalPosF', 'stZeroFLabelParallelPosF',
-                'stZeroFLabelPerimColor', 'stZeroFLabelPerimOn',
-                'stZeroFLabelPerimSpaceF', 'stZeroFLabelPerimThicknessF',
-                'stZeroFLabelSide', 'stZeroFLabelString', 'stZeroFLabelTextDirection',
-                'stZeroFLabelZone', 'tfDoNDCOverlay', 'tfPlotManagerOn',
-                'tfPolyDrawList', 'tfPolyDrawOrder', 'tiDeltaF', 'tiMainAngleF',
-                'tiMainConstantSpacingF', 'tiMainDirection', 'tiMainFont',
-                'tiMainFontAspectF', 'tiMainFontColor', 'tiMainFontHeightF',
-                'tiMainFontQuality', 'tiMainFontThicknessF', 'tiMainFuncCode',
-                'tiMainJust', 'tiMainOffsetXF', 'tiMainOffsetYF', 'tiMainOn',
-                'tiMainPosition', 'tiMainSide', 'tiMainString', 'tiUseMainAttributes',
-                'tiXAxisAngleF', 'tiXAxisConstantSpacingF', 'tiXAxisDirection',
-                'tiXAxisFont', 'tiXAxisFontAspectF', 'tiXAxisFontColor',
-                'tiXAxisFontHeightF', 'tiXAxisFontQuality', 'tiXAxisFontThicknessF',
-                'tiXAxisFuncCode', 'tiXAxisJust', 'tiXAxisOffsetXF',
-                'tiXAxisOffsetYF', 'tiXAxisOn', 'tiXAxisPosition', 'tiXAxisSide',
-                'tiXAxisString', 'tiYAxisAngleF', 'tiYAxisConstantSpacingF',
-                'tiYAxisDirection', 'tiYAxisFont', 'tiYAxisFontAspectF',
-                'tiYAxisFontColor', 'tiYAxisFontHeightF', 'tiYAxisFontQuality',
-                'tiYAxisFontThicknessF', 'tiYAxisFuncCode', 'tiYAxisJust',
-                'tiYAxisOffsetXF', 'tiYAxisOffsetYF', 'tiYAxisOn', 'tiYAxisPosition',
-                'tiYAxisSide', 'tiYAxisString', 'tmBorderLineColor',
-                'tmBorderThicknessF', 'tmEqualizeXYSizes', 'tmLabelAutoStride',
-                'tmSciNoteCutoff', 'tmXBAutoPrecision', 'tmXBBorderOn',
-                'tmXBDataLeftF', 'tmXBDataRightF', 'tmXBFormat', 'tmXBIrrTensionF',
-                'tmXBIrregularPoints', 'tmXBLabelAngleF', 'tmXBLabelConstantSpacingF',
-                'tmXBLabelDeltaF', 'tmXBLabelDirection', 'tmXBLabelFont',
-                'tmXBLabelFontAspectF', 'tmXBLabelFontColor', 'tmXBLabelFontHeightF',
-                'tmXBLabelFontQuality', 'tmXBLabelFontThicknessF',
-                'tmXBLabelFuncCode', 'tmXBLabelJust', 'tmXBLabelStride', 'tmXBLabels',
-                'tmXBLabelsOn', 'tmXBMajorLengthF', 'tmXBMajorLineColor',
-                'tmXBMajorOutwardLengthF', 'tmXBMajorThicknessF', 'tmXBMaxLabelLenF',
-                'tmXBMaxTicks', 'tmXBMinLabelSpacingF', 'tmXBMinorLengthF',
-                'tmXBMinorLineColor', 'tmXBMinorOn', 'tmXBMinorOutwardLengthF',
-                'tmXBMinorPerMajor', 'tmXBMinorThicknessF', 'tmXBMinorValues',
-                'tmXBMode', 'tmXBOn', 'tmXBPrecision', 'tmXBStyle', 'tmXBTickEndF',
-                'tmXBTickSpacingF', 'tmXBTickStartF', 'tmXBValues', 'tmXMajorGrid',
-                'tmXMajorGridLineColor', 'tmXMajorGridLineDashPattern',
-                'tmXMajorGridThicknessF', 'tmXMinorGrid', 'tmXMinorGridLineColor',
-                'tmXMinorGridLineDashPattern', 'tmXMinorGridThicknessF',
-                'tmXTAutoPrecision', 'tmXTBorderOn', 'tmXTDataLeftF',
-                'tmXTDataRightF', 'tmXTFormat', 'tmXTIrrTensionF',
-                'tmXTIrregularPoints', 'tmXTLabelAngleF', 'tmXTLabelConstantSpacingF',
-                'tmXTLabelDeltaF', 'tmXTLabelDirection', 'tmXTLabelFont',
-                'tmXTLabelFontAspectF', 'tmXTLabelFontColor', 'tmXTLabelFontHeightF',
-                'tmXTLabelFontQuality', 'tmXTLabelFontThicknessF',
-                'tmXTLabelFuncCode', 'tmXTLabelJust', 'tmXTLabelStride', 'tmXTLabels',
-                'tmXTLabelsOn', 'tmXTMajorLengthF', 'tmXTMajorLineColor',
-                'tmXTMajorOutwardLengthF', 'tmXTMajorThicknessF', 'tmXTMaxLabelLenF',
-                'tmXTMaxTicks', 'tmXTMinLabelSpacingF', 'tmXTMinorLengthF',
-                'tmXTMinorLineColor', 'tmXTMinorOn', 'tmXTMinorOutwardLengthF',
-                'tmXTMinorPerMajor', 'tmXTMinorThicknessF', 'tmXTMinorValues',
-                'tmXTMode', 'tmXTOn', 'tmXTPrecision', 'tmXTStyle', 'tmXTTickEndF',
-                'tmXTTickSpacingF', 'tmXTTickStartF', 'tmXTValues', 'tmXUseBottom',
-                'tmYLAutoPrecision', 'tmYLBorderOn', 'tmYLDataBottomF',
-                'tmYLDataTopF', 'tmYLFormat', 'tmYLIrrTensionF',
-                'tmYLIrregularPoints', 'tmYLLabelAngleF', 'tmYLLabelConstantSpacingF',
-                'tmYLLabelDeltaF', 'tmYLLabelDirection', 'tmYLLabelFont',
-                'tmYLLabelFontAspectF', 'tmYLLabelFontColor', 'tmYLLabelFontHeightF',
-                'tmYLLabelFontQuality', 'tmYLLabelFontThicknessF',
-                'tmYLLabelFuncCode', 'tmYLLabelJust', 'tmYLLabelStride', 'tmYLLabels',
-                'tmYLLabelsOn', 'tmYLMajorLengthF', 'tmYLMajorLineColor',
-                'tmYLMajorOutwardLengthF', 'tmYLMajorThicknessF', 'tmYLMaxLabelLenF',
-                'tmYLMaxTicks', 'tmYLMinLabelSpacingF', 'tmYLMinorLengthF',
-                'tmYLMinorLineColor', 'tmYLMinorOn', 'tmYLMinorOutwardLengthF',
-                'tmYLMinorPerMajor', 'tmYLMinorThicknessF', 'tmYLMinorValues',
-                'tmYLMode', 'tmYLOn', 'tmYLPrecision', 'tmYLStyle', 'tmYLTickEndF',
-                'tmYLTickSpacingF', 'tmYLTickStartF', 'tmYLValues', 'tmYMajorGrid',
-                'tmYMajorGridLineColor', 'tmYMajorGridLineDashPattern',
-                'tmYMajorGridThicknessF', 'tmYMinorGrid', 'tmYMinorGridLineColor',
-                'tmYMinorGridLineDashPattern', 'tmYMinorGridThicknessF',
-                'tmYRAutoPrecision', 'tmYRBorderOn', 'tmYRDataBottomF',
-                'tmYRDataTopF', 'tmYRFormat', 'tmYRIrrTensionF',
-                'tmYRIrregularPoints', 'tmYRLabelAngleF', 'tmYRLabelConstantSpacingF',
-                'tmYRLabelDeltaF', 'tmYRLabelDirection', 'tmYRLabelFont',
-                'tmYRLabelFontAspectF', 'tmYRLabelFontColor', 'tmYRLabelFontHeightF',
-                'tmYRLabelFontQuality', 'tmYRLabelFontThicknessF',
-                'tmYRLabelFuncCode', 'tmYRLabelJust', 'tmYRLabelStride', 'tmYRLabels',
-                'tmYRLabelsOn', 'tmYRMajorLengthF', 'tmYRMajorLineColor',
-                'tmYRMajorOutwardLengthF', 'tmYRMajorThicknessF', 'tmYRMaxLabelLenF',
-                'tmYRMaxTicks', 'tmYRMinLabelSpacingF', 'tmYRMinorLengthF',
-                'tmYRMinorLineColor', 'tmYRMinorOn', 'tmYRMinorOutwardLengthF',
-                'tmYRMinorPerMajor', 'tmYRMinorThicknessF', 'tmYRMinorValues',
-                'tmYRMode', 'tmYROn', 'tmYRPrecision', 'tmYRStyle', 'tmYRTickEndF',
-                'tmYRTickSpacingF', 'tmYRTickStartF', 'tmYRValues', 'tmYUseLeft',
-                'trGridType', 'trLineInterpolationOn',
-                'trXAxisType', 'trXCoordPoints', 'trXInterPoints', 'trXLog',
-                'trXMaxF', 'trXMinF', 'trXReverse', 'trXSamples', 'trXTensionF',
-                'trYAxisType', 'trYCoordPoints', 'trYInterPoints', 'trYLog',
-                'trYMaxF', 'trYMinF', 'trYReverse', 'trYSamples', 'trYTensionF',
-                'txAngleF', 'txBackgroundFillColor', 'txConstantSpacingF', 'txDirection',
-                'txFont', 'HLU-Fonts', 'txFontAspectF', 'txFontColor',
-                'txFontHeightF', 'txFontOpacityF', 'txFontQuality',
-                'txFontThicknessF', 'txFuncCode', 'txJust', 'txPerimColor',
-                'txPerimDashLengthF', 'txPerimDashPattern', 'txPerimOn',
-                'txPerimSpaceF', 'txPerimThicknessF', 'txPosXF', 'txPosYF',
-                'txString', 'vcExplicitLabelBarLabelsOn', 'vcFillArrowEdgeColor',
-                'vcFillArrowEdgeThicknessF', 'vcFillArrowFillColor',
-                'vcFillArrowHeadInteriorXF', 'vcFillArrowHeadMinFracXF',
-                'vcFillArrowHeadMinFracYF', 'vcFillArrowHeadXF', 'vcFillArrowHeadYF',
-                'vcFillArrowMinFracWidthF', 'vcFillArrowWidthF', 'vcFillArrowsOn',
-                'vcFillOverEdge', 'vcGlyphOpacityF', 'vcGlyphStyle',
-                'vcLabelBarEndLabelsOn', 'vcLabelFontColor', 'vcLabelFontHeightF',
-                'vcLabelsOn', 'vcLabelsUseVectorColor', 'vcLevelColors',
-                'vcLevelCount', 'vcLevelPalette', 'vcLevelSelectionMode',
-                'vcLevelSpacingF', 'vcLevels', 'vcLineArrowColor',
-                'vcLineArrowHeadMaxSizeF', 'vcLineArrowHeadMinSizeF',
-                'vcLineArrowThicknessF', 'vcMagnitudeFormat',
-                'vcMagnitudeScaleFactorF', 'vcMagnitudeScaleValueF',
-                'vcMagnitudeScalingMode', 'vcMapDirection', 'vcMaxLevelCount',
-                'vcMaxLevelValF', 'vcMaxMagnitudeF', 'vcMinAnnoAngleF',
-                'vcMinAnnoArrowAngleF', 'vcMinAnnoArrowEdgeColor',
-                'vcMinAnnoArrowFillColor', 'vcMinAnnoArrowLineColor',
-                'vcMinAnnoArrowMinOffsetF', 'vcMinAnnoArrowSpaceF',
-                'vcMinAnnoArrowUseVecColor', 'vcMinAnnoBackgroundColor',
-                'vcMinAnnoConstantSpacingF', 'vcMinAnnoExplicitMagnitudeF',
-                'vcMinAnnoFont', 'vcMinAnnoFontAspectF', 'vcMinAnnoFontColor',
-                'vcMinAnnoFontHeightF', 'vcMinAnnoFontQuality',
-                'vcMinAnnoFontThicknessF', 'vcMinAnnoFuncCode', 'vcMinAnnoJust',
-                'vcMinAnnoOn', 'vcMinAnnoOrientation', 'vcMinAnnoOrthogonalPosF',
-                'vcMinAnnoParallelPosF', 'vcMinAnnoPerimColor', 'vcMinAnnoPerimOn',
-                'vcMinAnnoPerimSpaceF', 'vcMinAnnoPerimThicknessF', 'vcMinAnnoSide',
-                'vcMinAnnoString1', 'vcMinAnnoString1On', 'vcMinAnnoString2',
-                'vcMinAnnoString2On', 'vcMinAnnoTextDirection', 'vcMinAnnoZone',
-                'vcMinDistanceF', 'vcMinFracLengthF', 'vcMinLevelValF',
-                'vcMinMagnitudeF', 'vcMonoFillArrowEdgeColor',
-                'vcMonoFillArrowFillColor', 'vcMonoLineArrowColor',
-                'vcMonoWindBarbColor', 'vcNoDataLabelOn', 'vcNoDataLabelString',
-                'vcPositionMode', 'vcRefAnnoAngleF', 'vcRefAnnoArrowAngleF',
-                'vcRefAnnoArrowEdgeColor', 'vcRefAnnoArrowFillColor',
-                'vcRefAnnoArrowLineColor', 'vcRefAnnoArrowMinOffsetF',
-                'vcRefAnnoArrowSpaceF', 'vcRefAnnoArrowUseVecColor',
-                'vcRefAnnoBackgroundColor', 'vcRefAnnoConstantSpacingF',
-                'vcRefAnnoExplicitMagnitudeF', 'vcRefAnnoFont',
-                'vcRefAnnoFontAspectF', 'vcRefAnnoFontColor', 'vcRefAnnoFontHeightF',
-                'vcRefAnnoFontQuality', 'vcRefAnnoFontThicknessF',
-                'vcRefAnnoFuncCode', 'vcRefAnnoJust', 'vcRefAnnoOn',
-                'vcRefAnnoOrientation', 'vcRefAnnoOrthogonalPosF',
-                'vcRefAnnoParallelPosF', 'vcRefAnnoPerimColor', 'vcRefAnnoPerimOn',
-                'vcRefAnnoPerimSpaceF', 'vcRefAnnoPerimThicknessF', 'vcRefAnnoSide',
-                'vcRefAnnoString1', 'vcRefAnnoString1On', 'vcRefAnnoString2',
-                'vcRefAnnoString2On', 'vcRefAnnoTextDirection', 'vcRefAnnoZone',
-                'vcRefLengthF', 'vcRefMagnitudeF', 'vcScalarFieldData',
-                'vcScalarMissingValColor', 'vcScalarValueFormat',
-                'vcScalarValueScaleFactorF', 'vcScalarValueScaleValueF',
-                'vcScalarValueScalingMode', 'vcSpanLevelPalette', 'vcUseRefAnnoRes',
-                'vcUseScalarArray', 'vcVectorDrawOrder', 'vcVectorFieldData',
-                'vcWindBarbCalmCircleSizeF', 'vcWindBarbColor',
-                'vcWindBarbLineThicknessF', 'vcWindBarbScaleFactorF',
-                'vcWindBarbTickAngleF', 'vcWindBarbTickLengthF',
-                'vcWindBarbTickSpacingF', 'vcZeroFLabelAngleF',
-                'vcZeroFLabelBackgroundColor', 'vcZeroFLabelConstantSpacingF',
-                'vcZeroFLabelFont', 'vcZeroFLabelFontAspectF',
-                'vcZeroFLabelFontColor', 'vcZeroFLabelFontHeightF',
-                'vcZeroFLabelFontQuality', 'vcZeroFLabelFontThicknessF',
-                'vcZeroFLabelFuncCode', 'vcZeroFLabelJust', 'vcZeroFLabelOn',
-                'vcZeroFLabelOrthogonalPosF', 'vcZeroFLabelParallelPosF',
-                'vcZeroFLabelPerimColor', 'vcZeroFLabelPerimOn',
-                'vcZeroFLabelPerimSpaceF', 'vcZeroFLabelPerimThicknessF',
-                'vcZeroFLabelSide', 'vcZeroFLabelString', 'vcZeroFLabelTextDirection',
-                'vcZeroFLabelZone', 'vfCopyData', 'vfDataArray',
-                'vfExchangeDimensions', 'vfExchangeUVData', 'vfMagMaxV', 'vfMagMinV',
-                'vfMissingUValueV', 'vfMissingVValueV', 'vfPolarData',
-                'vfSingleMissingValue', 'vfUDataArray', 'vfUMaxV', 'vfUMinV',
-                'vfVDataArray', 'vfVMaxV', 'vfVMinV', 'vfXArray', 'vfXCActualEndF',
-                'vfXCActualStartF', 'vfXCEndIndex', 'vfXCEndSubsetV', 'vfXCEndV',
-                'vfXCStartIndex', 'vfXCStartSubsetV', 'vfXCStartV', 'vfXCStride',
-                'vfYArray', 'vfYCActualEndF', 'vfYCActualStartF', 'vfYCEndIndex',
-                'vfYCEndSubsetV', 'vfYCEndV', 'vfYCStartIndex', 'vfYCStartSubsetV',
-                'vfYCStartV', 'vfYCStride', 'vpAnnoManagerId', 'vpClipOn',
-                'vpHeightF', 'vpKeepAspect', 'vpOn', 'vpUseSegments', 'vpWidthF',
-                'vpXF', 'vpYF', 'wkAntiAlias', 'wkBackgroundColor', 'wkBackgroundOpacityF',
-                'wkColorMapLen', 'wkColorMap', 'wkColorModel', 'wkDashTableLength',
-                'wkDefGraphicStyleId', 'wkDeviceLowerX', 'wkDeviceLowerY',
-                'wkDeviceUpperX', 'wkDeviceUpperY', 'wkFileName', 'wkFillTableLength',
-                'wkForegroundColor', 'wkFormat', 'wkFullBackground', 'wkGksWorkId',
-                'wkHeight', 'wkMarkerTableLength', 'wkMetaName', 'wkOrientation',
-                'wkPDFFileName', 'wkPDFFormat', 'wkPDFResolution', 'wkPSFileName',
-                'wkPSFormat', 'wkPSResolution', 'wkPaperHeightF', 'wkPaperSize',
-                'wkPaperWidthF', 'wkPause', 'wkTopLevelViews', 'wkViews',
-                'wkVisualType', 'wkWidth', 'wkWindowId', 'wkXColorMode', 'wsCurrentSize',
-                'wsMaximumSize', 'wsThresholdSize', 'xyComputeXMax',
-                'xyComputeXMin', 'xyComputeYMax', 'xyComputeYMin', 'xyCoordData',
-                'xyCoordDataSpec', 'xyCurveDrawOrder', 'xyDashPattern',
-                'xyDashPatterns', 'xyExplicitLabels', 'xyExplicitLegendLabels',
-                'xyLabelMode', 'xyLineColor', 'xyLineColors', 'xyLineDashSegLenF',
-                'xyLineLabelConstantSpacingF', 'xyLineLabelFont',
-                'xyLineLabelFontAspectF', 'xyLineLabelFontColor',
-                'xyLineLabelFontColors', 'xyLineLabelFontHeightF',
-                'xyLineLabelFontQuality', 'xyLineLabelFontThicknessF',
-                'xyLineLabelFuncCode', 'xyLineThicknessF', 'xyLineThicknesses',
-                'xyMarkLineMode', 'xyMarkLineModes', 'xyMarker', 'xyMarkerColor',
-                'xyMarkerColors', 'xyMarkerSizeF', 'xyMarkerSizes',
-                'xyMarkerThicknessF', 'xyMarkerThicknesses', 'xyMarkers',
-                'xyMonoDashPattern', 'xyMonoLineColor', 'xyMonoLineLabelFontColor',
-                'xyMonoLineThickness', 'xyMonoMarkLineMode', 'xyMonoMarker',
-                'xyMonoMarkerColor', 'xyMonoMarkerSize', 'xyMonoMarkerThickness',
-                'xyXIrrTensionF', 'xyXIrregularPoints', 'xyXStyle', 'xyYIrrTensionF',
-                'xyYIrregularPoints', 'xyYStyle'), prefix=r'\b'),
-             Name.Builtin),
-
-            # Booleans
-            (r'\.(True|False)\.', Name.Builtin),
-            # Comparing Operators
-            (r'\.(eq|ne|lt|le|gt|ge|not|and|or|xor)\.', Operator.Word),
-        ],
-
-        'strings': [
-            (r'(?s)"(\\\\|\\[0-7]+|\\.|[^"\\])*"', String.Double),
-        ],
-
-        'nums': [
-            (r'\d+(?![.e])(_[a-z]\w+)?', Number.Integer),
-            (r'[+-]?\d*\.\d+(e[-+]?\d+)?(_[a-z]\w+)?', Number.Float),
-            (r'[+-]?\d+\.\d*(e[-+]?\d+)?(_[a-z]\w+)?', Number.Float),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/nimrod.py b/venv/lib/python3.10/site-packages/pygments/lexers/nimrod.py
deleted file mode 100644
index 365a8dcca961a8b1934b02d33cfaa181ed97736f..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/nimrod.py
+++ /dev/null
@@ -1,199 +0,0 @@
-"""
-    pygments.lexers.nimrod
-    ~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for the Nim language (formerly known as Nimrod).
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, include, default, bygroups
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Error
-
-__all__ = ['NimrodLexer']
-
-
-class NimrodLexer(RegexLexer):
-    """
-    For Nim source code.
-    """
-
-    name = 'Nimrod'
-    url = 'http://nim-lang.org/'
-    aliases = ['nimrod', 'nim']
-    filenames = ['*.nim', '*.nimrod']
-    mimetypes = ['text/x-nim']
-    version_added = '1.5'
-
-    flags = re.MULTILINE | re.IGNORECASE
-
-    def underscorize(words):
-        newWords = []
-        new = []
-        for word in words:
-            for ch in word:
-                new.append(ch)
-                new.append("_?")
-            newWords.append(''.join(new))
-            new = []
-        return "|".join(newWords)
-
-    keywords = [
-        'addr', 'and', 'as', 'asm', 'bind', 'block', 'break', 'case',
-        'cast', 'concept', 'const', 'continue', 'converter', 'defer', 'discard',
-        'distinct', 'div', 'do', 'elif', 'else', 'end', 'enum', 'except',
-        'export', 'finally', 'for', 'if', 'in', 'yield', 'interface',
-        'is', 'isnot', 'iterator', 'let', 'mixin', 'mod',
-        'not', 'notin', 'object', 'of', 'or', 'out', 'ptr', 'raise',
-        'ref', 'return', 'shl', 'shr', 'static', 'try',
-        'tuple', 'type', 'using', 'when', 'while', 'xor'
-    ]
-
-    keywordsPseudo = [
-        'nil', 'true', 'false'
-    ]
-
-    opWords = [
-        'and', 'or', 'not', 'xor', 'shl', 'shr', 'div', 'mod', 'in',
-        'notin', 'is', 'isnot'
-    ]
-
-    types = [
-        'int', 'int8', 'int16', 'int32', 'int64', 'float', 'float32', 'float64',
-        'bool', 'char', 'range', 'array', 'seq', 'set', 'string'
-    ]
-
-    tokens = {
-        'root': [
-            # Comments
-            (r'##\[', String.Doc, 'doccomment'),
-            (r'##.*$', String.Doc),
-            (r'#\[', Comment.Multiline, 'comment'),
-            (r'#.*$', Comment),
-
-            # Pragmas
-            (r'\{\.', String.Other, 'pragma'),
-
-            # Operators
-            (r'[*=><+\-/@$~&%!?|\\\[\]]', Operator),
-            (r'\.\.|\.|,|\[\.|\.\]|\{\.|\.\}|\(\.|\.\)|\{|\}|\(|\)|:|\^|`|;',
-             Punctuation),
-
-            # Case statement branch
-            (r'(\n\s*)(of)(\s)', bygroups(Text.Whitespace, Keyword,
-                                          Text.Whitespace), 'casebranch'),
-
-            # Strings
-            (r'(?:[\w]+)"', String, 'rdqs'),
-            (r'"""', String.Double, 'tdqs'),
-            ('"', String, 'dqs'),
-
-            # Char
-            ("'", String.Char, 'chars'),
-
-            # Keywords
-            (rf'({underscorize(opWords)})\b', Operator.Word),
-            (r'(proc|func|method|macro|template)(\s)(?![(\[\]])',
-             bygroups(Keyword, Text.Whitespace), 'funcname'),
-            (rf'({underscorize(keywords)})\b', Keyword),
-            (r'({})\b'.format(underscorize(['from', 'import', 'include', 'export'])),
-             Keyword.Namespace),
-            (r'(v_?a_?r)\b', Keyword.Declaration),
-            (rf'({underscorize(types)})\b', Name.Builtin),
-            (rf'({underscorize(keywordsPseudo)})\b', Keyword.Pseudo),
-
-            # Identifiers
-            (r'\b((?![_\d])\w)(((?!_)\w)|(_(?!_)\w))*', Name),
-
-            # Numbers
-            (r'[0-9][0-9_]*(?=([e.]|\'f(32|64)))',
-             Number.Float, ('float-suffix', 'float-number')),
-            (r'0x[a-f0-9][a-f0-9_]*', Number.Hex, 'int-suffix'),
-            (r'0b[01][01_]*', Number.Bin, 'int-suffix'),
-            (r'0o[0-7][0-7_]*', Number.Oct, 'int-suffix'),
-            (r'[0-9][0-9_]*', Number.Integer, 'int-suffix'),
-
-            # Whitespace
-            (r'\s+', Text.Whitespace),
-            (r'.+$', Error),
-        ],
-        'chars': [
-            (r'\\([\\abcefnrtvl"\']|x[a-f0-9]{2}|[0-9]{1,3})', String.Escape),
-            (r"'", String.Char, '#pop'),
-            (r".", String.Char)
-        ],
-        'strings': [
-            (r'(?|>=|>>|>|<=|<<|<|\+|-|=|/|\*|%|\+=|-=|!|@', Operator),
-            (r'\(|\)|\[|\]|,|\.\.\.|\.\.|\.|::|:', Punctuation),
-            (r'`\{[^`]*`\}', Text),  # Extern blocks won't be Lexed by Nit
-            (r'[\r\n\t ]+', Text),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/nix.py b/venv/lib/python3.10/site-packages/pygments/lexers/nix.py
deleted file mode 100644
index 3fa88c65aafd4843fe77b5b6774060339cb4c821..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/nix.py
+++ /dev/null
@@ -1,144 +0,0 @@
-"""
-    pygments.lexers.nix
-    ~~~~~~~~~~~~~~~~~~~
-
-    Lexers for the NixOS Nix language.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, include
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Literal
-
-__all__ = ['NixLexer']
-
-
-class NixLexer(RegexLexer):
-    """
-    For the Nix language.
-    """
-
-    name = 'Nix'
-    url = 'http://nixos.org/nix/'
-    aliases = ['nixos', 'nix']
-    filenames = ['*.nix']
-    mimetypes = ['text/x-nix']
-    version_added = '2.0'
-
-    keywords = ['rec', 'with', 'let', 'in', 'inherit', 'assert', 'if',
-                'else', 'then', '...']
-    builtins = ['import', 'abort', 'baseNameOf', 'dirOf', 'isNull', 'builtins',
-                'map', 'removeAttrs', 'throw', 'toString', 'derivation']
-    operators = ['++', '+', '?', '.', '!', '//', '==', '/',
-                 '!=', '&&', '||', '->', '=', '<', '>', '*', '-']
-
-    punctuations = ["(", ")", "[", "]", ";", "{", "}", ":", ",", "@"]
-
-    tokens = {
-        'root': [
-            # comments starting with #
-            (r'#.*$', Comment.Single),
-
-            # multiline comments
-            (r'/\*', Comment.Multiline, 'comment'),
-
-            # whitespace
-            (r'\s+', Text),
-
-            # keywords
-            ('({})'.format('|'.join(re.escape(entry) + '\\b' for entry in keywords)), Keyword),
-
-            # highlight the builtins
-            ('({})'.format('|'.join(re.escape(entry) + '\\b' for entry in builtins)),
-             Name.Builtin),
-
-            (r'\b(true|false|null)\b', Name.Constant),
-
-            # floats
-            (r'-?(\d+\.\d*|\.\d+)([eE][-+]?\d+)?', Number.Float),
-
-            # integers
-            (r'-?[0-9]+', Number.Integer),
-
-            # paths
-            (r'[\w.+-]*(\/[\w.+-]+)+', Literal),
-            (r'~(\/[\w.+-]+)+', Literal),
-            (r'\<[\w.+-]+(\/[\w.+-]+)*\>', Literal),
-
-            # operators
-            ('({})'.format('|'.join(re.escape(entry) for entry in operators)),
-             Operator),
-
-            # word operators
-            (r'\b(or|and)\b', Operator.Word),
-
-            (r'\{', Punctuation, 'block'),
-
-            # punctuations
-            ('({})'.format('|'.join(re.escape(entry) for entry in punctuations)), Punctuation),
-
-            # strings
-            (r'"', String.Double, 'doublequote'),
-            (r"''", String.Multiline, 'multiline'),
-
-            # urls
-            (r'[a-zA-Z][a-zA-Z0-9\+\-\.]*\:[\w%/?:@&=+$,\\.!~*\'-]+', Literal),
-
-            # names of variables
-            (r'[\w-]+(?=\s*=)', String.Symbol),
-            (r'[a-zA-Z_][\w\'-]*', Text),
-
-            (r"\$\{", String.Interpol, 'antiquote'),
-        ],
-        'comment': [
-            (r'[^/*]+', Comment.Multiline),
-            (r'/\*', Comment.Multiline, '#push'),
-            (r'\*/', Comment.Multiline, '#pop'),
-            (r'[*/]', Comment.Multiline),
-        ],
-        'multiline': [
-            (r"''(\$|'|\\n|\\r|\\t|\\)", String.Escape),
-            (r"''", String.Multiline, '#pop'),
-            (r'\$\{', String.Interpol, 'antiquote'),
-            (r"[^'\$]+", String.Multiline),
-            (r"\$[^\{']", String.Multiline),
-            (r"'[^']", String.Multiline),
-            (r"\$(?=')", String.Multiline),
-        ],
-        'doublequote': [
-            (r'\\(\\|"|\$|n)', String.Escape),
-            (r'"', String.Double, '#pop'),
-            (r'\$\{', String.Interpol, 'antiquote'),
-            (r'[^"\\\$]+', String.Double),
-            (r'\$[^\{"]', String.Double),
-            (r'\$(?=")', String.Double),
-            (r'\\', String.Double),
-        ],
-        'antiquote': [
-            (r"\}", String.Interpol, '#pop'),
-            # TODO: we should probably escape also here ''${ \${
-            (r"\$\{", String.Interpol, '#push'),
-            include('root'),
-        ],
-        'block': [
-            (r"\}", Punctuation, '#pop'),
-            include('root'),
-        ],
-    }
-
-    def analyse_text(text):
-        rv = 0.0
-        # TODO: let/in
-        if re.search(r'import.+?<[^>]+>', text):
-            rv += 0.4
-        if re.search(r'mkDerivation\s+(\(|\{|rec)', text):
-            rv += 0.4
-        if re.search(r'=\s+mkIf\s+', text):
-            rv += 0.4
-        if re.search(r'\{[a-zA-Z,\s]+\}:', text):
-            rv += 0.1
-        return rv
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/numbair.py b/venv/lib/python3.10/site-packages/pygments/lexers/numbair.py
deleted file mode 100644
index 435863e132afcff16c70ed81bb51739848c1ecc0..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/numbair.py
+++ /dev/null
@@ -1,63 +0,0 @@
-"""
-    pygments.lexers.numbair
-    ~~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for other Numba Intermediate Representation.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, include, bygroups, words
-from pygments.token import Whitespace, Name, String,  Punctuation, Keyword, \
-    Operator, Number
-
-__all__ = ["NumbaIRLexer"]
-
-class NumbaIRLexer(RegexLexer):
-    """
-    Lexer for Numba IR
-    """
-    name = 'Numba_IR'
-    url = "https://numba.readthedocs.io/en/stable/developer/architecture.html#stage-2-generate-the-numba-ir"
-    aliases = ['numba_ir', 'numbair']
-    filenames = ['*.numba_ir']
-    mimetypes = ['text/x-numba_ir', 'text/x-numbair']
-    version_added = '2.19'
-
-    identifier = r'\$[a-zA-Z0-9._]+'
-    fun_or_var = r'([a-zA-Z_]+[a-zA-Z0-9]*)'
-
-    tokens = {
-        'root' : [
-            (r'(label)(\ [0-9]+)(:)$',
-                bygroups(Keyword, Name.Label, Punctuation)),
-
-            (r'=', Operator),
-            include('whitespace'),
-            include('keyword'),
-
-            (identifier, Name.Variable),
-            (fun_or_var + r'(\()',
-                bygroups(Name.Function, Punctuation)),
-            (fun_or_var + r'(\=)',
-                bygroups(Name.Attribute, Punctuation)),
-            (fun_or_var, Name.Constant),
-            (r'[0-9]+', Number),
-
-            # 
-            (r'<[^>\n]*>', String),
-
-            (r'[=<>{}\[\]()*.,!\':]|x\b', Punctuation)
-        ],
-
-        'keyword':[
-            (words((
-                'del', 'jump', 'call', 'branch',
-            ), suffix=' '), Keyword),
-        ],
-
-        'whitespace': [
-            (r'(\n|\s)+', Whitespace),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/oberon.py b/venv/lib/python3.10/site-packages/pygments/lexers/oberon.py
deleted file mode 100644
index 61f3c2d2879e4dc7fedade39e203c1c50d5bbd20..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/oberon.py
+++ /dev/null
@@ -1,120 +0,0 @@
-"""
-    pygments.lexers.oberon
-    ~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Oberon family languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, include, words
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation
-
-__all__ = ['ComponentPascalLexer']
-
-
-class ComponentPascalLexer(RegexLexer):
-    """
-    For Component Pascal source code.
-    """
-    name = 'Component Pascal'
-    aliases = ['componentpascal', 'cp']
-    filenames = ['*.cp', '*.cps']
-    mimetypes = ['text/x-component-pascal']
-    url = 'https://blackboxframework.org'
-    version_added = '2.1'
-
-    flags = re.MULTILINE | re.DOTALL
-
-    tokens = {
-        'root': [
-            include('whitespace'),
-            include('comments'),
-            include('punctuation'),
-            include('numliterals'),
-            include('strings'),
-            include('operators'),
-            include('builtins'),
-            include('identifiers'),
-        ],
-        'whitespace': [
-            (r'\n+', Text),  # blank lines
-            (r'\s+', Text),  # whitespace
-        ],
-        'comments': [
-            (r'\(\*([^$].*?)\*\)', Comment.Multiline),
-            # TODO: nested comments (* (* ... *) ... (* ... *) *) not supported!
-        ],
-        'punctuation': [
-            (r'[()\[\]{},.:;|]', Punctuation),
-        ],
-        'numliterals': [
-            (r'[0-9A-F]+X\b', Number.Hex),                 # char code
-            (r'[0-9A-F]+[HL]\b', Number.Hex),              # hexadecimal number
-            (r'[0-9]+\.[0-9]+E[+-][0-9]+', Number.Float),  # real number
-            (r'[0-9]+\.[0-9]+', Number.Float),             # real number
-            (r'[0-9]+', Number.Integer),                   # decimal whole number
-        ],
-        'strings': [
-            (r"'[^\n']*'", String),  # single quoted string
-            (r'"[^\n"]*"', String),  # double quoted string
-        ],
-        'operators': [
-            # Arithmetic Operators
-            (r'[+-]', Operator),
-            (r'[*/]', Operator),
-            # Relational Operators
-            (r'[=#<>]', Operator),
-            # Dereferencing Operator
-            (r'\^', Operator),
-            # Logical AND Operator
-            (r'&', Operator),
-            # Logical NOT Operator
-            (r'~', Operator),
-            # Assignment Symbol
-            (r':=', Operator),
-            # Range Constructor
-            (r'\.\.', Operator),
-            (r'\$', Operator),
-        ],
-        'identifiers': [
-            (r'([a-zA-Z_$][\w$]*)', Name),
-        ],
-        'builtins': [
-            (words((
-                'ANYPTR', 'ANYREC', 'BOOLEAN', 'BYTE', 'CHAR', 'INTEGER', 'LONGINT',
-                'REAL', 'SET', 'SHORTCHAR', 'SHORTINT', 'SHORTREAL'
-                ), suffix=r'\b'), Keyword.Type),
-            (words((
-                'ABS', 'ABSTRACT', 'ARRAY', 'ASH', 'ASSERT', 'BEGIN', 'BITS', 'BY',
-                'CAP', 'CASE', 'CHR', 'CLOSE', 'CONST', 'DEC', 'DIV', 'DO', 'ELSE',
-                'ELSIF', 'EMPTY', 'END', 'ENTIER', 'EXCL', 'EXIT', 'EXTENSIBLE', 'FOR',
-                'HALT', 'IF', 'IMPORT', 'IN', 'INC', 'INCL', 'IS', 'LEN', 'LIMITED',
-                'LONG', 'LOOP', 'MAX', 'MIN', 'MOD', 'MODULE', 'NEW', 'ODD', 'OF',
-                'OR', 'ORD', 'OUT', 'POINTER', 'PROCEDURE', 'RECORD', 'REPEAT', 'RETURN',
-                'SHORT', 'SHORTCHAR', 'SHORTINT', 'SIZE', 'THEN', 'TYPE', 'TO', 'UNTIL',
-                'VAR', 'WHILE', 'WITH'
-                ), suffix=r'\b'), Keyword.Reserved),
-            (r'(TRUE|FALSE|NIL|INF)\b', Keyword.Constant),
-        ]
-    }
-
-    def analyse_text(text):
-        """The only other lexer using .cp is the C++ one, so we check if for
-        a few common Pascal keywords here. Those are unfortunately quite
-        common across various business languages as well."""
-        result = 0
-        if 'BEGIN' in text:
-            result += 0.01
-        if 'END' in text:
-            result += 0.01
-        if 'PROCEDURE' in text:
-            result += 0.01
-        if 'END' in text:
-            result += 0.01
-
-        return result
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/objective.py b/venv/lib/python3.10/site-packages/pygments/lexers/objective.py
deleted file mode 100644
index 899c2c44c965248e98a705d9bb330d4e8158b9ac..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/objective.py
+++ /dev/null
@@ -1,513 +0,0 @@
-"""
-    pygments.lexers.objective
-    ~~~~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Objective-C family languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, include, bygroups, using, this, words, \
-    inherit, default
-from pygments.token import Text, Keyword, Name, String, Operator, \
-    Number, Punctuation, Literal, Comment, Whitespace
-
-from pygments.lexers.c_cpp import CLexer, CppLexer
-
-__all__ = ['ObjectiveCLexer', 'ObjectiveCppLexer', 'LogosLexer', 'SwiftLexer']
-
-
-def objective(baselexer):
-    """
-    Generate a subclass of baselexer that accepts the Objective-C syntax
-    extensions.
-    """
-
-    # Have to be careful not to accidentally match JavaDoc/Doxygen syntax here,
-    # since that's quite common in ordinary C/C++ files.  It's OK to match
-    # JavaDoc/Doxygen keywords that only apply to Objective-C, mind.
-    #
-    # The upshot of this is that we CANNOT match @class or @interface
-    _oc_keywords = re.compile(r'@(?:end|implementation|protocol)')
-
-    # Matches [ ? identifier  ( identifier ? ] |  identifier? : )
-    # (note the identifier is *optional* when there is a ':'!)
-    _oc_message = re.compile(r'\[\s*[a-zA-Z_]\w*\s+'
-                             r'(?:[a-zA-Z_]\w*\s*\]|'
-                             r'(?:[a-zA-Z_]\w*)?:)')
-
-    class GeneratedObjectiveCVariant(baselexer):
-        """
-        Implements Objective-C syntax on top of an existing C family lexer.
-        """
-
-        tokens = {
-            'statements': [
-                (r'@"', String, 'string'),
-                (r'@(YES|NO)', Number),
-                (r"@'(\\.|\\[0-7]{1,3}|\\x[a-fA-F0-9]{1,2}|[^\\\'\n])'", String.Char),
-                (r'@(\d+\.\d*|\.\d+|\d+)[eE][+-]?\d+[lL]?', Number.Float),
-                (r'@(\d+\.\d*|\.\d+|\d+[fF])[fF]?', Number.Float),
-                (r'@0x[0-9a-fA-F]+[Ll]?', Number.Hex),
-                (r'@0[0-7]+[Ll]?', Number.Oct),
-                (r'@\d+[Ll]?', Number.Integer),
-                (r'@\(', Literal, 'literal_number'),
-                (r'@\[', Literal, 'literal_array'),
-                (r'@\{', Literal, 'literal_dictionary'),
-                (words((
-                    '@selector', '@private', '@protected', '@public', '@encode',
-                    '@synchronized', '@try', '@throw', '@catch', '@finally',
-                    '@end', '@property', '@synthesize', '__bridge', '__bridge_transfer',
-                    '__autoreleasing', '__block', '__weak', '__strong', 'weak', 'strong',
-                    'copy', 'retain', 'assign', 'unsafe_unretained', 'atomic', 'nonatomic',
-                    'readonly', 'readwrite', 'setter', 'getter', 'typeof', 'in',
-                    'out', 'inout', 'release', 'class', '@dynamic', '@optional',
-                    '@required', '@autoreleasepool', '@import'), suffix=r'\b'),
-                 Keyword),
-                (words(('id', 'instancetype', 'Class', 'IMP', 'SEL', 'BOOL',
-                        'IBOutlet', 'IBAction', 'unichar'), suffix=r'\b'),
-                 Keyword.Type),
-                (r'@(true|false|YES|NO)\n', Name.Builtin),
-                (r'(YES|NO|nil|self|super)\b', Name.Builtin),
-                # Carbon types
-                (r'(Boolean|UInt8|SInt8|UInt16|SInt16|UInt32|SInt32)\b', Keyword.Type),
-                # Carbon built-ins
-                (r'(TRUE|FALSE)\b', Name.Builtin),
-                (r'(@interface|@implementation)(\s+)', bygroups(Keyword, Text),
-                 ('#pop', 'oc_classname')),
-                (r'(@class|@protocol)(\s+)', bygroups(Keyword, Text),
-                 ('#pop', 'oc_forward_classname')),
-                # @ can also prefix other expressions like @{...} or @(...)
-                (r'@', Punctuation),
-                inherit,
-            ],
-            'oc_classname': [
-                # interface definition that inherits
-                (r'([a-zA-Z$_][\w$]*)(\s*:\s*)([a-zA-Z$_][\w$]*)?(\s*)(\{)',
-                 bygroups(Name.Class, Text, Name.Class, Text, Punctuation),
-                 ('#pop', 'oc_ivars')),
-                (r'([a-zA-Z$_][\w$]*)(\s*:\s*)([a-zA-Z$_][\w$]*)?',
-                 bygroups(Name.Class, Text, Name.Class), '#pop'),
-                # interface definition for a category
-                (r'([a-zA-Z$_][\w$]*)(\s*)(\([a-zA-Z$_][\w$]*\))(\s*)(\{)',
-                 bygroups(Name.Class, Text, Name.Label, Text, Punctuation),
-                 ('#pop', 'oc_ivars')),
-                (r'([a-zA-Z$_][\w$]*)(\s*)(\([a-zA-Z$_][\w$]*\))',
-                 bygroups(Name.Class, Text, Name.Label), '#pop'),
-                # simple interface / implementation
-                (r'([a-zA-Z$_][\w$]*)(\s*)(\{)',
-                 bygroups(Name.Class, Text, Punctuation), ('#pop', 'oc_ivars')),
-                (r'([a-zA-Z$_][\w$]*)', Name.Class, '#pop')
-            ],
-            'oc_forward_classname': [
-                (r'([a-zA-Z$_][\w$]*)(\s*,\s*)',
-                 bygroups(Name.Class, Text), 'oc_forward_classname'),
-                (r'([a-zA-Z$_][\w$]*)(\s*;?)',
-                 bygroups(Name.Class, Text), '#pop')
-            ],
-            'oc_ivars': [
-                include('whitespace'),
-                include('statements'),
-                (';', Punctuation),
-                (r'\{', Punctuation, '#push'),
-                (r'\}', Punctuation, '#pop'),
-            ],
-            'root': [
-                # methods
-                (r'^([-+])(\s*)'                         # method marker
-                 r'(\(.*?\))?(\s*)'                      # return type
-                 r'([a-zA-Z$_][\w$]*:?)',        # begin of method name
-                 bygroups(Punctuation, Text, using(this),
-                          Text, Name.Function),
-                 'method'),
-                inherit,
-            ],
-            'method': [
-                include('whitespace'),
-                # TODO unsure if ellipses are allowed elsewhere, see
-                # discussion in Issue 789
-                (r',', Punctuation),
-                (r'\.\.\.', Punctuation),
-                (r'(\(.*?\))(\s*)([a-zA-Z$_][\w$]*)',
-                 bygroups(using(this), Text, Name.Variable)),
-                (r'[a-zA-Z$_][\w$]*:', Name.Function),
-                (';', Punctuation, '#pop'),
-                (r'\{', Punctuation, 'function'),
-                default('#pop'),
-            ],
-            'literal_number': [
-                (r'\(', Punctuation, 'literal_number_inner'),
-                (r'\)', Literal, '#pop'),
-                include('statement'),
-            ],
-            'literal_number_inner': [
-                (r'\(', Punctuation, '#push'),
-                (r'\)', Punctuation, '#pop'),
-                include('statement'),
-            ],
-            'literal_array': [
-                (r'\[', Punctuation, 'literal_array_inner'),
-                (r'\]', Literal, '#pop'),
-                include('statement'),
-            ],
-            'literal_array_inner': [
-                (r'\[', Punctuation, '#push'),
-                (r'\]', Punctuation, '#pop'),
-                include('statement'),
-            ],
-            'literal_dictionary': [
-                (r'\}', Literal, '#pop'),
-                include('statement'),
-            ],
-        }
-
-        def analyse_text(text):
-            if _oc_keywords.search(text):
-                return 1.0
-            elif '@"' in text:  # strings
-                return 0.8
-            elif re.search('@[0-9]+', text):
-                return 0.7
-            elif _oc_message.search(text):
-                return 0.8
-            return 0
-
-        def get_tokens_unprocessed(self, text, stack=('root',)):
-            from pygments.lexers._cocoa_builtins import COCOA_INTERFACES, \
-                COCOA_PROTOCOLS, COCOA_PRIMITIVES
-
-            for index, token, value in \
-                    baselexer.get_tokens_unprocessed(self, text, stack):
-                if token is Name or token is Name.Class:
-                    if value in COCOA_INTERFACES or value in COCOA_PROTOCOLS \
-                       or value in COCOA_PRIMITIVES:
-                        token = Name.Builtin.Pseudo
-
-                yield index, token, value
-
-    return GeneratedObjectiveCVariant
-
-
-class ObjectiveCLexer(objective(CLexer)):
-    """
-    For Objective-C source code with preprocessor directives.
-    """
-
-    name = 'Objective-C'
-    url = 'https://developer.apple.com/library/archive/documentation/Cocoa/Conceptual/ProgrammingWithObjectiveC/Introduction/Introduction.html'
-    aliases = ['objective-c', 'objectivec', 'obj-c', 'objc']
-    filenames = ['*.m', '*.h']
-    mimetypes = ['text/x-objective-c']
-    version_added = ''
-    priority = 0.05    # Lower than C
-
-
-class ObjectiveCppLexer(objective(CppLexer)):
-    """
-    For Objective-C++ source code with preprocessor directives.
-    """
-
-    name = 'Objective-C++'
-    aliases = ['objective-c++', 'objectivec++', 'obj-c++', 'objc++']
-    filenames = ['*.mm', '*.hh']
-    mimetypes = ['text/x-objective-c++']
-    version_added = ''
-    priority = 0.05    # Lower than C++
-
-
-class LogosLexer(ObjectiveCppLexer):
-    """
-    For Logos + Objective-C source code with preprocessor directives.
-    """
-
-    name = 'Logos'
-    aliases = ['logos']
-    filenames = ['*.x', '*.xi', '*.xm', '*.xmi']
-    mimetypes = ['text/x-logos']
-    version_added = '1.6'
-    priority = 0.25
-
-    tokens = {
-        'statements': [
-            (r'(%orig|%log)\b', Keyword),
-            (r'(%c)\b(\()(\s*)([a-zA-Z$_][\w$]*)(\s*)(\))',
-             bygroups(Keyword, Punctuation, Text, Name.Class, Text, Punctuation)),
-            (r'(%init)\b(\()',
-             bygroups(Keyword, Punctuation), 'logos_init_directive'),
-            (r'(%init)(?=\s*;)', bygroups(Keyword)),
-            (r'(%hook|%group)(\s+)([a-zA-Z$_][\w$]+)',
-             bygroups(Keyword, Text, Name.Class), '#pop'),
-            (r'(%subclass)(\s+)', bygroups(Keyword, Text),
-             ('#pop', 'logos_classname')),
-            inherit,
-        ],
-        'logos_init_directive': [
-            (r'\s+', Text),
-            (',', Punctuation, ('logos_init_directive', '#pop')),
-            (r'([a-zA-Z$_][\w$]*)(\s*)(=)(\s*)([^);]*)',
-             bygroups(Name.Class, Text, Punctuation, Text, Text)),
-            (r'([a-zA-Z$_][\w$]*)', Name.Class),
-            (r'\)', Punctuation, '#pop'),
-        ],
-        'logos_classname': [
-            (r'([a-zA-Z$_][\w$]*)(\s*:\s*)([a-zA-Z$_][\w$]*)?',
-             bygroups(Name.Class, Text, Name.Class), '#pop'),
-            (r'([a-zA-Z$_][\w$]*)', Name.Class, '#pop')
-        ],
-        'root': [
-            (r'(%subclass)(\s+)', bygroups(Keyword, Text),
-             'logos_classname'),
-            (r'(%hook|%group)(\s+)([a-zA-Z$_][\w$]+)',
-             bygroups(Keyword, Text, Name.Class)),
-            (r'(%config)(\s*\(\s*)(\w+)(\s*=)(.*?)(\)\s*)',
-             bygroups(Keyword, Text, Name.Variable, Text, String, Text)),
-            (r'(%ctor)(\s*)(\{)', bygroups(Keyword, Text, Punctuation),
-             'function'),
-            (r'(%new)(\s*)(\()(.*?)(\))',
-             bygroups(Keyword, Text, Keyword, String, Keyword)),
-            (r'(\s*)(%end)(\s*)', bygroups(Text, Keyword, Text)),
-            inherit,
-        ],
-    }
-
-    _logos_keywords = re.compile(r'%(?:hook|ctor|init|c\()')
-
-    def analyse_text(text):
-        if LogosLexer._logos_keywords.search(text):
-            return 1.0
-        return 0
-
-
-class SwiftLexer(RegexLexer):
-    """
-    For Swift source.
-    """
-    name = 'Swift'
-    url = 'https://www.swift.org/'
-    filenames = ['*.swift']
-    aliases = ['swift']
-    mimetypes = ['text/x-swift']
-    version_added = '2.0'
-
-    tokens = {
-        'root': [
-            # Whitespace and Comments
-            (r'\n', Text),
-            (r'\s+', Whitespace),
-            (r'//', Comment.Single, 'comment-single'),
-            (r'/\*', Comment.Multiline, 'comment-multi'),
-            (r'#(if|elseif|else|endif|available)\b', Comment.Preproc, 'preproc'),
-
-            # Keywords
-            include('keywords'),
-
-            # Global Types
-            (words((
-                'Array', 'AutoreleasingUnsafeMutablePointer', 'BidirectionalReverseView',
-                'Bit', 'Bool', 'CFunctionPointer', 'COpaquePointer', 'CVaListPointer',
-                'Character', 'ClosedInterval', 'CollectionOfOne', 'ContiguousArray',
-                'Dictionary', 'DictionaryGenerator', 'DictionaryIndex', 'Double',
-                'EmptyCollection', 'EmptyGenerator', 'EnumerateGenerator',
-                'EnumerateSequence', 'FilterCollectionView',
-                'FilterCollectionViewIndex', 'FilterGenerator', 'FilterSequenceView',
-                'Float', 'Float80', 'FloatingPointClassification', 'GeneratorOf',
-                'GeneratorOfOne', 'GeneratorSequence', 'HalfOpenInterval', 'HeapBuffer',
-                'HeapBufferStorage', 'ImplicitlyUnwrappedOptional', 'IndexingGenerator',
-                'Int', 'Int16', 'Int32', 'Int64', 'Int8', 'LazyBidirectionalCollection',
-                'LazyForwardCollection', 'LazyRandomAccessCollection',
-                'LazySequence', 'MapCollectionView', 'MapSequenceGenerator',
-                'MapSequenceView', 'MirrorDisposition', 'ObjectIdentifier', 'OnHeap',
-                'Optional', 'PermutationGenerator', 'QuickLookObject',
-                'RandomAccessReverseView', 'Range', 'RangeGenerator', 'RawByte', 'Repeat',
-                'ReverseBidirectionalIndex', 'ReverseRandomAccessIndex', 'SequenceOf',
-                'SinkOf', 'Slice', 'StaticString', 'StrideThrough', 'StrideThroughGenerator',
-                'StrideTo', 'StrideToGenerator', 'String', 'UInt', 'UInt16', 'UInt32',
-                'UInt64', 'UInt8', 'UTF16', 'UTF32', 'UTF8', 'UnicodeDecodingResult',
-                'UnicodeScalar', 'Unmanaged', 'UnsafeBufferPointer',
-                'UnsafeBufferPointerGenerator', 'UnsafeMutableBufferPointer',
-                'UnsafeMutablePointer', 'UnsafePointer', 'Zip2', 'ZipGenerator2',
-                # Protocols
-                'AbsoluteValuable', 'AnyObject', 'ArrayLiteralConvertible',
-                'BidirectionalIndexType', 'BitwiseOperationsType',
-                'BooleanLiteralConvertible', 'BooleanType', 'CVarArgType',
-                'CollectionType', 'Comparable', 'DebugPrintable',
-                'DictionaryLiteralConvertible', 'Equatable',
-                'ExtendedGraphemeClusterLiteralConvertible',
-                'ExtensibleCollectionType', 'FloatLiteralConvertible',
-                'FloatingPointType', 'ForwardIndexType', 'GeneratorType', 'Hashable',
-                'IntegerArithmeticType', 'IntegerLiteralConvertible', 'IntegerType',
-                'IntervalType', 'MirrorType', 'MutableCollectionType', 'MutableSliceable',
-                'NilLiteralConvertible', 'OutputStreamType', 'Printable',
-                'RandomAccessIndexType', 'RangeReplaceableCollectionType',
-                'RawOptionSetType', 'RawRepresentable', 'Reflectable', 'SequenceType',
-                'SignedIntegerType', 'SignedNumberType', 'SinkType', 'Sliceable',
-                'Streamable', 'Strideable', 'StringInterpolationConvertible',
-                'StringLiteralConvertible', 'UnicodeCodecType',
-                'UnicodeScalarLiteralConvertible', 'UnsignedIntegerType',
-                '_ArrayBufferType', '_BidirectionalIndexType', '_CocoaStringType',
-                '_CollectionType', '_Comparable', '_ExtensibleCollectionType',
-                '_ForwardIndexType', '_Incrementable', '_IntegerArithmeticType',
-                '_IntegerType', '_ObjectiveCBridgeable', '_RandomAccessIndexType',
-                '_RawOptionSetType', '_SequenceType', '_Sequence_Type',
-                '_SignedIntegerType', '_SignedNumberType', '_Sliceable', '_Strideable',
-                '_SwiftNSArrayRequiredOverridesType', '_SwiftNSArrayType',
-                '_SwiftNSCopyingType', '_SwiftNSDictionaryRequiredOverridesType',
-                '_SwiftNSDictionaryType', '_SwiftNSEnumeratorType',
-                '_SwiftNSFastEnumerationType', '_SwiftNSStringRequiredOverridesType',
-                '_SwiftNSStringType', '_UnsignedIntegerType',
-                # Variables
-                'C_ARGC', 'C_ARGV', 'Process',
-                # Typealiases
-                'Any', 'AnyClass', 'BooleanLiteralType', 'CBool', 'CChar', 'CChar16',
-                'CChar32', 'CDouble', 'CFloat', 'CInt', 'CLong', 'CLongLong', 'CShort',
-                'CSignedChar', 'CUnsignedInt', 'CUnsignedLong', 'CUnsignedShort',
-                'CWideChar', 'ExtendedGraphemeClusterType', 'Float32', 'Float64',
-                'FloatLiteralType', 'IntMax', 'IntegerLiteralType', 'StringLiteralType',
-                'UIntMax', 'UWord', 'UnicodeScalarType', 'Void', 'Word',
-                # Foundation/Cocoa
-                'NSErrorPointer', 'NSObjectProtocol', 'Selector'), suffix=r'\b'),
-             Name.Builtin),
-            # Functions
-            (words((
-                'abs', 'advance', 'alignof', 'alignofValue', 'assert', 'assertionFailure',
-                'contains', 'count', 'countElements', 'debugPrint', 'debugPrintln',
-                'distance', 'dropFirst', 'dropLast', 'dump', 'enumerate', 'equal',
-                'extend', 'fatalError', 'filter', 'find', 'first', 'getVaList', 'indices',
-                'insert', 'isEmpty', 'join', 'last', 'lazy', 'lexicographicalCompare',
-                'map', 'max', 'maxElement', 'min', 'minElement', 'numericCast', 'overlaps',
-                'partition', 'precondition', 'preconditionFailure', 'prefix', 'print',
-                'println', 'reduce', 'reflect', 'removeAll', 'removeAtIndex', 'removeLast',
-                'removeRange', 'reverse', 'sizeof', 'sizeofValue', 'sort', 'sorted',
-                'splice', 'split', 'startsWith', 'stride', 'strideof', 'strideofValue',
-                'suffix', 'swap', 'toDebugString', 'toString', 'transcode',
-                'underestimateCount', 'unsafeAddressOf', 'unsafeBitCast', 'unsafeDowncast',
-                'withExtendedLifetime', 'withUnsafeMutablePointer',
-                'withUnsafeMutablePointers', 'withUnsafePointer', 'withUnsafePointers',
-                'withVaList'), suffix=r'\b'),
-             Name.Builtin.Pseudo),
-
-            # Implicit Block Variables
-            (r'\$\d+', Name.Variable),
-
-            # Binary Literal
-            (r'0b[01_]+', Number.Bin),
-            # Octal Literal
-            (r'0o[0-7_]+', Number.Oct),
-            # Hexadecimal Literal
-            (r'0x[0-9a-fA-F_]+', Number.Hex),
-            # Decimal Literal
-            (r'[0-9][0-9_]*(\.[0-9_]+[eE][+\-]?[0-9_]+|'
-             r'\.[0-9_]*|[eE][+\-]?[0-9_]+)', Number.Float),
-            (r'[0-9][0-9_]*', Number.Integer),
-            # String Literal
-            (r'"""', String, 'string-multi'),
-            (r'"', String, 'string'),
-
-            # Operators and Punctuation
-            (r'[(){}\[\].,:;=@#`?]|->|[<&?](?=\w)|(?<=\w)[>!?]', Punctuation),
-            (r'[/=\-+!*%<>&|^?~]+', Operator),
-
-            # Identifier
-            (r'[a-zA-Z_]\w*', Name)
-        ],
-        'keywords': [
-            (words((
-                'as', 'async', 'await', 'break', 'case', 'catch', 'continue', 'default', 'defer',
-                'do', 'else', 'fallthrough', 'for', 'guard', 'if', 'in', 'is',
-                'repeat', 'return', '#selector', 'switch', 'throw', 'try',
-                'where', 'while'), suffix=r'\b'),
-             Keyword),
-            (r'@availability\([^)]+\)', Keyword.Reserved),
-            (words((
-                'associativity', 'convenience', 'dynamic', 'didSet', 'final',
-                'get', 'indirect', 'infix', 'inout', 'lazy', 'left', 'mutating',
-                'none', 'nonmutating', 'optional', 'override', 'postfix',
-                'precedence', 'prefix', 'Protocol', 'required', 'rethrows',
-                'right', 'set', 'throws', 'Type', 'unowned', 'weak', 'willSet',
-                '@availability', '@autoclosure', '@noreturn',
-                '@NSApplicationMain', '@NSCopying', '@NSManaged', '@objc',
-                '@UIApplicationMain', '@IBAction', '@IBDesignable',
-                '@IBInspectable', '@IBOutlet'), suffix=r'\b'),
-             Keyword.Reserved),
-            (r'(as|dynamicType|false|is|nil|self|Self|super|true|__COLUMN__'
-             r'|__FILE__|__FUNCTION__|__LINE__|_'
-             r'|#(?:file|line|column|function))\b', Keyword.Constant),
-            (r'import\b', Keyword.Declaration, 'module'),
-            (r'(class|enum|extension|struct|protocol)(\s+)([a-zA-Z_]\w*)',
-             bygroups(Keyword.Declaration, Whitespace, Name.Class)),
-            (r'(func)(\s+)([a-zA-Z_]\w*)',
-             bygroups(Keyword.Declaration, Whitespace, Name.Function)),
-            (r'(var|let)(\s+)([a-zA-Z_]\w*)', bygroups(Keyword.Declaration,
-             Whitespace, Name.Variable)),
-            (words((
-                'actor', 'associatedtype', 'class', 'deinit', 'enum', 'extension', 'func', 'import',
-                'init', 'internal', 'let', 'operator', 'private', 'protocol', 'public',
-                'static', 'struct', 'subscript', 'typealias', 'var'), suffix=r'\b'),
-             Keyword.Declaration)
-        ],
-        'comment': [
-            (r':param: [a-zA-Z_]\w*|:returns?:|(FIXME|MARK|TODO):',
-             Comment.Special)
-        ],
-
-        # Nested
-        'comment-single': [
-            (r'\n', Whitespace, '#pop'),
-            include('comment'),
-            (r'[^\n]+', Comment.Single)
-        ],
-        'comment-multi': [
-            include('comment'),
-            (r'[^*/]+', Comment.Multiline),
-            (r'/\*', Comment.Multiline, '#push'),
-            (r'\*/', Comment.Multiline, '#pop'),
-            (r'[*/]+', Comment.Multiline)
-        ],
-        'module': [
-            (r'\n', Whitespace, '#pop'),
-            (r'[a-zA-Z_]\w*', Name.Class),
-            include('root')
-        ],
-        'preproc': [
-            (r'\n', Whitespace, '#pop'),
-            include('keywords'),
-            (r'[A-Za-z]\w*', Comment.Preproc),
-            include('root')
-        ],
-        'string': [
-            (r'"', String, '#pop'),
-            include("string-common"),
-        ],
-        'string-multi': [
-            (r'"""', String, '#pop'),
-            include("string-common"),
-        ],
-        'string-common': [
-            (r'\\\(', String.Interpol, 'string-intp'),
-            (r"""\\['"\\nrt]|\\x[0-9a-fA-F]{2}|\\[0-7]{1,3}"""
-             r"""|\\u[0-9a-fA-F]{4}|\\U[0-9a-fA-F]{8}""", String.Escape),
-            (r'[^\\"]+', String),
-            (r'\\', String)
-        ],
-        'string-intp': [
-            (r'\(', String.Interpol, '#push'),
-            (r'\)', String.Interpol, '#pop'),
-            include('root')
-        ]
-    }
-
-    def get_tokens_unprocessed(self, text):
-        from pygments.lexers._cocoa_builtins import COCOA_INTERFACES, \
-            COCOA_PROTOCOLS, COCOA_PRIMITIVES
-
-        for index, token, value in \
-                RegexLexer.get_tokens_unprocessed(self, text):
-            if token is Name or token is Name.Class:
-                if value in COCOA_INTERFACES or value in COCOA_PROTOCOLS \
-                   or value in COCOA_PRIMITIVES:
-                    token = Name.Builtin.Pseudo
-
-            yield index, token, value
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/ooc.py b/venv/lib/python3.10/site-packages/pygments/lexers/ooc.py
deleted file mode 100644
index 8a990801a3c31535c8713c80a0b4d95215f4a6c8..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/ooc.py
+++ /dev/null
@@ -1,84 +0,0 @@
-"""
-    pygments.lexers.ooc
-    ~~~~~~~~~~~~~~~~~~~
-
-    Lexers for the Ooc language.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, bygroups, words
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation
-
-__all__ = ['OocLexer']
-
-
-class OocLexer(RegexLexer):
-    """
-    For Ooc source code
-    """
-    name = 'Ooc'
-    url = 'https://ooc-lang.github.io/'
-    aliases = ['ooc']
-    filenames = ['*.ooc']
-    mimetypes = ['text/x-ooc']
-    version_added = '1.2'
-
-    tokens = {
-        'root': [
-            (words((
-                'class', 'interface', 'implement', 'abstract', 'extends', 'from',
-                'this', 'super', 'new', 'const', 'final', 'static', 'import',
-                'use', 'extern', 'inline', 'proto', 'break', 'continue',
-                'fallthrough', 'operator', 'if', 'else', 'for', 'while', 'do',
-                'switch', 'case', 'as', 'in', 'version', 'return', 'true',
-                'false', 'null'), prefix=r'\b', suffix=r'\b'),
-             Keyword),
-            (r'include\b', Keyword, 'include'),
-            (r'(cover)([ \t]+)(from)([ \t]+)(\w+[*@]?)',
-             bygroups(Keyword, Text, Keyword, Text, Name.Class)),
-            (r'(func)((?:[ \t]|\\\n)+)(~[a-z_]\w*)',
-             bygroups(Keyword, Text, Name.Function)),
-            (r'\bfunc\b', Keyword),
-            # Note: %= not listed on https://ooc-lang.github.io/docs/lang/operators/
-            (r'//.*', Comment),
-            (r'(?s)/\*.*?\*/', Comment.Multiline),
-            (r'(==?|\+=?|-[=>]?|\*=?|/=?|:=|!=?|%=?|\?|>{1,3}=?|<{1,3}=?|\.\.|'
-             r'&&?|\|\|?|\^=?)', Operator),
-            (r'(\.)([ \t]*)([a-z]\w*)', bygroups(Operator, Text,
-                                                 Name.Function)),
-            (r'[A-Z][A-Z0-9_]+', Name.Constant),
-            (r'[A-Z]\w*([@*]|\[[ \t]*\])?', Name.Class),
-
-            (r'([a-z]\w*(?:~[a-z]\w*)?)((?:[ \t]|\\\n)*)(?=\()',
-             bygroups(Name.Function, Text)),
-            (r'[a-z]\w*', Name.Variable),
-
-            # : introduces types
-            (r'[:(){}\[\];,]', Punctuation),
-
-            (r'0x[0-9a-fA-F]+', Number.Hex),
-            (r'0c[0-9]+', Number.Oct),
-            (r'0b[01]+', Number.Bin),
-            (r'[0-9_]\.[0-9_]*(?!\.)', Number.Float),
-            (r'[0-9_]+', Number.Decimal),
-
-            (r'"(?:\\.|\\[0-7]{1,3}|\\x[a-fA-F0-9]{1,2}|[^\\"])*"',
-             String.Double),
-            (r"'(?:\\.|\\[0-9]{1,3}|\\x[a-fA-F0-9]{1,2}|[^\\\'\n])'",
-             String.Char),
-            (r'@', Punctuation),  # pointer dereference
-            (r'\.', Punctuation),  # imports or chain operator
-
-            (r'\\[ \t\n]', Text),
-            (r'[ \t]+', Text),
-        ],
-        'include': [
-            (r'[\w/]+', Name),
-            (r',', Punctuation),
-            (r'[ \t]', Text),
-            (r'[;\n]', Text, '#pop'),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/openscad.py b/venv/lib/python3.10/site-packages/pygments/lexers/openscad.py
deleted file mode 100644
index b06de227c7c485db854e5ba25616431a84129ba8..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/openscad.py
+++ /dev/null
@@ -1,96 +0,0 @@
-"""
-    pygments.lexers.openscad
-    ~~~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for the OpenSCAD languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, bygroups, words, include
-from pygments.token import Text, Comment, Punctuation, Operator, Keyword, Name, Number, Whitespace, Literal, String
-
-__all__ = ['OpenScadLexer']
-
-
-class OpenScadLexer(RegexLexer):
-    """For openSCAD code.
-    """
-    name = "OpenSCAD"
-    url = "https://openscad.org/"
-    aliases = ["openscad"]
-    filenames = ["*.scad"]
-    mimetypes = ["application/x-openscad"]
-    version_added = '2.16'
-
-    tokens = {
-        "root": [
-            (r"[^\S\n]+", Whitespace),
-            (r'//', Comment.Single, 'comment-single'),
-            (r'/\*', Comment.Multiline, 'comment-multi'),
-            (r"[{}\[\]\(\),;:]", Punctuation),
-            (r"[*!#%\-+=?/]", Operator),
-            (r"<=|<|==|!=|>=|>|&&|\|\|", Operator),
-            (r"\$(f[asn]|t|vp[rtd]|children)", Operator),
-            (r"(undef|PI)\b", Keyword.Constant),
-            (
-                r"(use|include)((?:\s|\\\\s)+)",
-                bygroups(Keyword.Namespace, Text),
-                "includes",
-            ),
-            (r"(module)(\s*)([^\s\(]+)",
-             bygroups(Keyword.Namespace, Whitespace, Name.Namespace)),
-            (r"(function)(\s*)([^\s\(]+)",
-             bygroups(Keyword.Declaration, Whitespace, Name.Function)),
-            (words(("true", "false"), prefix=r"\b", suffix=r"\b"), Literal),
-            (words((
-                "function", "module", "include", "use", "for",
-                "intersection_for", "if", "else", "return"
-                ), prefix=r"\b", suffix=r"\b"), Keyword
-            ),
-            (words((
-                "circle", "square", "polygon", "text", "sphere", "cube",
-                "cylinder", "polyhedron", "translate", "rotate", "scale",
-                "resize", "mirror", "multmatrix", "color", "offset", "hull",
-                "minkowski", "union", "difference", "intersection", "abs",
-                "sign", "sin", "cos", "tan", "acos", "asin", "atan", "atan2",
-                "floor", "round", "ceil", "ln", "log", "pow", "sqrt", "exp",
-                "rands", "min", "max", "concat", "lookup", "str", "chr",
-                "search", "version", "version_num", "norm", "cross",
-                "parent_module", "echo", "import", "import_dxf",
-                "dxf_linear_extrude", "linear_extrude", "rotate_extrude",
-                "surface", "projection", "render", "dxf_cross",
-                "dxf_dim", "let", "assign", "len"
-                ), prefix=r"\b", suffix=r"\b"),
-                Name.Builtin
-            ),
-            (r"\bchildren\b", Name.Builtin.Pseudo),
-            (r'""".*?"""', String.Double),
-            (r'"(\\\\|\\[^\\]|[^"\\])*"', String.Double),
-            (r"-?\d+(\.\d+)?(e[+-]?\d+)?", Number),
-            (r"\w+", Name),
-        ],
-        "includes": [
-            (
-                r"(<)([^>]*)(>)",
-                bygroups(Punctuation, Comment.PreprocFile, Punctuation),
-            ),
-        ],
-        'comment': [
-            (r':param: [a-zA-Z_]\w*|:returns?:|(FIXME|MARK|TODO):',
-             Comment.Special)
-        ],
-        'comment-single': [
-            (r'\n', Text, '#pop'),
-            include('comment'),
-            (r'[^\n]+', Comment.Single)
-        ],
-        'comment-multi': [
-            include('comment'),
-            (r'[^*/]+', Comment.Multiline),
-            (r'/\*', Comment.Multiline, '#push'),
-            (r'\*/', Comment.Multiline, '#pop'),
-            (r'[*/]', Comment.Multiline)
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/other.py b/venv/lib/python3.10/site-packages/pygments/lexers/other.py
deleted file mode 100644
index 2b7dfb4ab534f51b785f35a1c13478397debac6c..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/other.py
+++ /dev/null
@@ -1,41 +0,0 @@
-"""
-    pygments.lexers.other
-    ~~~~~~~~~~~~~~~~~~~~~
-
-    Just export lexer classes previously contained in this module.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-# ruff: noqa: F401
-from pygments.lexers.sql import SqlLexer, MySqlLexer, SqliteConsoleLexer
-from pygments.lexers.shell import BashLexer, BashSessionLexer, BatchLexer, \
-    TcshLexer
-from pygments.lexers.robotframework import RobotFrameworkLexer
-from pygments.lexers.testing import GherkinLexer
-from pygments.lexers.esoteric import BrainfuckLexer, BefungeLexer, RedcodeLexer
-from pygments.lexers.prolog import LogtalkLexer
-from pygments.lexers.snobol import SnobolLexer
-from pygments.lexers.rebol import RebolLexer
-from pygments.lexers.configs import KconfigLexer, Cfengine3Lexer
-from pygments.lexers.modeling import ModelicaLexer
-from pygments.lexers.scripting import AppleScriptLexer, MOOCodeLexer, \
-    HybrisLexer
-from pygments.lexers.graphics import PostScriptLexer, GnuplotLexer, \
-    AsymptoteLexer, PovrayLexer
-from pygments.lexers.business import ABAPLexer, OpenEdgeLexer, \
-    GoodDataCLLexer, MaqlLexer
-from pygments.lexers.automation import AutoItLexer, AutohotkeyLexer
-from pygments.lexers.dsls import ProtoBufLexer, BroLexer, PuppetLexer, \
-    MscgenLexer, VGLLexer
-from pygments.lexers.basic import CbmBasicV2Lexer
-from pygments.lexers.pawn import SourcePawnLexer, PawnLexer
-from pygments.lexers.ecl import ECLLexer
-from pygments.lexers.urbi import UrbiscriptLexer
-from pygments.lexers.smalltalk import SmalltalkLexer, NewspeakLexer
-from pygments.lexers.installers import NSISLexer, RPMSpecLexer
-from pygments.lexers.textedit import AwkLexer
-from pygments.lexers.smv import NuSMVLexer
-
-__all__ = []
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/parasail.py b/venv/lib/python3.10/site-packages/pygments/lexers/parasail.py
deleted file mode 100644
index 150d6a9c48213856c423f4ceb28fb1ab7400f137..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/parasail.py
+++ /dev/null
@@ -1,78 +0,0 @@
-"""
-    pygments.lexers.parasail
-    ~~~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for ParaSail.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, include
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Literal
-
-__all__ = ['ParaSailLexer']
-
-
-class ParaSailLexer(RegexLexer):
-    """
-    For ParaSail source code.
-    """
-
-    name = 'ParaSail'
-    url = 'http://www.parasail-lang.org'
-    aliases = ['parasail']
-    filenames = ['*.psi', '*.psl']
-    mimetypes = ['text/x-parasail']
-    version_added = '2.1'
-
-    flags = re.MULTILINE
-
-    tokens = {
-        'root': [
-            (r'[^\S\n]+', Text),
-            (r'//.*?\n', Comment.Single),
-            (r'\b(and|or|xor)=', Operator.Word),
-            (r'\b(and(\s+then)?|or(\s+else)?|xor|rem|mod|'
-             r'(is|not)\s+null)\b',
-             Operator.Word),
-            # Keywords
-            (r'\b(abs|abstract|all|block|class|concurrent|const|continue|'
-             r'each|end|exit|extends|exports|forward|func|global|implements|'
-             r'import|in|interface|is|lambda|locked|new|not|null|of|op|'
-             r'optional|private|queued|ref|return|reverse|separate|some|'
-             r'type|until|var|with|'
-             # Control flow
-             r'if|then|else|elsif|case|for|while|loop)\b',
-             Keyword.Reserved),
-            (r'(abstract\s+)?(interface|class|op|func|type)',
-             Keyword.Declaration),
-            # Literals
-            (r'"[^"]*"', String),
-            (r'\\[\'ntrf"0]', String.Escape),
-            (r'#[a-zA-Z]\w*', Literal),       # Enumeration
-            include('numbers'),
-            (r"'[^']'", String.Char),
-            (r'[a-zA-Z]\w*', Name),
-            # Operators and Punctuation
-            (r'(<==|==>|<=>|\*\*=|<\|=|<<=|>>=|==|!=|=\?|<=|>=|'
-             r'\*\*|<<|>>|=>|:=|\+=|-=|\*=|\|=|\||/=|\+|-|\*|/|'
-             r'\.\.|<\.\.|\.\.<|<\.\.<)',
-             Operator),
-            (r'(<|>|\[|\]|\(|\)|\||:|;|,|.|\{|\}|->)',
-             Punctuation),
-            (r'\n+', Text),
-        ],
-        'numbers': [
-            (r'\d[0-9_]*#[0-9a-fA-F][0-9a-fA-F_]*#', Number.Hex),  # any base
-            (r'0[xX][0-9a-fA-F][0-9a-fA-F_]*', Number.Hex),        # C-like hex
-            (r'0[bB][01][01_]*', Number.Bin),                      # C-like bin
-            (r'\d[0-9_]*\.\d[0-9_]*[eE][+-]\d[0-9_]*',             # float exp
-             Number.Float),
-            (r'\d[0-9_]*\.\d[0-9_]*', Number.Float),               # float
-            (r'\d[0-9_]*', Number.Integer),                        # integer
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/parsers.py b/venv/lib/python3.10/site-packages/pygments/lexers/parsers.py
deleted file mode 100644
index 7a4ed9d1df789fd8e559e3c6de2fb161c4620547..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/parsers.py
+++ /dev/null
@@ -1,798 +0,0 @@
-"""
-    pygments.lexers.parsers
-    ~~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for parser generators.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, DelegatingLexer, \
-    include, bygroups, using
-from pygments.token import Punctuation, Other, Text, Comment, Operator, \
-    Keyword, Name, String, Number, Whitespace
-from pygments.lexers.jvm import JavaLexer
-from pygments.lexers.c_cpp import CLexer, CppLexer
-from pygments.lexers.objective import ObjectiveCLexer
-from pygments.lexers.d import DLexer
-from pygments.lexers.dotnet import CSharpLexer
-from pygments.lexers.ruby import RubyLexer
-from pygments.lexers.python import PythonLexer
-from pygments.lexers.perl import PerlLexer
-
-__all__ = ['RagelLexer', 'RagelEmbeddedLexer', 'RagelCLexer', 'RagelDLexer',
-           'RagelCppLexer', 'RagelObjectiveCLexer', 'RagelRubyLexer',
-           'RagelJavaLexer', 'AntlrLexer', 'AntlrPythonLexer',
-           'AntlrPerlLexer', 'AntlrRubyLexer', 'AntlrCppLexer',
-           'AntlrCSharpLexer', 'AntlrObjectiveCLexer',
-           'AntlrJavaLexer', 'AntlrActionScriptLexer',
-           'TreetopLexer', 'EbnfLexer']
-
-
-class RagelLexer(RegexLexer):
-    """A pure `Ragel `_ lexer.  Use this
-    for fragments of Ragel.  For ``.rl`` files, use
-    :class:`RagelEmbeddedLexer` instead (or one of the
-    language-specific subclasses).
-
-    """
-
-    name = 'Ragel'
-    url = 'http://www.colm.net/open-source/ragel/'
-    aliases = ['ragel']
-    filenames = []
-    version_added = '1.1'
-
-    tokens = {
-        'whitespace': [
-            (r'\s+', Whitespace)
-        ],
-        'comments': [
-            (r'\#.*$', Comment),
-        ],
-        'keywords': [
-            (r'(access|action|alphtype)\b', Keyword),
-            (r'(getkey|write|machine|include)\b', Keyword),
-            (r'(any|ascii|extend|alpha|digit|alnum|lower|upper)\b', Keyword),
-            (r'(xdigit|cntrl|graph|print|punct|space|zlen|empty)\b', Keyword)
-        ],
-        'numbers': [
-            (r'0x[0-9A-Fa-f]+', Number.Hex),
-            (r'[+-]?[0-9]+', Number.Integer),
-        ],
-        'literals': [
-            (r'"(\\\\|\\[^\\]|[^"\\])*"', String.Double),
-            (r"'(\\\\|\\[^\\]|[^'\\])*'", String.Single),
-            (r'\[(\\\\|\\[^\\]|[^\\\]])*\]', String),          # square bracket literals
-            (r'/(?!\*)(\\\\|\\[^\\]|[^/\\])*/', String.Regex),  # regular expressions
-        ],
-        'identifiers': [
-            (r'[a-zA-Z_]\w*', Name.Variable),
-        ],
-        'operators': [
-            (r',', Operator),                           # Join
-            (r'\||&|--?', Operator),                    # Union, Intersection and Subtraction
-            (r'\.|<:|:>>?', Operator),                  # Concatention
-            (r':', Operator),                           # Label
-            (r'->', Operator),                          # Epsilon Transition
-            (r'(>|\$|%|<|@|<>)(/|eof\b)', Operator),    # EOF Actions
-            (r'(>|\$|%|<|@|<>)(!|err\b)', Operator),    # Global Error Actions
-            (r'(>|\$|%|<|@|<>)(\^|lerr\b)', Operator),  # Local Error Actions
-            (r'(>|\$|%|<|@|<>)(~|to\b)', Operator),     # To-State Actions
-            (r'(>|\$|%|<|@|<>)(\*|from\b)', Operator),  # From-State Actions
-            (r'>|@|\$|%', Operator),                    # Transition Actions and Priorities
-            (r'\*|\?|\+|\{[0-9]*,[0-9]*\}', Operator),  # Repetition
-            (r'!|\^', Operator),                        # Negation
-            (r'\(|\)', Operator),                       # Grouping
-        ],
-        'root': [
-            include('literals'),
-            include('whitespace'),
-            include('comments'),
-            include('keywords'),
-            include('numbers'),
-            include('identifiers'),
-            include('operators'),
-            (r'\{', Punctuation, 'host'),
-            (r'=', Operator),
-            (r';', Punctuation),
-        ],
-        'host': [
-            (r'(' + r'|'.join((  # keep host code in largest possible chunks
-                r'[^{}\'"/#]+',  # exclude unsafe characters
-                r'[^\\]\\[{}]',  # allow escaped { or }
-
-                # strings and comments may safely contain unsafe characters
-                r'"(\\\\|\\[^\\]|[^"\\])*"',
-                r"'(\\\\|\\[^\\]|[^'\\])*'",
-                r'//.*$\n?',            # single line comment
-                r'/\*(.|\n)*?\*/',      # multi-line javadoc-style comment
-                r'\#.*$\n?',            # ruby comment
-
-                # regular expression: There's no reason for it to start
-                # with a * and this stops confusion with comments.
-                r'/(?!\*)(\\\\|\\[^\\]|[^/\\])*/',
-
-                # / is safe now that we've handled regex and javadoc comments
-                r'/',
-            )) + r')+', Other),
-
-            (r'\{', Punctuation, '#push'),
-            (r'\}', Punctuation, '#pop'),
-        ],
-    }
-
-
-class RagelEmbeddedLexer(RegexLexer):
-    """
-    A lexer for Ragel embedded in a host language file.
-
-    This will only highlight Ragel statements. If you want host language
-    highlighting then call the language-specific Ragel lexer.
-    """
-
-    name = 'Embedded Ragel'
-    aliases = ['ragel-em']
-    filenames = ['*.rl']
-    url = 'http://www.colm.net/open-source/ragel/'
-    version_added = '1.1'
-
-    tokens = {
-        'root': [
-            (r'(' + r'|'.join((   # keep host code in largest possible chunks
-                r'[^%\'"/#]+',    # exclude unsafe characters
-                r'%(?=[^%]|$)',   # a single % sign is okay, just not 2 of them
-
-                # strings and comments may safely contain unsafe characters
-                r'"(\\\\|\\[^\\]|[^"\\])*"',
-                r"'(\\\\|\\[^\\]|[^'\\])*'",
-                r'/\*(.|\n)*?\*/',      # multi-line javadoc-style comment
-                r'//.*$\n?',  # single line comment
-                r'\#.*$\n?',  # ruby/ragel comment
-                r'/(?!\*)(\\\\|\\[^\\]|[^/\\])*/',  # regular expression
-
-                # / is safe now that we've handled regex and javadoc comments
-                r'/',
-            )) + r')+', Other),
-
-            # Single Line FSM.
-            # Please don't put a quoted newline in a single line FSM.
-            # That's just mean. It will break this.
-            (r'(%%)(?![{%])(.*)($|;)(\n?)', bygroups(Punctuation,
-                                                     using(RagelLexer),
-                                                     Punctuation, Text)),
-
-            # Multi Line FSM.
-            (r'(%%%%|%%)\{', Punctuation, 'multi-line-fsm'),
-        ],
-        'multi-line-fsm': [
-            (r'(' + r'|'.join((  # keep ragel code in largest possible chunks.
-                r'(' + r'|'.join((
-                    r'[^}\'"\[/#]',   # exclude unsafe characters
-                    r'\}(?=[^%]|$)',   # } is okay as long as it's not followed by %
-                    r'\}%(?=[^%]|$)',  # ...well, one %'s okay, just not two...
-                    r'[^\\]\\[{}]',   # ...and } is okay if it's escaped
-
-                    # allow / if it's preceded with one of these symbols
-                    # (ragel EOF actions)
-                    r'(>|\$|%|<|@|<>)/',
-
-                    # specifically allow regex followed immediately by *
-                    # so it doesn't get mistaken for a comment
-                    r'/(?!\*)(\\\\|\\[^\\]|[^/\\])*/\*',
-
-                    # allow / as long as it's not followed by another / or by a *
-                    r'/(?=[^/*]|$)',
-
-                    # We want to match as many of these as we can in one block.
-                    # Not sure if we need the + sign here,
-                    # does it help performance?
-                )) + r')+',
-
-                # strings and comments may safely contain unsafe characters
-                r'"(\\\\|\\[^\\]|[^"\\])*"',
-                r"'(\\\\|\\[^\\]|[^'\\])*'",
-                r"\[(\\\\|\\[^\\]|[^\]\\])*\]",  # square bracket literal
-                r'/\*(.|\n)*?\*/',          # multi-line javadoc-style comment
-                r'//.*$\n?',                # single line comment
-                r'\#.*$\n?',                # ruby/ragel comment
-            )) + r')+', using(RagelLexer)),
-
-            (r'\}%%', Punctuation, '#pop'),
-        ]
-    }
-
-    def analyse_text(text):
-        return '@LANG: indep' in text
-
-
-class RagelRubyLexer(DelegatingLexer):
-    """
-    A lexer for Ragel in a Ruby host file.
-    """
-
-    name = 'Ragel in Ruby Host'
-    aliases = ['ragel-ruby', 'ragel-rb']
-    filenames = ['*.rl']
-    url = 'http://www.colm.net/open-source/ragel/'
-    version_added = '1.1'
-
-    def __init__(self, **options):
-        super().__init__(RubyLexer, RagelEmbeddedLexer, **options)
-
-    def analyse_text(text):
-        return '@LANG: ruby' in text
-
-
-class RagelCLexer(DelegatingLexer):
-    """
-    A lexer for Ragel in a C host file.
-    """
-
-    name = 'Ragel in C Host'
-    aliases = ['ragel-c']
-    filenames = ['*.rl']
-    url = 'http://www.colm.net/open-source/ragel/'
-    version_added = '1.1'
-
-    def __init__(self, **options):
-        super().__init__(CLexer, RagelEmbeddedLexer, **options)
-
-    def analyse_text(text):
-        return '@LANG: c' in text
-
-
-class RagelDLexer(DelegatingLexer):
-    """
-    A lexer for Ragel in a D host file.
-    """
-
-    name = 'Ragel in D Host'
-    aliases = ['ragel-d']
-    filenames = ['*.rl']
-    url = 'http://www.colm.net/open-source/ragel/'
-    version_added = '1.1'
-
-    def __init__(self, **options):
-        super().__init__(DLexer, RagelEmbeddedLexer, **options)
-
-    def analyse_text(text):
-        return '@LANG: d' in text
-
-
-class RagelCppLexer(DelegatingLexer):
-    """
-    A lexer for Ragel in a C++ host file.
-    """
-
-    name = 'Ragel in CPP Host'
-    aliases = ['ragel-cpp']
-    filenames = ['*.rl']
-    url = 'http://www.colm.net/open-source/ragel/'
-    version_added = '1.1'
-
-    def __init__(self, **options):
-        super().__init__(CppLexer, RagelEmbeddedLexer, **options)
-
-    def analyse_text(text):
-        return '@LANG: c++' in text
-
-
-class RagelObjectiveCLexer(DelegatingLexer):
-    """
-    A lexer for Ragel in an Objective C host file.
-    """
-
-    name = 'Ragel in Objective C Host'
-    aliases = ['ragel-objc']
-    filenames = ['*.rl']
-    url = 'http://www.colm.net/open-source/ragel/'
-    version_added = '1.1'
-
-    def __init__(self, **options):
-        super().__init__(ObjectiveCLexer, RagelEmbeddedLexer, **options)
-
-    def analyse_text(text):
-        return '@LANG: objc' in text
-
-
-class RagelJavaLexer(DelegatingLexer):
-    """
-    A lexer for Ragel in a Java host file.
-    """
-
-    name = 'Ragel in Java Host'
-    aliases = ['ragel-java']
-    filenames = ['*.rl']
-    url = 'http://www.colm.net/open-source/ragel/'
-    version_added = '1.1'
-
-    def __init__(self, **options):
-        super().__init__(JavaLexer, RagelEmbeddedLexer, **options)
-
-    def analyse_text(text):
-        return '@LANG: java' in text
-
-
-class AntlrLexer(RegexLexer):
-    """
-    Generic ANTLR Lexer.
-    Should not be called directly, instead
-    use DelegatingLexer for your target language.
-    """
-
-    name = 'ANTLR'
-    aliases = ['antlr']
-    filenames = []
-    url = 'https://www.antlr.org'
-    version_added = '1.1'
-
-    _id = r'[A-Za-z]\w*'
-    _TOKEN_REF = r'[A-Z]\w*'
-    _RULE_REF = r'[a-z]\w*'
-    _STRING_LITERAL = r'\'(?:\\\\|\\\'|[^\']*)\''
-    _INT = r'[0-9]+'
-
-    tokens = {
-        'whitespace': [
-            (r'\s+', Whitespace),
-        ],
-        'comments': [
-            (r'//.*$', Comment),
-            (r'/\*(.|\n)*?\*/', Comment),
-        ],
-        'root': [
-            include('whitespace'),
-            include('comments'),
-
-            (r'(lexer|parser|tree)?(\s*)(grammar\b)(\s*)(' + _id + ')(;)',
-             bygroups(Keyword, Whitespace, Keyword, Whitespace, Name.Class,
-                      Punctuation)),
-            # optionsSpec
-            (r'options\b', Keyword, 'options'),
-            # tokensSpec
-            (r'tokens\b', Keyword, 'tokens'),
-            # attrScope
-            (r'(scope)(\s*)(' + _id + r')(\s*)(\{)',
-             bygroups(Keyword, Whitespace, Name.Variable, Whitespace,
-                      Punctuation), 'action'),
-            # exception
-            (r'(catch|finally)\b', Keyword, 'exception'),
-            # action
-            (r'(@' + _id + r')(\s*)(::)?(\s*)(' + _id + r')(\s*)(\{)',
-             bygroups(Name.Label, Whitespace, Punctuation, Whitespace,
-                      Name.Label, Whitespace, Punctuation), 'action'),
-            # rule
-            (r'((?:protected|private|public|fragment)\b)?(\s*)(' + _id + ')(!)?',
-             bygroups(Keyword, Whitespace, Name.Label, Punctuation),
-             ('rule-alts', 'rule-prelims')),
-        ],
-        'exception': [
-            (r'\n', Whitespace, '#pop'),
-            (r'\s', Whitespace),
-            include('comments'),
-
-            (r'\[', Punctuation, 'nested-arg-action'),
-            (r'\{', Punctuation, 'action'),
-        ],
-        'rule-prelims': [
-            include('whitespace'),
-            include('comments'),
-
-            (r'returns\b', Keyword),
-            (r'\[', Punctuation, 'nested-arg-action'),
-            (r'\{', Punctuation, 'action'),
-            # throwsSpec
-            (r'(throws)(\s+)(' + _id + ')',
-             bygroups(Keyword, Whitespace, Name.Label)),
-            (r'(,)(\s*)(' + _id + ')',
-             bygroups(Punctuation, Whitespace, Name.Label)),  # Additional throws
-            # optionsSpec
-            (r'options\b', Keyword, 'options'),
-            # ruleScopeSpec - scope followed by target language code or name of action
-            # TODO finish implementing other possibilities for scope
-            # L173 ANTLRv3.g from ANTLR book
-            (r'(scope)(\s+)(\{)', bygroups(Keyword, Whitespace, Punctuation),
-             'action'),
-            (r'(scope)(\s+)(' + _id + r')(\s*)(;)',
-             bygroups(Keyword, Whitespace, Name.Label, Whitespace, Punctuation)),
-            # ruleAction
-            (r'(@' + _id + r')(\s*)(\{)',
-             bygroups(Name.Label, Whitespace, Punctuation), 'action'),
-            # finished prelims, go to rule alts!
-            (r':', Punctuation, '#pop')
-        ],
-        'rule-alts': [
-            include('whitespace'),
-            include('comments'),
-
-            # These might need to go in a separate 'block' state triggered by (
-            (r'options\b', Keyword, 'options'),
-            (r':', Punctuation),
-
-            # literals
-            (r'"(\\\\|\\[^\\]|[^"\\])*"', String.Double),
-            (r"'(\\\\|\\[^\\]|[^'\\])*'", String.Single),
-            (r'<<([^>]|>[^>])>>', String),
-            # identifiers
-            # Tokens start with capital letter.
-            (r'\$?[A-Z_]\w*', Name.Constant),
-            # Rules start with small letter.
-            (r'\$?[a-z_]\w*', Name.Variable),
-            # operators
-            (r'(\+|\||->|=>|=|\(|\)|\.\.|\.|\?|\*|\^|!|\#|~)', Operator),
-            (r',', Punctuation),
-            (r'\[', Punctuation, 'nested-arg-action'),
-            (r'\{', Punctuation, 'action'),
-            (r';', Punctuation, '#pop')
-        ],
-        'tokens': [
-            include('whitespace'),
-            include('comments'),
-            (r'\{', Punctuation),
-            (r'(' + _TOKEN_REF + r')(\s*)(=)?(\s*)(' + _STRING_LITERAL
-             + r')?(\s*)(;)',
-             bygroups(Name.Label, Whitespace, Punctuation, Whitespace,
-                      String, Whitespace, Punctuation)),
-            (r'\}', Punctuation, '#pop'),
-        ],
-        'options': [
-            include('whitespace'),
-            include('comments'),
-            (r'\{', Punctuation),
-            (r'(' + _id + r')(\s*)(=)(\s*)(' +
-             '|'.join((_id, _STRING_LITERAL, _INT, r'\*')) + r')(\s*)(;)',
-             bygroups(Name.Variable, Whitespace, Punctuation, Whitespace,
-                      Text, Whitespace, Punctuation)),
-            (r'\}', Punctuation, '#pop'),
-        ],
-        'action': [
-            (r'(' + r'|'.join((    # keep host code in largest possible chunks
-                r'[^${}\'"/\\]+',  # exclude unsafe characters
-
-                # strings and comments may safely contain unsafe characters
-                r'"(\\\\|\\[^\\]|[^"\\])*"',
-                r"'(\\\\|\\[^\\]|[^'\\])*'",
-                r'//.*$\n?',            # single line comment
-                r'/\*(.|\n)*?\*/',      # multi-line javadoc-style comment
-
-                # regular expression: There's no reason for it to start
-                # with a * and this stops confusion with comments.
-                r'/(?!\*)(\\\\|\\[^\\]|[^/\\])*/',
-
-                # backslashes are okay, as long as we are not backslashing a %
-                r'\\(?!%)',
-
-                # Now that we've handled regex and javadoc comments
-                # it's safe to let / through.
-                r'/',
-            )) + r')+', Other),
-            (r'(\\)(%)', bygroups(Punctuation, Other)),
-            (r'(\$[a-zA-Z]+)(\.?)(text|value)?',
-             bygroups(Name.Variable, Punctuation, Name.Property)),
-            (r'\{', Punctuation, '#push'),
-            (r'\}', Punctuation, '#pop'),
-        ],
-        'nested-arg-action': [
-            (r'(' + r'|'.join((    # keep host code in largest possible chunks.
-                r'[^$\[\]\'"/]+',  # exclude unsafe characters
-
-                # strings and comments may safely contain unsafe characters
-                r'"(\\\\|\\[^\\]|[^"\\])*"',
-                r"'(\\\\|\\[^\\]|[^'\\])*'",
-                r'//.*$\n?',            # single line comment
-                r'/\*(.|\n)*?\*/',      # multi-line javadoc-style comment
-
-                # regular expression: There's no reason for it to start
-                # with a * and this stops confusion with comments.
-                r'/(?!\*)(\\\\|\\[^\\]|[^/\\])*/',
-
-                # Now that we've handled regex and javadoc comments
-                # it's safe to let / through.
-                r'/',
-            )) + r')+', Other),
-
-
-            (r'\[', Punctuation, '#push'),
-            (r'\]', Punctuation, '#pop'),
-            (r'(\$[a-zA-Z]+)(\.?)(text|value)?',
-             bygroups(Name.Variable, Punctuation, Name.Property)),
-            (r'(\\\\|\\\]|\\\[|[^\[\]])+', Other),
-        ]
-    }
-
-    def analyse_text(text):
-        return re.search(r'^\s*grammar\s+[a-zA-Z0-9]+\s*;', text, re.M)
-
-
-# http://www.antlr.org/wiki/display/ANTLR3/Code+Generation+Targets
-
-class AntlrCppLexer(DelegatingLexer):
-    """
-    ANTLR with C++ Target
-    """
-
-    name = 'ANTLR With CPP Target'
-    aliases = ['antlr-cpp']
-    filenames = ['*.G', '*.g']
-    url = 'https://www.antlr.org'
-    version_added = '1.1'
-
-    def __init__(self, **options):
-        super().__init__(CppLexer, AntlrLexer, **options)
-
-    def analyse_text(text):
-        return AntlrLexer.analyse_text(text) and \
-            re.search(r'^\s*language\s*=\s*C\s*;', text, re.M)
-
-
-class AntlrObjectiveCLexer(DelegatingLexer):
-    """
-    ANTLR with Objective-C Target
-    """
-
-    name = 'ANTLR With ObjectiveC Target'
-    aliases = ['antlr-objc']
-    filenames = ['*.G', '*.g']
-    url = 'https://www.antlr.org'
-    version_added = '1.1'
-
-    def __init__(self, **options):
-        super().__init__(ObjectiveCLexer, AntlrLexer, **options)
-
-    def analyse_text(text):
-        return AntlrLexer.analyse_text(text) and \
-            re.search(r'^\s*language\s*=\s*ObjC\s*;', text)
-
-
-class AntlrCSharpLexer(DelegatingLexer):
-    """
-    ANTLR with C# Target
-    """
-
-    name = 'ANTLR With C# Target'
-    aliases = ['antlr-csharp', 'antlr-c#']
-    filenames = ['*.G', '*.g']
-    url = 'https://www.antlr.org'
-    version_added = '1.1'
-
-    def __init__(self, **options):
-        super().__init__(CSharpLexer, AntlrLexer, **options)
-
-    def analyse_text(text):
-        return AntlrLexer.analyse_text(text) and \
-            re.search(r'^\s*language\s*=\s*CSharp2\s*;', text, re.M)
-
-
-class AntlrPythonLexer(DelegatingLexer):
-    """
-    ANTLR with Python Target
-    """
-
-    name = 'ANTLR With Python Target'
-    aliases = ['antlr-python']
-    filenames = ['*.G', '*.g']
-    url = 'https://www.antlr.org'
-    version_added = '1.1'
-
-    def __init__(self, **options):
-        super().__init__(PythonLexer, AntlrLexer, **options)
-
-    def analyse_text(text):
-        return AntlrLexer.analyse_text(text) and \
-            re.search(r'^\s*language\s*=\s*Python\s*;', text, re.M)
-
-
-class AntlrJavaLexer(DelegatingLexer):
-    """
-    ANTLR with Java Target
-    """
-
-    name = 'ANTLR With Java Target'
-    aliases = ['antlr-java']
-    filenames = ['*.G', '*.g']
-    url = 'https://www.antlr.org'
-    version_added = '1.1'
-
-    def __init__(self, **options):
-        super().__init__(JavaLexer, AntlrLexer, **options)
-
-    def analyse_text(text):
-        # Antlr language is Java by default
-        return AntlrLexer.analyse_text(text) and 0.9
-
-
-class AntlrRubyLexer(DelegatingLexer):
-    """
-    ANTLR with Ruby Target
-    """
-
-    name = 'ANTLR With Ruby Target'
-    aliases = ['antlr-ruby', 'antlr-rb']
-    filenames = ['*.G', '*.g']
-    url = 'https://www.antlr.org'
-    version_added = '1.1'
-
-    def __init__(self, **options):
-        super().__init__(RubyLexer, AntlrLexer, **options)
-
-    def analyse_text(text):
-        return AntlrLexer.analyse_text(text) and \
-            re.search(r'^\s*language\s*=\s*Ruby\s*;', text, re.M)
-
-
-class AntlrPerlLexer(DelegatingLexer):
-    """
-    ANTLR with Perl Target
-    """
-
-    name = 'ANTLR With Perl Target'
-    aliases = ['antlr-perl']
-    filenames = ['*.G', '*.g']
-    url = 'https://www.antlr.org'
-    version_added = '1.1'
-
-    def __init__(self, **options):
-        super().__init__(PerlLexer, AntlrLexer, **options)
-
-    def analyse_text(text):
-        return AntlrLexer.analyse_text(text) and \
-            re.search(r'^\s*language\s*=\s*Perl5\s*;', text, re.M)
-
-
-class AntlrActionScriptLexer(DelegatingLexer):
-    """
-    ANTLR with ActionScript Target
-    """
-
-    name = 'ANTLR With ActionScript Target'
-    aliases = ['antlr-actionscript', 'antlr-as']
-    filenames = ['*.G', '*.g']
-    url = 'https://www.antlr.org'
-    version_added = '1.1'
-
-    def __init__(self, **options):
-        from pygments.lexers.actionscript import ActionScriptLexer
-        super().__init__(ActionScriptLexer, AntlrLexer, **options)
-
-    def analyse_text(text):
-        return AntlrLexer.analyse_text(text) and \
-            re.search(r'^\s*language\s*=\s*ActionScript\s*;', text, re.M)
-
-
-class TreetopBaseLexer(RegexLexer):
-    """
-    A base lexer for `Treetop `_ grammars.
-    Not for direct use; use :class:`TreetopLexer` instead.
-
-    .. versionadded:: 1.6
-    """
-
-    tokens = {
-        'root': [
-            include('space'),
-            (r'require[ \t]+[^\n\r]+[\n\r]', Other),
-            (r'module\b', Keyword.Namespace, 'module'),
-            (r'grammar\b', Keyword, 'grammar'),
-        ],
-        'module': [
-            include('space'),
-            include('end'),
-            (r'module\b', Keyword, '#push'),
-            (r'grammar\b', Keyword, 'grammar'),
-            (r'[A-Z]\w*(?:::[A-Z]\w*)*', Name.Namespace),
-        ],
-        'grammar': [
-            include('space'),
-            include('end'),
-            (r'rule\b', Keyword, 'rule'),
-            (r'include\b', Keyword, 'include'),
-            (r'[A-Z]\w*', Name),
-        ],
-        'include': [
-            include('space'),
-            (r'[A-Z]\w*(?:::[A-Z]\w*)*', Name.Class, '#pop'),
-        ],
-        'rule': [
-            include('space'),
-            include('end'),
-            (r'"(\\\\|\\[^\\]|[^"\\])*"', String.Double),
-            (r"'(\\\\|\\[^\\]|[^'\\])*'", String.Single),
-            (r'([A-Za-z_]\w*)(:)', bygroups(Name.Label, Punctuation)),
-            (r'[A-Za-z_]\w*', Name),
-            (r'[()]', Punctuation),
-            (r'[?+*/&!~]', Operator),
-            (r'\[(?:\\.|\[:\^?[a-z]+:\]|[^\\\]])+\]', String.Regex),
-            (r'([0-9]*)(\.\.)([0-9]*)',
-             bygroups(Number.Integer, Operator, Number.Integer)),
-            (r'(<)([^>]+)(>)', bygroups(Punctuation, Name.Class, Punctuation)),
-            (r'\{', Punctuation, 'inline_module'),
-            (r'\.', String.Regex),
-        ],
-        'inline_module': [
-            (r'\{', Other, 'ruby'),
-            (r'\}', Punctuation, '#pop'),
-            (r'[^{}]+', Other),
-        ],
-        'ruby': [
-            (r'\{', Other, '#push'),
-            (r'\}', Other, '#pop'),
-            (r'[^{}]+', Other),
-        ],
-        'space': [
-            (r'[ \t\n\r]+', Whitespace),
-            (r'#[^\n]*', Comment.Single),
-        ],
-        'end': [
-            (r'end\b', Keyword, '#pop'),
-        ],
-    }
-
-
-class TreetopLexer(DelegatingLexer):
-    """
-    A lexer for Treetop grammars.
-    """
-
-    name = 'Treetop'
-    aliases = ['treetop']
-    filenames = ['*.treetop', '*.tt']
-    url = 'https://cjheath.github.io/treetop'
-    version_added = '1.6'
-
-    def __init__(self, **options):
-        super().__init__(RubyLexer, TreetopBaseLexer, **options)
-
-
-class EbnfLexer(RegexLexer):
-    """
-    Lexer for `ISO/IEC 14977 EBNF
-    `_
-    grammars.
-    """
-
-    name = 'EBNF'
-    aliases = ['ebnf']
-    filenames = ['*.ebnf']
-    mimetypes = ['text/x-ebnf']
-    url = 'https://en.wikipedia.org/wiki/Extended_Backus%E2%80%93Naur_Form'
-    version_added = '2.0'
-
-    tokens = {
-        'root': [
-            include('whitespace'),
-            include('comment_start'),
-            include('identifier'),
-            (r'=', Operator, 'production'),
-        ],
-        'production': [
-            include('whitespace'),
-            include('comment_start'),
-            include('identifier'),
-            (r'"[^"]*"', String.Double),
-            (r"'[^']*'", String.Single),
-            (r'(\?[^?]*\?)', Name.Entity),
-            (r'[\[\]{}(),|]', Punctuation),
-            (r'-', Operator),
-            (r';', Punctuation, '#pop'),
-            (r'\.', Punctuation, '#pop'),
-        ],
-        'whitespace': [
-            (r'\s+', Text),
-        ],
-        'comment_start': [
-            (r'\(\*', Comment.Multiline, 'comment'),
-        ],
-        'comment': [
-            (r'[^*)]', Comment.Multiline),
-            include('comment_start'),
-            (r'\*\)', Comment.Multiline, '#pop'),
-            (r'[*)]', Comment.Multiline),
-        ],
-        'identifier': [
-            (r'([a-zA-Z][\w \-]*)', Keyword),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/pascal.py b/venv/lib/python3.10/site-packages/pygments/lexers/pascal.py
deleted file mode 100644
index 5f40dcc86a1fc26e0277ffd311f6086d231fc33e..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/pascal.py
+++ /dev/null
@@ -1,644 +0,0 @@
-"""
-    pygments.lexers.pascal
-    ~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Pascal family languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import Lexer
-from pygments.util import get_bool_opt, get_list_opt
-from pygments.token import Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Error, Whitespace
-from pygments.scanner import Scanner
-
-# compatibility import
-from pygments.lexers.modula2 import Modula2Lexer # noqa: F401
-
-__all__ = ['DelphiLexer', 'PortugolLexer']
-
-
-class PortugolLexer(Lexer):
-    """For Portugol, a Pascal dialect with keywords in Portuguese."""
-    name = 'Portugol'
-    aliases = ['portugol']
-    filenames = ['*.alg', '*.portugol']
-    mimetypes = []
-    url = "https://www.apoioinformatica.inf.br/produtos/visualg/linguagem"
-    version_added = ''
-
-    def __init__(self, **options):
-        Lexer.__init__(self, **options)
-        self.lexer = DelphiLexer(**options, portugol=True)
-
-    def get_tokens_unprocessed(self, text):
-        return self.lexer.get_tokens_unprocessed(text)
-
-
-class DelphiLexer(Lexer):
-    """
-    For Delphi (Borland Object Pascal),
-    Turbo Pascal and Free Pascal source code.
-
-    Additional options accepted:
-
-    `turbopascal`
-        Highlight Turbo Pascal specific keywords (default: ``True``).
-    `delphi`
-        Highlight Borland Delphi specific keywords (default: ``True``).
-    `freepascal`
-        Highlight Free Pascal specific keywords (default: ``True``).
-    `units`
-        A list of units that should be considered builtin, supported are
-        ``System``, ``SysUtils``, ``Classes`` and ``Math``.
-        Default is to consider all of them builtin.
-    """
-    name = 'Delphi'
-    aliases = ['delphi', 'pas', 'pascal', 'objectpascal']
-    filenames = ['*.pas', '*.dpr']
-    mimetypes = ['text/x-pascal']
-    url = 'https://www.embarcadero.com/products/delphi'
-    version_added = ''
-
-    TURBO_PASCAL_KEYWORDS = (
-        'absolute', 'and', 'array', 'asm', 'begin', 'break', 'case',
-        'const', 'constructor', 'continue', 'destructor', 'div', 'do',
-        'downto', 'else', 'end', 'file', 'for', 'function', 'goto',
-        'if', 'implementation', 'in', 'inherited', 'inline', 'interface',
-        'label', 'mod', 'nil', 'not', 'object', 'of', 'on', 'operator',
-        'or', 'packed', 'procedure', 'program', 'record', 'reintroduce',
-        'repeat', 'self', 'set', 'shl', 'shr', 'string', 'then', 'to',
-        'type', 'unit', 'until', 'uses', 'var', 'while', 'with', 'xor'
-    )
-
-    DELPHI_KEYWORDS = (
-        'as', 'class', 'except', 'exports', 'finalization', 'finally',
-        'initialization', 'is', 'library', 'on', 'property', 'raise',
-        'threadvar', 'try'
-    )
-
-    FREE_PASCAL_KEYWORDS = (
-        'dispose', 'exit', 'false', 'new', 'true'
-    )
-
-    BLOCK_KEYWORDS = {
-        'begin', 'class', 'const', 'constructor', 'destructor', 'end',
-        'finalization', 'function', 'implementation', 'initialization',
-        'label', 'library', 'operator', 'procedure', 'program', 'property',
-        'record', 'threadvar', 'type', 'unit', 'uses', 'var'
-    }
-
-    FUNCTION_MODIFIERS = {
-        'alias', 'cdecl', 'export', 'inline', 'interrupt', 'nostackframe',
-        'pascal', 'register', 'safecall', 'softfloat', 'stdcall',
-        'varargs', 'name', 'dynamic', 'near', 'virtual', 'external',
-        'override', 'assembler'
-    }
-
-    # XXX: those aren't global. but currently we know no way for defining
-    #      them just for the type context.
-    DIRECTIVES = {
-        'absolute', 'abstract', 'assembler', 'cppdecl', 'default', 'far',
-        'far16', 'forward', 'index', 'oldfpccall', 'private', 'protected',
-        'published', 'public'
-    }
-
-    BUILTIN_TYPES = {
-        'ansichar', 'ansistring', 'bool', 'boolean', 'byte', 'bytebool',
-        'cardinal', 'char', 'comp', 'currency', 'double', 'dword',
-        'extended', 'int64', 'integer', 'iunknown', 'longbool', 'longint',
-        'longword', 'pansichar', 'pansistring', 'pbool', 'pboolean',
-        'pbyte', 'pbytearray', 'pcardinal', 'pchar', 'pcomp', 'pcurrency',
-        'pdate', 'pdatetime', 'pdouble', 'pdword', 'pextended', 'phandle',
-        'pint64', 'pinteger', 'plongint', 'plongword', 'pointer',
-        'ppointer', 'pshortint', 'pshortstring', 'psingle', 'psmallint',
-        'pstring', 'pvariant', 'pwidechar', 'pwidestring', 'pword',
-        'pwordarray', 'pwordbool', 'real', 'real48', 'shortint',
-        'shortstring', 'single', 'smallint', 'string', 'tclass', 'tdate',
-        'tdatetime', 'textfile', 'thandle', 'tobject', 'ttime', 'variant',
-        'widechar', 'widestring', 'word', 'wordbool'
-    }
-
-    BUILTIN_UNITS = {
-        'System': (
-            'abs', 'acquireexceptionobject', 'addr', 'ansitoutf8',
-            'append', 'arctan', 'assert', 'assigned', 'assignfile',
-            'beginthread', 'blockread', 'blockwrite', 'break', 'chdir',
-            'chr', 'close', 'closefile', 'comptocurrency', 'comptodouble',
-            'concat', 'continue', 'copy', 'cos', 'dec', 'delete',
-            'dispose', 'doubletocomp', 'endthread', 'enummodules',
-            'enumresourcemodules', 'eof', 'eoln', 'erase', 'exceptaddr',
-            'exceptobject', 'exclude', 'exit', 'exp', 'filepos', 'filesize',
-            'fillchar', 'finalize', 'findclasshinstance', 'findhinstance',
-            'findresourcehinstance', 'flush', 'frac', 'freemem',
-            'get8087cw', 'getdir', 'getlasterror', 'getmem',
-            'getmemorymanager', 'getmodulefilename', 'getvariantmanager',
-            'halt', 'hi', 'high', 'inc', 'include', 'initialize', 'insert',
-            'int', 'ioresult', 'ismemorymanagerset', 'isvariantmanagerset',
-            'length', 'ln', 'lo', 'low', 'mkdir', 'move', 'new', 'odd',
-            'olestrtostring', 'olestrtostrvar', 'ord', 'paramcount',
-            'paramstr', 'pi', 'pos', 'pred', 'ptr', 'pucs4chars', 'random',
-            'randomize', 'read', 'readln', 'reallocmem',
-            'releaseexceptionobject', 'rename', 'reset', 'rewrite', 'rmdir',
-            'round', 'runerror', 'seek', 'seekeof', 'seekeoln',
-            'set8087cw', 'setlength', 'setlinebreakstyle',
-            'setmemorymanager', 'setstring', 'settextbuf',
-            'setvariantmanager', 'sin', 'sizeof', 'slice', 'sqr', 'sqrt',
-            'str', 'stringofchar', 'stringtoolestr', 'stringtowidechar',
-            'succ', 'swap', 'trunc', 'truncate', 'typeinfo',
-            'ucs4stringtowidestring', 'unicodetoutf8', 'uniquestring',
-            'upcase', 'utf8decode', 'utf8encode', 'utf8toansi',
-            'utf8tounicode', 'val', 'vararrayredim', 'varclear',
-            'widecharlentostring', 'widecharlentostrvar',
-            'widechartostring', 'widechartostrvar',
-            'widestringtoucs4string', 'write', 'writeln'
-        ),
-        'SysUtils': (
-            'abort', 'addexitproc', 'addterminateproc', 'adjustlinebreaks',
-            'allocmem', 'ansicomparefilename', 'ansicomparestr',
-            'ansicomparetext', 'ansidequotedstr', 'ansiextractquotedstr',
-            'ansilastchar', 'ansilowercase', 'ansilowercasefilename',
-            'ansipos', 'ansiquotedstr', 'ansisamestr', 'ansisametext',
-            'ansistrcomp', 'ansistricomp', 'ansistrlastchar', 'ansistrlcomp',
-            'ansistrlicomp', 'ansistrlower', 'ansistrpos', 'ansistrrscan',
-            'ansistrscan', 'ansistrupper', 'ansiuppercase',
-            'ansiuppercasefilename', 'appendstr', 'assignstr', 'beep',
-            'booltostr', 'bytetocharindex', 'bytetocharlen', 'bytetype',
-            'callterminateprocs', 'changefileext', 'charlength',
-            'chartobyteindex', 'chartobytelen', 'comparemem', 'comparestr',
-            'comparetext', 'createdir', 'createguid', 'currentyear',
-            'currtostr', 'currtostrf', 'date', 'datetimetofiledate',
-            'datetimetostr', 'datetimetostring', 'datetimetosystemtime',
-            'datetimetotimestamp', 'datetostr', 'dayofweek', 'decodedate',
-            'decodedatefully', 'decodetime', 'deletefile', 'directoryexists',
-            'diskfree', 'disksize', 'disposestr', 'encodedate', 'encodetime',
-            'exceptionerrormessage', 'excludetrailingbackslash',
-            'excludetrailingpathdelimiter', 'expandfilename',
-            'expandfilenamecase', 'expanduncfilename', 'extractfiledir',
-            'extractfiledrive', 'extractfileext', 'extractfilename',
-            'extractfilepath', 'extractrelativepath', 'extractshortpathname',
-            'fileage', 'fileclose', 'filecreate', 'filedatetodatetime',
-            'fileexists', 'filegetattr', 'filegetdate', 'fileisreadonly',
-            'fileopen', 'fileread', 'filesearch', 'fileseek', 'filesetattr',
-            'filesetdate', 'filesetreadonly', 'filewrite', 'finalizepackage',
-            'findclose', 'findcmdlineswitch', 'findfirst', 'findnext',
-            'floattocurr', 'floattodatetime', 'floattodecimal', 'floattostr',
-            'floattostrf', 'floattotext', 'floattotextfmt', 'fmtloadstr',
-            'fmtstr', 'forcedirectories', 'format', 'formatbuf', 'formatcurr',
-            'formatdatetime', 'formatfloat', 'freeandnil', 'getcurrentdir',
-            'getenvironmentvariable', 'getfileversion', 'getformatsettings',
-            'getlocaleformatsettings', 'getmodulename', 'getpackagedescription',
-            'getpackageinfo', 'gettime', 'guidtostring', 'incamonth',
-            'includetrailingbackslash', 'includetrailingpathdelimiter',
-            'incmonth', 'initializepackage', 'interlockeddecrement',
-            'interlockedexchange', 'interlockedexchangeadd',
-            'interlockedincrement', 'inttohex', 'inttostr', 'isdelimiter',
-            'isequalguid', 'isleapyear', 'ispathdelimiter', 'isvalidident',
-            'languages', 'lastdelimiter', 'loadpackage', 'loadstr',
-            'lowercase', 'msecstotimestamp', 'newstr', 'nextcharindex', 'now',
-            'outofmemoryerror', 'quotedstr', 'raiselastoserror',
-            'raiselastwin32error', 'removedir', 'renamefile', 'replacedate',
-            'replacetime', 'safeloadlibrary', 'samefilename', 'sametext',
-            'setcurrentdir', 'showexception', 'sleep', 'stralloc', 'strbufsize',
-            'strbytetype', 'strcat', 'strcharlength', 'strcomp', 'strcopy',
-            'strdispose', 'strecopy', 'strend', 'strfmt', 'stricomp',
-            'stringreplace', 'stringtoguid', 'strlcat', 'strlcomp', 'strlcopy',
-            'strlen', 'strlfmt', 'strlicomp', 'strlower', 'strmove', 'strnew',
-            'strnextchar', 'strpas', 'strpcopy', 'strplcopy', 'strpos',
-            'strrscan', 'strscan', 'strtobool', 'strtobooldef', 'strtocurr',
-            'strtocurrdef', 'strtodate', 'strtodatedef', 'strtodatetime',
-            'strtodatetimedef', 'strtofloat', 'strtofloatdef', 'strtoint',
-            'strtoint64', 'strtoint64def', 'strtointdef', 'strtotime',
-            'strtotimedef', 'strupper', 'supports', 'syserrormessage',
-            'systemtimetodatetime', 'texttofloat', 'time', 'timestamptodatetime',
-            'timestamptomsecs', 'timetostr', 'trim', 'trimleft', 'trimright',
-            'tryencodedate', 'tryencodetime', 'tryfloattocurr', 'tryfloattodatetime',
-            'trystrtobool', 'trystrtocurr', 'trystrtodate', 'trystrtodatetime',
-            'trystrtofloat', 'trystrtoint', 'trystrtoint64', 'trystrtotime',
-            'unloadpackage', 'uppercase', 'widecomparestr', 'widecomparetext',
-            'widefmtstr', 'wideformat', 'wideformatbuf', 'widelowercase',
-            'widesamestr', 'widesametext', 'wideuppercase', 'win32check',
-            'wraptext'
-        ),
-        'Classes': (
-            'activateclassgroup', 'allocatehwnd', 'bintohex', 'checksynchronize',
-            'collectionsequal', 'countgenerations', 'deallocatehwnd', 'equalrect',
-            'extractstrings', 'findclass', 'findglobalcomponent', 'getclass',
-            'groupdescendantswith', 'hextobin', 'identtoint',
-            'initinheritedcomponent', 'inttoident', 'invalidpoint',
-            'isuniqueglobalcomponentname', 'linestart', 'objectbinarytotext',
-            'objectresourcetotext', 'objecttexttobinary', 'objecttexttoresource',
-            'pointsequal', 'readcomponentres', 'readcomponentresex',
-            'readcomponentresfile', 'rect', 'registerclass', 'registerclassalias',
-            'registerclasses', 'registercomponents', 'registerintegerconsts',
-            'registernoicon', 'registernonactivex', 'smallpoint', 'startclassgroup',
-            'teststreamformat', 'unregisterclass', 'unregisterclasses',
-            'unregisterintegerconsts', 'unregistermoduleclasses',
-            'writecomponentresfile'
-        ),
-        'Math': (
-            'arccos', 'arccosh', 'arccot', 'arccoth', 'arccsc', 'arccsch', 'arcsec',
-            'arcsech', 'arcsin', 'arcsinh', 'arctan2', 'arctanh', 'ceil',
-            'comparevalue', 'cosecant', 'cosh', 'cot', 'cotan', 'coth', 'csc',
-            'csch', 'cycletodeg', 'cycletograd', 'cycletorad', 'degtocycle',
-            'degtograd', 'degtorad', 'divmod', 'doubledecliningbalance',
-            'ensurerange', 'floor', 'frexp', 'futurevalue', 'getexceptionmask',
-            'getprecisionmode', 'getroundmode', 'gradtocycle', 'gradtodeg',
-            'gradtorad', 'hypot', 'inrange', 'interestpayment', 'interestrate',
-            'internalrateofreturn', 'intpower', 'isinfinite', 'isnan', 'iszero',
-            'ldexp', 'lnxp1', 'log10', 'log2', 'logn', 'max', 'maxintvalue',
-            'maxvalue', 'mean', 'meanandstddev', 'min', 'minintvalue', 'minvalue',
-            'momentskewkurtosis', 'netpresentvalue', 'norm', 'numberofperiods',
-            'payment', 'periodpayment', 'poly', 'popnstddev', 'popnvariance',
-            'power', 'presentvalue', 'radtocycle', 'radtodeg', 'radtograd',
-            'randg', 'randomrange', 'roundto', 'samevalue', 'sec', 'secant',
-            'sech', 'setexceptionmask', 'setprecisionmode', 'setroundmode',
-            'sign', 'simpleroundto', 'sincos', 'sinh', 'slndepreciation', 'stddev',
-            'sum', 'sumint', 'sumofsquares', 'sumsandsquares', 'syddepreciation',
-            'tan', 'tanh', 'totalvariance', 'variance'
-        )
-    }
-
-    ASM_REGISTERS = {
-        'ah', 'al', 'ax', 'bh', 'bl', 'bp', 'bx', 'ch', 'cl', 'cr0',
-        'cr1', 'cr2', 'cr3', 'cr4', 'cs', 'cx', 'dh', 'di', 'dl', 'dr0',
-        'dr1', 'dr2', 'dr3', 'dr4', 'dr5', 'dr6', 'dr7', 'ds', 'dx',
-        'eax', 'ebp', 'ebx', 'ecx', 'edi', 'edx', 'es', 'esi', 'esp',
-        'fs', 'gs', 'mm0', 'mm1', 'mm2', 'mm3', 'mm4', 'mm5', 'mm6',
-        'mm7', 'si', 'sp', 'ss', 'st0', 'st1', 'st2', 'st3', 'st4', 'st5',
-        'st6', 'st7', 'xmm0', 'xmm1', 'xmm2', 'xmm3', 'xmm4', 'xmm5',
-        'xmm6', 'xmm7'
-    }
-
-    ASM_INSTRUCTIONS = {
-        'aaa', 'aad', 'aam', 'aas', 'adc', 'add', 'and', 'arpl', 'bound',
-        'bsf', 'bsr', 'bswap', 'bt', 'btc', 'btr', 'bts', 'call', 'cbw',
-        'cdq', 'clc', 'cld', 'cli', 'clts', 'cmc', 'cmova', 'cmovae',
-        'cmovb', 'cmovbe', 'cmovc', 'cmovcxz', 'cmove', 'cmovg',
-        'cmovge', 'cmovl', 'cmovle', 'cmovna', 'cmovnae', 'cmovnb',
-        'cmovnbe', 'cmovnc', 'cmovne', 'cmovng', 'cmovnge', 'cmovnl',
-        'cmovnle', 'cmovno', 'cmovnp', 'cmovns', 'cmovnz', 'cmovo',
-        'cmovp', 'cmovpe', 'cmovpo', 'cmovs', 'cmovz', 'cmp', 'cmpsb',
-        'cmpsd', 'cmpsw', 'cmpxchg', 'cmpxchg486', 'cmpxchg8b', 'cpuid',
-        'cwd', 'cwde', 'daa', 'das', 'dec', 'div', 'emms', 'enter', 'hlt',
-        'ibts', 'icebp', 'idiv', 'imul', 'in', 'inc', 'insb', 'insd',
-        'insw', 'int', 'int01', 'int03', 'int1', 'int3', 'into', 'invd',
-        'invlpg', 'iret', 'iretd', 'iretw', 'ja', 'jae', 'jb', 'jbe',
-        'jc', 'jcxz', 'jcxz', 'je', 'jecxz', 'jg', 'jge', 'jl', 'jle',
-        'jmp', 'jna', 'jnae', 'jnb', 'jnbe', 'jnc', 'jne', 'jng', 'jnge',
-        'jnl', 'jnle', 'jno', 'jnp', 'jns', 'jnz', 'jo', 'jp', 'jpe',
-        'jpo', 'js', 'jz', 'lahf', 'lar', 'lcall', 'lds', 'lea', 'leave',
-        'les', 'lfs', 'lgdt', 'lgs', 'lidt', 'ljmp', 'lldt', 'lmsw',
-        'loadall', 'loadall286', 'lock', 'lodsb', 'lodsd', 'lodsw',
-        'loop', 'loope', 'loopne', 'loopnz', 'loopz', 'lsl', 'lss', 'ltr',
-        'mov', 'movd', 'movq', 'movsb', 'movsd', 'movsw', 'movsx',
-        'movzx', 'mul', 'neg', 'nop', 'not', 'or', 'out', 'outsb', 'outsd',
-        'outsw', 'pop', 'popa', 'popad', 'popaw', 'popf', 'popfd', 'popfw',
-        'push', 'pusha', 'pushad', 'pushaw', 'pushf', 'pushfd', 'pushfw',
-        'rcl', 'rcr', 'rdmsr', 'rdpmc', 'rdshr', 'rdtsc', 'rep', 'repe',
-        'repne', 'repnz', 'repz', 'ret', 'retf', 'retn', 'rol', 'ror',
-        'rsdc', 'rsldt', 'rsm', 'sahf', 'sal', 'salc', 'sar', 'sbb',
-        'scasb', 'scasd', 'scasw', 'seta', 'setae', 'setb', 'setbe',
-        'setc', 'setcxz', 'sete', 'setg', 'setge', 'setl', 'setle',
-        'setna', 'setnae', 'setnb', 'setnbe', 'setnc', 'setne', 'setng',
-        'setnge', 'setnl', 'setnle', 'setno', 'setnp', 'setns', 'setnz',
-        'seto', 'setp', 'setpe', 'setpo', 'sets', 'setz', 'sgdt', 'shl',
-        'shld', 'shr', 'shrd', 'sidt', 'sldt', 'smi', 'smint', 'smintold',
-        'smsw', 'stc', 'std', 'sti', 'stosb', 'stosd', 'stosw', 'str',
-        'sub', 'svdc', 'svldt', 'svts', 'syscall', 'sysenter', 'sysexit',
-        'sysret', 'test', 'ud1', 'ud2', 'umov', 'verr', 'verw', 'wait',
-        'wbinvd', 'wrmsr', 'wrshr', 'xadd', 'xbts', 'xchg', 'xlat',
-        'xlatb', 'xor'
-    }
-
-    PORTUGOL_KEYWORDS = (
-        'aleatorio',
-        'algoritmo',
-        'arquivo',
-        'ate',
-        'caso',
-        'cronometro',
-        'debug',
-        'e',
-        'eco',
-        'enquanto',
-        'entao',
-        'escolha',
-        'escreva',
-        'escreval',
-        'faca',
-        'falso',
-        'fimalgoritmo',
-        'fimenquanto',
-        'fimescolha',
-        'fimfuncao',
-        'fimpara',
-        'fimprocedimento',
-        'fimrepita',
-        'fimse',
-        'funcao',
-        'inicio',
-        'int',
-        'interrompa',
-        'leia',
-        'limpatela',
-        'mod',
-        'nao',
-        'ou',
-        'outrocaso',
-        'para',
-        'passo',
-        'pausa',
-        'procedimento',
-        'repita',
-        'retorne',
-        'se',
-        'senao',
-        'timer',
-        'var',
-        'vetor',
-        'verdadeiro',
-        'xou',
-        'div',
-        'mod',
-        'abs',
-        'arccos',
-        'arcsen',
-        'arctan',
-        'cos',
-        'cotan',
-        'Exp',
-        'grauprad',
-        'int',
-        'log',
-        'logn',
-        'pi',
-        'quad',
-        'radpgrau',
-        'raizq',
-        'rand',
-        'randi',
-        'sen',
-        'Tan',
-        'asc',
-        'carac',
-        'caracpnum',
-        'compr',
-        'copia',
-        'maiusc',
-        'minusc',
-        'numpcarac',
-        'pos',
-    )
-
-    PORTUGOL_BUILTIN_TYPES = {
-        'inteiro', 'real', 'caractere', 'logico'
-    }
-
-    def __init__(self, **options):
-        Lexer.__init__(self, **options)
-        self.keywords = set()
-        self.builtins = set()
-        if get_bool_opt(options, 'portugol', False):
-            self.keywords.update(self.PORTUGOL_KEYWORDS)
-            self.builtins.update(self.PORTUGOL_BUILTIN_TYPES)
-            self.is_portugol = True
-        else:
-            self.is_portugol = False
-
-            if get_bool_opt(options, 'turbopascal', True):
-                self.keywords.update(self.TURBO_PASCAL_KEYWORDS)
-            if get_bool_opt(options, 'delphi', True):
-                self.keywords.update(self.DELPHI_KEYWORDS)
-            if get_bool_opt(options, 'freepascal', True):
-                self.keywords.update(self.FREE_PASCAL_KEYWORDS)
-            for unit in get_list_opt(options, 'units', list(self.BUILTIN_UNITS)):
-                self.builtins.update(self.BUILTIN_UNITS[unit])
-
-    def get_tokens_unprocessed(self, text):
-        scanner = Scanner(text, re.DOTALL | re.MULTILINE | re.IGNORECASE)
-        stack = ['initial']
-        in_function_block = False
-        in_property_block = False
-        was_dot = False
-        next_token_is_function = False
-        next_token_is_property = False
-        collect_labels = False
-        block_labels = set()
-        brace_balance = [0, 0]
-
-        while not scanner.eos:
-            token = Error
-
-            if stack[-1] == 'initial':
-                if scanner.scan(r'\s+'):
-                    token = Whitespace
-                elif not self.is_portugol and scanner.scan(r'\{.*?\}|\(\*.*?\*\)'):
-                    if scanner.match.startswith('$'):
-                        token = Comment.Preproc
-                    else:
-                        token = Comment.Multiline
-                elif scanner.scan(r'//.*?$'):
-                    token = Comment.Single
-                elif self.is_portugol and scanner.scan(r'(<\-)|(>=)|(<=)|%|<|>|-|\+|\*|\=|(<>)|\/|\.|:|,'):
-                    token = Operator
-                elif not self.is_portugol and scanner.scan(r'[-+*\/=<>:;,.@\^]'):
-                    token = Operator
-                    # stop label highlighting on next ";"
-                    if collect_labels and scanner.match == ';':
-                        collect_labels = False
-                elif scanner.scan(r'[\(\)\[\]]+'):
-                    token = Punctuation
-                    # abort function naming ``foo = Function(...)``
-                    next_token_is_function = False
-                    # if we are in a function block we count the open
-                    # braces because ootherwise it's impossible to
-                    # determine the end of the modifier context
-                    if in_function_block or in_property_block:
-                        if scanner.match == '(':
-                            brace_balance[0] += 1
-                        elif scanner.match == ')':
-                            brace_balance[0] -= 1
-                        elif scanner.match == '[':
-                            brace_balance[1] += 1
-                        elif scanner.match == ']':
-                            brace_balance[1] -= 1
-                elif scanner.scan(r'[A-Za-z_][A-Za-z_0-9]*'):
-                    lowercase_name = scanner.match.lower()
-                    if lowercase_name == 'result':
-                        token = Name.Builtin.Pseudo
-                    elif lowercase_name in self.keywords:
-                        token = Keyword
-                        # if we are in a special block and a
-                        # block ending keyword occurs (and the parenthesis
-                        # is balanced) we end the current block context
-                        if self.is_portugol:
-                            if lowercase_name in ('funcao', 'procedimento'):
-                                in_function_block = True
-                                next_token_is_function = True
-                        else:
-                            if (in_function_block or in_property_block) and \
-                                    lowercase_name in self.BLOCK_KEYWORDS and \
-                                    brace_balance[0] <= 0 and \
-                                    brace_balance[1] <= 0:
-                                in_function_block = False
-                                in_property_block = False
-                                brace_balance = [0, 0]
-                                block_labels = set()
-                            if lowercase_name in ('label', 'goto'):
-                                collect_labels = True
-                            elif lowercase_name == 'asm':
-                                stack.append('asm')
-                            elif lowercase_name == 'property':
-                                in_property_block = True
-                                next_token_is_property = True
-                            elif lowercase_name in ('procedure', 'operator',
-                                                    'function', 'constructor',
-                                                    'destructor'):
-                                in_function_block = True
-                                next_token_is_function = True
-                    # we are in a function block and the current name
-                    # is in the set of registered modifiers. highlight
-                    # it as pseudo keyword
-                    elif not self.is_portugol and in_function_block and \
-                            lowercase_name in self.FUNCTION_MODIFIERS:
-                        token = Keyword.Pseudo
-                    # if we are in a property highlight some more
-                    # modifiers
-                    elif not self.is_portugol and in_property_block and \
-                            lowercase_name in ('read', 'write'):
-                        token = Keyword.Pseudo
-                        next_token_is_function = True
-                    # if the last iteration set next_token_is_function
-                    # to true we now want this name highlighted as
-                    # function. so do that and reset the state
-                    elif next_token_is_function:
-                        # Look if the next token is a dot. If yes it's
-                        # not a function, but a class name and the
-                        # part after the dot a function name
-                        if not self.is_portugol and scanner.test(r'\s*\.\s*'):
-                            token = Name.Class
-                        # it's not a dot, our job is done
-                        else:
-                            token = Name.Function
-                            next_token_is_function = False
-
-                            if self.is_portugol:
-                                block_labels.add(scanner.match.lower())
-
-                    # same for properties
-                    elif not self.is_portugol and next_token_is_property:
-                        token = Name.Property
-                        next_token_is_property = False
-                    # Highlight this token as label and add it
-                    # to the list of known labels
-                    elif not self.is_portugol and collect_labels:
-                        token = Name.Label
-                        block_labels.add(scanner.match.lower())
-                    # name is in list of known labels
-                    elif lowercase_name in block_labels:
-                        token = Name.Label
-                    elif self.is_portugol and lowercase_name in self.PORTUGOL_BUILTIN_TYPES:
-                        token = Keyword.Type
-                    elif not self.is_portugol and lowercase_name in self.BUILTIN_TYPES:
-                        token = Keyword.Type
-                    elif not self.is_portugol and lowercase_name in self.DIRECTIVES:
-                        token = Keyword.Pseudo
-                    # builtins are just builtins if the token
-                    # before isn't a dot
-                    elif not self.is_portugol and not was_dot and lowercase_name in self.builtins:
-                        token = Name.Builtin
-                    else:
-                        token = Name
-                elif self.is_portugol and scanner.scan(r"\""):
-                    token = String
-                    stack.append('string')
-                elif not self.is_portugol and scanner.scan(r"'"):
-                    token = String
-                    stack.append('string')
-                elif not self.is_portugol and scanner.scan(r'\#(\d+|\$[0-9A-Fa-f]+)'):
-                    token = String.Char
-                elif not self.is_portugol and scanner.scan(r'\$[0-9A-Fa-f]+'):
-                    token = Number.Hex
-                elif scanner.scan(r'\d+(?![eE]|\.[^.])'):
-                    token = Number.Integer
-                elif scanner.scan(r'\d+(\.\d+([eE][+-]?\d+)?|[eE][+-]?\d+)'):
-                    token = Number.Float
-                else:
-                    # if the stack depth is deeper than once, pop
-                    if len(stack) > 1:
-                        stack.pop()
-                    scanner.get_char()
-
-            elif stack[-1] == 'string':
-                if self.is_portugol:
-                    if scanner.scan(r"''"):
-                        token = String.Escape
-                    elif scanner.scan(r"\""):
-                        token = String
-                        stack.pop()
-                    elif scanner.scan(r"[^\"]*"):
-                        token = String
-                    else:
-                        scanner.get_char()
-                        stack.pop()
-                else:
-                    if scanner.scan(r"''"):
-                        token = String.Escape
-                    elif scanner.scan(r"'"):
-                        token = String
-                        stack.pop()
-                    elif scanner.scan(r"[^']*"):
-                        token = String
-                    else:
-                        scanner.get_char()
-                        stack.pop()
-            elif not self.is_portugol and stack[-1] == 'asm':
-                if scanner.scan(r'\s+'):
-                    token = Whitespace
-                elif scanner.scan(r'end'):
-                    token = Keyword
-                    stack.pop()
-                elif scanner.scan(r'\{.*?\}|\(\*.*?\*\)'):
-                    if scanner.match.startswith('$'):
-                        token = Comment.Preproc
-                    else:
-                        token = Comment.Multiline
-                elif scanner.scan(r'//.*?$'):
-                    token = Comment.Single
-                elif scanner.scan(r"'"):
-                    token = String
-                    stack.append('string')
-                elif scanner.scan(r'@@[A-Za-z_][A-Za-z_0-9]*'):
-                    token = Name.Label
-                elif scanner.scan(r'[A-Za-z_][A-Za-z_0-9]*'):
-                    lowercase_name = scanner.match.lower()
-                    if lowercase_name in self.ASM_INSTRUCTIONS:
-                        token = Keyword
-                    elif lowercase_name in self.ASM_REGISTERS:
-                        token = Name.Builtin
-                    else:
-                        token = Name
-                elif scanner.scan(r'[-+*\/=<>:;,.@\^]+'):
-                    token = Operator
-                elif scanner.scan(r'[\(\)\[\]]+'):
-                    token = Punctuation
-                elif scanner.scan(r'\$[0-9A-Fa-f]+'):
-                    token = Number.Hex
-                elif scanner.scan(r'\d+(?![eE]|\.[^.])'):
-                    token = Number.Integer
-                elif scanner.scan(r'\d+(\.\d+([eE][+-]?\d+)?|[eE][+-]?\d+)'):
-                    token = Number.Float
-                else:
-                    scanner.get_char()
-                    stack.pop()
-
-            # save the dot!!!11
-            if not self.is_portugol and scanner.match.strip():
-                was_dot = scanner.match == '.'
-
-            yield scanner.start_pos, token, scanner.match or ''
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/pawn.py b/venv/lib/python3.10/site-packages/pygments/lexers/pawn.py
deleted file mode 100644
index 99d9c963d1473206ca68c389b068ba783783ce92..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/pawn.py
+++ /dev/null
@@ -1,202 +0,0 @@
-"""
-    pygments.lexers.pawn
-    ~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for the Pawn languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation
-from pygments.util import get_bool_opt
-
-__all__ = ['SourcePawnLexer', 'PawnLexer']
-
-
-class SourcePawnLexer(RegexLexer):
-    """
-    For SourcePawn source code with preprocessor directives.
-    """
-    name = 'SourcePawn'
-    aliases = ['sp']
-    filenames = ['*.sp']
-    mimetypes = ['text/x-sourcepawn']
-    url = 'https://github.com/alliedmodders/sourcepawn'
-    version_added = '1.6'
-
-    #: optional Comment or Whitespace
-    _ws = r'(?:\s|//.*?\n|/\*.*?\*/)+'
-    #: only one /* */ style comment
-    _ws1 = r'\s*(?:/[*].*?[*]/\s*)*'
-
-    tokens = {
-        'root': [
-            # preprocessor directives: without whitespace
-            (r'^#if\s+0', Comment.Preproc, 'if0'),
-            ('^#', Comment.Preproc, 'macro'),
-            # or with whitespace
-            ('^' + _ws1 + r'#if\s+0', Comment.Preproc, 'if0'),
-            ('^' + _ws1 + '#', Comment.Preproc, 'macro'),
-            (r'\n', Text),
-            (r'\s+', Text),
-            (r'\\\n', Text),  # line continuation
-            (r'/(\\\n)?/(\n|(.|\n)*?[^\\]\n)', Comment.Single),
-            (r'/(\\\n)?\*(.|\n)*?\*(\\\n)?/', Comment.Multiline),
-            (r'[{}]', Punctuation),
-            (r'L?"', String, 'string'),
-            (r"L?'(\\.|\\[0-7]{1,3}|\\x[a-fA-F0-9]{1,2}|[^\\\'\n])'", String.Char),
-            (r'(\d+\.\d*|\.\d+|\d+)[eE][+-]?\d+[LlUu]*', Number.Float),
-            (r'(\d+\.\d*|\.\d+|\d+[fF])[fF]?', Number.Float),
-            (r'0x[0-9a-fA-F]+[LlUu]*', Number.Hex),
-            (r'0[0-7]+[LlUu]*', Number.Oct),
-            (r'\d+[LlUu]*', Number.Integer),
-            (r'[~!%^&*+=|?:<>/-]', Operator),
-            (r'[()\[\],.;]', Punctuation),
-            (r'(case|const|continue|native|'
-             r'default|else|enum|for|if|new|operator|'
-             r'public|return|sizeof|static|decl|struct|switch)\b', Keyword),
-            (r'(bool|Float)\b', Keyword.Type),
-            (r'(true|false)\b', Keyword.Constant),
-            (r'[a-zA-Z_]\w*', Name),
-        ],
-        'string': [
-            (r'"', String, '#pop'),
-            (r'\\([\\abfnrtv"\']|x[a-fA-F0-9]{2,4}|[0-7]{1,3})', String.Escape),
-            (r'[^\\"\n]+', String),  # all other characters
-            (r'\\\n', String),       # line continuation
-            (r'\\', String),         # stray backslash
-        ],
-        'macro': [
-            (r'[^/\n]+', Comment.Preproc),
-            (r'/\*(.|\n)*?\*/', Comment.Multiline),
-            (r'//.*?\n', Comment.Single, '#pop'),
-            (r'/', Comment.Preproc),
-            (r'(?<=\\)\n', Comment.Preproc),
-            (r'\n', Comment.Preproc, '#pop'),
-        ],
-        'if0': [
-            (r'^\s*#if.*?(?/-]', Operator),
-            (r'[()\[\],.;]', Punctuation),
-            (r'(switch|case|default|const|new|static|char|continue|break|'
-             r'if|else|for|while|do|operator|enum|'
-             r'public|return|sizeof|tagof|state|goto)\b', Keyword),
-            (r'(bool|Float)\b', Keyword.Type),
-            (r'(true|false)\b', Keyword.Constant),
-            (r'[a-zA-Z_]\w*', Name),
-        ],
-        'string': [
-            (r'"', String, '#pop'),
-            (r'\\([\\abfnrtv"\']|x[a-fA-F0-9]{2,4}|[0-7]{1,3})', String.Escape),
-            (r'[^\\"\n]+', String),  # all other characters
-            (r'\\\n', String),       # line continuation
-            (r'\\', String),         # stray backslash
-        ],
-        'macro': [
-            (r'[^/\n]+', Comment.Preproc),
-            (r'/\*(.|\n)*?\*/', Comment.Multiline),
-            (r'//.*?\n', Comment.Single, '#pop'),
-            (r'/', Comment.Preproc),
-            (r'(?<=\\)\n', Comment.Preproc),
-            (r'\n', Comment.Preproc, '#pop'),
-        ],
-        'if0': [
-            (r'^\s*#if.*?(?<-]', Operator),
-            (r'[a-zA-Z][a-zA-Z0-9_-]*', Name),
-            (r'\?[a-zA-Z][a-zA-Z0-9_-]*', Name.Variable),
-            (r'[0-9]+\.[0-9]+', Number.Float),
-            (r'[0-9]+', Number.Integer),
-        ],
-        'keywords': [
-            (words((
-                ':requirements', ':types', ':constants',
-                ':predicates', ':functions', ':action', ':agent',
-                ':parameters', ':precondition', ':effect',
-                ':durative-action', ':duration', ':condition',
-                ':derived', ':domain', ':objects', ':init',
-                ':goal', ':metric', ':length', ':serial', ':parallel',
-                # the following are requirements
-                ':strips', ':typing', ':negative-preconditions',
-                ':disjunctive-preconditions', ':equality',
-                ':existential-preconditions', ':universal-preconditions',
-                ':conditional-effects', ':fluents', ':numeric-fluents',
-                ':object-fluents', ':adl', ':durative-actions',
-                ':continuous-effects', ':derived-predicates',
-                ':time-intial-literals', ':preferences',
-                ':constraints', ':action-costs', ':multi-agent',
-                ':unfactored-privacy', ':factored-privacy',
-                ':non-deterministic'
-                ), suffix=r'\b'), Keyword)
-        ],
-        'builtins': [
-            (words((
-                'define', 'domain', 'object', 'either', 'and',
-                'forall', 'preference', 'imply', 'or', 'exists',
-                'not', 'when', 'assign', 'scale-up', 'scale-down',
-                'increase', 'decrease', 'at', 'over', 'start',
-                'end', 'all', 'problem', 'always', 'sometime',
-                'within', 'at-most-once', 'sometime-after',
-                'sometime-before', 'always-within', 'hold-during',
-                'hold-after', 'minimize', 'maximize',
-                'total-time', 'is-violated'), suffix=r'\b'),
-                Name.Builtin)
-        ]
-    }
-
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/perl.py b/venv/lib/python3.10/site-packages/pygments/lexers/perl.py
deleted file mode 100644
index 33f91f5800b7e40c8bb195f11974454eb3e6aabf..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/perl.py
+++ /dev/null
@@ -1,733 +0,0 @@
-"""
-    pygments.lexers.perl
-    ~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Perl, Raku and related languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, ExtendedRegexLexer, include, bygroups, \
-    using, this, default, words
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Whitespace
-from pygments.util import shebang_matches
-
-__all__ = ['PerlLexer', 'Perl6Lexer']
-
-
-class PerlLexer(RegexLexer):
-    """
-    For Perl source code.
-    """
-
-    name = 'Perl'
-    url = 'https://www.perl.org'
-    aliases = ['perl', 'pl']
-    filenames = ['*.pl', '*.pm', '*.t', '*.perl']
-    mimetypes = ['text/x-perl', 'application/x-perl']
-    version_added = ''
-
-    flags = re.DOTALL | re.MULTILINE
-    # TODO: give this to a perl guy who knows how to parse perl...
-    tokens = {
-        'balanced-regex': [
-            (r'/(\\\\|\\[^\\]|[^\\/])*/[egimosx]*', String.Regex, '#pop'),
-            (r'!(\\\\|\\[^\\]|[^\\!])*![egimosx]*', String.Regex, '#pop'),
-            (r'\\(\\\\|[^\\])*\\[egimosx]*', String.Regex, '#pop'),
-            (r'\{(\\\\|\\[^\\]|[^\\}])*\}[egimosx]*', String.Regex, '#pop'),
-            (r'<(\\\\|\\[^\\]|[^\\>])*>[egimosx]*', String.Regex, '#pop'),
-            (r'\[(\\\\|\\[^\\]|[^\\\]])*\][egimosx]*', String.Regex, '#pop'),
-            (r'\((\\\\|\\[^\\]|[^\\)])*\)[egimosx]*', String.Regex, '#pop'),
-            (r'@(\\\\|\\[^\\]|[^\\@])*@[egimosx]*', String.Regex, '#pop'),
-            (r'%(\\\\|\\[^\\]|[^\\%])*%[egimosx]*', String.Regex, '#pop'),
-            (r'\$(\\\\|\\[^\\]|[^\\$])*\$[egimosx]*', String.Regex, '#pop'),
-        ],
-        'root': [
-            (r'\A\#!.+?$', Comment.Hashbang),
-            (r'\#.*?$', Comment.Single),
-            (r'^=[a-zA-Z0-9]+\s+.*?\n=cut', Comment.Multiline),
-            (words((
-                'case', 'continue', 'do', 'else', 'elsif', 'for', 'foreach',
-                'if', 'last', 'my', 'next', 'our', 'redo', 'reset', 'then',
-                'unless', 'until', 'while', 'print', 'new', 'BEGIN',
-                'CHECK', 'INIT', 'END', 'return'), suffix=r'\b'),
-             Keyword),
-            (r'(format)(\s+)(\w+)(\s*)(=)(\s*\n)',
-             bygroups(Keyword, Whitespace, Name, Whitespace, Punctuation, Whitespace), 'format'),
-            (r'(eq|lt|gt|le|ge|ne|not|and|or|cmp)\b', Operator.Word),
-            # common delimiters
-            (r's/(\\\\|\\[^\\]|[^\\/])*/(\\\\|\\[^\\]|[^\\/])*/[egimosx]*',
-                String.Regex),
-            (r's!(\\\\|\\!|[^!])*!(\\\\|\\!|[^!])*![egimosx]*', String.Regex),
-            (r's\\(\\\\|[^\\])*\\(\\\\|[^\\])*\\[egimosx]*', String.Regex),
-            (r's@(\\\\|\\[^\\]|[^\\@])*@(\\\\|\\[^\\]|[^\\@])*@[egimosx]*',
-                String.Regex),
-            (r's%(\\\\|\\[^\\]|[^\\%])*%(\\\\|\\[^\\]|[^\\%])*%[egimosx]*',
-                String.Regex),
-            # balanced delimiters
-            (r's\{(\\\\|\\[^\\]|[^\\}])*\}\s*', String.Regex, 'balanced-regex'),
-            (r's<(\\\\|\\[^\\]|[^\\>])*>\s*', String.Regex, 'balanced-regex'),
-            (r's\[(\\\\|\\[^\\]|[^\\\]])*\]\s*', String.Regex,
-                'balanced-regex'),
-            (r's\((\\\\|\\[^\\]|[^\\)])*\)\s*', String.Regex,
-                'balanced-regex'),
-
-            (r'm?/(\\\\|\\[^\\]|[^\\/\n])*/[gcimosx]*', String.Regex),
-            (r'm(?=[/!\\{<\[(@%$])', String.Regex, 'balanced-regex'),
-            (r'((?<==~)|(?<=\())\s*/(\\\\|\\[^\\]|[^\\/])*/[gcimosx]*',
-                String.Regex),
-            (r'\s+', Whitespace),
-            (words((
-                'abs', 'accept', 'alarm', 'atan2', 'bind', 'binmode', 'bless', 'caller', 'chdir',
-                'chmod', 'chomp', 'chop', 'chown', 'chr', 'chroot', 'close', 'closedir', 'connect',
-                'continue', 'cos', 'crypt', 'dbmclose', 'dbmopen', 'defined', 'delete', 'die',
-                'dump', 'each', 'endgrent', 'endhostent', 'endnetent', 'endprotoent',
-                'endpwent', 'endservent', 'eof', 'eval', 'exec', 'exists', 'exit', 'exp', 'fcntl',
-                'fileno', 'flock', 'fork', 'format', 'formline', 'getc', 'getgrent', 'getgrgid',
-                'getgrnam', 'gethostbyaddr', 'gethostbyname', 'gethostent', 'getlogin',
-                'getnetbyaddr', 'getnetbyname', 'getnetent', 'getpeername', 'getpgrp',
-                'getppid', 'getpriority', 'getprotobyname', 'getprotobynumber',
-                'getprotoent', 'getpwent', 'getpwnam', 'getpwuid', 'getservbyname',
-                'getservbyport', 'getservent', 'getsockname', 'getsockopt', 'glob', 'gmtime',
-                'goto', 'grep', 'hex', 'import', 'index', 'int', 'ioctl', 'join', 'keys', 'kill', 'last',
-                'lc', 'lcfirst', 'length', 'link', 'listen', 'local', 'localtime', 'log', 'lstat',
-                'map', 'mkdir', 'msgctl', 'msgget', 'msgrcv', 'msgsnd', 'my', 'next', 'oct', 'open',
-                'opendir', 'ord', 'our', 'pack', 'pipe', 'pop', 'pos', 'printf',
-                'prototype', 'push', 'quotemeta', 'rand', 'read', 'readdir',
-                'readline', 'readlink', 'readpipe', 'recv', 'redo', 'ref', 'rename',
-                'reverse', 'rewinddir', 'rindex', 'rmdir', 'scalar', 'seek', 'seekdir',
-                'select', 'semctl', 'semget', 'semop', 'send', 'setgrent', 'sethostent', 'setnetent',
-                'setpgrp', 'setpriority', 'setprotoent', 'setpwent', 'setservent',
-                'setsockopt', 'shift', 'shmctl', 'shmget', 'shmread', 'shmwrite', 'shutdown',
-                'sin', 'sleep', 'socket', 'socketpair', 'sort', 'splice', 'split', 'sprintf', 'sqrt',
-                'srand', 'stat', 'study', 'substr', 'symlink', 'syscall', 'sysopen', 'sysread',
-                'sysseek', 'system', 'syswrite', 'tell', 'telldir', 'tie', 'tied', 'time', 'times', 'tr',
-                'truncate', 'uc', 'ucfirst', 'umask', 'undef', 'unlink', 'unpack', 'unshift', 'untie',
-                'utime', 'values', 'vec', 'wait', 'waitpid', 'wantarray', 'warn', 'write'), suffix=r'\b'),
-             Name.Builtin),
-            (r'((__(DATA|DIE|WARN)__)|(STD(IN|OUT|ERR)))\b', Name.Builtin.Pseudo),
-            (r'(<<)([\'"]?)([a-zA-Z_]\w*)(\2;?\n.*?\n)(\3)(\n)',
-             bygroups(String, String, String.Delimiter, String, String.Delimiter, Whitespace)),
-            (r'__END__', Comment.Preproc, 'end-part'),
-            (r'\$\^[ADEFHILMOPSTWX]', Name.Variable.Global),
-            (r"\$[\\\"\[\]'&`+*.,;=%~?@$!<>(^|/-](?!\w)", Name.Variable.Global),
-            (r'[$@%#]+', Name.Variable, 'varname'),
-            (r'0_?[0-7]+(_[0-7]+)*', Number.Oct),
-            (r'0x[0-9A-Fa-f]+(_[0-9A-Fa-f]+)*', Number.Hex),
-            (r'0b[01]+(_[01]+)*', Number.Bin),
-            (r'(?i)(\d*(_\d*)*\.\d+(_\d*)*|\d+(_\d*)*\.\d+(_\d*)*)(e[+-]?\d+)?',
-             Number.Float),
-            (r'(?i)\d+(_\d*)*e[+-]?\d+(_\d*)*', Number.Float),
-            (r'\d+(_\d+)*', Number.Integer),
-            (r"'(\\\\|\\[^\\]|[^'\\])*'", String),
-            (r'"(\\\\|\\[^\\]|[^"\\])*"', String),
-            (r'`(\\\\|\\[^\\]|[^`\\])*`', String.Backtick),
-            (r'<([^\s>]+)>', String.Regex),
-            (r'(q|qq|qw|qr|qx)\{', String.Other, 'cb-string'),
-            (r'(q|qq|qw|qr|qx)\(', String.Other, 'rb-string'),
-            (r'(q|qq|qw|qr|qx)\[', String.Other, 'sb-string'),
-            (r'(q|qq|qw|qr|qx)\<', String.Other, 'lt-string'),
-            (r'(q|qq|qw|qr|qx)([\W_])(.|\n)*?\2', String.Other),
-            (r'(package)(\s+)([a-zA-Z_]\w*(?:::[a-zA-Z_]\w*)*)',
-             bygroups(Keyword, Whitespace, Name.Namespace)),
-            (r'(use|require|no)(\s+)([a-zA-Z_]\w*(?:::[a-zA-Z_]\w*)*)',
-             bygroups(Keyword, Whitespace, Name.Namespace)),
-            (r'(sub)(\s+)', bygroups(Keyword, Whitespace), 'funcname'),
-            (words((
-                'no', 'package', 'require', 'use'), suffix=r'\b'),
-             Keyword),
-            (r'(\[\]|\*\*|::|<<|>>|>=|<=>|<=|={3}|!=|=~|'
-             r'!~|&&?|\|\||\.{1,3})', Operator),
-            (r'[-+/*%=<>&^|!\\~]=?', Operator),
-            (r'[()\[\]:;,<>/?{}]', Punctuation),  # yes, there's no shortage
-                                                  # of punctuation in Perl!
-            (r'(?=\w)', Name, 'name'),
-        ],
-        'format': [
-            (r'\.\n', String.Interpol, '#pop'),
-            (r'[^\n]*\n', String.Interpol),
-        ],
-        'varname': [
-            (r'\s+', Whitespace),
-            (r'\{', Punctuation, '#pop'),    # hash syntax?
-            (r'\)|,', Punctuation, '#pop'),  # argument specifier
-            (r'\w+::', Name.Namespace),
-            (r'[\w:]+', Name.Variable, '#pop'),
-        ],
-        'name': [
-            (r'[a-zA-Z_]\w*(::[a-zA-Z_]\w*)*(::)?(?=\s*->)', Name.Namespace, '#pop'),
-            (r'[a-zA-Z_]\w*(::[a-zA-Z_]\w*)*::', Name.Namespace, '#pop'),
-            (r'[\w:]+', Name, '#pop'),
-            (r'[A-Z_]+(?=\W)', Name.Constant, '#pop'),
-            (r'(?=\W)', Text, '#pop'),
-        ],
-        'funcname': [
-            (r'[a-zA-Z_]\w*[!?]?', Name.Function),
-            (r'\s+', Whitespace),
-            # argument declaration
-            (r'(\([$@%]*\))(\s*)', bygroups(Punctuation, Whitespace)),
-            (r';', Punctuation, '#pop'),
-            (r'.*?\{', Punctuation, '#pop'),
-        ],
-        'cb-string': [
-            (r'\\[{}\\]', String.Other),
-            (r'\\', String.Other),
-            (r'\{', String.Other, 'cb-string'),
-            (r'\}', String.Other, '#pop'),
-            (r'[^{}\\]+', String.Other)
-        ],
-        'rb-string': [
-            (r'\\[()\\]', String.Other),
-            (r'\\', String.Other),
-            (r'\(', String.Other, 'rb-string'),
-            (r'\)', String.Other, '#pop'),
-            (r'[^()]+', String.Other)
-        ],
-        'sb-string': [
-            (r'\\[\[\]\\]', String.Other),
-            (r'\\', String.Other),
-            (r'\[', String.Other, 'sb-string'),
-            (r'\]', String.Other, '#pop'),
-            (r'[^\[\]]+', String.Other)
-        ],
-        'lt-string': [
-            (r'\\[<>\\]', String.Other),
-            (r'\\', String.Other),
-            (r'\<', String.Other, 'lt-string'),
-            (r'\>', String.Other, '#pop'),
-            (r'[^<>]+', String.Other)
-        ],
-        'end-part': [
-            (r'.+', Comment.Preproc, '#pop')
-        ]
-    }
-
-    def analyse_text(text):
-        if shebang_matches(text, r'perl'):
-            return True
-
-        result = 0
-
-        if re.search(r'(?:my|our)\s+[$@%(]', text):
-            result += 0.9
-
-        if ':=' in text:
-            # := is not valid Perl, but it appears in unicon, so we should
-            # become less confident if we think we found Perl with :=
-            result /= 2
-
-        return result
-
-
-class Perl6Lexer(ExtendedRegexLexer):
-    """
-    For Raku (a.k.a. Perl 6) source code.
-    """
-
-    name = 'Perl6'
-    url = 'https://www.raku.org'
-    aliases = ['perl6', 'pl6', 'raku']
-    filenames = ['*.pl', '*.pm', '*.nqp', '*.p6', '*.6pl', '*.p6l', '*.pl6',
-                 '*.6pm', '*.p6m', '*.pm6', '*.t', '*.raku', '*.rakumod',
-                 '*.rakutest', '*.rakudoc']
-    mimetypes = ['text/x-perl6', 'application/x-perl6']
-    version_added = '2.0'
-    flags = re.MULTILINE | re.DOTALL
-
-    PERL6_IDENTIFIER_RANGE = r"['\w:-]"
-
-    PERL6_KEYWORDS = (
-        #Phasers
-        'BEGIN','CATCH','CHECK','CLOSE','CONTROL','DOC','END','ENTER','FIRST',
-        'INIT','KEEP','LAST','LEAVE','NEXT','POST','PRE','QUIT','UNDO',
-        #Keywords
-        'anon','augment','but','class','constant','default','does','else',
-        'elsif','enum','for','gather','given','grammar','has','if','import',
-        'is','let','loop','made','make','method','module','multi','my','need',
-        'orwith','our','proceed','proto','repeat','require','return',
-        'return-rw','returns','role','rule','state','sub','submethod','subset',
-        'succeed','supersede','token','try','unit','unless','until','use',
-        'when','while','with','without',
-        #Traits
-        'export','native','repr','required','rw','symbol',
-    )
-
-    PERL6_BUILTINS = (
-        'ACCEPTS','abs','abs2rel','absolute','accept','accessed','acos',
-        'acosec','acosech','acosh','acotan','acotanh','acquire','act','action',
-        'actions','add','add_attribute','add_enum_value','add_fallback',
-        'add_method','add_parent','add_private_method','add_role','add_trustee',
-        'adverb','after','all','allocate','allof','allowed','alternative-names',
-        'annotations','antipair','antipairs','any','anyof','app_lifetime',
-        'append','arch','archname','args','arity','Array','asec','asech','asin',
-        'asinh','ASSIGN-KEY','ASSIGN-POS','assuming','ast','at','atan','atan2',
-        'atanh','AT-KEY','atomic-assign','atomic-dec-fetch','atomic-fetch',
-        'atomic-fetch-add','atomic-fetch-dec','atomic-fetch-inc',
-        'atomic-fetch-sub','atomic-inc-fetch','AT-POS','attributes','auth',
-        'await','backtrace','Bag','BagHash','bail-out','base','basename',
-        'base-repeating','batch','BIND-KEY','BIND-POS','bind-stderr',
-        'bind-stdin','bind-stdout','bind-udp','bits','bless','block','Bool',
-        'bool-only','bounds','break','Bridge','broken','BUILD','build-date',
-        'bytes','cache','callframe','calling-package','CALL-ME','callsame',
-        'callwith','can','cancel','candidates','cando','can-ok','canonpath',
-        'caps','caption','Capture','cas','catdir','categorize','categorize-list',
-        'catfile','catpath','cause','ceiling','cglobal','changed','Channel',
-        'chars','chdir','child','child-name','child-typename','chmod','chomp',
-        'chop','chr','chrs','chunks','cis','classify','classify-list','cleanup',
-        'clone','close','closed','close-stdin','cmp-ok','code','codes','collate',
-        'column','comb','combinations','command','comment','compiler','Complex',
-        'compose','compose_type','composer','condition','config',
-        'configure_destroy','configure_type_checking','conj','connect',
-        'constraints','construct','contains','contents','copy','cos','cosec',
-        'cosech','cosh','cotan','cotanh','count','count-only','cpu-cores',
-        'cpu-usage','CREATE','create_type','cross','cue','curdir','curupdir','d',
-        'Date','DateTime','day','daycount','day-of-month','day-of-week',
-        'day-of-year','days-in-month','declaration','decode','decoder','deepmap',
-        'default','defined','DEFINITE','delayed','DELETE-KEY','DELETE-POS',
-        'denominator','desc','DESTROY','destroyers','devnull','diag',
-        'did-you-mean','die','dies-ok','dir','dirname','dir-sep','DISTROnames',
-        'do','does','does-ok','done','done-testing','duckmap','dynamic','e',
-        'eager','earlier','elems','emit','enclosing','encode','encoder',
-        'encoding','end','ends-with','enum_from_value','enum_value_list',
-        'enum_values','enums','eof','EVAL','eval-dies-ok','EVALFILE',
-        'eval-lives-ok','exception','excludes-max','excludes-min','EXISTS-KEY',
-        'EXISTS-POS','exit','exitcode','exp','expected','explicitly-manage',
-        'expmod','extension','f','fail','fails-like','fc','feature','file',
-        'filename','find_method','find_method_qualified','finish','first','flat',
-        'flatmap','flip','floor','flunk','flush','fmt','format','formatter',
-        'freeze','from','from-list','from-loop','from-posix','full',
-        'full-barrier','get','get_value','getc','gist','got','grab','grabpairs',
-        'grep','handle','handled','handles','hardware','has_accessor','Hash',
-        'head','headers','hh-mm-ss','hidden','hides','hour','how','hyper','id',
-        'illegal','im','in','indent','index','indices','indir','infinite',
-        'infix','infix:<+>','infix:<->','install_method_cache','Instant',
-        'instead','Int','int-bounds','interval','in-timezone','invalid-str',
-        'invert','invocant','IO','IO::Notification.watch-path','is_trusted',
-        'is_type','isa','is-absolute','isa-ok','is-approx','is-deeply',
-        'is-hidden','is-initial-thread','is-int','is-lazy','is-leap-year',
-        'isNaN','isnt','is-prime','is-relative','is-routine','is-setting',
-        'is-win','item','iterator','join','keep','kept','KERNELnames','key',
-        'keyof','keys','kill','kv','kxxv','l','lang','last','lastcall','later',
-        'lazy','lc','leading','level','like','line','lines','link','List',
-        'listen','live','lives-ok','local','lock','log','log10','lookup','lsb',
-        'made','MAIN','make','Map','match','max','maxpairs','merge','message',
-        'method','method_table','methods','migrate','min','minmax','minpairs',
-        'minute','misplaced','Mix','MixHash','mkdir','mode','modified','month',
-        'move','mro','msb','multi','multiness','my','name','named','named_names',
-        'narrow','nativecast','native-descriptor','nativesizeof','new','new_type',
-        'new-from-daycount','new-from-pairs','next','nextcallee','next-handle',
-        'nextsame','nextwith','NFC','NFD','NFKC','NFKD','nl-in','nl-out',
-        'nodemap','nok','none','norm','not','note','now','nude','Num',
-        'numerator','Numeric','of','offset','offset-in-hours','offset-in-minutes',
-        'ok','old','on-close','one','on-switch','open','opened','operation',
-        'optional','ord','ords','orig','os-error','osname','out-buffer','pack',
-        'package','package-kind','package-name','packages','pair','pairs',
-        'pairup','parameter','params','parent','parent-name','parents','parse',
-        'parse-base','parsefile','parse-names','parts','pass','path','path-sep',
-        'payload','peer-host','peer-port','periods','perl','permutations','phaser',
-        'pick','pickpairs','pid','placeholder','plan','plus','polar','poll',
-        'polymod','pop','pos','positional','posix','postfix','postmatch',
-        'precomp-ext','precomp-target','pred','prefix','prematch','prepend',
-        'print','printf','print-nl','print-to','private','private_method_table',
-        'proc','produce','Promise','prompt','protect','pull-one','push',
-        'push-all','push-at-least','push-exactly','push-until-lazy','put',
-        'qualifier-type','quit','r','race','radix','rand','range','Rat','raw',
-        're','read','readchars','readonly','ready','Real','reallocate','reals',
-        'reason','rebless','receive','recv','redispatcher','redo','reduce',
-        'rel2abs','relative','release','rename','repeated','replacement',
-        'report','reserved','resolve','restore','result','resume','rethrow',
-        'reverse','right','rindex','rmdir','role','roles_to_compose','rolish',
-        'roll','rootdir','roots','rotate','rotor','round','roundrobin',
-        'routine-type','run','rwx','s','samecase','samemark','samewith','say',
-        'schedule-on','scheduler','scope','sec','sech','second','seek','self',
-        'send','Set','set_hidden','set_name','set_package','set_rw','set_value',
-        'SetHash','set-instruments','setup_finalization','shape','share','shell',
-        'shift','sibling','sigil','sign','signal','signals','signature','sin',
-        'sinh','sink','sink-all','skip','skip-at-least','skip-at-least-pull-one',
-        'skip-one','skip-rest','sleep','sleep-timer','sleep-until','Slip','slurp',
-        'slurp-rest','slurpy','snap','snapper','so','socket-host','socket-port',
-        'sort','source','source-package','spawn','SPEC','splice','split',
-        'splitdir','splitpath','sprintf','spurt','sqrt','squish','srand','stable',
-        'start','started','starts-with','status','stderr','stdout','Str',
-        'sub_signature','subbuf','subbuf-rw','subname','subparse','subst',
-        'subst-mutate','substr','substr-eq','substr-rw','subtest','succ','sum',
-        'Supply','symlink','t','tail','take','take-rw','tan','tanh','tap',
-        'target','target-name','tc','tclc','tell','then','throttle','throw',
-        'throws-like','timezone','tmpdir','to','today','todo','toggle','to-posix',
-        'total','trailing','trans','tree','trim','trim-leading','trim-trailing',
-        'truncate','truncated-to','trusts','try_acquire','trying','twigil','type',
-        'type_captures','typename','uc','udp','uncaught_handler','unimatch',
-        'uniname','uninames','uniparse','uniprop','uniprops','unique','unival',
-        'univals','unlike','unlink','unlock','unpack','unpolar','unshift',
-        'unwrap','updir','USAGE','use-ok','utc','val','value','values','VAR',
-        'variable','verbose-config','version','VMnames','volume','vow','w','wait',
-        'warn','watch','watch-path','week','weekday-of-month','week-number',
-        'week-year','WHAT','when','WHERE','WHEREFORE','WHICH','WHO',
-        'whole-second','WHY','wordcase','words','workaround','wrap','write',
-        'write-to','x','yada','year','yield','yyyy-mm-dd','z','zip','zip-latest',
-
-    )
-
-    PERL6_BUILTIN_CLASSES = (
-        #Booleans
-        'False','True',
-        #Classes
-        'Any','Array','Associative','AST','atomicint','Attribute','Backtrace',
-        'Backtrace::Frame','Bag','Baggy','BagHash','Blob','Block','Bool','Buf',
-        'Callable','CallFrame','Cancellation','Capture','CArray','Channel','Code',
-        'compiler','Complex','ComplexStr','Cool','CurrentThreadScheduler',
-        'Cursor','Date','Dateish','DateTime','Distro','Duration','Encoding',
-        'Exception','Failure','FatRat','Grammar','Hash','HyperWhatever','Instant',
-        'Int','int16','int32','int64','int8','IntStr','IO','IO::ArgFiles',
-        'IO::CatHandle','IO::Handle','IO::Notification','IO::Path',
-        'IO::Path::Cygwin','IO::Path::QNX','IO::Path::Unix','IO::Path::Win32',
-        'IO::Pipe','IO::Socket','IO::Socket::Async','IO::Socket::INET','IO::Spec',
-        'IO::Spec::Cygwin','IO::Spec::QNX','IO::Spec::Unix','IO::Spec::Win32',
-        'IO::Special','Iterable','Iterator','Junction','Kernel','Label','List',
-        'Lock','Lock::Async','long','longlong','Macro','Map','Match',
-        'Metamodel::AttributeContainer','Metamodel::C3MRO','Metamodel::ClassHOW',
-        'Metamodel::EnumHOW','Metamodel::Finalization','Metamodel::MethodContainer',
-        'Metamodel::MROBasedMethodDispatch','Metamodel::MultipleInheritance',
-        'Metamodel::Naming','Metamodel::Primitives','Metamodel::PrivateMethodContainer',
-        'Metamodel::RoleContainer','Metamodel::Trusting','Method','Mix','MixHash',
-        'Mixy','Mu','NFC','NFD','NFKC','NFKD','Nil','Num','num32','num64',
-        'Numeric','NumStr','ObjAt','Order','Pair','Parameter','Perl','Pod::Block',
-        'Pod::Block::Code','Pod::Block::Comment','Pod::Block::Declarator',
-        'Pod::Block::Named','Pod::Block::Para','Pod::Block::Table','Pod::Heading',
-        'Pod::Item','Pointer','Positional','PositionalBindFailover','Proc',
-        'Proc::Async','Promise','Proxy','PseudoStash','QuantHash','Range','Rat',
-        'Rational','RatStr','Real','Regex','Routine','Scalar','Scheduler',
-        'Semaphore','Seq','Set','SetHash','Setty','Signature','size_t','Slip',
-        'Stash','Str','StrDistance','Stringy','Sub','Submethod','Supplier',
-        'Supplier::Preserving','Supply','Systemic','Tap','Telemetry',
-        'Telemetry::Instrument::Thread','Telemetry::Instrument::Usage',
-        'Telemetry::Period','Telemetry::Sampler','Thread','ThreadPoolScheduler',
-        'UInt','uint16','uint32','uint64','uint8','Uni','utf8','Variable',
-        'Version','VM','Whatever','WhateverCode','WrapHandle'
-    )
-
-    PERL6_OPERATORS = (
-        'X', 'Z', 'after', 'also', 'and', 'andthen', 'before', 'cmp', 'div',
-        'eq', 'eqv', 'extra', 'ff', 'fff', 'ge', 'gt', 'le', 'leg', 'lt', 'm',
-        'mm', 'mod', 'ne', 'or', 'orelse', 'rx', 's', 'tr', 'x', 'xor', 'xx',
-        '++', '--', '**', '!', '+', '-', '~', '?', '|', '||', '+^', '~^', '?^',
-        '^', '*', '/', '%', '%%', '+&', '+<', '+>', '~&', '~<', '~>', '?&',
-        'gcd', 'lcm', '+', '-', '+|', '+^', '~|', '~^', '?|', '?^',
-        '~', '&', '^', 'but', 'does', '<=>', '..', '..^', '^..', '^..^',
-        '!=', '==', '<', '<=', '>', '>=', '~~', '===', '!eqv',
-        '&&', '||', '^^', '//', 'min', 'max', '??', '!!', 'ff', 'fff', 'so',
-        'not', '<==', '==>', '<<==', '==>>','unicmp',
-    )
-
-    # Perl 6 has a *lot* of possible bracketing characters
-    # this list was lifted from STD.pm6 (https://github.com/perl6/std)
-    PERL6_BRACKETS = {
-        '\u0028': '\u0029', '\u003c': '\u003e', '\u005b': '\u005d',
-        '\u007b': '\u007d', '\u00ab': '\u00bb', '\u0f3a': '\u0f3b',
-        '\u0f3c': '\u0f3d', '\u169b': '\u169c', '\u2018': '\u2019',
-        '\u201a': '\u2019', '\u201b': '\u2019', '\u201c': '\u201d',
-        '\u201e': '\u201d', '\u201f': '\u201d', '\u2039': '\u203a',
-        '\u2045': '\u2046', '\u207d': '\u207e', '\u208d': '\u208e',
-        '\u2208': '\u220b', '\u2209': '\u220c', '\u220a': '\u220d',
-        '\u2215': '\u29f5', '\u223c': '\u223d', '\u2243': '\u22cd',
-        '\u2252': '\u2253', '\u2254': '\u2255', '\u2264': '\u2265',
-        '\u2266': '\u2267', '\u2268': '\u2269', '\u226a': '\u226b',
-        '\u226e': '\u226f', '\u2270': '\u2271', '\u2272': '\u2273',
-        '\u2274': '\u2275', '\u2276': '\u2277', '\u2278': '\u2279',
-        '\u227a': '\u227b', '\u227c': '\u227d', '\u227e': '\u227f',
-        '\u2280': '\u2281', '\u2282': '\u2283', '\u2284': '\u2285',
-        '\u2286': '\u2287', '\u2288': '\u2289', '\u228a': '\u228b',
-        '\u228f': '\u2290', '\u2291': '\u2292', '\u2298': '\u29b8',
-        '\u22a2': '\u22a3', '\u22a6': '\u2ade', '\u22a8': '\u2ae4',
-        '\u22a9': '\u2ae3', '\u22ab': '\u2ae5', '\u22b0': '\u22b1',
-        '\u22b2': '\u22b3', '\u22b4': '\u22b5', '\u22b6': '\u22b7',
-        '\u22c9': '\u22ca', '\u22cb': '\u22cc', '\u22d0': '\u22d1',
-        '\u22d6': '\u22d7', '\u22d8': '\u22d9', '\u22da': '\u22db',
-        '\u22dc': '\u22dd', '\u22de': '\u22df', '\u22e0': '\u22e1',
-        '\u22e2': '\u22e3', '\u22e4': '\u22e5', '\u22e6': '\u22e7',
-        '\u22e8': '\u22e9', '\u22ea': '\u22eb', '\u22ec': '\u22ed',
-        '\u22f0': '\u22f1', '\u22f2': '\u22fa', '\u22f3': '\u22fb',
-        '\u22f4': '\u22fc', '\u22f6': '\u22fd', '\u22f7': '\u22fe',
-        '\u2308': '\u2309', '\u230a': '\u230b', '\u2329': '\u232a',
-        '\u23b4': '\u23b5', '\u2768': '\u2769', '\u276a': '\u276b',
-        '\u276c': '\u276d', '\u276e': '\u276f', '\u2770': '\u2771',
-        '\u2772': '\u2773', '\u2774': '\u2775', '\u27c3': '\u27c4',
-        '\u27c5': '\u27c6', '\u27d5': '\u27d6', '\u27dd': '\u27de',
-        '\u27e2': '\u27e3', '\u27e4': '\u27e5', '\u27e6': '\u27e7',
-        '\u27e8': '\u27e9', '\u27ea': '\u27eb', '\u2983': '\u2984',
-        '\u2985': '\u2986', '\u2987': '\u2988', '\u2989': '\u298a',
-        '\u298b': '\u298c', '\u298d': '\u298e', '\u298f': '\u2990',
-        '\u2991': '\u2992', '\u2993': '\u2994', '\u2995': '\u2996',
-        '\u2997': '\u2998', '\u29c0': '\u29c1', '\u29c4': '\u29c5',
-        '\u29cf': '\u29d0', '\u29d1': '\u29d2', '\u29d4': '\u29d5',
-        '\u29d8': '\u29d9', '\u29da': '\u29db', '\u29f8': '\u29f9',
-        '\u29fc': '\u29fd', '\u2a2b': '\u2a2c', '\u2a2d': '\u2a2e',
-        '\u2a34': '\u2a35', '\u2a3c': '\u2a3d', '\u2a64': '\u2a65',
-        '\u2a79': '\u2a7a', '\u2a7d': '\u2a7e', '\u2a7f': '\u2a80',
-        '\u2a81': '\u2a82', '\u2a83': '\u2a84', '\u2a8b': '\u2a8c',
-        '\u2a91': '\u2a92', '\u2a93': '\u2a94', '\u2a95': '\u2a96',
-        '\u2a97': '\u2a98', '\u2a99': '\u2a9a', '\u2a9b': '\u2a9c',
-        '\u2aa1': '\u2aa2', '\u2aa6': '\u2aa7', '\u2aa8': '\u2aa9',
-        '\u2aaa': '\u2aab', '\u2aac': '\u2aad', '\u2aaf': '\u2ab0',
-        '\u2ab3': '\u2ab4', '\u2abb': '\u2abc', '\u2abd': '\u2abe',
-        '\u2abf': '\u2ac0', '\u2ac1': '\u2ac2', '\u2ac3': '\u2ac4',
-        '\u2ac5': '\u2ac6', '\u2acd': '\u2ace', '\u2acf': '\u2ad0',
-        '\u2ad1': '\u2ad2', '\u2ad3': '\u2ad4', '\u2ad5': '\u2ad6',
-        '\u2aec': '\u2aed', '\u2af7': '\u2af8', '\u2af9': '\u2afa',
-        '\u2e02': '\u2e03', '\u2e04': '\u2e05', '\u2e09': '\u2e0a',
-        '\u2e0c': '\u2e0d', '\u2e1c': '\u2e1d', '\u2e20': '\u2e21',
-        '\u3008': '\u3009', '\u300a': '\u300b', '\u300c': '\u300d',
-        '\u300e': '\u300f', '\u3010': '\u3011', '\u3014': '\u3015',
-        '\u3016': '\u3017', '\u3018': '\u3019', '\u301a': '\u301b',
-        '\u301d': '\u301e', '\ufd3e': '\ufd3f', '\ufe17': '\ufe18',
-        '\ufe35': '\ufe36', '\ufe37': '\ufe38', '\ufe39': '\ufe3a',
-        '\ufe3b': '\ufe3c', '\ufe3d': '\ufe3e', '\ufe3f': '\ufe40',
-        '\ufe41': '\ufe42', '\ufe43': '\ufe44', '\ufe47': '\ufe48',
-        '\ufe59': '\ufe5a', '\ufe5b': '\ufe5c', '\ufe5d': '\ufe5e',
-        '\uff08': '\uff09', '\uff1c': '\uff1e', '\uff3b': '\uff3d',
-        '\uff5b': '\uff5d', '\uff5f': '\uff60', '\uff62': '\uff63',
-    }
-
-    def _build_word_match(words, boundary_regex_fragment=None, prefix='', suffix=''):
-        if boundary_regex_fragment is None:
-            return r'\b(' + prefix + r'|'.join(re.escape(x) for x in words) + \
-                suffix + r')\b'
-        else:
-            return r'(? 0:
-                    next_open_pos = text.find(opening_chars, search_pos + n_chars)
-                    next_close_pos = text.find(closing_chars, search_pos + n_chars)
-
-                    if next_close_pos == -1:
-                        next_close_pos = len(text)
-                        nesting_level = 0
-                    elif next_open_pos != -1 and next_open_pos < next_close_pos:
-                        nesting_level += 1
-                        search_pos = next_open_pos
-                    else:  # next_close_pos < next_open_pos
-                        nesting_level -= 1
-                        search_pos = next_close_pos
-
-                end_pos = next_close_pos
-
-            if end_pos < 0:     # if we didn't find a closer, just highlight the
-                                # rest of the text in this class
-                end_pos = len(text)
-
-            if adverbs is not None and re.search(r':to\b', adverbs):
-                heredoc_terminator = text[match.start('delimiter') + n_chars:end_pos]
-                end_heredoc = re.search(r'^\s*' + re.escape(heredoc_terminator) +
-                                        r'\s*$', text[end_pos:], re.MULTILINE)
-
-                if end_heredoc:
-                    end_pos += end_heredoc.end()
-                else:
-                    end_pos = len(text)
-
-            yield match.start(), token_class, text[match.start():end_pos + n_chars]
-            context.pos = end_pos + n_chars
-
-        return callback
-
-    def opening_brace_callback(lexer, match, context):
-        stack = context.stack
-
-        yield match.start(), Text, context.text[match.start():match.end()]
-        context.pos = match.end()
-
-        # if we encounter an opening brace and we're one level
-        # below a token state, it means we need to increment
-        # the nesting level for braces so we know later when
-        # we should return to the token rules.
-        if len(stack) > 2 and stack[-2] == 'token':
-            context.perl6_token_nesting_level += 1
-
-    def closing_brace_callback(lexer, match, context):
-        stack = context.stack
-
-        yield match.start(), Text, context.text[match.start():match.end()]
-        context.pos = match.end()
-
-        # if we encounter a free closing brace and we're one level
-        # below a token state, it means we need to check the nesting
-        # level to see if we need to return to the token state.
-        if len(stack) > 2 and stack[-2] == 'token':
-            context.perl6_token_nesting_level -= 1
-            if context.perl6_token_nesting_level == 0:
-                stack.pop()
-
-    def embedded_perl6_callback(lexer, match, context):
-        context.perl6_token_nesting_level = 1
-        yield match.start(), Text, context.text[match.start():match.end()]
-        context.pos = match.end()
-        context.stack.append('root')
-
-    # If you're modifying these rules, be careful if you need to process '{' or '}'
-    # characters. We have special logic for processing these characters (due to the fact
-    # that you can nest Perl 6 code in regex blocks), so if you need to process one of
-    # them, make sure you also process the corresponding one!
-    tokens = {
-        'common': [
-            (r'#[`|=](?P(?P[' + ''.join(PERL6_BRACKETS) + r'])(?P=first_char)*)',
-             brackets_callback(Comment.Multiline)),
-            (r'#[^\n]*$', Comment.Single),
-            (r'^(\s*)=begin\s+(\w+)\b.*?^\1=end\s+\2', Comment.Multiline),
-            (r'^(\s*)=for.*?\n\s*?\n', Comment.Multiline),
-            (r'^=.*?\n\s*?\n', Comment.Multiline),
-            (r'(regex|token|rule)(\s*' + PERL6_IDENTIFIER_RANGE + '+:sym)',
-             bygroups(Keyword, Name), 'token-sym-brackets'),
-            (r'(regex|token|rule)(?!' + PERL6_IDENTIFIER_RANGE + r')(\s*' + PERL6_IDENTIFIER_RANGE + '+)?',
-             bygroups(Keyword, Name), 'pre-token'),
-            # deal with a special case in the Perl 6 grammar (role q { ... })
-            (r'(role)(\s+)(q)(\s*)', bygroups(Keyword, Whitespace, Name, Whitespace)),
-            (_build_word_match(PERL6_KEYWORDS, PERL6_IDENTIFIER_RANGE), Keyword),
-            (_build_word_match(PERL6_BUILTIN_CLASSES, PERL6_IDENTIFIER_RANGE, suffix='(?::[UD])?'),
-             Name.Builtin),
-            (_build_word_match(PERL6_BUILTINS, PERL6_IDENTIFIER_RANGE), Name.Builtin),
-            # copied from PerlLexer
-            (r'[$@%&][.^:?=!~]?' + PERL6_IDENTIFIER_RANGE + '+(?:<<.*?>>|<.*?>|«.*?»)*',
-             Name.Variable),
-            (r'\$[!/](?:<<.*?>>|<.*?>|«.*?»)*', Name.Variable.Global),
-            (r'::\?\w+', Name.Variable.Global),
-            (r'[$@%&]\*' + PERL6_IDENTIFIER_RANGE + '+(?:<<.*?>>|<.*?>|«.*?»)*',
-             Name.Variable.Global),
-            (r'\$(?:<.*?>)+', Name.Variable),
-            (r'(?:q|qq|Q)[a-zA-Z]?\s*(?P:[\w\s:]+)?\s*(?P(?P[^0-9a-zA-Z:\s])'
-             r'(?P=first_char)*)', brackets_callback(String)),
-            # copied from PerlLexer
-            (r'0_?[0-7]+(_[0-7]+)*', Number.Oct),
-            (r'0x[0-9A-Fa-f]+(_[0-9A-Fa-f]+)*', Number.Hex),
-            (r'0b[01]+(_[01]+)*', Number.Bin),
-            (r'(?i)(\d*(_\d*)*\.\d+(_\d*)*|\d+(_\d*)*\.\d+(_\d*)*)(e[+-]?\d+)?',
-             Number.Float),
-            (r'(?i)\d+(_\d*)*e[+-]?\d+(_\d*)*', Number.Float),
-            (r'\d+(_\d+)*', Number.Integer),
-            (r'(?<=~~)\s*/(?:\\\\|\\/|.)*?/', String.Regex),
-            (r'(?<=[=(,])\s*/(?:\\\\|\\/|.)*?/', String.Regex),
-            (r'm\w+(?=\()', Name),
-            (r'(?:m|ms|rx)\s*(?P:[\w\s:]+)?\s*(?P(?P[^\w:\s])'
-             r'(?P=first_char)*)', brackets_callback(String.Regex)),
-            (r'(?:s|ss|tr)\s*(?::[\w\s:]+)?\s*/(?:\\\\|\\/|.)*?/(?:\\\\|\\/|.)*?/',
-             String.Regex),
-            (r'<[^\s=].*?\S>', String),
-            (_build_word_match(PERL6_OPERATORS), Operator),
-            (r'\w' + PERL6_IDENTIFIER_RANGE + '*', Name),
-            (r"'(\\\\|\\[^\\]|[^'\\])*'", String),
-            (r'"(\\\\|\\[^\\]|[^"\\])*"', String),
-        ],
-        'root': [
-            include('common'),
-            (r'\{', opening_brace_callback),
-            (r'\}', closing_brace_callback),
-            (r'.+?', Text),
-        ],
-        'pre-token': [
-            include('common'),
-            (r'\{', Text, ('#pop', 'token')),
-            (r'.+?', Text),
-        ],
-        'token-sym-brackets': [
-            (r'(?P(?P[' + ''.join(PERL6_BRACKETS) + '])(?P=first_char)*)',
-             brackets_callback(Name), ('#pop', 'pre-token')),
-            default(('#pop', 'pre-token')),
-        ],
-        'token': [
-            (r'\}', Text, '#pop'),
-            (r'(?<=:)(?:my|our|state|constant|temp|let).*?;', using(this)),
-            # make sure that quotes in character classes aren't treated as strings
-            (r'<(?:[-!?+.]\s*)?\[.*?\]>', String.Regex),
-            # make sure that '#' characters in quotes aren't treated as comments
-            (r"(?my|our)\s+)?(?:module|class|role|enum|grammar)', line)
-            if class_decl:
-                if saw_perl_decl or class_decl.group('scope') is not None:
-                    return True
-                rating = 0.05
-                continue
-            break
-
-        if ':=' in text:
-            # Same logic as above for PerlLexer
-            rating /= 2
-
-        return rating
-
-    def __init__(self, **options):
-        super().__init__(**options)
-        self.encoding = options.get('encoding', 'utf-8')
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/phix.py b/venv/lib/python3.10/site-packages/pygments/lexers/phix.py
deleted file mode 100644
index f0b03775e27efedbcd60b6f8ef0df0c93bb1eebc..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/phix.py
+++ /dev/null
@@ -1,363 +0,0 @@
-"""
-    pygments.lexers.phix
-    ~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Phix.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, words
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Whitespace
-
-__all__ = ['PhixLexer']
-
-
-class PhixLexer(RegexLexer):
-    """
-    Pygments Lexer for Phix files (.exw).
-    See http://phix.x10.mx
-    """
-
-    name = 'Phix'
-    url = 'http://phix.x10.mx'
-    aliases = ['phix']
-    filenames = ['*.exw']
-    mimetypes = ['text/x-phix']
-    version_added = '2.14'
-
-    flags = re.MULTILINE    # nb: **NOT** re.DOTALL! (totally spanners comment handling)
-
-    preproc = (
-        'ifdef', 'elsifdef', 'elsedef'
-    )
-    # Note these lists are auto-generated by pwa/p2js.exw, when pwa\src\p2js_keywords.e (etc)
-    #     change, though of course subsequent copy/commit/pull requests are all manual steps.
-    types = (
-        'string', 'nullable_string', 'atom_string', 'atom', 'bool', 'boolean',
-        'cdCanvan', 'cdCanvas', 'complex', 'CURLcode', 'dictionary', 'int',
-        'integer', 'Ihandle', 'Ihandles', 'Ihandln', 'mpfr', 'mpq', 'mpz',
-        'mpz_or_string', 'number', 'rid_string', 'seq', 'sequence', 'timedate',
-        'object'
-    )
-    keywords = (
-        'abstract', 'class', 'continue', 'export', 'extends', 'nullable',
-        'private', 'public', 'static', 'struct', 'trace',
-        'and', 'break', 'by', 'case', 'catch', 'const', 'constant', 'debug',
-        'default', 'do', 'else', 'elsif', 'end', 'enum', 'exit', 'fallthru',
-        'fallthrough', 'for', 'forward', 'function', 'global', 'if', 'in',
-        'include', 'js', 'javascript', 'javascript_semantics', 'let', 'not',
-        'or', 'procedure', 'profile', 'profile_time', 'return', 'safe_mode',
-        'switch', 'then', 'to', 'try', 'type', 'type_check', 'until', 'warning',
-        'while', 'with', 'without', 'xor'
-    )
-    routines = (
-        'abort', 'abs', 'adjust_timedate', 'and_bits', 'and_bitsu', 'apply',
-        'append', 'arccos', 'arcsin', 'arctan', 'assert', 'atan2',
-        'atom_to_float32', 'atom_to_float64', 'bankers_rounding', 'beep',
-        'begins', 'binary_search', 'bits_to_int', 'bk_color', 'bytes_to_int',
-        'call_func', 'call_proc', 'cdCanvasActivate', 'cdCanvasArc',
-        'cdCanvasBegin', 'cdCanvasBox', 'cdCanvasChord', 'cdCanvasCircle',
-        'cdCanvasClear', 'cdCanvasEnd', 'cdCanvasFlush', 'cdCanvasFont',
-        'cdCanvasGetImageRGB', 'cdCanvasGetSize', 'cdCanvasGetTextAlignment',
-        'cdCanvasGetTextSize', 'cdCanvasLine', 'cdCanvasMark',
-        'cdCanvasMarkSize', 'cdCanvasMultiLineVectorText', 'cdCanvasPixel',
-        'cdCanvasRect', 'cdCanvasRoundedBox', 'cdCanvasRoundedRect',
-        'cdCanvasSector', 'cdCanvasSetAttribute', 'cdCanvasSetBackground',
-        'cdCanvasSetFillMode', 'cdCanvasSetForeground',
-        'cdCanvasSetInteriorStyle', 'cdCanvasSetLineStyle',
-        'cdCanvasSetLineWidth', 'cdCanvasSetTextAlignment', 'cdCanvasText',
-        'cdCanvasSetTextOrientation', 'cdCanvasGetTextOrientation',
-        'cdCanvasVectorText', 'cdCanvasVectorTextDirection',
-        'cdCanvasVectorTextSize', 'cdCanvasVertex', 'cdCreateCanvas',
-        'cdDecodeAlpha', 'cdDecodeColor', 'cdDecodeColorAlpha', 'cdEncodeAlpha',
-        'cdEncodeColor', 'cdEncodeColorAlpha', 'cdKillCanvas', 'cdVersion',
-        'cdVersionDate', 'ceil', 'change_timezone', 'choose', 'clear_screen',
-        'columnize', 'command_line', 'compare', 'complex_abs', 'complex_add',
-        'complex_arg', 'complex_conjugate', 'complex_cos', 'complex_cosh',
-        'complex_div', 'complex_exp', 'complex_imag', 'complex_inv',
-        'complex_log', 'complex_mul', 'complex_neg', 'complex_new',
-        'complex_norm', 'complex_power', 'complex_rho', 'complex_real',
-        'complex_round', 'complex_sin', 'complex_sinh', 'complex_sprint',
-        'complex_sqrt', 'complex_sub', 'complex_theta', 'concat', 'cos',
-        'crash', 'custom_sort', 'date', 'day_of_week', 'day_of_year',
-        'days_in_month', 'decode_base64', 'decode_flags', 'deep_copy', 'deld',
-        'deserialize', 'destroy_dict', 'destroy_queue', 'destroy_stack',
-        'dict_name', 'dict_size', 'elapsed', 'elapsed_short', 'encode_base64',
-        'equal', 'even', 'exp', 'extract', 'factorial', 'factors',
-        'file_size_k', 'find', 'find_all', 'find_any', 'find_replace', 'filter',
-        'flatten', 'float32_to_atom', 'float64_to_atom', 'floor',
-        'format_timedate', 'free_console', 'from_polar', 'gcd', 'get_file_base',
-        'get_file_extension', 'get_file_name', 'get_file_name_and_path',
-        'get_file_path', 'get_file_path_and_name', 'get_maxprime', 'get_prime',
-        'get_primes', 'get_primes_le', 'get_proper_dir', 'get_proper_path',
-        'get_rand', 'get_routine_info', 'get_test_abort', 'get_test_logfile',
-        'get_test_pause', 'get_test_verbosity', 'get_tzid', 'getd', 'getdd',
-        'getd_all_keys', 'getd_by_index', 'getd_index', 'getd_partial_key',
-        'glAttachShader', 'glBindBuffer', 'glBindTexture', 'glBufferData',
-        'glCanvasSpecialText', 'glClear', 'glClearColor', 'glColor',
-        'glCompileShader', 'glCreateBuffer', 'glCreateProgram',
-        'glCreateShader', 'glCreateTexture', 'glDeleteProgram',
-        'glDeleteShader', 'glDrawArrays', 'glEnable',
-        'glEnableVertexAttribArray', 'glFloat32Array', 'glInt32Array',
-        'glFlush', 'glGetAttribLocation', 'glGetError', 'glGetProgramInfoLog',
-        'glGetProgramParameter', 'glGetShaderInfoLog', 'glGetShaderParameter',
-        'glGetUniformLocation', 'glLinkProgram', 'glLoadIdentity',
-        'glMatrixMode', 'glOrtho', 'glRotatef', 'glShadeModel',
-        'glShaderSource', 'glSimpleA7texcoords', 'glTexImage2Dc',
-        'glTexParameteri', 'glTranslate', 'glUniform1f', 'glUniform1i',
-        'glUniformMatrix4fv', 'glUseProgram', 'glVertex',
-        'glVertexAttribPointer', 'glViewport', 'head', 'hsv_to_rgb', 'iff',
-        'iif', 'include_file', 'incl0de_file', 'insert', 'instance',
-        'int_to_bits', 'int_to_bytes', 'is_dict', 'is_integer', 's_leap_year',
-        'is_prime', 'is_prime2', 'islower', 'isupper', 'Icallback',
-        'iup_isdouble', 'iup_isprint', 'iup_XkeyBase', 'IupAppend', 'IupAlarm',
-        'IupBackgroundBox', 'IupButton', 'IupCalendar', 'IupCanvas',
-        'IupClipboard', 'IupClose', 'IupCloseOnEscape', 'IupControlsOpen',
-        'IupDatePick', 'IupDestroy', 'IupDialog', 'IupDrawArc', 'IupDrawBegin',
-        'IupDrawEnd', 'IupDrawGetSize', 'IupDrawGetTextSize', 'IupDrawLine',
-        'IupDrawRectangle', 'IupDrawText', 'IupExpander', 'IupFill',
-        'IupFlatLabel', 'IupFlatList', 'IupFlatTree', 'IupFlush', 'IupFrame',
-        'IupGetAttribute', 'IupGetAttributeId', 'IupGetAttributePtr',
-        'IupGetBrother', 'IupGetChild', 'IupGetChildCount', 'IupGetClassName',
-        'IupGetDialog', 'IupGetDialogChild', 'IupGetDouble', 'IupGetFocus',
-        'IupGetGlobal', 'IupGetGlobalInt', 'IupGetGlobalIntInt', 'IupGetInt',
-        'IupGetInt2', 'IupGetIntId', 'IupGetIntInt', 'IupGetParent',
-        'IupGLCanvas', 'IupGLCanvasOpen', 'IupGLMakeCurrent', 'IupGraph',
-        'IupHbox', 'IupHide', 'IupImage', 'IupImageRGBA', 'IupItem',
-        'iupKeyCodeToName', 'IupLabel', 'IupLink', 'IupList', 'IupMap',
-        'IupMenu', 'IupMenuItem', 'IupMessage', 'IupMessageDlg', 'IupMultiBox',
-        'IupMultiLine', 'IupNextField', 'IupNormaliser', 'IupOpen',
-        'IupPlayInput', 'IupPopup', 'IupPreviousField', 'IupProgressBar',
-        'IupRadio', 'IupRecordInput', 'IupRedraw', 'IupRefresh',
-        'IupRefreshChildren', 'IupSeparator', 'IupSetAttribute',
-        'IupSetAttributes', 'IupSetAttributeHandle', 'IupSetAttributeId',
-        'IupSetAttributePtr', 'IupSetCallback', 'IupSetCallbacks',
-        'IupSetDouble', 'IupSetFocus', 'IupSetGlobal', 'IupSetGlobalInt',
-        'IupSetGlobalFunction', 'IupSetHandle', 'IupSetInt',
-        'IupSetStrAttribute', 'IupSetStrGlobal', 'IupShow', 'IupShowXY',
-        'IupSplit', 'IupStoreAttribute', 'IupSubmenu', 'IupTable',
-        'IupTableClearSelected', 'IupTableClick_cb', 'IupTableGetSelected',
-        'IupTableResize_cb', 'IupTableSetData', 'IupTabs', 'IupText',
-        'IupTimer', 'IupToggle', 'IupTreeAddNodes', 'IupTreeView', 'IupUpdate',
-        'IupValuator', 'IupVbox', 'join', 'join_by', 'join_path', 'k_perm',
-        'largest', 'lcm', 'length', 'log', 'log10', 'log2', 'lower',
-        'm4_crossProduct', 'm4_inverse', 'm4_lookAt', 'm4_multiply',
-        'm4_normalize', 'm4_perspective', 'm4_subtractVectors', 'm4_xRotate',
-        'm4_yRotate', 'machine_bits', 'machine_word', 'match', 'match_all',
-        'match_replace', 'max', 'maxsq', 'min', 'minsq', 'mod', 'mpfr_add',
-        'mpfr_ceil', 'mpfr_cmp', 'mpfr_cmp_si', 'mpfr_const_pi', 'mpfr_div',
-        'mpfr_div_si', 'mpfr_div_z', 'mpfr_floor', 'mpfr_free', 'mpfr_get_d',
-        'mpfr_get_default_precision', 'mpfr_get_default_rounding_mode',
-        'mpfr_get_fixed', 'mpfr_get_precision', 'mpfr_get_si', 'mpfr_init',
-        'mpfr_inits', 'mpfr_init_set', 'mpfr_init_set_q', 'mpfr_init_set_z',
-        'mpfr_mul', 'mpfr_mul_si', 'mpfr_pow_si', 'mpfr_set', 'mpfr_set_d',
-        'mpfr_set_default_precision', 'mpfr_set_default_rounding_mode',
-        'mpfr_set_precision', 'mpfr_set_q', 'mpfr_set_si', 'mpfr_set_str',
-        'mpfr_set_z', 'mpfr_si_div', 'mpfr_si_sub', 'mpfr_sqrt', 'mpfr_sub',
-        'mpfr_sub_si', 'mpq_abs', 'mpq_add', 'mpq_add_si', 'mpq_canonicalize',
-        'mpq_cmp', 'mpq_cmp_si', 'mpq_div', 'mpq_div_2exp', 'mpq_free',
-        'mpq_get_den', 'mpq_get_num', 'mpq_get_str', 'mpq_init', 'mpq_init_set',
-        'mpq_init_set_si', 'mpq_init_set_str', 'mpq_init_set_z', 'mpq_inits',
-        'mpq_inv', 'mpq_mul', 'mpq_neg', 'mpq_set', 'mpq_set_si', 'mpq_set_str',
-        'mpq_set_z', 'mpq_sub', 'mpz_abs', 'mpz_add', 'mpz_addmul',
-        'mpz_addmul_ui', 'mpz_addmul_si', 'mpz_add_si', 'mpz_add_ui', 'mpz_and',
-        'mpz_bin_uiui', 'mpz_cdiv_q', 'mpz_cmp', 'mpz_cmp_si', 'mpz_divexact',
-        'mpz_divexact_ui', 'mpz_divisible_p', 'mpz_divisible_ui_p', 'mpz_even',
-        'mpz_fac_ui', 'mpz_factorstring', 'mpz_fdiv_q', 'mpz_fdiv_q_2exp',
-        'mpz_fdiv_q_ui', 'mpz_fdiv_qr', 'mpz_fdiv_r', 'mpz_fdiv_ui',
-        'mpz_fib_ui', 'mpz_fib2_ui', 'mpz_fits_atom', 'mpz_fits_integer',
-        'mpz_free', 'mpz_gcd', 'mpz_gcd_ui', 'mpz_get_atom', 'mpz_get_integer',
-        'mpz_get_short_str', 'mpz_get_str', 'mpz_init', 'mpz_init_set',
-        'mpz_inits', 'mpz_invert', 'mpz_lcm', 'mpz_lcm_ui', 'mpz_max',
-        'mpz_min', 'mpz_mod', 'mpz_mod_ui', 'mpz_mul', 'mpz_mul_2exp',
-        'mpz_mul_d', 'mpz_mul_si', 'mpz_neg', 'mpz_nthroot', 'mpz_odd',
-        'mpz_pollard_rho', 'mpz_pow_ui', 'mpz_powm', 'mpz_powm_ui', 'mpz_prime',
-        'mpz_prime_factors', 'mpz_prime_mr', 'mpz_rand', 'mpz_rand_ui',
-        'mpz_re_compose', 'mpz_remove', 'mpz_scan0', 'mpz_scan1', 'mpz_set',
-        'mpz_set_d', 'mpz_set_si', 'mpz_set_str', 'mpz_set_v', 'mpz_sign',
-        'mpz_sizeinbase', 'mpz_sqrt', 'mpz_sub', 'mpz_sub_si', 'mpz_sub_ui',
-        'mpz_si_sub', 'mpz_tdiv_q_2exp', 'mpz_tdiv_r_2exp', 'mpz_tstbit',
-        'mpz_ui_pow_ui', 'mpz_xor', 'named_dict', 'new_dict', 'new_queue',
-        'new_stack', 'not_bits', 'not_bitsu', 'odd', 'or_all', 'or_allu',
-        'or_bits', 'or_bitsu', 'ord', 'ordinal', 'ordinant',
-        'override_timezone', 'pad', 'pad_head', 'pad_tail', 'parse_date_string',
-        'papply', 'peep', 'peepn', 'peep_dict', 'permute', 'permutes',
-        'platform', 'pop', 'popn', 'pop_dict', 'power', 'pp', 'ppEx', 'ppExf',
-        'ppf', 'ppOpt', 'pq_add', 'pq_destroy', 'pq_empty', 'pq_new', 'pq_peek',
-        'pq_pop', 'pq_pop_data', 'pq_size', 'prepend', 'prime_factors',
-        'printf', 'product', 'proper', 'push', 'pushn', 'putd', 'puts',
-        'queue_empty', 'queue_size', 'rand', 'rand_range', 'reinstate',
-        'remainder', 'remove', 'remove_all', 'repeat', 'repeatch', 'replace',
-        'requires', 'reverse', 'rfind', 'rgb', 'rmatch', 'rmdr', 'rnd', 'round',
-        'routine_id', 'scanf', 'serialize', 'series', 'set_rand',
-        'set_test_abort', 'set_test_logfile', 'set_test_module',
-        'set_test_pause', 'set_test_verbosity', 'set_timedate_formats',
-        'set_timezone', 'setd', 'setd_default', 'shorten', 'sha256',
-        'shift_bits', 'shuffle', 'sign', 'sin', 'smallest', 'sort',
-        'sort_columns', 'speak', 'splice', 'split', 'split_any', 'split_by',
-        'sprint', 'sprintf', 'sq_abs', 'sq_add', 'sq_and', 'sq_and_bits',
-        'sq_arccos', 'sq_arcsin', 'sq_arctan', 'sq_atom', 'sq_ceil', 'sq_cmp',
-        'sq_cos', 'sq_div', 'sq_even', 'sq_eq', 'sq_floor', 'sq_floor_div',
-        'sq_ge', 'sq_gt', 'sq_int', 'sq_le', 'sq_log', 'sq_log10', 'sq_log2',
-        'sq_lt', 'sq_max', 'sq_min', 'sq_mod', 'sq_mul', 'sq_ne', 'sq_not',
-        'sq_not_bits', 'sq_odd', 'sq_or', 'sq_or_bits', 'sq_power', 'sq_rand',
-        'sq_remainder', 'sq_rmdr', 'sq_rnd', 'sq_round', 'sq_seq', 'sq_sign',
-        'sq_sin', 'sq_sqrt', 'sq_str', 'sq_sub', 'sq_tan', 'sq_trunc',
-        'sq_uminus', 'sq_xor', 'sq_xor_bits', 'sqrt', 'square_free',
-        'stack_empty', 'stack_size', 'substitute', 'substitute_all', 'sum',
-        'tail', 'tan', 'test_equal', 'test_fail', 'test_false',
-        'test_not_equal', 'test_pass', 'test_summary', 'test_true',
-        'text_color', 'throw', 'time', 'timedate_diff', 'timedelta',
-        'to_integer', 'to_number', 'to_rgb', 'to_string', 'traverse_dict',
-        'traverse_dict_partial_key', 'trim', 'trim_head', 'trim_tail', 'trunc',
-        'tagset', 'tagstart', 'typeof', 'unique', 'unix_dict', 'upper',
-        'utf8_to_utf32', 'utf32_to_utf8', 'version', 'vlookup', 'vslice',
-        'wglGetProcAddress', 'wildcard_file', 'wildcard_match', 'with_rho',
-        'with_theta', 'xml_new_doc', 'xml_new_element', 'xml_set_attribute',
-        'xml_sprint', 'xor_bits', 'xor_bitsu',
-        'accept', 'allocate', 'allocate_string', 'allow_break', 'ARM',
-        'atom_to_float80', 'c_func', 'c_proc', 'call_back', 'chdir',
-        'check_break', 'clearDib', 'close', 'closesocket', 'console',
-        'copy_file', 'create', 'create_directory', 'create_thread',
-        'curl_easy_cleanup', 'curl_easy_get_file', 'curl_easy_init',
-        'curl_easy_perform', 'curl_easy_perform_ex', 'curl_easy_setopt',
-        'curl_easy_strerror', 'curl_global_cleanup', 'curl_global_init',
-        'curl_slist_append', 'curl_slist_free_all', 'current_dir', 'cursor',
-        'define_c_func', 'define_c_proc', 'delete', 'delete_cs', 'delete_file',
-        'dir', 'DLL', 'drawDib', 'drawShadedPolygonToDib', 'ELF32', 'ELF64',
-        'enter_cs', 'eval', 'exit_thread', 'free', 'file_exists', 'final',
-        'float80_to_atom', 'format', 'get_bytes', 'get_file_date',
-        'get_file_size', 'get_file_type', 'get_interpreter', 'get_key',
-        'get_socket_error', 'get_text', 'get_thread_exitcode', 'get_thread_id',
-        'getc', 'getenv', 'gets', 'getsockaddr', 'glBegin', 'glCallList',
-        'glFrustum', 'glGenLists', 'glGetString', 'glLight', 'glMaterial',
-        'glNewList', 'glNormal', 'glPopMatrix', 'glPushMatrix', 'glRotate',
-        'glEnd', 'glEndList', 'glTexImage2D', 'goto', 'GUI', 'icons', 'ilASM',
-        'include_files', 'include_paths', 'init_cs', 'ip_to_string',
-        'IupConfig', 'IupConfigDialogClosed', 'IupConfigDialogShow',
-        'IupConfigGetVariableInt', 'IupConfigLoad', 'IupConfigSave',
-        'IupConfigSetVariableInt', 'IupExitLoop', 'IupFileDlg', 'IupFileList',
-        'IupGLSwapBuffers', 'IupHelp', 'IupLoopStep', 'IupMainLoop',
-        'IupNormalizer', 'IupPlot', 'IupPlotAdd', 'IupPlotBegin', 'IupPlotEnd',
-        'IupPlotInsert', 'IupSaveImage', 'IupTreeGetUserId', 'IupUser',
-        'IupVersion', 'IupVersionDate', 'IupVersionNumber', 'IupVersionShow',
-        'killDib', 'leave_cs', 'listen', 'manifest', 'mem_copy', 'mem_set',
-        'mpfr_gamma', 'mpfr_printf', 'mpfr_sprintf', 'mpz_export', 'mpz_import',
-        'namespace', 'new', 'newDib', 'open', 'open_dll', 'PE32', 'PE64',
-        'peek', 'peek_string', 'peek1s', 'peek1u', 'peek2s', 'peek2u', 'peek4s',
-        'peek4u', 'peek8s', 'peek8u', 'peekNS', 'peekns', 'peeknu', 'poke',
-        'poke2', 'poke4', 'poke8', 'pokeN', 'poke_string', 'poke_wstring',
-        'position', 'progress', 'prompt_number', 'prompt_string', 'read_file',
-        'read_lines', 'recv', 'resume_thread', 'seek', 'select', 'send',
-        'setHandler', 'shutdown', 'sleep', 'SO', 'sockaddr_in', 'socket',
-        'split_path', 'suspend_thread', 'system', 'system_exec', 'system_open',
-        'system_wait', 'task_clock_start', 'task_clock_stop', 'task_create',
-        'task_delay', 'task_list', 'task_schedule', 'task_self', 'task_status',
-        'task_suspend', 'task_yield', 'thread_safe_string', 'try_cs',
-        'utf8_to_utf16', 'utf16_to_utf8', 'utf16_to_utf32', 'utf32_to_utf16',
-        'video_config', 'WSACleanup', 'wait_thread', 'walk_dir', 'where',
-        'write_lines', 'wait_key'
-    )
-    constants = (
-        'ANY_QUEUE', 'ASCENDING', 'BLACK', 'BLOCK_CURSOR', 'BLUE',
-        'BRIGHT_CYAN', 'BRIGHT_BLUE', 'BRIGHT_GREEN', 'BRIGHT_MAGENTA',
-        'BRIGHT_RED', 'BRIGHT_WHITE', 'BROWN', 'C_DWORD', 'C_INT', 'C_POINTER',
-        'C_USHORT', 'C_WORD', 'CD_AMBER', 'CD_BLACK', 'CD_BLUE', 'CD_BOLD',
-        'CD_BOLD_ITALIC', 'CD_BOX', 'CD_CENTER', 'CD_CIRCLE', 'CD_CLOSED_LINES',
-        'CD_CONTINUOUS', 'CD_CUSTOM', 'CD_CYAN', 'CD_DARK_BLUE', 'CD_DARK_CYAN',
-        'CD_DARK_GRAY', 'CD_DARK_GREY', 'CD_DARK_GREEN', 'CD_DARK_MAGENTA',
-        'CD_DARK_RED', 'CD_DARK_YELLOW', 'CD_DASH_DOT', 'CD_DASH_DOT_DOT',
-        'CD_DASHED', 'CD_DBUFFER', 'CD_DEG2RAD', 'CD_DIAMOND', 'CD_DOTTED',
-        'CD_EAST', 'CD_EVENODD', 'CD_FILL', 'CD_GL', 'CD_GRAY', 'CD_GREY',
-        'CD_GREEN', 'CD_HATCH', 'CD_HOLLOW', 'CD_HOLLOW_BOX',
-        'CD_HOLLOW_CIRCLE', 'CD_HOLLOW_DIAMOND', 'CD_INDIGO', 'CD_ITALIC',
-        'CD_IUP', 'CD_IUPDBUFFER', 'CD_LIGHT_BLUE', 'CD_LIGHT_GRAY',
-        'CD_LIGHT_GREY', 'CD_LIGHT_GREEN', 'CD_LIGHT_PARCHMENT', 'CD_MAGENTA',
-        'CD_NAVY', 'CD_NORTH', 'CD_NORTH_EAST', 'CD_NORTH_WEST', 'CD_OLIVE',
-        'CD_OPEN_LINES', 'CD_ORANGE', 'CD_PARCHMENT', 'CD_PATTERN',
-        'CD_PRINTER', 'CD_PURPLE', 'CD_PLAIN', 'CD_PLUS', 'CD_QUERY',
-        'CD_RAD2DEG', 'CD_RED', 'CD_SILVER', 'CD_SOLID', 'CD_SOUTH_EAST',
-        'CD_SOUTH_WEST', 'CD_STAR', 'CD_STIPPLE', 'CD_STRIKEOUT',
-        'CD_UNDERLINE', 'CD_WEST', 'CD_WHITE', 'CD_WINDING', 'CD_VIOLET',
-        'CD_X', 'CD_YELLOW', 'CURLE_OK', 'CURLOPT_MAIL_FROM',
-        'CURLOPT_MAIL_RCPT', 'CURLOPT_PASSWORD', 'CURLOPT_READDATA',
-        'CURLOPT_READFUNCTION', 'CURLOPT_SSL_VERIFYPEER',
-        'CURLOPT_SSL_VERIFYHOST', 'CURLOPT_UPLOAD', 'CURLOPT_URL',
-        'CURLOPT_USE_SSL', 'CURLOPT_USERNAME', 'CURLOPT_VERBOSE',
-        'CURLOPT_WRITEFUNCTION', 'CURLUSESSL_ALL', 'CYAN', 'D_NAME',
-        'D_ATTRIBUTES', 'D_SIZE', 'D_YEAR', 'D_MONTH', 'D_DAY', 'D_HOUR',
-        'D_MINUTE', 'D_SECOND', 'D_CREATION', 'D_LASTACCESS', 'D_MODIFICATION',
-        'DT_YEAR', 'DT_MONTH', 'DT_DAY', 'DT_HOUR', 'DT_MINUTE', 'DT_SECOND',
-        'DT_DOW', 'DT_MSEC', 'DT_DOY', 'DT_GMT', 'EULER', 'E_CODE', 'E_ADDR',
-        'E_LINE', 'E_RTN', 'E_NAME', 'E_FILE', 'E_PATH', 'E_USER', 'false',
-        'False', 'FALSE', 'FIFO_QUEUE', 'FILETYPE_DIRECTORY', 'FILETYPE_FILE',
-        'GET_EOF', 'GET_FAIL', 'GET_IGNORE', 'GET_SUCCESS',
-        'GL_AMBIENT_AND_DIFFUSE', 'GL_ARRAY_BUFFER', 'GL_CLAMP',
-        'GL_CLAMP_TO_BORDER', 'GL_CLAMP_TO_EDGE', 'GL_COLOR_BUFFER_BIT',
-        'GL_COMPILE', 'GL_COMPILE_STATUS', 'GL_CULL_FACE',
-        'GL_DEPTH_BUFFER_BIT', 'GL_DEPTH_TEST', 'GL_EXTENSIONS', 'GL_FLAT',
-        'GL_FLOAT', 'GL_FRAGMENT_SHADER', 'GL_FRONT', 'GL_LIGHT0',
-        'GL_LIGHTING', 'GL_LINEAR', 'GL_LINK_STATUS', 'GL_MODELVIEW',
-        'GL_NEAREST', 'GL_NO_ERROR', 'GL_NORMALIZE', 'GL_POSITION',
-        'GL_PROJECTION', 'GL_QUAD_STRIP', 'GL_QUADS', 'GL_RENDERER',
-        'GL_REPEAT', 'GL_RGB', 'GL_RGBA', 'GL_SMOOTH', 'GL_STATIC_DRAW',
-        'GL_TEXTURE_2D', 'GL_TEXTURE_MAG_FILTER', 'GL_TEXTURE_MIN_FILTER',
-        'GL_TEXTURE_WRAP_S', 'GL_TEXTURE_WRAP_T', 'GL_TRIANGLES',
-        'GL_UNSIGNED_BYTE', 'GL_VENDOR', 'GL_VERSION', 'GL_VERTEX_SHADER',
-        'GRAY', 'GREEN', 'GT_LF_STRIPPED', 'GT_WHOLE_FILE', 'INVLN10',
-        'IUP_CLOSE', 'IUP_CONTINUE', 'IUP_DEFAULT', 'IUP_BLACK', 'IUP_BLUE',
-        'IUP_BUTTON1', 'IUP_BUTTON3', 'IUP_CENTER', 'IUP_CYAN', 'IUP_DARK_BLUE',
-        'IUP_DARK_CYAN', 'IUP_DARK_GRAY', 'IUP_DARK_GREY', 'IUP_DARK_GREEN',
-        'IUP_DARK_MAGENTA', 'IUP_DARK_RED', 'IUP_GRAY', 'IUP_GREY', 'IUP_GREEN',
-        'IUP_IGNORE', 'IUP_INDIGO', 'IUP_MAGENTA', 'IUP_MASK_INT',
-        'IUP_MASK_UINT', 'IUP_MOUSEPOS', 'IUP_NAVY', 'IUP_OLIVE', 'IUP_RECTEXT',
-        'IUP_RED', 'IUP_LIGHT_BLUE', 'IUP_LIGHT_GRAY', 'IUP_LIGHT_GREY',
-        'IUP_LIGHT_GREEN', 'IUP_ORANGE', 'IUP_PARCHMENT', 'IUP_PURPLE',
-        'IUP_SILVER', 'IUP_TEAL', 'IUP_VIOLET', 'IUP_WHITE', 'IUP_YELLOW',
-        'K_BS', 'K_cA', 'K_cC', 'K_cD', 'K_cF5', 'K_cK', 'K_cM', 'K_cN', 'K_cO',
-        'K_cP', 'K_cR', 'K_cS', 'K_cT', 'K_cW', 'K_CR', 'K_DEL', 'K_DOWN',
-        'K_END', 'K_ESC', 'K_F1', 'K_F2', 'K_F3', 'K_F4', 'K_F5', 'K_F6',
-        'K_F7', 'K_F8', 'K_F9', 'K_F10', 'K_F11', 'K_F12', 'K_HOME', 'K_INS',
-        'K_LEFT', 'K_MIDDLE', 'K_PGDN', 'K_PGUP', 'K_RIGHT', 'K_SP', 'K_TAB',
-        'K_UP', 'K_h', 'K_i', 'K_j', 'K_p', 'K_r', 'K_s', 'JS', 'LIFO_QUEUE',
-        'LINUX', 'MAX_HEAP', 'MAGENTA', 'MIN_HEAP', 'Nan', 'NO_CURSOR', 'null',
-        'NULL', 'PI', 'pp_Ascii', 'pp_Brkt', 'pp_Date', 'pp_File', 'pp_FltFmt',
-        'pp_Indent', 'pp_IntCh', 'pp_IntFmt', 'pp_Maxlen', 'pp_Nest',
-        'pp_Pause', 'pp_Q22', 'pp_StrFmt', 'RED', 'SEEK_OK', 'SLASH',
-        'TEST_ABORT', 'TEST_CRASH', 'TEST_PAUSE', 'TEST_PAUSE_FAIL',
-        'TEST_QUIET', 'TEST_SHOW_ALL', 'TEST_SHOW_FAILED', 'TEST_SUMMARY',
-        'true', 'True', 'TRUE', 'VC_SCRNLINES', 'WHITE', 'WINDOWS', 'YELLOW'
-    )
-
-    tokens = {
-        'root': [
-            (r"\s+", Whitespace),
-            (r'/\*|--/\*|#\[', Comment.Multiline, 'comment'),
-            (r'(?://|--|#!).*$', Comment.Single),
-#Alt:
-#           (r'//.*$|--.*$|#!.*$', Comment.Single),
-            (r'"([^"\\]|\\.)*"', String.Other),
-            (r'\'[^\']*\'', String.Other),
-            (r'`[^`]*`', String.Other),
-
-            (words(types, prefix=r'\b', suffix=r'\b'), Name.Function),
-            (words(routines, prefix=r'\b', suffix=r'\b'), Name.Function),
-            (words(preproc, prefix=r'\b', suffix=r'\b'), Keyword.Declaration),
-            (words(keywords, prefix=r'\b', suffix=r'\b'), Keyword.Declaration),
-            (words(constants, prefix=r'\b', suffix=r'\b'), Name.Constant),
-            # Aside: Phix only supports/uses the ascii/non-unicode tilde
-            (r'!=|==|<<|>>|:=|[-~+/*%=<>&^|\.(){},?:\[\]$\\;#]', Operator),
-            (r'[\w-]+', Text)
-        ],
-        'comment': [
-            (r'[^*/#]+', Comment.Multiline),
-            (r'/\*|#\[', Comment.Multiline, '#push'),
-            (r'\*/|#\]', Comment.Multiline, '#pop'),
-            (r'[*/#]', Comment.Multiline)
-        ]
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/php.py b/venv/lib/python3.10/site-packages/pygments/lexers/php.py
deleted file mode 100644
index 82d4aeb3d9d8eedca8cd33fa63443a5936dfc8e1..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/php.py
+++ /dev/null
@@ -1,334 +0,0 @@
-"""
-    pygments.lexers.php
-    ~~~~~~~~~~~~~~~~~~~
-
-    Lexers for PHP and related languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import Lexer, RegexLexer, include, bygroups, default, \
-    using, this, words, do_insertions, line_re
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Other, Generic
-from pygments.util import get_bool_opt, get_list_opt, shebang_matches
-
-__all__ = ['ZephirLexer', 'PsyshConsoleLexer', 'PhpLexer']
-
-
-class ZephirLexer(RegexLexer):
-    """
-    For Zephir language source code.
-
-    Zephir is a compiled high level language aimed
-    to the creation of C-extensions for PHP.
-    """
-
-    name = 'Zephir'
-    url = 'http://zephir-lang.com/'
-    aliases = ['zephir']
-    filenames = ['*.zep']
-    version_added = '2.0'
-
-    zephir_keywords = ['fetch', 'echo', 'isset', 'empty']
-    zephir_type = ['bit', 'bits', 'string']
-
-    flags = re.DOTALL | re.MULTILINE
-
-    tokens = {
-        'commentsandwhitespace': [
-            (r'\s+', Text),
-            (r'//.*?\n', Comment.Single),
-            (r'/\*.*?\*/', Comment.Multiline)
-        ],
-        'slashstartsregex': [
-            include('commentsandwhitespace'),
-            (r'/(\\.|[^[/\\\n]|\[(\\.|[^\]\\\n])*])+/'
-             r'([gim]+\b|\B)', String.Regex, '#pop'),
-            (r'/', Operator, '#pop'),
-            default('#pop')
-        ],
-        'badregex': [
-            (r'\n', Text, '#pop')
-        ],
-        'root': [
-            (r'^(?=\s|/)', Text, 'slashstartsregex'),
-            include('commentsandwhitespace'),
-            (r'\+\+|--|~|&&|\?|:|\|\||\\(?=\n)|'
-             r'(<<|>>>?|==?|!=?|->|[-<>+*%&|^/])=?', Operator, 'slashstartsregex'),
-            (r'[{(\[;,]', Punctuation, 'slashstartsregex'),
-            (r'[})\].]', Punctuation),
-            (r'(for|in|while|do|break|return|continue|switch|case|default|if|else|loop|'
-             r'require|inline|throw|try|catch|finally|new|delete|typeof|instanceof|void|'
-             r'namespace|use|extends|this|fetch|isset|unset|echo|fetch|likely|unlikely|'
-             r'empty)\b', Keyword, 'slashstartsregex'),
-            (r'(var|let|with|function)\b', Keyword.Declaration, 'slashstartsregex'),
-            (r'(abstract|boolean|bool|char|class|const|double|enum|export|extends|final|'
-             r'native|goto|implements|import|int|string|interface|long|ulong|char|uchar|'
-             r'float|unsigned|private|protected|public|short|static|self|throws|reverse|'
-             r'transient|volatile|readonly)\b', Keyword.Reserved),
-            (r'(true|false|null|undefined)\b', Keyword.Constant),
-            (r'(Array|Boolean|Date|_REQUEST|_COOKIE|_SESSION|'
-             r'_GET|_POST|_SERVER|this|stdClass|range|count|iterator|'
-             r'window)\b', Name.Builtin),
-            (r'[$a-zA-Z_][\w\\]*', Name.Other),
-            (r'[0-9][0-9]*\.[0-9]+([eE][0-9]+)?[fd]?', Number.Float),
-            (r'0x[0-9a-fA-F]+', Number.Hex),
-            (r'[0-9]+', Number.Integer),
-            (r'"(\\\\|\\[^\\]|[^"\\])*"', String.Double),
-            (r"'(\\\\|\\[^\\]|[^'\\])*'", String.Single),
-        ]
-    }
-
-
-class PsyshConsoleLexer(Lexer):
-    """
-    For PsySH console output, such as:
-
-    .. sourcecode:: psysh
-
-        >>> $greeting = function($name): string {
-        ...     return "Hello, {$name}";
-        ... };
-        => Closure($name): string {#2371 …3}
-        >>> $greeting('World')
-        => "Hello, World"
-    """
-    name = 'PsySH console session for PHP'
-    url = 'https://psysh.org/'
-    aliases = ['psysh']
-    version_added = '2.7'
-
-    def __init__(self, **options):
-        options['startinline'] = True
-        Lexer.__init__(self, **options)
-
-    def get_tokens_unprocessed(self, text):
-        phplexer = PhpLexer(**self.options)
-        curcode = ''
-        insertions = []
-        for match in line_re.finditer(text):
-            line = match.group()
-            if line.startswith('>>> ') or line.startswith('... '):
-                insertions.append((len(curcode),
-                                   [(0, Generic.Prompt, line[:4])]))
-                curcode += line[4:]
-            elif line.rstrip() == '...':
-                insertions.append((len(curcode),
-                                   [(0, Generic.Prompt, '...')]))
-                curcode += line[3:]
-            else:
-                if curcode:
-                    yield from do_insertions(
-                        insertions, phplexer.get_tokens_unprocessed(curcode))
-                    curcode = ''
-                    insertions = []
-                yield match.start(), Generic.Output, line
-        if curcode:
-            yield from do_insertions(insertions,
-                                     phplexer.get_tokens_unprocessed(curcode))
-
-
-class PhpLexer(RegexLexer):
-    """
-    For PHP source code.
-    For PHP embedded in HTML, use the `HtmlPhpLexer`.
-
-    Additional options accepted:
-
-    `startinline`
-        If given and ``True`` the lexer starts highlighting with
-        php code (i.e.: no starting ``>> from pygments.lexers._php_builtins import MODULES
-            >>> MODULES.keys()
-            ['PHP Options/Info', 'Zip', 'dba', ...]
-
-        In fact the names of those modules match the module names from
-        the php documentation.
-    """
-
-    name = 'PHP'
-    url = 'https://www.php.net/'
-    aliases = ['php', 'php3', 'php4', 'php5']
-    filenames = ['*.php', '*.php[345]', '*.inc']
-    mimetypes = ['text/x-php']
-    version_added = ''
-
-    # Note that a backslash is included, PHP uses a backslash as a namespace
-    # separator.
-    _ident_inner = r'(?:[\\_a-z]|[^\x00-\x7f])(?:[\\\w]|[^\x00-\x7f])*'
-    # But not inside strings.
-    _ident_nons = r'(?:[_a-z]|[^\x00-\x7f])(?:\w|[^\x00-\x7f])*'
-
-    flags = re.IGNORECASE | re.DOTALL | re.MULTILINE
-    tokens = {
-        'root': [
-            (r'<\?(php)?', Comment.Preproc, 'php'),
-            (r'[^<]+', Other),
-            (r'<', Other)
-        ],
-        'php': [
-            (r'\?>', Comment.Preproc, '#pop'),
-            (r'(<<<)([\'"]?)(' + _ident_nons + r')(\2\n.*?\n\s*)(\3)(;?)(\n)',
-             bygroups(String, String, String.Delimiter, String, String.Delimiter,
-                      Punctuation, Text)),
-            (r'\s+', Text),
-            (r'#\[', Punctuation, 'attribute'),
-            (r'#.*?\n', Comment.Single),
-            (r'//.*?\n', Comment.Single),
-            # put the empty comment here, it is otherwise seen as
-            # the start of a docstring
-            (r'/\*\*/', Comment.Multiline),
-            (r'/\*\*.*?\*/', String.Doc),
-            (r'/\*.*?\*/', Comment.Multiline),
-            (r'(->|::)(\s*)(' + _ident_nons + ')',
-             bygroups(Operator, Text, Name.Attribute)),
-            (r'[~!%^&*+=|:.<>/@-]+', Operator),
-            (r'\?', Operator),  # don't add to the charclass above!
-            (r'[\[\]{}();,]+', Punctuation),
-            (r'(new)(\s+)(class)\b', bygroups(Keyword, Text, Keyword)),
-            (r'(class)(\s+)', bygroups(Keyword, Text), 'classname'),
-            (r'(function)(\s*)(?=\()', bygroups(Keyword, Text)),
-            (r'(function)(\s+)(&?)(\s*)',
-             bygroups(Keyword, Text, Operator, Text), 'functionname'),
-            (r'(const)(\s+)(' + _ident_inner + ')',
-             bygroups(Keyword, Text, Name.Constant)),
-            (r'(and|E_PARSE|old_function|E_ERROR|or|as|E_WARNING|parent|'
-             r'eval|PHP_OS|break|exit|case|extends|PHP_VERSION|cfunction|'
-             r'FALSE|print|for|require|continue|foreach|require_once|'
-             r'declare|return|default|static|do|switch|die|stdClass|'
-             r'echo|else|TRUE|elseif|var|empty|if|xor|enddeclare|include|'
-             r'virtual|endfor|include_once|while|endforeach|global|'
-             r'endif|list|endswitch|new|endwhile|not|'
-             r'array|E_ALL|NULL|final|php_user_filter|interface|'
-             r'implements|public|private|protected|abstract|clone|try|'
-             r'catch|throw|this|use|namespace|trait|yield|'
-             r'finally|match)\b', Keyword),
-            (r'(true|false|null)\b', Keyword.Constant),
-            include('magicconstants'),
-            (r'\$\{', Name.Variable, 'variablevariable'),
-            (r'\$+' + _ident_inner, Name.Variable),
-            (_ident_inner, Name.Other),
-            (r'(\d+\.\d*|\d*\.\d+)(e[+-]?[0-9]+)?', Number.Float),
-            (r'\d+e[+-]?[0-9]+', Number.Float),
-            (r'0[0-7]+', Number.Oct),
-            (r'0x[a-f0-9]+', Number.Hex),
-            (r'\d+', Number.Integer),
-            (r'0b[01]+', Number.Bin),
-            (r"'([^'\\]*(?:\\.[^'\\]*)*)'", String.Single),
-            (r'`([^`\\]*(?:\\.[^`\\]*)*)`', String.Backtick),
-            (r'"', String.Double, 'string'),
-        ],
-        'variablevariable': [
-            (r'\}', Name.Variable, '#pop'),
-            include('php')
-        ],
-        'magicfuncs': [
-            # source: http://php.net/manual/en/language.oop5.magic.php
-            (words((
-                '__construct', '__destruct', '__call', '__callStatic', '__get', '__set',
-                '__isset', '__unset', '__sleep', '__wakeup', '__toString', '__invoke',
-                '__set_state', '__clone', '__debugInfo',), suffix=r'\b'),
-             Name.Function.Magic),
-        ],
-        'magicconstants': [
-            # source: http://php.net/manual/en/language.constants.predefined.php
-            (words((
-                '__LINE__', '__FILE__', '__DIR__', '__FUNCTION__', '__CLASS__',
-                '__TRAIT__', '__METHOD__', '__NAMESPACE__',),
-                suffix=r'\b'),
-             Name.Constant),
-        ],
-        'classname': [
-            (_ident_inner, Name.Class, '#pop')
-        ],
-        'functionname': [
-            include('magicfuncs'),
-            (_ident_inner, Name.Function, '#pop'),
-            default('#pop')
-        ],
-        'string': [
-            (r'"', String.Double, '#pop'),
-            (r'[^{$"\\]+', String.Double),
-            (r'\\([nrt"$\\]|[0-7]{1,3}|x[0-9a-f]{1,2})', String.Escape),
-            (r'\$' + _ident_nons + r'(\[\S+?\]|->' + _ident_nons + ')?',
-             String.Interpol),
-            (r'(\{\$\{)(.*?)(\}\})',
-             bygroups(String.Interpol, using(this, _startinline=True),
-                      String.Interpol)),
-            (r'(\{)(\$.*?)(\})',
-             bygroups(String.Interpol, using(this, _startinline=True),
-                      String.Interpol)),
-            (r'(\$\{)(\S+)(\})',
-             bygroups(String.Interpol, Name.Variable, String.Interpol)),
-            (r'[${\\]', String.Double)
-        ],
-        'attribute': [
-            (r'\]', Punctuation, '#pop'),
-            (r'\(', Punctuation, 'attributeparams'),
-            (_ident_inner, Name.Decorator),
-            include('php')
-        ],
-        'attributeparams': [
-            (r'\)', Punctuation, '#pop'),
-            include('php')
-        ],
-    }
-
-    def __init__(self, **options):
-        self.funcnamehighlighting = get_bool_opt(
-            options, 'funcnamehighlighting', True)
-        self.disabledmodules = get_list_opt(
-            options, 'disabledmodules', ['unknown'])
-        self.startinline = get_bool_opt(options, 'startinline', False)
-
-        # private option argument for the lexer itself
-        if '_startinline' in options:
-            self.startinline = options.pop('_startinline')
-
-        # collect activated functions in a set
-        self._functions = set()
-        if self.funcnamehighlighting:
-            from pygments.lexers._php_builtins import MODULES
-            for key, value in MODULES.items():
-                if key not in self.disabledmodules:
-                    self._functions.update(value)
-        RegexLexer.__init__(self, **options)
-
-    def get_tokens_unprocessed(self, text):
-        stack = ['root']
-        if self.startinline:
-            stack.append('php')
-        for index, token, value in \
-                RegexLexer.get_tokens_unprocessed(self, text, stack):
-            if token is Name.Other:
-                if value in self._functions:
-                    yield index, Name.Builtin, value
-                    continue
-            yield index, token, value
-
-    def analyse_text(text):
-        if shebang_matches(text, r'php'):
-            return True
-        rv = 0.0
-        if re.search(r'<\?(?!xml)', text):
-            rv += 0.3
-        return rv
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/pointless.py b/venv/lib/python3.10/site-packages/pygments/lexers/pointless.py
deleted file mode 100644
index adedb757d891928fe415b11dfae0917f430a7a6d..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/pointless.py
+++ /dev/null
@@ -1,70 +0,0 @@
-"""
-    pygments.lexers.pointless
-    ~~~~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Pointless.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, words
-from pygments.token import Comment, Error, Keyword, Name, Number, Operator, \
-    Punctuation, String, Text
-
-__all__ = ['PointlessLexer']
-
-
-class PointlessLexer(RegexLexer):
-    """
-    For Pointless source code.
-    """
-
-    name = 'Pointless'
-    url = 'https://ptls.dev'
-    aliases = ['pointless']
-    filenames = ['*.ptls']
-    version_added = '2.7'
-
-    ops = words([
-        "+", "-", "*", "/", "**", "%", "+=", "-=", "*=",
-        "/=", "**=", "%=", "|>", "=", "==", "!=", "<", ">",
-        "<=", ">=", "=>", "$", "++",
-    ])
-
-    keywords = words([
-        "if", "then", "else", "where", "with", "cond",
-        "case", "and", "or", "not", "in", "as", "for",
-        "requires", "throw", "try", "catch", "when",
-        "yield", "upval",
-    ], suffix=r'\b')
-
-    tokens = {
-        'root': [
-            (r'[ \n\r]+', Text),
-            (r'--.*$', Comment.Single),
-            (r'"""', String, 'multiString'),
-            (r'"', String, 'string'),
-            (r'[\[\](){}:;,.]', Punctuation),
-            (ops, Operator),
-            (keywords, Keyword),
-            (r'\d+|\d*\.\d+', Number),
-            (r'(true|false)\b', Name.Builtin),
-            (r'[A-Z][a-zA-Z0-9]*\b', String.Symbol),
-            (r'output\b', Name.Variable.Magic),
-            (r'(export|import)\b', Keyword.Namespace),
-            (r'[a-z][a-zA-Z0-9]*\b', Name.Variable)
-        ],
-        'multiString': [
-            (r'\\.', String.Escape),
-            (r'"""', String, '#pop'),
-            (r'"', String),
-            (r'[^\\"]+', String),
-        ],
-        'string': [
-            (r'\\.', String.Escape),
-            (r'"', String, '#pop'),
-            (r'\n', Error),
-            (r'[^\\"]+', String),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/pony.py b/venv/lib/python3.10/site-packages/pygments/lexers/pony.py
deleted file mode 100644
index 055423a4694c3aa9ecb5f3ff6e64d17d6487be93..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/pony.py
+++ /dev/null
@@ -1,93 +0,0 @@
-"""
-    pygments.lexers.pony
-    ~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Pony and related languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, bygroups, words
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation
-
-__all__ = ['PonyLexer']
-
-
-class PonyLexer(RegexLexer):
-    """
-    For Pony source code.
-    """
-
-    name = 'Pony'
-    aliases = ['pony']
-    filenames = ['*.pony']
-    url = 'https://www.ponylang.io'
-    version_added = '2.4'
-
-    _caps = r'(iso|trn|ref|val|box|tag)'
-
-    tokens = {
-        'root': [
-            (r'\n', Text),
-            (r'[^\S\n]+', Text),
-            (r'//.*\n', Comment.Single),
-            (r'/\*', Comment.Multiline, 'nested_comment'),
-            (r'"""(?:.|\n)*?"""', String.Doc),
-            (r'"', String, 'string'),
-            (r'\'.*\'', String.Char),
-            (r'=>|[]{}:().~;,|&!^?[]', Punctuation),
-            (words((
-                'addressof', 'and', 'as', 'consume', 'digestof', 'is', 'isnt',
-                'not', 'or'),
-                suffix=r'\b'),
-             Operator.Word),
-            (r'!=|==|<<|>>|[-+/*%=<>]', Operator),
-            (words((
-                'box', 'break', 'compile_error', 'compile_intrinsic',
-                'continue', 'do', 'else', 'elseif', 'embed', 'end', 'error',
-                'for', 'if', 'ifdef', 'in', 'iso', 'lambda', 'let', 'match',
-                'object', 'recover', 'ref', 'repeat', 'return', 'tag', 'then',
-                'this', 'trn', 'try', 'until', 'use', 'var', 'val', 'where',
-                'while', 'with', '#any', '#read', '#send', '#share'),
-                suffix=r'\b'),
-             Keyword),
-            (r'(actor|class|struct|primitive|interface|trait|type)((?:\s)+)',
-             bygroups(Keyword, Text), 'typename'),
-            (r'(new|fun|be)((?:\s)+)', bygroups(Keyword, Text), 'methodname'),
-            (words((
-                'I8', 'U8', 'I16', 'U16', 'I32', 'U32', 'I64', 'U64', 'I128',
-                'U128', 'ILong', 'ULong', 'ISize', 'USize', 'F32', 'F64',
-                'Bool', 'Pointer', 'None', 'Any', 'Array', 'String',
-                'Iterator'),
-                suffix=r'\b'),
-             Name.Builtin.Type),
-            (r'_?[A-Z]\w*', Name.Type),
-            (r'(\d+\.\d*|\.\d+|\d+)[eE][+-]?\d+', Number.Float),
-            (r'0x[0-9a-fA-F]+', Number.Hex),
-            (r'\d+', Number.Integer),
-            (r'(true|false)\b', Name.Builtin),
-            (r'_\d*', Name),
-            (r'_?[a-z][\w\']*', Name)
-        ],
-        'typename': [
-            (_caps + r'?((?:\s)*)(_?[A-Z]\w*)',
-             bygroups(Keyword, Text, Name.Class), '#pop')
-        ],
-        'methodname': [
-            (_caps + r'?((?:\s)*)(_?[a-z]\w*)',
-             bygroups(Keyword, Text, Name.Function), '#pop')
-        ],
-        'nested_comment': [
-            (r'[^*/]+', Comment.Multiline),
-            (r'/\*', Comment.Multiline, '#push'),
-            (r'\*/', Comment.Multiline, '#pop'),
-            (r'[*/]', Comment.Multiline)
-        ],
-        'string': [
-            (r'"', String, '#pop'),
-            (r'\\"', String),
-            (r'[^\\"]+', String)
-        ]
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/praat.py b/venv/lib/python3.10/site-packages/pygments/lexers/praat.py
deleted file mode 100644
index 054f5b61e56a53d2bf04d2bba2fc15c1fa85566d..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/praat.py
+++ /dev/null
@@ -1,303 +0,0 @@
-"""
-    pygments.lexers.praat
-    ~~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for Praat
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, words, bygroups, include
-from pygments.token import Name, Text, Comment, Keyword, String, Punctuation, \
-    Number, Operator, Whitespace
-
-__all__ = ['PraatLexer']
-
-
-class PraatLexer(RegexLexer):
-    """
-    For Praat scripts.
-    """
-
-    name = 'Praat'
-    url = 'http://www.praat.org'
-    aliases = ['praat']
-    filenames = ['*.praat', '*.proc', '*.psc']
-    version_added = '2.1'
-
-    keywords = (
-        'if', 'then', 'else', 'elsif', 'elif', 'endif', 'fi', 'for', 'from', 'to',
-        'endfor', 'endproc', 'while', 'endwhile', 'repeat', 'until', 'select', 'plus',
-        'minus', 'demo', 'assert', 'stopwatch', 'nocheck', 'nowarn', 'noprogress',
-        'editor', 'endeditor', 'clearinfo',
-    )
-
-    functions_string = (
-        'backslashTrigraphsToUnicode', 'chooseDirectory', 'chooseReadFile',
-        'chooseWriteFile', 'date', 'demoKey', 'do', 'environment', 'extractLine',
-        'extractWord', 'fixed', 'info', 'left', 'mid', 'percent', 'readFile', 'replace',
-        'replace_regex', 'right', 'selected', 'string', 'unicodeToBackslashTrigraphs',
-    )
-
-    functions_numeric = (
-        'abs', 'appendFile', 'appendFileLine', 'appendInfo', 'appendInfoLine', 'arccos',
-        'arccosh', 'arcsin', 'arcsinh', 'arctan', 'arctan2', 'arctanh', 'barkToHertz',
-        'beginPause', 'beginSendPraat', 'besselI', 'besselK', 'beta', 'beta2',
-        'binomialP', 'binomialQ', 'boolean', 'ceiling', 'chiSquareP', 'chiSquareQ',
-        'choice', 'comment', 'cos', 'cosh', 'createDirectory', 'deleteFile',
-        'demoClicked', 'demoClickedIn', 'demoCommandKeyPressed',
-        'demoExtraControlKeyPressed', 'demoInput', 'demoKeyPressed',
-        'demoOptionKeyPressed', 'demoShiftKeyPressed', 'demoShow', 'demoWaitForInput',
-        'demoWindowTitle', 'demoX', 'demoY', 'differenceLimensToPhon', 'do', 'editor',
-        'endPause', 'endSendPraat', 'endsWith', 'erb', 'erbToHertz', 'erf', 'erfc',
-        'exitScript', 'exp', 'extractNumber', 'fileReadable', 'fisherP', 'fisherQ',
-        'floor', 'gaussP', 'gaussQ', 'hertzToBark', 'hertzToErb', 'hertzToMel',
-        'hertzToSemitones', 'imax', 'imin', 'incompleteBeta', 'incompleteGammaP', 'index',
-        'index_regex', 'integer', 'invBinomialP', 'invBinomialQ', 'invChiSquareQ', 'invFisherQ',
-        'invGaussQ', 'invSigmoid', 'invStudentQ', 'length', 'ln', 'lnBeta', 'lnGamma',
-        'log10', 'log2', 'max', 'melToHertz', 'min', 'minusObject', 'natural', 'number',
-        'numberOfColumns', 'numberOfRows', 'numberOfSelected', 'objectsAreIdentical',
-        'option', 'optionMenu', 'pauseScript', 'phonToDifferenceLimens', 'plusObject',
-        'positive', 'randomBinomial', 'randomGauss', 'randomInteger', 'randomPoisson',
-        'randomUniform', 'real', 'readFile', 'removeObject', 'rindex', 'rindex_regex',
-        'round', 'runScript', 'runSystem', 'runSystem_nocheck', 'selectObject',
-        'selected', 'semitonesToHertz', 'sentence', 'sentencetext', 'sigmoid', 'sin', 'sinc',
-        'sincpi', 'sinh', 'soundPressureToPhon', 'sqrt', 'startsWith', 'studentP',
-        'studentQ', 'tan', 'tanh', 'text', 'variableExists', 'word', 'writeFile', 'writeFileLine',
-        'writeInfo', 'writeInfoLine',
-    )
-
-    functions_array = (
-        'linear', 'randomGauss', 'randomInteger', 'randomUniform', 'zero',
-    )
-
-    objects = (
-        'Activation', 'AffineTransform', 'AmplitudeTier', 'Art', 'Artword',
-        'Autosegment', 'BarkFilter', 'BarkSpectrogram', 'CCA', 'Categories',
-        'Cepstrogram', 'Cepstrum', 'Cepstrumc', 'ChebyshevSeries', 'ClassificationTable',
-        'Cochleagram', 'Collection', 'ComplexSpectrogram', 'Configuration', 'Confusion',
-        'ContingencyTable', 'Corpus', 'Correlation', 'Covariance',
-        'CrossCorrelationTable', 'CrossCorrelationTables', 'DTW', 'DataModeler',
-        'Diagonalizer', 'Discriminant', 'Dissimilarity', 'Distance', 'Distributions',
-        'DurationTier', 'EEG', 'ERP', 'ERPTier', 'EditCostsTable', 'EditDistanceTable',
-        'Eigen', 'Excitation', 'Excitations', 'ExperimentMFC', 'FFNet', 'FeatureWeights',
-        'FileInMemory', 'FilesInMemory', 'Formant', 'FormantFilter', 'FormantGrid',
-        'FormantModeler', 'FormantPoint', 'FormantTier', 'GaussianMixture', 'HMM',
-        'HMM_Observation', 'HMM_ObservationSequence', 'HMM_State', 'HMM_StateSequence',
-        'Harmonicity', 'ISpline', 'Index', 'Intensity', 'IntensityTier', 'IntervalTier',
-        'KNN', 'KlattGrid', 'KlattTable', 'LFCC', 'LPC', 'Label', 'LegendreSeries',
-        'LinearRegression', 'LogisticRegression', 'LongSound', 'Ltas', 'MFCC', 'MSpline',
-        'ManPages', 'Manipulation', 'Matrix', 'MelFilter', 'MelSpectrogram',
-        'MixingMatrix', 'Movie', 'Network', 'Object', 'OTGrammar', 'OTHistory', 'OTMulti',
-        'PCA', 'PairDistribution', 'ParamCurve', 'Pattern', 'Permutation', 'Photo',
-        'Pitch', 'PitchModeler', 'PitchTier', 'PointProcess', 'Polygon', 'Polynomial',
-        'PowerCepstrogram', 'PowerCepstrum', 'Procrustes', 'RealPoint', 'RealTier',
-        'ResultsMFC', 'Roots', 'SPINET', 'SSCP', 'SVD', 'Salience', 'ScalarProduct',
-        'Similarity', 'SimpleString', 'SortedSetOfString', 'Sound', 'Speaker',
-        'Spectrogram', 'Spectrum', 'SpectrumTier', 'SpeechSynthesizer', 'SpellingChecker',
-        'Strings', 'StringsIndex', 'Table', 'TableOfReal', 'TextGrid', 'TextInterval',
-        'TextPoint', 'TextTier', 'Tier', 'Transition', 'VocalTract', 'VocalTractTier',
-        'Weight', 'WordList',
-    )
-
-    variables_numeric = (
-        'macintosh', 'windows', 'unix', 'praatVersion', 'pi', 'e', 'undefined',
-    )
-
-    variables_string = (
-        'praatVersion', 'tab', 'shellDirectory', 'homeDirectory',
-        'preferencesDirectory', 'newline', 'temporaryDirectory',
-        'defaultDirectory',
-    )
-
-    object_attributes = (
-        'ncol', 'nrow', 'xmin', 'ymin', 'xmax', 'ymax', 'nx', 'ny', 'dx', 'dy',
-    )
-
-    tokens = {
-        'root': [
-            (r'(\s+)(#.*?$)',  bygroups(Whitespace, Comment.Single)),
-            (r'^#.*?$',        Comment.Single),
-            (r';[^\n]*',       Comment.Single),
-            (r'\s+',           Whitespace),
-
-            (r'\bprocedure\b', Keyword,       'procedure_definition'),
-            (r'\bcall\b',      Keyword,       'procedure_call'),
-            (r'@',             Name.Function, 'procedure_call'),
-
-            include('function_call'),
-
-            (words(keywords, suffix=r'\b'), Keyword),
-
-            (r'(\bform\b)(\s+)([^\n]+)',
-             bygroups(Keyword, Whitespace, String), 'old_form'),
-
-            (r'(print(?:line|tab)?|echo|exit|asserterror|pause|send(?:praat|socket)|'
-             r'include|execute|system(?:_nocheck)?)(\s+)',
-             bygroups(Keyword, Whitespace), 'string_unquoted'),
-
-            (r'(goto|label)(\s+)(\w+)', bygroups(Keyword, Whitespace, Name.Label)),
-
-            include('variable_name'),
-            include('number'),
-
-            (r'"', String, 'string'),
-
-            (words((objects), suffix=r'(?=\s+\S+\n)'), Name.Class, 'string_unquoted'),
-
-            (r'\b[A-Z]', Keyword, 'command'),
-            (r'(\.{3}|[)(,])', Punctuation),
-        ],
-        'command': [
-            (r'( ?[\w()-]+ ?)', Keyword),
-
-            include('string_interpolated'),
-
-            (r'\.{3}', Keyword, ('#pop', 'old_arguments')),
-            (r':', Keyword, ('#pop', 'comma_list')),
-            (r'\s', Whitespace, '#pop'),
-        ],
-        'procedure_call': [
-            (r'\s+', Whitespace),
-            (r'([\w.]+)(?:(:)|(?:(\s*)(\()))',
-             bygroups(Name.Function, Punctuation,
-                      Text.Whitespace, Punctuation), '#pop'),
-            (r'([\w.]+)', Name.Function, ('#pop', 'old_arguments')),
-        ],
-        'procedure_definition': [
-            (r'\s', Whitespace),
-            (r'([\w.]+)(\s*?[(:])',
-             bygroups(Name.Function, Whitespace), '#pop'),
-            (r'([\w.]+)([^\n]*)',
-             bygroups(Name.Function, Text), '#pop'),
-        ],
-        'function_call': [
-            (words(functions_string, suffix=r'\$(?=\s*[:(])'), Name.Function, 'function'),
-            (words(functions_array, suffix=r'#(?=\s*[:(])'),   Name.Function, 'function'),
-            (words(functions_numeric, suffix=r'(?=\s*[:(])'),  Name.Function, 'function'),
-        ],
-        'function': [
-            (r'\s+',   Whitespace),
-            (r':',     Punctuation, ('#pop', 'comma_list')),
-            (r'\s*\(', Punctuation, ('#pop', 'comma_list')),
-        ],
-        'comma_list': [
-            (r'(\s*\n\s*)(\.{3})', bygroups(Whitespace, Punctuation)),
-
-            (r'(\s*)(?:([)\]])|(\n))', bygroups(
-                Whitespace, Punctuation, Whitespace), '#pop'),
-
-            (r'\s+', Whitespace),
-            (r'"',   String, 'string'),
-            (r'\b(if|then|else|fi|endif)\b', Keyword),
-
-            include('function_call'),
-            include('variable_name'),
-            include('operator'),
-            include('number'),
-
-            (r'[()]', Text),
-            (r',', Punctuation),
-        ],
-        'old_arguments': [
-            (r'\n', Whitespace, '#pop'),
-
-            include('variable_name'),
-            include('operator'),
-            include('number'),
-
-            (r'"', String, 'string'),
-            (r'[^\n]', Text),
-        ],
-        'number': [
-            (r'\n', Whitespace, '#pop'),
-            (r'\b\d+(\.\d*)?([eE][-+]?\d+)?%?', Number),
-        ],
-        'object_reference': [
-            include('string_interpolated'),
-            (r'([a-z][a-zA-Z0-9_]*|\d+)', Name.Builtin),
-
-            (words(object_attributes, prefix=r'\.'), Name.Builtin, '#pop'),
-
-            (r'\$', Name.Builtin),
-            (r'\[', Text, '#pop'),
-        ],
-        'variable_name': [
-            include('operator'),
-            include('number'),
-
-            (words(variables_string,  suffix=r'\$'), Name.Variable.Global),
-            (words(variables_numeric,
-             suffix=r'(?=[^a-zA-Z0-9_."\'$#\[:(]|\s|^|$)'),
-             Name.Variable.Global),
-
-            (words(objects, prefix=r'\b', suffix=r"(_)"),
-             bygroups(Name.Builtin, Name.Builtin),
-             'object_reference'),
-
-            (r'\.?_?[a-z][\w.]*(\$|#)?', Text),
-            (r'[\[\]]', Punctuation, 'comma_list'),
-
-            include('string_interpolated'),
-        ],
-        'operator': [
-            (r'([+\/*<>=!-]=?|[&*|][&*|]?|\^|<>)',       Operator),
-            (r'(?', Punctuation),
-            (r'"(?:\\x[0-9a-fA-F]+\\|\\u[0-9a-fA-F]{4}|\\U[0-9a-fA-F]{8}|'
-             r'\\[0-7]+\\|\\["\\abcefnrstv]|[^\\"])*"', String.Double),
-            (r"'(?:''|[^'])*'", String.Atom),  # quoted atom
-            # Needs to not be followed by an atom.
-            # (r'=(?=\s|[a-zA-Z\[])', Operator),
-            (r'is\b', Operator),
-            (r'(<|>|=<|>=|==|=:=|=|/|//|\*|\+|-)(?=\s|[a-zA-Z0-9\[])',
-             Operator),
-            (r'(mod|div|not)\b', Operator),
-            (r'_', Keyword),  # The don't-care variable
-            (r'([a-z]+)(:)', bygroups(Name.Namespace, Punctuation)),
-            (r'([a-z\u00c0-\u1fff\u3040-\ud7ff\ue000-\uffef]'
-             r'[\w$\u00c0-\u1fff\u3040-\ud7ff\ue000-\uffef]*)'
-             r'(\s*)(:-|-->)',
-             bygroups(Name.Function, Text, Operator)),  # function defn
-            (r'([a-z\u00c0-\u1fff\u3040-\ud7ff\ue000-\uffef]'
-             r'[\w$\u00c0-\u1fff\u3040-\ud7ff\ue000-\uffef]*)'
-             r'(\s*)(\()',
-             bygroups(Name.Function, Text, Punctuation)),
-            (r'[a-z\u00c0-\u1fff\u3040-\ud7ff\ue000-\uffef]'
-             r'[\w$\u00c0-\u1fff\u3040-\ud7ff\ue000-\uffef]*',
-             String.Atom),  # atom, characters
-            # This one includes !
-            (r'[#&*+\-./:<=>?@\\^~\u00a1-\u00bf\u2010-\u303f]+',
-             String.Atom),  # atom, graphics
-            (r'[A-Z_]\w*', Name.Variable),
-            (r'\s+|[\u2000-\u200f\ufff0-\ufffe\uffef]', Text),
-        ],
-        'nested-comment': [
-            (r'\*/', Comment.Multiline, '#pop'),
-            (r'/\*', Comment.Multiline, '#push'),
-            (r'[^*/]+', Comment.Multiline),
-            (r'[*/]', Comment.Multiline),
-        ],
-    }
-
-    def analyse_text(text):
-        """Competes with IDL and Visual Prolog on *.pro"""
-        if ':-' in text:
-            # Visual Prolog also uses :-
-            return 0.5
-        else:
-            return 0
-
-
-class LogtalkLexer(RegexLexer):
-    """
-    For Logtalk source code.
-    """
-
-    name = 'Logtalk'
-    url = 'http://logtalk.org/'
-    aliases = ['logtalk']
-    filenames = ['*.lgt', '*.logtalk']
-    mimetypes = ['text/x-logtalk']
-    version_added = '0.10'
-
-    tokens = {
-        'root': [
-            # Directives
-            (r'^\s*:-\s', Punctuation, 'directive'),
-            # Comments
-            (r'%.*?\n', Comment),
-            (r'/\*(.|\n)*?\*/', Comment),
-            # Whitespace
-            (r'\n', Text),
-            (r'\s+', Text),
-            # Numbers
-            (r"0'[\\]?.", Number),
-            (r'0b[01]+', Number.Bin),
-            (r'0o[0-7]+', Number.Oct),
-            (r'0x[0-9a-fA-F]+', Number.Hex),
-            (r'\d+\.?\d*((e|E)(\+|-)?\d+)?', Number),
-            # Variables
-            (r'([A-Z_][a-zA-Z0-9_]*)', Name.Variable),
-            # Event handlers
-            (r'(after|before)(?=[(])', Keyword),
-            # Message forwarding handler
-            (r'forward(?=[(])', Keyword),
-            # Execution-context methods
-            (r'(context|parameter|this|se(lf|nder))(?=[(])', Keyword),
-            # Reflection
-            (r'(current_predicate|predicate_property)(?=[(])', Keyword),
-            # DCGs and term expansion
-            (r'(expand_(goal|term)|(goal|term)_expansion|phrase)(?=[(])', Keyword),
-            # Entity
-            (r'(abolish|c(reate|urrent))_(object|protocol|category)(?=[(])', Keyword),
-            (r'(object|protocol|category)_property(?=[(])', Keyword),
-            # Entity relations
-            (r'co(mplements_object|nforms_to_protocol)(?=[(])', Keyword),
-            (r'extends_(object|protocol|category)(?=[(])', Keyword),
-            (r'imp(lements_protocol|orts_category)(?=[(])', Keyword),
-            (r'(instantiat|specializ)es_class(?=[(])', Keyword),
-            # Events
-            (r'(current_event|(abolish|define)_events)(?=[(])', Keyword),
-            # Flags
-            (r'(create|current|set)_logtalk_flag(?=[(])', Keyword),
-            # Compiling, loading, and library paths
-            (r'logtalk_(compile|l(ibrary_path|oad|oad_context)|make(_target_action)?)(?=[(])', Keyword),
-            (r'\blogtalk_make\b', Keyword),
-            # Database
-            (r'(clause|retract(all)?)(?=[(])', Keyword),
-            (r'a(bolish|ssert(a|z))(?=[(])', Keyword),
-            # Control constructs
-            (r'(ca(ll|tch)|throw)(?=[(])', Keyword),
-            (r'(fa(il|lse)|true|(instantiation|system)_error)\b', Keyword),
-            (r'(uninstantiation|type|domain|existence|permission|representation|evaluation|resource|syntax)_error(?=[(])', Keyword),
-            # All solutions
-            (r'((bag|set)of|f(ind|or)all)(?=[(])', Keyword),
-            # Multi-threading predicates
-            (r'threaded(_(ca(ll|ncel)|once|ignore|exit|peek|wait|notify))?(?=[(])', Keyword),
-            # Engine predicates
-            (r'threaded_engine(_(create|destroy|self|next|next_reified|yield|post|fetch))?(?=[(])', Keyword),
-            # Term unification
-            (r'(subsumes_term|unify_with_occurs_check)(?=[(])', Keyword),
-            # Term creation and decomposition
-            (r'(functor|arg|copy_term|numbervars|term_variables)(?=[(])', Keyword),
-            # Evaluable functors
-            (r'(div|rem|m(ax|in|od)|abs|sign)(?=[(])', Keyword),
-            (r'float(_(integer|fractional)_part)?(?=[(])', Keyword),
-            (r'(floor|t(an|runcate)|round|ceiling)(?=[(])', Keyword),
-            # Other arithmetic functors
-            (r'(cos|a(cos|sin|tan|tan2)|exp|log|s(in|qrt)|xor)(?=[(])', Keyword),
-            # Term testing
-            (r'(var|atom(ic)?|integer|float|c(allable|ompound)|n(onvar|umber)|ground|acyclic_term)(?=[(])', Keyword),
-            # Term comparison
-            (r'compare(?=[(])', Keyword),
-            # Stream selection and control
-            (r'(curren|se)t_(in|out)put(?=[(])', Keyword),
-            (r'(open|close)(?=[(])', Keyword),
-            (r'flush_output(?=[(])', Keyword),
-            (r'(at_end_of_stream|flush_output)\b', Keyword),
-            (r'(stream_property|at_end_of_stream|set_stream_position)(?=[(])', Keyword),
-            # Character and byte input/output
-            (r'(nl|(get|peek|put)_(byte|c(har|ode)))(?=[(])', Keyword),
-            (r'\bnl\b', Keyword),
-            # Term input/output
-            (r'read(_term)?(?=[(])', Keyword),
-            (r'write(q|_(canonical|term))?(?=[(])', Keyword),
-            (r'(current_)?op(?=[(])', Keyword),
-            (r'(current_)?char_conversion(?=[(])', Keyword),
-            # Atomic term processing
-            (r'atom_(length|c(hars|o(ncat|des)))(?=[(])', Keyword),
-            (r'(char_code|sub_atom)(?=[(])', Keyword),
-            (r'number_c(har|ode)s(?=[(])', Keyword),
-            # Implementation defined hooks functions
-            (r'(se|curren)t_prolog_flag(?=[(])', Keyword),
-            (r'\bhalt\b', Keyword),
-            (r'halt(?=[(])', Keyword),
-            # Message sending operators
-            (r'(::|:|\^\^)', Operator),
-            # External call
-            (r'[{}]', Keyword),
-            # Logic and control
-            (r'(ignore|once)(?=[(])', Keyword),
-            (r'\brepeat\b', Keyword),
-            # Sorting
-            (r'(key)?sort(?=[(])', Keyword),
-            # Bitwise functors
-            (r'(>>|<<|/\\|\\\\|\\)', Operator),
-            # Predicate aliases
-            (r'\bas\b', Operator),
-            # Arithmetic evaluation
-            (r'\bis\b', Keyword),
-            # Arithmetic comparison
-            (r'(=:=|=\\=|<|=<|>=|>)', Operator),
-            # Term creation and decomposition
-            (r'=\.\.', Operator),
-            # Term unification
-            (r'(=|\\=)', Operator),
-            # Term comparison
-            (r'(==|\\==|@=<|@<|@>=|@>)', Operator),
-            # Evaluable functors
-            (r'(//|[-+*/])', Operator),
-            (r'\b(e|pi|div|mod|rem)\b', Operator),
-            # Other arithmetic functors
-            (r'\b\*\*\b', Operator),
-            # DCG rules
-            (r'-->', Operator),
-            # Control constructs
-            (r'([!;]|->)', Operator),
-            # Logic and control
-            (r'\\+', Operator),
-            # Mode operators
-            (r'[?@]', Operator),
-            # Existential quantifier
-            (r'\^', Operator),
-            # Punctuation
-            (r'[()\[\],.|]', Text),
-            # Atoms
-            (r"[a-z][a-zA-Z0-9_]*", Text),
-            (r"'", String, 'quoted_atom'),
-            # Double-quoted terms
-            (r'"', String, 'double_quoted_term'),
-        ],
-
-        'quoted_atom': [
-            (r"''", String),
-            (r"'", String, '#pop'),
-            (r'\\([\\abfnrtv"\']|(x[a-fA-F0-9]+|[0-7]+)\\)', String.Escape),
-            (r"[^\\'\n]+", String),
-            (r'\\', String),
-        ],
-
-        'double_quoted_term': [
-            (r'""', String),
-            (r'"', String, '#pop'),
-            (r'\\([\\abfnrtv"\']|(x[a-fA-F0-9]+|[0-7]+)\\)', String.Escape),
-            (r'[^\\"\n]+', String),
-            (r'\\', String),
-        ],
-
-        'directive': [
-            # Conditional compilation directives
-            (r'(el)?if(?=[(])', Keyword, 'root'),
-            (r'(e(lse|ndif))(?=[.])', Keyword, 'root'),
-            # Entity directives
-            (r'(category|object|protocol)(?=[(])', Keyword, 'entityrelations'),
-            (r'(end_(category|object|protocol))(?=[.])', Keyword, 'root'),
-            # Predicate scope directives
-            (r'(public|protected|private)(?=[(])', Keyword, 'root'),
-            # Other directives
-            (r'e(n(coding|sure_loaded)|xport)(?=[(])', Keyword, 'root'),
-            (r'in(clude|itialization|fo)(?=[(])', Keyword, 'root'),
-            (r'(built_in|dynamic|synchronized|threaded)(?=[.])', Keyword, 'root'),
-            (r'(alias|d(ynamic|iscontiguous)|m(eta_(non_terminal|predicate)|ode|ultifile)|s(et_(logtalk|prolog)_flag|ynchronized))(?=[(])', Keyword, 'root'),
-            (r'op(?=[(])', Keyword, 'root'),
-            (r'(c(alls|oinductive)|module|reexport|use(s|_module))(?=[(])', Keyword, 'root'),
-            (r'[a-z][a-zA-Z0-9_]*(?=[(])', Text, 'root'),
-            (r'[a-z][a-zA-Z0-9_]*(?=[.])', Text, 'root'),
-        ],
-
-        'entityrelations': [
-            (r'(complements|extends|i(nstantiates|mp(lements|orts))|specializes)(?=[(])', Keyword),
-            # Numbers
-            (r"0'[\\]?.", Number),
-            (r'0b[01]+', Number.Bin),
-            (r'0o[0-7]+', Number.Oct),
-            (r'0x[0-9a-fA-F]+', Number.Hex),
-            (r'\d+\.?\d*((e|E)(\+|-)?\d+)?', Number),
-            # Variables
-            (r'([A-Z_][a-zA-Z0-9_]*)', Name.Variable),
-            # Atoms
-            (r"[a-z][a-zA-Z0-9_]*", Text),
-            (r"'", String, 'quoted_atom'),
-            # Double-quoted terms
-            (r'"', String, 'double_quoted_term'),
-            # End of entity-opening directive
-            (r'([)]\.)', Text, 'root'),
-            # Scope operator
-            (r'(::)', Operator),
-            # Punctuation
-            (r'[()\[\],.|]', Text),
-            # Comments
-            (r'%.*?\n', Comment),
-            (r'/\*(.|\n)*?\*/', Comment),
-            # Whitespace
-            (r'\n', Text),
-            (r'\s+', Text),
-        ]
-    }
-
-    def analyse_text(text):
-        if ':- object(' in text:
-            return 1.0
-        elif ':- protocol(' in text:
-            return 1.0
-        elif ':- category(' in text:
-            return 1.0
-        elif re.search(r'^:-\s[a-z]', text, re.M):
-            return 0.9
-        else:
-            return 0.0
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/promql.py b/venv/lib/python3.10/site-packages/pygments/lexers/promql.py
deleted file mode 100644
index cad3c254a1202682f4a831d4c09472bce34c33f0..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/promql.py
+++ /dev/null
@@ -1,176 +0,0 @@
-"""
-    pygments.lexers.promql
-    ~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for Prometheus Query Language.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, bygroups, default, words
-from pygments.token import Comment, Keyword, Name, Number, Operator, \
-    Punctuation, String, Whitespace
-
-__all__ = ["PromQLLexer"]
-
-
-class PromQLLexer(RegexLexer):
-    """
-    For PromQL queries.
-
-    For details about the grammar see:
-    https://github.com/prometheus/prometheus/tree/master/promql/parser
-
-    .. versionadded: 2.7
-    """
-
-    name = "PromQL"
-    url = 'https://prometheus.io/docs/prometheus/latest/querying/basics/'
-    aliases = ["promql"]
-    filenames = ["*.promql"]
-    version_added = ''
-
-    base_keywords = (
-        words(
-            (
-                "bool",
-                "by",
-                "group_left",
-                "group_right",
-                "ignoring",
-                "offset",
-                "on",
-                "without",
-            ),
-            suffix=r"\b",
-        ),
-        Keyword,
-    )
-
-    aggregator_keywords = (
-        words(
-            (
-                "sum",
-                "min",
-                "max",
-                "avg",
-                "group",
-                "stddev",
-                "stdvar",
-                "count",
-                "count_values",
-                "bottomk",
-                "topk",
-                "quantile",
-            ),
-            suffix=r"\b",
-        ),
-        Keyword,
-    )
-
-    function_keywords = (
-        words(
-            (
-                "abs",
-                "absent",
-                "absent_over_time",
-                "avg_over_time",
-                "ceil",
-                "changes",
-                "clamp_max",
-                "clamp_min",
-                "count_over_time",
-                "day_of_month",
-                "day_of_week",
-                "days_in_month",
-                "delta",
-                "deriv",
-                "exp",
-                "floor",
-                "histogram_quantile",
-                "holt_winters",
-                "hour",
-                "idelta",
-                "increase",
-                "irate",
-                "label_join",
-                "label_replace",
-                "ln",
-                "log10",
-                "log2",
-                "max_over_time",
-                "min_over_time",
-                "minute",
-                "month",
-                "predict_linear",
-                "quantile_over_time",
-                "rate",
-                "resets",
-                "round",
-                "scalar",
-                "sort",
-                "sort_desc",
-                "sqrt",
-                "stddev_over_time",
-                "stdvar_over_time",
-                "sum_over_time",
-                "time",
-                "timestamp",
-                "vector",
-                "year",
-            ),
-            suffix=r"\b",
-        ),
-        Keyword.Reserved,
-    )
-
-    tokens = {
-        "root": [
-            (r"\n", Whitespace),
-            (r"\s+", Whitespace),
-            (r",", Punctuation),
-            # Keywords
-            base_keywords,
-            aggregator_keywords,
-            function_keywords,
-            # Offsets
-            (r"[1-9][0-9]*[smhdwy]", String),
-            # Numbers
-            (r"-?[0-9]+\.[0-9]+", Number.Float),
-            (r"-?[0-9]+", Number.Integer),
-            # Comments
-            (r"#.*?$", Comment.Single),
-            # Operators
-            (r"(\+|\-|\*|\/|\%|\^)", Operator),
-            (r"==|!=|>=|<=|<|>", Operator),
-            (r"and|or|unless", Operator.Word),
-            # Metrics
-            (r"[_a-zA-Z][a-zA-Z0-9_]+", Name.Variable),
-            # Params
-            (r'(["\'])(.*?)(["\'])', bygroups(Punctuation, String, Punctuation)),
-            # Other states
-            (r"\(", Operator, "function"),
-            (r"\)", Operator),
-            (r"\{", Punctuation, "labels"),
-            (r"\[", Punctuation, "range"),
-        ],
-        "labels": [
-            (r"\}", Punctuation, "#pop"),
-            (r"\n", Whitespace),
-            (r"\s+", Whitespace),
-            (r",", Punctuation),
-            (r'([_a-zA-Z][a-zA-Z0-9_]*?)(\s*?)(=~|!=|=|!~)(\s*?)("|\')(.*?)("|\')',
-             bygroups(Name.Label, Whitespace, Operator, Whitespace,
-                      Punctuation, String, Punctuation)),
-        ],
-        "range": [
-            (r"\]", Punctuation, "#pop"),
-            (r"[1-9][0-9]*[smhdwy]", String),
-        ],
-        "function": [
-            (r"\)", Operator, "#pop"),
-            (r"\(", Operator, "#push"),
-            default("#pop"),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/prql.py b/venv/lib/python3.10/site-packages/pygments/lexers/prql.py
deleted file mode 100644
index ee95d2d47491a45479aa79a0e7c2c24a4cef01cf..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/prql.py
+++ /dev/null
@@ -1,251 +0,0 @@
-"""
-    pygments.lexers.prql
-    ~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for the PRQL query language.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, combined, words, include, bygroups
-from pygments.token import Comment, Literal, Keyword, Name, Number, Operator, \
-    Punctuation, String, Text, Whitespace
-
-__all__ = ['PrqlLexer']
-
-
-class PrqlLexer(RegexLexer):
-    """
-    For PRQL source code.
-
-    grammar: https://github.com/PRQL/prql/tree/main/grammars
-    """
-
-    name = 'PRQL'
-    url = 'https://prql-lang.org/'
-    aliases = ['prql']
-    filenames = ['*.prql']
-    mimetypes = ['application/prql', 'application/x-prql']
-    version_added = '2.17'
-
-    builtinTypes = words((
-        "bool",
-        "int",
-        "int8", "int16", "int32", "int64", "int128",
-        "float",
-        "text",
-        "set"), suffix=r'\b')
-
-    def innerstring_rules(ttype):
-        return [
-            # the new style '{}'.format(...) string formatting
-            (r'\{'
-             r'((\w+)((\.\w+)|(\[[^\]]+\]))*)?'  # field name
-             r'(\:(.?[<>=\^])?[-+ ]?#?0?(\d+)?,?(\.\d+)?[E-GXb-gnosx%]?)?'
-             r'\}', String.Interpol),
-
-            (r'[^\\\'"%{\n]+', ttype),
-            (r'[\'"\\]', ttype),
-            (r'%|(\{{1,2})', ttype)
-        ]
-
-    def fstring_rules(ttype):
-        return [
-            (r'\}', String.Interpol),
-            (r'\{', String.Interpol, 'expr-inside-fstring'),
-            (r'[^\\\'"{}\n]+', ttype),
-            (r'[\'"\\]', ttype),
-        ]
-
-    tokens = {
-        'root': [
-
-            # Comments
-            (r'#!.*', String.Doc),
-            (r'#.*', Comment.Single),
-
-            # Whitespace
-            (r'\s+', Whitespace),
-
-            # Modules
-            (r'^(\s*)(module)(\s*)',
-             bygroups(Whitespace, Keyword.Namespace, Whitespace),
-             'imports'),
-
-            (builtinTypes, Keyword.Type),
-
-            # Main
-            (r'^prql ', Keyword.Reserved),
-
-            ('let', Keyword.Declaration),
-
-            include('keywords'),
-            include('expr'),
-
-            # Transforms
-            (r'^[A-Za-z_][a-zA-Z0-9_]*', Keyword),
-        ],
-        'expr': [
-            # non-raw f-strings
-            ('(f)(""")', bygroups(String.Affix, String.Double),
-             combined('fstringescape', 'tdqf')),
-            ("(f)(''')", bygroups(String.Affix, String.Single),
-             combined('fstringescape', 'tsqf')),
-            ('(f)(")', bygroups(String.Affix, String.Double),
-             combined('fstringescape', 'dqf')),
-            ("(f)(')", bygroups(String.Affix, String.Single),
-             combined('fstringescape', 'sqf')),
-
-            # non-raw s-strings
-            ('(s)(""")', bygroups(String.Affix, String.Double),
-             combined('stringescape', 'tdqf')),
-            ("(s)(''')", bygroups(String.Affix, String.Single),
-             combined('stringescape', 'tsqf')),
-            ('(s)(")', bygroups(String.Affix, String.Double),
-             combined('stringescape', 'dqf')),
-            ("(s)(')", bygroups(String.Affix, String.Single),
-             combined('stringescape', 'sqf')),
-
-            # raw strings
-            ('(?i)(r)(""")',
-             bygroups(String.Affix, String.Double), 'tdqs'),
-            ("(?i)(r)(''')",
-             bygroups(String.Affix, String.Single), 'tsqs'),
-            ('(?i)(r)(")',
-             bygroups(String.Affix, String.Double), 'dqs'),
-            ("(?i)(r)(')",
-             bygroups(String.Affix, String.Single), 'sqs'),
-
-            # non-raw strings
-            ('"""', String.Double, combined('stringescape', 'tdqs')),
-            ("'''", String.Single, combined('stringescape', 'tsqs')),
-            ('"', String.Double, combined('stringescape', 'dqs')),
-            ("'", String.Single, combined('stringescape', 'sqs')),
-
-            # Time and dates
-            (r'@\d{4}-\d{2}-\d{2}T\d{2}(:\d{2})?(:\d{2})?(\.\d{1,6})?(Z|[+-]\d{1,2}(:\d{1,2})?)?', Literal.Date),
-            (r'@\d{4}-\d{2}-\d{2}', Literal.Date),
-            (r'@\d{2}(:\d{2})?(:\d{2})?(\.\d{1,6})?(Z|[+-]\d{1,2}(:\d{1,2})?)?', Literal.Date),
-
-            (r'[^\S\n]+', Text),
-            include('numbers'),
-            (r'->|=>|==|!=|>=|<=|~=|&&|\|\||\?\?|\/\/', Operator),
-            (r'[-~+/*%=<>&^|.@]', Operator),
-            (r'[]{}:(),;[]', Punctuation),
-            include('functions'),
-
-            # Variable Names
-            (r'[A-Za-z_][a-zA-Z0-9_]*', Name.Variable),
-        ],
-        'numbers': [
-            (r'(\d(?:_?\d)*\.(?:\d(?:_?\d)*)?|(?:\d(?:_?\d)*)?\.\d(?:_?\d)*)'
-             r'([eE][+-]?\d(?:_?\d)*)?', Number.Float),
-            (r'\d(?:_?\d)*[eE][+-]?\d(?:_?\d)*j?', Number.Float),
-            (r'0[oO](?:_?[0-7])+', Number.Oct),
-            (r'0[bB](?:_?[01])+', Number.Bin),
-            (r'0[xX](?:_?[a-fA-F0-9])+', Number.Hex),
-            (r'\d(?:_?\d)*', Number.Integer),
-        ],
-        'fstringescape': [
-            include('stringescape'),
-        ],
-        'bytesescape': [
-            (r'\\([\\bfnrt"\']|\n|x[a-fA-F0-9]{2}|[0-7]{1,3})', String.Escape)
-        ],
-        'stringescape': [
-            (r'\\(N\{.*?\}|u\{[a-fA-F0-9]{1,6}\})', String.Escape),
-            include('bytesescape')
-        ],
-        'fstrings-single': fstring_rules(String.Single),
-        'fstrings-double': fstring_rules(String.Double),
-        'strings-single': innerstring_rules(String.Single),
-        'strings-double': innerstring_rules(String.Double),
-        'dqf': [
-            (r'"', String.Double, '#pop'),
-            (r'\\\\|\\"|\\\n', String.Escape),  # included here for raw strings
-            include('fstrings-double')
-        ],
-        'sqf': [
-            (r"'", String.Single, '#pop'),
-            (r"\\\\|\\'|\\\n", String.Escape),  # included here for raw strings
-            include('fstrings-single')
-        ],
-        'dqs': [
-            (r'"', String.Double, '#pop'),
-            (r'\\\\|\\"|\\\n', String.Escape),  # included here for raw strings
-            include('strings-double')
-        ],
-        'sqs': [
-            (r"'", String.Single, '#pop'),
-            (r"\\\\|\\'|\\\n", String.Escape),  # included here for raw strings
-            include('strings-single')
-        ],
-        'tdqf': [
-            (r'"""', String.Double, '#pop'),
-            include('fstrings-double'),
-            (r'\n', String.Double)
-        ],
-        'tsqf': [
-            (r"'''", String.Single, '#pop'),
-            include('fstrings-single'),
-            (r'\n', String.Single)
-        ],
-        'tdqs': [
-            (r'"""', String.Double, '#pop'),
-            include('strings-double'),
-            (r'\n', String.Double)
-        ],
-        'tsqs': [
-            (r"'''", String.Single, '#pop'),
-            include('strings-single'),
-            (r'\n', String.Single)
-        ],
-
-        'expr-inside-fstring': [
-            (r'[{([]', Punctuation, 'expr-inside-fstring-inner'),
-            # without format specifier
-            (r'(=\s*)?'         # debug (https://bugs.python.org/issue36817)
-             r'\}', String.Interpol, '#pop'),
-            # with format specifier
-            # we'll catch the remaining '}' in the outer scope
-            (r'(=\s*)?'         # debug (https://bugs.python.org/issue36817)
-             r':', String.Interpol, '#pop'),
-            (r'\s+', Whitespace),  # allow new lines
-            include('expr'),
-        ],
-        'expr-inside-fstring-inner': [
-            (r'[{([]', Punctuation, 'expr-inside-fstring-inner'),
-            (r'[])}]', Punctuation, '#pop'),
-            (r'\s+', Whitespace),  # allow new lines
-            include('expr'),
-        ],
-        'keywords': [
-            (words((
-                'into', 'case', 'type', 'module', 'internal',
-            ), suffix=r'\b'),
-                Keyword),
-            (words(('true', 'false', 'null'), suffix=r'\b'), Keyword.Constant),
-        ],
-        'functions': [
-            (words((
-                "min", "max", "sum", "average", "stddev", "every", "any",
-                "concat_array", "count", "lag", "lead", "first", "last",
-                "rank", "rank_dense", "row_number", "round", "as", "in",
-                "tuple_every", "tuple_map", "tuple_zip", "_eq", "_is_null",
-                "from_text", "lower", "upper", "read_parquet", "read_csv"),
-                suffix=r'\b'),
-             Name.Function),
-        ],
-
-        'comment': [
-            (r'-(?!\})', Comment.Multiline),
-            (r'\{-', Comment.Multiline, 'comment'),
-            (r'[^-}]', Comment.Multiline),
-            (r'-\}', Comment.Multiline, '#pop'),
-        ],
-
-        'imports': [
-            (r'\w+(\.\w+)*', Name.Class, '#pop'),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/ptx.py b/venv/lib/python3.10/site-packages/pygments/lexers/ptx.py
deleted file mode 100644
index 784ca13a6f0c0ba5464706f68c84eeae9d36c9e9..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/ptx.py
+++ /dev/null
@@ -1,119 +0,0 @@
-"""
-    pygments.lexers.ptx
-    ~~~~~~~~~~~~~~~~~~~
-
-    Lexer for other PTX language.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, include, words
-from pygments.token import Comment, Keyword, Name, String, Number, \
-    Punctuation, Whitespace, Operator
-
-__all__ = ["PtxLexer"]
-
-
-class PtxLexer(RegexLexer):
-    """
-    For NVIDIA `PTX `_
-    source.
-    """
-    name = 'PTX'
-    url = "https://docs.nvidia.com/cuda/parallel-thread-execution/"
-    filenames = ['*.ptx']
-    aliases = ['ptx']
-    mimetypes = ['text/x-ptx']
-    version_added = '2.16'
-
-    #: optional Comment or Whitespace
-    string = r'"[^"]*?"'
-    followsym = r'[a-zA-Z0-9_$]'
-    identifier = r'([-a-zA-Z$._][\w\-$.]*|' + string + ')'
-    block_label = r'(' + identifier + r'|(\d+))'
-
-    tokens = {
-        'root': [
-            include('whitespace'),
-
-            (block_label + r'\s*:', Name.Label),
-
-            include('keyword'),
-
-            (r'%' + identifier, Name.Variable),
-            (r'%\d+', Name.Variable.Anonymous),
-            (r'c?' + string, String),
-            (identifier, Name.Variable),
-            (r';', Punctuation),
-            (r'[*+-/]', Operator),
-
-            (r'0[xX][a-fA-F0-9]+', Number),
-            (r'-?\d+(?:[.]\d+)?(?:[eE][-+]?\d+(?:[.]\d+)?)?', Number),
-
-            (r'[=<>{}\[\]()*.,!]|x\b', Punctuation)
-
-        ],
-        'whitespace': [
-            (r'(\n|\s+)+', Whitespace),
-            (r'//.*?\n', Comment)
-        ],
-
-        'keyword': [
-            # Instruction keywords
-            (words((
-                'abs', 'discard', 'min', 'shf', 'vadd',
-                'activemask', 'div', 'mma', 'shfl', 'vadd2',
-                'add', 'dp2a', 'mov', 'shl', 'vadd4',
-                'addc', 'dp4a', 'movmatrix', 'shr', 'vavrg2',
-                'alloca', 'elect', 'mul', 'sin', 'vavrg4',
-                'and', 'ex2', 'mul24', 'slct', 'vmad',
-                'applypriority', 'exit', 'multimem', 'sqrt', 'vmax',
-                'atom', 'fence', 'nanosleep', 'st', 'vmax2',
-                'bar', 'fma', 'neg', 'stackrestore', 'vmax4',
-                'barrier', 'fns', 'not', 'stacksave', 'vmin',
-                'bfe', 'getctarank', 'or', 'stmatrix', 'vmin2',
-                'bfi', 'griddepcontrol', 'pmevent', 'sub', 'vmin4',
-                'bfind', 'isspacep', 'popc', 'subc', 'vote',
-                'bmsk', 'istypep', 'prefetch', 'suld', 'vset',
-                'bra', 'ld', 'prefetchu', 'suq', 'vset2',
-                'brev', 'ldmatrix', 'prmt', 'sured', 'vset4',
-                'brkpt', 'ldu', 'rcp', 'sust', 'vshl',
-                'brx', 'lg2', 'red', 'szext', 'vshr',
-                'call', 'lop3', 'redux', 'tanh', 'vsub',
-                'clz', 'mad', 'rem', 'testp', 'vsub2',
-                'cnot', 'mad24', 'ret', 'tex', 'vsub4',
-                'copysign', 'madc', 'rsqrt', 'tld4', 'wgmma',
-                'cos', 'mapa', 'sad', 'trap', 'wmma',
-                'cp', 'match', 'selp', 'txq', 'xor',
-                'createpolicy', 'max', 'set', 'vabsdiff', 'cvt',
-                'mbarrier', 'setmaxnreg', 'vabsdiff2', 'cvta',
-                'membar', 'setp', 'vabsdiff4')), Keyword),
-            # State Spaces and Suffixes
-            (words((
-                'reg', '.sreg', '.const', '.global',
-                '.local', '.param', '.shared', '.tex',
-                '.wide', '.loc'
-            )), Keyword.Pseudo),
-            # PTX Directives
-            (words((
-                '.address_size', '.explicitcluster', '.maxnreg', '.section',
-                '.alias', '.extern', '.maxntid', '.shared',
-                '.align', '.file', '.minnctapersm', '.sreg',
-                '.branchtargets', '.func', '.noreturn', '.target',
-                '.callprototype', '.global', '.param', '.tex',
-                '.calltargets', '.loc', '.pragma', '.version',
-                '.common', '.local', '.reg', '.visible',
-                '.const', '.maxclusterrank', '.reqnctapercluster', '.weak',
-                '.entry', '.maxnctapersm', '.reqntid')), Keyword.Reserved),
-            # Fundamental Types
-            (words((
-                '.s8', '.s16', '.s32', '.s64',
-                '.u8', '.u16', '.u32', '.u64',
-                '.f16', '.f16x2', '.f32', '.f64',
-                '.b8', '.b16', '.b32', '.b64',
-                '.pred'
-            )), Keyword.Type)
-        ],
-
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/python.py b/venv/lib/python3.10/site-packages/pygments/lexers/python.py
deleted file mode 100644
index 805f6ff2ac778176926cbdb40eca17ef4c93ee70..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/python.py
+++ /dev/null
@@ -1,1201 +0,0 @@
-"""
-    pygments.lexers.python
-    ~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Python and related languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import keyword
-
-from pygments.lexer import DelegatingLexer, RegexLexer, include, \
-    bygroups, using, default, words, combined, this
-from pygments.util import get_bool_opt, shebang_matches
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Generic, Other, Error, Whitespace
-from pygments import unistring as uni
-
-__all__ = ['PythonLexer', 'PythonConsoleLexer', 'PythonTracebackLexer',
-           'Python2Lexer', 'Python2TracebackLexer',
-           'CythonLexer', 'DgLexer', 'NumPyLexer']
-
-
-class PythonLexer(RegexLexer):
-    """
-    For Python source code (version 3.x).
-
-    .. versionchanged:: 2.5
-       This is now the default ``PythonLexer``.  It is still available as the
-       alias ``Python3Lexer``.
-    """
-
-    name = 'Python'
-    url = 'https://www.python.org'
-    aliases = ['python', 'py', 'sage', 'python3', 'py3', 'bazel', 'starlark', 'pyi']
-    filenames = [
-        '*.py',
-        '*.pyw',
-        # Type stubs
-        '*.pyi',
-        # Jython
-        '*.jy',
-        # Sage
-        '*.sage',
-        # SCons
-        '*.sc',
-        'SConstruct',
-        'SConscript',
-        # Skylark/Starlark (used by Bazel, Buck, and Pants)
-        '*.bzl',
-        'BUCK',
-        'BUILD',
-        'BUILD.bazel',
-        'WORKSPACE',
-        # Twisted Application infrastructure
-        '*.tac',
-    ]
-    mimetypes = ['text/x-python', 'application/x-python',
-                 'text/x-python3', 'application/x-python3']
-    version_added = '0.10'
-
-    uni_name = f"[{uni.xid_start}][{uni.xid_continue}]*"
-
-    def innerstring_rules(ttype):
-        return [
-            # the old style '%s' % (...) string formatting (still valid in Py3)
-            (r'%(\(\w+\))?[-#0 +]*([0-9]+|[*])?(\.([0-9]+|[*]))?'
-             '[hlL]?[E-GXc-giorsaux%]', String.Interpol),
-            # the new style '{}'.format(...) string formatting
-            (r'\{'
-             r'((\w+)((\.\w+)|(\[[^\]]+\]))*)?'  # field name
-             r'(\![sra])?'                       # conversion
-             r'(\:(.?[<>=\^])?[-+ ]?#?0?(\d+)?,?(\.\d+)?[E-GXb-gnosx%]?)?'
-             r'\}', String.Interpol),
-
-            # backslashes, quotes and formatting signs must be parsed one at a time
-            (r'[^\\\'"%{\n]+', ttype),
-            (r'[\'"\\]', ttype),
-            # unhandled string formatting sign
-            (r'%|(\{{1,2})', ttype)
-            # newlines are an error (use "nl" state)
-        ]
-
-    def fstring_rules(ttype):
-        return [
-            # Assuming that a '}' is the closing brace after format specifier.
-            # Sadly, this means that we won't detect syntax error. But it's
-            # more important to parse correct syntax correctly, than to
-            # highlight invalid syntax.
-            (r'\}', String.Interpol),
-            (r'\{', String.Interpol, 'expr-inside-fstring'),
-            # backslashes, quotes and formatting signs must be parsed one at a time
-            (r'[^\\\'"{}\n]+', ttype),
-            (r'[\'"\\]', ttype),
-            # newlines are an error (use "nl" state)
-        ]
-
-    tokens = {
-        'root': [
-            (r'\n', Whitespace),
-            (r'^(\s*)([rRuUbB]{,2})("""(?:.|\n)*?""")',
-             bygroups(Whitespace, String.Affix, String.Doc)),
-            (r"^(\s*)([rRuUbB]{,2})('''(?:.|\n)*?''')",
-             bygroups(Whitespace, String.Affix, String.Doc)),
-            (r'\A#!.+$', Comment.Hashbang),
-            (r'#.*$', Comment.Single),
-            (r'\\\n', Text),
-            (r'\\', Text),
-            include('keywords'),
-            include('soft-keywords'),
-            (r'(def)((?:\s|\\\s)+)', bygroups(Keyword, Whitespace), 'funcname'),
-            (r'(class)((?:\s|\\\s)+)', bygroups(Keyword, Whitespace), 'classname'),
-            (r'(from)((?:\s|\\\s)+)', bygroups(Keyword.Namespace, Whitespace),
-             'fromimport'),
-            (r'(import)((?:\s|\\\s)+)', bygroups(Keyword.Namespace, Whitespace),
-             'import'),
-            include('expr'),
-        ],
-        'expr': [
-            # raw f-strings
-            ('(?i)(rf|fr)(""")',
-             bygroups(String.Affix, String.Double),
-             combined('rfstringescape', 'tdqf')),
-            ("(?i)(rf|fr)(''')",
-             bygroups(String.Affix, String.Single),
-             combined('rfstringescape', 'tsqf')),
-            ('(?i)(rf|fr)(")',
-             bygroups(String.Affix, String.Double),
-             combined('rfstringescape', 'dqf')),
-            ("(?i)(rf|fr)(')",
-             bygroups(String.Affix, String.Single),
-             combined('rfstringescape', 'sqf')),
-            # non-raw f-strings
-            ('([fF])(""")', bygroups(String.Affix, String.Double),
-             combined('fstringescape', 'tdqf')),
-            ("([fF])(''')", bygroups(String.Affix, String.Single),
-             combined('fstringescape', 'tsqf')),
-            ('([fF])(")', bygroups(String.Affix, String.Double),
-             combined('fstringescape', 'dqf')),
-            ("([fF])(')", bygroups(String.Affix, String.Single),
-             combined('fstringescape', 'sqf')),
-            # raw bytes and strings
-            ('(?i)(rb|br|r)(""")',
-             bygroups(String.Affix, String.Double), 'tdqs'),
-            ("(?i)(rb|br|r)(''')",
-             bygroups(String.Affix, String.Single), 'tsqs'),
-            ('(?i)(rb|br|r)(")',
-             bygroups(String.Affix, String.Double), 'dqs'),
-            ("(?i)(rb|br|r)(')",
-             bygroups(String.Affix, String.Single), 'sqs'),
-            # non-raw strings
-            ('([uU]?)(""")', bygroups(String.Affix, String.Double),
-             combined('stringescape', 'tdqs')),
-            ("([uU]?)(''')", bygroups(String.Affix, String.Single),
-             combined('stringescape', 'tsqs')),
-            ('([uU]?)(")', bygroups(String.Affix, String.Double),
-             combined('stringescape', 'dqs')),
-            ("([uU]?)(')", bygroups(String.Affix, String.Single),
-             combined('stringescape', 'sqs')),
-            # non-raw bytes
-            ('([bB])(""")', bygroups(String.Affix, String.Double),
-             combined('bytesescape', 'tdqs')),
-            ("([bB])(''')", bygroups(String.Affix, String.Single),
-             combined('bytesescape', 'tsqs')),
-            ('([bB])(")', bygroups(String.Affix, String.Double),
-             combined('bytesescape', 'dqs')),
-            ("([bB])(')", bygroups(String.Affix, String.Single),
-             combined('bytesescape', 'sqs')),
-
-            (r'[^\S\n]+', Text),
-            include('numbers'),
-            (r'!=|==|<<|>>|:=|[-~+/*%=<>&^|.]', Operator),
-            (r'[]{}:(),;[]', Punctuation),
-            (r'(in|is|and|or|not)\b', Operator.Word),
-            include('expr-keywords'),
-            include('builtins'),
-            include('magicfuncs'),
-            include('magicvars'),
-            include('name'),
-        ],
-        'expr-inside-fstring': [
-            (r'[{([]', Punctuation, 'expr-inside-fstring-inner'),
-            # without format specifier
-            (r'(=\s*)?'         # debug (https://bugs.python.org/issue36817)
-             r'(\![sraf])?'     # conversion
-             r'\}', String.Interpol, '#pop'),
-            # with format specifier
-            # we'll catch the remaining '}' in the outer scope
-            (r'(=\s*)?'         # debug (https://bugs.python.org/issue36817)
-             r'(\![sraf])?'     # conversion
-             r':', String.Interpol, '#pop'),
-            (r'\s+', Whitespace),  # allow new lines
-            include('expr'),
-        ],
-        'expr-inside-fstring-inner': [
-            (r'[{([]', Punctuation, 'expr-inside-fstring-inner'),
-            (r'[])}]', Punctuation, '#pop'),
-            (r'\s+', Whitespace),  # allow new lines
-            include('expr'),
-        ],
-        'expr-keywords': [
-            # Based on https://docs.python.org/3/reference/expressions.html
-            (words((
-                'async for', 'await', 'else', 'for', 'if', 'lambda',
-                'yield', 'yield from'), suffix=r'\b'),
-             Keyword),
-            (words(('True', 'False', 'None'), suffix=r'\b'), Keyword.Constant),
-        ],
-        'keywords': [
-            (words((
-                'assert', 'async', 'await', 'break', 'continue', 'del', 'elif',
-                'else', 'except', 'finally', 'for', 'global', 'if', 'lambda',
-                'pass', 'raise', 'nonlocal', 'return', 'try', 'while', 'yield',
-                'yield from', 'as', 'with'), suffix=r'\b'),
-             Keyword),
-            (words(('True', 'False', 'None'), suffix=r'\b'), Keyword.Constant),
-        ],
-        'soft-keywords': [
-            # `match`, `case` and `_` soft keywords
-            (r'(^[ \t]*)'              # at beginning of line + possible indentation
-             r'(match|case)\b'         # a possible keyword
-             r'(?![ \t]*(?:'           # not followed by...
-             r'[:,;=^&|@~)\]}]|(?:' +  # characters and keywords that mean this isn't
-                                       # pattern matching (but None/True/False is ok)
-             r'|'.join(k for k in keyword.kwlist if k[0].islower()) + r')\b))',
-             bygroups(Text, Keyword), 'soft-keywords-inner'),
-        ],
-        'soft-keywords-inner': [
-            # optional `_` keyword
-            (r'(\s+)([^\n_]*)(_\b)', bygroups(Whitespace, using(this), Keyword)),
-            default('#pop')
-        ],
-        'builtins': [
-            (words((
-                '__import__', 'abs', 'aiter', 'all', 'any', 'bin', 'bool', 'bytearray',
-                'breakpoint', 'bytes', 'callable', 'chr', 'classmethod', 'compile',
-                'complex', 'delattr', 'dict', 'dir', 'divmod', 'enumerate', 'eval',
-                'filter', 'float', 'format', 'frozenset', 'getattr', 'globals',
-                'hasattr', 'hash', 'hex', 'id', 'input', 'int', 'isinstance',
-                'issubclass', 'iter', 'len', 'list', 'locals', 'map', 'max',
-                'memoryview', 'min', 'next', 'object', 'oct', 'open', 'ord', 'pow',
-                'print', 'property', 'range', 'repr', 'reversed', 'round', 'set',
-                'setattr', 'slice', 'sorted', 'staticmethod', 'str', 'sum', 'super',
-                'tuple', 'type', 'vars', 'zip'), prefix=r'(?>|[-~+/*%=<>&^|.]', Operator),
-            include('keywords'),
-            (r'(def)((?:\s|\\\s)+)', bygroups(Keyword, Whitespace), 'funcname'),
-            (r'(class)((?:\s|\\\s)+)', bygroups(Keyword, Whitespace), 'classname'),
-            (r'(from)((?:\s|\\\s)+)', bygroups(Keyword.Namespace, Whitespace),
-             'fromimport'),
-            (r'(import)((?:\s|\\\s)+)', bygroups(Keyword.Namespace, Whitespace),
-             'import'),
-            include('builtins'),
-            include('magicfuncs'),
-            include('magicvars'),
-            include('backtick'),
-            ('([rR]|[uUbB][rR]|[rR][uUbB])(""")',
-             bygroups(String.Affix, String.Double), 'tdqs'),
-            ("([rR]|[uUbB][rR]|[rR][uUbB])(''')",
-             bygroups(String.Affix, String.Single), 'tsqs'),
-            ('([rR]|[uUbB][rR]|[rR][uUbB])(")',
-             bygroups(String.Affix, String.Double), 'dqs'),
-            ("([rR]|[uUbB][rR]|[rR][uUbB])(')",
-             bygroups(String.Affix, String.Single), 'sqs'),
-            ('([uUbB]?)(""")', bygroups(String.Affix, String.Double),
-             combined('stringescape', 'tdqs')),
-            ("([uUbB]?)(''')", bygroups(String.Affix, String.Single),
-             combined('stringescape', 'tsqs')),
-            ('([uUbB]?)(")', bygroups(String.Affix, String.Double),
-             combined('stringescape', 'dqs')),
-            ("([uUbB]?)(')", bygroups(String.Affix, String.Single),
-             combined('stringescape', 'sqs')),
-            include('name'),
-            include('numbers'),
-        ],
-        'keywords': [
-            (words((
-                'assert', 'break', 'continue', 'del', 'elif', 'else', 'except',
-                'exec', 'finally', 'for', 'global', 'if', 'lambda', 'pass',
-                'print', 'raise', 'return', 'try', 'while', 'yield',
-                'yield from', 'as', 'with'), suffix=r'\b'),
-             Keyword),
-        ],
-        'builtins': [
-            (words((
-                '__import__', 'abs', 'all', 'any', 'apply', 'basestring', 'bin',
-                'bool', 'buffer', 'bytearray', 'bytes', 'callable', 'chr', 'classmethod',
-                'cmp', 'coerce', 'compile', 'complex', 'delattr', 'dict', 'dir', 'divmod',
-                'enumerate', 'eval', 'execfile', 'exit', 'file', 'filter', 'float',
-                'frozenset', 'getattr', 'globals', 'hasattr', 'hash', 'hex', 'id',
-                'input', 'int', 'intern', 'isinstance', 'issubclass', 'iter', 'len',
-                'list', 'locals', 'long', 'map', 'max', 'min', 'next', 'object',
-                'oct', 'open', 'ord', 'pow', 'property', 'range', 'raw_input', 'reduce',
-                'reload', 'repr', 'reversed', 'round', 'set', 'setattr', 'slice',
-                'sorted', 'staticmethod', 'str', 'sum', 'super', 'tuple', 'type',
-                'unichr', 'unicode', 'vars', 'xrange', 'zip'),
-                prefix=r'(?>> )(.*\n)', bygroups(Generic.Prompt, Other.Code), 'continuations'),
-            # This happens, e.g., when tracebacks are embedded in documentation;
-            # trailing whitespaces are often stripped in such contexts.
-            (r'(>>>)(\n)', bygroups(Generic.Prompt, Whitespace)),
-            (r'(\^C)?Traceback \(most recent call last\):\n', Other.Traceback, 'traceback'),
-            # SyntaxError starts with this
-            (r'  File "[^"]+", line \d+', Other.Traceback, 'traceback'),
-            (r'.*\n', Generic.Output),
-        ],
-        'continuations': [
-            (r'(\.\.\. )(.*\n)', bygroups(Generic.Prompt, Other.Code)),
-            # See above.
-            (r'(\.\.\.)(\n)', bygroups(Generic.Prompt, Whitespace)),
-            default('#pop'),
-        ],
-        'traceback': [
-            # As soon as we see a traceback, consume everything until the next
-            # >>> prompt.
-            (r'(?=>>>( |$))', Text, '#pop'),
-            (r'(KeyboardInterrupt)(\n)', bygroups(Name.Class, Whitespace)),
-            (r'.*\n', Other.Traceback),
-        ],
-    }
-
-
-class PythonConsoleLexer(DelegatingLexer):
-    """
-    For Python console output or doctests, such as:
-
-    .. sourcecode:: pycon
-
-        >>> a = 'foo'
-        >>> print(a)
-        foo
-        >>> 1 / 0
-        Traceback (most recent call last):
-          File "", line 1, in 
-        ZeroDivisionError: integer division or modulo by zero
-
-    Additional options:
-
-    `python3`
-        Use Python 3 lexer for code.  Default is ``True``.
-
-        .. versionadded:: 1.0
-        .. versionchanged:: 2.5
-           Now defaults to ``True``.
-    """
-
-    name = 'Python console session'
-    aliases = ['pycon', 'python-console']
-    mimetypes = ['text/x-python-doctest']
-    url = 'https://python.org'
-    version_added = ''
-
-    def __init__(self, **options):
-        python3 = get_bool_opt(options, 'python3', True)
-        if python3:
-            pylexer = PythonLexer
-            tblexer = PythonTracebackLexer
-        else:
-            pylexer = Python2Lexer
-            tblexer = Python2TracebackLexer
-        # We have two auxiliary lexers. Use DelegatingLexer twice with
-        # different tokens.  TODO: DelegatingLexer should support this
-        # directly, by accepting a tuplet of auxiliary lexers and a tuple of
-        # distinguishing tokens. Then we wouldn't need this intermediary
-        # class.
-        class _ReplaceInnerCode(DelegatingLexer):
-            def __init__(self, **options):
-                super().__init__(pylexer, _PythonConsoleLexerBase, Other.Code, **options)
-        super().__init__(tblexer, _ReplaceInnerCode, Other.Traceback, **options)
-
-
-class PythonTracebackLexer(RegexLexer):
-    """
-    For Python 3.x tracebacks, with support for chained exceptions.
-
-    .. versionchanged:: 2.5
-       This is now the default ``PythonTracebackLexer``.  It is still available
-       as the alias ``Python3TracebackLexer``.
-    """
-
-    name = 'Python Traceback'
-    aliases = ['pytb', 'py3tb']
-    filenames = ['*.pytb', '*.py3tb']
-    mimetypes = ['text/x-python-traceback', 'text/x-python3-traceback']
-    url = 'https://python.org'
-    version_added = '1.0'
-
-    tokens = {
-        'root': [
-            (r'\n', Whitespace),
-            (r'^(\^C)?Traceback \(most recent call last\):\n', Generic.Traceback, 'intb'),
-            (r'^During handling of the above exception, another '
-             r'exception occurred:\n\n', Generic.Traceback),
-            (r'^The above exception was the direct cause of the '
-             r'following exception:\n\n', Generic.Traceback),
-            (r'^(?=  File "[^"]+", line \d+)', Generic.Traceback, 'intb'),
-            (r'^.*\n', Other),
-        ],
-        'intb': [
-            (r'^(  File )("[^"]+")(, line )(\d+)(, in )(.+)(\n)',
-             bygroups(Text, Name.Builtin, Text, Number, Text, Name, Whitespace)),
-            (r'^(  File )("[^"]+")(, line )(\d+)(\n)',
-             bygroups(Text, Name.Builtin, Text, Number, Whitespace)),
-            (r'^(    )(.+)(\n)',
-             bygroups(Whitespace, using(PythonLexer), Whitespace), 'markers'),
-            (r'^([ \t]*)(\.\.\.)(\n)',
-             bygroups(Whitespace, Comment, Whitespace)),  # for doctests...
-            (r'^([^:]+)(: )(.+)(\n)',
-             bygroups(Generic.Error, Text, Name, Whitespace), '#pop'),
-            (r'^([a-zA-Z_][\w.]*)(:?\n)',
-             bygroups(Generic.Error, Whitespace), '#pop'),
-            default('#pop'),
-        ],
-        'markers': [
-            # Either `PEP 657 `
-            # error locations in Python 3.11+, or single-caret markers
-            # for syntax errors before that.
-            (r'^( {4,})([~^]+)(\n)',
-             bygroups(Whitespace, Punctuation.Marker, Whitespace),
-             '#pop'),
-            default('#pop'),
-        ],
-    }
-
-
-Python3TracebackLexer = PythonTracebackLexer
-
-
-class Python2TracebackLexer(RegexLexer):
-    """
-    For Python tracebacks.
-
-    .. versionchanged:: 2.5
-       This class has been renamed from ``PythonTracebackLexer``.
-       ``PythonTracebackLexer`` now refers to the Python 3 variant.
-    """
-
-    name = 'Python 2.x Traceback'
-    aliases = ['py2tb']
-    filenames = ['*.py2tb']
-    mimetypes = ['text/x-python2-traceback']
-    url = 'https://python.org'
-    version_added = '0.7'
-
-    tokens = {
-        'root': [
-            # Cover both (most recent call last) and (innermost last)
-            # The optional ^C allows us to catch keyboard interrupt signals.
-            (r'^(\^C)?(Traceback.*\n)',
-             bygroups(Text, Generic.Traceback), 'intb'),
-            # SyntaxError starts with this.
-            (r'^(?=  File "[^"]+", line \d+)', Generic.Traceback, 'intb'),
-            (r'^.*\n', Other),
-        ],
-        'intb': [
-            (r'^(  File )("[^"]+")(, line )(\d+)(, in )(.+)(\n)',
-             bygroups(Text, Name.Builtin, Text, Number, Text, Name, Whitespace)),
-            (r'^(  File )("[^"]+")(, line )(\d+)(\n)',
-             bygroups(Text, Name.Builtin, Text, Number, Whitespace)),
-            (r'^(    )(.+)(\n)',
-             bygroups(Text, using(Python2Lexer), Whitespace), 'marker'),
-            (r'^([ \t]*)(\.\.\.)(\n)',
-             bygroups(Text, Comment, Whitespace)),  # for doctests...
-            (r'^([^:]+)(: )(.+)(\n)',
-             bygroups(Generic.Error, Text, Name, Whitespace), '#pop'),
-            (r'^([a-zA-Z_]\w*)(:?\n)',
-             bygroups(Generic.Error, Whitespace), '#pop')
-        ],
-        'marker': [
-            # For syntax errors.
-            (r'( {4,})(\^)', bygroups(Text, Punctuation.Marker), '#pop'),
-            default('#pop'),
-        ],
-    }
-
-
-class CythonLexer(RegexLexer):
-    """
-    For Pyrex and Cython source code.
-    """
-
-    name = 'Cython'
-    url = 'https://cython.org'
-    aliases = ['cython', 'pyx', 'pyrex']
-    filenames = ['*.pyx', '*.pxd', '*.pxi']
-    mimetypes = ['text/x-cython', 'application/x-cython']
-    version_added = '1.1'
-
-    tokens = {
-        'root': [
-            (r'\n', Whitespace),
-            (r'^(\s*)("""(?:.|\n)*?""")', bygroups(Whitespace, String.Doc)),
-            (r"^(\s*)('''(?:.|\n)*?''')", bygroups(Whitespace, String.Doc)),
-            (r'[^\S\n]+', Text),
-            (r'#.*$', Comment),
-            (r'[]{}:(),;[]', Punctuation),
-            (r'\\\n', Whitespace),
-            (r'\\', Text),
-            (r'(in|is|and|or|not)\b', Operator.Word),
-            (r'(<)([a-zA-Z0-9.?]+)(>)',
-             bygroups(Punctuation, Keyword.Type, Punctuation)),
-            (r'!=|==|<<|>>|[-~+/*%=<>&^|.?]', Operator),
-            (r'(from)(\d+)(<=)(\s+)(<)(\d+)(:)',
-             bygroups(Keyword, Number.Integer, Operator, Whitespace, Operator,
-                      Name, Punctuation)),
-            include('keywords'),
-            (r'(def|property)(\s+)', bygroups(Keyword, Whitespace), 'funcname'),
-            (r'(cp?def)(\s+)', bygroups(Keyword, Whitespace), 'cdef'),
-            # (should actually start a block with only cdefs)
-            (r'(cdef)(:)', bygroups(Keyword, Punctuation)),
-            (r'(class|struct)(\s+)', bygroups(Keyword, Whitespace), 'classname'),
-            (r'(from)(\s+)', bygroups(Keyword, Whitespace), 'fromimport'),
-            (r'(c?import)(\s+)', bygroups(Keyword, Whitespace), 'import'),
-            include('builtins'),
-            include('backtick'),
-            ('(?:[rR]|[uU][rR]|[rR][uU])"""', String, 'tdqs'),
-            ("(?:[rR]|[uU][rR]|[rR][uU])'''", String, 'tsqs'),
-            ('(?:[rR]|[uU][rR]|[rR][uU])"', String, 'dqs'),
-            ("(?:[rR]|[uU][rR]|[rR][uU])'", String, 'sqs'),
-            ('[uU]?"""', String, combined('stringescape', 'tdqs')),
-            ("[uU]?'''", String, combined('stringescape', 'tsqs')),
-            ('[uU]?"', String, combined('stringescape', 'dqs')),
-            ("[uU]?'", String, combined('stringescape', 'sqs')),
-            include('name'),
-            include('numbers'),
-        ],
-        'keywords': [
-            (words((
-                'assert', 'async', 'await', 'break', 'by', 'continue', 'ctypedef', 'del', 'elif',
-                'else', 'except', 'except?', 'exec', 'finally', 'for', 'fused', 'gil',
-                'global', 'if', 'include', 'lambda', 'nogil', 'pass', 'print',
-                'raise', 'return', 'try', 'while', 'yield', 'as', 'with'), suffix=r'\b'),
-             Keyword),
-            (r'(DEF|IF|ELIF|ELSE)\b', Comment.Preproc),
-        ],
-        'builtins': [
-            (words((
-                '__import__', 'abs', 'all', 'any', 'apply', 'basestring', 'bin', 'bint',
-                'bool', 'buffer', 'bytearray', 'bytes', 'callable', 'chr',
-                'classmethod', 'cmp', 'coerce', 'compile', 'complex', 'delattr',
-                'dict', 'dir', 'divmod', 'enumerate', 'eval', 'execfile', 'exit',
-                'file', 'filter', 'float', 'frozenset', 'getattr', 'globals',
-                'hasattr', 'hash', 'hex', 'id', 'input', 'int', 'intern', 'isinstance',
-                'issubclass', 'iter', 'len', 'list', 'locals', 'long', 'map', 'max',
-                'min', 'next', 'object', 'oct', 'open', 'ord', 'pow', 'property', 'Py_ssize_t',
-                'range', 'raw_input', 'reduce', 'reload', 'repr', 'reversed',
-                'round', 'set', 'setattr', 'slice', 'sorted', 'staticmethod',
-                'str', 'sum', 'super', 'tuple', 'type', 'unichr', 'unicode', 'unsigned',
-                'vars', 'xrange', 'zip'), prefix=r'(??/\\:']?:)(\s*)(\{)",
-             bygroups(Name.Function, Whitespace, Operator, Whitespace, Punctuation),
-             "functions"),
-            # Variable Names
-            (r"([.]?[a-zA-Z][\w.]*)(\s*)([-.~=!@#$%^&*_+|,<>?/\\:']?:)",
-             bygroups(Name.Variable, Whitespace, Operator)),
-            # Functions
-            (r"\{", Punctuation, "functions"),
-            # Parentheses
-            (r"\(", Punctuation, "parentheses"),
-            # Brackets
-            (r"\[", Punctuation, "brackets"),
-            # Errors
-            (r"'`([a-zA-Z][\w.]*)?", Name.Exception),
-            # File Symbols
-            (r"`:([a-zA-Z/][\w./]*)?", String.Symbol),
-            # Symbols
-            (r"`([a-zA-Z][\w.]*)?", String.Symbol),
-            # Numbers
-            include("numbers"),
-            # Variable Names
-            (r"[a-zA-Z][\w.]*", Name),
-            # Operators
-            (r"[-=+*#$%@!~^&:.,<>'\\|/?_]", Operator),
-            # Punctuation
-            (r";", Punctuation),
-        ],
-        "functions": [
-            include("root"),
-            (r"\}", Punctuation, "#pop"),
-        ],
-        "parentheses": [
-            include("root"),
-            (r"\)", Punctuation, "#pop"),
-        ],
-        "brackets": [
-            include("root"),
-            (r"\]", Punctuation, "#pop"),
-        ],
-        "numbers": [
-            # Binary Values
-            (r"[01]+b", Number.Bin),
-            # Nulls/Infinities
-            (r"0[nNwW][cefghijmndzuvtp]?", Number),
-            # Timestamps
-            ((r"(?:[0-9]{4}[.][0-9]{2}[.][0-9]{2}|[0-9]+)"
-              "D(?:[0-9](?:[0-9](?::[0-9]{2}"
-              "(?::[0-9]{2}(?:[.][0-9]*)?)?)?)?)?"), Literal.Date),
-            # Datetimes
-            ((r"[0-9]{4}[.][0-9]{2}"
-              "(?:m|[.][0-9]{2}(?:T(?:[0-9]{2}:[0-9]{2}"
-              "(?::[0-9]{2}(?:[.][0-9]*)?)?)?)?)"), Literal.Date),
-            # Times
-            (r"[0-9]{2}:[0-9]{2}(?::[0-9]{2}(?:[.][0-9]{1,3})?)?",
-             Literal.Date),
-            # GUIDs
-            (r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}",
-             Number.Hex),
-            # Byte Vectors
-            (r"0x[0-9a-fA-F]+", Number.Hex),
-            # Floats
-            (r"([0-9]*[.]?[0-9]+|[0-9]+[.]?[0-9]*)[eE][+-]?[0-9]+[ef]?",
-             Number.Float),
-            (r"([0-9]*[.][0-9]+|[0-9]+[.][0-9]*)[ef]?", Number.Float),
-            (r"[0-9]+[ef]", Number.Float),
-            # Characters
-            (r"[0-9]+c", Number),
-            # Integers
-            (r"[0-9]+[ihtuv]", Number.Integer),
-            # Long Integers
-            (r"[0-9]+[jnp]?", Number.Integer.Long),
-        ],
-        "comments": [
-            (r"[^\\]+", Comment.Multiline),
-            (r"^\\", Comment.Multiline, "#pop"),
-            (r"\\", Comment.Multiline),
-        ],
-        "strings": [
-            (r'[^"\\]+', String.Double),
-            (r"\\.", String.Escape),
-            (r'"', String.Double, "#pop"),
-        ],
-    }
-
-
-class QLexer(KLexer):
-    """
-    For `Q `_ source code.
-    """
-
-    name = "Q"
-    aliases = ["q"]
-    filenames = ["*.q"]
-    version_added = '2.12'
-
-    tokens = {
-        "root": [
-            (words(("aj", "aj0", "ajf", "ajf0", "all", "and", "any", "asc",
-                    "asof", "attr", "avgs", "ceiling", "cols", "count", "cross",
-                    "csv", "cut", "deltas", "desc", "differ", "distinct", "dsave",
-                    "each", "ej", "ema", "eval", "except", "fby", "fills", "first",
-                    "fkeys", "flip", "floor", "get", "group", "gtime", "hclose",
-                    "hcount", "hdel", "hsym", "iasc", "idesc", "ij", "ijf",
-                    "inter", "inv", "key", "keys", "lj", "ljf", "load", "lower",
-                    "lsq", "ltime", "ltrim", "mavg", "maxs", "mcount", "md5",
-                    "mdev", "med", "meta", "mins", "mmax", "mmin", "mmu", "mod",
-                    "msum", "neg", "next", "not", "null", "or", "over", "parse",
-                    "peach", "pj", "prds", "prior", "prev", "rand", "rank", "ratios",
-                    "raze", "read0", "read1", "reciprocal", "reval", "reverse",
-                    "rload", "rotate", "rsave", "rtrim", "save", "scan", "scov",
-                    "sdev", "set", "show", "signum", "ssr", "string", "sublist",
-                    "sums", "sv", "svar", "system", "tables", "til", "trim", "txf",
-                    "type", "uj", "ujf", "ungroup", "union", "upper", "upsert",
-                    "value", "view", "views", "vs", "where", "wj", "wj1", "ww",
-                    "xasc", "xbar", "xcol", "xcols", "xdesc", "xgroup", "xkey",
-                    "xlog", "xprev", "xrank"),
-                    suffix=r"\b"), Name.Builtin,
-            ),
-            inherit,
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/qlik.py b/venv/lib/python3.10/site-packages/pygments/lexers/qlik.py
deleted file mode 100644
index a29f89f35a9b1f4e8be514da1259e42feb2a9642..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/qlik.py
+++ /dev/null
@@ -1,117 +0,0 @@
-"""
-    pygments.lexers.qlik
-    ~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for the qlik scripting language
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, include, bygroups, words
-from pygments.token import Comment, Keyword, Name, Number, Operator, \
-    Punctuation, String, Text
-from pygments.lexers._qlik_builtins import OPERATORS_LIST, STATEMENT_LIST, \
-    SCRIPT_FUNCTIONS, CONSTANT_LIST
-
-__all__ = ["QlikLexer"]
-
-
-class QlikLexer(RegexLexer):
-    """
-    Lexer for qlik code, including .qvs files
-    """
-
-    name = "Qlik"
-    aliases = ["qlik", "qlikview", "qliksense", "qlikscript"]
-    filenames = ["*.qvs", "*.qvw"]
-    url = "https://qlik.com"
-    version_added = '2.12'
-
-    flags = re.IGNORECASE
-
-    tokens = {
-        # Handle multi-line comments
-        "comment": [
-            (r"\*/", Comment.Multiline, "#pop"),
-            (r"[^*]+", Comment.Multiline),
-        ],
-        # Handle numbers
-        "numerics": [
-            (r"\b\d+\.\d+(e\d+)?[fd]?\b", Number.Float),
-            (r"\b\d+\b", Number.Integer),
-        ],
-        # Handle variable names in things
-        "interp": [
-            (
-                r"(\$\()(\w+)(\))",
-                bygroups(String.Interpol, Name.Variable, String.Interpol),
-            ),
-        ],
-        # Handle strings
-        "string": [
-            (r"'", String, "#pop"),
-            include("interp"),
-            (r"[^'$]+", String),
-            (r"\$", String),
-        ],
-        #
-        "assignment": [
-            (r";", Punctuation, "#pop"),
-            include("root"),
-        ],
-        "field_name_quote": [
-            (r'"', String.Symbol, "#pop"),
-            include("interp"),
-            (r"[^\"$]+", String.Symbol),
-            (r"\$", String.Symbol),
-        ],
-        "field_name_bracket": [
-            (r"\]", String.Symbol, "#pop"),
-            include("interp"),
-            (r"[^\]$]+", String.Symbol),
-            (r"\$", String.Symbol),
-        ],
-        "function": [(r"\)", Punctuation, "#pop"), include("root")],
-        "root": [
-            # Whitespace and comments
-            (r"\s+", Text.Whitespace),
-            (r"/\*", Comment.Multiline, "comment"),
-            (r"//.*\n", Comment.Single),
-            # variable assignment
-            (r"(let|set)(\s+)", bygroups(Keyword.Declaration, Text.Whitespace),
-             "assignment"),
-            # Word operators
-            (words(OPERATORS_LIST["words"], prefix=r"\b", suffix=r"\b"),
-             Operator.Word),
-            # Statements
-            (words(STATEMENT_LIST, suffix=r"\b"), Keyword),
-            # Table names
-            (r"[a-z]\w*:", Keyword.Declaration),
-            # Constants
-            (words(CONSTANT_LIST, suffix=r"\b"), Keyword.Constant),
-            # Functions
-            (words(SCRIPT_FUNCTIONS, suffix=r"(?=\s*\()"), Name.Builtin,
-             "function"),
-            # interpolation - e.g. $(variableName)
-            include("interp"),
-            # Quotes denote a field/file name
-            (r'"', String.Symbol, "field_name_quote"),
-            # Square brackets denote a field/file name
-            (r"\[", String.Symbol, "field_name_bracket"),
-            # Strings
-            (r"'", String, "string"),
-            # Numbers
-            include("numerics"),
-            # Operator symbols
-            (words(OPERATORS_LIST["symbols"]), Operator),
-            # Strings denoted by single quotes
-            (r"'.+?'", String),
-            # Words as text
-            (r"\b\w+\b", Text),
-            # Basic punctuation
-            (r"[,;.()\\/]", Punctuation),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/qvt.py b/venv/lib/python3.10/site-packages/pygments/lexers/qvt.py
deleted file mode 100644
index 302d1b6ed894ef6a928fd09ba7b679791d4ddaa7..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/qvt.py
+++ /dev/null
@@ -1,153 +0,0 @@
-"""
-    pygments.lexers.qvt
-    ~~~~~~~~~~~~~~~~~~~
-
-    Lexer for QVT Operational language.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, bygroups, include, combined, default, \
-    words
-from pygments.token import Text, Comment, Operator, Keyword, Punctuation, \
-    Name, String, Number
-
-__all__ = ['QVToLexer']
-
-
-class QVToLexer(RegexLexer):
-    """
-    For the QVT Operational Mapping language.
-
-    Reference for implementing this: «Meta Object Facility (MOF) 2.0
-    Query/View/Transformation Specification», Version 1.1 - January 2011
-    (https://www.omg.org/spec/QVT/1.1/), see §8.4, «Concrete Syntax» in
-    particular.
-
-    Notable tokens assignments:
-
-    - Name.Class is assigned to the identifier following any of the following
-      keywords: metamodel, class, exception, primitive, enum, transformation
-      or library
-
-    - Name.Function is assigned to the names of mappings and queries
-
-    - Name.Builtin.Pseudo is assigned to the pre-defined variables 'this',
-      'self' and 'result'.
-    """
-    # With obvious borrowings & inspiration from the Java, Python and C lexers
-
-    name = 'QVTO'
-    aliases = ['qvto', 'qvt']
-    filenames = ['*.qvto']
-    url = 'https://www.omg.org/spec/QVT/1.1'
-    version_added = ''
-
-    tokens = {
-        'root': [
-            (r'\n', Text),
-            (r'[^\S\n]+', Text),
-            (r'(--|//)(\s*)(directive:)?(.*)$',
-             bygroups(Comment, Comment, Comment.Preproc, Comment)),
-            # Uncomment the following if you want to distinguish between
-            # '/*' and '/**', à la javadoc
-            # (r'/[*]{2}(.|\n)*?[*]/', Comment.Multiline),
-            (r'/[*](.|\n)*?[*]/', Comment.Multiline),
-            (r'\\\n', Text),
-            (r'(and|not|or|xor|##?)\b', Operator.Word),
-            (r'(:{1,2}=|[-+]=)\b', Operator.Word),
-            (r'(@|<<|>>)\b', Keyword),  # stereotypes
-            (r'!=|<>|==|=|!->|->|>=|<=|[.]{3}|[+/*%=<>&|.~]', Operator),
-            (r'[]{}:(),;[]', Punctuation),
-            (r'(true|false|unlimited|null)\b', Keyword.Constant),
-            (r'(this|self|result)\b', Name.Builtin.Pseudo),
-            (r'(var)\b', Keyword.Declaration),
-            (r'(from|import)\b', Keyword.Namespace, 'fromimport'),
-            (r'(metamodel|class|exception|primitive|enum|transformation|'
-             r'library)(\s+)(\w+)',
-             bygroups(Keyword.Word, Text, Name.Class)),
-            (r'(exception)(\s+)(\w+)',
-             bygroups(Keyword.Word, Text, Name.Exception)),
-            (r'(main)\b', Name.Function),
-            (r'(mapping|helper|query)(\s+)',
-             bygroups(Keyword.Declaration, Text), 'operation'),
-            (r'(assert)(\s+)\b', bygroups(Keyword, Text), 'assert'),
-            (r'(Bag|Collection|Dict|OrderedSet|Sequence|Set|Tuple|List)\b',
-             Keyword.Type),
-            include('keywords'),
-            ('"', String, combined('stringescape', 'dqs')),
-            ("'", String, combined('stringescape', 'sqs')),
-            include('name'),
-            include('numbers'),
-            # (r'([a-zA-Z_]\w*)(::)([a-zA-Z_]\w*)',
-            # bygroups(Text, Text, Text)),
-        ],
-
-        'fromimport': [
-            (r'(?:[ \t]|\\\n)+', Text),
-            (r'[a-zA-Z_][\w.]*', Name.Namespace),
-            default('#pop'),
-        ],
-
-        'operation': [
-            (r'::', Text),
-            (r'(.*::)([a-zA-Z_]\w*)([ \t]*)(\()',
-             bygroups(Text, Name.Function, Text, Punctuation), '#pop')
-        ],
-
-        'assert': [
-            (r'(warning|error|fatal)\b', Keyword, '#pop'),
-            default('#pop'),  # all else: go back
-        ],
-
-        'keywords': [
-            (words((
-                'abstract', 'access', 'any', 'assert', 'blackbox', 'break',
-                'case', 'collect', 'collectNested', 'collectOne', 'collectselect',
-                'collectselectOne', 'composes', 'compute', 'configuration',
-                'constructor', 'continue', 'datatype', 'default', 'derived',
-                'disjuncts', 'do', 'elif', 'else', 'end', 'endif', 'except',
-                'exists', 'extends', 'forAll', 'forEach', 'forOne', 'from', 'if',
-                'implies', 'in', 'inherits', 'init', 'inout', 'intermediate',
-                'invresolve', 'invresolveIn', 'invresolveone', 'invresolveoneIn',
-                'isUnique', 'iterate', 'late', 'let', 'literal', 'log', 'map',
-                'merges', 'modeltype', 'new', 'object', 'one', 'ordered', 'out',
-                'package', 'population', 'property', 'raise', 'readonly',
-                'references', 'refines', 'reject', 'resolve', 'resolveIn',
-                'resolveone', 'resolveoneIn', 'return', 'select', 'selectOne',
-                'sortedBy', 'static', 'switch', 'tag', 'then', 'try', 'typedef',
-                'unlimited', 'uses', 'when', 'where', 'while', 'with', 'xcollect',
-                'xmap', 'xselect'), suffix=r'\b'), Keyword),
-        ],
-
-        # There is no need to distinguish between String.Single and
-        # String.Double: 'strings' is factorised for 'dqs' and 'sqs'
-        'strings': [
-            (r'[^\\\'"\n]+', String),
-            # quotes, percents and backslashes must be parsed one at a time
-            (r'[\'"\\]', String),
-        ],
-        'stringescape': [
-            (r'\\([\\btnfr"\']|u[0-3][0-7]{2}|u[0-7]{1,2})', String.Escape)
-        ],
-        'dqs': [  # double-quoted string
-            (r'"', String, '#pop'),
-            (r'\\\\|\\"', String.Escape),
-            include('strings')
-        ],
-        'sqs': [  # single-quoted string
-            (r"'", String, '#pop'),
-            (r"\\\\|\\'", String.Escape),
-            include('strings')
-        ],
-        'name': [
-            (r'[a-zA-Z_]\w*', Name),
-        ],
-        # numbers: excerpt taken from the python lexer
-        'numbers': [
-            (r'(\d+\.\d*|\d*\.\d+)([eE][+-]?[0-9]+)?', Number.Float),
-            (r'\d+[eE][+-]?[0-9]+', Number.Float),
-            (r'\d+', Number.Integer)
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/r.py b/venv/lib/python3.10/site-packages/pygments/lexers/r.py
deleted file mode 100644
index d3f65ba2d8644a6faa4f50526d2e31b08ed4d339..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/r.py
+++ /dev/null
@@ -1,196 +0,0 @@
-"""
-    pygments.lexers.r
-    ~~~~~~~~~~~~~~~~~
-
-    Lexers for the R/S languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import Lexer, RegexLexer, include, do_insertions
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Generic, Whitespace
-
-__all__ = ['RConsoleLexer', 'SLexer', 'RdLexer']
-
-
-line_re  = re.compile('.*?\n')
-
-
-class RConsoleLexer(Lexer):
-    """
-    For R console transcripts or R CMD BATCH output files.
-    """
-
-    name = 'RConsole'
-    aliases = ['rconsole', 'rout']
-    filenames = ['*.Rout']
-    url = 'https://www.r-project.org'
-    version_added = ''
-    _example = "rconsole/r-console-transcript.Rout"
-
-    def get_tokens_unprocessed(self, text):
-        slexer = SLexer(**self.options)
-
-        current_code_block = ''
-        insertions = []
-
-        for match in line_re.finditer(text):
-            line = match.group()
-            if line.startswith('>') or line.startswith('+'):
-                # Colorize the prompt as such,
-                # then put rest of line into current_code_block
-                insertions.append((len(current_code_block),
-                                   [(0, Generic.Prompt, line[:2])]))
-                current_code_block += line[2:]
-            else:
-                # We have reached a non-prompt line!
-                # If we have stored prompt lines, need to process them first.
-                if current_code_block:
-                    # Weave together the prompts and highlight code.
-                    yield from do_insertions(
-                        insertions, slexer.get_tokens_unprocessed(current_code_block))
-                    # Reset vars for next code block.
-                    current_code_block = ''
-                    insertions = []
-                # Now process the actual line itself, this is output from R.
-                yield match.start(), Generic.Output, line
-
-        # If we happen to end on a code block with nothing after it, need to
-        # process the last code block. This is neither elegant nor DRY so
-        # should be changed.
-        if current_code_block:
-            yield from do_insertions(
-                insertions, slexer.get_tokens_unprocessed(current_code_block))
-
-
-class SLexer(RegexLexer):
-    """
-    For S, S-plus, and R source code.
-    """
-
-    name = 'S'
-    aliases = ['splus', 's', 'r']
-    filenames = ['*.S', '*.R', '.Rhistory', '.Rprofile', '.Renviron']
-    mimetypes = ['text/S-plus', 'text/S', 'text/x-r-source', 'text/x-r',
-                 'text/x-R', 'text/x-r-history', 'text/x-r-profile']
-    url = 'https://www.r-project.org'
-    version_added = '0.10'
-
-    valid_name = r'`[^`\\]*(?:\\.[^`\\]*)*`|(?:[a-zA-Z]|\.[A-Za-z_.])[\w.]*|\.'
-    tokens = {
-        'comments': [
-            (r'#.*$', Comment.Single),
-        ],
-        'valid_name': [
-            (valid_name, Name),
-        ],
-        'function_name': [
-            (rf'({valid_name})\s*(?=\()', Name.Function),
-        ],
-        'punctuation': [
-            (r'\[{1,2}|\]{1,2}|\(|\)|;|,', Punctuation),
-        ],
-        'keywords': [
-            (r'(if|else|for|while|repeat|in|next|break|return|switch|function)'
-             r'(?![\w.])',
-             Keyword.Reserved),
-        ],
-        'operators': [
-            (r'<>?|-|==|<=|>=|\|>|<|>|&&?|!=|\|\|?|\?', Operator),
-            (r'\*|\+|\^|/|!|%[^%]*%|=|~|\$|@|:{1,3}', Operator),
-        ],
-        'builtin_symbols': [
-            (r'(NULL|NA(_(integer|real|complex|character)_)?|'
-             r'letters|LETTERS|Inf|TRUE|FALSE|NaN|pi|\.\.(\.|[0-9]+))'
-             r'(?![\w.])',
-             Keyword.Constant),
-            (r'(T|F)\b', Name.Builtin.Pseudo),
-        ],
-        'numbers': [
-            # hex number
-            (r'0[xX][a-fA-F0-9]+([pP][0-9]+)?[Li]?', Number.Hex),
-            # decimal number
-            (r'[+-]?([0-9]+(\.[0-9]+)?|\.[0-9]+|\.)([eE][+-]?[0-9]+)?[Li]?',
-             Number),
-        ],
-        'statements': [
-            include('comments'),
-            # whitespaces
-            (r'\s+', Whitespace),
-            (r'\'', String, 'string_squote'),
-            (r'\"', String, 'string_dquote'),
-            include('builtin_symbols'),
-            include('keywords'),
-            include('function_name'),
-            include('valid_name'),
-            include('numbers'),
-            include('punctuation'),
-            include('operators'),
-        ],
-        'root': [
-            # calls:
-            include('statements'),
-            # blocks:
-            (r'\{|\}', Punctuation),
-            # (r'\{', Punctuation, 'block'),
-            (r'.', Text),
-        ],
-        # 'block': [
-        #    include('statements'),
-        #    ('\{', Punctuation, '#push'),
-        #    ('\}', Punctuation, '#pop')
-        # ],
-        'string_squote': [
-            (r'([^\'\\]|\\.)*\'', String, '#pop'),
-        ],
-        'string_dquote': [
-            (r'([^"\\]|\\.)*"', String, '#pop'),
-        ],
-    }
-
-    def analyse_text(text):
-        if re.search(r'[a-z0-9_\])\s]<-(?!-)', text):
-            return 0.11
-
-
-class RdLexer(RegexLexer):
-    """
-    Pygments Lexer for R documentation (Rd) files
-
-    This is a very minimal implementation, highlighting little more
-    than the macros. A description of Rd syntax is found in `Writing R
-    Extensions `_
-    and `Parsing Rd files `_.
-    """
-    name = 'Rd'
-    aliases = ['rd']
-    filenames = ['*.Rd']
-    mimetypes = ['text/x-r-doc']
-    url = 'http://cran.r-project.org/doc/manuals/R-exts.html'
-    version_added = '1.6'
-
-    # To account for verbatim / LaTeX-like / and R-like areas
-    # would require parsing.
-    tokens = {
-        'root': [
-            # catch escaped brackets and percent sign
-            (r'\\[\\{}%]', String.Escape),
-            # comments
-            (r'%.*$', Comment),
-            # special macros with no arguments
-            (r'\\(?:cr|l?dots|R|tab)\b', Keyword.Constant),
-            # macros
-            (r'\\[a-zA-Z]+\b', Keyword),
-            # special preprocessor macros
-            (r'^\s*#(?:ifn?def|endif).*\b', Comment.Preproc),
-            # non-escaped brackets
-            (r'[{}]', Name.Builtin),
-            # everything else
-            (r'[^\\%\n{}]+', Text),
-            (r'.', Text),
-        ]
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/rdf.py b/venv/lib/python3.10/site-packages/pygments/lexers/rdf.py
deleted file mode 100644
index 4930c1b387a6f098074fc15e26a1f7a62f3554c1..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/rdf.py
+++ /dev/null
@@ -1,468 +0,0 @@
-"""
-    pygments.lexers.rdf
-    ~~~~~~~~~~~~~~~~~~~
-
-    Lexers for semantic web and RDF query languages and markup.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, bygroups, default
-from pygments.token import Keyword, Punctuation, String, Number, Operator, \
-    Generic, Whitespace, Name, Literal, Comment, Text
-
-__all__ = ['SparqlLexer', 'TurtleLexer', 'ShExCLexer']
-
-
-class SparqlLexer(RegexLexer):
-    """
-    Lexer for SPARQL query language.
-    """
-    name = 'SPARQL'
-    aliases = ['sparql']
-    filenames = ['*.rq', '*.sparql']
-    mimetypes = ['application/sparql-query']
-    url = 'https://www.w3.org/TR/sparql11-query'
-    version_added = '2.0'
-
-    # character group definitions ::
-
-    PN_CHARS_BASE_GRP = ('a-zA-Z'
-                         '\u00c0-\u00d6'
-                         '\u00d8-\u00f6'
-                         '\u00f8-\u02ff'
-                         '\u0370-\u037d'
-                         '\u037f-\u1fff'
-                         '\u200c-\u200d'
-                         '\u2070-\u218f'
-                         '\u2c00-\u2fef'
-                         '\u3001-\ud7ff'
-                         '\uf900-\ufdcf'
-                         '\ufdf0-\ufffd')
-
-    PN_CHARS_U_GRP = (PN_CHARS_BASE_GRP + '_')
-
-    PN_CHARS_GRP = (PN_CHARS_U_GRP +
-                    r'\-' +
-                    r'0-9' +
-                    '\u00b7' +
-                    '\u0300-\u036f' +
-                    '\u203f-\u2040')
-
-    HEX_GRP = '0-9A-Fa-f'
-
-    PN_LOCAL_ESC_CHARS_GRP = r' _~.\-!$&"()*+,;=/?#@%'
-
-    # terminal productions ::
-
-    PN_CHARS_BASE = '[' + PN_CHARS_BASE_GRP + ']'
-
-    PN_CHARS_U = '[' + PN_CHARS_U_GRP + ']'
-
-    PN_CHARS = '[' + PN_CHARS_GRP + ']'
-
-    HEX = '[' + HEX_GRP + ']'
-
-    PN_LOCAL_ESC_CHARS = '[' + PN_LOCAL_ESC_CHARS_GRP + ']'
-
-    IRIREF = r'<(?:[^<>"{}|^`\\\x00-\x20])*>'
-
-    BLANK_NODE_LABEL = '_:[0-9' + PN_CHARS_U_GRP + '](?:[' + PN_CHARS_GRP + \
-                       '.]*' + PN_CHARS + ')?'
-
-    PN_PREFIX = PN_CHARS_BASE + '(?:[' + PN_CHARS_GRP + '.]*' + PN_CHARS + ')?'
-
-    VARNAME = '[0-9' + PN_CHARS_U_GRP + '][' + PN_CHARS_U_GRP + \
-              '0-9\u00b7\u0300-\u036f\u203f-\u2040]*'
-
-    PERCENT = '%' + HEX + HEX
-
-    PN_LOCAL_ESC = r'\\' + PN_LOCAL_ESC_CHARS
-
-    PLX = '(?:' + PERCENT + ')|(?:' + PN_LOCAL_ESC + ')'
-
-    PN_LOCAL = ('(?:[' + PN_CHARS_U_GRP + ':0-9' + ']|' + PLX + ')' +
-                '(?:(?:[' + PN_CHARS_GRP + '.:]|' + PLX + ')*(?:[' +
-                PN_CHARS_GRP + ':]|' + PLX + '))?')
-
-    EXPONENT = r'[eE][+-]?\d+'
-
-    # Lexer token definitions ::
-
-    tokens = {
-        'root': [
-            (r'\s+', Text),
-            # keywords ::
-            (r'(?i)(select|construct|describe|ask|where|filter|group\s+by|minus|'
-             r'distinct|reduced|from\s+named|from|order\s+by|desc|asc|limit|'
-             r'offset|values|bindings|load|into|clear|drop|create|add|move|copy|'
-             r'insert\s+data|delete\s+data|delete\s+where|with|delete|insert|'
-             r'using\s+named|using|graph|default|named|all|optional|service|'
-             r'silent|bind|undef|union|not\s+in|in|as|having|to|prefix|base)\b', Keyword),
-            (r'(a)\b', Keyword),
-            # IRIs ::
-            ('(' + IRIREF + ')', Name.Label),
-            # blank nodes ::
-            ('(' + BLANK_NODE_LABEL + ')', Name.Label),
-            #  # variables ::
-            ('[?$]' + VARNAME, Name.Variable),
-            # prefixed names ::
-            (r'(' + PN_PREFIX + r')?(\:)(' + PN_LOCAL + r')?',
-             bygroups(Name.Namespace, Punctuation, Name.Tag)),
-            # function names ::
-            (r'(?i)(str|lang|langmatches|datatype|bound|iri|uri|bnode|rand|abs|'
-             r'ceil|floor|round|concat|strlen|ucase|lcase|encode_for_uri|'
-             r'contains|strstarts|strends|strbefore|strafter|year|month|day|'
-             r'hours|minutes|seconds|timezone|tz|now|uuid|struuid|md5|sha1|sha256|sha384|'
-             r'sha512|coalesce|if|strlang|strdt|sameterm|isiri|isuri|isblank|'
-             r'isliteral|isnumeric|regex|substr|replace|exists|not\s+exists|'
-             r'count|sum|min|max|avg|sample|group_concat|separator)\b',
-             Name.Function),
-            # boolean literals ::
-            (r'(true|false)', Keyword.Constant),
-            # double literals ::
-            (r'[+\-]?(\d+\.\d*' + EXPONENT + r'|\.?\d+' + EXPONENT + ')', Number.Float),
-            # decimal literals ::
-            (r'[+\-]?(\d+\.\d*|\.\d+)', Number.Float),
-            # integer literals ::
-            (r'[+\-]?\d+', Number.Integer),
-            # operators ::
-            (r'(\|\||&&|=|\*|\-|\+|/|!=|<=|>=|!|<|>)', Operator),
-            # punctuation characters ::
-            (r'[(){}.;,:^\[\]]', Punctuation),
-            # line comments ::
-            (r'#[^\n]*', Comment),
-            # strings ::
-            (r'"""', String, 'triple-double-quoted-string'),
-            (r'"', String, 'single-double-quoted-string'),
-            (r"'''", String, 'triple-single-quoted-string'),
-            (r"'", String, 'single-single-quoted-string'),
-        ],
-        'triple-double-quoted-string': [
-            (r'"""', String, 'end-of-string'),
-            (r'[^\\]+', String),
-            (r'\\', String, 'string-escape'),
-        ],
-        'single-double-quoted-string': [
-            (r'"', String, 'end-of-string'),
-            (r'[^"\\\n]+', String),
-            (r'\\', String, 'string-escape'),
-        ],
-        'triple-single-quoted-string': [
-            (r"'''", String, 'end-of-string'),
-            (r'[^\\]+', String),
-            (r'\\', String.Escape, 'string-escape'),
-        ],
-        'single-single-quoted-string': [
-            (r"'", String, 'end-of-string'),
-            (r"[^'\\\n]+", String),
-            (r'\\', String, 'string-escape'),
-        ],
-        'string-escape': [
-            (r'u' + HEX + '{4}', String.Escape, '#pop'),
-            (r'U' + HEX + '{8}', String.Escape, '#pop'),
-            (r'.', String.Escape, '#pop'),
-        ],
-        'end-of-string': [
-            (r'(@)([a-zA-Z]+(?:-[a-zA-Z0-9]+)*)',
-             bygroups(Operator, Name.Function), '#pop:2'),
-            (r'\^\^', Operator, '#pop:2'),
-            default('#pop:2'),
-        ],
-    }
-
-
-class TurtleLexer(RegexLexer):
-    """
-    Lexer for Turtle data language.
-    """
-    name = 'Turtle'
-    aliases = ['turtle']
-    filenames = ['*.ttl']
-    mimetypes = ['text/turtle', 'application/x-turtle']
-    url = 'https://www.w3.org/TR/turtle'
-    version_added = '2.1'
-
-    # character group definitions ::
-    PN_CHARS_BASE_GRP = ('a-zA-Z'
-                         '\u00c0-\u00d6'
-                         '\u00d8-\u00f6'
-                         '\u00f8-\u02ff'
-                         '\u0370-\u037d'
-                         '\u037f-\u1fff'
-                         '\u200c-\u200d'
-                         '\u2070-\u218f'
-                         '\u2c00-\u2fef'
-                         '\u3001-\ud7ff'
-                         '\uf900-\ufdcf'
-                         '\ufdf0-\ufffd')
-
-    PN_CHARS_U_GRP = (PN_CHARS_BASE_GRP + '_')
-
-    PN_CHARS_GRP = (PN_CHARS_U_GRP +
-                    r'\-' +
-                    r'0-9' +
-                    '\u00b7' +
-                    '\u0300-\u036f' +
-                    '\u203f-\u2040')
-
-    PN_CHARS = '[' + PN_CHARS_GRP + ']'
-
-    PN_CHARS_BASE = '[' + PN_CHARS_BASE_GRP + ']'
-
-    PN_PREFIX = PN_CHARS_BASE + '(?:[' + PN_CHARS_GRP + '.]*' + PN_CHARS + ')?'
-
-    HEX_GRP = '0-9A-Fa-f'
-
-    HEX = '[' + HEX_GRP + ']'
-
-    PERCENT = '%' + HEX + HEX
-
-    PN_LOCAL_ESC_CHARS_GRP = r' _~.\-!$&"()*+,;=/?#@%'
-
-    PN_LOCAL_ESC_CHARS = '[' + PN_LOCAL_ESC_CHARS_GRP + ']'
-
-    PN_LOCAL_ESC = r'\\' + PN_LOCAL_ESC_CHARS
-
-    PLX = '(?:' + PERCENT + ')|(?:' + PN_LOCAL_ESC + ')'
-
-    PN_LOCAL = ('(?:[' + PN_CHARS_U_GRP + ':0-9' + ']|' + PLX + ')' +
-                '(?:(?:[' + PN_CHARS_GRP + '.:]|' + PLX + ')*(?:[' +
-                PN_CHARS_GRP + ':]|' + PLX + '))?')
-
-    patterns = {
-        'PNAME_NS': r'((?:[a-zA-Z][\w-]*)?\:)',  # Simplified character range
-        'IRIREF': r'(<[^<>"{}|^`\\\x00-\x20]*>)'
-    }
-
-    tokens = {
-        'root': [
-            (r'\s+', Text),
-
-            # Base / prefix
-            (r'(@base|BASE)(\s+){IRIREF}(\s*)(\.?)'.format(**patterns),
-             bygroups(Keyword, Whitespace, Name.Variable, Whitespace,
-                      Punctuation)),
-            (r'(@prefix|PREFIX)(\s+){PNAME_NS}(\s+){IRIREF}(\s*)(\.?)'.format(**patterns),
-             bygroups(Keyword, Whitespace, Name.Namespace, Whitespace,
-                      Name.Variable, Whitespace, Punctuation)),
-
-            # The shorthand predicate 'a'
-            (r'(?<=\s)a(?=\s)', Keyword.Type),
-
-            # IRIREF
-            (r'{IRIREF}'.format(**patterns), Name.Variable),
-
-            # PrefixedName
-            (r'(' + PN_PREFIX + r')?(\:)(' + PN_LOCAL + r')?',
-             bygroups(Name.Namespace, Punctuation, Name.Tag)),
-
-            # BlankNodeLabel
-            (r'(_)(:)([' + PN_CHARS_U_GRP + r'0-9]([' + PN_CHARS_GRP + r'.]*' + PN_CHARS + ')?)',
-             bygroups(Name.Namespace, Punctuation, Name.Tag)),
-
-            # Comment
-            (r'#[^\n]+', Comment),
-
-            (r'\b(true|false)\b', Literal),
-            (r'[+\-]?\d*\.\d+', Number.Float),
-            (r'[+\-]?\d*(:?\.\d+)?E[+\-]?\d+', Number.Float),
-            (r'[+\-]?\d+', Number.Integer),
-            (r'[\[\](){}.;,:^]', Punctuation),
-
-            (r'"""', String, 'triple-double-quoted-string'),
-            (r'"', String, 'single-double-quoted-string'),
-            (r"'''", String, 'triple-single-quoted-string'),
-            (r"'", String, 'single-single-quoted-string'),
-        ],
-        'triple-double-quoted-string': [
-            (r'"""', String, 'end-of-string'),
-            (r'[^\\]+(?=""")', String),
-            (r'\\', String, 'string-escape'),
-        ],
-        'single-double-quoted-string': [
-            (r'"', String, 'end-of-string'),
-            (r'[^"\\\n]+', String),
-            (r'\\', String, 'string-escape'),
-        ],
-        'triple-single-quoted-string': [
-            (r"'''", String, 'end-of-string'),
-            (r"[^\\]+(?=''')", String),
-            (r'\\', String, 'string-escape'),
-        ],
-        'single-single-quoted-string': [
-            (r"'", String, 'end-of-string'),
-            (r"[^'\\\n]+", String),
-            (r'\\', String, 'string-escape'),
-        ],
-        'string-escape': [
-            (r'.', String, '#pop'),
-        ],
-        'end-of-string': [
-            (r'(@)([a-zA-Z]+(?:-[a-zA-Z0-9]+)*)',
-             bygroups(Operator, Generic.Emph), '#pop:2'),
-
-            (r'(\^\^){IRIREF}'.format(**patterns), bygroups(Operator, Generic.Emph), '#pop:2'),
-
-            default('#pop:2'),
-
-        ],
-    }
-
-    # Turtle and Tera Term macro files share the same file extension
-    # but each has a recognizable and distinct syntax.
-    def analyse_text(text):
-        for t in ('@base ', 'BASE ', '@prefix ', 'PREFIX '):
-            if re.search(rf'^\s*{t}', text):
-                return 0.80
-
-
-class ShExCLexer(RegexLexer):
-    """
-    Lexer for ShExC shape expressions language syntax.
-    """
-    name = 'ShExC'
-    aliases = ['shexc', 'shex']
-    filenames = ['*.shex']
-    mimetypes = ['text/shex']
-    url = 'https://shex.io/shex-semantics/#shexc'
-    version_added = ''
-
-    # character group definitions ::
-
-    PN_CHARS_BASE_GRP = ('a-zA-Z'
-                         '\u00c0-\u00d6'
-                         '\u00d8-\u00f6'
-                         '\u00f8-\u02ff'
-                         '\u0370-\u037d'
-                         '\u037f-\u1fff'
-                         '\u200c-\u200d'
-                         '\u2070-\u218f'
-                         '\u2c00-\u2fef'
-                         '\u3001-\ud7ff'
-                         '\uf900-\ufdcf'
-                         '\ufdf0-\ufffd')
-
-    PN_CHARS_U_GRP = (PN_CHARS_BASE_GRP + '_')
-
-    PN_CHARS_GRP = (PN_CHARS_U_GRP +
-                    r'\-' +
-                    r'0-9' +
-                    '\u00b7' +
-                    '\u0300-\u036f' +
-                    '\u203f-\u2040')
-
-    HEX_GRP = '0-9A-Fa-f'
-
-    PN_LOCAL_ESC_CHARS_GRP = r"_~.\-!$&'()*+,;=/?#@%"
-
-    # terminal productions ::
-
-    PN_CHARS_BASE = '[' + PN_CHARS_BASE_GRP + ']'
-
-    PN_CHARS_U = '[' + PN_CHARS_U_GRP + ']'
-
-    PN_CHARS = '[' + PN_CHARS_GRP + ']'
-
-    HEX = '[' + HEX_GRP + ']'
-
-    PN_LOCAL_ESC_CHARS = '[' + PN_LOCAL_ESC_CHARS_GRP + ']'
-
-    UCHAR_NO_BACKSLASH = '(?:u' + HEX + '{4}|U' + HEX + '{8})'
-
-    UCHAR = r'\\' + UCHAR_NO_BACKSLASH
-
-    IRIREF = r'<(?:[^\x00-\x20<>"{}|^`\\]|' + UCHAR + ')*>'
-
-    BLANK_NODE_LABEL = '_:[0-9' + PN_CHARS_U_GRP + '](?:[' + PN_CHARS_GRP + \
-                       '.]*' + PN_CHARS + ')?'
-
-    PN_PREFIX = PN_CHARS_BASE + '(?:[' + PN_CHARS_GRP + '.]*' + PN_CHARS + ')?'
-
-    PERCENT = '%' + HEX + HEX
-
-    PN_LOCAL_ESC = r'\\' + PN_LOCAL_ESC_CHARS
-
-    PLX = '(?:' + PERCENT + ')|(?:' + PN_LOCAL_ESC + ')'
-
-    PN_LOCAL = ('(?:[' + PN_CHARS_U_GRP + ':0-9' + ']|' + PLX + ')' +
-                '(?:(?:[' + PN_CHARS_GRP + '.:]|' + PLX + ')*(?:[' +
-                PN_CHARS_GRP + ':]|' + PLX + '))?')
-
-    EXPONENT = r'[eE][+-]?\d+'
-
-    # Lexer token definitions ::
-
-    tokens = {
-        'root': [
-            (r'\s+', Text),
-            # keywords ::
-            (r'(?i)(base|prefix|start|external|'
-             r'literal|iri|bnode|nonliteral|length|minlength|maxlength|'
-             r'mininclusive|minexclusive|maxinclusive|maxexclusive|'
-             r'totaldigits|fractiondigits|'
-             r'closed|extra)\b', Keyword),
-            (r'(a)\b', Keyword),
-            # IRIs ::
-            ('(' + IRIREF + ')', Name.Label),
-            # blank nodes ::
-            ('(' + BLANK_NODE_LABEL + ')', Name.Label),
-            # prefixed names ::
-            (r'(' + PN_PREFIX + r')?(\:)(' + PN_LOCAL + ')?',
-             bygroups(Name.Namespace, Punctuation, Name.Tag)),
-            # boolean literals ::
-            (r'(true|false)', Keyword.Constant),
-            # double literals ::
-            (r'[+\-]?(\d+\.\d*' + EXPONENT + r'|\.?\d+' + EXPONENT + ')', Number.Float),
-            # decimal literals ::
-            (r'[+\-]?(\d+\.\d*|\.\d+)', Number.Float),
-            # integer literals ::
-            (r'[+\-]?\d+', Number.Integer),
-            # operators ::
-            (r'[@|$&=*+?^\-~]', Operator),
-            # operator keywords ::
-            (r'(?i)(and|or|not)\b', Operator.Word),
-            # punctuation characters ::
-            (r'[(){}.;,:^\[\]]', Punctuation),
-            # line comments ::
-            (r'#[^\n]*', Comment),
-            # strings ::
-            (r'"""', String, 'triple-double-quoted-string'),
-            (r'"', String, 'single-double-quoted-string'),
-            (r"'''", String, 'triple-single-quoted-string'),
-            (r"'", String, 'single-single-quoted-string'),
-        ],
-        'triple-double-quoted-string': [
-            (r'"""', String, 'end-of-string'),
-            (r'[^\\]+', String),
-            (r'\\', String, 'string-escape'),
-        ],
-        'single-double-quoted-string': [
-            (r'"', String, 'end-of-string'),
-            (r'[^"\\\n]+', String),
-            (r'\\', String, 'string-escape'),
-        ],
-        'triple-single-quoted-string': [
-            (r"'''", String, 'end-of-string'),
-            (r'[^\\]+', String),
-            (r'\\', String.Escape, 'string-escape'),
-        ],
-        'single-single-quoted-string': [
-            (r"'", String, 'end-of-string'),
-            (r"[^'\\\n]+", String),
-            (r'\\', String, 'string-escape'),
-        ],
-        'string-escape': [
-            (UCHAR_NO_BACKSLASH, String.Escape, '#pop'),
-            (r'.', String.Escape, '#pop'),
-        ],
-        'end-of-string': [
-            (r'(@)([a-zA-Z]+(?:-[a-zA-Z0-9]+)*)',
-             bygroups(Operator, Name.Function), '#pop:2'),
-            (r'\^\^', Operator, '#pop:2'),
-            default('#pop:2'),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/rebol.py b/venv/lib/python3.10/site-packages/pygments/lexers/rebol.py
deleted file mode 100644
index 4b37a749450472b4adf9c6a613bc1b8ab503f5a2..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/rebol.py
+++ /dev/null
@@ -1,419 +0,0 @@
-"""
-    pygments.lexers.rebol
-    ~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for the REBOL and related languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, bygroups
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Generic, Whitespace
-
-__all__ = ['RebolLexer', 'RedLexer']
-
-
-class RebolLexer(RegexLexer):
-    """
-    A REBOL lexer.
-    """
-    name = 'REBOL'
-    aliases = ['rebol']
-    filenames = ['*.r', '*.r3', '*.reb']
-    mimetypes = ['text/x-rebol']
-    url = 'http://www.rebol.com'
-    version_added = '1.1'
-
-    flags = re.IGNORECASE | re.MULTILINE
-
-    escape_re = r'(?:\^\([0-9a-f]{1,4}\)*)'
-
-    def word_callback(lexer, match):
-        word = match.group()
-
-        if re.match(".*:$", word):
-            yield match.start(), Generic.Subheading, word
-        elif re.match(
-            r'(native|alias|all|any|as-string|as-binary|bind|bound\?|case|'
-            r'catch|checksum|comment|debase|dehex|exclude|difference|disarm|'
-            r'either|else|enbase|foreach|remove-each|form|free|get|get-env|if|'
-            r'in|intersect|loop|minimum-of|maximum-of|mold|new-line|'
-            r'new-line\?|not|now|prin|print|reduce|compose|construct|repeat|'
-            r'reverse|save|script\?|set|shift|switch|throw|to-hex|trace|try|'
-            r'type\?|union|unique|unless|unprotect|unset|until|use|value\?|'
-            r'while|compress|decompress|secure|open|close|read|read-io|'
-            r'write-io|write|update|query|wait|input\?|exp|log-10|log-2|'
-            r'log-e|square-root|cosine|sine|tangent|arccosine|arcsine|'
-            r'arctangent|protect|lowercase|uppercase|entab|detab|connected\?|'
-            r'browse|launch|stats|get-modes|set-modes|to-local-file|'
-            r'to-rebol-file|encloak|decloak|create-link|do-browser|bind\?|'
-            r'hide|draw|show|size-text|textinfo|offset-to-caret|'
-            r'caret-to-offset|local-request-file|rgb-to-hsv|hsv-to-rgb|'
-            r'crypt-strength\?|dh-make-key|dh-generate-key|dh-compute-key|'
-            r'dsa-make-key|dsa-generate-key|dsa-make-signature|'
-            r'dsa-verify-signature|rsa-make-key|rsa-generate-key|'
-            r'rsa-encrypt)$', word):
-            yield match.start(), Name.Builtin, word
-        elif re.match(
-            r'(add|subtract|multiply|divide|remainder|power|and~|or~|xor~|'
-            r'minimum|maximum|negate|complement|absolute|random|head|tail|'
-            r'next|back|skip|at|pick|first|second|third|fourth|fifth|sixth|'
-            r'seventh|eighth|ninth|tenth|last|path|find|select|make|to|copy\*|'
-            r'insert|remove|change|poke|clear|trim|sort|min|max|abs|cp|'
-            r'copy)$', word):
-            yield match.start(), Name.Function, word
-        elif re.match(
-            r'(error|source|input|license|help|install|echo|Usage|with|func|'
-            r'throw-on-error|function|does|has|context|probe|\?\?|as-pair|'
-            r'mod|modulo|round|repend|about|set-net|append|join|rejoin|reform|'
-            r'remold|charset|array|replace|move|extract|forskip|forall|alter|'
-            r'first+|also|take|for|forever|dispatch|attempt|what-dir|'
-            r'change-dir|clean-path|list-dir|dirize|rename|split-path|delete|'
-            r'make-dir|delete-dir|in-dir|confirm|dump-obj|upgrade|what|'
-            r'build-tag|process-source|build-markup|decode-cgi|read-cgi|'
-            r'write-user|save-user|set-user-name|protect-system|parse-xml|'
-            r'cvs-date|cvs-version|do-boot|get-net-info|desktop|layout|'
-            r'scroll-para|get-face|alert|set-face|uninstall|unfocus|'
-            r'request-dir|center-face|do-events|net-error|decode-url|'
-            r'parse-header|parse-header-date|parse-email-addrs|import-email|'
-            r'send|build-attach-body|resend|show-popup|hide-popup|open-events|'
-            r'find-key-face|do-face|viewtop|confine|find-window|'
-            r'insert-event-func|remove-event-func|inform|dump-pane|dump-face|'
-            r'flag-face|deflag-face|clear-fields|read-net|vbug|path-thru|'
-            r'read-thru|load-thru|do-thru|launch-thru|load-image|'
-            r'request-download|do-face-alt|set-font|set-para|get-style|'
-            r'set-style|make-face|stylize|choose|hilight-text|hilight-all|'
-            r'unlight-text|focus|scroll-drag|clear-face|reset-face|scroll-face|'
-            r'resize-face|load-stock|load-stock-block|notify|request|flash|'
-            r'request-color|request-pass|request-text|request-list|'
-            r'request-date|request-file|dbug|editor|link-relative-path|'
-            r'emailer|parse-error)$', word):
-            yield match.start(), Keyword.Namespace, word
-        elif re.match(
-            r'(halt|quit|do|load|q|recycle|call|run|ask|parse|view|unview|'
-            r'return|exit|break)$', word):
-            yield match.start(), Name.Exception, word
-        elif re.match('REBOL$', word):
-            yield match.start(), Generic.Heading, word
-        elif re.match("to-.*", word):
-            yield match.start(), Keyword, word
-        elif re.match(r'(\+|-|\*|/|//|\*\*|and|or|xor|=\?|=|==|<>|<|>|<=|>=)$',
-                      word):
-            yield match.start(), Operator, word
-        elif re.match(r".*\?$", word):
-            yield match.start(), Keyword, word
-        elif re.match(r".*\!$", word):
-            yield match.start(), Keyword.Type, word
-        elif re.match("'.*", word):
-            yield match.start(), Name.Variable.Instance, word  # lit-word
-        elif re.match("#.*", word):
-            yield match.start(), Name.Label, word  # issue
-        elif re.match("%.*", word):
-            yield match.start(), Name.Decorator, word  # file
-        else:
-            yield match.start(), Name.Variable, word
-
-    tokens = {
-        'root': [
-            (r'\s+', Text),
-            (r'#"', String.Char, 'char'),
-            (r'#\{[0-9a-f]*\}', Number.Hex),
-            (r'2#\{', Number.Hex, 'bin2'),
-            (r'64#\{[0-9a-z+/=\s]*\}', Number.Hex),
-            (r'"', String, 'string'),
-            (r'\{', String, 'string2'),
-            (r';#+.*\n', Comment.Special),
-            (r';\*+.*\n', Comment.Preproc),
-            (r';.*\n', Comment),
-            (r'%"', Name.Decorator, 'stringFile'),
-            (r'%[^(^{")\s\[\]]+', Name.Decorator),
-            (r'[+-]?([a-z]{1,3})?\$\d+(\.\d+)?', Number.Float),  # money
-            (r'[+-]?\d+\:\d+(\:\d+)?(\.\d+)?', String.Other),    # time
-            (r'\d+[\-/][0-9a-z]+[\-/]\d+(\/\d+\:\d+((\:\d+)?'
-             r'([.\d+]?([+-]?\d+:\d+)?)?)?)?', String.Other),   # date
-            (r'\d+(\.\d+)+\.\d+', Keyword.Constant),             # tuple
-            (r'\d+X\d+', Keyword.Constant),                   # pair
-            (r'[+-]?\d+(\'\d+)?([.,]\d*)?E[+-]?\d+', Number.Float),
-            (r'[+-]?\d+(\'\d+)?[.,]\d*', Number.Float),
-            (r'[+-]?\d+(\'\d+)?', Number),
-            (r'[\[\]()]', Generic.Strong),
-            (r'[a-z]+[^(^{"\s:)]*://[^(^{"\s)]*', Name.Decorator),  # url
-            (r'mailto:[^(^{"@\s)]+@[^(^{"@\s)]+', Name.Decorator),  # url
-            (r'[^(^{"@\s)]+@[^(^{"@\s)]+', Name.Decorator),         # email
-            (r'comment\s"', Comment, 'commentString1'),
-            (r'comment\s\{', Comment, 'commentString2'),
-            (r'comment\s\[', Comment, 'commentBlock'),
-            (r'comment\s[^(\s{"\[]+', Comment),
-            (r'/[^(^{")\s/[\]]*', Name.Attribute),
-            (r'([^(^{")\s/[\]]+)(?=[:({"\s/\[\]])', word_callback),
-            (r'<[\w:.-]*>', Name.Tag),
-            (r'<[^(<>\s")]+', Name.Tag, 'tag'),
-            (r'([^(^{")\s]+)', Text),
-        ],
-        'string': [
-            (r'[^(^")]+', String),
-            (escape_re, String.Escape),
-            (r'[(|)]+', String),
-            (r'\^.', String.Escape),
-            (r'"', String, '#pop'),
-        ],
-        'string2': [
-            (r'[^(^{})]+', String),
-            (escape_re, String.Escape),
-            (r'[(|)]+', String),
-            (r'\^.', String.Escape),
-            (r'\{', String, '#push'),
-            (r'\}', String, '#pop'),
-        ],
-        'stringFile': [
-            (r'[^(^")]+', Name.Decorator),
-            (escape_re, Name.Decorator),
-            (r'\^.', Name.Decorator),
-            (r'"', Name.Decorator, '#pop'),
-        ],
-        'char': [
-            (escape_re + '"', String.Char, '#pop'),
-            (r'\^."', String.Char, '#pop'),
-            (r'."', String.Char, '#pop'),
-        ],
-        'tag': [
-            (escape_re, Name.Tag),
-            (r'"', Name.Tag, 'tagString'),
-            (r'[^(<>\r\n")]+', Name.Tag),
-            (r'>', Name.Tag, '#pop'),
-        ],
-        'tagString': [
-            (r'[^(^")]+', Name.Tag),
-            (escape_re, Name.Tag),
-            (r'[(|)]+', Name.Tag),
-            (r'\^.', Name.Tag),
-            (r'"', Name.Tag, '#pop'),
-        ],
-        'tuple': [
-            (r'(\d+\.)+', Keyword.Constant),
-            (r'\d+', Keyword.Constant, '#pop'),
-        ],
-        'bin2': [
-            (r'\s+', Number.Hex),
-            (r'([01]\s*){8}', Number.Hex),
-            (r'\}', Number.Hex, '#pop'),
-        ],
-        'commentString1': [
-            (r'[^(^")]+', Comment),
-            (escape_re, Comment),
-            (r'[(|)]+', Comment),
-            (r'\^.', Comment),
-            (r'"', Comment, '#pop'),
-        ],
-        'commentString2': [
-            (r'[^(^{})]+', Comment),
-            (escape_re, Comment),
-            (r'[(|)]+', Comment),
-            (r'\^.', Comment),
-            (r'\{', Comment, '#push'),
-            (r'\}', Comment, '#pop'),
-        ],
-        'commentBlock': [
-            (r'\[', Comment, '#push'),
-            (r'\]', Comment, '#pop'),
-            (r'"', Comment, "commentString1"),
-            (r'\{', Comment, "commentString2"),
-            (r'[^(\[\]"{)]+', Comment),
-        ],
-    }
-
-    def analyse_text(text):
-        """
-        Check if code contains REBOL header and so it probably not R code
-        """
-        if re.match(r'^\s*REBOL\s*\[', text, re.IGNORECASE):
-            # The code starts with REBOL header
-            return 1.0
-        elif re.search(r'\s*REBOL\s*\[', text, re.IGNORECASE):
-            # The code contains REBOL header but also some text before it
-            return 0.5
-
-
-class RedLexer(RegexLexer):
-    """
-    A Red-language lexer.
-    """
-    name = 'Red'
-    aliases = ['red', 'red/system']
-    filenames = ['*.red', '*.reds']
-    mimetypes = ['text/x-red', 'text/x-red-system']
-    url = 'https://www.red-lang.org'
-    version_added = '2.0'
-
-    flags = re.IGNORECASE | re.MULTILINE
-
-    escape_re = r'(?:\^\([0-9a-f]{1,4}\)*)'
-
-    def word_callback(lexer, match):
-        word = match.group()
-
-        if re.match(".*:$", word):
-            yield match.start(), Generic.Subheading, word
-        elif re.match(r'(if|unless|either|any|all|while|until|loop|repeat|'
-                      r'foreach|forall|func|function|does|has|switch|'
-                      r'case|reduce|compose|get|set|print|prin|equal\?|'
-                      r'not-equal\?|strict-equal\?|lesser\?|greater\?|lesser-or-equal\?|'
-                      r'greater-or-equal\?|same\?|not|type\?|stats|'
-                      r'bind|union|replace|charset|routine)$', word):
-            yield match.start(), Name.Builtin, word
-        elif re.match(r'(make|random|reflect|to|form|mold|absolute|add|divide|multiply|negate|'
-                      r'power|remainder|round|subtract|even\?|odd\?|and~|complement|or~|xor~|'
-                      r'append|at|back|change|clear|copy|find|head|head\?|index\?|insert|'
-                      r'length\?|next|pick|poke|remove|reverse|select|sort|skip|swap|tail|tail\?|'
-                      r'take|trim|create|close|delete|modify|open|open\?|query|read|rename|'
-                      r'update|write)$', word):
-            yield match.start(), Name.Function, word
-        elif re.match(r'(yes|on|no|off|true|false|tab|cr|lf|newline|escape|slash|sp|space|null|'
-                      r'none|crlf|dot|null-byte)$', word):
-            yield match.start(), Name.Builtin.Pseudo, word
-        elif re.match(r'(#system-global|#include|#enum|#define|#either|#if|#import|#export|'
-                      r'#switch|#default|#get-definition)$', word):
-            yield match.start(), Keyword.Namespace, word
-        elif re.match(r'(system|halt|quit|quit-return|do|load|q|recycle|call|run|ask|parse|'
-                      r'raise-error|return|exit|break|alias|push|pop|probe|\?\?|spec-of|body-of|'
-                      r'quote|forever)$', word):
-            yield match.start(), Name.Exception, word
-        elif re.match(r'(action\?|block\?|char\?|datatype\?|file\?|function\?|get-path\?|zero\?|'
-                      r'get-word\?|integer\?|issue\?|lit-path\?|lit-word\?|logic\?|native\?|'
-                      r'op\?|paren\?|path\?|refinement\?|set-path\?|set-word\?|string\?|unset\?|'
-                      r'any-struct\?|none\?|word\?|any-series\?)$', word):
-            yield match.start(), Keyword, word
-        elif re.match(r'(JNICALL|stdcall|cdecl|infix)$', word):
-            yield match.start(), Keyword.Namespace, word
-        elif re.match("to-.*", word):
-            yield match.start(), Keyword, word
-        elif re.match(r'(\+|-\*\*|-|\*\*|//|/|\*|and|or|xor|=\?|===|==|=|<>|<=|>=|'
-                      r'<<<|>>>|<<|>>|<|>%)$', word):
-            yield match.start(), Operator, word
-        elif re.match(r".*\!$", word):
-            yield match.start(), Keyword.Type, word
-        elif re.match("'.*", word):
-            yield match.start(), Name.Variable.Instance, word  # lit-word
-        elif re.match("#.*", word):
-            yield match.start(), Name.Label, word  # issue
-        elif re.match("%.*", word):
-            yield match.start(), Name.Decorator, word  # file
-        elif re.match(":.*", word):
-            yield match.start(), Generic.Subheading, word  # get-word
-        else:
-            yield match.start(), Name.Variable, word
-
-    tokens = {
-        'root': [
-            (r'\s+', Text),
-            (r'#"', String.Char, 'char'),
-            (r'#\{[0-9a-f\s]*\}', Number.Hex),
-            (r'2#\{', Number.Hex, 'bin2'),
-            (r'64#\{[0-9a-z+/=\s]*\}', Number.Hex),
-            (r'([0-9a-f]+)(h)((\s)|(?=[\[\]{}"()]))',
-             bygroups(Number.Hex, Name.Variable, Whitespace)),
-            (r'"', String, 'string'),
-            (r'\{', String, 'string2'),
-            (r';#+.*\n', Comment.Special),
-            (r';\*+.*\n', Comment.Preproc),
-            (r';.*\n', Comment),
-            (r'%"', Name.Decorator, 'stringFile'),
-            (r'%[^(^{")\s\[\]]+', Name.Decorator),
-            (r'[+-]?([a-z]{1,3})?\$\d+(\.\d+)?', Number.Float),  # money
-            (r'[+-]?\d+\:\d+(\:\d+)?(\.\d+)?', String.Other),    # time
-            (r'\d+[\-/][0-9a-z]+[\-/]\d+(/\d+:\d+((:\d+)?'
-             r'([\.\d+]?([+-]?\d+:\d+)?)?)?)?', String.Other),   # date
-            (r'\d+(\.\d+)+\.\d+', Keyword.Constant),             # tuple
-            (r'\d+X\d+', Keyword.Constant),                   # pair
-            (r'[+-]?\d+(\'\d+)?([.,]\d*)?E[+-]?\d+', Number.Float),
-            (r'[+-]?\d+(\'\d+)?[.,]\d*', Number.Float),
-            (r'[+-]?\d+(\'\d+)?', Number),
-            (r'[\[\]()]', Generic.Strong),
-            (r'[a-z]+[^(^{"\s:)]*://[^(^{"\s)]*', Name.Decorator),  # url
-            (r'mailto:[^(^{"@\s)]+@[^(^{"@\s)]+', Name.Decorator),  # url
-            (r'[^(^{"@\s)]+@[^(^{"@\s)]+', Name.Decorator),         # email
-            (r'comment\s"', Comment, 'commentString1'),
-            (r'comment\s\{', Comment, 'commentString2'),
-            (r'comment\s\[', Comment, 'commentBlock'),
-            (r'comment\s[^(\s{"\[]+', Comment),
-            (r'/[^(^{^")\s/[\]]*', Name.Attribute),
-            (r'([^(^{^")\s/[\]]+)(?=[:({"\s/\[\]])', word_callback),
-            (r'<[\w:.-]*>', Name.Tag),
-            (r'<[^(<>\s")]+', Name.Tag, 'tag'),
-            (r'([^(^{")\s]+)', Text),
-        ],
-        'string': [
-            (r'[^(^")]+', String),
-            (escape_re, String.Escape),
-            (r'[(|)]+', String),
-            (r'\^.', String.Escape),
-            (r'"', String, '#pop'),
-        ],
-        'string2': [
-            (r'[^(^{})]+', String),
-            (escape_re, String.Escape),
-            (r'[(|)]+', String),
-            (r'\^.', String.Escape),
-            (r'\{', String, '#push'),
-            (r'\}', String, '#pop'),
-        ],
-        'stringFile': [
-            (r'[^(^")]+', Name.Decorator),
-            (escape_re, Name.Decorator),
-            (r'\^.', Name.Decorator),
-            (r'"', Name.Decorator, '#pop'),
-        ],
-        'char': [
-            (escape_re + '"', String.Char, '#pop'),
-            (r'\^."', String.Char, '#pop'),
-            (r'."', String.Char, '#pop'),
-        ],
-        'tag': [
-            (escape_re, Name.Tag),
-            (r'"', Name.Tag, 'tagString'),
-            (r'[^(<>\r\n")]+', Name.Tag),
-            (r'>', Name.Tag, '#pop'),
-        ],
-        'tagString': [
-            (r'[^(^")]+', Name.Tag),
-            (escape_re, Name.Tag),
-            (r'[(|)]+', Name.Tag),
-            (r'\^.', Name.Tag),
-            (r'"', Name.Tag, '#pop'),
-        ],
-        'tuple': [
-            (r'(\d+\.)+', Keyword.Constant),
-            (r'\d+', Keyword.Constant, '#pop'),
-        ],
-        'bin2': [
-            (r'\s+', Number.Hex),
-            (r'([01]\s*){8}', Number.Hex),
-            (r'\}', Number.Hex, '#pop'),
-        ],
-        'commentString1': [
-            (r'[^(^")]+', Comment),
-            (escape_re, Comment),
-            (r'[(|)]+', Comment),
-            (r'\^.', Comment),
-            (r'"', Comment, '#pop'),
-        ],
-        'commentString2': [
-            (r'[^(^{})]+', Comment),
-            (escape_re, Comment),
-            (r'[(|)]+', Comment),
-            (r'\^.', Comment),
-            (r'\{', Comment, '#push'),
-            (r'\}', Comment, '#pop'),
-        ],
-        'commentBlock': [
-            (r'\[', Comment, '#push'),
-            (r'\]', Comment, '#pop'),
-            (r'"', Comment, "commentString1"),
-            (r'\{', Comment, "commentString2"),
-            (r'[^(\[\]"{)]+', Comment),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/rego.py b/venv/lib/python3.10/site-packages/pygments/lexers/rego.py
deleted file mode 100644
index 6f2e3e9e669bcf30a05f6e0e8deef87d79f65bd2..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/rego.py
+++ /dev/null
@@ -1,57 +0,0 @@
-"""
-    pygments.lexers.rego
-    ~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for the Rego policy languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, words
-from pygments.token import Comment, Operator, Keyword, Name, String, Number, Punctuation, Whitespace
-
-class RegoLexer(RegexLexer):
-    """
-    For Rego source.
-    """
-    name = 'Rego'
-    url = 'https://www.openpolicyagent.org/docs/latest/policy-language/'
-    filenames = ['*.rego']
-    aliases = ['rego']
-    mimetypes = ['text/x-rego']
-    version_added = '2.19'
-
-    reserved_words = (
-        'as', 'contains', 'data', 'default', 'else', 'every', 'false',
-        'if', 'in', 'import', 'package', 'not', 'null',
-        'some', 'true', 'with'
-    )
-
-    builtins = (
-        # https://www.openpolicyagent.org/docs/latest/philosophy/#the-opa-document-model
-        'data',  # Global variable for accessing base and virtual documents
-        'input', # Represents synchronously pushed base documents
-    )
-
-    tokens = {
-        'root': [
-            (r'\n', Whitespace),
-            (r'\s+', Whitespace),
-            (r'#.*?$', Comment.Single),
-            (words(reserved_words, suffix=r'\b'), Keyword),
-            (words(builtins, suffix=r'\b'), Name.Builtin),
-            (r'[a-zA-Z_][a-zA-Z0-9_]*', Name),
-            (r'"(\\\\|\\"|[^"])*"', String.Double),
-            (r'`[^`]*`', String.Backtick),
-            (r'-?\d+(\.\d+)?', Number),
-            (r'(==|!=|<=|>=|:=)', Operator),  # Compound operators
-            (r'[=<>+\-*/%&|]', Operator),     # Single-character operators
-            (r'[\[\]{}(),.:;]', Punctuation),
-        ]
-    }
-
-__all__ = ['RegoLexer']
-
-
-
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/resource.py b/venv/lib/python3.10/site-packages/pygments/lexers/resource.py
deleted file mode 100644
index 9593c2124c1388aa2b7e821fada19347ab395d54..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/resource.py
+++ /dev/null
@@ -1,83 +0,0 @@
-"""
-    pygments.lexers.resource
-    ~~~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for resource definition files.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, bygroups, words
-from pygments.token import Comment, String, Number, Operator, Text, \
-    Keyword, Name
-
-__all__ = ['ResourceLexer']
-
-
-class ResourceLexer(RegexLexer):
-    """Lexer for ICU Resource bundles.
-    """
-    name = 'ResourceBundle'
-    aliases = ['resourcebundle', 'resource']
-    filenames = []
-    url = 'https://unicode-org.github.io/icu/userguide/locale/resources.html'
-    version_added = '2.0'
-
-    _types = (':table', ':array', ':string', ':bin', ':import', ':intvector',
-              ':int', ':alias')
-
-    flags = re.MULTILINE | re.IGNORECASE
-    tokens = {
-        'root': [
-            (r'//.*?$', Comment),
-            (r'"', String, 'string'),
-            (r'-?\d+', Number.Integer),
-            (r'[,{}]', Operator),
-            (r'([^\s{{:]+)(\s*)({}?)'.format('|'.join(_types)),
-             bygroups(Name, Text, Keyword)),
-            (r'\s+', Text),
-            (words(_types), Keyword),
-        ],
-        'string': [
-            (r'(\\x[0-9a-f]{2}|\\u[0-9a-f]{4}|\\U00[0-9a-f]{6}|'
-             r'\\[0-7]{1,3}|\\c.|\\[abtnvfre\'"?\\]|\\\{|[^"{\\])+', String),
-            (r'\{', String.Escape, 'msgname'),
-            (r'"', String, '#pop')
-        ],
-        'msgname': [
-            (r'([^{},]+)(\s*)', bygroups(Name, String.Escape), ('#pop', 'message'))
-        ],
-        'message': [
-            (r'\{', String.Escape, 'msgname'),
-            (r'\}', String.Escape, '#pop'),
-            (r'(,)(\s*)([a-z]+)(\s*\})',
-             bygroups(Operator, String.Escape, Keyword, String.Escape), '#pop'),
-            (r'(,)(\s*)([a-z]+)(\s*)(,)(\s*)(offset)(\s*)(:)(\s*)(-?\d+)(\s*)',
-             bygroups(Operator, String.Escape, Keyword, String.Escape, Operator,
-                      String.Escape, Operator.Word, String.Escape, Operator,
-                      String.Escape, Number.Integer, String.Escape), 'choice'),
-            (r'(,)(\s*)([a-z]+)(\s*)(,)(\s*)',
-             bygroups(Operator, String.Escape, Keyword, String.Escape, Operator,
-                      String.Escape), 'choice'),
-            (r'\s+', String.Escape)
-        ],
-        'choice': [
-            (r'(=|<|>|<=|>=|!=)(-?\d+)(\s*\{)',
-             bygroups(Operator, Number.Integer, String.Escape), 'message'),
-            (r'([a-z]+)(\s*\{)', bygroups(Keyword.Type, String.Escape), 'str'),
-            (r'\}', String.Escape, ('#pop', '#pop')),
-            (r'\s+', String.Escape)
-        ],
-        'str': [
-            (r'\}', String.Escape, '#pop'),
-            (r'\{', String.Escape, 'msgname'),
-            (r'[^{}]+', String)
-        ]
-    }
-
-    def analyse_text(text):
-        if text.startswith('root:table'):
-            return 1.0
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/ride.py b/venv/lib/python3.10/site-packages/pygments/lexers/ride.py
deleted file mode 100644
index 4d60c29cbb259a7a1d70cbf0a3ec08438f1f55b2..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/ride.py
+++ /dev/null
@@ -1,138 +0,0 @@
-"""
-    pygments.lexers.ride
-    ~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for the Ride programming language.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, words, include
-from pygments.token import Comment, Keyword, Name, Number, Punctuation, \
-    String, Text
-
-__all__ = ['RideLexer']
-
-
-class RideLexer(RegexLexer):
-    """
-    For Ride source code.
-    """
-
-    name = 'Ride'
-    aliases = ['ride']
-    filenames = ['*.ride']
-    mimetypes = ['text/x-ride']
-    url = 'https://docs.waves.tech/en/ride'
-    version_added = '2.6'
-
-    validName = r'[a-zA-Z_][a-zA-Z0-9_\']*'
-
-    builtinOps = (
-        '||', '|', '>=', '>', '==', '!',
-        '=', '<=', '<', '::', ':+', ':', '!=', '/',
-        '.', '=>', '-', '+', '*', '&&', '%', '++',
-    )
-
-    globalVariablesName = (
-        'NOALG', 'MD5', 'SHA1', 'SHA224', 'SHA256', 'SHA384', 'SHA512',
-        'SHA3224', 'SHA3256', 'SHA3384', 'SHA3512', 'nil', 'this', 'unit',
-        'height', 'lastBlock', 'Buy', 'Sell', 'CEILING', 'FLOOR', 'DOWN',
-        'HALFDOWN', 'HALFEVEN', 'HALFUP', 'UP',
-    )
-
-    typesName = (
-        'Unit', 'Int', 'Boolean', 'ByteVector', 'String', 'Address', 'Alias',
-        'Transfer', 'AssetPair', 'DataEntry', 'Order', 'Transaction',
-        'GenesisTransaction', 'PaymentTransaction', 'ReissueTransaction',
-        'BurnTransaction', 'MassTransferTransaction', 'ExchangeTransaction',
-        'TransferTransaction', 'SetAssetScriptTransaction',
-        'InvokeScriptTransaction', 'IssueTransaction', 'LeaseTransaction',
-        'LeaseCancelTransaction', 'CreateAliasTransaction',
-        'SetScriptTransaction', 'SponsorFeeTransaction', 'DataTransaction',
-        'WriteSet', 'AttachedPayment', 'ScriptTransfer', 'TransferSet',
-        'ScriptResult', 'Invocation', 'Asset', 'BlockInfo', 'Issue', 'Reissue',
-        'Burn', 'NoAlg', 'Md5', 'Sha1', 'Sha224', 'Sha256', 'Sha384', 'Sha512',
-        'Sha3224', 'Sha3256', 'Sha3384', 'Sha3512', 'BinaryEntry',
-        'BooleanEntry', 'IntegerEntry', 'StringEntry', 'List', 'Ceiling',
-        'Down', 'Floor', 'HalfDown', 'HalfEven', 'HalfUp', 'Up',
-    )
-
-    functionsName = (
-        'fraction', 'size', 'toBytes', 'take', 'drop', 'takeRight', 'dropRight',
-        'toString', 'isDefined', 'extract', 'throw', 'getElement', 'value',
-        'cons', 'toUtf8String', 'toInt', 'indexOf', 'lastIndexOf', 'split',
-        'parseInt', 'parseIntValue', 'keccak256', 'blake2b256', 'sha256',
-        'sigVerify', 'toBase58String', 'fromBase58String', 'toBase64String',
-        'fromBase64String', 'transactionById', 'transactionHeightById',
-        'getInteger', 'getBoolean', 'getBinary', 'getString',
-        'addressFromPublicKey', 'addressFromString', 'addressFromRecipient',
-        'assetBalance', 'wavesBalance', 'getIntegerValue', 'getBooleanValue',
-        'getBinaryValue', 'getStringValue', 'addressFromStringValue',
-        'assetInfo', 'rsaVerify', 'checkMerkleProof', 'median',
-        'valueOrElse', 'valueOrErrorMessage', 'contains', 'log', 'pow',
-        'toBase16String', 'fromBase16String', 'blockInfoByHeight',
-        'transferTransactionById',
-    )
-
-    reservedWords = words((
-        'match', 'case', 'else', 'func', 'if',
-        'let', 'then', '@Callable', '@Verifier',
-    ), suffix=r'\b')
-
-    tokens = {
-        'root': [
-            # Comments
-            (r'#.*', Comment.Single),
-            # Whitespace
-            (r'\s+', Text),
-            # Strings
-            (r'"', String, 'doublequote'),
-            (r'utf8\'', String, 'utf8quote'),
-            (r'base(58|64|16)\'', String, 'singlequote'),
-            # Keywords
-            (reservedWords, Keyword.Reserved),
-            (r'\{-#.*?#-\}', Keyword.Reserved),
-            (r'FOLD<\d+>', Keyword.Reserved),
-            # Types
-            (words(typesName), Keyword.Type),
-            # Main
-            # (specialName, Keyword.Reserved),
-            # Prefix Operators
-            (words(builtinOps, prefix=r'\(', suffix=r'\)'), Name.Function),
-            # Infix Operators
-            (words(builtinOps), Name.Function),
-            (words(globalVariablesName), Name.Function),
-            (words(functionsName), Name.Function),
-            # Numbers
-            include('numbers'),
-            # Variable Names
-            (validName, Name.Variable),
-            # Parens
-            (r'[,()\[\]{}]', Punctuation),
-        ],
-
-        'doublequote': [
-            (r'\\u[0-9a-fA-F]{4}', String.Escape),
-            (r'\\[nrfvb\\"]', String.Escape),
-            (r'[^"]', String),
-            (r'"', String, '#pop'),
-        ],
-
-        'utf8quote': [
-            (r'\\u[0-9a-fA-F]{4}', String.Escape),
-            (r'\\[nrfvb\\\']', String.Escape),
-            (r'[^\']', String),
-            (r'\'', String, '#pop'),
-        ],
-
-        'singlequote': [
-            (r'[^\']', String),
-            (r'\'', String, '#pop'),
-        ],
-
-        'numbers': [
-            (r'_?\d+', Number.Integer),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/rita.py b/venv/lib/python3.10/site-packages/pygments/lexers/rita.py
deleted file mode 100644
index 536aafffd2e88c9a5a4b6b4ff9d6cd42872a2162..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/rita.py
+++ /dev/null
@@ -1,42 +0,0 @@
-"""
-    pygments.lexers.rita
-    ~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for RITA language
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer
-from pygments.token import Comment, Operator, Keyword, Name, Literal, \
-    Punctuation, Whitespace
-
-__all__ = ['RitaLexer']
-
-
-class RitaLexer(RegexLexer):
-    """
-    Lexer for RITA.
-    """
-    name = 'Rita'
-    url = 'https://github.com/zaibacu/rita-dsl'
-    filenames = ['*.rita']
-    aliases = ['rita']
-    mimetypes = ['text/rita']
-    version_added = '2.11'
-
-    tokens = {
-        'root': [
-            (r'\n', Whitespace),
-            (r'\s+', Whitespace),
-            (r'#(.*?)\n', Comment.Single),
-            (r'@(.*?)\n', Operator),  # Yes, whole line as an operator
-            (r'"(\w|\d|\s|(\\")|[\'_\-./,\?\!])+?"', Literal),
-            (r'\'(\w|\d|\s|(\\\')|["_\-./,\?\!])+?\'', Literal),
-            (r'([A-Z_]+)', Keyword),
-            (r'([a-z0-9_]+)', Name),
-            (r'((->)|[!?+*|=])', Operator),
-            (r'[\(\),\{\}]', Punctuation)
-        ]
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/rnc.py b/venv/lib/python3.10/site-packages/pygments/lexers/rnc.py
deleted file mode 100644
index b7a06bb918390f59753808c2f4f99009da5c51b1..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/rnc.py
+++ /dev/null
@@ -1,66 +0,0 @@
-"""
-    pygments.lexers.rnc
-    ~~~~~~~~~~~~~~~~~~~
-
-    Lexer for Relax-NG Compact syntax
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Punctuation
-
-__all__ = ['RNCCompactLexer']
-
-
-class RNCCompactLexer(RegexLexer):
-    """
-    For RelaxNG-compact syntax.
-    """
-
-    name = 'Relax-NG Compact'
-    url = 'http://relaxng.org'
-    aliases = ['rng-compact', 'rnc']
-    filenames = ['*.rnc']
-    version_added = '2.2'
-
-    tokens = {
-        'root': [
-            (r'namespace\b', Keyword.Namespace),
-            (r'(?:default|datatypes)\b', Keyword.Declaration),
-            (r'##.*$', Comment.Preproc),
-            (r'#.*$', Comment.Single),
-            (r'"[^"]*"', String.Double),
-            # TODO single quoted strings and escape sequences outside of
-            # double-quoted strings
-            (r'(?:element|attribute|mixed)\b', Keyword.Declaration, 'variable'),
-            (r'(text\b|xsd:[^ ]+)', Keyword.Type, 'maybe_xsdattributes'),
-            (r'[,?&*=|~]|>>', Operator),
-            (r'[(){}]', Punctuation),
-            (r'.', Text),
-        ],
-
-        # a variable has been declared using `element` or `attribute`
-        'variable': [
-            (r'[^{]+', Name.Variable),
-            (r'\{', Punctuation, '#pop'),
-        ],
-
-        # after an xsd: declaration there may be attributes
-        'maybe_xsdattributes': [
-            (r'\{', Punctuation, 'xsdattributes'),
-            (r'\}', Punctuation, '#pop'),
-            (r'.', Text),
-        ],
-
-        # attributes take the form { key1 = value1 key2 = value2 ... }
-        'xsdattributes': [
-            (r'[^ =}]', Name.Attribute),
-            (r'=', Operator),
-            (r'"[^"]*"', String.Double),
-            (r'\}', Punctuation, '#pop'),
-            (r'.', Text),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/roboconf.py b/venv/lib/python3.10/site-packages/pygments/lexers/roboconf.py
deleted file mode 100644
index 31adba9f48f99a4f9c38941e3c0d73714051a994..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/roboconf.py
+++ /dev/null
@@ -1,81 +0,0 @@
-"""
-    pygments.lexers.roboconf
-    ~~~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Roboconf DSL.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, words, re
-from pygments.token import Text, Operator, Keyword, Name, Comment
-
-__all__ = ['RoboconfGraphLexer', 'RoboconfInstancesLexer']
-
-
-class RoboconfGraphLexer(RegexLexer):
-    """
-    Lexer for Roboconf graph files.
-    """
-    name = 'Roboconf Graph'
-    aliases = ['roboconf-graph']
-    filenames = ['*.graph']
-    url = 'https://roboconf.github.io/en/user-guide/graph-definition.html'
-    version_added = '2.1'
-
-    flags = re.IGNORECASE | re.MULTILINE
-    tokens = {
-        'root': [
-            # Skip white spaces
-            (r'\s+', Text),
-
-            # There is one operator
-            (r'=', Operator),
-
-            # Keywords
-            (words(('facet', 'import'), suffix=r'\s*\b', prefix=r'\b'), Keyword),
-            (words((
-                'installer', 'extends', 'exports', 'imports', 'facets',
-                'children'), suffix=r'\s*:?', prefix=r'\b'), Name),
-
-            # Comments
-            (r'#.*\n', Comment),
-
-            # Default
-            (r'[^#]', Text),
-            (r'.*\n', Text)
-        ]
-    }
-
-
-class RoboconfInstancesLexer(RegexLexer):
-    """
-    Lexer for Roboconf instances files.
-    """
-    name = 'Roboconf Instances'
-    aliases = ['roboconf-instances']
-    filenames = ['*.instances']
-    url = 'https://roboconf.github.io'
-    version_added = '2.1'
-
-    flags = re.IGNORECASE | re.MULTILINE
-    tokens = {
-        'root': [
-
-            # Skip white spaces
-            (r'\s+', Text),
-
-            # Keywords
-            (words(('instance of', 'import'), suffix=r'\s*\b', prefix=r'\b'), Keyword),
-            (words(('name', 'count'), suffix=r's*:?', prefix=r'\b'), Name),
-            (r'\s*[\w.-]+\s*:', Name),
-
-            # Comments
-            (r'#.*\n', Comment),
-
-            # Default
-            (r'[^#]', Text),
-            (r'.*\n', Text)
-        ]
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/robotframework.py b/venv/lib/python3.10/site-packages/pygments/lexers/robotframework.py
deleted file mode 100644
index f92d567503fbb371a224d531f504c59fbf77e1a8..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/robotframework.py
+++ /dev/null
@@ -1,551 +0,0 @@
-"""
-    pygments.lexers.robotframework
-    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for Robot Framework.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-#  Copyright 2012 Nokia Siemens Networks Oyj
-#
-#  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 re
-
-from pygments.lexer import Lexer
-from pygments.token import Token
-
-__all__ = ['RobotFrameworkLexer']
-
-
-HEADING = Token.Generic.Heading
-SETTING = Token.Keyword.Namespace
-IMPORT = Token.Name.Namespace
-TC_KW_NAME = Token.Generic.Subheading
-KEYWORD = Token.Name.Function
-ARGUMENT = Token.String
-VARIABLE = Token.Name.Variable
-COMMENT = Token.Comment
-SEPARATOR = Token.Punctuation
-SYNTAX = Token.Punctuation
-GHERKIN = Token.Generic.Emph
-ERROR = Token.Error
-
-
-def normalize(string, remove=''):
-    string = string.lower()
-    for char in remove + ' ':
-        if char in string:
-            string = string.replace(char, '')
-    return string
-
-
-class RobotFrameworkLexer(Lexer):
-    """
-    For Robot Framework test data.
-
-    Supports both space and pipe separated plain text formats.
-    """
-    name = 'RobotFramework'
-    url = 'http://robotframework.org'
-    aliases = ['robotframework']
-    filenames = ['*.robot', '*.resource']
-    mimetypes = ['text/x-robotframework']
-    version_added = '1.6'
-
-    def __init__(self, **options):
-        options['tabsize'] = 2
-        options['encoding'] = 'UTF-8'
-        Lexer.__init__(self, **options)
-
-    def get_tokens_unprocessed(self, text):
-        row_tokenizer = RowTokenizer()
-        var_tokenizer = VariableTokenizer()
-        index = 0
-        for row in text.splitlines():
-            for value, token in row_tokenizer.tokenize(row):
-                for value, token in var_tokenizer.tokenize(value, token):
-                    if value:
-                        yield index, token, str(value)
-                        index += len(value)
-
-
-class VariableTokenizer:
-
-    def tokenize(self, string, token):
-        var = VariableSplitter(string, identifiers='$@%&')
-        if var.start < 0 or token in (COMMENT, ERROR):
-            yield string, token
-            return
-        for value, token in self._tokenize(var, string, token):
-            if value:
-                yield value, token
-
-    def _tokenize(self, var, string, orig_token):
-        before = string[:var.start]
-        yield before, orig_token
-        yield var.identifier + '{', SYNTAX
-        yield from self.tokenize(var.base, VARIABLE)
-        yield '}', SYNTAX
-        if var.index is not None:
-            yield '[', SYNTAX
-            yield from self.tokenize(var.index, VARIABLE)
-            yield ']', SYNTAX
-        yield from self.tokenize(string[var.end:], orig_token)
-
-
-class RowTokenizer:
-
-    def __init__(self):
-        self._table = UnknownTable()
-        self._splitter = RowSplitter()
-        testcases = TestCaseTable()
-        settings = SettingTable(testcases.set_default_template)
-        variables = VariableTable()
-        keywords = KeywordTable()
-        self._tables = {'settings': settings, 'setting': settings,
-                        'metadata': settings,
-                        'variables': variables, 'variable': variables,
-                        'testcases': testcases, 'testcase': testcases,
-                        'tasks': testcases, 'task': testcases,
-                        'keywords': keywords, 'keyword': keywords,
-                        'userkeywords': keywords, 'userkeyword': keywords}
-
-    def tokenize(self, row):
-        commented = False
-        heading = False
-        for index, value in enumerate(self._splitter.split(row)):
-            # First value, and every second after that, is a separator.
-            index, separator = divmod(index-1, 2)
-            if value.startswith('#'):
-                commented = True
-            elif index == 0 and value.startswith('*'):
-                self._table = self._start_table(value)
-                heading = True
-            yield from self._tokenize(value, index, commented,
-                                      separator, heading)
-        self._table.end_row()
-
-    def _start_table(self, header):
-        name = normalize(header, remove='*')
-        return self._tables.get(name, UnknownTable())
-
-    def _tokenize(self, value, index, commented, separator, heading):
-        if commented:
-            yield value, COMMENT
-        elif separator:
-            yield value, SEPARATOR
-        elif heading:
-            yield value, HEADING
-        else:
-            yield from self._table.tokenize(value, index)
-
-
-class RowSplitter:
-    _space_splitter = re.compile('( {2,})')
-    _pipe_splitter = re.compile(r'((?:^| +)\|(?: +|$))')
-
-    def split(self, row):
-        splitter = (row.startswith('| ') and self._split_from_pipes
-                    or self._split_from_spaces)
-        yield from splitter(row)
-        yield '\n'
-
-    def _split_from_spaces(self, row):
-        yield ''  # Start with (pseudo)separator similarly as with pipes
-        yield from self._space_splitter.split(row)
-
-    def _split_from_pipes(self, row):
-        _, separator, rest = self._pipe_splitter.split(row, 1)
-        yield separator
-        while self._pipe_splitter.search(rest):
-            cell, separator, rest = self._pipe_splitter.split(rest, 1)
-            yield cell
-            yield separator
-        yield rest
-
-
-class Tokenizer:
-    _tokens = None
-
-    def __init__(self):
-        self._index = 0
-
-    def tokenize(self, value):
-        values_and_tokens = self._tokenize(value, self._index)
-        self._index += 1
-        if isinstance(values_and_tokens, type(Token)):
-            values_and_tokens = [(value, values_and_tokens)]
-        return values_and_tokens
-
-    def _tokenize(self, value, index):
-        index = min(index, len(self._tokens) - 1)
-        return self._tokens[index]
-
-    def _is_assign(self, value):
-        if value.endswith('='):
-            value = value[:-1].strip()
-        var = VariableSplitter(value, identifiers='$@&')
-        return var.start == 0 and var.end == len(value)
-
-
-class Comment(Tokenizer):
-    _tokens = (COMMENT,)
-
-
-class Setting(Tokenizer):
-    _tokens = (SETTING, ARGUMENT)
-    _keyword_settings = ('suitesetup', 'suiteprecondition', 'suiteteardown',
-                         'suitepostcondition', 'testsetup', 'tasksetup', 'testprecondition',
-                         'testteardown','taskteardown', 'testpostcondition', 'testtemplate', 'tasktemplate')
-    _import_settings = ('library', 'resource', 'variables')
-    _other_settings = ('documentation', 'metadata', 'forcetags', 'defaulttags',
-                       'testtimeout','tasktimeout')
-    _custom_tokenizer = None
-
-    def __init__(self, template_setter=None):
-        Tokenizer.__init__(self)
-        self._template_setter = template_setter
-
-    def _tokenize(self, value, index):
-        if index == 1 and self._template_setter:
-            self._template_setter(value)
-        if index == 0:
-            normalized = normalize(value)
-            if normalized in self._keyword_settings:
-                self._custom_tokenizer = KeywordCall(support_assign=False)
-            elif normalized in self._import_settings:
-                self._custom_tokenizer = ImportSetting()
-            elif normalized not in self._other_settings:
-                return ERROR
-        elif self._custom_tokenizer:
-            return self._custom_tokenizer.tokenize(value)
-        return Tokenizer._tokenize(self, value, index)
-
-
-class ImportSetting(Tokenizer):
-    _tokens = (IMPORT, ARGUMENT)
-
-
-class TestCaseSetting(Setting):
-    _keyword_settings = ('setup', 'precondition', 'teardown', 'postcondition',
-                         'template')
-    _import_settings = ()
-    _other_settings = ('documentation', 'tags', 'timeout')
-
-    def _tokenize(self, value, index):
-        if index == 0:
-            type = Setting._tokenize(self, value[1:-1], index)
-            return [('[', SYNTAX), (value[1:-1], type), (']', SYNTAX)]
-        return Setting._tokenize(self, value, index)
-
-
-class KeywordSetting(TestCaseSetting):
-    _keyword_settings = ('teardown',)
-    _other_settings = ('documentation', 'arguments', 'return', 'timeout', 'tags')
-
-
-class Variable(Tokenizer):
-    _tokens = (SYNTAX, ARGUMENT)
-
-    def _tokenize(self, value, index):
-        if index == 0 and not self._is_assign(value):
-            return ERROR
-        return Tokenizer._tokenize(self, value, index)
-
-
-class KeywordCall(Tokenizer):
-    _tokens = (KEYWORD, ARGUMENT)
-
-    def __init__(self, support_assign=True):
-        Tokenizer.__init__(self)
-        self._keyword_found = not support_assign
-        self._assigns = 0
-
-    def _tokenize(self, value, index):
-        if not self._keyword_found and self._is_assign(value):
-            self._assigns += 1
-            return SYNTAX  # VariableTokenizer tokenizes this later.
-        if self._keyword_found:
-            return Tokenizer._tokenize(self, value, index - self._assigns)
-        self._keyword_found = True
-        return GherkinTokenizer().tokenize(value, KEYWORD)
-
-
-class GherkinTokenizer:
-    _gherkin_prefix = re.compile('^(Given|When|Then|And|But) ', re.IGNORECASE)
-
-    def tokenize(self, value, token):
-        match = self._gherkin_prefix.match(value)
-        if not match:
-            return [(value, token)]
-        end = match.end()
-        return [(value[:end], GHERKIN), (value[end:], token)]
-
-
-class TemplatedKeywordCall(Tokenizer):
-    _tokens = (ARGUMENT,)
-
-
-class ForLoop(Tokenizer):
-
-    def __init__(self):
-        Tokenizer.__init__(self)
-        self._in_arguments = False
-
-    def _tokenize(self, value, index):
-        token = self._in_arguments and ARGUMENT or SYNTAX
-        if value.upper() in ('IN', 'IN RANGE'):
-            self._in_arguments = True
-        return token
-
-
-class _Table:
-    _tokenizer_class = None
-
-    def __init__(self, prev_tokenizer=None):
-        self._tokenizer = self._tokenizer_class()
-        self._prev_tokenizer = prev_tokenizer
-        self._prev_values_on_row = []
-
-    def tokenize(self, value, index):
-        if self._continues(value, index):
-            self._tokenizer = self._prev_tokenizer
-            yield value, SYNTAX
-        else:
-            yield from self._tokenize(value, index)
-        self._prev_values_on_row.append(value)
-
-    def _continues(self, value, index):
-        return value == '...' and all(self._is_empty(t)
-                                      for t in self._prev_values_on_row)
-
-    def _is_empty(self, value):
-        return value in ('', '\\')
-
-    def _tokenize(self, value, index):
-        return self._tokenizer.tokenize(value)
-
-    def end_row(self):
-        self.__init__(prev_tokenizer=self._tokenizer)
-
-
-class UnknownTable(_Table):
-    _tokenizer_class = Comment
-
-    def _continues(self, value, index):
-        return False
-
-
-class VariableTable(_Table):
-    _tokenizer_class = Variable
-
-
-class SettingTable(_Table):
-    _tokenizer_class = Setting
-
-    def __init__(self, template_setter, prev_tokenizer=None):
-        _Table.__init__(self, prev_tokenizer)
-        self._template_setter = template_setter
-
-    def _tokenize(self, value, index):
-        if index == 0 and normalize(value) == 'testtemplate':
-            self._tokenizer = Setting(self._template_setter)
-        return _Table._tokenize(self, value, index)
-
-    def end_row(self):
-        self.__init__(self._template_setter, prev_tokenizer=self._tokenizer)
-
-
-class TestCaseTable(_Table):
-    _setting_class = TestCaseSetting
-    _test_template = None
-    _default_template = None
-
-    @property
-    def _tokenizer_class(self):
-        if self._test_template or (self._default_template and
-                                   self._test_template is not False):
-            return TemplatedKeywordCall
-        return KeywordCall
-
-    def _continues(self, value, index):
-        return index > 0 and _Table._continues(self, value, index)
-
-    def _tokenize(self, value, index):
-        if index == 0:
-            if value:
-                self._test_template = None
-            return GherkinTokenizer().tokenize(value, TC_KW_NAME)
-        if index == 1 and self._is_setting(value):
-            if self._is_template(value):
-                self._test_template = False
-                self._tokenizer = self._setting_class(self.set_test_template)
-            else:
-                self._tokenizer = self._setting_class()
-        if index == 1 and self._is_for_loop(value):
-            self._tokenizer = ForLoop()
-        if index == 1 and self._is_empty(value):
-            return [(value, SYNTAX)]
-        return _Table._tokenize(self, value, index)
-
-    def _is_setting(self, value):
-        return value.startswith('[') and value.endswith(']')
-
-    def _is_template(self, value):
-        return normalize(value) == '[template]'
-
-    def _is_for_loop(self, value):
-        return value.startswith(':') and normalize(value, remove=':') == 'for'
-
-    def set_test_template(self, template):
-        self._test_template = self._is_template_set(template)
-
-    def set_default_template(self, template):
-        self._default_template = self._is_template_set(template)
-
-    def _is_template_set(self, template):
-        return normalize(template) not in ('', '\\', 'none', '${empty}')
-
-
-class KeywordTable(TestCaseTable):
-    _tokenizer_class = KeywordCall
-    _setting_class = KeywordSetting
-
-    def _is_template(self, value):
-        return False
-
-
-# Following code copied directly from Robot Framework 2.7.5.
-
-class VariableSplitter:
-
-    def __init__(self, string, identifiers):
-        self.identifier = None
-        self.base = None
-        self.index = None
-        self.start = -1
-        self.end = -1
-        self._identifiers = identifiers
-        self._may_have_internal_variables = False
-        try:
-            self._split(string)
-        except ValueError:
-            pass
-        else:
-            self._finalize()
-
-    def get_replaced_base(self, variables):
-        if self._may_have_internal_variables:
-            return variables.replace_string(self.base)
-        return self.base
-
-    def _finalize(self):
-        self.identifier = self._variable_chars[0]
-        self.base = ''.join(self._variable_chars[2:-1])
-        self.end = self.start + len(self._variable_chars)
-        if self._has_list_or_dict_variable_index():
-            self.index = ''.join(self._list_and_dict_variable_index_chars[1:-1])
-            self.end += len(self._list_and_dict_variable_index_chars)
-
-    def _has_list_or_dict_variable_index(self):
-        return self._list_and_dict_variable_index_chars\
-        and self._list_and_dict_variable_index_chars[-1] == ']'
-
-    def _split(self, string):
-        start_index, max_index = self._find_variable(string)
-        self.start = start_index
-        self._open_curly = 1
-        self._state = self._variable_state
-        self._variable_chars = [string[start_index], '{']
-        self._list_and_dict_variable_index_chars = []
-        self._string = string
-        start_index += 2
-        for index, char in enumerate(string[start_index:]):
-            index += start_index  # Giving start to enumerate only in Py 2.6+
-            try:
-                self._state(char, index)
-            except StopIteration:
-                return
-            if index  == max_index and not self._scanning_list_variable_index():
-                return
-
-    def _scanning_list_variable_index(self):
-        return self._state in [self._waiting_list_variable_index_state,
-                               self._list_variable_index_state]
-
-    def _find_variable(self, string):
-        max_end_index = string.rfind('}')
-        if max_end_index == -1:
-            raise ValueError('No variable end found')
-        if self._is_escaped(string, max_end_index):
-            return self._find_variable(string[:max_end_index])
-        start_index = self._find_start_index(string, 1, max_end_index)
-        if start_index == -1:
-            raise ValueError('No variable start found')
-        return start_index, max_end_index
-
-    def _find_start_index(self, string, start, end):
-        index = string.find('{', start, end) - 1
-        if index < 0:
-            return -1
-        if self._start_index_is_ok(string, index):
-            return index
-        return self._find_start_index(string, index+2, end)
-
-    def _start_index_is_ok(self, string, index):
-        return string[index] in self._identifiers\
-        and not self._is_escaped(string, index)
-
-    def _is_escaped(self, string, index):
-        escaped = False
-        while index > 0 and string[index-1] == '\\':
-            index -= 1
-            escaped = not escaped
-        return escaped
-
-    def _variable_state(self, char, index):
-        self._variable_chars.append(char)
-        if char == '}' and not self._is_escaped(self._string, index):
-            self._open_curly -= 1
-            if self._open_curly == 0:
-                if not self._is_list_or_dict_variable():
-                    raise StopIteration
-                self._state = self._waiting_list_variable_index_state
-        elif char in self._identifiers:
-            self._state = self._internal_variable_start_state
-
-    def _is_list_or_dict_variable(self):
-        return self._variable_chars[0] in ('@','&')
-
-    def _internal_variable_start_state(self, char, index):
-        self._state = self._variable_state
-        if char == '{':
-            self._variable_chars.append(char)
-            self._open_curly += 1
-            self._may_have_internal_variables = True
-        else:
-            self._variable_state(char, index)
-
-    def _waiting_list_variable_index_state(self, char, index):
-        if char != '[':
-            raise StopIteration
-        self._list_and_dict_variable_index_chars.append(char)
-        self._state = self._list_variable_index_state
-
-    def _list_variable_index_state(self, char, index):
-        self._list_and_dict_variable_index_chars.append(char)
-        if char == ']':
-            raise StopIteration
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/ruby.py b/venv/lib/python3.10/site-packages/pygments/lexers/ruby.py
deleted file mode 100644
index 72aaeb5fec08c7c14e3f827d9c297c485a087e67..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/ruby.py
+++ /dev/null
@@ -1,518 +0,0 @@
-"""
-    pygments.lexers.ruby
-    ~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Ruby and related languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import Lexer, RegexLexer, ExtendedRegexLexer, include, \
-    bygroups, default, LexerContext, do_insertions, words, line_re
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Error, Generic, Whitespace
-from pygments.util import shebang_matches
-
-__all__ = ['RubyLexer', 'RubyConsoleLexer', 'FancyLexer']
-
-
-RUBY_OPERATORS = (
-    '*', '**', '-', '+', '-@', '+@', '/', '%', '&', '|', '^', '`', '~',
-    '[]', '[]=', '<<', '>>', '<', '<>', '<=>', '>', '>=', '==', '==='
-)
-
-
-class RubyLexer(ExtendedRegexLexer):
-    """
-    For Ruby source code.
-    """
-
-    name = 'Ruby'
-    url = 'http://www.ruby-lang.org'
-    aliases = ['ruby', 'rb', 'duby']
-    filenames = ['*.rb', '*.rbw', 'Rakefile', '*.rake', '*.gemspec',
-                 '*.rbx', '*.duby', 'Gemfile', 'Vagrantfile']
-    mimetypes = ['text/x-ruby', 'application/x-ruby']
-    version_added = ''
-
-    flags = re.DOTALL | re.MULTILINE
-
-    def heredoc_callback(self, match, ctx):
-        # okay, this is the hardest part of parsing Ruby...
-        # match: 1 = <<[-~]?, 2 = quote? 3 = name 4 = quote? 5 = rest of line
-
-        start = match.start(1)
-        yield start, Operator, match.group(1)        # <<[-~]?
-        yield match.start(2), String.Heredoc, match.group(2)   # quote ", ', `
-        yield match.start(3), String.Delimiter, match.group(3) # heredoc name
-        yield match.start(4), String.Heredoc, match.group(4)   # quote again
-
-        heredocstack = ctx.__dict__.setdefault('heredocstack', [])
-        outermost = not bool(heredocstack)
-        heredocstack.append((match.group(1) in ('<<-', '<<~'), match.group(3)))
-
-        ctx.pos = match.start(5)
-        ctx.end = match.end(5)
-        # this may find other heredocs, so limit the recursion depth
-        if len(heredocstack) < 100:
-            yield from self.get_tokens_unprocessed(context=ctx)
-        else:
-            yield ctx.pos, String.Heredoc, match.group(5)
-        ctx.pos = match.end()
-
-        if outermost:
-            # this is the outer heredoc again, now we can process them all
-            for tolerant, hdname in heredocstack:
-                lines = []
-                for match in line_re.finditer(ctx.text, ctx.pos):
-                    if tolerant:
-                        check = match.group().strip()
-                    else:
-                        check = match.group().rstrip()
-                    if check == hdname:
-                        for amatch in lines:
-                            yield amatch.start(), String.Heredoc, amatch.group()
-                        yield match.start(), String.Delimiter, match.group()
-                        ctx.pos = match.end()
-                        break
-                    else:
-                        lines.append(match)
-                else:
-                    # end of heredoc not found -- error!
-                    for amatch in lines:
-                        yield amatch.start(), Error, amatch.group()
-            ctx.end = len(ctx.text)
-            del heredocstack[:]
-
-    def gen_rubystrings_rules():
-        def intp_regex_callback(self, match, ctx):
-            yield match.start(1), String.Regex, match.group(1)  # begin
-            nctx = LexerContext(match.group(3), 0, ['interpolated-regex'])
-            for i, t, v in self.get_tokens_unprocessed(context=nctx):
-                yield match.start(3)+i, t, v
-            yield match.start(4), String.Regex, match.group(4)  # end[mixounse]*
-            ctx.pos = match.end()
-
-        def intp_string_callback(self, match, ctx):
-            yield match.start(1), String.Other, match.group(1)
-            nctx = LexerContext(match.group(3), 0, ['interpolated-string'])
-            for i, t, v in self.get_tokens_unprocessed(context=nctx):
-                yield match.start(3)+i, t, v
-            yield match.start(4), String.Other, match.group(4)  # end
-            ctx.pos = match.end()
-
-        states = {}
-        states['strings'] = [
-            # easy ones
-            (r'\:@{0,2}[a-zA-Z_]\w*[!?]?', String.Symbol),
-            (words(RUBY_OPERATORS, prefix=r'\:@{0,2}'), String.Symbol),
-            (r":'(\\\\|\\[^\\]|[^'\\])*'", String.Symbol),
-            (r':"', String.Symbol, 'simple-sym'),
-            (r'([a-zA-Z_]\w*)(:)(?!:)',
-             bygroups(String.Symbol, Punctuation)),  # Since Ruby 1.9
-            (r'"', String.Double, 'simple-string-double'),
-            (r"'", String.Single, 'simple-string-single'),
-            (r'(?', '<>', 'ab'):
-            states[name+'-intp-string'] = [
-                (r'\\[\\' + bracecc + ']', String.Other),
-                (lbrace, String.Other, '#push'),
-                (rbrace, String.Other, '#pop'),
-                include('string-intp-escaped'),
-                (r'[\\#' + bracecc + ']', String.Other),
-                (r'[^\\#' + bracecc + ']+', String.Other),
-            ]
-            states['strings'].append((r'%[QWx]?' + lbrace, String.Other,
-                                      name+'-intp-string'))
-            states[name+'-string'] = [
-                (r'\\[\\' + bracecc + ']', String.Other),
-                (lbrace, String.Other, '#push'),
-                (rbrace, String.Other, '#pop'),
-                (r'[\\#' + bracecc + ']', String.Other),
-                (r'[^\\#' + bracecc + ']+', String.Other),
-            ]
-            states['strings'].append((r'%[qsw]' + lbrace, String.Other,
-                                      name+'-string'))
-            states[name+'-regex'] = [
-                (r'\\[\\' + bracecc + ']', String.Regex),
-                (lbrace, String.Regex, '#push'),
-                (rbrace + '[mixounse]*', String.Regex, '#pop'),
-                include('string-intp'),
-                (r'[\\#' + bracecc + ']', String.Regex),
-                (r'[^\\#' + bracecc + ']+', String.Regex),
-            ]
-            states['strings'].append((r'%r' + lbrace, String.Regex,
-                                      name+'-regex'))
-
-        # these must come after %!
-        states['strings'] += [
-            # %r regex
-            (r'(%r([\W_]))((?:\\\2|(?!\2).)*)(\2[mixounse]*)',
-             intp_regex_callback),
-            # regular fancy strings with qsw
-            (r'%[qsw]([\W_])((?:\\\1|(?!\1).)*)\1', String.Other),
-            (r'(%[QWx]([\W_]))((?:\\\2|(?!\2).)*)(\2)',
-             intp_string_callback),
-            # special forms of fancy strings after operators or
-            # in method calls with braces
-            (r'(?<=[-+/*%=<>&!^|~,(])(\s*)(%([\t ])(?:(?:\\\3|(?!\3).)*)\3)',
-             bygroups(Whitespace, String.Other, None)),
-            # and because of fixed width lookbehinds the whole thing a
-            # second time for line startings...
-            (r'^(\s*)(%([\t ])(?:(?:\\\3|(?!\3).)*)\3)',
-             bygroups(Whitespace, String.Other, None)),
-            # all regular fancy strings without qsw
-            (r'(%([^a-zA-Z0-9\s]))((?:\\\2|(?!\2).)*)(\2)',
-             intp_string_callback),
-        ]
-
-        return states
-
-    tokens = {
-        'root': [
-            (r'\A#!.+?$', Comment.Hashbang),
-            (r'#.*?$', Comment.Single),
-            (r'=begin\s.*?\n=end.*?$', Comment.Multiline),
-            # keywords
-            (words((
-                'BEGIN', 'END', 'alias', 'begin', 'break', 'case', 'defined?',
-                'do', 'else', 'elsif', 'end', 'ensure', 'for', 'if', 'in', 'next', 'redo',
-                'rescue', 'raise', 'retry', 'return', 'super', 'then', 'undef',
-                'unless', 'until', 'when', 'while', 'yield'), suffix=r'\b'),
-             Keyword),
-            # start of function, class and module names
-            (r'(module)(\s+)([a-zA-Z_]\w*'
-             r'(?:::[a-zA-Z_]\w*)*)',
-             bygroups(Keyword, Whitespace, Name.Namespace)),
-            (r'(def)(\s+)', bygroups(Keyword, Whitespace), 'funcname'),
-            (r'def(?=[*%&^`~+-/\[<>=])', Keyword, 'funcname'),
-            (r'(class)(\s+)', bygroups(Keyword, Whitespace), 'classname'),
-            # special methods
-            (words((
-                'initialize', 'new', 'loop', 'include', 'extend', 'raise', 'attr_reader',
-                'attr_writer', 'attr_accessor', 'attr', 'catch', 'throw', 'private',
-                'module_function', 'public', 'protected', 'true', 'false', 'nil'),
-                suffix=r'\b'),
-             Keyword.Pseudo),
-            (r'(not|and|or)\b', Operator.Word),
-            (words((
-                'autoload', 'block_given', 'const_defined', 'eql', 'equal', 'frozen', 'include',
-                'instance_of', 'is_a', 'iterator', 'kind_of', 'method_defined', 'nil',
-                'private_method_defined', 'protected_method_defined',
-                'public_method_defined', 'respond_to', 'tainted'), suffix=r'\?'),
-             Name.Builtin),
-            (r'(chomp|chop|exit|gsub|sub)!', Name.Builtin),
-            (words((
-                'Array', 'Float', 'Integer', 'String', '__id__', '__send__', 'abort',
-                'ancestors', 'at_exit', 'autoload', 'binding', 'callcc', 'caller',
-                'catch', 'chomp', 'chop', 'class_eval', 'class_variables',
-                'clone', 'const_defined?', 'const_get', 'const_missing', 'const_set',
-                'constants', 'display', 'dup', 'eval', 'exec', 'exit', 'extend', 'fail', 'fork',
-                'format', 'freeze', 'getc', 'gets', 'global_variables', 'gsub',
-                'hash', 'id', 'included_modules', 'inspect', 'instance_eval',
-                'instance_method', 'instance_methods',
-                'instance_variable_get', 'instance_variable_set', 'instance_variables',
-                'lambda', 'load', 'local_variables', 'loop',
-                'method', 'method_missing', 'methods', 'module_eval', 'name',
-                'object_id', 'open', 'p', 'print', 'printf', 'private_class_method',
-                'private_instance_methods',
-                'private_methods', 'proc', 'protected_instance_methods',
-                'protected_methods', 'public_class_method',
-                'public_instance_methods', 'public_methods',
-                'putc', 'puts', 'raise', 'rand', 'readline', 'readlines', 'require',
-                'scan', 'select', 'self', 'send', 'set_trace_func', 'singleton_methods', 'sleep',
-                'split', 'sprintf', 'srand', 'sub', 'syscall', 'system', 'taint',
-                'test', 'throw', 'to_a', 'to_s', 'trace_var', 'trap', 'untaint',
-                'untrace_var', 'warn'), prefix=r'(?~!:])|'
-             r'(?<=(?:\s|;)when\s)|'
-             r'(?<=(?:\s|;)or\s)|'
-             r'(?<=(?:\s|;)and\s)|'
-             r'(?<=\.index\s)|'
-             r'(?<=\.scan\s)|'
-             r'(?<=\.sub\s)|'
-             r'(?<=\.sub!\s)|'
-             r'(?<=\.gsub\s)|'
-             r'(?<=\.gsub!\s)|'
-             r'(?<=\.match\s)|'
-             r'(?<=(?:\s|;)if\s)|'
-             r'(?<=(?:\s|;)elsif\s)|'
-             r'(?<=^when\s)|'
-             r'(?<=^index\s)|'
-             r'(?<=^scan\s)|'
-             r'(?<=^sub\s)|'
-             r'(?<=^gsub\s)|'
-             r'(?<=^sub!\s)|'
-             r'(?<=^gsub!\s)|'
-             r'(?<=^match\s)|'
-             r'(?<=^if\s)|'
-             r'(?<=^elsif\s)'
-             r')(\s*)(/)', bygroups(Text, String.Regex), 'multiline-regex'),
-            # multiline regex (in method calls or subscripts)
-            (r'(?<=\(|,|\[)/', String.Regex, 'multiline-regex'),
-            # multiline regex (this time the funny no whitespace rule)
-            (r'(\s+)(/)(?![\s=])', bygroups(Whitespace, String.Regex),
-             'multiline-regex'),
-            # lex numbers and ignore following regular expressions which
-            # are division operators in fact (grrrr. i hate that. any
-            # better ideas?)
-            # since pygments 0.7 we also eat a "?" operator after numbers
-            # so that the char operator does not work. Chars are not allowed
-            # there so that you can use the ternary operator.
-            # stupid example:
-            #   x>=0?n[x]:""
-            (r'(0_?[0-7]+(?:_[0-7]+)*)(\s*)([/?])?',
-             bygroups(Number.Oct, Whitespace, Operator)),
-            (r'(0x[0-9A-Fa-f]+(?:_[0-9A-Fa-f]+)*)(\s*)([/?])?',
-             bygroups(Number.Hex, Whitespace, Operator)),
-            (r'(0b[01]+(?:_[01]+)*)(\s*)([/?])?',
-             bygroups(Number.Bin, Whitespace, Operator)),
-            (r'([\d]+(?:_\d+)*)(\s*)([/?])?',
-             bygroups(Number.Integer, Whitespace, Operator)),
-            # Names
-            (r'@@[a-zA-Z_]\w*', Name.Variable.Class),
-            (r'@[a-zA-Z_]\w*', Name.Variable.Instance),
-            (r'\$\w+', Name.Variable.Global),
-            (r'\$[!@&`\'+~=/\\,;.<>_*$?:"^-]', Name.Variable.Global),
-            (r'\$-[0adFiIlpvw]', Name.Variable.Global),
-            (r'::', Operator),
-            include('strings'),
-            # chars
-            (r'\?(\\[MC]-)*'  # modifiers
-             r'(\\([\\abefnrstv#"\']|x[a-fA-F0-9]{1,2}|[0-7]{1,3})|\S)'
-             r'(?!\w)',
-             String.Char),
-            (r'[A-Z]\w+', Name.Constant),
-            # this is needed because ruby attributes can look
-            # like keywords (class) or like this: ` ?!?
-            (words(RUBY_OPERATORS, prefix=r'(\.|::)'),
-             bygroups(Operator, Name.Operator)),
-            (r'(\.|::)([a-zA-Z_]\w*[!?]?|[*%&^`~+\-/\[<>=])',
-             bygroups(Operator, Name)),
-            (r'[a-zA-Z_]\w*[!?]?', Name),
-            (r'(\[|\]|\*\*|<>?|>=|<=|<=>|=~|={3}|'
-             r'!~|&&?|\|\||\.{1,3})', Operator),
-            (r'[-+/*%=<>&!^|~]=?', Operator),
-            (r'[(){};,/?:\\]', Punctuation),
-            (r'\s+', Whitespace)
-        ],
-        'funcname': [
-            (r'\(', Punctuation, 'defexpr'),
-            (r'(?:([a-zA-Z_]\w*)(\.))?'  # optional scope name, like "self."
-             r'('
-                r'[a-zA-Z\u0080-\uffff][a-zA-Z0-9_\u0080-\uffff]*[!?=]?'  # method name
-                r'|!=|!~|=~|\*\*?|[-+!~]@?|[/%&|^]|<=>|<[<=]?|>[>=]?|===?'  # or operator override
-                r'|\[\]=?'  # or element reference/assignment override
-                r'|`'  # or the undocumented backtick override
-             r')',
-             bygroups(Name.Class, Operator, Name.Function), '#pop'),
-            default('#pop')
-        ],
-        'classname': [
-            (r'\(', Punctuation, 'defexpr'),
-            (r'<<', Operator, '#pop'),
-            (r'[A-Z_]\w*', Name.Class, '#pop'),
-            default('#pop')
-        ],
-        'defexpr': [
-            (r'(\))(\.|::)?', bygroups(Punctuation, Operator), '#pop'),
-            (r'\(', Operator, '#push'),
-            include('root')
-        ],
-        'in-intp': [
-            (r'\{', String.Interpol, '#push'),
-            (r'\}', String.Interpol, '#pop'),
-            include('root'),
-        ],
-        'string-intp': [
-            (r'#\{', String.Interpol, 'in-intp'),
-            (r'#@@?[a-zA-Z_]\w*', String.Interpol),
-            (r'#\$[a-zA-Z_]\w*', String.Interpol)
-        ],
-        'string-intp-escaped': [
-            include('string-intp'),
-            (r'\\([\\abefnrstv#"\']|x[a-fA-F0-9]{1,2}|[0-7]{1,3})',
-             String.Escape)
-        ],
-        'interpolated-regex': [
-            include('string-intp'),
-            (r'[\\#]', String.Regex),
-            (r'[^\\#]+', String.Regex),
-        ],
-        'interpolated-string': [
-            include('string-intp'),
-            (r'[\\#]', String.Other),
-            (r'[^\\#]+', String.Other),
-        ],
-        'multiline-regex': [
-            include('string-intp'),
-            (r'\\\\', String.Regex),
-            (r'\\/', String.Regex),
-            (r'[\\#]', String.Regex),
-            (r'[^\\/#]+', String.Regex),
-            (r'/[mixounse]*', String.Regex, '#pop'),
-        ],
-        'end-part': [
-            (r'.+', Comment.Preproc, '#pop')
-        ]
-    }
-    tokens.update(gen_rubystrings_rules())
-
-    def analyse_text(text):
-        return shebang_matches(text, r'ruby(1\.\d)?')
-
-
-class RubyConsoleLexer(Lexer):
-    """
-    For Ruby interactive console (**irb**) output.
-    """
-    name = 'Ruby irb session'
-    aliases = ['rbcon', 'irb']
-    mimetypes = ['text/x-ruby-shellsession']
-    url = 'https://www.ruby-lang.org'
-    version_added = ''
-    _example = 'rbcon/console'
-
-    _prompt_re = re.compile(r'irb\([a-zA-Z_]\w*\):\d{3}:\d+[>*"\'] '
-                            r'|>> |\?> ')
-
-    def get_tokens_unprocessed(self, text):
-        rblexer = RubyLexer(**self.options)
-
-        curcode = ''
-        insertions = []
-        for match in line_re.finditer(text):
-            line = match.group()
-            m = self._prompt_re.match(line)
-            if m is not None:
-                end = m.end()
-                insertions.append((len(curcode),
-                                   [(0, Generic.Prompt, line[:end])]))
-                curcode += line[end:]
-            else:
-                if curcode:
-                    yield from do_insertions(
-                        insertions, rblexer.get_tokens_unprocessed(curcode))
-                    curcode = ''
-                    insertions = []
-                yield match.start(), Generic.Output, line
-        if curcode:
-            yield from do_insertions(
-                insertions, rblexer.get_tokens_unprocessed(curcode))
-
-
-class FancyLexer(RegexLexer):
-    """
-    Pygments Lexer For Fancy.
-
-    Fancy is a self-hosted, pure object-oriented, dynamic,
-    class-based, concurrent general-purpose programming language
-    running on Rubinius, the Ruby VM.
-    """
-    name = 'Fancy'
-    url = 'https://github.com/bakkdoor/fancy'
-    filenames = ['*.fy', '*.fancypack']
-    aliases = ['fancy', 'fy']
-    mimetypes = ['text/x-fancysrc']
-    version_added = '1.5'
-
-    tokens = {
-        # copied from PerlLexer:
-        'balanced-regex': [
-            (r'/(\\\\|\\[^\\]|[^/\\])*/[egimosx]*', String.Regex, '#pop'),
-            (r'!(\\\\|\\[^\\]|[^!\\])*![egimosx]*', String.Regex, '#pop'),
-            (r'\\(\\\\|[^\\])*\\[egimosx]*', String.Regex, '#pop'),
-            (r'\{(\\\\|\\[^\\]|[^}\\])*\}[egimosx]*', String.Regex, '#pop'),
-            (r'<(\\\\|\\[^\\]|[^>\\])*>[egimosx]*', String.Regex, '#pop'),
-            (r'\[(\\\\|\\[^\\]|[^\]\\])*\][egimosx]*', String.Regex, '#pop'),
-            (r'\((\\\\|\\[^\\]|[^)\\])*\)[egimosx]*', String.Regex, '#pop'),
-            (r'@(\\\\|\\[^\\]|[^@\\])*@[egimosx]*', String.Regex, '#pop'),
-            (r'%(\\\\|\\[^\\]|[^%\\])*%[egimosx]*', String.Regex, '#pop'),
-            (r'\$(\\\\|\\[^\\]|[^$\\])*\$[egimosx]*', String.Regex, '#pop'),
-        ],
-        'root': [
-            (r'\s+', Whitespace),
-
-            # balanced delimiters (copied from PerlLexer):
-            (r's\{(\\\\|\\[^\\]|[^}\\])*\}\s*', String.Regex, 'balanced-regex'),
-            (r's<(\\\\|\\[^\\]|[^>\\])*>\s*', String.Regex, 'balanced-regex'),
-            (r's\[(\\\\|\\[^\\]|[^\]\\])*\]\s*', String.Regex, 'balanced-regex'),
-            (r's\((\\\\|\\[^\\]|[^)\\])*\)\s*', String.Regex, 'balanced-regex'),
-            (r'm?/(\\\\|\\[^\\]|[^///\n])*/[gcimosx]*', String.Regex),
-            (r'm(?=[/!\\{<\[(@%$])', String.Regex, 'balanced-regex'),
-
-            # Comments
-            (r'#(.*?)\n', Comment.Single),
-            # Symbols
-            (r'\'([^\'\s\[\](){}]+|\[\])', String.Symbol),
-            # Multi-line DoubleQuotedString
-            (r'"""(\\\\|\\[^\\]|[^\\])*?"""', String),
-            # DoubleQuotedString
-            (r'"(\\\\|\\[^\\]|[^"\\])*"', String),
-            # keywords
-            (r'(def|class|try|catch|finally|retry|return|return_local|match|'
-             r'case|->|=>)\b', Keyword),
-            # constants
-            (r'(self|super|nil|false|true)\b', Name.Constant),
-            (r'[(){};,/?|:\\]', Punctuation),
-            # names
-            (words((
-                'Object', 'Array', 'Hash', 'Directory', 'File', 'Class', 'String',
-                'Number', 'Enumerable', 'FancyEnumerable', 'Block', 'TrueClass',
-                'NilClass', 'FalseClass', 'Tuple', 'Symbol', 'Stack', 'Set',
-                'FancySpec', 'Method', 'Package', 'Range'), suffix=r'\b'),
-             Name.Builtin),
-            # functions
-            (r'[a-zA-Z](\w|[-+?!=*/^><%])*:', Name.Function),
-            # operators, must be below functions
-            (r'[-+*/~,<>=&!?%^\[\].$]+', Operator),
-            (r'[A-Z]\w*', Name.Constant),
-            (r'@[a-zA-Z_]\w*', Name.Variable.Instance),
-            (r'@@[a-zA-Z_]\w*', Name.Variable.Class),
-            ('@@?', Operator),
-            (r'[a-zA-Z_]\w*', Name),
-            # numbers - / checks are necessary to avoid mismarking regexes,
-            # see comment in RubyLexer
-            (r'(0[oO]?[0-7]+(?:_[0-7]+)*)(\s*)([/?])?',
-             bygroups(Number.Oct, Whitespace, Operator)),
-            (r'(0[xX][0-9A-Fa-f]+(?:_[0-9A-Fa-f]+)*)(\s*)([/?])?',
-             bygroups(Number.Hex, Whitespace, Operator)),
-            (r'(0[bB][01]+(?:_[01]+)*)(\s*)([/?])?',
-             bygroups(Number.Bin, Whitespace, Operator)),
-            (r'([\d]+(?:_\d+)*)(\s*)([/?])?',
-             bygroups(Number.Integer, Whitespace, Operator)),
-            (r'\d+([eE][+-]?[0-9]+)|\d+\.\d+([eE][+-]?[0-9]+)?', Number.Float),
-            (r'\d+', Number.Integer)
-        ]
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/rust.py b/venv/lib/python3.10/site-packages/pygments/lexers/rust.py
deleted file mode 100644
index 63410475553d44d7aa080e075ac71180f3160d6c..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/rust.py
+++ /dev/null
@@ -1,222 +0,0 @@
-"""
-    pygments.lexers.rust
-    ~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for the Rust language.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, include, bygroups, words, default
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Whitespace
-
-__all__ = ['RustLexer']
-
-
-class RustLexer(RegexLexer):
-    """
-    Lexer for the Rust programming language (version 1.47).
-    """
-    name = 'Rust'
-    url = 'https://www.rust-lang.org/'
-    filenames = ['*.rs', '*.rs.in']
-    aliases = ['rust', 'rs']
-    mimetypes = ['text/rust', 'text/x-rust']
-    version_added = '1.6'
-
-    keyword_types = (words((
-        'u8', 'u16', 'u32', 'u64', 'u128', 'i8', 'i16', 'i32', 'i64', 'i128',
-        'usize', 'isize', 'f32', 'f64', 'char', 'str', 'bool',
-    ), suffix=r'\b'), Keyword.Type)
-
-    builtin_funcs_types = (words((
-        'Copy', 'Send', 'Sized', 'Sync', 'Unpin',
-        'Drop', 'Fn', 'FnMut', 'FnOnce', 'drop',
-        'Box', 'ToOwned', 'Clone',
-        'PartialEq', 'PartialOrd', 'Eq', 'Ord',
-        'AsRef', 'AsMut', 'Into', 'From', 'Default',
-        'Iterator', 'Extend', 'IntoIterator', 'DoubleEndedIterator',
-        'ExactSizeIterator',
-        'Option', 'Some', 'None',
-        'Result', 'Ok', 'Err',
-        'String', 'ToString', 'Vec',
-    ), suffix=r'\b'), Name.Builtin)
-
-    builtin_macros = (words((
-        'asm', 'assert', 'assert_eq', 'assert_ne', 'cfg', 'column',
-        'compile_error', 'concat', 'concat_idents', 'dbg', 'debug_assert',
-        'debug_assert_eq', 'debug_assert_ne', 'env', 'eprint', 'eprintln',
-        'file', 'format', 'format_args', 'format_args_nl', 'global_asm',
-        'include', 'include_bytes', 'include_str',
-        'is_aarch64_feature_detected',
-        'is_arm_feature_detected',
-        'is_mips64_feature_detected',
-        'is_mips_feature_detected',
-        'is_powerpc64_feature_detected',
-        'is_powerpc_feature_detected',
-        'is_x86_feature_detected',
-        'line', 'llvm_asm', 'log_syntax', 'macro_rules', 'matches',
-        'module_path', 'option_env', 'panic', 'print', 'println', 'stringify',
-        'thread_local', 'todo', 'trace_macros', 'unimplemented', 'unreachable',
-        'vec', 'write', 'writeln',
-    ), suffix=r'!'), Name.Function.Magic)
-
-    tokens = {
-        'root': [
-            # rust allows a file to start with a shebang, but if the first line
-            # starts with #![ then it's not a shebang but a crate attribute.
-            (r'#![^[\r\n].*$', Comment.Preproc),
-            default('base'),
-        ],
-        'base': [
-            # Whitespace and Comments
-            (r'\n', Whitespace),
-            (r'\s+', Whitespace),
-            (r'//!.*?\n', String.Doc),
-            (r'///(\n|[^/].*?\n)', String.Doc),
-            (r'//(.*?)\n', Comment.Single),
-            (r'/\*\*(\n|[^/*])', String.Doc, 'doccomment'),
-            (r'/\*!', String.Doc, 'doccomment'),
-            (r'/\*', Comment.Multiline, 'comment'),
-
-            # Macro parameters
-            (r"""\$([a-zA-Z_]\w*|\(,?|\),?|,?)""", Comment.Preproc),
-            # Keywords
-            (words(('as', 'async', 'await', 'box', 'const', 'crate', 'dyn',
-                    'else', 'extern', 'for', 'if', 'impl', 'in', 'loop',
-                    'match', 'move', 'mut', 'pub', 'ref', 'return', 'static',
-                    'super', 'trait', 'unsafe', 'use', 'where', 'while'),
-                   suffix=r'\b'), Keyword),
-            (words(('abstract', 'become', 'do', 'final', 'macro', 'override',
-                    'priv', 'typeof', 'try', 'unsized', 'virtual', 'yield'),
-                   suffix=r'\b'), Keyword.Reserved),
-            (r'(true|false)\b', Keyword.Constant),
-            (r'self\b', Name.Builtin.Pseudo),
-            (r'mod\b', Keyword, 'modname'),
-            (r'let\b', Keyword.Declaration),
-            (r'fn\b', Keyword, 'funcname'),
-            (r'(struct|enum|type|union)\b', Keyword, 'typename'),
-            (r'(default)(\s+)(type|fn)\b', bygroups(Keyword, Whitespace, Keyword)),
-            keyword_types,
-            (r'[sS]elf\b', Name.Builtin.Pseudo),
-            # Prelude (taken from Rust's src/libstd/prelude.rs)
-            builtin_funcs_types,
-            builtin_macros,
-            # Path separators, so types don't catch them.
-            (r'::\b', Punctuation),
-            # Types in positions.
-            (r'(?::|->)', Punctuation, 'typename'),
-            # Labels
-            (r'(break|continue)(\b\s*)(\'[A-Za-z_]\w*)?',
-             bygroups(Keyword, Text.Whitespace, Name.Label)),
-
-            # Character literals
-            (r"""'(\\['"\\nrt]|\\x[0-7][0-9a-fA-F]|\\0"""
-             r"""|\\u\{[0-9a-fA-F]{1,6}\}|.)'""",
-             String.Char),
-            (r"""b'(\\['"\\nrt]|\\x[0-9a-fA-F]{2}|\\0"""
-             r"""|\\u\{[0-9a-fA-F]{1,6}\}|.)'""",
-             String.Char),
-
-            # Binary literals
-            (r'0b[01_]+', Number.Bin, 'number_lit'),
-            # Octal literals
-            (r'0o[0-7_]+', Number.Oct, 'number_lit'),
-            # Hexadecimal literals
-            (r'0[xX][0-9a-fA-F_]+', Number.Hex, 'number_lit'),
-            # Decimal literals
-            (r'[0-9][0-9_]*(\.[0-9_]+[eE][+\-]?[0-9_]+|'
-             r'\.[0-9_]*(?!\.)|[eE][+\-]?[0-9_]+)', Number.Float,
-             'number_lit'),
-            (r'[0-9][0-9_]*', Number.Integer, 'number_lit'),
-
-            # String literals
-            (r'b"', String, 'bytestring'),
-            (r'"', String, 'string'),
-            (r'(?s)b?r(#*)".*?"\1', String),
-
-            # Lifetime names
-            (r"'", Operator, 'lifetime'),
-
-            # Operators and Punctuation
-            (r'\.\.=?', Operator),
-            (r'[{}()\[\],.;]', Punctuation),
-            (r'[+\-*/%&|<>^!~@=:?]', Operator),
-
-            # Identifiers
-            (r'[a-zA-Z_]\w*', Name),
-            # Raw identifiers
-            (r'r#[a-zA-Z_]\w*', Name),
-
-            # Attributes
-            (r'#!?\[', Comment.Preproc, 'attribute['),
-
-            # Misc
-            # Lone hashes: not used in Rust syntax, but allowed in macro
-            # arguments, most famously for quote::quote!()
-            (r'#', Punctuation),
-        ],
-        'comment': [
-            (r'[^*/]+', Comment.Multiline),
-            (r'/\*', Comment.Multiline, '#push'),
-            (r'\*/', Comment.Multiline, '#pop'),
-            (r'[*/]', Comment.Multiline),
-        ],
-        'doccomment': [
-            (r'[^*/]+', String.Doc),
-            (r'/\*', String.Doc, '#push'),
-            (r'\*/', String.Doc, '#pop'),
-            (r'[*/]', String.Doc),
-        ],
-        'modname': [
-            (r'\s+', Whitespace),
-            (r'[a-zA-Z_]\w*', Name.Namespace, '#pop'),
-            default('#pop'),
-        ],
-        'funcname': [
-            (r'\s+', Whitespace),
-            (r'[a-zA-Z_]\w*', Name.Function, '#pop'),
-            default('#pop'),
-        ],
-        'typename': [
-            (r'\s+', Whitespace),
-            (r'&', Keyword.Pseudo),
-            (r"'", Operator, 'lifetime'),
-            builtin_funcs_types,
-            keyword_types,
-            (r'[a-zA-Z_]\w*', Name.Class, '#pop'),
-            default('#pop'),
-        ],
-        'lifetime': [
-            (r"(static|_)", Name.Builtin),
-            (r"[a-zA-Z_]+\w*", Name.Attribute),
-            default('#pop'),
-        ],
-        'number_lit': [
-            (r'[ui](8|16|32|64|size)', Keyword, '#pop'),
-            (r'f(32|64)', Keyword, '#pop'),
-            default('#pop'),
-        ],
-        'string': [
-            (r'"', String, '#pop'),
-            (r"""\\['"\\nrt]|\\x[0-7][0-9a-fA-F]|\\0"""
-             r"""|\\u\{[0-9a-fA-F]{1,6}\}""", String.Escape),
-            (r'[^\\"]+', String),
-            (r'\\', String),
-        ],
-        'bytestring': [
-            (r"""\\x[89a-fA-F][0-9a-fA-F]""", String.Escape),
-            include('string'),
-        ],
-        'attribute_common': [
-            (r'"', String, 'string'),
-            (r'\[', Comment.Preproc, 'attribute['),
-        ],
-        'attribute[': [
-            include('attribute_common'),
-            (r'\]', Comment.Preproc, '#pop'),
-            (r'[^"\]\[]+', Comment.Preproc),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/sas.py b/venv/lib/python3.10/site-packages/pygments/lexers/sas.py
deleted file mode 100644
index 1b2ad432d231d21192c4f3653ef2e541588afd1d..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/sas.py
+++ /dev/null
@@ -1,227 +0,0 @@
-"""
-    pygments.lexers.sas
-    ~~~~~~~~~~~~~~~~~~~
-
-    Lexer for SAS.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-from pygments.lexer import RegexLexer, include, words
-from pygments.token import Comment, Keyword, Name, Number, String, Text, \
-    Other, Generic
-
-__all__ = ['SASLexer']
-
-
-class SASLexer(RegexLexer):
-    """
-    For SAS files.
-    """
-    # Syntax from syntax/sas.vim by James Kidd 
-
-    name      = 'SAS'
-    aliases   = ['sas']
-    filenames = ['*.SAS', '*.sas']
-    mimetypes = ['text/x-sas', 'text/sas', 'application/x-sas']
-    url = 'https://en.wikipedia.org/wiki/SAS_(software)'
-    version_added = '2.2'
-    flags     = re.IGNORECASE | re.MULTILINE
-
-    builtins_macros = (
-        "bquote", "nrbquote", "cmpres", "qcmpres", "compstor", "datatyp",
-        "display", "do", "else", "end", "eval", "global", "goto", "if",
-        "index", "input", "keydef", "label", "left", "length", "let",
-        "local", "lowcase", "macro", "mend", "nrquote",
-        "nrstr", "put", "qleft", "qlowcase", "qscan",
-        "qsubstr", "qsysfunc", "qtrim", "quote", "qupcase", "scan",
-        "str", "substr", "superq", "syscall", "sysevalf", "sysexec",
-        "sysfunc", "sysget", "syslput", "sysprod", "sysrc", "sysrput",
-        "then", "to", "trim", "unquote", "until", "upcase", "verify",
-        "while", "window"
-    )
-
-    builtins_conditionals = (
-        "do", "if", "then", "else", "end", "until", "while"
-    )
-
-    builtins_statements = (
-        "abort", "array", "attrib", "by", "call", "cards", "cards4",
-        "catname", "continue", "datalines", "datalines4", "delete", "delim",
-        "delimiter", "display", "dm", "drop", "endsas", "error", "file",
-        "filename", "footnote", "format", "goto", "in", "infile", "informat",
-        "input", "keep", "label", "leave", "length", "libname", "link",
-        "list", "lostcard", "merge", "missing", "modify", "options", "output",
-        "out", "page", "put", "redirect", "remove", "rename", "replace",
-        "retain", "return", "select", "set", "skip", "startsas", "stop",
-        "title", "update", "waitsas", "where", "window", "x", "systask"
-    )
-
-    builtins_sql = (
-        "add", "and", "alter", "as", "cascade", "check", "create",
-        "delete", "describe", "distinct", "drop", "foreign", "from",
-        "group", "having", "index", "insert", "into", "in", "key", "like",
-        "message", "modify", "msgtype", "not", "null", "on", "or",
-        "order", "primary", "references", "reset", "restrict", "select",
-        "set", "table", "unique", "update", "validate", "view", "where"
-    )
-
-    builtins_functions = (
-        "abs", "addr", "airy", "arcos", "arsin", "atan", "attrc",
-        "attrn", "band", "betainv", "blshift", "bnot", "bor",
-        "brshift", "bxor", "byte", "cdf", "ceil", "cexist", "cinv",
-        "close", "cnonct", "collate", "compbl", "compound",
-        "compress", "cos", "cosh", "css", "curobs", "cv", "daccdb",
-        "daccdbsl", "daccsl", "daccsyd", "dacctab", "dairy", "date",
-        "datejul", "datepart", "datetime", "day", "dclose", "depdb",
-        "depdbsl", "depsl", "depsyd",
-        "deptab", "dequote", "dhms", "dif", "digamma",
-        "dim", "dinfo", "dnum", "dopen", "doptname", "doptnum",
-        "dread", "dropnote", "dsname", "erf", "erfc", "exist", "exp",
-        "fappend", "fclose", "fcol", "fdelete", "fetch", "fetchobs",
-        "fexist", "fget", "fileexist", "filename", "fileref",
-        "finfo", "finv", "fipname", "fipnamel", "fipstate", "floor",
-        "fnonct", "fnote", "fopen", "foptname", "foptnum", "fpoint",
-        "fpos", "fput", "fread", "frewind", "frlen", "fsep", "fuzz",
-        "fwrite", "gaminv", "gamma", "getoption", "getvarc", "getvarn",
-        "hbound", "hms", "hosthelp", "hour", "ibessel", "index",
-        "indexc", "indexw", "input", "inputc", "inputn", "int",
-        "intck", "intnx", "intrr", "irr", "jbessel", "juldate",
-        "kurtosis", "lag", "lbound", "left", "length", "lgamma",
-        "libname", "libref", "log", "log10", "log2", "logpdf", "logpmf",
-        "logsdf", "lowcase", "max", "mdy", "mean", "min", "minute",
-        "mod", "month", "mopen", "mort", "n", "netpv", "nmiss",
-        "normal", "note", "npv", "open", "ordinal", "pathname",
-        "pdf", "peek", "peekc", "pmf", "point", "poisson", "poke",
-        "probbeta", "probbnml", "probchi", "probf", "probgam",
-        "probhypr", "probit", "probnegb", "probnorm", "probt",
-        "put", "putc", "putn", "qtr", "quote", "ranbin", "rancau",
-        "ranexp", "rangam", "range", "rank", "rannor", "ranpoi",
-        "rantbl", "rantri", "ranuni", "repeat", "resolve", "reverse",
-        "rewind", "right", "round", "saving", "scan", "sdf", "second",
-        "sign", "sin", "sinh", "skewness", "soundex", "spedis",
-        "sqrt", "std", "stderr", "stfips", "stname", "stnamel",
-        "substr", "sum", "symget", "sysget", "sysmsg", "sysprod",
-        "sysrc", "system", "tan", "tanh", "time", "timepart", "tinv",
-        "tnonct", "today", "translate", "tranwrd", "trigamma",
-        "trim", "trimn", "trunc", "uniform", "upcase", "uss", "var",
-        "varfmt", "varinfmt", "varlabel", "varlen", "varname",
-        "varnum", "varray", "varrayx", "vartype", "verify", "vformat",
-        "vformatd", "vformatdx", "vformatn", "vformatnx", "vformatw",
-        "vformatwx", "vformatx", "vinarray", "vinarrayx", "vinformat",
-        "vinformatd", "vinformatdx", "vinformatn", "vinformatnx",
-        "vinformatw", "vinformatwx", "vinformatx", "vlabel",
-        "vlabelx", "vlength", "vlengthx", "vname", "vnamex", "vtype",
-        "vtypex", "weekday", "year", "yyq", "zipfips", "zipname",
-        "zipnamel", "zipstate"
-    )
-
-    tokens = {
-        'root': [
-            include('comments'),
-            include('proc-data'),
-            include('cards-datalines'),
-            include('logs'),
-            include('general'),
-            (r'.', Text),
-        ],
-        # SAS is multi-line regardless, but * is ended by ;
-        'comments': [
-            (r'^\s*\*.*?;', Comment),
-            (r'/\*.*?\*/', Comment),
-            (r'^\s*\*(.|\n)*?;', Comment.Multiline),
-            (r'/[*](.|\n)*?[*]/', Comment.Multiline),
-        ],
-        # Special highlight for proc, data, quit, run
-        'proc-data': [
-            (r'(^|;)\s*(proc \w+|data|run|quit)[\s;]',
-             Keyword.Reserved),
-        ],
-        # Special highlight cards and datalines
-        'cards-datalines': [
-            (r'^\s*(datalines|cards)\s*;\s*$', Keyword, 'data'),
-        ],
-        'data': [
-            (r'(.|\n)*^\s*;\s*$', Other, '#pop'),
-        ],
-        # Special highlight for put NOTE|ERROR|WARNING (order matters)
-        'logs': [
-            (r'\n?^\s*%?put ', Keyword, 'log-messages'),
-        ],
-        'log-messages': [
-            (r'NOTE(:|-).*', Generic, '#pop'),
-            (r'WARNING(:|-).*', Generic.Emph, '#pop'),
-            (r'ERROR(:|-).*', Generic.Error, '#pop'),
-            include('general'),
-        ],
-        'general': [
-            include('keywords'),
-            include('vars-strings'),
-            include('special'),
-            include('numbers'),
-        ],
-        # Keywords, statements, functions, macros
-        'keywords': [
-            (words(builtins_statements,
-                   prefix = r'\b',
-                   suffix = r'\b'),
-             Keyword),
-            (words(builtins_sql,
-                   prefix = r'\b',
-                   suffix = r'\b'),
-             Keyword),
-            (words(builtins_conditionals,
-                   prefix = r'\b',
-                   suffix = r'\b'),
-             Keyword),
-            (words(builtins_macros,
-                   prefix = r'%',
-                   suffix = r'\b'),
-             Name.Builtin),
-            (words(builtins_functions,
-                   prefix = r'\b',
-                   suffix = r'\('),
-             Name.Builtin),
-        ],
-        # Strings and user-defined variables and macros (order matters)
-        'vars-strings': [
-            (r'&[a-z_]\w{0,31}\.?', Name.Variable),
-            (r'%[a-z_]\w{0,31}', Name.Function),
-            (r'\'', String, 'string_squote'),
-            (r'"', String, 'string_dquote'),
-        ],
-        'string_squote': [
-            ('\'', String, '#pop'),
-            (r'\\\\|\\"|\\\n', String.Escape),
-            # AFAIK, macro variables are not evaluated in single quotes
-            # (r'&', Name.Variable, 'validvar'),
-            (r'[^$\'\\]+', String),
-            (r'[$\'\\]', String),
-        ],
-        'string_dquote': [
-            (r'"', String, '#pop'),
-            (r'\\\\|\\"|\\\n', String.Escape),
-            (r'&', Name.Variable, 'validvar'),
-            (r'[^$&"\\]+', String),
-            (r'[$"\\]', String),
-        ],
-        'validvar': [
-            (r'[a-z_]\w{0,31}\.?', Name.Variable, '#pop'),
-        ],
-        # SAS numbers and special variables
-        'numbers': [
-            (r'\b[+-]?([0-9]+(\.[0-9]+)?|\.[0-9]+|\.)(E[+-]?[0-9]+)?i?\b',
-             Number),
-        ],
-        'special': [
-            (r'(null|missing|_all_|_automatic_|_character_|_n_|'
-             r'_infile_|_name_|_null_|_numeric_|_user_|_webout_)',
-             Keyword.Constant),
-        ],
-        # 'operators': [
-        #     (r'(-|=|<=|>=|<|>|<>|&|!=|'
-        #      r'\||\*|\+|\^|/|!|~|~=)', Operator)
-        # ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/savi.py b/venv/lib/python3.10/site-packages/pygments/lexers/savi.py
deleted file mode 100644
index 1e443ae302a647b6c401d9c0e6fe451e12732dde..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/savi.py
+++ /dev/null
@@ -1,171 +0,0 @@
-"""
-    pygments.lexers.savi
-    ~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for Savi.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, bygroups, include
-from pygments.token import Whitespace, Keyword, Name, String, Number, \
-  Operator, Punctuation, Comment, Generic, Error
-
-__all__ = ['SaviLexer']
-
-
-# The canonical version of this file can be found in the following repository,
-# where it is kept in sync with any language changes, as well as the other
-# pygments-like lexers that are maintained for use with other tools:
-# - https://github.com/savi-lang/savi/blob/main/tooling/pygments/lexers/savi.py
-#
-# If you're changing this file in the pygments repository, please ensure that
-# any changes you make are also propagated to the official Savi repository,
-# in order to avoid accidental clobbering of your changes later when an update
-# from the Savi repository flows forward into the pygments repository.
-#
-# If you're changing this file in the Savi repository, please ensure that
-# any changes you make are also reflected in the other pygments-like lexers
-# (rouge, vscode, etc) so that all of the lexers can be kept cleanly in sync.
-
-class SaviLexer(RegexLexer):
-    """
-    For Savi source code.
-
-    .. versionadded: 2.10
-    """
-
-    name = 'Savi'
-    url = 'https://github.com/savi-lang/savi'
-    aliases = ['savi']
-    filenames = ['*.savi']
-    version_added = ''
-
-    tokens = {
-      "root": [
-        # Line Comment
-        (r'//.*?$', Comment.Single),
-
-        # Doc Comment
-        (r'::.*?$', Comment.Single),
-
-        # Capability Operator
-        (r'(\')(\w+)(?=[^\'])', bygroups(Operator, Name)),
-
-        # Double-Quote String
-        (r'\w?"', String.Double, "string.double"),
-
-        # Single-Char String
-        (r"'", String.Char, "string.char"),
-
-        # Type Name
-        (r'(_?[A-Z]\w*)', Name.Class),
-
-        # Nested Type Name
-        (r'(\.)(\s*)(_?[A-Z]\w*)', bygroups(Punctuation, Whitespace, Name.Class)),
-
-        # Declare
-        (r'^([ \t]*)(:\w+)',
-          bygroups(Whitespace, Name.Tag),
-          "decl"),
-
-        # Error-Raising Calls/Names
-        (r'((\w+|\+|\-|\*)\!)', Generic.Deleted),
-
-        # Numeric Values
-        (r'\b\d([\d_]*(\.[\d_]+)?)\b', Number),
-
-        # Hex Numeric Values
-        (r'\b0x([0-9a-fA-F_]+)\b', Number.Hex),
-
-        # Binary Numeric Values
-        (r'\b0b([01_]+)\b', Number.Bin),
-
-        # Function Call (with braces)
-        (r'\w+(?=\()', Name.Function),
-
-        # Function Call (with receiver)
-        (r'(\.)(\s*)(\w+)', bygroups(Punctuation, Whitespace, Name.Function)),
-
-        # Function Call (with self receiver)
-        (r'(@)(\w+)', bygroups(Punctuation, Name.Function)),
-
-        # Parenthesis
-        (r'\(', Punctuation, "root"),
-        (r'\)', Punctuation, "#pop"),
-
-        # Brace
-        (r'\{', Punctuation, "root"),
-        (r'\}', Punctuation, "#pop"),
-
-        # Bracket
-        (r'\[', Punctuation, "root"),
-        (r'(\])(\!)', bygroups(Punctuation, Generic.Deleted), "#pop"),
-        (r'\]', Punctuation, "#pop"),
-
-        # Punctuation
-        (r'[,;:\.@]', Punctuation),
-
-        # Piping Operators
-        (r'(\|\>)', Operator),
-
-        # Branching Operators
-        (r'(\&\&|\|\||\?\?|\&\?|\|\?|\.\?)', Operator),
-
-        # Comparison Operators
-        (r'(\<\=\>|\=\~|\=\=|\<\=|\>\=|\<|\>)', Operator),
-
-        # Arithmetic Operators
-        (r'(\+|\-|\/|\*|\%)', Operator),
-
-        # Assignment Operators
-        (r'(\=)', Operator),
-
-        # Other Operators
-        (r'(\!|\<\<|\<|\&|\|)', Operator),
-
-        # Identifiers
-        (r'\b\w+\b', Name),
-
-        # Whitespace
-        (r'[ \t\r]+\n*|\n+', Whitespace),
-      ],
-
-      # Declare (nested rules)
-      "decl": [
-        (r'\b[a-z_]\w*\b(?!\!)', Keyword.Declaration),
-        (r':', Punctuation, "#pop"),
-        (r'\n', Whitespace, "#pop"),
-        include("root"),
-      ],
-
-      # Double-Quote String (nested rules)
-      "string.double": [
-        (r'\\\(', String.Interpol, "string.interpolation"),
-        (r'\\u[0-9a-fA-F]{4}', String.Escape),
-        (r'\\x[0-9a-fA-F]{2}', String.Escape),
-        (r'\\[bfnrt\\\']', String.Escape),
-        (r'\\"', String.Escape),
-        (r'"', String.Double, "#pop"),
-        (r'[^\\"]+', String.Double),
-        (r'.', Error),
-      ],
-
-      # Single-Char String (nested rules)
-      "string.char": [
-        (r'\\u[0-9a-fA-F]{4}', String.Escape),
-        (r'\\x[0-9a-fA-F]{2}', String.Escape),
-        (r'\\[bfnrt\\\']', String.Escape),
-        (r"\\'", String.Escape),
-        (r"'", String.Char, "#pop"),
-        (r"[^\\']+", String.Char),
-        (r'.', Error),
-      ],
-
-      # Interpolation inside String (nested rules)
-      "string.interpolation": [
-        (r"\)", String.Interpol, "#pop"),
-        include("root"),
-      ]
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/scdoc.py b/venv/lib/python3.10/site-packages/pygments/lexers/scdoc.py
deleted file mode 100644
index 8e850d02ed8ca77a90c9de217fac89dccd60b9ab..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/scdoc.py
+++ /dev/null
@@ -1,85 +0,0 @@
-"""
-    pygments.lexers.scdoc
-    ~~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for scdoc, a simple man page generator.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, include, bygroups, using, this
-from pygments.token import Text, Comment, Keyword, String, Generic
-
-__all__ = ['ScdocLexer']
-
-
-class ScdocLexer(RegexLexer):
-    """
-    `scdoc` is a simple man page generator for POSIX systems written in C99.
-    """
-    name = 'scdoc'
-    url = 'https://git.sr.ht/~sircmpwn/scdoc'
-    aliases = ['scdoc', 'scd']
-    filenames = ['*.scd', '*.scdoc']
-    version_added = '2.5'
-    flags = re.MULTILINE
-
-    tokens = {
-        'root': [
-            # comment
-            (r'^(;.+\n)', bygroups(Comment)),
-
-            # heading with pound prefix
-            (r'^(#)([^#].+\n)', bygroups(Generic.Heading, Text)),
-            (r'^(#{2})(.+\n)', bygroups(Generic.Subheading, Text)),
-            # bulleted lists
-            (r'^(\s*)([*-])(\s)(.+\n)',
-            bygroups(Text, Keyword, Text, using(this, state='inline'))),
-            # numbered lists
-            (r'^(\s*)(\.+\.)( .+\n)',
-            bygroups(Text, Keyword, using(this, state='inline'))),
-            # quote
-            (r'^(\s*>\s)(.+\n)', bygroups(Keyword, Generic.Emph)),
-            # text block
-            (r'^(```\n)([\w\W]*?)(^```$)', bygroups(String, Text, String)),
-
-            include('inline'),
-        ],
-        'inline': [
-            # escape
-            (r'\\.', Text),
-            # underlines
-            (r'(\s)(_[^_]+_)(\W|\n)', bygroups(Text, Generic.Emph, Text)),
-            # bold
-            (r'(\s)(\*[^*]+\*)(\W|\n)', bygroups(Text, Generic.Strong, Text)),
-            # inline code
-            (r'`[^`]+`', String.Backtick),
-
-            # general text, must come last!
-            (r'[^\\\s]+', Text),
-            (r'.', Text),
-        ],
-    }
-
-    def analyse_text(text):
-        """We checks for bold and underline text with * and _. Also
-        every scdoc file must start with a strictly defined first line."""
-        result = 0
-
-        if '*' in text:
-            result += 0.01
-
-        if '_' in text:
-            result += 0.01
-
-        # name(section) ["left_footer" ["center_header"]]
-        first_line = text.partition('\n')[0]
-        scdoc_preamble_pattern = r'^.*\([1-7]\)( "[^"]+"){0,2}$'
-
-        if re.search(scdoc_preamble_pattern, first_line):
-            result += 0.5
-
-        return result
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/scripting.py b/venv/lib/python3.10/site-packages/pygments/lexers/scripting.py
deleted file mode 100644
index 6e494c33b8ec2c939c6e55921a03aba232f5a456..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/scripting.py
+++ /dev/null
@@ -1,1616 +0,0 @@
-"""
-    pygments.lexers.scripting
-    ~~~~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for scripting and embedded languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import RegexLexer, include, bygroups, default, combined, \
-    words
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Error, Whitespace, Other
-from pygments.util import get_bool_opt, get_list_opt
-
-__all__ = ['LuaLexer', 'LuauLexer', 'MoonScriptLexer', 'ChaiscriptLexer', 'LSLLexer',
-           'AppleScriptLexer', 'RexxLexer', 'MOOCodeLexer', 'HybrisLexer',
-           'EasytrieveLexer', 'JclLexer', 'MiniScriptLexer']
-
-
-def all_lua_builtins():
-    from pygments.lexers._lua_builtins import MODULES
-    return [w for values in MODULES.values() for w in values]
-
-class LuaLexer(RegexLexer):
-    """
-    For Lua source code.
-
-    Additional options accepted:
-
-    `func_name_highlighting`
-        If given and ``True``, highlight builtin function names
-        (default: ``True``).
-    `disabled_modules`
-        If given, must be a list of module names whose function names
-        should not be highlighted. By default all modules are highlighted.
-
-        To get a list of allowed modules have a look into the
-        `_lua_builtins` module:
-
-        .. sourcecode:: pycon
-
-            >>> from pygments.lexers._lua_builtins import MODULES
-            >>> MODULES.keys()
-            ['string', 'coroutine', 'modules', 'io', 'basic', ...]
-    """
-
-    name = 'Lua'
-    url = 'https://www.lua.org/'
-    aliases = ['lua']
-    filenames = ['*.lua', '*.wlua']
-    mimetypes = ['text/x-lua', 'application/x-lua']
-    version_added = ''
-
-    _comment_multiline = r'(?:--\[(?P=*)\[[\w\W]*?\](?P=level)\])'
-    _comment_single = r'(?:--.*$)'
-    _space = r'(?:\s+(?!\s))'
-    _s = rf'(?:{_comment_multiline}|{_comment_single}|{_space})'
-    _name = r'(?:[^\W\d]\w*)'
-
-    tokens = {
-        'root': [
-            # Lua allows a file to start with a shebang.
-            (r'#!.*', Comment.Preproc),
-            default('base'),
-        ],
-        'ws': [
-            (_comment_multiline, Comment.Multiline),
-            (_comment_single, Comment.Single),
-            (_space, Whitespace),
-        ],
-        'base': [
-            include('ws'),
-
-            (r'(?i)0x[\da-f]*(\.[\da-f]*)?(p[+-]?\d+)?', Number.Hex),
-            (r'(?i)(\d*\.\d+|\d+\.\d*)(e[+-]?\d+)?', Number.Float),
-            (r'(?i)\d+e[+-]?\d+', Number.Float),
-            (r'\d+', Number.Integer),
-
-            # multiline strings
-            (r'(?s)\[(=*)\[.*?\]\1\]', String),
-
-            (r'::', Punctuation, 'label'),
-            (r'\.{3}', Punctuation),
-            (r'[=<>|~&+\-*/%#^]+|\.\.', Operator),
-            (r'[\[\]{}().,:;]+', Punctuation),
-            (r'(and|or|not)\b', Operator.Word),
-
-            (words([
-                'break', 'do', 'else', 'elseif', 'end', 'for', 'if', 'in',
-                'repeat', 'return', 'then', 'until', 'while'
-            ], suffix=r'\b'), Keyword.Reserved),
-            (r'goto\b', Keyword.Reserved, 'goto'),
-            (r'(local)\b', Keyword.Declaration),
-            (r'(true|false|nil)\b', Keyword.Constant),
-
-            (r'(function)\b', Keyword.Reserved, 'funcname'),
-
-            (words(all_lua_builtins(), suffix=r"\b"), Name.Builtin),
-            (fr'[A-Za-z_]\w*(?={_s}*[.:])', Name.Variable, 'varname'),
-            (fr'[A-Za-z_]\w*(?={_s}*\()', Name.Function),
-            (r'[A-Za-z_]\w*', Name.Variable),
-
-            ("'", String.Single, combined('stringescape', 'sqs')),
-            ('"', String.Double, combined('stringescape', 'dqs'))
-        ],
-
-        'varname': [
-            include('ws'),
-            (r'\.\.', Operator, '#pop'),
-            (r'[.:]', Punctuation),
-            (rf'{_name}(?={_s}*[.:])', Name.Property),
-            (rf'{_name}(?={_s}*\()', Name.Function, '#pop'),
-            (_name, Name.Property, '#pop'),
-        ],
-
-        'funcname': [
-            include('ws'),
-            (r'[.:]', Punctuation),
-            (rf'{_name}(?={_s}*[.:])', Name.Class),
-            (_name, Name.Function, '#pop'),
-            # inline function
-            (r'\(', Punctuation, '#pop'),
-        ],
-
-        'goto': [
-            include('ws'),
-            (_name, Name.Label, '#pop'),
-        ],
-
-        'label': [
-            include('ws'),
-            (r'::', Punctuation, '#pop'),
-            (_name, Name.Label),
-        ],
-
-        'stringescape': [
-            (r'\\([abfnrtv\\"\']|[\r\n]{1,2}|z\s*|x[0-9a-fA-F]{2}|\d{1,3}|'
-             r'u\{[0-9a-fA-F]+\})', String.Escape),
-        ],
-
-        'sqs': [
-            (r"'", String.Single, '#pop'),
-            (r"[^\\']+", String.Single),
-        ],
-
-        'dqs': [
-            (r'"', String.Double, '#pop'),
-            (r'[^\\"]+', String.Double),
-        ]
-    }
-
-    def __init__(self, **options):
-        self.func_name_highlighting = get_bool_opt(
-            options, 'func_name_highlighting', True)
-        self.disabled_modules = get_list_opt(options, 'disabled_modules', [])
-
-        self._functions = set()
-        if self.func_name_highlighting:
-            from pygments.lexers._lua_builtins import MODULES
-            for mod, func in MODULES.items():
-                if mod not in self.disabled_modules:
-                    self._functions.update(func)
-        RegexLexer.__init__(self, **options)
-
-    def get_tokens_unprocessed(self, text):
-        for index, token, value in \
-                RegexLexer.get_tokens_unprocessed(self, text):
-            if token is Name.Builtin and value not in self._functions:
-                if '.' in value:
-                    a, b = value.split('.')
-                    yield index, Name, a
-                    yield index + len(a), Punctuation, '.'
-                    yield index + len(a) + 1, Name, b
-                else:
-                    yield index, Name, value
-                continue
-            yield index, token, value
-
-def _luau_make_expression(should_pop, _s):
-    temp_list = [
-        (r'0[xX][\da-fA-F_]*', Number.Hex, '#pop'),
-        (r'0[bB][\d_]*', Number.Bin, '#pop'),
-        (r'\.?\d[\d_]*(?:\.[\d_]*)?(?:[eE][+-]?[\d_]+)?', Number.Float, '#pop'),
-
-        (words((
-            'true', 'false', 'nil'
-        ), suffix=r'\b'), Keyword.Constant, '#pop'),
-
-        (r'\[(=*)\[[.\n]*?\]\1\]', String, '#pop'),
-
-        (r'(\.)([a-zA-Z_]\w*)(?=%s*[({"\'])', bygroups(Punctuation, Name.Function), '#pop'),
-        (r'(\.)([a-zA-Z_]\w*)', bygroups(Punctuation, Name.Variable), '#pop'),
-
-        (rf'[a-zA-Z_]\w*(?:\.[a-zA-Z_]\w*)*(?={_s}*[({{"\'])', Name.Other, '#pop'),
-        (r'[a-zA-Z_]\w*(?:\.[a-zA-Z_]\w*)*', Name, '#pop'),
-    ]
-    if should_pop:
-        return temp_list
-    return [entry[:2] for entry in temp_list]
-
-def _luau_make_expression_special(should_pop):
-    temp_list = [
-        (r'\{', Punctuation, ('#pop', 'closing_brace_base', 'expression')),
-        (r'\(', Punctuation, ('#pop', 'closing_parenthesis_base', 'expression')),
-
-        (r'::?', Punctuation, ('#pop', 'type_end', 'type_start')),
-
-        (r"'", String.Single, ('#pop', 'string_single')),
-        (r'"', String.Double, ('#pop', 'string_double')),
-        (r'`', String.Backtick, ('#pop', 'string_interpolated')),
-    ]
-    if should_pop:
-        return temp_list
-    return [(entry[0], entry[1], entry[2][1:]) for entry in temp_list]
-
-class LuauLexer(RegexLexer):
-    """
-    For Luau source code.
-
-    Additional options accepted:
-
-    `include_luau_builtins`
-        If given and ``True``, automatically highlight Luau builtins
-        (default: ``True``).
-    `include_roblox_builtins`
-        If given and ``True``, automatically highlight Roblox-specific builtins
-        (default: ``False``).
-    `additional_builtins`
-        If given, must be a list of additional builtins to highlight.
-    `disabled_builtins`
-        If given, must be a list of builtins that will not be highlighted.
-    """
-
-    name = 'Luau'
-    url = 'https://luau-lang.org/'
-    aliases = ['luau']
-    filenames = ['*.luau']
-    version_added = '2.18'
-
-    _comment_multiline = r'(?:--\[(?P=*)\[[\w\W]*?\](?P=level)\])'
-    _comment_single = r'(?:--.*$)'
-    _s = r'(?:{}|{}|{})'.format(_comment_multiline, _comment_single, r'\s+')
-
-    tokens = {
-        'root': [
-            (r'#!.*', Comment.Hashbang, 'base'),
-            default('base'),
-        ],
-
-        'ws': [
-            (_comment_multiline, Comment.Multiline),
-            (_comment_single, Comment.Single),
-            (r'\s+', Whitespace),
-        ],
-
-        'base': [
-            include('ws'),
-
-            *_luau_make_expression_special(False),
-            (r'\.\.\.', Punctuation),
-
-            (rf'type\b(?={_s}+[a-zA-Z_])', Keyword.Reserved, 'type_declaration'),
-            (rf'export\b(?={_s}+[a-zA-Z_])', Keyword.Reserved),
-
-            (r'(?:\.\.|//|[+\-*\/%^<>=])=?', Operator, 'expression'),
-            (r'~=', Operator, 'expression'),
-
-            (words((
-                'and', 'or', 'not'
-            ), suffix=r'\b'), Operator.Word, 'expression'),
-
-            (words((
-                'elseif', 'for', 'if', 'in', 'repeat', 'return', 'until',
-                'while'), suffix=r'\b'), Keyword.Reserved, 'expression'),
-            (r'local\b', Keyword.Declaration, 'expression'),
-
-            (r'function\b', Keyword.Reserved, ('expression', 'func_name')),
-
-            (r'[\])};]+', Punctuation),
-
-            include('expression_static'),
-            *_luau_make_expression(False, _s),
-
-            (r'[\[.,]', Punctuation, 'expression'),
-        ],
-        'expression_static': [
-            (words((
-                'break', 'continue', 'do', 'else', 'elseif', 'end', 'for',
-                'if', 'in', 'repeat', 'return', 'then', 'until', 'while'),
-                suffix=r'\b'), Keyword.Reserved),
-        ],
-        'expression': [
-            include('ws'),
-
-            (r'if\b', Keyword.Reserved, ('ternary', 'expression')),
-
-            (r'local\b', Keyword.Declaration),
-            *_luau_make_expression_special(True),
-            (r'\.\.\.', Punctuation, '#pop'),
-
-            (r'function\b', Keyword.Reserved, 'func_name'),
-
-            include('expression_static'),
-            *_luau_make_expression(True, _s),
-
-            default('#pop'),
-        ],
-        'ternary': [
-            include('ws'),
-
-            (r'else\b', Keyword.Reserved, '#pop'),
-            (words((
-                'then', 'elseif',
-            ), suffix=r'\b'), Operator.Reserved, 'expression'),
-
-            default('#pop'),
-        ],
-
-        'closing_brace_pop': [
-            (r'\}', Punctuation, '#pop'),
-        ],
-        'closing_parenthesis_pop': [
-            (r'\)', Punctuation, '#pop'),
-        ],
-        'closing_gt_pop': [
-            (r'>', Punctuation, '#pop'),
-        ],
-
-        'closing_parenthesis_base': [
-            include('closing_parenthesis_pop'),
-            include('base'),
-        ],
-        'closing_parenthesis_type': [
-            include('closing_parenthesis_pop'),
-            include('type'),
-        ],
-        'closing_brace_base': [
-            include('closing_brace_pop'),
-            include('base'),
-        ],
-        'closing_brace_type': [
-            include('closing_brace_pop'),
-            include('type'),
-        ],
-        'closing_gt_type': [
-            include('closing_gt_pop'),
-            include('type'),
-        ],
-
-        'string_escape': [
-            (r'\\z\s*', String.Escape),
-            (r'\\(?:[abfnrtvz\\"\'`\{\n])|[\r\n]{1,2}|x[\da-fA-F]{2}|\d{1,3}|'
-             r'u\{\}[\da-fA-F]*\}', String.Escape),
-        ],
-        'string_single': [
-            include('string_escape'),
-
-            (r"'", String.Single, "#pop"),
-            (r"[^\\']+", String.Single),
-        ],
-        'string_double': [
-            include('string_escape'),
-
-            (r'"', String.Double, "#pop"),
-            (r'[^\\"]+', String.Double),
-        ],
-        'string_interpolated': [
-            include('string_escape'),
-
-            (r'\{', Punctuation, ('closing_brace_base', 'expression')),
-
-            (r'`', String.Backtick, "#pop"),
-            (r'[^\\`\{]+', String.Backtick),
-        ],
-
-        'func_name': [
-            include('ws'),
-
-            (r'[.:]', Punctuation),
-            (rf'[a-zA-Z_]\w*(?={_s}*[.:])', Name.Class),
-            (r'[a-zA-Z_]\w*', Name.Function),
-
-            (r'<', Punctuation, 'closing_gt_type'),
-
-            (r'\(', Punctuation, '#pop'),
-        ],
-
-        'type': [
-            include('ws'),
-
-            (r'\(', Punctuation, 'closing_parenthesis_type'),
-            (r'\{', Punctuation, 'closing_brace_type'),
-            (r'<', Punctuation, 'closing_gt_type'),
-
-            (r"'", String.Single, 'string_single'),
-            (r'"', String.Double, 'string_double'),
-
-            (r'[|&\.,\[\]:=]+', Punctuation),
-            (r'->', Punctuation),
-
-            (r'typeof\(', Name.Builtin, ('closing_parenthesis_base',
-                                         'expression')),
-            (r'[a-zA-Z_]\w*', Name.Class),
-        ],
-        'type_start': [
-            include('ws'),
-
-            (r'\(', Punctuation, ('#pop', 'closing_parenthesis_type')),
-            (r'\{', Punctuation, ('#pop', 'closing_brace_type')),
-            (r'<', Punctuation, ('#pop', 'closing_gt_type')),
-
-            (r"'", String.Single, ('#pop', 'string_single')),
-            (r'"', String.Double, ('#pop', 'string_double')),
-
-            (r'typeof\(', Name.Builtin, ('#pop', 'closing_parenthesis_base',
-                                         'expression')),
-            (r'[a-zA-Z_]\w*', Name.Class, '#pop'),
-        ],
-        'type_end': [
-            include('ws'),
-
-            (r'[|&\.]', Punctuation, 'type_start'),
-            (r'->', Punctuation, 'type_start'),
-
-            (r'<', Punctuation, 'closing_gt_type'),
-
-            default('#pop'),
-        ],
-        'type_declaration': [
-            include('ws'),
-
-            (r'[a-zA-Z_]\w*', Name.Class),
-            (r'<', Punctuation, 'closing_gt_type'),
-
-            (r'=', Punctuation, ('#pop', 'type_end', 'type_start')),
-        ],
-    }
-
-    def __init__(self, **options):
-        self.include_luau_builtins = get_bool_opt(
-            options, 'include_luau_builtins', True)
-        self.include_roblox_builtins = get_bool_opt(
-            options, 'include_roblox_builtins', False)
-        self.additional_builtins = get_list_opt(options, 'additional_builtins', [])
-        self.disabled_builtins = get_list_opt(options, 'disabled_builtins', [])
-
-        self._builtins = set(self.additional_builtins)
-        if self.include_luau_builtins:
-            from pygments.lexers._luau_builtins import LUAU_BUILTINS
-            self._builtins.update(LUAU_BUILTINS)
-        if self.include_roblox_builtins:
-            from pygments.lexers._luau_builtins import ROBLOX_BUILTINS
-            self._builtins.update(ROBLOX_BUILTINS)
-        if self.additional_builtins:
-            self._builtins.update(self.additional_builtins)
-        self._builtins.difference_update(self.disabled_builtins)
-
-        RegexLexer.__init__(self, **options)
-
-    def get_tokens_unprocessed(self, text):
-        for index, token, value in \
-                RegexLexer.get_tokens_unprocessed(self, text):
-            if token is Name or token is Name.Other:
-                split_value = value.split('.')
-                complete_value = []
-                new_index = index
-                for position in range(len(split_value), 0, -1):
-                    potential_string = '.'.join(split_value[:position])
-                    if potential_string in self._builtins:
-                        yield index, Name.Builtin, potential_string
-                        new_index += len(potential_string)
-
-                        if complete_value:
-                            yield new_index, Punctuation, '.'
-                            new_index += 1
-                        break
-                    complete_value.insert(0, split_value[position - 1])
-
-                for position, substring in enumerate(complete_value):
-                    if position + 1 == len(complete_value):
-                        if token is Name:
-                            yield new_index, Name.Variable, substring
-                            continue
-                        yield new_index, Name.Function, substring
-                        continue
-                    yield new_index, Name.Variable, substring
-                    new_index += len(substring)
-                    yield new_index, Punctuation, '.'
-                    new_index += 1
-
-                continue
-            yield index, token, value
-
-class MoonScriptLexer(LuaLexer):
-    """
-    For MoonScript source code.
-    """
-
-    name = 'MoonScript'
-    url = 'http://moonscript.org'
-    aliases = ['moonscript', 'moon']
-    filenames = ['*.moon']
-    mimetypes = ['text/x-moonscript', 'application/x-moonscript']
-    version_added = '1.5'
-
-    tokens = {
-        'root': [
-            (r'#!(.*?)$', Comment.Preproc),
-            default('base'),
-        ],
-        'base': [
-            ('--.*$', Comment.Single),
-            (r'(?i)(\d*\.\d+|\d+\.\d*)(e[+-]?\d+)?', Number.Float),
-            (r'(?i)\d+e[+-]?\d+', Number.Float),
-            (r'(?i)0x[0-9a-f]*', Number.Hex),
-            (r'\d+', Number.Integer),
-            (r'\n', Whitespace),
-            (r'[^\S\n]+', Text),
-            (r'(?s)\[(=*)\[.*?\]\1\]', String),
-            (r'(->|=>)', Name.Function),
-            (r':[a-zA-Z_]\w*', Name.Variable),
-            (r'(==|!=|~=|<=|>=|\.\.\.|\.\.|[=+\-*/%^<>#!.\\:])', Operator),
-            (r'[;,]', Punctuation),
-            (r'[\[\]{}()]', Keyword.Type),
-            (r'[a-zA-Z_]\w*:', Name.Variable),
-            (words((
-                'class', 'extends', 'if', 'then', 'super', 'do', 'with',
-                'import', 'export', 'while', 'elseif', 'return', 'for', 'in',
-                'from', 'when', 'using', 'else', 'and', 'or', 'not', 'switch',
-                'break'), suffix=r'\b'),
-             Keyword),
-            (r'(true|false|nil)\b', Keyword.Constant),
-            (r'(and|or|not)\b', Operator.Word),
-            (r'(self)\b', Name.Builtin.Pseudo),
-            (r'@@?([a-zA-Z_]\w*)?', Name.Variable.Class),
-            (r'[A-Z]\w*', Name.Class),  # proper name
-            (words(all_lua_builtins(), suffix=r"\b"), Name.Builtin),
-            (r'[A-Za-z_]\w*', Name),
-            ("'", String.Single, combined('stringescape', 'sqs')),
-            ('"', String.Double, combined('stringescape', 'dqs'))
-        ],
-        'stringescape': [
-            (r'''\\([abfnrtv\\"']|\d{1,3})''', String.Escape)
-        ],
-        'sqs': [
-            ("'", String.Single, '#pop'),
-            ("[^']+", String)
-        ],
-        'dqs': [
-            ('"', String.Double, '#pop'),
-            ('[^"]+', String)
-        ]
-    }
-
-    def get_tokens_unprocessed(self, text):
-        # set . as Operator instead of Punctuation
-        for index, token, value in LuaLexer.get_tokens_unprocessed(self, text):
-            if token == Punctuation and value == ".":
-                token = Operator
-            yield index, token, value
-
-
-class ChaiscriptLexer(RegexLexer):
-    """
-    For ChaiScript source code.
-    """
-
-    name = 'ChaiScript'
-    url = 'http://chaiscript.com/'
-    aliases = ['chaiscript', 'chai']
-    filenames = ['*.chai']
-    mimetypes = ['text/x-chaiscript', 'application/x-chaiscript']
-    version_added = '2.0'
-
-    flags = re.DOTALL | re.MULTILINE
-
-    tokens = {
-        'commentsandwhitespace': [
-            (r'\s+', Text),
-            (r'//.*?\n', Comment.Single),
-            (r'/\*.*?\*/', Comment.Multiline),
-            (r'^\#.*?\n', Comment.Single)
-        ],
-        'slashstartsregex': [
-            include('commentsandwhitespace'),
-            (r'/(\\.|[^[/\\\n]|\[(\\.|[^\]\\\n])*])+/'
-             r'([gim]+\b|\B)', String.Regex, '#pop'),
-            (r'(?=/)', Text, ('#pop', 'badregex')),
-            default('#pop')
-        ],
-        'badregex': [
-            (r'\n', Text, '#pop')
-        ],
-        'root': [
-            include('commentsandwhitespace'),
-            (r'\n', Text),
-            (r'[^\S\n]+', Text),
-            (r'\+\+|--|~|&&|\?|:|\|\||\\(?=\n)|\.\.'
-             r'(<<|>>>?|==?|!=?|[-<>+*%&|^/])=?', Operator, 'slashstartsregex'),
-            (r'[{(\[;,]', Punctuation, 'slashstartsregex'),
-            (r'[})\].]', Punctuation),
-            (r'[=+\-*/]', Operator),
-            (r'(for|in|while|do|break|return|continue|if|else|'
-             r'throw|try|catch'
-             r')\b', Keyword, 'slashstartsregex'),
-            (r'(var)\b', Keyword.Declaration, 'slashstartsregex'),
-            (r'(attr|def|fun)\b', Keyword.Reserved),
-            (r'(true|false)\b', Keyword.Constant),
-            (r'(eval|throw)\b', Name.Builtin),
-            (r'`\S+`', Name.Builtin),
-            (r'[$a-zA-Z_]\w*', Name.Other),
-            (r'[0-9][0-9]*\.[0-9]+([eE][0-9]+)?[fd]?', Number.Float),
-            (r'0x[0-9a-fA-F]+', Number.Hex),
-            (r'[0-9]+', Number.Integer),
-            (r'"', String.Double, 'dqstring'),
-            (r"'(\\\\|\\[^\\]|[^'\\])*'", String.Single),
-        ],
-        'dqstring': [
-            (r'\$\{[^"}]+?\}', String.Interpol),
-            (r'\$', String.Double),
-            (r'\\\\', String.Double),
-            (r'\\"', String.Double),
-            (r'[^\\"$]+', String.Double),
-            (r'"', String.Double, '#pop'),
-        ],
-    }
-
-
-class LSLLexer(RegexLexer):
-    """
-    For Second Life's Linden Scripting Language source code.
-    """
-
-    name = 'LSL'
-    aliases = ['lsl']
-    filenames = ['*.lsl']
-    mimetypes = ['text/x-lsl']
-    url = 'https://wiki.secondlife.com/wiki/Linden_Scripting_Language'
-    version_added = '2.0'
-
-    flags = re.MULTILINE
-
-    lsl_keywords = r'\b(?:do|else|for|if|jump|return|while)\b'
-    lsl_types = r'\b(?:float|integer|key|list|quaternion|rotation|string|vector)\b'
-    lsl_states = r'\b(?:(?:state)\s+\w+|default)\b'
-    lsl_events = r'\b(?:state_(?:entry|exit)|touch(?:_(?:start|end))?|(?:land_)?collision(?:_(?:start|end))?|timer|listen|(?:no_)?sensor|control|(?:not_)?at_(?:rot_)?target|money|email|run_time_permissions|changed|attach|dataserver|moving_(?:start|end)|link_message|(?:on|object)_rez|remote_data|http_re(?:sponse|quest)|path_update|transaction_result)\b'
-    lsl_functions_builtin = r'\b(?:ll(?:ReturnObjectsBy(?:ID|Owner)|Json(?:2List|[GS]etValue|ValueType)|Sin|Cos|Tan|Atan2|Sqrt|Pow|Abs|Fabs|Frand|Floor|Ceil|Round|Vec(?:Mag|Norm|Dist)|Rot(?:Between|2(?:Euler|Fwd|Left|Up))|(?:Euler|Axes)2Rot|Whisper|(?:Region|Owner)?Say|Shout|Listen(?:Control|Remove)?|Sensor(?:Repeat|Remove)?|Detected(?:Name|Key|Owner|Type|Pos|Vel|Grab|Rot|Group|LinkNumber)|Die|Ground|Wind|(?:[GS]et)(?:AnimationOverride|MemoryLimit|PrimMediaParams|ParcelMusicURL|Object(?:Desc|Name)|PhysicsMaterial|Status|Scale|Color|Alpha|Texture|Pos|Rot|Force|Torque)|ResetAnimationOverride|(?:Scale|Offset|Rotate)Texture|(?:Rot)?Target(?:Remove)?|(?:Stop)?MoveToTarget|Apply(?:Rotational)?Impulse|Set(?:KeyframedMotion|ContentType|RegionPos|(?:Angular)?Velocity|Buoyancy|HoverHeight|ForceAndTorque|TimerEvent|ScriptState|Damage|TextureAnim|Sound(?:Queueing|Radius)|Vehicle(?:Type|(?:Float|Vector|Rotation)Param)|(?:Touch|Sit)?Text|Camera(?:Eye|At)Offset|PrimitiveParams|ClickAction|Link(?:Alpha|Color|PrimitiveParams(?:Fast)?|Texture(?:Anim)?|Camera|Media)|RemoteScriptAccessPin|PayPrice|LocalRot)|ScaleByFactor|Get(?:(?:Max|Min)ScaleFactor|ClosestNavPoint|StaticPath|SimStats|Env|PrimitiveParams|Link(?:PrimitiveParams|Number(?:OfSides)?|Key|Name|Media)|HTTPHeader|FreeURLs|Object(?:Details|PermMask|PrimCount)|Parcel(?:MaxPrims|Details|Prim(?:Count|Owners))|Attached|(?:SPMax|Free|Used)Memory|Region(?:Name|TimeDilation|FPS|Corner|AgentCount)|Root(?:Position|Rotation)|UnixTime|(?:Parcel|Region)Flags|(?:Wall|GMT)clock|SimulatorHostname|BoundingBox|GeometricCenter|Creator|NumberOf(?:Prims|NotecardLines|Sides)|Animation(?:List)?|(?:Camera|Local)(?:Pos|Rot)|Vel|Accel|Omega|Time(?:stamp|OfDay)|(?:Object|CenterOf)?Mass|MassMKS|Energy|Owner|(?:Owner)?Key|SunDirection|Texture(?:Offset|Scale|Rot)|Inventory(?:Number|Name|Key|Type|Creator|PermMask)|Permissions(?:Key)?|StartParameter|List(?:Length|EntryType)|Date|Agent(?:Size|Info|Language|List)|LandOwnerAt|NotecardLine|Script(?:Name|State))|(?:Get|Reset|GetAndReset)Time|PlaySound(?:Slave)?|LoopSound(?:Master|Slave)?|(?:Trigger|Stop|Preload)Sound|(?:(?:Get|Delete)Sub|Insert)String|To(?:Upper|Lower)|Give(?:InventoryList|Money)|RezObject|(?:Stop)?LookAt|Sleep|CollisionFilter|(?:Take|Release)Controls|DetachFromAvatar|AttachToAvatar(?:Temp)?|InstantMessage|(?:GetNext)?Email|StopHover|MinEventDelay|RotLookAt|String(?:Length|Trim)|(?:Start|Stop)Animation|TargetOmega|RequestPermissions|(?:Create|Break)Link|BreakAllLinks|(?:Give|Remove)Inventory|Water|PassTouches|Request(?:Agent|Inventory)Data|TeleportAgent(?:Home|GlobalCoords)?|ModifyLand|CollisionSound|ResetScript|MessageLinked|PushObject|PassCollisions|AxisAngle2Rot|Rot2(?:Axis|Angle)|A(?:cos|sin)|AngleBetween|AllowInventoryDrop|SubStringIndex|List2(?:CSV|Integer|Json|Float|String|Key|Vector|Rot|List(?:Strided)?)|DeleteSubList|List(?:Statistics|Sort|Randomize|(?:Insert|Find|Replace)List)|EdgeOfWorld|AdjustSoundVolume|Key2Name|TriggerSoundLimited|EjectFromLand|(?:CSV|ParseString)2List|OverMyLand|SameGroup|UnSit|Ground(?:Slope|Normal|Contour)|GroundRepel|(?:Set|Remove)VehicleFlags|(?:AvatarOn)?(?:Link)?SitTarget|Script(?:Danger|Profiler)|Dialog|VolumeDetect|ResetOtherScript|RemoteLoadScriptPin|(?:Open|Close)RemoteDataChannel|SendRemoteData|RemoteDataReply|(?:Integer|String)ToBase64|XorBase64|Log(?:10)?|Base64To(?:String|Integer)|ParseStringKeepNulls|RezAtRoot|RequestSimulatorData|ForceMouselook|(?:Load|Release|(?:E|Une)scape)URL|ParcelMedia(?:CommandList|Query)|ModPow|MapDestination|(?:RemoveFrom|AddTo|Reset)Land(?:Pass|Ban)List|(?:Set|Clear)CameraParams|HTTP(?:Request|Response)|TextBox|DetectedTouch(?:UV|Face|Pos|(?:N|Bin)ormal|ST)|(?:MD5|SHA1|DumpList2)String|Request(?:Secure)?URL|Clear(?:Prim|Link)Media|(?:Link)?ParticleSystem|(?:Get|Request)(?:Username|DisplayName)|RegionSayTo|CastRay|GenerateKey|TransferLindenDollars|ManageEstateAccess|(?:Create|Delete)Character|ExecCharacterCmd|Evade|FleeFrom|NavigateTo|PatrolPoints|Pursue|UpdateCharacter|WanderWithin))\b'
-    lsl_constants_float = r'\b(?:DEG_TO_RAD|PI(?:_BY_TWO)?|RAD_TO_DEG|SQRT2|TWO_PI)\b'
-    lsl_constants_integer = r'\b(?:JSON_APPEND|STATUS_(?:PHYSICS|ROTATE_[XYZ]|PHANTOM|SANDBOX|BLOCK_GRAB(?:_OBJECT)?|(?:DIE|RETURN)_AT_EDGE|CAST_SHADOWS|OK|MALFORMED_PARAMS|TYPE_MISMATCH|BOUNDS_ERROR|NOT_(?:FOUND|SUPPORTED)|INTERNAL_ERROR|WHITELIST_FAILED)|AGENT(?:_(?:BY_(?:LEGACY_|USER)NAME|FLYING|ATTACHMENTS|SCRIPTED|MOUSELOOK|SITTING|ON_OBJECT|AWAY|WALKING|IN_AIR|TYPING|CROUCHING|BUSY|ALWAYS_RUN|AUTOPILOT|LIST_(?:PARCEL(?:_OWNER)?|REGION)))?|CAMERA_(?:PITCH|DISTANCE|BEHINDNESS_(?:ANGLE|LAG)|(?:FOCUS|POSITION)(?:_(?:THRESHOLD|LOCKED|LAG))?|FOCUS_OFFSET|ACTIVE)|ANIM_ON|LOOP|REVERSE|PING_PONG|SMOOTH|ROTATE|SCALE|ALL_SIDES|LINK_(?:ROOT|SET|ALL_(?:OTHERS|CHILDREN)|THIS)|ACTIVE|PASSIVE|SCRIPTED|CONTROL_(?:FWD|BACK|(?:ROT_)?(?:LEFT|RIGHT)|UP|DOWN|(?:ML_)?LBUTTON)|PERMISSION_(?:RETURN_OBJECTS|DEBIT|OVERRIDE_ANIMATIONS|SILENT_ESTATE_MANAGEMENT|TAKE_CONTROLS|TRIGGER_ANIMATION|ATTACH|CHANGE_LINKS|(?:CONTROL|TRACK)_CAMERA|TELEPORT)|INVENTORY_(?:TEXTURE|SOUND|OBJECT|SCRIPT|LANDMARK|CLOTHING|NOTECARD|BODYPART|ANIMATION|GESTURE|ALL|NONE)|CHANGED_(?:INVENTORY|COLOR|SHAPE|SCALE|TEXTURE|LINK|ALLOWED_DROP|OWNER|REGION(?:_START)?|TELEPORT|MEDIA)|OBJECT_(?:(?:PHYSICS|SERVER|STREAMING)_COST|UNKNOWN_DETAIL|CHARACTER_TIME|PHANTOM|PHYSICS|TEMP_ON_REZ|NAME|DESC|POS|PRIM_EQUIVALENCE|RETURN_(?:PARCEL(?:_OWNER)?|REGION)|ROO?T|VELOCITY|OWNER|GROUP|CREATOR|ATTACHED_POINT|RENDER_WEIGHT|PATHFINDING_TYPE|(?:RUNNING|TOTAL)_SCRIPT_COUNT|SCRIPT_(?:MEMORY|TIME))|TYPE_(?:INTEGER|FLOAT|STRING|KEY|VECTOR|ROTATION|INVALID)|(?:DEBUG|PUBLIC)_CHANNEL|ATTACH_(?:AVATAR_CENTER|CHEST|HEAD|BACK|PELVIS|MOUTH|CHIN|NECK|NOSE|BELLY|[LR](?:SHOULDER|HAND|FOOT|EAR|EYE|[UL](?:ARM|LEG)|HIP)|(?:LEFT|RIGHT)_PEC|HUD_(?:CENTER_[12]|TOP_(?:RIGHT|CENTER|LEFT)|BOTTOM(?:_(?:RIGHT|LEFT))?))|LAND_(?:LEVEL|RAISE|LOWER|SMOOTH|NOISE|REVERT)|DATA_(?:ONLINE|NAME|BORN|SIM_(?:POS|STATUS|RATING)|PAYINFO)|PAYMENT_INFO_(?:ON_FILE|USED)|REMOTE_DATA_(?:CHANNEL|REQUEST|REPLY)|PSYS_(?:PART_(?:BF_(?:ZERO|ONE(?:_MINUS_(?:DEST_COLOR|SOURCE_(ALPHA|COLOR)))?|DEST_COLOR|SOURCE_(ALPHA|COLOR))|BLEND_FUNC_(DEST|SOURCE)|FLAGS|(?:START|END)_(?:COLOR|ALPHA|SCALE|GLOW)|MAX_AGE|(?:RIBBON|WIND|INTERP_(?:COLOR|SCALE)|BOUNCE|FOLLOW_(?:SRC|VELOCITY)|TARGET_(?:POS|LINEAR)|EMISSIVE)_MASK)|SRC_(?:MAX_AGE|PATTERN|ANGLE_(?:BEGIN|END)|BURST_(?:RATE|PART_COUNT|RADIUS|SPEED_(?:MIN|MAX))|ACCEL|TEXTURE|TARGET_KEY|OMEGA|PATTERN_(?:DROP|EXPLODE|ANGLE(?:_CONE(?:_EMPTY)?)?)))|VEHICLE_(?:REFERENCE_FRAME|TYPE_(?:NONE|SLED|CAR|BOAT|AIRPLANE|BALLOON)|(?:LINEAR|ANGULAR)_(?:FRICTION_TIMESCALE|MOTOR_DIRECTION)|LINEAR_MOTOR_OFFSET|HOVER_(?:HEIGHT|EFFICIENCY|TIMESCALE)|BUOYANCY|(?:LINEAR|ANGULAR)_(?:DEFLECTION_(?:EFFICIENCY|TIMESCALE)|MOTOR_(?:DECAY_)?TIMESCALE)|VERTICAL_ATTRACTION_(?:EFFICIENCY|TIMESCALE)|BANKING_(?:EFFICIENCY|MIX|TIMESCALE)|FLAG_(?:NO_DEFLECTION_UP|LIMIT_(?:ROLL_ONLY|MOTOR_UP)|HOVER_(?:(?:WATER|TERRAIN|UP)_ONLY|GLOBAL_HEIGHT)|MOUSELOOK_(?:STEER|BANK)|CAMERA_DECOUPLED))|PRIM_(?:TYPE(?:_(?:BOX|CYLINDER|PRISM|SPHERE|TORUS|TUBE|RING|SCULPT))?|HOLE_(?:DEFAULT|CIRCLE|SQUARE|TRIANGLE)|MATERIAL(?:_(?:STONE|METAL|GLASS|WOOD|FLESH|PLASTIC|RUBBER))?|SHINY_(?:NONE|LOW|MEDIUM|HIGH)|BUMP_(?:NONE|BRIGHT|DARK|WOOD|BARK|BRICKS|CHECKER|CONCRETE|TILE|STONE|DISKS|GRAVEL|BLOBS|SIDING|LARGETILE|STUCCO|SUCTION|WEAVE)|TEXGEN_(?:DEFAULT|PLANAR)|SCULPT_(?:TYPE_(?:SPHERE|TORUS|PLANE|CYLINDER|MASK)|FLAG_(?:MIRROR|INVERT))|PHYSICS(?:_(?:SHAPE_(?:CONVEX|NONE|PRIM|TYPE)))?|(?:POS|ROT)_LOCAL|SLICE|TEXT|FLEXIBLE|POINT_LIGHT|TEMP_ON_REZ|PHANTOM|POSITION|SIZE|ROTATION|TEXTURE|NAME|OMEGA|DESC|LINK_TARGET|COLOR|BUMP_SHINY|FULLBRIGHT|TEXGEN|GLOW|MEDIA_(?:ALT_IMAGE_ENABLE|CONTROLS|(?:CURRENT|HOME)_URL|AUTO_(?:LOOP|PLAY|SCALE|ZOOM)|FIRST_CLICK_INTERACT|(?:WIDTH|HEIGHT)_PIXELS|WHITELIST(?:_ENABLE)?|PERMS_(?:INTERACT|CONTROL)|PARAM_MAX|CONTROLS_(?:STANDARD|MINI)|PERM_(?:NONE|OWNER|GROUP|ANYONE)|MAX_(?:URL_LENGTH|WHITELIST_(?:SIZE|COUNT)|(?:WIDTH|HEIGHT)_PIXELS)))|MASK_(?:BASE|OWNER|GROUP|EVERYONE|NEXT)|PERM_(?:TRANSFER|MODIFY|COPY|MOVE|ALL)|PARCEL_(?:MEDIA_COMMAND_(?:STOP|PAUSE|PLAY|LOOP|TEXTURE|URL|TIME|AGENT|UNLOAD|AUTO_ALIGN|TYPE|SIZE|DESC|LOOP_SET)|FLAG_(?:ALLOW_(?:FLY|(?:GROUP_)?SCRIPTS|LANDMARK|TERRAFORM|DAMAGE|CREATE_(?:GROUP_)?OBJECTS)|USE_(?:ACCESS_(?:GROUP|LIST)|BAN_LIST|LAND_PASS_LIST)|LOCAL_SOUND_ONLY|RESTRICT_PUSHOBJECT|ALLOW_(?:GROUP|ALL)_OBJECT_ENTRY)|COUNT_(?:TOTAL|OWNER|GROUP|OTHER|SELECTED|TEMP)|DETAILS_(?:NAME|DESC|OWNER|GROUP|AREA|ID|SEE_AVATARS))|LIST_STAT_(?:MAX|MIN|MEAN|MEDIAN|STD_DEV|SUM(?:_SQUARES)?|NUM_COUNT|GEOMETRIC_MEAN|RANGE)|PAY_(?:HIDE|DEFAULT)|REGION_FLAG_(?:ALLOW_DAMAGE|FIXED_SUN|BLOCK_TERRAFORM|SANDBOX|DISABLE_(?:COLLISIONS|PHYSICS)|BLOCK_FLY|ALLOW_DIRECT_TELEPORT|RESTRICT_PUSHOBJECT)|HTTP_(?:METHOD|MIMETYPE|BODY_(?:MAXLENGTH|TRUNCATED)|CUSTOM_HEADER|PRAGMA_NO_CACHE|VERBOSE_THROTTLE|VERIFY_CERT)|STRING_(?:TRIM(?:_(?:HEAD|TAIL))?)|CLICK_ACTION_(?:NONE|TOUCH|SIT|BUY|PAY|OPEN(?:_MEDIA)?|PLAY|ZOOM)|TOUCH_INVALID_FACE|PROFILE_(?:NONE|SCRIPT_MEMORY)|RC_(?:DATA_FLAGS|DETECT_PHANTOM|GET_(?:LINK_NUM|NORMAL|ROOT_KEY)|MAX_HITS|REJECT_(?:TYPES|AGENTS|(?:NON)?PHYSICAL|LAND))|RCERR_(?:CAST_TIME_EXCEEDED|SIM_PERF_LOW|UNKNOWN)|ESTATE_ACCESS_(?:ALLOWED_(?:AGENT|GROUP)_(?:ADD|REMOVE)|BANNED_AGENT_(?:ADD|REMOVE))|DENSITY|FRICTION|RESTITUTION|GRAVITY_MULTIPLIER|KFM_(?:COMMAND|CMD_(?:PLAY|STOP|PAUSE|SET_MODE)|MODE|FORWARD|LOOP|PING_PONG|REVERSE|DATA|ROTATION|TRANSLATION)|ERR_(?:GENERIC|PARCEL_PERMISSIONS|MALFORMED_PARAMS|RUNTIME_PERMISSIONS|THROTTLED)|CHARACTER_(?:CMD_(?:(?:SMOOTH_)?STOP|JUMP)|DESIRED_(?:TURN_)?SPEED|RADIUS|STAY_WITHIN_PARCEL|LENGTH|ORIENTATION|ACCOUNT_FOR_SKIPPED_FRAMES|AVOIDANCE_MODE|TYPE(?:_(?:[A-D]|NONE))?|MAX_(?:DECEL|TURN_RADIUS|(?:ACCEL|SPEED)))|PURSUIT_(?:OFFSET|FUZZ_FACTOR|GOAL_TOLERANCE|INTERCEPT)|REQUIRE_LINE_OF_SIGHT|FORCE_DIRECT_PATH|VERTICAL|HORIZONTAL|AVOID_(?:CHARACTERS|DYNAMIC_OBSTACLES|NONE)|PU_(?:EVADE_(?:HIDDEN|SPOTTED)|FAILURE_(?:DYNAMIC_PATHFINDING_DISABLED|INVALID_(?:GOAL|START)|NO_(?:NAVMESH|VALID_DESTINATION)|OTHER|TARGET_GONE|(?:PARCEL_)?UNREACHABLE)|(?:GOAL|SLOWDOWN_DISTANCE)_REACHED)|TRAVERSAL_TYPE(?:_(?:FAST|NONE|SLOW))?|CONTENT_TYPE_(?:ATOM|FORM|HTML|JSON|LLSD|RSS|TEXT|XHTML|XML)|GCNP_(?:RADIUS|STATIC)|(?:PATROL|WANDER)_PAUSE_AT_WAYPOINTS|OPT_(?:AVATAR|CHARACTER|EXCLUSION_VOLUME|LEGACY_LINKSET|MATERIAL_VOLUME|OTHER|STATIC_OBSTACLE|WALKABLE)|SIM_STAT_PCT_CHARS_STEPPED)\b'
-    lsl_constants_integer_boolean = r'\b(?:FALSE|TRUE)\b'
-    lsl_constants_rotation = r'\b(?:ZERO_ROTATION)\b'
-    lsl_constants_string = r'\b(?:EOF|JSON_(?:ARRAY|DELETE|FALSE|INVALID|NULL|NUMBER|OBJECT|STRING|TRUE)|NULL_KEY|TEXTURE_(?:BLANK|DEFAULT|MEDIA|PLYWOOD|TRANSPARENT)|URL_REQUEST_(?:GRANTED|DENIED))\b'
-    lsl_constants_vector = r'\b(?:TOUCH_INVALID_(?:TEXCOORD|VECTOR)|ZERO_VECTOR)\b'
-    lsl_invalid_broken = r'\b(?:LAND_(?:LARGE|MEDIUM|SMALL)_BRUSH)\b'
-    lsl_invalid_deprecated = r'\b(?:ATTACH_[LR]PEC|DATA_RATING|OBJECT_ATTACHMENT_(?:GEOMETRY_BYTES|SURFACE_AREA)|PRIM_(?:CAST_SHADOWS|MATERIAL_LIGHT|TYPE_LEGACY)|PSYS_SRC_(?:INNER|OUTER)ANGLE|VEHICLE_FLAG_NO_FLY_UP|ll(?:Cloud|Make(?:Explosion|Fountain|Smoke|Fire)|RemoteDataSetRegion|Sound(?:Preload)?|XorBase64Strings(?:Correct)?))\b'
-    lsl_invalid_illegal = r'\b(?:event)\b'
-    lsl_invalid_unimplemented = r'\b(?:CHARACTER_(?:MAX_ANGULAR_(?:ACCEL|SPEED)|TURN_SPEED_MULTIPLIER)|PERMISSION_(?:CHANGE_(?:JOINTS|PERMISSIONS)|RELEASE_OWNERSHIP|REMAP_CONTROLS)|PRIM_PHYSICS_MATERIAL|PSYS_SRC_OBJ_REL_MASK|ll(?:CollisionSprite|(?:Stop)?PointAt|(?:(?:Refresh|Set)Prim)URL|(?:Take|Release)Camera|RemoteLoadScript))\b'
-    lsl_reserved_godmode = r'\b(?:ll(?:GodLikeRezObject|Set(?:Inventory|Object)PermMask))\b'
-    lsl_reserved_log = r'\b(?:print)\b'
-    lsl_operators = r'\+\+|\-\-|<<|>>|&&?|\|\|?|\^|~|[!%<>=*+\-/]=?'
-
-    tokens = {
-        'root':
-        [
-            (r'//.*?\n',                          Comment.Single),
-            (r'/\*',                              Comment.Multiline, 'comment'),
-            (r'"',                                String.Double, 'string'),
-            (lsl_keywords,                        Keyword),
-            (lsl_types,                           Keyword.Type),
-            (lsl_states,                          Name.Class),
-            (lsl_events,                          Name.Builtin),
-            (lsl_functions_builtin,               Name.Function),
-            (lsl_constants_float,                 Keyword.Constant),
-            (lsl_constants_integer,               Keyword.Constant),
-            (lsl_constants_integer_boolean,       Keyword.Constant),
-            (lsl_constants_rotation,              Keyword.Constant),
-            (lsl_constants_string,                Keyword.Constant),
-            (lsl_constants_vector,                Keyword.Constant),
-            (lsl_invalid_broken,                  Error),
-            (lsl_invalid_deprecated,              Error),
-            (lsl_invalid_illegal,                 Error),
-            (lsl_invalid_unimplemented,           Error),
-            (lsl_reserved_godmode,                Keyword.Reserved),
-            (lsl_reserved_log,                    Keyword.Reserved),
-            (r'\b([a-zA-Z_]\w*)\b',     Name.Variable),
-            (r'(\d+\.\d*|\.\d+|\d+)[eE][+-]?\d*', Number.Float),
-            (r'(\d+\.\d*|\.\d+)',                 Number.Float),
-            (r'0[xX][0-9a-fA-F]+',                Number.Hex),
-            (r'\d+',                              Number.Integer),
-            (lsl_operators,                       Operator),
-            (r':=?',                              Error),
-            (r'[,;{}()\[\]]',                     Punctuation),
-            (r'\n+',                              Whitespace),
-            (r'\s+',                              Whitespace)
-        ],
-        'comment':
-        [
-            (r'[^*/]+',                           Comment.Multiline),
-            (r'/\*',                              Comment.Multiline, '#push'),
-            (r'\*/',                              Comment.Multiline, '#pop'),
-            (r'[*/]',                             Comment.Multiline)
-        ],
-        'string':
-        [
-            (r'\\([nt"\\])',                      String.Escape),
-            (r'"',                                String.Double, '#pop'),
-            (r'\\.',                              Error),
-            (r'[^"\\]+',                          String.Double),
-        ]
-    }
-
-
-class AppleScriptLexer(RegexLexer):
-    """
-    For AppleScript source code,
-    including `AppleScript Studio
-    `_.
-    Contributed by Andreas Amann .
-    """
-
-    name = 'AppleScript'
-    url = 'https://developer.apple.com/library/archive/documentation/AppleScript/Conceptual/AppleScriptLangGuide/introduction/ASLR_intro.html'
-    aliases = ['applescript']
-    filenames = ['*.applescript']
-    version_added = '1.0'
-
-    flags = re.MULTILINE | re.DOTALL
-
-    Identifiers = r'[a-zA-Z]\w*'
-
-    # XXX: use words() for all of these
-    Literals = ('AppleScript', 'current application', 'false', 'linefeed',
-                'missing value', 'pi', 'quote', 'result', 'return', 'space',
-                'tab', 'text item delimiters', 'true', 'version')
-    Classes = ('alias ', 'application ', 'boolean ', 'class ', 'constant ',
-               'date ', 'file ', 'integer ', 'list ', 'number ', 'POSIX file ',
-               'real ', 'record ', 'reference ', 'RGB color ', 'script ',
-               'text ', 'unit types', '(?:Unicode )?text', 'string')
-    BuiltIn = ('attachment', 'attribute run', 'character', 'day', 'month',
-               'paragraph', 'word', 'year')
-    HandlerParams = ('about', 'above', 'against', 'apart from', 'around',
-                     'aside from', 'at', 'below', 'beneath', 'beside',
-                     'between', 'for', 'given', 'instead of', 'on', 'onto',
-                     'out of', 'over', 'since')
-    Commands = ('ASCII (character|number)', 'activate', 'beep', 'choose URL',
-                'choose application', 'choose color', 'choose file( name)?',
-                'choose folder', 'choose from list',
-                'choose remote application', 'clipboard info',
-                'close( access)?', 'copy', 'count', 'current date', 'delay',
-                'delete', 'display (alert|dialog)', 'do shell script',
-                'duplicate', 'exists', 'get eof', 'get volume settings',
-                'info for', 'launch', 'list (disks|folder)', 'load script',
-                'log', 'make', 'mount volume', 'new', 'offset',
-                'open( (for access|location))?', 'path to', 'print', 'quit',
-                'random number', 'read', 'round', 'run( script)?',
-                'say', 'scripting components',
-                'set (eof|the clipboard to|volume)', 'store script',
-                'summarize', 'system attribute', 'system info',
-                'the clipboard', 'time to GMT', 'write', 'quoted form')
-    References = ('(in )?back of', '(in )?front of', '[0-9]+(st|nd|rd|th)',
-                  'first', 'second', 'third', 'fourth', 'fifth', 'sixth',
-                  'seventh', 'eighth', 'ninth', 'tenth', 'after', 'back',
-                  'before', 'behind', 'every', 'front', 'index', 'last',
-                  'middle', 'some', 'that', 'through', 'thru', 'where', 'whose')
-    Operators = ("and", "or", "is equal", "equals", "(is )?equal to", "is not",
-                 "isn't", "isn't equal( to)?", "is not equal( to)?",
-                 "doesn't equal", "does not equal", "(is )?greater than",
-                 "comes after", "is not less than or equal( to)?",
-                 "isn't less than or equal( to)?", "(is )?less than",
-                 "comes before", "is not greater than or equal( to)?",
-                 "isn't greater than or equal( to)?",
-                 "(is  )?greater than or equal( to)?", "is not less than",
-                 "isn't less than", "does not come before",
-                 "doesn't come before", "(is )?less than or equal( to)?",
-                 "is not greater than", "isn't greater than",
-                 "does not come after", "doesn't come after", "starts? with",
-                 "begins? with", "ends? with", "contains?", "does not contain",
-                 "doesn't contain", "is in", "is contained by", "is not in",
-                 "is not contained by", "isn't contained by", "div", "mod",
-                 "not", "(a  )?(ref( to)?|reference to)", "is", "does")
-    Control = ('considering', 'else', 'error', 'exit', 'from', 'if',
-               'ignoring', 'in', 'repeat', 'tell', 'then', 'times', 'to',
-               'try', 'until', 'using terms from', 'while', 'whith',
-               'with timeout( of)?', 'with transaction', 'by', 'continue',
-               'end', 'its?', 'me', 'my', 'return', 'of', 'as')
-    Declarations = ('global', 'local', 'prop(erty)?', 'set', 'get')
-    Reserved = ('but', 'put', 'returning', 'the')
-    StudioClasses = ('action cell', 'alert reply', 'application', 'box',
-                     'browser( cell)?', 'bundle', 'button( cell)?', 'cell',
-                     'clip view', 'color well', 'color-panel',
-                     'combo box( item)?', 'control',
-                     'data( (cell|column|item|row|source))?', 'default entry',
-                     'dialog reply', 'document', 'drag info', 'drawer',
-                     'event', 'font(-panel)?', 'formatter',
-                     'image( (cell|view))?', 'matrix', 'menu( item)?', 'item',
-                     'movie( view)?', 'open-panel', 'outline view', 'panel',
-                     'pasteboard', 'plugin', 'popup button',
-                     'progress indicator', 'responder', 'save-panel',
-                     'scroll view', 'secure text field( cell)?', 'slider',
-                     'sound', 'split view', 'stepper', 'tab view( item)?',
-                     'table( (column|header cell|header view|view))',
-                     'text( (field( cell)?|view))?', 'toolbar( item)?',
-                     'user-defaults', 'view', 'window')
-    StudioEvents = ('accept outline drop', 'accept table drop', 'action',
-                    'activated', 'alert ended', 'awake from nib', 'became key',
-                    'became main', 'begin editing', 'bounds changed',
-                    'cell value', 'cell value changed', 'change cell value',
-                    'change item value', 'changed', 'child of item',
-                    'choose menu item', 'clicked', 'clicked toolbar item',
-                    'closed', 'column clicked', 'column moved',
-                    'column resized', 'conclude drop', 'data representation',
-                    'deminiaturized', 'dialog ended', 'document nib name',
-                    'double clicked', 'drag( (entered|exited|updated))?',
-                    'drop', 'end editing', 'exposed', 'idle', 'item expandable',
-                    'item value', 'item value changed', 'items changed',
-                    'keyboard down', 'keyboard up', 'launched',
-                    'load data representation', 'miniaturized', 'mouse down',
-                    'mouse dragged', 'mouse entered', 'mouse exited',
-                    'mouse moved', 'mouse up', 'moved',
-                    'number of browser rows', 'number of items',
-                    'number of rows', 'open untitled', 'opened', 'panel ended',
-                    'parameters updated', 'plugin loaded', 'prepare drop',
-                    'prepare outline drag', 'prepare outline drop',
-                    'prepare table drag', 'prepare table drop',
-                    'read from file', 'resigned active', 'resigned key',
-                    'resigned main', 'resized( sub views)?',
-                    'right mouse down', 'right mouse dragged',
-                    'right mouse up', 'rows changed', 'scroll wheel',
-                    'selected tab view item', 'selection changed',
-                    'selection changing', 'should begin editing',
-                    'should close', 'should collapse item',
-                    'should end editing', 'should expand item',
-                    'should open( untitled)?',
-                    'should quit( after last window closed)?',
-                    'should select column', 'should select item',
-                    'should select row', 'should select tab view item',
-                    'should selection change', 'should zoom', 'shown',
-                    'update menu item', 'update parameters',
-                    'update toolbar item', 'was hidden', 'was miniaturized',
-                    'will become active', 'will close', 'will dismiss',
-                    'will display browser cell', 'will display cell',
-                    'will display item cell', 'will display outline cell',
-                    'will finish launching', 'will hide', 'will miniaturize',
-                    'will move', 'will open', 'will pop up', 'will quit',
-                    'will resign active', 'will resize( sub views)?',
-                    'will select tab view item', 'will show', 'will zoom',
-                    'write to file', 'zoomed')
-    StudioCommands = ('animate', 'append', 'call method', 'center',
-                      'close drawer', 'close panel', 'display',
-                      'display alert', 'display dialog', 'display panel', 'go',
-                      'hide', 'highlight', 'increment', 'item for',
-                      'load image', 'load movie', 'load nib', 'load panel',
-                      'load sound', 'localized string', 'lock focus', 'log',
-                      'open drawer', 'path for', 'pause', 'perform action',
-                      'play', 'register', 'resume', 'scroll', 'select( all)?',
-                      'show', 'size to fit', 'start', 'step back',
-                      'step forward', 'stop', 'synchronize', 'unlock focus',
-                      'update')
-    StudioProperties = ('accepts arrow key', 'action method', 'active',
-                        'alignment', 'allowed identifiers',
-                        'allows branch selection', 'allows column reordering',
-                        'allows column resizing', 'allows column selection',
-                        'allows customization',
-                        'allows editing text attributes',
-                        'allows empty selection', 'allows mixed state',
-                        'allows multiple selection', 'allows reordering',
-                        'allows undo', 'alpha( value)?', 'alternate image',
-                        'alternate increment value', 'alternate title',
-                        'animation delay', 'associated file name',
-                        'associated object', 'auto completes', 'auto display',
-                        'auto enables items', 'auto repeat',
-                        'auto resizes( outline column)?',
-                        'auto save expanded items', 'auto save name',
-                        'auto save table columns', 'auto saves configuration',
-                        'auto scroll', 'auto sizes all columns to fit',
-                        'auto sizes cells', 'background color', 'bezel state',
-                        'bezel style', 'bezeled', 'border rect', 'border type',
-                        'bordered', 'bounds( rotation)?', 'box type',
-                        'button returned', 'button type',
-                        'can choose directories', 'can choose files',
-                        'can draw', 'can hide',
-                        'cell( (background color|size|type))?', 'characters',
-                        'class', 'click count', 'clicked( data)? column',
-                        'clicked data item', 'clicked( data)? row',
-                        'closeable', 'collating', 'color( (mode|panel))',
-                        'command key down', 'configuration',
-                        'content(s| (size|view( margins)?))?', 'context',
-                        'continuous', 'control key down', 'control size',
-                        'control tint', 'control view',
-                        'controller visible', 'coordinate system',
-                        'copies( on scroll)?', 'corner view', 'current cell',
-                        'current column', 'current( field)?  editor',
-                        'current( menu)? item', 'current row',
-                        'current tab view item', 'data source',
-                        'default identifiers', 'delta (x|y|z)',
-                        'destination window', 'directory', 'display mode',
-                        'displayed cell', 'document( (edited|rect|view))?',
-                        'double value', 'dragged column', 'dragged distance',
-                        'dragged items', 'draws( cell)? background',
-                        'draws grid', 'dynamically scrolls', 'echos bullets',
-                        'edge', 'editable', 'edited( data)? column',
-                        'edited data item', 'edited( data)? row', 'enabled',
-                        'enclosing scroll view', 'ending page',
-                        'error handling', 'event number', 'event type',
-                        'excluded from windows menu', 'executable path',
-                        'expanded', 'fax number', 'field editor', 'file kind',
-                        'file name', 'file type', 'first responder',
-                        'first visible column', 'flipped', 'floating',
-                        'font( panel)?', 'formatter', 'frameworks path',
-                        'frontmost', 'gave up', 'grid color', 'has data items',
-                        'has horizontal ruler', 'has horizontal scroller',
-                        'has parent data item', 'has resize indicator',
-                        'has shadow', 'has sub menu', 'has vertical ruler',
-                        'has vertical scroller', 'header cell', 'header view',
-                        'hidden', 'hides when deactivated', 'highlights by',
-                        'horizontal line scroll', 'horizontal page scroll',
-                        'horizontal ruler view', 'horizontally resizable',
-                        'icon image', 'id', 'identifier',
-                        'ignores multiple clicks',
-                        'image( (alignment|dims when disabled|frame style|scaling))?',
-                        'imports graphics', 'increment value',
-                        'indentation per level', 'indeterminate', 'index',
-                        'integer value', 'intercell spacing', 'item height',
-                        'key( (code|equivalent( modifier)?|window))?',
-                        'knob thickness', 'label', 'last( visible)? column',
-                        'leading offset', 'leaf', 'level', 'line scroll',
-                        'loaded', 'localized sort', 'location', 'loop mode',
-                        'main( (bunde|menu|window))?', 'marker follows cell',
-                        'matrix mode', 'maximum( content)? size',
-                        'maximum visible columns',
-                        'menu( form representation)?', 'miniaturizable',
-                        'miniaturized', 'minimized image', 'minimized title',
-                        'minimum column width', 'minimum( content)? size',
-                        'modal', 'modified', 'mouse down state',
-                        'movie( (controller|file|rect))?', 'muted', 'name',
-                        'needs display', 'next state', 'next text',
-                        'number of tick marks', 'only tick mark values',
-                        'opaque', 'open panel', 'option key down',
-                        'outline table column', 'page scroll', 'pages across',
-                        'pages down', 'palette label', 'pane splitter',
-                        'parent data item', 'parent window', 'pasteboard',
-                        'path( (names|separator))?', 'playing',
-                        'plays every frame', 'plays selection only', 'position',
-                        'preferred edge', 'preferred type', 'pressure',
-                        'previous text', 'prompt', 'properties',
-                        'prototype cell', 'pulls down', 'rate',
-                        'released when closed', 'repeated',
-                        'requested print time', 'required file type',
-                        'resizable', 'resized column', 'resource path',
-                        'returns records', 'reuses columns', 'rich text',
-                        'roll over', 'row height', 'rulers visible',
-                        'save panel', 'scripts path', 'scrollable',
-                        'selectable( identifiers)?', 'selected cell',
-                        'selected( data)? columns?', 'selected data items?',
-                        'selected( data)? rows?', 'selected item identifier',
-                        'selection by rect', 'send action on arrow key',
-                        'sends action when done editing', 'separates columns',
-                        'separator item', 'sequence number', 'services menu',
-                        'shared frameworks path', 'shared support path',
-                        'sheet', 'shift key down', 'shows alpha',
-                        'shows state by', 'size( mode)?',
-                        'smart insert delete enabled', 'sort case sensitivity',
-                        'sort column', 'sort order', 'sort type',
-                        'sorted( data rows)?', 'sound', 'source( mask)?',
-                        'spell checking enabled', 'starting page', 'state',
-                        'string value', 'sub menu', 'super menu', 'super view',
-                        'tab key traverses cells', 'tab state', 'tab type',
-                        'tab view', 'table view', 'tag', 'target( printer)?',
-                        'text color', 'text container insert',
-                        'text container origin', 'text returned',
-                        'tick mark position', 'time stamp',
-                        'title(d| (cell|font|height|position|rect))?',
-                        'tool tip', 'toolbar', 'trailing offset', 'transparent',
-                        'treat packages as directories', 'truncated labels',
-                        'types', 'unmodified characters', 'update views',
-                        'use sort indicator', 'user defaults',
-                        'uses data source', 'uses ruler',
-                        'uses threaded animation',
-                        'uses title from previous column', 'value wraps',
-                        'version',
-                        'vertical( (line scroll|page scroll|ruler view))?',
-                        'vertically resizable', 'view',
-                        'visible( document rect)?', 'volume', 'width', 'window',
-                        'windows menu', 'wraps', 'zoomable', 'zoomed')
-
-    tokens = {
-        'root': [
-            (r'\s+', Text),
-            (r'¬\n', String.Escape),
-            (r"'s\s+", Text),  # This is a possessive, consider moving
-            (r'(--|#).*?$', Comment),
-            (r'\(\*', Comment.Multiline, 'comment'),
-            (r'[(){}!,.:]', Punctuation),
-            (r'(«)([^»]+)(»)',
-             bygroups(Text, Name.Builtin, Text)),
-            (r'\b((?:considering|ignoring)\s*)'
-             r'(application responses|case|diacriticals|hyphens|'
-             r'numeric strings|punctuation|white space)',
-             bygroups(Keyword, Name.Builtin)),
-            (r'(-|\*|\+|&|≠|>=?|<=?|=|≥|≤|/|÷|\^)', Operator),
-            (r"\b({})\b".format('|'.join(Operators)), Operator.Word),
-            (r'^(\s*(?:on|end)\s+)'
-             r'({})'.format('|'.join(StudioEvents[::-1])),
-             bygroups(Keyword, Name.Function)),
-            (r'^(\s*)(in|on|script|to)(\s+)', bygroups(Text, Keyword, Text)),
-            (r'\b(as )({})\b'.format('|'.join(Classes)),
-             bygroups(Keyword, Name.Class)),
-            (r'\b({})\b'.format('|'.join(Literals)), Name.Constant),
-            (r'\b({})\b'.format('|'.join(Commands)), Name.Builtin),
-            (r'\b({})\b'.format('|'.join(Control)), Keyword),
-            (r'\b({})\b'.format('|'.join(Declarations)), Keyword),
-            (r'\b({})\b'.format('|'.join(Reserved)), Name.Builtin),
-            (r'\b({})s?\b'.format('|'.join(BuiltIn)), Name.Builtin),
-            (r'\b({})\b'.format('|'.join(HandlerParams)), Name.Builtin),
-            (r'\b({})\b'.format('|'.join(StudioProperties)), Name.Attribute),
-            (r'\b({})s?\b'.format('|'.join(StudioClasses)), Name.Builtin),
-            (r'\b({})\b'.format('|'.join(StudioCommands)), Name.Builtin),
-            (r'\b({})\b'.format('|'.join(References)), Name.Builtin),
-            (r'"(\\\\|\\[^\\]|[^"\\])*"', String.Double),
-            (rf'\b({Identifiers})\b', Name.Variable),
-            (r'[-+]?(\d+\.\d*|\d*\.\d+)(E[-+][0-9]+)?', Number.Float),
-            (r'[-+]?\d+', Number.Integer),
-        ],
-        'comment': [
-            (r'\(\*', Comment.Multiline, '#push'),
-            (r'\*\)', Comment.Multiline, '#pop'),
-            ('[^*(]+', Comment.Multiline),
-            ('[*(]', Comment.Multiline),
-        ],
-    }
-
-
-class RexxLexer(RegexLexer):
-    """
-    Rexx is a scripting language available for
-    a wide range of different platforms with its roots found on mainframe
-    systems. It is popular for I/O- and data based tasks and can act as glue
-    language to bind different applications together.
-    """
-    name = 'Rexx'
-    url = 'http://www.rexxinfo.org/'
-    aliases = ['rexx', 'arexx']
-    filenames = ['*.rexx', '*.rex', '*.rx', '*.arexx']
-    mimetypes = ['text/x-rexx']
-    version_added = '2.0'
-    flags = re.IGNORECASE
-
-    tokens = {
-        'root': [
-            (r'\s+', Whitespace),
-            (r'/\*', Comment.Multiline, 'comment'),
-            (r'"', String, 'string_double'),
-            (r"'", String, 'string_single'),
-            (r'[0-9]+(\.[0-9]+)?(e[+-]?[0-9])?', Number),
-            (r'([a-z_]\w*)(\s*)(:)(\s*)(procedure)\b',
-             bygroups(Name.Function, Whitespace, Operator, Whitespace,
-                      Keyword.Declaration)),
-            (r'([a-z_]\w*)(\s*)(:)',
-             bygroups(Name.Label, Whitespace, Operator)),
-            include('function'),
-            include('keyword'),
-            include('operator'),
-            (r'[a-z_]\w*', Text),
-        ],
-        'function': [
-            (words((
-                'abbrev', 'abs', 'address', 'arg', 'b2x', 'bitand', 'bitor', 'bitxor',
-                'c2d', 'c2x', 'center', 'charin', 'charout', 'chars', 'compare',
-                'condition', 'copies', 'd2c', 'd2x', 'datatype', 'date', 'delstr',
-                'delword', 'digits', 'errortext', 'form', 'format', 'fuzz', 'insert',
-                'lastpos', 'left', 'length', 'linein', 'lineout', 'lines', 'max',
-                'min', 'overlay', 'pos', 'queued', 'random', 'reverse', 'right', 'sign',
-                'sourceline', 'space', 'stream', 'strip', 'substr', 'subword', 'symbol',
-                'time', 'trace', 'translate', 'trunc', 'value', 'verify', 'word',
-                'wordindex', 'wordlength', 'wordpos', 'words', 'x2b', 'x2c', 'x2d',
-                'xrange'), suffix=r'(\s*)(\()'),
-             bygroups(Name.Builtin, Whitespace, Operator)),
-        ],
-        'keyword': [
-            (r'(address|arg|by|call|do|drop|else|end|exit|for|forever|if|'
-             r'interpret|iterate|leave|nop|numeric|off|on|options|parse|'
-             r'pull|push|queue|return|say|select|signal|to|then|trace|until|'
-             r'while)\b', Keyword.Reserved),
-        ],
-        'operator': [
-            (r'(-|//|/|\(|\)|\*\*|\*|\\<<|\\<|\\==|\\=|\\>>|\\>|\\|\|\||\||'
-             r'&&|&|%|\+|<<=|<<|<=|<>|<|==|=|><|>=|>>=|>>|>|¬<<|¬<|¬==|¬=|'
-             r'¬>>|¬>|¬|\.|,)', Operator),
-        ],
-        'string_double': [
-            (r'[^"\n]+', String),
-            (r'""', String),
-            (r'"', String, '#pop'),
-            (r'\n', Text, '#pop'),  # Stray linefeed also terminates strings.
-        ],
-        'string_single': [
-            (r'[^\'\n]+', String),
-            (r'\'\'', String),
-            (r'\'', String, '#pop'),
-            (r'\n', Text, '#pop'),  # Stray linefeed also terminates strings.
-        ],
-        'comment': [
-            (r'[^*]+', Comment.Multiline),
-            (r'\*/', Comment.Multiline, '#pop'),
-            (r'\*', Comment.Multiline),
-        ]
-    }
-
-    def _c(s):
-        return re.compile(s, re.MULTILINE)
-    _ADDRESS_COMMAND_PATTERN = _c(r'^\s*address\s+command\b')
-    _ADDRESS_PATTERN = _c(r'^\s*address\s+')
-    _DO_WHILE_PATTERN = _c(r'^\s*do\s+while\b')
-    _IF_THEN_DO_PATTERN = _c(r'^\s*if\b.+\bthen\s+do\s*$')
-    _PROCEDURE_PATTERN = _c(r'^\s*([a-z_]\w*)(\s*)(:)(\s*)(procedure)\b')
-    _ELSE_DO_PATTERN = _c(r'\belse\s+do\s*$')
-    _PARSE_ARG_PATTERN = _c(r'^\s*parse\s+(upper\s+)?(arg|value)\b')
-    PATTERNS_AND_WEIGHTS = (
-        (_ADDRESS_COMMAND_PATTERN, 0.2),
-        (_ADDRESS_PATTERN, 0.05),
-        (_DO_WHILE_PATTERN, 0.1),
-        (_ELSE_DO_PATTERN, 0.1),
-        (_IF_THEN_DO_PATTERN, 0.1),
-        (_PROCEDURE_PATTERN, 0.5),
-        (_PARSE_ARG_PATTERN, 0.2),
-    )
-
-    def analyse_text(text):
-        """
-        Check for initial comment and patterns that distinguish Rexx from other
-        C-like languages.
-        """
-        if re.search(r'/\*\**\s*rexx', text, re.IGNORECASE):
-            # Header matches MVS Rexx requirements, this is certainly a Rexx
-            # script.
-            return 1.0
-        elif text.startswith('/*'):
-            # Header matches general Rexx requirements; the source code might
-            # still be any language using C comments such as C++, C# or Java.
-            lowerText = text.lower()
-            result = sum(weight
-                         for (pattern, weight) in RexxLexer.PATTERNS_AND_WEIGHTS
-                         if pattern.search(lowerText)) + 0.01
-            return min(result, 1.0)
-
-
-class MOOCodeLexer(RegexLexer):
-    """
-    For MOOCode (the MOO scripting language).
-    """
-    name = 'MOOCode'
-    url = 'http://www.moo.mud.org/'
-    filenames = ['*.moo']
-    aliases = ['moocode', 'moo']
-    mimetypes = ['text/x-moocode']
-    version_added = '0.9'
-
-    tokens = {
-        'root': [
-            # Numbers
-            (r'(0|[1-9][0-9_]*)', Number.Integer),
-            # Strings
-            (r'"(\\\\|\\[^\\]|[^"\\])*"', String),
-            # exceptions
-            (r'(E_PERM|E_DIV)', Name.Exception),
-            # db-refs
-            (r'((#[-0-9]+)|(\$\w+))', Name.Entity),
-            # Keywords
-            (r'\b(if|else|elseif|endif|for|endfor|fork|endfork|while'
-             r'|endwhile|break|continue|return|try'
-             r'|except|endtry|finally|in)\b', Keyword),
-            # builtins
-            (r'(random|length)', Name.Builtin),
-            # special variables
-            (r'(player|caller|this|args)', Name.Variable.Instance),
-            # skip whitespace
-            (r'\s+', Text),
-            (r'\n', Text),
-            # other operators
-            (r'([!;=,{}&|:.\[\]@()<>?]+)', Operator),
-            # function call
-            (r'(\w+)(\()', bygroups(Name.Function, Operator)),
-            # variables
-            (r'(\w+)', Text),
-        ]
-    }
-
-
-class HybrisLexer(RegexLexer):
-    """
-    For Hybris source code.
-    """
-
-    name = 'Hybris'
-    aliases = ['hybris']
-    filenames = ['*.hyb']
-    mimetypes = ['text/x-hybris', 'application/x-hybris']
-    url = 'https://github.com/evilsocket/hybris'
-    version_added = '1.4'
-
-    flags = re.MULTILINE | re.DOTALL
-
-    tokens = {
-        'root': [
-            # method names
-            (r'^(\s*(?:function|method|operator\s+)+?)'
-             r'([a-zA-Z_]\w*)'
-             r'(\s*)(\()', bygroups(Keyword, Name.Function, Text, Operator)),
-            (r'[^\S\n]+', Text),
-            (r'//.*?\n', Comment.Single),
-            (r'/\*.*?\*/', Comment.Multiline),
-            (r'@[a-zA-Z_][\w.]*', Name.Decorator),
-            (r'(break|case|catch|next|default|do|else|finally|for|foreach|of|'
-             r'unless|if|new|return|switch|me|throw|try|while)\b', Keyword),
-            (r'(extends|private|protected|public|static|throws|function|method|'
-             r'operator)\b', Keyword.Declaration),
-            (r'(true|false|null|__FILE__|__LINE__|__VERSION__|__LIB_PATH__|'
-             r'__INC_PATH__)\b', Keyword.Constant),
-            (r'(class|struct)(\s+)',
-             bygroups(Keyword.Declaration, Text), 'class'),
-            (r'(import|include)(\s+)',
-             bygroups(Keyword.Namespace, Text), 'import'),
-            (words((
-                'gc_collect', 'gc_mm_items', 'gc_mm_usage', 'gc_collect_threshold',
-                'urlencode', 'urldecode', 'base64encode', 'base64decode', 'sha1', 'crc32',
-                'sha2', 'md5', 'md5_file', 'acos', 'asin', 'atan', 'atan2', 'ceil', 'cos',
-                'cosh', 'exp', 'fabs', 'floor', 'fmod', 'log', 'log10', 'pow', 'sin',
-                'sinh', 'sqrt', 'tan', 'tanh', 'isint', 'isfloat', 'ischar', 'isstring',
-                'isarray', 'ismap', 'isalias', 'typeof', 'sizeof', 'toint', 'tostring',
-                'fromxml', 'toxml', 'binary', 'pack', 'load', 'eval', 'var_names',
-                'var_values', 'user_functions', 'dyn_functions', 'methods', 'call',
-                'call_method', 'mknod', 'mkfifo', 'mount', 'umount2', 'umount', 'ticks',
-                'usleep', 'sleep', 'time', 'strtime', 'strdate', 'dllopen', 'dlllink',
-                'dllcall', 'dllcall_argv', 'dllclose', 'env', 'exec', 'fork', 'getpid',
-                'wait', 'popen', 'pclose', 'exit', 'kill', 'pthread_create',
-                'pthread_create_argv', 'pthread_exit', 'pthread_join', 'pthread_kill',
-                'smtp_send', 'http_get', 'http_post', 'http_download', 'socket', 'bind',
-                'listen', 'accept', 'getsockname', 'getpeername', 'settimeout', 'connect',
-                'server', 'recv', 'send', 'close', 'print', 'println', 'printf', 'input',
-                'readline', 'serial_open', 'serial_fcntl', 'serial_get_attr',
-                'serial_get_ispeed', 'serial_get_ospeed', 'serial_set_attr',
-                'serial_set_ispeed', 'serial_set_ospeed', 'serial_write', 'serial_read',
-                'serial_close', 'xml_load', 'xml_parse', 'fopen', 'fseek', 'ftell',
-                'fsize', 'fread', 'fwrite', 'fgets', 'fclose', 'file', 'readdir',
-                'pcre_replace', 'size', 'pop', 'unmap', 'has', 'keys', 'values',
-                'length', 'find', 'substr', 'replace', 'split', 'trim', 'remove',
-                'contains', 'join'), suffix=r'\b'),
-             Name.Builtin),
-            (words((
-                'MethodReference', 'Runner', 'Dll', 'Thread', 'Pipe', 'Process',
-                'Runnable', 'CGI', 'ClientSocket', 'Socket', 'ServerSocket',
-                'File', 'Console', 'Directory', 'Exception'), suffix=r'\b'),
-             Keyword.Type),
-            (r'"(\\\\|\\[^\\]|[^"\\])*"', String),
-            (r"'\\.'|'[^\\]'|'\\u[0-9a-f]{4}'", String.Char),
-            (r'(\.)([a-zA-Z_]\w*)',
-             bygroups(Operator, Name.Attribute)),
-            (r'[a-zA-Z_]\w*:', Name.Label),
-            (r'[a-zA-Z_$]\w*', Name),
-            (r'[~^*!%&\[\](){}<>|+=:;,./?\-@]+', Operator),
-            (r'[0-9][0-9]*\.[0-9]+([eE][0-9]+)?[fd]?', Number.Float),
-            (r'0x[0-9a-f]+', Number.Hex),
-            (r'[0-9]+L?', Number.Integer),
-            (r'\n', Text),
-        ],
-        'class': [
-            (r'[a-zA-Z_]\w*', Name.Class, '#pop')
-        ],
-        'import': [
-            (r'[\w.]+\*?', Name.Namespace, '#pop')
-        ],
-    }
-
-    def analyse_text(text):
-        """public method and private method don't seem to be quite common
-        elsewhere."""
-        result = 0
-        if re.search(r'\b(?:public|private)\s+method\b', text):
-            result += 0.01
-        return result
-
-
-
-class EasytrieveLexer(RegexLexer):
-    """
-    Easytrieve Plus is a programming language for extracting, filtering and
-    converting sequential data. Furthermore it can layout data for reports.
-    It is mainly used on mainframe platforms and can access several of the
-    mainframe's native file formats. It is somewhat comparable to awk.
-    """
-    name = 'Easytrieve'
-    aliases = ['easytrieve']
-    filenames = ['*.ezt', '*.mac']
-    mimetypes = ['text/x-easytrieve']
-    url = 'https://www.broadcom.com/products/mainframe/application-development/easytrieve-report-generator'
-    version_added = '2.1'
-    flags = 0
-
-    # Note: We cannot use r'\b' at the start and end of keywords because
-    # Easytrieve Plus delimiter characters are:
-    #
-    #   * space ( )
-    #   * apostrophe (')
-    #   * period (.)
-    #   * comma (,)
-    #   * parenthesis ( and )
-    #   * colon (:)
-    #
-    # Additionally words end once a '*' appears, indicatins a comment.
-    _DELIMITERS = r' \'.,():\n'
-    _DELIMITERS_OR_COMENT = _DELIMITERS + '*'
-    _DELIMITER_PATTERN = '[' + _DELIMITERS + ']'
-    _DELIMITER_PATTERN_CAPTURE = '(' + _DELIMITER_PATTERN + ')'
-    _NON_DELIMITER_OR_COMMENT_PATTERN = '[^' + _DELIMITERS_OR_COMENT + ']'
-    _OPERATORS_PATTERN = '[.+\\-/=\\[\\](){}<>;,&%¬]'
-    _KEYWORDS = [
-        'AFTER-BREAK', 'AFTER-LINE', 'AFTER-SCREEN', 'AIM', 'AND', 'ATTR',
-        'BEFORE', 'BEFORE-BREAK', 'BEFORE-LINE', 'BEFORE-SCREEN', 'BUSHU',
-        'BY', 'CALL', 'CASE', 'CHECKPOINT', 'CHKP', 'CHKP-STATUS', 'CLEAR',
-        'CLOSE', 'COL', 'COLOR', 'COMMIT', 'CONTROL', 'COPY', 'CURSOR', 'D',
-        'DECLARE', 'DEFAULT', 'DEFINE', 'DELETE', 'DENWA', 'DISPLAY', 'DLI',
-        'DO', 'DUPLICATE', 'E', 'ELSE', 'ELSE-IF', 'END', 'END-CASE',
-        'END-DO', 'END-IF', 'END-PROC', 'ENDPAGE', 'ENDTABLE', 'ENTER', 'EOF',
-        'EQ', 'ERROR', 'EXIT', 'EXTERNAL', 'EZLIB', 'F1', 'F10', 'F11', 'F12',
-        'F13', 'F14', 'F15', 'F16', 'F17', 'F18', 'F19', 'F2', 'F20', 'F21',
-        'F22', 'F23', 'F24', 'F25', 'F26', 'F27', 'F28', 'F29', 'F3', 'F30',
-        'F31', 'F32', 'F33', 'F34', 'F35', 'F36', 'F4', 'F5', 'F6', 'F7',
-        'F8', 'F9', 'FETCH', 'FILE-STATUS', 'FILL', 'FINAL', 'FIRST',
-        'FIRST-DUP', 'FOR', 'GE', 'GET', 'GO', 'GOTO', 'GQ', 'GR', 'GT',
-        'HEADING', 'HEX', 'HIGH-VALUES', 'IDD', 'IDMS', 'IF', 'IN', 'INSERT',
-        'JUSTIFY', 'KANJI-DATE', 'KANJI-DATE-LONG', 'KANJI-TIME', 'KEY',
-        'KEY-PRESSED', 'KOKUGO', 'KUN', 'LAST-DUP', 'LE', 'LEVEL', 'LIKE',
-        'LINE', 'LINE-COUNT', 'LINE-NUMBER', 'LINK', 'LIST', 'LOW-VALUES',
-        'LQ', 'LS', 'LT', 'MACRO', 'MASK', 'MATCHED', 'MEND', 'MESSAGE',
-        'MOVE', 'MSTART', 'NE', 'NEWPAGE', 'NOMASK', 'NOPRINT', 'NOT',
-        'NOTE', 'NOVERIFY', 'NQ', 'NULL', 'OF', 'OR', 'OTHERWISE', 'PA1',
-        'PA2', 'PA3', 'PAGE-COUNT', 'PAGE-NUMBER', 'PARM-REGISTER',
-        'PATH-ID', 'PATTERN', 'PERFORM', 'POINT', 'POS', 'PRIMARY', 'PRINT',
-        'PROCEDURE', 'PROGRAM', 'PUT', 'READ', 'RECORD', 'RECORD-COUNT',
-        'RECORD-LENGTH', 'REFRESH', 'RELEASE', 'RENUM', 'REPEAT', 'REPORT',
-        'REPORT-INPUT', 'RESHOW', 'RESTART', 'RETRIEVE', 'RETURN-CODE',
-        'ROLLBACK', 'ROW', 'S', 'SCREEN', 'SEARCH', 'SECONDARY', 'SELECT',
-        'SEQUENCE', 'SIZE', 'SKIP', 'SOKAKU', 'SORT', 'SQL', 'STOP', 'SUM',
-        'SYSDATE', 'SYSDATE-LONG', 'SYSIN', 'SYSIPT', 'SYSLST', 'SYSPRINT',
-        'SYSSNAP', 'SYSTIME', 'TALLY', 'TERM-COLUMNS', 'TERM-NAME',
-        'TERM-ROWS', 'TERMINATION', 'TITLE', 'TO', 'TRANSFER', 'TRC',
-        'UNIQUE', 'UNTIL', 'UPDATE', 'UPPERCASE', 'USER', 'USERID', 'VALUE',
-        'VERIFY', 'W', 'WHEN', 'WHILE', 'WORK', 'WRITE', 'X', 'XDM', 'XRST'
-    ]
-
-    tokens = {
-        'root': [
-            (r'\*.*\n', Comment.Single),
-            (r'\n+', Whitespace),
-            # Macro argument
-            (r'&' + _NON_DELIMITER_OR_COMMENT_PATTERN + r'+\.', Name.Variable,
-             'after_macro_argument'),
-            # Macro call
-            (r'%' + _NON_DELIMITER_OR_COMMENT_PATTERN + r'+', Name.Variable),
-            (r'(FILE|MACRO|REPORT)(\s+)',
-             bygroups(Keyword.Declaration, Whitespace), 'after_declaration'),
-            (r'(JOB|PARM)' + r'(' + _DELIMITER_PATTERN + r')',
-             bygroups(Keyword.Declaration, Operator)),
-            (words(_KEYWORDS, suffix=_DELIMITER_PATTERN_CAPTURE),
-             bygroups(Keyword.Reserved, Operator)),
-            (_OPERATORS_PATTERN, Operator),
-            # Procedure declaration
-            (r'(' + _NON_DELIMITER_OR_COMMENT_PATTERN + r'+)(\s*)(\.?)(\s*)(PROC)(\s*\n)',
-             bygroups(Name.Function, Whitespace, Operator, Whitespace,
-                      Keyword.Declaration, Whitespace)),
-            (r'[0-9]+\.[0-9]*', Number.Float),
-            (r'[0-9]+', Number.Integer),
-            (r"'(''|[^'])*'", String),
-            (r'\s+', Whitespace),
-            # Everything else just belongs to a name
-            (_NON_DELIMITER_OR_COMMENT_PATTERN + r'+', Name),
-         ],
-        'after_declaration': [
-            (_NON_DELIMITER_OR_COMMENT_PATTERN + r'+', Name.Function),
-            default('#pop'),
-        ],
-        'after_macro_argument': [
-            (r'\*.*\n', Comment.Single, '#pop'),
-            (r'\s+', Whitespace, '#pop'),
-            (_OPERATORS_PATTERN, Operator, '#pop'),
-            (r"'(''|[^'])*'", String, '#pop'),
-            # Everything else just belongs to a name
-            (_NON_DELIMITER_OR_COMMENT_PATTERN + r'+', Name),
-        ],
-    }
-    _COMMENT_LINE_REGEX = re.compile(r'^\s*\*')
-    _MACRO_HEADER_REGEX = re.compile(r'^\s*MACRO')
-
-    def analyse_text(text):
-        """
-        Perform a structural analysis for basic Easytrieve constructs.
-        """
-        result = 0.0
-        lines = text.split('\n')
-        hasEndProc = False
-        hasHeaderComment = False
-        hasFile = False
-        hasJob = False
-        hasProc = False
-        hasParm = False
-        hasReport = False
-
-        def isCommentLine(line):
-            return EasytrieveLexer._COMMENT_LINE_REGEX.match(lines[0]) is not None
-
-        def isEmptyLine(line):
-            return not bool(line.strip())
-
-        # Remove possible empty lines and header comments.
-        while lines and (isEmptyLine(lines[0]) or isCommentLine(lines[0])):
-            if not isEmptyLine(lines[0]):
-                hasHeaderComment = True
-            del lines[0]
-
-        if EasytrieveLexer._MACRO_HEADER_REGEX.match(lines[0]):
-            # Looks like an Easytrieve macro.
-            result = 0.4
-            if hasHeaderComment:
-                result += 0.4
-        else:
-            # Scan the source for lines starting with indicators.
-            for line in lines:
-                words = line.split()
-                if (len(words) >= 2):
-                    firstWord = words[0]
-                    if not hasReport:
-                        if not hasJob:
-                            if not hasFile:
-                                if not hasParm:
-                                    if firstWord == 'PARM':
-                                        hasParm = True
-                                if firstWord == 'FILE':
-                                    hasFile = True
-                            if firstWord == 'JOB':
-                                hasJob = True
-                        elif firstWord == 'PROC':
-                            hasProc = True
-                        elif firstWord == 'END-PROC':
-                            hasEndProc = True
-                        elif firstWord == 'REPORT':
-                            hasReport = True
-
-            # Weight the findings.
-            if hasJob and (hasProc == hasEndProc):
-                if hasHeaderComment:
-                    result += 0.1
-                if hasParm:
-                    if hasProc:
-                        # Found PARM, JOB and PROC/END-PROC:
-                        # pretty sure this is Easytrieve.
-                        result += 0.8
-                    else:
-                        # Found PARAM and  JOB: probably this is Easytrieve
-                        result += 0.5
-                else:
-                    # Found JOB and possibly other keywords: might be Easytrieve
-                    result += 0.11
-                    if hasParm:
-                        # Note: PARAM is not a proper English word, so this is
-                        # regarded a much better indicator for Easytrieve than
-                        # the other words.
-                        result += 0.2
-                    if hasFile:
-                        result += 0.01
-                    if hasReport:
-                        result += 0.01
-        assert 0.0 <= result <= 1.0
-        return result
-
-
-class JclLexer(RegexLexer):
-    """
-    Job Control Language (JCL)
-    is a scripting language used on mainframe platforms to instruct the system
-    on how to run a batch job or start a subsystem. It is somewhat
-    comparable to MS DOS batch and Unix shell scripts.
-    """
-    name = 'JCL'
-    aliases = ['jcl']
-    filenames = ['*.jcl']
-    mimetypes = ['text/x-jcl']
-    url = 'https://en.wikipedia.org/wiki/Job_Control_Language'
-    version_added = '2.1'
-
-    flags = re.IGNORECASE
-
-    tokens = {
-        'root': [
-            (r'//\*.*\n', Comment.Single),
-            (r'//', Keyword.Pseudo, 'statement'),
-            (r'/\*', Keyword.Pseudo, 'jes2_statement'),
-            # TODO: JES3 statement
-            (r'.*\n', Other)  # Input text or inline code in any language.
-        ],
-        'statement': [
-            (r'\s*\n', Whitespace, '#pop'),
-            (r'([a-z]\w*)(\s+)(exec|job)(\s*)',
-             bygroups(Name.Label, Whitespace, Keyword.Reserved, Whitespace),
-             'option'),
-            (r'[a-z]\w*', Name.Variable, 'statement_command'),
-            (r'\s+', Whitespace, 'statement_command'),
-        ],
-        'statement_command': [
-            (r'\s+(command|cntl|dd|endctl|endif|else|include|jcllib|'
-             r'output|pend|proc|set|then|xmit)\s+', Keyword.Reserved, 'option'),
-            include('option')
-        ],
-        'jes2_statement': [
-            (r'\s*\n', Whitespace, '#pop'),
-            (r'\$', Keyword, 'option'),
-            (r'\b(jobparam|message|netacct|notify|output|priority|route|'
-             r'setup|signoff|xeq|xmit)\b', Keyword, 'option'),
-        ],
-        'option': [
-            # (r'\n', Text, 'root'),
-            (r'\*', Name.Builtin),
-            (r'[\[\](){}<>;,]', Punctuation),
-            (r'[-+*/=&%]', Operator),
-            (r'[a-z_]\w*', Name),
-            (r'\d+\.\d*', Number.Float),
-            (r'\.\d+', Number.Float),
-            (r'\d+', Number.Integer),
-            (r"'", String, 'option_string'),
-            (r'[ \t]+', Whitespace, 'option_comment'),
-            (r'\.', Punctuation),
-        ],
-        'option_string': [
-            (r"(\n)(//)", bygroups(Text, Keyword.Pseudo)),
-            (r"''", String),
-            (r"[^']", String),
-            (r"'", String, '#pop'),
-        ],
-        'option_comment': [
-            # (r'\n', Text, 'root'),
-            (r'.+', Comment.Single),
-        ]
-    }
-
-    _JOB_HEADER_PATTERN = re.compile(r'^//[a-z#$@][a-z0-9#$@]{0,7}\s+job(\s+.*)?$',
-                                     re.IGNORECASE)
-
-    def analyse_text(text):
-        """
-        Recognize JCL job by header.
-        """
-        result = 0.0
-        lines = text.split('\n')
-        if len(lines) > 0:
-            if JclLexer._JOB_HEADER_PATTERN.match(lines[0]):
-                result = 1.0
-        assert 0.0 <= result <= 1.0
-        return result
-
-
-class MiniScriptLexer(RegexLexer):
-    """
-    For MiniScript source code.
-    """
-
-    name = 'MiniScript'
-    url = 'https://miniscript.org'
-    aliases = ['miniscript', 'ms']
-    filenames = ['*.ms']
-    mimetypes = ['text/x-minicript', 'application/x-miniscript']
-    version_added = '2.6'
-
-    tokens = {
-        'root': [
-            (r'#!(.*?)$', Comment.Preproc),
-            default('base'),
-        ],
-        'base': [
-            ('//.*$', Comment.Single),
-            (r'(?i)(\d*\.\d+|\d+\.\d*)(e[+-]?\d+)?', Number),
-            (r'(?i)\d+e[+-]?\d+', Number),
-            (r'\d+', Number),
-            (r'\n', Text),
-            (r'[^\S\n]+', Text),
-            (r'"', String, 'string_double'),
-            (r'(==|!=|<=|>=|[=+\-*/%^<>.:])', Operator),
-            (r'[;,\[\]{}()]', Punctuation),
-            (words((
-                'break', 'continue', 'else', 'end', 'for', 'function', 'if',
-                'in', 'isa', 'then', 'repeat', 'return', 'while'), suffix=r'\b'),
-             Keyword),
-            (words((
-                'abs', 'acos', 'asin', 'atan', 'ceil', 'char', 'cos', 'floor',
-                'log', 'round', 'rnd', 'pi', 'sign', 'sin', 'sqrt', 'str', 'tan',
-                'hasIndex', 'indexOf', 'len', 'val', 'code', 'remove', 'lower',
-                'upper', 'replace', 'split', 'indexes', 'values', 'join', 'sum',
-                'sort', 'shuffle', 'push', 'pop', 'pull', 'range',
-                'print', 'input', 'time', 'wait', 'locals', 'globals', 'outer',
-                'yield'), suffix=r'\b'),
-             Name.Builtin),
-            (r'(true|false|null)\b', Keyword.Constant),
-            (r'(and|or|not|new)\b', Operator.Word),
-            (r'(self|super|__isa)\b', Name.Builtin.Pseudo),
-            (r'[a-zA-Z_]\w*', Name.Variable)
-        ],
-        'string_double': [
-            (r'[^"\n]+', String),
-            (r'""', String),
-            (r'"', String, '#pop'),
-            (r'\n', Text, '#pop'),  # Stray linefeed also terminates strings.
-        ]
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/sgf.py b/venv/lib/python3.10/site-packages/pygments/lexers/sgf.py
deleted file mode 100644
index f0e56cba55242ba7874e1e71186be35986537407..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/sgf.py
+++ /dev/null
@@ -1,59 +0,0 @@
-"""
-    pygments.lexers.sgf
-    ~~~~~~~~~~~~~~~~~~~
-
-    Lexer for Smart Game Format (sgf) file format.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, bygroups
-from pygments.token import Name, Literal, String, Punctuation, Whitespace
-
-__all__ = ["SmartGameFormatLexer"]
-
-
-class SmartGameFormatLexer(RegexLexer):
-    """
-    Lexer for Smart Game Format (sgf) file format.
-
-    The format is used to store game records of board games for two players
-    (mainly Go game).
-    """
-    name = 'SmartGameFormat'
-    url = 'https://www.red-bean.com/sgf/'
-    aliases = ['sgf']
-    filenames = ['*.sgf']
-    version_added = '2.4'
-
-    tokens = {
-        'root': [
-            (r'[():;]+', Punctuation),
-            # tokens:
-            (r'(A[BW]|AE|AN|AP|AR|AS|[BW]L|BM|[BW]R|[BW]S|[BW]T|CA|CH|CP|CR|'
-             r'DD|DM|DO|DT|EL|EV|EX|FF|FG|G[BW]|GC|GM|GN|HA|HO|ID|IP|IT|IY|KM|'
-             r'KO|LB|LN|LT|L|MA|MN|M|N|OB|OM|ON|OP|OT|OV|P[BW]|PC|PL|PM|RE|RG|'
-             r'RO|RU|SO|SC|SE|SI|SL|SO|SQ|ST|SU|SZ|T[BW]|TC|TE|TM|TR|UC|US|VW|'
-             r'V|[BW]|C)',
-             Name.Builtin),
-            # number:
-            (r'(\[)([0-9.]+)(\])',
-             bygroups(Punctuation, Literal.Number, Punctuation)),
-            # date:
-            (r'(\[)([0-9]{4}-[0-9]{2}-[0-9]{2})(\])',
-             bygroups(Punctuation, Literal.Date, Punctuation)),
-            # point:
-            (r'(\[)([a-z]{2})(\])',
-             bygroups(Punctuation, String, Punctuation)),
-            # double points:
-            (r'(\[)([a-z]{2})(:)([a-z]{2})(\])',
-             bygroups(Punctuation, String, Punctuation, String, Punctuation)),
-
-            (r'(\[)([\w\s#()+,\-.:?]+)(\])',
-             bygroups(Punctuation, String, Punctuation)),
-            (r'(\[)(\s.*)(\])',
-             bygroups(Punctuation, Whitespace, Punctuation)),
-            (r'\s+', Whitespace)
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/shell.py b/venv/lib/python3.10/site-packages/pygments/lexers/shell.py
deleted file mode 100644
index 744767a1d4464f838b7c062f63c850e3a2f18554..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/shell.py
+++ /dev/null
@@ -1,902 +0,0 @@
-"""
-    pygments.lexers.shell
-    ~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for various shells.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import re
-
-from pygments.lexer import Lexer, RegexLexer, do_insertions, bygroups, \
-    include, default, this, using, words, line_re
-from pygments.token import Punctuation, Whitespace, \
-    Text, Comment, Operator, Keyword, Name, String, Number, Generic
-from pygments.util import shebang_matches
-
-__all__ = ['BashLexer', 'BashSessionLexer', 'TcshLexer', 'BatchLexer',
-           'SlurmBashLexer', 'MSDOSSessionLexer', 'PowerShellLexer',
-           'PowerShellSessionLexer', 'TcshSessionLexer', 'FishShellLexer',
-           'ExeclineLexer']
-
-
-class BashLexer(RegexLexer):
-    """
-    Lexer for (ba|k|z|)sh shell scripts.
-    """
-
-    name = 'Bash'
-    aliases = ['bash', 'sh', 'ksh', 'zsh', 'shell', 'openrc']
-    filenames = ['*.sh', '*.ksh', '*.bash', '*.ebuild', '*.eclass',
-                 '*.exheres-0', '*.exlib', '*.zsh',
-                 '.bashrc', 'bashrc', '.bash_*', 'bash_*', 'zshrc', '.zshrc',
-                 '.kshrc', 'kshrc',
-                 'PKGBUILD']
-    mimetypes = ['application/x-sh', 'application/x-shellscript', 'text/x-shellscript']
-    url = 'https://en.wikipedia.org/wiki/Unix_shell'
-    version_added = '0.6'
-
-    tokens = {
-        'root': [
-            include('basic'),
-            (r'`', String.Backtick, 'backticks'),
-            include('data'),
-            include('interp'),
-        ],
-        'interp': [
-            (r'\$\(\(', Keyword, 'math'),
-            (r'\$\(', Keyword, 'paren'),
-            (r'\$\{#?', String.Interpol, 'curly'),
-            (r'\$[a-zA-Z_]\w*', Name.Variable),  # user variable
-            (r'\$(?:\d+|[#$?!_*@-])', Name.Variable),      # builtin
-            (r'\$', Text),
-        ],
-        'basic': [
-            (r'\b(if|fi|else|while|in|do|done|for|then|return|function|case|'
-             r'select|break|continue|until|esac|elif)(\s*)\b',
-             bygroups(Keyword, Whitespace)),
-            (r'\b(alias|bg|bind|builtin|caller|cd|command|compgen|'
-             r'complete|declare|dirs|disown|echo|enable|eval|exec|exit|'
-             r'export|false|fc|fg|getopts|hash|help|history|jobs|kill|let|'
-             r'local|logout|popd|printf|pushd|pwd|read|readonly|set|shift|'
-             r'shopt|source|suspend|test|time|times|trap|true|type|typeset|'
-             r'ulimit|umask|unalias|unset|wait)(?=[\s)`])',
-             Name.Builtin),
-            (r'\A#!.+\n', Comment.Hashbang),
-            (r'#.*\n', Comment.Single),
-            (r'\\[\w\W]', String.Escape),
-            (r'(\b\w+)(\s*)(\+?=)', bygroups(Name.Variable, Whitespace, Operator)),
-            (r'[\[\]{}()=]', Operator),
-            (r'<<<', Operator),  # here-string
-            (r'<<-?\s*(\'?)\\?(\w+)[\w\W]+?\2', String),
-            (r'&&|\|\|', Operator),
-        ],
-        'data': [
-            (r'(?s)\$?"(\\.|[^"\\$])*"', String.Double),
-            (r'"', String.Double, 'string'),
-            (r"(?s)\$'(\\\\|\\[0-7]+|\\.|[^'\\])*'", String.Single),
-            (r"(?s)'.*?'", String.Single),
-            (r';', Punctuation),
-            (r'&', Punctuation),
-            (r'\|', Punctuation),
-            (r'\s+', Whitespace),
-            (r'\d+\b', Number),
-            (r'[^=\s\[\]{}()$"\'`\\<&|;]+', Text),
-            (r'<', Text),
-        ],
-        'string': [
-            (r'"', String.Double, '#pop'),
-            (r'(?s)(\\\\|\\[0-7]+|\\.|[^"\\$])+', String.Double),
-            include('interp'),
-        ],
-        'curly': [
-            (r'\}', String.Interpol, '#pop'),
-            (r':-', Keyword),
-            (r'\w+', Name.Variable),
-            (r'[^}:"\'`$\\]+', Punctuation),
-            (r':', Punctuation),
-            include('root'),
-        ],
-        'paren': [
-            (r'\)', Keyword, '#pop'),
-            include('root'),
-        ],
-        'math': [
-            (r'\)\)', Keyword, '#pop'),
-            (r'\*\*|\|\||<<|>>|[-+*/%^|&<>]', Operator),
-            (r'\d+#[\da-zA-Z]+', Number),
-            (r'\d+#(?! )', Number),
-            (r'0[xX][\da-fA-F]+', Number),
-            (r'\d+', Number),
-            (r'[a-zA-Z_]\w*', Name.Variable),  # user variable
-            include('root'),
-        ],
-        'backticks': [
-            (r'`', String.Backtick, '#pop'),
-            include('root'),
-        ],
-    }
-
-    def analyse_text(text):
-        if shebang_matches(text, r'(ba|z|)sh'):
-            return 1
-        if text.startswith('$ '):
-            return 0.2
-
-
-class SlurmBashLexer(BashLexer):
-    """
-    Lexer for (ba|k|z|)sh Slurm scripts.
-    """
-
-    name = 'Slurm'
-    aliases = ['slurm', 'sbatch']
-    filenames = ['*.sl']
-    mimetypes = []
-    version_added = '2.4'
-    EXTRA_KEYWORDS = {'srun'}
-
-    def get_tokens_unprocessed(self, text):
-        for index, token, value in BashLexer.get_tokens_unprocessed(self, text):
-            if token is Text and value in self.EXTRA_KEYWORDS:
-                yield index, Name.Builtin, value
-            elif token is Comment.Single and 'SBATCH' in value:
-                yield index, Keyword.Pseudo, value
-            else:
-                yield index, token, value
-
-
-class ShellSessionBaseLexer(Lexer):
-    """
-    Base lexer for shell sessions.
-
-    .. versionadded:: 2.1
-    """
-
-    _bare_continuation = False
-    _venv = re.compile(r'^(\([^)]*\))(\s*)')
-
-    def get_tokens_unprocessed(self, text):
-        innerlexer = self._innerLexerCls(**self.options)
-
-        pos = 0
-        curcode = ''
-        insertions = []
-        backslash_continuation = False
-
-        for match in line_re.finditer(text):
-            line = match.group()
-
-            venv_match = self._venv.match(line)
-            if venv_match:
-                venv = venv_match.group(1)
-                venv_whitespace = venv_match.group(2)
-                insertions.append((len(curcode),
-                                   [(0, Generic.Prompt.VirtualEnv, venv)]))
-                if venv_whitespace:
-                    insertions.append((len(curcode),
-                                       [(0, Text, venv_whitespace)]))
-                line = line[venv_match.end():]
-
-            m = self._ps1rgx.match(line)
-            if m:
-                # To support output lexers (say diff output), the output
-                # needs to be broken by prompts whenever the output lexer
-                # changes.
-                if not insertions:
-                    pos = match.start()
-
-                insertions.append((len(curcode),
-                                   [(0, Generic.Prompt, m.group(1))]))
-                curcode += m.group(2)
-                backslash_continuation = curcode.endswith('\\\n')
-            elif backslash_continuation:
-                if line.startswith(self._ps2):
-                    insertions.append((len(curcode),
-                                       [(0, Generic.Prompt,
-                                         line[:len(self._ps2)])]))
-                    curcode += line[len(self._ps2):]
-                else:
-                    curcode += line
-                backslash_continuation = curcode.endswith('\\\n')
-            elif self._bare_continuation and line.startswith(self._ps2):
-                insertions.append((len(curcode),
-                                   [(0, Generic.Prompt,
-                                     line[:len(self._ps2)])]))
-                curcode += line[len(self._ps2):]
-            else:
-                if insertions:
-                    toks = innerlexer.get_tokens_unprocessed(curcode)
-                    for i, t, v in do_insertions(insertions, toks):
-                        yield pos+i, t, v
-                yield match.start(), Generic.Output, line
-                insertions = []
-                curcode = ''
-        if insertions:
-            for i, t, v in do_insertions(insertions,
-                                         innerlexer.get_tokens_unprocessed(curcode)):
-                yield pos+i, t, v
-
-
-class BashSessionLexer(ShellSessionBaseLexer):
-    """
-    Lexer for Bash shell sessions, i.e. command lines, including a
-    prompt, interspersed with output.
-    """
-
-    name = 'Bash Session'
-    aliases = ['console', 'shell-session']
-    filenames = ['*.sh-session', '*.shell-session']
-    mimetypes = ['application/x-shell-session', 'application/x-sh-session']
-    url = 'https://en.wikipedia.org/wiki/Unix_shell'
-    version_added = '1.1'
-    _example = "console/example.sh-session"
-
-    _innerLexerCls = BashLexer
-    _ps1rgx = re.compile(
-        r'^((?:(?:\[.*?\])|(?:\(\S+\))?(?:| |sh\S*?|\w+\S+[@:]\S+(?:\s+\S+)' \
-        r'?|\[\S+[@:][^\n]+\].+))\s*[$#%]\s*)(.*\n?)')
-    _ps2 = '> '
-
-
-class BatchLexer(RegexLexer):
-    """
-    Lexer for the DOS/Windows Batch file format.
-    """
-    name = 'Batchfile'
-    aliases = ['batch', 'bat', 'dosbatch', 'winbatch']
-    filenames = ['*.bat', '*.cmd']
-    mimetypes = ['application/x-dos-batch']
-    url = 'https://en.wikipedia.org/wiki/Batch_file'
-    version_added = '0.7'
-
-    flags = re.MULTILINE | re.IGNORECASE
-
-    _nl = r'\n\x1a'
-    _punct = r'&<>|'
-    _ws = r'\t\v\f\r ,;=\xa0'
-    _nlws = r'\s\x1a\xa0,;='
-    _space = rf'(?:(?:(?:\^[{_nl}])?[{_ws}])+)'
-    _keyword_terminator = (rf'(?=(?:\^[{_nl}]?)?[{_ws}+./:[\\\]]|[{_nl}{_punct}(])')
-    _token_terminator = rf'(?=\^?[{_ws}]|[{_punct}{_nl}])'
-    _start_label = rf'((?:(?<=^[^:])|^[^:]?)[{_ws}]*)(:)'
-    _label = rf'(?:(?:[^{_nlws}{_punct}+:^]|\^[{_nl}]?[\w\W])*)'
-    _label_compound = rf'(?:(?:[^{_nlws}{_punct}+:^)]|\^[{_nl}]?[^)])*)'
-    _number = rf'(?:-?(?:0[0-7]+|0x[\da-f]+|\d+){_token_terminator})'
-    _opword = r'(?:equ|geq|gtr|leq|lss|neq)'
-    _string = rf'(?:"[^{_nl}"]*(?:"|(?=[{_nl}])))'
-    _variable = (r'(?:(?:%(?:\*|(?:~[a-z]*(?:\$[^:]+:)?)?\d|'
-                 rf'[^%:{_nl}]+(?::(?:~(?:-?\d+)?(?:,(?:-?\d+)?)?|(?:[^%{_nl}^]|'
-                 rf'\^[^%{_nl}])[^={_nl}]*=(?:[^%{_nl}^]|\^[^%{_nl}])*)?)?%))|'
-                 rf'(?:\^?![^!:{_nl}]+(?::(?:~(?:-?\d+)?(?:,(?:-?\d+)?)?|(?:'
-                 rf'[^!{_nl}^]|\^[^!{_nl}])[^={_nl}]*=(?:[^!{_nl}^]|\^[^!{_nl}])*)?)?\^?!))')
-    _core_token = rf'(?:(?:(?:\^[{_nl}]?)?[^"{_nlws}{_punct}])+)'
-    _core_token_compound = rf'(?:(?:(?:\^[{_nl}]?)?[^"{_nlws}{_punct})])+)'
-    _token = rf'(?:[{_punct}]+|{_core_token})'
-    _token_compound = rf'(?:[{_punct}]+|{_core_token_compound})'
-    _stoken = (rf'(?:[{_punct}]+|(?:{_string}|{_variable}|{_core_token})+)')
-
-    def _make_begin_state(compound, _core_token=_core_token,
-                          _core_token_compound=_core_token_compound,
-                          _keyword_terminator=_keyword_terminator,
-                          _nl=_nl, _punct=_punct, _string=_string,
-                          _space=_space, _start_label=_start_label,
-                          _stoken=_stoken, _token_terminator=_token_terminator,
-                          _variable=_variable, _ws=_ws):
-        rest = '(?:{}|{}|[^"%{}{}{}])*'.format(_string, _variable, _nl, _punct,
-                                            ')' if compound else '')
-        rest_of_line = rf'(?:(?:[^{_nl}^]|\^[{_nl}]?[\w\W])*)'
-        rest_of_line_compound = rf'(?:(?:[^{_nl}^)]|\^[{_nl}]?[^)])*)'
-        set_space = rf'((?:(?:\^[{_nl}]?)?[^\S\n])*)'
-        suffix = ''
-        if compound:
-            _keyword_terminator = rf'(?:(?=\))|{_keyword_terminator})'
-            _token_terminator = rf'(?:(?=\))|{_token_terminator})'
-            suffix = '/compound'
-        return [
-            ((r'\)', Punctuation, '#pop') if compound else
-             (rf'\)((?=\()|{_token_terminator}){rest_of_line}',
-              Comment.Single)),
-            (rf'(?={_start_label})', Text, f'follow{suffix}'),
-            (_space, using(this, state='text')),
-            include(f'redirect{suffix}'),
-            (rf'[{_nl}]+', Text),
-            (r'\(', Punctuation, 'root/compound'),
-            (r'@+', Punctuation),
-            (rf'((?:for|if|rem)(?:(?=(?:\^[{_nl}]?)?/)|(?:(?!\^)|'
-             rf'(?<=m))(?:(?=\()|{_token_terminator})))({_space}?{_core_token_compound if compound else _core_token}?(?:\^[{_nl}]?)?/(?:\^[{_nl}]?)?\?)',
-             bygroups(Keyword, using(this, state='text')),
-             f'follow{suffix}'),
-            (rf'(goto{_keyword_terminator})({rest}(?:\^[{_nl}]?)?/(?:\^[{_nl}]?)?\?{rest})',
-             bygroups(Keyword, using(this, state='text')),
-             f'follow{suffix}'),
-            (words(('assoc', 'break', 'cd', 'chdir', 'cls', 'color', 'copy',
-                    'date', 'del', 'dir', 'dpath', 'echo', 'endlocal', 'erase',
-                    'exit', 'ftype', 'keys', 'md', 'mkdir', 'mklink', 'move',
-                    'path', 'pause', 'popd', 'prompt', 'pushd', 'rd', 'ren',
-                    'rename', 'rmdir', 'setlocal', 'shift', 'start', 'time',
-                    'title', 'type', 'ver', 'verify', 'vol'),
-                   suffix=_keyword_terminator), Keyword, f'follow{suffix}'),
-            (rf'(call)({_space}?)(:)',
-             bygroups(Keyword, using(this, state='text'), Punctuation),
-             f'call{suffix}'),
-            (rf'call{_keyword_terminator}', Keyword),
-            (rf'(for{_token_terminator}(?!\^))({_space})(/f{_token_terminator})',
-             bygroups(Keyword, using(this, state='text'), Keyword),
-             ('for/f', 'for')),
-            (rf'(for{_token_terminator}(?!\^))({_space})(/l{_token_terminator})',
-             bygroups(Keyword, using(this, state='text'), Keyword),
-             ('for/l', 'for')),
-            (rf'for{_token_terminator}(?!\^)', Keyword, ('for2', 'for')),
-            (rf'(goto{_keyword_terminator})({_space}?)(:?)',
-             bygroups(Keyword, using(this, state='text'), Punctuation),
-             f'label{suffix}'),
-            (rf'(if(?:(?=\()|{_token_terminator})(?!\^))({_space}?)((?:/i{_token_terminator})?)({_space}?)((?:not{_token_terminator})?)({_space}?)',
-             bygroups(Keyword, using(this, state='text'), Keyword,
-                      using(this, state='text'), Keyword,
-                      using(this, state='text')), ('(?', 'if')),
-            (rf'rem(((?=\()|{_token_terminator}){_space}?{_stoken}?.*|{_keyword_terminator}{rest_of_line_compound if compound else rest_of_line})',
-             Comment.Single, f'follow{suffix}'),
-            (rf'(set{_keyword_terminator}){set_space}(/a)',
-             bygroups(Keyword, using(this, state='text'), Keyword),
-             f'arithmetic{suffix}'),
-            (r'(set{}){}((?:/p)?){}((?:(?:(?:\^[{}]?)?[^"{}{}^={}]|'
-             r'\^[{}]?[^"=])+)?)((?:(?:\^[{}]?)?=)?)'.format(_keyword_terminator, set_space, set_space, _nl, _nl, _punct,
-              ')' if compound else '', _nl, _nl),
-             bygroups(Keyword, using(this, state='text'), Keyword,
-                      using(this, state='text'), using(this, state='variable'),
-                      Punctuation),
-             f'follow{suffix}'),
-            default(f'follow{suffix}')
-        ]
-
-    def _make_follow_state(compound, _label=_label,
-                           _label_compound=_label_compound, _nl=_nl,
-                           _space=_space, _start_label=_start_label,
-                           _token=_token, _token_compound=_token_compound,
-                           _ws=_ws):
-        suffix = '/compound' if compound else ''
-        state = []
-        if compound:
-            state.append((r'(?=\))', Text, '#pop'))
-        state += [
-            (rf'{_start_label}([{_ws}]*)({_label_compound if compound else _label})(.*)',
-             bygroups(Text, Punctuation, Text, Name.Label, Comment.Single)),
-            include(f'redirect{suffix}'),
-            (rf'(?=[{_nl}])', Text, '#pop'),
-            (r'\|\|?|&&?', Punctuation, '#pop'),
-            include('text')
-        ]
-        return state
-
-    def _make_arithmetic_state(compound, _nl=_nl, _punct=_punct,
-                               _string=_string, _variable=_variable,
-                               _ws=_ws, _nlws=_nlws):
-        op = r'=+\-*/!~'
-        state = []
-        if compound:
-            state.append((r'(?=\))', Text, '#pop'))
-        state += [
-            (r'0[0-7]+', Number.Oct),
-            (r'0x[\da-f]+', Number.Hex),
-            (r'\d+', Number.Integer),
-            (r'[(),]+', Punctuation),
-            (rf'([{op}]|%|\^\^)+', Operator),
-            (r'({}|{}|(\^[{}]?)?[^(){}%\^"{}{}]|\^[{}]?{})+'.format(_string, _variable, _nl, op, _nlws, _punct, _nlws,
-              r'[^)]' if compound else r'[\w\W]'),
-             using(this, state='variable')),
-            (r'(?=[\x00|&])', Text, '#pop'),
-            include('follow')
-        ]
-        return state
-
-    def _make_call_state(compound, _label=_label,
-                         _label_compound=_label_compound):
-        state = []
-        if compound:
-            state.append((r'(?=\))', Text, '#pop'))
-        state.append((r'(:?)(%s)' % (_label_compound if compound else _label),
-                      bygroups(Punctuation, Name.Label), '#pop'))
-        return state
-
-    def _make_label_state(compound, _label=_label,
-                          _label_compound=_label_compound, _nl=_nl,
-                          _punct=_punct, _string=_string, _variable=_variable):
-        state = []
-        if compound:
-            state.append((r'(?=\))', Text, '#pop'))
-        state.append((r'({}?)((?:{}|{}|\^[{}]?{}|[^"%^{}{}{}])*)'.format(_label_compound if compound else _label, _string,
-                       _variable, _nl, r'[^)]' if compound else r'[\w\W]', _nl,
-                       _punct, r')' if compound else ''),
-                      bygroups(Name.Label, Comment.Single), '#pop'))
-        return state
-
-    def _make_redirect_state(compound,
-                             _core_token_compound=_core_token_compound,
-                             _nl=_nl, _punct=_punct, _stoken=_stoken,
-                             _string=_string, _space=_space,
-                             _variable=_variable, _nlws=_nlws):
-        stoken_compound = (rf'(?:[{_punct}]+|(?:{_string}|{_variable}|{_core_token_compound})+)')
-        return [
-            (rf'((?:(?<=[{_nlws}])\d)?)(>>?&|<&)([{_nlws}]*)(\d)',
-             bygroups(Number.Integer, Punctuation, Text, Number.Integer)),
-            (rf'((?:(?<=[{_nlws}])(?>?|<)({_space}?{stoken_compound if compound else _stoken})',
-             bygroups(Number.Integer, Punctuation, using(this, state='text')))
-        ]
-
-    tokens = {
-        'root': _make_begin_state(False),
-        'follow': _make_follow_state(False),
-        'arithmetic': _make_arithmetic_state(False),
-        'call': _make_call_state(False),
-        'label': _make_label_state(False),
-        'redirect': _make_redirect_state(False),
-        'root/compound': _make_begin_state(True),
-        'follow/compound': _make_follow_state(True),
-        'arithmetic/compound': _make_arithmetic_state(True),
-        'call/compound': _make_call_state(True),
-        'label/compound': _make_label_state(True),
-        'redirect/compound': _make_redirect_state(True),
-        'variable-or-escape': [
-            (_variable, Name.Variable),
-            (rf'%%|\^[{_nl}]?(\^!|[\w\W])', String.Escape)
-        ],
-        'string': [
-            (r'"', String.Double, '#pop'),
-            (_variable, Name.Variable),
-            (r'\^!|%%', String.Escape),
-            (rf'[^"%^{_nl}]+|[%^]', String.Double),
-            default('#pop')
-        ],
-        'sqstring': [
-            include('variable-or-escape'),
-            (r'[^%]+|%', String.Single)
-        ],
-        'bqstring': [
-            include('variable-or-escape'),
-            (r'[^%]+|%', String.Backtick)
-        ],
-        'text': [
-            (r'"', String.Double, 'string'),
-            include('variable-or-escape'),
-            (rf'[^"%^{_nlws}{_punct}\d)]+|.', Text)
-        ],
-        'variable': [
-            (r'"', String.Double, 'string'),
-            include('variable-or-escape'),
-            (rf'[^"%^{_nl}]+|.', Name.Variable)
-        ],
-        'for': [
-            (rf'({_space})(in)({_space})(\()',
-             bygroups(using(this, state='text'), Keyword,
-                      using(this, state='text'), Punctuation), '#pop'),
-            include('follow')
-        ],
-        'for2': [
-            (r'\)', Punctuation),
-            (rf'({_space})(do{_token_terminator})',
-             bygroups(using(this, state='text'), Keyword), '#pop'),
-            (rf'[{_nl}]+', Text),
-            include('follow')
-        ],
-        'for/f': [
-            (rf'(")((?:{_variable}|[^"])*?")([{_nlws}]*)(\))',
-             bygroups(String.Double, using(this, state='string'), Text,
-                      Punctuation)),
-            (r'"', String.Double, ('#pop', 'for2', 'string')),
-            (rf"('(?:%%|{_variable}|[\w\W])*?')([{_nlws}]*)(\))",
-             bygroups(using(this, state='sqstring'), Text, Punctuation)),
-            (rf'(`(?:%%|{_variable}|[\w\W])*?`)([{_nlws}]*)(\))',
-             bygroups(using(this, state='bqstring'), Text, Punctuation)),
-            include('for2')
-        ],
-        'for/l': [
-            (r'-?\d+', Number.Integer),
-            include('for2')
-        ],
-        'if': [
-            (rf'((?:cmdextversion|errorlevel){_token_terminator})({_space})(\d+)',
-             bygroups(Keyword, using(this, state='text'),
-                      Number.Integer), '#pop'),
-            (rf'(defined{_token_terminator})({_space})({_stoken})',
-             bygroups(Keyword, using(this, state='text'),
-                      using(this, state='variable')), '#pop'),
-            (rf'(exist{_token_terminator})({_space}{_stoken})',
-             bygroups(Keyword, using(this, state='text')), '#pop'),
-            (rf'({_number}{_space})({_opword})({_space}{_number})',
-             bygroups(using(this, state='arithmetic'), Operator.Word,
-                      using(this, state='arithmetic')), '#pop'),
-            (_stoken, using(this, state='text'), ('#pop', 'if2')),
-        ],
-        'if2': [
-            (rf'({_space}?)(==)({_space}?{_stoken})',
-             bygroups(using(this, state='text'), Operator,
-                      using(this, state='text')), '#pop'),
-            (rf'({_space})({_opword})({_space}{_stoken})',
-             bygroups(using(this, state='text'), Operator.Word,
-                      using(this, state='text')), '#pop')
-        ],
-        '(?': [
-            (_space, using(this, state='text')),
-            (r'\(', Punctuation, ('#pop', 'else?', 'root/compound')),
-            default('#pop')
-        ],
-        'else?': [
-            (_space, using(this, state='text')),
-            (rf'else{_token_terminator}', Keyword, '#pop'),
-            default('#pop')
-        ]
-    }
-
-
-class MSDOSSessionLexer(ShellSessionBaseLexer):
-    """
-    Lexer for MS DOS shell sessions, i.e. command lines, including a
-    prompt, interspersed with output.
-    """
-
-    name = 'MSDOS Session'
-    aliases = ['doscon']
-    filenames = []
-    mimetypes = []
-    url = 'https://en.wikipedia.org/wiki/MS-DOS'
-    version_added = '2.1'
-    _example = "doscon/session"
-
-    _innerLexerCls = BatchLexer
-    _ps1rgx = re.compile(r'^([^>]*>)(.*\n?)')
-    _ps2 = 'More? '
-
-
-class TcshLexer(RegexLexer):
-    """
-    Lexer for tcsh scripts.
-    """
-
-    name = 'Tcsh'
-    aliases = ['tcsh', 'csh']
-    filenames = ['*.tcsh', '*.csh']
-    mimetypes = ['application/x-csh']
-    url = 'https://www.tcsh.org'
-    version_added = '0.10'
-
-    tokens = {
-        'root': [
-            include('basic'),
-            (r'\$\(', Keyword, 'paren'),
-            (r'\$\{#?', Keyword, 'curly'),
-            (r'`', String.Backtick, 'backticks'),
-            include('data'),
-        ],
-        'basic': [
-            (r'\b(if|endif|else|while|then|foreach|case|default|'
-             r'break|continue|goto|breaksw|end|switch|endsw)\s*\b',
-             Keyword),
-            (r'\b(alias|alloc|bg|bindkey|builtins|bye|caller|cd|chdir|'
-             r'complete|dirs|echo|echotc|eval|exec|exit|fg|filetest|getxvers|'
-             r'glob|getspath|hashstat|history|hup|inlib|jobs|kill|'
-             r'limit|log|login|logout|ls-F|migrate|newgrp|nice|nohup|notify|'
-             r'onintr|popd|printenv|pushd|rehash|repeat|rootnode|popd|pushd|'
-             r'set|shift|sched|setenv|setpath|settc|setty|setxvers|shift|'
-             r'source|stop|suspend|source|suspend|telltc|time|'
-             r'umask|unalias|uncomplete|unhash|universe|unlimit|unset|unsetenv|'
-             r'ver|wait|warp|watchlog|where|which)\s*\b',
-             Name.Builtin),
-            (r'#.*', Comment),
-            (r'\\[\w\W]', String.Escape),
-            (r'(\b\w+)(\s*)(=)', bygroups(Name.Variable, Text, Operator)),
-            (r'[\[\]{}()=]+', Operator),
-            (r'<<\s*(\'?)\\?(\w+)[\w\W]+?\2', String),
-            (r';', Punctuation),
-        ],
-        'data': [
-            (r'(?s)"(\\\\|\\[0-7]+|\\.|[^"\\])*"', String.Double),
-            (r"(?s)'(\\\\|\\[0-7]+|\\.|[^'\\])*'", String.Single),
-            (r'\s+', Text),
-            (r'[^=\s\[\]{}()$"\'`\\;#]+', Text),
-            (r'\d+(?= |\Z)', Number),
-            (r'\$#?(\w+|.)', Name.Variable),
-        ],
-        'curly': [
-            (r'\}', Keyword, '#pop'),
-            (r':-', Keyword),
-            (r'\w+', Name.Variable),
-            (r'[^}:"\'`$]+', Punctuation),
-            (r':', Punctuation),
-            include('root'),
-        ],
-        'paren': [
-            (r'\)', Keyword, '#pop'),
-            include('root'),
-        ],
-        'backticks': [
-            (r'`', String.Backtick, '#pop'),
-            include('root'),
-        ],
-    }
-
-
-class TcshSessionLexer(ShellSessionBaseLexer):
-    """
-    Lexer for Tcsh sessions, i.e. command lines, including a
-    prompt, interspersed with output.
-    """
-
-    name = 'Tcsh Session'
-    aliases = ['tcshcon']
-    filenames = []
-    mimetypes = []
-    url = 'https://www.tcsh.org'
-    version_added = '2.1'
-    _example = "tcshcon/session"
-
-    _innerLexerCls = TcshLexer
-    _ps1rgx = re.compile(r'^([^>]+>)(.*\n?)')
-    _ps2 = '? '
-
-
-class PowerShellLexer(RegexLexer):
-    """
-    For Windows PowerShell code.
-    """
-    name = 'PowerShell'
-    aliases = ['powershell', 'pwsh', 'posh', 'ps1', 'psm1']
-    filenames = ['*.ps1', '*.psm1']
-    mimetypes = ['text/x-powershell']
-    url = 'https://learn.microsoft.com/en-us/powershell'
-    version_added = '1.5'
-
-    flags = re.DOTALL | re.IGNORECASE | re.MULTILINE
-
-    keywords = (
-        'while validateset validaterange validatepattern validatelength '
-        'validatecount until trap switch return ref process param parameter in '
-        'if global: local: function foreach for finally filter end elseif else '
-        'dynamicparam do default continue cmdletbinding break begin alias \\? '
-        '% #script #private #local #global mandatory parametersetname position '
-        'valuefrompipeline valuefrompipelinebypropertyname '
-        'valuefromremainingarguments helpmessage try catch throw').split()
-
-    operators = (
-        'and as band bnot bor bxor casesensitive ccontains ceq cge cgt cle '
-        'clike clt cmatch cne cnotcontains cnotlike cnotmatch contains '
-        'creplace eq exact f file ge gt icontains ieq ige igt ile ilike ilt '
-        'imatch ine inotcontains inotlike inotmatch ireplace is isnot le like '
-        'lt match ne not notcontains notlike notmatch or regex replace '
-        'wildcard').split()
-
-    verbs = (
-        'write where watch wait use update unregister unpublish unprotect '
-        'unlock uninstall undo unblock trace test tee take sync switch '
-        'suspend submit stop step start split sort skip show set send select '
-        'search scroll save revoke resume restore restart resolve resize '
-        'reset request repair rename remove register redo receive read push '
-        'publish protect pop ping out optimize open new move mount merge '
-        'measure lock limit join invoke install initialize import hide group '
-        'grant get format foreach find export expand exit enter enable edit '
-        'dismount disconnect disable deny debug cxnew copy convertto '
-        'convertfrom convert connect confirm compress complete compare close '
-        'clear checkpoint block backup assert approve aggregate add').split()
-
-    aliases_ = (
-        'ac asnp cat cd cfs chdir clc clear clhy cli clp cls clv cnsn '
-        'compare copy cp cpi cpp curl cvpa dbp del diff dir dnsn ebp echo epal '
-        'epcsv epsn erase etsn exsn fc fhx fl foreach ft fw gal gbp gc gci gcm '
-        'gcs gdr ghy gi gjb gl gm gmo gp gps gpv group gsn gsnp gsv gu gv gwmi '
-        'h history icm iex ihy ii ipal ipcsv ipmo ipsn irm ise iwmi iwr kill lp '
-        'ls man md measure mi mount move mp mv nal ndr ni nmo npssc nsn nv ogv '
-        'oh popd ps pushd pwd r rbp rcjb rcsn rd rdr ren ri rjb rm rmdir rmo '
-        'rni rnp rp rsn rsnp rujb rv rvpa rwmi sajb sal saps sasv sbp sc select '
-        'set shcm si sl sleep sls sort sp spjb spps spsv start sujb sv swmi tee '
-        'trcm type wget where wjb write').split()
-
-    commenthelp = (
-        'component description example externalhelp forwardhelpcategory '
-        'forwardhelptargetname functionality inputs link '
-        'notes outputs parameter remotehelprunspace role synopsis').split()
-
-    tokens = {
-        'root': [
-            # we need to count pairs of parentheses for correct highlight
-            # of '$(...)' blocks in strings
-            (r'\(', Punctuation, 'child'),
-            (r'\s+', Text),
-            (r'^(\s*#[#\s]*)(\.(?:{}))([^\n]*$)'.format('|'.join(commenthelp)),
-             bygroups(Comment, String.Doc, Comment)),
-            (r'#[^\n]*?$', Comment),
-            (r'(<|<)#', Comment.Multiline, 'multline'),
-            (r'@"\n', String.Heredoc, 'heredoc-double'),
-            (r"@'\n.*?\n'@", String.Heredoc),
-            # escaped syntax
-            (r'`[\'"$@-]', Punctuation),
-            (r'"', String.Double, 'string'),
-            (r"'([^']|'')*'", String.Single),
-            (r'(\$|@@|@)((global|script|private|env):)?\w+',
-             Name.Variable),
-            (r'({})\b'.format('|'.join(keywords)), Keyword),
-            (r'-({})\b'.format('|'.join(operators)), Operator),
-            (r'({})-[a-z_]\w*\b'.format('|'.join(verbs)), Name.Builtin),
-            (r'({})\s'.format('|'.join(aliases_)), Name.Builtin),
-            (r'\[[a-z_\[][\w. `,\[\]]*\]', Name.Constant),  # .net [type]s
-            (r'-[a-z_]\w*', Name),
-            (r'\w+', Name),
-            (r'[.,;:@{}\[\]$()=+*/\\&%!~?^`|<>-]', Punctuation),
-        ],
-        'child': [
-            (r'\)', Punctuation, '#pop'),
-            include('root'),
-        ],
-        'multline': [
-            (r'[^#&.]+', Comment.Multiline),
-            (r'#(>|>)', Comment.Multiline, '#pop'),
-            (r'\.({})'.format('|'.join(commenthelp)), String.Doc),
-            (r'[#&.]', Comment.Multiline),
-        ],
-        'string': [
-            (r"`[0abfnrtv'\"$`]", String.Escape),
-            (r'[^$`"]+', String.Double),
-            (r'\$\(', Punctuation, 'child'),
-            (r'""', String.Double),
-            (r'[`$]', String.Double),
-            (r'"', String.Double, '#pop'),
-        ],
-        'heredoc-double': [
-            (r'\n"@', String.Heredoc, '#pop'),
-            (r'\$\(', Punctuation, 'child'),
-            (r'[^@\n]+"]', String.Heredoc),
-            (r".", String.Heredoc),
-        ]
-    }
-
-
-class PowerShellSessionLexer(ShellSessionBaseLexer):
-    """
-    Lexer for PowerShell sessions, i.e. command lines, including a
-    prompt, interspersed with output.
-    """
-
-    name = 'PowerShell Session'
-    aliases = ['pwsh-session', 'ps1con']
-    filenames = []
-    mimetypes = []
-    url = 'https://learn.microsoft.com/en-us/powershell'
-    version_added = '2.1'
-    _example = "pwsh-session/session"
-
-    _innerLexerCls = PowerShellLexer
-    _bare_continuation = True
-    _ps1rgx = re.compile(r'^((?:\[[^]]+\]: )?PS[^>]*> ?)(.*\n?)')
-    _ps2 = '> '
-
-
-class FishShellLexer(RegexLexer):
-    """
-    Lexer for Fish shell scripts.
-    """
-
-    name = 'Fish'
-    aliases = ['fish', 'fishshell']
-    filenames = ['*.fish', '*.load']
-    mimetypes = ['application/x-fish']
-    url = 'https://fishshell.com'
-    version_added = '2.1'
-
-    tokens = {
-        'root': [
-            include('basic'),
-            include('data'),
-            include('interp'),
-        ],
-        'interp': [
-            (r'\$\(\(', Keyword, 'math'),
-            (r'\(', Keyword, 'paren'),
-            (r'\$#?(\w+|.)', Name.Variable),
-        ],
-        'basic': [
-            (r'\b(begin|end|if|else|while|break|for|in|return|function|block|'
-             r'case|continue|switch|not|and|or|set|echo|exit|pwd|true|false|'
-             r'cd|count|test)(\s*)\b',
-             bygroups(Keyword, Text)),
-            (r'\b(alias|bg|bind|breakpoint|builtin|command|commandline|'
-             r'complete|contains|dirh|dirs|emit|eval|exec|fg|fish|fish_config|'
-             r'fish_indent|fish_pager|fish_prompt|fish_right_prompt|'
-             r'fish_update_completions|fishd|funced|funcsave|functions|help|'
-             r'history|isatty|jobs|math|mimedb|nextd|open|popd|prevd|psub|'
-             r'pushd|random|read|set_color|source|status|trap|type|ulimit|'
-             r'umask|vared|fc|getopts|hash|kill|printf|time|wait)\s*\b(?!\.)',
-             Name.Builtin),
-            (r'#.*\n', Comment),
-            (r'\\[\w\W]', String.Escape),
-            (r'(\b\w+)(\s*)(=)', bygroups(Name.Variable, Whitespace, Operator)),
-            (r'[\[\]()=]', Operator),
-            (r'<<-?\s*(\'?)\\?(\w+)[\w\W]+?\2', String),
-        ],
-        'data': [
-            (r'(?s)\$?"(\\\\|\\[0-7]+|\\.|[^"\\$])*"', String.Double),
-            (r'"', String.Double, 'string'),
-            (r"(?s)\$'(\\\\|\\[0-7]+|\\.|[^'\\])*'", String.Single),
-            (r"(?s)'.*?'", String.Single),
-            (r';', Punctuation),
-            (r'&|\||\^|<|>', Operator),
-            (r'\s+', Text),
-            (r'\d+(?= |\Z)', Number),
-            (r'[^=\s\[\]{}()$"\'`\\<&|;]+', Text),
-        ],
-        'string': [
-            (r'"', String.Double, '#pop'),
-            (r'(?s)(\\\\|\\[0-7]+|\\.|[^"\\$])+', String.Double),
-            include('interp'),
-        ],
-        'paren': [
-            (r'\)', Keyword, '#pop'),
-            include('root'),
-        ],
-        'math': [
-            (r'\)\)', Keyword, '#pop'),
-            (r'[-+*/%^|&]|\*\*|\|\|', Operator),
-            (r'\d+#\d+', Number),
-            (r'\d+#(?! )', Number),
-            (r'\d+', Number),
-            include('root'),
-        ],
-    }
-
-class ExeclineLexer(RegexLexer):
-    """
-    Lexer for Laurent Bercot's execline language.
-    """
-
-    name = 'execline'
-    aliases = ['execline']
-    filenames = ['*.exec']
-    url = 'https://skarnet.org/software/execline'
-    version_added = '2.7'
-
-    tokens = {
-        'root': [
-            include('basic'),
-            include('data'),
-            include('interp')
-        ],
-        'interp': [
-            (r'\$\{', String.Interpol, 'curly'),
-            (r'\$[\w@#]+', Name.Variable),  # user variable
-            (r'\$', Text),
-        ],
-        'basic': [
-            (r'\b(background|backtick|cd|define|dollarat|elgetopt|'
-             r'elgetpositionals|elglob|emptyenv|envfile|exec|execlineb|'
-             r'exit|export|fdblock|fdclose|fdmove|fdreserve|fdswap|'
-             r'forbacktickx|foreground|forstdin|forx|getcwd|getpid|heredoc|'
-             r'homeof|if|ifelse|ifte|ifthenelse|importas|loopwhilex|'
-             r'multidefine|multisubstitute|pipeline|piperw|posix-cd|'
-             r'redirfd|runblock|shift|trap|tryexec|umask|unexport|wait|'
-             r'withstdinas)\b', Name.Builtin),
-            (r'\A#!.+\n', Comment.Hashbang),
-            (r'#.*\n', Comment.Single),
-            (r'[{}]', Operator)
-        ],
-        'data': [
-            (r'(?s)"(\\.|[^"\\$])*"', String.Double),
-            (r'"', String.Double, 'string'),
-            (r'\s+', Text),
-            (r'[^\s{}$"\\]+', Text)
-        ],
-        'string': [
-            (r'"', String.Double, '#pop'),
-            (r'(?s)(\\\\|\\.|[^"\\$])+', String.Double),
-            include('interp'),
-        ],
-        'curly': [
-            (r'\}', String.Interpol, '#pop'),
-            (r'[\w#@]+', Name.Variable),
-            include('root')
-        ]
-
-    }
-
-    def analyse_text(text):
-        if shebang_matches(text, r'execlineb'):
-            return 1
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/sieve.py b/venv/lib/python3.10/site-packages/pygments/lexers/sieve.py
deleted file mode 100644
index fc48980c49921d811b5bed122c531cc8213be463..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/sieve.py
+++ /dev/null
@@ -1,78 +0,0 @@
-"""
-    pygments.lexers.sieve
-    ~~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for Sieve file format.
-
-    https://tools.ietf.org/html/rfc5228
-    https://tools.ietf.org/html/rfc5173
-    https://tools.ietf.org/html/rfc5229
-    https://tools.ietf.org/html/rfc5230
-    https://tools.ietf.org/html/rfc5232
-    https://tools.ietf.org/html/rfc5235
-    https://tools.ietf.org/html/rfc5429
-    https://tools.ietf.org/html/rfc8580
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, bygroups
-from pygments.token import Comment, Name, Literal, String, Text, Punctuation, \
-    Keyword
-
-__all__ = ["SieveLexer"]
-
-
-class SieveLexer(RegexLexer):
-    """
-    Lexer for sieve format.
-    """
-    name = 'Sieve'
-    filenames = ['*.siv', '*.sieve']
-    aliases = ['sieve']
-    url = 'https://en.wikipedia.org/wiki/Sieve_(mail_filtering_language)'
-    version_added = '2.6'
-
-    tokens = {
-        'root': [
-            (r'\s+', Text),
-            (r'[();,{}\[\]]', Punctuation),
-            # import:
-            (r'(?i)require',
-             Keyword.Namespace),
-            # tags:
-            (r'(?i)(:)(addresses|all|contains|content|create|copy|comparator|'
-             r'count|days|detail|domain|fcc|flags|from|handle|importance|is|'
-             r'localpart|length|lowerfirst|lower|matches|message|mime|options|'
-             r'over|percent|quotewildcard|raw|regex|specialuse|subject|text|'
-             r'under|upperfirst|upper|value)',
-             bygroups(Name.Tag, Name.Tag)),
-            # tokens:
-            (r'(?i)(address|addflag|allof|anyof|body|discard|elsif|else|envelope|'
-             r'ereject|exists|false|fileinto|if|hasflag|header|keep|'
-             r'notify_method_capability|notify|not|redirect|reject|removeflag|'
-             r'setflag|size|spamtest|stop|string|true|vacation|virustest)',
-             Name.Builtin),
-            (r'(?i)set',
-             Keyword.Declaration),
-            # number:
-            (r'([0-9.]+)([kmgKMG])?',
-             bygroups(Literal.Number, Literal.Number)),
-            # comment:
-            (r'#.*$',
-             Comment.Single),
-            (r'/\*.*\*/',
-             Comment.Multiline),
-            # string:
-            (r'"[^"]*?"',
-             String),
-            # text block:
-            (r'text:',
-             Name.Tag, 'text'),
-        ],
-        'text': [
-            (r'[^.].*?\n', String),
-            (r'^\.', Punctuation, "#pop"),
-        ]
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/slash.py b/venv/lib/python3.10/site-packages/pygments/lexers/slash.py
deleted file mode 100644
index 1c439d0db0fd51d6870f1294b2eb276aa64e8a6c..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/slash.py
+++ /dev/null
@@ -1,183 +0,0 @@
-"""
-    pygments.lexers.slash
-    ~~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for the Slash programming language.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import ExtendedRegexLexer, bygroups, DelegatingLexer
-from pygments.token import Name, Number, String, Comment, Punctuation, \
-    Other, Keyword, Operator, Whitespace
-
-__all__ = ['SlashLexer']
-
-
-class SlashLanguageLexer(ExtendedRegexLexer):
-    _nkw = r'(?=[^a-zA-Z_0-9])'
-
-    def move_state(new_state):
-        return ("#pop", new_state)
-
-    def right_angle_bracket(lexer, match, ctx):
-        if len(ctx.stack) > 1 and ctx.stack[-2] == "string":
-            ctx.stack.pop()
-        yield match.start(), String.Interpol, '}'
-        ctx.pos = match.end()
-        pass
-
-    tokens = {
-        "root": [
-            (r"<%=",        Comment.Preproc,    move_state("slash")),
-            (r"<%!!",       Comment.Preproc,    move_state("slash")),
-            (r"<%#.*?%>",   Comment.Multiline),
-            (r"<%",         Comment.Preproc,    move_state("slash")),
-            (r".|\n",       Other),
-        ],
-        "string": [
-            (r"\\",         String.Escape,      move_state("string_e")),
-            (r"\"",         String,             move_state("slash")),
-            (r"#\{",        String.Interpol,    "slash"),
-            (r'.|\n',       String),
-        ],
-        "string_e": [
-            (r'n',                  String.Escape,      move_state("string")),
-            (r't',                  String.Escape,      move_state("string")),
-            (r'r',                  String.Escape,      move_state("string")),
-            (r'e',                  String.Escape,      move_state("string")),
-            (r'x[a-fA-F0-9]{2}',    String.Escape,      move_state("string")),
-            (r'.',                  String.Escape,      move_state("string")),
-        ],
-        "regexp": [
-            (r'}[a-z]*',            String.Regex,       move_state("slash")),
-            (r'\\(.|\n)',           String.Regex),
-            (r'{',                  String.Regex,       "regexp_r"),
-            (r'.|\n',               String.Regex),
-        ],
-        "regexp_r": [
-            (r'}[a-z]*',            String.Regex,       "#pop"),
-            (r'\\(.|\n)',           String.Regex),
-            (r'{',                  String.Regex,       "regexp_r"),
-        ],
-        "slash": [
-            (r"%>",                     Comment.Preproc,    move_state("root")),
-            (r"\"",                     String,             move_state("string")),
-            (r"'[a-zA-Z0-9_]+",         String),
-            (r'%r{',                    String.Regex,       move_state("regexp")),
-            (r'/\*.*?\*/',              Comment.Multiline),
-            (r"(#|//).*?\n",            Comment.Single),
-            (r'-?[0-9]+e[+-]?[0-9]+',   Number.Float),
-            (r'-?[0-9]+\.[0-9]+(e[+-]?[0-9]+)?', Number.Float),
-            (r'-?[0-9]+',               Number.Integer),
-            (r'nil'+_nkw,               Name.Builtin),
-            (r'true'+_nkw,              Name.Builtin),
-            (r'false'+_nkw,             Name.Builtin),
-            (r'self'+_nkw,              Name.Builtin),
-            (r'(class)(\s+)([A-Z][a-zA-Z0-9_\']*)',
-                bygroups(Keyword, Whitespace, Name.Class)),
-            (r'class'+_nkw,             Keyword),
-            (r'extends'+_nkw,           Keyword),
-            (r'(def)(\s+)(self)(\s*)(\.)(\s*)([a-z_][a-zA-Z0-9_\']*=?|<<|>>|==|<=>|<=|<|>=|>|\+|-(self)?|~(self)?|\*|/|%|^|&&|&|\||\[\]=?)',
-                bygroups(Keyword, Whitespace, Name.Builtin, Whitespace, Punctuation, Whitespace, Name.Function)),
-            (r'(def)(\s+)([a-z_][a-zA-Z0-9_\']*=?|<<|>>|==|<=>|<=|<|>=|>|\+|-(self)?|~(self)?|\*|/|%|^|&&|&|\||\[\]=?)',
-                bygroups(Keyword, Whitespace, Name.Function)),
-            (r'def'+_nkw,               Keyword),
-            (r'if'+_nkw,                Keyword),
-            (r'elsif'+_nkw,             Keyword),
-            (r'else'+_nkw,              Keyword),
-            (r'unless'+_nkw,            Keyword),
-            (r'for'+_nkw,               Keyword),
-            (r'in'+_nkw,                Keyword),
-            (r'while'+_nkw,             Keyword),
-            (r'until'+_nkw,             Keyword),
-            (r'and'+_nkw,               Keyword),
-            (r'or'+_nkw,                Keyword),
-            (r'not'+_nkw,               Keyword),
-            (r'lambda'+_nkw,            Keyword),
-            (r'try'+_nkw,               Keyword),
-            (r'catch'+_nkw,             Keyword),
-            (r'return'+_nkw,            Keyword),
-            (r'next'+_nkw,              Keyword),
-            (r'last'+_nkw,              Keyword),
-            (r'throw'+_nkw,             Keyword),
-            (r'use'+_nkw,               Keyword),
-            (r'switch'+_nkw,            Keyword),
-            (r'\\',                     Keyword),
-            (r'λ',                      Keyword),
-            (r'__FILE__'+_nkw,          Name.Builtin.Pseudo),
-            (r'__LINE__'+_nkw,          Name.Builtin.Pseudo),
-            (r'[A-Z][a-zA-Z0-9_\']*'+_nkw, Name.Constant),
-            (r'[a-z_][a-zA-Z0-9_\']*'+_nkw, Name),
-            (r'@[a-z_][a-zA-Z0-9_\']*'+_nkw, Name.Variable.Instance),
-            (r'@@[a-z_][a-zA-Z0-9_\']*'+_nkw, Name.Variable.Class),
-            (r'\(',                     Punctuation),
-            (r'\)',                     Punctuation),
-            (r'\[',                     Punctuation),
-            (r'\]',                     Punctuation),
-            (r'\{',                     Punctuation),
-            (r'\}',                     right_angle_bracket),
-            (r';',                      Punctuation),
-            (r',',                      Punctuation),
-            (r'<<=',                    Operator),
-            (r'>>=',                    Operator),
-            (r'<<',                     Operator),
-            (r'>>',                     Operator),
-            (r'==',                     Operator),
-            (r'!=',                     Operator),
-            (r'=>',                     Operator),
-            (r'=',                      Operator),
-            (r'<=>',                    Operator),
-            (r'<=',                     Operator),
-            (r'>=',                     Operator),
-            (r'<',                      Operator),
-            (r'>',                      Operator),
-            (r'\+\+',                   Operator),
-            (r'\+=',                    Operator),
-            (r'-=',                     Operator),
-            (r'\*\*=',                  Operator),
-            (r'\*=',                    Operator),
-            (r'\*\*',                   Operator),
-            (r'\*',                     Operator),
-            (r'/=',                     Operator),
-            (r'\+',                     Operator),
-            (r'-',                      Operator),
-            (r'/',                      Operator),
-            (r'%=',                     Operator),
-            (r'%',                      Operator),
-            (r'^=',                     Operator),
-            (r'&&=',                    Operator),
-            (r'&=',                     Operator),
-            (r'&&',                     Operator),
-            (r'&',                      Operator),
-            (r'\|\|=',                  Operator),
-            (r'\|=',                    Operator),
-            (r'\|\|',                   Operator),
-            (r'\|',                     Operator),
-            (r'!',                      Operator),
-            (r'\.\.\.',                 Operator),
-            (r'\.\.',                   Operator),
-            (r'\.',                     Operator),
-            (r'::',                     Operator),
-            (r':',                      Operator),
-            (r'(\s|\n)+',               Whitespace),
-            (r'[a-z_][a-zA-Z0-9_\']*',  Name.Variable),
-        ],
-    }
-
-
-class SlashLexer(DelegatingLexer):
-    """
-    Lexer for the Slash programming language.
-    """
-
-    name = 'Slash'
-    aliases = ['slash']
-    filenames = ['*.sla']
-    url = 'https://github.com/arturadib/Slash-A'
-    version_added = '2.4'
-
-    def __init__(self, **options):
-        from pygments.lexers.web import HtmlLexer
-        super().__init__(HtmlLexer, SlashLanguageLexer, **options)
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/smalltalk.py b/venv/lib/python3.10/site-packages/pygments/lexers/smalltalk.py
deleted file mode 100644
index 674b7b4b345a4b968f080bcd1552979201da6164..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/smalltalk.py
+++ /dev/null
@@ -1,194 +0,0 @@
-"""
-    pygments.lexers.smalltalk
-    ~~~~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Smalltalk and related languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, include, bygroups, default
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation
-
-__all__ = ['SmalltalkLexer', 'NewspeakLexer']
-
-
-class SmalltalkLexer(RegexLexer):
-    """
-    For Smalltalk syntax.
-    Contributed by Stefan Matthias Aust.
-    Rewritten by Nils Winter.
-    """
-    name = 'Smalltalk'
-    url = 'http://www.smalltalk.org/'
-    filenames = ['*.st']
-    aliases = ['smalltalk', 'squeak', 'st']
-    mimetypes = ['text/x-smalltalk']
-    version_added = '0.10'
-
-    tokens = {
-        'root': [
-            (r'(<)(\w+:)(.*?)(>)', bygroups(Text, Keyword, Text, Text)),
-            include('squeak fileout'),
-            include('whitespaces'),
-            include('method definition'),
-            (r'(\|)([\w\s]*)(\|)', bygroups(Operator, Name.Variable, Operator)),
-            include('objects'),
-            (r'\^|\:=|\_', Operator),
-            # temporaries
-            (r'[\]({}.;!]', Text),
-        ],
-        'method definition': [
-            # Not perfect can't allow whitespaces at the beginning and the
-            # without breaking everything
-            (r'([a-zA-Z]+\w*:)(\s*)(\w+)',
-             bygroups(Name.Function, Text, Name.Variable)),
-            (r'^(\b[a-zA-Z]+\w*\b)(\s*)$', bygroups(Name.Function, Text)),
-            (r'^([-+*/\\~<>=|&!?,@%]+)(\s*)(\w+)(\s*)$',
-             bygroups(Name.Function, Text, Name.Variable, Text)),
-        ],
-        'blockvariables': [
-            include('whitespaces'),
-            (r'(:)(\s*)(\w+)',
-             bygroups(Operator, Text, Name.Variable)),
-            (r'\|', Operator, '#pop'),
-            default('#pop'),  # else pop
-        ],
-        'literals': [
-            (r"'(''|[^'])*'", String, 'afterobject'),
-            (r'\$.', String.Char, 'afterobject'),
-            (r'#\(', String.Symbol, 'parenth'),
-            (r'\)', Text, 'afterobject'),
-            (r'(\d+r)?-?\d+(\.\d+)?(e-?\d+)?', Number, 'afterobject'),
-        ],
-        '_parenth_helper': [
-            include('whitespaces'),
-            (r'(\d+r)?-?\d+(\.\d+)?(e-?\d+)?', Number),
-            (r'[-+*/\\~<>=|&#!?,@%\w:]+', String.Symbol),
-            # literals
-            (r"'(''|[^'])*'", String),
-            (r'\$.', String.Char),
-            (r'#*\(', String.Symbol, 'inner_parenth'),
-        ],
-        'parenth': [
-            # This state is a bit tricky since
-            # we can't just pop this state
-            (r'\)', String.Symbol, ('root', 'afterobject')),
-            include('_parenth_helper'),
-        ],
-        'inner_parenth': [
-            (r'\)', String.Symbol, '#pop'),
-            include('_parenth_helper'),
-        ],
-        'whitespaces': [
-            # skip whitespace and comments
-            (r'\s+', Text),
-            (r'"(""|[^"])*"', Comment),
-        ],
-        'objects': [
-            (r'\[', Text, 'blockvariables'),
-            (r'\]', Text, 'afterobject'),
-            (r'\b(self|super|true|false|nil|thisContext)\b',
-             Name.Builtin.Pseudo, 'afterobject'),
-            (r'\b[A-Z]\w*(?!:)\b', Name.Class, 'afterobject'),
-            (r'\b[a-z]\w*(?!:)\b', Name.Variable, 'afterobject'),
-            (r'#("(""|[^"])*"|[-+*/\\~<>=|&!?,@%]+|[\w:]+)',
-             String.Symbol, 'afterobject'),
-            include('literals'),
-        ],
-        'afterobject': [
-            (r'! !$', Keyword, '#pop'),  # squeak chunk delimiter
-            include('whitespaces'),
-            (r'\b(ifTrue:|ifFalse:|whileTrue:|whileFalse:|timesRepeat:)',
-             Name.Builtin, '#pop'),
-            (r'\b(new\b(?!:))', Name.Builtin),
-            (r'\:=|\_', Operator, '#pop'),
-            (r'\b[a-zA-Z]+\w*:', Name.Function, '#pop'),
-            (r'\b[a-zA-Z]+\w*', Name.Function),
-            (r'\w+:?|[-+*/\\~<>=|&!?,@%]+', Name.Function, '#pop'),
-            (r'\.', Punctuation, '#pop'),
-            (r';', Punctuation),
-            (r'[\])}]', Text),
-            (r'[\[({]', Text, '#pop'),
-        ],
-        'squeak fileout': [
-            # Squeak fileout format (optional)
-            (r'^"(""|[^"])*"!', Keyword),
-            (r"^'(''|[^'])*'!", Keyword),
-            (r'^(!)(\w+)( commentStamp: )(.*?)( prior: .*?!\n)(.*?)(!)',
-                bygroups(Keyword, Name.Class, Keyword, String, Keyword, Text, Keyword)),
-            (r"^(!)(\w+(?: class)?)( methodsFor: )('(?:''|[^'])*')(.*?!)",
-                bygroups(Keyword, Name.Class, Keyword, String, Keyword)),
-            (r'^(\w+)( subclass: )(#\w+)'
-             r'(\s+instanceVariableNames: )(.*?)'
-             r'(\s+classVariableNames: )(.*?)'
-             r'(\s+poolDictionaries: )(.*?)'
-             r'(\s+category: )(.*?)(!)',
-                bygroups(Name.Class, Keyword, String.Symbol, Keyword, String, Keyword,
-                         String, Keyword, String, Keyword, String, Keyword)),
-            (r'^(\w+(?: class)?)(\s+instanceVariableNames: )(.*?)(!)',
-                bygroups(Name.Class, Keyword, String, Keyword)),
-            (r'(!\n)(\].*)(! !)$', bygroups(Keyword, Text, Keyword)),
-            (r'! !$', Keyword),
-        ],
-    }
-
-
-class NewspeakLexer(RegexLexer):
-    """
-    For Newspeak syntax.
-    """
-    name = 'Newspeak'
-    url = 'http://newspeaklanguage.org/'
-    filenames = ['*.ns2']
-    aliases = ['newspeak', ]
-    mimetypes = ['text/x-newspeak']
-    version_added = '1.1'
-
-    tokens = {
-        'root': [
-            (r'\b(Newsqueak2)\b', Keyword.Declaration),
-            (r"'[^']*'", String),
-            (r'\b(class)(\s+)(\w+)(\s*)',
-             bygroups(Keyword.Declaration, Text, Name.Class, Text)),
-            (r'\b(mixin|self|super|private|public|protected|nil|true|false)\b',
-             Keyword),
-            (r'(\w+\:)(\s*)([a-zA-Z_]\w+)',
-             bygroups(Name.Function, Text, Name.Variable)),
-            (r'(\w+)(\s*)(=)',
-             bygroups(Name.Attribute, Text, Operator)),
-            (r'<\w+>', Comment.Special),
-            include('expressionstat'),
-            include('whitespace')
-        ],
-
-        'expressionstat': [
-            (r'(\d+\.\d*|\.\d+|\d+[fF])[fF]?', Number.Float),
-            (r'\d+', Number.Integer),
-            (r':\w+', Name.Variable),
-            (r'(\w+)(::)', bygroups(Name.Variable, Operator)),
-            (r'\w+:', Name.Function),
-            (r'\w+', Name.Variable),
-            (r'\(|\)', Punctuation),
-            (r'\[|\]', Punctuation),
-            (r'\{|\}', Punctuation),
-
-            (r'(\^|\+|\/|~|\*|<|>|=|@|%|\||&|\?|!|,|-|:)', Operator),
-            (r'\.|;', Punctuation),
-            include('whitespace'),
-            include('literals'),
-        ],
-        'literals': [
-            (r'\$.', String),
-            (r"'[^']*'", String),
-            (r"#'[^']*'", String.Symbol),
-            (r"#\w+:?", String.Symbol),
-            (r"#(\+|\/|~|\*|<|>|=|@|%|\||&|\?|!|,|-)+", String.Symbol)
-        ],
-        'whitespace': [
-            (r'\s+', Text),
-            (r'"[^"]*"', Comment)
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/smithy.py b/venv/lib/python3.10/site-packages/pygments/lexers/smithy.py
deleted file mode 100644
index bd479aec40c596032921b5cd24f249052058e04e..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/smithy.py
+++ /dev/null
@@ -1,77 +0,0 @@
-"""
-    pygments.lexers.smithy
-    ~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for the Smithy IDL.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, bygroups, words
-from pygments.token import Text, Comment, Keyword, Name, String, \
-    Number, Whitespace, Punctuation
-
-__all__ = ['SmithyLexer']
-
-
-class SmithyLexer(RegexLexer):
-    """
-    For Smithy IDL
-    """
-    name = 'Smithy'
-    url = 'https://awslabs.github.io/smithy/'
-    filenames = ['*.smithy']
-    aliases = ['smithy']
-    version_added = '2.10'
-
-    unquoted = r'[A-Za-z0-9_\.#$-]+'
-    identifier = r"[A-Za-z0-9_\.#$-]+"
-
-    simple_shapes = (
-        'use', 'byte', 'short', 'integer', 'long', 'float', 'document',
-        'double', 'bigInteger', 'bigDecimal', 'boolean', 'blob', 'string',
-        'timestamp',
-    )
-
-    aggregate_shapes = (
-       'apply', 'list', 'map', 'set', 'structure', 'union', 'resource',
-       'operation', 'service', 'trait'
-    )
-
-    tokens = {
-        'root': [
-            (r'///.*$', Comment.Multiline),
-            (r'//.*$', Comment),
-            (r'@[0-9a-zA-Z\.#-]*', Name.Decorator),
-            (r'(=)', Name.Decorator),
-            (r'^(\$version)(:)(.+)',
-                bygroups(Keyword.Declaration, Name.Decorator, Name.Class)),
-            (r'^(namespace)(\s+' + identifier + r')\b',
-                bygroups(Keyword.Declaration, Name.Class)),
-            (words(simple_shapes,
-                   prefix=r'^', suffix=r'(\s+' + identifier + r')\b'),
-                bygroups(Keyword.Declaration, Name.Class)),
-            (words(aggregate_shapes,
-                   prefix=r'^', suffix=r'(\s+' + identifier + r')'),
-                bygroups(Keyword.Declaration, Name.Class)),
-            (r'^(metadata)(\s+)((?:\S+)|(?:\"[^"]+\"))(\s*)(=)',
-                bygroups(Keyword.Declaration, Whitespace, Name.Class,
-                         Whitespace, Name.Decorator)),
-            (r"(true|false|null)", Keyword.Constant),
-            (r"(-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?)", Number),
-            (identifier + ":", Name.Label),
-            (identifier, Name.Variable.Class),
-            (r'\[', Text, "#push"),
-            (r'\]', Text, "#pop"),
-            (r'\(', Text, "#push"),
-            (r'\)', Text, "#pop"),
-            (r'\{', Text, "#push"),
-            (r'\}', Text, "#pop"),
-            (r'"{3}(\\\\|\n|\\")*"{3}', String.Doc),
-            (r'"(\\\\|\n|\\"|[^"])*"', String.Double),
-            (r"'(\\\\|\n|\\'|[^'])*'", String.Single),
-            (r'[:,]+', Punctuation),
-            (r'\s+', Whitespace),
-        ]
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/smv.py b/venv/lib/python3.10/site-packages/pygments/lexers/smv.py
deleted file mode 100644
index bf97b52a5ce1953909b43b66255f0324fd8f051e..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/smv.py
+++ /dev/null
@@ -1,78 +0,0 @@
-"""
-    pygments.lexers.smv
-    ~~~~~~~~~~~~~~~~~~~
-
-    Lexers for the SMV languages.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, words
-from pygments.token import Comment, Keyword, Name, Number, Operator, \
-    Punctuation, Text
-
-__all__ = ['NuSMVLexer']
-
-
-class NuSMVLexer(RegexLexer):
-    """
-    Lexer for the NuSMV language.
-    """
-
-    name = 'NuSMV'
-    aliases = ['nusmv']
-    filenames = ['*.smv']
-    mimetypes = []
-    url = 'https://nusmv.fbk.eu'
-    version_added = '2.2'
-
-    tokens = {
-        'root': [
-            # Comments
-            (r'(?s)\/\-\-.*?\-\-/', Comment),
-            (r'--.*\n', Comment),
-
-            # Reserved
-            (words(('MODULE', 'DEFINE', 'MDEFINE', 'CONSTANTS', 'VAR', 'IVAR',
-                    'FROZENVAR', 'INIT', 'TRANS', 'INVAR', 'SPEC', 'CTLSPEC',
-                    'LTLSPEC', 'PSLSPEC', 'COMPUTE', 'NAME', 'INVARSPEC',
-                    'FAIRNESS', 'JUSTICE', 'COMPASSION', 'ISA', 'ASSIGN',
-                    'CONSTRAINT', 'SIMPWFF', 'CTLWFF', 'LTLWFF', 'PSLWFF',
-                    'COMPWFF', 'IN', 'MIN', 'MAX', 'MIRROR', 'PRED',
-                    'PREDICATES'), suffix=r'(?![\w$#-])'),
-             Keyword.Declaration),
-            (r'process(?![\w$#-])', Keyword),
-            (words(('array', 'of', 'boolean', 'integer', 'real', 'word'),
-                   suffix=r'(?![\w$#-])'), Keyword.Type),
-            (words(('case', 'esac'), suffix=r'(?![\w$#-])'), Keyword),
-            (words(('word1', 'bool', 'signed', 'unsigned', 'extend', 'resize',
-                    'sizeof', 'uwconst', 'swconst', 'init', 'self', 'count',
-                    'abs', 'max', 'min'), suffix=r'(?![\w$#-])'),
-             Name.Builtin),
-            (words(('EX', 'AX', 'EF', 'AF', 'EG', 'AG', 'E', 'F', 'O', 'G',
-                    'H', 'X', 'Y', 'Z', 'A', 'U', 'S', 'V', 'T', 'BU', 'EBF',
-                    'ABF', 'EBG', 'ABG', 'next', 'mod', 'union', 'in', 'xor',
-                    'xnor'), suffix=r'(?![\w$#-])'),
-                Operator.Word),
-            (words(('TRUE', 'FALSE'), suffix=r'(?![\w$#-])'), Keyword.Constant),
-
-            # Names
-            (r'[a-zA-Z_][\w$#-]*', Name.Variable),
-
-            # Operators
-            (r':=', Operator),
-            (r'[-&|+*/<>!=]', Operator),
-
-            # Literals
-            (r'\-?\d+\b', Number.Integer),
-            (r'0[su][bB]\d*_[01_]+', Number.Bin),
-            (r'0[su][oO]\d*_[0-7_]+', Number.Oct),
-            (r'0[su][dD]\d*_[\d_]+', Number.Decimal),
-            (r'0[su][hH]\d*_[\da-fA-F_]+', Number.Hex),
-
-            # Whitespace, punctuation and the rest
-            (r'\s+', Text.Whitespace),
-            (r'[()\[\]{};?:.,]', Punctuation),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/snobol.py b/venv/lib/python3.10/site-packages/pygments/lexers/snobol.py
deleted file mode 100644
index bab51e9b1107c5e295febaa842c68ac93ee1b1fc..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/snobol.py
+++ /dev/null
@@ -1,82 +0,0 @@
-"""
-    pygments.lexers.snobol
-    ~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for the SNOBOL language.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, bygroups
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation
-
-__all__ = ['SnobolLexer']
-
-
-class SnobolLexer(RegexLexer):
-    """
-    Lexer for the SNOBOL4 programming language.
-
-    Recognizes the common ASCII equivalents of the original SNOBOL4 operators.
-    Does not require spaces around binary operators.
-    """
-
-    name = "Snobol"
-    aliases = ["snobol"]
-    filenames = ['*.snobol']
-    mimetypes = ['text/x-snobol']
-    url = 'https://www.regressive.org/snobol4'
-    version_added = '1.5'
-
-    tokens = {
-        # root state, start of line
-        # comments, continuation lines, and directives start in column 1
-        # as do labels
-        'root': [
-            (r'\*.*\n', Comment),
-            (r'[+.] ', Punctuation, 'statement'),
-            (r'-.*\n', Comment),
-            (r'END\s*\n', Name.Label, 'heredoc'),
-            (r'[A-Za-z$][\w$]*', Name.Label, 'statement'),
-            (r'\s+', Text, 'statement'),
-        ],
-        # statement state, line after continuation or label
-        'statement': [
-            (r'\s*\n', Text, '#pop'),
-            (r'\s+', Text),
-            (r'(?<=[^\w.])(LT|LE|EQ|NE|GE|GT|INTEGER|IDENT|DIFFER|LGT|SIZE|'
-             r'REPLACE|TRIM|DUPL|REMDR|DATE|TIME|EVAL|APPLY|OPSYN|LOAD|UNLOAD|'
-             r'LEN|SPAN|BREAK|ANY|NOTANY|TAB|RTAB|REM|POS|RPOS|FAIL|FENCE|'
-             r'ABORT|ARB|ARBNO|BAL|SUCCEED|INPUT|OUTPUT|TERMINAL)(?=[^\w.])',
-             Name.Builtin),
-            (r'[A-Za-z][\w.]*', Name),
-            # ASCII equivalents of original operators
-            # | for the EBCDIC equivalent, ! likewise
-            # \ for EBCDIC negation
-            (r'\*\*|[?$.!%*/#+\-@|&\\=]', Operator),
-            (r'"[^"]*"', String),
-            (r"'[^']*'", String),
-            # Accept SPITBOL syntax for real numbers
-            # as well as Macro SNOBOL4
-            (r'[0-9]+(?=[^.EeDd])', Number.Integer),
-            (r'[0-9]+(\.[0-9]*)?([EDed][-+]?[0-9]+)?', Number.Float),
-            # Goto
-            (r':', Punctuation, 'goto'),
-            (r'[()<>,;]', Punctuation),
-        ],
-        # Goto block
-        'goto': [
-            (r'\s*\n', Text, "#pop:2"),
-            (r'\s+', Text),
-            (r'F|S', Keyword),
-            (r'(\()([A-Za-z][\w.]*)(\))',
-             bygroups(Punctuation, Name.Label, Punctuation))
-        ],
-        # everything after the END statement is basically one
-        # big heredoc.
-        'heredoc': [
-            (r'.*\n', String.Heredoc)
-        ]
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/solidity.py b/venv/lib/python3.10/site-packages/pygments/lexers/solidity.py
deleted file mode 100644
index 3182a148a3830c6b55e6b4f194010121dfe0abc4..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/solidity.py
+++ /dev/null
@@ -1,87 +0,0 @@
-"""
-    pygments.lexers.solidity
-    ~~~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Solidity.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, bygroups, include, words
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Whitespace
-
-__all__ = ['SolidityLexer']
-
-
-class SolidityLexer(RegexLexer):
-    """
-    For Solidity source code.
-    """
-
-    name = 'Solidity'
-    aliases = ['solidity']
-    filenames = ['*.sol']
-    mimetypes = []
-    url = 'https://soliditylang.org'
-    version_added = '2.5'
-
-    datatype = (
-        r'\b(address|bool|(?:(?:bytes|hash|int|string|uint)(?:8|16|24|32|40|48|56|64'
-        r'|72|80|88|96|104|112|120|128|136|144|152|160|168|176|184|192|200|208'
-        r'|216|224|232|240|248|256)?))\b'
-    )
-
-    tokens = {
-        'root': [
-            include('whitespace'),
-            include('comments'),
-            (r'\bpragma\s+solidity\b', Keyword, 'pragma'),
-            (r'\b(contract)(\s+)([a-zA-Z_]\w*)',
-             bygroups(Keyword, Whitespace, Name.Entity)),
-            (datatype + r'(\s+)((?:external|public|internal|private)\s+)?' +
-             r'([a-zA-Z_]\w*)',
-             bygroups(Keyword.Type, Whitespace, Keyword, Name.Variable)),
-            (r'\b(enum|event|function|struct)(\s+)([a-zA-Z_]\w*)',
-             bygroups(Keyword.Type, Whitespace, Name.Variable)),
-            (r'\b(msg|block|tx)\.([A-Za-z_][a-zA-Z0-9_]*)\b', Keyword),
-            (words((
-                'block', 'break', 'constant', 'constructor', 'continue',
-                'contract', 'do', 'else', 'external', 'false', 'for',
-                'function', 'if', 'import', 'inherited', 'internal', 'is',
-                'library', 'mapping', 'memory', 'modifier', 'msg', 'new',
-                'payable', 'private', 'public', 'require', 'return',
-                'returns', 'struct', 'suicide', 'throw', 'this', 'true',
-                'tx', 'var', 'while'), prefix=r'\b', suffix=r'\b'),
-             Keyword.Type),
-            (words(('keccak256',), prefix=r'\b', suffix=r'\b'), Name.Builtin),
-            (datatype, Keyword.Type),
-            include('constants'),
-            (r'[a-zA-Z_]\w*', Text),
-            (r'[~!%^&*+=|?:<>/-]', Operator),
-            (r'[.;{}(),\[\]]', Punctuation)
-        ],
-        'comments': [
-            (r'//(\n|[\w\W]*?[^\\]\n)', Comment.Single),
-            (r'/(\\\n)?[*][\w\W]*?[*](\\\n)?/', Comment.Multiline),
-            (r'/(\\\n)?[*][\w\W]*', Comment.Multiline)
-        ],
-        'constants': [
-            (r'("(\\"|.)*?")', String.Double),
-            (r"('(\\'|.)*?')", String.Single),
-            (r'\b0[xX][0-9a-fA-F]+\b', Number.Hex),
-            (r'\b\d+\b', Number.Decimal),
-        ],
-        'pragma': [
-            include('whitespace'),
-            include('comments'),
-            (r'(\^|>=|<)(\s*)(\d+\.\d+\.\d+)',
-             bygroups(Operator, Whitespace, Keyword)),
-            (r';', Punctuation, '#pop')
-        ],
-        'whitespace': [
-            (r'\s+', Whitespace),
-            (r'\n', Whitespace)
-        ]
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/soong.py b/venv/lib/python3.10/site-packages/pygments/lexers/soong.py
deleted file mode 100644
index bbf204dd22543f06c65849d54613c3a573dda899..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/soong.py
+++ /dev/null
@@ -1,78 +0,0 @@
-"""
-    pygments.lexers.soong
-    ~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for Soong (Android.bp Blueprint) files.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, bygroups, include
-from pygments.token import Comment, Name, Number, Operator, Punctuation, \
-        String, Whitespace
-
-__all__ = ['SoongLexer']
-
-class SoongLexer(RegexLexer):
-    name = 'Soong'
-    version_added = '2.18'
-    url = 'https://source.android.com/docs/setup/reference/androidbp'
-    aliases = ['androidbp', 'bp', 'soong']
-    filenames = ['Android.bp']
-
-    tokens = {
-        'root': [
-            # A variable assignment
-            (r'(\w*)(\s*)(\+?=)(\s*)',
-             bygroups(Name.Variable, Whitespace, Operator, Whitespace),
-             'assign-rhs'),
-
-            # A top-level module
-            (r'(\w*)(\s*)(\{)',
-             bygroups(Name.Function, Whitespace, Punctuation),
-             'in-rule'),
-
-            # Everything else
-            include('comments'),
-            (r'\s+', Whitespace),  # newlines okay
-        ],
-        'assign-rhs': [
-            include('expr'),
-            (r'\n', Whitespace, '#pop'),
-        ],
-        'in-list': [
-            include('expr'),
-            include('comments'),
-            (r'\s+', Whitespace),  # newlines okay in a list
-            (r',', Punctuation),
-            (r'\]', Punctuation, '#pop'),
-        ],
-        'in-map': [
-            # A map key
-            (r'(\w+)(:)(\s*)', bygroups(Name, Punctuation, Whitespace)),
-
-            include('expr'),
-            include('comments'),
-            (r'\s+', Whitespace),  # newlines okay in a map
-            (r',', Punctuation),
-            (r'\}', Punctuation, '#pop'),
-        ],
-        'in-rule': [
-            # Just re-use map syntax
-            include('in-map'),
-        ],
-        'comments': [
-            (r'//.*', Comment.Single),
-            (r'/(\\\n)?[*](.|\n)*?[*](\\\n)?/', Comment.Multiline),
-        ],
-        'expr': [
-            (r'(true|false)\b', Name.Builtin),
-            (r'0x[0-9a-fA-F]+', Number.Hex),
-            (r'\d+', Number.Integer),
-            (r'".*?"', String),
-            (r'\{', Punctuation, 'in-map'),
-            (r'\[', Punctuation, 'in-list'),
-            (r'\w+', Name),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/sophia.py b/venv/lib/python3.10/site-packages/pygments/lexers/sophia.py
deleted file mode 100644
index 37fcec5c39bfe4e8f52bf735618971684ab27739..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/sophia.py
+++ /dev/null
@@ -1,102 +0,0 @@
-"""
-    pygments.lexers.sophia
-    ~~~~~~~~~~~~~~~~~~~~~~
-
-    Lexer for Sophia.
-
-    Derived from pygments/lexers/reason.py.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, include, default, words
-from pygments.token import Comment, Keyword, Name, Number, Operator, \
-    Punctuation, String, Text
-
-__all__ = ['SophiaLexer']
-
-class SophiaLexer(RegexLexer):
-    """
-    A Sophia lexer.
-    """
-
-    name = 'Sophia'
-    aliases = ['sophia']
-    filenames = ['*.aes']
-    mimetypes = []
-    url = 'https://docs.aeternity.com/aesophia'
-    version_added = '2.11'
-
-    keywords = (
-        'contract', 'include', 'let', 'switch', 'type', 'record', 'datatype',
-        'if', 'elif', 'else', 'function', 'stateful', 'payable', 'public',
-        'entrypoint', 'private', 'indexed', 'namespace', 'interface', 'main',
-        'using', 'as', 'for', 'hiding',
-    )
-
-    builtins = ('state', 'put', 'abort', 'require')
-
-    word_operators = ('mod', 'band', 'bor', 'bxor', 'bnot')
-
-    primitive_types = ('int', 'address', 'bool', 'bits', 'bytes', 'string',
-                       'list', 'option', 'char', 'unit', 'map', 'event',
-                       'hash', 'signature', 'oracle', 'oracle_query')
-
-    tokens = {
-        'escape-sequence': [
-            (r'\\[\\"\'ntbr]', String.Escape),
-            (r'\\[0-9]{3}', String.Escape),
-            (r'\\x[0-9a-fA-F]{2}', String.Escape),
-        ],
-        'root': [
-            (r'\s+', Text.Whitespace),
-            (r'(true|false)\b', Keyword.Constant),
-            (r'\b([A-Z][\w\']*)(?=\s*\.)', Name.Class, 'dotted'),
-            (r'\b([A-Z][\w\']*)', Name.Function),
-            (r'//.*?\n', Comment.Single),
-            (r'\/\*(?!/)', Comment.Multiline, 'comment'),
-
-            (r'0[xX][\da-fA-F][\da-fA-F_]*', Number.Hex),
-            (r'#[\da-fA-F][\da-fA-F_]*', Name.Label),
-            (r'\d[\d_]*', Number.Integer),
-
-            (words(keywords, suffix=r'\b'), Keyword),
-            (words(builtins, suffix=r'\b'), Name.Builtin),
-            (words(word_operators, prefix=r'\b', suffix=r'\b'), Operator.Word),
-            (words(primitive_types, prefix=r'\b', suffix=r'\b'), Keyword.Type),
-
-            (r'[=!<>+\\*/:&|?~@^-]', Operator.Word),
-            (r'[.;:{}(),\[\]]', Punctuation),
-
-            (r"(ak_|ok_|oq_|ct_)[\w']*", Name.Label),
-            (r"[^\W\d][\w']*", Name),
-
-            (r"'(?:(\\[\\\"'ntbr ])|(\\[0-9]{3})|(\\x[0-9a-fA-F]{2}))'",
-             String.Char),
-            (r"'.'", String.Char),
-            (r"'[a-z][\w]*", Name.Variable),
-
-            (r'"', String.Double, 'string')
-        ],
-        'comment': [
-            (r'[^/*]+', Comment.Multiline),
-            (r'\/\*', Comment.Multiline, '#push'),
-            (r'\*\/', Comment.Multiline, '#pop'),
-            (r'\*', Comment.Multiline),
-        ],
-        'string': [
-            (r'[^\\"]+', String.Double),
-            include('escape-sequence'),
-            (r'\\\n', String.Double),
-            (r'"', String.Double, '#pop'),
-        ],
-        'dotted': [
-            (r'\s+', Text),
-            (r'\.', Punctuation),
-            (r'[A-Z][\w\']*(?=\s*\.)', Name.Function),
-            (r'[A-Z][\w\']*', Name.Function, '#pop'),
-            (r'[a-z_][\w\']*', Name, '#pop'),
-            default('#pop'),
-        ],
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/special.py b/venv/lib/python3.10/site-packages/pygments/lexers/special.py
deleted file mode 100644
index 524946fc310a6acef08db483ac7f12f09a5e8706..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/special.py
+++ /dev/null
@@ -1,122 +0,0 @@
-"""
-    pygments.lexers.special
-    ~~~~~~~~~~~~~~~~~~~~~~~
-
-    Special lexers.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-import ast
-
-from pygments.lexer import Lexer, line_re
-from pygments.token import Token, Error, Text, Generic
-from pygments.util import get_choice_opt
-
-
-__all__ = ['TextLexer', 'OutputLexer', 'RawTokenLexer']
-
-
-class TextLexer(Lexer):
-    """
-    "Null" lexer, doesn't highlight anything.
-    """
-    name = 'Text only'
-    aliases = ['text']
-    filenames = ['*.txt']
-    mimetypes = ['text/plain']
-    url = ""
-    version_added = ''
-
-    priority = 0.01
-
-    def get_tokens_unprocessed(self, text):
-        yield 0, Text, text
-
-    def analyse_text(text):
-        return TextLexer.priority
-
-
-class OutputLexer(Lexer):
-    """
-    Simple lexer that highlights everything as ``Token.Generic.Output``.
-    """
-    name = 'Text output'
-    aliases = ['output']
-    url = ""
-    version_added = '2.10'
-    _example = "output/output"
-
-    def get_tokens_unprocessed(self, text):
-        yield 0, Generic.Output, text
-
-
-_ttype_cache = {}
-
-
-class RawTokenLexer(Lexer):
-    """
-    Recreate a token stream formatted with the `RawTokenFormatter`.
-
-    Additional options accepted:
-
-    `compress`
-        If set to ``"gz"`` or ``"bz2"``, decompress the token stream with
-        the given compression algorithm before lexing (default: ``""``).
-    """
-    name = 'Raw token data'
-    aliases = []
-    filenames = []
-    mimetypes = ['application/x-pygments-tokens']
-    url = 'https://pygments.org/docs/formatters/#RawTokenFormatter'
-    version_added = ''
-
-    def __init__(self, **options):
-        self.compress = get_choice_opt(options, 'compress',
-                                       ['', 'none', 'gz', 'bz2'], '')
-        Lexer.__init__(self, **options)
-
-    def get_tokens(self, text):
-        if self.compress:
-            if isinstance(text, str):
-                text = text.encode('latin1')
-            try:
-                if self.compress == 'gz':
-                    import gzip
-                    text = gzip.decompress(text)
-                elif self.compress == 'bz2':
-                    import bz2
-                    text = bz2.decompress(text)
-            except OSError:
-                yield Error, text.decode('latin1')
-        if isinstance(text, bytes):
-            text = text.decode('latin1')
-
-        # do not call Lexer.get_tokens() because stripping is not optional.
-        text = text.strip('\n') + '\n'
-        for i, t, v in self.get_tokens_unprocessed(text):
-            yield t, v
-
-    def get_tokens_unprocessed(self, text):
-        length = 0
-        for match in line_re.finditer(text):
-            try:
-                ttypestr, val = match.group().rstrip().split('\t', 1)
-                ttype = _ttype_cache.get(ttypestr)
-                if not ttype:
-                    ttype = Token
-                    ttypes = ttypestr.split('.')[1:]
-                    for ttype_ in ttypes:
-                        if not ttype_ or not ttype_[0].isupper():
-                            raise ValueError('malformed token name')
-                        ttype = getattr(ttype, ttype_)
-                    _ttype_cache[ttypestr] = ttype
-                val = ast.literal_eval(val)
-                if not isinstance(val, str):
-                    raise ValueError('expected str')
-            except (SyntaxError, ValueError):
-                val = match.group()
-                ttype = Error
-            yield length, ttype, val
-            length += len(val)
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/spice.py b/venv/lib/python3.10/site-packages/pygments/lexers/spice.py
deleted file mode 100644
index 9d2b1a1a81a47b9ac830b204df3ce015b8f66985..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/spice.py
+++ /dev/null
@@ -1,70 +0,0 @@
-"""
-    pygments.lexers.spice
-    ~~~~~~~~~~~~~~~~~~~~~
-
-    Lexers for the Spice programming language.
-
-    :copyright: Copyright 2006-2025 by the Pygments team, see AUTHORS.
-    :license: BSD, see LICENSE for details.
-"""
-
-from pygments.lexer import RegexLexer, bygroups, words
-from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
-    Number, Punctuation, Whitespace
-
-__all__ = ['SpiceLexer']
-
-
-class SpiceLexer(RegexLexer):
-    """
-    For Spice source.
-    """
-    name = 'Spice'
-    url = 'https://www.spicelang.com'
-    filenames = ['*.spice']
-    aliases = ['spice', 'spicelang']
-    mimetypes = ['text/x-spice']
-    version_added = '2.11'
-
-    tokens = {
-        'root': [
-            (r'\n', Whitespace),
-            (r'\s+', Whitespace),
-            (r'\\\n', Text),
-            # comments
-            (r'//(.*?)\n', Comment.Single),
-            (r'/(\\\n)?[*]{2}(.|\n)*?[*](\\\n)?/', String.Doc),
-            (r'/(\\\n)?[*](.|\n)*?[*](\\\n)?/', Comment.Multiline),
-            # keywords
-            (r'(import|as)\b', Keyword.Namespace),
-            (r'(f|p|type|struct|interface|enum|alias|operator)\b', Keyword.Declaration),
-            (words(('if', 'else', 'switch', 'case', 'default', 'for', 'foreach', 'do',
-                    'while', 'break', 'continue', 'fallthrough', 'return', 'assert',
-                    'unsafe', 'ext'), suffix=r'\b'), Keyword),
-            (words(('const', 'signed', 'unsigned', 'inline', 'public', 'heap', 'compose'),
-                   suffix=r'\b'), Keyword.Pseudo),
-            (words(('new', 'yield', 'stash', 'pick', 'sync', 'class'), suffix=r'\b'),
-                   Keyword.Reserved),
-            (r'(true|false|nil)\b', Keyword.Constant),
-            (words(('double', 'int', 'short', 'long', 'byte', 'char', 'string',
-                    'bool', 'dyn'), suffix=r'\b'), Keyword.Type),
-            (words(('printf', 'sizeof', 'alignof', 'len', 'panic'), suffix=r'\b(\()'),
-             bygroups(Name.Builtin, Punctuation)),
-            # numeric literals
-            (r'[-]?[0-9]*[.][0-9]+([eE][+-]?[0-9]+)?', Number.Double),
-            (r'0[bB][01]+[slu]?', Number.Bin),
-            (r'0[oO][0-7]+[slu]?', Number.Oct),
-            (r'0[xXhH][0-9a-fA-F]+[slu]?', Number.Hex),
-            (r'(0[dD])?[0-9]+[slu]?', Number.Integer),
-            # string literal
-            (r'"(\\\\|\\[^\\]|[^"\\])*"', String),
-            # char literal
-            (r'\'(\\\\|\\[^\\]|[^\'\\])\'', String.Char),
-            # tokens
-            (r'<<=|>>=|<<|>>|<=|>=|\+=|-=|\*=|/=|\%=|\|=|&=|\^=|&&|\|\||&|\||'
-             r'\+\+|--|\%|\^|\~|==|!=|->|::|[.]{3}|#!|#|[+\-*/&]', Operator),
-            (r'[|<>=!()\[\]{}.,;:\?]', Punctuation),
-            # identifiers
-            (r'[^\W\d]\w*', Name.Other),
-        ]
-    }
diff --git a/venv/lib/python3.10/site-packages/pygments/lexers/sql.py b/venv/lib/python3.10/site-packages/pygments/lexers/sql.py
deleted file mode 100644
index d3e6f17f39491b43cc621ca6981ff33ca90c380f..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/pygments/lexers/sql.py
+++ /dev/null
@@ -1,1109 +0,0 @@
-"""
-    pygments.lexers.sql
-    ~~~~~~~~~~~~~~~~~~~
-
-    Lexers for various SQL dialects and related interactive sessions.
-
-    Postgres specific lexers:
-
-    `PostgresLexer`
-        A SQL lexer for the PostgreSQL dialect. Differences w.r.t. the SQL
-        lexer are:
-
-        - keywords and data types list parsed from the PG docs (run the
-          `_postgres_builtins` module to update them);
-        - Content of $-strings parsed using a specific lexer, e.g. the content
-          of a PL/Python function is parsed using the Python lexer;
-        - parse PG specific constructs: E-strings, $-strings, U&-strings,
-          different operators and punctuation.
-
-    `PlPgsqlLexer`
-        A lexer for the PL/pgSQL language. Adds a few specific construct on
-        top of the PG SQL lexer (such as <{text}' if rule else text
-                    append(text)
-            else:
-                styles: Dict[str, int] = {}
-                for text, style, _ in Segment.filter_control(
-                    Segment.simplify(self._record_buffer)
-                ):
-                    text = escape(text)
-                    if style:
-                        rule = style.get_html_style(_theme)
-                        style_number = styles.setdefault(rule, len(styles) + 1)
-                        if style.link:
-                            text = f'{text}'
-                        else:
-                            text = f'{text}'
-                    append(text)
-                stylesheet_rules: List[str] = []
-                stylesheet_append = stylesheet_rules.append
-                for style_rule, style_number in styles.items():
-                    if style_rule:
-                        stylesheet_append(f".r{style_number} {{{style_rule}}}")
-                stylesheet = "\n".join(stylesheet_rules)
-
-            rendered_code = render_code_format.format(
-                code="".join(fragments),
-                stylesheet=stylesheet,
-                foreground=_theme.foreground_color.hex,
-                background=_theme.background_color.hex,
-            )
-            if clear:
-                del self._record_buffer[:]
-        return rendered_code
-
-    def save_html(
-        self,
-        path: str,
-        *,
-        theme: Optional[TerminalTheme] = None,
-        clear: bool = True,
-        code_format: str = CONSOLE_HTML_FORMAT,
-        inline_styles: bool = False,
-    ) -> None:
-        """Generate HTML from console contents and write to a file (requires record=True argument in constructor).
-
-        Args:
-            path (str): Path to write html file.
-            theme (TerminalTheme, optional): TerminalTheme object containing console colors.
-            clear (bool, optional): Clear record buffer after exporting. Defaults to ``True``.
-            code_format (str, optional): Format string to render HTML. In addition to '{foreground}',
-                '{background}', and '{code}', should contain '{stylesheet}' if inline_styles is ``False``.
-            inline_styles (bool, optional): If ``True`` styles will be inlined in to spans, which makes files
-                larger but easier to cut and paste markup. If ``False``, styles will be embedded in a style tag.
-                Defaults to False.
-
-        """
-        html = self.export_html(
-            theme=theme,
-            clear=clear,
-            code_format=code_format,
-            inline_styles=inline_styles,
-        )
-        with open(path, "w", encoding="utf-8") as write_file:
-            write_file.write(html)
-
-    def export_svg(
-        self,
-        *,
-        title: str = "Rich",
-        theme: Optional[TerminalTheme] = None,
-        clear: bool = True,
-        code_format: str = CONSOLE_SVG_FORMAT,
-        font_aspect_ratio: float = 0.61,
-        unique_id: Optional[str] = None,
-    ) -> str:
-        """
-        Generate an SVG from the console contents (requires record=True in Console constructor).
-
-        Args:
-            title (str, optional): The title of the tab in the output image
-            theme (TerminalTheme, optional): The ``TerminalTheme`` object to use to style the terminal
-            clear (bool, optional): Clear record buffer after exporting. Defaults to ``True``
-            code_format (str, optional): Format string used to generate the SVG. Rich will inject a number of variables
-                into the string in order to form the final SVG output. The default template used and the variables
-                injected by Rich can be found by inspecting the ``console.CONSOLE_SVG_FORMAT`` variable.
-            font_aspect_ratio (float, optional): The width to height ratio of the font used in the ``code_format``
-                string. Defaults to 0.61, which is the width to height ratio of Fira Code (the default font).
-                If you aren't specifying a different font inside ``code_format``, you probably don't need this.
-            unique_id (str, optional): unique id that is used as the prefix for various elements (CSS styles, node
-                ids). If not set, this defaults to a computed value based on the recorded content.
-        """
-
-        from rich.cells import cell_len
-
-        style_cache: Dict[Style, str] = {}
-
-        def get_svg_style(style: Style) -> str:
-            """Convert a Style to CSS rules for SVG."""
-            if style in style_cache:
-                return style_cache[style]
-            css_rules = []
-            color = (
-                _theme.foreground_color
-                if (style.color is None or style.color.is_default)
-                else style.color.get_truecolor(_theme)
-            )
-            bgcolor = (
-                _theme.background_color
-                if (style.bgcolor is None or style.bgcolor.is_default)
-                else style.bgcolor.get_truecolor(_theme)
-            )
-            if style.reverse:
-                color, bgcolor = bgcolor, color
-            if style.dim:
-                color = blend_rgb(color, bgcolor, 0.4)
-            css_rules.append(f"fill: {color.hex}")
-            if style.bold:
-                css_rules.append("font-weight: bold")
-            if style.italic:
-                css_rules.append("font-style: italic;")
-            if style.underline:
-                css_rules.append("text-decoration: underline;")
-            if style.strike:
-                css_rules.append("text-decoration: line-through;")
-
-            css = ";".join(css_rules)
-            style_cache[style] = css
-            return css
-
-        _theme = theme or SVG_EXPORT_THEME
-
-        width = self.width
-        char_height = 20
-        char_width = char_height * font_aspect_ratio
-        line_height = char_height * 1.22
-
-        margin_top = 1
-        margin_right = 1
-        margin_bottom = 1
-        margin_left = 1
-
-        padding_top = 40
-        padding_right = 8
-        padding_bottom = 8
-        padding_left = 8
-
-        padding_width = padding_left + padding_right
-        padding_height = padding_top + padding_bottom
-        margin_width = margin_left + margin_right
-        margin_height = margin_top + margin_bottom
-
-        text_backgrounds: List[str] = []
-        text_group: List[str] = []
-        classes: Dict[str, int] = {}
-        style_no = 1
-
-        def escape_text(text: str) -> str:
-            """HTML escape text and replace spaces with nbsp."""
-            return escape(text).replace(" ", " ")
-
-        def make_tag(
-            name: str, content: Optional[str] = None, **attribs: object
-        ) -> str:
-            """Make a tag from name, content, and attributes."""
-
-            def stringify(value: object) -> str:
-                if isinstance(value, (float)):
-                    return format(value, "g")
-                return str(value)
-
-            tag_attribs = " ".join(
-                f'{k.lstrip("_").replace("_", "-")}="{stringify(v)}"'
-                for k, v in attribs.items()
-            )
-            return (
-                f"<{name} {tag_attribs}>{content}"
-                if content
-                else f"<{name} {tag_attribs}/>"
-            )
-
-        with self._record_buffer_lock:
-            segments = list(Segment.filter_control(self._record_buffer))
-            if clear:
-                self._record_buffer.clear()
-
-        if unique_id is None:
-            unique_id = "terminal-" + str(
-                zlib.adler32(
-                    ("".join(repr(segment) for segment in segments)).encode(
-                        "utf-8",
-                        "ignore",
-                    )
-                    + title.encode("utf-8", "ignore")
-                )
-            )
-        y = 0
-        for y, line in enumerate(Segment.split_and_crop_lines(segments, length=width)):
-            x = 0
-            for text, style, _control in line:
-                style = style or Style()
-                rules = get_svg_style(style)
-                if rules not in classes:
-                    classes[rules] = style_no
-                    style_no += 1
-                class_name = f"r{classes[rules]}"
-
-                if style.reverse:
-                    has_background = True
-                    background = (
-                        _theme.foreground_color.hex
-                        if style.color is None
-                        else style.color.get_truecolor(_theme).hex
-                    )
-                else:
-                    bgcolor = style.bgcolor
-                    has_background = bgcolor is not None and not bgcolor.is_default
-                    background = (
-                        _theme.background_color.hex
-                        if style.bgcolor is None
-                        else style.bgcolor.get_truecolor(_theme).hex
-                    )
-
-                text_length = cell_len(text)
-                if has_background:
-                    text_backgrounds.append(
-                        make_tag(
-                            "rect",
-                            fill=background,
-                            x=x * char_width,
-                            y=y * line_height + 1.5,
-                            width=char_width * text_length,
-                            height=line_height + 0.25,
-                            shape_rendering="crispEdges",
-                        )
-                    )
-
-                if text != " " * len(text):
-                    text_group.append(
-                        make_tag(
-                            "text",
-                            escape_text(text),
-                            _class=f"{unique_id}-{class_name}",
-                            x=x * char_width,
-                            y=y * line_height + char_height,
-                            textLength=char_width * len(text),
-                            clip_path=f"url(#{unique_id}-line-{y})",
-                        )
-                    )
-                x += cell_len(text)
-
-        line_offsets = [line_no * line_height + 1.5 for line_no in range(y)]
-        lines = "\n".join(
-            f"""
-    {make_tag("rect", x=0, y=offset, width=char_width * width, height=line_height + 0.25)}
-            """
-            for line_no, offset in enumerate(line_offsets)
-        )
-
-        styles = "\n".join(
-            f".{unique_id}-r{rule_no} {{ {css} }}" for css, rule_no in classes.items()
-        )
-        backgrounds = "".join(text_backgrounds)
-        matrix = "".join(text_group)
-
-        terminal_width = ceil(width * char_width + padding_width)
-        terminal_height = (y + 1) * line_height + padding_height
-        chrome = make_tag(
-            "rect",
-            fill=_theme.background_color.hex,
-            stroke="rgba(255,255,255,0.35)",
-            stroke_width="1",
-            x=margin_left,
-            y=margin_top,
-            width=terminal_width,
-            height=terminal_height,
-            rx=8,
-        )
-
-        title_color = _theme.foreground_color.hex
-        if title:
-            chrome += make_tag(
-                "text",
-                escape_text(title),
-                _class=f"{unique_id}-title",
-                fill=title_color,
-                text_anchor="middle",
-                x=terminal_width // 2,
-                y=margin_top + char_height + 6,
-            )
-        chrome += f"""
-            
-            
-            
-            
-            
-        """
-
-        svg = code_format.format(
-            unique_id=unique_id,
-            char_width=char_width,
-            char_height=char_height,
-            line_height=line_height,
-            terminal_width=char_width * width - 1,
-            terminal_height=(y + 1) * line_height - 1,
-            width=terminal_width + margin_width,
-            height=terminal_height + margin_height,
-            terminal_x=margin_left + padding_left,
-            terminal_y=margin_top + padding_top,
-            styles=styles,
-            chrome=chrome,
-            backgrounds=backgrounds,
-            matrix=matrix,
-            lines=lines,
-        )
-        return svg
-
-    def save_svg(
-        self,
-        path: str,
-        *,
-        title: str = "Rich",
-        theme: Optional[TerminalTheme] = None,
-        clear: bool = True,
-        code_format: str = CONSOLE_SVG_FORMAT,
-        font_aspect_ratio: float = 0.61,
-        unique_id: Optional[str] = None,
-    ) -> None:
-        """Generate an SVG file from the console contents (requires record=True in Console constructor).
-
-        Args:
-            path (str): The path to write the SVG to.
-            title (str, optional): The title of the tab in the output image
-            theme (TerminalTheme, optional): The ``TerminalTheme`` object to use to style the terminal
-            clear (bool, optional): Clear record buffer after exporting. Defaults to ``True``
-            code_format (str, optional): Format string used to generate the SVG. Rich will inject a number of variables
-                into the string in order to form the final SVG output. The default template used and the variables
-                injected by Rich can be found by inspecting the ``console.CONSOLE_SVG_FORMAT`` variable.
-            font_aspect_ratio (float, optional): The width to height ratio of the font used in the ``code_format``
-                string. Defaults to 0.61, which is the width to height ratio of Fira Code (the default font).
-                If you aren't specifying a different font inside ``code_format``, you probably don't need this.
-            unique_id (str, optional): unique id that is used as the prefix for various elements (CSS styles, node
-                ids). If not set, this defaults to a computed value based on the recorded content.
-        """
-        svg = self.export_svg(
-            title=title,
-            theme=theme,
-            clear=clear,
-            code_format=code_format,
-            font_aspect_ratio=font_aspect_ratio,
-            unique_id=unique_id,
-        )
-        with open(path, "w", encoding="utf-8") as write_file:
-            write_file.write(svg)
-
-
-def _svg_hash(svg_main_code: str) -> str:
-    """Returns a unique hash for the given SVG main code.
-
-    Args:
-        svg_main_code (str): The content we're going to inject in the SVG envelope.
-
-    Returns:
-        str: a hash of the given content
-    """
-    return str(zlib.adler32(svg_main_code.encode()))
-
-
-if __name__ == "__main__":  # pragma: no cover
-    console = Console(record=True)
-
-    console.log(
-        "JSONRPC [i]request[/i]",
-        5,
-        1.3,
-        True,
-        False,
-        None,
-        {
-            "jsonrpc": "2.0",
-            "method": "subtract",
-            "params": {"minuend": 42, "subtrahend": 23},
-            "id": 3,
-        },
-    )
-
-    console.log("Hello, World!", "{'a': 1}", repr(console))
-
-    console.print(
-        {
-            "name": None,
-            "empty": [],
-            "quiz": {
-                "sport": {
-                    "answered": True,
-                    "q1": {
-                        "question": "Which one is correct team name in NBA?",
-                        "options": [
-                            "New York Bulls",
-                            "Los Angeles Kings",
-                            "Golden State Warriors",
-                            "Huston Rocket",
-                        ],
-                        "answer": "Huston Rocket",
-                    },
-                },
-                "maths": {
-                    "answered": False,
-                    "q1": {
-                        "question": "5 + 7 = ?",
-                        "options": [10, 11, 12, 13],
-                        "answer": 12,
-                    },
-                    "q2": {
-                        "question": "12 - 8 = ?",
-                        "options": [1, 2, 3, 4],
-                        "answer": 4,
-                    },
-                },
-            },
-        }
-    )
diff --git a/venv/lib/python3.10/site-packages/rich/constrain.py b/venv/lib/python3.10/site-packages/rich/constrain.py
deleted file mode 100644
index 65fdf56342e8b5b8e181914881025231684e1871..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/constrain.py
+++ /dev/null
@@ -1,37 +0,0 @@
-from typing import Optional, TYPE_CHECKING
-
-from .jupyter import JupyterMixin
-from .measure import Measurement
-
-if TYPE_CHECKING:
-    from .console import Console, ConsoleOptions, RenderableType, RenderResult
-
-
-class Constrain(JupyterMixin):
-    """Constrain the width of a renderable to a given number of characters.
-
-    Args:
-        renderable (RenderableType): A renderable object.
-        width (int, optional): The maximum width (in characters) to render. Defaults to 80.
-    """
-
-    def __init__(self, renderable: "RenderableType", width: Optional[int] = 80) -> None:
-        self.renderable = renderable
-        self.width = width
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        if self.width is None:
-            yield self.renderable
-        else:
-            child_options = options.update_width(min(self.width, options.max_width))
-            yield from console.render(self.renderable, child_options)
-
-    def __rich_measure__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "Measurement":
-        if self.width is not None:
-            options = options.update_width(self.width)
-        measurement = Measurement.get(console, options, self.renderable)
-        return measurement
diff --git a/venv/lib/python3.10/site-packages/rich/containers.py b/venv/lib/python3.10/site-packages/rich/containers.py
deleted file mode 100644
index 901ff8ba6ea0836481a015ed5c627889cc416c03..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/containers.py
+++ /dev/null
@@ -1,167 +0,0 @@
-from itertools import zip_longest
-from typing import (
-    TYPE_CHECKING,
-    Iterable,
-    Iterator,
-    List,
-    Optional,
-    TypeVar,
-    Union,
-    overload,
-)
-
-if TYPE_CHECKING:
-    from .console import (
-        Console,
-        ConsoleOptions,
-        JustifyMethod,
-        OverflowMethod,
-        RenderResult,
-        RenderableType,
-    )
-    from .text import Text
-
-from .cells import cell_len
-from .measure import Measurement
-
-T = TypeVar("T")
-
-
-class Renderables:
-    """A list subclass which renders its contents to the console."""
-
-    def __init__(
-        self, renderables: Optional[Iterable["RenderableType"]] = None
-    ) -> None:
-        self._renderables: List["RenderableType"] = (
-            list(renderables) if renderables is not None else []
-        )
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        """Console render method to insert line-breaks."""
-        yield from self._renderables
-
-    def __rich_measure__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "Measurement":
-        dimensions = [
-            Measurement.get(console, options, renderable)
-            for renderable in self._renderables
-        ]
-        if not dimensions:
-            return Measurement(1, 1)
-        _min = max(dimension.minimum for dimension in dimensions)
-        _max = max(dimension.maximum for dimension in dimensions)
-        return Measurement(_min, _max)
-
-    def append(self, renderable: "RenderableType") -> None:
-        self._renderables.append(renderable)
-
-    def __iter__(self) -> Iterable["RenderableType"]:
-        return iter(self._renderables)
-
-
-class Lines:
-    """A list subclass which can render to the console."""
-
-    def __init__(self, lines: Iterable["Text"] = ()) -> None:
-        self._lines: List["Text"] = list(lines)
-
-    def __repr__(self) -> str:
-        return f"Lines({self._lines!r})"
-
-    def __iter__(self) -> Iterator["Text"]:
-        return iter(self._lines)
-
-    @overload
-    def __getitem__(self, index: int) -> "Text":
-        ...
-
-    @overload
-    def __getitem__(self, index: slice) -> List["Text"]:
-        ...
-
-    def __getitem__(self, index: Union[slice, int]) -> Union["Text", List["Text"]]:
-        return self._lines[index]
-
-    def __setitem__(self, index: int, value: "Text") -> "Lines":
-        self._lines[index] = value
-        return self
-
-    def __len__(self) -> int:
-        return self._lines.__len__()
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        """Console render method to insert line-breaks."""
-        yield from self._lines
-
-    def append(self, line: "Text") -> None:
-        self._lines.append(line)
-
-    def extend(self, lines: Iterable["Text"]) -> None:
-        self._lines.extend(lines)
-
-    def pop(self, index: int = -1) -> "Text":
-        return self._lines.pop(index)
-
-    def justify(
-        self,
-        console: "Console",
-        width: int,
-        justify: "JustifyMethod" = "left",
-        overflow: "OverflowMethod" = "fold",
-    ) -> None:
-        """Justify and overflow text to a given width.
-
-        Args:
-            console (Console): Console instance.
-            width (int): Number of cells available per line.
-            justify (str, optional): Default justify method for text: "left", "center", "full" or "right". Defaults to "left".
-            overflow (str, optional): Default overflow for text: "crop", "fold", or "ellipsis". Defaults to "fold".
-
-        """
-        from .text import Text
-
-        if justify == "left":
-            for line in self._lines:
-                line.truncate(width, overflow=overflow, pad=True)
-        elif justify == "center":
-            for line in self._lines:
-                line.rstrip()
-                line.truncate(width, overflow=overflow)
-                line.pad_left((width - cell_len(line.plain)) // 2)
-                line.pad_right(width - cell_len(line.plain))
-        elif justify == "right":
-            for line in self._lines:
-                line.rstrip()
-                line.truncate(width, overflow=overflow)
-                line.pad_left(width - cell_len(line.plain))
-        elif justify == "full":
-            for line_index, line in enumerate(self._lines):
-                if line_index == len(self._lines) - 1:
-                    break
-                words = line.split(" ")
-                words_size = sum(cell_len(word.plain) for word in words)
-                num_spaces = len(words) - 1
-                spaces = [1 for _ in range(num_spaces)]
-                index = 0
-                if spaces:
-                    while words_size + num_spaces < width:
-                        spaces[len(spaces) - index - 1] += 1
-                        num_spaces += 1
-                        index = (index + 1) % len(spaces)
-                tokens: List[Text] = []
-                for index, (word, next_word) in enumerate(
-                    zip_longest(words, words[1:])
-                ):
-                    tokens.append(word)
-                    if index < len(spaces):
-                        style = word.get_style_at_offset(console, -1)
-                        next_style = next_word.get_style_at_offset(console, 0)
-                        space_style = style if style == next_style else line.style
-                        tokens.append(Text(" " * spaces[index], style=space_style))
-                self[line_index] = Text("").join(tokens)
diff --git a/venv/lib/python3.10/site-packages/rich/control.py b/venv/lib/python3.10/site-packages/rich/control.py
deleted file mode 100644
index a8a912553a94660e55d94c958a4483d3ba2884d4..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/control.py
+++ /dev/null
@@ -1,225 +0,0 @@
-import sys
-import time
-from typing import TYPE_CHECKING, Callable, Dict, Iterable, List, Union
-
-if sys.version_info >= (3, 8):
-    from typing import Final
-else:
-    from typing_extensions import Final  # pragma: no cover
-
-from .segment import ControlCode, ControlType, Segment
-
-if TYPE_CHECKING:
-    from .console import Console, ConsoleOptions, RenderResult
-
-STRIP_CONTROL_CODES: Final = [
-    7,  # Bell
-    8,  # Backspace
-    11,  # Vertical tab
-    12,  # Form feed
-    13,  # Carriage return
-]
-_CONTROL_STRIP_TRANSLATE: Final = {
-    _codepoint: None for _codepoint in STRIP_CONTROL_CODES
-}
-
-CONTROL_ESCAPE: Final = {
-    7: "\\a",
-    8: "\\b",
-    11: "\\v",
-    12: "\\f",
-    13: "\\r",
-}
-
-CONTROL_CODES_FORMAT: Dict[int, Callable[..., str]] = {
-    ControlType.BELL: lambda: "\x07",
-    ControlType.CARRIAGE_RETURN: lambda: "\r",
-    ControlType.HOME: lambda: "\x1b[H",
-    ControlType.CLEAR: lambda: "\x1b[2J",
-    ControlType.ENABLE_ALT_SCREEN: lambda: "\x1b[?1049h",
-    ControlType.DISABLE_ALT_SCREEN: lambda: "\x1b[?1049l",
-    ControlType.SHOW_CURSOR: lambda: "\x1b[?25h",
-    ControlType.HIDE_CURSOR: lambda: "\x1b[?25l",
-    ControlType.CURSOR_UP: lambda param: f"\x1b[{param}A",
-    ControlType.CURSOR_DOWN: lambda param: f"\x1b[{param}B",
-    ControlType.CURSOR_FORWARD: lambda param: f"\x1b[{param}C",
-    ControlType.CURSOR_BACKWARD: lambda param: f"\x1b[{param}D",
-    ControlType.CURSOR_MOVE_TO_COLUMN: lambda param: f"\x1b[{param+1}G",
-    ControlType.ERASE_IN_LINE: lambda param: f"\x1b[{param}K",
-    ControlType.CURSOR_MOVE_TO: lambda x, y: f"\x1b[{y+1};{x+1}H",
-    ControlType.SET_WINDOW_TITLE: lambda title: f"\x1b]0;{title}\x07",
-}
-
-
-class Control:
-    """A renderable that inserts a control code (non printable but may move cursor).
-
-    Args:
-        *codes (str): Positional arguments are either a :class:`~rich.segment.ControlType` enum or a
-            tuple of ControlType and an integer parameter
-    """
-
-    __slots__ = ["segment"]
-
-    def __init__(self, *codes: Union[ControlType, ControlCode]) -> None:
-        control_codes: List[ControlCode] = [
-            (code,) if isinstance(code, ControlType) else code for code in codes
-        ]
-        _format_map = CONTROL_CODES_FORMAT
-        rendered_codes = "".join(
-            _format_map[code](*parameters) for code, *parameters in control_codes
-        )
-        self.segment = Segment(rendered_codes, None, control_codes)
-
-    @classmethod
-    def bell(cls) -> "Control":
-        """Ring the 'bell'."""
-        return cls(ControlType.BELL)
-
-    @classmethod
-    def home(cls) -> "Control":
-        """Move cursor to 'home' position."""
-        return cls(ControlType.HOME)
-
-    @classmethod
-    def move(cls, x: int = 0, y: int = 0) -> "Control":
-        """Move cursor relative to current position.
-
-        Args:
-            x (int): X offset.
-            y (int): Y offset.
-
-        Returns:
-            ~Control: Control object.
-
-        """
-
-        def get_codes() -> Iterable[ControlCode]:
-            control = ControlType
-            if x:
-                yield (
-                    control.CURSOR_FORWARD if x > 0 else control.CURSOR_BACKWARD,
-                    abs(x),
-                )
-            if y:
-                yield (
-                    control.CURSOR_DOWN if y > 0 else control.CURSOR_UP,
-                    abs(y),
-                )
-
-        control = cls(*get_codes())
-        return control
-
-    @classmethod
-    def move_to_column(cls, x: int, y: int = 0) -> "Control":
-        """Move to the given column, optionally add offset to row.
-
-        Returns:
-            x (int): absolute x (column)
-            y (int): optional y offset (row)
-
-        Returns:
-            ~Control: Control object.
-        """
-
-        return (
-            cls(
-                (ControlType.CURSOR_MOVE_TO_COLUMN, x),
-                (
-                    ControlType.CURSOR_DOWN if y > 0 else ControlType.CURSOR_UP,
-                    abs(y),
-                ),
-            )
-            if y
-            else cls((ControlType.CURSOR_MOVE_TO_COLUMN, x))
-        )
-
-    @classmethod
-    def move_to(cls, x: int, y: int) -> "Control":
-        """Move cursor to absolute position.
-
-        Args:
-            x (int): x offset (column)
-            y (int): y offset (row)
-
-        Returns:
-            ~Control: Control object.
-        """
-        return cls((ControlType.CURSOR_MOVE_TO, x, y))
-
-    @classmethod
-    def clear(cls) -> "Control":
-        """Clear the screen."""
-        return cls(ControlType.CLEAR)
-
-    @classmethod
-    def show_cursor(cls, show: bool) -> "Control":
-        """Show or hide the cursor."""
-        return cls(ControlType.SHOW_CURSOR if show else ControlType.HIDE_CURSOR)
-
-    @classmethod
-    def alt_screen(cls, enable: bool) -> "Control":
-        """Enable or disable alt screen."""
-        if enable:
-            return cls(ControlType.ENABLE_ALT_SCREEN, ControlType.HOME)
-        else:
-            return cls(ControlType.DISABLE_ALT_SCREEN)
-
-    @classmethod
-    def title(cls, title: str) -> "Control":
-        """Set the terminal window title
-
-        Args:
-            title (str): The new terminal window title
-        """
-        return cls((ControlType.SET_WINDOW_TITLE, title))
-
-    def __str__(self) -> str:
-        return self.segment.text
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        if self.segment.text:
-            yield self.segment
-
-
-def strip_control_codes(
-    text: str, _translate_table: Dict[int, None] = _CONTROL_STRIP_TRANSLATE
-) -> str:
-    """Remove control codes from text.
-
-    Args:
-        text (str): A string possibly contain control codes.
-
-    Returns:
-        str: String with control codes removed.
-    """
-    return text.translate(_translate_table)
-
-
-def escape_control_codes(
-    text: str,
-    _translate_table: Dict[int, str] = CONTROL_ESCAPE,
-) -> str:
-    """Replace control codes with their "escaped" equivalent in the given text.
-    (e.g. "\b" becomes "\\b")
-
-    Args:
-        text (str): A string possibly containing control codes.
-
-    Returns:
-        str: String with control codes replaced with their escaped version.
-    """
-    return text.translate(_translate_table)
-
-
-if __name__ == "__main__":  # pragma: no cover
-    from rich.console import Console
-
-    console = Console()
-    console.print("Look at the title of your terminal window ^")
-    # console.print(Control((ControlType.SET_WINDOW_TITLE, "Hello, world!")))
-    for i in range(10):
-        console.set_window_title("🚀 Loading" + "." * i)
-        time.sleep(0.5)
diff --git a/venv/lib/python3.10/site-packages/rich/default_styles.py b/venv/lib/python3.10/site-packages/rich/default_styles.py
deleted file mode 100644
index 3a0ad83a29f0c65b58e501ad4f5ed10f0830d0a1..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/default_styles.py
+++ /dev/null
@@ -1,193 +0,0 @@
-from typing import Dict
-
-from .style import Style
-
-DEFAULT_STYLES: Dict[str, Style] = {
-    "none": Style.null(),
-    "reset": Style(
-        color="default",
-        bgcolor="default",
-        dim=False,
-        bold=False,
-        italic=False,
-        underline=False,
-        blink=False,
-        blink2=False,
-        reverse=False,
-        conceal=False,
-        strike=False,
-    ),
-    "dim": Style(dim=True),
-    "bright": Style(dim=False),
-    "bold": Style(bold=True),
-    "strong": Style(bold=True),
-    "code": Style(reverse=True, bold=True),
-    "italic": Style(italic=True),
-    "emphasize": Style(italic=True),
-    "underline": Style(underline=True),
-    "blink": Style(blink=True),
-    "blink2": Style(blink2=True),
-    "reverse": Style(reverse=True),
-    "strike": Style(strike=True),
-    "black": Style(color="black"),
-    "red": Style(color="red"),
-    "green": Style(color="green"),
-    "yellow": Style(color="yellow"),
-    "magenta": Style(color="magenta"),
-    "cyan": Style(color="cyan"),
-    "white": Style(color="white"),
-    "inspect.attr": Style(color="yellow", italic=True),
-    "inspect.attr.dunder": Style(color="yellow", italic=True, dim=True),
-    "inspect.callable": Style(bold=True, color="red"),
-    "inspect.async_def": Style(italic=True, color="bright_cyan"),
-    "inspect.def": Style(italic=True, color="bright_cyan"),
-    "inspect.class": Style(italic=True, color="bright_cyan"),
-    "inspect.error": Style(bold=True, color="red"),
-    "inspect.equals": Style(),
-    "inspect.help": Style(color="cyan"),
-    "inspect.doc": Style(dim=True),
-    "inspect.value.border": Style(color="green"),
-    "live.ellipsis": Style(bold=True, color="red"),
-    "layout.tree.row": Style(dim=False, color="red"),
-    "layout.tree.column": Style(dim=False, color="blue"),
-    "logging.keyword": Style(bold=True, color="yellow"),
-    "logging.level.notset": Style(dim=True),
-    "logging.level.debug": Style(color="green"),
-    "logging.level.info": Style(color="blue"),
-    "logging.level.warning": Style(color="yellow"),
-    "logging.level.error": Style(color="red", bold=True),
-    "logging.level.critical": Style(color="red", bold=True, reverse=True),
-    "log.level": Style.null(),
-    "log.time": Style(color="cyan", dim=True),
-    "log.message": Style.null(),
-    "log.path": Style(dim=True),
-    "repr.ellipsis": Style(color="yellow"),
-    "repr.indent": Style(color="green", dim=True),
-    "repr.error": Style(color="red", bold=True),
-    "repr.str": Style(color="green", italic=False, bold=False),
-    "repr.brace": Style(bold=True),
-    "repr.comma": Style(bold=True),
-    "repr.ipv4": Style(bold=True, color="bright_green"),
-    "repr.ipv6": Style(bold=True, color="bright_green"),
-    "repr.eui48": Style(bold=True, color="bright_green"),
-    "repr.eui64": Style(bold=True, color="bright_green"),
-    "repr.tag_start": Style(bold=True),
-    "repr.tag_name": Style(color="bright_magenta", bold=True),
-    "repr.tag_contents": Style(color="default"),
-    "repr.tag_end": Style(bold=True),
-    "repr.attrib_name": Style(color="yellow", italic=False),
-    "repr.attrib_equal": Style(bold=True),
-    "repr.attrib_value": Style(color="magenta", italic=False),
-    "repr.number": Style(color="cyan", bold=True, italic=False),
-    "repr.number_complex": Style(color="cyan", bold=True, italic=False),  # same
-    "repr.bool_true": Style(color="bright_green", italic=True),
-    "repr.bool_false": Style(color="bright_red", italic=True),
-    "repr.none": Style(color="magenta", italic=True),
-    "repr.url": Style(underline=True, color="bright_blue", italic=False, bold=False),
-    "repr.uuid": Style(color="bright_yellow", bold=False),
-    "repr.call": Style(color="magenta", bold=True),
-    "repr.path": Style(color="magenta"),
-    "repr.filename": Style(color="bright_magenta"),
-    "rule.line": Style(color="bright_green"),
-    "rule.text": Style.null(),
-    "json.brace": Style(bold=True),
-    "json.bool_true": Style(color="bright_green", italic=True),
-    "json.bool_false": Style(color="bright_red", italic=True),
-    "json.null": Style(color="magenta", italic=True),
-    "json.number": Style(color="cyan", bold=True, italic=False),
-    "json.str": Style(color="green", italic=False, bold=False),
-    "json.key": Style(color="blue", bold=True),
-    "prompt": Style.null(),
-    "prompt.choices": Style(color="magenta", bold=True),
-    "prompt.default": Style(color="cyan", bold=True),
-    "prompt.invalid": Style(color="red"),
-    "prompt.invalid.choice": Style(color="red"),
-    "pretty": Style.null(),
-    "scope.border": Style(color="blue"),
-    "scope.key": Style(color="yellow", italic=True),
-    "scope.key.special": Style(color="yellow", italic=True, dim=True),
-    "scope.equals": Style(color="red"),
-    "table.header": Style(bold=True),
-    "table.footer": Style(bold=True),
-    "table.cell": Style.null(),
-    "table.title": Style(italic=True),
-    "table.caption": Style(italic=True, dim=True),
-    "traceback.error": Style(color="red", italic=True),
-    "traceback.border.syntax_error": Style(color="bright_red"),
-    "traceback.border": Style(color="red"),
-    "traceback.text": Style.null(),
-    "traceback.title": Style(color="red", bold=True),
-    "traceback.exc_type": Style(color="bright_red", bold=True),
-    "traceback.exc_value": Style.null(),
-    "traceback.offset": Style(color="bright_red", bold=True),
-    "traceback.error_range": Style(underline=True, bold=True),
-    "traceback.note": Style(color="green", bold=True),
-    "traceback.group.border": Style(color="magenta"),
-    "bar.back": Style(color="grey23"),
-    "bar.complete": Style(color="rgb(249,38,114)"),
-    "bar.finished": Style(color="rgb(114,156,31)"),
-    "bar.pulse": Style(color="rgb(249,38,114)"),
-    "progress.description": Style.null(),
-    "progress.filesize": Style(color="green"),
-    "progress.filesize.total": Style(color="green"),
-    "progress.download": Style(color="green"),
-    "progress.elapsed": Style(color="yellow"),
-    "progress.percentage": Style(color="magenta"),
-    "progress.remaining": Style(color="cyan"),
-    "progress.data.speed": Style(color="red"),
-    "progress.spinner": Style(color="green"),
-    "status.spinner": Style(color="green"),
-    "tree": Style(),
-    "tree.line": Style(),
-    "markdown.paragraph": Style(),
-    "markdown.text": Style(),
-    "markdown.em": Style(italic=True),
-    "markdown.emph": Style(italic=True),  # For commonmark backwards compatibility
-    "markdown.strong": Style(bold=True),
-    "markdown.code": Style(bold=True, color="cyan", bgcolor="black"),
-    "markdown.code_block": Style(color="cyan", bgcolor="black"),
-    "markdown.block_quote": Style(color="magenta"),
-    "markdown.list": Style(color="cyan"),
-    "markdown.item": Style(),
-    "markdown.item.bullet": Style(color="yellow", bold=True),
-    "markdown.item.number": Style(color="yellow", bold=True),
-    "markdown.hr": Style(color="yellow"),
-    "markdown.h1.border": Style(),
-    "markdown.h1": Style(bold=True),
-    "markdown.h2": Style(bold=True, underline=True),
-    "markdown.h3": Style(bold=True),
-    "markdown.h4": Style(bold=True, dim=True),
-    "markdown.h5": Style(underline=True),
-    "markdown.h6": Style(italic=True),
-    "markdown.h7": Style(italic=True, dim=True),
-    "markdown.link": Style(color="bright_blue"),
-    "markdown.link_url": Style(color="blue", underline=True),
-    "markdown.s": Style(strike=True),
-    "iso8601.date": Style(color="blue"),
-    "iso8601.time": Style(color="magenta"),
-    "iso8601.timezone": Style(color="yellow"),
-}
-
-
-if __name__ == "__main__":  # pragma: no cover
-    import argparse
-    import io
-
-    from rich.console import Console
-    from rich.table import Table
-    from rich.text import Text
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--html", action="store_true", help="Export as HTML table")
-    args = parser.parse_args()
-    html: bool = args.html
-    console = Console(record=True, width=70, file=io.StringIO()) if html else Console()
-
-    table = Table("Name", "Styling")
-
-    for style_name, style in DEFAULT_STYLES.items():
-        table.add_row(Text(style_name, style=style), str(style))
-
-    console.print(table)
-    if html:
-        print(console.export_html(inline_styles=True))
diff --git a/venv/lib/python3.10/site-packages/rich/diagnose.py b/venv/lib/python3.10/site-packages/rich/diagnose.py
deleted file mode 100644
index 5d163877da5ce00cabbb09930aadd8f49c0de3b5..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/diagnose.py
+++ /dev/null
@@ -1,38 +0,0 @@
-import os
-import platform
-
-from rich import inspect
-from rich.console import Console, get_windows_console_features
-from rich.panel import Panel
-from rich.pretty import Pretty
-
-
-def report() -> None:  # pragma: no cover
-    """Print a report to the terminal with debugging information"""
-    console = Console()
-    inspect(console)
-    features = get_windows_console_features()
-    inspect(features)
-
-    env_names = (
-        "CLICOLOR",
-        "COLORTERM",
-        "COLUMNS",
-        "JPY_PARENT_PID",
-        "JUPYTER_COLUMNS",
-        "JUPYTER_LINES",
-        "LINES",
-        "NO_COLOR",
-        "TERM_PROGRAM",
-        "TERM",
-        "TTY_COMPATIBLE",
-        "VSCODE_VERBOSE_LOGGING",
-    )
-    env = {name: os.getenv(name) for name in env_names}
-    console.print(Panel.fit((Pretty(env)), title="[b]Environment Variables"))
-
-    console.print(f'platform="{platform.system()}"')
-
-
-if __name__ == "__main__":  # pragma: no cover
-    report()
diff --git a/venv/lib/python3.10/site-packages/rich/emoji.py b/venv/lib/python3.10/site-packages/rich/emoji.py
deleted file mode 100644
index d5a1062a965c0ee75d20ad08d0179a90fca1646d..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/emoji.py
+++ /dev/null
@@ -1,96 +0,0 @@
-import sys
-from typing import TYPE_CHECKING, Optional, Union
-
-from .jupyter import JupyterMixin
-from .segment import Segment
-from .style import Style
-from ._emoji_codes import EMOJI
-from ._emoji_replace import _emoji_replace
-
-if sys.version_info >= (3, 8):
-    from typing import Literal
-else:
-    from typing_extensions import Literal  # pragma: no cover
-
-
-if TYPE_CHECKING:
-    from .console import Console, ConsoleOptions, RenderResult
-
-
-EmojiVariant = Literal["emoji", "text"]
-
-
-class NoEmoji(Exception):
-    """No emoji by that name."""
-
-
-class Emoji(JupyterMixin):
-    __slots__ = ["name", "style", "_char", "variant"]
-
-    VARIANTS = {"text": "\uFE0E", "emoji": "\uFE0F"}
-
-    def __init__(
-        self,
-        name: str,
-        style: Union[str, Style] = "none",
-        variant: Optional[EmojiVariant] = None,
-    ) -> None:
-        """A single emoji character.
-
-        Args:
-            name (str): Name of emoji.
-            style (Union[str, Style], optional): Optional style. Defaults to None.
-
-        Raises:
-            NoEmoji: If the emoji doesn't exist.
-        """
-        self.name = name
-        self.style = style
-        self.variant = variant
-        try:
-            self._char = EMOJI[name]
-        except KeyError:
-            raise NoEmoji(f"No emoji called {name!r}")
-        if variant is not None:
-            self._char += self.VARIANTS.get(variant, "")
-
-    @classmethod
-    def replace(cls, text: str) -> str:
-        """Replace emoji markup with corresponding unicode characters.
-
-        Args:
-            text (str): A string with emojis codes, e.g. "Hello :smiley:!"
-
-        Returns:
-            str: A string with emoji codes replaces with actual emoji.
-        """
-        return _emoji_replace(text)
-
-    def __repr__(self) -> str:
-        return f""
-
-    def __str__(self) -> str:
-        return self._char
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        yield Segment(self._char, console.get_style(self.style))
-
-
-if __name__ == "__main__":  # pragma: no cover
-    import sys
-
-    from rich.columns import Columns
-    from rich.console import Console
-
-    console = Console(record=True)
-
-    columns = Columns(
-        (f":{name}: {name}" for name in sorted(EMOJI.keys()) if "\u200D" not in name),
-        column_first=True,
-    )
-
-    console.print(columns)
-    if len(sys.argv) > 1:
-        console.save_html(sys.argv[1])
diff --git a/venv/lib/python3.10/site-packages/rich/errors.py b/venv/lib/python3.10/site-packages/rich/errors.py
deleted file mode 100644
index 0bcbe53ef59373c608e62ea285536f8b22b47ecb..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/errors.py
+++ /dev/null
@@ -1,34 +0,0 @@
-class ConsoleError(Exception):
-    """An error in console operation."""
-
-
-class StyleError(Exception):
-    """An error in styles."""
-
-
-class StyleSyntaxError(ConsoleError):
-    """Style was badly formatted."""
-
-
-class MissingStyle(StyleError):
-    """No such style."""
-
-
-class StyleStackError(ConsoleError):
-    """Style stack is invalid."""
-
-
-class NotRenderableError(ConsoleError):
-    """Object is not renderable."""
-
-
-class MarkupError(ConsoleError):
-    """Markup was badly formatted."""
-
-
-class LiveError(ConsoleError):
-    """Error related to Live display."""
-
-
-class NoAltScreen(ConsoleError):
-    """Alt screen mode was required."""
diff --git a/venv/lib/python3.10/site-packages/rich/file_proxy.py b/venv/lib/python3.10/site-packages/rich/file_proxy.py
deleted file mode 100644
index 4b0b0da6c2a62b2b1468c35ddd69f1bbb9b91aa8..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/file_proxy.py
+++ /dev/null
@@ -1,57 +0,0 @@
-import io
-from typing import IO, TYPE_CHECKING, Any, List
-
-from .ansi import AnsiDecoder
-from .text import Text
-
-if TYPE_CHECKING:
-    from .console import Console
-
-
-class FileProxy(io.TextIOBase):
-    """Wraps a file (e.g. sys.stdout) and redirects writes to a console."""
-
-    def __init__(self, console: "Console", file: IO[str]) -> None:
-        self.__console = console
-        self.__file = file
-        self.__buffer: List[str] = []
-        self.__ansi_decoder = AnsiDecoder()
-
-    @property
-    def rich_proxied_file(self) -> IO[str]:
-        """Get proxied file."""
-        return self.__file
-
-    def __getattr__(self, name: str) -> Any:
-        return getattr(self.__file, name)
-
-    def write(self, text: str) -> int:
-        if not isinstance(text, str):
-            raise TypeError(f"write() argument must be str, not {type(text).__name__}")
-        buffer = self.__buffer
-        lines: List[str] = []
-        while text:
-            line, new_line, text = text.partition("\n")
-            if new_line:
-                lines.append("".join(buffer) + line)
-                buffer.clear()
-            else:
-                buffer.append(line)
-                break
-        if lines:
-            console = self.__console
-            with console:
-                output = Text("\n").join(
-                    self.__ansi_decoder.decode_line(line) for line in lines
-                )
-                console.print(output)
-        return len(text)
-
-    def flush(self) -> None:
-        output = "".join(self.__buffer)
-        if output:
-            self.__console.print(output)
-        del self.__buffer[:]
-
-    def fileno(self) -> int:
-        return self.__file.fileno()
diff --git a/venv/lib/python3.10/site-packages/rich/filesize.py b/venv/lib/python3.10/site-packages/rich/filesize.py
deleted file mode 100644
index 83bc9118d2bdb8983f863063687c2ea394a9abb1..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/filesize.py
+++ /dev/null
@@ -1,88 +0,0 @@
-"""Functions for reporting filesizes. Borrowed from https://github.com/PyFilesystem/pyfilesystem2
-
-The functions declared in this module should cover the different
-use cases needed to generate a string representation of a file size
-using several different units. Since there are many standards regarding
-file size units, three different functions have been implemented.
-
-See Also:
-    * `Wikipedia: Binary prefix `_
-
-"""
-
-__all__ = ["decimal"]
-
-from typing import Iterable, List, Optional, Tuple
-
-
-def _to_str(
-    size: int,
-    suffixes: Iterable[str],
-    base: int,
-    *,
-    precision: Optional[int] = 1,
-    separator: Optional[str] = " ",
-) -> str:
-    if size == 1:
-        return "1 byte"
-    elif size < base:
-        return f"{size:,} bytes"
-
-    for i, suffix in enumerate(suffixes, 2):  # noqa: B007
-        unit = base**i
-        if size < unit:
-            break
-    return "{:,.{precision}f}{separator}{}".format(
-        (base * size / unit),
-        suffix,
-        precision=precision,
-        separator=separator,
-    )
-
-
-def pick_unit_and_suffix(size: int, suffixes: List[str], base: int) -> Tuple[int, str]:
-    """Pick a suffix and base for the given size."""
-    for i, suffix in enumerate(suffixes):
-        unit = base**i
-        if size < unit * base:
-            break
-    return unit, suffix
-
-
-def decimal(
-    size: int,
-    *,
-    precision: Optional[int] = 1,
-    separator: Optional[str] = " ",
-) -> str:
-    """Convert a filesize in to a string (powers of 1000, SI prefixes).
-
-    In this convention, ``1000 B = 1 kB``.
-
-    This is typically the format used to advertise the storage
-    capacity of USB flash drives and the like (*256 MB* meaning
-    actually a storage capacity of more than *256 000 000 B*),
-    or used by **Mac OS X** since v10.6 to report file sizes.
-
-    Arguments:
-        int (size): A file size.
-        int (precision): The number of decimal places to include (default = 1).
-        str (separator): The string to separate the value from the units (default = " ").
-
-    Returns:
-        `str`: A string containing a abbreviated file size and units.
-
-    Example:
-        >>> filesize.decimal(30000)
-        '30.0 kB'
-        >>> filesize.decimal(30000, precision=2, separator="")
-        '30.00kB'
-
-    """
-    return _to_str(
-        size,
-        ("kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"),
-        1000,
-        precision=precision,
-        separator=separator,
-    )
diff --git a/venv/lib/python3.10/site-packages/rich/highlighter.py b/venv/lib/python3.10/site-packages/rich/highlighter.py
deleted file mode 100644
index e4c462e2b632a268ecfb4772320206b94a3d859c..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/highlighter.py
+++ /dev/null
@@ -1,232 +0,0 @@
-import re
-from abc import ABC, abstractmethod
-from typing import List, Union
-
-from .text import Span, Text
-
-
-def _combine_regex(*regexes: str) -> str:
-    """Combine a number of regexes in to a single regex.
-
-    Returns:
-        str: New regex with all regexes ORed together.
-    """
-    return "|".join(regexes)
-
-
-class Highlighter(ABC):
-    """Abstract base class for highlighters."""
-
-    def __call__(self, text: Union[str, Text]) -> Text:
-        """Highlight a str or Text instance.
-
-        Args:
-            text (Union[str, ~Text]): Text to highlight.
-
-        Raises:
-            TypeError: If not called with text or str.
-
-        Returns:
-            Text: A test instance with highlighting applied.
-        """
-        if isinstance(text, str):
-            highlight_text = Text(text)
-        elif isinstance(text, Text):
-            highlight_text = text.copy()
-        else:
-            raise TypeError(f"str or Text instance required, not {text!r}")
-        self.highlight(highlight_text)
-        return highlight_text
-
-    @abstractmethod
-    def highlight(self, text: Text) -> None:
-        """Apply highlighting in place to text.
-
-        Args:
-            text (~Text): A text object highlight.
-        """
-
-
-class NullHighlighter(Highlighter):
-    """A highlighter object that doesn't highlight.
-
-    May be used to disable highlighting entirely.
-
-    """
-
-    def highlight(self, text: Text) -> None:
-        """Nothing to do"""
-
-
-class RegexHighlighter(Highlighter):
-    """Applies highlighting from a list of regular expressions."""
-
-    highlights: List[str] = []
-    base_style: str = ""
-
-    def highlight(self, text: Text) -> None:
-        """Highlight :class:`rich.text.Text` using regular expressions.
-
-        Args:
-            text (~Text): Text to highlighted.
-
-        """
-
-        highlight_regex = text.highlight_regex
-        for re_highlight in self.highlights:
-            highlight_regex(re_highlight, style_prefix=self.base_style)
-
-
-class ReprHighlighter(RegexHighlighter):
-    """Highlights the text typically produced from ``__repr__`` methods."""
-
-    base_style = "repr."
-    highlights = [
-        r"(?P<)(?P[-\w.:|]*)(?P[\w\W]*)(?P>)",
-        r'(?P[\w_]{1,50})=(?P"?[\w_]+"?)?',
-        r"(?P[][{}()])",
-        _combine_regex(
-            r"(?P[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})",
-            r"(?P([A-Fa-f0-9]{1,4}::?){1,7}[A-Fa-f0-9]{1,4})",
-            r"(?P(?:[0-9A-Fa-f]{1,2}-){7}[0-9A-Fa-f]{1,2}|(?:[0-9A-Fa-f]{1,2}:){7}[0-9A-Fa-f]{1,2}|(?:[0-9A-Fa-f]{4}\.){3}[0-9A-Fa-f]{4})",
-            r"(?P(?:[0-9A-Fa-f]{1,2}-){5}[0-9A-Fa-f]{1,2}|(?:[0-9A-Fa-f]{1,2}:){5}[0-9A-Fa-f]{1,2}|(?:[0-9A-Fa-f]{4}\.){2}[0-9A-Fa-f]{4})",
-            r"(?P[a-fA-F0-9]{8}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{12})",
-            r"(?P[\w.]*?)\(",
-            r"\b(?PTrue)\b|\b(?PFalse)\b|\b(?PNone)\b",
-            r"(?P\.\.\.)",
-            r"(?P(?(?\B(/[-\w._+]+)*\/)(?P[-\w._+]*)?",
-            r"(?b?'''.*?(?(file|https|http|ws|wss)://[-0-9a-zA-Z$_+!`(),.?/;:&=%#~@]*)",
-        ),
-    ]
-
-
-class JSONHighlighter(RegexHighlighter):
-    """Highlights JSON"""
-
-    # Captures the start and end of JSON strings, handling escaped quotes
-    JSON_STR = r"(?b?\".*?(?[\{\[\(\)\]\}])",
-            r"\b(?Ptrue)\b|\b(?Pfalse)\b|\b(?Pnull)\b",
-            r"(?P(? None:
-        super().highlight(text)
-
-        # Additional work to handle highlighting JSON keys
-        plain = text.plain
-        append = text.spans.append
-        whitespace = self.JSON_WHITESPACE
-        for match in re.finditer(self.JSON_STR, plain):
-            start, end = match.span()
-            cursor = end
-            while cursor < len(plain):
-                char = plain[cursor]
-                cursor += 1
-                if char == ":":
-                    append(Span(start, end, "json.key"))
-                elif char in whitespace:
-                    continue
-                break
-
-
-class ISO8601Highlighter(RegexHighlighter):
-    """Highlights the ISO8601 date time strings.
-    Regex reference: https://www.oreilly.com/library/view/regular-expressions-cookbook/9781449327453/ch04s07.html
-    """
-
-    base_style = "iso8601."
-    highlights = [
-        #
-        # Dates
-        #
-        # Calendar month (e.g. 2008-08). The hyphen is required
-        r"^(?P[0-9]{4})-(?P1[0-2]|0[1-9])$",
-        # Calendar date w/o hyphens (e.g. 20080830)
-        r"^(?P(?P[0-9]{4})(?P1[0-2]|0[1-9])(?P3[01]|0[1-9]|[12][0-9]))$",
-        # Ordinal date (e.g. 2008-243). The hyphen is optional
-        r"^(?P(?P[0-9]{4})-?(?P36[0-6]|3[0-5][0-9]|[12][0-9]{2}|0[1-9][0-9]|00[1-9]))$",
-        #
-        # Weeks
-        #
-        # Week of the year (e.g., 2008-W35). The hyphen is optional
-        r"^(?P(?P[0-9]{4})-?W(?P5[0-3]|[1-4][0-9]|0[1-9]))$",
-        # Week date (e.g., 2008-W35-6). The hyphens are optional
-        r"^(?P(?P[0-9]{4})-?W(?P5[0-3]|[1-4][0-9]|0[1-9])-?(?P[1-7]))$",
-        #
-        # Times
-        #
-        # Hours and minutes (e.g., 17:21). The colon is optional
-        r"^(?P{text}'
-        append_fragment(text)
-
-    code = "".join(fragments)
-    html = JUPYTER_HTML_FORMAT.format(code=code)
-
-    return html
-
-
-def display(segments: Iterable[Segment], text: str) -> None:
-    """Render segments to Jupyter."""
-    html = _render_segments(segments)
-    jupyter_renderable = JupyterRenderable(html, text)
-    try:
-        from IPython.display import display as ipython_display
-
-        ipython_display(jupyter_renderable)
-    except ModuleNotFoundError:
-        # Handle the case where the Console has force_jupyter=True,
-        # but IPython is not installed.
-        pass
-
-
-def print(*args: Any, **kwargs: Any) -> None:
-    """Proxy for Console print."""
-    console = get_console()
-    return console.print(*args, **kwargs)
diff --git a/venv/lib/python3.10/site-packages/rich/layout.py b/venv/lib/python3.10/site-packages/rich/layout.py
deleted file mode 100644
index 7fa2852ea14b4b22c5a9cbc3acdfddc5c55c4200..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/layout.py
+++ /dev/null
@@ -1,442 +0,0 @@
-from abc import ABC, abstractmethod
-from itertools import islice
-from operator import itemgetter
-from threading import RLock
-from typing import (
-    TYPE_CHECKING,
-    Dict,
-    Iterable,
-    List,
-    NamedTuple,
-    Optional,
-    Sequence,
-    Tuple,
-    Union,
-)
-
-from ._ratio import ratio_resolve
-from .align import Align
-from .console import Console, ConsoleOptions, RenderableType, RenderResult
-from .highlighter import ReprHighlighter
-from .panel import Panel
-from .pretty import Pretty
-from .region import Region
-from .repr import Result, rich_repr
-from .segment import Segment
-from .style import StyleType
-
-if TYPE_CHECKING:
-    from rich.tree import Tree
-
-
-class LayoutRender(NamedTuple):
-    """An individual layout render."""
-
-    region: Region
-    render: List[List[Segment]]
-
-
-RegionMap = Dict["Layout", Region]
-RenderMap = Dict["Layout", LayoutRender]
-
-
-class LayoutError(Exception):
-    """Layout related error."""
-
-
-class NoSplitter(LayoutError):
-    """Requested splitter does not exist."""
-
-
-class _Placeholder:
-    """An internal renderable used as a Layout placeholder."""
-
-    highlighter = ReprHighlighter()
-
-    def __init__(self, layout: "Layout", style: StyleType = "") -> None:
-        self.layout = layout
-        self.style = style
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        width = options.max_width
-        height = options.height or options.size.height
-        layout = self.layout
-        title = (
-            f"{layout.name!r} ({width} x {height})"
-            if layout.name
-            else f"({width} x {height})"
-        )
-        yield Panel(
-            Align.center(Pretty(layout), vertical="middle"),
-            style=self.style,
-            title=self.highlighter(title),
-            border_style="blue",
-            height=height,
-        )
-
-
-class Splitter(ABC):
-    """Base class for a splitter."""
-
-    name: str = ""
-
-    @abstractmethod
-    def get_tree_icon(self) -> str:
-        """Get the icon (emoji) used in layout.tree"""
-
-    @abstractmethod
-    def divide(
-        self, children: Sequence["Layout"], region: Region
-    ) -> Iterable[Tuple["Layout", Region]]:
-        """Divide a region amongst several child layouts.
-
-        Args:
-            children (Sequence(Layout)): A number of child layouts.
-            region (Region): A rectangular region to divide.
-        """
-
-
-class RowSplitter(Splitter):
-    """Split a layout region in to rows."""
-
-    name = "row"
-
-    def get_tree_icon(self) -> str:
-        return "[layout.tree.row]⬌"
-
-    def divide(
-        self, children: Sequence["Layout"], region: Region
-    ) -> Iterable[Tuple["Layout", Region]]:
-        x, y, width, height = region
-        render_widths = ratio_resolve(width, children)
-        offset = 0
-        _Region = Region
-        for child, child_width in zip(children, render_widths):
-            yield child, _Region(x + offset, y, child_width, height)
-            offset += child_width
-
-
-class ColumnSplitter(Splitter):
-    """Split a layout region in to columns."""
-
-    name = "column"
-
-    def get_tree_icon(self) -> str:
-        return "[layout.tree.column]⬍"
-
-    def divide(
-        self, children: Sequence["Layout"], region: Region
-    ) -> Iterable[Tuple["Layout", Region]]:
-        x, y, width, height = region
-        render_heights = ratio_resolve(height, children)
-        offset = 0
-        _Region = Region
-        for child, child_height in zip(children, render_heights):
-            yield child, _Region(x, y + offset, width, child_height)
-            offset += child_height
-
-
-@rich_repr
-class Layout:
-    """A renderable to divide a fixed height in to rows or columns.
-
-    Args:
-        renderable (RenderableType, optional): Renderable content, or None for placeholder. Defaults to None.
-        name (str, optional): Optional identifier for Layout. Defaults to None.
-        size (int, optional): Optional fixed size of layout. Defaults to None.
-        minimum_size (int, optional): Minimum size of layout. Defaults to 1.
-        ratio (int, optional): Optional ratio for flexible layout. Defaults to 1.
-        visible (bool, optional): Visibility of layout. Defaults to True.
-    """
-
-    splitters = {"row": RowSplitter, "column": ColumnSplitter}
-
-    def __init__(
-        self,
-        renderable: Optional[RenderableType] = None,
-        *,
-        name: Optional[str] = None,
-        size: Optional[int] = None,
-        minimum_size: int = 1,
-        ratio: int = 1,
-        visible: bool = True,
-    ) -> None:
-        self._renderable = renderable or _Placeholder(self)
-        self.size = size
-        self.minimum_size = minimum_size
-        self.ratio = ratio
-        self.name = name
-        self.visible = visible
-        self.splitter: Splitter = self.splitters["column"]()
-        self._children: List[Layout] = []
-        self._render_map: RenderMap = {}
-        self._lock = RLock()
-
-    def __rich_repr__(self) -> Result:
-        yield "name", self.name, None
-        yield "size", self.size, None
-        yield "minimum_size", self.minimum_size, 1
-        yield "ratio", self.ratio, 1
-
-    @property
-    def renderable(self) -> RenderableType:
-        """Layout renderable."""
-        return self if self._children else self._renderable
-
-    @property
-    def children(self) -> List["Layout"]:
-        """Gets (visible) layout children."""
-        return [child for child in self._children if child.visible]
-
-    @property
-    def map(self) -> RenderMap:
-        """Get a map of the last render."""
-        return self._render_map
-
-    def get(self, name: str) -> Optional["Layout"]:
-        """Get a named layout, or None if it doesn't exist.
-
-        Args:
-            name (str): Name of layout.
-
-        Returns:
-            Optional[Layout]: Layout instance or None if no layout was found.
-        """
-        if self.name == name:
-            return self
-        else:
-            for child in self._children:
-                named_layout = child.get(name)
-                if named_layout is not None:
-                    return named_layout
-        return None
-
-    def __getitem__(self, name: str) -> "Layout":
-        layout = self.get(name)
-        if layout is None:
-            raise KeyError(f"No layout with name {name!r}")
-        return layout
-
-    @property
-    def tree(self) -> "Tree":
-        """Get a tree renderable to show layout structure."""
-        from rich.styled import Styled
-        from rich.table import Table
-        from rich.tree import Tree
-
-        def summary(layout: "Layout") -> Table:
-            icon = layout.splitter.get_tree_icon()
-
-            table = Table.grid(padding=(0, 1, 0, 0))
-
-            text: RenderableType = (
-                Pretty(layout) if layout.visible else Styled(Pretty(layout), "dim")
-            )
-            table.add_row(icon, text)
-            _summary = table
-            return _summary
-
-        layout = self
-        tree = Tree(
-            summary(layout),
-            guide_style=f"layout.tree.{layout.splitter.name}",
-            highlight=True,
-        )
-
-        def recurse(tree: "Tree", layout: "Layout") -> None:
-            for child in layout._children:
-                recurse(
-                    tree.add(
-                        summary(child),
-                        guide_style=f"layout.tree.{child.splitter.name}",
-                    ),
-                    child,
-                )
-
-        recurse(tree, self)
-        return tree
-
-    def split(
-        self,
-        *layouts: Union["Layout", RenderableType],
-        splitter: Union[Splitter, str] = "column",
-    ) -> None:
-        """Split the layout in to multiple sub-layouts.
-
-        Args:
-            *layouts (Layout): Positional arguments should be (sub) Layout instances.
-            splitter (Union[Splitter, str]): Splitter instance or name of splitter.
-        """
-        _layouts = [
-            layout if isinstance(layout, Layout) else Layout(layout)
-            for layout in layouts
-        ]
-        try:
-            self.splitter = (
-                splitter
-                if isinstance(splitter, Splitter)
-                else self.splitters[splitter]()
-            )
-        except KeyError:
-            raise NoSplitter(f"No splitter called {splitter!r}")
-        self._children[:] = _layouts
-
-    def add_split(self, *layouts: Union["Layout", RenderableType]) -> None:
-        """Add a new layout(s) to existing split.
-
-        Args:
-            *layouts (Union[Layout, RenderableType]): Positional arguments should be renderables or (sub) Layout instances.
-
-        """
-        _layouts = (
-            layout if isinstance(layout, Layout) else Layout(layout)
-            for layout in layouts
-        )
-        self._children.extend(_layouts)
-
-    def split_row(self, *layouts: Union["Layout", RenderableType]) -> None:
-        """Split the layout in to a row (layouts side by side).
-
-        Args:
-            *layouts (Layout): Positional arguments should be (sub) Layout instances.
-        """
-        self.split(*layouts, splitter="row")
-
-    def split_column(self, *layouts: Union["Layout", RenderableType]) -> None:
-        """Split the layout in to a column (layouts stacked on top of each other).
-
-        Args:
-            *layouts (Layout): Positional arguments should be (sub) Layout instances.
-        """
-        self.split(*layouts, splitter="column")
-
-    def unsplit(self) -> None:
-        """Reset splits to initial state."""
-        del self._children[:]
-
-    def update(self, renderable: RenderableType) -> None:
-        """Update renderable.
-
-        Args:
-            renderable (RenderableType): New renderable object.
-        """
-        with self._lock:
-            self._renderable = renderable
-
-    def refresh_screen(self, console: "Console", layout_name: str) -> None:
-        """Refresh a sub-layout.
-
-        Args:
-            console (Console): Console instance where Layout is to be rendered.
-            layout_name (str): Name of layout.
-        """
-        with self._lock:
-            layout = self[layout_name]
-            region, _lines = self._render_map[layout]
-            (x, y, width, height) = region
-            lines = console.render_lines(
-                layout, console.options.update_dimensions(width, height)
-            )
-            self._render_map[layout] = LayoutRender(region, lines)
-            console.update_screen_lines(lines, x, y)
-
-    def _make_region_map(self, width: int, height: int) -> RegionMap:
-        """Create a dict that maps layout on to Region."""
-        stack: List[Tuple[Layout, Region]] = [(self, Region(0, 0, width, height))]
-        push = stack.append
-        pop = stack.pop
-        layout_regions: List[Tuple[Layout, Region]] = []
-        append_layout_region = layout_regions.append
-        while stack:
-            append_layout_region(pop())
-            layout, region = layout_regions[-1]
-            children = layout.children
-            if children:
-                for child_and_region in layout.splitter.divide(children, region):
-                    push(child_and_region)
-
-        region_map = {
-            layout: region
-            for layout, region in sorted(layout_regions, key=itemgetter(1))
-        }
-        return region_map
-
-    def render(self, console: Console, options: ConsoleOptions) -> RenderMap:
-        """Render the sub_layouts.
-
-        Args:
-            console (Console): Console instance.
-            options (ConsoleOptions): Console options.
-
-        Returns:
-            RenderMap: A dict that maps Layout on to a tuple of Region, lines
-        """
-        render_width = options.max_width
-        render_height = options.height or console.height
-        region_map = self._make_region_map(render_width, render_height)
-        layout_regions = [
-            (layout, region)
-            for layout, region in region_map.items()
-            if not layout.children
-        ]
-        render_map: Dict["Layout", "LayoutRender"] = {}
-        render_lines = console.render_lines
-        update_dimensions = options.update_dimensions
-
-        for layout, region in layout_regions:
-            lines = render_lines(
-                layout.renderable, update_dimensions(region.width, region.height)
-            )
-            render_map[layout] = LayoutRender(region, lines)
-        return render_map
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        with self._lock:
-            width = options.max_width or console.width
-            height = options.height or console.height
-            render_map = self.render(console, options.update_dimensions(width, height))
-            self._render_map = render_map
-            layout_lines: List[List[Segment]] = [[] for _ in range(height)]
-            _islice = islice
-            for region, lines in render_map.values():
-                _x, y, _layout_width, layout_height = region
-                for row, line in zip(
-                    _islice(layout_lines, y, y + layout_height), lines
-                ):
-                    row.extend(line)
-
-            new_line = Segment.line()
-            for layout_row in layout_lines:
-                yield from layout_row
-                yield new_line
-
-
-if __name__ == "__main__":
-    from rich.console import Console
-
-    console = Console()
-    layout = Layout()
-
-    layout.split_column(
-        Layout(name="header", size=3),
-        Layout(ratio=1, name="main"),
-        Layout(size=10, name="footer"),
-    )
-
-    layout["main"].split_row(Layout(name="side"), Layout(name="body", ratio=2))
-
-    layout["body"].split_row(Layout(name="content", ratio=2), Layout(name="s2"))
-
-    layout["s2"].split_column(
-        Layout(name="top"), Layout(name="middle"), Layout(name="bottom")
-    )
-
-    layout["side"].split_column(Layout(layout.tree, name="left1"), Layout(name="left2"))
-
-    layout["content"].update("foo")
-
-    console.print(layout)
diff --git a/venv/lib/python3.10/site-packages/rich/live.py b/venv/lib/python3.10/site-packages/rich/live.py
deleted file mode 100644
index 8738cf09f494e12ab9a38119f3558ae9419de7e6..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/live.py
+++ /dev/null
@@ -1,375 +0,0 @@
-import sys
-from threading import Event, RLock, Thread
-from types import TracebackType
-from typing import IO, Any, Callable, List, Optional, TextIO, Type, cast
-
-from . import get_console
-from .console import Console, ConsoleRenderable, RenderableType, RenderHook
-from .control import Control
-from .file_proxy import FileProxy
-from .jupyter import JupyterMixin
-from .live_render import LiveRender, VerticalOverflowMethod
-from .screen import Screen
-from .text import Text
-
-
-class _RefreshThread(Thread):
-    """A thread that calls refresh() at regular intervals."""
-
-    def __init__(self, live: "Live", refresh_per_second: float) -> None:
-        self.live = live
-        self.refresh_per_second = refresh_per_second
-        self.done = Event()
-        super().__init__(daemon=True)
-
-    def stop(self) -> None:
-        self.done.set()
-
-    def run(self) -> None:
-        while not self.done.wait(1 / self.refresh_per_second):
-            with self.live._lock:
-                if not self.done.is_set():
-                    self.live.refresh()
-
-
-class Live(JupyterMixin, RenderHook):
-    """Renders an auto-updating live display of any given renderable.
-
-    Args:
-        renderable (RenderableType, optional): The renderable to live display. Defaults to displaying nothing.
-        console (Console, optional): Optional Console instance. Defaults to an internal Console instance writing to stdout.
-        screen (bool, optional): Enable alternate screen mode. Defaults to False.
-        auto_refresh (bool, optional): Enable auto refresh. If disabled, you will need to call `refresh()` or `update()` with refresh flag. Defaults to True
-        refresh_per_second (float, optional): Number of times per second to refresh the live display. Defaults to 4.
-        transient (bool, optional): Clear the renderable on exit (has no effect when screen=True). Defaults to False.
-        redirect_stdout (bool, optional): Enable redirection of stdout, so ``print`` may be used. Defaults to True.
-        redirect_stderr (bool, optional): Enable redirection of stderr. Defaults to True.
-        vertical_overflow (VerticalOverflowMethod, optional): How to handle renderable when it is too tall for the console. Defaults to "ellipsis".
-        get_renderable (Callable[[], RenderableType], optional): Optional callable to get renderable. Defaults to None.
-    """
-
-    def __init__(
-        self,
-        renderable: Optional[RenderableType] = None,
-        *,
-        console: Optional[Console] = None,
-        screen: bool = False,
-        auto_refresh: bool = True,
-        refresh_per_second: float = 4,
-        transient: bool = False,
-        redirect_stdout: bool = True,
-        redirect_stderr: bool = True,
-        vertical_overflow: VerticalOverflowMethod = "ellipsis",
-        get_renderable: Optional[Callable[[], RenderableType]] = None,
-    ) -> None:
-        assert refresh_per_second > 0, "refresh_per_second must be > 0"
-        self._renderable = renderable
-        self.console = console if console is not None else get_console()
-        self._screen = screen
-        self._alt_screen = False
-
-        self._redirect_stdout = redirect_stdout
-        self._redirect_stderr = redirect_stderr
-        self._restore_stdout: Optional[IO[str]] = None
-        self._restore_stderr: Optional[IO[str]] = None
-
-        self._lock = RLock()
-        self.ipy_widget: Optional[Any] = None
-        self.auto_refresh = auto_refresh
-        self._started: bool = False
-        self.transient = True if screen else transient
-
-        self._refresh_thread: Optional[_RefreshThread] = None
-        self.refresh_per_second = refresh_per_second
-
-        self.vertical_overflow = vertical_overflow
-        self._get_renderable = get_renderable
-        self._live_render = LiveRender(
-            self.get_renderable(), vertical_overflow=vertical_overflow
-        )
-
-    @property
-    def is_started(self) -> bool:
-        """Check if live display has been started."""
-        return self._started
-
-    def get_renderable(self) -> RenderableType:
-        renderable = (
-            self._get_renderable()
-            if self._get_renderable is not None
-            else self._renderable
-        )
-        return renderable or ""
-
-    def start(self, refresh: bool = False) -> None:
-        """Start live rendering display.
-
-        Args:
-            refresh (bool, optional): Also refresh. Defaults to False.
-        """
-        with self._lock:
-            if self._started:
-                return
-            self.console.set_live(self)
-            self._started = True
-            if self._screen:
-                self._alt_screen = self.console.set_alt_screen(True)
-            self.console.show_cursor(False)
-            self._enable_redirect_io()
-            self.console.push_render_hook(self)
-            if refresh:
-                try:
-                    self.refresh()
-                except Exception:
-                    # If refresh fails, we want to stop the redirection of sys.stderr,
-                    # so the error stacktrace is properly displayed in the terminal.
-                    # (or, if the code that calls Rich captures the exception and wants to display something,
-                    # let this be displayed in the terminal).
-                    self.stop()
-                    raise
-            if self.auto_refresh:
-                self._refresh_thread = _RefreshThread(self, self.refresh_per_second)
-                self._refresh_thread.start()
-
-    def stop(self) -> None:
-        """Stop live rendering display."""
-        with self._lock:
-            if not self._started:
-                return
-            self.console.clear_live()
-            self._started = False
-
-            if self.auto_refresh and self._refresh_thread is not None:
-                self._refresh_thread.stop()
-                self._refresh_thread = None
-            # allow it to fully render on the last even if overflow
-            self.vertical_overflow = "visible"
-            with self.console:
-                try:
-                    if not self._alt_screen and not self.console.is_jupyter:
-                        self.refresh()
-                finally:
-                    self._disable_redirect_io()
-                    self.console.pop_render_hook()
-                    if not self._alt_screen and self.console.is_terminal:
-                        self.console.line()
-                    self.console.show_cursor(True)
-                    if self._alt_screen:
-                        self.console.set_alt_screen(False)
-
-                    if self.transient and not self._alt_screen:
-                        self.console.control(self._live_render.restore_cursor())
-                    if self.ipy_widget is not None and self.transient:
-                        self.ipy_widget.close()  # pragma: no cover
-
-    def __enter__(self) -> "Live":
-        self.start(refresh=self._renderable is not None)
-        return self
-
-    def __exit__(
-        self,
-        exc_type: Optional[Type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
-    ) -> None:
-        self.stop()
-
-    def _enable_redirect_io(self) -> None:
-        """Enable redirecting of stdout / stderr."""
-        if self.console.is_terminal or self.console.is_jupyter:
-            if self._redirect_stdout and not isinstance(sys.stdout, FileProxy):
-                self._restore_stdout = sys.stdout
-                sys.stdout = cast("TextIO", FileProxy(self.console, sys.stdout))
-            if self._redirect_stderr and not isinstance(sys.stderr, FileProxy):
-                self._restore_stderr = sys.stderr
-                sys.stderr = cast("TextIO", FileProxy(self.console, sys.stderr))
-
-    def _disable_redirect_io(self) -> None:
-        """Disable redirecting of stdout / stderr."""
-        if self._restore_stdout:
-            sys.stdout = cast("TextIO", self._restore_stdout)
-            self._restore_stdout = None
-        if self._restore_stderr:
-            sys.stderr = cast("TextIO", self._restore_stderr)
-            self._restore_stderr = None
-
-    @property
-    def renderable(self) -> RenderableType:
-        """Get the renderable that is being displayed
-
-        Returns:
-            RenderableType: Displayed renderable.
-        """
-        renderable = self.get_renderable()
-        return Screen(renderable) if self._alt_screen else renderable
-
-    def update(self, renderable: RenderableType, *, refresh: bool = False) -> None:
-        """Update the renderable that is being displayed
-
-        Args:
-            renderable (RenderableType): New renderable to use.
-            refresh (bool, optional): Refresh the display. Defaults to False.
-        """
-        if isinstance(renderable, str):
-            renderable = self.console.render_str(renderable)
-        with self._lock:
-            self._renderable = renderable
-            if refresh:
-                self.refresh()
-
-    def refresh(self) -> None:
-        """Update the display of the Live Render."""
-        with self._lock:
-            self._live_render.set_renderable(self.renderable)
-            if self.console.is_jupyter:  # pragma: no cover
-                try:
-                    from IPython.display import display
-                    from ipywidgets import Output
-                except ImportError:
-                    import warnings
-
-                    warnings.warn('install "ipywidgets" for Jupyter support')
-                else:
-                    if self.ipy_widget is None:
-                        self.ipy_widget = Output()
-                        display(self.ipy_widget)
-
-                    with self.ipy_widget:
-                        self.ipy_widget.clear_output(wait=True)
-                        self.console.print(self._live_render.renderable)
-            elif self.console.is_terminal and not self.console.is_dumb_terminal:
-                with self.console:
-                    self.console.print(Control())
-            elif (
-                not self._started and not self.transient
-            ):  # if it is finished allow files or dumb-terminals to see final result
-                with self.console:
-                    self.console.print(Control())
-
-    def process_renderables(
-        self, renderables: List[ConsoleRenderable]
-    ) -> List[ConsoleRenderable]:
-        """Process renderables to restore cursor and display progress."""
-        self._live_render.vertical_overflow = self.vertical_overflow
-        if self.console.is_interactive:
-            # lock needs acquiring as user can modify live_render renderable at any time unlike in Progress.
-            with self._lock:
-                reset = (
-                    Control.home()
-                    if self._alt_screen
-                    else self._live_render.position_cursor()
-                )
-                renderables = [reset, *renderables, self._live_render]
-        elif (
-            not self._started and not self.transient
-        ):  # if it is finished render the final output for files or dumb_terminals
-            renderables = [*renderables, self._live_render]
-
-        return renderables
-
-
-if __name__ == "__main__":  # pragma: no cover
-    import random
-    import time
-    from itertools import cycle
-    from typing import Dict, List, Tuple
-
-    from .align import Align
-    from .console import Console
-    from .live import Live as Live
-    from .panel import Panel
-    from .rule import Rule
-    from .syntax import Syntax
-    from .table import Table
-
-    console = Console()
-
-    syntax = Syntax(
-        '''def loop_last(values: Iterable[T]) -> Iterable[Tuple[bool, T]]:
-    """Iterate and generate a tuple with a flag for last value."""
-    iter_values = iter(values)
-    try:
-        previous_value = next(iter_values)
-    except StopIteration:
-        return
-    for value in iter_values:
-        yield False, previous_value
-        previous_value = value
-    yield True, previous_value''',
-        "python",
-        line_numbers=True,
-    )
-
-    table = Table("foo", "bar", "baz")
-    table.add_row("1", "2", "3")
-
-    progress_renderables = [
-        "You can make the terminal shorter and taller to see the live table hide"
-        "Text may be printed while the progress bars are rendering.",
-        Panel("In fact, [i]any[/i] renderable will work"),
-        "Such as [magenta]tables[/]...",
-        table,
-        "Pretty printed structures...",
-        {"type": "example", "text": "Pretty printed"},
-        "Syntax...",
-        syntax,
-        Rule("Give it a try!"),
-    ]
-
-    examples = cycle(progress_renderables)
-
-    exchanges = [
-        "SGD",
-        "MYR",
-        "EUR",
-        "USD",
-        "AUD",
-        "JPY",
-        "CNH",
-        "HKD",
-        "CAD",
-        "INR",
-        "DKK",
-        "GBP",
-        "RUB",
-        "NZD",
-        "MXN",
-        "IDR",
-        "TWD",
-        "THB",
-        "VND",
-    ]
-    with Live(console=console) as live_table:
-        exchange_rate_dict: Dict[Tuple[str, str], float] = {}
-
-        for index in range(100):
-            select_exchange = exchanges[index % len(exchanges)]
-
-            for exchange in exchanges:
-                if exchange == select_exchange:
-                    continue
-                time.sleep(0.4)
-                if random.randint(0, 10) < 1:
-                    console.log(next(examples))
-                exchange_rate_dict[(select_exchange, exchange)] = 200 / (
-                    (random.random() * 320) + 1
-                )
-                if len(exchange_rate_dict) > len(exchanges) - 1:
-                    exchange_rate_dict.pop(list(exchange_rate_dict.keys())[0])
-                table = Table(title="Exchange Rates")
-
-                table.add_column("Source Currency")
-                table.add_column("Destination Currency")
-                table.add_column("Exchange Rate")
-
-                for (source, dest), exchange_rate in exchange_rate_dict.items():
-                    table.add_row(
-                        source,
-                        dest,
-                        Text(
-                            f"{exchange_rate:.4f}",
-                            style="red" if exchange_rate < 1.0 else "green",
-                        ),
-                    )
-
-                live_table.update(Align.center(table))
diff --git a/venv/lib/python3.10/site-packages/rich/live_render.py b/venv/lib/python3.10/site-packages/rich/live_render.py
deleted file mode 100644
index 4284cccc4b13d5b4a4d62e1838e0eae4a9b80471..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/live_render.py
+++ /dev/null
@@ -1,112 +0,0 @@
-import sys
-from typing import Optional, Tuple
-
-if sys.version_info >= (3, 8):
-    from typing import Literal
-else:
-    from typing_extensions import Literal  # pragma: no cover
-
-
-from ._loop import loop_last
-from .console import Console, ConsoleOptions, RenderableType, RenderResult
-from .control import Control
-from .segment import ControlType, Segment
-from .style import StyleType
-from .text import Text
-
-VerticalOverflowMethod = Literal["crop", "ellipsis", "visible"]
-
-
-class LiveRender:
-    """Creates a renderable that may be updated.
-
-    Args:
-        renderable (RenderableType): Any renderable object.
-        style (StyleType, optional): An optional style to apply to the renderable. Defaults to "".
-    """
-
-    def __init__(
-        self,
-        renderable: RenderableType,
-        style: StyleType = "",
-        vertical_overflow: VerticalOverflowMethod = "ellipsis",
-    ) -> None:
-        self.renderable = renderable
-        self.style = style
-        self.vertical_overflow = vertical_overflow
-        self._shape: Optional[Tuple[int, int]] = None
-
-    def set_renderable(self, renderable: RenderableType) -> None:
-        """Set a new renderable.
-
-        Args:
-            renderable (RenderableType): Any renderable object, including str.
-        """
-        self.renderable = renderable
-
-    def position_cursor(self) -> Control:
-        """Get control codes to move cursor to beginning of live render.
-
-        Returns:
-            Control: A control instance that may be printed.
-        """
-        if self._shape is not None:
-            _, height = self._shape
-            return Control(
-                ControlType.CARRIAGE_RETURN,
-                (ControlType.ERASE_IN_LINE, 2),
-                *(
-                    (
-                        (ControlType.CURSOR_UP, 1),
-                        (ControlType.ERASE_IN_LINE, 2),
-                    )
-                    * (height - 1)
-                )
-            )
-        return Control()
-
-    def restore_cursor(self) -> Control:
-        """Get control codes to clear the render and restore the cursor to its previous position.
-
-        Returns:
-            Control: A Control instance that may be printed.
-        """
-        if self._shape is not None:
-            _, height = self._shape
-            return Control(
-                ControlType.CARRIAGE_RETURN,
-                *((ControlType.CURSOR_UP, 1), (ControlType.ERASE_IN_LINE, 2)) * height
-            )
-        return Control()
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        renderable = self.renderable
-        style = console.get_style(self.style)
-        lines = console.render_lines(renderable, options, style=style, pad=False)
-        shape = Segment.get_shape(lines)
-
-        _, height = shape
-        if height > options.size.height:
-            if self.vertical_overflow == "crop":
-                lines = lines[: options.size.height]
-                shape = Segment.get_shape(lines)
-            elif self.vertical_overflow == "ellipsis":
-                lines = lines[: (options.size.height - 1)]
-                overflow_text = Text(
-                    "...",
-                    overflow="crop",
-                    justify="center",
-                    end="",
-                    style="live.ellipsis",
-                )
-                lines.append(list(console.render(overflow_text)))
-                shape = Segment.get_shape(lines)
-        self._shape = shape
-
-        new_line = Segment.line()
-        for last, line in loop_last(lines):
-            yield from line
-            if not last:
-                yield new_line
diff --git a/venv/lib/python3.10/site-packages/rich/logging.py b/venv/lib/python3.10/site-packages/rich/logging.py
deleted file mode 100644
index b2624cd54a9a967990eb1e9803ce47facca8e2e4..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/logging.py
+++ /dev/null
@@ -1,297 +0,0 @@
-import logging
-from datetime import datetime
-from logging import Handler, LogRecord
-from pathlib import Path
-from types import ModuleType
-from typing import ClassVar, Iterable, List, Optional, Type, Union
-
-from rich._null_file import NullFile
-
-from . import get_console
-from ._log_render import FormatTimeCallable, LogRender
-from .console import Console, ConsoleRenderable
-from .highlighter import Highlighter, ReprHighlighter
-from .text import Text
-from .traceback import Traceback
-
-
-class RichHandler(Handler):
-    """A logging handler that renders output with Rich. The time / level / message and file are displayed in columns.
-    The level is color coded, and the message is syntax highlighted.
-
-    Note:
-        Be careful when enabling console markup in log messages if you have configured logging for libraries not
-        under your control. If a dependency writes messages containing square brackets, it may not produce the intended output.
-
-    Args:
-        level (Union[int, str], optional): Log level. Defaults to logging.NOTSET.
-        console (:class:`~rich.console.Console`, optional): Optional console instance to write logs.
-            Default will use a global console instance writing to stdout.
-        show_time (bool, optional): Show a column for the time. Defaults to True.
-        omit_repeated_times (bool, optional): Omit repetition of the same time. Defaults to True.
-        show_level (bool, optional): Show a column for the level. Defaults to True.
-        show_path (bool, optional): Show the path to the original log call. Defaults to True.
-        enable_link_path (bool, optional): Enable terminal link of path column to file. Defaults to True.
-        highlighter (Highlighter, optional): Highlighter to style log messages, or None to use ReprHighlighter. Defaults to None.
-        markup (bool, optional): Enable console markup in log messages. Defaults to False.
-        rich_tracebacks (bool, optional): Enable rich tracebacks with syntax highlighting and formatting. Defaults to False.
-        tracebacks_width (Optional[int], optional): Number of characters used to render tracebacks, or None for full width. Defaults to None.
-        tracebacks_code_width (int, optional): Number of code characters used to render tracebacks, or None for full width. Defaults to 88.
-        tracebacks_extra_lines (int, optional): Additional lines of code to render tracebacks, or None for full width. Defaults to None.
-        tracebacks_theme (str, optional): Override pygments theme used in traceback.
-        tracebacks_word_wrap (bool, optional): Enable word wrapping of long tracebacks lines. Defaults to True.
-        tracebacks_show_locals (bool, optional): Enable display of locals in tracebacks. Defaults to False.
-        tracebacks_suppress (Sequence[Union[str, ModuleType]]): Optional sequence of modules or paths to exclude from traceback.
-        tracebacks_max_frames (int, optional): Optional maximum number of frames returned by traceback.
-        locals_max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
-            Defaults to 10.
-        locals_max_string (int, optional): Maximum length of string before truncating, or None to disable. Defaults to 80.
-        log_time_format (Union[str, TimeFormatterCallable], optional): If ``log_time`` is enabled, either string for strftime or callable that formats the time. Defaults to "[%x %X] ".
-        keywords (List[str], optional): List of words to highlight instead of ``RichHandler.KEYWORDS``.
-    """
-
-    KEYWORDS: ClassVar[Optional[List[str]]] = [
-        "GET",
-        "POST",
-        "HEAD",
-        "PUT",
-        "DELETE",
-        "OPTIONS",
-        "TRACE",
-        "PATCH",
-    ]
-    HIGHLIGHTER_CLASS: ClassVar[Type[Highlighter]] = ReprHighlighter
-
-    def __init__(
-        self,
-        level: Union[int, str] = logging.NOTSET,
-        console: Optional[Console] = None,
-        *,
-        show_time: bool = True,
-        omit_repeated_times: bool = True,
-        show_level: bool = True,
-        show_path: bool = True,
-        enable_link_path: bool = True,
-        highlighter: Optional[Highlighter] = None,
-        markup: bool = False,
-        rich_tracebacks: bool = False,
-        tracebacks_width: Optional[int] = None,
-        tracebacks_code_width: int = 88,
-        tracebacks_extra_lines: int = 3,
-        tracebacks_theme: Optional[str] = None,
-        tracebacks_word_wrap: bool = True,
-        tracebacks_show_locals: bool = False,
-        tracebacks_suppress: Iterable[Union[str, ModuleType]] = (),
-        tracebacks_max_frames: int = 100,
-        locals_max_length: int = 10,
-        locals_max_string: int = 80,
-        log_time_format: Union[str, FormatTimeCallable] = "[%x %X]",
-        keywords: Optional[List[str]] = None,
-    ) -> None:
-        super().__init__(level=level)
-        self.console = console or get_console()
-        self.highlighter = highlighter or self.HIGHLIGHTER_CLASS()
-        self._log_render = LogRender(
-            show_time=show_time,
-            show_level=show_level,
-            show_path=show_path,
-            time_format=log_time_format,
-            omit_repeated_times=omit_repeated_times,
-            level_width=None,
-        )
-        self.enable_link_path = enable_link_path
-        self.markup = markup
-        self.rich_tracebacks = rich_tracebacks
-        self.tracebacks_width = tracebacks_width
-        self.tracebacks_extra_lines = tracebacks_extra_lines
-        self.tracebacks_theme = tracebacks_theme
-        self.tracebacks_word_wrap = tracebacks_word_wrap
-        self.tracebacks_show_locals = tracebacks_show_locals
-        self.tracebacks_suppress = tracebacks_suppress
-        self.tracebacks_max_frames = tracebacks_max_frames
-        self.tracebacks_code_width = tracebacks_code_width
-        self.locals_max_length = locals_max_length
-        self.locals_max_string = locals_max_string
-        self.keywords = keywords
-
-    def get_level_text(self, record: LogRecord) -> Text:
-        """Get the level name from the record.
-
-        Args:
-            record (LogRecord): LogRecord instance.
-
-        Returns:
-            Text: A tuple of the style and level name.
-        """
-        level_name = record.levelname
-        level_text = Text.styled(
-            level_name.ljust(8), f"logging.level.{level_name.lower()}"
-        )
-        return level_text
-
-    def emit(self, record: LogRecord) -> None:
-        """Invoked by logging."""
-        message = self.format(record)
-        traceback = None
-        if (
-            self.rich_tracebacks
-            and record.exc_info
-            and record.exc_info != (None, None, None)
-        ):
-            exc_type, exc_value, exc_traceback = record.exc_info
-            assert exc_type is not None
-            assert exc_value is not None
-            traceback = Traceback.from_exception(
-                exc_type,
-                exc_value,
-                exc_traceback,
-                width=self.tracebacks_width,
-                code_width=self.tracebacks_code_width,
-                extra_lines=self.tracebacks_extra_lines,
-                theme=self.tracebacks_theme,
-                word_wrap=self.tracebacks_word_wrap,
-                show_locals=self.tracebacks_show_locals,
-                locals_max_length=self.locals_max_length,
-                locals_max_string=self.locals_max_string,
-                suppress=self.tracebacks_suppress,
-                max_frames=self.tracebacks_max_frames,
-            )
-            message = record.getMessage()
-            if self.formatter:
-                record.message = record.getMessage()
-                formatter = self.formatter
-                if hasattr(formatter, "usesTime") and formatter.usesTime():
-                    record.asctime = formatter.formatTime(record, formatter.datefmt)
-                message = formatter.formatMessage(record)
-
-        message_renderable = self.render_message(record, message)
-        log_renderable = self.render(
-            record=record, traceback=traceback, message_renderable=message_renderable
-        )
-        if isinstance(self.console.file, NullFile):
-            # Handles pythonw, where stdout/stderr are null, and we return NullFile
-            # instance from Console.file. In this case, we still want to make a log record
-            # even though we won't be writing anything to a file.
-            self.handleError(record)
-        else:
-            try:
-                self.console.print(log_renderable)
-            except Exception:
-                self.handleError(record)
-
-    def render_message(self, record: LogRecord, message: str) -> "ConsoleRenderable":
-        """Render message text in to Text.
-
-        Args:
-            record (LogRecord): logging Record.
-            message (str): String containing log message.
-
-        Returns:
-            ConsoleRenderable: Renderable to display log message.
-        """
-        use_markup = getattr(record, "markup", self.markup)
-        message_text = Text.from_markup(message) if use_markup else Text(message)
-
-        highlighter = getattr(record, "highlighter", self.highlighter)
-        if highlighter:
-            message_text = highlighter(message_text)
-
-        if self.keywords is None:
-            self.keywords = self.KEYWORDS
-
-        if self.keywords:
-            message_text.highlight_words(self.keywords, "logging.keyword")
-
-        return message_text
-
-    def render(
-        self,
-        *,
-        record: LogRecord,
-        traceback: Optional[Traceback],
-        message_renderable: "ConsoleRenderable",
-    ) -> "ConsoleRenderable":
-        """Render log for display.
-
-        Args:
-            record (LogRecord): logging Record.
-            traceback (Optional[Traceback]): Traceback instance or None for no Traceback.
-            message_renderable (ConsoleRenderable): Renderable (typically Text) containing log message contents.
-
-        Returns:
-            ConsoleRenderable: Renderable to display log.
-        """
-        path = Path(record.pathname).name
-        level = self.get_level_text(record)
-        time_format = None if self.formatter is None else self.formatter.datefmt
-        log_time = datetime.fromtimestamp(record.created)
-
-        log_renderable = self._log_render(
-            self.console,
-            [message_renderable] if not traceback else [message_renderable, traceback],
-            log_time=log_time,
-            time_format=time_format,
-            level=level,
-            path=path,
-            line_no=record.lineno,
-            link_path=record.pathname if self.enable_link_path else None,
-        )
-        return log_renderable
-
-
-if __name__ == "__main__":  # pragma: no cover
-    from time import sleep
-
-    FORMAT = "%(message)s"
-    # FORMAT = "%(asctime)-15s - %(levelname)s - %(message)s"
-    logging.basicConfig(
-        level="NOTSET",
-        format=FORMAT,
-        datefmt="[%X]",
-        handlers=[RichHandler(rich_tracebacks=True, tracebacks_show_locals=True)],
-    )
-    log = logging.getLogger("rich")
-
-    log.info("Server starting...")
-    log.info("Listening on http://127.0.0.1:8080")
-    sleep(1)
-
-    log.info("GET /index.html 200 1298")
-    log.info("GET /imgs/backgrounds/back1.jpg 200 54386")
-    log.info("GET /css/styles.css 200 54386")
-    log.warning("GET /favicon.ico 404 242")
-    sleep(1)
-
-    log.debug(
-        "JSONRPC request\n--> %r\n<-- %r",
-        {
-            "version": "1.1",
-            "method": "confirmFruitPurchase",
-            "params": [["apple", "orange", "mangoes", "pomelo"], 1.123],
-            "id": "194521489",
-        },
-        {"version": "1.1", "result": True, "error": None, "id": "194521489"},
-    )
-    log.debug(
-        "Loading configuration file /adasd/asdasd/qeqwe/qwrqwrqwr/sdgsdgsdg/werwerwer/dfgerert/ertertert/ertetert/werwerwer"
-    )
-    log.error("Unable to find 'pomelo' in database!")
-    log.info("POST /jsonrpc/ 200 65532")
-    log.info("POST /admin/ 401 42234")
-    log.warning("password was rejected for admin site.")
-
-    def divide() -> None:
-        number = 1
-        divisor = 0
-        foos = ["foo"] * 100
-        log.debug("in divide")
-        try:
-            number / divisor
-        except:
-            log.exception("An error of some kind occurred!")
-
-    divide()
-    sleep(1)
-    log.critical("Out of memory!")
-    log.info("Server exited with code=-1")
-    log.info("[bold]EXITING...[/bold]", extra=dict(markup=True))
diff --git a/venv/lib/python3.10/site-packages/rich/markdown.py b/venv/lib/python3.10/site-packages/rich/markdown.py
deleted file mode 100644
index 26c58d1558bc151489178b0cdd05d52990d6d493..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/markdown.py
+++ /dev/null
@@ -1,784 +0,0 @@
-from __future__ import annotations
-
-import sys
-from typing import ClassVar, Iterable
-
-from markdown_it import MarkdownIt
-from markdown_it.token import Token
-
-if sys.version_info >= (3, 8):
-    from typing import get_args
-else:
-    from typing_extensions import get_args  # pragma: no cover
-
-from rich.table import Table
-
-from . import box
-from ._loop import loop_first
-from ._stack import Stack
-from .console import Console, ConsoleOptions, JustifyMethod, RenderResult
-from .containers import Renderables
-from .jupyter import JupyterMixin
-from .panel import Panel
-from .rule import Rule
-from .segment import Segment
-from .style import Style, StyleStack
-from .syntax import Syntax
-from .text import Text, TextType
-
-
-class MarkdownElement:
-    new_line: ClassVar[bool] = True
-
-    @classmethod
-    def create(cls, markdown: Markdown, token: Token) -> MarkdownElement:
-        """Factory to create markdown element,
-
-        Args:
-            markdown (Markdown): The parent Markdown object.
-            token (Token): A node from markdown-it.
-
-        Returns:
-            MarkdownElement: A new markdown element
-        """
-        return cls()
-
-    def on_enter(self, context: MarkdownContext) -> None:
-        """Called when the node is entered.
-
-        Args:
-            context (MarkdownContext): The markdown context.
-        """
-
-    def on_text(self, context: MarkdownContext, text: TextType) -> None:
-        """Called when text is parsed.
-
-        Args:
-            context (MarkdownContext): The markdown context.
-        """
-
-    def on_leave(self, context: MarkdownContext) -> None:
-        """Called when the parser leaves the element.
-
-        Args:
-            context (MarkdownContext): [description]
-        """
-
-    def on_child_close(self, context: MarkdownContext, child: MarkdownElement) -> bool:
-        """Called when a child element is closed.
-
-        This method allows a parent element to take over rendering of its children.
-
-        Args:
-            context (MarkdownContext): The markdown context.
-            child (MarkdownElement): The child markdown element.
-
-        Returns:
-            bool: Return True to render the element, or False to not render the element.
-        """
-        return True
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        return ()
-
-
-class UnknownElement(MarkdownElement):
-    """An unknown element.
-
-    Hopefully there will be no unknown elements, and we will have a MarkdownElement for
-    everything in the document.
-
-    """
-
-
-class TextElement(MarkdownElement):
-    """Base class for elements that render text."""
-
-    style_name = "none"
-
-    def on_enter(self, context: MarkdownContext) -> None:
-        self.style = context.enter_style(self.style_name)
-        self.text = Text(justify="left")
-
-    def on_text(self, context: MarkdownContext, text: TextType) -> None:
-        self.text.append(text, context.current_style if isinstance(text, str) else None)
-
-    def on_leave(self, context: MarkdownContext) -> None:
-        context.leave_style()
-
-
-class Paragraph(TextElement):
-    """A Paragraph."""
-
-    style_name = "markdown.paragraph"
-    justify: JustifyMethod
-
-    @classmethod
-    def create(cls, markdown: Markdown, token: Token) -> Paragraph:
-        return cls(justify=markdown.justify or "left")
-
-    def __init__(self, justify: JustifyMethod) -> None:
-        self.justify = justify
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        self.text.justify = self.justify
-        yield self.text
-
-
-class Heading(TextElement):
-    """A heading."""
-
-    @classmethod
-    def create(cls, markdown: Markdown, token: Token) -> Heading:
-        return cls(token.tag)
-
-    def on_enter(self, context: MarkdownContext) -> None:
-        self.text = Text()
-        context.enter_style(self.style_name)
-
-    def __init__(self, tag: str) -> None:
-        self.tag = tag
-        self.style_name = f"markdown.{tag}"
-        super().__init__()
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        text = self.text
-        text.justify = "center"
-        if self.tag == "h1":
-            # Draw a border around h1s
-            yield Panel(
-                text,
-                box=box.HEAVY,
-                style="markdown.h1.border",
-            )
-        else:
-            # Styled text for h2 and beyond
-            if self.tag == "h2":
-                yield Text("")
-            yield text
-
-
-class CodeBlock(TextElement):
-    """A code block with syntax highlighting."""
-
-    style_name = "markdown.code_block"
-
-    @classmethod
-    def create(cls, markdown: Markdown, token: Token) -> CodeBlock:
-        node_info = token.info or ""
-        lexer_name = node_info.partition(" ")[0]
-        return cls(lexer_name or "text", markdown.code_theme)
-
-    def __init__(self, lexer_name: str, theme: str) -> None:
-        self.lexer_name = lexer_name
-        self.theme = theme
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        code = str(self.text).rstrip()
-        syntax = Syntax(
-            code, self.lexer_name, theme=self.theme, word_wrap=True, padding=1
-        )
-        yield syntax
-
-
-class BlockQuote(TextElement):
-    """A block quote."""
-
-    style_name = "markdown.block_quote"
-
-    def __init__(self) -> None:
-        self.elements: Renderables = Renderables()
-
-    def on_child_close(self, context: MarkdownContext, child: MarkdownElement) -> bool:
-        self.elements.append(child)
-        return False
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        render_options = options.update(width=options.max_width - 4)
-        lines = console.render_lines(self.elements, render_options, style=self.style)
-        style = self.style
-        new_line = Segment("\n")
-        padding = Segment("▌ ", style)
-        for line in lines:
-            yield padding
-            yield from line
-            yield new_line
-
-
-class HorizontalRule(MarkdownElement):
-    """A horizontal rule to divide sections."""
-
-    new_line = False
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        style = console.get_style("markdown.hr", default="none")
-        yield Rule(style=style)
-
-
-class TableElement(MarkdownElement):
-    """MarkdownElement corresponding to `table_open`."""
-
-    def __init__(self) -> None:
-        self.header: TableHeaderElement | None = None
-        self.body: TableBodyElement | None = None
-
-    def on_child_close(self, context: MarkdownContext, child: MarkdownElement) -> bool:
-        if isinstance(child, TableHeaderElement):
-            self.header = child
-        elif isinstance(child, TableBodyElement):
-            self.body = child
-        else:
-            raise RuntimeError("Couldn't process markdown table.")
-        return False
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        table = Table(box=box.SIMPLE_HEAVY)
-
-        if self.header is not None and self.header.row is not None:
-            for column in self.header.row.cells:
-                table.add_column(column.content)
-
-        if self.body is not None:
-            for row in self.body.rows:
-                row_content = [element.content for element in row.cells]
-                table.add_row(*row_content)
-
-        yield table
-
-
-class TableHeaderElement(MarkdownElement):
-    """MarkdownElement corresponding to `thead_open` and `thead_close`."""
-
-    def __init__(self) -> None:
-        self.row: TableRowElement | None = None
-
-    def on_child_close(self, context: MarkdownContext, child: MarkdownElement) -> bool:
-        assert isinstance(child, TableRowElement)
-        self.row = child
-        return False
-
-
-class TableBodyElement(MarkdownElement):
-    """MarkdownElement corresponding to `tbody_open` and `tbody_close`."""
-
-    def __init__(self) -> None:
-        self.rows: list[TableRowElement] = []
-
-    def on_child_close(self, context: MarkdownContext, child: MarkdownElement) -> bool:
-        assert isinstance(child, TableRowElement)
-        self.rows.append(child)
-        return False
-
-
-class TableRowElement(MarkdownElement):
-    """MarkdownElement corresponding to `tr_open` and `tr_close`."""
-
-    def __init__(self) -> None:
-        self.cells: list[TableDataElement] = []
-
-    def on_child_close(self, context: MarkdownContext, child: MarkdownElement) -> bool:
-        assert isinstance(child, TableDataElement)
-        self.cells.append(child)
-        return False
-
-
-class TableDataElement(MarkdownElement):
-    """MarkdownElement corresponding to `td_open` and `td_close`
-    and `th_open` and `th_close`."""
-
-    @classmethod
-    def create(cls, markdown: Markdown, token: Token) -> MarkdownElement:
-        style = str(token.attrs.get("style")) or ""
-
-        justify: JustifyMethod
-        if "text-align:right" in style:
-            justify = "right"
-        elif "text-align:center" in style:
-            justify = "center"
-        elif "text-align:left" in style:
-            justify = "left"
-        else:
-            justify = "default"
-
-        assert justify in get_args(JustifyMethod)
-        return cls(justify=justify)
-
-    def __init__(self, justify: JustifyMethod) -> None:
-        self.content: Text = Text("", justify=justify)
-        self.justify = justify
-
-    def on_text(self, context: MarkdownContext, text: TextType) -> None:
-        text = Text(text) if isinstance(text, str) else text
-        text.stylize(context.current_style)
-        self.content.append_text(text)
-
-
-class ListElement(MarkdownElement):
-    """A list element."""
-
-    @classmethod
-    def create(cls, markdown: Markdown, token: Token) -> ListElement:
-        return cls(token.type, int(token.attrs.get("start", 1)))
-
-    def __init__(self, list_type: str, list_start: int | None) -> None:
-        self.items: list[ListItem] = []
-        self.list_type = list_type
-        self.list_start = list_start
-
-    def on_child_close(self, context: MarkdownContext, child: MarkdownElement) -> bool:
-        assert isinstance(child, ListItem)
-        self.items.append(child)
-        return False
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        if self.list_type == "bullet_list_open":
-            for item in self.items:
-                yield from item.render_bullet(console, options)
-        else:
-            number = 1 if self.list_start is None else self.list_start
-            last_number = number + len(self.items)
-            for index, item in enumerate(self.items):
-                yield from item.render_number(
-                    console, options, number + index, last_number
-                )
-
-
-class ListItem(TextElement):
-    """An item in a list."""
-
-    style_name = "markdown.item"
-
-    def __init__(self) -> None:
-        self.elements: Renderables = Renderables()
-
-    def on_child_close(self, context: MarkdownContext, child: MarkdownElement) -> bool:
-        self.elements.append(child)
-        return False
-
-    def render_bullet(self, console: Console, options: ConsoleOptions) -> RenderResult:
-        render_options = options.update(width=options.max_width - 3)
-        lines = console.render_lines(self.elements, render_options, style=self.style)
-        bullet_style = console.get_style("markdown.item.bullet", default="none")
-
-        bullet = Segment(" • ", bullet_style)
-        padding = Segment(" " * 3, bullet_style)
-        new_line = Segment("\n")
-        for first, line in loop_first(lines):
-            yield bullet if first else padding
-            yield from line
-            yield new_line
-
-    def render_number(
-        self, console: Console, options: ConsoleOptions, number: int, last_number: int
-    ) -> RenderResult:
-        number_width = len(str(last_number)) + 2
-        render_options = options.update(width=options.max_width - number_width)
-        lines = console.render_lines(self.elements, render_options, style=self.style)
-        number_style = console.get_style("markdown.item.number", default="none")
-
-        new_line = Segment("\n")
-        padding = Segment(" " * number_width, number_style)
-        numeral = Segment(f"{number}".rjust(number_width - 1) + " ", number_style)
-        for first, line in loop_first(lines):
-            yield numeral if first else padding
-            yield from line
-            yield new_line
-
-
-class Link(TextElement):
-    @classmethod
-    def create(cls, markdown: Markdown, token: Token) -> MarkdownElement:
-        url = token.attrs.get("href", "#")
-        return cls(token.content, str(url))
-
-    def __init__(self, text: str, href: str):
-        self.text = Text(text)
-        self.href = href
-
-
-class ImageItem(TextElement):
-    """Renders a placeholder for an image."""
-
-    new_line = False
-
-    @classmethod
-    def create(cls, markdown: Markdown, token: Token) -> MarkdownElement:
-        """Factory to create markdown element,
-
-        Args:
-            markdown (Markdown): The parent Markdown object.
-            token (Any): A token from markdown-it.
-
-        Returns:
-            MarkdownElement: A new markdown element
-        """
-        return cls(str(token.attrs.get("src", "")), markdown.hyperlinks)
-
-    def __init__(self, destination: str, hyperlinks: bool) -> None:
-        self.destination = destination
-        self.hyperlinks = hyperlinks
-        self.link: str | None = None
-        super().__init__()
-
-    def on_enter(self, context: MarkdownContext) -> None:
-        self.link = context.current_style.link
-        self.text = Text(justify="left")
-        super().on_enter(context)
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        link_style = Style(link=self.link or self.destination or None)
-        title = self.text or Text(self.destination.strip("/").rsplit("/", 1)[-1])
-        if self.hyperlinks:
-            title.stylize(link_style)
-        text = Text.assemble("🌆 ", title, " ", end="")
-        yield text
-
-
-class MarkdownContext:
-    """Manages the console render state."""
-
-    def __init__(
-        self,
-        console: Console,
-        options: ConsoleOptions,
-        style: Style,
-        inline_code_lexer: str | None = None,
-        inline_code_theme: str = "monokai",
-    ) -> None:
-        self.console = console
-        self.options = options
-        self.style_stack: StyleStack = StyleStack(style)
-        self.stack: Stack[MarkdownElement] = Stack()
-
-        self._syntax: Syntax | None = None
-        if inline_code_lexer is not None:
-            self._syntax = Syntax("", inline_code_lexer, theme=inline_code_theme)
-
-    @property
-    def current_style(self) -> Style:
-        """Current style which is the product of all styles on the stack."""
-        return self.style_stack.current
-
-    def on_text(self, text: str, node_type: str) -> None:
-        """Called when the parser visits text."""
-        if node_type in {"fence", "code_inline"} and self._syntax is not None:
-            highlight_text = self._syntax.highlight(text)
-            highlight_text.rstrip()
-            self.stack.top.on_text(
-                self, Text.assemble(highlight_text, style=self.style_stack.current)
-            )
-        else:
-            self.stack.top.on_text(self, text)
-
-    def enter_style(self, style_name: str | Style) -> Style:
-        """Enter a style context."""
-        style = self.console.get_style(style_name, default="none")
-        self.style_stack.push(style)
-        return self.current_style
-
-    def leave_style(self) -> Style:
-        """Leave a style context."""
-        style = self.style_stack.pop()
-        return style
-
-
-class Markdown(JupyterMixin):
-    """A Markdown renderable.
-
-    Args:
-        markup (str): A string containing markdown.
-        code_theme (str, optional): Pygments theme for code blocks. Defaults to "monokai". See https://pygments.org/styles/ for code themes.
-        justify (JustifyMethod, optional): Justify value for paragraphs. Defaults to None.
-        style (Union[str, Style], optional): Optional style to apply to markdown.
-        hyperlinks (bool, optional): Enable hyperlinks. Defaults to ``True``.
-        inline_code_lexer: (str, optional): Lexer to use if inline code highlighting is
-            enabled. Defaults to None.
-        inline_code_theme: (Optional[str], optional): Pygments theme for inline code
-            highlighting, or None for no highlighting. Defaults to None.
-    """
-
-    elements: ClassVar[dict[str, type[MarkdownElement]]] = {
-        "paragraph_open": Paragraph,
-        "heading_open": Heading,
-        "fence": CodeBlock,
-        "code_block": CodeBlock,
-        "blockquote_open": BlockQuote,
-        "hr": HorizontalRule,
-        "bullet_list_open": ListElement,
-        "ordered_list_open": ListElement,
-        "list_item_open": ListItem,
-        "image": ImageItem,
-        "table_open": TableElement,
-        "tbody_open": TableBodyElement,
-        "thead_open": TableHeaderElement,
-        "tr_open": TableRowElement,
-        "td_open": TableDataElement,
-        "th_open": TableDataElement,
-    }
-
-    inlines = {"em", "strong", "code", "s"}
-
-    def __init__(
-        self,
-        markup: str,
-        code_theme: str = "monokai",
-        justify: JustifyMethod | None = None,
-        style: str | Style = "none",
-        hyperlinks: bool = True,
-        inline_code_lexer: str | None = None,
-        inline_code_theme: str | None = None,
-    ) -> None:
-        parser = MarkdownIt().enable("strikethrough").enable("table")
-        self.markup = markup
-        self.parsed = parser.parse(markup)
-        self.code_theme = code_theme
-        self.justify: JustifyMethod | None = justify
-        self.style = style
-        self.hyperlinks = hyperlinks
-        self.inline_code_lexer = inline_code_lexer
-        self.inline_code_theme = inline_code_theme or code_theme
-
-    def _flatten_tokens(self, tokens: Iterable[Token]) -> Iterable[Token]:
-        """Flattens the token stream."""
-        for token in tokens:
-            is_fence = token.type == "fence"
-            is_image = token.tag == "img"
-            if token.children and not (is_image or is_fence):
-                yield from self._flatten_tokens(token.children)
-            else:
-                yield token
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        """Render markdown to the console."""
-        style = console.get_style(self.style, default="none")
-        options = options.update(height=None)
-        context = MarkdownContext(
-            console,
-            options,
-            style,
-            inline_code_lexer=self.inline_code_lexer,
-            inline_code_theme=self.inline_code_theme,
-        )
-        tokens = self.parsed
-        inline_style_tags = self.inlines
-        new_line = False
-        _new_line_segment = Segment.line()
-
-        for token in self._flatten_tokens(tokens):
-            node_type = token.type
-            tag = token.tag
-
-            entering = token.nesting == 1
-            exiting = token.nesting == -1
-            self_closing = token.nesting == 0
-
-            if node_type == "text":
-                context.on_text(token.content, node_type)
-            elif node_type == "hardbreak":
-                context.on_text("\n", node_type)
-            elif node_type == "softbreak":
-                context.on_text(" ", node_type)
-            elif node_type == "link_open":
-                href = str(token.attrs.get("href", ""))
-                if self.hyperlinks:
-                    link_style = console.get_style("markdown.link_url", default="none")
-                    link_style += Style(link=href)
-                    context.enter_style(link_style)
-                else:
-                    context.stack.push(Link.create(self, token))
-            elif node_type == "link_close":
-                if self.hyperlinks:
-                    context.leave_style()
-                else:
-                    element = context.stack.pop()
-                    assert isinstance(element, Link)
-                    link_style = console.get_style("markdown.link", default="none")
-                    context.enter_style(link_style)
-                    context.on_text(element.text.plain, node_type)
-                    context.leave_style()
-                    context.on_text(" (", node_type)
-                    link_url_style = console.get_style(
-                        "markdown.link_url", default="none"
-                    )
-                    context.enter_style(link_url_style)
-                    context.on_text(element.href, node_type)
-                    context.leave_style()
-                    context.on_text(")", node_type)
-            elif (
-                tag in inline_style_tags
-                and node_type != "fence"
-                and node_type != "code_block"
-            ):
-                if entering:
-                    # If it's an opening inline token e.g. strong, em, etc.
-                    # Then we move into a style context i.e. push to stack.
-                    context.enter_style(f"markdown.{tag}")
-                elif exiting:
-                    # If it's a closing inline style, then we pop the style
-                    # off of the stack, to move out of the context of it...
-                    context.leave_style()
-                else:
-                    # If it's a self-closing inline style e.g. `code_inline`
-                    context.enter_style(f"markdown.{tag}")
-                    if token.content:
-                        context.on_text(token.content, node_type)
-                    context.leave_style()
-            else:
-                # Map the markdown tag -> MarkdownElement renderable
-                element_class = self.elements.get(token.type) or UnknownElement
-                element = element_class.create(self, token)
-
-                if entering or self_closing:
-                    context.stack.push(element)
-                    element.on_enter(context)
-
-                if exiting:  # CLOSING tag
-                    element = context.stack.pop()
-
-                    should_render = not context.stack or (
-                        context.stack
-                        and context.stack.top.on_child_close(context, element)
-                    )
-
-                    if should_render:
-                        if new_line:
-                            yield _new_line_segment
-
-                        yield from console.render(element, context.options)
-                elif self_closing:  # SELF-CLOSING tags (e.g. text, code, image)
-                    context.stack.pop()
-                    text = token.content
-                    if text is not None:
-                        element.on_text(context, text)
-
-                    should_render = (
-                        not context.stack
-                        or context.stack
-                        and context.stack.top.on_child_close(context, element)
-                    )
-                    if should_render:
-                        if new_line and node_type != "inline":
-                            yield _new_line_segment
-                        yield from console.render(element, context.options)
-
-                if exiting or self_closing:
-                    element.on_leave(context)
-                    new_line = element.new_line
-
-
-if __name__ == "__main__":  # pragma: no cover
-    import argparse
-    import sys
-
-    parser = argparse.ArgumentParser(
-        description="Render Markdown to the console with Rich"
-    )
-    parser.add_argument(
-        "path",
-        metavar="PATH",
-        help="path to markdown file, or - for stdin",
-    )
-    parser.add_argument(
-        "-c",
-        "--force-color",
-        dest="force_color",
-        action="store_true",
-        default=None,
-        help="force color for non-terminals",
-    )
-    parser.add_argument(
-        "-t",
-        "--code-theme",
-        dest="code_theme",
-        default="monokai",
-        help="pygments code theme",
-    )
-    parser.add_argument(
-        "-i",
-        "--inline-code-lexer",
-        dest="inline_code_lexer",
-        default=None,
-        help="inline_code_lexer",
-    )
-    parser.add_argument(
-        "-y",
-        "--hyperlinks",
-        dest="hyperlinks",
-        action="store_true",
-        help="enable hyperlinks",
-    )
-    parser.add_argument(
-        "-w",
-        "--width",
-        type=int,
-        dest="width",
-        default=None,
-        help="width of output (default will auto-detect)",
-    )
-    parser.add_argument(
-        "-j",
-        "--justify",
-        dest="justify",
-        action="store_true",
-        help="enable full text justify",
-    )
-    parser.add_argument(
-        "-p",
-        "--page",
-        dest="page",
-        action="store_true",
-        help="use pager to scroll output",
-    )
-    args = parser.parse_args()
-
-    from rich.console import Console
-
-    if args.path == "-":
-        markdown_body = sys.stdin.read()
-    else:
-        with open(args.path, encoding="utf-8") as markdown_file:
-            markdown_body = markdown_file.read()
-
-    markdown = Markdown(
-        markdown_body,
-        justify="full" if args.justify else "left",
-        code_theme=args.code_theme,
-        hyperlinks=args.hyperlinks,
-        inline_code_lexer=args.inline_code_lexer,
-    )
-    if args.page:
-        import io
-        import pydoc
-
-        fileio = io.StringIO()
-        console = Console(
-            file=fileio, force_terminal=args.force_color, width=args.width
-        )
-        console.print(markdown)
-        pydoc.pager(fileio.getvalue())
-
-    else:
-        console = Console(
-            force_terminal=args.force_color, width=args.width, record=True
-        )
-        console.print(markdown)
diff --git a/venv/lib/python3.10/site-packages/rich/markup.py b/venv/lib/python3.10/site-packages/rich/markup.py
deleted file mode 100644
index bd9c05a712a8077098507d6870ddc0b0bf3028b0..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/markup.py
+++ /dev/null
@@ -1,251 +0,0 @@
-import re
-from ast import literal_eval
-from operator import attrgetter
-from typing import Callable, Iterable, List, Match, NamedTuple, Optional, Tuple, Union
-
-from ._emoji_replace import _emoji_replace
-from .emoji import EmojiVariant
-from .errors import MarkupError
-from .style import Style
-from .text import Span, Text
-
-RE_TAGS = re.compile(
-    r"""((\\*)\[([a-z#/@][^[]*?)])""",
-    re.VERBOSE,
-)
-
-RE_HANDLER = re.compile(r"^([\w.]*?)(\(.*?\))?$")
-
-
-class Tag(NamedTuple):
-    """A tag in console markup."""
-
-    name: str
-    """The tag name. e.g. 'bold'."""
-    parameters: Optional[str]
-    """Any additional parameters after the name."""
-
-    def __str__(self) -> str:
-        return (
-            self.name if self.parameters is None else f"{self.name} {self.parameters}"
-        )
-
-    @property
-    def markup(self) -> str:
-        """Get the string representation of this tag."""
-        return (
-            f"[{self.name}]"
-            if self.parameters is None
-            else f"[{self.name}={self.parameters}]"
-        )
-
-
-_ReStringMatch = Match[str]  # regex match object
-_ReSubCallable = Callable[[_ReStringMatch], str]  # Callable invoked by re.sub
-_EscapeSubMethod = Callable[[_ReSubCallable, str], str]  # Sub method of a compiled re
-
-
-def escape(
-    markup: str,
-    _escape: _EscapeSubMethod = re.compile(r"(\\*)(\[[a-z#/@][^[]*?])").sub,
-) -> str:
-    """Escapes text so that it won't be interpreted as markup.
-
-    Args:
-        markup (str): Content to be inserted in to markup.
-
-    Returns:
-        str: Markup with square brackets escaped.
-    """
-
-    def escape_backslashes(match: Match[str]) -> str:
-        """Called by re.sub replace matches."""
-        backslashes, text = match.groups()
-        return f"{backslashes}{backslashes}\\{text}"
-
-    markup = _escape(escape_backslashes, markup)
-    if markup.endswith("\\") and not markup.endswith("\\\\"):
-        return markup + "\\"
-
-    return markup
-
-
-def _parse(markup: str) -> Iterable[Tuple[int, Optional[str], Optional[Tag]]]:
-    """Parse markup in to an iterable of tuples of (position, text, tag).
-
-    Args:
-        markup (str): A string containing console markup
-
-    """
-    position = 0
-    _divmod = divmod
-    _Tag = Tag
-    for match in RE_TAGS.finditer(markup):
-        full_text, escapes, tag_text = match.groups()
-        start, end = match.span()
-        if start > position:
-            yield start, markup[position:start], None
-        if escapes:
-            backslashes, escaped = _divmod(len(escapes), 2)
-            if backslashes:
-                # Literal backslashes
-                yield start, "\\" * backslashes, None
-                start += backslashes * 2
-            if escaped:
-                # Escape of tag
-                yield start, full_text[len(escapes) :], None
-                position = end
-                continue
-        text, equals, parameters = tag_text.partition("=")
-        yield start, None, _Tag(text, parameters if equals else None)
-        position = end
-    if position < len(markup):
-        yield position, markup[position:], None
-
-
-def render(
-    markup: str,
-    style: Union[str, Style] = "",
-    emoji: bool = True,
-    emoji_variant: Optional[EmojiVariant] = None,
-) -> Text:
-    """Render console markup in to a Text instance.
-
-    Args:
-        markup (str): A string containing console markup.
-        style: (Union[str, Style]): The style to use.
-        emoji (bool, optional): Also render emoji code. Defaults to True.
-        emoji_variant (str, optional): Optional emoji variant, either "text" or "emoji". Defaults to None.
-
-
-    Raises:
-        MarkupError: If there is a syntax error in the markup.
-
-    Returns:
-        Text: A test instance.
-    """
-    emoji_replace = _emoji_replace
-    if "[" not in markup:
-        return Text(
-            emoji_replace(markup, default_variant=emoji_variant) if emoji else markup,
-            style=style,
-        )
-    text = Text(style=style)
-    append = text.append
-    normalize = Style.normalize
-
-    style_stack: List[Tuple[int, Tag]] = []
-    pop = style_stack.pop
-
-    spans: List[Span] = []
-    append_span = spans.append
-
-    _Span = Span
-    _Tag = Tag
-
-    def pop_style(style_name: str) -> Tuple[int, Tag]:
-        """Pop tag matching given style name."""
-        for index, (_, tag) in enumerate(reversed(style_stack), 1):
-            if tag.name == style_name:
-                return pop(-index)
-        raise KeyError(style_name)
-
-    for position, plain_text, tag in _parse(markup):
-        if plain_text is not None:
-            # Handle open brace escapes, where the brace is not part of a tag.
-            plain_text = plain_text.replace("\\[", "[")
-            append(emoji_replace(plain_text) if emoji else plain_text)
-        elif tag is not None:
-            if tag.name.startswith("/"):  # Closing tag
-                style_name = tag.name[1:].strip()
-
-                if style_name:  # explicit close
-                    style_name = normalize(style_name)
-                    try:
-                        start, open_tag = pop_style(style_name)
-                    except KeyError:
-                        raise MarkupError(
-                            f"closing tag '{tag.markup}' at position {position} doesn't match any open tag"
-                        ) from None
-                else:  # implicit close
-                    try:
-                        start, open_tag = pop()
-                    except IndexError:
-                        raise MarkupError(
-                            f"closing tag '[/]' at position {position} has nothing to close"
-                        ) from None
-
-                if open_tag.name.startswith("@"):
-                    if open_tag.parameters:
-                        handler_name = ""
-                        parameters = open_tag.parameters.strip()
-                        handler_match = RE_HANDLER.match(parameters)
-                        if handler_match is not None:
-                            handler_name, match_parameters = handler_match.groups()
-                            parameters = (
-                                "()" if match_parameters is None else match_parameters
-                            )
-
-                        try:
-                            meta_params = literal_eval(parameters)
-                        except SyntaxError as error:
-                            raise MarkupError(
-                                f"error parsing {parameters!r} in {open_tag.parameters!r}; {error.msg}"
-                            )
-                        except Exception as error:
-                            raise MarkupError(
-                                f"error parsing {open_tag.parameters!r}; {error}"
-                            ) from None
-
-                        if handler_name:
-                            meta_params = (
-                                handler_name,
-                                meta_params
-                                if isinstance(meta_params, tuple)
-                                else (meta_params,),
-                            )
-
-                    else:
-                        meta_params = ()
-
-                    append_span(
-                        _Span(
-                            start, len(text), Style(meta={open_tag.name: meta_params})
-                        )
-                    )
-                else:
-                    append_span(_Span(start, len(text), str(open_tag)))
-
-            else:  # Opening tag
-                normalized_tag = _Tag(normalize(tag.name), tag.parameters)
-                style_stack.append((len(text), normalized_tag))
-
-    text_length = len(text)
-    while style_stack:
-        start, tag = style_stack.pop()
-        style = str(tag)
-        if style:
-            append_span(_Span(start, text_length, style))
-
-    text.spans = sorted(spans[::-1], key=attrgetter("start"))
-    return text
-
-
-if __name__ == "__main__":  # pragma: no cover
-    MARKUP = [
-        "[red]Hello World[/red]",
-        "[magenta]Hello [b]World[/b]",
-        "[bold]Bold[italic] bold and italic [/bold]italic[/italic]",
-        "Click [link=https://www.willmcgugan.com]here[/link] to visit my Blog",
-        ":warning-emoji: [bold red blink] DANGER![/]",
-    ]
-
-    from rich import print
-    from rich.table import Table
-
-    grid = Table("Markup", "Result", padding=(0, 1))
-
-    for markup in MARKUP:
-        grid.add_row(Text(markup), markup)
-
-    print(grid)
diff --git a/venv/lib/python3.10/site-packages/rich/measure.py b/venv/lib/python3.10/site-packages/rich/measure.py
deleted file mode 100644
index a508ffa80bd715b47c190ed9d747dbc388fa5b19..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/measure.py
+++ /dev/null
@@ -1,151 +0,0 @@
-from operator import itemgetter
-from typing import TYPE_CHECKING, Callable, NamedTuple, Optional, Sequence
-
-from . import errors
-from .protocol import is_renderable, rich_cast
-
-if TYPE_CHECKING:
-    from .console import Console, ConsoleOptions, RenderableType
-
-
-class Measurement(NamedTuple):
-    """Stores the minimum and maximum widths (in characters) required to render an object."""
-
-    minimum: int
-    """Minimum number of cells required to render."""
-    maximum: int
-    """Maximum number of cells required to render."""
-
-    @property
-    def span(self) -> int:
-        """Get difference between maximum and minimum."""
-        return self.maximum - self.minimum
-
-    def normalize(self) -> "Measurement":
-        """Get measurement that ensures that minimum <= maximum and minimum >= 0
-
-        Returns:
-            Measurement: A normalized measurement.
-        """
-        minimum, maximum = self
-        minimum = min(max(0, minimum), maximum)
-        return Measurement(max(0, minimum), max(0, max(minimum, maximum)))
-
-    def with_maximum(self, width: int) -> "Measurement":
-        """Get a RenderableWith where the widths are <= width.
-
-        Args:
-            width (int): Maximum desired width.
-
-        Returns:
-            Measurement: New Measurement object.
-        """
-        minimum, maximum = self
-        return Measurement(min(minimum, width), min(maximum, width))
-
-    def with_minimum(self, width: int) -> "Measurement":
-        """Get a RenderableWith where the widths are >= width.
-
-        Args:
-            width (int): Minimum desired width.
-
-        Returns:
-            Measurement: New Measurement object.
-        """
-        minimum, maximum = self
-        width = max(0, width)
-        return Measurement(max(minimum, width), max(maximum, width))
-
-    def clamp(
-        self, min_width: Optional[int] = None, max_width: Optional[int] = None
-    ) -> "Measurement":
-        """Clamp a measurement within the specified range.
-
-        Args:
-            min_width (int): Minimum desired width, or ``None`` for no minimum. Defaults to None.
-            max_width (int): Maximum desired width, or ``None`` for no maximum. Defaults to None.
-
-        Returns:
-            Measurement: New Measurement object.
-        """
-        measurement = self
-        if min_width is not None:
-            measurement = measurement.with_minimum(min_width)
-        if max_width is not None:
-            measurement = measurement.with_maximum(max_width)
-        return measurement
-
-    @classmethod
-    def get(
-        cls, console: "Console", options: "ConsoleOptions", renderable: "RenderableType"
-    ) -> "Measurement":
-        """Get a measurement for a renderable.
-
-        Args:
-            console (~rich.console.Console): Console instance.
-            options (~rich.console.ConsoleOptions): Console options.
-            renderable (RenderableType): An object that may be rendered with Rich.
-
-        Raises:
-            errors.NotRenderableError: If the object is not renderable.
-
-        Returns:
-            Measurement: Measurement object containing range of character widths required to render the object.
-        """
-        _max_width = options.max_width
-        if _max_width < 1:
-            return Measurement(0, 0)
-        if isinstance(renderable, str):
-            renderable = console.render_str(
-                renderable, markup=options.markup, highlight=False
-            )
-        renderable = rich_cast(renderable)
-        if is_renderable(renderable):
-            get_console_width: Optional[
-                Callable[["Console", "ConsoleOptions"], "Measurement"]
-            ] = getattr(renderable, "__rich_measure__", None)
-            if get_console_width is not None:
-                render_width = (
-                    get_console_width(console, options)
-                    .normalize()
-                    .with_maximum(_max_width)
-                )
-                if render_width.maximum < 1:
-                    return Measurement(0, 0)
-                return render_width.normalize()
-            else:
-                return Measurement(0, _max_width)
-        else:
-            raise errors.NotRenderableError(
-                f"Unable to get render width for {renderable!r}; "
-                "a str, Segment, or object with __rich_console__ method is required"
-            )
-
-
-def measure_renderables(
-    console: "Console",
-    options: "ConsoleOptions",
-    renderables: Sequence["RenderableType"],
-) -> "Measurement":
-    """Get a measurement that would fit a number of renderables.
-
-    Args:
-        console (~rich.console.Console): Console instance.
-        options (~rich.console.ConsoleOptions): Console options.
-        renderables (Iterable[RenderableType]): One or more renderable objects.
-
-    Returns:
-        Measurement: Measurement object containing range of character widths required to
-            contain all given renderables.
-    """
-    if not renderables:
-        return Measurement(0, 0)
-    get_measurement = Measurement.get
-    measurements = [
-        get_measurement(console, options, renderable) for renderable in renderables
-    ]
-    measured_width = Measurement(
-        max(measurements, key=itemgetter(0)).minimum,
-        max(measurements, key=itemgetter(1)).maximum,
-    )
-    return measured_width
diff --git a/venv/lib/python3.10/site-packages/rich/padding.py b/venv/lib/python3.10/site-packages/rich/padding.py
deleted file mode 100644
index d1aa01bee541c0b7899f137fe1c355ac19109525..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/padding.py
+++ /dev/null
@@ -1,141 +0,0 @@
-from typing import TYPE_CHECKING, List, Optional, Tuple, Union
-
-if TYPE_CHECKING:
-    from .console import (
-        Console,
-        ConsoleOptions,
-        RenderableType,
-        RenderResult,
-    )
-
-from .jupyter import JupyterMixin
-from .measure import Measurement
-from .segment import Segment
-from .style import Style
-
-PaddingDimensions = Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int, int]]
-
-
-class Padding(JupyterMixin):
-    """Draw space around content.
-
-    Example:
-        >>> print(Padding("Hello", (2, 4), style="on blue"))
-
-    Args:
-        renderable (RenderableType): String or other renderable.
-        pad (Union[int, Tuple[int]]): Padding for top, right, bottom, and left borders.
-            May be specified with 1, 2, or 4 integers (CSS style).
-        style (Union[str, Style], optional): Style for padding characters. Defaults to "none".
-        expand (bool, optional): Expand padding to fit available width. Defaults to True.
-    """
-
-    def __init__(
-        self,
-        renderable: "RenderableType",
-        pad: "PaddingDimensions" = (0, 0, 0, 0),
-        *,
-        style: Union[str, Style] = "none",
-        expand: bool = True,
-    ):
-        self.renderable = renderable
-        self.top, self.right, self.bottom, self.left = self.unpack(pad)
-        self.style = style
-        self.expand = expand
-
-    @classmethod
-    def indent(cls, renderable: "RenderableType", level: int) -> "Padding":
-        """Make padding instance to render an indent.
-
-        Args:
-            renderable (RenderableType): String or other renderable.
-            level (int): Number of characters to indent.
-
-        Returns:
-            Padding: A Padding instance.
-        """
-
-        return Padding(renderable, pad=(0, 0, 0, level), expand=False)
-
-    @staticmethod
-    def unpack(pad: "PaddingDimensions") -> Tuple[int, int, int, int]:
-        """Unpack padding specified in CSS style."""
-        if isinstance(pad, int):
-            return (pad, pad, pad, pad)
-        if len(pad) == 1:
-            _pad = pad[0]
-            return (_pad, _pad, _pad, _pad)
-        if len(pad) == 2:
-            pad_top, pad_right = pad
-            return (pad_top, pad_right, pad_top, pad_right)
-        if len(pad) == 4:
-            top, right, bottom, left = pad
-            return (top, right, bottom, left)
-        raise ValueError(f"1, 2 or 4 integers required for padding; {len(pad)} given")
-
-    def __repr__(self) -> str:
-        return f"Padding({self.renderable!r}, ({self.top},{self.right},{self.bottom},{self.left}))"
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        style = console.get_style(self.style)
-        if self.expand:
-            width = options.max_width
-        else:
-            width = min(
-                Measurement.get(console, options, self.renderable).maximum
-                + self.left
-                + self.right,
-                options.max_width,
-            )
-        render_options = options.update_width(width - self.left - self.right)
-        if render_options.height is not None:
-            render_options = render_options.update_height(
-                height=render_options.height - self.top - self.bottom
-            )
-        lines = console.render_lines(
-            self.renderable, render_options, style=style, pad=True
-        )
-        _Segment = Segment
-
-        left = _Segment(" " * self.left, style) if self.left else None
-        right = (
-            [_Segment(f'{" " * self.right}', style), _Segment.line()]
-            if self.right
-            else [_Segment.line()]
-        )
-        blank_line: Optional[List[Segment]] = None
-        if self.top:
-            blank_line = [_Segment(f'{" " * width}\n', style)]
-            yield from blank_line * self.top
-        if left:
-            for line in lines:
-                yield left
-                yield from line
-                yield from right
-        else:
-            for line in lines:
-                yield from line
-                yield from right
-        if self.bottom:
-            blank_line = blank_line or [_Segment(f'{" " * width}\n', style)]
-            yield from blank_line * self.bottom
-
-    def __rich_measure__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "Measurement":
-        max_width = options.max_width
-        extra_width = self.left + self.right
-        if max_width - extra_width < 1:
-            return Measurement(max_width, max_width)
-        measure_min, measure_max = Measurement.get(console, options, self.renderable)
-        measurement = Measurement(measure_min + extra_width, measure_max + extra_width)
-        measurement = measurement.with_maximum(max_width)
-        return measurement
-
-
-if __name__ == "__main__":  #  pragma: no cover
-    from rich import print
-
-    print(Padding("Hello, World", (2, 4), style="on blue"))
diff --git a/venv/lib/python3.10/site-packages/rich/pager.py b/venv/lib/python3.10/site-packages/rich/pager.py
deleted file mode 100644
index a3f7aa62af1ee2690e1e17ee41f3c368953625b8..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/pager.py
+++ /dev/null
@@ -1,34 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import Any
-
-
-class Pager(ABC):
-    """Base class for a pager."""
-
-    @abstractmethod
-    def show(self, content: str) -> None:
-        """Show content in pager.
-
-        Args:
-            content (str): Content to be displayed.
-        """
-
-
-class SystemPager(Pager):
-    """Uses the pager installed on the system."""
-
-    def _pager(self, content: str) -> Any:  #  pragma: no cover
-        return __import__("pydoc").pager(content)
-
-    def show(self, content: str) -> None:
-        """Use the same pager used by pydoc."""
-        self._pager(content)
-
-
-if __name__ == "__main__":  # pragma: no cover
-    from .__main__ import make_test_card
-    from .console import Console
-
-    console = Console()
-    with console.pager(styles=True):
-        console.print(make_test_card())
diff --git a/venv/lib/python3.10/site-packages/rich/palette.py b/venv/lib/python3.10/site-packages/rich/palette.py
deleted file mode 100644
index f29587945cb2214ab6d92856dac6d11c320a8ca1..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/palette.py
+++ /dev/null
@@ -1,100 +0,0 @@
-from math import sqrt
-from functools import lru_cache
-from typing import Sequence, Tuple, TYPE_CHECKING
-
-from .color_triplet import ColorTriplet
-
-if TYPE_CHECKING:
-    from rich.table import Table
-
-
-class Palette:
-    """A palette of available colors."""
-
-    def __init__(self, colors: Sequence[Tuple[int, int, int]]):
-        self._colors = colors
-
-    def __getitem__(self, number: int) -> ColorTriplet:
-        return ColorTriplet(*self._colors[number])
-
-    def __rich__(self) -> "Table":
-        from rich.color import Color
-        from rich.style import Style
-        from rich.text import Text
-        from rich.table import Table
-
-        table = Table(
-            "index",
-            "RGB",
-            "Color",
-            title="Palette",
-            caption=f"{len(self._colors)} colors",
-            highlight=True,
-            caption_justify="right",
-        )
-        for index, color in enumerate(self._colors):
-            table.add_row(
-                str(index),
-                repr(color),
-                Text(" " * 16, style=Style(bgcolor=Color.from_rgb(*color))),
-            )
-        return table
-
-    # This is somewhat inefficient and needs caching
-    @lru_cache(maxsize=1024)
-    def match(self, color: Tuple[int, int, int]) -> int:
-        """Find a color from a palette that most closely matches a given color.
-
-        Args:
-            color (Tuple[int, int, int]): RGB components in range 0 > 255.
-
-        Returns:
-            int: Index of closes matching color.
-        """
-        red1, green1, blue1 = color
-        _sqrt = sqrt
-        get_color = self._colors.__getitem__
-
-        def get_color_distance(index: int) -> float:
-            """Get the distance to a color."""
-            red2, green2, blue2 = get_color(index)
-            red_mean = (red1 + red2) // 2
-            red = red1 - red2
-            green = green1 - green2
-            blue = blue1 - blue2
-            return _sqrt(
-                (((512 + red_mean) * red * red) >> 8)
-                + 4 * green * green
-                + (((767 - red_mean) * blue * blue) >> 8)
-            )
-
-        min_index = min(range(len(self._colors)), key=get_color_distance)
-        return min_index
-
-
-if __name__ == "__main__":  # pragma: no cover
-    import colorsys
-    from typing import Iterable
-    from rich.color import Color
-    from rich.console import Console, ConsoleOptions
-    from rich.segment import Segment
-    from rich.style import Style
-
-    class ColorBox:
-        def __rich_console__(
-            self, console: Console, options: ConsoleOptions
-        ) -> Iterable[Segment]:
-            height = console.size.height - 3
-            for y in range(0, height):
-                for x in range(options.max_width):
-                    h = x / options.max_width
-                    l = y / (height + 1)
-                    r1, g1, b1 = colorsys.hls_to_rgb(h, l, 1.0)
-                    r2, g2, b2 = colorsys.hls_to_rgb(h, l + (1 / height / 2), 1.0)
-                    bgcolor = Color.from_rgb(r1 * 255, g1 * 255, b1 * 255)
-                    color = Color.from_rgb(r2 * 255, g2 * 255, b2 * 255)
-                    yield Segment("▄", Style(color=color, bgcolor=bgcolor))
-                yield Segment.line()
-
-    console = Console()
-    console.print(ColorBox())
diff --git a/venv/lib/python3.10/site-packages/rich/panel.py b/venv/lib/python3.10/site-packages/rich/panel.py
deleted file mode 100644
index d411e291533e45ab6d669c33a7875252dfbec3b6..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/panel.py
+++ /dev/null
@@ -1,318 +0,0 @@
-from typing import TYPE_CHECKING, Optional
-
-from .align import AlignMethod
-from .box import ROUNDED, Box
-from .cells import cell_len
-from .jupyter import JupyterMixin
-from .measure import Measurement, measure_renderables
-from .padding import Padding, PaddingDimensions
-from .segment import Segment
-from .style import Style, StyleType
-from .text import Text, TextType
-
-if TYPE_CHECKING:
-    from .console import Console, ConsoleOptions, RenderableType, RenderResult
-
-
-class Panel(JupyterMixin):
-    """A console renderable that draws a border around its contents.
-
-    Example:
-        >>> console.print(Panel("Hello, World!"))
-
-    Args:
-        renderable (RenderableType): A console renderable object.
-        box (Box): A Box instance that defines the look of the border (see :ref:`appendix_box`. Defaults to box.ROUNDED.
-        title (Optional[TextType], optional): Optional title displayed in panel header. Defaults to None.
-        title_align (AlignMethod, optional): Alignment of title. Defaults to "center".
-        subtitle (Optional[TextType], optional): Optional subtitle displayed in panel footer. Defaults to None.
-        subtitle_align (AlignMethod, optional): Alignment of subtitle. Defaults to "center".
-        safe_box (bool, optional): Disable box characters that don't display on windows legacy terminal with *raster* fonts. Defaults to True.
-        expand (bool, optional): If True the panel will stretch to fill the console width, otherwise it will be sized to fit the contents. Defaults to True.
-        style (str, optional): The style of the panel (border and contents). Defaults to "none".
-        border_style (str, optional): The style of the border. Defaults to "none".
-        width (Optional[int], optional): Optional width of panel. Defaults to None to auto-detect.
-        height (Optional[int], optional): Optional height of panel. Defaults to None to auto-detect.
-        padding (Optional[PaddingDimensions]): Optional padding around renderable. Defaults to 0.
-        highlight (bool, optional): Enable automatic highlighting of panel title (if str). Defaults to False.
-    """
-
-    def __init__(
-        self,
-        renderable: "RenderableType",
-        box: Box = ROUNDED,
-        *,
-        title: Optional[TextType] = None,
-        title_align: AlignMethod = "center",
-        subtitle: Optional[TextType] = None,
-        subtitle_align: AlignMethod = "center",
-        safe_box: Optional[bool] = None,
-        expand: bool = True,
-        style: StyleType = "none",
-        border_style: StyleType = "none",
-        width: Optional[int] = None,
-        height: Optional[int] = None,
-        padding: PaddingDimensions = (0, 1),
-        highlight: bool = False,
-    ) -> None:
-        self.renderable = renderable
-        self.box = box
-        self.title = title
-        self.title_align: AlignMethod = title_align
-        self.subtitle = subtitle
-        self.subtitle_align = subtitle_align
-        self.safe_box = safe_box
-        self.expand = expand
-        self.style = style
-        self.border_style = border_style
-        self.width = width
-        self.height = height
-        self.padding = padding
-        self.highlight = highlight
-
-    @classmethod
-    def fit(
-        cls,
-        renderable: "RenderableType",
-        box: Box = ROUNDED,
-        *,
-        title: Optional[TextType] = None,
-        title_align: AlignMethod = "center",
-        subtitle: Optional[TextType] = None,
-        subtitle_align: AlignMethod = "center",
-        safe_box: Optional[bool] = None,
-        style: StyleType = "none",
-        border_style: StyleType = "none",
-        width: Optional[int] = None,
-        height: Optional[int] = None,
-        padding: PaddingDimensions = (0, 1),
-        highlight: bool = False,
-    ) -> "Panel":
-        """An alternative constructor that sets expand=False."""
-        return cls(
-            renderable,
-            box,
-            title=title,
-            title_align=title_align,
-            subtitle=subtitle,
-            subtitle_align=subtitle_align,
-            safe_box=safe_box,
-            style=style,
-            border_style=border_style,
-            width=width,
-            height=height,
-            padding=padding,
-            highlight=highlight,
-            expand=False,
-        )
-
-    @property
-    def _title(self) -> Optional[Text]:
-        if self.title:
-            title_text = (
-                Text.from_markup(self.title)
-                if isinstance(self.title, str)
-                else self.title.copy()
-            )
-            title_text.end = ""
-            title_text.plain = title_text.plain.replace("\n", " ")
-            title_text.no_wrap = True
-            title_text.expand_tabs()
-            title_text.pad(1)
-            return title_text
-        return None
-
-    @property
-    def _subtitle(self) -> Optional[Text]:
-        if self.subtitle:
-            subtitle_text = (
-                Text.from_markup(self.subtitle)
-                if isinstance(self.subtitle, str)
-                else self.subtitle.copy()
-            )
-            subtitle_text.end = ""
-            subtitle_text.plain = subtitle_text.plain.replace("\n", " ")
-            subtitle_text.no_wrap = True
-            subtitle_text.expand_tabs()
-            subtitle_text.pad(1)
-            return subtitle_text
-        return None
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        _padding = Padding.unpack(self.padding)
-        renderable = (
-            Padding(self.renderable, _padding) if any(_padding) else self.renderable
-        )
-        style = console.get_style(self.style)
-        partial_border_style = console.get_style(self.border_style)
-        border_style = style + partial_border_style
-        width = (
-            options.max_width
-            if self.width is None
-            else min(options.max_width, self.width)
-        )
-
-        safe_box: bool = console.safe_box if self.safe_box is None else self.safe_box
-        box = self.box.substitute(options, safe=safe_box)
-
-        def align_text(
-            text: Text, width: int, align: str, character: str, style: Style
-        ) -> Text:
-            """Gets new aligned text.
-
-            Args:
-                text (Text): Title or subtitle text.
-                width (int): Desired width.
-                align (str): Alignment.
-                character (str): Character for alignment.
-                style (Style): Border style
-
-            Returns:
-                Text: New text instance
-            """
-            text = text.copy()
-            text.truncate(width)
-            excess_space = width - cell_len(text.plain)
-            if text.style:
-                text.stylize(console.get_style(text.style))
-
-            if excess_space:
-                if align == "left":
-                    return Text.assemble(
-                        text,
-                        (character * excess_space, style),
-                        no_wrap=True,
-                        end="",
-                    )
-                elif align == "center":
-                    left = excess_space // 2
-                    return Text.assemble(
-                        (character * left, style),
-                        text,
-                        (character * (excess_space - left), style),
-                        no_wrap=True,
-                        end="",
-                    )
-                else:
-                    return Text.assemble(
-                        (character * excess_space, style),
-                        text,
-                        no_wrap=True,
-                        end="",
-                    )
-            return text
-
-        title_text = self._title
-        if title_text is not None:
-            title_text.stylize_before(partial_border_style)
-
-        child_width = (
-            width - 2
-            if self.expand
-            else console.measure(
-                renderable, options=options.update_width(width - 2)
-            ).maximum
-        )
-        child_height = self.height or options.height or None
-        if child_height:
-            child_height -= 2
-        if title_text is not None:
-            child_width = min(
-                options.max_width - 2, max(child_width, title_text.cell_len + 2)
-            )
-
-        width = child_width + 2
-        child_options = options.update(
-            width=child_width, height=child_height, highlight=self.highlight
-        )
-        lines = console.render_lines(renderable, child_options, style=style)
-
-        line_start = Segment(box.mid_left, border_style)
-        line_end = Segment(f"{box.mid_right}", border_style)
-        new_line = Segment.line()
-        if title_text is None or width <= 4:
-            yield Segment(box.get_top([width - 2]), border_style)
-        else:
-            title_text = align_text(
-                title_text,
-                width - 4,
-                self.title_align,
-                box.top,
-                border_style,
-            )
-            yield Segment(box.top_left + box.top, border_style)
-            yield from console.render(title_text, child_options.update_width(width - 4))
-            yield Segment(box.top + box.top_right, border_style)
-
-        yield new_line
-        for line in lines:
-            yield line_start
-            yield from line
-            yield line_end
-            yield new_line
-
-        subtitle_text = self._subtitle
-        if subtitle_text is not None:
-            subtitle_text.stylize_before(partial_border_style)
-
-        if subtitle_text is None or width <= 4:
-            yield Segment(box.get_bottom([width - 2]), border_style)
-        else:
-            subtitle_text = align_text(
-                subtitle_text,
-                width - 4,
-                self.subtitle_align,
-                box.bottom,
-                border_style,
-            )
-            yield Segment(box.bottom_left + box.bottom, border_style)
-            yield from console.render(
-                subtitle_text, child_options.update_width(width - 4)
-            )
-            yield Segment(box.bottom + box.bottom_right, border_style)
-
-        yield new_line
-
-    def __rich_measure__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "Measurement":
-        _title = self._title
-        _, right, _, left = Padding.unpack(self.padding)
-        padding = left + right
-        renderables = [self.renderable, _title] if _title else [self.renderable]
-
-        if self.width is None:
-            width = (
-                measure_renderables(
-                    console,
-                    options.update_width(options.max_width - padding - 2),
-                    renderables,
-                ).maximum
-                + padding
-                + 2
-            )
-        else:
-            width = self.width
-        return Measurement(width, width)
-
-
-if __name__ == "__main__":  # pragma: no cover
-    from .console import Console
-
-    c = Console()
-
-    from .box import DOUBLE, ROUNDED
-    from .padding import Padding
-
-    p = Panel(
-        "Hello, World!",
-        title="rich.Panel",
-        style="white on blue",
-        box=DOUBLE,
-        padding=1,
-    )
-
-    c.print()
-    c.print(p)
diff --git a/venv/lib/python3.10/site-packages/rich/pretty.py b/venv/lib/python3.10/site-packages/rich/pretty.py
deleted file mode 100644
index 5c725c0c53b3d59149f2b69e4b0f11c236743437..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/pretty.py
+++ /dev/null
@@ -1,1016 +0,0 @@
-import builtins
-import collections
-import dataclasses
-import inspect
-import os
-import reprlib
-import sys
-from array import array
-from collections import Counter, UserDict, UserList, defaultdict, deque
-from dataclasses import dataclass, fields, is_dataclass
-from inspect import isclass
-from itertools import islice
-from types import MappingProxyType
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    DefaultDict,
-    Deque,
-    Dict,
-    Iterable,
-    List,
-    Optional,
-    Sequence,
-    Set,
-    Tuple,
-    Union,
-)
-
-from rich.repr import RichReprResult
-
-try:
-    import attr as _attr_module
-
-    _has_attrs = hasattr(_attr_module, "ib")
-except ImportError:  # pragma: no cover
-    _has_attrs = False
-
-from . import get_console
-from ._loop import loop_last
-from ._pick import pick_bool
-from .abc import RichRenderable
-from .cells import cell_len
-from .highlighter import ReprHighlighter
-from .jupyter import JupyterMixin, JupyterRenderable
-from .measure import Measurement
-from .text import Text
-
-if TYPE_CHECKING:
-    from .console import (
-        Console,
-        ConsoleOptions,
-        HighlighterType,
-        JustifyMethod,
-        OverflowMethod,
-        RenderResult,
-    )
-
-
-def _is_attr_object(obj: Any) -> bool:
-    """Check if an object was created with attrs module."""
-    return _has_attrs and _attr_module.has(type(obj))
-
-
-def _get_attr_fields(obj: Any) -> Sequence["_attr_module.Attribute[Any]"]:
-    """Get fields for an attrs object."""
-    return _attr_module.fields(type(obj)) if _has_attrs else []
-
-
-def _is_dataclass_repr(obj: object) -> bool:
-    """Check if an instance of a dataclass contains the default repr.
-
-    Args:
-        obj (object): A dataclass instance.
-
-    Returns:
-        bool: True if the default repr is used, False if there is a custom repr.
-    """
-    # Digging in to a lot of internals here
-    # Catching all exceptions in case something is missing on a non CPython implementation
-    try:
-        return obj.__repr__.__code__.co_filename in (
-            dataclasses.__file__,
-            reprlib.__file__,
-        )
-    except Exception:  # pragma: no coverage
-        return False
-
-
-_dummy_namedtuple = collections.namedtuple("_dummy_namedtuple", [])
-
-
-def _has_default_namedtuple_repr(obj: object) -> bool:
-    """Check if an instance of namedtuple contains the default repr
-
-    Args:
-        obj (object): A namedtuple
-
-    Returns:
-        bool: True if the default repr is used, False if there's a custom repr.
-    """
-    obj_file = None
-    try:
-        obj_file = inspect.getfile(obj.__repr__)
-    except (OSError, TypeError):
-        # OSError handles case where object is defined in __main__ scope, e.g. REPL - no filename available.
-        # TypeError trapped defensively, in case of object without filename slips through.
-        pass
-    default_repr_file = inspect.getfile(_dummy_namedtuple.__repr__)
-    return obj_file == default_repr_file
-
-
-def _ipy_display_hook(
-    value: Any,
-    console: Optional["Console"] = None,
-    overflow: "OverflowMethod" = "ignore",
-    crop: bool = False,
-    indent_guides: bool = False,
-    max_length: Optional[int] = None,
-    max_string: Optional[int] = None,
-    max_depth: Optional[int] = None,
-    expand_all: bool = False,
-) -> Union[str, None]:
-    # needed here to prevent circular import:
-    from .console import ConsoleRenderable
-
-    # always skip rich generated jupyter renderables or None values
-    if _safe_isinstance(value, JupyterRenderable) or value is None:
-        return None
-
-    console = console or get_console()
-
-    with console.capture() as capture:
-        # certain renderables should start on a new line
-        if _safe_isinstance(value, ConsoleRenderable):
-            console.line()
-        console.print(
-            (
-                value
-                if _safe_isinstance(value, RichRenderable)
-                else Pretty(
-                    value,
-                    overflow=overflow,
-                    indent_guides=indent_guides,
-                    max_length=max_length,
-                    max_string=max_string,
-                    max_depth=max_depth,
-                    expand_all=expand_all,
-                    margin=12,
-                )
-            ),
-            crop=crop,
-            new_line_start=True,
-            end="",
-        )
-    # strip trailing newline, not usually part of a text repr
-    # I'm not sure if this should be prevented at a lower level
-    return capture.get().rstrip("\n")
-
-
-def _safe_isinstance(
-    obj: object, class_or_tuple: Union[type, Tuple[type, ...]]
-) -> bool:
-    """isinstance can fail in rare cases, for example types with no __class__"""
-    try:
-        return isinstance(obj, class_or_tuple)
-    except Exception:
-        return False
-
-
-def install(
-    console: Optional["Console"] = None,
-    overflow: "OverflowMethod" = "ignore",
-    crop: bool = False,
-    indent_guides: bool = False,
-    max_length: Optional[int] = None,
-    max_string: Optional[int] = None,
-    max_depth: Optional[int] = None,
-    expand_all: bool = False,
-) -> None:
-    """Install automatic pretty printing in the Python REPL.
-
-    Args:
-        console (Console, optional): Console instance or ``None`` to use global console. Defaults to None.
-        overflow (Optional[OverflowMethod], optional): Overflow method. Defaults to "ignore".
-        crop (Optional[bool], optional): Enable cropping of long lines. Defaults to False.
-        indent_guides (bool, optional): Enable indentation guides. Defaults to False.
-        max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
-            Defaults to None.
-        max_string (int, optional): Maximum length of string before truncating, or None to disable. Defaults to None.
-        max_depth (int, optional): Maximum depth of nested data structures, or None for no maximum. Defaults to None.
-        expand_all (bool, optional): Expand all containers. Defaults to False.
-        max_frames (int): Maximum number of frames to show in a traceback, 0 for no maximum. Defaults to 100.
-    """
-    from rich import get_console
-
-    console = console or get_console()
-    assert console is not None
-
-    def display_hook(value: Any) -> None:
-        """Replacement sys.displayhook which prettifies objects with Rich."""
-        if value is not None:
-            assert console is not None
-            builtins._ = None  # type: ignore[attr-defined]
-            console.print(
-                (
-                    value
-                    if _safe_isinstance(value, RichRenderable)
-                    else Pretty(
-                        value,
-                        overflow=overflow,
-                        indent_guides=indent_guides,
-                        max_length=max_length,
-                        max_string=max_string,
-                        max_depth=max_depth,
-                        expand_all=expand_all,
-                    )
-                ),
-                crop=crop,
-            )
-            builtins._ = value  # type: ignore[attr-defined]
-
-    try:
-        ip = get_ipython()  # type: ignore[name-defined]
-    except NameError:
-        sys.displayhook = display_hook
-    else:
-        from IPython.core.formatters import BaseFormatter
-
-        class RichFormatter(BaseFormatter):  # type: ignore[misc]
-            pprint: bool = True
-
-            def __call__(self, value: Any) -> Any:
-                if self.pprint:
-                    return _ipy_display_hook(
-                        value,
-                        console=get_console(),
-                        overflow=overflow,
-                        indent_guides=indent_guides,
-                        max_length=max_length,
-                        max_string=max_string,
-                        max_depth=max_depth,
-                        expand_all=expand_all,
-                    )
-                else:
-                    return repr(value)
-
-        # replace plain text formatter with rich formatter
-        rich_formatter = RichFormatter()
-        ip.display_formatter.formatters["text/plain"] = rich_formatter
-
-
-class Pretty(JupyterMixin):
-    """A rich renderable that pretty prints an object.
-
-    Args:
-        _object (Any): An object to pretty print.
-        highlighter (HighlighterType, optional): Highlighter object to apply to result, or None for ReprHighlighter. Defaults to None.
-        indent_size (int, optional): Number of spaces in indent. Defaults to 4.
-        justify (JustifyMethod, optional): Justify method, or None for default. Defaults to None.
-        overflow (OverflowMethod, optional): Overflow method, or None for default. Defaults to None.
-        no_wrap (Optional[bool], optional): Disable word wrapping. Defaults to False.
-        indent_guides (bool, optional): Enable indentation guides. Defaults to False.
-        max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
-            Defaults to None.
-        max_string (int, optional): Maximum length of string before truncating, or None to disable. Defaults to None.
-        max_depth (int, optional): Maximum depth of nested data structures, or None for no maximum. Defaults to None.
-        expand_all (bool, optional): Expand all containers. Defaults to False.
-        margin (int, optional): Subtrace a margin from width to force containers to expand earlier. Defaults to 0.
-        insert_line (bool, optional): Insert a new line if the output has multiple new lines. Defaults to False.
-    """
-
-    def __init__(
-        self,
-        _object: Any,
-        highlighter: Optional["HighlighterType"] = None,
-        *,
-        indent_size: int = 4,
-        justify: Optional["JustifyMethod"] = None,
-        overflow: Optional["OverflowMethod"] = None,
-        no_wrap: Optional[bool] = False,
-        indent_guides: bool = False,
-        max_length: Optional[int] = None,
-        max_string: Optional[int] = None,
-        max_depth: Optional[int] = None,
-        expand_all: bool = False,
-        margin: int = 0,
-        insert_line: bool = False,
-    ) -> None:
-        self._object = _object
-        self.highlighter = highlighter or ReprHighlighter()
-        self.indent_size = indent_size
-        self.justify: Optional["JustifyMethod"] = justify
-        self.overflow: Optional["OverflowMethod"] = overflow
-        self.no_wrap = no_wrap
-        self.indent_guides = indent_guides
-        self.max_length = max_length
-        self.max_string = max_string
-        self.max_depth = max_depth
-        self.expand_all = expand_all
-        self.margin = margin
-        self.insert_line = insert_line
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        pretty_str = pretty_repr(
-            self._object,
-            max_width=options.max_width - self.margin,
-            indent_size=self.indent_size,
-            max_length=self.max_length,
-            max_string=self.max_string,
-            max_depth=self.max_depth,
-            expand_all=self.expand_all,
-        )
-        pretty_text = Text.from_ansi(
-            pretty_str,
-            justify=self.justify or options.justify,
-            overflow=self.overflow or options.overflow,
-            no_wrap=pick_bool(self.no_wrap, options.no_wrap),
-            style="pretty",
-        )
-        pretty_text = (
-            self.highlighter(pretty_text)
-            if pretty_text
-            else Text(
-                f"{type(self._object)}.__repr__ returned empty string",
-                style="dim italic",
-            )
-        )
-        if self.indent_guides and not options.ascii_only:
-            pretty_text = pretty_text.with_indent_guides(
-                self.indent_size, style="repr.indent"
-            )
-        if self.insert_line and "\n" in pretty_text:
-            yield ""
-        yield pretty_text
-
-    def __rich_measure__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "Measurement":
-        pretty_str = pretty_repr(
-            self._object,
-            max_width=options.max_width,
-            indent_size=self.indent_size,
-            max_length=self.max_length,
-            max_string=self.max_string,
-            max_depth=self.max_depth,
-            expand_all=self.expand_all,
-        )
-        text_width = (
-            max(cell_len(line) for line in pretty_str.splitlines()) if pretty_str else 0
-        )
-        return Measurement(text_width, text_width)
-
-
-def _get_braces_for_defaultdict(_object: DefaultDict[Any, Any]) -> Tuple[str, str, str]:
-    return (
-        f"defaultdict({_object.default_factory!r}, {{",
-        "})",
-        f"defaultdict({_object.default_factory!r}, {{}})",
-    )
-
-
-def _get_braces_for_deque(_object: Deque[Any]) -> Tuple[str, str, str]:
-    if _object.maxlen is None:
-        return ("deque([", "])", "deque()")
-    return (
-        "deque([",
-        f"], maxlen={_object.maxlen})",
-        f"deque(maxlen={_object.maxlen})",
-    )
-
-
-def _get_braces_for_array(_object: "array[Any]") -> Tuple[str, str, str]:
-    return (f"array({_object.typecode!r}, [", "])", f"array({_object.typecode!r})")
-
-
-_BRACES: Dict[type, Callable[[Any], Tuple[str, str, str]]] = {
-    os._Environ: lambda _object: ("environ({", "})", "environ({})"),
-    array: _get_braces_for_array,
-    defaultdict: _get_braces_for_defaultdict,
-    Counter: lambda _object: ("Counter({", "})", "Counter()"),
-    deque: _get_braces_for_deque,
-    dict: lambda _object: ("{", "}", "{}"),
-    UserDict: lambda _object: ("{", "}", "{}"),
-    frozenset: lambda _object: ("frozenset({", "})", "frozenset()"),
-    list: lambda _object: ("[", "]", "[]"),
-    UserList: lambda _object: ("[", "]", "[]"),
-    set: lambda _object: ("{", "}", "set()"),
-    tuple: lambda _object: ("(", ")", "()"),
-    MappingProxyType: lambda _object: ("mappingproxy({", "})", "mappingproxy({})"),
-}
-_CONTAINERS = tuple(_BRACES.keys())
-_MAPPING_CONTAINERS = (dict, os._Environ, MappingProxyType, UserDict)
-
-
-def is_expandable(obj: Any) -> bool:
-    """Check if an object may be expanded by pretty print."""
-    return (
-        _safe_isinstance(obj, _CONTAINERS)
-        or (is_dataclass(obj))
-        or (hasattr(obj, "__rich_repr__"))
-        or _is_attr_object(obj)
-    ) and not isclass(obj)
-
-
-@dataclass
-class Node:
-    """A node in a repr tree. May be atomic or a container."""
-
-    key_repr: str = ""
-    value_repr: str = ""
-    open_brace: str = ""
-    close_brace: str = ""
-    empty: str = ""
-    last: bool = False
-    is_tuple: bool = False
-    is_namedtuple: bool = False
-    children: Optional[List["Node"]] = None
-    key_separator: str = ": "
-    separator: str = ", "
-
-    def iter_tokens(self) -> Iterable[str]:
-        """Generate tokens for this node."""
-        if self.key_repr:
-            yield self.key_repr
-            yield self.key_separator
-        if self.value_repr:
-            yield self.value_repr
-        elif self.children is not None:
-            if self.children:
-                yield self.open_brace
-                if self.is_tuple and not self.is_namedtuple and len(self.children) == 1:
-                    yield from self.children[0].iter_tokens()
-                    yield ","
-                else:
-                    for child in self.children:
-                        yield from child.iter_tokens()
-                        if not child.last:
-                            yield self.separator
-                yield self.close_brace
-            else:
-                yield self.empty
-
-    def check_length(self, start_length: int, max_length: int) -> bool:
-        """Check the length fits within a limit.
-
-        Args:
-            start_length (int): Starting length of the line (indent, prefix, suffix).
-            max_length (int): Maximum length.
-
-        Returns:
-            bool: True if the node can be rendered within max length, otherwise False.
-        """
-        total_length = start_length
-        for token in self.iter_tokens():
-            total_length += cell_len(token)
-            if total_length > max_length:
-                return False
-        return True
-
-    def __str__(self) -> str:
-        repr_text = "".join(self.iter_tokens())
-        return repr_text
-
-    def render(
-        self, max_width: int = 80, indent_size: int = 4, expand_all: bool = False
-    ) -> str:
-        """Render the node to a pretty repr.
-
-        Args:
-            max_width (int, optional): Maximum width of the repr. Defaults to 80.
-            indent_size (int, optional): Size of indents. Defaults to 4.
-            expand_all (bool, optional): Expand all levels. Defaults to False.
-
-        Returns:
-            str: A repr string of the original object.
-        """
-        lines = [_Line(node=self, is_root=True)]
-        line_no = 0
-        while line_no < len(lines):
-            line = lines[line_no]
-            if line.expandable and not line.expanded:
-                if expand_all or not line.check_length(max_width):
-                    lines[line_no : line_no + 1] = line.expand(indent_size)
-            line_no += 1
-
-        repr_str = "\n".join(str(line) for line in lines)
-        return repr_str
-
-
-@dataclass
-class _Line:
-    """A line in repr output."""
-
-    parent: Optional["_Line"] = None
-    is_root: bool = False
-    node: Optional[Node] = None
-    text: str = ""
-    suffix: str = ""
-    whitespace: str = ""
-    expanded: bool = False
-    last: bool = False
-
-    @property
-    def expandable(self) -> bool:
-        """Check if the line may be expanded."""
-        return bool(self.node is not None and self.node.children)
-
-    def check_length(self, max_length: int) -> bool:
-        """Check this line fits within a given number of cells."""
-        start_length = (
-            len(self.whitespace) + cell_len(self.text) + cell_len(self.suffix)
-        )
-        assert self.node is not None
-        return self.node.check_length(start_length, max_length)
-
-    def expand(self, indent_size: int) -> Iterable["_Line"]:
-        """Expand this line by adding children on their own line."""
-        node = self.node
-        assert node is not None
-        whitespace = self.whitespace
-        assert node.children
-        if node.key_repr:
-            new_line = yield _Line(
-                text=f"{node.key_repr}{node.key_separator}{node.open_brace}",
-                whitespace=whitespace,
-            )
-        else:
-            new_line = yield _Line(text=node.open_brace, whitespace=whitespace)
-        child_whitespace = self.whitespace + " " * indent_size
-        tuple_of_one = node.is_tuple and len(node.children) == 1
-        for last, child in loop_last(node.children):
-            separator = "," if tuple_of_one else node.separator
-            line = _Line(
-                parent=new_line,
-                node=child,
-                whitespace=child_whitespace,
-                suffix=separator,
-                last=last and not tuple_of_one,
-            )
-            yield line
-
-        yield _Line(
-            text=node.close_brace,
-            whitespace=whitespace,
-            suffix=self.suffix,
-            last=self.last,
-        )
-
-    def __str__(self) -> str:
-        if self.last:
-            return f"{self.whitespace}{self.text}{self.node or ''}"
-        else:
-            return (
-                f"{self.whitespace}{self.text}{self.node or ''}{self.suffix.rstrip()}"
-            )
-
-
-def _is_namedtuple(obj: Any) -> bool:
-    """Checks if an object is most likely a namedtuple. It is possible
-    to craft an object that passes this check and isn't a namedtuple, but
-    there is only a minuscule chance of this happening unintentionally.
-
-    Args:
-        obj (Any): The object to test
-
-    Returns:
-        bool: True if the object is a namedtuple. False otherwise.
-    """
-    try:
-        fields = getattr(obj, "_fields", None)
-    except Exception:
-        # Being very defensive - if we cannot get the attr then its not a namedtuple
-        return False
-    return isinstance(obj, tuple) and isinstance(fields, tuple)
-
-
-def traverse(
-    _object: Any,
-    max_length: Optional[int] = None,
-    max_string: Optional[int] = None,
-    max_depth: Optional[int] = None,
-) -> Node:
-    """Traverse object and generate a tree.
-
-    Args:
-        _object (Any): Object to be traversed.
-        max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
-            Defaults to None.
-        max_string (int, optional): Maximum length of string before truncating, or None to disable truncating.
-            Defaults to None.
-        max_depth (int, optional): Maximum depth of data structures, or None for no maximum.
-            Defaults to None.
-
-    Returns:
-        Node: The root of a tree structure which can be used to render a pretty repr.
-    """
-
-    def to_repr(obj: Any) -> str:
-        """Get repr string for an object, but catch errors."""
-        if (
-            max_string is not None
-            and _safe_isinstance(obj, (bytes, str))
-            and len(obj) > max_string
-        ):
-            truncated = len(obj) - max_string
-            obj_repr = f"{obj[:max_string]!r}+{truncated}"
-        else:
-            try:
-                obj_repr = repr(obj)
-            except Exception as error:
-                obj_repr = f""
-        return obj_repr
-
-    visited_ids: Set[int] = set()
-    push_visited = visited_ids.add
-    pop_visited = visited_ids.remove
-
-    def _traverse(obj: Any, root: bool = False, depth: int = 0) -> Node:
-        """Walk the object depth first."""
-
-        obj_id = id(obj)
-        if obj_id in visited_ids:
-            # Recursion detected
-            return Node(value_repr="...")
-
-        obj_type = type(obj)
-        children: List[Node]
-        reached_max_depth = max_depth is not None and depth >= max_depth
-
-        def iter_rich_args(rich_args: Any) -> Iterable[Union[Any, Tuple[str, Any]]]:
-            for arg in rich_args:
-                if _safe_isinstance(arg, tuple):
-                    if len(arg) == 3:
-                        key, child, default = arg
-                        if default == child:
-                            continue
-                        yield key, child
-                    elif len(arg) == 2:
-                        key, child = arg
-                        yield key, child
-                    elif len(arg) == 1:
-                        yield arg[0]
-                else:
-                    yield arg
-
-        try:
-            fake_attributes = hasattr(
-                obj, "awehoi234_wdfjwljet234_234wdfoijsdfmmnxpi492"
-            )
-        except Exception:
-            fake_attributes = False
-
-        rich_repr_result: Optional[RichReprResult] = None
-        if not fake_attributes:
-            try:
-                if hasattr(obj, "__rich_repr__") and not isclass(obj):
-                    rich_repr_result = obj.__rich_repr__()
-            except Exception:
-                pass
-
-        if rich_repr_result is not None:
-            push_visited(obj_id)
-            angular = getattr(obj.__rich_repr__, "angular", False)
-            args = list(iter_rich_args(rich_repr_result))
-            class_name = obj.__class__.__name__
-
-            if args:
-                children = []
-                append = children.append
-
-                if reached_max_depth:
-                    if angular:
-                        node = Node(value_repr=f"<{class_name}...>")
-                    else:
-                        node = Node(value_repr=f"{class_name}(...)")
-                else:
-                    if angular:
-                        node = Node(
-                            open_brace=f"<{class_name} ",
-                            close_brace=">",
-                            children=children,
-                            last=root,
-                            separator=" ",
-                        )
-                    else:
-                        node = Node(
-                            open_brace=f"{class_name}(",
-                            close_brace=")",
-                            children=children,
-                            last=root,
-                        )
-                    for last, arg in loop_last(args):
-                        if _safe_isinstance(arg, tuple):
-                            key, child = arg
-                            child_node = _traverse(child, depth=depth + 1)
-                            child_node.last = last
-                            child_node.key_repr = key
-                            child_node.key_separator = "="
-                            append(child_node)
-                        else:
-                            child_node = _traverse(arg, depth=depth + 1)
-                            child_node.last = last
-                            append(child_node)
-            else:
-                node = Node(
-                    value_repr=f"<{class_name}>" if angular else f"{class_name}()",
-                    children=[],
-                    last=root,
-                )
-            pop_visited(obj_id)
-        elif _is_attr_object(obj) and not fake_attributes:
-            push_visited(obj_id)
-            children = []
-            append = children.append
-
-            attr_fields = _get_attr_fields(obj)
-            if attr_fields:
-                if reached_max_depth:
-                    node = Node(value_repr=f"{obj.__class__.__name__}(...)")
-                else:
-                    node = Node(
-                        open_brace=f"{obj.__class__.__name__}(",
-                        close_brace=")",
-                        children=children,
-                        last=root,
-                    )
-
-                    def iter_attrs() -> (
-                        Iterable[Tuple[str, Any, Optional[Callable[[Any], str]]]]
-                    ):
-                        """Iterate over attr fields and values."""
-                        for attr in attr_fields:
-                            if attr.repr:
-                                try:
-                                    value = getattr(obj, attr.name)
-                                except Exception as error:
-                                    # Can happen, albeit rarely
-                                    yield (attr.name, error, None)
-                                else:
-                                    yield (
-                                        attr.name,
-                                        value,
-                                        attr.repr if callable(attr.repr) else None,
-                                    )
-
-                    for last, (name, value, repr_callable) in loop_last(iter_attrs()):
-                        if repr_callable:
-                            child_node = Node(value_repr=str(repr_callable(value)))
-                        else:
-                            child_node = _traverse(value, depth=depth + 1)
-                        child_node.last = last
-                        child_node.key_repr = name
-                        child_node.key_separator = "="
-                        append(child_node)
-            else:
-                node = Node(
-                    value_repr=f"{obj.__class__.__name__}()", children=[], last=root
-                )
-            pop_visited(obj_id)
-        elif (
-            is_dataclass(obj)
-            and not _safe_isinstance(obj, type)
-            and not fake_attributes
-            and _is_dataclass_repr(obj)
-        ):
-            push_visited(obj_id)
-            children = []
-            append = children.append
-            if reached_max_depth:
-                node = Node(value_repr=f"{obj.__class__.__name__}(...)")
-            else:
-                node = Node(
-                    open_brace=f"{obj.__class__.__name__}(",
-                    close_brace=")",
-                    children=children,
-                    last=root,
-                    empty=f"{obj.__class__.__name__}()",
-                )
-
-                for last, field in loop_last(
-                    field
-                    for field in fields(obj)
-                    if field.repr and hasattr(obj, field.name)
-                ):
-                    child_node = _traverse(getattr(obj, field.name), depth=depth + 1)
-                    child_node.key_repr = field.name
-                    child_node.last = last
-                    child_node.key_separator = "="
-                    append(child_node)
-
-            pop_visited(obj_id)
-        elif _is_namedtuple(obj) and _has_default_namedtuple_repr(obj):
-            push_visited(obj_id)
-            class_name = obj.__class__.__name__
-            if reached_max_depth:
-                # If we've reached the max depth, we still show the class name, but not its contents
-                node = Node(
-                    value_repr=f"{class_name}(...)",
-                )
-            else:
-                children = []
-                append = children.append
-                node = Node(
-                    open_brace=f"{class_name}(",
-                    close_brace=")",
-                    children=children,
-                    empty=f"{class_name}()",
-                )
-                for last, (key, value) in loop_last(obj._asdict().items()):
-                    child_node = _traverse(value, depth=depth + 1)
-                    child_node.key_repr = key
-                    child_node.last = last
-                    child_node.key_separator = "="
-                    append(child_node)
-            pop_visited(obj_id)
-        elif _safe_isinstance(obj, _CONTAINERS):
-            for container_type in _CONTAINERS:
-                if _safe_isinstance(obj, container_type):
-                    obj_type = container_type
-                    break
-
-            push_visited(obj_id)
-
-            open_brace, close_brace, empty = _BRACES[obj_type](obj)
-
-            if reached_max_depth:
-                node = Node(value_repr=f"{open_brace}...{close_brace}")
-            elif obj_type.__repr__ != type(obj).__repr__:
-                node = Node(value_repr=to_repr(obj), last=root)
-            elif obj:
-                children = []
-                node = Node(
-                    open_brace=open_brace,
-                    close_brace=close_brace,
-                    children=children,
-                    last=root,
-                )
-                append = children.append
-                num_items = len(obj)
-                last_item_index = num_items - 1
-
-                if _safe_isinstance(obj, _MAPPING_CONTAINERS):
-                    iter_items = iter(obj.items())
-                    if max_length is not None:
-                        iter_items = islice(iter_items, max_length)
-                    for index, (key, child) in enumerate(iter_items):
-                        child_node = _traverse(child, depth=depth + 1)
-                        child_node.key_repr = to_repr(key)
-                        child_node.last = index == last_item_index
-                        append(child_node)
-                else:
-                    iter_values = iter(obj)
-                    if max_length is not None:
-                        iter_values = islice(iter_values, max_length)
-                    for index, child in enumerate(iter_values):
-                        child_node = _traverse(child, depth=depth + 1)
-                        child_node.last = index == last_item_index
-                        append(child_node)
-                if max_length is not None and num_items > max_length:
-                    append(Node(value_repr=f"... +{num_items - max_length}", last=True))
-            else:
-                node = Node(empty=empty, children=[], last=root)
-
-            pop_visited(obj_id)
-        else:
-            node = Node(value_repr=to_repr(obj), last=root)
-        node.is_tuple = type(obj) == tuple
-        node.is_namedtuple = _is_namedtuple(obj)
-        return node
-
-    node = _traverse(_object, root=True)
-    return node
-
-
-def pretty_repr(
-    _object: Any,
-    *,
-    max_width: int = 80,
-    indent_size: int = 4,
-    max_length: Optional[int] = None,
-    max_string: Optional[int] = None,
-    max_depth: Optional[int] = None,
-    expand_all: bool = False,
-) -> str:
-    """Prettify repr string by expanding on to new lines to fit within a given width.
-
-    Args:
-        _object (Any): Object to repr.
-        max_width (int, optional): Desired maximum width of repr string. Defaults to 80.
-        indent_size (int, optional): Number of spaces to indent. Defaults to 4.
-        max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
-            Defaults to None.
-        max_string (int, optional): Maximum length of string before truncating, or None to disable truncating.
-            Defaults to None.
-        max_depth (int, optional): Maximum depth of nested data structure, or None for no depth.
-            Defaults to None.
-        expand_all (bool, optional): Expand all containers regardless of available width. Defaults to False.
-
-    Returns:
-        str: A possibly multi-line representation of the object.
-    """
-
-    if _safe_isinstance(_object, Node):
-        node = _object
-    else:
-        node = traverse(
-            _object, max_length=max_length, max_string=max_string, max_depth=max_depth
-        )
-    repr_str: str = node.render(
-        max_width=max_width, indent_size=indent_size, expand_all=expand_all
-    )
-    return repr_str
-
-
-def pprint(
-    _object: Any,
-    *,
-    console: Optional["Console"] = None,
-    indent_guides: bool = True,
-    max_length: Optional[int] = None,
-    max_string: Optional[int] = None,
-    max_depth: Optional[int] = None,
-    expand_all: bool = False,
-) -> None:
-    """A convenience function for pretty printing.
-
-    Args:
-        _object (Any): Object to pretty print.
-        console (Console, optional): Console instance, or None to use default. Defaults to None.
-        max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
-            Defaults to None.
-        max_string (int, optional): Maximum length of strings before truncating, or None to disable. Defaults to None.
-        max_depth (int, optional): Maximum depth for nested data structures, or None for unlimited depth. Defaults to None.
-        indent_guides (bool, optional): Enable indentation guides. Defaults to True.
-        expand_all (bool, optional): Expand all containers. Defaults to False.
-    """
-    _console = get_console() if console is None else console
-    _console.print(
-        Pretty(
-            _object,
-            max_length=max_length,
-            max_string=max_string,
-            max_depth=max_depth,
-            indent_guides=indent_guides,
-            expand_all=expand_all,
-            overflow="ignore",
-        ),
-        soft_wrap=True,
-    )
-
-
-if __name__ == "__main__":  # pragma: no cover
-
-    class BrokenRepr:
-        def __repr__(self) -> str:
-            1 / 0
-            return "this will fail"
-
-    from typing import NamedTuple
-
-    class StockKeepingUnit(NamedTuple):
-        name: str
-        description: str
-        price: float
-        category: str
-        reviews: List[str]
-
-    d = defaultdict(int)
-    d["foo"] = 5
-    data = {
-        "foo": [
-            1,
-            "Hello World!",
-            100.123,
-            323.232,
-            432324.0,
-            {5, 6, 7, (1, 2, 3, 4), 8},
-        ],
-        "bar": frozenset({1, 2, 3}),
-        "defaultdict": defaultdict(
-            list, {"crumble": ["apple", "rhubarb", "butter", "sugar", "flour"]}
-        ),
-        "counter": Counter(
-            [
-                "apple",
-                "orange",
-                "pear",
-                "kumquat",
-                "kumquat",
-                "durian" * 100,
-            ]
-        ),
-        "atomic": (False, True, None),
-        "namedtuple": StockKeepingUnit(
-            "Sparkling British Spring Water",
-            "Carbonated spring water",
-            0.9,
-            "water",
-            ["its amazing!", "its terrible!"],
-        ),
-        "Broken": BrokenRepr(),
-    }
-    data["foo"].append(data)  # type: ignore[attr-defined]
-
-    from rich import print
-
-    print(Pretty(data, indent_guides=True, max_string=20))
-
-    class Thing:
-        def __repr__(self) -> str:
-            return "Hello\x1b[38;5;239m World!"
-
-    print(Pretty(Thing()))
diff --git a/venv/lib/python3.10/site-packages/rich/progress.py b/venv/lib/python3.10/site-packages/rich/progress.py
deleted file mode 100644
index 1e92eb6b1362515efa93395f3484088b57db8eb8..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/progress.py
+++ /dev/null
@@ -1,1715 +0,0 @@
-import io
-import sys
-import typing
-import warnings
-from abc import ABC, abstractmethod
-from collections import deque
-from dataclasses import dataclass, field
-from datetime import timedelta
-from io import RawIOBase, UnsupportedOperation
-from math import ceil
-from mmap import mmap
-from operator import length_hint
-from os import PathLike, stat
-from threading import Event, RLock, Thread
-from types import TracebackType
-from typing import (
-    Any,
-    BinaryIO,
-    Callable,
-    ContextManager,
-    Deque,
-    Dict,
-    Generic,
-    Iterable,
-    List,
-    NamedTuple,
-    NewType,
-    Optional,
-    Sequence,
-    TextIO,
-    Tuple,
-    Type,
-    TypeVar,
-    Union,
-)
-
-if sys.version_info >= (3, 8):
-    from typing import Literal
-else:
-    from typing_extensions import Literal  # pragma: no cover
-
-if sys.version_info >= (3, 11):
-    from typing import Self
-else:
-    from typing_extensions import Self  # pragma: no cover
-
-from . import filesize, get_console
-from .console import Console, Group, JustifyMethod, RenderableType
-from .highlighter import Highlighter
-from .jupyter import JupyterMixin
-from .live import Live
-from .progress_bar import ProgressBar
-from .spinner import Spinner
-from .style import StyleType
-from .table import Column, Table
-from .text import Text, TextType
-
-TaskID = NewType("TaskID", int)
-
-ProgressType = TypeVar("ProgressType")
-
-GetTimeCallable = Callable[[], float]
-
-
-_I = typing.TypeVar("_I", TextIO, BinaryIO)
-
-
-class _TrackThread(Thread):
-    """A thread to periodically update progress."""
-
-    def __init__(self, progress: "Progress", task_id: "TaskID", update_period: float):
-        self.progress = progress
-        self.task_id = task_id
-        self.update_period = update_period
-        self.done = Event()
-
-        self.completed = 0
-        super().__init__(daemon=True)
-
-    def run(self) -> None:
-        task_id = self.task_id
-        advance = self.progress.advance
-        update_period = self.update_period
-        last_completed = 0
-        wait = self.done.wait
-        while not wait(update_period) and self.progress.live.is_started:
-            completed = self.completed
-            if last_completed != completed:
-                advance(task_id, completed - last_completed)
-                last_completed = completed
-
-        self.progress.update(self.task_id, completed=self.completed, refresh=True)
-
-    def __enter__(self) -> "_TrackThread":
-        self.start()
-        return self
-
-    def __exit__(
-        self,
-        exc_type: Optional[Type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
-    ) -> None:
-        self.done.set()
-        self.join()
-
-
-def track(
-    sequence: Union[Sequence[ProgressType], Iterable[ProgressType]],
-    description: str = "Working...",
-    total: Optional[float] = None,
-    completed: int = 0,
-    auto_refresh: bool = True,
-    console: Optional[Console] = None,
-    transient: bool = False,
-    get_time: Optional[Callable[[], float]] = None,
-    refresh_per_second: float = 10,
-    style: StyleType = "bar.back",
-    complete_style: StyleType = "bar.complete",
-    finished_style: StyleType = "bar.finished",
-    pulse_style: StyleType = "bar.pulse",
-    update_period: float = 0.1,
-    disable: bool = False,
-    show_speed: bool = True,
-) -> Iterable[ProgressType]:
-    """Track progress by iterating over a sequence.
-
-    Args:
-        sequence (Iterable[ProgressType]): A sequence (must support "len") you wish to iterate over.
-        description (str, optional): Description of task show next to progress bar. Defaults to "Working".
-        total: (float, optional): Total number of steps. Default is len(sequence).
-        completed (int, optional): Number of steps completed so far. Defaults to 0.
-        auto_refresh (bool, optional): Automatic refresh, disable to force a refresh after each iteration. Default is True.
-        transient: (bool, optional): Clear the progress on exit. Defaults to False.
-        console (Console, optional): Console to write to. Default creates internal Console instance.
-        refresh_per_second (float): Number of times per second to refresh the progress information. Defaults to 10.
-        style (StyleType, optional): Style for the bar background. Defaults to "bar.back".
-        complete_style (StyleType, optional): Style for the completed bar. Defaults to "bar.complete".
-        finished_style (StyleType, optional): Style for a finished bar. Defaults to "bar.finished".
-        pulse_style (StyleType, optional): Style for pulsing bars. Defaults to "bar.pulse".
-        update_period (float, optional): Minimum time (in seconds) between calls to update(). Defaults to 0.1.
-        disable (bool, optional): Disable display of progress.
-        show_speed (bool, optional): Show speed if total isn't known. Defaults to True.
-    Returns:
-        Iterable[ProgressType]: An iterable of the values in the sequence.
-
-    """
-
-    columns: List["ProgressColumn"] = (
-        [TextColumn("[progress.description]{task.description}")] if description else []
-    )
-    columns.extend(
-        (
-            BarColumn(
-                style=style,
-                complete_style=complete_style,
-                finished_style=finished_style,
-                pulse_style=pulse_style,
-            ),
-            TaskProgressColumn(show_speed=show_speed),
-            TimeRemainingColumn(elapsed_when_finished=True),
-        )
-    )
-    progress = Progress(
-        *columns,
-        auto_refresh=auto_refresh,
-        console=console,
-        transient=transient,
-        get_time=get_time,
-        refresh_per_second=refresh_per_second or 10,
-        disable=disable,
-    )
-
-    with progress:
-        yield from progress.track(
-            sequence,
-            total=total,
-            completed=completed,
-            description=description,
-            update_period=update_period,
-        )
-
-
-class _Reader(RawIOBase, BinaryIO):
-    """A reader that tracks progress while it's being read from."""
-
-    def __init__(
-        self,
-        handle: BinaryIO,
-        progress: "Progress",
-        task: TaskID,
-        close_handle: bool = True,
-    ) -> None:
-        self.handle = handle
-        self.progress = progress
-        self.task = task
-        self.close_handle = close_handle
-        self._closed = False
-
-    def __enter__(self) -> "_Reader":
-        self.handle.__enter__()
-        return self
-
-    def __exit__(
-        self,
-        exc_type: Optional[Type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
-    ) -> None:
-        self.close()
-
-    def __iter__(self) -> BinaryIO:
-        return self
-
-    def __next__(self) -> bytes:
-        line = next(self.handle)
-        self.progress.advance(self.task, advance=len(line))
-        return line
-
-    @property
-    def closed(self) -> bool:
-        return self._closed
-
-    def fileno(self) -> int:
-        return self.handle.fileno()
-
-    def isatty(self) -> bool:
-        return self.handle.isatty()
-
-    @property
-    def mode(self) -> str:
-        return self.handle.mode
-
-    @property
-    def name(self) -> str:
-        return self.handle.name
-
-    def readable(self) -> bool:
-        return self.handle.readable()
-
-    def seekable(self) -> bool:
-        return self.handle.seekable()
-
-    def writable(self) -> bool:
-        return False
-
-    def read(self, size: int = -1) -> bytes:
-        block = self.handle.read(size)
-        self.progress.advance(self.task, advance=len(block))
-        return block
-
-    def readinto(self, b: Union[bytearray, memoryview, mmap]):  # type: ignore[no-untyped-def, override]
-        n = self.handle.readinto(b)  # type: ignore[attr-defined]
-        self.progress.advance(self.task, advance=n)
-        return n
-
-    def readline(self, size: int = -1) -> bytes:  # type: ignore[override]
-        line = self.handle.readline(size)
-        self.progress.advance(self.task, advance=len(line))
-        return line
-
-    def readlines(self, hint: int = -1) -> List[bytes]:
-        lines = self.handle.readlines(hint)
-        self.progress.advance(self.task, advance=sum(map(len, lines)))
-        return lines
-
-    def close(self) -> None:
-        if self.close_handle:
-            self.handle.close()
-        self._closed = True
-
-    def seek(self, offset: int, whence: int = 0) -> int:
-        pos = self.handle.seek(offset, whence)
-        self.progress.update(self.task, completed=pos)
-        return pos
-
-    def tell(self) -> int:
-        return self.handle.tell()
-
-    def write(self, s: Any) -> int:
-        raise UnsupportedOperation("write")
-
-    def writelines(self, lines: Iterable[Any]) -> None:
-        raise UnsupportedOperation("writelines")
-
-
-class _ReadContext(ContextManager[_I], Generic[_I]):
-    """A utility class to handle a context for both a reader and a progress."""
-
-    def __init__(self, progress: "Progress", reader: _I) -> None:
-        self.progress = progress
-        self.reader: _I = reader
-
-    def __enter__(self) -> _I:
-        self.progress.start()
-        return self.reader.__enter__()
-
-    def __exit__(
-        self,
-        exc_type: Optional[Type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
-    ) -> None:
-        self.progress.stop()
-        self.reader.__exit__(exc_type, exc_val, exc_tb)
-
-
-def wrap_file(
-    file: BinaryIO,
-    total: int,
-    *,
-    description: str = "Reading...",
-    auto_refresh: bool = True,
-    console: Optional[Console] = None,
-    transient: bool = False,
-    get_time: Optional[Callable[[], float]] = None,
-    refresh_per_second: float = 10,
-    style: StyleType = "bar.back",
-    complete_style: StyleType = "bar.complete",
-    finished_style: StyleType = "bar.finished",
-    pulse_style: StyleType = "bar.pulse",
-    disable: bool = False,
-) -> ContextManager[BinaryIO]:
-    """Read bytes from a file while tracking progress.
-
-    Args:
-        file (Union[str, PathLike[str], BinaryIO]): The path to the file to read, or a file-like object in binary mode.
-        total (int): Total number of bytes to read.
-        description (str, optional): Description of task show next to progress bar. Defaults to "Reading".
-        auto_refresh (bool, optional): Automatic refresh, disable to force a refresh after each iteration. Default is True.
-        transient: (bool, optional): Clear the progress on exit. Defaults to False.
-        console (Console, optional): Console to write to. Default creates internal Console instance.
-        refresh_per_second (float): Number of times per second to refresh the progress information. Defaults to 10.
-        style (StyleType, optional): Style for the bar background. Defaults to "bar.back".
-        complete_style (StyleType, optional): Style for the completed bar. Defaults to "bar.complete".
-        finished_style (StyleType, optional): Style for a finished bar. Defaults to "bar.finished".
-        pulse_style (StyleType, optional): Style for pulsing bars. Defaults to "bar.pulse".
-        disable (bool, optional): Disable display of progress.
-    Returns:
-        ContextManager[BinaryIO]: A context manager yielding a progress reader.
-
-    """
-
-    columns: List["ProgressColumn"] = (
-        [TextColumn("[progress.description]{task.description}")] if description else []
-    )
-    columns.extend(
-        (
-            BarColumn(
-                style=style,
-                complete_style=complete_style,
-                finished_style=finished_style,
-                pulse_style=pulse_style,
-            ),
-            DownloadColumn(),
-            TimeRemainingColumn(),
-        )
-    )
-    progress = Progress(
-        *columns,
-        auto_refresh=auto_refresh,
-        console=console,
-        transient=transient,
-        get_time=get_time,
-        refresh_per_second=refresh_per_second or 10,
-        disable=disable,
-    )
-
-    reader = progress.wrap_file(file, total=total, description=description)
-    return _ReadContext(progress, reader)
-
-
-@typing.overload
-def open(
-    file: Union[str, "PathLike[str]", bytes],
-    mode: Union[Literal["rt"], Literal["r"]],
-    buffering: int = -1,
-    encoding: Optional[str] = None,
-    errors: Optional[str] = None,
-    newline: Optional[str] = None,
-    *,
-    total: Optional[int] = None,
-    description: str = "Reading...",
-    auto_refresh: bool = True,
-    console: Optional[Console] = None,
-    transient: bool = False,
-    get_time: Optional[Callable[[], float]] = None,
-    refresh_per_second: float = 10,
-    style: StyleType = "bar.back",
-    complete_style: StyleType = "bar.complete",
-    finished_style: StyleType = "bar.finished",
-    pulse_style: StyleType = "bar.pulse",
-    disable: bool = False,
-) -> ContextManager[TextIO]:
-    pass
-
-
-@typing.overload
-def open(
-    file: Union[str, "PathLike[str]", bytes],
-    mode: Literal["rb"],
-    buffering: int = -1,
-    encoding: Optional[str] = None,
-    errors: Optional[str] = None,
-    newline: Optional[str] = None,
-    *,
-    total: Optional[int] = None,
-    description: str = "Reading...",
-    auto_refresh: bool = True,
-    console: Optional[Console] = None,
-    transient: bool = False,
-    get_time: Optional[Callable[[], float]] = None,
-    refresh_per_second: float = 10,
-    style: StyleType = "bar.back",
-    complete_style: StyleType = "bar.complete",
-    finished_style: StyleType = "bar.finished",
-    pulse_style: StyleType = "bar.pulse",
-    disable: bool = False,
-) -> ContextManager[BinaryIO]:
-    pass
-
-
-def open(
-    file: Union[str, "PathLike[str]", bytes],
-    mode: Union[Literal["rb"], Literal["rt"], Literal["r"]] = "r",
-    buffering: int = -1,
-    encoding: Optional[str] = None,
-    errors: Optional[str] = None,
-    newline: Optional[str] = None,
-    *,
-    total: Optional[int] = None,
-    description: str = "Reading...",
-    auto_refresh: bool = True,
-    console: Optional[Console] = None,
-    transient: bool = False,
-    get_time: Optional[Callable[[], float]] = None,
-    refresh_per_second: float = 10,
-    style: StyleType = "bar.back",
-    complete_style: StyleType = "bar.complete",
-    finished_style: StyleType = "bar.finished",
-    pulse_style: StyleType = "bar.pulse",
-    disable: bool = False,
-) -> Union[ContextManager[BinaryIO], ContextManager[TextIO]]:
-    """Read bytes from a file while tracking progress.
-
-    Args:
-        path (Union[str, PathLike[str], BinaryIO]): The path to the file to read, or a file-like object in binary mode.
-        mode (str): The mode to use to open the file. Only supports "r", "rb" or "rt".
-        buffering (int): The buffering strategy to use, see :func:`io.open`.
-        encoding (str, optional): The encoding to use when reading in text mode, see :func:`io.open`.
-        errors (str, optional): The error handling strategy for decoding errors, see :func:`io.open`.
-        newline (str, optional): The strategy for handling newlines in text mode, see :func:`io.open`
-        total: (int, optional): Total number of bytes to read. Must be provided if reading from a file handle. Default for a path is os.stat(file).st_size.
-        description (str, optional): Description of task show next to progress bar. Defaults to "Reading".
-        auto_refresh (bool, optional): Automatic refresh, disable to force a refresh after each iteration. Default is True.
-        transient: (bool, optional): Clear the progress on exit. Defaults to False.
-        console (Console, optional): Console to write to. Default creates internal Console instance.
-        refresh_per_second (float): Number of times per second to refresh the progress information. Defaults to 10.
-        style (StyleType, optional): Style for the bar background. Defaults to "bar.back".
-        complete_style (StyleType, optional): Style for the completed bar. Defaults to "bar.complete".
-        finished_style (StyleType, optional): Style for a finished bar. Defaults to "bar.finished".
-        pulse_style (StyleType, optional): Style for pulsing bars. Defaults to "bar.pulse".
-        disable (bool, optional): Disable display of progress.
-        encoding (str, optional): The encoding to use when reading in text mode.
-
-    Returns:
-        ContextManager[BinaryIO]: A context manager yielding a progress reader.
-
-    """
-
-    columns: List["ProgressColumn"] = (
-        [TextColumn("[progress.description]{task.description}")] if description else []
-    )
-    columns.extend(
-        (
-            BarColumn(
-                style=style,
-                complete_style=complete_style,
-                finished_style=finished_style,
-                pulse_style=pulse_style,
-            ),
-            DownloadColumn(),
-            TimeRemainingColumn(),
-        )
-    )
-    progress = Progress(
-        *columns,
-        auto_refresh=auto_refresh,
-        console=console,
-        transient=transient,
-        get_time=get_time,
-        refresh_per_second=refresh_per_second or 10,
-        disable=disable,
-    )
-
-    reader = progress.open(
-        file,
-        mode=mode,
-        buffering=buffering,
-        encoding=encoding,
-        errors=errors,
-        newline=newline,
-        total=total,
-        description=description,
-    )
-    return _ReadContext(progress, reader)  # type: ignore[return-value, type-var]
-
-
-class ProgressColumn(ABC):
-    """Base class for a widget to use in progress display."""
-
-    max_refresh: Optional[float] = None
-
-    def __init__(self, table_column: Optional[Column] = None) -> None:
-        self._table_column = table_column
-        self._renderable_cache: Dict[TaskID, Tuple[float, RenderableType]] = {}
-        self._update_time: Optional[float] = None
-
-    def get_table_column(self) -> Column:
-        """Get a table column, used to build tasks table."""
-        return self._table_column or Column()
-
-    def __call__(self, task: "Task") -> RenderableType:
-        """Called by the Progress object to return a renderable for the given task.
-
-        Args:
-            task (Task): An object containing information regarding the task.
-
-        Returns:
-            RenderableType: Anything renderable (including str).
-        """
-        current_time = task.get_time()
-        if self.max_refresh is not None and not task.completed:
-            try:
-                timestamp, renderable = self._renderable_cache[task.id]
-            except KeyError:
-                pass
-            else:
-                if timestamp + self.max_refresh > current_time:
-                    return renderable
-
-        renderable = self.render(task)
-        self._renderable_cache[task.id] = (current_time, renderable)
-        return renderable
-
-    @abstractmethod
-    def render(self, task: "Task") -> RenderableType:
-        """Should return a renderable object."""
-
-
-class RenderableColumn(ProgressColumn):
-    """A column to insert an arbitrary column.
-
-    Args:
-        renderable (RenderableType, optional): Any renderable. Defaults to empty string.
-    """
-
-    def __init__(
-        self, renderable: RenderableType = "", *, table_column: Optional[Column] = None
-    ):
-        self.renderable = renderable
-        super().__init__(table_column=table_column)
-
-    def render(self, task: "Task") -> RenderableType:
-        return self.renderable
-
-
-class SpinnerColumn(ProgressColumn):
-    """A column with a 'spinner' animation.
-
-    Args:
-        spinner_name (str, optional): Name of spinner animation. Defaults to "dots".
-        style (StyleType, optional): Style of spinner. Defaults to "progress.spinner".
-        speed (float, optional): Speed factor of spinner. Defaults to 1.0.
-        finished_text (TextType, optional): Text used when task is finished. Defaults to " ".
-    """
-
-    def __init__(
-        self,
-        spinner_name: str = "dots",
-        style: Optional[StyleType] = "progress.spinner",
-        speed: float = 1.0,
-        finished_text: TextType = " ",
-        table_column: Optional[Column] = None,
-    ):
-        self.spinner = Spinner(spinner_name, style=style, speed=speed)
-        self.finished_text = (
-            Text.from_markup(finished_text)
-            if isinstance(finished_text, str)
-            else finished_text
-        )
-        super().__init__(table_column=table_column)
-
-    def set_spinner(
-        self,
-        spinner_name: str,
-        spinner_style: Optional[StyleType] = "progress.spinner",
-        speed: float = 1.0,
-    ) -> None:
-        """Set a new spinner.
-
-        Args:
-            spinner_name (str): Spinner name, see python -m rich.spinner.
-            spinner_style (Optional[StyleType], optional): Spinner style. Defaults to "progress.spinner".
-            speed (float, optional): Speed factor of spinner. Defaults to 1.0.
-        """
-        self.spinner = Spinner(spinner_name, style=spinner_style, speed=speed)
-
-    def render(self, task: "Task") -> RenderableType:
-        text = (
-            self.finished_text
-            if task.finished
-            else self.spinner.render(task.get_time())
-        )
-        return text
-
-
-class TextColumn(ProgressColumn):
-    """A column containing text."""
-
-    def __init__(
-        self,
-        text_format: str,
-        style: StyleType = "none",
-        justify: JustifyMethod = "left",
-        markup: bool = True,
-        highlighter: Optional[Highlighter] = None,
-        table_column: Optional[Column] = None,
-    ) -> None:
-        self.text_format = text_format
-        self.justify: JustifyMethod = justify
-        self.style = style
-        self.markup = markup
-        self.highlighter = highlighter
-        super().__init__(table_column=table_column or Column(no_wrap=True))
-
-    def render(self, task: "Task") -> Text:
-        _text = self.text_format.format(task=task)
-        if self.markup:
-            text = Text.from_markup(_text, style=self.style, justify=self.justify)
-        else:
-            text = Text(_text, style=self.style, justify=self.justify)
-        if self.highlighter:
-            self.highlighter.highlight(text)
-        return text
-
-
-class BarColumn(ProgressColumn):
-    """Renders a visual progress bar.
-
-    Args:
-        bar_width (Optional[int], optional): Width of bar or None for full width. Defaults to 40.
-        style (StyleType, optional): Style for the bar background. Defaults to "bar.back".
-        complete_style (StyleType, optional): Style for the completed bar. Defaults to "bar.complete".
-        finished_style (StyleType, optional): Style for a finished bar. Defaults to "bar.finished".
-        pulse_style (StyleType, optional): Style for pulsing bars. Defaults to "bar.pulse".
-    """
-
-    def __init__(
-        self,
-        bar_width: Optional[int] = 40,
-        style: StyleType = "bar.back",
-        complete_style: StyleType = "bar.complete",
-        finished_style: StyleType = "bar.finished",
-        pulse_style: StyleType = "bar.pulse",
-        table_column: Optional[Column] = None,
-    ) -> None:
-        self.bar_width = bar_width
-        self.style = style
-        self.complete_style = complete_style
-        self.finished_style = finished_style
-        self.pulse_style = pulse_style
-        super().__init__(table_column=table_column)
-
-    def render(self, task: "Task") -> ProgressBar:
-        """Gets a progress bar widget for a task."""
-        return ProgressBar(
-            total=max(0, task.total) if task.total is not None else None,
-            completed=max(0, task.completed),
-            width=None if self.bar_width is None else max(1, self.bar_width),
-            pulse=not task.started,
-            animation_time=task.get_time(),
-            style=self.style,
-            complete_style=self.complete_style,
-            finished_style=self.finished_style,
-            pulse_style=self.pulse_style,
-        )
-
-
-class TimeElapsedColumn(ProgressColumn):
-    """Renders time elapsed."""
-
-    def render(self, task: "Task") -> Text:
-        """Show time elapsed."""
-        elapsed = task.finished_time if task.finished else task.elapsed
-        if elapsed is None:
-            return Text("-:--:--", style="progress.elapsed")
-        delta = timedelta(seconds=max(0, int(elapsed)))
-        return Text(str(delta), style="progress.elapsed")
-
-
-class TaskProgressColumn(TextColumn):
-    """Show task progress as a percentage.
-
-    Args:
-        text_format (str, optional): Format for percentage display. Defaults to "[progress.percentage]{task.percentage:>3.0f}%".
-        text_format_no_percentage (str, optional): Format if percentage is unknown. Defaults to "".
-        style (StyleType, optional): Style of output. Defaults to "none".
-        justify (JustifyMethod, optional): Text justification. Defaults to "left".
-        markup (bool, optional): Enable markup. Defaults to True.
-        highlighter (Optional[Highlighter], optional): Highlighter to apply to output. Defaults to None.
-        table_column (Optional[Column], optional): Table Column to use. Defaults to None.
-        show_speed (bool, optional): Show speed if total is unknown. Defaults to False.
-    """
-
-    def __init__(
-        self,
-        text_format: str = "[progress.percentage]{task.percentage:>3.0f}%",
-        text_format_no_percentage: str = "",
-        style: StyleType = "none",
-        justify: JustifyMethod = "left",
-        markup: bool = True,
-        highlighter: Optional[Highlighter] = None,
-        table_column: Optional[Column] = None,
-        show_speed: bool = False,
-    ) -> None:
-        self.text_format_no_percentage = text_format_no_percentage
-        self.show_speed = show_speed
-        super().__init__(
-            text_format=text_format,
-            style=style,
-            justify=justify,
-            markup=markup,
-            highlighter=highlighter,
-            table_column=table_column,
-        )
-
-    @classmethod
-    def render_speed(cls, speed: Optional[float]) -> Text:
-        """Render the speed in iterations per second.
-
-        Args:
-            task (Task): A Task object.
-
-        Returns:
-            Text: Text object containing the task speed.
-        """
-        if speed is None:
-            return Text("", style="progress.percentage")
-        unit, suffix = filesize.pick_unit_and_suffix(
-            int(speed),
-            ["", "×10³", "×10⁶", "×10⁹", "×10¹²"],
-            1000,
-        )
-        data_speed = speed / unit
-        return Text(f"{data_speed:.1f}{suffix} it/s", style="progress.percentage")
-
-    def render(self, task: "Task") -> Text:
-        if task.total is None and self.show_speed:
-            return self.render_speed(task.finished_speed or task.speed)
-        text_format = (
-            self.text_format_no_percentage if task.total is None else self.text_format
-        )
-        _text = text_format.format(task=task)
-        if self.markup:
-            text = Text.from_markup(_text, style=self.style, justify=self.justify)
-        else:
-            text = Text(_text, style=self.style, justify=self.justify)
-        if self.highlighter:
-            self.highlighter.highlight(text)
-        return text
-
-
-class TimeRemainingColumn(ProgressColumn):
-    """Renders estimated time remaining.
-
-    Args:
-        compact (bool, optional): Render MM:SS when time remaining is less than an hour. Defaults to False.
-        elapsed_when_finished (bool, optional): Render time elapsed when the task is finished. Defaults to False.
-    """
-
-    # Only refresh twice a second to prevent jitter
-    max_refresh = 0.5
-
-    def __init__(
-        self,
-        compact: bool = False,
-        elapsed_when_finished: bool = False,
-        table_column: Optional[Column] = None,
-    ):
-        self.compact = compact
-        self.elapsed_when_finished = elapsed_when_finished
-        super().__init__(table_column=table_column)
-
-    def render(self, task: "Task") -> Text:
-        """Show time remaining."""
-        if self.elapsed_when_finished and task.finished:
-            task_time = task.finished_time
-            style = "progress.elapsed"
-        else:
-            task_time = task.time_remaining
-            style = "progress.remaining"
-
-        if task.total is None:
-            return Text("", style=style)
-
-        if task_time is None:
-            return Text("--:--" if self.compact else "-:--:--", style=style)
-
-        # Based on https://github.com/tqdm/tqdm/blob/master/tqdm/std.py
-        minutes, seconds = divmod(int(task_time), 60)
-        hours, minutes = divmod(minutes, 60)
-
-        if self.compact and not hours:
-            formatted = f"{minutes:02d}:{seconds:02d}"
-        else:
-            formatted = f"{hours:d}:{minutes:02d}:{seconds:02d}"
-
-        return Text(formatted, style=style)
-
-
-class FileSizeColumn(ProgressColumn):
-    """Renders completed filesize."""
-
-    def render(self, task: "Task") -> Text:
-        """Show data completed."""
-        data_size = filesize.decimal(int(task.completed))
-        return Text(data_size, style="progress.filesize")
-
-
-class TotalFileSizeColumn(ProgressColumn):
-    """Renders total filesize."""
-
-    def render(self, task: "Task") -> Text:
-        """Show data completed."""
-        data_size = filesize.decimal(int(task.total)) if task.total is not None else ""
-        return Text(data_size, style="progress.filesize.total")
-
-
-class MofNCompleteColumn(ProgressColumn):
-    """Renders completed count/total, e.g. '  10/1000'.
-
-    Best for bounded tasks with int quantities.
-
-    Space pads the completed count so that progress length does not change as task progresses
-    past powers of 10.
-
-    Args:
-        separator (str, optional): Text to separate completed and total values. Defaults to "/".
-    """
-
-    def __init__(self, separator: str = "/", table_column: Optional[Column] = None):
-        self.separator = separator
-        super().__init__(table_column=table_column)
-
-    def render(self, task: "Task") -> Text:
-        """Show completed/total."""
-        completed = int(task.completed)
-        total = int(task.total) if task.total is not None else "?"
-        total_width = len(str(total))
-        return Text(
-            f"{completed:{total_width}d}{self.separator}{total}",
-            style="progress.download",
-        )
-
-
-class DownloadColumn(ProgressColumn):
-    """Renders file size downloaded and total, e.g. '0.5/2.3 GB'.
-
-    Args:
-        binary_units (bool, optional): Use binary units, KiB, MiB etc. Defaults to False.
-    """
-
-    def __init__(
-        self, binary_units: bool = False, table_column: Optional[Column] = None
-    ) -> None:
-        self.binary_units = binary_units
-        super().__init__(table_column=table_column)
-
-    def render(self, task: "Task") -> Text:
-        """Calculate common unit for completed and total."""
-        completed = int(task.completed)
-
-        unit_and_suffix_calculation_base = (
-            int(task.total) if task.total is not None else completed
-        )
-        if self.binary_units:
-            unit, suffix = filesize.pick_unit_and_suffix(
-                unit_and_suffix_calculation_base,
-                ["bytes", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB"],
-                1024,
-            )
-        else:
-            unit, suffix = filesize.pick_unit_and_suffix(
-                unit_and_suffix_calculation_base,
-                ["bytes", "kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"],
-                1000,
-            )
-        precision = 0 if unit == 1 else 1
-
-        completed_ratio = completed / unit
-        completed_str = f"{completed_ratio:,.{precision}f}"
-
-        if task.total is not None:
-            total = int(task.total)
-            total_ratio = total / unit
-            total_str = f"{total_ratio:,.{precision}f}"
-        else:
-            total_str = "?"
-
-        download_status = f"{completed_str}/{total_str} {suffix}"
-        download_text = Text(download_status, style="progress.download")
-        return download_text
-
-
-class TransferSpeedColumn(ProgressColumn):
-    """Renders human readable transfer speed."""
-
-    def render(self, task: "Task") -> Text:
-        """Show data transfer speed."""
-        speed = task.finished_speed or task.speed
-        if speed is None:
-            return Text("?", style="progress.data.speed")
-        data_speed = filesize.decimal(int(speed))
-        return Text(f"{data_speed}/s", style="progress.data.speed")
-
-
-class ProgressSample(NamedTuple):
-    """Sample of progress for a given time."""
-
-    timestamp: float
-    """Timestamp of sample."""
-    completed: float
-    """Number of steps completed."""
-
-
-@dataclass
-class Task:
-    """Information regarding a progress task.
-
-    This object should be considered read-only outside of the :class:`~Progress` class.
-
-    """
-
-    id: TaskID
-    """Task ID associated with this task (used in Progress methods)."""
-
-    description: str
-    """str: Description of the task."""
-
-    total: Optional[float]
-    """Optional[float]: Total number of steps in this task."""
-
-    completed: float
-    """float: Number of steps completed"""
-
-    _get_time: GetTimeCallable
-    """Callable to get the current time."""
-
-    finished_time: Optional[float] = None
-    """float: Time task was finished."""
-
-    visible: bool = True
-    """bool: Indicates if this task is visible in the progress display."""
-
-    fields: Dict[str, Any] = field(default_factory=dict)
-    """dict: Arbitrary fields passed in via Progress.update."""
-
-    start_time: Optional[float] = field(default=None, init=False, repr=False)
-    """Optional[float]: Time this task was started, or None if not started."""
-
-    stop_time: Optional[float] = field(default=None, init=False, repr=False)
-    """Optional[float]: Time this task was stopped, or None if not stopped."""
-
-    finished_speed: Optional[float] = None
-    """Optional[float]: The last speed for a finished task."""
-
-    _progress: Deque[ProgressSample] = field(
-        default_factory=lambda: deque(maxlen=1000), init=False, repr=False
-    )
-
-    _lock: RLock = field(repr=False, default_factory=RLock)
-    """Thread lock."""
-
-    def get_time(self) -> float:
-        """float: Get the current time, in seconds."""
-        return self._get_time()
-
-    @property
-    def started(self) -> bool:
-        """bool: Check if the task as started."""
-        return self.start_time is not None
-
-    @property
-    def remaining(self) -> Optional[float]:
-        """Optional[float]: Get the number of steps remaining, if a non-None total was set."""
-        if self.total is None:
-            return None
-        return self.total - self.completed
-
-    @property
-    def elapsed(self) -> Optional[float]:
-        """Optional[float]: Time elapsed since task was started, or ``None`` if the task hasn't started."""
-        if self.start_time is None:
-            return None
-        if self.stop_time is not None:
-            return self.stop_time - self.start_time
-        return self.get_time() - self.start_time
-
-    @property
-    def finished(self) -> bool:
-        """Check if the task has finished."""
-        return self.finished_time is not None
-
-    @property
-    def percentage(self) -> float:
-        """float: Get progress of task as a percentage. If a None total was set, returns 0"""
-        if not self.total:
-            return 0.0
-        completed = (self.completed / self.total) * 100.0
-        completed = min(100.0, max(0.0, completed))
-        return completed
-
-    @property
-    def speed(self) -> Optional[float]:
-        """Optional[float]: Get the estimated speed in steps per second."""
-        if self.start_time is None:
-            return None
-        with self._lock:
-            progress = self._progress
-            if not progress:
-                return None
-            total_time = progress[-1].timestamp - progress[0].timestamp
-            if total_time == 0:
-                return None
-            iter_progress = iter(progress)
-            next(iter_progress)
-            total_completed = sum(sample.completed for sample in iter_progress)
-            speed = total_completed / total_time
-            return speed
-
-    @property
-    def time_remaining(self) -> Optional[float]:
-        """Optional[float]: Get estimated time to completion, or ``None`` if no data."""
-        if self.finished:
-            return 0.0
-        speed = self.speed
-        if not speed:
-            return None
-        remaining = self.remaining
-        if remaining is None:
-            return None
-        estimate = ceil(remaining / speed)
-        return estimate
-
-    def _reset(self) -> None:
-        """Reset progress."""
-        self._progress.clear()
-        self.finished_time = None
-        self.finished_speed = None
-
-
-class Progress(JupyterMixin):
-    """Renders an auto-updating progress bar(s).
-
-    Args:
-        console (Console, optional): Optional Console instance. Defaults to an internal Console instance writing to stdout.
-        auto_refresh (bool, optional): Enable auto refresh. If disabled, you will need to call `refresh()`.
-        refresh_per_second (Optional[float], optional): Number of times per second to refresh the progress information or None to use default (10). Defaults to None.
-        speed_estimate_period: (float, optional): Period (in seconds) used to calculate the speed estimate. Defaults to 30.
-        transient: (bool, optional): Clear the progress on exit. Defaults to False.
-        redirect_stdout: (bool, optional): Enable redirection of stdout, so ``print`` may be used. Defaults to True.
-        redirect_stderr: (bool, optional): Enable redirection of stderr. Defaults to True.
-        get_time: (Callable, optional): A callable that gets the current time, or None to use Console.get_time. Defaults to None.
-        disable (bool, optional): Disable progress display. Defaults to False
-        expand (bool, optional): Expand tasks table to fit width. Defaults to False.
-    """
-
-    def __init__(
-        self,
-        *columns: Union[str, ProgressColumn],
-        console: Optional[Console] = None,
-        auto_refresh: bool = True,
-        refresh_per_second: float = 10,
-        speed_estimate_period: float = 30.0,
-        transient: bool = False,
-        redirect_stdout: bool = True,
-        redirect_stderr: bool = True,
-        get_time: Optional[GetTimeCallable] = None,
-        disable: bool = False,
-        expand: bool = False,
-    ) -> None:
-        assert refresh_per_second > 0, "refresh_per_second must be > 0"
-        self._lock = RLock()
-        self.columns = columns or self.get_default_columns()
-        self.speed_estimate_period = speed_estimate_period
-
-        self.disable = disable
-        self.expand = expand
-        self._tasks: Dict[TaskID, Task] = {}
-        self._task_index: TaskID = TaskID(0)
-        self.live = Live(
-            console=console or get_console(),
-            auto_refresh=auto_refresh,
-            refresh_per_second=refresh_per_second,
-            transient=transient,
-            redirect_stdout=redirect_stdout,
-            redirect_stderr=redirect_stderr,
-            get_renderable=self.get_renderable,
-        )
-        self.get_time = get_time or self.console.get_time
-        self.print = self.console.print
-        self.log = self.console.log
-
-    @classmethod
-    def get_default_columns(cls) -> Tuple[ProgressColumn, ...]:
-        """Get the default columns used for a new Progress instance:
-           - a text column for the description (TextColumn)
-           - the bar itself (BarColumn)
-           - a text column showing completion percentage (TextColumn)
-           - an estimated-time-remaining column (TimeRemainingColumn)
-        If the Progress instance is created without passing a columns argument,
-        the default columns defined here will be used.
-
-        You can also create a Progress instance using custom columns before
-        and/or after the defaults, as in this example:
-
-            progress = Progress(
-                SpinnerColumn(),
-                *Progress.get_default_columns(),
-                "Elapsed:",
-                TimeElapsedColumn(),
-            )
-
-        This code shows the creation of a Progress display, containing
-        a spinner to the left, the default columns, and a labeled elapsed
-        time column.
-        """
-        return (
-            TextColumn("[progress.description]{task.description}"),
-            BarColumn(),
-            TaskProgressColumn(),
-            TimeRemainingColumn(),
-        )
-
-    @property
-    def console(self) -> Console:
-        return self.live.console
-
-    @property
-    def tasks(self) -> List[Task]:
-        """Get a list of Task instances."""
-        with self._lock:
-            return list(self._tasks.values())
-
-    @property
-    def task_ids(self) -> List[TaskID]:
-        """A list of task IDs."""
-        with self._lock:
-            return list(self._tasks.keys())
-
-    @property
-    def finished(self) -> bool:
-        """Check if all tasks have been completed."""
-        with self._lock:
-            if not self._tasks:
-                return True
-            return all(task.finished for task in self._tasks.values())
-
-    def start(self) -> None:
-        """Start the progress display."""
-        if not self.disable:
-            self.live.start(refresh=True)
-
-    def stop(self) -> None:
-        """Stop the progress display."""
-        self.live.stop()
-        if not self.console.is_interactive and not self.console.is_jupyter:
-            self.console.print()
-
-    def __enter__(self) -> Self:
-        self.start()
-        return self
-
-    def __exit__(
-        self,
-        exc_type: Optional[Type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
-    ) -> None:
-        self.stop()
-
-    def track(
-        self,
-        sequence: Union[Iterable[ProgressType], Sequence[ProgressType]],
-        total: Optional[float] = None,
-        completed: int = 0,
-        task_id: Optional[TaskID] = None,
-        description: str = "Working...",
-        update_period: float = 0.1,
-    ) -> Iterable[ProgressType]:
-        """Track progress by iterating over a sequence.
-
-        Args:
-            sequence (Sequence[ProgressType]): A sequence of values you want to iterate over and track progress.
-            total: (float, optional): Total number of steps. Default is len(sequence).
-            completed (int, optional): Number of steps completed so far. Defaults to 0.
-            task_id: (TaskID): Task to track. Default is new task.
-            description: (str, optional): Description of task, if new task is created.
-            update_period (float, optional): Minimum time (in seconds) between calls to update(). Defaults to 0.1.
-
-        Returns:
-            Iterable[ProgressType]: An iterable of values taken from the provided sequence.
-        """
-        if total is None:
-            total = float(length_hint(sequence)) or None
-
-        if task_id is None:
-            task_id = self.add_task(description, total=total, completed=completed)
-        else:
-            self.update(task_id, total=total, completed=completed)
-
-        if self.live.auto_refresh:
-            with _TrackThread(self, task_id, update_period) as track_thread:
-                for value in sequence:
-                    yield value
-                    track_thread.completed += 1
-        else:
-            advance = self.advance
-            refresh = self.refresh
-            for value in sequence:
-                yield value
-                advance(task_id, 1)
-                refresh()
-
-    def wrap_file(
-        self,
-        file: BinaryIO,
-        total: Optional[int] = None,
-        *,
-        task_id: Optional[TaskID] = None,
-        description: str = "Reading...",
-    ) -> BinaryIO:
-        """Track progress file reading from a binary file.
-
-        Args:
-            file (BinaryIO): A file-like object opened in binary mode.
-            total (int, optional): Total number of bytes to read. This must be provided unless a task with a total is also given.
-            task_id (TaskID): Task to track. Default is new task.
-            description (str, optional): Description of task, if new task is created.
-
-        Returns:
-            BinaryIO: A readable file-like object in binary mode.
-
-        Raises:
-            ValueError: When no total value can be extracted from the arguments or the task.
-        """
-        # attempt to recover the total from the task
-        total_bytes: Optional[float] = None
-        if total is not None:
-            total_bytes = total
-        elif task_id is not None:
-            with self._lock:
-                total_bytes = self._tasks[task_id].total
-        if total_bytes is None:
-            raise ValueError(
-                f"unable to get the total number of bytes, please specify 'total'"
-            )
-
-        # update total of task or create new task
-        if task_id is None:
-            task_id = self.add_task(description, total=total_bytes)
-        else:
-            self.update(task_id, total=total_bytes)
-
-        return _Reader(file, self, task_id, close_handle=False)
-
-    @typing.overload
-    def open(
-        self,
-        file: Union[str, "PathLike[str]", bytes],
-        mode: Literal["rb"],
-        buffering: int = -1,
-        encoding: Optional[str] = None,
-        errors: Optional[str] = None,
-        newline: Optional[str] = None,
-        *,
-        total: Optional[int] = None,
-        task_id: Optional[TaskID] = None,
-        description: str = "Reading...",
-    ) -> BinaryIO:
-        pass
-
-    @typing.overload
-    def open(
-        self,
-        file: Union[str, "PathLike[str]", bytes],
-        mode: Union[Literal["r"], Literal["rt"]],
-        buffering: int = -1,
-        encoding: Optional[str] = None,
-        errors: Optional[str] = None,
-        newline: Optional[str] = None,
-        *,
-        total: Optional[int] = None,
-        task_id: Optional[TaskID] = None,
-        description: str = "Reading...",
-    ) -> TextIO:
-        pass
-
-    def open(
-        self,
-        file: Union[str, "PathLike[str]", bytes],
-        mode: Union[Literal["rb"], Literal["rt"], Literal["r"]] = "r",
-        buffering: int = -1,
-        encoding: Optional[str] = None,
-        errors: Optional[str] = None,
-        newline: Optional[str] = None,
-        *,
-        total: Optional[int] = None,
-        task_id: Optional[TaskID] = None,
-        description: str = "Reading...",
-    ) -> Union[BinaryIO, TextIO]:
-        """Track progress while reading from a binary file.
-
-        Args:
-            path (Union[str, PathLike[str]]): The path to the file to read.
-            mode (str): The mode to use to open the file. Only supports "r", "rb" or "rt".
-            buffering (int): The buffering strategy to use, see :func:`io.open`.
-            encoding (str, optional): The encoding to use when reading in text mode, see :func:`io.open`.
-            errors (str, optional): The error handling strategy for decoding errors, see :func:`io.open`.
-            newline (str, optional): The strategy for handling newlines in text mode, see :func:`io.open`.
-            total (int, optional): Total number of bytes to read. If none given, os.stat(path).st_size is used.
-            task_id (TaskID): Task to track. Default is new task.
-            description (str, optional): Description of task, if new task is created.
-
-        Returns:
-            BinaryIO: A readable file-like object in binary mode.
-
-        Raises:
-            ValueError: When an invalid mode is given.
-        """
-        # normalize the mode (always rb, rt)
-        _mode = "".join(sorted(mode, reverse=False))
-        if _mode not in ("br", "rt", "r"):
-            raise ValueError(f"invalid mode {mode!r}")
-
-        # patch buffering to provide the same behaviour as the builtin `open`
-        line_buffering = buffering == 1
-        if _mode == "br" and buffering == 1:
-            warnings.warn(
-                "line buffering (buffering=1) isn't supported in binary mode, the default buffer size will be used",
-                RuntimeWarning,
-            )
-            buffering = -1
-        elif _mode in ("rt", "r"):
-            if buffering == 0:
-                raise ValueError("can't have unbuffered text I/O")
-            elif buffering == 1:
-                buffering = -1
-
-        # attempt to get the total with `os.stat`
-        if total is None:
-            total = stat(file).st_size
-
-        # update total of task or create new task
-        if task_id is None:
-            task_id = self.add_task(description, total=total)
-        else:
-            self.update(task_id, total=total)
-
-        # open the file in binary mode,
-        handle = io.open(file, "rb", buffering=buffering)
-        reader = _Reader(handle, self, task_id, close_handle=True)
-
-        # wrap the reader in a `TextIOWrapper` if text mode
-        if mode in ("r", "rt"):
-            return io.TextIOWrapper(
-                reader,
-                encoding=encoding,
-                errors=errors,
-                newline=newline,
-                line_buffering=line_buffering,
-            )
-
-        return reader
-
-    def start_task(self, task_id: TaskID) -> None:
-        """Start a task.
-
-        Starts a task (used when calculating elapsed time). You may need to call this manually,
-        if you called ``add_task`` with ``start=False``.
-
-        Args:
-            task_id (TaskID): ID of task.
-        """
-        with self._lock:
-            task = self._tasks[task_id]
-            if task.start_time is None:
-                task.start_time = self.get_time()
-
-    def stop_task(self, task_id: TaskID) -> None:
-        """Stop a task.
-
-        This will freeze the elapsed time on the task.
-
-        Args:
-            task_id (TaskID): ID of task.
-        """
-        with self._lock:
-            task = self._tasks[task_id]
-            current_time = self.get_time()
-            if task.start_time is None:
-                task.start_time = current_time
-            task.stop_time = current_time
-
-    def update(
-        self,
-        task_id: TaskID,
-        *,
-        total: Optional[float] = None,
-        completed: Optional[float] = None,
-        advance: Optional[float] = None,
-        description: Optional[str] = None,
-        visible: Optional[bool] = None,
-        refresh: bool = False,
-        **fields: Any,
-    ) -> None:
-        """Update information associated with a task.
-
-        Args:
-            task_id (TaskID): Task id (returned by add_task).
-            total (float, optional): Updates task.total if not None.
-            completed (float, optional): Updates task.completed if not None.
-            advance (float, optional): Add a value to task.completed if not None.
-            description (str, optional): Change task description if not None.
-            visible (bool, optional): Set visible flag if not None.
-            refresh (bool): Force a refresh of progress information. Default is False.
-            **fields (Any): Additional data fields required for rendering.
-        """
-        with self._lock:
-            task = self._tasks[task_id]
-            completed_start = task.completed
-
-            if total is not None and total != task.total:
-                task.total = total
-                task._reset()
-            if advance is not None:
-                task.completed += advance
-            if completed is not None:
-                task.completed = completed
-            if description is not None:
-                task.description = description
-            if visible is not None:
-                task.visible = visible
-            task.fields.update(fields)
-            update_completed = task.completed - completed_start
-
-            current_time = self.get_time()
-            old_sample_time = current_time - self.speed_estimate_period
-            _progress = task._progress
-
-            popleft = _progress.popleft
-            while _progress and _progress[0].timestamp < old_sample_time:
-                popleft()
-            if update_completed > 0:
-                _progress.append(ProgressSample(current_time, update_completed))
-            if (
-                task.total is not None
-                and task.completed >= task.total
-                and task.finished_time is None
-            ):
-                task.finished_time = task.elapsed
-
-        if refresh:
-            self.refresh()
-
-    def reset(
-        self,
-        task_id: TaskID,
-        *,
-        start: bool = True,
-        total: Optional[float] = None,
-        completed: int = 0,
-        visible: Optional[bool] = None,
-        description: Optional[str] = None,
-        **fields: Any,
-    ) -> None:
-        """Reset a task so completed is 0 and the clock is reset.
-
-        Args:
-            task_id (TaskID): ID of task.
-            start (bool, optional): Start the task after reset. Defaults to True.
-            total (float, optional): New total steps in task, or None to use current total. Defaults to None.
-            completed (int, optional): Number of steps completed. Defaults to 0.
-            visible (bool, optional): Enable display of the task. Defaults to True.
-            description (str, optional): Change task description if not None. Defaults to None.
-            **fields (str): Additional data fields required for rendering.
-        """
-        current_time = self.get_time()
-        with self._lock:
-            task = self._tasks[task_id]
-            task._reset()
-            task.start_time = current_time if start else None
-            if total is not None:
-                task.total = total
-            task.completed = completed
-            if visible is not None:
-                task.visible = visible
-            if fields:
-                task.fields = fields
-            if description is not None:
-                task.description = description
-            task.finished_time = None
-        self.refresh()
-
-    def advance(self, task_id: TaskID, advance: float = 1) -> None:
-        """Advance task by a number of steps.
-
-        Args:
-            task_id (TaskID): ID of task.
-            advance (float): Number of steps to advance. Default is 1.
-        """
-        current_time = self.get_time()
-        with self._lock:
-            task = self._tasks[task_id]
-            completed_start = task.completed
-            task.completed += advance
-            update_completed = task.completed - completed_start
-            old_sample_time = current_time - self.speed_estimate_period
-            _progress = task._progress
-
-            popleft = _progress.popleft
-            while _progress and _progress[0].timestamp < old_sample_time:
-                popleft()
-            while len(_progress) > 1000:
-                popleft()
-            _progress.append(ProgressSample(current_time, update_completed))
-            if (
-                task.total is not None
-                and task.completed >= task.total
-                and task.finished_time is None
-            ):
-                task.finished_time = task.elapsed
-                task.finished_speed = task.speed
-
-    def refresh(self) -> None:
-        """Refresh (render) the progress information."""
-        if not self.disable and self.live.is_started:
-            self.live.refresh()
-
-    def get_renderable(self) -> RenderableType:
-        """Get a renderable for the progress display."""
-        renderable = Group(*self.get_renderables())
-        return renderable
-
-    def get_renderables(self) -> Iterable[RenderableType]:
-        """Get a number of renderables for the progress display."""
-        table = self.make_tasks_table(self.tasks)
-        yield table
-
-    def make_tasks_table(self, tasks: Iterable[Task]) -> Table:
-        """Get a table to render the Progress display.
-
-        Args:
-            tasks (Iterable[Task]): An iterable of Task instances, one per row of the table.
-
-        Returns:
-            Table: A table instance.
-        """
-        table_columns = (
-            (
-                Column(no_wrap=True)
-                if isinstance(_column, str)
-                else _column.get_table_column().copy()
-            )
-            for _column in self.columns
-        )
-        table = Table.grid(*table_columns, padding=(0, 1), expand=self.expand)
-
-        for task in tasks:
-            if task.visible:
-                table.add_row(
-                    *(
-                        (
-                            column.format(task=task)
-                            if isinstance(column, str)
-                            else column(task)
-                        )
-                        for column in self.columns
-                    )
-                )
-        return table
-
-    def __rich__(self) -> RenderableType:
-        """Makes the Progress class itself renderable."""
-        with self._lock:
-            return self.get_renderable()
-
-    def add_task(
-        self,
-        description: str,
-        start: bool = True,
-        total: Optional[float] = 100.0,
-        completed: int = 0,
-        visible: bool = True,
-        **fields: Any,
-    ) -> TaskID:
-        """Add a new 'task' to the Progress display.
-
-        Args:
-            description (str): A description of the task.
-            start (bool, optional): Start the task immediately (to calculate elapsed time). If set to False,
-                you will need to call `start` manually. Defaults to True.
-            total (float, optional): Number of total steps in the progress if known.
-                Set to None to render a pulsing animation. Defaults to 100.
-            completed (int, optional): Number of steps completed so far. Defaults to 0.
-            visible (bool, optional): Enable display of the task. Defaults to True.
-            **fields (str): Additional data fields required for rendering.
-
-        Returns:
-            TaskID: An ID you can use when calling `update`.
-        """
-        with self._lock:
-            task = Task(
-                self._task_index,
-                description,
-                total,
-                completed,
-                visible=visible,
-                fields=fields,
-                _get_time=self.get_time,
-                _lock=self._lock,
-            )
-            self._tasks[self._task_index] = task
-            if start:
-                self.start_task(self._task_index)
-            new_task_index = self._task_index
-            self._task_index = TaskID(int(self._task_index) + 1)
-        self.refresh()
-        return new_task_index
-
-    def remove_task(self, task_id: TaskID) -> None:
-        """Delete a task if it exists.
-
-        Args:
-            task_id (TaskID): A task ID.
-
-        """
-        with self._lock:
-            del self._tasks[task_id]
-
-
-if __name__ == "__main__":  # pragma: no coverage
-    import random
-    import time
-
-    from .panel import Panel
-    from .rule import Rule
-    from .syntax import Syntax
-    from .table import Table
-
-    syntax = Syntax(
-        '''def loop_last(values: Iterable[T]) -> Iterable[Tuple[bool, T]]:
-    """Iterate and generate a tuple with a flag for last value."""
-    iter_values = iter(values)
-    try:
-        previous_value = next(iter_values)
-    except StopIteration:
-        return
-    for value in iter_values:
-        yield False, previous_value
-        previous_value = value
-    yield True, previous_value''',
-        "python",
-        line_numbers=True,
-    )
-
-    table = Table("foo", "bar", "baz")
-    table.add_row("1", "2", "3")
-
-    progress_renderables = [
-        "Text may be printed while the progress bars are rendering.",
-        Panel("In fact, [i]any[/i] renderable will work"),
-        "Such as [magenta]tables[/]...",
-        table,
-        "Pretty printed structures...",
-        {"type": "example", "text": "Pretty printed"},
-        "Syntax...",
-        syntax,
-        Rule("Give it a try!"),
-    ]
-
-    from itertools import cycle
-
-    examples = cycle(progress_renderables)
-
-    console = Console(record=True)
-
-    with Progress(
-        SpinnerColumn(),
-        *Progress.get_default_columns(),
-        TimeElapsedColumn(),
-        console=console,
-        transient=False,
-    ) as progress:
-        task1 = progress.add_task("[red]Downloading", total=1000)
-        task2 = progress.add_task("[green]Processing", total=1000)
-        task3 = progress.add_task("[yellow]Thinking", total=None)
-
-        while not progress.finished:
-            progress.update(task1, advance=0.5)
-            progress.update(task2, advance=0.3)
-            time.sleep(0.01)
-            if random.randint(0, 100) < 1:
-                progress.log(next(examples))
diff --git a/venv/lib/python3.10/site-packages/rich/progress_bar.py b/venv/lib/python3.10/site-packages/rich/progress_bar.py
deleted file mode 100644
index 41794f767877ad6330541c076801c8e76e324bbe..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/progress_bar.py
+++ /dev/null
@@ -1,223 +0,0 @@
-import math
-from functools import lru_cache
-from time import monotonic
-from typing import Iterable, List, Optional
-
-from .color import Color, blend_rgb
-from .color_triplet import ColorTriplet
-from .console import Console, ConsoleOptions, RenderResult
-from .jupyter import JupyterMixin
-from .measure import Measurement
-from .segment import Segment
-from .style import Style, StyleType
-
-# Number of characters before 'pulse' animation repeats
-PULSE_SIZE = 20
-
-
-class ProgressBar(JupyterMixin):
-    """Renders a (progress) bar. Used by rich.progress.
-
-    Args:
-        total (float, optional): Number of steps in the bar. Defaults to 100. Set to None to render a pulsing animation.
-        completed (float, optional): Number of steps completed. Defaults to 0.
-        width (int, optional): Width of the bar, or ``None`` for maximum width. Defaults to None.
-        pulse (bool, optional): Enable pulse effect. Defaults to False. Will pulse if a None total was passed.
-        style (StyleType, optional): Style for the bar background. Defaults to "bar.back".
-        complete_style (StyleType, optional): Style for the completed bar. Defaults to "bar.complete".
-        finished_style (StyleType, optional): Style for a finished bar. Defaults to "bar.finished".
-        pulse_style (StyleType, optional): Style for pulsing bars. Defaults to "bar.pulse".
-        animation_time (Optional[float], optional): Time in seconds to use for animation, or None to use system time.
-    """
-
-    def __init__(
-        self,
-        total: Optional[float] = 100.0,
-        completed: float = 0,
-        width: Optional[int] = None,
-        pulse: bool = False,
-        style: StyleType = "bar.back",
-        complete_style: StyleType = "bar.complete",
-        finished_style: StyleType = "bar.finished",
-        pulse_style: StyleType = "bar.pulse",
-        animation_time: Optional[float] = None,
-    ):
-        self.total = total
-        self.completed = completed
-        self.width = width
-        self.pulse = pulse
-        self.style = style
-        self.complete_style = complete_style
-        self.finished_style = finished_style
-        self.pulse_style = pulse_style
-        self.animation_time = animation_time
-
-        self._pulse_segments: Optional[List[Segment]] = None
-
-    def __repr__(self) -> str:
-        return f""
-
-    @property
-    def percentage_completed(self) -> Optional[float]:
-        """Calculate percentage complete."""
-        if self.total is None:
-            return None
-        completed = (self.completed / self.total) * 100.0
-        completed = min(100, max(0.0, completed))
-        return completed
-
-    @lru_cache(maxsize=16)
-    def _get_pulse_segments(
-        self,
-        fore_style: Style,
-        back_style: Style,
-        color_system: str,
-        no_color: bool,
-        ascii: bool = False,
-    ) -> List[Segment]:
-        """Get a list of segments to render a pulse animation.
-
-        Returns:
-            List[Segment]: A list of segments, one segment per character.
-        """
-        bar = "-" if ascii else "━"
-        segments: List[Segment] = []
-        if color_system not in ("standard", "eight_bit", "truecolor") or no_color:
-            segments += [Segment(bar, fore_style)] * (PULSE_SIZE // 2)
-            segments += [Segment(" " if no_color else bar, back_style)] * (
-                PULSE_SIZE - (PULSE_SIZE // 2)
-            )
-            return segments
-
-        append = segments.append
-        fore_color = (
-            fore_style.color.get_truecolor()
-            if fore_style.color
-            else ColorTriplet(255, 0, 255)
-        )
-        back_color = (
-            back_style.color.get_truecolor()
-            if back_style.color
-            else ColorTriplet(0, 0, 0)
-        )
-        cos = math.cos
-        pi = math.pi
-        _Segment = Segment
-        _Style = Style
-        from_triplet = Color.from_triplet
-
-        for index in range(PULSE_SIZE):
-            position = index / PULSE_SIZE
-            fade = 0.5 + cos(position * pi * 2) / 2.0
-            color = blend_rgb(fore_color, back_color, cross_fade=fade)
-            append(_Segment(bar, _Style(color=from_triplet(color))))
-        return segments
-
-    def update(self, completed: float, total: Optional[float] = None) -> None:
-        """Update progress with new values.
-
-        Args:
-            completed (float): Number of steps completed.
-            total (float, optional): Total number of steps, or ``None`` to not change. Defaults to None.
-        """
-        self.completed = completed
-        self.total = total if total is not None else self.total
-
-    def _render_pulse(
-        self, console: Console, width: int, ascii: bool = False
-    ) -> Iterable[Segment]:
-        """Renders the pulse animation.
-
-        Args:
-            console (Console): Console instance.
-            width (int): Width in characters of pulse animation.
-
-        Returns:
-            RenderResult: [description]
-
-        Yields:
-            Iterator[Segment]: Segments to render pulse
-        """
-        fore_style = console.get_style(self.pulse_style, default="white")
-        back_style = console.get_style(self.style, default="black")
-
-        pulse_segments = self._get_pulse_segments(
-            fore_style, back_style, console.color_system, console.no_color, ascii=ascii
-        )
-        segment_count = len(pulse_segments)
-        current_time = (
-            monotonic() if self.animation_time is None else self.animation_time
-        )
-        segments = pulse_segments * (int(width / segment_count) + 2)
-        offset = int(-current_time * 15) % segment_count
-        segments = segments[offset : offset + width]
-        yield from segments
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        width = min(self.width or options.max_width, options.max_width)
-        ascii = options.legacy_windows or options.ascii_only
-        should_pulse = self.pulse or self.total is None
-        if should_pulse:
-            yield from self._render_pulse(console, width, ascii=ascii)
-            return
-
-        completed: Optional[float] = (
-            min(self.total, max(0, self.completed)) if self.total is not None else None
-        )
-
-        bar = "-" if ascii else "━"
-        half_bar_right = " " if ascii else "╸"
-        half_bar_left = " " if ascii else "╺"
-        complete_halves = (
-            int(width * 2 * completed / self.total)
-            if self.total and completed is not None
-            else width * 2
-        )
-        bar_count = complete_halves // 2
-        half_bar_count = complete_halves % 2
-        style = console.get_style(self.style)
-        is_finished = self.total is None or self.completed >= self.total
-        complete_style = console.get_style(
-            self.finished_style if is_finished else self.complete_style
-        )
-        _Segment = Segment
-        if bar_count:
-            yield _Segment(bar * bar_count, complete_style)
-        if half_bar_count:
-            yield _Segment(half_bar_right * half_bar_count, complete_style)
-
-        if not console.no_color:
-            remaining_bars = width - bar_count - half_bar_count
-            if remaining_bars and console.color_system is not None:
-                if not half_bar_count and bar_count:
-                    yield _Segment(half_bar_left, style)
-                    remaining_bars -= 1
-                if remaining_bars:
-                    yield _Segment(bar * remaining_bars, style)
-
-    def __rich_measure__(
-        self, console: Console, options: ConsoleOptions
-    ) -> Measurement:
-        return (
-            Measurement(self.width, self.width)
-            if self.width is not None
-            else Measurement(4, options.max_width)
-        )
-
-
-if __name__ == "__main__":  # pragma: no cover
-    console = Console()
-    bar = ProgressBar(width=50, total=100)
-
-    import time
-
-    console.show_cursor(False)
-    for n in range(0, 101, 1):
-        bar.update(n)
-        console.print(bar)
-        console.file.write("\r")
-        time.sleep(0.05)
-    console.show_cursor(True)
-    console.print()
diff --git a/venv/lib/python3.10/site-packages/rich/prompt.py b/venv/lib/python3.10/site-packages/rich/prompt.py
deleted file mode 100644
index c7cf25bae8554bf53173ed1eb78a3206f6e3c05d..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/prompt.py
+++ /dev/null
@@ -1,400 +0,0 @@
-from typing import Any, Generic, List, Optional, TextIO, TypeVar, Union, overload
-
-from . import get_console
-from .console import Console
-from .text import Text, TextType
-
-PromptType = TypeVar("PromptType")
-DefaultType = TypeVar("DefaultType")
-
-
-class PromptError(Exception):
-    """Exception base class for prompt related errors."""
-
-
-class InvalidResponse(PromptError):
-    """Exception to indicate a response was invalid. Raise this within process_response() to indicate an error
-    and provide an error message.
-
-    Args:
-        message (Union[str, Text]): Error message.
-    """
-
-    def __init__(self, message: TextType) -> None:
-        self.message = message
-
-    def __rich__(self) -> TextType:
-        return self.message
-
-
-class PromptBase(Generic[PromptType]):
-    """Ask the user for input until a valid response is received. This is the base class, see one of
-    the concrete classes for examples.
-
-    Args:
-        prompt (TextType, optional): Prompt text. Defaults to "".
-        console (Console, optional): A Console instance or None to use global console. Defaults to None.
-        password (bool, optional): Enable password input. Defaults to False.
-        choices (List[str], optional): A list of valid choices. Defaults to None.
-        case_sensitive (bool, optional): Matching of choices should be case-sensitive. Defaults to True.
-        show_default (bool, optional): Show default in prompt. Defaults to True.
-        show_choices (bool, optional): Show choices in prompt. Defaults to True.
-    """
-
-    response_type: type = str
-
-    validate_error_message = "[prompt.invalid]Please enter a valid value"
-    illegal_choice_message = (
-        "[prompt.invalid.choice]Please select one of the available options"
-    )
-    prompt_suffix = ": "
-
-    choices: Optional[List[str]] = None
-
-    def __init__(
-        self,
-        prompt: TextType = "",
-        *,
-        console: Optional[Console] = None,
-        password: bool = False,
-        choices: Optional[List[str]] = None,
-        case_sensitive: bool = True,
-        show_default: bool = True,
-        show_choices: bool = True,
-    ) -> None:
-        self.console = console or get_console()
-        self.prompt = (
-            Text.from_markup(prompt, style="prompt")
-            if isinstance(prompt, str)
-            else prompt
-        )
-        self.password = password
-        if choices is not None:
-            self.choices = choices
-        self.case_sensitive = case_sensitive
-        self.show_default = show_default
-        self.show_choices = show_choices
-
-    @classmethod
-    @overload
-    def ask(
-        cls,
-        prompt: TextType = "",
-        *,
-        console: Optional[Console] = None,
-        password: bool = False,
-        choices: Optional[List[str]] = None,
-        case_sensitive: bool = True,
-        show_default: bool = True,
-        show_choices: bool = True,
-        default: DefaultType,
-        stream: Optional[TextIO] = None,
-    ) -> Union[DefaultType, PromptType]:
-        ...
-
-    @classmethod
-    @overload
-    def ask(
-        cls,
-        prompt: TextType = "",
-        *,
-        console: Optional[Console] = None,
-        password: bool = False,
-        choices: Optional[List[str]] = None,
-        case_sensitive: bool = True,
-        show_default: bool = True,
-        show_choices: bool = True,
-        stream: Optional[TextIO] = None,
-    ) -> PromptType:
-        ...
-
-    @classmethod
-    def ask(
-        cls,
-        prompt: TextType = "",
-        *,
-        console: Optional[Console] = None,
-        password: bool = False,
-        choices: Optional[List[str]] = None,
-        case_sensitive: bool = True,
-        show_default: bool = True,
-        show_choices: bool = True,
-        default: Any = ...,
-        stream: Optional[TextIO] = None,
-    ) -> Any:
-        """Shortcut to construct and run a prompt loop and return the result.
-
-        Example:
-            >>> filename = Prompt.ask("Enter a filename")
-
-        Args:
-            prompt (TextType, optional): Prompt text. Defaults to "".
-            console (Console, optional): A Console instance or None to use global console. Defaults to None.
-            password (bool, optional): Enable password input. Defaults to False.
-            choices (List[str], optional): A list of valid choices. Defaults to None.
-            case_sensitive (bool, optional): Matching of choices should be case-sensitive. Defaults to True.
-            show_default (bool, optional): Show default in prompt. Defaults to True.
-            show_choices (bool, optional): Show choices in prompt. Defaults to True.
-            stream (TextIO, optional): Optional text file open for reading to get input. Defaults to None.
-        """
-        _prompt = cls(
-            prompt,
-            console=console,
-            password=password,
-            choices=choices,
-            case_sensitive=case_sensitive,
-            show_default=show_default,
-            show_choices=show_choices,
-        )
-        return _prompt(default=default, stream=stream)
-
-    def render_default(self, default: DefaultType) -> Text:
-        """Turn the supplied default in to a Text instance.
-
-        Args:
-            default (DefaultType): Default value.
-
-        Returns:
-            Text: Text containing rendering of default value.
-        """
-        return Text(f"({default})", "prompt.default")
-
-    def make_prompt(self, default: DefaultType) -> Text:
-        """Make prompt text.
-
-        Args:
-            default (DefaultType): Default value.
-
-        Returns:
-            Text: Text to display in prompt.
-        """
-        prompt = self.prompt.copy()
-        prompt.end = ""
-
-        if self.show_choices and self.choices:
-            _choices = "/".join(self.choices)
-            choices = f"[{_choices}]"
-            prompt.append(" ")
-            prompt.append(choices, "prompt.choices")
-
-        if (
-            default != ...
-            and self.show_default
-            and isinstance(default, (str, self.response_type))
-        ):
-            prompt.append(" ")
-            _default = self.render_default(default)
-            prompt.append(_default)
-
-        prompt.append(self.prompt_suffix)
-
-        return prompt
-
-    @classmethod
-    def get_input(
-        cls,
-        console: Console,
-        prompt: TextType,
-        password: bool,
-        stream: Optional[TextIO] = None,
-    ) -> str:
-        """Get input from user.
-
-        Args:
-            console (Console): Console instance.
-            prompt (TextType): Prompt text.
-            password (bool): Enable password entry.
-
-        Returns:
-            str: String from user.
-        """
-        return console.input(prompt, password=password, stream=stream)
-
-    def check_choice(self, value: str) -> bool:
-        """Check value is in the list of valid choices.
-
-        Args:
-            value (str): Value entered by user.
-
-        Returns:
-            bool: True if choice was valid, otherwise False.
-        """
-        assert self.choices is not None
-        if self.case_sensitive:
-            return value.strip() in self.choices
-        return value.strip().lower() in [choice.lower() for choice in self.choices]
-
-    def process_response(self, value: str) -> PromptType:
-        """Process response from user, convert to prompt type.
-
-        Args:
-            value (str): String typed by user.
-
-        Raises:
-            InvalidResponse: If ``value`` is invalid.
-
-        Returns:
-            PromptType: The value to be returned from ask method.
-        """
-        value = value.strip()
-        try:
-            return_value: PromptType = self.response_type(value)
-        except ValueError:
-            raise InvalidResponse(self.validate_error_message)
-
-        if self.choices is not None:
-            if not self.check_choice(value):
-                raise InvalidResponse(self.illegal_choice_message)
-
-            if not self.case_sensitive:
-                # return the original choice, not the lower case version
-                return_value = self.response_type(
-                    self.choices[
-                        [choice.lower() for choice in self.choices].index(value.lower())
-                    ]
-                )
-        return return_value
-
-    def on_validate_error(self, value: str, error: InvalidResponse) -> None:
-        """Called to handle validation error.
-
-        Args:
-            value (str): String entered by user.
-            error (InvalidResponse): Exception instance the initiated the error.
-        """
-        self.console.print(error)
-
-    def pre_prompt(self) -> None:
-        """Hook to display something before the prompt."""
-
-    @overload
-    def __call__(self, *, stream: Optional[TextIO] = None) -> PromptType:
-        ...
-
-    @overload
-    def __call__(
-        self, *, default: DefaultType, stream: Optional[TextIO] = None
-    ) -> Union[PromptType, DefaultType]:
-        ...
-
-    def __call__(self, *, default: Any = ..., stream: Optional[TextIO] = None) -> Any:
-        """Run the prompt loop.
-
-        Args:
-            default (Any, optional): Optional default value.
-
-        Returns:
-            PromptType: Processed value.
-        """
-        while True:
-            self.pre_prompt()
-            prompt = self.make_prompt(default)
-            value = self.get_input(self.console, prompt, self.password, stream=stream)
-            if value == "" and default != ...:
-                return default
-            try:
-                return_value = self.process_response(value)
-            except InvalidResponse as error:
-                self.on_validate_error(value, error)
-                continue
-            else:
-                return return_value
-
-
-class Prompt(PromptBase[str]):
-    """A prompt that returns a str.
-
-    Example:
-        >>> name = Prompt.ask("Enter your name")
-
-
-    """
-
-    response_type = str
-
-
-class IntPrompt(PromptBase[int]):
-    """A prompt that returns an integer.
-
-    Example:
-        >>> burrito_count = IntPrompt.ask("How many burritos do you want to order")
-
-    """
-
-    response_type = int
-    validate_error_message = "[prompt.invalid]Please enter a valid integer number"
-
-
-class FloatPrompt(PromptBase[float]):
-    """A prompt that returns a float.
-
-    Example:
-        >>> temperature = FloatPrompt.ask("Enter desired temperature")
-
-    """
-
-    response_type = float
-    validate_error_message = "[prompt.invalid]Please enter a number"
-
-
-class Confirm(PromptBase[bool]):
-    """A yes / no confirmation prompt.
-
-    Example:
-        >>> if Confirm.ask("Continue"):
-                run_job()
-
-    """
-
-    response_type = bool
-    validate_error_message = "[prompt.invalid]Please enter Y or N"
-    choices: List[str] = ["y", "n"]
-
-    def render_default(self, default: DefaultType) -> Text:
-        """Render the default as (y) or (n) rather than True/False."""
-        yes, no = self.choices
-        return Text(f"({yes})" if default else f"({no})", style="prompt.default")
-
-    def process_response(self, value: str) -> bool:
-        """Convert choices to a bool."""
-        value = value.strip().lower()
-        if value not in self.choices:
-            raise InvalidResponse(self.validate_error_message)
-        return value == self.choices[0]
-
-
-if __name__ == "__main__":  # pragma: no cover
-    from rich import print
-
-    if Confirm.ask("Run [i]prompt[/i] tests?", default=True):
-        while True:
-            result = IntPrompt.ask(
-                ":rocket: Enter a number between [b]1[/b] and [b]10[/b]", default=5
-            )
-            if result >= 1 and result <= 10:
-                break
-            print(":pile_of_poo: [prompt.invalid]Number must be between 1 and 10")
-        print(f"number={result}")
-
-        while True:
-            password = Prompt.ask(
-                "Please enter a password [cyan](must be at least 5 characters)",
-                password=True,
-            )
-            if len(password) >= 5:
-                break
-            print("[prompt.invalid]password too short")
-        print(f"password={password!r}")
-
-        fruit = Prompt.ask("Enter a fruit", choices=["apple", "orange", "pear"])
-        print(f"fruit={fruit!r}")
-
-        doggie = Prompt.ask(
-            "What's the best Dog? (Case INSENSITIVE)",
-            choices=["Border Terrier", "Collie", "Labradoodle"],
-            case_sensitive=False,
-        )
-        print(f"doggie={doggie!r}")
-
-    else:
-        print("[b]OK :loudly_crying_face:")
diff --git a/venv/lib/python3.10/site-packages/rich/protocol.py b/venv/lib/python3.10/site-packages/rich/protocol.py
deleted file mode 100644
index c6923dd7d870ec02decdb66b55889df9041574d2..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/protocol.py
+++ /dev/null
@@ -1,42 +0,0 @@
-from typing import Any, cast, Set, TYPE_CHECKING
-from inspect import isclass
-
-if TYPE_CHECKING:
-    from rich.console import RenderableType
-
-_GIBBERISH = """aihwerij235234ljsdnp34ksodfipwoe234234jlskjdf"""
-
-
-def is_renderable(check_object: Any) -> bool:
-    """Check if an object may be rendered by Rich."""
-    return (
-        isinstance(check_object, str)
-        or hasattr(check_object, "__rich__")
-        or hasattr(check_object, "__rich_console__")
-    )
-
-
-def rich_cast(renderable: object) -> "RenderableType":
-    """Cast an object to a renderable by calling __rich__ if present.
-
-    Args:
-        renderable (object): A potentially renderable object
-
-    Returns:
-        object: The result of recursively calling __rich__.
-    """
-    from rich.console import RenderableType
-
-    rich_visited_set: Set[type] = set()  # Prevent potential infinite loop
-    while hasattr(renderable, "__rich__") and not isclass(renderable):
-        # Detect object which claim to have all the attributes
-        if hasattr(renderable, _GIBBERISH):
-            return repr(renderable)
-        cast_method = getattr(renderable, "__rich__")
-        renderable = cast_method()
-        renderable_type = type(renderable)
-        if renderable_type in rich_visited_set:
-            break
-        rich_visited_set.add(renderable_type)
-
-    return cast(RenderableType, renderable)
diff --git a/venv/lib/python3.10/site-packages/rich/py.typed b/venv/lib/python3.10/site-packages/rich/py.typed
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/venv/lib/python3.10/site-packages/rich/region.py b/venv/lib/python3.10/site-packages/rich/region.py
deleted file mode 100644
index 75b3631c3879294549f1f27418859aefb63925a7..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/region.py
+++ /dev/null
@@ -1,10 +0,0 @@
-from typing import NamedTuple
-
-
-class Region(NamedTuple):
-    """Defines a rectangular region of the screen."""
-
-    x: int
-    y: int
-    width: int
-    height: int
diff --git a/venv/lib/python3.10/site-packages/rich/repr.py b/venv/lib/python3.10/site-packages/rich/repr.py
deleted file mode 100644
index 95331006eede0f7b4fa3a5e0ed8a95dc96cbec18..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/repr.py
+++ /dev/null
@@ -1,149 +0,0 @@
-import inspect
-from functools import partial
-from typing import (
-    Any,
-    Callable,
-    Iterable,
-    List,
-    Optional,
-    Tuple,
-    Type,
-    TypeVar,
-    Union,
-    overload,
-)
-
-T = TypeVar("T")
-
-
-Result = Iterable[Union[Any, Tuple[Any], Tuple[str, Any], Tuple[str, Any, Any]]]
-RichReprResult = Result
-
-
-class ReprError(Exception):
-    """An error occurred when attempting to build a repr."""
-
-
-@overload
-def auto(cls: Optional[Type[T]]) -> Type[T]:
-    ...
-
-
-@overload
-def auto(*, angular: bool = False) -> Callable[[Type[T]], Type[T]]:
-    ...
-
-
-def auto(
-    cls: Optional[Type[T]] = None, *, angular: Optional[bool] = None
-) -> Union[Type[T], Callable[[Type[T]], Type[T]]]:
-    """Class decorator to create __repr__ from __rich_repr__"""
-
-    def do_replace(cls: Type[T], angular: Optional[bool] = None) -> Type[T]:
-        def auto_repr(self: T) -> str:
-            """Create repr string from __rich_repr__"""
-            repr_str: List[str] = []
-            append = repr_str.append
-
-            angular: bool = getattr(self.__rich_repr__, "angular", False)  # type: ignore[attr-defined]
-            for arg in self.__rich_repr__():  # type: ignore[attr-defined]
-                if isinstance(arg, tuple):
-                    if len(arg) == 1:
-                        append(repr(arg[0]))
-                    else:
-                        key, value, *default = arg
-                        if key is None:
-                            append(repr(value))
-                        else:
-                            if default and default[0] == value:
-                                continue
-                            append(f"{key}={value!r}")
-                else:
-                    append(repr(arg))
-            if angular:
-                return f"<{self.__class__.__name__} {' '.join(repr_str)}>"
-            else:
-                return f"{self.__class__.__name__}({', '.join(repr_str)})"
-
-        def auto_rich_repr(self: Type[T]) -> Result:
-            """Auto generate __rich_rep__ from signature of __init__"""
-            try:
-                signature = inspect.signature(self.__init__)
-                for name, param in signature.parameters.items():
-                    if param.kind == param.POSITIONAL_ONLY:
-                        yield getattr(self, name)
-                    elif param.kind in (
-                        param.POSITIONAL_OR_KEYWORD,
-                        param.KEYWORD_ONLY,
-                    ):
-                        if param.default is param.empty:
-                            yield getattr(self, param.name)
-                        else:
-                            yield param.name, getattr(self, param.name), param.default
-            except Exception as error:
-                raise ReprError(
-                    f"Failed to auto generate __rich_repr__; {error}"
-                ) from None
-
-        if not hasattr(cls, "__rich_repr__"):
-            auto_rich_repr.__doc__ = "Build a rich repr"
-            cls.__rich_repr__ = auto_rich_repr  # type: ignore[attr-defined]
-
-        auto_repr.__doc__ = "Return repr(self)"
-        cls.__repr__ = auto_repr  # type: ignore[assignment]
-        if angular is not None:
-            cls.__rich_repr__.angular = angular  # type: ignore[attr-defined]
-        return cls
-
-    if cls is None:
-        return partial(do_replace, angular=angular)
-    else:
-        return do_replace(cls, angular=angular)
-
-
-@overload
-def rich_repr(cls: Optional[Type[T]]) -> Type[T]:
-    ...
-
-
-@overload
-def rich_repr(*, angular: bool = False) -> Callable[[Type[T]], Type[T]]:
-    ...
-
-
-def rich_repr(
-    cls: Optional[Type[T]] = None, *, angular: bool = False
-) -> Union[Type[T], Callable[[Type[T]], Type[T]]]:
-    if cls is None:
-        return auto(angular=angular)
-    else:
-        return auto(cls)
-
-
-if __name__ == "__main__":
-
-    @auto
-    class Foo:
-        def __rich_repr__(self) -> Result:
-            yield "foo"
-            yield "bar", {"shopping": ["eggs", "ham", "pineapple"]}
-            yield "buy", "hand sanitizer"
-
-    foo = Foo()
-    from rich.console import Console
-
-    console = Console()
-
-    console.rule("Standard repr")
-    console.print(foo)
-
-    console.print(foo, width=60)
-    console.print(foo, width=30)
-
-    console.rule("Angular repr")
-    Foo.__rich_repr__.angular = True  # type: ignore[attr-defined]
-
-    console.print(foo)
-
-    console.print(foo, width=60)
-    console.print(foo, width=30)
diff --git a/venv/lib/python3.10/site-packages/rich/rule.py b/venv/lib/python3.10/site-packages/rich/rule.py
deleted file mode 100644
index fb3d43271dc1c20cc79bd628df0544e92040e401..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/rule.py
+++ /dev/null
@@ -1,130 +0,0 @@
-from typing import Union
-
-from .align import AlignMethod
-from .cells import cell_len, set_cell_size
-from .console import Console, ConsoleOptions, RenderResult
-from .jupyter import JupyterMixin
-from .measure import Measurement
-from .style import Style
-from .text import Text
-
-
-class Rule(JupyterMixin):
-    """A console renderable to draw a horizontal rule (line).
-
-    Args:
-        title (Union[str, Text], optional): Text to render in the rule. Defaults to "".
-        characters (str, optional): Character(s) used to draw the line. Defaults to "─".
-        style (StyleType, optional): Style of Rule. Defaults to "rule.line".
-        end (str, optional): Character at end of Rule. defaults to "\\\\n"
-        align (str, optional): How to align the title, one of "left", "center", or "right". Defaults to "center".
-    """
-
-    def __init__(
-        self,
-        title: Union[str, Text] = "",
-        *,
-        characters: str = "─",
-        style: Union[str, Style] = "rule.line",
-        end: str = "\n",
-        align: AlignMethod = "center",
-    ) -> None:
-        if cell_len(characters) < 1:
-            raise ValueError(
-                "'characters' argument must have a cell width of at least 1"
-            )
-        if align not in ("left", "center", "right"):
-            raise ValueError(
-                f'invalid value for align, expected "left", "center", "right" (not {align!r})'
-            )
-        self.title = title
-        self.characters = characters
-        self.style = style
-        self.end = end
-        self.align = align
-
-    def __repr__(self) -> str:
-        return f"Rule({self.title!r}, {self.characters!r})"
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        width = options.max_width
-
-        characters = (
-            "-"
-            if (options.ascii_only and not self.characters.isascii())
-            else self.characters
-        )
-
-        chars_len = cell_len(characters)
-        if not self.title:
-            yield self._rule_line(chars_len, width)
-            return
-
-        if isinstance(self.title, Text):
-            title_text = self.title
-        else:
-            title_text = console.render_str(self.title, style="rule.text")
-
-        title_text.plain = title_text.plain.replace("\n", " ")
-        title_text.expand_tabs()
-
-        required_space = 4 if self.align == "center" else 2
-        truncate_width = max(0, width - required_space)
-        if not truncate_width:
-            yield self._rule_line(chars_len, width)
-            return
-
-        rule_text = Text(end=self.end)
-        if self.align == "center":
-            title_text.truncate(truncate_width, overflow="ellipsis")
-            side_width = (width - cell_len(title_text.plain)) // 2
-            left = Text(characters * (side_width // chars_len + 1))
-            left.truncate(side_width - 1)
-            right_length = width - cell_len(left.plain) - cell_len(title_text.plain)
-            right = Text(characters * (side_width // chars_len + 1))
-            right.truncate(right_length)
-            rule_text.append(left.plain + " ", self.style)
-            rule_text.append(title_text)
-            rule_text.append(" " + right.plain, self.style)
-        elif self.align == "left":
-            title_text.truncate(truncate_width, overflow="ellipsis")
-            rule_text.append(title_text)
-            rule_text.append(" ")
-            rule_text.append(characters * (width - rule_text.cell_len), self.style)
-        elif self.align == "right":
-            title_text.truncate(truncate_width, overflow="ellipsis")
-            rule_text.append(characters * (width - title_text.cell_len - 1), self.style)
-            rule_text.append(" ")
-            rule_text.append(title_text)
-
-        rule_text.plain = set_cell_size(rule_text.plain, width)
-        yield rule_text
-
-    def _rule_line(self, chars_len: int, width: int) -> Text:
-        rule_text = Text(self.characters * ((width // chars_len) + 1), self.style)
-        rule_text.truncate(width)
-        rule_text.plain = set_cell_size(rule_text.plain, width)
-        return rule_text
-
-    def __rich_measure__(
-        self, console: Console, options: ConsoleOptions
-    ) -> Measurement:
-        return Measurement(1, 1)
-
-
-if __name__ == "__main__":  # pragma: no cover
-    import sys
-
-    from rich.console import Console
-
-    try:
-        text = sys.argv[1]
-    except IndexError:
-        text = "Hello, World"
-    console = Console()
-    console.print(Rule(title=text))
-
-    console = Console()
-    console.print(Rule("foo"), width=4)
diff --git a/venv/lib/python3.10/site-packages/rich/scope.py b/venv/lib/python3.10/site-packages/rich/scope.py
deleted file mode 100644
index 36c06241f4618d4a8b74dfb1581f17514c693dc1..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/scope.py
+++ /dev/null
@@ -1,86 +0,0 @@
-from collections.abc import Mapping
-from typing import TYPE_CHECKING, Any, Optional, Tuple
-
-from .highlighter import ReprHighlighter
-from .panel import Panel
-from .pretty import Pretty
-from .table import Table
-from .text import Text, TextType
-
-if TYPE_CHECKING:
-    from .console import ConsoleRenderable
-
-
-def render_scope(
-    scope: "Mapping[str, Any]",
-    *,
-    title: Optional[TextType] = None,
-    sort_keys: bool = True,
-    indent_guides: bool = False,
-    max_length: Optional[int] = None,
-    max_string: Optional[int] = None,
-) -> "ConsoleRenderable":
-    """Render python variables in a given scope.
-
-    Args:
-        scope (Mapping): A mapping containing variable names and values.
-        title (str, optional): Optional title. Defaults to None.
-        sort_keys (bool, optional): Enable sorting of items. Defaults to True.
-        indent_guides (bool, optional): Enable indentation guides. Defaults to False.
-        max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
-            Defaults to None.
-        max_string (int, optional): Maximum length of string before truncating, or None to disable. Defaults to None.
-
-    Returns:
-        ConsoleRenderable: A renderable object.
-    """
-    highlighter = ReprHighlighter()
-    items_table = Table.grid(padding=(0, 1), expand=False)
-    items_table.add_column(justify="right")
-
-    def sort_items(item: Tuple[str, Any]) -> Tuple[bool, str]:
-        """Sort special variables first, then alphabetically."""
-        key, _ = item
-        return (not key.startswith("__"), key.lower())
-
-    items = sorted(scope.items(), key=sort_items) if sort_keys else scope.items()
-    for key, value in items:
-        key_text = Text.assemble(
-            (key, "scope.key.special" if key.startswith("__") else "scope.key"),
-            (" =", "scope.equals"),
-        )
-        items_table.add_row(
-            key_text,
-            Pretty(
-                value,
-                highlighter=highlighter,
-                indent_guides=indent_guides,
-                max_length=max_length,
-                max_string=max_string,
-            ),
-        )
-    return Panel.fit(
-        items_table,
-        title=title,
-        border_style="scope.border",
-        padding=(0, 1),
-    )
-
-
-if __name__ == "__main__":  # pragma: no cover
-    from rich import print
-
-    print()
-
-    def test(foo: float, bar: float) -> None:
-        list_of_things = [1, 2, 3, None, 4, True, False, "Hello World"]
-        dict_of_things = {
-            "version": "1.1",
-            "method": "confirmFruitPurchase",
-            "params": [["apple", "orange", "mangoes", "pomelo"], 1.123],
-            "id": "194521489",
-        }
-        print(render_scope(locals(), title="[i]locals", sort_keys=False))
-
-    test(20.3423, 3.1427)
-    print()
diff --git a/venv/lib/python3.10/site-packages/rich/screen.py b/venv/lib/python3.10/site-packages/rich/screen.py
deleted file mode 100644
index b4f7fd19de7ffc7e3c18702389e093402a633c5b..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/screen.py
+++ /dev/null
@@ -1,54 +0,0 @@
-from typing import Optional, TYPE_CHECKING
-
-from .segment import Segment
-from .style import StyleType
-from ._loop import loop_last
-
-
-if TYPE_CHECKING:
-    from .console import (
-        Console,
-        ConsoleOptions,
-        RenderResult,
-        RenderableType,
-        Group,
-    )
-
-
-class Screen:
-    """A renderable that fills the terminal screen and crops excess.
-
-    Args:
-        renderable (RenderableType): Child renderable.
-        style (StyleType, optional): Optional background style. Defaults to None.
-    """
-
-    renderable: "RenderableType"
-
-    def __init__(
-        self,
-        *renderables: "RenderableType",
-        style: Optional[StyleType] = None,
-        application_mode: bool = False,
-    ) -> None:
-        from rich.console import Group
-
-        self.renderable = Group(*renderables)
-        self.style = style
-        self.application_mode = application_mode
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        width, height = options.size
-        style = console.get_style(self.style) if self.style else None
-        render_options = options.update(width=width, height=height)
-        lines = console.render_lines(
-            self.renderable or "", render_options, style=style, pad=True
-        )
-        lines = Segment.set_shape(lines, width, height, style=style)
-        new_line = Segment("\n\r") if self.application_mode else Segment.line()
-        for last, line in loop_last(lines):
-            yield from line
-            if not last:
-                yield new_line
diff --git a/venv/lib/python3.10/site-packages/rich/segment.py b/venv/lib/python3.10/site-packages/rich/segment.py
deleted file mode 100644
index edcb52dd3f0c293daa3eb83c8aed69fc97639aff..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/segment.py
+++ /dev/null
@@ -1,752 +0,0 @@
-from enum import IntEnum
-from functools import lru_cache
-from itertools import filterfalse
-from logging import getLogger
-from operator import attrgetter
-from typing import (
-    TYPE_CHECKING,
-    Dict,
-    Iterable,
-    List,
-    NamedTuple,
-    Optional,
-    Sequence,
-    Tuple,
-    Type,
-    Union,
-)
-
-from .cells import (
-    _is_single_cell_widths,
-    cached_cell_len,
-    cell_len,
-    get_character_cell_size,
-    set_cell_size,
-)
-from .repr import Result, rich_repr
-from .style import Style
-
-if TYPE_CHECKING:
-    from .console import Console, ConsoleOptions, RenderResult
-
-log = getLogger("rich")
-
-
-class ControlType(IntEnum):
-    """Non-printable control codes which typically translate to ANSI codes."""
-
-    BELL = 1
-    CARRIAGE_RETURN = 2
-    HOME = 3
-    CLEAR = 4
-    SHOW_CURSOR = 5
-    HIDE_CURSOR = 6
-    ENABLE_ALT_SCREEN = 7
-    DISABLE_ALT_SCREEN = 8
-    CURSOR_UP = 9
-    CURSOR_DOWN = 10
-    CURSOR_FORWARD = 11
-    CURSOR_BACKWARD = 12
-    CURSOR_MOVE_TO_COLUMN = 13
-    CURSOR_MOVE_TO = 14
-    ERASE_IN_LINE = 15
-    SET_WINDOW_TITLE = 16
-
-
-ControlCode = Union[
-    Tuple[ControlType],
-    Tuple[ControlType, Union[int, str]],
-    Tuple[ControlType, int, int],
-]
-
-
-@rich_repr()
-class Segment(NamedTuple):
-    """A piece of text with associated style. Segments are produced by the Console render process and
-    are ultimately converted in to strings to be written to the terminal.
-
-    Args:
-        text (str): A piece of text.
-        style (:class:`~rich.style.Style`, optional): An optional style to apply to the text.
-        control (Tuple[ControlCode], optional): Optional sequence of control codes.
-
-    Attributes:
-        cell_length (int): The cell length of this Segment.
-    """
-
-    text: str
-    style: Optional[Style] = None
-    control: Optional[Sequence[ControlCode]] = None
-
-    @property
-    def cell_length(self) -> int:
-        """The number of terminal cells required to display self.text.
-
-        Returns:
-            int: A number of cells.
-        """
-        text, _style, control = self
-        return 0 if control else cell_len(text)
-
-    def __rich_repr__(self) -> Result:
-        yield self.text
-        if self.control is None:
-            if self.style is not None:
-                yield self.style
-        else:
-            yield self.style
-            yield self.control
-
-    def __bool__(self) -> bool:
-        """Check if the segment contains text."""
-        return bool(self.text)
-
-    @property
-    def is_control(self) -> bool:
-        """Check if the segment contains control codes."""
-        return self.control is not None
-
-    @classmethod
-    @lru_cache(1024 * 16)
-    def _split_cells(cls, segment: "Segment", cut: int) -> Tuple["Segment", "Segment"]:
-        """Split a segment in to two at a given cell position.
-
-        Note that splitting a double-width character, may result in that character turning
-        into two spaces.
-
-        Args:
-            segment (Segment): A segment to split.
-            cut (int): A cell position to cut on.
-
-        Returns:
-            A tuple of two segments.
-        """
-        text, style, control = segment
-        _Segment = Segment
-        cell_length = segment.cell_length
-        if cut >= cell_length:
-            return segment, _Segment("", style, control)
-
-        cell_size = get_character_cell_size
-
-        pos = int((cut / cell_length) * len(text))
-
-        while True:
-            before = text[:pos]
-            cell_pos = cell_len(before)
-            out_by = cell_pos - cut
-            if not out_by:
-                return (
-                    _Segment(before, style, control),
-                    _Segment(text[pos:], style, control),
-                )
-            if out_by == -1 and cell_size(text[pos]) == 2:
-                return (
-                    _Segment(text[:pos] + " ", style, control),
-                    _Segment(" " + text[pos + 1 :], style, control),
-                )
-            if out_by == +1 and cell_size(text[pos - 1]) == 2:
-                return (
-                    _Segment(text[: pos - 1] + " ", style, control),
-                    _Segment(" " + text[pos:], style, control),
-                )
-            if cell_pos < cut:
-                pos += 1
-            else:
-                pos -= 1
-
-    def split_cells(self, cut: int) -> Tuple["Segment", "Segment"]:
-        """Split segment in to two segments at the specified column.
-
-        If the cut point falls in the middle of a 2-cell wide character then it is replaced
-        by two spaces, to preserve the display width of the parent segment.
-
-        Args:
-            cut (int): Offset within the segment to cut.
-
-        Returns:
-            Tuple[Segment, Segment]: Two segments.
-        """
-        text, style, control = self
-        assert cut >= 0
-
-        if _is_single_cell_widths(text):
-            # Fast path with all 1 cell characters
-            if cut >= len(text):
-                return self, Segment("", style, control)
-            return (
-                Segment(text[:cut], style, control),
-                Segment(text[cut:], style, control),
-            )
-
-        return self._split_cells(self, cut)
-
-    @classmethod
-    def line(cls) -> "Segment":
-        """Make a new line segment."""
-        return cls("\n")
-
-    @classmethod
-    def apply_style(
-        cls,
-        segments: Iterable["Segment"],
-        style: Optional[Style] = None,
-        post_style: Optional[Style] = None,
-    ) -> Iterable["Segment"]:
-        """Apply style(s) to an iterable of segments.
-
-        Returns an iterable of segments where the style is replaced by ``style + segment.style + post_style``.
-
-        Args:
-            segments (Iterable[Segment]): Segments to process.
-            style (Style, optional): Base style. Defaults to None.
-            post_style (Style, optional): Style to apply on top of segment style. Defaults to None.
-
-        Returns:
-            Iterable[Segments]: A new iterable of segments (possibly the same iterable).
-        """
-        result_segments = segments
-        if style:
-            apply = style.__add__
-            result_segments = (
-                cls(text, None if control else apply(_style), control)
-                for text, _style, control in result_segments
-            )
-        if post_style:
-            result_segments = (
-                cls(
-                    text,
-                    (
-                        None
-                        if control
-                        else (_style + post_style if _style else post_style)
-                    ),
-                    control,
-                )
-                for text, _style, control in result_segments
-            )
-        return result_segments
-
-    @classmethod
-    def filter_control(
-        cls, segments: Iterable["Segment"], is_control: bool = False
-    ) -> Iterable["Segment"]:
-        """Filter segments by ``is_control`` attribute.
-
-        Args:
-            segments (Iterable[Segment]): An iterable of Segment instances.
-            is_control (bool, optional): is_control flag to match in search.
-
-        Returns:
-            Iterable[Segment]: And iterable of Segment instances.
-
-        """
-        if is_control:
-            return filter(attrgetter("control"), segments)
-        else:
-            return filterfalse(attrgetter("control"), segments)
-
-    @classmethod
-    def split_lines(cls, segments: Iterable["Segment"]) -> Iterable[List["Segment"]]:
-        """Split a sequence of segments in to a list of lines.
-
-        Args:
-            segments (Iterable[Segment]): Segments potentially containing line feeds.
-
-        Yields:
-            Iterable[List[Segment]]: Iterable of segment lists, one per line.
-        """
-        line: List[Segment] = []
-        append = line.append
-
-        for segment in segments:
-            if "\n" in segment.text and not segment.control:
-                text, style, _ = segment
-                while text:
-                    _text, new_line, text = text.partition("\n")
-                    if _text:
-                        append(cls(_text, style))
-                    if new_line:
-                        yield line
-                        line = []
-                        append = line.append
-            else:
-                append(segment)
-        if line:
-            yield line
-
-    @classmethod
-    def split_and_crop_lines(
-        cls,
-        segments: Iterable["Segment"],
-        length: int,
-        style: Optional[Style] = None,
-        pad: bool = True,
-        include_new_lines: bool = True,
-    ) -> Iterable[List["Segment"]]:
-        """Split segments in to lines, and crop lines greater than a given length.
-
-        Args:
-            segments (Iterable[Segment]): An iterable of segments, probably
-                generated from console.render.
-            length (int): Desired line length.
-            style (Style, optional): Style to use for any padding.
-            pad (bool): Enable padding of lines that are less than `length`.
-
-        Returns:
-            Iterable[List[Segment]]: An iterable of lines of segments.
-        """
-        line: List[Segment] = []
-        append = line.append
-
-        adjust_line_length = cls.adjust_line_length
-        new_line_segment = cls("\n")
-
-        for segment in segments:
-            if "\n" in segment.text and not segment.control:
-                text, segment_style, _ = segment
-                while text:
-                    _text, new_line, text = text.partition("\n")
-                    if _text:
-                        append(cls(_text, segment_style))
-                    if new_line:
-                        cropped_line = adjust_line_length(
-                            line, length, style=style, pad=pad
-                        )
-                        if include_new_lines:
-                            cropped_line.append(new_line_segment)
-                        yield cropped_line
-                        line.clear()
-            else:
-                append(segment)
-        if line:
-            yield adjust_line_length(line, length, style=style, pad=pad)
-
-    @classmethod
-    def adjust_line_length(
-        cls,
-        line: List["Segment"],
-        length: int,
-        style: Optional[Style] = None,
-        pad: bool = True,
-    ) -> List["Segment"]:
-        """Adjust a line to a given width (cropping or padding as required).
-
-        Args:
-            segments (Iterable[Segment]): A list of segments in a single line.
-            length (int): The desired width of the line.
-            style (Style, optional): The style of padding if used (space on the end). Defaults to None.
-            pad (bool, optional): Pad lines with spaces if they are shorter than `length`. Defaults to True.
-
-        Returns:
-            List[Segment]: A line of segments with the desired length.
-        """
-        line_length = sum(segment.cell_length for segment in line)
-        new_line: List[Segment]
-
-        if line_length < length:
-            if pad:
-                new_line = line + [cls(" " * (length - line_length), style)]
-            else:
-                new_line = line[:]
-        elif line_length > length:
-            new_line = []
-            append = new_line.append
-            line_length = 0
-            for segment in line:
-                segment_length = segment.cell_length
-                if line_length + segment_length < length or segment.control:
-                    append(segment)
-                    line_length += segment_length
-                else:
-                    text, segment_style, _ = segment
-                    text = set_cell_size(text, length - line_length)
-                    append(cls(text, segment_style))
-                    break
-        else:
-            new_line = line[:]
-        return new_line
-
-    @classmethod
-    def get_line_length(cls, line: List["Segment"]) -> int:
-        """Get the length of list of segments.
-
-        Args:
-            line (List[Segment]): A line encoded as a list of Segments (assumes no '\\\\n' characters),
-
-        Returns:
-            int: The length of the line.
-        """
-        _cell_len = cell_len
-        return sum(_cell_len(text) for text, style, control in line if not control)
-
-    @classmethod
-    def get_shape(cls, lines: List[List["Segment"]]) -> Tuple[int, int]:
-        """Get the shape (enclosing rectangle) of a list of lines.
-
-        Args:
-            lines (List[List[Segment]]): A list of lines (no '\\\\n' characters).
-
-        Returns:
-            Tuple[int, int]: Width and height in characters.
-        """
-        get_line_length = cls.get_line_length
-        max_width = max(get_line_length(line) for line in lines) if lines else 0
-        return (max_width, len(lines))
-
-    @classmethod
-    def set_shape(
-        cls,
-        lines: List[List["Segment"]],
-        width: int,
-        height: Optional[int] = None,
-        style: Optional[Style] = None,
-        new_lines: bool = False,
-    ) -> List[List["Segment"]]:
-        """Set the shape of a list of lines (enclosing rectangle).
-
-        Args:
-            lines (List[List[Segment]]): A list of lines.
-            width (int): Desired width.
-            height (int, optional): Desired height or None for no change.
-            style (Style, optional): Style of any padding added.
-            new_lines (bool, optional): Padded lines should include "\n". Defaults to False.
-
-        Returns:
-            List[List[Segment]]: New list of lines.
-        """
-        _height = height or len(lines)
-
-        blank = (
-            [cls(" " * width + "\n", style)] if new_lines else [cls(" " * width, style)]
-        )
-
-        adjust_line_length = cls.adjust_line_length
-        shaped_lines = lines[:_height]
-        shaped_lines[:] = [
-            adjust_line_length(line, width, style=style) for line in lines
-        ]
-        if len(shaped_lines) < _height:
-            shaped_lines.extend([blank] * (_height - len(shaped_lines)))
-        return shaped_lines
-
-    @classmethod
-    def align_top(
-        cls: Type["Segment"],
-        lines: List[List["Segment"]],
-        width: int,
-        height: int,
-        style: Style,
-        new_lines: bool = False,
-    ) -> List[List["Segment"]]:
-        """Aligns lines to top (adds extra lines to bottom as required).
-
-        Args:
-            lines (List[List[Segment]]): A list of lines.
-            width (int): Desired width.
-            height (int, optional): Desired height or None for no change.
-            style (Style): Style of any padding added.
-            new_lines (bool, optional): Padded lines should include "\n". Defaults to False.
-
-        Returns:
-            List[List[Segment]]: New list of lines.
-        """
-        extra_lines = height - len(lines)
-        if not extra_lines:
-            return lines[:]
-        lines = lines[:height]
-        blank = cls(" " * width + "\n", style) if new_lines else cls(" " * width, style)
-        lines = lines + [[blank]] * extra_lines
-        return lines
-
-    @classmethod
-    def align_bottom(
-        cls: Type["Segment"],
-        lines: List[List["Segment"]],
-        width: int,
-        height: int,
-        style: Style,
-        new_lines: bool = False,
-    ) -> List[List["Segment"]]:
-        """Aligns render to bottom (adds extra lines above as required).
-
-        Args:
-            lines (List[List[Segment]]): A list of lines.
-            width (int): Desired width.
-            height (int, optional): Desired height or None for no change.
-            style (Style): Style of any padding added. Defaults to None.
-            new_lines (bool, optional): Padded lines should include "\n". Defaults to False.
-
-        Returns:
-            List[List[Segment]]: New list of lines.
-        """
-        extra_lines = height - len(lines)
-        if not extra_lines:
-            return lines[:]
-        lines = lines[:height]
-        blank = cls(" " * width + "\n", style) if new_lines else cls(" " * width, style)
-        lines = [[blank]] * extra_lines + lines
-        return lines
-
-    @classmethod
-    def align_middle(
-        cls: Type["Segment"],
-        lines: List[List["Segment"]],
-        width: int,
-        height: int,
-        style: Style,
-        new_lines: bool = False,
-    ) -> List[List["Segment"]]:
-        """Aligns lines to middle (adds extra lines to above and below as required).
-
-        Args:
-            lines (List[List[Segment]]): A list of lines.
-            width (int): Desired width.
-            height (int, optional): Desired height or None for no change.
-            style (Style): Style of any padding added.
-            new_lines (bool, optional): Padded lines should include "\n". Defaults to False.
-
-        Returns:
-            List[List[Segment]]: New list of lines.
-        """
-        extra_lines = height - len(lines)
-        if not extra_lines:
-            return lines[:]
-        lines = lines[:height]
-        blank = cls(" " * width + "\n", style) if new_lines else cls(" " * width, style)
-        top_lines = extra_lines // 2
-        bottom_lines = extra_lines - top_lines
-        lines = [[blank]] * top_lines + lines + [[blank]] * bottom_lines
-        return lines
-
-    @classmethod
-    def simplify(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:
-        """Simplify an iterable of segments by combining contiguous segments with the same style.
-
-        Args:
-            segments (Iterable[Segment]): An iterable of segments.
-
-        Returns:
-            Iterable[Segment]: A possibly smaller iterable of segments that will render the same way.
-        """
-        iter_segments = iter(segments)
-        try:
-            last_segment = next(iter_segments)
-        except StopIteration:
-            return
-
-        _Segment = Segment
-        for segment in iter_segments:
-            if last_segment.style == segment.style and not segment.control:
-                last_segment = _Segment(
-                    last_segment.text + segment.text, last_segment.style
-                )
-            else:
-                yield last_segment
-                last_segment = segment
-        yield last_segment
-
-    @classmethod
-    def strip_links(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:
-        """Remove all links from an iterable of styles.
-
-        Args:
-            segments (Iterable[Segment]): An iterable segments.
-
-        Yields:
-            Segment: Segments with link removed.
-        """
-        for segment in segments:
-            if segment.control or segment.style is None:
-                yield segment
-            else:
-                text, style, _control = segment
-                yield cls(text, style.update_link(None) if style else None)
-
-    @classmethod
-    def strip_styles(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:
-        """Remove all styles from an iterable of segments.
-
-        Args:
-            segments (Iterable[Segment]): An iterable segments.
-
-        Yields:
-            Segment: Segments with styles replace with None
-        """
-        for text, _style, control in segments:
-            yield cls(text, None, control)
-
-    @classmethod
-    def remove_color(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:
-        """Remove all color from an iterable of segments.
-
-        Args:
-            segments (Iterable[Segment]): An iterable segments.
-
-        Yields:
-            Segment: Segments with colorless style.
-        """
-
-        cache: Dict[Style, Style] = {}
-        for text, style, control in segments:
-            if style:
-                colorless_style = cache.get(style)
-                if colorless_style is None:
-                    colorless_style = style.without_color
-                    cache[style] = colorless_style
-                yield cls(text, colorless_style, control)
-            else:
-                yield cls(text, None, control)
-
-    @classmethod
-    def divide(
-        cls, segments: Iterable["Segment"], cuts: Iterable[int]
-    ) -> Iterable[List["Segment"]]:
-        """Divides an iterable of segments in to portions.
-
-        Args:
-            cuts (Iterable[int]): Cell positions where to divide.
-
-        Yields:
-            [Iterable[List[Segment]]]: An iterable of Segments in List.
-        """
-        split_segments: List["Segment"] = []
-        add_segment = split_segments.append
-
-        iter_cuts = iter(cuts)
-
-        while True:
-            cut = next(iter_cuts, -1)
-            if cut == -1:
-                return
-            if cut != 0:
-                break
-            yield []
-        pos = 0
-
-        segments_clear = split_segments.clear
-        segments_copy = split_segments.copy
-
-        _cell_len = cached_cell_len
-        for segment in segments:
-            text, _style, control = segment
-            while text:
-                end_pos = pos if control else pos + _cell_len(text)
-                if end_pos < cut:
-                    add_segment(segment)
-                    pos = end_pos
-                    break
-
-                if end_pos == cut:
-                    add_segment(segment)
-                    yield segments_copy()
-                    segments_clear()
-                    pos = end_pos
-
-                    cut = next(iter_cuts, -1)
-                    if cut == -1:
-                        if split_segments:
-                            yield segments_copy()
-                        return
-
-                    break
-
-                else:
-                    before, segment = segment.split_cells(cut - pos)
-                    text, _style, control = segment
-                    add_segment(before)
-                    yield segments_copy()
-                    segments_clear()
-                    pos = cut
-
-                cut = next(iter_cuts, -1)
-                if cut == -1:
-                    if split_segments:
-                        yield segments_copy()
-                    return
-
-        yield segments_copy()
-
-
-class Segments:
-    """A simple renderable to render an iterable of segments. This class may be useful if
-    you want to print segments outside of a __rich_console__ method.
-
-    Args:
-        segments (Iterable[Segment]): An iterable of segments.
-        new_lines (bool, optional): Add new lines between segments. Defaults to False.
-    """
-
-    def __init__(self, segments: Iterable[Segment], new_lines: bool = False) -> None:
-        self.segments = list(segments)
-        self.new_lines = new_lines
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        if self.new_lines:
-            line = Segment.line()
-            for segment in self.segments:
-                yield segment
-                yield line
-        else:
-            yield from self.segments
-
-
-class SegmentLines:
-    def __init__(self, lines: Iterable[List[Segment]], new_lines: bool = False) -> None:
-        """A simple renderable containing a number of lines of segments. May be used as an intermediate
-        in rendering process.
-
-        Args:
-            lines (Iterable[List[Segment]]): Lists of segments forming lines.
-            new_lines (bool, optional): Insert new lines after each line. Defaults to False.
-        """
-        self.lines = list(lines)
-        self.new_lines = new_lines
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        if self.new_lines:
-            new_line = Segment.line()
-            for line in self.lines:
-                yield from line
-                yield new_line
-        else:
-            for line in self.lines:
-                yield from line
-
-
-if __name__ == "__main__":  # pragma: no cover
-    from rich.console import Console
-    from rich.syntax import Syntax
-    from rich.text import Text
-
-    code = """from rich.console import Console
-console = Console()
-text = Text.from_markup("Hello, [bold magenta]World[/]!")
-console.print(text)"""
-
-    text = Text.from_markup("Hello, [bold magenta]World[/]!")
-
-    console = Console()
-
-    console.rule("rich.Segment")
-    console.print(
-        "A Segment is the last step in the Rich render process before generating text with ANSI codes."
-    )
-    console.print("\nConsider the following code:\n")
-    console.print(Syntax(code, "python", line_numbers=True))
-    console.print()
-    console.print(
-        "When you call [b]print()[/b], Rich [i]renders[/i] the object in to the following:\n"
-    )
-    fragments = list(console.render(text))
-    console.print(fragments)
-    console.print()
-    console.print("The Segments are then processed to produce the following output:\n")
-    console.print(text)
-    console.print(
-        "\nYou will only need to know this if you are implementing your own Rich renderables."
-    )
diff --git a/venv/lib/python3.10/site-packages/rich/spinner.py b/venv/lib/python3.10/site-packages/rich/spinner.py
deleted file mode 100644
index 70570b6b0962e530e76d599d2f886d14d91bda33..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/spinner.py
+++ /dev/null
@@ -1,138 +0,0 @@
-from typing import cast, List, Optional, TYPE_CHECKING, Union
-
-from ._spinners import SPINNERS
-from .measure import Measurement
-from .table import Table
-from .text import Text
-
-if TYPE_CHECKING:
-    from .console import Console, ConsoleOptions, RenderResult, RenderableType
-    from .style import StyleType
-
-
-class Spinner:
-    """A spinner animation.
-
-    Args:
-        name (str): Name of spinner (run python -m rich.spinner).
-        text (RenderableType, optional): A renderable to display at the right of the spinner (str or Text typically). Defaults to "".
-        style (StyleType, optional): Style for spinner animation. Defaults to None.
-        speed (float, optional): Speed factor for animation. Defaults to 1.0.
-
-    Raises:
-        KeyError: If name isn't one of the supported spinner animations.
-    """
-
-    def __init__(
-        self,
-        name: str,
-        text: "RenderableType" = "",
-        *,
-        style: Optional["StyleType"] = None,
-        speed: float = 1.0,
-    ) -> None:
-        try:
-            spinner = SPINNERS[name]
-        except KeyError:
-            raise KeyError(f"no spinner called {name!r}")
-        self.text: "Union[RenderableType, Text]" = (
-            Text.from_markup(text) if isinstance(text, str) else text
-        )
-        self.name = name
-        self.frames = cast(List[str], spinner["frames"])[:]
-        self.interval = cast(float, spinner["interval"])
-        self.start_time: Optional[float] = None
-        self.style = style
-        self.speed = speed
-        self.frame_no_offset: float = 0.0
-        self._update_speed = 0.0
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        yield self.render(console.get_time())
-
-    def __rich_measure__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> Measurement:
-        text = self.render(0)
-        return Measurement.get(console, options, text)
-
-    def render(self, time: float) -> "RenderableType":
-        """Render the spinner for a given time.
-
-        Args:
-            time (float): Time in seconds.
-
-        Returns:
-            RenderableType: A renderable containing animation frame.
-        """
-        if self.start_time is None:
-            self.start_time = time
-
-        frame_no = ((time - self.start_time) * self.speed) / (
-            self.interval / 1000.0
-        ) + self.frame_no_offset
-        frame = Text(
-            self.frames[int(frame_no) % len(self.frames)], style=self.style or ""
-        )
-
-        if self._update_speed:
-            self.frame_no_offset = frame_no
-            self.start_time = time
-            self.speed = self._update_speed
-            self._update_speed = 0.0
-
-        if not self.text:
-            return frame
-        elif isinstance(self.text, (str, Text)):
-            return Text.assemble(frame, " ", self.text)
-        else:
-            table = Table.grid(padding=1)
-            table.add_row(frame, self.text)
-            return table
-
-    def update(
-        self,
-        *,
-        text: "RenderableType" = "",
-        style: Optional["StyleType"] = None,
-        speed: Optional[float] = None,
-    ) -> None:
-        """Updates attributes of a spinner after it has been started.
-
-        Args:
-            text (RenderableType, optional): A renderable to display at the right of the spinner (str or Text typically). Defaults to "".
-            style (StyleType, optional): Style for spinner animation. Defaults to None.
-            speed (float, optional): Speed factor for animation. Defaults to None.
-        """
-        if text:
-            self.text = Text.from_markup(text) if isinstance(text, str) else text
-        if style:
-            self.style = style
-        if speed:
-            self._update_speed = speed
-
-
-if __name__ == "__main__":  # pragma: no cover
-    from time import sleep
-
-    from .columns import Columns
-    from .panel import Panel
-    from .live import Live
-
-    all_spinners = Columns(
-        [
-            Spinner(spinner_name, text=Text(repr(spinner_name), style="green"))
-            for spinner_name in sorted(SPINNERS.keys())
-        ],
-        column_first=True,
-        expand=True,
-    )
-
-    with Live(
-        Panel(all_spinners, title="Spinners", border_style="blue"),
-        refresh_per_second=20,
-    ) as live:
-        while True:
-            sleep(0.1)
diff --git a/venv/lib/python3.10/site-packages/rich/status.py b/venv/lib/python3.10/site-packages/rich/status.py
deleted file mode 100644
index 65744838e3fcd7365ce6f268fe511d75d55e6a07..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/status.py
+++ /dev/null
@@ -1,131 +0,0 @@
-from types import TracebackType
-from typing import Optional, Type
-
-from .console import Console, RenderableType
-from .jupyter import JupyterMixin
-from .live import Live
-from .spinner import Spinner
-from .style import StyleType
-
-
-class Status(JupyterMixin):
-    """Displays a status indicator with a 'spinner' animation.
-
-    Args:
-        status (RenderableType): A status renderable (str or Text typically).
-        console (Console, optional): Console instance to use, or None for global console. Defaults to None.
-        spinner (str, optional): Name of spinner animation (see python -m rich.spinner). Defaults to "dots".
-        spinner_style (StyleType, optional): Style of spinner. Defaults to "status.spinner".
-        speed (float, optional): Speed factor for spinner animation. Defaults to 1.0.
-        refresh_per_second (float, optional): Number of refreshes per second. Defaults to 12.5.
-    """
-
-    def __init__(
-        self,
-        status: RenderableType,
-        *,
-        console: Optional[Console] = None,
-        spinner: str = "dots",
-        spinner_style: StyleType = "status.spinner",
-        speed: float = 1.0,
-        refresh_per_second: float = 12.5,
-    ):
-        self.status = status
-        self.spinner_style = spinner_style
-        self.speed = speed
-        self._spinner = Spinner(spinner, text=status, style=spinner_style, speed=speed)
-        self._live = Live(
-            self.renderable,
-            console=console,
-            refresh_per_second=refresh_per_second,
-            transient=True,
-        )
-
-    @property
-    def renderable(self) -> Spinner:
-        return self._spinner
-
-    @property
-    def console(self) -> "Console":
-        """Get the Console used by the Status objects."""
-        return self._live.console
-
-    def update(
-        self,
-        status: Optional[RenderableType] = None,
-        *,
-        spinner: Optional[str] = None,
-        spinner_style: Optional[StyleType] = None,
-        speed: Optional[float] = None,
-    ) -> None:
-        """Update status.
-
-        Args:
-            status (Optional[RenderableType], optional): New status renderable or None for no change. Defaults to None.
-            spinner (Optional[str], optional): New spinner or None for no change. Defaults to None.
-            spinner_style (Optional[StyleType], optional): New spinner style or None for no change. Defaults to None.
-            speed (Optional[float], optional): Speed factor for spinner animation or None for no change. Defaults to None.
-        """
-        if status is not None:
-            self.status = status
-        if spinner_style is not None:
-            self.spinner_style = spinner_style
-        if speed is not None:
-            self.speed = speed
-        if spinner is not None:
-            self._spinner = Spinner(
-                spinner, text=self.status, style=self.spinner_style, speed=self.speed
-            )
-            self._live.update(self.renderable, refresh=True)
-        else:
-            self._spinner.update(
-                text=self.status, style=self.spinner_style, speed=self.speed
-            )
-
-    def start(self) -> None:
-        """Start the status animation."""
-        self._live.start()
-
-    def stop(self) -> None:
-        """Stop the spinner animation."""
-        self._live.stop()
-
-    def __rich__(self) -> RenderableType:
-        return self.renderable
-
-    def __enter__(self) -> "Status":
-        self.start()
-        return self
-
-    def __exit__(
-        self,
-        exc_type: Optional[Type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
-    ) -> None:
-        self.stop()
-
-
-if __name__ == "__main__":  # pragma: no cover
-    from time import sleep
-
-    from .console import Console
-
-    console = Console()
-    with console.status("[magenta]Covid detector booting up") as status:
-        sleep(3)
-        console.log("Importing advanced AI")
-        sleep(3)
-        console.log("Advanced Covid AI Ready")
-        sleep(3)
-        status.update(status="[bold blue] Scanning for Covid", spinner="earth")
-        sleep(3)
-        console.log("Found 10,000,000,000 copies of Covid32.exe")
-        sleep(3)
-        status.update(
-            status="[bold red]Moving Covid32.exe to Trash",
-            spinner="bouncingBall",
-            spinner_style="yellow",
-        )
-        sleep(5)
-    console.print("[bold green]Covid deleted successfully")
diff --git a/venv/lib/python3.10/site-packages/rich/style.py b/venv/lib/python3.10/site-packages/rich/style.py
deleted file mode 100644
index 835d06f3eb2242159a65ac2dd8137c336f7ac324..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/style.py
+++ /dev/null
@@ -1,796 +0,0 @@
-import sys
-from functools import lru_cache
-from marshal import dumps, loads
-from random import randint
-from typing import Any, Dict, Iterable, List, Optional, Type, Union, cast
-
-from . import errors
-from .color import Color, ColorParseError, ColorSystem, blend_rgb
-from .repr import Result, rich_repr
-from .terminal_theme import DEFAULT_TERMINAL_THEME, TerminalTheme
-
-# Style instances and style definitions are often interchangeable
-StyleType = Union[str, "Style"]
-
-
-class _Bit:
-    """A descriptor to get/set a style attribute bit."""
-
-    __slots__ = ["bit"]
-
-    def __init__(self, bit_no: int) -> None:
-        self.bit = 1 << bit_no
-
-    def __get__(self, obj: "Style", objtype: Type["Style"]) -> Optional[bool]:
-        if obj._set_attributes & self.bit:
-            return obj._attributes & self.bit != 0
-        return None
-
-
-@rich_repr
-class Style:
-    """A terminal style.
-
-    A terminal style consists of a color (`color`), a background color (`bgcolor`), and a number of attributes, such
-    as bold, italic etc. The attributes have 3 states: they can either be on
-    (``True``), off (``False``), or not set (``None``).
-
-    Args:
-        color (Union[Color, str], optional): Color of terminal text. Defaults to None.
-        bgcolor (Union[Color, str], optional): Color of terminal background. Defaults to None.
-        bold (bool, optional): Enable bold text. Defaults to None.
-        dim (bool, optional): Enable dim text. Defaults to None.
-        italic (bool, optional): Enable italic text. Defaults to None.
-        underline (bool, optional): Enable underlined text. Defaults to None.
-        blink (bool, optional): Enabled blinking text. Defaults to None.
-        blink2 (bool, optional): Enable fast blinking text. Defaults to None.
-        reverse (bool, optional): Enabled reverse text. Defaults to None.
-        conceal (bool, optional): Enable concealed text. Defaults to None.
-        strike (bool, optional): Enable strikethrough text. Defaults to None.
-        underline2 (bool, optional): Enable doubly underlined text. Defaults to None.
-        frame (bool, optional): Enable framed text. Defaults to None.
-        encircle (bool, optional): Enable encircled text. Defaults to None.
-        overline (bool, optional): Enable overlined text. Defaults to None.
-        link (str, link): Link URL. Defaults to None.
-
-    """
-
-    _color: Optional[Color]
-    _bgcolor: Optional[Color]
-    _attributes: int
-    _set_attributes: int
-    _hash: Optional[int]
-    _null: bool
-    _meta: Optional[bytes]
-
-    __slots__ = [
-        "_color",
-        "_bgcolor",
-        "_attributes",
-        "_set_attributes",
-        "_link",
-        "_link_id",
-        "_ansi",
-        "_style_definition",
-        "_hash",
-        "_null",
-        "_meta",
-    ]
-
-    # maps bits on to SGR parameter
-    _style_map = {
-        0: "1",
-        1: "2",
-        2: "3",
-        3: "4",
-        4: "5",
-        5: "6",
-        6: "7",
-        7: "8",
-        8: "9",
-        9: "21",
-        10: "51",
-        11: "52",
-        12: "53",
-    }
-
-    STYLE_ATTRIBUTES = {
-        "dim": "dim",
-        "d": "dim",
-        "bold": "bold",
-        "b": "bold",
-        "italic": "italic",
-        "i": "italic",
-        "underline": "underline",
-        "u": "underline",
-        "blink": "blink",
-        "blink2": "blink2",
-        "reverse": "reverse",
-        "r": "reverse",
-        "conceal": "conceal",
-        "c": "conceal",
-        "strike": "strike",
-        "s": "strike",
-        "underline2": "underline2",
-        "uu": "underline2",
-        "frame": "frame",
-        "encircle": "encircle",
-        "overline": "overline",
-        "o": "overline",
-    }
-
-    def __init__(
-        self,
-        *,
-        color: Optional[Union[Color, str]] = None,
-        bgcolor: Optional[Union[Color, str]] = None,
-        bold: Optional[bool] = None,
-        dim: Optional[bool] = None,
-        italic: Optional[bool] = None,
-        underline: Optional[bool] = None,
-        blink: Optional[bool] = None,
-        blink2: Optional[bool] = None,
-        reverse: Optional[bool] = None,
-        conceal: Optional[bool] = None,
-        strike: Optional[bool] = None,
-        underline2: Optional[bool] = None,
-        frame: Optional[bool] = None,
-        encircle: Optional[bool] = None,
-        overline: Optional[bool] = None,
-        link: Optional[str] = None,
-        meta: Optional[Dict[str, Any]] = None,
-    ):
-        self._ansi: Optional[str] = None
-        self._style_definition: Optional[str] = None
-
-        def _make_color(color: Union[Color, str]) -> Color:
-            return color if isinstance(color, Color) else Color.parse(color)
-
-        self._color = None if color is None else _make_color(color)
-        self._bgcolor = None if bgcolor is None else _make_color(bgcolor)
-        self._set_attributes = sum(
-            (
-                bold is not None,
-                dim is not None and 2,
-                italic is not None and 4,
-                underline is not None and 8,
-                blink is not None and 16,
-                blink2 is not None and 32,
-                reverse is not None and 64,
-                conceal is not None and 128,
-                strike is not None and 256,
-                underline2 is not None and 512,
-                frame is not None and 1024,
-                encircle is not None and 2048,
-                overline is not None and 4096,
-            )
-        )
-        self._attributes = (
-            sum(
-                (
-                    bold and 1 or 0,
-                    dim and 2 or 0,
-                    italic and 4 or 0,
-                    underline and 8 or 0,
-                    blink and 16 or 0,
-                    blink2 and 32 or 0,
-                    reverse and 64 or 0,
-                    conceal and 128 or 0,
-                    strike and 256 or 0,
-                    underline2 and 512 or 0,
-                    frame and 1024 or 0,
-                    encircle and 2048 or 0,
-                    overline and 4096 or 0,
-                )
-            )
-            if self._set_attributes
-            else 0
-        )
-
-        self._link = link
-        self._meta = None if meta is None else dumps(meta)
-        self._link_id = (
-            f"{randint(0, 999999)}{hash(self._meta)}" if (link or meta) else ""
-        )
-        self._hash: Optional[int] = None
-        self._null = not (self._set_attributes or color or bgcolor or link or meta)
-
-    @classmethod
-    def null(cls) -> "Style":
-        """Create an 'null' style, equivalent to Style(), but more performant."""
-        return NULL_STYLE
-
-    @classmethod
-    def from_color(
-        cls, color: Optional[Color] = None, bgcolor: Optional[Color] = None
-    ) -> "Style":
-        """Create a new style with colors and no attributes.
-
-        Returns:
-            color (Optional[Color]): A (foreground) color, or None for no color. Defaults to None.
-            bgcolor (Optional[Color]): A (background) color, or None for no color. Defaults to None.
-        """
-        style: Style = cls.__new__(Style)
-        style._ansi = None
-        style._style_definition = None
-        style._color = color
-        style._bgcolor = bgcolor
-        style._set_attributes = 0
-        style._attributes = 0
-        style._link = None
-        style._link_id = ""
-        style._meta = None
-        style._null = not (color or bgcolor)
-        style._hash = None
-        return style
-
-    @classmethod
-    def from_meta(cls, meta: Optional[Dict[str, Any]]) -> "Style":
-        """Create a new style with meta data.
-
-        Returns:
-            meta (Optional[Dict[str, Any]]): A dictionary of meta data. Defaults to None.
-        """
-        style: Style = cls.__new__(Style)
-        style._ansi = None
-        style._style_definition = None
-        style._color = None
-        style._bgcolor = None
-        style._set_attributes = 0
-        style._attributes = 0
-        style._link = None
-        style._meta = dumps(meta)
-        style._link_id = f"{randint(0, 999999)}{hash(style._meta)}"
-        style._hash = None
-        style._null = not (meta)
-        return style
-
-    @classmethod
-    def on(cls, meta: Optional[Dict[str, Any]] = None, **handlers: Any) -> "Style":
-        """Create a blank style with meta information.
-
-        Example:
-            style = Style.on(click=self.on_click)
-
-        Args:
-            meta (Optional[Dict[str, Any]], optional): An optional dict of meta information.
-            **handlers (Any): Keyword arguments are translated in to handlers.
-
-        Returns:
-            Style: A Style with meta information attached.
-        """
-        meta = {} if meta is None else meta
-        meta.update({f"@{key}": value for key, value in handlers.items()})
-        return cls.from_meta(meta)
-
-    bold = _Bit(0)
-    dim = _Bit(1)
-    italic = _Bit(2)
-    underline = _Bit(3)
-    blink = _Bit(4)
-    blink2 = _Bit(5)
-    reverse = _Bit(6)
-    conceal = _Bit(7)
-    strike = _Bit(8)
-    underline2 = _Bit(9)
-    frame = _Bit(10)
-    encircle = _Bit(11)
-    overline = _Bit(12)
-
-    @property
-    def link_id(self) -> str:
-        """Get a link id, used in ansi code for links."""
-        return self._link_id
-
-    def __str__(self) -> str:
-        """Re-generate style definition from attributes."""
-        if self._style_definition is None:
-            attributes: List[str] = []
-            append = attributes.append
-            bits = self._set_attributes
-            if bits & 0b0000000001111:
-                if bits & 1:
-                    append("bold" if self.bold else "not bold")
-                if bits & (1 << 1):
-                    append("dim" if self.dim else "not dim")
-                if bits & (1 << 2):
-                    append("italic" if self.italic else "not italic")
-                if bits & (1 << 3):
-                    append("underline" if self.underline else "not underline")
-            if bits & 0b0000111110000:
-                if bits & (1 << 4):
-                    append("blink" if self.blink else "not blink")
-                if bits & (1 << 5):
-                    append("blink2" if self.blink2 else "not blink2")
-                if bits & (1 << 6):
-                    append("reverse" if self.reverse else "not reverse")
-                if bits & (1 << 7):
-                    append("conceal" if self.conceal else "not conceal")
-                if bits & (1 << 8):
-                    append("strike" if self.strike else "not strike")
-            if bits & 0b1111000000000:
-                if bits & (1 << 9):
-                    append("underline2" if self.underline2 else "not underline2")
-                if bits & (1 << 10):
-                    append("frame" if self.frame else "not frame")
-                if bits & (1 << 11):
-                    append("encircle" if self.encircle else "not encircle")
-                if bits & (1 << 12):
-                    append("overline" if self.overline else "not overline")
-            if self._color is not None:
-                append(self._color.name)
-            if self._bgcolor is not None:
-                append("on")
-                append(self._bgcolor.name)
-            if self._link:
-                append("link")
-                append(self._link)
-            self._style_definition = " ".join(attributes) or "none"
-        return self._style_definition
-
-    def __bool__(self) -> bool:
-        """A Style is false if it has no attributes, colors, or links."""
-        return not self._null
-
-    def _make_ansi_codes(self, color_system: ColorSystem) -> str:
-        """Generate ANSI codes for this style.
-
-        Args:
-            color_system (ColorSystem): Color system.
-
-        Returns:
-            str: String containing codes.
-        """
-
-        if self._ansi is None:
-            sgr: List[str] = []
-            append = sgr.append
-            _style_map = self._style_map
-            attributes = self._attributes & self._set_attributes
-            if attributes:
-                if attributes & 1:
-                    append(_style_map[0])
-                if attributes & 2:
-                    append(_style_map[1])
-                if attributes & 4:
-                    append(_style_map[2])
-                if attributes & 8:
-                    append(_style_map[3])
-                if attributes & 0b0000111110000:
-                    for bit in range(4, 9):
-                        if attributes & (1 << bit):
-                            append(_style_map[bit])
-                if attributes & 0b1111000000000:
-                    for bit in range(9, 13):
-                        if attributes & (1 << bit):
-                            append(_style_map[bit])
-            if self._color is not None:
-                sgr.extend(self._color.downgrade(color_system).get_ansi_codes())
-            if self._bgcolor is not None:
-                sgr.extend(
-                    self._bgcolor.downgrade(color_system).get_ansi_codes(
-                        foreground=False
-                    )
-                )
-            self._ansi = ";".join(sgr)
-        return self._ansi
-
-    @classmethod
-    @lru_cache(maxsize=1024)
-    def normalize(cls, style: str) -> str:
-        """Normalize a style definition so that styles with the same effect have the same string
-        representation.
-
-        Args:
-            style (str): A style definition.
-
-        Returns:
-            str: Normal form of style definition.
-        """
-        try:
-            return str(cls.parse(style))
-        except errors.StyleSyntaxError:
-            return style.strip().lower()
-
-    @classmethod
-    def pick_first(cls, *values: Optional[StyleType]) -> StyleType:
-        """Pick first non-None style."""
-        for value in values:
-            if value is not None:
-                return value
-        raise ValueError("expected at least one non-None style")
-
-    def __rich_repr__(self) -> Result:
-        yield "color", self.color, None
-        yield "bgcolor", self.bgcolor, None
-        yield "bold", self.bold, None,
-        yield "dim", self.dim, None,
-        yield "italic", self.italic, None
-        yield "underline", self.underline, None,
-        yield "blink", self.blink, None
-        yield "blink2", self.blink2, None
-        yield "reverse", self.reverse, None
-        yield "conceal", self.conceal, None
-        yield "strike", self.strike, None
-        yield "underline2", self.underline2, None
-        yield "frame", self.frame, None
-        yield "encircle", self.encircle, None
-        yield "link", self.link, None
-        if self._meta:
-            yield "meta", self.meta
-
-    def __eq__(self, other: Any) -> bool:
-        if not isinstance(other, Style):
-            return NotImplemented
-        return self.__hash__() == other.__hash__()
-
-    def __ne__(self, other: Any) -> bool:
-        if not isinstance(other, Style):
-            return NotImplemented
-        return self.__hash__() != other.__hash__()
-
-    def __hash__(self) -> int:
-        if self._hash is not None:
-            return self._hash
-        self._hash = hash(
-            (
-                self._color,
-                self._bgcolor,
-                self._attributes,
-                self._set_attributes,
-                self._link,
-                self._meta,
-            )
-        )
-        return self._hash
-
-    @property
-    def color(self) -> Optional[Color]:
-        """The foreground color or None if it is not set."""
-        return self._color
-
-    @property
-    def bgcolor(self) -> Optional[Color]:
-        """The background color or None if it is not set."""
-        return self._bgcolor
-
-    @property
-    def link(self) -> Optional[str]:
-        """Link text, if set."""
-        return self._link
-
-    @property
-    def transparent_background(self) -> bool:
-        """Check if the style specified a transparent background."""
-        return self.bgcolor is None or self.bgcolor.is_default
-
-    @property
-    def background_style(self) -> "Style":
-        """A Style with background only."""
-        return Style(bgcolor=self.bgcolor)
-
-    @property
-    def meta(self) -> Dict[str, Any]:
-        """Get meta information (can not be changed after construction)."""
-        return {} if self._meta is None else cast(Dict[str, Any], loads(self._meta))
-
-    @property
-    def without_color(self) -> "Style":
-        """Get a copy of the style with color removed."""
-        if self._null:
-            return NULL_STYLE
-        style: Style = self.__new__(Style)
-        style._ansi = None
-        style._style_definition = None
-        style._color = None
-        style._bgcolor = None
-        style._attributes = self._attributes
-        style._set_attributes = self._set_attributes
-        style._link = self._link
-        style._link_id = f"{randint(0, 999999)}" if self._link else ""
-        style._null = False
-        style._meta = None
-        style._hash = None
-        return style
-
-    @classmethod
-    @lru_cache(maxsize=4096)
-    def parse(cls, style_definition: str) -> "Style":
-        """Parse a style definition.
-
-        Args:
-            style_definition (str): A string containing a style.
-
-        Raises:
-            errors.StyleSyntaxError: If the style definition syntax is invalid.
-
-        Returns:
-            `Style`: A Style instance.
-        """
-        if style_definition.strip() == "none" or not style_definition:
-            return cls.null()
-
-        STYLE_ATTRIBUTES = cls.STYLE_ATTRIBUTES
-        color: Optional[str] = None
-        bgcolor: Optional[str] = None
-        attributes: Dict[str, Optional[Any]] = {}
-        link: Optional[str] = None
-
-        words = iter(style_definition.split())
-        for original_word in words:
-            word = original_word.lower()
-            if word == "on":
-                word = next(words, "")
-                if not word:
-                    raise errors.StyleSyntaxError("color expected after 'on'")
-                try:
-                    Color.parse(word)
-                except ColorParseError as error:
-                    raise errors.StyleSyntaxError(
-                        f"unable to parse {word!r} as background color; {error}"
-                    ) from None
-                bgcolor = word
-
-            elif word == "not":
-                word = next(words, "")
-                attribute = STYLE_ATTRIBUTES.get(word)
-                if attribute is None:
-                    raise errors.StyleSyntaxError(
-                        f"expected style attribute after 'not', found {word!r}"
-                    )
-                attributes[attribute] = False
-
-            elif word == "link":
-                word = next(words, "")
-                if not word:
-                    raise errors.StyleSyntaxError("URL expected after 'link'")
-                link = word
-
-            elif word in STYLE_ATTRIBUTES:
-                attributes[STYLE_ATTRIBUTES[word]] = True
-
-            else:
-                try:
-                    Color.parse(word)
-                except ColorParseError as error:
-                    raise errors.StyleSyntaxError(
-                        f"unable to parse {word!r} as color; {error}"
-                    ) from None
-                color = word
-        style = Style(color=color, bgcolor=bgcolor, link=link, **attributes)
-        return style
-
-    @lru_cache(maxsize=1024)
-    def get_html_style(self, theme: Optional[TerminalTheme] = None) -> str:
-        """Get a CSS style rule."""
-        theme = theme or DEFAULT_TERMINAL_THEME
-        css: List[str] = []
-        append = css.append
-
-        color = self.color
-        bgcolor = self.bgcolor
-        if self.reverse:
-            color, bgcolor = bgcolor, color
-        if self.dim:
-            foreground_color = (
-                theme.foreground_color if color is None else color.get_truecolor(theme)
-            )
-            color = Color.from_triplet(
-                blend_rgb(foreground_color, theme.background_color, 0.5)
-            )
-        if color is not None:
-            theme_color = color.get_truecolor(theme)
-            append(f"color: {theme_color.hex}")
-            append(f"text-decoration-color: {theme_color.hex}")
-        if bgcolor is not None:
-            theme_color = bgcolor.get_truecolor(theme, foreground=False)
-            append(f"background-color: {theme_color.hex}")
-        if self.bold:
-            append("font-weight: bold")
-        if self.italic:
-            append("font-style: italic")
-        if self.underline:
-            append("text-decoration: underline")
-        if self.strike:
-            append("text-decoration: line-through")
-        if self.overline:
-            append("text-decoration: overline")
-        return "; ".join(css)
-
-    @classmethod
-    def combine(cls, styles: Iterable["Style"]) -> "Style":
-        """Combine styles and get result.
-
-        Args:
-            styles (Iterable[Style]): Styles to combine.
-
-        Returns:
-            Style: A new style instance.
-        """
-        iter_styles = iter(styles)
-        return sum(iter_styles, next(iter_styles))
-
-    @classmethod
-    def chain(cls, *styles: "Style") -> "Style":
-        """Combine styles from positional argument in to a single style.
-
-        Args:
-            *styles (Iterable[Style]): Styles to combine.
-
-        Returns:
-            Style: A new style instance.
-        """
-        iter_styles = iter(styles)
-        return sum(iter_styles, next(iter_styles))
-
-    def copy(self) -> "Style":
-        """Get a copy of this style.
-
-        Returns:
-            Style: A new Style instance with identical attributes.
-        """
-        if self._null:
-            return NULL_STYLE
-        style: Style = self.__new__(Style)
-        style._ansi = self._ansi
-        style._style_definition = self._style_definition
-        style._color = self._color
-        style._bgcolor = self._bgcolor
-        style._attributes = self._attributes
-        style._set_attributes = self._set_attributes
-        style._link = self._link
-        style._link_id = f"{randint(0, 999999)}" if self._link else ""
-        style._hash = self._hash
-        style._null = False
-        style._meta = self._meta
-        return style
-
-    @lru_cache(maxsize=128)
-    def clear_meta_and_links(self) -> "Style":
-        """Get a copy of this style with link and meta information removed.
-
-        Returns:
-            Style: New style object.
-        """
-        if self._null:
-            return NULL_STYLE
-        style: Style = self.__new__(Style)
-        style._ansi = self._ansi
-        style._style_definition = self._style_definition
-        style._color = self._color
-        style._bgcolor = self._bgcolor
-        style._attributes = self._attributes
-        style._set_attributes = self._set_attributes
-        style._link = None
-        style._link_id = ""
-        style._hash = None
-        style._null = False
-        style._meta = None
-        return style
-
-    def update_link(self, link: Optional[str] = None) -> "Style":
-        """Get a copy with a different value for link.
-
-        Args:
-            link (str, optional): New value for link. Defaults to None.
-
-        Returns:
-            Style: A new Style instance.
-        """
-        style: Style = self.__new__(Style)
-        style._ansi = self._ansi
-        style._style_definition = self._style_definition
-        style._color = self._color
-        style._bgcolor = self._bgcolor
-        style._attributes = self._attributes
-        style._set_attributes = self._set_attributes
-        style._link = link
-        style._link_id = f"{randint(0, 999999)}" if link else ""
-        style._hash = None
-        style._null = False
-        style._meta = self._meta
-        return style
-
-    def render(
-        self,
-        text: str = "",
-        *,
-        color_system: Optional[ColorSystem] = ColorSystem.TRUECOLOR,
-        legacy_windows: bool = False,
-    ) -> str:
-        """Render the ANSI codes for the style.
-
-        Args:
-            text (str, optional): A string to style. Defaults to "".
-            color_system (Optional[ColorSystem], optional): Color system to render to. Defaults to ColorSystem.TRUECOLOR.
-
-        Returns:
-            str: A string containing ANSI style codes.
-        """
-        if not text or color_system is None:
-            return text
-        attrs = self._ansi or self._make_ansi_codes(color_system)
-        rendered = f"\x1b[{attrs}m{text}\x1b[0m" if attrs else text
-        if self._link and not legacy_windows:
-            rendered = (
-                f"\x1b]8;id={self._link_id};{self._link}\x1b\\{rendered}\x1b]8;;\x1b\\"
-            )
-        return rendered
-
-    def test(self, text: Optional[str] = None) -> None:
-        """Write text with style directly to terminal.
-
-        This method is for testing purposes only.
-
-        Args:
-            text (Optional[str], optional): Text to style or None for style name.
-
-        """
-        text = text or str(self)
-        sys.stdout.write(f"{self.render(text)}\n")
-
-    @lru_cache(maxsize=1024)
-    def _add(self, style: Optional["Style"]) -> "Style":
-        if style is None or style._null:
-            return self
-        if self._null:
-            return style
-        new_style: Style = self.__new__(Style)
-        new_style._ansi = None
-        new_style._style_definition = None
-        new_style._color = style._color or self._color
-        new_style._bgcolor = style._bgcolor or self._bgcolor
-        new_style._attributes = (self._attributes & ~style._set_attributes) | (
-            style._attributes & style._set_attributes
-        )
-        new_style._set_attributes = self._set_attributes | style._set_attributes
-        new_style._link = style._link or self._link
-        new_style._link_id = style._link_id or self._link_id
-        new_style._null = style._null
-        if self._meta and style._meta:
-            new_style._meta = dumps({**self.meta, **style.meta})
-        else:
-            new_style._meta = self._meta or style._meta
-        new_style._hash = None
-        return new_style
-
-    def __add__(self, style: Optional["Style"]) -> "Style":
-        combined_style = self._add(style)
-        return combined_style.copy() if combined_style.link else combined_style
-
-
-NULL_STYLE = Style()
-
-
-class StyleStack:
-    """A stack of styles."""
-
-    __slots__ = ["_stack"]
-
-    def __init__(self, default_style: "Style") -> None:
-        self._stack: List[Style] = [default_style]
-
-    def __repr__(self) -> str:
-        return f""
-
-    @property
-    def current(self) -> Style:
-        """Get the Style at the top of the stack."""
-        return self._stack[-1]
-
-    def push(self, style: Style) -> None:
-        """Push a new style on to the stack.
-
-        Args:
-            style (Style): New style to combine with current style.
-        """
-        self._stack.append(self._stack[-1] + style)
-
-    def pop(self) -> Style:
-        """Pop last style and discard.
-
-        Returns:
-            Style: New current style (also available as stack.current)
-        """
-        self._stack.pop()
-        return self._stack[-1]
diff --git a/venv/lib/python3.10/site-packages/rich/styled.py b/venv/lib/python3.10/site-packages/rich/styled.py
deleted file mode 100644
index 27243beb76bb22c2851b97227249e70068ec4420..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/styled.py
+++ /dev/null
@@ -1,42 +0,0 @@
-from typing import TYPE_CHECKING
-
-from .measure import Measurement
-from .segment import Segment
-from .style import StyleType
-
-if TYPE_CHECKING:
-    from .console import Console, ConsoleOptions, RenderResult, RenderableType
-
-
-class Styled:
-    """Apply a style to a renderable.
-
-    Args:
-        renderable (RenderableType): Any renderable.
-        style (StyleType): A style to apply across the entire renderable.
-    """
-
-    def __init__(self, renderable: "RenderableType", style: "StyleType") -> None:
-        self.renderable = renderable
-        self.style = style
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        style = console.get_style(self.style)
-        rendered_segments = console.render(self.renderable, options)
-        segments = Segment.apply_style(rendered_segments, style)
-        return segments
-
-    def __rich_measure__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> Measurement:
-        return Measurement.get(console, options, self.renderable)
-
-
-if __name__ == "__main__":  # pragma: no cover
-    from rich import print
-    from rich.panel import Panel
-
-    panel = Styled(Panel("hello"), "on blue")
-    print(panel)
diff --git a/venv/lib/python3.10/site-packages/rich/syntax.py b/venv/lib/python3.10/site-packages/rich/syntax.py
deleted file mode 100644
index cff8fd235dba57f0ae660b5cbf3409d431452d8c..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/syntax.py
+++ /dev/null
@@ -1,966 +0,0 @@
-import os.path
-import re
-import sys
-import textwrap
-from abc import ABC, abstractmethod
-from pathlib import Path
-from typing import (
-    Any,
-    Dict,
-    Iterable,
-    List,
-    NamedTuple,
-    Optional,
-    Sequence,
-    Set,
-    Tuple,
-    Type,
-    Union,
-)
-
-from pygments.lexer import Lexer
-from pygments.lexers import get_lexer_by_name, guess_lexer_for_filename
-from pygments.style import Style as PygmentsStyle
-from pygments.styles import get_style_by_name
-from pygments.token import (
-    Comment,
-    Error,
-    Generic,
-    Keyword,
-    Name,
-    Number,
-    Operator,
-    String,
-    Token,
-    Whitespace,
-)
-from pygments.util import ClassNotFound
-
-from rich.containers import Lines
-from rich.padding import Padding, PaddingDimensions
-
-from ._loop import loop_first
-from .cells import cell_len
-from .color import Color, blend_rgb
-from .console import Console, ConsoleOptions, JustifyMethod, RenderResult
-from .jupyter import JupyterMixin
-from .measure import Measurement
-from .segment import Segment, Segments
-from .style import Style, StyleType
-from .text import Text
-
-TokenType = Tuple[str, ...]
-
-WINDOWS = sys.platform == "win32"
-DEFAULT_THEME = "monokai"
-
-# The following styles are based on https://github.com/pygments/pygments/blob/master/pygments/formatters/terminal.py
-# A few modifications were made
-
-ANSI_LIGHT: Dict[TokenType, Style] = {
-    Token: Style(),
-    Whitespace: Style(color="white"),
-    Comment: Style(dim=True),
-    Comment.Preproc: Style(color="cyan"),
-    Keyword: Style(color="blue"),
-    Keyword.Type: Style(color="cyan"),
-    Operator.Word: Style(color="magenta"),
-    Name.Builtin: Style(color="cyan"),
-    Name.Function: Style(color="green"),
-    Name.Namespace: Style(color="cyan", underline=True),
-    Name.Class: Style(color="green", underline=True),
-    Name.Exception: Style(color="cyan"),
-    Name.Decorator: Style(color="magenta", bold=True),
-    Name.Variable: Style(color="red"),
-    Name.Constant: Style(color="red"),
-    Name.Attribute: Style(color="cyan"),
-    Name.Tag: Style(color="bright_blue"),
-    String: Style(color="yellow"),
-    Number: Style(color="blue"),
-    Generic.Deleted: Style(color="bright_red"),
-    Generic.Inserted: Style(color="green"),
-    Generic.Heading: Style(bold=True),
-    Generic.Subheading: Style(color="magenta", bold=True),
-    Generic.Prompt: Style(bold=True),
-    Generic.Error: Style(color="bright_red"),
-    Error: Style(color="red", underline=True),
-}
-
-ANSI_DARK: Dict[TokenType, Style] = {
-    Token: Style(),
-    Whitespace: Style(color="bright_black"),
-    Comment: Style(dim=True),
-    Comment.Preproc: Style(color="bright_cyan"),
-    Keyword: Style(color="bright_blue"),
-    Keyword.Type: Style(color="bright_cyan"),
-    Operator.Word: Style(color="bright_magenta"),
-    Name.Builtin: Style(color="bright_cyan"),
-    Name.Function: Style(color="bright_green"),
-    Name.Namespace: Style(color="bright_cyan", underline=True),
-    Name.Class: Style(color="bright_green", underline=True),
-    Name.Exception: Style(color="bright_cyan"),
-    Name.Decorator: Style(color="bright_magenta", bold=True),
-    Name.Variable: Style(color="bright_red"),
-    Name.Constant: Style(color="bright_red"),
-    Name.Attribute: Style(color="bright_cyan"),
-    Name.Tag: Style(color="bright_blue"),
-    String: Style(color="yellow"),
-    Number: Style(color="bright_blue"),
-    Generic.Deleted: Style(color="bright_red"),
-    Generic.Inserted: Style(color="bright_green"),
-    Generic.Heading: Style(bold=True),
-    Generic.Subheading: Style(color="bright_magenta", bold=True),
-    Generic.Prompt: Style(bold=True),
-    Generic.Error: Style(color="bright_red"),
-    Error: Style(color="red", underline=True),
-}
-
-RICH_SYNTAX_THEMES = {"ansi_light": ANSI_LIGHT, "ansi_dark": ANSI_DARK}
-NUMBERS_COLUMN_DEFAULT_PADDING = 2
-
-
-class SyntaxTheme(ABC):
-    """Base class for a syntax theme."""
-
-    @abstractmethod
-    def get_style_for_token(self, token_type: TokenType) -> Style:
-        """Get a style for a given Pygments token."""
-        raise NotImplementedError  # pragma: no cover
-
-    @abstractmethod
-    def get_background_style(self) -> Style:
-        """Get the background color."""
-        raise NotImplementedError  # pragma: no cover
-
-
-class PygmentsSyntaxTheme(SyntaxTheme):
-    """Syntax theme that delegates to Pygments theme."""
-
-    def __init__(self, theme: Union[str, Type[PygmentsStyle]]) -> None:
-        self._style_cache: Dict[TokenType, Style] = {}
-        if isinstance(theme, str):
-            try:
-                self._pygments_style_class = get_style_by_name(theme)
-            except ClassNotFound:
-                self._pygments_style_class = get_style_by_name("default")
-        else:
-            self._pygments_style_class = theme
-
-        self._background_color = self._pygments_style_class.background_color
-        self._background_style = Style(bgcolor=self._background_color)
-
-    def get_style_for_token(self, token_type: TokenType) -> Style:
-        """Get a style from a Pygments class."""
-        try:
-            return self._style_cache[token_type]
-        except KeyError:
-            try:
-                pygments_style = self._pygments_style_class.style_for_token(token_type)
-            except KeyError:
-                style = Style.null()
-            else:
-                color = pygments_style["color"]
-                bgcolor = pygments_style["bgcolor"]
-                style = Style(
-                    color="#" + color if color else "#000000",
-                    bgcolor="#" + bgcolor if bgcolor else self._background_color,
-                    bold=pygments_style["bold"],
-                    italic=pygments_style["italic"],
-                    underline=pygments_style["underline"],
-                )
-            self._style_cache[token_type] = style
-        return style
-
-    def get_background_style(self) -> Style:
-        return self._background_style
-
-
-class ANSISyntaxTheme(SyntaxTheme):
-    """Syntax theme to use standard colors."""
-
-    def __init__(self, style_map: Dict[TokenType, Style]) -> None:
-        self.style_map = style_map
-        self._missing_style = Style.null()
-        self._background_style = Style.null()
-        self._style_cache: Dict[TokenType, Style] = {}
-
-    def get_style_for_token(self, token_type: TokenType) -> Style:
-        """Look up style in the style map."""
-        try:
-            return self._style_cache[token_type]
-        except KeyError:
-            # Styles form a hierarchy
-            # We need to go from most to least specific
-            # e.g. ("foo", "bar", "baz") to ("foo", "bar")  to ("foo",)
-            get_style = self.style_map.get
-            token = tuple(token_type)
-            style = self._missing_style
-            while token:
-                _style = get_style(token)
-                if _style is not None:
-                    style = _style
-                    break
-                token = token[:-1]
-            self._style_cache[token_type] = style
-            return style
-
-    def get_background_style(self) -> Style:
-        return self._background_style
-
-
-SyntaxPosition = Tuple[int, int]
-
-
-class _SyntaxHighlightRange(NamedTuple):
-    """
-    A range to highlight in a Syntax object.
-    `start` and `end` are 2-integers tuples, where the first integer is the line number
-    (starting from 1) and the second integer is the column index (starting from 0).
-    """
-
-    style: StyleType
-    start: SyntaxPosition
-    end: SyntaxPosition
-    style_before: bool = False
-
-
-class Syntax(JupyterMixin):
-    """Construct a Syntax object to render syntax highlighted code.
-
-    Args:
-        code (str): Code to highlight.
-        lexer (Lexer | str): Lexer to use (see https://pygments.org/docs/lexers/)
-        theme (str, optional): Color theme, aka Pygments style (see https://pygments.org/docs/styles/#getting-a-list-of-available-styles). Defaults to "monokai".
-        dedent (bool, optional): Enable stripping of initial whitespace. Defaults to False.
-        line_numbers (bool, optional): Enable rendering of line numbers. Defaults to False.
-        start_line (int, optional): Starting number for line numbers. Defaults to 1.
-        line_range (Tuple[int | None, int | None], optional): If given should be a tuple of the start and end line to render.
-            A value of None in the tuple indicates the range is open in that direction.
-        highlight_lines (Set[int]): A set of line numbers to highlight.
-        code_width: Width of code to render (not including line numbers), or ``None`` to use all available width.
-        tab_size (int, optional): Size of tabs. Defaults to 4.
-        word_wrap (bool, optional): Enable word wrapping.
-        background_color (str, optional): Optional background color, or None to use theme color. Defaults to None.
-        indent_guides (bool, optional): Show indent guides. Defaults to False.
-        padding (PaddingDimensions): Padding to apply around the syntax. Defaults to 0 (no padding).
-    """
-
-    _pygments_style_class: Type[PygmentsStyle]
-    _theme: SyntaxTheme
-
-    @classmethod
-    def get_theme(cls, name: Union[str, SyntaxTheme]) -> SyntaxTheme:
-        """Get a syntax theme instance."""
-        if isinstance(name, SyntaxTheme):
-            return name
-        theme: SyntaxTheme
-        if name in RICH_SYNTAX_THEMES:
-            theme = ANSISyntaxTheme(RICH_SYNTAX_THEMES[name])
-        else:
-            theme = PygmentsSyntaxTheme(name)
-        return theme
-
-    def __init__(
-        self,
-        code: str,
-        lexer: Union[Lexer, str],
-        *,
-        theme: Union[str, SyntaxTheme] = DEFAULT_THEME,
-        dedent: bool = False,
-        line_numbers: bool = False,
-        start_line: int = 1,
-        line_range: Optional[Tuple[Optional[int], Optional[int]]] = None,
-        highlight_lines: Optional[Set[int]] = None,
-        code_width: Optional[int] = None,
-        tab_size: int = 4,
-        word_wrap: bool = False,
-        background_color: Optional[str] = None,
-        indent_guides: bool = False,
-        padding: PaddingDimensions = 0,
-    ) -> None:
-        self.code = code
-        self._lexer = lexer
-        self.dedent = dedent
-        self.line_numbers = line_numbers
-        self.start_line = start_line
-        self.line_range = line_range
-        self.highlight_lines = highlight_lines or set()
-        self.code_width = code_width
-        self.tab_size = tab_size
-        self.word_wrap = word_wrap
-        self.background_color = background_color
-        self.background_style = (
-            Style(bgcolor=background_color) if background_color else Style()
-        )
-        self.indent_guides = indent_guides
-        self.padding = padding
-
-        self._theme = self.get_theme(theme)
-        self._stylized_ranges: List[_SyntaxHighlightRange] = []
-
-    @classmethod
-    def from_path(
-        cls,
-        path: str,
-        encoding: str = "utf-8",
-        lexer: Optional[Union[Lexer, str]] = None,
-        theme: Union[str, SyntaxTheme] = DEFAULT_THEME,
-        dedent: bool = False,
-        line_numbers: bool = False,
-        line_range: Optional[Tuple[int, int]] = None,
-        start_line: int = 1,
-        highlight_lines: Optional[Set[int]] = None,
-        code_width: Optional[int] = None,
-        tab_size: int = 4,
-        word_wrap: bool = False,
-        background_color: Optional[str] = None,
-        indent_guides: bool = False,
-        padding: PaddingDimensions = 0,
-    ) -> "Syntax":
-        """Construct a Syntax object from a file.
-
-        Args:
-            path (str): Path to file to highlight.
-            encoding (str): Encoding of file.
-            lexer (str | Lexer, optional): Lexer to use. If None, lexer will be auto-detected from path/file content.
-            theme (str, optional): Color theme, aka Pygments style (see https://pygments.org/docs/styles/#getting-a-list-of-available-styles). Defaults to "emacs".
-            dedent (bool, optional): Enable stripping of initial whitespace. Defaults to True.
-            line_numbers (bool, optional): Enable rendering of line numbers. Defaults to False.
-            start_line (int, optional): Starting number for line numbers. Defaults to 1.
-            line_range (Tuple[int, int], optional): If given should be a tuple of the start and end line to render.
-            highlight_lines (Set[int]): A set of line numbers to highlight.
-            code_width: Width of code to render (not including line numbers), or ``None`` to use all available width.
-            tab_size (int, optional): Size of tabs. Defaults to 4.
-            word_wrap (bool, optional): Enable word wrapping of code.
-            background_color (str, optional): Optional background color, or None to use theme color. Defaults to None.
-            indent_guides (bool, optional): Show indent guides. Defaults to False.
-            padding (PaddingDimensions): Padding to apply around the syntax. Defaults to 0 (no padding).
-
-        Returns:
-            [Syntax]: A Syntax object that may be printed to the console
-        """
-        code = Path(path).read_text(encoding=encoding)
-
-        if not lexer:
-            lexer = cls.guess_lexer(path, code=code)
-
-        return cls(
-            code,
-            lexer,
-            theme=theme,
-            dedent=dedent,
-            line_numbers=line_numbers,
-            line_range=line_range,
-            start_line=start_line,
-            highlight_lines=highlight_lines,
-            code_width=code_width,
-            tab_size=tab_size,
-            word_wrap=word_wrap,
-            background_color=background_color,
-            indent_guides=indent_guides,
-            padding=padding,
-        )
-
-    @classmethod
-    def guess_lexer(cls, path: str, code: Optional[str] = None) -> str:
-        """Guess the alias of the Pygments lexer to use based on a path and an optional string of code.
-        If code is supplied, it will use a combination of the code and the filename to determine the
-        best lexer to use. For example, if the file is ``index.html`` and the file contains Django
-        templating syntax, then "html+django" will be returned. If the file is ``index.html``, and no
-        templating language is used, the "html" lexer will be used. If no string of code
-        is supplied, the lexer will be chosen based on the file extension..
-
-        Args:
-             path (AnyStr): The path to the file containing the code you wish to know the lexer for.
-             code (str, optional): Optional string of code that will be used as a fallback if no lexer
-                is found for the supplied path.
-
-        Returns:
-            str: The name of the Pygments lexer that best matches the supplied path/code.
-        """
-        lexer: Optional[Lexer] = None
-        lexer_name = "default"
-        if code:
-            try:
-                lexer = guess_lexer_for_filename(path, code)
-            except ClassNotFound:
-                pass
-
-        if not lexer:
-            try:
-                _, ext = os.path.splitext(path)
-                if ext:
-                    extension = ext.lstrip(".").lower()
-                    lexer = get_lexer_by_name(extension)
-            except ClassNotFound:
-                pass
-
-        if lexer:
-            if lexer.aliases:
-                lexer_name = lexer.aliases[0]
-            else:
-                lexer_name = lexer.name
-
-        return lexer_name
-
-    def _get_base_style(self) -> Style:
-        """Get the base style."""
-        default_style = self._theme.get_background_style() + self.background_style
-        return default_style
-
-    def _get_token_color(self, token_type: TokenType) -> Optional[Color]:
-        """Get a color (if any) for the given token.
-
-        Args:
-            token_type (TokenType): A token type tuple from Pygments.
-
-        Returns:
-            Optional[Color]: Color from theme, or None for no color.
-        """
-        style = self._theme.get_style_for_token(token_type)
-        return style.color
-
-    @property
-    def lexer(self) -> Optional[Lexer]:
-        """The lexer for this syntax, or None if no lexer was found.
-
-        Tries to find the lexer by name if a string was passed to the constructor.
-        """
-
-        if isinstance(self._lexer, Lexer):
-            return self._lexer
-        try:
-            return get_lexer_by_name(
-                self._lexer,
-                stripnl=False,
-                ensurenl=True,
-                tabsize=self.tab_size,
-            )
-        except ClassNotFound:
-            return None
-
-    @property
-    def default_lexer(self) -> Lexer:
-        """A Pygments Lexer to use if one is not specified or invalid."""
-        return get_lexer_by_name(
-            "text",
-            stripnl=False,
-            ensurenl=True,
-            tabsize=self.tab_size,
-        )
-
-    def highlight(
-        self,
-        code: str,
-        line_range: Optional[Tuple[Optional[int], Optional[int]]] = None,
-    ) -> Text:
-        """Highlight code and return a Text instance.
-
-        Args:
-            code (str): Code to highlight.
-            line_range(Tuple[int, int], optional): Optional line range to highlight.
-
-        Returns:
-            Text: A text instance containing highlighted syntax.
-        """
-
-        base_style = self._get_base_style()
-        justify: JustifyMethod = (
-            "default" if base_style.transparent_background else "left"
-        )
-
-        text = Text(
-            justify=justify,
-            style=base_style,
-            tab_size=self.tab_size,
-            no_wrap=not self.word_wrap,
-        )
-        _get_theme_style = self._theme.get_style_for_token
-
-        lexer = self.lexer or self.default_lexer
-
-        if lexer is None:
-            text.append(code)
-        else:
-            if line_range:
-                # More complicated path to only stylize a portion of the code
-                # This speeds up further operations as there are less spans to process
-                line_start, line_end = line_range
-
-                def line_tokenize() -> Iterable[Tuple[Any, str]]:
-                    """Split tokens to one per line."""
-                    assert lexer  # required to make MyPy happy - we know lexer is not None at this point
-
-                    for token_type, token in lexer.get_tokens(code):
-                        while token:
-                            line_token, new_line, token = token.partition("\n")
-                            yield token_type, line_token + new_line
-
-                def tokens_to_spans() -> Iterable[Tuple[str, Optional[Style]]]:
-                    """Convert tokens to spans."""
-                    tokens = iter(line_tokenize())
-                    line_no = 0
-                    _line_start = line_start - 1 if line_start else 0
-
-                    # Skip over tokens until line start
-                    while line_no < _line_start:
-                        try:
-                            _token_type, token = next(tokens)
-                        except StopIteration:
-                            break
-                        yield (token, None)
-                        if token.endswith("\n"):
-                            line_no += 1
-                    # Generate spans until line end
-                    for token_type, token in tokens:
-                        yield (token, _get_theme_style(token_type))
-                        if token.endswith("\n"):
-                            line_no += 1
-                            if line_end and line_no >= line_end:
-                                break
-
-                text.append_tokens(tokens_to_spans())
-
-            else:
-                text.append_tokens(
-                    (token, _get_theme_style(token_type))
-                    for token_type, token in lexer.get_tokens(code)
-                )
-            if self.background_color is not None:
-                text.stylize(f"on {self.background_color}")
-
-        if self._stylized_ranges:
-            self._apply_stylized_ranges(text)
-
-        return text
-
-    def stylize_range(
-        self,
-        style: StyleType,
-        start: SyntaxPosition,
-        end: SyntaxPosition,
-        style_before: bool = False,
-    ) -> None:
-        """
-        Adds a custom style on a part of the code, that will be applied to the syntax display when it's rendered.
-        Line numbers are 1-based, while column indexes are 0-based.
-
-        Args:
-            style (StyleType): The style to apply.
-            start (Tuple[int, int]): The start of the range, in the form `[line number, column index]`.
-            end (Tuple[int, int]): The end of the range, in the form `[line number, column index]`.
-            style_before (bool): Apply the style before any existing styles.
-        """
-        self._stylized_ranges.append(
-            _SyntaxHighlightRange(style, start, end, style_before)
-        )
-
-    def _get_line_numbers_color(self, blend: float = 0.3) -> Color:
-        background_style = self._theme.get_background_style() + self.background_style
-        background_color = background_style.bgcolor
-        if background_color is None or background_color.is_system_defined:
-            return Color.default()
-        foreground_color = self._get_token_color(Token.Text)
-        if foreground_color is None or foreground_color.is_system_defined:
-            return foreground_color or Color.default()
-        new_color = blend_rgb(
-            background_color.get_truecolor(),
-            foreground_color.get_truecolor(),
-            cross_fade=blend,
-        )
-        return Color.from_triplet(new_color)
-
-    @property
-    def _numbers_column_width(self) -> int:
-        """Get the number of characters used to render the numbers column."""
-        column_width = 0
-        if self.line_numbers:
-            column_width = (
-                len(str(self.start_line + self.code.count("\n")))
-                + NUMBERS_COLUMN_DEFAULT_PADDING
-            )
-        return column_width
-
-    def _get_number_styles(self, console: Console) -> Tuple[Style, Style, Style]:
-        """Get background, number, and highlight styles for line numbers."""
-        background_style = self._get_base_style()
-        if background_style.transparent_background:
-            return Style.null(), Style(dim=True), Style.null()
-        if console.color_system in ("256", "truecolor"):
-            number_style = Style.chain(
-                background_style,
-                self._theme.get_style_for_token(Token.Text),
-                Style(color=self._get_line_numbers_color()),
-                self.background_style,
-            )
-            highlight_number_style = Style.chain(
-                background_style,
-                self._theme.get_style_for_token(Token.Text),
-                Style(bold=True, color=self._get_line_numbers_color(0.9)),
-                self.background_style,
-            )
-        else:
-            number_style = background_style + Style(dim=True)
-            highlight_number_style = background_style + Style(dim=False)
-        return background_style, number_style, highlight_number_style
-
-    def __rich_measure__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "Measurement":
-        _, right, _, left = Padding.unpack(self.padding)
-        padding = left + right
-        if self.code_width is not None:
-            width = self.code_width + self._numbers_column_width + padding + 1
-            return Measurement(self._numbers_column_width, width)
-        lines = self.code.splitlines()
-        width = (
-            self._numbers_column_width
-            + padding
-            + (max(cell_len(line) for line in lines) if lines else 0)
-        )
-        if self.line_numbers:
-            width += 1
-        return Measurement(self._numbers_column_width, width)
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        segments = Segments(self._get_syntax(console, options))
-        if self.padding:
-            yield Padding(segments, style=self._get_base_style(), pad=self.padding)
-        else:
-            yield segments
-
-    def _get_syntax(
-        self,
-        console: Console,
-        options: ConsoleOptions,
-    ) -> Iterable[Segment]:
-        """
-        Get the Segments for the Syntax object, excluding any vertical/horizontal padding
-        """
-        transparent_background = self._get_base_style().transparent_background
-        code_width = (
-            (
-                (options.max_width - self._numbers_column_width - 1)
-                if self.line_numbers
-                else options.max_width
-            )
-            if self.code_width is None
-            else self.code_width
-        )
-
-        ends_on_nl, processed_code = self._process_code(self.code)
-        text = self.highlight(processed_code, self.line_range)
-
-        if not self.line_numbers and not self.word_wrap and not self.line_range:
-            if not ends_on_nl:
-                text.remove_suffix("\n")
-            # Simple case of just rendering text
-            style = (
-                self._get_base_style()
-                + self._theme.get_style_for_token(Comment)
-                + Style(dim=True)
-                + self.background_style
-            )
-            if self.indent_guides and not options.ascii_only:
-                text = text.with_indent_guides(self.tab_size, style=style)
-                text.overflow = "crop"
-            if style.transparent_background:
-                yield from console.render(
-                    text, options=options.update(width=code_width)
-                )
-            else:
-                syntax_lines = console.render_lines(
-                    text,
-                    options.update(width=code_width, height=None, justify="left"),
-                    style=self.background_style,
-                    pad=True,
-                    new_lines=True,
-                )
-                for syntax_line in syntax_lines:
-                    yield from syntax_line
-            return
-
-        start_line, end_line = self.line_range or (None, None)
-        line_offset = 0
-        if start_line:
-            line_offset = max(0, start_line - 1)
-        lines: Union[List[Text], Lines] = text.split("\n", allow_blank=ends_on_nl)
-        if self.line_range:
-            if line_offset > len(lines):
-                return
-            lines = lines[line_offset:end_line]
-
-        if self.indent_guides and not options.ascii_only:
-            style = (
-                self._get_base_style()
-                + self._theme.get_style_for_token(Comment)
-                + Style(dim=True)
-                + self.background_style
-            )
-            lines = (
-                Text("\n")
-                .join(lines)
-                .with_indent_guides(self.tab_size, style=style + Style(italic=False))
-                .split("\n", allow_blank=True)
-            )
-
-        numbers_column_width = self._numbers_column_width
-        render_options = options.update(width=code_width)
-
-        highlight_line = self.highlight_lines.__contains__
-        _Segment = Segment
-        new_line = _Segment("\n")
-
-        line_pointer = "> " if options.legacy_windows else "❱ "
-
-        (
-            background_style,
-            number_style,
-            highlight_number_style,
-        ) = self._get_number_styles(console)
-
-        for line_no, line in enumerate(lines, self.start_line + line_offset):
-            if self.word_wrap:
-                wrapped_lines = console.render_lines(
-                    line,
-                    render_options.update(height=None, justify="left"),
-                    style=background_style,
-                    pad=not transparent_background,
-                )
-            else:
-                segments = list(line.render(console, end=""))
-                if options.no_wrap:
-                    wrapped_lines = [segments]
-                else:
-                    wrapped_lines = [
-                        _Segment.adjust_line_length(
-                            segments,
-                            render_options.max_width,
-                            style=background_style,
-                            pad=not transparent_background,
-                        )
-                    ]
-
-            if self.line_numbers:
-                wrapped_line_left_pad = _Segment(
-                    " " * numbers_column_width + " ", background_style
-                )
-                for first, wrapped_line in loop_first(wrapped_lines):
-                    if first:
-                        line_column = str(line_no).rjust(numbers_column_width - 2) + " "
-                        if highlight_line(line_no):
-                            yield _Segment(line_pointer, Style(color="red"))
-                            yield _Segment(line_column, highlight_number_style)
-                        else:
-                            yield _Segment("  ", highlight_number_style)
-                            yield _Segment(line_column, number_style)
-                    else:
-                        yield wrapped_line_left_pad
-                    yield from wrapped_line
-                    yield new_line
-            else:
-                for wrapped_line in wrapped_lines:
-                    yield from wrapped_line
-                    yield new_line
-
-    def _apply_stylized_ranges(self, text: Text) -> None:
-        """
-        Apply stylized ranges to a text instance,
-        using the given code to determine the right portion to apply the style to.
-
-        Args:
-            text (Text): Text instance to apply the style to.
-        """
-        code = text.plain
-        newlines_offsets = [
-            # Let's add outer boundaries at each side of the list:
-            0,
-            # N.B. using "\n" here is much faster than using metacharacters such as "^" or "\Z":
-            *[
-                match.start() + 1
-                for match in re.finditer("\n", code, flags=re.MULTILINE)
-            ],
-            len(code) + 1,
-        ]
-
-        for stylized_range in self._stylized_ranges:
-            start = _get_code_index_for_syntax_position(
-                newlines_offsets, stylized_range.start
-            )
-            end = _get_code_index_for_syntax_position(
-                newlines_offsets, stylized_range.end
-            )
-            if start is not None and end is not None:
-                if stylized_range.style_before:
-                    text.stylize_before(stylized_range.style, start, end)
-                else:
-                    text.stylize(stylized_range.style, start, end)
-
-    def _process_code(self, code: str) -> Tuple[bool, str]:
-        """
-        Applies various processing to a raw code string
-        (normalises it so it always ends with a line return, dedents it if necessary, etc.)
-
-        Args:
-            code (str): The raw code string to process
-
-        Returns:
-            Tuple[bool, str]: the boolean indicates whether the raw code ends with a line return,
-                while the string is the processed code.
-        """
-        ends_on_nl = code.endswith("\n")
-        processed_code = code if ends_on_nl else code + "\n"
-        processed_code = (
-            textwrap.dedent(processed_code) if self.dedent else processed_code
-        )
-        processed_code = processed_code.expandtabs(self.tab_size)
-        return ends_on_nl, processed_code
-
-
-def _get_code_index_for_syntax_position(
-    newlines_offsets: Sequence[int], position: SyntaxPosition
-) -> Optional[int]:
-    """
-    Returns the index of the code string for the given positions.
-
-    Args:
-        newlines_offsets (Sequence[int]): The offset of each newline character found in the code snippet.
-        position (SyntaxPosition): The position to search for.
-
-    Returns:
-        Optional[int]: The index of the code string for this position, or `None`
-            if the given position's line number is out of range (if it's the column that is out of range
-            we silently clamp its value so that it reaches the end of the line)
-    """
-    lines_count = len(newlines_offsets)
-
-    line_number, column_index = position
-    if line_number > lines_count or len(newlines_offsets) < (line_number + 1):
-        return None  # `line_number` is out of range
-    line_index = line_number - 1
-    line_length = newlines_offsets[line_index + 1] - newlines_offsets[line_index] - 1
-    # If `column_index` is out of range: let's silently clamp it:
-    column_index = min(line_length, column_index)
-    return newlines_offsets[line_index] + column_index
-
-
-if __name__ == "__main__":  # pragma: no cover
-    import argparse
-    import sys
-
-    parser = argparse.ArgumentParser(
-        description="Render syntax to the console with Rich"
-    )
-    parser.add_argument(
-        "path",
-        metavar="PATH",
-        help="path to file, or - for stdin",
-    )
-    parser.add_argument(
-        "-c",
-        "--force-color",
-        dest="force_color",
-        action="store_true",
-        default=None,
-        help="force color for non-terminals",
-    )
-    parser.add_argument(
-        "-i",
-        "--indent-guides",
-        dest="indent_guides",
-        action="store_true",
-        default=False,
-        help="display indent guides",
-    )
-    parser.add_argument(
-        "-l",
-        "--line-numbers",
-        dest="line_numbers",
-        action="store_true",
-        help="render line numbers",
-    )
-    parser.add_argument(
-        "-w",
-        "--width",
-        type=int,
-        dest="width",
-        default=None,
-        help="width of output (default will auto-detect)",
-    )
-    parser.add_argument(
-        "-r",
-        "--wrap",
-        dest="word_wrap",
-        action="store_true",
-        default=False,
-        help="word wrap long lines",
-    )
-    parser.add_argument(
-        "-s",
-        "--soft-wrap",
-        action="store_true",
-        dest="soft_wrap",
-        default=False,
-        help="enable soft wrapping mode",
-    )
-    parser.add_argument(
-        "-t", "--theme", dest="theme", default="monokai", help="pygments theme"
-    )
-    parser.add_argument(
-        "-b",
-        "--background-color",
-        dest="background_color",
-        default=None,
-        help="Override background color",
-    )
-    parser.add_argument(
-        "-x",
-        "--lexer",
-        default=None,
-        dest="lexer_name",
-        help="Lexer name",
-    )
-    parser.add_argument(
-        "-p", "--padding", type=int, default=0, dest="padding", help="Padding"
-    )
-    parser.add_argument(
-        "--highlight-line",
-        type=int,
-        default=None,
-        dest="highlight_line",
-        help="The line number (not index!) to highlight",
-    )
-    args = parser.parse_args()
-
-    from rich.console import Console
-
-    console = Console(force_terminal=args.force_color, width=args.width)
-
-    if args.path == "-":
-        code = sys.stdin.read()
-        syntax = Syntax(
-            code=code,
-            lexer=args.lexer_name,
-            line_numbers=args.line_numbers,
-            word_wrap=args.word_wrap,
-            theme=args.theme,
-            background_color=args.background_color,
-            indent_guides=args.indent_guides,
-            padding=args.padding,
-            highlight_lines={args.highlight_line},
-        )
-    else:
-        syntax = Syntax.from_path(
-            args.path,
-            lexer=args.lexer_name,
-            line_numbers=args.line_numbers,
-            word_wrap=args.word_wrap,
-            theme=args.theme,
-            background_color=args.background_color,
-            indent_guides=args.indent_guides,
-            padding=args.padding,
-            highlight_lines={args.highlight_line},
-        )
-    console.print(syntax, soft_wrap=args.soft_wrap)
diff --git a/venv/lib/python3.10/site-packages/rich/table.py b/venv/lib/python3.10/site-packages/rich/table.py
deleted file mode 100644
index 942175dc3a53ffbfb0df3fc3f2dc921a148243fd..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/table.py
+++ /dev/null
@@ -1,1006 +0,0 @@
-from dataclasses import dataclass, field, replace
-from typing import (
-    TYPE_CHECKING,
-    Dict,
-    Iterable,
-    List,
-    NamedTuple,
-    Optional,
-    Sequence,
-    Tuple,
-    Union,
-)
-
-from . import box, errors
-from ._loop import loop_first_last, loop_last
-from ._pick import pick_bool
-from ._ratio import ratio_distribute, ratio_reduce
-from .align import VerticalAlignMethod
-from .jupyter import JupyterMixin
-from .measure import Measurement
-from .padding import Padding, PaddingDimensions
-from .protocol import is_renderable
-from .segment import Segment
-from .style import Style, StyleType
-from .text import Text, TextType
-
-if TYPE_CHECKING:
-    from .console import (
-        Console,
-        ConsoleOptions,
-        JustifyMethod,
-        OverflowMethod,
-        RenderableType,
-        RenderResult,
-    )
-
-
-@dataclass
-class Column:
-    """Defines a column within a ~Table.
-
-    Args:
-        title (Union[str, Text], optional): The title of the table rendered at the top. Defaults to None.
-        caption (Union[str, Text], optional): The table caption rendered below. Defaults to None.
-        width (int, optional): The width in characters of the table, or ``None`` to automatically fit. Defaults to None.
-        min_width (Optional[int], optional): The minimum width of the table, or ``None`` for no minimum. Defaults to None.
-        box (box.Box, optional): One of the constants in box.py used to draw the edges (see :ref:`appendix_box`), or ``None`` for no box lines. Defaults to box.HEAVY_HEAD.
-        safe_box (Optional[bool], optional): Disable box characters that don't display on windows legacy terminal with *raster* fonts. Defaults to True.
-        padding (PaddingDimensions, optional): Padding for cells (top, right, bottom, left). Defaults to (0, 1).
-        collapse_padding (bool, optional): Enable collapsing of padding around cells. Defaults to False.
-        pad_edge (bool, optional): Enable padding of edge cells. Defaults to True.
-        expand (bool, optional): Expand the table to fit the available space if ``True``, otherwise the table width will be auto-calculated. Defaults to False.
-        show_header (bool, optional): Show a header row. Defaults to True.
-        show_footer (bool, optional): Show a footer row. Defaults to False.
-        show_edge (bool, optional): Draw a box around the outside of the table. Defaults to True.
-        show_lines (bool, optional): Draw lines between every row. Defaults to False.
-        leading (int, optional): Number of blank lines between rows (precludes ``show_lines``). Defaults to 0.
-        style (Union[str, Style], optional): Default style for the table. Defaults to "none".
-        row_styles (List[Union, str], optional): Optional list of row styles, if more than one style is given then the styles will alternate. Defaults to None.
-        header_style (Union[str, Style], optional): Style of the header. Defaults to "table.header".
-        footer_style (Union[str, Style], optional): Style of the footer. Defaults to "table.footer".
-        border_style (Union[str, Style], optional): Style of the border. Defaults to None.
-        title_style (Union[str, Style], optional): Style of the title. Defaults to None.
-        caption_style (Union[str, Style], optional): Style of the caption. Defaults to None.
-        title_justify (str, optional): Justify method for title. Defaults to "center".
-        caption_justify (str, optional): Justify method for caption. Defaults to "center".
-        highlight (bool, optional): Highlight cell contents (if str). Defaults to False.
-    """
-
-    header: "RenderableType" = ""
-    """RenderableType: Renderable for the header (typically a string)"""
-
-    footer: "RenderableType" = ""
-    """RenderableType: Renderable for the footer (typically a string)"""
-
-    header_style: StyleType = ""
-    """StyleType: The style of the header."""
-
-    footer_style: StyleType = ""
-    """StyleType: The style of the footer."""
-
-    style: StyleType = ""
-    """StyleType: The style of the column."""
-
-    justify: "JustifyMethod" = "left"
-    """str: How to justify text within the column ("left", "center", "right", or "full")"""
-
-    vertical: "VerticalAlignMethod" = "top"
-    """str: How to vertically align content ("top", "middle", or "bottom")"""
-
-    overflow: "OverflowMethod" = "ellipsis"
-    """str: Overflow method."""
-
-    width: Optional[int] = None
-    """Optional[int]: Width of the column, or ``None`` (default) to auto calculate width."""
-
-    min_width: Optional[int] = None
-    """Optional[int]: Minimum width of column, or ``None`` for no minimum. Defaults to None."""
-
-    max_width: Optional[int] = None
-    """Optional[int]: Maximum width of column, or ``None`` for no maximum. Defaults to None."""
-
-    ratio: Optional[int] = None
-    """Optional[int]: Ratio to use when calculating column width, or ``None`` (default) to adapt to column contents."""
-
-    no_wrap: bool = False
-    """bool: Prevent wrapping of text within the column. Defaults to ``False``."""
-
-    highlight: bool = False
-    """bool: Apply highlighter to column. Defaults to ``False``."""
-
-    _index: int = 0
-    """Index of column."""
-
-    _cells: List["RenderableType"] = field(default_factory=list)
-
-    def copy(self) -> "Column":
-        """Return a copy of this Column."""
-        return replace(self, _cells=[])
-
-    @property
-    def cells(self) -> Iterable["RenderableType"]:
-        """Get all cells in the column, not including header."""
-        yield from self._cells
-
-    @property
-    def flexible(self) -> bool:
-        """Check if this column is flexible."""
-        return self.ratio is not None
-
-
-@dataclass
-class Row:
-    """Information regarding a row."""
-
-    style: Optional[StyleType] = None
-    """Style to apply to row."""
-
-    end_section: bool = False
-    """Indicated end of section, which will force a line beneath the row."""
-
-
-class _Cell(NamedTuple):
-    """A single cell in a table."""
-
-    style: StyleType
-    """Style to apply to cell."""
-    renderable: "RenderableType"
-    """Cell renderable."""
-    vertical: VerticalAlignMethod
-    """Cell vertical alignment."""
-
-
-class Table(JupyterMixin):
-    """A console renderable to draw a table.
-
-    Args:
-        *headers (Union[Column, str]): Column headers, either as a string, or :class:`~rich.table.Column` instance.
-        title (Union[str, Text], optional): The title of the table rendered at the top. Defaults to None.
-        caption (Union[str, Text], optional): The table caption rendered below. Defaults to None.
-        width (int, optional): The width in characters of the table, or ``None`` to automatically fit. Defaults to None.
-        min_width (Optional[int], optional): The minimum width of the table, or ``None`` for no minimum. Defaults to None.
-        box (box.Box, optional): One of the constants in box.py used to draw the edges (see :ref:`appendix_box`), or ``None`` for no box lines. Defaults to box.HEAVY_HEAD.
-        safe_box (Optional[bool], optional): Disable box characters that don't display on windows legacy terminal with *raster* fonts. Defaults to True.
-        padding (PaddingDimensions, optional): Padding for cells (top, right, bottom, left). Defaults to (0, 1).
-        collapse_padding (bool, optional): Enable collapsing of padding around cells. Defaults to False.
-        pad_edge (bool, optional): Enable padding of edge cells. Defaults to True.
-        expand (bool, optional): Expand the table to fit the available space if ``True``, otherwise the table width will be auto-calculated. Defaults to False.
-        show_header (bool, optional): Show a header row. Defaults to True.
-        show_footer (bool, optional): Show a footer row. Defaults to False.
-        show_edge (bool, optional): Draw a box around the outside of the table. Defaults to True.
-        show_lines (bool, optional): Draw lines between every row. Defaults to False.
-        leading (int, optional): Number of blank lines between rows (precludes ``show_lines``). Defaults to 0.
-        style (Union[str, Style], optional): Default style for the table. Defaults to "none".
-        row_styles (List[Union, str], optional): Optional list of row styles, if more than one style is given then the styles will alternate. Defaults to None.
-        header_style (Union[str, Style], optional): Style of the header. Defaults to "table.header".
-        footer_style (Union[str, Style], optional): Style of the footer. Defaults to "table.footer".
-        border_style (Union[str, Style], optional): Style of the border. Defaults to None.
-        title_style (Union[str, Style], optional): Style of the title. Defaults to None.
-        caption_style (Union[str, Style], optional): Style of the caption. Defaults to None.
-        title_justify (str, optional): Justify method for title. Defaults to "center".
-        caption_justify (str, optional): Justify method for caption. Defaults to "center".
-        highlight (bool, optional): Highlight cell contents (if str). Defaults to False.
-    """
-
-    columns: List[Column]
-    rows: List[Row]
-
-    def __init__(
-        self,
-        *headers: Union[Column, str],
-        title: Optional[TextType] = None,
-        caption: Optional[TextType] = None,
-        width: Optional[int] = None,
-        min_width: Optional[int] = None,
-        box: Optional[box.Box] = box.HEAVY_HEAD,
-        safe_box: Optional[bool] = None,
-        padding: PaddingDimensions = (0, 1),
-        collapse_padding: bool = False,
-        pad_edge: bool = True,
-        expand: bool = False,
-        show_header: bool = True,
-        show_footer: bool = False,
-        show_edge: bool = True,
-        show_lines: bool = False,
-        leading: int = 0,
-        style: StyleType = "none",
-        row_styles: Optional[Iterable[StyleType]] = None,
-        header_style: Optional[StyleType] = "table.header",
-        footer_style: Optional[StyleType] = "table.footer",
-        border_style: Optional[StyleType] = None,
-        title_style: Optional[StyleType] = None,
-        caption_style: Optional[StyleType] = None,
-        title_justify: "JustifyMethod" = "center",
-        caption_justify: "JustifyMethod" = "center",
-        highlight: bool = False,
-    ) -> None:
-        self.columns: List[Column] = []
-        self.rows: List[Row] = []
-        self.title = title
-        self.caption = caption
-        self.width = width
-        self.min_width = min_width
-        self.box = box
-        self.safe_box = safe_box
-        self._padding = Padding.unpack(padding)
-        self.pad_edge = pad_edge
-        self._expand = expand
-        self.show_header = show_header
-        self.show_footer = show_footer
-        self.show_edge = show_edge
-        self.show_lines = show_lines
-        self.leading = leading
-        self.collapse_padding = collapse_padding
-        self.style = style
-        self.header_style = header_style or ""
-        self.footer_style = footer_style or ""
-        self.border_style = border_style
-        self.title_style = title_style
-        self.caption_style = caption_style
-        self.title_justify: "JustifyMethod" = title_justify
-        self.caption_justify: "JustifyMethod" = caption_justify
-        self.highlight = highlight
-        self.row_styles: Sequence[StyleType] = list(row_styles or [])
-        append_column = self.columns.append
-        for header in headers:
-            if isinstance(header, str):
-                self.add_column(header=header)
-            else:
-                header._index = len(self.columns)
-                append_column(header)
-
-    @classmethod
-    def grid(
-        cls,
-        *headers: Union[Column, str],
-        padding: PaddingDimensions = 0,
-        collapse_padding: bool = True,
-        pad_edge: bool = False,
-        expand: bool = False,
-    ) -> "Table":
-        """Get a table with no lines, headers, or footer.
-
-        Args:
-            *headers (Union[Column, str]): Column headers, either as a string, or :class:`~rich.table.Column` instance.
-            padding (PaddingDimensions, optional): Get padding around cells. Defaults to 0.
-            collapse_padding (bool, optional): Enable collapsing of padding around cells. Defaults to True.
-            pad_edge (bool, optional): Enable padding around edges of table. Defaults to False.
-            expand (bool, optional): Expand the table to fit the available space if ``True``, otherwise the table width will be auto-calculated. Defaults to False.
-
-        Returns:
-            Table: A table instance.
-        """
-        return cls(
-            *headers,
-            box=None,
-            padding=padding,
-            collapse_padding=collapse_padding,
-            show_header=False,
-            show_footer=False,
-            show_edge=False,
-            pad_edge=pad_edge,
-            expand=expand,
-        )
-
-    @property
-    def expand(self) -> bool:
-        """Setting a non-None self.width implies expand."""
-        return self._expand or self.width is not None
-
-    @expand.setter
-    def expand(self, expand: bool) -> None:
-        """Set expand."""
-        self._expand = expand
-
-    @property
-    def _extra_width(self) -> int:
-        """Get extra width to add to cell content."""
-        width = 0
-        if self.box and self.show_edge:
-            width += 2
-        if self.box:
-            width += len(self.columns) - 1
-        return width
-
-    @property
-    def row_count(self) -> int:
-        """Get the current number of rows."""
-        return len(self.rows)
-
-    def get_row_style(self, console: "Console", index: int) -> StyleType:
-        """Get the current row style."""
-        style = Style.null()
-        if self.row_styles:
-            style += console.get_style(self.row_styles[index % len(self.row_styles)])
-        row_style = self.rows[index].style
-        if row_style is not None:
-            style += console.get_style(row_style)
-        return style
-
-    def __rich_measure__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> Measurement:
-        max_width = options.max_width
-        if self.width is not None:
-            max_width = self.width
-        if max_width < 0:
-            return Measurement(0, 0)
-
-        extra_width = self._extra_width
-        max_width = sum(
-            self._calculate_column_widths(
-                console, options.update_width(max_width - extra_width)
-            )
-        )
-        _measure_column = self._measure_column
-
-        measurements = [
-            _measure_column(console, options.update_width(max_width), column)
-            for column in self.columns
-        ]
-        minimum_width = (
-            sum(measurement.minimum for measurement in measurements) + extra_width
-        )
-        maximum_width = (
-            sum(measurement.maximum for measurement in measurements) + extra_width
-            if (self.width is None)
-            else self.width
-        )
-        measurement = Measurement(minimum_width, maximum_width)
-        measurement = measurement.clamp(self.min_width)
-        return measurement
-
-    @property
-    def padding(self) -> Tuple[int, int, int, int]:
-        """Get cell padding."""
-        return self._padding
-
-    @padding.setter
-    def padding(self, padding: PaddingDimensions) -> "Table":
-        """Set cell padding."""
-        self._padding = Padding.unpack(padding)
-        return self
-
-    def add_column(
-        self,
-        header: "RenderableType" = "",
-        footer: "RenderableType" = "",
-        *,
-        header_style: Optional[StyleType] = None,
-        highlight: Optional[bool] = None,
-        footer_style: Optional[StyleType] = None,
-        style: Optional[StyleType] = None,
-        justify: "JustifyMethod" = "left",
-        vertical: "VerticalAlignMethod" = "top",
-        overflow: "OverflowMethod" = "ellipsis",
-        width: Optional[int] = None,
-        min_width: Optional[int] = None,
-        max_width: Optional[int] = None,
-        ratio: Optional[int] = None,
-        no_wrap: bool = False,
-    ) -> None:
-        """Add a column to the table.
-
-        Args:
-            header (RenderableType, optional): Text or renderable for the header.
-                Defaults to "".
-            footer (RenderableType, optional): Text or renderable for the footer.
-                Defaults to "".
-            header_style (Union[str, Style], optional): Style for the header, or None for default. Defaults to None.
-            highlight (bool, optional): Whether to highlight the text. The default of None uses the value of the table (self) object.
-            footer_style (Union[str, Style], optional): Style for the footer, or None for default. Defaults to None.
-            style (Union[str, Style], optional): Style for the column cells, or None for default. Defaults to None.
-            justify (JustifyMethod, optional): Alignment for cells. Defaults to "left".
-            vertical (VerticalAlignMethod, optional): Vertical alignment, one of "top", "middle", or "bottom". Defaults to "top".
-            overflow (OverflowMethod): Overflow method: "crop", "fold", "ellipsis". Defaults to "ellipsis".
-            width (int, optional): Desired width of column in characters, or None to fit to contents. Defaults to None.
-            min_width (Optional[int], optional): Minimum width of column, or ``None`` for no minimum. Defaults to None.
-            max_width (Optional[int], optional): Maximum width of column, or ``None`` for no maximum. Defaults to None.
-            ratio (int, optional): Flexible ratio for the column (requires ``Table.expand`` or ``Table.width``). Defaults to None.
-            no_wrap (bool, optional): Set to ``True`` to disable wrapping of this column.
-        """
-
-        column = Column(
-            _index=len(self.columns),
-            header=header,
-            footer=footer,
-            header_style=header_style or "",
-            highlight=highlight if highlight is not None else self.highlight,
-            footer_style=footer_style or "",
-            style=style or "",
-            justify=justify,
-            vertical=vertical,
-            overflow=overflow,
-            width=width,
-            min_width=min_width,
-            max_width=max_width,
-            ratio=ratio,
-            no_wrap=no_wrap,
-        )
-        self.columns.append(column)
-
-    def add_row(
-        self,
-        *renderables: Optional["RenderableType"],
-        style: Optional[StyleType] = None,
-        end_section: bool = False,
-    ) -> None:
-        """Add a row of renderables.
-
-        Args:
-            *renderables (None or renderable): Each cell in a row must be a renderable object (including str),
-                or ``None`` for a blank cell.
-            style (StyleType, optional): An optional style to apply to the entire row. Defaults to None.
-            end_section (bool, optional): End a section and draw a line. Defaults to False.
-
-        Raises:
-            errors.NotRenderableError: If you add something that can't be rendered.
-        """
-
-        def add_cell(column: Column, renderable: "RenderableType") -> None:
-            column._cells.append(renderable)
-
-        cell_renderables: List[Optional["RenderableType"]] = list(renderables)
-
-        columns = self.columns
-        if len(cell_renderables) < len(columns):
-            cell_renderables = [
-                *cell_renderables,
-                *[None] * (len(columns) - len(cell_renderables)),
-            ]
-        for index, renderable in enumerate(cell_renderables):
-            if index == len(columns):
-                column = Column(_index=index, highlight=self.highlight)
-                for _ in self.rows:
-                    add_cell(column, Text(""))
-                self.columns.append(column)
-            else:
-                column = columns[index]
-            if renderable is None:
-                add_cell(column, "")
-            elif is_renderable(renderable):
-                add_cell(column, renderable)
-            else:
-                raise errors.NotRenderableError(
-                    f"unable to render {type(renderable).__name__}; a string or other renderable object is required"
-                )
-        self.rows.append(Row(style=style, end_section=end_section))
-
-    def add_section(self) -> None:
-        """Add a new section (draw a line after current row)."""
-
-        if self.rows:
-            self.rows[-1].end_section = True
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        if not self.columns:
-            yield Segment("\n")
-            return
-
-        max_width = options.max_width
-        if self.width is not None:
-            max_width = self.width
-
-        extra_width = self._extra_width
-        widths = self._calculate_column_widths(
-            console, options.update_width(max_width - extra_width)
-        )
-        table_width = sum(widths) + extra_width
-
-        render_options = options.update(
-            width=table_width, highlight=self.highlight, height=None
-        )
-
-        def render_annotation(
-            text: TextType, style: StyleType, justify: "JustifyMethod" = "center"
-        ) -> "RenderResult":
-            render_text = (
-                console.render_str(text, style=style, highlight=False)
-                if isinstance(text, str)
-                else text
-            )
-            return console.render(
-                render_text, options=render_options.update(justify=justify)
-            )
-
-        if self.title:
-            yield from render_annotation(
-                self.title,
-                style=Style.pick_first(self.title_style, "table.title"),
-                justify=self.title_justify,
-            )
-        yield from self._render(console, render_options, widths)
-        if self.caption:
-            yield from render_annotation(
-                self.caption,
-                style=Style.pick_first(self.caption_style, "table.caption"),
-                justify=self.caption_justify,
-            )
-
-    def _calculate_column_widths(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> List[int]:
-        """Calculate the widths of each column, including padding, not including borders."""
-        max_width = options.max_width
-        columns = self.columns
-        width_ranges = [
-            self._measure_column(console, options, column) for column in columns
-        ]
-        widths = [_range.maximum or 1 for _range in width_ranges]
-        get_padding_width = self._get_padding_width
-        extra_width = self._extra_width
-        if self.expand:
-            ratios = [col.ratio or 0 for col in columns if col.flexible]
-            if any(ratios):
-                fixed_widths = [
-                    0 if column.flexible else _range.maximum
-                    for _range, column in zip(width_ranges, columns)
-                ]
-                flex_minimum = [
-                    (column.width or 1) + get_padding_width(column._index)
-                    for column in columns
-                    if column.flexible
-                ]
-                flexible_width = max_width - sum(fixed_widths)
-                flex_widths = ratio_distribute(flexible_width, ratios, flex_minimum)
-                iter_flex_widths = iter(flex_widths)
-                for index, column in enumerate(columns):
-                    if column.flexible:
-                        widths[index] = fixed_widths[index] + next(iter_flex_widths)
-        table_width = sum(widths)
-
-        if table_width > max_width:
-            widths = self._collapse_widths(
-                widths,
-                [(column.width is None and not column.no_wrap) for column in columns],
-                max_width,
-            )
-            table_width = sum(widths)
-            # last resort, reduce columns evenly
-            if table_width > max_width:
-                excess_width = table_width - max_width
-                widths = ratio_reduce(excess_width, [1] * len(widths), widths, widths)
-                table_width = sum(widths)
-
-            width_ranges = [
-                self._measure_column(console, options.update_width(width), column)
-                for width, column in zip(widths, columns)
-            ]
-            widths = [_range.maximum or 0 for _range in width_ranges]
-
-        if (table_width < max_width and self.expand) or (
-            self.min_width is not None and table_width < (self.min_width - extra_width)
-        ):
-            _max_width = (
-                max_width
-                if self.min_width is None
-                else min(self.min_width - extra_width, max_width)
-            )
-            pad_widths = ratio_distribute(_max_width - table_width, widths)
-            widths = [_width + pad for _width, pad in zip(widths, pad_widths)]
-
-        return widths
-
-    @classmethod
-    def _collapse_widths(
-        cls, widths: List[int], wrapable: List[bool], max_width: int
-    ) -> List[int]:
-        """Reduce widths so that the total is under max_width.
-
-        Args:
-            widths (List[int]): List of widths.
-            wrapable (List[bool]): List of booleans that indicate if a column may shrink.
-            max_width (int): Maximum width to reduce to.
-
-        Returns:
-            List[int]: A new list of widths.
-        """
-        total_width = sum(widths)
-        excess_width = total_width - max_width
-        if any(wrapable):
-            while total_width and excess_width > 0:
-                max_column = max(
-                    width for width, allow_wrap in zip(widths, wrapable) if allow_wrap
-                )
-                second_max_column = max(
-                    width if allow_wrap and width != max_column else 0
-                    for width, allow_wrap in zip(widths, wrapable)
-                )
-                column_difference = max_column - second_max_column
-                ratios = [
-                    (1 if (width == max_column and allow_wrap) else 0)
-                    for width, allow_wrap in zip(widths, wrapable)
-                ]
-                if not any(ratios) or not column_difference:
-                    break
-                max_reduce = [min(excess_width, column_difference)] * len(widths)
-                widths = ratio_reduce(excess_width, ratios, max_reduce, widths)
-
-                total_width = sum(widths)
-                excess_width = total_width - max_width
-        return widths
-
-    def _get_cells(
-        self, console: "Console", column_index: int, column: Column
-    ) -> Iterable[_Cell]:
-        """Get all the cells with padding and optional header."""
-
-        collapse_padding = self.collapse_padding
-        pad_edge = self.pad_edge
-        padding = self.padding
-        any_padding = any(padding)
-
-        first_column = column_index == 0
-        last_column = column_index == len(self.columns) - 1
-
-        _padding_cache: Dict[Tuple[bool, bool], Tuple[int, int, int, int]] = {}
-
-        def get_padding(first_row: bool, last_row: bool) -> Tuple[int, int, int, int]:
-            cached = _padding_cache.get((first_row, last_row))
-            if cached:
-                return cached
-            top, right, bottom, left = padding
-
-            if collapse_padding:
-                if not first_column:
-                    left = max(0, left - right)
-                if not last_row:
-                    bottom = max(0, top - bottom)
-
-            if not pad_edge:
-                if first_column:
-                    left = 0
-                if last_column:
-                    right = 0
-                if first_row:
-                    top = 0
-                if last_row:
-                    bottom = 0
-            _padding = (top, right, bottom, left)
-            _padding_cache[(first_row, last_row)] = _padding
-            return _padding
-
-        raw_cells: List[Tuple[StyleType, "RenderableType"]] = []
-        _append = raw_cells.append
-        get_style = console.get_style
-        if self.show_header:
-            header_style = get_style(self.header_style or "") + get_style(
-                column.header_style
-            )
-            _append((header_style, column.header))
-        cell_style = get_style(column.style or "")
-        for cell in column.cells:
-            _append((cell_style, cell))
-        if self.show_footer:
-            footer_style = get_style(self.footer_style or "") + get_style(
-                column.footer_style
-            )
-            _append((footer_style, column.footer))
-
-        if any_padding:
-            _Padding = Padding
-            for first, last, (style, renderable) in loop_first_last(raw_cells):
-                yield _Cell(
-                    style,
-                    _Padding(renderable, get_padding(first, last)),
-                    getattr(renderable, "vertical", None) or column.vertical,
-                )
-        else:
-            for style, renderable in raw_cells:
-                yield _Cell(
-                    style,
-                    renderable,
-                    getattr(renderable, "vertical", None) or column.vertical,
-                )
-
-    def _get_padding_width(self, column_index: int) -> int:
-        """Get extra width from padding."""
-        _, pad_right, _, pad_left = self.padding
-        if self.collapse_padding:
-            if column_index > 0:
-                pad_left = max(0, pad_left - pad_right)
-        return pad_left + pad_right
-
-    def _measure_column(
-        self,
-        console: "Console",
-        options: "ConsoleOptions",
-        column: Column,
-    ) -> Measurement:
-        """Get the minimum and maximum width of the column."""
-
-        max_width = options.max_width
-        if max_width < 1:
-            return Measurement(0, 0)
-
-        padding_width = self._get_padding_width(column._index)
-
-        if column.width is not None:
-            # Fixed width column
-            return Measurement(
-                column.width + padding_width, column.width + padding_width
-            ).with_maximum(max_width)
-        # Flexible column, we need to measure contents
-        min_widths: List[int] = []
-        max_widths: List[int] = []
-        append_min = min_widths.append
-        append_max = max_widths.append
-        get_render_width = Measurement.get
-        for cell in self._get_cells(console, column._index, column):
-            _min, _max = get_render_width(console, options, cell.renderable)
-            append_min(_min)
-            append_max(_max)
-
-        measurement = Measurement(
-            max(min_widths) if min_widths else 1,
-            max(max_widths) if max_widths else max_width,
-        ).with_maximum(max_width)
-        measurement = measurement.clamp(
-            None if column.min_width is None else column.min_width + padding_width,
-            None if column.max_width is None else column.max_width + padding_width,
-        )
-        return measurement
-
-    def _render(
-        self, console: "Console", options: "ConsoleOptions", widths: List[int]
-    ) -> "RenderResult":
-        table_style = console.get_style(self.style or "")
-
-        border_style = table_style + console.get_style(self.border_style or "")
-        _column_cells = (
-            self._get_cells(console, column_index, column)
-            for column_index, column in enumerate(self.columns)
-        )
-        row_cells: List[Tuple[_Cell, ...]] = list(zip(*_column_cells))
-        _box = (
-            self.box.substitute(
-                options, safe=pick_bool(self.safe_box, console.safe_box)
-            )
-            if self.box
-            else None
-        )
-        _box = _box.get_plain_headed_box() if _box and not self.show_header else _box
-
-        new_line = Segment.line()
-
-        columns = self.columns
-        show_header = self.show_header
-        show_footer = self.show_footer
-        show_edge = self.show_edge
-        show_lines = self.show_lines
-        leading = self.leading
-
-        _Segment = Segment
-        if _box:
-            box_segments = [
-                (
-                    _Segment(_box.head_left, border_style),
-                    _Segment(_box.head_right, border_style),
-                    _Segment(_box.head_vertical, border_style),
-                ),
-                (
-                    _Segment(_box.mid_left, border_style),
-                    _Segment(_box.mid_right, border_style),
-                    _Segment(_box.mid_vertical, border_style),
-                ),
-                (
-                    _Segment(_box.foot_left, border_style),
-                    _Segment(_box.foot_right, border_style),
-                    _Segment(_box.foot_vertical, border_style),
-                ),
-            ]
-            if show_edge:
-                yield _Segment(_box.get_top(widths), border_style)
-                yield new_line
-        else:
-            box_segments = []
-
-        get_row_style = self.get_row_style
-        get_style = console.get_style
-
-        for index, (first, last, row_cell) in enumerate(loop_first_last(row_cells)):
-            header_row = first and show_header
-            footer_row = last and show_footer
-            row = (
-                self.rows[index - show_header]
-                if (not header_row and not footer_row)
-                else None
-            )
-            max_height = 1
-            cells: List[List[List[Segment]]] = []
-            if header_row or footer_row:
-                row_style = Style.null()
-            else:
-                row_style = get_style(
-                    get_row_style(console, index - 1 if show_header else index)
-                )
-            for width, cell, column in zip(widths, row_cell, columns):
-                render_options = options.update(
-                    width=width,
-                    justify=column.justify,
-                    no_wrap=column.no_wrap,
-                    overflow=column.overflow,
-                    height=None,
-                    highlight=column.highlight,
-                )
-                lines = console.render_lines(
-                    cell.renderable,
-                    render_options,
-                    style=get_style(cell.style) + row_style,
-                )
-                max_height = max(max_height, len(lines))
-                cells.append(lines)
-
-            row_height = max(len(cell) for cell in cells)
-
-            def align_cell(
-                cell: List[List[Segment]],
-                vertical: "VerticalAlignMethod",
-                width: int,
-                style: Style,
-            ) -> List[List[Segment]]:
-                if header_row:
-                    vertical = "bottom"
-                elif footer_row:
-                    vertical = "top"
-
-                if vertical == "top":
-                    return _Segment.align_top(cell, width, row_height, style)
-                elif vertical == "middle":
-                    return _Segment.align_middle(cell, width, row_height, style)
-                return _Segment.align_bottom(cell, width, row_height, style)
-
-            cells[:] = [
-                _Segment.set_shape(
-                    align_cell(
-                        cell,
-                        _cell.vertical,
-                        width,
-                        get_style(_cell.style) + row_style,
-                    ),
-                    width,
-                    max_height,
-                )
-                for width, _cell, cell, column in zip(widths, row_cell, cells, columns)
-            ]
-
-            if _box:
-                if last and show_footer:
-                    yield _Segment(
-                        _box.get_row(widths, "foot", edge=show_edge), border_style
-                    )
-                    yield new_line
-                left, right, _divider = box_segments[0 if first else (2 if last else 1)]
-
-                # If the column divider is whitespace also style it with the row background
-                divider = (
-                    _divider
-                    if _divider.text.strip()
-                    else _Segment(
-                        _divider.text, row_style.background_style + _divider.style
-                    )
-                )
-                for line_no in range(max_height):
-                    if show_edge:
-                        yield left
-                    for last_cell, rendered_cell in loop_last(cells):
-                        yield from rendered_cell[line_no]
-                        if not last_cell:
-                            yield divider
-                    if show_edge:
-                        yield right
-                    yield new_line
-            else:
-                for line_no in range(max_height):
-                    for rendered_cell in cells:
-                        yield from rendered_cell[line_no]
-                    yield new_line
-            if _box and first and show_header:
-                yield _Segment(
-                    _box.get_row(widths, "head", edge=show_edge), border_style
-                )
-                yield new_line
-            end_section = row and row.end_section
-            if _box and (show_lines or leading or end_section):
-                if (
-                    not last
-                    and not (show_footer and index >= len(row_cells) - 2)
-                    and not (show_header and header_row)
-                ):
-                    if leading:
-                        yield _Segment(
-                            _box.get_row(widths, "mid", edge=show_edge) * leading,
-                            border_style,
-                        )
-                    else:
-                        yield _Segment(
-                            _box.get_row(widths, "row", edge=show_edge), border_style
-                        )
-                    yield new_line
-
-        if _box and show_edge:
-            yield _Segment(_box.get_bottom(widths), border_style)
-            yield new_line
-
-
-if __name__ == "__main__":  # pragma: no cover
-    from rich.console import Console
-    from rich.highlighter import ReprHighlighter
-
-    from ._timer import timer
-
-    with timer("Table render"):
-        table = Table(
-            title="Star Wars Movies",
-            caption="Rich example table",
-            caption_justify="right",
-        )
-
-        table.add_column(
-            "Released", header_style="bright_cyan", style="cyan", no_wrap=True
-        )
-        table.add_column("Title", style="magenta")
-        table.add_column("Box Office", justify="right", style="green")
-
-        table.add_row(
-            "Dec 20, 2019",
-            "Star Wars: The Rise of Skywalker",
-            "$952,110,690",
-        )
-        table.add_row("May 25, 2018", "Solo: A Star Wars Story", "$393,151,347")
-        table.add_row(
-            "Dec 15, 2017",
-            "Star Wars Ep. V111: The Last Jedi",
-            "$1,332,539,889",
-            style="on black",
-            end_section=True,
-        )
-        table.add_row(
-            "Dec 16, 2016",
-            "Rogue One: A Star Wars Story",
-            "$1,332,439,889",
-        )
-
-        def header(text: str) -> None:
-            console.print()
-            console.rule(highlight(text))
-            console.print()
-
-        console = Console()
-        highlight = ReprHighlighter()
-        header("Example Table")
-        console.print(table, justify="center")
-
-        table.expand = True
-        header("expand=True")
-        console.print(table)
-
-        table.width = 50
-        header("width=50")
-
-        console.print(table, justify="center")
-
-        table.width = None
-        table.expand = False
-        table.row_styles = ["dim", "none"]
-        header("row_styles=['dim', 'none']")
-
-        console.print(table, justify="center")
-
-        table.width = None
-        table.expand = False
-        table.row_styles = ["dim", "none"]
-        table.leading = 1
-        header("leading=1, row_styles=['dim', 'none']")
-        console.print(table, justify="center")
-
-        table.width = None
-        table.expand = False
-        table.row_styles = ["dim", "none"]
-        table.show_lines = True
-        table.leading = 0
-        header("show_lines=True, row_styles=['dim', 'none']")
-        console.print(table, justify="center")
diff --git a/venv/lib/python3.10/site-packages/rich/terminal_theme.py b/venv/lib/python3.10/site-packages/rich/terminal_theme.py
deleted file mode 100644
index 565e9d960f8604c487e063ad9ed3f6f63027f3b4..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/terminal_theme.py
+++ /dev/null
@@ -1,153 +0,0 @@
-from typing import List, Optional, Tuple
-
-from .color_triplet import ColorTriplet
-from .palette import Palette
-
-_ColorTuple = Tuple[int, int, int]
-
-
-class TerminalTheme:
-    """A color theme used when exporting console content.
-
-    Args:
-        background (Tuple[int, int, int]): The background color.
-        foreground (Tuple[int, int, int]): The foreground (text) color.
-        normal (List[Tuple[int, int, int]]): A list of 8 normal intensity colors.
-        bright (List[Tuple[int, int, int]], optional): A list of 8 bright colors, or None
-            to repeat normal intensity. Defaults to None.
-    """
-
-    def __init__(
-        self,
-        background: _ColorTuple,
-        foreground: _ColorTuple,
-        normal: List[_ColorTuple],
-        bright: Optional[List[_ColorTuple]] = None,
-    ) -> None:
-        self.background_color = ColorTriplet(*background)
-        self.foreground_color = ColorTriplet(*foreground)
-        self.ansi_colors = Palette(normal + (bright or normal))
-
-
-DEFAULT_TERMINAL_THEME = TerminalTheme(
-    (255, 255, 255),
-    (0, 0, 0),
-    [
-        (0, 0, 0),
-        (128, 0, 0),
-        (0, 128, 0),
-        (128, 128, 0),
-        (0, 0, 128),
-        (128, 0, 128),
-        (0, 128, 128),
-        (192, 192, 192),
-    ],
-    [
-        (128, 128, 128),
-        (255, 0, 0),
-        (0, 255, 0),
-        (255, 255, 0),
-        (0, 0, 255),
-        (255, 0, 255),
-        (0, 255, 255),
-        (255, 255, 255),
-    ],
-)
-
-MONOKAI = TerminalTheme(
-    (12, 12, 12),
-    (217, 217, 217),
-    [
-        (26, 26, 26),
-        (244, 0, 95),
-        (152, 224, 36),
-        (253, 151, 31),
-        (157, 101, 255),
-        (244, 0, 95),
-        (88, 209, 235),
-        (196, 197, 181),
-        (98, 94, 76),
-    ],
-    [
-        (244, 0, 95),
-        (152, 224, 36),
-        (224, 213, 97),
-        (157, 101, 255),
-        (244, 0, 95),
-        (88, 209, 235),
-        (246, 246, 239),
-    ],
-)
-DIMMED_MONOKAI = TerminalTheme(
-    (25, 25, 25),
-    (185, 188, 186),
-    [
-        (58, 61, 67),
-        (190, 63, 72),
-        (135, 154, 59),
-        (197, 166, 53),
-        (79, 118, 161),
-        (133, 92, 141),
-        (87, 143, 164),
-        (185, 188, 186),
-        (136, 137, 135),
-    ],
-    [
-        (251, 0, 31),
-        (15, 114, 47),
-        (196, 112, 51),
-        (24, 109, 227),
-        (251, 0, 103),
-        (46, 112, 109),
-        (253, 255, 185),
-    ],
-)
-NIGHT_OWLISH = TerminalTheme(
-    (255, 255, 255),
-    (64, 63, 83),
-    [
-        (1, 22, 39),
-        (211, 66, 62),
-        (42, 162, 152),
-        (218, 170, 1),
-        (72, 118, 214),
-        (64, 63, 83),
-        (8, 145, 106),
-        (122, 129, 129),
-        (122, 129, 129),
-    ],
-    [
-        (247, 110, 110),
-        (73, 208, 197),
-        (218, 194, 107),
-        (92, 167, 228),
-        (105, 112, 152),
-        (0, 201, 144),
-        (152, 159, 177),
-    ],
-)
-
-SVG_EXPORT_THEME = TerminalTheme(
-    (41, 41, 41),
-    (197, 200, 198),
-    [
-        (75, 78, 85),
-        (204, 85, 90),
-        (152, 168, 75),
-        (208, 179, 68),
-        (96, 138, 177),
-        (152, 114, 159),
-        (104, 160, 179),
-        (197, 200, 198),
-        (154, 155, 153),
-    ],
-    [
-        (255, 38, 39),
-        (0, 130, 61),
-        (208, 132, 66),
-        (25, 132, 233),
-        (255, 44, 122),
-        (57, 130, 128),
-        (253, 253, 197),
-    ],
-)
diff --git a/venv/lib/python3.10/site-packages/rich/text.py b/venv/lib/python3.10/site-packages/rich/text.py
deleted file mode 100644
index b57d77c276c5e52f11f736ef6ab3c58b31e23045..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/text.py
+++ /dev/null
@@ -1,1361 +0,0 @@
-import re
-from functools import partial, reduce
-from math import gcd
-from operator import itemgetter
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    Iterable,
-    List,
-    NamedTuple,
-    Optional,
-    Pattern,
-    Tuple,
-    Union,
-)
-
-from ._loop import loop_last
-from ._pick import pick_bool
-from ._wrap import divide_line
-from .align import AlignMethod
-from .cells import cell_len, set_cell_size
-from .containers import Lines
-from .control import strip_control_codes
-from .emoji import EmojiVariant
-from .jupyter import JupyterMixin
-from .measure import Measurement
-from .segment import Segment
-from .style import Style, StyleType
-
-if TYPE_CHECKING:  # pragma: no cover
-    from .console import Console, ConsoleOptions, JustifyMethod, OverflowMethod
-
-DEFAULT_JUSTIFY: "JustifyMethod" = "default"
-DEFAULT_OVERFLOW: "OverflowMethod" = "fold"
-
-
-_re_whitespace = re.compile(r"\s+$")
-
-TextType = Union[str, "Text"]
-"""A plain string or a :class:`Text` instance."""
-
-GetStyleCallable = Callable[[str], Optional[StyleType]]
-
-
-class Span(NamedTuple):
-    """A marked up region in some text."""
-
-    start: int
-    """Span start index."""
-    end: int
-    """Span end index."""
-    style: Union[str, Style]
-    """Style associated with the span."""
-
-    def __repr__(self) -> str:
-        return f"Span({self.start}, {self.end}, {self.style!r})"
-
-    def __bool__(self) -> bool:
-        return self.end > self.start
-
-    def split(self, offset: int) -> Tuple["Span", Optional["Span"]]:
-        """Split a span in to 2 from a given offset."""
-
-        if offset < self.start:
-            return self, None
-        if offset >= self.end:
-            return self, None
-
-        start, end, style = self
-        span1 = Span(start, min(end, offset), style)
-        span2 = Span(span1.end, end, style)
-        return span1, span2
-
-    def move(self, offset: int) -> "Span":
-        """Move start and end by a given offset.
-
-        Args:
-            offset (int): Number of characters to add to start and end.
-
-        Returns:
-            TextSpan: A new TextSpan with adjusted position.
-        """
-        start, end, style = self
-        return Span(start + offset, end + offset, style)
-
-    def right_crop(self, offset: int) -> "Span":
-        """Crop the span at the given offset.
-
-        Args:
-            offset (int): A value between start and end.
-
-        Returns:
-            Span: A new (possibly smaller) span.
-        """
-        start, end, style = self
-        if offset >= end:
-            return self
-        return Span(start, min(offset, end), style)
-
-    def extend(self, cells: int) -> "Span":
-        """Extend the span by the given number of cells.
-
-        Args:
-            cells (int): Additional space to add to end of span.
-
-        Returns:
-            Span: A span.
-        """
-        if cells:
-            start, end, style = self
-            return Span(start, end + cells, style)
-        else:
-            return self
-
-
-class Text(JupyterMixin):
-    """Text with color / style.
-
-    Args:
-        text (str, optional): Default unstyled text. Defaults to "".
-        style (Union[str, Style], optional): Base style for text. Defaults to "".
-        justify (str, optional): Justify method: "left", "center", "full", "right". Defaults to None.
-        overflow (str, optional): Overflow method: "crop", "fold", "ellipsis". Defaults to None.
-        no_wrap (bool, optional): Disable text wrapping, or None for default. Defaults to None.
-        end (str, optional): Character to end text with. Defaults to "\\\\n".
-        tab_size (int): Number of spaces per tab, or ``None`` to use ``console.tab_size``. Defaults to None.
-        spans (List[Span], optional). A list of predefined style spans. Defaults to None.
-    """
-
-    __slots__ = [
-        "_text",
-        "style",
-        "justify",
-        "overflow",
-        "no_wrap",
-        "end",
-        "tab_size",
-        "_spans",
-        "_length",
-    ]
-
-    def __init__(
-        self,
-        text: str = "",
-        style: Union[str, Style] = "",
-        *,
-        justify: Optional["JustifyMethod"] = None,
-        overflow: Optional["OverflowMethod"] = None,
-        no_wrap: Optional[bool] = None,
-        end: str = "\n",
-        tab_size: Optional[int] = None,
-        spans: Optional[List[Span]] = None,
-    ) -> None:
-        sanitized_text = strip_control_codes(text)
-        self._text = [sanitized_text]
-        self.style = style
-        self.justify: Optional["JustifyMethod"] = justify
-        self.overflow: Optional["OverflowMethod"] = overflow
-        self.no_wrap = no_wrap
-        self.end = end
-        self.tab_size = tab_size
-        self._spans: List[Span] = spans or []
-        self._length: int = len(sanitized_text)
-
-    def __len__(self) -> int:
-        return self._length
-
-    def __bool__(self) -> bool:
-        return bool(self._length)
-
-    def __str__(self) -> str:
-        return self.plain
-
-    def __repr__(self) -> str:
-        return f""
-
-    def __add__(self, other: Any) -> "Text":
-        if isinstance(other, (str, Text)):
-            result = self.copy()
-            result.append(other)
-            return result
-        return NotImplemented
-
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, Text):
-            return NotImplemented
-        return self.plain == other.plain and self._spans == other._spans
-
-    def __contains__(self, other: object) -> bool:
-        if isinstance(other, str):
-            return other in self.plain
-        elif isinstance(other, Text):
-            return other.plain in self.plain
-        return False
-
-    def __getitem__(self, slice: Union[int, slice]) -> "Text":
-        def get_text_at(offset: int) -> "Text":
-            _Span = Span
-            text = Text(
-                self.plain[offset],
-                spans=[
-                    _Span(0, 1, style)
-                    for start, end, style in self._spans
-                    if end > offset >= start
-                ],
-                end="",
-            )
-            return text
-
-        if isinstance(slice, int):
-            return get_text_at(slice)
-        else:
-            start, stop, step = slice.indices(len(self.plain))
-            if step == 1:
-                lines = self.divide([start, stop])
-                return lines[1]
-            else:
-                # This would be a bit of work to implement efficiently
-                # For now, its not required
-                raise TypeError("slices with step!=1 are not supported")
-
-    @property
-    def cell_len(self) -> int:
-        """Get the number of cells required to render this text."""
-        return cell_len(self.plain)
-
-    @property
-    def markup(self) -> str:
-        """Get console markup to render this Text.
-
-        Returns:
-            str: A string potentially creating markup tags.
-        """
-        from .markup import escape
-
-        output: List[str] = []
-
-        plain = self.plain
-        markup_spans = [
-            (0, False, self.style),
-            *((span.start, False, span.style) for span in self._spans),
-            *((span.end, True, span.style) for span in self._spans),
-            (len(plain), True, self.style),
-        ]
-        markup_spans.sort(key=itemgetter(0, 1))
-        position = 0
-        append = output.append
-        for offset, closing, style in markup_spans:
-            if offset > position:
-                append(escape(plain[position:offset]))
-                position = offset
-            if style:
-                append(f"[/{style}]" if closing else f"[{style}]")
-        markup = "".join(output)
-        return markup
-
-    @classmethod
-    def from_markup(
-        cls,
-        text: str,
-        *,
-        style: Union[str, Style] = "",
-        emoji: bool = True,
-        emoji_variant: Optional[EmojiVariant] = None,
-        justify: Optional["JustifyMethod"] = None,
-        overflow: Optional["OverflowMethod"] = None,
-        end: str = "\n",
-    ) -> "Text":
-        """Create Text instance from markup.
-
-        Args:
-            text (str): A string containing console markup.
-            style (Union[str, Style], optional): Base style for text. Defaults to "".
-            emoji (bool, optional): Also render emoji code. Defaults to True.
-            emoji_variant (str, optional): Optional emoji variant, either "text" or "emoji". Defaults to None.
-            justify (str, optional): Justify method: "left", "center", "full", "right". Defaults to None.
-            overflow (str, optional): Overflow method: "crop", "fold", "ellipsis". Defaults to None.
-            end (str, optional): Character to end text with. Defaults to "\\\\n".
-
-        Returns:
-            Text: A Text instance with markup rendered.
-        """
-        from .markup import render
-
-        rendered_text = render(text, style, emoji=emoji, emoji_variant=emoji_variant)
-        rendered_text.justify = justify
-        rendered_text.overflow = overflow
-        rendered_text.end = end
-        return rendered_text
-
-    @classmethod
-    def from_ansi(
-        cls,
-        text: str,
-        *,
-        style: Union[str, Style] = "",
-        justify: Optional["JustifyMethod"] = None,
-        overflow: Optional["OverflowMethod"] = None,
-        no_wrap: Optional[bool] = None,
-        end: str = "\n",
-        tab_size: Optional[int] = 8,
-    ) -> "Text":
-        """Create a Text object from a string containing ANSI escape codes.
-
-        Args:
-            text (str): A string containing escape codes.
-            style (Union[str, Style], optional): Base style for text. Defaults to "".
-            justify (str, optional): Justify method: "left", "center", "full", "right". Defaults to None.
-            overflow (str, optional): Overflow method: "crop", "fold", "ellipsis". Defaults to None.
-            no_wrap (bool, optional): Disable text wrapping, or None for default. Defaults to None.
-            end (str, optional): Character to end text with. Defaults to "\\\\n".
-            tab_size (int): Number of spaces per tab, or ``None`` to use ``console.tab_size``. Defaults to None.
-        """
-        from .ansi import AnsiDecoder
-
-        joiner = Text(
-            "\n",
-            justify=justify,
-            overflow=overflow,
-            no_wrap=no_wrap,
-            end=end,
-            tab_size=tab_size,
-            style=style,
-        )
-        decoder = AnsiDecoder()
-        result = joiner.join(line for line in decoder.decode(text))
-        return result
-
-    @classmethod
-    def styled(
-        cls,
-        text: str,
-        style: StyleType = "",
-        *,
-        justify: Optional["JustifyMethod"] = None,
-        overflow: Optional["OverflowMethod"] = None,
-    ) -> "Text":
-        """Construct a Text instance with a pre-applied styled. A style applied in this way won't be used
-        to pad the text when it is justified.
-
-        Args:
-            text (str): A string containing console markup.
-            style (Union[str, Style]): Style to apply to the text. Defaults to "".
-            justify (str, optional): Justify method: "left", "center", "full", "right". Defaults to None.
-            overflow (str, optional): Overflow method: "crop", "fold", "ellipsis". Defaults to None.
-
-        Returns:
-            Text: A text instance with a style applied to the entire string.
-        """
-        styled_text = cls(text, justify=justify, overflow=overflow)
-        styled_text.stylize(style)
-        return styled_text
-
-    @classmethod
-    def assemble(
-        cls,
-        *parts: Union[str, "Text", Tuple[str, StyleType]],
-        style: Union[str, Style] = "",
-        justify: Optional["JustifyMethod"] = None,
-        overflow: Optional["OverflowMethod"] = None,
-        no_wrap: Optional[bool] = None,
-        end: str = "\n",
-        tab_size: int = 8,
-        meta: Optional[Dict[str, Any]] = None,
-    ) -> "Text":
-        """Construct a text instance by combining a sequence of strings with optional styles.
-        The positional arguments should be either strings, or a tuple of string + style.
-
-        Args:
-            style (Union[str, Style], optional): Base style for text. Defaults to "".
-            justify (str, optional): Justify method: "left", "center", "full", "right". Defaults to None.
-            overflow (str, optional): Overflow method: "crop", "fold", "ellipsis". Defaults to None.
-            no_wrap (bool, optional): Disable text wrapping, or None for default. Defaults to None.
-            end (str, optional): Character to end text with. Defaults to "\\\\n".
-            tab_size (int): Number of spaces per tab, or ``None`` to use ``console.tab_size``. Defaults to None.
-            meta (Dict[str, Any], optional). Meta data to apply to text, or None for no meta data. Default to None
-
-        Returns:
-            Text: A new text instance.
-        """
-        text = cls(
-            style=style,
-            justify=justify,
-            overflow=overflow,
-            no_wrap=no_wrap,
-            end=end,
-            tab_size=tab_size,
-        )
-        append = text.append
-        _Text = Text
-        for part in parts:
-            if isinstance(part, (_Text, str)):
-                append(part)
-            else:
-                append(*part)
-        if meta:
-            text.apply_meta(meta)
-        return text
-
-    @property
-    def plain(self) -> str:
-        """Get the text as a single string."""
-        if len(self._text) != 1:
-            self._text[:] = ["".join(self._text)]
-        return self._text[0]
-
-    @plain.setter
-    def plain(self, new_text: str) -> None:
-        """Set the text to a new value."""
-        if new_text != self.plain:
-            sanitized_text = strip_control_codes(new_text)
-            self._text[:] = [sanitized_text]
-            old_length = self._length
-            self._length = len(sanitized_text)
-            if old_length > self._length:
-                self._trim_spans()
-
-    @property
-    def spans(self) -> List[Span]:
-        """Get a reference to the internal list of spans."""
-        return self._spans
-
-    @spans.setter
-    def spans(self, spans: List[Span]) -> None:
-        """Set spans."""
-        self._spans = spans[:]
-
-    def blank_copy(self, plain: str = "") -> "Text":
-        """Return a new Text instance with copied metadata (but not the string or spans)."""
-        copy_self = Text(
-            plain,
-            style=self.style,
-            justify=self.justify,
-            overflow=self.overflow,
-            no_wrap=self.no_wrap,
-            end=self.end,
-            tab_size=self.tab_size,
-        )
-        return copy_self
-
-    def copy(self) -> "Text":
-        """Return a copy of this instance."""
-        copy_self = Text(
-            self.plain,
-            style=self.style,
-            justify=self.justify,
-            overflow=self.overflow,
-            no_wrap=self.no_wrap,
-            end=self.end,
-            tab_size=self.tab_size,
-        )
-        copy_self._spans[:] = self._spans
-        return copy_self
-
-    def stylize(
-        self,
-        style: Union[str, Style],
-        start: int = 0,
-        end: Optional[int] = None,
-    ) -> None:
-        """Apply a style to the text, or a portion of the text.
-
-        Args:
-            style (Union[str, Style]): Style instance or style definition to apply.
-            start (int): Start offset (negative indexing is supported). Defaults to 0.
-            end (Optional[int], optional): End offset (negative indexing is supported), or None for end of text. Defaults to None.
-        """
-        if style:
-            length = len(self)
-            if start < 0:
-                start = length + start
-            if end is None:
-                end = length
-            if end < 0:
-                end = length + end
-            if start >= length or end <= start:
-                # Span not in text or not valid
-                return
-            self._spans.append(Span(start, min(length, end), style))
-
-    def stylize_before(
-        self,
-        style: Union[str, Style],
-        start: int = 0,
-        end: Optional[int] = None,
-    ) -> None:
-        """Apply a style to the text, or a portion of the text. Styles will be applied before other styles already present.
-
-        Args:
-            style (Union[str, Style]): Style instance or style definition to apply.
-            start (int): Start offset (negative indexing is supported). Defaults to 0.
-            end (Optional[int], optional): End offset (negative indexing is supported), or None for end of text. Defaults to None.
-        """
-        if style:
-            length = len(self)
-            if start < 0:
-                start = length + start
-            if end is None:
-                end = length
-            if end < 0:
-                end = length + end
-            if start >= length or end <= start:
-                # Span not in text or not valid
-                return
-            self._spans.insert(0, Span(start, min(length, end), style))
-
-    def apply_meta(
-        self, meta: Dict[str, Any], start: int = 0, end: Optional[int] = None
-    ) -> None:
-        """Apply metadata to the text, or a portion of the text.
-
-        Args:
-            meta (Dict[str, Any]): A dict of meta information.
-            start (int): Start offset (negative indexing is supported). Defaults to 0.
-            end (Optional[int], optional): End offset (negative indexing is supported), or None for end of text. Defaults to None.
-
-        """
-        style = Style.from_meta(meta)
-        self.stylize(style, start=start, end=end)
-
-    def on(self, meta: Optional[Dict[str, Any]] = None, **handlers: Any) -> "Text":
-        """Apply event handlers (used by Textual project).
-
-        Example:
-            >>> from rich.text import Text
-            >>> text = Text("hello world")
-            >>> text.on(click="view.toggle('world')")
-
-        Args:
-            meta (Dict[str, Any]): Mapping of meta information.
-            **handlers: Keyword args are prefixed with "@" to defined handlers.
-
-        Returns:
-            Text: Self is returned to method may be chained.
-        """
-        meta = {} if meta is None else meta
-        meta.update({f"@{key}": value for key, value in handlers.items()})
-        self.stylize(Style.from_meta(meta))
-        return self
-
-    def remove_suffix(self, suffix: str) -> None:
-        """Remove a suffix if it exists.
-
-        Args:
-            suffix (str): Suffix to remove.
-        """
-        if self.plain.endswith(suffix):
-            self.right_crop(len(suffix))
-
-    def get_style_at_offset(self, console: "Console", offset: int) -> Style:
-        """Get the style of a character at give offset.
-
-        Args:
-            console (~Console): Console where text will be rendered.
-            offset (int): Offset in to text (negative indexing supported)
-
-        Returns:
-            Style: A Style instance.
-        """
-        # TODO: This is a little inefficient, it is only used by full justify
-        if offset < 0:
-            offset = len(self) + offset
-        get_style = console.get_style
-        style = get_style(self.style).copy()
-        for start, end, span_style in self._spans:
-            if end > offset >= start:
-                style += get_style(span_style, default="")
-        return style
-
-    def extend_style(self, spaces: int) -> None:
-        """Extend the Text given number of spaces where the spaces have the same style as the last character.
-
-        Args:
-            spaces (int): Number of spaces to add to the Text.
-        """
-        if spaces <= 0:
-            return
-        spans = self.spans
-        new_spaces = " " * spaces
-        if spans:
-            end_offset = len(self)
-            self._spans[:] = [
-                span.extend(spaces) if span.end >= end_offset else span
-                for span in spans
-            ]
-            self._text.append(new_spaces)
-            self._length += spaces
-        else:
-            self.plain += new_spaces
-
-    def highlight_regex(
-        self,
-        re_highlight: Union[Pattern[str], str],
-        style: Optional[Union[GetStyleCallable, StyleType]] = None,
-        *,
-        style_prefix: str = "",
-    ) -> int:
-        """Highlight text with a regular expression, where group names are
-        translated to styles.
-
-        Args:
-            re_highlight (Union[re.Pattern, str]): A regular expression object or string.
-            style (Union[GetStyleCallable, StyleType]): Optional style to apply to whole match, or a callable
-                which accepts the matched text and returns a style. Defaults to None.
-            style_prefix (str, optional): Optional prefix to add to style group names.
-
-        Returns:
-            int: Number of regex matches
-        """
-        count = 0
-        append_span = self._spans.append
-        _Span = Span
-        plain = self.plain
-        if isinstance(re_highlight, str):
-            re_highlight = re.compile(re_highlight)
-        for match in re_highlight.finditer(plain):
-            get_span = match.span
-            if style:
-                start, end = get_span()
-                match_style = style(plain[start:end]) if callable(style) else style
-                if match_style is not None and end > start:
-                    append_span(_Span(start, end, match_style))
-
-            count += 1
-            for name in match.groupdict().keys():
-                start, end = get_span(name)
-                if start != -1 and end > start:
-                    append_span(_Span(start, end, f"{style_prefix}{name}"))
-        return count
-
-    def highlight_words(
-        self,
-        words: Iterable[str],
-        style: Union[str, Style],
-        *,
-        case_sensitive: bool = True,
-    ) -> int:
-        """Highlight words with a style.
-
-        Args:
-            words (Iterable[str]): Words to highlight.
-            style (Union[str, Style]): Style to apply.
-            case_sensitive (bool, optional): Enable case sensitive matching. Defaults to True.
-
-        Returns:
-            int: Number of words highlighted.
-        """
-        re_words = "|".join(re.escape(word) for word in words)
-        add_span = self._spans.append
-        count = 0
-        _Span = Span
-        for match in re.finditer(
-            re_words, self.plain, flags=0 if case_sensitive else re.IGNORECASE
-        ):
-            start, end = match.span(0)
-            add_span(_Span(start, end, style))
-            count += 1
-        return count
-
-    def rstrip(self) -> None:
-        """Strip whitespace from end of text."""
-        self.plain = self.plain.rstrip()
-
-    def rstrip_end(self, size: int) -> None:
-        """Remove whitespace beyond a certain width at the end of the text.
-
-        Args:
-            size (int): The desired size of the text.
-        """
-        text_length = len(self)
-        if text_length > size:
-            excess = text_length - size
-            whitespace_match = _re_whitespace.search(self.plain)
-            if whitespace_match is not None:
-                whitespace_count = len(whitespace_match.group(0))
-                self.right_crop(min(whitespace_count, excess))
-
-    def set_length(self, new_length: int) -> None:
-        """Set new length of the text, clipping or padding is required."""
-        length = len(self)
-        if length != new_length:
-            if length < new_length:
-                self.pad_right(new_length - length)
-            else:
-                self.right_crop(length - new_length)
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> Iterable[Segment]:
-        tab_size: int = console.tab_size if self.tab_size is None else self.tab_size
-        justify = self.justify or options.justify or DEFAULT_JUSTIFY
-
-        overflow = self.overflow or options.overflow or DEFAULT_OVERFLOW
-
-        lines = self.wrap(
-            console,
-            options.max_width,
-            justify=justify,
-            overflow=overflow,
-            tab_size=tab_size or 8,
-            no_wrap=pick_bool(self.no_wrap, options.no_wrap, False),
-        )
-        all_lines = Text("\n").join(lines)
-        yield from all_lines.render(console, end=self.end)
-
-    def __rich_measure__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> Measurement:
-        text = self.plain
-        lines = text.splitlines()
-        max_text_width = max(cell_len(line) for line in lines) if lines else 0
-        words = text.split()
-        min_text_width = (
-            max(cell_len(word) for word in words) if words else max_text_width
-        )
-        return Measurement(min_text_width, max_text_width)
-
-    def render(self, console: "Console", end: str = "") -> Iterable["Segment"]:
-        """Render the text as Segments.
-
-        Args:
-            console (Console): Console instance.
-            end (Optional[str], optional): Optional end character.
-
-        Returns:
-            Iterable[Segment]: Result of render that may be written to the console.
-        """
-        _Segment = Segment
-        text = self.plain
-        if not self._spans:
-            yield Segment(text)
-            if end:
-                yield _Segment(end)
-            return
-        get_style = partial(console.get_style, default=Style.null())
-
-        enumerated_spans = list(enumerate(self._spans, 1))
-        style_map = {index: get_style(span.style) for index, span in enumerated_spans}
-        style_map[0] = get_style(self.style)
-
-        spans = [
-            (0, False, 0),
-            *((span.start, False, index) for index, span in enumerated_spans),
-            *((span.end, True, index) for index, span in enumerated_spans),
-            (len(text), True, 0),
-        ]
-        spans.sort(key=itemgetter(0, 1))
-
-        stack: List[int] = []
-        stack_append = stack.append
-        stack_pop = stack.remove
-
-        style_cache: Dict[Tuple[Style, ...], Style] = {}
-        style_cache_get = style_cache.get
-        combine = Style.combine
-
-        def get_current_style() -> Style:
-            """Construct current style from stack."""
-            styles = tuple(style_map[_style_id] for _style_id in sorted(stack))
-            cached_style = style_cache_get(styles)
-            if cached_style is not None:
-                return cached_style
-            current_style = combine(styles)
-            style_cache[styles] = current_style
-            return current_style
-
-        for (offset, leaving, style_id), (next_offset, _, _) in zip(spans, spans[1:]):
-            if leaving:
-                stack_pop(style_id)
-            else:
-                stack_append(style_id)
-            if next_offset > offset:
-                yield _Segment(text[offset:next_offset], get_current_style())
-        if end:
-            yield _Segment(end)
-
-    def join(self, lines: Iterable["Text"]) -> "Text":
-        """Join text together with this instance as the separator.
-
-        Args:
-            lines (Iterable[Text]): An iterable of Text instances to join.
-
-        Returns:
-            Text: A new text instance containing join text.
-        """
-
-        new_text = self.blank_copy()
-
-        def iter_text() -> Iterable["Text"]:
-            if self.plain:
-                for last, line in loop_last(lines):
-                    yield line
-                    if not last:
-                        yield self
-            else:
-                yield from lines
-
-        extend_text = new_text._text.extend
-        append_span = new_text._spans.append
-        extend_spans = new_text._spans.extend
-        offset = 0
-        _Span = Span
-
-        for text in iter_text():
-            extend_text(text._text)
-            if text.style:
-                append_span(_Span(offset, offset + len(text), text.style))
-            extend_spans(
-                _Span(offset + start, offset + end, style)
-                for start, end, style in text._spans
-            )
-            offset += len(text)
-        new_text._length = offset
-        return new_text
-
-    def expand_tabs(self, tab_size: Optional[int] = None) -> None:
-        """Converts tabs to spaces.
-
-        Args:
-            tab_size (int, optional): Size of tabs. Defaults to 8.
-
-        """
-        if "\t" not in self.plain:
-            return
-        if tab_size is None:
-            tab_size = self.tab_size
-        if tab_size is None:
-            tab_size = 8
-
-        new_text: List[Text] = []
-        append = new_text.append
-
-        for line in self.split("\n", include_separator=True):
-            if "\t" not in line.plain:
-                append(line)
-            else:
-                cell_position = 0
-                parts = line.split("\t", include_separator=True)
-                for part in parts:
-                    if part.plain.endswith("\t"):
-                        part._text[-1] = part._text[-1][:-1] + " "
-                        cell_position += part.cell_len
-                        tab_remainder = cell_position % tab_size
-                        if tab_remainder:
-                            spaces = tab_size - tab_remainder
-                            part.extend_style(spaces)
-                            cell_position += spaces
-                    else:
-                        cell_position += part.cell_len
-                    append(part)
-
-        result = Text("").join(new_text)
-
-        self._text = [result.plain]
-        self._length = len(self.plain)
-        self._spans[:] = result._spans
-
-    def truncate(
-        self,
-        max_width: int,
-        *,
-        overflow: Optional["OverflowMethod"] = None,
-        pad: bool = False,
-    ) -> None:
-        """Truncate text if it is longer that a given width.
-
-        Args:
-            max_width (int): Maximum number of characters in text.
-            overflow (str, optional): Overflow method: "crop", "fold", or "ellipsis". Defaults to None, to use self.overflow.
-            pad (bool, optional): Pad with spaces if the length is less than max_width. Defaults to False.
-        """
-        _overflow = overflow or self.overflow or DEFAULT_OVERFLOW
-        if _overflow != "ignore":
-            length = cell_len(self.plain)
-            if length > max_width:
-                if _overflow == "ellipsis":
-                    self.plain = set_cell_size(self.plain, max_width - 1) + "…"
-                else:
-                    self.plain = set_cell_size(self.plain, max_width)
-            if pad and length < max_width:
-                spaces = max_width - length
-                self._text = [f"{self.plain}{' ' * spaces}"]
-                self._length = len(self.plain)
-
-    def _trim_spans(self) -> None:
-        """Remove or modify any spans that are over the end of the text."""
-        max_offset = len(self.plain)
-        _Span = Span
-        self._spans[:] = [
-            (
-                span
-                if span.end < max_offset
-                else _Span(span.start, min(max_offset, span.end), span.style)
-            )
-            for span in self._spans
-            if span.start < max_offset
-        ]
-
-    def pad(self, count: int, character: str = " ") -> None:
-        """Pad left and right with a given number of characters.
-
-        Args:
-            count (int): Width of padding.
-            character (str): The character to pad with. Must be a string of length 1.
-        """
-        assert len(character) == 1, "Character must be a string of length 1"
-        if count:
-            pad_characters = character * count
-            self.plain = f"{pad_characters}{self.plain}{pad_characters}"
-            _Span = Span
-            self._spans[:] = [
-                _Span(start + count, end + count, style)
-                for start, end, style in self._spans
-            ]
-
-    def pad_left(self, count: int, character: str = " ") -> None:
-        """Pad the left with a given character.
-
-        Args:
-            count (int): Number of characters to pad.
-            character (str, optional): Character to pad with. Defaults to " ".
-        """
-        assert len(character) == 1, "Character must be a string of length 1"
-        if count:
-            self.plain = f"{character * count}{self.plain}"
-            _Span = Span
-            self._spans[:] = [
-                _Span(start + count, end + count, style)
-                for start, end, style in self._spans
-            ]
-
-    def pad_right(self, count: int, character: str = " ") -> None:
-        """Pad the right with a given character.
-
-        Args:
-            count (int): Number of characters to pad.
-            character (str, optional): Character to pad with. Defaults to " ".
-        """
-        assert len(character) == 1, "Character must be a string of length 1"
-        if count:
-            self.plain = f"{self.plain}{character * count}"
-
-    def align(self, align: AlignMethod, width: int, character: str = " ") -> None:
-        """Align text to a given width.
-
-        Args:
-            align (AlignMethod): One of "left", "center", or "right".
-            width (int): Desired width.
-            character (str, optional): Character to pad with. Defaults to " ".
-        """
-        self.truncate(width)
-        excess_space = width - cell_len(self.plain)
-        if excess_space:
-            if align == "left":
-                self.pad_right(excess_space, character)
-            elif align == "center":
-                left = excess_space // 2
-                self.pad_left(left, character)
-                self.pad_right(excess_space - left, character)
-            else:
-                self.pad_left(excess_space, character)
-
-    def append(
-        self, text: Union["Text", str], style: Optional[Union[str, "Style"]] = None
-    ) -> "Text":
-        """Add text with an optional style.
-
-        Args:
-            text (Union[Text, str]): A str or Text to append.
-            style (str, optional): A style name. Defaults to None.
-
-        Returns:
-            Text: Returns self for chaining.
-        """
-
-        if not isinstance(text, (str, Text)):
-            raise TypeError("Only str or Text can be appended to Text")
-
-        if len(text):
-            if isinstance(text, str):
-                sanitized_text = strip_control_codes(text)
-                self._text.append(sanitized_text)
-                offset = len(self)
-                text_length = len(sanitized_text)
-                if style:
-                    self._spans.append(Span(offset, offset + text_length, style))
-                self._length += text_length
-            elif isinstance(text, Text):
-                _Span = Span
-                if style is not None:
-                    raise ValueError(
-                        "style must not be set when appending Text instance"
-                    )
-                text_length = self._length
-                if text.style:
-                    self._spans.append(
-                        _Span(text_length, text_length + len(text), text.style)
-                    )
-                self._text.append(text.plain)
-                self._spans.extend(
-                    _Span(start + text_length, end + text_length, style)
-                    for start, end, style in text._spans.copy()
-                )
-                self._length += len(text)
-        return self
-
-    def append_text(self, text: "Text") -> "Text":
-        """Append another Text instance. This method is more performant that Text.append, but
-        only works for Text.
-
-        Args:
-            text (Text): The Text instance to append to this instance.
-
-        Returns:
-            Text: Returns self for chaining.
-        """
-        _Span = Span
-        text_length = self._length
-        if text.style:
-            self._spans.append(_Span(text_length, text_length + len(text), text.style))
-        self._text.append(text.plain)
-        self._spans.extend(
-            _Span(start + text_length, end + text_length, style)
-            for start, end, style in text._spans.copy()
-        )
-        self._length += len(text)
-        return self
-
-    def append_tokens(
-        self, tokens: Iterable[Tuple[str, Optional[StyleType]]]
-    ) -> "Text":
-        """Append iterable of str and style. Style may be a Style instance or a str style definition.
-
-        Args:
-            tokens (Iterable[Tuple[str, Optional[StyleType]]]): An iterable of tuples containing str content and style.
-
-        Returns:
-            Text: Returns self for chaining.
-        """
-        append_text = self._text.append
-        append_span = self._spans.append
-        _Span = Span
-        offset = len(self)
-        for content, style in tokens:
-            content = strip_control_codes(content)
-            append_text(content)
-            if style:
-                append_span(_Span(offset, offset + len(content), style))
-            offset += len(content)
-        self._length = offset
-        return self
-
-    def copy_styles(self, text: "Text") -> None:
-        """Copy styles from another Text instance.
-
-        Args:
-            text (Text): A Text instance to copy styles from, must be the same length.
-        """
-        self._spans.extend(text._spans)
-
-    def split(
-        self,
-        separator: str = "\n",
-        *,
-        include_separator: bool = False,
-        allow_blank: bool = False,
-    ) -> Lines:
-        """Split rich text in to lines, preserving styles.
-
-        Args:
-            separator (str, optional): String to split on. Defaults to "\\\\n".
-            include_separator (bool, optional): Include the separator in the lines. Defaults to False.
-            allow_blank (bool, optional): Return a blank line if the text ends with a separator. Defaults to False.
-
-        Returns:
-            List[RichText]: A list of rich text, one per line of the original.
-        """
-        assert separator, "separator must not be empty"
-
-        text = self.plain
-        if separator not in text:
-            return Lines([self.copy()])
-
-        if include_separator:
-            lines = self.divide(
-                match.end() for match in re.finditer(re.escape(separator), text)
-            )
-        else:
-
-            def flatten_spans() -> Iterable[int]:
-                for match in re.finditer(re.escape(separator), text):
-                    start, end = match.span()
-                    yield start
-                    yield end
-
-            lines = Lines(
-                line for line in self.divide(flatten_spans()) if line.plain != separator
-            )
-
-        if not allow_blank and text.endswith(separator):
-            lines.pop()
-
-        return lines
-
-    def divide(self, offsets: Iterable[int]) -> Lines:
-        """Divide text in to a number of lines at given offsets.
-
-        Args:
-            offsets (Iterable[int]): Offsets used to divide text.
-
-        Returns:
-            Lines: New RichText instances between offsets.
-        """
-        _offsets = list(offsets)
-
-        if not _offsets:
-            return Lines([self.copy()])
-
-        text = self.plain
-        text_length = len(text)
-        divide_offsets = [0, *_offsets, text_length]
-        line_ranges = list(zip(divide_offsets, divide_offsets[1:]))
-
-        style = self.style
-        justify = self.justify
-        overflow = self.overflow
-        _Text = Text
-        new_lines = Lines(
-            _Text(
-                text[start:end],
-                style=style,
-                justify=justify,
-                overflow=overflow,
-            )
-            for start, end in line_ranges
-        )
-        if not self._spans:
-            return new_lines
-
-        _line_appends = [line._spans.append for line in new_lines._lines]
-        line_count = len(line_ranges)
-        _Span = Span
-
-        for span_start, span_end, style in self._spans:
-            lower_bound = 0
-            upper_bound = line_count
-            start_line_no = (lower_bound + upper_bound) // 2
-
-            while True:
-                line_start, line_end = line_ranges[start_line_no]
-                if span_start < line_start:
-                    upper_bound = start_line_no - 1
-                elif span_start > line_end:
-                    lower_bound = start_line_no + 1
-                else:
-                    break
-                start_line_no = (lower_bound + upper_bound) // 2
-
-            if span_end < line_end:
-                end_line_no = start_line_no
-            else:
-                end_line_no = lower_bound = start_line_no
-                upper_bound = line_count
-
-                while True:
-                    line_start, line_end = line_ranges[end_line_no]
-                    if span_end < line_start:
-                        upper_bound = end_line_no - 1
-                    elif span_end > line_end:
-                        lower_bound = end_line_no + 1
-                    else:
-                        break
-                    end_line_no = (lower_bound + upper_bound) // 2
-
-            for line_no in range(start_line_no, end_line_no + 1):
-                line_start, line_end = line_ranges[line_no]
-                new_start = max(0, span_start - line_start)
-                new_end = min(span_end - line_start, line_end - line_start)
-                if new_end > new_start:
-                    _line_appends[line_no](_Span(new_start, new_end, style))
-
-        return new_lines
-
-    def right_crop(self, amount: int = 1) -> None:
-        """Remove a number of characters from the end of the text."""
-        max_offset = len(self.plain) - amount
-        _Span = Span
-        self._spans[:] = [
-            (
-                span
-                if span.end < max_offset
-                else _Span(span.start, min(max_offset, span.end), span.style)
-            )
-            for span in self._spans
-            if span.start < max_offset
-        ]
-        self._text = [self.plain[:-amount]]
-        self._length -= amount
-
-    def wrap(
-        self,
-        console: "Console",
-        width: int,
-        *,
-        justify: Optional["JustifyMethod"] = None,
-        overflow: Optional["OverflowMethod"] = None,
-        tab_size: int = 8,
-        no_wrap: Optional[bool] = None,
-    ) -> Lines:
-        """Word wrap the text.
-
-        Args:
-            console (Console): Console instance.
-            width (int): Number of cells available per line.
-            justify (str, optional): Justify method: "default", "left", "center", "full", "right". Defaults to "default".
-            overflow (str, optional): Overflow method: "crop", "fold", or "ellipsis". Defaults to None.
-            tab_size (int, optional): Default tab size. Defaults to 8.
-            no_wrap (bool, optional): Disable wrapping, Defaults to False.
-
-        Returns:
-            Lines: Number of lines.
-        """
-        wrap_justify = justify or self.justify or DEFAULT_JUSTIFY
-        wrap_overflow = overflow or self.overflow or DEFAULT_OVERFLOW
-
-        no_wrap = pick_bool(no_wrap, self.no_wrap, False) or overflow == "ignore"
-
-        lines = Lines()
-        for line in self.split(allow_blank=True):
-            if "\t" in line:
-                line.expand_tabs(tab_size)
-            if no_wrap:
-                new_lines = Lines([line])
-            else:
-                offsets = divide_line(str(line), width, fold=wrap_overflow == "fold")
-                new_lines = line.divide(offsets)
-            for line in new_lines:
-                line.rstrip_end(width)
-            if wrap_justify:
-                new_lines.justify(
-                    console, width, justify=wrap_justify, overflow=wrap_overflow
-                )
-            for line in new_lines:
-                line.truncate(width, overflow=wrap_overflow)
-            lines.extend(new_lines)
-        return lines
-
-    def fit(self, width: int) -> Lines:
-        """Fit the text in to given width by chopping in to lines.
-
-        Args:
-            width (int): Maximum characters in a line.
-
-        Returns:
-            Lines: Lines container.
-        """
-        lines: Lines = Lines()
-        append = lines.append
-        for line in self.split():
-            line.set_length(width)
-            append(line)
-        return lines
-
-    def detect_indentation(self) -> int:
-        """Auto-detect indentation of code.
-
-        Returns:
-            int: Number of spaces used to indent code.
-        """
-
-        _indentations = {
-            len(match.group(1))
-            for match in re.finditer(r"^( *)(.*)$", self.plain, flags=re.MULTILINE)
-        }
-
-        try:
-            indentation = (
-                reduce(gcd, [indent for indent in _indentations if not indent % 2]) or 1
-            )
-        except TypeError:
-            indentation = 1
-
-        return indentation
-
-    def with_indent_guides(
-        self,
-        indent_size: Optional[int] = None,
-        *,
-        character: str = "│",
-        style: StyleType = "dim green",
-    ) -> "Text":
-        """Adds indent guide lines to text.
-
-        Args:
-            indent_size (Optional[int]): Size of indentation, or None to auto detect. Defaults to None.
-            character (str, optional): Character to use for indentation. Defaults to "│".
-            style (Union[Style, str], optional): Style of indent guides.
-
-        Returns:
-            Text: New text with indentation guides.
-        """
-
-        _indent_size = self.detect_indentation() if indent_size is None else indent_size
-
-        text = self.copy()
-        text.expand_tabs()
-        indent_line = f"{character}{' ' * (_indent_size - 1)}"
-
-        re_indent = re.compile(r"^( *)(.*)$")
-        new_lines: List[Text] = []
-        add_line = new_lines.append
-        blank_lines = 0
-        for line in text.split(allow_blank=True):
-            match = re_indent.match(line.plain)
-            if not match or not match.group(2):
-                blank_lines += 1
-                continue
-            indent = match.group(1)
-            full_indents, remaining_space = divmod(len(indent), _indent_size)
-            new_indent = f"{indent_line * full_indents}{' ' * remaining_space}"
-            line.plain = new_indent + line.plain[len(new_indent) :]
-            line.stylize(style, 0, len(new_indent))
-            if blank_lines:
-                new_lines.extend([Text(new_indent, style=style)] * blank_lines)
-                blank_lines = 0
-            add_line(line)
-        if blank_lines:
-            new_lines.extend([Text("", style=style)] * blank_lines)
-
-        new_text = text.blank_copy("\n").join(new_lines)
-        return new_text
-
-
-if __name__ == "__main__":  # pragma: no cover
-    from rich.console import Console
-
-    text = Text(
-        """\nLorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.\n"""
-    )
-    text.highlight_words(["Lorem"], "bold")
-    text.highlight_words(["ipsum"], "italic")
-
-    console = Console()
-
-    console.rule("justify='left'")
-    console.print(text, style="red")
-    console.print()
-
-    console.rule("justify='center'")
-    console.print(text, style="green", justify="center")
-    console.print()
-
-    console.rule("justify='right'")
-    console.print(text, style="blue", justify="right")
-    console.print()
-
-    console.rule("justify='full'")
-    console.print(text, style="magenta", justify="full")
-    console.print()
diff --git a/venv/lib/python3.10/site-packages/rich/theme.py b/venv/lib/python3.10/site-packages/rich/theme.py
deleted file mode 100644
index 227f1d8635f8ba915153b21a6b925643a11d286e..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/theme.py
+++ /dev/null
@@ -1,115 +0,0 @@
-import configparser
-from typing import IO, Dict, List, Mapping, Optional
-
-from .default_styles import DEFAULT_STYLES
-from .style import Style, StyleType
-
-
-class Theme:
-    """A container for style information, used by :class:`~rich.console.Console`.
-
-    Args:
-        styles (Dict[str, Style], optional): A mapping of style names on to styles. Defaults to None for a theme with no styles.
-        inherit (bool, optional): Inherit default styles. Defaults to True.
-    """
-
-    styles: Dict[str, Style]
-
-    def __init__(
-        self, styles: Optional[Mapping[str, StyleType]] = None, inherit: bool = True
-    ):
-        self.styles = DEFAULT_STYLES.copy() if inherit else {}
-        if styles is not None:
-            self.styles.update(
-                {
-                    name: style if isinstance(style, Style) else Style.parse(style)
-                    for name, style in styles.items()
-                }
-            )
-
-    @property
-    def config(self) -> str:
-        """Get contents of a config file for this theme."""
-        config = "[styles]\n" + "\n".join(
-            f"{name} = {style}" for name, style in sorted(self.styles.items())
-        )
-        return config
-
-    @classmethod
-    def from_file(
-        cls, config_file: IO[str], source: Optional[str] = None, inherit: bool = True
-    ) -> "Theme":
-        """Load a theme from a text mode file.
-
-        Args:
-            config_file (IO[str]): An open conf file.
-            source (str, optional): The filename of the open file. Defaults to None.
-            inherit (bool, optional): Inherit default styles. Defaults to True.
-
-        Returns:
-            Theme: A New theme instance.
-        """
-        config = configparser.ConfigParser()
-        config.read_file(config_file, source=source)
-        styles = {name: Style.parse(value) for name, value in config.items("styles")}
-        theme = Theme(styles, inherit=inherit)
-        return theme
-
-    @classmethod
-    def read(
-        cls, path: str, inherit: bool = True, encoding: Optional[str] = None
-    ) -> "Theme":
-        """Read a theme from a path.
-
-        Args:
-            path (str): Path to a config file readable by Python configparser module.
-            inherit (bool, optional): Inherit default styles. Defaults to True.
-            encoding (str, optional): Encoding of the config file. Defaults to None.
-
-        Returns:
-            Theme: A new theme instance.
-        """
-        with open(path, encoding=encoding) as config_file:
-            return cls.from_file(config_file, source=path, inherit=inherit)
-
-
-class ThemeStackError(Exception):
-    """Base exception for errors related to the theme stack."""
-
-
-class ThemeStack:
-    """A stack of themes.
-
-    Args:
-        theme (Theme): A theme instance
-    """
-
-    def __init__(self, theme: Theme) -> None:
-        self._entries: List[Dict[str, Style]] = [theme.styles]
-        self.get = self._entries[-1].get
-
-    def push_theme(self, theme: Theme, inherit: bool = True) -> None:
-        """Push a theme on the top of the stack.
-
-        Args:
-            theme (Theme): A Theme instance.
-            inherit (boolean, optional): Inherit styles from current top of stack.
-        """
-        styles: Dict[str, Style]
-        styles = (
-            {**self._entries[-1], **theme.styles} if inherit else theme.styles.copy()
-        )
-        self._entries.append(styles)
-        self.get = self._entries[-1].get
-
-    def pop_theme(self) -> None:
-        """Pop (and discard) the top-most theme."""
-        if len(self._entries) == 1:
-            raise ThemeStackError("Unable to pop base theme")
-        self._entries.pop()
-        self.get = self._entries[-1].get
-
-
-if __name__ == "__main__":  # pragma: no cover
-    theme = Theme()
-    print(theme.config)
diff --git a/venv/lib/python3.10/site-packages/rich/themes.py b/venv/lib/python3.10/site-packages/rich/themes.py
deleted file mode 100644
index bf6db104a2c4fd4f3dc699e85f2b262c3d31e9a0..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/themes.py
+++ /dev/null
@@ -1,5 +0,0 @@
-from .default_styles import DEFAULT_STYLES
-from .theme import Theme
-
-
-DEFAULT = Theme(DEFAULT_STYLES)
diff --git a/venv/lib/python3.10/site-packages/rich/traceback.py b/venv/lib/python3.10/site-packages/rich/traceback.py
deleted file mode 100644
index b2cc6304048b6149d55c05550c597774241ddd0d..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/traceback.py
+++ /dev/null
@@ -1,884 +0,0 @@
-import inspect
-import linecache
-import os
-import sys
-from dataclasses import dataclass, field
-from itertools import islice
-from traceback import walk_tb
-from types import ModuleType, TracebackType
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Iterable,
-    List,
-    Optional,
-    Sequence,
-    Tuple,
-    Type,
-    Union,
-)
-
-from pygments.lexers import guess_lexer_for_filename
-from pygments.token import Comment, Keyword, Name, Number, Operator, String
-from pygments.token import Text as TextToken
-from pygments.token import Token
-from pygments.util import ClassNotFound
-
-from . import pretty
-from ._loop import loop_first_last, loop_last
-from .columns import Columns
-from .console import (
-    Console,
-    ConsoleOptions,
-    ConsoleRenderable,
-    Group,
-    RenderResult,
-    group,
-)
-from .constrain import Constrain
-from .highlighter import RegexHighlighter, ReprHighlighter
-from .panel import Panel
-from .scope import render_scope
-from .style import Style
-from .syntax import Syntax, SyntaxPosition
-from .text import Text
-from .theme import Theme
-
-WINDOWS = sys.platform == "win32"
-
-LOCALS_MAX_LENGTH = 10
-LOCALS_MAX_STRING = 80
-
-
-def _iter_syntax_lines(
-    start: SyntaxPosition, end: SyntaxPosition
-) -> Iterable[Tuple[int, int, int]]:
-    """Yield start and end positions per line.
-
-    Args:
-        start: Start position.
-        end: End position.
-
-    Returns:
-        Iterable of (LINE, COLUMN1, COLUMN2).
-    """
-
-    line1, column1 = start
-    line2, column2 = end
-
-    if line1 == line2:
-        yield line1, column1, column2
-    else:
-        for first, last, line_no in loop_first_last(range(line1, line2 + 1)):
-            if first:
-                yield line_no, column1, -1
-            elif last:
-                yield line_no, 0, column2
-            else:
-                yield line_no, 0, -1
-
-
-def install(
-    *,
-    console: Optional[Console] = None,
-    width: Optional[int] = 100,
-    code_width: Optional[int] = 88,
-    extra_lines: int = 3,
-    theme: Optional[str] = None,
-    word_wrap: bool = False,
-    show_locals: bool = False,
-    locals_max_length: int = LOCALS_MAX_LENGTH,
-    locals_max_string: int = LOCALS_MAX_STRING,
-    locals_hide_dunder: bool = True,
-    locals_hide_sunder: Optional[bool] = None,
-    indent_guides: bool = True,
-    suppress: Iterable[Union[str, ModuleType]] = (),
-    max_frames: int = 100,
-) -> Callable[[Type[BaseException], BaseException, Optional[TracebackType]], Any]:
-    """Install a rich traceback handler.
-
-    Once installed, any tracebacks will be printed with syntax highlighting and rich formatting.
-
-
-    Args:
-        console (Optional[Console], optional): Console to write exception to. Default uses internal Console instance.
-        width (Optional[int], optional): Width (in characters) of traceback. Defaults to 100.
-        code_width (Optional[int], optional): Code width (in characters) of traceback. Defaults to 88.
-        extra_lines (int, optional): Extra lines of code. Defaults to 3.
-        theme (Optional[str], optional): Pygments theme to use in traceback. Defaults to ``None`` which will pick
-            a theme appropriate for the platform.
-        word_wrap (bool, optional): Enable word wrapping of long lines. Defaults to False.
-        show_locals (bool, optional): Enable display of local variables. Defaults to False.
-        locals_max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
-            Defaults to 10.
-        locals_max_string (int, optional): Maximum length of string before truncating, or None to disable. Defaults to 80.
-        locals_hide_dunder (bool, optional): Hide locals prefixed with double underscore. Defaults to True.
-        locals_hide_sunder (bool, optional): Hide locals prefixed with single underscore. Defaults to False.
-        indent_guides (bool, optional): Enable indent guides in code and locals. Defaults to True.
-        suppress (Sequence[Union[str, ModuleType]]): Optional sequence of modules or paths to exclude from traceback.
-
-    Returns:
-        Callable: The previous exception handler that was replaced.
-
-    """
-    traceback_console = Console(stderr=True) if console is None else console
-
-    locals_hide_sunder = (
-        True
-        if (traceback_console.is_jupyter and locals_hide_sunder is None)
-        else locals_hide_sunder
-    )
-
-    def excepthook(
-        type_: Type[BaseException],
-        value: BaseException,
-        traceback: Optional[TracebackType],
-    ) -> None:
-        exception_traceback = Traceback.from_exception(
-            type_,
-            value,
-            traceback,
-            width=width,
-            code_width=code_width,
-            extra_lines=extra_lines,
-            theme=theme,
-            word_wrap=word_wrap,
-            show_locals=show_locals,
-            locals_max_length=locals_max_length,
-            locals_max_string=locals_max_string,
-            locals_hide_dunder=locals_hide_dunder,
-            locals_hide_sunder=bool(locals_hide_sunder),
-            indent_guides=indent_guides,
-            suppress=suppress,
-            max_frames=max_frames,
-        )
-        traceback_console.print(exception_traceback)
-
-    def ipy_excepthook_closure(ip: Any) -> None:  # pragma: no cover
-        tb_data = {}  # store information about showtraceback call
-        default_showtraceback = ip.showtraceback  # keep reference of default traceback
-
-        def ipy_show_traceback(*args: Any, **kwargs: Any) -> None:
-            """wrap the default ip.showtraceback to store info for ip._showtraceback"""
-            nonlocal tb_data
-            tb_data = kwargs
-            default_showtraceback(*args, **kwargs)
-
-        def ipy_display_traceback(
-            *args: Any, is_syntax: bool = False, **kwargs: Any
-        ) -> None:
-            """Internally called traceback from ip._showtraceback"""
-            nonlocal tb_data
-            exc_tuple = ip._get_exc_info()
-
-            # do not display trace on syntax error
-            tb: Optional[TracebackType] = None if is_syntax else exc_tuple[2]
-
-            # determine correct tb_offset
-            compiled = tb_data.get("running_compiled_code", False)
-            tb_offset = tb_data.get("tb_offset", 1 if compiled else 0)
-            # remove ipython internal frames from trace with tb_offset
-            for _ in range(tb_offset):
-                if tb is None:
-                    break
-                tb = tb.tb_next
-
-            excepthook(exc_tuple[0], exc_tuple[1], tb)
-            tb_data = {}  # clear data upon usage
-
-        # replace _showtraceback instead of showtraceback to allow ipython features such as debugging to work
-        # this is also what the ipython docs recommends to modify when subclassing InteractiveShell
-        ip._showtraceback = ipy_display_traceback
-        # add wrapper to capture tb_data
-        ip.showtraceback = ipy_show_traceback
-        ip.showsyntaxerror = lambda *args, **kwargs: ipy_display_traceback(
-            *args, is_syntax=True, **kwargs
-        )
-
-    try:  # pragma: no cover
-        # if within ipython, use customized traceback
-        ip = get_ipython()  # type: ignore[name-defined]
-        ipy_excepthook_closure(ip)
-        return sys.excepthook
-    except Exception:
-        # otherwise use default system hook
-        old_excepthook = sys.excepthook
-        sys.excepthook = excepthook
-        return old_excepthook
-
-
-@dataclass
-class Frame:
-    filename: str
-    lineno: int
-    name: str
-    line: str = ""
-    locals: Optional[Dict[str, pretty.Node]] = None
-    last_instruction: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] = None
-
-
-@dataclass
-class _SyntaxError:
-    offset: int
-    filename: str
-    line: str
-    lineno: int
-    msg: str
-    notes: List[str] = field(default_factory=list)
-
-
-@dataclass
-class Stack:
-    exc_type: str
-    exc_value: str
-    syntax_error: Optional[_SyntaxError] = None
-    is_cause: bool = False
-    frames: List[Frame] = field(default_factory=list)
-    notes: List[str] = field(default_factory=list)
-    is_group: bool = False
-    exceptions: List["Trace"] = field(default_factory=list)
-
-
-@dataclass
-class Trace:
-    stacks: List[Stack]
-
-
-class PathHighlighter(RegexHighlighter):
-    highlights = [r"(?P.*/)(?P.+)"]
-
-
-class Traceback:
-    """A Console renderable that renders a traceback.
-
-    Args:
-        trace (Trace, optional): A `Trace` object produced from `extract`. Defaults to None, which uses
-            the last exception.
-        width (Optional[int], optional): Number of characters used to traceback. Defaults to 100.
-        code_width (Optional[int], optional): Number of code characters used to traceback. Defaults to 88.
-        extra_lines (int, optional): Additional lines of code to render. Defaults to 3.
-        theme (str, optional): Override pygments theme used in traceback.
-        word_wrap (bool, optional): Enable word wrapping of long lines. Defaults to False.
-        show_locals (bool, optional): Enable display of local variables. Defaults to False.
-        indent_guides (bool, optional): Enable indent guides in code and locals. Defaults to True.
-        locals_max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
-            Defaults to 10.
-        locals_max_string (int, optional): Maximum length of string before truncating, or None to disable. Defaults to 80.
-        locals_hide_dunder (bool, optional): Hide locals prefixed with double underscore. Defaults to True.
-        locals_hide_sunder (bool, optional): Hide locals prefixed with single underscore. Defaults to False.
-        suppress (Sequence[Union[str, ModuleType]]): Optional sequence of modules or paths to exclude from traceback.
-        max_frames (int): Maximum number of frames to show in a traceback, 0 for no maximum. Defaults to 100.
-
-    """
-
-    LEXERS = {
-        "": "text",
-        ".py": "python",
-        ".pxd": "cython",
-        ".pyx": "cython",
-        ".pxi": "pyrex",
-    }
-
-    def __init__(
-        self,
-        trace: Optional[Trace] = None,
-        *,
-        width: Optional[int] = 100,
-        code_width: Optional[int] = 88,
-        extra_lines: int = 3,
-        theme: Optional[str] = None,
-        word_wrap: bool = False,
-        show_locals: bool = False,
-        locals_max_length: int = LOCALS_MAX_LENGTH,
-        locals_max_string: int = LOCALS_MAX_STRING,
-        locals_hide_dunder: bool = True,
-        locals_hide_sunder: bool = False,
-        indent_guides: bool = True,
-        suppress: Iterable[Union[str, ModuleType]] = (),
-        max_frames: int = 100,
-    ):
-        if trace is None:
-            exc_type, exc_value, traceback = sys.exc_info()
-            if exc_type is None or exc_value is None or traceback is None:
-                raise ValueError(
-                    "Value for 'trace' required if not called in except: block"
-                )
-            trace = self.extract(
-                exc_type, exc_value, traceback, show_locals=show_locals
-            )
-        self.trace = trace
-        self.width = width
-        self.code_width = code_width
-        self.extra_lines = extra_lines
-        self.theme = Syntax.get_theme(theme or "ansi_dark")
-        self.word_wrap = word_wrap
-        self.show_locals = show_locals
-        self.indent_guides = indent_guides
-        self.locals_max_length = locals_max_length
-        self.locals_max_string = locals_max_string
-        self.locals_hide_dunder = locals_hide_dunder
-        self.locals_hide_sunder = locals_hide_sunder
-
-        self.suppress: Sequence[str] = []
-        for suppress_entity in suppress:
-            if not isinstance(suppress_entity, str):
-                assert (
-                    suppress_entity.__file__ is not None
-                ), f"{suppress_entity!r} must be a module with '__file__' attribute"
-                path = os.path.dirname(suppress_entity.__file__)
-            else:
-                path = suppress_entity
-            path = os.path.normpath(os.path.abspath(path))
-            self.suppress.append(path)
-        self.max_frames = max(4, max_frames) if max_frames > 0 else 0
-
-    @classmethod
-    def from_exception(
-        cls,
-        exc_type: Type[Any],
-        exc_value: BaseException,
-        traceback: Optional[TracebackType],
-        *,
-        width: Optional[int] = 100,
-        code_width: Optional[int] = 88,
-        extra_lines: int = 3,
-        theme: Optional[str] = None,
-        word_wrap: bool = False,
-        show_locals: bool = False,
-        locals_max_length: int = LOCALS_MAX_LENGTH,
-        locals_max_string: int = LOCALS_MAX_STRING,
-        locals_hide_dunder: bool = True,
-        locals_hide_sunder: bool = False,
-        indent_guides: bool = True,
-        suppress: Iterable[Union[str, ModuleType]] = (),
-        max_frames: int = 100,
-    ) -> "Traceback":
-        """Create a traceback from exception info
-
-        Args:
-            exc_type (Type[BaseException]): Exception type.
-            exc_value (BaseException): Exception value.
-            traceback (TracebackType): Python Traceback object.
-            width (Optional[int], optional): Number of characters used to traceback. Defaults to 100.
-            code_width (Optional[int], optional): Number of code characters used to traceback. Defaults to 88.
-            extra_lines (int, optional): Additional lines of code to render. Defaults to 3.
-            theme (str, optional): Override pygments theme used in traceback.
-            word_wrap (bool, optional): Enable word wrapping of long lines. Defaults to False.
-            show_locals (bool, optional): Enable display of local variables. Defaults to False.
-            indent_guides (bool, optional): Enable indent guides in code and locals. Defaults to True.
-            locals_max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
-                Defaults to 10.
-            locals_max_string (int, optional): Maximum length of string before truncating, or None to disable. Defaults to 80.
-            locals_hide_dunder (bool, optional): Hide locals prefixed with double underscore. Defaults to True.
-            locals_hide_sunder (bool, optional): Hide locals prefixed with single underscore. Defaults to False.
-            suppress (Iterable[Union[str, ModuleType]]): Optional sequence of modules or paths to exclude from traceback.
-            max_frames (int): Maximum number of frames to show in a traceback, 0 for no maximum. Defaults to 100.
-
-        Returns:
-            Traceback: A Traceback instance that may be printed.
-        """
-        rich_traceback = cls.extract(
-            exc_type,
-            exc_value,
-            traceback,
-            show_locals=show_locals,
-            locals_max_length=locals_max_length,
-            locals_max_string=locals_max_string,
-            locals_hide_dunder=locals_hide_dunder,
-            locals_hide_sunder=locals_hide_sunder,
-        )
-
-        return cls(
-            rich_traceback,
-            width=width,
-            code_width=code_width,
-            extra_lines=extra_lines,
-            theme=theme,
-            word_wrap=word_wrap,
-            show_locals=show_locals,
-            indent_guides=indent_guides,
-            locals_max_length=locals_max_length,
-            locals_max_string=locals_max_string,
-            locals_hide_dunder=locals_hide_dunder,
-            locals_hide_sunder=locals_hide_sunder,
-            suppress=suppress,
-            max_frames=max_frames,
-        )
-
-    @classmethod
-    def extract(
-        cls,
-        exc_type: Type[BaseException],
-        exc_value: BaseException,
-        traceback: Optional[TracebackType],
-        *,
-        show_locals: bool = False,
-        locals_max_length: int = LOCALS_MAX_LENGTH,
-        locals_max_string: int = LOCALS_MAX_STRING,
-        locals_hide_dunder: bool = True,
-        locals_hide_sunder: bool = False,
-    ) -> Trace:
-        """Extract traceback information.
-
-        Args:
-            exc_type (Type[BaseException]): Exception type.
-            exc_value (BaseException): Exception value.
-            traceback (TracebackType): Python Traceback object.
-            show_locals (bool, optional): Enable display of local variables. Defaults to False.
-            locals_max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
-                Defaults to 10.
-            locals_max_string (int, optional): Maximum length of string before truncating, or None to disable. Defaults to 80.
-            locals_hide_dunder (bool, optional): Hide locals prefixed with double underscore. Defaults to True.
-            locals_hide_sunder (bool, optional): Hide locals prefixed with single underscore. Defaults to False.
-
-        Returns:
-            Trace: A Trace instance which you can use to construct a `Traceback`.
-        """
-
-        stacks: List[Stack] = []
-        is_cause = False
-
-        from rich import _IMPORT_CWD
-
-        notes: List[str] = getattr(exc_value, "__notes__", None) or []
-
-        def safe_str(_object: Any) -> str:
-            """Don't allow exceptions from __str__ to propagate."""
-            try:
-                return str(_object)
-            except Exception:
-                return ""
-
-        while True:
-            stack = Stack(
-                exc_type=safe_str(exc_type.__name__),
-                exc_value=safe_str(exc_value),
-                is_cause=is_cause,
-                notes=notes,
-            )
-
-            if sys.version_info >= (3, 11):
-                if isinstance(exc_value, (BaseExceptionGroup, ExceptionGroup)):
-                    stack.is_group = True
-                    for exception in exc_value.exceptions:
-                        stack.exceptions.append(
-                            Traceback.extract(
-                                type(exception),
-                                exception,
-                                exception.__traceback__,
-                                show_locals=show_locals,
-                                locals_max_length=locals_max_length,
-                                locals_hide_dunder=locals_hide_dunder,
-                                locals_hide_sunder=locals_hide_sunder,
-                            )
-                        )
-
-            if isinstance(exc_value, SyntaxError):
-                stack.syntax_error = _SyntaxError(
-                    offset=exc_value.offset or 0,
-                    filename=exc_value.filename or "?",
-                    lineno=exc_value.lineno or 0,
-                    line=exc_value.text or "",
-                    msg=exc_value.msg,
-                    notes=notes,
-                )
-
-            stacks.append(stack)
-            append = stack.frames.append
-
-            def get_locals(
-                iter_locals: Iterable[Tuple[str, object]],
-            ) -> Iterable[Tuple[str, object]]:
-                """Extract locals from an iterator of key pairs."""
-                if not (locals_hide_dunder or locals_hide_sunder):
-                    yield from iter_locals
-                    return
-                for key, value in iter_locals:
-                    if locals_hide_dunder and key.startswith("__"):
-                        continue
-                    if locals_hide_sunder and key.startswith("_"):
-                        continue
-                    yield key, value
-
-            for frame_summary, line_no in walk_tb(traceback):
-                filename = frame_summary.f_code.co_filename
-
-                last_instruction: Optional[Tuple[Tuple[int, int], Tuple[int, int]]]
-                last_instruction = None
-                if sys.version_info >= (3, 11):
-                    instruction_index = frame_summary.f_lasti // 2
-                    instruction_position = next(
-                        islice(
-                            frame_summary.f_code.co_positions(),
-                            instruction_index,
-                            instruction_index + 1,
-                        )
-                    )
-                    (
-                        start_line,
-                        end_line,
-                        start_column,
-                        end_column,
-                    ) = instruction_position
-                    if (
-                        start_line is not None
-                        and end_line is not None
-                        and start_column is not None
-                        and end_column is not None
-                    ):
-                        last_instruction = (
-                            (start_line, start_column),
-                            (end_line, end_column),
-                        )
-
-                if filename and not filename.startswith("<"):
-                    if not os.path.isabs(filename):
-                        filename = os.path.join(_IMPORT_CWD, filename)
-                if frame_summary.f_locals.get("_rich_traceback_omit", False):
-                    continue
-
-                frame = Frame(
-                    filename=filename or "?",
-                    lineno=line_no,
-                    name=frame_summary.f_code.co_name,
-                    locals=(
-                        {
-                            key: pretty.traverse(
-                                value,
-                                max_length=locals_max_length,
-                                max_string=locals_max_string,
-                            )
-                            for key, value in get_locals(frame_summary.f_locals.items())
-                            if not (inspect.isfunction(value) or inspect.isclass(value))
-                        }
-                        if show_locals
-                        else None
-                    ),
-                    last_instruction=last_instruction,
-                )
-                append(frame)
-                if frame_summary.f_locals.get("_rich_traceback_guard", False):
-                    del stack.frames[:]
-
-            cause = getattr(exc_value, "__cause__", None)
-            if cause:
-                exc_type = cause.__class__
-                exc_value = cause
-                # __traceback__ can be None, e.g. for exceptions raised by the
-                # 'multiprocessing' module
-                traceback = cause.__traceback__
-                is_cause = True
-                continue
-
-            cause = exc_value.__context__
-            if cause and not getattr(exc_value, "__suppress_context__", False):
-                exc_type = cause.__class__
-                exc_value = cause
-                traceback = cause.__traceback__
-                is_cause = False
-                continue
-            # No cover, code is reached but coverage doesn't recognize it.
-            break  # pragma: no cover
-
-        trace = Trace(stacks=stacks)
-
-        return trace
-
-    def __rich_console__(
-        self, console: Console, options: ConsoleOptions
-    ) -> RenderResult:
-        theme = self.theme
-        background_style = theme.get_background_style()
-        token_style = theme.get_style_for_token
-
-        traceback_theme = Theme(
-            {
-                "pretty": token_style(TextToken),
-                "pygments.text": token_style(Token),
-                "pygments.string": token_style(String),
-                "pygments.function": token_style(Name.Function),
-                "pygments.number": token_style(Number),
-                "repr.indent": token_style(Comment) + Style(dim=True),
-                "repr.str": token_style(String),
-                "repr.brace": token_style(TextToken) + Style(bold=True),
-                "repr.number": token_style(Number),
-                "repr.bool_true": token_style(Keyword.Constant),
-                "repr.bool_false": token_style(Keyword.Constant),
-                "repr.none": token_style(Keyword.Constant),
-                "scope.border": token_style(String.Delimiter),
-                "scope.equals": token_style(Operator),
-                "scope.key": token_style(Name),
-                "scope.key.special": token_style(Name.Constant) + Style(dim=True),
-            },
-            inherit=False,
-        )
-
-        highlighter = ReprHighlighter()
-
-        @group()
-        def render_stack(stack: Stack, last: bool) -> RenderResult:
-            if stack.frames:
-                stack_renderable: ConsoleRenderable = Panel(
-                    self._render_stack(stack),
-                    title="[traceback.title]Traceback [dim](most recent call last)",
-                    style=background_style,
-                    border_style="traceback.border",
-                    expand=True,
-                    padding=(0, 1),
-                )
-                stack_renderable = Constrain(stack_renderable, self.width)
-                with console.use_theme(traceback_theme):
-                    yield stack_renderable
-
-            if stack.syntax_error is not None:
-                with console.use_theme(traceback_theme):
-                    yield Constrain(
-                        Panel(
-                            self._render_syntax_error(stack.syntax_error),
-                            style=background_style,
-                            border_style="traceback.border.syntax_error",
-                            expand=True,
-                            padding=(0, 1),
-                            width=self.width,
-                        ),
-                        self.width,
-                    )
-                yield Text.assemble(
-                    (f"{stack.exc_type}: ", "traceback.exc_type"),
-                    highlighter(stack.syntax_error.msg),
-                )
-            elif stack.exc_value:
-                yield Text.assemble(
-                    (f"{stack.exc_type}: ", "traceback.exc_type"),
-                    highlighter(stack.exc_value),
-                )
-            else:
-                yield Text.assemble((f"{stack.exc_type}", "traceback.exc_type"))
-
-            for note in stack.notes:
-                yield Text.assemble(("[NOTE] ", "traceback.note"), highlighter(note))
-
-            if stack.is_group:
-                for group_no, group_exception in enumerate(stack.exceptions, 1):
-                    grouped_exceptions: List[Group] = []
-                    for group_last, group_stack in loop_last(group_exception.stacks):
-                        grouped_exceptions.append(render_stack(group_stack, group_last))
-                    yield ""
-                    yield Constrain(
-                        Panel(
-                            Group(*grouped_exceptions),
-                            title=f"Sub-exception #{group_no}",
-                            border_style="traceback.group.border",
-                        ),
-                        self.width,
-                    )
-
-            if not last:
-                if stack.is_cause:
-                    yield Text.from_markup(
-                        "\n[i]The above exception was the direct cause of the following exception:\n",
-                    )
-                else:
-                    yield Text.from_markup(
-                        "\n[i]During handling of the above exception, another exception occurred:\n",
-                    )
-
-        for last, stack in loop_last(reversed(self.trace.stacks)):
-            yield render_stack(stack, last)
-
-    @group()
-    def _render_syntax_error(self, syntax_error: _SyntaxError) -> RenderResult:
-        highlighter = ReprHighlighter()
-        path_highlighter = PathHighlighter()
-        if syntax_error.filename != "":
-            if os.path.exists(syntax_error.filename):
-                text = Text.assemble(
-                    (f" {syntax_error.filename}", "pygments.string"),
-                    (":", "pygments.text"),
-                    (str(syntax_error.lineno), "pygments.number"),
-                    style="pygments.text",
-                )
-                yield path_highlighter(text)
-        syntax_error_text = highlighter(syntax_error.line.rstrip())
-        syntax_error_text.no_wrap = True
-        offset = min(syntax_error.offset - 1, len(syntax_error_text))
-        syntax_error_text.stylize("bold underline", offset, offset)
-        syntax_error_text += Text.from_markup(
-            "\n" + " " * offset + "[traceback.offset]▲[/]",
-            style="pygments.text",
-        )
-        yield syntax_error_text
-
-    @classmethod
-    def _guess_lexer(cls, filename: str, code: str) -> str:
-        ext = os.path.splitext(filename)[-1]
-        if not ext:
-            # No extension, look at first line to see if it is a hashbang
-            # Note, this is an educated guess and not a guarantee
-            # If it fails, the only downside is that the code is highlighted strangely
-            new_line_index = code.index("\n")
-            first_line = code[:new_line_index] if new_line_index != -1 else code
-            if first_line.startswith("#!") and "python" in first_line.lower():
-                return "python"
-        try:
-            return cls.LEXERS.get(ext) or guess_lexer_for_filename(filename, code).name
-        except ClassNotFound:
-            return "text"
-
-    @group()
-    def _render_stack(self, stack: Stack) -> RenderResult:
-        path_highlighter = PathHighlighter()
-        theme = self.theme
-
-        def render_locals(frame: Frame) -> Iterable[ConsoleRenderable]:
-            if frame.locals:
-                yield render_scope(
-                    frame.locals,
-                    title="locals",
-                    indent_guides=self.indent_guides,
-                    max_length=self.locals_max_length,
-                    max_string=self.locals_max_string,
-                )
-
-        exclude_frames: Optional[range] = None
-        if self.max_frames != 0:
-            exclude_frames = range(
-                self.max_frames // 2,
-                len(stack.frames) - self.max_frames // 2,
-            )
-
-        excluded = False
-        for frame_index, frame in enumerate(stack.frames):
-            if exclude_frames and frame_index in exclude_frames:
-                excluded = True
-                continue
-
-            if excluded:
-                assert exclude_frames is not None
-                yield Text(
-                    f"\n... {len(exclude_frames)} frames hidden ...",
-                    justify="center",
-                    style="traceback.error",
-                )
-                excluded = False
-
-            first = frame_index == 0
-            frame_filename = frame.filename
-            suppressed = any(frame_filename.startswith(path) for path in self.suppress)
-
-            if os.path.exists(frame.filename):
-                text = Text.assemble(
-                    path_highlighter(Text(frame.filename, style="pygments.string")),
-                    (":", "pygments.text"),
-                    (str(frame.lineno), "pygments.number"),
-                    " in ",
-                    (frame.name, "pygments.function"),
-                    style="pygments.text",
-                )
-            else:
-                text = Text.assemble(
-                    "in ",
-                    (frame.name, "pygments.function"),
-                    (":", "pygments.text"),
-                    (str(frame.lineno), "pygments.number"),
-                    style="pygments.text",
-                )
-            if not frame.filename.startswith("<") and not first:
-                yield ""
-            yield text
-            if frame.filename.startswith("<"):
-                yield from render_locals(frame)
-                continue
-            if not suppressed:
-                try:
-                    code_lines = linecache.getlines(frame.filename)
-                    code = "".join(code_lines)
-                    if not code:
-                        # code may be an empty string if the file doesn't exist, OR
-                        # if the traceback filename is generated dynamically
-                        continue
-                    lexer_name = self._guess_lexer(frame.filename, code)
-                    syntax = Syntax(
-                        code,
-                        lexer_name,
-                        theme=theme,
-                        line_numbers=True,
-                        line_range=(
-                            frame.lineno - self.extra_lines,
-                            frame.lineno + self.extra_lines,
-                        ),
-                        highlight_lines={frame.lineno},
-                        word_wrap=self.word_wrap,
-                        code_width=self.code_width,
-                        indent_guides=self.indent_guides,
-                        dedent=False,
-                    )
-                    yield ""
-                except Exception as error:
-                    yield Text.assemble(
-                        (f"\n{error}", "traceback.error"),
-                    )
-                else:
-                    if frame.last_instruction is not None:
-                        start, end = frame.last_instruction
-
-                        # Stylize a line at a time
-                        # So that indentation isn't underlined (which looks bad)
-                        for line1, column1, column2 in _iter_syntax_lines(start, end):
-                            try:
-                                if column1 == 0:
-                                    line = code_lines[line1 - 1]
-                                    column1 = len(line) - len(line.lstrip())
-                                if column2 == -1:
-                                    column2 = len(code_lines[line1 - 1])
-                            except IndexError:
-                                # Being defensive here
-                                # If last_instruction reports a line out-of-bounds, we don't want to crash
-                                continue
-
-                            syntax.stylize_range(
-                                style="traceback.error_range",
-                                start=(line1, column1),
-                                end=(line1, column2),
-                            )
-                    yield (
-                        Columns(
-                            [
-                                syntax,
-                                *render_locals(frame),
-                            ],
-                            padding=1,
-                        )
-                        if frame.locals
-                        else syntax
-                    )
-
-
-if __name__ == "__main__":  # pragma: no cover
-    install(show_locals=True)
-    import sys
-
-    def bar(
-        a: Any,
-    ) -> None:  # 这是对亚洲语言支持的测试。面对模棱两可的想法，拒绝猜测的诱惑
-        one = 1
-        print(one / a)
-
-    def foo(a: Any) -> None:
-        _rich_traceback_guard = True
-        zed = {
-            "characters": {
-                "Paul Atreides",
-                "Vladimir Harkonnen",
-                "Thufir Hawat",
-                "Duncan Idaho",
-            },
-            "atomic_types": (None, False, True),
-        }
-        bar(a)
-
-    def error() -> None:
-        foo(0)
-
-    error()
diff --git a/venv/lib/python3.10/site-packages/rich/tree.py b/venv/lib/python3.10/site-packages/rich/tree.py
deleted file mode 100644
index 9a87d60ded6fc3c8a6e678644dbf2931005dc71f..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/rich/tree.py
+++ /dev/null
@@ -1,257 +0,0 @@
-from typing import Iterator, List, Optional, Tuple
-
-from ._loop import loop_first, loop_last
-from .console import Console, ConsoleOptions, RenderableType, RenderResult
-from .jupyter import JupyterMixin
-from .measure import Measurement
-from .segment import Segment
-from .style import Style, StyleStack, StyleType
-from .styled import Styled
-
-GuideType = Tuple[str, str, str, str]
-
-
-class Tree(JupyterMixin):
-    """A renderable for a tree structure.
-
-    Attributes:
-        ASCII_GUIDES (GuideType): Guide lines used when Console.ascii_only is True.
-        TREE_GUIDES (List[GuideType, GuideType, GuideType]): Default guide lines.
-
-    Args:
-        label (RenderableType): The renderable or str for the tree label.
-        style (StyleType, optional): Style of this tree. Defaults to "tree".
-        guide_style (StyleType, optional): Style of the guide lines. Defaults to "tree.line".
-        expanded (bool, optional): Also display children. Defaults to True.
-        highlight (bool, optional): Highlight renderable (if str). Defaults to False.
-        hide_root (bool, optional): Hide the root node. Defaults to False.
-    """
-
-    ASCII_GUIDES = ("    ", "|   ", "+-- ", "`-- ")
-    TREE_GUIDES = [
-        ("    ", "│   ", "├── ", "└── "),
-        ("    ", "┃   ", "┣━━ ", "┗━━ "),
-        ("    ", "║   ", "╠══ ", "╚══ "),
-    ]
-
-    def __init__(
-        self,
-        label: RenderableType,
-        *,
-        style: StyleType = "tree",
-        guide_style: StyleType = "tree.line",
-        expanded: bool = True,
-        highlight: bool = False,
-        hide_root: bool = False,
-    ) -> None:
-        self.label = label
-        self.style = style
-        self.guide_style = guide_style
-        self.children: List[Tree] = []
-        self.expanded = expanded
-        self.highlight = highlight
-        self.hide_root = hide_root
-
-    def add(
-        self,
-        label: RenderableType,
-        *,
-        style: Optional[StyleType] = None,
-        guide_style: Optional[StyleType] = None,
-        expanded: bool = True,
-        highlight: Optional[bool] = False,
-    ) -> "Tree":
-        """Add a child tree.
-
-        Args:
-            label (RenderableType): The renderable or str for the tree label.
-            style (StyleType, optional): Style of this tree. Defaults to "tree".
-            guide_style (StyleType, optional): Style of the guide lines. Defaults to "tree.line".
-            expanded (bool, optional): Also display children. Defaults to True.
-            highlight (Optional[bool], optional): Highlight renderable (if str). Defaults to False.
-
-        Returns:
-            Tree: A new child Tree, which may be further modified.
-        """
-        node = Tree(
-            label,
-            style=self.style if style is None else style,
-            guide_style=self.guide_style if guide_style is None else guide_style,
-            expanded=expanded,
-            highlight=self.highlight if highlight is None else highlight,
-        )
-        self.children.append(node)
-        return node
-
-    def __rich_console__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "RenderResult":
-        stack: List[Iterator[Tuple[bool, Tree]]] = []
-        pop = stack.pop
-        push = stack.append
-        new_line = Segment.line()
-
-        get_style = console.get_style
-        null_style = Style.null()
-        guide_style = get_style(self.guide_style, default="") or null_style
-        SPACE, CONTINUE, FORK, END = range(4)
-
-        _Segment = Segment
-
-        def make_guide(index: int, style: Style) -> Segment:
-            """Make a Segment for a level of the guide lines."""
-            if options.ascii_only:
-                line = self.ASCII_GUIDES[index]
-            else:
-                guide = 1 if style.bold else (2 if style.underline2 else 0)
-                line = self.TREE_GUIDES[0 if options.legacy_windows else guide][index]
-            return _Segment(line, style)
-
-        levels: List[Segment] = [make_guide(CONTINUE, guide_style)]
-        push(iter(loop_last([self])))
-
-        guide_style_stack = StyleStack(get_style(self.guide_style))
-        style_stack = StyleStack(get_style(self.style))
-        remove_guide_styles = Style(bold=False, underline2=False)
-
-        depth = 0
-
-        while stack:
-            stack_node = pop()
-            try:
-                last, node = next(stack_node)
-            except StopIteration:
-                levels.pop()
-                if levels:
-                    guide_style = levels[-1].style or null_style
-                    levels[-1] = make_guide(FORK, guide_style)
-                    guide_style_stack.pop()
-                    style_stack.pop()
-                continue
-            push(stack_node)
-            if last:
-                levels[-1] = make_guide(END, levels[-1].style or null_style)
-
-            guide_style = guide_style_stack.current + get_style(node.guide_style)
-            style = style_stack.current + get_style(node.style)
-            prefix = levels[(2 if self.hide_root else 1) :]
-            renderable_lines = console.render_lines(
-                Styled(node.label, style),
-                options.update(
-                    width=options.max_width
-                    - sum(level.cell_length for level in prefix),
-                    highlight=self.highlight,
-                    height=None,
-                ),
-                pad=options.justify is not None,
-            )
-
-            if not (depth == 0 and self.hide_root):
-                for first, line in loop_first(renderable_lines):
-                    if prefix:
-                        yield from _Segment.apply_style(
-                            prefix,
-                            style.background_style,
-                            post_style=remove_guide_styles,
-                        )
-                    yield from line
-                    yield new_line
-                    if first and prefix:
-                        prefix[-1] = make_guide(
-                            SPACE if last else CONTINUE, prefix[-1].style or null_style
-                        )
-
-            if node.expanded and node.children:
-                levels[-1] = make_guide(
-                    SPACE if last else CONTINUE, levels[-1].style or null_style
-                )
-                levels.append(
-                    make_guide(END if len(node.children) == 1 else FORK, guide_style)
-                )
-                style_stack.push(get_style(node.style))
-                guide_style_stack.push(get_style(node.guide_style))
-                push(iter(loop_last(node.children)))
-                depth += 1
-
-    def __rich_measure__(
-        self, console: "Console", options: "ConsoleOptions"
-    ) -> "Measurement":
-        stack: List[Iterator[Tree]] = [iter([self])]
-        pop = stack.pop
-        push = stack.append
-        minimum = 0
-        maximum = 0
-        measure = Measurement.get
-        level = 0
-        while stack:
-            iter_tree = pop()
-            try:
-                tree = next(iter_tree)
-            except StopIteration:
-                level -= 1
-                continue
-            push(iter_tree)
-            min_measure, max_measure = measure(console, options, tree.label)
-            indent = level * 4
-            minimum = max(min_measure + indent, minimum)
-            maximum = max(max_measure + indent, maximum)
-            if tree.expanded and tree.children:
-                push(iter(tree.children))
-                level += 1
-        return Measurement(minimum, maximum)
-
-
-if __name__ == "__main__":  # pragma: no cover
-    from rich.console import Group
-    from rich.markdown import Markdown
-    from rich.panel import Panel
-    from rich.syntax import Syntax
-    from rich.table import Table
-
-    table = Table(row_styles=["", "dim"])
-
-    table.add_column("Released", style="cyan", no_wrap=True)
-    table.add_column("Title", style="magenta")
-    table.add_column("Box Office", justify="right", style="green")
-
-    table.add_row("Dec 20, 2019", "Star Wars: The Rise of Skywalker", "$952,110,690")
-    table.add_row("May 25, 2018", "Solo: A Star Wars Story", "$393,151,347")
-    table.add_row("Dec 15, 2017", "Star Wars Ep. V111: The Last Jedi", "$1,332,539,889")
-    table.add_row("Dec 16, 2016", "Rogue One: A Star Wars Story", "$1,332,439,889")
-
-    code = """\
-class Segment(NamedTuple):
-    text: str = ""
-    style: Optional[Style] = None
-    is_control: bool = False
-"""
-    syntax = Syntax(code, "python", theme="monokai", line_numbers=True)
-
-    markdown = Markdown(
-        """\
-### example.md
-> Hello, World!
->
-> Markdown _all_ the things
-"""
-    )
-
-    root = Tree("🌲 [b green]Rich Tree", highlight=True, hide_root=True)
-
-    node = root.add(":file_folder: Renderables", guide_style="red")
-    simple_node = node.add(":file_folder: [bold yellow]Atomic", guide_style="uu green")
-    simple_node.add(Group("📄 Syntax", syntax))
-    simple_node.add(Group("📄 Markdown", Panel(markdown, border_style="green")))
-
-    containers_node = node.add(
-        ":file_folder: [bold magenta]Containers", guide_style="bold magenta"
-    )
-    containers_node.expanded = True
-    panel = Panel.fit("Just a panel", border_style="red")
-    containers_node.add(Group("📄 Panels", panel))
-
-    containers_node.add(Group("📄 [b magenta]Table", table))
-
-    console = Console()
-
-    console.print(root)
diff --git a/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/INSTALLER
deleted file mode 100644
index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/INSTALLER
+++ /dev/null
@@ -1 +0,0 @@
-pip
diff --git a/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/LICENSE b/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/LICENSE
deleted file mode 100644
index b9077766e9b9bdcae49ea5c8fced750ed13ec8f7..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/LICENSE
+++ /dev/null
@@ -1,13 +0,0 @@
-Copyright (c) 2018, Tzu-ping Chung 
-
-Permission to use, copy, modify, and distribute this software for any
-purpose with or without fee is hereby granted, provided that the above
-copyright notice and this permission notice appear in all copies.
-
-THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
-WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
-MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
-ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
-WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
-ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
-OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
diff --git a/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/METADATA b/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/METADATA
deleted file mode 100644
index 52118f1e5c83bd7ef39196a749651fc87d176812..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/METADATA
+++ /dev/null
@@ -1,106 +0,0 @@
-Metadata-Version: 2.1
-Name: shellingham
-Version: 1.5.4
-Summary: Tool to Detect Surrounding Shell
-Home-page: https://github.com/sarugaku/shellingham
-Author: Tzu-ping Chung
-Author-email: uranusjr@gmail.com
-License: ISC License
-Keywords: shell
-Classifier: Development Status :: 3 - Alpha
-Classifier: Environment :: Console
-Classifier: Intended Audience :: Developers
-Classifier: License :: OSI Approved :: ISC License (ISCL)
-Classifier: Operating System :: OS Independent
-Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Python: >=3.7
-Description-Content-Type: text/x-rst
-License-File: LICENSE
-
-=============================================
-Shellingham: Tool to Detect Surrounding Shell
-=============================================
-
-.. image:: https://img.shields.io/pypi/v/shellingham.svg
-    :target: https://pypi.org/project/shellingham/
-
-Shellingham detects what shell the current Python executable is running in.
-
-
-Usage
-=====
-
-.. code-block:: python
-
-    >>> import shellingham
-    >>> shellingham.detect_shell()
-    ('bash', '/bin/bash')
-
-``detect_shell`` pokes around the process's running environment to determine
-what shell it is run in. It returns a 2-tuple:
-
-* The shell name, always lowercased.
-* The command used to run the shell.
-
-``ShellDetectionFailure`` is raised if ``detect_shell`` fails to detect the
-surrounding shell.
-
-
-Notes
-=====
-
-* The shell name is always lowercased.
-* On Windows, the shell name is the name of the executable, minus the file
-  extension.
-
-
-Notes for Application Developers
-================================
-
-Remember, your application's user is not necessarily using a shell.
-Shellingham raises ``ShellDetectionFailure`` if there is no shell to detect,
-but *your application should almost never do this to your user*.
-
-A practical approach to this is to wrap ``detect_shell`` in a try block, and
-provide a sane default on failure
-
-.. code-block:: python
-
-    try:
-        shell = shellingham.detect_shell()
-    except shellingham.ShellDetectionFailure:
-        shell = provide_default()
-
-
-There are a few choices for you to choose from.
-
-* The POSIX standard mandates the environment variable ``SHELL`` to refer to
-  "the user's preferred command language interpreter". This is always available
-  (even if the user is not in an interactive session), and likely the correct
-  choice to launch an interactive sub-shell with.
-* A command ``sh`` is almost guaranteed to exist, likely at ``/bin/sh``, since
-  several POSIX tools rely on it. This should be suitable if you want to run a
-  (possibly non-interactive) script.
-* All versions of DOS and Windows have an environment variable ``COMSPEC``.
-  This can always be used to launch a usable command prompt (e.g. `cmd.exe` on
-  Windows).
-
-Here's a simple implementation to provide a default shell
-
-.. code-block:: python
-
-    import os
-
-    def provide_default():
-        if os.name == 'posix':
-            return os.environ['SHELL']
-        elif os.name == 'nt':
-            return os.environ['COMSPEC']
-        raise NotImplementedError(f'OS {os.name!r} support not available')
diff --git a/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/RECORD b/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/RECORD
deleted file mode 100644
index a03e85926af1ddd60a801034a230b4812a4b9ffd..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/RECORD
+++ /dev/null
@@ -1,21 +0,0 @@
-shellingham-1.5.4.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
-shellingham-1.5.4.dist-info/LICENSE,sha256=84j9OMrRMRLB3A9mm76A5_hFQe26-3LzAw0sp2QsPJ0,751
-shellingham-1.5.4.dist-info/METADATA,sha256=GD2AIgo3STJieVc53TV8xbs_Sb05DMkZjVGA5UUaB_o,3461
-shellingham-1.5.4.dist-info/RECORD,,
-shellingham-1.5.4.dist-info/WHEEL,sha256=iYlv5fX357PQyRT2o6tw1bN-YcKFFHKqB_LwHO5wP-g,110
-shellingham-1.5.4.dist-info/top_level.txt,sha256=uKMQL5AKxPi4O9_Rbd838QeEs4ImpGQKNbEDZYqgBgk,12
-shellingham-1.5.4.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-shellingham/__init__.py,sha256=pAKXUPKUdwyErC0ZjS-5w-fRdSbmdcfvnpt_x1yWqtA,635
-shellingham/__pycache__/__init__.cpython-310.pyc,,
-shellingham/__pycache__/_core.cpython-310.pyc,,
-shellingham/__pycache__/nt.cpython-310.pyc,,
-shellingham/_core.py,sha256=v-CTr_7F7cJAtNnzpa1N_Hl8afkY5yiDA4joGmsUBu0,300
-shellingham/nt.py,sha256=m6J6SuwyqVVlxXT9Bc-9F_1x-T5u0gCFFrRAF2LIkeg,4516
-shellingham/posix/__init__.py,sha256=pB69qtvZJ_yIf48nl4-ZfS3wLwwuXuknXOZhBnC2T1o,3129
-shellingham/posix/__pycache__/__init__.cpython-310.pyc,,
-shellingham/posix/__pycache__/_core.cpython-310.pyc,,
-shellingham/posix/__pycache__/proc.cpython-310.pyc,,
-shellingham/posix/__pycache__/ps.cpython-310.pyc,,
-shellingham/posix/_core.py,sha256=_v18UaXbzr4muNhr3-mH1FdSdjZ_dOXQrtUyomIbKYQ,81
-shellingham/posix/proc.py,sha256=nSUxIuQSotvaDW76i0oTQAM9aZ9PXBLFAEktWljSKCo,2659
-shellingham/posix/ps.py,sha256=NGmDKCukhNp0lahwYCaMXphBYaVbhbiR9BtE0OkT8qU,1770
diff --git a/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/WHEEL b/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/WHEEL
deleted file mode 100644
index c34f1162ef9a50c355df1261ef6194ffc1b39975..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/WHEEL
+++ /dev/null
@@ -1,6 +0,0 @@
-Wheel-Version: 1.0
-Generator: bdist_wheel (0.41.2)
-Root-Is-Purelib: true
-Tag: py2-none-any
-Tag: py3-none-any
-
diff --git a/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/top_level.txt
deleted file mode 100644
index d4e44ce0299bb38463f8491ec8850910235c2709..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/top_level.txt
+++ /dev/null
@@ -1 +0,0 @@
-shellingham
diff --git a/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/zip-safe b/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/zip-safe
deleted file mode 100644
index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/shellingham-1.5.4.dist-info/zip-safe
+++ /dev/null
@@ -1 +0,0 @@
-
diff --git a/venv/lib/python3.10/site-packages/shellingham/__init__.py b/venv/lib/python3.10/site-packages/shellingham/__init__.py
deleted file mode 100644
index 15f7a90cbd02e5c2cc933cf6aa0374cca68035f1..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/shellingham/__init__.py
+++ /dev/null
@@ -1,23 +0,0 @@
-import importlib
-import os
-
-from ._core import ShellDetectionFailure
-
-__version__ = "1.5.4"
-
-
-def detect_shell(pid=None, max_depth=10):
-    name = os.name
-    try:
-        impl = importlib.import_module(".{}".format(name), __name__)
-    except ImportError:
-        message = "Shell detection not implemented for {0!r}".format(name)
-        raise RuntimeError(message)
-    try:
-        get_shell = impl.get_shell
-    except AttributeError:
-        raise RuntimeError("get_shell not implemented for {0!r}".format(name))
-    shell = get_shell(pid, max_depth=max_depth)
-    if shell:
-        return shell
-    raise ShellDetectionFailure()
diff --git a/venv/lib/python3.10/site-packages/shellingham/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/shellingham/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 79511e0ac5e7afad5608168dbe045fe6c415cee8..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/shellingham/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/shellingham/__pycache__/_core.cpython-310.pyc b/venv/lib/python3.10/site-packages/shellingham/__pycache__/_core.cpython-310.pyc
deleted file mode 100644
index ee9dd8d3372574e0bd2669920edda38a9b43130e..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/shellingham/__pycache__/_core.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/shellingham/__pycache__/nt.cpython-310.pyc b/venv/lib/python3.10/site-packages/shellingham/__pycache__/nt.cpython-310.pyc
deleted file mode 100644
index bdd56526596135d7be9d4e3a451c0ab280cf64f4..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/shellingham/__pycache__/nt.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/shellingham/_core.py b/venv/lib/python3.10/site-packages/shellingham/_core.py
deleted file mode 100644
index 13b65417c733b54e48b120e37f573c2baa6ef72b..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/shellingham/_core.py
+++ /dev/null
@@ -1,11 +0,0 @@
-SHELL_NAMES = (
-    {"sh", "bash", "dash", "ash"}  # Bourne.
-    | {"csh", "tcsh"}  # C.
-    | {"ksh", "zsh", "fish"}  # Common alternatives.
-    | {"cmd", "powershell", "pwsh"}  # Microsoft.
-    | {"elvish", "xonsh", "nu"}  # More exotic.
-)
-
-
-class ShellDetectionFailure(EnvironmentError):
-    pass
diff --git a/venv/lib/python3.10/site-packages/shellingham/nt.py b/venv/lib/python3.10/site-packages/shellingham/nt.py
deleted file mode 100644
index 389551b223a761fa2f97e929b60bf3ca5baed94c..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/shellingham/nt.py
+++ /dev/null
@@ -1,163 +0,0 @@
-import contextlib
-import ctypes
-import os
-
-from ctypes.wintypes import (
-    BOOL,
-    CHAR,
-    DWORD,
-    HANDLE,
-    LONG,
-    LPWSTR,
-    MAX_PATH,
-    PDWORD,
-    ULONG,
-)
-
-from shellingham._core import SHELL_NAMES
-
-
-INVALID_HANDLE_VALUE = HANDLE(-1).value
-ERROR_NO_MORE_FILES = 18
-ERROR_INSUFFICIENT_BUFFER = 122
-TH32CS_SNAPPROCESS = 2
-PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
-
-
-kernel32 = ctypes.windll.kernel32
-
-
-def _check_handle(error_val=0):
-    def check(ret, func, args):
-        if ret == error_val:
-            raise ctypes.WinError()
-        return ret
-
-    return check
-
-
-def _check_expected(expected):
-    def check(ret, func, args):
-        if ret:
-            return True
-        code = ctypes.GetLastError()
-        if code == expected:
-            return False
-        raise ctypes.WinError(code)
-
-    return check
-
-
-class ProcessEntry32(ctypes.Structure):
-    _fields_ = (
-        ("dwSize", DWORD),
-        ("cntUsage", DWORD),
-        ("th32ProcessID", DWORD),
-        ("th32DefaultHeapID", ctypes.POINTER(ULONG)),
-        ("th32ModuleID", DWORD),
-        ("cntThreads", DWORD),
-        ("th32ParentProcessID", DWORD),
-        ("pcPriClassBase", LONG),
-        ("dwFlags", DWORD),
-        ("szExeFile", CHAR * MAX_PATH),
-    )
-
-
-kernel32.CloseHandle.argtypes = [HANDLE]
-kernel32.CloseHandle.restype = BOOL
-
-kernel32.CreateToolhelp32Snapshot.argtypes = [DWORD, DWORD]
-kernel32.CreateToolhelp32Snapshot.restype = HANDLE
-kernel32.CreateToolhelp32Snapshot.errcheck = _check_handle(  # type: ignore
-    INVALID_HANDLE_VALUE,
-)
-
-kernel32.Process32First.argtypes = [HANDLE, ctypes.POINTER(ProcessEntry32)]
-kernel32.Process32First.restype = BOOL
-kernel32.Process32First.errcheck = _check_expected(  # type: ignore
-    ERROR_NO_MORE_FILES,
-)
-
-kernel32.Process32Next.argtypes = [HANDLE, ctypes.POINTER(ProcessEntry32)]
-kernel32.Process32Next.restype = BOOL
-kernel32.Process32Next.errcheck = _check_expected(  # type: ignore
-    ERROR_NO_MORE_FILES,
-)
-
-kernel32.GetCurrentProcessId.argtypes = []
-kernel32.GetCurrentProcessId.restype = DWORD
-
-kernel32.OpenProcess.argtypes = [DWORD, BOOL, DWORD]
-kernel32.OpenProcess.restype = HANDLE
-kernel32.OpenProcess.errcheck = _check_handle(  # type: ignore
-    INVALID_HANDLE_VALUE,
-)
-
-kernel32.QueryFullProcessImageNameW.argtypes = [HANDLE, DWORD, LPWSTR, PDWORD]
-kernel32.QueryFullProcessImageNameW.restype = BOOL
-kernel32.QueryFullProcessImageNameW.errcheck = _check_expected(  # type: ignore
-    ERROR_INSUFFICIENT_BUFFER,
-)
-
-
-@contextlib.contextmanager
-def _handle(f, *args, **kwargs):
-    handle = f(*args, **kwargs)
-    try:
-        yield handle
-    finally:
-        kernel32.CloseHandle(handle)
-
-
-def _iter_processes():
-    f = kernel32.CreateToolhelp32Snapshot
-    with _handle(f, TH32CS_SNAPPROCESS, 0) as snap:
-        entry = ProcessEntry32()
-        entry.dwSize = ctypes.sizeof(entry)
-        ret = kernel32.Process32First(snap, entry)
-        while ret:
-            yield entry
-            ret = kernel32.Process32Next(snap, entry)
-
-
-def _get_full_path(proch):
-    size = DWORD(MAX_PATH)
-    while True:
-        path_buff = ctypes.create_unicode_buffer("", size.value)
-        if kernel32.QueryFullProcessImageNameW(proch, 0, path_buff, size):
-            return path_buff.value
-        size.value *= 2
-
-
-def get_shell(pid=None, max_depth=10):
-    proc_map = {
-        proc.th32ProcessID: (proc.th32ParentProcessID, proc.szExeFile)
-        for proc in _iter_processes()
-    }
-    pid = pid or os.getpid()
-
-    for _ in range(0, max_depth + 1):
-        try:
-            ppid, executable = proc_map[pid]
-        except KeyError:  # No such process? Give up.
-            break
-
-        # The executable name would be encoded with the current code page if
-        # we're in ANSI mode (usually). Try to decode it into str/unicode,
-        # replacing invalid characters to be safe (not thoeratically necessary,
-        # I think). Note that we need to use 'mbcs' instead of encoding
-        # settings from sys because this is from the Windows API, not Python
-        # internals (which those settings reflect). (pypa/pipenv#3382)
-        if isinstance(executable, bytes):
-            executable = executable.decode("mbcs", "replace")
-
-        name = executable.rpartition(".")[0].lower()
-        if name not in SHELL_NAMES:
-            pid = ppid
-            continue
-
-        key = PROCESS_QUERY_LIMITED_INFORMATION
-        with _handle(kernel32.OpenProcess, key, 0, pid) as proch:
-            return (name, _get_full_path(proch))
-
-    return None
diff --git a/venv/lib/python3.10/site-packages/shellingham/posix/__init__.py b/venv/lib/python3.10/site-packages/shellingham/posix/__init__.py
deleted file mode 100644
index 5bd2070db27189e62a1867e4de49f16f8c8841ff..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/shellingham/posix/__init__.py
+++ /dev/null
@@ -1,112 +0,0 @@
-import os
-import re
-
-from .._core import SHELL_NAMES, ShellDetectionFailure
-from . import proc, ps
-
-# Based on QEMU docs: https://www.qemu.org/docs/master/user/main.html
-QEMU_BIN_REGEX = re.compile(
-    r"""qemu-
-        (alpha
-        |armeb
-        |arm
-        |m68k
-        |cris
-        |i386
-        |x86_64
-        |microblaze
-        |mips
-        |mipsel
-        |mips64
-        |mips64el
-        |mipsn32
-        |mipsn32el
-        |nios2
-        |ppc64
-        |ppc
-        |sh4eb
-        |sh4
-        |sparc
-        |sparc32plus
-        |sparc64
-    )""",
-    re.VERBOSE,
-)
-
-
-def _iter_process_parents(pid, max_depth=10):
-    """Select a way to obtain process information from the system.
-
-    * `/proc` is used if supported.
-    * The system `ps` utility is used as a fallback option.
-    """
-    for impl in (proc, ps):
-        try:
-            iterator = impl.iter_process_parents(pid, max_depth)
-        except EnvironmentError:
-            continue
-        return iterator
-    raise ShellDetectionFailure("compatible proc fs or ps utility is required")
-
-
-def _get_login_shell(proc_cmd):
-    """Form shell information from SHELL environ if possible."""
-    login_shell = os.environ.get("SHELL", "")
-    if login_shell:
-        proc_cmd = login_shell
-    else:
-        proc_cmd = proc_cmd[1:]
-    return (os.path.basename(proc_cmd).lower(), proc_cmd)
-
-
-_INTERPRETER_SHELL_NAMES = [
-    (re.compile(r"^python(\d+(\.\d+)?)?$"), {"xonsh"}),
-]
-
-
-def _get_interpreter_shell(proc_name, proc_args):
-    """Get shell invoked via an interpreter.
-
-    Some shells are implemented on, and invoked with an interpreter, e.g. xonsh
-    is commonly executed with an executable Python script. This detects what
-    script the interpreter is actually running, and check whether that looks
-    like a shell.
-
-    See sarugaku/shellingham#26 for rational.
-    """
-    for pattern, shell_names in _INTERPRETER_SHELL_NAMES:
-        if not pattern.match(proc_name):
-            continue
-        for arg in proc_args:
-            name = os.path.basename(arg).lower()
-            if os.path.isfile(arg) and name in shell_names:
-                return (name, arg)
-    return None
-
-
-def _get_shell(cmd, *args):
-    if cmd.startswith("-"):  # Login shell! Let's use this.
-        return _get_login_shell(cmd)
-    name = os.path.basename(cmd).lower()
-    if name == "rosetta" or QEMU_BIN_REGEX.fullmatch(name):
-        # If the current process is Rosetta or QEMU, this likely is a
-        # containerized process. Parse out the actual command instead.
-        cmd = args[0]
-        args = args[1:]
-        name = os.path.basename(cmd).lower()
-    if name in SHELL_NAMES:  # Command looks like a shell.
-        return (name, cmd)
-    shell = _get_interpreter_shell(name, args)
-    if shell:
-        return shell
-    return None
-
-
-def get_shell(pid=None, max_depth=10):
-    """Get the shell that the supplied pid or os.getpid() is running in."""
-    pid = str(pid or os.getpid())
-    for proc_args, _, _ in _iter_process_parents(pid, max_depth):
-        shell = _get_shell(*proc_args)
-        if shell:
-            return shell
-    return None
diff --git a/venv/lib/python3.10/site-packages/shellingham/posix/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/shellingham/posix/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 3f09864eb8c32d31253f2214d287f1a1d3d17a9f..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/shellingham/posix/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/shellingham/posix/__pycache__/_core.cpython-310.pyc b/venv/lib/python3.10/site-packages/shellingham/posix/__pycache__/_core.cpython-310.pyc
deleted file mode 100644
index 1a13533803dd13cb9c8cbce30a6f7a64b895a6a1..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/shellingham/posix/__pycache__/_core.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/shellingham/posix/__pycache__/proc.cpython-310.pyc b/venv/lib/python3.10/site-packages/shellingham/posix/__pycache__/proc.cpython-310.pyc
deleted file mode 100644
index 78e8c1718cb54ec52ea1b6f78b9c8a3a89873381..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/shellingham/posix/__pycache__/proc.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/shellingham/posix/__pycache__/ps.cpython-310.pyc b/venv/lib/python3.10/site-packages/shellingham/posix/__pycache__/ps.cpython-310.pyc
deleted file mode 100644
index 9759e33c200dce16e2887b84ceb99f58f5f2a531..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/shellingham/posix/__pycache__/ps.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/shellingham/posix/_core.py b/venv/lib/python3.10/site-packages/shellingham/posix/_core.py
deleted file mode 100644
index adc49e6e7a9d3edf062c55e0078136899f78d30d..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/shellingham/posix/_core.py
+++ /dev/null
@@ -1,3 +0,0 @@
-import collections
-
-Process = collections.namedtuple("Process", "args pid ppid")
diff --git a/venv/lib/python3.10/site-packages/shellingham/posix/proc.py b/venv/lib/python3.10/site-packages/shellingham/posix/proc.py
deleted file mode 100644
index 950f63228e5b328f82b70da8851ec60c6a2ff029..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/shellingham/posix/proc.py
+++ /dev/null
@@ -1,83 +0,0 @@
-import io
-import os
-import re
-import sys
-
-from ._core import Process
-
-# FreeBSD: https://www.freebsd.org/cgi/man.cgi?query=procfs
-# NetBSD: https://man.netbsd.org/NetBSD-9.3-STABLE/mount_procfs.8
-# DragonFlyBSD: https://www.dragonflybsd.org/cgi/web-man?command=procfs
-BSD_STAT_PPID = 2
-
-# See https://docs.kernel.org/filesystems/proc.html
-LINUX_STAT_PPID = 3
-
-STAT_PATTERN = re.compile(r"\(.+\)|\S+")
-
-
-def detect_proc():
-    """Detect /proc filesystem style.
-
-    This checks the /proc/{pid} directory for possible formats. Returns one of
-    the following as str:
-
-    * `stat`: Linux-style, i.e. ``/proc/{pid}/stat``.
-    * `status`: BSD-style, i.e. ``/proc/{pid}/status``.
-    """
-    pid = os.getpid()
-    for name in ("stat", "status"):
-        if os.path.exists(os.path.join("/proc", str(pid), name)):
-            return name
-    raise ProcFormatError("unsupported proc format")
-
-
-def _use_bsd_stat_format():
-    try:
-        return os.uname().sysname.lower() in ("freebsd", "netbsd", "dragonfly")
-    except Exception:
-        return False
-
-
-def _get_ppid(pid, name):
-    path = os.path.join("/proc", str(pid), name)
-    with io.open(path, encoding="ascii", errors="replace") as f:
-        parts = STAT_PATTERN.findall(f.read())
-    # We only care about TTY and PPID -- both are numbers.
-    if _use_bsd_stat_format():
-        return parts[BSD_STAT_PPID]
-    return parts[LINUX_STAT_PPID]
-
-
-def _get_cmdline(pid):
-    path = os.path.join("/proc", str(pid), "cmdline")
-    encoding = sys.getfilesystemencoding() or "utf-8"
-    with io.open(path, encoding=encoding, errors="replace") as f:
-        # XXX: Command line arguments can be arbitrary byte sequences, not
-        # necessarily decodable. For Shellingham's purpose, however, we don't
-        # care. (pypa/pipenv#2820)
-        # cmdline appends an extra NULL at the end, hence the [:-1].
-        return tuple(f.read().split("\0")[:-1])
-
-
-class ProcFormatError(EnvironmentError):
-    pass
-
-
-def iter_process_parents(pid, max_depth=10):
-    """Try to look up the process tree via the /proc interface."""
-    stat_name = detect_proc()
-
-    # Inner generator function so we correctly throw an error eagerly if proc
-    # is not supported, rather than on the first call to the iterator. This
-    # allows the call site detects the correct implementation.
-    def _iter_process_parents(pid, max_depth):
-        for _ in range(max_depth):
-            ppid = _get_ppid(pid, stat_name)
-            args = _get_cmdline(pid)
-            yield Process(args=args, pid=pid, ppid=ppid)
-            if ppid == "0":
-                break
-            pid = ppid
-
-    return _iter_process_parents(pid, max_depth)
diff --git a/venv/lib/python3.10/site-packages/shellingham/posix/ps.py b/venv/lib/python3.10/site-packages/shellingham/posix/ps.py
deleted file mode 100644
index 3bc39a74a56390c263e63bfead028f6bce4df3cb..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/shellingham/posix/ps.py
+++ /dev/null
@@ -1,51 +0,0 @@
-import errno
-import subprocess
-import sys
-
-from ._core import Process
-
-
-class PsNotAvailable(EnvironmentError):
-    pass
-
-
-def iter_process_parents(pid, max_depth=10):
-    """Try to look up the process tree via the output of `ps`."""
-    try:
-        cmd = ["ps", "-ww", "-o", "pid=", "-o", "ppid=", "-o", "args="]
-        output = subprocess.check_output(cmd)
-    except OSError as e:  # Python 2-compatible FileNotFoundError.
-        if e.errno != errno.ENOENT:
-            raise
-        raise PsNotAvailable("ps not found")
-    except subprocess.CalledProcessError as e:
-        # `ps` can return 1 if the process list is completely empty.
-        # (sarugaku/shellingham#15)
-        if not e.output.strip():
-            return
-        raise
-    if not isinstance(output, str):
-        encoding = sys.getfilesystemencoding() or sys.getdefaultencoding()
-        output = output.decode(encoding)
-
-    processes_mapping = {}
-    for line in output.split("\n"):
-        try:
-            _pid, ppid, args = line.strip().split(None, 2)
-            # XXX: This is not right, but we are really out of options.
-            # ps does not offer a sane way to decode the argument display,
-            # and this is "Good Enough" for obtaining shell names. Hopefully
-            # people don't name their shell with a space, or have something
-            # like "/usr/bin/xonsh is uber". (sarugaku/shellingham#14)
-            args = tuple(a.strip() for a in args.split(" "))
-        except ValueError:
-            continue
-        processes_mapping[_pid] = Process(args=args, pid=_pid, ppid=ppid)
-
-    for _ in range(max_depth):
-        try:
-            process = processes_mapping[pid]
-        except KeyError:
-            return
-        yield process
-        pid = process.ppid
diff --git a/venv/lib/python3.10/site-packages/simsimd.cpython-310-x86_64-linux-gnu.so b/venv/lib/python3.10/site-packages/simsimd.cpython-310-x86_64-linux-gnu.so
deleted file mode 100644
index e208faf61afcc7f218dee263e57d4a9cdcd8bb0a..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/simsimd.cpython-310-x86_64-linux-gnu.so
+++ /dev/null
@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:dcda5a17dbd92e30fd71c5258b546fe234a0f006191fa3f5fa7586cde41908d9
-size 3166321
diff --git a/venv/lib/python3.10/site-packages/six.py b/venv/lib/python3.10/site-packages/six.py
deleted file mode 100644
index 3de5969b1ad3b973342e5e88ee1770fa7c798152..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/six.py
+++ /dev/null
@@ -1,1003 +0,0 @@
-# Copyright (c) 2010-2024 Benjamin Peterson
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-"""Utilities for writing code that runs on Python 2 and 3"""
-
-from __future__ import absolute_import
-
-import functools
-import itertools
-import operator
-import sys
-import types
-
-__author__ = "Benjamin Peterson "
-__version__ = "1.17.0"
-
-
-# Useful for very coarse version differentiation.
-PY2 = sys.version_info[0] == 2
-PY3 = sys.version_info[0] == 3
-PY34 = sys.version_info[0:2] >= (3, 4)
-
-if PY3:
-    string_types = str,
-    integer_types = int,
-    class_types = type,
-    text_type = str
-    binary_type = bytes
-
-    MAXSIZE = sys.maxsize
-else:
-    string_types = basestring,
-    integer_types = (int, long)
-    class_types = (type, types.ClassType)
-    text_type = unicode
-    binary_type = str
-
-    if sys.platform.startswith("java"):
-        # Jython always uses 32 bits.
-        MAXSIZE = int((1 << 31) - 1)
-    else:
-        # It's possible to have sizeof(long) != sizeof(Py_ssize_t).
-        class X(object):
-
-            def __len__(self):
-                return 1 << 31
-        try:
-            len(X())
-        except OverflowError:
-            # 32-bit
-            MAXSIZE = int((1 << 31) - 1)
-        else:
-            # 64-bit
-            MAXSIZE = int((1 << 63) - 1)
-        del X
-
-if PY34:
-    from importlib.util import spec_from_loader
-else:
-    spec_from_loader = None
-
-
-def _add_doc(func, doc):
-    """Add documentation to a function."""
-    func.__doc__ = doc
-
-
-def _import_module(name):
-    """Import module, returning the module after the last dot."""
-    __import__(name)
-    return sys.modules[name]
-
-
-class _LazyDescr(object):
-
-    def __init__(self, name):
-        self.name = name
-
-    def __get__(self, obj, tp):
-        result = self._resolve()
-        setattr(obj, self.name, result)  # Invokes __set__.
-        try:
-            # This is a bit ugly, but it avoids running this again by
-            # removing this descriptor.
-            delattr(obj.__class__, self.name)
-        except AttributeError:
-            pass
-        return result
-
-
-class MovedModule(_LazyDescr):
-
-    def __init__(self, name, old, new=None):
-        super(MovedModule, self).__init__(name)
-        if PY3:
-            if new is None:
-                new = name
-            self.mod = new
-        else:
-            self.mod = old
-
-    def _resolve(self):
-        return _import_module(self.mod)
-
-    def __getattr__(self, attr):
-        _module = self._resolve()
-        value = getattr(_module, attr)
-        setattr(self, attr, value)
-        return value
-
-
-class _LazyModule(types.ModuleType):
-
-    def __init__(self, name):
-        super(_LazyModule, self).__init__(name)
-        self.__doc__ = self.__class__.__doc__
-
-    def __dir__(self):
-        attrs = ["__doc__", "__name__"]
-        attrs += [attr.name for attr in self._moved_attributes]
-        return attrs
-
-    # Subclasses should override this
-    _moved_attributes = []
-
-
-class MovedAttribute(_LazyDescr):
-
-    def __init__(self, name, old_mod, new_mod, old_attr=None, new_attr=None):
-        super(MovedAttribute, self).__init__(name)
-        if PY3:
-            if new_mod is None:
-                new_mod = name
-            self.mod = new_mod
-            if new_attr is None:
-                if old_attr is None:
-                    new_attr = name
-                else:
-                    new_attr = old_attr
-            self.attr = new_attr
-        else:
-            self.mod = old_mod
-            if old_attr is None:
-                old_attr = name
-            self.attr = old_attr
-
-    def _resolve(self):
-        module = _import_module(self.mod)
-        return getattr(module, self.attr)
-
-
-class _SixMetaPathImporter(object):
-
-    """
-    A meta path importer to import six.moves and its submodules.
-
-    This class implements a PEP302 finder and loader. It should be compatible
-    with Python 2.5 and all existing versions of Python3
-    """
-
-    def __init__(self, six_module_name):
-        self.name = six_module_name
-        self.known_modules = {}
-
-    def _add_module(self, mod, *fullnames):
-        for fullname in fullnames:
-            self.known_modules[self.name + "." + fullname] = mod
-
-    def _get_module(self, fullname):
-        return self.known_modules[self.name + "." + fullname]
-
-    def find_module(self, fullname, path=None):
-        if fullname in self.known_modules:
-            return self
-        return None
-
-    def find_spec(self, fullname, path, target=None):
-        if fullname in self.known_modules:
-            return spec_from_loader(fullname, self)
-        return None
-
-    def __get_module(self, fullname):
-        try:
-            return self.known_modules[fullname]
-        except KeyError:
-            raise ImportError("This loader does not know module " + fullname)
-
-    def load_module(self, fullname):
-        try:
-            # in case of a reload
-            return sys.modules[fullname]
-        except KeyError:
-            pass
-        mod = self.__get_module(fullname)
-        if isinstance(mod, MovedModule):
-            mod = mod._resolve()
-        else:
-            mod.__loader__ = self
-        sys.modules[fullname] = mod
-        return mod
-
-    def is_package(self, fullname):
-        """
-        Return true, if the named module is a package.
-
-        We need this method to get correct spec objects with
-        Python 3.4 (see PEP451)
-        """
-        return hasattr(self.__get_module(fullname), "__path__")
-
-    def get_code(self, fullname):
-        """Return None
-
-        Required, if is_package is implemented"""
-        self.__get_module(fullname)  # eventually raises ImportError
-        return None
-    get_source = get_code  # same as get_code
-
-    def create_module(self, spec):
-        return self.load_module(spec.name)
-
-    def exec_module(self, module):
-        pass
-
-_importer = _SixMetaPathImporter(__name__)
-
-
-class _MovedItems(_LazyModule):
-
-    """Lazy loading of moved objects"""
-    __path__ = []  # mark as package
-
-
-_moved_attributes = [
-    MovedAttribute("cStringIO", "cStringIO", "io", "StringIO"),
-    MovedAttribute("filter", "itertools", "builtins", "ifilter", "filter"),
-    MovedAttribute("filterfalse", "itertools", "itertools", "ifilterfalse", "filterfalse"),
-    MovedAttribute("input", "__builtin__", "builtins", "raw_input", "input"),
-    MovedAttribute("intern", "__builtin__", "sys"),
-    MovedAttribute("map", "itertools", "builtins", "imap", "map"),
-    MovedAttribute("getcwd", "os", "os", "getcwdu", "getcwd"),
-    MovedAttribute("getcwdb", "os", "os", "getcwd", "getcwdb"),
-    MovedAttribute("getoutput", "commands", "subprocess"),
-    MovedAttribute("range", "__builtin__", "builtins", "xrange", "range"),
-    MovedAttribute("reload_module", "__builtin__", "importlib" if PY34 else "imp", "reload"),
-    MovedAttribute("reduce", "__builtin__", "functools"),
-    MovedAttribute("shlex_quote", "pipes", "shlex", "quote"),
-    MovedAttribute("StringIO", "StringIO", "io"),
-    MovedAttribute("UserDict", "UserDict", "collections", "IterableUserDict", "UserDict"),
-    MovedAttribute("UserList", "UserList", "collections"),
-    MovedAttribute("UserString", "UserString", "collections"),
-    MovedAttribute("xrange", "__builtin__", "builtins", "xrange", "range"),
-    MovedAttribute("zip", "itertools", "builtins", "izip", "zip"),
-    MovedAttribute("zip_longest", "itertools", "itertools", "izip_longest", "zip_longest"),
-    MovedModule("builtins", "__builtin__"),
-    MovedModule("configparser", "ConfigParser"),
-    MovedModule("collections_abc", "collections", "collections.abc" if sys.version_info >= (3, 3) else "collections"),
-    MovedModule("copyreg", "copy_reg"),
-    MovedModule("dbm_gnu", "gdbm", "dbm.gnu"),
-    MovedModule("dbm_ndbm", "dbm", "dbm.ndbm"),
-    MovedModule("_dummy_thread", "dummy_thread", "_dummy_thread" if sys.version_info < (3, 9) else "_thread"),
-    MovedModule("http_cookiejar", "cookielib", "http.cookiejar"),
-    MovedModule("http_cookies", "Cookie", "http.cookies"),
-    MovedModule("html_entities", "htmlentitydefs", "html.entities"),
-    MovedModule("html_parser", "HTMLParser", "html.parser"),
-    MovedModule("http_client", "httplib", "http.client"),
-    MovedModule("email_mime_base", "email.MIMEBase", "email.mime.base"),
-    MovedModule("email_mime_image", "email.MIMEImage", "email.mime.image"),
-    MovedModule("email_mime_multipart", "email.MIMEMultipart", "email.mime.multipart"),
-    MovedModule("email_mime_nonmultipart", "email.MIMENonMultipart", "email.mime.nonmultipart"),
-    MovedModule("email_mime_text", "email.MIMEText", "email.mime.text"),
-    MovedModule("BaseHTTPServer", "BaseHTTPServer", "http.server"),
-    MovedModule("CGIHTTPServer", "CGIHTTPServer", "http.server"),
-    MovedModule("SimpleHTTPServer", "SimpleHTTPServer", "http.server"),
-    MovedModule("cPickle", "cPickle", "pickle"),
-    MovedModule("queue", "Queue"),
-    MovedModule("reprlib", "repr"),
-    MovedModule("socketserver", "SocketServer"),
-    MovedModule("_thread", "thread", "_thread"),
-    MovedModule("tkinter", "Tkinter"),
-    MovedModule("tkinter_dialog", "Dialog", "tkinter.dialog"),
-    MovedModule("tkinter_filedialog", "FileDialog", "tkinter.filedialog"),
-    MovedModule("tkinter_scrolledtext", "ScrolledText", "tkinter.scrolledtext"),
-    MovedModule("tkinter_simpledialog", "SimpleDialog", "tkinter.simpledialog"),
-    MovedModule("tkinter_tix", "Tix", "tkinter.tix"),
-    MovedModule("tkinter_ttk", "ttk", "tkinter.ttk"),
-    MovedModule("tkinter_constants", "Tkconstants", "tkinter.constants"),
-    MovedModule("tkinter_dnd", "Tkdnd", "tkinter.dnd"),
-    MovedModule("tkinter_colorchooser", "tkColorChooser",
-                "tkinter.colorchooser"),
-    MovedModule("tkinter_commondialog", "tkCommonDialog",
-                "tkinter.commondialog"),
-    MovedModule("tkinter_tkfiledialog", "tkFileDialog", "tkinter.filedialog"),
-    MovedModule("tkinter_font", "tkFont", "tkinter.font"),
-    MovedModule("tkinter_messagebox", "tkMessageBox", "tkinter.messagebox"),
-    MovedModule("tkinter_tksimpledialog", "tkSimpleDialog",
-                "tkinter.simpledialog"),
-    MovedModule("urllib_parse", __name__ + ".moves.urllib_parse", "urllib.parse"),
-    MovedModule("urllib_error", __name__ + ".moves.urllib_error", "urllib.error"),
-    MovedModule("urllib", __name__ + ".moves.urllib", __name__ + ".moves.urllib"),
-    MovedModule("urllib_robotparser", "robotparser", "urllib.robotparser"),
-    MovedModule("xmlrpc_client", "xmlrpclib", "xmlrpc.client"),
-    MovedModule("xmlrpc_server", "SimpleXMLRPCServer", "xmlrpc.server"),
-]
-# Add windows specific modules.
-if sys.platform == "win32":
-    _moved_attributes += [
-        MovedModule("winreg", "_winreg"),
-    ]
-
-for attr in _moved_attributes:
-    setattr(_MovedItems, attr.name, attr)
-    if isinstance(attr, MovedModule):
-        _importer._add_module(attr, "moves." + attr.name)
-del attr
-
-_MovedItems._moved_attributes = _moved_attributes
-
-moves = _MovedItems(__name__ + ".moves")
-_importer._add_module(moves, "moves")
-
-
-class Module_six_moves_urllib_parse(_LazyModule):
-
-    """Lazy loading of moved objects in six.moves.urllib_parse"""
-
-
-_urllib_parse_moved_attributes = [
-    MovedAttribute("ParseResult", "urlparse", "urllib.parse"),
-    MovedAttribute("SplitResult", "urlparse", "urllib.parse"),
-    MovedAttribute("parse_qs", "urlparse", "urllib.parse"),
-    MovedAttribute("parse_qsl", "urlparse", "urllib.parse"),
-    MovedAttribute("urldefrag", "urlparse", "urllib.parse"),
-    MovedAttribute("urljoin", "urlparse", "urllib.parse"),
-    MovedAttribute("urlparse", "urlparse", "urllib.parse"),
-    MovedAttribute("urlsplit", "urlparse", "urllib.parse"),
-    MovedAttribute("urlunparse", "urlparse", "urllib.parse"),
-    MovedAttribute("urlunsplit", "urlparse", "urllib.parse"),
-    MovedAttribute("quote", "urllib", "urllib.parse"),
-    MovedAttribute("quote_plus", "urllib", "urllib.parse"),
-    MovedAttribute("unquote", "urllib", "urllib.parse"),
-    MovedAttribute("unquote_plus", "urllib", "urllib.parse"),
-    MovedAttribute("unquote_to_bytes", "urllib", "urllib.parse", "unquote", "unquote_to_bytes"),
-    MovedAttribute("urlencode", "urllib", "urllib.parse"),
-    MovedAttribute("splitquery", "urllib", "urllib.parse"),
-    MovedAttribute("splittag", "urllib", "urllib.parse"),
-    MovedAttribute("splituser", "urllib", "urllib.parse"),
-    MovedAttribute("splitvalue", "urllib", "urllib.parse"),
-    MovedAttribute("uses_fragment", "urlparse", "urllib.parse"),
-    MovedAttribute("uses_netloc", "urlparse", "urllib.parse"),
-    MovedAttribute("uses_params", "urlparse", "urllib.parse"),
-    MovedAttribute("uses_query", "urlparse", "urllib.parse"),
-    MovedAttribute("uses_relative", "urlparse", "urllib.parse"),
-]
-for attr in _urllib_parse_moved_attributes:
-    setattr(Module_six_moves_urllib_parse, attr.name, attr)
-del attr
-
-Module_six_moves_urllib_parse._moved_attributes = _urllib_parse_moved_attributes
-
-_importer._add_module(Module_six_moves_urllib_parse(__name__ + ".moves.urllib_parse"),
-                      "moves.urllib_parse", "moves.urllib.parse")
-
-
-class Module_six_moves_urllib_error(_LazyModule):
-
-    """Lazy loading of moved objects in six.moves.urllib_error"""
-
-
-_urllib_error_moved_attributes = [
-    MovedAttribute("URLError", "urllib2", "urllib.error"),
-    MovedAttribute("HTTPError", "urllib2", "urllib.error"),
-    MovedAttribute("ContentTooShortError", "urllib", "urllib.error"),
-]
-for attr in _urllib_error_moved_attributes:
-    setattr(Module_six_moves_urllib_error, attr.name, attr)
-del attr
-
-Module_six_moves_urllib_error._moved_attributes = _urllib_error_moved_attributes
-
-_importer._add_module(Module_six_moves_urllib_error(__name__ + ".moves.urllib.error"),
-                      "moves.urllib_error", "moves.urllib.error")
-
-
-class Module_six_moves_urllib_request(_LazyModule):
-
-    """Lazy loading of moved objects in six.moves.urllib_request"""
-
-
-_urllib_request_moved_attributes = [
-    MovedAttribute("urlopen", "urllib2", "urllib.request"),
-    MovedAttribute("install_opener", "urllib2", "urllib.request"),
-    MovedAttribute("build_opener", "urllib2", "urllib.request"),
-    MovedAttribute("pathname2url", "urllib", "urllib.request"),
-    MovedAttribute("url2pathname", "urllib", "urllib.request"),
-    MovedAttribute("getproxies", "urllib", "urllib.request"),
-    MovedAttribute("Request", "urllib2", "urllib.request"),
-    MovedAttribute("OpenerDirector", "urllib2", "urllib.request"),
-    MovedAttribute("HTTPDefaultErrorHandler", "urllib2", "urllib.request"),
-    MovedAttribute("HTTPRedirectHandler", "urllib2", "urllib.request"),
-    MovedAttribute("HTTPCookieProcessor", "urllib2", "urllib.request"),
-    MovedAttribute("ProxyHandler", "urllib2", "urllib.request"),
-    MovedAttribute("BaseHandler", "urllib2", "urllib.request"),
-    MovedAttribute("HTTPPasswordMgr", "urllib2", "urllib.request"),
-    MovedAttribute("HTTPPasswordMgrWithDefaultRealm", "urllib2", "urllib.request"),
-    MovedAttribute("AbstractBasicAuthHandler", "urllib2", "urllib.request"),
-    MovedAttribute("HTTPBasicAuthHandler", "urllib2", "urllib.request"),
-    MovedAttribute("ProxyBasicAuthHandler", "urllib2", "urllib.request"),
-    MovedAttribute("AbstractDigestAuthHandler", "urllib2", "urllib.request"),
-    MovedAttribute("HTTPDigestAuthHandler", "urllib2", "urllib.request"),
-    MovedAttribute("ProxyDigestAuthHandler", "urllib2", "urllib.request"),
-    MovedAttribute("HTTPHandler", "urllib2", "urllib.request"),
-    MovedAttribute("HTTPSHandler", "urllib2", "urllib.request"),
-    MovedAttribute("FileHandler", "urllib2", "urllib.request"),
-    MovedAttribute("FTPHandler", "urllib2", "urllib.request"),
-    MovedAttribute("CacheFTPHandler", "urllib2", "urllib.request"),
-    MovedAttribute("UnknownHandler", "urllib2", "urllib.request"),
-    MovedAttribute("HTTPErrorProcessor", "urllib2", "urllib.request"),
-    MovedAttribute("urlretrieve", "urllib", "urllib.request"),
-    MovedAttribute("urlcleanup", "urllib", "urllib.request"),
-    MovedAttribute("proxy_bypass", "urllib", "urllib.request"),
-    MovedAttribute("parse_http_list", "urllib2", "urllib.request"),
-    MovedAttribute("parse_keqv_list", "urllib2", "urllib.request"),
-]
-if sys.version_info[:2] < (3, 14):
-    _urllib_request_moved_attributes.extend(
-        [
-            MovedAttribute("URLopener", "urllib", "urllib.request"),
-            MovedAttribute("FancyURLopener", "urllib", "urllib.request"),
-        ]
-    )
-for attr in _urllib_request_moved_attributes:
-    setattr(Module_six_moves_urllib_request, attr.name, attr)
-del attr
-
-Module_six_moves_urllib_request._moved_attributes = _urllib_request_moved_attributes
-
-_importer._add_module(Module_six_moves_urllib_request(__name__ + ".moves.urllib.request"),
-                      "moves.urllib_request", "moves.urllib.request")
-
-
-class Module_six_moves_urllib_response(_LazyModule):
-
-    """Lazy loading of moved objects in six.moves.urllib_response"""
-
-
-_urllib_response_moved_attributes = [
-    MovedAttribute("addbase", "urllib", "urllib.response"),
-    MovedAttribute("addclosehook", "urllib", "urllib.response"),
-    MovedAttribute("addinfo", "urllib", "urllib.response"),
-    MovedAttribute("addinfourl", "urllib", "urllib.response"),
-]
-for attr in _urllib_response_moved_attributes:
-    setattr(Module_six_moves_urllib_response, attr.name, attr)
-del attr
-
-Module_six_moves_urllib_response._moved_attributes = _urllib_response_moved_attributes
-
-_importer._add_module(Module_six_moves_urllib_response(__name__ + ".moves.urllib.response"),
-                      "moves.urllib_response", "moves.urllib.response")
-
-
-class Module_six_moves_urllib_robotparser(_LazyModule):
-
-    """Lazy loading of moved objects in six.moves.urllib_robotparser"""
-
-
-_urllib_robotparser_moved_attributes = [
-    MovedAttribute("RobotFileParser", "robotparser", "urllib.robotparser"),
-]
-for attr in _urllib_robotparser_moved_attributes:
-    setattr(Module_six_moves_urllib_robotparser, attr.name, attr)
-del attr
-
-Module_six_moves_urllib_robotparser._moved_attributes = _urllib_robotparser_moved_attributes
-
-_importer._add_module(Module_six_moves_urllib_robotparser(__name__ + ".moves.urllib.robotparser"),
-                      "moves.urllib_robotparser", "moves.urllib.robotparser")
-
-
-class Module_six_moves_urllib(types.ModuleType):
-
-    """Create a six.moves.urllib namespace that resembles the Python 3 namespace"""
-    __path__ = []  # mark as package
-    parse = _importer._get_module("moves.urllib_parse")
-    error = _importer._get_module("moves.urllib_error")
-    request = _importer._get_module("moves.urllib_request")
-    response = _importer._get_module("moves.urllib_response")
-    robotparser = _importer._get_module("moves.urllib_robotparser")
-
-    def __dir__(self):
-        return ['parse', 'error', 'request', 'response', 'robotparser']
-
-_importer._add_module(Module_six_moves_urllib(__name__ + ".moves.urllib"),
-                      "moves.urllib")
-
-
-def add_move(move):
-    """Add an item to six.moves."""
-    setattr(_MovedItems, move.name, move)
-
-
-def remove_move(name):
-    """Remove item from six.moves."""
-    try:
-        delattr(_MovedItems, name)
-    except AttributeError:
-        try:
-            del moves.__dict__[name]
-        except KeyError:
-            raise AttributeError("no such move, %r" % (name,))
-
-
-if PY3:
-    _meth_func = "__func__"
-    _meth_self = "__self__"
-
-    _func_closure = "__closure__"
-    _func_code = "__code__"
-    _func_defaults = "__defaults__"
-    _func_globals = "__globals__"
-else:
-    _meth_func = "im_func"
-    _meth_self = "im_self"
-
-    _func_closure = "func_closure"
-    _func_code = "func_code"
-    _func_defaults = "func_defaults"
-    _func_globals = "func_globals"
-
-
-try:
-    advance_iterator = next
-except NameError:
-    def advance_iterator(it):
-        return it.next()
-next = advance_iterator
-
-
-try:
-    callable = callable
-except NameError:
-    def callable(obj):
-        return any("__call__" in klass.__dict__ for klass in type(obj).__mro__)
-
-
-if PY3:
-    def get_unbound_function(unbound):
-        return unbound
-
-    create_bound_method = types.MethodType
-
-    def create_unbound_method(func, cls):
-        return func
-
-    Iterator = object
-else:
-    def get_unbound_function(unbound):
-        return unbound.im_func
-
-    def create_bound_method(func, obj):
-        return types.MethodType(func, obj, obj.__class__)
-
-    def create_unbound_method(func, cls):
-        return types.MethodType(func, None, cls)
-
-    class Iterator(object):
-
-        def next(self):
-            return type(self).__next__(self)
-
-    callable = callable
-_add_doc(get_unbound_function,
-         """Get the function out of a possibly unbound function""")
-
-
-get_method_function = operator.attrgetter(_meth_func)
-get_method_self = operator.attrgetter(_meth_self)
-get_function_closure = operator.attrgetter(_func_closure)
-get_function_code = operator.attrgetter(_func_code)
-get_function_defaults = operator.attrgetter(_func_defaults)
-get_function_globals = operator.attrgetter(_func_globals)
-
-
-if PY3:
-    def iterkeys(d, **kw):
-        return iter(d.keys(**kw))
-
-    def itervalues(d, **kw):
-        return iter(d.values(**kw))
-
-    def iteritems(d, **kw):
-        return iter(d.items(**kw))
-
-    def iterlists(d, **kw):
-        return iter(d.lists(**kw))
-
-    viewkeys = operator.methodcaller("keys")
-
-    viewvalues = operator.methodcaller("values")
-
-    viewitems = operator.methodcaller("items")
-else:
-    def iterkeys(d, **kw):
-        return d.iterkeys(**kw)
-
-    def itervalues(d, **kw):
-        return d.itervalues(**kw)
-
-    def iteritems(d, **kw):
-        return d.iteritems(**kw)
-
-    def iterlists(d, **kw):
-        return d.iterlists(**kw)
-
-    viewkeys = operator.methodcaller("viewkeys")
-
-    viewvalues = operator.methodcaller("viewvalues")
-
-    viewitems = operator.methodcaller("viewitems")
-
-_add_doc(iterkeys, "Return an iterator over the keys of a dictionary.")
-_add_doc(itervalues, "Return an iterator over the values of a dictionary.")
-_add_doc(iteritems,
-         "Return an iterator over the (key, value) pairs of a dictionary.")
-_add_doc(iterlists,
-         "Return an iterator over the (key, [values]) pairs of a dictionary.")
-
-
-if PY3:
-    def b(s):
-        return s.encode("latin-1")
-
-    def u(s):
-        return s
-    unichr = chr
-    import struct
-    int2byte = struct.Struct(">B").pack
-    del struct
-    byte2int = operator.itemgetter(0)
-    indexbytes = operator.getitem
-    iterbytes = iter
-    import io
-    StringIO = io.StringIO
-    BytesIO = io.BytesIO
-    del io
-    _assertCountEqual = "assertCountEqual"
-    if sys.version_info[1] <= 1:
-        _assertRaisesRegex = "assertRaisesRegexp"
-        _assertRegex = "assertRegexpMatches"
-        _assertNotRegex = "assertNotRegexpMatches"
-    else:
-        _assertRaisesRegex = "assertRaisesRegex"
-        _assertRegex = "assertRegex"
-        _assertNotRegex = "assertNotRegex"
-else:
-    def b(s):
-        return s
-    # Workaround for standalone backslash
-
-    def u(s):
-        return unicode(s.replace(r'\\', r'\\\\'), "unicode_escape")
-    unichr = unichr
-    int2byte = chr
-
-    def byte2int(bs):
-        return ord(bs[0])
-
-    def indexbytes(buf, i):
-        return ord(buf[i])
-    iterbytes = functools.partial(itertools.imap, ord)
-    import StringIO
-    StringIO = BytesIO = StringIO.StringIO
-    _assertCountEqual = "assertItemsEqual"
-    _assertRaisesRegex = "assertRaisesRegexp"
-    _assertRegex = "assertRegexpMatches"
-    _assertNotRegex = "assertNotRegexpMatches"
-_add_doc(b, """Byte literal""")
-_add_doc(u, """Text literal""")
-
-
-def assertCountEqual(self, *args, **kwargs):
-    return getattr(self, _assertCountEqual)(*args, **kwargs)
-
-
-def assertRaisesRegex(self, *args, **kwargs):
-    return getattr(self, _assertRaisesRegex)(*args, **kwargs)
-
-
-def assertRegex(self, *args, **kwargs):
-    return getattr(self, _assertRegex)(*args, **kwargs)
-
-
-def assertNotRegex(self, *args, **kwargs):
-    return getattr(self, _assertNotRegex)(*args, **kwargs)
-
-
-if PY3:
-    exec_ = getattr(moves.builtins, "exec")
-
-    def reraise(tp, value, tb=None):
-        try:
-            if value is None:
-                value = tp()
-            if value.__traceback__ is not tb:
-                raise value.with_traceback(tb)
-            raise value
-        finally:
-            value = None
-            tb = None
-
-else:
-    def exec_(_code_, _globs_=None, _locs_=None):
-        """Execute code in a namespace."""
-        if _globs_ is None:
-            frame = sys._getframe(1)
-            _globs_ = frame.f_globals
-            if _locs_ is None:
-                _locs_ = frame.f_locals
-            del frame
-        elif _locs_ is None:
-            _locs_ = _globs_
-        exec("""exec _code_ in _globs_, _locs_""")
-
-    exec_("""def reraise(tp, value, tb=None):
-    try:
-        raise tp, value, tb
-    finally:
-        tb = None
-""")
-
-
-if sys.version_info[:2] > (3,):
-    exec_("""def raise_from(value, from_value):
-    try:
-        raise value from from_value
-    finally:
-        value = None
-""")
-else:
-    def raise_from(value, from_value):
-        raise value
-
-
-print_ = getattr(moves.builtins, "print", None)
-if print_ is None:
-    def print_(*args, **kwargs):
-        """The new-style print function for Python 2.4 and 2.5."""
-        fp = kwargs.pop("file", sys.stdout)
-        if fp is None:
-            return
-
-        def write(data):
-            if not isinstance(data, basestring):
-                data = str(data)
-            # If the file has an encoding, encode unicode with it.
-            if (isinstance(fp, file) and
-                    isinstance(data, unicode) and
-                    fp.encoding is not None):
-                errors = getattr(fp, "errors", None)
-                if errors is None:
-                    errors = "strict"
-                data = data.encode(fp.encoding, errors)
-            fp.write(data)
-        want_unicode = False
-        sep = kwargs.pop("sep", None)
-        if sep is not None:
-            if isinstance(sep, unicode):
-                want_unicode = True
-            elif not isinstance(sep, str):
-                raise TypeError("sep must be None or a string")
-        end = kwargs.pop("end", None)
-        if end is not None:
-            if isinstance(end, unicode):
-                want_unicode = True
-            elif not isinstance(end, str):
-                raise TypeError("end must be None or a string")
-        if kwargs:
-            raise TypeError("invalid keyword arguments to print()")
-        if not want_unicode:
-            for arg in args:
-                if isinstance(arg, unicode):
-                    want_unicode = True
-                    break
-        if want_unicode:
-            newline = unicode("\n")
-            space = unicode(" ")
-        else:
-            newline = "\n"
-            space = " "
-        if sep is None:
-            sep = space
-        if end is None:
-            end = newline
-        for i, arg in enumerate(args):
-            if i:
-                write(sep)
-            write(arg)
-        write(end)
-if sys.version_info[:2] < (3, 3):
-    _print = print_
-
-    def print_(*args, **kwargs):
-        fp = kwargs.get("file", sys.stdout)
-        flush = kwargs.pop("flush", False)
-        _print(*args, **kwargs)
-        if flush and fp is not None:
-            fp.flush()
-
-_add_doc(reraise, """Reraise an exception.""")
-
-if sys.version_info[0:2] < (3, 4):
-    # This does exactly the same what the :func:`py3:functools.update_wrapper`
-    # function does on Python versions after 3.2. It sets the ``__wrapped__``
-    # attribute on ``wrapper`` object and it doesn't raise an error if any of
-    # the attributes mentioned in ``assigned`` and ``updated`` are missing on
-    # ``wrapped`` object.
-    def _update_wrapper(wrapper, wrapped,
-                        assigned=functools.WRAPPER_ASSIGNMENTS,
-                        updated=functools.WRAPPER_UPDATES):
-        for attr in assigned:
-            try:
-                value = getattr(wrapped, attr)
-            except AttributeError:
-                continue
-            else:
-                setattr(wrapper, attr, value)
-        for attr in updated:
-            getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
-        wrapper.__wrapped__ = wrapped
-        return wrapper
-    _update_wrapper.__doc__ = functools.update_wrapper.__doc__
-
-    def wraps(wrapped, assigned=functools.WRAPPER_ASSIGNMENTS,
-              updated=functools.WRAPPER_UPDATES):
-        return functools.partial(_update_wrapper, wrapped=wrapped,
-                                 assigned=assigned, updated=updated)
-    wraps.__doc__ = functools.wraps.__doc__
-
-else:
-    wraps = functools.wraps
-
-
-def with_metaclass(meta, *bases):
-    """Create a base class with a metaclass."""
-    # This requires a bit of explanation: the basic idea is to make a dummy
-    # metaclass for one level of class instantiation that replaces itself with
-    # the actual metaclass.
-    class metaclass(type):
-
-        def __new__(cls, name, this_bases, d):
-            if sys.version_info[:2] >= (3, 7):
-                # This version introduced PEP 560 that requires a bit
-                # of extra care (we mimic what is done by __build_class__).
-                resolved_bases = types.resolve_bases(bases)
-                if resolved_bases is not bases:
-                    d['__orig_bases__'] = bases
-            else:
-                resolved_bases = bases
-            return meta(name, resolved_bases, d)
-
-        @classmethod
-        def __prepare__(cls, name, this_bases):
-            return meta.__prepare__(name, bases)
-    return type.__new__(metaclass, 'temporary_class', (), {})
-
-
-def add_metaclass(metaclass):
-    """Class decorator for creating a class with a metaclass."""
-    def wrapper(cls):
-        orig_vars = cls.__dict__.copy()
-        slots = orig_vars.get('__slots__')
-        if slots is not None:
-            if isinstance(slots, str):
-                slots = [slots]
-            for slots_var in slots:
-                orig_vars.pop(slots_var)
-        orig_vars.pop('__dict__', None)
-        orig_vars.pop('__weakref__', None)
-        if hasattr(cls, '__qualname__'):
-            orig_vars['__qualname__'] = cls.__qualname__
-        return metaclass(cls.__name__, cls.__bases__, orig_vars)
-    return wrapper
-
-
-def ensure_binary(s, encoding='utf-8', errors='strict'):
-    """Coerce **s** to six.binary_type.
-
-    For Python 2:
-      - `unicode` -> encoded to `str`
-      - `str` -> `str`
-
-    For Python 3:
-      - `str` -> encoded to `bytes`
-      - `bytes` -> `bytes`
-    """
-    if isinstance(s, binary_type):
-        return s
-    if isinstance(s, text_type):
-        return s.encode(encoding, errors)
-    raise TypeError("not expecting type '%s'" % type(s))
-
-
-def ensure_str(s, encoding='utf-8', errors='strict'):
-    """Coerce *s* to `str`.
-
-    For Python 2:
-      - `unicode` -> encoded to `str`
-      - `str` -> `str`
-
-    For Python 3:
-      - `str` -> `str`
-      - `bytes` -> decoded to `str`
-    """
-    # Optimization: Fast return for the common case.
-    if type(s) is str:
-        return s
-    if PY2 and isinstance(s, text_type):
-        return s.encode(encoding, errors)
-    elif PY3 and isinstance(s, binary_type):
-        return s.decode(encoding, errors)
-    elif not isinstance(s, (text_type, binary_type)):
-        raise TypeError("not expecting type '%s'" % type(s))
-    return s
-
-
-def ensure_text(s, encoding='utf-8', errors='strict'):
-    """Coerce *s* to six.text_type.
-
-    For Python 2:
-      - `unicode` -> `unicode`
-      - `str` -> `unicode`
-
-    For Python 3:
-      - `str` -> `str`
-      - `bytes` -> decoded to `str`
-    """
-    if isinstance(s, binary_type):
-        return s.decode(encoding, errors)
-    elif isinstance(s, text_type):
-        return s
-    else:
-        raise TypeError("not expecting type '%s'" % type(s))
-
-
-def python_2_unicode_compatible(klass):
-    """
-    A class decorator that defines __unicode__ and __str__ methods under Python 2.
-    Under Python 3 it does nothing.
-
-    To support Python 2 and 3 with a single code base, define a __str__ method
-    returning text and apply this decorator to the class.
-    """
-    if PY2:
-        if '__str__' not in klass.__dict__:
-            raise ValueError("@python_2_unicode_compatible cannot be applied "
-                             "to %s because it doesn't define __str__()." %
-                             klass.__name__)
-        klass.__unicode__ = klass.__str__
-        klass.__str__ = lambda self: self.__unicode__().encode('utf-8')
-    return klass
-
-
-# Complete the moves implementation.
-# This code is at the end of this module to speed up module loading.
-# Turn this module into a package.
-__path__ = []  # required for PEP 302 and PEP 451
-__package__ = __name__  # see PEP 366 @ReservedAssignment
-if globals().get("__spec__") is not None:
-    __spec__.submodule_search_locations = []  # PEP 451 @UndefinedVariable
-# Remove other six meta path importers, since they cause problems. This can
-# happen if six is removed from sys.modules and then reloaded. (Setuptools does
-# this for some reason.)
-if sys.meta_path:
-    for i, importer in enumerate(sys.meta_path):
-        # Here's some real nastiness: Another "instance" of the six module might
-        # be floating around. Therefore, we can't use isinstance() to check for
-        # the six meta path importer, since the other six instance will have
-        # inserted an importer with different class.
-        if (type(importer).__name__ == "_SixMetaPathImporter" and
-                importer.name == __name__):
-            del sys.meta_path[i]
-            break
-    del i, importer
-# Finally, add the importer to the meta path import hook.
-sys.meta_path.append(_importer)
diff --git a/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/INSTALLER
deleted file mode 100644
index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/INSTALLER
+++ /dev/null
@@ -1 +0,0 @@
-pip
diff --git a/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/LICENSE b/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/LICENSE
deleted file mode 100644
index 710010f1fe0311afd9823e319ff958827a419964..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/LICENSE
+++ /dev/null
@@ -1,30 +0,0 @@
-Copyright (C) 2010, 2011 Sebastian Thiel and contributors
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without 
-modification, are permitted provided that the following conditions 
-are met:
-
-* Redistributions of source code must retain the above copyright 
-notice, this list of conditions and the following disclaimer.
-
-* Redistributions in binary form must reproduce the above copyright 
-notice, this list of conditions and the following disclaimer in the 
-documentation and/or other materials provided with the distribution.
-
-* Neither the name of the async project nor the names of 
-its contributors may be used to endorse or promote products derived 
-from this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 
-"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 
-LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 
-A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 
-OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 
-SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED 
-TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 
-PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF 
-LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING 
-NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 
-SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
diff --git a/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/METADATA b/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/METADATA
deleted file mode 100644
index ce0537bde0f6d4cfee9549b001156bd5cabc107a..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/METADATA
+++ /dev/null
@@ -1,113 +0,0 @@
-Metadata-Version: 2.1
-Name: smmap
-Version: 5.0.2
-Summary: A pure Python implementation of a sliding window memory map manager
-Home-page: https://github.com/gitpython-developers/smmap
-Author: Sebastian Thiel
-Author-email: byronimo@gmail.com
-License: BSD-3-Clause
-Platform: any
-Classifier: Development Status :: 5 - Production/Stable
-Classifier: Environment :: Console
-Classifier: Intended Audience :: Developers
-Classifier: License :: OSI Approved :: BSD License
-Classifier: Operating System :: OS Independent
-Classifier: Operating System :: POSIX
-Classifier: Operating System :: Microsoft :: Windows
-Classifier: Operating System :: MacOS :: MacOS X
-Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Classifier: Programming Language :: Python :: 3 :: Only
-Requires-Python: >=3.7
-Description-Content-Type: text/markdown
-License-File: LICENSE
-
-## Motivation
-
-When reading from many possibly large files in a fashion similar to random access, it is usually the fastest and most efficient to use memory maps.
-
-Although memory maps have many advantages, they represent a very limited system resource as every map uses one file descriptor, whose amount is limited per process. On 32 bit systems, the amount of memory you can have mapped at a time is naturally limited to theoretical 4GB of memory, which may not be enough for some applications.
-
-
-## Limitations
-
-* **System resources (file-handles) are likely to be leaked!** This is due to the library authors reliance on a deterministic `__del__()` destructor.
-* The memory access is read-only by design.
-
-
-## Overview
-
-![Python package](https://github.com/gitpython-developers/smmap/workflows/Python%20package/badge.svg)
-
-Smmap wraps an interface around mmap and tracks the mapped files as well as the amount of clients who use it. If the system runs out of resources, or if a memory limit is reached, it will automatically unload unused maps to allow continued operation.
-
-To allow processing large files even on 32 bit systems, it allows only portions of the file to be mapped. Once the user reads beyond the mapped region, smmap will automatically map the next required region, unloading unused regions using a LRU algorithm.
-
-Although the library can be used most efficiently with its native interface, a Buffer implementation is provided to hide these details behind a simple string-like interface.
-
-For performance critical 64 bit applications, a simplified version of memory mapping is provided which always maps the whole file, but still provides the benefit of unloading unused mappings on demand.
-
-
-
-## Prerequisites
-
-* Python 3.7+
-* OSX, Windows or Linux
-
-The package was tested on all of the previously mentioned configurations.
-
-## Installing smmap
-
-[![Documentation Status](https://readthedocs.org/projects/smmap/badge/?version=latest)](https://readthedocs.org/projects/smmap/?badge=latest)
-
-Its easiest to install smmap using the [pip](http://www.pip-installer.org/en/latest) program:
-
-```bash
-$ pip install smmap
-```
-
-As the command will install smmap in your respective python distribution, you will most likely need root permissions to authorize the required changes.
-
-If you have downloaded the source archive, the package can be installed by running the `setup.py` script:
-
-```bash
-$ python setup.py install
-```
-
-It is advised to have a look at the **Usage Guide** for a brief introduction on the different database implementations.
-
-
-
-## Homepage and Links
-
-The project is home on github at https://github.com/gitpython-developers/smmap .
-
-The latest source can be cloned from github as well:
-
-* git://github.com/gitpython-developers/smmap.git
-
-
-For support, please use the git-python mailing list:
-
-* http://groups.google.com/group/git-python
-
-
-Issues can be filed on github:
-
-* https://github.com/gitpython-developers/smmap/issues
-
-A link to the pypi page related to this repository:
-
-* https://pypi.org/project/smmap/
-
-
-## License Information
-
-*smmap* is licensed under the New BSD License.
-
diff --git a/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/RECORD b/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/RECORD
deleted file mode 100644
index 5e5685f8fff0c6cfa171ef4f5e11b8e269bae405..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/RECORD
+++ /dev/null
@@ -1,27 +0,0 @@
-smmap-5.0.2.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
-smmap-5.0.2.dist-info/LICENSE,sha256=iOnZP3CNEQsyioNDAt0dXGr72lMOdyHRXYCzUR2G8jU,1519
-smmap-5.0.2.dist-info/METADATA,sha256=3BUpM_oN3bljLC5WDH9bxeo5JqXV-AGVQ4Epw1C5h2Y,4306
-smmap-5.0.2.dist-info/RECORD,,
-smmap-5.0.2.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-smmap-5.0.2.dist-info/top_level.txt,sha256=h0Tp1UdaROCy2bjWNha1MFpwgVghxEUQsaLHbZs96Gw,6
-smmap-5.0.2.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-smmap/__init__.py,sha256=Bf9KWm-zr5__SoSEEXI20pqKOJFpUApidvardLulRLY,343
-smmap/__pycache__/__init__.cpython-310.pyc,,
-smmap/__pycache__/buf.cpython-310.pyc,,
-smmap/__pycache__/mman.cpython-310.pyc,,
-smmap/__pycache__/util.cpython-310.pyc,,
-smmap/buf.py,sha256=wElndUNa4Njhnes6JYmI_Xp4zwaohirLdr7eVR9XYNU,5762
-smmap/mman.py,sha256=HYxohllZqpPY4KgRF_D_lNp4fx0akZ-NQosutLDrQ60,24022
-smmap/test/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-smmap/test/__pycache__/__init__.cpython-310.pyc,,
-smmap/test/__pycache__/lib.cpython-310.pyc,,
-smmap/test/__pycache__/test_buf.cpython-310.pyc,,
-smmap/test/__pycache__/test_mman.cpython-310.pyc,,
-smmap/test/__pycache__/test_tutorial.cpython-310.pyc,,
-smmap/test/__pycache__/test_util.cpython-310.pyc,,
-smmap/test/lib.py,sha256=jYIT4rKckHw2s7-l5lwysft1jqrwDFkjtyNLKNpu_j0,1489
-smmap/test/test_buf.py,sha256=BCg9FbusakiYNs5MUNoJSiyuTP5HcMYl65HzvpY9SNo,5439
-smmap/test/test_mman.py,sha256=o3cPg5idyYJpq32zXgDy1qtD10RA_e88-1zPYBWi4Ms,10822
-smmap/test/test_tutorial.py,sha256=ZwCphCbKAGYt_fn_CiOJ9lWwpWwpqE4-zBUp-v-t9eM,3174
-smmap/test/test_util.py,sha256=3hJyW9Km7k7XSgxxtDOkG8eVagk3lIzP4H2pR0S2ewg,3468
-smmap/util.py,sha256=_Sl3Sd2KWdd-5nmhUegFMfT6L1j5BGchRT4z53Sobug,7476
diff --git a/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/WHEEL b/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/WHEEL
deleted file mode 100644
index ae527e7d64811439e61b93aa375defb30e06edfe..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/WHEEL
+++ /dev/null
@@ -1,5 +0,0 @@
-Wheel-Version: 1.0
-Generator: setuptools (75.6.0)
-Root-Is-Purelib: true
-Tag: py3-none-any
-
diff --git a/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/top_level.txt
deleted file mode 100644
index 2765cefa1c4976ec31a903f3ffb7b2f9f92531f3..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/top_level.txt
+++ /dev/null
@@ -1 +0,0 @@
-smmap
diff --git a/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/zip-safe b/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/zip-safe
deleted file mode 100644
index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap-5.0.2.dist-info/zip-safe
+++ /dev/null
@@ -1 +0,0 @@
-
diff --git a/venv/lib/python3.10/site-packages/smmap/__init__.py b/venv/lib/python3.10/site-packages/smmap/__init__.py
deleted file mode 100644
index 657c8c5af20f538203ef618e0b38c12bf6c998bc..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap/__init__.py
+++ /dev/null
@@ -1,11 +0,0 @@
-"""Initialize the smmap package"""
-
-__author__ = "Sebastian Thiel"
-__contact__ = "byronimo@gmail.com"
-__homepage__ = "https://github.com/gitpython-developers/smmap"
-version_info = (5, 0, 2)
-__version__ = '.'.join(str(i) for i in version_info)
-
-# make everything available in root package for convenience
-from .mman import *
-from .buf import *
diff --git a/venv/lib/python3.10/site-packages/smmap/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/smmap/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index e8288829fbbd58ee292026974e3fc6bbc91f295b..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/smmap/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/smmap/__pycache__/buf.cpython-310.pyc b/venv/lib/python3.10/site-packages/smmap/__pycache__/buf.cpython-310.pyc
deleted file mode 100644
index e8471eb5c1287b19f50e05230f9058839861e15f..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/smmap/__pycache__/buf.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/smmap/__pycache__/mman.cpython-310.pyc b/venv/lib/python3.10/site-packages/smmap/__pycache__/mman.cpython-310.pyc
deleted file mode 100644
index 151538085c6e46a8d439eff4bff1f41d9d88edc8..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/smmap/__pycache__/mman.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/smmap/__pycache__/util.cpython-310.pyc b/venv/lib/python3.10/site-packages/smmap/__pycache__/util.cpython-310.pyc
deleted file mode 100644
index 907fc1594088dcb417382db016f31b42fff2e910..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/smmap/__pycache__/util.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/smmap/buf.py b/venv/lib/python3.10/site-packages/smmap/buf.py
deleted file mode 100644
index 731b0644b5a269b1f941dc7cd49163e6754878a6..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap/buf.py
+++ /dev/null
@@ -1,143 +0,0 @@
-"""Module with a simple buffer implementation using the memory manager"""
-import sys
-
-__all__ = ["SlidingWindowMapBuffer"]
-
-
-class SlidingWindowMapBuffer:
-
-    """A buffer like object which allows direct byte-wise object and slicing into
-    memory of a mapped file. The mapping is controlled by the provided cursor.
-
-    The buffer is relative, that is if you map an offset, index 0 will map to the
-    first byte at the offset you used during initialization or begin_access
-
-    **Note:** Although this type effectively hides the fact that there are mapped windows
-    underneath, it can unfortunately not be used in any non-pure python method which
-    needs a buffer or string"""
-    __slots__ = (
-        '_c',           # our cursor
-        '_size',        # our supposed size
-    )
-
-    def __init__(self, cursor=None, offset=0, size=sys.maxsize, flags=0):
-        """Initialize the instance to operate on the given cursor.
-        :param cursor: if not None, the associated cursor to the file you want to access
-            If None, you have call begin_access before using the buffer and provide a cursor
-        :param offset: absolute offset in bytes
-        :param size: the total size of the mapping. Defaults to the maximum possible size
-            From that point on, the __len__ of the buffer will be the given size or the file size.
-            If the size is larger than the mappable area, you can only access the actually available
-            area, although the length of the buffer is reported to be your given size.
-            Hence it is in your own interest to provide a proper size !
-        :param flags: Additional flags to be passed to os.open
-        :raise ValueError: if the buffer could not achieve a valid state"""
-        self._c = cursor
-        if cursor and not self.begin_access(cursor, offset, size, flags):
-            raise ValueError("Failed to allocate the buffer - probably the given offset is out of bounds")
-        # END handle offset
-
-    def __del__(self):
-        self.end_access()
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        self.end_access()
-
-    def __len__(self):
-        return self._size
-
-    def __getitem__(self, i):
-        if isinstance(i, slice):
-            return self.__getslice__(i.start or 0, i.stop or self._size)
-        c = self._c
-        assert c.is_valid()
-        if i < 0:
-            i = self._size + i
-        if not c.includes_ofs(i):
-            c.use_region(i, 1)
-        # END handle region usage
-        return c.buffer()[i - c.ofs_begin()]
-
-    def __getslice__(self, i, j):
-        c = self._c
-        # fast path, slice fully included - safes a concatenate operation and
-        # should be the default
-        assert c.is_valid()
-        if i < 0:
-            i = self._size + i
-        if j == sys.maxsize:
-            j = self._size
-        if j < 0:
-            j = self._size + j
-        if (c.ofs_begin() <= i) and (j < c.ofs_end()):
-            b = c.ofs_begin()
-            return c.buffer()[i - b:j - b]
-        else:
-            l = j - i                 # total length
-            ofs = i
-            # It's fastest to keep tokens and join later, especially in py3, which was 7 times slower
-            # in the previous iteration of this code
-            md = list()
-            while l:
-                c.use_region(ofs, l)
-                assert c.is_valid()
-                d = c.buffer()[:l]
-                ofs += len(d)
-                l -= len(d)
-                # Make sure we don't keep references, as c.use_region() might attempt to free resources, but
-                # can't unless we use pure bytes
-                if hasattr(d, 'tobytes'):
-                    d = d.tobytes()
-                md.append(d)
-            # END while there are bytes to read
-            return b''.join(md)
-        # END fast or slow path
-    #{ Interface
-
-    def begin_access(self, cursor=None, offset=0, size=sys.maxsize, flags=0):
-        """Call this before the first use of this instance. The method was already
-        called by the constructor in case sufficient information was provided.
-
-        For more information no the parameters, see the __init__ method
-        :param path: if cursor is None the existing one will be used.
-        :return: True if the buffer can be used"""
-        if cursor:
-            self._c = cursor
-        # END update our cursor
-
-        # reuse existing cursors if possible
-        if self._c is not None and self._c.is_associated():
-            res = self._c.use_region(offset, size, flags).is_valid()
-            if res:
-                # if given size is too large or default, we computer a proper size
-                # If its smaller, we assume the combination between offset and size
-                # as chosen by the user is correct and use it !
-                # If not, the user is in trouble.
-                if size > self._c.file_size():
-                    size = self._c.file_size() - offset
-                # END handle size
-                self._size = size
-            # END set size
-            return res
-        # END use our cursor
-        return False
-
-    def end_access(self):
-        """Call this method once you are done using the instance. It is automatically
-        called on destruction, and should be called just in time to allow system
-        resources to be freed.
-
-        Once you called end_access, you must call begin access before reusing this instance!"""
-        self._size = 0
-        if self._c is not None:
-            self._c.unuse_region()
-        # END unuse region
-
-    def cursor(self):
-        """:return: the currently set cursor which provides access to the data"""
-        return self._c
-
-    #}END interface
diff --git a/venv/lib/python3.10/site-packages/smmap/mman.py b/venv/lib/python3.10/site-packages/smmap/mman.py
deleted file mode 100644
index 873f687fdc6f0d78ee24a6d25f7720708d48b6e5..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap/mman.py
+++ /dev/null
@@ -1,588 +0,0 @@
-"""Module containing a memory memory manager which provides a sliding window on a number of memory mapped files"""
-from .util import (
-    MapWindow,
-    MapRegion,
-    MapRegionList,
-    is_64_bit,
-)
-
-import sys
-from functools import reduce
-
-__all__ = ["StaticWindowMapManager", "SlidingWindowMapManager", "WindowCursor"]
-#{ Utilities
-
-#}END utilities
-
-
-class WindowCursor:
-
-    """
-    Pointer into the mapped region of the memory manager, keeping the map
-    alive until it is destroyed and no other client uses it.
-
-    Cursors should not be created manually, but are instead returned by the SlidingWindowMapManager
-
-    **Note:**: The current implementation is suited for static and sliding window managers, but it also means
-    that it must be suited for the somewhat quite different sliding manager. It could be improved, but
-    I see no real need to do so."""
-    __slots__ = (
-        '_manager',  # the manager keeping all file regions
-        '_rlist',   # a regions list with regions for our file
-        '_region',  # our current class:`MapRegion` or None
-        '_ofs',     # relative offset from the actually mapped area to our start area
-        '_size'     # maximum size we should provide
-    )
-
-    def __init__(self, manager=None, regions=None):
-        self._manager = manager
-        self._rlist = regions
-        self._region = None
-        self._ofs = 0
-        self._size = 0
-
-    def __del__(self):
-        self._destroy()
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        self._destroy()
-
-    def _destroy(self):
-        """Destruction code to decrement counters"""
-        self.unuse_region()
-
-        if self._rlist is not None:
-            # Actual client count, which doesn't include the reference kept by the manager, nor ours
-            # as we are about to be deleted
-            try:
-                if len(self._rlist) == 0:
-                    # Free all resources associated with the mapped file
-                    self._manager._fdict.pop(self._rlist.path_or_fd())
-                # END remove regions list from manager
-            except (TypeError, KeyError):
-                # sometimes, during shutdown, getrefcount is None. Its possible
-                # to re-import it, however, its probably better to just ignore
-                # this python problem (for now).
-                # The next step is to get rid of the error prone getrefcount altogether.
-                pass
-            # END exception handling
-        # END handle regions
-
-    def _copy_from(self, rhs):
-        """Copy all data from rhs into this instance, handles usage count"""
-        self._manager = rhs._manager
-        self._rlist = type(rhs._rlist)(rhs._rlist)
-        self._region = rhs._region
-        self._ofs = rhs._ofs
-        self._size = rhs._size
-
-        for region in self._rlist:
-            region.increment_client_count()
-
-        if self._region is not None:
-            self._region.increment_client_count()
-        # END handle regions
-
-    def __copy__(self):
-        """copy module interface"""
-        cpy = type(self)()
-        cpy._copy_from(self)
-        return cpy
-
-    #{ Interface
-    def assign(self, rhs):
-        """Assign rhs to this instance. This is required in order to get a real copy.
-        Alternatively, you can copy an existing instance using the copy module"""
-        self._destroy()
-        self._copy_from(rhs)
-
-    def use_region(self, offset=0, size=0, flags=0):
-        """Assure we point to a window which allows access to the given offset into the file
-
-        :param offset: absolute offset in bytes into the file
-        :param size: amount of bytes to map. If 0, all available bytes will be mapped
-        :param flags: additional flags to be given to os.open in case a file handle is initially opened
-            for mapping. Has no effect if a region can actually be reused.
-        :return: this instance - it should be queried for whether it points to a valid memory region.
-            This is not the case if the mapping failed because we reached the end of the file
-
-        **Note:**: The size actually mapped may be smaller than the given size. If that is the case,
-        either the file has reached its end, or the map was created between two existing regions"""
-        need_region = True
-        man = self._manager
-        fsize = self._rlist.file_size()
-        size = min(size or fsize, man.window_size() or fsize)   # clamp size to window size
-
-        if self._region is not None:
-            if self._region.includes_ofs(offset):
-                need_region = False
-            else:
-                self.unuse_region()
-            # END handle existing region
-        # END check existing region
-
-        # offset too large ?
-        if offset >= fsize:
-            return self
-        # END handle offset
-
-        if need_region:
-            self._region = man._obtain_region(self._rlist, offset, size, flags, False)
-            self._region.increment_client_count()
-        # END need region handling
-
-        self._ofs = offset - self._region._b
-        self._size = min(size, self._region.ofs_end() - offset)
-
-        return self
-
-    def unuse_region(self):
-        """Unuse the current region. Does nothing if we have no current region
-
-        **Note:** the cursor unuses the region automatically upon destruction. It is recommended
-        to un-use the region once you are done reading from it in persistent cursors as it
-        helps to free up resource more quickly"""
-        if self._region is not None:
-            self._region.increment_client_count(-1)
-        self._region = None
-        # note: should reset ofs and size, but we spare that for performance. Its not
-        # allowed to query information if we are not valid !
-
-    def buffer(self):
-        """Return a buffer object which allows access to our memory region from our offset
-        to the window size. Please note that it might be smaller than you requested when calling use_region()
-
-        **Note:** You can only obtain a buffer if this instance is_valid() !
-
-        **Note:** buffers should not be cached passed the duration of your access as it will
-        prevent resources from being freed even though they might not be accounted for anymore !"""
-        return memoryview(self._region.buffer())[self._ofs:self._ofs+self._size]
-
-    def map(self):
-        """
-        :return: the underlying raw memory map. Please not that the offset and size is likely to be different
-            to what you set as offset and size. Use it only if you are sure about the region it maps, which is the whole
-            file in case of StaticWindowMapManager"""
-        return self._region.map()
-
-    def is_valid(self):
-        """:return: True if we have a valid and usable region"""
-        return self._region is not None
-
-    def is_associated(self):
-        """:return: True if we are associated with a specific file already"""
-        return self._rlist is not None
-
-    def ofs_begin(self):
-        """:return: offset to the first byte pointed to by our cursor
-
-        **Note:** only if is_valid() is True"""
-        return self._region._b + self._ofs
-
-    def ofs_end(self):
-        """:return: offset to one past the last available byte"""
-        # unroll method calls for performance !
-        return self._region._b + self._ofs + self._size
-
-    def size(self):
-        """:return: amount of bytes we point to"""
-        return self._size
-
-    def region(self):
-        """:return: our mapped region, or None if nothing is mapped yet
-        :raise AssertionError: if we have no current region. This is only useful for debugging"""
-        return self._region
-
-    def includes_ofs(self, ofs):
-        """:return: True if the given absolute offset is contained in the cursors
-            current region
-
-        **Note:** cursor must be valid for this to work"""
-        # unroll methods
-        return (self._region._b + self._ofs) <= ofs < (self._region._b + self._ofs + self._size)
-
-    def file_size(self):
-        """:return: size of the underlying file"""
-        return self._rlist.file_size()
-
-    def path_or_fd(self):
-        """:return: path or file descriptor of the underlying mapped file"""
-        return self._rlist.path_or_fd()
-
-    def path(self):
-        """:return: path of the underlying mapped file
-        :raise ValueError: if attached path is not a path"""
-        if isinstance(self._rlist.path_or_fd(), int):
-            raise ValueError("Path queried although mapping was applied to a file descriptor")
-        # END handle type
-        return self._rlist.path_or_fd()
-
-    def fd(self):
-        """:return: file descriptor used to create the underlying mapping.
-
-        **Note:** it is not required to be valid anymore
-        :raise ValueError: if the mapping was not created by a file descriptor"""
-        if isinstance(self._rlist.path_or_fd(), str):
-            raise ValueError("File descriptor queried although mapping was generated from path")
-        # END handle type
-        return self._rlist.path_or_fd()
-
-    #} END interface
-
-
-class StaticWindowMapManager:
-
-    """Provides a manager which will produce single size cursors that are allowed
-    to always map the whole file.
-
-    Clients must be written to specifically know that they are accessing their data
-    through a StaticWindowMapManager, as they otherwise have to deal with their window size.
-
-    These clients would have to use a SlidingWindowMapBuffer to hide this fact.
-
-    This type will always use a maximum window size, and optimize certain methods to
-    accommodate this fact"""
-
-    __slots__ = [
-        '_fdict',           # mapping of path -> StorageHelper (of some kind
-        '_window_size',     # maximum size of a window
-        '_max_memory_size',  # maximum amount of memory we may allocate
-        '_max_handle_count',        # maximum amount of handles to keep open
-        '_memory_size',     # currently allocated memory size
-        '_handle_count',        # amount of currently allocated file handles
-    ]
-
-    #{ Configuration
-    MapRegionListCls = MapRegionList
-    MapWindowCls = MapWindow
-    MapRegionCls = MapRegion
-    WindowCursorCls = WindowCursor
-    #} END configuration
-
-    _MB_in_bytes = 1024 * 1024
-
-    def __init__(self, window_size=0, max_memory_size=0, max_open_handles=sys.maxsize):
-        """initialize the manager with the given parameters.
-        :param window_size: if -1, a default window size will be chosen depending on
-            the operating system's architecture. It will internally be quantified to a multiple of the page size
-            If 0, the window may have any size, which basically results in mapping the whole file at one
-        :param max_memory_size: maximum amount of memory we may map at once before releasing mapped regions.
-            If 0, a viable default will be set depending on the system's architecture.
-            It is a soft limit that is tried to be kept, but nothing bad happens if we have to over-allocate
-        :param max_open_handles: if not maxint, limit the amount of open file handles to the given number.
-            Otherwise the amount is only limited by the system itself. If a system or soft limit is hit,
-            the manager will free as many handles as possible"""
-        self._fdict = dict()
-        self._window_size = window_size
-        self._max_memory_size = max_memory_size
-        self._max_handle_count = max_open_handles
-        self._memory_size = 0
-        self._handle_count = 0
-
-        if window_size < 0:
-            coeff = 64
-            if is_64_bit():
-                coeff = 1024
-            # END handle arch
-            self._window_size = coeff * self._MB_in_bytes
-        # END handle max window size
-
-        if max_memory_size == 0:
-            coeff = 1024
-            if is_64_bit():
-                coeff = 8192
-            # END handle arch
-            self._max_memory_size = coeff * self._MB_in_bytes
-        # END handle max memory size
-
-    #{ Internal Methods
-
-    def _collect_lru_region(self, size):
-        """Unmap the region which was least-recently used and has no client
-        :param size: size of the region we want to map next (assuming its not already mapped partially or full
-            if 0, we try to free any available region
-        :return: Amount of freed regions
-
-        .. Note::
-            We don't raise exceptions anymore, in order to keep the system working, allowing temporary overallocation.
-            If the system runs out of memory, it will tell.
-
-        .. TODO::
-            implement a case where all unusued regions are discarded efficiently.
-            Currently its only brute force
-        """
-        num_found = 0
-        while (size == 0) or (self._memory_size + size > self._max_memory_size):
-            lru_region = None
-            lru_list = None
-            for regions in self._fdict.values():
-                for region in regions:
-                    # check client count - if it's 1, it's just us
-                    if (region.client_count() == 1 and
-                            (lru_region is None or region._uc < lru_region._uc)):
-                        lru_region = region
-                        lru_list = regions
-                    # END update lru_region
-                # END for each region
-            # END for each regions list
-
-            if lru_region is None:
-                break
-            # END handle region not found
-
-            num_found += 1
-            del(lru_list[lru_list.index(lru_region)])
-            lru_region.increment_client_count(-1)
-            self._memory_size -= lru_region.size()
-            self._handle_count -= 1
-        # END while there is more memory to free
-        return num_found
-
-    def _obtain_region(self, a, offset, size, flags, is_recursive):
-        """Utility to create a new region - for more information on the parameters,
-        see MapCursor.use_region.
-        :param a: A regions (a)rray
-        :return: The newly created region"""
-        if self._memory_size + size > self._max_memory_size:
-            self._collect_lru_region(size)
-        # END handle collection
-
-        r = None
-        if a:
-            assert len(a) == 1
-            r = a[0]
-        else:
-            try:
-                r = self.MapRegionCls(a.path_or_fd(), 0, sys.maxsize, flags)
-            except Exception:
-                # apparently we are out of system resources or hit a limit
-                # As many more operations are likely to fail in that condition (
-                # like reading a file from disk, etc) we free up as much as possible
-                # As this invalidates our insert position, we have to recurse here
-                if is_recursive:
-                    # we already tried this, and still have no success in obtaining
-                    # a mapping. This is an exception, so we propagate it
-                    raise
-                # END handle existing recursion
-                self._collect_lru_region(0)
-                return self._obtain_region(a, offset, size, flags, True)
-            # END handle exceptions
-
-            self._handle_count += 1
-            self._memory_size += r.size()
-            a.append(r)
-        # END handle array
-
-        assert r.includes_ofs(offset)
-        return r
-
-    #}END internal methods
-
-    #{ Interface
-    def make_cursor(self, path_or_fd):
-        """
-        :return: a cursor pointing to the given path or file descriptor.
-            It can be used to map new regions of the file into memory
-
-        **Note:** if a file descriptor is given, it is assumed to be open and valid,
-        but may be closed afterwards. To refer to the same file, you may reuse
-        your existing file descriptor, but keep in mind that new windows can only
-        be mapped as long as it stays valid. This is why the using actual file paths
-        are preferred unless you plan to keep the file descriptor open.
-
-        **Note:** file descriptors are problematic as they are not necessarily unique, as two
-        different files opened and closed in succession might have the same file descriptor id.
-
-        **Note:** Using file descriptors directly is faster once new windows are mapped as it
-        prevents the file to be opened again just for the purpose of mapping it."""
-        regions = self._fdict.get(path_or_fd)
-        if regions is None:
-            regions = self.MapRegionListCls(path_or_fd)
-            self._fdict[path_or_fd] = regions
-        # END obtain region for path
-        return self.WindowCursorCls(self, regions)
-
-    def collect(self):
-        """Collect all available free-to-collect mapped regions
-        :return: Amount of freed handles"""
-        return self._collect_lru_region(0)
-
-    def num_file_handles(self):
-        """:return: amount of file handles in use. Each mapped region uses one file handle"""
-        return self._handle_count
-
-    def num_open_files(self):
-        """Amount of opened files in the system"""
-        return reduce(lambda x, y: x + y, (1 for rlist in self._fdict.values() if len(rlist) > 0), 0)
-
-    def window_size(self):
-        """:return: size of each window when allocating new regions"""
-        return self._window_size
-
-    def mapped_memory_size(self):
-        """:return: amount of bytes currently mapped in total"""
-        return self._memory_size
-
-    def max_file_handles(self):
-        """:return: maximum amount of handles we may have opened"""
-        return self._max_handle_count
-
-    def max_mapped_memory_size(self):
-        """:return: maximum amount of memory we may allocate"""
-        return self._max_memory_size
-
-    #} END interface
-
-    #{ Special Purpose Interface
-
-    def force_map_handle_removal_win(self, base_path):
-        """ONLY AVAILABLE ON WINDOWS
-        On windows removing files is not allowed if anybody still has it opened.
-        If this process is ourselves, and if the whole process uses this memory
-        manager (as far as the parent framework is concerned) we can enforce
-        closing all memory maps whose path matches the given base path to
-        allow the respective operation after all.
-        The respective system must NOT access the closed memory regions anymore !
-        This really may only be used if you know that the items which keep
-        the cursors alive will not be using it anymore. They need to be recreated !
-        :return: Amount of closed handles
-
-        **Note:** does nothing on non-windows platforms"""
-        if sys.platform != 'win32':
-            return
-        # END early bailout
-
-        num_closed = 0
-        for path, rlist in self._fdict.items():
-            if path.startswith(base_path):
-                for region in rlist:
-                    region.release()
-                    num_closed += 1
-            # END path matches
-        # END for each path
-        return num_closed
-    #} END special purpose interface
-
-
-class SlidingWindowMapManager(StaticWindowMapManager):
-
-    """Maintains a list of ranges of mapped memory regions in one or more files and allows to easily
-    obtain additional regions assuring there is no overlap.
-    Once a certain memory limit is reached globally, or if there cannot be more open file handles
-    which result from each mmap call, the least recently used, and currently unused mapped regions
-    are unloaded automatically.
-
-    **Note:** currently not thread-safe !
-
-    **Note:** in the current implementation, we will automatically unload windows if we either cannot
-        create more memory maps (as the open file handles limit is hit) or if we have allocated more than
-        a safe amount of memory already, which would possibly cause memory allocations to fail as our address
-        space is full."""
-
-    __slots__ = tuple()
-
-    def __init__(self, window_size=-1, max_memory_size=0, max_open_handles=sys.maxsize):
-        """Adjusts the default window size to -1"""
-        super().__init__(window_size, max_memory_size, max_open_handles)
-
-    def _obtain_region(self, a, offset, size, flags, is_recursive):
-        # bisect to find an existing region. The c++ implementation cannot
-        # do that as it uses a linked list for regions.
-        r = None
-        lo = 0
-        hi = len(a)
-        while lo < hi:
-            mid = (lo + hi) // 2
-            ofs = a[mid]._b
-            if ofs <= offset:
-                if a[mid].includes_ofs(offset):
-                    r = a[mid]
-                    break
-                # END have region
-                lo = mid + 1
-            else:
-                hi = mid
-            # END handle position
-        # END while bisecting
-
-        if r is None:
-            window_size = self._window_size
-            left = self.MapWindowCls(0, 0)
-            mid = self.MapWindowCls(offset, size)
-            right = self.MapWindowCls(a.file_size(), 0)
-
-            # we want to honor the max memory size, and assure we have anough
-            # memory available
-            # Save calls !
-            if self._memory_size + window_size > self._max_memory_size:
-                self._collect_lru_region(window_size)
-            # END handle collection
-
-            # we assume the list remains sorted by offset
-            insert_pos = 0
-            len_regions = len(a)
-            if len_regions == 1:
-                if a[0]._b <= offset:
-                    insert_pos = 1
-                # END maintain sort
-            else:
-                # find insert position
-                insert_pos = len_regions
-                for i, region in enumerate(a):
-                    if region._b > offset:
-                        insert_pos = i
-                        break
-                    # END if insert position is correct
-                # END for each region
-            # END obtain insert pos
-
-            # adjust the actual offset and size values to create the largest
-            # possible mapping
-            if insert_pos == 0:
-                if len_regions:
-                    right = self.MapWindowCls.from_region(a[insert_pos])
-                # END adjust right side
-            else:
-                if insert_pos != len_regions:
-                    right = self.MapWindowCls.from_region(a[insert_pos])
-                # END adjust right window
-                left = self.MapWindowCls.from_region(a[insert_pos - 1])
-            # END adjust surrounding windows
-
-            mid.extend_left_to(left, window_size)
-            mid.extend_right_to(right, window_size)
-            mid.align()
-
-            # it can happen that we align beyond the end of the file
-            if mid.ofs_end() > right.ofs:
-                mid.size = right.ofs - mid.ofs
-            # END readjust size
-
-            # insert new region at the right offset to keep the order
-            try:
-                if self._handle_count >= self._max_handle_count:
-                    raise Exception
-                # END assert own imposed max file handles
-                r = self.MapRegionCls(a.path_or_fd(), mid.ofs, mid.size, flags)
-            except Exception:
-                # apparently we are out of system resources or hit a limit
-                # As many more operations are likely to fail in that condition (
-                # like reading a file from disk, etc) we free up as much as possible
-                # As this invalidates our insert position, we have to recurse here
-                if is_recursive:
-                    # we already tried this, and still have no success in obtaining
-                    # a mapping. This is an exception, so we propagate it
-                    raise
-                # END handle existing recursion
-                self._collect_lru_region(0)
-                return self._obtain_region(a, offset, size, flags, True)
-            # END handle exceptions
-
-            self._handle_count += 1
-            self._memory_size += r.size()
-            a.insert(insert_pos, r)
-        # END create new region
-        return r
diff --git a/venv/lib/python3.10/site-packages/smmap/test/__init__.py b/venv/lib/python3.10/site-packages/smmap/test/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/venv/lib/python3.10/site-packages/smmap/test/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/smmap/test/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index be9503332031d7677754d682dda44d08684bc753..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/smmap/test/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/smmap/test/__pycache__/lib.cpython-310.pyc b/venv/lib/python3.10/site-packages/smmap/test/__pycache__/lib.cpython-310.pyc
deleted file mode 100644
index e4975041a35c9f98b2c9b7fab02f3ca8951e94a4..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/smmap/test/__pycache__/lib.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/smmap/test/__pycache__/test_buf.cpython-310.pyc b/venv/lib/python3.10/site-packages/smmap/test/__pycache__/test_buf.cpython-310.pyc
deleted file mode 100644
index d22a49223897bf467bbb149eb9816a4e853551ca..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/smmap/test/__pycache__/test_buf.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/smmap/test/__pycache__/test_mman.cpython-310.pyc b/venv/lib/python3.10/site-packages/smmap/test/__pycache__/test_mman.cpython-310.pyc
deleted file mode 100644
index a75b52dbe210653d42b7e307456fa2dfe927e3b5..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/smmap/test/__pycache__/test_mman.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/smmap/test/__pycache__/test_tutorial.cpython-310.pyc b/venv/lib/python3.10/site-packages/smmap/test/__pycache__/test_tutorial.cpython-310.pyc
deleted file mode 100644
index 942ca8dcb93b791bfe8dcad6e598da39cf3f51a5..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/smmap/test/__pycache__/test_tutorial.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/smmap/test/__pycache__/test_util.cpython-310.pyc b/venv/lib/python3.10/site-packages/smmap/test/__pycache__/test_util.cpython-310.pyc
deleted file mode 100644
index 41faf261819c85506ab08a7b6cd1b43f49cb85ac..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/smmap/test/__pycache__/test_util.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/smmap/test/lib.py b/venv/lib/python3.10/site-packages/smmap/test/lib.py
deleted file mode 100644
index b15b0ec6a8eca03b9b01f0f469f8687300039df6..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap/test/lib.py
+++ /dev/null
@@ -1,72 +0,0 @@
-"""Provide base classes for the test system"""
-from unittest import TestCase
-import os
-import tempfile
-
-__all__ = ['TestBase', 'FileCreator']
-
-
-#{ Utilities
-
-class FileCreator:
-
-    """A instance which creates a temporary file with a prefix and a given size
-    and provides this info to the user.
-    Once it gets deleted, it will remove the temporary file as well."""
-    __slots__ = ("_size", "_path")
-
-    def __init__(self, size, prefix=''):
-        assert size, "Require size to be larger 0"
-
-        self._size = size
-
-        with tempfile.NamedTemporaryFile("wb", prefix=prefix, delete=False) as file:
-            self._path = file.name
-            file.seek(size - 1)
-            file.write(b'1')
-
-        assert os.path.getsize(self.path) == size
-
-    def __del__(self):
-        try:
-            os.remove(self.path)
-        except OSError:
-            pass
-        # END exception handling
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        self.__del__()
-
-    @property
-    def path(self):
-        return self._path
-
-    @property
-    def size(self):
-        return self._size
-
-#} END utilities
-
-
-class TestBase(TestCase):
-
-    """Foundation used by all tests"""
-
-    #{ Configuration
-    k_window_test_size = 1000 * 1000 * 8 + 5195
-    #} END configuration
-
-    #{ Overrides
-    @classmethod
-    def setUpAll(cls):
-        # nothing for now
-        pass
-
-    # END overrides
-
-    #{ Interface
-
-    #} END interface
diff --git a/venv/lib/python3.10/site-packages/smmap/test/test_buf.py b/venv/lib/python3.10/site-packages/smmap/test/test_buf.py
deleted file mode 100644
index f0a86fb643b512454c60451de56b63e171e60960..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap/test/test_buf.py
+++ /dev/null
@@ -1,126 +0,0 @@
-from .lib import TestBase, FileCreator
-
-from smmap.mman import (
-    SlidingWindowMapManager,
-    StaticWindowMapManager
-)
-from smmap.buf import SlidingWindowMapBuffer
-
-from random import randint
-from time import time
-import sys
-import os
-
-
-man_optimal = SlidingWindowMapManager()
-man_worst_case = SlidingWindowMapManager(
-    window_size=TestBase.k_window_test_size // 100,
-    max_memory_size=TestBase.k_window_test_size // 3,
-    max_open_handles=15)
-static_man = StaticWindowMapManager()
-
-
-class TestBuf(TestBase):
-
-    def test_basics(self):
-        with FileCreator(self.k_window_test_size, "buffer_test") as fc:
-
-            # invalid paths fail upon construction
-            c = man_optimal.make_cursor(fc.path)
-            self.assertRaises(ValueError, SlidingWindowMapBuffer, type(c)())            # invalid cursor
-            self.assertRaises(ValueError, SlidingWindowMapBuffer, c, fc.size)       # offset too large
-
-            buf = SlidingWindowMapBuffer()                                              # can create uninitialized buffers
-            assert buf.cursor() is None
-
-            # can call end access any time
-            buf.end_access()
-            buf.end_access()
-            assert len(buf) == 0
-
-            # begin access can revive it, if the offset is suitable
-            offset = 100
-            assert buf.begin_access(c, fc.size) == False
-            assert buf.begin_access(c, offset) == True
-            assert len(buf) == fc.size - offset
-            assert buf.cursor().is_valid()
-
-            # empty begin access keeps it valid on the same path, but alters the offset
-            assert buf.begin_access() == True
-            assert len(buf) == fc.size
-            assert buf.cursor().is_valid()
-
-            # simple access
-            with open(fc.path, 'rb') as fp:
-                data = fp.read()
-            assert data[offset] == buf[0]
-            assert data[offset:offset * 2] == buf[0:offset]
-
-            # negative indices, partial slices
-            assert buf[-1] == buf[len(buf) - 1]
-            assert buf[-10:] == buf[len(buf) - 10:len(buf)]
-
-            # end access makes its cursor invalid
-            buf.end_access()
-            assert not buf.cursor().is_valid()
-            assert buf.cursor().is_associated()         # but it remains associated
-
-            # an empty begin access fixes it up again
-            assert buf.begin_access() == True and buf.cursor().is_valid()
-            del(buf)        # ends access automatically
-            del(c)
-
-            assert man_optimal.num_file_handles() == 1
-
-            # PERFORMANCE
-            # blast away with random access and a full mapping - we don't want to
-            # exaggerate the manager's overhead, but measure the buffer overhead
-            # We do it once with an optimal setting, and with a worse manager which
-            # will produce small mappings only !
-            max_num_accesses = 100
-            fd = os.open(fc.path, os.O_RDONLY)
-            for item in (fc.path, fd):
-                for manager, man_id in ((man_optimal, 'optimal'),
-                                        (man_worst_case, 'worst case'),
-                                        (static_man, 'static optimal')):
-                    buf = SlidingWindowMapBuffer(manager.make_cursor(item))
-                    assert manager.num_file_handles() == 1
-                    for access_mode in range(2):    # single, multi
-                        num_accesses_left = max_num_accesses
-                        num_bytes = 0
-                        fsize = fc.size
-
-                        st = time()
-                        buf.begin_access()
-                        while num_accesses_left:
-                            num_accesses_left -= 1
-                            if access_mode:  # multi
-                                ofs_start = randint(0, fsize)
-                                ofs_end = randint(ofs_start, fsize)
-                                d = buf[ofs_start:ofs_end]
-                                assert len(d) == ofs_end - ofs_start
-                                assert d == data[ofs_start:ofs_end]
-                                num_bytes += len(d)
-                                del d
-                            else:
-                                pos = randint(0, fsize)
-                                assert buf[pos] == data[pos]
-                                num_bytes += 1
-                            # END handle mode
-                        # END handle num accesses
-
-                        buf.end_access()
-                        assert manager.num_file_handles()
-                        assert manager.collect()
-                        assert manager.num_file_handles() == 0
-                        elapsed = max(time() - st, 0.001)  # prevent zero division errors on windows
-                        mb = float(1000 * 1000)
-                        mode_str = (access_mode and "slice") or "single byte"
-                        print("%s: Made %i random %s accesses to buffer created from %s reading a total of %f mb in %f s (%f mb/s)"
-                              % (man_id, max_num_accesses, mode_str, type(item), num_bytes / mb, elapsed, (num_bytes / mb) / elapsed),
-                              file=sys.stderr)
-                    # END handle access mode
-                    del buf
-                # END for each manager
-            # END for each input
-            os.close(fd)
diff --git a/venv/lib/python3.10/site-packages/smmap/test/test_mman.py b/venv/lib/python3.10/site-packages/smmap/test/test_mman.py
deleted file mode 100644
index 7a5f4092c11e7abb6a086a87a481a9ffa0f818aa..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap/test/test_mman.py
+++ /dev/null
@@ -1,224 +0,0 @@
-from .lib import TestBase, FileCreator
-
-from smmap.mman import (
-    WindowCursor,
-    SlidingWindowMapManager,
-    StaticWindowMapManager
-)
-from smmap.util import align_to_mmap
-
-from random import randint
-from time import time
-import os
-import sys
-from copy import copy
-
-
-class TestMMan(TestBase):
-
-    def test_cursor(self):
-        with FileCreator(self.k_window_test_size, "cursor_test") as fc:
-            man = SlidingWindowMapManager()
-            ci = WindowCursor(man)  # invalid cursor
-            assert not ci.is_valid()
-            assert not ci.is_associated()
-            assert ci.size() == 0       # this is cached, so we can query it in invalid state
-
-            cv = man.make_cursor(fc.path)
-            assert not cv.is_valid()    # no region mapped yet
-            assert cv.is_associated()  # but it know where to map it from
-            assert cv.file_size() == fc.size
-            assert cv.path() == fc.path
-
-        # copy module
-        cio = copy(cv)
-        assert not cio.is_valid() and cio.is_associated()
-
-        # assign method
-        assert not ci.is_associated()
-        ci.assign(cv)
-        assert not ci.is_valid() and ci.is_associated()
-
-        # unuse non-existing region is fine
-        cv.unuse_region()
-        cv.unuse_region()
-
-        # destruction is fine (even multiple times)
-        cv._destroy()
-        WindowCursor(man)._destroy()
-
-    def test_memory_manager(self):
-        slide_man = SlidingWindowMapManager()
-        static_man = StaticWindowMapManager()
-
-        for man in (static_man, slide_man):
-            assert man.num_file_handles() == 0
-            assert man.num_open_files() == 0
-            winsize_cmp_val = 0
-            if isinstance(man, StaticWindowMapManager):
-                winsize_cmp_val = -1
-            # END handle window size
-            assert man.window_size() > winsize_cmp_val
-            assert man.mapped_memory_size() == 0
-            assert man.max_mapped_memory_size() > 0
-
-            # collection doesn't raise in 'any' mode
-            man._collect_lru_region(0)
-            # doesn't raise if we are within the limit
-            man._collect_lru_region(10)
-
-            # doesn't fail if we over-allocate
-            assert man._collect_lru_region(sys.maxsize) == 0
-
-            # use a region, verify most basic functionality
-            with FileCreator(self.k_window_test_size, "manager_test") as fc:
-                fd = os.open(fc.path, os.O_RDONLY)
-                try:
-                    for item in (fc.path, fd):
-                        c = man.make_cursor(item)
-                        assert c.path_or_fd() is item
-                        assert c.use_region(10, 10).is_valid()
-                        assert c.ofs_begin() == 10
-                        assert c.size() == 10
-                        with open(fc.path, 'rb') as fp:
-                            assert c.buffer()[:] == fp.read(20)[10:]
-
-                    if isinstance(item, int):
-                        self.assertRaises(ValueError, c.path)
-                    else:
-                        self.assertRaises(ValueError, c.fd)
-                    # END handle value error
-                # END for each input
-                finally:
-                    os.close(fd)
-        # END for each manasger type
-
-    def test_memman_operation(self):
-        # test more access, force it to actually unmap regions
-        with FileCreator(self.k_window_test_size, "manager_operation_test") as fc:
-            with open(fc.path, 'rb') as fp:
-                data = fp.read()
-            fd = os.open(fc.path, os.O_RDONLY)
-            try:
-                max_num_handles = 15
-                # small_size =
-                for mtype, args in ((StaticWindowMapManager, (0, fc.size // 3, max_num_handles)),
-                                    (SlidingWindowMapManager, (fc.size // 100, fc.size // 3, max_num_handles)),):
-                    for item in (fc.path, fd):
-                        assert len(data) == fc.size
-
-                        # small windows, a reasonable max memory. Not too many regions at once
-                        man = mtype(window_size=args[0], max_memory_size=args[1], max_open_handles=args[2])
-                        c = man.make_cursor(item)
-
-                        # still empty (more about that is tested in test_memory_manager()
-                        assert man.num_open_files() == 0
-                        assert man.mapped_memory_size() == 0
-
-                        base_offset = 5000
-                        # window size is 0 for static managers, hence size will be 0. We take that into consideration
-                        size = man.window_size() // 2
-                        assert c.use_region(base_offset, size).is_valid()
-                        rr = c.region()
-                        assert rr.client_count() == 2  # the manager and the cursor and us
-
-                        assert man.num_open_files() == 1
-                        assert man.num_file_handles() == 1
-                        assert man.mapped_memory_size() == rr.size()
-
-                        # assert c.size() == size        # the cursor may overallocate in its static version
-                        assert c.ofs_begin() == base_offset
-                        assert rr.ofs_begin() == 0        # it was aligned and expanded
-                        if man.window_size():
-                            # but isn't larger than the max window (aligned)
-                            assert rr.size() == align_to_mmap(man.window_size(), True)
-                        else:
-                            assert rr.size() == fc.size
-                        # END ignore static managers which dont use windows and are aligned to file boundaries
-
-                        assert c.buffer()[:] == data[base_offset:base_offset + (size or c.size())]
-
-                        # obtain second window, which spans the first part of the file - it is a still the same window
-                        nsize = (size or fc.size) - 10
-                        assert c.use_region(0, nsize).is_valid()
-                        assert c.region() == rr
-                        assert man.num_file_handles() == 1
-                        assert c.size() == nsize
-                        assert c.ofs_begin() == 0
-                        assert c.buffer()[:] == data[:nsize]
-
-                        # map some part at the end, our requested size cannot be kept
-                        overshoot = 4000
-                        base_offset = fc.size - (size or c.size()) + overshoot
-                        assert c.use_region(base_offset, size).is_valid()
-                        if man.window_size():
-                            assert man.num_file_handles() == 2
-                            assert c.size() < size
-                            assert c.region() is not rr  # old region is still available, but has not cursor ref anymore
-                            assert rr.client_count() == 1  # only held by manager
-                        else:
-                            assert c.size() < fc.size
-                        # END ignore static managers which only have one handle per file
-                        rr = c.region()
-                        assert rr.client_count() == 2  # manager + cursor
-                        assert rr.ofs_begin() < c.ofs_begin()  # it should have extended itself to the left
-                        assert rr.ofs_end() <= fc.size  # it cannot be larger than the file
-                        assert c.buffer()[:] == data[base_offset:base_offset + (size or c.size())]
-
-                        # unising a region makes the cursor invalid
-                        c.unuse_region()
-                        assert not c.is_valid()
-                        if man.window_size():
-                            # but doesn't change anything regarding the handle count - we cache it and only
-                            # remove mapped regions if we have to
-                            assert man.num_file_handles() == 2
-                        # END ignore this for static managers
-
-                        # iterate through the windows, verify data contents
-                        # this will trigger map collection after a while
-                        max_random_accesses = 5000
-                        num_random_accesses = max_random_accesses
-                        memory_read = 0
-                        st = time()
-
-                        # cache everything to get some more performance
-                        includes_ofs = c.includes_ofs
-                        max_mapped_memory_size = man.max_mapped_memory_size()
-                        max_file_handles = man.max_file_handles()
-                        mapped_memory_size = man.mapped_memory_size
-                        num_file_handles = man.num_file_handles
-                        while num_random_accesses:
-                            num_random_accesses -= 1
-                            base_offset = randint(0, fc.size - 1)
-
-                            # precondition
-                            if man.window_size():
-                                assert max_mapped_memory_size >= mapped_memory_size()
-                            # END statistics will overshoot, which is fine
-                            assert max_file_handles >= num_file_handles()
-                            assert c.use_region(base_offset, (size or c.size())).is_valid()
-                            csize = c.size()
-                            assert c.buffer()[:] == data[base_offset:base_offset + csize]
-                            memory_read += csize
-
-                            assert includes_ofs(base_offset)
-                            assert includes_ofs(base_offset + csize - 1)
-                            assert not includes_ofs(base_offset + csize)
-                        # END while we should do an access
-                        elapsed = max(time() - st, 0.001)  # prevent zero division errors on windows
-                        mb = float(1000 * 1000)
-                        print("%s: Read %i mb of memory with %i random on cursor initialized with %s accesses in %fs (%f mb/s)\n"
-                              % (mtype, memory_read / mb, max_random_accesses, type(item), elapsed, (memory_read / mb) / elapsed),
-                              file=sys.stderr)
-
-                        # an offset as large as the size doesn't work !
-                        assert not c.use_region(fc.size, size).is_valid()
-
-                        # collection - it should be able to collect all
-                        assert man.num_file_handles()
-                        assert man.collect()
-                        assert man.num_file_handles() == 0
-                    # END for each item
-                # END for each manager type
-            finally:
-                os.close(fd)
diff --git a/venv/lib/python3.10/site-packages/smmap/test/test_tutorial.py b/venv/lib/python3.10/site-packages/smmap/test/test_tutorial.py
deleted file mode 100644
index 31c272abbf09361b6c5dc2077e853a8d47958458..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap/test/test_tutorial.py
+++ /dev/null
@@ -1,75 +0,0 @@
-from .lib import TestBase
-
-
-class TestTutorial(TestBase):
-
-    def test_example(self):
-        # Memory Managers
-        ##################
-        import smmap
-        # This instance should be globally available in your application
-        # It is configured to be well suitable for 32-bit or 64 bit applications.
-        mman = smmap.SlidingWindowMapManager()
-
-        # the manager provides much useful information about its current state
-        # like the amount of open file handles or the amount of mapped memory
-        assert mman.num_file_handles() == 0
-        assert mman.mapped_memory_size() == 0
-        # and many more ...
-
-        # Cursors
-        ##########
-        import smmap.test.lib
-        with smmap.test.lib.FileCreator(1024 * 1024 * 8, "test_file") as fc:
-            # obtain a cursor to access some file.
-            c = mman.make_cursor(fc.path)
-
-            # the cursor is now associated with the file, but not yet usable
-            assert c.is_associated()
-            assert not c.is_valid()
-
-            # before you can use the cursor, you have to specify a window you want to
-            # access. The following just says you want as much data as possible starting
-            # from offset 0.
-            # To be sure your region could be mapped, query for validity
-            assert c.use_region().is_valid()        # use_region returns self
-
-            # once a region was mapped, you must query its dimension regularly
-            # to assure you don't try to access its buffer out of its bounds
-            assert c.size()
-            c.buffer()[0]           # first byte
-            c.buffer()[1:10]            # first 9 bytes
-            c.buffer()[c.size() - 1]  # last byte
-
-            # you can query absolute offsets, and check whether an offset is included
-            # in the cursor's data.
-            assert c.ofs_begin() < c.ofs_end()
-            assert c.includes_ofs(100)
-
-            # If you are over out of bounds with one of your region requests, the
-            # cursor will be come invalid. It cannot be used in that state
-            assert not c.use_region(fc.size, 100).is_valid()
-            # map as much as possible after skipping the first 100 bytes
-            assert c.use_region(100).is_valid()
-
-            # You can explicitly free cursor resources by unusing the cursor's region
-            c.unuse_region()
-            assert not c.is_valid()
-
-            # Buffers
-            #########
-            # Create a default buffer which can operate on the whole file
-            buf = smmap.SlidingWindowMapBuffer(mman.make_cursor(fc.path))
-
-            # you can use it right away
-            assert buf.cursor().is_valid()
-
-            buf[0]  # access the first byte
-            buf[-1]  # access the last ten bytes on the file
-            buf[-10:]  # access the last ten bytes
-
-            # If you want to keep the instance between different accesses, use the
-            # dedicated methods
-            buf.end_access()
-            assert not buf.cursor().is_valid()  # you cannot use the buffer anymore
-            assert buf.begin_access(offset=10)  # start using the buffer at an offset
diff --git a/venv/lib/python3.10/site-packages/smmap/test/test_util.py b/venv/lib/python3.10/site-packages/smmap/test/test_util.py
deleted file mode 100644
index e6ac10f35cb0841947b17f08a49ce52176ded51b..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap/test/test_util.py
+++ /dev/null
@@ -1,105 +0,0 @@
-from .lib import TestBase, FileCreator
-
-from smmap.util import (
-    MapWindow,
-    MapRegion,
-    MapRegionList,
-    ALLOCATIONGRANULARITY,
-    is_64_bit,
-    align_to_mmap
-)
-
-import os
-import sys
-
-
-class TestMMan(TestBase):
-
-    def test_window(self):
-        wl = MapWindow(0, 1)        # left
-        wc = MapWindow(1, 1)        # center
-        wc2 = MapWindow(10, 5)      # another center
-        wr = MapWindow(8000, 50)    # right
-
-        assert wl.ofs_end() == 1
-        assert wc.ofs_end() == 2
-        assert wr.ofs_end() == 8050
-
-        # extension does nothing if already in place
-        maxsize = 100
-        wc.extend_left_to(wl, maxsize)
-        assert wc.ofs == 1 and wc.size == 1
-        wl.extend_right_to(wc, maxsize)
-        wl.extend_right_to(wc, maxsize)
-        assert wl.ofs == 0 and wl.size == 1
-
-        # an actual left extension
-        pofs_end = wc2.ofs_end()
-        wc2.extend_left_to(wc, maxsize)
-        assert wc2.ofs == wc.ofs_end() and pofs_end == wc2.ofs_end()
-
-        # respects maxsize
-        wc.extend_right_to(wr, maxsize)
-        assert wc.ofs == 1 and wc.size == maxsize
-        wc.extend_right_to(wr, maxsize)
-        assert wc.ofs == 1 and wc.size == maxsize
-
-        # without maxsize
-        wc.extend_right_to(wr, sys.maxsize)
-        assert wc.ofs_end() == wr.ofs and wc.ofs == 1
-
-        # extend left
-        wr.extend_left_to(wc2, maxsize)
-        wr.extend_left_to(wc2, maxsize)
-        assert wr.size == maxsize
-
-        wr.extend_left_to(wc2, sys.maxsize)
-        assert wr.ofs == wc2.ofs_end()
-
-        wc.align()
-        assert wc.ofs == 0 and wc.size == align_to_mmap(wc.size, True)
-
-    def test_region(self):
-        with FileCreator(self.k_window_test_size, "window_test") as fc:
-            half_size = fc.size // 2
-            rofs = align_to_mmap(4200, False)
-            rfull = MapRegion(fc.path, 0, fc.size)
-            rhalfofs = MapRegion(fc.path, rofs, fc.size)
-            rhalfsize = MapRegion(fc.path, 0, half_size)
-
-            # offsets
-            assert rfull.ofs_begin() == 0 and rfull.size() == fc.size
-            assert rfull.ofs_end() == fc.size   # if this method works, it works always
-
-            assert rhalfofs.ofs_begin() == rofs and rhalfofs.size() == fc.size - rofs
-            assert rhalfsize.ofs_begin() == 0 and rhalfsize.size() == half_size
-
-            assert rfull.includes_ofs(0) and rfull.includes_ofs(fc.size - 1) and rfull.includes_ofs(half_size)
-            assert not rfull.includes_ofs(-1) and not rfull.includes_ofs(sys.maxsize)
-
-        # auto-refcount
-        assert rfull.client_count() == 1
-        rfull2 = rfull
-        assert rfull.client_count() == 1, "no auto-counting"
-
-        # window constructor
-        w = MapWindow.from_region(rfull)
-        assert w.ofs == rfull.ofs_begin() and w.ofs_end() == rfull.ofs_end()
-
-    def test_region_list(self):
-        with FileCreator(100, "sample_file") as fc:
-            fd = os.open(fc.path, os.O_RDONLY)
-            try:
-                for item in (fc.path, fd):
-                    ml = MapRegionList(item)
-
-                    assert len(ml) == 0
-                    assert ml.path_or_fd() == item
-                    assert ml.file_size() == fc.size
-            finally:
-                os.close(fd)
-
-    def test_util(self):
-        assert isinstance(is_64_bit(), bool)    # just call it
-        assert align_to_mmap(1, False) == 0
-        assert align_to_mmap(1, True) == ALLOCATIONGRANULARITY
diff --git a/venv/lib/python3.10/site-packages/smmap/util.py b/venv/lib/python3.10/site-packages/smmap/util.py
deleted file mode 100644
index fbb3872118e644ee0a920914fa5835cf2bfab44d..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/smmap/util.py
+++ /dev/null
@@ -1,222 +0,0 @@
-"""Module containing a memory memory manager which provides a sliding window on a number of memory mapped files"""
-import os
-import sys
-
-from mmap import mmap, ACCESS_READ
-from mmap import ALLOCATIONGRANULARITY
-
-__all__ = ["align_to_mmap", "is_64_bit",
-           "MapWindow", "MapRegion", "MapRegionList", "ALLOCATIONGRANULARITY"]
-
-#{ Utilities
-
-
-def align_to_mmap(num, round_up):
-    """
-    Align the given integer number to the closest page offset, which usually is 4096 bytes.
-
-    :param round_up: if True, the next higher multiple of page size is used, otherwise
-        the lower page_size will be used (i.e. if True, 1 becomes 4096, otherwise it becomes 0)
-    :return: num rounded to closest page"""
-    res = (num // ALLOCATIONGRANULARITY) * ALLOCATIONGRANULARITY
-    if round_up and (res != num):
-        res += ALLOCATIONGRANULARITY
-    # END handle size
-    return res
-
-
-def is_64_bit():
-    """:return: True if the system is 64 bit. Otherwise it can be assumed to be 32 bit"""
-    return sys.maxsize > (1 << 32) - 1
-
-#}END utilities
-
-
-#{ Utility Classes
-
-class MapWindow:
-
-    """Utility type which is used to snap windows towards each other, and to adjust their size"""
-    __slots__ = (
-        'ofs',      # offset into the file in bytes
-        'size'              # size of the window in bytes
-    )
-
-    def __init__(self, offset, size):
-        self.ofs = offset
-        self.size = size
-
-    def __repr__(self):
-        return "MapWindow(%i, %i)" % (self.ofs, self.size)
-
-    @classmethod
-    def from_region(cls, region):
-        """:return: new window from a region"""
-        return cls(region._b, region.size())
-
-    def ofs_end(self):
-        return self.ofs + self.size
-
-    def align(self):
-        """Assures the previous window area is contained in the new one"""
-        nofs = align_to_mmap(self.ofs, 0)
-        self.size += self.ofs - nofs    # keep size constant
-        self.ofs = nofs
-        self.size = align_to_mmap(self.size, 1)
-
-    def extend_left_to(self, window, max_size):
-        """Adjust the offset to start where the given window on our left ends if possible,
-        but don't make yourself larger than max_size.
-        The resize will assure that the new window still contains the old window area"""
-        rofs = self.ofs - window.ofs_end()
-        nsize = rofs + self.size
-        rofs -= nsize - min(nsize, max_size)
-        self.ofs -= rofs
-        self.size += rofs
-
-    def extend_right_to(self, window, max_size):
-        """Adjust the size to make our window end where the right window begins, but don't
-        get larger than max_size"""
-        self.size = min(self.size + (window.ofs - self.ofs_end()), max_size)
-
-
-class MapRegion:
-
-    """Defines a mapped region of memory, aligned to pagesizes
-
-    **Note:** deallocates used region automatically on destruction"""
-    __slots__ = [
-        '_b',   # beginning of mapping
-        '_mf',  # mapped memory chunk (as returned by mmap)
-        '_uc',  # total amount of usages
-        '_size',  # cached size of our memory map
-        '__weakref__'
-    ]
-
-    #{ Configuration
-    #} END configuration
-
-    def __init__(self, path_or_fd, ofs, size, flags=0):
-        """Initialize a region, allocate the memory map
-        :param path_or_fd: path to the file to map, or the opened file descriptor
-        :param ofs: **aligned** offset into the file to be mapped
-        :param size: if size is larger then the file on disk, the whole file will be
-            allocated the the size automatically adjusted
-        :param flags: additional flags to be given when opening the file.
-        :raise Exception: if no memory can be allocated"""
-        self._b = ofs
-        self._size = 0
-        self._uc = 0
-
-        if isinstance(path_or_fd, int):
-            fd = path_or_fd
-        else:
-            fd = os.open(path_or_fd, os.O_RDONLY | getattr(os, 'O_BINARY', 0) | flags)
-        # END handle fd
-
-        try:
-            kwargs = dict(access=ACCESS_READ, offset=ofs)
-            corrected_size = size
-            sizeofs = ofs
-
-            # have to correct size, otherwise (instead of the c version) it will
-            # bark that the size is too large ... many extra file accesses because
-            # if this ... argh !
-            actual_size = min(os.fstat(fd).st_size - sizeofs, corrected_size)
-            self._mf = mmap(fd, actual_size, **kwargs)
-            # END handle memory mode
-
-            self._size = len(self._mf)
-        finally:
-            if isinstance(path_or_fd, str):
-                os.close(fd)
-            # END only close it if we opened it
-        # END close file handle
-        # We assume the first one to use us keeps us around
-        self.increment_client_count()
-
-    def __repr__(self):
-        return "MapRegion<%i, %i>" % (self._b, self.size())
-
-    #{ Interface
-
-    def buffer(self):
-        """:return: a buffer containing the memory"""
-        return self._mf
-
-    def map(self):
-        """:return: a memory map containing the memory"""
-        return self._mf
-
-    def ofs_begin(self):
-        """:return: absolute byte offset to the first byte of the mapping"""
-        return self._b
-
-    def size(self):
-        """:return: total size of the mapped region in bytes"""
-        return self._size
-
-    def ofs_end(self):
-        """:return: Absolute offset to one byte beyond the mapping into the file"""
-        return self._b + self._size
-
-    def includes_ofs(self, ofs):
-        """:return: True if the given offset can be read in our mapped region"""
-        return self._b <= ofs < self._b + self._size
-
-    def client_count(self):
-        """:return: number of clients currently using this region"""
-        return self._uc
-
-    def increment_client_count(self, ofs = 1):
-        """Adjust the usage count by the given positive or negative offset.
-        If usage count equals 0, we will auto-release our resources
-        :return: True if we released resources, False otherwise. In the latter case, we can still be used"""
-        self._uc += ofs
-        assert self._uc > -1, "Increments must match decrements, usage counter negative: %i" % self._uc
-
-        if self.client_count() == 0:
-            self.release()
-            return True
-        else:
-            return False
-        # end handle release
-
-    def release(self):
-        """Release all resources this instance might hold. Must only be called if there usage_count() is zero"""
-        self._mf.close()
-
-    #} END interface
-
-
-class MapRegionList(list):
-
-    """List of MapRegion instances associating a path with a list of regions."""
-    __slots__ = (
-        '_path_or_fd',  # path or file descriptor which is mapped by all our regions
-        '_file_size'    # total size of the file we map
-    )
-
-    def __new__(cls, path):
-        return super().__new__(cls)
-
-    def __init__(self, path_or_fd):
-        self._path_or_fd = path_or_fd
-        self._file_size = None
-
-    def path_or_fd(self):
-        """:return: path or file descriptor we are attached to"""
-        return self._path_or_fd
-
-    def file_size(self):
-        """:return: size of file we manager"""
-        if self._file_size is None:
-            if isinstance(self._path_or_fd, str):
-                self._file_size = os.stat(self._path_or_fd).st_size
-            else:
-                self._file_size = os.fstat(self._path_or_fd).st_size
-            # END handle path type
-        # END update file size
-        return self._file_size
-
-#} END utility classes
diff --git a/venv/lib/python3.10/site-packages/soundfile.py b/venv/lib/python3.10/site-packages/soundfile.py
deleted file mode 100644
index 61db56cbd2534be8924f4915f162183c1b25ad08..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/soundfile.py
+++ /dev/null
@@ -1,1674 +0,0 @@
-"""python-soundfile is an audio library based on libsndfile, CFFI and NumPy.
-
-Sound files can be read or written directly using the functions
-`read()` and `write()`.
-To read a sound file in a block-wise fashion, use `blocks()`.
-Alternatively, sound files can be opened as `SoundFile` objects.
-
-For further information, see https://python-soundfile.readthedocs.io/.
-
-"""
-__version__ = "0.13.1"
-
-import os as _os
-import sys as _sys
-from os import SEEK_SET, SEEK_CUR, SEEK_END
-from ctypes.util import find_library as _find_library
-from _soundfile import ffi as _ffi
-
-try:
-    _unicode = unicode  # doesn't exist in Python 3.x
-except NameError:
-    _unicode = str
-
-
-_str_types = {
-    'title':       0x01,
-    'copyright':   0x02,
-    'software':    0x03,
-    'artist':      0x04,
-    'comment':     0x05,
-    'date':        0x06,
-    'album':       0x07,
-    'license':     0x08,
-    'tracknumber': 0x09,
-    'genre':       0x10,
-}
-
-_formats = {
-    'WAV':   0x010000,  # Microsoft WAV format (little endian default).
-    'AIFF':  0x020000,  # Apple/SGI AIFF format (big endian).
-    'AU':    0x030000,  # Sun/NeXT AU format (big endian).
-    'RAW':   0x040000,  # RAW PCM data.
-    'PAF':   0x050000,  # Ensoniq PARIS file format.
-    'SVX':   0x060000,  # Amiga IFF / SVX8 / SV16 format.
-    'NIST':  0x070000,  # Sphere NIST format.
-    'VOC':   0x080000,  # VOC files.
-    'IRCAM': 0x0A0000,  # Berkeley/IRCAM/CARL
-    'W64':   0x0B0000,  # Sonic Foundry's 64 bit RIFF/WAV
-    'MAT4':  0x0C0000,  # Matlab (tm) V4.2 / GNU Octave 2.0
-    'MAT5':  0x0D0000,  # Matlab (tm) V5.0 / GNU Octave 2.1
-    'PVF':   0x0E0000,  # Portable Voice Format
-    'XI':    0x0F0000,  # Fasttracker 2 Extended Instrument
-    'HTK':   0x100000,  # HMM Tool Kit format
-    'SDS':   0x110000,  # Midi Sample Dump Standard
-    'AVR':   0x120000,  # Audio Visual Research
-    'WAVEX': 0x130000,  # MS WAVE with WAVEFORMATEX
-    'SD2':   0x160000,  # Sound Designer 2
-    'FLAC':  0x170000,  # FLAC lossless file format
-    'CAF':   0x180000,  # Core Audio File format
-    'WVE':   0x190000,  # Psion WVE format
-    'OGG':   0x200000,  # Xiph OGG container
-    'MPC2K': 0x210000,  # Akai MPC 2000 sampler
-    'RF64':  0x220000,  # RF64 WAV file
-    'MP3':   0x230000,  # MPEG-1/2 audio stream
-}
-
-_subtypes = {
-    'PCM_S8':         0x0001,  # Signed 8 bit data
-    'PCM_16':         0x0002,  # Signed 16 bit data
-    'PCM_24':         0x0003,  # Signed 24 bit data
-    'PCM_32':         0x0004,  # Signed 32 bit data
-    'PCM_U8':         0x0005,  # Unsigned 8 bit data (WAV and RAW only)
-    'FLOAT':          0x0006,  # 32 bit float data
-    'DOUBLE':         0x0007,  # 64 bit float data
-    'ULAW':           0x0010,  # U-Law encoded.
-    'ALAW':           0x0011,  # A-Law encoded.
-    'IMA_ADPCM':      0x0012,  # IMA ADPCM.
-    'MS_ADPCM':       0x0013,  # Microsoft ADPCM.
-    'GSM610':         0x0020,  # GSM 6.10 encoding.
-    'VOX_ADPCM':      0x0021,  # OKI / Dialogix ADPCM
-    'NMS_ADPCM_16':   0x0022,  # 16kbs NMS G721-variant encoding.
-    'NMS_ADPCM_24':   0x0023,  # 24kbs NMS G721-variant encoding.
-    'NMS_ADPCM_32':   0x0024,  # 32kbs NMS G721-variant encoding.
-    'G721_32':        0x0030,  # 32kbs G721 ADPCM encoding.
-    'G723_24':        0x0031,  # 24kbs G723 ADPCM encoding.
-    'G723_40':        0x0032,  # 40kbs G723 ADPCM encoding.
-    'DWVW_12':        0x0040,  # 12 bit Delta Width Variable Word encoding.
-    'DWVW_16':        0x0041,  # 16 bit Delta Width Variable Word encoding.
-    'DWVW_24':        0x0042,  # 24 bit Delta Width Variable Word encoding.
-    'DWVW_N':         0x0043,  # N bit Delta Width Variable Word encoding.
-    'DPCM_8':         0x0050,  # 8 bit differential PCM (XI only)
-    'DPCM_16':        0x0051,  # 16 bit differential PCM (XI only)
-    'VORBIS':         0x0060,  # Xiph Vorbis encoding.
-    'OPUS':           0x0064,  # Xiph/Skype Opus encoding.
-    'ALAC_16':        0x0070,  # Apple Lossless Audio Codec (16 bit).
-    'ALAC_20':        0x0071,  # Apple Lossless Audio Codec (20 bit).
-    'ALAC_24':        0x0072,  # Apple Lossless Audio Codec (24 bit).
-    'ALAC_32':        0x0073,  # Apple Lossless Audio Codec (32 bit).
-    'MPEG_LAYER_I':   0x0080,  # MPEG-1 Audio Layer I.
-    'MPEG_LAYER_II':  0x0081,  # MPEG-1 Audio Layer II.
-    'MPEG_LAYER_III': 0x0082,  # MPEG-2 Audio Layer III.
-}
-
-_endians = {
-    'FILE':   0x00000000,  # Default file endian-ness.
-    'LITTLE': 0x10000000,  # Force little endian-ness.
-    'BIG':    0x20000000,  # Force big endian-ness.
-    'CPU':    0x30000000,  # Force CPU endian-ness.
-}
-
-# libsndfile doesn't specify default subtypes, these are somehow arbitrary:
-_default_subtypes = {
-    'WAV':   'PCM_16',
-    'AIFF':  'PCM_16',
-    'AU':    'PCM_16',
-    # 'RAW':  # subtype must be explicit!
-    'PAF':   'PCM_16',
-    'SVX':   'PCM_16',
-    'NIST':  'PCM_16',
-    'VOC':   'PCM_16',
-    'IRCAM': 'PCM_16',
-    'W64':   'PCM_16',
-    'MAT4':  'DOUBLE',
-    'MAT5':  'DOUBLE',
-    'PVF':   'PCM_16',
-    'XI':    'DPCM_16',
-    'HTK':   'PCM_16',
-    'SDS':   'PCM_16',
-    'AVR':   'PCM_16',
-    'WAVEX': 'PCM_16',
-    'SD2':   'PCM_16',
-    'FLAC':  'PCM_16',
-    'CAF':   'PCM_16',
-    'WVE':   'ALAW',
-    'OGG':   'VORBIS',
-    'MPC2K': 'PCM_16',
-    'RF64':  'PCM_16',
-    'MP3':   'MPEG_LAYER_III',
-}
-
-_ffi_types = {
-    'float64': 'double',
-    'float32': 'float',
-    'int32': 'int',
-    'int16': 'short'
-}
-
-_bitrate_modes = {
-    'CONSTANT': 0,
-    'AVERAGE': 1,
-    'VARIABLE': 2,
-}
-
-try:  # packaged lib (in _soundfile_data which should be on python path)
-    if _sys.platform == 'darwin':
-        from platform import machine as _machine
-        _packaged_libname = 'libsndfile_' + _machine() + '.dylib'
-    elif _sys.platform == 'win32':
-        from platform import architecture as _architecture
-        from platform import machine as _machine
-        # this check can not be completed correctly: for x64 binaries running on
-        # arm64 Windows report the same values as arm64 binaries. For now, neither
-        # numpy nor cffi are available for arm64, so we can safely assume we're
-        # in x86 land:
-        if _architecture()[0] == '64bit':
-            _packaged_libname = 'libsndfile_x64.dll'
-        elif _architecture()[0] == '32bit':
-            _packaged_libname = 'libsndfile_x86.dll'
-        else:
-            raise OSError('no packaged library for Windows {} {}'
-                          .format(_architecture(), _machine()))
-    elif _sys.platform == 'linux':
-        from platform import machine as _machine
-        if _machine() in ["aarch64", "aarch64_be", "armv8b", "armv8l"]:
-            _packaged_libname = 'libsndfile_arm64.so'
-        else:
-            _packaged_libname = 'libsndfile_' + _machine() + '.so'
-    else:
-        raise OSError('no packaged library for this platform')
-
-    import _soundfile_data  # ImportError if this doesn't exist
-    _path = _os.path.dirname(_soundfile_data.__file__)  # TypeError if __file__ is None
-    _full_path = _os.path.join(_path, _packaged_libname)
-    _snd = _ffi.dlopen(_full_path)  # OSError if file doesn't exist or can't be loaded
-
-except (OSError, ImportError, TypeError):
-    try:  # system-wide libsndfile:
-        _libname = _find_library('sndfile')
-        if _libname is None:
-            raise OSError('sndfile library not found using ctypes.util.find_library')
-        _snd = _ffi.dlopen(_libname)
-
-    except OSError:
-        # Try explicit file name, if the general does not work (e.g. on nixos)
-        if _sys.platform == 'darwin':
-            _explicit_libname = 'libsndfile.dylib'
-        elif _sys.platform == 'win32':
-            _explicit_libname = 'libsndfile.dll'
-        elif _sys.platform == 'linux':
-            _explicit_libname = 'libsndfile.so'
-        else:
-            raise
-
-        # Homebrew on Apple M1 uses a `/opt/homebrew/lib` instead of
-        # `/usr/local/lib`. We are making sure we pick that up.
-        from platform import machine as _machine
-        if _sys.platform == 'darwin' and _machine() == 'arm64':
-            _hbrew_path = '/opt/homebrew/lib/' if _os.path.isdir('/opt/homebrew/lib/') \
-                else '/usr/local/lib/'
-            _snd = _ffi.dlopen(_os.path.join(_hbrew_path, _explicit_libname))
-        else:
-            _snd = _ffi.dlopen(_explicit_libname)
-
-__libsndfile_version__ = _ffi.string(_snd.sf_version_string()).decode('utf-8', 'replace')
-if __libsndfile_version__.startswith('libsndfile-'):
-    __libsndfile_version__ = __libsndfile_version__[len('libsndfile-'):]
-
-
-def read(file, frames=-1, start=0, stop=None, dtype='float64', always_2d=False,
-         fill_value=None, out=None, samplerate=None, channels=None,
-         format=None, subtype=None, endian=None, closefd=True):
-    """Provide audio data from a sound file as NumPy array.
-
-    By default, the whole file is read from the beginning, but the
-    position to start reading can be specified with *start* and the
-    number of frames to read can be specified with *frames*.
-    Alternatively, a range can be specified with *start* and *stop*.
-
-    If there is less data left in the file than requested, the rest of
-    the frames are filled with *fill_value*.
-    If no *fill_value* is specified, a smaller array is returned.
-
-    Parameters
-    ----------
-    file : str or int or file-like object
-        The file to read from.  See `SoundFile` for details.
-    frames : int, optional
-        The number of frames to read. If *frames* is negative, the whole
-        rest of the file is read.  Not allowed if *stop* is given.
-    start : int, optional
-        Where to start reading.  A negative value counts from the end.
-    stop : int, optional
-        The index after the last frame to be read.  A negative value
-        counts from the end.  Not allowed if *frames* is given.
-    dtype : {'float64', 'float32', 'int32', 'int16'}, optional
-        Data type of the returned array, by default ``'float64'``.
-        Floating point audio data is typically in the range from
-        ``-1.0`` to ``1.0``.  Integer data is in the range from
-        ``-2**15`` to ``2**15-1`` for ``'int16'`` and from ``-2**31`` to
-        ``2**31-1`` for ``'int32'``.
-
-        .. note:: Reading int values from a float file will *not*
-            scale the data to [-1.0, 1.0). If the file contains
-            ``np.array([42.6], dtype='float32')``, you will read
-            ``np.array([43], dtype='int32')`` for ``dtype='int32'``.
-
-    Returns
-    -------
-    audiodata : `numpy.ndarray` or type(out)
-        A two-dimensional (frames x channels) NumPy array is returned.
-        If the sound file has only one channel, a one-dimensional array
-        is returned.  Use ``always_2d=True`` to return a two-dimensional
-        array anyway.
-
-        If *out* was specified, it is returned.  If *out* has more
-        frames than available in the file (or if *frames* is smaller
-        than the length of *out*) and no *fill_value* is given, then
-        only a part of *out* is overwritten and a view containing all
-        valid frames is returned.
-    samplerate : int
-        The sample rate of the audio file.
-
-    Other Parameters
-    ----------------
-    always_2d : bool, optional
-        By default, reading a mono sound file will return a
-        one-dimensional array.  With ``always_2d=True``, audio data is
-        always returned as a two-dimensional array, even if the audio
-        file has only one channel.
-    fill_value : float, optional
-        If more frames are requested than available in the file, the
-        rest of the output is be filled with *fill_value*.  If
-        *fill_value* is not specified, a smaller array is returned.
-    out : `numpy.ndarray` or subclass, optional
-        If *out* is specified, the data is written into the given array
-        instead of creating a new array.  In this case, the arguments
-        *dtype* and *always_2d* are silently ignored!  If *frames* is
-        not given, it is obtained from the length of *out*.
-    samplerate, channels, format, subtype, endian, closefd
-        See `SoundFile`.
-
-    Examples
-    --------
-    >>> import soundfile as sf
-    >>> data, samplerate = sf.read('stereo_file.wav')
-    >>> data
-    array([[ 0.71329652,  0.06294799],
-           [-0.26450912, -0.38874483],
-           ...
-           [ 0.67398441, -0.11516333]])
-    >>> samplerate
-    44100
-
-    """
-    with SoundFile(file, 'r', samplerate, channels,
-                   subtype, endian, format, closefd) as f:
-        frames = f._prepare_read(start, stop, frames)
-        data = f.read(frames, dtype, always_2d, fill_value, out)
-    return data, f.samplerate
-
-
-def write(file, data, samplerate, subtype=None, endian=None, format=None,
-          closefd=True, compression_level=None, bitrate_mode=None):
-    """Write data to a sound file.
-
-    .. note:: If *file* exists, it will be truncated and overwritten!
-
-    Parameters
-    ----------
-    file : str or int or file-like object
-        The file to write to.  See `SoundFile` for details.
-    data : array_like
-        The data to write.  Usually two-dimensional (frames x channels),
-        but one-dimensional *data* can be used for mono files.
-        Only the data types ``'float64'``, ``'float32'``, ``'int32'``
-        and ``'int16'`` are supported.
-
-        .. note:: The data type of *data* does **not** select the data
-                  type of the written file. Audio data will be
-                  converted to the given *subtype*. Writing int values
-                  to a float file will *not* scale the values to
-                  [-1.0, 1.0). If you write the value ``np.array([42],
-                  dtype='int32')``, to a ``subtype='FLOAT'`` file, the
-                  file will then contain ``np.array([42.],
-                  dtype='float32')``.
-
-    samplerate : int
-        The sample rate of the audio data.
-    subtype : str, optional
-        See `default_subtype()` for the default value and
-        `available_subtypes()` for all possible values.
-
-    Other Parameters
-    ----------------
-    format, endian, closefd, compression_level, bitrate_mode
-        See `SoundFile`.
-
-    Examples
-    --------
-    Write 10 frames of random data to a new file:
-
-    >>> import numpy as np
-    >>> import soundfile as sf
-    >>> sf.write('stereo_file.wav', np.random.randn(10, 2), 44100, 'PCM_24')
-
-    """
-    import numpy as np
-    data = np.asarray(data)
-    if data.ndim == 1:
-        channels = 1
-    else:
-        channels = data.shape[1]
-    with SoundFile(file, 'w', samplerate, channels,
-                   subtype, endian, format, closefd,
-                   compression_level, bitrate_mode) as f:
-        f.write(data)
-
-
-def blocks(file, blocksize=None, overlap=0, frames=-1, start=0, stop=None,
-           dtype='float64', always_2d=False, fill_value=None, out=None,
-           samplerate=None, channels=None,
-           format=None, subtype=None, endian=None, closefd=True):
-    """Return a generator for block-wise reading.
-
-    By default, iteration starts at the beginning and stops at the end
-    of the file.  Use *start* to start at a later position and *frames*
-    or *stop* to stop earlier.
-
-    If you stop iterating over the generator before it's exhausted,
-    the sound file is not closed. This is normally not a problem
-    because the file is opened in read-only mode. To close the file
-    properly, the generator's ``close()`` method can be called.
-
-    Parameters
-    ----------
-    file : str or int or file-like object
-        The file to read from.  See `SoundFile` for details.
-    blocksize : int
-        The number of frames to read per block.
-        Either this or *out* must be given.
-    overlap : int, optional
-        The number of frames to rewind between each block.
-
-    Yields
-    ------
-    `numpy.ndarray` or type(out)
-        Blocks of audio data.
-        If *out* was given, and the requested frames are not an integer
-        multiple of the length of *out*, and no *fill_value* was given,
-        the last block will be a smaller view into *out*.
-
-    Other Parameters
-    ----------------
-    frames, start, stop
-        See `read()`.
-    dtype : {'float64', 'float32', 'int32', 'int16'}, optional
-        See `read()`.
-    always_2d, fill_value, out
-        See `read()`.
-    samplerate, channels, format, subtype, endian, closefd
-        See `SoundFile`.
-
-    Examples
-    --------
-    >>> import soundfile as sf
-    >>> for block in sf.blocks('stereo_file.wav', blocksize=1024):
-    >>>     pass  # do something with 'block'
-
-    """
-    with SoundFile(file, 'r', samplerate, channels,
-                   subtype, endian, format, closefd) as f:
-        frames = f._prepare_read(start, stop, frames)
-        for block in f.blocks(blocksize, overlap, frames,
-                              dtype, always_2d, fill_value, out):
-            yield block
-
-
-class _SoundFileInfo(object):
-    """Information about a SoundFile"""
-
-    def __init__(self, file, verbose):
-        self.verbose = verbose
-        with SoundFile(file) as f:
-            self.name = f.name
-            self.samplerate = f.samplerate
-            self.channels = f.channels
-            self.frames = f.frames
-            self.duration = float(self.frames)/f.samplerate
-            self.format = f.format
-            self.subtype = f.subtype
-            self.endian = f.endian
-            self.format_info = f.format_info
-            self.subtype_info = f.subtype_info
-            self.sections = f.sections
-            self.extra_info = f.extra_info
-
-    @property
-    def _duration_str(self):
-        hours, rest = divmod(self.duration, 3600)
-        minutes, seconds = divmod(rest, 60)
-        if hours >= 1:
-            duration = "{0:.0g}:{1:02.0g}:{2:05.3f} h".format(hours, minutes, seconds)
-        elif minutes >= 1:
-            duration = "{0:02.0g}:{1:05.3f} min".format(minutes, seconds)
-        elif seconds <= 1:
-            duration = "{0:d} samples".format(self.frames)
-        else:
-            duration = "{0:.3f} s".format(seconds)
-        return duration
-
-    def __repr__(self):
-        info = "\n".join(
-            ["{0.name}",
-             "samplerate: {0.samplerate} Hz",
-             "channels: {0.channels}",
-             "duration: {0._duration_str}",
-             "format: {0.format_info} [{0.format}]",
-             "subtype: {0.subtype_info} [{0.subtype}]"])
-        if self.verbose:
-            info += "\n".join(
-                ["\nendian: {0.endian}",
-                 "sections: {0.sections}",
-                 "frames: {0.frames}",
-                 'extra_info: """',
-                 '    {1}"""'])
-        indented_extra_info = ("\n"+" "*4).join(self.extra_info.split("\n"))
-        return info.format(self, indented_extra_info)
-
-
-def info(file, verbose=False):
-    """Returns an object with information about a `SoundFile`.
-
-    Parameters
-    ----------
-    verbose : bool
-        Whether to print additional information.
-    """
-    return _SoundFileInfo(file, verbose)
-
-
-def available_formats():
-    """Return a dictionary of available major formats.
-
-    Examples
-    --------
-    >>> import soundfile as sf
-    >>> sf.available_formats()
-    {'FLAC': 'FLAC (FLAC Lossless Audio Codec)',
-     'OGG': 'OGG (OGG Container format)',
-     'WAV': 'WAV (Microsoft)',
-     'AIFF': 'AIFF (Apple/SGI)',
-     ...
-     'WAVEX': 'WAVEX (Microsoft)',
-     'RAW': 'RAW (header-less)',
-     'MAT5': 'MAT5 (GNU Octave 2.1 / Matlab 5.0)'}
-
-    """
-    return dict(_available_formats_helper(_snd.SFC_GET_FORMAT_MAJOR_COUNT,
-                                          _snd.SFC_GET_FORMAT_MAJOR))
-
-
-def available_subtypes(format=None):
-    """Return a dictionary of available subtypes.
-
-    Parameters
-    ----------
-    format : str
-        If given, only compatible subtypes are returned.
-
-    Examples
-    --------
-    >>> import soundfile as sf
-    >>> sf.available_subtypes('FLAC')
-    {'PCM_24': 'Signed 24 bit PCM',
-     'PCM_16': 'Signed 16 bit PCM',
-     'PCM_S8': 'Signed 8 bit PCM'}
-
-    """
-    subtypes = _available_formats_helper(_snd.SFC_GET_FORMAT_SUBTYPE_COUNT,
-                                         _snd.SFC_GET_FORMAT_SUBTYPE)
-    return dict((subtype, name) for subtype, name in subtypes
-                if format is None or check_format(format, subtype))
-
-
-def check_format(format, subtype=None, endian=None):
-    """Check if the combination of format/subtype/endian is valid.
-
-    Examples
-    --------
-    >>> import soundfile as sf
-    >>> sf.check_format('WAV', 'PCM_24')
-    True
-    >>> sf.check_format('FLAC', 'VORBIS')
-    False
-
-    """
-    try:
-        return bool(_format_int(format, subtype, endian))
-    except (ValueError, TypeError):
-        return False
-
-
-def default_subtype(format):
-    """Return the default subtype for a given format.
-
-    Examples
-    --------
-    >>> import soundfile as sf
-    >>> sf.default_subtype('WAV')
-    'PCM_16'
-    >>> sf.default_subtype('MAT5')
-    'DOUBLE'
-
-    """
-    _check_format(format)
-    return _default_subtypes.get(format.upper())
-
-
-class SoundFile(object):
-    """A sound file.
-
-    For more documentation see the __init__() docstring (which is also
-    used for the online documentation (https://python-soundfile.readthedocs.io/).
-
-    """
-
-    def __init__(self, file, mode='r', samplerate=None, channels=None,
-                 subtype=None, endian=None, format=None, closefd=True,
-                 compression_level=None, bitrate_mode=None):
-        """Open a sound file.
-
-        If a file is opened with `mode` ``'r'`` (the default) or
-        ``'r+'``, no sample rate, channels or file format need to be
-        given because the information is obtained from the file. An
-        exception is the ``'RAW'`` data format, which always requires
-        these data points.
-
-        File formats consist of three case-insensitive strings:
-
-        * a *major format* which is by default obtained from the
-          extension of the file name (if known) and which can be
-          forced with the format argument (e.g. ``format='WAVEX'``).
-        * a *subtype*, e.g. ``'PCM_24'``. Most major formats have a
-          default subtype which is used if no subtype is specified.
-        * an *endian-ness*, which doesn't have to be specified at all in
-          most cases.
-
-        A `SoundFile` object is a *context manager*, which means
-        if used in a "with" statement, `close()` is automatically
-        called when reaching the end of the code block inside the "with"
-        statement.
-
-        Parameters
-        ----------
-        file : str or int or file-like object
-            The file to open.  This can be a file name, a file
-            descriptor or a Python file object (or a similar object with
-            the methods ``read()``/``readinto()``, ``write()``,
-            ``seek()`` and ``tell()``).
-        mode : {'r', 'r+', 'w', 'w+', 'x', 'x+'}, optional
-            Open mode.  Has to begin with one of these three characters:
-            ``'r'`` for reading, ``'w'`` for writing (truncates *file*)
-            or ``'x'`` for writing (raises an error if *file* already
-            exists).  Additionally, it may contain ``'+'`` to open
-            *file* for both reading and writing.
-            The character ``'b'`` for *binary mode* is implied because
-            all sound files have to be opened in this mode.
-            If *file* is a file descriptor or a file-like object,
-            ``'w'`` doesn't truncate and ``'x'`` doesn't raise an error.
-        samplerate : int
-            The sample rate of the file.  If `mode` contains ``'r'``,
-            this is obtained from the file (except for ``'RAW'`` files).
-        channels : int
-            The number of channels of the file.
-            If `mode` contains ``'r'``, this is obtained from the file
-            (except for ``'RAW'`` files).
-        subtype : str, sometimes optional
-            The subtype of the sound file.  If `mode` contains ``'r'``,
-            this is obtained from the file (except for ``'RAW'``
-            files), if not, the default value depends on the selected
-            `format` (see `default_subtype()`).
-            See `available_subtypes()` for all possible subtypes for
-            a given `format`.
-        endian : {'FILE', 'LITTLE', 'BIG', 'CPU'}, sometimes optional
-            The endian-ness of the sound file.  If `mode` contains
-            ``'r'``, this is obtained from the file (except for
-            ``'RAW'`` files), if not, the default value is ``'FILE'``,
-            which is correct in most cases.
-        format : str, sometimes optional
-            The major format of the sound file.  If `mode` contains
-            ``'r'``, this is obtained from the file (except for
-            ``'RAW'`` files), if not, the default value is determined
-            from the file extension.  See `available_formats()` for
-            all possible values.
-        closefd : bool, optional
-            Whether to close the file descriptor on `close()`. Only
-            applicable if the *file* argument is a file descriptor.
-        compression_level : float, optional
-            The compression level on 'write()'. The compression level
-            should be between 0.0 (minimum compression level) and 1.0
-            (highest compression level).
-            See `libsndfile document `__.
-        bitrate_mode : {'CONSTANT', 'AVERAGE', 'VARIABLE'}, optional
-            The bitrate mode on 'write()'. 
-            See `libsndfile document `__.
-
-        Examples
-        --------
-        >>> from soundfile import SoundFile
-
-        Open an existing file for reading:
-
-        >>> myfile = SoundFile('existing_file.wav')
-        >>> # do something with myfile
-        >>> myfile.close()
-
-        Create a new sound file for reading and writing using a with
-        statement:
-
-        >>> with SoundFile('new_file.wav', 'x+', 44100, 2) as myfile:
-        >>>     # do something with myfile
-        >>>     # ...
-        >>>     assert not myfile.closed
-        >>>     # myfile.close() is called automatically at the end
-        >>> assert myfile.closed
-
-        """
-        # resolve PathLike objects (see PEP519 for details):
-        # can be replaced with _os.fspath(file) for Python >= 3.6
-        file = file.__fspath__() if hasattr(file, '__fspath__') else file
-        self._name = file
-        if mode is None:
-            mode = getattr(file, 'mode', None)
-        mode_int = _check_mode(mode)
-        self._mode = mode
-        self._compression_level = compression_level
-        self._bitrate_mode = bitrate_mode
-        self._info = _create_info_struct(file, mode, samplerate, channels,
-                                         format, subtype, endian)
-        self._file = self._open(file, mode_int, closefd)
-        if set(mode).issuperset('r+') and self.seekable():
-            # Move write position to 0 (like in Python file objects)
-            self.seek(0)
-        _snd.sf_command(self._file, _snd.SFC_SET_CLIPPING, _ffi.NULL,
-                        _snd.SF_TRUE)
-        
-        # set compression setting
-        if self._compression_level is not None:
-            # needs to be called before set_bitrate_mode
-            self._set_compression_level(self._compression_level)
-            if self._bitrate_mode is not None:
-                self._set_bitrate_mode(self._bitrate_mode)
-
-    name = property(lambda self: self._name)
-    """The file name of the sound file."""
-    mode = property(lambda self: self._mode)
-    """The open mode the sound file was opened with."""
-    samplerate = property(lambda self: self._info.samplerate)
-    """The sample rate of the sound file."""
-    frames = property(lambda self: self._info.frames)
-    """The number of frames in the sound file."""
-    channels = property(lambda self: self._info.channels)
-    """The number of channels in the sound file."""
-    format = property(
-        lambda self: _format_str(self._info.format & _snd.SF_FORMAT_TYPEMASK))
-    """The major format of the sound file."""
-    subtype = property(
-        lambda self: _format_str(self._info.format & _snd.SF_FORMAT_SUBMASK))
-    """The subtype of data in the the sound file."""
-    endian = property(
-        lambda self: _format_str(self._info.format & _snd.SF_FORMAT_ENDMASK))
-    """The endian-ness of the data in the sound file."""
-    format_info = property(
-        lambda self: _format_info(self._info.format &
-                                  _snd.SF_FORMAT_TYPEMASK)[1])
-    """A description of the major format of the sound file."""
-    subtype_info = property(
-        lambda self: _format_info(self._info.format &
-                                  _snd.SF_FORMAT_SUBMASK)[1])
-    """A description of the subtype of the sound file."""
-    sections = property(lambda self: self._info.sections)
-    """The number of sections of the sound file."""
-    closed = property(lambda self: self._file is None)
-    """Whether the sound file is closed or not."""
-    _errorcode = property(lambda self: _snd.sf_error(self._file))
-    """A pending sndfile error code."""
-    compression_level = property(lambda self: self._compression_level)
-    """The compression level on 'write()'"""
-    bitrate_mode = property(lambda self: self._bitrate_mode)
-    """The bitrate mode on 'write()'"""
-
-    @property
-    def extra_info(self):
-        """Retrieve the log string generated when opening the file."""
-        info = _ffi.new("char[]", 2**14)
-        _snd.sf_command(self._file, _snd.SFC_GET_LOG_INFO,
-                        info, _ffi.sizeof(info))
-        return _ffi.string(info).decode('utf-8', 'replace')
-
-    # avoid confusion if something goes wrong before assigning self._file:
-    _file = None
-
-    def __repr__(self):
-        compression_setting = (", compression_level={0}".format(self.compression_level) 
-                               if self.compression_level is not None else "")
-        compression_setting += (", bitrate_mode='{0}'".format(self.bitrate_mode) 
-                                if self.bitrate_mode is not None else "")
-        return ("SoundFile({0.name!r}, mode={0.mode!r}, "
-                "samplerate={0.samplerate}, channels={0.channels}, "
-                "format={0.format!r}, subtype={0.subtype!r}, "
-                "endian={0.endian!r}{1})".format(self, compression_setting))
-
-    def __del__(self):
-        self.close()
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, *args):
-        self.close()
-
-    def __setattr__(self, name, value):
-        """Write text meta-data in the sound file through properties."""
-        if name in _str_types:
-            self._check_if_closed()
-            err = _snd.sf_set_string(self._file, _str_types[name],
-                                     value.encode())
-            _error_check(err)
-        else:
-            object.__setattr__(self, name, value)
-
-    def __getattr__(self, name):
-        """Read text meta-data in the sound file through properties."""
-        if name in _str_types:
-            self._check_if_closed()
-            data = _snd.sf_get_string(self._file, _str_types[name])
-            return _ffi.string(data).decode('utf-8', 'replace') if data else ""
-        else:
-            raise AttributeError(
-                "'SoundFile' object has no attribute {0!r}".format(name))
-
-    def __len__(self):
-        # Note: This is deprecated and will be removed at some point,
-        # see https://github.com/bastibe/python-soundfile/issues/199
-        return self._info.frames
-
-    def __bool__(self):
-        # Note: This is temporary until __len__ is removed, afterwards it
-        # can (and should) be removed without change of behavior
-        return True
-
-    def __nonzero__(self):
-        # Note: This is only for compatibility with Python 2 and it shall be
-        # removed at the same time as __bool__().
-        return self.__bool__()
-
-    def seekable(self):
-        """Return True if the file supports seeking."""
-        return self._info.seekable == _snd.SF_TRUE
-
-    def seek(self, frames, whence=SEEK_SET):
-        """Set the read/write position.
-
-        Parameters
-        ----------
-        frames : int
-            The frame index or offset to seek.
-        whence : {SEEK_SET, SEEK_CUR, SEEK_END}, optional
-            By default (``whence=SEEK_SET``), *frames* are counted from
-            the beginning of the file.
-            ``whence=SEEK_CUR`` seeks from the current position
-            (positive and negative values are allowed for *frames*).
-            ``whence=SEEK_END`` seeks from the end (use negative value
-            for *frames*).
-
-        Returns
-        -------
-        int
-            The new absolute read/write position in frames.
-
-        Examples
-        --------
-        >>> from soundfile import SoundFile, SEEK_END
-        >>> myfile = SoundFile('stereo_file.wav')
-
-        Seek to the beginning of the file:
-
-        >>> myfile.seek(0)
-        0
-
-        Seek to the end of the file:
-
-        >>> myfile.seek(0, SEEK_END)
-        44100  # this is the file length
-
-        """
-        self._check_if_closed()
-        position = _snd.sf_seek(self._file, frames, whence)
-        _error_check(self._errorcode)
-        return position
-
-    def tell(self):
-        """Return the current read/write position."""
-        return self.seek(0, SEEK_CUR)
-
-    def read(self, frames=-1, dtype='float64', always_2d=False,
-             fill_value=None, out=None):
-        """Read from the file and return data as NumPy array.
-
-        Reads the given number of frames in the given data format
-        starting at the current read/write position.  This advances the
-        read/write position by the same number of frames.
-        By default, all frames from the current read/write position to
-        the end of the file are returned.
-        Use `seek()` to move the current read/write position.
-
-        Parameters
-        ----------
-        frames : int, optional
-            The number of frames to read. If ``frames < 0``, the whole
-            rest of the file is read.
-        dtype : {'float64', 'float32', 'int32', 'int16'}, optional
-            Data type of the returned array, by default ``'float64'``.
-            Floating point audio data is typically in the range from
-            ``-1.0`` to ``1.0``. Integer data is in the range from
-            ``-2**15`` to ``2**15-1`` for ``'int16'`` and from
-            ``-2**31`` to ``2**31-1`` for ``'int32'``.
-
-            .. note:: Reading int values from a float file will *not*
-                scale the data to [-1.0, 1.0). If the file contains
-                ``np.array([42.6], dtype='float32')``, you will read
-                ``np.array([43], dtype='int32')`` for
-                ``dtype='int32'``.
-
-        Returns
-        -------
-        audiodata : `numpy.ndarray` or type(out)
-            A two-dimensional NumPy (frames x channels) array is
-            returned. If the sound file has only one channel, a
-            one-dimensional array is returned. Use ``always_2d=True``
-            to return a two-dimensional array anyway.
-
-            If *out* was specified, it is returned. If *out* has more
-            frames than available in the file (or if *frames* is
-            smaller than the length of *out*) and no *fill_value* is
-            given, then only a part of *out* is overwritten and a view
-            containing all valid frames is returned.
-
-        Other Parameters
-        ----------------
-        always_2d : bool, optional
-            By default, reading a mono sound file will return a
-            one-dimensional array. With ``always_2d=True``, audio data
-            is always returned as a two-dimensional array, even if the
-            audio file has only one channel.
-        fill_value : float, optional
-            If more frames are requested than available in the file,
-            the rest of the output is be filled with *fill_value*. If
-            *fill_value* is not specified, a smaller array is
-            returned.
-        out : `numpy.ndarray` or subclass, optional
-            If *out* is specified, the data is written into the given
-            array instead of creating a new array. In this case, the
-            arguments *dtype* and *always_2d* are silently ignored! If
-            *frames* is not given, it is obtained from the length of
-            *out*.
-
-        Examples
-        --------
-        >>> from soundfile import SoundFile
-        >>> myfile = SoundFile('stereo_file.wav')
-
-        Reading 3 frames from a stereo file:
-
-        >>> myfile.read(3)
-        array([[ 0.71329652,  0.06294799],
-               [-0.26450912, -0.38874483],
-               [ 0.67398441, -0.11516333]])
-        >>> myfile.close()
-
-        See Also
-        --------
-        buffer_read, .write
-
-        """
-        if out is None:
-            frames = self._check_frames(frames, fill_value)
-            out = self._create_empty_array(frames, always_2d, dtype)
-        else:
-            if frames < 0 or frames > len(out):
-                frames = len(out)
-        frames = self._array_io('read', out, frames)
-        if len(out) > frames:
-            if fill_value is None:
-                out = out[:frames]
-            else:
-                out[frames:] = fill_value
-        return out
-
-    def buffer_read(self, frames=-1, dtype=None):
-        """Read from the file and return data as buffer object.
-
-        Reads the given number of *frames* in the given data format
-        starting at the current read/write position.  This advances the
-        read/write position by the same number of frames.
-        By default, all frames from the current read/write position to
-        the end of the file are returned.
-        Use `seek()` to move the current read/write position.
-
-        Parameters
-        ----------
-        frames : int, optional
-            The number of frames to read. If ``frames < 0``, the whole
-            rest of the file is read.
-        dtype : {'float64', 'float32', 'int32', 'int16'}
-            Audio data will be converted to the given data type.
-
-        Returns
-        -------
-        buffer
-            A buffer containing the read data.
-
-        See Also
-        --------
-        buffer_read_into, .read, buffer_write
-
-        """
-        frames = self._check_frames(frames, fill_value=None)
-        ctype = self._check_dtype(dtype)
-        cdata = _ffi.new(ctype + '[]', frames * self.channels)
-        read_frames = self._cdata_io('read', cdata, ctype, frames)
-        assert read_frames == frames
-        return _ffi.buffer(cdata)
-
-    def buffer_read_into(self, buffer, dtype):
-        """Read from the file into a given buffer object.
-
-        Fills the given *buffer* with frames in the given data format
-        starting at the current read/write position (which can be
-        changed with `seek()`) until the buffer is full or the end
-        of the file is reached.  This advances the read/write position
-        by the number of frames that were read.
-
-        Parameters
-        ----------
-        buffer : writable buffer
-            Audio frames from the file are written to this buffer.
-        dtype : {'float64', 'float32', 'int32', 'int16'}
-            The data type of *buffer*.
-
-        Returns
-        -------
-        int
-            The number of frames that were read from the file.
-            This can be less than the size of *buffer*.
-            The rest of the buffer is not filled with meaningful data.
-
-        See Also
-        --------
-        buffer_read, .read
-
-        """
-        ctype = self._check_dtype(dtype)
-        cdata, frames = self._check_buffer(buffer, ctype)
-        frames = self._cdata_io('read', cdata, ctype, frames)
-        return frames
-
-    def write(self, data):
-        """Write audio data from a NumPy array to the file.
-
-        Writes a number of frames at the read/write position to the
-        file. This also advances the read/write position by the same
-        number of frames and enlarges the file if necessary.
-
-        Note that writing int values to a float file will *not* scale
-        the values to [-1.0, 1.0). If you write the value
-        ``np.array([42], dtype='int32')``, to a ``subtype='FLOAT'``
-        file, the file will then contain ``np.array([42.],
-        dtype='float32')``.
-
-        Parameters
-        ----------
-        data : array_like
-            The data to write. Usually two-dimensional (frames x
-            channels), but one-dimensional *data* can be used for mono
-            files. Only the data types ``'float64'``, ``'float32'``,
-            ``'int32'`` and ``'int16'`` are supported.
-
-            .. note:: The data type of *data* does **not** select the
-                  data type of the written file. Audio data will be
-                  converted to the given *subtype*. Writing int values
-                  to a float file will *not* scale the values to
-                  [-1.0, 1.0). If you write the value ``np.array([42],
-                  dtype='int32')``, to a ``subtype='FLOAT'`` file, the
-                  file will then contain ``np.array([42.],
-                  dtype='float32')``.
-
-        Examples
-        --------
-        >>> import numpy as np
-        >>> from soundfile import SoundFile
-        >>> myfile = SoundFile('stereo_file.wav')
-
-        Write 10 frames of random data to a new file:
-
-        >>> with SoundFile('stereo_file.wav', 'w', 44100, 2, 'PCM_24') as f:
-        >>>     f.write(np.random.randn(10, 2))
-
-        See Also
-        --------
-        buffer_write, .read
-
-        """
-        import numpy as np
-
-        # no copy is made if data has already the correct memory layout:
-        data = np.ascontiguousarray(data)
-        written = self._array_io('write', data, len(data))
-        assert written == len(data)
-        self._update_frames(written)
-
-    def buffer_write(self, data, dtype):
-        """Write audio data from a buffer/bytes object to the file.
-
-        Writes the contents of *data* to the file at the current
-        read/write position.
-        This also advances the read/write position by the number of
-        frames that were written and enlarges the file if necessary.
-
-        Parameters
-        ----------
-        data : buffer or bytes
-            A buffer or bytes object containing the audio data to be
-            written.
-        dtype : {'float64', 'float32', 'int32', 'int16'}
-            The data type of the audio data stored in *data*.
-
-        See Also
-        --------
-        .write, buffer_read
-
-        """
-        ctype = self._check_dtype(dtype)
-        cdata, frames = self._check_buffer(data, ctype)
-        written = self._cdata_io('write', cdata, ctype, frames)
-        assert written == frames
-        self._update_frames(written)
-
-    def blocks(self, blocksize=None, overlap=0, frames=-1, dtype='float64',
-               always_2d=False, fill_value=None, out=None):
-        """Return a generator for block-wise reading.
-
-        By default, the generator yields blocks of the given
-        *blocksize* (using a given *overlap*) until the end of the file
-        is reached; *frames* can be used to stop earlier.
-
-        Parameters
-        ----------
-        blocksize : int
-            The number of frames to read per block. Either this or *out*
-            must be given.
-        overlap : int, optional
-            The number of frames to rewind between each block.
-        frames : int, optional
-            The number of frames to read.
-            If ``frames < 0``, the file is read until the end.
-        dtype : {'float64', 'float32', 'int32', 'int16'}, optional
-            See `read()`.
-
-        Yields
-        ------
-        `numpy.ndarray` or type(out)
-            Blocks of audio data.
-            If *out* was given, and the requested frames are not an
-            integer multiple of the length of *out*, and no
-            *fill_value* was given, the last block will be a smaller
-            view into *out*.
-
-
-        Other Parameters
-        ----------------
-        always_2d, fill_value, out
-            See `read()`.
-        fill_value : float, optional
-            See `read()`.
-        out : `numpy.ndarray` or subclass, optional
-            If *out* is specified, the data is written into the given
-            array instead of creating a new array. In this case, the
-            arguments *dtype* and *always_2d* are silently ignored!
-
-        Examples
-        --------
-        >>> from soundfile import SoundFile
-        >>> with SoundFile('stereo_file.wav') as f:
-        >>>     for block in f.blocks(blocksize=1024):
-        >>>         pass  # do something with 'block'
-
-        """
-        import numpy as np
-
-        if 'r' not in self.mode and '+' not in self.mode:
-            raise SoundFileRuntimeError("blocks() is not allowed in write-only mode")
-
-        frames = self._check_frames(frames, fill_value)
-        if out is None:
-            if blocksize is None:
-                raise TypeError("One of {blocksize, out} must be specified")
-            out_size = blocksize if fill_value is not None else min(blocksize, frames)
-            out = self._create_empty_array(out_size, always_2d, dtype)
-            copy_out = True
-        else:
-            if blocksize is not None:
-                raise TypeError(
-                    "Only one of {blocksize, out} may be specified")
-            blocksize = len(out)
-            copy_out = False
-
-        overlap_memory = None
-        while frames > 0:
-            if overlap_memory is None:
-                output_offset = 0
-            else:
-                output_offset = len(overlap_memory)
-                out[:output_offset] = overlap_memory
-
-            toread = min(blocksize - output_offset, frames)
-            self.read(toread, dtype, always_2d, fill_value, out[output_offset:])
-
-            if overlap:
-                if overlap_memory is None:
-                    overlap_memory = np.copy(out[-overlap:])
-                else:
-                    overlap_memory[:] = out[-overlap:]
-
-            if blocksize > frames + overlap and fill_value is None:
-                block = out[:frames + overlap]
-            else:
-                block = out
-            yield np.copy(block) if copy_out else block
-            frames -= toread
-
-    def truncate(self, frames=None):
-        """Truncate the file to a given number of frames.
-
-        After this command, the read/write position will be at the new
-        end of the file.
-
-        Parameters
-        ----------
-        frames : int, optional
-            Only the data before *frames* is kept, the rest is deleted.
-            If not specified, the current read/write position is used.
-
-        """
-        if frames is None:
-            frames = self.tell()
-        err = _snd.sf_command(self._file, _snd.SFC_FILE_TRUNCATE,
-                              _ffi.new("sf_count_t*", frames),
-                              _ffi.sizeof("sf_count_t"))
-        if err:
-            # get the actual error code
-            err = _snd.sf_error(self._file)
-            raise LibsndfileError(err, "Error truncating the file")
-        self._info.frames = frames
-
-    def flush(self):
-        """Write unwritten data to the file system.
-
-        Data written with `write()` is not immediately written to
-        the file system but buffered in memory to be written at a later
-        time.  Calling `flush()` makes sure that all changes are
-        actually written to the file system.
-
-        This has no effect on files opened in read-only mode.
-
-        """
-        self._check_if_closed()
-        _snd.sf_write_sync(self._file)
-
-    def close(self):
-        """Close the file.  Can be called multiple times."""
-        if not self.closed:
-            # be sure to flush data to disk before closing the file
-            self.flush()
-            err = _snd.sf_close(self._file)
-            self._file = None
-            _error_check(err)
-
-    def _open(self, file, mode_int, closefd):
-        """Call the appropriate sf_open*() function from libsndfile."""
-        if isinstance(file, (_unicode, bytes)):
-            if _os.path.isfile(file):
-                if 'x' in self.mode:
-                    raise OSError("File exists: {0!r}".format(self.name))
-                elif set(self.mode).issuperset('w+'):
-                    # truncate the file, because SFM_RDWR doesn't:
-                    _os.close(_os.open(file, _os.O_WRONLY | _os.O_TRUNC))
-            openfunction = _snd.sf_open
-            if isinstance(file, _unicode):
-                if _sys.platform == 'win32':
-                    openfunction = _snd.sf_wchar_open
-                else:
-                    file = file.encode(_sys.getfilesystemencoding())
-            file_ptr = openfunction(file, mode_int, self._info)
-        elif isinstance(file, int):
-            file_ptr = _snd.sf_open_fd(file, mode_int, self._info, closefd)
-        elif _has_virtual_io_attrs(file, mode_int):
-            file_ptr = _snd.sf_open_virtual(self._init_virtual_io(file),
-                                            mode_int, self._info, _ffi.NULL)
-        else:
-            raise TypeError("Invalid file: {0!r}".format(self.name))
-        if file_ptr == _ffi.NULL:
-            # get the actual error code
-            err = _snd.sf_error(file_ptr)
-            raise LibsndfileError(err, prefix="Error opening {0!r}: ".format(self.name))
-        if mode_int == _snd.SFM_WRITE:
-            # Due to a bug in libsndfile version <= 1.0.25, frames != 0
-            # when opening a named pipe in SFM_WRITE mode.
-            # See http://github.com/erikd/libsndfile/issues/77.
-            self._info.frames = 0
-            # This is not necessary for "normal" files (because
-            # frames == 0 in this case), but it doesn't hurt, either.
-        return file_ptr
-
-    def _init_virtual_io(self, file):
-        """Initialize callback functions for sf_open_virtual()."""
-        @_ffi.callback("sf_vio_get_filelen")
-        def vio_get_filelen(user_data):
-            curr = file.tell()
-            file.seek(0, SEEK_END)
-            size = file.tell()
-            file.seek(curr, SEEK_SET)
-            return size
-
-        @_ffi.callback("sf_vio_seek")
-        def vio_seek(offset, whence, user_data):
-            file.seek(offset, whence)
-            return file.tell()
-
-        @_ffi.callback("sf_vio_read")
-        def vio_read(ptr, count, user_data):
-            # first try readinto(), if not available fall back to read()
-            try:
-                buf = _ffi.buffer(ptr, count)
-                data_read = file.readinto(buf)
-            except AttributeError:
-                data = file.read(count)
-                data_read = len(data)
-                buf = _ffi.buffer(ptr, data_read)
-                buf[0:data_read] = data
-            return data_read
-
-        @_ffi.callback("sf_vio_write")
-        def vio_write(ptr, count, user_data):
-            buf = _ffi.buffer(ptr, count)
-            data = buf[:]
-            written = file.write(data)
-            # write() returns None for file objects in Python <= 2.7:
-            if written is None:
-                written = count
-            return written
-
-        @_ffi.callback("sf_vio_tell")
-        def vio_tell(user_data):
-            return file.tell()
-
-        # Note: the callback functions must be kept alive!
-        self._virtual_io = {'get_filelen': vio_get_filelen,
-                            'seek': vio_seek,
-                            'read': vio_read,
-                            'write': vio_write,
-                            'tell': vio_tell}
-
-        return _ffi.new("SF_VIRTUAL_IO*", self._virtual_io)
-
-    def _getAttributeNames(self):
-        """Return all attributes used in __setattr__ and __getattr__.
-
-        This is useful for auto-completion (e.g. IPython).
-
-        """
-        return _str_types
-
-    def _check_if_closed(self):
-        """Check if the file is closed and raise an error if it is.
-
-        This should be used in every method that uses self._file.
-
-        """
-        if self.closed:
-            raise SoundFileRuntimeError("I/O operation on closed file")
-
-    def _check_frames(self, frames, fill_value):
-        """Reduce frames to no more than are available in the file."""
-        if self.seekable():
-            remaining_frames = self.frames - self.tell()
-            if frames < 0 or (frames > remaining_frames and
-                              fill_value is None):
-                frames = remaining_frames
-        elif frames < 0:
-            raise ValueError("frames must be specified for non-seekable files")
-        return frames
-
-    def _check_buffer(self, data, ctype):
-        """Convert buffer to cdata and check for valid size."""
-        assert ctype in _ffi_types.values()
-        if not isinstance(data, bytes):
-            data = _ffi.from_buffer(data)
-        frames, remainder = divmod(len(data),
-                                   self.channels * _ffi.sizeof(ctype))
-        if remainder:
-            raise ValueError("Data size must be a multiple of frame size")
-        return data, frames
-
-    def _create_empty_array(self, frames, always_2d, dtype):
-        """Create an empty array with appropriate shape."""
-        import numpy as np
-        if always_2d or self.channels > 1:
-            shape = frames, self.channels
-        else:
-            shape = frames,
-        return np.empty(shape, dtype, order='C')
-
-    def _check_dtype(self, dtype):
-        """Check if dtype string is valid and return ctype string."""
-        try:
-            return _ffi_types[dtype]
-        except KeyError:
-            raise ValueError("dtype must be one of {0!r} and not {1!r}".format(
-                sorted(_ffi_types.keys()), dtype))
-
-    def _array_io(self, action, array, frames):
-        """Check array and call low-level IO function."""
-        if array.ndim not in (1,2):
-            raise ValueError("Invalid shape: {0!r} ({1})".format(array.shape, "0 dimensions not supported" if array.ndim < 1 else "too many dimensions"))
-        array_channels = 1 if array.ndim == 1 else array.shape[1]
-        if array_channels != self.channels:
-            raise ValueError("Invalid shape: {0!r} (Expected {1} channels, got {2})".format(array.shape, self.channels, array_channels))
-        if not array.flags.c_contiguous:
-            raise ValueError("Data must be C-contiguous")
-        ctype = self._check_dtype(array.dtype.name)
-        assert array.dtype.itemsize == _ffi.sizeof(ctype)
-        cdata = _ffi.cast(ctype + '*', array.__array_interface__['data'][0])
-        return self._cdata_io(action, cdata, ctype, frames)
-
-    def _cdata_io(self, action, data, ctype, frames):
-        """Call one of libsndfile's read/write functions."""
-        assert ctype in _ffi_types.values()
-        self._check_if_closed()
-        if self.seekable():
-            curr = self.tell()
-        func = getattr(_snd, 'sf_' + action + 'f_' + ctype)
-        frames = func(self._file, data, frames)
-        _error_check(self._errorcode)
-        if self.seekable():
-            self.seek(curr + frames, SEEK_SET)  # Update read & write position
-        return frames
-
-    def _update_frames(self, written):
-        """Update self.frames after writing."""
-        if self.seekable():
-            curr = self.tell()
-            self._info.frames = self.seek(0, SEEK_END)
-            self.seek(curr, SEEK_SET)
-        else:
-            self._info.frames += written
-
-    def _prepare_read(self, start, stop, frames):
-        """Seek to start frame and calculate length."""
-        if start != 0 and not self.seekable():
-            raise ValueError("start is only allowed for seekable files")
-        if frames >= 0 and stop is not None:
-            raise TypeError("Only one of {frames, stop} may be used")
-
-        start, stop, _ = slice(start, stop).indices(self.frames)
-        if stop < start:
-            stop = start
-        if frames < 0:
-            frames = stop - start
-        if self.seekable():
-            self.seek(start, SEEK_SET)
-        return frames
-
-    def copy_metadata(self):
-        """Get all metadata present in this SoundFile
-
-        Returns
-        -------
-
-        metadata: dict[str, str]
-            A dict with all metadata. Possible keys are: 'title', 'copyright',
-            'software', 'artist', 'comment', 'date', 'album', 'license',
-            'tracknumber' and 'genre'.
-        """
-        strs = {}
-        for strtype, strid in _str_types.items():
-            data = _snd.sf_get_string(self._file, strid)
-            if data:
-                strs[strtype] = _ffi.string(data).decode('utf-8', 'replace')
-        return strs
-    
-    def _set_bitrate_mode(self, bitrate_mode):
-        """Call libsndfile's set bitrate mode function."""
-        assert bitrate_mode in _bitrate_modes
-
-        pointer_bitrate_mode = _ffi.new("int[1]")
-        pointer_bitrate_mode[0] = _bitrate_modes[bitrate_mode]
-        err = _snd.sf_command(self._file, _snd.SFC_SET_BITRATE_MODE, pointer_bitrate_mode, _ffi.sizeof(pointer_bitrate_mode))
-        if err != _snd.SF_TRUE:
-            err = _snd.sf_error(self._file)
-            raise LibsndfileError(err, f"Error set bitrate mode {bitrate_mode}")
-
-        
-    def _set_compression_level(self, compression_level):
-        """Call libsndfile's set compression level function."""
-        if not (0 <= compression_level <= 1):
-            raise ValueError("Compression level must be in range [0..1]")
-
-        pointer_compression_level = _ffi.new("double[1]")
-        pointer_compression_level[0] = compression_level
-        err = _snd.sf_command(self._file, _snd.SFC_SET_COMPRESSION_LEVEL, pointer_compression_level, _ffi.sizeof(pointer_compression_level))
-        if err != _snd.SF_TRUE:
-            err = _snd.sf_error(self._file)
-            raise LibsndfileError(err, f"Error set compression level {compression_level}")
-
-
-def _error_check(err, prefix=""):
-    """Raise LibsndfileError if there is an error."""
-    if err != 0:
-        raise LibsndfileError(err, prefix=prefix)
-
-
-def _format_int(format, subtype, endian):
-    """Return numeric ID for given format|subtype|endian combo."""
-    result = _check_format(format)
-    if subtype is None:
-        subtype = default_subtype(format)
-        if subtype is None:
-            raise TypeError(
-                "No default subtype for major format {0!r}".format(format))
-    elif not isinstance(subtype, (_unicode, str)):
-        raise TypeError("Invalid subtype: {0!r}".format(subtype))
-    try:
-        result |= _subtypes[subtype.upper()]
-    except KeyError:
-        raise ValueError("Unknown subtype: {0!r}".format(subtype))
-    if endian is None:
-        endian = 'FILE'
-    elif not isinstance(endian, (_unicode, str)):
-        raise TypeError("Invalid endian-ness: {0!r}".format(endian))
-    try:
-        result |= _endians[endian.upper()]
-    except KeyError:
-        raise ValueError("Unknown endian-ness: {0!r}".format(endian))
-
-    info = _ffi.new("SF_INFO*")
-    info.format = result
-    info.channels = 1
-    if _snd.sf_format_check(info) == _snd.SF_FALSE:
-        raise ValueError(
-            "Invalid combination of format, subtype and endian")
-    return result
-
-
-def _check_mode(mode):
-    """Check if mode is valid and return its integer representation."""
-    if not isinstance(mode, (_unicode, str)):
-        raise TypeError("Invalid mode: {0!r}".format(mode))
-    mode_set = set(mode)
-    if mode_set.difference('xrwb+') or len(mode) > len(mode_set):
-        raise ValueError("Invalid mode: {0!r}".format(mode))
-    if len(mode_set.intersection('xrw')) != 1:
-        raise ValueError("mode must contain exactly one of 'xrw'")
-
-    if '+' in mode_set:
-        mode_int = _snd.SFM_RDWR
-    elif 'r' in mode_set:
-        mode_int = _snd.SFM_READ
-    else:
-        mode_int = _snd.SFM_WRITE
-    return mode_int
-
-
-def _create_info_struct(file, mode, samplerate, channels,
-                        format, subtype, endian):
-    """Check arguments and create SF_INFO struct."""
-    original_format = format
-    if format is None:
-        format = _get_format_from_filename(file, mode)
-        assert isinstance(format, (_unicode, str))
-    else:
-        _check_format(format)
-
-    info = _ffi.new("SF_INFO*")
-    if 'r' not in mode or format.upper() == 'RAW':
-        if samplerate is None:
-            raise TypeError("samplerate must be specified")
-        info.samplerate = samplerate
-        if channels is None:
-            raise TypeError("channels must be specified")
-        info.channels = channels
-        info.format = _format_int(format, subtype, endian)
-    else:
-        if any(arg is not None for arg in (
-                samplerate, channels, original_format, subtype, endian)):
-            raise TypeError("Not allowed for existing files (except 'RAW'): "
-                            "samplerate, channels, format, subtype, endian")
-    return info
-
-
-def _get_format_from_filename(file, mode):
-    """Return a format string obtained from file (or file.name).
-
-    If file already exists (= read mode), an empty string is returned on
-    error.  If not, an exception is raised.
-    The return type will always be str or unicode (even if
-    file/file.name is a bytes object).
-
-    """
-    format = ''
-    file = getattr(file, 'name', file)
-    try:
-        # This raises an exception if file is not a (Unicode/byte) string:
-        format = _os.path.splitext(file)[-1][1:]
-        # Convert bytes to unicode (raises AttributeError on Python 3 str):
-        format = format.decode('utf-8', 'replace')
-    except Exception:
-        pass
-    if format.upper() not in _formats and 'r' not in mode:
-        raise TypeError("No format specified and unable to get format from "
-                        "file extension: {0!r}".format(file))
-    return format
-
-
-def _format_str(format_int):
-    """Return the string representation of a given numeric format."""
-    for dictionary in _formats, _subtypes, _endians:
-        for k, v in dictionary.items():
-            if v == format_int:
-                return k
-    else:
-        return 'n/a'
-
-
-def _format_info(format_int, format_flag=_snd.SFC_GET_FORMAT_INFO):
-    """Return the ID and short description of a given format."""
-    format_info = _ffi.new("SF_FORMAT_INFO*")
-    format_info.format = format_int
-    _snd.sf_command(_ffi.NULL, format_flag, format_info,
-                    _ffi.sizeof("SF_FORMAT_INFO"))
-    name = format_info.name
-    return (_format_str(format_info.format),
-            _ffi.string(name).decode('utf-8', 'replace') if name else "")
-
-
-def _available_formats_helper(count_flag, format_flag):
-    """Helper for available_formats() and available_subtypes()."""
-    count = _ffi.new("int*")
-    _snd.sf_command(_ffi.NULL, count_flag, count, _ffi.sizeof("int"))
-    for format_int in range(count[0]):
-        yield _format_info(format_int, format_flag)
-
-
-def _check_format(format_str):
-    """Check if `format_str` is valid and return format ID."""
-    if not isinstance(format_str, (_unicode, str)):
-        raise TypeError("Invalid format: {0!r}".format(format_str))
-    try:
-        format_int = _formats[format_str.upper()]
-    except KeyError:
-        raise ValueError("Unknown format: {0!r}".format(format_str))
-    return format_int
-
-
-def _has_virtual_io_attrs(file, mode_int):
-    """Check if file has all the necessary attributes for virtual IO."""
-    readonly = mode_int == _snd.SFM_READ
-    writeonly = mode_int == _snd.SFM_WRITE
-    return all([
-        hasattr(file, 'seek'),
-        hasattr(file, 'tell'),
-        hasattr(file, 'write') or readonly,
-        hasattr(file, 'read') or hasattr(file, 'readinto') or writeonly,
-    ])
-
-
-class SoundFileError(Exception):
-    """Base class for all soundfile-specific errors."""
-    pass
-
-class SoundFileRuntimeError(SoundFileError, RuntimeError):
-    """soundfile module runtime error.
-
-    Errors that used to be `RuntimeError`."""
-    pass
-
-class LibsndfileError(SoundFileRuntimeError):
-    """libsndfile errors.
-
-
-    Attributes
-    ----------
-    code
-        libsndfile internal error number.
-    """
-    def __init__(self, code, prefix=""):
-        SoundFileRuntimeError.__init__(self, code, prefix)
-        self.code = code
-        self.prefix = prefix
-
-    @property
-    def error_string(self):
-        """Raw libsndfile error message."""
-        if self.code:
-            err_str = _snd.sf_error_number(self.code)
-            return _ffi.string(err_str).decode('utf-8', 'replace')
-        else:
-            # Due to race conditions, if used concurrently, sf_error() may
-            # return 0 (= no error) even if an error has happened.
-            # See https://github.com/erikd/libsndfile/issues/610 for details.
-            return "(Garbled error message from libsndfile)"
-
-    def __str__(self):
-        return self.prefix + self.error_string
diff --git a/venv/lib/python3.10/site-packages/stringzilla.cpython-310-x86_64-linux-gnu.so b/venv/lib/python3.10/site-packages/stringzilla.cpython-310-x86_64-linux-gnu.so
deleted file mode 100644
index ce51888d5fa4f6ea0eb514dfe10b06e8f11d0ca5..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/stringzilla.cpython-310-x86_64-linux-gnu.so
+++ /dev/null
@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:57d013b7ca3360af9a9a242c223fd81168676cd8a870a16b35e2bd66e525f652
-size 842368
diff --git a/venv/lib/python3.10/site-packages/test/__init__.py b/venv/lib/python3.10/site-packages/test/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/venv/lib/python3.10/site-packages/test/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/test/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 0b020216823c4c833083c33336ae3edffd9bc9ab..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/test/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/test/__pycache__/api_test.cpython-310.pyc b/venv/lib/python3.10/site-packages/test/__pycache__/api_test.cpython-310.pyc
deleted file mode 100644
index f9a8b50bf10f1255b67f847e89f83407c4db5028..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/test/__pycache__/api_test.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/test/__pycache__/client_test.cpython-310.pyc b/venv/lib/python3.10/site-packages/test/__pycache__/client_test.cpython-310.pyc
deleted file mode 100644
index a797a743b4f7eb610521b510cb613244fc5fed02..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/test/__pycache__/client_test.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/test/__pycache__/response_examples.cpython-310.pyc b/venv/lib/python3.10/site-packages/test/__pycache__/response_examples.cpython-310.pyc
deleted file mode 100644
index cdbc00be4e1f2619d402f2acb4987bc52e921f35..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/test/__pycache__/response_examples.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/test/__pycache__/user_test.cpython-310.pyc b/venv/lib/python3.10/site-packages/test/__pycache__/user_test.cpython-310.pyc
deleted file mode 100644
index 8c6878f9c24239c692542d58c0353ee4eb7cdcf3..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/test/__pycache__/user_test.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/test/api_test.py b/venv/lib/python3.10/site-packages/test/api_test.py
deleted file mode 100644
index 90175918bdeb39bccd903b1f5f3a0424954aedd5..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/test/api_test.py
+++ /dev/null
@@ -1,507 +0,0 @@
-import responses
-import pytest
-import json
-from copy import deepcopy
-from matrix_client import client, api
-from matrix_client.errors import MatrixRequestError, MatrixError, MatrixHttpLibError
-from matrix_client import __version__ as lib_version
-from . import response_examples
-MATRIX_V2_API_PATH = "/_matrix/client/r0"
-
-
-class TestTagsApi:
-    cli = client.MatrixClient("http://example.com")
-    user_id = "@user:matrix.org"
-    room_id = "#foo:matrix.org"
-
-    @responses.activate
-    def test_get_user_tags(self):
-        tags_url = "http://example.com" \
-            "/_matrix/client/r0/user/@user:matrix.org/rooms/#foo:matrix.org/tags"
-        responses.add(responses.GET, tags_url, body='{}')
-        self.cli.api.get_user_tags(self.user_id, self.room_id)
-        req = responses.calls[0].request
-        assert req.url == tags_url
-        assert req.method == 'GET'
-
-    @responses.activate
-    def test_add_user_tags(self):
-        tags_url = "http://example.com" \
-            "/_matrix/client/r0/user/@user:matrix.org/rooms/#foo:matrix.org/tags/foo"
-        responses.add(responses.PUT, tags_url, body='{}')
-        self.cli.api.add_user_tag(self.user_id, self.room_id, "foo", body={"order": "5"})
-        req = responses.calls[0].request
-        assert req.url == tags_url
-        assert req.method == 'PUT'
-
-    @responses.activate
-    def test_remove_user_tags(self):
-        tags_url = "http://example.com" \
-            "/_matrix/client/r0/user/@user:matrix.org/rooms/#foo:matrix.org/tags/foo"
-        responses.add(responses.DELETE, tags_url, body='{}')
-        self.cli.api.remove_user_tag(self.user_id, self.room_id, "foo")
-        req = responses.calls[0].request
-        assert req.url == tags_url
-        assert req.method == 'DELETE'
-
-
-class TestAccountDataApi:
-    cli = client.MatrixClient("http://example.com")
-    user_id = "@user:matrix.org"
-    room_id = "#foo:matrix.org"
-
-    @responses.activate
-    def test_set_account_data(self):
-        account_data_url = "http://example.com" \
-            "/_matrix/client/r0/user/@user:matrix.org/account_data/foo"
-        responses.add(responses.PUT, account_data_url, body='{}')
-        self.cli.api.set_account_data(self.user_id, 'foo', {'bar': 1})
-        req = responses.calls[0].request
-        assert req.url == account_data_url
-        assert req.method == 'PUT'
-
-    @responses.activate
-    def test_set_room_account_data(self):
-        account_data_url = "http://example.com/_matrix/client/r0/user" \
-            "/@user:matrix.org/rooms/#foo:matrix.org/account_data/foo"
-        responses.add(responses.PUT, account_data_url, body='{}')
-        self.cli.api.set_room_account_data(self.user_id, self.room_id, 'foo', {'bar': 1})
-        req = responses.calls[0].request
-        assert req.url == account_data_url
-        assert req.method == 'PUT'
-
-
-class TestUnbanApi:
-    cli = client.MatrixClient("http://example.com")
-    user_id = "@user:matrix.org"
-    room_id = "#foo:matrix.org"
-
-    @responses.activate
-    def test_unban(self):
-        unban_url = "http://example.com" \
-                    "/_matrix/client/r0/rooms/#foo:matrix.org/unban"
-        responses.add(responses.POST, unban_url, body='{}')
-        self.cli.api.unban_user(self.room_id, self.user_id)
-        req = responses.calls[0].request
-        assert req.url == unban_url
-        assert req.method == 'POST'
-
-
-class TestDeviceApi:
-    cli = client.MatrixClient("http://example.com")
-    device_id = "QBUAZIFURK"
-    display_name = "test_name"
-    auth_body = {
-        "auth": {
-            "type": "example.type.foo",
-            "session": "xxxxx",
-            "example_credential": "verypoorsharedsecret"
-        }
-    }
-
-    @responses.activate
-    def test_get_devices(self):
-        get_devices_url = "http://example.com/_matrix/client/r0/devices"
-        responses.add(responses.GET, get_devices_url, body='{}')
-        self.cli.api.get_devices()
-        req = responses.calls[0].request
-        assert req.url == get_devices_url
-        assert req.method == 'GET'
-
-    @responses.activate
-    def test_get_device(self):
-        get_device_url = "http://example.com/_matrix/client/r0/devices/QBUAZIFURK"
-        responses.add(responses.GET, get_device_url, body='{}')
-        self.cli.api.get_device(self.device_id)
-        req = responses.calls[0].request
-        assert req.url == get_device_url
-        assert req.method == 'GET'
-
-    @responses.activate
-    def test_update_device_info(self):
-        update_url = "http://example.com/_matrix/client/r0/devices/QBUAZIFURK"
-        responses.add(responses.PUT, update_url, body='{}')
-        self.cli.api.update_device_info(self.device_id, self.display_name)
-        req = responses.calls[0].request
-        assert req.url == update_url
-        assert req.method == 'PUT'
-
-    @responses.activate
-    def test_delete_device(self):
-        delete_device_url = "http://example.com/_matrix/client/r0/devices/QBUAZIFURK"
-        responses.add(responses.DELETE, delete_device_url, body='{}')
-        # Test for 401 status code of User-Interactive Auth API
-        responses.add(responses.DELETE, delete_device_url, body='{}', status=401)
-        self.cli.api.delete_device(self.auth_body, self.device_id)
-        req = responses.calls[0].request
-        assert req.url == delete_device_url
-        assert req.method == 'DELETE'
-
-        with pytest.raises(MatrixRequestError):
-            self.cli.api.delete_device(self.auth_body, self.device_id)
-
-    @responses.activate
-    def test_delete_devices(self):
-        delete_devices_url = "http://example.com/_matrix/client/r0/delete_devices"
-        responses.add(responses.POST, delete_devices_url, body='{}')
-        # Test for 401 status code of User-Interactive Auth API
-        responses.add(responses.POST, delete_devices_url, body='{}', status=401)
-        self.cli.api.delete_devices(self.auth_body, [self.device_id])
-        req = responses.calls[0].request
-        assert req.url == delete_devices_url
-        assert req.method == 'POST'
-
-        with pytest.raises(MatrixRequestError):
-            self.cli.api.delete_devices(self.auth_body, [self.device_id])
-
-
-class TestKeysApi:
-    cli = client.MatrixClient("http://example.com")
-    user_id = "@alice:matrix.org"
-    device_id = "JLAFKJWSCS"
-    one_time_keys = {"curve25519:AAAAAQ": "/qyvZvwjiTxGdGU0RCguDCLeR+nmsb3FfNG3/Ve4vU8"}
-    device_keys = {
-        "user_id": "@alice:example.com",
-        "device_id": "JLAFKJWSCS",
-        "algorithms": [
-            "m.olm.curve25519-aes-sha256",
-            "m.megolm.v1.aes-sha"
-        ],
-        "keys": {
-            "curve25519:JLAFKJWSCS": "3C5BFWi2Y8MaVvjM8M22DBmh24PmgR0nPvJOIArzgyI",
-            "ed25519:JLAFKJWSCS": "lEuiRJBit0IG6nUf5pUzWTUEsRVVe/HJkoKuEww9ULI"
-        },
-        "signatures": {
-            "@alice:example.com": {
-                "ed25519:JLAFKJWSCS": ("dSO80A01XiigH3uBiDVx/EjzaoycHcjq9lfQX0uWsqxl2gi"
-                                       "MIiSPR8a4d291W1ihKJL/a+myXS367WT6NAIcBA")
-            }
-        }
-    }
-
-    @responses.activate
-    @pytest.mark.parametrize("args", [
-        {},
-        {'device_keys': device_keys},
-        {'one_time_keys': one_time_keys}
-    ])
-    def test_upload_keys(self, args):
-        upload_keys_url = "http://example.com/_matrix/client/r0/keys/upload"
-        responses.add(responses.POST, upload_keys_url, body='{}')
-        self.cli.api.upload_keys(**args)
-        req = responses.calls[0].request
-        assert req.url == upload_keys_url
-        assert req.method == 'POST'
-
-    @responses.activate
-    def test_query_keys(self):
-        query_user_keys_url = "http://example.com/_matrix/client/r0/keys/query"
-        responses.add(responses.POST, query_user_keys_url, body='{}')
-        self.cli.api.query_keys({self.user_id: self.device_id}, timeout=10)
-        req = responses.calls[0].request
-        assert req.url == query_user_keys_url
-        assert req.method == 'POST'
-
-    @responses.activate
-    def test_claim_keys(self):
-        claim_keys_url = "http://example.com/_matrix/client/r0/keys/claim"
-        responses.add(responses.POST, claim_keys_url, body='{}')
-        self.cli.api.claim_keys({self.user_id: {self.device_id: "algo"}}, timeout=1000)
-        req = responses.calls[0].request
-        assert req.url == claim_keys_url
-        assert req.method == 'POST'
-
-    @responses.activate
-    def test_key_changes(self):
-        key_changes_url = "http://example.com/_matrix/client/r0/keys/changes"
-        responses.add(responses.GET, key_changes_url, body='{}')
-        self.cli.api.key_changes('s72594_4483_1934', 's75689_5632_2435')
-        req = responses.calls[0].request
-        assert req.url.split('?')[0] == key_changes_url
-        assert req.method == 'GET'
-
-
-class TestSendToDeviceApi:
-    cli = client.MatrixClient("http://example.com")
-    user_id = "@alice:matrix.org"
-    device_id = "JLAFKJWSCS"
-
-    @responses.activate
-    def test_send_to_device(self):
-        txn_id = self.cli.api._make_txn_id()
-        send_to_device_url = \
-            "http://example.com/_matrix/client/r0/sendToDevice/m.new_device/" + txn_id
-        responses.add(responses.PUT, send_to_device_url, body='{}')
-        payload = {self.user_id: {self.device_id: {"test": 1}}}
-        self.cli.api.send_to_device("m.new_device", payload, txn_id)
-        req = responses.calls[0].request
-        assert req.url == send_to_device_url
-        assert req.method == 'PUT'
-
-
-class TestMainApi:
-    user_id = "@alice:matrix.org"
-    token = "Dp0YKRXwx0iWDhFj7lg3DVjwsWzGcUIgARljgyAip2JD8qd5dSaW" \
-        "cxowTKEFetPulfLijAhv8eOmUSScyGcWgZyNMRTBmoJ0RFc0HotPvTBZ" \
-        "U98yKRLtat7V43aCpFmK"
-    test_path = "/account/whoami"
-
-    @responses.activate
-    def test_send_token_header(self):
-        mapi = api.MatrixHttpApi("http://example.com", token=self.token)
-        responses.add(
-            responses.GET,
-            mapi._base_url+MATRIX_V2_API_PATH+self.test_path,
-            body='{"application/json": {"user_id": "%s"}}' % self.user_id
-        )
-        mapi._send("GET", self.test_path)
-        req = responses.calls[0].request
-        assert req.method == 'GET'
-        assert req.headers['Authorization'] == 'Bearer %s' % self.token
-
-    @responses.activate
-    def test_send_user_agent_header(self):
-        mapi = api.MatrixHttpApi("http://example.com")
-        responses.add(
-            responses.GET,
-            mapi._base_url+MATRIX_V2_API_PATH+self.test_path,
-            body='{"application/json": {"user_id": "%s"}}' % self.user_id
-        )
-        mapi._send("GET", self.test_path)
-        req = responses.calls[0].request
-        assert req.method == 'GET'
-        assert req.headers['User-Agent'] == 'matrix-python-sdk/%s' % lib_version
-
-    @responses.activate
-    def test_send_token_query(self):
-        mapi = api.MatrixHttpApi(
-            "http://example.com",
-            token=self.token,
-            use_authorization_header=False
-        )
-        responses.add(
-            responses.GET,
-            mapi._base_url+MATRIX_V2_API_PATH+self.test_path,
-            body='{"application/json": {"user_id": "%s"}}' % self.user_id
-        )
-        mapi._send("GET", self.test_path)
-        req = responses.calls[0].request
-        assert req.method == 'GET'
-        assert self.token in req.url
-
-    @responses.activate
-    def test_send_user_id(self):
-        mapi = api.MatrixHttpApi(
-            "http://example.com",
-            token=self.token,
-            identity=self.user_id
-        )
-        responses.add(
-            responses.GET,
-            mapi._base_url+MATRIX_V2_API_PATH+self.test_path,
-            body='{"application/json": {"user_id": "%s"}}' % self.user_id
-        )
-        mapi._send("GET", self.test_path)
-        req = responses.calls[0].request
-        assert "user_id" in req.url
-
-    @responses.activate
-    def test_send_unsup_method(self):
-        mapi = api.MatrixHttpApi("http://example.com")
-        with pytest.raises(MatrixError):
-            mapi._send("GOT", self.test_path)
-
-    @responses.activate
-    def test_send_request_error(self):
-        mapi = api.MatrixHttpApi("http://example.com")
-        with pytest.raises(MatrixHttpLibError):
-            mapi._send("GET", self.test_path)
-
-
-class TestMediaApi:
-    cli = client.MatrixClient("http://example.com")
-    user_id = "@alice:example.com"
-    mxcurl = "mxc://example.com/OonjUOmcuVpUnmOWKtzPmAFe"
-
-    @responses.activate
-    def test_media_download(self):
-        media_url = \
-            "http://example.com/_matrix/media/r0/download/" + self.mxcurl[6:]
-        with open('test/response_examples.py', 'rb') as fil:
-            responses.add(
-                responses.GET, media_url,
-                content_type='application/python',
-                body=fil.read(), status=200, stream=True
-            )
-        resp = self.cli.api.media_download(self.mxcurl, allow_remote=False)
-        resp.raw.decode_content = True
-        req = responses.calls[0].request
-        assert req.url.split('?')[0] == media_url
-        assert req.method == 'GET'
-
-    def test_media_download_wrong_url(self):
-        with pytest.raises(ValueError):
-            self.cli.api.media_download(self.mxcurl[6:])
-
-    @responses.activate
-    def test_get_thumbnail(self):
-        media_url = \
-            "http://example.com/_matrix/media/r0/thumbnail/" + self.mxcurl[6:]
-        with open('test/response_examples.py', 'rb') as fil:
-            responses.add(
-                responses.GET, media_url,
-                content_type='application/python',
-                body=fil.read(), status=200, stream=True
-            )
-        resp = self.cli.api.get_thumbnail(
-            self.mxcurl, 28, 28, allow_remote=False
-        )
-        resp.raw.decode_content = True
-        req = responses.calls[0].request
-        assert req.url.split('?')[0] == media_url
-        assert req.method == 'GET'
-
-    def test_get_thumbnail_wrong_url(self):
-        with pytest.raises(ValueError):
-            self.cli.api.get_thumbnail(self.mxcurl[6:], 28, 28)
-
-    def test_get_thumbnail_wrong_method(self):
-        with pytest.raises(ValueError):
-            self.cli.api.get_thumbnail(self.mxcurl, 28, 28, 'cut')
-
-    @responses.activate
-    def test_get_url_preview(self):
-        media_url = \
-            "http://example.com/_matrix/media/r0/preview_url"
-        preview_url = deepcopy(response_examples.example_preview_url)
-        responses.add(
-            responses.GET, media_url,
-            body=json.dumps(preview_url)
-        )
-        self.cli.api.get_url_preview("https://google.com/", 1510610716656)
-        req = responses.calls[0].request
-        assert req.url.split('?')[0] == media_url
-        assert req.method == 'GET'
-
-
-class TestRoomApi:
-    cli = client.MatrixClient("http://example.com")
-    user_id = "@user:matrix.org"
-    room_id = "#foo:matrix.org"
-
-    @responses.activate
-    def test_create_room_visibility_public(self):
-        create_room_url = "http://example.com" \
-            "/_matrix/client/r0/createRoom"
-        responses.add(
-            responses.POST,
-            create_room_url,
-            json='{"room_id": "!sefiuhWgwghwWgh:example.com"}'
-        )
-        self.cli.api.create_room(
-            name="test",
-            alias="#test:example.com",
-            is_public=True
-        )
-        req = responses.calls[0].request
-        assert req.url == create_room_url
-        assert req.method == 'POST'
-        j = json.loads(req.body)
-        assert j["room_alias_name"] == "#test:example.com"
-        assert j["visibility"] == "public"
-        assert j["name"] == "test"
-
-    @responses.activate
-    def test_create_room_visibility_private(self):
-        create_room_url = "http://example.com" \
-            "/_matrix/client/r0/createRoom"
-        responses.add(
-            responses.POST,
-            create_room_url,
-            json='{"room_id": "!sefiuhWgwghwWgh:example.com"}'
-        )
-        self.cli.api.create_room(
-            name="test",
-            alias="#test:example.com",
-            is_public=False
-        )
-        req = responses.calls[0].request
-        assert req.url == create_room_url
-        assert req.method == 'POST'
-        j = json.loads(req.body)
-        assert j["room_alias_name"] == "#test:example.com"
-        assert j["visibility"] == "private"
-        assert j["name"] == "test"
-
-    @responses.activate
-    def test_create_room_federate_true(self):
-        create_room_url = "http://example.com" \
-            "/_matrix/client/r0/createRoom"
-        responses.add(
-            responses.POST,
-            create_room_url,
-            json='{"room_id": "!sefiuhWgwghwWgh:example.com"}'
-        )
-        self.cli.api.create_room(
-            name="test2",
-            alias="#test2:example.com",
-            federate=True
-        )
-        req = responses.calls[0].request
-        assert req.url == create_room_url
-        assert req.method == 'POST'
-        j = json.loads(req.body)
-        assert j["creation_content"]["m.federate"]
-
-    @responses.activate
-    def test_create_room_federate_false(self):
-        create_room_url = "http://example.com" \
-            "/_matrix/client/r0/createRoom"
-        responses.add(
-            responses.POST,
-            create_room_url,
-            json='{"room_id": "!sefiuhWgwghwWgh:example.com"}'
-        )
-        self.cli.api.create_room(
-            name="test",
-            alias="#test:example.com",
-            federate=False
-        )
-        req = responses.calls[0].request
-        assert req.url == create_room_url
-        assert req.method == 'POST'
-        j = json.loads(req.body)
-        assert not j["creation_content"]["m.federate"]
-
-
-class TestWhoamiQuery:
-    user_id = "@alice:example.com"
-    token = "Dp0YKRXwx0iWDhFj7lg3DVjwsWzGcUIgARljgyAip2JD8qd5dSaW" \
-            "cxowTKEFetPulfLijAhv8eOmUSScyGcWgZyNMRTBmoJ0RFc0HotPvTBZ" \
-            "U98yKRLtat7V43aCpFmK"
-
-    @responses.activate
-    def test_whoami(self):
-        mapi = api.MatrixHttpApi("http://example.com", token=self.token)
-        whoami_url = "http://example.com/_matrix/client/r0/account/whoami"
-        responses.add(
-            responses.GET,
-            whoami_url,
-            body='{"user_id": "%s"}' % self.user_id
-        )
-        mapi.whoami()
-        req = responses.calls[0].request
-        assert req.method == 'GET'
-        assert whoami_url in req.url
-
-    @responses.activate
-    def test_whoami_unauth(self):
-        mapi = api.MatrixHttpApi("http://example.com")
-        whoami_url = "http://example.com/_matrix/client/r0/account/whoami"
-        responses.add(
-            responses.GET,
-            whoami_url,
-            body='{"user_id": "%s"}' % self.user_id
-        )
-        with pytest.raises(MatrixError):
-            mapi.whoami()
diff --git a/venv/lib/python3.10/site-packages/test/client_test.py b/venv/lib/python3.10/site-packages/test/client_test.py
deleted file mode 100644
index 406082522195352ad34a51cff6180a6ca4314abe..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/test/client_test.py
+++ /dev/null
@@ -1,573 +0,0 @@
-import pytest
-import responses
-import json
-from copy import deepcopy
-from matrix_client.client import MatrixClient, Room, User, CACHE
-from matrix_client.api import MATRIX_V2_API_PATH
-from . import response_examples
-try:
-    from urllib import quote
-except ImportError:
-    from urllib.parse import quote
-
-HOSTNAME = "http://example.com"
-
-
-def test_create_client():
-    MatrixClient("http://example.com")
-
-
-@responses.activate
-def test_create_client_with_token():
-    user_id = "@alice:example.com"
-    token = "Dp0YKRXwx0iWDhFj7lg3DVjwsWzGcUIgARljgyAip2JD8qd5dSaW" \
-            "cxowTKEFetPulfLijAhv8eOmUSScyGcWgZyNMRTBmoJ0RFc0HotPvTBZ" \
-            "U98yKRLtat7V43aCpFmK"
-    whoami_url = HOSTNAME+MATRIX_V2_API_PATH+"/account/whoami"
-    responses.add(
-        responses.GET,
-        whoami_url,
-        body='{"user_id": "%s"}' % user_id
-    )
-    sync_response = deepcopy(response_examples.example_sync)
-    response_body = json.dumps(sync_response)
-    sync_url = HOSTNAME + MATRIX_V2_API_PATH + "/sync"
-    responses.add(responses.GET, sync_url, body=response_body)
-    MatrixClient(HOSTNAME, token=token)
-    req = responses.calls[0].request
-    assert req.method == 'GET'
-    assert whoami_url in req.url
-
-
-def test_sync_token():
-    client = MatrixClient("http://example.com")
-    assert client.get_sync_token() is None
-    client.set_sync_token("FAKE_TOKEN")
-    assert client.get_sync_token() == "FAKE_TOKEN"
-
-
-def test__mkroom():
-    client = MatrixClient("http://example.com")
-
-    roomId = "!UcYsUzyxTGDxLBEvLz:matrix.org"
-    goodRoom = client._mkroom(roomId)
-
-    assert isinstance(goodRoom, Room)
-    assert goodRoom.room_id is roomId
-
-    with pytest.raises(ValueError):
-        client._mkroom("BAD_ROOM:matrix.org")
-        client._mkroom("!BAD_ROOMmatrix.org")
-        client._mkroom("!BAD_ROOM::matrix.org")
-
-
-def test_get_rooms():
-    client = MatrixClient("http://example.com")
-    rooms = client.get_rooms()
-    assert isinstance(rooms, dict)
-    assert len(rooms) == 0
-
-    client = MatrixClient("http://example.com")
-
-    client._mkroom("!abc:matrix.org")
-    client._mkroom("!def:matrix.org")
-    client._mkroom("!ghi:matrix.org")
-
-    rooms = client.get_rooms()
-    assert isinstance(rooms, dict)
-    assert len(rooms) == 3
-
-
-def test_bad_state_events():
-    client = MatrixClient("http://example.com")
-    room = client._mkroom("!abc:matrix.org")
-
-    ev = {
-        "tomato": False
-    }
-
-    room._process_state_event(ev)
-
-
-def test_state_event():
-    client = MatrixClient("http://example.com")
-    room = client._mkroom("!abc:matrix.org")
-
-    room.name = False
-    room.topic = False
-    room.aliases = False
-
-    ev = {
-        "type": "m.room.name",
-        "content": {},
-        "event_id": "$10000000000000AAAAA:matrix.org"
-    }
-
-    room._process_state_event(ev)
-    assert room.name is None
-
-    ev["content"]["name"] = "TestName"
-    room._process_state_event(ev)
-    assert room.name == "TestName"
-
-    ev["type"] = "m.room.topic"
-    room._process_state_event(ev)
-    assert room.topic is None
-
-    ev["content"]["topic"] = "TestTopic"
-    room._process_state_event(ev)
-    assert room.topic == "TestTopic"
-
-    ev["type"] = "m.room.aliases"
-    room._process_state_event(ev)
-    assert room.aliases is None
-
-    aliases = ["#foo:matrix.org", "#bar:matrix.org"]
-    ev["content"]["aliases"] = aliases
-    room._process_state_event(ev)
-    assert room.aliases is aliases
-
-    # test member join event
-    ev["type"] = "m.room.member"
-    ev["content"] = {'membership': 'join', 'displayname': 'stereo'}
-    ev["state_key"] = "@stereo:xxx.org"
-    room._process_state_event(ev)
-    assert len(room._members) == 1
-    assert room._members["@stereo:xxx.org"]
-    # test member leave event
-    ev["content"]['membership'] = 'leave'
-    room._process_state_event(ev)
-    assert len(room._members) == 0
-
-    # test join_rules
-    room.invite_only = False
-    ev["type"] = "m.room.join_rules"
-    ev["content"] = {"join_rule": "invite"}
-    room._process_state_event(ev)
-    assert room.invite_only
-
-    # test guest_access
-    room.guest_access = False
-    ev["type"] = "m.room.guest_access"
-    ev["content"] = {"guest_access": "can_join"}
-    room._process_state_event(ev)
-    assert room.guest_access
-
-    # test malformed event (check does not throw exception)
-    room.guest_access = False
-    ev["type"] = "m.room.guest_access"
-    ev["content"] = {}
-    room._process_state_event(ev)
-    assert not room.guest_access
-
-    # test encryption
-    room.encrypted = False
-    ev["type"] = "m.room.encryption"
-    ev["content"] = {"algorithm": "m.megolm.v1.aes-sha2"}
-    room._process_state_event(ev)
-    assert room.encrypted
-    # encrypted flag must not be cleared on configuration change
-    ev["content"] = {"algorithm": None}
-    room._process_state_event(ev)
-    assert room.encrypted
-
-
-def test_get_user():
-    client = MatrixClient("http://example.com")
-
-    assert isinstance(client.get_user("@foobar:matrix.org"), User)
-
-    with pytest.raises(ValueError):
-        client.get_user("badfoobar:matrix.org")
-        client.get_user("@badfoobarmatrix.org")
-        client.get_user("@badfoobar:::matrix.org")
-
-
-def test_get_download_url():
-    client = MatrixClient("http://example.com")
-    real_url = "http://example.com/_matrix/media/r0/download/foobar"
-    assert client.api.get_download_url("mxc://foobar") == real_url
-
-    with pytest.raises(ValueError):
-        client.api.get_download_url("http://foobar")
-
-
-def test_remove_listener():
-    def dummy_listener():
-        pass
-
-    client = MatrixClient("http://example.com")
-    handler = client.add_listener(dummy_listener)
-
-    found_listener = False
-    for listener in client.listeners:
-        if listener["uid"] == handler:
-            found_listener = True
-            break
-
-    assert found_listener, "listener was not added properly"
-
-    client.remove_listener(handler)
-    found_listener = False
-    for listener in client.listeners:
-        if listener["uid"] == handler:
-            found_listener = True
-            break
-
-    assert not found_listener, "listener was not removed properly"
-
-
-class TestClientRegister:
-    cli = MatrixClient(HOSTNAME)
-
-    @responses.activate
-    def test_register_as_guest(self):
-        cli = self.cli
-
-        def _sync(self):
-            self._sync_called = True
-        cli.__dict__[_sync.__name__] = _sync.__get__(cli, cli.__class__)
-        register_guest_url = HOSTNAME + MATRIX_V2_API_PATH + "/register"
-        response_body = json.dumps({
-            'access_token': 'EXAMPLE_ACCESS_TOKEN',
-            'device_id': 'guest_device',
-            'home_server': 'example.com',
-            'user_id': '@455:example.com'
-        })
-        responses.add(responses.POST, register_guest_url, body=response_body)
-        cli.register_as_guest()
-        assert cli.token == cli.api.token == 'EXAMPLE_ACCESS_TOKEN'
-        assert cli.hs == 'example.com'
-        assert cli.user_id == '@455:example.com'
-        assert cli._sync_called
-
-
-def test_get_rooms_display_name():
-
-    def add_members(api, room, num):
-        for i in range(num):
-            room._add_member('@frho%s:matrix.org' % i, 'ho%s' % i)
-
-    client = MatrixClient("http://example.com")
-    client.user_id = "@frho0:matrix.org"
-    room1 = client._mkroom("!abc:matrix.org")
-    add_members(client.api, room1, 1)
-    room2 = client._mkroom("!def:matrix.org")
-    add_members(client.api, room2, 2)
-    room3 = client._mkroom("!ghi:matrix.org")
-    add_members(client.api, room3, 3)
-    room4 = client._mkroom("!rfi:matrix.org")
-    add_members(client.api, room4, 30)
-
-    rooms = client.get_rooms()
-    assert len(rooms) == 4
-    assert room1.display_name == "Empty room"
-    assert room2.display_name == "ho1"
-    assert room3.display_name == "ho1 and ho2"
-    assert room4.display_name == "ho1 and 28 others"
-
-
-@responses.activate
-def test_presence_listener():
-    client = MatrixClient("http://example.com")
-    accumulator = []
-
-    def dummy_callback(event):
-        accumulator.append(event)
-    presence_events = [
-        {
-            "content": {
-                "avatar_url": "mxc://localhost:wefuiwegh8742w",
-                "currently_active": False,
-                "last_active_ago": 2478593,
-                "presence": "online",
-                "user_id": "@example:localhost"
-            },
-            "event_id": "$WLGTSEFSEF:localhost",
-            "type": "m.presence"
-        },
-        {
-            "content": {
-                "avatar_url": "mxc://localhost:weaugwe742w",
-                "currently_active": True,
-                "last_active_ago": 1478593,
-                "presence": "online",
-                "user_id": "@example2:localhost"
-            },
-            "event_id": "$CIGTXEFREF:localhost",
-            "type": "m.presence"
-        },
-        {
-            "content": {
-                "avatar_url": "mxc://localhost:wefudweg13742w",
-                "currently_active": False,
-                "last_active_ago": 24795,
-                "presence": "offline",
-                "user_id": "@example3:localhost"
-            },
-            "event_id": "$ZEGASEDSEF:localhost",
-            "type": "m.presence"
-        },
-    ]
-    sync_response = deepcopy(response_examples.example_sync)
-    sync_response["presence"]["events"] = presence_events
-    response_body = json.dumps(sync_response)
-    sync_url = HOSTNAME + MATRIX_V2_API_PATH + "/sync"
-
-    responses.add(responses.GET, sync_url, body=response_body)
-    callback_uid = client.add_presence_listener(dummy_callback)
-    client._sync()
-    assert accumulator == presence_events
-
-    responses.add(responses.GET, sync_url, body=response_body)
-    client.remove_presence_listener(callback_uid)
-    accumulator = []
-    client._sync()
-    assert accumulator == []
-
-
-@responses.activate
-def test_changing_user_power_levels():
-    client = MatrixClient(HOSTNAME)
-    room_id = "!UcYsUzyxTGDxLBEvLz:matrix.org"
-    room = client._mkroom(room_id)
-    PL_state_path = HOSTNAME + MATRIX_V2_API_PATH + \
-        "/rooms/" + quote(room_id) + "/state/m.room.power_levels"
-
-    # Code should first get current power_levels and then modify them
-    responses.add(responses.GET, PL_state_path,
-                  json=response_examples.example_pl_event["content"])
-    responses.add(responses.PUT, PL_state_path,
-                  json=response_examples.example_event_response)
-    # Removes user from user and adds user to to users list
-    assert room.modify_user_power_levels(users={"@example:localhost": None,
-                                                "@foobar:example.com": 49})
-
-    expected_request = deepcopy(response_examples.example_pl_event["content"])
-    del expected_request["users"]["@example:localhost"]
-    expected_request["users"]["@foobar:example.com"] = 49
-
-    assert json.loads(responses.calls[1].request.body) == expected_request
-
-
-@responses.activate
-def test_changing_default_power_level():
-    client = MatrixClient(HOSTNAME)
-    room_id = "!UcYsUzyxTGDxLBEvLz:matrix.org"
-    room = client._mkroom(room_id)
-    PL_state_path = HOSTNAME + MATRIX_V2_API_PATH + \
-        "/rooms/" + quote(room_id) + "/state/m.room.power_levels"
-
-    # Code should first get current power_levels and then modify them
-    responses.add(responses.GET, PL_state_path,
-                  json=response_examples.example_pl_event["content"])
-    responses.add(responses.PUT, PL_state_path,
-                  json=response_examples.example_event_response)
-    assert room.modify_user_power_levels(users_default=23)
-
-    expected_request = deepcopy(response_examples.example_pl_event["content"])
-    expected_request["users_default"] = 23
-
-    assert json.loads(responses.calls[1].request.body) == expected_request
-
-
-@responses.activate
-def test_changing_event_required_power_levels():
-    client = MatrixClient(HOSTNAME)
-    room_id = "!UcYsUzyxTGDxLBEvLz:matrix.org"
-    room = client._mkroom(room_id)
-    PL_state_path = HOSTNAME + MATRIX_V2_API_PATH + \
-        "/rooms/" + quote(room_id) + "/state/m.room.power_levels"
-
-    # Code should first get current power_levels and then modify them
-    responses.add(responses.GET, PL_state_path,
-                  json=response_examples.example_pl_event["content"])
-    responses.add(responses.PUT, PL_state_path,
-                  json=response_examples.example_event_response)
-    # Remove event from events and adds new controlled event
-    assert room.modify_required_power_levels(events={"m.room.name": None,
-                                                     "example.event": 51})
-
-    expected_request = deepcopy(response_examples.example_pl_event["content"])
-    del expected_request["events"]["m.room.name"]
-    expected_request["events"]["example.event"] = 51
-
-    assert json.loads(responses.calls[1].request.body) == expected_request
-
-
-@responses.activate
-def test_changing_other_required_power_levels():
-    client = MatrixClient(HOSTNAME)
-    room_id = "!UcYsUzyxTGDxLBEvLz:matrix.org"
-    room = client._mkroom(room_id)
-    PL_state_path = HOSTNAME + MATRIX_V2_API_PATH + \
-        "/rooms/" + quote(room_id) + "/state/m.room.power_levels"
-
-    # Code should first get current power_levels and then modify them
-    responses.add(responses.GET, PL_state_path,
-                  json=response_examples.example_pl_event["content"])
-    responses.add(responses.PUT, PL_state_path,
-                  json=response_examples.example_event_response)
-    # Remove event from events and adds new controlled event
-    assert room.modify_required_power_levels(kick=53, redact=2,
-                                             state_default=None)
-
-    expected_request = deepcopy(response_examples.example_pl_event["content"])
-    expected_request["kick"] = 53
-    expected_request["redact"] = 2
-    del expected_request["state_default"]
-
-    assert json.loads(responses.calls[1].request.body) == expected_request
-
-
-@responses.activate
-def test_cache():
-    m_none = MatrixClient("http://example.com", cache_level=CACHE.NONE)
-    m_some = MatrixClient("http://example.com", cache_level=CACHE.SOME)
-    m_all = MatrixClient("http://example.com", cache_level=CACHE.ALL)
-    sync_url = HOSTNAME + MATRIX_V2_API_PATH + "/sync"
-    room_id = "!726s6s6q:example.com"
-    room_name = "The FooBar"
-    sync_response = deepcopy(response_examples.example_sync)
-
-    with pytest.raises(ValueError):
-        MatrixClient("http://example.com", cache_level=1)
-        MatrixClient("http://example.com", cache_level=5)
-        MatrixClient("http://example.com", cache_level=0.5)
-        MatrixClient("http://example.com", cache_level=-5)
-        MatrixClient("http://example.com", cache_level="foo")
-        MatrixClient("http://example.com", cache_level=0.0)
-
-    sync_response["rooms"]["join"][room_id]["state"]["events"].append(
-        {
-            "sender": "@alice:example.com",
-            "type": "m.room.name",
-            "state_key": "",
-            "content": {"name": room_name},
-        }
-    )
-
-    responses.add(responses.GET, sync_url, json.dumps(sync_response))
-    m_none._sync()
-    responses.add(responses.GET, sync_url, json.dumps(sync_response))
-    m_some._sync()
-    responses.add(responses.GET, sync_url, json.dumps(sync_response))
-    m_all._sync()
-
-    assert m_none.rooms[room_id].name is None
-    assert m_some.rooms[room_id].name == room_name
-    assert m_all.rooms[room_id].name == room_name
-
-    assert m_none.rooms[room_id]._members == m_some.rooms[room_id]._members == {}
-    assert len(m_all.rooms[room_id]._members) == 2
-    assert m_all.rooms[room_id]._members["@alice:example.com"]
-
-
-@responses.activate
-def test_room_join_rules():
-    client = MatrixClient(HOSTNAME)
-    room_id = "!UcYsUzyxTGDxLBEvLz:matrix.org"
-    room = client._mkroom(room_id)
-    assert room.invite_only is None
-    join_rules_state_path = HOSTNAME + MATRIX_V2_API_PATH + \
-        "/rooms/" + quote(room_id) + "/state/m.room.join_rules"
-
-    responses.add(responses.PUT, join_rules_state_path,
-                  json=response_examples.example_event_response)
-
-    assert room.set_invite_only(True)
-    assert room.invite_only
-
-
-@responses.activate
-def test_room_guest_access():
-    client = MatrixClient(HOSTNAME)
-    room_id = "!UcYsUzyxTGDxLBEvLz:matrix.org"
-    room = client._mkroom(room_id)
-    assert room.guest_access is None
-    guest_access_state_path = HOSTNAME + MATRIX_V2_API_PATH + \
-        "/rooms/" + quote(room_id) + "/state/m.room.guest_access"
-
-    responses.add(responses.PUT, guest_access_state_path,
-                  json=response_examples.example_event_response)
-
-    assert room.set_guest_access(True)
-    assert room.guest_access
-
-
-@responses.activate
-def test_enable_encryption():
-    pytest.importorskip('olm')
-    client = MatrixClient(HOSTNAME, encryption=True)
-
-    login_path = HOSTNAME + MATRIX_V2_API_PATH + "/login"
-    responses.add(responses.POST, login_path,
-                  json=response_examples.example_success_login_response)
-
-    upload_path = HOSTNAME + MATRIX_V2_API_PATH + '/keys/upload'
-    responses.add(responses.POST, upload_path, body='{"one_time_key_counts": {}}')
-
-    client.login("@example:localhost", "password", sync=False)
-
-    assert client.olm_device
-
-
-@responses.activate
-def test_enable_encryption_in_room():
-    pytest.importorskip('olm')
-    client = MatrixClient(HOSTNAME)
-    room_id = "!UcYsUzyxTGDxLBEvLz:matrix.org"
-    room = client._mkroom(room_id)
-    assert not room.encrypted
-    encryption_state_path = HOSTNAME + MATRIX_V2_API_PATH + \
-        "/rooms/" + quote(room_id) + "/state/m.room.encryption"
-
-    responses.add(responses.PUT, encryption_state_path,
-                  json=response_examples.example_event_response)
-
-    assert room.enable_encryption()
-    assert room.encrypted
-
-
-@responses.activate
-def test_detect_encryption_state():
-    pytest.importorskip('olm')
-    client = MatrixClient(HOSTNAME, encryption=True)
-    room_id = "!UcYsUzyxTGDxLBEvLz:matrix.org"
-
-    encryption_state_path = HOSTNAME + MATRIX_V2_API_PATH + \
-        "/rooms/" + quote(room_id) + "/state/m.room.encryption"
-    responses.add(responses.GET, encryption_state_path,
-                  json={"algorithm": "m.megolm.v1.aes-sha2"})
-    responses.add(responses.GET, encryption_state_path,
-                  json={}, status=404)
-
-    room = client._mkroom(room_id)
-    assert room.encrypted
-
-    room = client._mkroom(room_id)
-    assert not room.encrypted
-
-
-@responses.activate
-def test_one_time_keys_sync():
-    pytest.importorskip('olm')
-    client = MatrixClient(HOSTNAME, encryption=True)
-    sync_url = HOSTNAME + MATRIX_V2_API_PATH + "/sync"
-    sync_response = deepcopy(response_examples.example_sync)
-    payload = {'dummy': 1}
-    sync_response["device_one_time_keys_count"] = payload
-    sync_response['rooms']['join'] = {}
-
-    class DummyDevice:
-
-        def update_one_time_key_counts(self, payload):
-            self.payload = payload
-
-    device = DummyDevice()
-    client.olm_device = device
-
-    responses.add(responses.GET, sync_url, json=sync_response)
-
-    client._sync()
-    assert device.payload == payload
diff --git a/venv/lib/python3.10/site-packages/test/response_examples.py b/venv/lib/python3.10/site-packages/test/response_examples.py
deleted file mode 100644
index 2d45aa8612698a0da45fe6f7d8d72760288e8f7d..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/test/response_examples.py
+++ /dev/null
@@ -1,186 +0,0 @@
-example_sync = {
-    "next_batch": "s72595_4483_1934",
-    "presence": {
-        "events": [
-            {
-                "sender": "@alice:example.com",
-                "type": "m.presence",
-                "content": {
-                    "presence": "online"
-                }
-            }
-        ]
-    },
-    "account_data": {
-        "events": [
-            {
-                "type": "org.example.custom.config",
-                "content": {
-                    "custom_config_key": "custom_config_value"
-                }
-            }
-        ]
-    },
-    "rooms": {
-        "join": {
-            "!726s6s6q:example.com": {
-                "state": {
-                    "events": [
-                        {
-                            "sender": "@alice:example.com",
-                            "type": "m.room.member",
-                            "state_key": "@alice:example.com",
-                            "content": {
-                                "membership": "join"
-                            },
-                            "origin_server_ts": 1417731086795,
-                            "event_id": "$66697273743031:example.com"
-                        }
-                    ]
-                },
-                "timeline": {
-                    "events": [
-                        {
-                            "sender": "@bob:example.com",
-                            "type": "m.room.member",
-                            "state_key": "@bob:example.com",
-                            "content": {
-                                "membership": "join"
-                            },
-                            "prev_content": {
-                                "membership": "invite"
-                            },
-                            "origin_server_ts": 1417731086795,
-                            "event_id": "$7365636s6r6432:example.com"
-                        },
-                        {
-                            "sender": "@alice:example.com",
-                            "type": "m.room.message",
-                            "age": 124524,
-                            "txn_id": "1234",
-                            "content": {
-                                "body": "I am a fish",
-                                "msgtype": "m.text"
-                            },
-                            "origin_server_ts": 1417731086797,
-                            "event_id": "$74686972643033:example.com"
-                        }
-                    ],
-                    "limited": True,
-                    "prev_batch": "t34-23535_0_0"
-                },
-                "ephemeral": {
-                    "events": [
-                        {
-                            "type": "m.typing",
-                            "content": {
-                                "user_ids": [
-                                    "@alice:example.com"
-                                ]
-                            }
-                        }
-                    ]
-                },
-                "account_data": {
-                    "events": [
-                        {
-                            "type": "m.tag",
-                            "content": {
-                                "tags": {
-                                    "work": {
-                                        "order": 1
-                                    }
-                                }
-                            }
-                        },
-                        {
-                            "type": "org.example.custom.room.config",
-                            "content": {
-                                "custom_config_key": "custom_config_value"
-                            }
-                        }
-                    ]
-                }
-            }
-        },
-        "invite": {
-            "!696r7674:example.com": {
-                "invite_state": {
-                    "events": [
-                        {
-                            "sender": "@alice:example.com",
-                            "type": "m.room.name",
-                            "state_key": "",
-                            "content": {
-                                "name": "My Room Name"
-                            }
-                        },
-                        {
-                            "sender": "@alice:example.com",
-                            "type": "m.room.member",
-                            "state_key": "@bob:example.com",
-                            "content": {
-                                "membership": "invite"
-                            }
-                        }
-                    ]
-                }
-            }
-        },
-        "leave": {}
-    }
-}
-
-example_pl_event = {
-    "age": 242352,
-    "content": {
-        "ban": 50,
-        "events": {
-            "m.room.name": 100,
-            "m.room.power_levels": 100
-        },
-        "events_default": 0,
-        "invite": 50,
-        "kick": 50,
-        "redact": 50,
-        "state_default": 50,
-        "users": {
-            "@example:localhost": 100
-        },
-        "users_default": 0
-    },
-    "event_id": "$WLGTSEFSEF:localhost",
-    "origin_server_ts": 1431961217939,
-    "room_id": "!Cuyf34gef24t:localhost",
-    "sender": "@example:localhost",
-    "state_key": "",
-    "type": "m.room.power_levels"
-}
-
-example_event_response = {
-    "event_id": "YUwRidLecu"
-}
-
-example_key_upload_response = {
-    "one_time_key_counts": {
-        "curve25519": 10,
-        "signed_curve25519": 20
-    }
-}
-
-example_success_login_response = {
-    "user_id": "@example:localhost",
-    "access_token": "abc123",
-    "home_server": "matrix.org",
-    "device_id": "GHTYAJCE"
-}
-
-example_preview_url = {
-    "matrix:image:size": 102400,
-    "og:description": "This is a really cool blog post from matrix.org",
-    "og:image": "mxc://example.com/ascERGshawAWawugaAcauga",
-    "og:image:height": 48,
-    "og:image:type": "image/png",
-    "og:image:width": 48,
-    "og:title": "Matrix Blog Post"
-}
diff --git a/venv/lib/python3.10/site-packages/test/user_test.py b/venv/lib/python3.10/site-packages/test/user_test.py
deleted file mode 100644
index db5bae820c2331294d3627d80cae678e64478e7d..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/test/user_test.py
+++ /dev/null
@@ -1,52 +0,0 @@
-import pytest
-import responses
-
-from matrix_client.api import MATRIX_V2_API_PATH
-from matrix_client.client import MatrixClient
-from matrix_client.user import User
-
-HOSTNAME = "http://localhost"
-
-
-class TestUser:
-    cli = MatrixClient(HOSTNAME)
-    user_id = "@test:localhost"
-    room_id = "!test:localhost"
-
-    @pytest.fixture()
-    def user(self):
-        return User(self.cli.api, self.user_id)
-
-    @pytest.fixture()
-    def room(self):
-        return self.cli._mkroom(self.room_id)
-
-    @responses.activate
-    def test_get_display_name(self, user, room):
-        displayname_url = HOSTNAME + MATRIX_V2_API_PATH + \
-            "/profile/{}/displayname".format(user.user_id)
-        displayname = 'test'
-        room_displayname = 'room_test'
-
-        # No displayname
-        assert user.get_display_name(room) == user.user_id
-        responses.add(responses.GET, displayname_url, json={})
-        assert user.get_display_name() == user.user_id
-        assert len(responses.calls) == 1
-
-        # Get global displayname
-        responses.replace(responses.GET, displayname_url,
-                          json={"displayname": displayname})
-        assert user.get_display_name() == displayname
-        assert len(responses.calls) == 2
-
-        # Global displayname already present
-        assert user.get_display_name() == displayname
-        # No new request
-        assert len(responses.calls) == 2
-
-        # Per-room displayname
-        room.members_displaynames[user.user_id] = room_displayname
-        assert user.get_display_name(room) == room_displayname
-        # No new request
-        assert len(responses.calls) == 2
diff --git a/venv/lib/python3.10/site-packages/threadpoolctl.py b/venv/lib/python3.10/site-packages/threadpoolctl.py
deleted file mode 100644
index aea33509dc5512c0e845d387e01577c22f027676..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/threadpoolctl.py
+++ /dev/null
@@ -1,1292 +0,0 @@
-"""threadpoolctl
-
-This module provides utilities to introspect native libraries that relies on
-thread pools (notably BLAS and OpenMP implementations) and dynamically set the
-maximal number of threads they can use.
-"""
-
-# License: BSD 3-Clause
-
-# The code to introspect dynamically loaded libraries on POSIX systems is
-# adapted from code by Intel developer @anton-malakhov available at
-# https://github.com/IntelPython/smp (Copyright (c) 2017, Intel Corporation)
-# and also published under the BSD 3-Clause license
-import os
-import re
-import sys
-import ctypes
-import itertools
-import textwrap
-from typing import final
-import warnings
-from ctypes.util import find_library
-from abc import ABC, abstractmethod
-from functools import lru_cache
-from contextlib import ContextDecorator
-
-__version__ = "3.6.0"
-__all__ = [
-    "threadpool_limits",
-    "threadpool_info",
-    "ThreadpoolController",
-    "LibController",
-    "register",
-]
-
-
-# One can get runtime errors or even segfaults due to multiple OpenMP libraries
-# loaded simultaneously which can happen easily in Python when importing and
-# using compiled extensions built with different compilers and therefore
-# different OpenMP runtimes in the same program. In particular libiomp (used by
-# Intel ICC) and libomp used by clang/llvm tend to crash. This can happen for
-# instance when calling BLAS inside a prange. Setting the following environment
-# variable allows multiple OpenMP libraries to be loaded. It should not degrade
-# performances since we manually take care of potential over-subscription
-# performance issues, in sections of the code where nested OpenMP loops can
-# happen, by dynamically reconfiguring the inner OpenMP runtime to temporarily
-# disable it while under the scope of the outer OpenMP parallel section.
-os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "True")
-
-# Structure to cast the info on dynamically loaded library. See
-# https://linux.die.net/man/3/dl_iterate_phdr for more details.
-_SYSTEM_UINT = ctypes.c_uint64 if sys.maxsize > 2**32 else ctypes.c_uint32
-_SYSTEM_UINT_HALF = ctypes.c_uint32 if sys.maxsize > 2**32 else ctypes.c_uint16
-
-
-class _dl_phdr_info(ctypes.Structure):
-    _fields_ = [
-        ("dlpi_addr", _SYSTEM_UINT),  # Base address of object
-        ("dlpi_name", ctypes.c_char_p),  # path to the library
-        ("dlpi_phdr", ctypes.c_void_p),  # pointer on dlpi_headers
-        ("dlpi_phnum", _SYSTEM_UINT_HALF),  # number of elements in dlpi_phdr
-    ]
-
-
-# The RTLD_NOLOAD flag for loading shared libraries is not defined on Windows.
-try:
-    _RTLD_NOLOAD = os.RTLD_NOLOAD
-except AttributeError:
-    _RTLD_NOLOAD = ctypes.DEFAULT_MODE
-
-
-class LibController(ABC):
-    """Abstract base class for the individual library controllers
-
-    A library controller must expose the following class attributes:
-        - user_api : str
-            Usually the name of the library or generic specification the library
-            implements, e.g. "blas" is a specification with different implementations.
-        - internal_api : str
-            Usually the name of the library or concrete implementation of some
-            specification, e.g. "openblas" is an implementation of the "blas"
-            specification.
-        - filename_prefixes : tuple
-            Possible prefixes of the shared library's filename that allow to
-            identify the library. e.g. "libopenblas" for libopenblas.so.
-
-    and implement the following methods: `get_num_threads`, `set_num_threads` and
-    `get_version`.
-
-    Threadpoolctl loops through all the loaded shared libraries and tries to match
-    the filename of each library with the `filename_prefixes`. If a match is found, a
-    controller is instantiated and a handler to the library is stored in the `dynlib`
-    attribute as a `ctypes.CDLL` object. It can be used to access the necessary symbols
-    of the shared library to implement the above methods.
-
-    The following information will be exposed in the info dictionary:
-      - user_api : standardized API, if any, or a copy of internal_api.
-      - internal_api : implementation-specific API.
-      - num_threads : the current thread limit.
-      - prefix : prefix of the shared library's filename.
-      - filepath : path to the loaded shared library.
-      - version : version of the library (if available).
-
-    In addition, each library controller may expose internal API specific entries. They
-    must be set as attributes in the `set_additional_attributes` method.
-    """
-
-    @final
-    def __init__(self, *, filepath=None, prefix=None, parent=None):
-        """This is not meant to be overriden by subclasses."""
-        self.parent = parent
-        self.prefix = prefix
-        self.filepath = filepath
-        self.dynlib = ctypes.CDLL(filepath, mode=_RTLD_NOLOAD)
-        self._symbol_prefix, self._symbol_suffix = self._find_affixes()
-        self.version = self.get_version()
-        self.set_additional_attributes()
-
-    def info(self):
-        """Return relevant info wrapped in a dict"""
-        hidden_attrs = ("dynlib", "parent", "_symbol_prefix", "_symbol_suffix")
-        return {
-            "user_api": self.user_api,
-            "internal_api": self.internal_api,
-            "num_threads": self.num_threads,
-            **{k: v for k, v in vars(self).items() if k not in hidden_attrs},
-        }
-
-    def set_additional_attributes(self):
-        """Set additional attributes meant to be exposed in the info dict"""
-
-    @property
-    def num_threads(self):
-        """Exposes the current thread limit as a dynamic property
-
-        This is not meant to be used or overriden by subclasses.
-        """
-        return self.get_num_threads()
-
-    @abstractmethod
-    def get_num_threads(self):
-        """Return the maximum number of threads available to use"""
-
-    @abstractmethod
-    def set_num_threads(self, num_threads):
-        """Set the maximum number of threads to use"""
-
-    @abstractmethod
-    def get_version(self):
-        """Return the version of the shared library"""
-
-    def _find_affixes(self):
-        """Return the affixes for the symbols of the shared library"""
-        return "", ""
-
-    def _get_symbol(self, name):
-        """Return the symbol of the shared library accounding for the affixes"""
-        return getattr(
-            self.dynlib, f"{self._symbol_prefix}{name}{self._symbol_suffix}", None
-        )
-
-
-class OpenBLASController(LibController):
-    """Controller class for OpenBLAS"""
-
-    user_api = "blas"
-    internal_api = "openblas"
-    filename_prefixes = ("libopenblas", "libblas", "libscipy_openblas")
-
-    _symbol_prefixes = ("", "scipy_")
-    _symbol_suffixes = ("", "64_", "_64")
-
-    # All variations of "openblas_get_num_threads", accounting for the affixes
-    check_symbols = tuple(
-        f"{prefix}openblas_get_num_threads{suffix}"
-        for prefix, suffix in itertools.product(_symbol_prefixes, _symbol_suffixes)
-    )
-
-    def _find_affixes(self):
-        for prefix, suffix in itertools.product(
-            self._symbol_prefixes, self._symbol_suffixes
-        ):
-            if hasattr(self.dynlib, f"{prefix}openblas_get_num_threads{suffix}"):
-                return prefix, suffix
-
-    def set_additional_attributes(self):
-        self.threading_layer = self._get_threading_layer()
-        self.architecture = self._get_architecture()
-
-    def get_num_threads(self):
-        get_num_threads_func = self._get_symbol("openblas_get_num_threads")
-        if get_num_threads_func is not None:
-            return get_num_threads_func()
-        return None
-
-    def set_num_threads(self, num_threads):
-        set_num_threads_func = self._get_symbol("openblas_set_num_threads")
-        if set_num_threads_func is not None:
-            return set_num_threads_func(num_threads)
-        return None
-
-    def get_version(self):
-        # None means OpenBLAS is not loaded or version < 0.3.4, since OpenBLAS
-        # did not expose its version before that.
-        get_version_func = self._get_symbol("openblas_get_config")
-        if get_version_func is not None:
-            get_version_func.restype = ctypes.c_char_p
-            config = get_version_func().split()
-            if config[0] == b"OpenBLAS":
-                return config[1].decode("utf-8")
-            return None
-        return None
-
-    def _get_threading_layer(self):
-        """Return the threading layer of OpenBLAS"""
-        get_threading_layer_func = self._get_symbol("openblas_get_parallel")
-        if get_threading_layer_func is not None:
-            threading_layer = get_threading_layer_func()
-            if threading_layer == 2:
-                return "openmp"
-            elif threading_layer == 1:
-                return "pthreads"
-            return "disabled"
-        return "unknown"
-
-    def _get_architecture(self):
-        """Return the architecture detected by OpenBLAS"""
-        get_architecture_func = self._get_symbol("openblas_get_corename")
-        if get_architecture_func is not None:
-            get_architecture_func.restype = ctypes.c_char_p
-            return get_architecture_func().decode("utf-8")
-        return None
-
-
-class BLISController(LibController):
-    """Controller class for BLIS"""
-
-    user_api = "blas"
-    internal_api = "blis"
-    filename_prefixes = ("libblis", "libblas")
-    check_symbols = (
-        "bli_thread_get_num_threads",
-        "bli_thread_set_num_threads",
-        "bli_info_get_version_str",
-        "bli_info_get_enable_openmp",
-        "bli_info_get_enable_pthreads",
-        "bli_arch_query_id",
-        "bli_arch_string",
-    )
-
-    def set_additional_attributes(self):
-        self.threading_layer = self._get_threading_layer()
-        self.architecture = self._get_architecture()
-
-    def get_num_threads(self):
-        get_func = getattr(self.dynlib, "bli_thread_get_num_threads", lambda: None)
-        num_threads = get_func()
-        # by default BLIS is single-threaded and get_num_threads
-        # returns -1. We map it to 1 for consistency with other libraries.
-        return 1 if num_threads == -1 else num_threads
-
-    def set_num_threads(self, num_threads):
-        set_func = getattr(
-            self.dynlib, "bli_thread_set_num_threads", lambda num_threads: None
-        )
-        return set_func(num_threads)
-
-    def get_version(self):
-        get_version_ = getattr(self.dynlib, "bli_info_get_version_str", None)
-        if get_version_ is None:
-            return None
-
-        get_version_.restype = ctypes.c_char_p
-        return get_version_().decode("utf-8")
-
-    def _get_threading_layer(self):
-        """Return the threading layer of BLIS"""
-        if getattr(self.dynlib, "bli_info_get_enable_openmp", lambda: False)():
-            return "openmp"
-        elif getattr(self.dynlib, "bli_info_get_enable_pthreads", lambda: False)():
-            return "pthreads"
-        return "disabled"
-
-    def _get_architecture(self):
-        """Return the architecture detected by BLIS"""
-        bli_arch_query_id = getattr(self.dynlib, "bli_arch_query_id", None)
-        bli_arch_string = getattr(self.dynlib, "bli_arch_string", None)
-        if bli_arch_query_id is None or bli_arch_string is None:
-            return None
-
-        # the true restype should be BLIS' arch_t (enum) but int should work
-        # for us:
-        bli_arch_query_id.restype = ctypes.c_int
-        bli_arch_string.restype = ctypes.c_char_p
-        return bli_arch_string(bli_arch_query_id()).decode("utf-8")
-
-
-class FlexiBLASController(LibController):
-    """Controller class for FlexiBLAS"""
-
-    user_api = "blas"
-    internal_api = "flexiblas"
-    filename_prefixes = ("libflexiblas",)
-    check_symbols = (
-        "flexiblas_get_num_threads",
-        "flexiblas_set_num_threads",
-        "flexiblas_get_version",
-        "flexiblas_list",
-        "flexiblas_list_loaded",
-        "flexiblas_current_backend",
-    )
-
-    @property
-    def loaded_backends(self):
-        return self._get_backend_list(loaded=True)
-
-    @property
-    def current_backend(self):
-        return self._get_current_backend()
-
-    def info(self):
-        """Return relevant info wrapped in a dict"""
-        # We override the info method because the loaded and current backends
-        # are dynamic properties
-        exposed_attrs = super().info()
-        exposed_attrs["loaded_backends"] = self.loaded_backends
-        exposed_attrs["current_backend"] = self.current_backend
-
-        return exposed_attrs
-
-    def set_additional_attributes(self):
-        self.available_backends = self._get_backend_list(loaded=False)
-
-    def get_num_threads(self):
-        get_func = getattr(self.dynlib, "flexiblas_get_num_threads", lambda: None)
-        num_threads = get_func()
-        # by default BLIS is single-threaded and get_num_threads
-        # returns -1. We map it to 1 for consistency with other libraries.
-        return 1 if num_threads == -1 else num_threads
-
-    def set_num_threads(self, num_threads):
-        set_func = getattr(
-            self.dynlib, "flexiblas_set_num_threads", lambda num_threads: None
-        )
-        return set_func(num_threads)
-
-    def get_version(self):
-        get_version_ = getattr(self.dynlib, "flexiblas_get_version", None)
-        if get_version_ is None:
-            return None
-
-        major = ctypes.c_int()
-        minor = ctypes.c_int()
-        patch = ctypes.c_int()
-        get_version_(ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch))
-        return f"{major.value}.{minor.value}.{patch.value}"
-
-    def _get_backend_list(self, loaded=False):
-        """Return the list of available backends for FlexiBLAS.
-
-        If loaded is False, return the list of available backends from the FlexiBLAS
-        configuration. If loaded is True, return the list of actually loaded backends.
-        """
-        func_name = f"flexiblas_list{'_loaded' if loaded else ''}"
-        get_backend_list_ = getattr(self.dynlib, func_name, None)
-        if get_backend_list_ is None:
-            return None
-
-        n_backends = get_backend_list_(None, 0, 0)
-
-        backends = []
-        for i in range(n_backends):
-            backend_name = ctypes.create_string_buffer(1024)
-            get_backend_list_(backend_name, 1024, i)
-            if backend_name.value.decode("utf-8") != "__FALLBACK__":
-                # We don't know when to expect __FALLBACK__ but it is not a real
-                # backend and does not show up when running flexiblas list.
-                backends.append(backend_name.value.decode("utf-8"))
-        return backends
-
-    def _get_current_backend(self):
-        """Return the backend of FlexiBLAS"""
-        get_backend_ = getattr(self.dynlib, "flexiblas_current_backend", None)
-        if get_backend_ is None:
-            return None
-
-        backend = ctypes.create_string_buffer(1024)
-        get_backend_(backend, ctypes.sizeof(backend))
-        return backend.value.decode("utf-8")
-
-    def switch_backend(self, backend):
-        """Switch the backend of FlexiBLAS
-
-        Parameters
-        ----------
-        backend : str
-            The name or the path to the shared library of the backend to switch to. If
-            the backend is not already loaded, it will be loaded first.
-        """
-        if backend not in self.loaded_backends:
-            if backend in self.available_backends:
-                load_func = getattr(self.dynlib, "flexiblas_load_backend", lambda _: -1)
-            else:  # assume backend is a path to a shared library
-                load_func = getattr(
-                    self.dynlib, "flexiblas_load_backend_library", lambda _: -1
-                )
-            res = load_func(str(backend).encode("utf-8"))
-            if res == -1:
-                raise RuntimeError(
-                    f"Failed to load backend {backend!r}. It must either be the name of"
-                    " a backend available in the FlexiBLAS configuration "
-                    f"{self.available_backends} or the path to a valid shared library."
-                )
-
-            # Trigger a new search of loaded shared libraries since loading a new
-            # backend caused a dlopen.
-            self.parent._load_libraries()
-
-        switch_func = getattr(self.dynlib, "flexiblas_switch", lambda _: -1)
-        idx = self.loaded_backends.index(backend)
-        res = switch_func(idx)
-        if res == -1:
-            raise RuntimeError(f"Failed to switch to backend {backend!r}.")
-
-
-class MKLController(LibController):
-    """Controller class for MKL"""
-
-    user_api = "blas"
-    internal_api = "mkl"
-    filename_prefixes = ("libmkl_rt", "mkl_rt", "libblas")
-    check_symbols = (
-        "MKL_Get_Max_Threads",
-        "MKL_Set_Num_Threads",
-        "MKL_Get_Version_String",
-        "MKL_Set_Threading_Layer",
-    )
-
-    def set_additional_attributes(self):
-        self.threading_layer = self._get_threading_layer()
-
-    def get_num_threads(self):
-        get_func = getattr(self.dynlib, "MKL_Get_Max_Threads", lambda: None)
-        return get_func()
-
-    def set_num_threads(self, num_threads):
-        set_func = getattr(self.dynlib, "MKL_Set_Num_Threads", lambda num_threads: None)
-        return set_func(num_threads)
-
-    def get_version(self):
-        if not hasattr(self.dynlib, "MKL_Get_Version_String"):
-            return None
-
-        res = ctypes.create_string_buffer(200)
-        self.dynlib.MKL_Get_Version_String(res, 200)
-
-        version = res.value.decode("utf-8")
-        group = re.search(r"Version ([^ ]+) ", version)
-        if group is not None:
-            version = group.groups()[0]
-        return version.strip()
-
-    def _get_threading_layer(self):
-        """Return the threading layer of MKL"""
-        # The function mkl_set_threading_layer returns the current threading
-        # layer. Calling it with an invalid threading layer allows us to safely
-        # get the threading layer
-        set_threading_layer = getattr(
-            self.dynlib, "MKL_Set_Threading_Layer", lambda layer: -1
-        )
-        layer_map = {
-            0: "intel",
-            1: "sequential",
-            2: "pgi",
-            3: "gnu",
-            4: "tbb",
-            -1: "not specified",
-        }
-        return layer_map[set_threading_layer(-1)]
-
-
-class OpenMPController(LibController):
-    """Controller class for OpenMP"""
-
-    user_api = "openmp"
-    internal_api = "openmp"
-    filename_prefixes = ("libiomp", "libgomp", "libomp", "vcomp")
-    check_symbols = (
-        "omp_get_max_threads",
-        "omp_get_num_threads",
-    )
-
-    def get_num_threads(self):
-        get_func = getattr(self.dynlib, "omp_get_max_threads", lambda: None)
-        return get_func()
-
-    def set_num_threads(self, num_threads):
-        set_func = getattr(self.dynlib, "omp_set_num_threads", lambda num_threads: None)
-        return set_func(num_threads)
-
-    def get_version(self):
-        # There is no way to get the version number programmatically in OpenMP.
-        return None
-
-
-# Controllers for the libraries that we'll look for in the loaded libraries.
-# Third party libraries can register their own controllers.
-_ALL_CONTROLLERS = [
-    OpenBLASController,
-    BLISController,
-    MKLController,
-    OpenMPController,
-    FlexiBLASController,
-]
-
-# Helpers for the doc and test names
-_ALL_USER_APIS = list(set(lib.user_api for lib in _ALL_CONTROLLERS))
-_ALL_INTERNAL_APIS = [lib.internal_api for lib in _ALL_CONTROLLERS]
-_ALL_PREFIXES = list(
-    set(prefix for lib in _ALL_CONTROLLERS for prefix in lib.filename_prefixes)
-)
-_ALL_BLAS_LIBRARIES = [
-    lib.internal_api for lib in _ALL_CONTROLLERS if lib.user_api == "blas"
-]
-_ALL_OPENMP_LIBRARIES = OpenMPController.filename_prefixes
-
-
-def register(controller):
-    """Register a new controller"""
-    _ALL_CONTROLLERS.append(controller)
-    _ALL_USER_APIS.append(controller.user_api)
-    _ALL_INTERNAL_APIS.append(controller.internal_api)
-    _ALL_PREFIXES.extend(controller.filename_prefixes)
-
-
-def _format_docstring(*args, **kwargs):
-    def decorator(o):
-        if o.__doc__ is not None:
-            o.__doc__ = o.__doc__.format(*args, **kwargs)
-        return o
-
-    return decorator
-
-
-@lru_cache(maxsize=10000)
-def _realpath(filepath):
-    """Small caching wrapper around os.path.realpath to limit system calls"""
-    return os.path.realpath(filepath)
-
-
-@_format_docstring(USER_APIS=list(_ALL_USER_APIS), INTERNAL_APIS=_ALL_INTERNAL_APIS)
-def threadpool_info():
-    """Return the maximal number of threads for each detected library.
-
-    Return a list with all the supported libraries that have been found. Each
-    library is represented by a dict with the following information:
-
-      - "user_api" : user API. Possible values are {USER_APIS}.
-      - "internal_api": internal API. Possible values are {INTERNAL_APIS}.
-      - "prefix" : filename prefix of the specific implementation.
-      - "filepath": path to the loaded library.
-      - "version": version of the library (if available).
-      - "num_threads": the current thread limit.
-
-    In addition, each library may contain internal_api specific entries.
-    """
-    return ThreadpoolController().info()
-
-
-class _ThreadpoolLimiter:
-    """The guts of ThreadpoolController.limit
-
-    Refer to the docstring of ThreadpoolController.limit for more details.
-
-    It will only act on the library controllers held by the provided `controller`.
-    Using the default constructor sets the limits right away such that it can be used as
-    a callable. Setting the limits can be delayed by using the `wrap` class method such
-    that it can be used as a decorator.
-    """
-
-    def __init__(self, controller, *, limits=None, user_api=None):
-        self._controller = controller
-        self._limits, self._user_api, self._prefixes = self._check_params(
-            limits, user_api
-        )
-        self._original_info = self._controller.info()
-        self._set_threadpool_limits()
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, type, value, traceback):
-        self.restore_original_limits()
-
-    @classmethod
-    def wrap(cls, controller, *, limits=None, user_api=None):
-        """Return an instance of this class that can be used as a decorator"""
-        return _ThreadpoolLimiterDecorator(
-            controller=controller, limits=limits, user_api=user_api
-        )
-
-    def restore_original_limits(self):
-        """Set the limits back to their original values"""
-        for lib_controller, original_info in zip(
-            self._controller.lib_controllers, self._original_info
-        ):
-            lib_controller.set_num_threads(original_info["num_threads"])
-
-    # Alias of `restore_original_limits` for backward compatibility
-    unregister = restore_original_limits
-
-    def get_original_num_threads(self):
-        """Original num_threads from before calling threadpool_limits
-
-        Return a dict `{user_api: num_threads}`.
-        """
-        num_threads = {}
-        warning_apis = []
-
-        for user_api in self._user_api:
-            limits = [
-                lib_info["num_threads"]
-                for lib_info in self._original_info
-                if lib_info["user_api"] == user_api
-            ]
-            limits = set(limits)
-            n_limits = len(limits)
-
-            if n_limits == 1:
-                limit = limits.pop()
-            elif n_limits == 0:
-                limit = None
-            else:
-                limit = min(limits)
-                warning_apis.append(user_api)
-
-            num_threads[user_api] = limit
-
-        if warning_apis:
-            warnings.warn(
-                "Multiple value possible for following user apis: "
-                + ", ".join(warning_apis)
-                + ". Returning the minimum."
-            )
-
-        return num_threads
-
-    def _check_params(self, limits, user_api):
-        """Suitable values for the _limits, _user_api and _prefixes attributes"""
-
-        if isinstance(limits, str) and limits == "sequential_blas_under_openmp":
-            (
-                limits,
-                user_api,
-            ) = self._controller._get_params_for_sequential_blas_under_openmp().values()
-
-        if limits is None or isinstance(limits, int):
-            if user_api is None:
-                user_api = _ALL_USER_APIS
-            elif user_api in _ALL_USER_APIS:
-                user_api = [user_api]
-            else:
-                raise ValueError(
-                    f"user_api must be either in {_ALL_USER_APIS} or None. Got "
-                    f"{user_api} instead."
-                )
-
-            if limits is not None:
-                limits = {api: limits for api in user_api}
-            prefixes = []
-        else:
-            if isinstance(limits, list):
-                # This should be a list of dicts of library info, for
-                # compatibility with the result from threadpool_info.
-                limits = {
-                    lib_info["prefix"]: lib_info["num_threads"] for lib_info in limits
-                }
-            elif isinstance(limits, ThreadpoolController):
-                # To set the limits from the library controllers of a
-                # ThreadpoolController object.
-                limits = {
-                    lib_controller.prefix: lib_controller.num_threads
-                    for lib_controller in limits.lib_controllers
-                }
-
-            if not isinstance(limits, dict):
-                raise TypeError(
-                    "limits must either be an int, a list, a dict, or "
-                    f"'sequential_blas_under_openmp'. Got {type(limits)} instead"
-                )
-
-            # With a dictionary, can set both specific limit for given
-            # libraries and global limit for user_api. Fetch each separately.
-            prefixes = [prefix for prefix in limits if prefix in _ALL_PREFIXES]
-            user_api = [api for api in limits if api in _ALL_USER_APIS]
-
-        return limits, user_api, prefixes
-
-    def _set_threadpool_limits(self):
-        """Change the maximal number of threads in selected thread pools.
-
-        Return a list with all the supported libraries that have been found
-        matching `self._prefixes` and `self._user_api`.
-        """
-        if self._limits is None:
-            return
-
-        for lib_controller in self._controller.lib_controllers:
-            # self._limits is a dict {key: num_threads} where key is either
-            # a prefix or a user_api. If a library matches both, the limit
-            # corresponding to the prefix is chosen.
-            if lib_controller.prefix in self._limits:
-                num_threads = self._limits[lib_controller.prefix]
-            elif lib_controller.user_api in self._limits:
-                num_threads = self._limits[lib_controller.user_api]
-            else:
-                continue
-
-            if num_threads is not None:
-                lib_controller.set_num_threads(num_threads)
-
-
-class _ThreadpoolLimiterDecorator(_ThreadpoolLimiter, ContextDecorator):
-    """Same as _ThreadpoolLimiter but to be used as a decorator"""
-
-    def __init__(self, controller, *, limits=None, user_api=None):
-        self._limits, self._user_api, self._prefixes = self._check_params(
-            limits, user_api
-        )
-        self._controller = controller
-
-    def __enter__(self):
-        # we need to set the limits here and not in the __init__ because we want the
-        # limits to be set when calling the decorated function, not when creating the
-        # decorator.
-        self._original_info = self._controller.info()
-        self._set_threadpool_limits()
-        return self
-
-
-@_format_docstring(
-    USER_APIS=", ".join(f'"{api}"' for api in _ALL_USER_APIS),
-    BLAS_LIBS=", ".join(_ALL_BLAS_LIBRARIES),
-    OPENMP_LIBS=", ".join(_ALL_OPENMP_LIBRARIES),
-)
-class threadpool_limits(_ThreadpoolLimiter):
-    """Change the maximal number of threads that can be used in thread pools.
-
-    This object can be used either as a callable (the construction of this object
-    limits the number of threads), as a context manager in a `with` block to
-    automatically restore the original state of the controlled libraries when exiting
-    the block, or as a decorator through its `wrap` method.
-
-    Set the maximal number of threads that can be used in thread pools used in
-    the supported libraries to `limit`. This function works for libraries that
-    are already loaded in the interpreter and can be changed dynamically.
-
-    This effect is global and impacts the whole Python process. There is no thread level
-    isolation as these libraries do not offer thread-local APIs to configure the number
-    of threads to use in nested parallel calls.
-
-    Parameters
-    ----------
-    limits : int, dict, 'sequential_blas_under_openmp' or None (default=None)
-        The maximal number of threads that can be used in thread pools
-
-        - If int, sets the maximum number of threads to `limits` for each
-          library selected by `user_api`.
-
-        - If it is a dictionary `{{key: max_threads}}`, this function sets a
-          custom maximum number of threads for each `key` which can be either a
-          `user_api` or a `prefix` for a specific library.
-
-        - If 'sequential_blas_under_openmp', it will chose the appropriate `limits`
-          and `user_api` parameters for the specific use case of sequential BLAS
-          calls within an OpenMP parallel region. The `user_api` parameter is
-          ignored.
-
-        - If None, this function does not do anything.
-
-    user_api : {USER_APIS} or None (default=None)
-        APIs of libraries to limit. Used only if `limits` is an int.
-
-        - If "blas", it will only limit BLAS supported libraries ({BLAS_LIBS}).
-
-        - If "openmp", it will only limit OpenMP supported libraries
-          ({OPENMP_LIBS}). Note that it can affect the number of threads used
-          by the BLAS libraries if they rely on OpenMP.
-
-        - If None, this function will apply to all supported libraries.
-    """
-
-    def __init__(self, limits=None, user_api=None):
-        super().__init__(ThreadpoolController(), limits=limits, user_api=user_api)
-
-    @classmethod
-    def wrap(cls, limits=None, user_api=None):
-        return super().wrap(ThreadpoolController(), limits=limits, user_api=user_api)
-
-
-class ThreadpoolController:
-    """Collection of LibController objects for all loaded supported libraries
-
-    Attributes
-    ----------
-    lib_controllers : list of `LibController` objects
-        The list of library controllers of all loaded supported libraries.
-    """
-
-    # Cache for libc under POSIX and a few system libraries under Windows.
-    # We use a class level cache instead of an instance level cache because
-    # it's very unlikely that a shared library will be unloaded and reloaded
-    # during the lifetime of a program.
-    _system_libraries = dict()
-
-    def __init__(self):
-        self.lib_controllers = []
-        self._load_libraries()
-        self._warn_if_incompatible_openmp()
-
-    @classmethod
-    def _from_controllers(cls, lib_controllers):
-        new_controller = cls.__new__(cls)
-        new_controller.lib_controllers = lib_controllers
-        return new_controller
-
-    def info(self):
-        """Return lib_controllers info as a list of dicts"""
-        return [lib_controller.info() for lib_controller in self.lib_controllers]
-
-    def select(self, **kwargs):
-        """Return a ThreadpoolController containing a subset of its current
-        library controllers
-
-        It will select all libraries matching at least one pair (key, value) from kwargs
-        where key is an entry of the library info dict (like "user_api", "internal_api",
-        "prefix", ...) and value is the value or a list of acceptable values for that
-        entry.
-
-        For instance, `ThreadpoolController().select(internal_api=["blis", "openblas"])`
-        will select all library controllers whose internal_api is either "blis" or
-        "openblas".
-        """
-        for key, vals in kwargs.items():
-            kwargs[key] = [vals] if not isinstance(vals, list) else vals
-
-        lib_controllers = [
-            lib_controller
-            for lib_controller in self.lib_controllers
-            if any(
-                getattr(lib_controller, key, None) in vals
-                for key, vals in kwargs.items()
-            )
-        ]
-
-        return ThreadpoolController._from_controllers(lib_controllers)
-
-    def _get_params_for_sequential_blas_under_openmp(self):
-        """Return appropriate params to use for a sequential BLAS call in an OpenMP loop
-
-        This function takes into account the unexpected behavior of OpenBLAS with the
-        OpenMP threading layer.
-        """
-        if self.select(
-            internal_api="openblas", threading_layer="openmp"
-        ).lib_controllers:
-            return {"limits": None, "user_api": None}
-        return {"limits": 1, "user_api": "blas"}
-
-    @_format_docstring(
-        USER_APIS=", ".join('"{}"'.format(api) for api in _ALL_USER_APIS),
-        BLAS_LIBS=", ".join(_ALL_BLAS_LIBRARIES),
-        OPENMP_LIBS=", ".join(_ALL_OPENMP_LIBRARIES),
-    )
-    def limit(self, *, limits=None, user_api=None):
-        """Change the maximal number of threads that can be used in thread pools.
-
-        This function returns an object that can be used either as a callable (the
-        construction of this object limits the number of threads) or as a context
-        manager, in a `with` block to automatically restore the original state of the
-        controlled libraries when exiting the block.
-
-        Set the maximal number of threads that can be used in thread pools used in
-        the supported libraries to `limits`. This function works for libraries that
-        are already loaded in the interpreter and can be changed dynamically.
-
-        This effect is global and impacts the whole Python process. There is no thread
-        level isolation as these libraries do not offer thread-local APIs to configure
-        the number of threads to use in nested parallel calls.
-
-        Parameters
-        ----------
-        limits : int, dict, 'sequential_blas_under_openmp' or None (default=None)
-            The maximal number of threads that can be used in thread pools
-
-            - If int, sets the maximum number of threads to `limits` for each
-              library selected by `user_api`.
-
-            - If it is a dictionary `{{key: max_threads}}`, this function sets a
-              custom maximum number of threads for each `key` which can be either a
-              `user_api` or a `prefix` for a specific library.
-
-            - If 'sequential_blas_under_openmp', it will chose the appropriate `limits`
-              and `user_api` parameters for the specific use case of sequential BLAS
-              calls within an OpenMP parallel region. The `user_api` parameter is
-              ignored.
-
-            - If None, this function does not do anything.
-
-        user_api : {USER_APIS} or None (default=None)
-            APIs of libraries to limit. Used only if `limits` is an int.
-
-            - If "blas", it will only limit BLAS supported libraries ({BLAS_LIBS}).
-
-            - If "openmp", it will only limit OpenMP supported libraries
-              ({OPENMP_LIBS}). Note that it can affect the number of threads used
-              by the BLAS libraries if they rely on OpenMP.
-
-            - If None, this function will apply to all supported libraries.
-        """
-        return _ThreadpoolLimiter(self, limits=limits, user_api=user_api)
-
-    @_format_docstring(
-        USER_APIS=", ".join('"{}"'.format(api) for api in _ALL_USER_APIS),
-        BLAS_LIBS=", ".join(_ALL_BLAS_LIBRARIES),
-        OPENMP_LIBS=", ".join(_ALL_OPENMP_LIBRARIES),
-    )
-    def wrap(self, *, limits=None, user_api=None):
-        """Change the maximal number of threads that can be used in thread pools.
-
-        This function returns an object that can be used as a decorator.
-
-        Set the maximal number of threads that can be used in thread pools used in
-        the supported libraries to `limits`. This function works for libraries that
-        are already loaded in the interpreter and can be changed dynamically.
-
-        Parameters
-        ----------
-        limits : int, dict or None (default=None)
-            The maximal number of threads that can be used in thread pools
-
-            - If int, sets the maximum number of threads to `limits` for each
-              library selected by `user_api`.
-
-            - If it is a dictionary `{{key: max_threads}}`, this function sets a
-              custom maximum number of threads for each `key` which can be either a
-              `user_api` or a `prefix` for a specific library.
-
-            - If None, this function does not do anything.
-
-        user_api : {USER_APIS} or None (default=None)
-            APIs of libraries to limit. Used only if `limits` is an int.
-
-            - If "blas", it will only limit BLAS supported libraries ({BLAS_LIBS}).
-
-            - If "openmp", it will only limit OpenMP supported libraries
-              ({OPENMP_LIBS}). Note that it can affect the number of threads used
-              by the BLAS libraries if they rely on OpenMP.
-
-            - If None, this function will apply to all supported libraries.
-        """
-        return _ThreadpoolLimiter.wrap(self, limits=limits, user_api=user_api)
-
-    def __len__(self):
-        return len(self.lib_controllers)
-
-    def _load_libraries(self):
-        """Loop through loaded shared libraries and store the supported ones"""
-        if sys.platform == "darwin":
-            self._find_libraries_with_dyld()
-        elif sys.platform == "win32":
-            self._find_libraries_with_enum_process_module_ex()
-        elif "pyodide" in sys.modules:
-            self._find_libraries_pyodide()
-        else:
-            self._find_libraries_with_dl_iterate_phdr()
-
-    def _find_libraries_with_dl_iterate_phdr(self):
-        """Loop through loaded libraries and return binders on supported ones
-
-        This function is expected to work on POSIX system only.
-        This code is adapted from code by Intel developer @anton-malakhov
-        available at https://github.com/IntelPython/smp
-
-        Copyright (c) 2017, Intel Corporation published under the BSD 3-Clause
-        license
-        """
-        libc = self._get_libc()
-        if not hasattr(libc, "dl_iterate_phdr"):  # pragma: no cover
-            warnings.warn(
-                "Could not find dl_iterate_phdr in the C standard library.",
-                RuntimeWarning,
-            )
-            return []
-
-        # Callback function for `dl_iterate_phdr` which is called for every
-        # library loaded in the current process until it returns 1.
-        def match_library_callback(info, size, data):
-            # Get the path of the current library
-            filepath = info.contents.dlpi_name
-            if filepath:
-                filepath = filepath.decode("utf-8")
-
-                # Store the library controller if it is supported and selected
-                self._make_controller_from_path(filepath)
-            return 0
-
-        c_func_signature = ctypes.CFUNCTYPE(
-            ctypes.c_int,  # Return type
-            ctypes.POINTER(_dl_phdr_info),
-            ctypes.c_size_t,
-            ctypes.c_char_p,
-        )
-        c_match_library_callback = c_func_signature(match_library_callback)
-
-        data = ctypes.c_char_p(b"")
-        libc.dl_iterate_phdr(c_match_library_callback, data)
-
-    def _find_libraries_with_dyld(self):
-        """Loop through loaded libraries and return binders on supported ones
-
-        This function is expected to work on OSX system only
-        """
-        libc = self._get_libc()
-        if not hasattr(libc, "_dyld_image_count"):  # pragma: no cover
-            warnings.warn(
-                "Could not find _dyld_image_count in the C standard library.",
-                RuntimeWarning,
-            )
-            return []
-
-        n_dyld = libc._dyld_image_count()
-        libc._dyld_get_image_name.restype = ctypes.c_char_p
-
-        for i in range(n_dyld):
-            filepath = ctypes.string_at(libc._dyld_get_image_name(i))
-            filepath = filepath.decode("utf-8")
-
-            # Store the library controller if it is supported and selected
-            self._make_controller_from_path(filepath)
-
-    def _find_libraries_with_enum_process_module_ex(self):
-        """Loop through loaded libraries and return binders on supported ones
-
-        This function is expected to work on windows system only.
-        This code is adapted from code by Philipp Hagemeister @phihag available
-        at https://stackoverflow.com/questions/17474574
-        """
-        from ctypes.wintypes import DWORD, HMODULE, MAX_PATH
-
-        PROCESS_QUERY_INFORMATION = 0x0400
-        PROCESS_VM_READ = 0x0010
-
-        LIST_LIBRARIES_ALL = 0x03
-
-        ps_api = self._get_windll("Psapi")
-        kernel_32 = self._get_windll("kernel32")
-
-        h_process = kernel_32.OpenProcess(
-            PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, False, os.getpid()
-        )
-        if not h_process:  # pragma: no cover
-            raise OSError(f"Could not open PID {os.getpid()}")
-
-        try:
-            buf_count = 256
-            needed = DWORD()
-            # Grow the buffer until it becomes large enough to hold all the
-            # module headers
-            while True:
-                buf = (HMODULE * buf_count)()
-                buf_size = ctypes.sizeof(buf)
-                if not ps_api.EnumProcessModulesEx(
-                    h_process,
-                    ctypes.byref(buf),
-                    buf_size,
-                    ctypes.byref(needed),
-                    LIST_LIBRARIES_ALL,
-                ):
-                    raise OSError("EnumProcessModulesEx failed")
-                if buf_size >= needed.value:
-                    break
-                buf_count = needed.value // (buf_size // buf_count)
-
-            count = needed.value // (buf_size // buf_count)
-            h_modules = map(HMODULE, buf[:count])
-
-            # Loop through all the module headers and get the library path
-            # Allocate a buffer for the path 10 times the size of MAX_PATH to take
-            # into account long path names.
-            max_path = 10 * MAX_PATH
-            buf = ctypes.create_unicode_buffer(max_path)
-            n_size = DWORD()
-            for h_module in h_modules:
-                # Get the path of the current module
-                if not ps_api.GetModuleFileNameExW(
-                    h_process, h_module, ctypes.byref(buf), ctypes.byref(n_size)
-                ):
-                    raise OSError("GetModuleFileNameEx failed")
-                filepath = buf.value
-
-                if len(filepath) == max_path:  # pragma: no cover
-                    warnings.warn(
-                        "Could not get the full path of a dynamic library (path too "
-                        "long). This library will be ignored and threadpoolctl might "
-                        "not be able to control or display information about all "
-                        f"loaded libraries. Here's the truncated path: {filepath!r}",
-                        RuntimeWarning,
-                    )
-                else:
-                    # Store the library controller if it is supported and selected
-                    self._make_controller_from_path(filepath)
-        finally:
-            kernel_32.CloseHandle(h_process)
-
-    def _find_libraries_pyodide(self):
-        """Pyodide specific implementation for finding loaded libraries.
-
-        Adapted from suggestion in https://github.com/joblib/threadpoolctl/pull/169#issuecomment-1946696449.
-
-        One day, we may have a simpler solution. libc dl_iterate_phdr needs to
-        be implemented in Emscripten and exposed in Pyodide, see
-        https://github.com/emscripten-core/emscripten/issues/21354 for more
-        details.
-        """
-        try:
-            from pyodide_js._module import LDSO
-        except ImportError:
-            warnings.warn(
-                "Unable to import LDSO from pyodide_js._module. This should never "
-                "happen."
-            )
-            return
-
-        for filepath in LDSO.loadedLibsByName.as_object_map():
-            # Some libraries are duplicated by Pyodide and do not exist in the
-            # filesystem, so we first check for the existence of the file. For
-            # more details, see
-            # https://github.com/joblib/threadpoolctl/pull/169#issuecomment-1947946728
-            if os.path.exists(filepath):
-                self._make_controller_from_path(filepath)
-
-    def _make_controller_from_path(self, filepath):
-        """Store a library controller if it is supported and selected"""
-        # Required to resolve symlinks
-        filepath = _realpath(filepath)
-        # `lower` required to take account of OpenMP dll case on Windows
-        # (vcomp, VCOMP, Vcomp, ...)
-        filename = os.path.basename(filepath).lower()
-
-        # Loop through supported libraries to find if this filename corresponds
-        # to a supported one.
-        for controller_class in _ALL_CONTROLLERS:
-            # check if filename matches a supported prefix
-            prefix = self._check_prefix(filename, controller_class.filename_prefixes)
-
-            # filename does not match any of the prefixes of the candidate
-            # library. move to next library.
-            if prefix is None:
-                continue
-
-            # workaround for BLAS libraries packaged by conda-forge on windows, which
-            # are all renamed "libblas.dll". We thus have to check to which BLAS
-            # implementation it actually corresponds looking for implementation
-            # specific symbols.
-            if prefix == "libblas":
-                if filename.endswith(".dll"):
-                    libblas = ctypes.CDLL(filepath, _RTLD_NOLOAD)
-                    if not any(
-                        hasattr(libblas, func)
-                        for func in controller_class.check_symbols
-                    ):
-                        continue
-                else:
-                    # We ignore libblas on other platforms than windows because there
-                    # might be a libblas dso comming with openblas for instance that
-                    # can't be used to instantiate a pertinent LibController (many
-                    # symbols are missing) and would create confusion by making a
-                    # duplicate entry in threadpool_info.
-                    continue
-
-            # filename matches a prefix. Now we check if the library has the symbols we
-            # are looking for. If none of the symbols exists, it's very likely not the
-            # expected library (e.g. a library having a common prefix with one of the
-            # our supported libraries). Otherwise, create and store the library
-            # controller.
-            lib_controller = controller_class(
-                filepath=filepath, prefix=prefix, parent=self
-            )
-
-            if filepath in (lib.filepath for lib in self.lib_controllers):
-                # We already have a controller for this library.
-                continue
-
-            if not hasattr(controller_class, "check_symbols") or any(
-                hasattr(lib_controller.dynlib, func)
-                for func in controller_class.check_symbols
-            ):
-                self.lib_controllers.append(lib_controller)
-
-    def _check_prefix(self, library_basename, filename_prefixes):
-        """Return the prefix library_basename starts with
-
-        Return None if none matches.
-        """
-        for prefix in filename_prefixes:
-            if library_basename.startswith(prefix):
-                return prefix
-        return None
-
-    def _warn_if_incompatible_openmp(self):
-        """Raise a warning if llvm-OpenMP and intel-OpenMP are both loaded"""
-        prefixes = [lib_controller.prefix for lib_controller in self.lib_controllers]
-        msg = textwrap.dedent(
-            """
-            Found Intel OpenMP ('libiomp') and LLVM OpenMP ('libomp') loaded at
-            the same time. Both libraries are known to be incompatible and this
-            can cause random crashes or deadlocks on Linux when loaded in the
-            same Python program.
-            Using threadpoolctl may cause crashes or deadlocks. For more
-            information and possible workarounds, please see
-                https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md
-            """
-        )
-        if "libomp" in prefixes and "libiomp" in prefixes:
-            warnings.warn(msg, RuntimeWarning)
-
-    @classmethod
-    def _get_libc(cls):
-        """Load the lib-C for unix systems."""
-        libc = cls._system_libraries.get("libc")
-        if libc is None:
-            # Remark: If libc is statically linked or if Python is linked against an
-            # alternative implementation of libc like musl, find_library will return
-            # None and CDLL will load the main program itself which should contain the
-            # libc symbols. We still name it libc for convenience.
-            # If the main program does not contain the libc symbols, it's ok because
-            # we check their presence later anyway.
-            libc = ctypes.CDLL(find_library("c"), mode=_RTLD_NOLOAD)
-            cls._system_libraries["libc"] = libc
-        return libc
-
-    @classmethod
-    def _get_windll(cls, dll_name):
-        """Load a windows DLL"""
-        dll = cls._system_libraries.get(dll_name)
-        if dll is None:
-            dll = ctypes.WinDLL(f"{dll_name}.dll")
-            cls._system_libraries[dll_name] = dll
-        return dll
-
-
-def _main():
-    """Commandline interface to display thread-pool information and exit."""
-    import argparse
-    import importlib
-    import json
-    import sys
-
-    parser = argparse.ArgumentParser(
-        usage="python -m threadpoolctl -i numpy scipy.linalg xgboost",
-        description="Display thread-pool information and exit.",
-    )
-    parser.add_argument(
-        "-i",
-        "--import",
-        dest="modules",
-        nargs="*",
-        default=(),
-        help="Python modules to import before introspecting thread-pools.",
-    )
-    parser.add_argument(
-        "-c",
-        "--command",
-        help="a Python statement to execute before introspecting thread-pools.",
-    )
-
-    options = parser.parse_args(sys.argv[1:])
-    for module in options.modules:
-        try:
-            importlib.import_module(module, package=None)
-        except ImportError:
-            print("WARNING: could not import", module, file=sys.stderr)
-
-    if options.command:
-        exec(options.command)
-
-    print(json.dumps(threadpool_info(), indent=2))
-
-
-if __name__ == "__main__":
-    _main()
diff --git a/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/INSTALLER
deleted file mode 100644
index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/INSTALLER
+++ /dev/null
@@ -1 +0,0 @@
-pip
diff --git a/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/LICENSE b/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/LICENSE
deleted file mode 100644
index 5010e3075e6b150fdfc0042ca4adb3d41335e08d..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/LICENSE
+++ /dev/null
@@ -1,27 +0,0 @@
-The MIT License
-
-Copyright 2013-2019 William Pearson
-Copyright 2015-2016 Julien Enselme
-Copyright 2016 Google Inc.
-Copyright 2017 Samuel Vasko
-Copyright 2017 Nate Prewitt
-Copyright 2017 Jack Evans
-Copyright 2019 Filippo Broggini
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in
-all copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-THE SOFTWARE.
\ No newline at end of file
diff --git a/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/METADATA b/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/METADATA
deleted file mode 100644
index 6f2635ce4d573a350041792c4e5fe26fed51e986..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/METADATA
+++ /dev/null
@@ -1,255 +0,0 @@
-Metadata-Version: 2.1
-Name: toml
-Version: 0.10.2
-Summary: Python Library for Tom's Obvious, Minimal Language
-Home-page: https://github.com/uiri/toml
-Author: William Pearson
-Author-email: uiri@xqz.ca
-License: MIT
-Platform: UNKNOWN
-Classifier: Development Status :: 5 - Production/Stable
-Classifier: Intended Audience :: Developers
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Operating System :: OS Independent
-Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 2
-Classifier: Programming Language :: Python :: 2.6
-Classifier: Programming Language :: Python :: 2.7
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.3
-Classifier: Programming Language :: Python :: 3.4
-Classifier: Programming Language :: Python :: 3.5
-Classifier: Programming Language :: Python :: 3.6
-Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: Implementation :: CPython
-Classifier: Programming Language :: Python :: Implementation :: PyPy
-Requires-Python: >=2.6, !=3.0.*, !=3.1.*, !=3.2.*
-
-****
-TOML
-****
-
-.. image:: https://img.shields.io/pypi/v/toml
-    :target: https://pypi.org/project/toml/
-
-.. image:: https://travis-ci.org/uiri/toml.svg?branch=master
-    :target: https://travis-ci.org/uiri/toml
-
-.. image:: https://img.shields.io/pypi/pyversions/toml.svg
-    :target: https://pypi.org/project/toml/
-
-
-A Python library for parsing and creating `TOML `_.
-
-The module passes `the TOML test suite `_.
-
-See also:
-
-* `The TOML Standard `_
-* `The currently supported TOML specification `_
-
-Installation
-============
-
-To install the latest release on `PyPI `_,
-simply run:
-
-::
-
-  pip install toml
-
-Or to install the latest development version, run:
-
-::
-
-  git clone https://github.com/uiri/toml.git
-  cd toml
-  python setup.py install
-
-Quick Tutorial
-==============
-
-*toml.loads* takes in a string containing standard TOML-formatted data and
-returns a dictionary containing the parsed data.
-
-.. code:: pycon
-
-  >>> import toml
-  >>> toml_string = """
-  ... # This is a TOML document.
-  ...
-  ... title = "TOML Example"
-  ...
-  ... [owner]
-  ... name = "Tom Preston-Werner"
-  ... dob = 1979-05-27T07:32:00-08:00 # First class dates
-  ...
-  ... [database]
-  ... server = "192.168.1.1"
-  ... ports = [ 8001, 8001, 8002 ]
-  ... connection_max = 5000
-  ... enabled = true
-  ...
-  ... [servers]
-  ...
-  ...   # Indentation (tabs and/or spaces) is allowed but not required
-  ...   [servers.alpha]
-  ...   ip = "10.0.0.1"
-  ...   dc = "eqdc10"
-  ...
-  ...   [servers.beta]
-  ...   ip = "10.0.0.2"
-  ...   dc = "eqdc10"
-  ...
-  ... [clients]
-  ... data = [ ["gamma", "delta"], [1, 2] ]
-  ...
-  ... # Line breaks are OK when inside arrays
-  ... hosts = [
-  ...   "alpha",
-  ...   "omega"
-  ... ]
-  ... """
-  >>> parsed_toml = toml.loads(toml_string)
-
-
-*toml.dumps* takes a dictionary and returns a string containing the
-corresponding TOML-formatted data.
-
-.. code:: pycon
-
-  >>> new_toml_string = toml.dumps(parsed_toml)
-  >>> print(new_toml_string)
-  title = "TOML Example"
-  [owner]
-  name = "Tom Preston-Werner"
-  dob = 1979-05-27T07:32:00Z
-  [database]
-  server = "192.168.1.1"
-  ports = [ 8001, 8001, 8002,]
-  connection_max = 5000
-  enabled = true
-  [clients]
-  data = [ [ "gamma", "delta",], [ 1, 2,],]
-  hosts = [ "alpha", "omega",]
-  [servers.alpha]
-  ip = "10.0.0.1"
-  dc = "eqdc10"
-  [servers.beta]
-  ip = "10.0.0.2"
-  dc = "eqdc10"
-
-*toml.dump* takes a dictionary and a file descriptor and returns a string containing the
-corresponding TOML-formatted data.
-
-.. code:: pycon
-
-  >>> with open('new_toml_file.toml', 'w') as f:
-  ...     new_toml_string = toml.dump(parsed_toml, f)
-  >>> print(new_toml_string)
-  title = "TOML Example"
-  [owner]
-  name = "Tom Preston-Werner"
-  dob = 1979-05-27T07:32:00Z
-  [database]
-  server = "192.168.1.1"
-  ports = [ 8001, 8001, 8002,]
-  connection_max = 5000
-  enabled = true
-  [clients]
-  data = [ [ "gamma", "delta",], [ 1, 2,],]
-  hosts = [ "alpha", "omega",]
-  [servers.alpha]
-  ip = "10.0.0.1"
-  dc = "eqdc10"
-  [servers.beta]
-  ip = "10.0.0.2"
-  dc = "eqdc10"
-
-For more functions, view the API Reference below.
-
-Note
-----
-
-For Numpy users, by default the data types ``np.floatX`` will not be translated to floats by toml, but will instead be encoded as strings. To get around this, specify the ``TomlNumpyEncoder`` when saving your data.
-
-.. code:: pycon
-
-  >>> import toml
-  >>> import numpy as np
-  >>> a = np.arange(0, 10, dtype=np.double)
-  >>> output = {'a': a}
-  >>> toml.dumps(output)
-  'a = [ "0.0", "1.0", "2.0", "3.0", "4.0", "5.0", "6.0", "7.0", "8.0", "9.0",]\n'
-  >>> toml.dumps(output, encoder=toml.TomlNumpyEncoder())
-  'a = [ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0,]\n'
-
-API Reference
-=============
-
-``toml.load(f, _dict=dict)``
-  Parse a file or a list of files as TOML and return a dictionary.
-
-  :Args:
-    * ``f``: A path to a file, list of filepaths (to be read into single
-      object) or a file descriptor
-    * ``_dict``: The class of the dictionary object to be returned
-
-  :Returns:
-    A dictionary (or object ``_dict``) containing parsed TOML data
-
-  :Raises:
-    * ``TypeError``: When ``f`` is an invalid type or is a list containing
-      invalid types
-    * ``TomlDecodeError``: When an error occurs while decoding the file(s)
-
-``toml.loads(s, _dict=dict)``
-  Parse a TOML-formatted string to a dictionary.
-
-  :Args:
-    * ``s``: The TOML-formatted string to be parsed
-    * ``_dict``: Specifies the class of the returned toml dictionary
-
-  :Returns:
-    A dictionary (or object ``_dict``) containing parsed TOML data
-
-  :Raises:
-    * ``TypeError``: When a non-string object is passed
-    * ``TomlDecodeError``: When an error occurs while decoding the
-      TOML-formatted string
-
-``toml.dump(o, f, encoder=None)``
-  Write a dictionary to a file containing TOML-formatted data
-
-  :Args:
-    * ``o``: An object to be converted into TOML
-    * ``f``: A File descriptor where the TOML-formatted output should be stored
-    * ``encoder``: An instance of ``TomlEncoder`` (or subclass) for encoding the object. If ``None``, will default to ``TomlEncoder``
-
-  :Returns:
-    A string containing the TOML-formatted data corresponding to object ``o``
-
-  :Raises:
-    * ``TypeError``: When anything other than file descriptor is passed
-
-``toml.dumps(o, encoder=None)``
-  Create a TOML-formatted string from an input object
-
-  :Args:
-    * ``o``: An object to be converted into TOML
-    * ``encoder``: An instance of ``TomlEncoder`` (or subclass) for encoding the object. If ``None``, will default to ``TomlEncoder``
-
-  :Returns:
-    A string containing the TOML-formatted data corresponding to object ``o``
-
-
-
-Licensing
-=========
-
-This project is released under the terms of the MIT Open Source License. View
-*LICENSE.txt* for more information.
-
-
diff --git a/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/RECORD b/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/RECORD
deleted file mode 100644
index 3b057433f03fc9a1ffbbe3bdfd832b9ae4d709f7..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/RECORD
+++ /dev/null
@@ -1,17 +0,0 @@
-toml-0.10.2.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
-toml-0.10.2.dist-info/LICENSE,sha256=LZKUgj32yJNXyL5JJ_znk2HWVh5e51MtWSbmOTmqpTY,1252
-toml-0.10.2.dist-info/METADATA,sha256=n_YkspvEihd_QXLIZZ50WVSFz3rZ_k7jQP-OU1WUpWY,7142
-toml-0.10.2.dist-info/RECORD,,
-toml-0.10.2.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-toml-0.10.2.dist-info/WHEEL,sha256=ADKeyaGyKF5DwBNE0sRE5pvW-bSkFMJfBuhzZ3rceP4,110
-toml-0.10.2.dist-info/top_level.txt,sha256=2BO8ZRNnvJWgXyiQv66LBb_v87qBzcoUtEBefA75Ouk,5
-toml/__init__.py,sha256=Au3kqCwKD0cjbf4yJGOpUFwpsY0WHsC1ZRGvWgIKmpc,723
-toml/__pycache__/__init__.cpython-310.pyc,,
-toml/__pycache__/decoder.cpython-310.pyc,,
-toml/__pycache__/encoder.cpython-310.pyc,,
-toml/__pycache__/ordered.cpython-310.pyc,,
-toml/__pycache__/tz.cpython-310.pyc,,
-toml/decoder.py,sha256=hSGTLf-2WBDZ_ddoCHWFy6N647XyMSh1o3rN2o4dEFg,38942
-toml/encoder.py,sha256=XjBc8ayvvlsLyd_qDA4tMWDNmMFRS4DpwtuDSWBq7zo,9940
-toml/ordered.py,sha256=mz03lZmV0bmc9lsYRIUOuj7Dsu5Ptwq-UtGVq5FdVZ4,354
-toml/tz.py,sha256=-5vg8wkg_atnVi2TnEveexIVE7T_FxBVr_-2WVfO1oA,701
diff --git a/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/REQUESTED
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/WHEEL b/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/WHEEL
deleted file mode 100644
index 6d38aa0601b31c7f4c47ff3016173426df4e1d53..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/WHEEL
+++ /dev/null
@@ -1,6 +0,0 @@
-Wheel-Version: 1.0
-Generator: bdist_wheel (0.35.1)
-Root-Is-Purelib: true
-Tag: py2-none-any
-Tag: py3-none-any
-
diff --git a/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/top_level.txt
deleted file mode 100644
index bd79a658fe794e5e3fc28407e0cf0867e1d551a9..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/toml-0.10.2.dist-info/top_level.txt
+++ /dev/null
@@ -1 +0,0 @@
-toml
diff --git a/venv/lib/python3.10/site-packages/toml/__init__.py b/venv/lib/python3.10/site-packages/toml/__init__.py
deleted file mode 100644
index 7719ac23a73c5b5fbe73c814229e594850ab9ae0..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/toml/__init__.py
+++ /dev/null
@@ -1,25 +0,0 @@
-"""Python module which parses and emits TOML.
-
-Released under the MIT license.
-"""
-
-from toml import encoder
-from toml import decoder
-
-__version__ = "0.10.2"
-_spec_ = "0.5.0"
-
-load = decoder.load
-loads = decoder.loads
-TomlDecoder = decoder.TomlDecoder
-TomlDecodeError = decoder.TomlDecodeError
-TomlPreserveCommentDecoder = decoder.TomlPreserveCommentDecoder
-
-dump = encoder.dump
-dumps = encoder.dumps
-TomlEncoder = encoder.TomlEncoder
-TomlArraySeparatorEncoder = encoder.TomlArraySeparatorEncoder
-TomlPreserveInlineDictEncoder = encoder.TomlPreserveInlineDictEncoder
-TomlNumpyEncoder = encoder.TomlNumpyEncoder
-TomlPreserveCommentEncoder = encoder.TomlPreserveCommentEncoder
-TomlPathlibEncoder = encoder.TomlPathlibEncoder
diff --git a/venv/lib/python3.10/site-packages/toml/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/toml/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index b5b3989ed072dcae36b669b26c38fe55c262b7ae..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/toml/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/toml/__pycache__/decoder.cpython-310.pyc b/venv/lib/python3.10/site-packages/toml/__pycache__/decoder.cpython-310.pyc
deleted file mode 100644
index 795c5477c91f065a9d25542c611d2c9117e3569b..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/toml/__pycache__/decoder.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/toml/__pycache__/encoder.cpython-310.pyc b/venv/lib/python3.10/site-packages/toml/__pycache__/encoder.cpython-310.pyc
deleted file mode 100644
index fd85f8f1ec9b58ce99465bbcd336cc1e96b9fbf5..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/toml/__pycache__/encoder.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/toml/__pycache__/ordered.cpython-310.pyc b/venv/lib/python3.10/site-packages/toml/__pycache__/ordered.cpython-310.pyc
deleted file mode 100644
index ec264c35496bc9626a39fc7fe5aed15660307741..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/toml/__pycache__/ordered.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/toml/__pycache__/tz.cpython-310.pyc b/venv/lib/python3.10/site-packages/toml/__pycache__/tz.cpython-310.pyc
deleted file mode 100644
index 727abf4fd8a04469061d8502c97ecea0e5d87212..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/toml/__pycache__/tz.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/toml/decoder.py b/venv/lib/python3.10/site-packages/toml/decoder.py
deleted file mode 100644
index bf400e97611327e2092f31ac16a4d620d3502b0a..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/toml/decoder.py
+++ /dev/null
@@ -1,1057 +0,0 @@
-import datetime
-import io
-from os import linesep
-import re
-import sys
-
-from toml.tz import TomlTz
-
-if sys.version_info < (3,):
-    _range = xrange  # noqa: F821
-else:
-    unicode = str
-    _range = range
-    basestring = str
-    unichr = chr
-
-
-def _detect_pathlib_path(p):
-    if (3, 4) <= sys.version_info:
-        import pathlib
-        if isinstance(p, pathlib.PurePath):
-            return True
-    return False
-
-
-def _ispath(p):
-    if isinstance(p, (bytes, basestring)):
-        return True
-    return _detect_pathlib_path(p)
-
-
-def _getpath(p):
-    if (3, 6) <= sys.version_info:
-        import os
-        return os.fspath(p)
-    if _detect_pathlib_path(p):
-        return str(p)
-    return p
-
-
-try:
-    FNFError = FileNotFoundError
-except NameError:
-    FNFError = IOError
-
-
-TIME_RE = re.compile(r"([0-9]{2}):([0-9]{2}):([0-9]{2})(\.([0-9]{3,6}))?")
-
-
-class TomlDecodeError(ValueError):
-    """Base toml Exception / Error."""
-
-    def __init__(self, msg, doc, pos):
-        lineno = doc.count('\n', 0, pos) + 1
-        colno = pos - doc.rfind('\n', 0, pos)
-        emsg = '{} (line {} column {} char {})'.format(msg, lineno, colno, pos)
-        ValueError.__init__(self, emsg)
-        self.msg = msg
-        self.doc = doc
-        self.pos = pos
-        self.lineno = lineno
-        self.colno = colno
-
-
-# Matches a TOML number, which allows underscores for readability
-_number_with_underscores = re.compile('([0-9])(_([0-9]))*')
-
-
-class CommentValue(object):
-    def __init__(self, val, comment, beginline, _dict):
-        self.val = val
-        separator = "\n" if beginline else " "
-        self.comment = separator + comment
-        self._dict = _dict
-
-    def __getitem__(self, key):
-        return self.val[key]
-
-    def __setitem__(self, key, value):
-        self.val[key] = value
-
-    def dump(self, dump_value_func):
-        retstr = dump_value_func(self.val)
-        if isinstance(self.val, self._dict):
-            return self.comment + "\n" + unicode(retstr)
-        else:
-            return unicode(retstr) + self.comment
-
-
-def _strictly_valid_num(n):
-    n = n.strip()
-    if not n:
-        return False
-    if n[0] == '_':
-        return False
-    if n[-1] == '_':
-        return False
-    if "_." in n or "._" in n:
-        return False
-    if len(n) == 1:
-        return True
-    if n[0] == '0' and n[1] not in ['.', 'o', 'b', 'x']:
-        return False
-    if n[0] == '+' or n[0] == '-':
-        n = n[1:]
-        if len(n) > 1 and n[0] == '0' and n[1] != '.':
-            return False
-    if '__' in n:
-        return False
-    return True
-
-
-def load(f, _dict=dict, decoder=None):
-    """Parses named file or files as toml and returns a dictionary
-
-    Args:
-        f: Path to the file to open, array of files to read into single dict
-           or a file descriptor
-        _dict: (optional) Specifies the class of the returned toml dictionary
-        decoder: The decoder to use
-
-    Returns:
-        Parsed toml file represented as a dictionary
-
-    Raises:
-        TypeError -- When f is invalid type
-        TomlDecodeError: Error while decoding toml
-        IOError / FileNotFoundError -- When an array with no valid (existing)
-        (Python 2 / Python 3)          file paths is passed
-    """
-
-    if _ispath(f):
-        with io.open(_getpath(f), encoding='utf-8') as ffile:
-            return loads(ffile.read(), _dict, decoder)
-    elif isinstance(f, list):
-        from os import path as op
-        from warnings import warn
-        if not [path for path in f if op.exists(path)]:
-            error_msg = "Load expects a list to contain filenames only."
-            error_msg += linesep
-            error_msg += ("The list needs to contain the path of at least one "
-                          "existing file.")
-            raise FNFError(error_msg)
-        if decoder is None:
-            decoder = TomlDecoder(_dict)
-        d = decoder.get_empty_table()
-        for l in f:  # noqa: E741
-            if op.exists(l):
-                d.update(load(l, _dict, decoder))
-            else:
-                warn("Non-existent filename in list with at least one valid "
-                     "filename")
-        return d
-    else:
-        try:
-            return loads(f.read(), _dict, decoder)
-        except AttributeError:
-            raise TypeError("You can only load a file descriptor, filename or "
-                            "list")
-
-
-_groupname_re = re.compile(r'^[A-Za-z0-9_-]+$')
-
-
-def loads(s, _dict=dict, decoder=None):
-    """Parses string as toml
-
-    Args:
-        s: String to be parsed
-        _dict: (optional) Specifies the class of the returned toml dictionary
-
-    Returns:
-        Parsed toml file represented as a dictionary
-
-    Raises:
-        TypeError: When a non-string is passed
-        TomlDecodeError: Error while decoding toml
-    """
-
-    implicitgroups = []
-    if decoder is None:
-        decoder = TomlDecoder(_dict)
-    retval = decoder.get_empty_table()
-    currentlevel = retval
-    if not isinstance(s, basestring):
-        raise TypeError("Expecting something like a string")
-
-    if not isinstance(s, unicode):
-        s = s.decode('utf8')
-
-    original = s
-    sl = list(s)
-    openarr = 0
-    openstring = False
-    openstrchar = ""
-    multilinestr = False
-    arrayoftables = False
-    beginline = True
-    keygroup = False
-    dottedkey = False
-    keyname = 0
-    key = ''
-    prev_key = ''
-    line_no = 1
-
-    for i, item in enumerate(sl):
-        if item == '\r' and sl[i + 1] == '\n':
-            sl[i] = ' '
-            continue
-        if keyname:
-            key += item
-            if item == '\n':
-                raise TomlDecodeError("Key name found without value."
-                                      " Reached end of line.", original, i)
-            if openstring:
-                if item == openstrchar:
-                    oddbackslash = False
-                    k = 1
-                    while i >= k and sl[i - k] == '\\':
-                        oddbackslash = not oddbackslash
-                        k += 1
-                    if not oddbackslash:
-                        keyname = 2
-                        openstring = False
-                        openstrchar = ""
-                continue
-            elif keyname == 1:
-                if item.isspace():
-                    keyname = 2
-                    continue
-                elif item == '.':
-                    dottedkey = True
-                    continue
-                elif item.isalnum() or item == '_' or item == '-':
-                    continue
-                elif (dottedkey and sl[i - 1] == '.' and
-                      (item == '"' or item == "'")):
-                    openstring = True
-                    openstrchar = item
-                    continue
-            elif keyname == 2:
-                if item.isspace():
-                    if dottedkey:
-                        nextitem = sl[i + 1]
-                        if not nextitem.isspace() and nextitem != '.':
-                            keyname = 1
-                    continue
-                if item == '.':
-                    dottedkey = True
-                    nextitem = sl[i + 1]
-                    if not nextitem.isspace() and nextitem != '.':
-                        keyname = 1
-                    continue
-            if item == '=':
-                keyname = 0
-                prev_key = key[:-1].rstrip()
-                key = ''
-                dottedkey = False
-            else:
-                raise TomlDecodeError("Found invalid character in key name: '" +
-                                      item + "'. Try quoting the key name.",
-                                      original, i)
-        if item == "'" and openstrchar != '"':
-            k = 1
-            try:
-                while sl[i - k] == "'":
-                    k += 1
-                    if k == 3:
-                        break
-            except IndexError:
-                pass
-            if k == 3:
-                multilinestr = not multilinestr
-                openstring = multilinestr
-            else:
-                openstring = not openstring
-            if openstring:
-                openstrchar = "'"
-            else:
-                openstrchar = ""
-        if item == '"' and openstrchar != "'":
-            oddbackslash = False
-            k = 1
-            tripquote = False
-            try:
-                while sl[i - k] == '"':
-                    k += 1
-                    if k == 3:
-                        tripquote = True
-                        break
-                if k == 1 or (k == 3 and tripquote):
-                    while sl[i - k] == '\\':
-                        oddbackslash = not oddbackslash
-                        k += 1
-            except IndexError:
-                pass
-            if not oddbackslash:
-                if tripquote:
-                    multilinestr = not multilinestr
-                    openstring = multilinestr
-                else:
-                    openstring = not openstring
-            if openstring:
-                openstrchar = '"'
-            else:
-                openstrchar = ""
-        if item == '#' and (not openstring and not keygroup and
-                            not arrayoftables):
-            j = i
-            comment = ""
-            try:
-                while sl[j] != '\n':
-                    comment += s[j]
-                    sl[j] = ' '
-                    j += 1
-            except IndexError:
-                break
-            if not openarr:
-                decoder.preserve_comment(line_no, prev_key, comment, beginline)
-        if item == '[' and (not openstring and not keygroup and
-                            not arrayoftables):
-            if beginline:
-                if len(sl) > i + 1 and sl[i + 1] == '[':
-                    arrayoftables = True
-                else:
-                    keygroup = True
-            else:
-                openarr += 1
-        if item == ']' and not openstring:
-            if keygroup:
-                keygroup = False
-            elif arrayoftables:
-                if sl[i - 1] == ']':
-                    arrayoftables = False
-            else:
-                openarr -= 1
-        if item == '\n':
-            if openstring or multilinestr:
-                if not multilinestr:
-                    raise TomlDecodeError("Unbalanced quotes", original, i)
-                if ((sl[i - 1] == "'" or sl[i - 1] == '"') and (
-                        sl[i - 2] == sl[i - 1])):
-                    sl[i] = sl[i - 1]
-                    if sl[i - 3] == sl[i - 1]:
-                        sl[i - 3] = ' '
-            elif openarr:
-                sl[i] = ' '
-            else:
-                beginline = True
-            line_no += 1
-        elif beginline and sl[i] != ' ' and sl[i] != '\t':
-            beginline = False
-            if not keygroup and not arrayoftables:
-                if sl[i] == '=':
-                    raise TomlDecodeError("Found empty keyname. ", original, i)
-                keyname = 1
-                key += item
-    if keyname:
-        raise TomlDecodeError("Key name found without value."
-                              " Reached end of file.", original, len(s))
-    if openstring:  # reached EOF and have an unterminated string
-        raise TomlDecodeError("Unterminated string found."
-                              " Reached end of file.", original, len(s))
-    s = ''.join(sl)
-    s = s.split('\n')
-    multikey = None
-    multilinestr = ""
-    multibackslash = False
-    pos = 0
-    for idx, line in enumerate(s):
-        if idx > 0:
-            pos += len(s[idx - 1]) + 1
-
-        decoder.embed_comments(idx, currentlevel)
-
-        if not multilinestr or multibackslash or '\n' not in multilinestr:
-            line = line.strip()
-        if line == "" and (not multikey or multibackslash):
-            continue
-        if multikey:
-            if multibackslash:
-                multilinestr += line
-            else:
-                multilinestr += line
-            multibackslash = False
-            closed = False
-            if multilinestr[0] == '[':
-                closed = line[-1] == ']'
-            elif len(line) > 2:
-                closed = (line[-1] == multilinestr[0] and
-                          line[-2] == multilinestr[0] and
-                          line[-3] == multilinestr[0])
-            if closed:
-                try:
-                    value, vtype = decoder.load_value(multilinestr)
-                except ValueError as err:
-                    raise TomlDecodeError(str(err), original, pos)
-                currentlevel[multikey] = value
-                multikey = None
-                multilinestr = ""
-            else:
-                k = len(multilinestr) - 1
-                while k > -1 and multilinestr[k] == '\\':
-                    multibackslash = not multibackslash
-                    k -= 1
-                if multibackslash:
-                    multilinestr = multilinestr[:-1]
-                else:
-                    multilinestr += "\n"
-            continue
-        if line[0] == '[':
-            arrayoftables = False
-            if len(line) == 1:
-                raise TomlDecodeError("Opening key group bracket on line by "
-                                      "itself.", original, pos)
-            if line[1] == '[':
-                arrayoftables = True
-                line = line[2:]
-                splitstr = ']]'
-            else:
-                line = line[1:]
-                splitstr = ']'
-            i = 1
-            quotesplits = decoder._get_split_on_quotes(line)
-            quoted = False
-            for quotesplit in quotesplits:
-                if not quoted and splitstr in quotesplit:
-                    break
-                i += quotesplit.count(splitstr)
-                quoted = not quoted
-            line = line.split(splitstr, i)
-            if len(line) < i + 1 or line[-1].strip() != "":
-                raise TomlDecodeError("Key group not on a line by itself.",
-                                      original, pos)
-            groups = splitstr.join(line[:-1]).split('.')
-            i = 0
-            while i < len(groups):
-                groups[i] = groups[i].strip()
-                if len(groups[i]) > 0 and (groups[i][0] == '"' or
-                                           groups[i][0] == "'"):
-                    groupstr = groups[i]
-                    j = i + 1
-                    while ((not groupstr[0] == groupstr[-1]) or
-                           len(groupstr) == 1):
-                        j += 1
-                        if j > len(groups) + 2:
-                            raise TomlDecodeError("Invalid group name '" +
-                                                  groupstr + "' Something " +
-                                                  "went wrong.", original, pos)
-                        groupstr = '.'.join(groups[i:j]).strip()
-                    groups[i] = groupstr[1:-1]
-                    groups[i + 1:j] = []
-                else:
-                    if not _groupname_re.match(groups[i]):
-                        raise TomlDecodeError("Invalid group name '" +
-                                              groups[i] + "'. Try quoting it.",
-                                              original, pos)
-                i += 1
-            currentlevel = retval
-            for i in _range(len(groups)):
-                group = groups[i]
-                if group == "":
-                    raise TomlDecodeError("Can't have a keygroup with an empty "
-                                          "name", original, pos)
-                try:
-                    currentlevel[group]
-                    if i == len(groups) - 1:
-                        if group in implicitgroups:
-                            implicitgroups.remove(group)
-                            if arrayoftables:
-                                raise TomlDecodeError("An implicitly defined "
-                                                      "table can't be an array",
-                                                      original, pos)
-                        elif arrayoftables:
-                            currentlevel[group].append(decoder.get_empty_table()
-                                                       )
-                        else:
-                            raise TomlDecodeError("What? " + group +
-                                                  " already exists?" +
-                                                  str(currentlevel),
-                                                  original, pos)
-                except TypeError:
-                    currentlevel = currentlevel[-1]
-                    if group not in currentlevel:
-                        currentlevel[group] = decoder.get_empty_table()
-                        if i == len(groups) - 1 and arrayoftables:
-                            currentlevel[group] = [decoder.get_empty_table()]
-                except KeyError:
-                    if i != len(groups) - 1:
-                        implicitgroups.append(group)
-                    currentlevel[group] = decoder.get_empty_table()
-                    if i == len(groups) - 1 and arrayoftables:
-                        currentlevel[group] = [decoder.get_empty_table()]
-                currentlevel = currentlevel[group]
-                if arrayoftables:
-                    try:
-                        currentlevel = currentlevel[-1]
-                    except KeyError:
-                        pass
-        elif line[0] == "{":
-            if line[-1] != "}":
-                raise TomlDecodeError("Line breaks are not allowed in inline"
-                                      "objects", original, pos)
-            try:
-                decoder.load_inline_object(line, currentlevel, multikey,
-                                           multibackslash)
-            except ValueError as err:
-                raise TomlDecodeError(str(err), original, pos)
-        elif "=" in line:
-            try:
-                ret = decoder.load_line(line, currentlevel, multikey,
-                                        multibackslash)
-            except ValueError as err:
-                raise TomlDecodeError(str(err), original, pos)
-            if ret is not None:
-                multikey, multilinestr, multibackslash = ret
-    return retval
-
-
-def _load_date(val):
-    microsecond = 0
-    tz = None
-    try:
-        if len(val) > 19:
-            if val[19] == '.':
-                if val[-1].upper() == 'Z':
-                    subsecondval = val[20:-1]
-                    tzval = "Z"
-                else:
-                    subsecondvalandtz = val[20:]
-                    if '+' in subsecondvalandtz:
-                        splitpoint = subsecondvalandtz.index('+')
-                        subsecondval = subsecondvalandtz[:splitpoint]
-                        tzval = subsecondvalandtz[splitpoint:]
-                    elif '-' in subsecondvalandtz:
-                        splitpoint = subsecondvalandtz.index('-')
-                        subsecondval = subsecondvalandtz[:splitpoint]
-                        tzval = subsecondvalandtz[splitpoint:]
-                    else:
-                        tzval = None
-                        subsecondval = subsecondvalandtz
-                if tzval is not None:
-                    tz = TomlTz(tzval)
-                microsecond = int(int(subsecondval) *
-                                  (10 ** (6 - len(subsecondval))))
-            else:
-                tz = TomlTz(val[19:])
-    except ValueError:
-        tz = None
-    if "-" not in val[1:]:
-        return None
-    try:
-        if len(val) == 10:
-            d = datetime.date(
-                int(val[:4]), int(val[5:7]),
-                int(val[8:10]))
-        else:
-            d = datetime.datetime(
-                int(val[:4]), int(val[5:7]),
-                int(val[8:10]), int(val[11:13]),
-                int(val[14:16]), int(val[17:19]), microsecond, tz)
-    except ValueError:
-        return None
-    return d
-
-
-def _load_unicode_escapes(v, hexbytes, prefix):
-    skip = False
-    i = len(v) - 1
-    while i > -1 and v[i] == '\\':
-        skip = not skip
-        i -= 1
-    for hx in hexbytes:
-        if skip:
-            skip = False
-            i = len(hx) - 1
-            while i > -1 and hx[i] == '\\':
-                skip = not skip
-                i -= 1
-            v += prefix
-            v += hx
-            continue
-        hxb = ""
-        i = 0
-        hxblen = 4
-        if prefix == "\\U":
-            hxblen = 8
-        hxb = ''.join(hx[i:i + hxblen]).lower()
-        if hxb.strip('0123456789abcdef'):
-            raise ValueError("Invalid escape sequence: " + hxb)
-        if hxb[0] == "d" and hxb[1].strip('01234567'):
-            raise ValueError("Invalid escape sequence: " + hxb +
-                             ". Only scalar unicode points are allowed.")
-        v += unichr(int(hxb, 16))
-        v += unicode(hx[len(hxb):])
-    return v
-
-
-# Unescape TOML string values.
-
-# content after the \
-_escapes = ['0', 'b', 'f', 'n', 'r', 't', '"']
-# What it should be replaced by
-_escapedchars = ['\0', '\b', '\f', '\n', '\r', '\t', '\"']
-# Used for substitution
-_escape_to_escapedchars = dict(zip(_escapes, _escapedchars))
-
-
-def _unescape(v):
-    """Unescape characters in a TOML string."""
-    i = 0
-    backslash = False
-    while i < len(v):
-        if backslash:
-            backslash = False
-            if v[i] in _escapes:
-                v = v[:i - 1] + _escape_to_escapedchars[v[i]] + v[i + 1:]
-            elif v[i] == '\\':
-                v = v[:i - 1] + v[i:]
-            elif v[i] == 'u' or v[i] == 'U':
-                i += 1
-            else:
-                raise ValueError("Reserved escape sequence used")
-            continue
-        elif v[i] == '\\':
-            backslash = True
-        i += 1
-    return v
-
-
-class InlineTableDict(object):
-    """Sentinel subclass of dict for inline tables."""
-
-
-class TomlDecoder(object):
-
-    def __init__(self, _dict=dict):
-        self._dict = _dict
-
-    def get_empty_table(self):
-        return self._dict()
-
-    def get_empty_inline_table(self):
-        class DynamicInlineTableDict(self._dict, InlineTableDict):
-            """Concrete sentinel subclass for inline tables.
-            It is a subclass of _dict which is passed in dynamically at load
-            time
-
-            It is also a subclass of InlineTableDict
-            """
-
-        return DynamicInlineTableDict()
-
-    def load_inline_object(self, line, currentlevel, multikey=False,
-                           multibackslash=False):
-        candidate_groups = line[1:-1].split(",")
-        groups = []
-        if len(candidate_groups) == 1 and not candidate_groups[0].strip():
-            candidate_groups.pop()
-        while len(candidate_groups) > 0:
-            candidate_group = candidate_groups.pop(0)
-            try:
-                _, value = candidate_group.split('=', 1)
-            except ValueError:
-                raise ValueError("Invalid inline table encountered")
-            value = value.strip()
-            if ((value[0] == value[-1] and value[0] in ('"', "'")) or (
-                    value[0] in '-0123456789' or
-                    value in ('true', 'false') or
-                    (value[0] == "[" and value[-1] == "]") or
-                    (value[0] == '{' and value[-1] == '}'))):
-                groups.append(candidate_group)
-            elif len(candidate_groups) > 0:
-                candidate_groups[0] = (candidate_group + "," +
-                                       candidate_groups[0])
-            else:
-                raise ValueError("Invalid inline table value encountered")
-        for group in groups:
-            status = self.load_line(group, currentlevel, multikey,
-                                    multibackslash)
-            if status is not None:
-                break
-
-    def _get_split_on_quotes(self, line):
-        doublequotesplits = line.split('"')
-        quoted = False
-        quotesplits = []
-        if len(doublequotesplits) > 1 and "'" in doublequotesplits[0]:
-            singlequotesplits = doublequotesplits[0].split("'")
-            doublequotesplits = doublequotesplits[1:]
-            while len(singlequotesplits) % 2 == 0 and len(doublequotesplits):
-                singlequotesplits[-1] += '"' + doublequotesplits[0]
-                doublequotesplits = doublequotesplits[1:]
-                if "'" in singlequotesplits[-1]:
-                    singlequotesplits = (singlequotesplits[:-1] +
-                                         singlequotesplits[-1].split("'"))
-            quotesplits += singlequotesplits
-        for doublequotesplit in doublequotesplits:
-            if quoted:
-                quotesplits.append(doublequotesplit)
-            else:
-                quotesplits += doublequotesplit.split("'")
-                quoted = not quoted
-        return quotesplits
-
-    def load_line(self, line, currentlevel, multikey, multibackslash):
-        i = 1
-        quotesplits = self._get_split_on_quotes(line)
-        quoted = False
-        for quotesplit in quotesplits:
-            if not quoted and '=' in quotesplit:
-                break
-            i += quotesplit.count('=')
-            quoted = not quoted
-        pair = line.split('=', i)
-        strictly_valid = _strictly_valid_num(pair[-1])
-        if _number_with_underscores.match(pair[-1]):
-            pair[-1] = pair[-1].replace('_', '')
-        while len(pair[-1]) and (pair[-1][0] != ' ' and pair[-1][0] != '\t' and
-                                 pair[-1][0] != "'" and pair[-1][0] != '"' and
-                                 pair[-1][0] != '[' and pair[-1][0] != '{' and
-                                 pair[-1].strip() != 'true' and
-                                 pair[-1].strip() != 'false'):
-            try:
-                float(pair[-1])
-                break
-            except ValueError:
-                pass
-            if _load_date(pair[-1]) is not None:
-                break
-            if TIME_RE.match(pair[-1]):
-                break
-            i += 1
-            prev_val = pair[-1]
-            pair = line.split('=', i)
-            if prev_val == pair[-1]:
-                raise ValueError("Invalid date or number")
-            if strictly_valid:
-                strictly_valid = _strictly_valid_num(pair[-1])
-        pair = ['='.join(pair[:-1]).strip(), pair[-1].strip()]
-        if '.' in pair[0]:
-            if '"' in pair[0] or "'" in pair[0]:
-                quotesplits = self._get_split_on_quotes(pair[0])
-                quoted = False
-                levels = []
-                for quotesplit in quotesplits:
-                    if quoted:
-                        levels.append(quotesplit)
-                    else:
-                        levels += [level.strip() for level in
-                                   quotesplit.split('.')]
-                    quoted = not quoted
-            else:
-                levels = pair[0].split('.')
-            while levels[-1] == "":
-                levels = levels[:-1]
-            for level in levels[:-1]:
-                if level == "":
-                    continue
-                if level not in currentlevel:
-                    currentlevel[level] = self.get_empty_table()
-                currentlevel = currentlevel[level]
-            pair[0] = levels[-1].strip()
-        elif (pair[0][0] == '"' or pair[0][0] == "'") and \
-                (pair[0][-1] == pair[0][0]):
-            pair[0] = _unescape(pair[0][1:-1])
-        k, koffset = self._load_line_multiline_str(pair[1])
-        if k > -1:
-            while k > -1 and pair[1][k + koffset] == '\\':
-                multibackslash = not multibackslash
-                k -= 1
-            if multibackslash:
-                multilinestr = pair[1][:-1]
-            else:
-                multilinestr = pair[1] + "\n"
-            multikey = pair[0]
-        else:
-            value, vtype = self.load_value(pair[1], strictly_valid)
-        try:
-            currentlevel[pair[0]]
-            raise ValueError("Duplicate keys!")
-        except TypeError:
-            raise ValueError("Duplicate keys!")
-        except KeyError:
-            if multikey:
-                return multikey, multilinestr, multibackslash
-            else:
-                currentlevel[pair[0]] = value
-
-    def _load_line_multiline_str(self, p):
-        poffset = 0
-        if len(p) < 3:
-            return -1, poffset
-        if p[0] == '[' and (p.strip()[-1] != ']' and
-                            self._load_array_isstrarray(p)):
-            newp = p[1:].strip().split(',')
-            while len(newp) > 1 and newp[-1][0] != '"' and newp[-1][0] != "'":
-                newp = newp[:-2] + [newp[-2] + ',' + newp[-1]]
-            newp = newp[-1]
-            poffset = len(p) - len(newp)
-            p = newp
-        if p[0] != '"' and p[0] != "'":
-            return -1, poffset
-        if p[1] != p[0] or p[2] != p[0]:
-            return -1, poffset
-        if len(p) > 5 and p[-1] == p[0] and p[-2] == p[0] and p[-3] == p[0]:
-            return -1, poffset
-        return len(p) - 1, poffset
-
-    def load_value(self, v, strictly_valid=True):
-        if not v:
-            raise ValueError("Empty value is invalid")
-        if v == 'true':
-            return (True, "bool")
-        elif v.lower() == 'true':
-            raise ValueError("Only all lowercase booleans allowed")
-        elif v == 'false':
-            return (False, "bool")
-        elif v.lower() == 'false':
-            raise ValueError("Only all lowercase booleans allowed")
-        elif v[0] == '"' or v[0] == "'":
-            quotechar = v[0]
-            testv = v[1:].split(quotechar)
-            triplequote = False
-            triplequotecount = 0
-            if len(testv) > 1 and testv[0] == '' and testv[1] == '':
-                testv = testv[2:]
-                triplequote = True
-            closed = False
-            for tv in testv:
-                if tv == '':
-                    if triplequote:
-                        triplequotecount += 1
-                    else:
-                        closed = True
-                else:
-                    oddbackslash = False
-                    try:
-                        i = -1
-                        j = tv[i]
-                        while j == '\\':
-                            oddbackslash = not oddbackslash
-                            i -= 1
-                            j = tv[i]
-                    except IndexError:
-                        pass
-                    if not oddbackslash:
-                        if closed:
-                            raise ValueError("Found tokens after a closed " +
-                                             "string. Invalid TOML.")
-                        else:
-                            if not triplequote or triplequotecount > 1:
-                                closed = True
-                            else:
-                                triplequotecount = 0
-            if quotechar == '"':
-                escapeseqs = v.split('\\')[1:]
-                backslash = False
-                for i in escapeseqs:
-                    if i == '':
-                        backslash = not backslash
-                    else:
-                        if i[0] not in _escapes and (i[0] != 'u' and
-                                                     i[0] != 'U' and
-                                                     not backslash):
-                            raise ValueError("Reserved escape sequence used")
-                        if backslash:
-                            backslash = False
-                for prefix in ["\\u", "\\U"]:
-                    if prefix in v:
-                        hexbytes = v.split(prefix)
-                        v = _load_unicode_escapes(hexbytes[0], hexbytes[1:],
-                                                  prefix)
-                v = _unescape(v)
-            if len(v) > 1 and v[1] == quotechar and (len(v) < 3 or
-                                                     v[1] == v[2]):
-                v = v[2:-2]
-            return (v[1:-1], "str")
-        elif v[0] == '[':
-            return (self.load_array(v), "array")
-        elif v[0] == '{':
-            inline_object = self.get_empty_inline_table()
-            self.load_inline_object(v, inline_object)
-            return (inline_object, "inline_object")
-        elif TIME_RE.match(v):
-            h, m, s, _, ms = TIME_RE.match(v).groups()
-            time = datetime.time(int(h), int(m), int(s), int(ms) if ms else 0)
-            return (time, "time")
-        else:
-            parsed_date = _load_date(v)
-            if parsed_date is not None:
-                return (parsed_date, "date")
-            if not strictly_valid:
-                raise ValueError("Weirdness with leading zeroes or "
-                                 "underscores in your number.")
-            itype = "int"
-            neg = False
-            if v[0] == '-':
-                neg = True
-                v = v[1:]
-            elif v[0] == '+':
-                v = v[1:]
-            v = v.replace('_', '')
-            lowerv = v.lower()
-            if '.' in v or ('x' not in v and ('e' in v or 'E' in v)):
-                if '.' in v and v.split('.', 1)[1] == '':
-                    raise ValueError("This float is missing digits after "
-                                     "the point")
-                if v[0] not in '0123456789':
-                    raise ValueError("This float doesn't have a leading "
-                                     "digit")
-                v = float(v)
-                itype = "float"
-            elif len(lowerv) == 3 and (lowerv == 'inf' or lowerv == 'nan'):
-                v = float(v)
-                itype = "float"
-            if itype == "int":
-                v = int(v, 0)
-            if neg:
-                return (0 - v, itype)
-            return (v, itype)
-
-    def bounded_string(self, s):
-        if len(s) == 0:
-            return True
-        if s[-1] != s[0]:
-            return False
-        i = -2
-        backslash = False
-        while len(s) + i > 0:
-            if s[i] == "\\":
-                backslash = not backslash
-                i -= 1
-            else:
-                break
-        return not backslash
-
-    def _load_array_isstrarray(self, a):
-        a = a[1:-1].strip()
-        if a != '' and (a[0] == '"' or a[0] == "'"):
-            return True
-        return False
-
-    def load_array(self, a):
-        atype = None
-        retval = []
-        a = a.strip()
-        if '[' not in a[1:-1] or "" != a[1:-1].split('[')[0].strip():
-            strarray = self._load_array_isstrarray(a)
-            if not a[1:-1].strip().startswith('{'):
-                a = a[1:-1].split(',')
-            else:
-                # a is an inline object, we must find the matching parenthesis
-                # to define groups
-                new_a = []
-                start_group_index = 1
-                end_group_index = 2
-                open_bracket_count = 1 if a[start_group_index] == '{' else 0
-                in_str = False
-                while end_group_index < len(a[1:]):
-                    if a[end_group_index] == '"' or a[end_group_index] == "'":
-                        if in_str:
-                            backslash_index = end_group_index - 1
-                            while (backslash_index > -1 and
-                                   a[backslash_index] == '\\'):
-                                in_str = not in_str
-                                backslash_index -= 1
-                        in_str = not in_str
-                    if not in_str and a[end_group_index] == '{':
-                        open_bracket_count += 1
-                    if in_str or a[end_group_index] != '}':
-                        end_group_index += 1
-                        continue
-                    elif a[end_group_index] == '}' and open_bracket_count > 1:
-                        open_bracket_count -= 1
-                        end_group_index += 1
-                        continue
-
-                    # Increase end_group_index by 1 to get the closing bracket
-                    end_group_index += 1
-
-                    new_a.append(a[start_group_index:end_group_index])
-
-                    # The next start index is at least after the closing
-                    # bracket, a closing bracket can be followed by a comma
-                    # since we are in an array.
-                    start_group_index = end_group_index + 1
-                    while (start_group_index < len(a[1:]) and
-                           a[start_group_index] != '{'):
-                        start_group_index += 1
-                    end_group_index = start_group_index + 1
-                a = new_a
-            b = 0
-            if strarray:
-                while b < len(a) - 1:
-                    ab = a[b].strip()
-                    while (not self.bounded_string(ab) or
-                           (len(ab) > 2 and
-                            ab[0] == ab[1] == ab[2] and
-                            ab[-2] != ab[0] and
-                            ab[-3] != ab[0])):
-                        a[b] = a[b] + ',' + a[b + 1]
-                        ab = a[b].strip()
-                        if b < len(a) - 2:
-                            a = a[:b + 1] + a[b + 2:]
-                        else:
-                            a = a[:b + 1]
-                    b += 1
-        else:
-            al = list(a[1:-1])
-            a = []
-            openarr = 0
-            j = 0
-            for i in _range(len(al)):
-                if al[i] == '[':
-                    openarr += 1
-                elif al[i] == ']':
-                    openarr -= 1
-                elif al[i] == ',' and not openarr:
-                    a.append(''.join(al[j:i]))
-                    j = i + 1
-            a.append(''.join(al[j:]))
-        for i in _range(len(a)):
-            a[i] = a[i].strip()
-            if a[i] != '':
-                nval, ntype = self.load_value(a[i])
-                if atype:
-                    if ntype != atype:
-                        raise ValueError("Not a homogeneous array")
-                else:
-                    atype = ntype
-                retval.append(nval)
-        return retval
-
-    def preserve_comment(self, line_no, key, comment, beginline):
-        pass
-
-    def embed_comments(self, idx, currentlevel):
-        pass
-
-
-class TomlPreserveCommentDecoder(TomlDecoder):
-
-    def __init__(self, _dict=dict):
-        self.saved_comments = {}
-        super(TomlPreserveCommentDecoder, self).__init__(_dict)
-
-    def preserve_comment(self, line_no, key, comment, beginline):
-        self.saved_comments[line_no] = (key, comment, beginline)
-
-    def embed_comments(self, idx, currentlevel):
-        if idx not in self.saved_comments:
-            return
-
-        key, comment, beginline = self.saved_comments[idx]
-        currentlevel[key] = CommentValue(currentlevel[key], comment, beginline,
-                                         self._dict)
diff --git a/venv/lib/python3.10/site-packages/toml/encoder.py b/venv/lib/python3.10/site-packages/toml/encoder.py
deleted file mode 100644
index bf17a72b62366870e2b73f2de8ff6c898105cf64..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/toml/encoder.py
+++ /dev/null
@@ -1,304 +0,0 @@
-import datetime
-import re
-import sys
-from decimal import Decimal
-
-from toml.decoder import InlineTableDict
-
-if sys.version_info >= (3,):
-    unicode = str
-
-
-def dump(o, f, encoder=None):
-    """Writes out dict as toml to a file
-
-    Args:
-        o: Object to dump into toml
-        f: File descriptor where the toml should be stored
-        encoder: The ``TomlEncoder`` to use for constructing the output string
-
-    Returns:
-        String containing the toml corresponding to dictionary
-
-    Raises:
-        TypeError: When anything other than file descriptor is passed
-    """
-
-    if not f.write:
-        raise TypeError("You can only dump an object to a file descriptor")
-    d = dumps(o, encoder=encoder)
-    f.write(d)
-    return d
-
-
-def dumps(o, encoder=None):
-    """Stringifies input dict as toml
-
-    Args:
-        o: Object to dump into toml
-        encoder: The ``TomlEncoder`` to use for constructing the output string
-
-    Returns:
-        String containing the toml corresponding to dict
-
-    Examples:
-        ```python
-        >>> import toml
-        >>> output = {
-        ... 'a': "I'm a string",
-        ... 'b': ["I'm", "a", "list"],
-        ... 'c': 2400
-        ... }
-        >>> toml.dumps(output)
-        'a = "I\'m a string"\nb = [ "I\'m", "a", "list",]\nc = 2400\n'
-        ```
-    """
-
-    retval = ""
-    if encoder is None:
-        encoder = TomlEncoder(o.__class__)
-    addtoretval, sections = encoder.dump_sections(o, "")
-    retval += addtoretval
-    outer_objs = [id(o)]
-    while sections:
-        section_ids = [id(section) for section in sections.values()]
-        for outer_obj in outer_objs:
-            if outer_obj in section_ids:
-                raise ValueError("Circular reference detected")
-        outer_objs += section_ids
-        newsections = encoder.get_empty_table()
-        for section in sections:
-            addtoretval, addtosections = encoder.dump_sections(
-                sections[section], section)
-
-            if addtoretval or (not addtoretval and not addtosections):
-                if retval and retval[-2:] != "\n\n":
-                    retval += "\n"
-                retval += "[" + section + "]\n"
-                if addtoretval:
-                    retval += addtoretval
-            for s in addtosections:
-                newsections[section + "." + s] = addtosections[s]
-        sections = newsections
-    return retval
-
-
-def _dump_str(v):
-    if sys.version_info < (3,) and hasattr(v, 'decode') and isinstance(v, str):
-        v = v.decode('utf-8')
-    v = "%r" % v
-    if v[0] == 'u':
-        v = v[1:]
-    singlequote = v.startswith("'")
-    if singlequote or v.startswith('"'):
-        v = v[1:-1]
-    if singlequote:
-        v = v.replace("\\'", "'")
-        v = v.replace('"', '\\"')
-    v = v.split("\\x")
-    while len(v) > 1:
-        i = -1
-        if not v[0]:
-            v = v[1:]
-        v[0] = v[0].replace("\\\\", "\\")
-        # No, I don't know why != works and == breaks
-        joinx = v[0][i] != "\\"
-        while v[0][:i] and v[0][i] == "\\":
-            joinx = not joinx
-            i -= 1
-        if joinx:
-            joiner = "x"
-        else:
-            joiner = "u00"
-        v = [v[0] + joiner + v[1]] + v[2:]
-    return unicode('"' + v[0] + '"')
-
-
-def _dump_float(v):
-    return "{}".format(v).replace("e+0", "e+").replace("e-0", "e-")
-
-
-def _dump_time(v):
-    utcoffset = v.utcoffset()
-    if utcoffset is None:
-        return v.isoformat()
-    # The TOML norm specifies that it's local time thus we drop the offset
-    return v.isoformat()[:-6]
-
-
-class TomlEncoder(object):
-
-    def __init__(self, _dict=dict, preserve=False):
-        self._dict = _dict
-        self.preserve = preserve
-        self.dump_funcs = {
-            str: _dump_str,
-            unicode: _dump_str,
-            list: self.dump_list,
-            bool: lambda v: unicode(v).lower(),
-            int: lambda v: v,
-            float: _dump_float,
-            Decimal: _dump_float,
-            datetime.datetime: lambda v: v.isoformat().replace('+00:00', 'Z'),
-            datetime.time: _dump_time,
-            datetime.date: lambda v: v.isoformat()
-        }
-
-    def get_empty_table(self):
-        return self._dict()
-
-    def dump_list(self, v):
-        retval = "["
-        for u in v:
-            retval += " " + unicode(self.dump_value(u)) + ","
-        retval += "]"
-        return retval
-
-    def dump_inline_table(self, section):
-        """Preserve inline table in its compact syntax instead of expanding
-        into subsection.
-
-        https://github.com/toml-lang/toml#user-content-inline-table
-        """
-        retval = ""
-        if isinstance(section, dict):
-            val_list = []
-            for k, v in section.items():
-                val = self.dump_inline_table(v)
-                val_list.append(k + " = " + val)
-            retval += "{ " + ", ".join(val_list) + " }\n"
-            return retval
-        else:
-            return unicode(self.dump_value(section))
-
-    def dump_value(self, v):
-        # Lookup function corresponding to v's type
-        dump_fn = self.dump_funcs.get(type(v))
-        if dump_fn is None and hasattr(v, '__iter__'):
-            dump_fn = self.dump_funcs[list]
-        # Evaluate function (if it exists) else return v
-        return dump_fn(v) if dump_fn is not None else self.dump_funcs[str](v)
-
-    def dump_sections(self, o, sup):
-        retstr = ""
-        if sup != "" and sup[-1] != ".":
-            sup += '.'
-        retdict = self._dict()
-        arraystr = ""
-        for section in o:
-            section = unicode(section)
-            qsection = section
-            if not re.match(r'^[A-Za-z0-9_-]+$', section):
-                qsection = _dump_str(section)
-            if not isinstance(o[section], dict):
-                arrayoftables = False
-                if isinstance(o[section], list):
-                    for a in o[section]:
-                        if isinstance(a, dict):
-                            arrayoftables = True
-                if arrayoftables:
-                    for a in o[section]:
-                        arraytabstr = "\n"
-                        arraystr += "[[" + sup + qsection + "]]\n"
-                        s, d = self.dump_sections(a, sup + qsection)
-                        if s:
-                            if s[0] == "[":
-                                arraytabstr += s
-                            else:
-                                arraystr += s
-                        while d:
-                            newd = self._dict()
-                            for dsec in d:
-                                s1, d1 = self.dump_sections(d[dsec], sup +
-                                                            qsection + "." +
-                                                            dsec)
-                                if s1:
-                                    arraytabstr += ("[" + sup + qsection +
-                                                    "." + dsec + "]\n")
-                                    arraytabstr += s1
-                                for s1 in d1:
-                                    newd[dsec + "." + s1] = d1[s1]
-                            d = newd
-                        arraystr += arraytabstr
-                else:
-                    if o[section] is not None:
-                        retstr += (qsection + " = " +
-                                   unicode(self.dump_value(o[section])) + '\n')
-            elif self.preserve and isinstance(o[section], InlineTableDict):
-                retstr += (qsection + " = " +
-                           self.dump_inline_table(o[section]))
-            else:
-                retdict[qsection] = o[section]
-        retstr += arraystr
-        return (retstr, retdict)
-
-
-class TomlPreserveInlineDictEncoder(TomlEncoder):
-
-    def __init__(self, _dict=dict):
-        super(TomlPreserveInlineDictEncoder, self).__init__(_dict, True)
-
-
-class TomlArraySeparatorEncoder(TomlEncoder):
-
-    def __init__(self, _dict=dict, preserve=False, separator=","):
-        super(TomlArraySeparatorEncoder, self).__init__(_dict, preserve)
-        if separator.strip() == "":
-            separator = "," + separator
-        elif separator.strip(' \t\n\r,'):
-            raise ValueError("Invalid separator for arrays")
-        self.separator = separator
-
-    def dump_list(self, v):
-        t = []
-        retval = "["
-        for u in v:
-            t.append(self.dump_value(u))
-        while t != []:
-            s = []
-            for u in t:
-                if isinstance(u, list):
-                    for r in u:
-                        s.append(r)
-                else:
-                    retval += " " + unicode(u) + self.separator
-            t = s
-        retval += "]"
-        return retval
-
-
-class TomlNumpyEncoder(TomlEncoder):
-
-    def __init__(self, _dict=dict, preserve=False):
-        import numpy as np
-        super(TomlNumpyEncoder, self).__init__(_dict, preserve)
-        self.dump_funcs[np.float16] = _dump_float
-        self.dump_funcs[np.float32] = _dump_float
-        self.dump_funcs[np.float64] = _dump_float
-        self.dump_funcs[np.int16] = self._dump_int
-        self.dump_funcs[np.int32] = self._dump_int
-        self.dump_funcs[np.int64] = self._dump_int
-
-    def _dump_int(self, v):
-        return "{}".format(int(v))
-
-
-class TomlPreserveCommentEncoder(TomlEncoder):
-
-    def __init__(self, _dict=dict, preserve=False):
-        from toml.decoder import CommentValue
-        super(TomlPreserveCommentEncoder, self).__init__(_dict, preserve)
-        self.dump_funcs[CommentValue] = lambda v: v.dump(self.dump_value)
-
-
-class TomlPathlibEncoder(TomlEncoder):
-
-    def _dump_pathlib_path(self, v):
-        return _dump_str(str(v))
-
-    def dump_value(self, v):
-        if (3, 4) <= sys.version_info:
-            import pathlib
-            if isinstance(v, pathlib.PurePath):
-                v = str(v)
-        return super(TomlPathlibEncoder, self).dump_value(v)
diff --git a/venv/lib/python3.10/site-packages/toml/ordered.py b/venv/lib/python3.10/site-packages/toml/ordered.py
deleted file mode 100644
index 9c20c41a1b9bcdb641dc52865bce04ccdb8f046c..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/toml/ordered.py
+++ /dev/null
@@ -1,15 +0,0 @@
-from collections import OrderedDict
-from toml import TomlEncoder
-from toml import TomlDecoder
-
-
-class TomlOrderedDecoder(TomlDecoder):
-
-    def __init__(self):
-        super(self.__class__, self).__init__(_dict=OrderedDict)
-
-
-class TomlOrderedEncoder(TomlEncoder):
-
-    def __init__(self):
-        super(self.__class__, self).__init__(_dict=OrderedDict)
diff --git a/venv/lib/python3.10/site-packages/toml/tz.py b/venv/lib/python3.10/site-packages/toml/tz.py
deleted file mode 100644
index bf20593a2647cb2cf0a858fb19734874867352ef..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/toml/tz.py
+++ /dev/null
@@ -1,24 +0,0 @@
-from datetime import tzinfo, timedelta
-
-
-class TomlTz(tzinfo):
-    def __init__(self, toml_offset):
-        if toml_offset == "Z":
-            self._raw_offset = "+00:00"
-        else:
-            self._raw_offset = toml_offset
-        self._sign = -1 if self._raw_offset[0] == '-' else 1
-        self._hours = int(self._raw_offset[1:3])
-        self._minutes = int(self._raw_offset[4:6])
-
-    def __deepcopy__(self, memo):
-        return self.__class__(self._raw_offset)
-
-    def tzname(self, dt):
-        return "UTC" + self._raw_offset
-
-    def utcoffset(self, dt):
-        return self._sign * timedelta(hours=self._hours, minutes=self._minutes)
-
-    def dst(self, dt):
-        return timedelta(0)
diff --git a/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/INSTALLER
deleted file mode 100644
index a1b589e38a32041e49332e5e81c2d363dc418d68..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/INSTALLER
+++ /dev/null
@@ -1 +0,0 @@
-pip
diff --git a/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/METADATA b/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/METADATA
deleted file mode 100644
index 1de4a67d210059b2ce00534323b0dad953ea10ed..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/METADATA
+++ /dev/null
@@ -1,422 +0,0 @@
-Metadata-Version: 2.1
-Name: typer
-Version: 0.16.0
-Summary: Typer, build great CLIs. Easy to code. Based on Python type hints.
-Author-Email: =?utf-8?q?Sebasti=C3=A1n_Ram=C3=ADrez?= 
-Classifier: Intended Audience :: Information Technology
-Classifier: Intended Audience :: System Administrators
-Classifier: Operating System :: OS Independent
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python
-Classifier: Topic :: Software Development :: Libraries :: Application Frameworks
-Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Classifier: Topic :: Software Development :: Libraries
-Classifier: Topic :: Software Development
-Classifier: Typing :: Typed
-Classifier: Development Status :: 4 - Beta
-Classifier: Intended Audience :: Developers
-Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Classifier: Programming Language :: Python :: 3.13
-Classifier: License :: OSI Approved :: MIT License
-Project-URL: Homepage, https://github.com/fastapi/typer
-Project-URL: Documentation, https://typer.tiangolo.com
-Project-URL: Repository, https://github.com/fastapi/typer
-Project-URL: Issues, https://github.com/fastapi/typer/issues
-Project-URL: Changelog, https://typer.tiangolo.com/release-notes/
-Requires-Python: >=3.7
-Requires-Dist: click>=8.0.0
-Requires-Dist: typing-extensions>=3.7.4.3
-Requires-Dist: shellingham>=1.3.0
-Requires-Dist: rich>=10.11.0
-Description-Content-Type: text/markdown
-
-

-  Typer
-
-

-

-    Typer, build great CLIs. Easy to code. Based on Python type hints.
-

-

-
-    Test
-
-
-    Publish
-
-
-    Coverage
-
-    Package version
-
-

-
----
-
-**Documentation**: https://typer.tiangolo.com
-
-**Source Code**: https://github.com/fastapi/typer
-
----
-
-Typer is a library for building CLI applications that users will **love using** and developers will **love creating**. Based on Python type hints.
-
-It's also a command line tool to run scripts, automatically converting them to CLI applications.
-
-The key features are:
-
-* **Intuitive to write**: Great editor support. Completion everywhere. Less time debugging. Designed to be easy to use and learn. Less time reading docs.
-* **Easy to use**: It's easy to use for the final users. Automatic help, and automatic completion for all shells.
-* **Short**: Minimize code duplication. Multiple features from each parameter declaration. Fewer bugs.
-* **Start simple**: The simplest example adds only 2 lines of code to your app: **1 import, 1 function call**.
-* **Grow large**: Grow in complexity as much as you want, create arbitrarily complex trees of commands and groups of subcommands, with options and arguments.
-* **Run scripts**: Typer includes a `typer` command/program that you can use to run scripts, automatically converting them to CLIs, even if they don't use Typer internally.
-
-## FastAPI of CLIs
-
-**Typer** is FastAPI's little sibling, it's the FastAPI of CLIs.
-
-## Installation
-
-Create and activate a virtual environment and then install **Typer**:
-
-

-
-```console
-$ pip install typer
----> 100%
-Successfully installed typer rich shellingham
-```
-
-

-
-## Example
-
-### The absolute minimum
-
-* Create a file `main.py` with:
-
-```Python
-def main(name: str):
-    print(f"Hello {name}")
-```
-
-This script doesn't even use Typer internally. But you can use the `typer` command to run it as a CLI application.
-
-### Run it
-
-Run your application with the `typer` command:
-
-

-
-```console
-// Run your application
-$ typer main.py run
-
-// You get a nice error, you are missing NAME
-Usage: typer [PATH_OR_MODULE] run [OPTIONS] NAME
-Try 'typer [PATH_OR_MODULE] run --help' for help.
-╭─ Error ───────────────────────────────────────────╮
-│ Missing argument 'NAME'.                          │
-╰───────────────────────────────────────────────────╯
-
-
-// You get a --help for free
-$ typer main.py run --help
-
-Usage: typer [PATH_OR_MODULE] run [OPTIONS] NAME
-
-Run the provided Typer app.
-
-╭─ Arguments ───────────────────────────────────────╮
-│ *    name      TEXT  [default: None] [required]   |
-╰───────────────────────────────────────────────────╯
-╭─ Options ─────────────────────────────────────────╮
-│ --help          Show this message and exit.       │
-╰───────────────────────────────────────────────────╯
-
-// Now pass the NAME argument
-$ typer main.py run Camila
-
-Hello Camila
-
-// It works! 🎉
-```
-
-

-
-This is the simplest use case, not even using Typer internally, but it can already be quite useful for simple scripts.
-
-**Note**: auto-completion works when you create a Python package and run it with `--install-completion` or when you use the `typer` command.
-
-## Use Typer in your code
-
-Now let's start using Typer in your own code, update `main.py` with:
-
-```Python
-import typer
-
-
-def main(name: str):
-    print(f"Hello {name}")
-
-
-if __name__ == "__main__":
-    typer.run(main)
-```
-
-Now you could run it with Python directly:
-
-

-
-```console
-// Run your application
-$ python main.py
-
-// You get a nice error, you are missing NAME
-Usage: main.py [OPTIONS] NAME
-Try 'main.py --help' for help.
-╭─ Error ───────────────────────────────────────────╮
-│ Missing argument 'NAME'.                          │
-╰───────────────────────────────────────────────────╯
-
-
-// You get a --help for free
-$ python main.py --help
-
-Usage: main.py [OPTIONS] NAME
-
-╭─ Arguments ───────────────────────────────────────╮
-│ *    name      TEXT  [default: None] [required]   |
-╰───────────────────────────────────────────────────╯
-╭─ Options ─────────────────────────────────────────╮
-│ --help          Show this message and exit.       │
-╰───────────────────────────────────────────────────╯
-
-// Now pass the NAME argument
-$ python main.py Camila
-
-Hello Camila
-
-// It works! 🎉
-```
-
-

-
-**Note**: you can also call this same script with the `typer` command, but you don't need to.
-
-## Example upgrade
-
-This was the simplest example possible.
-
-Now let's see one a bit more complex.
-
-### An example with two subcommands
-
-Modify the file `main.py`.
-
-Create a `typer.Typer()` app, and create two subcommands with their parameters.
-
-```Python hl_lines="3  6  11  20"
-import typer
-
-app = typer.Typer()
-
-
-@app.command()
-def hello(name: str):
-    print(f"Hello {name}")
-
-
-@app.command()
-def goodbye(name: str, formal: bool = False):
-    if formal:
-        print(f"Goodbye Ms. {name}. Have a good day.")
-    else:
-        print(f"Bye {name}!")
-
-
-if __name__ == "__main__":
-    app()
-```
-
-And that will:
-
-* Explicitly create a `typer.Typer` app.
-    * The previous `typer.run` actually creates one implicitly for you.
-* Add two subcommands with `@app.command()`.
-* Execute the `app()` itself, as if it was a function (instead of `typer.run`).
-
-### Run the upgraded example
-
-Check the new help:
-
-

-
-```console
-$ python main.py --help
-
- Usage: main.py [OPTIONS] COMMAND [ARGS]...
-
-╭─ Options ─────────────────────────────────────────╮
-│ --install-completion          Install completion  │
-│                               for the current     │
-│                               shell.              │
-│ --show-completion             Show completion for │
-│                               the current shell,  │
-│                               to copy it or       │
-│                               customize the       │
-│                               installation.       │
-│ --help                        Show this message   │
-│                               and exit.           │
-╰───────────────────────────────────────────────────╯
-╭─ Commands ────────────────────────────────────────╮
-│ goodbye                                           │
-│ hello                                             │
-╰───────────────────────────────────────────────────╯
-
-// When you create a package you get ✨ auto-completion ✨ for free, installed with --install-completion
-
-// You have 2 subcommands (the 2 functions): goodbye and hello
-```
-
-

-
-Now check the help for the `hello` command:
-
-

-
-```console
-$ python main.py hello --help
-
- Usage: main.py hello [OPTIONS] NAME
-
-╭─ Arguments ───────────────────────────────────────╮
-│ *    name      TEXT  [default: None] [required]   │
-╰───────────────────────────────────────────────────╯
-╭─ Options ─────────────────────────────────────────╮
-│ --help          Show this message and exit.       │
-╰───────────────────────────────────────────────────╯
-```
-
-

-
-And now check the help for the `goodbye` command:
-
-

-
-```console
-$ python main.py goodbye --help
-
- Usage: main.py goodbye [OPTIONS] NAME
-
-╭─ Arguments ───────────────────────────────────────╮
-│ *    name      TEXT  [default: None] [required]   │
-╰───────────────────────────────────────────────────╯
-╭─ Options ─────────────────────────────────────────╮
-│ --formal    --no-formal      [default: no-formal] │
-│ --help                       Show this message    │
-│                              and exit.            │
-╰───────────────────────────────────────────────────╯
-
-// Automatic --formal and --no-formal for the bool option 🎉
-```
-
-

-
-Now you can try out the new command line application:
-
-

-
-```console
-// Use it with the hello command
-
-$ python main.py hello Camila
-
-Hello Camila
-
-// And with the goodbye command
-
-$ python main.py goodbye Camila
-
-Bye Camila!
-
-// And with --formal
-
-$ python main.py goodbye --formal Camila
-
-Goodbye Ms. Camila. Have a good day.
-```
-
-

-
-### Recap
-
-In summary, you declare **once** the types of parameters (*CLI arguments* and *CLI options*) as function parameters.
-
-You do that with standard modern Python types.
-
-You don't have to learn a new syntax, the methods or classes of a specific library, etc.
-
-Just standard **Python**.
-
-For example, for an `int`:
-
-```Python
-total: int
-```
-
-or for a `bool` flag:
-
-```Python
-force: bool
-```
-
-And similarly for **files**, **paths**, **enums** (choices), etc. And there are tools to create **groups of subcommands**, add metadata, extra **validation**, etc.
-
-**You get**: great editor support, including **completion** and **type checks** everywhere.
-
-**Your users get**: automatic **`--help`**, **auto-completion** in their terminal (Bash, Zsh, Fish, PowerShell) when they install your package or when using the `typer` command.
-
-For a more complete example including more features, see the Tutorial - User Guide.
-
-## Dependencies
-
-**Typer** stands on the shoulders of a giant. Its only internal required dependency is Click.
-
-By default it also comes with extra standard dependencies:
-
-* rich: to show nicely formatted errors automatically.
-* shellingham: to automatically detect the current shell when installing completion.
-    * With `shellingham` you can just use `--install-completion`.
-    * Without `shellingham`, you have to pass the name of the shell to install completion for, e.g. `--install-completion bash`.
-
-### `typer-slim`
-
-If you don't want the extra standard optional dependencies, install `typer-slim` instead.
-
-When you install with:
-
-```bash
-pip install typer
-```
-
-...it includes the same code and dependencies as:
-
-```bash
-pip install "typer-slim[standard]"
-```
-
-The `standard` extra dependencies are `rich` and `shellingham`.
-
-**Note**: The `typer` command is only included in the `typer` package.
-
-## License
-
-This project is licensed under the terms of the MIT license.
diff --git a/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/RECORD b/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/RECORD
deleted file mode 100644
index b80218804fb4a41f552a44bf19c3201789c51587..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/RECORD
+++ /dev/null
@@ -1,41 +0,0 @@
-../../../bin/typer,sha256=P2GrxFgKtN7Jd8wTJYKjQ307MTvKCnRbuQz6Oz-UALY,226
-typer-0.16.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
-typer-0.16.0.dist-info/METADATA,sha256=4n4FU9l0gR5JhaGE3ERjZn_aHAklo64CkFk_gzI1kdM,15721
-typer-0.16.0.dist-info/RECORD,,
-typer-0.16.0.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-typer-0.16.0.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
-typer-0.16.0.dist-info/entry_points.txt,sha256=YO13ByiqWeuas9V0JADLUARZFUe_cwU_7wmTNvxBYQ8,57
-typer-0.16.0.dist-info/licenses/LICENSE,sha256=WJks68-N-25AxOIRLtEhJsJDZm3KORKj14t-ysSFnUk,1086
-typer/__init__.py,sha256=D3u-F2ltL-fo4S8GGp0g-OEaovfTSp-W6eAv1jKqBU8,1596
-typer/__main__.py,sha256=bYt9eEaoRQWdejEHFD8REx9jxVEdZptECFsV7F49Ink,30
-typer/__pycache__/__init__.cpython-310.pyc,,
-typer/__pycache__/__main__.cpython-310.pyc,,
-typer/__pycache__/_completion_classes.cpython-310.pyc,,
-typer/__pycache__/_completion_shared.cpython-310.pyc,,
-typer/__pycache__/_types.cpython-310.pyc,,
-typer/__pycache__/_typing.cpython-310.pyc,,
-typer/__pycache__/cli.cpython-310.pyc,,
-typer/__pycache__/colors.cpython-310.pyc,,
-typer/__pycache__/completion.cpython-310.pyc,,
-typer/__pycache__/core.cpython-310.pyc,,
-typer/__pycache__/main.cpython-310.pyc,,
-typer/__pycache__/models.cpython-310.pyc,,
-typer/__pycache__/params.cpython-310.pyc,,
-typer/__pycache__/rich_utils.cpython-310.pyc,,
-typer/__pycache__/testing.cpython-310.pyc,,
-typer/__pycache__/utils.cpython-310.pyc,,
-typer/_completion_classes.py,sha256=Px9bV56Y4J_F9S_sUI2iPpUh_i--Chocxrk4HAut2HE,7385
-typer/_completion_shared.py,sha256=4lFOUhXSry2bIZR9al7Uq33itpgSzg6eagquAysNmOE,8758
-typer/_types.py,sha256=kSLxhKmX37YzizQjqYUAWmr_JFcCW5vhEc4YshDTC9Q,1031
-typer/_typing.py,sha256=nsZ-TKcMlGAtqiWXM60r97rqtWMZxdhwj_YkLo8_neM,3001
-typer/cli.py,sha256=YaXpDud7wRtDCsJsWkE1L0BPWhpIHAaPbvZNLTJ854w,9779
-typer/colors.py,sha256=e42j8uB520hLpX5C_0fiR3OOoIFMbhO3ADZvv6hlAV8,430
-typer/completion.py,sha256=d1AiptrsEwUeSPQaIA5oiv4T3Xy3MDq5o3VIjK39Qeg,4810
-typer/core.py,sha256=TsdtoJYOOa4-sppkYybsAbLJ6gN3EJyPSXUJu5ib5rA,26530
-typer/main.py,sha256=Upb2hrXidyzCkSZQ5BdDipG3c4KRU5v7L-fV-9b19Mw,42043
-typer/models.py,sha256=Q6v9BQYutNlH44i7fn0YbZ-OeRXsgib1o_tR4gTmqow,17188
-typer/params.py,sha256=MRVCwRPzNMkOdYU6VNVGkawX_gAoYzbiCfL_tYcR6x8,14929
-typer/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-typer/rich_utils.py,sha256=svygTTSUOs94K3r_VNnmMjO1GOXGVa0GqX59k2c77aU,25315
-typer/testing.py,sha256=Mb_HqTkpPw24qsVYxCQrDJpjq_oOHlgqZpauWofxkq0,874
-typer/utils.py,sha256=G0qddDX06YtHuMJNCmj-frLJYkYxfUa7iwO6KOTX2FI,7368
diff --git a/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/REQUESTED
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/WHEEL b/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/WHEEL
deleted file mode 100644
index 45ec8c4e6bf53687a824041168247bae0953dd7c..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/WHEEL
+++ /dev/null
@@ -1,4 +0,0 @@
-Wheel-Version: 1.0
-Generator: pdm-backend (2.4.4)
-Root-Is-Purelib: true
-Tag: py3-none-any
diff --git a/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/entry_points.txt b/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/entry_points.txt
deleted file mode 100644
index ca44c05a071b87fe02a3724bd895838305fe4e5d..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/entry_points.txt
+++ /dev/null
@@ -1,5 +0,0 @@
-[console_scripts]
-typer = typer.cli:main
-
-[gui_scripts]
-
diff --git a/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/licenses/LICENSE b/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/licenses/LICENSE
deleted file mode 100644
index a7694736cf37716aafec14b24aa8d6316ebe07a3..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer-0.16.0.dist-info/licenses/LICENSE
+++ /dev/null
@@ -1,21 +0,0 @@
-The MIT License (MIT)
-
-Copyright (c) 2019 Sebastián Ramírez
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in
-all copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-THE SOFTWARE.
diff --git a/venv/lib/python3.10/site-packages/typer/__init__.py b/venv/lib/python3.10/site-packages/typer/__init__.py
deleted file mode 100644
index 067b36d4899927c6b2b4429634cd3458df4f0ac8..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer/__init__.py
+++ /dev/null
@@ -1,39 +0,0 @@
-"""Typer, build great CLIs. Easy to code. Based on Python type hints."""
-
-__version__ = "0.16.0"
-
-from shutil import get_terminal_size as get_terminal_size
-
-from click.exceptions import Abort as Abort
-from click.exceptions import BadParameter as BadParameter
-from click.exceptions import Exit as Exit
-from click.termui import clear as clear
-from click.termui import confirm as confirm
-from click.termui import echo_via_pager as echo_via_pager
-from click.termui import edit as edit
-from click.termui import getchar as getchar
-from click.termui import pause as pause
-from click.termui import progressbar as progressbar
-from click.termui import prompt as prompt
-from click.termui import secho as secho
-from click.termui import style as style
-from click.termui import unstyle as unstyle
-from click.utils import echo as echo
-from click.utils import format_filename as format_filename
-from click.utils import get_app_dir as get_app_dir
-from click.utils import get_binary_stream as get_binary_stream
-from click.utils import get_text_stream as get_text_stream
-from click.utils import open_file as open_file
-
-from . import colors as colors
-from .main import Typer as Typer
-from .main import launch as launch
-from .main import run as run
-from .models import CallbackParam as CallbackParam
-from .models import Context as Context
-from .models import FileBinaryRead as FileBinaryRead
-from .models import FileBinaryWrite as FileBinaryWrite
-from .models import FileText as FileText
-from .models import FileTextWrite as FileTextWrite
-from .params import Argument as Argument
-from .params import Option as Option
diff --git a/venv/lib/python3.10/site-packages/typer/__main__.py b/venv/lib/python3.10/site-packages/typer/__main__.py
deleted file mode 100644
index 4e28416e104515e90fca4b69cc60d0c61fd15d61..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer/__main__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from .cli import main
-
-main()
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index ab375b7b761b0544f5068e0acd4675185501c4f0..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/__main__.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/__main__.cpython-310.pyc
deleted file mode 100644
index 6076738e3bb2de28977a600d48121dd90e79c396..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/__main__.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/_completion_classes.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/_completion_classes.cpython-310.pyc
deleted file mode 100644
index 5ba0542e972e8ecc6d159c3cc80df66bf64cc3b2..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/_completion_classes.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/_completion_shared.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/_completion_shared.cpython-310.pyc
deleted file mode 100644
index ff1a2d176d714f4041f2b75d76542797a9f83d36..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/_completion_shared.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/_types.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/_types.cpython-310.pyc
deleted file mode 100644
index bca2d9ebb218815847f1a39ce458814fda19844c..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/_types.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/_typing.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/_typing.cpython-310.pyc
deleted file mode 100644
index 7445fc632d628c63047d2c92dbac16d3584fcde3..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/_typing.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/cli.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/cli.cpython-310.pyc
deleted file mode 100644
index b0982b575784aead331d2be666bdac164a5dfeec..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/cli.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/colors.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/colors.cpython-310.pyc
deleted file mode 100644
index 030ea1d606cc6c4de76af4363ecfdf9711ed01c7..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/colors.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/completion.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/completion.cpython-310.pyc
deleted file mode 100644
index bc186e7f038056cab30d889529b52f3e47563253..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/completion.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/core.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/core.cpython-310.pyc
deleted file mode 100644
index e67441eda7dd53f3685b6a78dc86b9485f984a33..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/core.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/main.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/main.cpython-310.pyc
deleted file mode 100644
index 663479137de81071d1553431f65dd3d5090fefab..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/main.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/models.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/models.cpython-310.pyc
deleted file mode 100644
index b68721f88f2b54ce180361d0a1a9664cee087d46..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/models.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/params.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/params.cpython-310.pyc
deleted file mode 100644
index 2beb9d176739c1f24cc7eff3cfb1fe7dc03700d9..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/params.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/rich_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/rich_utils.cpython-310.pyc
deleted file mode 100644
index 0fd68815f15ae82ec96bf22f90d8c281e3e6c0c7..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/rich_utils.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/testing.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/testing.cpython-310.pyc
deleted file mode 100644
index 53b6ea5829374e24d70f7c82ae5372caeaad3bd2..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/testing.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/typer/__pycache__/utils.cpython-310.pyc
deleted file mode 100644
index a7e78f3491824511b25f1e8e22edd0f07fa8c35f..0000000000000000000000000000000000000000
Binary files a/venv/lib/python3.10/site-packages/typer/__pycache__/utils.cpython-310.pyc and /dev/null differ
diff --git a/venv/lib/python3.10/site-packages/typer/_completion_classes.py b/venv/lib/python3.10/site-packages/typer/_completion_classes.py
deleted file mode 100644
index 5980248afe4863048ecdc42e6b83281f2f870622..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer/_completion_classes.py
+++ /dev/null
@@ -1,211 +0,0 @@
-import importlib.util
-import os
-import re
-import sys
-from typing import Any, Dict, List, Tuple
-
-import click
-import click.parser
-import click.shell_completion
-
-from ._completion_shared import (
-    COMPLETION_SCRIPT_BASH,
-    COMPLETION_SCRIPT_FISH,
-    COMPLETION_SCRIPT_POWER_SHELL,
-    COMPLETION_SCRIPT_ZSH,
-    Shells,
-)
-
-try:
-    from click.shell_completion import split_arg_string as click_split_arg_string
-except ImportError:  # pragma: no cover
-    # TODO: when removing support for Click < 8.2, remove this import
-    from click.parser import (  # type: ignore[no-redef]
-        split_arg_string as click_split_arg_string,
-    )
-
-try:
-    import shellingham
-except ImportError:  # pragma: no cover
-    shellingham = None
-
-
-def _sanitize_help_text(text: str) -> str:
-    """Sanitizes the help text by removing rich tags"""
-    if not importlib.util.find_spec("rich"):
-        return text
-    from . import rich_utils
-
-    return rich_utils.rich_render_text(text)
-
-
-class BashComplete(click.shell_completion.BashComplete):
-    name = Shells.bash.value
-    source_template = COMPLETION_SCRIPT_BASH
-
-    def source_vars(self) -> Dict[str, Any]:
-        return {
-            "complete_func": self.func_name,
-            "autocomplete_var": self.complete_var,
-            "prog_name": self.prog_name,
-        }
-
-    def get_completion_args(self) -> Tuple[List[str], str]:
-        cwords = click_split_arg_string(os.environ["COMP_WORDS"])
-        cword = int(os.environ["COMP_CWORD"])
-        args = cwords[1:cword]
-
-        try:
-            incomplete = cwords[cword]
-        except IndexError:
-            incomplete = ""
-
-        return args, incomplete
-
-    def format_completion(self, item: click.shell_completion.CompletionItem) -> str:
-        # TODO: Explore replicating the new behavior from Click, with item types and
-        # triggering completion for files and directories
-        # return f"{item.type},{item.value}"
-        return f"{item.value}"
-
-    def complete(self) -> str:
-        args, incomplete = self.get_completion_args()
-        completions = self.get_completions(args, incomplete)
-        out = [self.format_completion(item) for item in completions]
-        return "\n".join(out)
-
-
-class ZshComplete(click.shell_completion.ZshComplete):
-    name = Shells.zsh.value
-    source_template = COMPLETION_SCRIPT_ZSH
-
-    def source_vars(self) -> Dict[str, Any]:
-        return {
-            "complete_func": self.func_name,
-            "autocomplete_var": self.complete_var,
-            "prog_name": self.prog_name,
-        }
-
-    def get_completion_args(self) -> Tuple[List[str], str]:
-        completion_args = os.getenv("_TYPER_COMPLETE_ARGS", "")
-        cwords = click_split_arg_string(completion_args)
-        args = cwords[1:]
-        if args and not completion_args.endswith(" "):
-            incomplete = args[-1]
-            args = args[:-1]
-        else:
-            incomplete = ""
-        return args, incomplete
-
-    def format_completion(self, item: click.shell_completion.CompletionItem) -> str:
-        def escape(s: str) -> str:
-            return (
-                s.replace('"', '""')
-                .replace("'", "''")
-                .replace("$", "\\$")
-                .replace("`", "\\`")
-                .replace(":", r"\\:")
-            )
-
-        # TODO: Explore replicating the new behavior from Click, pay attention to
-        # the difference with and without escape
-        # return f"{item.type}\n{item.value}\n{item.help if item.help else '_'}"
-        if item.help:
-            return f'"{escape(item.value)}":"{_sanitize_help_text(escape(item.help))}"'
-        else:
-            return f'"{escape(item.value)}"'
-
-    def complete(self) -> str:
-        args, incomplete = self.get_completion_args()
-        completions = self.get_completions(args, incomplete)
-        res = [self.format_completion(item) for item in completions]
-        if res:
-            args_str = "\n".join(res)
-            return f"_arguments '*: :(({args_str}))'"
-        else:
-            return "_files"
-
-
-class FishComplete(click.shell_completion.FishComplete):
-    name = Shells.fish.value
-    source_template = COMPLETION_SCRIPT_FISH
-
-    def source_vars(self) -> Dict[str, Any]:
-        return {
-            "complete_func": self.func_name,
-            "autocomplete_var": self.complete_var,
-            "prog_name": self.prog_name,
-        }
-
-    def get_completion_args(self) -> Tuple[List[str], str]:
-        completion_args = os.getenv("_TYPER_COMPLETE_ARGS", "")
-        cwords = click_split_arg_string(completion_args)
-        args = cwords[1:]
-        if args and not completion_args.endswith(" "):
-            incomplete = args[-1]
-            args = args[:-1]
-        else:
-            incomplete = ""
-        return args, incomplete
-
-    def format_completion(self, item: click.shell_completion.CompletionItem) -> str:
-        # TODO: Explore replicating the new behavior from Click, pay attention to
-        # the difference with and without formatted help
-        # if item.help:
-        #     return f"{item.type},{item.value}\t{item.help}"
-
-        # return f"{item.type},{item.value}
-        if item.help:
-            formatted_help = re.sub(r"\s", " ", item.help)
-            return f"{item.value}\t{_sanitize_help_text(formatted_help)}"
-        else:
-            return f"{item.value}"
-
-    def complete(self) -> str:
-        complete_action = os.getenv("_TYPER_COMPLETE_FISH_ACTION", "")
-        args, incomplete = self.get_completion_args()
-        completions = self.get_completions(args, incomplete)
-        show_args = [self.format_completion(item) for item in completions]
-        if complete_action == "get-args":
-            if show_args:
-                return "\n".join(show_args)
-        elif complete_action == "is-args":
-            if show_args:
-                # Activate complete args (no files)
-                sys.exit(0)
-            else:
-                # Deactivate complete args (allow files)
-                sys.exit(1)
-        return ""  # pragma: no cover
-
-
-class PowerShellComplete(click.shell_completion.ShellComplete):
-    name = Shells.powershell.value
-    source_template = COMPLETION_SCRIPT_POWER_SHELL
-
-    def source_vars(self) -> Dict[str, Any]:
-        return {
-            "complete_func": self.func_name,
-            "autocomplete_var": self.complete_var,
-            "prog_name": self.prog_name,
-        }
-
-    def get_completion_args(self) -> Tuple[List[str], str]:
-        completion_args = os.getenv("_TYPER_COMPLETE_ARGS", "")
-        incomplete = os.getenv("_TYPER_COMPLETE_WORD_TO_COMPLETE", "")
-        cwords = click_split_arg_string(completion_args)
-        args = cwords[1:-1] if incomplete else cwords[1:]
-        return args, incomplete
-
-    def format_completion(self, item: click.shell_completion.CompletionItem) -> str:
-        return f"{item.value}:::{_sanitize_help_text(item.help) if item.help else ' '}"
-
-
-def completion_init() -> None:
-    click.shell_completion.add_completion_class(BashComplete, Shells.bash.value)
-    click.shell_completion.add_completion_class(ZshComplete, Shells.zsh.value)
-    click.shell_completion.add_completion_class(FishComplete, Shells.fish.value)
-    click.shell_completion.add_completion_class(
-        PowerShellComplete, Shells.powershell.value
-    )
-    click.shell_completion.add_completion_class(PowerShellComplete, Shells.pwsh.value)
diff --git a/venv/lib/python3.10/site-packages/typer/_completion_shared.py b/venv/lib/python3.10/site-packages/typer/_completion_shared.py
deleted file mode 100644
index cc0add992c722c6044342d912f2772aedca86538..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer/_completion_shared.py
+++ /dev/null
@@ -1,240 +0,0 @@
-import os
-import re
-import subprocess
-from enum import Enum
-from pathlib import Path
-from typing import Optional, Tuple
-
-import click
-
-try:
-    import shellingham
-except ImportError:  # pragma: no cover
-    shellingham = None
-
-
-class Shells(str, Enum):
-    bash = "bash"
-    zsh = "zsh"
-    fish = "fish"
-    powershell = "powershell"
-    pwsh = "pwsh"
-
-
-COMPLETION_SCRIPT_BASH = """
-%(complete_func)s() {
-    local IFS=$'\n'
-    COMPREPLY=( $( env COMP_WORDS="${COMP_WORDS[*]}" \\
-                   COMP_CWORD=$COMP_CWORD \\
-                   %(autocomplete_var)s=complete_bash $1 ) )
-    return 0
-}
-
-complete -o default -F %(complete_func)s %(prog_name)s
-"""
-
-COMPLETION_SCRIPT_ZSH = """
-#compdef %(prog_name)s
-
-%(complete_func)s() {
-  eval $(env _TYPER_COMPLETE_ARGS="${words[1,$CURRENT]}" %(autocomplete_var)s=complete_zsh %(prog_name)s)
-}
-
-compdef %(complete_func)s %(prog_name)s
-"""
-
-COMPLETION_SCRIPT_FISH = 'complete --command %(prog_name)s --no-files --arguments "(env %(autocomplete_var)s=complete_fish _TYPER_COMPLETE_FISH_ACTION=get-args _TYPER_COMPLETE_ARGS=(commandline -cp) %(prog_name)s)" --condition "env %(autocomplete_var)s=complete_fish _TYPER_COMPLETE_FISH_ACTION=is-args _TYPER_COMPLETE_ARGS=(commandline -cp) %(prog_name)s"'
-
-COMPLETION_SCRIPT_POWER_SHELL = """
-Import-Module PSReadLine
-Set-PSReadLineKeyHandler -Chord Tab -Function MenuComplete
-$scriptblock = {
-    param($wordToComplete, $commandAst, $cursorPosition)
-    $Env:%(autocomplete_var)s = "complete_powershell"
-    $Env:_TYPER_COMPLETE_ARGS = $commandAst.ToString()
-    $Env:_TYPER_COMPLETE_WORD_TO_COMPLETE = $wordToComplete
-    %(prog_name)s | ForEach-Object {
-        $commandArray = $_ -Split ":::"
-        $command = $commandArray[0]
-        $helpString = $commandArray[1]
-        [System.Management.Automation.CompletionResult]::new(
-            $command, $command, 'ParameterValue', $helpString)
-    }
-    $Env:%(autocomplete_var)s = ""
-    $Env:_TYPER_COMPLETE_ARGS = ""
-    $Env:_TYPER_COMPLETE_WORD_TO_COMPLETE = ""
-}
-Register-ArgumentCompleter -Native -CommandName %(prog_name)s -ScriptBlock $scriptblock
-"""
-
-_completion_scripts = {
-    "bash": COMPLETION_SCRIPT_BASH,
-    "zsh": COMPLETION_SCRIPT_ZSH,
-    "fish": COMPLETION_SCRIPT_FISH,
-    "powershell": COMPLETION_SCRIPT_POWER_SHELL,
-    "pwsh": COMPLETION_SCRIPT_POWER_SHELL,
-}
-
-# TODO: Probably refactor this, copied from Click 7.x
-_invalid_ident_char_re = re.compile(r"[^a-zA-Z0-9_]")
-
-
-def get_completion_script(*, prog_name: str, complete_var: str, shell: str) -> str:
-    cf_name = _invalid_ident_char_re.sub("", prog_name.replace("-", "_"))
-    script = _completion_scripts.get(shell)
-    if script is None:
-        click.echo(f"Shell {shell} not supported.", err=True)
-        raise click.exceptions.Exit(1)
-    return (
-        script
-        % {
-            "complete_func": f"_{cf_name}_completion",
-            "prog_name": prog_name,
-            "autocomplete_var": complete_var,
-        }
-    ).strip()
-
-
-def install_bash(*, prog_name: str, complete_var: str, shell: str) -> Path:
-    # Ref: https://github.com/scop/bash-completion#faq
-    # It seems bash-completion is the official completion system for bash:
-    # Ref: https://www.gnu.org/software/bash/manual/html_node/A-Programmable-Completion-Example.html
-    # But installing in the locations from the docs doesn't seem to have effect
-    completion_path = Path.home() / ".bash_completions" / f"{prog_name}.sh"
-    rc_path = Path.home() / ".bashrc"
-    rc_path.parent.mkdir(parents=True, exist_ok=True)
-    rc_content = ""
-    if rc_path.is_file():
-        rc_content = rc_path.read_text()
-    completion_init_lines = [f"source '{completion_path}'"]
-    for line in completion_init_lines:
-        if line not in rc_content:  # pragma: no cover
-            rc_content += f"\n{line}"
-    rc_content += "\n"
-    rc_path.write_text(rc_content)
-    # Install completion
-    completion_path.parent.mkdir(parents=True, exist_ok=True)
-    script_content = get_completion_script(
-        prog_name=prog_name, complete_var=complete_var, shell=shell
-    )
-    completion_path.write_text(script_content)
-    return completion_path
-
-
-def install_zsh(*, prog_name: str, complete_var: str, shell: str) -> Path:
-    # Setup Zsh and load ~/.zfunc
-    zshrc_path = Path.home() / ".zshrc"
-    zshrc_path.parent.mkdir(parents=True, exist_ok=True)
-    zshrc_content = ""
-    if zshrc_path.is_file():
-        zshrc_content = zshrc_path.read_text()
-    completion_line = "fpath+=~/.zfunc; autoload -Uz compinit; compinit"
-    if completion_line not in zshrc_content:
-        zshrc_content += f"\n{completion_line}\n"
-    style_line = "zstyle ':completion:*' menu select"
-    # TODO: consider setting the style only for the current program
-    # style_line = f"zstyle ':completion:*:*:{prog_name}:*' menu select"
-    # Install zstyle completion config only if the user doesn't have a customization
-    if "zstyle" not in zshrc_content:
-        zshrc_content += f"\n{style_line}\n"
-    zshrc_content = f"{zshrc_content.strip()}\n"
-    zshrc_path.write_text(zshrc_content)
-    # Install completion under ~/.zfunc/
-    path_obj = Path.home() / f".zfunc/_{prog_name}"
-    path_obj.parent.mkdir(parents=True, exist_ok=True)
-    script_content = get_completion_script(
-        prog_name=prog_name, complete_var=complete_var, shell=shell
-    )
-    path_obj.write_text(script_content)
-    return path_obj
-
-
-def install_fish(*, prog_name: str, complete_var: str, shell: str) -> Path:
-    path_obj = Path.home() / f".config/fish/completions/{prog_name}.fish"
-    parent_dir: Path = path_obj.parent
-    parent_dir.mkdir(parents=True, exist_ok=True)
-    script_content = get_completion_script(
-        prog_name=prog_name, complete_var=complete_var, shell=shell
-    )
-    path_obj.write_text(f"{script_content}\n")
-    return path_obj
-
-
-def install_powershell(*, prog_name: str, complete_var: str, shell: str) -> Path:
-    subprocess.run(
-        [
-            shell,
-            "-Command",
-            "Set-ExecutionPolicy",
-            "Unrestricted",
-            "-Scope",
-            "CurrentUser",
-        ]
-    )
-    result = subprocess.run(
-        [shell, "-NoProfile", "-Command", "echo", "$profile"],
-        check=True,
-        stdout=subprocess.PIPE,
-    )
-    if result.returncode != 0:  # pragma: no cover
-        click.echo("Couldn't get PowerShell user profile", err=True)
-        raise click.exceptions.Exit(result.returncode)
-    path_str = ""
-    if isinstance(result.stdout, str):  # pragma: no cover
-        path_str = result.stdout
-    if isinstance(result.stdout, bytes):
-        for encoding in ["windows-1252", "utf8", "cp850"]:
-            try:
-                path_str = result.stdout.decode(encoding)
-                break
-            except UnicodeDecodeError:  # pragma: no cover
-                pass
-        if not path_str:  # pragma: no cover
-            click.echo("Couldn't decode the path automatically", err=True)
-            raise click.exceptions.Exit(1)
-    path_obj = Path(path_str.strip())
-    parent_dir: Path = path_obj.parent
-    parent_dir.mkdir(parents=True, exist_ok=True)
-    script_content = get_completion_script(
-        prog_name=prog_name, complete_var=complete_var, shell=shell
-    )
-    with path_obj.open(mode="a") as f:
-        f.write(f"{script_content}\n")
-    return path_obj
-
-
-def install(
-    shell: Optional[str] = None,
-    prog_name: Optional[str] = None,
-    complete_var: Optional[str] = None,
-) -> Tuple[str, Path]:
-    prog_name = prog_name or click.get_current_context().find_root().info_name
-    assert prog_name
-    if complete_var is None:
-        complete_var = "_{}_COMPLETE".format(prog_name.replace("-", "_").upper())
-    test_disable_detection = os.getenv("_TYPER_COMPLETE_TEST_DISABLE_SHELL_DETECTION")
-    if shell is None and shellingham is not None and not test_disable_detection:
-        shell, _ = shellingham.detect_shell()
-    if shell == "bash":
-        installed_path = install_bash(
-            prog_name=prog_name, complete_var=complete_var, shell=shell
-        )
-        return shell, installed_path
-    elif shell == "zsh":
-        installed_path = install_zsh(
-            prog_name=prog_name, complete_var=complete_var, shell=shell
-        )
-        return shell, installed_path
-    elif shell == "fish":
-        installed_path = install_fish(
-            prog_name=prog_name, complete_var=complete_var, shell=shell
-        )
-        return shell, installed_path
-    elif shell in {"powershell", "pwsh"}:
-        installed_path = install_powershell(
-            prog_name=prog_name, complete_var=complete_var, shell=shell
-        )
-        return shell, installed_path
-    else:
-        click.echo(f"Shell {shell} is not supported.")
-        raise click.exceptions.Exit(1)
diff --git a/venv/lib/python3.10/site-packages/typer/_types.py b/venv/lib/python3.10/site-packages/typer/_types.py
deleted file mode 100644
index 045e36b815652debaeea1804b1e334218422dcb4..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer/_types.py
+++ /dev/null
@@ -1,27 +0,0 @@
-from enum import Enum
-from typing import Generic, TypeVar, Union
-
-import click
-
-ParamTypeValue = TypeVar("ParamTypeValue")
-
-
-class TyperChoice(click.Choice, Generic[ParamTypeValue]):  # type: ignore[type-arg]
-    def normalize_choice(
-        self, choice: ParamTypeValue, ctx: Union[click.Context, None]
-    ) -> str:
-        # Click 8.2.0 added a new method `normalize_choice` to the `Choice` class
-        # to support enums, but it uses the enum names, while Typer has always used the
-        # enum values.
-        # This class overrides that method to maintain the previous behavior.
-        # In Click:
-        # normed_value = choice.name if isinstance(choice, Enum) else str(choice)
-        normed_value = str(choice.value) if isinstance(choice, Enum) else str(choice)
-
-        if ctx is not None and ctx.token_normalize_func is not None:
-            normed_value = ctx.token_normalize_func(normed_value)
-
-        if not self.case_sensitive:
-            normed_value = normed_value.casefold()
-
-        return normed_value
diff --git a/venv/lib/python3.10/site-packages/typer/_typing.py b/venv/lib/python3.10/site-packages/typer/_typing.py
deleted file mode 100644
index c4b3f7c8dd6d8692f41d9d28dce3fec2627617b6..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer/_typing.py
+++ /dev/null
@@ -1,113 +0,0 @@
-# Copied from pydantic 1.9.2 (the latest version to support python 3.6.)
-# https://github.com/pydantic/pydantic/blob/v1.9.2/pydantic/typing.py
-# Reduced drastically to only include Typer-specific 3.7+ functionality
-# mypy: ignore-errors
-
-import sys
-from typing import (
-    Any,
-    Callable,
-    Optional,
-    Tuple,
-    Type,
-    Union,
-)
-
-if sys.version_info >= (3, 9):
-    from typing import Annotated, Literal, get_args, get_origin, get_type_hints
-else:
-    from typing_extensions import (
-        Annotated,
-        Literal,
-        get_args,
-        get_origin,
-        get_type_hints,
-    )
-
-if sys.version_info < (3, 10):
-
-    def is_union(tp: Optional[Type[Any]]) -> bool:
-        return tp is Union
-
-else:
-    import types
-
-    def is_union(tp: Optional[Type[Any]]) -> bool:
-        return tp is Union or tp is types.UnionType  # noqa: E721
-
-
-__all__ = (
-    "NoneType",
-    "is_none_type",
-    "is_callable_type",
-    "is_literal_type",
-    "all_literal_values",
-    "is_union",
-    "Annotated",
-    "Literal",
-    "get_args",
-    "get_origin",
-    "get_type_hints",
-)
-
-
-NoneType = None.__class__
-
-
-NONE_TYPES: Tuple[Any, Any, Any] = (None, NoneType, Literal[None])
-
-
-if sys.version_info < (3, 8):
-    # Even though this implementation is slower, we need it for python 3.7:
-    # In python 3.7 "Literal" is not a builtin type and uses a different
-    # mechanism.
-    # for this reason `Literal[None] is Literal[None]` evaluates to `False`,
-    # breaking the faster implementation used for the other python versions.
-
-    def is_none_type(type_: Any) -> bool:
-        return type_ in NONE_TYPES
-
-elif sys.version_info[:2] == (3, 8):
-    # We can use the fast implementation for 3.8 but there is a very weird bug
-    # where it can fail for `Literal[None]`.
-    # We just need to redefine a useless `Literal[None]` inside the function body to fix this
-
-    def is_none_type(type_: Any) -> bool:
-        Literal[None]  # fix edge case
-        for none_type in NONE_TYPES:
-            if type_ is none_type:
-                return True
-        return False
-
-else:
-
-    def is_none_type(type_: Any) -> bool:
-        for none_type in NONE_TYPES:
-            if type_ is none_type:
-                return True
-        return False
-
-
-def is_callable_type(type_: Type[Any]) -> bool:
-    return type_ is Callable or get_origin(type_) is Callable
-
-
-def is_literal_type(type_: Type[Any]) -> bool:
-    return Literal is not None and get_origin(type_) is Literal
-
-
-def literal_values(type_: Type[Any]) -> Tuple[Any, ...]:
-    return get_args(type_)
-
-
-def all_literal_values(type_: Type[Any]) -> Tuple[Any, ...]:
-    """
-    This method is used to retrieve all Literal values as
-    Literal can be used recursively (see https://www.python.org/dev/peps/pep-0586)
-    e.g. `Literal[Literal[Literal[1, 2, 3], "foo"], 5, None]`
-    """
-    if not is_literal_type(type_):
-        return (type_,)
-
-    values = literal_values(type_)
-    return tuple(x for value in values for x in all_literal_values(value))
diff --git a/venv/lib/python3.10/site-packages/typer/cli.py b/venv/lib/python3.10/site-packages/typer/cli.py
deleted file mode 100644
index 3fe3d3ee7fbd6fedeade47ceb22a83cf258f7945..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer/cli.py
+++ /dev/null
@@ -1,315 +0,0 @@
-import importlib.util
-import re
-import sys
-from pathlib import Path
-from typing import Any, List, Optional
-
-import click
-import typer
-import typer.core
-from click import Command, Group, Option
-
-from . import __version__
-
-try:
-    import rich
-
-    has_rich = True
-    from . import rich_utils
-
-except ImportError:  # pragma: no cover
-    has_rich = False
-    rich = None  # type: ignore
-
-default_app_names = ("app", "cli", "main")
-default_func_names = ("main", "cli", "app")
-
-app = typer.Typer()
-utils_app = typer.Typer(help="Extra utility commands for Typer apps.")
-app.add_typer(utils_app, name="utils")
-
-
-class State:
-    def __init__(self) -> None:
-        self.app: Optional[str] = None
-        self.func: Optional[str] = None
-        self.file: Optional[Path] = None
-        self.module: Optional[str] = None
-
-
-state = State()
-
-
-def maybe_update_state(ctx: click.Context) -> None:
-    path_or_module = ctx.params.get("path_or_module")
-    if path_or_module:
-        file_path = Path(path_or_module)
-        if file_path.exists() and file_path.is_file():
-            state.file = file_path
-        else:
-            if not re.fullmatch(r"[a-zA-Z_]\w*(\.[a-zA-Z_]\w*)*", path_or_module):
-                typer.echo(
-                    f"Not a valid file or Python module: {path_or_module}", err=True
-                )
-                sys.exit(1)
-            state.module = path_or_module
-    app_name = ctx.params.get("app")
-    if app_name:
-        state.app = app_name
-    func_name = ctx.params.get("func")
-    if func_name:
-        state.func = func_name
-
-
-class TyperCLIGroup(typer.core.TyperGroup):
-    def list_commands(self, ctx: click.Context) -> List[str]:
-        self.maybe_add_run(ctx)
-        return super().list_commands(ctx)
-
-    def get_command(self, ctx: click.Context, name: str) -> Optional[Command]:
-        self.maybe_add_run(ctx)
-        return super().get_command(ctx, name)
-
-    def invoke(self, ctx: click.Context) -> Any:
-        self.maybe_add_run(ctx)
-        return super().invoke(ctx)
-
-    def maybe_add_run(self, ctx: click.Context) -> None:
-        maybe_update_state(ctx)
-        maybe_add_run_to_cli(self)
-
-
-def get_typer_from_module(module: Any) -> Optional[typer.Typer]:
-    # Try to get defined app
-    if state.app:
-        obj = getattr(module, state.app, None)
-        if not isinstance(obj, typer.Typer):
-            typer.echo(f"Not a Typer object: --app {state.app}", err=True)
-            sys.exit(1)
-        return obj
-    # Try to get defined function
-    if state.func:
-        func_obj = getattr(module, state.func, None)
-        if not callable(func_obj):
-            typer.echo(f"Not a function: --func {state.func}", err=True)
-            sys.exit(1)
-        sub_app = typer.Typer()
-        sub_app.command()(func_obj)
-        return sub_app
-    # Iterate and get a default object to use as CLI
-    local_names = dir(module)
-    local_names_set = set(local_names)
-    # Try to get a default Typer app
-    for name in default_app_names:
-        if name in local_names_set:
-            obj = getattr(module, name, None)
-            if isinstance(obj, typer.Typer):
-                return obj
-    # Try to get any Typer app
-    for name in local_names_set - set(default_app_names):
-        obj = getattr(module, name)
-        if isinstance(obj, typer.Typer):
-            return obj
-    # Try to get a default function
-    for func_name in default_func_names:
-        func_obj = getattr(module, func_name, None)
-        if callable(func_obj):
-            sub_app = typer.Typer()
-            sub_app.command()(func_obj)
-            return sub_app
-    # Try to get any func app
-    for func_name in local_names_set - set(default_func_names):
-        func_obj = getattr(module, func_name)
-        if callable(func_obj):
-            sub_app = typer.Typer()
-            sub_app.command()(func_obj)
-            return sub_app
-    return None
-
-
-def get_typer_from_state() -> Optional[typer.Typer]:
-    spec = None
-    if state.file:
-        module_name = state.file.name
-        spec = importlib.util.spec_from_file_location(module_name, str(state.file))
-    elif state.module:
-        spec = importlib.util.find_spec(state.module)
-    if spec is None:
-        if state.file:
-            typer.echo(f"Could not import as Python file: {state.file}", err=True)
-        else:
-            typer.echo(f"Could not import as Python module: {state.module}", err=True)
-        sys.exit(1)
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)  # type: ignore
-    obj = get_typer_from_module(module)
-    return obj
-
-
-def maybe_add_run_to_cli(cli: click.Group) -> None:
-    if "run" not in cli.commands:
-        if state.file or state.module:
-            obj = get_typer_from_state()
-            if obj:
-                obj._add_completion = False
-                click_obj = typer.main.get_command(obj)
-                click_obj.name = "run"
-                if not click_obj.help:
-                    click_obj.help = "Run the provided Typer app."
-                cli.add_command(click_obj)
-
-
-def print_version(ctx: click.Context, param: Option, value: bool) -> None:
-    if not value or ctx.resilient_parsing:
-        return
-    typer.echo(f"Typer version: {__version__}")
-    raise typer.Exit()
-
-
-@app.callback(cls=TyperCLIGroup, no_args_is_help=True)
-def callback(
-    ctx: typer.Context,
-    *,
-    path_or_module: str = typer.Argument(None),
-    app: str = typer.Option(None, help="The typer app object/variable to use."),
-    func: str = typer.Option(None, help="The function to convert to Typer."),
-    version: bool = typer.Option(
-        False,
-        "--version",
-        help="Print version and exit.",
-        callback=print_version,
-    ),
-) -> None:
-    """
-    Run Typer scripts with completion, without having to create a package.
-
-    You probably want to install completion for the typer command:
-
-    $ typer --install-completion
-
-    https://typer.tiangolo.com/
-    """
-    maybe_update_state(ctx)
-
-
-def get_docs_for_click(
-    *,
-    obj: Command,
-    ctx: typer.Context,
-    indent: int = 0,
-    name: str = "",
-    call_prefix: str = "",
-    title: Optional[str] = None,
-) -> str:
-    docs = "#" * (1 + indent)
-    command_name = name or obj.name
-    if call_prefix:
-        command_name = f"{call_prefix} {command_name}"
-    if not title:
-        title = f"`{command_name}`" if command_name else "CLI"
-    docs += f" {title}\n\n"
-    if obj.help:
-        docs += f"{_parse_html(obj.help)}\n\n"
-    usage_pieces = obj.collect_usage_pieces(ctx)
-    if usage_pieces:
-        docs += "**Usage**:\n\n"
-        docs += "```console\n"
-        docs += "$ "
-        if command_name:
-            docs += f"{command_name} "
-        docs += f"{' '.join(usage_pieces)}\n"
-        docs += "```\n\n"
-    args = []
-    opts = []
-    for param in obj.get_params(ctx):
-        rv = param.get_help_record(ctx)
-        if rv is not None:
-            if param.param_type_name == "argument":
-                args.append(rv)
-            elif param.param_type_name == "option":
-                opts.append(rv)
-    if args:
-        docs += "**Arguments**:\n\n"
-        for arg_name, arg_help in args:
-            docs += f"* `{arg_name}`"
-            if arg_help:
-                docs += f": {_parse_html(arg_help)}"
-            docs += "\n"
-        docs += "\n"
-    if opts:
-        docs += "**Options**:\n\n"
-        for opt_name, opt_help in opts:
-            docs += f"* `{opt_name}`"
-            if opt_help:
-                docs += f": {_parse_html(opt_help)}"
-            docs += "\n"
-        docs += "\n"
-    if obj.epilog:
-        docs += f"{obj.epilog}\n\n"
-    if isinstance(obj, Group):
-        group = obj
-        commands = group.list_commands(ctx)
-        if commands:
-            docs += "**Commands**:\n\n"
-            for command in commands:
-                command_obj = group.get_command(ctx, command)
-                assert command_obj
-                docs += f"* `{command_obj.name}`"
-                command_help = command_obj.get_short_help_str()
-                if command_help:
-                    docs += f": {_parse_html(command_help)}"
-                docs += "\n"
-            docs += "\n"
-        for command in commands:
-            command_obj = group.get_command(ctx, command)
-            assert command_obj
-            use_prefix = ""
-            if command_name:
-                use_prefix += f"{command_name}"
-            docs += get_docs_for_click(
-                obj=command_obj, ctx=ctx, indent=indent + 1, call_prefix=use_prefix
-            )
-    return docs
-
-
-def _parse_html(input_text: str) -> str:
-    if not has_rich:  # pragma: no cover
-        return input_text
-    return rich_utils.rich_to_html(input_text)
-
-
-@utils_app.command()
-def docs(
-    ctx: typer.Context,
-    name: str = typer.Option("", help="The name of the CLI program to use in docs."),
-    output: Optional[Path] = typer.Option(
-        None,
-        help="An output file to write docs to, like README.md.",
-        file_okay=True,
-        dir_okay=False,
-    ),
-    title: Optional[str] = typer.Option(
-        None,
-        help="The title for the documentation page. If not provided, the name of "
-        "the program is used.",
-    ),
-) -> None:
-    """
-    Generate Markdown docs for a Typer app.
-    """
-    typer_obj = get_typer_from_state()
-    if not typer_obj:
-        typer.echo("No Typer app found", err=True)
-        raise typer.Abort()
-    click_obj = typer.main.get_command(typer_obj)
-    docs = get_docs_for_click(obj=click_obj, ctx=ctx, name=name, title=title)
-    clean_docs = f"{docs.strip()}\n"
-    if output:
-        output.write_text(clean_docs)
-        typer.echo(f"Docs saved to: {output}")
-    else:
-        typer.echo(clean_docs)
-
-
-def main() -> Any:
-    return app()
diff --git a/venv/lib/python3.10/site-packages/typer/colors.py b/venv/lib/python3.10/site-packages/typer/colors.py
deleted file mode 100644
index 54e7b166cb1de83321a4965cc4915824b47a7f4f..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer/colors.py
+++ /dev/null
@@ -1,20 +0,0 @@
-# Variable names to colors, just for completion
-BLACK = "black"
-RED = "red"
-GREEN = "green"
-YELLOW = "yellow"
-BLUE = "blue"
-MAGENTA = "magenta"
-CYAN = "cyan"
-WHITE = "white"
-
-RESET = "reset"
-
-BRIGHT_BLACK = "bright_black"
-BRIGHT_RED = "bright_red"
-BRIGHT_GREEN = "bright_green"
-BRIGHT_YELLOW = "bright_yellow"
-BRIGHT_BLUE = "bright_blue"
-BRIGHT_MAGENTA = "bright_magenta"
-BRIGHT_CYAN = "bright_cyan"
-BRIGHT_WHITE = "bright_white"
diff --git a/venv/lib/python3.10/site-packages/typer/completion.py b/venv/lib/python3.10/site-packages/typer/completion.py
deleted file mode 100644
index c355baa78182a6e7a6ba692db660fca162ef79f9..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer/completion.py
+++ /dev/null
@@ -1,149 +0,0 @@
-import os
-import sys
-from typing import Any, MutableMapping, Tuple
-
-import click
-
-from ._completion_classes import completion_init
-from ._completion_shared import Shells, get_completion_script, install
-from .models import ParamMeta
-from .params import Option
-from .utils import get_params_from_function
-
-try:
-    import shellingham
-except ImportError:  # pragma: no cover
-    shellingham = None
-
-
-_click_patched = False
-
-
-def get_completion_inspect_parameters() -> Tuple[ParamMeta, ParamMeta]:
-    completion_init()
-    test_disable_detection = os.getenv("_TYPER_COMPLETE_TEST_DISABLE_SHELL_DETECTION")
-    if shellingham and not test_disable_detection:
-        parameters = get_params_from_function(_install_completion_placeholder_function)
-    else:
-        parameters = get_params_from_function(
-            _install_completion_no_auto_placeholder_function
-        )
-    install_param, show_param = parameters.values()
-    return install_param, show_param
-
-
-def install_callback(ctx: click.Context, param: click.Parameter, value: Any) -> Any:
-    if not value or ctx.resilient_parsing:
-        return value  # pragma: no cover
-    if isinstance(value, str):
-        shell, path = install(shell=value)
-    else:
-        shell, path = install()
-    click.secho(f"{shell} completion installed in {path}", fg="green")
-    click.echo("Completion will take effect once you restart the terminal")
-    sys.exit(0)
-
-
-def show_callback(ctx: click.Context, param: click.Parameter, value: Any) -> Any:
-    if not value or ctx.resilient_parsing:
-        return value  # pragma: no cover
-    prog_name = ctx.find_root().info_name
-    assert prog_name
-    complete_var = "_{}_COMPLETE".format(prog_name.replace("-", "_").upper())
-    shell = ""
-    test_disable_detection = os.getenv("_TYPER_COMPLETE_TEST_DISABLE_SHELL_DETECTION")
-    if isinstance(value, str):
-        shell = value
-    elif shellingham and not test_disable_detection:
-        shell, _ = shellingham.detect_shell()
-    script_content = get_completion_script(
-        prog_name=prog_name, complete_var=complete_var, shell=shell
-    )
-    click.echo(script_content)
-    sys.exit(0)
-
-
-# Create a fake command function to extract the completion parameters
-def _install_completion_placeholder_function(
-    install_completion: bool = Option(
-        None,
-        "--install-completion",
-        callback=install_callback,
-        expose_value=False,
-        help="Install completion for the current shell.",
-    ),
-    show_completion: bool = Option(
-        None,
-        "--show-completion",
-        callback=show_callback,
-        expose_value=False,
-        help="Show completion for the current shell, to copy it or customize the installation.",
-    ),
-) -> Any:
-    pass  # pragma: no cover
-
-
-def _install_completion_no_auto_placeholder_function(
-    install_completion: Shells = Option(
-        None,
-        callback=install_callback,
-        expose_value=False,
-        help="Install completion for the specified shell.",
-    ),
-    show_completion: Shells = Option(
-        None,
-        callback=show_callback,
-        expose_value=False,
-        help="Show completion for the specified shell, to copy it or customize the installation.",
-    ),
-) -> Any:
-    pass  # pragma: no cover
-
-
-# Re-implement Click's shell_complete to add error message with:
-# Invalid completion instruction
-# To use 7.x instruction style for compatibility
-# And to add extra error messages, for compatibility with Typer in previous versions
-# This is only called in new Command method, only used by Click 8.x+
-def shell_complete(
-    cli: click.Command,
-    ctx_args: MutableMapping[str, Any],
-    prog_name: str,
-    complete_var: str,
-    instruction: str,
-) -> int:
-    import click
-    import click.shell_completion
-
-    if "_" not in instruction:
-        click.echo("Invalid completion instruction.", err=True)
-        return 1
-
-    # Click 8 changed the order/style of shell instructions from e.g.
-    # source_bash to bash_source
-    # Typer override to preserve the old style for compatibility
-    # Original in Click 8.x commented:
-    # shell, _, instruction = instruction.partition("_")
-    instruction, _, shell = instruction.partition("_")
-    # Typer override end
-
-    comp_cls = click.shell_completion.get_completion_class(shell)
-
-    if comp_cls is None:
-        click.echo(f"Shell {shell} not supported.", err=True)
-        return 1
-
-    comp = comp_cls(cli, ctx_args, prog_name, complete_var)
-
-    if instruction == "source":
-        click.echo(comp.source())
-        return 0
-
-    # Typer override to print the completion help msg with Rich
-    if instruction == "complete":
-        click.echo(comp.complete())
-        return 0
-    # Typer override end
-
-    click.echo(f'Completion instruction "{instruction}" not supported.', err=True)
-    return 1
diff --git a/venv/lib/python3.10/site-packages/typer/core.py b/venv/lib/python3.10/site-packages/typer/core.py
deleted file mode 100644
index f6c4f72e8a5bd2ab8d117f87895b9d1200cb6e9a..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer/core.py
+++ /dev/null
@@ -1,781 +0,0 @@
-import errno
-import inspect
-import os
-import sys
-from enum import Enum
-from gettext import gettext as _
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    List,
-    MutableMapping,
-    Optional,
-    Sequence,
-    TextIO,
-    Tuple,
-    Union,
-    cast,
-)
-
-import click
-import click.core
-import click.formatting
-import click.parser
-import click.shell_completion
-import click.types
-import click.utils
-
-from ._typing import Literal
-
-MarkupMode = Literal["markdown", "rich", None]
-
-try:
-    import rich
-
-    from . import rich_utils
-
-    DEFAULT_MARKUP_MODE: MarkupMode = "rich"
-
-except ImportError:  # pragma: no cover
-    rich = None  # type: ignore
-    DEFAULT_MARKUP_MODE = None
-
-
-# Copy from click.parser._split_opt
-def _split_opt(opt: str) -> Tuple[str, str]:
-    first = opt[:1]
-    if first.isalnum():
-        return "", opt
-    if opt[1:2] == first:
-        return opt[:2], opt[2:]
-    return first, opt[1:]
-
-
-def _typer_param_setup_autocompletion_compat(
-    self: click.Parameter,
-    *,
-    autocompletion: Optional[
-        Callable[[click.Context, List[str], str], List[Union[Tuple[str, str], str]]]
-    ] = None,
-) -> None:
-    if self._custom_shell_complete is not None:
-        import warnings
-
-        warnings.warn(
-            "In Typer, only the parameter 'autocompletion' is supported. "
-            "The support for 'shell_complete' is deprecated and will be removed in upcoming versions. ",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-
-    if autocompletion is not None:
-
-        def compat_autocompletion(
-            ctx: click.Context, param: click.core.Parameter, incomplete: str
-        ) -> List["click.shell_completion.CompletionItem"]:
-            from click.shell_completion import CompletionItem
-
-            out = []
-
-            for c in autocompletion(ctx, [], incomplete):
-                if isinstance(c, tuple):
-                    use_completion = CompletionItem(c[0], help=c[1])
-                else:
-                    assert isinstance(c, str)
-                    use_completion = CompletionItem(c)
-
-                if use_completion.value.startswith(incomplete):
-                    out.append(use_completion)
-
-            return out
-
-        self._custom_shell_complete = compat_autocompletion
-
-
-def _get_default_string(
-    obj: Union["TyperArgument", "TyperOption"],
-    *,
-    ctx: click.Context,
-    show_default_is_str: bool,
-    default_value: Union[List[Any], Tuple[Any, ...], str, Callable[..., Any], Any],
-) -> str:
-    # Extracted from click.core.Option.get_help_record() to be reused by
-    # rich_utils avoiding RegEx hacks
-    if show_default_is_str:
-        default_string = f"({obj.show_default})"
-    elif isinstance(default_value, (list, tuple)):
-        default_string = ", ".join(
-            _get_default_string(
-                obj, ctx=ctx, show_default_is_str=show_default_is_str, default_value=d
-            )
-            for d in default_value
-        )
-    elif isinstance(default_value, Enum):
-        default_string = str(default_value.value)
-    elif inspect.isfunction(default_value):
-        default_string = _("(dynamic)")
-    elif isinstance(obj, TyperOption) and obj.is_bool_flag and obj.secondary_opts:
-        # For boolean flags that have distinct True/False opts,
-        # use the opt without prefix instead of the value.
-        # Typer override, original commented
-        # default_string = click.parser.split_opt(
-        #     (self.opts if self.default else self.secondary_opts)[0]
-        # )[1]
-        if obj.default:
-            if obj.opts:
-                default_string = _split_opt(obj.opts[0])[1]
-            else:
-                default_string = str(default_value)
-        else:
-            default_string = _split_opt(obj.secondary_opts[0])[1]
-        # Typer override end
-    elif (
-        isinstance(obj, TyperOption)
-        and obj.is_bool_flag
-        and not obj.secondary_opts
-        and not default_value
-    ):
-        default_string = ""
-    else:
-        default_string = str(default_value)
-    return default_string
-
-
-def _extract_default_help_str(
-    obj: Union["TyperArgument", "TyperOption"], *, ctx: click.Context
-) -> Optional[Union[Any, Callable[[], Any]]]:
-    # Extracted from click.core.Option.get_help_record() to be reused by
-    # rich_utils avoiding RegEx hacks
-    # Temporarily enable resilient parsing to avoid type casting
-    # failing for the default. Might be possible to extend this to
-    # help formatting in general.
-    resilient = ctx.resilient_parsing
-    ctx.resilient_parsing = True
-
-    try:
-        default_value = obj.get_default(ctx, call=False)
-    finally:
-        ctx.resilient_parsing = resilient
-    return default_value
-
-
-def _main(
-    self: click.Command,
-    *,
-    args: Optional[Sequence[str]] = None,
-    prog_name: Optional[str] = None,
-    complete_var: Optional[str] = None,
-    standalone_mode: bool = True,
-    windows_expand_args: bool = True,
-    rich_markup_mode: MarkupMode = DEFAULT_MARKUP_MODE,
-    **extra: Any,
-) -> Any:
-    # Typer override, duplicated from click.main() to handle custom rich exceptions
-    # Verify that the environment is configured correctly, or reject
-    # further execution to avoid a broken script.
-    if args is None:
-        args = sys.argv[1:]
-
-        # Covered in Click tests
-        if os.name == "nt" and windows_expand_args:  # pragma: no cover
-            args = click.utils._expand_args(args)
-    else:
-        args = list(args)
-
-    if prog_name is None:
-        prog_name = click.utils._detect_program_name()
-
-    # Process shell completion requests and exit early.
-    self._main_shell_completion(extra, prog_name, complete_var)
-
-    try:
-        try:
-            with self.make_context(prog_name, args, **extra) as ctx:
-                rv = self.invoke(ctx)
-                if not standalone_mode:
-                    return rv
-                # it's not safe to `ctx.exit(rv)` here!
-                # note that `rv` may actually contain data like "1" which
-                # has obvious effects
-                # more subtle case: `rv=[None, None]` can come out of
-                # chained commands which all returned `None` -- so it's not
-                # even always obvious that `rv` indicates success/failure
-                # by its truthiness/falsiness
-                ctx.exit()
-        except EOFError as e:
-            click.echo(file=sys.stderr)
-            raise click.Abort() from e
-        except KeyboardInterrupt as e:
-            raise click.exceptions.Exit(130) from e
-        except click.ClickException as e:
-            if not standalone_mode:
-                raise
-            # Typer override
-            if rich and rich_markup_mode is not None:
-                rich_utils.rich_format_error(e)
-            else:
-                e.show()
-            # Typer override end
-            sys.exit(e.exit_code)
-        except OSError as e:
-            if e.errno == errno.EPIPE:
-                sys.stdout = cast(TextIO, click.utils.PacifyFlushWrapper(sys.stdout))
-                sys.stderr = cast(TextIO, click.utils.PacifyFlushWrapper(sys.stderr))
-                sys.exit(1)
-            else:
-                raise
-    except click.exceptions.Exit as e:
-        if standalone_mode:
-            sys.exit(e.exit_code)
-        else:
-            # in non-standalone mode, return the exit code
-            # note that this is only reached if `self.invoke` above raises
-            # an Exit explicitly -- thus bypassing the check there which
-            # would return its result
-            # the results of non-standalone execution may therefore be
-            # somewhat ambiguous: if there are codepaths which lead to
-            # `ctx.exit(1)` and to `return 1`, the caller won't be able to
-            # tell the difference between the two
-            return e.exit_code
-    except click.Abort:
-        if not standalone_mode:
-            raise
-        # Typer override
-        if rich and rich_markup_mode is not None:
-            rich_utils.rich_abort_error()
-        else:
-            click.echo(_("Aborted!"), file=sys.stderr)
-        # Typer override end
-        sys.exit(1)
-
-
-class TyperArgument(click.core.Argument):
-    def __init__(
-        self,
-        *,
-        # Parameter
-        param_decls: List[str],
-        type: Optional[Any] = None,
-        required: Optional[bool] = None,
-        default: Optional[Any] = None,
-        callback: Optional[Callable[..., Any]] = None,
-        nargs: Optional[int] = None,
-        metavar: Optional[str] = None,
-        expose_value: bool = True,
-        is_eager: bool = False,
-        envvar: Optional[Union[str, List[str]]] = None,
-        # Note that shell_complete is not fully supported and will be removed in future versions
-        # TODO: Remove shell_complete in a future version (after 0.16.0)
-        shell_complete: Optional[
-            Callable[
-                [click.Context, click.Parameter, str],
-                Union[List["click.shell_completion.CompletionItem"], List[str]],
-            ]
-        ] = None,
-        autocompletion: Optional[Callable[..., Any]] = None,
-        # TyperArgument
-        show_default: Union[bool, str] = True,
-        show_choices: bool = True,
-        show_envvar: bool = True,
-        help: Optional[str] = None,
-        hidden: bool = False,
-        # Rich settings
-        rich_help_panel: Union[str, None] = None,
-    ):
-        self.help = help
-        self.show_default = show_default
-        self.show_choices = show_choices
-        self.show_envvar = show_envvar
-        self.hidden = hidden
-        self.rich_help_panel = rich_help_panel
-
-        super().__init__(
-            param_decls=param_decls,
-            type=type,
-            required=required,
-            default=default,
-            callback=callback,
-            nargs=nargs,
-            metavar=metavar,
-            expose_value=expose_value,
-            is_eager=is_eager,
-            envvar=envvar,
-            shell_complete=shell_complete,
-        )
-        _typer_param_setup_autocompletion_compat(self, autocompletion=autocompletion)
-
-    def _get_default_string(
-        self,
-        *,
-        ctx: click.Context,
-        show_default_is_str: bool,
-        default_value: Union[List[Any], Tuple[Any, ...], str, Callable[..., Any], Any],
-    ) -> str:
-        return _get_default_string(
-            self,
-            ctx=ctx,
-            show_default_is_str=show_default_is_str,
-            default_value=default_value,
-        )
-
-    def _extract_default_help_str(
-        self, *, ctx: click.Context
-    ) -> Optional[Union[Any, Callable[[], Any]]]:
-        return _extract_default_help_str(self, ctx=ctx)
-
-    def get_help_record(self, ctx: click.Context) -> Optional[Tuple[str, str]]:
-        # Modified version of click.core.Option.get_help_record()
-        # to support Arguments
-        if self.hidden:
-            return None
-        name = self.make_metavar(ctx=ctx)
-        help = self.help or ""
-        extra = []
-        if self.show_envvar:
-            envvar = self.envvar
-            # allow_from_autoenv is currently not supported in Typer for CLI Arguments
-            if envvar is not None:
-                var_str = (
-                    ", ".join(str(d) for d in envvar)
-                    if isinstance(envvar, (list, tuple))
-                    else envvar
-                )
-                extra.append(f"env var: {var_str}")
-
-        # Typer override:
-        # Extracted to _extract_default_help_str() to allow re-using it in rich_utils
-        default_value = self._extract_default_help_str(ctx=ctx)
-        # Typer override end
-
-        show_default_is_str = isinstance(self.show_default, str)
-
-        if show_default_is_str or (
-            default_value is not None and (self.show_default or ctx.show_default)
-        ):
-            # Typer override:
-            # Extracted to _get_default_string() to allow re-using it in rich_utils
-            default_string = self._get_default_string(
-                ctx=ctx,
-                show_default_is_str=show_default_is_str,
-                default_value=default_value,
-            )
-            # Typer override end
-            if default_string:
-                extra.append(_("default: {default}").format(default=default_string))
-        if self.required:
-            extra.append(_("required"))
-        if extra:
-            extra_str = "; ".join(extra)
-            extra_str = f"[{extra_str}]"
-            if rich is not None:
-                # This is needed for when we want to export to HTML
-                extra_str = rich.markup.escape(extra_str).strip()
-
-            help = f"{help}  {extra_str}" if help else f"{extra_str}"
-        return name, help
-
-    def make_metavar(self, ctx: Union[click.Context, None] = None) -> str:
-        # Modified version of click.core.Argument.make_metavar()
-        # to include Argument name
-        if self.metavar is not None:
-            return self.metavar
-        var = (self.name or "").upper()
-        if not self.required:
-            var = f"[{var}]"
-        # TODO: When deprecating Click < 8.2, remove this
-        signature = inspect.signature(self.type.get_metavar)
-        if "ctx" in signature.parameters:
-            # Click >= 8.2
-            type_var = self.type.get_metavar(self, ctx=ctx)  # type: ignore[arg-type]
-        else:
-            # Click < 8.2
-            type_var = self.type.get_metavar(self)  # type: ignore[call-arg]
-        # TODO: /When deprecating Click < 8.2, remove this, uncomment the line below
-        # type_var = self.type.get_metavar(self, ctx=ctx)
-        if type_var:
-            var += f":{type_var}"
-        if self.nargs != 1:
-            var += "..."
-        return var
-
-
-class TyperOption(click.core.Option):
-    def __init__(
-        self,
-        *,
-        # Parameter
-        param_decls: List[str],
-        type: Optional[Union[click.types.ParamType, Any]] = None,
-        required: Optional[bool] = None,
-        default: Optional[Any] = None,
-        callback: Optional[Callable[..., Any]] = None,
-        nargs: Optional[int] = None,
-        metavar: Optional[str] = None,
-        expose_value: bool = True,
-        is_eager: bool = False,
-        envvar: Optional[Union[str, List[str]]] = None,
-        # Note that shell_complete is not fully supported and will be removed in future versions
-        # TODO: Remove shell_complete in a future version (after 0.16.0)
-        shell_complete: Optional[
-            Callable[
-                [click.Context, click.Parameter, str],
-                Union[List["click.shell_completion.CompletionItem"], List[str]],
-            ]
-        ] = None,
-        autocompletion: Optional[Callable[..., Any]] = None,
-        # Option
-        show_default: Union[bool, str] = False,
-        prompt: Union[bool, str] = False,
-        confirmation_prompt: Union[bool, str] = False,
-        prompt_required: bool = True,
-        hide_input: bool = False,
-        is_flag: Optional[bool] = None,
-        multiple: bool = False,
-        count: bool = False,
-        allow_from_autoenv: bool = True,
-        help: Optional[str] = None,
-        hidden: bool = False,
-        show_choices: bool = True,
-        show_envvar: bool = False,
-        # Rich settings
-        rich_help_panel: Union[str, None] = None,
-    ):
-        super().__init__(
-            param_decls=param_decls,
-            type=type,
-            required=required,
-            default=default,
-            callback=callback,
-            nargs=nargs,
-            metavar=metavar,
-            expose_value=expose_value,
-            is_eager=is_eager,
-            envvar=envvar,
-            show_default=show_default,
-            prompt=prompt,
-            confirmation_prompt=confirmation_prompt,
-            hide_input=hide_input,
-            is_flag=is_flag,
-            multiple=multiple,
-            count=count,
-            allow_from_autoenv=allow_from_autoenv,
-            help=help,
-            hidden=hidden,
-            show_choices=show_choices,
-            show_envvar=show_envvar,
-            prompt_required=prompt_required,
-            shell_complete=shell_complete,
-        )
-        _typer_param_setup_autocompletion_compat(self, autocompletion=autocompletion)
-        self.rich_help_panel = rich_help_panel
-
-    def _get_default_string(
-        self,
-        *,
-        ctx: click.Context,
-        show_default_is_str: bool,
-        default_value: Union[List[Any], Tuple[Any, ...], str, Callable[..., Any], Any],
-    ) -> str:
-        return _get_default_string(
-            self,
-            ctx=ctx,
-            show_default_is_str=show_default_is_str,
-            default_value=default_value,
-        )
-
-    def _extract_default_help_str(
-        self, *, ctx: click.Context
-    ) -> Optional[Union[Any, Callable[[], Any]]]:
-        return _extract_default_help_str(self, ctx=ctx)
-
-    def make_metavar(self, ctx: Union[click.Context, None] = None) -> str:
-        signature = inspect.signature(super().make_metavar)
-        if "ctx" in signature.parameters:
-            # Click >= 8.2
-            return super().make_metavar(ctx=ctx)  # type: ignore[arg-type]
-        # Click < 8.2
-        return super().make_metavar()  # type: ignore[call-arg]
-
-    def get_help_record(self, ctx: click.Context) -> Optional[Tuple[str, str]]:
-        # Duplicate all of Click's logic only to modify a single line, to allow boolean
-        # flags with only names for False values as it's currently supported by Typer
-        # Ref: https://typer.tiangolo.com/tutorial/parameter-types/bool/#only-names-for-false
-        if self.hidden:
-            return None
-
-        any_prefix_is_slash = False
-
-        def _write_opts(opts: Sequence[str]) -> str:
-            nonlocal any_prefix_is_slash
-
-            rv, any_slashes = click.formatting.join_options(opts)
-
-            if any_slashes:
-                any_prefix_is_slash = True
-
-            if not self.is_flag and not self.count:
-                rv += f" {self.make_metavar(ctx=ctx)}"
-
-            return rv
-
-        rv = [_write_opts(self.opts)]
-
-        if self.secondary_opts:
-            rv.append(_write_opts(self.secondary_opts))
-
-        help = self.help or ""
-        extra = []
-
-        if self.show_envvar:
-            envvar = self.envvar
-
-            if envvar is None:
-                if (
-                    self.allow_from_autoenv
-                    and ctx.auto_envvar_prefix is not None
-                    and self.name is not None
-                ):
-                    envvar = f"{ctx.auto_envvar_prefix}_{self.name.upper()}"
-
-            if envvar is not None:
-                var_str = (
-                    envvar
-                    if isinstance(envvar, str)
-                    else ", ".join(str(d) for d in envvar)
-                )
-                extra.append(_("env var: {var}").format(var=var_str))
-
-        # Typer override:
-        # Extracted to _extract_default() to allow re-using it in rich_utils
-        default_value = self._extract_default_help_str(ctx=ctx)
-        # Typer override end
-
-        show_default_is_str = isinstance(self.show_default, str)
-
-        if show_default_is_str or (
-            default_value is not None and (self.show_default or ctx.show_default)
-        ):
-            # Typer override:
-            # Extracted to _get_default_string() to allow re-using it in rich_utils
-            default_string = self._get_default_string(
-                ctx=ctx,
-                show_default_is_str=show_default_is_str,
-                default_value=default_value,
-            )
-            # Typer override end
-            if default_string:
-                extra.append(_("default: {default}").format(default=default_string))
-
-        if isinstance(self.type, click.types._NumberRangeBase):
-            range_str = self.type._describe_range()
-
-            if range_str:
-                extra.append(range_str)
-
-        if self.required:
-            extra.append(_("required"))
-
-        if extra:
-            extra_str = "; ".join(extra)
-            extra_str = f"[{extra_str}]"
-            if rich is not None:
-                # This is needed for when we want to export to HTML
-                extra_str = rich.markup.escape(extra_str).strip()
-
-            help = f"{help}  {extra_str}" if help else f"{extra_str}"
-
-        return ("; " if any_prefix_is_slash else " / ").join(rv), help
-
-
-def _typer_format_options(
-    self: click.core.Command, *, ctx: click.Context, formatter: click.HelpFormatter
-) -> None:
-    args = []
-    opts = []
-    for param in self.get_params(ctx):
-        rv = param.get_help_record(ctx)
-        if rv is not None:
-            if param.param_type_name == "argument":
-                args.append(rv)
-            elif param.param_type_name == "option":
-                opts.append(rv)
-
-    if args:
-        with formatter.section(_("Arguments")):
-            formatter.write_dl(args)
-    if opts:
-        with formatter.section(_("Options")):
-            formatter.write_dl(opts)
-
-
-def _typer_main_shell_completion(
-    self: click.core.Command,
-    *,
-    ctx_args: MutableMapping[str, Any],
-    prog_name: str,
-    complete_var: Optional[str] = None,
-) -> None:
-    if complete_var is None:
-        complete_var = f"_{prog_name}_COMPLETE".replace("-", "_").upper()
-
-    instruction = os.environ.get(complete_var)
-
-    if not instruction:
-        return
-
-    from .completion import shell_complete
-
-    rv = shell_complete(self, ctx_args, prog_name, complete_var, instruction)
-    sys.exit(rv)
-
-
-class TyperCommand(click.core.Command):
-    def __init__(
-        self,
-        name: Optional[str],
-        *,
-        context_settings: Optional[Dict[str, Any]] = None,
-        callback: Optional[Callable[..., Any]] = None,
-        params: Optional[List[click.Parameter]] = None,
-        help: Optional[str] = None,
-        epilog: Optional[str] = None,
-        short_help: Optional[str] = None,
-        options_metavar: Optional[str] = "[OPTIONS]",
-        add_help_option: bool = True,
-        no_args_is_help: bool = False,
-        hidden: bool = False,
-        deprecated: bool = False,
-        # Rich settings
-        rich_markup_mode: MarkupMode = DEFAULT_MARKUP_MODE,
-        rich_help_panel: Union[str, None] = None,
-    ) -> None:
-        super().__init__(
-            name=name,
-            context_settings=context_settings,
-            callback=callback,
-            params=params,
-            help=help,
-            epilog=epilog,
-            short_help=short_help,
-            options_metavar=options_metavar,
-            add_help_option=add_help_option,
-            no_args_is_help=no_args_is_help,
-            hidden=hidden,
-            deprecated=deprecated,
-        )
-        self.rich_markup_mode: MarkupMode = rich_markup_mode
-        self.rich_help_panel = rich_help_panel
-
-    def format_options(
-        self, ctx: click.Context, formatter: click.HelpFormatter
-    ) -> None:
-        _typer_format_options(self, ctx=ctx, formatter=formatter)
-
-    def _main_shell_completion(
-        self,
-        ctx_args: MutableMapping[str, Any],
-        prog_name: str,
-        complete_var: Optional[str] = None,
-    ) -> None:
-        _typer_main_shell_completion(
-            self, ctx_args=ctx_args, prog_name=prog_name, complete_var=complete_var
-        )
-
-    def main(
-        self,
-        args: Optional[Sequence[str]] = None,
-        prog_name: Optional[str] = None,
-        complete_var: Optional[str] = None,
-        standalone_mode: bool = True,
-        windows_expand_args: bool = True,
-        **extra: Any,
-    ) -> Any:
-        return _main(
-            self,
-            args=args,
-            prog_name=prog_name,
-            complete_var=complete_var,
-            standalone_mode=standalone_mode,
-            windows_expand_args=windows_expand_args,
-            rich_markup_mode=self.rich_markup_mode,
-            **extra,
-        )
-
-    def format_help(self, ctx: click.Context, formatter: click.HelpFormatter) -> None:
-        if not rich or self.rich_markup_mode is None:
-            return super().format_help(ctx, formatter)
-        return rich_utils.rich_format_help(
-            obj=self,
-            ctx=ctx,
-            markup_mode=self.rich_markup_mode,
-        )
-
-
-class TyperGroup(click.core.Group):
-    def __init__(
-        self,
-        *,
-        name: Optional[str] = None,
-        commands: Optional[
-            Union[Dict[str, click.Command], Sequence[click.Command]]
-        ] = None,
-        # Rich settings
-        rich_markup_mode: MarkupMode = DEFAULT_MARKUP_MODE,
-        rich_help_panel: Union[str, None] = None,
-        **attrs: Any,
-    ) -> None:
-        super().__init__(name=name, commands=commands, **attrs)
-        self.rich_markup_mode: MarkupMode = rich_markup_mode
-        self.rich_help_panel = rich_help_panel
-
-    def format_options(
-        self, ctx: click.Context, formatter: click.HelpFormatter
-    ) -> None:
-        _typer_format_options(self, ctx=ctx, formatter=formatter)
-        self.format_commands(ctx, formatter)
-
-    def _main_shell_completion(
-        self,
-        ctx_args: MutableMapping[str, Any],
-        prog_name: str,
-        complete_var: Optional[str] = None,
-    ) -> None:
-        _typer_main_shell_completion(
-            self, ctx_args=ctx_args, prog_name=prog_name, complete_var=complete_var
-        )
-
-    def main(
-        self,
-        args: Optional[Sequence[str]] = None,
-        prog_name: Optional[str] = None,
-        complete_var: Optional[str] = None,
-        standalone_mode: bool = True,
-        windows_expand_args: bool = True,
-        **extra: Any,
-    ) -> Any:
-        return _main(
-            self,
-            args=args,
-            prog_name=prog_name,
-            complete_var=complete_var,
-            standalone_mode=standalone_mode,
-            windows_expand_args=windows_expand_args,
-            rich_markup_mode=self.rich_markup_mode,
-            **extra,
-        )
-
-    def format_help(self, ctx: click.Context, formatter: click.HelpFormatter) -> None:
-        if not rich or self.rich_markup_mode is None:
-            return super().format_help(ctx, formatter)
-        return rich_utils.rich_format_help(
-            obj=self,
-            ctx=ctx,
-            markup_mode=self.rich_markup_mode,
-        )
-
-    def list_commands(self, ctx: click.Context) -> List[str]:
-        """Returns a list of subcommand names.
-        Note that in Click's Group class, these are sorted.
-        In Typer, we wish to maintain the original order of creation (cf Issue #933)"""
-        return [n for n, c in self.commands.items()]
diff --git a/venv/lib/python3.10/site-packages/typer/main.py b/venv/lib/python3.10/site-packages/typer/main.py
deleted file mode 100644
index 59e22c77aaac18f3f4194201e2a118dfa136b8a8..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer/main.py
+++ /dev/null
@@ -1,1145 +0,0 @@
-import inspect
-import os
-import platform
-import shutil
-import subprocess
-import sys
-import traceback
-from datetime import datetime
-from enum import Enum
-from functools import update_wrapper
-from pathlib import Path
-from traceback import FrameSummary, StackSummary
-from types import TracebackType
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type, Union
-from uuid import UUID
-
-import click
-from typer._types import TyperChoice
-
-from ._typing import get_args, get_origin, is_union
-from .completion import get_completion_inspect_parameters
-from .core import (
-    DEFAULT_MARKUP_MODE,
-    MarkupMode,
-    TyperArgument,
-    TyperCommand,
-    TyperGroup,
-    TyperOption,
-)
-from .models import (
-    AnyType,
-    ArgumentInfo,
-    CommandFunctionType,
-    CommandInfo,
-    Default,
-    DefaultPlaceholder,
-    DeveloperExceptionConfig,
-    FileBinaryRead,
-    FileBinaryWrite,
-    FileText,
-    FileTextWrite,
-    NoneType,
-    OptionInfo,
-    ParameterInfo,
-    ParamMeta,
-    Required,
-    TyperInfo,
-    TyperPath,
-)
-from .utils import get_params_from_function
-
-try:
-    import rich
-    from rich.traceback import Traceback
-
-    from . import rich_utils
-
-    console_stderr = rich_utils._get_rich_console(stderr=True)
-
-except ImportError:  # pragma: no cover
-    rich = None  # type: ignore
-
-_original_except_hook = sys.excepthook
-_typer_developer_exception_attr_name = "__typer_developer_exception__"
-
-
-def except_hook(
-    exc_type: Type[BaseException], exc_value: BaseException, tb: Optional[TracebackType]
-) -> None:
-    exception_config: Union[DeveloperExceptionConfig, None] = getattr(
-        exc_value, _typer_developer_exception_attr_name, None
-    )
-    standard_traceback = os.getenv("_TYPER_STANDARD_TRACEBACK")
-    if (
-        standard_traceback
-        or not exception_config
-        or not exception_config.pretty_exceptions_enable
-    ):
-        _original_except_hook(exc_type, exc_value, tb)
-        return
-    typer_path = os.path.dirname(__file__)
-    click_path = os.path.dirname(click.__file__)
-    supress_internal_dir_names = [typer_path, click_path]
-    exc = exc_value
-    if rich:
-        from .rich_utils import MAX_WIDTH
-
-        rich_tb = Traceback.from_exception(
-            type(exc),
-            exc,
-            exc.__traceback__,
-            show_locals=exception_config.pretty_exceptions_show_locals,
-            suppress=supress_internal_dir_names,
-            width=MAX_WIDTH,
-        )
-        console_stderr.print(rich_tb)
-        return
-    tb_exc = traceback.TracebackException.from_exception(exc)
-    stack: List[FrameSummary] = []
-    for frame in tb_exc.stack:
-        if any(frame.filename.startswith(path) for path in supress_internal_dir_names):
-            if not exception_config.pretty_exceptions_short:
-                # Hide the line for internal libraries, Typer and Click
-                stack.append(
-                    traceback.FrameSummary(
-                        filename=frame.filename,
-                        lineno=frame.lineno,
-                        name=frame.name,
-                        line="",
-                    )
-                )
-        else:
-            stack.append(frame)
-    # Type ignore ref: https://github.com/python/typeshed/pull/8244
-    final_stack_summary = StackSummary.from_list(stack)
-    tb_exc.stack = final_stack_summary
-    for line in tb_exc.format():
-        print(line, file=sys.stderr)
-    return
-
-
-def get_install_completion_arguments() -> Tuple[click.Parameter, click.Parameter]:
-    install_param, show_param = get_completion_inspect_parameters()
-    click_install_param, _ = get_click_param(install_param)
-    click_show_param, _ = get_click_param(show_param)
-    return click_install_param, click_show_param
-
-
-class Typer:
-    def __init__(
-        self,
-        *,
-        name: Optional[str] = Default(None),
-        cls: Optional[Type[TyperGroup]] = Default(None),
-        invoke_without_command: bool = Default(False),
-        no_args_is_help: bool = Default(False),
-        subcommand_metavar: Optional[str] = Default(None),
-        chain: bool = Default(False),
-        result_callback: Optional[Callable[..., Any]] = Default(None),
-        # Command
-        context_settings: Optional[Dict[Any, Any]] = Default(None),
-        callback: Optional[Callable[..., Any]] = Default(None),
-        help: Optional[str] = Default(None),
-        epilog: Optional[str] = Default(None),
-        short_help: Optional[str] = Default(None),
-        options_metavar: str = Default("[OPTIONS]"),
-        add_help_option: bool = Default(True),
-        hidden: bool = Default(False),
-        deprecated: bool = Default(False),
-        add_completion: bool = True,
-        # Rich settings
-        rich_markup_mode: MarkupMode = Default(DEFAULT_MARKUP_MODE),
-        rich_help_panel: Union[str, None] = Default(None),
-        pretty_exceptions_enable: bool = True,
-        pretty_exceptions_show_locals: bool = True,
-        pretty_exceptions_short: bool = True,
-    ):
-        self._add_completion = add_completion
-        self.rich_markup_mode: MarkupMode = rich_markup_mode
-        self.rich_help_panel = rich_help_panel
-        self.pretty_exceptions_enable = pretty_exceptions_enable
-        self.pretty_exceptions_show_locals = pretty_exceptions_show_locals
-        self.pretty_exceptions_short = pretty_exceptions_short
-        self.info = TyperInfo(
-            name=name,
-            cls=cls,
-            invoke_without_command=invoke_without_command,
-            no_args_is_help=no_args_is_help,
-            subcommand_metavar=subcommand_metavar,
-            chain=chain,
-            result_callback=result_callback,
-            context_settings=context_settings,
-            callback=callback,
-            help=help,
-            epilog=epilog,
-            short_help=short_help,
-            options_metavar=options_metavar,
-            add_help_option=add_help_option,
-            hidden=hidden,
-            deprecated=deprecated,
-        )
-        self.registered_groups: List[TyperInfo] = []
-        self.registered_commands: List[CommandInfo] = []
-        self.registered_callback: Optional[TyperInfo] = None
-
-    def callback(
-        self,
-        *,
-        cls: Optional[Type[TyperGroup]] = Default(None),
-        invoke_without_command: bool = Default(False),
-        no_args_is_help: bool = Default(False),
-        subcommand_metavar: Optional[str] = Default(None),
-        chain: bool = Default(False),
-        result_callback: Optional[Callable[..., Any]] = Default(None),
-        # Command
-        context_settings: Optional[Dict[Any, Any]] = Default(None),
-        help: Optional[str] = Default(None),
-        epilog: Optional[str] = Default(None),
-        short_help: Optional[str] = Default(None),
-        options_metavar: str = Default("[OPTIONS]"),
-        add_help_option: bool = Default(True),
-        hidden: bool = Default(False),
-        deprecated: bool = Default(False),
-        # Rich settings
-        rich_help_panel: Union[str, None] = Default(None),
-    ) -> Callable[[CommandFunctionType], CommandFunctionType]:
-        def decorator(f: CommandFunctionType) -> CommandFunctionType:
-            self.registered_callback = TyperInfo(
-                cls=cls,
-                invoke_without_command=invoke_without_command,
-                no_args_is_help=no_args_is_help,
-                subcommand_metavar=subcommand_metavar,
-                chain=chain,
-                result_callback=result_callback,
-                context_settings=context_settings,
-                callback=f,
-                help=help,
-                epilog=epilog,
-                short_help=short_help,
-                options_metavar=options_metavar,
-                add_help_option=add_help_option,
-                hidden=hidden,
-                deprecated=deprecated,
-                rich_help_panel=rich_help_panel,
-            )
-            return f
-
-        return decorator
-
-    def command(
-        self,
-        name: Optional[str] = None,
-        *,
-        cls: Optional[Type[TyperCommand]] = None,
-        context_settings: Optional[Dict[Any, Any]] = None,
-        help: Optional[str] = None,
-        epilog: Optional[str] = None,
-        short_help: Optional[str] = None,
-        options_metavar: str = "[OPTIONS]",
-        add_help_option: bool = True,
-        no_args_is_help: bool = False,
-        hidden: bool = False,
-        deprecated: bool = False,
-        # Rich settings
-        rich_help_panel: Union[str, None] = Default(None),
-    ) -> Callable[[CommandFunctionType], CommandFunctionType]:
-        if cls is None:
-            cls = TyperCommand
-
-        def decorator(f: CommandFunctionType) -> CommandFunctionType:
-            self.registered_commands.append(
-                CommandInfo(
-                    name=name,
-                    cls=cls,
-                    context_settings=context_settings,
-                    callback=f,
-                    help=help,
-                    epilog=epilog,
-                    short_help=short_help,
-                    options_metavar=options_metavar,
-                    add_help_option=add_help_option,
-                    no_args_is_help=no_args_is_help,
-                    hidden=hidden,
-                    deprecated=deprecated,
-                    # Rich settings
-                    rich_help_panel=rich_help_panel,
-                )
-            )
-            return f
-
-        return decorator
-
-    def add_typer(
-        self,
-        typer_instance: "Typer",
-        *,
-        name: Optional[str] = Default(None),
-        cls: Optional[Type[TyperGroup]] = Default(None),
-        invoke_without_command: bool = Default(False),
-        no_args_is_help: bool = Default(False),
-        subcommand_metavar: Optional[str] = Default(None),
-        chain: bool = Default(False),
-        result_callback: Optional[Callable[..., Any]] = Default(None),
-        # Command
-        context_settings: Optional[Dict[Any, Any]] = Default(None),
-        callback: Optional[Callable[..., Any]] = Default(None),
-        help: Optional[str] = Default(None),
-        epilog: Optional[str] = Default(None),
-        short_help: Optional[str] = Default(None),
-        options_metavar: str = Default("[OPTIONS]"),
-        add_help_option: bool = Default(True),
-        hidden: bool = Default(False),
-        deprecated: bool = Default(False),
-        # Rich settings
-        rich_help_panel: Union[str, None] = Default(None),
-    ) -> None:
-        self.registered_groups.append(
-            TyperInfo(
-                typer_instance,
-                name=name,
-                cls=cls,
-                invoke_without_command=invoke_without_command,
-                no_args_is_help=no_args_is_help,
-                subcommand_metavar=subcommand_metavar,
-                chain=chain,
-                result_callback=result_callback,
-                context_settings=context_settings,
-                callback=callback,
-                help=help,
-                epilog=epilog,
-                short_help=short_help,
-                options_metavar=options_metavar,
-                add_help_option=add_help_option,
-                hidden=hidden,
-                deprecated=deprecated,
-                rich_help_panel=rich_help_panel,
-            )
-        )
-
-    def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        if sys.excepthook != except_hook:
-            sys.excepthook = except_hook
-        try:
-            return get_command(self)(*args, **kwargs)
-        except Exception as e:
-            # Set a custom attribute to tell the hook to show nice exceptions for user
-            # code. An alternative/first implementation was a custom exception with
-            # raise custom_exc from e
-            # but that means the last error shown is the custom exception, not the
-            # actual error. This trick improves developer experience by showing the
-            # actual error last.
-            setattr(
-                e,
-                _typer_developer_exception_attr_name,
-                DeveloperExceptionConfig(
-                    pretty_exceptions_enable=self.pretty_exceptions_enable,
-                    pretty_exceptions_show_locals=self.pretty_exceptions_show_locals,
-                    pretty_exceptions_short=self.pretty_exceptions_short,
-                ),
-            )
-            raise e
-
-
-def get_group(typer_instance: Typer) -> TyperGroup:
-    group = get_group_from_info(
-        TyperInfo(typer_instance),
-        pretty_exceptions_short=typer_instance.pretty_exceptions_short,
-        rich_markup_mode=typer_instance.rich_markup_mode,
-    )
-    return group
-
-
-def get_command(typer_instance: Typer) -> click.Command:
-    if typer_instance._add_completion:
-        click_install_param, click_show_param = get_install_completion_arguments()
-    if (
-        typer_instance.registered_callback
-        or typer_instance.info.callback
-        or typer_instance.registered_groups
-        or len(typer_instance.registered_commands) > 1
-    ):
-        # Create a Group
-        click_command: click.Command = get_group(typer_instance)
-        if typer_instance._add_completion:
-            click_command.params.append(click_install_param)
-            click_command.params.append(click_show_param)
-        return click_command
-    elif len(typer_instance.registered_commands) == 1:
-        # Create a single Command
-        single_command = typer_instance.registered_commands[0]
-
-        if not single_command.context_settings and not isinstance(
-            typer_instance.info.context_settings, DefaultPlaceholder
-        ):
-            single_command.context_settings = typer_instance.info.context_settings
-
-        click_command = get_command_from_info(
-            single_command,
-            pretty_exceptions_short=typer_instance.pretty_exceptions_short,
-            rich_markup_mode=typer_instance.rich_markup_mode,
-        )
-        if typer_instance._add_completion:
-            click_command.params.append(click_install_param)
-            click_command.params.append(click_show_param)
-        return click_command
-    raise RuntimeError(
-        "Could not get a command for this Typer instance"
-    )  # pragma: no cover
-
-
-def solve_typer_info_help(typer_info: TyperInfo) -> str:
-    # Priority 1: Explicit value was set in app.add_typer()
-    if not isinstance(typer_info.help, DefaultPlaceholder):
-        return inspect.cleandoc(typer_info.help or "")
-    # Priority 2: Explicit value was set in sub_app.callback()
-    try:
-        callback_help = typer_info.typer_instance.registered_callback.help
-        if not isinstance(callback_help, DefaultPlaceholder):
-            return inspect.cleandoc(callback_help or "")
-    except AttributeError:
-        pass
-    # Priority 3: Explicit value was set in sub_app = typer.Typer()
-    try:
-        instance_help = typer_info.typer_instance.info.help
-        if not isinstance(instance_help, DefaultPlaceholder):
-            return inspect.cleandoc(instance_help or "")
-    except AttributeError:
-        pass
-    # Priority 4: Implicit inference from callback docstring in app.add_typer()
-    if typer_info.callback:
-        doc = inspect.getdoc(typer_info.callback)
-        if doc:
-            return doc
-    # Priority 5: Implicit inference from callback docstring in @app.callback()
-    try:
-        callback = typer_info.typer_instance.registered_callback.callback
-        if not isinstance(callback, DefaultPlaceholder):
-            doc = inspect.getdoc(callback or "")
-            if doc:
-                return doc
-    except AttributeError:
-        pass
-    # Priority 6: Implicit inference from callback docstring in typer.Typer()
-    try:
-        instance_callback = typer_info.typer_instance.info.callback
-        if not isinstance(instance_callback, DefaultPlaceholder):
-            doc = inspect.getdoc(instance_callback)
-            if doc:
-                return doc
-    except AttributeError:
-        pass
-    # Value not set, use the default
-    return typer_info.help.value
-
-
-def solve_typer_info_defaults(typer_info: TyperInfo) -> TyperInfo:
-    values: Dict[str, Any] = {}
-    for name, value in typer_info.__dict__.items():
-        # Priority 1: Value was set in app.add_typer()
-        if not isinstance(value, DefaultPlaceholder):
-            values[name] = value
-            continue
-        # Priority 2: Value was set in @subapp.callback()
-        try:
-            callback_value = getattr(
-                typer_info.typer_instance.registered_callback,  # type: ignore
-                name,
-            )
-            if not isinstance(callback_value, DefaultPlaceholder):
-                values[name] = callback_value
-                continue
-        except AttributeError:
-            pass
-        # Priority 3: Value set in subapp = typer.Typer()
-        try:
-            instance_value = getattr(
-                typer_info.typer_instance.info,  # type: ignore
-                name,
-            )
-            if not isinstance(instance_value, DefaultPlaceholder):
-                values[name] = instance_value
-                continue
-        except AttributeError:
-            pass
-        # Value not set, use the default
-        values[name] = value.value
-    values["help"] = solve_typer_info_help(typer_info)
-    return TyperInfo(**values)
-
-
-def get_group_from_info(
-    group_info: TyperInfo,
-    *,
-    pretty_exceptions_short: bool,
-    rich_markup_mode: MarkupMode,
-) -> TyperGroup:
-    assert group_info.typer_instance, (
-        "A Typer instance is needed to generate a Click Group"
-    )
-    commands: Dict[str, click.Command] = {}
-    for command_info in group_info.typer_instance.registered_commands:
-        command = get_command_from_info(
-            command_info=command_info,
-            pretty_exceptions_short=pretty_exceptions_short,
-            rich_markup_mode=rich_markup_mode,
-        )
-        if command.name:
-            commands[command.name] = command
-    for sub_group_info in group_info.typer_instance.registered_groups:
-        sub_group = get_group_from_info(
-            sub_group_info,
-            pretty_exceptions_short=pretty_exceptions_short,
-            rich_markup_mode=rich_markup_mode,
-        )
-        if sub_group.name:
-            commands[sub_group.name] = sub_group
-        else:
-            if sub_group.callback:
-                import warnings
-
-                warnings.warn(
-                    "The 'callback' parameter is not supported by Typer when using `add_typer` without a name",
-                    stacklevel=5,
-                )
-            for sub_command_name, sub_command in sub_group.commands.items():
-                commands[sub_command_name] = sub_command
-    solved_info = solve_typer_info_defaults(group_info)
-    (
-        params,
-        convertors,
-        context_param_name,
-    ) = get_params_convertors_ctx_param_name_from_function(solved_info.callback)
-    cls = solved_info.cls or TyperGroup
-    assert issubclass(cls, TyperGroup), f"{cls} should be a subclass of {TyperGroup}"
-    group = cls(
-        name=solved_info.name or "",
-        commands=commands,
-        invoke_without_command=solved_info.invoke_without_command,
-        no_args_is_help=solved_info.no_args_is_help,
-        subcommand_metavar=solved_info.subcommand_metavar,
-        chain=solved_info.chain,
-        result_callback=solved_info.result_callback,
-        context_settings=solved_info.context_settings,
-        callback=get_callback(
-            callback=solved_info.callback,
-            params=params,
-            convertors=convertors,
-            context_param_name=context_param_name,
-            pretty_exceptions_short=pretty_exceptions_short,
-        ),
-        params=params,
-        help=solved_info.help,
-        epilog=solved_info.epilog,
-        short_help=solved_info.short_help,
-        options_metavar=solved_info.options_metavar,
-        add_help_option=solved_info.add_help_option,
-        hidden=solved_info.hidden,
-        deprecated=solved_info.deprecated,
-        rich_markup_mode=rich_markup_mode,
-        # Rich settings
-        rich_help_panel=solved_info.rich_help_panel,
-    )
-    return group
-
-
-def get_command_name(name: str) -> str:
-    return name.lower().replace("_", "-")
-
-
-def get_params_convertors_ctx_param_name_from_function(
-    callback: Optional[Callable[..., Any]],
-) -> Tuple[List[Union[click.Argument, click.Option]], Dict[str, Any], Optional[str]]:
-    params = []
-    convertors = {}
-    context_param_name = None
-    if callback:
-        parameters = get_params_from_function(callback)
-        for param_name, param in parameters.items():
-            if lenient_issubclass(param.annotation, click.Context):
-                context_param_name = param_name
-                continue
-            click_param, convertor = get_click_param(param)
-            if convertor:
-                convertors[param_name] = convertor
-            params.append(click_param)
-    return params, convertors, context_param_name
-
-
-def get_command_from_info(
-    command_info: CommandInfo,
-    *,
-    pretty_exceptions_short: bool,
-    rich_markup_mode: MarkupMode,
-) -> click.Command:
-    assert command_info.callback, "A command must have a callback function"
-    name = command_info.name or get_command_name(command_info.callback.__name__)
-    use_help = command_info.help
-    if use_help is None:
-        use_help = inspect.getdoc(command_info.callback)
-    else:
-        use_help = inspect.cleandoc(use_help)
-    (
-        params,
-        convertors,
-        context_param_name,
-    ) = get_params_convertors_ctx_param_name_from_function(command_info.callback)
-    cls = command_info.cls or TyperCommand
-    command = cls(
-        name=name,
-        context_settings=command_info.context_settings,
-        callback=get_callback(
-            callback=command_info.callback,
-            params=params,
-            convertors=convertors,
-            context_param_name=context_param_name,
-            pretty_exceptions_short=pretty_exceptions_short,
-        ),
-        params=params,  # type: ignore
-        help=use_help,
-        epilog=command_info.epilog,
-        short_help=command_info.short_help,
-        options_metavar=command_info.options_metavar,
-        add_help_option=command_info.add_help_option,
-        no_args_is_help=command_info.no_args_is_help,
-        hidden=command_info.hidden,
-        deprecated=command_info.deprecated,
-        rich_markup_mode=rich_markup_mode,
-        # Rich settings
-        rich_help_panel=command_info.rich_help_panel,
-    )
-    return command
-
-
-def determine_type_convertor(type_: Any) -> Optional[Callable[[Any], Any]]:
-    convertor: Optional[Callable[[Any], Any]] = None
-    if lenient_issubclass(type_, Path):
-        convertor = param_path_convertor
-    if lenient_issubclass(type_, Enum):
-        convertor = generate_enum_convertor(type_)
-    return convertor
-
-
-def param_path_convertor(value: Optional[str] = None) -> Optional[Path]:
-    if value is not None:
-        return Path(value)
-    return None
-
-
-def generate_enum_convertor(enum: Type[Enum]) -> Callable[[Any], Any]:
-    val_map = {str(val.value): val for val in enum}
-
-    def convertor(value: Any) -> Any:
-        if value is not None:
-            val = str(value)
-            if val in val_map:
-                key = val_map[val]
-                return enum(key)
-
-    return convertor
-
-
-def generate_list_convertor(
-    convertor: Optional[Callable[[Any], Any]], default_value: Optional[Any]
-) -> Callable[[Sequence[Any]], Optional[List[Any]]]:
-    def internal_convertor(value: Sequence[Any]) -> Optional[List[Any]]:
-        if default_value is None and len(value) == 0:
-            return None
-        return [convertor(v) if convertor else v for v in value]
-
-    return internal_convertor
-
-
-def generate_tuple_convertor(
-    types: Sequence[Any],
-) -> Callable[[Optional[Tuple[Any, ...]]], Optional[Tuple[Any, ...]]]:
-    convertors = [determine_type_convertor(type_) for type_ in types]
-
-    def internal_convertor(
-        param_args: Optional[Tuple[Any, ...]],
-    ) -> Optional[Tuple[Any, ...]]:
-        if param_args is None:
-            return None
-        return tuple(
-            convertor(arg) if convertor else arg
-            for (convertor, arg) in zip(convertors, param_args)
-        )
-
-    return internal_convertor
-
-
-def get_callback(
-    *,
-    callback: Optional[Callable[..., Any]] = None,
-    params: Sequence[click.Parameter] = [],
-    convertors: Optional[Dict[str, Callable[[str], Any]]] = None,
-    context_param_name: Optional[str] = None,
-    pretty_exceptions_short: bool,
-) -> Optional[Callable[..., Any]]:
-    use_convertors = convertors or {}
-    if not callback:
-        return None
-    parameters = get_params_from_function(callback)
-    use_params: Dict[str, Any] = {}
-    for param_name in parameters:
-        use_params[param_name] = None
-    for param in params:
-        if param.name:
-            use_params[param.name] = param.default
-
-    def wrapper(**kwargs: Any) -> Any:
-        _rich_traceback_guard = pretty_exceptions_short  # noqa: F841
-        for k, v in kwargs.items():
-            if k in use_convertors:
-                use_params[k] = use_convertors[k](v)
-            else:
-                use_params[k] = v
-        if context_param_name:
-            use_params[context_param_name] = click.get_current_context()
-        return callback(**use_params)
-
-    update_wrapper(wrapper, callback)
-    return wrapper
-
-
-def get_click_type(
-    *, annotation: Any, parameter_info: ParameterInfo
-) -> click.ParamType:
-    if parameter_info.click_type is not None:
-        return parameter_info.click_type
-
-    elif parameter_info.parser is not None:
-        return click.types.FuncParamType(parameter_info.parser)
-
-    elif annotation is str:
-        return click.STRING
-    elif annotation is int:
-        if parameter_info.min is not None or parameter_info.max is not None:
-            min_ = None
-            max_ = None
-            if parameter_info.min is not None:
-                min_ = int(parameter_info.min)
-            if parameter_info.max is not None:
-                max_ = int(parameter_info.max)
-            return click.IntRange(min=min_, max=max_, clamp=parameter_info.clamp)
-        else:
-            return click.INT
-    elif annotation is float:
-        if parameter_info.min is not None or parameter_info.max is not None:
-            return click.FloatRange(
-                min=parameter_info.min,
-                max=parameter_info.max,
-                clamp=parameter_info.clamp,
-            )
-        else:
-            return click.FLOAT
-    elif annotation is bool:
-        return click.BOOL
-    elif annotation == UUID:
-        return click.UUID
-    elif annotation == datetime:
-        return click.DateTime(formats=parameter_info.formats)
-    elif (
-        annotation == Path
-        or parameter_info.allow_dash
-        or parameter_info.path_type
-        or parameter_info.resolve_path
-    ):
-        return TyperPath(
-            exists=parameter_info.exists,
-            file_okay=parameter_info.file_okay,
-            dir_okay=parameter_info.dir_okay,
-            writable=parameter_info.writable,
-            readable=parameter_info.readable,
-            resolve_path=parameter_info.resolve_path,
-            allow_dash=parameter_info.allow_dash,
-            path_type=parameter_info.path_type,
-        )
-    elif lenient_issubclass(annotation, FileTextWrite):
-        return click.File(
-            mode=parameter_info.mode or "w",
-            encoding=parameter_info.encoding,
-            errors=parameter_info.errors,
-            lazy=parameter_info.lazy,
-            atomic=parameter_info.atomic,
-        )
-    elif lenient_issubclass(annotation, FileText):
-        return click.File(
-            mode=parameter_info.mode or "r",
-            encoding=parameter_info.encoding,
-            errors=parameter_info.errors,
-            lazy=parameter_info.lazy,
-            atomic=parameter_info.atomic,
-        )
-    elif lenient_issubclass(annotation, FileBinaryRead):
-        return click.File(
-            mode=parameter_info.mode or "rb",
-            encoding=parameter_info.encoding,
-            errors=parameter_info.errors,
-            lazy=parameter_info.lazy,
-            atomic=parameter_info.atomic,
-        )
-    elif lenient_issubclass(annotation, FileBinaryWrite):
-        return click.File(
-            mode=parameter_info.mode or "wb",
-            encoding=parameter_info.encoding,
-            errors=parameter_info.errors,
-            lazy=parameter_info.lazy,
-            atomic=parameter_info.atomic,
-        )
-    elif lenient_issubclass(annotation, Enum):
-        # The custom TyperChoice is only needed for Click < 8.2.0, to parse the
-        # command line values matching them to the enum values. Click 8.2.0 added
-        # support for enum values but reading enum names.
-        # Passing here the list of enum values (instead of just the enum) accounts for
-        # Click < 8.2.0.
-        return TyperChoice(
-            [item.value for item in annotation],
-            case_sensitive=parameter_info.case_sensitive,
-        )
-    raise RuntimeError(f"Type not yet supported: {annotation}")  # pragma: no cover
-
-
-def lenient_issubclass(
-    cls: Any, class_or_tuple: Union[AnyType, Tuple[AnyType, ...]]
-) -> bool:
-    return isinstance(cls, type) and issubclass(cls, class_or_tuple)
-
-
-def get_click_param(
-    param: ParamMeta,
-) -> Tuple[Union[click.Argument, click.Option], Any]:
-    # First, find out what will be:
-    # * ParamInfo (ArgumentInfo or OptionInfo)
-    # * default_value
-    # * required
-    default_value = None
-    required = False
-    if isinstance(param.default, ParameterInfo):
-        parameter_info = param.default
-        if parameter_info.default == Required:
-            required = True
-        else:
-            default_value = parameter_info.default
-    elif param.default == Required or param.default is param.empty:
-        required = True
-        parameter_info = ArgumentInfo()
-    else:
-        default_value = param.default
-        parameter_info = OptionInfo()
-    annotation: Any
-    if param.annotation is not param.empty:
-        annotation = param.annotation
-    else:
-        annotation = str
-    main_type = annotation
-    is_list = False
-    is_tuple = False
-    parameter_type: Any = None
-    is_flag = None
-    origin = get_origin(main_type)
-
-    if origin is not None:
-        # Handle SomeType | None and Optional[SomeType]
-        if is_union(origin):
-            types = []
-            for type_ in get_args(main_type):
-                if type_ is NoneType:
-                    continue
-                types.append(type_)
-            assert len(types) == 1, "Typer Currently doesn't support Union types"
-            main_type = types[0]
-            origin = get_origin(main_type)
-        # Handle Tuples and Lists
-        if lenient_issubclass(origin, List):
-            main_type = get_args(main_type)[0]
-            assert not get_origin(main_type), (
-                "List types with complex sub-types are not currently supported"
-            )
-            is_list = True
-        elif lenient_issubclass(origin, Tuple):  # type: ignore
-            types = []
-            for type_ in get_args(main_type):
-                assert not get_origin(type_), (
-                    "Tuple types with complex sub-types are not currently supported"
-                )
-                types.append(
-                    get_click_type(annotation=type_, parameter_info=parameter_info)
-                )
-            parameter_type = tuple(types)
-            is_tuple = True
-    if parameter_type is None:
-        parameter_type = get_click_type(
-            annotation=main_type, parameter_info=parameter_info
-        )
-    convertor = determine_type_convertor(main_type)
-    if is_list:
-        convertor = generate_list_convertor(
-            convertor=convertor, default_value=default_value
-        )
-    if is_tuple:
-        convertor = generate_tuple_convertor(get_args(main_type))
-    if isinstance(parameter_info, OptionInfo):
-        if main_type is bool:
-            is_flag = True
-            # Click doesn't accept a flag of type bool, only None, and then it sets it
-            # to bool internally
-            parameter_type = None
-        default_option_name = get_command_name(param.name)
-        if is_flag:
-            default_option_declaration = (
-                f"--{default_option_name}/--no-{default_option_name}"
-            )
-        else:
-            default_option_declaration = f"--{default_option_name}"
-        param_decls = [param.name]
-        if parameter_info.param_decls:
-            param_decls.extend(parameter_info.param_decls)
-        else:
-            param_decls.append(default_option_declaration)
-        return (
-            TyperOption(
-                # Option
-                param_decls=param_decls,
-                show_default=parameter_info.show_default,
-                prompt=parameter_info.prompt,
-                confirmation_prompt=parameter_info.confirmation_prompt,
-                prompt_required=parameter_info.prompt_required,
-                hide_input=parameter_info.hide_input,
-                is_flag=is_flag,
-                multiple=is_list,
-                count=parameter_info.count,
-                allow_from_autoenv=parameter_info.allow_from_autoenv,
-                type=parameter_type,
-                help=parameter_info.help,
-                hidden=parameter_info.hidden,
-                show_choices=parameter_info.show_choices,
-                show_envvar=parameter_info.show_envvar,
-                # Parameter
-                required=required,
-                default=default_value,
-                callback=get_param_callback(
-                    callback=parameter_info.callback, convertor=convertor
-                ),
-                metavar=parameter_info.metavar,
-                expose_value=parameter_info.expose_value,
-                is_eager=parameter_info.is_eager,
-                envvar=parameter_info.envvar,
-                shell_complete=parameter_info.shell_complete,
-                autocompletion=get_param_completion(parameter_info.autocompletion),
-                # Rich settings
-                rich_help_panel=parameter_info.rich_help_panel,
-            ),
-            convertor,
-        )
-    elif isinstance(parameter_info, ArgumentInfo):
-        param_decls = [param.name]
-        nargs = None
-        if is_list:
-            nargs = -1
-        return (
-            TyperArgument(
-                # Argument
-                param_decls=param_decls,
-                type=parameter_type,
-                required=required,
-                nargs=nargs,
-                # TyperArgument
-                show_default=parameter_info.show_default,
-                show_choices=parameter_info.show_choices,
-                show_envvar=parameter_info.show_envvar,
-                help=parameter_info.help,
-                hidden=parameter_info.hidden,
-                # Parameter
-                default=default_value,
-                callback=get_param_callback(
-                    callback=parameter_info.callback, convertor=convertor
-                ),
-                metavar=parameter_info.metavar,
-                expose_value=parameter_info.expose_value,
-                is_eager=parameter_info.is_eager,
-                envvar=parameter_info.envvar,
-                shell_complete=parameter_info.shell_complete,
-                autocompletion=get_param_completion(parameter_info.autocompletion),
-                # Rich settings
-                rich_help_panel=parameter_info.rich_help_panel,
-            ),
-            convertor,
-        )
-    raise AssertionError("A click.Parameter should be returned")  # pragma: no cover
-
-
-def get_param_callback(
-    *,
-    callback: Optional[Callable[..., Any]] = None,
-    convertor: Optional[Callable[..., Any]] = None,
-) -> Optional[Callable[..., Any]]:
-    if not callback:
-        return None
-    parameters = get_params_from_function(callback)
-    ctx_name = None
-    click_param_name = None
-    value_name = None
-    untyped_names: List[str] = []
-    for param_name, param_sig in parameters.items():
-        if lenient_issubclass(param_sig.annotation, click.Context):
-            ctx_name = param_name
-        elif lenient_issubclass(param_sig.annotation, click.Parameter):
-            click_param_name = param_name
-        else:
-            untyped_names.append(param_name)
-    # Extract value param name first
-    if untyped_names:
-        value_name = untyped_names.pop()
-    # If context and Click param were not typed (old/Click callback style) extract them
-    if untyped_names:
-        if ctx_name is None:
-            ctx_name = untyped_names.pop(0)
-        if click_param_name is None:
-            if untyped_names:
-                click_param_name = untyped_names.pop(0)
-        if untyped_names:
-            raise click.ClickException(
-                "Too many CLI parameter callback function parameters"
-            )
-
-    def wrapper(ctx: click.Context, param: click.Parameter, value: Any) -> Any:
-        use_params: Dict[str, Any] = {}
-        if ctx_name:
-            use_params[ctx_name] = ctx
-        if click_param_name:
-            use_params[click_param_name] = param
-        if value_name:
-            if convertor:
-                use_value = convertor(value)
-            else:
-                use_value = value
-            use_params[value_name] = use_value
-        return callback(**use_params)
-
-    update_wrapper(wrapper, callback)
-    return wrapper
-
-
-def get_param_completion(
-    callback: Optional[Callable[..., Any]] = None,
-) -> Optional[Callable[..., Any]]:
-    if not callback:
-        return None
-    parameters = get_params_from_function(callback)
-    ctx_name = None
-    args_name = None
-    incomplete_name = None
-    unassigned_params = list(parameters.values())
-    for param_sig in unassigned_params[:]:
-        origin = get_origin(param_sig.annotation)
-        if lenient_issubclass(param_sig.annotation, click.Context):
-            ctx_name = param_sig.name
-            unassigned_params.remove(param_sig)
-        elif lenient_issubclass(origin, List):
-            args_name = param_sig.name
-            unassigned_params.remove(param_sig)
-        elif lenient_issubclass(param_sig.annotation, str):
-            incomplete_name = param_sig.name
-            unassigned_params.remove(param_sig)
-    # If there are still unassigned parameters (not typed), extract by name
-    for param_sig in unassigned_params[:]:
-        if ctx_name is None and param_sig.name == "ctx":
-            ctx_name = param_sig.name
-            unassigned_params.remove(param_sig)
-        elif args_name is None and param_sig.name == "args":
-            args_name = param_sig.name
-            unassigned_params.remove(param_sig)
-        elif incomplete_name is None and param_sig.name == "incomplete":
-            incomplete_name = param_sig.name
-            unassigned_params.remove(param_sig)
-    # Extract value param name first
-    if unassigned_params:
-        show_params = " ".join([param.name for param in unassigned_params])
-        raise click.ClickException(
-            f"Invalid autocompletion callback parameters: {show_params}"
-        )
-
-    def wrapper(ctx: click.Context, args: List[str], incomplete: Optional[str]) -> Any:
-        use_params: Dict[str, Any] = {}
-        if ctx_name:
-            use_params[ctx_name] = ctx
-        if args_name:
-            use_params[args_name] = args
-        if incomplete_name:
-            use_params[incomplete_name] = incomplete
-        return callback(**use_params)
-
-    update_wrapper(wrapper, callback)
-    return wrapper
-
-
-def run(function: Callable[..., Any]) -> None:
-    app = Typer(add_completion=False)
-    app.command()(function)
-    app()
-
-
-def _is_macos() -> bool:
-    return platform.system() == "Darwin"
-
-
-def _is_linux_or_bsd() -> bool:
-    if platform.system() == "Linux":
-        return True
-
-    return "BSD" in platform.system()
-
-
-def launch(url: str, wait: bool = False, locate: bool = False) -> int:
-    """This function launches the given URL (or filename) in the default
-    viewer application for this file type.  If this is an executable, it
-    might launch the executable in a new session.  The return value is
-    the exit code of the launched application.  Usually, ``0`` indicates
-    success.
-
-    This function handles url in different operating systems separately:
-    - On macOS (Darwin), it uses the 'open' command.
-    - On Linux and BSD, it uses 'xdg-open' if available.
-    - On Windows (and other OSes), it uses the standard webbrowser module.
-
-    The function avoids, when possible, using the webbrowser module on Linux and macOS
-    to prevent spammy terminal messages from some browsers (e.g., Chrome).
-
-    Examples::
-
-        typer.launch("https://typer.tiangolo.com/")
-        typer.launch("/my/downloaded/file", locate=True)
-
-    :param url: URL or filename of the thing to launch.
-    :param wait: Wait for the program to exit before returning. This
-        only works if the launched program blocks. In particular,
-        ``xdg-open`` on Linux does not block.
-    :param locate: if this is set to `True` then instead of launching the
-                   application associated with the URL it will attempt to
-                   launch a file manager with the file located.  This
-                   might have weird effects if the URL does not point to
-                   the filesystem.
-    """
-
-    if url.startswith("http://") or url.startswith("https://"):
-        if _is_macos():
-            return subprocess.Popen(
-                ["open", url], stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT
-            ).wait()
-
-        has_xdg_open = _is_linux_or_bsd() and shutil.which("xdg-open") is not None
-
-        if has_xdg_open:
-            return subprocess.Popen(
-                ["xdg-open", url], stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT
-            ).wait()
-
-        import webbrowser
-
-        webbrowser.open(url)
-
-        return 0
-
-    else:
-        return click.launch(url)
diff --git a/venv/lib/python3.10/site-packages/typer/models.py b/venv/lib/python3.10/site-packages/typer/models.py
deleted file mode 100644
index e0bddb965be67ec2e8fd7c045a9e53baa887a915..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer/models.py
+++ /dev/null
@@ -1,544 +0,0 @@
-import inspect
-import io
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    List,
-    Optional,
-    Sequence,
-    Type,
-    TypeVar,
-    Union,
-)
-
-import click
-import click.shell_completion
-
-if TYPE_CHECKING:  # pragma: no cover
-    from .core import TyperCommand, TyperGroup
-    from .main import Typer
-
-
-NoneType = type(None)
-
-AnyType = Type[Any]
-
-Required = ...
-
-
-class Context(click.Context):
-    pass
-
-
-class FileText(io.TextIOWrapper):
-    pass
-
-
-class FileTextWrite(FileText):
-    pass
-
-
-class FileBinaryRead(io.BufferedReader):
-    pass
-
-
-class FileBinaryWrite(io.BufferedWriter):
-    pass
-
-
-class CallbackParam(click.Parameter):
-    pass
-
-
-class DefaultPlaceholder:
-    """
-    You shouldn't use this class directly.
-
-    It's used internally to recognize when a default value has been overwritten, even
-    if the new value is `None`.
-    """
-
-    def __init__(self, value: Any):
-        self.value = value
-
-    def __bool__(self) -> bool:
-        return bool(self.value)
-
-
-DefaultType = TypeVar("DefaultType")
-
-CommandFunctionType = TypeVar("CommandFunctionType", bound=Callable[..., Any])
-
-
-def Default(value: DefaultType) -> DefaultType:
-    """
-    You shouldn't use this function directly.
-
-    It's used internally to recognize when a default value has been overwritten, even
-    if the new value is `None`.
-    """
-    return DefaultPlaceholder(value)  # type: ignore
-
-
-class CommandInfo:
-    def __init__(
-        self,
-        name: Optional[str] = None,
-        *,
-        cls: Optional[Type["TyperCommand"]] = None,
-        context_settings: Optional[Dict[Any, Any]] = None,
-        callback: Optional[Callable[..., Any]] = None,
-        help: Optional[str] = None,
-        epilog: Optional[str] = None,
-        short_help: Optional[str] = None,
-        options_metavar: str = "[OPTIONS]",
-        add_help_option: bool = True,
-        no_args_is_help: bool = False,
-        hidden: bool = False,
-        deprecated: bool = False,
-        # Rich settings
-        rich_help_panel: Union[str, None] = None,
-    ):
-        self.name = name
-        self.cls = cls
-        self.context_settings = context_settings
-        self.callback = callback
-        self.help = help
-        self.epilog = epilog
-        self.short_help = short_help
-        self.options_metavar = options_metavar
-        self.add_help_option = add_help_option
-        self.no_args_is_help = no_args_is_help
-        self.hidden = hidden
-        self.deprecated = deprecated
-        # Rich settings
-        self.rich_help_panel = rich_help_panel
-
-
-class TyperInfo:
-    def __init__(
-        self,
-        typer_instance: Optional["Typer"] = Default(None),
-        *,
-        name: Optional[str] = Default(None),
-        cls: Optional[Type["TyperGroup"]] = Default(None),
-        invoke_without_command: bool = Default(False),
-        no_args_is_help: bool = Default(False),
-        subcommand_metavar: Optional[str] = Default(None),
-        chain: bool = Default(False),
-        result_callback: Optional[Callable[..., Any]] = Default(None),
-        # Command
-        context_settings: Optional[Dict[Any, Any]] = Default(None),
-        callback: Optional[Callable[..., Any]] = Default(None),
-        help: Optional[str] = Default(None),
-        epilog: Optional[str] = Default(None),
-        short_help: Optional[str] = Default(None),
-        options_metavar: str = Default("[OPTIONS]"),
-        add_help_option: bool = Default(True),
-        hidden: bool = Default(False),
-        deprecated: bool = Default(False),
-        # Rich settings
-        rich_help_panel: Union[str, None] = Default(None),
-    ):
-        self.typer_instance = typer_instance
-        self.name = name
-        self.cls = cls
-        self.invoke_without_command = invoke_without_command
-        self.no_args_is_help = no_args_is_help
-        self.subcommand_metavar = subcommand_metavar
-        self.chain = chain
-        self.result_callback = result_callback
-        self.context_settings = context_settings
-        self.callback = callback
-        self.help = help
-        self.epilog = epilog
-        self.short_help = short_help
-        self.options_metavar = options_metavar
-        self.add_help_option = add_help_option
-        self.hidden = hidden
-        self.deprecated = deprecated
-        self.rich_help_panel = rich_help_panel
-
-
-class ParameterInfo:
-    def __init__(
-        self,
-        *,
-        default: Optional[Any] = None,
-        param_decls: Optional[Sequence[str]] = None,
-        callback: Optional[Callable[..., Any]] = None,
-        metavar: Optional[str] = None,
-        expose_value: bool = True,
-        is_eager: bool = False,
-        envvar: Optional[Union[str, List[str]]] = None,
-        # Note that shell_complete is not fully supported and will be removed in future versions
-        # TODO: Remove shell_complete in a future version (after 0.16.0)
-        shell_complete: Optional[
-            Callable[
-                [click.Context, click.Parameter, str],
-                Union[List["click.shell_completion.CompletionItem"], List[str]],
-            ]
-        ] = None,
-        autocompletion: Optional[Callable[..., Any]] = None,
-        default_factory: Optional[Callable[[], Any]] = None,
-        # Custom type
-        parser: Optional[Callable[[str], Any]] = None,
-        click_type: Optional[click.ParamType] = None,
-        # TyperArgument
-        show_default: Union[bool, str] = True,
-        show_choices: bool = True,
-        show_envvar: bool = True,
-        help: Optional[str] = None,
-        hidden: bool = False,
-        # Choice
-        case_sensitive: bool = True,
-        # Numbers
-        min: Optional[Union[int, float]] = None,
-        max: Optional[Union[int, float]] = None,
-        clamp: bool = False,
-        # DateTime
-        formats: Optional[List[str]] = None,
-        # File
-        mode: Optional[str] = None,
-        encoding: Optional[str] = None,
-        errors: Optional[str] = "strict",
-        lazy: Optional[bool] = None,
-        atomic: bool = False,
-        # Path
-        exists: bool = False,
-        file_okay: bool = True,
-        dir_okay: bool = True,
-        writable: bool = False,
-        readable: bool = True,
-        resolve_path: bool = False,
-        allow_dash: bool = False,
-        path_type: Union[None, Type[str], Type[bytes]] = None,
-        # Rich settings
-        rich_help_panel: Union[str, None] = None,
-    ):
-        # Check if user has provided multiple custom parsers
-        if parser and click_type:
-            raise ValueError(
-                "Multiple custom type parsers provided. "
-                "`parser` and `click_type` may not both be provided."
-            )
-
-        self.default = default
-        self.param_decls = param_decls
-        self.callback = callback
-        self.metavar = metavar
-        self.expose_value = expose_value
-        self.is_eager = is_eager
-        self.envvar = envvar
-        self.shell_complete = shell_complete
-        self.autocompletion = autocompletion
-        self.default_factory = default_factory
-        # Custom type
-        self.parser = parser
-        self.click_type = click_type
-        # TyperArgument
-        self.show_default = show_default
-        self.show_choices = show_choices
-        self.show_envvar = show_envvar
-        self.help = help
-        self.hidden = hidden
-        # Choice
-        self.case_sensitive = case_sensitive
-        # Numbers
-        self.min = min
-        self.max = max
-        self.clamp = clamp
-        # DateTime
-        self.formats = formats
-        # File
-        self.mode = mode
-        self.encoding = encoding
-        self.errors = errors
-        self.lazy = lazy
-        self.atomic = atomic
-        # Path
-        self.exists = exists
-        self.file_okay = file_okay
-        self.dir_okay = dir_okay
-        self.writable = writable
-        self.readable = readable
-        self.resolve_path = resolve_path
-        self.allow_dash = allow_dash
-        self.path_type = path_type
-        # Rich settings
-        self.rich_help_panel = rich_help_panel
-
-
-class OptionInfo(ParameterInfo):
-    def __init__(
-        self,
-        *,
-        # ParameterInfo
-        default: Optional[Any] = None,
-        param_decls: Optional[Sequence[str]] = None,
-        callback: Optional[Callable[..., Any]] = None,
-        metavar: Optional[str] = None,
-        expose_value: bool = True,
-        is_eager: bool = False,
-        envvar: Optional[Union[str, List[str]]] = None,
-        # Note that shell_complete is not fully supported and will be removed in future versions
-        # TODO: Remove shell_complete in a future version (after 0.16.0)
-        shell_complete: Optional[
-            Callable[
-                [click.Context, click.Parameter, str],
-                Union[List["click.shell_completion.CompletionItem"], List[str]],
-            ]
-        ] = None,
-        autocompletion: Optional[Callable[..., Any]] = None,
-        default_factory: Optional[Callable[[], Any]] = None,
-        # Custom type
-        parser: Optional[Callable[[str], Any]] = None,
-        click_type: Optional[click.ParamType] = None,
-        # Option
-        show_default: Union[bool, str] = True,
-        prompt: Union[bool, str] = False,
-        confirmation_prompt: bool = False,
-        prompt_required: bool = True,
-        hide_input: bool = False,
-        # TODO: remove is_flag and flag_value in a future release
-        is_flag: Optional[bool] = None,
-        flag_value: Optional[Any] = None,
-        count: bool = False,
-        allow_from_autoenv: bool = True,
-        help: Optional[str] = None,
-        hidden: bool = False,
-        show_choices: bool = True,
-        show_envvar: bool = True,
-        # Choice
-        case_sensitive: bool = True,
-        # Numbers
-        min: Optional[Union[int, float]] = None,
-        max: Optional[Union[int, float]] = None,
-        clamp: bool = False,
-        # DateTime
-        formats: Optional[List[str]] = None,
-        # File
-        mode: Optional[str] = None,
-        encoding: Optional[str] = None,
-        errors: Optional[str] = "strict",
-        lazy: Optional[bool] = None,
-        atomic: bool = False,
-        # Path
-        exists: bool = False,
-        file_okay: bool = True,
-        dir_okay: bool = True,
-        writable: bool = False,
-        readable: bool = True,
-        resolve_path: bool = False,
-        allow_dash: bool = False,
-        path_type: Union[None, Type[str], Type[bytes]] = None,
-        # Rich settings
-        rich_help_panel: Union[str, None] = None,
-    ):
-        super().__init__(
-            default=default,
-            param_decls=param_decls,
-            callback=callback,
-            metavar=metavar,
-            expose_value=expose_value,
-            is_eager=is_eager,
-            envvar=envvar,
-            shell_complete=shell_complete,
-            autocompletion=autocompletion,
-            default_factory=default_factory,
-            # Custom type
-            parser=parser,
-            click_type=click_type,
-            # TyperArgument
-            show_default=show_default,
-            show_choices=show_choices,
-            show_envvar=show_envvar,
-            help=help,
-            hidden=hidden,
-            # Choice
-            case_sensitive=case_sensitive,
-            # Numbers
-            min=min,
-            max=max,
-            clamp=clamp,
-            # DateTime
-            formats=formats,
-            # File
-            mode=mode,
-            encoding=encoding,
-            errors=errors,
-            lazy=lazy,
-            atomic=atomic,
-            # Path
-            exists=exists,
-            file_okay=file_okay,
-            dir_okay=dir_okay,
-            writable=writable,
-            readable=readable,
-            resolve_path=resolve_path,
-            allow_dash=allow_dash,
-            path_type=path_type,
-            # Rich settings
-            rich_help_panel=rich_help_panel,
-        )
-        if is_flag is not None or flag_value is not None:
-            import warnings
-
-            warnings.warn(
-                "The 'is_flag' and 'flag_value' parameters are not supported by Typer "
-                "and will be removed entirely in a future release.",
-                DeprecationWarning,
-                stacklevel=2,
-            )
-        self.prompt = prompt
-        self.confirmation_prompt = confirmation_prompt
-        self.prompt_required = prompt_required
-        self.hide_input = hide_input
-        self.count = count
-        self.allow_from_autoenv = allow_from_autoenv
-
-
-class ArgumentInfo(ParameterInfo):
-    def __init__(
-        self,
-        *,
-        # ParameterInfo
-        default: Optional[Any] = None,
-        param_decls: Optional[Sequence[str]] = None,
-        callback: Optional[Callable[..., Any]] = None,
-        metavar: Optional[str] = None,
-        expose_value: bool = True,
-        is_eager: bool = False,
-        envvar: Optional[Union[str, List[str]]] = None,
-        # Note that shell_complete is not fully supported and will be removed in future versions
-        # TODO: Remove shell_complete in a future version (after 0.16.0)
-        shell_complete: Optional[
-            Callable[
-                [click.Context, click.Parameter, str],
-                Union[List["click.shell_completion.CompletionItem"], List[str]],
-            ]
-        ] = None,
-        autocompletion: Optional[Callable[..., Any]] = None,
-        default_factory: Optional[Callable[[], Any]] = None,
-        # Custom type
-        parser: Optional[Callable[[str], Any]] = None,
-        click_type: Optional[click.ParamType] = None,
-        # TyperArgument
-        show_default: Union[bool, str] = True,
-        show_choices: bool = True,
-        show_envvar: bool = True,
-        help: Optional[str] = None,
-        hidden: bool = False,
-        # Choice
-        case_sensitive: bool = True,
-        # Numbers
-        min: Optional[Union[int, float]] = None,
-        max: Optional[Union[int, float]] = None,
-        clamp: bool = False,
-        # DateTime
-        formats: Optional[List[str]] = None,
-        # File
-        mode: Optional[str] = None,
-        encoding: Optional[str] = None,
-        errors: Optional[str] = "strict",
-        lazy: Optional[bool] = None,
-        atomic: bool = False,
-        # Path
-        exists: bool = False,
-        file_okay: bool = True,
-        dir_okay: bool = True,
-        writable: bool = False,
-        readable: bool = True,
-        resolve_path: bool = False,
-        allow_dash: bool = False,
-        path_type: Union[None, Type[str], Type[bytes]] = None,
-        # Rich settings
-        rich_help_panel: Union[str, None] = None,
-    ):
-        super().__init__(
-            default=default,
-            param_decls=param_decls,
-            callback=callback,
-            metavar=metavar,
-            expose_value=expose_value,
-            is_eager=is_eager,
-            envvar=envvar,
-            shell_complete=shell_complete,
-            autocompletion=autocompletion,
-            default_factory=default_factory,
-            # Custom type
-            parser=parser,
-            click_type=click_type,
-            # TyperArgument
-            show_default=show_default,
-            show_choices=show_choices,
-            show_envvar=show_envvar,
-            help=help,
-            hidden=hidden,
-            # Choice
-            case_sensitive=case_sensitive,
-            # Numbers
-            min=min,
-            max=max,
-            clamp=clamp,
-            # DateTime
-            formats=formats,
-            # File
-            mode=mode,
-            encoding=encoding,
-            errors=errors,
-            lazy=lazy,
-            atomic=atomic,
-            # Path
-            exists=exists,
-            file_okay=file_okay,
-            dir_okay=dir_okay,
-            writable=writable,
-            readable=readable,
-            resolve_path=resolve_path,
-            allow_dash=allow_dash,
-            path_type=path_type,
-            # Rich settings
-            rich_help_panel=rich_help_panel,
-        )
-
-
-class ParamMeta:
-    empty = inspect.Parameter.empty
-
-    def __init__(
-        self,
-        *,
-        name: str,
-        default: Any = inspect.Parameter.empty,
-        annotation: Any = inspect.Parameter.empty,
-    ) -> None:
-        self.name = name
-        self.default = default
-        self.annotation = annotation
-
-
-class DeveloperExceptionConfig:
-    def __init__(
-        self,
-        *,
-        pretty_exceptions_enable: bool = True,
-        pretty_exceptions_show_locals: bool = True,
-        pretty_exceptions_short: bool = True,
-    ) -> None:
-        self.pretty_exceptions_enable = pretty_exceptions_enable
-        self.pretty_exceptions_show_locals = pretty_exceptions_show_locals
-        self.pretty_exceptions_short = pretty_exceptions_short
-
-
-class TyperPath(click.Path):
-    # Overwrite Click's behaviour to be compatible with Typer's autocompletion system
-    def shell_complete(
-        self, ctx: click.Context, param: click.Parameter, incomplete: str
-    ) -> List[click.shell_completion.CompletionItem]:
-        """Return an empty list so that the autocompletion functionality
-        will work properly from the commandline.
-        """
-        return []
diff --git a/venv/lib/python3.10/site-packages/typer/params.py b/venv/lib/python3.10/site-packages/typer/params.py
deleted file mode 100644
index 66c2b32d3e35e313454ed3dcbaac8ac9bf71d14d..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer/params.py
+++ /dev/null
@@ -1,479 +0,0 @@
-from typing import TYPE_CHECKING, Any, Callable, List, Optional, Type, Union, overload
-
-import click
-
-from .models import ArgumentInfo, OptionInfo
-
-if TYPE_CHECKING:  # pragma: no cover
-    import click.shell_completion
-
-
-# Overload for Option created with custom type 'parser'
-@overload
-def Option(
-    # Parameter
-    default: Optional[Any] = ...,
-    *param_decls: str,
-    callback: Optional[Callable[..., Any]] = None,
-    metavar: Optional[str] = None,
-    expose_value: bool = True,
-    is_eager: bool = False,
-    envvar: Optional[Union[str, List[str]]] = None,
-    # Note that shell_complete is not fully supported and will be removed in future versions
-    # TODO: Remove shell_complete in a future version (after 0.16.0)
-    shell_complete: Optional[
-        Callable[
-            [click.Context, click.Parameter, str],
-            Union[List["click.shell_completion.CompletionItem"], List[str]],
-        ]
-    ] = None,
-    autocompletion: Optional[Callable[..., Any]] = None,
-    default_factory: Optional[Callable[[], Any]] = None,
-    # Custom type
-    parser: Optional[Callable[[str], Any]] = None,
-    # Option
-    show_default: Union[bool, str] = True,
-    prompt: Union[bool, str] = False,
-    confirmation_prompt: bool = False,
-    prompt_required: bool = True,
-    hide_input: bool = False,
-    # TODO: remove is_flag and flag_value in a future release
-    is_flag: Optional[bool] = None,
-    flag_value: Optional[Any] = None,
-    count: bool = False,
-    allow_from_autoenv: bool = True,
-    help: Optional[str] = None,
-    hidden: bool = False,
-    show_choices: bool = True,
-    show_envvar: bool = True,
-    # Choice
-    case_sensitive: bool = True,
-    # Numbers
-    min: Optional[Union[int, float]] = None,
-    max: Optional[Union[int, float]] = None,
-    clamp: bool = False,
-    # DateTime
-    formats: Optional[List[str]] = None,
-    # File
-    mode: Optional[str] = None,
-    encoding: Optional[str] = None,
-    errors: Optional[str] = "strict",
-    lazy: Optional[bool] = None,
-    atomic: bool = False,
-    # Path
-    exists: bool = False,
-    file_okay: bool = True,
-    dir_okay: bool = True,
-    writable: bool = False,
-    readable: bool = True,
-    resolve_path: bool = False,
-    allow_dash: bool = False,
-    path_type: Union[None, Type[str], Type[bytes]] = None,
-    # Rich settings
-    rich_help_panel: Union[str, None] = None,
-) -> Any: ...
-
-
-# Overload for Option created with custom type 'click_type'
-@overload
-def Option(
-    # Parameter
-    default: Optional[Any] = ...,
-    *param_decls: str,
-    callback: Optional[Callable[..., Any]] = None,
-    metavar: Optional[str] = None,
-    expose_value: bool = True,
-    is_eager: bool = False,
-    envvar: Optional[Union[str, List[str]]] = None,
-    # Note that shell_complete is not fully supported and will be removed in future versions
-    # TODO: Remove shell_complete in a future version (after 0.16.0)
-    shell_complete: Optional[
-        Callable[
-            [click.Context, click.Parameter, str],
-            Union[List["click.shell_completion.CompletionItem"], List[str]],
-        ]
-    ] = None,
-    autocompletion: Optional[Callable[..., Any]] = None,
-    default_factory: Optional[Callable[[], Any]] = None,
-    # Custom type
-    click_type: Optional[click.ParamType] = None,
-    # Option
-    show_default: Union[bool, str] = True,
-    prompt: Union[bool, str] = False,
-    confirmation_prompt: bool = False,
-    prompt_required: bool = True,
-    hide_input: bool = False,
-    # TODO: remove is_flag and flag_value in a future release
-    is_flag: Optional[bool] = None,
-    flag_value: Optional[Any] = None,
-    count: bool = False,
-    allow_from_autoenv: bool = True,
-    help: Optional[str] = None,
-    hidden: bool = False,
-    show_choices: bool = True,
-    show_envvar: bool = True,
-    # Choice
-    case_sensitive: bool = True,
-    # Numbers
-    min: Optional[Union[int, float]] = None,
-    max: Optional[Union[int, float]] = None,
-    clamp: bool = False,
-    # DateTime
-    formats: Optional[List[str]] = None,
-    # File
-    mode: Optional[str] = None,
-    encoding: Optional[str] = None,
-    errors: Optional[str] = "strict",
-    lazy: Optional[bool] = None,
-    atomic: bool = False,
-    # Path
-    exists: bool = False,
-    file_okay: bool = True,
-    dir_okay: bool = True,
-    writable: bool = False,
-    readable: bool = True,
-    resolve_path: bool = False,
-    allow_dash: bool = False,
-    path_type: Union[None, Type[str], Type[bytes]] = None,
-    # Rich settings
-    rich_help_panel: Union[str, None] = None,
-) -> Any: ...
-
-
-def Option(
-    # Parameter
-    default: Optional[Any] = ...,
-    *param_decls: str,
-    callback: Optional[Callable[..., Any]] = None,
-    metavar: Optional[str] = None,
-    expose_value: bool = True,
-    is_eager: bool = False,
-    envvar: Optional[Union[str, List[str]]] = None,
-    # Note that shell_complete is not fully supported and will be removed in future versions
-    # TODO: Remove shell_complete in a future version (after 0.16.0)
-    shell_complete: Optional[
-        Callable[
-            [click.Context, click.Parameter, str],
-            Union[List["click.shell_completion.CompletionItem"], List[str]],
-        ]
-    ] = None,
-    autocompletion: Optional[Callable[..., Any]] = None,
-    default_factory: Optional[Callable[[], Any]] = None,
-    # Custom type
-    parser: Optional[Callable[[str], Any]] = None,
-    click_type: Optional[click.ParamType] = None,
-    # Option
-    show_default: Union[bool, str] = True,
-    prompt: Union[bool, str] = False,
-    confirmation_prompt: bool = False,
-    prompt_required: bool = True,
-    hide_input: bool = False,
-    # TODO: remove is_flag and flag_value in a future release
-    is_flag: Optional[bool] = None,
-    flag_value: Optional[Any] = None,
-    count: bool = False,
-    allow_from_autoenv: bool = True,
-    help: Optional[str] = None,
-    hidden: bool = False,
-    show_choices: bool = True,
-    show_envvar: bool = True,
-    # Choice
-    case_sensitive: bool = True,
-    # Numbers
-    min: Optional[Union[int, float]] = None,
-    max: Optional[Union[int, float]] = None,
-    clamp: bool = False,
-    # DateTime
-    formats: Optional[List[str]] = None,
-    # File
-    mode: Optional[str] = None,
-    encoding: Optional[str] = None,
-    errors: Optional[str] = "strict",
-    lazy: Optional[bool] = None,
-    atomic: bool = False,
-    # Path
-    exists: bool = False,
-    file_okay: bool = True,
-    dir_okay: bool = True,
-    writable: bool = False,
-    readable: bool = True,
-    resolve_path: bool = False,
-    allow_dash: bool = False,
-    path_type: Union[None, Type[str], Type[bytes]] = None,
-    # Rich settings
-    rich_help_panel: Union[str, None] = None,
-) -> Any:
-    return OptionInfo(
-        # Parameter
-        default=default,
-        param_decls=param_decls,
-        callback=callback,
-        metavar=metavar,
-        expose_value=expose_value,
-        is_eager=is_eager,
-        envvar=envvar,
-        shell_complete=shell_complete,
-        autocompletion=autocompletion,
-        default_factory=default_factory,
-        # Custom type
-        parser=parser,
-        click_type=click_type,
-        # Option
-        show_default=show_default,
-        prompt=prompt,
-        confirmation_prompt=confirmation_prompt,
-        prompt_required=prompt_required,
-        hide_input=hide_input,
-        is_flag=is_flag,
-        flag_value=flag_value,
-        count=count,
-        allow_from_autoenv=allow_from_autoenv,
-        help=help,
-        hidden=hidden,
-        show_choices=show_choices,
-        show_envvar=show_envvar,
-        # Choice
-        case_sensitive=case_sensitive,
-        # Numbers
-        min=min,
-        max=max,
-        clamp=clamp,
-        # DateTime
-        formats=formats,
-        # File
-        mode=mode,
-        encoding=encoding,
-        errors=errors,
-        lazy=lazy,
-        atomic=atomic,
-        # Path
-        exists=exists,
-        file_okay=file_okay,
-        dir_okay=dir_okay,
-        writable=writable,
-        readable=readable,
-        resolve_path=resolve_path,
-        allow_dash=allow_dash,
-        path_type=path_type,
-        # Rich settings
-        rich_help_panel=rich_help_panel,
-    )
-
-
-# Overload for Argument created with custom type 'parser'
-@overload
-def Argument(
-    # Parameter
-    default: Optional[Any] = ...,
-    *,
-    callback: Optional[Callable[..., Any]] = None,
-    metavar: Optional[str] = None,
-    expose_value: bool = True,
-    is_eager: bool = False,
-    envvar: Optional[Union[str, List[str]]] = None,
-    # Note that shell_complete is not fully supported and will be removed in future versions
-    # TODO: Remove shell_complete in a future version (after 0.16.0)
-    shell_complete: Optional[
-        Callable[
-            [click.Context, click.Parameter, str],
-            Union[List["click.shell_completion.CompletionItem"], List[str]],
-        ]
-    ] = None,
-    autocompletion: Optional[Callable[..., Any]] = None,
-    default_factory: Optional[Callable[[], Any]] = None,
-    # Custom type
-    parser: Optional[Callable[[str], Any]] = None,
-    # TyperArgument
-    show_default: Union[bool, str] = True,
-    show_choices: bool = True,
-    show_envvar: bool = True,
-    help: Optional[str] = None,
-    hidden: bool = False,
-    # Choice
-    case_sensitive: bool = True,
-    # Numbers
-    min: Optional[Union[int, float]] = None,
-    max: Optional[Union[int, float]] = None,
-    clamp: bool = False,
-    # DateTime
-    formats: Optional[List[str]] = None,
-    # File
-    mode: Optional[str] = None,
-    encoding: Optional[str] = None,
-    errors: Optional[str] = "strict",
-    lazy: Optional[bool] = None,
-    atomic: bool = False,
-    # Path
-    exists: bool = False,
-    file_okay: bool = True,
-    dir_okay: bool = True,
-    writable: bool = False,
-    readable: bool = True,
-    resolve_path: bool = False,
-    allow_dash: bool = False,
-    path_type: Union[None, Type[str], Type[bytes]] = None,
-    # Rich settings
-    rich_help_panel: Union[str, None] = None,
-) -> Any: ...
-
-
-# Overload for Argument created with custom type 'click_type'
-@overload
-def Argument(
-    # Parameter
-    default: Optional[Any] = ...,
-    *,
-    callback: Optional[Callable[..., Any]] = None,
-    metavar: Optional[str] = None,
-    expose_value: bool = True,
-    is_eager: bool = False,
-    envvar: Optional[Union[str, List[str]]] = None,
-    # Note that shell_complete is not fully supported and will be removed in future versions
-    # TODO: Remove shell_complete in a future version (after 0.16.0)
-    shell_complete: Optional[
-        Callable[
-            [click.Context, click.Parameter, str],
-            Union[List["click.shell_completion.CompletionItem"], List[str]],
-        ]
-    ] = None,
-    autocompletion: Optional[Callable[..., Any]] = None,
-    default_factory: Optional[Callable[[], Any]] = None,
-    # Custom type
-    click_type: Optional[click.ParamType] = None,
-    # TyperArgument
-    show_default: Union[bool, str] = True,
-    show_choices: bool = True,
-    show_envvar: bool = True,
-    help: Optional[str] = None,
-    hidden: bool = False,
-    # Choice
-    case_sensitive: bool = True,
-    # Numbers
-    min: Optional[Union[int, float]] = None,
-    max: Optional[Union[int, float]] = None,
-    clamp: bool = False,
-    # DateTime
-    formats: Optional[List[str]] = None,
-    # File
-    mode: Optional[str] = None,
-    encoding: Optional[str] = None,
-    errors: Optional[str] = "strict",
-    lazy: Optional[bool] = None,
-    atomic: bool = False,
-    # Path
-    exists: bool = False,
-    file_okay: bool = True,
-    dir_okay: bool = True,
-    writable: bool = False,
-    readable: bool = True,
-    resolve_path: bool = False,
-    allow_dash: bool = False,
-    path_type: Union[None, Type[str], Type[bytes]] = None,
-    # Rich settings
-    rich_help_panel: Union[str, None] = None,
-) -> Any: ...
-
-
-def Argument(
-    # Parameter
-    default: Optional[Any] = ...,
-    *,
-    callback: Optional[Callable[..., Any]] = None,
-    metavar: Optional[str] = None,
-    expose_value: bool = True,
-    is_eager: bool = False,
-    envvar: Optional[Union[str, List[str]]] = None,
-    # Note that shell_complete is not fully supported and will be removed in future versions
-    # TODO: Remove shell_complete in a future version (after 0.16.0)
-    shell_complete: Optional[
-        Callable[
-            [click.Context, click.Parameter, str],
-            Union[List["click.shell_completion.CompletionItem"], List[str]],
-        ]
-    ] = None,
-    autocompletion: Optional[Callable[..., Any]] = None,
-    default_factory: Optional[Callable[[], Any]] = None,
-    # Custom type
-    parser: Optional[Callable[[str], Any]] = None,
-    click_type: Optional[click.ParamType] = None,
-    # TyperArgument
-    show_default: Union[bool, str] = True,
-    show_choices: bool = True,
-    show_envvar: bool = True,
-    help: Optional[str] = None,
-    hidden: bool = False,
-    # Choice
-    case_sensitive: bool = True,
-    # Numbers
-    min: Optional[Union[int, float]] = None,
-    max: Optional[Union[int, float]] = None,
-    clamp: bool = False,
-    # DateTime
-    formats: Optional[List[str]] = None,
-    # File
-    mode: Optional[str] = None,
-    encoding: Optional[str] = None,
-    errors: Optional[str] = "strict",
-    lazy: Optional[bool] = None,
-    atomic: bool = False,
-    # Path
-    exists: bool = False,
-    file_okay: bool = True,
-    dir_okay: bool = True,
-    writable: bool = False,
-    readable: bool = True,
-    resolve_path: bool = False,
-    allow_dash: bool = False,
-    path_type: Union[None, Type[str], Type[bytes]] = None,
-    # Rich settings
-    rich_help_panel: Union[str, None] = None,
-) -> Any:
-    return ArgumentInfo(
-        # Parameter
-        default=default,
-        # Arguments can only have one param declaration
-        # it will be generated from the param name
-        param_decls=None,
-        callback=callback,
-        metavar=metavar,
-        expose_value=expose_value,
-        is_eager=is_eager,
-        envvar=envvar,
-        shell_complete=shell_complete,
-        autocompletion=autocompletion,
-        default_factory=default_factory,
-        # Custom type
-        parser=parser,
-        click_type=click_type,
-        # TyperArgument
-        show_default=show_default,
-        show_choices=show_choices,
-        show_envvar=show_envvar,
-        help=help,
-        hidden=hidden,
-        # Choice
-        case_sensitive=case_sensitive,
-        # Numbers
-        min=min,
-        max=max,
-        clamp=clamp,
-        # DateTime
-        formats=formats,
-        # File
-        mode=mode,
-        encoding=encoding,
-        errors=errors,
-        lazy=lazy,
-        atomic=atomic,
-        # Path
-        exists=exists,
-        file_okay=file_okay,
-        dir_okay=dir_okay,
-        writable=writable,
-        readable=readable,
-        resolve_path=resolve_path,
-        allow_dash=allow_dash,
-        path_type=path_type,
-        # Rich settings
-        rich_help_panel=rich_help_panel,
-    )
diff --git a/venv/lib/python3.10/site-packages/typer/py.typed b/venv/lib/python3.10/site-packages/typer/py.typed
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/venv/lib/python3.10/site-packages/typer/rich_utils.py b/venv/lib/python3.10/site-packages/typer/rich_utils.py
deleted file mode 100644
index 4b6c5a840f2aa9b52c788ea998d906e196c4224a..0000000000000000000000000000000000000000
--- a/venv/lib/python3.10/site-packages/typer/rich_utils.py
+++ /dev/null
@@ -1,741 +0,0 @@
-# Extracted and modified from https://github.com/ewels/rich-click
-
-import inspect
-import io
-import sys
-from collections import defaultdict
-from gettext import gettext as _
-from os import getenv
-from typing import Any, DefaultDict, Dict, Iterable, List, Optional, Union
-
-import click
-from rich import box
-from rich.align import Align
-from rich.columns import Columns
-from rich.console import Console, RenderableType, group
-from rich.emoji import Emoji
-from rich.highlighter import RegexHighlighter
-from rich.markdown import Markdown
-from rich.padding import Padding
-from rich.panel import Panel
-from rich.table import Table
-from rich.text import Text
-from rich.theme import Theme
-
-if sys.version_info >= (3, 9):
-    from typing import Literal
-else:
-    from typing_extensions import Literal
-
-# Default styles
-STYLE_OPTION = "bold cyan"
-STYLE_SWITCH = "bold green"
-STYLE_NEGATIVE_OPTION = "bold magenta"
-STYLE_NEGATIVE_SWITCH = "bold red"
-STYLE_METAVAR = "bold yellow"
-STYLE_METAVAR_SEPARATOR = "dim"
-STYLE_USAGE = "yellow"
-STYLE_USAGE_COMMAND = "bold"
-STYLE_DEPRECATED = "red"
-STYLE_DEPRECATED_COMMAND = "dim"
-STYLE_HELPTEXT_FIRST_LINE = ""
-STYLE_HELPTEXT = "dim"
-STYLE_OPTION_HELP = ""
-STYLE_OPTION_DEFAULT = "dim"
-STYLE_OPTION_ENVVAR = "dim yellow"
-STYLE_REQUIRED_SHORT = "red"
-STYLE_REQUIRED_LONG = "dim red"
-STYLE_OPTIONS_PANEL_BORDER = "dim"
-ALIGN_OPTIONS_PANEL: Literal["left", "center", "right"] = "left"
-STYLE_OPTIONS_TABLE_SHOW_LINES = False
-STYLE_OPTIONS_TABLE_LEADING = 0
-STYLE_OPTIONS_TABLE_PAD_EDGE = False
-STYLE_OPTIONS_TABLE_PADDING = (0, 1)
-STYLE_OPTIONS_TABLE_BOX = ""
-STYLE_OPTIONS_TABLE_ROW_STYLES = None
-STYLE_OPTIONS_TABLE_BORDER_STYLE = None
-STYLE_COMMANDS_PANEL_BORDER = "dim"
-ALIGN_COMMANDS_PANEL: Literal["left", "center", "right"] = "left"
-STYLE_COMMANDS_TABLE_SHOW_LINES = False
-STYLE_COMMANDS_TABLE_LEADING = 0
-STYLE_COMMANDS_TABLE_PAD_EDGE = False
-STYLE_COMMANDS_TABLE_PADDING = (0, 1)
-STYLE_COMMANDS_TABLE_BOX = ""
-STYLE_COMMANDS_TABLE_ROW_STYLES = None
-STYLE_COMMANDS_TABLE_BORDER_STYLE = None
-STYLE_COMMANDS_TABLE_FIRST_COLUMN = "bold cyan"
-STYLE_ERRORS_PANEL_BORDER = "red"
-ALIGN_ERRORS_PANEL: Literal["left", "center", "right"] = "left"
-STYLE_ERRORS_SUGGESTION = "dim"
-STYLE_ABORTED = "red"
-_TERMINAL_WIDTH = getenv("TERMINAL_WIDTH")
-MAX_WIDTH = int(_TERMINAL_WIDTH) if _TERMINAL_WIDTH else None
-COLOR_SYSTEM: Optional[Literal["auto", "standard", "256", "truecolor", "windows"]] = (
-    "auto"  # Set to None to disable colors
-)
-_TYPER_FORCE_DISABLE_TERMINAL = getenv("_TYPER_FORCE_DISABLE_TERMINAL")
-FORCE_TERMINAL = (
-    True
-    if getenv("GITHUB_ACTIONS") or getenv("FORCE_COLOR") or getenv("PY_COLORS")
-    else None
-)
-if _TYPER_FORCE_DISABLE_TERMINAL:
-    FORCE_TERMINAL = False
-
-# Fixed strings
-DEPRECATED_STRING = _("(deprecated) ")
-DEFAULT_STRING = _("[default: {}]")
-ENVVAR_STRING = _("[env var: {}]")
-REQUIRED_SHORT_STRING = "*"
-REQUIRED_LONG_STRING = _("[required]")
-RANGE_STRING = " [{}]"
-ARGUMENTS_PANEL_TITLE = _("Arguments")
-OPTIONS_PANEL_TITLE = _("Options")
-COMMANDS_PANEL_TITLE = _("Commands")
-ERRORS_PANEL_TITLE = _("Error")
-ABORTED_TEXT = _("Aborted.")
-RICH_HELP = _("Try [blue]'{command_path} {help_option}'[/] for help.")
-
-MARKUP_MODE_MARKDOWN = "markdown"
-MARKUP_MODE_RICH = "rich"
-_RICH_HELP_PANEL_NAME = "rich_help_panel"
-
-MarkupMode = Literal["markdown", "rich", None]
-
-
-# Rich regex highlighter
-class OptionHighlighter(RegexHighlighter):
-    """Highlights our special options."""
-
-    highlights = [
-        r"(^|\W)(?P\-\w+)(?![a-zA-Z0-9])",
-        r"(^|\W)(?P

' : '\U0001d4ab', - '\\' : '\U0001d4ac', - '\\' : '\U0000211b', - '\\' : '\U0001d4ae', - '\\' : '\U0001d4af', - '\\' : '\U0001d4b0', - '\\' : '\U0001d4b1', - '\\' : '\U0001d4b2', - '\\' : '\U0001d4b3', - '\\' : '\U0001d4b4', - '\\' : '\U0001d4b5', - '\\' : '\U0001d5ba', - '\\' : '\U0001d5bb', - '\\' : '\U0001d5bc', - '\\' : '\U0001d5bd', - '\\' : '\U0001d5be', - '\\' : '\U0001d5bf', - '\\' : '\U0001d5c0', - '\\' : '\U0001d5c1', - '\\' : '\U0001d5c2', - '\\' : '\U0001d5c3', - '\\' : '\U0001d5c4', - '\\' : '\U0001d5c5', - '\\' : '\U0001d5c6', - '\\' : '\U0001d5c7', - '\\' : '\U0001d5c8', - '\\